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)