pyextal.jaxEngine.hdf5io

HDF5-based I/O for JAX refinement data.

Provides load_for_refinement / save_refinement / list_refinement_runs.

The /refinement group inside the HDF5 file holds one subgroup per optimization run, so multiple configurations can be compared within a single dataset file.

Layout of each run group

refinement/run_NNN/

attrs: timestamp, notes, + scalar OptConfig fields optimized_sf (n_reflections, 4) [xfre, xfim, afre, afim] gvectors (n_reflections, 3) thickness_ratio (n_layers,) simulated_data (n_beams, n_rows, n_cols) loss_history (n_steps,) history/

steps (n_snapshots,) optimizer step each snapshot was taken at structure_factor (n_snapshots, n_reflections, 4) physical units thickness (n_snapshots, n_layers) Angstrom intensity_scale, background, probe_sigma, probe_gamma (n_snapshots,)

metadata/

sg_batch, bigkz_batch, bigkn_batch, ig_batch, output_indices_reordered, mask, static_ughr, static_ughi, reflection_list, phase_real, phase_imag, bounds_json (attr — JSON string for upperbound/lowerbound/norm_mask)

Functions

load_for_refinement(h5path)

Load pre-refinement data from HDF5 ready for the JAX optimizer.

save_refinement(h5path, config, optimized_param, ...)

Append a new run_NNN group to /refinement/ with all JAX optimization results.

list_refinement_runs(h5path)

Return a summary list of all refinement runs stored in the HDF5 file.

_ds(grp, name, data)

_jsonable(obj)

Recursively convert a nested structure to JSON-serialisable form.

snapshot_steps(num_steps[, offset, every])

The optimizer step index of each snapshot optimize() keeps.

_save_param_history(grp, param_hist, metadata, steps, ...)

Write the parameter trajectory under history/ as plain arrays.

Module Contents

pyextal.jaxEngine.hdf5io.load_for_refinement(h5path)

Load pre-refinement data from HDF5 ready for the JAX optimizer.

Returns:

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

Return type:

dict with keys

pyextal.jaxEngine.hdf5io.save_refinement(h5path, config, optimized_param, metadata, loss_history, param_hist, sim, *, notes='', snapshot_steps=None)

Append a new run_NNN group to /refinement/ with all JAX optimization results.

Each call creates the next run group (run_001, run_002, …) so multiple configurations can be saved to the same dataset file and compared.

Parameters:

  • h5path (path-like)

  • config (OptConfig namedtuple)

  • optimized_param (dict — output of optimize())

  • metadata (dict — output of setup_optimization_params())

  • loss_history (list of floats)

  • param_hist (list of param dicts (snapshots every N steps))

  • sim (array (n_rows, n_cols, n_beams) — output of simulate())

  • notes (str — human-readable label for this run (e.g. "3layers_chi2_lr1e-3"))

  • snapshot_steps (sequence of int, optional — the optimizer step each entry of) – param_hist was taken at, for the x axis of a trajectory plot. Use snapshot_steps() to build it. Falls back to even spacing across loss_history when omitted, which is right only if the snapshots are evenly spaced.

Return type:

str — the run key used (e.g. “run_003”)

pyextal.jaxEngine.hdf5io.list_refinement_runs(h5path)

Return a summary list of all refinement runs stored in the HDF5 file.

Returns:

list of dicts, each with keys

Return type:

run, timestamp, notes, num_steps, loss_final

pyextal.jaxEngine.hdf5io._ds(grp, name, data)
pyextal.jaxEngine.hdf5io._jsonable(obj)

Recursively convert a nested structure to JSON-serialisable form.

pyextal.jaxEngine.hdf5io.snapshot_steps(num_steps, offset=0, every=5)

The optimizer step index of each snapshot optimize() keeps.

optimize appends the starting parameters once, then appends again after every every-th step, so a run of num_steps steps produces snapshots at steps 0, 1, every+1, 2*every+1, … Pass offset when the run is one segment of a longer refinement (the geometry loop calls optimize several times), so the returned indices are absolute rather than per-segment.

Parameters:

  • num_steps (int) – Steps in this call to optimize.

  • offset (int) – Steps already completed before it.

  • every (int) – optimize’s snapshot interval.

Returns:

One entry per snapshot, same length as the returned param_history.

Return type:

list of int

pyextal.jaxEngine.hdf5io._save_param_history(grp, param_hist, metadata, steps, n_steps)

Write the parameter trajectory under history/ as plain arrays.

The snapshots optimize() returns are nested dicts of JAX arrays holding normalised values, so np.array(param_hist) builds an object array that h5py refuses – which is why this used to land in its caller’s except every single time and no run ever carried a history. Denormalise each snapshot and split it into one array per quantity instead.

Thickness is stored in Angstrom: the forward model uses the parameter as a ratio of metadata[‘initial_thickness’] (see model.propagate’s caller), so the physical value is denormalised_ratio * initial_thickness.

Parameters:

  • grp (h5py.Group) – The run group.

  • param_hist (list) – Snapshots from optimize().

  • metadata (dict) – From setup_optimization_params; supplies the bounds, the norm mask and initial_thickness.

  • steps (sequence or None) – Step index per snapshot; even spacing across n_steps when None.

  • n_steps (int) – Length of the loss history, for the fallback spacing.