slate: add test for checking count robustness

firedrakeproject · Jan 31, 2025 · 5aa2a35 · 5aa2a35
1 parent 844d3ac
commit 5aa2a35
Showing 1 changed file with 27 additions and 0 deletions.
diff --git a/tests/firedrake/slate/test_slate_hybridization.py b/tests/firedrake/slate/test_slate_hybridization.py
@@ -612,3 +612,30 @@ def trace_nullsp(T):
 
     uD_solver.solve()
     b_solver.solve()
+
+
+@pytest.mark.parametrize('counts', [(10001, 10002), (10002, 10001)])
+def test_slate_hybridization_count_safe(counts):
+    g_count, c_count = counts
+    mesh = UnitTriangleMesh()
+    BDM = FunctionSpace(mesh, "BDM", 2)
+    DG = FunctionSpace(mesh, "DG", 1)
+    V = BDM * DG
+    VectorDG = VectorFunctionSpace(mesh, 'DG', 0)
+    u = TrialFunction(V)
+    v = TestFunction(V)
+    g = Function(VectorDG, count=g_count)
+    c = Function(DG, count=c_count)
+    a = (
+        inner(u, v) * dx
+        + inner(g[0] * u[0], v[0]) * dx
+        + inner(c * grad(u[0]), grad(v[0])) * dx
+    )
+    sol = Function(V)
+    solver_parameters = {
+        'mat_type': 'matfree',
+        'ksp_type': 'preonly',
+        'pc_type': 'python',
+        'pc_python_type': 'firedrake.HybridizationPC',
+    }
+    solve(lhs(a) == rhs(a), sol, solver_parameters=solver_parameters)