Skip to content

Fitting parameters

This page provides an overview of fitting the learnable parameters of a model to observed data with ds.inference.optsolve. Throughout, we fit synthetic data with additive Gaussian noise generated by solving the heat equation

\[ \frac{\partial u}{\partial t} = D\,\frac{\partial^2 u}{\partial x^2}, \qquad x \in (0, 1), \]

with \(u(0, t) = u(1, t) = 0\), \(u(x, 0) = x(1 - x)\), and \(D = 0.3\), sampled at five times with noise of scale \(10^{-3}\):

import jax

jax.config.update("jax_enable_x64", True)

import jax.numpy as jnp
import numpy as np

import diff_soup as ds
from diff_soup.domains import cartesian_grid_1d

n = 32
spec = ds.problem("heat")
spec.domain("rod", cartesian_grid_1d(n, 0.0, 1.0, method="fvm"))
spec.field("u", mass="differential", scale=1.0, init=lambda x: x * (1.0 - x))
spec.param("D", value=0.3)                                            # the generating diffusivity


@spec.equation(writes="u", reads=("u",))
def heat(q):
    return q.p.D * q.ops.laplacian(q.u)


spec.bc.dirichlet("u", node_indices=np.array([0, n - 1]), value=0.0)
problem_true = spec.build()
y0 = spec.assemble_state({})
solver = ds.JIDA()

ts = jnp.linspace(0.02, 0.1, 5)
sol_true = ds.pdesolve(problem_true, solver, (0.0, 0.1), y0, saveat=ds.SaveAt(ts=ts))
data = sol_true.ys + 1e-3 * np.random.default_rng(0).standard_normal(sol_true.ys.shape)

The least-squares problem

Here we parameterize the solution of the heat equation, \(\hat u(\theta)\), with a learnable diffusivity denoted by the parameter \(\theta\),

\[ \frac{\partial \hat u}{\partial t} = \theta\,\frac{\partial^2 \hat u}{\partial x^2}, \qquad x \in (0, 1), \]

with the same boundary and initial conditions. Given the observed data \(u_{\text{data}}\) at the observation times, we fit \(\theta\) by minimizing the least-squares loss

\[ \mathcal{L}(\theta) = \tfrac{1}{2}\,\big\|\hat u(\theta) - u_{\text{data}}\big\|_2^2 . \]

The full fit:

spec = ds.problem("heat")
spec.domain("rod", cartesian_grid_1d(n, 0.0, 1.0, method="fvm"))
spec.field("u", mass="differential", scale=1.0, init=lambda x: x * (1.0 - x))
D = spec.param("D", value=0.1, learnable=True, bounds=(0.01, 2.0))    # the learnable diffusivity


@spec.equation(writes="u", reads=("u",))
def heat(q):
    return q.p.D * q.ops.laplacian(q.u)


spec.bc.dirichlet("u", node_indices=np.array([0, n - 1]), value=0.0)
problem = spec.build()
y0 = spec.assemble_state({})

measurement = ds.inference.Measurement("u", ts, data, weight=1e6)   # declare the data
experiment = ds.inference.Experiment(problem, solver, (0.0, 0.1), y0, (measurement,))                # bind the data to one solve
infproblem = ds.inference.InferenceProblem(experiment, parameters=D)                                  # pose the least-squares problem
result = ds.inference.optsolve(infproblem, ds.inference.LevenbergMarquardt())                        # solve it
result[D]                                                                                            # 0.2993, the fitted diffusivity

Solving the least-squares problem

To solve this problem, we first register the observed data using ds.inference.Measurement, bind it to one solve of the model using ds.inference.Experiment, define the optimization problem using the ds.inference.InferenceProblem API, and solve the least-squares problem using ds.inference.optsolve with an optimization algorithm, here ds.inference.LevenbergMarquardt().

Declare the learnable parameter

spec.param declares the diffusivity as a decision variable with learnable=True and nominal starting value 0.1. The optional bounds= argument determines whether the optimization problem is solved as a constrained or unconstrained optimization problem. Without bounds, the problem is solved as an unconstrained problem. With bounds=(lo, hi), or the one-sided bounds=(lo, None) and bounds=(None, hi), the problem is solved as a box-constrained optimization problem as discussed further in box bounds and optimization solvers.

D = spec.param("D", value=0.1, learnable=True, bounds=(0.01, 2.0))

Declare the data

ds.inference.Measurement(name, times, values, *, weight=None, likelihood=None, observe=None, mask=None, shape=None, dtype=None) declares one stream of observed values together with its weight or its noise model. An experiment takes a tuple of measurements.

measurement = ds.inference.Measurement("u", ts, data, weight=1e6)
  • name labels the measurement in the diagnostics. Names are unique within an experiment.
  • times holds the observation times, and values the observed values, one row per time.
  • weight= scales the residual of the measurement. The residual rows are \(\sqrt{w}\,(u_{\text{data}} - \hat u)\) with weight=1 by default, and weight=1e6 above is \(1/\sigma^2\) for a noise scale of \(10^{-3}\). It takes a scalar or an array broadcastable against values.
  • likelihood= states a noise model instead, for statistical estimation. A measurement takes weight= or likelihood=, not both. Without a likelihood the fit is weighted least squares and the result reports log_likelihood=None; information, laplace, profile, identifiability, residual_report, practical_report, log_posterior, posterior_target, and the MAP objective refuse a measurement without one, naming it. bootstrap and cluster_covariance do not need one.
  • observe= maps the state at a time to the observed quantity. The default observes the state itself. The accepted forms are covered under observables below.
  • mask= selects the times that carry a measurement, one True or False per time. The default keeps every time.
  • values may also be a function of no arguments that returns the array, loaded once when the fit first needs it, with shape= and dtype= stating the shape and type it returns (multi-GPU parallelism).

Observations that reach the state through boundary traces or averaged images are covered on data assimilation.

Declare the experiment

ds.inference.Experiment(model, solver, t_span, initial_state, measurements, inputs=None, id=None, trajectory=None) binds the measurements to one solve of the model.

experiment = ds.inference.Experiment(problem, solver, (0.0, 0.1), y0, (measurement,))
  • model is the built problem. A Recipe template is accepted in its place (hierarchical fits).
  • solver is the time-stepping solver used for every solve during the fit (ODE/DAE solvers).
  • t_span=(t0, t1) is the time interval of the solve. Every observation time lies inside it.
  • initial_state is the state the solve starts from at t0.
  • measurements is the tuple of measurements the solve is compared against.
  • inputs= fixes named parameter values for this experiment alone (hierarchical fits).
  • id= names the experiment in the result. The default numbers the experiments in order, experiment-0 onwards.
  • trajectory= decides how the solve is split in time. The default ds.inference.SingleShoot() runs one solve over the interval. ds.inference.Multishoot(windows) splits it into windows (multiple shooting).
  • schedule= is an operating protocol the solve follows, a ds.Schedule (operating protocols).

Pose the inference problem

ds.inference.InferenceProblem(experiments, *, parameters, objective=ds.inference.MLE(), clusters=None, layout=None, trajectory=None, regularization=None) poses the optimization problem. It takes one experiment or a tuple of them, and the parameters the fit moves.

infproblem = ds.inference.InferenceProblem(experiment, parameters=D)
  • experiments is the experiment to fit, or a tuple of experiments fitted together (hierarchical fits).
  • parameters= names the parameters and functions the fit moves, by the handles that spec.param and spec.function return. It takes one handle, a tuple of handles, or ds.ParameterSet(shared=, varying=) for a parameter tied across experiments or taking one value per cluster (hierarchical fits). The handles are listed in the order their fitted values take in result.params. An inference problem without parameters= is refused.
  • objective= is the objective the fit minimizes. The default ds.inference.MLE() is the maximum-likelihood objective, which on measurements without a likelihood is the weighted least-squares loss above. ds.inference.MAP() adds the priors declared on the parameters (statistical estimation).
  • clusters= labels each experiment by the cluster it belongs to, for a parameter that takes one value per cluster (hierarchical fits).
  • layout= places the experiments on a device mesh (multi-GPU parallelism).
  • trajectory= applies one trajectory setting to every experiment (multiple shooting).
  • regularization= adds a penalty on the parameters to the loss. The default adds none (regularization).

Run the fit

ds.inference.optsolve(problem, optimizer=None, *, callbacks=(), diagnostics=True) solves the inference problem and returns the solution as a ds.inference.OptResult. The optimization algorithms and their options are covered in optimization solvers. Solving many problems at once across several devices is covered in multi-GPU parallelism, and the further options of multiple-shooting fits in multiple shooting.

result = ds.inference.optsolve(infproblem, ds.inference.LevenbergMarquardt())
  • optimizer= is the optimization algorithm. The default is ds.inference.NewtonTrustRegion(), which steps on the Gauss-Newton matrix of the residual rows; NewtonTrustRegion(hessian="true") steps on the exact Hessian. ds.inference.LevenbergMarquardt() is the damped least-squares choice (optimization solvers).
  • callbacks= lets you monitor the progress of the optimization run. ds.inference.LogCallback() prints one line per step, and ds.inference.LogCallback(log_file=...) also writes those lines to a file. ds.inference.CheckpointCallback(directory, every=k) saves the run so far to directory after every k-th accepted step, and every= defaults to 1. A custom callback subclasses ds.inference.Callback.
  • diagnostics=True provides a diagnostic report on the quality of the optimization solution (fit diagnostics). False skips the report.

Analyze the result

The result object stores all the information about the optimization run. It carries the fitted values, the loss and its history, the flags that say how the run ended, and the diagnostics. result[D] returns the fitted value of the parameter D, result.params returns every fitted value as one flat array, and result.opt_params holds the fitted parameter values in the form ds.pdesolve takes through theta=.

result[D]                                                                        # 0.2993, the fitted value by handle
result.params      # array([0.2993]), the flat vector of fitted values
result.status      # TerminationStatus.COST_TOL
result.cost
sol_fit = ds.pdesolve(problem, solver, (0.0, 0.1), y0, theta=result.opt_params,
                      saveat=ds.SaveAt(ts=ts))

The result also carries:

  • initial_params, the starting values.
  • cost, the final value of the loss, and cost_history, its value after every step.
  • status, the flag that says how the run ended, converged, whether that flag counts as success, and message, the same flag as text.
  • n_steps and wall_time, how many steps the optimizer took and how long it ran.
  • diagnostics, the fit diagnostics.
  • experiments, the fitted values of each experiment by its id, result.experiments["experiment-0"][D] here.

Multiple-shooting, hierarchical, and statistical fits fill further fields, covered on their pages.

The status flags:

  • COST_TOL, STEP_TOL, and GRAD_TOL are convergence on the change in the loss, the size of the step, or the gradient.
  • MAX_STEPS is the step budget running out.
  • STALLED is a fit that stopped while its model still predicted a decrease, or after too many invalid trial steps in a row. A stalled fit has not converged. NewtonTrustRegion and LevenbergMarquardt stall when their trust radius or damping shrinks the step below its tolerance, and record why in result.info["stall_reason"].
  • DERIVATIVE_ACCURACY is ds.inference.TrustRegion converging to the accuracy of its gradient, which the time-stepping solves compute to about their relative tolerance. A tighter rtol lets the fit go further.
  • ACCEPTED_STEPS is ds.inference.TrustRegion reaching its target_accepted_steps= (hierarchical fits).
  • DIVERGED is ds.inference.BFGS meeting a loss or gradient that is not finite.
  • SOLVER_FAILED is the time-stepping solve at the solution failing before the end of the time interval.
  • UNKNOWN is the flag before any run.
  • converged is True for SUCCESS, GRAD_TOL, STEP_TOL, COST_TOL, and DERIVATIVE_ACCURACY. It means a local stopping test passed, not that the fit found the best parameters. With a likelihood declared, check every fit with ds.inference.residual_report(infproblem, result) (goodness of fit).

For example, a run limited to fifty steps, with the progress printed:

result = ds.inference.optsolve(infproblem, ds.inference.LevenbergMarquardt(max_steps=50),
                          callbacks=(ds.inference.LogCallback(),))
result.converged     # True
result.message       # "cost_tolerance"
result.n_steps       # 7

The result is also the starting point for uncertainty quantification methods discussed in statistical estimation and posterior sampling.

Observables

An observable is the quantity the data measures, written as a function of the state. In our example the observable is the state variable \(u\) itself, the temperature along the rod. In practice the data often measure a derived quantity instead. For the heat equation that could be the energy density throughout the rod, the heat flux through the boundaries, or the spatial mean of the temperature. The accepted forms of observe=:

  • The default observes the state itself.
  • A function of the state, such as lambda state: state.mean(). The function may also take the keyword arguments t, the observation time, and theta, the current parameter values.
  • The name of a declared observation, such as observe="mean_u" for a spec.observe declaration named mean_u, read in its dimensional units.
  • The predict method of a subclass of ds.inference.data.ObservationOperator, which defines predict(self, u, t, theta), as in the second example below. u is the state, t the observation time, and theta the current parameter values.

To observe one field at a few probe cells, add the field's offset in the state to the cell indices:

probes = problem.field_slice("u").start + jnp.array([5, 15, 25])
probe_measurement = ds.inference.Measurement("probes", ts, data[:, [5, 15, 25]],
                                             observe=lambda state: state[probes], weight=1e6)

For example, to add a second measurement that observes the spatial mean of the temperature at the later times only:

mean_measurement = ds.inference.Measurement("mean", ts, data.mean(axis=1), observe=lambda state: state.mean(),
                                            weight=1e6, mask=ts > 0.05)
experiment_two = ds.inference.Experiment(problem, solver, (0.0, 0.1), y0, (measurement, mean_measurement))
result = ds.inference.optsolve(ds.inference.InferenceProblem(experiment_two, parameters=D))
result.diagnostics.series_names      # ("u", "mean")
result.diagnostics.per_series_cost

The same spatial mean written as a class:

class SpatialMean(ds.inference.data.ObservationOperator):
    def predict(self, u, t, theta):
        return jnp.mean(u, keepdims=True)


mean_measurement = ds.inference.Measurement("mean", ts, data.mean(axis=1, keepdims=True),
                                            observe=SpatialMean().predict, weight=1e6)

Regularization

Regularization adds a penalty on the parameter values to the loss, which keeps the fit from drifting to large or poorly determined values. ds.inference.Tikhonov(lambda_=0.01, weights=None) adds the penalty

\[ \tfrac{1}{2}\,\lambda \sum_k w_k\,\theta_k^2 \]

to the loss.

  • lambda_= sets the strength \(\lambda\).
  • weights= sets one weight \(w_k\) per parameter. The default weights a single parameter by one, and the coefficients of a Legendre expansion by their degree (function approximation).

For example, to add a small ridge on \(D\):

infproblem = ds.inference.InferenceProblem(experiment, parameters=D,
                                           regularization=ds.inference.Tikhonov(lambda_=1e-2))
result = ds.inference.optsolve(infproblem)
result[D]        # 0.2993, the ridge is negligible against the data

Box bounds

Bounds declared with spec.param(..., bounds=(lo, hi)) apply whenever that parameter is fitted. The optimizer keeps every trial inside the box, and convergence is tested on the projected gradient, which is zero on an active bound. A parameter fitted across several experiments declares the same bounds in each of them (hierarchical fits).

For example, to bound \(D\) below the generating value:

spec_bounded = ds.problem("heat_bounded")
spec_bounded.domain("rod", cartesian_grid_1d(n, 0.0, 1.0, method="fvm"))
spec_bounded.field("u", mass="differential", scale=1.0, init=lambda x: x * (1.0 - x))
D_bounded = spec_bounded.param("D", value=0.1, learnable=True, bounds=(0.05, 0.25))


@spec_bounded.equation(writes="u", reads=("u",))
def heat_bounded(q):
    return q.p.D * q.ops.laplacian(q.u)


spec_bounded.bc.dirichlet("u", node_indices=np.array([0, n - 1]), value=0.0)
problem_bounded = spec_bounded.build()
result = ds.inference.optsolve(ds.inference.InferenceProblem(
    ds.inference.Experiment(problem_bounded, solver, (0.0, 0.1), y0, (measurement,)), parameters=D_bounded))
result[D_bounded]             # 0.25, the upper bound is active
result.status                 # TerminationStatus.GRAD_TOL
result.projected_grad_norm    # 0.0, stationary on the bound

result.projected_grad_norm is the convergence measure with bounds active. It is zero on an active bound, where result.grad_norm is not.

Fit diagnostics

result.diagnostics is a report on the quality of the solution, computed from the sensitivity of the residual to each parameter at the solution. It says whether the data determine every parameter.

diag = result.diagnostics
diag.theta_singular_values   # singular values of the parameter block of the Jacobian
diag.jacobian_cond           # their ratio, 1.0 for one parameter
diag.per_series_cost         # cost per stream, parallel to diag.series_names
diag.param_group_cond        # condition number per group, parallel to diag.param_group_names
diag.warnings                # identifiability notes, empty here
  • theta_singular_values and jacobian_cond measure how well the data determine the parameters. A condition number near one is ideal, and inf marks a parameter the data do not depend on.
  • per_series_cost splits the loss by data stream.
  • param_group_cond gives the condition number per parameter group.
  • warnings lists the identifiability notes as text.

Multiple-shooting fits fill further fields (multiple shooting).

A condition number above \(10^8\), or a parameter the data do not depend on, raises ds.inference.fitting.IdentifiabilityWarning naming the parameters. For example, two parameters that enter only as a product cannot be told apart:

spec_two = ds.problem("heat_product")
spec_two.domain("rod", cartesian_grid_1d(n, 0.0, 1.0, method="fvm"))
spec_two.field("u", mass="differential", scale=1.0, init=lambda x: x * (1.0 - x))
D = spec_two.param("D", value=0.1, learnable=True)
E = spec_two.param("E", value=1.0, learnable=True)


@spec_two.equation(writes="u", reads=("u",))
def heat_product(q):
    return q.p.D * q.p.E * q.ops.laplacian(q.u)


spec_two.bc.dirichlet("u", node_indices=np.array([0, n - 1]), value=0.0)
problem_two = spec_two.build()
result = ds.inference.optsolve(ds.inference.InferenceProblem(
    ds.inference.Experiment(problem_two, solver, (0.0, 0.1), y0, (measurement,)), parameters=(D, E)))
# IdentifiabilityWarning: optsolve: the learnable-parameter Jacobian is rank deficient: theta-block condition number 5.40e+14 exceeds 1e+08. ...
result.diagnostics.jacobian_cond           # 5.4e14
result.diagnostics.theta_singular_values   # one singular value near zero

The weights of a learnable function (spec.function(..., learnable=True)) are left out of the test, since a network's weights are not determined one at a time. A fit of weights alone records a note in warnings instead, and theta_singular_values and jacobian_cond cover the other parameters only.

The Fisher information, profile likelihoods, and design sensitivities behind these numbers are covered on identifiability.