def __init__(self,
dcoll: DiscretizationCollection,
remote_rank,
vol_field,
tag=None):
self.tag = self.base_tag
if tag is not None:
self.tag += tag
self.dcoll = dcoll
self.array_context = vol_field.array_context
self.remote_btag = BTAG_PARTITION(remote_rank)
self.bdry_discr = dcoll.discr_from_dd(self.remote_btag)
from grudge.op import project
self.local_dof_array = project(dcoll, "vol", self.remote_btag,
vol_field)
local_data = self.array_context.to_numpy(flatten(self.local_dof_array))
comm = self.dcoll.mpi_communicator
self.send_req = comm.Isend(local_data, remote_rank, tag=self.tag)
self.remote_data_host = np.empty_like(local_data)
self.recv_req = comm.Irecv(self.remote_data_host, remote_rank,
self.tag)
def cross_rank_trace_pairs(discrwb, vec, tag=None):
if (isinstance(vec, np.ndarray)
and vec.dtype.char == "O"
and not isinstance(vec, DOFArray)):
n, = vec.shape
result = {}
for ivec in range(n):
for rank_tpair in _cross_rank_trace_pairs_scalar_field(
discrwb, vec[ivec]):
assert isinstance(rank_tpair.dd.domain_tag, sym.DTAG_BOUNDARY)
assert isinstance(rank_tpair.dd.domain_tag.tag, BTAG_PARTITION)
result[rank_tpair.dd.domain_tag.tag.part_nr, ivec] = rank_tpair
return [
TracePair(
dd=sym.as_dofdesc(sym.DTAG_BOUNDARY(BTAG_PARTITION(remote_rank))),
interior=make_obj_array([
result[remote_rank, i].int for i in range(n)]),
exterior=make_obj_array([
result[remote_rank, i].ext for i in range(n)])
)
for remote_rank in discrwb.connected_ranks()]
else:
return _cross_rank_trace_pairs_scalar_field(discrwb, vec, tag=tag)
def _cross_rank_trace_pairs_scalar_field(dcoll, vec, tag=None):
if isinstance(vec, Number):
return [
TracePair(BTAG_PARTITION(remote_rank), interior=vec, exterior=vec)
for remote_rank in connected_ranks(dcoll)
]
else:
rbcomms = [
_RankBoundaryCommunication(dcoll, remote_rank, vec, tag=tag)
for remote_rank in connected_ranks(dcoll)
]
return [rbcomm.finish() for rbcomm in rbcomms]
def cross_rank_trace_pairs(dcoll: DiscretizationCollection,
ary,
tag=None) -> list:
r"""Get a :class:`list` of *ary* trace pairs for each partition boundary.
For each partition boundary, the field data values in *ary* are
communicated to/from the neighboring partition. Presumably, this
communication is MPI (but strictly speaking, may not be, and this
routine is agnostic to the underlying communication).
For each face on each partition boundary, a
:class:`TracePair` is created with the locally, and
remotely owned partition boundary face data as the `internal`, and `external`
components, respectively. Each of the TracePair components are structured
like *ary*.
:arg ary: a single :class:`~meshmode.dof_array.DOFArray`, or an object
array of :class:`~meshmode.dof_array.DOFArray`\ s
of arbitrary shape.
:returns: a :class:`list` of :class:`TracePair` objects.
"""
if isinstance(ary, np.ndarray):
oshape = ary.shape
comm_vec = ary.flatten()
n, = comm_vec.shape
result = {}
# FIXME: Batch this communication rather than
# doing it in sequence.
for ivec in range(n):
for rank_tpair in _cross_rank_trace_pairs_scalar_field(
dcoll, comm_vec[ivec]):
assert isinstance(rank_tpair.dd.domain_tag,
dof_desc.DTAG_BOUNDARY)
assert isinstance(rank_tpair.dd.domain_tag.tag, BTAG_PARTITION)
result[rank_tpair.dd.domain_tag.tag.part_nr, ivec] = rank_tpair
return [
TracePair(dd=dof_desc.as_dofdesc(
dof_desc.DTAG_BOUNDARY(BTAG_PARTITION(remote_rank))),
interior=make_obj_array(
[result[remote_rank, i].int
for i in range(n)]).reshape(oshape),
exterior=make_obj_array(
[result[remote_rank, i].ext
for i in range(n)]).reshape(oshape))
for remote_rank in connected_ranks(dcoll)
]
else:
return _cross_rank_trace_pairs_scalar_field(dcoll, ary, tag=tag)
def __init__(self,
dcoll: DiscretizationCollection,
array_container: ArrayOrContainerT,
remote_rank, tag=None):
actx = get_container_context_recursively(array_container)
btag = BTAG_PARTITION(remote_rank)
local_bdry_data = project(dcoll, "vol", btag, array_container)
comm = dcoll.mpi_communicator
self.dcoll = dcoll
self.array_context = actx
self.remote_btag = btag
self.bdry_discr = dcoll.discr_from_dd(btag)
self.local_bdry_data = local_bdry_data
self.local_bdry_data_np = \
to_numpy(flatten(self.local_bdry_data, actx), actx)
self.tag = self.base_tag
if tag is not None:
self.tag += tag
# Here, we initialize both send and recieve operations through
# mpi4py `Request` (MPI_Request) instances for comm.Isend (MPI_Isend)
# and comm.Irecv (MPI_Irecv) respectively. These initiate non-blocking
# point-to-point communication requests and require explicit management
# via the use of wait (MPI_Wait, MPI_Waitall, MPI_Waitany, MPI_Waitsome),
# test (MPI_Test, MPI_Testall, MPI_Testany, MPI_Testsome), and cancel
# (MPI_Cancel). The rank-local data `self.local_bdry_data_np` will have its
# associated memory buffer sent across connected ranks and must not be
# modified at the Python level during this process. Completion of the
# requests is handled in :meth:`finish`.
#
# For more details on the mpi4py semantics, see:
# https://mpi4py.readthedocs.io/en/stable/overview.html#nonblocking-communications
#
# NOTE: mpi4py currently (2021-11-03) holds a reference to the send
# memory buffer for (i.e. `self.local_bdry_data_np`) until the send
# requests is complete, however it is not clear that this is documented
# behavior. We hold on to the buffer (via the instance attribute)
# as well, just in case.
self.send_req = comm.Isend(self.local_bdry_data_np,
remote_rank,
tag=self.tag)
self.remote_data_host_numpy = np.empty_like(self.local_bdry_data_np)
self.recv_req = comm.Irecv(self.remote_data_host_numpy,
remote_rank,
tag=self.tag)
def map_operator_binding(self, expr):
from meshmode.mesh import BTAG_PARTITION
from meshmode.discretization.connection import (FACE_RESTR_ALL,
FACE_RESTR_INTERIOR)
if (isinstance(expr.op, op.ProjectionOperator)
and expr.op.dd_in.domain_tag is FACE_RESTR_INTERIOR
and expr.op.dd_out.domain_tag is FACE_RESTR_ALL):
distributed_work = 0
for i_remote_part in self.connected_parts:
mapped_field = RankGeometryChanger(i_remote_part)(expr.field)
btag_part = BTAG_PARTITION(i_remote_part)
distributed_work += op.ProjectionOperator(
dd_in=btag_part, dd_out=expr.op.dd_out)(mapped_field)
return expr + distributed_work
else:
return IdentityMapper.map_operator_binding(self, expr)
def __init__(self, discrwb, remote_rank, vol_field, tag=None):
self.tag = self.base_tag
if tag is not None:
self.tag += tag
self.discrwb = discrwb
self.array_context = vol_field.array_context
self.remote_btag = BTAG_PARTITION(remote_rank)
self.bdry_discr = discrwb.discr_from_dd(self.remote_btag)
self.local_dof_array = discrwb.project("vol", self.remote_btag, vol_field)
local_data = self.array_context.to_numpy(flatten(self.local_dof_array))
comm = self.discrwb.mpi_communicator
self.send_req = comm.Isend(
local_data, remote_rank, tag=self.tag)
self.remote_data_host = np.empty_like(local_data)
self.recv_req = comm.Irecv(self.remote_data_host, remote_rank, self.tag)
def _set_up_distributed_communication(self, mpi_communicator,
array_context):
from_dd = DOFDesc("vol", DISCR_TAG_BASE)
boundary_connections = {}
from meshmode.distributed import get_connected_partitions
connected_parts = get_connected_partitions(self._volume_discr.mesh)
if connected_parts:
if mpi_communicator is None:
raise RuntimeError(
"must supply an MPI communicator when using a "
"distributed mesh")
grp_factory = \
self.group_factory_for_discretization_tag(DISCR_TAG_BASE)
local_boundary_connections = {}
for i_remote_part in connected_parts:
local_boundary_connections[
i_remote_part] = self.connection_from_dds(
from_dd,
DOFDesc(BTAG_PARTITION(i_remote_part), DISCR_TAG_BASE))
from meshmode.distributed import MPIBoundaryCommSetupHelper
with MPIBoundaryCommSetupHelper(mpi_communicator, array_context,
local_boundary_connections,
grp_factory) as bdry_setup_helper:
while True:
conns = bdry_setup_helper.complete_some()
if not conns:
break
for i_remote_part, conn in conns.items():
boundary_connections[i_remote_part] = conn
return boundary_connections
def cross_rank_trace_pairs(
dcoll: DiscretizationCollection, ary, tag=None) -> list:
r"""Get a :class:`list` of *ary* trace pairs for each partition boundary.
For each partition boundary, the field data values in *ary* are
communicated to/from the neighboring partition. Presumably, this
communication is MPI (but strictly speaking, may not be, and this
routine is agnostic to the underlying communication).
For each face on each partition boundary, a
:class:`TracePair` is created with the locally, and
remotely owned partition boundary face data as the `internal`, and `external`
components, respectively. Each of the TracePair components are structured
like *ary*.
If *ary* is a number, rather than a
:class:`~meshmode.dof_array.DOFArray` or an
:class:`~arraycontext.container.ArrayContainer` of them, it is assumed
that the same number is being communicated on every rank.
:arg ary: a :class:`~meshmode.dof_array.DOFArray` or an
:class:`~arraycontext.container.ArrayContainer` of them.
:returns: a :class:`list` of :class:`TracePair` objects.
"""
if isinstance(ary, Number):
# NOTE: Assumed that the same number is passed on every rank
return [TracePair(BTAG_PARTITION(remote_rank), interior=ary, exterior=ary)
for remote_rank in connected_ranks(dcoll)]
# Initialize and post all sends/receives
rank_bdry_communcators = [
_RankBoundaryCommunication(dcoll, ary, remote_rank, tag=tag)
for remote_rank in connected_ranks(dcoll)
]
# Complete send/receives and return communicated data
return [rc.finish() for rc in rank_bdry_communcators]
def test_partition_interpolation(actx_factory, dim, mesh_pars, num_parts,
num_groups, part_method):
np.random.seed(42)
group_factory = PolynomialWarpAndBlendGroupFactory
actx = actx_factory()
order = 4
def f(x):
return 10. * actx.np.sin(50. * x)
for n in mesh_pars:
from meshmode.mesh.generation import generate_warped_rect_mesh
base_mesh = generate_warped_rect_mesh(dim, order=order, n=n)
if num_groups > 1:
from meshmode.mesh.processing import split_mesh_groups
# Group every Nth element
element_flags = np.arange(
base_mesh.nelements,
dtype=base_mesh.element_id_dtype) % num_groups
mesh = split_mesh_groups(base_mesh, element_flags)
else:
mesh = base_mesh
if part_method == "random":
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, connectivity=part_method)
from meshmode.mesh.processing import partition_mesh
part_meshes = [
partition_mesh(mesh, part_per_element, i)[0]
for i in range(num_parts)
]
connected_parts = set()
for i_local_part, part_mesh in enumerate(part_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.discretization import Discretization
vol_discrs = [
Discretization(actx, part_meshes[i], group_factory(order))
for i in range(num_parts)
]
from meshmode.mesh import BTAG_PARTITION
from meshmode.discretization.connection import (
make_face_restriction, make_partition_connection, check_connection)
for i_local_part, i_remote_part in connected_parts:
# Mark faces within local_mesh that are connected to remote_mesh
local_bdry_conn = make_face_restriction(
actx, vol_discrs[i_local_part], group_factory(order),
BTAG_PARTITION(i_remote_part))
# Mark faces within remote_mesh that are connected to local_mesh
remote_bdry_conn = make_face_restriction(
actx, vol_discrs[i_remote_part], group_factory(order),
BTAG_PARTITION(i_local_part))
bdry_nelements = sum(grp.nelements
for grp in local_bdry_conn.to_discr.groups)
remote_bdry_nelements = sum(
grp.nelements for grp in remote_bdry_conn.to_discr.groups)
assert bdry_nelements == remote_bdry_nelements, \
"partitions do not have the same number of connected elements"
local_bdry = local_bdry_conn.to_discr
remote_bdry = remote_bdry_conn.to_discr
from meshmode.distributed import make_remote_group_infos
remote_to_local_conn = make_partition_connection(
actx,
local_bdry_conn=local_bdry_conn,
i_local_part=i_local_part,
remote_bdry_discr=remote_bdry,
remote_group_infos=make_remote_group_infos(
actx, remote_bdry_conn))
# Connect from local mesh to remote mesh
local_to_remote_conn = make_partition_connection(
actx,
local_bdry_conn=remote_bdry_conn,
i_local_part=i_remote_part,
remote_bdry_discr=local_bdry,
remote_group_infos=make_remote_group_infos(
actx, local_bdry_conn))
check_connection(actx, remote_to_local_conn)
check_connection(actx, local_to_remote_conn)
true_local_points = f(thaw(actx, local_bdry.nodes()[0]))
remote_points = local_to_remote_conn(true_local_points)
local_points = remote_to_local_conn(remote_points)
err = actx.np.linalg.norm(true_local_points - local_points, np.inf)
# Can't currently expect exact results due to limitations of
# interpolation "snapping" in DirectDiscretizationConnection's
# _resample_point_pick_indices
assert err < 1e-11
def _test_mpi_boundary_swap(dim, order, num_groups):
from meshmode.distributed import MPIMeshDistributor, MPIBoundaryCommSetupHelper
from mpi4py import MPI
mpi_comm = MPI.COMM_WORLD
i_local_part = mpi_comm.Get_rank()
num_parts = mpi_comm.Get_size()
mesh_dist = MPIMeshDistributor(mpi_comm)
if mesh_dist.is_mananger_rank():
np.random.seed(42)
from meshmode.mesh.generation import generate_warped_rect_mesh
meshes = [
generate_warped_rect_mesh(dim, order=order, n=4)
for _ in range(num_groups)
]
if num_groups > 1:
from meshmode.mesh.processing import merge_disjoint_meshes
mesh = merge_disjoint_meshes(meshes)
else:
mesh = meshes[0]
part_per_element = np.random.randint(num_parts, size=mesh.nelements)
local_mesh = mesh_dist.send_mesh_parts(mesh, part_per_element,
num_parts)
else:
local_mesh = mesh_dist.receive_mesh_part()
group_factory = PolynomialWarpAndBlendGroupFactory(order)
from meshmode.array_context import PyOpenCLArrayContext
cl_ctx = cl.create_some_context()
queue = cl.CommandQueue(cl_ctx)
actx = PyOpenCLArrayContext(queue)
from meshmode.discretization import Discretization
vol_discr = Discretization(actx, local_mesh, group_factory)
from meshmode.distributed import get_connected_partitions
connected_parts = get_connected_partitions(local_mesh)
assert i_local_part not in connected_parts
bdry_setup_helpers = {}
local_bdry_conns = {}
from meshmode.discretization.connection import make_face_restriction
from meshmode.mesh import BTAG_PARTITION
for i_remote_part in connected_parts:
local_bdry_conns[i_remote_part] = make_face_restriction(
actx, vol_discr, group_factory, BTAG_PARTITION(i_remote_part))
setup_helper = bdry_setup_helpers[i_remote_part] = \
MPIBoundaryCommSetupHelper(
mpi_comm, actx, local_bdry_conns[i_remote_part],
i_remote_part, bdry_grp_factory=group_factory)
setup_helper.post_sends()
remote_to_local_bdry_conns = {}
from meshmode.discretization.connection import check_connection
while bdry_setup_helpers:
for i_remote_part, setup_helper in bdry_setup_helpers.items():
if setup_helper.is_setup_ready():
assert bdry_setup_helpers.pop(i_remote_part) is setup_helper
conn = setup_helper.complete_setup()
check_connection(actx, conn)
remote_to_local_bdry_conns[i_remote_part] = conn
break
# FIXME: Not ideal, busy-waits
_test_data_transfer(mpi_comm, actx, local_bdry_conns,
remote_to_local_bdry_conns, connected_parts)
logger.debug("Rank %d exiting", i_local_part)
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 partition_mesh(mesh, part_per_element, part_num):
"""
:arg mesh: A :class:`meshmode.mesh.Mesh` to be partitioned.
:arg part_per_element: A :class:`numpy.ndarray` containing one
integer per element of *mesh* indicating which part of the
partitioned mesh the element is to become a part of.
:arg part_num: The part number of the mesh to return.
:returns: A tuple ``(part_mesh, part_to_global)``, where *part_mesh*
is a :class:`meshmode.mesh.Mesh` that is a partition of mesh, and
*part_to_global* is a :class:`numpy.ndarray` mapping element
numbers on *part_mesh* to ones in *mesh*.
.. versionadded:: 2017.1
"""
assert len(part_per_element) == mesh.nelements, (
"part_per_element must have shape (mesh.nelements,)")
# Contains the indices of the elements requested.
queried_elems = np.where(np.array(part_per_element) == part_num)[0]
num_groups = len(mesh.groups)
new_indices = []
new_nodes = []
# The set of vertex indices we need.
# NOTE: There are two methods for producing required_indices.
# Optimizations may come from further exploring these options.
#index_set = np.array([], dtype=int)
index_sets = np.array([], dtype=set)
skip_groups = []
num_prev_elems = 0
start_idx = 0
for group_num in range(num_groups):
mesh_group = mesh.groups[group_num]
# Find the index of first element in the next group.
end_idx = len(queried_elems)
for idx in range(start_idx, len(queried_elems)):
if queried_elems[idx] - num_prev_elems >= mesh_group.nelements:
end_idx = idx
break
if start_idx == end_idx:
skip_groups.append(group_num)
new_indices.append(np.array([]))
new_nodes.append(np.array([]))
num_prev_elems += mesh_group.nelements
continue
elems = queried_elems[start_idx:end_idx] - num_prev_elems
new_indices.append(mesh_group.vertex_indices[elems])
new_nodes.append(
np.zeros((mesh.ambient_dim, end_idx - start_idx,
mesh_group.nunit_nodes)))
for i in range(mesh.ambient_dim):
for j in range(start_idx, end_idx):
elems = queried_elems[j] - num_prev_elems
new_idx = j - start_idx
new_nodes[group_num][i,
new_idx, :] = mesh_group.nodes[i,
elems, :]
#index_set = np.append(index_set, new_indices[group_num].ravel())
index_sets = np.append(index_sets, set(new_indices[group_num].ravel()))
num_prev_elems += mesh_group.nelements
start_idx = end_idx
# A sorted np.array of vertex indices we need (without duplicates).
#required_indices = np.unique(np.sort(index_set))
required_indices = np.array(list(set.union(*index_sets)))
new_vertices = np.zeros((mesh.ambient_dim, len(required_indices)))
for dim in range(mesh.ambient_dim):
new_vertices[dim] = mesh.vertices[dim][required_indices]
# Our indices need to be in range [0, len(mesh.nelements)].
for group_num in range(num_groups):
if group_num not in skip_groups:
for i in range(len(new_indices[group_num])):
for j in range(len(new_indices[group_num][0])):
original_index = new_indices[group_num][i, j]
new_indices[group_num][i, j] = np.where(
required_indices == original_index)[0]
new_mesh_groups = []
for group_num, mesh_group in enumerate(mesh.groups):
if group_num not in skip_groups:
new_mesh_groups.append(
type(mesh_group)(mesh_group.order,
new_indices[group_num],
new_nodes[group_num],
unit_nodes=mesh_group.unit_nodes))
from meshmode.mesh import BTAG_ALL, BTAG_PARTITION
boundary_tags = [BTAG_PARTITION(n) for n in np.unique(part_per_element)]
from meshmode.mesh import Mesh
part_mesh = Mesh(new_vertices,
new_mesh_groups,
facial_adjacency_groups=None,
boundary_tags=boundary_tags,
is_conforming=mesh.is_conforming)
adj_data = [[] for _ in range(len(part_mesh.groups))]
for igrp, grp in enumerate(part_mesh.groups):
elem_base = grp.element_nr_base
boundary_adj = part_mesh.facial_adjacency_groups[igrp][None]
boundary_elems = boundary_adj.elements
boundary_faces = boundary_adj.element_faces
p_meshwide_elems = queried_elems[boundary_elems + elem_base]
parent_igrps = find_group_indices(mesh.groups, p_meshwide_elems)
for adj_idx, elem in enumerate(boundary_elems):
face = boundary_faces[adj_idx]
tag = -boundary_adj.neighbors[adj_idx]
assert tag >= 0, "Expected boundary tag in adjacency group."
parent_igrp = parent_igrps[adj_idx]
parent_elem_base = mesh.groups[parent_igrp].element_nr_base
parent_elem = p_meshwide_elems[adj_idx] - parent_elem_base
parent_adj = mesh.facial_adjacency_groups[parent_igrp]
for parent_facial_group in parent_adj.values():
indices, = np.nonzero(
parent_facial_group.elements == parent_elem)
for idx in indices:
if (parent_facial_group.neighbors[idx] >= 0 and
parent_facial_group.element_faces[idx] == face):
rank_neighbor = (parent_facial_group.neighbors[idx] +
parent_elem_base)
n_face = parent_facial_group.neighbor_faces[idx]
n_part_num = part_per_element[rank_neighbor]
tag = tag & ~part_mesh.boundary_tag_bit(BTAG_ALL)
tag = tag | part_mesh.boundary_tag_bit(
BTAG_PARTITION(n_part_num))
boundary_adj.neighbors[adj_idx] = -tag
# Find the neighbor element from the other partition.
n_meshwide_elem = np.count_nonzero(
part_per_element[:rank_neighbor] == n_part_num)
adj_data[igrp].append(
(elem, face, n_part_num, n_meshwide_elem, n_face))
connected_mesh = part_mesh.copy()
from meshmode.mesh import InterPartitionAdjacencyGroup
for igrp, adj in enumerate(adj_data):
if adj:
bdry = connected_mesh.facial_adjacency_groups[igrp][None]
# Initialize connections
n_parts = np.zeros_like(bdry.elements)
n_parts.fill(-1)
global_n_elems = np.copy(n_parts)
n_faces = np.copy(n_parts)
# Sort both sets of elements so that we can quickly merge
# the two data structures
bdry_perm = np.lexsort([bdry.element_faces, bdry.elements])
elems = bdry.elements[bdry_perm]
faces = bdry.element_faces[bdry_perm]
neighbors = bdry.neighbors[bdry_perm]
adj_elems, adj_faces, adj_n_parts, adj_gl_n_elems, adj_n_faces =\
np.array(adj).T
adj_perm = np.lexsort([adj_faces, adj_elems])
adj_elems = adj_elems[adj_perm]
adj_faces = adj_faces[adj_perm]
adj_n_parts = adj_n_parts[adj_perm]
adj_gl_n_elems = adj_gl_n_elems[adj_perm]
adj_n_faces = adj_n_faces[adj_perm]
# Merge interpartition adjacency data with FacialAdjacencyGroup
adj_idx = 0
for bdry_idx in range(len(elems)):
if adj_idx >= len(adj_elems):
break
if (adj_elems[adj_idx] == elems[bdry_idx]
and adj_faces[adj_idx] == faces[bdry_idx]):
n_parts[bdry_idx] = adj_n_parts[adj_idx]
global_n_elems[bdry_idx] = adj_gl_n_elems[adj_idx]
n_faces[bdry_idx] = adj_n_faces[adj_idx]
adj_idx += 1
connected_mesh.facial_adjacency_groups[igrp][None] =\
InterPartitionAdjacencyGroup(elements=elems,
element_faces=faces,
neighbors=neighbors,
igroup=bdry.igroup,
ineighbor_group=None,
neighbor_partitions=n_parts,
global_neighbors=global_n_elems,
neighbor_faces=n_faces)
return connected_mesh, queried_elems
def test_partition_interpolation(ctx_factory, dim, mesh_pars, num_parts,
num_groups, scramble_partitions):
np.random.seed(42)
group_factory = PolynomialWarpAndBlendGroupFactory
cl_ctx = ctx_factory()
queue = cl.CommandQueue(cl_ctx)
order = 4
from pytools.convergence import EOCRecorder
eoc_rec = dict()
for i in range(num_parts):
for j in range(num_parts):
if i == j:
continue
eoc_rec[i, j] = EOCRecorder()
def f(x):
return 10. * cl.clmath.sin(50. * x)
for n in mesh_pars:
from meshmode.mesh.generation import generate_warped_rect_mesh
meshes = [
generate_warped_rect_mesh(dim, order=order, n=n)
for _ in range(num_groups)
]
if num_groups > 1:
from meshmode.mesh.processing import merge_disjoint_meshes
mesh = merge_disjoint_meshes(meshes)
else:
mesh = meshes[0]
if scramble_partitions:
part_per_element = np.random.randint(num_parts,
size=mesh.nelements)
else:
from pymetis import part_graph
_, p = part_graph(
num_parts,
xadj=mesh.nodal_adjacency.neighbors_starts.tolist(),
adjncy=mesh.nodal_adjacency.neighbors.tolist())
part_per_element = np.array(p)
from meshmode.mesh.processing import partition_mesh
part_meshes = [
partition_mesh(mesh, part_per_element, i)[0]
for i in range(num_parts)
]
from meshmode.discretization import Discretization
vol_discrs = [
Discretization(cl_ctx, part_meshes[i], group_factory(order))
for i in range(num_parts)
]
from meshmode.mesh import BTAG_PARTITION
from meshmode.discretization.connection import (
make_face_restriction, make_partition_connection, check_connection)
for i_local_part, i_remote_part in eoc_rec.keys():
if eoc_rec[i_local_part, i_remote_part] is None:
continue
# Mark faces within local_mesh that are connected to remote_mesh
local_bdry_conn = make_face_restriction(
vol_discrs[i_local_part], group_factory(order),
BTAG_PARTITION(i_remote_part))
# If these parts are not connected, don't bother checking the error
bdry_nodes = local_bdry_conn.to_discr.nodes()
if bdry_nodes.size == 0:
eoc_rec[i_local_part, i_remote_part] = None
continue
# Mark faces within remote_mesh that are connected to local_mesh
remote_bdry_conn = make_face_restriction(
vol_discrs[i_remote_part], group_factory(order),
BTAG_PARTITION(i_local_part))
assert bdry_nodes.size == remote_bdry_conn.to_discr.nodes().size, \
"partitions do not have the same number of connected nodes"
# Gather just enough information for the connection
local_bdry = local_bdry_conn.to_discr
local_mesh = part_meshes[i_local_part]
local_adj_groups = [
local_mesh.facial_adjacency_groups[i][None]
for i in range(len(local_mesh.groups))
]
local_batches = [
local_bdry_conn.groups[i].batches
for i in range(len(local_mesh.groups))
]
local_from_elem_faces = [[
batch.to_element_face for batch in grp_batches
] for grp_batches in local_batches]
local_from_elem_indices = [[
batch.to_element_indices.get(queue=queue)
for batch in grp_batches
] for grp_batches in local_batches]
remote_bdry = remote_bdry_conn.to_discr
remote_mesh = part_meshes[i_remote_part]
remote_adj_groups = [
remote_mesh.facial_adjacency_groups[i][None]
for i in range(len(remote_mesh.groups))
]
remote_batches = [
remote_bdry_conn.groups[i].batches
for i in range(len(remote_mesh.groups))
]
remote_from_elem_faces = [[
batch.to_element_face for batch in grp_batches
] for grp_batches in remote_batches]
remote_from_elem_indices = [[
batch.to_element_indices.get(queue=queue)
for batch in grp_batches
] for grp_batches in remote_batches]
# Connect from remote_mesh to local_mesh
remote_to_local_conn = make_partition_connection(
local_bdry_conn, i_local_part, remote_bdry, remote_adj_groups,
remote_from_elem_faces, remote_from_elem_indices)
# Connect from local mesh to remote mesh
local_to_remote_conn = make_partition_connection(
remote_bdry_conn, i_remote_part, local_bdry, local_adj_groups,
local_from_elem_faces, local_from_elem_indices)
check_connection(remote_to_local_conn)
check_connection(local_to_remote_conn)
true_local_points = f(local_bdry.nodes()[0].with_queue(queue))
remote_points = local_to_remote_conn(queue, true_local_points)
local_points = remote_to_local_conn(queue, remote_points)
err = la.norm((true_local_points - local_points).get(), np.inf)
eoc_rec[i_local_part, i_remote_part].add_data_point(1. / n, err)
for (i, j), e in eoc_rec.items():
if e is not None:
print("Error of connection from part %i to part %i." % (i, j))
print(e)
assert (e.order_estimate() >= order - 0.5 or e.max_error() < 1e-11)
def test_partition_interpolation(ctx_factory, dim, mesh_pars,
num_parts, num_groups, part_method):
np.random.seed(42)
group_factory = PolynomialWarpAndBlendGroupFactory
cl_ctx = ctx_factory()
queue = cl.CommandQueue(cl_ctx)
actx = PyOpenCLArrayContext(queue)
order = 4
def f(x):
return 10.*actx.np.sin(50.*x)
for n in mesh_pars:
from meshmode.mesh.generation import generate_warped_rect_mesh
base_mesh = generate_warped_rect_mesh(dim, order=order, n=n)
if num_groups > 1:
from meshmode.mesh.processing import split_mesh_groups
# Group every Nth element
element_flags = np.arange(base_mesh.nelements,
dtype=base_mesh.element_id_dtype) % num_groups
mesh = split_mesh_groups(base_mesh, element_flags)
else:
mesh = base_mesh
if part_method == "random":
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,
connectivity=part_method)
from meshmode.mesh.processing import partition_mesh
part_meshes = [
partition_mesh(mesh, part_per_element, i)[0] for i in range(num_parts)]
connected_parts = set()
for i_local_part, part_mesh in enumerate(part_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.discretization import Discretization
vol_discrs = [Discretization(actx, part_meshes[i], group_factory(order))
for i in range(num_parts)]
from meshmode.mesh import BTAG_PARTITION
from meshmode.discretization.connection import (make_face_restriction,
make_partition_connection,
check_connection)
for i_local_part, i_remote_part in connected_parts:
# Mark faces within local_mesh that are connected to remote_mesh
local_bdry_conn = make_face_restriction(actx, vol_discrs[i_local_part],
group_factory(order),
BTAG_PARTITION(i_remote_part))
# Mark faces within remote_mesh that are connected to local_mesh
remote_bdry_conn = make_face_restriction(actx, vol_discrs[i_remote_part],
group_factory(order),
BTAG_PARTITION(i_local_part))
bdry_nelements = sum(
grp.nelements for grp in local_bdry_conn.to_discr.groups)
remote_bdry_nelements = sum(
grp.nelements for grp in remote_bdry_conn.to_discr.groups)
assert bdry_nelements == remote_bdry_nelements, \
"partitions do not have the same number of connected elements"
# Gather just enough information for the connection
local_bdry = local_bdry_conn.to_discr
local_mesh = part_meshes[i_local_part]
local_adj_groups = [local_mesh.facial_adjacency_groups[i][None]
for i in range(len(local_mesh.groups))]
local_batches = [local_bdry_conn.groups[i].batches
for i in range(len(local_mesh.groups))]
local_from_elem_faces = [[batch.to_element_face
for batch in grp_batches]
for grp_batches in local_batches]
local_from_elem_indices = [[batch.to_element_indices.get(queue=queue)
for batch in grp_batches]
for grp_batches in local_batches]
remote_bdry = remote_bdry_conn.to_discr
remote_mesh = part_meshes[i_remote_part]
remote_adj_groups = [remote_mesh.facial_adjacency_groups[i][None]
for i in range(len(remote_mesh.groups))]
remote_batches = [remote_bdry_conn.groups[i].batches
for i in range(len(remote_mesh.groups))]
remote_from_elem_faces = [[batch.to_element_face
for batch in grp_batches]
for grp_batches in remote_batches]
remote_from_elem_indices = [[batch.to_element_indices.get(queue=queue)
for batch in grp_batches]
for grp_batches in remote_batches]
# Connect from remote_mesh to local_mesh
remote_to_local_conn = make_partition_connection(
actx, local_bdry_conn, i_local_part, remote_bdry,
remote_adj_groups, remote_from_elem_faces,
remote_from_elem_indices)
# Connect from local mesh to remote mesh
local_to_remote_conn = make_partition_connection(
actx, remote_bdry_conn, i_remote_part, local_bdry,
local_adj_groups, local_from_elem_faces,
local_from_elem_indices)
check_connection(actx, remote_to_local_conn)
check_connection(actx, local_to_remote_conn)
true_local_points = f(thaw(actx, local_bdry.nodes()[0]))
remote_points = local_to_remote_conn(true_local_points)
local_points = remote_to_local_conn(remote_points)
err = flat_norm(true_local_points - local_points, np.inf)
# Can't currently expect exact results due to limitations of
# interpolation 'snapping' in DirectDiscretizationConnection's
# _resample_point_pick_indices
assert err < 1e-11
def partition_mesh(mesh, part_per_element, part_num):
"""
:arg mesh: A :class:`~meshmode.mesh.Mesh` to be partitioned.
:arg part_per_element: A :class:`numpy.ndarray` containing one
integer per element of *mesh* indicating which part of the
partitioned mesh the element is to become a part of.
:arg part_num: The part number of the mesh to return.
:returns: A tuple ``(part_mesh, part_to_global)``, where *part_mesh*
is a :class:`~meshmode.mesh.Mesh` that is a partition of mesh, and
*part_to_global* is a :class:`numpy.ndarray` mapping element
numbers on *part_mesh* to ones in *mesh*.
.. versionadded:: 2017.1
"""
assert len(part_per_element) == mesh.nelements, (
"part_per_element must have shape (mesh.nelements,)")
# Contains the indices of the elements requested.
queried_elems = np.where(np.array(part_per_element) == part_num)[0]
global_elem_to_part_elem = _compute_global_elem_to_part_elem(
part_per_element, {part_num}, mesh.element_id_dtype)
# Create new mesh groups that mimick the original mesh's groups but only contain
# the local partition's elements
part_mesh_groups, global_group_to_part_group, required_vertex_indices =\
_filter_mesh_groups(mesh.groups, queried_elems, mesh.vertex_id_dtype)
part_vertices = np.zeros((mesh.ambient_dim, len(required_vertex_indices)))
for dim in range(mesh.ambient_dim):
part_vertices[dim] = mesh.vertices[dim][required_vertex_indices]
part_mesh_group_elem_base = [0 for _ in part_mesh_groups]
el_nr = 0
for i_part_grp, grp in enumerate(part_mesh_groups):
part_mesh_group_elem_base[i_part_grp] = el_nr
el_nr += grp.nelements
local_to_local_adj_groups = _create_local_to_local_adjacency_groups(
mesh, global_elem_to_part_elem, part_mesh_groups,
global_group_to_part_group, part_mesh_group_elem_base)
nonlocal_adj_data = _collect_nonlocal_adjacency_data(
mesh, np.array(part_per_element), global_elem_to_part_elem,
part_mesh_groups, global_group_to_part_group,
part_mesh_group_elem_base)
bdry_data = _collect_bdry_data(mesh, global_elem_to_part_elem,
part_mesh_groups,
global_group_to_part_group,
part_mesh_group_elem_base)
group_neighbor_parts = [
adj.neighbor_parts for adj in nonlocal_adj_data if adj is not None
]
all_neighbor_parts = set(np.unique(np.concatenate(group_neighbor_parts))) if\
group_neighbor_parts else set()
boundary_tags = mesh.boundary_tags[:]
btag_to_index = {tag: i for i, tag in enumerate(boundary_tags)}
def boundary_tag_bit(boundary_tag):
from meshmode.mesh import _boundary_tag_bit
return _boundary_tag_bit(boundary_tags, btag_to_index, boundary_tag)
from meshmode.mesh import BTAG_PARTITION
for i_neighbor_part in all_neighbor_parts:
part_tag = BTAG_PARTITION(i_neighbor_part)
boundary_tags.append(part_tag)
btag_to_index[part_tag] = len(boundary_tags) - 1
inter_partition_adj_groups = _create_inter_partition_adjacency_groups(
mesh, part_per_element, part_mesh_groups, all_neighbor_parts,
nonlocal_adj_data, bdry_data, boundary_tag_bit)
# Combine local and inter-partition/boundary adjacency groups
part_facial_adj_groups = local_to_local_adj_groups
for igrp, facial_adj in enumerate(inter_partition_adj_groups):
part_facial_adj_groups[igrp][None] = facial_adj
from meshmode.mesh import Mesh
part_mesh = Mesh(part_vertices,
part_mesh_groups,
facial_adjacency_groups=part_facial_adj_groups,
boundary_tags=boundary_tags,
is_conforming=mesh.is_conforming)
return part_mesh, queried_elems
def _create_inter_partition_adjacency_groups(mesh, part_per_element,
part_mesh_groups,
all_neighbor_parts,
nonlocal_adj_data, bdry_data,
boundary_tag_bit):
"""
Combine non-local adjacency data and boundary data into inter-partition
adjacency groups.
:arg mesh: A :class:`~meshmode.mesh.Mesh` representing the unpartitioned mesh.
:arg part_per_element: A :class:`numpy.ndarray` mapping element indices to
partition numbers.
:arg part_mesh_groups: An array of `~meshmode.mesh.ElementGroup` instances
representing the partitioned mesh groups.
:arg all_neighbor_parts: A `set` containing all partition numbers that neighbor
the current one.
:arg nonlocal_adj_data: A list of :class:`_NonLocalAdjacencyData` instances, one
for each group in *part_mesh_groups*, containing non-local adjacency data.
:arg bdry_data: A list of :class:`_BoundaryData` instances, one for each group
in *part_mesh_groups*, containing boundary data.
:returns: A list of `~meshmode.mesh.InterPartitionAdjacencyGroup` instances, one
for each group in *part_mesh_groups*, containing the aggregated non-local
adjacency and boundary information from *nonlocal_adj_data* and *bdry_data*.
"""
global_elem_to_neighbor_elem = _compute_global_elem_to_part_elem(
part_per_element, all_neighbor_parts, mesh.element_id_dtype)
inter_partition_adj_groups = []
for i_part_grp in range(len(part_mesh_groups)):
nl = nonlocal_adj_data[i_part_grp]
bdry = bdry_data[i_part_grp]
if nl is None and bdry is None:
# Neither non-local adjacency nor boundary
elements = np.array([], dtype=mesh.element_id_dtype)
element_faces = np.array([], dtype=mesh.face_id_dtype)
neighbor_parts = np.array([], dtype=np.int32)
neighbors = np.array([], dtype=mesh.element_id_dtype)
neighbor_elements = np.array([], dtype=mesh.element_id_dtype)
neighbor_faces = np.array([], dtype=mesh.face_id_dtype)
elif bdry is None:
# Non-local adjacency only
elements = nl.elements
element_faces = nl.element_faces
neighbor_parts = nl.neighbor_parts
neighbors = np.empty_like(elements)
for inonlocal in range(len(neighbors)):
i_neighbor_part = neighbor_parts[inonlocal]
from meshmode.mesh import BTAG_REALLY_ALL, BTAG_PARTITION
neighbors[inonlocal] = -(boundary_tag_bit(BTAG_REALLY_ALL)
| boundary_tag_bit(
BTAG_PARTITION(i_neighbor_part)))
neighbor_elements = global_elem_to_neighbor_elem[
nl.global_neighbors]
neighbor_faces = nl.neighbor_faces
elif nl is None:
# Boundary only
nelems = len(bdry.elements)
elements = bdry.elements
element_faces = bdry.element_faces
neighbor_parts = np.empty(nelems, dtype=np.int32)
neighbor_parts.fill(-1)
neighbors = bdry.neighbors
neighbor_elements = np.empty(nelems, dtype=mesh.element_id_dtype)
neighbor_elements.fill(-1)
neighbor_faces = np.empty(nelems, dtype=mesh.face_id_dtype)
neighbor_faces.fill(-1)
else:
# Both; need to merge together
nnonlocal = len(nl.elements)
nbdry = len(bdry.elements)
nelems = nnonlocal + nbdry
elements = np.empty(nelems, dtype=mesh.element_id_dtype)
element_faces = np.empty(nelems, dtype=mesh.face_id_dtype)
neighbor_parts = np.empty(nelems, dtype=np.int32)
neighbors = np.empty(nelems, dtype=mesh.element_id_dtype)
neighbor_elements = np.empty(nelems, dtype=mesh.element_id_dtype)
neighbor_faces = np.empty(nelems, dtype=mesh.face_id_dtype)
# Combine lists of elements/faces and sort to assist in merging
combined_elements = np.concatenate((nl.elements, bdry.elements))
combined_element_faces = np.concatenate(
(nl.element_faces, bdry.element_faces))
perm = np.lexsort([combined_element_faces, combined_elements])
# Merge non-local part
nonlocal_indices = np.where(perm < nnonlocal)[0]
elements[nonlocal_indices] = nl.elements
element_faces[nonlocal_indices] = nl.element_faces
neighbor_parts[nonlocal_indices] = nl.neighbor_parts
for imerged in nonlocal_indices:
i_neighbor_part = neighbor_parts[imerged]
from meshmode.mesh import BTAG_REALLY_ALL, BTAG_PARTITION
neighbors[imerged] = -(boundary_tag_bit(BTAG_REALLY_ALL)
| boundary_tag_bit(
BTAG_PARTITION(i_neighbor_part)))
neighbor_elements[nonlocal_indices] = global_elem_to_neighbor_elem[
nl.global_neighbors]
neighbor_faces[nonlocal_indices] = nl.neighbor_faces
# Merge boundary part
bdry_indices = np.where(perm >= nnonlocal)[0]
elements[bdry_indices] = bdry.elements
element_faces[bdry_indices] = bdry.element_faces
neighbors[bdry_indices] = bdry.neighbors
neighbor_parts[bdry_indices] = -1
neighbor_elements[bdry_indices] = -1
neighbor_faces[bdry_indices] = -1
from meshmode.mesh import InterPartitionAdjacencyGroup
inter_partition_adj_groups.append(
InterPartitionAdjacencyGroup(igroup=i_part_grp,
ineighbor_group=None,
elements=elements,
element_faces=element_faces,
neighbors=neighbors,
neighbor_partitions=neighbor_parts,
partition_neighbors=neighbor_elements,
neighbor_faces=neighbor_faces))
return inter_partition_adj_groups
def _test_mpi_boundary_swap(dim, order, num_groups):
from meshmode.distributed import MPIMeshDistributor, MPIBoundaryCommSetupHelper
from mpi4py import MPI
mpi_comm = MPI.COMM_WORLD
i_local_part = mpi_comm.Get_rank()
num_parts = mpi_comm.Get_size()
mesh_dist = MPIMeshDistributor(mpi_comm)
if mesh_dist.is_mananger_rank():
np.random.seed(42)
from meshmode.mesh.generation import generate_warped_rect_mesh
meshes = [generate_warped_rect_mesh(dim, order=order, nelements_side=4)
for _ in range(num_groups)]
if num_groups > 1:
from meshmode.mesh.processing import merge_disjoint_meshes
mesh = merge_disjoint_meshes(meshes)
else:
mesh = meshes[0]
part_per_element = np.random.randint(num_parts, size=mesh.nelements)
local_mesh = mesh_dist.send_mesh_parts(mesh, part_per_element, num_parts)
else:
local_mesh = mesh_dist.receive_mesh_part()
group_factory = PolynomialWarpAndBlendGroupFactory(order)
from arraycontext import PyOpenCLArrayContext
cl_ctx = cl.create_some_context()
queue = cl.CommandQueue(cl_ctx)
actx = PyOpenCLArrayContext(queue)
from meshmode.discretization import Discretization
vol_discr = Discretization(actx, local_mesh, group_factory)
from meshmode.distributed import get_connected_partitions
connected_parts = get_connected_partitions(local_mesh)
# Check that the connectivity makes sense before doing any communication
_test_connected_parts(mpi_comm, connected_parts)
from meshmode.discretization.connection import make_face_restriction
from meshmode.mesh import BTAG_PARTITION
local_bdry_conns = {}
for i_remote_part in connected_parts:
local_bdry_conns[i_remote_part] = make_face_restriction(
actx, vol_discr, group_factory, BTAG_PARTITION(i_remote_part))
remote_to_local_bdry_conns = {}
with MPIBoundaryCommSetupHelper(mpi_comm, actx, local_bdry_conns,
bdry_grp_factory=group_factory) as bdry_setup_helper:
from meshmode.discretization.connection import check_connection
while True:
conns = bdry_setup_helper.complete_some()
if not conns:
break
for i_remote_part, conn in conns.items():
check_connection(actx, conn)
remote_to_local_bdry_conns[i_remote_part] = conn
_test_data_transfer(mpi_comm,
actx,
local_bdry_conns,
remote_to_local_bdry_conns,
connected_parts)
logger.debug("Rank %d exiting", i_local_part)
def __init__(self, i_remote_part):
from meshmode.discretization.connection import FACE_RESTR_INTERIOR
from meshmode.mesh import BTAG_PARTITION
self.prev_dd = dof_desc.as_dofdesc(FACE_RESTR_INTERIOR)
self.new_dd = dof_desc.as_dofdesc(BTAG_PARTITION(i_remote_part))
