{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Risø Conference Proceedings 2026 - grazing-incidence case\n", "This notebook contains the accompanying processing code for the grazing-indidence case to support the Risø Conference Proceedings 2026 paper:\n", "\n", "Ball, J. A. D., Andreasen, J. W., Angelis, S. D., Wright, J. P., & Detlefs, C. (2026, July 9). Multi-Beam 3DXRD. IOP Conference Series: Materials Science and Engineering. 46th Risø International Symposium on Materials Science: Characterization of evolving microstructures in metals, DTU Risø Campus, Roskilde, Denmark. Accepted for publication." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "# don't hog GPU memory (if you have one)\n", "os.environ[\"XLA_PYTHON_CLIENT_PREALLOCATE\"] = \"0\"\n", "\n", "import jax\n", "jax.config.update(\"jax_enable_x64\", True) # required for accuracy, especially strains\n", "import jax.numpy as jnp\n", "import optax\n", "\n", "import time\n", "from functools import partial\n", "\n", "import scipy\n", "from scipy.spatial.transform import Rotation\n", "from matplotlib import pyplot as plt\n", "\n", "import ImageD11.grain\n", "import ImageD11.unitcell\n", "import ImageD11.indexing\n", "from xfab.symmetry import ROTATIONS, Umis\n", "\n", "from ImageD11.nbGui.nb_utils import plot_grain_positions, plot_all_ipfs\n", "\n", "import anri.crystal, anri.diffract, anri.geom, anri.fwd\n", "\n", "start = time.time()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "### BEAM\n", "\n", "# Energies in eV\n", "beam_ev = 12_000\n", "\n", "def eV_to_wavelength_A(energy_eV):\n", " # L = hc/E\n", " # Must convert eV to Joules first\n", " wavelength_A = (scipy.constants.h * scipy.constants.c)/(scipy.constants.e*energy_eV)*1e10\n", " return wavelength_A\n", "\n", "wavelength = eV_to_wavelength_A(beam_ev)\n", "\n", "wavelength" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Crystallography\n", "We'll use cubic Si here" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "struc = anri.crystal.Structure.from_cif(\"../../../tests/data/cif/Si.cif\")\n", "struc.lattice_parameters" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We generate some HKLs." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "scrolled": true }, "outputs": [], "source": [ "struc.make_hkls(dsmax=1.0, wavelength=wavelength)\n", "struc.rings_table" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Phantom sample\n", "\n", "Let's define our phantom sample. \n", "We define everything in the sample coordinate system (on top of the goniometer). \n", "We'll randomly generate some grains within a box." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "rng = 54 # chosen by fair dice roll, guaranteed to be random\n", "key = jax.random.key(rng)\n", "\n", "n_grains = 100\n", "\n", "### Positions\n", "sample_width = 1000.0\n", "sample_height = 10.0\n", "translations_sample_xy = jax.random.uniform(key, shape=(n_grains,2,), minval=-sample_width/2, maxval=sample_width/2)\n", "translations_sample_z = jax.random.uniform(key, shape=(n_grains,), minval=-sample_height/2, maxval=sample_height/2)\n", "translations_sample = jnp.column_stack((translations_sample_xy, translations_sample_z))\n", "\n", "### Orientations\n", "U_matrices = Rotation.random(n_grains, rng=rng).as_matrix()\n", "UB_matrices = U_matrices @ struc.B\n", "UBI_matrices = jnp.linalg.inv(UB_matrices)\n", "\n", "### Volumes - just for visualisation purposes!\n", "radii_sigma = 0.2\n", "radii_mean = 50.0\n", "radii = jax.random.lognormal(key, shape=(n_grains,), sigma=radii_sigma) * radii_mean\n", "volumes = (4./3)*jnp.pi*(radii**3)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We can use ImageD11 to plot this phantom." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "ref_unitcell = ImageD11.unitcell.unitcell(struc.lattice_parameters, symmetry=struc.sgno)\n", "grains = [ImageD11.grain.grain(UBI_matrices[i], translation=translations_sample[i]) for i in range(n_grains)]\n", "\n", "for i, g in enumerate(grains):\n", " g.ref_unitcell = ref_unitcell\n", " g.intensity_info = f\"mean = {volumes[i]}\"\n", "\n", "plot_grain_positions(grains, 'z', size_scaling=0.1)\n", "plot_all_ipfs(grains)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Forward projection\n", "Now we can generate our scattering vectors and translate them into the lab frame, then into the detector. \n", "This will yield centroids, which are arrays of `[slow, fast, omega]` which represent the centre-of-mass positions of the peaks on the detector surface. \n", "We must define our goniometer and detector positions:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "### Goniometer\n", "\n", "wedge = -0.5 # anri now right-handed!\n", "chi = 0.0\n", "y0 = 0.0\n", "\n", "### Detector\n", "y_center = 1024.0\n", "z_center = -1024.0\n", "y_size = 50.0\n", "z_size = 50.0\n", "tilt_x = 0.0\n", "tilt_y = jnp.radians(-60.0)\n", "tilt_z = 0.0\n", "distance = 180e3\n", "o11 = 1\n", "o12 = 0\n", "o21 = 0\n", "o22 = 1\n", "\n", "# not official parameters but useful for plotting later\n", "det_size_s = 2048\n", "det_size_f = 2048\n", "\n", "# Get change-of-basis parameters to go from lab to detector space:\n", "det_trans, beam_cen_shift, x_distance_shift = anri.geom.detector_transforms(\n", " y_center,\n", " y_size,\n", " tilt_y,\n", " z_center,\n", " z_size,\n", " tilt_z,\n", " tilt_x,\n", " distance,\n", " o11,\n", " o12,\n", " o21,\n", " o22\n", ")\n", "\n", "# Get detector unit vectors (slow, fast directions) in lab frame:\n", "sc_lab, fc_lab, norm_lab = anri.geom.detector_basis_vectors_lab(det_trans, beam_cen_shift, x_distance_shift)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# we get plus/minus Friedel pairs\n", "# we also get boolean masks for validity - sometimes no solution exists for the Ewald condition, so we never see the peak.\n", "\n", "k_in_direct = jnp.array([1., 0, 0]) # direct beam wavevector\n", "k_in_refl = anri.geom.rot_y(2*wedge) @ k_in_direct # reflected beam wavevector\n", "\n", "\n", "centroid_dir_p, valid_dir_p = anri.fwd.get_centroid_box_all(UBI_matrices, translations_sample,\n", " struc.ringhkls_arr,\n", " 1.0, wavelength, k_in_direct, 0, 0,\n", " wedge, chi,\n", " sc_lab, fc_lab, norm_lab)\n", "centroid_dir_m, valid_dir_m = anri.fwd.get_centroid_box_all(UBI_matrices, translations_sample,\n", " struc.ringhkls_arr,\n", " -1.0, wavelength, k_in_direct, 0, 0,\n", " wedge, chi,\n", " sc_lab, fc_lab, norm_lab)\n", "\n", "centroid_refl_p, valid_refl_p = anri.fwd.get_centroid_box_all(UBI_matrices, translations_sample,\n", " struc.ringhkls_arr,\n", " 1.0, wavelength, k_in_refl, 0, 0,\n", " wedge, chi,\n", " sc_lab, fc_lab, norm_lab)\n", "centroid_refl_m, valid_refl_m = anri.fwd.get_centroid_box_all(UBI_matrices, translations_sample,\n", " struc.ringhkls_arr,\n", " -1.0, wavelength, k_in_refl, 0, 0,\n", " wedge, chi,\n", " sc_lab, fc_lab, norm_lab)\n", "\n", "# join Friedel pairs into single datasets:\n", "centroid_dir = jnp.concatenate([centroid_dir_p[valid_dir_p], centroid_dir_m[valid_dir_m]])\n", "centroid_refl = jnp.concatenate([centroid_refl_p[valid_refl_p], centroid_refl_m[valid_refl_m]])\n", "\n", "# flatten into (N, 3)\n", "centroid_dir = centroid_dir.reshape(-1, 3)\n", "centroid_refl = centroid_refl.reshape(-1, 3)\n", "\n", "# mask to detector pixel range\n", "m_dir = (centroid_dir[:, 0] > 0) & (centroid_dir[:, 0] < det_size_f) & (centroid_dir[:, 1] > 0) & (centroid_dir[:, 1] < det_size_s)\n", "m_refl = (centroid_refl[:, 0] > 0) & (centroid_refl[:, 0] < det_size_f) & (centroid_refl[:, 1] > 0) & (centroid_refl[:, 1] < det_size_s)\n", "centroid_dir = centroid_dir[m_dir]\n", "centroid_refl = centroid_refl[m_refl]" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, axs = plt.subplots(1, 2, figsize=(12,20), constrained_layout=True)\n", "axs[0].scatter(centroid_dir[:, 1], centroid_dir[:, 0], label=r'Direct peaks', s=2)\n", "axs[0].scatter(centroid_refl[:, 1], centroid_refl[:, 0], label=r'Reflected peaks', s=2)\n", "axs[0].set_aspect(1)\n", "axs[0].set(xlabel='Detector fast', ylabel='Detector slow', title='Whole detector view', xlim=(0, det_size_f), ylim=(0, det_size_s))\n", "axs[0].legend(loc='upper right')\n", "\n", "axs[1].scatter(centroid_dir[:, 1], centroid_dir[:, 0], label=r'Direct peaks')\n", "axs[1].scatter(centroid_refl[:, 1], centroid_refl[:, 0], label=r'Reflected peaks')\n", "axs[1].set_aspect(1)\n", "axs[1].set(xlabel='Detector fast', ylabel='Detector slow', xlim=(300, 450), ylim=(850, 1000), title='Detail view')\n", "axs[1].legend(loc='upper right')\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Scattering vector identification\n", "With the peaks forward-projected onto the detector, we now 'forget' which wavelength each peak came from. \n", "We compute scattering vectors in the sample frame for *all* peaks, twice, assuming each wavelength. \n", "We also forget the origin of diffraction of each centroid." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "centroid = jnp.concatenate([centroid_dir, centroid_refl])" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Simulating experimental error\n", "We now make a sensible effort to \"spoil\" the measured centroids to account for experimental errors. \n", "This can be done more accurately (and will be in the future when we simulate intensity profiles on the detector) but to first order this should be a reasonable approach." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def spoil_centroids(centroids, width_px, width_omega):\n", " npks = centroids.shape[0]\n", " px_error_vector = jax.random.uniform(key, shape=(npks,2), minval=-width_px/2, maxval=width_px/2)\n", " omega_error_vector = jax.random.uniform(key, shape=(npks,), minval=-width_omega/2, maxval=width_omega/2)\n", " error_vector = jnp.column_stack((px_error_vector, omega_error_vector))\n", " centroids = centroids + error_vector\n", " return centroids\n", "\n", "ostep = 0.1 # realistic omega step\n", "centroid_with_error = spoil_centroids(centroid, 1.0, ostep)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, axs = plt.subplots(1, 2, figsize=(12,20), constrained_layout=True)\n", "axs[0].scatter(centroid[:, 1], centroid[:, 0], label=r'Peaks without error')\n", "axs[0].scatter(centroid_with_error[:, 1], centroid_with_error[:, 0], s=10, label=r'Peaks with error')\n", "axs[0].set_aspect(1)\n", "axs[0].set(xlabel='Detector fast', ylabel='Detector slow', title='Whole detector view')\n", "axs[0].legend(loc='upper right')\n", "\n", "axs[1].scatter(centroid[:, 1], centroid[:, 0], label=r'Peaks without error')\n", "axs[1].scatter(centroid_with_error[:, 1], centroid_with_error[:, 0], s=10, label=r'Peaks with error')\n", "axs[1].set_aspect(1)\n", "axs[1].set(xlabel='Detector fast', ylabel='Detector slow', xlim=(300, 450), ylim=(850, 1000), title='Detail view')\n", "axs[1].legend(loc='upper right')\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# All these primitive functions are written for single vectors\n", "# We write the overall function for a single vector, then vmap it over many vectors for speed:\n", "\n", "@jax.jit\n", "def detector_to_q(slow, fast, omega, wavelength, k_in_lab_hat, origin_sample):\n", " # peak vector in lab frame\n", " peak_lab = anri.geom.det_to_lab(slow, fast, det_trans, beam_cen_shift, x_distance_shift)\n", " # rotate origin into sample frame\n", " origin_lab = anri.geom.sample_to_lab(origin_sample, omega, wedge, chi, 0.0, 0.0)\n", " \n", " # convert peak vector to k_out, subtracts off the origin\n", " k_out_lab_norm = anri.diffract.peak_lab_to_k_out(peak_lab, origin_lab, wavelength)\n", " # normalise k_in by wavelength\n", " k_in_lab_norm = anri.diffract.scale_norm_k(k_in_lab_hat, wavelength)\n", " # simply q = k_out - k_in\n", " q_lab = anri.diffract.k_to_q_lab(k_in_lab_norm, k_out_lab_norm)\n", " # rotate into sample frame\n", " q_sample = anri.geom.lab_to_sample(q_lab, omega, wedge, chi, 0.0, 0.0)\n", " return q_sample\n", "\n", "# the vmap operation\n", "detector_to_q_vec = jax.vmap(detector_to_q, in_axes=(0, 0, 0, 0, 0, None))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Compute scattering vectors for all peaks assuming each beam direction\n", "\n", "kvecs_direct = jnp.broadcast_to(k_in_direct, centroid_with_error.shape)\n", "kvecs_refl = jnp.broadcast_to(k_in_refl, centroid_with_error.shape)\n", "wavelengths = jnp.full(centroid_with_error.shape[0], wavelength)\n", "\n", "q_sample_direct = detector_to_q_vec(centroid_with_error[:, 0], centroid_with_error[:, 1], centroid_with_error[:, 2], wavelengths, kvecs_direct, jnp.array([0., 0., 0.]))\n", "q_sample_refl = detector_to_q_vec(centroid_with_error[:, 0], centroid_with_error[:, 1], centroid_with_error[:, 2], wavelengths, kvecs_refl, jnp.array([0., 0., 0.]))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## *k*-d tree search\n", "We now construct a *k*-d tree in sample space. \n", "The idea is that half of the observed scattering vectors will be correctly computed with the direct beam k_in vector, and the other half will be correctly computed with the reflected beam k_in vector. \n", "As these duplicated scattering vectors come from a single set of reciprocal lattice points, the two datasets should 'overlap' if and only if the incoming wavevector was correctly chosen for a given scattering vector.\n", "We look for 'overlapping' (i.e. duplicate) scattering vectors in sample space using a *k*-d tree." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "kd_refl = scipy.spatial.cKDTree(q_sample_refl)\n", "\n", "# Find the pairs\n", "distances, indices = kd_refl.query(q_sample_direct, k=1, distance_upper_bound=0.01)\n", "valid_mask = jnp.isfinite(distances)\n", "valid_distances = distances[valid_mask]\n", "\n", "fig, ax = plt.subplots()\n", "ax.hist(valid_distances, bins=100)\n", "ax.set(xlabel='Distance in g-vector space', ylabel='Count', title='Histogram of g-vector neighbour distances')\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We can perhaps discern that a sensible cutoff would be `0.0015` to get the first spike" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Find the pairs\n", "distances, indices = kd_refl.query(q_sample_direct, k=1, distance_upper_bound=0.0015)\n", "valid_mask = jnp.isfinite(distances)\n", "valid_distances = distances[valid_mask]\n", "\n", "fig, ax = plt.subplots()\n", "ax.hist(valid_distances, bins=50)\n", "ax.set(xlabel='Distance in g-vector space', ylabel='Count', title='Histogram of g-vector neighbour distances')\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now we need to average or combine our observations of the g-vectors.\n", "\n", "We have three sources of g-vectors we can possibly index:\n", "\n", "- G-vectors only from $K_{\\alpha_{1}}$\n", "- G-vectors only from $K_{\\alpha_{2}}$\n", "- Averaged/combined observations of both" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "mask_dir = valid_mask\n", "mask_refl = indices[mask_dir]\n", "\n", "q_sample_direct_paired = q_sample_direct[mask_dir]\n", "q_sample_refl_paired = q_sample_refl[mask_refl]\n", "\n", "# concatenate\n", "q_sample_cat = jnp.concatenate((q_sample_direct_paired, q_sample_refl_paired))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We can confirm that the deduplication succeeded by masking the centroids array in the same way:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, axs = plt.subplots(1, 2, figsize=(12,20), constrained_layout=True)\n", "axs[0].scatter(centroid_with_error[:, 1], centroid_with_error[:, 0], s=50, label=r'Direct and reflected beam peaks')\n", "axs[0].scatter(centroid_with_error[:, 1][mask_dir], centroid_with_error[:, 0][mask_dir], s=10, label=r'Deduplicated peaks')\n", "axs[0].set_aspect(1)\n", "axs[0].set(xlabel='Detector fast', ylabel='Detector slow', title='Whole detector view', xlim=(0, det_size_f), ylim=(0, det_size_s))\n", "axs[0].legend(loc='upper right')\n", "\n", "axs[1].scatter(centroid_with_error[:, 1], centroid_with_error[:, 0], s=50, label=r'Direct and reflected beam peaks')\n", "axs[1].scatter(centroid_with_error[:, 1][mask_dir], centroid_with_error[:, 0][mask_dir], s=10, label=r'Deduplicated peaks')\n", "axs[1].set_aspect(1)\n", "axs[1].set(xlabel='Detector fast', ylabel='Detector slow', xlim=(300, 450), ylim=(850, 1000), title='Detail view')\n", "axs[1].legend(loc='upper right')\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Indexing the recovered scattering vectors\n", "### Just direct beam" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "scrolled": true }, "outputs": [], "source": [ "# make an indexer\n", "idx_dir = ImageD11.indexing.indexer(unitcell=ref_unitcell, gv=q_sample_direct_paired, minpks=15)\n", "idx_dir.ds_tol = 0.005\n", "idx_dir.assigntorings()\n", "idx_dir.hkl_tol = 0.025\n", "idx_dir.cosine_tol = 0.002\n", "idx_dir.score_all_pairs()\n", "idx_dir.saveubis('grains_found_direct.ubi')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Just reflected beam" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "scrolled": true }, "outputs": [], "source": [ "# make an indexer\n", "idx_refl = ImageD11.indexing.indexer(unitcell=ref_unitcell, gv=q_sample_refl_paired, minpks=15)\n", "idx_refl.ds_tol = 0.005\n", "idx_refl.assigntorings()\n", "idx_refl.hkl_tol = 0.025\n", "idx_refl.cosine_tol = 0.002\n", "idx_refl.score_all_pairs()\n", "idx_refl.saveubis('grains_found_refl.ubi')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Combined scattering vectors" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "scrolled": true }, "outputs": [], "source": [ "# make an indexer\n", "idx_both = ImageD11.indexing.indexer(unitcell=ref_unitcell, gv=q_sample_cat, minpks=30)\n", "idx_both.ds_tol = 0.005\n", "idx_both.assigntorings()\n", "idx_both.hkl_tol = 0.025\n", "idx_both.cosine_tol = 0.002\n", "idx_both.score_all_pairs()\n", "idx_both.saveubis('grains_found_both.ubi')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Grain position and UBI refinement\n", "We can now try to refine the positions and UBIs of the grains we've found. \n", "We do this in three ways: \n", "- Indexed with direct beam g-vectors, refined with direct beam centroids\n", "- Indexed with reflected beam g-vectors, refined with reflected beam centroids\n", "- Indexed with combined g-vectors, refined with both direct and reflected beam centroids (alternating)\n", "\n", "We implement a gradient-aware minimiser using ADAM from the Optax library. The loss function itself is identical to `makemap.py` in ImageD11." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Peak to grain assignment\n", "We have to establish static peak to grain assignments before we can refine, so grains do not compete for peaks. We do this just like ImageD11 - each peak is assigned to the grain that best indexes it (yields $hkl$ values closest to integer)." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "@jax.jit\n", "def assign(ubi, q_sample):\n", " hklf = (ubi @ q_sample.T).T\n", " hkli = jnp.rint(hklf)\n", " hkle = jnp.linalg.norm(hklf - hkli)\n", " return hkle, hkli\n", "\n", "assign_peaks = jax.vmap(assign, in_axes=[None, 0])\n", "assign_grains = jax.vmap(assign_peaks, in_axes=[0, None])\n", "\n", "# compute hkl errors\n", "hkle_dir, hkli_dir = assign_grains(jnp.array(idx_dir.ubis), q_sample_direct[mask_dir])\n", "hkle_refl, hkli_refl = assign_grains(jnp.array(idx_refl.ubis), q_sample_refl[mask_refl])\n", "hkle_both, hkli_both = assign_grains(jnp.array(idx_both.ubis), q_sample_cat)\n", "\n", "# get grain assignments (minimised error per peak)\n", "ga_dir = jnp.argmin(hkle_dir, axis=0)\n", "ga_refl = jnp.argmin(hkle_refl, axis=0)\n", "ga_both = jnp.argmin(hkle_both, axis=0)\n", "\n", "# get hkli per peak\n", "hkli_dir = jnp.squeeze(jnp.take_along_axis(hkli_dir, ga_dir[None, :, None], axis=0))\n", "hkli_refl = jnp.squeeze(jnp.take_along_axis(hkli_refl, ga_refl[None, :, None], axis=0))\n", "hkli_both = jnp.squeeze(jnp.take_along_axis(hkli_both, ga_both[None, :, None], axis=0))\n", "\n", "# max grains per peak\n", "M = max(jnp.unique(ga_dir, return_counts=True)[1].max(), jnp.unique(ga_refl, return_counts=True)[1].max(), jnp.unique(ga_both, return_counts=True)[1].max())\n", "M" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "@jax.jit\n", "def solve_ub_analytical(hkl_int, q_sample, mask):\n", " \"\"\"q = UB @ h => UB = (Q^T H)(H^T H)^-1, masked rows contributing zero.\"\"\"\n", " m = mask.astype(q_sample.dtype)[:, None]\n", "\n", " H = hkl_int * m\n", " Qs = q_sample * m\n", "\n", " HHT = H.T @ H # (3, 3), symmetric\n", " QsHT = Qs.T @ H\n", "\n", " # degenerate grains would otherwise NaN one row of the vmapped batch\n", " HHT = HHT + (1e-12 * jnp.trace(HHT) + 1e-15) * jnp.eye(3, dtype=HHT.dtype)\n", "\n", " return jnp.linalg.solve(HHT, QsHT.T).T\n", "\n", "\n", "@jax.jit\n", "def get_grain_loss_function(ubi, origin_sample, cen_obs, hkl_int, wavelength, kvecs, mask):\n", " # NOTE: `ubi` is deliberately unused — UB is profiled out analytically below.\n", " q_sample = detector_to_q_vec(cen_obs[:, 0], cen_obs[:, 1], cen_obs[:, 2],\n", " wavelength, kvecs, origin_sample)\n", "\n", " ub_fit = solve_ub_analytical(hkl_int, q_sample, mask)\n", " ubi_fit = jnp.linalg.inv(ub_fit)\n", "\n", " d = jnp.linalg.solve(ub_fit, q_sample.T).T - hkl_int\n", " hkle = jnp.sqrt(jnp.sum(d * d, axis=1) + 1e-24) # NaN-safe at zero residual\n", "\n", " mf = mask.astype(hkle.dtype)\n", " return jnp.sum(hkle * mf) / (jnp.sum(mf) + 1e-10) + 1e-10, ubi_fit\n", "\n", "\n", "@partial(jax.jit, static_argnums=(7,))\n", "def refine_grain(initial_ubi, initial_origin, cen_obs, hkl_int, wavelength, kvecs, mask,\n", " num_steps=100, learning_rate=1e-2, scaling_factor=1000.0):\n", "\n", " x0 = initial_origin / scaling_factor \n", "\n", " optimizer = optax.adam(learning_rate=learning_rate)\n", " opt_state0 = optimizer.init(x0)\n", "\n", " def objective(scaled_origin):\n", " loss, ubi = get_grain_loss_function(\n", " initial_ubi, scaled_origin * scaling_factor,\n", " cen_obs, hkl_int, wavelength, kvecs, mask)\n", " return loss, ubi\n", "\n", " grad_fn = jax.value_and_grad(objective, has_aux=True)\n", "\n", " def scan_body(carry, _):\n", " x, opt_state, best_x, best_loss = carry\n", "\n", " (loss, _), grads = grad_fn(x)\n", "\n", " better = loss < best_loss\n", " best_x = jnp.where(better, x, best_x)\n", " best_loss = jnp.where(better, loss, best_loss)\n", "\n", " updates, opt_state = optimizer.update(grads, opt_state)\n", " x = optax.apply_updates(x, updates)\n", "\n", " return (x, opt_state, best_x, best_loss), loss\n", "\n", " init_carry = (x0, opt_state0, x0, jnp.asarray(jnp.inf, x0.dtype))\n", " (_, _, best_x, _), loss_history = jax.lax.scan(\n", " scan_body, init_carry, None, length=num_steps)\n", "\n", " final_origin = best_x * scaling_factor\n", "\n", " final_loss, final_ubi = get_grain_loss_function(\n", " initial_ubi, final_origin, cen_obs, hkl_int, wavelength, kvecs, mask)\n", "\n", " loss_history = jnp.concatenate([loss_history, final_loss[None]])\n", "\n", " return final_ubi, final_origin, loss_history\n", "\n", "refine_vmap = jax.vmap(\n", " refine_grain, \n", " in_axes=(0, 0, 0, 0, 0, 0, 0, None, None, None)\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Just direct beam" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "scrolled": true }, "outputs": [], "source": [ "%%time\n", "\n", "all_masks_dir = jnp.array([ga_dir == gid for gid in range(len(idx_dir.ubis))])\n", "all_origins_dir = jnp.zeros((len(idx_dir.ubis), 3))\n", "\n", "def get_compressed_grain_data(grain_id, assignments, cen, hkl, wave, kvec):\n", " # Create a boolean mask for this specific grain\n", " grain_mask = (assignments == grain_id)\n", " \n", " # Use top_k to find the indices of the True values.\n", " # This identifies exactly which rows in the filtered arrays belong to this grain.\n", " # Since we want static shapes, we always take M indices.\n", " _, indices = jax.lax.top_k(grain_mask.astype(kvec.dtype), M)\n", " \n", " # Gather the data for this grain\n", " # If the grain has < M peaks, top_k will pad with the last found index, \n", " # but the 'mask' will correctly be False for those duplicates.\n", " return {\n", " \"cen\": cen[indices],\n", " \"hkl\": hkl[indices],\n", " \"wave\": wave[indices],\n", " \"kvec\": kvec[indices],\n", " \"mask\": grain_mask[indices]\n", " }\n", "\n", "# Vectorize the packing over all grain IDs\n", "v_pack = jax.vmap(\n", " get_compressed_grain_data, \n", " in_axes=(0, None, None, None, None, None)\n", ")\n", "\n", "# Pack the data into grain-specific buffers (num_grains, M, ...)\n", "grain_ids_dir = jnp.arange(len(idx_dir.ubis))\n", "compressed_dir = v_pack(grain_ids_dir, ga_dir, centroid_with_error[mask_dir], hkli_dir, wavelengths, kvecs_direct)\n", "# …likewise compressed_refl, compressed_both\n", "\n", "# Run the refinement\n", "# All inputs (except scalars) are now (num_grains, ...) so in_axes are all 0\n", "ubis_fit_dir, origins_sample_fit_dir, all_losses_dir = refine_vmap(\n", " jnp.array(idx_dir.ubis),\n", " all_origins_dir,\n", " compressed_dir[\"cen\"],\n", " compressed_dir[\"hkl\"],\n", " compressed_dir[\"wave\"],\n", " compressed_dir[\"kvec\"],\n", " compressed_dir[\"mask\"],\n", " 100,\n", " 1e-2,\n", " 1000.0\n", ")\n", "\n", "ubis_fit_dir.block_until_ready()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, axs = plt.subplots(2,1,sharex=True,sharey=True)\n", "axs[0].plot(all_losses_dir.T)\n", "axs[1].plot(all_losses_dir.mean(axis=0), label='dir')\n", "axs[0].set(yscale='log')\n", "axs[1].set(yscale='log')\n", "axs[0].set(title='Loss over all grains')\n", "axs[1].set(xlabel='Epoch', ylabel='Loss',title='Mean loss')\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Just reflected beam" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "scrolled": true }, "outputs": [], "source": [ "%%time\n", "\n", "all_masks_refl = jnp.array([ga_refl == gid for gid in range(len(idx_refl.ubis))])\n", "all_origins_refl = jnp.zeros((len(idx_refl.ubis), 3))\n", "\n", "# Pack the data into grain-specific buffers (num_grains, M, ...)\n", "grain_ids_refl = jnp.arange(len(idx_refl.ubis))\n", "compressed_refl = v_pack(grain_ids_refl, ga_refl, centroid_with_error[mask_refl], hkli_refl, wavelengths, kvecs_refl)\n", "\n", "# 2. Run the refinement\n", "# All inputs (except scalars) are now (num_grains, ...) so in_axes are all 0\n", "ubis_fit_refl, origins_sample_fit_refl, all_losses_refl = refine_vmap(\n", " jnp.array(idx_refl.ubis),\n", " all_origins_refl,\n", " compressed_refl[\"cen\"],\n", " compressed_refl[\"hkl\"],\n", " compressed_refl[\"wave\"],\n", " compressed_refl[\"kvec\"],\n", " compressed_refl[\"mask\"],\n", " 100,\n", " 1e-2,\n", " 1000.0\n", ")\n", "\n", "ubis_fit_refl.block_until_ready()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Both direct and reflected beams" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "scrolled": true }, "outputs": [], "source": [ "%%time\n", "\n", "all_masks_both = jnp.array([ga_both == gid for gid in range(len(idx_both.ubis))])\n", "all_origins_both = jnp.zeros((len(idx_both.ubis), 3))\n", "centroids_both = jnp.concatenate([centroid_with_error[mask_dir], centroid_with_error[mask_refl]])\n", "wavelengths_both = jnp.concatenate([wavelengths[mask_dir], wavelengths[mask_refl]])\n", "kvecs_both = jnp.concatenate([kvecs_direct[mask_dir], kvecs_refl[mask_refl]])\n", "\n", "# Pack the data into grain-specific buffers (num_grains, M, ...)\n", "grain_ids_both = jnp.arange(len(idx_both.ubis))\n", "compressed_both = v_pack(grain_ids_both, ga_both, centroids_both, hkli_both, wavelengths_both, kvecs_both)\n", "\n", "# Run the refinement\n", "ubis_fit_both, origins_sample_fit_both, all_losses_both = refine_vmap(\n", " jnp.array(idx_both.ubis),\n", " all_origins_both,\n", " compressed_both[\"cen\"],\n", " compressed_both[\"hkl\"],\n", " compressed_both[\"wave\"],\n", " compressed_both[\"kvec\"],\n", " compressed_both[\"mask\"],\n", " 100,\n", " 1e-2,\n", " 1000.0\n", ")\n", "\n", "ubis_fit_both.block_until_ready()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, ax = plt.subplots()\n", "ax.plot(all_losses_dir.mean(axis=0), label='Direct beam')\n", "ax.plot(all_losses_refl.mean(axis=0), label='Reflected beam')\n", "ax.plot(all_losses_both.mean(axis=0), label='Both beams')\n", "ax.set(yscale='log')\n", "ax.legend()\n", "ax.set(xlabel='Epoch', ylabel='loss', title='Loss over all grains')\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Match results to ground truth\n", "As the indexer returns UBIs in a different order, we need to match the refinement results to the ground-truth grains. \n", "First, we make grain lists from the refined results. \n", "Then, we map each of the grain lists (including the ground truth grains) into the fundamental zone by maximising traces. \n", "Then we look for matches with a 6D translation-orientation feature vector $k$-d tree. This is highly biased towards orientations because they should more more reliable than translations." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "SYM = jnp.array(ROTATIONS[7])\n", "\n", "def cast_to_fundamental_zone(U, symmetry_ops):\n", " best_U = U\n", " max_trace = jnp.trace(U)\n", " \n", " for S in symmetry_ops:\n", "\n", " U_equiv = jnp.matmul(U, S)\n", " \n", " current_trace = jnp.trace(U_equiv)\n", " \n", " if current_trace > max_trace:\n", " max_trace = current_trace\n", " best_U = U_equiv\n", " \n", " return best_U\n", "\n", "def move_grains_to_fz(gl, sym_ops):\n", " newgl = []\n", " for g in gl:\n", " best_U = cast_to_fundamental_zone(g.U, sym_ops)\n", " new_UB = best_U @ g.B\n", " new_UBI = jnp.linalg.inv(new_UB)\n", " newg = ImageD11.grain.grain(new_UBI, translation=g.translation)\n", " newg.ref_unitcell = ref_unitcell\n", " newgl.append(newg)\n", "\n", " return newgl\n", "\n", "def match_grains(grains1, grains2, W=50, dist_tol=100):\n", " \"\"\"Returns (distance, index into grains1, index into grains2), all same length.\"\"\"\n", " def desc(gl):\n", " return jnp.column_stack([\n", " jnp.array([jnp.asarray(g.translation) for g in gl]),\n", " W * jnp.array([Rotation.from_matrix(jnp.asarray(g.U)).as_rotvec() for g in gl])])\n", " d, m = scipy.spatial.cKDTree(desc(grains2)).query(desc(grains1), k=1,\n", " distance_upper_bound=dist_tol)\n", " valid = jnp.isfinite(d)\n", " return d[valid], jnp.nonzero(valid)[0], m[valid]\n", "\n", "@jax.jit\n", "def gt_ubi_aligned(U_gt, U_fit, B):\n", " \"\"\"GT UBI in the same symmetry setting as the fitted grain.\"\"\"\n", " M = U_fit.T @ U_gt\n", " S = SYM[jnp.argmax(jnp.trace(M @ SYM, axis1=1, axis2=2))]\n", " return jnp.linalg.inv((U_gt @ S) @ B)\n", "\n", "@jax.jit\n", "def loss_at_fixed_ubi(ubi, origin, cen, hkl, wave, kvec, mask):\n", " \"\"\"Mean |dhkl| for a GIVEN ubi — no analytical UB refit.\"\"\"\n", " q = detector_to_q_vec(cen[:, 0], cen[:, 1], cen[:, 2], wave, kvec, origin)\n", " d = (ubi @ q.T).T - hkl\n", " e = jnp.sqrt(jnp.sum(d * d, axis=1) + 1e-24)\n", " mf = mask.astype(e.dtype)\n", " return jnp.sum(e * mf) / (jnp.sum(mf) + 1e-10) + 1e-10\n", "\n", "gt_loss_vmap = jax.vmap(loss_at_fixed_ubi, in_axes=(0, 0, 0, 0, 0, 0, 0))\n", "\n", "# ---------- refined grain lists ----------\n", "\n", "grains_dir = [ImageD11.grain.grain(ubis_fit_dir[g], translation=origins_sample_fit_dir[g]) for g in range(len(idx_dir.ubis))]\n", "grains_refl = [ImageD11.grain.grain(ubis_fit_refl[g], translation=origins_sample_fit_refl[g]) for g in range(len(idx_refl.ubis))]\n", "grains_both = [ImageD11.grain.grain(ubis_fit_both[g], translation=origins_sample_fit_both[g]) for g in range(len(idx_both.ubis))]\n", "\n", "oriens_dir = jnp.stack([g.U for g in grains_dir])\n", "oriens_refl = jnp.stack([g.U for g in grains_refl])\n", "oriens_both = jnp.stack([g.U for g in grains_both])\n", "\n", "# ---------- match ----------\n", "\n", "grains_fz = move_grains_to_fz(grains, ROTATIONS[7])\n", "grains_dir_fz = move_grains_to_fz(grains_dir, ROTATIONS[7])\n", "grains_refl_fz = move_grains_to_fz(grains_refl, ROTATIONS[7])\n", "grains_both_fz = move_grains_to_fz(grains_both, ROTATIONS[7])\n", "\n", "dist_dir, fit_dir, gt_dir = match_grains(grains_dir_fz, grains_fz)\n", "dist_refl, fit_refl, gt_refl = match_grains(grains_refl_fz, grains_fz)\n", "dist_both, fit_both, gt_both = match_grains(grains_both_fz, grains_fz)\n", "\n", "for tag, fit, gt, n in [('a1', fit_dir, gt_dir, len(grains_dir)),\n", " ('a2', fit_refl, gt_refl, len(grains_refl)),\n", " ('both', fit_both, gt_both, len(grains_both))]:\n", " print(f\"{tag}: {len(fit)}/{n} matched, {len(gt) - len(jnp.unique(gt))} GT grains claimed twice\")\n", "\n", "# ---------- ground-truth losses ----------\n", "\n", "def gt_losses(gt_idx, fit_idx, oriens, packed):\n", " ubi_gt = jax.vmap(gt_ubi_aligned, in_axes=(0, 0, None))(\n", " U_matrices[gt_idx], oriens[fit_idx], struc.B)\n", " return gt_loss_vmap(ubi_gt, translations_sample[gt_idx],\n", " packed[\"cen\"][fit_idx], packed[\"hkl\"][fit_idx],\n", " packed[\"wave\"][fit_idx], packed[\"kvec\"][fit_idx],\n", " packed[\"mask\"][fit_idx])\n", "\n", "gt_losses_dir = gt_losses(gt_dir, fit_dir, oriens_dir, compressed_dir)\n", "gt_losses_refl = gt_losses(gt_refl, fit_refl, oriens_refl, compressed_refl)\n", "gt_losses_both = gt_losses(gt_both, fit_both, oriens_both, compressed_both)\n", "\n", "# ---------- metrics ----------\n", "\n", "pos_diff_dir = jnp.linalg.norm(translations_sample[gt_dir] - origins_sample_fit_dir[fit_dir], axis=1)\n", "pos_diff_refl = jnp.linalg.norm(translations_sample[gt_refl] - origins_sample_fit_refl[fit_refl], axis=1)\n", "pos_diff_both = jnp.linalg.norm(translations_sample[gt_both] - origins_sample_fit_both[fit_both], axis=1)\n", "\n", "misorien_dir = jnp.array([jnp.min(Umis(U_matrices[g], oriens_dir[f], 7)[:, 1]) for g, f in zip(gt_dir, fit_dir)])\n", "misorien_refl = jnp.array([jnp.min(Umis(U_matrices[g], oriens_refl[f], 7)[:, 1]) for g, f in zip(gt_refl, fit_refl)])\n", "misorien_both = jnp.array([jnp.min(Umis(U_matrices[g], oriens_both[f], 7)[:, 1]) for g, f in zip(gt_both, fit_both)])\n", "\n", "strains_dir = jnp.array([grains_dir[f].eps_sample_matrix(struc.lattice_parameters) for f in fit_dir])\n", "strains_refl = jnp.array([grains_refl[f].eps_sample_matrix(struc.lattice_parameters) for f in fit_refl])\n", "strains_both = jnp.array([grains_both[f].eps_sample_matrix(struc.lattice_parameters) for f in fit_both])\n", "\n", "strains_norm_dir = jnp.sqrt((strains_dir**2).sum(axis=(1,2)))\n", "strains_norm_refl = jnp.sqrt((strains_refl**2).sum(axis=(1,2)))\n", "strains_norm_both = jnp.sqrt((strains_both**2).sum(axis=(1,2)))\n", "\n", "loss_diff_dir = all_losses_dir[fit_dir, -1] - gt_losses_dir\n", "loss_diff_refl = all_losses_refl[fit_refl, -1] - gt_losses_refl\n", "loss_diff_both = all_losses_both[fit_both, -1] - gt_losses_both\n", "\n", "print(\"refined:\", float(all_losses_dir[fit_dir, -1].mean()),\n", " \" ground truth:\", float(gt_losses_dir.mean()))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, axs = plt.subplots(3,1,sharex=True,sharey=True,constrained_layout=True)\n", "axs[0].hist(pos_diff_dir, bins=20)\n", "axs[1].hist(pos_diff_refl, bins=20)\n", "axs[2].hist(pos_diff_both, bins=20)\n", "axs[0].set_title(r'Just direct beam')\n", "axs[1].set_title(r'Just reflected beam')\n", "axs[2].set_title(r'Direct and reflected beams')\n", "fig.supxlabel(r'Position error to ground truth (μm)')\n", "fig.supylabel('Counts')\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, axs = plt.subplots(3,1,sharex=True,sharey=True,constrained_layout=True)\n", "axs[0].hist(strains_norm_dir*1e3, bins=20)\n", "axs[1].hist(strains_norm_refl*1e3, bins=20)\n", "axs[2].hist(strains_norm_both*1e3, bins=20)\n", "axs[0].set_title(r'Just direct beam')\n", "axs[1].set_title(r'Just reflected beam')\n", "axs[2].set_title(r'Direct and reflected beams')\n", "fig.supxlabel(r'Strain norm (should be 0) (x1e-3)')\n", "fig.supylabel('Counts')\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, axs = plt.subplots(3,1,sharex=True,sharey=True,constrained_layout=True)\n", "axs[0].hist(misorien_dir, bins=20)\n", "axs[1].hist(misorien_refl, bins=20)\n", "axs[2].hist(misorien_both, bins=20)\n", "axs[0].set_title(r'Just direct beam')\n", "axs[1].set_title(r'Just reflected beam')\n", "axs[2].set_title(r'Direct and reflected beams')\n", "fig.supxlabel('Misorientation to ground truth (deg)')\n", "fig.supylabel('Counts')\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, axs = plt.subplots(3,1,sharex=True,sharey=True,constrained_layout=True)\n", "axs[0].hist(all_losses_dir[:, -1], bins=20, alpha=0.5)\n", "axs[0].hist(gt_losses_dir, bins=20, alpha=0.5)\n", "axs[1].hist(all_losses_refl[:, -1], bins=20, alpha=0.5)\n", "axs[1].hist(gt_losses_refl, bins=20, alpha=0.5)\n", "axs[2].hist(all_losses_both[:, -1], bins=20,alpha=0.5)\n", "axs[2].hist(gt_losses_both, bins=20, alpha=0.5)\n", "axs[0].set_title(r'Just direct beam')\n", "axs[1].set_title(r'Just reflected beam')\n", "axs[2].set_title(r'Direct and reflected beams')\n", "fig.supxlabel('Final loss function')\n", "fig.supylabel('Counts')\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, axs = plt.subplots(3,1,sharex=True,sharey=True,constrained_layout=True)\n", "axs[0].hist(loss_diff_dir, bins=20)\n", "axs[1].hist(loss_diff_refl, bins=20)\n", "axs[2].hist(loss_diff_both, bins=20)\n", "axs[0].set_title(r'Just direct beam')\n", "axs[1].set_title(r'Just reflected beam')\n", "axs[2].set_title(r'Direct and reflected beams')\n", "fig.supxlabel('Difference in loss function to ground truth')\n", "fig.supylabel('Counts')\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def print_grain_metrics(\n", " pos_diff_dir, pos_diff_refl, pos_diff_both,\n", " misorien_dir, misorien_refl, misorien_both,\n", " strains_norm_dir, strains_norm_refl, strains_norm_both,\n", " loss_diff_dir, loss_diff_refl, loss_diff_both\n", "):\n", " # Defining metrics with their specific scaling factors\n", " # (Label, datdir, datrefl, databoth, scale_factor, unit_label)\n", " metrics = [\n", " (\"Loss Diff\", loss_diff_dir, loss_diff_refl, loss_diff_both, 1e6, \"(x1e-6)\"),\n", " (\"Positional Diff\", pos_diff_dir, pos_diff_refl, pos_diff_both, 1.0, \"\"),\n", " (\"Misorientation\", misorien_dir, misorien_refl, misorien_both, 1e3, \"(x1e-3)\"),\n", " (\"Strain Norm\", strains_norm_dir, strains_norm_refl, strains_norm_both, 1e4, \"(x1e-4)\"),\n", " ]\n", "\n", " header = f\"{'Metric':<25} | {'dir Mean':<12} | {'refl Mean':<12} | {'Both Mean':<12} | {'Comb. Frac':<12}\"\n", " print(header)\n", " print(\"-\" * len(header))\n", "\n", " for label, d1, d2, db, scale, unit in metrics:\n", " # Apply scaling and calculate means\n", " m1 = jnp.mean(d1).item() * scale\n", " m2 = jnp.mean(d2).item() * scale\n", " mb = jnp.mean(db).item() * scale\n", " \n", " m_avg_indiv = (m1 + m2) / 2\n", " \n", " # The fraction remains the same regardless of scaling\n", " fraction = mb / m_avg_indiv if m_avg_indiv != 0 else 0\n", " \n", " full_label = f\"{label} {unit}\".strip()\n", " print(f\"{full_label:<25} | {m1:>12.6f} | {m2:>12.6f} | {mb:>12.6f} | {fraction:>12.4f}\")\n", "\n", "print_grain_metrics(\n", " pos_diff_dir, pos_diff_refl, pos_diff_both,\n", " misorien_dir, misorien_refl, misorien_both,\n", " strains_norm_dir, strains_norm_refl, strains_norm_both,\n", " loss_diff_dir, loss_diff_refl, loss_diff_both\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# total peaks (dir and refl)\n", "mask_dir.sum(), mask_dir.sum()*2, mask_dir.sum()/n_grains, mask_dir.sum()*2/n_grains" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Plots for paper" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, ax = plt.subplots(figsize=(3.54, 3.54), layout='constrained')\n", "ax.scatter(centroid_with_error[:mask_dir.sum()][:, 1], centroid_with_error[:mask_dir.sum()][:, 0], label=r'Direct peaks', s=.5)\n", "ax.scatter(centroid_with_error[mask_dir.sum():][:, 1], centroid_with_error[mask_dir.sum():][:, 0], label=r'Reflected peaks', s=.5)\n", "ax.set_aspect(1)\n", "ax.set(title='Grazing incidence')\n", "ax.set_xticks([])\n", "ax.set_yticks([])\n", "ax.legend(loc='upper right')\n", "plt.show()\n", "plt.savefig('graz_peaks.png', dpi=600)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "end = time.time()\n", "print(f'Took {end-start:.0f} seconds')" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (main)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.9" }, "widgets": { "application/vnd.jupyter.widget-state+json": { "state": {}, "version_major": 2, "version_minor": 0 } } }, "nbformat": 4, "nbformat_minor": 4 }