diff --git a/firedrake/function.py b/firedrake/function.py index fa321e0352..d732d9440a 100644 --- a/firedrake/function.py +++ b/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() diff --git a/firedrake/functionspaceimpl.py b/firedrake/functionspaceimpl.py index 2f7ebeb5a0..a564abc287 100644 --- a/firedrake/functionspaceimpl.py +++ b/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 diff --git a/firedrake/ufl_expr.py b/firedrake/ufl_expr.py index b4ee6e5349..0d0b39638e 100644 --- a/firedrake/ufl_expr.py +++ b/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 diff --git a/tests/regression/test_assemble.py b/tests/regression/test_assemble.py index 3b671d50af..7751f95355 100644 --- a/tests/regression/test_assemble.py +++ b/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) diff --git a/tests/regression/test_function.py b/tests/regression/test_function.py index 20bb701c40..86703625ba 100644 --- a/tests/regression/test_function.py +++ b/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) diff --git a/tests/regression/test_scaled_mass.py b/tests/regression/test_scaled_mass.py index c66c9a905b..646e07529d 100644 --- a/tests/regression/test_scaled_mass.py +++ b/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) diff --git a/tests/regression/test_zero_forms.py b/tests/regression/test_zero_forms.py index 0260236f2e..73c1554f8b 100644 --- a/tests/regression/test_zero_forms.py +++ b/tests/regression/test_zero_forms.py @@ -10,11 +10,6 @@ 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), @@ -22,31 +17,31 @@ def one(mesh): (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) diff --git a/tests/vertexonly/test_interpolation_from_parent.py b/tests/vertexonly/test_interpolation_from_parent.py index a8afe84150..5ff999076c 100644 --- a/tests/vertexonly/test_interpolation_from_parent.py +++ b/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) diff --git a/tests/vertexonly/test_vertex_only_fs.py b/tests/vertexonly/test_vertex_only_fs.py index c7f8e1127c..22b20a548f 100644 --- a/tests/vertexonly/test_vertex_only_fs.py +++ b/tests/vertexonly/test_vertex_only_fs.py @@ -103,7 +103,7 @@ def functionspace_tests(vm): # Assembly works as expected - global assembly (integration) of a # constant on a vertex only mesh is evaluation of that constant # num_vertices (globally) times - f.interpolate(Constant(2, domain=vm)) + f.interpolate(Constant(2)) assert np.isclose(assemble(f*dx), 2*num_cells_mpi_global) if "input_ordering" in vm.name: assert vm.input_ordering is None @@ -131,7 +131,7 @@ def functionspace_tests(vm): assert np.allclose(h2.dat.data_ro_with_halos[idxs_to_include], h.dat.data_ro_with_halos[idxs_to_include]) # check we can interpolate expressions h2 = Function(W) - h2.interpolate(2*g*Constant(1, domain=vm)) + h2.interpolate(2*g) assert np.allclose(h2.dat.data_ro_with_halos[idxs_to_include], 2*np.prod(vm.input_ordering.coordinates.dat.data_ro_with_halos[idxs_to_include].reshape(-1, vm.input_ordering.geometric_dimension()), axis=1)) # Check that the opposite works g.dat.data_wo_with_halos[:] = -1 @@ -145,7 +145,7 @@ def functionspace_tests(vm): h = I_io.interpolate(g) assert np.allclose(h.dat.data_ro_with_halos[idxs_to_include], np.prod(vm.input_ordering.coordinates.dat.data_ro_with_halos[idxs_to_include].reshape(-1, vm.input_ordering.geometric_dimension()), axis=1)) assert np.all(h.dat.data_ro_with_halos[~idxs_to_include] == 0) - I2_io = Interpolator(2*TestFunction(V)*Constant(1, domain=vm), W) + I2_io = Interpolator(2*TestFunction(V), W) h2 = I2_io.interpolate(g) assert np.allclose(h2.dat.data_ro_with_halos[idxs_to_include], 2*np.prod(vm.input_ordering.coordinates.dat.data_ro_with_halos[idxs_to_include].reshape(-1, vm.input_ordering.geometric_dimension()), axis=1)) g = I_io.interpolate(h, transpose=True) @@ -156,7 +156,7 @@ def functionspace_tests(vm): assert np.allclose(g2.dat.data_ro_with_halos, 2*np.prod(vm.coordinates.dat.data_ro_with_halos.reshape(-1, vm.geometric_dimension()), axis=1)) I_io_transpose = Interpolator(TestFunction(W), V) - I2_io_transpose = Interpolator(2*TestFunction(W)*Constant(1, domain=vm.input_ordering), V) + I2_io_transpose = Interpolator(2*TestFunction(W), V) h = I_io_transpose.interpolate(g, transpose=True) assert np.allclose(h.dat.data_ro_with_halos[idxs_to_include], np.prod(vm.input_ordering.coordinates.dat.data_ro_with_halos[idxs_to_include].reshape(-1, vm.input_ordering.geometric_dimension()), axis=1)) assert np.all(h.dat.data_ro_with_halos[~idxs_to_include] == 0) @@ -202,7 +202,9 @@ def vectorfunctionspace_tests(vm): # num_vertices (globally) times. Note that we get a vertex cell for # each geometric dimension so we have to sum over geometric # dimension too. - f.interpolate(Constant([1] * gdim, domain=vm)) + R = VectorFunctionSpace(vm, "R", dim=gdim) + ones = Function(R).assign(1) + f.interpolate(ones) assert np.isclose(assemble(inner(f, f)*dx), num_cells_mpi_global*gdim) if "input_ordering" in vm.name: assert vm.input_ordering is None @@ -230,7 +232,7 @@ def vectorfunctionspace_tests(vm): assert np.allclose(h2.dat.data_ro_with_halos[idxs_to_include], h.dat.data_ro_with_halos[idxs_to_include]) # check we can interpolate expressions h2 = Function(W) - h2.interpolate(2*g*Constant(1, domain=vm)) + h2.interpolate(2*g) assert np.allclose(h2.dat.data_ro[idxs_to_include], 4*vm.input_ordering.coordinates.dat.data_ro_with_halos[idxs_to_include]) # Check that the opposite works g.dat.data_wo_with_halos[:] = -1 @@ -244,7 +246,7 @@ def vectorfunctionspace_tests(vm): h = I_io.interpolate(g) assert np.allclose(h.dat.data_ro[idxs_to_include], 2*vm.input_ordering.coordinates.dat.data_ro_with_halos[idxs_to_include]) assert np.all(h.dat.data_ro_with_halos[~idxs_to_include] == 0) - I2_io = Interpolator(2*TestFunction(V)*Constant(1, domain=vm), W) + I2_io = Interpolator(2*TestFunction(V), W) h2 = I2_io.interpolate(g) assert np.allclose(h2.dat.data_ro[idxs_to_include], 4*vm.input_ordering.coordinates.dat.data_ro_with_halos[idxs_to_include]) g = I_io.interpolate(h, transpose=True) @@ -255,7 +257,7 @@ def vectorfunctionspace_tests(vm): assert np.allclose(g2.dat.data_ro_with_halos, 4*vm.coordinates.dat.data_ro_with_halos) I_io_transpose = Interpolator(TestFunction(W), V) - I2_io_transpose = Interpolator(2*TestFunction(W)*Constant(1, domain=vm.input_ordering), V) + I2_io_transpose = Interpolator(2*TestFunction(W), V) h = I_io_transpose.interpolate(g, transpose=True) assert np.allclose(h.dat.data_ro[idxs_to_include], 2*vm.input_ordering.coordinates.dat.data_ro_with_halos[idxs_to_include]) assert np.all(h.dat.data_ro_with_halos[~idxs_to_include] == 0)