clean up stretch factors

Author: alexfikl

test/test_symbolic.py

@@ -500,28 +500,45 @@ def make_gmsh_torus(order: int, cls: type):
             )
 
 
-def make_simplex_stretch_factors(ambient_dim):
+def metric_from_form1(form1, metric_type: str):
+    from pytential.symbolic.primitives import _small_mat_eigenvalues
+    s0, s1 = _small_mat_eigenvalues(4 * form1)
+
+    if metric_type == "singvals":
+        return np.array([sym.sqrt(s0), sym.sqrt(s1)], dtype=object)
+    elif metric_type == "det":
+        return np.array([s0 * s1], dtype=object)
+    elif metric_type == "trace":
+        return np.array([s0 + s1], dtype=object)
+    elif metric_type == "norm":
+        return np.array([sym.sqrt(s0**2 + s1**2)], dtype=object)
+    elif metric_type == "aspect":
+        return np.array([
+            (s0 * s1)**(2/3) / (s0**2 + s1**2),
+            ], dtype=object)
+    elif metric_type == "condition":
+        import pymbolic.primitives as prim
+        return np.array([
+            prim.Max((s0, s1)) / prim.Min((s0, s1))
+            ], dtype=object)
+    else:
+        raise ValueError(f"unknown metric type: '{metric_type}'")
+
+
+def make_simplex_stretch_factors(ambient_dim: int, metric_type: str):
     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))
-            ]
+    return metric_from_form1(equi_form1, metric_type)
 
 
-def make_quad_stretch_factors(ambient_dim):
+def make_quad_stretch_factors(ambient_dim: int, metric_type: str):
     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))
-            ]
+    return metric_from_form1(form1, metric_type)
 
 
 @pytest.mark.parametrize("order", [4, 8])
@@ -569,42 +586,22 @@ def test_stretch_factor(actx_factory, order,
     print(f"simplex_discr.ndofs: {simplex_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}'")
+    sym_simplex = make_simplex_stretch_factors(ambient_dim, metric_type)
+    sym_quad = make_quad_stretch_factors(ambient_dim, metric_type)
+
+    s = bind(simplex_discr, sym_simplex)(actx)
+    q = bind(quad_discr, sym_quad)(actx)
 
-    s0, s1 = bind(simplex_discr, sym_simplex)(actx)
-    q0, q1 = bind(quad_discr, sym_quad)(actx)
+    def print_bounds(x, name):
+        for i, si in enumerate(x):
+            print("{}{} [{:.12e}, {:.12e}]".format(
+                name, i,
+                actx.to_numpy(actx.np.min(si))[()],
+                actx.to_numpy(actx.np.min(si))[()]
+                ))
 
-    print("s0")
-    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("s1")
-    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))[()])
+    print_bounds(s, "s")
+    print_bounds(q, "q")
 
     if not visualize:
         return
@@ -615,17 +612,23 @@ def test_stretch_factor(actx_factory, order,
 
     from meshmode.discretization.visualization import make_visualizer
     vis = make_visualizer(actx, simplex_discr, order, force_equidistant=True)
-    vis.write_vtk_file(f"simplex_{suffix}.vtu", [
-        ("s0", s0), ("s1", s1),
-        ], overwrite=True, use_high_order=True)
+    vis.write_vtk_file(f"simplex_{suffix}.vtu",
+            [(f"s{i}", si) for i, si in enumerate(s)],
+            overwrite=True, use_high_order=True)
 
     vis = make_visualizer(actx, quad_discr, order, force_equidistant=True)
-    vis.write_vtk_file(f"quad_{suffix}.vtu", [
-        ("s0", q0), ("s1", q1),
-        ], overwrite=True, use_high_order=True)
+    vis.write_vtk_file(f"quad_{suffix}.vtu",
+            [(f"q{i}", qi) for i, qi in enumerate(q)],
+            overwrite=True, use_high_order=True)
 
     # }}}
 
+    if s.size != 2:
+        return
+
+    s0, s1 = s
+    q0, q1 = q
+
     # {{{ plot reference simplex
 
     if quad_discr.mesh.nelements <= 2: