include stretch factor test on torus

Author: alexfikl

test/test_symbolic.py

@@ -337,50 +337,87 @@ def test_node_reduction(actx_factory):
 
 # {{{ test_stretch_factor
 
-def make_simplex_mesh(order):
+def make_stretch_mesh(order, cls):
     vertices = np.array([
         [-1, -1, 0], [1, -1, 0], [-1, 1, 0], [1, 1, 0], [3, -1, 0], [3, 1, 0],
         ], dtype=np.float64).T
-    vertex_indices = np.array([
-        (0, 1, 2), (1, 3, 2),
-        (1, 4, 3), (4, 5, 3),
-        ], dtype=np.int32)
 
-    from meshmode.mesh import SimplexElementGroup
+    import meshmode.mesh as mm
+    if issubclass(cls, mm.SimplexElementGroup):
+        vertex_indices = np.array([
+            (0, 1, 2), (1, 3, 2),
+            # (1, 4, 3), (4, 5, 3),
+            ], dtype=np.int32)
+    elif issubclass(cls, mm.TensorProductElementGroup):
+        vertex_indices = np.array([
+            (0, 1, 2, 3),
+            # (1, 4, 3, 5)
+            ], dtype=np.int32)
+    else:
+        raise ValueError
+
     from meshmode.mesh.generation import make_group_from_vertices
     grp = make_group_from_vertices(
             vertices, vertex_indices, order,
             unit_nodes=None,
-            group_cls=SimplexElementGroup)
+            group_cls=cls)
 
-    from meshmode.mesh import Mesh
-    return Mesh(vertices, [grp], is_conforming=True)
+    return mm.Mesh(vertices, [grp], is_conforming=True)
 
 
-def make_square_mesh(order):
-    vertices = np.array([
-        [-1, -1, 0], [1, -1, 0], [-1, 1, 0], [1, 1, 0],
-        [3, -1, 0], [3, 1, 0],
-        ], dtype=np.float64).T
-    vertex_indices = np.array([
-        (0, 1, 2, 3), (1, 4, 3, 5)
-        ], dtype=np.int32)
+def make_torus_mesh(order, cls, a=2.0, b=1.0, n_major=12, n_minor=6):
+    u, v = np.mgrid[0:2*np.pi:2*np.pi/n_major, 0:2*np.pi:2*np.pi/n_minor]
+
+    x = np.cos(u) * (a + b*np.cos(v))
+    y = np.sin(u) * (a + b*np.cos(v))
+    z = b * np.sin(v)
+    del u, v
+
+    vertices = (
+            np.vstack((x[np.newaxis], y[np.newaxis], z[np.newaxis]))
+            .transpose(0, 2, 1).copy().reshape(3, -1))
+
+    def idx(i, j):
+        return (i % n_major) + (j % n_minor) * n_major
+
+    import meshmode.mesh as mm
+    # i, j = 0, 0
+    i, j = 0, n_minor // 3
+    if issubclass(cls, mm.SimplexElementGroup):
+        vertex_indices = [
+                (idx(i, j), idx(i+1, j), idx(i, j+1)),
+                (idx(i+1, j), idx(i+1, j+1), idx(i, j+1)),
+                ]
+    elif issubclass(cls, mm.TensorProductElementGroup):
+        vertex_indices = [(idx(i, j), idx(i+1, j), idx(i, j+1), idx(i+1, j+1))]
+    else:
+        raise TypeError(f"unsupported 'group_cls': {cls}")
 
-    from meshmode.mesh import TensorProductElementGroup
     from meshmode.mesh.generation import make_group_from_vertices
+    vertex_indices = np.array(vertex_indices, dtype=np.int32)
     grp = make_group_from_vertices(
             vertices, vertex_indices, order,
-            unit_nodes=None,
-            group_cls=TensorProductElementGroup)
+            group_cls=cls)
+
+    # NOTE: project the nodes back to the torus surface
+    u = np.arctan2(grp.nodes[1], grp.nodes[0])
+    v = np.arctan2(
+            grp.nodes[2],
+            grp.nodes[0] * np.cos(u) + grp.nodes[1] * np.sin(u) - a)
 
-    from meshmode.mesh import Mesh
-    return Mesh(vertices, [grp], is_conforming=True)
+    nodes = np.empty_like(grp.nodes)
+    nodes[0] = np.cos(u) * (a + b*np.cos(v))
+    nodes[1] = np.sin(u) * (a + b*np.cos(v))
+    nodes[2] = b * np.sin(v)
+
+    return mm.Mesh(vertices, [grp.copy(nodes=nodes)], is_conforming=True)
 
 
 def make_simplex_stretch_factors(ambient_dim):
     from pytential.symbolic.primitives import \
             _equilateral_parametrization_derivative_matrix
     equi_pder = _equilateral_parametrization_derivative_matrix(ambient_dim)
+    # equi_pder = sym.parametrization_derivative_matrix(ambient_dim, ambient_dim - 1)
     equi_form1 = sym.cse(equi_pder.T @ equi_pder, "pd_mat_jtj")
 
     from pytential.symbolic.primitives import _small_mat_eigenvalues
@@ -406,16 +443,29 @@ def test_stretch_factor(actx_factory, order, visualize=False):
     actx = actx_factory()
 
     def wobble(x):
+        rx, ry = 2, 0.5
+        theta = np.pi / 4
+
         result = np.empty_like(x)
-        result[0] = x[0]
-        result[1] = x[1]
-        result[2] = np.sin(x[1]) * np.sin(x[0])
+        result[0] = rx * (np.cos(theta) * x[0] - np.sin(theta) * x[1])
+        result[1] = np.sin(ry * (np.sin(theta) * x[0] + np.cos(theta) * x[1]))
+        result[2] = x[2]
+        # result[2] = np.sin(x[1]) * np.sin(x[0])
 
         return result
 
-    from meshmode.mesh.processing import map_mesh
-    square_mesh = map_mesh(make_square_mesh(order), wobble)
-    simplex_mesh = map_mesh(make_simplex_mesh(order), wobble)
+    from meshmode.mesh import SimplexElementGroup, TensorProductElementGroup
+    if True:
+        square_mesh = make_torus_mesh(order, TensorProductElementGroup)
+        simplex_mesh = make_torus_mesh(order, SimplexElementGroup)
+    else:
+        from meshmode.mesh.processing import map_mesh
+        square_mesh = map_mesh(
+                make_stretch_mesh(order, TensorProductElementGroup),
+                wobble)
+        simplex_mesh = map_mesh(
+                make_stretch_mesh(order, SimplexElementGroup),
+                wobble)
 
     from meshmode.discretization import Discretization
     import meshmode.discretization.poly_element as mpoly
@@ -429,6 +479,14 @@ def wobble(x):
     q0, q1 = bind(square_discr,
             make_quad_stretch_factors(square_discr.ambient_dim))(actx)
 
+    # s0 = (1 + (np.sqrt(3) / 4) / 4) * s0
+    # s1 = (1 + (np.sqrt(3) / 4) / 4) * s1
+
+    print(actx.to_numpy(actx.np.min(s0))[()], actx.to_numpy(actx.np.max(s0))[()])
+    print(actx.to_numpy(actx.np.min(q0))[()], actx.to_numpy(actx.np.max(q0))[()])
+    print(actx.to_numpy(actx.np.min(s1))[()], actx.to_numpy(actx.np.max(s1))[()])
+    print(actx.to_numpy(actx.np.min(q1))[()], actx.to_numpy(actx.np.max(q1))[()])
+
     if not visualize:
         return
 
@@ -450,13 +508,10 @@ def wobble(x):
 
     xi = simplex_discr.groups[0].unit_nodes
     for name, sv in zip(["s0", "s1"], [s0, s1]):
-        sv0 = actx.to_numpy(sv[0])[0].ravel()
-        sv1 = actx.to_numpy(sv[0])[0].ravel()
-        print(sv0, sv1)
+        sv = actx.to_numpy(sv[0])[0].ravel()
 
         ax = fig.gca()
-        im = ax.tricontourf(xi[0], xi[1], sv0, levels=32)
-        im = ax.tricontourf(-xi[0], -xi[1], sv1, levels=32)
+        im = ax.tricontourf(xi[0], xi[1], sv, levels=32)
         fig.colorbar(im, ax=ax)
         fig.savefig(f"stretch_factor_simplex_{order:02d}_{name}")
         fig.clf()
@@ -468,7 +523,6 @@ def wobble(x):
     xi = square_discr.groups[0].unit_nodes
     for name, sv in zip(["q0", "q1"], [q0, q1]):
         sv = actx.to_numpy(sv[0])[0]
-        print(sv)
 
         ax = fig.gca()
         im = ax.tricontourf(xi[0], xi[1], sv, levels=32)