pyextal.jaxGeo.simulate

A memory-safe jaxEngine.utils.simulate.

Functions

simulate_safely(params, metadata, config, constants[, ...])

Runs simulate, sizing the batch to fit and halving it on OOM.

Module Contents

pyextal.jaxGeo.simulate.simulate_safely(params, metadata, config, constants, batch=None, memory_gb=1.5)

Runs simulate, sizing the batch to fit and halving it on OOM.

The forward model holds a (batch, nbeams, nbeams) structure matrix, so the footprint is batch * nbeams**2 * 16 bytes – with 565 beams that is 5.1 MB per pixel, and simulate’s default batch of 4096 asks for 20 GB. That is what fails after a refinement that itself ran fine, because training chunks are far smaller than the simulation batch.

Parameters:

  • params (dict) – Normalised parameters.

  • metadata (dict) – From setup_optimization_params.

  • config (OptConfig) – The configuration used.

  • constants (dict) – Bloch constants; supplies nbeams.

  • batch (int, optional) – Pixels per batch. Sized from memory_gb when None.

  • memory_gb (float) – GB to aim for per batch.

Returns:

The simulated intensities.

Return type:

np.ndarray

Raises:

RuntimeError – If even the smallest batch runs out of memory.