laue_torch

A fully differentiable PyTorch forward model for white-beam Laue micro-diffraction, plus a per-voxel orientation- and strain-distribution function (ODF / SDF) recovery pipeline that runs on top of any conventional Laue indexer.

The package is built as a drop-in refinement stage downstream of the existing LaueMatching pipeline: LaueMatching's coarse-grid indexer (or any equivalent) gives an approximate mean orientation per voxel; laue_torch refines the distribution parameters around that mean from peak shape, with analytic gradients through every step of the forward map.

Highlights

• Differentiable forward: parity-tested against the reference NumPy/C simulator (94/94 spot match on the canonical test case); passes a strict torch.autograd.gradcheck on the geometry path for every parameter group (orientation, strain, lattice, detector pose).
• Distribution-level recovery: tangent-Gaussian on SO(3) for unimodal mosaic; mixture for twins / sub-grain modes; multivariate Gaussian on Voigt-6 strain.
• Tensor GND density: the lattice curvature κ of the recovered per-voxel orientation field and Nye's tensor α = κᵀ − tr(κ) I from it (laue_torch.nye.lattice_curvature, nye_alpha); FCC slip-system projection helper (minimum-norm) included.
• Posterior uncertainty: Laplace approximation at convergence, returned with its Hessian eigenvalues, condition number, effective rank and a positive-definiteness flag. Read those before any marginal sigma (see below).
• Hessian eigenanalysis: explicit identification of the physical degeneracies of polychromatic Laue (hydrostatic-strain null, lattice scale null, $P_z!\leftrightarrow!$lattice trade-off, etc.).
• Real-data adapter: reads LaueMatching post-processed H5 output and runs the per-voxel ODF refinement with one function call.

Quick start

import torch
from laue_torch import LaueForwardModel, generate_hkls, parse_params

p = parse_params("simulation/params_sim.txt")
hkls = generate_hkls(p.sg_num, p.lattice, p.E_hi)
t = p.to_tensors()

model = LaueForwardModel(
    hkls=hkls, n_pix=t["n_pix"], px_size=t["px_size"],
    psf_sigma=t["psf_sigma"],
    rotation="quat", strain_mode="voigt",
)

U = torch.tensor([[1.0, 0.0, 0.0, 0.0]], dtype=torch.float64)   # (G, 4) quaternion
img = model(U, t["lattice"], t["P"], t["R"])                    # (Nx, Ny)

Axis order. The model renders img[X, Y] (X = detector column first). A real detector frame, a background built from one, and the LaueMatching indexer are image[row, col] = [Y, X], the transpose. Transpose a render (img.T) before writing it as a frame for the indexer or comparing it with a real frame; on a square detector nothing fails if you forget, and every spot lands on the wrong pixel. laue-torch CLI output carries /entry1/axis_order = b"XY" for this reason (and warns); pass -axisOrder YX to write detector-layout frames the indexer can read directly. laue_torch.io.to_model_layout converts and shape-checks.

Per-voxel ODF refinement on real data

from laue_torch import parse_params
from laue_torch.realdata import LaueScanLoader, VoxelODFRefiner, plot_sigma_map

params = parse_params("simulation/params_sim.txt")
refiner = VoxelODFRefiner(params, sigma_init_deg=1.0,
                          M_render=128, n_steps=500)

results = []
for voxel in LaueScanLoader("/path/to/scan/"):
    results.append(refiner.refine(voxel))

plot_sigma_map(results, grid_shape=(20, 20),
               out_path="sigma_map.png", show_posterior=True)

What the loader and refiner do with the data:

• Orientation seeds come from the indexer's OrientMatrix columns, chosen by the solution table's column count: 34 (RunImage, cols 22..30) or 35 (stream, cols 23..31). Any other layout raises; pass orientation_columns=(lo, hi).
• Frames stay as stored (image[row, col]) in VoxelMeasurement.image, labelled axis_order="YX" (or whatever an <entry>/axis_order marker in the file says). VoxelODFRefiner and MultiGrainVoxelRefiner transpose to the model layout on entry and check the shape against (NrPxX, NrPxY). There is no default layout: a hand-built VoxelMeasurement, or MultiGrainVoxelRefiner.refine(image, U), must say axis_order="YX" (real frame) or "XY" (a laue_torch render), or the refiner raises. (A default could not be right for both: VoxelODFRefiner did not transpose before 0.1.4, and on a square detector a wrong guess is silent.)
• Energy band: fits render in the parameter file's Elo..Ehi. A file without them is refused by every refiner (the forward CLI still defaults to 5-30 keV), and make_lauematching_params without a band flags it so write_lauematching_config refuses to write it.
• Intensity scale and background: VoxelODFRefiner fits a * render + b to the frame with (a, b) solved in closed form at every step (and in the posterior residual), so the recovered spread does not depend on the counts or the pedestal; (a, b) are in result.metadata["intensity_scale" / "intensity_offset"]. Before 0.1.5 the loss was the raw MSE against a unit-intensity render.
• Harmonics: reflections that share a seed pixel ((111), (222), ...) are rendered once (the lowest order), as in MultiGrainVoxelRefiner; before 0.1.5 such a spot was predicted n times brighter than a single reflection.

Each result is a VoxelODFResult carrying:

• U_mean: refined mean orientation (3×3)
• sigma_U_deg: recovered isotropic mosaic spread
• posterior_sigma_U_deg: a one-number summary of the Laplace posterior on sigma_U_deg, and posterior, the full LaplacePosterior
• final_loss, n_steps, dt_s, ...

Before quoting any posterior sigma, read posterior.eigvals, posterior.cond_number, posterior.rank_eff and posterior.is_positive_definite. A non-positive-definite Hessian (an unconverged fit, or a saddle) makes the affected sigmas NaN rather than small. Orientation spread and deviatoric strain are strongly coupled in white-beam Laue, so a small marginal sigma on one of them can sit on a near-null direction the eigenvalues expose.

Several grains in one voxel

MultiGrainVoxelRefiner(params, mode=...) fits K modes to one frame: "orient_only", "strain_voigt" (6-component strain mean) or "strain_deviatoric", which fits the 5 trace-free components (e11, e22, e23, e13, e12) with e33 = -(e11 + e22). Prefer it to strain_voigt: a pure hydrostatic strain moves no Laue spot (it only shifts the Bragg energy), so the 6th Voigt direction is exactly unidentifiable from spot positions. compute_posterior=True returns result.posterior with result.posterior_param_names; with refine_means=False it is conditional on the fixed orientations (metadata["posterior_conditional_on_fixed_means"]), which understates strain uncertainty.

For multi-voxel scans, the plots module supplies plot_sigma_map, plot_orientation_map (pass space_group=params.sg_num; it was cubic-only), and plot_gnd_map, which computes the lattice curvature of the recovered orientation field by central differences, Nye's tensor from it, and ||α||_F / b with the voxel spacing converted from um to m (before 0.1.5 it plotted ||κ||_F / b with the spacing left in um).

Tutorial

The single-script tutorial examples/tutorial_per_voxel_odf.py walks through the full synthetic pipeline end-to-end (forward → render → Adam → Nye tensor → Laplace posterior) in $\sim$3 min on CPU.

Synthetic experiments

The numerical experiments that characterised this package (mosaic-spread recovery, twin-variant identification, intragranular gradients and the Nye tensor, Laplace-posterior calibration, basin-of-convergence and Hessian eigenanalysis studies) are research scripts and are not distributed with the package.

What ships instead is the part that transfers: the two runnable tutorials in examples/, and the test suite, which exercises every one of those code paths against synthetic ground truth.

Tests

cd packages/laue_torch
pip install -e '.[dev]'
KMP_DUPLICATE_LIB_OK=TRUE pytest

343 tests cover parity against the NumPy/C reference, gradient flow / gradcheck on every parameter group, calibration recovery, distribution moments, mixture / Nye correctness, coded-aperture and joint-fit paths, and a contract test pinning the midas-stress misorientation convention.

The four parity tests need a LaueMatching checkout (they read simulation/params_sim.txt and the reference simulator). Run from an installed wheel instead and they skip cleanly rather than failing.

Documentation

• docs/torch-forward-model.md — math + API reference: the forward map, every differentiable parameter group, and the parameterisation choices.
• examples/ — two runnable tutorials: per-voxel ODF recovery, and the coded-aperture voxel workflow.
• The test suite is the other reference. tests/test_parity.py pins the forward model against the repository's NumPy/C simulator, and tests/test_grad.py gradchecks every parameter group.

Dependencies

Runtime (installed automatically):

• torch >= 2.0, numpy >= 1.20, scipy >= 1.7, h5py >= 3.0, Pillow >= 9.0
• midas-stress >= 0.11.0 — symmetry-reduced misorientation. Note it returns radians; laue_torch.symmetry converts to degrees in one place.
• midas-hkls >= 0.9.0 — lattice, space group, hkl generation, absorption
• midas-invert >= 0.1.1 — Laplace uncertainty, fit / loss primitives

Optional: matplotlib >= 3.5 (pip install 'laue-torch[viz]') for the plotting helpers in laue_torch.realdata.plots.

Citation

If you use this package, please cite the LaueMatching pipeline that produces the indexer seed orientations:

H. Sharma, D. Sheyfer, R. Harder, J.Z. Tischler. LaueMatching: an approach for rapid and robust indexing of Laue diffraction patterns. J. Appl. Cryst. 59, 552–563 (2026).

License

BSD-3-Clause (UChicago Argonne, LLC) — see LICENSE.