# -*- coding: utf-8 -*-
from __future__ import (absolute_import, division, print_function)
import numpy as np
[docs]def banded_jacobian(y, x, ml, mu):
""" Calculates a banded version of the jacobian
Compatible with the format requested by
:func:`scipy.integrate.ode` (for SciPy >= v0.15).
Parameters
----------
y: array_like of expressions
x: array_like of symbols
ml: int
number of lower bands
mu: int
number of upper bands
Returns
-------
2D array of shape ``(1+ml+mu, len(y))``
"""
ny = len(y)
nx = len(x)
packed = np.zeros((mu+ml+1, nx), dtype=object)
def set(ri, ci, val):
packed[ri-ci+mu, ci] = val
for ri in range(ny):
for ci in range(max(0, ri-ml), min(nx, ri+mu+1)):
set(ri, ci, y[ri].diff(x[ci]))
return packed