add tests

Author: pbrubeck

firedrake/mg/embedded.py

@@ -97,7 +97,7 @@ def cache(self, V):
 
     def cache_key(self, V):
         _, level = get_level(V.mesh())
-        return level
+        return (level, )
 
     def V_dof_weights(self, V):
         """Dof weights for averaging projection.

firedrake/preconditioners/pmg.py

@@ -1208,7 +1208,8 @@ def __init__(self, Vc, Vf, Vc_bcs, Vf_bcs):
             self.Vc_bcs = [bc.reconstruct(V=self.Vc, g=0) for bc in self.Vc_bcs]
             self.Vf_bcs = [bc.reconstruct(V=self.Vf, g=0) for bc in self.Vf_bcs]
 
-    def work_function(self, V):
+    @staticmethod
+    def work_function(V):
         if isinstance(V, firedrake.Function):
             return V
         cache = V.mesh()._shared_data_cache["pmg_work_function"]

tests/firedrake/multigrid/test_p_multigrid.py

@@ -73,6 +73,17 @@ def test_reconstruct_degree(tp_mesh, mixed_family):
         assert e == PMGPC.reconstruct_degree(elist[0], degree)
 
 
+def test_work_function_cache(tp_mesh):
+    from firedrake.preconditioners.pmg import StandaloneInterpolationMatrix
+
+    V1 = FunctionSpace(tp_mesh, "Lagrange", 1)
+    V2 = FunctionSpace(tp_mesh, "Lagrange", 1)
+    assert V1 is not V2
+    w1 = StandaloneInterpolationMatrix.work_function(V1)
+    w2 = StandaloneInterpolationMatrix.work_function(V2)
+    assert w1 is w2
+
+
 @pytest.mark.parametrize("family", ["Q", "NCE", "NCF", "DQ"])
 def test_prolong_basic(tp_mesh, family):
     """ Interpolate a constant function between low-order and high-order spaces

tests/firedrake/multigrid/test_transfer_manager.py

@@ -131,3 +131,73 @@ def test_transfer_manager_dat_version_cache(action, transfer_op, spaces):
 
     else:
         raise ValueError(f"Unrecognized action {action}")
+
+
+@pytest.fixture
+def DG_spaces(hierarchy):
+    return tuple(VectorFunctionSpace(mesh, "DG", 1) for mesh in hierarchy)
+
+
+@pytest.fixture
+def RT_spaces(hierarchy):
+    return tuple(RestrictedFunctionSpace(FunctionSpace(mesh, "RT", 1), ["on_boundary"]) for mesh in hierarchy)
+
+
+@pytest.mark.parametrize("action", [
+    "DG_work_function",
+    "DG_work_cofunction",
+    "V_dof_weights",
+    "work_vec",
+    "V_DG_mass",
+    "DG_inv_mass",
+    "V_approx_inv_mass",
+    "V_inv_mass_ksp",
+])
+def test_transfer_manager_cache(action, DG_spaces, RT_spaces):
+    V1 = RT_spaces[0]
+    DG1 = DG_spaces[0]
+
+    V2 = V1.reconstruct(name="V2")
+    DG2 = DG1.reconstruct(name="DG2")
+    assert V2 is not V1
+    assert DG2 is not DG1
+
+    if action == "DG_work_function":
+        w1 = transfer.DG_work(DG1)
+        w2 = transfer.DG_work(DG2)
+        assert w1 is w2
+
+    elif action == "DG_work_cofunction":
+        w1 = transfer.DG_work(DG1.dual())
+        w2 = transfer.DG_work(DG2.dual())
+        assert w1 is w2
+
+    elif action == "V_dof_weights":
+        w1 = transfer.V_dof_weights(V1)
+        w2 = transfer.V_dof_weights(V2)
+        assert w1 is w2
+
+    elif action == "work_vec":
+        w1 = transfer.work_vec(V1)
+        w2 = transfer.work_vec(V2)
+        assert w1 is w2
+
+    elif action == "V_DG_mass":
+        M1 = transfer.V_DG_mass(V1, DG1)
+        M2 = transfer.V_DG_mass(V2, DG2)
+        assert M1 is M2
+
+    elif action == "DG_inv_mass":
+        M1 = transfer.DG_inv_mass(DG1)
+        M2 = transfer.DG_inv_mass(DG2)
+        assert M1 is M2
+
+    elif action == "V_approx_inv_mass":
+        M1 = transfer.V_approx_inv_mass(V1, DG1)
+        M2 = transfer.V_approx_inv_mass(V2, DG2)
+        assert M1 is M2
+
+    elif action == "V_inv_mass_ksp":
+        ksp1 = transfer.V_inv_mass_ksp(V1)
+        ksp2 = transfer.V_inv_mass_ksp(V2)
+        assert ksp1 is ksp2