Vectorised second order ode solving in python

Question

I am trying to write a python program that simulates the motion of a large number of particles by numerically integrating a second order ordinary differential equation. I first split the ODE into two coupled first order ODEs and solve using scipy.integrate.solve_ivp following a similar method to what is described in the answer to this question. However, my problem is that I want to solve this system for a large number of particles, each with different initial conditions. Naively I could do this with a for loop, but I'm sure numpy and scipy must have a way of vectorising this operation that would be much faster.

I have had a look at the documentation for scipy.integrate.solve_ivp and it talks about vectorisation, but not in the way I want. I would expect for a system with two coupled first order ODEs and n particles you could input the initial conditions as an array with size (2,n) but this is not the case.

Is there a way to solve for multiple initial conditions without resorting to a slow python for loop?

For reference the system of ODEs I want to solve looks like

dx/dt = v
dv/dt = F(x,v)

With initial conditions in an array like

initialConditions = [[x0,v0],[x1,v1],...,[xN,vN]]

Answer 1

Is there a way to solve for multiple initial conditions without resorting to a slow python for loop?

I think the solution to your problem might be parallelization, not vectorization, unfortunately. While Python is going to be slower than C or C++, the for loop in question isn't actually a part of the computation, it's a different problem numerically because of the different initial conditions. Besides, there's so much data involved in the computation there wouldn't be any spare room on the cache lines anyhow.

As far as vectorization is concerned, I'd be willing to bet there's not much else that can be done to improve the code in that respect. The integration algorithms are implemented in Fortran, which is faster than even C and C++ a lot of times because of its strict aliasing rules, which allow for more aggressive optimizations.

As far as the Python code itself is concerned, in particular your loop, you might even be hard-pressed to get a measurable improvement in execution speed by writing it in C, just because of how negligible of an impact that part of the code has in comparison to the actual heavy lifting being done during the integration, and that's not even factoring in how long it would take you to convert Python to C.

The upside is that the AVX-512 instruction set came out long enough ago by now that you might get lucky enough to be able to run your code on a vectorization register that's twice as big as anything I've ever gotten to stand in the presence of, and maybe even some solid GPUs to boot. It might make enough of a difference, you never know; unfortunately optimization problems are always by definition super specific, so the answer to most questions is usually just "it depends." Sorry I couldn't be of more help.

Usually, the standard advice here would be to write try running a few trials, profile the execution, identify where the bottleneck is, and focus all your efforts on that, as even a few percentage points of improvement could mean dozens (or even hundreds) of hours of precious computing time. Unfortunately, the scipy code is probably not getting much better, so in this case I would probably try to somehow reduce the sheer number of initial conditions to be tested analytically. I'm thinking of something like averaging if you were working with the Navier-Stokes equations, for example.