Merge branch 'connorjward/pyop3' into indiamai/fuse_pyop3_merge

Author: indiamai

.github/workflows/core.yml

@@ -1,4 +1,8 @@
 # Main Firedrake CI workflow run for pull requests or pushes to main or release
+<<<<<<< HEAD
+=======
+
+>>>>>>> connorjward/pyop3
 on:
   workflow_call:
     inputs:

firedrake/assemble.py

@@ -28,7 +28,7 @@
 from firedrake.ufl_expr import extract_domains
 from firedrake.bcs import DirichletBC, EquationBC, EquationBCSplit
 from firedrake.matrix import MatrixBase, Matrix, ImplicitMatrix
-from firedrake.functionspaceimpl import WithGeometry, FunctionSpace, FiredrakeDualSpace
+from firedrake.functionspaceimpl import WithGeometry, FunctionSpace, FiredrakeDualSpace, is_mixed
 from firedrake.interpolation import get_interpolator
 from firedrake.pack import pack, modified_lgmaps
 from firedrake.petsc import PETSc, local_submat
@@ -1092,6 +1092,11 @@ def assemble(self, tensor=None, current_state=None):
             else:
                 parloop(**{self._tensor_name[local_kernel]: subtensor}, compiler_parameters=pyop3_compiler_parameters)
 
+        # FIXME: This is necessary for test_submesh_solve_simple to pass for the moment
+        # This is unsatisfying because in theory this isn't required - we can stash up the
+        # pending increments and only apply them lazily. Something is going wrong somewhere.
+        subtensor.assemble()
+
         for bc in self._bcs:
             self._apply_bc(tensor, bc, u=current_state)
 
@@ -1508,7 +1513,7 @@ def _sub_mat_type(self) -> str | None:
     @staticmethod
     def _make_sparsity(test, trial, mat_spec, maps_and_regions):
         # Is this overly restrictive?
-        if any(len(a.function_space()) > 1 for a in [test, trial]) and mat_spec.mat_type == "baij":
+        if any(is_mixed(a) for a in [test, trial]) and mat_spec.mat_type == "baij":
             raise ValueError("BAIJ matrix type makes no sense for mixed spaces, use 'aij'")
 
         sparsity = op3.Mat.sparsity(
@@ -1526,7 +1531,7 @@ def _make_sparsity(test, trial, mat_spec, maps_and_regions):
         for loop_info, (test_index, trial_index) in maps_and_regions:
             # If indices are 'None' then this means all to allocate for all spaces
             if test_index is None:
-                if len(test.function_space()) > 1:
+                if is_mixed(test):
                     test_spaces = tuple(test.function_space())
                     test_indices = test.function_space().field_axis.component_labels
                 else:
@@ -1537,7 +1542,7 @@ def _make_sparsity(test, trial, mat_spec, maps_and_regions):
                 test_index = test.function_space().field_axis.component_labels[test_index]
                 test_indices = (test_index,)
             if trial_index is None:
-                if len(trial.function_space()) > 1:
+                if is_mixed(trial):
                     trial_spaces = tuple(trial.function_space())
                     trial_indices = trial.function_space().field_axis.component_labels
                 else:

firedrake/cofunction.py

@@ -113,7 +113,7 @@ def _analyze_form_arguments(self):
     def subfunctions(self):
         r"""Extract any sub :class:`Cofunction`\s defined on the component spaces
         of this this :class:`Cofunction`'s :class:`.FunctionSpace`."""
-        if len(self.function_space()) > 1:
+        if functionspaceimpl.is_mixed(self.function_space()):
             subfuncs = []
             for i, component in enumerate(self.dat.axes.trees[0].root.components):
                 subspace = self.function_space().sub(i)

firedrake/cython/dmcommon.pyx

@@ -2949,10 +2949,6 @@ def compute_dm_renumbering(
     # CHKERR(ISInvertPermutation(ordering_is.iset, -1, &renumbering_is.iset))
     # return renumbering_is
 
-    # return ordering_is.invertPermutation()
-    return ordering_is
-    # CHKERR(ISInvertPermutation(ordering_is.iset, -1, &renumbering_is.iset))
-    # return renumbering_is
 
 def partition_renumbering(PETSc.DM dm, PETSc.IS serial_new_to_old_renumbering) -> PETSc.IS:
     """Partition a serial point renumbering into owned and ghost points."""
@@ -4055,7 +4051,6 @@ cdef int DMPlexGetAdjacency_Facet_Support(PETSc.PetscDM dm,
             if numAdj > maxAdjSize:
                 CHKERR(PETSC_ERR_LIB)
     CHKERR(DMPlexRestoreTransitiveClosure(dm, point, PETSC_TRUE, &closureSize, &closure))
-    CHKERR(DMPlexRestoreTransitiveClosure(dm, p, PETSC_FALSE, &starSize, &star))
     adjSize[0] = numAdj
     return 0
 
@@ -4422,11 +4417,13 @@ def submesh_create(PETSc.DM dm,
                                              ignoreHalo=ignore_label_halo,
                                              sanitizeSubMesh=PETSC_TRUE,
                                              comm=comm)
+
     # Destroy temp_label.
     dm.removeLabel(temp_label_name)
     subdm.removeLabel(temp_label_name)
     submesh_update_facet_labels(dm, subdm)
     submesh_correct_entity_classes(dm, subdm, ownership_transfer_sf)
+
     return subdm
 
 
@@ -4466,6 +4463,7 @@ def submesh_correct_entity_classes(PETSc.DM dm,
     CHKERR(DMPlexGetChart(subdm.dm, &subpStart, &subpEnd))
     assert pStart == 0
     assert subpStart == 0
+
     CHKERR(DMGetLabel(subdm.dm, b"firedrake_is_ghost", &is_ghost))
     CHKERR(DMLabelCreateIndex(is_ghost, subpStart, subpEnd))
 
@@ -4499,7 +4497,7 @@ def submesh_correct_entity_classes(PETSc.DM dm,
             if ownership_gain[p] == 1:
                 CHKERR(DMLabelHasPoint(is_ghost, subp, &has))
                 assert has
-                CHKERR(DMLabelSetValue(is_ghost, subp, 0))
+                CHKERR(DMLabelClearValue(is_ghost, subp, 1))
 
         CHKERR(ISRestoreIndices(subpoint_is.iset, &subpoint_indices))
     CHKERR(DMLabelDestroyIndex(is_ghost))

firedrake/function.py

@@ -33,7 +33,7 @@
 from firedrake import utils
 from firedrake.adjoint_utils import FunctionMixin
 from firedrake.petsc import PETSc
-from firedrake.functionspaceimpl import MixedFunctionSpace, parse_component_indices
+from firedrake.functionspaceimpl import MixedFunctionSpace, parse_component_indices, is_mixed
 from firedrake.mesh import MeshGeometry, VertexOnlyMesh, extract_mesh_topologies
 from firedrake.functionspace import FunctionSpace, VectorFunctionSpace, TensorFunctionSpace
 from firedrake.exceptions import PointNotInDomainError
@@ -118,7 +118,7 @@ def ufl_id(self):
     def subfunctions(self):
         r"""Extract any sub :class:`Function`\s defined on the component spaces
         of this this :class:`Function`'s :class:`.FunctionSpace`."""
-        if isinstance(self.function_space(), MixedFunctionSpace):
+        if is_mixed(self.function_space()):
             # NOTE: This is quite tricky for fieldsplit. Previously the fields would
             # be renumbered when split, but now we retain the labels in the dat but
             # not the function space.
@@ -317,7 +317,7 @@ def __dir__(self):
     def subfunctions(self):
         r"""Extract any sub :class:`Function`\s defined on the component spaces
         of this this :class:`Function`'s :class:`.FunctionSpace`."""
-        if isinstance(self.function_space().topological, MixedFunctionSpace):
+        if is_mixed(self.function_space()):
             return tuple(
                 type(self)(self.function_space().sub(i), val)
                 for (i, val) in zip(range(len(self.function_space())), self.topological.subfunctions))
@@ -355,7 +355,7 @@ def sub(self, indices: tuple[int] | int) -> "Function":
         subfunctions
 
         """
-        if type(self.function_space().ufl_element()) is MixedElement:
+        if is_mixed(self.function_space()):
             return self.subfunctions[indices]
         elif not self.function_space().shape:
             # TODO: Decide if this is acceptable usage
@@ -538,12 +538,17 @@ def __itruediv__(self, expr):
         return self
 
     def __float__(self):
-
         if (
             self.ufl_element().family() == "Real"
             and self.function_space().shape == ()
         ):
-            return float(self.dat.data_ro[0])
+            self.dat.assemble()
+            with op3.mpi.temp_internal_comm(self.comm) as icomm:
+                if icomm.rank == 0:
+                    value = icomm.bcast(utils.just_one(self.dat.data_ro))
+                else:
+                    value = icomm.bcast(None)
+            return float(value)
         else:
             raise ValueError("Can only cast scalar 'Real' Functions to float.")
 

firedrake/functionspaceimpl.py

@@ -660,7 +660,10 @@ def local_section(self):
     @cached_property
     @deprecated("axes.template_vec")
     def template_vec(self):
-        block_shape = self.shape if len(self) == 1 else ()
+        if is_mixed(self):
+            block_shape = ()
+        else:
+            block_shape = self.shape
         return self.layout_axes.template_vec(block_shape)
 
     @cached_method()
@@ -682,7 +685,7 @@ def lgmap(self, bcs: Iterable[DirichletBC] = (), index: int | None = None) -> PE
 
         """
         lgmap_axes = self.axes
-        if len(self) > 1 or any(bc.function_space().component is not None for bc in bcs):
+        if is_mixed(self) or any(bc.function_space().component is not None for bc in bcs):
             block_size = 1
         else:
             lgmap_axes = lgmap_axes.blocked(self.shape)
@@ -692,7 +695,7 @@ def lgmap(self, bcs: Iterable[DirichletBC] = (), index: int | None = None) -> PE
         # track which BCs are used so we can warn if any are missed
         unused_bcs = set(bcs)
         if index is None:  # The lgmap is for the full space
-            if len(self) > 1:
+            if is_mixed(self):
                 split_bcs = []
                 for subspace in self:
                     matching_bcs = []
@@ -781,7 +784,7 @@ def interior_facet_node_map_dat(self) -> op3.Dat:
     @property
     @deprecated("cell_node_map_dat.data_ro")
     def cell_node_list(self) -> numpy.ndarray:
-        if len(self) > 1 or self.parent:
+        if is_mixed(self) or self.parent:
             warnings.warn(
                 "For mixed spaces it is no longer the case that offset=node*bsize, "
                 "use a DoF list instead."
@@ -791,7 +794,7 @@ def cell_node_list(self) -> numpy.ndarray:
     @property
     @deprecated("exterior_facet_node_map_dat.data_ro")
     def exterior_facet_node_list(self) -> numpy.ndarray:
-        if len(self) > 1 or self.parent:
+        if is_mixed(self) or self.parent:
             warnings.warn(
                 "For mixed spaces it is no longer the case that offset=node*bsize, "
                 "use a DoF list instead."
@@ -801,7 +804,7 @@ def exterior_facet_node_list(self) -> numpy.ndarray:
     @property
     @deprecated("interior_facet_node_map_dat.data_ro")
     def interior_facet_node_list(self) -> numpy.ndarray:
-        if len(self) > 1 or self.parent:
+        if is_mixed(self) or self.parent:
             warnings.warn(
                 "For mixed spaces it is no longer the case that offset=node*bsize "
                 "because the strides between nodes are not constant as they are "
@@ -898,7 +901,7 @@ def _iterset_to_dof_map_dat(
         #     └──➤ {field: [{functionspace0: 1}, {functionspace1: 1}]}
         #          ├──➤ {some_label: 2}
         #          └──➤ {some_label: 1}
-        if len(self) > 1:
+        if is_mixed(self):
             map_dof_axes = iterset_axes.add_axis(None, packed_offsets.axes.root)
             for label, subspace in zip(self._labels, self):
                 path = iterset_axes.leaf_path | {"field": label}
@@ -1606,7 +1609,7 @@ def entity_node_map(self, iteration_spec):
             Entity node map.
 
         """
-        if len(self) > 1:
+        if is_mixed(self):
             raise NotImplementedError("will this work?")
 
         iter_mesh = iteration_spec.mesh
@@ -2558,6 +2561,96 @@ def _(index: int, shape: tuple[int, ...]) -> tuple[int, ...]:
 
 
 
+def merge_axis_constraints(root_axis: op3.Axis, axis_constraintss: Sequence[Sequence[AxisConstraint]]) -> tuple[AxisConstraint]:
+    # start by collecting like axes
+    axis_info: defaultdict[op3.Axis, dict[op3.ComponentLabelT, idict]] = defaultdict(dict)
+    for root_component, constraints in zip(root_axis.components, axis_constraintss, strict=True):
+        for constraint in constraints:
+            axis_info[constraint.axis][root_component.label] = constraint.within_axes
+
+    # Now build the new set of constraints. To do this we inspect the
+    # per-component constraints for each axis: if the constraints are all the
+    # same then it is not necessary to specialise by component, otherwise an
+    # extra constraint is needed. For example:
+    #
+    # * Consider the "dof" axis for a mixed space with identical subspaces:
+    #
+    #     {dof_axis: {0: {"mesh": None}, 1: {"mesh": None}}}
+    #
+    #   Here this is saying that 'dof_axis' exists under root components 0 and 1
+    #   and each time must satisfy the constraint of having '{"mesh": None}'
+    #   above them.
+    #
+    #   Since the constraints are identical for all components they do not need
+    #   to be specialised. The final constraint is thus:
+    #
+    #     AxisConstraint(dof_axis, {"mesh": None})
+    #
+    # * Alternatively consider a mixed space of CG1 x Real:
+    #
+    #     {mesh_axis: {0: {}}}
+    #
+    #   The "mesh" axis only exists for the CG1 subspace and so a new constraint
+    #   is needed:
+    #
+    #     AxisConstraint(mesh_axis, {"field": 0})
+    constraints = [AxisConstraint(root_axis)]
+    for axis, per_component_info in axis_info.items():
+        if (
+            per_component_info.keys() == set(root_axis.component_labels)
+            and utils.is_single_valued(per_component_info.values())
+        ):
+            # Axis present for all components and constraints match: use as is
+            within_axes = utils.single_valued(per_component_info.values())
+            constraints.append(AxisConstraint(axis, within_axes))
+        else:
+            # Constraint mismatch: need to specialise by component
+            for component_label, orig_within_axes in per_component_info.items():
+                within_axes = orig_within_axes | {root_axis.label: component_label}
+                constraints.append(AxisConstraint(axis, within_axes))
+    return tuple(constraints)
+
+
+@functools.singledispatch
+def parse_component_indices(indices: Any, shape: tuple[int, ...]) -> tuple[int, ...]:
+    raise TypeError
+
+
+@parse_component_indices.register(tuple)
+def _(indices: tuple[int, ...], shape: tuple[int, ...]) -> tuple[int, ...]:
+    return indices
+
+
+@parse_component_indices.register(int)
+def _(index: int, shape: tuple[int, ...]) -> tuple[int, ...]:
+    # Historically tensor-valued spaces would be addressed using a flat index
+    # instead of a tuple. Here we convert the old-style flat index to a
+    # nested one. Eventually we should be able to remove this and simply cast
+    # an integer index to a tuple (e.g. '3' to '(3,)').
+    if len(shape) > 1:
+        warnings.warn(
+            "Scalar indexing of a tensor-valued space is no longer recommended "
+            "practice, please pass a tuple instead",
+            FutureWarning,
+        )
+    return list(numpy.ndindex(shape))[index]
+
+    for component in axis.components:
+        subconstraints_ = tuple(
+            subconstraint
+            for subconstraint in subconstraints
+            if axis.label not in subconstraint.within_axes
+            or (axis.label, component.label) in subconstraint.within_axes.items()
+        )
+        if subconstraints_:
+            # FIXME: Not doing anything with visited_axes
+            subnest = _axis_nest_from_constraints(subconstraints_, visited_axes)
+            axis_nest[axis].append(subnest)
+
+    return idict(axis_nest) if axis_nest else axis
+
+
+
 def merge_axis_constraints(root_axis: op3.Axis, axis_constraintss: Sequence[Sequence[AxisConstraint]]) -> tuple[AxisConstraint]:
     # start by collecting like axes
     axis_info: defaultdict[op3.Axis, dict[op3.ComponentLabelT, idict]] = defaultdict(dict)
@@ -2667,3 +2760,23 @@ def entity_permutations_key(entity_permutations):
         key.append(tuple(sub_key))
     key = tuple(key)
     return key
+
+
+@functools.singledispatch
+def is_mixed(obj) -> bool:
+    raise TypeError
+
+
+@is_mixed.register
+def _(elem: finat.ufl.FiniteElementBase) -> bool:
+    return type(elem) is finat.ufl.MixedElement
+
+
+@is_mixed.register
+def _(V: WithGeometryBase | AbstractFunctionSpace) -> bool:
+    return is_mixed(V.ufl_element())
+
+
+@is_mixed.register
+def _(arg: ufl.Argument, /) -> bool:
+    return is_mixed(arg.ufl_function_space())