pyextal.jaxEngine.postProcess
=============================

.. py:module:: pyextal.jaxEngine.postProcess


Attributes
----------

.. autoapisummary::

   pyextal.jaxEngine.postProcess.convolved_2d_vmap


Functions
---------

.. autoapisummary::

   pyextal.jaxEngine.postProcess.create_gaussian_kernel
   pyextal.jaxEngine.postProcess.create_pseudo_voigt_kernel
   pyextal.jaxEngine.postProcess.post_process_intensity
   pyextal.jaxEngine.postProcess.excluded_outputs
   pyextal.jaxEngine.postProcess.apply_oversubtraction
   pyextal.jaxEngine.postProcess.clip_experimental


Module Contents
---------------

.. py:function:: create_gaussian_kernel(radius: int, sigma: float) -> jax.numpy.ndarray

   Create 2D Gaussian convolution kernel


.. py:function:: create_pseudo_voigt_kernel(radius: int, sigma: float, gamma: float) -> jax.numpy.ndarray

   Create 2D Pseudo-Voigt convolution kernel

   :param radius: Kernel radius (kernel will be (2*radius+1) x (2*radius+1))
   :param sigma: Gaussian width parameter
   :param gamma: Lorentzian width parameter

   :returns: 2D Pseudo-Voigt kernel normalized to sum to 1


.. py:data:: convolved_2d_vmap

.. py:function:: 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


.. py:function:: 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.

   :param indices: (n_pixels, n_out) positions of the measured beams.
   :param ndiag: Number of diagonalised beams.
   :param config: Supplies ``use_voigt`` and ``kernel_radius``.
   :param chunk_shape: (rows, cols) of the chunk.

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


.. py:function:: 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.

   :param simulated: (..., n_beams) intensities after scale and background.
   :param params: Denormalised parameters holding ``oversub_offset`` and
                  ``oversub_leak``.
   :param config: Supplies the channel lists.
   :param experimental: (..., n_beams) measured intensities on the same pixels.

   :returns: Same shape as ``simulated``.


.. py:function:: 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.

   :param experimental: (..., n_beams) measured intensities.
   :param config: Supplies ``oversub_indices``.

   :returns: Same shape.