io: adapt to simplified plex interface

Author: ksagiyam

firedrake/checkpointing.py

@@ -1114,7 +1114,7 @@ def load_mesh(self, name=DEFAULT_MESH_NAME, reorder=None, distribution_parameter
             coordinates = self._load_function_topology(tmesh, coord_element, coord_name)
             mesh = make_mesh_from_coordinates(coordinates, name)
             # Load plex coordinates for a complete representation of plex.
-            tmesh.topology_dm.coordinatesLoad(self.viewer, tmesh.sfXC)
+            tmesh.topology_dm.coordinatesLoad(self.viewer)
             # Load cell_orientations for immersed meshes.
             path = self._path_to_mesh_immersed(tmesh.name, name)
             if path in self.h5pyfile:
@@ -1192,9 +1192,9 @@ def _load_mesh_topology(self, tmesh_name, reorder, distribution_parameters):
         format = ViewerHDF5.Format.HDF5_PETSC
         self.viewer.pushFormat(format=format)
         plex.distributionSetName(distribution_name)
-        sfXB = plex.topologyLoad(self.viewer)
+        plex.topologyLoad(self.viewer)
         plex.distributionSetName(None)
-        plex.labelsLoad(self.viewer, sfXB)
+        plex.labelsLoad(self.viewer)
         self.viewer.popFormat()
         # These labels are distribution dependent.
         # We should be able to save/load labels selectively.
@@ -1223,7 +1223,7 @@ def _load_mesh_topology(self, tmesh_name, reorder, distribution_parameters):
         # -- Construct Mesh (Topology) --
         # Use public API so pass user comm (self.comm)
         tmesh = MeshTopology(plex, name=plex.getName(), reorder=reorder,
-                             distribution_parameters=distribution_parameters, sfXB=sfXB, perm_is=perm_is,
+                             distribution_parameters=distribution_parameters, perm_is=perm_is,
                              distribution_name=distribution_name, permutation_name=permutation_name,
                              comm=self.comm)
         return tmesh
@@ -1283,9 +1283,8 @@ def _load_function_space_topology(self, tmesh, element):
             section.setPermutation(tmesh._dm_renumbering)
             dm.setSection(section)
             base_tmesh = tmesh._base_mesh if isinstance(tmesh, ExtrudedMeshTopology) else tmesh
-            sfXC = base_tmesh.sfXC
             topology_dm.setName(tmesh.name)
-            gsf, lsf = topology_dm.sectionLoad(self.viewer, dm, sfXC)
+            gsf, lsf = topology_dm.sectionLoad(self.viewer, dm)
             topology_dm.setName(base_tmesh.name)
             nodes_per_entity, real_tensorproduct, block_size = sd_key
             # Don't cache if the section has been expanded by block_size

firedrake/mesh.py

@@ -486,7 +486,7 @@ class AbstractMeshTopology(object, metaclass=abc.ABCMeta):
     """A representation of an abstract mesh topology without a concrete
         PETSc DM implementation"""
 
-    def __init__(self, topology_dm, name, reorder, sfXB, perm_is, distribution_name, permutation_name, comm, submesh_parent=None):
+    def __init__(self, topology_dm, name, reorder, perm_is, distribution_name, permutation_name, comm, submesh_parent=None):
         """Initialise a mesh topology.
 
         Parameters
@@ -497,11 +497,6 @@ def __init__(self, topology_dm, name, reorder, sfXB, perm_is, distribution_name,
             Name of the mesh topology.
         reorder : bool
             Whether to reorder the mesh entities.
-        sfXB : PETSc.PetscSF
-            `PETSc.SF` that pushes forward the global point number
-            slab ``[0, NX)`` to input (naive) plex (only significant when
-            the mesh topology is loaded from file and only passed from inside
-            `~.CheckpointFile`).
         perm_is : PETSc.IS
             `PETSc.IS` that is used as ``_dm_renumbering``; only
             makes sense if we know the exact parallel distribution of ``plex``
@@ -522,10 +517,6 @@ def __init__(self, topology_dm, name, reorder, sfXB, perm_is, distribution_name,
         topology_dm.setFromOptions()
         self.topology_dm = topology_dm
         r"The PETSc DM representation of the mesh topology."
-        self.sfBC = None
-        r"The PETSc SF that pushes the input (naive) plex to current (good) plex."
-        self.sfXB = sfXB
-        r"The PETSc SF that pushes the global point number slab [0, NX) to input (naive) plex."
         self.submesh_parent = submesh_parent
         # User comm
         self.user_comm = comm
@@ -536,8 +527,6 @@ def __init__(self, topology_dm, name, reorder, sfXB, perm_is, distribution_name,
         self._grown_halos = False
         if self.comm.size > 1:
             self._add_overlap()
-        if self.sfXB is not None:
-            self.sfXC = sfXB.compose(self.sfBC) if self.sfBC else self.sfXB
         dmcommon.label_facets(self.topology_dm)
         dmcommon.complete_facet_labels(self.topology_dm)
         # TODO: Allow users to set distribution name if they want to save
@@ -1057,7 +1046,6 @@ def __init__(
         name,
         reorder,
         distribution_parameters,
-        sfXB=None,
         perm_is=None,
         distribution_name=None,
         permutation_name=None,
@@ -1076,11 +1064,6 @@ def __init__(
             Whether to reorder the mesh entities.
         distribution_parameters : dict
             Options controlling mesh distribution; see `Mesh` for details.
-        sfXB : PETSc.PetscSF
-            `PETSc.SF` that pushes forward the global point number
-            slab ``[0, NX)`` to input (naive) plex (only significant when
-            the mesh topology is loaded from file and only passed from inside
-            `~.CheckpointFile`).
         perm_is : PETSc.IS
             `PETSc.IS` that is used as ``_dm_renumbering``; only
             makes sense if we know the exact parallel distribution of ``plex``
@@ -1111,7 +1094,7 @@ def __init__(
         # Disable auto distribution and reordering before setFromOptions is called.
         plex.distributeSetDefault(False)
         plex.reorderSetDefault(PETSc.DMPlex.ReorderDefaultFlag.FALSE)
-        super().__init__(plex, name, reorder, sfXB, perm_is, distribution_name, permutation_name, comm, submesh_parent=submesh_parent)
+        super().__init__(plex, name, reorder, perm_is, distribution_name, permutation_name, comm, submesh_parent=submesh_parent)
 
     def _distribute(self):
         # Distribute/redistribute the dm to all ranks
@@ -1122,9 +1105,8 @@ def _distribute(self):
             # refine this mesh in parallel.  Later, when we actually use
             # it, we grow the halo.
             original_name = plex.getName()
-            sfBC = plex.distribute(overlap=0)
+            _ = plex.distribute(overlap=0)
             plex.setName(original_name)
-            self.sfBC = sfBC
             # plex carries a new dm after distribute, which
             # does not inherit partitioner from the old dm.
             # It probably makes sense as chaco does not work
@@ -1140,17 +1122,15 @@ def _add_overlap(self):
         elif overlap_type in [DistributedMeshOverlapType.FACET, DistributedMeshOverlapType.RIDGE]:
             dmcommon.set_adjacency_callback(self.topology_dm, overlap_type)
             original_name = self.topology_dm.getName()
-            sfBC = self.topology_dm.distributeOverlap(overlap)
+            _ = self.topology_dm.distributeOverlap(overlap)
             self.topology_dm.setName(original_name)
-            self.sfBC = self.sfBC.compose(sfBC) if self.sfBC else sfBC
             dmcommon.clear_adjacency_callback(self.topology_dm)
             self._grown_halos = True
         elif overlap_type == DistributedMeshOverlapType.VERTEX:
             # Default is FEM (vertex star) adjacency.
             original_name = self.topology_dm.getName()
-            sfBC = self.topology_dm.distributeOverlap(overlap)
+            _ = self.topology_dm.distributeOverlap(overlap)
             self.topology_dm.setName(original_name)
-            self.sfBC = self.sfBC.compose(sfBC) if self.sfBC else sfBC
             self._grown_halos = True
         else:
             raise ValueError("Unknown overlap type %r" % overlap_type)