Scipy odeint with banded Jacobian matrix

Question

I'm integrating a system of stiff ODE's using SciPy's integrate.odeint function. As the integration is non-trivial and time consuming I'm also using the corresponding jacobian. By rearranging the equations I can define the jacobian to be a banded matrix. Following the API documentation I would like to define the shape with the mu and ml params. Unfortunately, the documentation is a bit ambiguous so that I was not able to figure out how to implement my jacobian function.

In order to verify how odeint has to be called, I've been using the following (somewhat silly) code:

from scipy.integrate import odeint
lmax = 5

def f1(y, t):
    ydot = np.zeros(lmax)
    for i in range(lmax):
        ydot[i] = y[i]**2-y[i]**3
    return ydot

def fulljac(y, t,):
    J = np.zeros((lmax, lmax))
    J[0,0]=-3*y[0]**2 + 2*y[0]
    J[1,1]=-3*y[1]**2 + 2*y[1]
    J[2,2]=-3*y[2]**2 + 2*y[2]
    J[3,3]=-3*y[3]**2 + 2*y[3]
    J[4,4]=-3*y[4]**2 + 2*y[4]
    return J

## initial conditions and output times
delta = 0.0001;
yini  = np.array([delta]*lmax)
times = np.linspace(0, 2/delta, 100)

y, infodict = odeint(f1, yini, times, Dfun=fulljac, full_output=1)
print("f1: nst: {0}, nfe: {1}, nje: {2}".format(infodict["nst"][-1],
                                                infodict["nfe"][-1],
                                                infodict["nje"][-1]))

Using the full NxN jacobian matrix the integration is successfull. Using only the diagonal and mu=0 and ml=0 the integration succeeds as well.

To test the banded matrix use case I'm creating an artificial 3xN banded jacobian using mu=1 and ml=1, where all the derivatives off the diagonal are zero. This causes a weird behavior of the solver (similar to what I see in my original problem where the off-diagonals are non zero).

def bandjac(y, t):
    J = np.zeros((lmax, 3))
    J[0,1]=-3*y[0]**2 + 2*y[0]
    J[1,1]=-3*y[1]**2 + 2*y[1]
    J[2,1]=-3*y[2]**2 + 2*y[2]
    J[3,1]=-3*y[3]**2 + 2*y[3]
    J[4,1]=-3*y[4]**2 + 2*y[4]
    return J

y, infodict = odeint(f1, yini, times, Dfun=bandjac, full_output=1, mu=1, ml=1)
print("f1: nst: {0}, nfe: {1}, nje: {2}".format(infodict["nst"][-1],
                                                infodict["nfe"][-1],
                                                infodict["nje"][-1]))

What is the proper way to use the banded jacobian option with SciPy's integrate.odeint?

Answer 1

For the sake of completeness I’m replying to my own question.

As pointed out by Warren Weckesser there is a bug in Scipy <0.14.0 on how odeint handles banded jacobians.

The current documentation of odeint states:

Dfun should return a matrix whose rows contain the non-zero bands (starting with the lowest diagonal).

Which I believe is incorrect. Instead it should start with the highest diagonal.

The following code snippet shows how Dfun should return the jacobian (derived from the test_integrate.py unit test):

def func(y, t, c):
    return c.dot(y)

def jac(y, t, c):
    return c

def bjac_rows(y, t, c):
    return np.array([[   0,    75,       1,  0.2], # mu - upper 
                     [ -50,  -0.1,  -1e-04,  -40], # diagonal
                     [ 0.2,   0.1,    -0.2,   0]]) # lower - ml

c = array([[-50,    75,     0,   0],
            [0.2, -0.1,     1,   0],
            [0,    0.1, -1e-4,   0.2],
            [0,      0,   -0.2, -40]])

y0 = arange(4)

t = np.linspace(0, 50, 6)

# using the full jacobian
sol0, info0 = odeint(func, y0, t, args=(c,), full_output=True,Dfun=jac)
print("f1: nst: {0}, nfe: {1}, nje: {2}".format(info0["nst"][-1],
                                                info0["nfe"][-1],
                                                info0["nje"][-1]))

# using the row based banded jacobian
sol2, info2 = odeint(func, y0, t, args=(c,), full_output=True,Dfun=bjac_rows, ml=1, mu=1)
print("f1: nst: {0}, nfe: {1}, nje: {2}".format(info2["nst"][-1],
                                                info2["nfe"][-1],
                                                info2["nje"][-1]))

Note: the transposed banded matrix does not seem to work with col_deriv=True

Answer 2

I made a function to convert a jacobian matrix to banded form as expected by odeint, as well as the mu and ml parameters. The input is expected to be an np.array. For best performance, you probably want to turn the result matrix into an np.array, and then have the jacobian function return its transpose, and use col_der=True. I think BrainGrylls is correct that you should start with the highest diagonal first, contrary to the current API documentation.

def make_banded_jacobian(matrix):
    '''returns a banded jacobian list (in odeint's format), along with mu and ml parameters'''

    # first find the values of mu and ml
    dims = matrix.shape[0]
    assert dims == matrix.shape[1]
    mu = 0
    ml = 0

    for row in xrange(dims):
        for col in xrange(dims):
            if matrix[row][col] != 0:
                if col > row:
                    dif = col - row
                    mu = max(mu, dif)
                else:
                    dif = row - col
                    ml = max(ml, dif)

    banded = []

    for yoffset in xrange(-mu, ml+1):
        row = []

        for diag in xrange(dims):
            x_index = diag
            y_index = diag + yoffset

            if y_index < 0 or y_index >= dims:
                row.append(0.0)
            else:
                row.append(matrix[y_index][x_index])

        banded.append(row)

    return (banded, mu, ml)