plot stretch factors using high-order vtk

Author: alexfikl

test/test_symbolic.py

@@ -422,30 +422,6 @@ def idx(i, j):
     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_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, "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))
-            ]
-
-
 def make_gmsh_sphere(order: int, cls: type):
     from meshmode.mesh.io import ScriptSource
     from meshmode.mesh import SimplexElementGroup, TensorProductElementGroup
@@ -524,49 +500,104 @@ def make_gmsh_torus(order: int, cls: type):
             )
 
 
+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, "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, mesh_name="torus", visualize=False):
+def test_stretch_factor(actx_factory, order,
+        mesh_name="torus", metric_type="eigs",
+        visualize=False):
     logging.basicConfig(level=logging.INFO)
     actx = actx_factory()
 
     from meshmode.mesh import SimplexElementGroup, TensorProductElementGroup
     if mesh_name == "torus":
-        square_mesh = make_torus_mesh(order, TensorProductElementGroup)
+        quad_mesh = make_torus_mesh(order, TensorProductElementGroup)
         simplex_mesh = make_torus_mesh(order, SimplexElementGroup)
     elif mesh_name == "twisted":
-        square_mesh = make_twisted_mesh(order, TensorProductElementGroup)
+        quad_mesh = make_twisted_mesh(order, TensorProductElementGroup)
         simplex_mesh = make_twisted_mesh(order, SimplexElementGroup)
     elif mesh_name == "sphere":
         from meshmode.mesh.generation import generate_sphere
         simplex_mesh = generate_sphere(1, order=order,
                 uniform_refinement_rounds=1,
                 group_cls=SimplexElementGroup)
-        square_mesh = generate_sphere(1, order=order,
+        quad_mesh = generate_sphere(1, order=order,
                 uniform_refinement_rounds=1,
                 group_cls=TensorProductElementGroup)
     elif mesh_name == "gmsh_sphere":
         simplex_mesh = make_gmsh_sphere(order, cls=SimplexElementGroup)
-        square_mesh = make_gmsh_sphere(order, cls=TensorProductElementGroup)
+        quad_mesh = make_gmsh_sphere(order, cls=TensorProductElementGroup)
     elif mesh_name == "gmsh_torus":
         simplex_mesh = make_gmsh_torus(order, cls=SimplexElementGroup)
-        square_mesh = make_gmsh_torus(order, cls=TensorProductElementGroup)
+        quad_mesh = make_gmsh_torus(order, cls=TensorProductElementGroup)
     else:
         raise ValueError(f"unknown mesh: '{mesh_name}'")
 
+    ambient_dim = 3
+    assert simplex_mesh.ambient_dim == ambient_dim
+    assert quad_mesh.ambient_dim == ambient_dim
+
     from meshmode.discretization import Discretization
     import meshmode.discretization.poly_element as mpoly
     simplex_discr = Discretization(actx, simplex_mesh,
-            mpoly.InterpolatoryQuadratureGroupFactory(order))
-    square_discr = Discretization(actx, square_mesh,
-            mpoly.InterpolatoryQuadratureGroupFactory(order))
+            mpoly.InterpolatoryEquidistantGroupFactory(order))
+    quad_discr = Discretization(actx, quad_mesh,
+            mpoly.InterpolatoryEquidistantGroupFactory(order))
 
     print(f"simplex_discr.ndofs: {simplex_discr.ndofs}")
-    print(f"square_discr.ndofs:  {square_discr.ndofs}")
+    print(f"quad_discr.ndofs:  {quad_discr.ndofs}")
+
+    if metric_type == "eigs":
+        sym_simplex = make_simplex_stretch_factors(ambient_dim)
+        sym_quad = make_quad_stretch_factors(ambient_dim)
+    elif metric_type == "det":
+        form1 = sym.first_fundamental_form(ambient_dim)
+        sym_simplex = np.array([
+            form1[0, 0] * form1[1, 1] - form1[0, 1] * form1[1, 0],
+            sym.Ones(),
+            ], dtype=object)
+        sym_quad = sym_simplex
+    elif metric_type == "trace":
+        form1 = sym.first_fundamental_form(ambient_dim)
+        sym_simplex = np.array([
+            form1[0, 0] + form1[1, 1],
+            sym.Ones(),
+            ], dtype=object)
+        sym_quad = sym_simplex
+    elif metric_type == "norm":
+        form1 = sym.first_fundamental_form(ambient_dim)
+        sym_simplex = np.array([
+            sym.sqrt(sum(form1 * form1)),
+            sym.Ones(),
+            ], dtype=object)
+        sym_quad = sym_simplex
+    else:
+        raise ValueError(f"unknown metric type: '{metric_type}'")
 
-    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, s1 = bind(simplex_discr, sym_simplex)(actx)
+    q0, q1 = bind(quad_discr, sym_quad)(actx)
 
     print("s0")
     print(actx.to_numpy(actx.np.min(s0))[()], actx.to_numpy(actx.np.max(s0))[()])
@@ -578,26 +609,26 @@ def test_stretch_factor(actx_factory, order, mesh_name="torus", visualize=False)
     if not visualize:
         return
 
-    suffix = f"{mesh_name}_stretch_factors_{order:02d}"
+    suffix = f"{mesh_name}_{metric_type}_{order:02d}"
 
     # {{{ plot vtk
 
     from meshmode.discretization.visualization import make_visualizer
-    vis = make_visualizer(actx, simplex_discr, order)
+    vis = make_visualizer(actx, simplex_discr, order, force_equidistant=True)
     vis.write_vtk_file(f"simplex_{suffix}.vtu", [
         ("s0", s0), ("s1", s1),
-        ], overwrite=True)
+        ], overwrite=True, use_high_order=True)
 
-    vis = make_visualizer(actx, square_discr, order)
-    vis.write_vtk_file(f"square_{suffix}.vtu", [
+    vis = make_visualizer(actx, quad_discr, order, force_equidistant=True)
+    vis.write_vtk_file(f"quad_{suffix}.vtu", [
         ("s0", q0), ("s1", q1),
-        ], overwrite=True)
+        ], overwrite=True, use_high_order=True)
 
     # }}}
 
     # {{{ plot reference simplex
 
-    if square_discr.mesh.nelements <= 2:
+    if quad_discr.mesh.nelements <= 2:
         import matplotlib.pyplot as plt
         fig = plt.figure(figsize=(12, 10), dpi=300)
 
@@ -613,17 +644,17 @@ def test_stretch_factor(actx_factory, order, mesh_name="torus", visualize=False)
 
     # }}}
 
-    # {{{ plot reference square
+    # {{{ plot reference quad
 
-    if square_discr.mesh.nelements <= 2:
-        xi = square_discr.groups[0].unit_nodes
+    if quad_discr.mesh.nelements <= 2:
+        xi = quad_discr.groups[0].unit_nodes
         for name, sv in zip(["q0", "q1"], [q0, q1]):
             sv = actx.to_numpy(sv[0])[0]
 
             ax = fig.gca()
             im = ax.tricontourf(xi[0], xi[1], sv, levels=32)
             fig.colorbar(im, ax=ax)
-            fig.savefig(f"square_{suffix}_{name}")
+            fig.savefig(f"quad_{suffix}_{name}")
             fig.clf()
 
     # }}}