Cleanup

Author: pbrubeck

docs/notebooks/03-elasticity.ipynb

@@ -328,7 +328,7 @@
       "\u001b[36mFile \u001b[39m\u001b[32mpetsc4py/PETSc/Log.pyx:188\u001b[39m, in \u001b[36mpetsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func\u001b[39m\u001b[34m()\u001b[39m\n",
       "\u001b[36mFile \u001b[39m\u001b[32mpetsc4py/PETSc/Log.pyx:189\u001b[39m, in \u001b[36mpetsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func\u001b[39m\u001b[34m()\u001b[39m\n",
       "\u001b[36mFile \u001b[39m\u001b[32m~/src/firedrake-pyadjoint/firedrake/firedrake/adjoint_utils/variational_solver.py:108\u001b[39m, in \u001b[36mNonlinearVariationalSolverMixin._ad_annotate_solve.<locals>.wrapper\u001b[39m\u001b[34m(self, **kwargs)\u001b[39m\n\u001b[32m    105\u001b[39m     tape.add_block(block)\n\u001b[32m    107\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m stop_annotating():\n\u001b[32m--> \u001b[39m\u001b[32m108\u001b[39m     out = \u001b[43msolve\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    110\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m annotate:\n\u001b[32m    111\u001b[39m     block.add_output(\u001b[38;5;28mself\u001b[39m._ad_problem._ad_u.create_block_variable())\n",
-      "\u001b[36mFile \u001b[39m\u001b[32m~/src/firedrake-pyadjoint/firedrake/firedrake/variational_solver.py:361\u001b[39m, in \u001b[36mNonlinearVariationalSolver.solve\u001b[39m\u001b[34m(self, bounds)\u001b[39m\n\u001b[32m    359\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m problem.restrict:\n\u001b[32m    360\u001b[39m     problem.u.interpolate(problem.u_restrict)\n\u001b[32m--> \u001b[39m\u001b[32m361\u001b[39m \u001b[43msolving_utils\u001b[49m\u001b[43m.\u001b[49m\u001b[43mcheck_snes_convergence\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43msnes\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    363\u001b[39m \u001b[38;5;66;03m# Grab the comm associated with the `_problem` and call PETSc's garbage cleanup routine\u001b[39;00m\n\u001b[32m    364\u001b[39m comm = \u001b[38;5;28mself\u001b[39m._problem.u_restrict.function_space().mesh()._comm\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m~/src/firedrake-pyadjoint/firedrake/firedrake/variational_solver.py:361\u001b[39m, in \u001b[36mNonlinearVariationalSolver.solve\u001b[39m\u001b[34m(self, bounds)\u001b[39m\n\u001b[32m    359\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m problem.restrict:\n\u001b[32m    360\u001b[39m     problem.u.interpolate(problem.u_restrict)\n\u001b[32m--> \u001b[39m\u001b[32m361\u001b[39m \u001b[43msolving_utils\u001b[49m\u001b[43m.\u001b[49m\u001b[43mcheck_snes_convergence\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43msnes\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    363\u001b[39m \u001b[38;5;66;03m# Grab the comm associated with the `_problem` and call PETSc's garbage cleanup routine\u001b[39;00m\n\u001b[32m    364\u001b[39m comm = \u001b[38;5;28mself\u001b[39m._problem.u_restrict.function_space().mesh().comm\n",
       "\u001b[36mFile \u001b[39m\u001b[32m~/src/firedrake-pyadjoint/firedrake/firedrake/solving_utils.py:128\u001b[39m, in \u001b[36mcheck_snes_convergence\u001b[39m\u001b[34m(snes)\u001b[39m\n\u001b[32m    126\u001b[39m         \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m    127\u001b[39m             msg = reason\n\u001b[32m--> \u001b[39m\u001b[32m128\u001b[39m         \u001b[38;5;28;01mraise\u001b[39;00m ConvergenceError(\u001b[33mr\u001b[39m\u001b[33m\"\"\"\u001b[39m\u001b[33mNonlinear solve failed to converge after \u001b[39m\u001b[38;5;132;01m%d\u001b[39;00m\u001b[33m nonlinear iterations.\u001b[39m\n\u001b[32m    129\u001b[39m \u001b[33mReason:\u001b[39m\n\u001b[32m    130\u001b[39m \u001b[33m   \u001b[39m\u001b[38;5;132;01m%s\u001b[39;00m\u001b[33m\"\"\"\u001b[39m % (snes.getIterationNumber(), msg))\n",
       "\u001b[31mConvergenceError\u001b[39m: Nonlinear solve failed to converge after 0 nonlinear iterations.\nReason:\n   DIVERGED_LINEAR_SOLVE"
      ]

docs/source/parallelism.rst

@@ -96,9 +96,5 @@ different simulations on the two halves we would write.
 To access the communicator a mesh was created on, we can use the
 ``mesh.comm`` property, or the function ``mesh.mpi_comm``.
 
-.. warning::
-  Do not use the internal ``mesh._comm`` attribute for communication.
-  This communicator is for internal Firedrake MPI communication only.
-
 .. _MPI: http://mpi-forum.org/
 .. _STREAMS: http://www.cs.virginia.edu/stream/

firedrake/__init__.py

@@ -8,6 +8,7 @@
 # TODO RELEASE
 # PETSC_SUPPORTED_VERSIONS = ">=3.25"
 
+
 def init_petsc():
     import os
     import sys

firedrake/adjoint/__init__.py

@@ -10,13 +10,6 @@
 import pyadjoint
 __version__ = pyadjoint.__version__
 
-import sys
-if 'backend' not in sys.modules:
-    import firedrake
-    sys.modules['backend'] = firedrake
-else:
-    raise ImportError("'backend' module already exists?")
-
 from pyadjoint.tape import Tape, set_working_tape, get_working_tape, \
                             pause_annotation, continue_annotation, \
                             stop_annotating, annotate_tape  # noqa F401

firedrake/adjoint_utils/function.py

@@ -1,4 +1,5 @@
 from functools import wraps
+from pyop2.mpi import temp_internal_comm
 import ufl
 from ufl.domain import extract_unique_domain
 from pyadjoint.overloaded_type import create_overloaded_object, FloatingType
@@ -280,8 +281,9 @@ def _ad_assign_numpy(dst, src, offset):
         m_a_local = src[offset + range_begin:offset + range_end]
         if dst.function_space().ufl_element().family() == "Real":
             # Real space keeps a redundant copy of the data on every rank
-            comm = dst.function_space().mesh()._comm
-            dst.dat.data_wo[...] = comm.bcast(m_a_local, root=0)
+            comm = dst.function_space().mesh().comm
+            with temp_internal_comm(comm) as icomm:
+                dst.dat.data_wo[...] = icomm.bcast(m_a_local, root=0)
         else:
             dst.dat.data_wo[...] = m_a_local.reshape(dst.dat.data_wo.shape)
         offset += dst.dat.dataset.layout_vec.size

firedrake/assemble.py

@@ -615,7 +615,7 @@ def base_form_assembly_visitor(self, expr, tensor, bcs, *args):
             if rank > 2:
                 raise ValueError("Cannot assemble an Interpolate with more than two arguments")
             interpolator = get_interpolator(expr)
-            return interpolator.assemble(tensor=tensor, bcs=bcs)
+            return interpolator.assemble(tensor=tensor, bcs=bcs, mat_type=self._mat_type, sub_mat_type=self._sub_mat_type)
         elif tensor and isinstance(expr, (firedrake.Function, firedrake.Cofunction, firedrake.MatrixBase)):
             return tensor.assign(expr)
         elif tensor and isinstance(expr, ufl.ZeroBaseForm):
@@ -855,6 +855,15 @@ def restructure_base_form(expr, visited=None):
         if isinstance(expr, ufl.FormSum) and all(ufl.duals.is_dual(a.function_space()) for a in expr.arguments()):
             # Return ufl.Sum if we are assembling a FormSum with Coarguments (a primal expression)
             return sum(w*c for w, c in zip(expr.weights(), expr.components()))
+
+        # If F: V3 x V2 -> R, then
+        # Interpolate(TestFunction(V1), F) <=> Action(Interpolate(TestFunction(V1), TrialFunction(V2.dual())), F).
+        # The result is a two-form V3 x V1 -> R.
+        if isinstance(expr, ufl.Interpolate) and isinstance(expr.argument_slots()[0], ufl.form.Form):
+            form, operand = expr.argument_slots()
+            vstar = firedrake.Argument(form.arguments()[0].function_space().dual(), 1)
+            expr = expr._ufl_expr_reconstruct_(operand, v=vstar)
+            return ufl.action(expr, form)
         return expr
 
     @staticmethod
@@ -1159,13 +1168,13 @@ def __init__(self, form, form_compiler_parameters=None):
 
     def allocate(self):
         # Getting the comm attribute of a form isn't straightforward
-        # form.ufl_domains()[0]._comm seems the most robust method
+        # form.ufl_domains()[0].comm seems the most robust method
         # revisit in a refactor
         return op2.Global(
             1,
             [0.0],
             dtype=utils.ScalarType,
-            comm=self._form.ufl_domains()[0]._comm
+            comm=self._form.ufl_domains()[0].comm
         )
 
     def _apply_bc(self, tensor, bc, u=None):

firedrake/checkpointing.py

@@ -3,7 +3,7 @@
 from petsc4py.PETSc import ViewerHDF5
 import finat.ufl
 from pyop2 import op2
-from pyop2.mpi import COMM_WORLD, internal_comm, MPI
+from pyop2.mpi import COMM_WORLD, MPI
 from petsctools import OptionsManager
 from firedrake.cython import hdf5interface as h5i
 from firedrake.cython import dmcommon
@@ -104,7 +104,6 @@ def __init__(self, basename, single_file=True,
             warnings.warn("DumbCheckpoint class will soon be deprecated; use CheckpointFile class instead.",
                           DeprecationWarning)
         self.comm = comm or COMM_WORLD
-        self._comm = internal_comm(self.comm, self)
         self.mode = mode
 
         self._single = single_file
@@ -195,7 +194,7 @@ def new_file(self, name=None):
         if mode == FILE_UPDATE and not exists:
             mode = FILE_CREATE
         self._vwr = PETSc.ViewerHDF5().create(name, mode=mode,
-                                              comm=self._comm)
+                                              comm=self.comm)
         if self.mode == FILE_READ:
             nprocs = self.read_attribute("/", "nprocs")
             if nprocs != self.comm.size:
@@ -379,7 +378,6 @@ def __init__(self, filename, file_mode, comm=None):
             warnings.warn("HDF5File class will soon be deprecated; use CheckpointFile class instead.",
                           DeprecationWarning)
         self.comm = comm or COMM_WORLD
-        self._comm = internal_comm(self.comm, self)
 
         self._filename = filename
         self._mode = file_mode
@@ -397,7 +395,7 @@ def __init__(self, filename, file_mode, comm=None):
 
         # Try to use MPI
         try:
-            self._h5file = h5py.File(filename, file_mode, driver="mpio", comm=self._comm)
+            self._h5file = h5py.File(filename, file_mode, driver="mpio", comm=self.comm)
         except NameError:  # the error you get if h5py isn't compiled against parallel HDF5
             raise RuntimeError("h5py *must* be installed with MPI support")
 
@@ -527,10 +525,9 @@ def __init__(self, filename, mode, comm=COMM_WORLD):
         self.viewer = ViewerHDF5()
         self.filename = filename
         self.comm = comm
-        self._comm = internal_comm(comm, self)
         r"""The neme of the checkpoint file."""
-        self.viewer.create(filename, mode=mode, comm=self._comm)
-        self.commkey = self._comm.py2f()
+        self.viewer.create(filename, mode=mode, comm=self.comm)
+        self.commkey = self.comm.py2f()
         assert self.commkey != MPI.COMM_NULL.py2f()
         self._function_spaces = {}
         self._function_load_utils = {}
@@ -597,7 +594,7 @@ def save_mesh(self, mesh, distribution_name=None, permutation_name=None):
                     layers_tV = impl.FunctionSpace(base_tmesh, element)
                     self._save_function_space_topology(layers_tV)
                     # Note that _cell_numbering coincides with DG0 section, so we can use tmesh.layers directly.
-                    layers_iset = PETSc.IS().createGeneral(tmesh.layers[:tmesh.cell_set.size, :], comm=tmesh._comm)
+                    layers_iset = PETSc.IS().createGeneral(tmesh.layers[:tmesh.cell_set.size, :], comm=tmesh.comm)
                     layers_iset.setName("_".join([PREFIX_EXTRUDED, "layers_iset"]))
                     self.viewer.pushGroup(path)
                     layers_iset.view(self.viewer)
@@ -661,7 +658,7 @@ def save_mesh(self, mesh, distribution_name=None, permutation_name=None):
                     reflected = o_r_map[tmesh.entity_orientations[:tmesh.cell_set.size, -1]]
                     reflected_indices = (reflected == 1)
                     canonical_cell_orientations[reflected_indices] = 1 - canonical_cell_orientations[reflected_indices]
-                    cell_orientations_iset = PETSc.IS().createGeneral(canonical_cell_orientations, comm=tmesh._comm)
+                    cell_orientations_iset = PETSc.IS().createGeneral(canonical_cell_orientations, comm=tmesh.comm)
                     cell_orientations_iset.setName("_".join([PREFIX_IMMERSED, "cell_orientations_iset"]))
                     self.viewer.pushGroup(path)
                     cell_orientations_iset.view(self.viewer)
@@ -1065,7 +1062,7 @@ def load_mesh(self, name=DEFAULT_MESH_NAME, reorder=None, distribution_parameter
                 _, _, lsf = self._function_load_utils[base_tmesh_key + sd_key]
                 nroots, _, _ = lsf.getGraph()
                 layers_a = np.empty(nroots, dtype=utils.IntType)
-                layers_a_iset = PETSc.IS().createGeneral(layers_a, comm=self._comm)
+                layers_a_iset = PETSc.IS().createGeneral(layers_a, comm=self.comm)
                 layers_a_iset.setName("_".join([PREFIX_EXTRUDED, "layers_iset"]))
                 self.viewer.pushGroup(path)
                 layers_a_iset.load(self.viewer)
@@ -1128,7 +1125,7 @@ def load_mesh(self, name=DEFAULT_MESH_NAME, reorder=None, distribution_parameter
                 _, _, lsf = self._function_load_utils[tmesh_key + sd_key]
                 nroots, _, _ = lsf.getGraph()
                 cell_orientations_a = np.empty(nroots, dtype=utils.IntType)
-                cell_orientations_a_iset = PETSc.IS().createGeneral(cell_orientations_a, comm=self._comm)
+                cell_orientations_a_iset = PETSc.IS().createGeneral(cell_orientations_a, comm=self.comm)
                 cell_orientations_a_iset.setName("_".join([PREFIX_IMMERSED, "cell_orientations_iset"]))
                 self.viewer.pushGroup(path)
                 cell_orientations_a_iset.load(self.viewer)
@@ -1165,7 +1162,7 @@ def _load_mesh_topology(self, tmesh_name, reorder, distribution_parameters):
             _distribution_name, = self.h5pyfile[path].keys()
             path = self._path_to_distribution(tmesh_name, _distribution_name)
             _comm_size = self.get_attr(path, "comm_size")
-            if _comm_size == self._comm.size and \
+            if _comm_size == self.comm.size and \
                distribution_parameters is None and reorder is None:
                 load_distribution_permutation = True
         if load_distribution_permutation:
@@ -1185,7 +1182,7 @@ def _load_mesh_topology(self, tmesh_name, reorder, distribution_parameters):
             permutation_name = None
             perm_is = None
         plex = PETSc.DMPlex()
-        plex.create(comm=self._comm)
+        plex.create(comm=self.comm)
         plex.setName(tmesh_name)
         # Check format
         path = os.path.join(self._path_to_topology(tmesh_name), "topology")
@@ -1203,15 +1200,15 @@ def _load_mesh_topology(self, tmesh_name, reorder, distribution_parameters):
         plex.removeLabel("pyop2_ghost")
         if load_distribution_permutation:
             chart_size = np.empty(1, dtype=utils.IntType)
-            chart_sizes_iset = PETSc.IS().createGeneral(chart_size, comm=self._comm)
+            chart_sizes_iset = PETSc.IS().createGeneral(chart_size, comm=self.comm)
             chart_sizes_iset.setName("chart_sizes")
             path = self._path_to_distribution(tmesh_name, distribution_name)
             self.viewer.pushGroup(path)
             chart_sizes_iset.load(self.viewer)
             self.viewer.popGroup()
             chart_size = chart_sizes_iset.getIndices().item()
             perm = np.empty(chart_size, dtype=utils.IntType)
-            perm_is = PETSc.IS().createGeneral(perm, comm=self._comm)
+            perm_is = PETSc.IS().createGeneral(perm, comm=self.comm)
             path = self._path_to_permutation(tmesh_name, distribution_name, permutation_name)
             self.viewer.pushGroup(path)
             perm_is.setName("permutation")
@@ -1276,10 +1273,10 @@ def _load_function_space_topology(self, tmesh, element):
         sd_key = self._get_shared_data_key_for_checkpointing(tmesh, element)
         if tmesh_key + sd_key not in self._function_load_utils:
             topology_dm = tmesh.topology_dm
-            dm = PETSc.DMShell().create(comm=tmesh._comm)
+            dm = PETSc.DMShell().create(comm=tmesh.comm)
             dm.setName(self._get_dm_name_for_checkpointing(tmesh, element))
             dm.setPointSF(topology_dm.getPointSF())
-            section = PETSc.Section().create(comm=tmesh._comm)
+            section = PETSc.Section().create(comm=tmesh.comm)
             section.setPermutation(tmesh._dm_renumbering)
             dm.setSection(section)
             base_tmesh = tmesh._base_mesh if isinstance(tmesh, ExtrudedMeshTopology) else tmesh
@@ -1446,7 +1443,7 @@ def _get_dm_for_checkpointing(self, tV):
             nodes_per_entity, real_tensorproduct, block_size = sd_key
             global_numbering, _ = tV.mesh().create_section(nodes_per_entity, real_tensorproduct, block_size=block_size)
             topology_dm = tV.mesh().topology_dm
-            dm = PETSc.DMShell().create(tV.mesh()._comm)
+            dm = PETSc.DMShell().create(tV.mesh().comm)
             dm.setPointSF(topology_dm.getPointSF())
             dm.setSection(global_numbering)
         else:

firedrake/cofunction.py

@@ -3,7 +3,7 @@
 import ufl
 
 from ufl.form import BaseForm
-from pyop2 import op2, mpi
+from pyop2 import op2
 from pyadjoint.tape import stop_annotating, annotate_tape, get_working_tape
 from finat.ufl import MixedElement
 import firedrake.assemble
@@ -65,18 +65,16 @@ def __init__(self, function_space, val=None, name=None, dtype=ScalarType,
 
         # User comm
         self.comm = V.comm
-        # Internal comm
-        self._comm = mpi.internal_comm(V.comm, self)
         self._function_space = V
-        self.uid = utils._new_uid(self._comm)
+        self.uid = utils._new_uid(self.comm)
         self._name = name or 'cofunction_%d' % self.uid
         self._label = "a cofunction"
 
         if isinstance(val, Cofunction):
             val = val.dat
 
         if isinstance(val, (op2.Dat, op2.DatView, op2.MixedDat, op2.Global)):
-            assert val.comm == self._comm
+            assert val.comm == self.comm
             self.dat = val
         else:
             self.dat = function_space.make_dat(val, dtype, self.name())

firedrake/cython/dmcommon.pyx

@@ -1263,7 +1263,7 @@ def create_section(mesh, nodes_per_entity, on_base=False, block_size=1, boundary
                 nodes = sum(nodes_per_entity[:, i]*(mesh.layers - i) for i in range(2)).reshape(dimension + 1, -1)
     else:
         nodes = nodes_per_entity.reshape(dimension + 1, -1)
-    section = PETSc.Section().create(comm=mesh._comm)
+    section = PETSc.Section().create(comm=mesh.comm)
     get_chart(dm.dm, &pStart, &pEnd)
     section.setChart(pStart, pEnd)
 

firedrake/dmhooks.py

@@ -393,7 +393,7 @@ def create_subdm(dm, fields, *args, **kwargs):
         # Index set mapping from W into subspace.
         iset = PETSc.IS().createGeneral(numpy.concatenate([W._ises[f].indices
                                                            for f in fields]),
-                                        comm=W._comm)
+                                        comm=W.comm)
         if ctx is not None:
             ctx, = ctx.split([fields])
             add_hook(parent, setup=partial(push_appctx, subspace.dm, ctx),