pyextal.jaxEngine.optimize
==========================

.. py:module:: pyextal.jaxEngine.optimize


Classes
-------

.. autoapisummary::

   pyextal.jaxEngine.optimize.BatchGenerator
   pyextal.jaxEngine.optimize.ChunkingGenerator
   pyextal.jaxEngine.optimize.RandomGenerator


Functions
---------

.. autoapisummary::

   pyextal.jaxEngine.optimize.get_loss_fn
   pyextal.jaxEngine.optimize._freeze_layers
   pyextal.jaxEngine.optimize.create_freeze_mask
   pyextal.jaxEngine.optimize.optimize


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

.. py:class:: BatchGenerator(config: pyextal.jaxEngine.setParams.OptConfig, tilts: jax.numpy.ndarray, experimental_data: jax.numpy.ndarray, variance: jax.numpy.ndarray, metadata: Dict)

   .. py:attribute:: config


   .. py:attribute:: tilts


   .. py:attribute:: experimental_data


   .. py:attribute:: variance


   .. py:attribute:: metadata


   .. py:attribute:: n_samples


   .. py:method:: get_batches(epoch: int) -> List[Dict]
      :abstractmethod:



.. py:class:: ChunkingGenerator(config: pyextal.jaxEngine.setParams.OptConfig, tilts: jax.numpy.ndarray, experimental_data: jax.numpy.ndarray, variance: jax.numpy.ndarray, metadata: Dict)

   Bases: :py:obj:`BatchGenerator`


   .. py:attribute:: _cached_batches
      :value: None



   .. py:method:: get_batches(epoch: int) -> List[Dict]


   .. py:method:: create_2d_tilt_chunks() -> List[Tuple[slice, slice]]


.. py:class:: RandomGenerator(config: pyextal.jaxEngine.setParams.OptConfig, tilts: jax.numpy.ndarray, experimental_data: jax.numpy.ndarray, variance: jax.numpy.ndarray, metadata: Dict, batch_size: int = 64)

   Bases: :py:obj:`BatchGenerator`


   .. py:attribute:: batch_size
      :value: 64



   .. py:attribute:: experimental_data


   .. py:attribute:: variance


   .. py:method:: get_batches(epoch: int) -> List[Dict]


.. py:function:: get_loss_fn(config: pyextal.jaxEngine.setParams.OptConfig, constants: Dict, metadata: Dict, model_fn: callable = None) -> Callable

   Create loss function with model_fn bound via closure.

   :param config: Optimization configuration
   :param constants: Bloch constants
   :param metadata: Pre-calculated metadata
   :param model_fn: The model function (for crystal parameter refinement)


.. py:function:: _freeze_layers(mask: Dict, config: pyextal.jaxEngine.setParams.OptConfig) -> None

   Freezes every structure factor of the layers in ``config.frozen_layers``.


.. py:function:: create_freeze_mask(config: pyextal.jaxEngine.setParams.OptConfig, params: Dict) -> Dict

   Create mask for freezing/unfreezing parameters during optimization.


.. py:function:: optimize(tilts: jax.numpy.ndarray, experimental_data: jax.numpy.ndarray, variance: jax.numpy.ndarray, config: pyextal.jaxEngine.setParams.OptConfig, constants: Dict, initial_params: Dict, metadata: Dict, opt_state: optax.OptState = None) -> Tuple[Dict, List[float], List[Dict], optax.OptState]