import numpy
import ufl
from ufl.corealg.map_dag import map_expr_dag
from ufl.corealg.multifunction import MultiFunction
from ufl.domain import extract_unique_domain
from functools import singledispatch, partial
import firedrake
from firedrake.petsc import PETSc
from . import utils
__all__ = ["coarsen"]
class CoarseningError(Exception):
"""Exception raised when coarsening symbolic information fails."""
pass
class CoarsenIntegrand(MultiFunction):
"""'Coarsen' a :class:`ufl.Expr` by replacing coefficients,
arguments and domain data with coarse mesh equivalents."""
def __init__(self, coarsen, coefficient_mapping):
self.coefficient_mapping = coefficient_mapping or {}
self.coarsen = coarsen
super(CoarsenIntegrand, self).__init__()
expr = MultiFunction.reuse_if_untouched
def argument(self, o):
V = self.coarsen(o.function_space(), self.coarsen)
return o.reconstruct(V)
def coefficient(self, o):
return self.coarsen(o, self.coarsen, coefficient_mapping=self.coefficient_mapping)
def geometric_quantity(self, o):
return type(o)(self.coarsen(extract_unique_domain(o), self.coarsen))
def circumradius(self, o):
mesh = self.coarsen(extract_unique_domain(o), self.coarsen)
return firedrake.Circumradius(mesh)
def facet_normal(self, o):
mesh = self.coarsen(extract_unique_domain(o), self.coarsen)
return firedrake.FacetNormal(mesh)
[docs]
@singledispatch
def coarsen(expr, self, coefficient_mapping=None):
# Default, just send it back
return expr
@coarsen.register(ufl.Mesh)
def coarsen_mesh(mesh, self, coefficient_mapping=None):
hierarchy, level = utils.get_level(mesh)
if hierarchy is None:
raise CoarseningError("No mesh hierarchy available")
return hierarchy[level - 1]
@coarsen.register(ufl.classes.Expr)
def coarse_expr(expr, self, coefficient_mapping=None):
if expr is None:
return None
mapper = CoarsenIntegrand(self, coefficient_mapping)
return map_expr_dag(mapper, expr)
@coarsen.register(ufl.Form)
def coarsen_form(form, self, coefficient_mapping=None):
"""Return a coarse mesh version of a form
:arg form: The :class:`~ufl.classes.Form` to coarsen.
:kwarg mapping: an optional map from coefficients to their
coarsened equivalents.
This maps over the form and replaces coefficients and arguments
with their coarse mesh equivalents."""
if form is None:
return None
mapper = CoarsenIntegrand(self, coefficient_mapping)
integrals = []
for it in form.integrals():
integrand = map_expr_dag(mapper, it.integrand())
mesh = it.ufl_domain()
hierarchy, level = utils.get_level(mesh)
new_mesh = hierarchy[level-1]
if isinstance(integrand, ufl.classes.Zero):
continue
if it.subdomain_data() is not None:
raise CoarseningError("Don't know how to coarsen subdomain data")
new_itg = it.reconstruct(integrand=integrand,
domain=new_mesh)
integrals.append(new_itg)
form = ufl.Form(integrals)
form._cache["coefficient_mapping"] = coefficient_mapping
return form
@coarsen.register(firedrake.DirichletBC)
def coarsen_bc(bc, self, coefficient_mapping=None):
V = self(bc.function_space(), self, coefficient_mapping=coefficient_mapping)
val = self(bc.function_arg, self, coefficient_mapping=coefficient_mapping)
subdomain = bc.sub_domain
return type(bc)(V, val, subdomain)
@coarsen.register(firedrake.functionspaceimpl.FunctionSpace)
@coarsen.register(firedrake.functionspaceimpl.WithGeometry)
def coarsen_function_space(V, self, coefficient_mapping=None):
if hasattr(V, "_coarse"):
return V._coarse
from firedrake.dmhooks import get_parent, push_parent, pop_parent, add_hook
fine = V
indices = []
while True:
if V.index is not None:
indices.append(V.index)
if V.component is not None:
indices.append(V.component)
if V.parent is not None:
V = V.parent
else:
break
mesh = self(V.mesh(), self)
V = firedrake.FunctionSpace(mesh, V.ufl_element())
for i in reversed(indices):
V = V.sub(i)
V._fine = fine
fine._coarse = V
# FIXME: This replicates some code from dmhooks.coarsen, but we
# can't do things there because that code calls this code.
# We need to move these operators over here because if we have
# fieldsplits + MG with auxiliary coefficients on spaces other
# than which we do the MG, dm.coarsen is never called, so the
# hooks are not attached. Instead we just call (say) inject which
# coarsens the functionspace.
cdm = V.dm
parent = get_parent(fine.dm)
try:
add_hook(parent, setup=partial(push_parent, cdm, parent), teardown=partial(pop_parent, cdm, parent),
call_setup=True)
except ValueError:
# Not in an add_hooks context
pass
return V
@coarsen.register(firedrake.Function)
def coarsen_function(expr, self, coefficient_mapping=None):
if coefficient_mapping is None:
coefficient_mapping = {}
new = coefficient_mapping.get(expr)
if new is None:
V = expr.function_space()
manager = firedrake.dmhooks.get_transfer_manager(expr.function_space().dm)
V = self(expr.function_space(), self)
new = firedrake.Function(V, name="coarse_%s" % expr.name())
manager.inject(expr, new)
return new
@coarsen.register(firedrake.Constant)
def coarsen_constant(expr, self, coefficient_mapping=None):
if coefficient_mapping is None:
coefficient_mapping = {}
new = coefficient_mapping.get(expr)
if new is None:
mesh = self(extract_unique_domain(expr), self)
new = firedrake.Constant(numpy.zeros(expr.ufl_shape,
dtype=expr.dat.dtype),
domain=mesh)
# Share data pointer
new.dat = expr.dat
return new
@coarsen.register(firedrake.NonlinearVariationalProblem)
def coarsen_nlvp(problem, self, coefficient_mapping=None):
# Build set of coefficients we need to coarsen
seen = set()
coefficients = problem.F.coefficients() + problem.J.coefficients()
if problem.Jp is not None:
coefficients = coefficients + problem.Jp.coefficients()
# Coarsen them, and remember where from.
if coefficient_mapping is None:
coefficient_mapping = {}
for c in coefficients:
if c not in seen:
coefficient_mapping[c] = self(c, self, coefficient_mapping=coefficient_mapping)
seen.add(c)
u = coefficient_mapping[problem.u]
bcs = [self(bc, self) for bc in problem.bcs]
J = self(problem.J, self, coefficient_mapping=coefficient_mapping)
Jp = self(problem.Jp, self, coefficient_mapping=coefficient_mapping)
F = self(problem.F, self, coefficient_mapping=coefficient_mapping)
problem = firedrake.NonlinearVariationalProblem(F, u, bcs=bcs, J=J, Jp=Jp,
form_compiler_parameters=problem.form_compiler_parameters)
return problem
@coarsen.register(firedrake.VectorSpaceBasis)
def coarsen_vectorspacebasis(basis, self, coefficient_mapping=None):
coarse_vecs = [self(vec, self, coefficient_mapping=coefficient_mapping) for vec in basis._vecs]
vsb = firedrake.VectorSpaceBasis(coarse_vecs, constant=basis._constant)
vsb.orthonormalize()
return vsb
@coarsen.register(firedrake.MixedVectorSpaceBasis)
def coarsen_mixedvectorspacebasis(mspbasis, self, coefficient_mapping=None):
coarse_V = self(mspbasis._function_space, self, coefficient_mapping=coefficient_mapping)
coarse_bases = []
for basis in mspbasis._bases:
if isinstance(basis, firedrake.VectorSpaceBasis):
coarse_bases.append(self(basis, self, coefficient_mapping=coefficient_mapping))
elif basis.index is not None:
coarse_bases.append(coarse_V.sub(basis.index))
else:
raise RuntimeError("MixedVectorSpaceBasis can only contain vector space bases or indexed function spaces")
return firedrake.MixedVectorSpaceBasis(coarse_V, coarse_bases)
@coarsen.register(firedrake.solving_utils._SNESContext)
def coarsen_snescontext(context, self, coefficient_mapping=None):
if coefficient_mapping is None:
coefficient_mapping = {}
# Have we already done this?
coarse = context._coarse
if coarse is not None:
return coarse
problem = self(context._problem, self, coefficient_mapping=coefficient_mapping)
appctx = context.appctx
new_appctx = {}
for k in sorted(appctx.keys()):
v = appctx[k]
if k != "state":
# Constructor makes this one.
try:
new_appctx[k] = self(v, self, coefficient_mapping=coefficient_mapping)
except CoarseningError:
# Assume not something that needs coarsening (e.g. float)
new_appctx[k] = v
coarse = type(context)(problem,
mat_type=context.mat_type,
pmat_type=context.pmat_type,
appctx=new_appctx,
transfer_manager=context.transfer_manager)
coarse._fine = context
context._coarse = coarse
# Now that we have the coarse snescontext, push it to the coarsened DMs
# Otherwise they won't have the right transfer manager when they are
# coarsened in turn
from firedrake.dmhooks import get_appctx, push_appctx, pop_appctx
from firedrake.dmhooks import add_hook, get_parent
from itertools import chain
for val in chain(coefficient_mapping.values(), (bc.function_arg for bc in problem.bcs)):
if isinstance(val, firedrake.function.Function):
V = val.function_space()
coarseneddm = V.dm
parentdm = get_parent(context._problem.u.function_space().dm)
# Now attach the hook to the parent DM
if get_appctx(coarseneddm) is None:
push_appctx(coarseneddm, coarse)
teardown = partial(pop_appctx, coarseneddm, coarse)
add_hook(parentdm, teardown=teardown)
ises = problem.J.arguments()[0].function_space()._ises
coarse._nullspace = self(context._nullspace, self, coefficient_mapping=coefficient_mapping)
coarse.set_nullspace(coarse._nullspace, ises, transpose=False, near=False)
coarse._nullspace_T = self(context._nullspace_T, self, coefficient_mapping=coefficient_mapping)
coarse.set_nullspace(coarse._nullspace_T, ises, transpose=True, near=False)
coarse._near_nullspace = self(context._near_nullspace, self, coefficient_mapping=coefficient_mapping)
coarse.set_nullspace(coarse._near_nullspace, ises, transpose=False, near=True)
return coarse
class Interpolation(object):
def __init__(self, cfn, ffn, manager, cbcs=None, fbcs=None):
self.cfn = cfn
self.ffn = ffn
self.cbcs = cbcs or []
self.fbcs = fbcs or []
self.manager = manager
def mult(self, mat, x, y, inc=False):
with self.cfn.dat.vec_wo as v:
x.copy(v)
self.manager.prolong(self.cfn, self.ffn)
for bc in self.fbcs:
bc.zero(self.ffn)
with self.ffn.dat.vec_ro as v:
if inc:
y.axpy(1.0, v)
else:
v.copy(y)
def multAdd(self, mat, x, y, w):
if y.handle == w.handle:
self.mult(mat, x, w, inc=True)
else:
self.mult(mat, x, w)
w.axpy(1.0, y)
def multTranspose(self, mat, x, y, inc=False):
with self.ffn.dat.vec_wo as v:
x.copy(v)
self.manager.restrict(self.ffn, self.cfn)
for bc in self.cbcs:
bc.zero(self.cfn)
with self.cfn.dat.vec_ro as v:
if inc:
y.axpy(1.0, v)
else:
v.copy(y)
def multTransposeAdd(self, mat, x, y, w):
if y.handle == w.handle:
self.multTranspose(mat, x, w, inc=True)
else:
self.multTranspose(mat, x, w)
w.axpy(1.0, y)
class Injection(object):
def __init__(self, cfn, ffn, manager, cbcs=None):
self.cfn = cfn
self.ffn = ffn
self.cbcs = cbcs or []
self.manager = manager
def multTranspose(self, mat, x, y):
with self.ffn.dat.vec_wo as v:
x.copy(v)
self.manager.inject(self.ffn, self.cfn)
for bc in self.cbcs:
bc.apply(self.cfn)
with self.cfn.dat.vec_ro as v:
v.copy(y)
def create_interpolation(dmc, dmf):
cctx = firedrake.dmhooks.get_appctx(dmc)
fctx = firedrake.dmhooks.get_appctx(dmf)
manager = firedrake.dmhooks.get_transfer_manager(dmf)
V_c = cctx._problem.u.function_space()
V_f = fctx._problem.u.function_space()
row_size = V_f.dof_dset.layout_vec.getSizes()
col_size = V_c.dof_dset.layout_vec.getSizes()
cfn = firedrake.Function(V_c)
ffn = firedrake.Function(V_f)
cbcs = cctx._problem.bcs
fbcs = fctx._problem.bcs
ctx = Interpolation(cfn, ffn, manager, cbcs, fbcs)
mat = PETSc.Mat().create(comm=dmc.comm)
mat.setSizes((row_size, col_size))
mat.setType(mat.Type.PYTHON)
mat.setPythonContext(ctx)
mat.setUp()
return mat, None
def create_injection(dmc, dmf):
cctx = firedrake.dmhooks.get_appctx(dmc)
fctx = firedrake.dmhooks.get_appctx(dmf)
manager = firedrake.dmhooks.get_transfer_manager(dmf)
V_c = cctx._problem.u.function_space()
V_f = fctx._problem.u.function_space()
row_size = V_f.dof_dset.layout_vec.getSizes()
col_size = V_c.dof_dset.layout_vec.getSizes()
cfn = firedrake.Function(V_c)
ffn = firedrake.Function(V_f)
cbcs = cctx._problem.bcs
ctx = Injection(cfn, ffn, manager, cbcs)
mat = PETSc.Mat().create(comm=dmc.comm)
mat.setSizes((row_size, col_size))
mat.setType(mat.Type.PYTHON)
mat.setPythonContext(ctx)
mat.setUp()
return mat
