pyextal.jaxEngine.utils

Functions

reshape_tilts_to_2d(→ jax.numpy.ndarray)

Reshape flat tilt array to 2D grid

debug_plot(data[, num, title])

result_plot(save_path, prefix, results, config, tilts, ...)

simulate(→ jax.numpy.ndarray)

Simulate 2D intensity map by processing in batches.

simulate_with_gradients(→ Tuple[jax.numpy.ndarray, Dict])

Simulate diffraction patterns AND compute gradients w.r.t. structure factors.

matchIdx(beams, gInclude)

Module Contents

pyextal.jaxEngine.utils.reshape_tilts_to_2d(tilts: jax.numpy.ndarray, tilt_shape: Tuple[int, int]) → jax.numpy.ndarray

Reshape flat tilt array to 2D grid

Parameters:

  • tilts – Flat tilt array, shape (n_tilts, 3)

  • tilt_shape – 2D grid shape (ny, nx)

Returns:

Reshaped tilts, shape (ny, nx, 3)

pyextal.jaxEngine.utils.debug_plot(data, num=1, title=None)
pyextal.jaxEngine.utils.result_plot(save_path, prefix, results, config, tilts, experimental_data, variance)
pyextal.jaxEngine.utils.simulate(params: Dict, metadata: Dict, config: pyextal.jaxEngine.setParams.OptConfig, constants: Dict, batch_size: int = 4096) → jax.numpy.ndarray

Simulate 2D intensity map by processing in batches. Convolution is applied on the full image if enabled.

Parameters:

  • params – Optimization parameters (can be normalized or denormalized)

  • metadata – Pre-calculated metadata containing sg_batch, ig_batch, bigkn_batch

  • config – Optimization configuration

  • constants – Bloch wave constants

  • batch_size – Number of tilts per batch

Returns:

Simulated intensities, shape (ny, nx, n_beams)

pyextal.jaxEngine.utils.simulate_with_gradients(params: Dict, metadata: Dict, config: pyextal.jaxEngine.setParams.OptConfig, constants: Dict, batch_size: int = 4096) → Tuple[jax.numpy.ndarray, Dict]

Simulate diffraction patterns AND compute gradients w.r.t. structure factors.

Uses JAX’s automatic differentiation to compute how each output intensity depends on the structure factor parameters.

Parameters:

  • params – Optimization parameters

  • metadata – Pre-calculated metadata

  • config – Optimization configuration

  • constants – Bloch wave constants

  • batch_size – Number of tilts per batch (smaller than simulate due to gradient memory)

  • grad_wrt – What to compute gradients with respect to: ‘beams’ - structure factors only ‘thickness’ - layer thicknesses only ‘all’ - all parameters in params

Returns:

  • intensities: Simulated intensities, shape (ny, nx, n_beams)

  • gradients: Dict with same structure as params, containing gradients

    Each gradient leaf has shape (ny, nx, n_beams, *original_shape)

Return type:

Tuple of

pyextal.jaxEngine.utils.matchIdx(beams, gInclude)