diff --git a/gusto/core/labels.py b/gusto/core/labels.py index b49bbefad..995f03b70 100644 --- a/gusto/core/labels.py +++ b/gusto/core/labels.py @@ -92,6 +92,8 @@ def __call__(self, target, value=None): # ---------------------------------------------------------------------------- # implicit = Label("implicit") explicit = Label("explicit") +horizontal_transport = Label("horizontal_transport") +vertical_transport = Label("vertical_transport") source_label = Label("source_label") transporting_velocity = Label("transporting_velocity", validator=lambda value: type(value) in [Function, ufl.tensors.ListTensor, ufl.indexed.Indexed]) prognostic = Label("prognostic", validator=lambda value: type(value) == str) diff --git a/gusto/equations/common_forms.py b/gusto/equations/common_forms.py index 8bb606511..a96c84b8d 100644 --- a/gusto/equations/common_forms.py +++ b/gusto/equations/common_forms.py @@ -7,14 +7,15 @@ from firedrake.fml import subject, drop from gusto.core.configuration import TransportEquationType from gusto.core.labels import (transport, transporting_velocity, diffusion, - prognostic, linearisation) + prognostic, linearisation, horizontal_transport, + vertical_transport) __all__ = ["advection_form", "advection_form_1d", "continuity_form", "continuity_form_1d", "vector_invariant_form", "kinetic_energy_form", "advection_equation_circulation_form", "diffusion_form", "diffusion_form_1d", "linear_advection_form", "linear_continuity_form", - "split_continuity_form", "tracer_conservative_form"] + "split_continuity_form", "tracer_conservative_form", "split_hv_advective_form"] def advection_form(test, q, ubar): @@ -346,3 +347,126 @@ def tracer_conservative_form(test, q, rho, ubar): form = transporting_velocity(L, ubar) return transport(form, TransportEquationType.tracer_conservative) + + +def split_advection_form(test, q, ubar, ubar_full): + u""" + The form corresponding to the advective transport operator in either horzontal + or vertical directions (dependent on ubar). + + This describes either u_h.(∇)q or w dq/dz, 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 in a subset of dimensions. + ubar_full (:class:`ufl.Expr`): the transporting velocity in all dimensions. + + Returns: + class:`LabelledForm`: a labelled transport form. + """ + + L = inner(test, dot(ubar, grad(q)))*dx + form = transporting_velocity(L, ubar_full) + + return transport(form, TransportEquationType.advective) + + +def split_linear_advection_form(test, qbar, ubar, ubar_full): + """ + The form corresponding to the linearised advective transport operator in + either horzontal or vertical directions (dependent on ubar). + + Args: + test (:class:`TestFunction`): the test function. + qbar (:class:`ufl.Expr`): the variable to be transported. + ubar (:class:`ufl.Expr`): the transporting velocity in a subset of dimensions. + ubar_full (:class:`ufl.Expr`): the transporting velocity in all dimensions. + + Returns: + :class:`LabelledForm`: a labelled transport form. + """ + + L = test*dot(ubar, grad(qbar))*dx + form = transporting_velocity(L, ubar_full) + + return transport(form, TransportEquationType.advective) + + +def split_hv_advective_form(equation, field_name): + u""" + Splits advective term into horizontal and vertical terms. + This describes splitting u.∇(q) terms into u_h.(∇)q and w dq/dz, + for transporting velocity u and transported q. + Args: + equation (:class:`PrognosticEquation`): the model's equation. + Returns: + :class:`PrognosticEquation`: the model's equation. + """ + k = equation.domain.k # vertical unit vector + for t in equation.residual: + if (t.get(transport) == TransportEquationType.advective and t.get(prognostic) == field_name): + # Get fields and test functions + subj = t.get(subject) + + # u is either a prognostic or prescribed field + if (hasattr(equation, "field_names") + and 'u' in equation.field_names): + idx = equation.field_names.index(field_name) + W = equation.function_space + test = TestFunctions(W)[idx] + q = split(subj)[idx] + u_idx = equation.field_names.index('u') + uadv = split(equation.X)[u_idx] + elif 'u' in equation.prescribed_fields._field_names: + uadv = equation.prescribed_fields('u') + q = subj + W = equation.function_space + test = TestFunction(W) + else: + raise ValueError('Cannot get velocity field') + + # Create new advective and divergence terms + u_vertical = k*inner(uadv, k) + u_horizontal = uadv - u_vertical + vertical_adv_term = prognostic( + vertical_transport( + split_advection_form(test, q, u_vertical, uadv) + ), + field_name + ) + horizontal_adv_term = prognostic( + horizontal_transport( + split_advection_form(test, q, u_horizontal, uadv) + ), + field_name + ) + + # Add linearisations of new terms if required + if (t.has_label(linearisation)): + u_trial = TrialFunctions(W)[u_idx] + u_trial_vert = k*inner(u_trial, k) + u_trial_horiz = u_trial - u_trial_vert + qbar = split(equation.X_ref)[idx] + # Add linearisations + linear_hori_term = horizontal_transport( + split_linear_advection_form(test, qbar, u_trial_horiz, u_trial) + ) + adv_horiz_term = linearisation(horizontal_adv_term, linear_hori_term) + + linear_vert_term = vertical_transport( + split_linear_advection_form(test, qbar, u_trial_vert, u_trial) + ) + adv_vert_term = linearisation(vertical_adv_term, linear_vert_term) + else: + adv_vert_term = vertical_adv_term + adv_horiz_term = horizontal_adv_term + # Drop old term + equation.residual = equation.residual.label_map( + lambda t: t.get(transport) == TransportEquationType.advective and t.get(prognostic) == field_name, + map_if_true=drop) + + # Add new terms onto residual + equation.residual += subject(adv_horiz_term, subj) + subject(adv_vert_term, subj) + + return equation diff --git a/gusto/spatial_methods/spatial_methods.py b/gusto/spatial_methods/spatial_methods.py index 6ad214dcd..2fff98728 100644 --- a/gusto/spatial_methods/spatial_methods.py +++ b/gusto/spatial_methods/spatial_methods.py @@ -4,7 +4,7 @@ """ from firedrake import split -from firedrake.fml import Term, keep, drop +from firedrake.fml import Term, keep, drop, all_terms from gusto.core.labels import prognostic __all__ = ['SpatialMethod'] @@ -15,19 +15,19 @@ class SpatialMethod(object): The base object for describing a spatial discretisation of some term. """ - def __init__(self, equation, variable, term_label): + def __init__(self, equation, variable, *term_labels): """ Args: equation (:class:`PrognosticEquation`): the equation, which includes the original type of this term. 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. + term_labels (:class:`Label`): One or more labels specifying which type(s) + of terms should be discretized. """ self.equation = equation self.variable = variable self.domain = self.equation.domain - self.term_label = term_label + self.term_labels = list(term_labels) if hasattr(equation, "field_names"): # Equation with multiple prognostic variables @@ -38,14 +38,37 @@ def __init__(self, equation, variable, term_label): self.field = equation.X self.test = equation.test - # 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) - - num_terms = len(self.original_form.terms) - assert num_terms == 1, f'Unable to find {term_label.label} term ' \ - + f'for {variable}. {num_terms} found' + if (len(self.term_labels) == 1): + # Most cases only have one term to be replaced + self.term_label = self.term_labels[0] + self.original_form = equation.residual.label_map( + lambda t: t.has_label(self.term_label) and t.get(prognostic) == variable, + map_if_true=keep, + map_if_false=drop + ) + # Check that the original form has the correct number of terms + num_terms = len(self.original_form.terms) + assert num_terms == 1, f'Unable to find {self.term_label.label} term ' \ + + f'for {variable}. {num_terms} found' + else: + # Multiple terms to be replaced. Find the original terms to be used + self.term_label = self.term_labels[0] + self.original_form = equation.residual.label_map( + all_terms, + map_if_true=drop + ) + for term in self.term_labels: + original_form = equation.residual.label_map( + lambda t: t.has_label(term) and t.get(prognostic) == variable, + map_if_true=keep, + map_if_false=drop + ) + # Check that the original form has the correct number of terms + num_terms = len(original_form.terms) + assert num_terms == 1, f'Unable to find {term.label} term ' \ + + f'for {variable}. {num_terms} found' + # Add the terms form to original forms + self.original_form += original_form def replace_form(self, equation): """ diff --git a/gusto/spatial_methods/transport_methods.py b/gusto/spatial_methods/transport_methods.py index 8d194796f..edfd79a5c 100644 --- a/gusto/spatial_methods/transport_methods.py +++ b/gusto/spatial_methods/transport_methods.py @@ -10,12 +10,12 @@ from gusto.core.configuration import IntegrateByParts, TransportEquationType from gusto.core.labels import ( prognostic, transport, transporting_velocity, ibp_label, mass_weighted, - all_but_last + all_but_last, horizontal_transport, vertical_transport ) from gusto.core.logging import logger from gusto.spatial_methods.spatial_methods import SpatialMethod -__all__ = ["DefaultTransport", "DGUpwind"] +__all__ = ["DefaultTransport", "DGUpwind", "SplitDGUpwind"] # ---------------------------------------------------------------------------- # @@ -26,16 +26,21 @@ class TransportMethod(SpatialMethod): The base object for describing a transport scheme. """ - def __init__(self, equation, variable): + def __init__(self, equation, variable, *term_labels): """ Args: equation (:class:`PrognosticEquation`): the equation, which includes a transport term. variable (str): name of the variable to set the transport scheme for + term_labels (:class:`Label`): One or more labels specifying which type(s) + of terms should be discretized. """ # Inherited init method extracts original term to be replaced - super().__init__(equation, variable, transport) + if not term_labels: + super().__init__(equation, variable, transport) + else: + super().__init__(equation, variable, *term_labels) # If this is term has a mass_weighted label, then we need to # use the tracer_conservative version of the transport method. @@ -111,8 +116,45 @@ def replace_form(self, equation): map_if_true=lambda t: new_term) else: - raise RuntimeError('Found multiple transport terms for ' - + f'{self.variable}. {len(original_form.terms)} found') + horizontal_form = equation.residual.label_map( + lambda t: t.has_label(transport) and t.has_label(horizontal_transport) and t.get(prognostic) == self.variable, + map_if_true=keep, map_if_false=drop + ) + vertical_form = equation.residual.label_map( + lambda t: t.has_label(transport) and t.has_label(vertical_transport) and t.get(prognostic) == self.variable, + map_if_true=keep, map_if_false=drop + ) + if len(horizontal_form.terms) == 1 and len(vertical_form.terms) == 1: + + # Replace forms + horizontal_term = horizontal_form.terms[0] + vertical_term = vertical_form.terms[0] + + # Update transporting velocity + new_horizontal_transporting_velocity = self.form_h.terms[0].get(transporting_velocity) + new_vertical_transporting_velocity = self.form_v.terms[0].get(transporting_velocity) + horizontal_term = transporting_velocity.update_value(horizontal_term, new_horizontal_transporting_velocity) + vertical_term = transporting_velocity.update_value(vertical_term, new_vertical_transporting_velocity) + + # Create new terms + new_horizontal_term = Term(self.form_h.form, horizontal_term.labels) + new_vertical_term = Term(self.form_v.form, vertical_term.labels) + + # Check if this is a conservative transport + if horizontal_term.has_label(mass_weighted) or vertical_term.has_label(mass_weighted): + raise NotImplementedError('Mass weighted transport terms not yet supported for multiple terms') + + # Replace original terms with new terms + equation.residual = equation.residual.label_map( + lambda t: t.has_label(transport) and t.has_label(horizontal_transport) and t.get(prognostic) == self.variable, + map_if_true=lambda _: new_horizontal_term) + + equation.residual = equation.residual.label_map( + lambda t: t.has_label(transport) and t.has_label(vertical_transport) and t.get(prognostic) == self.variable, + map_if_true=lambda _: new_vertical_term) + else: + raise RuntimeError('Found multiple transport terms for the same ' + 'variable in the equation where there should only be one') # ---------------------------------------------------------------------------- # @@ -277,9 +319,130 @@ def __init__(self, equation, variable, ibp=IntegrateByParts.ONCE, self.form = form +class SplitDGUpwind(TransportMethod): + """ + The Discontinuous Galerkin Upwind transport scheme applied separately in the + horizontal and vertical directions. + 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, horizontal_transport, vertical_transport) + self.ibp = ibp + self.vector_manifold_correction = vector_manifold_correction + self.outflow = outflow + + # -------------------------------------------------------------------- # + # Determine appropriate form to use + # -------------------------------------------------------------------- # + # first check for 1d mesh and scalar velocity space + if equation.domain.mesh.topological_dimension() == 1 and len(equation.domain.spaces("HDiv").shape) == 0: + assert not vector_manifold_correction + raise ValueError('You cannot do horizontal and vertical splitting in 1D') + else: + if self.transport_equation_type == TransportEquationType.advective: + + form_h, form_v = split_upwind_advection_form(self.domain, self.test, + self.field, + ibp=ibp, outflow=outflow) + + else: + raise NotImplementedError('Split hv Upwind transport scheme has not been ' + + 'implemented for this transport equation type') + self.form_v = form_v + self.form_h = form_h + + # ---------------------------------------------------------------------------- # # Forms for DG Upwind transport # ---------------------------------------------------------------------------- # +def split_upwind_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): + u""" + The forms corresponding to the DG upwind advective transport operator in + the horizontal and vertical directions. + This discretises u_h.(∇_h)q and w dq/dz, 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") + k = domain.k + quad = domain.max_quad_degree + dS_ = (dS_v(degree=quad) + dS_h(degree=quad)) if Vu.extruded else dS + ubar = Function(Vu) + ubar_v = k*inner(ubar, k) + ubar_h = ubar - ubar_v + + if ibp == IntegrateByParts.ONCE: + L_h = -inner(div(outer(test, ubar_h)), q)*dx(degree=quad) + L_v = -inner(div(outer(test, ubar_v)), q)*dx(degree=quad) + else: + L_h = inner(outer(test, ubar_h), grad(q))*dx(degree=quad) + L_v = inner(outer(test, ubar_v), grad(q))*dx(degree=quad) + + if ibp != IntegrateByParts.NEVER: + n = FacetNormal(domain.mesh) + un_h = 0.5*(dot(ubar_h, n) + abs(dot(ubar_h, n))) + + L_h += dot(jump(test), (un_h('+')*q('+') - un_h('-')*q('-')))*dS_ + + un_v = 0.5*(dot(ubar_v, n) + abs(dot(ubar_v, n))) + + L_v += dot(jump(test), (un_v('+')*q('+') - un_v('-')*q('-')))*dS_ + + if ibp == IntegrateByParts.TWICE: + L_h -= (inner(test('+'), dot(ubar_h('+'), n('+')) * q('+')) + + inner(test('-'), dot(ubar_h('-'), n('-')) * q('-'))) * dS_ + + L_v -= (inner(test('+'), dot(ubar_v('+'), n('+')) * q('+')) + + inner(test('-'), dot(ubar_v('-'), n('-')) * q('-'))) * dS_ + + if outflow: + n = FacetNormal(domain.mesh) + un_h = 0.5*(dot(ubar_h, n) + abs(dot(ubar_h, n))) + L_h += test*un_h*q*(ds_v + ds_t + ds_b) + + un_v = 0.5*(dot(ubar_v, n) + abs(dot(ubar_v, n))) + L_v += test*un_v*q*(ds_v + ds_t + ds_b) + + form_h = transporting_velocity(L_h, ubar) + form_v = transporting_velocity(L_v, ubar) + labelled_form_h = ibp_label(transport(form_h, TransportEquationType.advective), ibp) + labelled_form_v = ibp_label(transport(form_v, TransportEquationType.advective), ibp) + return labelled_form_h, labelled_form_v + + def upwind_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): u""" The form corresponding to the DG upwind advective transport operator. diff --git a/gusto/time_discretisation/sdc.py b/gusto/time_discretisation/sdc.py index 3e634f17a..7c419cb21 100644 --- a/gusto/time_discretisation/sdc.py +++ b/gusto/time_discretisation/sdc.py @@ -338,7 +338,7 @@ def res(self, m): r_source_k = self.residual.label_map( lambda t: t.has_label(source_label), - map_if_true=replace_subject(self.source_Ukp1[i+1], old_idx=self.idx), + map_if_true=replace_subject(self.source_Uk[i+1], old_idx=self.idx), map_if_false=drop) r_source_k = r_source_k.label_map( all_terms, diff --git a/integration-tests/transport/test_split_dg_transport.py b/integration-tests/transport/test_split_dg_transport.py new file mode 100644 index 000000000..5fe1030df --- /dev/null +++ b/integration-tests/transport/test_split_dg_transport.py @@ -0,0 +1,64 @@ +""" +Tests the split horizontal and vertical DG upwind transport scheme for +advective form transport equation. This tests that the +field is transported to the correct place. +""" + +from firedrake import norm, VectorFunctionSpace, as_vector +from gusto import * + + +def run(timestepper, tmax, f_end): + timestepper.run(0, tmax) + return norm(timestepper.fields("f") - f_end) / norm(f_end) + + +def test_split_dg_transport_scalar(tmpdir, tracer_setup): + setup = tracer_setup(tmpdir, "slice") + domain = setup.domain + V = domain.spaces("DG") + + eqn = AdvectionEquation(domain, V, "f") + eqn = split_hv_advective_form(eqn, "f") + + transport_method = SplitDGUpwind(eqn, "f") + transport_scheme = SSPRK3(domain) + + time_varying_velocity = False + timestepper = PrescribedTransport( + eqn, transport_scheme, setup.io, time_varying_velocity, transport_method + ) + + # Initial conditions + timestepper.fields("f").interpolate(setup.f_init) + timestepper.fields("u").project(setup.uexpr) + + error = run(timestepper, setup.tmax, setup.f_end) + assert error < setup.tol, \ + 'The transport error is greater than the permitted tolerance' + + +def test_split_dg_transport_vector(tmpdir, tracer_setup): + setup = tracer_setup(tmpdir, "slice") + domain = setup.domain + gdim = domain.mesh.geometric_dimension() + f_init = as_vector([setup.f_init]*gdim) + V = VectorFunctionSpace(domain.mesh, "DG", 1) + eqn = AdvectionEquation(domain, V, "f") + eqn = split_hv_advective_form(eqn, "f") + + transport_scheme = SSPRK3(domain) + transport_method = SplitDGUpwind(eqn, "f") + + time_varying_velocity = False + timestepper = PrescribedTransport( + eqn, transport_scheme, setup.io, time_varying_velocity, transport_method + ) + + # Initial conditions + timestepper.fields("f").interpolate(f_init) + timestepper.fields("u").project(setup.uexpr) + f_end = as_vector([setup.f_end]*gdim) + error = run(timestepper, setup.tmax, f_end) + assert error < setup.tol, \ + 'The transport error is greater than the permitted tolerance'