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

.. py:module:: pyextal.jaxJoint.refine

.. autoapi-nested-parse::

   The joint refinement: shared structure factors, per-dataset geometry.

   This is `jaxGeo.refine.refine` for several datasets of one crystal at once.

   What is shared and what is not
   ------------------------------
   The structure factors are one set of parameters for all datasets. Everything
   that describes a particular measurement is per dataset: the geometry (``gl``,
   ``rotation``, ``tiltx``, ``tilty``), the thickness (per layer), the intensity
   scale and background, the probe, and the over-subtraction model. Each dataset
   keeps its own ``.dat``, beam list, reflection choice, tiling and geometry-step
   settings; the options that decide the shared parameters or the one optimiser's
   schedule (`jaxJoint.settings.JOINT_FIELDS`: ``imag``, ``symmetry``, ``layers``,
   ``steps``, ``lr``, ``geo_start``, ...) are common.

   Which structure factors (`jaxJoint.link`)
   -----------------------------------------
   Each dataset's ``reflections`` choice ('exported', 'review' or a list) is mapped
   to canonical keys, and the union is refined. A refined U_g is used in every
   dataset whose structure matrix holds it -- measured there or not -- so the
   crystal has one potential. The engine labels a structure factor by a member of
   its orbit that depends on the beam list, and equivalent members differ by the
   phase of the space group's translations (a sign, for a glide or screw) or by
   conjugation (Friedel mate); each dataset's labels are linked to the shared value
   by those maps, which are checked against the engine's own structure factors
   before anything runs. The run starts with a table of the links.

   The loss
   --------
   Adam minimises ``sum_d w_d chi2_d``, each ``chi2_d`` the dataset's mean chi^2
   exactly as `jaxGeo` computes it. ``weighting='pixels'`` (default) makes ``w_d``
   proportional to the dataset's data points, so the loss is the chi^2 of all of
   them pooled; ``'equal'`` averages the datasets; explicit weights are allowed.
   The log prints the pooled loss and each dataset's own chi^2.

   Geometry
   --------
   At ``geo_start`` and every ``geo_every`` steps Adam stops and each dataset's
   geometry is refined in turn by `jaxGeo`'s Levenberg-Marquardt, on that
   dataset's loss alone (the others do not depend on it), with the structure
   factors and everything else held. Tiling, Hessian method and the gradient check
   are chosen per dataset.

   The engine
   ----------
   The Fortran engine holds one ``.dat`` at a time, and ``beamcheck`` -- run at
   setup and at every geometry rebuild -- reads it. Every such step first
   re-parses its dataset's ``.dat`` (`jaxJoint.engine`). Paths over 120 characters
   are reached through a short symlink.

   Output
   ------
   Each dataset gets a new ``run_NNN`` in its own file, with its own loss history,
   its parameters as that dataset sees them (its labels, its signs), the
   simulation and the figures `jaxGeo` writes; its notes name the joint run and
   the other datasets. ``joint_<stamp>.json`` in ``out_dir`` holds the settings,
   the datasets and their runs, the weights, the loss histories, the geometries
   and the shared structure-factor table.



Functions
---------

.. autoapisummary::

   pyextal.jaxJoint.refine._select_cpu
   pyextal.jaxJoint.refine.refine_geometry
   pyextal.jaxJoint.refine._release_oversub
   pyextal.jaxJoint.refine._figures
   pyextal.jaxJoint.refine._finish
   pyextal.jaxJoint.refine.refine


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

.. py:function:: _select_cpu()

   Forces the CPU backend, before or after JAX has been imported.


.. py:function:: refine_geometry(ds, G, engine, done)

   One Levenberg-Marquardt geometry refinement of dataset ``ds``.

   :param ds: Its ``x``, ``metadata`` and ``tilts`` are updated.
   :type ds: Dataset
   :param G: The joint tree, held fixed.
   :type G: dict
   :param engine: Switched to ``ds`` for the ``beamcheck`` rebuilds.
   :type engine: EngineSwitch
   :param done: The Adam step, for the history.
   :type done: int


.. py:function:: _release_oversub(ds, G)

   Fits ``ds``'s over-subtraction model and writes it into ``G``.


.. py:function:: _figures(ds, run, params, sim, exp_fig, common)

   The comparison, history and layer/phase output `jaxGeo` writes for a run.


.. py:function:: _finish(ds, G, snaps, hist_steps, stamp, others, common)

   Simulates, saves and draws one dataset; returns its run key or None.


.. 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.