Search code examples
pythonscipy

solving vectors of ODE's in python


I would like to know if it is possible to solve a vector of ODE's in scipy.

Here is an example from scipy's page:

import numpy as np

from scipy.integrate import solve_ivp

def lotkavolterra(t, z, a, b, c, d):
    x, y = z
    return [a*x - b*x*y, -c*y + d*x*y]


sol = solve_ivp(lotkavolterra, [0, 15], [10, 5], args=(1.5, 1, 3, 1),

                dense_output=True)


t = np.linspace(0, 15, 300)

z = sol.sol(t)

Can I solve this equation if I have a vector of initial conditions and parameters in a single solve_ivp call?

ICs = np.random.rand(10, 2)
args = np.random.rand(10, 4)

Solution

  • Yes, this is possible, but it is not recommended.

    1. Because you are varying the parameters, the solution of the ODEs will be limited by the stiffest ODE (since they are all being evaluated at the same times). This can result in slower code.
    2. This is not the intended functionality but is more of an abuse of code. It requires being careful with how your variables are created and used

    That said, I'll still show you how you can do it. Before we begin, we must be aware of the requirements of the solver.

    1. y0 must be 1D, so you will need to flatten your variables and reshape them in the function.
    2. args will split up the variables, so they must be put in a single numpy array that is placed inside a tuple, which is passed as one argument.
    3. You must use the same t_span and function for all of the simulations.

    For creating y0 and args to be compatible, you can do the following:

    rng = np.random.default_rng(0)
    
    N = 2
    y0 = rng.random((N,2)).flat
    args = (rng.random((N,4)),)
    

    Here, N is the number of simulations you are doing at once. Each simulation has 2 variables and 4 arguments. Notice how y0 is flattened and args is placed inside a tuple.

    For the function, you need to now accept a single argument that will be split up inside. And, as mentioned before, you will need to reshape the solution vector. With the form I used above, we need to flatten it such that each simulation solution is kept together, so that will require transposing the before it is flattened. Undoing the transpose is also needed when extracting the variables, which is why the reshape and transpose is done when defining x and y.

    def lotkavolterra(t, z, args):
        a, b, c, d = args[0].T  # transposed because the columns are the parameters
        x, y = z.reshape(-1, 2).T
        return np.array([a*x - b*x*y, -c*y + d*x*y]).T.flat
    

    All together, we have:

    import numpy as np
    from scipy.integrate import solve_ivp
    import matplotlib.pyplot as plt
    
    plt.close("all")
    
    def lotkavolterra(t, z, args):
        a, b, c, d = args[0].T
        x, y = z.reshape(-1, 2).T
        return np.array([a*x - b*x*y, -c*y + d*x*y]).T.flat
    
    
    rng = np.random.default_rng(0)
    
    N = 2
    t0 = 0
    tn = 20
    y0 = rng.random((N,2)).flat
    args = (rng.random((N,4)),)
    sol = solve_ivp(lotkavolterra, [0, 15], y0, args=args, dense_output=True)
    
    
    t = np.linspace(0, 15, 300)
    z = sol.sol(t)
    
    fig, ax = plt.subplots()
    for i in range(N):
        x, y = z[2*i:2*i+2]
        ax.plot(t, x, label=f"x{i}")
        ax.plot(t, y, label=f"y{i}")
        
    ax.legend()
    ax.set_title("Lotka-Volterra System")
    fig.show()
    

    Result:

    Again, this is not the right way to do this and probably shouldn't be used. But in the end, you're the programmer and can decided what you want to do.