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

.. py:module:: pyextal.jaxJoint

.. autoapi-nested-parse::

   Joint LARBED refinement of several datasets on top of `jaxGeo`.

   Several datasets of one crystal -- other scans, other camera lengths, other zone
   axes -- are refined together. Their structure factors are one shared set of
   parameters, while each dataset keeps its own geometry (``gl``, ``rotation``,
   ``tiltx``, ``tilty``), thickness, scale, background, probe and over-subtraction
   model. Adam minimises the pooled chi^2 of all datasets; between Adam segments
   each dataset's geometry is refined by `jaxGeo`'s Levenberg-Marquardt, on its own
   loss, with everything else held fixed.

   Usage from Python::

       from pyextal.jaxJoint import JointSettings, refine
       from pyextal.jaxGeo import RefineSettings
       result = refine(JointSettings(
           common=RefineSettings(steps=800, tile=180, reflections='exported'),
           datasets=[dict(h5=R5), dict(h5=R4, subsample=2)]))

   or from a script, with the full command line::

       from pyextal.jaxJoint import cli
       raise SystemExit(cli.main(doc=__doc__, datasets=[dict(h5=R5), dict(h5=R4)],
                                 steps=800, tile=180))

   See `jaxJoint.refine` for the method.

   Importing this package sets the same JAX allocator defaults as `jaxGeo` (it
   imports `jaxGeo` first), so import ``jaxJoint`` before JAX.



Submodules
----------

.. toctree::
   :maxdepth: 1

   /autoapi/pyextal/jaxJoint/cli/index
   /autoapi/pyextal/jaxJoint/datasets/index
   /autoapi/pyextal/jaxJoint/engine/index
   /autoapi/pyextal/jaxJoint/link/index
   /autoapi/pyextal/jaxJoint/optimize/index
   /autoapi/pyextal/jaxJoint/params/index
   /autoapi/pyextal/jaxJoint/refine/index
   /autoapi/pyextal/jaxJoint/report/index
   /autoapi/pyextal/jaxJoint/settings/index


Attributes
----------

.. autoapisummary::

   pyextal.jaxJoint.JOINT_FIELDS


Classes
-------

.. autoapisummary::

   pyextal.jaxJoint.JointSettings


Functions
---------

.. autoapisummary::

   pyextal.jaxJoint.refine


Package Contents
----------------

.. py:class:: JointSettings

   Options for `jaxJoint.refine.refine`.

   :param common: Options every dataset starts from, including all
                  of `JOINT_FIELDS`.
   :type common: RefineSettings
   :param datasets: One entry per dataset: a dict of the fields that
                    differ from ``common`` (at least ``h5``), a full `RefineSettings`,
                    or just the dataset path.
   :type datasets: sequence
   :param names: Display names; default `display_names`.
   :type names: sequence, optional
   :param weighting: ``'pixels'`` weights each dataset's mean chi^2 by its number
                     of data points, so the loss is the chi^2 of all of them pooled;
                     ``'equal'`` averages the datasets' chi^2; a sequence gives explicit
                     relative weights.
   :param out_dir: Where the joint summary goes; default the
                   first dataset's folder.
   :type out_dir: Path, optional


   .. py:attribute:: common
      :type:  pyextal.jaxGeo.settings.RefineSettings


   .. py:attribute:: datasets
      :type:  Tuple
      :value: ()



   .. py:attribute:: names
      :type:  Optional[Tuple[str, ...]]
      :value: None



   .. py:attribute:: weighting
      :type:  Union[str, Tuple[float, ...]]
      :value: 'pixels'



   .. py:attribute:: out_dir
      :type:  Optional[pathlib.Path]
      :value: None



   .. py:method:: resolved()

      Checks the options and returns a canonical copy.

      :returns: ``common`` resolved, ``datasets`` as a tuple of
                resolved `RefineSettings`, ``names`` filled in and ``weighting`` as
                'pixels', 'equal' or a tuple of floats.
      :rtype: JointSettings

      :raises SetupError: On an inconsistent option, a dataset setting a joint
          field, or the same file twice.



.. py:data:: JOINT_FIELDS
   :value: ('symmetry', 'imag', 'layers', 'refine_layers', 'steps', 'warmup', 'lr', 'geometry',...


.. py:function:: refine(settings=None, **overrides)

   Runs a joint refinement, or validates its setup with ``dry_run``.

   :param settings: The options.
   :type settings: JointSettings, optional
   :param \*\*overrides: Fields to change on ``settings``.

   :returns: ``G`` (the joint tree), ``datasets`` (the `Dataset` objects),
             ``table`` (shared structure factors), ``engine`` and, unless
             ``dry_run``, ``loss_history``, ``runs`` (dataset name -> run key) and
             ``summary`` (the JSON path).
   :rtype: dict

   :raises SetupError: If the datasets cannot be refined together.