fitting¶

def fit_parameters(
    make_model: Callable[..., tuple[EquationOfMotion, Pulse]],
    params0: PyTree,
    *,
    loss_fn: Callable[..., ArrayLike],
    optimizer: optax.GradientTransformation,
    n_steps: int = 200,
    conditions: Sequence[PyTree] | None = None,
    save_spec: SaveSpec | None = None,
    t_max: ArrayLike | None = None,
    config: SolverConfig | None = None,
    adjoint: diffrax.AbstractAdjoint | None = None,
    max_backtracks: int = 5,
    step_callback: Callable[[int, PyTree, float], None] | None = None,
    log_every: int = 25,
) -> FitResult:
    r"""Fit model parameters by differentiating through the ODE solve.

    `fit_parameters` minimises the mean loss over the conditions,

    $$
    L(\theta) = \frac{1}{N} \sum_{c=1}^{N}
    \ell\big(\mathrm{simulate}(\mathrm{make\_model}(p, c)),\, c\big),
    \qquad p = \mathrm{unwrap}(\theta),
    $$

    where $\theta$ holds the optimiser's coordinates (see
    [`Parameter`][jbubble.fitting.Parameter]), $\ell$ is `loss_fn`, and
    $\mathrm{simulate}$ is
    [`run_simulation`][jbubble.simulation.run_simulation] with `save_spec`,
    `t_max`, `config`, and `adjoint`. Without `conditions`, $N = 1$ and
    neither `make_model` nor `loss_fn` receives a condition. Each step
    computes the update $\Delta\theta$ with `optimizer` and tries

    $$
    \theta_k = \theta + 2^{-k}\,\Delta\theta, \qquad k = 0, 1, \ldots,
    \texttt{max\_backtracks},
    $$

    accepting the first $\theta_k$ at which every solve converges and both
    $L$ and its gradient are finite.

    `fit_parameters` runs a Python loop around compiled functions, so you
    can't call it under `jax.jit` or `jax.vmap`, and each call compiles
    again. For many fits, such as multi-start fits, write the training loop
    yourself with [`unwrap`][jbubble.fitting.unwrap].

    Parameters
    ----------
    make_model : callable
        `make_model(params) -> (eom, pulse)`, or
        `make_model(params, condition) -> (eom, pulse)` when you pass
        `conditions`. `params` has the structure of `params0`, in physical
        units. Must be JAX-traceable.
    params0 : PyTree
        Initial parameters: a float, a `Parameter`, a dict, list, or tuple
        of them, or an Equinox module such as a
        [`NeuralProperty`][jbubble.bubble.property.NeuralProperty]. What
        gets fitted:

        - A `Parameter`, on its scaled coordinate, unless it's `fixed`.
        - A Python float (or `np.float64`) in `params0` itself or in a dict,
          list, or tuple: `1e-9` is shorthand for `Parameter(1e-9)`.
        - A floating-point JAX or NumPy array, including a NumPy scalar
          other than `np.float64`, in its own units, for example neural
          network weights.

        Everything else is held fixed: integers, booleans, strings,
        callables, and Python floats inside an Equinox module. If you set a
        Python float inside a module, `fit_parameters` warns and names it.
    loss_fn : callable
        `loss_fn(result) -> scalar`, or `loss_fn(result, condition) ->
        scalar` when you pass `conditions`. `result` is a
        [`SimulationResult`][jbubble.simulation.SimulationResult].
    optimizer : optax.GradientTransformation
        For example `optax.adam(0.05)`. For `Parameter` and float values,
        the learning rate is a relative step size: 0.05 changes a value by
        about 5 % per step. `optimizer.update` receives the current
        coordinates as `params`, so transformations such as `optax.adamw`
        work, but act on the coordinates: weight decay pulls a lower-bounded
        `Parameter` toward its initial value, not toward zero.
    n_steps : int
        Number of accepted steps. Default: `200`.
    conditions : list or tuple, optional
        One entry per experimental condition, for example per driving
        pressure or per recording. Each entry can be any pytree, and
        `fit_parameters` passes it to `make_model` and `loss_fn`. All
        conditions share `params`. Conditions with the same structure and
        leaf shapes run in parallel under `jax.vmap`; others run one after
        another. In parallel, each number in a condition, including a
        Python int, reaches `make_model` and `loss_fn` as a traced array:
        use it in `jnp.where` or `jax.lax.cond`, not in an `if` or as a
        slice bound. A condition that holds a Python bool or a string runs
        one after another, so these stay Python values that you can use as
        flags.
    save_spec : SaveSpec, optional
        Output sampling. `None` uses [`SaveSpec()`][jbubble.solver.SaveSpec].
    t_max : float, optional
        Integration end time [s]. `None` uses each pulse's `t_end`.
    config : SolverConfig, optional
        ODE solver settings. `None` uses
        [`SolverConfig()`][jbubble.solver.SolverConfig], the same as
        [`run_simulation`][jbubble.simulation.run_simulation], so the model
        you fit is the model you simulate.
    adjoint : diffrax.AbstractAdjoint, optional
        How the gradient is computed through the solve. `None` uses
        `diffrax.RecursiveCheckpointAdjoint()`. `fit_parameters`
        differentiates in reverse mode, so it can't use
        `diffrax.ForwardMode()`.
    max_backtracks : int
        How many times to halve a step that fails before stopping. If every
        halved step fails, the fit warns and returns the last accepted
        parameters. Default: `5`.
    step_callback : callable, optional
        `step_callback(step, params, loss)`, called outside JIT for
        `params0` (`step == 0`) and after each accepted step. `params` is in
        physical units, and `loss` is the loss at `params`. Raise
        `StopIteration` to end the fit. Any other exception propagates and
        discards the fit, so save checkpoints from the callback if you need
        them.
    log_every : int
        Print progress every `log_every` steps. `0` turns off logging.
        Default: `25`.

    Returns
    -------
    FitResult
        Fitted parameters, loss history, the simulation at the fitted
        parameters, and how the fit ended.

    Raises
    ------
    ValueError
        If `params0` has nothing to fit, a Python float in `params0` is zero,
        or `conditions` is empty.
    TypeError
        If `adjoint` is `diffrax.ForwardMode()`, `conditions` isn't a list
        or tuple, `make_model` or `loss_fn` takes the wrong number of
        arguments, `make_model` doesn't return an `(eom, pulse)` tuple, or
        `loss_fn` doesn't return a scalar.
    RuntimeError
        If a solve fails, or the loss or its gradient isn't finite, at
        `params0`.

    Warns
    -----
    UserWarning
        If you set a Python float inside an Equinox module in `params0`,
        which stays fixed, or if the first step is larger than a
        coordinate's scale or smaller than a millionth of it.
    RuntimeWarning
        If every halved step fails and the fit stops early.

    Examples
    --------
    ```python
    import optax
    from jbubble import fit_parameters
    from jbubble.fitting import Parameter
    from jbubble.metrics import normalised_mse_radius

    fit = fit_parameters(
        lambda p: (make_eom(kappa_s=p["kappa_s"], chi=p["chi"]), pulse),
        {
            "kappa_s": Parameter(1e-9, lower=0.0),  # [N s/m]
            "chi": Parameter(0.3, lower=0.0, upper=1.5),  # [N/m]
        },
        loss_fn=lambda r: normalised_mse_radius(r.radius, measured, R0),
        optimizer=optax.adam(0.05),
    )
    ```
    """
    config = SolverConfig() if config is None else config
    adjoint = diffrax.RecursiveCheckpointAdjoint() if adjoint is None else adjoint
    if isinstance(adjoint, diffrax.ForwardMode):
        raise TypeError(
            "fit_parameters: diffrax.ForwardMode() gives only forward-mode "
            "derivatives, and fit_parameters differentiates in reverse mode. Keep "
            "the default diffrax.RecursiveCheckpointAdjoint(). For a least-squares "
            "fit with forward-mode Jacobians, pass adjoint=diffrax.ForwardMode() "
            "to run_simulation instead, as the fitting guide shows."
        )

    if conditions is not None:
        if not isinstance(conditions, (list, tuple)):
            raise TypeError(
                "fit_parameters: conditions must be a list or tuple with one entry "
                "per condition, for example [{'pressure': 50e3, 'radius': r50}, ...]; "
                f"got {type(conditions).__name__}."
            )
        if not conditions:
            raise ValueError("fit_parameters: conditions is empty.")
        conditions = list(conditions)
    for fn, name, first in (
        (make_model, "make_model", "params"),
        (loss_fn, "loss_fn", "result"),
    ):
        _check_arity(fn, name, (first,) if conditions is None else (first, "condition"))
    stacked = None if conditions is None else _stack(conditions)

    spec = _canonicalise(params0)
    _warn_held_floats(spec)
    trainable, static = eqx.partition(spec, _trainable_filter(spec))
    # Strong dtypes: a weakly typed array, such as jnp.array(1e-9), comes back
    # strongly typed from the first update and would make `evaluate` compile
    # a second time.
    trainable = jtu.tree_map(
        lambda x: jnp.asarray(x, dtype=jnp.result_type(x)), trainable
    )
    if not jtu.tree_leaves(trainable):
        raise ValueError(
            "fit_parameters: params0 has nothing to fit. Use Parameter or Python "
            "floats for physical values, and floating-point arrays for network "
            "weights. Integers, booleans, fixed Parameters, and Python floats "
            "inside an Equinox module are held fixed."
        )

    def one(params: PyTree, condition: Any) -> tuple[jax.Array, SimulationResult]:
        model = (
            make_model(params) if conditions is None else make_model(params, condition)
        )
        if not _is_model(model):
            raise TypeError(
                "fit_parameters: make_model must return an (EquationOfMotion, Pulse) "
                f"tuple; got {type(model).__name__}."
            )
        result = _simulate(
            model[0],
            model[1],
            save_spec=save_spec,
            t_max=t_max,
            config=config,
            adjoint=adjoint,
        )
        loss = jnp.asarray(
            loss_fn(result) if conditions is None else loss_fn(result, condition)
        )
        if loss.ndim != 0:
            raise TypeError(
                f"fit_parameters: loss_fn must return a scalar; got shape {loss.shape}."
            )
        return loss, result

    def objective(tr: PyTree, st: PyTree, conds: Any):
        params = unwrap(eqx.combine(tr, st))
        if conditions is None:
            loss, result = one(params, None)
            return loss, (result, result.converged[None], loss[None])
        if stacked is not None:
            losses, results = eqx.filter_vmap(lambda c: one(params, c))(conds)
            return jnp.mean(losses), (results, results.converged, losses)
        outs = [one(params, c) for c in conds]
        losses = jnp.stack([loss for loss, _ in outs])
        results = [r for _, r in outs]
        converged = jnp.stack([r.converged for r in results])
        return jnp.mean(losses), (results, converged, losses)

    @eqx.filter_jit
    def evaluate(tr: PyTree, st: PyTree, conds: Any):
        (loss, (results, converged, losses)), grads = jax.value_and_grad(
            objective, has_aux=True
        )(tr, st, conds)
        grads_finite = _all_finite(grads)
        ok = jnp.all(converged) & jnp.isfinite(loss) & grads_finite
        return loss, grads, ok, (converged, losses, grads_finite), results

    @eqx.filter_jit
    def propose(tr: PyTree, opt_state: Any, grads: PyTree):
        return optimizer.update(grads, opt_state, tr)

    @eqx.filter_jit
    def apply(tr: PyTree, updates: PyTree, factor: jax.Array):
        return eqx.apply_updates(tr, jtu.tree_map(lambda u: factor * u, updates))

    conds_arg = stacked if stacked is not None else conditions
    loss, grads, ok, diagnostics, results = evaluate(trainable, static, conds_arg)
    if not bool(ok):
        converged, losses, grads_finite = diagnostics
        problems = []
        for i, (ok_i, loss_i) in enumerate(
            zip(converged.tolist(), losses.tolist(), strict=True)
        ):
            where = "" if conditions is None else f"condition {i}: "
            if not ok_i:
                problems.append(f"{where}the ODE solve did not converge")
            elif not jnp.isfinite(loss_i):
                problems.append(f"{where}the loss is {loss_i}")
        if not problems and not bool(grads_finite):
            problems.append("the gradient is not finite")
        if not all(converged.tolist()):
            hint = (
                "Check params0 and the driving pressure, raise "
                "SolverConfig.max_steps, or tighten the solver tolerances."
            )
        else:
            hint = (
                "The solve converged, so check that loss_fn returns a finite "
                "value with a finite gradient at params0."
            )
        raise RuntimeError(
            "fit_parameters: can't start at params0, because "
            + "; ".join(problems)
            + ". "
            + hint
        )

    def physical(tr: PyTree) -> PyTree:
        return unwrap(eqx.combine(tr, static))

    # Typical magnitude of each coordinate, for the first-step check.
    reference = eqx.partition(
        jtu.tree_map(
            lambda x: (
                eqx.tree_at(lambda q: q.raw, x, jnp.ones_like(x.raw))
                if _is_parameter(x)
                else x
            ),
            spec,
            is_leaf=_is_parameter,
        ),
        _trainable_filter(spec),
    )[0]
    opt_state = optimizer.init(trainable)
    history = [float(loss)]
    num_rejected = 0
    stopped_early = False
    message = f"completed {n_steps} steps"

    def report(step: int) -> None:
        if log_every > 0 and (step % log_every == 0 or step == n_steps):
            print(
                f"  step {step:>4} / {n_steps}  loss = {history[-1]:.4e}  "
                f"{_format_params(physical(trainable))}"
                + (f"  (rejected {num_rejected})" if num_rejected else "")
            )

    def callback(step: int) -> bool:
        if step_callback is None:
            return False
        try:
            step_callback(step, physical(trainable), history[-1])
        except StopIteration:
            return True
        return False

    report(0)
    stop = callback(0)
    if stop:
        stopped_early = n_steps > 0
        message = "step_callback raised StopIteration at step 0"
    step = 0
    while not stop and step < n_steps:
        updates, new_opt_state = propose(trainable, opt_state, grads)
        if step == 0:
            _warn_on_first_step(reference, updates)
        for k in range(max_backtracks + 1):
            candidate = apply(trainable, updates, jnp.asarray(0.5**k))
            c_loss, c_grads, c_ok, _, c_results = evaluate(candidate, static, conds_arg)
            if bool(c_ok):
                break
            num_rejected += 1
        else:
            stopped_early = True
            message = (
                f"stopped after {step} of {n_steps} steps: {max_backtracks + 1} "
                "successively halved steps all made a solve fail or gave a non-finite "
                "loss or gradient"
            )
            warnings.warn(
                f"fit_parameters: {message}. Returning the last accepted parameters. "
                "Lower the learning rate, bound the parameters with Parameter, or "
                "raise SolverConfig.max_steps.",
                RuntimeWarning,
                stacklevel=2,
            )
            break
        trainable, opt_state = candidate, new_opt_state
        loss, grads, results = c_loss, c_grads, c_results
        step += 1
        history.append(float(loss))
        report(step)
        stop = callback(step)
        if stop:
            stopped_early = step < n_steps
            message = f"step_callback raised StopIteration after step {step}"

    if stacked is not None:
        num_conditions = len(jtu.tree_leaves(stacked)[0])
        results = [
            jtu.tree_map(lambda x, i=i: x[i], results) for i in range(num_conditions)
        ]

    return FitResult(
        params=physical(trainable),
        loss_history=jnp.asarray(history),
        result=results,
        num_rejected=num_rejected,
        stopped_early=stopped_early,
        message=message,
    )

def unwrap(tree: PyTree) -> PyTree:
    """Replace every [`Parameter`][jbubble.fitting.Parameter] in `tree` by its value.

    [`fit_parameters`][jbubble.fitting.fit_parameters] calls this for you
    before it calls `make_model`. Call it yourself in a hand-written training
    loop.

    Parameters
    ----------
    tree : PyTree
        Any pytree, such as a dict of `Parameter`s or an Equinox module with
        `Parameter` fields.

    Returns
    -------
    PyTree
        `tree`, with each `Parameter` replaced by its physical value, a
        `jax.Array`.
    """
    return jtu.tree_map(
        lambda x: x.value if _is_parameter(x) else x, tree, is_leaf=_is_parameter
    )