Indiamai/fuse mesh cell

.github/actions/install/action.yml

@@ -156,7 +156,7 @@ runs:
       with:
         repository: firedrakeproject/fiat
         path: fiat-repo
-        ref: refs/heads/indiamai/integrate_fuse
+        ref: refs/heads/indiamai/fuse_cell_change
 
     - name: Install checked out FIAT
       shell: bash

.github/workflows/core.yml

@@ -139,14 +139,12 @@ jobs:
         # UNDO ME
         timeout-minutes: 180
 
-      - name: Run tests (nprocs = 2)
-        if: success() || steps.install.conclusion == 'success'
-        run: |
-          . venv/bin/activate
-          firedrake-run-split-tests 2 4 "$EXTRA_PYTEST_ARGS" firedrake-repo/tests/firedrake
-        timeout-minutes: 60
-
-     # - name: Run tests (nprocs = 3)
+     # - name: Run tests (nprocs = 2)
+     #   if: success() || steps.install.conclusion == 'success'
+     #   run: |
+     #     . venv/bin/activate
+     #     firedrake-run-split-tests 2 4 "$EXTRA_PYTEST_ARGS" firedrake-repo/tests/firedrake
+     #   timeout-minutes: 60
      #   if: success() || steps.install.conclusion == 'success'
      #   run: |
      #     . venv/bin/activate
@@ -210,7 +208,8 @@ jobs:
      #   if: success() || steps.install.conclusion == 'success'
      #   run: |
      #     . venv/bin/activate
-     #     git clone --depth 1 https://github.com/thetisproject/thetis.git thetis-repo --branch ${{ inputs.base_ref }}
+     #     : # UNDO ME
+     #     git clone --depth 1 https://github.com/thetisproject/thetis.git thetis-repo --branch connorjward/pyop3
      #     pip install --verbose ./thetis-repo
      #     python -m pytest -n 8 --verbose thetis-repo/test_adjoint/test_swe_adjoint.py
      #   timeout-minutes: 10
@@ -228,7 +227,8 @@ jobs:
      #   if: success() || steps.install.conclusion == 'success'
      #   run: |
      #     . venv/bin/activate
-     #     git clone --depth 1 https://github.com/g-adopt/g-adopt.git g-adopt-repo --branch ${{ inputs.base_ref }}
+     #     : # UNDO ME
+     #     git clone --depth 1 https://github.com/connorjward/g-adopt.git g-adopt-repo --branch connorjward/pyop3
      #     pip install --verbose ./g-adopt-repo
      #     make -C g-adopt-repo/demos/mantle_convection/base_case check
      #   timeout-minutes: 5

firedrake/checkpointing.py

@@ -1520,6 +1520,8 @@ def _load_function_topology(self, tmesh, element, tf_name, idx=None):
             self.viewer.pushTimestepping()
             self.viewer.setTimestep(idx)
         if element.family() == "Real":
+            if tV.comm.size > 1:
+                raise NotImplementedError
             assert not isinstance(element, (finat.ufl.VectorElement, finat.ufl.TensorElement))
             value = self.get_attr(path, "_".join([PREFIX, "value" if idx is None else "value_" + str(idx)]))
             tf.dat.data[...] = value

firedrake/cython/dmcommon.pyx

@@ -2587,65 +2587,6 @@ def mark_entity_classes_using_cell_dm(PETSc.DM swarm):
     CHKERR(PetscFree(plex_cell_classes))
 
 
-@cython.boundscheck(False)
-@cython.wraparound(False)
-def get_cell_markers(PETSc.DM dm, np.ndarray cell_numbering,
-                     subdomain_id):
-    """Get the cells marked by a given subdomain_id.
-
-    :arg dm: The DM for the mesh topology
-    :arg cell_numbering: Section mapping dm cell points to firedrake cell indices.
-    :arg subdomain_id: The subdomain_id to look for.
-
-    :raises ValueError: if the subdomain_id is not valid.
-    :returns: A numpy array (possibly empty) of the cell ids.
-    """
-    cdef:
-        PetscInt i, j, n, offset, c, cStart, cEnd, ncells
-        np.ndarray cells
-        np.ndarray indices
-
-    if not dm.hasLabel(CELL_SETS_LABEL):
-        return np.empty(0, dtype=IntType)
-    vals = dm.getLabelIdIS(CELL_SETS_LABEL).indices
-    comm = dm.comm.tompi4py()
-
-    def merge_ids(x, y, datatype):
-        return x.union(y)
-
-    op = MPI.Op.Create(merge_ids, commute=True)
-
-    all_ids = np.asarray(sorted(comm.allreduce(set(vals), op=op)),
-                         dtype=IntType)
-    op.Free()
-    if subdomain_id not in all_ids:
-        raise ValueError("Invalid subdomain_id %d not in %s" % (subdomain_id, vals))
-
-    if subdomain_id not in vals:
-        return np.empty(0, dtype=IntType)
-
-    n = dm.getStratumSize(CELL_SETS_LABEL, subdomain_id)
-    if n == 0:
-        return np.empty(0, dtype=IntType)
-    indices = dm.getStratumIS(CELL_SETS_LABEL, subdomain_id).indices
-    ncells = 0
-    get_height_stratum(dm.dm, 0, &cStart, &cEnd)
-    for i in range(n):
-        c = indices[i]
-        if cStart <= c < cEnd:
-            ncells += 1
-    if ncells == 0:
-        return np.empty(0, dtype=IntType)
-    cells = np.empty(ncells, dtype=IntType)
-    j = 0
-    for i in range(n):
-        c = indices[i]
-        if cStart <= c < cEnd:
-            cells[j] = cell_numbering[c]
-            j += 1
-    return cells
-
-
 @cython.boundscheck(False)
 @cython.wraparound(False)
 def get_facet_ordering(PETSc.DM plex, PETSc.Section facet_numbering):
@@ -4748,8 +4689,8 @@ def filter_is(is_: PETSc.IS, start: IntType, end: IntType) -> PETSc.IS:
     cdef:
         PETSc.IS filtered_is
 
-    # empty ISes remain empty
-    if not is_ or is_.size == 0:
+    # cast null ISes to empty ones
+    if not is_:
         return PETSc.IS().createGeneral(np.empty(0, dtype=IntType))
 
     filtered_is = is_.duplicate()

firedrake/extrusion_utils.py

@@ -16,7 +16,7 @@
 
 
 @PETSc.Log.EventDecorator()
-@with_heavy_caches(lambda extr_top, *a, **kw: {extr_top})
+@with_heavy_caches(lambda extr_top, *a, **kw: [extr_top])
 def make_extruded_coords(extruded_topology, base_coords, ext_coords,
                          layer_height, extrusion_type='uniform', kernel=None):
     """
@@ -69,7 +69,8 @@ def make_extruded_coords(extruded_topology, base_coords, ext_coords,
         layer_height = numpy.cumsum(numpy.concatenate(([0], layer_height)))
 
     layer_heights = layer_height.size
-    layer_height = op3.Dat.from_array(layer_height)
+    # NOTE: Not sure about this in parallel, needs an SF?
+    layer_height = op3.Dat.from_array(layer_height, comm=extruded_topology.comm)
 
     if kernel is not None:
         raise NotImplementedError
@@ -319,7 +320,7 @@ def is_real_tensor_product_element(element):
     assert not isinstance(element, finat.TensorFiniteElement)
 
     if isinstance(element, finat.TensorProductElement):
-        _, factor = element.factors
+        factor = element.factors[-1]
         return isinstance(factor, finat.Real)
     else:
         return False

firedrake/functionspace.py

@@ -6,7 +6,6 @@
 """
 import itertools
 import ufl
-from ufl.cell import as_cell
 import finat.ufl
 
 from pyop3.pyop2_utils import flatten
@@ -69,7 +68,7 @@ def make_scalar_element(mesh, family, degree, vfamily, vdegree, variant, quad_sc
                                      quad_scheme=quad_scheme)
         # If second element was passed in, use it
         lb = finat.ufl.FiniteElement(vfamily,
-                                     cell=ufl.interval,
+                                     cell=finat.ufl.as_cell("interval"),
                                      degree=vdegree,
                                      variant=variant,
                                      quad_scheme=quad_scheme)

firedrake/functionspaceimpl.py

@@ -49,7 +49,7 @@
 from firedrake.utils import IntType, deprecated
 
 
-def check_element(element, top=True):
+def check_element(element, top=True, use_fuse=False):
     """Run some checks on the provided element.
 
     The :class:`finat.ufl.mixedelement.VectorElement` and
@@ -82,7 +82,7 @@ def check_element(element, top=True):
        type(element) is not finat.ufl.MixedElement:
         raise ValueError("MixedElement modifier must be outermost")
     if element.cell.cellname == "hexahedron" and \
-       element.family() not in ["Q", "DQ", "Real"]:
+       element.family() not in ["Q", "DQ", "Real", "IT"] and not use_fuse:
         raise NotImplementedError("Currently can only use 'Q', 'DQ', and/or 'Real' elements on hexahedral meshes, not", element.family())
     if type(element) in (finat.ufl.BrokenElement, finat.ufl.RestrictedElement,
                          finat.ufl.HDivElement, finat.ufl.HCurlElement):
@@ -99,11 +99,11 @@ def check_element(element, top=True):
     else:
         inner = ()
     for e in inner:
-        check_element(e, top=False)
+        check_element(e, top=False, use_fuse=use_fuse)
 
 
-def create_element(ufl_element):
-    finat_element = _create_element(ufl_element)
+def create_element(ufl_element, use_fuse=False):
+    finat_element = _create_element(ufl_element, use_fuse=use_fuse)
     if isinstance(finat_element, finat.TensorFiniteElement):
         # Retrieve scalar element
         finat_element = finat_element.base_element
@@ -482,6 +482,8 @@ def collapse(self):
     def make_function_space(cls, mesh, element, name=None, _labels=None, **kwargs):
         r"""Factory method for :class:`.WithGeometryBase`."""
         topology = mesh.topology
+        if hasattr(element, 'triple') and not mesh._use_fuse:
+            raise NotImplementedError("FUSE defined element not on fuse enabled mesh.")
         # Create a new abstract (Mixed/Real)FunctionSpace, these are neither primal nor dual.
         if type(element) is finat.ufl.MixedElement:
             if isinstance(mesh, MeshGeometry):
@@ -497,7 +499,7 @@ def make_function_space(cls, mesh, element, name=None, _labels=None, **kwargs):
             if isinstance(mesh, MeshSequenceGeometry):
                 raise TypeError(f"mesh must not be MeshSequenceGeometry: got {mesh}")
             # Check that any Vector/Tensor/Mixed modifiers are outermost.
-            check_element(element)
+            check_element(element, use_fuse=mesh._use_fuse)
             if element.family() == "Real":
                 new = RealFunctionSpace(topology, element, name=name, **kwargs)
             else:
@@ -1078,7 +1080,7 @@ def __init__(self, mesh, element, name=None, *, layout=None):
         self.comm = mesh.comm
 
         self.element = element
-        self.finat_element = create_element(element)
+        self.finat_element = create_element(element, use_fuse=mesh._use_fuse)
 
         entity_dofs = self.finat_element.entity_dofs()
         nodes_per_entity = tuple(len(entity_dofs[d][0]) for d in sorted(entity_dofs))

firedrake/interpolation.py

@@ -7,9 +7,10 @@
 import tempfile
 import abc
 
+import dataclasses
 from functools import cached_property, partial
 from typing import Hashable, Literal, Callable, Iterable
-from dataclasses import asdict, dataclass
+from dataclasses import dataclass
 from numbers import Number
 
 from ufl.algorithms import extract_arguments, replace
@@ -151,7 +152,13 @@ def __init__(self, expr: Expr, V: WithGeometry | BaseForm, **kwargs):
     def _ufl_expr_reconstruct_(
             self, expr: Expr, v: WithGeometry | BaseForm | None = None, **interp_data
     ):
-        interp_data = interp_data or asdict(self.options)
+        # Note that we can't use dataclasses.asdict here because we can't deepcopy
+        # PETSc objects, this is the recommended workaround.
+        options = {
+            field.name: getattr(self.options, field.name)
+            for field in dataclasses.fields(self.options)
+        }
+        interp_data = options | interp_data
         return UFLInterpolate._ufl_expr_reconstruct_(self, expr, v=v, **interp_data)
 
     @property
@@ -1073,12 +1080,10 @@ def _build_interpolation_callables(
         target_element = runtime_quadrature_element(source_mesh, target_element,
                                                     rt_var_name=rt_var_name)
 
-    iter_spec = get_iteration_spec(target_mesh, "cell")
-
-    if not (subset is None or subset is Ellipsis):
-        raise NotImplementedError
-        assert subset.superset == cell_set
-        cell_set = subset
+    if subset is not None:
+        iter_spec = get_iteration_spec(target_mesh, "cell", subdomain_id=subset)
+    else:
+        iter_spec = get_iteration_spec(target_mesh, "cell")
 
     parameters = {}
     parameters['scalar_type'] = ScalarType
@@ -1667,8 +1672,12 @@ def _build_matnest(
         """Return a PETSc nested matrix built from sub-interpolator matrices."""
         shape = tuple(len(a.function_space()) for a in self.interpolate_args)
         blocks = numpy.full(shape, PETSc.Mat(), dtype=object)
-        for indices, (interp, sub_bcs) in Isub.items():
-            blocks[indices] = interp._get_callable(bcs=sub_bcs, mat_type=sub_mat_type)()
+        for (ridx, cidx), (interp, sub_bcs) in Isub.items():
+            if ridx is None:
+                ridx = 0
+            if cidx is None:
+                cidx = 0
+            blocks[ridx, cidx] = interp._get_callable(bcs=sub_bcs, mat_type=sub_mat_type)()
         return PETSc.Mat().createNest(blocks)
 
     def _build_aij(