PyExtal-JAX: Gradient-Based Refinement

pyextal.jaxEngine is the autodifferentiation refinement engine for PyExtal. Install it with pip install "pyextal[jax]" (or "pyextal[cuda]" for NVIDIA GPUs). It performs gradient-based optimisation of crystal structure factors, sample thickness, probe parameters, and intensity scaling against experimental LARBED data using JAX and the Adam optimiser.

---

Complete Workflow

The full pipeline has two sequential stages:

Raw LARBED data
      │
      ▼
┌─────────────────────────────────────────────┐
│  Stage 1 — Pre-refinement  (pyextal)        │
│                                             │
│  LARBEDDiffractionInfo + LARBEDROI          │
│    → CoarseOptimize  (thickness, gl, tilt)  │
│    → FineOptimize    (ROI geometry, scale)  │
│    → save_pre_refinement(h5path, ...)       │
└───────────────┬─────────────────────────────┘
                │  writes  {sample}.h5
                ▼
┌─────────────────────────────────────────────┐
│  Stage 2 — JAX refinement  (jaxEngine)     │
│                                             │
│  load_for_refinement(h5path)                │
│    → OptConfig + setup_optimization_params  │
│    → optimize()                             │
│    → simulate()                             │
│    → save_refinement(h5path, ..., notes="") │
└─────────────────────────────────────────────┘

Both stages share a single HDF5 file per dataset. Multiple refinement runs (different OptConfig settings) accumulate as run_001, run_002, … in the same file — no stale directory trees or prefixed .npy files.

---

Stage 1: Pre-refinement

The pre-refinement notebook docs/source/notebooks/si110zoneaxis0430.ipynb demonstrates the full process on a silicon [110] zone-axis dataset. The key steps are:

1. Load the raw LARBED stack, variance maps, and g-vector indices.

2. Create LARBEDDiffractionInfo with the crystal .dat file, initial thickness, and tilt.

3. Create LARBEDROI — selects which reflections to include, applies rotation, sets the sampling grid.

4. CoarseOptimize.optimizeOrientationThickness() — grid-search for thickness and tilt via template cross-correlation.

5. FineOptimize.optimize_roi() — refines scale, background, rotation, and ROI geometry.

6. Save everything to HDF5.

from pyextal.hdf5io import save_raw, save_pre_refinement

h5path = box_path / f"{prefix}.h5"
# store the raw measurement (optional but recommended for provenance)
save_raw(h5path, dp, variance, gindex, prefix=prefix, crystal_dat="si110.dat")
# store all pre-refinement outputs (simgrid, templates, geometry, …)
save_pre_refinement(h5path, roi, fine, gInclude)

save_pre_refinement reads the following attributes directly from the roi and fine objects — no manual array packing:

Written to HDF5	Source
simgrid (n_rows, n_cols, 3)	roi.simGrid
templates (n_beams, n_rows, n_cols)	roi.templates
variance (n_beams, n_rows, n_cols)	roi.variance
roi_regions	roi.regions
g_include	the gInclude list
allshift	roi.allshift
attrs: thickness, scale, background, gl, rotation	fine.thickness, fine.errorFunc.scale/background, roi.gl, roi.rotation

---

Stage 2: JAX Refinement

Loading data

from pyextal.jaxEngine.hdf5io import load_for_refinement

data = load_for_refinement("sample.h5")

# data keys:
#   tilts             jnp.array  (n_rows, n_cols, 3)
#   experimental_data jnp.array  (n_rows, n_cols, n_beams)
#   variance          jnp.array  (n_rows, n_cols, n_beams)
#   geometry          dict — 'thickness', 'scale', 'background', 'gl', 'rotation', 'allshift'

# optional 2× spatial subsampling for speed
stride = 2
tilts             = data['tilts'][::stride, ::stride, :].reshape(-1, 3)
experimental_data = data['experimental_data'][::stride, ::stride, :]
variance          = data['variance'][::stride, ::stride, :]
geometry          = data['geometry']

Configuring OptConfig

OptConfig is a NamedTuple in pyextal.jaxEngine.setParams. Every field has a default; only the parameters you want to override need to be specified.

from pyextal.jaxEngine.setParams import OptConfig, get_bloch_constants, setup_optimization_params
from pyextal.callBloch import bloch_parse

bloch_param = bloch_parse('crystal.dat')
constants   = get_bloch_constants(bloch_param.alpha)

output_shape  = experimental_data.shape[:2]
kernel_radius = 4
stride        = 16

config = OptConfig(
    # --- data dimensions ---
    output_indices = tuple(range(experimental_data.shape[-1])),
    output_shape   = output_shape,

    # --- batching (chunking required for Voigt convolution) ---
    batch_mode   = 'chunking',
    chunk_shape  = (stride + kernel_radius*2, stride + kernel_radius*2),
    stride       = (stride, stride),
    update_mode  = 'per_epoch',

    # --- optimiser ---
    loss_function  = 'chi2',
    learning_rate  = 1e-3,
    num_steps      = 400,
    warmup_steps   = 150,   # freeze structure factors for first N steps
    fix_after      = 16,    # freeze beams beyond this index during warmup

    # --- structure factors ---
    nlayers             = 3,
    refined_entity      = 'structure_factor',
    reflections         = 50,
    sf_mask             = (True, True, True, True),
    symmetry_constraint = 'SpaceGroup',

    # --- thickness ---
    thickness_optimize = True,
    initial_thickness  = geometry['thickness'],
    thickness_range    = (0, 1),
    thickness_normalize= 100,

    # --- probe (Voigt PSF) ---
    probe_optimize     = True,
    use_voigt          = True,
    kernel_radius      = kernel_radius,
    initial_sigma_probe= 0.5,
    initial_gamma_probe= 0.5,
    probe_sigma_range  = (0.2, 0.8),
    probe_gamma_range  = (0.2, 0.8),
    probe_normalize    = 100,

    # --- intensity scaling & background ---
    initial_scale      = geometry['scale'],
    scale_range        = (geometry['scale']*0.9, geometry['scale']*1.1),
    initial_background = geometry['background'],
    background_range   = (geometry['background']*0.5, geometry['background']*1.5),
    intensity_scale_normalize = 100,
    background_normalize      = 100,
)

Running the optimisation

from pyextal.jaxEngine.optimize import optimize
from pyextal.jaxEngine.utils import simulate

params, metadata, config = setup_optimization_params(config, tilts, constants)

optimized_param, loss_history, param_hist, opt_state = optimize(
    tilts, experimental_data, variance, config, constants, params, metadata
)

sim = simulate(optimized_param, metadata, config, constants)

Saving results

from pyextal.jaxEngine.hdf5io import save_refinement

run_key = save_refinement(
    "sample.h5", config, optimized_param, metadata,
    loss_history, param_hist, sim,
    notes="3layers_50ref_SpaceGroup_chi2",
)
print(f"Saved as {run_key}")  # e.g. "run_001"

Each call to save_refinement appends a new run_NNN group to /refinement/ in the same HDF5 file. The group stores the full OptConfig (scalar fields as attributes), the optimised structure factors, loss curve, simulated intensities, and tilt-dependent metadata.

Pass snapshot_steps= so the parameter trajectory under history/ gets a real x axis; hdf5io.snapshot_steps(num_steps, offset) builds it, and the offset is what makes it right when optimize is called once per segment:

from pyextal.jaxEngine.hdf5io import save_refinement, snapshot_steps

run_key = save_refinement(
    "sample.h5", config, optimized_param, metadata,
    loss_history, param_hist, sim,
    notes="3layers_50ref_SpaceGroup_chi2",
    snapshot_steps=snapshot_steps(config.num_steps),
)

Everything under history/ is denormalised, so it is directly plottable — thickness in Angstrom (the raw parameter is a ratio of metadata['initial_thickness']), structure factors as U_g rather than as positions between their bounds. Runs written before this existed carry a param_history_note attribute instead: np.array() on a list of nested param dicts raises, so no trajectory was ever stored.

Comparing runs

from pyextal.jaxEngine.hdf5io import list_refinement_runs

for run in list_refinement_runs("sample.h5"):
    print(run['run'], run['timestamp'], f"loss={run['loss_final']:.4f}", run['notes'])

---

HDF5 File Layout

{sample}.h5
├── attrs:  crystal_dat, acquisition_prefix, date
│
├── raw/
│   ├── dp              [n_g, ny, nx]
│   ├── variance        [n_g, ny, nx]
│   └── g_vectors       [n_g, 3]
│
├── pre_refinement/
│   ├── attrs:          thickness, scale, background, gl, rotation, grid_shape
│   ├── simgrid         [n_rows, n_cols, 3]
│   ├── templates       [n_beams, n_rows, n_cols]
│   ├── variance        [n_beams, n_rows, n_cols]
│   ├── roi_regions     [n_regions, 4, 2]
│   ├── g_include       [n_beams, 3]
│   └── allshift        [n_shifts, 2]
│
└── refinement/
    ├── run_001/
    │   ├── attrs:          timestamp, notes, num_steps, learning_rate, …
    │   ├── optimized_sf    [n_ref, 4]   — [xfre, xfim, afre, afim]
    │   ├── gvectors        [n_ref, 3]
    │   ├── thickness_ratio [n_layers]
    │   ├── loss_history    [n_steps]
    │   ├── simulated_data  [n_beams, n_rows, n_cols]
    │   ├── history/                       — parameter trajectory, physical units
    │   │   ├── steps            [n_snap]  — optimizer step per snapshot
    │   │   ├── structure_factor [n_snap, n_ref, 4]
    │   │   ├── thickness        [n_snap, n_layers]   (Angstrom)
    │   │   └── intensity_scale, background, probe_sigma, probe_gamma  [n_snap]
    │   └── metadata/
    │       ├── sg_batch, bigkz_batch, bigkn_batch, ig_batch
    │       ├── phase_real, phase_imag    [n_ref, nbeams, nbeams]
    │       ├── mask, static_ughr, static_ughi, reflection_list
    │       └── attrs: bounds_json        (upperbound/lowerbound/norm_mask as JSON)
    ├── run_002/
    └── …

---

OptConfig Parameter Reference

General optimisation

Parameter	Type	Default	Description
loss_function	str	'chi2'	'chi2': weighted squared error (requires variance). 'cosine': ignores absolute scale, matches pattern shape.
batch_mode	str	'chunking'	'chunking': 2D spatial tiles — required for Voigt convolution. 'random': random pixel sampling, faster but no convolution.
update_mode	str	'per_epoch'	'per_batch': update after every tile. 'per_epoch': accumulate gradients over all tiles then update — more stable.
learning_rate	float	1e-3	Adam step size.
num_steps	int	1000	Total optimisation iterations.
warmup_steps	int	0	Steps where structure factors are frozen; only thickness, probe, and scale are refined.
fix_after	int	None	During warmup, freeze beams beyond this index (by reciprocal-vector length).

Structure factors

Parameter	Type	Description
reflections	int | tuple	int: optimise top N shortest-g reflections. tuple: explicit list of (h,k,l) tuples.
sf_mask	tuple[bool, bool, bool, bool]	Which components to optimise: (xfre, xfim, afre, afim). xf = elastic \(U_g\), af = absorptive \(U'_g\).
symmetry_constraint	str	'Friedel': \(U_{-g}=U_g^*\). 'SpaceGroup': full space-group constraints.
nlayers	int	Number of layers for multi-layer structures.
frozen_layers	tuple[int, ...]	0-based layers whose structure factors are held fixed (thicknesses still refine). Honoured after warm-up too.
diagonalise_measured	bool	Default True: widen the diagonalised set (ndiag) so every measured beam is computed exactly at every tilt. False: a measured beam that beamcheck ranks weak at its tilt is simulated as 0 and left out of the loss.
refined_entity	str	'structure_factor' (default) or 'crystal_parameter' (parameterised model).
power_regularization_weight	float	L2 regularisation on total potential.

Thickness

Parameter	Type	Description
thickness_optimize	bool	Refine sample thickness.
initial_thickness	float	Starting value in Å — use geometry['thickness'] from pre-refinement.
thickness_range	(min, max)	Bounds on the thickness as a ratio of initial_thickness — e.g. (0.8, 1.2) for ±20 %. One Adam step moves the ratio by about learning_rate × (max − min), so bounds given in Å make every step initial_thickness times too large; setup warns when they look like Å.
thickness_normalize	float	Internal scaling factor (100 → parameter lives in [0, 1] range).

Probe

Parameter	Type	Description
probe_optimize	bool	Refine the point-spread function.
use_voigt	bool	Apply Pseudo-Voigt convolution. Requires batch_mode='chunking'.
kernel_radius	int	Convolution kernel radius in pixels; kernel is \((2r+1)\times(2r+1)\).
initial_sigma_probe	float	Initial Gaussian width \(\sigma\).
initial_gamma_probe	float	Initial Lorentzian width \(\gamma\).
probe_sigma_range	(min, max)	Bounds for \(\sigma\).
probe_gamma_range	(min, max)	Bounds for \(\gamma\).

Intensity & background

Parameter	Type	Description
initial_scale	float	Starting global scale \(S\) in \(I_{obs} = S \cdot I_{sim} + B\). Use geometry['scale'].
scale_range	(min, max)	Bounds for \(S\).
initial_background	float	Starting background \(B\). Use geometry['background'].
background_range	(min, max)	Bounds for \(B\).
intensity_scale_normalize, background_normalize	float	Internal scaling (typically 100).

Spatial grid & chunking

Parameter	Type	Description
output_indices	tuple[int]	Beam indices to include in the simulation output.
output_shape	(ny, nx)	Full scan dimensions after any spatial subsampling.
chunk_shape	(cy, cx)	Tile size in chunking mode; use (stride + kernel_radius*2, …) to include convolution overlap. In random mode chunk_shape[0] is the batch size.
stride	(sy, sx)	Tile step in chunking mode. stride < chunk_shape produces overlapping tiles.

---

Batching strategies

Chunking (batch_mode='chunking')

Required when use_voigt=True. The scan is tiled into overlapping rectangular windows.

batch_mode    = 'chunking',
output_shape  = (100, 100),
kernel_radius = 4,
stride        = 16,
chunk_shape   = (stride + kernel_radius*2, stride + kernel_radius*2),  # = (24, 24)

The overlap (kernel_radius pixels on each side) prevents edge artefacts in the Voigt convolution. The stride controls how far each tile advances; the optimize() loop assembles the full map from tile contributions.

The tile does not set the GPU memory: within a chunk the forward model runs the tilts constants['tilt_batch'] at a time (default jaxEngine.model.TILT_BATCH, 1024), each sub-batch rematerialised. Larger tiles therefore cost no memory and waste less on the overlap; a chunk covering the whole scan is fine. jaxEngine.model.vmapped_forward_reference keeps the previous all-tilts-at-once evaluation for comparison.

Random (batch_mode='random')

No convolution — useful for a fast initial structure-factor pass.

batch_mode  = 'random',
use_voigt   = False,
chunk_shape = (1024, 1),   # 1024 = batch size; second dim is ignored

---

Saving and loading OptConfig

# save for reproducibility
config.to_yaml("opt_config.yaml")

# restore
config = OptConfig.from_yaml("opt_config.yaml")

The YAML round-trip converts all internal NumPy types to plain Python scalars. The config is also embedded verbatim in the HDF5 run group as individual attributes, so every saved result carries its own configuration.