clean up stretch factor example test

Author: alexfikl

test/test_symbolic.py

@@ -337,21 +337,19 @@ def test_node_reduction(actx_factory):
 
 # {{{ test_stretch_factor
 
-def make_stretch_mesh(order, cls):
+def make_twisted_mesh(order, cls):
     vertices = np.array([
-        [-1, -1, 0], [1, -1, 0], [-1, 1, 0], [1, 1, 0], [3, -1, 0], [3, 1, 0],
+        [-1, -1, 0], [1, -1, 0], [-1, 1, 0], [1, 1, 0],
         ], dtype=np.float64).T
 
     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
@@ -362,7 +360,21 @@ def make_stretch_mesh(order, cls):
             unit_nodes=None,
             group_cls=cls)
 
-    return mm.Mesh(vertices, [grp], is_conforming=True)
+    def wobble(x):
+        rx, ry = 2, 0.5
+        theta = np.pi / 4
+
+        result = np.empty_like(x)
+        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
+    mesh = mm.Mesh(vertices, [grp], is_conforming=True)
+    return map_mesh(mesh, wobble)
 
 
 def make_torus_mesh(order, cls, a=2.0, b=1.0, n_major=12, n_minor=6):
@@ -417,7 +429,6 @@ 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
@@ -429,7 +440,7 @@ def make_simplex_stretch_factors(ambient_dim):
 
 def make_quad_stretch_factors(ambient_dim):
     pder = sym.parametrization_derivative_matrix(ambient_dim, ambient_dim - 1)
-    form1 = sym.cse(pder.T @ pder + 1.0e-14, "pd_mat_jtj")
+    form1 = sym.cse(pder.T @ pder, "pd_mat_jtj")
 
     from pytential.symbolic.primitives import _small_mat_eigenvalues
     return [
@@ -442,46 +453,29 @@ def make_quad_stretch_factors(ambient_dim):
 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] = 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 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)
+        square_mesh = make_twisted_mesh(order, TensorProductElementGroup)
+        simplex_mesh = make_twisted_mesh(order, SimplexElementGroup)
 
     from meshmode.discretization import Discretization
     import meshmode.discretization.poly_element as mpoly
     simplex_discr = Discretization(actx, simplex_mesh,
-            mpoly.InterpolatoryEdgeClusteredGroupFactory(order))
+            mpoly.InterpolatoryQuadratureGroupFactory(order))
     square_discr = Discretization(actx, square_mesh,
-            mpoly.InterpolatoryEdgeClusteredGroupFactory(order))
+            mpoly.InterpolatoryQuadratureGroupFactory(order))
+
+    print(f"simplex_discr.ndofs: {simplex_discr.ndofs}")
+    print(f"square_discr.ndofs:  {square_discr.ndofs}")
 
     s0, s1 = bind(simplex_discr,
             make_simplex_stretch_factors(simplex_discr.ambient_dim))(actx)
     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))[()])
@@ -490,18 +484,24 @@ def wobble(x):
     if not visualize:
         return
 
+    suffix = f"stretch_factors_{order:02d}"
+
+    # {{{ plot vtk
+
     from meshmode.discretization.visualization import make_visualizer
     vis = make_visualizer(actx, simplex_discr, order)
-    vis.write_vtk_file(f"stretch_factor_simplex_{order:02d}.vtu", [
+    vis.write_vtk_file(f"simplex_{suffix}.vtu", [
         ("s0", s0), ("s1", s1),
         ], overwrite=True)
 
     vis = make_visualizer(actx, square_discr, order)
-    vis.write_vtk_file(f"stretch_factor_square_{order:02d}.vtu", [
+    vis.write_vtk_file(f"square_{suffix}.vtu", [
         ("s0", q0), ("s1", q1),
         ], overwrite=True)
 
-    # {{{ plot simplex
+    # }}}
+
+    # {{{ plot reference simplex
 
     import matplotlib.pyplot as plt
     fig = plt.figure(figsize=(12, 10), dpi=300)
@@ -513,12 +513,12 @@ def wobble(x):
         ax = fig.gca()
         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.savefig(f"simplex_{suffix}_{name}")
         fig.clf()
 
     # }}}
 
-    # {{{ plot square
+    # {{{ plot reference square
 
     xi = square_discr.groups[0].unit_nodes
     for name, sv in zip(["q0", "q1"], [q0, q1]):
@@ -527,7 +527,7 @@ def wobble(x):
         ax = fig.gca()
         im = ax.tricontourf(xi[0], xi[1], sv, levels=32)
         fig.colorbar(im, ax=ax)
-        fig.savefig(f"stretch_factor_square_{order:02d}_{name}")
+        fig.savefig(f"square_{suffix}_{name}")
         fig.clf()
 
     # }}}