From 1a56e1b819bb684035340bb24beb045aca620094 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Thu, 7 Sep 2023 14:49:24 +0100 Subject: [PATCH 01/22] add unit test --- integration-tests/model/test_time_discretisation.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/integration-tests/model/test_time_discretisation.py b/integration-tests/model/test_time_discretisation.py index 5be78c871..e4170be49 100644 --- a/integration-tests/model/test_time_discretisation.py +++ b/integration-tests/model/test_time_discretisation.py @@ -8,7 +8,7 @@ def run(timestepper, tmax, f_end): return norm(timestepper.fields("f") - f_end) / norm(f_end) -@pytest.mark.parametrize("scheme", ["ssprk", "implicit_midpoint", +@pytest.mark.parametrize("scheme", ["ssprk", "TrapeziumRule", "ImplicitMidpoint", "RK4", "Heun", "BDF2", "TR_BDF2", "AdamsBashforth", "Leapfrog", "AdamsMoulton"]) def test_time_discretisation(tmpdir, scheme, tracer_setup): if (scheme == "AdamsBashforth"): @@ -27,8 +27,10 @@ def test_time_discretisation(tmpdir, scheme, tracer_setup): if scheme == "ssprk": transport_scheme = SSPRK3(domain) - elif scheme == "implicit_midpoint": + elif scheme == "TrapeziumRule": transport_scheme = TrapeziumRule(domain) + elif scheme == "ImplicitMidpoint": + transport_scheme = ImplicitMidpoint(domain) elif scheme == "RK4": transport_scheme = RK4(domain) elif scheme == "Heun": From fae09a959f292ad73f0ff68b5cc9b01af3aa4084 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Thu, 14 Sep 2023 13:46:59 +0100 Subject: [PATCH 02/22] Needs tidying, but IMEX_Euler3 is a seemingly working version of IMEX_Multistage --- gusto/labels.py | 2 + gusto/time_discretisation.py | 544 ++++++++++++++++++++++++++++++++++- gusto/timeloop.py | 26 +- 3 files changed, 556 insertions(+), 16 deletions(-) diff --git a/gusto/labels.py b/gusto/labels.py index 74e458729..230c4ecb0 100644 --- a/gusto/labels.py +++ b/gusto/labels.py @@ -11,6 +11,8 @@ # ---------------------------------------------------------------------------- # time_derivative = Label("time_derivative") +implicit = Label("implicit") +explicit = Label("explicit") transport = Label("transport", validator=lambda value: type(value) == TransportEquationType) diffusion = Label("diffusion") physics = Label("physics", validator=lambda value: type(value) == MethodType) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index dce51c963..6a226d679 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -7,7 +7,7 @@ from abc import ABCMeta, abstractmethod, abstractproperty from firedrake import (Function, TestFunction, NonlinearVariationalProblem, - NonlinearVariationalSolver, DirichletBC, split) + NonlinearVariationalSolver, DirichletBC, split, Constant) from firedrake.formmanipulation import split_form from firedrake.utils import cached_property @@ -15,7 +15,7 @@ from gusto.fml import ( replace_subject, replace_test_function, Term, all_terms, drop ) -from gusto.labels import time_derivative, prognostic, physics +from gusto.labels import time_derivative, prognostic, physics, transport, implicit, explicit from gusto.logging import logger, DEBUG, logging_ksp_monitor_true_residual from gusto.wrappers import * import numpy as np @@ -23,7 +23,8 @@ __all__ = ["ForwardEuler", "BackwardEuler", "ExplicitMultistage", "ImplicitMultistage", "SSPRK3", "RK4", "Heun", "ThetaMethod", "TrapeziumRule", "BDF2", "TR_BDF2", - "Leapfrog", "AdamsMoulton", "AdamsBashforth", "ImplicitMidpoint", "QinZhang"] + "Leapfrog", "AdamsMoulton", "AdamsBashforth", "ImplicitMidpoint", "QinZhang", "IMEX_Euler", + "IMEX_Euler2","IMEX_Euler3"] def wrapper_apply(original_apply): @@ -234,6 +235,524 @@ def apply(self, x_out, x_in): pass +class IMEXMultistage(TimeDiscretisation): + """ + A class for implementing general diagonally implicit multistage (Runge-Kutta) + methods based on its Butcher tableau. + + A Butcher tableau is formed in the following way for a s-th order + diagonally implicit scheme: + + c_0 | a_00 a_01 . a_0s + c_1 | a_10 a_11 a_1s + . | . . . . + . | . . . . + c_s | a_s0 a_s1 . a_ss + ------------------------- + | b_1 b_2 ... b_s + + The gradient function has no time-dependence, so c elements are not needed. + A square 'butcher_matrix' is defined in each class that uses + the ImplicitMultiStage structure, + + [a_00 0 . 0 ] + [a_10 a_11 . 0 ] + [ . . . . ] + [ b_0 b_1 . b_s] + + Unlike the explicit method, all upper diagonal a_ij elements are non-zero for implicit methods. + + There are three steps to move from the current solution, y^n, to the new one, y^{n+1} + + For an s stage method, + At iteration i (from 0 to s-1) + An intermediate location is computed as y_i = y^n + sum{over j less than or equal to i} (dt*a_ij*k_i) + Compute the gradient at the intermediate location, k_i = F(y_i) + + At the last stage, compute the new solution by: + y^{n+1} = y^n + sum_{j from 0 to s} (dt*b_i*k_i) + + """ + + def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None, butcher_imp=None, butcher_exp=None): + super().__init__(domain, field_name=field_name, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + if butcher_imp is not None: + self.butcher_imp = butcher_imp + self.nStages = int(np.shape(self.butcher_imp)[1]) + if butcher_exp is not None: + self.butcher_exp = butcher_exp + self.nStages = int(np.shape(self.butcher_exp)[1]) + + def setup(self, equation, apply_bcs=True, *active_labels): + """ + Set up the time discretisation based on the equation. + + Args: + equation (:class:`PrognosticEquation`): the model's equation. + *active_labels (:class:`Label`): labels indicating which terms of + the equation to include. + """ + + super().setup(equation, apply_bcs, *active_labels) + + self.xs = [Function(self.fs) for i in range(self.nStages)] + + def lhs(self): + return super().lhs + + def rhs(self): + return super().rhs + + def solvers(self, stage): + + mass_form = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=drop) + residual = mass_form.label_map(all_terms, + replace_subject(self.x_out, self.idx)) + residual =residual - mass_form.label_map(all_terms, + replace_subject(self.x1, self.idx)) + for i in range(stage): + # r_imp = self.residual.label_map( + # lambda t: any(t.has_label(time_derivative, transport)), + # map_if_true=drop, + # map_if_false=replace_subject(self.xs[i], self.idx), + # ) + r_exp = self.residual.label_map( + lambda t: t.has_label(transport), + map_if_true=replace_subject(self.x1, self.idx), + map_if_false=drop, + ) + print("imp1",self.butcher_imp[stage, i]) + + print("exp1",self.butcher_exp[stage, i]) + # residual = residual + r_exp.label_map( + # all_terms, + # lambda t: self.dt*t) + + print("exp1",self.butcher_exp[stage, i]) + residual = residual + r_exp.label_map( + all_terms, + lambda t: self.dt*t) + print("imp2",self.butcher_imp[stage, stage]) + r_imp = self.residual.label_map( + lambda t: t.has_label(time_derivative, transport), + map_if_true=drop, + map_if_false=replace_subject(self.x_out, self.idx), + ) + residual = residual + r_imp.label_map( + all_terms, + map_if_true=lambda t: self.dt*t) + + problem = NonlinearVariationalProblem(residual.form, self.x_out, bcs=self.bcs) + + solver_name = self.field_name+self.__class__.__name__ + "%s" % (stage) + return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, + options_prefix=solver_name) + @cached_property + def solverf(self): + mass_form = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=drop) + residual = mass_form.label_map(all_terms, + replace_subject(self.x_out, self.idx)) + residual -= mass_form.label_map(all_terms, + replace_subject(self.x1, self.idx)) + for i in range(self.nStages): + r_exp = self.residual.label_map( + lambda t: t.has_label(transport), + map_if_true=replace_subject(self.xs[i], self.idx), + map_if_false=drop, + ) + # residual = residual + r_exp.label_map( + # all_terms, + # lambda t: Constant(self.butcher_exp[self.nStages, i])*self.dt*t) + + r_imp = self.residual.label_map( + lambda t: t.has_label(time_derivative, transport), + map_if_true=drop, + map_if_false=replace_subject(self.xs[i], self.idx), + ) + residual = residual + r_imp.label_map( + all_terms, + map_if_true=lambda t: Constant(self.butcher_imp[self.nStages, i])*self.dt*t) + problem = NonlinearVariationalProblem(residual.form, self.x_out, bcs=self.bcs) + + solver_name = self.field_name+self.__class__.__name__ + return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, + options_prefix=solver_name) + + def solve_stage(self, x0, stage): + self.x1.assign(x0) + + self.solvers(stage).solve() + + self.xs[stage].assign(self.x_out) + + print(np.max(self.xs[stage].dat.data[:]), np.min(self.xs[stage].dat.data[:])) + print(np.max(self.x1.dat.data[:]), np.min(self.x1.dat.data[:])) + + if self.limiter is not None: + self.limiter.apply(self.xs[stage]) + + def apply(self, x_out, x_in): + self.xs[0].assign(x_in) + for i in range(1,self.nStages): + print("stage:", i) + self.solve_stage(x_in, i) + print("final_solve") + self.x1.assign(x_in) + self.solverf.solve() + x_out.assign(self.x_out) + + if self.limiter is not None: + self.limiter.apply(x_out) + + +class IMEX_Euler3(TimeDiscretisation): + + def setup(self, equation, apply_bcs=True, *active_labels): + """ + Set up the time discretisation based on the equation. + + Args: + equation (:class:`PrognosticEquation`): the model's equation. + *active_labels (:class:`Label`): labels indicating which terms of + the equation to include. + """ + + super().setup(equation, apply_bcs, *active_labels) + + self.residual = self.residual.label_map( + lambda t: any(t.has_label(time_derivative, transport)), + map_if_false=lambda t: implicit(t)) + + self.residual = self.residual.label_map( + lambda t: t.has_label(transport), + map_if_true=lambda t: explicit(t)) + + self.butcher_imp = np.array([[0., 0.], [0., 1.], [0., 1.]]) + self.butcher_exp = np.array([[0., 0.], [1., 0.], [1., 0.]]) + + self.butcher_imp = np.array([[0., 0.], [0.5, 0.5], [0.5, 0.5]]) + self.butcher_exp = np.array([[0., 0.], [1., 0.], [0.5, 0.5]]) + + # ARS3 + g = (3. + np.sqrt(3.))/6. + + self.butcher_imp = np.array([[0., 0., 0.], [0., g, 0.], [0., 1-2.*g, g], [0., 0.5, 0.5]]) + self.butcher_exp = np.array([[0., 0., 0.], [g, 0., 0.], [g-1., 2.*(1.-g), 0.], [0., 0.5, 0.5]]) + + # Trap2 + self.butcher_imp = np.array([[0., 0., 0.], [0.5, 0.5, 0.], [0.5, 0., 0.5], [0.5, 0., 0.5]]) + self.butcher_exp = np.array([[0., 0., 0.], [1., 0., 0.], [0.5, 0.5, 0.], [0.5, 0.5, 0.]]) + + # ARK2 + g = 1. - 1./np.sqrt(2.) + d = 1./(2.*np.sqrt(2.)) + a = 1./6.*(3. + 2.*np.sqrt(2.)) + self.butcher_imp = np.array([[0., 0., 0.], [g, g, 0.], [d, d, g], [d, d, g]]) + self.butcher_exp = np.array([[0., 0., 0.], [2.*g, 0., 0.], [1.-a, a, 0.], [d, d, g]]) + + self.nStages = int(np.shape(self.butcher_imp)[1]) + self.xs = [Function(self.fs) for i in range(self.nStages)] + + @property + def lhs(self): + l = self.residual.label_map( + lambda t: any(t.has_label(implicit, time_derivative)), + map_if_true=replace_subject(self.x_out, old_idx=self.idx), + map_if_false=drop) + l = l.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: self.dt*t) + return l.form + + @property + def rhs(self): + r = self.residual.label_map( + lambda t: any(t.has_label(explicit, time_derivative)), + map_if_true=replace_subject(self.x1, old_idx=self.idx), + map_if_false=drop) + r = r.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -self.dt*t) + + return r.form + + @property + def resval0(self): + stage =0 + mass_form = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=drop) + res = mass_form.label_map(all_terms, + map_if_true=replace_subject(self.x_out, old_idx=self.idx)) + res -= mass_form.label_map(all_terms, + map_if_true=replace_subject(self.x1, old_idx=self.idx)) + for i in range(stage): + r_exp = self.residual.label_map( + lambda t: t.has_label(explicit), + map_if_true=replace_subject(self.xs[i], old_idx=self.idx), + map_if_false=drop) + r_exp = r_exp.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: Constant(self.butcher_exp[stage, i])*self.dt*t) + r_imp = self.residual.label_map( + lambda t: t.has_label(implicit), + map_if_true=replace_subject(self.xs[i], old_idx=self.idx), + map_if_false=drop) + r_imp = r_imp.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: Constant(self.butcher_imp[stage, i])*self.dt*t) + res += r_imp + res +=r_exp + r_imp = self.residual.label_map( + lambda t: t.has_label(implicit), + map_if_true=replace_subject(self.x_out, old_idx=self.idx), + map_if_false=drop) + r_imp = r_imp.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: Constant(self.butcher_imp[stage, stage])*self.dt*t) + res += r_imp + return res.form + + def resval(self, stage): + #stage =1 + print("stage:", stage) + mass_form = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=drop) + res = mass_form.label_map(all_terms, + map_if_true=replace_subject(self.x_out, old_idx=self.idx)) + res -= mass_form.label_map(all_terms, + map_if_true=replace_subject(self.x1, old_idx=self.idx)) + for i in range(stage): + r_exp = self.residual.label_map( + lambda t: t.has_label(explicit), + map_if_true=replace_subject(self.xs[i], old_idx=self.idx), + map_if_false=drop) + print("stage, i, a_exp:", stage, i, self.butcher_exp[stage, i]) + r_exp = r_exp.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: Constant(self.butcher_exp[stage, i])*self.dt*t) + print("stage, i, a_imp:", stage, i, self.butcher_imp[stage, i]) + r_imp = self.residual.label_map( + lambda t: t.has_label(implicit), + map_if_true=replace_subject(self.xs[i], old_idx=self.idx), + map_if_false=drop) + r_imp = r_imp.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: Constant(self.butcher_imp[stage, i])*self.dt*t) + res += r_imp + res +=r_exp + print("stage, i, a_imp:", stage, stage, self.butcher_imp[stage, stage]) + r_imp = self.residual.label_map( + lambda t: t.has_label(implicit), + map_if_true=replace_subject(self.x_out, old_idx=self.idx), + map_if_false=drop) + r_imp = r_imp.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: Constant(self.butcher_imp[stage, stage])*self.dt*t) + res += r_imp + return res.form + + @property + def resfin(self): + mass_form = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=drop) + res = mass_form.label_map(all_terms, + map_if_true=replace_subject(self.x_out, old_idx=self.idx)) + res -= mass_form.label_map(all_terms, + map_if_true=replace_subject(self.x1, old_idx=self.idx)) + for i in range(self.nStages): + r_exp = self.residual.label_map( + lambda t: t.has_label(explicit), + map_if_true=replace_subject(self.xs[i], old_idx=self.idx), + map_if_false=drop) + r_exp = r_exp.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: Constant(self.butcher_exp[self.nStages, i])*self.dt*t) + r_imp = self.residual.label_map( + lambda t: t.has_label(implicit), + map_if_true=replace_subject(self.xs[i], old_idx=self.idx), + map_if_false=drop) + r_imp = r_imp.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: Constant(self.butcher_imp[self.nStages, i])*self.dt*t) + res += r_imp + res +=r_exp + return res.form + + def solver(self,stage): + # setup solver using lhs and rhs defined in derived class + problem = NonlinearVariationalProblem(self.resval(stage), self.x_out, bcs=self.bcs) + solver_name = self.field_name+self.__class__.__name__+"%s"%(stage) + return NonlinearVariationalSolver(problem, options_prefix=solver_name) + + @cached_property + def solvers(self): + solvers = [] + for stage in range(self.nStages): + # setup solver using lhs and rhs defined in derived class + problem = NonlinearVariationalProblem(self.resval(stage), self.x_out, bcs=self.bcs) + solver_name = self.field_name+self.__class__.__name__+"%s"%(stage) + solvers.append(NonlinearVariationalSolver(problem, options_prefix=solver_name)) + return solvers + + @cached_property + def solver_fin(self): + # setup solver using lhs and rhs defined in derived class + problem = NonlinearVariationalProblem(self.resfin, self.x_out, bcs=self.bcs) + solver_name = self.field_name+self.__class__.__name__ + return NonlinearVariationalSolver(problem, options_prefix=solver_name) + + + def apply(self, x_out, x_in): + self.x1.assign(x_in) + solver_list = self.solvers + + for stage in range(self.nStages): + solver = solver_list[stage] + solver.solve() + self.xs[stage].assign(self.x_out) + + self.solver_fin.solve() + x_out.assign(self.x_out) + + +class IMEX_Euler2(TimeDiscretisation): + """ + A class for implementing general diagonally implicit multistage (Runge-Kutta) + methods based on its Butcher tableau. + + A Butcher tableau is formed in the following way for a s-th order + diagonally implicit scheme: + + c_0 | a_00 a_01 . a_0s + c_1 | a_10 a_11 a_1s + . | . . . . + . | . . . . + c_s | a_s0 a_s1 . a_ss + ------------------------- + | b_1 b_2 ... b_s + + The gradient function has no time-dependence, so c elements are not needed. + A square 'butcher_matrix' is defined in each class that uses + the ImplicitMultiStage structure, + + [a_00 0 . 0 ] + [a_10 a_11 . 0 ] + [ . . . . ] + [ b_0 b_1 . b_s] + + Unlike the explicit method, all upper diagonal a_ij elements are non-zero for implicit methods. + + There are three steps to move from the current solution, y^n, to the new one, y^{n+1} + + For an s stage method, + At iteration i (from 0 to s-1) + An intermediate location is computed as y_i = y^n + sum{over j less than or equal to i} (dt*a_ij*k_i) + Compute the gradient at the intermediate location, k_i = F(y_i) + + At the last stage, compute the new solution by: + y^{n+1} = y^n + sum_{j from 0 to s} (dt*b_i*k_i) + + """ + + def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None): + super().__init__(domain, field_name=field_name, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + + def setup(self, equation, apply_bcs=True, *active_labels): + """ + Set up the time discretisation based on the equation. + + Args: + equation (:class:`PrognosticEquation`): the model's equation. + *active_labels (:class:`Label`): labels indicating which terms of + the equation to include. + """ + + super().setup(equation, apply_bcs, *active_labels) + + def lhs(self): + return super().lhs + + def rhs(self): + return super().rhs + + @cached_property + def solver(self): + + r_exp = self.residual.label_map( + lambda t: any(t.has_label(transport, time_derivative)), + map_if_true=replace_subject(self.x1, self.idx), + map_if_false=drop, + ) + r_exp = r_exp.label_map( + lambda t: t.has_label(transport), + map_if_true=lambda t: -self.dt*t) + r_imp = self.residual.label_map( + lambda t: t.has_label(transport), + map_if_true=drop, + map_if_false=replace_subject(self.x_out, self.idx), + ) + r_imp= r_imp.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: self.dt*t) + + problem = NonlinearVariationalProblem(r_imp.form - r_exp.form, self.x_out, bcs=self.bcs) + + solver_name = self.field_name+self.__class__.__name__ + return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, + options_prefix=solver_name) + + def solvers(self, stage): + mass_form = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=drop) + residual = mass_form.label_map(all_terms, + replace_subject(self.x_out, self.idx)) + residual =residual - mass_form.label_map(all_terms, + replace_subject(self.x1, self.idx)) + + r_exp = self.residual.label_map( + lambda t: t.has_label(transport), + map_if_true=replace_subject(self.x1, self.idx), + map_if_false=drop, + ) + residual = residual + r_exp.label_map( + all_terms, + lambda t: self.dt*t) + r_imp = self.residual.label_map( + lambda t: any(t.has_label(time_derivative, transport)), + map_if_true=drop, + map_if_false=replace_subject(self.x_out, self.idx), + ) + residual = residual + r_imp.label_map( + all_terms, + map_if_true=lambda t: self.dt*t) + + problem = NonlinearVariationalProblem(residual.form, self.x_out, bcs=self.bcs) + + solver_name = self.field_name+self.__class__.__name__ + "%s" % (stage) + return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, + options_prefix=solver_name) + + def apply(self, x_out, x_in): + self.x1.assign(x_in) + self.solver.solve() + x_out.assign(self.x_out) + + if self.limiter is not None: + self.limiter.apply(x_out) + + class ImplicitMultistage(TimeDiscretisation): """ A class for implementing general diagonally implicit multistage (Runge-Kutta) @@ -1495,3 +2014,22 @@ def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None limiter=limiter, options=options) self.butcher_matrix = np.array([[0.25, 0], [0.5, 0.25], [0.5, 0.5]]) self.nStages = int(np.shape(self.butcher_matrix)[1]) + +class IMEX_Euler(IMEXMultistage): + u""" + Implements Qin and Zhang's two-stage, 2nd order, implicit Runge–Kutta method. + + The method, for solving + ∂y/∂t = F(y), can be written as: + + k0 = F[y^n + 0.25*dt*k0] + k1 = F[y^n + 0.5*dt*k0 + 0.25*dt*k1] + y^(n+1) = y^n + 0.5*dt*(k0 + k1) + """ + def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None, butcher_imp=None, butcher_exp=None): + super().__init__(domain, field_name, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + self.butcher_imp = np.array([[0., 0.], [0., 1.], [0., 1.]]) + self.butcher_exp = np.array([[0., 0.], [1., 0.], [1., 0.]]) + self.nStages = int(np.shape(self.butcher_imp)[1]) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 1b5024f28..3a2715da0 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -152,13 +152,13 @@ def run(self, t, tmax, pick_up=False): # Set up diagnostics, which may set up some fields necessary to pick up self.io.setup_diagnostics(self.fields) - self.io.setup_log_courant(self.fields) - if self.equation.domain.mesh.extruded: - self.io.setup_log_courant(self.fields, component='horizontal') - self.io.setup_log_courant(self.fields, component='vertical') - if self.transporting_velocity != "prognostic": - self.io.setup_log_courant(self.fields, name='transporting_velocity', - expression=self.transporting_velocity) + # self.io.setup_log_courant(self.fields) + # if self.equation.domain.mesh.extruded: + # self.io.setup_log_courant(self.fields, component='horizontal') + # self.io.setup_log_courant(self.fields, component='vertical') + # if self.transporting_velocity != "prognostic": + # self.io.setup_log_courant(self.fields, name='transporting_velocity', + # expression=self.transporting_velocity) if pick_up: # Pick up fields, and return other info to be picked up @@ -178,10 +178,10 @@ def run(self, t, tmax, pick_up=False): self.x.update() - self.io.log_courant(self.fields) - if self.equation.domain.mesh.extruded: - self.io.log_courant(self.fields, component='horizontal', message='horizontal') - self.io.log_courant(self.fields, component='vertical', message='vertical') + # self.io.log_courant(self.fields) + # if self.equation.domain.mesh.extruded: + # self.io.log_courant(self.fields, component='horizontal', message='horizontal') + # self.io.log_courant(self.fields, component='vertical', message='vertical') self.timestep() @@ -563,8 +563,8 @@ def timestep(self): for k in range(self.maxk): with timed_stage("Transport"): - self.io.log_courant(self.fields, 'transporting_velocity', - message=f'transporting velocity, outer iteration {k}') + # self.io.log_courant(self.fields, 'transporting_velocity', + # message=f'transporting velocity, outer iteration {k}') for name, scheme in self.active_transport: # transports a field from xstar and puts result in xp scheme.apply(xp(name), xstar(name)) From 97ff98feaf000392b46af66b36143de78d1233f2 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Mon, 18 Sep 2023 09:50:34 +0100 Subject: [PATCH 03/22] Tidy and adding new schemes --- gusto/time_discretisation.py | 418 +++++------------------------------ 1 file changed, 53 insertions(+), 365 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 6a226d679..588153dfa 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -21,10 +21,10 @@ import numpy as np -__all__ = ["ForwardEuler", "BackwardEuler", "ExplicitMultistage", "ImplicitMultistage", +__all__ = ["ForwardEuler", "BackwardEuler", "ExplicitMultistage", "ImplicitMultistage", "IMEXMultistage", "SSPRK3", "RK4", "Heun", "ThetaMethod", "TrapeziumRule", "BDF2", "TR_BDF2", "Leapfrog", "AdamsMoulton", "AdamsBashforth", "ImplicitMidpoint", "QinZhang", "IMEX_Euler", - "IMEX_Euler2","IMEX_Euler3"] + "ARS3","ARK2"] def wrapper_apply(original_apply): @@ -234,185 +234,21 @@ def apply(self, x_out, x_in): """ pass - class IMEXMultistage(TimeDiscretisation): - """ - A class for implementing general diagonally implicit multistage (Runge-Kutta) - methods based on its Butcher tableau. - - A Butcher tableau is formed in the following way for a s-th order - diagonally implicit scheme: - - c_0 | a_00 a_01 . a_0s - c_1 | a_10 a_11 a_1s - . | . . . . - . | . . . . - c_s | a_s0 a_s1 . a_ss - ------------------------- - | b_1 b_2 ... b_s - The gradient function has no time-dependence, so c elements are not needed. - A square 'butcher_matrix' is defined in each class that uses - the ImplicitMultiStage structure, - - [a_00 0 . 0 ] - [a_10 a_11 . 0 ] - [ . . . . ] - [ b_0 b_1 . b_s] - - Unlike the explicit method, all upper diagonal a_ij elements are non-zero for implicit methods. - - There are three steps to move from the current solution, y^n, to the new one, y^{n+1} - - For an s stage method, - At iteration i (from 0 to s-1) - An intermediate location is computed as y_i = y^n + sum{over j less than or equal to i} (dt*a_ij*k_i) - Compute the gradient at the intermediate location, k_i = F(y_i) - - At the last stage, compute the new solution by: - y^{n+1} = y^n + sum_{j from 0 to s} (dt*b_i*k_i) - - """ - - def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None, butcher_imp=None, butcher_exp=None): + def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, + options=None, butcher_imp=None, butcher_exp=None): super().__init__(domain, field_name=field_name, - solver_parameters=solver_parameters, - limiter=limiter, options=options) + solver_parameters=solver_parameters, + limiter=limiter, options=options) if butcher_imp is not None: self.butcher_imp = butcher_imp self.nStages = int(np.shape(self.butcher_imp)[1]) + if butcher_exp is not None: self.butcher_exp = butcher_exp self.nStages = int(np.shape(self.butcher_exp)[1]) - def setup(self, equation, apply_bcs=True, *active_labels): - """ - Set up the time discretisation based on the equation. - - Args: - equation (:class:`PrognosticEquation`): the model's equation. - *active_labels (:class:`Label`): labels indicating which terms of - the equation to include. - """ - - super().setup(equation, apply_bcs, *active_labels) - - self.xs = [Function(self.fs) for i in range(self.nStages)] - - def lhs(self): - return super().lhs - - def rhs(self): - return super().rhs - - def solvers(self, stage): - - mass_form = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_false=drop) - residual = mass_form.label_map(all_terms, - replace_subject(self.x_out, self.idx)) - residual =residual - mass_form.label_map(all_terms, - replace_subject(self.x1, self.idx)) - for i in range(stage): - # r_imp = self.residual.label_map( - # lambda t: any(t.has_label(time_derivative, transport)), - # map_if_true=drop, - # map_if_false=replace_subject(self.xs[i], self.idx), - # ) - r_exp = self.residual.label_map( - lambda t: t.has_label(transport), - map_if_true=replace_subject(self.x1, self.idx), - map_if_false=drop, - ) - print("imp1",self.butcher_imp[stage, i]) - - print("exp1",self.butcher_exp[stage, i]) - # residual = residual + r_exp.label_map( - # all_terms, - # lambda t: self.dt*t) - - print("exp1",self.butcher_exp[stage, i]) - residual = residual + r_exp.label_map( - all_terms, - lambda t: self.dt*t) - print("imp2",self.butcher_imp[stage, stage]) - r_imp = self.residual.label_map( - lambda t: t.has_label(time_derivative, transport), - map_if_true=drop, - map_if_false=replace_subject(self.x_out, self.idx), - ) - residual = residual + r_imp.label_map( - all_terms, - map_if_true=lambda t: self.dt*t) - - problem = NonlinearVariationalProblem(residual.form, self.x_out, bcs=self.bcs) - - solver_name = self.field_name+self.__class__.__name__ + "%s" % (stage) - return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, - options_prefix=solver_name) - @cached_property - def solverf(self): - mass_form = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_false=drop) - residual = mass_form.label_map(all_terms, - replace_subject(self.x_out, self.idx)) - residual -= mass_form.label_map(all_terms, - replace_subject(self.x1, self.idx)) - for i in range(self.nStages): - r_exp = self.residual.label_map( - lambda t: t.has_label(transport), - map_if_true=replace_subject(self.xs[i], self.idx), - map_if_false=drop, - ) - # residual = residual + r_exp.label_map( - # all_terms, - # lambda t: Constant(self.butcher_exp[self.nStages, i])*self.dt*t) - - r_imp = self.residual.label_map( - lambda t: t.has_label(time_derivative, transport), - map_if_true=drop, - map_if_false=replace_subject(self.xs[i], self.idx), - ) - residual = residual + r_imp.label_map( - all_terms, - map_if_true=lambda t: Constant(self.butcher_imp[self.nStages, i])*self.dt*t) - problem = NonlinearVariationalProblem(residual.form, self.x_out, bcs=self.bcs) - - solver_name = self.field_name+self.__class__.__name__ - return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, - options_prefix=solver_name) - - def solve_stage(self, x0, stage): - self.x1.assign(x0) - - self.solvers(stage).solve() - - self.xs[stage].assign(self.x_out) - - print(np.max(self.xs[stage].dat.data[:]), np.min(self.xs[stage].dat.data[:])) - print(np.max(self.x1.dat.data[:]), np.min(self.x1.dat.data[:])) - - if self.limiter is not None: - self.limiter.apply(self.xs[stage]) - - def apply(self, x_out, x_in): - self.xs[0].assign(x_in) - for i in range(1,self.nStages): - print("stage:", i) - self.solve_stage(x_in, i) - print("final_solve") - self.x1.assign(x_in) - self.solverf.solve() - x_out.assign(self.x_out) - - if self.limiter is not None: - self.limiter.apply(x_out) - - -class IMEX_Euler3(TimeDiscretisation): - def setup(self, equation, apply_bcs=True, *active_labels): """ Set up the time discretisation based on the equation. @@ -433,29 +269,6 @@ def setup(self, equation, apply_bcs=True, *active_labels): lambda t: t.has_label(transport), map_if_true=lambda t: explicit(t)) - self.butcher_imp = np.array([[0., 0.], [0., 1.], [0., 1.]]) - self.butcher_exp = np.array([[0., 0.], [1., 0.], [1., 0.]]) - - self.butcher_imp = np.array([[0., 0.], [0.5, 0.5], [0.5, 0.5]]) - self.butcher_exp = np.array([[0., 0.], [1., 0.], [0.5, 0.5]]) - - # ARS3 - g = (3. + np.sqrt(3.))/6. - - self.butcher_imp = np.array([[0., 0., 0.], [0., g, 0.], [0., 1-2.*g, g], [0., 0.5, 0.5]]) - self.butcher_exp = np.array([[0., 0., 0.], [g, 0., 0.], [g-1., 2.*(1.-g), 0.], [0., 0.5, 0.5]]) - - # Trap2 - self.butcher_imp = np.array([[0., 0., 0.], [0.5, 0.5, 0.], [0.5, 0., 0.5], [0.5, 0., 0.5]]) - self.butcher_exp = np.array([[0., 0., 0.], [1., 0., 0.], [0.5, 0.5, 0.], [0.5, 0.5, 0.]]) - - # ARK2 - g = 1. - 1./np.sqrt(2.) - d = 1./(2.*np.sqrt(2.)) - a = 1./6.*(3. + 2.*np.sqrt(2.)) - self.butcher_imp = np.array([[0., 0., 0.], [g, g, 0.], [d, d, g], [d, d, g]]) - self.butcher_exp = np.array([[0., 0., 0.], [2.*g, 0., 0.], [1.-a, a, 0.], [d, d, g]]) - self.nStages = int(np.shape(self.butcher_imp)[1]) self.xs = [Function(self.fs) for i in range(self.nStages)] @@ -482,46 +295,7 @@ def rhs(self): return r.form - @property - def resval0(self): - stage =0 - mass_form = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_false=drop) - res = mass_form.label_map(all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - res -= mass_form.label_map(all_terms, - map_if_true=replace_subject(self.x1, old_idx=self.idx)) - for i in range(stage): - r_exp = self.residual.label_map( - lambda t: t.has_label(explicit), - map_if_true=replace_subject(self.xs[i], old_idx=self.idx), - map_if_false=drop) - r_exp = r_exp.label_map( - lambda t: t.has_label(time_derivative), - map_if_false=lambda t: Constant(self.butcher_exp[stage, i])*self.dt*t) - r_imp = self.residual.label_map( - lambda t: t.has_label(implicit), - map_if_true=replace_subject(self.xs[i], old_idx=self.idx), - map_if_false=drop) - r_imp = r_imp.label_map( - lambda t: t.has_label(time_derivative), - map_if_false=lambda t: Constant(self.butcher_imp[stage, i])*self.dt*t) - res += r_imp - res +=r_exp - r_imp = self.residual.label_map( - lambda t: t.has_label(implicit), - map_if_true=replace_subject(self.x_out, old_idx=self.idx), - map_if_false=drop) - r_imp = r_imp.label_map( - lambda t: t.has_label(time_derivative), - map_if_false=lambda t: Constant(self.butcher_imp[stage, stage])*self.dt*t) - res += r_imp - return res.form - def resval(self, stage): - #stage =1 - print("stage:", stage) mass_form = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_false=drop) @@ -534,11 +308,9 @@ def resval(self, stage): lambda t: t.has_label(explicit), map_if_true=replace_subject(self.xs[i], old_idx=self.idx), map_if_false=drop) - print("stage, i, a_exp:", stage, i, self.butcher_exp[stage, i]) r_exp = r_exp.label_map( lambda t: t.has_label(time_derivative), map_if_false=lambda t: Constant(self.butcher_exp[stage, i])*self.dt*t) - print("stage, i, a_imp:", stage, i, self.butcher_imp[stage, i]) r_imp = self.residual.label_map( lambda t: t.has_label(implicit), map_if_true=replace_subject(self.xs[i], old_idx=self.idx), @@ -548,7 +320,6 @@ def resval(self, stage): map_if_false=lambda t: Constant(self.butcher_imp[stage, i])*self.dt*t) res += r_imp res +=r_exp - print("stage, i, a_imp:", stage, stage, self.butcher_imp[stage, stage]) r_imp = self.residual.label_map( lambda t: t.has_label(implicit), map_if_true=replace_subject(self.x_out, old_idx=self.idx), @@ -624,135 +395,6 @@ def apply(self, x_out, x_in): x_out.assign(self.x_out) -class IMEX_Euler2(TimeDiscretisation): - """ - A class for implementing general diagonally implicit multistage (Runge-Kutta) - methods based on its Butcher tableau. - - A Butcher tableau is formed in the following way for a s-th order - diagonally implicit scheme: - - c_0 | a_00 a_01 . a_0s - c_1 | a_10 a_11 a_1s - . | . . . . - . | . . . . - c_s | a_s0 a_s1 . a_ss - ------------------------- - | b_1 b_2 ... b_s - - The gradient function has no time-dependence, so c elements are not needed. - A square 'butcher_matrix' is defined in each class that uses - the ImplicitMultiStage structure, - - [a_00 0 . 0 ] - [a_10 a_11 . 0 ] - [ . . . . ] - [ b_0 b_1 . b_s] - - Unlike the explicit method, all upper diagonal a_ij elements are non-zero for implicit methods. - - There are three steps to move from the current solution, y^n, to the new one, y^{n+1} - - For an s stage method, - At iteration i (from 0 to s-1) - An intermediate location is computed as y_i = y^n + sum{over j less than or equal to i} (dt*a_ij*k_i) - Compute the gradient at the intermediate location, k_i = F(y_i) - - At the last stage, compute the new solution by: - y^{n+1} = y^n + sum_{j from 0 to s} (dt*b_i*k_i) - - """ - - def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None): - super().__init__(domain, field_name=field_name, - solver_parameters=solver_parameters, - limiter=limiter, options=options) - - def setup(self, equation, apply_bcs=True, *active_labels): - """ - Set up the time discretisation based on the equation. - - Args: - equation (:class:`PrognosticEquation`): the model's equation. - *active_labels (:class:`Label`): labels indicating which terms of - the equation to include. - """ - - super().setup(equation, apply_bcs, *active_labels) - - def lhs(self): - return super().lhs - - def rhs(self): - return super().rhs - - @cached_property - def solver(self): - - r_exp = self.residual.label_map( - lambda t: any(t.has_label(transport, time_derivative)), - map_if_true=replace_subject(self.x1, self.idx), - map_if_false=drop, - ) - r_exp = r_exp.label_map( - lambda t: t.has_label(transport), - map_if_true=lambda t: -self.dt*t) - r_imp = self.residual.label_map( - lambda t: t.has_label(transport), - map_if_true=drop, - map_if_false=replace_subject(self.x_out, self.idx), - ) - r_imp= r_imp.label_map( - lambda t: t.has_label(time_derivative), - map_if_false=lambda t: self.dt*t) - - problem = NonlinearVariationalProblem(r_imp.form - r_exp.form, self.x_out, bcs=self.bcs) - - solver_name = self.field_name+self.__class__.__name__ - return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, - options_prefix=solver_name) - - def solvers(self, stage): - mass_form = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_false=drop) - residual = mass_form.label_map(all_terms, - replace_subject(self.x_out, self.idx)) - residual =residual - mass_form.label_map(all_terms, - replace_subject(self.x1, self.idx)) - - r_exp = self.residual.label_map( - lambda t: t.has_label(transport), - map_if_true=replace_subject(self.x1, self.idx), - map_if_false=drop, - ) - residual = residual + r_exp.label_map( - all_terms, - lambda t: self.dt*t) - r_imp = self.residual.label_map( - lambda t: any(t.has_label(time_derivative, transport)), - map_if_true=drop, - map_if_false=replace_subject(self.x_out, self.idx), - ) - residual = residual + r_imp.label_map( - all_terms, - map_if_true=lambda t: self.dt*t) - - problem = NonlinearVariationalProblem(residual.form, self.x_out, bcs=self.bcs) - - solver_name = self.field_name+self.__class__.__name__ + "%s" % (stage) - return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, - options_prefix=solver_name) - - def apply(self, x_out, x_in): - self.x1.assign(x_in) - self.solver.solve() - x_out.assign(self.x_out) - - if self.limiter is not None: - self.limiter.apply(x_out) - - class ImplicitMultistage(TimeDiscretisation): """ A class for implementing general diagonally implicit multistage (Runge-Kutta) @@ -2033,3 +1675,49 @@ def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None self.butcher_imp = np.array([[0., 0.], [0., 1.], [0., 1.]]) self.butcher_exp = np.array([[0., 0.], [1., 0.], [1., 0.]]) self.nStages = int(np.shape(self.butcher_imp)[1]) + +class ARS3(IMEXMultistage): + u""" + Implements Qin and Zhang's two-stage, 2nd order, implicit Runge–Kutta method. + + The method, for solving + ∂y/∂t = F(y), can be written as: + + k0 = F[y^n + 0.25*dt*k0] + k1 = F[y^n + 0.5*dt*k0 + 0.25*dt*k1] + y^(n+1) = y^n + 0.5*dt*(k0 + k1) + """ + def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None, butcher_imp=None, butcher_exp=None): + super().__init__(domain, field_name, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + # ARS3 + g = (3. + np.sqrt(3.))/6. + + self.butcher_imp = np.array([[0., 0., 0.], [0., g, 0.], [0., 1-2.*g, g], [0., 0.5, 0.5]]) + self.butcher_exp = np.array([[0., 0., 0.], [g, 0., 0.], [g-1., 2.*(1.-g), 0.], [0., 0.5, 0.5]]) + + self.nStages = int(np.shape(self.butcher_imp)[1]) + +class ARK2(IMEXMultistage): + u""" + Implements Qin and Zhang's two-stage, 2nd order, implicit Runge–Kutta method. + + The method, for solving + ∂y/∂t = F(y), can be written as: + + k0 = F[y^n + 0.25*dt*k0] + k1 = F[y^n + 0.5*dt*k0 + 0.25*dt*k1] + y^(n+1) = y^n + 0.5*dt*(k0 + k1) + """ + def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None, butcher_imp=None, butcher_exp=None): + super().__init__(domain, field_name, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + # ARK2 + g = 1. - 1./np.sqrt(2.) + d = 1./(2.*np.sqrt(2.)) + a = 1./6.*(3. + 2.*np.sqrt(2.)) + self.butcher_imp = np.array([[0., 0., 0.], [g, g, 0.], [d, d, g], [d, d, g]]) + self.butcher_exp = np.array([[0., 0., 0.], [2.*g, 0., 0.], [1.-a, a, 0.], [d, d, g]]) + self.nStages = int(np.shape(self.butcher_imp)[1]) From aa2d599484fc0ac495c8f47c5cd9f94654f3816f Mon Sep 17 00:00:00 2001 From: atb1995 Date: Tue, 19 Sep 2023 10:56:05 +0100 Subject: [PATCH 04/22] Add back in courant logging, tidy time disc --- gusto/time_discretisation.py | 44 +++++++++++++++++++++++++++++++++++- gusto/timeloop.py | 26 ++++++++++----------- 2 files changed, 56 insertions(+), 14 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 588153dfa..c319f79af 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -24,7 +24,7 @@ __all__ = ["ForwardEuler", "BackwardEuler", "ExplicitMultistage", "ImplicitMultistage", "IMEXMultistage", "SSPRK3", "RK4", "Heun", "ThetaMethod", "TrapeziumRule", "BDF2", "TR_BDF2", "Leapfrog", "AdamsMoulton", "AdamsBashforth", "ImplicitMidpoint", "QinZhang", "IMEX_Euler", - "ARS3","ARK2"] + "ARS3","ARK2", "Trap2", "SSP3"] def wrapper_apply(original_apply): @@ -1721,3 +1721,45 @@ def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None self.butcher_imp = np.array([[0., 0., 0.], [g, g, 0.], [d, d, g], [d, d, g]]) self.butcher_exp = np.array([[0., 0., 0.], [2.*g, 0., 0.], [1.-a, a, 0.], [d, d, g]]) self.nStages = int(np.shape(self.butcher_imp)[1]) + +class SSP3(IMEXMultistage): + u""" + Implements Qin and Zhang's two-stage, 2nd order, implicit Runge–Kutta method. + + The method, for solving + ∂y/∂t = F(y), can be written as: + + k0 = F[y^n + 0.25*dt*k0] + k1 = F[y^n + 0.5*dt*k0 + 0.25*dt*k1] + y^(n+1) = y^n + 0.5*dt*(k0 + k1) + """ + def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None, butcher_imp=None, butcher_exp=None): + super().__init__(domain, field_name, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + # SSP3(3,3,2) + g = 1. - (1./np.sqrt(2.)) + self.butcher_imp = np.array([[g, 0., 0.], [1-2.*g, g, 0.], [0.5-g, 0., g], [(1./6.),(1./6.), (2./6.)]]) + self.butcher_exp = np.array([[0., 0., 0.], [1., 0., 0.], [0.25, 0.25, 0.], [(1./6.),(1./6.), (2./6.)]]) + self.nStages = int(np.shape(self.butcher_imp)[1]) + +class Trap2(IMEXMultistage): + u""" + Implements Qin and Zhang's two-stage, 2nd order, implicit Runge–Kutta method. + + The method, for solving + ∂y/∂t = F(y), can be written as: + + k0 = F[y^n + 0.25*dt*k0] + k1 = F[y^n + 0.5*dt*k0 + 0.25*dt*k1] + y^(n+1) = y^n + 0.5*dt*(k0 + k1) + """ + def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None, butcher_imp=None, butcher_exp=None): + super().__init__(domain, field_name, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + # Trap2(2+e,3,2) + e = 0. + self.butcher_imp = np.array([[0., 0., 0.,0,], [e, 0., 0.,0.], [0.5, 0., 0.5, 0.], [0.5, 0., 0., 0.5], [0.5, 0., 0., 0.5]]) + self.butcher_exp= np.array([[0., 0., 0.,0,], [1., 0., 0.,0.], [0.5, 0.5, 0., 0.], [0.5, 0., 0.5, 0.], [0.5, 0., 0.5, 0.]]) + self.nStages = int(np.shape(self.butcher_imp)[1]) \ No newline at end of file diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 3a2715da0..1b5024f28 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -152,13 +152,13 @@ def run(self, t, tmax, pick_up=False): # Set up diagnostics, which may set up some fields necessary to pick up self.io.setup_diagnostics(self.fields) - # self.io.setup_log_courant(self.fields) - # if self.equation.domain.mesh.extruded: - # self.io.setup_log_courant(self.fields, component='horizontal') - # self.io.setup_log_courant(self.fields, component='vertical') - # if self.transporting_velocity != "prognostic": - # self.io.setup_log_courant(self.fields, name='transporting_velocity', - # expression=self.transporting_velocity) + self.io.setup_log_courant(self.fields) + if self.equation.domain.mesh.extruded: + self.io.setup_log_courant(self.fields, component='horizontal') + self.io.setup_log_courant(self.fields, component='vertical') + if self.transporting_velocity != "prognostic": + self.io.setup_log_courant(self.fields, name='transporting_velocity', + expression=self.transporting_velocity) if pick_up: # Pick up fields, and return other info to be picked up @@ -178,10 +178,10 @@ def run(self, t, tmax, pick_up=False): self.x.update() - # self.io.log_courant(self.fields) - # if self.equation.domain.mesh.extruded: - # self.io.log_courant(self.fields, component='horizontal', message='horizontal') - # self.io.log_courant(self.fields, component='vertical', message='vertical') + self.io.log_courant(self.fields) + if self.equation.domain.mesh.extruded: + self.io.log_courant(self.fields, component='horizontal', message='horizontal') + self.io.log_courant(self.fields, component='vertical', message='vertical') self.timestep() @@ -563,8 +563,8 @@ def timestep(self): for k in range(self.maxk): with timed_stage("Transport"): - # self.io.log_courant(self.fields, 'transporting_velocity', - # message=f'transporting velocity, outer iteration {k}') + self.io.log_courant(self.fields, 'transporting_velocity', + message=f'transporting velocity, outer iteration {k}') for name, scheme in self.active_transport: # transports a field from xstar and puts result in xp scheme.apply(xp(name), xstar(name)) From b1b4cf5e897ec86df268fffd96d2d7510e37ee70 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Fri, 29 Sep 2023 11:05:00 +0100 Subject: [PATCH 05/22] change to split advective and div term in SWE --- gusto/equations.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/gusto/equations.py b/gusto/equations.py index dd4a9ceb4..32846e1f7 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -680,7 +680,8 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, raise ValueError("Invalid u_transport_option: %s" % u_transport_option) # Depth transport term - D_adv = prognostic(continuity_form(phi, D, u), 'D') + #D_adv = prognostic(continuity_form(phi, D, u), 'D') + D_adv = prognostic(advection_form(phi, D, u), 'D') # Transport term needs special linearisation if self.linearisation_map(D_adv.terms[0]): linear_D_adv = linear_continuity_form(phi, H, u_trial) @@ -711,6 +712,10 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, residual = (mass_form + adv_form + pressure_gradient_form) + geo_grad_form = subject(prognostic(phi*D*div(u)*dx), self.X) + + residual += geo_grad_form + # -------------------------------------------------------------------- # # Extra Terms (Coriolis, Topography and Thermal) # -------------------------------------------------------------------- # From 00f9966c00811e7247a44c2a267b86c14910ee2c Mon Sep 17 00:00:00 2001 From: atb1995 Date: Thu, 5 Oct 2023 12:32:10 +0100 Subject: [PATCH 06/22] SSP3 correction --- gusto/time_discretisation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index c319f79af..8590a1aff 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -1739,8 +1739,8 @@ def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None limiter=limiter, options=options) # SSP3(3,3,2) g = 1. - (1./np.sqrt(2.)) - self.butcher_imp = np.array([[g, 0., 0.], [1-2.*g, g, 0.], [0.5-g, 0., g], [(1./6.),(1./6.), (2./6.)]]) - self.butcher_exp = np.array([[0., 0., 0.], [1., 0., 0.], [0.25, 0.25, 0.], [(1./6.),(1./6.), (2./6.)]]) + self.butcher_imp = np.array([[g, 0., 0.], [1-2.*g, g, 0.], [0.5-g, 0., g], [(1./6.),(1./6.), (2./3.)]]) + self.butcher_exp = np.array([[0., 0., 0.], [1., 0., 0.], [0.25, 0.25, 0.], [(1./6.),(1./6.), (2./3.)]]) self.nStages = int(np.shape(self.butcher_imp)[1]) class Trap2(IMEXMultistage): From a6b28a1e9b1943171b79b9a0412ac7ed94cfa47e Mon Sep 17 00:00:00 2001 From: atb1995 Date: Thu, 19 Oct 2023 15:52:11 +0100 Subject: [PATCH 07/22] Tidy of code, lint corrections --- gusto/equations.py | 16 +- gusto/time_discretisation.py | 562 +++++++++++++++++------------------ 2 files changed, 276 insertions(+), 302 deletions(-) diff --git a/gusto/equations.py b/gusto/equations.py index 32846e1f7..47f7a4bd5 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -586,7 +586,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): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -680,8 +680,11 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, raise ValueError("Invalid u_transport_option: %s" % u_transport_option) # Depth transport term - #D_adv = prognostic(continuity_form(phi, D, u), 'D') - D_adv = prognostic(advection_form(phi, D, u), 'D') + if (conservative_depth): + D_adv = prognostic(continuity_form(phi, D, u), 'D') + else: + D_adv = prognostic(advection_form(phi, D, u), 'D') + # Transport term needs special linearisation if self.linearisation_map(D_adv.terms[0]): linear_D_adv = linear_continuity_form(phi, H, u_trial) @@ -712,9 +715,10 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, residual = (mass_form + adv_form + pressure_gradient_form) - geo_grad_form = subject(prognostic(phi*D*div(u)*dx), self.X) - - residual += geo_grad_form + # Add divergence term + if (not conservative_depth): + geo_grad_form = subject(prognostic(phi*D*div(u)*dx), self.X) + residual += geo_grad_form # -------------------------------------------------------------------- # # Extra Terms (Coriolis, Topography and Thermal) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 8590a1aff..71fb0de4a 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -7,7 +7,7 @@ from abc import ABCMeta, abstractmethod, abstractproperty from firedrake import (Function, TestFunction, NonlinearVariationalProblem, - NonlinearVariationalSolver, DirichletBC, split, Constant) + NonlinearVariationalSolver, DirichletBC, Constant) from firedrake.formmanipulation import split_form from firedrake.utils import cached_property @@ -21,10 +21,10 @@ import numpy as np -__all__ = ["ForwardEuler", "BackwardEuler", "ExplicitMultistage", "ImplicitMultistage", "IMEXMultistage", +__all__ = ["ForwardEuler", "BackwardEuler", "ExplicitMultistage", "IMEXMultistage", "SSPRK3", "RK4", "Heun", "ThetaMethod", "TrapeziumRule", "BDF2", "TR_BDF2", - "Leapfrog", "AdamsMoulton", "AdamsBashforth", "ImplicitMidpoint", "QinZhang", "IMEX_Euler", - "ARS3","ARK2", "Trap2", "SSP3"] + "Leapfrog", "AdamsMoulton", "AdamsBashforth", "IMEX_Euler", + "ARS3", "ARK2", "Trap2", "SSP3"] def wrapper_apply(original_apply): @@ -234,20 +234,78 @@ def apply(self, x_out, x_in): """ pass + class IMEXMultistage(TimeDiscretisation): + """ + A class for implementing general IMEX multistage (Runge-Kutta) + methods based on two Butcher tableaus, to solve \n + + ∂y/∂t = F(y) + S(y) \n - def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, - options=None, butcher_imp=None, butcher_exp=None): + Where F are implicit fast terms, and S are explicit slow terms. \n + + There are three steps to move from the current solution, y^n, to the new one, y^{n+1} + + For each i = 1, s in an s stage method + we compute the intermediate solutions: \n + y_i = y^n + dt*(a_i1*F(y_1) + a_i2*F(y_2)+ ... + a_ii*F(y_i)) \n + + dt*(d_i1*S(y_1) + d_i2*S(y_2)+ ... + d_{i,i-1}*S(y_{i-1})) + + At the last stage, compute the new solution by: \n + y^{n+1} = y^n + dt*(b_1*F(y_1) + b_2*F(y_2) + .... + b_s*F(y_s)) \n + + dt*(e_1*S(y_1) + e_2*S(y_2) + .... + e_s*S(y_s)) \n + + """ + # --------------------------------------------------------------------------- + # Butcher tableaus for a s-th order + # diagonally implicit scheme (left) and explicit scheme (right): + # c_0 | a_00 0 . 0 f_0 | 0 0 . 0 + # c_1 | a_10 a_11 . 0 f_1 | d_10 0 . 0 + # . | . . . . . | . . . . + # . | . . . . . | . . . . + # c_s | a_s0 a_s1 . a_ss f_s | d_s0 d_s1 . 0 + # ------------------------- ------------------------- + # | b_1 b_2 ... b_s | b_1 b_2 ... b_s + # + # + # The corresponding square 'butcher_imp' and 'butcher_exp' matrices are: + # + # [a_00 0 0 . 0 ] [ 0 0 0 . 0 ] + # [a_10 a_11 0 . 0 ] [d_10 0 0 . 0 ] + # [a_20 a_21 a_22 . 0 ] [d_20 d_21 0 . 0 ] + # [ . . . . . ] [ . . . . . ] + # [ b_0 b_1 . b_s] [ e_0 e_1 . . e_s] + # --------------------------------------------------------------------------- + + def __init__(self, domain, butcher_imp, butcher_exp, field_name=None, + solver_parameters=None, limiter=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + 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 + 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. + Defaults to None. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ super().__init__(domain, field_name=field_name, - solver_parameters=solver_parameters, - limiter=limiter, options=options) - if butcher_imp is not None: - self.butcher_imp = butcher_imp - self.nStages = int(np.shape(self.butcher_imp)[1]) - - if butcher_exp is not None: - self.butcher_exp = butcher_exp - self.nStages = int(np.shape(self.butcher_exp)[1]) + solver_parameters=solver_parameters, + limiter=limiter, options=options) + self.butcher_imp = butcher_imp + self.butcher_exp = butcher_exp + self.nStages = int(np.shape(self.butcher_imp)[1]) def setup(self, equation, apply_bcs=True, *active_labels): """ @@ -269,40 +327,26 @@ def setup(self, equation, apply_bcs=True, *active_labels): lambda t: t.has_label(transport), map_if_true=lambda t: explicit(t)) - self.nStages = int(np.shape(self.butcher_imp)[1]) self.xs = [Function(self.fs) for i in range(self.nStages)] - @property + @cached_property def lhs(self): - l = self.residual.label_map( - lambda t: any(t.has_label(implicit, time_derivative)), - map_if_true=replace_subject(self.x_out, old_idx=self.idx), - map_if_false=drop) - l = l.label_map( - lambda t: t.has_label(time_derivative), - map_if_false=lambda t: self.dt*t) - return l.form + """Set up the discretisation's left hand side (the time derivative).""" + return super(IMEXMultistage, self).lhs - @property + @cached_property def rhs(self): - r = self.residual.label_map( - lambda t: any(t.has_label(explicit, time_derivative)), - map_if_true=replace_subject(self.x1, old_idx=self.idx), - map_if_false=drop) - r = r.label_map( - lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -self.dt*t) + """Set up the discretisation's right hand side (the time derivative).""" + return super(IMEXMultistage, self).rhs - return r.form - - def resval(self, stage): + def res(self, stage): mass_form = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_false=drop) - res = mass_form.label_map(all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - res -= mass_form.label_map(all_terms, - map_if_true=replace_subject(self.x1, old_idx=self.idx)) + residual = mass_form.label_map(all_terms, + 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): r_exp = self.residual.label_map( lambda t: t.has_label(explicit), @@ -318,27 +362,26 @@ def resval(self, stage): r_imp = r_imp.label_map( lambda t: t.has_label(time_derivative), map_if_false=lambda t: Constant(self.butcher_imp[stage, i])*self.dt*t) - res += r_imp - res +=r_exp + residual += r_imp + residual += r_exp r_imp = self.residual.label_map( - lambda t: t.has_label(implicit), - map_if_true=replace_subject(self.x_out, old_idx=self.idx), - map_if_false=drop) + lambda t: t.has_label(implicit), + map_if_true=replace_subject(self.x_out, old_idx=self.idx), + map_if_false=drop) r_imp = r_imp.label_map( lambda t: t.has_label(time_derivative), map_if_false=lambda t: Constant(self.butcher_imp[stage, stage])*self.dt*t) - res += r_imp - return res.form - + residual += r_imp + return residual.form + @property - def resfin(self): - mass_form = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_false=drop) - res = mass_form.label_map(all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - res -= mass_form.label_map(all_terms, - map_if_true=replace_subject(self.x1, old_idx=self.idx)) + def final_res(self): + mass_form = self.residual.label_map(lambda t: t.has_label(time_derivative), + map_if_false=drop) + residual = mass_form.label_map(all_terms, + 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(self.nStages): r_exp = self.residual.label_map( lambda t: t.has_label(explicit), @@ -354,162 +397,40 @@ def resfin(self): r_imp = r_imp.label_map( lambda t: t.has_label(time_derivative), map_if_false=lambda t: Constant(self.butcher_imp[self.nStages, i])*self.dt*t) - res += r_imp - res +=r_exp - return res.form + residual += r_imp + residual += r_exp + return residual.form - def solver(self,stage): - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.resval(stage), self.x_out, bcs=self.bcs) - solver_name = self.field_name+self.__class__.__name__+"%s"%(stage) - return NonlinearVariationalSolver(problem, options_prefix=solver_name) - @cached_property def solvers(self): solvers = [] for stage in range(self.nStages): - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.resval(stage), self.x_out, bcs=self.bcs) - solver_name = self.field_name+self.__class__.__name__+"%s"%(stage) + # setup solver using residual defined in derived class + problem = NonlinearVariationalProblem(self.res(stage), self.x_out, bcs=self.bcs) + solver_name = self.field_name+self.__class__.__name__ + "%s" % (stage) solvers.append(NonlinearVariationalSolver(problem, options_prefix=solver_name)) return solvers - + @cached_property - def solver_fin(self): + def final_solver(self): # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.resfin, self.x_out, bcs=self.bcs) + problem = NonlinearVariationalProblem(self.final_res, self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__ return NonlinearVariationalSolver(problem, options_prefix=solver_name) - def apply(self, x_out, x_in): self.x1.assign(x_in) solver_list = self.solvers - + for stage in range(self.nStages): - solver = solver_list[stage] - solver.solve() + self.solver = solver_list[stage] + self.solver.solve() self.xs[stage].assign(self.x_out) - self.solver_fin.solve() + self.final_solver.solve() x_out.assign(self.x_out) -class ImplicitMultistage(TimeDiscretisation): - """ - A class for implementing general diagonally implicit multistage (Runge-Kutta) - methods based on its Butcher tableau. - - A Butcher tableau is formed in the following way for a s-th order - diagonally implicit scheme: - - c_0 | a_00 a_01 . a_0s - c_1 | a_10 a_11 a_1s - . | . . . . - . | . . . . - c_s | a_s0 a_s1 . a_ss - ------------------------- - | b_1 b_2 ... b_s - - The gradient function has no time-dependence, so c elements are not needed. - A square 'butcher_matrix' is defined in each class that uses - the ImplicitMultiStage structure, - - [a_00 0 . 0 ] - [a_10 a_11 . 0 ] - [ . . . . ] - [ b_0 b_1 . b_s] - - Unlike the explicit method, all upper diagonal a_ij elements are non-zero for implicit methods. - - There are three steps to move from the current solution, y^n, to the new one, y^{n+1} - - For an s stage method, - At iteration i (from 0 to s-1) - An intermediate location is computed as y_i = y^n + sum{over j less than or equal to i} (dt*a_ij*k_i) - Compute the gradient at the intermediate location, k_i = F(y_i) - - At the last stage, compute the new solution by: - y^{n+1} = y^n + sum_{j from 0 to s} (dt*b_i*k_i) - - """ - - def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None, butcher_matrix=None): - super().__init__(domain, field_name=field_name, - solver_parameters=solver_parameters, - limiter=limiter, options=options) - if butcher_matrix is not None: - self.butcher_matrix = butcher_matrix - self.nStages = int(np.shape(self.butcher_matrix)[1]) - - def setup(self, equation, apply_bcs=True, *active_labels): - """ - Set up the time discretisation based on the equation. - - Args: - equation (:class:`PrognosticEquation`): the model's equation. - *active_labels (:class:`Label`): labels indicating which terms of - the equation to include. - """ - - super().setup(equation, apply_bcs, *active_labels) - - self.k = [Function(self.fs) for i in range(self.nStages)] - - def lhs(self): - return super().lhs - - def rhs(self): - return super().rhs - - def solvers(self, stage): - residual = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=drop, - map_if_false=replace_subject(self.xnph, self.idx), - ) - mass_form = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_false=drop) - residual += mass_form.label_map(all_terms, - replace_subject(self.x_out, self.idx)) - - problem = NonlinearVariationalProblem(residual.form, self.x_out, bcs=self.bcs) - - solver_name = self.field_name+self.__class__.__name__ + "%s" % (stage) - return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, - options_prefix=solver_name) - - def solve_stage(self, x0, stage): - self.x1.assign(x0) - for i in range(stage): - self.x1.assign(self.x1 + self.butcher_matrix[stage, i]*self.dt*self.k[i]) - - if self.limiter is not None: - self.limiter.apply(self.x1) - - if self.idx is None and len(self.fs) > 1: - self.xnph = tuple([self.dt*self.butcher_matrix[stage, stage]*a + b for a, b in zip(split(self.x_out), split(self.x1))]) - else: - self.xnph = self.x1 + self.butcher_matrix[stage, stage]*self.dt*self.x_out - - self.solvers(stage).solve() - - self.k[stage].assign(self.x_out) - - def apply(self, x_out, x_in): - - for i in range(self.nStages): - self.solve_stage(x_in, i) - - x_out.assign(x_in) - for i in range(self.nStages): - x_out.assign(x_out + self.butcher_matrix[self.nStages, i]*self.dt*self.k[i]) - - if self.limiter is not None: - self.limiter.apply(x_out) - - class ExplicitTimeDiscretisation(TimeDiscretisation): """Base class for explicit time discretisations.""" @@ -1621,145 +1542,194 @@ def apply(self, x_out, *x_in): x_out.assign(self.x_out) -class ImplicitMidpoint(ImplicitMultistage): - u""" - Implements the Implicit Midpoint method as a 1-stage Runge Kutta method. - - The method, for solving - ∂y/∂t = F(y), can be written as: - - k0 = F[y^n + 0.5*dt*k0] - y^(n+1) = y^n + dt*k0 - """ - def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None, butcher_matrix=None): - super().__init__(domain, field_name, - solver_parameters=solver_parameters, - limiter=limiter, options=options) - self.butcher_matrix = np.array([[0.5], [1.]]) - self.nStages = int(np.shape(self.butcher_matrix)[1]) - - -class QinZhang(ImplicitMultistage): - u""" - Implements Qin and Zhang's two-stage, 2nd order, implicit Runge–Kutta method. - - The method, for solving - ∂y/∂t = F(y), can be written as: - - k0 = F[y^n + 0.25*dt*k0] - k1 = F[y^n + 0.5*dt*k0 + 0.25*dt*k1] - y^(n+1) = y^n + 0.5*dt*(k0 + k1) - """ - def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None, butcher_matrix=None): - super().__init__(domain, field_name, - solver_parameters=solver_parameters, - limiter=limiter, options=options) - self.butcher_matrix = np.array([[0.25, 0], [0.5, 0.25], [0.5, 0.5]]) - self.nStages = int(np.shape(self.butcher_matrix)[1]) - class IMEX_Euler(IMEXMultistage): u""" - Implements Qin and Zhang's two-stage, 2nd order, implicit Runge–Kutta method. + Implements IMEX Euler one-stage method. - The method, for solving - ∂y/∂t = F(y), can be written as: + The method, for solving \n + ∂y/∂t = F(y) + S(y), can be written as: \n - k0 = F[y^n + 0.25*dt*k0] - k1 = F[y^n + 0.5*dt*k0 + 0.25*dt*k1] - y^(n+1) = y^n + 0.5*dt*(k0 + k1) + y_0 = y^n \n + y_1 = y^n + dt*F[y_1] + dt*S[y_0] \n + y^(n+1) = y^n + dt*F[y_1] + dt*S[y_0] \n """ - def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None, butcher_imp=None, butcher_exp=None): - super().__init__(domain, field_name, + def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ + butcher_imp = np.array([[0., 0.], [0., 1.], [0., 1.]]) + butcher_exp = np.array([[0., 0.], [1., 0.], [1., 0.]]) + super().__init__(domain, butcher_imp, butcher_exp, field_name, solver_parameters=solver_parameters, limiter=limiter, options=options) - self.butcher_imp = np.array([[0., 0.], [0., 1.], [0., 1.]]) - self.butcher_exp = np.array([[0., 0.], [1., 0.], [1., 0.]]) - self.nStages = int(np.shape(self.butcher_imp)[1]) + class ARS3(IMEXMultistage): u""" - Implements Qin and Zhang's two-stage, 2nd order, implicit Runge–Kutta method. - - The method, for solving - ∂y/∂t = F(y), can be written as: - - k0 = F[y^n + 0.25*dt*k0] - k1 = F[y^n + 0.5*dt*k0 + 0.25*dt*k1] - y^(n+1) = y^n + 0.5*dt*(k0 + k1) + Implements ARS3(2,3,3) two-stage IMEX Runge–Kutta method + from RK IMEX for HEVI (Weller et al 2013). + Where g = (3 + sqrt(3))/6. + + The method, for solving \n + ∂y/∂t = F(y) + S(y), can be written as: \n + + y_0 = y^n \n + y_1 = y^n + dt*g*F[y_1] + dt*g*S[y_0] \n + y_2 = y^n + dt*((1-2g)*F[y_1]+g*F[y_2]) \n + + dt*((g-1)*S[y_0]+2(g-1)*S[y_1]) \n + y^(n+1) = y^n + dt*(g*F[y_1]+(1-g)*F[y_2]) \n + + dt*(0.5*S[y_1]+0.5*S[y_2]) \n """ - def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None, butcher_imp=None, butcher_exp=None): - super().__init__(domain, field_name, - solver_parameters=solver_parameters, - limiter=limiter, options=options) - # ARS3 + def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ g = (3. + np.sqrt(3.))/6. + butcher_imp = np.array([[0., 0., 0.], [0., g, 0.], [0., 1-2.*g, g], [0., 0.5, 0.5]]) + butcher_exp = np.array([[0., 0., 0.], [g, 0., 0.], [g-1., 2.*(1.-g), 0.], [0., 0.5, 0.5]]) - self.butcher_imp = np.array([[0., 0., 0.], [0., g, 0.], [0., 1-2.*g, g], [0., 0.5, 0.5]]) - self.butcher_exp = np.array([[0., 0., 0.], [g, 0., 0.], [g-1., 2.*(1.-g), 0.], [0., 0.5, 0.5]]) + super().__init__(domain, butcher_imp, butcher_exp, field_name, + solver_parameters=solver_parameters, + limiter=limiter, options=options) - self.nStages = int(np.shape(self.butcher_imp)[1]) class ARK2(IMEXMultistage): u""" - Implements Qin and Zhang's two-stage, 2nd order, implicit Runge–Kutta method. - - The method, for solving - ∂y/∂t = F(y), can be written as: - - k0 = F[y^n + 0.25*dt*k0] - k1 = F[y^n + 0.5*dt*k0 + 0.25*dt*k1] - y^(n+1) = y^n + 0.5*dt*(k0 + k1) + Implements ARK2(2,3,2) two-stage IMEX Runge–Kutta method from + RK IMEX for HEVI (Weller et al 2013). + Where g = 1 - 1/sqrt(2), a = 1/6(3 + 2sqrt(2)), d = 1/2sqrt(2). + + The method, for solving \n + ∂y/∂t = F(y) + S(y), can be written as: \n + + y_0 = y^n \n + y_1 = y^n + dt*(g*F[y_0]+g*F[y_1]) + 2*dt*g*S[y_0] \n + y_2 = y^n + dt*(d*F[y_0]+d*F[y_1]+g*F[y_2]) \n + + dt*((1-a)*S[y_0]+a*S[y_1]) \n + y^(n+1) = y^n + dt*(d*F[y_0]+d*F[y_1]+g*F[y_2]) \n + + dt*(d*S[y_0]+d*S[y_1]+g*S[y_2]) \n """ - def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None, butcher_imp=None, butcher_exp=None): - super().__init__(domain, field_name, - solver_parameters=solver_parameters, - limiter=limiter, options=options) - # ARK2 + def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ g = 1. - 1./np.sqrt(2.) d = 1./(2.*np.sqrt(2.)) a = 1./6.*(3. + 2.*np.sqrt(2.)) - self.butcher_imp = np.array([[0., 0., 0.], [g, g, 0.], [d, d, g], [d, d, g]]) - self.butcher_exp = np.array([[0., 0., 0.], [2.*g, 0., 0.], [1.-a, a, 0.], [d, d, g]]) - self.nStages = int(np.shape(self.butcher_imp)[1]) + butcher_imp = np.array([[0., 0., 0.], [g, g, 0.], [d, d, g], [d, d, g]]) + butcher_exp = np.array([[0., 0., 0.], [2.*g, 0., 0.], [1.-a, a, 0.], [d, d, g]]) + super().__init__(domain, butcher_imp, butcher_exp, field_name, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + class SSP3(IMEXMultistage): u""" - Implements Qin and Zhang's two-stage, 2nd order, implicit Runge–Kutta method. + Implements SSP3(3,3,2) three-stage IMEX Runge–Kutta method from RK IMEX for HEVI (Weller et al 2013). + Where g = 1 - 1/sqrt(2) - The method, for solving - ∂y/∂t = F(y), can be written as: + The method, for solving \n + ∂y/∂t = F(y) + S(y), can be written as: \n - k0 = F[y^n + 0.25*dt*k0] - k1 = F[y^n + 0.5*dt*k0 + 0.25*dt*k1] - y^(n+1) = y^n + 0.5*dt*(k0 + k1) + y_1 = y^n + dt*g*F[y_1] \n + y_2 = y^n + dt*((1-2g)*F[y_1]+g*F[y_2]) + dt*S[y_1] \n + y_3 = y^n + dt*((0.5-g)*F[y_1]+g*F[y_3]) + dt*(0.25*S[y_1]+0.25*S[y_2]) \n + y^(n+1) = y^n + dt*(1/6*F[y_1]+1/6*F[y_2]+2/3*F[y_3]) \n + + dt*(1/6*S[y_1]+1/6*S[y_2]+2/3*S[y_3]) \n """ - def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None, butcher_imp=None, butcher_exp=None): - super().__init__(domain, field_name, + def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ + g = 1. - (1./np.sqrt(2.)) + butcher_imp = np.array([[g, 0., 0.], [1-2.*g, g, 0.], [0.5-g, 0., g], [(1./6.), (1./6.), (2./3.)]]) + butcher_exp = np.array([[0., 0., 0.], [1., 0., 0.], [0.25, 0.25, 0.], [(1./6.), (1./6.), (2./3.)]]) + super().__init__(domain, butcher_imp, butcher_exp, field_name, solver_parameters=solver_parameters, limiter=limiter, options=options) - # SSP3(3,3,2) - g = 1. - (1./np.sqrt(2.)) - self.butcher_imp = np.array([[g, 0., 0.], [1-2.*g, g, 0.], [0.5-g, 0., g], [(1./6.),(1./6.), (2./3.)]]) - self.butcher_exp = np.array([[0., 0., 0.], [1., 0., 0.], [0.25, 0.25, 0.], [(1./6.),(1./6.), (2./3.)]]) - self.nStages = int(np.shape(self.butcher_imp)[1]) + class Trap2(IMEXMultistage): u""" - Implements Qin and Zhang's two-stage, 2nd order, implicit Runge–Kutta method. + Implements Trap2(2+e,3,2) three-stage IMEX Runge–Kutta method from RK IMEX for HEVI (Weller et al 2013). + For e = 1 or 0. - The method, for solving - ∂y/∂t = F(y), can be written as: + The method, for solving \n + ∂y/∂t = F(y) + S(y), can be written as: \n - k0 = F[y^n + 0.25*dt*k0] - k1 = F[y^n + 0.5*dt*k0 + 0.25*dt*k1] - y^(n+1) = y^n + 0.5*dt*(k0 + k1) + y_0 = y^n \n + y_1 = y^n + dt*e*F[y_0] + dt*S[y_0] \n + y_2 = y^n + dt*(0.5*F[y_0]+0.5*F[y_2]) + dt*(0.5*S[y_0]+0.5*S[y_1]) \n + y_3 = y^n + dt*(0.5*F[y_0]+0.5*F[y_3]) + dt*(0.5*S[y_0]+0.5*S[y_2]) \n + y^(n+1) = y^n + dt*(0.5*F[y_0]+0.5*F[y_3]) + dt*(0.5*S[y_0] + 0.5*S[y_2]) \n """ - def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None, butcher_imp=None, butcher_exp=None): - super().__init__(domain, field_name, + def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + field_name (str, optional): name of the field to be evolved. + Defaults to None. + solver_parameters (dict, optional): dictionary of parameters to + pass to the underlying solver. Defaults to None. + limiter (:class:`Limiter` object, optional): a limiter to apply to + the evolving field to enforce monotonicity. Defaults to None. + options (:class:`AdvectionOptions`, optional): an object containing + options to either be passed to the spatial discretisation, or + to control the "wrapper" methods, such as Embedded DG or a + recovery method. Defaults to None. + """ + e = 0. + butcher_imp = np.array([[0., 0., 0., 0.], [e, 0., 0., 0.], [0.5, 0., 0.5, 0.], [0.5, 0., 0., 0.5], [0.5, 0., 0., 0.5]]) + butcher_exp = np.array([[0., 0., 0., 0.], [1., 0., 0., 0.], [0.5, 0.5, 0., 0.], [0.5, 0., 0.5, 0.], [0.5, 0., 0.5, 0.]]) + super().__init__(domain, butcher_imp, butcher_exp, field_name, solver_parameters=solver_parameters, limiter=limiter, options=options) - # Trap2(2+e,3,2) - e = 0. - self.butcher_imp = np.array([[0., 0., 0.,0,], [e, 0., 0.,0.], [0.5, 0., 0.5, 0.], [0.5, 0., 0., 0.5], [0.5, 0., 0., 0.5]]) - self.butcher_exp= np.array([[0., 0., 0.,0,], [1., 0., 0.,0.], [0.5, 0.5, 0., 0.], [0.5, 0., 0.5, 0.], [0.5, 0., 0.5, 0.]]) - self.nStages = int(np.shape(self.butcher_imp)[1]) \ No newline at end of file From c8207bc330df9ab356c0c004ce0f5a47f99910d1 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Thu, 19 Oct 2023 16:05:29 +0100 Subject: [PATCH 08/22] revert explicitmultistage changes --- gusto/time_discretisation.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 8255ae4e9..73ed4b5cf 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -542,25 +542,23 @@ class ExplicitMultistage(ExplicitTimeDiscretisation): A Butcher tableau is formed in the following way for a s-th order explicit scheme: - c_0 | a_00 a_01 . a_0s - c_1 | a_10 a_11 a_1s + c_0 | a_00 + c_1 | a_10 a_11 + . | a_20 a_21 a_22 . | . . . . - . | . . . . - c_s | a_s0 a_s1 . a_ss ------------------------- - | b_1 b_2 ... b_s + c_s | b_1 b_2 ... b_s The gradient function has no time-dependence, so c elements are not needed. A square 'butcher_matrix' is defined in each class that uses the ExplicitMultiStage structure, - [a_10 0 . 0 ] - [a_20 a_21 . 0 ] + [a_00 0 . 0 ] + [a_10 a_11 . 0 ] [ . . . . ] - [ b_0 b_1 . b_s] + [ b_0 b_1 . b_{s-1}] - All upper diagonal a_ij elements are zero for explicit methods. We exclude the first - row of the butcher tableau from our butcher matrix as the row is always zeros. + All upper diagonal a_ij elements are zero for explicit methods. There are three steps to move from the current solution, y^n, to the new one, y^{n+1} From 78f5557b9e26c9395a0364cd482c7553213701e5 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Thu, 19 Oct 2023 16:07:36 +0100 Subject: [PATCH 09/22] revert test changes --- integration-tests/model/test_time_discretisation.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/integration-tests/model/test_time_discretisation.py b/integration-tests/model/test_time_discretisation.py index 6b96b401b..4912b26bf 100644 --- a/integration-tests/model/test_time_discretisation.py +++ b/integration-tests/model/test_time_discretisation.py @@ -8,8 +8,9 @@ def run(timestepper, tmax, f_end): return norm(timestepper.fields("f") - f_end) / norm(f_end) -@pytest.mark.parametrize("scheme", ["ssprk", "TrapeziumRule", "ImplicitMidpoint", - "RK4", "Heun", "BDF2", "TR_BDF2", "AdamsBashforth", "Leapfrog", "AdamsMoulton"]) +@pytest.mark.parametrize("scheme", ["ssprk", "TrapeziumRule", + "RK4", "Heun", "BDF2", "TR_BDF2", "AdamsBashforth", + "Leapfrog", "AdamsMoulton"]) def test_time_discretisation(tmpdir, scheme, tracer_setup): if (scheme == "AdamsBashforth"): # Tighter stability constraints @@ -29,8 +30,6 @@ def test_time_discretisation(tmpdir, scheme, tracer_setup): transport_scheme = SSPRK3(domain) elif scheme == "TrapeziumRule": transport_scheme = TrapeziumRule(domain) - elif scheme == "ImplicitMidpoint": - transport_scheme = ImplicitMidpoint(domain) elif scheme == "RK4": transport_scheme = RK4(domain) elif scheme == "Heun": From 177c308449724fac4cae08af878bf75d378f1091 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Wed, 25 Oct 2023 10:55:01 +0100 Subject: [PATCH 10/22] adressed most review comments --- gusto/time_discretisation.py | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 73ed4b5cf..1817d7417 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -287,18 +287,14 @@ def __init__(self, domain, butcher_imp, butcher_exp, field_name=None, Args: domain (:class:`Domain`): the model's domain object, containing the mesh and the compatible function spaces. - butcher_imp (numpy array): A matrix containing the coefficients of + butcher_imp (:class:`numpy.ndarray`): 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 + butcher_exp (:class:`numpy.ndarray`): A matrix containing the coefficients of 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. - Defaults to None. solver_parameters (dict, optional): dictionary of parameters to pass to the underlying solver. Defaults to None. - limiter (:class:`Limiter` object, optional): a limiter to apply to - the evolving field to enforce monotonicity. Defaults to None. options (:class:`AdvectionOptions`, optional): an object containing options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a @@ -306,7 +302,7 @@ def __init__(self, domain, butcher_imp, butcher_exp, field_name=None, """ super().__init__(domain, field_name=field_name, solver_parameters=solver_parameters, - limiter=limiter, options=options) + options=options) self.butcher_imp = butcher_imp self.butcher_exp = butcher_exp self.nStages = int(np.shape(self.butcher_imp)[1]) @@ -344,6 +340,8 @@ def rhs(self): return super(IMEXMultistage, self).rhs def res(self, stage): + """Set up the discretisation's residual for a given stage.""" + # Add time derivative terms y_s - y^n for stage s mass_form = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_false=drop) @@ -351,6 +349,10 @@ def res(self, stage): 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)) + # Loop through stages up to s-1 and calcualte/sum + # dt*(a_s1*F(y_1) + a_s2*F(y_2)+ ... + a_{s,s-1}*F(y_{s-1})) + # and + # dt*(d_s1*S(y_1) + d_s2*S(y_2)+ ... + d_{s,s-1}*S(y_{s-1})) for i in range(stage): r_exp = self.residual.label_map( lambda t: t.has_label(explicit), @@ -368,6 +370,7 @@ def res(self, stage): map_if_false=lambda t: Constant(self.butcher_imp[stage, i])*self.dt*t) residual += r_imp residual += r_exp + # Calculate and add on dt*a_ss*F(y_s) r_imp = self.residual.label_map( lambda t: t.has_label(implicit), map_if_true=replace_subject(self.x_out, old_idx=self.idx), @@ -380,12 +383,18 @@ def res(self, stage): @property def final_res(self): + """Set up the discretisation's final residual.""" + # Add time derivative terms y^{n+1} - y^n mass_form = self.residual.label_map(lambda t: t.has_label(time_derivative), map_if_false=drop) residual = mass_form.label_map(all_terms, 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)) + # Loop through stages up to s-1 and calcualte/sum + # dt*(b_1*F(y_1) + b_2*F(y_2) + .... + b_s*F(y_s)) + # and + # dt*(e_1*S(y_1) + e_2*S(y_2) + .... + e_s*S(y_s)) for i in range(self.nStages): r_exp = self.residual.label_map( lambda t: t.has_label(explicit), @@ -407,6 +416,7 @@ def final_res(self): @cached_property def solvers(self): + """Set up a list of solvers for each problem at a stage.""" solvers = [] for stage in range(self.nStages): # setup solver using residual defined in derived class @@ -417,6 +427,7 @@ def solvers(self): @cached_property def final_solver(self): + """Set up a solver for the final solve to evaluate time level n+1.""" # setup solver using lhs and rhs defined in derived class problem = NonlinearVariationalProblem(self.final_res, self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__ From 3a34def092cb336a0adb1241348af0baa72f2e26 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Wed, 25 Oct 2023 15:38:51 +0100 Subject: [PATCH 11/22] Warning for explicit implicit splitting added --- gusto/time_discretisation.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 1817d7417..50f8e74e3 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -260,7 +260,7 @@ class IMEXMultistage(TimeDiscretisation): + dt*(e_1*S(y_1) + e_2*S(y_2) + .... + e_s*S(y_s)) \n """ - # --------------------------------------------------------------------------- + # -------------------------------------------------------------------------- # Butcher tableaus for a s-th order # diagonally implicit scheme (left) and explicit scheme (right): # c_0 | a_00 0 . 0 f_0 | 0 0 . 0 @@ -279,7 +279,8 @@ class IMEXMultistage(TimeDiscretisation): # [a_20 a_21 a_22 . 0 ] [d_20 d_21 0 . 0 ] # [ . . . . . ] [ . . . . . ] # [ b_0 b_1 . b_s] [ e_0 e_1 . . e_s] - # --------------------------------------------------------------------------- + # + # -------------------------------------------------------------------------- def __init__(self, domain, butcher_imp, butcher_exp, field_name=None, solver_parameters=None, limiter=None, options=None): @@ -326,6 +327,8 @@ def setup(self, equation, apply_bcs=True, *active_labels): self.residual = self.residual.label_map( lambda t: t.has_label(transport), map_if_true=lambda t: explicit(t)) + + logger.warning("Default IMEX Multistage treats transport terms explicitly, and all other terms implicitly") self.xs = [Function(self.fs) for i in range(self.nStages)] @@ -341,7 +344,7 @@ def rhs(self): def res(self, stage): """Set up the discretisation's residual for a given stage.""" - # Add time derivative terms y_s - y^n for stage s + # Add time derivative terms y_s - y^n for stage s mass_form = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_false=drop) @@ -384,7 +387,7 @@ def res(self, stage): @property def final_res(self): """Set up the discretisation's final residual.""" - # Add time derivative terms y^{n+1} - y^n + # Add time derivative terms y^{n+1} - y^n mass_form = self.residual.label_map(lambda t: t.has_label(time_derivative), map_if_false=drop) residual = mass_form.label_map(all_terms, From 564cafba73fc15173881eb614011eb4e01c03267 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Thu, 26 Oct 2023 16:17:46 +0100 Subject: [PATCH 12/22] Comments added and equations refactor moved into time_discretisation. Code currently does not work --- gusto/equations.py | 12 ++---------- gusto/time_discretisation.py | 35 ++++++++++++++++++++++++++++++----- 2 files changed, 32 insertions(+), 15 deletions(-) diff --git a/gusto/equations.py b/gusto/equations.py index d97a992dc..5a32cc430 100644 --- a/gusto/equations.py +++ b/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, conservative_depth=True): + thermal=False): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -676,10 +676,7 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, raise ValueError("Invalid u_transport_option: %s" % u_transport_option) # Depth transport term - if (conservative_depth): - D_adv = prognostic(continuity_form(phi, D, u), 'D') - else: - D_adv = prognostic(advection_form(phi, D, u), 'D') + D_adv = prognostic(continuity_form(phi, D, u), 'D') # Transport term needs special linearisation if self.linearisation_map(D_adv.terms[0]): @@ -711,11 +708,6 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, residual = (mass_form + adv_form + pressure_gradient_form) - # Add divergence term - if (not conservative_depth): - geo_grad_form = subject(prognostic(phi*D*div(u)*dx), self.X) - residual += geo_grad_form - # -------------------------------------------------------------------- # # Extra Terms (Coriolis, Topography and Thermal) # -------------------------------------------------------------------- # diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 50f8e74e3..df2fad6ca 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -7,15 +7,18 @@ from abc import ABCMeta, abstractmethod, abstractproperty from firedrake import (Function, TestFunction, NonlinearVariationalProblem, - NonlinearVariationalSolver, DirichletBC, Constant) + NonlinearVariationalSolver, DirichletBC, Constant, split, + div, dx) from firedrake.formmanipulation import split_form from firedrake.utils import cached_property -from gusto.configuration import EmbeddedDGOptions, RecoveryOptions +from gusto.configuration import EmbeddedDGOptions, RecoveryOptions, TransportEquationType from gusto.fml import ( - replace_subject, replace_test_function, Term, all_terms, drop + replace_subject, replace_test_function, Term, all_terms, drop, subject ) -from gusto.labels import time_derivative, prognostic, physics_label, transport, implicit, explicit +from gusto.labels import (time_derivative, prognostic, physics_label, + transport, implicit, explicit, transporting_velocity) +from gusto.common_forms import advection_form from gusto.logging import logger, DEBUG, logging_ksp_monitor_true_residual from gusto.wrappers import * import numpy as np @@ -252,7 +255,7 @@ class IMEXMultistage(TimeDiscretisation): For each i = 1, s in an s stage method we compute the intermediate solutions: \n - y_i = y^n + dt*(a_i1*F(y_1) + a_i2*F(y_2)+ ... + a_ii*F(y_i)) \n + y_i = y^n + dt*(a_i1*F(y_1) + a_i2*F(y_2)+ ... + a_ii*F(y_i)) \n + dt*(d_i1*S(y_1) + d_i2*S(y_2)+ ... + d_{i,i-1}*S(y_{i-1})) At the last stage, compute the new solution by: \n @@ -320,6 +323,28 @@ 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: + if(t.get(transport) == TransportEquationType.conservative): + # Split continuity form term + test = t.form.arguments()[0] + subj = t.get(subject) + prognostic_field_name = t.get(prognostic) + idx = self.equation.field_names.index(prognostic_field_name) + transported_field = split(subj)[idx] + # u_idx = self.equation.field_names.index('u') + # uadv = split(self.equation.X)[u_idx] + # breakpoint() + uadv = t.get(transporting_velocity) + breakpoint() + new_transport_term = prognostic(subject(advection_form(test, transported_field, uadv) + test*transported_field*div(uadv)*dx, subj, prognostic_field_name)) + # Add onto residual and drop old term + self.residual = self.residual.label_map( + lambda t: t.get(transport) == TransportEquationType.conservative, + map_if_true=drop) + self.residual += new_transport_term.form + + # Label transport terms as explicit, all other terms as implicit self.residual = self.residual.label_map( lambda t: any(t.has_label(time_derivative, transport)), map_if_false=lambda t: implicit(t)) From 31ddf80dabb5f156d9ca3b8bc547b60fea47c45f Mon Sep 17 00:00:00 2001 From: atb1995 Date: Fri, 27 Oct 2023 14:00:33 +0100 Subject: [PATCH 13/22] Splitting on conservative term added. Need to make it work in case of linear equation set.. --- gusto/time_discretisation.py | 47 ++++++++++++++++++------------------ 1 file changed, 24 insertions(+), 23 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index df2fad6ca..2cfe9d2d2 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -6,9 +6,9 @@ """ from abc import ABCMeta, abstractmethod, abstractproperty -from firedrake import (Function, TestFunction, NonlinearVariationalProblem, - NonlinearVariationalSolver, DirichletBC, Constant, split, - div, dx) +from firedrake import (Function, TestFunction, TestFunctions, + NonlinearVariationalProblem, NonlinearVariationalSolver, + DirichletBC, Constant, split, div, dx) from firedrake.formmanipulation import split_form from firedrake.utils import cached_property @@ -17,7 +17,7 @@ replace_subject, replace_test_function, Term, all_terms, drop, subject ) from gusto.labels import (time_derivative, prognostic, physics_label, - transport, implicit, explicit, transporting_velocity) + transport, implicit, explicit) from gusto.common_forms import advection_form from gusto.logging import logger, DEBUG, logging_ksp_monitor_true_residual from gusto.wrappers import * @@ -325,24 +325,25 @@ def setup(self, equation, apply_bcs=True, *active_labels): # Get continuity form transport term for t in self.residual: - if(t.get(transport) == TransportEquationType.conservative): - # Split continuity form term - test = t.form.arguments()[0] - subj = t.get(subject) - prognostic_field_name = t.get(prognostic) - idx = self.equation.field_names.index(prognostic_field_name) - transported_field = split(subj)[idx] - # u_idx = self.equation.field_names.index('u') - # uadv = split(self.equation.X)[u_idx] - # breakpoint() - uadv = t.get(transporting_velocity) - breakpoint() - new_transport_term = prognostic(subject(advection_form(test, transported_field, uadv) + test*transported_field*div(uadv)*dx, subj, prognostic_field_name)) - # Add onto residual and drop old term - self.residual = self.residual.label_map( - lambda t: t.get(transport) == TransportEquationType.conservative, - map_if_true=drop) - self.residual += new_transport_term.form + if (t.get(transport) == TransportEquationType.conservative): + # Split continuity form term + subj = t.get(subject) + prognostic_field_name = t.get(prognostic) + idx = self.equation.field_names.index(prognostic_field_name) + W = self.fs + test = TestFunctions(W)[idx] + transported_field = split(subj)[idx] + u_idx = self.equation.field_names.index('u') + uadv = split(self.equation.X)[u_idx] + new_transport_term = prognostic(advection_form(test, transported_field, uadv), prognostic_field_name) + div_term = prognostic(test*transported_field*div(uadv)*dx, prognostic_field_name) + # Add onto residual + self.residual += subject(new_transport_term + div_term, subj) + + # Drop old term + self.residual = self.residual.label_map( + lambda t: t.get(transport) == TransportEquationType.conservative, + map_if_true=drop) # Label transport terms as explicit, all other terms as implicit self.residual = self.residual.label_map( @@ -352,7 +353,7 @@ def setup(self, equation, apply_bcs=True, *active_labels): self.residual = self.residual.label_map( lambda t: t.has_label(transport), map_if_true=lambda t: explicit(t)) - + logger.warning("Default IMEX Multistage treats transport terms explicitly, and all other terms implicitly") self.xs = [Function(self.fs) for i in range(self.nStages)] From 85e57a74d434353693b21899370f62c74902fca0 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Fri, 27 Oct 2023 17:35:15 +0100 Subject: [PATCH 14/22] Lint fix --- gusto/time_discretisation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index aaaff6ffd..8c6cb353f 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -27,7 +27,7 @@ __all__ = ["ForwardEuler", "BackwardEuler", "ExplicitMultistage", "IMEXMultistage", "SSPRK3", "RK4", "Heun", "ThetaMethod", - "TrapeziumRule", "BDF2", "TR_BDF2", "Leapfrog","AdamsMoulton", + "TrapeziumRule", "BDF2", "TR_BDF2", "Leapfrog", "AdamsMoulton", "AdamsBashforth", "ImplicitMidpoint", "QinZhang", "IMEX_Euler", "ARS3", "ARK2", "Trap2", "SSP3"] From 10ed951fc7960d1c619643cec56fa352e0eb501e Mon Sep 17 00:00:00 2001 From: atb1995 Date: Wed, 1 Nov 2023 10:25:49 +0000 Subject: [PATCH 15/22] Labelling of implicit and explicit terms moved outside of time_disc, splitting of terms added to common_forms, new test for IMEX schemes added --- gusto/common_forms.py | 58 ++++++++++++++++++++++++++-- gusto/time_discretisation.py | 47 +++++----------------- integration-tests/model/test_IMEX.py | 41 ++++++++++++++++++++ 3 files changed, 106 insertions(+), 40 deletions(-) create mode 100644 integration-tests/model/test_IMEX.py diff --git a/gusto/common_forms.py b/gusto/common_forms.py index bb9de1d94..5d0d2b740 100644 --- a/gusto/common_forms.py +++ b/gusto/common_forms.py @@ -2,13 +2,17 @@ Provides some basic forms for discretising various common terms in equations for geophysical fluid dynamics.""" -from firedrake import dx, dot, grad, div, inner, outer, cross, curl +from firedrake import (dx, dot, grad, div, inner, outer, cross, curl, split, + TestFunctions) from gusto.configuration import TransportEquationType -from gusto.labels import transport, transporting_velocity, diffusion +from gusto.labels import (transport, transporting_velocity, diffusion, + prognostic) +from gusto.fml import subject, drop __all__ = ["advection_form", "continuity_form", "vector_invariant_form", "kinetic_energy_form", "advection_equation_circulation_form", - "diffusion_form", "linear_advection_form", "linear_continuity_form"] + "diffusion_form", "linear_advection_form", "linear_continuity_form", + "split_continuity_form"] def advection_form(test, q, ubar): @@ -194,3 +198,51 @@ def diffusion_form(test, q, kappa): form = -inner(test, div(kappa*grad(q)))*dx return diffusion(form) + + +def split_continuity_form(equation): + u""" + Loops through terms in a given equation, and splits continuity terms into + advective and divergence terms. + + This describes splitting ∇.(u*q) into u.∇q and q(∇.u), + for transporting velocity u and transported q. + + Args: + equation (:class:`PrognosticEquation`): the model's equation. + + Returns: + equation (:class:`PrognosticEquation`): the model's equation. + """ + + for t in equation.residual: + if (t.get(transport) == TransportEquationType.conservative): + # Split continuity form term + subj = t.get(subject) + prognostic_field_name = t.get(prognostic) + if hasattr(equation, "field_names"): + idx = equation.field_names.index(prognostic_field_name) + else: + idx = equation.field_name.index(prognostic_field_name) + W = equation.function_space + test = TestFunctions(W)[idx] + q = split(subj)[idx] + # u is either a prognostic or prescribed field + if (hasattr(equation, "field_names") + and 'u' in equation.field_names): + u_idx = equation.field_names.index('u') + uadv = split(equation.X)[u_idx] + elif 'u' in equation.prescribed_fields._field_names: + uadv = equation.prescribed_fields('u') + else: + raise ValueError('Cannot get velocity field') + adv_term = prognostic(advection_form(test, q, uadv), prognostic_field_name) + div_term = prognostic(test*q*div(uadv)*dx, prognostic_field_name) + # Add onto residual + equation.residual += subject(adv_term + div_term, subj) + # Drop old term + equation.residual = equation.residual.label_map( + lambda t: t.get(transport) == TransportEquationType.conservative, + map_if_true=drop) + + return equation diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 8c6cb353f..2e77450fb 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -6,19 +6,18 @@ """ from abc import ABCMeta, abstractmethod, abstractproperty -from firedrake import (Function, TestFunction, TestFunctions, +from firedrake import (Function, TestFunction, NonlinearVariationalProblem, NonlinearVariationalSolver, - DirichletBC, Constant, split, div, dx) + DirichletBC, Constant, split) from firedrake.formmanipulation import split_form from firedrake.utils import cached_property -from gusto.configuration import EmbeddedDGOptions, RecoveryOptions, TransportEquationType +from gusto.configuration import EmbeddedDGOptions, RecoveryOptions from gusto.fml import ( - replace_subject, replace_test_function, Term, all_terms, drop, subject + replace_subject, replace_test_function, Term, all_terms, drop ) from gusto.labels import (time_derivative, prognostic, physics_label, - transport, implicit, explicit) -from gusto.common_forms import advection_form + implicit, explicit) from gusto.logging import logger, DEBUG, logging_ksp_monitor_true_residual from gusto.wrappers import * import math @@ -329,38 +328,12 @@ def setup(self, equation, apply_bcs=True, *active_labels): super().setup(equation, apply_bcs, *active_labels) - # Get continuity form transport term + # Check all terms are labeled implicit, exlicit for t in self.residual: - if (t.get(transport) == TransportEquationType.conservative): - # Split continuity form term - subj = t.get(subject) - prognostic_field_name = t.get(prognostic) - idx = self.equation.field_names.index(prognostic_field_name) - W = self.fs - test = TestFunctions(W)[idx] - transported_field = split(subj)[idx] - u_idx = self.equation.field_names.index('u') - uadv = split(self.equation.X)[u_idx] - new_transport_term = prognostic(advection_form(test, transported_field, uadv), prognostic_field_name) - div_term = prognostic(test*transported_field*div(uadv)*dx, prognostic_field_name) - # Add onto residual - self.residual += subject(new_transport_term + div_term, subj) - - # Drop old term - self.residual = self.residual.label_map( - lambda t: t.get(transport) == TransportEquationType.conservative, - map_if_true=drop) - - # Label transport terms as explicit, all other terms as implicit - self.residual = self.residual.label_map( - lambda t: any(t.has_label(time_derivative, transport)), - map_if_false=lambda t: implicit(t)) - - self.residual = self.residual.label_map( - lambda t: t.has_label(transport), - map_if_true=lambda t: explicit(t)) - - logger.warning("Default IMEX Multistage treats transport terms explicitly, and all other terms implicitly") + if ((not t.has_label(implicit)) and (not t.has_label(explicit)) + and (not t.has_label(time_derivative))): + logger.error("Non time-derivative terms must be labeled as implicit or explicit") + raise NotImplementedError self.xs = [Function(self.fs) for i in range(self.nStages)] diff --git a/integration-tests/model/test_IMEX.py b/integration-tests/model/test_IMEX.py new file mode 100644 index 000000000..9320a70cb --- /dev/null +++ b/integration-tests/model/test_IMEX.py @@ -0,0 +1,41 @@ +from firedrake import norm +from gusto import * +import pytest + + +def run(timestepper, tmax, f_end): + timestepper.run(0, tmax) + return norm(timestepper.fields("f") - f_end) / norm(f_end) + +@pytest.mark.parametrize("scheme", ["ssp3","ark2","ars3", "trap2", "euler"]) +def test_time_discretisation(tmpdir, scheme, tracer_setup): + + setup = tracer_setup(tmpdir, "sphere") + domain = setup.domain + V = domain.spaces("DG") + eqn = ContinuityEquation(domain, V, "f") + # Split continuity term + eqn = split_continuity_form(eqn) + #Label terms are implicit and explicit + eqn.label_terms(lambda t: not any(t.has_label(time_derivative, transport)), implicit) + eqn.label_terms(lambda t: t.has_label(transport), explicit) + + if scheme == "ssp3": + transport_scheme = SSP3(domain) + elif scheme == "ark2": + transport_scheme = ARK2(domain) + elif scheme == "ars3": + transport_scheme = ARS3(domain) + elif scheme == "trap2": + transport_scheme = Trap2(domain) + elif scheme == "euler": + transport_scheme = IMEX_Euler(domain) + + transport_method = DGUpwind(eqn, "f") + + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io, transport_method) + + # Initial conditions + timestepper.fields("f").interpolate(setup.f_init) + timestepper.fields("u").project(setup.uexpr) + assert run(timestepper, setup.tmax, setup.f_end) < setup.tol From e8661933b570332eac81537b592084c9b52bd096 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Wed, 1 Nov 2023 10:30:01 +0000 Subject: [PATCH 16/22] change error message --- gusto/time_discretisation.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 2e77450fb..44b8bc4dc 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -332,8 +332,7 @@ def setup(self, equation, apply_bcs=True, *active_labels): for t in self.residual: if ((not t.has_label(implicit)) and (not t.has_label(explicit)) and (not t.has_label(time_derivative))): - logger.error("Non time-derivative terms must be labeled as implicit or explicit") - raise NotImplementedError + raise NotImplementedError("Non time-derivative terms must be labeled as implicit or explicit") self.xs = [Function(self.fs) for i in range(self.nStages)] From 5878f23f6268cc3e3bef65afb7348916359752fd Mon Sep 17 00:00:00 2001 From: atb1995 Date: Wed, 1 Nov 2023 10:52:31 +0000 Subject: [PATCH 17/22] lint fix + firedrake fml move --- gusto/common_forms.py | 2 +- integration-tests/model/test_IMEX.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/gusto/common_forms.py b/gusto/common_forms.py index 5d0d2b740..17183ae02 100644 --- a/gusto/common_forms.py +++ b/gusto/common_forms.py @@ -4,10 +4,10 @@ from firedrake import (dx, dot, grad, div, inner, outer, cross, curl, split, TestFunctions) +from firedrake.fml import subject, drop from gusto.configuration import TransportEquationType from gusto.labels import (transport, transporting_velocity, diffusion, prognostic) -from gusto.fml import subject, drop __all__ = ["advection_form", "continuity_form", "vector_invariant_form", "kinetic_energy_form", "advection_equation_circulation_form", diff --git a/integration-tests/model/test_IMEX.py b/integration-tests/model/test_IMEX.py index 9320a70cb..d4ce82c22 100644 --- a/integration-tests/model/test_IMEX.py +++ b/integration-tests/model/test_IMEX.py @@ -7,7 +7,8 @@ def run(timestepper, tmax, f_end): timestepper.run(0, tmax) return norm(timestepper.fields("f") - f_end) / norm(f_end) -@pytest.mark.parametrize("scheme", ["ssp3","ark2","ars3", "trap2", "euler"]) + +@pytest.mark.parametrize("scheme", ["ssp3", "ark2", "ars3", "trap2", "euler"]) def test_time_discretisation(tmpdir, scheme, tracer_setup): setup = tracer_setup(tmpdir, "sphere") @@ -16,7 +17,7 @@ def test_time_discretisation(tmpdir, scheme, tracer_setup): eqn = ContinuityEquation(domain, V, "f") # Split continuity term eqn = split_continuity_form(eqn) - #Label terms are implicit and explicit + # Label terms are implicit and explicit eqn.label_terms(lambda t: not any(t.has_label(time_derivative, transport)), implicit) eqn.label_terms(lambda t: t.has_label(transport), explicit) From 958e002c69ae50afcc4136d756c4f30ad797a0ac Mon Sep 17 00:00:00 2001 From: atb1995 Date: Wed, 1 Nov 2023 11:45:19 +0000 Subject: [PATCH 18/22] Error switched to Runtime --- gusto/time_discretisation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 2cdfb4b55..79c8af94e 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -95,7 +95,7 @@ def __init__(self, domain, field_name=None, solver_parameters=None, elif self.wrapper_name == "supg": self.wrapper = SUPGWrapper(self, options) else: - raise NotImplementedError( + raise RuntimeError( f'Time discretisation: wrapper {self.wrapper_name} not implemented') else: self.wrapper = None From 6daae5075e22a21468ab54d0518ea28d22042920 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Wed, 1 Nov 2023 13:35:51 +0000 Subject: [PATCH 19/22] linearisation terms added --- gusto/common_forms.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/gusto/common_forms.py b/gusto/common_forms.py index 17183ae02..f389c9091 100644 --- a/gusto/common_forms.py +++ b/gusto/common_forms.py @@ -3,11 +3,11 @@ geophysical fluid dynamics.""" from firedrake import (dx, dot, grad, div, inner, outer, cross, curl, split, - TestFunctions) + TestFunctions, TrialFunction) from firedrake.fml import subject, drop from gusto.configuration import TransportEquationType from gusto.labels import (transport, transporting_velocity, diffusion, - prognostic) + prognostic, linearisation) __all__ = ["advection_form", "continuity_form", "vector_invariant_form", "kinetic_energy_form", "advection_equation_circulation_form", @@ -217,7 +217,7 @@ def split_continuity_form(equation): for t in equation.residual: if (t.get(transport) == TransportEquationType.conservative): - # Split continuity form term + # Get fields and test functions subj = t.get(subject) prognostic_field_name = t.get(prognostic) if hasattr(equation, "field_names"): @@ -236,9 +236,23 @@ def split_continuity_form(equation): uadv = equation.prescribed_fields('u') else: raise ValueError('Cannot get velocity field') + + # Create new advective and divergence terms adv_term = prognostic(advection_form(test, q, uadv), prognostic_field_name) div_term = prognostic(test*q*div(uadv)*dx, prognostic_field_name) - # Add onto residual + + # Add linearisations of new terms if required + if (t.has_label(linearisation)): + u_trial = TrialFunction(W)[u_idx] + qbar = split(equation.X_ref)[idx] + # Add linearisation to adv_term + linear_adv_term = linear_advection_form(test, qbar, u_trial) + adv_term = linearisation(adv_term, linear_adv_term) + # Add linearisation to div_term + linear_div_term = transporting_velocity(qbar*test*div(u_trial)*dx, u_trial) + div_term = linearisation(div_term, linear_div_term) + + # Add new terms onto residual equation.residual += subject(adv_term + div_term, subj) # Drop old term equation.residual = equation.residual.label_map( From 91919014f6396df8b01463fe1f8f6ec3ccec64b9 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Wed, 1 Nov 2023 14:30:06 +0000 Subject: [PATCH 20/22] bug fix to test func / q --- gusto/common_forms.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/gusto/common_forms.py b/gusto/common_forms.py index f389c9091..e385600fb 100644 --- a/gusto/common_forms.py +++ b/gusto/common_forms.py @@ -3,7 +3,7 @@ geophysical fluid dynamics.""" from firedrake import (dx, dot, grad, div, inner, outer, cross, curl, split, - TestFunctions, TrialFunction) + TestFunction, TestFunctions, TrialFunction) from firedrake.fml import subject, drop from gusto.configuration import TransportEquationType from gusto.labels import (transport, transporting_velocity, diffusion, @@ -222,11 +222,13 @@ def split_continuity_form(equation): prognostic_field_name = t.get(prognostic) if hasattr(equation, "field_names"): idx = equation.field_names.index(prognostic_field_name) + W = equation.function_space + test = TestFunctions(W)[idx] + q = split(subj)[idx] else: - idx = equation.field_name.index(prognostic_field_name) - W = equation.function_space - test = TestFunctions(W)[idx] - q = split(subj)[idx] + W = equation.function_space + test = TestFunction(W) + q = subj # u is either a prognostic or prescribed field if (hasattr(equation, "field_names") and 'u' in equation.field_names): From d29d1373f9d7ebea71a12e6efd83607ee3972361 Mon Sep 17 00:00:00 2001 From: atb1995 <81297297+atb1995@users.noreply.github.com> Date: Tue, 7 Nov 2023 16:54:10 +0000 Subject: [PATCH 21/22] Update gusto/common_forms.py Co-authored-by: tommbendall --- gusto/common_forms.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gusto/common_forms.py b/gusto/common_forms.py index e385600fb..0a36ad5ed 100644 --- a/gusto/common_forms.py +++ b/gusto/common_forms.py @@ -212,7 +212,7 @@ def split_continuity_form(equation): equation (:class:`PrognosticEquation`): the model's equation. Returns: - equation (:class:`PrognosticEquation`): the model's equation. + :class:`PrognosticEquation`: the model's equation. """ for t in equation.residual: From 7a14129a4850033fe64015711ff1376f5853f0b3 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Tue, 7 Nov 2023 16:56:27 +0000 Subject: [PATCH 22/22] Change in description --- gusto/common_forms.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/gusto/common_forms.py b/gusto/common_forms.py index 0a36ad5ed..67e3b1e96 100644 --- a/gusto/common_forms.py +++ b/gusto/common_forms.py @@ -202,10 +202,10 @@ def diffusion_form(test, q, kappa): def split_continuity_form(equation): u""" - Loops through terms in a given equation, and splits continuity terms into - advective and divergence terms. + Loops through terms in a given equation, and splits all continuity terms + into advective and divergence terms. - This describes splitting ∇.(u*q) into u.∇q and q(∇.u), + This describes splitting ∇.(u*q) terms into u.∇q and q(∇.u), for transporting velocity u and transported q. Args: