Fitter

Fitter#

class braintools.fit.Fitter(simulate, objective, space, *, init=None, optimizer=None, backend='grad', transient=None)#

Fit model parameters to data by minimising an objective.

Functional mode (this constructor) fits a pure function; module mode (from_module()) fits the trainable parameters of a brainstate.nn.Module.

Parameters:

simulate (Callable[[Dict[str, Any]], Any]) – simulate(params) -> prediction, a pure JAX function of a dict {name: value} (values carry the units of space). It must be traceable by jax.jit/jax.vmap/jax.grad ('grad' needs gradients; the other backends only need vmap).
objective (Callable[[Any, Any], Any]) – objective(prediction, target) -> scalar, e.g. an Objective or a loss from braintools.fit.losses.
space (Dict[str, Any]) – {name: (low, high)} box bounds (scalars, arrays or Quantities; low or high may be None for a one-sided bound) or {name: None} for an unbounded parameter whose value comes from init. The shape of a parameter is the broadcast shape of its bounds and initial value.
init (Dict[str, Any] | None) – Initial values {name: value}, used for unbounded parameters (required there), for one-sided bounds (required there), and as the single start when n_starts=1 and no key is given (box parameters default to the midpoint of their box).
optimizer (Any) – Backend options. A string names the method ('lbfgs'/'adam'/ 'gd' for 'grad', a scipy.optimize.minimize() method for 'scipy', a Nevergrad optimiser such as 'DE' or 'CMA' for 'nevergrad'). A dict holds 'method' plus options: 'tol' and 'learning_rate' for 'grad'; 'tol' plus scipy.optimize.minimize() options for 'scipy'; 'n_sample' (population per generation, default n_starts) plus NevergradOptimizer keyword arguments for 'nevergrad'. For 'grad' an optax transformation or a braintools.optim optax optimiser (its tx) is used as the update rule. Defaults: L-BFGS, 'L-BFGS-B', 'DE'.
backend (str) – Optimiser backend.
transient (int | None) – Number of leading samples (axis 0 of every prediction leaf) dropped before the objective is evaluated. fit() may override it.

Raises:

ValueError – If backend is unknown, or space/init are inconsistent.
TypeError – If simulate or objective is not callable, or optimizer has the wrong type for the backend.

See also

FitResult: The result of fit().
braintools.optim.multistart_minimize: The 'grad' backend.
sweep: Scan parameters on a grid instead.
hessian_ci: Confidence intervals at the optimum.

Notes

Parameter spaces. The 'grad' and 'scipy' backends optimise in an unconstrained space: a bounded parameter is mapped through brainstate.nn.SigmoidT (SoftplusT for one-sided bounds), a module parameter through its own transform. 'nevergrad' searches the constrained box directly, so every parameter needs finite bounds.

Starts. With n_starts=1 and no key, the single start is the initial value (init, the box midpoint, or the module’s current value). Otherwise key is required and every bounded parameter is drawn uniformly from its box (unbounded ones use their initial value). x0 in fit() gives the starts explicitly. For 'grad' and 'scipy', a start within 1e-3 of the box width of a bound is moved that far inside, so a start on a bound (including a module parameter sitting at its transform’s bound) maps to a finite unconstrained value where the transform is not flat. The 'grad' convergence test then acts on the unconstrained gradient (see multistart_minimize() for the test itself).

Randomness. All random draws derive from key; 'nevergrad' is seeded from it and does not touch NumPy’s global random generator.

Purity. Functional mode is pure. In module mode, the model is not touched while fitting: the parameter values are fed into a functionalised copy of predict (see from_module()), and the best values are written into the parameters only after fit() finishes.

Transformations. fit() drives the search from the host (it compiles its own programs: one jit-ed, vmap-ped multi-start loop for 'grad', a jit-ed loss and gradient for 'scipy', a jit-ed, vmap-ped batch loss for 'nevergrad'), so it cannot itself be called inside jax.jit or jax.vmap.

Examples

>>> import jax
>>> import jax.numpy as jnp
>>> from braintools.fit import Fitter
>>> t = jnp.linspace(0., 1., 50)
>>> simulate = lambda p: p['a'] * jnp.exp(-t / p['tau'])
>>> target = simulate({'a': 2.0, 'tau': 0.3})
>>> mse = lambda pred, y: jnp.mean((pred - y) ** 2)
>>> fitter = Fitter(simulate, mse, {'a': (0., 5.), 'tau': (0.05, 1.)})
>>> res = fitter.fit(target, n_starts=4, key=jax.random.PRNGKey(0))
>>> print(round(float(res.best_params['a']), 3), round(float(res.best_params['tau']), 3))
2.0 0.3

fit(target, *, n_starts=1, key=None, n_steps=100, transient=None, x0=None)[source]#

Run the optimisation.

Parameters:

target (Any) – Data compared with the prediction by the objective.
n_starts (int) – Number of starts S: parallel local searches for 'grad', sequential SciPy runs for 'scipy', the population per generation for 'nevergrad' (unless 'n_sample' is set in the options).
key (Any) – Random key for drawing starts (required when n_starts > 1 and x0 is not given, and always for 'nevergrad').
n_steps (int) – Maximum iterations per start ('grad', 'scipy') or number of generations ('nevergrad').
transient (int | None) – Overrides the constructor’s transient.
x0 (Dict[str, Any] | None) – Explicit starts {name: values} in the constrained space, each with a leading axis S (which then sets n_starts). Not used by 'nevergrad'.

Returns:

Best and per-start parameters and losses.

Return type:

FitResult

Raises:

ValueError – If n_starts or n_steps is not a positive integer, a key is needed but missing, x0 is malformed, or 'nevergrad' is used with a parameter that has no finite box.

classmethod from_module(model, predict, objective, *, optimizer=None, backend='grad', transient=None, search_space=None)[source]#

Fit the trainable brainstate.nn.Param parameters of a module.

Parameters:

model (Any) – The model. Every brainstate.nn.Param with fit=True is fitted, named by its dotted attribute path (e.g. 'w_ee', 'layer.tau').
predict (Callable[[Any], Any]) – predict(model) -> prediction; runs the model and reads its output. It is traced once per fit() call.
objective (Callable[[Any, Any], Any]) – objective(prediction, target) -> scalar.
optimizer (Any) – As for the Fitter constructor.
backend (str) – Optimiser backend.
transient (int | None) – Leading prediction samples dropped before the objective.
search_space (Dict[str, Any] | None) – {name: (low, high)} constrained bounds overriding the box implied by a parameter’s transform (SigmoidT, TanhT, SoftsignT, ScaledSigmoidT, ClipT). Parameters with a box have their starts drawn from it; 'nevergrad' needs a box for every parameter.

Returns:

A fitter in module mode.

Return type:

Fitter

Raises:

ValueError – If the model has no fit=True parameter, or search_space names an unknown parameter.

Notes

predict(model) (plus model.reg_loss() when the model defines it, as brainmass does) is turned into a pure function of the parameter values with brainstate.transform.StatefulFunction: the unconstrained values are substituted for the parameters’ ParamState values in the traced program, so starts and candidates are vmap-ped and the model’s states are never written during the search. The parameter’s own transform maps the unconstrained value to the constrained one. After fit() returns, the model holds the best parameters.

Examples

>>> import brainstate
>>> import jax
>>> import jax.numpy as jnp
>>> from braintools.fit import Fitter
>>> class Decay(brainstate.nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.tau = brainstate.nn.Param(0.5, t=brainstate.nn.SigmoidT(0.05, 2.))
...     def __call__(self, t):
...         return jnp.exp(-t / self.tau.value())
>>> t = jnp.linspace(0., 1., 40)
>>> model = Decay()
>>> mse = lambda p, y: jnp.mean((p - y) ** 2)
>>> fitter = Fitter.from_module(model, lambda m: m(t), mse)
>>> res = fitter.fit(jnp.exp(-t / 0.2))
>>> print(round(float(res.best_params['tau']), 3), round(float(model.tau.value()), 3))
0.2 0.2

property names: Tuple[str, ...]#: Names of the fitted parameters.

Fitter

Contents

Fitter#