Calculating First-order Sensitivity using diffrax

[sn248]

Hi Patrick

I am trying to calculate the first order sensitivity w.r.t. to all the three kinetic parameters for the Robertson Equations using diffrax. I am using the stiff ODE example provided here.

As far as I understand, the only change I have to make is that I have to provide an option to calculate adjoints, so my diffeqsolve looks like the following

sol = diffrax.diffeqsolve(
        terms,
        solver,
        t0,
        t1,
        dt0,
        y0,
        saveat=saveat,
        stepsize_controller=stepsize_controller,
       adoints = diffrax.RecursiveCheckpointAdjoint(checkpoints=None)
    )

The code runs without any errors but I am not able to extract the sensitivities, and sol.ys provides values for the species, but not their sensitivities. I am trying to compare the results with the sensitivities provided by the SUNDIALS CVODES solver, see their results here.

Any help would be highly appreciated!

[patrick-kidger]

Diffrax provides gradients via JAX's standard API for doing so -- jax.grad, jax.value_and_grad and friends. Wrap your diffeqsolve in one of those.

For an introductory example take a look at https://docs.kidger.site/diffrax/examples/forcing/, which uses jax.grad to compute the derivative of a differential equation with respect to a forcing term.