def test_partition_mesh(mesh_size, num_parts, num_groups, dim, scramble_partitions): np.random.seed(42) n = (mesh_size, ) * dim from meshmode.mesh.generation import generate_regular_rect_mesh meshes = [ generate_regular_rect_mesh(a=(0 + i, ) * dim, b=(1 + i, ) * dim, n=n) for i in range(num_groups) ] from meshmode.mesh.processing import merge_disjoint_meshes mesh = merge_disjoint_meshes(meshes) if scramble_partitions: part_per_element = np.random.randint(num_parts, size=mesh.nelements) else: pytest.importorskip("pymetis") from meshmode.distributed import get_partition_by_pymetis part_per_element = get_partition_by_pymetis(mesh, num_parts) from meshmode.mesh.processing import partition_mesh # TODO: The same part_per_element array must be used to partition each mesh. # Maybe the interface should be changed to guarantee this. new_meshes = [ partition_mesh(mesh, part_per_element, i) for i in range(num_parts) ] assert mesh.nelements == np.sum( [new_meshes[i][0].nelements for i in range(num_parts)]), \ "part_mesh has the wrong number of elements" assert count_tags(mesh, BTAG_ALL) == np.sum( [count_tags(new_meshes[i][0], BTAG_ALL) for i in range(num_parts)]), \ "part_mesh has the wrong number of BTAG_ALL boundaries" connected_parts = set() for i_local_part, (part_mesh, _) in enumerate(new_meshes): from meshmode.distributed import get_connected_partitions neighbors = get_connected_partitions(part_mesh) for i_remote_part in neighbors: connected_parts.add((i_local_part, i_remote_part)) from meshmode.mesh import BTAG_PARTITION, InterPartitionAdjacencyGroup from meshmode.mesh.processing import find_group_indices num_tags = np.zeros((num_parts, )) index_lookup_table = dict() for ipart, (m, _) in enumerate(new_meshes): for igrp in range(len(m.groups)): adj = m.facial_adjacency_groups[igrp][None] if not isinstance(adj, InterPartitionAdjacencyGroup): # This group is not connected to another partition. continue for i, (elem, face) in enumerate(zip(adj.elements, adj.element_faces)): index_lookup_table[ipart, igrp, elem, face] = i for part_num in range(num_parts): part, part_to_global = new_meshes[part_num] for grp_num in range(len(part.groups)): adj = part.facial_adjacency_groups[grp_num][None] tags = -part.facial_adjacency_groups[grp_num][None].neighbors assert np.all(tags >= 0) if not isinstance(adj, InterPartitionAdjacencyGroup): # This group is not connected to another partition. continue elem_base = part.groups[grp_num].element_nr_base for idx in range(len(adj.elements)): if adj.partition_neighbors[idx] == -1: continue elem = adj.elements[idx] face = adj.element_faces[idx] n_part_num = adj.neighbor_partitions[idx] n_meshwide_elem = adj.partition_neighbors[idx] n_face = adj.neighbor_faces[idx] num_tags[n_part_num] += 1 n_part, n_part_to_global = new_meshes[n_part_num] # Hack: find_igrps expects a numpy.ndarray and returns # a numpy.ndarray. But if a single integer is fed # into find_igrps, an integer is returned. n_grp_num = int( find_group_indices(n_part.groups, n_meshwide_elem)) n_adj = n_part.facial_adjacency_groups[n_grp_num][None] n_elem_base = n_part.groups[n_grp_num].element_nr_base n_elem = n_meshwide_elem - n_elem_base n_idx = index_lookup_table[n_part_num, n_grp_num, n_elem, n_face] assert (part_num == n_adj.neighbor_partitions[n_idx] and elem + elem_base == n_adj.partition_neighbors[n_idx] and face == n_adj.neighbor_faces[n_idx]),\ "InterPartitionAdjacencyGroup is not consistent" _, n_part_to_global = new_meshes[n_part_num] p_meshwide_elem = part_to_global[elem + elem_base] p_meshwide_n_elem = n_part_to_global[n_elem + n_elem_base] p_grp_num = find_group_indices(mesh.groups, p_meshwide_elem) p_n_grp_num = find_group_indices(mesh.groups, p_meshwide_n_elem) p_elem_base = mesh.groups[p_grp_num].element_nr_base p_n_elem_base = mesh.groups[p_n_grp_num].element_nr_base p_elem = p_meshwide_elem - p_elem_base p_n_elem = p_meshwide_n_elem - p_n_elem_base f_groups = mesh.facial_adjacency_groups[p_grp_num] for p_bnd_adj in f_groups.values(): for idx in range(len(p_bnd_adj.elements)): if (p_elem == p_bnd_adj.elements[idx] and face == p_bnd_adj.element_faces[idx]): assert p_n_elem == p_bnd_adj.neighbors[idx],\ "Tag does not give correct neighbor" assert n_face == p_bnd_adj.neighbor_faces[idx],\ "Tag does not give correct neighbor" for i_remote_part in range(num_parts): tag_sum = 0 for i_local_part, (mesh, _) in enumerate(new_meshes): if (i_local_part, i_remote_part) in connected_parts: tag_sum += count_tags(mesh, BTAG_PARTITION(i_remote_part)) assert num_tags[i_remote_part] == tag_sum,\ "part_mesh has the wrong number of BTAG_PARTITION boundaries"
def make_partition_connection(actx, *, local_bdry_conn, i_local_part, remote_bdry_discr, remote_group_infos): """ Connects ``local_bdry_conn`` to a neighboring partition. :arg local_bdry_conn: A :class:`DiscretizationConnection` of the local partition. :arg i_local_part: The partition number of the local partition. :arg remote_bdry_discr: A :class:`~meshmode.discretization.Discretization` of the boundary of the remote partition. :arg remote_group_infos: An array of :class:`meshmode.distributed.RemoteGroupInfo` instances, one per remote volume element group. :returns: A :class:`DirectDiscretizationConnection` that performs data exchange across faces from the remote partition to partition `i_local_part`. .. versionadded:: 2017.1 .. warning:: Interface is not final. """ from meshmode.mesh.processing import find_group_indices from meshmode.discretization.connection import ( DirectDiscretizationConnection, DiscretizationConnectionElementGroup) local_vol_groups = local_bdry_conn.from_discr.mesh.groups part_batches = [[] for _ in local_vol_groups] assert len(local_vol_groups) == len(local_bdry_conn.to_discr.groups) # We need a nested loop over remote and local groups here. # The code assumes that there is the same number of volume and surface groups. # # A weak reason to choose remote as the outer loop is because # InterPartitionAdjacency refers to neighbors by global volume element # numbers, and we only have enough information to resolve those to (group, # group_local_el_nr) for local elements (whereas we have no information # about remote volume elements). # # (See the find_group_indices below.) for rgi in remote_group_infos: rem_ipag = rgi.inter_partition_adj_group indices = (i_local_part == rem_ipag.neighbor_partitions) if not np.any(indices): # Skip because remote group is not connected to i_local_part. continue i_remote_vol_elems = rem_ipag.elements[indices] i_remote_faces = rem_ipag.element_faces[indices] i_local_vol_elems = rem_ipag.partition_neighbors[indices] i_local_faces = rem_ipag.neighbor_faces[indices] del indices i_local_grps = find_group_indices(local_vol_groups, i_local_vol_elems) # {{{ make remote_vol_to_bdry remote_approx_vol_nelements = np.max(rgi.vol_elem_indices) + 1 remote_approx_nfaces = np.max(rgi.bdry_faces) + 1 remote_vol_to_bdry = np.full( (remote_approx_vol_nelements, remote_approx_nfaces), -1, dtype=remote_bdry_discr.mesh.element_id_dtype) remote_vol_to_bdry[rgi.vol_elem_indices, rgi.bdry_faces] = \ rgi.bdry_elem_indices # }}} for i_local_grp in np.unique(i_local_grps): # {{{ come up with matched_{local,remote}_bdry_el_indices local_vol_to_bdry = _make_bdry_el_lookup_table( actx, local_bdry_conn, i_local_grp) local_indices = np.where(i_local_grps == i_local_grp)[0] if len(local_indices) == 0: continue local_grp_vol_elems = ( i_local_vol_elems[local_indices] - local_vol_groups[i_local_grp].element_nr_base) # These are group-local. remote_grp_vol_elems = i_remote_vol_elems[local_indices] matched_local_bdry_el_indices = local_vol_to_bdry[ local_grp_vol_elems, i_local_faces[local_indices]] assert (matched_local_bdry_el_indices >= 0).all() matched_remote_bdry_el_indices = remote_vol_to_bdry[ remote_grp_vol_elems, i_remote_faces[local_indices]] assert (matched_remote_bdry_el_indices >= 0).all() # }}} grp_batches = _make_cross_face_batches( actx, local_bdry_conn.to_discr, remote_bdry_discr, i_local_grp, rem_ipag.igroup, matched_local_bdry_el_indices, matched_remote_bdry_el_indices) part_batches[i_local_grp].extend(grp_batches) return DirectDiscretizationConnection( from_discr=remote_bdry_discr, to_discr=local_bdry_conn.to_discr, groups=[ DiscretizationConnectionElementGroup(batches=grp_batches) for grp_batches in part_batches ], is_surjective=True)
def make_partition_connection(local_bdry_conn, i_local_part, remote_bdry, remote_adj_groups, remote_from_elem_faces, remote_from_elem_indices): """ Connects ``local_bdry_conn`` to a neighboring partition. :arg local_bdry_conn: A :class:`DiscretizationConnection` of the local partition. :arg i_local_part: The partition number of the local partition. :arg remote_adj_groups: A list of :class:`InterPartitionAdjacency`` of the remote partition. :arg remote_bdry: A :class:`Discretization` of the boundary of the remote partition. :arg remote_from_elem_faces: `remote_from_elem_faces[igrp][idx]` gives the face that batch `idx` interpolates from in group `igrp`. :arg remote_from_elem_indices: `remote_from_elem_indices[igrp][idx]` gives a :class:`np.array` of element indices that batch `idx` interpolates from in group `igrp`. :returns: A :class:`DirectDiscretizationConnection` that performs data exchange across faces from the remote partition to partition `i_local_part`. .. versionadded:: 2017.1 .. warning:: Interface is not final. """ from meshmode.mesh.processing import find_group_indices from meshmode.discretization.connection import ( DirectDiscretizationConnection, DiscretizationConnectionElementGroup) local_bdry = local_bdry_conn.to_discr local_groups = local_bdry_conn.from_discr.mesh.groups part_batches = [[] for _ in local_groups] with cl.CommandQueue(local_bdry_conn.cl_context) as queue: for i_remote_grp, adj in enumerate(remote_adj_groups): indices = (i_local_part == adj.neighbor_partitions) if not np.any(indices): # Skip because i_remote_grp is not connected to i_local_part. continue i_remote_faces = adj.element_faces[indices] i_local_meshwide_elems = adj.global_neighbors[indices] i_local_faces = adj.neighbor_faces[indices] i_local_grps = find_group_indices(local_groups, i_local_meshwide_elems) for i_local_grp in np.unique(i_local_grps): elem_base = local_groups[i_local_grp].element_nr_base local_el_lookup = _make_bdry_el_lookup_table( queue, local_bdry_conn, i_local_grp) for i_remote_face in i_remote_faces: index_flags = np.logical_and( i_local_grps == i_local_grp, i_remote_faces == i_remote_face) if not np.any(index_flags): continue remote_bdry_indices = None for idxs, face in zip( remote_from_elem_indices[i_remote_grp], remote_from_elem_faces[i_remote_grp]): if face == i_remote_face: remote_bdry_indices = idxs break assert remote_bdry_indices is not None elems = i_local_meshwide_elems[index_flags] - elem_base faces = i_local_faces[index_flags] local_bdry_indices = local_el_lookup[elems, faces] batches = _make_cross_face_batches( queue, local_bdry, remote_bdry, i_local_grp, i_remote_grp, local_bdry_indices, remote_bdry_indices) part_batches[i_local_grp].extend(batches) return DirectDiscretizationConnection( from_discr=remote_bdry, to_discr=local_bdry, groups=[ DiscretizationConnectionElementGroup(batches=grp_batches) for grp_batches in part_batches ], is_surjective=True)
def make_partition_connection(local_bdry_conn, i_local_part, remote_bdry, remote_adj_groups, remote_from_elem_faces, remote_from_elem_indices): """ Connects ``local_bdry_conn`` to a neighboring partition. :arg local_bdry_conn: A :class:`DiscretizationConnection` of the local partition. :arg i_local_part: The partition number of the local partition. :arg remote_adj_groups: A list of :class:`InterPartitionAdjacency`` of the remote partition. :arg remote_bdry: A :class:`Discretization` of the boundary of the remote partition. :arg remote_from_elem_faces: `remote_from_elem_faces[igrp][idx]` gives the face that batch `idx` interpolates from in group `igrp`. :arg remote_from_elem_indices: `remote_from_elem_indices[igrp][idx]` gives a :class:`np.array` of element indices that batch `idx` interpolates from in group `igrp`. :returns: A :class:`DirectDiscretizationConnection` that performs data exchange across faces from the remote partition to partition `i_local_part`. .. versionadded:: 2017.1 .. warning:: Interface is not final. """ from meshmode.mesh.processing import find_group_indices from meshmode.discretization.connection import ( DirectDiscretizationConnection, DiscretizationConnectionElementGroup) local_bdry = local_bdry_conn.to_discr local_groups = local_bdry_conn.from_discr.mesh.groups part_batches = [[] for _ in local_groups] with cl.CommandQueue(local_bdry_conn.cl_context) as queue: for i_remote_grp, adj in enumerate(remote_adj_groups): indices = (i_local_part == adj.neighbor_partitions) if not np.any(indices): # Skip because i_remote_grp is not connected to i_local_part. continue i_remote_faces = adj.element_faces[indices] i_local_meshwide_elems = adj.global_neighbors[indices] i_local_faces = adj.neighbor_faces[indices] i_local_grps = find_group_indices(local_groups, i_local_meshwide_elems) for i_local_grp in np.unique(i_local_grps): elem_base = local_groups[i_local_grp].element_nr_base local_el_lookup = _make_bdry_el_lookup_table(queue, local_bdry_conn, i_local_grp) for i_remote_face in i_remote_faces: index_flags = np.logical_and(i_local_grps == i_local_grp, i_remote_faces == i_remote_face) if not np.any(index_flags): continue remote_bdry_indices = None for idxs, face in zip(remote_from_elem_indices[i_remote_grp], remote_from_elem_faces[i_remote_grp]): if face == i_remote_face: remote_bdry_indices = idxs break assert remote_bdry_indices is not None elems = i_local_meshwide_elems[index_flags] - elem_base faces = i_local_faces[index_flags] local_bdry_indices = local_el_lookup[elems, faces] batches = _make_cross_face_batches(queue, local_bdry, remote_bdry, i_local_grp, i_remote_grp, local_bdry_indices, remote_bdry_indices) part_batches[i_local_grp].extend(batches) return DirectDiscretizationConnection( from_discr=remote_bdry, to_discr=local_bdry, groups=[DiscretizationConnectionElementGroup(batches=grp_batches) for grp_batches in part_batches], is_surjective=True)