From 40e8bafd18d7e6e6850220332e028eb0bca93ad2 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Mon, 24 Jul 2023 21:31:06 +0100 Subject: [PATCH 1/4] begin reorganising function spaces --- gusto/function_spaces.py | 72 +++++++++++++++++++++++++--------------- 1 file changed, 46 insertions(+), 26 deletions(-) diff --git a/gusto/function_spaces.py b/gusto/function_spaces.py index 4af0f814b..a50ac112f 100644 --- a/gusto/function_spaces.py +++ b/gusto/function_spaces.py @@ -3,16 +3,21 @@ used by the model. """ -from firedrake import (HDiv, FunctionSpace, FiniteElement, TensorProductElement, - interval) +from firedrake import (HCurl, HDiv, FunctionSpace, FiniteElement, + TensorProductElement, interval) -# TODO: there is danger here for confusion about degree, particularly for the CG -# spaces -- does a "CG" space with degree = 1 mean the "CG" space in the de Rham -# complex of degree 1 ("CG3"), or "CG1"? -# TODO: would it be better to separate creation of specific named spaces from -# the creation of the de Rham complex spaces? -# TODO: how do we create HCurl spaces if we want them? +__all__ = ["Spaces", "check_degree_args"] +# HDiv spaces are keys, HCurl spaces are values +hdiv_hcurl_dict = {'RT': 'RTE', + 'RTF': 'RTE' + 'BDM': 'BDME', + 'BDMF': 'BDME', + 'RTCF': 'RTCE', + 'CG': 'CG'} + +# Can't just reverse the other dictionary as values are not necessarily unique +hcurl_hdiv_dict = {'BDME': 'RT'} class Spaces(object): """Object to create and hold the model's finite element spaces.""" @@ -63,6 +68,12 @@ def __call__(self, name, family=None, degree=None, :class:`FunctionSpace`: the desired function space. """ + implemented_families = ["DG", "CG", "RT", "RTF", "RTE", "RTCF", "RTCE", + "BDM", "BDMF", "BDME"] + if family not in [None]+implemented_families: + raise NotImplementedError(f'family {family} either not recognised ' + + 'or implemented in Gusto') + if hasattr(self, name) and (V is None or not overwrite_space): # We have requested a space that should already have been created if V is not None: @@ -89,8 +100,10 @@ def __call__(self, name, family=None, degree=None, vertical_degree = degree if vertical_degree is None else vertical_degree # Loop through name and family combinations - if name == "HDiv" and family in ["BDM", "RT", "CG", "RTCF"]: + if name == "HDiv" and family in ["BDM", "RT", "CG", "RTCF", "RTF", "BDMF"]: value = self.build_hdiv_space(family, horizontal_degree, vertical_degree) + elif name == "HCurl" and family in ["BDM", "RT", "CG", "RTCE", "RTE", "BDME"]: + value = self.build_hcurl_space(family, horizontal_degree, vertical_degree) elif name == "theta": value = self.build_theta_space(horizontal_degree, vertical_degree) elif family == "DG": @@ -108,9 +121,9 @@ def build_compatible_spaces(self, family, horizontal_degree, Builds the sequence of compatible finite element spaces for the mesh. If the mesh is not extruded, this builds and returns the spaces: \n - (HDiv, DG). \n + (H1, HCurl, HDiv, DG). \n If the mesh is extruded, this builds and returns the spaces: \n - (HDiv, DG, theta). \n + (H1, HCurl, HDiv, DG, theta). \n The 'theta' space corresponds to the vertical component of the velocity. Args: @@ -125,19 +138,26 @@ def build_compatible_spaces(self, family, horizontal_degree, tuple: the created compatible :class:`FunctionSpace` objects. """ if self.extruded_mesh and not self._initialised_base_spaces: + # Base spaces need building, while horizontal and vertical degrees + # need specifying separately self.build_base_spaces(family, horizontal_degree, vertical_degree) + Vcg = self.build_h1_space(cg_degree(family, horizontal_degree), + cg_degree(family, vertical_degree), name='H1') + setattr(self, "H1", Vcg) + Vcurl = self.build_hcurl_space(family, horizontal_degree, vertical_degree) + setattr(self, "HCurl", Vcurl) Vu = self.build_hdiv_space(family, horizontal_degree, vertical_degree) setattr(self, "HDiv", Vu) - Vdg = self.build_dg_space(horizontal_degree, vertical_degree, name='DG') - setattr(self, "DG", Vdg) + Vdg = self.build_l2_space(horizontal_degree, vertical_degree, name='L2') + setattr(self, "L2", Vdg) Vth = self.build_theta_space(horizontal_degree, vertical_degree) setattr(self, "theta", Vth) return Vu, Vdg, Vth else: Vu = self.build_hdiv_space(family, horizontal_degree+1) setattr(self, "HDiv", Vu) - Vdg = self.build_dg_space(horizontal_degree, vertical_degree, name='DG') - setattr(self, "DG", Vdg) + Vdg = self.build_dg_space(horizontal_degree, vertical_degree, name='L2') + setattr(self, "L2", Vdg) return Vu, Vdg def build_base_spaces(self, family, horizontal_degree, vertical_degree): @@ -192,30 +212,30 @@ def build_hdiv_space(self, family, horizontal_degree, vertical_degree=None): V_elt = FiniteElement(family, cell, horizontal_degree) return FunctionSpace(self.mesh, V_elt, name='HDiv') - def build_dg_space(self, horizontal_degree, vertical_degree=None, variant=None, name='DG'): + def build_l2_space(self, horizontal_degree, vertical_degree=None, variant=None, name='L2'): """ - Builds and returns the DG :class:`FunctionSpace`. + Builds and returns the discontinuous L2 :class:`FunctionSpace`. Args: horizontal_degree (int): the polynomial degree of the horizontal - part of the DG space. + part of the L2 space. vertical_degree (int, optional): the polynomial degree of the - vertical part of the DG space. Defaults to None. Must be + vertical part of the L2 space. Defaults to None. Must be specified if the mesh is extruded. variant (str, optional): the variant of the underlying :class:`FiniteElement` to use. Defaults to None, which will call the default variant. name (str, optional): name to assign to the function space. Default - is "DG". + is "L2". Returns: - :class:`FunctionSpace`: the DG space. + :class:`FunctionSpace`: the L2 space. """ assert not hasattr(self, name), f'There already exists a function space with name {name}' if self.extruded_mesh: if vertical_degree is None: - raise ValueError('vertical_degree must be specified to create DG space on an extruded mesh') + raise ValueError('vertical_degree must be specified to create L2 space on an extruded mesh') if not self._initialised_base_spaces or self.T1.degree() != vertical_degree or self.T1.variant() != variant: cell = self.mesh._base_mesh.ufl_cell().cellname() S2 = FiniteElement("DG", cell, horizontal_degree, variant=variant) @@ -240,9 +260,9 @@ def build_theta_space(self, horizontal_degree, vertical_degree): Args: horizontal_degree (int): the polynomial degree of the horizontal - part of the DG space from the de Rham complex. + part of the L2 space from the de Rham complex. vertical_degree (int): the polynomial degree of the vertical part of - the DG space from the de Rham complex. + the L2 space from the de Rham complex. Raises: AssertionError: the mesh is not extruded. @@ -258,7 +278,7 @@ def build_theta_space(self, horizontal_degree, vertical_degree): V_elt = TensorProductElement(self.S2, self.T0) return FunctionSpace(self.mesh, V_elt, name='theta') - def build_cg_space(self, horizontal_degree, vertical_degree=None, name='CG'): + def build_h1_space(self, horizontal_degree, vertical_degree=None, name='H1'): """ Builds the continuous scalar space at the top of the de Rham complex. @@ -269,7 +289,7 @@ def build_cg_space(self, horizontal_degree, vertical_degree=None, name='CG'): vertical part of the the CG space. Defaults to None. Must be specified if the mesh is extruded. name (str, optional): name to assign to the function space. Default - is "CG". + is "H1". Returns: :class:`FunctionSpace`: the continuous space. From f253414a21adef71ac0db8ad0c0d39f1960f5473 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Wed, 26 Jul 2023 13:23:06 +0100 Subject: [PATCH 2/4] further work on function spaces --- gusto/function_spaces.py | 133 ++++++++++++++++++++++++++++++--------- 1 file changed, 104 insertions(+), 29 deletions(-) diff --git a/gusto/function_spaces.py b/gusto/function_spaces.py index a50ac112f..0a15a1558 100644 --- a/gusto/function_spaces.py +++ b/gusto/function_spaces.py @@ -10,14 +10,37 @@ # HDiv spaces are keys, HCurl spaces are values hdiv_hcurl_dict = {'RT': 'RTE', - 'RTF': 'RTE' + 'RTF': 'RTE', 'BDM': 'BDME', 'BDMF': 'BDME', 'RTCF': 'RTCE', 'CG': 'CG'} +# HCurl spaces are keys, HDiv spaces are values # Can't just reverse the other dictionary as values are not necessarily unique -hcurl_hdiv_dict = {'BDME': 'RT'} +hcurl_hdiv_dict = {'RT': 'RTF', + 'RTE': 'RTF', + 'BDM': 'BDMF', + 'BDME': 'BDMF', + 'RTCE': 'RTCF', + 'CG': 'CG'} + +# Degree to use for H1 space for a particular family +def h1_degree(family, l2_degree): + """ + Return the degree of the H1 space, for a particular de Rham complex. + + Args: + family (str): the family of the HDiv or HCurl elements. + l2_degree (int): the degree of the L2 space at the bottom of the complex + + Returns: + int: the degree of the H1 space at the top of the complex. + """ + if family in ['CG', 'RT', 'RTE', 'RTF', 'RTCE', 'RTCF']: + return l2_degree + 1 + elif family in ['BDM', 'BDME', 'BDMF']: + return l2_degree + 2 class Spaces(object): """Object to create and hold the model's finite element spaces.""" @@ -90,7 +113,7 @@ def __call__(self, name, family=None, degree=None, elif name == "DG1_equispaced": # Special case as no degree arguments need providing - value = self.build_dg_space(1, 1, variant='equispaced', name='DG1_equispaced') + value = self.build_l2_space(1, 1, variant='equispaced', name='DG1_equispaced') else: check_degree_args('Spaces', self.mesh, degree, horizontal_degree, vertical_degree) @@ -101,15 +124,17 @@ def __call__(self, name, family=None, degree=None, # Loop through name and family combinations if name == "HDiv" and family in ["BDM", "RT", "CG", "RTCF", "RTF", "BDMF"]: - value = self.build_hdiv_space(family, horizontal_degree, vertical_degree) + hdiv_family = hcurl_hdiv_dict[family] + value = self.build_hdiv_space(hdiv_family, horizontal_degree, vertical_degree) elif name == "HCurl" and family in ["BDM", "RT", "CG", "RTCE", "RTE", "BDME"]: - value = self.build_hcurl_space(family, horizontal_degree, vertical_degree) + hcurl_family = hdiv_hcurl_dict[family] + value = self.build_hcurl_space(hcurl_family, horizontal_degree, vertical_degree) elif name == "theta": value = self.build_theta_space(horizontal_degree, vertical_degree) elif family == "DG": - value = self.build_dg_space(horizontal_degree, vertical_degree, name=name) + value = self.build_l2_space(horizontal_degree, vertical_degree, name=name) elif family == "CG": - value = self.build_cg_space(horizontal_degree, vertical_degree, name=name) + value = self.build_h1_space(horizontal_degree, vertical_degree, name=name) else: raise ValueError(f'There is no space corresponding to {name}') setattr(self, name, value) @@ -129,9 +154,9 @@ def build_compatible_spaces(self, family, horizontal_degree, Args: family (str): the family of the horizontal part of the HDiv space. horizontal_degree (int): the polynomial degree of the horizontal - part of the DG space. + part of the L2 space. vertical_degree (int, optional): the polynomial degree of the - vertical part of the DG space. Defaults to None. Must be + vertical part of the L2 space. Defaults to None. Must be specified if the mesh is extruded. Returns: @@ -139,26 +164,47 @@ def build_compatible_spaces(self, family, horizontal_degree, """ if self.extruded_mesh and not self._initialised_base_spaces: # Base spaces need building, while horizontal and vertical degrees - # need specifying separately + # need specifying separately. Vtheta needs returning. self.build_base_spaces(family, horizontal_degree, vertical_degree) - Vcg = self.build_h1_space(cg_degree(family, horizontal_degree), - cg_degree(family, vertical_degree), name='H1') + Vcg = self.build_h1_space(h1_degree(family, horizontal_degree), + h1_degree(family, vertical_degree), name='H1') setattr(self, "H1", Vcg) - Vcurl = self.build_hcurl_space(family, horizontal_degree, vertical_degree) + hcurl_family = hdiv_hcurl_dict[family] + Vcurl = self.build_hcurl_space(hcurl_family, horizontal_degree, vertical_degree) setattr(self, "HCurl", Vcurl) - Vu = self.build_hdiv_space(family, horizontal_degree, vertical_degree) + hdiv_family = hcurl_hdiv_dict[family] + Vu = self.build_hdiv_space(hdiv_family, horizontal_degree, vertical_degree) setattr(self, "HDiv", Vu) Vdg = self.build_l2_space(horizontal_degree, vertical_degree, name='L2') setattr(self, "L2", Vdg) Vth = self.build_theta_space(horizontal_degree, vertical_degree) setattr(self, "theta", Vth) - return Vu, Vdg, Vth - else: + return Vcg, Vcurl, Vu, Vdg, Vth + elif self.mesh.topological_dimension() > 1: + # 2D: two de Rham complexes (hcurl or hdiv) with 3 spaces + # 3D: one de Rham complexes with 4 spaces + # either way, build all spaces + Vcg = self.build_h1_space(h1_degree(family, horizontal_degree), name='H1') + setattr(self, "H1", Vcg) + hcurl_family = hdiv_hcurl_dict[family] + Vcurl = self.build_hcurl_space(hcurl_family, horizontal_degree+1) + setattr(self, "HCurl", Vcurl) + hdiv_family = hcurl_hdiv_dict[family] Vu = self.build_hdiv_space(family, horizontal_degree+1) setattr(self, "HDiv", Vu) - Vdg = self.build_dg_space(horizontal_degree, vertical_degree, name='L2') + Vdg = self.build_l2_space(horizontal_degree, vertical_degree, name='L2') + setattr(self, "L2", Vdg) + return Vcg, Vcurl, Vu, Vdg + else: + # 1D domain, de Rham complex has 2 spaces + # CG, hdiv and hcurl spaces should be the same + Vcg = self.build_h1_space(horizontal_degree+1, name='H1') + setattr(self, "H1", Vcg) + setattr(self, "HCurl", Vcurl) + setattr(self, "HDiv", Vu) + Vdg = self.build_l2_space(horizontal_degree, name='L2') setattr(self, "L2", Vdg) - return Vu, Vdg + return Vcg, Vdg def build_base_spaces(self, family, horizontal_degree, vertical_degree): """ @@ -167,22 +213,51 @@ def build_base_spaces(self, family, horizontal_degree, vertical_degree): Args: family (str): the family of the horizontal part of the HDiv space. horizontal_degree (int): the polynomial degree of the horizontal - part of the DG space. + part of the L2 space. vertical_degree (int): the polynomial degree of the vertical part of - the DG space. + the L2 space. """ cell = self.mesh._base_mesh.ufl_cell().cellname() # horizontal base spaces - self.S1 = FiniteElement(family, cell, horizontal_degree+1) - self.S2 = FiniteElement("DG", cell, horizontal_degree) + self.base_elt_hori_hdiv = FiniteElement(hdiv_family, cell, horizontal_degree+1) + self.base_elt_hori_hcurl = FiniteElement(hcurl_family, cell, horizontal_degree+1) + self.base_elt_hori_dg = FiniteElement("DG", cell, horizontal_degree) # vertical base spaces - self.T0 = FiniteElement("CG", interval, vertical_degree+1) - self.T1 = FiniteElement("DG", interval, vertical_degree) + self.base_elt_vert_cg = FiniteElement("CG", interval, vertical_degree+1) + self.base_elt_vert_dg = FiniteElement("DG", interval, vertical_degree) self._initialised_base_spaces = True + def build_hcurl_space(self, family, horizontal_degree, vertical_degree=None): + """ + Builds and returns the HCurl :class:`FunctionSpace`. + + Args: + family (str): the family of the horizontal part of the HCurl space. + horizontal_degree (int): the polynomial degree of the horizontal + part of the L2 space from the de Rham complex. + vertical_degree (int, optional): the polynomial degree of the + vertical part of the L2 space from the de Rham complex. + Defaults to None. Must be specified if the mesh is extruded. + + Returns: + :class:`FunctionSpace`: the HCurl space. + """ + if self.extruded_mesh: + if not self._initialised_base_spaces: + if vertical_degree is None: + raise ValueError('vertical_degree must be specified to create HCurl space on an extruded mesh') + self.build_base_spaces(family, horizontal_degree, vertical_degree) + Vh_elt = HCurl(TensorProductElement(self.S1, self.T1)) + Vv_elt = HCurl(TensorProductElement(self.S2, self.T0)) + V_elt = Vh_elt + Vv_elt + else: + cell = self.mesh.ufl_cell().cellname() + V_elt = FiniteElement(family, cell, horizontal_degree) + return FunctionSpace(self.mesh, V_elt, name='HCurl') + def build_hdiv_space(self, family, horizontal_degree, vertical_degree=None): """ Builds and returns the HDiv :class:`FunctionSpace`. @@ -190,9 +265,9 @@ def build_hdiv_space(self, family, horizontal_degree, vertical_degree=None): Args: family (str): the family of the horizontal part of the HDiv space. horizontal_degree (int): the polynomial degree of the horizontal - part of the DG space from the de Rham complex. + part of the L2 space from the de Rham complex. vertical_degree (int, optional): the polynomial degree of the - vertical part of the the DG space from the de Rham complex. + vertical part of the L2 space from the de Rham complex. Defaults to None. Must be specified if the mesh is extruded. Returns: @@ -284,9 +359,9 @@ def build_h1_space(self, horizontal_degree, vertical_degree=None, name='H1'): Args: horizontal_degree (int): the polynomial degree of the horizontal - part of the CG space. + part of the H1 space. vertical_degree (int, optional): the polynomial degree of the - vertical part of the the CG space. Defaults to None. Must be + vertical part of the H1 space. Defaults to None. Must be specified if the mesh is extruded. name (str, optional): name to assign to the function space. Default is "H1". @@ -298,7 +373,7 @@ def build_h1_space(self, horizontal_degree, vertical_degree=None, name='H1'): if self.extruded_mesh: if vertical_degree is None: - raise ValueError('vertical_degree must be specified to create CG space on an extruded mesh') + raise ValueError('vertical_degree must be specified to create H1 space on an extruded mesh') cell = self.mesh._base_mesh.ufl_cell().cellname() CG_hori = FiniteElement("CG", cell, horizontal_degree) CG_vert = FiniteElement("CG", interval, vertical_degree) From 6146c05edfe7bb4a98aba0a040f8246c8ed17251 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Thu, 27 Jul 2023 07:44:09 +0100 Subject: [PATCH 3/4] make whole de rham complex --- gusto/diagnostics.py | 10 +- gusto/domain.py | 4 +- gusto/equations.py | 188 ++++++++++++++++++----------- gusto/function_spaces.py | 141 +++++++++++++++++----- unit-tests/test_function_spaces.py | 148 +++++++++++++++++++++++ 5 files changed, 385 insertions(+), 106 deletions(-) create mode 100644 unit-tests/test_function_spaces.py diff --git a/gusto/diagnostics.py b/gusto/diagnostics.py index 05a43670e..34762418c 100644 --- a/gusto/diagnostics.py +++ b/gusto/diagnostics.py @@ -1,6 +1,6 @@ """Common diagnostic fields.""" -from firedrake import op2, assemble, dot, dx, FunctionSpace, Function, sqrt, \ +from firedrake import op2, assemble, dot, dx, Function, sqrt, \ TestFunction, TrialFunction, Constant, grad, inner, curl, \ LinearVariationalProblem, LinearVariationalSolver, FacetNormal, \ ds_b, ds_v, ds_t, dS_v, div, avg, jump, \ @@ -1368,13 +1368,7 @@ def setup(self, domain, state_fields, vorticity_type=None): vorticity_types = ["relative", "absolute", "potential"] if vorticity_type not in vorticity_types: raise ValueError(f"vorticity type must be one of {vorticity_types}, not {vorticity_type}") - try: - space = domain.spaces("CG") - except ValueError: - dgspace = domain.spaces("DG") - # TODO: should this be degree + 1? - cg_degree = dgspace.ufl_element().degree() + 2 - space = FunctionSpace(domain.mesh, "CG", cg_degree, name=f"CG{cg_degree}") + space = domain.spaces("H1") u = state_fields("u") if vorticity_type in ["absolute", "potential"]: diff --git a/gusto/domain.py b/gusto/domain.py index 85f2308d4..3e92c11a9 100644 --- a/gusto/domain.py +++ b/gusto/domain.py @@ -72,8 +72,8 @@ def __init__(self, mesh, dt, family, degree=None, self.mesh = mesh self.family = family self.spaces = Spaces(mesh) - # Build and store compatible spaces - self.compatible_spaces = [space for space in self.spaces.build_compatible_spaces(self.family, self.horizontal_degree, self.vertical_degree)] + self.spaces.build_compatible_spaces(self.family, self.horizontal_degree, + self.vertical_degree) # -------------------------------------------------------------------- # # Determine some useful aspects of domain diff --git a/gusto/equations.py b/gusto/equations.py index 6b40aee68..aeeb9ab4d 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -200,14 +200,18 @@ class PrognosticEquationSet(PrognosticEquation, metaclass=ABCMeta): contains common routines for these equation sets. """ - def __init__(self, field_names, domain, linearisation_map=None, - no_normal_flow_bc_ids=None, active_tracers=None): + def __init__(self, field_names, domain, space_names, + linearisation_map=None, no_normal_flow_bc_ids=None, + active_tracers=None): """ Args: field_names (list): a list of strings for names of the prognostic variables for the equation set. domain (:class:`Domain`): the model's domain object, containing the mesh and the compatible function spaces. + space_names (dict): a dictionary of strings for names of the + function spaces to use for the spatial discretisation. The keys + are the names of the prognostic variables. linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. Defaults to None. no_normal_flow_bc_ids (list, optional): a list of IDs of domain @@ -219,16 +223,13 @@ def __init__(self, field_names, domain, linearisation_map=None, """ self.field_names = field_names + self.space_names = space_names self.active_tracers = active_tracers self.linearisation_map = lambda t: False if linearisation_map is None else linearisation_map(t) # Build finite element spaces - # TODO: this implies order of spaces matches order of variables. - # It also assumes that if one additional field is required then it - # should live on the DG space. - self.spaces = [space for space in domain.compatible_spaces] - if len(self.field_names) - len(self.spaces) == 1: - self.spaces.append(domain.spaces("DG")) + self.spaces = [domain.spaces(space_name) for space_name in + [self.space_names[field_name] for field_name in self.field_names]] # Add active tracers to the list of prognostics if active_tracers is None: @@ -413,17 +414,26 @@ def add_tracers_to_prognostics(self, domain, active_tracers): self.field_names.append(tracer.name) else: raise ValueError(f'There is already a field named {tracer.name}') + + # Add name of space to self.space_names, but check for conflict + # with the tracer's name + if tracer.name in self.space_names: + assert self.space_names[tracer.name] == tracer.space, \ + 'space_name dict provided to equation has space ' \ + + f'{self.space_names[tracer.name]} for tracer ' \ + + f'{tracer.name} which conflicts with the space ' \ + + f'{tracer.space} specified in the ActiveTracer object' + else: + self.space_names[tracer.name] = tracer.space self.spaces.append(domain.spaces(tracer.space)) else: raise TypeError(f'Tracers must be ActiveTracer objects, not {type(tracer)}') - def generate_tracer_transport_terms(self, domain, active_tracers): + def generate_tracer_transport_terms(self, active_tracers): """ Adds the transport forms for the active tracers to the equation set. Args: - domain (:class:`Domain`): the model's domain object, containing the - mesh and the compatible function spaces. active_tracers (list): A list of :class:`ActiveTracer` objects that encode the metadata for the active tracers. @@ -499,6 +509,7 @@ def __init__(self, domain, function_space, field_name, Vu=None, """ self.field_names = [field_name] + self.space_names = {field_name: function_space.name} self.active_tracers = active_tracers self.terms_to_linearise = {} @@ -533,7 +544,7 @@ def __init__(self, domain, function_space, field_name, Vu=None, # Add transport of tracers if len(active_tracers) > 0: - self.residual += self.generate_tracer_transport_terms(domain, active_tracers) + self.residual += self.generate_tracer_transport_terms(active_tracers) # ============================================================================ # # Specified Equation Sets @@ -550,7 +561,7 @@ class ShallowWaterEquations(PrognosticEquationSet): """ def __init__(self, domain, parameters, fexpr=None, bexpr=None, - linearisation_map='default', + space_names=None, linearisation_map='default', u_transport_option='vector_invariant_form', no_normal_flow_bc_ids=None, active_tracers=None, thermal=False): @@ -564,6 +575,11 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, parameter. Defaults to None. bexpr (:class:`ufl.Expr`, optional): an expression for the bottom surface of the fluid. Defaults to None. + space_names (dict, optional): a dictionary of strings for names of + the function spaces to use for the spatial discretisation. The + keys are the names of the prognostic variables. Defaults to None + in which case the spaces are taken from the de Rham complex. Any + buoyancy variable is taken by default to lie in the L2 space. linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', @@ -589,22 +605,27 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, """ self.thermal = thermal - field_names = ["u", "D"] + field_names = ['u', 'D'] + + if space_names is None: + space_names = {'u': 'HDiv', 'D': 'L2'} if active_tracers is None: active_tracers = [] if self.thermal: - field_names.append("b") + field_names.append('b') + if 'b' not in space_names.keys(): + space_names['b'] = 'L2' if linearisation_map == 'default': # Default linearisation is time derivatives, pressure gradient and # transport term from depth equation. Don't include active tracers linearisation_map = lambda t: \ - t.get(prognostic) in ["u", "D"] \ + 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))) - super().__init__(field_names, domain, + or (t.get(prognostic) in ['D', 'b'] and t.has_label(transport))) + super().__init__(field_names, domain, space_names, linearisation_map=linearisation_map, no_normal_flow_bc_ids=no_normal_flow_bc_ids, active_tracers=active_tracers) @@ -627,17 +648,17 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=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") + 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") + 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") - u_adv = prognostic(advection_equation_circulation_form(domain, w, u, u), "u") + ke_form + ke_form = prognostic(kinetic_energy_form(w, u, u), 'u') + 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) # Depth transport term - D_adv = prognostic(continuity_form(phi, D, u), "D") + D_adv = prognostic(continuity_form(phi, D, u), 'D') # Transport term needs special linearisation if self.linearisation_map(D_adv.terms[0]): linear_D_adv = linear_continuity_form(phi, H, u_trial) @@ -648,12 +669,12 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, # Add transport of tracers if len(active_tracers) > 0: - adv_form += self.generate_tracer_transport_terms(domain, active_tracers) + adv_form += self.generate_tracer_transport_terms(active_tracers) # Add transport of buoyancy, if thermal shallow water equations if self.thermal: gamma = self.tests[2] b = split(self.X)[2] - b_adv = prognostic(advection_form(gamma, b, u), "b") + b_adv = prognostic(advection_form(gamma, b, u), 'b') adv_form += subject(b_adv, self.X) # -------------------------------------------------------------------- # @@ -664,7 +685,7 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, residual = (mass_form + adv_form) else: pressure_gradient_form = pressure_gradient( - subject(prognostic(-g*div(w)*D*dx, "u"), self.X)) + subject(prognostic(-g*div(w)*D*dx, 'u'), self.X)) residual = (mass_form + adv_form + pressure_gradient_form) @@ -676,17 +697,17 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, # the equation, and initialised when the equation is if fexpr is not None: - V = FunctionSpace(domain.mesh, "CG", 1) - f = self.prescribed_fields("coriolis", V).interpolate(fexpr) + V = FunctionSpace(domain.mesh, 'CG', 1) + f = self.prescribed_fields('coriolis', V).interpolate(fexpr) coriolis_form = coriolis(subject( - prognostic(f*inner(domain.perp(u), w)*dx, "u"), self.X)) + prognostic(f*inner(domain.perp(u), w)*dx, 'u'), self.X)) if not domain.on_sphere: coriolis_form = perp(coriolis_form, domain.perp) # Add linearisation if self.linearisation_map(coriolis_form.terms[0]): linear_coriolis = perp( coriolis( - subject(prognostic(f*inner(domain.perp(u_trial), w)*dx, "u"), self.X) + subject(prognostic(f*inner(domain.perp(u_trial), w)*dx, 'u'), self.X) ), domain.perp) if not domain.on_sphere: linear_coriolis = perp(linear_coriolis, domain.perp) @@ -694,16 +715,16 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, residual += coriolis_form if bexpr is not None: - topography = self.prescribed_fields("topography", domain.spaces("DG")).interpolate(bexpr) + topography = self.prescribed_fields('topography', domain.spaces('DG')).interpolate(bexpr) if self.thermal: n = FacetNormal(domain.mesh) topography_form = subject(prognostic (-topography*div(b*w)*dx + jump(b*w, n)*avg(topography)*dS, - "u"), self.X) + 'u'), self.X) else: topography_form = subject(prognostic - (-g*div(w)*topography*dx, "u"), + (-g*div(w)*topography*dx, 'u'), self.X) residual += topography_form @@ -714,7 +735,7 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, - 0.5*b*div(D*w)*dx + jump(b*w, n)*avg(D)*dS + 0.5*jump(D*w, n)*avg(b)*dS, - "u"), self.X) + 'u'), self.X) residual += source_form # -------------------------------------------------------------------- # @@ -737,9 +758,10 @@ class LinearShallowWaterEquations(ShallowWaterEquations): """ def __init__(self, domain, parameters, fexpr=None, bexpr=None, - linearisation_map='default', + space_names=None, linearisation_map='default', u_transport_option="vector_invariant_form", - no_normal_flow_bc_ids=None, active_tracers=None): + no_normal_flow_bc_ids=None, active_tracers=None, + thermal=False): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -750,6 +772,11 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, parameter. Defaults to None. bexpr (:class:`ufl.Expr`, optional): an expression for the bottom surface of the fluid. Defaults to None. + space_names (dict, optional): a dictionary of strings for names of + the function spaces to use for the spatial discretisation. The + keys are the names of the prognostic variables. Defaults to None + in which case the spaces are taken from the de Rham complex. Any + buoyancy variable is taken by default to lie in the L2 space. linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', @@ -766,6 +793,9 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=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. + thermal (flag, optional): specifies whether the equations have a + thermal or buoyancy variable that feeds back on the momentum. + Defaults to False. """ if linearisation_map == 'default': @@ -773,14 +803,14 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, # 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))) + or (t.get(prognostic) in ['D', 'b'] and t.has_label(transport))) super().__init__(domain, parameters, - fexpr=fexpr, bexpr=bexpr, + fexpr=fexpr, bexpr=bexpr, space_names=space_names, linearisation_map=linearisation_map, u_transport_option=u_transport_option, no_normal_flow_bc_ids=no_normal_flow_bc_ids, - active_tracers=active_tracers) + active_tracers=active_tracers, thermal=thermal) # Use the underlying routine to do a first linearisation of the equations self.linearise_equation_set() @@ -800,7 +830,8 @@ class CompressibleEulerEquations(PrognosticEquationSet): """ def __init__(self, domain, parameters, Omega=None, sponge=None, - extra_terms=None, linearisation_map='default', + extra_terms=None, space_names=None, + linearisation_map='default', u_transport_option="vector_invariant_form", diffusion_options=None, no_normal_flow_bc_ids=None, @@ -817,6 +848,10 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, layer. Defaults to None. extra_terms (:class:`ufl.Expr`, optional): any extra terms to be included in the equation set. Defaults to None. + space_names (dict, optional): a dictionary of strings for names of + the function spaces to use for the spatial discretisation. The + keys are the names of the prognostic variables. Defaults to None + in which case the spaces are taken from the de Rham complex. linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', @@ -843,6 +878,9 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, field_names = ['u', 'rho', 'theta'] + if space_names is None: + space_names = {'u': 'HDiv', 'rho': 'L2', 'theta': 'theta'} + if active_tracers is None: active_tracers = [] @@ -853,7 +891,7 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, t.get(prognostic) in ['u', 'rho', 'theta'] \ and (t.has_label(time_derivative) or (t.get(prognostic) != 'u' and t.has_label(transport))) - super().__init__(field_names, domain, + super().__init__(field_names, domain, space_names, linearisation_map=linearisation_map, no_normal_flow_bc_ids=no_normal_flow_bc_ids, active_tracers=active_tracers) @@ -880,24 +918,24 @@ 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") + 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") + 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") - u_adv = prognostic(advection_equation_circulation_form(domain, w, u, u), "u") + ke_form + ke_form = prognostic(kinetic_energy_form(w, u, u), 'u') + 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) # Density transport (conservative form) - rho_adv = prognostic(continuity_form(phi, rho, u), "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(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") + 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) @@ -907,7 +945,7 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, # Add transport of tracers if len(active_tracers) > 0: - adv_form += self.generate_tracer_transport_terms(domain, active_tracers) + adv_form += self.generate_tracer_transport_terms(active_tracers) # -------------------------------------------------------------------- # # Pressure Gradient Term @@ -924,12 +962,12 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, pressure_gradient_form = name(subject(prognostic( cp*(-div(theta_v*w)*exner*dx - + jump(theta_v*w, n)*avg(exner)*dS_v), "u"), self.X), "pressure_gradient") + + jump(theta_v*w, n)*avg(exner)*dS_v), 'u'), self.X), "pressure_gradient") # -------------------------------------------------------------------- # # Gravitational Term # -------------------------------------------------------------------- # - gravity_form = subject(prognostic(Term(g*inner(domain.k, w)*dx), "u"), self.X) + gravity_form = subject(prognostic(Term(g*inner(domain.k, w)*dx), 'u'), self.X) residual = (mass_form + adv_form + pressure_gradient_form + gravity_form) @@ -965,7 +1003,7 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, residual += subject(prognostic( gamma * theta * div(u) - * (R_m / c_vml - (R_d * c_pml) / (cp * c_vml))*dx, "theta"), self.X) + * (R_m / c_vml - (R_d * c_pml) / (cp * c_vml))*dx, 'theta'), self.X) # -------------------------------------------------------------------- # # Extra Terms (Coriolis, Sponge, Diffusion and others) @@ -981,7 +1019,8 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, z = x[len(x)-1] H = sponge.H zc = sponge.z_level - assert float(zc) < float(H), "you have set the sponge level above the height of your domain" + assert float(zc) < float(H), \ + "The sponge level is set above the height the your domain" mubar = sponge.mubar muexpr = conditional(z <= zc, 0.0, @@ -989,7 +1028,7 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, self.mu = self.prescribed_fields("sponge", W_DG).interpolate(muexpr) residual += name(subject(prognostic( - self.mu*inner(w, domain.k)*inner(u, domain.k)*dx, "u"), self.X), "sponge") + self.mu*inner(w, domain.k)*inner(u, domain.k)*dx, 'u'), self.X), "sponge") if diffusion_options is not None: for field, diffusion in diffusion_options: @@ -1033,7 +1072,7 @@ class HydrostaticCompressibleEulerEquations(CompressibleEulerEquations): """ def __init__(self, domain, parameters, Omega=None, sponge=None, - extra_terms=None, linearisation_map='default', + extra_terms=None, space_names=None, linearisation_map='default', u_transport_option="vector_invariant_form", diffusion_options=None, no_normal_flow_bc_ids=None, @@ -1050,6 +1089,10 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, layer. Defaults to None. extra_terms (:class:`ufl.Expr`, optional): any extra terms to be included in the equation set. Defaults to None. + space_names (dict, optional): a dictionary of strings for names of + the function spaces to use for the spatial discretisation. The + keys are the names of the prognostic variables. Defaults to None + in which case the spaces are taken from the de Rham complex. linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', @@ -1075,7 +1118,7 @@ def __init__(self, domain, parameters, Omega=None, sponge=None, """ super().__init__(domain, parameters, Omega=Omega, sponge=sponge, - extra_terms=extra_terms, + extra_terms=extra_terms, space_names=space_names, linearisation_map=linearisation_map, u_transport_option=u_transport_option, diffusion_options=diffusion_options, @@ -1114,7 +1157,9 @@ def hydrostatic_projection(self, t): """ # TODO: make this more general, i.e. should work on the sphere - assert not self.domain.on_sphere, "the hydrostatic projection is not yet implemented for spherical geometry" + if self.domain.on_sphere: + raise NotImplementedError("The hydrostatic projection is not yet " + + "implemented for spherical geometry") k = Constant((*self.domain.k, 0, 0)) X = t.get(subject) @@ -1136,7 +1181,7 @@ class IncompressibleBoussinesqEquations(PrognosticEquationSet): """ def __init__(self, domain, parameters, Omega=None, - linearisation_map='default', + space_names=None, linearisation_map='default', u_transport_option="vector_invariant_form", no_normal_flow_bc_ids=None, active_tracers=None): @@ -1148,6 +1193,10 @@ def __init__(self, domain, parameters, Omega=None, the model's physical parameters. Omega (:class:`ufl.Expr`, optional): an expression for the planet's rotation vector. Defaults to None. + space_names (dict, optional): a dictionary of strings for names of + the function spaces to use for the spatial discretisation. The + keys are the names of the prognostic variables. Defaults to None + in which case the spaces are taken from the de Rham complex. linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. If None is specified then no terms are linearised. Defaults to the string 'default', @@ -1171,6 +1220,9 @@ def __init__(self, domain, parameters, Omega=None, field_names = ['u', 'p', 'b'] + if space_names is None: + space_names = {'u': 'HDiv', 'p': 'L2', 'b': 'theta'} + if active_tracers is not None: raise NotImplementedError('Tracers not implemented for Boussinesq equations') @@ -1185,7 +1237,7 @@ def __init__(self, domain, parameters, Omega=None, and (t.has_label(time_derivative) or (t.get(prognostic) not in ['u', 'p'] and t.has_label(transport))) - super().__init__(field_names, domain, + super().__init__(field_names, domain, space_names, linearisation_map=linearisation_map, no_normal_flow_bc_ids=no_normal_flow_bc_ids, active_tracers=active_tracers) @@ -1207,17 +1259,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") + 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") + 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") - u_adv = prognostic(advection_equation_circulation_form(domain, w, u, u), "u") + ke_form + ke_form = prognostic(kinetic_energy_form(w, u, u), 'u') + 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) # Buoyancy transport - b_adv = prognostic(advection_form(gamma, b, u), "b") + 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) @@ -1226,17 +1278,17 @@ def __init__(self, domain, parameters, Omega=None, # Add transport of tracers if len(active_tracers) > 0: - adv_form += self.generate_tracer_transport_terms(domain, active_tracers) + adv_form += self.generate_tracer_transport_terms(active_tracers) # -------------------------------------------------------------------- # # Pressure Gradient Term # -------------------------------------------------------------------- # - pressure_gradient_form = subject(prognostic(-div(w)*p*dx, "u"), self.X) + pressure_gradient_form = subject(prognostic(-div(w)*p*dx, 'u'), self.X) # -------------------------------------------------------------------- # # Gravitational Term # -------------------------------------------------------------------- # - gravity_form = subject(prognostic(-b*inner(w, domain.k)*dx, "u"), self.X) + gravity_form = subject(prognostic(-b*inner(w, domain.k)*dx, 'u'), self.X) # -------------------------------------------------------------------- # # Divergence Term @@ -1245,7 +1297,7 @@ def __init__(self, domain, parameters, Omega=None, # The p features here so that the div(u) evaluated in the "forcing" step # replaces the whole pressure field, rather than merely providing an # increment to it. - divergence_form = name(subject(prognostic(phi*(p-div(u))*dx, "p"), self.X), + divergence_form = name(subject(prognostic(phi*(p-div(u))*dx, 'p'), self.X), "incompressibility") residual = (mass_form + adv_form + divergence_form @@ -1257,7 +1309,7 @@ def __init__(self, domain, parameters, Omega=None, if Omega is not None: # TODO: add linearisation and label for this residual += subject(prognostic( - inner(w, cross(2*Omega, u))*dx, "u"), self.X) + inner(w, cross(2*Omega, u))*dx, 'u'), self.X) # -------------------------------------------------------------------- # # Linearise equations diff --git a/gusto/function_spaces.py b/gusto/function_spaces.py index 0a15a1558..8f23f2a0d 100644 --- a/gusto/function_spaces.py +++ b/gusto/function_spaces.py @@ -3,6 +3,7 @@ used by the model. """ +from gusto import logger from firedrake import (HCurl, HDiv, FunctionSpace, FiniteElement, TensorProductElement, interval) @@ -10,20 +11,29 @@ # HDiv spaces are keys, HCurl spaces are values hdiv_hcurl_dict = {'RT': 'RTE', + 'RTE': 'RTE', 'RTF': 'RTE', 'BDM': 'BDME', + 'BDME': 'BDME', 'BDMF': 'BDME', 'RTCF': 'RTCE', - 'CG': 'CG'} + 'RTCE': 'RTCE', + 'CG': 'DG', + 'BDFM': None} # HCurl spaces are keys, HDiv spaces are values # Can't just reverse the other dictionary as values are not necessarily unique hcurl_hdiv_dict = {'RT': 'RTF', 'RTE': 'RTF', + 'RTF': 'RTF', 'BDM': 'BDMF', 'BDME': 'BDMF', + 'BDMF': 'BDMF', 'RTCE': 'RTCF', - 'CG': 'CG'} + 'RTCF': 'RTCF', + 'CG': 'CG', + 'BDFM': 'BDFM'} + # Degree to use for H1 space for a particular family def h1_degree(family, l2_degree): @@ -41,6 +51,11 @@ def h1_degree(family, l2_degree): return l2_degree + 1 elif family in ['BDM', 'BDME', 'BDMF']: return l2_degree + 2 + elif family == 'BDFM': + return l2_degree + 1 + else: + raise ValueError(f'family {family} not recognised') + class Spaces(object): """Object to create and hold the model's finite element spaces.""" @@ -92,7 +107,7 @@ def __call__(self, name, family=None, degree=None, """ implemented_families = ["DG", "CG", "RT", "RTF", "RTE", "RTCF", "RTCE", - "BDM", "BDMF", "BDME"] + "BDM", "BDMF", "BDME", "BDFM"] if family not in [None]+implemented_families: raise NotImplementedError(f'family {family} either not recognised ' + 'or implemented in Gusto') @@ -146,9 +161,9 @@ def build_compatible_spaces(self, family, horizontal_degree, Builds the sequence of compatible finite element spaces for the mesh. If the mesh is not extruded, this builds and returns the spaces: \n - (H1, HCurl, HDiv, DG). \n + (H1, HCurl, HDiv, L2). \n If the mesh is extruded, this builds and returns the spaces: \n - (H1, HCurl, HDiv, DG, theta). \n + (H1, HCurl, HDiv, L2, theta). \n The 'theta' space corresponds to the vertical component of the velocity. Args: @@ -162,7 +177,10 @@ def build_compatible_spaces(self, family, horizontal_degree, Returns: tuple: the created compatible :class:`FunctionSpace` objects. """ - if self.extruded_mesh and not self._initialised_base_spaces: + if self._initialised_base_spaces: + pass + + elif self.extruded_mesh and not self._initialised_base_spaces: # Base spaces need building, while horizontal and vertical degrees # need specifying separately. Vtheta needs returning. self.build_base_spaces(family, horizontal_degree, vertical_degree) @@ -177,9 +195,12 @@ def build_compatible_spaces(self, family, horizontal_degree, setattr(self, "HDiv", Vu) Vdg = self.build_l2_space(horizontal_degree, vertical_degree, name='L2') setattr(self, "L2", Vdg) + setattr(self, "DG", Vdg) # Register this as "L2" and "DG" Vth = self.build_theta_space(horizontal_degree, vertical_degree) setattr(self, "theta", Vth) + return Vcg, Vcurl, Vu, Vdg, Vth + elif self.mesh.topological_dimension() > 1: # 2D: two de Rham complexes (hcurl or hdiv) with 3 spaces # 3D: one de Rham complexes with 4 spaces @@ -194,16 +215,21 @@ def build_compatible_spaces(self, family, horizontal_degree, setattr(self, "HDiv", Vu) Vdg = self.build_l2_space(horizontal_degree, vertical_degree, name='L2') setattr(self, "L2", Vdg) + setattr(self, "DG", Vdg) # Register this as "L2" and "DG" + return Vcg, Vcurl, Vu, Vdg + else: # 1D domain, de Rham complex has 2 spaces # CG, hdiv and hcurl spaces should be the same Vcg = self.build_h1_space(horizontal_degree+1, name='H1') setattr(self, "H1", Vcg) - setattr(self, "HCurl", Vcurl) - setattr(self, "HDiv", Vu) + setattr(self, "HCurl", None) + setattr(self, "HDiv", Vcg) Vdg = self.build_l2_space(horizontal_degree, name='L2') setattr(self, "L2", Vdg) + setattr(self, "DG", Vdg) # Register this as "L2" and "DG" + return Vcg, Vdg def build_base_spaces(self, family, horizontal_degree, vertical_degree): @@ -211,18 +237,30 @@ def build_base_spaces(self, family, horizontal_degree, vertical_degree): Builds the :class:`FiniteElement` objects for the base mesh. Args: - family (str): the family of the horizontal part of the HDiv space. + family (str): the family of the horizontal part of either the HDiv + or HCurl space. horizontal_degree (int): the polynomial degree of the horizontal part of the L2 space. vertical_degree (int): the polynomial degree of the vertical part of the L2 space. """ + + if family == 'BDFM': + # Need a special implementation of base spaces here as it does not + # fit the same pattern as other spaces + self.build_bdfm_base_spaces(horizontal_degree, vertical_degree) + return + cell = self.mesh._base_mesh.ufl_cell().cellname() + hdiv_family = hcurl_hdiv_dict[family] + hcurl_family = hdiv_hcurl_dict[family] + # horizontal base spaces self.base_elt_hori_hdiv = FiniteElement(hdiv_family, cell, horizontal_degree+1) self.base_elt_hori_hcurl = FiniteElement(hcurl_family, cell, horizontal_degree+1) self.base_elt_hori_dg = FiniteElement("DG", cell, horizontal_degree) + self.base_elt_hori_cg = FiniteElement("CG", cell, h1_degree(family, horizontal_degree)) # vertical base spaces self.base_elt_vert_cg = FiniteElement("CG", interval, vertical_degree+1) @@ -245,17 +283,23 @@ def build_hcurl_space(self, family, horizontal_degree, vertical_degree=None): Returns: :class:`FunctionSpace`: the HCurl space. """ + if family is None: + logger.warning('There is no HCurl space for this family. Not creating one') + return None + if self.extruded_mesh: if not self._initialised_base_spaces: if vertical_degree is None: raise ValueError('vertical_degree must be specified to create HCurl space on an extruded mesh') self.build_base_spaces(family, horizontal_degree, vertical_degree) - Vh_elt = HCurl(TensorProductElement(self.S1, self.T1)) - Vv_elt = HCurl(TensorProductElement(self.S2, self.T0)) + Vh_elt = HCurl(TensorProductElement(self.base_elt_hori_hcurl, self.base_elt_vert_cg)) + Vv_elt = HCurl(TensorProductElement(self.base_elt_hori_cg, self.base_elt_vert_dg)) V_elt = Vh_elt + Vv_elt else: cell = self.mesh.ufl_cell().cellname() - V_elt = FiniteElement(family, cell, horizontal_degree) + hcurl_family = hdiv_hcurl_dict[family] + V_elt = FiniteElement(hcurl_family, cell, horizontal_degree) + return FunctionSpace(self.mesh, V_elt, name='HCurl') def build_hdiv_space(self, family, horizontal_degree, vertical_degree=None): @@ -278,13 +322,14 @@ def build_hdiv_space(self, family, horizontal_degree, vertical_degree=None): if vertical_degree is None: raise ValueError('vertical_degree must be specified to create HDiv space on an extruded mesh') self.build_base_spaces(family, horizontal_degree, vertical_degree) - Vh_elt = HDiv(TensorProductElement(self.S1, self.T1)) - Vt_elt = TensorProductElement(self.S2, self.T0) + Vh_elt = HDiv(TensorProductElement(self.base_elt_hori_hdiv, self.base_elt_vert_dg)) + Vt_elt = TensorProductElement(self.base_elt_hori_dg, self.base_elt_vert_cg) Vv_elt = HDiv(Vt_elt) V_elt = Vh_elt + Vv_elt else: cell = self.mesh.ufl_cell().cellname() - V_elt = FiniteElement(family, cell, horizontal_degree) + hdiv_family = hcurl_hdiv_dict[family] + V_elt = FiniteElement(hdiv_family, cell, horizontal_degree) return FunctionSpace(self.mesh, V_elt, name='HDiv') def build_l2_space(self, horizontal_degree, vertical_degree=None, variant=None, name='L2'): @@ -311,14 +356,18 @@ def build_l2_space(self, horizontal_degree, vertical_degree=None, variant=None, if self.extruded_mesh: if vertical_degree is None: raise ValueError('vertical_degree must be specified to create L2 space on an extruded mesh') - if not self._initialised_base_spaces or self.T1.degree() != vertical_degree or self.T1.variant() != variant: + if (not self._initialised_base_spaces + or self.base_elt_vert_dg.degree() != vertical_degree + or self.base_elt_vert_dg.variant() != variant + or self.base_elt_hori_dg.degree() != horizontal_degree + or self.base_elt_hori_dg.degree() != variant): cell = self.mesh._base_mesh.ufl_cell().cellname() - S2 = FiniteElement("DG", cell, horizontal_degree, variant=variant) - T1 = FiniteElement("DG", interval, vertical_degree, variant=variant) + base_elt_hori_dg = FiniteElement("DG", cell, horizontal_degree, variant=variant) + base_elt_vert_dg = FiniteElement("DG", interval, vertical_degree, variant=variant) else: - S2 = self.S2 - T1 = self.T1 - V_elt = TensorProductElement(S2, T1) + base_elt_hori_dg = self.base_elt_hori_dg + base_elt_vert_dg = self.base_elt_vert_dg + V_elt = TensorProductElement(base_elt_hori_dg, base_elt_vert_dg) else: cell = self.mesh.ufl_cell().cellname() V_elt = FiniteElement("DG", cell, horizontal_degree, variant=variant) @@ -348,9 +397,9 @@ def build_theta_space(self, horizontal_degree, vertical_degree): assert self.extruded_mesh, 'Cannot create theta space if mesh is not extruded' if not self._initialised_base_spaces: cell = self.mesh._base_mesh.ufl_cell().cellname() - self.S2 = FiniteElement("DG", cell, horizontal_degree) - self.T0 = FiniteElement("CG", interval, vertical_degree+1) - V_elt = TensorProductElement(self.S2, self.T0) + self.base_elt_hori_dg = FiniteElement("DG", cell, horizontal_degree) + self.base_elt_vert_cg = FiniteElement("CG", interval, vertical_degree+1) + V_elt = TensorProductElement(self.base_elt_hori_dg, self.base_elt_vert_cg) return FunctionSpace(self.mesh, V_elt, name='theta') def build_h1_space(self, horizontal_degree, vertical_degree=None, name='H1'): @@ -374,16 +423,52 @@ def build_h1_space(self, horizontal_degree, vertical_degree=None, name='H1'): if self.extruded_mesh: if vertical_degree is None: raise ValueError('vertical_degree must be specified to create H1 space on an extruded mesh') - cell = self.mesh._base_mesh.ufl_cell().cellname() - CG_hori = FiniteElement("CG", cell, horizontal_degree) - CG_vert = FiniteElement("CG", interval, vertical_degree) - V_elt = TensorProductElement(CG_hori, CG_vert) + if (not self._initialised_base_spaces + or self.base_elt_vert_cg.degree() != vertical_degree + or self.base_elt_hori_cg.degree() != horizontal_degree): + cell = self.mesh._base_mesh.ufl_cell().cellname() + base_elt_hori_cg = FiniteElement("CG", cell, horizontal_degree) + base_elt_vert_cg = FiniteElement("CG", interval, vertical_degree) + else: + base_elt_hori_cg = self.base_elt_hori_cg + base_elt_vert_cg = self.base_elt_vert_cg + V_elt = TensorProductElement(base_elt_hori_cg, base_elt_vert_cg) else: cell = self.mesh.ufl_cell().cellname() V_elt = FiniteElement("CG", cell, horizontal_degree) return FunctionSpace(self.mesh, V_elt, name=name) + def build_bdfm_base_spaces(self, horizontal_degree, vertical_degree): + """ + Builds the :class:`FiniteElement` objects for the base mesh when using + the . + + Args: + horizontal_degree (int): the polynomial degree of the horizontal + part of the L2 space. + vertical_degree (int): the polynomial degree of the vertical part of + the L2 space. + """ + + cell = self.mesh._base_mesh.ufl_cell().cellname() + + hdiv_family = 'BDFM' + + # horizontal base spaces + self.base_elt_hori_hdiv = FiniteElement(hdiv_family, cell, horizontal_degree+1) + self.base_elt_hori_dg = FiniteElement("DG", cell, horizontal_degree) + + # Add bubble space + self.base_elt_hori_cg = FiniteElement("CG", cell, horizontal_degree+1) + self.base_elt_hori_cg += FiniteElement("Bubble", cell, horizontal_degree+2) + + # vertical base spaces + self.base_elt_vert_cg = FiniteElement("CG", interval, vertical_degree+1) + self.base_elt_vert_dg = FiniteElement("DG", interval, vertical_degree) + + self._initialised_base_spaces = True + def check_degree_args(name, mesh, degree, horizontal_degree, vertical_degree): """ diff --git a/unit-tests/test_function_spaces.py b/unit-tests/test_function_spaces.py new file mode 100644 index 000000000..e72789f28 --- /dev/null +++ b/unit-tests/test_function_spaces.py @@ -0,0 +1,148 @@ +""" +Tests the building of different spaces in the de Rham complex and other useful +spaces in Gusto. +""" + +from firedrake import UnitIntervalMesh, ExtrudedMesh, UnitSquareMesh +from gusto import Spaces +import pytest + +# List all allowed families for a particular domain +# Don't need to bother testing sphere +domain_family_dict = {'interval': ['CG'], + 'vertical_slice': ['CG'], + 'plane': ['BDM', 'BDMF', 'BDME', 'BDFM', + 'RT', 'RTF', 'RTE', 'RTCF', 'RTCE'], + 'extruded_plane': ['BDM', 'BDMF', 'BDME', 'BDFM', + 'RT', 'RTF', 'RTE', 'RTCF', 'RTCE']} + +reduced_domain_family_dict = {'interval': ['CG'], + 'vertical_slice': ['CG'], + 'plane': ['BDM', 'BDFM', 'RT', 'RTCF'], + 'extruded_plane': ['BDM', 'BDFM', 'RT', 'RTCF']} + + +# Routine to form all combinations of domains and families +def combos(domains_and_families): + all_combos = [] + for domain, families in domains_and_families.items(): + for family in families: + all_combos.append((domain, family)) + + return all_combos + + +def set_up_mesh(domain, family): + + if family in ['BDM', 'BDMF', 'BDME', 'BDFM', 'RT', 'RTE', 'RTF']: + quadrilateral = False + elif family in ['RTCF', 'RTCE']: + quadrilateral = True + + if domain == 'interval': + mesh = UnitIntervalMesh(3) + elif domain == 'vertical_slice': + m = UnitIntervalMesh(3) + mesh = ExtrudedMesh(m, 3, 3) + elif domain == 'plane': + mesh = UnitSquareMesh(3, 3, quadrilateral=quadrilateral) + elif domain == 'extruded_plane': + m = UnitSquareMesh(3, 3, quadrilateral=quadrilateral) + mesh = ExtrudedMesh(m, 3, 3) + else: + raise ValueError(f'domain {domain} not recognised') + + return mesh + + +# ---------------------------------------------------------------------------- # +# Test creation of full de Rham complex +# ---------------------------------------------------------------------------- # +@pytest.mark.parametrize("domain, family", combos(domain_family_dict)) +def test_de_rham_spaces(domain, family): + + mesh = set_up_mesh(domain, family) + spaces = Spaces(mesh) + + if domain in ['vertical_slice', 'extruded_plane']: + # Need horizontal and vertical degrees + degree = (1, 2) + spaces.build_compatible_spaces(family, degree[0], degree[1]) + else: + degree = 1 + spaces.build_compatible_spaces(family, degree) + + # Work out correct CG degree + if family in ['BDM', 'BDME', 'BDMF']: + if domain in ['vertical_slice', 'extruded_plane']: + cg_degree = (degree[0]+2, degree[1]+2) + else: + cg_degree = degree + 2 + elif domain in ['vertical_slice', 'extruded_plane']: + cg_degree = (degree[0]+1, degree[1]+1) + else: + cg_degree = degree + 1 + + # Check that H1 spaces and L2 spaces have the correct degrees + cg_space = spaces('H1') + elt = cg_space.ufl_element() + assert elt.degree() == cg_degree, '"H1" space does not seem to be degree ' \ + + f'{cg_degree}. Found degree {elt.degree()}' + + dg_space = spaces('L2') + elt = dg_space.ufl_element() + assert elt.degree() == degree, '"L2" space does not seem to be degree ' \ + + f'{degree}. Found degree {elt.degree()}' + + +# ---------------------------------------------------------------------------- # +# Test creation of DG1 equispaced +# ---------------------------------------------------------------------------- # +@pytest.mark.parametrize("domain, family", combos(reduced_domain_family_dict)) +def test_dg_equispaced(domain, family): + + mesh = set_up_mesh(domain, family) + spaces = Spaces(mesh) + + DG1 = spaces('DG1_equispaced') + elt = DG1.ufl_element() + assert elt.degree() in [1, (1, 1)], '"DG1 equispaced" does not seem to be ' \ + + f'degree 1. Found degree {elt.degree()}' + assert elt.variant() == "equispaced", '"DG1 equispaced" does not seem to ' \ + + f'be equispaced variant. Found variant {elt.variant()}' + + +# ---------------------------------------------------------------------------- # +# Test creation of DG0 space +# ---------------------------------------------------------------------------- # +@pytest.mark.parametrize("domain, family", combos(reduced_domain_family_dict)) +def test_dg0(domain, family): + + mesh = set_up_mesh(domain, family) + spaces = Spaces(mesh) + + DG0 = spaces('DG0', 'DG', degree=0) + elt = DG0.ufl_element() + assert elt.degree() in [0, (0, 0)], '"DG0" space does not seem to be' \ + + f'degree 0. Found degree {elt.degree()}' + + +# ---------------------------------------------------------------------------- # +# Test creation of a general CG space +# ---------------------------------------------------------------------------- # +@pytest.mark.parametrize("domain, family", combos(reduced_domain_family_dict)) +def test_cg(domain, family): + + mesh = set_up_mesh(domain, family) + spaces = Spaces(mesh) + + if domain in ['vertical_slice', 'extruded_plane']: + degree = (1, 2) + CG = spaces('CG', 'CG', horizontal_degree=degree[0], vertical_degree=degree[1]) + else: + degree = 3 + CG = spaces('CG', 'CG', degree=degree) + + elt = CG.ufl_element() + assert elt.degree() == degree, '"CG" space does not seem to be degree ' \ + + f'{degree}. Found degree {elt.degree()}' From 79aa19cc17032f0dfe7f545cf64c276423d52791 Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Thu, 27 Jul 2023 08:42:54 +0100 Subject: [PATCH 4/4] fix replace_perp test --- unit-tests/fml_tests/test_replace_perp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unit-tests/fml_tests/test_replace_perp.py b/unit-tests/fml_tests/test_replace_perp.py index 967780696..cd7d91f7b 100644 --- a/unit-tests/fml_tests/test_replace_perp.py +++ b/unit-tests/fml_tests/test_replace_perp.py @@ -19,7 +19,7 @@ def test_replace_perp(): Nx = 5 mesh = UnitSquareMesh(Nx, Nx) domain = Domain(mesh, 0.1, "BDM", 1) - spaces = [space for space in domain.compatible_spaces] + spaces = [domain.spaces('HDiv'), domain.spaces('L2')] W = MixedFunctionSpace(spaces) # set up labelled form with subject u