diff --git a/meshmode/mesh/processing.py b/meshmode/mesh/processing.py
index ff74ac81abdc6e414bd5296b27e8545a047acd25..37e4ac264c28f509623465c1e967fc909c070113 100644
--- a/meshmode/mesh/processing.py
+++ b/meshmode/mesh/processing.py
@@ -31,6 +31,7 @@ import modepy as mp
__doc__ = """
+.. autofunction:: partition_mesh
.. autofunction:: find_volume_mesh_element_orientations
.. autofunction:: perform_flips
.. autofunction:: find_bounding_box
@@ -40,6 +41,109 @@ __doc__ = """
"""
+def partition_mesh(mesh, part_per_element, part_nr):
+ """
+ :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_nr: 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
+
+ .. warning:: Interface is not final. Connectivity between elements
+ across groups needs to be added.
+ """
+ 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_nr)[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_nr in range(num_groups):
+ mesh_group = mesh.groups[group_nr]
+
+ # 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_nr)
+ 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_nr][i, new_idx, :] = mesh_group.nodes[i, elems, :]
+
+ #index_set = np.append(index_set, new_indices[group_nr].ravel())
+ index_sets = np.append(index_sets, set(new_indices[group_nr].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_nr in range(num_groups):
+ if group_nr not in skip_groups:
+ for i in range(len(new_indices[group_nr])):
+ for j in range(len(new_indices[group_nr][0])):
+ original_index = new_indices[group_nr][i, j]
+ new_indices[group_nr][i, j] = np.where(
+ required_indices == original_index)[0]
+
+ new_mesh_groups = []
+ for group_nr in range(num_groups):
+ if group_nr not in skip_groups:
+ mesh_group = mesh.groups[group_nr]
+ new_mesh_groups.append(
+ type(mesh_group)(mesh_group.order, new_indices[group_nr],
+ new_nodes[group_nr], unit_nodes=mesh_group.unit_nodes))
+
+ from meshmode.mesh import Mesh
+ part_mesh = Mesh(new_vertices, new_mesh_groups)
+
+ return (part_mesh, queried_elems)
+
+
# {{{ orientations
def find_volume_mesh_element_group_orientation(vertices, grp):
diff --git a/requirements.txt b/requirements.txt
index 6ed102e6d1309b0e65dcfec56e15aac08448a7d4..f5bb34647a3c1fe4e8989462447a78416879ba4d 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -9,3 +9,6 @@ git+https://github.com/inducer/loopy.git
git+https://github.com/inducer/boxtree.git
git+https://github.com/inducer/sumpy.git
git+https://github.com/inducer/pytential.git
+
+# requires pymetis for tests for partition_mesh
+git+https://github.com/inducer/pymetis.git
diff --git a/test/test_meshmode.py b/test/test_meshmode.py
index 3dabe0cfacde742dd4085293a9206c5a97ae6b7c..a793743ce6ff50cc67aea4b31225a1b9e4df4e3d 100644
--- a/test/test_meshmode.py
+++ b/test/test_meshmode.py
@@ -48,6 +48,55 @@ import logging
logger = logging.getLogger(__name__)
+# {{{ partition_mesh
+
+def test_partition_torus_mesh():
+ from meshmode.mesh.generation import generate_torus
+ my_mesh = generate_torus(2, 1, n_outer=2, n_inner=2)
+
+ part_per_element = np.array([0, 1, 2, 1, 1, 2, 1, 0])
+
+ from meshmode.mesh.processing import partition_mesh
+ (part_mesh0, _) = partition_mesh(my_mesh, part_per_element, 0)
+ (part_mesh1, _) = partition_mesh(my_mesh, part_per_element, 1)
+ (part_mesh2, _) = partition_mesh(my_mesh, part_per_element, 2)
+
+ assert part_mesh0.nelements == 2
+ assert part_mesh1.nelements == 4
+ assert part_mesh2.nelements == 2
+
+
+def test_partition_boxes_mesh():
+ n = 5
+ num_parts = 7
+ from meshmode.mesh.generation import generate_regular_rect_mesh
+ mesh1 = generate_regular_rect_mesh(a=(0, 0, 0), b=(1, 1, 1), n=(n, n, n))
+ mesh2 = generate_regular_rect_mesh(a=(2, 2, 2), b=(3, 3, 3), n=(n, n, n))
+
+ from meshmode.mesh.processing import merge_disjoint_meshes
+ mesh = merge_disjoint_meshes([mesh1, mesh2])
+
+ adjacency_list = np.zeros((mesh.nelements,), dtype=set)
+ for elem in range(mesh.nelements):
+ adjacency_list[elem] = set()
+ starts = mesh.nodal_adjacency.neighbors_starts
+ for n in range(starts[elem], starts[elem + 1]):
+ adjacency_list[elem].add(mesh.nodal_adjacency.neighbors[n])
+
+ from pymetis import part_graph
+ (_, p) = part_graph(num_parts, adjacency=adjacency_list)
+ part_per_element = np.array(p)
+
+ from meshmode.mesh.processing import partition_mesh
+ new_meshes = [
+ partition_mesh(mesh, part_per_element, i)[0] for i in range(num_parts)]
+
+ assert mesh.nelements == np.sum(
+ [new_meshes[i].nelements for i in range(num_parts)])
+
+# }}}
+
+
# {{{ circle mesh
def test_circle_mesh(do_plot=False):