From 55afd1dbb0963f706c11de9bd04ea048efeb85f0 Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Tue, 24 May 2016 18:49:40 +0200
Subject: [PATCH] Add nd_quad_submesh

---
 meshmode/mesh/tools.py | 47 ++++++++++++++++++++++++++++++++++++++++++
 test/test_meshmode.py  | 24 +++++++++++++++++++++
 2 files changed, 71 insertions(+)

diff --git a/meshmode/mesh/tools.py b/meshmode/mesh/tools.py
index a67a0fb..4b46bb8 100644
--- a/meshmode/mesh/tools.py
+++ b/meshmode/mesh/tools.py
@@ -28,6 +28,8 @@ import numpy as np
 from pytools.spatial_btree import SpatialBinaryTreeBucket
 
 
+# {{{ make_element_lookup_tree
+
 def make_element_lookup_tree(mesh, eps=1e-12):
     from meshmode.mesh.processing import find_bounding_box
     bbox_min, bbox_max = find_bounding_box(mesh)
@@ -46,3 +48,48 @@ def make_element_lookup_tree(mesh, eps=1e-12):
             tree.insert((igrp, iel_grp), (el_bbox_min, el_bbox_max))
 
     return tree
+
+# }}}
+
+
+# {{{ nd_quad_submesh
+
+def nd_quad_submesh(node_tuples):
+    """Return a list of tuples of indices into the node list that
+    generate a tesselation of the reference element.
+
+    :arg node_tuples: A list of tuples *(i, j, ...)* of integers
+        indicating node positions inside the unit element. The
+        returned list references indices in this list.
+
+        :func:`pytools.generate_nonnegative_integer_tuples_below`
+        may be used to generate *node_tuples*.
+
+    See also :func:`modepy.tools.simplex_submesh`.
+    """
+
+    from pytools import single_valued, add_tuples
+    dims = single_valued(len(nt) for nt in node_tuples)
+
+    node_dict = dict(
+            (ituple, idx)
+            for idx, ituple in enumerate(node_tuples))
+
+    from pytools import generate_nonnegative_integer_tuples_below as gnitb
+
+    result = []
+    for current in node_tuples:
+        try:
+            result.append(tuple(
+                    node_dict[add_tuples(current, offset)]
+                    for offset in gnitb(2, dims)))
+
+        except KeyError:
+            pass
+
+    return result
+
+# }}}
+
+
+# vim: foldmethod=marker
diff --git a/test/test_meshmode.py b/test/test_meshmode.py
index 2c4a5b4..5eba7fe 100644
--- a/test/test_meshmode.py
+++ b/test/test_meshmode.py
@@ -778,6 +778,30 @@ def test_lookup_tree(do_plot=False):
 # }}}
 
 
+# {{{ test_nd_quad_submesh
+
+@pytest.mark.parametrize("dims", [2, 3, 4])
+def test_nd_quad_submesh(dims):
+    from meshmode.mesh.tools import nd_quad_submesh
+    from pytools import generate_nonnegative_integer_tuples_below as gnitb
+
+    node_tuples = list(gnitb(3, dims))
+
+    for i, nt in enumerate(node_tuples):
+        print(i, nt)
+
+    assert len(node_tuples) == 3**dims
+
+    elements = nd_quad_submesh(node_tuples)
+
+    for e in elements:
+        print(e)
+
+    assert len(elements) == 2**dims
+
+# }}}
+
+
 if __name__ == "__main__":
     import sys
     if len(sys.argv) > 1:
-- 
GitLab