Adjoint Sensitivity

Adjoint sensitivity analysis efficiently computes the gradient of a scalar objective with respect to many parameters. For an implicit solver, you need to specify:

1. the right-hand side together with its state Jacobian-vector product and the two negative-transpose products used by the backward adjoint solver.
2. the initial condition together with the parameter sensitivity negative-transpose product.

For the logistic equation

$$\frac{dy}{dt} = r y (1 - y/K),$$

the state Jacobian-vector product is

$$Jv = r v (1 - 2y/K),$$

and the adjoint closure must return \(-J^T v\):

$$-J^T v = -r v (1 - 2y/K).$$

The parameter-adjoint closure similarly returns \(-J_p^T v\):

$$-J_p^T v = \begin{bmatrix} -v y (1-y/K) \\ -v r y^2/K^2 \end{bmatrix}.$$

Use rhs_adjoint_implicit for these four right-hand-side operations and init_adjoint for the initial state and its negative parameter transpose. The logistic initial state below is constant, so its parameter-adjoint product is zero.

use crate::{C, M, T, V};
use diffsol::{OdeBuilder, OdeEquationsImplicitAdjoint, OdeSolverProblem};

pub fn problem_adjoint_sens(
) -> OdeSolverProblem<impl OdeEquationsImplicitAdjoint<M = M, V = V, T = T, C = C>> {
    OdeBuilder::<M>::new()
        .p(vec![1.0, 10.0])
        .rhs_adjoint_implicit(
            |x, p, _t, y| y[0] = p[0] * x[0] * (1.0 - x[0] / p[1]),
            |x, p, _t, v, y| y[0] = p[0] * v[0] * (1.0 - 2.0 * x[0] / p[1]),
            |x, p, _t, v, y| y[0] = -p[0] * v[0] * (1.0 - 2.0 * x[0] / p[1]),
            |x, p, _t, v, y| {
                y[0] = -v[0] * x[0] * (1.0 - x[0] / p[1]);
                y[1] = -v[0] * p[0] * x[0] * x[0] / (p[1] * p[1]);
            },
        )
        .init_adjoint(
            |_p, _t, y| y[0] = 0.1,
            |_p, _t, _v, y| {
                y[0] = 0.0;
                y[1] = 0.0;
            },
            1,
        )
        .build()
        .unwrap()
}

The resulting problem implements OdeEquationsImplicitAdjoint and can be used to create an adjoint solver after a checkpointed forward solve.