Fitting

ds.inference.optsolve fits an InferenceProblem, an experiment or a study paired with an objective, and returns an OptResult. See fitting parameters.

Recommended defaults

Declare a Measurement and an Experiment, pose an InferenceProblem, and fit it with ds.inference.optsolve. Start with ds.inference.MLE() and a Gaussian likelihood, and move to MAP when every fitted handle has a prior. Bound every run with the optimizer's max_steps.

diff_soup.inference.optsolve

optsolve(problem: Any, optimizer: Any = None, *, callbacks: tuple = (), diagnostics: bool = True, run_restricted_initialization: bool = True, restricted_callbacks: tuple = (), resume: Any = None, checkpoint: Any = None) -> OptResult

Fit a problem with an optimizer and return the result.

problem is an InferenceProblem, a study lowered with ds.lower(study), or an OptimizationProblem. When a time-stepping solve at the solution fails, the result has status SOLVER_FAILED and a warning is raised.

Parameters:

Name	Type	Description	Default
problem	Any	The problem to fit.	required
optimizer	Any	The optimizer, ds.inference.NewtonTrustRegion() when None, which steps on the Gauss-Newton matrix of the residual rows.	None
callbacks	tuple	Callback objects called during the fit.	()
diagnostics	bool	If True, attach the fit diagnostics to the result.	True
run_restricted_initialization	bool	If True and the trajectory sets warm_start_steps, first fit the parameters alone for that many steps.	True
restricted_callbacks	tuple	Callbacks for that first fit.	()
resume	Any	An OptResult, a saved result file, or a checkpoint directory to continue from.	None
checkpoint	Any	A directory the fit is saved to as it runs, and resumed from when it already holds a checkpoint.	None

Returns:

Type	Description
OptResult	The OptResult of the fit.

Raises:

Type	Description
ConfigurationError	If the optimizer requires a structure the problem does not provide.
ValueError	If resume or checkpoint holds a fit of a different problem.

diff_soup.inference.InferenceProblem

An experiment or a study paired with the objective a fit minimizes.

Pass it to ds.inference.optsolve or to the statistics methods. ds.lower(study) builds the batched form of the same problem.

Parameters:

Name	Type	Description	Default
experiments	Any	An Experiment, a sequence of them, or a Study, which already declares its parameters and clusters.	required
parameters	Any	The handles the fit moves: one handle, a tuple of them, a ParameterSet or a Hierarchy.	None
objective	MLE | MAP	MLE, or MAP, which adds the declared priors.	MLE()
clusters	Any	The cluster label of each experiment, for a parameter with one value per cluster.	None
layout	Any	The placement of the experiments on a device mesh.	None
trajectory	Any	A SingleShoot or Multishoot applied to every experiment.	None
regularization	Any	A penalty on the parameters added to the objective, such as Tikhonov.	None
trial_predicate	Any	A function of the parameters returning whether a trial point may be evaluated.	None

placement property

placement: dict

A mapping from experiment id to the pair (index of the batched solve, process index) that evaluates it.

inventory

inventory() -> dict

Describe the layout of the parameters and the data as one JSON-serializable dict.

Returns:

Type	Description
dict	A dict with n_parameters, the residual_rows by kind, the parameter_columns by group, the clusters as owners, and one entry per experiment.

values

values(x: Any) -> Any

Return the parameter values a flat vector in this problem's fit order holds.

values = problem.values(info.params)
values[rate]

Parameters:

Name	Type	Description	Default
x	Any	A vector in the fit's order, such as fit.params, StepInfo.params or a checkpoint's vector.	required

Returns:

Type	Description
Any	A read-only mapping from each fitted parameter to its value, as fit[rate] reads it.

diff_soup.inference.MLE dataclass

The maximum-likelihood objective, which minimizes the negative log likelihood of the data.

Pass it as objective= of an InferenceProblem, where it is the default.

diff_soup.inference.MAP dataclass

The maximum-a-posteriori objective, which minimizes the negative log likelihood of the data plus the negative log prior.

Pass it as objective= of an InferenceProblem. Every fitted handle needs a prior, and every measurement a likelihood.

Parameters:

Name	Type	Description	Default
state_prior	Any	A Gaussian with one positive scale, the prior on the window starts of a multiple-shooting fit, or None.	None

---

Objective and trajectory

A trajectory sets how each experiment's solve is split in time (multiple shooting).

diff_soup.inference.LeastSquares

The least-squares objective, half the squared norm of the residual.

diff_soup.inference.SingleShoot

Single shooting, one solve over the experiment's whole time span, the default trajectory.

diff_soup.inference.Multishoot

Multiple shooting, one solve per window of the time span, each from its own start state.

Pass it as trajectory= of an Experiment, or of an InferenceProblem to apply it to every experiment.

Parameters:

Name	Type	Description	Default
windows	Windows	The Windows that split the time span.	required
continuity	Optional[Continuity]	The weight of the continuity term between fitted windows, Continuity() when None.	None
states	Optional[StateBasis]	The StateBasis of the fitted window starts, or None to fit every differential entry.	None
warm_start_steps	int	The number of optimizer steps on the parameters alone, with the window starts held, before the full fit.	0
first_step	Optional[RestartStep]	A FirstStep policy for the first time step of each window's solve, or None to leave it to the solver.	None

diff_soup.inference.Windows

The windows of a Multishoot trajectory, built by one of its constructors.

from_data and from_states start every window from a known state and fit only the parameters. at and uniform fit the window starts beside the parameters.

from_data classmethod

from_data(source: Optional[str] = None) -> 'Windows'

Start every window from a measured full state, one window per consecutive pair of frames.

Parameters:

Name	Type	Description	Default
source	Optional[str]	The name of the measurement that seeds the windows, needed when the experiment has two.	None

from_states classmethod

from_states(times: Any, states: Any) -> 'Windows'

Start every window from a given full state, one window per consecutive pair of frames.

The frame times must equal those of the experiment's one measurement.

Parameters:

Name	Type	Description	Default
times	Any	The strictly increasing frame times, at least two.	required
states	Any	The full state at each frame time, shape (n_frames, n_state).	required

uniform classmethod

uniform(n: int) -> 'Windows'

Split the span from the start time to the last frame into n equal windows and fit their starts.

Parameters:

Name	Type	Description	Default
n	int	The number of windows, at least one.	required

at classmethod

at(ts: Any) -> 'Windows'

Place window boundaries at the times ts and fit the window starts.

Parameters:

Name	Type	Description	Default
ts	Any	At least two boundary times, reaching from the start of the time span to the last frame.	required

diff_soup.inference.Continuity

The weight of the continuity term that joins fitted multiple-shooting windows.

Attributes:

Name	Type	Description
weight	Union[str, float]	"auto" scales the term from the solver tolerances and the observation weights, and a positive float W weights every row by sqrt(W).

diff_soup.inference.StateBasis

The coordinates in which multiple shooting fits each window start.

StateBasis.full() fits every differential entry. StateBasis.from_modes(modes) fits only coefficients c, with each window start its initial value plus modes @ c.

full classmethod

full() -> 'StateBasis'

Fit every differential entry of each window start.

from_modes classmethod

from_modes(modes: Any) -> 'StateBasis'

Fit each window start in the span of the columns of modes.

Parameters:

Name	Type	Description	Default
modes	Any	A column-orthonormal (n_diff, n_modes) array with one row per differential entry of the state.	required

Raises:

Type	Description
ValueError	If modes is not two-dimensional, finite and column-orthonormal.

diff_soup.inference.FixedInitialState

A known initial state, the form a bare array passed as initial_state= takes.

Attributes:

Name	Type	Description
y0	Any	The state at the start of the experiment's time span.
time	float	The start time, which the experiment sets to the start of its span.

diff_soup.inference.LearnedInitialState

An initial state fitted together with the parameters, starting from guess.

Pass it as initial_state= of an Experiment. The fitted state follows the parameters in result.params, and Windows.from_data refuses it.

Attributes:

Name	Type	Description
guess	Any	The starting value of the full state, a 1-D array.
time	float	The start time, which the experiment sets to the start of its span.

---

Bounds and regularization

diff_soup.inference.Bounds dataclass

Box bounds on the fitted parameters, one (lower, upper) pair per parameter group.

A study builds one from the bounds declared with spec.param(..., bounds=...). A None side is unbounded, and multiple-shooting window starts are never bounded.

Parameters:

Name	Type	Description	Default
spec	Optional[dict[str, tuple[Optional[float], Optional[float]]]]	A mapping from group name, the handle's name or name@member for a per-cluster parameter, to its (lower, upper) pair.	None
default	tuple[Optional[float], Optional[float]]	The (lower, upper) pair of every group spec does not name, (None, None) by default.	_UNBOUNDED

diff_soup.inference.Tikhonov dataclass

Adds the penalty 0.5 * lambda_ * sum(weights * x**2) on the fitted values to the objective.

Pass it as regularization= of an InferenceProblem.

Parameters:

Name	Type	Description	Default
lambda_	float	The strength of the penalty.	0.01
weights	Optional[ndarray]	One weight per penalized value, or None to weight the k-th value by 1 + k(k+1), one for a single parameter.	None
on	Any	A handle whose values alone are penalized, or None to penalize every fitted parameter.	None

diff_soup.inference.NoReg

No penalty, used by the solvers when a problem has no regularization.

---

Results

diff_soup.inference.OptResult

The result of ds.inference.optsolve.

result[handle] returns the fitted value of a parameter handle. A field the solver does not fill keeps its default.

Attributes:

Name	Type	Description
params	ndarray	The fitted values as one flat array.
cost	float	The final value of the objective.
status	TerminationStatus	The TerminationStatus of the run.
message	str	How the run ended, as text.
n_steps	int	The number of steps the optimizer took.
diagnostics	Any	The OptDiagnostics of the fit, or None with diagnostics=False.
log_likelihood	Optional[float]	The log-likelihood at the solution, or None when no measurement declares a likelihood.
log_prior	Optional[float]	The log-prior at the solution under a MAP objective, or None.
opt_params	Any	The fitted parameter values in the form ds.pdesolve takes through theta=.
initial_params	ndarray	The starting values.
grad_norm	float	The gradient norm at the solution.
projected_grad_norm	Optional[float]	The gradient norm projected onto the bounds, zero on an active bound.
cost_history	ndarray	The objective after every step.
wall_time	float	The run time in seconds.
experiments	Mapping	The fitted values of each experiment, by its id.
window_states	Any	The fitted window starts of a multiple-shooting fit, or None.
info	Mapping	Further records of the solver's settings and work.

params class-attribute instance-attribute

params: ndarray = field(default_factory=lambda: np.array([]))

cost class-attribute instance-attribute

cost: float = float('inf')

status class-attribute instance-attribute

status: TerminationStatus = eqx.field(static=True, default=TerminationStatus.UNKNOWN)

message class-attribute instance-attribute

message: str = eqx.field(static=True, default='')

converged property

converged: bool

Whether status counts as convergence.

n_steps class-attribute instance-attribute

n_steps: int = eqx.field(static=True, default=0)

diagnostics class-attribute instance-attribute

diagnostics: Any = eqx.field(static=True, default=None)

log_likelihood class-attribute instance-attribute

log_likelihood: Optional[float] = eqx.field(static=True, default=None)

log_prior class-attribute instance-attribute

log_prior: Optional[float] = eqx.field(static=True, default=None)

save

save(path: Any) -> None

Save the result to an .npz file that OptResult.load reads.

opt_params, diagnostics, experiments and wall_time are not saved.

Parameters:

Name	Type	Description	Default
path	Any	The file path, with .npz appended when missing.	required

load classmethod

load(path: Any, *, problem: Any = None) -> 'OptResult'

Load a result saved with OptResult.save.

Parameters:

Name	Type	Description	Default
path	Any	The file path, with .npz appended when missing.	required
problem	Any	A problem of the same parameter layout, used to restore result[handle]. A fit under another solver tolerance or data set loads.	None

Returns:

Type	Description
'OptResult'	The loaded OptResult.

Raises:

Type	Description
ValueError	If problem is given and the file's columns hold other handles, or its objective is another.

diff_soup.inference.TerminationStatus

How an optimizer run ended, at result.status.

status.converged is True for SUCCESS, GRAD_TOL, STEP_TOL, COST_TOL and DERIVATIVE_ACCURACY.

Attributes:

Name	Type	Description
SUCCESS		Convergence without a named test, which no built-in solver reports.
GRAD_TOL		The gradient norm fell below its tolerance.
STEP_TOL		The step fell below its tolerance.
COST_TOL		The change in the objective fell below its tolerance.
ACCEPTED_STEPS		TrustRegion reached its target_accepted_steps.
MAX_STEPS		The step or evaluation budget ran out.
SOLVER_FAILED		A time-stepping solve at the solution did not reach the end of its span.
DIVERGED		BFGS met an objective or gradient that is not finite.
STOPPED		A callback raised StopFit.
UNKNOWN		The status before any run.
STALLED		The fit stopped without converging: TrustRegion after failed trials in a row or while its model still predicted a decrease, and NewtonTrustRegion or LevenbergMarquardt when a step or cost test passed on a step the trust radius or damping held short.
DERIVATIVE_ACCURACY		TrustRegion converged to the accuracy of its derivatives, which a tighter solver rtol improves.

diff_soup.inference.OptDiagnostics dataclass

The quality report of a fit, at result.diagnostics, from the residual Jacobian at the solution.

A condition number above 1e8, or a parameter that does not move the residual, also issues an IdentifiabilityWarning.

Attributes:

Name	Type	Description
theta_singular_values	Optional[ndarray]	The singular values of the Jacobian with respect to the model parameters, learnable-function weights left out.
jacobian_cond	float	The ratio of the largest to the smallest of them, inf when the smallest is zero.
window_obs_rank	Optional[tuple]	Per fitted window start, the rank of the window's observations with respect to it.
continuity_defect_wrms	Optional[ndarray]	Per window boundary, the root mean square of the continuity mismatch, in state units.
continuity_unrepresented	Optional[ndarray]	Per window boundary, the weighted continuity mismatch a reduced StateBasis cannot represent.
per_series_cost	tuple	The cost of each measurement stream, or of each window under Windows.from_data.
series_names	Optional[tuple]	The stream names, parallel to per_series_cost, or None.
param_group_names	tuple	The names of the parameter groups.
param_group_cond	tuple	The condition number of each parameter group, parallel to param_group_names, nan for learnable-function weights.
ad_mode	Optional[str]	The derivative mode of the gradient.
hessian	Optional[str]	The curvature the optimizer used, "true" or "gauss-newton".
hessian_composition	Optional[str]	The autodiff composition of the exact Hessian.
subproblem	Optional[str]	How the optimizer solved its step subproblem.
linear_solver	Optional[str]	The factorization the optimizer used.
warnings	tuple	The identifiability and window notes, as text.

---

Callbacks

diff_soup.inference.Callback

Base class for a callback that the optimizer calls during a fit.

A subclass implements on_step, and the other methods do nothing by default. Raise StopFit from on_step to end the fit early.

bind

bind(problem: Any) -> None

Receive the problem before the fit starts.

Parameters:

Name	Type	Description	Default
problem	Any	The problem passed to ds.inference.optsolve.	required

on_start

on_start(n_params: int, n_residuals: int, max_steps: int) -> None

Called once before the first step.

Parameters:

Name	Type	Description	Default
n_params	int	The number of fitted values.	required
n_residuals	int	The number of residual rows, or 0 when the solver has none.	required
max_steps	int	The optimizer's step budget.	required

on_step abstractmethod

on_step(info: StepInfo) -> None

Called after every step.

Parameters:

Name	Type	Description	Default
info	StepInfo	The step's record, with step, cost, grad_norm, params and accepted among its fields.	required

on_end

on_end(result: OptResult) -> None

Called once with the result after the last step.

Parameters:

Name	Type	Description	Default
result	OptResult	The OptResult of the fit.	required

diff_soup.inference.LogCallback

Prints one line per step with the cost, the gradient norm and the step norm.

Parameters:

Name	Type	Description	Default
log_file	Optional[str]	A path that also receives the lines, overwritten when it exists.	None

diff_soup.inference.CheckpointCallback

Saves a fit to a directory as it runs, so that ds.inference.optsolve can resume it.

Each step appends a line to events.jsonl, and every every-th step is saved as an .npz result. Pass the directory as resume= or checkpoint= to continue the fit.

Parameters:

Name	Type	Description	Default
directory	Any	The directory to write to, created when missing.	required
every	int	Save every every-th accepted step, or every every-th attempt under a resumable solver such as TrustRegion.	1
keep	str	"last" keeps the newest saved attempt, "best" also the one of lowest cost, and "all" every one.	'last'

diff_soup.inference.StopFit

Raised from Callback.on_step to end a fit after the current step, with status STOPPED.

Every built-in solver catches it except Optax.