Source code for movement.monge_ampere

"""
Mesh movement based on solutions of equations of Monge-Ampère type.
"""

import abc
from collections.abc import Iterable
from warnings import warn

import firedrake
import firedrake.exceptions as fexc
import numpy as np
import ufl
from firedrake.petsc import PETSc
from pyadjoint import no_annotations

import movement.solver_parameters as solver_parameters
from movement.math import equation_of_hyperplane
from movement.mover import PrimeMover

__all__ = [
    "MongeAmpereMover_Relaxation",
    "MongeAmpereMover_QuasiNewton",
    "MongeAmpereMover",
]


[docs] def MongeAmpereMover(mesh, monitor_function, method="relaxation", **kwargs): r""" Factory function for generating Monge-Ampère mesh movers. Movement of a *mesh* is determined by a *monitor_function* :math:`m` and an equation of Monge-Ampère type, .. math:: m(\mathbf{x})\det(\mathbf{I} + \mathbf{H}(\phi)) = \theta, for a convex scalar potential :math:`\phi=\phi(\boldsymbol{\xi})`, which is a function of the coordinates of the *computational* mesh. Here :math:`m=m(\mathbf{x})` is a function of the coordinates of the *physical* mesh, :math:`\mathbf{I}` is the identity matrix, :math:`\theta` is a normalisation coefficient and :math:`\mathbf{H}(\phi)` denotes the Hessian of :math:`\phi` with respect to :math:`\boldsymbol{\xi}`. Different implementations solve the Monge-Ampère equation in different ways. If the `method` argument is set to `"relaxation"` then it is solved in parabolised form in :class:`~.MongeAmpereMover_Relaxation`. If the argument is set to `"quasi_newton"` then it is solved in its elliptic form using a quasi-Newton method in :class:`~.MongeAmpereMover_QuasiNewton`. Descriptions of both methods may be found in :cite:`MCB:18`. The physical mesh coordinates :math:`\mathbf{x}` are updated according to .. math:: \mathbf{x} = \boldsymbol{\xi} + \nabla_{\boldsymbol{\xi}}\phi. :arg mesh: the physical mesh :type mesh: :class:`firedrake.mesh.MeshGeometry` :arg monitor_function: a Python function which takes a mesh as input :type monitor_function: :class:`~.Callable` :kwarg method: choose from 'relaxation' and 'quasi_newton' :type method: :class:`str` :kwarg phi_init: initial guess for the scalar potential :type phi_init: :class:`firedrake.function.Function` :kwarg H_init: initial guess for the Hessian :type H_init: :class:`firedrake.function.Function` :kwarg maxiter: maximum number of iterations for the solver :type maxiter: :class:`int` :kwarg rtol: relative tolerance for the residual :type rtol: :class:`float` :kwarg dtol: divergence tolerance for the residual :type dtol: :class:`float` :kwarg pseudo_timestep: pseudo-timestep (only relevant to relaxation method) :type pseudo_timestep: :class:`float` :kwarg fixed_boundary_segments: labels corresponding to boundary segments to be fixed with a zero Dirichlet condition. The 'on_boundary' label indicates the whole domain boundary :type fixed_boundary_segments: :class:`list` of :class:`str` or :class:`int` :return: the Monge-Ampère Mover object :rtype: :class:`MongeAmpereMover_Relaxation` or :class:`MongeAmpereMover_QuasiNewton` """ implemented_methods = { "relaxation": MongeAmpereMover_Relaxation, "quasi_newton": MongeAmpereMover_QuasiNewton, } try: return implemented_methods[method](mesh, monitor_function, **kwargs) except KeyError as e: raise ValueError(f"Method '{method}' not recognised.") from e
def tangential(v, n): """ Return component of `v` perpendicular to `n` (assumed normalised). This is used to project vectors onto the tangent plane of a boundary. :arg v: the vector to project :type v: :class:`ufl.Expr` :arg n: the normal vector :type n: :class:`ufl.Expr` """ return v - ufl.dot(v, n) * n class MongeAmpereMover_Base(PrimeMover, metaclass=abc.ABCMeta): """ Base class for mesh movers based on the solution of Monge-Ampère type equations. Currently implemented subclasses: :class:`~.MongeAmpereMover_Relaxation` and :class:`~.MongeAmpereMover_QuasiNewton`. Descriptions of both methods may be found in :cite:`MCB:18`. """ def __init__(self, mesh, monitor_function, **kwargs): """ :arg mesh: the physical mesh :type mesh: :class:`firedrake.mesh.MeshGeometry` :arg monitor_function: a Python function which takes a mesh as input :type monitor_function: :class:`~.Callable` :kwarg phi_init: initial guess for the scalar potential :type phi_init: :class:`firedrake.function.Function` :kwarg H_init: initial guess for the Hessian :type H_init: :class:`firedrake.function.Function` :kwarg maxiter: maximum number of iterations for the relaxation :type maxiter: :class:`int` :kwarg rtol: relative tolerance for the residual :type rtol: :class:`float` :kwarg dtol: divergence tolerance for the residual :type dtol: :class:`float` :kwarg fixed_boundary_segments: labels corresponding to boundary segments to be fixed with a zero Dirichlet condition. The 'on_boundary' label indicates the whole domain boundary :type fixed_boundary_segments: :class:`list` of :class:`str` or :class:`int` """ if monitor_function is None: raise ValueError("Please supply a monitor function.") if mesh.coordinates.ufl_element().degree() != 1: raise NotImplementedError( f"{type(self).__name__} not implemented on curved meshes." ) # TODO: (#107) # Collect parameters before calling super self.maxiter = kwargs.pop("maxiter", 1000) self.rtol = kwargs.pop("rtol", 1.0e-08) self.dtol = kwargs.pop("dtol", 2.0) self.fixed_boundary_segments = kwargs.pop("fixed_boundary_segments", []) super().__init__(mesh, monitor_function=monitor_function, **kwargs) self.theta = firedrake.Constant(0.0) # Handle boundary segments where zero Dirichlet conditions are applied if self.fixed_boundary_segments == "on_boundary": self.fixed_boundary_segments = self._all_boundary_segments elif not isinstance(self.fixed_boundary_segments, Iterable): self.fixed_boundary_segments = [self.fixed_boundary_segments] if len(self._all_boundary_segments) == 0: warn( "Provided mesh has no boundary segments with Physical ID tags. If the " "boundaries aren't fully periodic then this will likely cause errors.", stacklevel=1, ) elif ( len(self.fixed_boundary_segments) == 1 and self.fixed_boundary_segments[0] == "on_boundary" ): self.fixed_boundary_segments = self._all_boundary_segments for boundary_tag in self.fixed_boundary_segments: if boundary_tag not in self._all_boundary_segments: raise ValueError(f"Provided boundary_tag '{boundary_tag}' is invalid.") def _create_function_spaces(self): super()._create_function_spaces() self.P1 = firedrake.FunctionSpace(self.mesh, "CG", 1) self.P1_vec = firedrake.VectorFunctionSpace(self.mesh, "CG", 1) self.P1_ten = firedrake.TensorFunctionSpace(self.mesh, "CG", 1) @abc.abstractmethod def _create_functions(self): super()._create_functions() self.monitor = firedrake.Function(self.P1, name="Monitor function") self._grad_phi = firedrake.Function(self.P1_vec) self.grad_phi = firedrake.Function(self.mesh.coordinates) self.monitor.interpolate(self.monitor_function(self.mesh)) self.original_volume = firedrake.Function(self.volume) self.total_volume = firedrake.assemble(firedrake.Constant(1.0) * self.dx) self.L_P0 = firedrake.TestFunction(self.P0) * self.monitor * self.dx @PETSc.Log.EventDecorator() def apply_initial_guess(self, phi_init, H_init): """ Initialise the approximations to the scalar potential and its Hessian with an initial guess. By default, both are initialised to zero, which corresponds to the case where the computational and physical meshes coincide. :arg phi_init: initial guess for the scalar potential :type phi_init: :class:`firedrake.function.Function` :arg H_init: initial guess for the Hessian :type H_init: :class:`firedrake.function.Function` """ if phi_init is not None and H_init is not None: if self.dim == 1: self.phi.interpolate(phi_init) self.H.interpolate(H_init) else: self.phi.project(phi_init) self.H.project(H_init) self.phi_old.assign(self.phi) self.H_old.assign(self.H) elif phi_init is not None or H_init is not None: raise ValueError("Need to initialise both phi *and* H.") @property def relative_l2_residual(self): """ :return: the relative :math:`L^2` norm residual. :rtype: :class:`float` """ assert hasattr(self, "_residual_l2_form") assert hasattr(self, "_norm_l2_form") return ( firedrake.assemble(self._residual_l2_form).dat.norm / firedrake.assemble(self._norm_l2_form).dat.norm ) @PETSc.Log.EventDecorator() def _update_physical_coordinates(self): r""" Update the physical coordinates :math:`\mathbf{x}` using the recovered gradient: .. math:: \mathbf{x} = \boldsymbol{\xi} + \nabla_{\boldsymbol{\xi}}\phi. After updating the coordinates, this method also checks for mesh tangling if this is turned on. (It will be turned on by default in the 2D case.) """ try: self.grad_phi.assign(self._grad_phi) except Exception: self.grad_phi.interpolate(self._grad_phi) self.x.assign(self.xi + self.grad_phi) # Check if the mesh has become tangled if hasattr(self, "tangling_checker"): self.tangling_checker.check() def _l2_projector_bcs(self, boundary_tag): r""" Determine boundary conditions to apply for the :math:`L^2` projection step for a given boundary segment. Note that the boundary segment *must* be a straight line or plane. :arg boundary_tag: tag for the boundary segment in consideration :type boundary_tag: :class:`str` :returns: tuple of boundary conditions :rtype: :class:`tuple` of :class:`~.DirichletBC`\s """ zero_bc = firedrake.DirichletBC(self.P1_vec, 0, boundary_tag) if (boundary_tag in self.fixed_boundary_segments) or self.dim == 1: return (zero_bc,) # If the boundary segment is axis-aligned, it is straightforward to avoid # movement in the normal direction while allowing tangential movement - simply # fix the component for the appropriate coordinate n = ufl.FacetNormal(self.mesh) ds = self.ds(boundary_tag) iszero = [np.isclose(firedrake.assemble(abs(ni) * ds), 0.0) for ni in n] nzero = sum(iszero) assert nzero < self.dim if nzero == self.dim - 1: idx = iszero.index(False) return (firedrake.DirichletBC(self.P1_vec.sub(idx), 0, boundary_tag),) # Create an equation boundary condition for enforcing no mesh movement normal to # domain boundaries u_cts = firedrake.TrialFunction(self.P1_vec) v_cts = firedrake.TestFunction(self.P1_vec) a_bc = ufl.dot(v_cts, n) * ufl.dot(u_cts, n) * ds L_bc = ufl.dot(v_cts, n) * firedrake.Constant(0.0) * ds bc1 = firedrake.EquationBC(a_bc == L_bc, self._grad_phi, boundary_tag) # Determine the 'corner' vertices which are at the intersection of two boundary # segments and create a Dirichlet condition for fixing them under mesh movement facet_indices = set(self._all_boundary_segments) ffacet_indices = [ (tag, boundary_tag) for tag in facet_indices.difference([boundary_tag]) ] ffacet_bc = firedrake.DirichletBC(self.P1_vec, 0, ffacet_indices) # Check the current boundary segment is a straight line or plane # TODO: Only do the following check in debugging mode (#94) assert self.coord_space.ufl_element().degree() == 1 ffacets = list(self.mesh.coordinates.dat.data_with_halos[ffacet_bc.nodes]) hyperplane = equation_of_hyperplane(*ffacets) for x in self.mesh.coordinates.dat.data_with_halos[zero_bc.nodes]: if not np.isclose(float(hyperplane(*x)), 0): raise ValueError( f"Boundary segment '{boundary_tag}' is not" f" {'linear' if self.dim == 2 else 'planar'}." ) # Create an equation boundary condition which allows tangential movement, but # only up until the 'corner' vertices where boundary segments meet a_bc = ufl.dot(tangential(v_cts, n), tangential(u_cts, n)) * ds L_bc = ufl.dot(tangential(v_cts, n), tangential(ufl.grad(self.phi_old), n)) * ds bc2 = firedrake.EquationBC( a_bc == L_bc, self._grad_phi, boundary_tag, bcs=ffacet_bc ) return bc1, bc2 @property @PETSc.Log.EventDecorator() def l2_projector(self): """ Create a linear solver for obtaining the gradient of the potential using an :math:`L^2` projection. :return: the linear solver :rtype: :class:`~.LinearVariationalSolver` """ if hasattr(self, "_l2_projector"): return self._l2_projector u_cts = firedrake.TrialFunction(self.P1_vec) v_cts = firedrake.TestFunction(self.P1_vec) # Domain interior a = ufl.inner(v_cts, u_cts) * self.dx L = ufl.inner(v_cts, ufl.grad(self.phi_old)) * self.dx # Enforce no movement normal to boundary bcs = [ dirichlet_bc for boundary_tag in self._all_boundary_segments for dirichlet_bc in self._l2_projector_bcs(boundary_tag) ] # Create solver problem = firedrake.LinearVariationalProblem(a, L, self._grad_phi, bcs=bcs) self._l2_projector = firedrake.LinearVariationalSolver( problem, solver_parameters=solver_parameters.cg_ilu ) return self._l2_projector
[docs] class MongeAmpereMover_Relaxation(MongeAmpereMover_Base): r""" The standard, elliptic form of the Monge-Ampère equation used for mesh movement is: .. math:: m(\mathbf{x})\det(\mathbf{I} + \mathbf{H}(\phi)) = \theta, for a convex scalar potential :math:`\phi=\phi(\boldsymbol{\xi})`, which is a function of the coordinates of the *computational* mesh. Here :math:`m=m(\mathbf{x})` is a user-provided monitor function, which is a function of the coordinates of the *physical* mesh. :math:`\mathbf{I}` is the identity matrix, :math:`\theta` is a normalisation coefficient and :math:`\mathbf{H}(\phi)` denotes the Hessian of :math:`\phi` with respect to :math:`\boldsymbol{\xi}`. In this mesh mover, the Monge-Ampère equation is instead solved in a parabolised form using a pseudo-time relaxation, .. math:: -\frac\partial{\partial\tau}\Delta\phi = m(\mathbf{x})\det(\mathbf{I} + \mathbf{H}(\phi)) - \theta, where :math:`\tau` is the pseudo-time variable. Forward Euler is used for the pseudo-time integration (see :cite:`MCB:18` for details). This approach typically takes tens or hundreds of iterations to converge, but each iteration is relatively cheap. """ @PETSc.Log.EventDecorator() def __init__(self, mesh, monitor_function, phi_init=None, H_init=None, **kwargs): """ :arg mesh: the physical mesh :type mesh: :class:`firedrake.mesh.MeshGeometry` :arg monitor_function: a Python function which takes a mesh as input :type monitor_function: :class:`~.Callable` :kwarg phi_init: initial guess for the scalar potential :type phi_init: :class:`firedrake.function.Function` :kwarg H_init: initial guess for the Hessian :type H_init: :class:`firedrake.function.Function` :kwarg pseudo_timestep: pseudo-timestep to use for the relaxation :type pseudo_timestep: :class:`float` :kwarg maxiter: maximum number of iterations for the relaxation :type maxiter: :class:`int` :kwarg rtol: relative tolerance for the residual :type rtol: :class:`float` :kwarg dtol: divergence tolerance for the residual :type dtol: :class:`float` :kwarg fixed_boundary_segments: labels corresponding to boundary segments to be fixed with a zero Dirichlet condition. The 'on_boundary' label indicates the whole domain boundary :type fixed_boundary_segments: :class:`list` of :class:`str` or :class:`int` """ self.pseudo_dt = firedrake.Constant(kwargs.pop("pseudo_timestep", 0.1)) super().__init__(mesh, monitor_function=monitor_function, **kwargs) # Initialise phi and H if phi_init or H_init: self.apply_initial_guess(phi_init, H_init) # Setup residuals I = ufl.Identity(self.dim) self.theta_form = self.monitor * ufl.det(I + self.H_old) * self.dx self.residual = self.monitor * ufl.det(I + self.H_old) - self.theta psi = firedrake.TestFunction(self.P1) self._residual_l2_form = psi * self.residual * self.dx self._norm_l2_form = psi * self.theta * self.dx def _create_functions(self): super()._create_functions() self.phi = firedrake.Function(self.P1) self.H = firedrake.Function(self.P1_ten) self.phi_old = firedrake.Function(self.P1) self.H_old = firedrake.Function(self.P1_ten) @property @PETSc.Log.EventDecorator() def pseudotimestepper(self): """ Setup the pseudo-timestepper for the relaxation method. Forward Euler is used for the pseudo-time integration (see :cite:`MCB:18` for details). The pseudo-timestep may be set through the `pseudo_timestep` keyword argument to the constructor. :return: the pseudo-timestepper :rtype: :class:`~.LinearVariationalSolver` """ if hasattr(self, "_pseudotimestepper"): return self._pseudotimestepper phi = firedrake.TrialFunction(self.P1) psi = firedrake.TestFunction(self.P1) a = ufl.inner(ufl.grad(psi), ufl.grad(phi)) * self.dx L = ( ufl.inner(ufl.grad(psi), ufl.grad(self.phi_old)) * self.dx + self.pseudo_dt * psi * self.residual * self.dx ) problem = firedrake.LinearVariationalProblem(a, L, self.phi) nullspace = firedrake.VectorSpaceBasis(constant=True) self._pseudotimestepper = firedrake.LinearVariationalSolver( problem, solver_parameters=solver_parameters.cg_gamg, nullspace=nullspace, transpose_nullspace=nullspace, ) return self._pseudotimestepper @property @PETSc.Log.EventDecorator() def equidistributor(self): r""" Setup the equidistributor for the relaxation method. The equidistributor solves the following equation: .. math:: \int_{\Omega} \tau : \mathbf{H} \, \mathrm{d}x = -\int_{\Omega} (\nabla \cdot \tau) \cdot (\nabla \phi) \, \mathrm{d}x + \int_{\partial \Omega} ((\nabla \phi \cdot \widehat{\mathbf{n}}) \cdot \tau) \cdot \widehat{\mathbf{n}} \, \mathrm{d}s, \quad \forall \tau \in \mathbb{P}1^{d \times d} for the Hessian :math:`\mathbf{H}`, where :math:`d` is the spatial dimension and :math:`\widehat{\mathbf{n}}` is a normal vector to the boundary. :return: the equidistributor :rtype: :class:`~.LinearVariationalSolver` """ if hasattr(self, "_equidistributor"): return self._equidistributor n = ufl.FacetNormal(self.mesh) tau = firedrake.TestFunction(self.P1_ten) a = ufl.inner(tau, firedrake.TrialFunction(self.P1_ten)) * self.dx L = ( -ufl.dot(ufl.div(tau), ufl.grad(self.phi)) * self.dx + ufl.dot(ufl.dot(tangential(ufl.grad(self.phi), n), tau), n) * self.ds ) problem = firedrake.LinearVariationalProblem(a, L, self.H) self._equidistributor = firedrake.LinearVariationalSolver( problem, solver_parameters=solver_parameters.cg_ilu ) return self._equidistributor
[docs] @PETSc.Log.EventDecorator() def move(self): r""" Run the relaxation method to convergence and update the mesh. :return: the iteration count :rtype: :class:`int` """ # Switch to computational coordinates self.to_computational_coordinates() # Take iterations of the relaxed system until reaching convergence for i in range(self.maxiter): self.l2_projector.solve() self._update_physical_coordinates() # Update monitor function self.to_physical_coordinates() self.monitor.interpolate(self.monitor_function(self.mesh)) firedrake.assemble(self.L_P0, tensor=self.volume) self.volume.interpolate(self.volume / self.original_volume) self.to_computational_coordinates() # Evaluate normalisation coefficient self.theta.assign(firedrake.assemble(self.theta_form) / self.total_volume) # Check convergence criteria residual = self.relative_l2_residual if i == 0: initial_norm = residual PETSc.Sys.Print( f"{i:4d}" f" Volume ratio {self.volume_ratio:5.2f}" f" Variation (σ/μ) {self.coefficient_of_variation:8.2e}" f" Residual {residual:8.2e}" ) if residual < self.rtol: self._convergence_message(i + 1) break if residual > self.dtol * initial_norm: self._divergence_error(i + 1) if i == self.maxiter - 1: self._convergence_error(i + 1) # Apply pseudotimestepper and equidistributor self.pseudotimestepper.solve() self.equidistributor.solve() self.phi_old.assign(self.phi) self.H_old.assign(self.H) self.to_physical_coordinates() return i
[docs] class MongeAmpereMover_QuasiNewton(MongeAmpereMover_Base): r""" The standard, elliptic form of the Monge-Ampère equation used for mesh movement is: .. math:: m(\mathbf{x})\det(\mathbf{I} + \mathbf{H}(\phi)) = \theta, for a convex scalar potential :math:`\phi=\phi(\boldsymbol{\xi})`, which is a function of the coordinates of the *computational* mesh. Here :math:`m=m(\mathbf{x})` is a user-provided monitor function, which is a function of the coordinates of the *physical* mesh. :math:`\mathbf{I}` is the identity matrix, :math:`\theta` is a normalisation coefficient and :math:`\mathbf{H}(\phi)` denotes the Hessian of :math:`\phi` with respect to :math:`\boldsymbol{\xi}`. In this mesh mover, the elliptic Monge-Ampère equation is solved using a quasi-Newton method (see :cite:`MCB:18` for details). This approach typically takes fewer than ten iterations to converge, but each iteration is relatively expensive. """ @PETSc.Log.EventDecorator() def __init__(self, mesh, monitor_function, phi_init=None, H_init=None, **kwargs): """ :arg mesh: the physical mesh :type mesh: :class:`firedrake.mesh.MeshGeometry` :arg monitor_function: a Python function which takes a mesh as input :type monitor_function: :class:`~.Callable` :kwarg phi_init: initial guess for the scalar potential :type phi_init: :class:`firedrake.function.Function` :kwarg H_init: initial guess for the Hessian :type H_init: :class:`firedrake.function.Function` :kwarg maxiter: maximum number of iterations for the Quasi-Newton solver :type maxiter: :class:`int` :kwarg rtol: relative tolerance for the residual :type rtol: :class:`float` :kwarg dtol: divergence tolerance for the residual :type dtol: :class:`float` """ super().__init__(mesh, monitor_function=monitor_function, **kwargs) # Initialise phi and H if phi_init or H_init: self.apply_initial_guess(phi_init, H_init) # Setup residuals I = ufl.Identity(self.dim) self.theta_form = self.monitor * ufl.det(I + self.H_old) * self.dx self.residual = self.monitor * ufl.det(I + self.H_old) - self.theta psi = firedrake.TestFunction(self.P1) self._residual_l2_form = psi * self.residual * self.dx self._norm_l2_form = psi * self.theta * self.dx def _create_functions(self): super()._create_functions() self.V = self.P1 * self.P1_ten self.phi_and_hessian = firedrake.Function(self.V) self.phi, self.H = self.phi_and_hessian.subfunctions self.phi_and_hessian_old = firedrake.Function(self.V) self.phi_old, self.H_old = self.phi_and_hessian_old.subfunctions @property @PETSc.Log.EventDecorator() def equidistributor(self): r""" Setup the equidistributor for the quasi-newton method. The equidistributor solves .. math:: \int_{\Omega} \boldsymbol{\tau} \cdot \mathbf{H} \, \mathrm{d}x + \int_{\Omega} (\nabla \cdot \boldsymbol{\tau}) \cdot (\nabla \phi) \, \mathrm{d}x - \int_{\partial \Omega} (((\nabla \phi) \cdot \widehat{\mathbf{n}}) \cdot \boldsymbol{\tau}) \cdot \widehat{\mathbf{n}} \, \mathrm{d}s - \int_{\Omega} \psi (m \det(\mathbf{I} + \mathbf{H}) - \theta) \, \mathrm{d}x = 0, \quad \forall \boldsymbol{\tau} \in \mathbb{P}1^{d \times d}, \quad \forall \psi \in \mathbb{P}1, for the potential :math:`\phi` and its Hessian :math:`\mathbf{H}`, where :math:`d` is the spatial dimension and :math:`\widehat{\mathbf{n}}` is a normal vector to the boundary. :return: the equidistributor :rtype: :class:`~.NonlinearVariationalSolver` """ if hasattr(self, "_equidistributor"): return self._equidistributor n = ufl.FacetNormal(self.mesh) I = ufl.Identity(self.dim) phi, H = firedrake.split(self.phi_and_hessian) psi, tau = firedrake.TestFunctions(self.V) F = ( ufl.inner(tau, H) * self.dx + ufl.dot(ufl.div(tau), ufl.grad(phi)) * self.dx - ufl.dot(ufl.dot(tangential(ufl.grad(phi), n), tau), n) * self.ds - psi * (self.monitor * ufl.det(I + H) - self.theta) * self.dx ) phi, H = firedrake.TrialFunctions(self.V) @PETSc.Log.EventDecorator("MongeAmpereMover.update_monitor") def update_monitor(cursol): """ Callback for updating the monitor function. """ with self.phi_and_hessian_old.dat.vec as v: cursol.copy(v) self.l2_projector.solve() self._update_physical_coordinates() self.to_physical_coordinates() self.monitor.interpolate(self.monitor_function(self.mesh)) self.to_computational_coordinates() self.theta.assign( firedrake.assemble(self.theta_form) * self.total_volume ** (-1) ) # Setup the variational problem # ============================= # We use a custom preconditioner Jp, chosen to approximate the Jacobian of the # system. It includes terms that represent the inner product of tau and H, # the product of phi and psi, and the inner product of the gradients of phi and # psi. This helps in stabilising the solver and improving convergence. Jp = ( ufl.inner(tau, H) * self.dx + phi * psi * self.dx + ufl.inner(ufl.grad(phi), ufl.grad(psi)) * self.dx ) problem = firedrake.NonlinearVariationalProblem(F, self.phi_and_hessian, Jp=Jp) # Setup the variational solver # ============================ # A nullspace is defined to handle the invariance of the solution under certain # transformations. The first component is a constant vector space basis, since # constant shifts in phi do not affect the solution. nullspace = firedrake.MixedVectorSpaceBasis( self.V, [firedrake.VectorSpaceBasis(constant=True), self.V.sub(1)] ) # Note that different solver parameters are used for serial and parallel runs sp = ( solver_parameters.serial_qn if firedrake.COMM_WORLD.size == 1 else solver_parameters.parallel_qn ) sp["snes_atol"] = self.rtol sp["snes_max_it"] = self.maxiter self._equidistributor = firedrake.NonlinearVariationalSolver( problem, nullspace=nullspace, transpose_nullspace=nullspace, pre_function_callback=update_monitor, pre_jacobian_callback=update_monitor, solver_parameters=sp, ) @no_annotations @PETSc.Log.EventDecorator("MongeAmpereMover.callback") def callback(snes, i, rnorm): """ Print progress of the optimisation to screen. Note that convergence is not actually checked. """ cursol = snes.getSolution() update_monitor(cursol) self.to_physical_coordinates() firedrake.assemble(self.L_P0, tensor=self.volume) self.volume.interpolate(self.volume / self.original_volume) self.to_computational_coordinates() PETSc.Sys.Print( f"{i:4d}" f" Volume ratio {self.volume_ratio:5.2f}" f" Variation (σ/μ) {self.coefficient_of_variation:8.2e}" f" Residual {self.relative_l2_residual:8.2e}" ) self.snes = self._equidistributor.snes self.snes.setMonitor(callback) return self._equidistributor
[docs] @PETSc.Log.EventDecorator() def move(self): r""" Run the quasi-Newton method to convergence and update the mesh. :return: the iteration count :rtype: :class:`int` """ # Solve equidistribution problem, handling convergence errors according to # desired behaviour try: self.equidistributor.solve() self._convergence_message(self.snes.getIterationNumber()) except fexc.ConvergenceError as conv_err: self._convergence_error(self.snes.getIterationNumber(), exception=conv_err) # Update mesh coordinates accordingly self.to_physical_coordinates() return self.snes.getIterationNumber()