pyextal.jaxEngine.model
=======================

.. py:module:: pyextal.jaxEngine.model


Attributes
----------

.. autoapisummary::

   pyextal.jaxEngine.model.TILT_BATCH
   pyextal.jaxEngine.model._EXPM_THETA
   pyextal.jaxEngine.model._EXPM_MAX_SQUARINGS
   pyextal.jaxEngine.model._PADE7
   pyextal.jaxEngine.model.vmapped_forward
   pyextal.jaxEngine.model.vmapped_forward_reference
   pyextal.jaxEngine.model.vmapped_grad_forward


Functions
---------

.. autoapisummary::

   pyextal.jaxEngine.model.construct_ugh
   pyextal.jaxEngine.model.select_outputs
   pyextal.jaxEngine.model.build_structure_matrix_single
   pyextal.jaxEngine.model.propagate
   pyextal.jaxEngine.model.forward_model_single
   pyextal.jaxEngine.model.structure_matrix_slab
   pyextal.jaxEngine.model._dot
   pyextal.jaxEngine.model.expm_batch
   pyextal.jaxEngine.model._forward_sub_batch
   pyextal.jaxEngine.model.batched_forward
   pyextal.jaxEngine.model.forward_model
   pyextal.jaxEngine.model.forward_model_crystalparams


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

.. py:function:: construct_ugh(params: Dict, metadata: Dict, layer_idx: int = 0, delta=None) -> Tuple[jax.numpy.ndarray, jax.numpy.ndarray]

   Update Ugh matrices with optimizable structure factors.

   ``delta`` -- optional ``(dr, di)``, each (nbeams, nbeams), added to the
   static (non-refined) part before the refined reflections: a model-driven
   change of every coupling the static matrix holds (e.g. a multipole model's
   U_g against the IAM the static matrix was built from). None leaves the
   static matrix as it is.


.. py:function:: select_outputs(phi: jax.numpy.ndarray, indices: jax.numpy.ndarray, axis: int = 1) -> jax.numpy.ndarray

   Gathers the measured beams from the diagonalised ones; 0 where a beam was not computed.

   ``phi`` holds the ``ndiag`` beams diagonalised at each tilt. A measured beam
   that ``beamcheck`` ranks as weak at some tilt sits at a position
   ``>= ndiag`` there, so it has no amplitude, and an out-of-range gather would
   return NaN -- which also poisons every gradient. The index is clipped
   instead and the value set to zero; see `postProcess.excluded_outputs` for
   keeping those entries out of the loss.

   :param phi: (..., ndiag) intensities along ``axis``.
   :param indices: Positions of the measured beams in the diagonalised order.
   :param axis: The beam axis.

   :returns: Intensities of the measured beams, shaped like ``indices``.


.. py:function:: build_structure_matrix_single(sg: jax.numpy.ndarray, ig: jax.numpy.ndarray, ughr: jax.numpy.ndarray, ughi: jax.numpy.ndarray, constants: Dict) -> jax.numpy.ndarray

   Build structure matrix A for a single tilt.

   :param sg: Pre-calculated excitation errors for this tilt, shape (nbeams,)
   :param ig: Pre-calculated ig indices for this tilt, shape (nbeams,)
   :param ughr: Real part of Ugh matrix, shape (nbeams, nbeams)
   :param ughi: Imaginary part of Ugh matrix, shape (nbeams, nbeams)
   :param constants: Dictionary containing Bloch constants

   :returns: Structure matrix, shape (ndiag, ndiag)
   :rtype: A


.. py:function:: propagate(bigkn: float, A: jax.numpy.ndarray, phi_0: jax.numpy.ndarray, thickness: float) -> jax.numpy.ndarray

   Propagate wave through scattering matrix


.. py:function:: forward_model_single(sg: jax.numpy.ndarray, ig: jax.numpy.ndarray, bigkn: float, sf_params: Dict, thickness_params: Dict, metadata: Dict, constants: Dict) -> jax.numpy.ndarray

   Forward model for a SINGLE tilt point.

   :param sg: Excitation errors for this tilt, shape (nbeams,)
   :param ig: ig indices for this tilt, shape (nbeams,)
   :param bigkn: bigkn value for this tilt, scalar
   :param params: Optimization parameters
   :param metadata: Metadata containing pre-calculated values
   :param constants: Bloch wave constants
   :param phi_0: Initial state, shape (ndiag,)

   :returns: Output wave function, shape (ndiag,)
   :rtype: phi_t


.. py:data:: TILT_BATCH
   :value: 1024


.. py:data:: _EXPM_THETA
   :value: 3.92572478313866


.. py:data:: _EXPM_MAX_SQUARINGS
   :value: 16


.. py:data:: _PADE7
   :value: (17297280.0, 8648640.0, 1995840.0, 277200.0, 25200.0, 1512.0, 56.0, 1.0)


.. py:function:: structure_matrix_slab(sg: jax.numpy.ndarray, ig: jax.numpy.ndarray, ughr: jax.numpy.ndarray, ughi: jax.numpy.ndarray, constants: Dict) -> jax.numpy.ndarray

   `build_structure_matrix_single`, gathering only the blocks it uses.

   The Bethe sum reads the diagonalised rows against every beam and the weak
   rows against the diagonalised columns; the weak-weak block, most of the
   matrix, is never needed.


.. py:function:: _dot(a, b)

.. py:function:: expm_batch(X: jax.numpy.ndarray) -> jax.numpy.ndarray

   exp of each matrix in ``X`` (b, n, n), sharing one squaring count.

   The count is the largest any matrix in the batch needs; a matrix needing
   fewer is scaled further, which Pade-7 handles at least as accurately. The
   branch on it is real (it is one scalar), so unused squarings cost nothing.
   Returns NaN if a matrix needs more than 16 squarings, as `expm` does.


.. py:function:: _forward_sub_batch(sg, ig, bigkn, ughs, thicknesses, constants)

   |phi|^2 for one sub-batch of tilts, (b, ndiag).


.. py:function:: batched_forward(sg_batch, ig_batch, bigkn_batch, sf_params, thickness_params, metadata, constants, ugh_delta=None)

   `forward_model_single` over every tilt, in rematerialised sub-batches.

   Same signature and result as ``jax.vmap(forward_model_single, ...)``; see
   the notes above `TILT_BATCH`. ``ugh_delta`` -- optional per-layer
   ``(dr, di)`` changes of the static Ugh matrix (`construct_ugh`).


.. py:data:: vmapped_forward

.. py:data:: vmapped_forward_reference

.. py:data:: vmapped_grad_forward

.. py:function:: forward_model(params: Dict, metadata: Dict, constants: Dict) -> jax.numpy.ndarray

   Vectorized forward model for multiple tilts.

   Uses vmap over forward_model_single for efficient batch processing.


.. py:function:: forward_model_crystalparams(params: Dict, metadata: Dict, constants: Dict, model_fn: callable) -> jax.numpy.ndarray

   Forward model for refining crystal parameters.

   :param params: Parameters including 'model_param' and 'beams'
   :param metadata: Metadata
   :param constants: Bloch wave constants
   :param phi_0_batch: Initial states
   :param model_fn: Model function that computes structure factors from crystal params

   :returns: Output wave functions, shape (n_tilts, ndiag)
   :rtype: phi_batch