dsl: Introduce abstractions for multi-stage time integrators #2599

fernanvr · 2025-05-05T16:59:28Z

Implementation of the ideas of @mloubout, as discussed in the Slack channel #timestepping.

It were created two new classes:

MultiStage: Each instance represents a combination of a PDE and its associated time integrator.
RK: Instances encapsulate a Butcher Tableau. The class methods define specific Runge-Kutta schemes, and it includes logic to expand a single PDE into multiple stage equations according to the RK method. It only contains the classic RK44, of fourth order and four stages, as a prove of concept.

To integrate the MultiStage into Devito’s pipeline, we modified the _lower(...) function in operator.py by adding the line:
expressions = cls._lower_multistage(expressions, **kwargs),
and modified _sanitize_exprs(cls, expressions, **kwargs) to recognize MultiStage instances.

We also created a new function:
_lower_multistage(),
in the same file. This function handles the expansion of MultiStage instances into their corresponding stage equations.

FabioLuporini

Hi, many thanks for this!

As you may see from the "Files changed" tab (https://github.com/devitocodes/devito/pull/2599/files), there are several new files that shouldn't be there. The whole .idea folder has probably been pushed erroneously. You will have to rebase your git history and erase it.

The new file new_classes should have a proper name and be placed in a different folder, @EdCaunt may have some ideas about that.

The test you have added has to be placed within a function def test_.... otherwise it won't be picked up by pytest. And we are going to need more tests, one isn't enough.
Also, a test is a test, it shouldn't plot anything. Instead, the numerical output should be compared to either an analytical solution or to one or more precomputed ("by hand") norms.

FabioLuporini · 2025-05-06T08:00:30Z

devito/operator/operator.py

@@ -271,6 +273,9 @@ def _lower(cls, expressions, **kwargs):
        # expression for which a partial or complete lowering is desired
        kwargs['rcompile'] = cls._rcompile_wrapper(**kwargs)

+        # [MultiStage] -> [Eqs]
+        expressions = cls._lower_multistage(expressions, **kwargs)


this should be called inside _lower_dsl, which in turn is called within _lower_exprs

I changed it to be called within _lower_exprs, but I couldn't find the _lower_dsl function...

EdCaunt

Lots of comments, but looks like a good first attempt at multistage timestepping!

devito/operator/new_classes.py

devito/operator/operator.py

devito/operator/new_classes.py

tests/test_multistage.py

EdCaunt · 2025-05-06T13:10:15Z

tests/test_multistage.py

+# Set logging level for debugging
+dv.configuration['log-level'] = 'DEBUG'
+
+# Parameters


Wants to use pytest going forward. If you look at the existing tests, you will see that a lot of them are very granular, which aids bugfixing and test coverage.

It's worth noting that tests (especially the very elementary ones) can be physically nonsensical, in order to check that some aspect of the API or compiler functions as intended. For example, it would be useful to test that a range of Butcher tableaus get correctly assembled into the corresponding timestepping schemes, even if those timestepping schemes themselves are meaningless.

Examples of simplifications used for unit tests include using 1D where possible and using trivial Eqs (i.e. Eq(f, 1))

Other tests that we include for essentially every new class include ensuring that it can be pickled (see test_pickle.py) and that it can be rebuilt correctly using its _rebuild method (also used for pickling). This is what the __rargs__ and __rkwargs__ attached to various classes are for.

I’ve been working in that direction, but I think it’s still a bit rough

ZoeLeibowitz · 2025-05-13T10:19:17Z

tests/test_multistage.py

+# Construct and run operator
+op = dv.Operator(pdes + [rec], subs=grid.spacing_map)
+op(dt=dt0, time=nt)
+


In the generated C code, I noticed the following:

k00[x + 2][y + 2] = v[t0][x + 2][y + 2]; k01[x + 2][y + 2] = v[t0][x + 2][y + 2]; k02[x + 2][y + 2] = v[t0][x + 2][y + 2]; k03[x + 2][y + 2] = v[t0][x + 2][y + 2];

Is this intended? It doesn't seem optimal since they are all assigned the same value

ZoeLeibowitz · 2025-05-13T10:21:38Z

tests/test_multistage.py

+src_temporal = (1 - 2 * (np.pi * f0 * (t * dt - 1/f0))**2) * sym.exp(-(np.pi * f0 * (t * dt - 1/f0))**2)
+
+# PDE system (2D acoustic)
+system_eqs = [U[1],


Maybe a simple heat equation would be a good example/test?

I incorporated for the tests

mloubout

Thanks for the contribution @fernanvr I think it's a great start!

I'm gonna leave some generic comments to start with and will refine as it gets updated. So here is a few main comments.

As @EdCaunt , the MultiStage object should inherit from Eq as it represents an equation, or in this case a set of equations. It should also be moved to devito/types/equations.py or make a new devito/types/multistage.py file.
Instead of that generic class with a self-referencing constructor, you should instead try to have a top level abstract class then specific cases for it. I.e something like:

class MultiStage(Eq):
    def __new__(cls, eq, method, target):

where target is the target field the same way as solve so that you can easily have forward and backward equations.

Then for RK you would have

class RK(MultiStage):
    a
    b
    c

a,b,c need to as class attributes so you can easily create each case as their own class so you can differentiate them if needed. e.g. RK4, RK3, etc. This way you can also have a MultiStage class that is not RK but something else, like Adams-Bashforth or whatever.

class RK4(RK)
    a = [0, 1/2, 1/2, 1]
    b = [1/6, 1/3, 1/3, 1/6]
    c = [0, 1/2, 1/2, 1]

You can also set it up as you did but the

   @classmethod
    def RK44(cls):

needs to be moved out of the class and made into a plain function.

The MultiStage should implement .evaluate(expand=False) then you won't need the lower_multistage as the evaluate will directly return the list of equations needed.
THe "test" you have is more of an example and should be made into a tutorial. The tests you be testing each piece separately such as the plain construction of each object and the results of the evaluation and lowering of it.
Optionally. I would abstract it away and directly allow something like eq= solve(pde, u.forward, method='RK4') which would automatically create the right MultiStage object. That would allow for the user to easily switch between different time steppers.

Feel free to ask questions or discuss on slack as well

…tegrator the merge is necessary to maintaing the development of multistages

fernanvr · 2025-06-13T22:17:11Z

Thank you all for the comprehensive review comments. They have been incredibly helpful in improving my coding skills and deepening my understanding of some of Devito’s workflows. Overall, I believe I have incorporated most of the suggestions to enhance the code. However, there are still a few points where I could use some additional guidance.

EdCaunt

More comments, but generally looking better

EdCaunt · 2025-06-16T12:24:42Z

devito/ir/equations/algorithms.py

+        * If the object is MultiStage, it creates the stages of the method.
+    """
+    lowered = []
+    for i, eq in enumerate(as_tuple(expressions)):


Rather than having the the as_tuple here, why not have a dispatch for _lower_multistage that dispatches on iterable types as per _concretize_subdims?

done, I think..

EdCaunt · 2025-06-16T12:27:42Z

devito/ir/equations/algorithms.py

+    return [expr]
+
+
+@_lower_multistage.register


I would personally tweak this for consistency with other uses of singledispatch in the codebase

EdCaunt · 2025-06-16T12:29:30Z

devito/operations/solve.py

@@ -15,7 +17,7 @@ class SolveError(Exception):
    pass


-def solve(eq, target, **kwargs):
+def solve(eq, target, method = None, eq_num = 0, **kwargs):


Kwargs should not have spaces around them. Furthermore, can method and eq_num simply be folded into **kwargs?

EdCaunt · 2025-06-16T12:29:55Z

devito/operations/solve.py

@@ -56,9 +58,15 @@ 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)


Should have whitespace around operator. Same below

EdCaunt · 2025-06-16T12:40:34Z

devito/operations/solve.py

+        sols_temp=sols[0]
+
+    if method is not None:
+        method_cls = MultiStage._resolve_method(method)


To me, this implies thatMultiStage._resolve_method should not be a class method. Perhaps the method kwarg should be dropped, and instead you should have something like:

method_cls = eq._resolve_method() # Possibly make this a property? if method_cls is None: return sols_temp return method_cls(sols_temp, target)._evaluate(eq_num=eq_num)

or even just

return eq._resolve_method(sols_temp, target)._evaluate(eq_num=eq_num)

where _resolve_method is some abstract method of Eq and its subclasses which defaults to a no-op or similar.

As a side note, why is the eq_num kwarg required?

EdCaunt · 2025-06-16T12:59:15Z

devito/types/multistage.py

+
+        # Create temporary Functions to hold each stage
+        # k = [Array(name=f'k{eq_num}{i}', dimensions=grid.shape, grid=grid, dtype=u.dtype) for i in range(self.s)]  # Trying Array
+        k = [Function(name=f'k{eq_num}{i}', grid=grid, space_order=u.space_order, dtype=u.dtype)


These are internal to Devito, should not appear in operator arguments, and should not be touched by the user, and so should use Array, not Function

EdCaunt · 2025-06-16T13:00:33Z

devito/types/multistage.py

+            - `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)`
+        """
+        base_eq=self.eq


With the __new__ specified above, these would just be u = self.lhs.function etc, which would clean things up

EdCaunt · 2025-06-16T13:01:33Z

devito/types/multistage.py

+        assert len(self.a) == self.s, f"'a'={a} must have {self.s} rows"
+        assert len(self.c) == self.s, f"'c'={c} must have {self.s} elements"
+
+    def _evaluate(self, eq_num=0):


Should have **kwargs for consistency with Eq._evaluate()

EdCaunt · 2025-06-16T13:05:29Z

devito/types/multistage.py

+    c : list[float]
+        Time positions of intermediate stages.
+    """
+    a = [[0, 0, 0, 0],


I would set these as tuples in the __init__. Definitely should not be mutable if set on a class level.

I would personally instead have a

def __init__(self): a = (... b = (... c = (... super.__init__(a=a, b=b, c=c)

EdCaunt · 2025-06-16T13:07:52Z

tests/test_multistage.py

+from devito.types.multistage import MultiStage
+
+
+def test_multistage_solve(time_int='RK44'):


These tests are looking better. I would suggest making some very granular tests to verify particular functionalities within MultiStage lowering

ZoeLeibowitz · 2025-06-24T16:36:09Z

tests/test_multistage.py

+                   src_spatial * src_temporal]
+
+    # Time integration scheme
+    return [solve(system_eqs_rhs[i] - U[i], U[i], method=time_int, eq_num=i) for i in range(2)]


You don’t want to use return in your pytest tests. Instead, you should use assertions to verify that the result matches what you expect

EdCaunt · 2025-06-25T14:42:03Z

devito/types/multistage.py

-        # k = [Array(name=f'k{eq_num}{i}', dimensions=grid.shape, grid=grid, dtype=u.dtype) for i in range(self.s)]  # Trying Array
-        k = [Function(name=f'k{eq_num}{i}', grid=grid, space_order=u.space_order, dtype=u.dtype)
+        # 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)


kwargs.get('sregistry').make_name(prefix='k') wants to be inside this loop to ensure that all names are unique

EdCaunt

Some more comments. Improving steadily once again

EdCaunt · 2025-06-25T14:44:30Z

devito/ir/equations/algorithms.py

+    """
+    Handle iterables of expressions.
+    """
+    lowered = []


Could you do return [_lower_multistage(expr, **kwargs) for i in exprs for expr in i]?

EdCaunt · 2025-06-25T15:04:51Z

devito/operations/solve.py

-        return sols[0]
+        sols_temp = sols[0]
+
+    method = kwargs.get("method", None)


Is method a string here, or is it the class for the method? In the latter case, it would remove the need to have the method_registry mapper. Furthermore, it would allow you to have method.resolve(target, sols_temp) here, which is tidier

EdCaunt · 2025-06-25T15:07:11Z

devito/types/multistage.py

+        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


This should be either a class or a callable imo. Alternatively, it should be entirely omitted and set by defining some method/_evaluate method. Of these two, I prefer the latter as it results in cleaner code and a simpler API.

In general, if you are using a string comparison, there is probably a better (and safer) way to achieve your aim.

EdCaunt · 2025-06-25T15:09:13Z

devito/types/multistage.py

+    def __init__(self, **kwargs):
+        self.a, self.b, self.c = self._validate(**kwargs)
+
+    def _validate(self, **kwargs):


Why not just use concrete args with type hinting rather than kwargs?

EdCaunt · 2025-06-25T15:09:37Z

devito/types/multistage.py

+            raise ValueError("RK subclass must define class attributes of the Butcher's array a, b, and c")
+        return a, b, c
+
+    @cached_property


Nitpick: Probably doesn't need caching

EdCaunt · 2025-06-26T09:42:35Z

devito/types/multistage.py

+        Number of stages in the RK method, inferred from `b`.
+    """
+
+    def __init__(self, a=None, b=None, c=None, **kwargs):


This can just be def __init__(self, a: list[list[float | np.number]], b: list[float | np.number], c: list[float | np.number], **kwargs) -> None:, avoiding the need for the validation function

…compatibility with pickles, and checking numerical convergence of three Runge-Kutta methods

codecov · 2025-07-10T16:09:16Z

Codecov Report

❌ Patch coverage is 25.52632% with 283 lines in your changes missing coverage. Please review.
✅ Project coverage is 45.80%. Comparing base (ed3585a) to head (ac1da7e).
⚠️ Report is 235 commits behind head on main.

Files with missing lines	Patch %	Lines
tests/test_multistage.py	10.09%	196 Missing ⚠️
devito/types/multistage.py	40.27%	86 Missing ⚠️
devito/ir/equations/algorithms.py	91.66%	1 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (ed3585a) and HEAD (ac1da7e). Click for more details.

HEAD has 15 uploads less than BASE

Flag BASE (ed3585a) HEAD (ac1da7e)

17 4

pytest-gpu-nvc-nvidiaX 1 0

pytest-gpu-aomp-amdgpuX 1 0

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #2599       +/-   ##
===========================================
- Coverage   92.00%   45.80%   -46.20%     
===========================================
  Files         245      250        +5     
  Lines       48727    50032     +1305     
  Branches     4294     4376       +82     
===========================================
- Hits        44830    22916    -21914     
- Misses       3209    26154    +22945     
- Partials      688      962      +274

Flag	Coverage Δ
pytest-gpu-aomp-amdgpuX	`?`
pytest-gpu-nvc-nvidiaX	`?`

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

EdCaunt · 2025-07-10T15:54:30Z

devito/types/multistage.py

@@ -0,0 +1,357 @@
+from .equation import Eq


Make these imports absolute (devito.types.equation)

EdCaunt · 2025-07-10T15:57:54Z

devito/types/multistage.py

+
+method_registry = {}
+
+def register_method(cls):


I'm not a fan of using string matching here.

I'm also not sure why this function is needed, especially when the registry itself is just a regular dict

EdCaunt · 2025-07-10T15:59:06Z

devito/types/multistage.py

+
+class MultiStage(Eq):
+    """
+       Abstract base class for multi-stage time integration methods


Good docstring, but is it overindented by a level?

EdCaunt · 2025-07-10T15:59:54Z

devito/types/multistage.py

+       of update expressions for each stage in the integration process.
+       """
+
+    def __new__(cls, lhs, rhs=0, subdomain=None, coefficients=None, implicit_dims=None, **kwargs):


This __new__ doesn't seem to do anything. Can it be removed?

EdCaunt · 2025-07-10T16:04:24Z

devito/types/multistage.py

+        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], **kwargs) -> None:


list[list[float | number]] should probably be defined as a type within the class to improve readability:

CoeffsBC = list[float | number] CoeffsA = list[CoeffBC] def __init__(self, a: CoeffsA, b: CoeffsBC, c: CoeffsBC, **kwargs) -> None:

EdCaunt · 2025-07-10T16:05:18Z

devito/types/multistage.py

+        Returns
+        -------
+        list of Eq
+            A list of SymPy Eq objects representing:


Nitpick: they will be Devito Eq objects

EdCaunt · 2025-07-10T16:07:54Z

devito/types/multistage.py

+    c : list[float]
+        Time positions of intermediate stages.
+    """
+    a = [[0, 0, 0, 0, 0, 0, 0, 0, 0],


Possibly a massive pain to do, but would it be possible to have some kind of calculation to get these coefficients rather than having all these tableaus defined in separate classes? In that way, you would just need to specify order and number of stages required

EdCaunt · 2025-07-10T16:08:43Z

devito/types/multistage.py

+
+
+@register_method
+class HORK(MultiStage):


This should just be the default Runge-Kutta unless there's something I'm overlooking imo

EdCaunt · 2025-07-11T08:14:08Z

devito/types/multistage.py

+    """
+
+
+    def ssprk_alpha(mu=1, **kwargs):


Either this doesn't need to live in the class, it should refer to self (perhaps using self.degree), or it should be a staticmethod/classmethod

EdCaunt · 2025-07-11T08:15:26Z

devito/types/multistage.py

+
+        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]


(mu * list(range(1, i))) does this definitely do what it is intended to do?

EdCaunt · 2025-07-11T08:16:12Z

devito/types/multistage.py

+            alpha[1:i] = 1 / (mu * list(range(1, i))) * alpha[:i - 1]
+            alpha[0] = 1 - sum(alpha[1:i + 1])
+
+        return alpha


Doesn't look to be an array getting returned

EdCaunt · 2025-07-11T08:17:49Z

devito/types/multistage.py

+            - 1 final update equation of the form `u.forward = u + dt * sum(b_i * k_i)`
+        """
+
+        u = self.lhs.function


There looks to be some repeated code in here vs other _evaluate methods. Perhaps a generic _evaluate for MultiStage or at least RK classes is in order with a suitable hook (or hooks) to grab the operations specific to that method

EdCaunt · 2025-07-11T08:18:08Z

devito/types/multistage.py

+        an_eq = range(len(U0))
+
+        # Compute SSPRK coefficients
+        alpha = np.array(ssprk_alpha(mu, degree), dtype=np.float64)


Move the array into ssprk_alpha. dtype should be pulled from u

EdCaunt · 2025-07-11T08:48:26Z

tests/test_multistage.py

+    # Time integration scheme
+    pde = [resolve_method(time_int)(u_multi_stage, eq_rhs)]
+    op = Operator(pde, subs=grid.spacing_map)
+    op(dt=0.01, time=1)


Again should assert a norm. Can also be consolidated with the previous test via parameterisation

EdCaunt · 2025-07-11T08:49:20Z

tests/test_multistage.py

+    op(dt=0.01, time=1)
+
+
+@pytest.mark.parametrize('time_int', ['RK44', 'RK32', 'RK97'])


Good! I would use the classes themselves rather than strings

EdCaunt · 2025-07-11T08:50:29Z

tests/test_multistage.py

+    pdes = [Eq(U[i].forward, solve(Eq(U[i].dt-eq_rhs[i]), U[i].forward)) for i in range(len(fun_labels))]
+    op = Operator(pdes, subs=grid.spacing_map)
+    op(dt=dt0, time=nt)
+    assert max(abs(U[0].data[0,:]-U_multi_stage[0].data[0,:]))<10**-5, "the method is not converging to the solution"


A better check would be to compare the L2 norms with np.isclose. See other tests in the codebase for examples

EdCaunt · 2025-07-11T08:52:04Z

tests/test_multistage.py

+
+
+@pytest.mark.parametrize('time_int', ['RK44', 'RK32', 'RK97'])
+def test_multistage_methods_convergence(time_int):


Slightly misleading test name unless I'm missing something? This sounds like a convergence test, but seems to actually test whether the multistage timestepper generates a solution which matches a more simple timestepping scheme

EdCaunt · 2025-07-11T08:52:30Z

tests/test_multistage.py

+    assert max(abs(U[0].data[0,:]-U_multi_stage[0].data[0,:]))<10**-5, "the method is not converging to the solution"
+
+
+


Nitpick: Too many blank lines at end of file

EdCaunt

Some more comments. More granular testing is still needed

EdCaunt · 2025-08-18T09:14:20Z

devito/types/multistage.py

@@ -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."""


For immutability, these should be tuple. Ditto with the next property. They should probably be made into tuples as early as possible

EdCaunt · 2025-08-18T09:14:49Z

devito/types/multistage.py

+
+    @property
+    def src(self):
+        """Return list of right-hand sides."""


Repeated docstring?

EdCaunt · 2025-08-18T09:15:12Z

devito/types/multistage.py

@@ -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]


Does number superclass np.number?

EdCaunt · 2025-08-18T09:16:07Z

devito/types/multistage.py

@@ -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)


Needs whitespace for flake8

EdCaunt · 2025-08-18T09:19:41Z

devito/types/multistage.py

            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


grids = {f.grid for f in u} if not len(grids) == 1: raise ValueError("Cannot construct multi-stage time integrator for Functions on disparate grids") grid = grids.pop() t = grid.time_dim

would be safer

EdCaunt · 2025-08-18T11:26:49Z

tests/test_multistage.py

+    # Source definition
+    src_spatial = Function(name="src_spat", grid=grid, space_order=2, dtype=np.float64)
+    src_spatial.data[100, 100] = 1
+    import sympy as sym


Rogue import

EdCaunt · 2025-08-18T11:27:00Z

tests/test_multistage.py

+    # Time integration scheme
+    pdes = resolve_method(time_int)(U_multi_stage, system_eqs_rhs, source=src, degree=4)
+    op = Operator(pdes, subs=grid.spacing_map)
+    op(dt=0.001, time=2000)


Needs an assertion

EdCaunt · 2025-08-18T11:28:13Z

tests/test_multistage.py

+    assert np.max(np.abs(U[0].data[0,:]-U_multi_stage[0].data[0,:]))<10**-5, "the method is not converging to the solution"
+
+
+def test_multistage_coupled_op_computing_exp(time_int='HORK_EXP'):


Should pass a class, not a string.

EdCaunt · 2025-08-18T11:28:54Z

tests/test_multistage.py

@@ -219,22 +324,22 @@ def test_multistage_coupled_op_computing(time_int='RK44'):
    op(dt=0.01, time=1)


Should also assert something

EdCaunt · 2025-08-18T11:29:20Z

tests/test_saving_multistage.pkl

Should never use pickling for long-term file storage. Remove this file

JDBetteridge

It would be nice (once this has been iterated a bit further and the API concertised) to see an example or demonstrative notebook that showcases the new functionality.

JDBetteridge · 2025-09-09T22:19:02Z

devito/operator/operator.py

@@ -36,7 +36,6 @@
                          disk_layer)
 from devito.types.dimension import Thickness

-


Please run the linter (flake8) 🙂

JDBetteridge · 2025-09-09T22:22:15Z

devito/types/multistage.py

+            f"_evaluate() must be implemented in the subclass {self.__class__.__name__}")
+
+
+class RK(MultiStage):


Suggested change

class RK(MultiStage):

class RungeKutta(MultiStage):

JDBetteridge · 2025-09-09T22:24:26Z

devito/types/multistage.py

+@register_method
+class RK44(RK):


I think RK4 and indeed all RK methods should be instances of the RK Class

Then you no longer need all of the boilerplate code below, which is just setting up Butcher tableau

Or the coefficients should be class attributes and set by the child class

JDBetteridge · 2025-09-09T22:30:07Z

devito/types/multistage.py



 @register_method
-class HORK(MultiStage):
+class HORK_EXP(MultiStage):


HighOrderRungeKuttaExplicit perhaps

Be explicit!

Use CamelCase for class definitions and names_with_underscores for variables and functions.

JDBetteridge · 2025-09-09T22:31:16Z

devito/types/multistage.py



 @register_method
-class HORK(MultiStage):
+class HORK_EXP(MultiStage):


Also if we want to start distinguishing between explicit and implicit timesteppers we should be using a class hierarchy to do so

JDBetteridge · 2025-09-09T22:34:35Z

devito/types/multistage.py

I think this file should be moved to somewhere like devito/timestepping/rungekutta.py or devito/timestepping/explicitmultistage.py that way additional timesteppers can be contributed as new files. (I'm thinking about implicit multistage, backward difference formulae etc...)

JDBetteridge · 2025-09-09T22:35:58Z

tests/test_multistage.py

+    # import matplotlib.pyplot as plt
+    # import numpy as np
+    # t=np.linspace(0,2000,1000)
+    # plt.plot(t,np.exp(1 - 2 * (t - 1)**2))


implementation of multi-stage time integrators

1d830b8

FabioLuporini requested changes May 6, 2025

View reviewed changes

FabioLuporini reviewed May 6, 2025

View reviewed changes

FabioLuporini added the API api (symbolics, types, ...) label May 6, 2025

FabioLuporini requested review from mloubout, EdCaunt, JDBetteridge, ZoeLeibowitz and georgebisbas May 6, 2025 08:04

EdCaunt changed the title ~~implementation of multi-stage time integrators #timestepping~~ dsl: Introduce abstractions for multi-stage time integrators May 6, 2025

EdCaunt reviewed May 6, 2025

View reviewed changes

ZoeLeibowitz reviewed May 13, 2025

View reviewed changes

mloubout requested changes May 15, 2025

View reviewed changes

fernanvr added 3 commits June 13, 2025 18:00

Merge remote-tracking branch 'upstream/main' into multi-stage-time-in…

7f087b3

…tegrator the merge is necessary to maintaing the development of multistages

Return of first PR comments

214d882

Return of first PR comments

d6c4d4a

EdCaunt reviewed Jun 16, 2025

View reviewed changes

fernanvr added 2 commits June 23, 2025 13:44

2nd PR revision

78f8a0b

2nd PR revision

1c9d517

ZoeLeibowitz reviewed Jun 24, 2025

View reviewed changes

fernanvr added 2 commits June 25, 2025 07:57

2rd PR, updating tests and suggestions of 2nd PR revision

11db48b

3rd PR, updating tests and suggestions of 2nd PR revision

83dfb04

EdCaunt reviewed Jun 25, 2025

View reviewed changes

fernanvr added 2 commits June 25, 2025 19:23

4th PR revision, code refining and improving tests

d47a106

4th PR revision, code refining and improving tests

1f93a45

EdCaunt reviewed Jun 26, 2025

View reviewed changes

fernanvr added 2 commits June 26, 2025 18:10

5th PR revision, one suggestion from EdC and improving tests

eea3a52

including two more Runge-Kutta methods and improving tests: checking …

11d1429

…compatibility with pickles, and checking numerical convergence of three Runge-Kutta methods

fernanvr requested review from FabioLuporini, EdCaunt, ZoeLeibowitz and mloubout July 9, 2025 15:59

EdCaunt reviewed Jul 10, 2025

View reviewed changes

EdCaunt reviewed Jul 11, 2025

View reviewed changes

fernanvr added 2 commits July 16, 2025 09:34

changes to consider coupled Multistage equations

4637ac2

Improvements of the HORK_EXP

ac1da7e

EdCaunt reviewed Aug 18, 2025

View reviewed changes

JDBetteridge requested changes Sep 9, 2025

View reviewed changes

		from devito.types.multistage import MultiStage


		def test_multistage_solve(time_int='RK44'):

		op(dt=0.01, time=1)


		@pytest.mark.parametrize('time_int', ['RK44', 'RK32', 'RK97'])



		@pytest.mark.parametrize('time_int', ['RK44', 'RK32', 'RK97'])
		def test_multistage_methods_convergence(time_int):

		assert max(abs(U[0].data[0,:]-U_multi_stage[0].data[0,:]))<10**-5, "the method is not converging to the solution"

		assert np.max(np.abs(U[0].data[0,:]-U_multi_stage[0].data[0,:]))<10**-5, "the method is not converging to the solution"


		def test_multistage_coupled_op_computing_exp(time_int='HORK_EXP'):

		@@ -219,22 +324,22 @@ def test_multistage_coupled_op_computing(time_int='RK44'):
		op(dt=0.01, time=1)

		@@ -36,7 +36,6 @@
		disk_layer)
		from devito.types.dimension import Thickness

		f"_evaluate() must be implemented in the subclass {self.__class__.__name__}")


		class RK(MultiStage):

dsl: Introduce abstractions for multi-stage time integrators #2599

Are you sure you want to change the base?

dsl: Introduce abstractions for multi-stage time integrators #2599

Uh oh!

Conversation

fernanvr commented May 5, 2025

Uh oh!

FabioLuporini left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

EdCaunt left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mloubout left a comment

Choose a reason for hiding this comment

Uh oh!

fernanvr commented Jun 13, 2025

Uh oh!

EdCaunt left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

EdCaunt left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

codecov bot commented Jul 10, 2025 •

edited

Loading

EdCaunt Jul 11, 2025 •

edited

Loading