Using odeint on multidimensional array

Question

I would like to apply odeint to an array of initial conditions and return a derivative that has the same size as those initial conditions. I could loop over each initial condition, but I think that will be very slow with higher N. The example below is from the docs for odeint. sol_1 works as expected, but sol_2 gives the error ValueError: Initial condition y0 must be one-dimensional.

Does anyone have a clever solution for how to make sol_2 run without just looping over each initial condition? Thanks.

import numpy as np
from scipy.integrate import odeint

def pend(y, t, b, c):
    theta, omega = y
    dydt = [omega, -b*omega - c*np.sin(theta)]
    return dydt

b = 0.25
c = 5.0
t = np.linspace(0, 10, 101)

#   0D initial values, directly from docs
y0_1 = [np.pi - 0.1, 0]
sol_1 = odeint(pend, y0_1, t, args=(b, c))

#   1D initial values, directly from docs
y0_2 = [np.ones((3)) * (np.pi - 0.1), np.zeros((3))]

#   Error here
sol2 = odeint(pend, y0_2, t, args=(b, c))

Answer 1

This is essentially the same answer written by OP, but I think it might be a little cleaner. First, I used the flatten method, which isn't as verbose as the np.ndarray.flatten function. Second, I flatten the results to pass to the ODE function, reshape inside the function for extracting values and using, and then flatten the result for the return. At the very end, I use np.split to extract each of the solutions.

import numpy as np
from scipy.integrate import odeint

def pend(y, t, b, c, N):
    theta, omega = y.reshape(-1,N)
    dydt = np.stack([omega, -b*omega - c*np.sin(theta)])
    return dydt.flatten()

b = 0.25
c = 5.0
t = np.linspace(0, 10, 101)

# Number of inputs
N = 3

# 1D initial values, directly from docs
theta0 = np.ones((N,)) * (np.pi - 0.1)
omega0 = np.zeros((N,))
y0_2 = np.stack([theta0, omega0])

#   Flattened version of the inputs
y0_2f = y0_2.flatten()

# Run altered pend (pend2) with flattened inputs
sol_2 = odeint(pend, y0_2f, t, args=(b, c, N))
sols_split = np.split(sol_2, N, axis=1)

Answer 2

Using the comment from hpaulj I came up with a solution. It's not very pretty because I'm passing N, but it does work and it's almost definitely better than looping. This should work N dimensional inputs as well, it would just require reshaping back to the original shape at the end.

import numpy as np
from scipy.integrate import odeint

def pend(y, t, b, c):
    theta, omega = y
    dydt = [omega, -b*omega - c*np.sin(theta)]
    return dydt

def pend2(y, t, b, c, N):
    theta = y[:N]
    omega = y[N:]
    dydt = [omega, -b*omega - c*np.sin(theta)]
    dydt = np.ndarray.flatten(np.array(dydt))
    return dydt

b = 0.25
c = 5.0
t = np.linspace(0, 10, 101)

#   0D initial values, directly from docs
y0_1 = [np.pi - 0.1, 0]
sol_1 = odeint(pend, y0_1, t, args=(b, c))

#   1D initial values, directly from docs
theta0 = np.ones((3)) * (np.pi - 0.1)
omega0 = np.zeros((3))
y0_2 = [theta0, omega0]

#   Flattened version of the inputs
y0_2f = np.ndarray.flatten(np.array(y0_2))

#   Number of inputs
N = len(theta0)

#   Run altered pend (pend2) with flattened inputs
sol_2 = odeint(pend2, y0_2f, t, args=(b, c, N))