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

.. py:module:: pyextal.jaxJoint.params

.. autoapi-nested-parse::

   The joint parameter tree, and each dataset's view of it.

   ::

       G = {'shared':   [OrderedDict(key -> {xfre, xfim, afre, afim})  per layer],
            'datasets': [{thickness, intensity_scale, background, probe_sigma,
                          probe_gamma[, oversub_offset, oversub_leak]}  per dataset]}

   ``shared`` holds one entry per canonical key (`jaxJoint.link`), in physical
   units: the structure factors are not normalised (``reflections_normalize`` is
   off), so a dataset's values are a fixed linear map of the shared ones. The
   per-dataset scalars stay normalised with that dataset's own bounds.

   `local_params` rebuilds exactly the tree `setup_optimization_params` made for a
   dataset -- same keys, same order -- so its loss, the geometry objective, the
   simulation and `save_refinement` all take it unchanged. It is a JAX function of
   ``G``, so a dataset's gradient with respect to ``G`` comes straight through it.



Attributes
----------

.. autoapisummary::

   pyextal.jaxJoint.params.COMPONENTS


Functions
---------

.. autoapisummary::

   pyextal.jaxJoint.params._term
   pyextal.jaxJoint.params.apply_link
   pyextal.jaxJoint.params.local_params
   pyextal.jaxJoint.params.shared_value
   pyextal.jaxJoint.params.freeze_mask


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

.. py:data:: COMPONENTS
   :value: ('xfre', 'xfim', 'afre', 'afim')


.. py:function:: _term(coef, x)

.. py:function:: apply_link(m, value)

   ``M @ value`` for one structure factor.

   Zero entries are skipped and +-1 entries are a copy or a negation, so a
   phase of 0 or pi links values exactly and a single dataset's refinement is
   bit for bit the one `jaxGeo` runs.

   :param m: (4, 4) link map.
   :type m: sequence
   :param value: ``xfre``/``xfim``/``afre``/``afim``.
   :type value: dict

   :returns: OrderedDict


.. py:function:: local_params(G, index, links)

   Dataset ``index``'s parameter tree, as `setup_optimization_params` builds it.

   :param G: The joint tree.
   :type G: dict
   :param index: The dataset.
   :type index: int
   :param links: ``(local key, canonical key, M)`` in the dataset's own
                 reflection order (``metadata['reflection_list']``).
   :type links: list

   :returns: ``beams`` (one OrderedDict per layer) and the dataset's scalars.
   :rtype: dict


.. py:function:: shared_value(m, value)

   The canonical value, ``M.T @ value`` (M is orthogonal).


.. py:function:: freeze_mask(G, configs, locals_, sf_mask, frozen_layers)

   Which leaves of ``G`` Adam holds (True) -- the joint `create_freeze_mask`.

   :param G: The joint tree.
   :type G: dict
   :param configs: Each dataset's OptConfig.
   :type configs: list
   :param locals_: Each dataset's `local_params` (for the scalar flags).
   :type locals_: list
   :param sf_mask: Which of (xfre, xfim, afre, afim) are free.
   :type sf_mask: tuple
   :param frozen_layers: 0-based layers whose structure factors are held.
   :type frozen_layers: tuple

   :returns: Same structure as ``G``.
   :rtype: dict