def step(self, snes, x, f, y):
push_appctx(self.plex, self.ctx)
x.copy(y)
self.patch.solve(snes.vec_rhs or self.dummy, y)
y.axpy(-1, x)
y.scale(-1)
snes.setConvergedReason(self.patch.getConvergedReason())
pop_appctx(self.plex)
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 step(self, snes, x, f, y):
push_appctx(self.plex, self.ctx)
x.copy(y)
self.patch.solve(snes.vec_rhs or self.dummy, y)
y.axpy(-1, x)
y.scale(-1)
snes.setConvergedReason(self.patch.getConvergedReason())
pop_appctx(self.plex)
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 applyTranspose(self, pc, x, y):
dm = pc.getDM()
push_appctx(dm, self._ctx_ref)
self.pc.applyTranspose(x, y)
pop_appctx(dm)
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 applyTranspose(self, pc, x, y):
dm = pc.getDM()
push_appctx(dm, self._ctx_ref)
self.pc.applyTranspose(x, y)
pop_appctx(dm)
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