WIP, now losing boundary markers

Author: connorjward

firedrake/cython/dmcommon.pyx

@@ -3921,15 +3921,16 @@ def submesh_create_cell_closure_cell_submesh(PETSc.DM subdm,
     return subcell_closure
 
 
-def set_coordinate_section_and_such(mesh, coordinates):
+def make_geometry_dm(coordinates) -> PETSc.DM:
     cdef:
-        PETSc.DM plex
+        PETSc.DM topology_dm, geometry_dm
         PETSc.Section vector_section
         PetscInt gdim
 
-    plex = mesh.topology_dm
+    topology_dm = coordinates.function_space().mesh().topology_dm
+    geometry_dm = topology_dm.clone()
     coordinate_element = coordinates.ufl_element()
-    gdim = plex.getCoordinateDim()
+    gdim = topology_dm.getCoordinateDim()
 
     # the existing section has the correct numbering but is scalar, expand to gdim
     scalar_dm = coordinates.function_space().dm
@@ -3950,11 +3951,67 @@ def set_coordinate_section_and_such(mesh, coordinates):
 
     if coordinate_element.family() == "Lagrange":
         # apparently gdim ignored in this call and set explicitly below
-        CHKERR(DMSetCoordinateSection(plex.dm, gdim, vector_section.sec))
-        plex.setCoordinateDim(gdim)
-        plex.setCoordinatesLocal(coordinates.dat._vec)
+        CHKERR(DMSetCoordinateSection(geometry_dm.dm, gdim, vector_section.sec))
+        geometry_dm.setCoordinateDim(gdim)
+        geometry_dm.setCoordinatesLocal(coordinates.dat._vec)
     else:
-        plex.setCellCoordinateDM(plex.getCoordinateDM().clone())
-        plex.setCellCoordinateSection(gdim, vector_section)  # gdim ignored
-        plex.setCoordinateDim(gdim)
-        plex.setCellCoordinatesLocal(coordinates.dat._vec)
+        geometry_dm.setCellCoordinateDM(geometry_dm.getCoordinateDM().clone())
+        geometry_dm.setCellCoordinateSection(gdim, vector_section)  # gdim ignored
+        geometry_dm.setCoordinateDim(gdim)
+        geometry_dm.setCellCoordinatesLocal(coordinates.dat._vec)
+
+    return geometry_dm
+
+
+def dmplex_migrate(PETSc.DM dm, PETSc.SF sf) -> PETSc.DM:
+    cdef:
+        PETSc.DM migrated_dm
+
+    migrated_dm = PETSc.DMPlex().create(comm=dm.comm)
+    CHKERR(DMPlexMigrate(dm.dm, sf.sf, migrated_dm.dm))
+    return migrated_dm
+
+
+def densify_sf(PETSc.DM topology_dm, PETSc.SF sparse_sf) -> PETSc.SF:
+    cdef:
+        PETSc.SF dense_sf
+        PetscInt dof_nroots, dof_nleaves
+        PetscInt *ilocal_new = NULL
+        PetscSFNode *iremote_new = NULL
+        PETSc.SF point_sf
+        PetscInt nroots, nleaves
+        const PetscInt *ilocal = NULL
+        const PetscSFNode *iremote = NULL
+        PETSc.Section local_sec
+        PetscInt pStart, pEnd, p, dof, off, m, n, i, j
+        np.ndarray local_offsets
+        np.ndarray remote_offsets
+
+    pStart, pEnd = topology_dm.getChart()
+    nPoints = pEnd - pStart
+    CHKERR(PetscSFGetGraph(sparse_sf.sf, &nroots, &nleaves, &ilocal, &iremote))
+
+    CHKERR(PetscMalloc1(nPoints, &ilocal_new))
+    CHKERR(PetscMalloc1(nPoints, &iremote_new))
+    for p in range(pStart, pEnd):
+        ilocal_new[p] = p
+        iremote_new[p].rank = topology_dm.comm.rank
+        iremote_new[p].index = p
+
+    for i in range(nleaves):
+        p = ilocal[i] if ilocal else i
+        iremote_new[p].rank = iremote[i].rank
+        iremote_new[p].index = iremote[i].index
+
+    dense_sf = PETSc.SF().create(comm=sparse_sf.comm)
+    CHKERR(PetscSFSetGraph(dense_sf.sf, nPoints, nPoints, ilocal_new, PETSC_OWN_POINTER, iremote_new, PETSC_OWN_POINTER))
+    return dense_sf
+
+
+def dmplex_create_overlap_migration_sf(PETSc.DM topology_dm, PETSc.SF overlap_sf) -> PETSc.SF:
+    cdef:
+        PETSc.SF migration_sf
+
+    migration_sf = PETSc.SF().create(comm=topology_dm.comm)
+    CHKERR(DMPlexCreateOverlapMigrationSF(topology_dm.dm, overlap_sf.sf, &migration_sf.sf))
+    return migration_sf

firedrake/cython/petschdr.pxi

@@ -57,6 +57,9 @@ cdef extern from "petscdmplex.h" nogil:
     int DMPlexSetSubpointMap(PETSc.PetscDM,PETSc.PetscDMLabel)
 
     int DMSetCoordinateSection(PETSc.PetscDM,PetscInt,PETSc.PetscSection)
+    int DMPlexMigrate(PETSc.PetscDM,PETSc.PetscSF,PETSc.PetscDM)
+    int DMPlexCreateOverlapMigrationSF(PETSc.PetscDM,PETSc.PetscSF,PETSc.PetscSF*)
+
 
 cdef extern from "petscdmlabel.h" nogil:
     struct _n_DMLabel

firedrake/mesh.py

@@ -543,7 +543,7 @@ def _from_cell_list(dim, cells, coords, comm, name=None):
     return plex
 
 
-class AbstractMeshTopology(object, metaclass=abc.ABCMeta):
+class AbstractMeshTopology(metaclass=abc.ABCMeta):
     """A representation of an abstract mesh topology without a concrete
         PETSc DM implementation"""
 
@@ -1186,6 +1186,7 @@ def _add_overlap(self):
             dmcommon.set_adjacency_callback(self.topology_dm)
             original_name = self.topology_dm.getName()
             sfBC = self.topology_dm.distributeOverlap(overlap)
+            self.overlap_sf = sfBC
             self.topology_dm.setName(original_name)
             self.sfBC = self.sfBC.compose(sfBC) if self.sfBC else sfBC
             dmcommon.clear_adjacency_callback(self.topology_dm)
@@ -1194,6 +1195,7 @@ def _add_overlap(self):
             # Default is FEM (vertex star) adjacency.
             original_name = self.topology_dm.getName()
             sfBC = self.topology_dm.distributeOverlap(overlap)
+            self.overlap_sf = sfBC
             self.topology_dm.setName(original_name)
             self.sfBC = self.sfBC.compose(sfBC) if self.sfBC else sfBC
             self._grown_halos = True
@@ -2242,6 +2244,7 @@ def __new__(cls, element, comm):
         cargo = MeshGeometryCargo(uid)
         assert isinstance(element, finat.ufl.FiniteElementBase)
         ufl.Mesh.__init__(mesh, element, ufl_id=mesh.uid, cargo=cargo)
+
         return mesh
 
     @MeshGeometryMixin._ad_annotate_init
@@ -2253,9 +2256,15 @@ def __init__(self, coordinates):
         coordinates : CoordinatelessFunction
             The `CoordinatelessFunction` containing the coordinates.
 
+        Notes
+        -----
+        At the time this constructor is called the mesh is fully distributed.
+
         """
         topology = coordinates.function_space().mesh()
 
+        self.geometry_dm = dmcommon.make_geometry_dm(coordinates)
+
         # this is codegen information so we attach it to the MeshGeometry rather than its cargo
         self.extruded = isinstance(topology, ExtrudedMeshTopology)
         self.variable_layers = self.extruded and topology.variable_layers
@@ -2272,9 +2281,9 @@ def __init__(self, coordinates):
         self._spatial_index = None
         self._saved_coordinate_dat_version = coordinates.dat.dat_version
 
-        # NOTE: This is only ever called with a fully distributed, overlapping
-        # mesh
-        dmcommon.set_coordinate_section_and_such(self, coordinates)
+    @property
+    def cell_closure_list(self):
+        return self.dm.getCoordinateDM().getAttr("cell_closure_list")
 
     def _ufl_signature_data_(self, *args, **kwargs):
         return (type(self), self.extruded, self.variable_layers,
@@ -2288,7 +2297,7 @@ def init(self):
         if hasattr(self, '_callback'):
             self._callback(self)
 
-    def _init_topology(self, topology):
+    def _init_topology(self, topology, geometry_dm):
         """Initialise the topology.
 
         :arg topology: The :class:`.MeshTopology` object.
@@ -2302,12 +2311,24 @@ def _init_topology(self, topology):
         import firedrake.function as function
 
         self._topology = topology
+        self._input_geometry_dm = geometry_dm
 
         def callback(self):
             """Finish initialisation."""
             del self._callback
             # Finish the initialisation of mesh topology
             self.topology.init()
+            # geometry_dm = self._input_geometry_dm.migrate(self.topology.overlap_sf)
+
+            # make the overlap SF dense
+            # dense_overlap_sf = dmcommon.densify_sf(self.topology.topology_dm, self.topology.overlap_sf)
+            # breakpoint()
+            # migration_sf = dmcommon.dmplex_create_overlap_migration_sf(self.topology.topology_dm, self.topology.overlap_sf)
+            # geometry_dm = dmcommon.dmplex_migrate(self._input_geometry_dm, dense_overlap_sf)
+            # geometry_dm = dmcommon.dmplex_migrate(self._input_geometry_dm, self.topology.overlap_sf)
+            # geometry_dm = dmcommon.dmplex_migrate(self._input_geometry_dm, migration_sf)
+            geometry_dm = dmcommon.dmplex_migrate(self._input_geometry_dm, self.topology.sfBC)
+
             coordinates_fs = functionspace.FunctionSpace(self.topology, self.ufl_coordinate_element())
 
             # reordered_coords only needs to reorder the coordinates
@@ -2316,7 +2337,7 @@ def callback(self):
             # anything. Otherwise, we have the set the correct section.
             # TODO: Ask Matt if there is a nice way to do this.
 
-            coordinates_data = dmcommon.reordered_coords(topology.topology_dm, coordinates_fs.dm.getDefaultSection(),
+            coordinates_data = dmcommon.reordered_coords(geometry_dm, coordinates_fs.dm.getDefaultSection(),
                                                          (self.num_vertices(), self.geometric_dimension()))
             coordinates = function.CoordinatelessFunction(coordinates_fs,
                                                           val=coordinates_data,
@@ -2879,8 +2900,9 @@ def make_mesh_from_coordinates(coordinates, name, tolerance=0.5):
     return mesh
 
 
-def make_mesh_from_mesh_topology(topology, name, tolerance=0.5):
-    """Make mesh from tpology.
+# Should be make_mesh_from_geometry_dm
+def make_mesh_from_mesh_topology(topology, geometry_dm, name, tolerance=0.5):
+    """Make mesh from topology.
 
     Parameters
     ----------
@@ -2901,13 +2923,13 @@ def make_mesh_from_mesh_topology(topology, name, tolerance=0.5):
     # TODO: meshfile might indicates higher-order coordinate element
     cell = topology.ufl_cell()
     geometric_dim = topology.topology_dm.getCoordinateDim()
-    if not topology.topology_dm.getCoordinatesLocalized():
+    if not geometry_dm.getCoordinatesLocalized():
         element = finat.ufl.VectorElement("Lagrange", cell, 1, dim=geometric_dim)
     else:
         element = finat.ufl.VectorElement("DQ" if cell in [ufl.quadrilateral, ufl.hexahedron] else "DG", cell, 1, dim=geometric_dim, variant="equispaced")
     # Create mesh object
     mesh = MeshGeometry.__new__(MeshGeometry, element, topology.comm)
-    mesh._init_topology(topology)
+    mesh._init_topology(topology, geometry_dm)
     mesh.name = name
     mesh._tolerance = tolerance
     return mesh
@@ -3105,17 +3127,23 @@ def Mesh(meshfile, **kwargs):
             raise RuntimeError("Mesh file %s has unknown format '%s'."
                                % (meshfile, ext[1:]))
         plex.setName(_generate_default_mesh_topology_name(name))
+
     # Create mesh topology
-    topology = MeshTopology(plex, name=plex.getName(), reorder=reorder,
+    # Pass the coordinate DM because we only want the topology here
+    topology = MeshTopology(plex.getCoordinateDM(), name=plex.getName(), reorder=reorder,
                             distribution_parameters=distribution_parameters,
                             distribution_name=kwargs.get("distribution_name"),
                             permutation_name=kwargs.get("permutation_name"),
                             comm=user_comm)
+
+    # distributed_plex = dmcommon.dmplex_migrate(plex, topology.sfBC)
+    distributed_plex = plex
+
     if netgen and isinstance(meshfile, netgen.libngpy._meshing.Mesh):
         netgen_firedrake_mesh.createFromTopology(topology, name=name, comm=user_comm)
         mesh = netgen_firedrake_mesh.firedrakeMesh
     else:
-        mesh = make_mesh_from_mesh_topology(topology, name)
+        mesh = make_mesh_from_mesh_topology(topology, distributed_plex, name)
     mesh._tolerance = tolerance
     return mesh
 
@@ -4576,13 +4604,13 @@ def RelabeledMesh(mesh, indicator_functions, subdomain_ids, **kwargs):
     distribution_parameters_noop = {"partition": False,
                                     "overlap_type": (DistributedMeshOverlapType.NONE, 0)}
     reorder_noop = None
-    tmesh1 = MeshTopology(plex1, name=plex1.getName(), reorder=reorder_noop,
+    tmesh1 = MeshTopology(plex1.getCoordinateDM(), name=plex1.getName(), reorder=reorder_noop,
                           distribution_parameters=distribution_parameters_noop,
                           perm_is=tmesh._dm_renumbering,
                           distribution_name=tmesh._distribution_name,
                           permutation_name=tmesh._permutation_name,
                           comm=tmesh.comm)
-    return make_mesh_from_mesh_topology(tmesh1, name1)
+    return make_mesh_from_mesh_topology(tmesh1, plex1, name1)
 
 
 @PETSc.Log.EventDecorator()

firedrake/mg/mesh.py

@@ -131,7 +131,7 @@ def MeshHierarchy(mesh, refinement_levels,
             raise RuntimeError("Cannot create a NetgenHierarchy from a mesh that has not been generated by\
                                 Netgen.")
 
-    cdm = mesh.topology_dm
+    cdm = mesh._input_geometry_dm
     cdm.setRefinementUniform(True)
     dms = []
     if mesh.comm.size > 1 and mesh._grown_halos: