From b3a06424dc46bc65f5dee7ff656f0fd2f57bcaa1 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Fri, 23 Jun 2023 13:19:04 +0100 Subject: [PATCH 01/20] pull out wrappers into their own object and file --- gusto/__init__.py | 1 + gusto/configuration.py | 9 +- gusto/time_discretisation.py | 210 +++++------------------- gusto/wrappers.py | 307 +++++++++++++++++++++++++++++++++++ 4 files changed, 352 insertions(+), 175 deletions(-) create mode 100644 gusto/wrappers.py diff --git a/gusto/__init__.py b/gusto/__init__.py index d541ba7a9..c20609f0f 100644 --- a/gusto/__init__.py +++ b/gusto/__init__.py @@ -18,3 +18,4 @@ from gusto.time_discretisation import * # noqa from gusto.timeloop import * # noqa from gusto.transport_forms import * # noqa +from gusto.wrappers import * # noqa diff --git a/gusto/configuration.py b/gusto/configuration.py index cc77843e3..5aeaad7e8 100644 --- a/gusto/configuration.py +++ b/gusto/configuration.py @@ -153,7 +153,7 @@ class ShallowWaterParameters(Configuration): H = None # mean depth -class TransportOptions(Configuration, metaclass=ABCMeta): +class WrapperOptions(Configuration, metaclass=ABCMeta): """Base class for specifying options for a transport scheme.""" @abstractproperty @@ -161,14 +161,15 @@ def name(self): pass -class EmbeddedDGOptions(TransportOptions): +class EmbeddedDGOptions(WrapperOptions): """Specifies options for an embedded DG method.""" name = "embedded_dg" + project_back_method = 'project' embedding_space = None -class RecoveryOptions(TransportOptions): +class RecoveryOptions(WrapperOptions): """Specifies options for a recovery wrapper method.""" name = "recovered" @@ -181,7 +182,7 @@ class RecoveryOptions(TransportOptions): broken_method = 'interpolate' -class SUPGOptions(TransportOptions): +class SUPGOptions(WrapperOptions): """Specifies options for an SUPG scheme.""" name = "supg" diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 1be341108..9584d68d0 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -19,47 +19,26 @@ from gusto.labels import (time_derivative, transporting_velocity, prognostic, subject, physics, transport, ibp_label, replace_subject, replace_test_function) -from gusto.recovery import Recoverer, ReversibleRecoverer from gusto.fml.form_manipulation_labelling import Term, all_terms, drop from gusto.transport_forms import advection_form, continuity_form +from gusto.wrappers import * __all__ = ["ForwardEuler", "BackwardEuler", "SSPRK3", "RK4", "Heun", "ThetaMethod", "ImplicitMidpoint", "BDF2", "TR_BDF2", "Leapfrog", "AdamsMoulton", "AdamsBashforth"] -def is_cg(V): - """ - Checks if a :class:`FunctionSpace` is continuous. - - Function to check if a given space, V, is CG. Broken elements are always - discontinuous; for vector elements we check the names of the Sobolev spaces - of the subelements and for all other elements we just check the Sobolev - space name. - - Args: - V (:class:`FunctionSpace`): the space to check. - """ - ele = V.ufl_element() - if isinstance(ele, BrokenElement): - return False - elif type(ele) == VectorElement: - return all([e.sobolev_space().name == "H1" for e in ele._sub_elements]) - else: - return V.ufl_element().sobolev_space().name == "H1" - - -def embedded_dg(original_apply): - """Decorator to add steps for embedded DG method.""" +def wrapper_apply(original_apply): + """Decorator to add steps for using a wrapper around the apply method.""" def get_apply(self, x_out, x_in): - if self.discretisation_option in ["embedded_dg", "recovered"]: + if self.wrapper is not None: def new_apply(self, x_out, x_in): - self.pre_apply(x_in, self.discretisation_option) - original_apply(self, self.xdg_out, self.xdg_in) - self.post_apply(x_out, self.discretisation_option) + self.wrapper.pre_apply(x_in) + original_apply(self, self.wrapper.x_out, self.wrapper.x_in) + self.wrapper.post_apply(x_out) return new_apply(self, x_out, x_in) @@ -95,14 +74,23 @@ def __init__(self, domain, field_name=None, solver_parameters=None, self.equation = None self.dt = domain.dt - + self.options = options self.limiter = limiter - self.options = options if options is not None: - self.discretisation_option = options.name + self.wrapper_name = options.name + if self.wrapper_name == "embedded_dg": + self.wrapper = EmbeddedDGWrapper(self, options) + elif self.wrapper_name == "recovered": + self.wrapper = RecoveryWrapper(self, options) + elif self.wrapper_name == "supg": + self.wrapper = SUPGWrapper(self, options) + else: + raise NotImplementedError(f'Time discretisation: wrapper ' + + '{self.wrapper_name} not implemented') else: - self.discretisation_option = None + self.wrapper = None + self.wrapper_name = None # get default solver options if none passed in if solver_parameters is None: @@ -114,6 +102,7 @@ def __init__(self, domain, field_name=None, solver_parameters=None, if logger.isEnabledFor(DEBUG): self.solver_parameters["ksp_monitor_true_residual"] = None + def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): """ Set up the time discretisation based on the equation. @@ -157,8 +146,6 @@ def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): if t.has_label(physics): self.evaluate_source.append(t.get(physics)) - options = self.options - # -------------------------------------------------------------------- # # Routines relating to transport # -------------------------------------------------------------------- # @@ -168,26 +155,23 @@ def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): self.replace_transporting_velocity(uadv) # -------------------------------------------------------------------- # - # Wrappers for embedded / recovery methods + # Set up Wrappers # -------------------------------------------------------------------- # - if self.discretisation_option in ["embedded_dg", "recovered"]: - # construct the embedding space if not specified - if options.embedding_space is None: - V_elt = BrokenElement(self.fs.ufl_element()) - self.fs = FunctionSpace(self.domain.mesh, V_elt) - else: - self.fs = options.embedding_space - self.xdg_in = Function(self.fs) - self.xdg_out = Function(self.fs) - if self.idx is None: - self.x_projected = Function(equation.function_space) - else: - self.x_projected = Function(equation.spaces[self.idx]) - new_test = TestFunction(self.fs) - parameters = {'ksp_type': 'cg', - 'pc_type': 'bjacobi', - 'sub_pc_type': 'ilu'} + if self.wrapper is not None: + self.wrapper.setup() + self.fs = self.wrapper.function_space + if self.solver_parameters is None: + self.solver_parameters = self.wrapper.solver_parameters + new_test = TestFunction(self.wrapper.test_space) + # TODO: this needs moving if SUPG becomes a transport scheme + if self.wrapper_name == "supg": + new_test = new_test + dot(dot(uadv, self.wrapper.tau), grad(new_test)) + + # Replace the original test function with the one from the wrapper + self.residual = self.residual.label_map( + all_terms, + map_if_true=replace_test_function(new_test)) # -------------------------------------------------------------------- # # Make boundary conditions @@ -195,135 +179,19 @@ def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): if not apply_bcs: self.bcs = None - elif self.discretisation_option in ["embedded_dg", "recovered"]: + elif self.wrapper is not None: # Transfer boundary conditions onto test function space self.bcs = [DirichletBC(self.fs, bc.function_arg, bc.sub_domain) for bc in bcs] else: self.bcs = bcs # -------------------------------------------------------------------- # - # Modify test function for SUPG methods + # Make the required functions # -------------------------------------------------------------------- # - if self.discretisation_option == "supg": - # construct tau, if it is not specified - dim = self.domain.mesh.topological_dimension() - if options.tau is not None: - # if tau is provided, check that is has the right size - tau = options.tau - assert as_ufl(tau).ufl_shape == (dim, dim), "Provided tau has incorrect shape!" - else: - # create tuple of default values of size dim - default_vals = [options.default*self.dt]*dim - # check for directions is which the space is discontinuous - # so that we don't apply supg in that direction - if is_cg(self.fs): - vals = default_vals - else: - space = self.fs.ufl_element().sobolev_space() - if space.name in ["HDiv", "DirectionalH"]: - vals = [default_vals[i] if space[i].name == "H1" - else 0. for i in range(dim)] - else: - raise ValueError("I don't know what to do with space %s" % space) - tau = Constant(tuple([ - tuple( - [vals[j] if i == j else 0. for i, v in enumerate(vals)] - ) for j in range(dim)]) - ) - self.solver_parameters = {'ksp_type': 'gmres', - 'pc_type': 'bjacobi', - 'sub_pc_type': 'ilu'} - - test = TestFunction(self.fs) - new_test = test + dot(dot(uadv, tau), grad(test)) - - if self.discretisation_option is not None: - # replace the original test function with one defined on - # the embedding space, as this is the space where the - # the problem will be solved - self.residual = self.residual.label_map( - all_terms, - map_if_true=replace_test_function(new_test)) - - if self.discretisation_option == "embedded_dg": - self.interp_back = False - if self.limiter is None: - self.x_out_projector = Projector(self.xdg_out, self.x_projected, - solver_parameters=parameters) - else: - self.x_out_projector = Recoverer(self.xdg_out, self.x_projected) - - if self.discretisation_option == "recovered": - # set up the necessary functions - if self.idx is not None: - self.x_in = Function(equation.spaces[self.idx]) - else: - self.x_in = Function(equation.function_space) - - # Operator to recover to higher discontinuous space - self.x_recoverer = ReversibleRecoverer(self.x_in, self.xdg_in, options) - - self.interp_back = (options.project_low_method == 'interpolate') - if options.project_low_method == 'interpolate': - self.x_out_projector = Interpolator(self.xdg_out, self.x_projected) - elif options.project_low_method == 'project': - self.x_out_projector = Projector(self.xdg_out, self.x_projected) - elif options.project_low_method == 'recover': - self.x_out_projector = Recoverer(self.xdg_out, self.x_projected, method=options.broken_method) - - if self.limiter is not None and options.project_low_method != 'recover': - logger.warning('A limiter has been requested for a recovered transport scheme, but the method for projecting back is not recovery') - - # setup required functions self.x_out = Function(self.fs) self.x1 = Function(self.fs) - def pre_apply(self, x_in, discretisation_option): - """ - Extra steps to the discretisation if using a "wrapper" method. - - Performs extra steps before the generic apply method if the whole method - is a "wrapper" around some existing discretisation. For instance, if - using an embedded or recovered method this routine performs the - transformation to the function space in which the discretisation is - computed. - - Args: - x_in (:class:`Function`): the original field to be evolved. - discretisation_option (str): specifies the "wrapper" method. - """ - if discretisation_option == "embedded_dg": - try: - self.xdg_in.interpolate(x_in) - except NotImplementedError: - self.xdg_in.project(x_in) - - elif discretisation_option == "recovered": - self.x_in.assign(x_in) - self.x_recoverer.project() - - else: - raise ValueError( - f'discretisation_option {discretisation_option} not recognised') - - def post_apply(self, x_out, discretisation_option): - """ - Extra steps to the discretisation if using a "wrapper" method. - - Performs projection steps after the generic apply method if the whole - method is a "wrapper" around some existing discretisation. This - generally returns a field to its original space. For the case of the - recovered scheme, there are two options dependent on whether - the scheme is limited or not. - - Args: - x_out (:class:`Function`): the outgoing field to be computed. - discretisation_option (str): specifies the "wrapper" method. - """ - self.x_out_projector.interpolate() if self.interp_back else self.x_out_projector.project() - x_out.assign(self.x_projected) - @property def nlevels(self): return 1 @@ -503,7 +371,7 @@ def apply_cycle(self, x_out, x_in): """ pass - @embedded_dg + @wrapper_apply def apply(self, x_out, x_in): """ Apply the time discretisation to advance one whole time step. diff --git a/gusto/wrappers.py b/gusto/wrappers.py new file mode 100644 index 000000000..86e19eadd --- /dev/null +++ b/gusto/wrappers.py @@ -0,0 +1,307 @@ +""" +Wrappers are objects that wrap around particular time discretisations, applying +some generic operation before and after a standard time discretisation is +called. +""" + +from abc import ABCMeta, abstractmethod +from firedrake import (FunctionSpace, Function, BrokenElement, Projector, + Interpolator, VectorElement, Constant, as_ufl) +from gusto.configuration import EmbeddedDGOptions, RecoveryOptions, SUPGOptions +from gusto.recovery import Recoverer, ReversibleRecoverer + +__all__ = ["EmbeddedDGWrapper", "RecoveryWrapper", "SUPGWrapper"] + +class Wrapper(object, metaclass=ABCMeta): + """Base class for time discretisation wrapper objects.""" + + def __init__(self, time_discretisation, wrapper_options): + """ + Args: + time_discretisation (:class:`TimeDiscretisation`): the time + discretisation that this wrapper is to be used with. + wrapper_options (:class:`WrapperOptions`): configuration object + holding the options specific to this `Wrapper`. + """ + + self.time_discretisation = time_discretisation + self.options = wrapper_options + self.solver_parameters = None + + @abstractmethod + def setup(self): + """ + Performs standard set up routines, and is to be called by the setup + method of the underlying time discretisation. + """ + pass + + @abstractmethod + def pre_apply(self): + """Generic steps to be done before time discretisation apply method.""" + pass + + @abstractmethod + def post_apply(self): + """Generic steps to be done after time discretisation apply method.""" + pass + + +class EmbeddedDGWrapper(Wrapper): + """ + Wrapper for computing a time discretisation with the Embedded DG method, in + which a field is converted to an embedding discontinuous space, then evolved + using a method suitable for this space, before projecting the field back to + the original space. + """ + + def setup(self): + """Sets up function spaces and fields needed for this wrapper.""" + + assert isinstance(self.options, EmbeddedDGOptions), \ + 'Embedded DG wrapper can only be used with Embedded DG Options' + + original_space = self.time_discretisation.fs + domain = self.time_discretisation.domain + equation = self.time_discretisation.equation + + # -------------------------------------------------------------------- # + # Set up spaces to be used with wrapper + # -------------------------------------------------------------------- # + + if self.options.embedding_space is None: + V_elt = BrokenElement(original_space.ufl_element()) + self.function_space = FunctionSpace(domain.mesh, V_elt) + else: + self.function_space = self.options.embedding_space + + self.test_space = self.function_space + + # -------------------------------------------------------------------- # + # Internal variables to be used + # -------------------------------------------------------------------- # + + self.x_in = Function(self.function_space) + self.x_out = Function(self.function_space) + if self.time_discretisation.idx is None: + self.x_projected = Function(equation.function_space) + else: + self.x_projected = Function(equation.spaces[self.idx]) + + if self.options.project_back_method == 'project': + self.x_out_projector = Projector(self.x_out, self.x_projected) + elif self.options.project_back_method == 'recover': + self.x_out_projector = Recoverer(self.x_out, self.x_projected) + else: + raise NotImplementedError(f'EmbeddedDG Wrapper: project_back_method' + + ' {self.options.project_back_method} is not implemented') + + self.parameters = {'ksp_type': 'cg', + 'pc_type': 'bjacobi', + 'sub_pc_type': 'ilu'} + + def pre_apply(self, x_in): + """ + Extra pre-apply steps for the embedded DG method. Interpolates or + projects x_in to the embedding space. + + Args: + x_in (:class:`Function`): the original input field. + """ + + try: + self.x_in.interpolate(x_in) + except NotImplementedError: + self.x_in.project(x_in) + + def post_apply(self, x_out): + """ + Extra post-apply steps for the embedded DG method. Projects the output + field from the embedding space to the original space. + + Args: + x_out (:class:`Function`): the output field in the original space. + """ + + self.x_out_projector.project() + x_out.assign(self.x_projected) + + +class RecoveryWrapper(Wrapper): + """ + Wrapper for computing a time discretisation with the "recovered" method, in + which a field is converted to higher-order function space space. The field + is then evolved in this higher-order function space to obtain an increased + order of accuracy over evolving the field in the lower-order space. The + field is then returned to the original space. + """ + + def setup(self): + """Sets up function spaces and fields needed for this wrapper.""" + + assert isinstance(self.options, RecoveryOptions), \ + 'Embedded DG wrapper can only be used with Recovery Options' + + original_space = self.time_discretisation.fs + domain = self.time_discretisation.domain + equation = self.time_discretisation.equation + + # -------------------------------------------------------------------- # + # Set up spaces to be used with wrapper + # -------------------------------------------------------------------- # + + if self.options.embedding_space is None: + V_elt = BrokenElement(original_space.ufl_element()) + self.function_space = FunctionSpace(domain.mesh, V_elt) + else: + self.function_space = self.options.embedding_space + + self.test_space = self.function_space + + # -------------------------------------------------------------------- # + # Internal variables to be used + # -------------------------------------------------------------------- # + + self.x_in_tmp = Function(self.time_discretisation.fs) + self.x_in = Function(self.function_space) + self.x_out = Function(self.function_space) + if self.time_discretisation.idx is None: + self.x_projected = Function(equation.function_space) + else: + self.x_projected = Function(equation.spaces[self.idx]) + + # Operator to recover to higher discontinuous space + self.x_recoverer = ReversibleRecoverer(self.x_in_tmp, self.x_in, self.options) + + # Operators for projecting back + self.interp_back = (self.options.project_low_method == 'interpolate') + if self.options.project_low_method == 'interpolate': + self.x_out_projector = Interpolator(self.x_out, self.x_projected) + elif self.options.project_low_method == 'project': + self.x_out_projector = Projector(self.x_out, self.x_projected) + elif self.options.project_low_method == 'recover': + self.x_out_projector = Recoverer(self.x_out, self.x_projected, + method=self.options.broken_method) + else: + raise NotImplementedError(f'Recovery Wrapper: project_back_method' + + ' {self.options.project_back_method} is not implemented') + + def pre_apply(self, x_in): + """ + Extra pre-apply steps for the recovered method. Interpolates or projects + x_in to the embedding space. + + Args: + x_in (:class:`Function`): the original input field. + """ + + self.x_in_tmp.assign(x_in) + self.x_recoverer.project() + + def post_apply(self, x_out): + """ + Extra post-apply steps for the recovered method. Projects the output + field from the embedding space to the original space. + + Args: + x_out (:class:`Function`): the output field in the original space. + """ + + if self.interp_back: + self.x_out_projector.interpolate() + else: + self.x_out_projector.project() + x_out.assign(self.x_projected) + + +def is_cg(V): + """ + Checks if a :class:`FunctionSpace` is continuous. + + Function to check if a given space, V, is CG. Broken elements are always + discontinuous; for vector elements we check the names of the Sobolev spaces + of the subelements and for all other elements we just check the Sobolev + space name. + + Args: + V (:class:`FunctionSpace`): the space to check. + """ + ele = V.ufl_element() + if isinstance(ele, BrokenElement): + return False + elif type(ele) == VectorElement: + return all([e.sobolev_space().name == "H1" for e in ele._sub_elements]) + else: + return V.ufl_element().sobolev_space().name == "H1" + + +class SUPGWrapper(Wrapper): + """ + Wrapper for computing a time discretisation with SUPG, which adjust the + test function space that is used to solve the problem. + """ + + def setup(self): + """Sets up function spaces and fields needed for this wrapper.""" + + assert isinstance(self.options, SUPGOptions), \ + 'SUPG wrapper can only be used with SUPG Options' + + domain = self.time_discretisation.domain + self.function_space = self.time_discretisation.fs + self.test_space = self.function_space + self.x_out = Function(self.function_space) + + # -------------------------------------------------------------------- # + # Work out SUPG parameter + # -------------------------------------------------------------------- # + + # construct tau, if it is not specified + dim = domain.mesh.topological_dimension() + if self.options.tau is not None: + # if tau is provided, check that is has the right size + self.tau = self.options.tau + assert as_ufl(self.tau).ufl_shape == (dim, dim), "Provided tau has incorrect shape!" + else: + # create tuple of default values of size dim + default_vals = [self.options.default*self.time_discretisation.dt]*dim + # check for directions is which the space is discontinuous + # so that we don't apply supg in that direction + if is_cg(self.function_space): + vals = default_vals + else: + space = self.function_space.ufl_element().sobolev_space() + if space.name in ["HDiv", "DirectionalH"]: + vals = [default_vals[i] if space[i].name == "H1" + else 0. for i in range(dim)] + else: + raise ValueError("I don't know what to do with space %s" % space) + self.tau = Constant(tuple([ + tuple( + [vals[j] if i == j else 0. for i, v in enumerate(vals)] + ) for j in range(dim)]) + ) + self.solver_parameters = {'ksp_type': 'gmres', + 'pc_type': 'bjacobi', + 'sub_pc_type': 'ilu'} + + + def pre_apply(self, x_in): + """ + Does nothing for SUPG, just sets the input field. + + Args: + x_in (:class:`Function`): the original input field. + """ + + self.x_in = x_in + + def post_apply(self, x_out): + """ + Does nothing for SUPG, just sets the output field. + + Args: + x_out (:class:`Function`): the output field in the original space. + """ + + x_out.assign(self.x_out) From 069ebd98909e679f306c36785c891d9d7118cfd1 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Fri, 23 Jun 2023 13:21:55 +0100 Subject: [PATCH 02/20] lint fixes --- gusto/time_discretisation.py | 9 +++------ gusto/wrappers.py | 12 +++++++----- 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 9584d68d0..2da2c40a0 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -7,9 +7,7 @@ from abc import ABCMeta, abstractmethod, abstractproperty from firedrake import (Function, NonlinearVariationalProblem, split, - NonlinearVariationalSolver, Projector, Interpolator, - BrokenElement, VectorElement, FunctionSpace, - TestFunction, Constant, dot, grad, as_ufl, + NonlinearVariationalSolver, TestFunction, dot, grad, DirichletBC) from firedrake.formmanipulation import split_form from firedrake.utils import cached_property @@ -86,8 +84,8 @@ def __init__(self, domain, field_name=None, solver_parameters=None, elif self.wrapper_name == "supg": self.wrapper = SUPGWrapper(self, options) else: - raise NotImplementedError(f'Time discretisation: wrapper ' - + '{self.wrapper_name} not implemented') + raise NotImplementedError( + f'Time discretisation: wrapper {self.wrapper_name} not implemented') else: self.wrapper = None self.wrapper_name = None @@ -102,7 +100,6 @@ def __init__(self, domain, field_name=None, solver_parameters=None, if logger.isEnabledFor(DEBUG): self.solver_parameters["ksp_monitor_true_residual"] = None - def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): """ Set up the time discretisation based on the equation. diff --git a/gusto/wrappers.py b/gusto/wrappers.py index 86e19eadd..e78706470 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -12,6 +12,7 @@ __all__ = ["EmbeddedDGWrapper", "RecoveryWrapper", "SUPGWrapper"] + class Wrapper(object, metaclass=ABCMeta): """Base class for time discretisation wrapper objects.""" @@ -93,8 +94,9 @@ def setup(self): elif self.options.project_back_method == 'recover': self.x_out_projector = Recoverer(self.x_out, self.x_projected) else: - raise NotImplementedError(f'EmbeddedDG Wrapper: project_back_method' - + ' {self.options.project_back_method} is not implemented') + raise NotImplementedError( + 'EmbeddedDG Wrapper: project_back_method' + + f' {self.options.project_back_method} is not implemented') self.parameters = {'ksp_type': 'cg', 'pc_type': 'bjacobi', @@ -183,8 +185,9 @@ def setup(self): self.x_out_projector = Recoverer(self.x_out, self.x_projected, method=self.options.broken_method) else: - raise NotImplementedError(f'Recovery Wrapper: project_back_method' - + ' {self.options.project_back_method} is not implemented') + raise NotImplementedError( + 'Recovery Wrapper: project_back_method' + + f' {self.options.project_back_method} is not implemented') def pre_apply(self, x_in): """ @@ -285,7 +288,6 @@ def setup(self): 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'} - def pre_apply(self, x_in): """ Does nothing for SUPG, just sets the input field. From 55eaf41fe8062eb6f9132ea8b084caa00f97820b Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Fri, 23 Jun 2023 15:10:42 +0100 Subject: [PATCH 03/20] small fix to where we get idx from for mixed function spaces --- gusto/wrappers.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gusto/wrappers.py b/gusto/wrappers.py index e78706470..af131ae60 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -87,7 +87,7 @@ def setup(self): if self.time_discretisation.idx is None: self.x_projected = Function(equation.function_space) else: - self.x_projected = Function(equation.spaces[self.idx]) + self.x_projected = Function(equation.spaces[self.time_discretisation.idx]) if self.options.project_back_method == 'project': self.x_out_projector = Projector(self.x_out, self.x_projected) @@ -170,7 +170,7 @@ def setup(self): if self.time_discretisation.idx is None: self.x_projected = Function(equation.function_space) else: - self.x_projected = Function(equation.spaces[self.idx]) + self.x_projected = Function(equation.spaces[self.time_discretisation.idx]) # Operator to recover to higher discontinuous space self.x_recoverer = ReversibleRecoverer(self.x_in_tmp, self.x_in, self.options) From 795f31b120428e42eed4333ecd69429e05ed7b44 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sun, 25 Jun 2023 18:16:30 +0100 Subject: [PATCH 04/20] try to pull out transport scheme into its own object. Not giving sensible answers --- gusto/__init__.py | 1 + gusto/time_discretisation.py | 77 --------- gusto/timeloop.py | 44 +++-- gusto/transport_schemes.py | 155 ++++++++++++++++++ .../balance/test_compressible_balance.py | 3 + 5 files changed, 191 insertions(+), 89 deletions(-) create mode 100644 gusto/transport_schemes.py diff --git a/gusto/__init__.py b/gusto/__init__.py index c20609f0f..0747b6230 100644 --- a/gusto/__init__.py +++ b/gusto/__init__.py @@ -18,4 +18,5 @@ from gusto.time_discretisation import * # noqa from gusto.timeloop import * # noqa from gusto.transport_forms import * # noqa +from gusto.transport_schemes import * # noqa from gusto.wrappers import * # noqa diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 2da2c40a0..4b7dc1e2f 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -143,14 +143,6 @@ def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): if t.has_label(physics): self.evaluate_source.append(t.get(physics)) - # -------------------------------------------------------------------- # - # Routines relating to transport - # -------------------------------------------------------------------- # - - if hasattr(self.options, 'ibp'): - self.replace_transport_term() - self.replace_transporting_velocity(uadv) - # -------------------------------------------------------------------- # # Set up Wrappers # -------------------------------------------------------------------- # @@ -216,75 +208,6 @@ def rhs(self): return r.form - def replace_transport_term(self): - """ - Replaces a transport term with some other transport term. - - This routine allows the default transport term to be replaced with a - different one, specified through the transport options. This is - necessary because when the prognostic equations are declared, - the particular transport scheme is not known. The details of the new - transport term are stored in the time discretisation's options object. - """ - # Extract transport term of equation - old_transport_term_list = self.residual.label_map( - lambda t: t.has_label(transport), map_if_false=drop) - - # If there are more transport terms, extract only the one for this variable - if len(old_transport_term_list.terms) > 1: - raise NotImplementedError('Cannot replace transport terms when there are more than one') - - # Then we should only have one transport term - old_transport_term = old_transport_term_list.terms[0] - - # If the transport term has an ibp label, then it could be replaced - if old_transport_term.has_label(ibp_label) and hasattr(self.options, 'ibp'): - # Do the options specify a different ibp to the old transport term? - if old_transport_term.labels['ibp'] != self.options.ibp: - # Set up a new transport term - if self.idx is not None: - field = self.equation.X.split()[self.idx] - else: - field = self.equation.X - test = TestFunction(self.fs) - - # Set up new transport term (depending on the type of transport equation) - if old_transport_term.labels['transport'] == TransportEquationType.advective: - new_transport_term = advection_form(self.domain, test, field, ibp=self.options.ibp) - elif old_transport_term.labels['transport'] == TransportEquationType.conservative: - new_transport_term = continuity_form(self.domain, test, field, ibp=self.options.ibp) - else: - raise NotImplementedError(f'Replacement of transport term not implemented yet for {old_transport_term.labels["transport"]}') - - # Finally, drop the old transport term and add the new one - self.residual = self.residual.label_map( - lambda t: t.has_label(transport), map_if_true=drop) - self.residual += subject(new_transport_term, field) - - def replace_transporting_velocity(self, uadv): - """ - Replace the transport velocity. - - Args: - uadv (:class:`ufl.Expr`): the new transporting velocity. - """ - # replace the transporting velocity in any terms that contain it - if any([t.has_label(transporting_velocity) for t in self.residual]): - assert uadv is not None - if uadv == "prognostic": - self.residual = self.residual.label_map( - lambda t: t.has_label(transporting_velocity), - map_if_true=lambda t: Term(ufl.replace( - t.form, {t.get(transporting_velocity): split(t.get(subject))[0]}), t.labels) - ) - else: - self.residual = self.residual.label_map( - lambda t: t.has_label(transporting_velocity), - map_if_true=lambda t: Term(ufl.replace( - t.form, {t.get(transporting_velocity): uadv}), t.labels) - ) - self.residual = transporting_velocity.update_value(self.residual, uadv) - @cached_property def solver(self): """Set up the problem and the solver.""" diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 726d19004..8c5d02259 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -6,12 +6,13 @@ from gusto.configuration import logger from gusto.equations import PrognosticEquationSet from gusto.forcing import Forcing -from gusto.fml.form_manipulation_labelling import drop, Label +from gusto.fml.form_manipulation_labelling import drop, Label, Term from gusto.labels import (transport, diffusion, time_derivative, linearisation, prognostic, physics) from gusto.linear_solvers import LinearTimesteppingSolver from gusto.fields import TimeLevelFields, StateFields from gusto.time_discretisation import ExplicitTimeDiscretisation +from gusto.transport_schemes import transport_discretisation __all__ = ["Timestepper", "SplitPhysicsTimestepper", "SemiImplicitQuasiNewton", "PrescribedTransport"] @@ -72,6 +73,22 @@ def get_initial_timesteps(self): # Return None if this is not applicable return self.scheme.initial_timesteps if can_get else None + def transport_discretisation_setup(self, equation=None): + equation_to_setup = self.equation if equation is None else equation + # Call the setup method using the transporting velocity + for term in equation_to_setup.residual: + if term.has_label(transport_discretisation): + term.get(transport_discretisation).setup(self.transporting_velocity) + elif term.has_label(transport): + logger.warning('Transport term detected without transport discretisation') + raise ValueError('Transport term detected without transport discretisation') + + # Replace standard transport term with the transport discretisation term + equation_to_setup.residual = equation_to_setup.residual.label_map( + lambda t: t.has_label(transport_discretisation), + map_if_true=lambda t: Term(t.get(transport_discretisation).labelled_form.form, t.labels) + ) + def run(self, t, tmax, pick_up=False): """ Runs the model for the specified time, from t to tmax @@ -176,7 +193,8 @@ def setup_fields(self): *self.io.output.dumplist) def setup_scheme(self): - self.scheme.setup(self.equation, self.transporting_velocity) + self.scheme.setup(self.equation) + self.transport_discretisation_setup() def timestep(self): """ @@ -221,18 +239,19 @@ def __init__(self, equation, scheme, io, physics_schemes=None): # check that the supplied schemes for physics are explicit assert isinstance(phys_scheme, ExplicitTimeDiscretisation), "Only explicit schemes can be used for physics" apply_bcs = False - phys_scheme.setup(equation, self.transporting_velocity, apply_bcs, physics) + phys_scheme.setup(equation, apply_bcs, physics) @property def transporting_velocity(self): return "prognostic" def setup_scheme(self): + self.transport_discretisation_setup() # Go through and label all non-physics terms with a "dynamics" label dynamics = Label('dynamics') self.equation.label_terms(lambda t: not any(t.has_label(time_derivative, physics)), dynamics) apply_bcs = True - self.scheme.setup(self.equation, self.transporting_velocity, apply_bcs, dynamics) + self.scheme.setup(self.equation, apply_bcs, dynamics) def timestep(self): @@ -326,7 +345,8 @@ def __init__(self, equation_set, io, transport_schemes, super().__init__(equation_set, io) for aux_eqn, aux_scheme in self.auxiliary_equations_and_schemes: - aux_scheme.setup(aux_eqn, self.transporting_velocity) + self.transport_discretisation_setup(aux_eqn) + aux_scheme.setup(aux_eqn) self.tracers_to_copy = [] for name in equation_set.field_names: @@ -383,14 +403,17 @@ def setup_scheme(self): # TODO: apply_bcs should be False for advection but this means # tests with KGOs fail apply_bcs = True + import pdb; pdb.set_trace() + self.transport_discretisation_setup() + import pdb; pdb.set_trace() for _, scheme in self.active_transport: - scheme.setup(self.equation, self.transporting_velocity, apply_bcs, transport) + scheme.setup(self.equation, apply_bcs, transport) apply_bcs = True for _, scheme in self.diffusion_schemes: - scheme.setup(self.equation, self.transporting_velocity, apply_bcs, diffusion) + scheme.setup(self.equation, apply_bcs, diffusion) for _, scheme in self.physics_schemes: apply_bcs = True - scheme.setup(self.equation, self.transporting_velocity, apply_bcs, physics) + scheme.setup(self.equation, apply_bcs, physics) def copy_active_tracers(self, x_in, x_out): """ @@ -512,7 +535,7 @@ def __init__(self, equation, scheme, io, physics_schemes=None, # check that the supplied schemes for physics are explicit assert isinstance(scheme, ExplicitTimeDiscretisation), "Only explicit schemes can be used for physics" apply_bcs = False - scheme.setup(equation, self.transporting_velocity, apply_bcs, physics) + scheme.setup(equation, apply_bcs, physics) if prescribed_transporting_velocity is not None: self.velocity_projection = Projector( @@ -530,9 +553,6 @@ def setup_fields(self): self.fields = StateFields(self.x, self.equation.prescribed_fields, *self.io.output.dumplist) - def setup_scheme(self): - self.scheme.setup(self.equation, self.transporting_velocity) - def timestep(self): if self.velocity_projection is not None: self.velocity_projection.project() diff --git a/gusto/transport_schemes.py b/gusto/transport_schemes.py new file mode 100644 index 000000000..6194b0769 --- /dev/null +++ b/gusto/transport_schemes.py @@ -0,0 +1,155 @@ +""" +Defines TransportScheme objects, which are used to solve a transport problem. +""" + +from firedrake import split +from gusto.configuration import IntegrateByParts, TransportEquationType +from gusto.fml import Term, keep, drop, Label, LabelledForm +from gusto.labels import transporting_velocity, prognostic, transport, subject +from gusto.transport_forms import * +import ufl + +__all__ = ["transport_discretisation", "DGUpwind", "SUPGTransport"] + +transport_discretisation = Label("transport_discretisation", + validator=lambda value: isinstance(value, TransportScheme)) + +class TransportScheme(object): + """ + The base object for describing a transport scheme. + """ + + def __init__(self, equation, variable): + """ + Args: + equation (:class:`PrognosticEquation`): the equation, which includes + a transport term. + variable (str): name of the variable to set the transport scheme for + """ + self.equation = equation + self.variable = variable + self.domain = self.equation.domain + + # TODO: how do we deal with plain transport equation? + variable_idx = equation.field_names.index(variable) + self.test = equation.tests[variable_idx] + self.field = split(equation.X)[variable_idx] + + # Find the original transport term to be used + self.original_form = equation.residual.label_map( + lambda t: t.has_label(transport) and t.get(prognostic) == variable, + map_if_true=keep, map_if_false=drop) + + num_terms = len(self.original_form.terms) + assert num_terms == 1, \ + f'Unable to find transport term for {variable}. {num_terms} found' + + def add_transport_form(self, labelled_form): + """ + Adds the form for the transport discretisation to the appropriate term + in the equation. + + Args: + form (:class:`LabelledForm`): the form used by this discretisation + of the transport term. + """ + + self.labelled_form = labelled_form + + # Add the form to the equation + self.equation.residual = self.equation.residual.label_map( + lambda t: t.has_label(transport) and t.get(prognostic) == self.variable, + map_if_true=lambda t: transport_discretisation(t, self)) + + + def setup(self, uadv): + """ + Set up the transport scheme by replacing the transporting velocity used + in the form. + + Args: + uadv (:class:`ufl.Expr`, optional): the transporting velocity. + Defaults to None. + """ + + assert self.labelled_form.terms[0].has_label(transporting_velocity), \ + 'Cannot set up transport scheme on a term that has no transporting velocity' + + if uadv == "prognostic": + # Find prognostic wind field + uadv = split(self.original_form.terms[0].get(subject))[0] + + self.labelled_form = self.labelled_form.label_map( + lambda t: t.has_label(transporting_velocity), + map_if_true=lambda t: + Term(ufl.replace(t.form, {t.get(transporting_velocity): uadv}), t.labels) + ) + + self.labelled_form = transporting_velocity.update_value(self.labelled_form, uadv) + # Add form to equation residual + self.add_transport_form(self.labelled_form) + + +class DGUpwind(TransportScheme): + """ + The Discontinuous Galerkin Upwind transport scheme. + + Discretises the gradient of a field weakly, taking the upwind value of the + transported variable at facets. + """ + def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, + vector_manifold_correction=False, outflow=False): + """ + Args: + equation (:class:`PrognosticEquation`): the equation, which includes + a transport term. + variable (str): name of the variable to set the transport scheme for + ibp (:class:`IntegrateByParts`, optional): an enumerator for how + many times to integrate by parts. Defaults to `ONCE`. + vector_manifold_correction (bool, optional): whether to include a + vector manifold correction term. Defaults to False. + outflow (bool, optional): whether to include outflow at the domain + boundaries, through exterior facet terms. Defaults to False. + """ + + super().__init__(equation, variable) + self.ibp = ibp + self.vector_manifold_correction = vector_manifold_correction + self.outflow = outflow + + # -------------------------------------------------------------------- # + # Determine appropriate form to use + # -------------------------------------------------------------------- # + + if self.original_form.terms[0].get(transport) == TransportEquationType.advective: + if vector_manifold_correction: + form = vector_manifold_advection_form(self.domain, self.test, + self.field, ibp=ibp, + outflow=outflow) + else: + form = advection_form(self.domain, self.test, self.field, + ibp=ibp, outflow=outflow) + + elif self.original_form.terms[0].get(transport) == TransportEquationType.conservative: + if vector_manifold_correction: + form = vector_manifold_continuity_form(self.domain, self.test, + self.field, ibp=ibp, + outflow=outflow) + else: + form = continuity_form(self.domain, self.test, self.field, + ibp=ibp, outflow=outflow) + + elif self.original_form.terms[0].get(transport) == TransportEquationType.vector_invariant: + if outflow: + raise NotImplementedError('Outflow not implemented for upwind vector invariant') + form = vector_invariant_form(self.domain, self.test, self.field, ibp=ibp) + + else: + raise NotImplementedError('Upwind transport scheme has not been ' + + 'implemented for this transport equation type') + + self.add_transport_form(form) + + +class SUPGTransport(TransportScheme): + pass \ No newline at end of file diff --git a/integration-tests/balance/test_compressible_balance.py b/integration-tests/balance/test_compressible_balance.py index f8efda306..d51de9875 100644 --- a/integration-tests/balance/test_compressible_balance.py +++ b/integration-tests/balance/test_compressible_balance.py @@ -42,6 +42,9 @@ def setup_balance(dirname): transported_fields = [ImplicitMidpoint(domain, "u"), SSPRK3(domain, "rho"), SSPRK3(domain, "theta", options=EmbeddedDGOptions())] + DGUpwind(eqns, 'u') + DGUpwind(eqns, 'rho') + DGUpwind(eqns, 'theta') # Set up linear solver linear_solver = CompressibleSolver(eqns) From c24af5eddc85fd2e22eafac37b314b203050a74e Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Fri, 30 Jun 2023 17:41:40 +0100 Subject: [PATCH 05/20] reorganise transport discretisations and start to debug --- gusto/timeloop.py | 82 +++++++++++++------ gusto/transport_schemes.py | 32 ++++---- .../balance/test_compressible_balance.py | 7 +- 3 files changed, 73 insertions(+), 48 deletions(-) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 8c5d02259..1702aeb83 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -12,7 +12,6 @@ from gusto.linear_solvers import LinearTimesteppingSolver from gusto.fields import TimeLevelFields, StateFields from gusto.time_discretisation import ExplicitTimeDiscretisation -from gusto.transport_schemes import transport_discretisation __all__ = ["Timestepper", "SplitPhysicsTimestepper", "SemiImplicitQuasiNewton", "PrescribedTransport"] @@ -73,21 +72,27 @@ def get_initial_timesteps(self): # Return None if this is not applicable return self.scheme.initial_timesteps if can_get else None - def transport_discretisation_setup(self, equation=None): - equation_to_setup = self.equation if equation is None else equation - # Call the setup method using the transporting velocity - for term in equation_to_setup.residual: - if term.has_label(transport_discretisation): - term.get(transport_discretisation).setup(self.transporting_velocity) - elif term.has_label(transport): - logger.warning('Transport term detected without transport discretisation') - raise ValueError('Transport term detected without transport discretisation') - - # Replace standard transport term with the transport discretisation term - equation_to_setup.residual = equation_to_setup.residual.label_map( - lambda t: t.has_label(transport_discretisation), - map_if_true=lambda t: Term(t.get(transport_discretisation).labelled_form.form, t.labels) - ) + def transport_discretisation_setup(self, equation): + """ + Sets up the transport discretisation, by the setting the transporting + velocity and setting the form used for transport in the equation. + + Args: + equation (:class:`PrognosticEquation`): the equation that the + transport discretisation is to be applied to. + """ + if self.transport_discretisations is not None: + for transport_discretisation in self.transport_discretisations: + transport_discretisation.setup(self.transporting_velocity) + logger.warning('') + transport_terms = equation.residual.label_map( + lambda t: t.has_label(transport), map_if_false=drop + ) + logger.warning(transport_terms.form.__str__()) + transport_discretisation.replace_transport_form(equation) + # TODO: raise a warning here if transport discretisations aren't provided + logger.warning('') + logger.warning(transport_terms.form.__str__()) def run(self, t, tmax, pick_up=False): """ @@ -172,15 +177,23 @@ class Timestepper(BaseTimestepper): Implements a timeloop by applying a scheme to a prognostic equation. """ - def __init__(self, equation, scheme, io): + def __init__(self, equation, scheme, io, transport_discretisations=None): """ Args: equation (:class:`PrognosticEquation`): the prognostic equation scheme (:class:`TimeDiscretisation`): the scheme to use to timestep the prognostic equation io (:class:`IO`): the model's object for controlling input/output. + transport_discretisations (iter,optional): a list of objects + describing the discretisations of transport terms to be used. + Defaults to None. """ self.scheme = scheme + if transport_discretisations is not None: + self.transport_discretisations = transport_discretisations + else: + self.transport_discretisations = [] + super().__init__(equation=equation, io=io) @property @@ -193,8 +206,8 @@ def setup_fields(self): *self.io.output.dumplist) def setup_scheme(self): + self.transport_discretisation_setup(self.equation) self.scheme.setup(self.equation) - self.transport_discretisation_setup() def timestep(self): """ @@ -214,13 +227,17 @@ class SplitPhysicsTimestepper(Timestepper): scheme to be applied to the physics terms than the prognostic equation. """ - def __init__(self, equation, scheme, io, physics_schemes=None): + def __init__(self, equation, scheme, io, transport_discretisations=None, + physics_schemes=None): """ Args: equation (:class:`PrognosticEquation`): the prognostic equation scheme (:class:`TimeDiscretisation`): the scheme to use to timestep the prognostic equation io (:class:`IO`): the model's object for controlling input/output. + transport_discretisations (iter,optional): a list of objects + describing the discretisations of transport terms to be used. + Defaults to None. physics_schemes: (list, optional): a list of :class:`Physics` and :class:`TimeDiscretisation` options describing physical parametrisations and timestepping schemes to use for each. @@ -228,7 +245,8 @@ def __init__(self, equation, scheme, io, physics_schemes=None): None. """ - super().__init__(equation, scheme, io) + super().__init__(equation, scheme, io, + transport_discretisations=transport_discretisations) if physics_schemes is not None: self.physics_schemes = physics_schemes @@ -246,7 +264,7 @@ def transporting_velocity(self): return "prognostic" def setup_scheme(self): - self.transport_discretisation_setup() + self.transport_discretisation_setup(self.equation) # Go through and label all non-physics terms with a "dynamics" label dynamics = Label('dynamics') self.equation.label_terms(lambda t: not any(t.has_label(time_derivative, physics)), dynamics) @@ -273,6 +291,7 @@ class SemiImplicitQuasiNewton(BaseTimestepper): """ def __init__(self, equation_set, io, transport_schemes, + transport_discretisations, auxiliary_equations_and_schemes=None, linear_solver=None, diffusion_schemes=None, @@ -286,6 +305,8 @@ def __init__(self, equation_set, io, transport_schemes, transport_schemes: iterable of ``(field_name, scheme)`` pairs indicating the name of the field (str) to transport, and the :class:`TimeDiscretisation` to use + transport_discretisations (iter): a list of objects describing the + spatial discretisations of transport terms to be used. auxiliary_equations_and_schemes: iterable of ``(equation, scheme)`` pairs indicating any additional equations to be solved and the scheme to use to solve them. Defaults to None. @@ -310,6 +331,8 @@ def __init__(self, equation_set, io, transport_schemes, if kwargs: raise ValueError("unexpected kwargs: %s" % list(kwargs.keys())) + self.transport_discretisations = transport_discretisations + if physics_schemes is not None: self.physics_schemes = physics_schemes else: @@ -403,9 +426,7 @@ def setup_scheme(self): # TODO: apply_bcs should be False for advection but this means # tests with KGOs fail apply_bcs = True - import pdb; pdb.set_trace() - self.transport_discretisation_setup() - import pdb; pdb.set_trace() + self.transport_discretisation_setup(self.equation) for _, scheme in self.active_transport: scheme.setup(self.equation, apply_bcs, transport) apply_bcs = True @@ -449,6 +470,9 @@ def timestep(self): for name, scheme in self.active_transport: # transports a field from xstar and puts result in xp scheme.apply(xp(name), xstar(name)) + # TODO: to remove + import numpy as np + logger.warning(f'{name}: {np.linalg.norm(xp(name).dat.data[:])} {np.linalg.norm(xstar(name).dat.data[:])}') xrhs.assign(0.) # xrhs is the residual which goes in the linear solve @@ -503,14 +527,17 @@ class PrescribedTransport(Timestepper): """ Implements a timeloop with a prescibed transporting velocity """ - def __init__(self, equation, scheme, io, physics_schemes=None, - prescribed_transporting_velocity=None): + def __init__(self, equation, scheme, io, transport_discretisations=None, + physics_schemes=None, prescribed_transporting_velocity=None): """ Args: equation (:class:`PrognosticEquation`): the prognostic equation scheme (:class:`TimeDiscretisation`): the scheme to use to timestep the prognostic equation io (:class:`IO`): the model's object for controlling input/output. + transport_discretisations (iter,optional): a list of objects + describing the discretisations of transport terms to be used. + Defaults to None. physics_schemes: (list, optional): a list of :class:`Physics` and :class:`TimeDiscretisation` options describing physical parametrisations and timestepping schemes to use for each. @@ -524,7 +551,8 @@ def __init__(self, equation, scheme, io, physics_schemes=None, updated. Defaults to None. """ - super().__init__(equation, scheme, io) + super().__init__(equation, scheme, io, + transport_discretisations=transport_discretisations) if physics_schemes is not None: self.physics_schemes = physics_schemes diff --git a/gusto/transport_schemes.py b/gusto/transport_schemes.py index 6194b0769..8d8452487 100644 --- a/gusto/transport_schemes.py +++ b/gusto/transport_schemes.py @@ -9,10 +9,9 @@ from gusto.transport_forms import * import ufl -__all__ = ["transport_discretisation", "DGUpwind", "SUPGTransport"] +__all__ = ["DGUpwind", "SUPGTransport"] -transport_discretisation = Label("transport_discretisation", - validator=lambda value: isinstance(value, TransportScheme)) +# TODO: settle on name: discretisation vs scheme class TransportScheme(object): """ @@ -36,6 +35,7 @@ def __init__(self, equation, variable): self.field = split(equation.X)[variable_idx] # Find the original transport term to be used + # TODO: do we need this? self.original_form = equation.residual.label_map( lambda t: t.has_label(transport) and t.get(prognostic) == variable, map_if_true=keep, map_if_false=drop) @@ -44,22 +44,20 @@ def __init__(self, equation, variable): assert num_terms == 1, \ f'Unable to find transport term for {variable}. {num_terms} found' - def add_transport_form(self, labelled_form): + def replace_transport_form(self, equation): """ - Adds the form for the transport discretisation to the appropriate term - in the equation. + Replaces the form for the transport term in the equation with the + form for the transport discretisation. Args: - form (:class:`LabelledForm`): the form used by this discretisation + equation (:class:`PrognosticEquation`): the form used by this discretisation of the transport term. """ - self.labelled_form = labelled_form - # Add the form to the equation - self.equation.residual = self.equation.residual.label_map( + equation.residual = equation.residual.label_map( lambda t: t.has_label(transport) and t.get(prognostic) == self.variable, - map_if_true=lambda t: transport_discretisation(t, self)) + map_if_true=lambda t: Term(self.form.form, t.labels)) def setup(self, uadv): @@ -72,22 +70,20 @@ def setup(self, uadv): Defaults to None. """ - assert self.labelled_form.terms[0].has_label(transporting_velocity), \ + assert self.form.terms[0].has_label(transporting_velocity), \ 'Cannot set up transport scheme on a term that has no transporting velocity' if uadv == "prognostic": # Find prognostic wind field uadv = split(self.original_form.terms[0].get(subject))[0] - self.labelled_form = self.labelled_form.label_map( + self.form = self.form.label_map( lambda t: t.has_label(transporting_velocity), map_if_true=lambda t: Term(ufl.replace(t.form, {t.get(transporting_velocity): uadv}), t.labels) ) - self.labelled_form = transporting_velocity.update_value(self.labelled_form, uadv) - # Add form to equation residual - self.add_transport_form(self.labelled_form) + self.form = transporting_velocity.update_value(self.form, uadv) class DGUpwind(TransportScheme): @@ -128,7 +124,7 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, outflow=outflow) else: form = advection_form(self.domain, self.test, self.field, - ibp=ibp, outflow=outflow) + ibp=ibp, outflow=outflow) elif self.original_form.terms[0].get(transport) == TransportEquationType.conservative: if vector_manifold_correction: @@ -148,7 +144,7 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, raise NotImplementedError('Upwind transport scheme has not been ' + 'implemented for this transport equation type') - self.add_transport_form(form) + self.form = form class SUPGTransport(TransportScheme): diff --git a/integration-tests/balance/test_compressible_balance.py b/integration-tests/balance/test_compressible_balance.py index d51de9875..96b0d34c1 100644 --- a/integration-tests/balance/test_compressible_balance.py +++ b/integration-tests/balance/test_compressible_balance.py @@ -42,15 +42,16 @@ def setup_balance(dirname): transported_fields = [ImplicitMidpoint(domain, "u"), SSPRK3(domain, "rho"), SSPRK3(domain, "theta", options=EmbeddedDGOptions())] - DGUpwind(eqns, 'u') - DGUpwind(eqns, 'rho') - DGUpwind(eqns, 'theta') + transport_discretisations = [DGUpwind(eqns, 'u'), + DGUpwind(eqns, 'rho'), + DGUpwind(eqns, 'theta')] # Set up linear solver linear_solver = CompressibleSolver(eqns) # build time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + transport_discretisations, linear_solver=linear_solver) # ------------------------------------------------------------------------ # From d704740241152eb3169a6b8f41317a04d5c83486 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sun, 2 Jul 2023 10:23:31 +0100 Subject: [PATCH 06/20] reorganisation to try to replicate old order of replace operations --- gusto/time_discretisation.py | 1 + gusto/timeloop.py | 61 +++++++++++++++++++++++++----------- gusto/transport_schemes.py | 16 ++++++---- 3 files changed, 54 insertions(+), 24 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 4b7dc1e2f..657a0ea97 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -712,6 +712,7 @@ def apply(self, x_out, x_in): x_out (:class:`Function`): the output field to be computed. x_in (:class:`Function`): the input field. """ + import pdb; pdb.set_trace() self.x1.assign(x_in) self.solver.solve() x_out.assign(self.x_out) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 1702aeb83..1af0e084a 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -72,10 +72,10 @@ def get_initial_timesteps(self): # Return None if this is not applicable return self.scheme.initial_timesteps if can_get else None - def transport_discretisation_setup(self, equation): + def transport_discretisation_setup_equation(self, equation): """ - Sets up the transport discretisation, by the setting the transporting - velocity and setting the form used for transport in the equation. + Sets up the transport discretisation for an equation, by the setting the + form used for transport in the equation. Args: equation (:class:`PrognosticEquation`): the equation that the @@ -83,16 +83,21 @@ def transport_discretisation_setup(self, equation): """ if self.transport_discretisations is not None: for transport_discretisation in self.transport_discretisations: - transport_discretisation.setup(self.transporting_velocity) - logger.warning('') - transport_terms = equation.residual.label_map( - lambda t: t.has_label(transport), map_if_false=drop - ) - logger.warning(transport_terms.form.__str__()) transport_discretisation.replace_transport_form(equation) # TODO: raise a warning here if transport discretisations aren't provided - logger.warning('') - logger.warning(transport_terms.form.__str__()) + + def transport_discretisation_setup_scheme(self, transport_discretisation, scheme): + """ + Sets up the transport discretisation, by the setting the transporting + velocity used in the time discretisation. + + Args: + transport_discretisation (:class:`TransportScheme`): the transport + discretisation to be set up. + scheme (:class:`TimeDiscretisation`): the time discretisation used + with this transport scheme. + """ + transport_discretisation.setup(self.transporting_velocity, scheme) def run(self, t, tmax, pick_up=False): """ @@ -206,8 +211,10 @@ def setup_fields(self): *self.io.output.dumplist) def setup_scheme(self): - self.transport_discretisation_setup(self.equation) + self.transport_discretisation_setup_equation(self.equation) self.scheme.setup(self.equation) + for transport_discretisation in self.transport_discretisations: + self.transport_discretisation_setup_scheme(transport_discretisation, self.scheme) def timestep(self): """ @@ -264,12 +271,14 @@ def transporting_velocity(self): return "prognostic" def setup_scheme(self): - self.transport_discretisation_setup(self.equation) + self.transport_discretisation_setup_equation(self.equation) # Go through and label all non-physics terms with a "dynamics" label dynamics = Label('dynamics') self.equation.label_terms(lambda t: not any(t.has_label(time_derivative, physics)), dynamics) apply_bcs = True self.scheme.setup(self.equation, apply_bcs, dynamics) + for transport_discretisation in self.transport_discretisations: + self.transport_discretisation_setup_scheme(transport_discretisation, self.scheme) def timestep(self): @@ -331,7 +340,7 @@ def __init__(self, equation_set, io, transport_schemes, if kwargs: raise ValueError("unexpected kwargs: %s" % list(kwargs.keys())) - self.transport_discretisations = transport_discretisations + self.transport_discretisations = [] if physics_schemes is not None: self.physics_schemes = physics_schemes @@ -345,7 +354,15 @@ def __init__(self, equation_set, io, transport_schemes, for scheme in transport_schemes: assert scheme.nlevels == 1, "multilevel schemes not supported as part of this timestepping loop" assert scheme.field_name in equation_set.field_names - self.active_transport.append((scheme.field_name, scheme)) + # Active transport will be a tuple of (field_name, scheme, discretisation) + discretisation_found = False + for transport_discretisation in transport_discretisations: + if scheme.field_name == transport_discretisation.variable: + discretisation_found = True + self.transport_discretisations.append(transport_discretisation) + self.active_transport.append((scheme.field_name, scheme, transport_discretisation)) + assert discretisation_found, 'No transport discretisation found ' \ + + f'for variable {scheme.field_name}' self.diffusion_schemes = [] if diffusion_schemes is not None: @@ -370,6 +387,12 @@ def __init__(self, equation_set, io, transport_schemes, for aux_eqn, aux_scheme in self.auxiliary_equations_and_schemes: self.transport_discretisation_setup(aux_eqn) aux_scheme.setup(aux_eqn) + # Auxiliary transport discretisations + for transport_discretisation in transport_discretisations: + if transport_discretisation.variable == aux_scheme.field_name: + self.transport_discretisation_setup_scheme(transport_discretisation, aux_scheme) + # TODO: add a check to make sure that auxiliary eqns have transport + # discretisations, but only if they need them! self.tracers_to_copy = [] for name in equation_set.field_names: @@ -426,9 +449,11 @@ def setup_scheme(self): # TODO: apply_bcs should be False for advection but this means # tests with KGOs fail apply_bcs = True - self.transport_discretisation_setup(self.equation) - for _, scheme in self.active_transport: + self.transport_discretisation_setup_equation(self.equation) + for _, scheme, transport_discretisation in self.active_transport: scheme.setup(self.equation, apply_bcs, transport) + self.transport_discretisation_setup_scheme(transport_discretisation, scheme) + apply_bcs = True for _, scheme in self.diffusion_schemes: scheme.setup(self.equation, apply_bcs, diffusion) @@ -467,7 +492,7 @@ def timestep(self): for k in range(self.maxk): with timed_stage("Transport"): - for name, scheme in self.active_transport: + for name, scheme, _ in self.active_transport: # transports a field from xstar and puts result in xp scheme.apply(xp(name), xstar(name)) # TODO: to remove diff --git a/gusto/transport_schemes.py b/gusto/transport_schemes.py index 8d8452487..bc5d0ed15 100644 --- a/gusto/transport_schemes.py +++ b/gusto/transport_schemes.py @@ -50,8 +50,9 @@ def replace_transport_form(self, equation): form for the transport discretisation. Args: - equation (:class:`PrognosticEquation`): the form used by this discretisation - of the transport term. + equation (:class:`PrognosticEquation`): the equation or scheme whose + transport term should be replaced with the transport term of + this discretisation. """ # Add the form to the equation @@ -60,7 +61,7 @@ def replace_transport_form(self, equation): map_if_true=lambda t: Term(self.form.form, t.labels)) - def setup(self, uadv): + def setup(self, uadv, equation): """ Set up the transport scheme by replacing the transporting velocity used in the form. @@ -68,6 +69,9 @@ def setup(self, uadv): Args: uadv (:class:`ufl.Expr`, optional): the transporting velocity. Defaults to None. + equation (:class:`PrognosticEquation`): the equation or scheme whose + transport term should be replaced with the transport term of + this discretisation. """ assert self.form.terms[0].has_label(transporting_velocity), \ @@ -77,13 +81,13 @@ def setup(self, uadv): # Find prognostic wind field uadv = split(self.original_form.terms[0].get(subject))[0] - self.form = self.form.label_map( + equation.residual = equation.residual.label_map( lambda t: t.has_label(transporting_velocity), map_if_true=lambda t: Term(ufl.replace(t.form, {t.get(transporting_velocity): uadv}), t.labels) - ) + ) - self.form = transporting_velocity.update_value(self.form, uadv) + equation.residual = transporting_velocity.update_value(equation.residual, uadv) class DGUpwind(TransportScheme): From 34d873b52c14694e69a039d5cff6e45a409c1fae Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sun, 2 Jul 2023 10:52:33 +0100 Subject: [PATCH 07/20] transport schemes now working (only one example tried) -- still needs neatening! --- gusto/time_discretisation.py | 5 +---- gusto/timeloop.py | 4 +--- gusto/transport_schemes.py | 5 +++++ 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 657a0ea97..36d921ee9 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -100,14 +100,12 @@ def __init__(self, domain, field_name=None, solver_parameters=None, if logger.isEnabledFor(DEBUG): self.solver_parameters["ksp_monitor_true_residual"] = None - def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): + 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. - uadv (:class:`ufl.Expr`, optional): the transporting velocity. - Defaults to None. apply_bcs (bool, optional): whether to apply the equation's boundary conditions. Defaults to True. *active_labels (:class:`Label`): labels indicating which terms of @@ -712,7 +710,6 @@ def apply(self, x_out, x_in): x_out (:class:`Function`): the output field to be computed. x_in (:class:`Function`): the input field. """ - import pdb; pdb.set_trace() self.x1.assign(x_in) self.solver.solve() x_out.assign(self.x_out) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 1af0e084a..7a832f5ac 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -451,6 +451,7 @@ def setup_scheme(self): apply_bcs = True self.transport_discretisation_setup_equation(self.equation) for _, scheme, transport_discretisation in self.active_transport: + print(f'Setting up transport for {scheme.field_name} {transport_discretisation.variable}') scheme.setup(self.equation, apply_bcs, transport) self.transport_discretisation_setup_scheme(transport_discretisation, scheme) @@ -495,9 +496,6 @@ def timestep(self): for name, scheme, _ in self.active_transport: # transports a field from xstar and puts result in xp scheme.apply(xp(name), xstar(name)) - # TODO: to remove - import numpy as np - logger.warning(f'{name}: {np.linalg.norm(xp(name).dat.data[:])} {np.linalg.norm(xstar(name).dat.data[:])}') xrhs.assign(0.) # xrhs is the residual which goes in the linear solve diff --git a/gusto/transport_schemes.py b/gusto/transport_schemes.py index bc5d0ed15..3b1c2eb3a 100644 --- a/gusto/transport_schemes.py +++ b/gusto/transport_schemes.py @@ -81,14 +81,19 @@ def setup(self, uadv, equation): # Find prognostic wind field uadv = split(self.original_form.terms[0].get(subject))[0] + import pdb; pdb.set_trace() + equation.residual = equation.residual.label_map( lambda t: t.has_label(transporting_velocity), map_if_true=lambda t: Term(ufl.replace(t.form, {t.get(transporting_velocity): uadv}), t.labels) ) + import pdb; pdb.set_trace() + equation.residual = transporting_velocity.update_value(equation.residual, uadv) + import pdb; pdb.set_trace() class DGUpwind(TransportScheme): """ From 0f3d736977717e75e0577c8e3129c8d976f56b50 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Thu, 6 Jul 2023 23:13:38 +0100 Subject: [PATCH 08/20] rename and reorganise things -- example working --- gusto/__init__.py | 2 +- gusto/equations.py | 95 ++++++----- gusto/labels.py | 2 +- gusto/timeloop.py | 153 ++++++++++-------- gusto/transport_forms.py | 54 ++++++- ...nsport_schemes.py => transport_methods.py} | 71 +++----- 6 files changed, 201 insertions(+), 176 deletions(-) rename gusto/{transport_schemes.py => transport_methods.py} (64%) diff --git a/gusto/__init__.py b/gusto/__init__.py index 0747b6230..99f5a6be3 100644 --- a/gusto/__init__.py +++ b/gusto/__init__.py @@ -18,5 +18,5 @@ from gusto.time_discretisation import * # noqa from gusto.timeloop import * # noqa from gusto.transport_forms import * # noqa -from gusto.transport_schemes import * # noqa +from gusto.transport_methods import * # noqa from gusto.wrappers import * # noqa diff --git a/gusto/equations.py b/gusto/equations.py index fac7ce954..ca0c25d66 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -15,7 +15,6 @@ from gusto.thermodynamics import exner_pressure from gusto.transport_forms import (advection_form, continuity_form, vector_invariant_form, - vector_manifold_advection_form, kinetic_energy_form, advection_equation_circulation_form, linear_continuity_form, @@ -70,7 +69,7 @@ def label_terms(self, term_filter, label): class AdvectionEquation(PrognosticEquation): u"""Discretises the advection equation, ∂q/∂t + (u.∇)q = 0""" - def __init__(self, domain, function_space, field_name, Vu=None, **kwargs): + def __init__(self, domain, function_space, field_name, Vu=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -89,21 +88,20 @@ def __init__(self, domain, function_space, field_name, Vu=None, **kwargs): V = domain.spaces("HDiv", V=Vu, overwrite_space=True) else: V = domain.spaces("HDiv") - self.prescribed_fields("u", V) + u = self.prescribed_fields("u", V) test = TestFunction(function_space) q = Function(function_space) mass_form = time_derivative(inner(q, test)*dx) + transport_form = advection_form(test, q, u) - self.residual = subject( - mass_form + advection_form(domain, test, q, **kwargs), q - ) + self.residual = subject(mass_form + transport_form, q) class ContinuityEquation(PrognosticEquation): u"""Discretises the continuity equation, ∂q/∂t + ∇.(u*q) = 0""" - def __init__(self, domain, function_space, field_name, Vu=None, **kwargs): + def __init__(self, domain, function_space, field_name, Vu=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -114,7 +112,6 @@ def __init__(self, domain, function_space, field_name, Vu=None, **kwargs): Vu (:class:`FunctionSpace`, optional): the function space for the velocity field. If this is not specified, uses the HDiv spaces set up by the domain. Defaults to None. - **kwargs: any keyword arguments to be passed to the advection form. """ super().__init__(domain, function_space, field_name) @@ -122,15 +119,14 @@ def __init__(self, domain, function_space, field_name, Vu=None, **kwargs): V = domain.spaces("HDiv", V=Vu, overwrite_space=True) else: V = domain.spaces("HDiv") - self.prescribed_fields("u", V) + u = self.prescribed_fields("u", V) test = TestFunction(function_space) q = Function(function_space) mass_form = time_derivative(inner(q, test)*dx) + transport_form = continuity_form(test, q, u) - self.residual = subject( - mass_form + continuity_form(domain, test, q, **kwargs), q - ) + self.residual = subject(mass_form + transport_form, q) class DiffusionEquation(PrognosticEquation): @@ -165,7 +161,7 @@ class AdvectionDiffusionEquation(PrognosticEquation): u"""The advection-diffusion equation, ∂q/∂t + (u.∇)q = ∇.(κ∇q)""" def __init__(self, domain, function_space, field_name, Vu=None, - diffusion_parameters=None, **kwargs): + diffusion_parameters=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -177,7 +173,6 @@ def __init__(self, domain, function_space, field_name, Vu=None, velocity field. If this is Defaults to None. diffusion_parameters (:class:`DiffusionParameters`, optional): parameters describing the diffusion to be applied. - **kwargs: any keyword arguments to be passed to the advection form. """ super().__init__(domain, function_space, field_name) @@ -186,15 +181,16 @@ def __init__(self, domain, function_space, field_name, Vu=None, V = domain.spaces("HDiv", V=Vu, overwrite_space=True) else: V = domain.spaces("HDiv") - self.prescribed_fields("u", V) + u = self.prescribed_fields("u", V) test = TestFunction(function_space) q = Function(function_space) mass_form = time_derivative(inner(q, test)*dx) + transport_form = advection_form(test, q, u) self.residual = subject( mass_form - + advection_form(domain, test, q, **kwargs) + + transport_form + interior_penalty_diffusion_form( domain, test, q, diffusion_parameters), q ) @@ -448,16 +444,22 @@ def generate_tracer_transport_terms(self, domain, active_tracers): # By default return None if no tracers are to be transported adv_form = None no_tracer_transported = True + u_idx = self.field_names.index('u') + u = split(self.X)[u_idx] - for i, tracer in enumerate(active_tracers): + for _, tracer in enumerate(active_tracers): if tracer.transport_eqn != TransportEquationType.no_transport: idx = self.field_names.index(tracer.name) tracer_prog = split(self.X)[idx] tracer_test = self.tests[idx] if tracer.transport_eqn == TransportEquationType.advective: - tracer_adv = prognostic(advection_form(domain, tracer_test, tracer_prog), tracer.name) + tracer_adv = prognostic( + advection_form(tracer_test, tracer_prog, u), + tracer.name) elif tracer.transport_eqn == TransportEquationType.conservative: - tracer_adv = prognostic(continuity_form(domain, tracer_test, tracer_prog), tracer.name) + tracer_adv = prognostic( + continuity_form(tracer_test, tracer_prog, u), + tracer.name) else: raise ValueError(f'Transport eqn {tracer.transport_eqn} not recognised') @@ -492,7 +494,6 @@ def __init__(self, domain, function_space, field_name, Vu=None, active_tracers (list, optional): a list of `ActiveTracer` objects that encode the metadata for any active tracers to be included in the equations. Defaults to None. - **kwargs: any keyword arguments to be passed to the advection form. """ self.field_names = [field_name] @@ -518,16 +519,15 @@ def __init__(self, domain, function_space, field_name, Vu=None, V = domain.spaces("HDiv", V=Vu, overwrite_space=True) else: V = domain.spaces("HDiv") - self.prescribed_fields("u", V) + u = self.prescribed_fields("u", V) self.tests = TestFunctions(W) self.X = Function(W) mass_form = self.generate_mass_terms() + transport_form = advection_form(self.tests[0], split(self.X)[0], u) - self.residual = subject( - mass_form + advection_form(domain, self.tests[0], split(self.X)[0], **kwargs), self.X - ) + self.residual = subject(mass_form + transport_form, self.X) # Add transport of tracers if len(active_tracers) > 0: @@ -569,8 +569,8 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, the 'D' transport term and the pressure gradient term. u_transport_option (str, optional): specifies the transport term used for the velocity equation. Supported options are: - 'vector_invariant_form', 'vector_advection_form', - 'vector_manifold_advection_form' and 'circulation_form'. + 'vector_invariant_form', 'vector_advection_form', and + 'circulation_form'. Defaults to 'vector_invariant_form'. no_normal_flow_bc_ids (list, optional): a list of IDs of domain boundaries at which no normal flow will be enforced. Defaults to @@ -625,12 +625,12 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, # -------------------------------------------------------------------- # # Velocity transport term -- depends on formulation if u_transport_option == "vector_invariant_form": + raise NotImplementedError u_adv = prognostic(vector_invariant_form(domain, w, u), "u") elif u_transport_option == "vector_advection_form": - u_adv = prognostic(advection_form(domain, w, u), "u") - elif u_transport_option == "vector_manifold_advection_form": - u_adv = prognostic(vector_manifold_advection_form(domain, w, u), "u") + u_adv = prognostic(advection_form(w, u, u), "u") elif u_transport_option == "circulation_form": + raise NotImplementedError ke_form = prognostic(kinetic_energy_form(domain, w, u), "u") ke_form = transport.remove(ke_form) ke_form = ke_form.label_map( @@ -643,7 +643,7 @@ 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(domain, phi, D), "D") + 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(domain, phi, H).label_map( @@ -662,7 +662,7 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, if self.thermal: gamma = self.tests[2] b = split(self.X)[2] - b_adv = prognostic(advection_form(domain, gamma, b), "b") + b_adv = prognostic(inner(gamma, inner(u, grad(b)))*dx, "b") # TODO: implement correct linearisation adv_form += subject(b_adv, self.X) @@ -767,8 +767,8 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, the 'D' transport term, pressure gradient and Coriolis terms. u_transport_option (str, optional): specifies the transport term used for the velocity equation. Supported options are: - 'vector_invariant_form', 'vector_advection_form', - 'vector_manifold_advection_form' and 'circulation_form'. + 'vector_invariant_form', 'vector_advection_form' and + 'circulation_form'. Defaults to 'vector_invariant_form'. no_normal_flow_bc_ids (list, optional): a list of IDs of domain boundaries at which no normal flow will be enforced. Defaults to @@ -834,8 +834,8 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, scalar transport terms. u_transport_option (str, optional): specifies the transport term used for the velocity equation. Supported options are: - 'vector_invariant_form', 'vector_advection_form', - 'vector_manifold_advection_form' and 'circulation_form'. + 'vector_invariant_form', 'vector_advection_form' and + 'circulation_form'. Defaults to 'vector_invariant_form'. diffusion_options (:class:`DiffusionOptions`, optional): any options to specify for applying diffusion terms to variables. Defaults @@ -892,9 +892,7 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, if u_transport_option == "vector_invariant_form": u_adv = prognostic(vector_invariant_form(domain, w, u), "u") elif u_transport_option == "vector_advection_form": - u_adv = prognostic(advection_form(domain, w, u), "u") - elif u_transport_option == "vector_manifold_advection_form": - u_adv = prognostic(vector_manifold_advection_form(domain, w, u), "u") + u_adv = prognostic(advection_form(w, u, u), "u") elif u_transport_option == "circulation_form": ke_form = prognostic(kinetic_energy_form(domain, w, u), "u") ke_form = transport.remove(ke_form) @@ -904,11 +902,12 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, t.form, {t.get(transporting_velocity): u}), t.labels)) ke_form = transporting_velocity.remove(ke_form) u_adv = prognostic(advection_equation_circulation_form(domain, w, u), "u") + ke_form + raise NotImplementedError else: raise ValueError("Invalid u_transport_option: %s" % u_transport_option) # Density transport (conservative form) - rho_adv = prognostic(continuity_form(domain, phi, rho), "rho") + rho_adv = prognostic(continuity_form(phi, rho, u), "rho") # Transport term needs special linearisation if self.linearisation_map(rho_adv.terms[0]): linear_rho_adv = linear_continuity_form(domain, phi, rho_bar).label_map( @@ -918,7 +917,7 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, rho_adv = linearisation(rho_adv, linear_rho_adv) # Potential temperature transport (advective form) - theta_adv = prognostic(advection_form(domain, gamma, theta), "theta") + theta_adv = prognostic(advection_form(gamma, theta, u), "theta") # Transport term needs special linearisation if self.linearisation_map(theta_adv.terms[0]): linear_theta_adv = linear_advection_form(domain, gamma, theta_bar).label_map( @@ -1081,8 +1080,8 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, scalar transport terms. u_transport_option (str, optional): specifies the transport term used for the velocity equation. Supported options are: - 'vector_invariant_form', 'vector_advection_form', - 'vector_manifold_advection_form' and 'circulation_form'. + 'vector_invariant_form', 'vector_advection_form' and + 'circulation_form'. Defaults to 'vector_invariant_form'. diffusion_options (:class:`DiffusionOptions`, optional): any options to specify for applying diffusion terms to variables. Defaults @@ -1179,8 +1178,8 @@ def __init__(self, domain, parameters, Omega=None, scalar transport terms. u_transport_option (str, optional): specifies the transport term used for the velocity equation. Supported options are: - 'vector_invariant_form', 'vector_advection_form', - 'vector_manifold_advection_form' and 'circulation_form'. + 'vector_invariant_form', 'vector_advection_form' and + 'circulation_form'. Defaults to 'vector_invariant_form'. no_normal_flow_bc_ids (list, optional): a list of IDs of domain boundaries at which no normal flow will be enforced. Defaults to @@ -1232,10 +1231,9 @@ def __init__(self, domain, parameters, Omega=None, # Velocity transport term -- depends on formulation if u_transport_option == "vector_invariant_form": u_adv = prognostic(vector_invariant_form(domain, w, u), "u") + raise NotImplementedError elif u_transport_option == "vector_advection_form": - u_adv = prognostic(advection_form(domain, w, u), "u") - elif u_transport_option == "vector_manifold_advection_form": - u_adv = prognostic(vector_manifold_advection_form(domain, w, u), "u") + u_adv = prognostic(advection_form(w, u, u), "u") elif u_transport_option == "circulation_form": ke_form = prognostic(kinetic_energy_form(domain, w, u), "u") ke_form = transport.remove(ke_form) @@ -1245,11 +1243,12 @@ def __init__(self, domain, parameters, Omega=None, t.form, {t.get(transporting_velocity): u}), t.labels)) ke_form = transporting_velocity.remove(ke_form) u_adv = prognostic(advection_equation_circulation_form(domain, w, u), "u") + ke_form + raise NotImplementedError else: raise ValueError("Invalid u_transport_option: %s" % u_transport_option) # Buoyancy transport - b_adv = prognostic(advection_form(domain, gamma, b), "b") + b_adv = prognostic(advection_form(gamma, b, u), "b") if self.linearisation_map(b_adv.terms[0]): linear_b_adv = linear_advection_form(domain, gamma, b_bar).label_map( lambda t: t.has_label(transporting_velocity), diff --git a/gusto/labels.py b/gusto/labels.py index 59ed09235..479b09dfb 100644 --- a/gusto/labels.py +++ b/gusto/labels.py @@ -225,7 +225,7 @@ def repl(t): transport = Label("transport", validator=lambda value: type(value) == TransportEquationType) diffusion = Label("diffusion") physics = Label("physics", validator=lambda value: type(value) == MethodType) -transporting_velocity = Label("transporting_velocity", validator=lambda value: type(value) == Function) +transporting_velocity = Label("transporting_velocity", validator=lambda value: type(value) in [Function, ufl.tensors.ListTensor]) subject = Label("subject", validator=lambda value: type(value) == Function) prognostic = Label("prognostic", validator=lambda value: type(value) == str) pressure_gradient = Label("pressure_gradient") diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 7a832f5ac..b448f7a33 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -1,17 +1,18 @@ """Classes for controlling the timestepping loop.""" from abc import ABCMeta, abstractmethod, abstractproperty -from firedrake import Function, Projector, Constant +from firedrake import Function, Projector, Constant, split from pyop2.profiling import timed_stage from gusto.configuration import logger from gusto.equations import PrognosticEquationSet from gusto.forcing import Forcing from gusto.fml.form_manipulation_labelling import drop, Label, Term -from gusto.labels import (transport, diffusion, time_derivative, - linearisation, prognostic, physics) +from gusto.labels import (transport, diffusion, time_derivative, linearisation, + prognostic, physics, transporting_velocity) from gusto.linear_solvers import LinearTimesteppingSolver from gusto.fields import TimeLevelFields, StateFields from gusto.time_discretisation import ExplicitTimeDiscretisation +import ufl __all__ = ["Timestepper", "SplitPhysicsTimestepper", "SemiImplicitQuasiNewton", "PrescribedTransport"] @@ -72,32 +73,57 @@ def get_initial_timesteps(self): # Return None if this is not applicable return self.scheme.initial_timesteps if can_get else None - def transport_discretisation_setup_equation(self, equation): + def setup_equation_transport(self, equation): """ - Sets up the transport discretisation for an equation, by the setting the - form used for transport in the equation. + Sets up the transport methods for an equation, by the setting the + forms used for transport in the equation. Args: equation (:class:`PrognosticEquation`): the equation that the - transport discretisation is to be applied to. + transport method is to be applied to. """ - if self.transport_discretisations is not None: - for transport_discretisation in self.transport_discretisations: - transport_discretisation.replace_transport_form(equation) - # TODO: raise a warning here if transport discretisations aren't provided - - def transport_discretisation_setup_scheme(self, transport_discretisation, scheme): + # Check that we have specified all transport terms + transport_residual = equation.residual.label_map( + lambda t: t.has_label(transport), map_if_false=drop + ) + transported_variables = [t.get(prognostic) for t in transport_residual.terms] + transport_method_variables = [method.variable for method in self.transport_methods] + for variable in transported_variables: + if variable not in transport_method_variables: + logger.warning(f'Variable {variable} has a transport term but ' + + 'no transport method has been specified. ' + + 'Using default transport form.') + + # Replace transport forms in equation + if self.transport_methods is not None: + for transport_method in self.transport_methods: + transport_method.replace_transport_form(equation) + + def setup_transporting_velocity(self, scheme): """ - Sets up the transport discretisation, by the setting the transporting - velocity used in the time discretisation. + Set up the time discretisation by replacing the transporting velocity + used by the appropriate one for this time loop. Args: - transport_discretisation (:class:`TransportScheme`): the transport - discretisation to be set up. - scheme (:class:`TimeDiscretisation`): the time discretisation used - with this transport scheme. + scheme (:class:`TimeDiscretisation`): the time discretisation whose + transport term should be replaced with the transport term of + this discretisation. """ - transport_discretisation.setup(self.transporting_velocity, scheme) + + if self.transporting_velocity == "prognostic": + # Use the prognostic wind variable as the transporting velocity + u_idx = self.equation.field_names.index('u') + uadv = split(self.equation.X)[u_idx] + else: + uadv = self.transporting_velocity + + scheme.residual = scheme.residual.label_map( + lambda t: t.has_label(transporting_velocity), + map_if_true=lambda t: + Term(ufl.replace(t.form, {t.get(transporting_velocity): uadv}), t.labels) + ) + + scheme.residual = transporting_velocity.update_value(scheme.residual, uadv) def run(self, t, tmax, pick_up=False): """ @@ -182,22 +208,23 @@ class Timestepper(BaseTimestepper): Implements a timeloop by applying a scheme to a prognostic equation. """ - def __init__(self, equation, scheme, io, transport_discretisations=None): + def __init__(self, equation, scheme, io, transport_methods=None): """ Args: equation (:class:`PrognosticEquation`): the prognostic equation scheme (:class:`TimeDiscretisation`): the scheme to use to timestep the prognostic equation io (:class:`IO`): the model's object for controlling input/output. - transport_discretisations (iter,optional): a list of objects - describing the discretisations of transport terms to be used. - Defaults to None. + transport_methods (iter,optional): a list of objects describing the + methods to use for discretising transport terms for each + transported variable. Defaults to None, in which case the + transport term follows the discretisation in the equation. """ self.scheme = scheme - if transport_discretisations is not None: - self.transport_discretisations = transport_discretisations + if transport_methods is not None: + self.transport_methods = transport_methods else: - self.transport_discretisations = [] + self.transport_methods = [] super().__init__(equation=equation, io=io) @@ -211,10 +238,9 @@ def setup_fields(self): *self.io.output.dumplist) def setup_scheme(self): - self.transport_discretisation_setup_equation(self.equation) + self.setup_equation_transport(self.equation) self.scheme.setup(self.equation) - for transport_discretisation in self.transport_discretisations: - self.transport_discretisation_setup_scheme(transport_discretisation, self.scheme) + self.setup_transporting_velocity(self.scheme) def timestep(self): """ @@ -234,7 +260,7 @@ class SplitPhysicsTimestepper(Timestepper): scheme to be applied to the physics terms than the prognostic equation. """ - def __init__(self, equation, scheme, io, transport_discretisations=None, + def __init__(self, equation, scheme, io, transport_methods=None, physics_schemes=None): """ Args: @@ -242,9 +268,10 @@ def __init__(self, equation, scheme, io, transport_discretisations=None, scheme (:class:`TimeDiscretisation`): the scheme to use to timestep the prognostic equation io (:class:`IO`): the model's object for controlling input/output. - transport_discretisations (iter,optional): a list of objects - describing the discretisations of transport terms to be used. - Defaults to None. + transport_methods (iter,optional): a list of objects describing the + methods to use for discretising transport terms for each + transported variable. Defaults to None, in which case the + transport term follows the discretisation in the equation. physics_schemes: (list, optional): a list of :class:`Physics` and :class:`TimeDiscretisation` options describing physical parametrisations and timestepping schemes to use for each. @@ -253,7 +280,7 @@ def __init__(self, equation, scheme, io, transport_discretisations=None, """ super().__init__(equation, scheme, io, - transport_discretisations=transport_discretisations) + transport_methods=transport_methods) if physics_schemes is not None: self.physics_schemes = physics_schemes @@ -271,14 +298,13 @@ def transporting_velocity(self): return "prognostic" def setup_scheme(self): - self.transport_discretisation_setup_equation(self.equation) + self.setup_equation_transport(self.equation) # Go through and label all non-physics terms with a "dynamics" label dynamics = Label('dynamics') self.equation.label_terms(lambda t: not any(t.has_label(time_derivative, physics)), dynamics) apply_bcs = True self.scheme.setup(self.equation, apply_bcs, dynamics) - for transport_discretisation in self.transport_discretisations: - self.transport_discretisation_setup_scheme(transport_discretisation, self.scheme) + self.setup_transporting_velocity(self.scheme) def timestep(self): @@ -300,7 +326,7 @@ class SemiImplicitQuasiNewton(BaseTimestepper): """ def __init__(self, equation_set, io, transport_schemes, - transport_discretisations, + transport_methods, auxiliary_equations_and_schemes=None, linear_solver=None, diffusion_schemes=None, @@ -314,8 +340,8 @@ def __init__(self, equation_set, io, transport_schemes, transport_schemes: iterable of ``(field_name, scheme)`` pairs indicating the name of the field (str) to transport, and the :class:`TimeDiscretisation` to use - transport_discretisations (iter): a list of objects describing the - spatial discretisations of transport terms to be used. + transport_methods (iter): a list of objects describing the spatial + discretisations of transport terms to be used. auxiliary_equations_and_schemes: iterable of ``(equation, scheme)`` pairs indicating any additional equations to be solved and the scheme to use to solve them. Defaults to None. @@ -340,7 +366,7 @@ def __init__(self, equation_set, io, transport_schemes, if kwargs: raise ValueError("unexpected kwargs: %s" % list(kwargs.keys())) - self.transport_discretisations = [] + self.transport_methods = [] if physics_schemes is not None: self.physics_schemes = physics_schemes @@ -354,15 +380,13 @@ def __init__(self, equation_set, io, transport_schemes, for scheme in transport_schemes: assert scheme.nlevels == 1, "multilevel schemes not supported as part of this timestepping loop" assert scheme.field_name in equation_set.field_names - # Active transport will be a tuple of (field_name, scheme, discretisation) - discretisation_found = False - for transport_discretisation in transport_discretisations: - if scheme.field_name == transport_discretisation.variable: - discretisation_found = True - self.transport_discretisations.append(transport_discretisation) - self.active_transport.append((scheme.field_name, scheme, transport_discretisation)) - assert discretisation_found, 'No transport discretisation found ' \ - + f'for variable {scheme.field_name}' + # Check that there is a corresponding transport method + method_found = False + for transport_method in transport_methods: + if scheme.field_name == transport_method.variable: + method_found = True + self.transport_methods.append(transport_method) + assert method_found, f'No transport method found for variable {scheme.field_name}' self.diffusion_schemes = [] if diffusion_schemes is not None: @@ -385,14 +409,9 @@ def __init__(self, equation_set, io, transport_schemes, super().__init__(equation_set, io) for aux_eqn, aux_scheme in self.auxiliary_equations_and_schemes: - self.transport_discretisation_setup(aux_eqn) + self.transport_method_setup(aux_eqn) aux_scheme.setup(aux_eqn) - # Auxiliary transport discretisations - for transport_discretisation in transport_discretisations: - if transport_discretisation.variable == aux_scheme.field_name: - self.transport_discretisation_setup_scheme(transport_discretisation, aux_scheme) - # TODO: add a check to make sure that auxiliary eqns have transport - # discretisations, but only if they need them! + self.setup_transporting_velocity(aux_scheme) self.tracers_to_copy = [] for name in equation_set.field_names: @@ -449,11 +468,10 @@ def setup_scheme(self): # TODO: apply_bcs should be False for advection but this means # tests with KGOs fail apply_bcs = True - self.transport_discretisation_setup_equation(self.equation) - for _, scheme, transport_discretisation in self.active_transport: - print(f'Setting up transport for {scheme.field_name} {transport_discretisation.variable}') + self.setup_equation_transport(self.equation) + for _, scheme in self.active_transport: scheme.setup(self.equation, apply_bcs, transport) - self.transport_discretisation_setup_scheme(transport_discretisation, scheme) + self.setup_transporting_velocity(scheme) apply_bcs = True for _, scheme in self.diffusion_schemes: @@ -550,7 +568,7 @@ class PrescribedTransport(Timestepper): """ Implements a timeloop with a prescibed transporting velocity """ - def __init__(self, equation, scheme, io, transport_discretisations=None, + def __init__(self, equation, scheme, io, transport_methods=None, physics_schemes=None, prescribed_transporting_velocity=None): """ Args: @@ -558,9 +576,10 @@ def __init__(self, equation, scheme, io, transport_discretisations=None, scheme (:class:`TimeDiscretisation`): the scheme to use to timestep the prognostic equation io (:class:`IO`): the model's object for controlling input/output. - transport_discretisations (iter,optional): a list of objects - describing the discretisations of transport terms to be used. - Defaults to None. + transport_methods (iter,optional): a list of objects describing the + methods to use for discretising transport terms for each + transported variable. Defaults to None, in which case the + transport term follows the discretisation in the equation. physics_schemes: (list, optional): a list of :class:`Physics` and :class:`TimeDiscretisation` options describing physical parametrisations and timestepping schemes to use for each. @@ -575,7 +594,7 @@ def __init__(self, equation, scheme, io, transport_discretisations=None, """ super().__init__(equation, scheme, io, - transport_discretisations=transport_discretisations) + transport_methods=transport_methods) if physics_schemes is not None: self.physics_schemes = physics_schemes diff --git a/gusto/transport_forms.py b/gusto/transport_forms.py index 2d438ca4f..63a6bce3f 100644 --- a/gusto/transport_forms.py +++ b/gusto/transport_forms.py @@ -10,7 +10,8 @@ __all__ = ["advection_form", "continuity_form", "vector_invariant_form", "vector_manifold_advection_form", "kinetic_energy_form", - "advection_equation_circulation_form", "linear_continuity_form"] + "advection_equation_circulation_form", "linear_continuity_form", + "upwind_advection_form", "upwind_continuity_form"] def linear_advection_form(domain, test, qbar): @@ -68,11 +69,50 @@ def linear_continuity_form(domain, test, qbar, facet_term=False): return transport(form, TransportEquationType.conservative) - -def advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): +def advection_form(test, q, ubar): u""" The form corresponding to the advective transport operator. + This describes (u.∇)q, for transporting velocity u and transported q. + + Args: + test (:class:`TestFunction`): the test function. + q (:class:`ufl.Expr`): the variable to be transported. + ubar (:class:`ufl.Expr`): the transporting velocity. + + Returns: + class:`LabelledForm`: a labelled transport form. + """ + + L = inner(test, inner(ubar, grad(q)))*dx + form = transporting_velocity(L, ubar) + + return transport(form, TransportEquationType.advective) + +def continuity_form(test, q, ubar): + u""" + The form corresponding to the continuity transport operator. + + This describes ∇.(u*q), for transporting velocity u and transported q. + + Args: + test (:class:`TestFunction`): the test function. + q (:class:`ufl.Expr`): the variable to be transported. + ubar (:class:`ufl.Expr`): the transporting velocity. + + Returns: + class:`LabelledForm`: a labelled transport form. + """ + + L = inner(test, div(outer(q, ubar)))*dx + form = transporting_velocity(L, ubar) + + return transport(form, TransportEquationType.conservative) + +def upwind_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): + u""" + The form corresponding to the DG upwind advective transport operator. + This discretises (u.∇)q, for transporting velocity u and transported variable q. An upwind discretisation is used for the facet terms when the form is integrated by parts. @@ -127,9 +167,9 @@ def advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): return ibp_label(transport(form, TransportEquationType.advective), ibp) -def continuity_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): +def upwind_continuity_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): u""" - The form corresponding to the continuity transport operator. + The form corresponding to the DG upwind continuity transport operator. This discretises ∇.(u*q), for transporting velocity u and transported variable q. An upwind discretisation is used for the facet terms when the @@ -209,7 +249,7 @@ def vector_manifold_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, o class:`LabelledForm`: a labelled transport form. """ - L = advection_form(domain, test, q, ibp, outflow) + L = upwind_advection_form(domain, test, q, ibp, outflow) # TODO: there should maybe be a restriction on IBP here Vu = domain.spaces("HDiv") @@ -247,7 +287,7 @@ def vector_manifold_continuity_form(domain, test, q, ibp=IntegrateByParts.ONCE, class:`LabelledForm`: a labelled transport form. """ - L = continuity_form(domain, test, q, ibp, outflow) + L = upwind_continuity_form(domain, test, q, ibp, outflow) Vu = domain.spaces("HDiv") dS_ = (dS_v + dS_h) if Vu.extruded else dS diff --git a/gusto/transport_schemes.py b/gusto/transport_methods.py similarity index 64% rename from gusto/transport_schemes.py rename to gusto/transport_methods.py index 3b1c2eb3a..1bf955479 100644 --- a/gusto/transport_schemes.py +++ b/gusto/transport_methods.py @@ -1,19 +1,18 @@ """ -Defines TransportScheme objects, which are used to solve a transport problem. +Defines TransportMethod objects, which are used to solve a transport problem. """ from firedrake import split from gusto.configuration import IntegrateByParts, TransportEquationType -from gusto.fml import Term, keep, drop, Label, LabelledForm -from gusto.labels import transporting_velocity, prognostic, transport, subject +from gusto.fml import Term, keep, drop +from gusto.labels import prognostic, transport from gusto.transport_forms import * -import ufl __all__ = ["DGUpwind", "SUPGTransport"] # TODO: settle on name: discretisation vs scheme -class TransportScheme(object): +class TransportMethod(object): """ The base object for describing a transport scheme. """ @@ -34,16 +33,18 @@ def __init__(self, equation, variable): self.test = equation.tests[variable_idx] self.field = split(equation.X)[variable_idx] - # Find the original transport term to be used - # TODO: do we need this? - self.original_form = equation.residual.label_map( + # Find the original transport term to be used, which we use to extract + # information about the transport equation type + original_form = equation.residual.label_map( lambda t: t.has_label(transport) and t.get(prognostic) == variable, map_if_true=keep, map_if_false=drop) - num_terms = len(self.original_form.terms) + num_terms = len(original_form.terms) assert num_terms == 1, \ f'Unable to find transport term for {variable}. {num_terms} found' + self.transport_equation_type = original_form.terms[0].get(transport) + def replace_transport_form(self, equation): """ Replaces the form for the transport term in the equation with the @@ -61,41 +62,7 @@ def replace_transport_form(self, equation): map_if_true=lambda t: Term(self.form.form, t.labels)) - def setup(self, uadv, equation): - """ - Set up the transport scheme by replacing the transporting velocity used - in the form. - - Args: - uadv (:class:`ufl.Expr`, optional): the transporting velocity. - Defaults to None. - equation (:class:`PrognosticEquation`): the equation or scheme whose - transport term should be replaced with the transport term of - this discretisation. - """ - - assert self.form.terms[0].has_label(transporting_velocity), \ - 'Cannot set up transport scheme on a term that has no transporting velocity' - - if uadv == "prognostic": - # Find prognostic wind field - uadv = split(self.original_form.terms[0].get(subject))[0] - - import pdb; pdb.set_trace() - - equation.residual = equation.residual.label_map( - lambda t: t.has_label(transporting_velocity), - map_if_true=lambda t: - Term(ufl.replace(t.form, {t.get(transporting_velocity): uadv}), t.labels) - ) - - import pdb; pdb.set_trace() - - equation.residual = transporting_velocity.update_value(equation.residual, uadv) - - import pdb; pdb.set_trace() - -class DGUpwind(TransportScheme): +class DGUpwind(TransportMethod): """ The Discontinuous Galerkin Upwind transport scheme. @@ -126,25 +93,25 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, # Determine appropriate form to use # -------------------------------------------------------------------- # - if self.original_form.terms[0].get(transport) == TransportEquationType.advective: + if self.transport_equation_type == TransportEquationType.advective: if vector_manifold_correction: form = vector_manifold_advection_form(self.domain, self.test, self.field, ibp=ibp, outflow=outflow) else: - form = advection_form(self.domain, self.test, self.field, - ibp=ibp, outflow=outflow) + form = upwind_advection_form(self.domain, self.test, self.field, + ibp=ibp, outflow=outflow) - elif self.original_form.terms[0].get(transport) == TransportEquationType.conservative: + elif self.transport_equation_type == TransportEquationType.conservative: if vector_manifold_correction: form = vector_manifold_continuity_form(self.domain, self.test, self.field, ibp=ibp, outflow=outflow) else: - form = continuity_form(self.domain, self.test, self.field, - ibp=ibp, outflow=outflow) + form = upwind_continuity_form(self.domain, self.test, self.field, + ibp=ibp, outflow=outflow) - elif self.original_form.terms[0].get(transport) == TransportEquationType.vector_invariant: + elif self.transport_equation_type == TransportEquationType.vector_invariant: if outflow: raise NotImplementedError('Outflow not implemented for upwind vector invariant') form = vector_invariant_form(self.domain, self.test, self.field, ibp=ibp) @@ -156,5 +123,5 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, self.form = form -class SUPGTransport(TransportScheme): +class SUPGTransport(TransportMethod): pass \ No newline at end of file From 7935f9cbafb176e146535063e5423dd753008df9 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Fri, 7 Jul 2023 15:37:06 +0100 Subject: [PATCH 09/20] try to get transport-only tests working, still a bug in form used in scheme --- gusto/equations.py | 42 +++++++++-------- gusto/time_discretisation.py | 32 ++++++------- gusto/timeloop.py | 15 +++---- gusto/transport_forms.py | 45 +++++++++++++++++-- gusto/transport_methods.py | 38 +++++++++++----- .../transport/test_dg_transport.py | 6 ++- 6 files changed, 115 insertions(+), 63 deletions(-) diff --git a/gusto/equations.py b/gusto/equations.py index ca0c25d66..c53d4a1b7 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -49,6 +49,9 @@ def __init__(self, domain, function_space, field_name): assert hasattr(self, "field_names") for fname in self.field_names: self.bcs[fname] = [] + else: + # To avoid confusion, only add "self.test" when not mixed FS + self.test = TestFunction(function_space) self.bcs[field_name] = [] @@ -90,12 +93,12 @@ def __init__(self, domain, function_space, field_name, Vu=None): V = domain.spaces("HDiv") u = self.prescribed_fields("u", V) - test = TestFunction(function_space) + test = self.test q = Function(function_space) mass_form = time_derivative(inner(q, test)*dx) transport_form = advection_form(test, q, u) - self.residual = subject(mass_form + transport_form, q) + self.residual = prognostic(subject(mass_form + transport_form, q), field_name) class ContinuityEquation(PrognosticEquation): @@ -121,12 +124,12 @@ def __init__(self, domain, function_space, field_name, Vu=None): V = domain.spaces("HDiv") u = self.prescribed_fields("u", V) - test = TestFunction(function_space) + test = self.test q = Function(function_space) mass_form = time_derivative(inner(q, test)*dx) transport_form = continuity_form(test, q, u) - self.residual = subject(mass_form + transport_form, q) + self.residual = prognostic(subject(mass_form + transport_form, q), field_name) class DiffusionEquation(PrognosticEquation): @@ -146,15 +149,13 @@ def __init__(self, domain, function_space, field_name, """ super().__init__(domain, function_space, field_name) - test = TestFunction(function_space) + test = self.test q = Function(function_space) mass_form = time_derivative(inner(q, test)*dx) + diffusion_form = interior_penalty_diffusion_form(domain, test, q, + diffusion_parameters) - self.residual = subject( - mass_form - + interior_penalty_diffusion_form( - domain, test, q, diffusion_parameters), q - ) + self.residual = prognostic(subject(mass_form + diffusion_form, q), field_name) class AdvectionDiffusionEquation(PrognosticEquation): @@ -183,17 +184,15 @@ def __init__(self, domain, function_space, field_name, Vu=None, V = domain.spaces("HDiv") u = self.prescribed_fields("u", V) - test = TestFunction(function_space) + test = self.test q = Function(function_space) mass_form = time_derivative(inner(q, test)*dx) transport_form = advection_form(test, q, u) + diffusion_form = interior_penalty_diffusion_form(domain, test, q, + diffusion_parameters) - self.residual = subject( - mass_form - + transport_form - + interior_penalty_diffusion_form( - domain, test, q, diffusion_parameters), q - ) + self.residual = prognostic(subject( + mass_form + transport_form + diffusion_form, q), field_name) class PrognosticEquationSet(PrognosticEquation, metaclass=ABCMeta): @@ -525,7 +524,7 @@ def __init__(self, domain, function_space, field_name, Vu=None, self.X = Function(W) mass_form = self.generate_mass_terms() - transport_form = advection_form(self.tests[0], split(self.X)[0], u) + transport_form = prognostic(advection_form(self.tests[0], split(self.X)[0], u), field_name) self.residual = subject(mass_form + transport_form, self.X) @@ -625,8 +624,7 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, # -------------------------------------------------------------------- # # Velocity transport term -- depends on formulation if u_transport_option == "vector_invariant_form": - raise NotImplementedError - u_adv = prognostic(vector_invariant_form(domain, w, u), "u") + u_adv = prognostic(vector_invariant_form(domain, w, u, u), "u") elif u_transport_option == "vector_advection_form": u_adv = prognostic(advection_form(w, u, u), "u") elif u_transport_option == "circulation_form": @@ -890,7 +888,7 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, # -------------------------------------------------------------------- # # Velocity transport term -- depends on formulation if u_transport_option == "vector_invariant_form": - u_adv = prognostic(vector_invariant_form(domain, w, u), "u") + u_adv = prognostic(vector_invariant_form(domain, w, u, u), "u") elif u_transport_option == "vector_advection_form": u_adv = prognostic(advection_form(w, u, u), "u") elif u_transport_option == "circulation_form": @@ -1230,7 +1228,7 @@ def __init__(self, domain, parameters, Omega=None, # -------------------------------------------------------------------- # # Velocity transport term -- depends on formulation if u_transport_option == "vector_invariant_form": - u_adv = prognostic(vector_invariant_form(domain, w, u), "u") + u_adv = prognostic(vector_invariant_form(domain, w, u, u), "u") raise NotImplementedError elif u_transport_option == "vector_advection_form": u_adv = prognostic(advection_form(w, u, u), "u") diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 36d921ee9..a0638eedf 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -254,18 +254,18 @@ def __init__(self, domain, field_name=None, subcycles=None, self.subcycles = subcycles - def setup(self, equation, uadv, apply_bcs=True, *active_labels): + 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. - uadv (:class:`ufl.Expr`, optional): the transporting velocity. - Defaults to None. + apply_bcs (bool, optional): whether boundary conditions are to be + applied. Defaults to True. *active_labels (:class:`Label`): labels indicating which terms of the equation to include. """ - super().setup(equation, uadv, apply_bcs, *active_labels) + super().setup(equation, apply_bcs, *active_labels) # if user has specified a number of subcycles, then save this # and rescale dt accordingly; else perform just one cycle using dt @@ -419,18 +419,16 @@ class RK4(ExplicitTimeDiscretisation): where superscripts indicate the time-level. """ - def setup(self, equation, uadv, *active_labels): + def setup(self, equation, *active_labels): """ Set up the time discretisation based on the equation. Args: equation (:class:`PrognosticEquation`): the model's equation. - uadv (:class:`ufl.Expr`, optional): the transporting velocity. - Defaults to None. *active_labels (:class:`Label`): labels indicating which terms of the equation to include. """ - super().setup(equation, uadv, *active_labels) + super().setup(equation, *active_labels) self.k1 = Function(self.fs) self.k2 = Function(self.fs) @@ -775,9 +773,8 @@ def __init__(self, domain, field_name=None, solver_parameters=None, def nlevels(self): pass - def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): - super().setup(equation=equation, uadv=uadv, apply_bcs=apply_bcs, - *active_labels) + def setup(self, equation, apply_bcs=True, *active_labels): + super().setup(equation=equation, apply_bcs=apply_bcs, *active_labels) for n in range(self.nlevels, 1, -1): setattr(self, "xnm%i" % (n-1), Function(self.fs)) @@ -921,8 +918,8 @@ def __init__(self, domain, gamma, field_name=None, self.gamma = gamma - def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): - super().setup(equation, uadv, apply_bcs, *active_labels) + def setup(self, equation, apply_bcs=True, *active_labels): + super().setup(equation, apply_bcs, *active_labels) self.xnpg = Function(self.fs) self.xn = Function(self.fs) @@ -1118,8 +1115,8 @@ def __init__(self, domain, order, field_name=None, self.order = order - def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): - super().setup(equation=equation, uadv=uadv, apply_bcs=apply_bcs, + def setup(self, equation, apply_bcs=True, *active_labels): + super().setup(equation=equation, apply_bcs=apply_bcs, *active_labels) self.x = [Function(self.fs) for i in range(self.nlevels)] @@ -1248,9 +1245,8 @@ def __init__(self, domain, order, field_name=None, self.order = order - def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): - super().setup(equation=equation, uadv=uadv, apply_bcs=apply_bcs, - *active_labels) + def setup(self, equation, apply_bcs=True, *active_labels): + super().setup(equation=equation, apply_bcs=apply_bcs, *active_labels) self.x = [Function(self.fs) for i in range(self.nlevels)] diff --git a/gusto/timeloop.py b/gusto/timeloop.py index b448f7a33..0d716ccdf 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -251,6 +251,7 @@ def timestep(self): x_in = [x(name) for x in self.x.previous[-self.scheme.nlevels:]] self.scheme.apply(xnp1(name), *x_in) + import pdb; pdb.set_trace() class SplitPhysicsTimestepper(Timestepper): @@ -566,20 +567,18 @@ def run(self, t, tmax, pick_up=False): class PrescribedTransport(Timestepper): """ - Implements a timeloop with a prescibed transporting velocity + Implements a timeloop with a prescibed transporting velocity. """ - def __init__(self, equation, scheme, io, transport_methods=None, + def __init__(self, equation, scheme, transport_method, io, physics_schemes=None, prescribed_transporting_velocity=None): """ Args: equation (:class:`PrognosticEquation`): the prognostic equation scheme (:class:`TimeDiscretisation`): the scheme to use to timestep - the prognostic equation + the prognostic equation. + transport_method (:class:`TransportMethod`): describes the method + used for discretising the transport term. io (:class:`IO`): the model's object for controlling input/output. - transport_methods (iter,optional): a list of objects describing the - methods to use for discretising transport terms for each - transported variable. Defaults to None, in which case the - transport term follows the discretisation in the equation. physics_schemes: (list, optional): a list of :class:`Physics` and :class:`TimeDiscretisation` options describing physical parametrisations and timestepping schemes to use for each. @@ -594,7 +593,7 @@ def __init__(self, equation, scheme, io, transport_methods=None, """ super().__init__(equation, scheme, io, - transport_methods=transport_methods) + transport_methods=[transport_method]) if physics_schemes is not None: self.physics_schemes = physics_schemes diff --git a/gusto/transport_forms.py b/gusto/transport_forms.py index 63a6bce3f..a8d653b9c 100644 --- a/gusto/transport_forms.py +++ b/gusto/transport_forms.py @@ -11,7 +11,8 @@ __all__ = ["advection_form", "continuity_form", "vector_invariant_form", "vector_manifold_advection_form", "kinetic_energy_form", "advection_equation_circulation_form", "linear_continuity_form", - "upwind_advection_form", "upwind_continuity_form"] + "upwind_advection_form", "upwind_continuity_form", + "upwind_vector_invariant_form"] def linear_advection_form(domain, test, qbar): @@ -301,8 +302,7 @@ def vector_manifold_continuity_form(domain, test, q, ibp=IntegrateByParts.ONCE, return transport(form) - -def vector_invariant_form(domain, test, q, ibp=IntegrateByParts.ONCE): +def vector_invariant_form(domain, test, q, ubar): u""" The form corresponding to the vector invariant transport operator. @@ -314,6 +314,45 @@ def vector_invariant_form(domain, test, q, ibp=IntegrateByParts.ONCE): this as: (u.∇)q = (∇×q)×u + (1/2)∇(u.q) + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + test (:class:`TestFunction`): the test function. + q (:class:`ufl.Expr`): the variable to be transported. + ubar (:class:`ufl.Expr`): the transporting velocity. + + Raises: + NotImplementedError: the specified integration by parts is not 'once'. + + Returns: + class:`LabelledForm`: a labelled transport form. + """ + + if domain.mesh.topological_dimension() == 3: + L = inner(test, cross(curl(q), ubar))*dx + + else: + perp = domain.perp + L = inner(test, div(perp(q))*perp(ubar))*dx + + # Add K.E. term + L -= 0.5*div(test)*inner(q, ubar)*dx + form = transporting_velocity(L, ubar) + + return transport(form, TransportEquationType.vector_invariant) + +def upwind_vector_invariant_form(domain, test, q, ibp=IntegrateByParts.ONCE): + u""" + The form corresponding to the DG upwind vector invariant transport operator. + + The self-transporting transport operator for a vector-valued field u can be + written as circulation and kinetic energy terms: + (u.∇)u = (∇×u)×u + (1/2)∇u^2 + + When the transporting field u and transported field q are similar, we write + this as: + (u.∇)q = (∇×q)×u + (1/2)∇(u.q) + This form discretises this final equation, using an upwind discretisation when integrating by parts. diff --git a/gusto/transport_methods.py b/gusto/transport_methods.py index 1bf955479..4f1ce319e 100644 --- a/gusto/transport_methods.py +++ b/gusto/transport_methods.py @@ -5,13 +5,11 @@ from firedrake import split from gusto.configuration import IntegrateByParts, TransportEquationType from gusto.fml import Term, keep, drop -from gusto.labels import prognostic, transport +from gusto.labels import prognostic, transport, transporting_velocity from gusto.transport_forms import * __all__ = ["DGUpwind", "SUPGTransport"] -# TODO: settle on name: discretisation vs scheme - class TransportMethod(object): """ The base object for describing a transport scheme. @@ -28,10 +26,14 @@ def __init__(self, equation, variable): self.variable = variable self.domain = self.equation.domain - # TODO: how do we deal with plain transport equation? - variable_idx = equation.field_names.index(variable) - self.test = equation.tests[variable_idx] - self.field = split(equation.X)[variable_idx] + if hasattr(equation, "field_names"): + # Equation with multiple prognostic variables + variable_idx = equation.field_names.index(variable) + self.test = equation.tests[variable_idx] + self.field = split(equation.X)[variable_idx] + else: + self.field = equation.X + self.test = equation.test # Find the original transport term to be used, which we use to extract # information about the transport equation type @@ -56,10 +58,26 @@ def replace_transport_form(self, equation): this discretisation. """ - # Add the form to the equation + # We need to take care to replace the term with all the same labels, + # except the label for the transporting velocity + # This is easiest to do by extracting the transport term itself + original_form = equation.residual.label_map( + lambda t: t.has_label(transport) and t.get(prognostic) == self.variable, + map_if_true=keep, map_if_false=drop + ) + original_term = original_form.terms[0] + + # Update transporting velocity + new_transporting_velocity = self.form.terms[0].get(transporting_velocity) + original_term = transporting_velocity.update_value(original_term, new_transporting_velocity) + + # Create new term + new_term = Term(self.form.form, original_term.labels) + + # Replace original term with new term equation.residual = equation.residual.label_map( lambda t: t.has_label(transport) and t.get(prognostic) == self.variable, - map_if_true=lambda t: Term(self.form.form, t.labels)) + map_if_true=lambda t: new_term) class DGUpwind(TransportMethod): @@ -114,7 +132,7 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, elif self.transport_equation_type == TransportEquationType.vector_invariant: if outflow: raise NotImplementedError('Outflow not implemented for upwind vector invariant') - form = vector_invariant_form(self.domain, self.test, self.field, ibp=ibp) + form = upwind_vector_invariant_form(self.domain, self.test, self.field, ibp=ibp) else: raise NotImplementedError('Upwind transport scheme has not been ' diff --git a/integration-tests/transport/test_dg_transport.py b/integration-tests/transport/test_dg_transport.py index 9b4abe75a..3b499e56f 100644 --- a/integration-tests/transport/test_dg_transport.py +++ b/integration-tests/transport/test_dg_transport.py @@ -26,8 +26,9 @@ def test_dg_transport_scalar(tmpdir, geometry, equation_form, tracer_setup): eqn = ContinuityEquation(domain, V, "f") transport_scheme = SSPRK3(domain) + transport_method = DGUpwind(eqn, "f") - timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) + timestepper = PrescribedTransport(eqn, transport_scheme, transport_method, setup.io) # Initial conditions timestepper.fields("f").interpolate(setup.f_init) @@ -52,8 +53,9 @@ def test_dg_transport_vector(tmpdir, geometry, equation_form, tracer_setup): eqn = ContinuityEquation(domain, V, "f") transport_scheme = SSPRK3(domain) + transport_method = DGUpwind(eqn, "f") - timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) + timestepper = PrescribedTransport(eqn, transport_scheme, transport_method, setup.io) # Initial conditions timestepper.fields("f").interpolate(f_init) From c9fd05390ee8f9c97923e8a4cc71f1ca146f2c70 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sat, 8 Jul 2023 08:21:21 +0100 Subject: [PATCH 10/20] get transport-only tests working properly. Still got to fix SUPG --- gusto/equations.py | 8 ++++---- gusto/timeloop.py | 1 - gusto/transport_forms.py | 2 +- integration-tests/transport/test_embedded_dg_advection.py | 8 +++++--- integration-tests/transport/test_limiters.py | 4 +++- integration-tests/transport/test_recovered_transport.py | 3 ++- integration-tests/transport/test_subcycling.py | 3 ++- integration-tests/transport/test_supg_transport.py | 7 +++++-- .../transport/test_vector_recovered_space.py | 3 ++- 9 files changed, 24 insertions(+), 15 deletions(-) diff --git a/gusto/equations.py b/gusto/equations.py index c53d4a1b7..6568639c9 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -94,7 +94,7 @@ def __init__(self, domain, function_space, field_name, Vu=None): u = self.prescribed_fields("u", V) test = self.test - q = Function(function_space) + q = self.X mass_form = time_derivative(inner(q, test)*dx) transport_form = advection_form(test, q, u) @@ -125,7 +125,7 @@ def __init__(self, domain, function_space, field_name, Vu=None): u = self.prescribed_fields("u", V) test = self.test - q = Function(function_space) + q = self.X mass_form = time_derivative(inner(q, test)*dx) transport_form = continuity_form(test, q, u) @@ -150,7 +150,7 @@ def __init__(self, domain, function_space, field_name, super().__init__(domain, function_space, field_name) test = self.test - q = Function(function_space) + q = self.X mass_form = time_derivative(inner(q, test)*dx) diffusion_form = interior_penalty_diffusion_form(domain, test, q, diffusion_parameters) @@ -185,7 +185,7 @@ def __init__(self, domain, function_space, field_name, Vu=None, u = self.prescribed_fields("u", V) test = self.test - q = Function(function_space) + q = self.X mass_form = time_derivative(inner(q, test)*dx) transport_form = advection_form(test, q, u) diffusion_form = interior_penalty_diffusion_form(domain, test, q, diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 0d716ccdf..29a989817 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -251,7 +251,6 @@ def timestep(self): x_in = [x(name) for x in self.x.previous[-self.scheme.nlevels:]] self.scheme.apply(xnp1(name), *x_in) - import pdb; pdb.set_trace() class SplitPhysicsTimestepper(Timestepper): diff --git a/gusto/transport_forms.py b/gusto/transport_forms.py index a8d653b9c..05b71c180 100644 --- a/gusto/transport_forms.py +++ b/gusto/transport_forms.py @@ -85,7 +85,7 @@ def advection_form(test, q, ubar): class:`LabelledForm`: a labelled transport form. """ - L = inner(test, inner(ubar, grad(q)))*dx + L = inner(test, dot(ubar, grad(q)))*dx form = transporting_velocity(L, ubar) return transport(form, TransportEquationType.advective) diff --git a/integration-tests/transport/test_embedded_dg_advection.py b/integration-tests/transport/test_embedded_dg_advection.py index d760a82f6..0cc023385 100644 --- a/integration-tests/transport/test_embedded_dg_advection.py +++ b/integration-tests/transport/test_embedded_dg_advection.py @@ -28,12 +28,14 @@ def test_embedded_dg_advection_scalar(tmpdir, ibp, equation_form, space, opts = EmbeddedDGOptions(embedding_space=domain.spaces("DG")) if equation_form == "advective": - eqn = AdvectionEquation(domain, V, "f", ibp=ibp) + eqn = AdvectionEquation(domain, V, "f") else: - eqn = ContinuityEquation(domain, V, "f", ibp=ibp) + eqn = ContinuityEquation(domain, V, "f") transport_schemes = SSPRK3(domain, options=opts) - timestepper = PrescribedTransport(eqn, transport_schemes, setup.io) + transport_method = DGUpwind(eqn, "f", ibp=ibp) + + timestepper = PrescribedTransport(eqn, transport_schemes, transport_method, setup.io) # Initial conditions timestepper.fields("f").interpolate(setup.f_init) diff --git a/integration-tests/transport/test_limiters.py b/integration-tests/transport/test_limiters.py index e98a0842e..b0ce302a5 100644 --- a/integration-tests/transport/test_limiters.py +++ b/integration-tests/transport/test_limiters.py @@ -87,8 +87,10 @@ def setup_limiters(dirname, space): else: raise NotImplementedError + transport_method = DGUpwind(eqn, "tracer") + # Build time stepper - stepper = PrescribedTransport(eqn, transport_schemes, io) + stepper = PrescribedTransport(eqn, transport_schemes, transport_method, io) # ------------------------------------------------------------------------ # # Initial condition diff --git a/integration-tests/transport/test_recovered_transport.py b/integration-tests/transport/test_recovered_transport.py index 87160db95..862714c3f 100644 --- a/integration-tests/transport/test_recovered_transport.py +++ b/integration-tests/transport/test_recovered_transport.py @@ -36,8 +36,9 @@ def test_recovered_space_setup(tmpdir, geometry, tracer_setup): boundary_method=BoundaryMethod.taylor) transport_scheme = SSPRK3(domain, options=recovery_opts) + transport_method = DGUpwind(eqn, "f") - timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) + timestepper = PrescribedTransport(eqn, transport_scheme, transport_method, setup.io) # Initialise fields timestepper.fields("f").interpolate(setup.f_init) diff --git a/integration-tests/transport/test_subcycling.py b/integration-tests/transport/test_subcycling.py index 09915ba18..97a5a31db 100644 --- a/integration-tests/transport/test_subcycling.py +++ b/integration-tests/transport/test_subcycling.py @@ -25,8 +25,9 @@ def test_subcyling(tmpdir, equation_form, tracer_setup): eqn = ContinuityEquation(domain, V, "f") transport_scheme = SSPRK3(domain, subcycles=2) + transport_method = DGUpwind(eqn, "f") - timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) + timestepper = PrescribedTransport(eqn, transport_scheme, transport_method, setup.io) # Initial conditions timestepper.fields("f").interpolate(setup.f_init) diff --git a/integration-tests/transport/test_supg_transport.py b/integration-tests/transport/test_supg_transport.py index 9475234f0..b3c60035a 100644 --- a/integration-tests/transport/test_supg_transport.py +++ b/integration-tests/transport/test_supg_transport.py @@ -41,7 +41,8 @@ def test_supg_transport_scalar(tmpdir, equation_form, scheme, space, elif scheme == "implicit_midpoint": transport_scheme = ImplicitMidpoint(domain, options=opts) - timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) + transport_method = SUPGTransport() + timestepper = PrescribedTransport(eqn, transport_scheme, transport_method, setup.io) # Initial conditions timestepper.fields("f").interpolate(setup.f_init) @@ -82,7 +83,9 @@ def test_supg_transport_vector(tmpdir, equation_form, scheme, space, elif scheme == "implicit_midpoint": transport_scheme = ImplicitMidpoint(domain, options=opts) - timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) + transport_method = SUPGTransport() + + timestepper = PrescribedTransport(eqn, transport_scheme, transport_method, setup.io) # Initial conditions f = timestepper.fields("f") diff --git a/integration-tests/transport/test_vector_recovered_space.py b/integration-tests/transport/test_vector_recovered_space.py index 4b4c15221..1c67bc3a9 100644 --- a/integration-tests/transport/test_vector_recovered_space.py +++ b/integration-tests/transport/test_vector_recovered_space.py @@ -41,7 +41,8 @@ def test_vector_recovered_space_setup(tmpdir, geometry, tracer_setup): # Make equation eqn = AdvectionEquation(domain, Vu, "f") transport_scheme = SSPRK3(domain, options=rec_opts) - timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) + transport_method = DGUpwind(eqn, "f") + timestepper = PrescribedTransport(eqn, transport_scheme, transport_method, setup.io) # Initialise fields f_init = as_vector([setup.f_init]*gdim) From c4c767df7961e37edc7d686bc6b23050471fbfcb Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sat, 8 Jul 2023 10:07:31 +0100 Subject: [PATCH 11/20] implement SUPG but there is still a bug --- gusto/time_discretisation.py | 13 ++--- gusto/transport_methods.py | 6 +- gusto/wrappers.py | 55 ++++++++++++++++++- .../transport/test_supg_transport.py | 5 +- 4 files changed, 62 insertions(+), 17 deletions(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index a0638eedf..ba7b9b6ef 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -12,13 +12,10 @@ from firedrake.formmanipulation import split_form from firedrake.utils import cached_property import ufl -from gusto.configuration import (logger, DEBUG, TransportEquationType, - EmbeddedDGOptions, RecoveryOptions) -from gusto.labels import (time_derivative, transporting_velocity, prognostic, - subject, physics, transport, ibp_label, +from gusto.configuration import (logger, DEBUG, EmbeddedDGOptions, RecoveryOptions) +from gusto.labels import (time_derivative, prognostic, physics, replace_subject, replace_test_function) from gusto.fml.form_manipulation_labelling import Term, all_terms, drop -from gusto.transport_forms import advection_form, continuity_form from gusto.wrappers import * @@ -151,15 +148,17 @@ def setup(self, equation, apply_bcs=True, *active_labels): if self.solver_parameters is None: self.solver_parameters = self.wrapper.solver_parameters new_test = TestFunction(self.wrapper.test_space) - # TODO: this needs moving if SUPG becomes a transport scheme + # SUPG has a special wrapper if self.wrapper_name == "supg": - new_test = new_test + dot(dot(uadv, self.wrapper.tau), grad(new_test)) + new_test = self.wrapper.test # Replace the original test function with the one from the wrapper self.residual = self.residual.label_map( all_terms, map_if_true=replace_test_function(new_test)) + self.residual = self.wrapper.label_terms(self.residual) + # -------------------------------------------------------------------- # # Make boundary conditions # -------------------------------------------------------------------- # diff --git a/gusto/transport_methods.py b/gusto/transport_methods.py index 4f1ce319e..34e7ffdce 100644 --- a/gusto/transport_methods.py +++ b/gusto/transport_methods.py @@ -8,7 +8,7 @@ from gusto.labels import prognostic, transport, transporting_velocity from gusto.transport_forms import * -__all__ = ["DGUpwind", "SUPGTransport"] +__all__ = ["DGUpwind"] class TransportMethod(object): """ @@ -139,7 +139,3 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, + 'implemented for this transport equation type') self.form = form - - -class SUPGTransport(TransportMethod): - pass \ No newline at end of file diff --git a/gusto/wrappers.py b/gusto/wrappers.py index af131ae60..f867c816c 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -6,9 +6,13 @@ from abc import ABCMeta, abstractmethod from firedrake import (FunctionSpace, Function, BrokenElement, Projector, - Interpolator, VectorElement, Constant, as_ufl) + Interpolator, VectorElement, Constant, as_ufl, dot, grad, + TestFunction) from gusto.configuration import EmbeddedDGOptions, RecoveryOptions, SUPGOptions +from gusto.fml import Term from gusto.recovery import Recoverer, ReversibleRecoverer +from gusto.labels import transporting_velocity +import ufl __all__ = ["EmbeddedDGWrapper", "RecoveryWrapper", "SUPGWrapper"] @@ -47,6 +51,21 @@ def post_apply(self): """Generic steps to be done after time discretisation apply method.""" pass + def label_terms(self, residual): + """ + A base method to allow labels to be updated or extra labels added to + the form that will be used with the wrapper. This base method does + nothing but there may be implementations in child classes. + + Args: + residual (:class:`LabelledForm`): the labelled form to update. + + Returns: + :class:`LabelledForm`: the updated labelled form. + """ + + return residual + class EmbeddedDGWrapper(Wrapper): """ @@ -240,7 +259,7 @@ def is_cg(V): class SUPGWrapper(Wrapper): """ - Wrapper for computing a time discretisation with SUPG, which adjust the + Wrapper for computing a time discretisation with SUPG, which adjusts the test function space that is used to solve the problem. """ @@ -288,6 +307,15 @@ def setup(self): 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'} + # -------------------------------------------------------------------- # + # Set up test function + # -------------------------------------------------------------------- # + + test = TestFunction(self.test_space) + uadv = Function(domain.spaces('HDiv')) + self.test = test + dot(dot(uadv, self.tau), grad(test)) + self.transporting_velocity = uadv + def pre_apply(self, x_in): """ Does nothing for SUPG, just sets the input field. @@ -307,3 +335,26 @@ def post_apply(self, x_out): """ x_out.assign(self.x_out) + + def label_terms(self, residual): + """ + A base method to allow labels to be updated or extra labels added to + the form that will be used with the wrapper. + + Args: + residual (:class:`LabelledForm`): the labelled form to update. + + Returns: + :class:`LabelledForm`: the updated labelled form. + """ + + new_residual = residual.label_map( + lambda t: t.has_label(transporting_velocity), + # Update and replace transporting velocity in any terms + map_if_true=lambda t: + Term(ufl.replace(t.form, {t.get(transporting_velocity): self.transporting_velocity}), t.labels), + # Add new label to other terms + map_if_false=lambda t: transporting_velocity(t, self.transporting_velocity) + ) + + return new_residual diff --git a/integration-tests/transport/test_supg_transport.py b/integration-tests/transport/test_supg_transport.py index b3c60035a..84642d3ec 100644 --- a/integration-tests/transport/test_supg_transport.py +++ b/integration-tests/transport/test_supg_transport.py @@ -41,7 +41,7 @@ def test_supg_transport_scalar(tmpdir, equation_form, scheme, space, elif scheme == "implicit_midpoint": transport_scheme = ImplicitMidpoint(domain, options=opts) - transport_method = SUPGTransport() + transport_method = DGUpwind(eqn, "f", ibp=ibp) timestepper = PrescribedTransport(eqn, transport_scheme, transport_method, setup.io) # Initial conditions @@ -83,8 +83,7 @@ def test_supg_transport_vector(tmpdir, equation_form, scheme, space, elif scheme == "implicit_midpoint": transport_scheme = ImplicitMidpoint(domain, options=opts) - transport_method = SUPGTransport() - + transport_method = DGUpwind(eqn, "f", ibp=ibp) timestepper = PrescribedTransport(eqn, transport_scheme, transport_method, setup.io) # Initial conditions From 3948d468226677b4fddf8d621fdaac13b93b3860 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Sat, 8 Jul 2023 19:15:19 +0100 Subject: [PATCH 12/20] make spatial methods general, and roll out changes to integration tests (not working) --- gusto/__init__.py | 5 +- gusto/common_forms.py | 168 ++++++ gusto/{diffusion.py => diffusion_methods.py} | 44 +- gusto/equations.py | 64 +-- gusto/physics.py | 2 +- gusto/spatial_methods.py | 64 +++ gusto/timeloop.py | 129 +++-- gusto/transport_forms.py | 493 ------------------ gusto/transport_methods.py | 228 +++++++- gusto/wrappers.py | 2 + .../balance/test_compressible_balance.py | 8 +- .../balance/test_saturated_balance.py | 6 + .../balance/test_unsaturated_balance.py | 7 + integration-tests/diffusion/test_diffusion.py | 6 +- .../equations/test_advection_diffusion.py | 4 +- .../equations/test_dry_compressible.py | 4 + .../equations/test_forced_advection.py | 3 +- .../equations/test_incompressible.py | 3 + .../equations/test_moist_compressible.py | 4 + integration-tests/equations/test_sw_fplane.py | 9 +- .../equations/test_sw_linear_triangle.py | 3 +- .../equations/test_sw_triangle.py | 16 +- .../equations/test_thermal_sw.py | 6 +- integration-tests/model/test_checkpointing.py | 14 +- integration-tests/model/test_nc_outputting.py | 3 +- .../model/test_passive_tracer.py | 4 +- .../model/test_prescribed_transport.py | 5 +- .../model/test_time_discretisation.py | 4 +- 28 files changed, 643 insertions(+), 665 deletions(-) create mode 100644 gusto/common_forms.py rename gusto/{diffusion.py => diffusion_methods.py} (50%) create mode 100644 gusto/spatial_methods.py delete mode 100644 gusto/transport_forms.py diff --git a/gusto/__init__.py b/gusto/__init__.py index 99f5a6be3..c932d6811 100644 --- a/gusto/__init__.py +++ b/gusto/__init__.py @@ -1,8 +1,9 @@ from gusto.active_tracers import * # noqa +from gusto.common_forms import * # noqa from gusto.configuration import * # noqa from gusto.domain import * # noqa from gusto.diagnostics import * # noqa -from gusto.diffusion import * # noqa +from gusto.diffusion_methods import * # noqa from gusto.equations import * # noqa from gusto.fml import * # noqa from gusto.forcing import * # noqa @@ -15,8 +16,8 @@ from gusto.physics import * # noqa from gusto.preconditioners import * # noqa from gusto.recovery import * # noqa +from gusto.spatial_methods import * # noqa from gusto.time_discretisation import * # noqa from gusto.timeloop import * # noqa -from gusto.transport_forms import * # noqa from gusto.transport_methods import * # noqa from gusto.wrappers import * # noqa diff --git a/gusto/common_forms.py b/gusto/common_forms.py new file mode 100644 index 000000000..8089b9ad7 --- /dev/null +++ b/gusto/common_forms.py @@ -0,0 +1,168 @@ +""" +Provides some basic forms for discretising various common terms in equations for +geophysical fluid dynamics.""" + +from firedrake import Function, dx, dot, grad, div, inner, outer, cross, curl +from gusto.configuration import IntegrateByParts, TransportEquationType +from gusto.labels import transport, transporting_velocity, diffusion + +__all__ = ["advection_form", "continuity_form", "vector_invariant_form", + "kinetic_energy_form", "advection_equation_circulation_form", + "diffusion_form"] + +def advection_form(test, q, ubar): + u""" + The form corresponding to the advective transport operator. + + This describes (u.∇)q, for transporting velocity u and transported q. + + Args: + test (:class:`TestFunction`): the test function. + q (:class:`ufl.Expr`): the variable to be transported. + ubar (:class:`ufl.Expr`): the transporting velocity. + + Returns: + class:`LabelledForm`: a labelled transport form. + """ + + L = inner(test, dot(ubar, grad(q)))*dx + form = transporting_velocity(L, ubar) + + return transport(form, TransportEquationType.advective) + +def continuity_form(test, q, ubar): + u""" + The form corresponding to the continuity transport operator. + + This describes ∇.(u*q), for transporting velocity u and transported q. + + Args: + test (:class:`TestFunction`): the test function. + q (:class:`ufl.Expr`): the variable to be transported. + ubar (:class:`ufl.Expr`): the transporting velocity. + + Returns: + class:`LabelledForm`: a labelled transport form. + """ + + L = inner(test, div(outer(q, ubar)))*dx + form = transporting_velocity(L, ubar) + + return transport(form, TransportEquationType.conservative) + +def vector_invariant_form(domain, test, q, ubar): + u""" + The form corresponding to the vector invariant transport operator. + + The self-transporting transport operator for a vector-valued field u can be + written as circulation and kinetic energy terms: + (u.∇)u = (∇×u)×u + (1/2)∇u^2 + + When the transporting field u and transported field q are similar, we write + this as: + (u.∇)q = (∇×q)×u + (1/2)∇(u.q) + + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + test (:class:`TestFunction`): the test function. + q (:class:`ufl.Expr`): the variable to be transported. + ubar (:class:`ufl.Expr`): the transporting velocity. + + Raises: + NotImplementedError: the specified integration by parts is not 'once'. + + Returns: + class:`LabelledForm`: a labelled transport form. + """ + + if domain.mesh.topological_dimension() == 3: + L = inner(test, cross(curl(q), ubar))*dx + + else: + perp = domain.perp + L = inner(test, div(perp(q))*perp(ubar))*dx + + # Add K.E. term + L -= 0.5*div(test)*inner(q, ubar)*dx + form = transporting_velocity(L, ubar) + + return transport(form, TransportEquationType.vector_invariant) + +def kinetic_energy_form(domain, test, q): + u""" + The form corresponding to the kinetic energy term. + + Writing the kinetic energy term as (1/2)∇u^2, if the transported variable + q is similar to the transporting variable u then this can be written as: + (1/2)∇(u.q). + + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + test (:class:`TestFunction`): the test function. + q (:class:`ufl.Expr`): the variable to be transported. + + Returns: + class:`LabelledForm`: a labelled transport form. + """ + + ubar = Function(domain.spaces("HDiv")) + L = 0.5*div(test)*inner(q, ubar)*dx + + form = transporting_velocity(L, ubar) + + return transport(form, TransportEquationType.vector_invariant) + +def advection_equation_circulation_form(domain, test, q, + ibp=IntegrateByParts.ONCE): + u""" + The circulation term in the transport of a vector-valued field. + + The self-transporting transport operator for a vector-valued field u can be + written as circulation and kinetic energy terms: + (u.∇)u = (∇×u)×u + (1/2)∇u^2 + + When the transporting field u and transported field q are similar, we write + this as: + (u.∇)q = (∇×q)×u + (1/2)∇(u.q) + + The form returned by this function corresponds to the (∇×q)×u circulation + term. An an upwind discretisation is used when integrating by parts. + + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + test (:class:`TestFunction`): the test function. + q (:class:`ufl.Expr`): the variable to be transported. + ibp (:class:`IntegrateByParts`, optional): an enumerator representing + the number of times to integrate by parts. Defaults to + `IntegrateByParts.ONCE`. + + Raises: + NotImplementedError: the specified integration by parts is not 'once'. + + Returns: + class:`LabelledForm`: a labelled transport form. + """ + + form = ( + vector_invariant_form(domain, test, q, ibp=ibp) + - kinetic_energy_form(domain, test, q) + ) + + return form + +def diffusion_form(test, q, kappa): + u""" + The diffusion form, ∇.(κ∇q) for diffusivity κ and variable q. + + Args: + test (:class:`TestFunction`): the test function. + q (:class:`ufl.Expr`): the variable to be diffused. + kappa: (:class:`ufl.Expr`): the diffusivity value. + """ + + form = inner(test, div(kappa*grad(q)))*dx + + return diffusion(form) diff --git a/gusto/diffusion.py b/gusto/diffusion_methods.py similarity index 50% rename from gusto/diffusion.py rename to gusto/diffusion_methods.py index 175038e85..6ce112574 100644 --- a/gusto/diffusion.py +++ b/gusto/diffusion_methods.py @@ -1,11 +1,28 @@ -"""Provides forms for describing diffusion terms.""" +"""Provides discretisations for diffusion terms.""" -from firedrake import (inner, outer, grad, avg, dx, dS_h, dS_v, dS, - FacetNormal) +from firedrake import inner, outer, grad, avg, dx, dS_h, dS_v, dS, FacetNormal from gusto.labels import diffusion +from gusto.spatial_methods import SpatialMethod -__all__ = ["interior_penalty_diffusion_form"] +__all__ = ["InteriorPenaltyDiffusion"] + + +class DiffusionMethod(SpatialMethod): + """ + The base object for describing a spatial discretisation of diffusion terms. + """ + + def __init__(self, equation, variable): + """ + Args: + equation (:class:`PrognosticEquation`): the equation, which includes + a diffusion term. + variable (str): name of the variable to set the diffusion scheme for + """ + + # Inherited init method extracts original term to be replaced + super.__init__(self, equation, variable, diffusion) def interior_penalty_diffusion_form(domain, test, q, parameters): @@ -57,3 +74,22 @@ def get_flux_form(dS, M): form += get_flux_form(dS_, kappa) return diffusion(form) + +class InteriorPenaltyDiffusion(DiffusionMethod): + """The interior penalty method for discretising the diffusion term.""" + + def __init__(self, equation, variable, diffusion_parameters): + """ + Args: + equation (:class:`PrognosticEquation`): the equation, which includes + a transport term. + variable (str): name of the variable to set the diffusion method for + diffusion_parameters (:class:`DiffusionParameters`): object + containing metadata describing the diffusion term. Includes + the kappa and mu constants. + """ + + super.__init__(equation, variable) + + self.form = interior_penalty_diffusion_form(equation.domain, self.test, + self.field, diffusion_parameters) diff --git a/gusto/equations.py b/gusto/equations.py index 6568639c9..169a9a4e1 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -13,13 +13,9 @@ name, pressure_gradient, coriolis, perp, replace_trial_function, hydrostatic) from gusto.thermodynamics import exner_pressure -from gusto.transport_forms import (advection_form, continuity_form, - vector_invariant_form, - kinetic_energy_form, - advection_equation_circulation_form, - linear_continuity_form, - linear_advection_form) -from gusto.diffusion import interior_penalty_diffusion_form +from gusto.common_forms import (advection_form, continuity_form, + vector_invariant_form, kinetic_energy_form, + advection_equation_circulation_form) from gusto.active_tracers import ActiveTracer, Phases, TracerVariableType from gusto.configuration import TransportEquationType import ufl @@ -152,10 +148,9 @@ def __init__(self, domain, function_space, field_name, test = self.test q = self.X mass_form = time_derivative(inner(q, test)*dx) - diffusion_form = interior_penalty_diffusion_form(domain, test, q, - diffusion_parameters) + diffusive_form = diffusion_form(test, q, diffusion_parameters.kappa) - self.residual = prognostic(subject(mass_form + diffusion_form, q), field_name) + self.residual = prognostic(subject(mass_form + diffusive_form, q), field_name) class AdvectionDiffusionEquation(PrognosticEquation): @@ -188,11 +183,10 @@ def __init__(self, domain, function_space, field_name, Vu=None, q = self.X mass_form = time_derivative(inner(q, test)*dx) transport_form = advection_form(test, q, u) - diffusion_form = interior_penalty_diffusion_form(domain, test, q, - diffusion_parameters) + diffusive_form = diffusion_form(test, q, diffusion_parameters.kappa) self.residual = prognostic(subject( - mass_form + transport_form + diffusion_form, q), field_name) + mass_form + transport_form + diffusive_form, q), field_name) class PrognosticEquationSet(PrognosticEquation, metaclass=ABCMeta): @@ -599,8 +593,7 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, # transport term from depth equation. Don't include active tracers linearisation_map = lambda t: \ t.get(prognostic) in ["u", "D"] \ - and (any(t.has_label(time_derivative, pressure_gradient)) - or (t.get(prognostic) == "D" and t.has_label(transport))) + and (any(t.has_label(time_derivative, pressure_gradient, transport))) super().__init__(field_names, domain, linearisation_map=linearisation_map, no_normal_flow_bc_ids=no_normal_flow_bc_ids, @@ -608,7 +601,6 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, self.parameters = parameters g = parameters.g - H = parameters.H w, phi = self.tests[0:2] u, D = split(self.X)[0:2] @@ -642,14 +634,6 @@ 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(domain, phi, H).label_map( - lambda t: t.has_label(transporting_velocity), - lambda t: Term(ufl.replace( - t.form, {t.get(transporting_velocity): u_trial}), t.labels)) - # Add linearisation to D_adv - D_adv = linearisation(D_adv, linear_D_adv) adv_form = subject(u_adv + D_adv, self.X) @@ -661,7 +645,6 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, gamma = self.tests[2] b = split(self.X)[2] b_adv = prognostic(inner(gamma, inner(u, grad(b)))*dx, "b") - # TODO: implement correct linearisation adv_form += subject(b_adv, self.X) # -------------------------------------------------------------------- # @@ -780,8 +763,7 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, # Default linearisation is time derivatives, pressure gradient, # Coriolis and transport term from depth equation linearisation_map = lambda t: \ - (any(t.has_label(time_derivative, pressure_gradient, coriolis)) - or (t.get(prognostic) == "D" and t.has_label(transport))) + (any(t.has_label(time_derivative, pressure_gradient, coriolis, transport))) super().__init__(domain, parameters, fexpr=fexpr, bexpr=bexpr, @@ -872,8 +854,6 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, w, phi, gamma = self.tests[0:3] u, rho, theta = split(self.X)[0:3] - u_trial = split(self.trials)[0] - _, rho_bar, theta_bar = split(self.X_ref)[0:3] zero_expr = Constant(0.0)*theta exner = exner_pressure(parameters, rho, theta) n = FacetNormal(domain.mesh) @@ -906,23 +886,9 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, # Density transport (conservative form) rho_adv = prognostic(continuity_form(phi, rho, u), "rho") - # Transport term needs special linearisation - if self.linearisation_map(rho_adv.terms[0]): - linear_rho_adv = linear_continuity_form(domain, phi, rho_bar).label_map( - lambda t: t.has_label(transporting_velocity), - lambda t: Term(ufl.replace( - t.form, {t.get(transporting_velocity): u_trial}), t.labels)) - rho_adv = linearisation(rho_adv, linear_rho_adv) # Potential temperature transport (advective form) theta_adv = prognostic(advection_form(gamma, theta, u), "theta") - # Transport term needs special linearisation - if self.linearisation_map(theta_adv.terms[0]): - linear_theta_adv = linear_advection_form(domain, gamma, theta_bar).label_map( - lambda t: t.has_label(transporting_velocity), - lambda t: Term(ufl.replace( - t.form, {t.get(transporting_velocity): u_trial}), t.labels)) - theta_adv = linearisation(theta_adv, linear_theta_adv) adv_form = subject(u_adv + rho_adv + theta_adv, self.X) @@ -1018,8 +984,8 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, test = self.tests[idx] fn = split(self.X)[idx] residual += subject( - prognostic(interior_penalty_diffusion_form( - domain, test, fn, diffusion), field), self.X) + prognostic(diffusion_form(test, fn, diffusion.kappa), field), + self.X) if extra_terms is not None: for field, term in extra_terms: @@ -1215,8 +1181,6 @@ def __init__(self, domain, parameters, Omega=None, w, phi, gamma = self.tests[0:3] u, p, b = split(self.X) - u_trial = split(self.trials)[0] - b_bar = split(self.X_ref)[2] # -------------------------------------------------------------------- # # Time Derivative Terms @@ -1247,12 +1211,6 @@ def __init__(self, domain, parameters, Omega=None, # Buoyancy transport b_adv = prognostic(advection_form(gamma, b, u), "b") - if self.linearisation_map(b_adv.terms[0]): - linear_b_adv = linear_advection_form(domain, gamma, b_bar).label_map( - lambda t: t.has_label(transporting_velocity), - lambda t: Term(ufl.replace( - t.form, {t.get(transporting_velocity): u_trial}), t.labels)) - b_adv = linearisation(b_adv, linear_b_adv) adv_form = subject(u_adv + b_adv, self.X) diff --git a/gusto/physics.py b/gusto/physics.py index 7062f33e3..9f55588c3 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -11,7 +11,7 @@ from gusto.active_tracers import Phases from gusto.recovery import Recoverer, BoundaryMethod from gusto.equations import CompressibleEulerEquations -from gusto.transport_forms import advection_form +from gusto.common_forms import advection_form from gusto.fml import identity, Term from gusto.labels import subject, physics, transporting_velocity from gusto.configuration import logger diff --git a/gusto/spatial_methods.py b/gusto/spatial_methods.py new file mode 100644 index 000000000..33a0a3bef --- /dev/null +++ b/gusto/spatial_methods.py @@ -0,0 +1,64 @@ +""" +This module defines the SpatialMethod base object, which is used to define a +spatial discretisation of some term. +""" + +from firedrake import split +from gusto.fml import Term, keep, drop +from gusto.labels import prognostic + +__all__ = ['SpatialMethod'] + +class SpatialMethod(object): + """ + The base object for describing a spatial discretisation of some term. + """ + + def __init__(self, equation, variable, term_label): + """ + Args: + equation (:class:`PrognosticEquation`): the equation, which includes + the original type of this term. + variable (str): name of the variable to set the transport scheme for + term_label (:class:`Label`): the label specifying which type of term + to be discretised. + """ + self.equation = equation + self.variable = variable + self.domain = self.equation.domain + self.term_label = term_label + + if hasattr(equation, "field_names"): + # Equation with multiple prognostic variables + variable_idx = equation.field_names.index(variable) + self.test = equation.tests[variable_idx] + self.field = split(equation.X)[variable_idx] + else: + self.field = equation.X + self.test = equation.test + + # Find the original transport term to be used, which we use to extract + # information about the transport equation type + self.original_form = equation.residual.label_map( + lambda t: t.has_label(term_label) and t.get(prognostic) == variable, + map_if_true=keep, map_if_false=drop) + + num_terms = len(self.original_form.terms) + assert num_terms == 1, f'Unable to find {term_label.name} term ' \ + + f'for {variable}. {num_terms} found' + + def replace_form(self, equation): + """ + Replaces the form for the transport term in the equation with the + form for the transport discretisation. + + Args: + equation (:class:`PrognosticEquation`): the equation or scheme whose + transport term should be replaced with the transport term of + this discretisation. + """ + + # Replace original term with new term + equation.residual = equation.residual.label_map( + lambda t: t.has_label(self.term_label) and t.get(prognostic) == self.variable, + map_if_true=lambda t: Term(self.form.form, t.labels)) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 29a989817..4bcb9b64f 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -73,31 +73,43 @@ def get_initial_timesteps(self): # Return None if this is not applicable return self.scheme.initial_timesteps if can_get else None - def setup_equation_transport(self, equation): + def setup_equation(self, equation): """ - Sets up the transport methods for an equation, by the setting the - forms used for transport in the equation. + Sets up the spatial methods for an equation, by the setting the + forms used for transport/diffusion in the equation. Args: equation (:class:`PrognosticEquation`): the equation that the transport method is to be applied to. """ - # Check that we have specified all transport terms - transport_residual = equation.residual.label_map( - lambda t: t.has_label(transport), map_if_false=drop - ) - transported_variables = [t.get(prognostic) for t in transport_residual.terms] - transport_method_variables = [method.variable for method in self.transport_methods] - for variable in transported_variables: - if variable not in transport_method_variables: - logger.warning(f'Variable {variable} has a transport term but ' - + 'no transport method has been specified. ' - + 'Using default transport form.') - - # Replace transport forms in equation - if self.transport_methods is not None: - for transport_method in self.transport_methods: - transport_method.replace_transport_form(equation) + + # For now, we only have methods for transport and diffusion + for term_label in [transport, diffusion]: + # ---------------------------------------------------------------- # + # Check that appropriate methods have been provided + # ---------------------------------------------------------------- # + # Extract all terms corresponding to this type of term + residual = equation.residual.label_map( + lambda t: t.has_label(term_label), map_if_false=drop + ) + active_variables = [t.get(prognostic) for t in residual.terms] + active_methods = list(filter(lambda t: t.term_label == term_label, + self.spatial_methods)) + method_variables = [method.variable for method in active_methods] + for variable in active_variables: + if variable not in method_variables: + message = f'Variable {variable} has a {term_label.name} ' \ + + 'but no method for this has been specified. Using ' \ + + 'default form for this term' + logger.warning(message) + + # -------------------------------------------------------------------- # + # Check that appropriate methods have been provided + # -------------------------------------------------------------------- # + # Replace forms in equation + if self.spatial_methods is not None: + for method in self.spatial_methods: + method.replace_form(equation) def setup_transporting_velocity(self, scheme): """ @@ -208,23 +220,24 @@ class Timestepper(BaseTimestepper): Implements a timeloop by applying a scheme to a prognostic equation. """ - def __init__(self, equation, scheme, io, transport_methods=None): + def __init__(self, equation, scheme, io, spatial_methods=None): """ Args: equation (:class:`PrognosticEquation`): the prognostic equation scheme (:class:`TimeDiscretisation`): the scheme to use to timestep the prognostic equation io (:class:`IO`): the model's object for controlling input/output. - transport_methods (iter,optional): a list of objects describing the - methods to use for discretising transport terms for each - transported variable. Defaults to None, in which case the - transport term follows the discretisation in the equation. + spatial_methods (iter,optional): a list of objects describing the + methods to use for discretising transport or diffusion terms + for each transported/diffused variable. Defaults to None, + in which case the terms follow the original discretisation in + the equation. """ self.scheme = scheme - if transport_methods is not None: - self.transport_methods = transport_methods + if spatial_methods is not None: + self.spatial_methods = spatial_methods else: - self.transport_methods = [] + self.spatial_methods = [] super().__init__(equation=equation, io=io) @@ -238,7 +251,7 @@ def setup_fields(self): *self.io.output.dumplist) def setup_scheme(self): - self.setup_equation_transport(self.equation) + self.setup_equation(self.equation) self.scheme.setup(self.equation) self.setup_transporting_velocity(self.scheme) @@ -260,7 +273,7 @@ class SplitPhysicsTimestepper(Timestepper): scheme to be applied to the physics terms than the prognostic equation. """ - def __init__(self, equation, scheme, io, transport_methods=None, + def __init__(self, equation, scheme, io, spatial_methods=None, physics_schemes=None): """ Args: @@ -268,10 +281,11 @@ def __init__(self, equation, scheme, io, transport_methods=None, scheme (:class:`TimeDiscretisation`): the scheme to use to timestep the prognostic equation io (:class:`IO`): the model's object for controlling input/output. - transport_methods (iter,optional): a list of objects describing the - methods to use for discretising transport terms for each - transported variable. Defaults to None, in which case the - transport term follows the discretisation in the equation. + spatial_methods (iter,optional): a list of objects describing the + methods to use for discretising transport or diffusion terms + for each transported/diffused variable. Defaults to None, + in which case the terms follow the original discretisation in + the equation. physics_schemes: (list, optional): a list of :class:`Physics` and :class:`TimeDiscretisation` options describing physical parametrisations and timestepping schemes to use for each. @@ -279,8 +293,7 @@ def __init__(self, equation, scheme, io, transport_methods=None, None. """ - super().__init__(equation, scheme, io, - transport_methods=transport_methods) + super().__init__(equation, scheme, io, spatial_methods=spatial_methods) if physics_schemes is not None: self.physics_schemes = physics_schemes @@ -298,7 +311,7 @@ def transporting_velocity(self): return "prognostic" def setup_scheme(self): - self.setup_equation_transport(self.equation) + self.setup_equation(self.equation) # Go through and label all non-physics terms with a "dynamics" label dynamics = Label('dynamics') self.equation.label_terms(lambda t: not any(t.has_label(time_derivative, physics)), dynamics) @@ -325,12 +338,9 @@ class SemiImplicitQuasiNewton(BaseTimestepper): terms. """ - def __init__(self, equation_set, io, transport_schemes, - transport_methods, - auxiliary_equations_and_schemes=None, - linear_solver=None, - diffusion_schemes=None, - physics_schemes=None, **kwargs): + def __init__(self, equation_set, io, transport_schemes, spatial_methods, + auxiliary_equations_and_schemes=None, linear_solver=None, + diffusion_schemes=None, physics_schemes=None, **kwargs): """ Args: @@ -340,16 +350,16 @@ def __init__(self, equation_set, io, transport_schemes, transport_schemes: iterable of ``(field_name, scheme)`` pairs indicating the name of the field (str) to transport, and the :class:`TimeDiscretisation` to use - transport_methods (iter): a list of objects describing the spatial - discretisations of transport terms to be used. + spatial_methods (iter): a list of objects describing the spatial + discretisations of transport or diffusion terms to be used. auxiliary_equations_and_schemes: iterable of ``(equation, scheme)`` pairs indicating any additional equations to be solved and the scheme to use to solve them. Defaults to None. - linear_solver: a :class:`.TimesteppingSolver` object. Defaults to - None. - diffusion_schemes: optional iterable of ``(field_name, scheme)`` - pairs indicating the fields to diffuse, and the - :class:`~.Diffusion` to use. Defaults to None. + linear_solver (:class:`TimesteppingSolver`, optional): the object + to use for the linear solve. Defaults to None. + diffusion_schemes (iter, optional): iterable of pairs of the form + ``(field_name, scheme)`` indicating the fields to diffuse, and the + the :class:`~.TimeDiscretisation` to use. Defaults to None. physics_schemes: (list, optional): a list of :class:`Physics` and :class:`TimeDiscretisation` options describing physical parametrisations and timestepping schemes to use for each. @@ -366,7 +376,7 @@ def __init__(self, equation_set, io, transport_schemes, if kwargs: raise ValueError("unexpected kwargs: %s" % list(kwargs.keys())) - self.transport_methods = [] + self.spatial_methods = [] if physics_schemes is not None: self.physics_schemes = physics_schemes @@ -382,10 +392,10 @@ def __init__(self, equation_set, io, transport_schemes, assert scheme.field_name in equation_set.field_names # Check that there is a corresponding transport method method_found = False - for transport_method in transport_methods: - if scheme.field_name == transport_method.variable: + for method in spatial_methods: + if scheme.field_name == method.variable and method.term_label == transport: method_found = True - self.transport_methods.append(transport_method) + self.spatial_methods.append(method) assert method_found, f'No transport method found for variable {scheme.field_name}' self.diffusion_schemes = [] @@ -394,6 +404,13 @@ def __init__(self, equation_set, io, transport_schemes, assert scheme.nlevels == 1, "multilevel schemes not supported as part of this timestepping loop" assert scheme.field_name in equation_set.field_names self.diffusion_schemes.append((scheme.field_name, scheme)) + # Check that there is a corresponding transport method + method_found = False + for method in spatial_methods: + if scheme.field_name == method.variable and method.term_label == diffusion: + method_found = True + self.diffusion_methods.append(method) + assert method_found, f'No diffusion method found for variable {scheme.field_name}' if auxiliary_equations_and_schemes is not None: for eqn, scheme in auxiliary_equations_and_schemes: @@ -409,7 +426,7 @@ def __init__(self, equation_set, io, transport_schemes, super().__init__(equation_set, io) for aux_eqn, aux_scheme in self.auxiliary_equations_and_schemes: - self.transport_method_setup(aux_eqn) + self.setup_equation(aux_eqn) aux_scheme.setup(aux_eqn) self.setup_transporting_velocity(aux_scheme) @@ -468,7 +485,7 @@ def setup_scheme(self): # TODO: apply_bcs should be False for advection but this means # tests with KGOs fail apply_bcs = True - self.setup_equation_transport(self.equation) + self.setup_equation(self.equation) for _, scheme in self.active_transport: scheme.setup(self.equation, apply_bcs, transport) self.setup_transporting_velocity(scheme) @@ -592,7 +609,7 @@ def __init__(self, equation, scheme, transport_method, io, """ super().__init__(equation, scheme, io, - transport_methods=[transport_method]) + spatial_methods=[transport_method]) if physics_schemes is not None: self.physics_schemes = physics_schemes diff --git a/gusto/transport_forms.py b/gusto/transport_forms.py deleted file mode 100644 index 05b71c180..000000000 --- a/gusto/transport_forms.py +++ /dev/null @@ -1,493 +0,0 @@ -"""Provides forms for different transport operators.""" - -from firedrake import (Function, FacetNormal, - dx, dot, grad, div, jump, avg, dS, dS_v, dS_h, inner, - ds_v, ds_t, ds_b, - outer, sign, cross, curl) -from gusto.configuration import IntegrateByParts, TransportEquationType -from gusto.labels import transport, transporting_velocity, ibp_label - - -__all__ = ["advection_form", "continuity_form", "vector_invariant_form", - "vector_manifold_advection_form", "kinetic_energy_form", - "advection_equation_circulation_form", "linear_continuity_form", - "upwind_advection_form", "upwind_continuity_form", - "upwind_vector_invariant_form"] - - -def linear_advection_form(domain, test, qbar): - """ - The form corresponding to the linearised advective transport operator. - - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - test (:class:`TestFunction`): the test function. - qbar (:class:`ufl.Expr`): the variable to be transported. - - Returns: - :class:`LabelledForm`: a labelled transport form. - """ - - ubar = Function(domain.spaces("HDiv")) - - # TODO: why is there a k here? - L = test*dot(ubar, domain.k)*dot(domain.k, grad(qbar))*dx - - form = transporting_velocity(L, ubar) - - return transport(form, TransportEquationType.advective) - - -def linear_continuity_form(domain, test, qbar, facet_term=False): - """ - The form corresponding to the linearised continuity transport operator. - - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - test (:class:`TestFunction`): the test function. - qbar (:class:`ufl.Expr`): the variable to be transported. - facet_term (bool, optional): whether to include interior facet terms. - Defaults to False. - - Returns: - :class:`LabelledForm`: a labelled transport form. - """ - - Vu = domain.spaces("HDiv") - ubar = Function(Vu) - - L = qbar*test*div(ubar)*dx - - if facet_term: - n = FacetNormal(domain.mesh) - Vu = domain.spaces("HDiv") - dS_ = (dS_v + dS_h) if Vu.extruded else dS - L += jump(ubar*test, n)*avg(qbar)*dS_ - - form = transporting_velocity(L, ubar) - - return transport(form, TransportEquationType.conservative) - -def advection_form(test, q, ubar): - u""" - The form corresponding to the advective transport operator. - - This describes (u.∇)q, for transporting velocity u and transported q. - - Args: - test (:class:`TestFunction`): the test function. - q (:class:`ufl.Expr`): the variable to be transported. - ubar (:class:`ufl.Expr`): the transporting velocity. - - Returns: - class:`LabelledForm`: a labelled transport form. - """ - - L = inner(test, dot(ubar, grad(q)))*dx - form = transporting_velocity(L, ubar) - - return transport(form, TransportEquationType.advective) - -def continuity_form(test, q, ubar): - u""" - The form corresponding to the continuity transport operator. - - This describes ∇.(u*q), for transporting velocity u and transported q. - - Args: - test (:class:`TestFunction`): the test function. - q (:class:`ufl.Expr`): the variable to be transported. - ubar (:class:`ufl.Expr`): the transporting velocity. - - Returns: - class:`LabelledForm`: a labelled transport form. - """ - - L = inner(test, div(outer(q, ubar)))*dx - form = transporting_velocity(L, ubar) - - return transport(form, TransportEquationType.conservative) - -def upwind_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): - u""" - The form corresponding to the DG upwind advective transport operator. - - This discretises (u.∇)q, for transporting velocity u and transported - variable q. An upwind discretisation is used for the facet terms when the - form is integrated by parts. - - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - test (:class:`TestFunction`): the test function. - q (:class:`ufl.Expr`): the variable to be transported. - ibp (:class:`IntegrateByParts`, optional): an enumerator representing - the number of times to integrate by parts. Defaults to - `IntegrateByParts.ONCE`. - outflow (bool, optional): whether to include outflow at the domain - boundaries, through exterior facet terms. Defaults to False. - - Raises: - ValueError: Can only use outflow option when the integration by parts - option is not "never". - - Returns: - class:`LabelledForm`: a labelled transport form. - """ - - if outflow and ibp == IntegrateByParts.NEVER: - raise ValueError("outflow is True and ibp is None are incompatible options") - Vu = domain.spaces("HDiv") - dS_ = (dS_v + dS_h) if Vu.extruded else dS - ubar = Function(Vu) - - if ibp == IntegrateByParts.ONCE: - L = -inner(div(outer(test, ubar)), q)*dx - else: - L = inner(outer(test, ubar), grad(q))*dx - - if ibp != IntegrateByParts.NEVER: - n = FacetNormal(domain.mesh) - un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) - - L += dot(jump(test), (un('+')*q('+') - un('-')*q('-')))*dS_ - - if ibp == IntegrateByParts.TWICE: - L -= (inner(test('+'), dot(ubar('+'), n('+'))*q('+')) - + inner(test('-'), dot(ubar('-'), n('-'))*q('-')))*dS_ - - if outflow: - n = FacetNormal(domain.mesh) - un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) - L += test*un*q*(ds_v + ds_t + ds_b) - - form = transporting_velocity(L, ubar) - - return ibp_label(transport(form, TransportEquationType.advective), ibp) - - -def upwind_continuity_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): - u""" - The form corresponding to the DG upwind continuity transport operator. - - This discretises ∇.(u*q), for transporting velocity u and transported - variable q. An upwind discretisation is used for the facet terms when the - form is integrated by parts. - - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - test (:class:`TestFunction`): the test function. - q (:class:`ufl.Expr`): the variable to be transported. - ibp (:class:`IntegrateByParts`, optional): an enumerator representing - the number of times to integrate by parts. Defaults to - `IntegrateByParts.ONCE`. - outflow (bool, optional): whether to include outflow at the domain - boundaries, through exterior facet terms. Defaults to False. - - Raises: - ValueError: Can only use outflow option when the integration by parts - option is not "never". - - Returns: - class:`LabelledForm`: a labelled transport form. - """ - - if outflow and ibp == IntegrateByParts.NEVER: - raise ValueError("outflow is True and ibp is None are incompatible options") - Vu = domain.spaces("HDiv") - dS_ = (dS_v + dS_h) if Vu.extruded else dS - ubar = Function(Vu) - - if ibp == IntegrateByParts.ONCE: - L = -inner(grad(test), outer(q, ubar))*dx - else: - L = inner(test, div(outer(q, ubar)))*dx - - if ibp != IntegrateByParts.NEVER: - n = FacetNormal(domain.mesh) - un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) - - L += dot(jump(test), (un('+')*q('+') - un('-')*q('-')))*dS_ - - if ibp == IntegrateByParts.TWICE: - L -= (inner(test('+'), dot(ubar('+'), n('+'))*q('+')) - + inner(test('-'), dot(ubar('-'), n('-'))*q('-')))*dS_ - - if outflow: - n = FacetNormal(domain.mesh) - un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) - L += test*un*q*(ds_v + ds_t + ds_b) - - form = transporting_velocity(L, ubar) - - return ibp_label(transport(form, TransportEquationType.conservative), ibp) - - -def vector_manifold_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): - """ - Form for advective transport operator including vector manifold correction. - - This creates the form corresponding to the advective transport operator, but - also includes a correction for the treatment of facet terms when the - transported field is vector-valued and the mesh is curved. This correction - is based on that of Bernard, Remacle et al (2009). - - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - test (:class:`TestFunction`): the test function. - q (:class:`ufl.Expr`): the variable to be transported. - ibp (:class:`IntegrateByParts`, optional): an enumerator representing - the number of times to integrate by parts. Defaults to - `IntegrateByParts.ONCE`. - outflow (bool, optional): whether to include outflow at the domain - boundaries, through exterior facet terms. Defaults to False. - - Returns: - class:`LabelledForm`: a labelled transport form. - """ - - L = upwind_advection_form(domain, test, q, ibp, outflow) - - # TODO: there should maybe be a restriction on IBP here - Vu = domain.spaces("HDiv") - dS_ = (dS_v + dS_h) if Vu.extruded else dS - ubar = Function(Vu) - n = FacetNormal(domain.mesh) - un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) - L += un('+')*inner(test('-'), n('+')+n('-'))*inner(q('+'), n('+'))*dS_ - L += un('-')*inner(test('+'), n('+')+n('-'))*inner(q('-'), n('-'))*dS_ - - return L - - -def vector_manifold_continuity_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): - """ - Form for continuity transport operator including vector manifold correction. - - This creates the form corresponding to the continuity transport operator, - but also includes a correction for the treatment of facet terms when the - transported field is vector-valued and the mesh is curved. This correction - is based on that of Bernard, Remacle et al (2009). - - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - test (:class:`TestFunction`): the test function. - q (:class:`ufl.Expr`): the variable to be transported. - ibp (:class:`IntegrateByParts`, optional): an enumerator representing - the number of times to integrate by parts. Defaults to - `IntegrateByParts.ONCE`. - outflow (bool, optional): whether to include outflow at the domain - boundaries, through exterior facet terms. Defaults to False. - - Returns: - class:`LabelledForm`: a labelled transport form. - """ - - L = upwind_continuity_form(domain, test, q, ibp, outflow) - - Vu = domain.spaces("HDiv") - dS_ = (dS_v + dS_h) if Vu.extruded else dS - ubar = Function(Vu) - n = FacetNormal(domain.mesh) - un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) - L += un('+')*inner(test('-'), n('+')+n('-'))*inner(q('+'), n('+'))*dS_ - L += un('-')*inner(test('+'), n('+')+n('-'))*inner(q('-'), n('-'))*dS_ - - form = transporting_velocity(L, ubar) - - return transport(form) - -def vector_invariant_form(domain, test, q, ubar): - u""" - The form corresponding to the vector invariant transport operator. - - The self-transporting transport operator for a vector-valued field u can be - written as circulation and kinetic energy terms: - (u.∇)u = (∇×u)×u + (1/2)∇u^2 - - When the transporting field u and transported field q are similar, we write - this as: - (u.∇)q = (∇×q)×u + (1/2)∇(u.q) - - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - test (:class:`TestFunction`): the test function. - q (:class:`ufl.Expr`): the variable to be transported. - ubar (:class:`ufl.Expr`): the transporting velocity. - - Raises: - NotImplementedError: the specified integration by parts is not 'once'. - - Returns: - class:`LabelledForm`: a labelled transport form. - """ - - if domain.mesh.topological_dimension() == 3: - L = inner(test, cross(curl(q), ubar))*dx - - else: - perp = domain.perp - L = inner(test, div(perp(q))*perp(ubar))*dx - - # Add K.E. term - L -= 0.5*div(test)*inner(q, ubar)*dx - form = transporting_velocity(L, ubar) - - return transport(form, TransportEquationType.vector_invariant) - -def upwind_vector_invariant_form(domain, test, q, ibp=IntegrateByParts.ONCE): - u""" - The form corresponding to the DG upwind vector invariant transport operator. - - The self-transporting transport operator for a vector-valued field u can be - written as circulation and kinetic energy terms: - (u.∇)u = (∇×u)×u + (1/2)∇u^2 - - When the transporting field u and transported field q are similar, we write - this as: - (u.∇)q = (∇×q)×u + (1/2)∇(u.q) - - This form discretises this final equation, using an upwind discretisation - when integrating by parts. - - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - test (:class:`TestFunction`): the test function. - q (:class:`ufl.Expr`): the variable to be transported. - ibp (:class:`IntegrateByParts`, optional): an enumerator representing - the number of times to integrate by parts. Defaults to - `IntegrateByParts.ONCE`. - - Raises: - NotImplementedError: the specified integration by parts is not 'once'. - - Returns: - class:`LabelledForm`: a labelled transport form. - """ - - Vu = domain.spaces("HDiv") - dS_ = (dS_v + dS_h) if Vu.extruded else dS - ubar = Function(Vu) - n = FacetNormal(domain.mesh) - Upwind = 0.5*(sign(dot(ubar, n))+1) - - if domain.mesh.topological_dimension() == 3: - - if ibp != IntegrateByParts.ONCE: - raise NotImplementedError - - # - # = - - # = - - # <> - - both = lambda u: 2*avg(u) - - L = ( - inner(q, curl(cross(ubar, test)))*dx - - inner(both(Upwind*q), - both(cross(n, cross(ubar, test))))*dS_ - ) - - else: - - perp = domain.perp - if domain.on_sphere: - outward_normals = domain.outward_normals - perp_u_upwind = lambda q: Upwind('+')*cross(outward_normals('+'), q('+')) + Upwind('-')*cross(outward_normals('-'), q('-')) - else: - perp_u_upwind = lambda q: Upwind('+')*perp(q('+')) + Upwind('-')*perp(q('-')) - - if ibp == IntegrateByParts.ONCE: - L = ( - -inner(perp(grad(inner(test, perp(ubar)))), q)*dx - - inner(jump(inner(test, perp(ubar)), n), - perp_u_upwind(q))*dS_ - ) - else: - L = ( - (-inner(test, div(perp(q))*perp(ubar)))*dx - - inner(jump(inner(test, perp(ubar)), n), - perp_u_upwind(q))*dS_ - + jump(inner(test, - perp(ubar))*perp(q), n)*dS_ - ) - - L -= 0.5*div(test)*inner(q, ubar)*dx - - form = transporting_velocity(L, ubar) - - return transport(form, TransportEquationType.vector_invariant) - - -def kinetic_energy_form(domain, test, q): - u""" - The form corresponding to the kinetic energy term. - - Writing the kinetic energy term as (1/2)∇u^2, if the transported variable - q is similar to the transporting variable u then this can be written as: - (1/2)∇(u.q). - - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - test (:class:`TestFunction`): the test function. - q (:class:`ufl.Expr`): the variable to be transported. - - Returns: - class:`LabelledForm`: a labelled transport form. - """ - - ubar = Function(domain.spaces("HDiv")) - L = 0.5*div(test)*inner(q, ubar)*dx - - form = transporting_velocity(L, ubar) - - return transport(form, TransportEquationType.vector_invariant) - - -def advection_equation_circulation_form(domain, test, q, - ibp=IntegrateByParts.ONCE): - u""" - The circulation term in the transport of a vector-valued field. - - The self-transporting transport operator for a vector-valued field u can be - written as circulation and kinetic energy terms: - (u.∇)u = (∇×u)×u + (1/2)∇u^2 - - When the transporting field u and transported field q are similar, we write - this as: - (u.∇)q = (∇×q)×u + (1/2)∇(u.q) - - The form returned by this function corresponds to the (∇×q)×u circulation - term. An an upwind discretisation is used when integrating by parts. - - Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. - test (:class:`TestFunction`): the test function. - q (:class:`ufl.Expr`): the variable to be transported. - ibp (:class:`IntegrateByParts`, optional): an enumerator representing - the number of times to integrate by parts. Defaults to - `IntegrateByParts.ONCE`. - - Raises: - NotImplementedError: the specified integration by parts is not 'once'. - - Returns: - class:`LabelledForm`: a labelled transport form. - """ - - form = ( - vector_invariant_form(domain, test, q, ibp=ibp) - - kinetic_energy_form(domain, test, q) - ) - - return form diff --git a/gusto/transport_methods.py b/gusto/transport_methods.py index 34e7ffdce..6b8970dd1 100644 --- a/gusto/transport_methods.py +++ b/gusto/transport_methods.py @@ -2,15 +2,16 @@ Defines TransportMethod objects, which are used to solve a transport problem. """ -from firedrake import split +from firedrake import (dx, dS, dS_v, dS_h, ds_t, ds_b, dot, inner, outer, jump, + grad, div, FacetNormal, Function) from gusto.configuration import IntegrateByParts, TransportEquationType from gusto.fml import Term, keep, drop -from gusto.labels import prognostic, transport, transporting_velocity -from gusto.transport_forms import * +from gusto.labels import prognostic, transport, transporting_velocity, ibp_label +from gusto.spatial_methods import SpatialMethod __all__ = ["DGUpwind"] -class TransportMethod(object): +class TransportMethod(SpatialMethod): """ The base object for describing a transport scheme. """ @@ -22,32 +23,13 @@ def __init__(self, equation, variable): a transport term. variable (str): name of the variable to set the transport scheme for """ - self.equation = equation - self.variable = variable - self.domain = self.equation.domain - - if hasattr(equation, "field_names"): - # Equation with multiple prognostic variables - variable_idx = equation.field_names.index(variable) - self.test = equation.tests[variable_idx] - self.field = split(equation.X)[variable_idx] - else: - self.field = equation.X - self.test = equation.test - - # Find the original transport term to be used, which we use to extract - # information about the transport equation type - original_form = equation.residual.label_map( - lambda t: t.has_label(transport) and t.get(prognostic) == variable, - map_if_true=keep, map_if_false=drop) - num_terms = len(original_form.terms) - assert num_terms == 1, \ - f'Unable to find transport term for {variable}. {num_terms} found' + # Inherited init method extracts original term to be replaced + super().__init__(equation, variable, transport) - self.transport_equation_type = original_form.terms[0].get(transport) + self.transport_equation_type = self.original_form.terms[0].get(transport) - def replace_transport_form(self, equation): + def replace_form(self, equation): """ Replaces the form for the transport term in the equation with the form for the transport discretisation. @@ -80,6 +62,198 @@ def replace_transport_form(self, equation): map_if_true=lambda t: new_term) +def upwind_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): + u""" + The form corresponding to the DG upwind advective transport operator. + + This discretises (u.∇)q, for transporting velocity u and transported + variable q. An upwind discretisation is used for the facet terms when the + form is integrated by parts. + + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + test (:class:`TestFunction`): the test function. + q (:class:`ufl.Expr`): the variable to be transported. + ibp (:class:`IntegrateByParts`, optional): an enumerator representing + the number of times to integrate by parts. Defaults to + `IntegrateByParts.ONCE`. + outflow (bool, optional): whether to include outflow at the domain + boundaries, through exterior facet terms. Defaults to False. + + Raises: + ValueError: Can only use outflow option when the integration by parts + option is not "never". + + Returns: + class:`LabelledForm`: a labelled transport form. + """ + + if outflow and ibp == IntegrateByParts.NEVER: + raise ValueError("outflow is True and ibp is None are incompatible options") + Vu = domain.spaces("HDiv") + dS_ = (dS_v + dS_h) if Vu.extruded else dS + ubar = Function(Vu) + + if ibp == IntegrateByParts.ONCE: + L = -inner(div(outer(test, ubar)), q)*dx + else: + L = inner(outer(test, ubar), grad(q))*dx + + if ibp != IntegrateByParts.NEVER: + n = FacetNormal(domain.mesh) + un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) + + L += dot(jump(test), (un('+')*q('+') - un('-')*q('-')))*dS_ + + if ibp == IntegrateByParts.TWICE: + L -= (inner(test('+'), dot(ubar('+'), n('+'))*q('+')) + + inner(test('-'), dot(ubar('-'), n('-'))*q('-')))*dS_ + + if outflow: + n = FacetNormal(domain.mesh) + un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) + L += test*un*q*(ds_v + ds_t + ds_b) + + form = transporting_velocity(L, ubar) + + return ibp_label(transport(form, TransportEquationType.advective), ibp) + + +def upwind_continuity_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): + u""" + The form corresponding to the DG upwind continuity transport operator. + + This discretises ∇.(u*q), for transporting velocity u and transported + variable q. An upwind discretisation is used for the facet terms when the + form is integrated by parts. + + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + test (:class:`TestFunction`): the test function. + q (:class:`ufl.Expr`): the variable to be transported. + ibp (:class:`IntegrateByParts`, optional): an enumerator representing + the number of times to integrate by parts. Defaults to + `IntegrateByParts.ONCE`. + outflow (bool, optional): whether to include outflow at the domain + boundaries, through exterior facet terms. Defaults to False. + + Raises: + ValueError: Can only use outflow option when the integration by parts + option is not "never". + + Returns: + class:`LabelledForm`: a labelled transport form. + """ + + if outflow and ibp == IntegrateByParts.NEVER: + raise ValueError("outflow is True and ibp is None are incompatible options") + Vu = domain.spaces("HDiv") + dS_ = (dS_v + dS_h) if Vu.extruded else dS + ubar = Function(Vu) + + if ibp == IntegrateByParts.ONCE: + L = -inner(grad(test), outer(q, ubar))*dx + else: + L = inner(test, div(outer(q, ubar)))*dx + + if ibp != IntegrateByParts.NEVER: + n = FacetNormal(domain.mesh) + un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) + + L += dot(jump(test), (un('+')*q('+') - un('-')*q('-')))*dS_ + + if ibp == IntegrateByParts.TWICE: + L -= (inner(test('+'), dot(ubar('+'), n('+'))*q('+')) + + inner(test('-'), dot(ubar('-'), n('-'))*q('-')))*dS_ + + if outflow: + n = FacetNormal(domain.mesh) + un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) + L += test*un*q*(ds_v + ds_t + ds_b) + + form = transporting_velocity(L, ubar) + + return ibp_label(transport(form, TransportEquationType.conservative), ibp) + + +def vector_manifold_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): + """ + Form for advective transport operator including vector manifold correction. + + This creates the form corresponding to the advective transport operator, but + also includes a correction for the treatment of facet terms when the + transported field is vector-valued and the mesh is curved. This correction + is based on that of Bernard, Remacle et al (2009). + + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + test (:class:`TestFunction`): the test function. + q (:class:`ufl.Expr`): the variable to be transported. + ibp (:class:`IntegrateByParts`, optional): an enumerator representing + the number of times to integrate by parts. Defaults to + `IntegrateByParts.ONCE`. + outflow (bool, optional): whether to include outflow at the domain + boundaries, through exterior facet terms. Defaults to False. + + Returns: + class:`LabelledForm`: a labelled transport form. + """ + + L = upwind_advection_form(domain, test, q, ibp, outflow) + + # TODO: there should maybe be a restriction on IBP here + Vu = domain.spaces("HDiv") + dS_ = (dS_v + dS_h) if Vu.extruded else dS + ubar = Function(Vu) + n = FacetNormal(domain.mesh) + un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) + L += un('+')*inner(test('-'), n('+')+n('-'))*inner(q('+'), n('+'))*dS_ + L += un('-')*inner(test('+'), n('+')+n('-'))*inner(q('-'), n('-'))*dS_ + + return L + +def vector_manifold_continuity_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): + """ + Form for continuity transport operator including vector manifold correction. + + This creates the form corresponding to the continuity transport operator, + but also includes a correction for the treatment of facet terms when the + transported field is vector-valued and the mesh is curved. This correction + is based on that of Bernard, Remacle et al (2009). + + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + test (:class:`TestFunction`): the test function. + q (:class:`ufl.Expr`): the variable to be transported. + ibp (:class:`IntegrateByParts`, optional): an enumerator representing + the number of times to integrate by parts. Defaults to + `IntegrateByParts.ONCE`. + outflow (bool, optional): whether to include outflow at the domain + boundaries, through exterior facet terms. Defaults to False. + + Returns: + class:`LabelledForm`: a labelled transport form. + """ + + L = upwind_continuity_form(domain, test, q, ibp, outflow) + + Vu = domain.spaces("HDiv") + dS_ = (dS_v + dS_h) if Vu.extruded else dS + ubar = Function(Vu) + n = FacetNormal(domain.mesh) + un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) + L += un('+')*inner(test('-'), n('+')+n('-'))*inner(q('+'), n('+'))*dS_ + L += un('-')*inner(test('+'), n('+')+n('-'))*inner(q('-'), n('-'))*dS_ + + form = transporting_velocity(L, ubar) + + return transport(form) + + class DGUpwind(TransportMethod): """ The Discontinuous Galerkin Upwind transport scheme. diff --git a/gusto/wrappers.py b/gusto/wrappers.py index f867c816c..79b650200 100644 --- a/gusto/wrappers.py +++ b/gusto/wrappers.py @@ -357,4 +357,6 @@ def label_terms(self, residual): map_if_false=lambda t: transporting_velocity(t, self.transporting_velocity) ) + new_residual = transporting_velocity.update_value(new_residual, self.transporting_velocity) + return new_residual diff --git a/integration-tests/balance/test_compressible_balance.py b/integration-tests/balance/test_compressible_balance.py index 96b0d34c1..03ecc285b 100644 --- a/integration-tests/balance/test_compressible_balance.py +++ b/integration-tests/balance/test_compressible_balance.py @@ -42,16 +42,16 @@ def setup_balance(dirname): transported_fields = [ImplicitMidpoint(domain, "u"), SSPRK3(domain, "rho"), SSPRK3(domain, "theta", options=EmbeddedDGOptions())] - transport_discretisations = [DGUpwind(eqns, 'u'), - DGUpwind(eqns, 'rho'), - DGUpwind(eqns, 'theta')] + transport_methods = [DGUpwind(eqns, 'u'), + DGUpwind(eqns, 'rho'), + DGUpwind(eqns, 'theta')] # Set up linear solver linear_solver = CompressibleSolver(eqns) # build time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, - transport_discretisations, + transport_methods, linear_solver=linear_solver) # ------------------------------------------------------------------------ # diff --git a/integration-tests/balance/test_saturated_balance.py b/integration-tests/balance/test_saturated_balance.py index 41dce0f18..810bd4d70 100644 --- a/integration-tests/balance/test_saturated_balance.py +++ b/integration-tests/balance/test_saturated_balance.py @@ -84,6 +84,12 @@ def setup_saturated(dirname, recovered): else: transported_fields.append(ImplicitMidpoint(domain, 'u')) + transport_methods = [DGUpwind(eqns, 'u'), + DGUpwind(eqns, 'rho'), + DGUpwind(eqns, 'theta'), + DGUpwind(eqns, 'water_vapour'), + DGUpwind(eqns, 'cloud_water')] + # Linear solver linear_solver = CompressibleSolver(eqns) diff --git a/integration-tests/balance/test_unsaturated_balance.py b/integration-tests/balance/test_unsaturated_balance.py index e2377ceed..6f29f8c9d 100644 --- a/integration-tests/balance/test_unsaturated_balance.py +++ b/integration-tests/balance/test_unsaturated_balance.py @@ -79,6 +79,12 @@ def setup_unsaturated(dirname, recovered): else: transported_fields.append(ImplicitMidpoint(domain, "u")) + transport_methods = [DGUpwind(eqns, 'u'), + DGUpwind(eqns, 'rho'), + DGUpwind(eqns, 'theta'), + DGUpwind(eqns, 'water_vapour'), + DGUpwind(eqns, 'cloud_water')] + # Linear solver linear_solver = CompressibleSolver(eqns) @@ -87,6 +93,7 @@ def setup_unsaturated(dirname, recovered): # Time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + transport_methods, linear_solver=linear_solver, physics_schemes=physics_schemes) diff --git a/integration-tests/diffusion/test_diffusion.py b/integration-tests/diffusion/test_diffusion.py index 9cc43b38f..e946052cc 100644 --- a/integration-tests/diffusion/test_diffusion.py +++ b/integration-tests/diffusion/test_diffusion.py @@ -35,7 +35,8 @@ def test_scalar_diffusion(tmpdir, DG, tracer_setup): diffusion_params = DiffusionParameters(kappa=kappa, mu=mu) eqn = DiffusionEquation(domain, V, "f", diffusion_parameters=diffusion_params) diffusion_scheme = BackwardEuler(domain) - timestepper = Timestepper(eqn, diffusion_scheme, setup.io) + diffusion_methods = [InteriorPenaltyDiffusion(eqn, "f", diffusion_params)] + timestepper = Timestepper(eqn, diffusion_scheme, setup.io, spatial_methods=diffusion_methods) # Initial conditions timestepper.fields("f").interpolate(f_init) @@ -68,7 +69,8 @@ def test_vector_diffusion(tmpdir, DG, tracer_setup): diffusion_params = DiffusionParameters(kappa=kappa, mu=mu) eqn = DiffusionEquation(domain, V, "f", diffusion_parameters=diffusion_params) diffusion_scheme = BackwardEuler(domain) - timestepper = Timestepper(eqn, diffusion_scheme, setup.io) + diffusion_methods = [InteriorPenaltyDiffusion(eqn, "f", diffusion_params)] + timestepper = Timestepper(eqn, diffusion_scheme, setup.io, spatial_methods=diffusion_methods) # Initial conditions if DG: diff --git a/integration-tests/equations/test_advection_diffusion.py b/integration-tests/equations/test_advection_diffusion.py index 20e7076b5..7153e075a 100644 --- a/integration-tests/equations/test_advection_diffusion.py +++ b/integration-tests/equations/test_advection_diffusion.py @@ -28,13 +28,15 @@ def run_advection_diffusion(tmpdir): equation = AdvectionDiffusionEquation(domain, V, "f", Vu=Vu, diffusion_parameters=diffusion_params) + spatial_methods = [DGUpwind(equation, "f"), + InteriorPenaltyDiffusion(equation, "f", diffusion_params)] # I/O output = OutputParameters(dirname=str(tmpdir), dumpfreq=25) io = IO(domain, output) # Time stepper - stepper = PrescribedTransport(equation, SSPRK3(domain), io) + stepper = PrescribedTransport(equation, SSPRK3(domain), io, spatial_methods) # ------------------------------------------------------------------------ # # Initial conditions diff --git a/integration-tests/equations/test_dry_compressible.py b/integration-tests/equations/test_dry_compressible.py index 83635a920..2a4d20d56 100644 --- a/integration-tests/equations/test_dry_compressible.py +++ b/integration-tests/equations/test_dry_compressible.py @@ -42,12 +42,16 @@ def run_dry_compressible(tmpdir): transported_fields = [ImplicitMidpoint(domain, "u"), SSPRK3(domain, "rho"), SSPRK3(domain, "theta")] + transport_methods = [DGUpwind(eqn, 'u'), + DGUpwind(eqn, 'rho'), + DGUpwind(eqn, 'theta')] # Linear solver linear_solver = CompressibleSolver(eqn) # Time stepper stepper = SemiImplicitQuasiNewton(eqn, io, transported_fields, + transport_methods, linear_solver=linear_solver) # ------------------------------------------------------------------------ # diff --git a/integration-tests/equations/test_forced_advection.py b/integration-tests/equations/test_forced_advection.py index 6c15a67c3..cbc0aa295 100644 --- a/integration-tests/equations/test_forced_advection.py +++ b/integration-tests/equations/test_forced_advection.py @@ -49,6 +49,7 @@ def run_forced_advection(tmpdir): transport_eqn=TransportEquationType.no_transport) meqn = ForcedAdvectionEquation(domain, VD, field_name="water_vapour", Vu=Vu, active_tracers=[rain]) + transport_method = DGUpwind(meqn, "water_vapour") physics_schemes = [(InstantRain(meqn, msat, rain_name="rain", parameters=None), ForwardEuler(domain))] @@ -58,7 +59,7 @@ def run_forced_advection(tmpdir): io = IO(domain, output, diagnostic_fields=diagnostic_fields) # Time Stepper - stepper = PrescribedTransport(meqn, RK4(domain), io, + stepper = PrescribedTransport(meqn, RK4(domain), io, transport_method, physics_schemes=physics_schemes) # ------------------------------------------------------------------------ # diff --git a/integration-tests/equations/test_incompressible.py b/integration-tests/equations/test_incompressible.py index d904d0335..cdc932cdd 100644 --- a/integration-tests/equations/test_incompressible.py +++ b/integration-tests/equations/test_incompressible.py @@ -41,12 +41,15 @@ def run_incompressible(tmpdir): b_opts = SUPGOptions() transported_fields = [ImplicitMidpoint(domain, "u"), SSPRK3(domain, "b", options=b_opts)] + transport_methods = [DGUpwind(eqn, "u"), + DGUpwind(eqn, "b", ibp=b_opts.ibp)] # Linear solver linear_solver = IncompressibleSolver(eqn) # Time stepper stepper = SemiImplicitQuasiNewton(eqn, io, transported_fields, + transport_methods, linear_solver=linear_solver) # ------------------------------------------------------------------------ # diff --git a/integration-tests/equations/test_moist_compressible.py b/integration-tests/equations/test_moist_compressible.py index 9b5526fc0..e5613d242 100644 --- a/integration-tests/equations/test_moist_compressible.py +++ b/integration-tests/equations/test_moist_compressible.py @@ -43,12 +43,16 @@ def run_moist_compressible(tmpdir): transported_fields = [ImplicitMidpoint(domain, "u"), SSPRK3(domain, "rho"), SSPRK3(domain, "theta")] + transport_methods = [DGUpwind(eqn, "u"), + DGUpwind(eqn, "rho"), + DGUpwind(eqn, "theta")] # Linear solver linear_solver = CompressibleSolver(eqn) # Time stepper stepper = SemiImplicitQuasiNewton(eqn, io, transported_fields, + transport_methods, linear_solver=linear_solver) # ------------------------------------------------------------------------ # diff --git a/integration-tests/equations/test_sw_fplane.py b/integration-tests/equations/test_sw_fplane.py index 00923cc5c..eb9dc5c9c 100644 --- a/integration-tests/equations/test_sw_fplane.py +++ b/integration-tests/equations/test_sw_fplane.py @@ -36,12 +36,13 @@ def run_sw_fplane(tmpdir): io = IO(domain, output, diagnostic_fields=[CourantNumber()]) # Transport schemes - transported_fields = [] - transported_fields.append((ImplicitMidpoint(domain, "u"))) - transported_fields.append((SSPRK3(domain, "D"))) + transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "D")] + transport_methods = [DGUpwind(eqns, "u"), + DGUpwind(eqns, "D")] # Time stepper - stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields) + stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, transport_methods) # ------------------------------------------------------------------------ # # Initial conditions diff --git a/integration-tests/equations/test_sw_linear_triangle.py b/integration-tests/equations/test_sw_linear_triangle.py index 6afd088ca..e66b072bf 100644 --- a/integration-tests/equations/test_sw_linear_triangle.py +++ b/integration-tests/equations/test_sw_linear_triangle.py @@ -42,9 +42,10 @@ def setup_sw(dirname): # Transport schemes transport_schemes = [ForwardEuler(domain, "D")] + transport_methods = [] # Time stepper - stepper = SemiImplicitQuasiNewton(eqns, io, transport_schemes) + stepper = SemiImplicitQuasiNewton(eqns, io, transport_schemes, transport_methods) # ------------------------------------------------------------------------ # # Initial conditions diff --git a/integration-tests/equations/test_sw_triangle.py b/integration-tests/equations/test_sw_triangle.py index 272407b6f..dd637e365 100644 --- a/integration-tests/equations/test_sw_triangle.py +++ b/integration-tests/equations/test_sw_triangle.py @@ -156,12 +156,15 @@ def test_sw_setup(tmpdir, u_transport_option): domain, eqns, io = setup_sw(dirname, dt, u_transport_option) # Transport schemes - transported_fields = [] - transported_fields.append((ImplicitMidpoint(domain, "u"))) - transported_fields.append((SSPRK3(domain, "D"))) + transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "D")] + + transport_methods = [DGUpwind(eqns, 'u'), + DGUpwind(eqns, 'D')] # Time stepper - stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields) + stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + transport_methods) # Initial conditions set_up_initial_conditions(domain, eqns, stepper) @@ -181,7 +184,10 @@ def test_sw_ssprk3(tmpdir, u_transport_option): dt = 100 domain, eqns, io = setup_sw(dirname, dt, u_transport_option) - stepper = Timestepper(eqns, SSPRK3(domain), io) + transport_methods = [DGUpwind(eqns, 'u'), + DGUpwind(eqns, 'D')] + + stepper = Timestepper(eqns, SSPRK3(domain), io, spatial_methods=transport_methods) # Initial conditions set_up_initial_conditions(domain, eqns, stepper) diff --git a/integration-tests/equations/test_thermal_sw.py b/integration-tests/equations/test_thermal_sw.py index 75a4190b7..c48f97939 100644 --- a/integration-tests/equations/test_thermal_sw.py +++ b/integration-tests/equations/test_thermal_sw.py @@ -116,7 +116,11 @@ def test_sw_ssprk3(tmpdir): dt = 100 domain, eqns, io = setup_sw(dirname, dt, u_transport_option) - stepper = Timestepper(eqns, SSPRK3(domain), io) + transport_methods = [DGUpwind(eqns, 'u'), + DGUpwind(eqns, 'D'), + DGUpwind(eqns, 'b')] + + stepper = Timestepper(eqns, SSPRK3(domain), io, transport_methods) # Initial conditions set_up_initial_conditions(domain, eqns, stepper) diff --git a/integration-tests/model/test_checkpointing.py b/integration-tests/model/test_checkpointing.py index 529342f49..be7c1eda8 100644 --- a/integration-tests/model/test_checkpointing.py +++ b/integration-tests/model/test_checkpointing.py @@ -24,23 +24,27 @@ def set_up_model_objects(mesh, dt, output, stepper_type): io = IO(domain, output, diagnostic_fields=diagnostic_fields) + transport_methods = [DGUpwind(eqns, 'u'), + DGUpwind(eqns, 'rho'), + DGUpwind(eqns, 'theta')] + if stepper_type == 'semi_implicit': # Set up transport schemes - transported_fields = [] - transported_fields.append(SSPRK3(domain, "u")) - transported_fields.append(SSPRK3(domain, "rho")) - transported_fields.append(SSPRK3(domain, "theta")) + transported_fields = [SSPRK3(domain, "u"), + SSPRK3(domain, "rho"), + SSPRK3(domain, "theta")] # Set up linear solver linear_solver = CompressibleSolver(eqns) # build time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + transport_methods, linear_solver=linear_solver) elif stepper_type == 'multi_level': scheme = AdamsBashforth(domain, order=2) - stepper = Timestepper(eqns, scheme, io) + stepper = Timestepper(eqns, scheme, io, spatial_methods=transport_methods) else: raise ValueError(f'stepper_type {stepper_type} not recognised') diff --git a/integration-tests/model/test_nc_outputting.py b/integration-tests/model/test_nc_outputting.py index 0e4b76955..fe26c18d4 100644 --- a/integration-tests/model/test_nc_outputting.py +++ b/integration-tests/model/test_nc_outputting.py @@ -71,6 +71,7 @@ def test_nc_outputting(tmpdir, geometry, domain_and_mesh_details): else: eqn = AdvectionEquation(domain, V, 'f') transport_scheme = ForwardEuler(domain) + transport_method = DGUpwind(eqn, 'f') output = OutputParameters(dirname=dirname, dumpfreq=1, dump_nc=True, dumplist=['f'], log_level='INFO', checkpoint=False) @@ -89,7 +90,7 @@ def test_nc_outputting(tmpdir, geometry, domain_and_mesh_details): diagnostic_fields = [ZonalComponent('u'), MeridionalComponent('u'), RadialComponent('u')] io = IO(domain, output, diagnostic_fields=diagnostic_fields) - stepper = PrescribedTransport(eqn, transport_scheme, io) + stepper = PrescribedTransport(eqn, transport_scheme, transport_method, io) # ------------------------------------------------------------------------ # # Initialise fields diff --git a/integration-tests/model/test_passive_tracer.py b/integration-tests/model/test_passive_tracer.py index 63f1f5d5f..41bd5608e 100644 --- a/integration-tests/model/test_passive_tracer.py +++ b/integration-tests/model/test_passive_tracer.py @@ -43,9 +43,11 @@ def run_tracer(setup): # Set up tracer transport tracer_transport = [(tracer_eqn, SSPRK3(domain))] + transport_methods = [DGUpwind(eqns, "D"), DGUpwind(tracer_eqn, "tracer")] + # build time stepper stepper = SemiImplicitQuasiNewton( - eqns, io, transport_schemes, + eqns, io, transport_schemes, transport_methods, auxiliary_equations_and_schemes=tracer_transport) # ------------------------------------------------------------------------ # diff --git a/integration-tests/model/test_prescribed_transport.py b/integration-tests/model/test_prescribed_transport.py index 758a6349d..dc64acc82 100644 --- a/integration-tests/model/test_prescribed_transport.py +++ b/integration-tests/model/test_prescribed_transport.py @@ -31,9 +31,10 @@ def u_evaluation(t): sin(2*pi*t/setup.tmax)*sin(pi*z)]) transport_scheme = SSPRK3(domain) + transport_method = DGUpwind(eqn, 'f') - timestepper = PrescribedTransport(eqn, transport_scheme, setup.io, - prescribed_transporting_velocity=u_evaluation) + timestepper = PrescribedTransport(eqn, transport_scheme, transport_method, + setup.io, prescribed_transporting_velocity=u_evaluation) # Initial conditions timestepper.fields("f").interpolate(setup.f_init) diff --git a/integration-tests/model/test_time_discretisation.py b/integration-tests/model/test_time_discretisation.py index 0b6874c97..150f5c842 100644 --- a/integration-tests/model/test_time_discretisation.py +++ b/integration-tests/model/test_time_discretisation.py @@ -47,7 +47,9 @@ def test_time_discretisation(tmpdir, scheme, tracer_setup): elif scheme == "AdamsMoulton": transport_scheme = AdamsMoulton(domain, order=2) - timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) + transport_method = DGUpwind(eqn, 'f') + + timestepper = PrescribedTransport(eqn, transport_scheme, transport_method, setup.io) # Initial conditions timestepper.fields("f").interpolate(setup.f_init) From 40bde2698cd775e4bd85416e25e1d6f77fbfccc3 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Tue, 11 Jul 2023 18:55:26 +0100 Subject: [PATCH 13/20] apply new transport approach to lots of examples. Integration tests failing --- .../compressible/dcmip_3_1_meanflow_quads.py | 3 +- examples/compressible/dry_bryan_fritsch.py | 3 + examples/compressible/moist_bryan_fritsch.py | 3 + examples/compressible/mountain_hydrostatic.py | 4 + .../compressible/mountain_nonhydrostatic.py | 4 + examples/compressible/robert_bubble.py | 12 ++- .../skamarock_klemp_hydrostatic.py | 13 ++- .../compressible/skamarock_klemp_nonlinear.py | 4 + examples/compressible/straka_bubble.py | 6 ++ examples/compressible/unsaturated_bubble.py | 5 +- gusto/common_forms.py | 19 ++-- gusto/diffusion_methods.py | 4 +- gusto/equations.py | 32 ++++--- gusto/physics.py | 10 ++- gusto/spatial_methods.py | 2 +- gusto/timeloop.py | 4 +- gusto/transport_methods.py | 87 ++++++++++++++++++- .../balance/test_saturated_balance.py | 1 + .../equations/test_dry_compressible.py | 2 +- .../equations/test_incompressible.py | 2 +- .../equations/test_moist_compressible.py | 2 +- integration-tests/equations/test_sw_fplane.py | 2 +- .../physics/test_precipitation.py | 6 +- 23 files changed, 178 insertions(+), 52 deletions(-) diff --git a/examples/compressible/dcmip_3_1_meanflow_quads.py b/examples/compressible/dcmip_3_1_meanflow_quads.py index 32e168f1b..82611ef04 100644 --- a/examples/compressible/dcmip_3_1_meanflow_quads.py +++ b/examples/compressible/dcmip_3_1_meanflow_quads.py @@ -87,12 +87,13 @@ transported_fields = [ImplicitMidpoint(domain, "u"), SSPRK3(domain, "rho", subcycles=2), SSPRK3(domain, "theta", options=SUPGOptions(), subcycles=2)] +transport_methods = [DGUpwind(eqns, field) for field in ["u", "rho", "theta"]] # Linear solver linear_solver = CompressibleSolver(eqns) # Time stepper -stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, +stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, transport_methods, linear_solver=linear_solver) # ---------------------------------------------------------------------------- # diff --git a/examples/compressible/dry_bryan_fritsch.py b/examples/compressible/dry_bryan_fritsch.py index 2e3fc1f0d..265fef490 100644 --- a/examples/compressible/dry_bryan_fritsch.py +++ b/examples/compressible/dry_bryan_fritsch.py @@ -78,11 +78,14 @@ SSPRK3(domain, "theta", options=theta_opts), SSPRK3(domain, "u", options=u_opts)] +transport_methods = [DGUpwind(eqns, field) for field in ["u", "rho", "theta"]] + # Linear solver linear_solver = CompressibleSolver(eqns) # Time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + transport_methods, linear_solver=linear_solver) # ---------------------------------------------------------------------------- # diff --git a/examples/compressible/moist_bryan_fritsch.py b/examples/compressible/moist_bryan_fritsch.py index 7e3e65493..41f9930a9 100644 --- a/examples/compressible/moist_bryan_fritsch.py +++ b/examples/compressible/moist_bryan_fritsch.py @@ -64,6 +64,8 @@ SSPRK3(domain, "cloud_water", options=EmbeddedDGOptions()), ImplicitMidpoint(domain, "u")] +transport_methods = [DGUpwind(eqns, field) for field in ["u", "rho", "theta", "water_vapour", "cloud_water"]] + # Linear solver linear_solver = CompressibleSolver(eqns) @@ -72,6 +74,7 @@ # Time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + transport_methods, linear_solver=linear_solver, physics_schemes=physics_schemes) diff --git a/examples/compressible/mountain_hydrostatic.py b/examples/compressible/mountain_hydrostatic.py index c6077cfe9..a57b2c530 100644 --- a/examples/compressible/mountain_hydrostatic.py +++ b/examples/compressible/mountain_hydrostatic.py @@ -74,6 +74,9 @@ transported_fields = [ImplicitMidpoint(domain, "u"), SSPRK3(domain, "rho"), SSPRK3(domain, "theta", options=theta_opts)] +transport_methods = [DGUpwind(eqns, "u"), + DGUpwind(eqns, "rho"), + DGUpwind(eqns, "theta", ibp=theta_opts.ibp)] # Linear solver params = {'mat_type': 'matfree', @@ -100,6 +103,7 @@ # Time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + transport_methods, linear_solver=linear_solver, alpha=alpha) diff --git a/examples/compressible/mountain_nonhydrostatic.py b/examples/compressible/mountain_nonhydrostatic.py index 7d8cc0931..890f95ba2 100644 --- a/examples/compressible/mountain_nonhydrostatic.py +++ b/examples/compressible/mountain_nonhydrostatic.py @@ -73,12 +73,16 @@ transported_fields = [ImplicitMidpoint(domain, "u"), SSPRK3(domain, "rho"), SSPRK3(domain, "theta", options=theta_opts)] +transport_methods = [DGUpwind(eqns, "u"), + DGUpwind(eqns, "rho"), + DGUpwind(eqns, "theta", ibp=theta_opts.ibp)] # Linear solver linear_solver = CompressibleSolver(eqns) # Time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + transport_methods, linear_solver=linear_solver) # ---------------------------------------------------------------------------- # diff --git a/examples/compressible/robert_bubble.py b/examples/compressible/robert_bubble.py index 52c305a65..b89f6aa33 100644 --- a/examples/compressible/robert_bubble.py +++ b/examples/compressible/robert_bubble.py @@ -52,16 +52,20 @@ # Transport schemes theta_opts = EmbeddedDGOptions() -transported_fields = [] -transported_fields.append(ImplicitMidpoint(domain, "u")) -transported_fields.append(SSPRK3(domain, "rho")) -transported_fields.append(SSPRK3(domain, "theta", options=theta_opts)) +transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "rho"), + SSPRK3(domain, "theta", options=theta_opts)] + +transport_methods = [DGUpwind(eqns, "u"), + DGUpwind(eqns, "rho"), + DGUpwind(eqns, "theta")] # Linear solver linear_solver = CompressibleSolver(eqns) # Time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + transport_methods, linear_solver=linear_solver) # ---------------------------------------------------------------------------- # diff --git a/examples/compressible/skamarock_klemp_hydrostatic.py b/examples/compressible/skamarock_klemp_hydrostatic.py index 7e2e12255..bfdaf4107 100644 --- a/examples/compressible/skamarock_klemp_hydrostatic.py +++ b/examples/compressible/skamarock_klemp_hydrostatic.py @@ -55,16 +55,21 @@ io = IO(domain, output, diagnostic_fields=diagnostic_fields) # Transport schemes -transported_fields = [] -transported_fields.append(ImplicitMidpoint(domain, "u")) -transported_fields.append(SSPRK3(domain, "rho")) -transported_fields.append(SSPRK3(domain, "theta", options=SUPGOptions())) +theta_opts = SUPGOptions() +transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "rho"), + SSPRK3(domain, "theta", options=theta_opts)] + +transport_methods = [DGUpwind(eqns, "u"), + DGUpwind(eqns, "rho"), + DGUpwind(eqns, "theta", ibp=theta_opts.ibp)] # Linear solver linear_solver = CompressibleSolver(eqns) # Time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + transport_methods, linear_solver=linear_solver) # ---------------------------------------------------------------------------- # diff --git a/examples/compressible/skamarock_klemp_nonlinear.py b/examples/compressible/skamarock_klemp_nonlinear.py index 15d947a68..e84c01d64 100644 --- a/examples/compressible/skamarock_klemp_nonlinear.py +++ b/examples/compressible/skamarock_klemp_nonlinear.py @@ -67,12 +67,16 @@ transported_fields = [ImplicitMidpoint(domain, "u"), SSPRK3(domain, "rho"), SSPRK3(domain, "theta", options=theta_opts)] +transport_methods = [DGUpwind(eqns, "u"), + DGUpwind(eqns, "rho"), + DGUpwind(eqns, "theta", ibp=theta_opts.ibp)] # Linear solver linear_solver = CompressibleSolver(eqns) # Time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + transport_methods, linear_solver=linear_solver) # ---------------------------------------------------------------------------- # diff --git a/examples/compressible/straka_bubble.py b/examples/compressible/straka_bubble.py index 972ab491e..c696733ed 100644 --- a/examples/compressible/straka_bubble.py +++ b/examples/compressible/straka_bubble.py @@ -65,6 +65,9 @@ transported_fields = [ImplicitMidpoint(domain, "u"), SSPRK3(domain, "rho"), SSPRK3(domain, "theta", options=theta_opts)] + transport_methods = [DGUpwind(eqns, "u"), + DGUpwind(eqns, "rho"), + DGUpwind(eqns, "theta", ibp=theta_opts.ibp)] # Linear solver linear_solver = CompressibleSolver(eqns) @@ -72,9 +75,12 @@ # Diffusion schemes diffusion_schemes = [BackwardEuler(domain, "u"), BackwardEuler(domain, "theta")] + diffusion_methods = [InteriorPenaltyDiffusion(eqns, "u", diffusion_options[0]), + InteriorPenaltyDiffusion(eqns, "theta", diffusion_options[1])] # Time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + spatial_methods=transport_methods+diffusion_methods, linear_solver=linear_solver, diffusion_schemes=diffusion_schemes) diff --git a/examples/compressible/unsaturated_bubble.py b/examples/compressible/unsaturated_bubble.py index 6e23d343f..8b338a556 100644 --- a/examples/compressible/unsaturated_bubble.py +++ b/examples/compressible/unsaturated_bubble.py @@ -80,12 +80,15 @@ SSPRK3(domain, "cloud_water", options=theta_opts, limiter=limiter), SSPRK3(domain, "rain", options=theta_opts, limiter=limiter)] +transport_methods = [DGUpwind(eqns, field) for field in ["u", "rho", "theta", "water_vapour", "cloud_water", "rain"]] + # Linear solver linear_solver = CompressibleSolver(eqns) # Physics schemes # NB: to use wrapper options with Fallout, need to pass field name to time discretisation -physics_schemes = [(Fallout(eqns, 'rain', domain), SSPRK3(domain, field_name='rain', options=theta_opts, limiter=limiter)), +rainfall_method = DGUpwind(eqns, 'rain', outflow=True) +physics_schemes = [(Fallout(eqns, 'rain', domain, rainfall_method), SSPRK3(domain, field_name='rain', options=theta_opts, limiter=limiter)), (Coalescence(eqns), ForwardEuler(domain)), (EvaporationOfRain(eqns), ForwardEuler(domain)), (SaturationAdjustment(eqns), ForwardEuler(domain))] diff --git a/gusto/common_forms.py b/gusto/common_forms.py index 8089b9ad7..8ea6e72c8 100644 --- a/gusto/common_forms.py +++ b/gusto/common_forms.py @@ -89,7 +89,7 @@ def vector_invariant_form(domain, test, q, ubar): return transport(form, TransportEquationType.vector_invariant) -def kinetic_energy_form(domain, test, q): +def kinetic_energy_form(test, q, ubar): u""" The form corresponding to the kinetic energy term. @@ -98,24 +98,21 @@ def kinetic_energy_form(domain, test, q): (1/2)∇(u.q). Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. test (:class:`TestFunction`): the test function. q (:class:`ufl.Expr`): the variable to be transported. + ubar (:class:`ufl.Expr`): the transporting velocity. Returns: class:`LabelledForm`: a labelled transport form. """ - ubar = Function(domain.spaces("HDiv")) L = 0.5*div(test)*inner(q, ubar)*dx form = transporting_velocity(L, ubar) return transport(form, TransportEquationType.vector_invariant) -def advection_equation_circulation_form(domain, test, q, - ibp=IntegrateByParts.ONCE): +def advection_equation_circulation_form(domain, test, q, ubar): u""" The circulation term in the transport of a vector-valued field. @@ -131,13 +128,9 @@ def advection_equation_circulation_form(domain, test, q, term. An an upwind discretisation is used when integrating by parts. Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. test (:class:`TestFunction`): the test function. q (:class:`ufl.Expr`): the variable to be transported. - ibp (:class:`IntegrateByParts`, optional): an enumerator representing - the number of times to integrate by parts. Defaults to - `IntegrateByParts.ONCE`. + ubar (:class:`ufl.Expr`): the transporting velocity. Raises: NotImplementedError: the specified integration by parts is not 'once'. @@ -147,8 +140,8 @@ def advection_equation_circulation_form(domain, test, q, """ form = ( - vector_invariant_form(domain, test, q, ibp=ibp) - - kinetic_energy_form(domain, test, q) + vector_invariant_form(domain, test, q, ubar) + - kinetic_energy_form(test, q, ubar) ) return form diff --git a/gusto/diffusion_methods.py b/gusto/diffusion_methods.py index 6ce112574..eed339bcb 100644 --- a/gusto/diffusion_methods.py +++ b/gusto/diffusion_methods.py @@ -22,7 +22,7 @@ def __init__(self, equation, variable): """ # Inherited init method extracts original term to be replaced - super.__init__(self, equation, variable, diffusion) + super().__init__(equation, variable, diffusion) def interior_penalty_diffusion_form(domain, test, q, parameters): @@ -89,7 +89,7 @@ def __init__(self, equation, variable, diffusion_parameters): the kappa and mu constants. """ - super.__init__(equation, variable) + super().__init__(equation, variable) self.form = interior_penalty_diffusion_form(equation.domain, self.test, self.field, diffusion_parameters) diff --git a/gusto/equations.py b/gusto/equations.py index 169a9a4e1..4f3c372b6 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -15,7 +15,8 @@ from gusto.thermodynamics import exner_pressure from gusto.common_forms import (advection_form, continuity_form, vector_invariant_form, kinetic_energy_form, - advection_equation_circulation_form) + advection_equation_circulation_form, + diffusion_form) from gusto.active_tracers import ActiveTracer, Phases, TracerVariableType from gusto.configuration import TransportEquationType import ufl @@ -437,8 +438,15 @@ def generate_tracer_transport_terms(self, domain, active_tracers): # By default return None if no tracers are to be transported adv_form = None no_tracer_transported = True - u_idx = self.field_names.index('u') - u = split(self.X)[u_idx] + if 'u' in self.field_names: + u_idx = self.field_names.index('u') + u = split(self.X)[u_idx] + elif 'u' in self.prescribed_fields._field_names: + u = self.prescribed_fields('u') + else: + raise ValueError('Cannot generate tracer transport terms ' + +'as there is no velocity field') + for _, tracer in enumerate(active_tracers): if tracer.transport_eqn != TransportEquationType.no_transport: @@ -620,15 +628,14 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, elif u_transport_option == "vector_advection_form": u_adv = prognostic(advection_form(w, u, u), "u") elif u_transport_option == "circulation_form": - raise NotImplementedError - ke_form = prognostic(kinetic_energy_form(domain, w, u), "u") + ke_form = prognostic(kinetic_energy_form(w, u, u), "u") ke_form = transport.remove(ke_form) ke_form = ke_form.label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u}), t.labels)) ke_form = transporting_velocity.remove(ke_form) - u_adv = prognostic(advection_equation_circulation_form(domain, w, u), "u") + ke_form + u_adv = prognostic(advection_equation_circulation_form(domain, w, u, u), "u") + ke_form else: raise ValueError("Invalid u_transport_option: %s" % u_transport_option) @@ -644,7 +651,7 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, if self.thermal: gamma = self.tests[2] b = split(self.X)[2] - b_adv = prognostic(inner(gamma, inner(u, grad(b)))*dx, "b") + b_adv = prognostic(advection_form(gamma, b, u), "b") adv_form += subject(b_adv, self.X) # -------------------------------------------------------------------- # @@ -872,15 +879,14 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, elif u_transport_option == "vector_advection_form": u_adv = prognostic(advection_form(w, u, u), "u") elif u_transport_option == "circulation_form": - ke_form = prognostic(kinetic_energy_form(domain, w, u), "u") + ke_form = prognostic(kinetic_energy_form(w, u, u), "u") ke_form = transport.remove(ke_form) ke_form = ke_form.label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u}), t.labels)) ke_form = transporting_velocity.remove(ke_form) - u_adv = prognostic(advection_equation_circulation_form(domain, w, u), "u") + ke_form - raise NotImplementedError + u_adv = prognostic(advection_equation_circulation_form(domain, w, u, u), "u") + ke_form else: raise ValueError("Invalid u_transport_option: %s" % u_transport_option) @@ -1193,19 +1199,17 @@ def __init__(self, domain, parameters, Omega=None, # Velocity transport term -- depends on formulation if u_transport_option == "vector_invariant_form": u_adv = prognostic(vector_invariant_form(domain, w, u, u), "u") - raise NotImplementedError elif u_transport_option == "vector_advection_form": u_adv = prognostic(advection_form(w, u, u), "u") elif u_transport_option == "circulation_form": - ke_form = prognostic(kinetic_energy_form(domain, w, u), "u") + ke_form = prognostic(kinetic_energy_form(w, u, u), "u") ke_form = transport.remove(ke_form) ke_form = ke_form.label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u}), t.labels)) ke_form = transporting_velocity.remove(ke_form) - u_adv = prognostic(advection_equation_circulation_form(domain, w, u), "u") + ke_form - raise NotImplementedError + u_adv = prognostic(advection_equation_circulation_form(domain, w, u, u), "u") + ke_form else: raise ValueError("Invalid u_transport_option: %s" % u_transport_option) diff --git a/gusto/physics.py b/gusto/physics.py index 9f55588c3..19b3c1476 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -257,7 +257,8 @@ class Fallout(Physics): for Cartesian geometry. """ - def __init__(self, equation, rain_name, domain, moments=AdvectedMoments.M3): + def __init__(self, equation, rain_name, domain, transport_method, + moments=AdvectedMoments.M3): """ Args: equation (:class:`PrognosticEquationSet`): the model's equation. @@ -265,6 +266,8 @@ def __init__(self, equation, rain_name, domain, moments=AdvectedMoments.M3): 'rain'. domain (:class:`Domain`): the model's domain object, containing the mesh and the compatible function spaces. + transport_method (:class:`TransportMethod`): the spatial method + used for transporting the rain. moments (int, optional): an :class:`AdvectedMoments` enumerator, representing the number of moments of the size distribution of raindrops to be transported. Defaults to `AdvectedMoments.M3`. @@ -279,7 +282,6 @@ def __init__(self, equation, rain_name, domain, moments=AdvectedMoments.M3): rain_idx = equation.field_names.index(rain_name) rain = self.X.split()[rain_idx] - test = equation.tests[rain_idx] Vu = domain.spaces("HDiv") # TODO: there must be a better way than forcing this into the equation @@ -289,7 +291,9 @@ def __init__(self, equation, rain_name, domain, moments=AdvectedMoments.M3): # Create physics term -- which is actually a transport term # -------------------------------------------------------------------- # - adv_term = advection_form(domain, test, rain, outflow=True) + assert transport_method.outflow, \ + 'Fallout requires a transport method with outflow=True' + adv_term = transport_method.form.terms[0] # Add rainfall velocity by replacing transport_velocity in term adv_term = adv_term.label_map(identity, map_if_true=lambda t: Term( diff --git a/gusto/spatial_methods.py b/gusto/spatial_methods.py index 33a0a3bef..2ff9bd7aa 100644 --- a/gusto/spatial_methods.py +++ b/gusto/spatial_methods.py @@ -44,7 +44,7 @@ def __init__(self, equation, variable, term_label): map_if_true=keep, map_if_false=drop) num_terms = len(self.original_form.terms) - assert num_terms == 1, f'Unable to find {term_label.name} term ' \ + assert num_terms == 1, f'Unable to find {term_label.label} term ' \ + f'for {variable}. {num_terms} found' def replace_form(self, equation): diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 4bcb9b64f..15a83c916 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -98,7 +98,7 @@ def setup_equation(self, equation): method_variables = [method.variable for method in active_methods] for variable in active_variables: if variable not in method_variables: - message = f'Variable {variable} has a {term_label.name} ' \ + message = f'Variable {variable} has a {term_label.label} ' \ + 'but no method for this has been specified. Using ' \ + 'default form for this term' logger.warning(message) @@ -122,7 +122,7 @@ def setup_transporting_velocity(self, scheme): this discretisation. """ - if self.transporting_velocity == "prognostic": + if self.transporting_velocity == "prognostic" and "u" in self.fields._field_names: # Use the prognostic wind variable as the transporting velocity u_idx = self.equation.field_names.index('u') uadv = split(self.equation.X)[u_idx] diff --git a/gusto/transport_methods.py b/gusto/transport_methods.py index 6b8970dd1..60f52f8ea 100644 --- a/gusto/transport_methods.py +++ b/gusto/transport_methods.py @@ -3,7 +3,7 @@ """ from firedrake import (dx, dS, dS_v, dS_h, ds_t, ds_b, dot, inner, outer, jump, - grad, div, FacetNormal, Function) + grad, div, FacetNormal, Function, sign, avg, cross) from gusto.configuration import IntegrateByParts, TransportEquationType from gusto.fml import Term, keep, drop from gusto.labels import prognostic, transport, transporting_velocity, ibp_label @@ -253,6 +253,91 @@ def vector_manifold_continuity_form(domain, test, q, ibp=IntegrateByParts.ONCE, return transport(form) +def upwind_vector_invariant_form(domain, test, q, ibp=IntegrateByParts.ONCE): + u""" + The form corresponding to the DG upwind vector invariant transport operator. + + The self-transporting transport operator for a vector-valued field u can be + written as circulation and kinetic energy terms: + (u.∇)u = (∇×u)×u + (1/2)∇u^2 + + When the transporting field u and transported field q are similar, we write + this as: + (u.∇)q = (∇×q)×u + (1/2)∇(u.q) + + This form discretises this final equation, using an upwind discretisation + when integrating by parts. + + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + test (:class:`TestFunction`): the test function. + q (:class:`ufl.Expr`): the variable to be transported. + ibp (:class:`IntegrateByParts`, optional): an enumerator representing + the number of times to integrate by parts. Defaults to + `IntegrateByParts.ONCE`. + + Raises: + NotImplementedError: the specified integration by parts is not 'once'. + + Returns: + class:`LabelledForm`: a labelled transport form. + """ + + Vu = domain.spaces("HDiv") + dS_ = (dS_v + dS_h) if Vu.extruded else dS + ubar = Function(Vu) + n = FacetNormal(domain.mesh) + Upwind = 0.5*(sign(dot(ubar, n))+1) + + if domain.mesh.topological_dimension() == 3: + + if ibp != IntegrateByParts.ONCE: + raise NotImplementedError + + # + # = - + # = - + # <> + + both = lambda u: 2*avg(u) + + L = ( + inner(q, curl(cross(ubar, test)))*dx + - inner(both(Upwind*q), + both(cross(n, cross(ubar, test))))*dS_ + ) + + else: + + perp = domain.perp + if domain.on_sphere: + outward_normals = domain.outward_normals + perp_u_upwind = lambda q: Upwind('+')*cross(outward_normals('+'), q('+')) + Upwind('-')*cross(outward_normals('-'), q('-')) + else: + perp_u_upwind = lambda q: Upwind('+')*perp(q('+')) + Upwind('-')*perp(q('-')) + + if ibp == IntegrateByParts.ONCE: + L = ( + -inner(perp(grad(inner(test, perp(ubar)))), q)*dx + - inner(jump(inner(test, perp(ubar)), n), + perp_u_upwind(q))*dS_ + ) + else: + L = ( + (-inner(test, div(perp(q))*perp(ubar)))*dx + - inner(jump(inner(test, perp(ubar)), n), + perp_u_upwind(q))*dS_ + + jump(inner(test, + perp(ubar))*perp(q), n)*dS_ + ) + + L -= 0.5*div(test)*inner(q, ubar)*dx + + form = transporting_velocity(L, ubar) + + return transport(form, TransportEquationType.vector_invariant) + class DGUpwind(TransportMethod): """ diff --git a/integration-tests/balance/test_saturated_balance.py b/integration-tests/balance/test_saturated_balance.py index 810bd4d70..59b16fe3e 100644 --- a/integration-tests/balance/test_saturated_balance.py +++ b/integration-tests/balance/test_saturated_balance.py @@ -98,6 +98,7 @@ def setup_saturated(dirname, recovered): # Time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + transport_methods, linear_solver=linear_solver, physics_schemes=physics_schemes) diff --git a/integration-tests/equations/test_dry_compressible.py b/integration-tests/equations/test_dry_compressible.py index 2a4d20d56..53377335a 100644 --- a/integration-tests/equations/test_dry_compressible.py +++ b/integration-tests/equations/test_dry_compressible.py @@ -98,7 +98,7 @@ def run_dry_compressible(tmpdir): check_domain = Domain(check_mesh, dt, "CG", 1) check_eqn = CompressibleEulerEquations(check_domain, parameters) check_io = IO(check_domain, check_output) - check_stepper = SemiImplicitQuasiNewton(check_eqn, check_io, []) + check_stepper = SemiImplicitQuasiNewton(check_eqn, check_io, [], []) check_stepper.io.pick_up_from_checkpoint(check_stepper.fields) return stepper, check_stepper diff --git a/integration-tests/equations/test_incompressible.py b/integration-tests/equations/test_incompressible.py index cdc932cdd..12260d1c9 100644 --- a/integration-tests/equations/test_incompressible.py +++ b/integration-tests/equations/test_incompressible.py @@ -92,7 +92,7 @@ def run_incompressible(tmpdir): check_domain = Domain(check_mesh, dt, "CG", 1) check_eqn = IncompressibleBoussinesqEquations(check_domain, parameters) check_io = IO(check_domain, check_output) - check_stepper = SemiImplicitQuasiNewton(check_eqn, check_io, []) + check_stepper = SemiImplicitQuasiNewton(check_eqn, check_io, [], []) check_stepper.io.pick_up_from_checkpoint(check_stepper.fields) return stepper, check_stepper diff --git a/integration-tests/equations/test_moist_compressible.py b/integration-tests/equations/test_moist_compressible.py index e5613d242..222914334 100644 --- a/integration-tests/equations/test_moist_compressible.py +++ b/integration-tests/equations/test_moist_compressible.py @@ -104,7 +104,7 @@ def run_moist_compressible(tmpdir): check_domain = Domain(check_mesh, dt, "CG", 1) check_eqn = CompressibleEulerEquations(check_domain, parameters, active_tracers=tracers) check_io = IO(check_domain, output=check_output) - check_stepper = SemiImplicitQuasiNewton(check_eqn, check_io, []) + check_stepper = SemiImplicitQuasiNewton(check_eqn, check_io, [], []) check_stepper.io.pick_up_from_checkpoint(check_stepper.fields) return stepper, check_stepper diff --git a/integration-tests/equations/test_sw_fplane.py b/integration-tests/equations/test_sw_fplane.py index eb9dc5c9c..b508ca43d 100644 --- a/integration-tests/equations/test_sw_fplane.py +++ b/integration-tests/equations/test_sw_fplane.py @@ -109,7 +109,7 @@ def run_sw_fplane(tmpdir): check_domain = Domain(check_mesh, dt, 'RTCF', 1) check_eqn = ShallowWaterEquations(check_domain, parameters, fexpr=fexpr) check_io = IO(check_domain, output=check_output) - check_stepper = SemiImplicitQuasiNewton(check_eqn, check_io, []) + check_stepper = SemiImplicitQuasiNewton(check_eqn, check_io, [], []) check_stepper.io.pick_up_from_checkpoint(check_stepper.fields) return stepper, check_stepper diff --git a/integration-tests/physics/test_precipitation.py b/integration-tests/physics/test_precipitation.py index aded26c70..720c1d5f7 100644 --- a/integration-tests/physics/test_precipitation.py +++ b/integration-tests/physics/test_precipitation.py @@ -34,6 +34,7 @@ def setup_fallout(dirname): Vrho = domain.spaces("DG1_equispaced") active_tracers = [Rain(space='DG1_equispaced')] eqn = ForcedAdvectionEquation(domain, Vrho, "rho", active_tracers=active_tracers) + transport_method = DGUpwind(eqn, "rho") # I/O output = OutputParameters(dirname=dirname+"/fallout", dumpfreq=10, dumplist=['rain']) @@ -41,11 +42,12 @@ def setup_fallout(dirname): io = IO(domain, output, diagnostic_fields=diagnostic_fields) # Physics schemes - physics_schemes = [(Fallout(eqn, 'rain', domain), SSPRK3(domain, 'rain'))] + rainfall_method = DGUpwind(eqn, 'rain', outflow=True) + physics_schemes = [(Fallout(eqn, 'rain', domain, rainfall_method), SSPRK3(domain, 'rain'))] # build time stepper scheme = ForwardEuler(domain) - stepper = PrescribedTransport(eqn, scheme, io, + stepper = PrescribedTransport(eqn, scheme, io, transport_method, physics_schemes=physics_schemes) # ------------------------------------------------------------------------ # From 485dafcd476d580b0e24623ad677c0f9a0facc42 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Tue, 11 Jul 2023 20:45:02 +0100 Subject: [PATCH 14/20] fix remaining couple of bugs and lint, hopefully mostly working --- .../incompressible/skamarock_klemp_incompressible.py | 2 ++ examples/shallow_water/linear_williamson_2.py | 3 ++- examples/shallow_water/thermal_williamson2.py | 5 ++++- examples/shallow_water/williamson_2.py | 3 ++- examples/shallow_water/williamson_5.py | 3 ++- gusto/common_forms.py | 10 ++++++++-- gusto/diffusion_methods.py | 1 + gusto/equations.py | 4 ++-- gusto/fields.py | 2 ++ gusto/physics.py | 3 +-- gusto/spatial_methods.py | 1 + gusto/time_discretisation.py | 6 ++---- gusto/timeloop.py | 5 +++-- gusto/transport_methods.py | 8 ++++++-- integration-tests/model/test_nc_outputting.py | 3 ++- integration-tests/physics/test_instant_rain.py | 3 ++- integration-tests/physics/test_precipitation.py | 2 +- 17 files changed, 43 insertions(+), 21 deletions(-) diff --git a/examples/incompressible/skamarock_klemp_incompressible.py b/examples/incompressible/skamarock_klemp_incompressible.py index e2296967e..26efb5f52 100644 --- a/examples/incompressible/skamarock_klemp_incompressible.py +++ b/examples/incompressible/skamarock_klemp_incompressible.py @@ -56,12 +56,14 @@ b_opts = SUPGOptions() transported_fields = [ImplicitMidpoint(domain, "u"), SSPRK3(domain, "b", options=b_opts)] +transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "b", ibp=b_opts.ibp)] # Linear solver linear_solver = IncompressibleSolver(eqns) # Time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + transport_methods, linear_solver=linear_solver) # ---------------------------------------------------------------------------- # diff --git a/examples/shallow_water/linear_williamson_2.py b/examples/shallow_water/linear_williamson_2.py index 2a2c7905d..d5c72828d 100644 --- a/examples/shallow_water/linear_williamson_2.py +++ b/examples/shallow_water/linear_williamson_2.py @@ -53,9 +53,10 @@ # Transport schemes transport_schemes = [ForwardEuler(domain, "D")] +transport_methods = [DGUpwind(eqns, "D")] # Time stepper -stepper = SemiImplicitQuasiNewton(eqns, io, transport_schemes) +stepper = SemiImplicitQuasiNewton(eqns, io, transport_schemes, transport_methods) # ---------------------------------------------------------------------------- # # Initial conditions diff --git a/examples/shallow_water/thermal_williamson2.py b/examples/shallow_water/thermal_williamson2.py index 1dedef3af..052f65ea7 100644 --- a/examples/shallow_water/thermal_williamson2.py +++ b/examples/shallow_water/thermal_williamson2.py @@ -47,9 +47,12 @@ ShallowWaterPotentialEnstrophy(), SteadyStateError('u'), SteadyStateError('D')] io = IO(domain, output, diagnostic_fields=diagnostic_fields) +transport_methods = [DGUpwind(eqns, "u"), + DGUpwind(eqns, "D"), + DGUpwind(eqns, "b")] # Time stepper -stepper = Timestepper(eqns, RK4(domain), io) +stepper = Timestepper(eqns, RK4(domain), io, spatial_methods=transport_methods) # ----------------------------------------------------------------- # # Initial conditions diff --git a/examples/shallow_water/williamson_2.py b/examples/shallow_water/williamson_2.py index 37306d472..1c2a8f8ec 100644 --- a/examples/shallow_water/williamson_2.py +++ b/examples/shallow_water/williamson_2.py @@ -68,9 +68,10 @@ # Transport schemes transported_fields = [ImplicitMidpoint(domain, "u"), SSPRK3(domain, "D", subcycles=2)] + transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "D")] # Time stepper - stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields) + stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, transport_methods) # ------------------------------------------------------------------------ # # Initial conditions diff --git a/examples/shallow_water/williamson_5.py b/examples/shallow_water/williamson_5.py index d16e66555..bb6b5d9b3 100644 --- a/examples/shallow_water/williamson_5.py +++ b/examples/shallow_water/williamson_5.py @@ -72,9 +72,10 @@ # Transport schemes transported_fields = [ImplicitMidpoint(domain, "u"), SSPRK3(domain, "D")] + transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "D")] # Time stepper - stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields) + stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, transport_methods) # ------------------------------------------------------------------------ # # Initial conditions diff --git a/gusto/common_forms.py b/gusto/common_forms.py index 8ea6e72c8..0f4c33747 100644 --- a/gusto/common_forms.py +++ b/gusto/common_forms.py @@ -2,14 +2,15 @@ Provides some basic forms for discretising various common terms in equations for geophysical fluid dynamics.""" -from firedrake import Function, dx, dot, grad, div, inner, outer, cross, curl -from gusto.configuration import IntegrateByParts, TransportEquationType +from firedrake import dx, dot, grad, div, inner, outer, cross, curl +from gusto.configuration import TransportEquationType from gusto.labels import transport, transporting_velocity, diffusion __all__ = ["advection_form", "continuity_form", "vector_invariant_form", "kinetic_energy_form", "advection_equation_circulation_form", "diffusion_form"] + def advection_form(test, q, ubar): u""" The form corresponding to the advective transport operator. @@ -30,6 +31,7 @@ def advection_form(test, q, ubar): return transport(form, TransportEquationType.advective) + def continuity_form(test, q, ubar): u""" The form corresponding to the continuity transport operator. @@ -50,6 +52,7 @@ def continuity_form(test, q, ubar): return transport(form, TransportEquationType.conservative) + def vector_invariant_form(domain, test, q, ubar): u""" The form corresponding to the vector invariant transport operator. @@ -89,6 +92,7 @@ def vector_invariant_form(domain, test, q, ubar): return transport(form, TransportEquationType.vector_invariant) + def kinetic_energy_form(test, q, ubar): u""" The form corresponding to the kinetic energy term. @@ -112,6 +116,7 @@ def kinetic_energy_form(test, q, ubar): return transport(form, TransportEquationType.vector_invariant) + def advection_equation_circulation_form(domain, test, q, ubar): u""" The circulation term in the transport of a vector-valued field. @@ -146,6 +151,7 @@ def advection_equation_circulation_form(domain, test, q, ubar): return form + def diffusion_form(test, q, kappa): u""" The diffusion form, ∇.(κ∇q) for diffusivity κ and variable q. diff --git a/gusto/diffusion_methods.py b/gusto/diffusion_methods.py index eed339bcb..42221952d 100644 --- a/gusto/diffusion_methods.py +++ b/gusto/diffusion_methods.py @@ -75,6 +75,7 @@ def get_flux_form(dS, M): return diffusion(form) + class InteriorPenaltyDiffusion(DiffusionMethod): """The interior penalty method for discretising the diffusion term.""" diff --git a/gusto/equations.py b/gusto/equations.py index 4f3c372b6..2e3e21e5a 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -438,6 +438,7 @@ def generate_tracer_transport_terms(self, domain, active_tracers): # By default return None if no tracers are to be transported adv_form = None no_tracer_transported = True + if 'u' in self.field_names: u_idx = self.field_names.index('u') u = split(self.X)[u_idx] @@ -445,8 +446,7 @@ def generate_tracer_transport_terms(self, domain, active_tracers): u = self.prescribed_fields('u') else: raise ValueError('Cannot generate tracer transport terms ' - +'as there is no velocity field') - + + 'as there is no velocity field') for _, tracer in enumerate(active_tracers): if tracer.transport_eqn != TransportEquationType.no_transport: diff --git a/gusto/fields.py b/gusto/fields.py index d203a30b9..cfa70cb08 100644 --- a/gusto/fields.py +++ b/gusto/fields.py @@ -58,6 +58,7 @@ class PrescribedFields(Fields): """Object to hold and create a specified set of prescribed fields.""" def __init__(self): self.fields = [] + self._field_names = [] def __call__(self, name, space=None): """ @@ -80,6 +81,7 @@ def __call__(self, name, space=None): else: # Create field self.add_field(name, space) + self._field_names.append(name) return getattr(self, name) def __iter__(self): diff --git a/gusto/physics.py b/gusto/physics.py index 19b3c1476..f8240d4a1 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -11,7 +11,6 @@ from gusto.active_tracers import Phases from gusto.recovery import Recoverer, BoundaryMethod from gusto.equations import CompressibleEulerEquations -from gusto.common_forms import advection_form from gusto.fml import identity, Term from gusto.labels import subject, physics, transporting_velocity from gusto.configuration import logger @@ -293,7 +292,7 @@ def __init__(self, equation, rain_name, domain, transport_method, assert transport_method.outflow, \ 'Fallout requires a transport method with outflow=True' - adv_term = transport_method.form.terms[0] + adv_term = transport_method.form # Add rainfall velocity by replacing transport_velocity in term adv_term = adv_term.label_map(identity, map_if_true=lambda t: Term( diff --git a/gusto/spatial_methods.py b/gusto/spatial_methods.py index 2ff9bd7aa..432b0c588 100644 --- a/gusto/spatial_methods.py +++ b/gusto/spatial_methods.py @@ -9,6 +9,7 @@ __all__ = ['SpatialMethod'] + class SpatialMethod(object): """ The base object for describing a spatial discretisation of some term. diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index ba7b9b6ef..fda83bb82 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -6,12 +6,10 @@ """ from abc import ABCMeta, abstractmethod, abstractproperty -from firedrake import (Function, NonlinearVariationalProblem, split, - NonlinearVariationalSolver, TestFunction, dot, grad, - DirichletBC) +from firedrake import (Function, TestFunction, NonlinearVariationalProblem, + NonlinearVariationalSolver, DirichletBC) from firedrake.formmanipulation import split_form from firedrake.utils import cached_property -import ufl from gusto.configuration import (logger, DEBUG, EmbeddedDGOptions, RecoveryOptions) from gusto.labels import (time_derivative, prognostic, physics, replace_subject, replace_test_function) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 15a83c916..7f5714a8f 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -91,7 +91,7 @@ def setup_equation(self, equation): # Extract all terms corresponding to this type of term residual = equation.residual.label_map( lambda t: t.has_label(term_label), map_if_false=drop - ) + ) active_variables = [t.get(prognostic) for t in residual.terms] active_methods = list(filter(lambda t: t.term_label == term_label, self.spatial_methods)) @@ -390,6 +390,7 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, for scheme in transport_schemes: assert scheme.nlevels == 1, "multilevel schemes not supported as part of this timestepping loop" assert scheme.field_name in equation_set.field_names + self.active_transport.append((scheme.field_name, scheme)) # Check that there is a corresponding transport method method_found = False for method in spatial_methods: @@ -528,7 +529,7 @@ def timestep(self): for k in range(self.maxk): with timed_stage("Transport"): - for name, scheme, _ in self.active_transport: + for name, scheme in self.active_transport: # transports a field from xstar and puts result in xp scheme.apply(xp(name), xstar(name)) diff --git a/gusto/transport_methods.py b/gusto/transport_methods.py index 60f52f8ea..247278835 100644 --- a/gusto/transport_methods.py +++ b/gusto/transport_methods.py @@ -2,8 +2,9 @@ Defines TransportMethod objects, which are used to solve a transport problem. """ -from firedrake import (dx, dS, dS_v, dS_h, ds_t, ds_b, dot, inner, outer, jump, - grad, div, FacetNormal, Function, sign, avg, cross) +from firedrake import (dx, dS, dS_v, dS_h, ds_t, ds_b, ds_v, dot, inner, outer, + jump, grad, div, FacetNormal, Function, sign, avg, cross, + curl) from gusto.configuration import IntegrateByParts, TransportEquationType from gusto.fml import Term, keep, drop from gusto.labels import prognostic, transport, transporting_velocity, ibp_label @@ -11,6 +12,7 @@ __all__ = ["DGUpwind"] + class TransportMethod(SpatialMethod): """ The base object for describing a transport scheme. @@ -215,6 +217,7 @@ def vector_manifold_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, o return L + def vector_manifold_continuity_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): """ Form for continuity transport operator including vector manifold correction. @@ -253,6 +256,7 @@ def vector_manifold_continuity_form(domain, test, q, ibp=IntegrateByParts.ONCE, return transport(form) + def upwind_vector_invariant_form(domain, test, q, ibp=IntegrateByParts.ONCE): u""" The form corresponding to the DG upwind vector invariant transport operator. diff --git a/integration-tests/model/test_nc_outputting.py b/integration-tests/model/test_nc_outputting.py index fe26c18d4..2b3a0a4da 100644 --- a/integration-tests/model/test_nc_outputting.py +++ b/integration-tests/model/test_nc_outputting.py @@ -9,7 +9,8 @@ as_vector, exp) from gusto import (Domain, IO, PrescribedTransport, AdvectionEquation, ForwardEuler, OutputParameters, VelocityX, VelocityY, - VelocityZ, MeridionalComponent, ZonalComponent, RadialComponent) + VelocityZ, MeridionalComponent, ZonalComponent, + RadialComponent, DGUpwind) from netCDF4 import Dataset import pytest diff --git a/integration-tests/physics/test_instant_rain.py b/integration-tests/physics/test_instant_rain.py index 8f9f00308..04bf3c736 100644 --- a/integration-tests/physics/test_instant_rain.py +++ b/integration-tests/physics/test_instant_rain.py @@ -46,6 +46,7 @@ def run_instant_rain(dirname): dumplist=['vapour', "rain"]) diagnostic_fields = [CourantNumber()] io = IO(domain, output, diagnostic_fields=diagnostic_fields) + transport_method = DGUpwind(eqns, "water_vapour") # Physics schemes # define saturation function @@ -54,7 +55,7 @@ def run_instant_rain(dirname): ForwardEuler(domain))] # Time stepper - stepper = PrescribedTransport(eqns, RK4(domain), io, + stepper = PrescribedTransport(eqns, RK4(domain), transport_method, io, physics_schemes=physics_schemes) # ------------------------------------------------------------------------ # diff --git a/integration-tests/physics/test_precipitation.py b/integration-tests/physics/test_precipitation.py index 720c1d5f7..037b0ef1a 100644 --- a/integration-tests/physics/test_precipitation.py +++ b/integration-tests/physics/test_precipitation.py @@ -47,7 +47,7 @@ def setup_fallout(dirname): # build time stepper scheme = ForwardEuler(domain) - stepper = PrescribedTransport(eqn, scheme, io, transport_method, + stepper = PrescribedTransport(eqn, scheme, transport_method, io, physics_schemes=physics_schemes) # ------------------------------------------------------------------------ # From 5481bcf4e2beade8ea762315fd0a917d595e7236 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Wed, 12 Jul 2023 20:32:53 +0100 Subject: [PATCH 15/20] reorder arguments for PrescribedTransport stepper --- gusto/timeloop.py | 12 +- gusto/transport_methods.py | 164 +++++++++++------- .../equations/test_linear_sw_wave.py | 3 +- .../equations/test_sw_linear_triangle.py | 2 +- integration-tests/model/test_nc_outputting.py | 2 +- .../model/test_prescribed_transport.py | 5 +- .../model/test_time_discretisation.py | 2 +- .../physics/test_instant_rain.py | 2 +- .../physics/test_precipitation.py | 2 +- .../transport/test_dg_transport.py | 4 +- .../transport/test_embedded_dg_advection.py | 2 +- integration-tests/transport/test_limiters.py | 2 +- .../transport/test_recovered_transport.py | 2 +- .../transport/test_subcycling.py | 2 +- .../transport/test_supg_transport.py | 4 +- .../transport/test_vector_recovered_space.py | 2 +- 16 files changed, 130 insertions(+), 82 deletions(-) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 7f5714a8f..46321cf43 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -12,6 +12,7 @@ from gusto.linear_solvers import LinearTimesteppingSolver from gusto.fields import TimeLevelFields, StateFields from gusto.time_discretisation import ExplicitTimeDiscretisation +from gusto.transport_methods import TransportMethod import ufl __all__ = ["Timestepper", "SplitPhysicsTimestepper", "SemiImplicitQuasiNewton", @@ -586,7 +587,7 @@ class PrescribedTransport(Timestepper): """ Implements a timeloop with a prescibed transporting velocity. """ - def __init__(self, equation, scheme, transport_method, io, + def __init__(self, equation, scheme, io, transport_method, physics_schemes=None, prescribed_transporting_velocity=None): """ Args: @@ -609,8 +610,13 @@ def __init__(self, equation, scheme, transport_method, io, updated. Defaults to None. """ - super().__init__(equation, scheme, io, - spatial_methods=[transport_method]) + if isinstance(transport_method, TransportMethod): + transport_methods = [transport_method] + else: + # Assume an iterable has been provided + transport_methods = transport_method + + super().__init__(equation, scheme, io, spatial_methods=transport_methods) if physics_schemes is not None: self.physics_schemes = physics_schemes diff --git a/gusto/transport_methods.py b/gusto/transport_methods.py index 247278835..74f18b654 100644 --- a/gusto/transport_methods.py +++ b/gusto/transport_methods.py @@ -10,9 +10,12 @@ from gusto.labels import prognostic, transport, transporting_velocity, ibp_label from gusto.spatial_methods import SpatialMethod -__all__ = ["DGUpwind"] +__all__ = ["DefaultTransport", "DGUpwind"] +# ---------------------------------------------------------------------------- # +# Base TransportMethod class +# ---------------------------------------------------------------------------- # class TransportMethod(SpatialMethod): """ The base object for describing a transport scheme. @@ -64,6 +67,104 @@ def replace_form(self, equation): map_if_true=lambda t: new_term) +# ---------------------------------------------------------------------------- # +# TransportMethod for using underlying default transport form +# ---------------------------------------------------------------------------- # +class DefaultTransport(TransportMethod): + """ + Performs no manipulation of the transport form, so the scheme is simply + based on the transport terms that are declared when the equation is set up. + """ + def __init__(self, equation, variable): + """ + Args: + equation (:class:`PrognosticEquation`): the equation, which includes + a transport term. + variable (str): name of the variable to set the transport scheme for + """ + + super().__init__(equation, variable) + + def replace_form(self, equation): + """ + In theory replaces the transport form in the equation, but in this case + does nothing. + + Args: + equation (:class:`PrognosticEquation`): the equation or scheme whose + transport term should (not!) be replaced with the transport term + of this discretisation. + """ + pass + + +# ---------------------------------------------------------------------------- # +# Class for DG Upwind transport methods +# ---------------------------------------------------------------------------- # +class DGUpwind(TransportMethod): + """ + The Discontinuous Galerkin Upwind transport scheme. + + Discretises the gradient of a field weakly, taking the upwind value of the + transported variable at facets. + """ + def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, + vector_manifold_correction=False, outflow=False): + """ + Args: + equation (:class:`PrognosticEquation`): the equation, which includes + a transport term. + variable (str): name of the variable to set the transport scheme for + ibp (:class:`IntegrateByParts`, optional): an enumerator for how + many times to integrate by parts. Defaults to `ONCE`. + vector_manifold_correction (bool, optional): whether to include a + vector manifold correction term. Defaults to False. + outflow (bool, optional): whether to include outflow at the domain + boundaries, through exterior facet terms. Defaults to False. + """ + + super().__init__(equation, variable) + self.ibp = ibp + self.vector_manifold_correction = vector_manifold_correction + self.outflow = outflow + + # -------------------------------------------------------------------- # + # Determine appropriate form to use + # -------------------------------------------------------------------- # + + if self.transport_equation_type == TransportEquationType.advective: + if vector_manifold_correction: + form = vector_manifold_advection_form(self.domain, self.test, + self.field, ibp=ibp, + outflow=outflow) + else: + form = upwind_advection_form(self.domain, self.test, self.field, + ibp=ibp, outflow=outflow) + + elif self.transport_equation_type == TransportEquationType.conservative: + if vector_manifold_correction: + form = vector_manifold_continuity_form(self.domain, self.test, + self.field, ibp=ibp, + outflow=outflow) + else: + form = upwind_continuity_form(self.domain, self.test, self.field, + ibp=ibp, outflow=outflow) + + elif self.transport_equation_type == TransportEquationType.vector_invariant: + if outflow: + raise NotImplementedError('Outflow not implemented for upwind vector invariant') + form = upwind_vector_invariant_form(self.domain, self.test, self.field, ibp=ibp) + + else: + raise NotImplementedError('Upwind transport scheme has not been ' + + 'implemented for this transport equation type') + + self.form = form + + +# ---------------------------------------------------------------------------- # +# Forms for DG Upwind transport +# ---------------------------------------------------------------------------- # def upwind_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): u""" The form corresponding to the DG upwind advective transport operator. @@ -341,64 +442,3 @@ def upwind_vector_invariant_form(domain, test, q, ibp=IntegrateByParts.ONCE): form = transporting_velocity(L, ubar) return transport(form, TransportEquationType.vector_invariant) - - -class DGUpwind(TransportMethod): - """ - The Discontinuous Galerkin Upwind transport scheme. - - Discretises the gradient of a field weakly, taking the upwind value of the - transported variable at facets. - """ - def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, - vector_manifold_correction=False, outflow=False): - """ - Args: - equation (:class:`PrognosticEquation`): the equation, which includes - a transport term. - variable (str): name of the variable to set the transport scheme for - ibp (:class:`IntegrateByParts`, optional): an enumerator for how - many times to integrate by parts. Defaults to `ONCE`. - vector_manifold_correction (bool, optional): whether to include a - vector manifold correction term. Defaults to False. - outflow (bool, optional): whether to include outflow at the domain - boundaries, through exterior facet terms. Defaults to False. - """ - - super().__init__(equation, variable) - self.ibp = ibp - self.vector_manifold_correction = vector_manifold_correction - self.outflow = outflow - - # -------------------------------------------------------------------- # - # Determine appropriate form to use - # -------------------------------------------------------------------- # - - if self.transport_equation_type == TransportEquationType.advective: - if vector_manifold_correction: - form = vector_manifold_advection_form(self.domain, self.test, - self.field, ibp=ibp, - outflow=outflow) - else: - form = upwind_advection_form(self.domain, self.test, self.field, - ibp=ibp, outflow=outflow) - - elif self.transport_equation_type == TransportEquationType.conservative: - if vector_manifold_correction: - form = vector_manifold_continuity_form(self.domain, self.test, - self.field, ibp=ibp, - outflow=outflow) - else: - form = upwind_continuity_form(self.domain, self.test, self.field, - ibp=ibp, outflow=outflow) - - elif self.transport_equation_type == TransportEquationType.vector_invariant: - if outflow: - raise NotImplementedError('Outflow not implemented for upwind vector invariant') - form = upwind_vector_invariant_form(self.domain, self.test, self.field, ibp=ibp) - - else: - raise NotImplementedError('Upwind transport scheme has not been ' - + 'implemented for this transport equation type') - - self.form = form diff --git a/integration-tests/equations/test_linear_sw_wave.py b/integration-tests/equations/test_linear_sw_wave.py index c146668d9..ffce84309 100644 --- a/integration-tests/equations/test_linear_sw_wave.py +++ b/integration-tests/equations/test_linear_sw_wave.py @@ -37,9 +37,10 @@ def run_linear_sw_wave(tmpdir): dumpfreq=1, log_level='INFO') io = IO(domain, output) + transport_methods = [DefaultTransport(eqns, "D")] # Timestepper - stepper = Timestepper(eqns, RK4(domain), io) + stepper = Timestepper(eqns, RK4(domain), io, transport_methods) # ---------------------------------------------------------------------- # # Initial conditions diff --git a/integration-tests/equations/test_sw_linear_triangle.py b/integration-tests/equations/test_sw_linear_triangle.py index e66b072bf..d533b3286 100644 --- a/integration-tests/equations/test_sw_linear_triangle.py +++ b/integration-tests/equations/test_sw_linear_triangle.py @@ -42,7 +42,7 @@ def setup_sw(dirname): # Transport schemes transport_schemes = [ForwardEuler(domain, "D")] - transport_methods = [] + transport_methods = [DefaultTransport(eqns, "D")] # Time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transport_schemes, transport_methods) diff --git a/integration-tests/model/test_nc_outputting.py b/integration-tests/model/test_nc_outputting.py index 2b3a0a4da..ca9519a6c 100644 --- a/integration-tests/model/test_nc_outputting.py +++ b/integration-tests/model/test_nc_outputting.py @@ -91,7 +91,7 @@ def test_nc_outputting(tmpdir, geometry, domain_and_mesh_details): diagnostic_fields = [ZonalComponent('u'), MeridionalComponent('u'), RadialComponent('u')] io = IO(domain, output, diagnostic_fields=diagnostic_fields) - stepper = PrescribedTransport(eqn, transport_scheme, transport_method, io) + stepper = PrescribedTransport(eqn, transport_scheme, io, transport_method) # ------------------------------------------------------------------------ # # Initialise fields diff --git a/integration-tests/model/test_prescribed_transport.py b/integration-tests/model/test_prescribed_transport.py index dc64acc82..cf31b2bfc 100644 --- a/integration-tests/model/test_prescribed_transport.py +++ b/integration-tests/model/test_prescribed_transport.py @@ -33,8 +33,9 @@ def u_evaluation(t): transport_scheme = SSPRK3(domain) transport_method = DGUpwind(eqn, 'f') - timestepper = PrescribedTransport(eqn, transport_scheme, transport_method, - setup.io, prescribed_transporting_velocity=u_evaluation) + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io, + transport_method, + prescribed_transporting_velocity=u_evaluation) # Initial conditions timestepper.fields("f").interpolate(setup.f_init) diff --git a/integration-tests/model/test_time_discretisation.py b/integration-tests/model/test_time_discretisation.py index 150f5c842..1b1f04106 100644 --- a/integration-tests/model/test_time_discretisation.py +++ b/integration-tests/model/test_time_discretisation.py @@ -49,7 +49,7 @@ def test_time_discretisation(tmpdir, scheme, tracer_setup): transport_method = DGUpwind(eqn, 'f') - timestepper = PrescribedTransport(eqn, transport_scheme, transport_method, setup.io) + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io, transport_method) # Initial conditions timestepper.fields("f").interpolate(setup.f_init) diff --git a/integration-tests/physics/test_instant_rain.py b/integration-tests/physics/test_instant_rain.py index 04bf3c736..77b95c6d2 100644 --- a/integration-tests/physics/test_instant_rain.py +++ b/integration-tests/physics/test_instant_rain.py @@ -55,7 +55,7 @@ def run_instant_rain(dirname): ForwardEuler(domain))] # Time stepper - stepper = PrescribedTransport(eqns, RK4(domain), transport_method, io, + stepper = PrescribedTransport(eqns, RK4(domain), io, transport_method, physics_schemes=physics_schemes) # ------------------------------------------------------------------------ # diff --git a/integration-tests/physics/test_precipitation.py b/integration-tests/physics/test_precipitation.py index 037b0ef1a..720c1d5f7 100644 --- a/integration-tests/physics/test_precipitation.py +++ b/integration-tests/physics/test_precipitation.py @@ -47,7 +47,7 @@ def setup_fallout(dirname): # build time stepper scheme = ForwardEuler(domain) - stepper = PrescribedTransport(eqn, scheme, transport_method, io, + stepper = PrescribedTransport(eqn, scheme, io, transport_method, physics_schemes=physics_schemes) # ------------------------------------------------------------------------ # diff --git a/integration-tests/transport/test_dg_transport.py b/integration-tests/transport/test_dg_transport.py index 3b499e56f..1117c6253 100644 --- a/integration-tests/transport/test_dg_transport.py +++ b/integration-tests/transport/test_dg_transport.py @@ -28,7 +28,7 @@ def test_dg_transport_scalar(tmpdir, geometry, equation_form, tracer_setup): transport_scheme = SSPRK3(domain) transport_method = DGUpwind(eqn, "f") - timestepper = PrescribedTransport(eqn, transport_scheme, transport_method, setup.io) + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io, transport_method) # Initial conditions timestepper.fields("f").interpolate(setup.f_init) @@ -55,7 +55,7 @@ def test_dg_transport_vector(tmpdir, geometry, equation_form, tracer_setup): transport_scheme = SSPRK3(domain) transport_method = DGUpwind(eqn, "f") - timestepper = PrescribedTransport(eqn, transport_scheme, transport_method, setup.io) + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io, transport_method) # Initial conditions timestepper.fields("f").interpolate(f_init) diff --git a/integration-tests/transport/test_embedded_dg_advection.py b/integration-tests/transport/test_embedded_dg_advection.py index 0cc023385..dd9ec357b 100644 --- a/integration-tests/transport/test_embedded_dg_advection.py +++ b/integration-tests/transport/test_embedded_dg_advection.py @@ -35,7 +35,7 @@ def test_embedded_dg_advection_scalar(tmpdir, ibp, equation_form, space, transport_schemes = SSPRK3(domain, options=opts) transport_method = DGUpwind(eqn, "f", ibp=ibp) - timestepper = PrescribedTransport(eqn, transport_schemes, transport_method, setup.io) + timestepper = PrescribedTransport(eqn, transport_schemes, setup.io, transport_method) # Initial conditions timestepper.fields("f").interpolate(setup.f_init) diff --git a/integration-tests/transport/test_limiters.py b/integration-tests/transport/test_limiters.py index b0ce302a5..f68d45a2d 100644 --- a/integration-tests/transport/test_limiters.py +++ b/integration-tests/transport/test_limiters.py @@ -90,7 +90,7 @@ def setup_limiters(dirname, space): transport_method = DGUpwind(eqn, "tracer") # Build time stepper - stepper = PrescribedTransport(eqn, transport_schemes, transport_method, io) + stepper = PrescribedTransport(eqn, transport_schemes, io, transport_method) # ------------------------------------------------------------------------ # # Initial condition diff --git a/integration-tests/transport/test_recovered_transport.py b/integration-tests/transport/test_recovered_transport.py index 862714c3f..c8cebcda6 100644 --- a/integration-tests/transport/test_recovered_transport.py +++ b/integration-tests/transport/test_recovered_transport.py @@ -38,7 +38,7 @@ def test_recovered_space_setup(tmpdir, geometry, tracer_setup): transport_scheme = SSPRK3(domain, options=recovery_opts) transport_method = DGUpwind(eqn, "f") - timestepper = PrescribedTransport(eqn, transport_scheme, transport_method, setup.io) + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io, transport_method) # Initialise fields timestepper.fields("f").interpolate(setup.f_init) diff --git a/integration-tests/transport/test_subcycling.py b/integration-tests/transport/test_subcycling.py index 97a5a31db..5894050b5 100644 --- a/integration-tests/transport/test_subcycling.py +++ b/integration-tests/transport/test_subcycling.py @@ -27,7 +27,7 @@ def test_subcyling(tmpdir, equation_form, tracer_setup): transport_scheme = SSPRK3(domain, subcycles=2) transport_method = DGUpwind(eqn, "f") - timestepper = PrescribedTransport(eqn, transport_scheme, transport_method, setup.io) + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io, transport_method) # Initial conditions timestepper.fields("f").interpolate(setup.f_init) diff --git a/integration-tests/transport/test_supg_transport.py b/integration-tests/transport/test_supg_transport.py index 84642d3ec..30422d8e6 100644 --- a/integration-tests/transport/test_supg_transport.py +++ b/integration-tests/transport/test_supg_transport.py @@ -42,7 +42,7 @@ def test_supg_transport_scalar(tmpdir, equation_form, scheme, space, transport_scheme = ImplicitMidpoint(domain, options=opts) transport_method = DGUpwind(eqn, "f", ibp=ibp) - timestepper = PrescribedTransport(eqn, transport_scheme, transport_method, setup.io) + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io, transport_method) # Initial conditions timestepper.fields("f").interpolate(setup.f_init) @@ -84,7 +84,7 @@ def test_supg_transport_vector(tmpdir, equation_form, scheme, space, transport_scheme = ImplicitMidpoint(domain, options=opts) transport_method = DGUpwind(eqn, "f", ibp=ibp) - timestepper = PrescribedTransport(eqn, transport_scheme, transport_method, setup.io) + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io, transport_method) # Initial conditions f = timestepper.fields("f") diff --git a/integration-tests/transport/test_vector_recovered_space.py b/integration-tests/transport/test_vector_recovered_space.py index 1c67bc3a9..b419b39fb 100644 --- a/integration-tests/transport/test_vector_recovered_space.py +++ b/integration-tests/transport/test_vector_recovered_space.py @@ -42,7 +42,7 @@ def test_vector_recovered_space_setup(tmpdir, geometry, tracer_setup): eqn = AdvectionEquation(domain, Vu, "f") transport_scheme = SSPRK3(domain, options=rec_opts) transport_method = DGUpwind(eqn, "f") - timestepper = PrescribedTransport(eqn, transport_scheme, transport_method, setup.io) + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io, transport_method) # Initialise fields f_init = as_vector([setup.f_init]*gdim) From 19629601d63aa0368257e2dde031ed5906a71f64 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Wed, 12 Jul 2023 21:11:53 +0100 Subject: [PATCH 16/20] fix linearisations (by restoring them) --- gusto/common_forms.py | 40 +++++++++++++++++++++++++++++++++++++++- gusto/equations.py | 29 ++++++++++++++++++++++++++--- 2 files changed, 65 insertions(+), 4 deletions(-) diff --git a/gusto/common_forms.py b/gusto/common_forms.py index 0f4c33747..56f17d2c5 100644 --- a/gusto/common_forms.py +++ b/gusto/common_forms.py @@ -8,7 +8,7 @@ __all__ = ["advection_form", "continuity_form", "vector_invariant_form", "kinetic_energy_form", "advection_equation_circulation_form", - "diffusion_form"] + "diffusion_form", "linear_advection_form", "linear_continuity_form"] def advection_form(test, q, ubar): @@ -32,6 +32,25 @@ def advection_form(test, q, ubar): return transport(form, TransportEquationType.advective) +def linear_advection_form(test, qbar, ubar): + """ + The form corresponding to the linearised advective transport operator. + + Args: + test (:class:`TestFunction`): the test function. + qbar (:class:`ufl.Expr`): the variable to be transported. + ubar (:class:`ufl.Expr`): the transporting velocity. + + Returns: + :class:`LabelledForm`: a labelled transport form. + """ + + L = test*dot(ubar, grad(qbar))*dx + form = transporting_velocity(L, ubar) + + return transport(form, TransportEquationType.advective) + + def continuity_form(test, q, ubar): u""" The form corresponding to the continuity transport operator. @@ -53,6 +72,25 @@ def continuity_form(test, q, ubar): return transport(form, TransportEquationType.conservative) +def linear_continuity_form(test, qbar, ubar): + """ + The form corresponding to the linearised continuity transport operator. + + Args: + test (:class:`TestFunction`): the test function. + qbar (:class:`ufl.Expr`): the variable to be transported. + ubar (:class:`ufl.Expr`): the transporting velocity. + + Returns: + :class:`LabelledForm`: a labelled transport form. + """ + + L = qbar*test*div(ubar)*dx + form = transporting_velocity(L, ubar) + + return transport(form, TransportEquationType.conservative) + + def vector_invariant_form(domain, test, q, ubar): u""" The form corresponding to the vector invariant transport operator. diff --git a/gusto/equations.py b/gusto/equations.py index 2e3e21e5a..20d20bde4 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -16,7 +16,7 @@ from gusto.common_forms import (advection_form, continuity_form, vector_invariant_form, kinetic_energy_form, advection_equation_circulation_form, - diffusion_form) + diffusion_form, linear_continuity_form) from gusto.active_tracers import ActiveTracer, Phases, TracerVariableType from gusto.configuration import TransportEquationType import ufl @@ -601,7 +601,8 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, # transport term from depth equation. Don't include active tracers linearisation_map = lambda t: \ t.get(prognostic) in ["u", "D"] \ - and (any(t.has_label(time_derivative, pressure_gradient, transport))) + and (any(t.has_label(time_derivative, pressure_gradient)) + or (t.get(prognostic) == "D" and t.has_label(transport))) super().__init__(field_names, domain, linearisation_map=linearisation_map, no_normal_flow_bc_ids=no_normal_flow_bc_ids, @@ -609,6 +610,7 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, self.parameters = parameters g = parameters.g + H = parameters.H w, phi = self.tests[0:2] u, D = split(self.X)[0:2] @@ -641,6 +643,11 @@ 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) + # Add linearisation to D_adv + D_adv = linearisation(D_adv, linear_D_adv) adv_form = subject(u_adv + D_adv, self.X) @@ -770,7 +777,8 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, # Default linearisation is time derivatives, pressure gradient, # Coriolis and transport term from depth equation linearisation_map = lambda t: \ - (any(t.has_label(time_derivative, pressure_gradient, coriolis, transport))) + (any(t.has_label(time_derivative, pressure_gradient, coriolis)) + or (t.get(prognostic) == "D" and t.has_label(transport))) super().__init__(domain, parameters, fexpr=fexpr, bexpr=bexpr, @@ -861,6 +869,8 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, w, phi, gamma = self.tests[0:3] u, rho, theta = split(self.X)[0:3] + u_trial = split(self.trials)[0] + _, rho_bar, theta_bar = split(self.X_ref)[0:3] zero_expr = Constant(0.0)*theta exner = exner_pressure(parameters, rho, theta) n = FacetNormal(domain.mesh) @@ -892,9 +902,17 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, # Density transport (conservative form) rho_adv = prognostic(continuity_form(phi, rho, u), "rho") + # Transport term needs special linearisation + if self.linearisation_map(rho_adv.terms[0]): + linear_rho_adv = linear_continuity_form(phi, rho_bar, u_trial) + rho_adv = linearisation(rho_adv, linear_rho_adv) # Potential temperature transport (advective form) theta_adv = prognostic(advection_form(gamma, theta, u), "theta") + # Transport term needs special linearisation + if self.linearisation_map(theta_adv.terms[0]): + linear_theta_adv = linear_advection_form(gamma, theta_bar, u_trial) + theta_adv = linearisation(theta_adv, linear_theta_adv) adv_form = subject(u_adv + rho_adv + theta_adv, self.X) @@ -1187,6 +1205,8 @@ def __init__(self, domain, parameters, Omega=None, w, phi, gamma = self.tests[0:3] u, p, b = split(self.X) + u_trial = split(self.trials)[0] + b_bar = split(self.X_ref)[2] # -------------------------------------------------------------------- # # Time Derivative Terms @@ -1215,6 +1235,9 @@ def __init__(self, domain, parameters, Omega=None, # Buoyancy transport b_adv = prognostic(advection_form(gamma, b, u), "b") + if self.linearisation_map(b_adv.terms[0]): + linear_b_adv = linear_advection_form(gamma, b_bar, u_trial) + b_adv = linearisation(b_adv, linear_b_adv) adv_form = subject(u_adv + b_adv, self.X) From 71dd858a35cee4163be6957ccc13b3baa8a0ada4 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Thu, 13 Jul 2023 08:47:06 +0100 Subject: [PATCH 17/20] tie up final set of loose ends. More clearly separate circulation and vector invariant types --- gusto/common_forms.py | 29 ++++++------- gusto/configuration.py | 2 + gusto/equations.py | 23 ++-------- gusto/timeloop.py | 5 +-- gusto/transport_methods.py | 88 ++++++++++++++++++++++++++++++-------- 5 files changed, 89 insertions(+), 58 deletions(-) diff --git a/gusto/common_forms.py b/gusto/common_forms.py index 56f17d2c5..22f467b06 100644 --- a/gusto/common_forms.py +++ b/gusto/common_forms.py @@ -117,12 +117,7 @@ def vector_invariant_form(domain, test, q, ubar): class:`LabelledForm`: a labelled transport form. """ - if domain.mesh.topological_dimension() == 3: - L = inner(test, cross(curl(q), ubar))*dx - - else: - perp = domain.perp - L = inner(test, div(perp(q))*perp(ubar))*dx + L = advection_equation_circulation_form(domain, test, q, ubar).terms[0].form # Add K.E. term L -= 0.5*div(test)*inner(q, ubar)*dx @@ -145,14 +140,12 @@ def kinetic_energy_form(test, q, ubar): ubar (:class:`ufl.Expr`): the transporting velocity. Returns: - class:`LabelledForm`: a labelled transport form. + class:`ufl.Form`: the kinetic energy form. """ - L = 0.5*div(test)*inner(q, ubar)*dx - - form = transporting_velocity(L, ubar) + L = -0.5*div(test)*inner(q, ubar)*dx - return transport(form, TransportEquationType.vector_invariant) + return L def advection_equation_circulation_form(domain, test, q, ubar): @@ -182,12 +175,16 @@ def advection_equation_circulation_form(domain, test, q, ubar): class:`LabelledForm`: a labelled transport form. """ - form = ( - vector_invariant_form(domain, test, q, ubar) - - kinetic_energy_form(test, q, ubar) - ) + if domain.mesh.topological_dimension() == 3: + L = inner(test, cross(curl(q), ubar))*dx + + else: + perp = domain.perp + L = inner(test, div(perp(q))*perp(ubar))*dx + + form = transporting_velocity(L, ubar) - return form + return transport(form, TransportEquationType.circulation) def diffusion_form(test, q, kappa): diff --git a/gusto/configuration.py b/gusto/configuration.py index 5aeaad7e8..db9260ae1 100644 --- a/gusto/configuration.py +++ b/gusto/configuration.py @@ -50,12 +50,14 @@ class TransportEquationType(Enum): advective: ∂q/∂t + (u.∇)q = 0 \n conservative: ∂q/∂t + ∇.(u*q) = 0 \n vector_invariant: ∂q/∂t + (∇×q)×u + (1/2)∇(q.u) + (1/2)[(∇q).u -(∇u).q)] = 0 + circulation: ∂q/∂t + (∇×q)×u + non-transport terms = 0 """ no_transport = 702 advective = 19 conservative = 291 vector_invariant = 9081 + circulation = 512 class Configuration(object): diff --git a/gusto/equations.py b/gusto/equations.py index 20d20bde4..6b40aee68 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -9,14 +9,15 @@ from gusto.fields import PrescribedFields from gusto.fml.form_manipulation_labelling import Term, all_terms, keep, drop, Label from gusto.labels import (subject, time_derivative, transport, prognostic, - transporting_velocity, replace_subject, linearisation, + replace_subject, linearisation, name, pressure_gradient, coriolis, perp, replace_trial_function, hydrostatic) from gusto.thermodynamics import exner_pressure from gusto.common_forms import (advection_form, continuity_form, vector_invariant_form, kinetic_energy_form, advection_equation_circulation_form, - diffusion_form, linear_continuity_form) + diffusion_form, linear_continuity_form, + linear_advection_form) from gusto.active_tracers import ActiveTracer, Phases, TracerVariableType from gusto.configuration import TransportEquationType import ufl @@ -631,12 +632,6 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, u_adv = prognostic(advection_form(w, u, u), "u") elif u_transport_option == "circulation_form": ke_form = prognostic(kinetic_energy_form(w, u, u), "u") - ke_form = transport.remove(ke_form) - ke_form = ke_form.label_map( - lambda t: t.has_label(transporting_velocity), - lambda t: Term(ufl.replace( - t.form, {t.get(transporting_velocity): u}), t.labels)) - ke_form = transporting_velocity.remove(ke_form) u_adv = prognostic(advection_equation_circulation_form(domain, w, u, u), "u") + ke_form else: raise ValueError("Invalid u_transport_option: %s" % u_transport_option) @@ -890,12 +885,6 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, u_adv = prognostic(advection_form(w, u, u), "u") elif u_transport_option == "circulation_form": ke_form = prognostic(kinetic_energy_form(w, u, u), "u") - ke_form = transport.remove(ke_form) - ke_form = ke_form.label_map( - lambda t: t.has_label(transporting_velocity), - lambda t: Term(ufl.replace( - t.form, {t.get(transporting_velocity): u}), t.labels)) - ke_form = transporting_velocity.remove(ke_form) u_adv = prognostic(advection_equation_circulation_form(domain, w, u, u), "u") + ke_form else: raise ValueError("Invalid u_transport_option: %s" % u_transport_option) @@ -1223,12 +1212,6 @@ def __init__(self, domain, parameters, Omega=None, u_adv = prognostic(advection_form(w, u, u), "u") elif u_transport_option == "circulation_form": ke_form = prognostic(kinetic_energy_form(w, u, u), "u") - ke_form = transport.remove(ke_form) - ke_form = ke_form.label_map( - lambda t: t.has_label(transporting_velocity), - lambda t: Term(ufl.replace( - t.form, {t.get(transporting_velocity): u}), t.labels)) - ke_form = transporting_velocity.remove(ke_form) u_adv = prognostic(advection_equation_circulation_form(domain, w, u, u), "u") + ke_form else: raise ValueError("Invalid u_transport_option: %s" % u_transport_option) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 46321cf43..581ceebcb 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -377,7 +377,7 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, if kwargs: raise ValueError("unexpected kwargs: %s" % list(kwargs.keys())) - self.spatial_methods = [] + self.spatial_methods = spatial_methods if physics_schemes is not None: self.physics_schemes = physics_schemes @@ -397,7 +397,6 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, for method in spatial_methods: if scheme.field_name == method.variable and method.term_label == transport: method_found = True - self.spatial_methods.append(method) assert method_found, f'No transport method found for variable {scheme.field_name}' self.diffusion_schemes = [] @@ -411,7 +410,6 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, for method in spatial_methods: if scheme.field_name == method.variable and method.term_label == diffusion: method_found = True - self.diffusion_methods.append(method) assert method_found, f'No diffusion method found for variable {scheme.field_name}' if auxiliary_equations_and_schemes is not None: @@ -420,6 +418,7 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, self.auxiliary_schemes = [ (eqn.field_name, scheme) for eqn, scheme in auxiliary_equations_and_schemes] + else: auxiliary_equations_and_schemes = [] self.auxiliary_schemes = [] diff --git a/gusto/transport_methods.py b/gusto/transport_methods.py index 74f18b654..e639bf84a 100644 --- a/gusto/transport_methods.py +++ b/gusto/transport_methods.py @@ -52,19 +52,30 @@ def replace_form(self, equation): lambda t: t.has_label(transport) and t.get(prognostic) == self.variable, map_if_true=keep, map_if_false=drop ) - original_term = original_form.terms[0] - # Update transporting velocity - new_transporting_velocity = self.form.terms[0].get(transporting_velocity) - original_term = transporting_velocity.update_value(original_term, new_transporting_velocity) + if len(original_form.terms) == 0: + # This is likely not the appropriate equation so skip + pass - # Create new term - new_term = Term(self.form.form, original_term.labels) + elif len(original_form.terms) == 1: + # Replace form + original_term = original_form.terms[0] - # Replace original term with new term - equation.residual = equation.residual.label_map( - lambda t: t.has_label(transport) and t.get(prognostic) == self.variable, - map_if_true=lambda t: new_term) + # Update transporting velocity + new_transporting_velocity = self.form.terms[0].get(transporting_velocity) + original_term = transporting_velocity.update_value(original_term, new_transporting_velocity) + + # Create new term + new_term = Term(self.form.form, original_term.labels) + + # Replace original term with new term + equation.residual = equation.residual.label_map( + lambda t: t.has_label(transport) and t.get(prognostic) == self.variable, + map_if_true=lambda t: new_term) + + else: + raise RuntimeError('Unable to find single transport term for ' + + f'{self.variable}. {len(original_form.terms)} found') # ---------------------------------------------------------------------------- # @@ -150,6 +161,11 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, form = upwind_continuity_form(self.domain, self.test, self.field, ibp=ibp, outflow=outflow) + elif self.transport_equation_type == TransportEquationType.circulation: + if outflow: + raise NotImplementedError('Outflow not implemented for upwind circulation') + form = upwind_circulation_form(self.domain, self.test, self.field, ibp=ibp) + elif self.transport_equation_type == TransportEquationType.vector_invariant: if outflow: raise NotImplementedError('Outflow not implemented for upwind vector invariant') @@ -358,20 +374,16 @@ def vector_manifold_continuity_form(domain, test, q, ibp=IntegrateByParts.ONCE, return transport(form) -def upwind_vector_invariant_form(domain, test, q, ibp=IntegrateByParts.ONCE): +def upwind_circulation_form(domain, test, q, ibp=IntegrateByParts.ONCE): u""" - The form corresponding to the DG upwind vector invariant transport operator. + The form corresponding to the DG upwind vector circulation operator. The self-transporting transport operator for a vector-valued field u can be written as circulation and kinetic energy terms: (u.∇)u = (∇×u)×u + (1/2)∇u^2 - When the transporting field u and transported field q are similar, we write - this as: - (u.∇)q = (∇×q)×u + (1/2)∇(u.q) - - This form discretises this final equation, using an upwind discretisation - when integrating by parts. + This form discretises the first term in this equation, (∇×u)×u, using an + upwind discretisation when integrating by parts. Args: domain (:class:`Domain`): the model's domain object, containing the @@ -437,8 +449,46 @@ def upwind_vector_invariant_form(domain, test, q, ibp=IntegrateByParts.ONCE): perp(ubar))*perp(q), n)*dS_ ) - L -= 0.5*div(test)*inner(q, ubar)*dx + form = transporting_velocity(L, ubar) + + return transport(form, TransportEquationType.circulation) + + +def upwind_vector_invariant_form(domain, test, q, ibp=IntegrateByParts.ONCE): + u""" + The form corresponding to the DG upwind vector invariant transport operator. + + The self-transporting transport operator for a vector-valued field u can be + written as circulation and kinetic energy terms: + (u.∇)u = (∇×u)×u + (1/2)∇u^2 + + When the transporting field u and transported field q are similar, we write + this as: + (u.∇)q = (∇×q)×u + (1/2)∇(u.q) + + This form discretises this final equation, using an upwind discretisation + when integrating by parts. + + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + test (:class:`TestFunction`): the test function. + q (:class:`ufl.Expr`): the variable to be transported. + ibp (:class:`IntegrateByParts`, optional): an enumerator representing + the number of times to integrate by parts. Defaults to + `IntegrateByParts.ONCE`. + + Raises: + NotImplementedError: the specified integration by parts is not 'once'. + + Returns: + class:`LabelledForm`: a labelled transport form. + """ + + circulation_form = upwind_circulation_form(domain, test, q, ibp=ibp) + ubar = circulation_form.terms[0].get(transporting_velocity) + L = circulation_form.terms[0].form - 0.5*div(test)*inner(q, ubar)*dx form = transporting_velocity(L, ubar) return transport(form, TransportEquationType.vector_invariant) From e035286decea8959a063a2919675e92a9da76a3c Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Thu, 13 Jul 2023 10:07:31 +0100 Subject: [PATCH 18/20] fix final two failing examples --- examples/compressible/straka_bubble.py | 10 +++++----- examples/compressible/unsaturated_bubble.py | 1 + 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/examples/compressible/straka_bubble.py b/examples/compressible/straka_bubble.py index c696733ed..cbe32800d 100644 --- a/examples/compressible/straka_bubble.py +++ b/examples/compressible/straka_bubble.py @@ -44,9 +44,9 @@ # Equation parameters = CompressibleParameters() - diffusion_options = [ - ("u", DiffusionParameters(kappa=75., mu=10./delta)), - ("theta", DiffusionParameters(kappa=75., mu=10./delta))] + u_diffusion_opts = DiffusionParameters(kappa=75., mu=10./delta) + theta_diffusion_opts = DiffusionParameters(kappa=75., mu=10./delta) + diffusion_options = [("u", u_diffusion_opts), ("theta", theta_diffusion_opts)] eqns = CompressibleEulerEquations(domain, parameters, diffusion_options=diffusion_options) @@ -75,8 +75,8 @@ # Diffusion schemes diffusion_schemes = [BackwardEuler(domain, "u"), BackwardEuler(domain, "theta")] - diffusion_methods = [InteriorPenaltyDiffusion(eqns, "u", diffusion_options[0]), - InteriorPenaltyDiffusion(eqns, "theta", diffusion_options[1])] + diffusion_methods = [InteriorPenaltyDiffusion(eqns, "u", u_diffusion_opts), + InteriorPenaltyDiffusion(eqns, "theta", theta_diffusion_opts)] # Time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, diff --git a/examples/compressible/unsaturated_bubble.py b/examples/compressible/unsaturated_bubble.py index 8b338a556..d6d8afd4c 100644 --- a/examples/compressible/unsaturated_bubble.py +++ b/examples/compressible/unsaturated_bubble.py @@ -95,6 +95,7 @@ # Time stepper stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + transport_methods, linear_solver=linear_solver, physics_schemes=physics_schemes) From 8dda4aa5017c3ed3f299b8d03080fce2e9b06913 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Mon, 24 Jul 2023 10:49:00 +0100 Subject: [PATCH 19/20] improve comments --- gusto/common_forms.py | 8 +------- gusto/spatial_methods.py | 13 ++++++------- gusto/timeloop.py | 8 ++++---- gusto/transport_methods.py | 7 ++++--- 4 files changed, 15 insertions(+), 21 deletions(-) diff --git a/gusto/common_forms.py b/gusto/common_forms.py index 22f467b06..0a3790147 100644 --- a/gusto/common_forms.py +++ b/gusto/common_forms.py @@ -110,9 +110,6 @@ def vector_invariant_form(domain, test, q, ubar): q (:class:`ufl.Expr`): the variable to be transported. ubar (:class:`ufl.Expr`): the transporting velocity. - Raises: - NotImplementedError: the specified integration by parts is not 'once'. - Returns: class:`LabelledForm`: a labelled transport form. """ @@ -161,16 +158,13 @@ def advection_equation_circulation_form(domain, test, q, ubar): (u.∇)q = (∇×q)×u + (1/2)∇(u.q) The form returned by this function corresponds to the (∇×q)×u circulation - term. An an upwind discretisation is used when integrating by parts. + term. Args: test (:class:`TestFunction`): the test function. q (:class:`ufl.Expr`): the variable to be transported. ubar (:class:`ufl.Expr`): the transporting velocity. - Raises: - NotImplementedError: the specified integration by parts is not 'once'. - Returns: class:`LabelledForm`: a labelled transport form. """ diff --git a/gusto/spatial_methods.py b/gusto/spatial_methods.py index 432b0c588..bb050f3c5 100644 --- a/gusto/spatial_methods.py +++ b/gusto/spatial_methods.py @@ -20,7 +20,7 @@ def __init__(self, equation, variable, term_label): Args: equation (:class:`PrognosticEquation`): the equation, which includes the original type of this term. - variable (str): name of the variable to set the transport scheme for + variable (str): name of the variable to set the method for term_label (:class:`Label`): the label specifying which type of term to be discretised. """ @@ -38,8 +38,7 @@ def __init__(self, equation, variable, term_label): self.field = equation.X self.test = equation.test - # Find the original transport term to be used, which we use to extract - # information about the transport equation type + # Find the original term to be used self.original_form = equation.residual.label_map( lambda t: t.has_label(term_label) and t.get(prognostic) == variable, map_if_true=keep, map_if_false=drop) @@ -50,13 +49,13 @@ def __init__(self, equation, variable, term_label): def replace_form(self, equation): """ - Replaces the form for the transport term in the equation with the - form for the transport discretisation. + Replaces the form for the term in the equation with the form for the + specific discretisation. Args: equation (:class:`PrognosticEquation`): the equation or scheme whose - transport term should be replaced with the transport term of - this discretisation. + term should be replaced with the specific term of this + discretisation. """ # Replace original term with new term diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 581ceebcb..d0235aa92 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -93,11 +93,11 @@ def setup_equation(self, equation): residual = equation.residual.label_map( lambda t: t.has_label(term_label), map_if_false=drop ) - active_variables = [t.get(prognostic) for t in residual.terms] - active_methods = list(filter(lambda t: t.term_label == term_label, + variables = [t.get(prognostic) for t in residual.terms] + methods = list(filter(lambda t: t.term_label == term_label, self.spatial_methods)) - method_variables = [method.variable for method in active_methods] - for variable in active_variables: + method_variables = [method.variable for method in methods] + for variable in variables: if variable not in method_variables: message = f'Variable {variable} has a {term_label.label} ' \ + 'but no method for this has been specified. Using ' \ diff --git a/gusto/transport_methods.py b/gusto/transport_methods.py index e639bf84a..5e21c4964 100644 --- a/gusto/transport_methods.py +++ b/gusto/transport_methods.py @@ -5,7 +5,7 @@ from firedrake import (dx, dS, dS_v, dS_h, ds_t, ds_b, ds_v, dot, inner, outer, jump, grad, div, FacetNormal, Function, sign, avg, cross, curl) -from gusto.configuration import IntegrateByParts, TransportEquationType +from gusto.configuration import IntegrateByParts, TransportEquationType, logger from gusto.fml import Term, keep, drop from gusto.labels import prognostic, transport, transporting_velocity, ibp_label from gusto.spatial_methods import SpatialMethod @@ -55,7 +55,8 @@ def replace_form(self, equation): if len(original_form.terms) == 0: # This is likely not the appropriate equation so skip - pass + logger.warning(f'No transport term found for {self.variable} in ' + + 'this equation. Skipping.') elif len(original_form.terms) == 1: # Replace form @@ -74,7 +75,7 @@ def replace_form(self, equation): map_if_true=lambda t: new_term) else: - raise RuntimeError('Unable to find single transport term for ' + raise RuntimeError('Found multiple transport terms for ' + f'{self.variable}. {len(original_form.terms)} found') From 3c2863a94dd30b9cd17c5532e58c91edb14b68ee Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Mon, 24 Jul 2023 13:27:16 +0100 Subject: [PATCH 20/20] fix lint --- gusto/timeloop.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index d0235aa92..02d11b74a 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -95,7 +95,7 @@ def setup_equation(self, equation): ) variables = [t.get(prognostic) for t in residual.terms] methods = list(filter(lambda t: t.term_label == term_label, - self.spatial_methods)) + self.spatial_methods)) method_variables = [method.variable for method in methods] for variable in variables: if variable not in method_variables: