braintools.fit module

Fit brain models to data.

braintools.fit is model-agnostic: it fits any pure simulate(params) function, or any brainstate module whose parameters are brainstate.nn.Param(..., fit=True), to a target.

• Objectives: Term wraps a loss fn(prediction, target); terms compose with + and scalar * into an Objective. Ready-made losses live in braintools.fit.losses.

• Fitting: Fitter (functional mode, or Fitter.from_module()) with the 'grad' (vmapped multi-start, see braintools.optim.multistart_minimize()), 'scipy' and 'nevergrad' backends; results come back as a FitResult.

• Exploration: sweep() evaluates a function over a parameter grid in memory-bounded vmapped chunks.

• Uncertainty: hessian_ci() gives curvature-based confidence intervals at the optimum (HessianCI).

Examples

>>> import jax
>>> import jax.numpy as jnp
>>> from braintools.fit import Fitter, Term, losses
>>> t = jnp.linspace(0., 1., 50)
>>> simulate = lambda p: p['a'] * jnp.exp(-t / p['tau'])
>>> target = simulate({'a': 2.0, 'tau': 0.3})
>>> objective = Term(losses.mse) + 0.1 * Term(losses.correlation_loss)
>>> fitter = Fitter(simulate, objective, {'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

See also

For a runnable walkthrough, see the tutorial Fitting models to data.

Objectives

A Term wraps one loss fn(prediction, target) with a weight and an optional key that selects one output of a dict-valued prediction. Terms compose with + and scalar * into an Objective, whose Objective.breakdown() reports each term separately.

Term	A weighted loss term of an Objective.
Objective	A weighted sum of loss terms.

Fitting

Fitter fits either a pure simulate(params) function over a search space, or a brainstate module through Fitter.from_module() (its brainstate.nn.Param(..., fit=True) parameters and their bounding transforms). Backends: 'grad' (all starts minimized at once by braintools.optim.multistart_minimize()), 'scipy' and 'nevergrad'. In module mode the best parameters are written back to the model only when Fitter.fit() finishes.

Fitter	Fit model parameters to data by minimising an objective.
FitResult	Outcome of Fitter.fit().

Parameter Sweeps and Uncertainty

sweep	Evaluate a pure function over a grid of parameter values.
hessian_ci	Compute confidence intervals from the Hessian of a loss at its minimum.

HessianCI

Result of hessian_ci().

Losses

Losses for fitting models to data.

Every loss has the signature loss(prediction, target, **options) -> scalar and can be wrapped in a braintools.fit.Term. Each docstring states whether the loss is differentiable.

• Pointwise: mse(), rmse(), mae(), correlation_loss().

• Functional connectivity: fc_correlation_loss(), fc_rmse(), fcd_ks(), fcd_wasserstein().

• Distributions: ks_distance(), wasserstein_1d().

• Summary statistics: psd_log_distance(), rate_loss(), cv_loss().

• Spike trains: van_rossum_loss() (differentiable on dense rasters).

Time steps given as plain numbers follow the convention of the matching analysis module: dt is in seconds for psd_log_distance() (as in braintools.analysis.spectral) but in milliseconds for rate_loss() and van_rossum_loss() (as in braintools.analysis.spike and brainstate). Passing dt as a brainunit Quantity avoids the ambiguity.

Examples

>>> import jax.numpy as jnp
>>> from braintools.fit import losses
>>> float(losses.mse(jnp.array([1.0, 2.0]), jnp.array([1.0, 4.0])))
2.0

mse	Compute the mean squared error between two time series.
rmse	Compute the root-mean-square error between two time series.
mae	Compute the mean absolute error between two time series.
correlation_loss	Compute one minus the Pearson correlation between two time series.
fc_correlation_loss	Compute one minus the correlation between two functional-connectivity matrices.
fc_rmse	Compute the RMSE between the upper triangles of two functional-connectivity matrices.
fcd_ks	Compute the Kolmogorov-Smirnov distance between two FCD value distributions.
fcd_wasserstein	Compute the Wasserstein-1 distance between two FCD value distributions.
ks_distance	Compute the Kolmogorov-Smirnov statistic between two densities on a shared grid.
wasserstein_1d	Compute the Wasserstein-1 distance between two densities on a shared uniform grid.
psd_log_distance	Compute the mean squared difference of the log10 power spectra of two signals.
rate_loss	Compute the squared difference of mean firing rates (in Hz).
cv_loss	Compute the squared difference of the mean inter-spike-interval CV.
van_rossum_loss	Compute a differentiable van Rossum distance between two dense spike rasters.