Implement Real VectorFunctionSpace (#3284)

Also fail hard if a user attempts to create an Argument in one of these spaces.

Co-authored-by: Connor Ward <[email protected]>

* test_vector_real_space_assign_constant

---------

Co-authored-by: Connor Ward <[email protected]>

firedrake/function.py

@@ -6,8 +6,9 @@
 from ufl.domain import extract_unique_domain
 import cachetools
 import ctypes
-from collections import OrderedDict
 from ctypes import POINTER, c_int, c_double, c_void_p
+from collections.abc import Collection
+from numbers import Number
 
 from pyop2 import op2, mpi
 from pyop2.exceptions import DataTypeError, DataValueError
@@ -310,7 +311,7 @@ def __getattr__(self, name):
 
     def __dir__(self):
         current = super(Function, self).__dir__()
-        return list(OrderedDict.fromkeys(dir(self._data) + current))
+        return list(dict.fromkeys(dir(self._data) + current))
 
     @utils.cached_property
     @FunctionMixin._ad_annotate_subfunctions
@@ -452,14 +453,13 @@ def assign(self, expr, subset=None):
             expressions (e.g. involving the product of functions) :meth:`.Function.interpolate`
             should be used.
         """
-        if expr == 0:
-            self.dat.zero(subset=subset)
-        elif self.ufl_element().family() == "Real":
+        if self.ufl_element().family() == "Real" and isinstance(expr, (Number, Collection)):
             try:
                 self.dat.data_wo[...] = expr
-                return self
             except (DataTypeError, DataValueError) as e:
                 raise ValueError(e)
+        elif expr == 0:
+            self.dat.zero(subset=subset)
         else:
             from firedrake.assign import Assigner
             Assigner(self, expr, subset).assign()

firedrake/functionspaceimpl.py

@@ -475,7 +475,6 @@ def __init__(self, mesh, element, name=None):
         super(FunctionSpace, self).__init__()
         if type(element) is finat.ufl.MixedElement:
             raise ValueError("Can't create FunctionSpace for MixedElement")
-        sdata = get_shared_data(mesh, element)
         # The function space shape is the number of dofs per node,
         # hence it is not always the value_shape.  Vector and Tensor
         # element modifiers *must* live on the outside!
@@ -496,7 +495,6 @@ def __init__(self, mesh, element, name=None):
         else:
             self.shape = ()
         self._ufl_function_space = ufl.FunctionSpace(mesh.ufl_mesh(), element)
-        self._shared_data = sdata
         self._mesh = mesh
 
         self.rank = len(self.shape)
@@ -512,28 +510,37 @@ def __init__(self, mesh, element, name=None):
         space node."""
         self.name = name
         r"""The (optional) descriptive name for this space."""
-        self.node_set = sdata.node_set
-        r"""A :class:`pyop2.types.set.Set` representing the function space nodes."""
-        self.dof_dset = op2.DataSet(self.node_set, self.shape or 1,
-                                    name="%s_nodes_dset" % self.name)
-        r"""A :class:`pyop2.types.dataset.DataSet` representing the function space
-        degrees of freedom."""
-
         # User comm
         self.comm = mesh.comm
+        self.set_shared_data()
+        self.dof_dset = self.make_dof_dset()
+        r"""A :class:`pyop2.types.dataset.DataSet` representing the function space
+        degrees of freedom."""
+        self.node_set = self.dof_dset.set
+        r"""A :class:`pyop2.types.set.Set` representing the function space nodes."""
         # Internal comm
         self._comm = mpi.internal_comm(self.node_set.comm)
+
+    def set_shared_data(self):
+        element = self.ufl_element()
+        sdata = get_shared_data(self._mesh, element)
         # Need to create finat element again as sdata does not
         # want to carry finat_element.
         self.finat_element = create_element(element)
         # Used for reconstruction of mixed/component spaces.
         # sdata carries real_tensorproduct.
+        self._shared_data = sdata
         self.real_tensorproduct = sdata.real_tensorproduct
         self.extruded = sdata.extruded
         self.offset = sdata.offset
         self.offset_quotient = sdata.offset_quotient
         self.cell_boundary_masks = sdata.cell_boundary_masks
         self.interior_facet_boundary_masks = sdata.interior_facet_boundary_masks
+        self.global_numbering = sdata.global_numbering
+
+    def make_dof_dset(self):
+        return op2.DataSet(self._shared_data.node_set, self.shape or 1,
+                           name=f"{self.name}_nodes_dset")
 
     def __del__(self):
         if hasattr(self, "_comm"):
@@ -584,7 +591,7 @@ def _dm(self):
         _, level = get_level(self.mesh())
         dmhooks.attach_hooks(dm, level=level,
                              sf=self.mesh().topology_dm.getPointSF(),
-                             section=self._shared_data.global_numbering)
+                             section=self.global_numbering)
         # Remember the function space so we can get from DM back to FunctionSpace.
         dmhooks.set_function_space(dm, self)
         return dm
@@ -1101,17 +1108,7 @@ class RealFunctionSpace(FunctionSpace):
     """
 
     finat_element = None
-    rank = 0
-    shape = ()
-    value_size = 1
-
-    def __init__(self, mesh, element, name):
-        self._ufl_function_space = ufl.FunctionSpace(mesh.ufl_mesh(), element)
-        self.name = name
-        self.comm = mesh.comm
-        self._mesh = mesh
-        self.dof_dset = op2.GlobalDataSet(self.make_dat())
-        self.node_set = self.dof_dset.set
+    global_numbering = None
 
     def __eq__(self, other):
         if not isinstance(other, RealFunctionSpace):
@@ -1126,16 +1123,11 @@ def __ne__(self, other):
     def __hash__(self):
         return hash((self.mesh(), self.ufl_element()))
 
-    def _dm(self):
-        from firedrake.mg.utils import get_level
-        dm = self.dof_dset.dm
-        _, level = get_level(self.mesh())
-        dmhooks.attach_hooks(dm, level=level,
-                             sf=self.mesh().topology_dm.getPointSF(),
-                             section=None)
-        # Remember the function space so we can get from DM back to FunctionSpace.
-        dmhooks.set_function_space(dm, self)
-        return dm
+    def set_shared_data(self):
+        pass
+
+    def make_dof_dset(self):
+        return op2.GlobalDataSet(self.make_dat())
 
     def make_dat(self, val=None, valuetype=None, name=None):
         r"""Return a newly allocated :class:`pyop2.types.glob.Global` representing the
@@ -1162,9 +1154,6 @@ def top_nodes(self):
         ":class:`RealFunctionSpace` objects have no bottom nodes."
         return None
 
-    def dim(self):
-        return 1
-
     def local_to_global_map(self, bcs, lgmap=None):
         assert len(bcs) == 0
         return None

firedrake/ufl_expr.py

@@ -38,6 +38,8 @@ def __new__(cls, *args, **kwargs):
         return super().__new__(cls, *args, **kwargs)
 
     def __init__(self, function_space, number, part=None):
+        if function_space.ufl_element().family() == "Real" and function_space.shape != ():
+            raise NotImplementedError(f"{type(self).__name__} on a vector-valued Real space is not supported.")
         super(Argument, self).__init__(function_space.ufl_function_space(),
                                        number, part=part)
         self._function_space = function_space
@@ -88,6 +90,8 @@ class Coargument(ufl.argument.Coargument):
     """
 
     def __init__(self, function_space, number, part=None):
+        if function_space.ufl_element().family() == "Real" and function_space.shape != ():
+            raise NotImplementedError(f"{type(self).__name__} on a vector-valued Real space is not supported.")
         super(Coargument, self).__init__(function_space.ufl_function_space(),
                                          number, part=part)
         self._function_space = function_space

tests/regression/test_assemble.py

@@ -282,3 +282,12 @@ def test_3125():
     F = inner(z, tst)*dx + inner(u, v)/(d+p)*dx(2, degree=10)
     # should run without error
     solve(F == 0, z)
+
+
+@pytest.mark.xfail(reason="Arguments on vector-valued R spaces are not supported")
+def test_assemble_vector_rspace_one_form(mesh):
+    V = VectorFunctionSpace(mesh, "Real", 0, dim=2)
+    u = Function(V)
+    U = inner(u, u)*dx
+    L = derivative(U, u)
+    assemble(L)

tests/regression/test_function.py

@@ -24,6 +24,12 @@ def W_nonstandard_shape():
     return W_nonstandard_shape
 
 
+@pytest.fixture
+def Rvector():
+    mesh = UnitSquareMesh(5, 5)
+    return VectorFunctionSpace(mesh, "R", 0, dim=4)
+
+
 def test_firedrake_scalar_function(V):
     f = Function(V)
     f.interpolate(Constant(1))
@@ -198,3 +204,38 @@ def test_tensor_function_zero_with_subset(W):
     f.zero(subset=subset)
     assert np.allclose(f.dat.data_ro[:3], 0.0)
     assert np.allclose(f.dat.data_ro[3:], 1.0)
+
+
+@pytest.mark.parametrize("value", [
+    1,
+    1.0,
+    (1, 2, 3, 4),
+    [1, 2, 3, 4],
+    np.array([5, 6, 7, 8]),
+    range(4)], ids=type)
+def test_vector_real_space_assign(Rvector, value):
+    f = Function(Rvector)
+    f.assign(value)
+    assert np.allclose(f.dat.data_ro, value)
+
+
+def test_vector_real_space_assign_function(Rvector):
+    value = [9, 10, 11, 12]
+    fvalue = Function(Rvector, val=value)
+    f = Function(Rvector)
+    f.assign(fvalue)
+    assert np.allclose(f.dat.data_ro, value)
+
+
+def test_vector_real_space_assign_constant(Rvector):
+    value = [9, 10, 11, 12]
+    fvalue = Constant(value)
+    f = Function(Rvector)
+    f.assign(fvalue)
+    assert np.allclose(f.dat.data_ro, value)
+
+
+def test_vector_real_space_assign_zero(Rvector):
+    f = Function(Rvector, val=[9, 10, 11, 12])
+    f.assign(zero())
+    assert np.allclose(f.dat.data_ro, 0)

tests/regression/test_scaled_mass.py

@@ -44,13 +44,13 @@ def test_math_functions(mesh, expr, value, typ, fs_type):
             f = inner(f, f)
     elif typ == 'Constant':
         if fs_type == 'vector':
-            f = Constant([value, value], domain=mesh)
+            f = Constant([value, value])
             f = dot(f, f)
         elif fs_type == 'tensor':
-            f = Constant([[value, value], [value, value]], domain=mesh)
+            f = Constant([[value, value], [value, value]])
             f = inner(f, f)
         else:
-            f = Constant(value, domain=mesh)
+            f = Constant(value)
 
     H = FunctionSpace(mesh, 'CG', 1)
     u = TrialFunction(H)

tests/regression/test_zero_forms.py

@@ -10,43 +10,38 @@ def mesh(request):
     return UnitSquareMesh(10, 10, quadrilateral=quadrilateral)
 
 
-@pytest.fixture(scope='module')
-def one(mesh):
-    return Constant(1, domain=mesh)
-
-
 domains = [(1, 2),
            (2, 3),
            (3, 4),
            (4, 1),
            (1, 2, 3, 4)]
 
 
-def test_ds_dx(one):
-    assert np.allclose(assemble(one*dx + one*ds), 5.0)
+def test_ds_dx(mesh):
+    assert np.allclose(assemble(1*dx(domain=mesh) + 1*ds(domain=mesh)), 5.0)
 
 
 @pytest.mark.parametrize('domains', domains)
-def test_dsn(one, domains):
+def test_dsn(mesh, domains):
 
-    assert np.allclose(assemble(one*ds(domains)), len(domains))
+    assert np.allclose(assemble(1*ds(domains, domain=mesh)), len(domains))
 
-    form = one*ds(domains[0])
+    form = 1*ds(domains[0], domain=mesh)
 
     for d in domains[1:]:
-        form += one*ds(d)
+        form += 1*ds(d, domain=mesh)
     assert np.allclose(assemble(form), len(domains))
 
 
 @pytest.mark.parallel
-def test_dsn_parallel(one):
+def test_dsn_parallel(mesh):
     for d in domains:
-        assert np.allclose(assemble(one*ds(d)), len(d))
+        assert np.allclose(assemble(1*ds(d, domain=mesh)), len(d))
 
     for domain in domains:
-        form = one*ds(domain[0])
+        form = 1*ds(domain[0], domain=mesh)
         for d in domain[1:]:
-            form += one*ds(d)
+            form += 1*ds(d, domain=mesh)
         assert np.allclose(assemble(form), len(domain))
 
 
@@ -67,27 +62,22 @@ def test_dsn_parallel(one):
                                            ['scalar', 'vector', 'tensor']))
 def test_math_functions(mesh, expr, value, typ, fs_type):
     if typ == 'function':
-        if fs_type == "vector":
-            V = VectorFunctionSpace(mesh, 'CG', 1)
-        elif fs_type == "tensor":
-            V = TensorFunctionSpace(mesh, 'CG', 1)
-        else:
-            V = FunctionSpace(mesh, 'CG', 1)
-        f = Function(V)
-        f.assign(value)
-        if fs_type == "vector":
-            f = dot(f, f)
-        elif fs_type == "tensor":
-            f = inner(f, f)
+        family, degree = 'CG', 1
     elif typ == 'constant':
-        if fs_type == "vector":
-            f = Constant([value, value], domain=mesh)
-            f = dot(f, f)
-        elif fs_type == "tensor":
-            f = Constant([[value, value], [value, value]], domain=mesh)
-            f = inner(f, f)
-        else:
-            f = Constant(value, domain=mesh)
+        family, degree = 'Real', 0
+
+    if fs_type == "vector":
+        V = VectorFunctionSpace(mesh, family, degree)
+    elif fs_type == "tensor":
+        V = TensorFunctionSpace(mesh, family, degree)
+    else:
+        V = FunctionSpace(mesh, family, degree)
+    f = Function(V)
+    f.assign(value)
+    if fs_type == "vector":
+        f = dot(f, f)
+    elif fs_type == "tensor":
+        f = inner(f, f)
 
     actual = assemble(eval(expr)*dx)
 

tests/vertexonly/test_interpolation_from_parent.py

@@ -194,7 +194,7 @@ def test_scalar_function_interpolation(parentmesh, vertexcoords, fs):
     A_w.interpolate(v, output=w_v)
     assert np.allclose(w_v.dat.data_ro, np.sum(vertexcoords, axis=1))
     # use it again for a different Function in V
-    v = Function(V).assign(Constant(2, domain=parentmesh))
+    v = Function(V).assign(Constant(2))
     A_w.interpolate(v, output=w_v)
     assert np.allclose(w_v.dat.data_ro, 2)
 
@@ -343,9 +343,9 @@ def test_scalar_real_interpolation(parentmesh, vertexcoords):
     # Remove below when interpolating constant onto Real works for extruded
     if type(parentmesh.topology) is mesh.ExtrudedMeshTopology:
         with pytest.raises(ValueError):
-            interpolate(Constant(1.0, domain=parentmesh), V)
+            interpolate(Constant(1), V)
         return
-    v = interpolate(Constant(1.0, domain=parentmesh), V)
+    v = interpolate(Constant(1), V)
     w_v = interpolate(v, W)
     assert np.allclose(w_v.dat.data_ro, 1.)
 
@@ -358,9 +358,9 @@ def test_scalar_real_interpolator(parentmesh, vertexcoords):
     # Remove below when interpolating constant onto Real works for extruded
     if type(parentmesh.topology) is mesh.ExtrudedMeshTopology:
         with pytest.raises(ValueError):
-            interpolate(Constant(1.0, domain=parentmesh), V)
+            interpolate(Constant(1), V)
         return
-    v = interpolate(Constant(1.0, domain=parentmesh), V)
+    v = interpolate(Constant(1), V)
     A_w = Interpolator(TestFunction(V), W)
     w_v = Function(W)
     A_w.interpolate(v, output=w_v)