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,42 @@ 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.
+    """
+    lowered = []
+    for expr in exprs:
+        lowered.extend(_lower_multistage(expr, **kwargs))
+    return lowered
+
+
 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,191 @@
+from .equation import Eq
+from .dense import Function
+from devito.symbolics import uxreplace
+
+from functools import cached_property
+
+# 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 wraps a symbolic equation of the form `target = rhs` and
+    provides a mechanism to associate a time integration scheme via the
+    `method` argument. Subclasses must implement the `_evaluate` method to
+    generate stage-wise update expressions.
+
+    Parameters
+    ----------
+    rhs : expr-like
+        The right-hand side of the equation to integrate.
+    target : Function
+        The time-updated symbol on the left-hand side, e.g., `u` or `u.forward`.
+    method : str or None
+        A string identifying the time integration method (e.g., 'RK44'),
+        which must correspond to a class defined in the global scope and
+        implementing `_evaluate`. If None, no method is applied.
+
+    Attributes
+    ----------
+    eq : Eq
+        The symbolic equation `target = rhs`.
+    method : class
+        The integration method class resolved from the `method` string.
+    """
+
+    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, **kwargs):
+        self.a, self.b, self.c = self._validate(**kwargs)
+
+    def _validate(self, **kwargs):
+        a = kwargs.get('a', None)
+        b = kwargs.get('b', None)
+        c = kwargs.get('c', None)
+        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
+
+    @cached_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`.
+
+        Parameters
+        ----------
+        eq_num : int, optional
+            An identifier index used to uniquely name the intermediate stage variables
+            (`k{eq_num}i`) in case of multiple equations being expanded.
+
+        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)`
+        """
+        stage_id = kwargs.get('sregistry').make_name(prefix='k')
+
+        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'{stage_id}{i}', dimensions=grid.shape, grid=grid, dtype=u.dtype) for i in range(self.s)]  # Trying Array
+        k = [Function(name=f'{stage_id}{i}', 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):
+        kwargs.setdefault('a', self.a)
+        kwargs.setdefault('b', self.b)
+        kwargs.setdefault('c', self.c)
+        super().__init__(**kwargs)
+