pyextal.jaxGeo.lm
=================

.. py:module:: pyextal.jaxGeo.lm

.. autoapi-nested-parse::

   Levenberg-Marquardt on the geometry, everything else held fixed.

   Why Levenberg-Marquardt and not Adam
   ------------------------------------
   There are four parameters of very different stiffness -- on tilted SrTiO3 data
   the Hessian diagonal spans a factor of ~400 per unit, ``tiltx`` stiffest and
   ``rotation`` softest -- and some are coupled (``gl`` with ``tilty`` at +0.6
   there; ``gl`` also trades against thickness, which is held fixed here). A
   first-order method with one learning rate handles that badly. With four
   parameters the full Hessian is cheap -- nine gradient evaluations by central
   differences, or four Hessian-vector products exactly -- and a damped Newton step
   then converges in a handful of iterations, including along correlated
   directions, where a pattern search, stepping one axis at a time, is slowest.

   The damping is the usual one: solve ``(H + lambda diag|H|) dz = -g``, accept the
   step if the loss falls (and shrink ``lambda``), otherwise grow ``lambda`` and try
   again. Steps are additionally capped at ``trust`` units per parameter so the
   frozen partition stays representative. The parameters are scaled so that one
   unit is a chosen step (`jaxGeo.geometry.units_for`), which puts them on
   comparable footing for both the cap and the solve.

   Freezing the beam partition is the standard active-set treatment. The partition
   only changes when a beam crosses the excitation cut, which moves the loss by a
   Bethe-perturbation-sized amount, so it is re-run (``beamcheck``) only when the
   geometry has moved ``refresh`` units, and once more at the end. The result after
   that final refresh is the one that is kept, and it is kept only if it beats the
   starting loss. Acceptance is judged on the training tiling, since a different
   tiling averages over slightly different pixels.

   The Hessian at the end gives the parameter correlations, which are printed:
   strongly correlated pairs are the ones the data does not separate.



Functions
---------

.. autoapisummary::

   pyextal.jaxGeo.lm.print_correlations
   pyextal.jaxGeo.lm.geometry_refine
   pyextal.jaxGeo.lm.check_gradient


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

.. py:function:: print_correlations(hess, names)

   Prints the parameter correlation matrix implied by a Hessian.

   :param hess: (k, k) Hessian in scaled units.
   :type hess: np.ndarray
   :param names: Parameter names, in order.
   :type names: list


.. py:function:: geometry_refine(x0, metadata, constants, make_obj, rebuild, epoch_loss_at, active, *, hessian='exact', lam=0.001, iters=8, trust=1.0, tol=1e-05, refresh=2.0)

   Levenberg-Marquardt on the geometry, everything else held fixed.

   :param x0: (4,) current geometry vector.
   :type x0: np.ndarray
   :param metadata: Metadata for ``x0``.
   :type metadata: dict
   :param constants: Bloch constants; ``ndiag`` is updated by rebuilds and
                     restored if the result is rejected.
   :type constants: dict
   :param make_obj: ``make_obj(x_ref, metadata) -> GeometryObjective``,
                    with the current parameters held fixed.
   :type make_obj: callable
   :param rebuild: ``rebuild(x) -> metadata`` -- re-runs ``beamcheck``
                   for geometry ``x``.
   :type rebuild: callable
   :param epoch_loss_at: ``epoch_loss_at(x, metadata) -> float`` -- the
                         training loss, on the training tiling.
   :type epoch_loss_at: callable
   :param active: Indices of the refined parameters.
   :type active: np.ndarray
   :param hessian: ``'exact'`` or ``'fd'``.
   :type hessian: str
   :param lam: Initial damping.
   :type lam: float
   :param iters: Iteration limit.
   :type iters: int
   :param trust: Largest step per iteration, units per parameter.
   :type trust: float
   :param tol: Loss improvement below which the refinement stops, and by
               which the result must beat the start to be kept.
   :type tol: float
   :param refresh: Re-run ``beamcheck`` once the geometry has moved this
                   many units from where the partition was built.
   :type refresh: float

   :returns: ``(x, metadata, loss, evals)``, ``loss`` on the training tiling
             and ``evals`` counting forward, gradient, HVP and rebuild calls.
   :rtype: tuple


.. py:function:: check_gradient(obj, make_obj, rebuild, names, h=0.05)

   Compares the frozen-partition gradient with the real (rebuilt) loss.

   Central differences of the full pipeline -- ``beamcheck`` rebuild included --
   at ``h`` units either side, against the JAX gradient at the same point, both
   on the objective's tiling. The two agree if freezing the partition is
   harmless at this scale.

   :param obj: Built at the current geometry.
   :type obj: GeometryObjective
   :param make_obj: As for :func:`geometry_refine`.
   :type make_obj: callable
   :param rebuild: As for :func:`geometry_refine`.
   :type rebuild: callable
   :param names: Parameter names.
   :type names: list
   :param h: Difference step, in units.
   :type h: float