{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Risø Conference Proceedings 2026 - lab source - LM optimizer\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.\n", "\n", "This notebook is very similar to the general lab case, but we use a Levenberg–Marquardt approach to optimize much faster." ] }, { "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", "\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": "markdown", "metadata": {}, "source": [ "## Constants" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "### BEAM\n", "\n", "# Energies in eV\n", "# https://xdb.lbl.gov/Section1/Table_1-2.pdf\n", "k_alpha_1_eV = 24209.7 \n", "k_alpha_2_eV = 24002.0\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", "k_alpha_1_A = eV_to_wavelength_A(k_alpha_1_eV)\n", "k_alpha_2_A = eV_to_wavelength_A(k_alpha_2_eV)\n", "\n", "k_alpha_1_A, k_alpha_2_A" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Crystallography\n", "We'll start with HCP Ti" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "struc = anri.crystal.Structure.from_cif(\"../../../tests/data/cif/Ti.cif\")\n", "struc.lattice_parameters" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We generate some HKLs. This needs a single wavelength supplied to compute the $2\\theta$ angles for the reflections table, but it won't be directly used here." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "struc.make_hkls(dsmax=1.0, wavelength=k_alpha_1_A)\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 = 28 # chosen by fair dice roll, guaranteed to be random\n", "key = jax.random.key(rng)\n", "\n", "n_grains = 450\n", "\n", "### Positions\n", "sample_width = 1000.0\n", "translations_sample = jax.random.uniform(key, shape=(n_grains,3,), minval=-sample_width/2, maxval=sample_width/2)\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 = 100.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.0\n", "chi = 0.0\n", "y0 = 0.0\n", "\n", "### Detector\n", "y_center = 2048.0\n", "z_center = 2048.0\n", "y_size = 100.0\n", "z_size = 100.0\n", "tilt_x = 0.0\n", "tilt_y = 0.0\n", "tilt_z = 0.0\n", "distance = 515e3\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 = 4096\n", "det_size_f = 4096\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", "# unit beam wavevector, lab frame\n", "k_in_lab_hat = jnp.array([1., 0, 0,])\n", "\n", "centroid_a1_p, valid_a1_p = anri.fwd.get_centroid_box_all(UBI_matrices, translations_sample,\n", " struc.ringhkls_arr,\n", " 1.0, k_alpha_1_A, k_in_lab_hat, 0, 0,\n", " wedge, chi,\n", " sc_lab, fc_lab, norm_lab)\n", "centroid_a1_m, valid_a1_m = anri.fwd.get_centroid_box_all(UBI_matrices, translations_sample,\n", " struc.ringhkls_arr,\n", " -1.0, k_alpha_1_A, k_in_lab_hat, 0, 0,\n", " wedge, chi,\n", " sc_lab, fc_lab, norm_lab)\n", "\n", "centroid_a2_p, valid_a2_p = anri.fwd.get_centroid_box_all(UBI_matrices, translations_sample,\n", " struc.ringhkls_arr,\n", " 1.0, k_alpha_2_A, k_in_lab_hat, 0, 0,\n", " wedge, chi,\n", " sc_lab, fc_lab, norm_lab)\n", "centroid_a2_m, valid_a2_m = anri.fwd.get_centroid_box_all(UBI_matrices, translations_sample,\n", " struc.ringhkls_arr,\n", " -1.0, k_alpha_2_A, k_in_lab_hat, 0, 0,\n", " wedge, chi,\n", " sc_lab, fc_lab, norm_lab)\n", "\n", "# join Friedel pairs into single datasets:\n", "centroid_a1 = jnp.concatenate([centroid_a1_p[valid_a1_p], centroid_a1_m[valid_a1_m]])\n", "centroid_a2 = jnp.concatenate([centroid_a2_p[valid_a2_p], centroid_a2_m[valid_a2_m]])\n", "\n", "# flatten into (N, 3)\n", "centroid_a1 = centroid_a1.reshape(-1, 3)\n", "centroid_a2 = centroid_a2.reshape(-1, 3)\n", "\n", "# mask to detector pixel range\n", "m_a1 = (centroid_a1[:, 0] > 0) & (centroid_a1[:, 0] < det_size_s) & (centroid_a1[:, 1] > 0) & (centroid_a1[:, 1] < det_size_f)\n", "m_a2 = (centroid_a2[:, 0] > 0) & (centroid_a2[:, 0] < det_size_s) & (centroid_a2[:, 1] > 0) & (centroid_a2[:, 1] < det_size_f)\n", "centroid_a1 = centroid_a1[m_a1]\n", "centroid_a2 = centroid_a2[m_a2]" ] }, { "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_a1[:, 1], centroid_a1[:, 0], label=r'$K_{\\alpha_{1}}$ peaks', s=2)\n", "axs[0].scatter(centroid_a2[:, 1], centroid_a2[:, 0], label=r'$K_{\\alpha_{2}}$ peaks', s=2)\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", "axs[1].scatter(centroid_a1[:, 1], centroid_a1[:, 0], label=r'$K_{\\alpha_{1}}$ peaks', s=2)\n", "axs[1].scatter(centroid_a2[:, 1], centroid_a2[:, 0], label=r'$K_{\\alpha_{2}}$ peaks', s=2)\n", "axs[1].set_aspect(1)\n", "axs[1].set(xlabel='Detector fast', ylabel='Detector slow', xlim=(890, 940), ylim=(890, 940), 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_a1, centroid_a2])" ] }, { "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", "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=(890, 940), ylim=(890, 940), title='Detail view')\n", "axs[1].legend(loc='upper right')\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Backward projection\n", "We can now project backwards from detector space to scattering vectors in sample space $\\mathbf{g_s}$" ] }, { "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 wavelength\n", "\n", "kvecs = jnp.broadcast_to(k_in_lab_hat, centroid_with_error.shape)\n", "wavelengths_a1 = jnp.full(centroid_with_error.shape[0], k_alpha_1_A)\n", "wavelengths_a2 = jnp.full(centroid_with_error.shape[0], k_alpha_2_A)\n", "\n", "q_sample_a1 = detector_to_q_vec(centroid_with_error[:, 0], centroid_with_error[:, 1], centroid_with_error[:, 2], wavelengths_a1, kvecs, jnp.array([0., 0., 0.]))\n", "q_sample_a2 = detector_to_q_vec(centroid_with_error[:, 0], centroid_with_error[:, 1], centroid_with_error[:, 2], wavelengths_a2, kvecs, 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 $\\lambda = K_{\\alpha_{1}}$, and the other half will be correctly computed with $\\lambda = K_{\\alpha_{2}}$. \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 wavelength 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.\n", "\n", "In practice, we first do this via a sparse distance matrix with a large distance tolerance, to determine where the sensible cutoff should be:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "kd_a2 = scipy.spatial.cKDTree(q_sample_a2)\n", "\n", "# Find the pairs\n", "distances, indices = kd_a2.query(q_sample_a1, k=1, distance_upper_bound=0.03)\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_a2.query(q_sample_a1, 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=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": [ "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_a1 = valid_mask\n", "mask_a2 = indices[mask_a1]\n", "\n", "q_sample_a1_paired = q_sample_a1[mask_a1]\n", "q_sample_a2_paired = q_sample_a2[mask_a2]\n", "\n", "# concatenate\n", "q_sample_cat = jnp.concatenate((q_sample_a1_paired, q_sample_a2_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'$K_{\\alpha_{1}}$ and $K_{\\alpha_{1}}$ peaks')\n", "axs[0].scatter(centroid_with_error[:, 1][mask_a1], centroid_with_error[:, 0][mask_a1], 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')\n", "axs[0].legend(loc='upper right')\n", "\n", "axs[1].scatter(centroid_with_error[:, 1], centroid_with_error[:, 0], s=50, label=r'$K_{\\alpha_{1}}$ and $K_{\\alpha_{1}}$ peaks')\n", "axs[1].scatter(centroid_with_error[:, 1][mask_a1], centroid_with_error[:, 0][mask_a1], s=10, label=r'Deduplicated peaks')\n", "axs[1].set_aspect(1)\n", "axs[1].set(xlabel='Detector fast', ylabel='Detector slow', xlim=(890, 940), ylim=(890, 940), 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 $K_{\\alpha_{1}}$" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "scrolled": true }, "outputs": [], "source": [ "# make an indexer\n", "idx_a1 = ImageD11.indexing.indexer(unitcell=ref_unitcell, gv=q_sample_a1_paired, minpks=80)\n", "idx_a1.ds_tol = 0.005\n", "idx_a1.assigntorings()\n", "idx_a1.hkl_tol = 0.02\n", "idx_a1.cosine_tol = 0.002\n", "idx_a1.score_all_pairs()\n", "idx_a1.saveubis('grains_found_a1.ubi')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Just $K_{\\alpha_{2}}$" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "scrolled": true }, "outputs": [], "source": [ "# make an indexer\n", "idx_a2 = ImageD11.indexing.indexer(unitcell=ref_unitcell, gv=q_sample_a2_paired, minpks=80)\n", "idx_a2.ds_tol = 0.005\n", "idx_a2.assigntorings()\n", "idx_a2.hkl_tol = 0.02\n", "idx_a2.cosine_tol = 0.002\n", "idx_a2.score_all_pairs()\n", "idx_a2.saveubis('grains_found_a2.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=160)\n", "idx_both.ds_tol = 0.005\n", "idx_both.assigntorings()\n", "idx_both.hkl_tol = 0.02\n", "idx_both.cosine_tol = 0.002\n", "idx_both.score_all_pairs()\n", "idx_both.saveubis('grains_found_avg.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 $K_{\\alpha_{1}}$ g-vectors, refined with $K_{\\alpha_{1}}$ centroids\n", "- Indexed with $K_{\\alpha_{2}}$ g-vectors, refined with $K_{\\alpha_{2}}$ centroids\n", "- Indexed with merged g-vectors, refined with both $K_{\\alpha_{1}}$ and $K_{\\alpha_{2}}$ centroids\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_a1, hkli_a1 = assign_grains(jnp.array(idx_a1.ubis), q_sample_a1[mask_a1])\n", "hkle_a2, hkli_a2 = assign_grains(jnp.array(idx_a2.ubis), q_sample_a2[mask_a2])\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_a1 = jnp.argmin(hkle_a1, axis=0)\n", "ga_a2 = jnp.argmin(hkle_a2, axis=0)\n", "ga_both = jnp.argmin(hkle_both, axis=0)\n", "\n", "# get hkli per peak\n", "hkli_a1 = jnp.squeeze(jnp.take_along_axis(hkli_a1, ga_a1[None, :, None], axis=0))\n", "hkli_a2 = jnp.squeeze(jnp.take_along_axis(hkli_a2, ga_a2[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_a1, return_counts=True)[1].max(), jnp.unique(ga_a2, return_counts=True)[1].max(), jnp.unique(ga_both, return_counts=True)[1].max())\n", "M" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Refinement code" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "@jax.jit\n", "def solve_ub_analytical(hkl_int, q_sample, mask):\n", " \"\"\"UB = (Qs^T H) (H^T H)^-1, with masked rows zeroed out.\"\"\"\n", " m = mask.astype(q_sample.dtype)[:, None]\n", "\n", " H = hkl_int * m # (N, 3)\n", " Qs = q_sample * m # (N, 3)\n", "\n", " HHT = H.T @ H # (3, 3), symmetric\n", " QsHT = Qs.T @ H # (3, 3)\n", "\n", " # ridge: grains with <3 usable peaks (or all-padding rows from top_k)\n", " # would otherwise return NaN and poison the whole vmap batch\n", " eps = 1e-12 * jnp.trace(HHT) + 1e-15\n", " HHT = HHT + eps * jnp.eye(3, dtype=HHT.dtype)\n", "\n", " # X = QsHT @ inv(HHT); HHT symmetric => X = solve(HHT, QsHT.T).T\n", " return jnp.linalg.solve(HHT, QsHT.T).T\n", "\n", "\n", "@jax.jit\n", "def per_peak_err(ubi, origin, cen, hkl, wave, kvec, mask):\n", " \"\"\"|Δhkl| for every peak in a grain's buffer. Masked rows -> inf.\"\"\"\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", " return jnp.where(mask, e, jnp.inf)\n", "\n", "per_peak_err_vmap = jax.vmap(per_peak_err, in_axes=(0, 0, 0, 0, 0, 0, 0))\n", "\n", "@partial(jax.jit, static_argnums=(7,))\n", "def refine_grain(\n", " initial_ubi, # kept for call-site compatibility; unused (as before)\n", " initial_origin,\n", " cen_obs,\n", " hkl_int,\n", " wavelength,\n", " kvecs,\n", " mask,\n", " num_steps=20,\n", " init_damping=1e-2, # was learning_rate; same positional slot\n", " scaling_factor=1000.0,\n", "):\n", " mf = mask.astype(cen_obs.dtype)\n", " npk = jnp.sum(mf) + 1e-10\n", "\n", " def residuals(x):\n", " \"\"\"Mask-weighted hkl residual, flattened, for a trial origin.\"\"\"\n", " origin = x * scaling_factor\n", " q_sample = detector_to_q_vec(\n", " cen_obs[:, 0], cen_obs[:, 1], cen_obs[:, 2],\n", " wavelength, kvecs, origin,\n", " )\n", " ub = solve_ub_analytical(hkl_int, q_sample, mask)\n", " hklf = jnp.linalg.solve(ub, q_sample.T).T # (N, 3)\n", " return ((hklf - hkl_int) * mf[:, None]).ravel() # (3N,)\n", "\n", " def mean_hkl_error(r):\n", " d = r.reshape(-1, 3)\n", " return jnp.sum(jnp.sqrt(jnp.sum(d * d, axis=1) + 1e-24)) / npk\n", "\n", " jac = jax.jacfwd(residuals) # (3N, 3) — only 3 JVPs\n", "\n", " def lm_step(carry, _):\n", " x, lam = carry\n", "\n", " r = residuals(x)\n", " J = jac(x)\n", "\n", " JTJ = J.T @ J\n", " g = J.T @ r\n", " I3 = jnp.eye(3, dtype=JTJ.dtype)\n", "\n", " damp = lam * (jnp.diag(jnp.diag(JTJ)) + 1e-12 * I3)\n", " dx = jnp.linalg.solve(JTJ + damp, -g)\n", "\n", " x_try = x + dx\n", " r_try = residuals(x_try)\n", "\n", " # NaN in r_try makes this False, so a bad step is rejected automatically\n", " improved = jnp.sum(r_try * r_try) < jnp.sum(r * r)\n", " x_next = jnp.where(improved, x_try, x)\n", " lam_next = jnp.clip(jnp.where(improved, lam / 3.0, lam * 5.0), 1e-10, 1e10)\n", "\n", " return (x_next, lam_next), mean_hkl_error(r)\n", "\n", " x0 = initial_origin / scaling_factor\n", " lam0 = jnp.asarray(init_damping, x0.dtype)\n", "\n", " (x_final, _), loss_history = jax.lax.scan(\n", " lm_step, (x0, lam0), None, length=num_steps\n", " )\n", "\n", " final_origin = x_final * scaling_factor\n", " q_sample = detector_to_q_vec(\n", " cen_obs[:, 0], cen_obs[:, 1], cen_obs[:, 2],\n", " wavelength, kvecs, final_origin,\n", " )\n", " final_ubi = jnp.linalg.inv(solve_ub_analytical(hkl_int, q_sample, mask))\n", "\n", " final_loss = mean_hkl_error(residuals(x_final))\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 $K_{\\alpha_{1}}$" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%%time\n", "\n", "all_origins_a1 = jnp.zeros((len(idx_a1.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(jnp.int32), M)\n", " \n", " # Gather the data for this grain\n", " # If the grain has < M peaks, top_k pads with the lowest-index zero\n", " # entries (i.e. other grains' peaks); 'mask' is correctly False there.\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_a1 = jnp.arange(len(idx_a1.ubis))\n", "compressed_a1 = v_pack(grain_ids_a1, ga_a1, centroid_with_error[mask_a1], hkli_a1, wavelengths_a1[mask_a1], kvecs[mask_a1])\n", "\n", "mask_r = compressed_a1[\"mask\"]\n", "\n", "for it in range(3):\n", " ubis, origins, losses = refine_vmap(\n", " jnp.array(idx_a1.ubis), all_origins_a1,\n", " compressed_a1[\"cen\"], compressed_a1[\"hkl\"], compressed_a1[\"wave\"],\n", " compressed_a1[\"kvec\"], mask_r,\n", " 5, 1e-2, 1000.0)\n", "\n", " # outlier rejection - unassign peaks with more than 4x the median hkl error\n", " # LM is very sensitive!\n", " e = per_peak_err_vmap(ubis, origins, compressed_a1[\"cen\"], compressed_a1[\"hkl\"],\n", " compressed_a1[\"wave\"], compressed_a1[\"kvec\"], mask_r)\n", " med = jnp.nanmedian(jnp.where(mask_r, e, jnp.nan), axis=1)\n", "\n", " keep = mask_r & (e < 4.0 * med[:, None])\n", " print(f\"iter {it}: dropped {int((mask_r & ~keep).sum())} peaks, \"\n", " f\"min/grain now {int(keep.sum(axis=1).min())}\")\n", " mask_r = keep\n", "\n", "ubis_fit_a1, origins_sample_fit_a1, all_losses_a1 = ubis, origins, losses" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# convergence check\n", "\n", "print(\"mask_r peaks:\", int(mask_r.sum()),\n", " \" of\", int(compressed_a1[\"mask\"].sum()))\n", "\n", "pert = jnp.full_like(all_origins_a1, 1.0 / 1000.0) # 1 um real (scaled units)\n", "\n", "for n in (1, 2, 5, 20):\n", " _, o_n, _ = refine_vmap(jnp.array(idx_a1.ubis), all_origins_a1,\n", " compressed_a1[\"cen\"], compressed_a1[\"hkl\"],\n", " compressed_a1[\"wave\"], compressed_a1[\"kvec\"],\n", " mask_r, n, 1e-2, 1000.0)\n", " _, o_p, _ = refine_vmap(jnp.array(idx_a1.ubis), all_origins_a1 + pert,\n", " compressed_a1[\"cen\"], compressed_a1[\"hkl\"],\n", " compressed_a1[\"wave\"], compressed_a1[\"kvec\"],\n", " mask_r, n, 1e-2, 1000.0)\n", " s = jnp.linalg.norm(o_n - o_p, axis=1)\n", " print(n, \"steps -> shift p50/max:\",\n", " float(jnp.percentile(s, 50)), float(s.max()))" ] }, { "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_a1.T)\n", "axs[1].plot(all_losses_a1.mean(axis=0), label='a1')\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 $K_{\\alpha_{2}}$" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%%time\n", "\n", "all_origins_a2 = jnp.zeros((len(idx_a2.ubis), 3))\n", "\n", "# Pack the data into grain-specific buffers (num_grains, M, ...)\n", "grain_ids_a2 = jnp.arange(len(idx_a2.ubis))\n", "compressed_a2 = v_pack(grain_ids_a2, ga_a2, centroid_with_error[mask_a2], hkli_a2, wavelengths_a2[mask_a2], kvecs[mask_a2])\n", "\n", "mask_r = compressed_a2[\"mask\"]\n", "\n", "for it in range(3):\n", " ubis, origins, losses = refine_vmap(\n", " jnp.array(idx_a2.ubis), all_origins_a2,\n", " compressed_a2[\"cen\"], compressed_a2[\"hkl\"], compressed_a2[\"wave\"],\n", " compressed_a2[\"kvec\"], mask_r,\n", " 5, 1e-2, 1000.0)\n", "\n", " e = per_peak_err_vmap(ubis, origins, compressed_a2[\"cen\"], compressed_a2[\"hkl\"],\n", " compressed_a2[\"wave\"], compressed_a2[\"kvec\"], mask_r)\n", " med = jnp.nanmedian(jnp.where(mask_r, e, jnp.nan), axis=1)\n", "\n", " keep = mask_r & (e < 4.0 * med[:, None])\n", " print(f\"iter {it}: dropped {int((mask_r & ~keep).sum())} peaks, \"\n", " f\"min/grain now {int(keep.sum(axis=1).min())}\")\n", " mask_r = keep\n", "\n", "ubis_fit_a2, origins_sample_fit_a2, all_losses_a2 = ubis, origins, losses" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Both $K_{\\alpha_{1}}$ and $K_{\\alpha_{2}}$" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%%time\n", "\n", "all_origins_both = jnp.zeros((len(idx_both.ubis), 3))\n", "centroids_both = jnp.concatenate([centroid_with_error[mask_a1], centroid_with_error[mask_a2]])\n", "wavelengths_both = jnp.concatenate([wavelengths_a1[mask_a1], wavelengths_a2[mask_a2]])\n", "kvecs_both = jnp.concatenate([kvecs[mask_a1], kvecs[mask_a2]])\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", "mask_r = compressed_both[\"mask\"]\n", "\n", "for it in range(3):\n", " ubis, origins, losses = refine_vmap(\n", " jnp.array(idx_both.ubis), all_origins_both,\n", " compressed_both[\"cen\"], compressed_both[\"hkl\"], compressed_both[\"wave\"],\n", " compressed_both[\"kvec\"], mask_r,\n", " 5, 1e-2, 1000.0)\n", "\n", " e = per_peak_err_vmap(ubis, origins, compressed_both[\"cen\"], compressed_both[\"hkl\"],\n", " compressed_both[\"wave\"], compressed_both[\"kvec\"], mask_r)\n", " med = jnp.nanmedian(jnp.where(mask_r, e, jnp.nan), axis=1)\n", "\n", " keep = mask_r & (e < 4.0 * med[:, None])\n", " print(f\"iter {it}: dropped {int((mask_r & ~keep).sum())} peaks, \"\n", " f\"min/grain now {int(keep.sum(axis=1).min())}\")\n", " mask_r = keep\n", "\n", "ubis_fit_both, origins_sample_fit_both, all_losses_both = ubis, origins, losses" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, ax = plt.subplots()\n", "ax.plot(all_losses_a1.mean(axis=0), label='a1')\n", "ax.plot(all_losses_a2.mean(axis=0), label='a2')\n", "ax.plot(all_losses_both.mean(axis=0), label='both')\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[6]) # hexagonal crystal-frame ops\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_a1 = [ImageD11.grain.grain(ubis_fit_a1[g], translation=origins_sample_fit_a1[g]) for g in range(len(idx_a1.ubis))]\n", "grains_a2 = [ImageD11.grain.grain(ubis_fit_a2[g], translation=origins_sample_fit_a2[g]) for g in range(len(idx_a2.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_a1 = jnp.stack([g.U for g in grains_a1])\n", "oriens_a2 = jnp.stack([g.U for g in grains_a2])\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[6])\n", "grains_a1_fz = move_grains_to_fz(grains_a1, ROTATIONS[6])\n", "grains_a2_fz = move_grains_to_fz(grains_a2, ROTATIONS[6])\n", "grains_both_fz = move_grains_to_fz(grains_both, ROTATIONS[6])\n", "\n", "dist_a1, fit_a1, gt_a1 = match_grains(grains_a1_fz, grains_fz)\n", "dist_a2, fit_a2, gt_a2 = match_grains(grains_a2_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_a1, gt_a1, len(grains_a1)),\n", " ('a2', fit_a2, gt_a2, len(grains_a2)),\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_a1 = gt_losses(gt_a1, fit_a1, oriens_a1, compressed_a1)\n", "gt_losses_a2 = gt_losses(gt_a2, fit_a2, oriens_a2, compressed_a2)\n", "gt_losses_both = gt_losses(gt_both, fit_both, oriens_both, compressed_both)\n", "\n", "# ---------- metrics ----------\n", "\n", "pos_diff_a1 = jnp.linalg.norm(translations_sample[gt_a1] - origins_sample_fit_a1[fit_a1], axis=1)\n", "pos_diff_a2 = jnp.linalg.norm(translations_sample[gt_a2] - origins_sample_fit_a2[fit_a2], axis=1)\n", "pos_diff_both = jnp.linalg.norm(translations_sample[gt_both] - origins_sample_fit_both[fit_both], axis=1)\n", "\n", "misorien_a1 = jnp.array([jnp.min(Umis(U_matrices[g], oriens_a1[f], 6)[:, 1]) for g, f in zip(gt_a1, fit_a1)])\n", "misorien_a2 = jnp.array([jnp.min(Umis(U_matrices[g], oriens_a2[f], 6)[:, 1]) for g, f in zip(gt_a2, fit_a2)])\n", "misorien_both = jnp.array([jnp.min(Umis(U_matrices[g], oriens_both[f], 6)[:, 1]) for g, f in zip(gt_both, fit_both)])\n", "\n", "strains_a1 = jnp.array([grains_a1[f].eps_sample_matrix(struc.lattice_parameters) for f in fit_a1])\n", "strains_a2 = jnp.array([grains_a2[f].eps_sample_matrix(struc.lattice_parameters) for f in fit_a2])\n", "strains_both = jnp.array([grains_both[f].eps_sample_matrix(struc.lattice_parameters) for f in fit_both])\n", "\n", "strains_norm_a1 = jnp.sqrt((strains_a1**2).sum(axis=(1,2)))\n", "strains_norm_a2 = jnp.sqrt((strains_a2**2).sum(axis=(1,2)))\n", "strains_norm_both = jnp.sqrt((strains_both**2).sum(axis=(1,2)))\n", "\n", "loss_diff_a1 = all_losses_a1[fit_a1, -1] - gt_losses_a1\n", "loss_diff_a2 = all_losses_a2[fit_a2, -1] - gt_losses_a2\n", "loss_diff_both = all_losses_both[fit_both, -1] - gt_losses_both\n", "\n", "print(\"refined:\", float(all_losses_a1[fit_a1, -1].mean()),\n", " \" ground truth:\", float(gt_losses_a1.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_a1, bins=20)\n", "axs[1].hist(pos_diff_a2, bins=20)\n", "axs[2].hist(pos_diff_both, bins=20)\n", "axs[0].set_title(r'Just $K_{\\alpha_1}$')\n", "axs[1].set_title(r'Just $K_{\\alpha_2}$')\n", "axs[2].set_title(r'$K_{\\alpha_1}$ and $K_{\\alpha_2}$')\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_a1*1e3, bins=20)\n", "axs[1].hist(strains_norm_a2*1e3, bins=20)\n", "axs[2].hist(strains_norm_both*1e3, bins=20)\n", "axs[0].set_title(r'Just $K_{\\alpha_1}$')\n", "axs[1].set_title(r'Just $K_{\\alpha_2}$')\n", "axs[2].set_title(r'$K_{\\alpha_1}$ and $K_{\\alpha_2}$')\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_a1, bins=20)\n", "axs[1].hist(misorien_a2, bins=20)\n", "axs[2].hist(misorien_both, bins=20)\n", "axs[0].set_title(r'Just $K_{\\alpha_1}$')\n", "axs[1].set_title(r'Just $K_{\\alpha_2}$')\n", "axs[2].set_title(r'$K_{\\alpha_1}$ and $K_{\\alpha_2}$')\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_a1[fit_a1, -1], bins=20, alpha=0.5, label='refined')\n", "axs[0].hist(gt_losses_a1, bins=20, alpha=0.5, label='ground truth')\n", "axs[1].hist(all_losses_a2[fit_a2, -1], bins=20, alpha=0.5)\n", "axs[1].hist(gt_losses_a2, bins=20, alpha=0.5)\n", "axs[2].hist(all_losses_both[fit_both, -1], bins=20,alpha=0.5)\n", "axs[2].hist(gt_losses_both, bins=20, alpha=0.5)\n", "axs[0].legend()\n", "axs[0].set_title(r'Just $K_{\\alpha_1}$')\n", "axs[1].set_title(r'Just $K_{\\alpha_2}$')\n", "axs[2].set_title(r'$K_{\\alpha_1}$ and $K_{\\alpha_2}$')\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_a1, bins=20)\n", "axs[1].hist(loss_diff_a2, bins=20)\n", "axs[2].hist(loss_diff_both, bins=20)\n", "axs[0].set_title(r'Just $K_{\\alpha_1}$')\n", "axs[1].set_title(r'Just $K_{\\alpha_2}$')\n", "axs[2].set_title(r'$K_{\\alpha_1}$ and $K_{\\alpha_2}$')\n", "fig.supxlabel('Difference in loss function to ground truth')\n", "fig.supylabel('Counts')\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Statistics for paper" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def print_grain_metrics(\n", " pos_diff_a1, pos_diff_a2, pos_diff_both,\n", " misorien_a1, misorien_a2, misorien_both,\n", " strains_norm_a1, strains_norm_a2, strains_norm_both,\n", " loss_diff_a1, loss_diff_a2, loss_diff_both\n", "):\n", " # Defining metrics with their specific scaling factors\n", " # (Label, data1, data2, databoth, scale_factor, unit_label)\n", " metrics = [\n", " (\"Loss Diff\", loss_diff_a1, loss_diff_a2, loss_diff_both, 1e6, \"(x1e-6)\"),\n", " (\"Positional Diff\", pos_diff_a1, pos_diff_a2, pos_diff_both, 1.0, \"\"),\n", " (\"Misorientation\", misorien_a1, misorien_a2, misorien_both, 1e3, \"(x1e-3)\"),\n", " (\"Strain Norm\", strains_norm_a1, strains_norm_a2, strains_norm_both, 1e4, \"(x1e-4)\"),\n", " ]\n", "\n", " header = f\"{'Metric':<25} | {'a1 Mean':<12} | {'a2 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_a1, pos_diff_a2, pos_diff_both,\n", " misorien_a1, misorien_a2, misorien_both,\n", " strains_norm_a1, strains_norm_a2, strains_norm_both,\n", " loss_diff_a1, loss_diff_a2, loss_diff_both\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# total peaks (a1 and a2)\n", "mask_a1.sum(), mask_a1.sum()*2, mask_a1.sum()/n_grains, mask_a1.sum()*2/n_grains" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "p = jnp.asarray(pos_diff_a1)\n", "print(f\"median {jnp.median(p):.2f} p90 {jnp.percentile(p,90):.2f} \"\n", " f\"p99 {jnp.percentile(p,99):.2f} max {p.max():.2f}\")" ] }, { "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_a1.sum()][:, 1], centroid_with_error[:mask_a1.sum()][:, 0], label=r'$K_{\\alpha_{1}}$ peaks', s=.5)\n", "ax.scatter(centroid_with_error[mask_a1.sum():][:, 1], centroid_with_error[mask_a1.sum():][:, 0], label=r'$K_{\\alpha_{2}}$ peaks', s=.5)\n", "ax.set_aspect(1)\n", "ax.set(title='Dual-wavelength')\n", "ax.set_xticks([])\n", "ax.set_yticks([])\n", "ax.legend(loc='upper right')\n", "plt.show()\n", "plt.savefig('dual_peaks.png', dpi=600)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "end = time.time()\n", "print(f'Took {end-start:.0f} seconds')" ] } ], "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 }