Add parameter-aware scaling for forward-sensitivity error control

[BradyPlanden]

Ran into this when testing sensitivities on low-magnitude parameters, such as particle diffusivities (magnitude ~5e-14) in electrochemical particle models.

---

Forward-sensitivity error control is not invariant under an equivalent rescaling of a model parameter. For sufficiently small unitful parameters, BDF can reach the minimum timestep even though the same primal problem expressed with dimensionless parameter solves normally.

sens_atol currently supplies one state-sized absolute-tolerance vector that is reused for every sensitivity vector. Since

S_j = dy / dp_j

has units [y] / [p_j], sensitivities for parameters with different scales need different absolute error weights. A single shared vector cannot represent that for a mixed-scale parameter set.

MWE

Tested the current state of main.

Cargo.toml:

[package]
name = "diffsol-sens-scaling"
version = "0.1.0"
edition = "2021"

[dependencies]
diffsol = "0.14.2"

src/main.rs:

use diffsol::{NalgebraLU, NalgebraMat, OdeBuilder, SensitivitiesOdeSolverMethod};

const PHYSICAL_PARAMETER: f64 = 3.3e-14;
const FREQUENCY: f64 = 1e3;
const STIFFNESS: f64 = 1e3;
const TOLERANCE: f64 = 1e-7;

fn solve(label: &str, parameter: f64, parameter_scale: f64) {
    let problem = OdeBuilder::<NalgebraMat<f64>>::new()
        .p([parameter])
        .rtol(TOLERANCE)
        .atol([TOLERANCE])
        .sens_rtol(TOLERANCE)
        .sens_atol([TOLERANCE])
        .rhs_sens_implicit(
            move |x, p, t, y| {
                y[0] = STIFFNESS * (parameter_scale * p[0] * (FREQUENCY * t).sin() - x[0]);
            },
            move |_x, _p, _t, v, y| {
                y[0] = -STIFFNESS * v[0];
            },
            move |_x, _p, t, v, y| {
                y[0] = STIFFNESS * parameter_scale * (FREQUENCY * t).sin() * v[0];
            },
        )
        .init_sens(
            |_p, _t, y| {
                y[0] = 0.0;
            },
            |_p, _t, _v, y| {
                y[0] = 0.0;
            },
            1,
        )
        .build()
        .unwrap();

    let mut solver = problem.bdf_sens::<NalgebraLU<f64>>().unwrap();
    let result = solver.solve_dense_sensitivities(&[0.0, 1e-3]);
    let stats = solver.get_statistics();
    println!(
        "{label}: result={:?}, steps={}, error_failures={}",
        result.as_ref().map(|_| ()),
        stats.number_of_steps,
        stats.number_of_error_test_failures,
    );
}

fn main() {
    // p = 3.3e-14, with the RHS depending on p / 3.3e-14.
    solve(
        "raw parameter",
        PHYSICAL_PARAMETER,
        1.0 / PHYSICAL_PARAMETER,
    );

    // q = p / 3.3e-14. This produces the same primal equation.
    solve("dimensionless parameter", 1.0, 1.0);
}

Output:

raw parameter: result=Err(OdeSolverError(StepSizeTooSmall {
    time: 6.324555847383075e-8
})), steps=2, error_failures=19
dimensionless parameter: result=Ok(()), steps=48, error_failures=0

The primal equation is identical in both cases. Only the parameterization and therefore the scale of dy/dp differ. For one raw parameter, this can be worked around manually by setting its sensitivity absolute tolerance to approximately

state_atol / abs(parameter)

but this is not possible for multiple differently scaled parameters that affect the same state because sens_atol is shared across all sensitivity vectors.

Possible Solutions

1. parameter scales such as sens_param_scales / pbar, used to apply the state error weights to pbar[j] * S[j]; or
2. one absolute-tolerance vector per sensitivity parameter?

The first option would mirror SUNDIALS IDAS estimated sensitivity tolerances: atol_s[j][i] = atol_y[i] / abs(pbar[j]).

Reference:
https://sundials.readthedocs.io/en/v7.4.0/idas/Mathematics_link.html#selection-of-the-absolute-tolerances-for-sensitivity-variables

[martinjrobins]

the 1st option sounds good and is probably what I should have done originally. There are already param_atol and param_rtol that are used for adjoints, could just use these for forward sensitivities as well since they are of length equal to the number of parameters
UPDATE: sorry I guess you don't want so supply tolerances, you want param scales, so using param_atol and param_rtol is a bit weird.

[martinjrobins]

I guess to match sundials we could change the name of stens_atol to sens_scales, with length num parameters, and then the sensitivity absolute tolerances are set to the state absolute tolerances divided by these scalses: atol / abs(sens_scales). The meaning and usage of sens_rtol would stay the same

[martinjrobins]

wait up, don't we already have the value of the parameter, why can't we just use this for the scale? Could then just keep things as they are, but divide sens_atol by abs(parameter)?

[BradyPlanden]

I think that's reasonable, but if the parameter is zero there would be numeric issues, right? I believe Sundials gates on finite, non-zero parameter values for their pbar vector for this reason -- just thinking a bit more, a parameter identification process that has a parameter converge to zero would result in the tolerances changing throughout the process 🤔

[martinjrobins]

yea, that is a good point. perhaps we will need a separate param_scales as well on the problem, which default to 1.0 but the user can customise....

[martinjrobins]

I've drafted this up in the PR to see how it would look