add test case for stretch factors

Author: alexfikl

pytential/symbolic/primitives.py

@@ -858,7 +858,7 @@ def _small_mat_eigenvalues(mat):
         det_mat = a * d - b * c
 
         # solutions to lambda**2 - tr(A) * lambda + det(A)
-        sqrt_discriminant = cse(sqrt(tr_mat**2 - 4*det_mat))
+        sqrt_discriminant = cse(sqrt(abs(tr_mat**2 - 4*det_mat)))
         return make_obj_array([
             (tr_mat - sqrt_discriminant) / 2,
             (tr_mat + sqrt_discriminant) / 2,

test/test_symbolic.py

@@ -335,6 +335,152 @@ def test_node_reduction(actx_factory):
 # }}}
 
 
+# {{{ test_stretch_factor
+
+def make_simplex_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), (1, 3, 2),
+        (1, 4, 3), (4, 5, 3),
+        ], dtype=np.int32)
+
+    from meshmode.mesh import SimplexElementGroup
+    from meshmode.mesh.generation import make_group_from_vertices
+    grp = make_group_from_vertices(
+            vertices, vertex_indices, order,
+            unit_nodes=None,
+            group_cls=SimplexElementGroup)
+
+    from meshmode.mesh import Mesh
+    return 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)
+
+    from meshmode.mesh import TensorProductElementGroup
+    from meshmode.mesh.generation import make_group_from_vertices
+    grp = make_group_from_vertices(
+            vertices, vertex_indices, order,
+            unit_nodes=None,
+            group_cls=TensorProductElementGroup)
+
+    from meshmode.mesh import Mesh
+    return Mesh(vertices, [grp], 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_form1 = sym.cse(equi_pder.T @ equi_pder, "pd_mat_jtj")
+
+    from pytential.symbolic.primitives import _small_mat_eigenvalues
+    return [
+            sym.cse(sym.sqrt(s), f"mapping_singval_{i}")
+            for i, s in enumerate(_small_mat_eigenvalues(4 * equi_form1))
+            ]
+
+
+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")
+
+    from pytential.symbolic.primitives import _small_mat_eigenvalues
+    return [
+            sym.cse(sym.sqrt(s), f"mapping_singval_{i}")
+            for i, s in enumerate(_small_mat_eigenvalues(4 * form1))
+            ]
+
+
+@pytest.mark.parametrize("order", [4, 8])
+def test_stretch_factor(actx_factory, order, visualize=False):
+    actx = actx_factory()
+
+    def wobble(x):
+        result = np.empty_like(x)
+        result[0] = x[0]
+        result[1] = x[1]
+        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.discretization import Discretization
+    import meshmode.discretization.poly_element as mpoly
+    simplex_discr = Discretization(actx, simplex_mesh,
+            mpoly.InterpolatoryEdgeClusteredGroupFactory(order))
+    square_discr = Discretization(actx, square_mesh,
+            mpoly.InterpolatoryEdgeClusteredGroupFactory(order))
+
+    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)
+
+    if not visualize:
+        return
+
+    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", [
+        ("s0", s0), ("s1", s1),
+        ], overwrite=True)
+
+    vis = make_visualizer(actx, square_discr, order)
+    vis.write_vtk_file(f"stretch_factor_square_{order:02d}.vtu", [
+        ("s0", q0), ("s1", q1),
+        ], overwrite=True)
+
+    # {{{ plot simplex
+
+    import matplotlib.pyplot as plt
+    fig = plt.figure(figsize=(12, 10), dpi=300)
+
+    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)
+
+        ax = fig.gca()
+        im = ax.tricontourf(xi[0], xi[1], sv0, levels=32)
+        im = ax.tricontourf(-xi[0], -xi[1], sv1, levels=32)
+        fig.colorbar(im, ax=ax)
+        fig.savefig(f"stretch_factor_simplex_{order:02d}_{name}")
+        fig.clf()
+
+    # }}}
+
+    # {{{ plot square
+
+    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)
+        fig.colorbar(im, ax=ax)
+        fig.savefig(f"stretch_factor_square_{order:02d}_{name}")
+        fig.clf()
+
+    # }}}
+
+# }}}
+
+
 # You can test individual routines by typing
 # $ python test_symbolic.py 'test_routine()'