dsl: Introduce abstractions for multi-stage time integrators

devito/types/multistage.py

@@ -3,6 +3,8 @@
 from devito.symbolics import uxreplace
 from numpy import number
 from devito.types.array import Array
+from devito.types.dense import Function
+from devito.types.constant import Constant
 from types import MappingProxyType
 
 method_registry = {}
@@ -51,7 +53,7 @@ class MultiStage(Eq):
     of update expressions for each stage in the integration process.
     """
 
-    def __new__(cls, lhs, rhs, **kwargs):
+    def __new__(cls, lhs, rhs, source=None, degree=6, **kwargs):
         if not isinstance(lhs, list):
             lhs=[lhs]
             rhs=[rhs]
@@ -61,6 +63,8 @@ def __new__(cls, lhs, rhs, **kwargs):
         obj._eq = [Eq(lhs[i], rhs[i]) for i in range(len(lhs))]
         obj._lhs = lhs
         obj._rhs = rhs
+        obj._deg = degree
+        obj._src = source
 
         return obj
 
@@ -79,6 +83,16 @@ def rhs(self):
         """Return list of right-hand sides."""
         return self._rhs
 
+    @property
+    def deg(self):
+        """Return list of right-hand sides."""
+        return self._deg
+
+    @property
+    def src(self):
+        """Return list of right-hand sides."""
+        return self._src
+
     def _evaluate(self, **kwargs):
         raise NotImplementedError(
             f"_evaluate() must be implemented in the subclass {self.__class__.__name__}")
@@ -115,7 +129,9 @@ class RK(MultiStage):
         Number of stages in the RK method, inferred from `b`.
     """
 
-    def __init__(self, a: list[list[float | number]], b: list[float | number], c: list[float | number], lhs, rhs, **kwargs) -> None:
+    CoeffsBC = list[float | number]
+    CoeffsA = list[CoeffsBC]
+    def __init__(self, a: CoeffsA, b: CoeffsBC, c: CoeffsBC, lhs, rhs, **kwargs) -> None:
         self.a, self.b, self.c = a, b, c
 
     @property
@@ -132,19 +148,18 @@ def _evaluate(self, **kwargs):
 
         Returns
         -------
-        list of Eq
+        list of Devito Eq objects
             A list of SymPy Eq objects representing:
             - `s` stage equations of the form `k_i = rhs evaluated at intermediate state`
             - 1 final update equation of the form `u.forward = u + dt * sum(b_i * k_i)`
         """
         n_eq=len(self.eq)
         u = [i.function for i in self.lhs]
-        grid = [u[i].grid for i in range(n_eq)]
-        t = grid[0].time_dim
+        t = u[0].grid.time_dim
         dt = t.spacing
 
         # Create temporary Functions to hold each stage
-        k = [[Array(name=f'{kwargs.get('sregistry').make_name(prefix='k')}', dimensions=grid[j].dimensions, grid=grid[j], dtype=u[j].dtype) for i in range(self.s)]
+        k = [[Array(name=f'{kwargs.get('sregistry').make_name(prefix='k')}', dimensions=u[j].grid.dimensions, grid=u[j].grid, dtype=u[j].dtype) for i in range(self.s)]
              for j in range(n_eq)]
 
         stage_eqs = []
@@ -214,8 +229,8 @@ class RK32(RK):
     b = [0, 0, 1]
     c = [0, 1/2, 1/2]
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(a=self.a, b=self.b, c=self.c, **kwargs)
+    def __init__(self, lhs, rhs, **kwargs):
+        super().__init__(a=self.a, b=self.b, c=self.c, lhs=lhs, rhs=rhs, **kwargs)
 
 
 @register_method
@@ -249,12 +264,12 @@ class RK97(RK):
          5963949/25894400, 50000000000/599799373173, 28487/712800]
     c = [0, 4/63, 2/21, 1/7, 7/17, 13/24, 7/9, 91/100, 1]
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(a=self.a, b=self.b, c=self.c, **kwargs)
+    def __init__(self, lhs, rhs, **kwargs):
+        super().__init__(a=self.a, b=self.b, c=self.c, lhs=lhs, rhs=rhs, **kwargs)
 
 
 @register_method
-class HORK(MultiStage):
+class HORK_EXP(MultiStage):
     # In construction
     """
     n stages Runge-Kutta (HORK) time integration method.
@@ -271,8 +286,19 @@ class HORK(MultiStage):
         Time positions of intermediate stages.
     """
 
+    def source_derivatives(self, src_index, t, **kwargs):
+
+        # Compute the base wavelet function
+        f_deriv = [[self.src[i][1] for i in range(len(self.src))]]
+
+        # Compute derivatives up to order p
+        for _ in range(self.deg - 1):
+            f_deriv.append([f_deriv[-1][i].diff(t) for i in range(len(src_index))])
+
+        f_deriv.reverse()
+        return f_deriv
 
-    def ssprk_alpha(mu=1, **kwargs):
+    def ssprk_alpha(self, mu=1):
         """
         Computes the coefficients for the Strong Stability Preserving Runge-Kutta (SSPRK) method.
 
@@ -287,18 +313,33 @@ def ssprk_alpha(mu=1, **kwargs):
         numpy.ndarray
             Array of SSPRK coefficients.
         """
-        degree=kwargs.get('degree')
 
-        alpha = [0]*degree
+        alpha = [0]*self.deg
         alpha[0] = 1.0  # Initial coefficient
 
-        for i in range(1, degree):
-            alpha[i] = 1 / (mu * (i + 1)) * alpha[i - 1]
-            alpha[1:i] = 1 / (mu * list(range(1, i))) * alpha[:i - 1]
+        for i in range(1, self.deg):
+            alpha[i] = 1/(mu*(i+1))*alpha[i-1]
+            alpha[1:i] = [1/(mu*j)*alpha[j-1] for j in range(1,i)]
             alpha[0] = 1 - sum(alpha[1:i + 1])
 
         return alpha
 
+
+    def source_inclusion(self, u, k, src_index, src_deriv, e_p, t, dt, mu, n_eq):
+
+        src_lhs = [uxreplace(self.rhs[i], {u[m]: k[m] for m in range(n_eq)}) for i in range(n_eq)]
+
+        p = len(src_deriv)
+
+        for i in range(p):
+            if e_p[i] != 0:
+                for j in range(len(src_index)):
+                    src_lhs[src_index[j]] += self.src[j][0]*src_deriv[i][j].subs({t: t * dt})*e_p[i]
+        e_p = [e_p[i]+mu*dt*e_p[i + 1] for i in range(p - 1)]+[e_p[-1]]
+
+        return src_lhs, e_p
+
+
     def _evaluate(self, **kwargs):
         """
         Generate the stage-wise equations for a Runge-Kutta time integration method.
@@ -315,66 +356,52 @@ def _evaluate(self, **kwargs):
             - 1 final update equation of the form `u.forward = u + dt * sum(b_i * k_i)`
         """
 
-        u = self.lhs.function
-        rhs = self.rhs
-        grid = u.grid
-        t = grid.time_dim
+        n_eq=len(self.eq)
+        u = [i.function for i in self.lhs]
+        t = u[0].grid.time_dim
         dt = t.spacing
 
-        an_eq = range(len(U0))
+        # Create a temporary Array for each variable to save the time stages
+        # k = [Array(name=f'{kwargs.get('sregistry').make_name(prefix='k')}', dimensions=u[i].grid.dimensions, grid=u[i].grid, dtype=u[i].dtype) for i in range(n_eq)]
+        k = [Function(name=f'{kwargs.get('sregistry').make_name(prefix='k')}', grid=u[i].grid, space_order=2, time_order=1, dtype=u[i].dtype) for i in range(n_eq)]
+        k_old = [Function(name=f'{kwargs.get('sregistry').make_name(prefix='k')}', grid=u[i].grid, space_order=2, time_order=1, dtype=u[i].dtype) for i in range(n_eq)]
 
         # Compute SSPRK coefficients
-        alpha = np.array(ssprk_alpha(mu, degree), dtype=np.float64)
+        mu = 1
+        alpha = self.ssprk_alpha(mu=mu)
 
         # Initialize symbolic differentiation for source terms
-        t_var = sym.Symbol('t_var')
-        src_deriv = aux_fun.derivates_f(degree, f0)
+        src_index_map={val: i for i, val in enumerate(u)}
+        src_index = [src_index_map[val] for val in [self.src[i][2] for i in range(len(self.src))]]
+        src_deriv = self.source_derivatives(src_index, t, **kwargs)
 
         # Expansion coefficients for stability control
-        e_p = [0] * degree
+        e_p = [0] * self.deg
+        eta = 1
         e_p[-1] = 1 / eta
 
-        # Initialize approximation and auxiliary variable
-        approx = [U0[i] * alpha[0] for i in n_eq]
-        aux = U0
-
-        # Perform Runge-Kutta steps
-        for i in range(1, degree - 1):
-            system_op, e_p = sys_op_extended(aux, x, y, z, param_fun, system, fd_order, src_spat, src_deriv, t, dt, t_var, e_p)
-            aux = [aux[j] + mu * dt * system_op[j] for j in n_eq]
-            approx = [approx[j] + aux[j] * alpha[i] for j in n_eq]
-
-        # Final Runge-Kutta updates
-        system_op, e_p = sys_op_extended(aux, x, y, z, param_fun, system, fd_order, src_spat, src_deriv, t, dt, t_var, e_p)
-        aux = [aux[i] + mu * dt * system_op[i] for i in n_eq]
-        system_op, e_p = sys_op_extended(aux, x, y, z, param_fun, system, fd_order, src_spat, src_deriv, t, dt, t_var, e_p)
-        aux = [aux[i] + mu * dt * system_op[i] for i in n_eq]
-
-        # Compute final approximation
-        approx = [approx[i] + aux[i] * alpha[degree - 1] for i in n_eq]
-
-        # Generate final PDE system
-        return [dv.Eq(U0[i].forward, approx[i]) for i in n_eq]
 
-        # Create temporary Functions to hold each stage
-        # k = [Array(name=f'{kwargs.get('sregistry').make_name(prefix='k')}', dimensions=grid.shape, grid=grid, dtype=u.dtype) for i in range(self.s)]  # Trying Array
-        k = [Function(name=f'{kwargs.get('sregistry').make_name(prefix='k')}', grid=grid, space_order=u.space_order, dtype=u.dtype)
-             for i in range(self.s)]
-
-        stage_eqs = []
+        stage_eqs = [Eq(k[i], u[i]) for i in range(n_eq)]
+        [stage_eqs.append(Eq(u[i].forward, u[i]*alpha[0])) for i in range(n_eq)]
 
         # Build each stage
-        for i in range(self.s):
-            u_temp = u + dt * sum(aij * kj for aij, kj in zip(self.a[i][:i], k[:i]))
-            t_shift = t + self.c[i] * dt
+        for i in range(1, self.deg-1):
+            [stage_eqs.append(Eq(k_old[j], k[j])) for j in range(n_eq)]
+            src_lhs, e_p = self.source_inclusion(u, k_old, src_index, src_deriv, e_p, t, dt, mu, n_eq)
+            [stage_eqs.append(Eq(k[j], k_old[j]+mu*dt*src_lhs[j])) for j in range(n_eq)]
+            [stage_eqs.append(Eq(u[j].forward, u[j].forward+k[j]*alpha[i])) for j in range(n_eq)]
 
-            # Evaluate RHS at intermediate value
-            stage_rhs = uxreplace(rhs, {u: u_temp, t: t_shift})
-            stage_eqs.append(Eq(k[i], stage_rhs))
+        # Final Runge-Kutta updates
+        [stage_eqs.append(Eq(k_old[j], k[j])) for j in range(n_eq)]
+        src_lhs, e_p = self.source_inclusion(u, k_old, src_index, src_deriv, e_p, t, dt, mu, n_eq)
+        [stage_eqs.append(Eq(k[j], k_old[j]+mu*dt*src_lhs[j])) for j in range(n_eq)]
 
-        # Final update: u.forward = u + dt * sum(b_i * k_i)
-        u_next = u + dt * sum(bi * ki for bi, ki in zip(self.b, k))
-        stage_eqs.append(Eq(u.forward, u_next))
+        [stage_eqs.append(Eq(k_old[j], k[j])) for j in range(n_eq)]
+        src_lhs, _ = self.source_inclusion(u, k_old, src_index, src_deriv, e_p, t, dt, mu, n_eq)
+        [stage_eqs.append(Eq(k[j], k_old[j]+mu*dt*src_lhs[j])) for j in range(n_eq)]
+
+        # Compute final approximation
+        [stage_eqs.append(Eq(u[j].forward, u[j].forward+k[j]*alpha[self.deg-1])) for j in range(n_eq)]
 
         return stage_eqs