Testing banded Jacobian LSODA in scipy

band_lsoda_results

# running with n=1, ml=1, mu=2
Solver     Full/Band    With/out Jac      # Func. Eval    # Jac. Eval
---------  -----------  --------------  --------------  -------------
solve_ivp  full         no jac                    1101             63
solve_ivp  full         jac                        722             48
solve_ivp  band         no jac                    1038             63
solve_ivp  band         jac                        722             48
odeint     full         no jac                     983             52
odeint     full         jac                        726             56
odeint     band         no jac                     931             52
odeint     band         jac                        726             56

# running with n=10, ml=1, mu=2

Solver     Full/Band    With/out Jac      # Func. Eval    # Jac. Eval
---------  -----------  --------------  --------------  -------------
solve_ivp  full         no jac                    3108             48
solve_ivp  full         jac                        720             53
solve_ivp  band         no jac                    1038             63
solve_ivp  band         jac                        722             48
odeint     full         no jac                    3377             53
odeint     full         jac                        721             47
odeint     band         no jac                     931             52
odeint     band         jac                        726             56


# running with n=10, ml=5, mu=5
Solver     Full/Band    With/out Jac      # Func. Eval    # Jac. Eval
---------  -----------  --------------  --------------  -------------
solve_ivp  full         no jac                    3108             48
solve_ivp  full         jac                        720             53
solve_ivp  band         no jac                    1479             63
solve_ivp  band         jac                        722             48
odeint     full         no jac                    3377             53
odeint     full         jac                        721             47
odeint     band         no jac                    1295             52
odeint     band         jac                        726             56

test.py

from scipy.integrate import solve_ivp, odeint
import numpy as np
from tabulate import tabulate

# parameters to change
# ml and mu must also be smaller than approximately 5*n/2
n = 10  # size of ode (5*n equations)
ml = 1  # must be >=1
mu = 2  # must be >=2


# f1, jac1, and band_jac1 are the original n=1 cases for comparison

def f1(t, y):
    return np.array([
        y[0],
        -0.04 * y[1] + 1e4 * y[2] * y[3],
        0.04 * y[1] - 1e4 * y[2] * y[3] - 3e7 * y[2]**2,
        3e7 * y[2]**2,
        y[4]
    ])

def jac1(t, y):
    return np.array([
        [1,     0,                    0,         0, 0],
        [0, -0.04,             1e4*y[3],  1e4*y[2], 0],
        [0,  0.04, -1e4*y[3]-3e7*2*y[2], -1e4*y[2], 0],
        [0,     0,           3e7*2*y[2],         0, 0],
        [0,     0,                    0,         0, 1]
    ])

def band_jac1(t, y):
    return np.array([
        [0,     0,                    0,  1e4*y[2], 0],
        [0,     0,             1e4*y[3], -1e4*y[2], 0],
        [1, -0.04, -1e4*y[3]-3e7*2*y[2],         0, 1],
        [0,  0.04,           3e7*2*y[2],         0, 0]
    ])


def gen_f(n):
    """Tile f n times, resulting in 5*n equations."""
    def f(t, y):
        arr = np.zeros(5*n)
        for i in range(n):
            arr[5*i+0:5*i+5] = [
                y[0+5*i],
                -0.04 * y[1+5*i] + 1e4 * y[2+5*i] * y[3+5*i],
                0.04 * y[1+5*i] - 1e4 * y[2+5*i] * y[3+5*i] - 3e7 * y[2+5*i]**2,
                3e7 * y[2+5*i]**2,
                y[4+5*i]
            ]
        return arr
    return f

def gen_jac(n):
    "Tile jac n times, resulting in (5*n, 5*n) jacobian."
    def jac(t, y):
        arr = np.zeros((5*n, 5*n))

        for i in range(n):
            arr[5*i, 5*i] = 1
            arr[5*i+1, 5*i+1] = -0.04
            arr[5*i+1, 5*i+2] = 1e4*y[3+5*i]
            arr[5*i+1, 5*i+3] = 1e4*y[2+5*i]
            arr[5*i+2, 5*i+1] = 0.04
            arr[5*i+2, 5*i+2] = -1e4*y[3+5*i]-3e7*2*y[2+5*i]
            arr[5*i+2, 5*i+3] = -1e4*y[2+5*i]
            arr[5*i+3, 5*i+2] = 3e7*2*y[2+5*i]
            arr[5*i+4, 5*i+4] = 1.0
        return arr
    return jac

def gen_band_jac(jac_fun, n, kl, ku):
    """"Make fully compressed banded jacobian, shape (kl+ku+1, n)."""
    def jac(t, y):
        j = jac_fun(t, y)
        m, n0 = j.shape
        assert m == n
        assert n0 == n

        compressed_band_jac = np.zeros((kl + ku + 1, n))

        for k, i in enumerate(reversed(range(ku))):
            diag = np.diagonal(j, i+1)
            compressed_band_jac[k] = np.hstack((np.zeros(n-diag.size), diag))

        compressed_band_jac[ku] = np.diagonal(j, 0)

        for k, i in enumerate(range(kl)):
            diag = np.diagonal(j, -i-1)
            compressed_band_jac[k+ku+1] = np.hstack((diag, np.zeros(n-diag.size)))

        return compressed_band_jac
    return jac

def gen_expanded_band_jac(jac_fun, n, kl, ku):
    """Add kl blank rows at end to fully compressed banded jacobian, shape (2*kl+ku+1, n)."""
    band_jac = gen_band_jac(jac_fun, n, kl, ku)
    def jac(t, y):
        j = band_jac(t, y)
        arr = np.zeros((kl*2 + ku + 1, n))
        arr[:kl+ku+1, :] = j
        return arr
    return jac

# Make sure that the generated cases match expectations
f = gen_f(1)
jac = gen_jac(1)
y0 = 1e-4
y = np.ones(5) * y0
assert np.allclose(f(0, y), f1(0, y))
assert np.allclose(jac(0, y), jac1(0, y))

band_jac = gen_band_jac(jac, 5, 1, 2)
assert np.allclose(band_jac(0, y), band_jac1(0, y))

# Solutions
rtol = 1e-8
atol = 1e-10

tspan = (0, 1000)
t_eval = np.linspace(*tspan, 100)

assert ml >= 1
assert mu >= 2

y0 = np.tile(np.array([0, 1, 0, 0, 0]), n)

sol_full_no_jac = solve_ivp(gen_f(n), tspan, y0, t_eval=t_eval, method='LSODA',
 rtol=rtol, atol=atol)

sol_full_jac = solve_ivp(gen_f(n), tspan, y0, t_eval=t_eval, method='LSODA',
 rtol=rtol, atol=atol, jac=gen_jac(n))

sol_band_no_jac = solve_ivp(gen_f(n), tspan, y0, t_eval=t_eval, method='LSODA',
 rtol=rtol, atol=atol, lband=ml, uband=mu)

sol_band_jac = solve_ivp(gen_f(n), tspan, y0, t_eval=t_eval, method='LSODA',
 rtol=rtol, atol=atol, jac=gen_expanded_band_jac(gen_jac(n), n*5, ml, mu), lband=ml, uband=mu)

_, odeint_full_no_jac = odeint(gen_f(n), y0, t_eval, tfirst=True, full_output=True, rtol=rtol, atol=atol)
_, odeint_full_jac = odeint(gen_f(n), y0, t_eval, tfirst=True, full_output=True, rtol=rtol, atol=atol, Dfun=gen_jac(n))
_, odeint_band_no_jac = odeint(gen_f(n), y0, t_eval, tfirst=True, full_output=True, rtol=rtol, atol=atol, ml=ml, mu=mu)
_, odeint_band_jac = odeint(gen_f(n), y0, t_eval, tfirst=True, full_output=True, rtol=rtol, atol=atol, Dfun=gen_band_jac(gen_jac(n), n*5, ml, mu), ml=ml, mu=mu)

print(tabulate(
    [
        ["solve_ivp", "full", "no jac", sol_full_no_jac.nfev, sol_full_no_jac.njev],
        ["solve_ivp", "full", "jac", sol_full_jac.nfev, sol_full_jac.njev],
        ["solve_ivp", "band", "no jac", sol_band_no_jac.nfev, sol_band_no_jac.njev],
        ["solve_ivp", "band", "jac", sol_band_jac.nfev, sol_band_jac.njev],
        ["odeint", "full", "no jac", odeint_full_no_jac["nfe"][-1], odeint_full_no_jac["nje"][-1]],
        ["odeint", "full", "jac", odeint_full_jac["nfe"][-1], odeint_full_jac["nje"][-1]],
        ["odeint", "band", "no jac", odeint_band_no_jac["nfe"][-1], odeint_band_no_jac["nje"][-1]],
        ["odeint", "band", "jac", odeint_band_jac["nfe"][-1], odeint_band_jac["nje"][-1]],

    ],
    headers = ["Solver", "Full/Band", "With/out Jac", "# Func. Eval", "# Jac. Eval"]
))