IMEX Multistage #453

atb1995 · 2023-10-19T14:58:52Z

This code change adds IMEX Multistage class which allows for general split diagonally implicit and explicit Runge-Kutta schemes by passing in a pair of Butcher Tableaus. It also adds IMEX Euler, ARK2, Trap2, SSP3 and ARS3. The butcher tableaus are of the form:

[a_00 0 . 0 ]
[a_10 a_11 . 0 ]
[ . . . . ]
[ b_0 b_1 . b_s]

and

[0 0 . 0 ]
[d_10 0 . 0 ]
[ . . . . ]
[ e_0 e_1 . e_s]

I ran convergence tests for Williamson 5 & 6 (see attached). This pull request is blocked by #440.

…_Multistage

atb1995 · 2023-10-19T14:59:24Z

tommbendall

Looks very exciting, thanks Alex. My suggestions are mainly about adding comments

tommbendall · 2023-10-24T17:02:58Z

integration-tests/model/test_time_discretisation.py

-@pytest.mark.parametrize("scheme", ["ssprk", "implicit_midpoint",
-                                    "RK4", "Heun", "BDF2", "TR_BDF2", "AdamsBashforth", "Leapfrog", "AdamsMoulton"])
+@pytest.mark.parametrize("scheme", ["ssprk", "TrapeziumRule",
+                                    "RK4", "Heun", "BDF2", "TR_BDF2", "AdamsBashforth",


Should we add some of the IMEX schemes here to test them?

The issue here is the test is a prescribed transport test, hence no option for implicit and explicit terms.. Could change or add another test?

We discussed offline: adding a new test_imex.py which could test these on something simple like a wave equation with an analytic solution

I've added test_IMEX.py. It uses a continuity form and splits the term to treat q*div(u) implicitly and u.grad(q) explicitly.

tommbendall · 2023-10-24T17:03:13Z

gusto/labels.py

@@ -90,6 +90,8 @@ def __call__(self, target, value=None):
 # ---------------------------------------------------------------------------- #

 time_derivative = Label("time_derivative")
+implicit = Label("implicit")


Nice to see these labels added!

tommbendall · 2023-10-24T17:03:38Z

gusto/equations.py

@@ -582,7 +582,7 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None,
                 space_names=None, linearisation_map='default',
                 u_transport_option='vector_invariant_form',
                 no_normal_flow_bc_ids=None, active_tracers=None,
-                 thermal=False):
+                 thermal=False, conservative_depth=True):


As discussed offline, I don't think we want to add this argument to the equation

tommbendall · 2023-10-24T17:04:38Z

gusto/time_discretisation.py

+                a butcher tableau defining a given explicit Runge Kutta time discretisation.
+            field_name (str, optional): name of the field to be evolved.
+                Defaults to None.
+            subcycles (int, optional): the number of sub-steps to perform.


You don't seem to have included this argument

tommbendall · 2023-10-24T17:06:04Z

gusto/time_discretisation.py

+                mesh and the compatible function spaces.
+            butcher_imp (numpy array): A matrix containing the coefficients of
+                a butcher tableau defining a given implicit Runge Kutta time discretisation.
+            butcher_exp (numpy array): A matrix containing the coefficients of


I think the type should be :class:`numpy.ndarray`

tommbendall · 2023-10-24T17:07:58Z

gusto/time_discretisation.py

+        super().setup(equation, apply_bcs, *active_labels)
+
+        self.residual = self.residual.label_map(
+            lambda t: any(t.has_label(time_derivative, transport)),


Should transport be included here?

Yes, it is a map if false, so all non-time derivative and transport terms are implicit

We've discussed this offline: we realised that we can label terms as explicit or implicit in any example file using the equation.label_map(...) routine.

But here we should have a check that throws an error if a term hasn't been labelled as implicit/explicit (or is the time_derivative)

An error checker has been added. I raised a "NotImplementedError" - is this the correct choice..?

tommbendall · 2023-10-24T17:08:20Z

gusto/time_discretisation.py

+            lambda t: any(t.has_label(time_derivative, transport)),
+            map_if_false=lambda t: implicit(t))
+
+        self.residual = self.residual.label_map(


Do we need to enforce transport terms being explicit?

This is a good point, it could be left up to the user to label terms as explicit or implicit at the script level.. What do you think?

I think that in the long term we would like users to be able to specify which terms they would like to make implicit / explicit and this could be via a dictionary that is passed in. But for now this seems like the most sensible option. I suggested to @atb1995 that he should add a warning message using the logger to tell people that transport terms are treated explicitly and all others are treated implicitly.

tommbendall · 2023-10-24T17:11:06Z

gusto/time_discretisation.py

+                                       map_if_true=replace_subject(self.x_out, old_idx=self.idx))
+        residual -= mass_form.label_map(all_terms,
+                                        map_if_true=replace_subject(self.x1, old_idx=self.idx))
+        for i in range(stage):


I can follow what is going on in this for loop but maybe you could add some comments to explain it?

Good point! Comments added

tommbendall · 2023-10-24T17:11:22Z

gusto/time_discretisation.py

+        return super(IMEXMultistage, self).rhs
+
+    def res(self, stage):
+        mass_form = self.residual.label_map(


Could you add a docstring here to explain this routine?

I will do! Will also add for final_res and the various solvers

tommbendall · 2023-10-24T17:11:57Z

gusto/time_discretisation.py

+
+    @property
+    def final_res(self):
+        mass_form = self.residual.label_map(lambda t: t.has_label(time_derivative),


I can't spot what is different about this compared with the normal res routine!

The final residual loops across all stages, and also loops up to the final stage for explicit terms.

… Code currently does not work

jshipton · 2023-10-27T07:56:24Z

gusto/time_discretisation.py

+                    self.residual = self.residual.label_map(
+                        lambda t: t.get(transport) == TransportEquationType.conservative,
+                        map_if_true=drop)
+                    self.residual += new_transport_term.form


you should just add the labelled form to the residual, not just the form part, i.e. self.residual += new_transport_term

… linear equation set..

tommbendall

Sorry to be asking for more changes!

tommbendall · 2023-10-31T13:29:14Z

gusto/time_discretisation.py

+        super().setup(equation, apply_bcs, *active_labels)
+
+        self.residual = self.residual.label_map(
+            lambda t: any(t.has_label(time_derivative, transport)),


We've discussed this offline: we realised that we can label terms as explicit or implicit in any example file using the equation.label_map(...) routine.

But here we should have a check that throws an error if a term hasn't been labelled as implicit/explicit (or is the time_derivative)

tommbendall · 2023-10-31T13:32:46Z

gusto/time_discretisation.py

@@ -323,13 +323,37 @@ def setup(self, equation, apply_bcs=True, *active_labels):

        super().setup(equation, apply_bcs, *active_labels)

+        # Get continuity form transport term
+        for t in self.residual:


As we discussed offline: it might be best to move this code to break up the continuity form out of time discretistations to a new routine (we suggested that it would be in common_forms.py). Then a TimeDiscretisation can be agnostic to the terms that it works on, and this code than then be reused for other things.

tommbendall · 2023-10-31T13:34:26Z

integration-tests/model/test_time_discretisation.py

-@pytest.mark.parametrize("scheme", ["ssprk", "implicit_midpoint",
-                                    "RK4", "Heun", "BDF2", "TR_BDF2", "AdamsBashforth", "Leapfrog", "AdamsMoulton"])
+@pytest.mark.parametrize("scheme", ["ssprk", "TrapeziumRule",
+                                    "RK4", "Heun", "BDF2", "TR_BDF2", "AdamsBashforth",


We discussed offline: adding a new test_imex.py which could test these on something simple like a wave equation with an analytic solution

…splitting of terms added to common_forms, new test for IMEX schemes added

tommbendall

Test looks really good, just some questions about what the split_continuity_form routine does when we don't have a mixed function space

tommbendall · 2023-11-01T14:20:09Z

gusto/common_forms.py

+            if hasattr(equation, "field_names"):
+                idx = equation.field_names.index(prognostic_field_name)
+            else:
+                idx = equation.field_name.index(prognostic_field_name)


I think this won't give what we expect! Is field_name a string and not a list?

In the case that we don't have field_names, we don't want an index and don't want to split anything up (because we won't have a mixed function space)

tommbendall

Two tiny changes and then I think that this is ready!

tommbendall · 2023-11-07T16:37:33Z

gusto/common_forms.py

+
+def split_continuity_form(equation):
+    u"""
+    Loops through terms in a given equation, and splits continuity terms into


Could you mention that this will split up all conservative terms in the equation?

Done! Thanks Tom

gusto/common_forms.py

Co-authored-by: tommbendall <[email protected]>

jshipton · 2023-11-08T09:15:48Z

[celebrate] Shipton, Jemma reacted to your message:

…

________________________________ From: tommbendall ***@***.***> Sent: Wednesday, November 8, 2023 9:12:30 AM To: firedrakeproject/gusto ***@***.***> Cc: Shipton, Jemma ***@***.***>; Comment ***@***.***> Subject: Re: [firedrakeproject/gusto] IMEX Multistage (PR #453) CAUTION: This email originated from outside of the organisation. Do not click links or open attachments unless you recognise the sender and know the content is safe. Merged #453<#453> into main. — Reply to this email directly, view it on GitHub<#453 (comment)>, or unsubscribe<https://github.com/notifications/unsubscribe-auth/ABUOTSTX6ZYJWRK4673LCH3YDNEH5AVCNFSM6AAAAAA6HLOQWSVHI2DSMVQWIX3LMV45UABCJFZXG5LFIV3GK3TUJZXXI2LGNFRWC5DJN5XDWMJQHA4TONZUG4ZDAMA>. You are receiving this because you commented.Message ID: ***@***.***>

alexbrown1995 added 8 commits September 7, 2023 14:49

add unit test

1a56e1b

Needs tidying, but IMEX_Euler3 is a seemingly working version of IMEX…

fae09a9

…_Multistage

Tidy and adding new schemes

97ff98f

Merge branch 'main' into IMEX_multistage

1a1fdc0

Add back in courant logging, tidy time disc

aa2d599

change to split advective and div term in SWE

b1b4cf5

SSP3 correction

00f9966

Tidy of code, lint corrections

a6b28a1

alexbrown1995 added 3 commits October 19, 2023 16:01

Merge branch 'main' into IMEX_multistage

9fef4b2

revert explicitmultistage changes

c8207bc

revert test changes

78f5557

tommbendall requested changes Oct 24, 2023

View reviewed changes

alexbrown1995 added 3 commits October 25, 2023 10:55

adressed most review comments

177c308

Warning for explicit implicit splitting added

3a34def

Comments added and equations refactor moved into time_discretisation.…

564cafb

… Code currently does not work

jshipton reviewed Oct 27, 2023

View reviewed changes

alexbrown1995 added 3 commits October 27, 2023 14:00

Splitting on conservative term added. Need to make it work in case of…

31ddf80

… linear equation set..

Merge to main

0c6fba4

Lint fix

85e57a7

tommbendall requested changes Oct 31, 2023

View reviewed changes

alexbrown1995 added 6 commits November 1, 2023 10:25

Labelling of implicit and explicit terms moved outside of time_disc, …

10ed951

…splitting of terms added to common_forms, new test for IMEX schemes added

change error message

e866193

Merge to main

cccc433

lint fix + firedrake fml move

5878f23

Error switched to Runtime

958e002

linearisation terms added

6daae50

tommbendall requested changes Nov 1, 2023

View reviewed changes

bug fix to test func / q

9191901

tommbendall requested changes Nov 7, 2023

View reviewed changes

atb1995 and others added 2 commits November 7, 2023 16:54

Update gusto/common_forms.py

d29d137

Co-authored-by: tommbendall <[email protected]>

Change in description

7a14129

tommbendall approved these changes Nov 8, 2023

View reviewed changes

tommbendall merged commit c209deb into main Nov 8, 2023
4 checks passed

tommbendall deleted the IMEX_multistage branch November 8, 2023 09:12

IMEX Multistage #453

IMEX Multistage #453

Conversation

atb1995 commented Oct 19, 2023

atb1995 commented Oct 19, 2023 • edited Loading

tommbendall left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tommbendall Oct 24, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tommbendall left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tommbendall left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tommbendall left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jshipton commented Nov 8, 2023 via email

atb1995 commented Oct 19, 2023 •

edited

Loading

tommbendall Oct 24, 2023 •

edited

Loading