From 0cde87b0adb4a626eb956fe8cf8ace7586832e1d Mon Sep 17 00:00:00 2001 From: Anurag Rao Date: Thu, 28 Aug 2025 08:51:23 -0700 Subject: [PATCH 01/14] Adaptive Multigrid --- ngsPETSc/__init__.py | 4 +- .../utils/firedrake/adaptive_hierarchy.py | 307 +++++++++++++ .../firedrake/adaptive_transfer_manager.py | 119 +++++ tests/test_adaptive_multigrid.py | 407 ++++++++++++++++++ 4 files changed, 836 insertions(+), 1 deletion(-) create mode 100644 ngsPETSc/utils/firedrake/adaptive_hierarchy.py create mode 100644 ngsPETSc/utils/firedrake/adaptive_transfer_manager.py create mode 100644 tests/test_adaptive_multigrid.py diff --git a/ngsPETSc/__init__.py b/ngsPETSc/__init__.py index 76fb68b0..ed1b28b2 100644 --- a/ngsPETSc/__init__.py +++ b/ngsPETSc/__init__.py @@ -14,7 +14,9 @@ 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 00000000..eca71c41 --- /dev/null +++ b/ngsPETSc/utils/firedrake/adaptive_hierarchy.py @@ -0,0 +1,307 @@ +from netgen.occ import * +import numpy as np + +from firedrake.mg import HierarchyBase +from firedrake import * +from fractions import Fraction +from firedrake.mg.utils import set_level, get_level +from collections import defaultdict + +__all__ = ["AdaptiveMeshHierarchy"] + + +class AdaptiveMeshHierarchy(HierarchyBase): + 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 = None + self.fine_to_coarse_cells = 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): + 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) + + if len(self.meshes) <= 2: + self.coarse_to_fine_cells = {} + self.fine_to_coarse_cells = {} + self.fine_to_coarse_cells[Fraction(0, 1)] = None + + # extract parent child relationships from netgen meshes, + (c2f, f2c) = get_c2f_f2c_fd(mesh, coarse_mesh) + + self.coarse_to_fine_cells[Fraction(len(self.meshes) - 2, 1)] = c2f + self.fine_to_coarse_cells[Fraction(len(self.meshes) - 1, 1)] = 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_mesh = RelabeledMesh( + coarse_mesh, + [coarse_splits[i] for i in range(1, max_children + 1)], + list(range(1, max_children + 1)), + 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) + + mesh = RelabeledMesh( + mesh, + [fine_splits[i] for i in range(1, max_children + 1)], + list(range(1, max_children + 1)), + ) + 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 + n = [ + len([el for el in c2f if len(el) == j]) for j in range(1, max_children + 1) + ] + c2f_adjusted = { + j: np.zeros((n, j)) + for n, j in zip(n, list(range(1, max_children + 1))) + if n != 0 + } + f2c_adjusted = { + j: np.zeros((n * j, 1)) + for n, j in zip(n, list(range(1, max_children + 1))) + if n != 0 + } + + coarse_full_to_sub_map = { + i: full_to_sub(coarse_mesh, c_subm[i], i) for i in c_subm + } + fine_full_to_sub_map = { + j: full_to_sub(mesh, f_subm[j], int("10" + str(j))) for j in f_subm + } + + for i in range(len(c2f)): + n = len(c2f[i]) + if 1 <= n <= max_children: + c2f_adjusted[n][coarse_full_to_sub_map[n](i)] = fine_full_to_sub_map[n]( + np.array(c2f[i]) + ) + + for j in range(len(f2c)): + n = int(num_children[f2c[j]].item()) + if 1 <= n <= max_children: + f2c_adjusted[n][fine_full_to_sub_map[n](j), 0] = coarse_full_to_sub_map[ + n + ](f2c[j].item()) + + 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): + # Input a list corresponding to which elements get refined + ngmesh = self.meshes[-1].netgen_mesh + for l, el in enumerate(ngmesh.Elements2D()): + el.refine = 0 + if refinements[l] == 1: + el.refine = 1 + + ngmesh.Refine(adaptive=True) + mesh = Mesh(ngmesh) + self.add_mesh(mesh) + + def adapt(self, eta, theta): + # Implement Dorfler marking, returns new mesh + 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 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] + u_corr_space = Function( + V.reconstruct( + hierarchy_dict[[*hierarchy_dict][0]].meshes[ind].submesh_parent + ), + 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 of DoFs across submeshes, to fix restriction before recombinatoin + 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 to 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 fd + V = FunctionSpace(mesh, "DG", 0) + V2 = FunctionSpace(coarse_mesh, "DG", 0) + ngmesh = mesh.netgen_mesh + num_parents = coarse_mesh.num_cells() + + if mesh.topology_dm.getDimension() == 2: + parents = ngmesh.parentsurfaceelements.NumPy() + elements = ngmesh.Elements2D() + if mesh.topology_dm.getDimension() == 3: + parents = ngmesh.parentelements.NumPy() + elements = ngmesh.Elements3D() + fine_mapping = lambda x: mesh._cell_numbering.getOffset(x) + coarse_mapping = lambda x: coarse_mesh._cell_numbering.getOffset(x) + + c2f = [[] for _ in range(num_parents)] + f2c = [[] for _ in range(mesh.num_cells())] + + if parents.shape[0] == 0: + raise Exception("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[fine_mapping(l)].append(coarse_mapping(l)) + c2f[coarse_mapping(l)].append(fine_mapping(l)) + + elif parents[l][0] < num_parents: + f2c[fine_mapping(l)].append(coarse_mapping(parents[l][0])) + c2f[coarse_mapping(parents[l][0])].append(fine_mapping(l)) + + else: + a = parents[parents[l][0]][0] + while a >= num_parents: + a = parents[a][0] + + f2c[fine_mapping(l)].append(coarse_mapping(a)) + c2f[coarse_mapping(a)].append(fine_mapping(l)) + + return c2f, np.array(f2c).astype(int) + + +def split_to_submesh(mesh, coarse_mesh, c2f, f2c): + # Splits mesh into element numberings to generate submeshes + 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 in range(len(c2f)): + n = len(c2f[i]) + 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, label): + # returns the submesh element number associated with the full mesh numbering + 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.interpolate(u2, subset=mesh.cell_subset(label)) + + return lambda x: u1.dat.data[x].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 00000000..35451862 --- /dev/null +++ b/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py @@ -0,0 +1,119 @@ +from enum import IntEnum +from firedrake import * +from firedrake.mg.embedded import TransferManager +from firedrake.mg.utils import get_level + +import time + +from mpi4py import MPI + + +__all__ = ("AdaptiveTransferManager",) + + +native_families = frozenset( + ["Lagrange", "Discontinuous Lagrange", "Real", "Q", "DQ", "BrokenElement"] +) +alfeld_families = frozenset( + [ + "Hsieh-Clough-Tocher", + "Reduced-Hsieh-Clough-Tocher", + "Johnson-Mercier", + "Alfeld-Sorokina", + "Arnold-Qin", + "Reduced-Arnold-Qin", + "Christiansen-Hu", + "Guzman-Neilan", + "Guzman-Neilan Bubble", + ] +) +non_native_variants = frozenset(["integral", "fdm", "alfeld"]) + + +class Op(IntEnum): + PROLONG = 0 + RESTRICT = 1 + INJECT = 2 + + +class AdaptiveTransferManager(TransferManager): + 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 = {} + + def generic_transfer(self, source, target, transfer_op): + # determine which meshes to iterate over + 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(): + 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): + 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): + 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 prolong(self, uc, uf): + self.generic_transfer(uc, uf, transfer_op=self.tm.prolong) + + def inject(self, uf, uc): + self.generic_transfer(uf, uc, transfer_op=self.tm.inject) + + def restrict(self, source, target): + 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 00000000..5998c5ec --- /dev/null +++ b/tests/test_adaptive_multigrid.py @@ -0,0 +1,407 @@ +import pytest +from firedrake import * +from netgen.occ import * +from ngsPETSc import AdaptiveMeshHierarchy, AdaptiveTransferManager +import random +from firedrake.mg.ufl_utils import coarsen +from firedrake.dmhooks import get_appctx +from firedrake import dmhooks +from firedrake.solving_utils import _SNESContext + +@pytest.fixture +def amh(): + random.seed(1234) + 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) + amh = AdaptiveMeshHierarchy([base]) + for i in range(2): + for l, el in enumerate(ngmesh.Elements2D()): + el.refine = 0 + if random.random() < 0.5: + el.refine = 1 + ngmesh.Refine(adaptive=True) + mesh = Mesh(ngmesh) + amh.add_mesh(mesh) + return amh + +@pytest.fixture +def amh3D(): + random.seed(1234) + cube = Box(Pnt(0,0,0), Pnt(2,2,2)) + geo = OCCGeometry(cube, dim=3) + + maxh = 1 + ngmesh = geo.GenerateMesh(maxh=maxh) + base = Mesh(ngmesh) + amh3D = AdaptiveMeshHierarchy([base]) + for i in range(2): + for l, el in enumerate(ngmesh.Elements3D()): + el.refine = 0 + if random.random() < 0.5: + el.refine = 1 + ngmesh.Refine(adaptive=True) + mesh = Mesh(ngmesh) + amh3D.add_mesh(mesh) + return amh3D + +@pytest.fixture +def mh_res(): + 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 = AdaptiveMeshHierarchy([base]) + for i in range(2): + refs = np.ones(len(ngmesh.Elements2D())) + amh.refine(refs) + + mh = MeshHierarchy(mesh2, 2) + + return amh, mh + +@pytest.fixture +def atm(): + return AdaptiveTransferManager() + +@pytest.fixture +def tm(): + return TransferManager() + + +@pytest.mark.parametrize("operator", ["prolong", "inject"]) +def test_DG0_2D(amh, atm, operator): + + 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, yc = SpatialCoordinate(V_coarse.mesh()) + stepc = conditional(ge(xc, 0), 1, 0) + xf, yf = 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_DG0_3D(amh3D, atm, operator): + V_coarse = FunctionSpace(amh3D[0], "DG", 0) + V_fine = FunctionSpace(amh3D[-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_2D(amh, atm, operator): + + 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, yc = SpatialCoordinate(V_coarse.mesh()) + xf, yf = 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 + +@pytest.mark.parametrize("operator", ["prolong", "inject"]) +def test_CG1_3D(amh3D, atm, operator): + + V_coarse = FunctionSpace(amh3D[0], "CG", 1) + V_fine = FunctionSpace(amh3D[-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): + amh = mh_res[0] + mh = mh_res[1] + + 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) + + # 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_2D(amh, atm): + 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_CG1_3D(amh3D, atm): + V_coarse = FunctionSpace(amh3D[0], "CG", 1) + V_fine = FunctionSpace(amh3D[-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_2D(amh, atm): + 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_restrict_DG0_3D(amh3D, atm): + V_coarse = FunctionSpace(amh3D[0], "DG", 0) + V_fine = FunctionSpace(amh3D[-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): + V_J = FunctionSpace(amh[-1], "CG", 1) + (x,y) = SpatialCoordinate(amh[-1]) + u_ex = Function(V_J, name="u_fine_real").interpolate(sin(2 * pi * x) * sin(2 * pi * y)) + u = Function(V_J) + v = TestFunction(V_J) + bc = DirichletBC(V_J, Constant(0), "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) + dm = u.function_space().dm + old_appctx = get_appctx(dm) + mat_type = "aij" + appctx = _SNESContext(problem, mat_type, mat_type, old_appctx) + appctx.transfer_manager = atm + solver = NonlinearVariationalSolver(problem, solver_parameters=params) + solver.set_transfer_manager(atm) + with dmhooks.add_hooks(dm, solver, appctx=appctx, save=False): + coarsen(problem, coarsen) + + solver.solve() + assert errornorm(u_ex, u) <= 1e-8 + +def test_mg_patch(amh, atm): + def solve_sys(params): + V_J = FunctionSpace(amh[-1], "CG", 1) + (x,y) = SpatialCoordinate(amh[-1]) + u_ex = Function(V_J, name="u_fine_real").interpolate(sin(2 * pi * x) * sin(2 * pi * y)) + u = Function(V_J) + v = TestFunction(V_J) + bc = DirichletBC(V_J, Constant(0), "on_boundary") + F = inner(grad(u - u_ex), grad(v)) * dx + + problem = NonlinearVariationalProblem(F, u, bc) + + dm = u.function_space().dm + old_appctx = get_appctx(dm) + mat_type = "aij" + appctx = _SNESContext(problem, mat_type, mat_type, old_appctx) + appctx.transfer_manager = atm + + solver = NonlinearVariationalSolver(problem, solver_parameters=params) + solver.set_transfer_manager(atm) + with dmhooks.add_hooks(dm, solver, appctx=appctx, save=False): + coarsen(problem, coarsen) + + solver.solve() + assert errornorm(u_ex, u) <= 1e-8 + + lu = { + "ksp_type": "preonly", + "pc_type": "lu" + } + assembled_lu = { + "ksp_type": "preonly", + "pc_type": "python", + "pc_python_type": "firedrake.AssembledPC", + "assembled": lu + } + + def mg_params(relax, mat_type="aij"): + if mat_type == "aij": + coarse = lu + else: + coarse = assembled_lu + + return { + "mat_type": mat_type, + "ksp_type": "cg", + "pc_type": "mg", + "mg_levels": { + "ksp_type": "chebyshev", + "ksp_max_it": 1, + **relax + }, + "mg_coarse": coarse + } + jacobi_relax = mg_params({"pc_type": "jacobi"}, mat_type="matfree") + + patch_relax = mg_params({ + "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"}}, + mat_type="matfree") + + asm_relax = mg_params({ + "pc_type": "python", + "pc_python_type": "firedrake.ASMStarPC", + "pc_star_backend_type": "tinyasm"}) + + names = {"Jacobi": jacobi_relax, + "Patch": patch_relax, + "ASM Star": asm_relax} + + for name, params in names.items(): + print(name) + solve_sys(params) + From e524aea0ea67aaeb95b8033ca6fe56242c024ac0 Mon Sep 17 00:00:00 2001 From: Anurag Rao Date: Thu, 28 Aug 2025 09:00:23 -0700 Subject: [PATCH 02/14] add test to yml --- .github/workflows/ngsPETSc.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/ngsPETSc.yml b/.github/workflows/ngsPETSc.yml index e754bb23..e08d62f5 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: | From 7bbbe47eb939747c003ba671ba7cc75e42110557 Mon Sep 17 00:00:00 2001 From: Anurag Rao Date: Thu, 28 Aug 2025 11:17:28 -0700 Subject: [PATCH 03/14] style improvements/parametrized test --- .../utils/firedrake/adaptive_hierarchy.py | 53 +++-- .../firedrake/adaptive_transfer_manager.py | 58 +++-- tests/test_adaptive_multigrid.py | 211 +++++++----------- 3 files changed, 140 insertions(+), 182 deletions(-) diff --git a/ngsPETSc/utils/firedrake/adaptive_hierarchy.py b/ngsPETSc/utils/firedrake/adaptive_hierarchy.py index eca71c41..fd35d6d1 100644 --- a/ngsPETSc/utils/firedrake/adaptive_hierarchy.py +++ b/ngsPETSc/utils/firedrake/adaptive_hierarchy.py @@ -1,16 +1,20 @@ -from netgen.occ import * +""" +This module contains the class for the AdaptiveMeshHierarchy and related helper functions +""" import numpy as np - -from firedrake.mg import HierarchyBase -from firedrake import * from fractions import Fraction -from firedrake.mg.utils import set_level, get_level from collections import defaultdict +from firedrake.mg import HierarchyBase +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) @@ -23,6 +27,11 @@ def __init__(self, mesh, refinements_per_level=1, nested=True): 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: @@ -39,7 +48,7 @@ def add_mesh(self, mesh): self.fine_to_coarse_cells = {} self.fine_to_coarse_cells[Fraction(0, 1)] = None - # extract parent child relationships from netgen meshes, + # extract parent child relationships from netgen meshes (c2f, f2c) = get_c2f_f2c_fd(mesh, coarse_mesh) self.coarse_to_fine_cells[Fraction(len(self.meshes) - 2, 1)] = c2f @@ -128,7 +137,9 @@ def add_mesh(self, mesh): self.submesh_hierarchies.append(hierarchy_dict) def refine(self, refinements): - # Input a list corresponding to which elements get refined + """ + Refines and adds mesh if input a boolean vector corresponding to cells + """ ngmesh = self.meshes[-1].netgen_mesh for l, el in enumerate(ngmesh.Elements2D()): el.refine = 0 @@ -140,7 +151,9 @@ def refine(self, refinements): self.add_mesh(mesh) def adapt(self, eta, theta): - # Implement Dorfler marking, returns new mesh + """ + Implement Dorfler marking, refines mesh from error estimator + """ mesh = self.meshes[-1] W = FunctionSpace(mesh, "DG", 0) markers = Function(W) @@ -156,7 +169,9 @@ def adapt(self, eta, theta): return refined_mesh def split_function(self, u, child=True): - # Split function across submeshes + """ + Split input function across submeshes + """ V = u.function_space() full_mesh = V.mesh() _, level = get_level(full_mesh) @@ -190,7 +205,9 @@ def split_function(self, u, child=True): return split_functions def use_weight(self, V, child): - # counts of DoFs across submeshes, to fix restriction before recombinatoin + """ + Counts DoFs across submeshes, computes partition of unity + """ w = Function(V).assign(1) splits = self.split_function(w, child) @@ -200,7 +217,9 @@ def use_weight(self, V, child): return w def recombine(self, split_funcs, f, child=True): - # Recombines functions on submeshes to full mesh + """ + Recombines functions on submeshes back full mesh + """ V = f.function_space() f.zero() parent_mesh = ( @@ -227,7 +246,9 @@ def recombine(self, split_funcs, f, child=True): def get_c2f_f2c_fd(mesh, coarse_mesh): - # Construct coarse -> fine and fine -> coarse relations by mapping netgen elements to fd + """ + Construct coarse->fine and fine->coarse relations by mapping netgen elements to firedrake ones + """ V = FunctionSpace(mesh, "DG", 0) V2 = FunctionSpace(coarse_mesh, "DG", 0) ngmesh = mesh.netgen_mesh @@ -268,7 +289,9 @@ def get_c2f_f2c_fd(mesh, coarse_mesh): def split_to_submesh(mesh, coarse_mesh, c2f, f2c): - # Splits mesh into element numberings to generate submeshes + """ + Computes submesh split from full mesh + """ if mesh.topological_dimension() <= 2: max_children = 4 else: @@ -296,7 +319,9 @@ def split_to_submesh(mesh, coarse_mesh, c2f, f2c): def full_to_sub(mesh, submesh, label): - # returns the submesh element number associated with the full mesh numbering + """ + 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) diff --git a/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py b/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py index 35451862..f0eaa2d1 100644 --- a/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py +++ b/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py @@ -1,42 +1,17 @@ -from enum import IntEnum -from firedrake import * +""" +This module contains the AdaptiveTransferManager used to perform transfer operations on AdaptiveMeshHierarchies +""" +from firedrake import Function from firedrake.mg.embedded import TransferManager from firedrake.mg.utils import get_level -import time - -from mpi4py import MPI - __all__ = ("AdaptiveTransferManager",) - -native_families = frozenset( - ["Lagrange", "Discontinuous Lagrange", "Real", "Q", "DQ", "BrokenElement"] -) -alfeld_families = frozenset( - [ - "Hsieh-Clough-Tocher", - "Reduced-Hsieh-Clough-Tocher", - "Johnson-Mercier", - "Alfeld-Sorokina", - "Arnold-Qin", - "Reduced-Arnold-Qin", - "Christiansen-Hu", - "Guzman-Neilan", - "Guzman-Neilan Bubble", - ] -) -non_native_variants = frozenset(["integral", "fdm", "alfeld"]) - - -class Op(IntEnum): - PROLONG = 0 - RESTRICT = 1 - INJECT = 2 - - 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() @@ -44,7 +19,9 @@ def __init__(self, *, native_transfers=None, use_averaging=True): self.work_function_cache = {} def generic_transfer(self, source, target, transfer_op): - # determine which meshes to iterate over + """ + Generalized implementation of transfer operations wrapping the operations from TransferManager() + """ amh, source_level = get_level(source.function_space().mesh()) _, target_level = get_level(target.function_space().mesh()) @@ -95,12 +72,18 @@ def generic_transfer(self, source, target, transfer_op): 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: @@ -110,10 +93,19 @@ def get_weight(self, V_source): ) 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 index 5998c5ec..90db9ee2 100644 --- a/tests/test_adaptive_multigrid.py +++ b/tests/test_adaptive_multigrid.py @@ -1,27 +1,49 @@ import pytest -from firedrake import * -from netgen.occ import * -from ngsPETSc import AdaptiveMeshHierarchy, AdaptiveTransferManager import random +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 from firedrake.mg.ufl_utils import coarsen from firedrake.dmhooks import get_appctx from firedrake import dmhooks from firedrake.solving_utils import _SNESContext -@pytest.fixture -def amh(): +@pytest.fixture(params=[2,3]) +def amh(request): + """ + Generate AdaptiveMeshHierarchies + """ + dim = request.param random.seed(1234) - wp = WorkPlane() - wp.Rectangle(2,2) - face = wp.Face() + 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 = 1 - geo = OCCGeometry(face, dim=2) - maxh = 0.5 ngmesh = geo.GenerateMesh(maxh=maxh) base = Mesh(ngmesh) amh = AdaptiveMeshHierarchy([base]) + + if dim == 2: + els = ngmesh.Elements2D() + else: + els = ngmesh.Elements3D() + for i in range(2): - for l, el in enumerate(ngmesh.Elements2D()): + for l, el in enumerate(els): el.refine = 0 if random.random() < 0.5: el.refine = 1 @@ -30,28 +52,12 @@ def amh(): amh.add_mesh(mesh) return amh -@pytest.fixture -def amh3D(): - random.seed(1234) - cube = Box(Pnt(0,0,0), Pnt(2,2,2)) - geo = OCCGeometry(cube, dim=3) - - maxh = 1 - ngmesh = geo.GenerateMesh(maxh=maxh) - base = Mesh(ngmesh) - amh3D = AdaptiveMeshHierarchy([base]) - for i in range(2): - for l, el in enumerate(ngmesh.Elements3D()): - el.refine = 0 - if random.random() < 0.5: - el.refine = 1 - ngmesh.Refine(adaptive=True) - mesh = Mesh(ngmesh) - amh3D.add_mesh(mesh) - return amh3D @pytest.fixture def mh_res(): + """ + Generate MeshHierarchy for reference + """ wp = WorkPlane() wp.Rectangle(2,2) face = wp.Face() @@ -79,41 +85,17 @@ def tm(): @pytest.mark.parametrize("operator", ["prolong", "inject"]) -def test_DG0_2D(amh, atm, operator): - +def test_DG0(amh, atm, operator): + """ + 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, yc = SpatialCoordinate(V_coarse.mesh()) - stepc = conditional(ge(xc, 0), 1, 0) - xf, yf = 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_DG0_3D(amh3D, atm, operator): - V_coarse = FunctionSpace(amh3D[0], "DG", 0) - V_fine = FunctionSpace(amh3D[-1], "DG", 0) - u_coarse = Function(V_coarse) - u_fine = Function(V_fine) - xc, _,_ = SpatialCoordinate(V_coarse.mesh()) + xc, *_ = SpatialCoordinate(V_coarse.mesh()) stepc = conditional(ge(xc, 0), 1, 0) - xf, _, _ = SpatialCoordinate(V_fine.mesh()) + xf, *_ = SpatialCoordinate(V_fine.mesh()) stepf = conditional(ge(xf, 0), 1, 0) if operator == "prolong": @@ -132,39 +114,16 @@ def test_DG0_3D(amh3D, atm, operator): @pytest.mark.parametrize("operator", ["prolong", "inject"]) -def test_CG1_2D(amh, atm, operator): - +def test_CG1(amh, atm, operator): + """ + 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, yc = SpatialCoordinate(V_coarse.mesh()) - xf, yf = 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 - -@pytest.mark.parametrize("operator", ["prolong", "inject"]) -def test_CG1_3D(amh3D, atm, operator): - - V_coarse = FunctionSpace(amh3D[0], "CG", 1) - V_fine = FunctionSpace(amh3D[-1], "CG", 1) - u_coarse = Function(V_coarse) - u_fine = Function(V_fine) - xc, _, _ = SpatialCoordinate(V_coarse.mesh()) - xf, _, _ = SpatialCoordinate(V_fine.mesh()) + xc, *_ = SpatialCoordinate(V_coarse.mesh()) + xf, *_ = SpatialCoordinate(V_fine.mesh()) if operator == "prolong": @@ -182,6 +141,9 @@ def test_CG1_3D(amh3D, atm, operator): assert errornorm(xc, u_coarse) <= 1e-12 def test_restrict_consistency(mh_res, atm, tm): + """ + Test restriction consistency of amh with uniform refinement vs mh + """ amh = mh_res[0] mh = mh_res[1] @@ -216,28 +178,15 @@ def test_restrict_consistency(mh_res, atm, tm): 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_2D(amh, atm): +def test_restrict_CG1(amh, atm): + """ + 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_CG1_3D(amh3D, atm): - V_coarse = FunctionSpace(amh3D[0], "CG", 1) - V_fine = FunctionSpace(amh3D[-1], "CG", 1) - u_coarse = Function(V_coarse) - u_fine = Function(V_fine) - xc, _, _ = SpatialCoordinate(V_coarse.mesh()) + xc, *_ = SpatialCoordinate(V_coarse.mesh()) u_coarse.interpolate(xc) atm.prolong(u_coarse, u_fine) @@ -248,12 +197,15 @@ def test_restrict_CG1_3D(amh3D, atm): assert np.allclose(assemble(action(rc, u_coarse)), assemble(action(rf, u_fine)), rtol=1e-12) -def test_restrict_DG0_2D(amh, atm): +def test_restrict_DG0(amh, atm): + """ + 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()) + xc, *_ = SpatialCoordinate(V_coarse.mesh()) u_coarse.interpolate(xc) atm.prolong(u_coarse, u_fine) @@ -264,30 +216,16 @@ def test_restrict_DG0_2D(amh, atm): assert np.allclose(assemble(action(rc, u_coarse)), assemble(action(rf, u_fine)), rtol=1e-12) -def test_restrict_DG0_3D(amh3D, atm): - V_coarse = FunctionSpace(amh3D[0], "DG", 0) - V_fine = FunctionSpace(amh3D[-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): + """ + Test multigrid with jacobi smoothers + """ V_J = FunctionSpace(amh[-1], "CG", 1) - (x,y) = SpatialCoordinate(amh[-1]) - u_ex = Function(V_J, name="u_fine_real").interpolate(sin(2 * pi * x) * sin(2 * pi * y)) + 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, Constant(0), "on_boundary") + bc = DirichletBC(V_J, u_ex, "on_boundary") F = inner(grad(u - u_ex), grad(v)) * dx params = { @@ -322,13 +260,16 @@ def test_mg_jacobi(amh, atm): assert errornorm(u_ex, u) <= 1e-8 def test_mg_patch(amh, atm): + """ + Test multigrid with patch relaxation + """ def solve_sys(params): V_J = FunctionSpace(amh[-1], "CG", 1) - (x,y) = SpatialCoordinate(amh[-1]) - u_ex = Function(V_J, name="u_fine_real").interpolate(sin(2 * pi * x) * sin(2 * pi * y)) + 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, Constant(0), "on_boundary") + bc = DirichletBC(V_J, u_ex, "on_boundary") F = inner(grad(u - u_ex), grad(v)) * dx problem = NonlinearVariationalProblem(F, u, bc) @@ -346,6 +287,7 @@ def solve_sys(params): solver.solve() assert errornorm(u_ex, u) <= 1e-8 + lu = { "ksp_type": "preonly", @@ -395,13 +337,12 @@ def mg_params(relax, mat_type="aij"): asm_relax = mg_params({ "pc_type": "python", "pc_python_type": "firedrake.ASMStarPC", - "pc_star_backend_type": "tinyasm"}) + "pc_star_backend": "tinyasm"}) names = {"Jacobi": jacobi_relax, "Patch": patch_relax, "ASM Star": asm_relax} - for name, params in names.items(): - print(name) + for _, params in names.items(): solve_sys(params) From e16d7d84e2d1f4f06fd083ead6507807691bc325 Mon Sep 17 00:00:00 2001 From: Anurag Rao Date: Thu, 28 Aug 2025 11:34:23 -0700 Subject: [PATCH 04/14] improved style, split heavily composed codeline from Pablo --- ngsPETSc/__init__.py | 3 +- .../utils/firedrake/adaptive_hierarchy.py | 41 +++++++++---------- .../firedrake/adaptive_transfer_manager.py | 6 ++- 3 files changed, 25 insertions(+), 25 deletions(-) diff --git a/ngsPETSc/__init__.py b/ngsPETSc/__init__.py index ed1b28b2..135e4e9e 100644 --- a/ngsPETSc/__init__.py +++ b/ngsPETSc/__init__.py @@ -16,7 +16,8 @@ from ngsPETSc.utils.firedrake.hierarchies import * from ngsPETSc.utils.firedrake.adaptive_hierarchy import * from ngsPETSc.utils.firedrake.adaptive_transfer_manager import * - __all__ = __all__ + ["FiredrakeMesh", "NetgenHierarchy","AdaptiveMeshHierarchy","AdaptiveTransferManager"] + __all__ = __all__ + ["FiredrakeMesh", "NetgenHierarchy", + "AdaptiveMeshHierarchy","AdaptiveTransferManager"] #FEniCSx try: diff --git a/ngsPETSc/utils/firedrake/adaptive_hierarchy.py b/ngsPETSc/utils/firedrake/adaptive_hierarchy.py index fd35d6d1..551bf7b2 100644 --- a/ngsPETSc/utils/firedrake/adaptive_hierarchy.py +++ b/ngsPETSc/utils/firedrake/adaptive_hierarchy.py @@ -1,11 +1,12 @@ """ This module contains the class for the AdaptiveMeshHierarchy and related helper functions """ -import numpy as np from fractions import Fraction from collections import defaultdict -from firedrake.mg import HierarchyBase -from firedrake import Function, Cofunction, Mesh, Submesh, RelabeledMesh, HierarchyBase, FunctionSpace, conditional, gt +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"] @@ -87,7 +88,8 @@ def add_mesh(self, mesh): } set_level(mesh, self, level - 1) - # update c2f and f2c for submeshes by mapping numberings on full mesh to numberings on coarse mesh + # update c2f and f2c for submeshes by mapping numberings + # on full mesh to numberings on coarse mesh n = [ len([el for el in c2f if len(el) == j]) for j in range(1, max_children + 1) ] @@ -178,12 +180,9 @@ def split_function(self, u, child=True): ind = 1 if child else 0 hierarchy_dict = self.submesh_hierarchies[int(level) - ind] - u_corr_space = Function( - V.reconstruct( - hierarchy_dict[[*hierarchy_dict][0]].meshes[ind].submesh_parent - ), - val=u.dat, - ) + 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] @@ -249,41 +248,39 @@ def get_c2f_f2c_fd(mesh, coarse_mesh): """ Construct coarse->fine and fine->coarse relations by mapping netgen elements to firedrake ones """ - V = FunctionSpace(mesh, "DG", 0) - V2 = FunctionSpace(coarse_mesh, "DG", 0) + _ = FunctionSpace(mesh, "DG", 0) + _ = FunctionSpace(coarse_mesh, "DG", 0) ngmesh = mesh.netgen_mesh num_parents = coarse_mesh.num_cells() if mesh.topology_dm.getDimension() == 2: parents = ngmesh.parentsurfaceelements.NumPy() elements = ngmesh.Elements2D() - if mesh.topology_dm.getDimension() == 3: + else: parents = ngmesh.parentelements.NumPy() elements = ngmesh.Elements3D() - fine_mapping = lambda x: mesh._cell_numbering.getOffset(x) - coarse_mapping = lambda x: coarse_mesh._cell_numbering.getOffset(x) c2f = [[] for _ in range(num_parents)] f2c = [[] for _ in range(mesh.num_cells())] if parents.shape[0] == 0: - raise Exception("Added mesh has not refined any cells from previous mesh") + 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[fine_mapping(l)].append(coarse_mapping(l)) - c2f[coarse_mapping(l)].append(fine_mapping(l)) + 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: - f2c[fine_mapping(l)].append(coarse_mapping(parents[l][0])) - c2f[coarse_mapping(parents[l][0])].append(fine_mapping(l)) + f2c[mesh._cell_numbering.getOffset(l)].append(coarse_mesh._cell_numbering.getOffset(parents[l][0])) + c2f[coarse_mesh._cell_numbering.getOffset(parents[l][0])].append(mesh._cell_numbering.getOffset(l)) else: a = parents[parents[l][0]][0] while a >= num_parents: a = parents[a][0] - f2c[fine_mapping(l)].append(coarse_mapping(a)) - c2f[coarse_mapping(a)].append(fine_mapping(l)) + 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) diff --git a/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py b/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py index f0eaa2d1..662c4757 100644 --- a/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py +++ b/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py @@ -1,5 +1,6 @@ """ -This module contains the AdaptiveTransferManager used to perform transfer operations on AdaptiveMeshHierarchies +This module contains the AdaptiveTransferManager used to perform +transfer operations on AdaptiveMeshHierarchies """ from firedrake import Function from firedrake.mg.embedded import TransferManager @@ -20,7 +21,8 @@ def __init__(self, *, native_transfers=None, use_averaging=True): def generic_transfer(self, source, target, transfer_op): """ - Generalized implementation of transfer operations wrapping the operations from TransferManager() + 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()) From b6e23e210b6224d765c4bed93f2b5d0dd8f7a304 Mon Sep 17 00:00:00 2001 From: Anurag Rao Date: Thu, 28 Aug 2025 12:27:12 -0700 Subject: [PATCH 05/14] removed range len loops --- .../utils/firedrake/adaptive_hierarchy.py | 26 ++++++++++--------- 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/ngsPETSc/utils/firedrake/adaptive_hierarchy.py b/ngsPETSc/utils/firedrake/adaptive_hierarchy.py index 551bf7b2..bf055bbf 100644 --- a/ngsPETSc/utils/firedrake/adaptive_hierarchy.py +++ b/ngsPETSc/utils/firedrake/adaptive_hierarchy.py @@ -4,7 +4,7 @@ from fractions import Fraction from collections import defaultdict import numpy as np -from firedrake import (Function, Cofunction, Mesh, Submesh, RelabeledMesh, +from firedrake import (Function, Cofunction, Mesh, Submesh, RelabeledMesh, HierarchyBase, FunctionSpace, conditional, gt ) from firedrake.mg.utils import set_level, get_level @@ -88,7 +88,7 @@ def add_mesh(self, mesh): } set_level(mesh, self, level - 1) - # update c2f and f2c for submeshes by mapping numberings + # update c2f and f2c for submeshes by mapping numberings # on full mesh to numberings on coarse mesh n = [ len([el for el in c2f if len(el) == j]) for j in range(1, max_children + 1) @@ -111,19 +111,19 @@ def add_mesh(self, mesh): j: full_to_sub(mesh, f_subm[j], int("10" + str(j))) for j in f_subm } - for i in range(len(c2f)): - n = len(c2f[i]) + for i, children in enumerate(c2f): + n = len(children) if 1 <= n <= max_children: c2f_adjusted[n][coarse_full_to_sub_map[n](i)] = fine_full_to_sub_map[n]( - np.array(c2f[i]) + np.array(children) ) - for j in range(len(f2c)): - n = int(num_children[f2c[j]].item()) + for j, parent in enumerate(f2c): + n = int(num_children[parent].item()) if 1 <= n <= max_children: f2c_adjusted[n][fine_full_to_sub_map[n](j), 0] = coarse_full_to_sub_map[ n - ](f2c[j].item()) + ](parent.item()) c2f_subm = { i: {Fraction(0, 1): c2f_adjusted[i].astype(int)} for i in c2f_adjusted @@ -271,8 +271,10 @@ def get_c2f_f2c_fd(mesh, coarse_mesh): c2f[coarse_mesh._cell_numbering.getOffset(l)].append(mesh._cell_numbering.getOffset(l)) elif parents[l][0] < num_parents: - f2c[mesh._cell_numbering.getOffset(l)].append(coarse_mesh._cell_numbering.getOffset(parents[l][0])) - c2f[coarse_mesh._cell_numbering.getOffset(parents[l][0])].append(mesh._cell_numbering.getOffset(l)) + 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] @@ -303,8 +305,8 @@ def split_to_submesh(mesh, coarse_mesh, c2f, f2c): } num_children = np.zeros((len(c2f))) - for i in range(len(c2f)): - n = len(c2f[i]) + 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 From 7a30a456532bdcf02321389ab61b30f4bf5d0d47 Mon Sep 17 00:00:00 2001 From: Anurag Rao Date: Thu, 28 Aug 2025 12:48:13 -0700 Subject: [PATCH 06/14] updated test style --- tests/test_adaptive_multigrid.py | 241 ++++++++++++++++++------------- 1 file changed, 143 insertions(+), 98 deletions(-) diff --git a/tests/test_adaptive_multigrid.py b/tests/test_adaptive_multigrid.py index 90db9ee2..00c36405 100644 --- a/tests/test_adaptive_multigrid.py +++ b/tests/test_adaptive_multigrid.py @@ -1,20 +1,42 @@ -import pytest +""" +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 from firedrake.mg.ufl_utils import coarsen from firedrake.dmhooks import get_appctx from firedrake import dmhooks from firedrake.solving_utils import _SNESContext +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]) +@pytest.fixture(params=[2, 3]) def amh(request): """ Generate AdaptiveMeshHierarchies @@ -23,13 +45,12 @@ def amh(request): random.seed(1234) if dim == 2: wp = WorkPlane() - wp.Rectangle(2,2) + 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)) + cube = Box(Pnt(0, 0, 0), Pnt(2, 2, 2)) geo = OCCGeometry(cube, dim=3) maxh = 1 @@ -42,11 +63,11 @@ def amh(request): else: els = ngmesh.Elements3D() - for i in range(2): - for l, el in enumerate(els): + for _ in range(2): + for _, el in enumerate(els): el.refine = 0 if random.random() < 0.5: - el.refine = 1 + el.refine = 1 ngmesh.Refine(adaptive=True) mesh = Mesh(ngmesh) amh.add_mesh(mesh) @@ -59,28 +80,31 @@ def mh_res(): Generate MeshHierarchy for reference """ wp = WorkPlane() - wp.Rectangle(2,2) + 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 = AdaptiveMeshHierarchy([base]) + amh_unif = AdaptiveMeshHierarchy([base]) for i in range(2): refs = np.ones(len(ngmesh.Elements2D())) - amh.refine(refs) - + amh_unif.refine(refs) mh = MeshHierarchy(mesh2, 2) - return amh, mh + 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() @@ -104,7 +128,6 @@ def test_DG0(amh, atm, operator): 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 @@ -125,14 +148,12 @@ def test_CG1(amh, atm, operator): 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 @@ -140,6 +161,7 @@ def test_CG1(amh, atm, operator): atm.inject(u_fine, u_coarse) assert errornorm(xc, u_coarse) <= 1e-12 + def test_restrict_consistency(mh_res, atm, tm): """ Test restriction consistency of amh with uniform refinement vs mh @@ -156,27 +178,34 @@ def test_restrict_consistency(mh_res, atm, tm): u_coarse.interpolate(xc) atm.prolong(u_coarse, u_fine) - rf = assemble(TestFunction(V_fine)*dx) + 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 = 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 + 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): """ @@ -191,11 +220,16 @@ def test_restrict_CG1(amh, atm): u_coarse.interpolate(xc) atm.prolong(u_coarse, u_fine) - rf = assemble(TestFunction(V_fine)*dx) - rc = Cofunction(V_coarse.dual()) + 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) + + assert np.allclose( + assemble(action(rc, u_coarse)), + assemble(action(rf, u_fine)), + rtol=1e-12 + ) + def test_restrict_DG0(amh, atm): """ @@ -210,11 +244,16 @@ def test_restrict_DG0(amh, atm): u_coarse.interpolate(xc) atm.prolong(u_coarse, u_fine) - rf = assemble(TestFunction(V_fine)*dx) - rc = Cofunction(V_coarse.dual()) + 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) + + assert np.allclose( + assemble(action(rc, u_coarse)), + assemble(action(rf, u_fine)), + rtol=1e-12 + ) + def test_mg_jacobi(amh, atm): """ @@ -222,29 +261,31 @@ def test_mg_jacobi(amh, atm): """ 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_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" - } - + "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) dm = u.function_space().dm old_appctx = get_appctx(dm) @@ -255,18 +296,22 @@ def test_mg_jacobi(amh, atm): solver.set_transfer_manager(atm) with dmhooks.add_hooks(dm, solver, appctx=appctx, save=False): coarsen(problem, coarsen) - + solver.solve() assert errornorm(u_ex, u) <= 1e-8 + def test_mg_patch(amh, atm): """ - Test multigrid with patch relaxation + Test multigrid with patch relaxation """ + def solve_sys(params): 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_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") @@ -284,22 +329,18 @@ def solve_sys(params): solver.set_transfer_manager(atm) with dmhooks.add_hooks(dm, solver, appctx=appctx, save=False): coarsen(problem, coarsen) - + solver.solve() assert errornorm(u_ex, u) <= 1e-8 - - lu = { - "ksp_type": "preonly", - "pc_type": "lu" - } + lu = {"ksp_type": "preonly", "pc_type": "lu"} assembled_lu = { "ksp_type": "preonly", "pc_type": "python", "pc_python_type": "firedrake.AssembledPC", - "assembled": lu + "assembled": lu, } - + def mg_params(relax, mat_type="aij"): if mat_type == "aij": coarse = lu @@ -310,39 +351,43 @@ def mg_params(relax, mat_type="aij"): "mat_type": mat_type, "ksp_type": "cg", "pc_type": "mg", - "mg_levels": { - "ksp_type": "chebyshev", - "ksp_max_it": 1, - **relax - }, - "mg_coarse": coarse + "mg_levels": {"ksp_type": "chebyshev", "ksp_max_it": 1, **relax}, + "mg_coarse": coarse, } + jacobi_relax = mg_params({"pc_type": "jacobi"}, mat_type="matfree") - patch_relax = mg_params({ - "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"}}, - mat_type="matfree") - - asm_relax = mg_params({ - "pc_type": "python", - "pc_python_type": "firedrake.ASMStarPC", - "pc_star_backend": "tinyasm"}) - - names = {"Jacobi": jacobi_relax, - "Patch": patch_relax, - "ASM Star": asm_relax} - + patch_relax = mg_params( + { + "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", + }, + }, + mat_type="matfree", + ) + + asm_relax = mg_params( + { + "pc_type": "python", + "pc_python_type": "firedrake.ASMStarPC", + "pc_star_backend": "tinyasm", + } + ) + + names = {"Jacobi": jacobi_relax, + "Patch": patch_relax, + "ASM Star": asm_relax} + for _, params in names.items(): solve_sys(params) - From 8f0adb4499c0a2cad9e23a4375c8b5a439916311 Mon Sep 17 00:00:00 2001 From: Anurag Rao Date: Thu, 28 Aug 2025 12:57:16 -0700 Subject: [PATCH 07/14] change to fixtures --- tests/test_adaptive_multigrid.py | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/tests/test_adaptive_multigrid.py b/tests/test_adaptive_multigrid.py index 00c36405..88816426 100644 --- a/tests/test_adaptive_multigrid.py +++ b/tests/test_adaptive_multigrid.py @@ -56,7 +56,7 @@ def amh(request): ngmesh = geo.GenerateMesh(maxh=maxh) base = Mesh(ngmesh) - amh = AdaptiveMeshHierarchy([base]) + amh_test = AdaptiveMeshHierarchy([base]) if dim == 2: els = ngmesh.Elements2D() @@ -70,8 +70,8 @@ def amh(request): el.refine = 1 ngmesh.Refine(adaptive=True) mesh = Mesh(ngmesh) - amh.add_mesh(mesh) - return amh + amh_test.add_mesh(mesh) + return amh_test @pytest.fixture @@ -88,7 +88,7 @@ def mh_res(): base = Mesh(ngmesh) mesh2 = Mesh(ngmesh) amh_unif = AdaptiveMeshHierarchy([base]) - for i in range(2): + for _ in range(2): refs = np.ones(len(ngmesh.Elements2D())) amh_unif.refine(refs) mh = MeshHierarchy(mesh2, 2) @@ -109,7 +109,7 @@ def tm(): @pytest.mark.parametrize("operator", ["prolong", "inject"]) -def test_DG0(amh, atm, operator): +def test_DG0(amh, atm, operator): # pylint: disable=W0621 """ Prolongation & Injection test for DG0 """ @@ -137,7 +137,7 @@ def test_DG0(amh, atm, operator): @pytest.mark.parametrize("operator", ["prolong", "inject"]) -def test_CG1(amh, atm, operator): +def test_CG1(amh, atm, operator): # pylint: disable=W0621 """ Prolongation & Injection test for CG1 """ @@ -162,7 +162,7 @@ def test_CG1(amh, atm, operator): assert errornorm(xc, u_coarse) <= 1e-12 -def test_restrict_consistency(mh_res, atm, tm): +def test_restrict_consistency(mh_res, atm, tm): # pylint: disable=W0621 """ Test restriction consistency of amh with uniform refinement vs mh """ @@ -198,16 +198,16 @@ def test_restrict_consistency(mh_res, atm, tm): tm.restrict(mhrf, mhrc) assert ( - (assemble(action(mhrc, mhuc)) - assemble(action(mhrf, mhuf))) + (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(rc, u_coarse)) - assemble(action(mhrc, mhuc))) / assemble(action(mhrc, mhuc)) ) <= 1e-12 -def test_restrict_CG1(amh, atm): +def test_restrict_CG1(amh, atm): # pylint: disable=W0621 """ Test restriction with CG1 """ @@ -231,7 +231,7 @@ def test_restrict_CG1(amh, atm): ) -def test_restrict_DG0(amh, atm): +def test_restrict_DG0(amh, atm): # pylint: disable=W0621 """ Test restriction with DG0 """ @@ -255,7 +255,7 @@ def test_restrict_DG0(amh, atm): ) -def test_mg_jacobi(amh, atm): +def test_mg_jacobi(amh, atm): # pylint: disable=W0621 """ Test multigrid with jacobi smoothers """ @@ -301,7 +301,7 @@ def test_mg_jacobi(amh, atm): assert errornorm(u_ex, u) <= 1e-8 -def test_mg_patch(amh, atm): +def test_mg_patch(amh, atm): # pylint: disable=W0621 """ Test multigrid with patch relaxation """ @@ -385,7 +385,7 @@ def mg_params(relax, mat_type="aij"): } ) - names = {"Jacobi": jacobi_relax, + names = {"Jacobi": jacobi_relax, "Patch": patch_relax, "ASM Star": asm_relax} From 52b3ff4bf73abff784c3571f34f2ae18a0f2f975 Mon Sep 17 00:00:00 2001 From: Anurag Rao Date: Thu, 28 Aug 2025 13:02:57 -0700 Subject: [PATCH 08/14] fixed redefined amh --- tests/test_adaptive_multigrid.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_adaptive_multigrid.py b/tests/test_adaptive_multigrid.py index 88816426..b468159f 100644 --- a/tests/test_adaptive_multigrid.py +++ b/tests/test_adaptive_multigrid.py @@ -166,11 +166,11 @@ def test_restrict_consistency(mh_res, atm, tm): # pylint: disable=W0621 """ Test restriction consistency of amh with uniform refinement vs mh """ - amh = mh_res[0] + amh_unif = mh_res[0] mh = mh_res[1] - V_coarse = FunctionSpace(amh[0], "DG", 0) - V_fine = FunctionSpace(amh[-1], "DG", 0) + 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()) From 9963383fe79d2f61a55850371ec0fda988dc6711 Mon Sep 17 00:00:00 2001 From: Anurag Rao Date: Fri, 29 Aug 2025 04:41:52 -0700 Subject: [PATCH 09/14] updated Pablo's recs --- .../utils/firedrake/adaptive_hierarchy.py | 139 ++++++++------ tests/test_adaptive_multigrid.py | 176 +++++++++--------- 2 files changed, 175 insertions(+), 140 deletions(-) diff --git a/ngsPETSc/utils/firedrake/adaptive_hierarchy.py b/ngsPETSc/utils/firedrake/adaptive_hierarchy.py index bf055bbf..5bfa7721 100644 --- a/ngsPETSc/utils/firedrake/adaptive_hierarchy.py +++ b/ngsPETSc/utils/firedrake/adaptive_hierarchy.py @@ -1,11 +1,21 @@ """ -This module contains the class for the AdaptiveMeshHierarchy and related helper functions +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 import ( + Function, + Cofunction, + Mesh, + Submesh, + RelabeledMesh, + HierarchyBase, + FunctionSpace, + conditional, + gt, ) from firedrake.mg.utils import set_level, get_level @@ -16,12 +26,14 @@ 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 = None - self.fine_to_coarse_cells = None + 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) @@ -44,16 +56,12 @@ def add_mesh(self, mesh): set_level(self.meshes[-1], self, level - 1) self._shared_data_cache = defaultdict(dict) - if len(self.meshes) <= 2: - self.coarse_to_fine_cells = {} - self.fine_to_coarse_cells = {} - self.fine_to_coarse_cells[Fraction(0, 1)] = None - # extract parent child relationships from netgen meshes (c2f, f2c) = get_c2f_f2c_fd(mesh, coarse_mesh) - - self.coarse_to_fine_cells[Fraction(len(self.meshes) - 2, 1)] = c2f - self.fine_to_coarse_cells[Fraction(len(self.meshes) - 1, 1)] = np.array(f2c) + 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( @@ -63,10 +71,15 @@ def add_mesh(self, mesh): 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_splits[i] for i in range(1, max_children + 1)], - list(range(1, max_children + 1)), + coarse_indicators, + coarse_labels, name="Relabeled_coarse", ) c_subm = { @@ -76,10 +89,15 @@ def add_mesh(self, mesh): } 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_splits[i] for i in range(1, max_children + 1)], - list(range(1, max_children + 1)), + fine_indicators, + fine_labels, ) f_subm = { int(str(j)[-2:]): Submesh(mesh, mesh.topology_dm.getDimension(), j) @@ -90,43 +108,45 @@ def add_mesh(self, mesh): # update c2f and f2c for submeshes by mapping numberings # on full mesh to numberings on coarse mesh - n = [ - len([el for el in c2f if len(el) == j]) for j in range(1, max_children + 1) - ] + 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((n, j)) - for n, j in zip(n, list(range(1, max_children + 1))) - if n != 0 + 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((n * j, 1)) - for n, j in zip(n, list(range(1, max_children + 1))) - if n != 0 + 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], i) for i in c_subm + 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], int("10" + str(j))) for j in f_subm + 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: - c2f_adjusted[n][coarse_full_to_sub_map[n](i)] = fine_full_to_sub_map[n]( - np.array(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 = int(num_children[parent].item()) - if 1 <= n <= max_children: - f2c_adjusted[n][fine_full_to_sub_map[n](j), 0] = coarse_full_to_sub_map[ - n - ](parent.item()) + n = num_children[parent].item() + 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 + 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} @@ -143,10 +163,8 @@ def refine(self, refinements): Refines and adds mesh if input a boolean vector corresponding to cells """ ngmesh = self.meshes[-1].netgen_mesh - for l, el in enumerate(ngmesh.Elements2D()): - el.refine = 0 - if refinements[l] == 1: - el.refine = 1 + for i, el in enumerate(ngmesh.Elements2D()): + el.refine = refinements[i] ngmesh.Refine(adaptive=True) mesh = Mesh(ngmesh) @@ -198,7 +216,10 @@ def split_function(self, u, child=True): V_split.mesh().submesh_parent == u_corr_space.function_space().mesh() ) - f = split_functions.setdefault(i, Function(V_split, name=str(i))) + f = split_functions.setdefault( + i, + Function(V_split, name=str(i)) + ) f.assign(u_corr_space) return split_functions @@ -248,17 +269,17 @@ def get_c2f_f2c_fd(mesh, coarse_mesh): """ Construct coarse->fine and fine->coarse relations by mapping netgen elements to firedrake ones """ - _ = FunctionSpace(mesh, "DG", 0) - _ = FunctionSpace(coarse_mesh, "DG", 0) ngmesh = mesh.netgen_mesh num_parents = coarse_mesh.num_cells() if mesh.topology_dm.getDimension() == 2: parents = ngmesh.parentsurfaceelements.NumPy() elements = ngmesh.Elements2D() - else: + 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())] @@ -267,8 +288,12 @@ def get_c2f_f2c_fd(mesh, coarse_mesh): 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)) + 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) @@ -281,15 +306,21 @@ def get_c2f_f2c_fd(mesh, coarse_mesh): 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)) + 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 + 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 @@ -317,7 +348,7 @@ def split_to_submesh(mesh, coarse_mesh, c2f, f2c): return coarse_splits, fine_splits, num_children -def full_to_sub(mesh, submesh, label): +def full_to_sub(mesh, submesh): """ Returns the submesh element id associated with the full mesh element id """ @@ -326,6 +357,6 @@ def full_to_sub(mesh, submesh, label): u1 = Function(V1) u2 = Function(V2) u2.dat.data[:] = np.arange(len(u2.dat.data)) - u1.interpolate(u2, subset=mesh.cell_subset(label)) + u1.assign(u2, allow_missing_dofs=True) - return lambda x: u1.dat.data[x].astype(int) + return u1.dat.data.astype(int) diff --git a/tests/test_adaptive_multigrid.py b/tests/test_adaptive_multigrid.py index b468159f..79cffed6 100644 --- a/tests/test_adaptive_multigrid.py +++ b/tests/test_adaptive_multigrid.py @@ -2,6 +2,7 @@ Tests for AdaptiveMeshHierarchy and AdaptiveTransferManager """ + import random import pytest import numpy as np @@ -98,18 +99,18 @@ def mh_res(): @pytest.fixture def atm(): - """ atm used in tests""" + """atm used in tests""" return AdaptiveTransferManager() @pytest.fixture def tm(): - """ tm used for restrict consistency""" + """tm used for restrict consistency""" return TransferManager() @pytest.mark.parametrize("operator", ["prolong", "inject"]) -def test_DG0(amh, atm, operator): # pylint: disable=W0621 +def test_DG0(amh, atm, operator): # pylint: disable=W0621 """ Prolongation & Injection test for DG0 """ @@ -137,7 +138,7 @@ def test_DG0(amh, atm, operator): # pylint: disable=W0621 @pytest.mark.parametrize("operator", ["prolong", "inject"]) -def test_CG1(amh, atm, operator): # pylint: disable=W0621 +def test_CG1(amh, atm, operator): # pylint: disable=W0621 """ Prolongation & Injection test for CG1 """ @@ -162,7 +163,7 @@ def test_CG1(amh, atm, operator): # pylint: disable=W0621 assert errornorm(xc, u_coarse) <= 1e-12 -def test_restrict_consistency(mh_res, atm, tm): # pylint: disable=W0621 +def test_restrict_consistency(mh_res, atm, tm): # pylint: disable=W0621 """ Test restriction consistency of amh with uniform refinement vs mh """ @@ -207,7 +208,7 @@ def test_restrict_consistency(mh_res, atm, tm): # pylint: disable=W0621 ) <= 1e-12 -def test_restrict_CG1(amh, atm): # pylint: disable=W0621 +def test_restrict_CG1(amh, atm): # pylint: disable=W0621 """ Test restriction with CG1 """ @@ -231,7 +232,7 @@ def test_restrict_CG1(amh, atm): # pylint: disable=W0621 ) -def test_restrict_DG0(amh, atm): # pylint: disable=W0621 +def test_restrict_DG0(amh, atm): # pylint: disable=W0621 """ Test restriction with DG0 """ @@ -255,7 +256,7 @@ def test_restrict_DG0(amh, atm): # pylint: disable=W0621 ) -def test_mg_jacobi(amh, atm): # pylint: disable=W0621 +def test_mg_jacobi(amh, atm): # pylint: disable=W0621 """ Test multigrid with jacobi smoothers """ @@ -301,93 +302,96 @@ def test_mg_jacobi(amh, atm): # pylint: disable=W0621 assert errornorm(u_ex, u) <= 1e-8 -def test_mg_patch(amh, atm): # pylint: disable=W0621 +@pytest.mark.parametrize("params", ["jacobi", "asm", "patch"]) +def test_mg_patch(amh, atm, params): # pylint: disable=W0621 """ Test multigrid with patch relaxation """ - - def solve_sys(params): - 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) - - dm = u.function_space().dm - old_appctx = get_appctx(dm) - mat_type = "aij" - appctx = _SNESContext(problem, mat_type, mat_type, old_appctx) - appctx.transfer_manager = atm - - solver = NonlinearVariationalSolver(problem, solver_parameters=params) - solver.set_transfer_manager(atm) - with dmhooks.add_hooks(dm, solver, appctx=appctx, save=False): - coarsen(problem, coarsen) - - solver.solve() - assert errornorm(u_ex, u) <= 1e-8 - - lu = {"ksp_type": "preonly", "pc_type": "lu"} - assembled_lu = { - "ksp_type": "preonly", - "pc_type": "python", - "pc_python_type": "firedrake.AssembledPC", - "assembled": lu, - } - - def mg_params(relax, mat_type="aij"): - if mat_type == "aij": - coarse = lu - else: - coarse = assembled_lu - - return { - "mat_type": mat_type, + if params == "jacobi": + solver_params = { + "mat_type": "matfree", "ksp_type": "cg", "pc_type": "mg", - "mg_levels": {"ksp_type": "chebyshev", "ksp_max_it": 1, **relax}, - "mg_coarse": coarse, + "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"}, + }, } - - jacobi_relax = mg_params({"pc_type": "jacobi"}, mat_type="matfree") - - patch_relax = mg_params( - { - "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, + if 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", }, - "sub_ksp_type": "preonly", - "sub_pc_type": "lu", }, - }, - mat_type="matfree", - ) - - asm_relax = mg_params( - { - "pc_type": "python", - "pc_python_type": "firedrake.ASMStarPC", - "pc_star_backend": "tinyasm", + "mg_coarse": { + "ksp_type": "preonly", + "pc_type": "python", + "pc_python_type": "firedrake.AssembledPC", + "assembled": {"ksp_type": "preonly", "pc_type": "lu"}, + }, + } + if params == "asm": + 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 - names = {"Jacobi": jacobi_relax, - "Patch": patch_relax, - "ASM Star": asm_relax} + problem = NonlinearVariationalProblem(F, u, bc) - for _, params in names.items(): - solve_sys(params) + dm = u.function_space().dm + old_appctx = get_appctx(dm) + mat_type = "aij" + appctx = _SNESContext(problem, mat_type, mat_type, old_appctx) + appctx.transfer_manager = atm + + solver = NonlinearVariationalSolver(problem, + solver_parameters=solver_params) + solver.set_transfer_manager(atm) + with dmhooks.add_hooks(dm, solver, appctx=appctx, save=False): + coarsen(problem, coarsen) + + solver.solve() + assert errornorm(u_ex, u) <= 1e-8 From 66ca8a4e1aac5d81866fa48e29df4528d719056e Mon Sep 17 00:00:00 2001 From: Anurag Rao Date: Fri, 29 Aug 2025 04:56:33 -0700 Subject: [PATCH 10/14] fixed style issue of test --- tests/test_adaptive_multigrid.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_adaptive_multigrid.py b/tests/test_adaptive_multigrid.py index 79cffed6..f9c27c49 100644 --- a/tests/test_adaptive_multigrid.py +++ b/tests/test_adaptive_multigrid.py @@ -324,7 +324,7 @@ def test_mg_patch(amh, atm, params): # pylint: disable=W0621 "assembled": {"ksp_type": "preonly", "pc_type": "lu"}, }, } - if params == "patch": + elif params == "patch": solver_params = { "mat_type": "matfree", "ksp_type": "cg", @@ -354,7 +354,7 @@ def test_mg_patch(amh, atm, params): # pylint: disable=W0621 "assembled": {"ksp_type": "preonly", "pc_type": "lu"}, }, } - if params == "asm": + else: solver_params = { "mat_type": "aij", "ksp_type": "cg", From 8bf630691e40d600818db619426a501e80f758f6 Mon Sep 17 00:00:00 2001 From: Anurag Rao Date: Sun, 31 Aug 2025 18:12:55 -0700 Subject: [PATCH 11/14] fixed bug in 3D case --- ngsPETSc/utils/firedrake/adaptive_hierarchy.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/ngsPETSc/utils/firedrake/adaptive_hierarchy.py b/ngsPETSc/utils/firedrake/adaptive_hierarchy.py index 5bfa7721..1cfb57fa 100644 --- a/ngsPETSc/utils/firedrake/adaptive_hierarchy.py +++ b/ngsPETSc/utils/firedrake/adaptive_hierarchy.py @@ -134,15 +134,17 @@ def add_mesh(self, mesh): for i, children in enumerate(c2f): n = len(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 + 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() - 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 + 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)} From 06d9ad000e24f4fe2a9ee00914e78aabd2de686e Mon Sep 17 00:00:00 2001 From: Anurag Rao Date: Tue, 2 Sep 2025 07:41:45 -0700 Subject: [PATCH 12/14] fixed max_children for 3D --- .../utils/firedrake/adaptive_hierarchy.py | 4 ++-- .../firedrake/adaptive_transfer_manager.py | 1 + tests/test_adaptive_multigrid.py | 22 ------------------- 3 files changed, 3 insertions(+), 24 deletions(-) diff --git a/ngsPETSc/utils/firedrake/adaptive_hierarchy.py b/ngsPETSc/utils/firedrake/adaptive_hierarchy.py index 1cfb57fa..7d288101 100644 --- a/ngsPETSc/utils/firedrake/adaptive_hierarchy.py +++ b/ngsPETSc/utils/firedrake/adaptive_hierarchy.py @@ -48,7 +48,7 @@ def add_mesh(self, mesh): if mesh.topological_dimension() <= 2: max_children = 4 else: - max_children = 16 + max_children = 8 self._meshes += tuple(mesh) self.meshes += tuple(mesh) coarse_mesh = self.meshes[-2] @@ -327,7 +327,7 @@ def split_to_submesh(mesh, coarse_mesh, c2f, f2c): if mesh.topological_dimension() <= 2: max_children = 4 else: - max_children = 16 + max_children = 8 V = FunctionSpace(mesh, "DG", 0) V2 = FunctionSpace(coarse_mesh, "DG", 0) coarse_splits = { diff --git a/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py b/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py index 662c4757..18563071 100644 --- a/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py +++ b/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py @@ -65,6 +65,7 @@ def generic_transfer(self, source, target, transfer_op): target_function_splits = amh.split_function(curr_target, child=False) for split_label, _ in source_function_splits.items(): + transfer_op( source_function_splits[split_label], target_function_splits[split_label], diff --git a/tests/test_adaptive_multigrid.py b/tests/test_adaptive_multigrid.py index f9c27c49..fd59fd60 100644 --- a/tests/test_adaptive_multigrid.py +++ b/tests/test_adaptive_multigrid.py @@ -6,10 +6,6 @@ import random import pytest import numpy as np -from firedrake.mg.ufl_utils import coarsen -from firedrake.dmhooks import get_appctx -from firedrake import dmhooks -from firedrake.solving_utils import _SNESContext from firedrake import ( Mesh, MeshHierarchy, @@ -288,16 +284,8 @@ def test_mg_jacobi(amh, atm): # pylint: disable=W0621 } problem = NonlinearVariationalProblem(F, u, bc) - dm = u.function_space().dm - old_appctx = get_appctx(dm) - mat_type = "aij" - appctx = _SNESContext(problem, mat_type, mat_type, old_appctx) - appctx.transfer_manager = atm solver = NonlinearVariationalSolver(problem, solver_parameters=params) solver.set_transfer_manager(atm) - with dmhooks.add_hooks(dm, solver, appctx=appctx, save=False): - coarsen(problem, coarsen) - solver.solve() assert errornorm(u_ex, u) <= 1e-8 @@ -380,18 +368,8 @@ def test_mg_patch(amh, atm, params): # pylint: disable=W0621 F = inner(grad(u - u_ex), grad(v)) * dx problem = NonlinearVariationalProblem(F, u, bc) - - dm = u.function_space().dm - old_appctx = get_appctx(dm) - mat_type = "aij" - appctx = _SNESContext(problem, mat_type, mat_type, old_appctx) - appctx.transfer_manager = atm - solver = NonlinearVariationalSolver(problem, solver_parameters=solver_params) solver.set_transfer_manager(atm) - with dmhooks.add_hooks(dm, solver, appctx=appctx, save=False): - coarsen(problem, coarsen) - solver.solve() assert errornorm(u_ex, u) <= 1e-8 From 11fbd51ccc2b16db1f923e13ee43bb69abf2c695 Mon Sep 17 00:00:00 2001 From: Anurag Rao Date: Tue, 2 Sep 2025 10:56:38 -0700 Subject: [PATCH 13/14] fixed max_children back to 16 to fix 3D bug, added assignment so we don't perform transfer for unsplit submeshes --- .../utils/firedrake/adaptive_hierarchy.py | 5 +- .../firedrake/adaptive_transfer_manager.py | 47 ++++++++++++++++--- tests/test_adaptive_multigrid.py | 5 +- 3 files changed, 47 insertions(+), 10 deletions(-) diff --git a/ngsPETSc/utils/firedrake/adaptive_hierarchy.py b/ngsPETSc/utils/firedrake/adaptive_hierarchy.py index 7d288101..6ab8fbd6 100644 --- a/ngsPETSc/utils/firedrake/adaptive_hierarchy.py +++ b/ngsPETSc/utils/firedrake/adaptive_hierarchy.py @@ -48,7 +48,7 @@ def add_mesh(self, mesh): if mesh.topological_dimension() <= 2: max_children = 4 else: - max_children = 8 + max_children = 16 self._meshes += tuple(mesh) self.meshes += tuple(mesh) coarse_mesh = self.meshes[-2] @@ -146,6 +146,7 @@ def add_mesh(self, mesh): 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 @@ -327,7 +328,7 @@ def split_to_submesh(mesh, coarse_mesh, c2f, f2c): if mesh.topological_dimension() <= 2: max_children = 4 else: - max_children = 8 + max_children = 16 V = FunctionSpace(mesh, "DG", 0) V2 = FunctionSpace(coarse_mesh, "DG", 0) coarse_splits = { diff --git a/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py b/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py index 18563071..333ab435 100644 --- a/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py +++ b/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py @@ -1,7 +1,8 @@ """ -This module contains the AdaptiveTransferManager used to perform +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 @@ -9,6 +10,7 @@ __all__ = ("AdaptiveTransferManager",) + class AdaptiveTransferManager(TransferManager): """ TransferManager for adaptively refined mesh hierarchies @@ -18,6 +20,7 @@ def __init__(self, *, native_transfers=None, use_averaging=True): self.tm = TransferManager() self.weight_cache = {} self.work_function_cache = {} + self.perm_cache = {} def generic_transfer(self, source, target, transfer_op): """ @@ -65,11 +68,22 @@ def generic_transfer(self, source, target, transfer_op): target_function_splits = amh.split_function(curr_target, child=False) for split_label, _ in source_function_splits.items(): - - transfer_op( - source_function_splits[split_label], - target_function_splits[split_label], - ) + 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 @@ -94,6 +108,27 @@ def get_weight(self, V_source): 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): """ diff --git a/tests/test_adaptive_multigrid.py b/tests/test_adaptive_multigrid.py index fd59fd60..6fd40431 100644 --- a/tests/test_adaptive_multigrid.py +++ b/tests/test_adaptive_multigrid.py @@ -49,11 +49,11 @@ def amh(request): else: cube = Box(Pnt(0, 0, 0), Pnt(2, 2, 2)) geo = OCCGeometry(cube, dim=3) - maxh = 1 + maxh = 0.5 ngmesh = geo.GenerateMesh(maxh=maxh) base = Mesh(ngmesh) - amh_test = AdaptiveMeshHierarchy([base]) + amh_test = AdaptiveMeshHierarchy(base) if dim == 2: els = ngmesh.Elements2D() @@ -371,5 +371,6 @@ def test_mg_patch(amh, atm, params): # pylint: disable=W0621 solver = NonlinearVariationalSolver(problem, solver_parameters=solver_params) solver.set_transfer_manager(atm) + solver.solve() assert errornorm(u_ex, u) <= 1e-8 From 3c76bf426b955d2099b21152808bec7f69d1d95f Mon Sep 17 00:00:00 2001 From: Anurag Rao Date: Tue, 2 Sep 2025 11:01:01 -0700 Subject: [PATCH 14/14] clean trailing whitespace --- ngsPETSc/utils/firedrake/adaptive_transfer_manager.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py b/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py index 333ab435..3880c739 100644 --- a/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py +++ b/ngsPETSc/utils/firedrake/adaptive_transfer_manager.py @@ -108,7 +108,7 @@ def get_weight(self, V_source): 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 @@ -117,7 +117,6 @@ def get_perm(self, unsplit_source, unsplit_target, transfer_op): """ key = (unsplit_source.function_space(), unsplit_target.function_space()) - try: return self.perm_cache[key] except KeyError: