diff --git a/gusto/common_forms.py b/gusto/common_forms.py index bb9de1d94..67e3b1e96 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, + TestFunction, TestFunctions, TrialFunction) +from firedrake.fml import subject, drop from gusto.configuration import TransportEquationType -from gusto.labels import transport, transporting_velocity, diffusion +from gusto.labels import (transport, transporting_velocity, diffusion, + prognostic, linearisation) __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,67 @@ 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 all continuity terms + into advective and divergence terms. + + This describes splitting ∇.(u*q) terms into u.∇q and q(∇.u), + for transporting velocity u and transported q. + + Args: + equation (:class:`PrognosticEquation`): the model's equation. + + Returns: + :class:`PrognosticEquation`: the model's equation. + """ + + for t in equation.residual: + if (t.get(transport) == TransportEquationType.conservative): + # Get fields and test functions + subj = t.get(subject) + 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: + 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): + 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') + + # 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 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( + lambda t: t.get(transport) == TransportEquationType.conservative, + map_if_true=drop) + + return equation diff --git a/gusto/equations.py b/gusto/equations.py index c49d28ccf..bd61e795b 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -682,6 +682,7 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, # Depth transport term D_adv = prognostic(continuity_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) diff --git a/gusto/labels.py b/gusto/labels.py index 5aa658d85..ec0fb4af4 100644 --- a/gusto/labels.py +++ b/gusto/labels.py @@ -90,6 +90,8 @@ def __call__(self, target, value=None): # ---------------------------------------------------------------------------- # time_derivative = Label("time_derivative") +implicit = Label("implicit") +explicit = Label("explicit") transport = Label("transport", validator=lambda value: type(value) == TransportEquationType) diffusion = Label("diffusion") transporting_velocity = Label("transporting_velocity", validator=lambda value: type(value) in [Function, ufl.tensors.ListTensor]) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index c8e237428..79c8af94e 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -20,14 +20,17 @@ from firedrake.utils import cached_property from gusto.configuration import EmbeddedDGOptions, RecoveryOptions -from gusto.labels import time_derivative, prognostic, physics_label +from gusto.labels import (time_derivative, prognostic, physics_label, + implicit, explicit) from gusto.logging import logger, DEBUG, logging_ksp_monitor_true_residual from gusto.wrappers import * -__all__ = ["ForwardEuler", "BackwardEuler", "ExplicitMultistage", "ImplicitMultistage", - "SSPRK3", "RK4", "Heun", "ThetaMethod", "TrapeziumRule", "BDF2", "TR_BDF2", - "Leapfrog", "AdamsMoulton", "AdamsBashforth", "ImplicitMidpoint", "QinZhang"] +__all__ = ["ForwardEuler", "BackwardEuler", "ExplicitMultistage", + "IMEXMultistage", "SSPRK3", "RK4", "Heun", "ThetaMethod", + "TrapeziumRule", "BDF2", "TR_BDF2", "Leapfrog", "AdamsMoulton", + "AdamsBashforth", "ImplicitMidpoint", "QinZhang", + "IMEX_Euler", "ARS3", "ARK2", "Trap2", "SSP3"] def wrapper_apply(original_apply): @@ -92,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 @@ -246,58 +249,63 @@ def apply(self, x_out, x_in): pass -class ImplicitMultistage(TimeDiscretisation): +class IMEXMultistage(TimeDiscretisation): """ - A class for implementing general diagonally implicit multistage (Runge-Kutta) - methods based on its Butcher tableau. + A class for implementing general IMEX multistage (Runge-Kutta) + methods based on two Butcher tableaus, to solve \n - Unlike the explicit method, all upper diagonal a_ij elements are non-zero for implicit methods. + ∂y/∂t = F(y) + S(y) \n + + 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 have the intermediate solutions: \n - y_i = y^n + dt*(a_i1*k_1 + a_i2*k_2 + ... + a_ii*k_i) \n - We compute the gradient at the intermediate location, k_i = F(y_i) \n + 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*k_1 + b_2*k_2 + .... + b_s*k_s) \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 tableau for a s-th order - # diagonally implicit scheme: - # c_0 | a_00 0 . 0 - # c_1 | a_10 a_11 . 0 - # . | . . . . - # . | . . . . - # c_s | a_s0 a_s1 . a_ss - # ------------------------- - # | b_1 b_2 ... b_s + # -------------------------------------------------------------------------- + # 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_matrix' is: + # The corresponding square 'butcher_imp' and 'butcher_exp' matrices are: # - # [a_00 0 . 0 ] - # [a_10 a_11 . 0 ] - # [ . . . . ] - # [ b_0 b_1 . b_s] - # --------------------------------------------------------------------------- + # [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_matrix, field_name=None, - solver_parameters=None, limiter=None, options=None,): + 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_matrix (numpy array): A matrix containing the coefficients of - a butcher tableau defining a given Runge Kutta time discretisation. + butcher_imp (:class:`numpy.ndarray`): A matrix containing the coefficients of + a butcher tableau defining a given implicit Runge Kutta time discretisation. + 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. 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 @@ -305,9 +313,10 @@ def __init__(self, domain, butcher_matrix, field_name=None, """ super().__init__(domain, field_name=field_name, solver_parameters=solver_parameters, - limiter=limiter, options=options) - self.butcher_matrix = butcher_matrix - self.nStages = int(np.shape(self.butcher_matrix)[1]) + 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): """ @@ -321,67 +330,129 @@ def setup(self, equation, apply_bcs=True, *active_labels): super().setup(equation, apply_bcs, *active_labels) - self.k = [Function(self.fs) for i in range(self.nStages)] + # Check all terms are labeled implicit, exlicit + for t in self.residual: + if ((not t.has_label(implicit)) and (not t.has_label(explicit)) + and (not t.has_label(time_derivative))): + raise NotImplementedError("Non time-derivative terms must be labeled as implicit or explicit") + + self.xs = [Function(self.fs) for i in range(self.nStages)] + @cached_property def lhs(self): - return super().lhs + """Set up the discretisation's left hand side (the time derivative).""" + return super(IMEXMultistage, self).lhs + @cached_property def rhs(self): - return super().rhs + """Set up the discretisation's right hand side (the time derivative).""" + return super(IMEXMultistage, self).rhs - def solver(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), - ) + 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) - residual += mass_form.label_map(all_terms, - replace_subject(self.x_out, self.idx)) - - problem = NonlinearVariationalProblem(residual.form, self.x_out, bcs=self.bcs) + 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*(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), + 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) + 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), + 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) + residual += r_imp + return residual.form - solver_name = self.field_name+self.__class__.__name__ + "%s" % (stage) - return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, - options_prefix=solver_name) + @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), + 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) + residual += r_imp + residual += r_exp + return residual.form @cached_property def solvers(self): + """Set up a list of solvers for each problem at a stage.""" solvers = [] for stage in range(self.nStages): - solvers.append(self.solver(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 - 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 - solver = self.solvers[stage] - solver.solve() - - self.k[stage].assign(self.x_out) + @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__ + return NonlinearVariationalSolver(problem, options_prefix=solver_name) def apply(self, x_out, x_in): + self.x1.assign(x_in) + solver_list = self.solvers - 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]) + for stage in range(self.nStages): + self.solver = solver_list[stage] + self.solver.solve() + self.xs[stage].assign(self.x_out) - if self.limiter is not None: - self.limiter.apply(x_out) + self.final_solver.solve() + x_out.assign(self.x_out) class ExplicitTimeDiscretisation(TimeDiscretisation): @@ -501,6 +572,144 @@ def apply(self, x_out, x_in): x_out.assign(self.x1) +class ImplicitMultistage(TimeDiscretisation): + """ + A class for implementing general diagonally implicit multistage (Runge-Kutta) + methods based on its Butcher tableau. + + 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 each i = 1, s in an s stage method + we have the intermediate solutions: \n + y_i = y^n + dt*(a_i1*k_1 + a_i2*k_2 + ... + a_ii*k_i) \n + We compute the gradient at the intermediate location, k_i = F(y_i) \n + + At the last stage, compute the new solution by: \n + y^{n+1} = y^n + dt*(b_1*k_1 + b_2*k_2 + .... + b_s*k_s) \n + + """ + # --------------------------------------------------------------------------- + # Butcher tableau for a s-th order + # diagonally implicit scheme: + # c_0 | a_00 0 . 0 + # c_1 | a_10 a_11 . 0 + # . | . . . . + # . | . . . . + # c_s | a_s0 a_s1 . a_ss + # ------------------------- + # | b_1 b_2 ... b_s + # + # + # The corresponding square 'butcher_matrix' is: + # + # [a_00 0 . 0 ] + # [a_10 a_11 . 0 ] + # [ . . . . ] + # [ b_0 b_1 . b_s] + # --------------------------------------------------------------------------- + + def __init__(self, domain, butcher_matrix, 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_matrix (numpy array): A matrix containing the coefficients of + a butcher tableau defining a given Runge Kutta time discretisation. + 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. + """ + super().__init__(domain, field_name=field_name, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + 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 solver(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) + + @cached_property + def solvers(self): + solvers = [] + for stage in range(self.nStages): + solvers.append(self.solver(stage)) + return solvers + + 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 + solver = self.solvers[stage] + solver.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 ExplicitMultistage(ExplicitTimeDiscretisation): """ A class for implementing general explicit multistage (Runge-Kutta) @@ -508,8 +717,7 @@ class ExplicitMultistage(ExplicitTimeDiscretisation): A Butcher tableau is formed in the following way for a s-th order explicit scheme: \n - 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} @@ -1725,3 +1933,196 @@ def __init__(self, domain, field_name=None, solver_parameters=None, super().__init__(domain, butcher_matrix, field_name, solver_parameters=solver_parameters, limiter=limiter, options=options) + + +class IMEX_Euler(IMEXMultistage): + u""" + Implements IMEX Euler one-stage method. + + 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*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): + """ + 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) + + +class ARS3(IMEXMultistage): + u""" + 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): + """ + 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]]) + + super().__init__(domain, butcher_imp, butcher_exp, field_name, + solver_parameters=solver_parameters, + limiter=limiter, options=options) + + +class ARK2(IMEXMultistage): + u""" + 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): + """ + 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.)) + 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 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 \n + ∂y/∂t = F(y) + S(y), can be written as: \n + + 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): + """ + 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) + + +class Trap2(IMEXMultistage): + u""" + 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 \n + ∂y/∂t = F(y) + S(y), can be written as: \n + + 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): + """ + 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) diff --git a/integration-tests/model/test_IMEX.py b/integration-tests/model/test_IMEX.py new file mode 100644 index 000000000..d4ce82c22 --- /dev/null +++ b/integration-tests/model/test_IMEX.py @@ -0,0 +1,42 @@ +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