SFI.inference.optimizers module

Optimizer back-ends for parametric inference methods.

Provides L-BFGS-B (via SciPy) and Adam (via optax) wrappers with logging, best-parameter tracking, and a unified result interface.

SFI.inference.optimizers.optimize_adam(loss, loss_grad, theta0_flat, *, maxiter, learning_rate, lr_schedule, loss_grad_batch=None, batch_rng_seed=0, batch_schedule=None)

Adam via optax (returns a namespace mimicking SciPy OptimizeResult).

When loss_grad_batch is provided, mini-batch stochastic gradients are used for parameter updates while the full-data loss is still used for tracking / best-parameter selection.

Parameters:

• loss_grad_batch (callable, optional) – loss_grad_batch(theta, rng_key) -> grad. When given, each Adam step uses a stochastic gradient instead of the full-data gradient. Ignored when batch_schedule is set.

• batch_rng_seed (int) – Seed for the mini-batch PRNG stream.

• batch_schedule (list of (float, callable), optional) – Batch-size annealing schedule. Each entry is (step_fraction, grad_fn) where grad_fn has the signature grad_fn(theta, rng_key) -> grad. The list must be sorted by ascending step_fraction; the last entry should have fraction 1.0. At each step the active gradient function is the first whose fraction exceeds step / maxiter.

SFI.inference.optimizers.optimize_lbfgsb(loss, loss_grad, theta0_flat, *, maxiter, tol)

L-BFGS-B via SciPy (returns a SciPy OptimizeResult).