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

        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

I think now it's done

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

devito/operations/solve.py



-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

    # 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?

Yes, eq_num it is for the counting in order to not repeat names. The think is that I'm having some troubles while implementing the SymbolRegistry

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

The thing is, when I try using Array, the second test fails, but if I use Functions, it works without any errors. I'm not sure if I'm doing something wrong.

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)

I did something like that... could you check?

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

This is still a work in progress, I just need a bit more time.

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

I did something like that...

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

It's a string. The idea is that the user provides a string to identify which time integrator to apply.

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.

Sorry for that, I removed 'method' from the class, but didn't update the docstring... fixing it now

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

oh, nice...

…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

The idea was to automatically map each time integration method's string representation to its corresponding class, eliminating the need to manually populate the dictionary. This allows users to specify the time integrator using only a string, without needing to explicitly import the corresponding class. Do you recommend another way to do this?

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?

Yes, I change it now, for the stages coupling

Yes, I changed a little bit now to consider stage coupling.

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

That’s a good idea, but these coefficients are computed based on different conditions. And there can be multiple Runge–Kutta methods with the same order and number of stages, each satisfying different constraints. So, in general, there isn’t a single formula to compute them...

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

Related to the previous comment, this is a specific class of Runge–Kutta methods that may have the SSP (Strong Stability Preserving) property. However, they do not encompass the other lower-order Runge–Kutta methods defined by Butcher tableaus in the other classes.

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

Yes, the HORK class is still experimental, and there are some issues in its interaction with Devito that may require code changes. Once I have a stable, fully functional version, I will take your suggestion into account.

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

changed, but I used lists instead, to not have to depend on numpy.

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

The same, the idea is to check if op.apply() executes without an error...

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

Do you think is best for the user to import the class than using an string to reference it?

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

changed to relative L2 norm

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

Well, the idea is to see if matches a method that we know it is correct. In that sense is a convergence test, even if I'm not testing the convergence order of the method. Isn't it?

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


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

    """

-    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?

no, is the same, I changed it to np.number to avoid confusion

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

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)`
        """
        n_eq=len(self.eq)


Needs whitespace for flake8

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

devito/types/multistage.py

        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

Maybe I'm wrong, but I think the functions should be allowed to be defined in different grids, like different staggered grids... the method should work also for those cases.

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.

I see your point; however, I think specifying the time integrator as a string is more convenient for users compared to having them import a specific class for their chosen method.

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

tests/test_multistage.py

    # Time integration scheme
    pdes = resolve_method(time_int)(U_multi_stage, system_eqs_rhs)
    op = Operator(pdes, subs=grid.spacing_map)
    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

                          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

I think I didn't understand well. Aren’t they already implemented as described in the second option of your comment?

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

I renamed the class to HighOrderRungeKuttaExponential. I realize the name might be confusing since this particular Runge-Kutta is explicit, but “EXP” was intended to highlight the exponential aspect. I’ve also updated the other class names based on your suggestions.

Regarding the file location, it’s currently in /types as recommended by @mloubout (see suggestion). Personally, I think both /timestepping and /types are reasonable options. Perhaps we can discuss this with @EdCaunt and @FabioLuporini to reach a consensus.

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


Actually, those tests are still under construction (have errors), so, the code is to help with the debugging :)

…tegrator

- Restored from commit 7d526db48 (multi-stage-time-integrator-rebase) - Clean imports and method registry system - Improved class names: RungeKutta44, RungeKutta32, RungeKutta97, HighOrderRungeKuttaExponential - Tuple-based coefficients implementation - Extensive test suite improvements (568 lines) - All flake8 fixes and code quality improvements

review-notebook-app · 2025-10-23T19:31:46Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

- Enhanced HORK_EXP method with improved source term handling - Fixed source inclusion and expansion coefficient computation - Added test coverage for HORK methods - Updated test structure with better parameterization

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 EdCaunt, JDBetteridge, ZoeLeibowitz, georgebisbas and mloubout 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 EdCaunt, FabioLuporini, 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 3 commits July 16, 2025 09:34

changes to consider coupled Multistage equations

4637ac2

Improvements of the HORK_EXP

ac1da7e

Merge branch 'main' into multi-stage-time-integrator

e9b3533

EdCaunt reviewed Aug 18, 2025

View reviewed changes

JDBetteridge requested changes Sep 9, 2025

View reviewed changes

fernanvr force-pushed the multi-stage-time-integrator branch 2 times, most recently from 67823fa to e9b3533 Compare October 7, 2025 23:59

fernanvr added 4 commits October 7, 2025 21:13

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

dc3dd77

…tegrator

improving spacing in some tests

a0c45c1

Add MFE time stepping Jupyter notebook

93c6e3f

fernanvr added 2 commits October 23, 2025 16:45

Remove MFE_time_size.ipynb notebook

ef8d1ac

Update multistage implementation and tests

fa5acac

- Enhanced HORK_EXP method with improved source term handling - Fixed source inclusion and expansion coefficient computation - Added test coverage for HORK methods - Updated test structure with better parameterization



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

		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'):

		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!

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!