def create_injection(self, dmc, dmf):
prefix = dmc.getOptionsPrefix()
mat_type = PETSc.Options(prefix).getString(
"mg_levels_transfer_mat_type", default="matfree")
I = self.create_transfer(get_appctx(dmf), get_appctx(dmc), mat_type,
False, False)
return PETSc.Mat().createTranspose(I)
def create_interpolation(self, dmc, dmf):
prefix = dmc.getOptionsPrefix()
mat_type = PETSc.Options(prefix).getString(
"mg_levels_transfer_mat_type", default="matfree")
I = self.create_transfer(get_appctx(dmc), get_appctx(dmf), mat_type,
True, False)
return I, None
def create_interpolation(self, dmc, dmf):
# This should be generalised to work for arbitrary function
# spaces. Currently I think it only works for CG/DG on simplices.
# I used the same code as firedrake.P1PC.
cctx = get_appctx(dmc)
fctx = get_appctx(dmf)
cV = cctx.J.arguments()[0].function_space()
fV = fctx.J.arguments()[0].function_space()
cbcs = cctx._problem.bcs
fbcs = fctx._problem.bcs
I = prolongation_matrix(fV, cV, fbcs, cbcs)
R = PETSc.Mat().createTranspose(I)
return R, None
def update(self, w):
if not self.separate_wind:
return
if isinstance(w, Function):
self.wind.assign(w)
else:
with self.wind.dat.vec_wo as wind_v:
w.copy(wind_v)
# Need to update your coarse grid mates, too.
# There's probably a better way to do this.
dm = self.V.dm
wind = self.wind
while get_level(get_function_space(dm).mesh())[1] > 0:
cdm = dm.coarsen()
cctx = get_appctx(cdm)
if cctx is None:
break
cwind = cctx.F._cache['coefficient_mapping'][wind]
inject(wind, cwind)
dm = cdm
wind = cwind
def compute_operators(ksp, J, P):
"""Form the Jacobian for this problem
:arg ksp: a PETSc KSP object
:arg J: the Jacobian (a Mat)
:arg P: the preconditioner matrix (a Mat)
"""
from firedrake import inject
dm = ksp.getDM()
ctx = dmhooks.get_appctx(dm)
problem = ctx._problem
assert J.handle == ctx._jac.petscmat.handle
if problem._constant_jacobian and ctx._jacobian_assembled:
# Don't need to do any work with a constant jacobian
# that's already assembled
return
ctx._jacobian_assembled = True
fine = ctx._fine
if fine is not None:
inject(fine._x, ctx._x)
for bc in ctx._problem.bcs:
bc.apply(ctx._x)
ctx._assemble_jac()
ctx._jac.force_evaluation()
if ctx.Jp is not None:
assert P.handle == ctx._pjac.petscmat.handle
ctx._assemble_pjac()
ctx._pjac.force_evaluation()
def form_jacobian(snes, X, J, P):
"""Form the Jacobian for this problem
:arg snes: a PETSc SNES object
:arg X: the current guess (a Vec)
:arg J: the Jacobian (a Mat)
:arg P: the preconditioner matrix (a Mat)
"""
dm = snes.getDM()
ctx = dmhooks.get_appctx(dm)
problem = ctx._problem
assert J.handle == ctx._jac.petscmat.handle
if problem._constant_jacobian and ctx._jacobian_assembled:
# Don't need to do any work with a constant jacobian
# that's already assembled
return
ctx._jacobian_assembled = True
# X may not be the same vector as the vec behind self._x, so
# copy guess in from X.
with ctx._x.dat.vec as v:
X.copy(v)
if ctx._pre_jacobian_callback is not None:
ctx._pre_jacobian_callback(X)
ctx._assemble_jac()
ctx._jac.force_evaluation()
if ctx.Jp is not None:
assert P.handle == ctx._pjac.petscmat.handle
ctx._assemble_pjac()
ctx._pjac.force_evaluation()
def form_function(snes, X, F):
"""Form the residual for this problem
:arg snes: a PETSc SNES object
:arg X: the current guess (a Vec)
:arg F: the residual at X (a Vec)
"""
dm = snes.getDM()
ctx = dmhooks.get_appctx(dm)
problem = ctx._problem
# X may not be the same vector as the vec behind self._x, so
# copy guess in from X.
with ctx._x.dat.vec as v:
X.copy(v)
if ctx._pre_function_callback is not None:
ctx._pre_function_callback(X)
ctx._assemble_residual()
# no mat_type -- it's a vector!
for bc in problem.bcs:
bc.zero(ctx._F)
# F may not be the same vector as self._F, so copy
# residual out to F.
with ctx._F.dat.vec_ro as v:
v.copy(F)
def initialize(self, obj):
if isinstance(obj, PETSc.PC):
A, P = obj.getOperators()
else:
raise ValueError("Not a PC?")
ctx = get_appctx(obj.getDM())
if ctx is None:
raise ValueError("No context found on form")
if not isinstance(ctx, _SNESContext):
raise ValueError("Don't know how to get form from %r", ctx)
if P.getType() == "python":
ictx = P.getPythonContext()
if ictx is None:
raise ValueError("No context found on matrix")
if not isinstance(ictx, ImplicitMatrixContext):
raise ValueError("Don't know how to get form from %r", ictx)
J = ictx.a
bcs = ictx.row_bcs
if bcs != ictx.col_bcs:
raise NotImplementedError("Row and column bcs must match")
else:
J = ctx.Jp or ctx.J
bcs = ctx._problem.bcs
mesh = J.ufl_domain()
self.plex = mesh._topology_dm
self.ctx = ctx
if mesh.cell_set._extruded:
raise NotImplementedError("Not implemented on extruded meshes")
if "overlap_type" not in mesh._distribution_parameters:
if mesh.comm.size > 1:
# Want to do
# warnings.warn("You almost surely want to set an overlap_type in your mesh's distribution_parameters.")
# but doesn't warn!
PETSc.Sys.Print("Warning: you almost surely want to set an overlap_type in your mesh's distribution_parameters.")
patch = obj.__class__().create(comm=obj.comm)
patch.setOptionsPrefix(obj.getOptionsPrefix() + "matpatch_")
self.configure_patch(patch, obj)
patch.setType("matpatch")
Jstate = None
V, _ = map(operator.methodcaller("function_space"), J.arguments())
patch.setDM(self.plex)
self.plex.setDefaultSection(V.dm.getDefaultSection())
self.plex.setDefaultGlobalSection(V.dm.getDefaultGlobalSection())
# pgraph = self.plex.getDefaultSF().getGraph()
# vgraph = V.dm.getDefaultSF().getGraph()
patch.setAttr("ctx", ctx)
patch.incrementTabLevel(1, parent=obj)
patch.setFromOptions()
patch.setUp()
self.patch = patch
def form_function(snes, X, F):
r"""Form the residual for this problem
:arg snes: a PETSc SNES object
:arg X: the current guess (a Vec)
:arg F: the residual at X (a Vec)
"""
dm = snes.getDM()
ctx = dmhooks.get_appctx(dm)
# X may not be the same vector as the vec behind self._x, so
# copy guess in from X.
with ctx._x.dat.vec_wo as v:
X.copy(v)
if ctx._pre_function_callback is not None:
ctx._pre_function_callback(X)
ctx._assemble_residual()
if ctx._post_function_callback is not None:
with ctx._F.dat.vec as F_:
ctx._post_function_callback(X, F_)
# F may not be the same vector as self._F, so copy
# residual out to F.
with ctx._F.dat.vec_ro as v:
v.copy(F)
def form_function(snes, X, F):
"""Form the residual for this problem
:arg snes: a PETSc SNES object
:arg X: the current guess (a Vec)
:arg F: the residual at X (a Vec)
"""
dm = snes.getDM()
ctx = dmhooks.get_appctx(dm)
problem = ctx._problem
# X may not be the same vector as the vec behind self._x, so
# copy guess in from X.
with ctx._x.dat.vec_wo as v:
X.copy(v)
if ctx._pre_function_callback is not None:
ctx._pre_function_callback(X)
ctx._assemble_residual()
# no mat_type -- it's a vector!
for bc in problem.bcs:
bc.zero(ctx._F)
# F may not be the same vector as self._F, so copy
# residual out to F.
with ctx._F.dat.vec_ro as v:
v.copy(F)
def form_function(ts, t, X, Xdot, F):
r"""Form the residual for this problem
:arg ts: a PETSc TS object
:arg t: the time at step/stage being solved
:arg X: state vector
:arg Xdot: time derivative of state vector
:arg F: function vector
"""
dm = ts.getDM()
ctx = dmhooks.get_appctx(dm)
# X may not be the same vector as the vec behind self._x, so
# copy guess in from X.
with ctx._x.dat.vec_wo as v:
X.copy(v)
with ctx._xdot.dat.vec_wo as v:
Xdot.copy(v)
ctx._time.assign(t)
if ctx._pre_function_callback is not None:
ctx._pre_function_callback(X, Xdot)
ctx._assemble_residual()
if ctx._post_function_callback is not None:
with ctx._F.dat.vec as F_:
ctx._post_function_callback(X, Xdot, F_)
# F may not be the same vector as self._F, so copy
# residual out to F.
with ctx._F.dat.vec_ro as v:
v.copy(F)
def form_function(ts, t, X, F):
r"""Form the residual for this problem
:arg ts: a PETSc TS object
:arg t: the time at step/stage being solved
:arg X: state vector
:arg F: function vector
"""
dm = ts.getDM()
ctx = dmhooks.get_appctx(dm)
# X may not be the same vector as the vec behind self._problem.phi_n, so
# copy guess in from X.
with ctx._problem.phi_n.dat.vec_wo as v:
X.copy(v)
b1 = assemble(ctx._problem.L1)
ctx._problem.solver1.solve(ctx._problem.p1, b1)
b2 = assemble(ctx._problem.L2)
ctx._problem.solver2.solve(ctx._problem.p2, b2)
b3 = assemble(ctx._problem.Lb)
ctx._problem.solver_b.solve(ctx._F, b3)
# F may not be the same vector as self._F, so copy
# residual out to F.
with ctx._F.dat.vec_ro as v:
v.copy(F)
def compute_operators(ksp, J, P):
"""Form the Jacobian for this problem
:arg ksp: a PETSc KSP object
:arg J: the Jacobian (a Mat)
:arg P: the preconditioner matrix (a Mat)
"""
from firedrake import inject
dm = ksp.getDM()
ctx = dmhooks.get_appctx(dm)
problem = ctx._problem
assert J.handle == ctx._jac.petscmat.handle
if problem._constant_jacobian and ctx._jacobian_assembled:
# Don't need to do any work with a constant jacobian
# that's already assembled
return
ctx._jacobian_assembled = True
fine = ctx._fine
if fine is not None:
inject(fine._x, ctx._x)
for bc in ctx._problem.bcs:
bc.apply(ctx._x)
ctx._assemble_jac()
ctx._jac.force_evaluation()
if ctx.Jp is not None:
assert P.handle == ctx._pjac.petscmat.handle
ctx._assemble_pjac()
ctx._pjac.force_evaluation()
def form_jacobian(snes, X, J, P):
"""Form the Jacobian for this problem
:arg snes: a PETSc SNES object
:arg X: the current guess (a Vec)
:arg J: the Jacobian (a Mat)
:arg P: the preconditioner matrix (a Mat)
"""
dm = snes.getDM()
ctx = dmhooks.get_appctx(dm)
problem = ctx._problem
assert J.handle == ctx._jac.petscmat.handle
if problem._constant_jacobian and ctx._jacobian_assembled:
# Don't need to do any work with a constant jacobian
# that's already assembled
return
ctx._jacobian_assembled = True
# X may not be the same vector as the vec behind self._x, so
# copy guess in from X.
with ctx._x.dat.vec_wo as v:
X.copy(v)
if ctx._pre_jacobian_callback is not None:
ctx._pre_jacobian_callback(X)
ctx._assemble_jac()
ctx._jac.force_evaluation()
if ctx.Jp is not None:
assert P.handle == ctx._pjac.petscmat.handle
ctx._assemble_pjac()
ctx._pjac.force_evaluation()
def compute_operators(ksp, J, P):
r"""Form the Jacobian for this problem
:arg ksp: a PETSc KSP object
:arg J: the Jacobian (a Mat)
:arg P: the preconditioner matrix (a Mat)
"""
from firedrake.bcs import DirichletBC
dm = ksp.getDM()
ctx = dmhooks.get_appctx(dm)
problem = ctx._problem
assert J.handle == ctx._jac.petscmat.handle
# TODO: Check how to use constant jacobian properly for TS
if problem._constant_jacobian and ctx._jacobian_assembled:
# Don't need to do any work with a constant jacobian
# that's already assembled
return
ctx._jacobian_assembled = True
fine = ctx._fine
if fine is not None:
manager = dmhooks.get_transfer_operators(
fine._x.function_space().dm)
manager.inject(fine._x, ctx._x)
for bc in itertools.chain(*ctx._problem.bcs):
if isinstance(bc, DirichletBC):
bc.apply(ctx._x)
ctx._assemble_jac()
if ctx.Jp is not None:
assert P.handle == ctx._pjac.petscmat.handle
ctx._assemble_pjac()
def step(self, snes, x, f, y):
ctx = get_appctx(snes.dm)
push_appctx(self.ppc.dm, ctx)
x.copy(y)
self.ppc.solve(snes.vec_rhs or self.dummy, y)
y.aypx(-1, x)
snes.setConvergedReason(self.ppc.getConvergedReason())
pop_appctx(self.ppc.dm)
def form(self, pc, test, trial):
_, P = pc.getOperators()
if P.getType() == "python":
context = P.getPythonContext()
return (context.a, context.row_bcs)
else:
context = dmhooks.get_appctx(pc.getDM())
return (context.Jp or context.J, context._problem.bcs)
def initialize(self, pc):
# Make a new DM.
# Hook up a (new) coarsen routine on that DM.
# Make a new PC, of type MG.
# Assign the DM to that PC.
odm = pc.getDM()
ctx = get_appctx(odm)
test, trial = ctx.J.arguments()
if test.function_space() != trial.function_space():
raise NotImplementedError("test and trial spaces must be the same")
prefix = pc.getOptionsPrefix()
options_prefix = prefix + "pmg_"
pdm = PETSc.DMShell().create(comm=pc.comm)
pdm.setOptionsPrefix(options_prefix)
# Get the coarse degree from PETSc options
self.coarse_degree = PETSc.Options(options_prefix).getInt("mg_coarse_degree", default=1)
# Construct a list with the elements we'll be using
V = test.function_space()
ele = V.ufl_element()
elements = [ele]
while True:
try:
ele_ = self.coarsen_element(ele)
assert ele_.value_shape() == ele.value_shape()
ele = ele_
except ValueError:
break
elements.append(ele)
sf = odm.getPointSF()
section = odm.getDefaultSection()
attach_hooks(pdm, level=len(elements)-1, sf=sf, section=section)
# Now overwrite some routines on the DM
pdm.setRefine(None)
pdm.setCoarsen(self.coarsen)
pdm.setCreateInterpolation(self.create_interpolation)
# We need this for p-FAS
pdm.setCreateInjection(self.create_injection)
pdm.setSNESJacobian(_SNESContext.form_jacobian)
pdm.setSNESFunction(_SNESContext.form_function)
pdm.setKSPComputeOperators(_SNESContext.compute_operators)
set_function_space(pdm, get_function_space(odm))
parent = get_parent(odm)
assert parent is not None
add_hook(parent, setup=partial(push_parent, pdm, parent), teardown=partial(pop_parent, pdm, parent), call_setup=True)
add_hook(parent, setup=partial(push_appctx, pdm, ctx), teardown=partial(pop_appctx, pdm, ctx), call_setup=True)
self.ppc = self.configure_pmg(pc, pdm)
self.ppc.setFromOptions()
self.ppc.setUp()
def form(self, pc, test, trial):
_, P = pc.getOperators()
if P.getType() == "python":
context = P.getPythonContext()
return (context.a, context.row_bcs)
else:
from firedrake.dmhooks import get_appctx
context = get_appctx(pc.getDM())
return (context.Jp or context.J, context._problem.bcs)
def create_interpolation(self, dmc, dmf):
cctx = get_appctx(dmc)
fctx = get_appctx(dmf)
cV = cctx.J.arguments()[0].function_space()
fV = fctx.J.arguments()[0].function_space()
cbcs = cctx._problem.bcs
fbcs = fctx._problem.bcs
prefix = dmc.getOptionsPrefix()
mattype = PETSc.Options(prefix).getString("mg_levels_transfer_mat_type", default="matfree")
if mattype == "matfree":
I = prolongation_matrix_matfree(fV, cV, fbcs, cbcs)
elif mattype == "aij":
I = prolongation_matrix_aij(fV, cV, fbcs, cbcs)
else:
raise ValueError("Unknown matrix type")
R = PETSc.Mat().createTranspose(I)
return R, None
def form_jacobian(ts, t, X, Xdot, shift, J, P):
r"""Form the Jacobian for this problem
:arg ts: a PETSc TS object
:arg t: the time at step/stage being solved
:arg X: state vector
:arg Xdot: time derivative of state vector
:arg J: the Jacobian (a Mat)
:arg P: the preconditioner matrix (a Mat)
"""
dm = ts.getDM()
ctx = dmhooks.get_appctx(dm)
problem = ctx._problem
assert J.handle == ctx._jac.petscmat.handle
# TODO: Check how to use constant jacobian properly for TS
if problem._constant_jacobian and ctx._jacobian_assembled:
# Don't need to do any work with a constant jacobian
# that's already assembled
return
ctx._jacobian_assembled = True
# X may not be the same vector as the vec behind self._x, so
# copy guess in from X.
with ctx._x.dat.vec_wo as v:
X.copy(v)
with ctx._xdot.dat.vec_wo as v:
Xdot.copy(v)
ctx._time.assign(t)
if ctx._pre_jacobian_callback is not None:
ctx._pre_jacobian_callback(X, Xdot)
ctx._shift.assign(shift)
ctx._assemble_jac()
if ctx._post_jacobian_callback is not None:
ctx._post_jacobian_callback(X, Xdot, J)
if ctx.Jp is not None:
assert P.handle == ctx._pjac.petscmat.handle
ctx._assemble_pjac()
ises = problem.J.arguments()[0].function_space()._ises
ctx.set_nullspace(ctx._nullspace, ises, transpose=False, near=False)
ctx.set_nullspace(ctx._nullspace_T, ises, transpose=True, near=False)
ctx.set_nullspace(ctx._near_nullspace,
ises,
transpose=False,
near=True)
def form_jacobian(snes, X, J, P):
r"""Form the Jacobian for this problem
:arg snes: a PETSc SNES object
:arg X: the current guess (a Vec)
:arg J: the Jacobian (a Mat)
:arg P: the preconditioner matrix (a Mat)
"""
dm = snes.getDM()
ctx = dmhooks.get_appctx(dm)
problem = ctx._problem
assert J.handle == ctx._jac.petscmat.handle
if problem._constant_jacobian and ctx._jacobian_assembled:
# Don't need to do any work with a constant jacobian
# that's already assembled
return
ctx._jacobian_assembled = True
# X may not be the same vector as the vec behind self._x, so
# copy guess in from X.
with ctx._x.dat.vec_wo as v:
X.copy(v)
if ctx._pre_jacobian_callback is not None:
ctx._pre_jacobian_callback(X)
ctx._assemble_jac()
if ctx._post_jacobian_callback is not None:
ctx._post_jacobian_callback(X, J)
if ctx.Jp is not None:
assert P.handle == ctx._pjac.petscmat.handle
ctx._assemble_pjac()
ises = problem.J.arguments()[0].function_space()._ises
ctx.set_nullspace(ctx._nullspace, ises, transpose=False, near=False)
ctx.set_nullspace(ctx._nullspace_T, ises, transpose=True, near=False)
ctx.set_nullspace(ctx._near_nullspace,
ises,
transpose=False,
near=True)
def new_snes_ctx(pc, op, bcs, mat_type, fcp=None):
""" Create a new SNES contex for nested preconditioning
"""
from firedrake.variational_solver import NonlinearVariationalProblem
from firedrake.function import Function
from firedrake.ufl_expr import action
from firedrake.dmhooks import get_appctx
from firedrake.solving_utils import _SNESContext
dm = pc.getDM()
old_appctx = get_appctx(dm).appctx
u = Function(op.arguments()[-1].function_space())
F = action(op, u)
nprob = NonlinearVariationalProblem(F,
u,
bcs=bcs,
J=op,
form_compiler_parameters=fcp)
nctx = _SNESContext(nprob, mat_type, mat_type, old_appctx)
return nctx
def initialize(self, pc):
A, P = pc.getOperators()
if P.getType() == "python":
from firedrake.matrix_free.operators import ImplicitMatrixContext
ctx = P.getPythonContext()
if ctx is None:
raise ValueError("No context found on matrix")
if not isinstance(ctx, ImplicitMatrixContext):
raise ValueError("Don't know how to get form from %r", ctx)
J = ctx.a
bcs = ctx.row_bcs
if bcs != ctx.col_bcs:
raise NotImplementedError("Row and column bcs must match")
else:
from firedrake.dmhooks import get_appctx
from firedrake.solving_utils import _SNESContext
ctx = get_appctx(pc.getDM())
if ctx is None:
raise ValueError("No context found on form")
if not isinstance(ctx, _SNESContext):
raise ValueError("Don't know how to get form from %r", ctx)
J = ctx.Jp or ctx.J
bcs = ctx._problem.bcs
mesh = J.ufl_domain()
if mesh.cell_set._extruded:
raise NotImplementedError("Not implemented on extruded meshes")
combiner = PETSc.PC().create(comm=pc.comm)
combiner.setOptionsPrefix(pc.getOptionsPrefix() + "ssc_")
combiner.setOperators(A, P)
combiner.setType(PETSc.PC.Type.COMPOSITE)
combiner.addCompositePC("patch")
#combiner.addCompositePC("python")
#divprolong = combiner.getCompositePC(0)
#divprolong.setPythonContext(DivProlong(J, bcs))
patch = combiner.getCompositePC(0)
patch = setup_patch_pc(patch, J, bcs)
self.combiner = combiner
self.combiner.setFromOptions()
def initialize(self, pc):
A, P = pc.getOperators()
if P.getType() == "python":
from firedrake.matrix_free.operators import ImplicitMatrixContext
ctx = P.getPythonContext()
if ctx is None:
raise ValueError("No context found on matrix")
if not isinstance(ctx, ImplicitMatrixContext):
raise ValueError("Don't know how to get form from %r", ctx)
J = ctx.a
bcs = ctx.row_bcs
if bcs != ctx.col_bcs:
raise NotImplementedError("Row and column bcs must match")
else:
from firedrake.dmhooks import get_appctx
from firedrake.solving_utils import _SNESContext
ctx = get_appctx(pc.getDM())
if ctx is None:
raise ValueError("No context found on form")
if not isinstance(ctx, _SNESContext):
raise ValueError("Don't know how to get form from %r", ctx)
J = ctx.Jp or ctx.J
bcs = ctx._problem.bcs
mesh = J.ufl_domain()
if mesh.cell_set._extruded:
raise NotImplementedError("Not implemented on extruded meshes")
patch = PETSc.PC().create(comm=pc.comm)
patch.setOptionsPrefix(pc.getOptionsPrefix() + "patch_")
patch.setOperators(A, P)
patch.setType("patch")
patch = setup_patch_pc(patch, J, bcs)
patch.setAttr("ctx", ctx)
patch.incrementTabLevel(1, parent=pc)
patch.setFromOptions()
self.patch = patch
def initialize(self, obj):
if isinstance(obj, PETSc.PC):
A, P = obj.getOperators()
elif isinstance(obj, PETSc.SNES):
A, P = obj.ksp.pc.getOperators()
else:
raise ValueError("Not a PC or SNES?")
ctx = get_appctx(obj.getDM())
if ctx is None:
raise ValueError("No context found on form")
if not isinstance(ctx, _SNESContext):
raise ValueError("Don't know how to get form from %r", ctx)
if P.getType() == "python":
ictx = P.getPythonContext()
if ictx is None:
raise ValueError("No context found on matrix")
if not isinstance(ictx, ImplicitMatrixContext):
raise ValueError("Don't know how to get form from %r", ictx)
J = ictx.a
bcs = ictx.row_bcs
if bcs != ictx.col_bcs:
raise NotImplementedError("Row and column bcs must match")
else:
J = ctx.Jp or ctx.J
bcs = ctx._problem.bcs
mesh = J.ufl_domain()
self.plex = mesh._plex
self.ctx = ctx
if mesh.cell_set._extruded:
raise NotImplementedError("Not implemented on extruded meshes")
if "overlap_type" not in mesh._distribution_parameters:
if mesh.comm.size > 1:
# Want to do
# warnings.warn("You almost surely want to set an overlap_type in your mesh's distribution_parameters.")
# but doesn't warn!
PETSc.Sys.Print(
"Warning: you almost surely want to set an overlap_type in your mesh's distribution_parameters."
)
patch = obj.__class__().create(comm=obj.comm)
patch.setOptionsPrefix(obj.getOptionsPrefix() + "patch_")
self.configure_patch(patch, obj)
patch.setType("patch")
if isinstance(obj, PETSc.SNES):
Jstate = ctx._problem.u
else:
Jstate = None
V, _ = map(operator.methodcaller("function_space"), J.arguments())
if len(bcs) > 0:
ghost_bc_nodes = numpy.unique(
numpy.concatenate([bcdofs(bc, ghost=True) for bc in bcs]))
global_bc_nodes = numpy.unique(
numpy.concatenate([bcdofs(bc, ghost=False) for bc in bcs]))
else:
ghost_bc_nodes = numpy.empty(0, dtype=PETSc.IntType)
global_bc_nodes = numpy.empty(0, dtype=PETSc.IntType)
Jfunptr, Jkinfo = matrix_funptr(J, Jstate)
Jop_coeffs = [mesh.coordinates]
for n in Jkinfo.coefficient_map:
Jop_coeffs.append(J.coefficients()[n])
Jop_args = []
Jop_state_slot = None
for c in Jop_coeffs:
if c is Jstate:
Jop_state_slot = len(Jop_args)
Jop_args.append(None)
Jop_args.append(None)
continue
for c_ in c.split():
Jop_args.append(c_.dat._data.ctypes.data)
c_map = c_.cell_node_map()
if c_map is not None:
Jop_args.append(c_map._values.ctypes.data)
def Jop(obj, point, vec, mat, cellIS, cell_dofmap, cell_dofmapWithAll):
cells = cellIS.indices
ncell = len(cells)
dofs = cell_dofmap.ctypes.data
if cell_dofmapWithAll is not None:
dofsWithAll = cell_dofmapWithAll.ctypes.data
else:
dofsWithAll = None
mat.zeroEntries()
if Jop_state_slot is not None:
assert dofsWithAll is not None
Jop_args[Jop_state_slot] = vec.array_r.ctypes.data
Jop_args[Jop_state_slot + 1] = dofsWithAll
Jfunptr(0, ncell, cells.ctypes.data, mat.handle, dofs, dofs,
*Jop_args)
mat.assemble()
patch.setPatchComputeOperator(Jop)
if hasattr(ctx, "F") and isinstance(obj, PETSc.SNES):
F = ctx.F
Fstate = ctx._problem.u
Ffunptr, Fkinfo = residual_funptr(F, Fstate)
Fop_coeffs = [mesh.coordinates]
for n in Fkinfo.coefficient_map:
Fop_coeffs.append(F.coefficients()[n])
assert any(
c is Fstate for c in
Fop_coeffs), "Couldn't find state vector in F.coefficients()"
Fop_args = []
Fop_state_slot = None
for c in Fop_coeffs:
if c is Fstate:
Fop_state_slot = len(Fop_args)
Fop_args.append(None)
Fop_args.append(None)
continue
for c_ in c.split():
Fop_args.append(c_.dat._data.ctypes.data)
c_map = c_.cell_node_map()
if c_map is not None:
Fop_args.append(c_map._values.ctypes.data)
assert Fop_state_slot is not None
def Fop(pc, point, vec, out, cellIS, cell_dofmap,
cell_dofmapWithAll):
cells = cellIS.indices
ncell = len(cells)
dofs = cell_dofmap.ctypes.data
dofsWithAll = cell_dofmapWithAll.ctypes.data
out.set(0)
outdata = out.array
Fop_args[Fop_state_slot] = vec.array_r.ctypes.data
Fop_args[Fop_state_slot + 1] = dofsWithAll
Ffunptr(0, ncell, cells.ctypes.data, outdata.ctypes.data, dofs,
*Fop_args)
patch.setPatchComputeFunction(Fop)
patch.setDM(self.plex)
patch.setPatchCellNumbering(mesh._cell_numbering)
offsets = numpy.append([0], numpy.cumsum([W.dof_count for W in V
])).astype(PETSc.IntType)
patch.setPatchDiscretisationInfo([W.dm for W in V],
numpy.array([W.value_size for W in V],
dtype=PETSc.IntType),
[W.cell_node_list
for W in V], offsets, ghost_bc_nodes,
global_bc_nodes)
patch.setPatchConstructType(PETSc.PC.PatchConstructType.PYTHON,
operator=self.user_construction_op)
patch.setAttr("ctx", ctx)
patch.incrementTabLevel(1, parent=obj)
patch.setFromOptions()
patch.setUp()
self.patch = patch
def initialize(self, obj):
if isinstance(obj, PETSc.PC):
A, P = obj.getOperators()
elif isinstance(obj, PETSc.SNES):
A, P = obj.ksp.pc.getOperators()
else:
raise ValueError("Not a PC or SNES?")
ctx = get_appctx(obj.getDM())
if ctx is None:
raise ValueError("No context found on form")
if not isinstance(ctx, _SNESContext):
raise ValueError("Don't know how to get form from %r", ctx)
if P.getType() == "python":
ictx = P.getPythonContext()
if ictx is None:
raise ValueError("No context found on matrix")
if not isinstance(ictx, ImplicitMatrixContext):
raise ValueError("Don't know how to get form from %r", ictx)
J = ictx.a
bcs = ictx.row_bcs
if bcs != ictx.col_bcs:
raise NotImplementedError("Row and column bcs must match")
else:
J = ctx.Jp or ctx.J
bcs = ctx._problem.bcs
mesh = J.ufl_domain()
self.plex = mesh._topology_dm
# We need to attach the mesh to the plex, so that
# PlaneSmoothers (and any other user-customised patch
# constructors) can use firedrake's opinion of what
# the coordinates are, rather than plex's.
self.plex.setAttr("__firedrake_mesh__", weakref.proxy(mesh))
self.ctx = ctx
if mesh.cell_set._extruded:
raise NotImplementedError("Not implemented on extruded meshes")
if "overlap_type" not in mesh._distribution_parameters:
if mesh.comm.size > 1:
# Want to do
# warnings.warn("You almost surely want to set an overlap_type in your mesh's distribution_parameters.")
# but doesn't warn!
PETSc.Sys.Print(
"Warning: you almost surely want to set an overlap_type in your mesh's distribution_parameters."
)
patch = obj.__class__().create(comm=obj.comm)
patch.setOptionsPrefix(obj.getOptionsPrefix() + "patch_")
self.configure_patch(patch, obj)
patch.setType("patch")
if isinstance(obj, PETSc.SNES):
Jstate = ctx._problem.u
is_snes = True
else:
Jstate = None
is_snes = False
V, _ = map(operator.methodcaller("function_space"), J.arguments())
if len(bcs) > 0:
ghost_bc_nodes = numpy.unique(
numpy.concatenate([bcdofs(bc, ghost=True) for bc in bcs]))
global_bc_nodes = numpy.unique(
numpy.concatenate([bcdofs(bc, ghost=False) for bc in bcs]))
else:
ghost_bc_nodes = numpy.empty(0, dtype=PETSc.IntType)
global_bc_nodes = numpy.empty(0, dtype=PETSc.IntType)
Jcell_kernels, Jint_facet_kernels = matrix_funptr(J, Jstate)
Jcell_kernel, = Jcell_kernels
Jop_data_args, Jop_map_args = make_c_arguments(
J, Jcell_kernel, Jstate, operator.methodcaller("cell_node_map"))
code, Struct = make_jacobian_wrapper(Jop_data_args, Jop_map_args)
Jop_function = load_c_function(code, "ComputeJacobian", obj.comm)
Jop_struct = make_c_struct(Jop_data_args, Jop_map_args,
Jcell_kernel.funptr, Struct)
Jhas_int_facet_kernel = False
if len(Jint_facet_kernels) > 0:
Jint_facet_kernel, = Jint_facet_kernels
Jhas_int_facet_kernel = True
facet_Jop_data_args, facet_Jop_map_args = make_c_arguments(
J,
Jint_facet_kernel,
Jstate,
operator.methodcaller("interior_facet_node_map"),
require_facet_number=True)
code, Struct = make_jacobian_wrapper(facet_Jop_data_args,
facet_Jop_map_args)
facet_Jop_function = load_c_function(code, "ComputeJacobian",
obj.comm)
point2facet = J.ufl_domain(
).interior_facets.point2facetnumber.ctypes.data
facet_Jop_struct = make_c_struct(facet_Jop_data_args,
facet_Jop_map_args,
Jint_facet_kernel.funptr,
Struct,
point2facet=point2facet)
set_residual = hasattr(ctx, "F") and isinstance(obj, PETSc.SNES)
if set_residual:
F = ctx.F
Fstate = ctx._problem.u
Fcell_kernels, Fint_facet_kernels = residual_funptr(F, Fstate)
Fcell_kernel, = Fcell_kernels
Fop_data_args, Fop_map_args = make_c_arguments(
F,
Fcell_kernel,
Fstate,
operator.methodcaller("cell_node_map"),
require_state=True)
code, Struct = make_residual_wrapper(Fop_data_args, Fop_map_args)
Fop_function = load_c_function(code, "ComputeResidual", obj.comm)
Fop_struct = make_c_struct(Fop_data_args, Fop_map_args,
Fcell_kernel.funptr, Struct)
Fhas_int_facet_kernel = False
if len(Fint_facet_kernels) > 0:
Fint_facet_kernel, = Fint_facet_kernels
Fhas_int_facet_kernel = True
facet_Fop_data_args, facet_Fop_map_args = make_c_arguments(
F,
Fint_facet_kernel,
Fstate,
operator.methodcaller("interior_facet_node_map"),
require_state=True,
require_facet_number=True)
code, Struct = make_jacobian_wrapper(facet_Fop_data_args,
facet_Fop_map_args)
facet_Fop_function = load_c_function(code, "ComputeResidual",
obj.comm)
point2facet = F.ufl_domain(
).interior_facets.point2facetnumber.ctypes.data
facet_Fop_struct = make_c_struct(facet_Fop_data_args,
facet_Fop_map_args,
Fint_facet_kernel.funptr,
Struct,
point2facet=point2facet)
patch.setDM(self.plex)
patch.setPatchCellNumbering(mesh._cell_numbering)
offsets = numpy.append([0], numpy.cumsum([W.dof_count for W in V
])).astype(PETSc.IntType)
patch.setPatchDiscretisationInfo([W.dm for W in V],
numpy.array([W.value_size for W in V],
dtype=PETSc.IntType),
[W.cell_node_list
for W in V], offsets, ghost_bc_nodes,
global_bc_nodes)
self.Jop_struct = Jop_struct
set_patch_jacobian(patch,
ctypes.cast(Jop_function, ctypes.c_voidp).value,
ctypes.addressof(Jop_struct),
is_snes=is_snes)
if Jhas_int_facet_kernel:
self.facet_Jop_struct = facet_Jop_struct
set_patch_jacobian(patch,
ctypes.cast(facet_Jop_function,
ctypes.c_voidp).value,
ctypes.addressof(facet_Jop_struct),
is_snes=is_snes,
interior_facets=True)
if set_residual:
self.Fop_struct = Fop_struct
set_patch_residual(patch,
ctypes.cast(Fop_function, ctypes.c_voidp).value,
ctypes.addressof(Fop_struct),
is_snes=is_snes)
if Fhas_int_facet_kernel:
set_patch_residual(patch,
ctypes.cast(facet_Fop_function,
ctypes.c_voidp).value,
ctypes.addressof(facet_Fop_struct),
is_snes=is_snes,
interior_facets=True)
patch.setPatchConstructType(PETSc.PC.PatchConstructType.PYTHON,
operator=self.user_construction_op)
patch.setAttr("ctx", ctx)
patch.incrementTabLevel(1, parent=obj)
patch.setFromOptions()
patch.setUp()
self.patch = patch
def initialize(self, pc):
from firedrake.assemble import allocate_matrix, create_assembly_callable
_, P = pc.getOperators()
if pc.getType() != "python":
raise ValueError("Expecting PC type python")
opc = pc
appctx = self.get_appctx(pc)
fcp = appctx.get("form_compiler_parameters")
V = get_function_space(pc.getDM())
if len(V) == 1:
V = FunctionSpace(V.mesh(), V.ufl_element())
else:
V = MixedFunctionSpace([V_ for V_ in V])
test = TestFunction(V)
trial = TrialFunction(V)
if P.type == "python":
context = P.getPythonContext()
# It only makes sense to preconditioner/invert a diagonal
# block in general. That's all we're going to allow.
if not context.on_diag:
raise ValueError("Only makes sense to invert diagonal block")
prefix = pc.getOptionsPrefix()
options_prefix = prefix + self._prefix
mat_type = PETSc.Options().getString(options_prefix + "mat_type", "aij")
(a, bcs) = self.form(pc, test, trial)
self.P = allocate_matrix(a, bcs=bcs,
form_compiler_parameters=fcp,
mat_type=mat_type,
options_prefix=options_prefix)
self._assemble_P = create_assembly_callable(a, tensor=self.P,
bcs=bcs,
form_compiler_parameters=fcp,
mat_type=mat_type)
self._assemble_P()
self.P.force_evaluation()
# Transfer nullspace over
Pmat = self.P.petscmat
Pmat.setNullSpace(P.getNullSpace())
tnullsp = P.getTransposeNullSpace()
if tnullsp.handle != 0:
Pmat.setTransposeNullSpace(tnullsp)
# Internally, we just set up a PC object that the user can configure
# however from the PETSc command line. Since PC allows the user to specify
# a KSP, we can do iterative by -assembled_pc_type ksp.
pc = PETSc.PC().create(comm=opc.comm)
pc.incrementTabLevel(1, parent=opc)
# We set a DM and an appropriate SNESContext on the constructed PC so one
# can do e.g. multigrid or patch solves.
from firedrake.variational_solver import NonlinearVariationalProblem
from firedrake.solving_utils import _SNESContext
dm = opc.getDM()
octx = get_appctx(dm)
oproblem = octx._problem
nproblem = NonlinearVariationalProblem(oproblem.F, oproblem.u, bcs, J=a, form_compiler_parameters=fcp)
nctx = _SNESContext(nproblem, mat_type, mat_type, octx.appctx)
push_appctx(dm, nctx)
self._ctx_ref = nctx
pc.setDM(dm)
pc.setOptionsPrefix(options_prefix)
pc.setOperators(Pmat, Pmat)
pc.setFromOptions()
pc.setUp()
self.pc = pc
pop_appctx(dm)
def initialize(self, obj):
if isinstance(obj, PETSc.PC):
A, P = obj.getOperators()
elif isinstance(obj, PETSc.SNES):
A, P = obj.ksp.pc.getOperators()
else:
raise ValueError("Not a PC or SNES?")
ctx = get_appctx(obj.getDM())
if ctx is None:
raise ValueError("No context found on form")
if not isinstance(ctx, _SNESContext):
raise ValueError("Don't know how to get form from %r", ctx)
if P.getType() == "python":
ictx = P.getPythonContext()
if ictx is None:
raise ValueError("No context found on matrix")
if not isinstance(ictx, ImplicitMatrixContext):
raise ValueError("Don't know how to get form from %r", ictx)
J = ictx.a
bcs = ictx.row_bcs
if bcs != ictx.col_bcs:
raise NotImplementedError("Row and column bcs must match")
else:
J = ctx.Jp or ctx.J
bcs = ctx._problem.bcs
mesh = J.ufl_domain()
self.plex = mesh._plex
self.ctx = ctx
if mesh.cell_set._extruded:
raise NotImplementedError("Not implemented on extruded meshes")
if "overlap_type" not in mesh._distribution_parameters:
if mesh.comm.size > 1:
# Want to do
# warnings.warn("You almost surely want to set an overlap_type in your mesh's distribution_parameters.")
# but doesn't warn!
PETSc.Sys.Print("Warning: you almost surely want to set an overlap_type in your mesh's distribution_parameters.")
patch = obj.__class__().create(comm=obj.comm)
patch.setOptionsPrefix(obj.getOptionsPrefix() + "patch_")
self.configure_patch(patch, obj)
patch.setType("patch")
if isinstance(obj, PETSc.SNES):
Jstate = ctx._problem.u
else:
Jstate = None
V, _ = map(operator.methodcaller("function_space"), J.arguments())
if len(bcs) > 0:
ghost_bc_nodes = numpy.unique(numpy.concatenate([bcdofs(bc, ghost=True)
for bc in bcs]))
global_bc_nodes = numpy.unique(numpy.concatenate([bcdofs(bc, ghost=False)
for bc in bcs]))
else:
ghost_bc_nodes = numpy.empty(0, dtype=PETSc.IntType)
global_bc_nodes = numpy.empty(0, dtype=PETSc.IntType)
Jcell_kernels, Jint_facet_kernels = matrix_funptr(J, Jstate)
Jop_coeffs = [mesh.coordinates]
Jcell_kernel, = Jcell_kernels
if Jcell_kernel.kinfo.oriented:
Jop_coeffs.append(J.ufl_domain().cell_orientations())
for n in Jcell_kernel.kinfo.coefficient_map:
Jop_coeffs.append(J.coefficients()[n])
Jop_data_args = []
Jop_map_args = []
seen = set()
Jop_state_data_slot = None
Jop_state_map_slot = None
for c in Jop_coeffs:
if c is Jstate:
Jop_state_data_slot = len(Jop_data_args)
Jop_state_map_slot = len(Jop_map_args)
Jop_data_args.append(None)
Jop_map_args.append(None)
continue
Jop_data_args.extend(c.dat._kernel_args_)
map_ = c.cell_node_map()
if map_ is not None:
for k in map_._kernel_args_:
if k in seen:
continue
Jop_map_args.append(k)
seen.add(k)
def Jop(obj, point, vec, mat, cellIS, cell_dofmap, cell_dofmapWithAll):
cells = cellIS.indices
ncell = len(cells)
dofs = cell_dofmap.ctypes.data
if cell_dofmapWithAll is not None:
dofsWithAll = cell_dofmapWithAll.ctypes.data
else:
dofsWithAll = None
if Jop_state_data_slot is not None:
assert dofsWithAll is not None
Jop_data_args[Jop_state_data_slot] = vec.array_r.ctypes.data
Jop_map_args[Jop_state_map_slot] = dofsWithAll
Jcell_kernel.funptr(0, ncell, cells.ctypes.data, mat.handle, *Jop_data_args,
dofs, *Jop_map_args)
Jhas_int_facet_kernel = False
if len(Jint_facet_kernels) > 0:
Jint_facet_kernel, = Jint_facet_kernels
Jhas_int_facet_kernel = True
facet_Jop_coeffs = [mesh.coordinates]
if Jint_facet_kernel.kinfo.oriented:
facet_Jop_coeffs.append(J.ufl_domain().cell_orientations())
for n in Jint_facet_kernel.kinfo.coefficient_map:
facet_Jop_coeffs.append(J.coefficients()[n])
facet_Jop_data_args = []
facet_Jop_map_args = []
facet_Jop_state_data_slot = None
facet_Jop_state_map_slot = None
seen = set()
for c in facet_Jop_coeffs:
if c is Jstate:
facet_Jop_state_data_slot = len(facet_Jop_data_args)
facet_Jop_state_map_slot = len(facet_Jop_map_args)
facet_Jop_data_args.append(None)
facet_Jop_map_args.append(None)
continue
facet_Jop_data_args.extend(c.dat._kernel_args_)
map_ = c.interior_facet_node_map()
if map_ is not None:
for k in map_._kernel_args_:
if k in seen:
continue
facet_Jop_map_args.append(k)
seen.add(k)
facet_Jop_data_args.extend(J.ufl_domain().interior_facets.local_facet_dat._kernel_args_)
point2facetnumber = J.ufl_domain().interior_facets.point2facetnumber
def Jfacet_op(pc, point, vec, mat, facetIS, facet_dofmap, facet_dofmapWithAll):
facets = numpy.asarray(list(map(point2facetnumber.__getitem__, facetIS.indices)),
dtype=IntType)
nfacet = len(facets)
dofs = facet_dofmap.ctypes.data
if facet_Jop_state_data_slot is not None:
assert facet_dofmapWithAll is not None
facet_Jop_data_args[facet_Jop_state_data_slot] = vec.array_r.ctypes.data
facet_Jop_map_args[facet_Jop_state_map_slot] = facet_dofmapWithAll.ctypes.data
Jint_facet_kernel.funptr(0, nfacet, facets.ctypes.data, mat.handle, *facet_Jop_data_args,
dofs, *facet_Jop_map_args)
set_residual = hasattr(ctx, "F") and isinstance(obj, PETSc.SNES)
if set_residual:
F = ctx.F
Fstate = ctx._problem.u
Fcell_kernels, Fint_facet_kernels = residual_funptr(F, Fstate)
Fop_coeffs = [mesh.coordinates]
Fcell_kernel, = Fcell_kernels
if Fcell_kernel.kinfo.oriented:
Fop_coeffs.append(F.ufl_domain().cell_orientations())
for n in Fcell_kernel.kinfo.coefficient_map:
Fop_coeffs.append(F.coefficients()[n])
assert any(c is Fstate for c in Fop_coeffs), "Couldn't find state vector in F.coefficients()"
Fop_data_args = []
Fop_map_args = []
seen = set()
Fop_state_data_slot = None
Fop_state_map_slot = None
for c in Fop_coeffs:
if c is Fstate:
Fop_state_data_slot = len(Fop_data_args)
Fop_state_map_slot = len(Fop_map_args)
Fop_data_args.append(None)
Fop_map_args.append(None)
continue
Fop_data_args.extend(c.dat._kernel_args_)
map_ = c.cell_node_map()
if map_ is not None:
for k in map_._kernel_args_:
if k in seen:
continue
Fop_map_args.append(k)
seen.add(k)
assert Fop_state_data_slot is not None
def Fop(pc, point, vec, out, cellIS, cell_dofmap, cell_dofmapWithAll):
cells = cellIS.indices
ncell = len(cells)
dofs = cell_dofmap.ctypes.data
dofsWithAll = cell_dofmapWithAll.ctypes.data
out.set(0)
outdata = out.array
Fop_data_args[Fop_state_data_slot] = vec.array_r.ctypes.data
Fop_map_args[Fop_state_map_slot] = dofsWithAll
Fcell_kernel.funptr(0, ncell, cells.ctypes.data, outdata.ctypes.data,
*Fop_data_args, dofs, *Fop_map_args)
Fhas_int_facet_kernel = False
if len(Fint_facet_kernels) > 0:
Fint_facet_kernel, = Fint_facet_kernels
Fhas_int_facet_kernel = True
facet_Fop_coeffs = [mesh.coordinates]
if Fint_facet_kernel.kinfo.oriented:
facet_Fop_coeffs.append(F.ufl_domain().cell_orientations())
for n in Fint_facet_kernel.kinfo.coefficient_map:
facet_Fop_coeffs.append(J.coefficients()[n])
facet_Fop_data_args = []
facet_Fop_map_args = []
facet_Fop_state_data_slot = None
facet_Fop_state_map_slot = None
seen = set()
for c in facet_Fop_coeffs:
if c is Fstate:
facet_Fop_state_data_slot = len(Fop_data_args)
facet_Fop_state_map_slot = len(Fop_map_args)
facet_Fop_data_args.append(None)
facet_Fop_map_args.append(None)
continue
facet_Fop_data_args.extend(c.dat._kernel_args_)
map_ = c.interior_facet_node_map()
if map_ is not None:
for k in map_._kernel_args_:
if k in seen:
continue
facet_Fop_map_args.append(k)
seen.add(k)
facet_Fop_data_args.extend(F.ufl_domain().interior_facets.local_facet_dat._kernel_args_)
point2facetnumber = F.ufl_domain().interior_facets.point2facetnumber
def Ffacet_op(pc, point, vec, mat, facetIS, facet_dofmap, facet_dofmapWithAll):
facets = numpy.asarray(list(map(point2facetnumber.__getitem__, facetIS.indices)),
dtype=IntType)
nfacet = len(facets)
dofs = facet_dofmap.ctypes.data
if facet_Fop_state_data_slot is not None:
assert facet_dofmapWithAll is not None
facet_Fop_data_args[facet_Fop_state_data_slot] = vec.array_r.ctypes.data
facet_Fop_map_args[facet_Fop_state_map_slot] = facet_dofmapWithAll.ctypes.data
Fint_facet_kernel.funptr(0, nfacet, facets.ctypes.data, mat.handle, *facet_Fop_data_args,
dofs, *facet_Fop_map_args)
patch.setDM(self.plex)
patch.setPatchCellNumbering(mesh._cell_numbering)
offsets = numpy.append([0], numpy.cumsum([W.dof_count
for W in V])).astype(PETSc.IntType)
patch.setPatchDiscretisationInfo([W.dm for W in V],
numpy.array([W.value_size for
W in V], dtype=PETSc.IntType),
[W.cell_node_list for W in V],
offsets,
ghost_bc_nodes,
global_bc_nodes)
patch.setPatchComputeOperator(Jop)
if Jhas_int_facet_kernel:
patch.setPatchComputeOperatorInteriorFacets(Jfacet_op)
if set_residual:
patch.setPatchComputeFunction(Fop)
if Fhas_int_facet_kernel:
patch.setPatchComputeFunctionInteriorFacets(Ffacet_op)
patch.setPatchConstructType(PETSc.PC.PatchConstructType.PYTHON, operator=self.user_construction_op)
patch.setAttr("ctx", ctx)
patch.incrementTabLevel(1, parent=obj)
patch.setFromOptions()
patch.setUp()
self.patch = patch
def get_appctx(pc):
from firedrake.dmhooks import get_appctx
return get_appctx(pc.getDM()).appctx
def bform(self, rhs):
a = get_appctx(rhs.function_space().dm).J
return action(a, rhs)
def form(self, V):
a = get_appctx(V.dm).J
return a
def initialize(self, pc):
from firedrake import TestFunction, parameters
from firedrake.assemble import allocate_matrix, create_assembly_callable
from firedrake.interpolation import Interpolator
from firedrake.solving_utils import _SNESContext
from firedrake.matrix_free.operators import ImplicitMatrixContext
_, P = pc.getOperators()
appctx = self.get_appctx(pc)
fcp = appctx.get("form_compiler_parameters")
if pc.getType() != "python":
raise ValueError("Expecting PC type python")
ctx = dmhooks.get_appctx(pc.getDM())
if ctx is None:
raise ValueError("No context found.")
if not isinstance(ctx, _SNESContext):
raise ValueError("Don't know how to get form from %r", ctx)
prefix = pc.getOptionsPrefix()
options_prefix = prefix + self._prefix
opts = PETSc.Options()
# Handle the fine operator if type is python
if P.getType() == "python":
ictx = P.getPythonContext()
if ictx is None:
raise ValueError("No context found on matrix")
if not isinstance(ictx, ImplicitMatrixContext):
raise ValueError("Don't know how to get form from %r", ictx)
fine_operator = ictx.a
fine_bcs = ictx.row_bcs
if fine_bcs != ictx.col_bcs:
raise NotImplementedError("Row and column bcs must match")
fine_mat_type = opts.getString(options_prefix + "mat_type",
parameters["default_matrix_type"])
self.fine_op = allocate_matrix(fine_operator,
bcs=fine_bcs,
form_compiler_parameters=fcp,
mat_type=fine_mat_type,
options_prefix=options_prefix)
self._assemble_fine_op = create_assembly_callable(
fine_operator,
tensor=self.fine_op,
bcs=fine_bcs,
form_compiler_parameters=fcp,
mat_type=fine_mat_type)
self._assemble_fine_op()
fine_petscmat = self.fine_op.petscmat
else:
fine_petscmat = P
# Transfer fine operator null space
fine_petscmat.setNullSpace(P.getNullSpace())
fine_transpose_nullspace = P.getTransposeNullSpace()
if fine_transpose_nullspace.handle != 0:
fine_petscmat.setTransposeNullSpace(fine_transpose_nullspace)
# Handle the coarse operator
coarse_options_prefix = options_prefix + "mg_coarse"
coarse_mat_type = opts.getString(coarse_options_prefix + "mat_type",
parameters["default_matrix_type"])
get_coarse_space = appctx.get("get_coarse_space", None)
if not get_coarse_space:
raise ValueError(
"Need to provide a callback which provides the coarse space.")
coarse_space = get_coarse_space()
get_coarse_operator = appctx.get("get_coarse_operator", None)
if not get_coarse_operator:
raise ValueError(
"Need to provide a callback which provides the coarse operator."
)
coarse_operator = get_coarse_operator()
coarse_space_bcs = appctx.get("coarse_space_bcs", None)
# These should be callbacks which return the relevant nullspaces
get_coarse_nullspace = appctx.get("get_coarse_op_nullspace", None)
get_coarse_transpose_nullspace = appctx.get(
"get_coarse_op_transpose_nullspace", None)
self.coarse_op = allocate_matrix(coarse_operator,
bcs=coarse_space_bcs,
form_compiler_parameters=fcp,
mat_type=coarse_mat_type,
options_prefix=coarse_options_prefix)
self._assemble_coarse_op = create_assembly_callable(
coarse_operator,
tensor=self.coarse_op,
bcs=coarse_space_bcs,
form_compiler_parameters=fcp)
self._assemble_coarse_op()
coarse_opmat = self.coarse_op.petscmat
# Set nullspace if provided
if get_coarse_nullspace:
nsp = get_coarse_nullspace()
coarse_opmat.setNullSpace(nsp.nullspace(comm=pc.comm))
if get_coarse_transpose_nullspace:
tnsp = get_coarse_transpose_nullspace()
coarse_opmat.setTransposeNullSpace(tnsp.nullspace(comm=pc.comm))
interp_petscmat = appctx.get("interpolation_matrix", None)
if interp_petscmat is None:
# Create interpolation matrix from coarse space to fine space
fine_space = ctx.J.arguments()[0].function_space()
interpolator = Interpolator(TestFunction(coarse_space), fine_space)
interpolation_matrix = interpolator.callable()
interp_petscmat = interpolation_matrix.handle
# We set up a PCMG object that uses the constructed interpolation
# matrix to generate the restriction/prolongation operators.
# This is a two-level multigrid preconditioner.
pcmg = PETSc.PC().create(comm=pc.comm)
pcmg.incrementTabLevel(1, parent=pc)
pcmg.setType(pc.Type.MG)
pcmg.setOptionsPrefix(options_prefix)
pcmg.setMGLevels(2)
pcmg.setMGCycleType(pc.MGCycleType.V)
pcmg.setMGInterpolation(1, interp_petscmat)
pcmg.setOperators(A=fine_petscmat, P=fine_petscmat)
# Create new appctx
self._ctx_ref = self.new_snes_ctx(pc, coarse_operator,
coarse_space_bcs, coarse_mat_type,
fcp)
coarse_solver = pcmg.getMGCoarseSolve()
coarse_solver.setOperators(A=coarse_opmat, P=coarse_opmat)
# coarse space dm
coarse_dm = coarse_space.dm
coarse_solver.setDM(coarse_dm)
coarse_solver.setDMActive(False)
pcmg.setDM(coarse_dm)
pcmg.setFromOptions()
self.pc = pcmg
self._dm = coarse_dm
with dmhooks.add_hooks(coarse_dm,
self,
appctx=self._ctx_ref,
save=False):
coarse_solver.setFromOptions()
def initialize(self, pc):
# Make a new DM.
# Hook up a (new) coarsen routine on that DM.
# Make a new PC, of type MG.
# Assign the DM to that PC.
odm = pc.getDM()
ctx = get_appctx(odm)
test, trial = ctx.J.arguments()
if test.function_space() != trial.function_space():
raise NotImplementedError("test and trial spaces must be the same")
# Construct a list with the elements we'll be using
V = test.function_space()
ele = V.ufl_element()
elements = [ele]
while True:
try:
ele_ = self.coarsen_element(ele)
assert ele_.value_shape() == ele.value_shape()
ele = ele_
except ValueError:
break
elements.append(ele)
pdm = PETSc.DMShell().create(comm=pc.comm)
sf = odm.getPointSF()
section = odm.getDefaultSection()
attach_hooks(pdm, level=len(elements) - 1, sf=sf, section=section)
# Now overwrite some routines on the DM
pdm.setRefine(None)
pdm.setCoarsen(self.coarsen)
pdm.setCreateInterpolation(self.create_interpolation)
pdm.setOptionsPrefix(pc.getOptionsPrefix() + "pmg_")
set_function_space(pdm, get_function_space(odm))
parent = get_parent(odm)
assert parent is not None
add_hook(parent,
setup=partial(push_parent, pdm, parent),
teardown=partial(pop_parent, pdm, parent),
call_setup=True)
add_hook(parent,
setup=partial(push_appctx, pdm, ctx),
teardown=partial(pop_appctx, pdm, ctx),
call_setup=True)
ppc = PETSc.PC().create(comm=pc.comm)
ppc.setOptionsPrefix(pc.getOptionsPrefix() + "pmg_")
ppc.setType("mg")
ppc.setOperators(*pc.getOperators())
ppc.setDM(pdm)
ppc.incrementTabLevel(1, parent=pc)
# PETSc unfortunately requires us to make an ugly hack.
# We would like to use GMG for the coarse solve, at least
# sometimes. But PETSc will use this p-DM's getRefineLevels()
# instead of the getRefineLevels() of the MeshHierarchy to
# decide how many levels it should use for PCMG applied to
# the p-MG's coarse problem. So we need to set an option
# for the user, if they haven't already; I don't know any
# other way to get PETSc to know this at the right time.
opts = PETSc.Options(pc.getOptionsPrefix() + "pmg_")
if "mg_coarse_pc_mg_levels" not in opts:
opts["mg_coarse_pc_mg_levels"] = odm.getRefineLevel() + 1
ppc.setFromOptions()
ppc.setUp()
self.ppc = ppc
def coarsen(self, fdm, comm):
fctx = get_appctx(fdm)
test, trial = fctx.J.arguments()
fV = test.function_space()
fu = fctx._problem.u
cele = self.coarsen_element(fV.ufl_element())
cV = firedrake.FunctionSpace(fV.mesh(), cele)
cdm = cV.dm
cu = firedrake.Function(cV)
interpolators = tuple(
firedrake.Interpolator(fus, cus)
for fus, cus in zip(fu.split(), cu.split()))
def inject_state(interpolators):
for interpolator in interpolators:
interpolator.interpolate()
parent = get_parent(fdm)
assert parent is not None
add_hook(parent,
setup=partial(push_parent, cdm, parent),
teardown=partial(pop_parent, cdm, parent),
call_setup=True)
replace_d = {
fu: cu,
test: firedrake.TestFunction(cV),
trial: firedrake.TrialFunction(cV)
}
cJ = replace(fctx.J, replace_d)
cF = replace(fctx.F, replace_d)
if fctx.Jp is not None:
cJp = replace(fctx.Jp, replace_d)
else:
cJp = None
cbcs = []
for bc in fctx._problem.bcs:
# Don't actually need the value, since it's only used for
# killing parts of the matrix. This should be generalised
# for p-FAS, if anyone ever wants to do that
cV_ = cV
for index in bc._indices:
cV_ = cV_.sub(index)
cbcs.append(
firedrake.DirichletBC(cV_,
firedrake.zero(cV_.shape),
bc.sub_domain,
method=bc.method))
fcp = fctx._problem.form_compiler_parameters
cproblem = firedrake.NonlinearVariationalProblem(
cF,
cu,
cbcs,
cJ,
Jp=cJp,
form_compiler_parameters=fcp,
is_linear=fctx._problem.is_linear)
cctx = _SNESContext(
cproblem,
fctx.mat_type,
fctx.pmat_type,
appctx=fctx.appctx,
pre_jacobian_callback=fctx._pre_jacobian_callback,
pre_function_callback=fctx._pre_function_callback,
post_jacobian_callback=fctx._post_jacobian_callback,
post_function_callback=fctx._post_function_callback,
options_prefix=fctx.options_prefix,
transfer_manager=fctx.transfer_manager)
add_hook(parent,
setup=partial(push_appctx, cdm, cctx),
teardown=partial(pop_appctx, cdm, cctx),
call_setup=True)
add_hook(parent,
setup=partial(inject_state, interpolators),
call_setup=True)
cdm.setKSPComputeOperators(_SNESContext.compute_operators)
cdm.setCreateInterpolation(self.create_interpolation)
cdm.setOptionsPrefix(fdm.getOptionsPrefix())
# If we're the coarsest grid of the p-hierarchy, don't
# overwrite the coarsen routine; this is so that you can
# use geometric multigrid for the p-coarse problem
try:
self.coarsen_element(cele)
cdm.setCoarsen(self.coarsen)
except ValueError:
pass
return cdm
def initialize(self, pc):
from firedrake.assemble import allocate_matrix, create_assembly_callable
_, P = pc.getOperators()
if pc.getType() != "python":
raise ValueError("Expecting PC type python")
opc = pc
appctx = self.get_appctx(pc)
fcp = appctx.get("form_compiler_parameters")
V = get_function_space(pc.getDM())
if len(V) == 1:
V = FunctionSpace(V.mesh(), V.ufl_element())
else:
V = MixedFunctionSpace([V_ for V_ in V])
test = TestFunction(V)
trial = TrialFunction(V)
if P.type == "python":
context = P.getPythonContext()
# It only makes sense to preconditioner/invert a diagonal
# block in general. That's all we're going to allow.
if not context.on_diag:
raise ValueError("Only makes sense to invert diagonal block")
prefix = pc.getOptionsPrefix()
options_prefix = prefix + self._prefix
mat_type = PETSc.Options().getString(options_prefix + "mat_type", "aij")
(a, bcs) = self.form(pc, test, trial)
self.P = allocate_matrix(a, bcs=bcs,
form_compiler_parameters=fcp,
mat_type=mat_type,
options_prefix=options_prefix)
self._assemble_P = create_assembly_callable(a, tensor=self.P,
bcs=bcs,
form_compiler_parameters=fcp,
mat_type=mat_type)
self._assemble_P()
self.P.force_evaluation()
# Transfer nullspace over
Pmat = self.P.petscmat
Pmat.setNullSpace(P.getNullSpace())
tnullsp = P.getTransposeNullSpace()
if tnullsp.handle != 0:
Pmat.setTransposeNullSpace(tnullsp)
# Internally, we just set up a PC object that the user can configure
# however from the PETSc command line. Since PC allows the user to specify
# a KSP, we can do iterative by -assembled_pc_type ksp.
pc = PETSc.PC().create(comm=opc.comm)
pc.incrementTabLevel(1, parent=opc)
# We set a DM and an appropriate SNESContext on the constructed PC so one
# can do e.g. multigrid or patch solves.
from firedrake.variational_solver import NonlinearVariationalProblem
from firedrake.solving_utils import _SNESContext
dm = opc.getDM()
octx = get_appctx(dm)
oproblem = octx._problem
nproblem = NonlinearVariationalProblem(oproblem.F, oproblem.u, bcs, J=a, form_compiler_parameters=fcp)
nctx = _SNESContext(nproblem, mat_type, mat_type, octx.appctx)
push_appctx(dm, nctx)
pc.setDM(dm)
pc.setOptionsPrefix(options_prefix)
pc.setOperators(Pmat, Pmat)
pc.setFromOptions()
pc.setUp()
self.pc = pc
pop_appctx(dm)
def get_appctx(pc):
from firedrake.dmhooks import get_appctx
return get_appctx(pc.getDM()).appctx
