From 4109283bf24afa06a821bde11a993ac7c81f8e75 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Thu, 9 Jan 2025 19:24:14 +0000 Subject: [PATCH] Fix test for mixed-ness in Function.sub, and fix Cofunction.sub to match (#3961) * fix test for mixed-ness in Function.sub, and fix Cofunction.sub to match * allow FunctionSpace.sub to take negative indices * cofunc docstring --- firedrake/cofunction.py | 28 +++++++++++++++---- firedrake/function.py | 13 +++------ firedrake/functionspaceimpl.py | 8 +----- .../regression/test_vfs_component_bcs.py | 2 +- 4 files changed, 28 insertions(+), 23 deletions(-) diff --git a/firedrake/cofunction.py b/firedrake/cofunction.py index 4878e6da59..ecfb386dd9 100644 --- a/firedrake/cofunction.py +++ b/firedrake/cofunction.py @@ -4,6 +4,7 @@ from ufl.form import BaseForm from pyop2 import op2, mpi from pyadjoint.tape import stop_annotating, annotate_tape, get_working_tape +from finat.ufl import MixedElement import firedrake.assemble import firedrake.functionspaceimpl as functionspaceimpl from firedrake import utils, vector, ufl_expr @@ -119,11 +120,14 @@ def split(self): @utils.cached_property def _components(self): - if self.function_space().value_size == 1: + if self.function_space().rank == 0: return (self, ) else: - return tuple(type(self)(self.function_space().sub(i), val=op2.DatView(self.dat, i)) - for i in range(self.function_space().value_size)) + if self.dof_dset.cdim == 1: + return (type(self)(self.function_space().sub(0), val=self.dat),) + else: + return tuple(type(self)(self.function_space().sub(i), val=op2.DatView(self.dat, j)) + for i, j in enumerate(np.ndindex(self.dof_dset.dim))) @PETSc.Log.EventDecorator() def sub(self, i): @@ -137,9 +141,9 @@ def sub(self, i): :func:`~.VectorFunctionSpace` or :func:`~.TensorFunctionSpace` this returns a proxy object indexing the ith component of the space, suitable for use in boundary condition application.""" - if len(self.function_space()) == 1: - return self._components[i] - return self.subfunctions[i] + mixed = type(self.function_space().ufl_element()) is MixedElement + data = self.subfunctions if mixed else self._components + return data[i] def function_space(self): r"""Return the :class:`.FunctionSpace`, or :class:`.MixedFunctionSpace` @@ -321,6 +325,12 @@ def vector(self): :class:`Cofunction`""" return vector.Vector(self) + @property + def cell_set(self): + r"""The :class:`pyop2.types.set.Set` of cells for the mesh on which this + :class:`Cofunction` is defined.""" + return self.function_space()._mesh.cell_set + @property def node_set(self): r"""A :class:`pyop2.types.set.Set` containing the nodes of this @@ -330,6 +340,12 @@ def node_set(self): """ return self.function_space().node_set + @property + def dof_dset(self): + r"""A :class:`pyop2.types.dataset.DataSet` containing the degrees of freedom of + this :class:`Cofunction`.""" + return self.function_space().dof_dset + def ufl_id(self): return self.uid diff --git a/firedrake/function.py b/firedrake/function.py index da4d264971..97a08be816 100644 --- a/firedrake/function.py +++ b/firedrake/function.py @@ -15,6 +15,7 @@ from pyop2 import op2, mpi from pyop2.exceptions import DataTypeError, DataValueError +from finat.ufl import MixedElement from firedrake.utils import ScalarType, IntType, as_ctypes from firedrake import functionspaceimpl @@ -147,11 +148,8 @@ def sub(self, i): rank-n :class:`~.FunctionSpace`, this returns a proxy object indexing the ith component of the space, suitable for use in boundary condition application.""" - mixed = len(self.function_space()) != 1 + mixed = type(self.function_space().ufl_element()) is MixedElement data = self.subfunctions if mixed else self._components - bound = len(data) - if i < 0 or i >= bound: - raise IndexError(f"Invalid component {i}, not in [0, {bound})") return data[i] @property @@ -352,11 +350,8 @@ def sub(self, i): :func:`~.VectorFunctionSpace` or :func:`~.TensorFunctionSpace` this returns a proxy object indexing the ith component of the space, suitable for use in boundary condition application.""" - mixed = len(self.function_space()) != 1 + mixed = type(self.function_space().ufl_element()) is MixedElement data = self.subfunctions if mixed else self._components - bound = len(data) - if i < 0 or i >= bound: - raise IndexError(f"Invalid component {i}, not in [0, {bound})") return data[i] @PETSc.Log.EventDecorator() @@ -672,7 +667,7 @@ def single_eval(x, buf): value_shape = self.ufl_shape subfunctions = self.subfunctions - mixed = len(subfunctions) != 1 + mixed = type(self.function_space().ufl_element()) is MixedElement # Local evaluation l_result = [] diff --git a/firedrake/functionspaceimpl.py b/firedrake/functionspaceimpl.py index 8fc81244f7..c366c03b66 100644 --- a/firedrake/functionspaceimpl.py +++ b/firedrake/functionspaceimpl.py @@ -183,11 +183,8 @@ def _components(self): @PETSc.Log.EventDecorator() def sub(self, i): - mixed = len(self) != 1 + mixed = type(self.ufl_element()) is finat.ufl.MixedElement data = self.subfunctions if mixed else self._components - bound = len(data) - if i < 0 or i >= bound: - raise IndexError(f"Invalid component {i}, not in [0, {bound})") return data[i] @utils.cached_property @@ -664,9 +661,6 @@ def _components(self): def sub(self, i): r"""Return a view into the ith component.""" - bound = len(self._components) - if i < 0 or i >= bound: - raise IndexError(f"Invalid component {i}, not in [0, {bound})") return self._components[i] def __mul__(self, other): diff --git a/tests/firedrake/regression/test_vfs_component_bcs.py b/tests/firedrake/regression/test_vfs_component_bcs.py index 76cbd69ea5..6b6a0ea938 100644 --- a/tests/firedrake/regression/test_vfs_component_bcs.py +++ b/tests/firedrake/regression/test_vfs_component_bcs.py @@ -149,7 +149,7 @@ def test_cant_integrate_subscripted_VFS(V): @pytest.mark.parametrize("cmpt", - [-1, 2]) + [-3, 2]) def test_cant_subscript_outside_components(V, cmpt): with pytest.raises(IndexError): return V.sub(cmpt)