Merge branch 'main' into pbrubeck/matis

Author: pbrubeck

firedrake/assemble.py

@@ -1666,6 +1666,7 @@ def __init__(self, form, local_knl, subdomain_id, all_integer_subdomain_ids, dia
         self._constants = _FormHandler.iter_constants(form, local_knl.kinfo)
         self._active_exterior_facets = _FormHandler.iter_active_exterior_facets(form, local_knl.kinfo)
         self._active_interior_facets = _FormHandler.iter_active_interior_facets(form, local_knl.kinfo)
+        self._active_orientations_cell = _FormHandler.iter_active_orientations_cell(form, local_knl.kinfo)
         self._active_orientations_exterior_facet = _FormHandler.iter_active_orientations_exterior_facet(form, local_knl.kinfo)
         self._active_orientations_interior_facet = _FormHandler.iter_active_orientations_interior_facet(form, local_knl.kinfo)
 
@@ -1688,6 +1689,7 @@ def build(self):
         assert_empty(self._constants)
         assert_empty(self._active_exterior_facets)
         assert_empty(self._active_interior_facets)
+        assert_empty(self._active_orientations_cell)
         assert_empty(self._active_orientations_exterior_facet)
         assert_empty(self._active_orientations_interior_facet)
 
@@ -1885,6 +1887,17 @@ def _as_global_kernel_arg_interior_facet(_, self):
         return op2.DatKernelArg((2,), m._global_kernel_arg)
 
 
+@_as_global_kernel_arg.register(kernel_args.OrientationsCellKernelArg)
+def _(_, self):
+    mesh = next(self._active_orientations_cell)
+    if mesh is self._mesh:
+        return op2.DatKernelArg((1,))
+    else:
+        m, integral_type = mesh.topology.trans_mesh_entity_map(self._mesh.topology, self._integral_type, self._subdomain_id, self._all_integer_subdomain_ids)
+        assert integral_type == "cell"
+        return op2.DatKernelArg((1,), m._global_kernel_arg)
+
+
 @_as_global_kernel_arg.register(kernel_args.OrientationsExteriorFacetKernelArg)
 def _(_, self):
     mesh = next(self._active_orientations_exterior_facet)
@@ -1956,6 +1969,7 @@ def __init__(self, form, bcs, local_knl, subdomain_id,
         self._constants = _FormHandler.iter_constants(form, local_knl.kinfo)
         self._active_exterior_facets = _FormHandler.iter_active_exterior_facets(form, local_knl.kinfo)
         self._active_interior_facets = _FormHandler.iter_active_interior_facets(form, local_knl.kinfo)
+        self._active_orientations_cell = _FormHandler.iter_active_orientations_cell(form, local_knl.kinfo)
         self._active_orientations_exterior_facet = _FormHandler.iter_active_orientations_exterior_facet(form, local_knl.kinfo)
         self._active_orientations_interior_facet = _FormHandler.iter_active_orientations_interior_facet(form, local_knl.kinfo)
 
@@ -2223,6 +2237,17 @@ def _as_parloop_arg_interior_facet(_, self):
     return op2.DatParloopArg(mesh.interior_facets.local_facet_dat, m)
 
 
+@_as_parloop_arg.register(kernel_args.OrientationsCellKernelArg)
+def _(_, self):
+    mesh = next(self._active_orientations_cell)
+    if mesh is self._mesh:
+        m = None
+    else:
+        m, integral_type = mesh.topology.trans_mesh_entity_map(self._mesh.topology, self._integral_type, self._subdomain_id, self._all_integer_subdomain_ids)
+        assert integral_type == "cell"
+    return op2.DatParloopArg(mesh.local_cell_orientation_dat, m)
+
+
 @_as_parloop_arg.register(kernel_args.OrientationsExteriorFacetKernelArg)
 def _(_, self):
     mesh = next(self._active_orientations_exterior_facet)
@@ -2319,6 +2344,14 @@ def iter_active_interior_facets(form, kinfo):
             mesh = all_meshes[i]
             yield mesh
 
+    @staticmethod
+    def iter_active_orientations_cell(form, kinfo):
+        """Yield the form cell orientations referenced in ``kinfo``."""
+        all_meshes = extract_domains(form)
+        for i in kinfo.active_domain_numbers.orientations_cell:
+            mesh = all_meshes[i]
+            yield mesh
+
     @staticmethod
     def iter_active_orientations_exterior_facet(form, kinfo):
         """Yield the form exterior facet orientations referenced in ``kinfo``."""

firedrake/assign.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 from pyadjoint.tape import annotate_tape
+from pyop2 import op2
 from pyop2.utils import cached_property
 import pytools
 import finat.ufl
@@ -12,19 +13,22 @@
 from ufl.corealg.multifunction import MultiFunction
 from ufl.domain import extract_unique_domain
 
+from firedrake.cofunction import Cofunction
 from firedrake.constant import Constant
 from firedrake.function import Function
 from firedrake.petsc import PETSc
 from firedrake.utils import ScalarType, split_by
 
+from mpi4py import MPI
+
 
 def _isconstant(expr):
     return isinstance(expr, Constant) or \
-        (isinstance(expr, Function) and expr.ufl_element().family() == "Real")
+        (isinstance(expr, (Function, Cofunction)) and expr.ufl_element().family() == "Real")
 
 
 def _isfunction(expr):
-    return isinstance(expr, Function) and expr.ufl_element().family() != "Real"
+    return isinstance(expr, (Function, Cofunction)) and expr.ufl_element().family() != "Real"
 
 
 class CoefficientCollector(MultiFunction):
@@ -99,6 +103,9 @@ def component_tensor(self, o, a, _):
     def coefficient(self, o):
         return ((o, 1),)
 
+    def cofunction(self, o):
+        return ((o, 1),)
+
     def constant_value(self, o):
         return ((o, 1),)
 
@@ -130,34 +137,56 @@ def _as_scalar(self, weighted_coefficients):
 
 
 class Assigner:
-    """Class performing pointwise assignment of an expression to a :class:`firedrake.function.Function`.
+    """Class performing pointwise assignment of an expression to a function or a cofunction.
+
+    Parameters
+    ----------
+    assignee : firedrake.function.Function or firedrake.cofunction.Cofunction
+        Function or Cofunction being assigned to.
+    expression : ufl.core.expr.Expr or ufl.form.BaseForm
+        Expression to be assigned.
+    subset : pyop2.types.set.Set or pyop2.types.set.Subset or pyop2.types.set.MixedSet
+        Subset to apply the assignment over.
 
-    :param assignee: The :class:`~.firedrake.function.Function` being assigned to.
-    :param expression: The :class:`ufl.core.expr.Expr` to evaluate.
-    :param subset: Optional subset (:class:`pyop2.types.set.Subset`) to apply the assignment over.
     """
     symbol = "="
 
     _coefficient_collector = CoefficientCollector()
 
     def __init__(self, assignee, expression, subset=None):
         expression = as_ufl(expression)
-
+        source_meshes = set()
         for coeff in extract_coefficients(expression):
-            if isinstance(coeff, Function) and coeff.ufl_element().family() != "Real":
+            if isinstance(coeff, (Function, Cofunction)) and coeff.ufl_element().family() != "Real":
                 if coeff.ufl_element() != assignee.ufl_element():
                     raise ValueError("All functions in the expression must have the same "
                                      "element as the assignee")
-                if extract_unique_domain(coeff) != extract_unique_domain(assignee):
-                    raise ValueError("All functions in the expression must use the same "
-                                     "mesh as the assignee")
-
-        if (subset and type(assignee.ufl_element()) == finat.ufl.MixedElement
-                and any(el.family() == "Real"
-                        for el in assignee.ufl_element().sub_elements)):
-            raise ValueError("Subset is not a valid argument for assigning to a mixed "
-                             "element including a real element")
-
+                source_meshes.add(extract_unique_domain(coeff))
+        if len(source_meshes) == 0:
+            pass
+        elif len(source_meshes) == 1:
+            target_mesh = extract_unique_domain(assignee)
+            source_mesh, = source_meshes
+            if target_mesh.submesh_youngest_common_ancester(source_mesh) is None:
+                raise ValueError(
+                    "All functions in the expression must be defined on a single domain "
+                    "that is in the same submesh family as domain of the assignee"
+                )
+        else:
+            raise ValueError(
+                "All functions in the expression must be defined on a single domain"
+            )
+        if subset is None:
+            subset = tuple(None for _ in assignee.function_space())
+        if len(subset) != len(assignee.function_space()):
+            raise ValueError(f"Provided subset ({subset}) incompatible with assignee ({assignee})")
+        if type(assignee.ufl_element()) == finat.ufl.MixedElement:
+            for subs, el in zip(subset, assignee.function_space().ufl_element().sub_elements):
+                if subs is not None and el.family() == "Real":
+                    raise ValueError(
+                        "Subset is not a valid argument for assigning to a mixed "
+                        "element including a real element"
+                    )
         self._assignee = assignee
         self._expression = expression
         self._subset = subset
@@ -169,14 +198,21 @@ def __repr__(self):
         return f"{self.__class__.__name__}({self._assignee!r}, {self._expression!r})"
 
     @PETSc.Log.EventDecorator()
-    def assign(self):
-        """Perform the assignment."""
+    def assign(self, allow_missing_dofs=False):
+        """Perform the assignment.
+
+        Parameters
+        ----------
+        allow_missing_dofs : bool
+            Permit assignment between objects with mismatching nodes. If `True` then
+            assignee nodes with no matching assigner nodes are ignored.
+
+        """
         if annotate_tape():
             raise NotImplementedError(
                 "Taping with explicit Assigner objects is not supported yet. "
                 "Use Function.assign instead."
             )
-
         # To minimize communication during assignment we perform a number of tricks:
         # * If we are not assigning to a subset then we can always write to the
         #   halo. The validity of the original assignee dat halo does not matter
@@ -191,28 +227,111 @@ def assign(self):
         #   end up doing a lot of halo exchanges for the expression just to avoid
         #   a single halo exchange for the assignee.
         # * If we do write to the halo then the resulting halo will never be dirty.
-
-        func_halos_valid = all(f.dat.halo_valid for f in self._functions)
-        assign_to_halos = (
-            func_halos_valid and (not self._subset or self._assignee.dat.halo_valid))
-
+        # If mixed, loop over individual components
+        for lhs_func, subset, *funcs in zip(self._assignee.subfunctions, self._subset, *(f.subfunctions for f in self._functions)):
+            target_mesh = extract_unique_domain(lhs_func)
+            target_V = lhs_func.function_space()
+            # Validate / Process subset.
+            if subset is not None:
+                if subset is target_V.node_set:
+                    # The whole set.
+                    subset = None
+                elif subset.superset is target_V.node_set:
+                    # op2.Subset of target_V.node_set
+                    pass
+                else:
+                    raise ValueError(f"subset ({subset}) not a subset of target_V.node_set ({target_V.node_set})")
+            source_meshes = set(extract_unique_domain(f) for f in funcs)
+            if len(source_meshes) == 0:
+                # Assign constants only.
+                single_mesh_assign = True
+            elif len(source_meshes) == 1:
+                source_mesh, = source_meshes
+                if target_mesh is source_mesh:
+                    # Assign (co)functions from one mesh to the same mesh.
+                    single_mesh_assign = True
+                else:
+                    # Assign (co)functions between a submesh and the parent or between two submeshes.
+                    single_mesh_assign = False
+            else:
+                raise ValueError("All functions in the expression must be defined on a single domain")
+            if single_mesh_assign:
+                self._assign_single_mesh(lhs_func, subset, funcs, operator)
+            else:
+                self._assign_multi_mesh(lhs_func, subset, funcs, operator, allow_missing_dofs)
+
+    def _assign_single_mesh(self, lhs_func, subset, funcs, operator):
+        assign_to_halos = all(f.dat.halo_valid for f in funcs) and (lhs_func.dat.halo_valid or subset is None)
         if assign_to_halos:
-            subset_indices = self._subset.indices if self._subset else ...
+            subset_indices = ... if subset is None else subset.indices
             data_ro = operator.attrgetter("data_ro_with_halos")
         else:
-            subset_indices = self._subset.owned_indices if self._subset else ...
+            subset_indices = ... if subset is None else subset.owned_indices
             data_ro = operator.attrgetter("data_ro")
-
-        # If mixed, loop over individual components
-        for lhs_dat, *func_dats in zip(self._assignee.dat.split,
-                                       *(f.dat.split for f in self._functions)):
-            func_data = np.array([data_ro(f)[subset_indices] for f in func_dats])
-            rvalue = self._compute_rvalue(func_data)
-            self._assign_single_dat(lhs_dat, subset_indices, rvalue, assign_to_halos)
-
-        # if we have bothered writing to halo it naturally must not be dirty
+        func_data = np.array([data_ro(f.dat)[subset_indices] for f in funcs])
+        rvalue = self._compute_rvalue(func_data)
+        self._assign_single_dat(lhs_func.dat, subset_indices, rvalue, assign_to_halos)
         if assign_to_halos:
-            self._assignee.dat.halo_valid = True
+            lhs_func.dat.halo_valid = True
+
+    def _assign_multi_mesh(self, lhs_func, subset, funcs, operator, allow_missing_dofs):
+        target_mesh = extract_unique_domain(lhs_func)
+        target_V = lhs_func.function_space()
+        source_V, = set(f.function_space() for f in funcs)
+        composed_map = source_V.topological.entity_node_map(target_mesh.topology, "cell", "everywhere", None)
+        indices_active = composed_map.indices_active_with_halo
+        indices_active_all = indices_active.all()
+        indices_active_all = target_mesh.comm.allreduce(indices_active_all, op=MPI.LAND)
+        if subset is None:
+            if not indices_active_all and not allow_missing_dofs:
+                raise ValueError("Found assignee nodes with no matching assigner nodes: run with `allow_missing_dofs=True`")
+            subset_indices_target = target_V.cell_node_map().values_with_halo[indices_active, :].flatten()
+            subset_indices_source = composed_map.values_with_halo[indices_active, :].flatten()
+        else:
+            subset_indices_target, perm, _ = np.intersect1d(
+                target_V.cell_node_map().values_with_halo[indices_active, :].flatten(),
+                subset.indices,
+                return_indices=True,
+            )
+            if len(subset.indices) > len(subset_indices_target) and not allow_missing_dofs:
+                raise ValueError("Found assignee nodes with no matching assigner nodes: run with `allow_missing_dofs=True`")
+            subset_indices_source = composed_map.values_with_halo[indices_active, :].flatten()[perm]
+        # Use buffer array to make sure that owned DoFs are updated upon assigning.
+        # The following example illustrates the issue that a naive assignment would cause.
+        #
+        # Consider the following target/source meshes distributed over 2 processes
+        # with no partition overlap:
+        #
+        #                0----0----0----1----1
+        #                |         |         |
+        # target         0    0    0    1    1
+        # (parent mesh)  |         |         |
+        #                0----0----0----1----1  (owning ranks are shown)
+        #
+        #                          1----1----1
+        #                          |         |
+        # source                   1    1    1
+        # (submesh)                |         |
+        #                          1----1----1  (owning ranks are shown)
+        #
+        # Consider CG1 functions f (on parent) and fsub (on submesh). By a naive
+        # f.assign(fsub, subset=...), the DoFs shared by rank 0 and rank 1 would
+        # only be updated on rank 1, which sees those DoFs as ghost, and those
+        # updated values on rank 1 would be overridden by the old values on rank 0
+        # upon a halo exchange.
+        #
+        # TODO: Use work array for buffer?
+        buffer = type(lhs_func)(target_V)
+        finfo = np.finfo(lhs_func.dat.dtype)
+        buffer.dat._data[:] = finfo.max
+        func_data = np.array([f.dat.data_ro_with_halos[subset_indices_source] for f in funcs])
+        rvalue = self._compute_rvalue(func_data)
+        self._assign_single_dat(buffer.dat, subset_indices_target, rvalue, True)
+        # Make all owned DoFs up-to-date; ghost DoFs may or may not be up-to-date after this.
+        buffer.dat.local_to_global_begin(op2.MIN)
+        buffer.dat.local_to_global_end(op2.MIN)
+        indices = np.where(buffer.dat.data_ro_with_halos < finfo.max * 0.999999999999)
+        lhs_func.dat.data_wo_with_halos[indices] = buffer.dat.data_ro_with_halos[indices]
 
     @cached_property
     def _constants(self):