# 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.

```python
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

```python
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.

```python
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

```python
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

```python
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:

```python
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

```python
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.

```python
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.

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

---

## Saving and loading `OptConfig`

```python
# 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.