New base class DIRK_IMEX (#112)

* New base class DIRK_IMEX

Created new base class DIRK_IMEX which now ARS_DIRK_IMEX inherits from.
Implicit tableau in ARS methods are now padded so sizing is uniform with
methods requiring an implicit first step.

Modified DIRKIMEXMethod to allow for implicit first stages. Also combined loops
for efficiency

Added additional ARS schemes for completeness

Added some SSPk schemes from Pareschi and Russo to demonstrate new functionality
in SSPK_DIRK_IMEX

Updated tests and demos accordingly

* Changes to resolve pull request issues

Fixed a typo in test suite

Updated version in docs

Changed __init__.py to match style

* Revisions to pull request

Dictionaries for SSPk and ARS schemes are now stored with their classes

Explicit parts of SSPk schemes are now available as a tableau via a new factory.
These tableau were also added to the testing in test_explicit.py

Testing in test_imex.py is now parameterized by the Butcher tableau

demos/monodomain/demo_monodomain_FHN_dirkimex.py.rst

@@ -30,7 +30,7 @@ We start with standard Firedrake/Irksome imports::
   from firedrake import (And, Constant, File, Function, FunctionSpace,
                          RectangleMesh, SpatialCoordinate, TestFunctions,
                          as_matrix, conditional, dx, grad, inner, split)
-  from irksome import Dt, MeshConstant, DIRK_IMEX, TimeStepper
+  from irksome import Dt, MeshConstant, TimeStepper, ARS_DIRK_IMEX
 
 And we set up the mesh and function space.::
   
@@ -75,7 +75,7 @@ This sets up the Butcher tableau.  Here, we use the DIRK-IMEX methods proposed
 by Ascher, Ruuth, and Spiteri in their 1997 Applied Numerical Mathematics paper.
 For this case, We use a four-stage method.::
   
-  butcher_tableau = DIRK_IMEX(4, 4, 3)
+  butcher_tableau = ARS_DIRK_IMEX(4, 4, 3)
   ns = butcher_tableau.num_stages
 
 To access an IMEX method, we need to separately specify the implicit and explicit parts of the operator.

irksome/__init__.py

@@ -9,6 +9,9 @@
 from .pep_explicit_rk import PEPRK          # noqa: F401
 from .deriv import Dt                       # noqa: F401
 from .dirk_imex_tableaux import DIRK_IMEX   # noqa: F401
+from .ars_dirk_imex_tableaux import ARS_DIRK_IMEX  # noqa: F401
+from .sspk_tableau import SSPK_DIRK_IMEX  # noqa: F401
+from .sspk_tableau import SSPButcherTableau  # noqa: F401
 from .dirk_stepper import DIRKTimeStepper   # noqa: F401
 from .getForm import getForm                # noqa: F401
 from .imex import RadauIIAIMEXMethod        # noqa: F401

irksome/ars_dirk_imex_tableaux.py

@@ -0,0 +1,111 @@
+import numpy as np
+
+from .dirk_imex_tableaux import DIRK_IMEX
+
+# Butcher tableau based on Ascher, Ruuth, and Spiteri Applied Numerical Mathematics 1997 (ARS)
+
+# ARS tableau assume a zero first column of the implicit A matrix, so only the lower right s x s
+# block is given for the implicit scheme. b and c are also of length s.
+
+# Butcher tableau for s = 1
+ars111A = np.array([[1.0]])
+ars111A_hat = np.array([[1.0]])
+ars111b = np.array([1.0])
+ars111b_hat = np.array([1.0, 0.0])
+ars111c = np.array([1.0])
+ars111c_hat = np.array([0.0, 1.0])
+
+ars121A = np.array([[1.0]])
+ars121A_hat = np.array([[1.0]])
+ars121b = np.array([1.0])
+ars121b_hat = np.array([0.0, 1.0])
+ars121c = np.array([1.0])
+ars121c_hat = np.array([0.0, 1.0])
+
+ars122A = np.array([[0.5]])
+ars122A_hat = np.array([[0.5]])
+ars122b = np.array([1.0])
+ars122b_hat = np.array([0.0, 1.0])
+ars122c = np.array([0.5])
+ars122c_hat = np.array([0.0, 0.5])
+
+# Butcher tableau for s = 2
+gamma233 = (3 + np.sqrt(3))/6
+ars233A = np.array([[gamma233, 0], [1 - 2*gamma233, gamma233]])
+ars233A_hat = np.array([[gamma233, 0], [gamma233 - 1, 2*(1 - gamma233)]])
+ars233b = np.array([0.5, 0.5])
+ars233b_hat = np.array([0, 0.5, 0.5])
+ars233c = np.array([gamma233, 1 - gamma233])
+ars233c_hat = np.array([0, gamma233, 1 - gamma233])
+
+gamma232 = (2 - np.sqrt(2)) / 2
+delta232 = -2 * np.sqrt(2) / 3
+ars232A = np.array([[gamma232, 0], [1 - gamma232, gamma232]])
+ars232A_hat = np.array([[gamma232, 0], [delta232, 1 - delta232]])
+ars232b = np.array([1 - gamma232, gamma232])
+ars232b_hat = np.array([0, 1 - gamma232, gamma232])
+ars232c = np.array([gamma232, 1.0])
+ars232c_hat = np.array([0, gamma232, 1.0])
+
+gamma222 = gamma232
+delta222 = 1 - 1/(2*gamma222)
+ars222A = np.array([[gamma222, 0], [1 - gamma222, gamma222]])
+ars222A_hat = np.array([[gamma222, 0], [delta222, 1 - delta222]])
+ars222b = np.array([1 - gamma222, gamma222])
+ars222b_hat = np.array([delta222, 1 - delta222, 0])
+ars222c = np.array([gamma222, 1.0])
+ars222c_hat = np.array([0, gamma222, 1.0])
+
+# Butcher tableau for s = 3
+ars343A = np.array([[0.4358665215, 0, 0], [0.2820667392, 0.4358665215, 0], [1.208496649, -0.644363171, 0.4358665215]])
+ars343A_hat = np.array([[0.4358665215, 0, 0], [0.3212788860, 0.3966543747, 0], [-0.105858296, 0.5529291479, 0.5529291479]])
+ars343b = np.array([1.208496649, -0.644363171, 0.4358665215])
+ars343b_hat = np.array([0, 1.208496649, -0.644363171, 0.4358665215])
+ars343c = np.array([0.4358665215, 0.7179332608, 1])
+ars343c_hat = np.array([0, 0.4358665215, 0.7179332608, 1.0])
+
+# Butcher tableau for s = 4
+ars443A = np.array([[1/2, 0, 0, 0],
+                    [1/6, 1/2, 0, 0],
+                    [-1/2, 1/2, 1/2, 0],
+                    [3/2, -3/2, 1/2, 1/2]])
+ars443A_hat = np.array([[1/2, 0, 0, 0],
+                        [11/18, 1/18, 0, 0],
+                        [5/6, -5/6, 1/2, 0],
+                        [1/4, 7/4, 3/4, -7/4]])
+ars443b = np.array([3/2, -3/2, 1/2, 1/2])
+ars443b_hat = np.array([1/4, 7/4, 3/4, -7/4, 0])
+ars443c = np.array([1/2, 2/3, 1/2, 1])
+ars443c_hat = np.array([0, 1/2, 2/3, 1/2, 1])
+
+ars_dict = {
+    (1, 1, 1): (ars111A, ars111b, ars111c, ars111A_hat, ars111b_hat, ars111c_hat),
+    (1, 2, 1): (ars121A, ars121b, ars121c, ars121A_hat, ars121b_hat, ars121c_hat),
+    (1, 2, 2): (ars122A, ars122b, ars122c, ars122A_hat, ars122b_hat, ars122c_hat),
+    (2, 2, 2): (ars222A, ars222b, ars222c, ars222A_hat, ars222b_hat, ars222c_hat),
+    (2, 3, 2): (ars232A, ars232b, ars232c, ars232A_hat, ars232b_hat, ars232c_hat),
+    (2, 3, 3): (ars233A, ars233b, ars233c, ars233A_hat, ars233b_hat, ars233c_hat),
+    (3, 4, 3): (ars343A, ars343b, ars343c, ars343A_hat, ars343b_hat, ars343c_hat),
+    (4, 4, 3): (ars443A, ars443b, ars443c, ars443A_hat, ars443b_hat, ars443c_hat)
+}
+
+
+class ARS_DIRK_IMEX(DIRK_IMEX):
+    """Class to generate IMEX tableaux based on Ascher, Ruuth, and Spiteri (ARS). It has members
+
+    :arg ns_imp: number of implicit stages
+    :arg ns_exp: number of explicit stages
+    :arg order: the (integer) former order of accuracy of the method
+    """
+    def __init__(self, ns_imp, ns_exp, order):
+        try:
+            A, b, c, A_hat, b_hat, c_hat = ars_dict[ns_imp, ns_exp, order]
+        except KeyError:
+            raise NotImplementedError("No ARS DIRK-IMEX method for that combination of implicit and explicit stages and order")
+
+        # Expand A, b, c with assumed zeros in ARS tableaux
+        A = self._pad_matrix(A, "lr")
+        b = np.append(np.zeros(1), b)
+        c = np.append(np.zeros(1), c)
+
+        super(ARS_DIRK_IMEX, self).__init__(A, b, c, A_hat, b_hat, c_hat, order)

irksome/dirk_imex_tableaux.py

@@ -1,71 +1,58 @@
 from .ButcherTableaux import ButcherTableau
 import numpy as np
 
-# For the implicit scheme, the full Butcher Table is given as A, b, c.
-
-# For the explicit scheme, the full b_hat and c_hat are given, but (to
-# avoid a lot of offset-by-ones in the code we store only the
-# lower-left ns x ns block of A_hat
-
-# IMEX Butcher tableau for 1 stage
-imex111A = np.array([[1.0]])
-imex111A_hat = np.array([[1.0]])
-imex111b = np.array([1.0])
-imex111b_hat = np.array([1.0, 0.0])
-imex111c = np.array([1.0])
-imex111c_hat = np.array([0.0, 1.0])
-
-
-# IMEX Butcher tableau for s = 2
-gamma = (2 - np.sqrt(2)) / 2
-delta = -2 * np.sqrt(2) / 3
-imex232A = np.array([[gamma, 0], [1 - gamma, gamma]])
-imex232A_hat = np.array([[gamma, 0], [delta, 1 - delta]])
-imex232b = np.array([1 - gamma, gamma])
-imex232b_hat = np.array([0, 1 - gamma, gamma])
-imex232c = np.array([gamma, 1.0])
-imex232c_hat = np.array([0, gamma, 1.0])
-
-# IMEX Butcher tableau for 3 stages
-imex343A = np.array([[0.4358665215, 0, 0], [0.2820667392, 0.4358665215, 0], [1.208496649, -0.644363171, 0.4358665215]])
-imex343A_hat = np.array([[0.4358665215, 0, 0], [0.3212788860, 0.3966543747, 0], [-0.105858296, 0.5529291479, 0.5529291479]])
-imex343b = np.array([1.208496649, -0.644363171, 0.4358665215])
-imex343b_hat = np.array([0, 1.208496649, -0.644363171, 0.4358665215])
-imex343c = np.array([0.4358665215, 0.7179332608, 1])
-imex343c_hat = np.array([0, 0.4358665215, 0.7179332608, 1.0])
-
-
-# IMEX Butcher tableau for 4 stages
-imex443A = np.array([[1/2, 0, 0, 0],
-                     [1/6, 1/2, 0, 0],
-                     [-1/2, 1/2, 1/2, 0],
-                     [3/2, -3/2, 1/2, 1/2]])
-imex443A_hat = np.array([[1/2, 0, 0, 0],
-                         [11/18, 1/18, 0, 0],
-                         [5/6, -5/6, 1/2, 0],
-                         [1/4, 7/4, 3/4, -7/4]])
-imex443b = np.array([3/2, -3/2, 1/2, 1/2])
-imex443b_hat = np.array([1/4, 7/4, 3/4, -7/4, 0])
-imex443c = np.array([1/2, 2/3, 1/2, 1])
-imex443c_hat = np.array([0, 1/2, 2/3, 1/2, 1])
-
-dirk_imex_dict = {
-    (1, 1, 1): (imex111A, imex111b, imex111c, imex111A_hat, imex111b_hat, imex111c_hat),
-    (2, 3, 2): (imex232A, imex232b, imex232c, imex232A_hat, imex232b_hat, imex232c_hat),
-    (3, 4, 3): (imex343A, imex343b, imex343c, imex343A_hat, imex343b_hat, imex343c_hat),
-    (4, 4, 3): (imex443A, imex443b, imex443c, imex443A_hat, imex443b_hat, imex443c_hat)
-}
-
 
 class DIRK_IMEX(ButcherTableau):
-    def __init__(self, ns_imp, ns_exp, order):
-        try:
-            A, b, c, A_hat, b_hat, c_hat = dirk_imex_dict[ns_imp, ns_exp, order]
-        except KeyError:
-            raise NotImplementedError("No DIRK-IMEX method for that combination of implicit and explicit stages and order")
-        self.order = order
-        super(DIRK_IMEX, self).__init__(A, b, None, c, order, None, None)
-        self.A_hat = A_hat
+    """Top-level class representing a pair of Butcher tableau encoding an implicit-explicit
+    additive Runge-Kutta method. Since the explicit Butcher matrix is strictly lower triangular,
+    only the lower-left (ns - 1)x(ns - 1) block is given. However, the full b_hat and c_hat are
+    given. It has members
+
+    :arg A: a 2d array containing the implicit Butcher matrix
+    :arg b: a 1d array giving weights assigned to each implicit stage when
+            computing the solution at time n+1.
+    :arg c: a 1d array containing weights at which time-dependent
+            implicit terms are evaluated.
+    :arg A_hat: a 2d array containing the explicit Butcher matrix (lower-left block only)
+    :arg b_hat: a 1d array giving weights assigned to each explicit stage when
+            computing the solution at time n+1.
+    :arg c_hat: a 1d array containing weights at which time-dependent
+            explicit terms are evaluated.
+    :arg order: the (integer) formal order of accuracy of the method
+    """
+
+    def __init__(self, A: np.ndarray, b: np.ndarray, c: np.ndarray, A_hat: np.ndarray,
+                 b_hat: np.ndarray, c_hat: np.ndarray, order: int = None):
+
+        # Number of stages
+        ns = A.shape[0]
+        assert ns == A.shape[1], "A must be square"
+        assert A_hat.shape == (ns - 1, ns - 1), "A_hat must have one fewer row and column than A"
+        assert ns == len(b) == len(b_hat), \
+            "b and b_hat must have the same length as the number of stages"
+        assert ns == len(c) == len(c_hat), \
+            "c and c_hat must have the same length as the number of stages"
+
+        super().__init__(A, b, None, c, order, None, None)
+        self.A_hat = self._pad_matrix(A_hat, "ll")
         self.b_hat = b_hat
         self.c_hat = c_hat
         self.is_dirk_imex = True  # Mark this as a DIRK-IMEX scheme
+
+    @staticmethod
+    def _pad_matrix(mat: np.ndarray, loc: str):
+        """Zero pads a matrix"""
+        n = mat.shape[0]
+        assert n == mat.shape[1], "Matrix must be square"
+        padded = np.zeros((n+1, n+1), dtype=mat.dtype)
+
+        if loc == "ll":
+            # Lower left corner
+            padded[1:, :-1] = mat
+        elif loc == "lr":
+            # Lower right corner
+            padded[1:, 1:] = mat
+        else:
+            raise ValueError("Location must be ll (lower left) or lr (lower right)")
+
+        return padded