Compute propagator for each block of block-diagonal system matrix (#70)

* compute propagator for each block of block-diagonal system matrix * update GSL on GitHub Actions CI * update GSL on GitHub Actions CI * update GSL on GitHub Actions CI * update GSL on GitHub Actions CI * update GSL on GitHub Actions CI * add note about PyGSL PyPI release * compute propagator for each block of block-diagonal system matrix * compute propagator for each block of block-diagonal system matrix * compute propagator for each block of block-diagonal system matrix * compute propagator for each block of block-diagonal system matrix * compute propagator for each block of block-diagonal system matrix * compute propagator for each block of block-diagonal system matrix --------- Co-authored-by: C.A.P. Linssen <[email protected]>
nest · Dec 4, 2024 · 6ddb472 · 6ddb472
1 parent 068c578
commit 6ddb472
Show file tree

Hide file tree

Showing 5 changed files with 55 additions and 42 deletions.
diff --git a/.github/workflows/ode-toolbox-build.yml b/.github/workflows/ode-toolbox-build.yml
@@ -43,7 +43,7 @@ jobs:
       fail-fast: false
       matrix:
         with_gsl: ["0", "1"]
-        sympy_version: ["==1.4", "!=1.5,!=1.5.*,!=1.6rc1,!=1.6rc2,!=1.6,!=1.6.*,!=1.10,!=1.10.*,!=1.13,!=1.13.*"]
+        sympy_version: ["==1.4", ""]  # empty string for "latest"
 
     steps:
       - name: Checkout ODE-toolbox code

diff --git a/doc/index.rst b/doc/index.rst
@@ -41,8 +41,6 @@ ODE-toolbox is written in Python and leverages SymPy for the symbolic manipulati
 Installation
 ------------
 
-.. Attention:: SymPy releases after 1.4 introduce a regression in runtime performance when computing matrix exponentials. Unless this conflicts with other version requirements, we recommend to use SymPy 1.4 for now (for example, by editing ``requirements.txt`` to read ``sympy==1.4``). This issue is being tracked at https://github.com/sympy/sympy/issues/23417.
-
 .. Attention:: To perform solver benchmarking, ODE-toolbox relies on GSL and PyGSL. Currently, the latest PyGSL release is not compatible with GSL. We recommend to use GSL 2.7 for now. This issue is being tracked at https://github.com/pygsl/pygsl/issues/62.
 
 
@@ -472,6 +470,12 @@ If the imaginary unit :math:`i` is found in any of the entries in :math:`\mathbf
 
 In some cases, elements of :math:`\mathbf{P}` may contain fractions that have a factor of the form :python:`param1 - param2` in their denominator. If at a later stage, the numerical value of :python:`param1` is chosen equal to that of :python:`param2`, a numerical singularity (division by zero) occurs. To avoid this issue, it is necessary to eliminate either :python:`param1` or :python:`param2` in the input, before the propagator matrix is generated. ODE-toolbox will detect conditions (in this example, :python:`param1 = param2`) under which these singularities can occur. If any conditions were found, log warning messages will be emitted during the computation of the propagator matrix. A condition is only reported if the system matrix :math:`A` is defined under that condition, ensuring that only those conditions are returned that are purely an artifact of the propagator computation.
 
+To speed up processing, the final system matrix :math:`\mathbf{A}` is rewritten as a block-diagonal matrix :math:`\mathbf{A} = \text{diag}(\mathbf{A}_1, \mathbf{A}_2, \dots, \mathbf{A}_k)`, where each of :math:`\mathbf{A}_1, \mathbf{A}_2, \dots, \mathbf{A}_k` is square. Then, the propagator matrix is computed for each individual block separately, making use of the following identity:
+
+.. math::
+
+   \exp(\mathbf{A}\cdot h) = \text{diag}(\exp(\mathbf{A}_1 \cdot h), \exp(\mathbf{A}_2 \cdot h), \dots, \exp(\mathbf{A}_k \cdot h))
+
 
 Computing the update expressions
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -570,26 +574,6 @@ The file `test_mixed_integrator_numeric.py <https://github.com/nest/ode-toolbox/
    <img src="https://raw.githubusercontent.com/nest/ode-toolbox/master/doc/fig/test_mixed_integrator_numeric.png" alt="g_in, g_in__d, g_ex, g_ex__d, V_m timeseries plots" width="620" height="451">
 
 
-Caching of results
-------------------
-
-.. admonition:: TODO
-
-   Not implemented yet.
-
-Some operations on SymPy expressions can be quite slow (see the section :ref:`Working with large expressions`\ ).
-
-Even dynamical systems of moderate size can require a few minutes of processing time, in large part due to SymPy calls, and solver selection.
-
-To speed up processing, a caching mechanism analyses the final system matrix :math:`\mathbf{A}` and rewrites it as a block-diagonal matrix :math:`\mathbf{A} = \text{diag}(\mathbf{B}_1, \mathbf{B}_2, \dots, \mathbf{B}_k)`, where each of :math:`\mathbf{B}_1, \mathbf{B}_2, \dots, \mathbf{B}_k` is square.
-
-For propagators, we note that
-
-.. math::
-
-   e^{\mathbf{A}t} = \text{diag}(e^{\mathbf{B}_1t}, e^{\mathbf{B}_2t}, \dots, e^{\mathbf{B}_kt})
-
-
 API documentation
 -----------------
 

diff --git a/odetoolbox/system_of_shapes.py b/odetoolbox/system_of_shapes.py
@@ -24,6 +24,7 @@
 import logging
 import numpy as np
 import scipy
+import scipy.linalg
 import scipy.sparse
 import sympy
 import sympy.matrices
@@ -34,11 +35,26 @@
 from .sympy_helpers import _custom_simplify_expr, _is_zero
 
 
-def off_diagonal_is_zero(row: int, A) -> bool:
-    for col in range(A.shape[1]):
-        if col != row and not _is_zero(A[row, col]):
-            return False
-    return True
+def get_block_diagonal_blocks(A):
+    assert A.shape[0] == A.shape[1], "matrix A should be square"
+
+    A_mirrored = (A + A.T) != 0   # make the matrix symmetric so we only have to check one triangle
+
+    graph_components = scipy.sparse.csgraph.connected_components(A_mirrored)[1]
+
+    assert all(np.diff(graph_components) >= 0), "Matrix is not ordered"
+
+    blocks = []
+    for i in np.unique(graph_components):
+        idx = np.where(graph_components == i)[0]
+        assert all(np.diff(idx) > 0)
+        assert len(idx) == 1 or (len(np.unique(np.diff(idx))) == 1 and np.unique(np.diff(idx))[0] == 1)
+        idx_min = np.amin(idx)
+        idx_max = np.amax(idx)
+        block = A[idx_min:idx_max + 1, idx_min:idx_max + 1]
+        blocks.append(block)
+
+    return blocks
 
 
 class PropagatorGenerationException(Exception):
@@ -178,22 +194,26 @@ def get_sub_system(self, symbols):
         return SystemOfShapes(x_sub, A_sub, b_sub, c_sub, shapes_sub)
 
 
-    def generate_propagator_solver(self):
-        r"""
-        Generate the propagator matrix and symbolic expressions for propagator-based updates; return as JSON.
+    def _generate_propagator_matrix(self, A):
+        r"""Generate the propagator matrix by matrix exponentiation.
+
+        XXX: the default custom simplification expression does not work well with sympy 1.4 here. Consider replacing sympy.simplify() with _custom_simplify_expr() if sympy 1.4 support is dropped.
         """
 
-        #
-        #   generate the propagator matrix
-        #
+        # naive: calculate propagators in one step
+        # P_naive = sympy.simplify(sympy.exp(A * sympy.Symbol(Config().output_timestep_symbol)))
 
-        P = sympy.simplify(sympy.exp(self.A_ * sympy.Symbol(Config().output_timestep_symbol)))    # XXX: the default custom simplification expression does not work well with sympy 1.4 here. Consider replacing sympy.simplify() with _custom_simplify_expr() if sympy 1.4 support is dropped.
+        # optimized: be explicit about block diagonal elements; much faster!
+        blocks = get_block_diagonal_blocks(np.array(A))
+        propagators = [sympy.simplify(sympy.exp(sympy.Matrix(block) * sympy.Symbol(Config().output_timestep_symbol))) for block in blocks]
+        P = sympy.Matrix(scipy.linalg.block_diag(*propagators))
 
+        # check the result
         if sympy.I in sympy.preorder_traversal(P):
             raise PropagatorGenerationException("The imaginary unit was found in the propagator matrix. This can happen if the dynamical system that was passed to ode-toolbox is unstable, i.e. one or more state variables will diverge to minus or positive infinity.")
 
         try:
-            condition = SingularityDetection.find_singularities(P, self.A_)
+            condition = SingularityDetection.find_singularities(P, A)
             if condition:
                 logging.warning("Under certain conditions, the propagator matrix is singular (contains infinities).")
                 logging.warning("List of all conditions that result in a singular propagator:")
@@ -208,6 +228,15 @@ def generate_propagator_solver(self):
         logging.debug("b = " + str(self.b_))
         logging.debug("c = " + str(self.c_))
 
+        return P
+
+
+    def generate_propagator_solver(self):
+        r"""
+        Generate the propagator matrix and symbolic expressions for propagator-based updates; return as JSON.
+        """
+
+        P = self._generate_propagator_matrix(self.A_)
 
         #
         #   generate symbols for each nonzero entry of the propagator matrix

diff --git a/requirements.txt b/requirements.txt
@@ -1,4 +1,4 @@
-sympy!=1.5,!=1.5.*,!=1.6rc1,!=1.6rc2,!=1.6,!=1.6.*,!=1.10,!=1.10.*,!=1.13,!=1.13.*
+sympy
 scipy
 numpy>=1.8.2
 pytest

diff --git a/tests/test_system_matrix_construction.py b/tests/test_system_matrix_construction.py
@@ -34,9 +34,9 @@ def test_system_matrix_construction(self):
         shapes, parameters = _from_json_to_shapes(indict)
         sigma, beta = sympy.symbols("sigma beta")
         shape_sys = SystemOfShapes.from_shapes(shapes, parameters=parameters)
-        assert shape_sys.A_ == sympy.Matrix([[-sigma, sigma, 0.0],
-                                             [0.0, 0.0, 0.0],
-                                             [0.0, 0.0, -beta]])
+        assert shape_sys.A_ == sympy.Matrix([[-sigma, sigma, 0],
+                                             [0, 0, 0],
+                                             [0, 0, -beta]])
         x, y, z = sympy.symbols("x y z")
         assert shape_sys.c_ == sympy.Matrix([[0],
                                              [3 * z * x**2 - x * y],
@@ -63,10 +63,10 @@ def test_morris_lecar(self):
         shape_sys = SystemOfShapes.from_shapes(shapes, parameters=parameters)
         C_m, g_Ca, g_K, g_L, E_Ca, E_K, E_L, I_ext = sympy.symbols("C_m g_Ca g_K g_L E_Ca E_K E_L I_ext")
         assert shape_sys.A_ == sympy.Matrix([[sympy.parsing.sympy_parser.parse_expr("-500.0 * g_Ca / C_m - 1000.0 * g_L / C_m"), sympy.parsing.sympy_parser.parse_expr("1000.0 * E_K * g_K / C_m")],
-                                            [sympy.parsing.sympy_parser.parse_expr("0.0"), sympy.parsing.sympy_parser.parse_expr("0.0")]])
+                                            [sympy.parsing.sympy_parser.parse_expr("0"), sympy.parsing.sympy_parser.parse_expr("0")]])
 
         V, W = sympy.symbols("V W")
         assert shape_sys.b_ == sympy.Matrix([[sympy.parsing.sympy_parser.parse_expr("500.0 * E_Ca * g_Ca / C_m + 1000.0 * E_L * g_L / C_m + 1000.0 * I_ext / C_m")],
-                                             [sympy.parsing.sympy_parser.parse_expr("0.0")]])
+                                             [sympy.parsing.sympy_parser.parse_expr("0")]])
         assert shape_sys.c_ == sympy.Matrix([[sympy.parsing.sympy_parser.parse_expr("500.0 * E_Ca * g_Ca * tanh(V / 15 + 1 / 15) / C_m - 1000.0 * V * W * g_K / C_m - 500.0 * V * g_Ca * tanh(V / 15 + 1 / 15) / C_m")],
                                              [sympy.parsing.sympy_parser.parse_expr("-200.0 * W * cosh(V / 60) + 100.0 * cosh(V / 60) * tanh(V / 30) + 100.0 * cosh(V / 60)")]])