diff --git a/.github/workflows/ngsPETSc.yml b/.github/workflows/ngsPETSc.yml index e754bb2..e08d62f 100644 --- a/.github/workflows/ngsPETSc.yml +++ b/.github/workflows/ngsPETSc.yml @@ -54,6 +54,7 @@ jobs: pytest -v tests/test_pc.py pytest -v tests/test_eps.py pytest -v tests/test_snes.py + pytest -v tests/test_adaptive_multigrid.py - name: Run test suite in parallel run: | diff --git a/ngsPETSc/__init__.py b/ngsPETSc/__init__.py index fdfec7c..7f9a1a8 100644 --- a/ngsPETSc/__init__.py +++ b/ngsPETSc/__init__.py @@ -20,7 +20,10 @@ if firedrake: from ngsPETSc.utils.firedrake.meshes import * from ngsPETSc.utils.firedrake.hierarchies import * - __all__ = __all__ + ["FiredrakeMesh", "NetgenHierarchy"] + from ngsPETSc.utils.firedrake.adaptive_hierarchy import * + from ngsPETSc.utils.firedrake.adaptive_transfer_manager import * + __all__ = __all__ + ["FiredrakeMesh", "NetgenHierarchy", + "AdaptiveMeshHierarchy","AdaptiveTransferManager"] #FEniCSx try: diff --git a/ngsPETSc/utils/firedrake/adaptive_hierarchy.py b/ngsPETSc/utils/firedrake/adaptive_hierarchy.py new file mode 100644 index 0000000..6ab8fbd --- /dev/null +++ b/ngsPETSc/utils/firedrake/adaptive_hierarchy.py @@ -0,0 +1,365 @@ +""" +This module contains the class for the AdaptiveMeshHierarchy and +related helper functions +""" + +from fractions import Fraction +from collections import defaultdict +import numpy as np +from firedrake import ( + Function, + Cofunction, + Mesh, + Submesh, + RelabeledMesh, + HierarchyBase, + FunctionSpace, + conditional, + gt, +) +from firedrake.mg.utils import set_level, get_level + +__all__ = ["AdaptiveMeshHierarchy"] + + +class AdaptiveMeshHierarchy(HierarchyBase): + """ + HierarchyBase for hierarchies of adaptively refined meshes + """ + + def __init__(self, mesh, refinements_per_level=1, nested=True): + self.meshes = tuple(mesh) + self._meshes = tuple(mesh) + self.submesh_hierarchies = [] + self.coarse_to_fine_cells = {} + self.fine_to_coarse_cells = {} + self.fine_to_coarse_cells[Fraction(0, 1)] = None + self.refinements_per_level = refinements_per_level + self.nested = nested + set_level(mesh[0], self, 0) + self.split_cache = {} + + def add_mesh(self, mesh): + """ + Adds newly refined mesh into hierarchy. + Then computes the coarse_to_fine and fine_to_coarse mappings. + Constructs intermediate submesh hierarchies with this. + """ + if mesh.topological_dimension() <= 2: + max_children = 4 + else: + max_children = 16 + self._meshes += tuple(mesh) + self.meshes += tuple(mesh) + coarse_mesh = self.meshes[-2] + level = len(self.meshes) + set_level(self.meshes[-1], self, level - 1) + self._shared_data_cache = defaultdict(dict) + + # extract parent child relationships from netgen meshes + (c2f, f2c) = get_c2f_f2c_fd(mesh, coarse_mesh) + c2f_global_key = Fraction(len(self.meshes) - 2, 1) + f2c_global_key = Fraction(len(self.meshes) - 1, 1) + self.coarse_to_fine_cells[c2f_global_key] = c2f + self.fine_to_coarse_cells[f2c_global_key] = np.array(f2c) + + # split both the fine and coarse meshes into the submeshes + (coarse_splits, fine_splits, num_children) = split_to_submesh( + mesh, coarse_mesh, c2f, f2c + ) + for i in range(1, max_children + 1): + coarse_mesh.mark_entities(coarse_splits[i], i) + mesh.mark_entities(fine_splits[i], int(f"10{i}")) + + coarse_indicators = [ + coarse_splits[i] + for i in range(1, max_children + 1) + ] + coarse_labels = list(range(1, max_children + 1)) + coarse_mesh = RelabeledMesh( + coarse_mesh, + coarse_indicators, + coarse_labels, + name="Relabeled_coarse", + ) + c_subm = { + j: Submesh(coarse_mesh, coarse_mesh.topology_dm.getDimension(), j) + for j in range(1, max_children + 1) + if any(num_children == j) + } + set_level(coarse_mesh, self, level - 2) + + fine_indicators = [ + fine_splits[i] + for i in range(1, max_children + 1) + ] + fine_labels = list(range(1, max_children + 1)) + mesh = RelabeledMesh( + mesh, + fine_indicators, + fine_labels, + ) + f_subm = { + int(str(j)[-2:]): Submesh(mesh, mesh.topology_dm.getDimension(), j) + for j in [int("10" + str(i)) for i in range(1, max_children + 1)] + if any(num_children == int(str(j)[-2:])) + } + set_level(mesh, self, level - 1) + + # update c2f and f2c for submeshes by mapping numberings + # on full mesh to numberings on coarse mesh + parents_per_child_count = [ + len([el for el in c2f if len(el) == j]) + for j in range(1, max_children + 1) + ] # stores number of parents for each amount of children + c2f_adjusted = { + j: np.zeros((num_parents, j)) + for j, num_parents in enumerate(parents_per_child_count, 1) + if num_parents != 0 + } + f2c_adjusted = { + j: np.zeros((num_parents * j, 1)) + for j, num_parents in enumerate(parents_per_child_count, 1) + if num_parents != 0 + } + + coarse_full_to_sub_map = { + i: full_to_sub(coarse_mesh, c_subm[i]) + for i in c_subm + } + fine_full_to_sub_map = { + j: full_to_sub(mesh, f_subm[j]) + for j in f_subm + } + + for i, children in enumerate(c2f): + n = len(children) + if 1 <= n <= max_children: + coarse_id_sub = coarse_full_to_sub_map[n][i] + fine_id_sub = fine_full_to_sub_map[n][np.array(children)] + c2f_adjusted[n][coarse_id_sub] = fine_id_sub + + for j, parent in enumerate(f2c): + n = num_children[parent].item() + if 1 <= n <= max_children: + fine_id_sub = fine_full_to_sub_map[n][j] + coarse_id_sub = coarse_full_to_sub_map[n][parent.item()] + f2c_adjusted[n][fine_id_sub, 0] = coarse_id_sub + + + c2f_subm = { + i: {Fraction(0, 1): c2f_adjusted[i].astype(int)} + for i in c2f_adjusted + } + f2c_subm = {i: {Fraction(1, 1): f2c_adjusted[i]} for i in f2c_adjusted} + + hierarchy_dict = { + i: HierarchyBase( + [c_subm[i], f_subm[i]], c2f_subm[i], f2c_subm[i], nested=True + ) + for i in c_subm + } + self.submesh_hierarchies.append(hierarchy_dict) + + def refine(self, refinements): + """ + Refines and adds mesh if input a boolean vector corresponding to cells + """ + ngmesh = self.meshes[-1].netgen_mesh + for i, el in enumerate(ngmesh.Elements2D()): + el.refine = refinements[i] + + ngmesh.Refine(adaptive=True) + mesh = Mesh(ngmesh) + self.add_mesh(mesh) + + def adapt(self, eta, theta): + """ + Implement Dorfler marking, refines mesh from error estimator + """ + mesh = self.meshes[-1] + W = FunctionSpace(mesh, "DG", 0) + markers = Function(W) + + with eta.dat.vec_ro as eta_: + eta_max = eta_.max()[1] + + should_refine = conditional(gt(eta, theta * eta_max), 1, 0) + markers.interpolate(should_refine) + + refined_mesh = mesh.refine_marked_elements(markers) + self.add_mesh(refined_mesh) + return refined_mesh + + def split_function(self, u, child=True): + """ + Split input function across submeshes + """ + V = u.function_space() + full_mesh = V.mesh() + _, level = get_level(full_mesh) + + ind = 1 if child else 0 + hierarchy_dict = self.submesh_hierarchies[int(level) - ind] + parent_mesh = hierarchy_dict[[*hierarchy_dict][0]].meshes[ind].submesh_parent + parent_space = V.reconstruct(parent_mesh) + u_corr_space = Function(parent_space, val=u.dat) + key = (u, child) + try: + split_functions = self.split_cache[key] + except KeyError: + split_functions = self.split_cache.setdefault(key, {}) + + for i in hierarchy_dict: + try: + f = split_functions[i].zero() + except KeyError: + V_split = V.reconstruct(mesh=hierarchy_dict[i].meshes[ind]) + assert ( + V_split.mesh().submesh_parent + == u_corr_space.function_space().mesh() + ) + f = split_functions.setdefault( + i, + Function(V_split, name=str(i)) + ) + + f.assign(u_corr_space) + return split_functions + + def use_weight(self, V, child): + """ + Counts DoFs across submeshes, computes partition of unity + """ + w = Function(V).assign(1) + splits = self.split_function(w, child) + + self.recombine(splits, w, child) + with w.dat.vec as wvec: + wvec.reciprocal() + return w + + def recombine(self, split_funcs, f, child=True): + """ + Recombines functions on submeshes back full mesh + """ + V = f.function_space() + f.zero() + parent_mesh = ( + split_funcs[[*split_funcs][0]].function_space().mesh().submesh_parent + ) + V_label = V.reconstruct(mesh=parent_mesh) + if isinstance(f, Function): + f_label = Function(V_label, val=f.dat) + elif isinstance(f, Cofunction): + f_label = Cofunction(V_label, val=f.dat) + + for split_label, val in split_funcs.items(): + assert val.function_space().mesh().submesh_parent == parent_mesh + if child: + split_label = int("10" + str(split_label)) + if isinstance(f_label, Function): + f_label.assign(val, allow_missing_dofs=True) + else: + curr = Function(f_label.function_space()).assign( + val, allow_missing_dofs=True + ) + f_label.assign(f_label + curr) # partition of unity for restriction + return f + + +def get_c2f_f2c_fd(mesh, coarse_mesh): + """ + Construct coarse->fine and fine->coarse relations by mapping netgen elements to firedrake ones + """ + ngmesh = mesh.netgen_mesh + num_parents = coarse_mesh.num_cells() + + if mesh.topology_dm.getDimension() == 2: + parents = ngmesh.parentsurfaceelements.NumPy() + elements = ngmesh.Elements2D() + elif mesh.topology_dm.getDimension() == 3: + parents = ngmesh.parentelements.NumPy() + elements = ngmesh.Elements3D() + else: + raise RuntimeError("Adaptivity not implemented in dimension of mesh") + + c2f = [[] for _ in range(num_parents)] + f2c = [[] for _ in range(mesh.num_cells())] + + if parents.shape[0] == 0: + raise RuntimeError("Added mesh has not refined any cells from previous mesh") + for l, _ in enumerate(elements): + if parents[l][0] == -1 or l < num_parents: + f2c[mesh._cell_numbering.getOffset(l)].append( + coarse_mesh._cell_numbering.getOffset(l) + ) + c2f[coarse_mesh._cell_numbering.getOffset(l)].append( + mesh._cell_numbering.getOffset(l) + ) + + elif parents[l][0] < num_parents: + fine_ind = mesh._cell_numbering.getOffset(l) + coarse_ind = coarse_mesh._cell_numbering.getOffset(parents[l][0]) + f2c[fine_ind].append(coarse_ind) + c2f[coarse_ind].append(fine_ind) + + else: + a = parents[parents[l][0]][0] + while a >= num_parents: + a = parents[a][0] + + f2c[mesh._cell_numbering.getOffset(l)].append( + coarse_mesh._cell_numbering.getOffset(a) + ) + c2f[coarse_mesh._cell_numbering.getOffset(a)].append( + mesh._cell_numbering.getOffset(l) + ) + + return c2f, np.array(f2c).astype(int) + + +def split_to_submesh(mesh, coarse_mesh, c2f, f2c): + """ + Computes submesh split from full mesh. + Returns splits which are Functions denoting whether elements + belong to the corresponing submesh (bool) + """ + if mesh.topological_dimension() <= 2: + max_children = 4 + else: + max_children = 16 + V = FunctionSpace(mesh, "DG", 0) + V2 = FunctionSpace(coarse_mesh, "DG", 0) + coarse_splits = { + i: Function(V2, name=f"{i}_elements") for i in range(1, max_children + 1) + } + fine_splits = { + i: Function(V, name=f"{i}_elements") for i in range(1, max_children + 1) + } + num_children = np.zeros((len(c2f))) + + for i, children in enumerate(c2f): + n = len(children) + if 1 <= n <= max_children: + coarse_splits[n].dat.data[i] = 1 + num_children[i] = n + + for i in range(1, max_children + 1): + fine_splits[i].dat.data[num_children[f2c.squeeze()] == i] = 1 + + return coarse_splits, fine_splits, num_children + + +def full_to_sub(mesh, submesh): + """ + Returns the submesh element id associated with the full mesh element id + """ + V1 = FunctionSpace(mesh, "DG", 0) + V2 = FunctionSpace(submesh, "DG", 0) + u1 = Function(V1) + u2 = Function(V2) + u2.dat.data[:] = np.arange(len(u2.dat.data)) + u1.assign(u2, allow_missing_dofs=True) + + return u1.dat.data.astype(int) diff --git a/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py b/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py new file mode 100644 index 0000000..3880c73 --- /dev/null +++ b/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py @@ -0,0 +1,148 @@ +""" +This module contains the AdaptiveTransferManager used to perform +transfer operations on AdaptiveMeshHierarchies +""" +import numpy as np +from firedrake import Function +from firedrake.mg.embedded import TransferManager +from firedrake.mg.utils import get_level + + +__all__ = ("AdaptiveTransferManager",) + + +class AdaptiveTransferManager(TransferManager): + """ + TransferManager for adaptively refined mesh hierarchies + """ + def __init__(self, *, native_transfers=None, use_averaging=True): + super().__init__(native_transfers=native_transfers, use_averaging=use_averaging) + self.tm = TransferManager() + self.weight_cache = {} + self.work_function_cache = {} + self.perm_cache = {} + + def generic_transfer(self, source, target, transfer_op): + """ + Generalized implementation of transfer operations by wrapping + transfer operations from TransferManager() + """ + amh, source_level = get_level(source.function_space().mesh()) + _, target_level = get_level(target.function_space().mesh()) + + # decide order of iteration depending on coarse -> fine or fine -> coarse + order = 1 + if target_level < source_level: + order = -1 + + curr_source = source + if source_level == target_level: + target.assign(source) + return + + for level in range(source_level, target_level, order): + if level + order == target_level: + curr_target = target + else: + target_mesh = amh.meshes[level + order] + curr_space = curr_source.function_space() + target_space = curr_space.reconstruct(mesh=target_mesh) + curr_target = self.get_work_function(target_space) + + if transfer_op == self.tm.restrict: + w = self.get_weight(curr_source.function_space()) + wsource = self.get_work_function(curr_source.function_space()) + with ( + curr_source.dat.vec as svec, + w.dat.vec as wvec, + wsource.dat.vec as wsvec, + ): + wsvec.pointwiseMult(svec, wvec) + curr_source = wsource + + if order == 1: + source_function_splits = amh.split_function(curr_source, child=False) + target_function_splits = amh.split_function(curr_target, child=True) + else: + source_function_splits = amh.split_function(curr_source, child=True) + target_function_splits = amh.split_function(curr_target, child=False) + + for split_label, _ in source_function_splits.items(): + if split_label == 1: + # we don't want to transfer across unsplit parts, + # instead we copy dofs + us_func = source_function_splits[1] + ut_func = target_function_splits[1] + permutations = self.get_perm( + us_func, + ut_func, + transfer_op + ) + ut_func.dat.data_wo[permutations] = us_func.dat.data_ro + else: + transfer_op( + source_function_splits[split_label], + target_function_splits[split_label], + ) + + amh.recombine(target_function_splits, curr_target, child=order + 1) + curr_source = curr_target + + def get_work_function(self, func_space): + """ + Cache for function on function space + """ + try: + return self.work_function_cache[func_space] + except KeyError: + return self.work_function_cache.setdefault(func_space, Function(func_space)) + + def get_weight(self, V_source): + """ + Cache for weights from partition of unity used during restriction + """ + try: + return self.weight_cache[V_source] + except KeyError: + amh, _ = get_level(V_source.mesh()) + return self.weight_cache.setdefault( + V_source, amh.use_weight(V_source, child=True) + ) + + def get_perm(self, unsplit_source, unsplit_target, transfer_op): + """ + Cache permutations of DoFs from unsplit source + to unsplit target. This is used to skip transfer + across unsplit mesh hierarchies + """ + key = (unsplit_source.function_space(), + unsplit_target.function_space()) + try: + return self.perm_cache[key] + except KeyError: + source_nodes = Function(key[0]) + permutation = Function(key[1]) + source_nodes.dat.data[:] = np.arange(len(source_nodes.dat.data)) + transfer_op(source_nodes, permutation) + return self.perm_cache.setdefault( + key, np.rint(permutation.dat.data_ro).astype(int) + ) + + + def prolong(self, uc, uf): + """ + Prolongation of AdaptiveMeshHierarchy + """ + self.generic_transfer(uc, uf, transfer_op=self.tm.prolong) + + def inject(self, uf, uc): + """ + Injection of AdaptiveMeshHierarchy + """ + self.generic_transfer(uf, uc, transfer_op=self.tm.inject) + + def restrict(self, source, target): + """ + Restriction of AdaptiveMeshHierarchy + """ + self.generic_transfer(source, target, transfer_op=self.tm.restrict) diff --git a/tests/test_adaptive_multigrid.py b/tests/test_adaptive_multigrid.py new file mode 100644 index 0000000..6fd4043 --- /dev/null +++ b/tests/test_adaptive_multigrid.py @@ -0,0 +1,376 @@ +""" +Tests for AdaptiveMeshHierarchy +and AdaptiveTransferManager +""" + +import random +import pytest +import numpy as np +from firedrake import ( + Mesh, + MeshHierarchy, + TransferManager, + FunctionSpace, + Function, + SpatialCoordinate, + conditional, + ge, + errornorm, + TestFunction, + assemble, + dx, + Cofunction, + action, + sin, + pi, + DirichletBC, + inner, + grad, + NonlinearVariationalProblem, + NonlinearVariationalSolver, +) +from netgen.occ import WorkPlane, OCCGeometry, Box, Pnt +from ngsPETSc import AdaptiveMeshHierarchy, AdaptiveTransferManager + + +@pytest.fixture(params=[2, 3]) +def amh(request): + """ + Generate AdaptiveMeshHierarchies + """ + dim = request.param + random.seed(1234) + if dim == 2: + wp = WorkPlane() + wp.Rectangle(2, 2) + face = wp.Face() + geo = OCCGeometry(face, dim=2) + maxh = 0.5 + else: + cube = Box(Pnt(0, 0, 0), Pnt(2, 2, 2)) + geo = OCCGeometry(cube, dim=3) + maxh = 0.5 + + ngmesh = geo.GenerateMesh(maxh=maxh) + base = Mesh(ngmesh) + amh_test = AdaptiveMeshHierarchy(base) + + if dim == 2: + els = ngmesh.Elements2D() + else: + els = ngmesh.Elements3D() + + for _ in range(2): + for _, el in enumerate(els): + el.refine = 0 + if random.random() < 0.5: + el.refine = 1 + ngmesh.Refine(adaptive=True) + mesh = Mesh(ngmesh) + amh_test.add_mesh(mesh) + return amh_test + + +@pytest.fixture +def mh_res(): + """ + Generate MeshHierarchy for reference + """ + wp = WorkPlane() + wp.Rectangle(2, 2) + face = wp.Face() + geo = OCCGeometry(face, dim=2) + maxh = 0.5 + ngmesh = geo.GenerateMesh(maxh=maxh) + base = Mesh(ngmesh) + mesh2 = Mesh(ngmesh) + amh_unif = AdaptiveMeshHierarchy([base]) + for _ in range(2): + refs = np.ones(len(ngmesh.Elements2D())) + amh_unif.refine(refs) + mh = MeshHierarchy(mesh2, 2) + + return amh_unif, mh + + +@pytest.fixture +def atm(): + """atm used in tests""" + return AdaptiveTransferManager() + + +@pytest.fixture +def tm(): + """tm used for restrict consistency""" + return TransferManager() + + +@pytest.mark.parametrize("operator", ["prolong", "inject"]) +def test_DG0(amh, atm, operator): # pylint: disable=W0621 + """ + Prolongation & Injection test for DG0 + """ + V_coarse = FunctionSpace(amh[0], "DG", 0) + V_fine = FunctionSpace(amh[-1], "DG", 0) + u_coarse = Function(V_coarse) + u_fine = Function(V_fine) + xc, *_ = SpatialCoordinate(V_coarse.mesh()) + stepc = conditional(ge(xc, 0), 1, 0) + xf, *_ = SpatialCoordinate(V_fine.mesh()) + stepf = conditional(ge(xf, 0), 1, 0) + + if operator == "prolong": + u_coarse.interpolate(stepc) + assert errornorm(stepc, u_coarse) <= 1e-12 + + atm.prolong(u_coarse, u_fine) + assert errornorm(stepf, u_fine) <= 1e-12 + if operator == "inject": + u_fine.interpolate(stepf) + assert errornorm(stepf, u_fine) <= 1e-12 + + atm.inject(u_fine, u_coarse) + assert errornorm(stepc, u_coarse) <= 1e-12 + + +@pytest.mark.parametrize("operator", ["prolong", "inject"]) +def test_CG1(amh, atm, operator): # pylint: disable=W0621 + """ + Prolongation & Injection test for CG1 + """ + V_coarse = FunctionSpace(amh[0], "CG", 1) + V_fine = FunctionSpace(amh[-1], "CG", 1) + u_coarse = Function(V_coarse) + u_fine = Function(V_fine) + xc, *_ = SpatialCoordinate(V_coarse.mesh()) + xf, *_ = SpatialCoordinate(V_fine.mesh()) + + if operator == "prolong": + u_coarse.interpolate(xc) + assert errornorm(xc, u_coarse) <= 1e-12 + + atm.prolong(u_coarse, u_fine) + assert errornorm(xf, u_fine) <= 1e-12 + if operator == "inject": + u_fine.interpolate(xf) + assert errornorm(xf, u_fine) <= 1e-12 + + atm.inject(u_fine, u_coarse) + assert errornorm(xc, u_coarse) <= 1e-12 + + +def test_restrict_consistency(mh_res, atm, tm): # pylint: disable=W0621 + """ + Test restriction consistency of amh with uniform refinement vs mh + """ + amh_unif = mh_res[0] + mh = mh_res[1] + + V_coarse = FunctionSpace(amh_unif[0], "DG", 0) + V_fine = FunctionSpace(amh_unif[-1], "DG", 0) + u_coarse = Function(V_coarse) + u_fine = Function(V_fine) + xc, _ = SpatialCoordinate(V_coarse.mesh()) + + u_coarse.interpolate(xc) + atm.prolong(u_coarse, u_fine) + + rf = assemble(TestFunction(V_fine) * dx) + rc = Cofunction(V_coarse.dual()) + atm.restrict(rf, rc) + + # compare with mesh_hierarchy + xcoarse, _ = SpatialCoordinate(mh[0]) + Vcoarse = FunctionSpace(mh[0], "DG", 0) + Vfine = FunctionSpace(mh[-1], "DG", 0) + + mhuc = Function(Vcoarse) + mhuc.interpolate(xcoarse) + mhuf = Function(Vfine) + tm.prolong(mhuc, mhuf) + + mhrf = assemble(TestFunction(Vfine) * dx) + mhrc = Cofunction(Vcoarse.dual()) + + tm.restrict(mhrf, mhrc) + + assert ( + (assemble(action(mhrc, mhuc)) - assemble(action(mhrf, mhuf))) + / assemble(action(mhrf, mhuf)) + ) <= 1e-12 + assert ( + (assemble(action(rc, u_coarse)) - assemble(action(mhrc, mhuc))) + / assemble(action(mhrc, mhuc)) + ) <= 1e-12 + + +def test_restrict_CG1(amh, atm): # pylint: disable=W0621 + """ + Test restriction with CG1 + """ + V_coarse = FunctionSpace(amh[0], "CG", 1) + V_fine = FunctionSpace(amh[-1], "CG", 1) + u_coarse = Function(V_coarse) + u_fine = Function(V_fine) + xc, *_ = SpatialCoordinate(V_coarse.mesh()) + + u_coarse.interpolate(xc) + atm.prolong(u_coarse, u_fine) + + rf = assemble(TestFunction(V_fine) * dx) + rc = Cofunction(V_coarse.dual()) + atm.restrict(rf, rc) + + assert np.allclose( + assemble(action(rc, u_coarse)), + assemble(action(rf, u_fine)), + rtol=1e-12 + ) + + +def test_restrict_DG0(amh, atm): # pylint: disable=W0621 + """ + Test restriction with DG0 + """ + V_coarse = FunctionSpace(amh[0], "DG", 0) + V_fine = FunctionSpace(amh[-1], "DG", 0) + u_coarse = Function(V_coarse) + u_fine = Function(V_fine) + xc, *_ = SpatialCoordinate(V_coarse.mesh()) + + u_coarse.interpolate(xc) + atm.prolong(u_coarse, u_fine) + + rf = assemble(TestFunction(V_fine) * dx) + rc = Cofunction(V_coarse.dual()) + atm.restrict(rf, rc) + + assert np.allclose( + assemble(action(rc, u_coarse)), + assemble(action(rf, u_fine)), + rtol=1e-12 + ) + + +def test_mg_jacobi(amh, atm): # pylint: disable=W0621 + """ + Test multigrid with jacobi smoothers + """ + V_J = FunctionSpace(amh[-1], "CG", 1) + x = SpatialCoordinate(amh[-1]) + u_ex = Function(V_J, name="u_fine_real").interpolate( + sin(2 * pi * x[0]) * sin(2 * pi * x[1]) + ) + u = Function(V_J) + v = TestFunction(V_J) + bc = DirichletBC(V_J, u_ex, "on_boundary") + F = inner(grad(u - u_ex), grad(v)) * dx + + params = { + "snes_type": "ksponly", + "ksp_max_it": 20, + "ksp_type": "cg", + "ksp_norm_type": "unpreconditioned", + "ksp_rtol": 1e-8, + "ksp_atol": 1e-8, + "pc_type": "mg", + "mg_levels_pc_type": "jacobi", + "mg_levels_ksp_type": "chebyshev", + "mg_levels_ksp_max_it": 2, + "mg_levels_ksp_richardson_scale": 1 / 3, + "mg_coarse_ksp_type": "preonly", + "mg_coarse_pc_type": "lu", + "mg_coarse_pc_factor_mat_solver_type": "mumps", + } + + problem = NonlinearVariationalProblem(F, u, bc) + solver = NonlinearVariationalSolver(problem, solver_parameters=params) + solver.set_transfer_manager(atm) + solver.solve() + assert errornorm(u_ex, u) <= 1e-8 + + +@pytest.mark.parametrize("params", ["jacobi", "asm", "patch"]) +def test_mg_patch(amh, atm, params): # pylint: disable=W0621 + """ + Test multigrid with patch relaxation + """ + if params == "jacobi": + solver_params = { + "mat_type": "matfree", + "ksp_type": "cg", + "pc_type": "mg", + "mg_levels": { + "ksp_type": "chebyshev", + "ksp_max_it": 1, + "pc_type": "jacobi", + }, + "mg_coarse": { + "ksp_type": "preonly", + "pc_type": "python", + "pc_python_type": "firedrake.AssembledPC", + "assembled": {"ksp_type": "preonly", "pc_type": "lu"}, + }, + } + elif params == "patch": + solver_params = { + "mat_type": "matfree", + "ksp_type": "cg", + "pc_type": "mg", + "mg_levels": { + "ksp_type": "chebyshev", + "ksp_max_it": 1, + "pc_type": "python", + "pc_python_type": "firedrake.PatchPC", + "patch": { + "pc_patch": { + "construct_type": "star", + "construct_dim": 0, + "sub_mat_type": "seqdense", + "dense_inverse": True, + "save_operators": True, + "precompute_element_tensors": True, + }, + "sub_ksp_type": "preonly", + "sub_pc_type": "lu", + }, + }, + "mg_coarse": { + "ksp_type": "preonly", + "pc_type": "python", + "pc_python_type": "firedrake.AssembledPC", + "assembled": {"ksp_type": "preonly", "pc_type": "lu"}, + }, + } + else: + solver_params = { + "mat_type": "aij", + "ksp_type": "cg", + "pc_type": "mg", + "mg_levels": { + "ksp_type": "chebyshev", + "ksp_max_it": 1, + "pc_type": "python", + "pc_python_type": "firedrake.ASMStarPC", + "pc_star_backend": "tinyasm", + }, + "mg_coarse": {"ksp_type": "preonly", "pc_type": "lu"}, + } + + V_J = FunctionSpace(amh[-1], "CG", 1) + x = SpatialCoordinate(amh[-1]) + u_ex = Function(V_J, name="u_fine_real").interpolate( + sin(2 * pi * x[0]) * sin(2 * pi * x[1]) + ) + u = Function(V_J) + v = TestFunction(V_J) + bc = DirichletBC(V_J, u_ex, "on_boundary") + F = inner(grad(u - u_ex), grad(v)) * dx + + problem = NonlinearVariationalProblem(F, u, bc) + solver = NonlinearVariationalSolver(problem, + solver_parameters=solver_params) + solver.set_transfer_manager(atm) + + solver.solve() + assert errornorm(u_ex, u) <= 1e-8