From f4fdbb887ef763b84a6b8368d762b003b3f7dbf2 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 21 Oct 2024 16:15:00 +0100 Subject: [PATCH 1/2] Fix coarsen DM --- firedrake/mg/ufl_utils.py | 34 +++++++++++++++++++--------------- 1 file changed, 19 insertions(+), 15 deletions(-) diff --git a/firedrake/mg/ufl_utils.py b/firedrake/mg/ufl_utils.py index 0408a6d2a6..5a121eb04f 100644 --- a/firedrake/mg/ufl_utils.py +++ b/firedrake/mg/ufl_utils.py @@ -268,20 +268,23 @@ def coarsen_snescontext(context, self, coefficient_mapping=None): coarse._fine = context context._coarse = coarse - # Now that we have the coarse snescontext, push it to the coarsened DMs - # Otherwise they won't have the right transfer manager when they are - # coarsened in turn - for val in chain(coefficient_mapping.values(), (bc.function_arg for bc in problem.bcs)): - if isinstance(val, (firedrake.Function, firedrake.Cofunction)): - V = val.function_space() - coarseneddm = V.dm - parentdm = get_parent(context._problem.u.function_space().dm) - - # Now attach the hook to the parent DM - if get_appctx(coarseneddm) is None: - push_appctx(coarseneddm, coarse) - teardown = partial(pop_appctx, coarseneddm, coarse) - add_hook(parentdm, teardown=teardown) + solution = context._problem.u + solutiondm = solution.function_space().dm + parentdm = get_parent(solutiondm) + if parentdm != solutiondm: + # Now that we have the coarse snescontext, push it to the coarsened DMs + # Otherwise they won't have the right transfer manager when they are + # coarsened in turn + for val in chain(coefficient_mapping.values(), (bc.function_arg for bc in problem.bcs)): + if isinstance(val, (firedrake.Function, firedrake.Cofunction)): + V = val.function_space() + coarseneddm = V.dm + + # Now attach the hook to the parent DM + if get_appctx(coarseneddm) is None: + push_appctx(coarseneddm, coarse) + teardown = partial(pop_appctx, coarseneddm, coarse) + add_hook(parentdm, teardown=teardown) ises = problem.J.arguments()[0].function_space()._ises coarse._nullspace = self(context._nullspace, self, coefficient_mapping=coefficient_mapping) @@ -384,7 +387,8 @@ def create_interpolation(dmc, dmf): mat.setType(mat.Type.PYTHON) mat.setPythonContext(ctx) mat.setUp() - return mat, None + rscale = mat.createVecLeft() if row_size == col_size else None + return mat, rscale def create_injection(dmc, dmf): From 0d1e1e839141ff6e1869ba2734aeb46814b797cb Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 21 Oct 2024 17:45:00 +0100 Subject: [PATCH 2/2] add a test --- firedrake/mg/ufl_utils.py | 30 +++++++++++------------- tests/multigrid/test_transfer_manager.py | 23 ++++++++++++++++++ 2 files changed, 37 insertions(+), 16 deletions(-) diff --git a/firedrake/mg/ufl_utils.py b/firedrake/mg/ufl_utils.py index 5a121eb04f..02649b57b9 100644 --- a/firedrake/mg/ufl_utils.py +++ b/firedrake/mg/ufl_utils.py @@ -268,23 +268,21 @@ def coarsen_snescontext(context, self, coefficient_mapping=None): coarse._fine = context context._coarse = coarse - solution = context._problem.u - solutiondm = solution.function_space().dm + solutiondm = context._problem.u.function_space().dm parentdm = get_parent(solutiondm) - if parentdm != solutiondm: - # Now that we have the coarse snescontext, push it to the coarsened DMs - # Otherwise they won't have the right transfer manager when they are - # coarsened in turn - for val in chain(coefficient_mapping.values(), (bc.function_arg for bc in problem.bcs)): - if isinstance(val, (firedrake.Function, firedrake.Cofunction)): - V = val.function_space() - coarseneddm = V.dm - - # Now attach the hook to the parent DM - if get_appctx(coarseneddm) is None: - push_appctx(coarseneddm, coarse) - teardown = partial(pop_appctx, coarseneddm, coarse) - add_hook(parentdm, teardown=teardown) + # Now that we have the coarse snescontext, push it to the coarsened DMs + # Otherwise they won't have the right transfer manager when they are + # coarsened in turn + for val in chain(coefficient_mapping.values(), (bc.function_arg for bc in problem.bcs)): + if isinstance(val, (firedrake.Function, firedrake.Cofunction)): + V = val.function_space() + coarseneddm = V.dm + + # Now attach the hook to the parent DM + if get_appctx(coarseneddm) is None: + push_appctx(coarseneddm, coarse) + if parentdm.getAttr("__setup_hooks__"): + add_hook(parentdm, teardown=partial(pop_appctx, coarseneddm, coarse)) ises = problem.J.arguments()[0].function_space()._ises coarse._nullspace = self(context._nullspace, self, coefficient_mapping=coefficient_mapping) diff --git a/tests/multigrid/test_transfer_manager.py b/tests/multigrid/test_transfer_manager.py index d11bceb41f..a3049af242 100644 --- a/tests/multigrid/test_transfer_manager.py +++ b/tests/multigrid/test_transfer_manager.py @@ -1,5 +1,6 @@ import pytest import numpy +import warnings from firedrake import * from firedrake.mg.ufl_utils import coarsen from firedrake.utils import complex_mode @@ -131,3 +132,25 @@ def test_transfer_manager_dat_version_cache(action, transfer_op, spaces): else: raise ValueError(f"Unrecognized action {action}") + + +@pytest.mark.parametrize("family, degree", [("CG", 1), ("R", 0)]) +def test_cached_transfer(family, degree): + # Test that we can properly reuse transfers within solve + sp = {"mat_type": "matfree", + "pc_type": "mg", + "mg_coarse_pc_type": "none", + "mg_levels_pc_type": "none"} + + base = UnitSquareMesh(1, 1) + hierarchy = MeshHierarchy(base, 3) + mesh = hierarchy[-1] + + V = FunctionSpace(mesh, family, degree) + u = Function(V) + F = inner(u - 1, TestFunction(V)) * dx + + # This test will fail if we raise this warning + with warnings.catch_warnings(): + warnings.filterwarnings("error", "Creating new TransferManager", RuntimeWarning) + solve(F == 0, u, solver_parameters=sp)