Merge pull request #95 from bjodah/sparse_jac

Sparse jac

conda-recipe/meta.yaml

@@ -21,7 +21,7 @@ requirements:
   run:
     - python
     - numpy
-    - sym >=0.3.3
+    - sym >=0.3.4
     - sympy >=1.1.1,!=1.2
     - scipy
     - matplotlib

examples/_sparse_matrix.ipynb

@@ -0,0 +1,131 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "import sympy\n",
+    "from pyodesys.symbolic import SymbolicSys\n",
+    "from pyodesys.tests._robertson import get_ode_exprs\n",
+    "sympy.init_printing()\n",
+    "%matplotlib inline"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "logc, logt, reduced, powsimp = False, False, False, False\n",
+    "odesys1 = SymbolicSys.from_callback(get_ode_exprs(logc, logt, reduced)[0],\n",
+    "                                   2 if reduced else 3, 6 if reduced else 3)\n",
+    "odesys1.exprs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "odesys1._jac"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Converting to sparse representation of jacobian:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sparse = SymbolicSys.from_other(odesys1, sparse=True)\n",
+    "sparse.nnz, sparse._rowvals, sparse._colptrs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sparse._jac"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Converting back to dense representation of jacobian:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "odesys2 = SymbolicSys.from_other(sparse, sparse=False)\n",
+    "assert odesys2.nnz == -1\n",
+    "odesys2._jac"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Solving the initial value problem numerically:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def solve_and_plot(odesys, integrate_kw, tout=[0, 1e11], y0=[1,0,0],\n",
+    "                   params=[0.04, 1e4, 3e7], ax=None):\n",
+    "    result = odesys.integrate(tout, y0, params, nsteps=5000,\n",
+    "                              integrator='cvode', **integrate_kw)\n",
+    "    result.plot(ax=ax, title_info=2)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pycvodes\n",
+    "has_klu = pycvodes.config['NO_KLU'] == '0'\n",
+    "systems = [odesys1, sparse, odesys2]\n",
+    "kws = [{}, dict(linear_solver='klu'), {}]\n",
+    "fig, axes = plt.subplots(1, len(systems), figsize=(5*len(systems), 4),\n",
+    "                         subplot_kw=dict(xscale='log', yscale='log'))\n",
+    "for odesys, ax, kw in zip(systems, axes, kws):\n",
+    "    if kw.get('linear_solver', '') == 'klu' and not has_klu:\n",
+    "        continue\n",
+    "    solve_and_plot(odesys, kw, ax=ax)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

pyodesys/core.py

@@ -48,9 +48,15 @@ class ODESys(object):
         dependent variable (x). Signature is any of:
             - ``rhs(x, y[:]) -> f[:]``
             - ``rhs(x, y[:], p[:]) -> f[:]``
-            - ``rhs(x, y[:], p[:], backend=math) -> f[:]``
+            - ``rhs(x, y[:], p[:], backend=math) -> f[:]``.
     jac : callback
         Jacobian matrix (dfdy). Required for implicit methods.
+        Signature should be either of:
+            - ``jac(x, y[:]) -> J``
+            - ``jac(x, y[:], p[:]) -J``.
+        If ``nnz < 0``, ``J`` should be a 2D array-like object if ``nnz < 0``
+        corresponding to a dense or banded jacobian (see also ``band``).
+        If ``nnz >= 0``, ``J`` should be an instance of ``scipy.sparse.csc_matrix``.
     dfdx : callback
         Signature ``dfdx(x, y[:], p[:]) -> out[:]`` (used by e.g. GSL)
     jtimes : callback
@@ -100,6 +106,9 @@ class ODESys(object):
         Describes whether the independent variable appears in the rhs expressions.
         If set to ``True`` the underlying solver is allowed to shift the
         independent variable during integration.
+    nnz : int (default : -1)
+        Maximum number of non-zero entries in sparse jacobian. When
+        non-negative, jacobian is assumed to be dense or banded.
 
     Attributes
     ----------
@@ -115,6 +124,7 @@ class ODESys(object):
         For calculating the first step based on x0, y0 & p.
     roots_cb : callback
     nroots : int
+    nnz : int
     names : tuple of strings
     param_names : tuple of strings
     description : str
@@ -151,7 +161,7 @@ def __init__(self, f, jac=None, dfdx=None, jtimes=None, first_step_cb=None, root
                  par_by_name=False, latex_names=(), latex_param_names=(), latex_indep_name=None,
                  taken_names=None, pre_processors=None, post_processors=None, append_iv=False,
                  autonomous_interface=None, to_arrays_callbacks=None, autonomous_exprs=None,
-                 _indep_autonomous_key=None, numpy=None, **kwargs):
+                 _indep_autonomous_key=None, numpy=None, nnz=-1, **kwargs):
         self.f_cb = _ensure_4args(f)
         self.j_cb = _ensure_4args(jac) if jac is not None else None
         self.jtimes_cb = _ensure_4args(jtimes) if jtimes is not None else None
@@ -163,6 +173,7 @@ def __init__(self, f, jac=None, dfdx=None, jtimes=None, first_step_cb=None, root
             if not band[0] >= 0 or not band[1] >= 0:
                 raise ValueError("bands needs to be > 0 if provided")
         self.band = band
+        self.nnz = nnz
         self.names = tuple(names or ())
         self.param_names = tuple(param_names or ())
         self.indep_name = indep_name
@@ -577,16 +588,29 @@ def _f(x, y, fout):
             if with_jacobian is None:
                 raise Exception("Must pass with_jacobian")
             elif with_jacobian is True:
-                def _j(x, y, jout, dfdx_out=None, fy=None):
-                    if len(_p) > 0:
-                        jout[:, :] = np.asarray(self.j_cb(x, y, _p))
-                    else:
-                        jout[:, :] = np.asarray(self.j_cb(x, y))
-                    if dfdx_out is not None:
+                if self.nnz >= 0:
+                    def _j(x, y, data, colptrs, rowvals):
+                        if len(_p) > 0:
+                            J = self.j_cb(x, y, _p)
+                        else:
+                            J = self.j_cb(x, y)
+                        J = J.asformat("csc")
+                        data[:] = J.data
+                        colptrs[:] = J.indptr
+                        rowvals[:] = J.indices
+
+                    new_kwargs['nnz'] = self.nnz
+                else:
+                    def _j(x, y, jout, dfdx_out=None, fy=None):
                         if len(_p) > 0:
-                            dfdx_out[:] = np.asarray(self.dfdx_cb(x, y, _p))
+                            jout[:, :] = np.asarray(self.j_cb(x, y, _p))
                         else:
-                            dfdx_out[:] = np.asarray(self.dfdx_cb(x, y))
+                            jout[:, :] = np.asarray(self.j_cb(x, y))
+                        if dfdx_out is not None:
+                            if len(_p) > 0:
+                                dfdx_out[:] = np.asarray(self.dfdx_cb(x, y, _p))
+                            else:
+                                dfdx_out[:] = np.asarray(self.dfdx_cb(x, y))
             else:
                 _j = None
 

pyodesys/native/_base.py

@@ -125,11 +125,16 @@ def variables(self):
         all_invar = tuple(self.odesys.all_invariants())
         ninvar = len(all_invar)
         jac = self.odesys.get_jac()
+        nnz = self.odesys.nnz
         all_exprs = self.odesys.exprs + all_invar
-        if jac is not False:
+        if jac is not False and nnz < 0:
             jac_dfdx = list(reduce(add, jac.tolist() + self.odesys.get_dfdx().tolist()))
             all_exprs += tuple(jac_dfdx)
             nj = len(jac_dfdx)
+        elif jac is not False and nnz >= 0:
+            jac_dfdx = list(reduce(add, jac.tolist()))
+            all_exprs += tuple(jac_dfdx)
+            nj = len(jac_dfdx)
         else:
             nj = 0
 
@@ -233,12 +238,18 @@ def _ccode(expr):
                 'cses': [(symb.name, _ccode(expr)) for symb, expr in jtimes_cses],
                 'exprs': list(map(_ccode, jtimes_exprs))
             },
-            p_jac=None if jac is False else {
+            p_jac_dense=None if jac is False or nnz >= 0 else {
                 'cses': [(symb.name, _ccode(expr)) for symb, expr in jac_cses],
                 'exprs': {(idx//ny, idx % ny): _ccode(expr)
                           for idx, expr in enumerate(jac_exprs[:ny*ny])},
                 'dfdt_exprs': list(map(_ccode, jac_exprs[ny*ny:]))
             },
+            p_jac_sparse=None if jac is False or nnz < 0 else {
+                'cses': [(symb.name, _ccode(expr)) for symb, expr in jac_cses],
+                'exprs': list(map(_ccode, jac_exprs[:nj])),
+                'colptrs': self.odesys._colptrs,
+                'rowvals': self.odesys._rowvals
+            },
             p_first_step=None if first_step is None else {
                 'cses': first_step_cses,
                 'expr': _ccode(first_step_exprs[0]),

pyodesys/native/sources/odesys_anyode_iterative.hpp

@@ -19,6 +19,7 @@ namespace odesys_anyode {
                bool, Real_t, std::vector<Real_t>);
         int nrev=0;  // number of calls to roots
         Index_t get_ny() const override;
+        Index_t get_nnz() const override;
         int get_nquads() const override;
         int get_nroots() const override;
         Real_t get_dx0(Real_t, const Real_t * const) override;
@@ -39,6 +40,12 @@ namespace odesys_anyode {
                                       Real_t * const __restrict__ jac,
                                       long int ldim,
                                       Real_t * const __restrict__ dfdt=nullptr) override;
+        AnyODE::Status sparse_jac_csc(Real_t t,
+                                      const Real_t * const __restrict__ y,
+                                      const Real_t * const __restrict__ fy,
+                                      Real_t * const __restrict__ data,
+                                      Index_t * const __restrict__ colptrs,
+                                      Index_t * const __restrict__ rowvals) override;
         AnyODE::Status jtimes(const Real_t * const __restrict__ vec,
                               Real_t * const __restrict__ out,
                               Real_t t,