diff --git a/demos/adaptive_multigrid/adaptive_convergence.png b/demos/adaptive_multigrid/adaptive_convergence.png new file mode 100644 index 0000000000..295d8de024 Binary files /dev/null and b/demos/adaptive_multigrid/adaptive_convergence.png differ diff --git a/demos/adaptive_multigrid/adaptive_multigrid.py.rst b/demos/adaptive_multigrid/adaptive_multigrid.py.rst new file mode 100644 index 0000000000..44ff4eac0f --- /dev/null +++ b/demos/adaptive_multigrid/adaptive_multigrid.py.rst @@ -0,0 +1,255 @@ +Adaptive Multigrid Methods using AdaptiveMeshHierarchy +====================================================== + + +Contributed by Anurag Rao. + +The purpose of this demo is to show how to use Firedrake's multigrid solver on a hierarchy of adaptively refined Netgen meshes. +We will first have a look at how to use the :class:`~.AdaptiveMeshHierarchy` to construct the mesh hierarchy with Netgen meshes, then we will consider a solution to the Poisson problem on an L-shaped domain. +Finally, we will show how to use the :class:`~.AdaptiveMeshHierarchy` and :class:`~.AdaptiveTransferManager` to construct a scalable solver. The :class:`~.AdaptiveMeshHierarchy` contains information of the mesh hierarchy and the parent child relations between the meshes. +The :class:`~.AdaptiveTransferManager` deals with the transfer operator logic across any given levels in the hierarchy. +We begin by importing the necessary libraries :: + + from firedrake import * + from netgen.occ import * + import numpy + +Constructing the Mesh Hierarchy +------------------------------- +We first must construct the domain over which we will solve the problem. For a more comprehensive demo on how to use Open Cascade Technology (OCC) and Constructive Solid Geometry (CSG), +see `Netgen integration in Firedrake `_. +We begin with the L-shaped domain, which we build as the union of two rectangles: :: + + rect1 = WorkPlane(Axes((0,0,0), n=Z, h=X)).Rectangle(1,2).Face() + rect2 = WorkPlane(Axes((0,1,0), n=Z, h=X)).Rectangle(2,1).Face() + L = rect1 + rect2 + + geo = OCCGeometry(L, dim=2) + ngmsh = geo.GenerateMesh(maxh=0.5) + mesh = Mesh(ngmsh) + +It is important to convert the initial Netgen mesh into a Firedrake mesh before constructing the :class:`~.AdaptiveMeshHierarchy`. To call the constructor to the hierarchy, we must pass the initial mesh. Our initial mesh looks like this: + +.. figure:: initial_mesh.png + :align: center + :alt: Initial mesh. + +We will also initialize the :class:`~.AdaptiveTransferManager` here: :: + + amh = AdaptiveMeshHierarchy(mesh) + atm = AdaptiveTransferManager() + +Poisson Problem +--------------- +Now we can define a simple Poisson problem + +.. math:: + + - \nabla^2 u = f \text{ in } \Omega, \quad u = 0 \text{ on } \partial \Omega. + +Our approach strongly follows the similar problem in this `lecture course `_. We define the function ``solve_poisson``. The first lines correspond to finding a solution in the CG1 space. The right-hand side is set to be the constant function equal to 1. Since we want Dirichlet boundary conditions, we construct the :class:`~.DirichletBC` object and apply it to the entire boundary: :: + + def solve_poisson(mesh, params): + V = FunctionSpace(mesh, "CG", 1) + v = TestFunction(V) + u = TrialFunction(V) + uh = Function(V, name="solution") + bcs = [DirichletBC(V, 0, "on_boundary")] + f = Constant(1) + + a = inner(grad(u), grad(v))*dx + L = inner(f, v)*dx + + problem = LinearVariationalProblem(a, L, uh, bcs) + solver = LinearVariationalSolver(problem, solver_parameters=params) + + solver.set_transfer_manager(atm) + solver.solve() + + its = solver.snes.getLinearSolveIterations() + return uh, its + +Note the code after the construction of the :class:`~.LinearVariationalProblem`. To use the :class:`~.AdaptiveMeshHierarchy` with the existing Firedrake solver, we have to set the :class:`~.AdaptiveTransferManager` as the transfer manager of the multigrid solver. +Since we are using linear Lagrange elements, we will employ Jacobi as the multigrid relaxation, which we define with :: + + solver_params = { + "mat_type": "matfree", + "ksp_type": "cg", + "pc_type": "mg", + "mg_levels": { + "ksp_type": "chebyshev", + "ksp_max_it": 1, + "pc_type": "jacobi", + }, + "mg_coarse": { + "mat_type": "aij", + "pc_type": "lu", + }, + } + +Alternatively for high-order CG elements, it is recommended to use patch relaxation +to achieve degree-independent multigrid convergence. +For more information +see :doc:`Using patch relaxation for multigrid `. +The initial solution is shown below. + +.. figure:: solution_l1.png + :align: center + :alt: Initial Solution from multigrid with initial mesh. + + +Adaptive Mesh Refinement +------------------------ +In this section we will discuss how to adaptively refine select elements and add the newly refined mesh into the :class:`~.AdaptiveMeshHierarchy`. +For this problem, we will be using the Babuška-Rheinbolt a-posteriori estimate for an element: + +.. math:: + \eta_K^2 = h_K^2 \int_K \| f + \nabla^2 u_h \|^2 \mathrm{d}x + \frac{h_K}{2} \int_{\partial K \setminus \partial \Omega} \left[[ \nabla u_h \cdot \mathbf{n} \right]]^2 \mathrm{d}s, + +where :math:`K` is the element, :math:`h_K` is the diameter of the element, :math:`\mathbf{n}` is the outward-facing normal, and :math:`\left[[ \cdot \right]]` is the jump operator. The a-posteriori estimator is computed using the solution at the current level :math:`h`. Integrating over the domain and using the fact that the components of the estimator are piecewise constant on each cell, we can transform the above estimator into the variational problem + +.. math:: + \int_\Omega \eta_K^2 q \,\mathrm{d}x = \int_\Omega \sum_K h_K^2 \int_K (f + \text{div} (\text{grad} u_h) )^2 \,\mathrm{d}x q \,\mathrm{d}x + \int_\Omega \sum_K \frac{h_K}{2} \int_{\partial K \setminus \partial \Omega} \left[[ \nabla u_h \cdot \mathbf{n} \right]]^2 \,\mathrm{d}s q \,\mathrm{d}x \quad \forall\, q \in \mathrm{DG}_0 + +Our approach will be to compute the estimator over all elements and selectively choose to refine only those that contribute most to the error. To compute the error estimator, we use the function below to solve the variational formulation of the error estimator. Since our estimator is a constant per element, we use a DG0 function space. :: + + def estimate_error(mesh, uh): + Q = FunctionSpace(mesh, "DG", 0) + eta_sq = Function(Q) + p = TrialFunction(Q) + q = TestFunction(Q) + f = Constant(1) + residual = f + div(grad(uh)) + + # symbols for mesh quantities + h = CellDiameter(mesh) + n = FacetNormal(mesh) + vol = CellVolume(mesh) + + # compute cellwise error estimator + a = inner(p, q / vol) * dx + L = (inner(residual**2, q * h**2) * dx + + inner(jump(grad(uh), n)**2, avg(q * h)) * dS + ) + + sp = {"mat_type": "matfree", "ksp_type": "preonly", "pc_type": "jacobi"} + solve(a == L, eta_sq, solver_parameters=sp) + + # compute eta from eta^2 + eta = Function(Q).interpolate(sqrt(eta_sq)) + + # compute estimate for error in energy norm + with eta.dat.vec_ro as eta_: + error_est = eta_.norm() + return eta, error_est + +The next step is to choose which elements to refine. For this we use a simplified variant of Dörfler marking :cite:`Dorfler1996`: + +.. math:: + \eta_K \geq \theta \text{max}_L \eta_L + +The logic is to select an element :math:`K` to refine if the estimator is greater than some factor :math:`\theta` of the maximum error estimate of the mesh, where :math:`\theta` ranges from 0 to 1. In our code we choose :math:`\theta=0.5`. +With these helper functions complete, we can solve the system iteratively. In the max_iterations is the number of total levels we want to perform multigrid on. We will solve for 15 levels. At every level :math:`l`, we first compute the solution using multigrid up to level :math:`l`. We then use the current approximation of the solution to estimate the error across the mesh. Finally, we adaptively refine the mesh and repeat. :: + + theta = 0.5 + refinements = 15 + est_errors = [] + sqrt_dofs = [] + mg_iterations = [] + for level in range(refinements): + print(f"level {level}") + + mesh = amh[-1] + uh, its = solve_poisson(mesh, solver_params) + VTKFile(f"output/adaptive_loop_{level}.pvd").write(uh) + + (eta, error_est) = estimate_error(mesh, uh) + VTKFile(f"output/eta_{level}.pvd").write(eta) + + est_errors.append(error_est) + sqrt_dofs.append(uh.function_space().dim() ** 0.5) + mg_iterations.append(its) + + print(f" ||u - u_h|| <= C * {error_est}") + if len(est_errors) > 1: + rates = -numpy.diff(numpy.log(est_errors)) / numpy.diff(numpy.log(sqrt_dofs)) + print(f" rate = {rates[-1]}") + + if i != refinements - 1: + amh.adapt(eta, theta) + +To perform Dörfler marking, refine the current mesh, and add the mesh to the :class:`~.AdaptiveMeshHierarchy`, we use the ``amh.adapt(eta, theta)`` method. In this method the input is the recently computed error estimator ``eta`` and the Dörfler marking parameter ``theta``. The method always performs this on the current fine mesh in the hierarchy. There is another method for adding a mesh to the hierarchy: ``amh.add_mesh(mesh)``. In this method, refinement on the mesh is performed externally by some custom procedure and the resulting mesh directly gets added to the hierarchy. +The meshes now refine according to the error estimator. The error estimators at levels 3,5, and 15 are shown below. Zooming into the vertex of the L-shape at level 15 shows the error indicator remains strongest there. Further refinements will focus on that area. + ++-------------------------------+-------------------------------+-------------------------------+ +| .. figure:: eta_l3.png | .. figure:: eta_l6.png | .. figure:: eta_l15.png | +| :align: center | :align: center | :align: center | +| :height: 250px | :height: 250px | :height: 250px | +| :alt: Eta at level 3 | :alt: Eta at level 6 | :alt: Eta at level 15 | +| | | | +| *Level 3* | *Level 6* | *Level 15* | ++-------------------------------+-------------------------------+-------------------------------+ + +The solutions at level 4 and 15 are shown below. + ++------------------------------------+------------------------------------+ +| .. figure:: solution_l4.png | .. figure:: solution_l15.png | +| :align: center | :align: center | +| :height: 300px | :height: 300px | +| :alt: Solution, level 4 | :alt: Solution, level 15 | +| | | +| *MG solution at level 4* | *MG solution at level 15* | ++------------------------------------+------------------------------------+ + +The convergence follows the expected optimal behavior: :: + + from matplotlib import pyplot as plt + + dofs = numpy.array(sqrt_dofs) ** 2 + opt_errors = est_errors[0] * (sqrt_dofs[0] / numpy.array(sqrt_dofs)) + plt.loglog(dofs, est_errors, '-o', markersize = 3, label="Estimated error") + plt.loglog(dofs, opt_errors, '--', markersize = 3, label="Optimal convergence") + plt.ylabel("Error estimate of the energy norm") + plt.xlabel("Number of degrees of freedom") + plt.legend() + plt.savefig("output/adaptive_convergence.png") + +.. figure:: adaptive_convergence.png + :align: center + :alt: Convergence of the error estimator. + +Moreover, the multigrid iteration count is robust to the level of refinement :: + + print(" Level\t | Iterations") + print("---------------------") + for level, its in enumerate(mg_iterations): + print(f" {level}\t | {its}") + +.. + +======== ================ + Level Iterations +======== ================ + 0 2 + 1 8 + 2 8 + 3 8 + 4 8 + 5 8 + 6 8 + 7 8 + 8 8 + 9 9 + 10 9 + 11 9 + 12 9 + 13 9 + 14 9 +======== ================ + +A runnable python version of this demo can be found :demo:`here`. + +.. rubric:: References + +.. bibliography:: demo_references.bib + :filter: docname in docnames diff --git a/demos/adaptive_multigrid/eta_l15.png b/demos/adaptive_multigrid/eta_l15.png new file mode 100644 index 0000000000..05efad50c3 Binary files /dev/null and b/demos/adaptive_multigrid/eta_l15.png differ diff --git a/demos/adaptive_multigrid/eta_l3.png b/demos/adaptive_multigrid/eta_l3.png new file mode 100644 index 0000000000..ccc80c30f4 Binary files /dev/null and b/demos/adaptive_multigrid/eta_l3.png differ diff --git a/demos/adaptive_multigrid/eta_l6.png b/demos/adaptive_multigrid/eta_l6.png new file mode 100644 index 0000000000..f71d2c5788 Binary files /dev/null and b/demos/adaptive_multigrid/eta_l6.png differ diff --git a/demos/adaptive_multigrid/initial_mesh.png b/demos/adaptive_multigrid/initial_mesh.png new file mode 100644 index 0000000000..7ea4bb85bc Binary files /dev/null and b/demos/adaptive_multigrid/initial_mesh.png differ diff --git a/demos/adaptive_multigrid/solution_l1.png b/demos/adaptive_multigrid/solution_l1.png new file mode 100644 index 0000000000..ff6e6ed997 Binary files /dev/null and b/demos/adaptive_multigrid/solution_l1.png differ diff --git a/demos/adaptive_multigrid/solution_l15.png b/demos/adaptive_multigrid/solution_l15.png new file mode 100644 index 0000000000..e2436d74c0 Binary files /dev/null and b/demos/adaptive_multigrid/solution_l15.png differ diff --git a/demos/adaptive_multigrid/solution_l4.png b/demos/adaptive_multigrid/solution_l4.png new file mode 100644 index 0000000000..d457e29979 Binary files /dev/null and b/demos/adaptive_multigrid/solution_l4.png differ diff --git a/demos/demo_references.bib b/demos/demo_references.bib index 901abaa917..ce29d3c2e8 100644 --- a/demos/demo_references.bib +++ b/demos/demo_references.bib @@ -370,6 +370,17 @@ @article{Brubeck2024 year = {2024} } +@article{Dorfler1996, + title={A convergent adaptive algorithm for Poisson’s equation}, + author={D{\"o}rfler, Willy}, + journal={SIAM Journal on Numerical Analysis}, + volume={33}, + number={3}, + pages={1106--1124}, + year={1996}, + publisher={SIAM} +} + @Article{Farrell2015, author = {Patrick E. Farrell and \'Asgeir Birkisson and Simon W. Funke}, title = {{Deflation techniques for finding distinct solutions of nonlinear partial differential equations}}, diff --git a/docs/source/advanced_tut.rst b/docs/source/advanced_tut.rst index 29d8d093a4..5007d12f4b 100644 --- a/docs/source/advanced_tut.rst +++ b/docs/source/advanced_tut.rst @@ -26,6 +26,7 @@ element systems. Full-waveform inversion: spatial and wave sources parallelism. 1D Vlasov-Poisson equation using vertical independent function spaces. Degree-independent multigrid convergence using patch relaxation. + Multigrid on adaptively-refined mesh hierarchies. Monolithic multigrid with Vanka relaxation for Stokes. Vertex/edge star multigrid relaxation for H(div). Auxiliary space patch relaxation multigrid for H(curl). diff --git a/firedrake/__init__.py b/firedrake/__init__.py index de15331684..e0009c2b0e 100644 --- a/firedrake/__init__.py +++ b/firedrake/__init__.py @@ -101,7 +101,8 @@ def init_petsc(): HierarchyBase, MeshHierarchy, ExtrudedMeshHierarchy, NonNestedHierarchy, SemiCoarsenedExtrudedHierarchy, prolong, restrict, inject, TransferManager, - OpenCascadeMeshHierarchy + OpenCascadeMeshHierarchy, AdaptiveMeshHierarchy, + AdaptiveTransferManager ) from firedrake.norms import errornorm, norm # noqa: F401 from firedrake.nullspace import VectorSpaceBasis, MixedVectorSpaceBasis # noqa: F401 diff --git a/firedrake/mg/__init__.py b/firedrake/mg/__init__.py index 3aa6c7e021..12f8ec7c36 100644 --- a/firedrake/mg/__init__.py +++ b/firedrake/mg/__init__.py @@ -7,3 +7,5 @@ ) from firedrake.mg.embedded import TransferManager # noqa F401 from firedrake.mg.opencascade_mh import OpenCascadeMeshHierarchy # noqa F401 +from firedrake.mg.adaptive_hierarchy import AdaptiveMeshHierarchy # noqa F401 +from firedrake.mg.adaptive_transfer_manager import AdaptiveTransferManager # noqa: F401 diff --git a/firedrake/mg/adaptive_hierarchy.py b/firedrake/mg/adaptive_hierarchy.py new file mode 100644 index 0000000000..fc6c2f728a --- /dev/null +++ b/firedrake/mg/adaptive_hierarchy.py @@ -0,0 +1,85 @@ +from firedrake.mesh import MeshGeometry +from firedrake.cofunction import Cofunction +from firedrake.function import Function +from firedrake.mg import HierarchyBase +from firedrake.mg.utils import set_level + +__all__ = ["AdaptiveMeshHierarchy"] + + +class AdaptiveMeshHierarchy(HierarchyBase): + """ + HierarchyBase for hierarchies of adaptively refined meshes. + + Parameters + ---------- + base_mesh + The coarsest mesh in the hierarchy. + nested: bool + A flag to indicate whether the meshes are nested. + + """ + def __init__(self, base_mesh: MeshGeometry, nested: bool = True): + self.meshes = [] + self._meshes = [] + self.nested = nested + self.add_mesh(base_mesh) + + def add_mesh(self, mesh: MeshGeometry): + """ + Adds a mesh into the hierarchy. + + Parameters + ---------- + mesh + The mesh to be added to the finest level. + """ + level = len(self.meshes) + self._meshes.append(mesh) + self.meshes.append(mesh) + set_level(mesh, self, level) + + def adapt(self, eta: Function | Cofunction, theta: float): + """ + Adds a new mesh to the hierarchy by locally refining the finest mesh + with a simplified variant of Dorfler marking. The finest mesh must + come from a netgen mesh. + + Parameters + ---------- + eta + A DG0 :class:`~firedrake.function.Function` with the local error estimator. + theta + The threshold for marking as a fraction of the maximum error. + + Note + ---- + Dorfler marking involves sorting all of the elements by decreasing + error estimator and taking the minimal set that exceeds some fixed + fraction of the total error. What this code implements is the simpler + variant that doesn't have a proof of convergence (as far as I know) + but works as well in practice. + + """ + if not isinstance(eta, (Function, Cofunction)): + raise TypeError(f"eta must be a Function or Cofunction, not a {type(eta).__name__}") + M = eta.function_space() + if M.finat_element.space_dimension() != 1: + raise ValueError("eta must be a Function or Cofunction in DG0") + mesh = self.meshes[-1] + if M.mesh() is not mesh: + raise ValueError("eta must be defined on the finest mesh of the hierarchy") + + # Take the maximum over all processes + with eta.dat.vec_ro as evec: + _, eta_max = evec.max() + + threshold = theta * eta_max + should_refine = eta.dat.data_ro > threshold + + markers = Function(M) + markers.dat.data_wo[should_refine] = 1 + + refined_mesh = mesh.refine_marked_elements(markers) + self.add_mesh(refined_mesh) + return refined_mesh diff --git a/firedrake/mg/adaptive_transfer_manager.py b/firedrake/mg/adaptive_transfer_manager.py new file mode 100644 index 0000000000..149906b845 --- /dev/null +++ b/firedrake/mg/adaptive_transfer_manager.py @@ -0,0 +1,55 @@ +""" +This module contains the AdaptiveTransferManager used to perform +transfer operations on AdaptiveMeshHierarchies +""" +from firedrake.mg.embedded import TransferManager +from firedrake.ufl_expr import action, TrialFunction +from firedrake.interpolation import interpolate + + +__all__ = ("AdaptiveTransferManager",) + + +class AdaptiveTransferManager(TransferManager): + """ + TransferManager for adaptively refined mesh hierarchies + """ + def __init__(self, *, native_transfers=None, use_averaging=True): + if native_transfers is not None: + raise NotImplementedError("Custom transfers not implemented.") + super().__init__(native_transfers=native_transfers, use_averaging=use_averaging) + self.cache = {} + + def get_interpolator(self, Vc, Vf): + from firedrake.assemble import assemble + key = (Vc, Vf) + try: + return self.cache[key] + except KeyError: + Iexpr = interpolate(TrialFunction(Vc), Vf) + # TODO reusable matfree Interpolator + I = assemble(Iexpr, mat_type="aij") + return self.cache.setdefault(key, I) + + def forward(self, uc, uf): + from firedrake.assemble import assemble + Vc = uc.function_space() + Vf = uf.function_space() + I = self.get_interpolator(Vc, Vf) + return assemble(action(I, uc), tensor=uf) + + def adjoint(self, uf, uc): + from firedrake.assemble import assemble + Vc = uc.function_space().dual() + Vf = uf.function_space().dual() + I = self.get_interpolator(Vc, Vf) + return assemble(action(uf, I), tensor=uc) + + def prolong(self, uf, uc): + return self.forward(uf, uc) + + def inject(self, uc, uf): + return self.forward(uc, uf) + + def restrict(self, uc, uf): + return self.adjoint(uc, uf) diff --git a/tests/firedrake/demos/test_demos_run.py b/tests/firedrake/demos/test_demos_run.py index bd46952e55..7202ddbc5f 100644 --- a/tests/firedrake/demos/test_demos_run.py +++ b/tests/firedrake/demos/test_demos_run.py @@ -18,6 +18,7 @@ DEMO_DIR = join(CWD, "..", "..", "..", "demos") SERIAL_DEMOS = [ + Demo(("adaptive_multigrid", "adaptive_multigrid"), ["matplotlib", "netgen", "vtk"]), Demo(("benney_luke", "benney_luke"), ["vtk"]), Demo(("boussinesq", "boussinesq"), []), Demo(("burgers", "burgers"), ["vtk"]), diff --git a/tests/firedrake/multigrid/test_adaptive_multigrid.py b/tests/firedrake/multigrid/test_adaptive_multigrid.py new file mode 100644 index 0000000000..1235c95bd5 --- /dev/null +++ b/tests/firedrake/multigrid/test_adaptive_multigrid.py @@ -0,0 +1,369 @@ +""" +Tests for AdaptiveMeshHierarchy +and AdaptiveTransferManager +""" + +import pytest +import numpy as np +from firedrake import * + + +@pytest.fixture(params=[2, 3]) +def amh(request): + """ + Generate AdaptiveMeshHierarchies + """ + from netgen.occ import WorkPlane, OCCGeometry, Box, Pnt + dim = request.param + if dim == 2: + wp = WorkPlane() + wp.Rectangle(1, 1) + face = wp.Face() + geo = OCCGeometry(face, dim=2) + maxh = 0.5 + else: + cube = Box(Pnt(0, 0, 0), Pnt(1, 1, 1)) + geo = OCCGeometry(cube, dim=3) + maxh = 0.5 + + ngmesh = geo.GenerateMesh(maxh=maxh) + + dparams = {"overlap_type": (DistributedMeshOverlapType.VERTEX, 1)} + base = Mesh(ngmesh, distribution_parameters=dparams) + amh_test = AdaptiveMeshHierarchy(base) + + rg = RandomGenerator(PCG64(seed=0)) + for l in range(2): + mesh = amh_test[-1] + DG = FunctionSpace(mesh, "DG", 0) + should_refine = rg.uniform(DG, 0, 1).dat.global_data + + ngmesh = mesh.netgen_mesh + if dim == 2: + els = ngmesh.Elements2D() + else: + els = ngmesh.Elements3D() + for i, el in enumerate(els): + el.refine = 1 if should_refine[i] < 0.5 else 0 + + ngmesh.Refine(adaptive=True) + mesh = Mesh(ngmesh, distribution_parameters=dparams) + amh_test.add_mesh(mesh) + return amh_test + + +@pytest.fixture +def mh_uniform(): + """ + Generate MeshHierarchy for reference + """ + from netgen.occ import WorkPlane, OCCGeometry + wp = WorkPlane() + wp.Rectangle(2, 2) + face = wp.Face() + geo = OCCGeometry(face, dim=2) + maxh = 0.5 + ngmesh = geo.GenerateMesh(maxh=maxh) + + dparams = {"overlap_type": (DistributedMeshOverlapType.VERTEX, 1)} + base1 = Mesh(ngmesh, distribution_parameters=dparams) + mh = MeshHierarchy(base1, 2) + + base2 = Mesh(ngmesh, distribution_parameters=dparams) + amh = AdaptiveMeshHierarchy(base2) + for _ in range(2): + mesh = amh[-1] + ngmesh = mesh.netgen_mesh + ngmesh.Refine() + mesh = Mesh(ngmesh, distribution_parameters=dparams) + amh.add_mesh(mesh) + return amh, mh + + +@pytest.fixture +def atm(): + """atm used in tests""" + return AdaptiveTransferManager() + + +@pytest.fixture +def tm(): + """tm used for restrict consistency""" + return TransferManager() + + +@pytest.mark.parallel([1, 2]) +@pytest.mark.skipnetgen +@pytest.mark.parametrize("operator", ["prolong", "inject"]) +def test_DG0(amh, atm, operator): # pylint: disable=W0621 + """ + Prolongation & Injection test for DG0 + """ + V_coarse = FunctionSpace(amh[0], "DG", 0) + V_fine = FunctionSpace(amh[-1], "DG", 0) + u_coarse = Function(V_coarse) + u_fine = Function(V_fine) + xc, *_ = SpatialCoordinate(V_coarse.mesh()) + stepc = conditional(ge(xc, 0), 1, 0) + xf, *_ = SpatialCoordinate(V_fine.mesh()) + stepf = conditional(ge(xf, 0), 1, 0) + + if operator == "prolong": + u_coarse.interpolate(stepc) + assert errornorm(stepc, u_coarse) <= 1e-12 + + atm.prolong(u_coarse, u_fine) + assert errornorm(stepf, u_fine) <= 1e-12 + if operator == "inject": + u_fine.interpolate(stepf) + assert errornorm(stepf, u_fine) <= 1e-12 + + atm.inject(u_fine, u_coarse) + assert errornorm(stepc, u_coarse) <= 1e-12 + + +@pytest.mark.parallel([1, 2]) +@pytest.mark.skipnetgen +@pytest.mark.parametrize("operator", ["prolong", "inject"]) +def test_CG1(amh, atm, operator): # pylint: disable=W0621 + """ + Prolongation & Injection test for CG1 + """ + V_coarse = FunctionSpace(amh[0], "CG", 1) + V_fine = FunctionSpace(amh[-1], "CG", 1) + u_coarse = Function(V_coarse) + u_fine = Function(V_fine) + xc, *_ = SpatialCoordinate(V_coarse.mesh()) + xf, *_ = SpatialCoordinate(V_fine.mesh()) + + if operator == "prolong": + u_coarse.interpolate(xc) + assert errornorm(xc, u_coarse) <= 1e-12 + + atm.prolong(u_coarse, u_fine) + assert errornorm(xf, u_fine) <= 1e-12 + if operator == "inject": + u_fine.interpolate(xf) + assert errornorm(xf, u_fine) <= 1e-12 + + atm.inject(u_fine, u_coarse) + assert errornorm(xc, u_coarse) <= 1e-12 + + +@pytest.mark.parallel([1, 2]) +@pytest.mark.skipnetgen +def test_restrict_consistency(mh_uniform, atm, tm): # pylint: disable=W0621 + """ + Test restriction consistency of amh with uniform refinement vs mh + """ + amh_unif, mh = mh_uniform + + V_coarse = FunctionSpace(amh_unif[0], "DG", 0) + V_fine = FunctionSpace(amh_unif[-1], "DG", 0) + u_coarse = Function(V_coarse) + u_fine = Function(V_fine) + xc, _ = SpatialCoordinate(V_coarse.mesh()) + + u_coarse.interpolate(xc) + atm.prolong(u_coarse, u_fine) + + rf = assemble(conj(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(conj(TestFunction(Vfine)) * dx) + mhrc = Cofunction(Vcoarse.dual()) + + tm.restrict(mhrf, mhrc) + + assert abs( + (assemble(action(mhrc, mhuc)) - assemble(action(mhrf, mhuf))) + / assemble(action(mhrf, mhuf)) + ) <= 1e-12 + assert abs( + (assemble(action(rc, u_coarse)) - assemble(action(mhrc, mhuc))) + / assemble(action(mhrc, mhuc)) + ) <= 1e-12 + + +@pytest.mark.parallel([1, 2]) +@pytest.mark.skipnetgen +def test_restrict_CG1(amh, atm): # pylint: disable=W0621 + """ + Test restriction with CG1 + """ + V_coarse = FunctionSpace(amh[0], "CG", 1) + V_fine = FunctionSpace(amh[-1], "CG", 1) + u_coarse = Function(V_coarse) + u_fine = Function(V_fine) + xc, *_ = SpatialCoordinate(V_coarse.mesh()) + + u_coarse.interpolate(xc) + atm.prolong(u_coarse, u_fine) + + rf = assemble(conj(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 + ) + + +@pytest.mark.parallel([1, 2]) +@pytest.mark.skipnetgen +def test_restrict_DG0(amh, atm): # pylint: disable=W0621 + """ + Test restriction with DG0 + """ + V_coarse = FunctionSpace(amh[0], "DG", 0) + V_fine = FunctionSpace(amh[-1], "DG", 0) + u_coarse = Function(V_coarse) + u_fine = Function(V_fine) + xc, *_ = SpatialCoordinate(V_coarse.mesh()) + + u_coarse.interpolate(xc) + atm.prolong(u_coarse, u_fine) + + rf = assemble(conj(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 + ) + + +@pytest.mark.parallel([1, 2]) +@pytest.mark.skipnetgen +def test_mg_jacobi(amh, atm): # pylint: disable=W0621 + """ + Test multigrid with jacobi smoothers + """ + V = FunctionSpace(amh[-1], "CG", 1) + x = SpatialCoordinate(amh[-1]) + u_ex = Function(V).interpolate(sin(2 * pi * x[0]) * sin(2 * pi * x[1])) + u = Function(V) + v = TestFunction(V) + bc = DirichletBC(V, 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_coarse_ksp_type": "preonly", + "mg_coarse_pc_type": "lu", + "mg_coarse_pc_factor_mat_solver_type": "mumps", + } + + problem = NonlinearVariationalProblem(F, u, bc) + solver = NonlinearVariationalSolver(problem, solver_parameters=params) + solver.set_transfer_manager(atm) + solver.solve() + assert errornorm(u_ex, u) <= 1e-8 + + +@pytest.mark.parallel([1, 2]) +@pytest.mark.skipnetgen +@pytest.mark.parametrize("params", ["jacobi", "asm", "patch"]) +def test_mg_patch(amh, atm, params): # pylint: disable=W0621 + """ + Test multigrid with patch relaxation + """ + if params == "jacobi": + solver_params = { + "mat_type": "matfree", + "ksp_type": "cg", + "pc_type": "mg", + "mg_levels": { + "ksp_type": "chebyshev", + "ksp_max_it": 1, + "pc_type": "jacobi", + }, + "mg_coarse": { + "mat_type": "aij", + "ksp_type": "preonly", + "pc_type": "lu", + }, + } + elif params == "patch": + solver_params = { + "mat_type": "matfree", + "ksp_type": "cg", + "pc_type": "mg", + "mg_levels": { + "ksp_type": "chebyshev", + "ksp_max_it": 1, + "pc_type": "python", + "pc_python_type": "firedrake.PatchPC", + "patch": { + "pc_patch": { + "construct_type": "star", + "construct_dim": 0, + "sub_mat_type": "seqdense", + "dense_inverse": True, + "save_operators": True, + "precompute_element_tensors": True, + }, + "sub_ksp_type": "preonly", + "sub_pc_type": "lu", + }, + }, + "mg_coarse": { + "mat_type": "aij", + "ksp_type": "preonly", + "pc_type": "lu", + }, + } + else: + solver_params = { + "mat_type": "aij", + "ksp_type": "cg", + "pc_type": "mg", + "mg_levels": { + "ksp_type": "chebyshev", + "ksp_max_it": 1, + "pc_type": "python", + "pc_python_type": "firedrake.ASMStarPC", + "pc_star_backend": "tinyasm", + }, + "mg_coarse": {"ksp_type": "preonly", "pc_type": "lu"}, + } + + V = FunctionSpace(amh[-1], "CG", 1) + x = SpatialCoordinate(amh[-1]) + u_ex = Function(V).interpolate(sin(2 * pi * x[0]) * sin(2 * pi * x[1])) + u = Function(V) + v = TestFunction(V) + bc = DirichletBC(V, u_ex, "on_boundary") + F = inner(grad(u - u_ex), grad(v)) * dx + + problem = NonlinearVariationalProblem(F, u, bc) + solver = NonlinearVariationalSolver(problem, + solver_parameters=solver_params) + solver.set_transfer_manager(atm) + + solver.solve() + assert errornorm(u_ex, u) <= 1e-8