pyextal.jaxGeo.lm

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

print_correlations(hess, names)

Prints the parameter correlation matrix implied by a Hessian.

geometry_refine(x0, metadata, constants, make_obj, ...)

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

check_gradient(obj, make_obj, rebuild, names[, h])

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

Module Contents

pyextal.jaxGeo.lm.print_correlations(hess, names)

Prints the parameter correlation matrix implied by a Hessian.

Parameters:

  • hess (np.ndarray) – (k, k) Hessian in scaled units.

  • names (list) – Parameter names, in order.

pyextal.jaxGeo.lm.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.

Parameters:

  • x0 (np.ndarray) – (4,) current geometry vector.

  • metadata (dict) – Metadata for x0.

  • constants (dict) – Bloch constants; ndiag is updated by rebuilds and restored if the result is rejected.

  • make_obj (callable) – make_obj(x_ref, metadata) -> GeometryObjective, with the current parameters held fixed.

  • rebuild (callable) – rebuild(x) -> metadata – re-runs beamcheck for geometry x.

  • epoch_loss_at (callable) – epoch_loss_at(x, metadata) -> float – the training loss, on the training tiling.

  • active (np.ndarray) – Indices of the refined parameters.

  • hessian (str) – 'exact' or 'fd'.

  • lam (float) – Initial damping.

  • iters (int) – Iteration limit.

  • trust (float) – Largest step per iteration, units per parameter.

  • tol (float) – Loss improvement below which the refinement stops, and by which the result must beat the start to be kept.

  • refresh (float) – Re-run beamcheck once the geometry has moved this many units from where the partition was built.

Returns:

(x, metadata, loss, evals), loss on the training tiling and evals counting forward, gradient, HVP and rebuild calls.

Return type:

tuple

pyextal.jaxGeo.lm.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.

Parameters: