pyextal.jaxEngine.postProcess

Attributes

convolved_2d_vmap

Functions

create_gaussian_kernel(→ jax.numpy.ndarray)

Create 2D Gaussian convolution kernel

create_pseudo_voigt_kernel(→ jax.numpy.ndarray)

Create 2D Pseudo-Voigt convolution kernel

post_process_intensity(→ jax.numpy.ndarray)

Convert beam amplitudes to 2D intensity map with optional gaussian convolution

excluded_outputs(→ jax.numpy.ndarray)

Entries of a chunk the loss must ignore because a beam was not simulated.

apply_oversubtraction(→ jax.numpy.ndarray)

Applies the integration's over-subtraction and clip to the simulation.

clip_experimental(→ jax.numpy.ndarray)

Clips the over-subtracted channels of the data at zero.

Module Contents

pyextal.jaxEngine.postProcess.create_gaussian_kernel(radius: int, sigma: float) → jax.numpy.ndarray

Create 2D Gaussian convolution kernel

pyextal.jaxEngine.postProcess.create_pseudo_voigt_kernel(radius: int, sigma: float, gamma: float) → jax.numpy.ndarray

Create 2D Pseudo-Voigt convolution kernel

Parameters:

  • radius – Kernel radius (kernel will be (2*radius+1) x (2*radius+1))

  • sigma – Gaussian width parameter

  • gamma – Lorentzian width parameter

Returns:

2D Pseudo-Voigt kernel normalized to sum to 1

pyextal.jaxEngine.postProcess.convolved_2d_vmap
pyextal.jaxEngine.postProcess.post_process_intensity(phi_g_batch: jax.numpy.ndarray, params: jax.numpy.ndarray, config: pyextal.jaxEngine.utils.OptConfig, metadata: dict, chunk_shape: tuple[int, int]) → jax.numpy.ndarray

Convert beam amplitudes to 2D intensity map with optional gaussian convolution

pyextal.jaxEngine.postProcess.excluded_outputs(indices: jax.numpy.ndarray, ndiag: int, config: pyextal.jaxEngine.utils.OptConfig, chunk_shape: tuple[int, int]) → jax.numpy.ndarray

Entries of a chunk the loss must ignore because a beam was not simulated.

Where a measured beam falls outside the diagonalised set at its tilt, model.select_outputs gives 0 – not a prediction, so it is kept out of the loss. With the probe convolution on, every entry within kernel_radius of such a pixel has that zero in its convolution and is excluded as well.

Parameters:

  • indices – (n_pixels, n_out) positions of the measured beams.

  • ndiag – Number of diagonalised beams.

  • config – Supplies use_voigt and kernel_radius.

  • chunk_shape – (rows, cols) of the chunk.

Returns:

(rows, cols, n_out) bool, True where excluded.

pyextal.jaxEngine.postProcess.apply_oversubtraction(simulated: jax.numpy.ndarray, params: dict, config: pyextal.jaxEngine.utils.OptConfig, experimental: jax.numpy.ndarray) → jax.numpy.ndarray

Applies the integration’s over-subtraction and clip to the simulation.

For each channel in config.oversub_indices the observable is max(sim - offset - leak * I_leak, 0): the background the integration subtracted was too large by a constant plus a share of a neighbour’s intensity, and ApertureSum clipped the result at zero. I_leak is the measured intensity of channel config.oversub_leak_index – the light that actually reached the annulus. Where the data is zero and the model is at or below zero the residual vanishes, which is the right treatment of a censored measurement. Other channels are returned unchanged.

Parameters:

  • simulated – (…, n_beams) intensities after scale and background.

  • params – Denormalised parameters holding oversub_offset and oversub_leak.

  • config – Supplies the channel lists.

  • experimental – (…, n_beams) measured intensities on the same pixels.

Returns:

Same shape as simulated.

pyextal.jaxEngine.postProcess.clip_experimental(experimental: jax.numpy.ndarray, config: pyextal.jaxEngine.utils.OptConfig) → jax.numpy.ndarray

Clips the over-subtracted channels of the data at zero.

The integration clipped them at zero; any negative values left are from resampling across that edge, and the model cannot go below zero.

Parameters:

  • experimental – (…, n_beams) measured intensities.

  • config – Supplies oversub_indices.

Returns:

Same shape.