Optimizer

class nyx.infer.optimizer.Optimizer

Bases: object

Fit any pytree of Parameters with a minimiser or least-squares solver.

Trains every non-frozen Parameter reachable from model. Per-obs parameters get independent per-observation gradients, global ones the sum over observations.

Parameters:

• fn (callable) – (model) -> scalar for a minimiser, or (model) -> residuals for a least-squares solver. A non-scalar fn given to a minimiser is wrapped with sum-of-squares.

• solver (optimistix.AbstractIterativeSolver) – E.g. optx.LBFGS or optx.LevenbergMarquardt.

Raises:

TypeError – If a scalar fn is paired with a least-squares solver.

Methods Summary

errors(fitted, *[, batch_size, ...])	1-sigma errors on every trainable Parameter at fitted.
init_state(model)	Initial solver state for manual stepping.
loss(model)	Scalar loss at model, without taking a solver step.
profile(model, scan, *[, max_steps, callback])	Profile likelihood over one or two parameters.
run(model, *[, max_steps, throw])	Run the solver to convergence, inside a single compiled loop.
step(model, state)	One solver step.
summary(model[, sol])	Goodness of fit at model.

Methods Documentation

errors(fitted, *, batch_size=8, reduced_chi2=False, rcond=None)

1-sigma errors on every trainable Parameter at fitted.

Valid only when fn returns residuals; run() need not have been called.

Parameters:

• fitted (pytree)

• batch_size (int, optional) – Jacobian columns evaluated per kernel launch.

• reduced_chi2 (bool, optional) – Whether to apply the (r^T r) / (m - n) factor. False assumes residuals pre-scaled by 1-sigma uncertainties.

• rcond (float or None, optional) – Pseudoinverse cutoff, relative to the largest singular value.

Returns:

sqrt(diag(pinv(J^T J))), shaped like the trainable subset of fitted.

Return type:

pytree

init_state(model)

Initial solver state for manual stepping.

Parameters:

model (pytree)

Returns:

Pass as state to the first step() call.

Return type:

optimistix solver state

loss(model)

Scalar loss at model, without taking a solver step.

For a residuals fn this is sum(r ** 2).

Parameters:

model (pytree)

Return type:

jax.Array

profile(model, scan, *, max_steps=256, callback=None)

Profile likelihood over one or two parameters.

At each grid point the scanned parameters are held and the rest refitted.

Parameters:

• model (pytree) – With its nuisance parameters already unfrozen.

• scan (dict of str to array-like) – One or two parameter names, each with the values to scan.

• max_steps (int, optional) – Maximum solver iterations per refit.

• callback (callable, optional) – (index, point, chi2) -> None, for a progress bar.

Return type:

ProfileGrid

run(model, *, max_steps=256, throw=True)

Run the solver to convergence, inside a single compiled loop.

Parameters:

• model (pytree)

• max_steps (int, optional) – Maximum solver iterations.

• throw (bool, optional) – Whether to raise on non-successful termination; False leaves sol.result to inspect.

Returns:

• model (pytree) – Fitted model.

• sol (optimistix.Solution)

step(model, state)

One solver step.

Parameters:

• model (pytree)

• state (optimistix solver state) – From init_state() or a previous step().

Returns:

• model (pytree) – Updated model.

• loss (jax.Array) – Scalar loss at model before the step was applied.

• state (optimistix solver state)

summary(model, sol=None)

Goodness of fit at model.

Parameters:

• model (pytree)

• sol (optimistix.Solution, optional) – Adds the solver’s status and step count.

Return type:

FitSummary

__init__(fn, solver)

classmethod __new__(*args, **kwargs)