Coupled ode system in cmdstan 2.18

In cmdstan 2.17 there was ode_system which

can be used to provide analytic Jacobians via partial template specialisation

Is coupled_ode_system in cmdstan 2.18 can be used to provide analytic Jacobians?

Yes, that’s where all the work is done.

So 2.18 supports analytic Jacobians? I didn’t know that. That’s awesome!

So if I understand correctly, we still have to forward solve a system of size N \cdot (K+1), but at least the Jacobian terms that go in to that system don’t need to be computed with autodiff. How much of a speedup does it usually provide?

No, analytic Jacobians are not supported by Stan in any “official” way. The coupled_ode_system though is what you need to tinker with to get them.

I have created a while ago a branch which even generates the Jacobian in symbolic form for you and injects it into the stan system (look for the rstanode branch in the rstan repo). Those codes are >1y out of date, but I intend to resurrect them. This won’t become an official feature any time soon (if at all).

The speedups you can expect, like always, depend a lot on details. I tested it on a PK model which has 3 states and 4 parameters. I recall getting a 2-3x speed increase with analytic Jacobians. So for the usual long inference times of ODE models this is really nice. However, these speedups are easy to get by now using brute force via map_rect and many cores (hierarchical ODE problems do scale essentially linearly with the # of cores).

I am working on those. Will let how to use in cmdstan 2.18.

Well, results with brute force are not very encouraging. cmdstan 2.17 with analytic Jacobians is twice faster+ than cmdstan 2.18 MPI with 20 cores with stan’s AD. As I understand jacobians are calculated with stan’s AD unless specifically supplied. I am almost there…

Steps to supply analytical jacobian are quite simple and I have automated them using R:

  • create system of ODEs using dMod library using addReaction

  • generate C code for the system using funC and save in ode.c

  • generate Jy and Jtheta and save in states.c and params.c correspondingly

  • remove all crap and save text for Jy in states.txt and text for Jtheta in params.txt

  • copy the attached file (after modifying namespace and typedef - only two lines) into user_header.hpp

  • compile model

  • make

This file is for cmdstan 2.17. For cmdstan 2.18 I have to modify it to correspond to coupled_ode_system. I am not very good with it but will try unless wds15 decides to help me out. I have R code to share.

pbpkauto_jacobianCmdStan.txt (2.82 KB)

1 Like

Another speedup gain when using stan’s AD is with ncp formulation - at least for my problem…

Hahaha I’m impressed you’ve gotten things this far. That takes some effort, good stuff.

I am impressed as well… It seems there is nothing sophisticated -
just modify few methods in coupled_ode_system to incorporate
analytical jacobian instead stan’s AD and save as user_header.hpp. I
don’t know yet which coupled_ode_system because there are two functors
in prim and rev subfolders. Since analytical jacobian is already there
it should not be so hard for somebody who knows C++ well. Would like
to give a hand?

I am on my way - last week I tried a different route which turned out to not work.

I am just a bit short on time this week, but let’s see - I don’t mind if others are faster.

I would be very interested in the speedups you are getting with the analytic Jacobians for your problems. These must be large!

Frankly I think I hacked almost everything from the code you have provided so all IP belongs to you. My C++ is bit rusty so it make take some time. Yes, I am trying to get this to make PBPK models a reality. Otherwise waiting one day for 8 subjects with analytical Jacobian is bit too much.

Possible suggestion for the next release - add functors for Jy, Jtheta, and Jinit to coupled_ode_systems. In this case no hacking is needed. Maybe I am missing something…

The initial attempt is attached. While compiles and links it breaks. Any ideas why? I copied <stan/math/rev/arr/functor/coupled_ode_system.hpp> and modified to incorporate Jy and Jtheta.
jac_cp.hpp (6.5 KB)

Hard to say – what’s the link error?

There is no link error. It works on Win 10 on a small example. Stan code to simulate the data, fit the model, r code to run, and analytical jacobian are attached. I am not sure why similar logics doesn’t run (MCMC gives straight lines for each chain) for real model on Win 10 - compiled without threading or MPI. Now I am not sure if it breaks on Linux either (compiled with MPI). I guess I need someone to doublecheck the code (jac.hpp).

sho.stan (623 Bytes)

sho-sim.stan (626 Bytes)

ho.r (3.71 KB)

jac.hpp (6.44 KB)

You should run cmdstan in the diagnose mode to check if your JAcobian trick worked fine (but run the diagnose first on a model which uses AD).

It will take some time. I still don’t understand why it works on toy problems and Jacobian trick produces the same results as AD. For real problem, however, AD is sampling (I wrote a code which reads partial files and outputs trace) while with analytic Jacobian the same program produces straight lines for each chain.

It appears trick is not ready yet. There is jacobian_states method still. I don’t know in how many places the jacobian is used.

It seems I found the missing piece. Even when initial conditions are know and parameters are unknown (what corresponds to coupled_ode_system<F, double, var> the jacobian_states method creates coupled_ode_system<F, var, double> so operators of coupled_ode_system<F, double, var> and coupled_ode_system<F, var, double>have to be modified to handle analytic jacobian. I still didn’t understand the reason why this is done so. However, the analytic jacobian trick cuts down the time 300% compared to AD for a toy problem. I will check if it works for real problem but just in case the code is attached.

myCoupledODESystem.hpp (19.5 KB)

It looks like I made one typo. The code I am using is not too elegant but looks like it works on a real problem which is running right now. I had to update one Stan’s .hpp and add few hpp/cpp files. If interested I can explain how to provide analytic jacobian for 2.18.

myCoupledODESystem.hpp (5.65 KB)

help.cpp (985 Bytes)

help.hpp (579 Bytes)

my_cvodes_ode_data.hpp (7.27 KB)

1 Like