record information about function space continuity in Spaces

gusto/core/function_spaces.py

@@ -71,6 +71,7 @@ def __init__(self, mesh):
         self.mesh = mesh
         self.extruded_mesh = hasattr(mesh, "_base_mesh")
         self.de_rham_complex = {}
+        self.continuity = {}
 
     def __call__(self, name):
         """
@@ -89,7 +90,7 @@ def __call__(self, name):
         else:
             raise ValueError(f'The space container has no space {name}')
 
-    def add_space(self, name, space, overwrite_space=False):
+    def add_space(self, name, space, overwrite_space=False, continuity=None):
         """
         Adds a function space to the container.
 
@@ -100,6 +101,10 @@ def add_space(self, name, space, overwrite_space=False):
             overwrite_space (bool, optional): Logical to allow space existing in
                 container to be overwritten by an incoming space. Defaults to
                 False.
+            continuity: Whether the space is continuous or not. For spaces on
+                extruded meshes, must be a dictionary with entries for the
+                'horizontal' and 'vertical' continuity.
+                If None, defaults to value of is_cg(space).
         """
 
         if hasattr(self, name) and not overwrite_space:
@@ -109,9 +114,29 @@ def add_space(self, name, space, overwrite_space=False):
 
         setattr(self, name, space)
 
-    def create_space(self, name, family, degree, overwrite_space=False):
+        # set the continuity of the space - for extruded meshes specify both directions
+        if continuity is None:
+            from gusto.time_discretisation.wrappers import is_cg
+            continuity = is_cg(space)
+            if self.extruded_mesh:
+                self.continuity[name] = {
+                    'horizontal': continuity,
+                    'vertical': continuity
+                }
+            else:
+                self.continuity[name] = continuity
+        else:
+            if self.extruded_mesh:
+                self.continuity[name] = {
+                    'horizontal': continuity['horizontal'],
+                    'vertical': continuity['vertical']
+                }
+            else:
+                self.continuity[name] = continuity
+
+    def create_space(self, name, family, degree, overwrite_space=False, continuity=None):
         """
-        Creates a space and adds it to the container..
+        Creates a space and adds it to the container.
 
         Args:
             name (str): the name to give to the space.
@@ -120,18 +145,18 @@ def create_space(self, name, family, degree, overwrite_space=False):
             overwrite_space (bool, optional): Logical to allow space existing in
                 container to be overwritten by an incoming space. Defaults to
                 False.
+            continuity: Whether the space is continuous or not. For spaces on
+                extruded meshes, must be a tuple for the (horizontal, vertical)
+                continuity. If None, defaults to value of is_cg(space).
 
         Returns:
             :class:`FunctionSpace`: the desired function space.
         """
-
-        if hasattr(self, name) and not overwrite_space:
-            raise RuntimeError(f'Space {name} already exists. If you really '
-                               + 'to create it then set `overwrite_space` as '
-                               + 'to be True')
-
         space = FunctionSpace(self.mesh, family, degree, name=name)
-        setattr(self, name, space)
+
+        self.add_space(name, space,
+                       overwrite_space=overwrite_space,
+                       continuity=continuity)
         return space
 
     def build_compatible_spaces(self, family, horizontal_degree,
@@ -181,9 +206,15 @@ def build_compatible_spaces(self, family, horizontal_degree,
         setattr(self, "L2"+complex_name, de_rham_complex.L2)
         # Register L2 space as DG also
         setattr(self, "DG"+complex_name, de_rham_complex.L2)
-        if hasattr(de_rham_complex, "theta"+complex_name):
+        if hasattr(de_rham_complex, "theta"):
             setattr(self, "theta"+complex_name, de_rham_complex.theta)
 
+        # Grab the continuity information from the complex
+        for space_type in ("H1, HCurl", "HDiv", "L2", "DG", "theta"):
+            space_name = space_type + complex_name
+            if hasattr(de_rham_complex, space_type):
+                self.continuity[space_name] = de_rham_complex.continuity[space_type]
+
     def build_dg1_equispaced(self):
         """
         Builds the equispaced variant of the DG1 function space, which is used in
@@ -198,12 +229,15 @@ def build_dg1_equispaced(self):
             hori_elt = FiniteElement('DG', cell, 1, variant='equispaced')
             vert_elt = FiniteElement('DG', interval, 1, variant='equispaced')
             V_elt = TensorProductElement(hori_elt, vert_elt)
+            continuity = {'horizontal': False, 'vertical': False}
         else:
             cell = self.mesh.ufl_cell().cellname()
             V_elt = FiniteElement('DG', cell, 1, variant='equispaced')
+            continuity = False
 
         space = FunctionSpace(self.mesh, V_elt, name='DG1_equispaced')
-        setattr(self, 'DG1_equispaced', space)
+
+        self.add_space('DG1_equispaced', space, continuity=continuity)
         return space
 
 
@@ -234,6 +268,7 @@ def __init__(self, mesh, family, horizontal_degree, vertical_degree=None,
         self.extruded_mesh = hasattr(mesh, '_base_mesh')
         self.family = family
         self.complex_name = complex_name
+        self.continuity = {}
         self.build_base_spaces(family, horizontal_degree, vertical_degree)
         self.build_compatible_spaces()
 
@@ -303,15 +338,24 @@ def build_compatible_spaces(self):
         if self.extruded_mesh:
             # Horizontal and vertical degrees
             # need specifying separately. Vtheta needs returning.
-            Vcg = self.build_h1_space()
+            Vcg, continuity = self.build_h1_space()
+            self.continuity["H1"] = continuity
             setattr(self, "H1", Vcg)
-            Vcurl = self.build_hcurl_space()
+
+            Vcurl, continuity = self.build_hcurl_space()
+            self.continuity["HCurl"] = continuity
             setattr(self, "HCurl", Vcurl)
-            Vu = self.build_hdiv_space()
+
+            Vu, continuity = self.build_hdiv_space()
+            self.continuity["HDiv"] = continuity
             setattr(self, "HDiv", Vu)
-            Vdg = self.build_l2_space()
+
+            Vdg, continuity = self.build_l2_space()
+            self.continuity["L2"] = continuity
             setattr(self, "L2", Vdg)
-            Vth = self.build_theta_space()
+
+            Vth, continuity = self.build_theta_space()
+            self.continuity["theta"] = continuity
             setattr(self, "theta", Vth)
 
             return Vcg, Vcurl, Vu, Vdg, Vth
@@ -320,25 +364,39 @@ def build_compatible_spaces(self):
             # 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()
+            Vcg, continuity = self.build_h1_space()
+            self.continuity["H1"] = continuity
             setattr(self, "H1", Vcg)
-            Vcurl = self.build_hcurl_space()
+
+            Vcurl, continuity = self.build_hcurl_space()
+            self.continuity["HCurl"] = continuity
             setattr(self, "HCurl", Vcurl)
-            Vu = self.build_hdiv_space()
+
+            Vu, continuity = self.build_hdiv_space()
+            self.continuity["HDiv"] = continuity
             setattr(self, "HDiv", Vu)
-            Vdg = self.build_l2_space()
+
+            Vdg, continuity = self.build_l2_space()
+            self.continuity["L2"] = continuity
             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()
+            Vcg, continuity = self.build_h1_space()
+
+            self.continuity["H1"] = continuity
             setattr(self, "H1", Vcg)
+
             setattr(self, "HCurl", None)
+
+            self.continuity["HDiv"] = continuity
             setattr(self, "HDiv", Vcg)
-            Vdg = self.build_l2_space()
+
+            Vdg, continuity = self.build_l2_space()
+            self.continuity["L2"] = continuity
             setattr(self, "L2", Vdg)
 
             return Vcg, Vdg
@@ -360,18 +418,21 @@ def build_hcurl_space(self):
         """
         if hdiv_hcurl_dict[self.family] is None:
             logger.warning('There is no HCurl space for this family. Not creating one')
-            return None
+            return None, None
 
         if self.extruded_mesh:
             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
+            continuity = {'horizontal': True, 'vertical': True}
         else:
             V_elt = self.base_elt_hori_hcurl
+            continuity = True
 
-        return FunctionSpace(self.mesh, V_elt, name='HCurl'+self.complex_name)
+        space_name = 'HCurl'+self.complex_name
+        return FunctionSpace(self.mesh, V_elt, name=space_name), continuity
 
     def build_hdiv_space(self):
         """
@@ -387,9 +448,13 @@ def build_hdiv_space(self):
                                           self.base_elt_vert_cg)
             Vv_elt = HDiv(Vt_elt)
             V_elt = Vh_elt + Vv_elt
+            continuity = {'horizontal': True, 'vertical': True}
         else:
             V_elt = self.base_elt_hori_hdiv
-        return FunctionSpace(self.mesh, V_elt, name='HDiv'+self.complex_name)
+            continuity = True
+
+        space_name = 'HDiv'+self.complex_name
+        return FunctionSpace(self.mesh, V_elt, name=space_name), continuity
 
     def build_l2_space(self):
         """
@@ -401,10 +466,13 @@ def build_l2_space(self):
 
         if self.extruded_mesh:
             V_elt = TensorProductElement(self.base_elt_hori_dg, self.base_elt_vert_dg)
+            continuity = {'horizontal': False, 'vertical': False}
         else:
             V_elt = self.base_elt_hori_dg
+            continuity = False
 
-        return FunctionSpace(self.mesh, V_elt, name='L2'+self.complex_name)
+        space_name = 'L2'+self.complex_name
+        return FunctionSpace(self.mesh, V_elt, name=space_name), continuity
 
     def build_theta_space(self):
         """
@@ -423,8 +491,10 @@ def build_theta_space(self):
         assert self.extruded_mesh, 'Cannot create theta space if mesh is not extruded'
 
         V_elt = TensorProductElement(self.base_elt_hori_dg, self.base_elt_vert_cg)
+        continuity = {'horizontal': False, 'vertical': True}
 
-        return FunctionSpace(self.mesh, V_elt, name='theta'+self.complex_name)
+        space_name = 'theta'+self.complex_name
+        return FunctionSpace(self.mesh, V_elt, name=space_name), continuity
 
     def build_h1_space(self):
         """
@@ -440,11 +510,13 @@ def build_h1_space(self):
 
         if self.extruded_mesh:
             V_elt = TensorProductElement(self.base_elt_hori_cg, self.base_elt_vert_cg)
-
+            continuity = {'horizontal': True, 'vertical': True}
         else:
             V_elt = self.base_elt_hori_cg
+            continuity = True
 
-        return FunctionSpace(self.mesh, V_elt, name='H1'+self.complex_name)
+        space_name = 'H1'+self.complex_name
+        return FunctionSpace(self.mesh, V_elt, name=space_name), continuity
 
 
 def check_degree_args(name, mesh, degree, horizontal_degree, vertical_degree):

unit-tests/test_function_spaces.py

@@ -100,6 +100,27 @@ def test_de_rham_spaces(domain, family):
     assert elt.degree() == degree, '"L2" space does not seem to be degree ' \
         + f'{degree}. Found degree {elt.degree()}'
 
+    # Check that continuities have been recorded correctly
+    if hasattr(mesh, "_base_mesh"):
+        expected_continuity = {
+            "H1": {'horizontal': True, 'vertical': True},
+            "L2": {'horizontal': False, 'vertical': False},
+            "HDiv": {'horizontal': True, 'vertical': True},
+            "HCurl": {'horizontal': True, 'vertical': True},
+            "theta": {'horizontal': False, 'vertical': True}
+        }
+    else:
+        expected_continuity = {
+            "H1": True,
+            "L2": False,
+            "HDiv": True,
+            "HCurl": True
+        }
+
+    for space, continuity in expected_continuity.items():
+        if space in spaces.continuity:
+            assert spaces.continuity[space] == continuity
+
 
 # ---------------------------------------------------------------------------- #
 # Test creation of DG1 equispaced
@@ -118,6 +139,13 @@ def test_dg_equispaced(domain, family):
     assert elt.variant() == "equispaced", '"DG1 equispaced" does not seem to ' \
         + f'be equispaced variant. Found variant {elt.variant()}'
 
+    if hasattr(mesh, "_base_mesh"):
+        expected_continuity = {'horizontal': False, 'vertical': False}
+    else:
+        expected_continuity = False
+
+    assert spaces.continuity['DG1_equispaced'] == expected_continuity, "DG is discontinuous"
+
 
 # ---------------------------------------------------------------------------- #
 # Test creation of DG0 space
@@ -133,6 +161,13 @@ def test_dg0(domain, family):
     assert elt.degree() in [0, (0, 0)], '"DG0" space does not seem to be' \
         + f'degree 0. Found degree {elt.degree()}'
 
+    if hasattr(mesh, "_base_mesh"):
+        expected_continuity = {'horizontal': False, 'vertical': False}
+    else:
+        expected_continuity = False
+
+    assert spaces.continuity['DG0'] == expected_continuity, "DG is discontinuous"
+
 
 # ---------------------------------------------------------------------------- #
 # Test creation of a general CG space
@@ -150,3 +185,10 @@ def test_cg(domain, family):
     assert elt.degree() == degree or elt.degree() == (degree, degree), \
         (f'"CG" space does not seem to be degree {degree}. '
          + f'Found degree {elt.degree()}')
+
+    if hasattr(mesh, "_base_mesh"):
+        expected_continuity = {'horizontal': True, 'vertical': True}
+    else:
+        expected_continuity = True
+
+    assert spaces.continuity['CG'] == expected_continuity, "CG is continuous"