Looking for a fast implementation of many trajectories using odeint

[alfercorral]

I am using jax.experimental.odeint for solving a coupled ODE with many possible trajectories (including different initial conditions for space and time) in a MonteCarlo fashion, that is, generating millions of random events within a probability distribution and solving many of them to see how the ending distribution looks like. For every event, I will find different initial conditions for space configuration y0: (6,) and time t0: (1,), but all trajectories end in the same final time t_f.

I found that odeint is very fast when working with same time expansion, that is, if all events started in the same t0. In that case, no further coding needs to be done, as odeint can take care of many initial positions in space configuration. However, when trying to work with many different time expansions, arising from:

t = [ jnp.linspace(t_i,t_f,n_steps) for t_i in t_0 ]
t = jnp.array(t)

I cannot find an optimised version of the code that allows for solutions as fast as in the same initial time expansion. The working time of the program goes from seconds for thousands of trajectories to minutes for hundreds of them with different initial conditions.

I have also tried implementing a jitted vmapped version of the code:

def odeint(y0,t):
return jax.vmap(ode.odeint, in_axes = (None,0,0))(force,y0,t)

odeint = jax.jit(odeint)

Or even a pmapped version of it:

def ode_int(y0,t):
return ode.odeint(force,y0,t)

inte = jax.jit(jax.pmap(ode_int, in_axes = (0, 0)))

But I find that for different reasons, these are not great solutions.

Is there any nice way of having thousands of trajectories being solved within the reach of odeint or maybe in Diffrax that I don't know of?

[alfercorral]

Also, I looked at the suggested version of Diffrax for similar code in a jitted, vmapped manner:

term = ODETerm(force)
solver = Tsit5() #Dopri5(scan_stages=True) # #

def sol(t_0, y0):
return diffeqsolve(term, t0 = t_0, t1 = t_f , dt0=0.01, y0 = y0.reshape(-1,6),
solver= solver, stepsize_controller=PIDController(rtol=1.4e-8, atol=1.4e-8),
adjoint=BacksolveAdjoint(), max_steps=None )

sol = jax.jit(jax.vmap(sol, in_axes = (0,0)))

Still this solution is not outperforming normal odeint solution.

Any suggestion from developer @patrick-kidger ?

[patrick-kidger]

What you've written looks reasonable to me.

Can you provide a complete MWE demonstrating the difference? (Both the fast and the slow version.)

If you encapsulate your code in three backticks it will be rendered as code, by the way. Like this:

```python
Your code here
```

The "python" gives appropriate syntax highlighting.

[alfercorral]

Thanks for your response.

I noticed now that for many trajectories (10000) at once, Diffrax can outperform regular jitted odeint. Here is a MWE having the three different options: 1st only one initial condition with ode.odeint, 2nd having several with ode.odeint and 3rd several initial conditions for Diffrax:

import numpy as np
from typing import Callable, List
import jax
from jax import numpy as jnp
from jax.config import config; config.update("jax_enable_x64", True)
import time
import math
import time
from jax.experimental import ode

def nuc_pot(Z, xyz, atom_pos):
    return Z / jnp.linalg.norm(xyz - atom_pos, axis = -1)

###  And take space derivative from it

def nuc_field(Z, xyz, atom_pos):
    def nuc_pot_x(xyz):
        return nuc_pot(Z,xyz,atom_pos)
    return jax.vmap(jax.jacrev(nuc_pot_x), in_axes = 0)(xyz)

### Get the total electric field by adding up the contributions of electric fields

def total_nuc_field(Z, xyz, atom_pos):
    return jnp.sum(jnp.array([nuc_field(Z[i], xyz, atom_pos[i]) for i in range(len(Z))]), axis=0)


Z = jnp.array([8,6,16])
atom_pos = jnp.array([[ 0.   ,  3.,  0.   ],
       [ 0.   ,  0.,  0.   ],
       [ 0.   , -1.,  0.   ]])

dt = 1e-2
n_events = 1000
## Use only one initial time for all trajectories (Dont bother for numerics, they are physically motivated quantities)

t_0 = -1.19947351
t_f =  jnp.pi* 658.4723441615452

## Start with random conditions $(\vec{v}, \vec{x})$ for a first order coupled Newton equation
y_0 = jnp.array(np.random.random((n_events,6)))

t = jnp.arange(t_0,t_f,dt) 

def force(xyz, t):
    x = xyz[:,3:]
    acc =  - total_nuc_field(Z, x, atom_pos)
    dx_dt = xyz[:, :3]
    return jnp.hstack((acc,dx_dt))    ### Derivatives are $(\vec{a}, \vec{v})$

force = jax.jit(force)

def ode_int(y_0, t):
    return ode.odeint(force,y_0,t)

ode_int = jax.jit(ode_int)

%time sol  = ode_int(y_0,t)

%time sol  = ode_int(y_0,t)

### Result is roughly 8 s.

## Using many random initial conditions

t_0 = jnp.array(np.random.random((n_events,)))

n_steps = int( (t_f - jnp.min(t_0) )/dt)

t = [ jnp.linspace(t_i,t_f,n_steps) for t_i in t_0 ]
t = jnp.array(t)

def odeint(y0,t):
    return ode.odeint(force,y0.reshape(1,6),t)

odeint = jax.jit(jax.vmap(odeint, in_axes = (0,0)))

%time sol  = odeint(y_0,t)

%time sol  = odeint(y_0,t)

### Results are of about 30 seconds

### Last, using Diffrax

from diffrax import diffeqsolve, ODETerm, Dopri5, Tsit5, PIDController, BacksolveAdjoint

def force(t, xyz, args):
    x = xyz[:,3:]
    acc = - total_nuc_field(Z, x, atom_pos) 
    dx_dt = xyz[:,:3]
    return jnp.hstack((acc,dx_dt))



term = ODETerm(force)
solver = Tsit5() #Dopri5(scan_stages=True) # #


def sol(t_0, y0):
    return diffeqsolve(term, t0 = t_0, t1 = t_f , dt0=0.01, y0 = y0.reshape(-1,6), \
solver= solver, stepsize_controller=PIDController(rtol=1.4e-8, atol=1.4e-8), \
adjoint=BacksolveAdjoint(), max_steps=None )

sol = jax.jit(jax.vmap(sol, in_axes = (0,0)))

%time sol(t_0[:2], y_0[:2])
%time solution_ = sol(t_0, y_0)


#### Roughly 2 seconds

I don't know if you could suggest of a better way of coding such implementation. Usually I don't work with an analytical function for force but rather with an interpolated one, which also slows down the procedure.

[patrick-kidger]

So under-the-hood, odeint and Diffrax use different strategies for vmap'ing. I'll skip the technical details, but generally speaking Diffrax's is more efficient. That might explain why Diffrax is so much faster than odeint when ran on a batch.