dsl: Introduce abstractions for multi-stage time integrators

devito/ir/equations/algorithms.py

@@ -6,14 +6,15 @@
 from devito.tools import (Ordering, as_tuple, flatten, filter_sorted, filter_ordered,
                           frozendict)
 from devito.types import (Dimension, Eq, IgnoreDimSort, SubDimension,
-                          ConditionalDimension)
+                          ConditionalDimension, MultiStage)
 from devito.types.array import Array
 from devito.types.basic import AbstractFunction
 from devito.types.dimension import MultiSubDimension, Thickness
 from devito.data.allocators import DataReference
 from devito.logger import warning
 
-__all__ = ['dimension_sort', 'lower_exprs', 'concretize_subdims']
+
+__all__ = ['dimension_sort', 'lower_multistage', 'lower_exprs', 'concretize_subdims']
 
 
 def dimension_sort(expr):
@@ -95,6 +96,39 @@ def handle_indexed(indexed):
     return ordering
 
 
+def lower_multistage(expressions, **kwargs):
+    """
+    Separating the multi-stage time-integrator scheme in stages:
+        * If the object is MultiStage, it creates the stages of the method.
+    """
+    return _lower_multistage(expressions, **kwargs)
+
+
+@singledispatch
+def _lower_multistage(expr, **kwargs):
+    """
+    Default handler for expressions that are not MultiStage.
+    Simply return them in a list.
+    """
+    return [expr]
+
+
+@_lower_multistage.register(MultiStage)
+def _(expr, **kwargs):
+    """
+    Specialized handler for MultiStage expressions.
+    """
+    return expr._evaluate(**kwargs)
+
+
+@_lower_multistage.register(Iterable)
+def _(exprs, **kwargs):
+    """
+    Handle iterables of expressions.
+    """
+    return sum([_lower_multistage(expr, **kwargs) for expr in exprs], [])
+
+
 def lower_exprs(expressions, subs=None, **kwargs):
     """
     Lowering an expression consists of the following passes:

devito/operations/solve.py

@@ -7,6 +7,8 @@
 from devito.finite_differences.derivative import Derivative
 from devito.tools import as_tuple
 
+from devito.types.multistage import resolve_method
+
 __all__ = ['solve', 'linsolve']
 
 
@@ -56,9 +58,12 @@ def solve(eq, target, **kwargs):
 
     # We need to rebuild the vector/tensor as sympy.solve outputs a tuple of solutions
     if len(sols) > 1:
-        return target.new_from_mat(sols)
+        sols_temp = target.new_from_mat(sols)
     else:
-        return sols[0]
+        sols_temp = sols[0]
+
+    method = kwargs.get("method", None)
+    return sols_temp if method is None else resolve_method(method)(target, sols_temp)
 
 
 def linsolve(expr, target, **kwargs):

devito/operator/operator.py

@@ -15,7 +15,7 @@
                                InvalidOperator)
 from devito.logger import (debug, info, perf, warning, is_log_enabled_for,
                            switch_log_level)
-from devito.ir.equations import LoweredEq, lower_exprs, concretize_subdims
+from devito.ir.equations import LoweredEq, lower_multistage, lower_exprs, concretize_subdims
 from devito.ir.clusters import ClusterGroup, clusterize
 from devito.ir.iet import (Callable, CInterface, EntryFunction, FindSymbols,
                            MetaCall, derive_parameters, iet_build)
@@ -36,7 +36,6 @@
                           disk_layer)
 from devito.types.dimension import Thickness
 
-
 __all__ = ['Operator']
 
 
@@ -327,6 +326,8 @@ def _lower_exprs(cls, expressions, **kwargs):
             * Apply substitution rules;
             * Shift indices for domain alignment.
         """
+        expressions = lower_multistage(expressions, **kwargs)
+
         expand = kwargs['options'].get('expand', True)
 
         # Specialization is performed on unevaluated expressions

devito/types/__init__.py

@@ -22,3 +22,5 @@
 from .relational import *  # noqa
 from .sparse import *  # noqa
 from .tensor import *  # noqa
+
+from .multistage import *  # noqa

devito/types/multistage.py

@@ -0,0 +1,179 @@
+from .equation import Eq
+from .dense import Function
+from devito.symbolics import uxreplace
+
+# from devito.ir.support import SymbolRegistry
+
+from .array import Array  # Trying Array
+
+
+method_registry = {}
+
+def register_method(cls):
+    method_registry[cls.__name__] = cls
+    return cls
+
+
+def resolve_method(method):
+    try:
+        return method_registry[method]
+    except KeyError:
+        raise ValueError(f"The time integrator '{method}' is not implemented.")
+
+
+class MultiStage(Eq):
+    """
+       Abstract base class for multi-stage time integration methods
+       (e.g., Runge-Kutta schemes) in Devito.
+
+       This class represents a symbolic equation of the form `target = rhs`
+       and provides a mechanism to associate it with a time integration
+       scheme. The specific integration behavior must be implemented by
+       subclasses via the `_evaluate` method.
+
+       Parameters
+       ----------
+       lhs : expr-like
+           The left-hand side of the equation, typically a time-updated Function
+           (e.g., `u.forward`).
+       rhs : expr-like, optional
+           The right-hand side of the equation to integrate. Defaults to 0.
+       subdomain : SubDomain, optional
+           A subdomain over which the equation applies.
+       coefficients : dict, optional
+           Optional dictionary of symbolic coefficients for the integration.
+       implicit_dims : tuple, optional
+           Additional dimensions that should be treated implicitly in the equation.
+       **kwargs : dict
+           Additional keyword arguments, such as time integration method selection.
+
+       Notes
+       -----
+       Subclasses must override the `_evaluate()` method to return a sequence
+       of update expressions for each stage in the integration process.
+       """
+
+    def __new__(cls, lhs, rhs=0, subdomain=None, coefficients=None, implicit_dims=None, **kwargs):
+        return super().__new__(cls, lhs, rhs=rhs, subdomain=subdomain, coefficients=coefficients, implicit_dims=implicit_dims, **kwargs)
+
+    def _evaluate(self, **kwargs):
+        raise NotImplementedError(
+            f"_evaluate() must be implemented in the subclass {self.__class__.__name__}")
+
+
+class RK(MultiStage):
+    """
+    Base class for explicit Runge-Kutta (RK) time integration methods defined
+    via a Butcher tableau.
+
+    This class handles the general structure of RK schemes by using
+    the Butcher coefficients (`a`, `b`, `c`) to expand a single equation into
+    a series of intermediate stages followed by a final update. Subclasses
+    must define `a`, `b`, and `c` as class attributes.
+
+    Parameters
+    ----------
+    a : list of list of float
+        The coefficient matrix representing stage dependencies.
+    b : list of float
+        The weights for the final combination step.
+    c : list of float
+        The time shifts for each intermediate stage (often the row sums of `a`).
+
+    Attributes
+    ----------
+    a : list[list[float]]
+        Butcher tableau `a` coefficients (stage coupling).
+    b : list[float]
+        Butcher tableau `b` coefficients (weights for combining stages).
+    c : list[float]
+        Butcher tableau `c` coefficients (stage time positions).
+    s : int
+        Number of stages in the RK method, inferred from `b`.
+    """
+
+    def __init__(self, a=None, b=None, c=None, **kwargs):
+        self.a, self.b, self.c = self._validate(a, b, c)
+
+    def _validate(self, a, b, c):
+        if a is None or b is None or c is None:
+            raise ValueError("RK subclass must define class attributes of the Butcher's array a, b, and c")
+        return a, b, c
+
+    @property
+    def s(self):
+        return len(self.b)
+
+    def _evaluate(self, **kwargs):
+        """
+        Generate the stage-wise equations for a Runge-Kutta time integration method.
+
+        This method takes a single equation of the form `Eq(u.forward, rhs)` and
+        expands it into a sequence of intermediate stage evaluations and a final
+        update equation according to the Runge-Kutta coefficients `a`, `b`, and `c`.
+
+        Returns
+        -------
+        list of Eq
+            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)`
+        """
+
+        u = self.lhs.function
+        rhs = self.rhs
+        grid = u.grid
+        t = 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.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 = []
+
+        # 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
+
+            # Evaluate RHS at intermediate value
+            stage_rhs = uxreplace(rhs, {u: u_temp, t: t_shift})
+            stage_eqs.append(Eq(k[i], stage_rhs))
+
+        # 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))
+
+        return stage_eqs
+
+
+@register_method
+class RK44(RK):
+    """
+    Classic 4th-order Runge-Kutta (RK4) time integration method.
+
+    This class implements the classic explicit Runge-Kutta method of order 4 (RK44).
+    It uses four intermediate stages and specific Butcher coefficients to achieve
+    high accuracy while remaining explicit.
+
+    Attributes
+    ----------
+    a : list[list[float]]
+        Coefficients of the `a` matrix for intermediate stage coupling.
+    b : list[float]
+        Weights for final combination.
+    c : list[float]
+        Time positions of intermediate stages.
+    """
+    a = [[0, 0, 0, 0],
+         [1/2, 0, 0, 0],
+         [0, 1/2, 0, 0],
+         [0, 0, 1, 0]]
+    b = [1/6, 1/3, 1/3, 1/6]
+    c = [0, 1/2, 1/2, 1]
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(a=self.a, b=self.b, c=self.c)
+