def create_subdm(cls, dm, fields, *args, **kwargs):
W = dm.getAttr('__fs__')()
if len(fields) == 1:
# Subspace is just a single FunctionSpace.
field = fields[0]
subdm = W[field]._dm
iset = W._ises[field]
return iset, subdm
else:
try:
# Look up the subspace in the cache
iset, subspace = W._subspaces[tuple(fields)]
return iset, subspace._dm
except KeyError:
pass
# Need to build an MFS for the subspace
subspace = MixedFunctionSpace([W[f] for f in fields])
# Index set mapping from W into subspace.
iset = PETSc.IS().createGeneral(
numpy.concatenate([W._ises[f].indices for f in fields]))
# Keep hold of strong reference to created subspace (given we
# only hold a weakref in the shell DM), and so we can
# reuse it later.
W._subspaces[tuple(fields)] = iset, subspace
return iset, subspace._dm
def _split_arguments(self):
"""Splits the function space and stores the component
spaces determined by the indices.
"""
from firedrake.functionspace import FunctionSpace, MixedFunctionSpace
from firedrake.ufl_expr import Argument
tensor, = self.operands
nargs = []
for i, arg in enumerate(tensor.arguments()):
V = arg.function_space()
V_is = V.split()
idx = as_tuple(self._blocks[i])
if len(idx) == 1:
fidx, = idx
W = V_is[fidx]
W = FunctionSpace(W.mesh(), W.ufl_element())
else:
W = MixedFunctionSpace([V_is[fidx] for fidx in idx])
nargs.append(Argument(W, arg.number(), part=arg.part()))
return tuple(nargs)
def __mul__(self, other):
r"""Create a :class:`.MixedFunctionSpace` composed of this
:class:`.FunctionSpace` and other"""
from firedrake.functionspace import MixedFunctionSpace
return MixedFunctionSpace((self, other))
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)