pyextal.jaxJoint.optimize
=========================

.. py:module:: pyextal.jaxJoint.optimize

.. autoapi-nested-parse::

   Adam on the joint parameter tree.

   `jaxEngine.optimize.optimize` for several datasets: each step evaluates every
   chunk of every dataset, averages each dataset's chunk losses and gradients as
   the single-dataset ``per_epoch`` mode does, and combines the datasets with their
   weights into one loss and one gradient with respect to ``G``. A dataset's
   gradient reaches the shared structure factors through `local_params` (its
   links) and touches only its own scalars, so one Adam update moves the shared
   U_g by the pooled evidence and each dataset's thickness, scale, background and
   probe by its own.

   The schedule mirrors `optimize` exactly -- the Adam state is reset at step
   ``warmup`` and snapshots are kept at the same steps -- so a joint run of a
   single dataset is the `jaxGeo` run (`jaxEngine.hdf5io.snapshot_steps` applies).



Functions
---------

.. autoapisummary::

   pyextal.jaxJoint.optimize._dataset_step_fn
   pyextal.jaxJoint.optimize._chunks
   pyextal.jaxJoint.optimize.build_optimizer
   pyextal.jaxJoint.optimize.joint_optimize


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

.. py:function:: _dataset_step_fn(ds)

   Jitted value-and-gradient, with respect to ``G``, of one chunk of ``ds``.


.. py:function:: _chunks(ds)

   The dataset's training chunks, as `optimize`'s ChunkingGenerator cuts them.


.. py:function:: build_optimizer(G, datasets, lr, sf_mask, frozen_layers)

   Adam with the joint freeze mask.


.. py:function:: joint_optimize(G, datasets, *, steps, lr, sf_mask, frozen_layers, warmup=0, offset=0, opt_state=None)

   Runs ``steps`` Adam steps on the pooled loss.

   :param G: The joint tree.
   :type G: dict
   :param datasets: `Dataset` objects with current ``config``,
                    ``metadata`` and ``tilts``; each one's ``weight`` scales its loss.
   :type datasets: list
   :param steps: Steps to run.
   :type steps: int
   :param lr: Adam step size.
   :type lr: float
   :param sf_mask: Free components of the structure factors.
   :type sf_mask: tuple
   :param frozen_layers: Layers whose structure factors are held.
   :type frozen_layers: tuple
   :param warmup: The Adam state is reset at this (absolute) step, as
                  `optimize` does; 0 for never.
   :type warmup: int
   :param offset: Steps already run, for the warmup and the log.
   :type offset: int
   :param opt_state: Adam state to continue from; a fresh one when None.

   :returns: ``(G, loss_history, per_dataset_history, snapshots, opt_state)``.
             ``per_dataset_history`` holds each dataset's own (unweighted) loss.
   :rtype: tuple