Change comm on submesh

Hi,

I’m trying to “change” the comm on a submesh from the comm it inherits from the full mesh to a split (or otherwise smaller sub-) comm. The full mesh has been partitioned by region so that certain ranks are guaranteed to have no cells (ghost mode none) and I’d like the submesh comm to only include those ranks with cells. What is the best way of doing this?

I include a somewhat contrived minimal almost working example below, run using the python API with dolfinx 0.8.0 but I can easily upgrade if necessary:

  1. imports
import gmsh
import dolfinx as df
import numpy as np
from mpi4py import MPI
comm_world = MPI.COMM_WORLD
  1. create a gmsh mesh
# Create a gmsh mesh of a unit square with various regions marked as physical groups
gmsh.initialize()

gmsh.model.add("square")
lc = 3e-2
gmsh.model.geo.addPoint(0, 0, 0, lc, 1)
gmsh.model.geo.addPoint(1, 0, 0, lc, 2)
gmsh.model.geo.addPoint(1, 0.4, 0, lc, 3)
gmsh.model.geo.addPoint(1, 0.8, 0, lc, 4)
gmsh.model.geo.addPoint(1, 1, 0, lc, 5)
gmsh.model.geo.addPoint(0.5, 1, 0, lc, 6)
gmsh.model.geo.addPoint(0, 1, 0, lc, 7)
gmsh.model.geo.addPoint(0, 0.8, 0, lc, 8)
gmsh.model.geo.addPoint(0, 0.4, 0, lc, 9)
gmsh.model.geo.addPoint(0.5, 0.4, 0, lc, 10)
gmsh.model.geo.addPoint(0.5, 0.8, 0, lc, 11)

gmsh.model.geo.addLine(1, 2, 1)
gmsh.model.geo.addLine(2, 3, 2)
gmsh.model.geo.addLine(3, 4, 3)
gmsh.model.geo.addLine(4, 5, 4)
gmsh.model.geo.addLine(5, 6, 5)
gmsh.model.geo.addLine(6, 7, 6)
gmsh.model.geo.addLine(7, 8, 7)
gmsh.model.geo.addLine(8, 9, 8)
gmsh.model.geo.addLine(9, 1, 9)
gmsh.model.geo.addLine(9, 10, 10)
gmsh.model.geo.addLine(10, 3, 11)
gmsh.model.geo.addLine(10, 11, 12)
gmsh.model.geo.addLine(8, 11, 13)
gmsh.model.geo.addLine(11, 4, 14)
gmsh.model.geo.addLine(11, 6, 15)

gmsh.model.geo.addCurveLoop([1, 2, -11, -10, 9], 1)
gmsh.model.geo.addCurveLoop([11, 3, -14, -12], 2)
gmsh.model.geo.addCurveLoop([14, 4, 5, -15], 3)
gmsh.model.geo.addCurveLoop([13, 15, 6, 7], 4)
gmsh.model.geo.addCurveLoop([10, 12, -13, 8], 5)
gmsh.model.geo.addPlaneSurface([1], 1)
gmsh.model.geo.addPlaneSurface([2], 2)
gmsh.model.geo.addPlaneSurface([3], 3)
gmsh.model.geo.addPlaneSurface([4], 4)
gmsh.model.geo.addPlaneSurface([5], 5)

gmsh.model.geo.synchronize()

gmsh.model.addPhysicalGroup(1, [7, 8, 9], 1)
gmsh.model.addPhysicalGroup(1, [2, 3, 4], 2)
gmsh.model.addPhysicalGroup(1, [1], 3)
gmsh.model.addPhysicalGroup(1, [5, 6], 4)
gmsh.model.addPhysicalGroup(2, [1], 1)
gmsh.model.addPhysicalGroup(2, [2], 2)
gmsh.model.addPhysicalGroup(2, [3], 3)
gmsh.model.addPhysicalGroup(2, [4], 4)
gmsh.model.addPhysicalGroup(2, [5], 5)

# Generate the mesh on rank 0
if comm_world.rank == 0:
    gmsh.model.mesh.generate(2)
  1. get cell region id information
# Get the cell information from the gmsh model
if comm_world.rank == 0:
    topologies = df.io.gmshio.extract_topology_and_markers(gmsh.model)

    num_cell_types = len(topologies.keys())
    cell_information = dict()
    cell_dimensions = np.zeros(num_cell_types, dtype=np.int32)
    for i, element in enumerate(topologies.keys()):
        _, dim, _, num_nodes, _, _ = gmsh.model.mesh.getElementProperties(element)
        cell_information[i] = {"id": element, "dim": dim, "num_nodes": num_nodes}
        cell_dimensions[i] = dim

    # Sort elements by ascending dimension
    perm_sort = np.argsort(cell_dimensions)
    cell_id = cell_information[perm_sort[-1]]["id"]
    num_nodes = cell_information[perm_sort[-1]]["num_nodes"]
    num_nodes = comm_world.bcast([num_nodes], root=0)

    cells = np.asarray(topologies[cell_id]["topology"], dtype=np.int64)
    cell_values = np.asarray(topologies[cell_id]["cell_data"], dtype=np.int32)
else:
    num_nodes = comm_world.bcast([None], root=0)[0]
    cells = np.empty([0, num_nodes], dtype=np.int32)
    cell_values = np.empty((0,), dtype=np.int32)
  1. setup a crude region partitioner
# convert an array to an adjacency list
def to_adj(cells, dtype):
    cflat = []
    coff = [0]
    for c in cells:
        cflat += list(c)
        cc = coff[-1] + len(c)
        coff += [cc]
    adj = df.graph.adjacencylist(np.array(cflat, dtype=dtype), 
                                 np.array(coff, dtype=dtype))
    return adj

# choose which region ids to group together (if sufficient ranks are available)
rid_groups = [[1], [2, 5], [3, 4]]

# create a partitioner that partitions the mesh depending on the cell region id
def partitioner(comm, commsize, celltype, topo):
    if commsize == 1:
        # trivial case
        al = df.cpp.graph.AdjacencyList_int32(np.zeros(len(topo), dtype=np.int32))
    elif commsize == 2:
        # special case: not a large enough comm so 
        # group the last two rid groups together
        groups = [rid_groups[0], rid_groups[1] + rid_groups[2]]
        procs = np.asarray([i for v in cell_values for i, g in enumerate(groups) if v in g], dtype=np.int32)
        al = df.cpp.graph.AdjacencyList_int32(procs)
    else:
        # arbitrary case
        groups = rid_groups
        if comm_world.rank == 0:
            # figure out which cells belong together
            cellindices = [None]*len(topo)
            cellgroups = [[] for group in groups]
            i = 0
            for g, group in enumerate(groups):
                for c, v in enumerate(cell_values):
                    if v in group:
                        cellgroups[g].append(topo.links(c))
                        cellindices[c] = i # remember the order
                        i += 1
            groupsizes = [max(1, int(len(cellgroup)*commsize/len(topo))) for cellgroup in cellgroups]
            # correct the logic above so that sum(groupsizes) == commsize
            while sum(groupsizes) < commsize:
                for g in [1, 0, 2]:
                    groupsizes[g] = groupsizes[g] + 1
                    if sum(groupsizes) == commsize: break
            while sum(groupsizes) > commsize:
                for g in [2, 0, 1]:
                    if groupsizes[g] > 1: groupsizes[g] = groupsizes[g] - 1
                    if sum(groupsizes) == commsize: break
            # communicate the groupsizes
            groupsizes = comm.bcast(groupsizes, root=0)
        else:
            cellgroups = [[] for group in groups]
            cellindices = []
            groupsizes = comm.bcast([None]*len(groups), root=0)
        if sum(groupsizes) != commsize:
            raise Exception("sum(groupsizes) != commsize")
        values = np.empty((0,), dtype=np.int32)
        # for each group, partition the cells using a cell partitioner
        for g in range(len(groups)):
            lal = df.mesh.create_cell_partitioner(df.mesh.GhostMode.none)(
                comm, groupsizes[g], celltype, to_adj(cellgroups[g], np.int64))
            values = np.concatenate((values, lal.array + sum(groupsizes[:g])), axis=0)
        al = df.graph.adjacencylist(values[cellindices])
    return al
  1. generate the mesh
# convert the gmsh mesh to a dolfinx mesh
mesh, ct, ft = df.io.gmshio.model_to_mesh(gmsh.model, comm_world, 0, gdim=2, partitioner=partitioner)

At this stage the mesh (and on 3 processes the mesh per process) looks like:

  1. create submeshes:
tdim = mesh.topology.dim
cells_0 = np.concatenate([ct.find(rid) for rid in rid_groups[0]])
submesh_0, _, _, _ = df.mesh.create_submesh(mesh, tdim, cells_0)
cells_1 = np.concatenate([ct.find(rid) for rid in rid_groups[1]])
submesh_1, _, _, _ = df.mesh.create_submesh(mesh, tdim, cells_1)

All the above works, and I’d now like to create submeshes that are on different comms, which include subsets of comm_world. Something like:
7. split comms:

comm_0 = comm_world.Split(submesh_0.topology.index_map(2).size_local>0)
comm_1 = comm_world.Split(submesh_1.topology.index_map(2).size_local>0)

and then create a new (geometrically and topologically identical) submesh_0 with comm comm_0 and similarly for submesh_1 with comm_1.

What would be the best way of doing this?

Many thanks!
Cian

You would have to make new a new Topology and Geometry by hand, which can be passed to the relevant constructors.
This can be seen in create_submesh:

as the Mesh itself takes in an MPI communicator and a Topology and a Geometry.
The Topology and Geometry constructor takes in:

Thank you @dokken!

I hadn’t thought of the comms on the index maps so pointing that out was especially helpful.

I had to upgrade to v0.9.0 to gain access to the Geometry constructors in python but after doing that everything now seems to work.

The code above (in v0.9.0) is only slightly modified (mostly the custom partitioner had to change a bit):

# imports
import gmsh
import dolfinx as df
import numpy as np
from mpi4py import MPI
comm_world = MPI.COMM_WORLD

# Create a gmsh mesh of a unit square with various regions marked as physical groups
gmsh.initialize()

gmsh.model.add("square")
lc = 3e-2
gmsh.model.geo.addPoint(0, 0, 0, lc, 1)
gmsh.model.geo.addPoint(1, 0, 0, lc, 2)
gmsh.model.geo.addPoint(1, 0.4, 0, lc, 3)
gmsh.model.geo.addPoint(1, 0.8, 0, lc, 4)
gmsh.model.geo.addPoint(1, 1, 0, lc, 5)
gmsh.model.geo.addPoint(0.5, 1, 0, lc, 6)
gmsh.model.geo.addPoint(0, 1, 0, lc, 7)
gmsh.model.geo.addPoint(0, 0.8, 0, lc, 8)
gmsh.model.geo.addPoint(0, 0.4, 0, lc, 9)
gmsh.model.geo.addPoint(0.5, 0.4, 0, lc, 10)
gmsh.model.geo.addPoint(0.5, 0.8, 0, lc, 11)

gmsh.model.geo.addLine(1, 2, 1)
gmsh.model.geo.addLine(2, 3, 2)
gmsh.model.geo.addLine(3, 4, 3)
gmsh.model.geo.addLine(4, 5, 4)
gmsh.model.geo.addLine(5, 6, 5)
gmsh.model.geo.addLine(6, 7, 6)
gmsh.model.geo.addLine(7, 8, 7)
gmsh.model.geo.addLine(8, 9, 8)
gmsh.model.geo.addLine(9, 1, 9)
gmsh.model.geo.addLine(9, 10, 10)
gmsh.model.geo.addLine(10, 3, 11)
gmsh.model.geo.addLine(10, 11, 12)
gmsh.model.geo.addLine(8, 11, 13)
gmsh.model.geo.addLine(11, 4, 14)
gmsh.model.geo.addLine(11, 6, 15)

gmsh.model.geo.addCurveLoop([1, 2, -11, -10, 9], 1)
gmsh.model.geo.addCurveLoop([11, 3, -14, -12], 2)
gmsh.model.geo.addCurveLoop([14, 4, 5, -15], 3)
gmsh.model.geo.addCurveLoop([13, 15, 6, 7], 4)
gmsh.model.geo.addCurveLoop([10, 12, -13, 8], 5)
gmsh.model.geo.addPlaneSurface([1], 1)
gmsh.model.geo.addPlaneSurface([2], 2)
gmsh.model.geo.addPlaneSurface([3], 3)
gmsh.model.geo.addPlaneSurface([4], 4)
gmsh.model.geo.addPlaneSurface([5], 5)

gmsh.model.geo.synchronize()

gmsh.model.addPhysicalGroup(1, [7, 8, 9], 1)
gmsh.model.addPhysicalGroup(1, [2, 3, 4], 2)
gmsh.model.addPhysicalGroup(1, [1], 3)
gmsh.model.addPhysicalGroup(1, [5, 6], 4)
gmsh.model.addPhysicalGroup(2, [1], 1)
gmsh.model.addPhysicalGroup(2, [2], 2)
gmsh.model.addPhysicalGroup(2, [3], 3)
gmsh.model.addPhysicalGroup(2, [4], 4)
gmsh.model.addPhysicalGroup(2, [5], 5)

# Generate the mesh on rank 0
if comm_world.rank == 0:
    gmsh.model.mesh.generate(2)

# Get the cell information from the gmsh model
if comm_world.rank == 0:
    topologies = df.io.gmshio.extract_topology_and_markers(gmsh.model)

    num_cell_types = len(topologies.keys())
    cell_information = dict()
    cell_dimensions = np.zeros(num_cell_types, dtype=np.int32)
    for i, element in enumerate(topologies.keys()):
        _, dim, _, num_nodes, _, _ = gmsh.model.mesh.getElementProperties(element)
        cell_information[i] = {"id": element, "dim": dim, "num_nodes": num_nodes}
        cell_dimensions[i] = dim

    # Sort elements by ascending dimension
    perm_sort = np.argsort(cell_dimensions)
    cell_id = cell_information[perm_sort[-1]]["id"]
    num_nodes = cell_information[perm_sort[-1]]["num_nodes"]
    num_nodes = comm_world.bcast([num_nodes], root=0)

    #cells = np.asarray(topologies[cell_id]["topology"], dtype=np.int64)
    cell_values = [np.asarray(topologies[cell_id]["cell_data"], dtype=np.int32)]
else:
    num_nodes = comm_world.bcast([None], root=0)[0]
    #cells = np.empty([0, num_nodes], dtype=np.int32)
    cell_values = [np.empty((0,), dtype=np.int32)]

# choose which region ids to group together (if sufficient ranks are available)
rid_groups = [[1], [2, 5], [3, 4]]

# create a partitioner that partitions the mesh depending on the cell region id
def partitioner(comm, commsize, celltypes, topos):
    nverts = [abs(int(celltype.value)) for celltype in celltypes]
    ncells = sum([int(len(topo)/nverts[t]) for t, topo in enumerate(topos)])
    if commsize == 1:
        # trivial case
        procs = np.zeros(ncells, dtype=np.int32)
        al = df.graph.adjacencylist(procs)
    elif commsize == 2:
        # special case: not a large enough comm so 
        # group the last two rid groups together
        groups = [rid_groups[0], rid_groups[1] + rid_groups[2]]
        procs = np.asarray([g for t in range(len(topos)) for val in cell_values[t] for g, group in enumerate(groups) if val in group], dtype=np.int32)
        al = df.graph.adjacencylist(procs)
    else:
        # arbitrary case
        groups = rid_groups
        if comm.rank == 0:
            # figure out which cells belong together
            cellorder = [None]*ncells
            lgrouptopos = [[[] for topo in topos] for group in groups]
            i = 0
            for g, group in enumerate(groups):
                for t, topo in enumerate(topos):
                    for c, val in enumerate(cell_values[t]):
                        if val in group:
                            lgrouptopos[g][t] += topo[c*nverts[t]:(c+1)*nverts[t]].tolist()
                            cellorder[c] = i # remember the order
                            i += 1
            grouptopos = [[np.array(grouptopo[t], dtype=np.int64) for t in range(len(topos))] for grouptopo in lgrouptopos]
            groupsizes = [max(1, int(sum([len(grouptopo[t])/nverts[t] for t in range(len(topos))])*commsize/ncells)) for grouptopo in grouptopos]
            # correct the logic above so that sum(groupsizes) == commsize
            while sum(groupsizes) < commsize:
                for g in [1, 0, 2]:
                    groupsizes[g] = groupsizes[g] + 1
                    if sum(groupsizes) == commsize: break
            while sum(groupsizes) > commsize:
                for g in [2, 0, 1]:
                    if groupsizes[g] > 1: groupsizes[g] = groupsizes[g] - 1
                    if sum(groupsizes) == commsize: break
            # communicate the groupsizes
            groupsizes = comm.bcast(groupsizes, root=0)
        else:
            grouptopos = [[np.empty((0,), dtype=np.int64) for t in range(len(topos))] for g in range(len(groups))]
            cellorder = []
            groupsizes = comm.bcast([None]*len(groups), root=0)
        if sum(groupsizes) != commsize:
            raise Exception("sum(groupsizes) != commsize")
        values = np.empty((0,), dtype=np.int32)
        # for each group, partition the cells using a cell partitioner
        for g in range(len(groups)):
            lal = df.mesh.create_cell_partitioner(df.mesh.GhostMode.none)(
                comm, groupsizes[g], celltypes, grouptopos[g])
            values = np.concatenate((values, lal.array + sum(groupsizes[:g])), axis=0)
        al = df.graph.adjacencylist(values[cellorder])
    return al

# convert the gmsh mesh to a dolfinx mesh
mesh, ct, ft = df.io.gmshio.model_to_mesh(gmsh.model, comm_world, 0, gdim=2, partitioner=partitioner)

After that, I define a submesh creation function that goes through the steps of setting up the index maps, topology and geometry as you suggested:

def create_submesh_split_comm(mesh, cell_indices, cell_tags=None, facet_tags=None):
    """
    Function to return a submesh based on the cell indices provided using a split comm
    that excludes ranks with no cells.

    Arguments:
      * mesh         - original (parent) mesh
      * cell_indices - cell indices of the parent mesh to include in the submesh

    Keyword Arguments:
      * cell_tags    - cell tags on parent mesh that will be mapped and returned relative to submesh (default=None)
      * facet_tags   - facet tags on parent mesh that will be mapped and returned relative to submesh (default=None)

    Returns:
      * submesh            - submesh of mesh given cell_indices
      * submesh_cell_tags  - cell tags relative to submesh if cell_tags provided (otherwise None)
      * submesh_facet_tags - facet tags relative to submesh if facet_tags provided (otherwise None)
      * submesh_cell_map   - map from submesh cells to parent cells
    """
    tdim = mesh.topology.dim
    fdim = tdim-1

    # create the submesh
    submesh, submesh_cell_map, submesh_vertex_map, submesh_geom_map = df.mesh.create_submesh(mesh, tdim, cell_indices)

    # create a split comm
    new_comm = mesh.comm.Split(submesh.topology.index_map(tdim).size_local>0)

    # map from the rank on the original comm to the new split comm
    # used to map index map owners to their new ranks
    # this will include undefined negative numbers on ranks with no corresponding entry on the split comm that should cause failure if the mesh has ghosts on unexpected ranks
    rank_map = np.asarray(submesh.comm.group.Translate_ranks(None, new_comm.group), dtype=np.int32)

    # create a new topology on the new comm
    new_topo = df.cpp.mesh.Topology(new_comm, submesh.topology.cell_type)

    new_topo_im_0 = df.common.IndexMap(new_comm, 
                                    submesh.topology.index_map(0).size_local, 
                                    submesh.topology.index_map(0).ghosts, 
                                    rank_map[submesh.topology.index_map(0).owners])
    new_topo.set_index_map(0, new_topo_im_0)

    new_topo_im_tdim = df.common.IndexMap(new_comm, 
                                    submesh.topology.index_map(tdim).size_local, 
                                    submesh.topology.index_map(tdim).ghosts, 
                                    rank_map[submesh.topology.index_map(tdim).owners])
    new_topo.set_index_map(tdim, new_topo_im_tdim)

    new_topo.set_connectivity(submesh.topology.connectivity(0,0), 0, 0)
    new_topo.set_connectivity(submesh.topology.connectivity(tdim, 0), tdim, 0)

    gdim = submesh.geometry.dim
    new_geom_im = df.common.IndexMap(new_comm,
                                    submesh.geometry.index_map().size_local,
                                    submesh.geometry.index_map().ghosts,
                                    rank_map[submesh.geometry.index_map().owners])
    new_geom = type(submesh.geometry._cpp_object)(new_geom_im, submesh.geometry.dofmap, 
                                    submesh.geometry.cmap._cpp_object, 
                                    submesh.geometry.x[:,:gdim], 
                                    submesh.geometry.input_global_indices)
    
    # set up a new submesh
    new_submesh = df.mesh.Mesh(type(submesh._cpp_object)(new_comm, new_topo, new_geom), submesh.ufl_domain())
    new_submesh.topology.create_connectivity(fdim, tdim)

    # if cell_tags are provided then map to the new submesh
    new_submesh_cell_tags = None
    if cell_tags is not None:
        submesh_cell_tags_indices = []
        submesh_cell_tags_values  = []
        # loop over the submesh cells, checking if they're included in
        # the parent cell_tags
        for i,parentind in enumerate(submesh_cell_map):
            parent_cell_tags_indices = np.argwhere(cell_tags.indices==parentind)
            if parent_cell_tags_indices.shape[0]>0:
                submesh_cell_tags_indices.append(i)
                submesh_cell_tags_values.append(cell_tags.values[parent_cell_tags_indices[0][0]])
        submesh_cell_tags_indices = np.asarray(submesh_cell_tags_indices)
        submesh_cell_tags_values  = np.asarray(submesh_cell_tags_values)

        # create a new meshtags object
        # indices should already be sorted by construction
        new_submesh_cell_tags = df.mesh.meshtags(new_submesh, tdim, 
                                             submesh_cell_tags_indices, 
                                             submesh_cell_tags_values)
            

    # if facet_tags are provided then map to the new submesh
    new_submesh_facet_tags = None
    if facet_tags is not None:
        # parent facet to vertices adjacency list
        f2vs = mesh.topology.connectivity(fdim, 0)

        # submesh facet to vertices adjaceny list
        new_submesh.topology.create_connectivity(fdim, 0)
        submesh_f2vs = new_submesh.topology.connectivity(fdim, 0)
        # create a map from the parent vertices to the submesh facets
        # (only for the facets that exist in the submesh)
        submesh_parentvs2subf = dict()
        for i in range(submesh_f2vs.num_nodes):
            submesh_parentvs2subf[tuple(sorted([submesh_vertex_map[j] for j in submesh_f2vs.links(i)]))] = i

        # loop over the facet_tags and map from the parent facet to the submesh facet
        # via the vertices, copying over the facet_tag values
        submesh_facet_tags_indices = []
        submesh_facet_tags_values  = []
        for i,parentind in enumerate(facet_tags.indices):
            subind = submesh_parentvs2subf.get(tuple(sorted(f2vs.links(parentind))), None)
            if subind is not None:
                submesh_facet_tags_indices.append(subind)
                submesh_facet_tags_values.append(facet_tags.values[i])
        submesh_facet_tags_indices = np.asarray(submesh_facet_tags_indices)
        submesh_facet_tags_values  = np.asarray(submesh_facet_tags_values)

        perm = np.argsort(submesh_facet_tags_indices)
        new_submesh_facet_tags = df.mesh.meshtags(new_submesh, fdim, 
                                              submesh_facet_tags_indices[perm], 
                                              submesh_facet_tags_values[perm])
    
    return new_submesh, new_submesh_cell_tags, new_submesh_facet_tags
    

This also optionally sets up cell and facet tags on the new com too.

Then to get the submeshes on the new comm I just call this function:

submesh_0, submesh_0_ct, submesh_0_ft = create_submesh_split_comm(mesh, 
                                                    np.concatenate([ct.find(rid) for rid in rid_groups[0]]),
                                                    cell_tags=ct, facet_tags=ft)
submesh_1, submesh_1_ct, submesh_1_ft = create_submesh_split_comm(mesh, 
                                                    np.concatenate([ct.find(rid) for rid in rid_groups[1]]),
                                                    cell_tags=ct, facet_tags=ft)

And, so far, these meshes seem to work. Will update if I find anything needs to change.

Thanks again,
Cian