pyextal.jaxEngine.utils
=======================

.. py:module:: pyextal.jaxEngine.utils


Functions
---------

.. autoapisummary::

   pyextal.jaxEngine.utils.reshape_tilts_to_2d
   pyextal.jaxEngine.utils.debug_plot
   pyextal.jaxEngine.utils.result_plot
   pyextal.jaxEngine.utils.simulate
   pyextal.jaxEngine.utils.simulate_with_gradients
   pyextal.jaxEngine.utils.matchIdx


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

.. py:function:: reshape_tilts_to_2d(tilts: jax.numpy.ndarray, tilt_shape: Tuple[int, int]) -> jax.numpy.ndarray

   Reshape flat tilt array to 2D grid

   :param tilts: Flat tilt array, shape (n_tilts, 3)
   :param tilt_shape: 2D grid shape (ny, nx)

   :returns: Reshaped tilts, shape (ny, nx, 3)


.. py:function:: debug_plot(data, num=1, title=None)

.. py:function:: result_plot(save_path, prefix, results, config, tilts, experimental_data, variance)

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

   :param params: Optimization parameters (can be normalized or denormalized)
   :param metadata: Pre-calculated metadata containing sg_batch, ig_batch, bigkn_batch
   :param config: Optimization configuration
   :param constants: Bloch wave constants
   :param batch_size: Number of tilts per batch

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


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

   :param params: Optimization parameters
   :param metadata: Pre-calculated metadata
   :param config: Optimization configuration
   :param constants: Bloch wave constants
   :param batch_size: Number of tilts per batch (smaller than simulate due to gradient memory)
   :param 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)
   :rtype: Tuple of


.. py:function:: matchIdx(beams, gInclude)