"""
Firedrake uses PETSc for its linear and nonlinear solvers. The
interaction is carried out through DM objects. These carry around any
user-defined application context and can be used to inform the solvers
how to create field decompositions (for fieldsplit preconditioning) as
well as creating sub-DMs (which only contain some fields), along with
multilevel information (for geometric multigrid)
The way Firedrake interacts with these DMs is, broadly, as follows:
A DM is tied to a :class:`~.FunctionSpace` and remembers what function
space that is. To avoid reference cycles defeating the garbage
collector, the DM holds a weakref to the FunctionSpace (which holds a
strong reference to the DM). Use :func:`get_function_space` to get
the function space attached to the DM, and :func:`set_function_space`
to attach it.
Similarly, when a DM is used in a solver, an application context is
attached to it, such that when PETSc calls back into Firedrake, we can
grab the relevant information (how to make the Jacobian, etc...).
This functions in a similar way using :func:`push_appctx` and
:func:`get_appctx` on the DM. You can set whatever you like in here,
but most of the rest of Firedrake expects to find either ``None`` or
else a :class:`firedrake.solving_utils._SNESContext` object.
A crucial part of this, for composition with multi-level solvers
(``-pc_type mg`` and ``-snes_type fas``) is decomposing the DMs. When
a field decomposition is created, the callback
:func:`create_field_decomposition` checks to see if an application
context exists. If so, it splits it apart (one for each of fields)
and attaches these split contexts to the subdms returned to PETSc.
This facilitates runtime composition with multilevel solvers. When
coarsening a DM, the application context is coarsened and transferred
to the coarse DM. The combination of these two symbolic transfer
operations allow us to nest geometric multigrid preconditioning inside
fieldsplit preconditioning, without having to set everything up in
advance.
"""
import weakref
import numpy
from functools import partial
import firedrake
from firedrake.petsc import PETSc
from firedrake.mesh import MeshSequenceGeometry
[docs]
@PETSc.Log.EventDecorator()
def get_function_space(dm):
"""Get the :class:`~.FunctionSpace` attached to this DM.
:arg dm: The DM to get the function space from.
:raises RuntimeError: if no function space was found.
"""
info = dm.getAttr("__fs_info__")
meshref_tuple, element, indices, (name, names), boundary_sets = info
if len(meshref_tuple) == 1:
mesh = meshref_tuple[0]()
else:
mesh = MeshSequenceGeometry([meshref() for meshref in meshref_tuple])
if mesh is None:
raise RuntimeError("Somehow your mesh was collected, this should never happen")
V = firedrake.FunctionSpace(mesh, element, name=name)
if any(boundary_sets):
V = firedrake.bcs.restricted_function_space(V, boundary_sets)
if len(V) > 1:
for V_, name in zip(V, names):
V_.topological.name = name
for index in indices:
V = V.sub(index)
return V
[docs]
@PETSc.Log.EventDecorator()
def set_function_space(dm, V):
"""Set the :class:`~.FunctionSpace` on this DM.
:arg dm: The DM
:arg V: The function space.
.. note::
This stores the information necessary to make a function space given a DM.
"""
indices = []
names = []
while V.parent is not None:
if V.index is not None:
assert V.component is None
indices.append(V.index)
if V.component is not None:
assert V.index is None
indices.append(V.component)
V = V.parent
mesh = V.mesh()
if len(V) > 1:
names = tuple(V_.name for V_ in V)
element = V.ufl_element()
boundary_sets = tuple(V_.boundary_set for V_ in V)
info = (tuple(weakref.ref(m) for m in mesh), element, tuple(reversed(indices)), (V.name, names), boundary_sets)
dm.setAttr("__fs_info__", info)
# Attribute management on DMs. Since they are reused in multiple
# places, use a stack.
[docs]
def push_attr(attr, dm, obj):
stack = dm.getAttr(attr)
if stack is None:
stack = []
dm.setAttr(attr, stack)
stack.append(obj)
[docs]
def pop_attr(attr, dm, match=None):
stack = dm.getAttr(attr)
if not stack:
return None
obj = stack.pop()
if match is not None and obj != match:
stack.append(obj)
return None
else:
return obj
[docs]
def get_attr(attr, dm, default=None):
stack = dm.getAttr(attr)
if not stack:
return default
return stack[-1]
[docs]
class SetupHooks(object):
"""Hooks run for setup and teardown of DMs inside solvers.
Used for transferring problem-specific data onto subproblems.
You probably don't want to use this directly, instead see
:class:`~add_hooks` or :func:`add_hook`."""
def __init__(self):
self._setup = []
self._teardown = []
[docs]
def add_setup(self, f):
self._setup.append(f)
[docs]
def add_teardown(self, f):
self._teardown.append(f)
[docs]
def setup(self):
for f in self._setup:
f()
[docs]
def teardown(self):
# We push onto a stack, so to pop off in the correct order, go backwards.
for f in reversed(self._teardown):
f()
[docs]
def add_hook(dm, setup=None, teardown=None, call_setup=False, call_teardown=False):
"""Add a hook to a DM to be called for setup/teardown of
subproblems.
:arg dm: The DM to save the hooks on. This is normally the DM
associated with the Firedrake solver.
:arg setup: function of no arguments to call to set up subproblem
data.
:arg teardown: function of no arguments to call to remove
subproblem data.
:arg call_setup: Should the setup function be called now?
:arg call_teardown: Should the teardown function be called now?
See also :class:`add_hooks` which provides a context manager which
manages everything."""
stack = dm.getAttr("__setup_hooks__")
if not stack:
raise ValueError("Expecting non-empty stack")
obj = stack[-1]
if setup is not None:
obj.add_setup(setup)
if call_setup:
setup()
if teardown is not None:
obj.add_teardown(teardown)
if call_teardown:
teardown()
[docs]
class add_hooks(object):
"""Context manager for adding subproblem setup hooks to a DM.
:arg DM: The DM to remember setup/teardown for.
:arg obj: The object that we're going to setup, typically a solver
of some kind: this is where the hooks are saved.
:arg save: Save this round of setup? Set this to False if all
you're going to do is setFromOptions.
:arg appctx: An application context to attach to the top-level DM
that describes the problem-specific data.
This is your normal entry-point for setting up problem specific
data on subdms. You would likely do something like, for a Python PC.
.. code::
# In setup
pc = ...
pc.setDM(dm)
with dmhooks.add_hooks(dm, self, appctx=ctx, save=False):
pc.setFromOptions()
...
# in apply
dm = pc.getDM()
with dmhooks.add_hooks(dm, self, appctx=self.ctx):
pc.apply(...)
"""
def __init__(self, dm, obj, *, save=True, appctx=None):
self.dm = dm
self.obj = obj
self.first_time = not hasattr(obj, "setup_hooks")
self.save = save
self.appctx = appctx
if not (self.save or self.first_time):
raise ValueError("Can't have save=False for non-first-time usage")
def __enter__(self):
if not self.first_time:
# We've already run setup, so just attach the data to the subdms.
hooks = self.obj.setup_hooks
push_attr("__setup_hooks__", self.dm, hooks)
hooks.setup()
else:
# Not yet seen, let's save the relevant information.
hooks = SetupHooks()
if self.save:
# Remember it for later
self.obj.setup_hooks = hooks
push_attr("__setup_hooks__", self.dm, hooks)
if self.appctx is not None:
add_hook(self.dm, setup=partial(push_appctx, self.dm, self.appctx),
teardown=partial(pop_appctx, self.dm, self.appctx),
call_setup=True)
def __exit__(self, typ, value, traceback):
hooks = pop_attr("__setup_hooks__", self.dm)
if self.first_time:
assert hooks is not None
else:
assert hooks == self.obj.setup_hooks
hooks.teardown()
# Things we're going to transfer around DMs
push_parent = partial(push_attr, "__parent__")
pop_parent = partial(pop_attr, "__parent__")
[docs]
def get_parent(dm):
return get_attr("__parent__", dm, default=dm)
push_appctx = partial(push_attr, "__appctx__")
pop_appctx = partial(pop_attr, "__appctx__")
get_appctx = partial(get_attr, "__appctx__")
[docs]
def get_transfer_manager(dm):
appctx = get_appctx(dm)
if appctx is None:
# We're not in a solve, so all we can do is make a new one (not cached)
import warnings
warnings.warn("Creating new TransferManager to transfer data to coarse grids. "
"This might be slow (you probably want to save it on an appctx)", RuntimeWarning)
transfer = firedrake.TransferManager()
else:
transfer = appctx.transfer_manager
return transfer
push_ctx_coarsener = partial(push_attr, "__ctx_coarsener__")
pop_ctx_coarsener = partial(pop_attr, "__ctx_coarsener__")
[docs]
def get_ctx_coarsener(dm):
from firedrake.mg.ufl_utils import coarsen
return get_attr("__ctx_coarsener__", dm, default=coarsen)
[docs]
class ctx_coarsener(object):
def __init__(self, V, coarsen=None):
from firedrake.mg.ufl_utils import coarsen as symbolic_coarsen
self.V = V
if coarsen is None:
coarsen = symbolic_coarsen
self.coarsen = coarsen
def __enter__(self):
push_ctx_coarsener(self.V.dm, self.coarsen)
def __exit__(self, typ, value, traceback):
pop_ctx_coarsener(self.V.dm, self.coarsen)
[docs]
def create_matrix(dm):
"""
Callback to create a matrix from this DM.
:arg DM: The DM.
.. note::
This only works if an application context is set, in which case
it returns the stored Jacobian. This *does not* make a new
matrix.
"""
ctx = get_appctx(dm)
if ctx is None:
raise ValueError("Cannot create matrix from DM with no AppCtx")
# TODO, should make new matrix and change solver
# form_function/jacobian to be able to assemble into a provided
# PETSc matrix.
return ctx._jac.petscmat
[docs]
@PETSc.Log.EventDecorator()
def create_field_decomposition(dm, *args, **kwargs):
"""Callback to decompose a DM.
:arg DM: The DM.
This grabs the function space in the DM, splits it apart (only
makes sense for mixed function spaces) and returns the DMs on each
of the subspaces. If an application context is present on the
input DM, it is split into individual field contexts and set on
the appropriate subdms as well.
"""
W = get_function_space(dm)
# Don't pass split number if name is None (this way the
# recursively created splits have the names you want)
names = [s.name for s in W]
dms = [V.dm for V in W]
ctx = get_appctx(dm)
coarsen = get_ctx_coarsener(dm)
parent = get_parent(dm)
for d in dms:
add_hook(parent, setup=partial(push_parent, d, parent), teardown=partial(pop_parent, d, parent),
call_setup=True)
if ctx is not None and len(W) > 1:
ctxs = ctx.split([(i, ) for i in range(len(W))])
for d, c in zip(dms, ctxs):
add_hook(parent, setup=partial(push_appctx, d, c), teardown=partial(pop_appctx, d, c),
call_setup=True)
add_hook(parent, setup=partial(push_ctx_coarsener, d, coarsen), teardown=partial(pop_ctx_coarsener, d, coarsen),
call_setup=True)
return names, W._ises, dms
[docs]
@PETSc.Log.EventDecorator()
def create_subdm(dm, fields, *args, **kwargs):
"""Callback to create a sub-DM describing the specified fields.
:arg DM: The DM.
:arg fields: The fields in the new sub-DM.
"""
W = get_function_space(dm)
ctx = get_appctx(dm)
coarsen = get_ctx_coarsener(dm)
parent = get_parent(dm)
if len(fields) == 1:
# Subspace is just a single FunctionSpace.
idx, = fields
subdm = W[idx].dm
iset = W._ises[idx]
add_hook(parent, setup=partial(push_parent, subdm, parent), teardown=partial(pop_parent, subdm, parent),
call_setup=True)
if ctx is not None:
ctx, = ctx.split([(idx, )])
add_hook(parent, setup=partial(push_appctx, subdm, ctx), teardown=partial(pop_appctx, subdm, ctx),
call_setup=True)
add_hook(parent, setup=partial(push_ctx_coarsener, subdm, coarsen), teardown=partial(pop_ctx_coarsener, subdm, coarsen),
call_setup=True)
return iset, subdm
else:
# Need to build an MFS for the subspace
subspace = firedrake.MixedFunctionSpace([W[f] for f in fields])
add_hook(parent, setup=partial(push_parent, subspace.dm, parent), teardown=partial(pop_parent, subspace.dm, parent),
call_setup=True)
# Index set mapping from W into subspace.
iset = PETSc.IS().createGeneral(numpy.concatenate([W._ises[f].indices
for f in fields]),
comm=W.comm)
if ctx is not None:
ctx, = ctx.split([fields])
add_hook(parent, setup=partial(push_appctx, subspace.dm, ctx),
teardown=partial(pop_appctx, subspace.dm, ctx),
call_setup=True)
add_hook(parent, setup=partial(push_ctx_coarsener, subspace.dm, coarsen),
teardown=partial(pop_ctx_coarsener, subspace.dm, coarsen),
call_setup=True)
return iset, subspace.dm
[docs]
def set_ksp_operators(dm: PETSc.DM) -> None:
"""Make the KSPs on a DM assemble their operators from the application context.
Parameters
----------
dm
The DM that receives the operator callbacks of :class:`~firedrake.solving_utils._SNESContext`.
"""
dm.setKSPCreateOperators(firedrake.solving_utils._SNESContext.create_operators)
dm.setKSPComputeOperators(firedrake.solving_utils._SNESContext.compute_operators)
[docs]
@PETSc.Log.EventDecorator()
def coarsen(dm, comm):
"""Callback to coarsen a DM.
:arg DM: The DM to coarsen.
:arg comm: The communicator for the new DM (ignored)
This transfers a coarse application context over to the coarsened
DM (if found on the input DM).
"""
from firedrake.mg.utils import get_level
V = get_function_space(dm)
mesh = V.mesh()
hierarchy, level = get_level(mesh)
if level < 1:
raise RuntimeError("Cannot coarsen coarsest DM")
coarsen = get_ctx_coarsener(dm)
Vc = coarsen(V, coarsen)
cdm = Vc.dm
parent = get_parent(dm)
add_hook(parent, setup=partial(push_parent, cdm, parent), teardown=partial(pop_parent, cdm, parent),
call_setup=True)
if len(V) > 1:
for V_, Vc_ in zip(V, Vc):
add_hook(parent, setup=partial(push_parent, Vc_.dm, parent), teardown=partial(pop_parent, Vc_.dm, parent),
call_setup=True)
add_hook(parent, setup=partial(push_ctx_coarsener, cdm, coarsen),
teardown=partial(pop_ctx_coarsener, cdm, coarsen),
call_setup=True)
ctx = get_appctx(dm)
if ctx is not None:
cctx = coarsen(ctx, coarsen)
add_hook(parent, setup=partial(push_appctx, cdm, cctx),
teardown=partial(pop_appctx, cdm, cctx),
call_setup=True)
# These callbacks tell the KSPs on the coarse level how to assemble
# their matrices. cdm inherits them from dm, and a different solver on
# this function space might overwrite them, so we set the callbacks on
# dm before each solve.
add_hook(parent, setup=partial(set_ksp_operators, dm), call_setup=True)
return cdm
def _refine_adaptive(dm):
"""
Return the DM of the `_SNESContext` reconstructed on the adaptively-refined
mesh using `_SNESContext.marking_callback` to mark the cells to be refined.
"""
from firedrake.mg.mesh import MeshHierarchy
from firedrake.mg.ufl_utils import refine
from firedrake.mg.utils import get_level
# DMAdaptorAdapt() unconditionally destroys its input DM, and each
# adapted input DM remains a level in the mesh hierarchy.
# Increase the reference count so the coarse DM survives.
dm.incRef()
ctx = get_appctx(dm)
if ctx is None:
raise RuntimeError("No _SNESContext found on DM")
current_solution = ctx._x
mesh = current_solution.function_space().mesh()
hierarchy, level = get_level(mesh)
if hierarchy is None:
hierarchy = MeshHierarchy(mesh)
level = 0
if level+1 != len(hierarchy):
raise RuntimeError("Adaptive SNES refinement can only add a mesh on top of the finest level")
if ctx._marking_callback is None:
raise RuntimeError("Adaptive SNES refinement requires setting a marking_callback")
markers = ctx._marking_callback(ctx, current_solution)
if not isinstance(markers, (firedrake.Function, firedrake.Cofunction)):
raise TypeError(
f"marking callback must return a Function or Cofunction, not a {type(markers).__name__}"
)
M = markers.function_space()
if M.mesh() is not mesh:
raise ValueError("marking callback must return markers on the current solution mesh")
num_dofs_per_cell = M.finat_element.space_dimension()
if num_dofs_per_cell != 1:
raise ValueError("marking callback must return a DG0 Function or Cofunction")
hierarchy.add_mesh(mesh.refine_marked_elements(markers))
coefficient_mapping = {}
refined_ctx = refine(ctx, refine, coefficient_mapping=coefficient_mapping)
parent = get_parent(dm)
coarsener = get_ctx_coarsener(dm)
# Get all DMs from the refined problem
dms = [refined_ctx._problem.u_restrict.function_space().dm]
for value in coefficient_mapping.values():
if isinstance(value, (firedrake.Function, firedrake.Cofunction)):
value_dm = value.function_space().dm
if value_dm not in dms:
dms.append(value_dm)
# Attach refined context
for refined_dm in dms:
add_hook(parent, setup=partial(push_parent, refined_dm, parent),
teardown=partial(pop_parent, refined_dm, parent),
call_setup=True)
add_hook(parent, setup=partial(push_ctx_coarsener, refined_dm, coarsener),
teardown=partial(pop_ctx_coarsener, refined_dm, coarsener),
call_setup=True)
add_hook(parent, setup=partial(push_appctx, refined_dm, refined_ctx),
teardown=partial(pop_appctx, refined_dm, refined_ctx),
call_setup=True)
return refined_ctx._problem.dm
def _refine_from_hierarchy(dm):
"""Return the DM of the refined function space."""
from firedrake.mg.ufl_utils import refine
Vcoarse = get_function_space(dm)
if Vcoarse is None:
raise RuntimeError("No FunctionSpace found on DM")
Vfine = refine(Vcoarse, refine)
return Vfine.dm
[docs]
@PETSc.Log.EventDecorator()
def refine(dm, comm):
"""Callback to refine a DM.
:arg DM: The DM to refine.
:arg comm: The communicator for the new DM (ignored)
"""
from firedrake.mg.utils import get_level
hierarchy, level = get_level(get_function_space(dm).mesh())
if hierarchy is not None and level+1 < len(hierarchy):
return _refine_from_hierarchy(dm)
else:
return _refine_adaptive(dm)
[docs]
def attach_hooks(dm, level=None, sf=None, section=None):
"""Attach callback hooks to a DM.
:arg DM: The DM to attach callbacks to.
:arg level: Optional refinement level.
:arg sf: Optional PETSc SF object describing the DM's ``points``.
:arg section: Optional PETSc Section object describing the DM's
data layout.
"""
from firedrake.mg.ufl_utils import create_interpolation, create_injection
# Data layout
if sf is not None:
dm.setPointSF(sf)
if section is not None:
dm.setLocalSection(section)
# Multilevel hierarchies
dm.setRefine(refine)
dm.setCoarsen(coarsen)
dm.setCreateMatrix(create_matrix)
dm.setCreateInterpolation(create_interpolation)
dm.setCreateInjection(create_injection)
if level is not None:
dm.setRefineLevel(level)
# Field splitting (these will never be called if the DM references
# a non-mixed space)
dm.setCreateFieldDecomposition(create_field_decomposition)
dm.setCreateSubDM(create_subdm)
