pyextal.jaxEngine.hdf5io
========================

.. py:module:: pyextal.jaxEngine.hdf5io

.. autoapi-nested-parse::

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

.. autoapisummary::

   pyextal.jaxEngine.hdf5io.load_for_refinement
   pyextal.jaxEngine.hdf5io.save_refinement
   pyextal.jaxEngine.hdf5io.list_refinement_runs
   pyextal.jaxEngine.hdf5io._ds
   pyextal.jaxEngine.hdf5io._jsonable
   pyextal.jaxEngine.hdf5io.snapshot_steps
   pyextal.jaxEngine.hdf5io._save_param_history


Module Contents
---------------

.. py:function:: 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
   :rtype: dict with keys


.. py:function:: 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.

   :param h5path:
   :type h5path: path-like
   :param config:
   :type config: OptConfig namedtuple
   :param optimized_param:
   :type optimized_param: dict — output of optimize()
   :param metadata:
   :type metadata: dict — output of setup_optimization_params()
   :param loss_history:
   :type loss_history: list of floats
   :param param_hist:
   :type param_hist: list of param dicts (snapshots every N steps)
   :param sim:
   :type sim: array (n_rows, n_cols, n_beams) — output of simulate()
   :param notes:
   :type notes: str — human-readable label for this run (e.g. "3layers_chi2_lr1e-3")
   :param snapshot_steps: `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.
   :type snapshot_steps: sequence of int, optional — the optimizer step each entry of

   :rtype: str — the run key used (e.g. "run_003")


.. py:function:: list_refinement_runs(h5path)

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

   :returns: **list of dicts, each with keys**
   :rtype: run, timestamp, notes, num_steps, loss_final


.. py:function:: _ds(grp, name, data)

.. py:function:: _jsonable(obj)

   Recursively convert a nested structure to JSON-serialisable form.


.. py:function:: 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.

   :param num_steps: Steps in this call to `optimize`.
   :type num_steps: int
   :param offset: Steps already completed before it.
   :type offset: int
   :param every: `optimize`'s snapshot interval.
   :type every: int

   :returns: One entry per snapshot, same length as the returned
             `param_history`.
   :rtype: list of int


.. py:function:: _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`.

   :param grp: The run group.
   :type grp: h5py.Group
   :param param_hist: Snapshots from `optimize()`.
   :type param_hist: list
   :param metadata: From `setup_optimization_params`; supplies the bounds,
                    the norm mask and `initial_thickness`.
   :type metadata: dict
   :param steps: Step index per snapshot; even spacing across
                 `n_steps` when None.
   :type steps: sequence or None
   :param n_steps: Length of the loss history, for the fallback spacing.
   :type n_steps: int