ptyrax.simulate

ptyrax.simulate#

Functions

post_simulation(model, simulated_dataset[, ...])

Post-process and persist simulation outputs.

simulate(output_file[, ...])

Run a forward simulation of a ptychography model and save outputs.

simulate_model(model, *[, key])

Generate simulated diffraction patterns from a model.
ptyrax.simulate.post_simulation(model, simulated_dataset, output_file=None, save_equinox_model=True, postprocess_functions=())[source]#

Post-process and persist simulation outputs.

Applies optional post-processing transforms and saves the simulated dataset to HDF5. Optionally saves the ground-truth model as both .eqx and HDF5 files alongside the output.

Parameters:

  • model (PtychographyModel) – Ground-truth model used for simulation.

  • simulated_dataset (ImageDataset) – Simulated dataset to save.

  • output_file (Path | None) – Destination path for the simulated dataset HDF5.

  • save_equinox_model (bool) – Whether to also save model weights.

  • postprocess_functions (list[Callable[[ImageDataset], ImageDataset]]) – Transforms applied to the dataset before saving.

Return type:

None

ptyrax.simulate.simulate(output_file, preprocess_functions=(), dataset_load_fn=<function from_hdf5>, continue_from_reconstruction=None, model_type=<class 'ptyrax.models.ptychography.PtychographyModel'>, sweep_id=None, log_dir=None, *, dataset_path=None, key)[source]#

Run a forward simulation of a ptychography model and save outputs.

Initializes a model (optionally from an existing dataset or reconstruction), simulates diffraction patterns for all scanning positions, and writes the resulting dataset to output_file.

Parameters:

  • output_file (Path) – Path to write the simulated dataset HDF5 file.

  • preprocess_functions (list[Callable[[ImageDataset], ImageDataset]]) – Preprocessing transforms applied to the loaded dataset before model initialization.

  • dataset_load_fn (Callable[[str], ImageDataset]) – Callable that loads an ImageDataset from a path.

  • continue_from_reconstruction (str | None) – Optional path to a prior reconstruction from which to load model parameters.

  • model_type (type[ImagePredictionModel]) – Model class to instantiate.

  • sweep_id (str) – Optional sweep identifier for experiment tracking.

  • log_dir (str) – If provided, TensorBoard logs are written here.

  • dataset_path (str | None) – If provided, the model is initialized to mimic the geometry of this existing dataset.

  • key (Key) – JAX PRNG key.

Returns:

The simulated ImageDataset.

Return type:

ImageDataset

ptyrax.simulate.simulate_model(model, *, key=Array([0, 0], dtype=uint32))[source]#

Generate simulated diffraction patterns from a model.

Iterates over all scanning indices, resolves parametrizations for each, and collects the predicted images into an ImageDataset.

Parameters:

  • model (ImagePredictionModel) – Initialized model with resolved geometry.

  • key (Key) – JAX PRNG key (unused but kept for API consistency).

Returns:

Simulated dataset containing predicted diffraction images.

Return type:

ImageDataset

Contents