pyextal.jaxEngine.optimize¶

Classes¶

`BatchGenerator`
`ChunkingGenerator`
`RandomGenerator`

Functions¶

`get_loss_fn`(→ Callable)	Create loss function with model_fn bound via closure.
`_freeze_layers`(→ None)	Freezes every structure factor of the layers in `config.frozen_layers`.
`create_freeze_mask`(→ Dict)	Create mask for freezing/unfreezing parameters during optimization.
`optimize`(→ Tuple[Dict, List[float], List[Dict], ...)

Module Contents¶

class pyextal.jaxEngine.optimize.BatchGenerator(config: pyextal.jaxEngine.setParams.OptConfig, tilts: jax.numpy.ndarray, experimental_data: jax.numpy.ndarray, variance: jax.numpy.ndarray, metadata: Dict)¶

config¶

tilts¶

experimental_data¶

variance¶

metadata¶

n_samples¶

abstractmethod get_batches(epoch: int) → List[Dict]¶

class pyextal.jaxEngine.optimize.ChunkingGenerator(config: pyextal.jaxEngine.setParams.OptConfig, tilts: jax.numpy.ndarray, experimental_data: jax.numpy.ndarray, variance: jax.numpy.ndarray, metadata: Dict)¶

Bases: BatchGenerator

_cached_batches = None¶

get_batches(epoch: int) → List[Dict]¶

create_2d_tilt_chunks() → List[Tuple[slice, slice]]¶

class pyextal.jaxEngine.optimize.RandomGenerator(config: pyextal.jaxEngine.setParams.OptConfig, tilts: jax.numpy.ndarray, experimental_data: jax.numpy.ndarray, variance: jax.numpy.ndarray, metadata: Dict, batch_size: int = 64)¶

Bases: BatchGenerator

batch_size = 64¶

experimental_data¶

variance¶

get_batches(epoch: int) → List[Dict]¶

pyextal.jaxEngine.optimize.get_loss_fn(config: pyextal.jaxEngine.setParams.OptConfig, constants: Dict, metadata: Dict, model_fn: callable = None) → Callable¶

Create loss function with model_fn bound via closure.

Parameters:

config – Optimization configuration
constants – Bloch constants
metadata – Pre-calculated metadata
model_fn – The model function (for crystal parameter refinement)

pyextal.jaxEngine.optimize._freeze_layers(mask: Dict, config: pyextal.jaxEngine.setParams.OptConfig) → None¶: Freezes every structure factor of the layers in config.frozen_layers.

pyextal.jaxEngine.optimize.create_freeze_mask(config: pyextal.jaxEngine.setParams.OptConfig, params: Dict) → Dict¶: Create mask for freezing/unfreezing parameters during optimization.

pyextal.jaxEngine.optimize.optimize(tilts: jax.numpy.ndarray, experimental_data: jax.numpy.ndarray, variance: jax.numpy.ndarray, config: pyextal.jaxEngine.setParams.OptConfig, constants: Dict, initial_params: Dict, metadata: Dict, opt_state: optax.OptState = None) → Tuple[Dict, List[float], List[Dict], optax.OptState]¶