def sc_solve(self, pc):
"""Solve the condensed linear system for the
condensed field.
:arg pc: a Preconditioner instance.
"""
dm = self.condensed_ksp.getDM()
with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref):
with self.condensed_rhs.dat.vec_ro as rhs:
if self.condensed_ksp.getInitialGuessNonzero():
acc = self.solution.split()[self.c_field].dat.vec
else:
acc = self.solution.split()[self.c_field].dat.vec_wo
with acc as sol:
self.condensed_ksp.solve(rhs, sol)
def adjoint_solve(self):
r"""Solve the adjoint problem.
"""
# Make sure appcontext is attached to the DM before the adjoint solve.
dm = self.ts.getDM()
with ExitStack() as stack:
for ctx in chain(
(
self.inserted_options(),
dmhooks.add_hooks(dm, self, appctx=self._ctx),
),
self._transfer_operators,
):
stack.enter_context(ctx)
self.ts.adjointSolve()
check_ts_convergence(self.ts)
def sc_solve(self, pc):
"""Solve the condensed linear system for the
condensed field.
:arg pc: a Preconditioner instance.
"""
dm = self.trace_ksp.getDM()
with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref):
# Solve the system for the Lagrange multipliers
with self.schur_rhs.dat.vec_ro as b:
if self.trace_ksp.getInitialGuessNonzero():
acc = self.trace_solution.dat.vec
else:
acc = self.trace_solution.dat.vec_wo
with acc as x_trace:
self.trace_ksp.solve(b, x_trace)
def solve(self, bounds=None):
r"""Solve the variational problem.
:arg bounds: Optional bounds on the solution (lower, upper).
``lower`` and ``upper`` must both be
:class:`~.Function`\s. or :class:`~.Vector`\s.
.. note::
If bounds are provided the ``snes_type`` must be set to
``vinewtonssls`` or ``vinewtonrsls``.
"""
# Make sure the DM has this solver's callback functions
self._ctx.set_function(self.snes)
self._ctx.set_jacobian(self.snes)
# Make sure appcontext is attached to the DM before we solve.
dm = self.snes.getDM()
for dbc in self._problem.dirichlet_bcs():
dbc.apply(self._problem.u)
if bounds is not None:
lower, upper = bounds
with lower.dat.vec_ro as lb, upper.dat.vec_ro as ub:
self.snes.setVariableBounds(lb, ub)
work = self._work
with self._problem.u.dat.vec as u:
u.copy(work)
with ExitStack() as stack:
# Ensure options database has full set of options (so monitors
# work right)
for ctx in chain(
(self.inserted_options(),
dmhooks.add_hooks(dm, self, appctx=self._ctx)),
self._transfer_operators):
stack.enter_context(ctx)
self.snes.solve(None, work)
work.copy(u)
self._setup = True
solving_utils.check_snes_convergence(self.snes)
def solve(self, x, b):
if not isinstance(x, (function.Function, vector.Vector)):
raise TypeError(
"Provided solution is a '%s', not a Function or Vector" %
type(x).__name__)
if isinstance(b, vector.Vector):
b = b.function
if not isinstance(b, function.Function):
raise TypeError("Provided RHS is a '%s', not a Function" %
type(b).__name__)
if len(self.trial_space) > 1 and self.nullspace is not None:
self.nullspace._apply(self.trial_space.dof_dset.field_ises)
if len(self.test_space) > 1 and self.transpose_nullspace is not None:
self.transpose_nullspace._apply(
self.test_space.dof_dset.field_ises, transpose=True)
if len(self.trial_space) > 1 and self.near_nullspace is not None:
self.near_nullspace._apply(self.trial_space.dof_dset.field_ises,
near=True)
if self.A.has_bcs:
b = self._lifted(b)
if self.ksp.getInitialGuessNonzero():
acc = x.dat.vec
else:
acc = x.dat.vec_wo
with self.inserted_options(
), b.dat.vec_ro as rhs, acc as solution, dmhooks.add_hooks(
self.ksp.dm, self):
self.ksp.solve(rhs, solution)
r = self.ksp.getConvergedReason()
if r < 0:
raise ConvergenceError(
"LinearSolver failed to converge after %d iterations with reason: %s",
self.ksp.getIterationNumber(), solving_utils.KSPReasons[r])
def _la_solve(A, x, b, **kwargs):
r"""Solve a linear algebra problem.
:arg A: the assembled bilinear form, a :class:`.Matrix`.
:arg x: the :class:`.Function` to write the solution into.
:arg b: the :class:`.Function` defining the right hand side values.
:kwarg solver_parameters: optional solver parameters.
:kwarg nullspace: an optional :class:`.VectorSpaceBasis` (or
:class:`.MixedVectorSpaceBasis`) spanning the null space of
the operator.
:kwarg transpose_nullspace: as for the nullspace, but used to
make the right hand side consistent.
:kwarg near_nullspace: as for the nullspace, but used to add
the near nullspace.
:kwarg options_prefix: an optional prefix used to distinguish
PETSc options. If not provided a unique prefix will be
created. Use this option if you want to pass options
to the solver from the command line in addition to
through the ``solver_parameters`` dict.
.. note::
This function no longer accepts :class:`.DirichletBC`\s or
:class:`.EquationBC`\s as arguments.
Any boundary conditions must be applied when assembling the
bilinear form as:
.. code-block:: python
A = assemble(a, bcs=[bc1])
solve(A, x, b)
Example usage:
.. code-block:: python
_la_solve(A, x, b, solver_parameters=parameters_dict)."""
bcs, solver_parameters, nullspace, nullspace_T, near_nullspace, \
options_prefix = _extract_linear_solver_args(A, x, b, **kwargs)
if bcs is not None:
raise RuntimeError(
"It is no longer possible to apply or change boundary conditions after assembling the matrix `A`; pass any necessary boundary conditions to `assemble` when assembling `A`."
)
solver = ls.LinearSolver(A,
solver_parameters=solver_parameters,
nullspace=nullspace,
transpose_nullspace=nullspace_T,
near_nullspace=near_nullspace,
options_prefix=options_prefix)
if isinstance(x, firedrake.Vector):
x = x.function
# linear MG doesn't need RHS, supply zero.
lvp = vs.LinearVariationalProblem(a=A.a, L=0, u=x, bcs=A.bcs)
mat_type = A.mat_type
appctx = solver_parameters.get("appctx", {})
ctx = solving_utils._SNESContext(lvp,
mat_type=mat_type,
pmat_type=mat_type,
appctx=appctx,
options_prefix=options_prefix)
dm = solver.ksp.dm
with dmhooks.add_hooks(dm, solver, appctx=ctx):
solver.solve(x, b)
def __init__(self, problem, **kwargs):
r"""
:arg problem: A :class:`NonlinearVariationalProblem` to solve.
:kwarg nullspace: an optional :class:`.VectorSpaceBasis` (or
:class:`.MixedVectorSpaceBasis`) spanning the null
space of the operator.
:kwarg transpose_nullspace: as for the nullspace, but used to
make the right hand side consistent.
:kwarg near_nullspace: as for the nullspace, but used to
specify the near nullspace (for multigrid solvers).
:kwarg solver_parameters: Solver parameters to pass to PETSc.
This should be a dict mapping PETSc options to values.
:kwarg appctx: A dictionary containing application context that
is passed to the preconditioner if matrix-free.
:kwarg options_prefix: an optional prefix used to distinguish
PETSc options. If not provided a unique prefix will be
created. Use this option if you want to pass options
to the solver from the command line in addition to
through the ``solver_parameters`` dict.
:kwarg pre_jacobian_callback: A user-defined function that will
be called immediately before Jacobian assembly. This can
be used, for example, to update a coefficient function
that has a complicated dependence on the unknown solution.
:kwarg pre_function_callback: As above, but called immediately
before residual assembly
Example usage of the ``solver_parameters`` option: to set the
nonlinear solver type to just use a linear solver, use
.. code-block:: python
{'snes_type': 'ksponly'}
PETSc flag options (where the presence of the option means something) should
be specified with ``None``.
For example:
.. code-block:: python
{'snes_monitor': None}
To use the ``pre_jacobian_callback`` or ``pre_function_callback``
functionality, the user-defined function must accept the current
solution as a petsc4py Vec. Example usage is given below:
.. code-block:: python
def update_diffusivity(current_solution):
with cursol.dat.vec_wo as v:
current_solution.copy(v)
solve(trial*test*dx == dot(grad(cursol), grad(test))*dx, diffusivity)
solver = NonlinearVariationalSolver(problem,
pre_jacobian_callback=update_diffusivity)
"""
assert isinstance(problem, NonlinearVariationalProblem)
parameters = kwargs.get("solver_parameters")
if "parameters" in kwargs:
raise TypeError("Use solver_parameters, not parameters")
nullspace = kwargs.get("nullspace")
nullspace_T = kwargs.get("transpose_nullspace")
near_nullspace = kwargs.get("near_nullspace")
options_prefix = kwargs.get("options_prefix")
pre_j_callback = kwargs.get("pre_jacobian_callback")
pre_f_callback = kwargs.get("pre_function_callback")
super(NonlinearVariationalSolver, self).__init__(parameters, options_prefix)
# Allow anything, interpret "matfree" as matrix_free.
mat_type = self.parameters.get("mat_type")
pmat_type = self.parameters.get("pmat_type")
matfree = mat_type == "matfree"
pmatfree = pmat_type == "matfree"
appctx = kwargs.get("appctx")
ctx = solving_utils._SNESContext(problem,
mat_type=mat_type,
pmat_type=pmat_type,
appctx=appctx,
pre_jacobian_callback=pre_j_callback,
pre_function_callback=pre_f_callback,
options_prefix=self.options_prefix)
# No preconditioner by default for matrix-free
if (problem.Jp is not None and pmatfree) or matfree:
self.set_default_parameter("pc_type", "none")
elif ctx.is_mixed:
# Mixed problem, use jacobi pc if user has not supplied
# one.
self.set_default_parameter("pc_type", "jacobi")
self.snes = PETSc.SNES().create(comm=problem.dm.comm)
self._problem = problem
self._ctx = ctx
self._work = problem.u.dof_dset.layout_vec.duplicate()
self.snes.setDM(problem.dm)
ctx.set_function(self.snes)
ctx.set_jacobian(self.snes)
ctx.set_nullspace(nullspace, problem.J.arguments()[0].function_space()._ises,
transpose=False, near=False)
ctx.set_nullspace(nullspace_T, problem.J.arguments()[1].function_space()._ises,
transpose=True, near=False)
ctx.set_nullspace(near_nullspace, problem.J.arguments()[0].function_space()._ises,
transpose=False, near=True)
ctx._nullspace = nullspace
ctx._nullspace_T = nullspace_T
ctx._near_nullspace = near_nullspace
# Set from options now, so that people who want to noodle with
# the snes object directly (mostly Patrick), can. We need the
# DM with an app context in place so that if the DM is active
# on a subKSP the context is available.
dm = self.snes.getDM()
with dmhooks.add_hooks(dm, self, appctx=self._ctx, save=False):
self.set_from_options(self.snes)
# Used for custom grid transfer.
self._transfer_operators = ()
self._setup = False
def applyTranspose(self, pc, X, Y):
dm = self._dm
with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref):
self.pc.applyTranspose(X, Y)
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):
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()
# 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)
self._ctx_ref = _SNESContext(nproblem,
mat_type,
mat_type,
octx.appctx,
options_prefix=options_prefix)
pc.setDM(dm)
pc.setOptionsPrefix(options_prefix)
pc.setOperators(Pmat, Pmat)
self.pc = pc
with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref, save=False):
pc.setFromOptions()
def applyTranspose(self, pc, x, y):
dm = pc.getDM()
with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref):
self.pc.applyTranspose(x, y)
def initialize(self, pc):
"""Set up the problem context. This takes the incoming
three-field system and constructs the static
condensation operators using Slate expressions.
A KSP is created for the reduced system. The eliminated
variables are recovered via back-substitution.
"""
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.bcs import DirichletBC
from firedrake.function import Function
from firedrake.functionspace import FunctionSpace
from firedrake.interpolation import interpolate
prefix = pc.getOptionsPrefix() + "condensed_field_"
A, P = pc.getOperators()
self.cxt = A.getPythonContext()
if not isinstance(self.cxt, ImplicitMatrixContext):
raise ValueError("Context must be an ImplicitMatrixContext")
self.bilinear_form = self.cxt.a
# Retrieve the mixed function space
W = self.bilinear_form.arguments()[0].function_space()
if len(W) > 3:
raise NotImplementedError("Only supports up to three function spaces.")
elim_fields = PETSc.Options().getString(pc.getOptionsPrefix()
+ "pc_sc_eliminate_fields",
None)
if elim_fields:
elim_fields = [int(i) for i in elim_fields.split(',')]
else:
# By default, we condense down to the last field in the
# mixed space.
elim_fields = [i for i in range(0, len(W) - 1)]
condensed_fields = list(set(range(len(W))) - set(elim_fields))
if len(condensed_fields) != 1:
raise NotImplementedError("Cannot condense to more than one field")
c_field, = condensed_fields
# Need to duplicate a space which is NOT
# associated with a subspace of a mixed space.
Vc = FunctionSpace(W.mesh(), W[c_field].ufl_element())
bcs = []
cxt_bcs = self.cxt.row_bcs
for bc in cxt_bcs:
if bc.function_space().index != c_field:
raise NotImplementedError("Strong BC set on unsupported space")
if isinstance(bc.function_arg, Function):
bc_arg = interpolate(bc.function_arg, Vc)
else:
# Constants don't need to be interpolated
bc_arg = bc.function_arg
bcs.append(DirichletBC(Vc, bc_arg, bc.sub_domain))
mat_type = PETSc.Options().getString(prefix + "mat_type", "aij")
self.c_field = c_field
self.condensed_rhs = Function(Vc)
self.residual = Function(W)
self.solution = Function(W)
# Get expressions for the condensed linear system
A_tensor = Tensor(self.bilinear_form)
reduced_sys = self.condensed_system(A_tensor, self.residual, elim_fields)
S_expr = reduced_sys.lhs
r_expr = reduced_sys.rhs
# Construct the condensed right-hand side
self._assemble_Srhs = create_assembly_callable(
r_expr,
tensor=self.condensed_rhs,
form_compiler_parameters=self.cxt.fc_params)
# Allocate and set the condensed operator
self.S = allocate_matrix(S_expr,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type,
options_prefix=prefix,
appctx=self.get_appctx(pc))
self._assemble_S = create_assembly_callable(
S_expr,
tensor=self.S,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type)
self._assemble_S()
Smat = self.S.petscmat
# If a different matrix is used for preconditioning,
# assemble this as well
if A != P:
self.cxt_pc = P.getPythonContext()
P_tensor = Tensor(self.cxt_pc.a)
P_reduced_sys = self.condensed_system(P_tensor,
self.residual,
elim_fields)
S_pc_expr = P_reduced_sys.lhs
self.S_pc_expr = S_pc_expr
# Allocate and set the condensed operator
self.S_pc = allocate_matrix(S_expr,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type,
options_prefix=prefix,
appctx=self.get_appctx(pc))
self._assemble_S_pc = create_assembly_callable(
S_pc_expr,
tensor=self.S_pc,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type)
self._assemble_S_pc()
Smat_pc = self.S_pc.petscmat
else:
self.S_pc_expr = S_expr
Smat_pc = Smat
# Get nullspace for the condensed operator (if any).
# This is provided as a user-specified callback which
# returns the basis for the nullspace.
nullspace = self.cxt.appctx.get("condensed_field_nullspace", None)
if nullspace is not None:
nsp = nullspace(Vc)
Smat.setNullSpace(nsp.nullspace(comm=pc.comm))
# Create a SNESContext for the DM associated with the trace problem
self._ctx_ref = self.new_snes_ctx(pc,
S_expr,
bcs,
mat_type,
self.cxt.fc_params,
options_prefix=prefix)
# Push new context onto the dm associated with the condensed problem
c_dm = Vc.dm
# Set up ksp for the condensed problem
c_ksp = PETSc.KSP().create(comm=pc.comm)
c_ksp.incrementTabLevel(1, parent=pc)
# Set the dm for the condensed solver
c_ksp.setDM(c_dm)
c_ksp.setDMActive(False)
c_ksp.setOptionsPrefix(prefix)
c_ksp.setOperators(A=Smat, P=Smat_pc)
self.condensed_ksp = c_ksp
with dmhooks.add_hooks(c_dm, self,
appctx=self._ctx_ref,
save=False):
c_ksp.setFromOptions()
# Set up local solvers for backwards substitution
self.local_solvers = self.local_solver_calls(A_tensor,
self.residual,
self.solution,
elim_fields)
def initialize(self, pc):
"""Set up the problem context. Take the original
mixed problem and reformulate the problem as a
hybridized mixed system.
A KSP is created for the Lagrange multiplier system.
"""
from firedrake import (FunctionSpace, Function, Constant,
TrialFunction, TrialFunctions, TestFunction,
DirichletBC)
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.formmanipulation import split_form
from ufl.algorithms.replace import replace
# Extract the problem context
prefix = pc.getOptionsPrefix() + "hybridization_"
_, P = pc.getOperators()
self.ctx = P.getPythonContext()
if not isinstance(self.ctx, ImplicitMatrixContext):
raise ValueError("The python context must be an ImplicitMatrixContext")
test, trial = self.ctx.a.arguments()
V = test.function_space()
mesh = V.mesh()
if len(V) != 2:
raise ValueError("Expecting two function spaces.")
if all(Vi.ufl_element().value_shape() for Vi in V):
raise ValueError("Expecting an H(div) x L2 pair of spaces.")
# Automagically determine which spaces are vector and scalar
for i, Vi in enumerate(V):
if Vi.ufl_element().sobolev_space().name == "HDiv":
self.vidx = i
else:
assert Vi.ufl_element().sobolev_space().name == "L2"
self.pidx = i
# Create the space of approximate traces.
W = V[self.vidx]
if W.ufl_element().family() == "Brezzi-Douglas-Marini":
tdegree = W.ufl_element().degree()
else:
try:
# If we have a tensor product element
h_deg, v_deg = W.ufl_element().degree()
tdegree = (h_deg - 1, v_deg - 1)
except TypeError:
tdegree = W.ufl_element().degree() - 1
TraceSpace = FunctionSpace(mesh, "HDiv Trace", tdegree)
# Break the function spaces and define fully discontinuous spaces
broken_elements = ufl.MixedElement([ufl.BrokenElement(Vi.ufl_element()) for Vi in V])
V_d = FunctionSpace(mesh, broken_elements)
# Set up the functions for the original, hybridized
# and schur complement systems
self.broken_solution = Function(V_d)
self.broken_residual = Function(V_d)
self.trace_solution = Function(TraceSpace)
self.unbroken_solution = Function(V)
self.unbroken_residual = Function(V)
shapes = (V[self.vidx].finat_element.space_dimension(),
np.prod(V[self.vidx].shape))
domain = "{[i,j]: 0 <= i < %d and 0 <= j < %d}" % shapes
instructions = """
for i, j
w[i,j] = w[i,j] + 1
end
"""
self.weight = Function(V[self.vidx])
par_loop((domain, instructions), ufl.dx, {"w": (self.weight, INC)},
is_loopy_kernel=True)
instructions = """
for i, j
vec_out[i,j] = vec_out[i,j] + vec_in[i,j]/w[i,j]
end
"""
self.average_kernel = (domain, instructions)
# Create the symbolic Schur-reduction:
# Original mixed operator replaced with "broken"
# arguments
arg_map = {test: TestFunction(V_d),
trial: TrialFunction(V_d)}
Atilde = Tensor(replace(self.ctx.a, arg_map))
gammar = TestFunction(TraceSpace)
n = ufl.FacetNormal(mesh)
sigma = TrialFunctions(V_d)[self.vidx]
if mesh.cell_set._extruded:
Kform = (gammar('+') * ufl.jump(sigma, n=n) * ufl.dS_h
+ gammar('+') * ufl.jump(sigma, n=n) * ufl.dS_v)
else:
Kform = (gammar('+') * ufl.jump(sigma, n=n) * ufl.dS)
# Here we deal with boundaries. If there are Neumann
# conditions (which should be enforced strongly for
# H(div)xL^2) then we need to add jump terms on the exterior
# facets. If there are Dirichlet conditions (which should be
# enforced weakly) then we need to zero out the trace
# variables there as they are not active (otherwise the hybrid
# problem is not well-posed).
# If boundary conditions are contained in the ImplicitMatrixContext:
if self.ctx.row_bcs:
# Find all the subdomains with neumann BCS
# These are Dirichlet BCs on the vidx space
neumann_subdomains = set()
for bc in self.ctx.row_bcs:
if bc.function_space().index == self.pidx:
raise NotImplementedError("Dirichlet conditions for scalar variable not supported. Use a weak bc")
if bc.function_space().index != self.vidx:
raise NotImplementedError("Dirichlet bc set on unsupported space.")
# append the set of sub domains
subdom = bc.sub_domain
if isinstance(subdom, str):
neumann_subdomains |= set([subdom])
else:
neumann_subdomains |= set(as_tuple(subdom, numbers.Integral))
# separate out the top and bottom bcs
extruded_neumann_subdomains = neumann_subdomains & {"top", "bottom"}
neumann_subdomains = neumann_subdomains - extruded_neumann_subdomains
integrand = gammar * ufl.dot(sigma, n)
measures = []
trace_subdomains = []
if mesh.cell_set._extruded:
ds = ufl.ds_v
for subdomain in sorted(extruded_neumann_subdomains):
measures.append({"top": ufl.ds_t, "bottom": ufl.ds_b}[subdomain])
trace_subdomains.extend(sorted({"top", "bottom"} - extruded_neumann_subdomains))
else:
ds = ufl.ds
if "on_boundary" in neumann_subdomains:
measures.append(ds)
else:
measures.extend((ds(sd) for sd in sorted(neumann_subdomains)))
markers = [int(x) for x in mesh.exterior_facets.unique_markers]
dirichlet_subdomains = set(markers) - neumann_subdomains
trace_subdomains.extend(sorted(dirichlet_subdomains))
for measure in measures:
Kform += integrand*measure
trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), subdomain) for subdomain in trace_subdomains]
else:
# No bcs were provided, we assume weak Dirichlet conditions.
# We zero out the contribution of the trace variables on
# the exterior boundary. Extruded cells will have both
# horizontal and vertical facets
trace_subdomains = ["on_boundary"]
if mesh.cell_set._extruded:
trace_subdomains.extend(["bottom", "top"])
trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), subdomain) for subdomain in trace_subdomains]
# Make a SLATE tensor from Kform
K = Tensor(Kform)
# Assemble the Schur complement operator and right-hand side
self.schur_rhs = Function(TraceSpace)
self._assemble_Srhs = create_assembly_callable(
K * Atilde.inv * AssembledVector(self.broken_residual),
tensor=self.schur_rhs,
form_compiler_parameters=self.ctx.fc_params)
mat_type = PETSc.Options().getString(prefix + "mat_type", "aij")
schur_comp = K * Atilde.inv * K.T
self.S = allocate_matrix(schur_comp, bcs=trace_bcs,
form_compiler_parameters=self.ctx.fc_params,
mat_type=mat_type,
options_prefix=prefix,
appctx=self.get_appctx(pc))
self._assemble_S = create_assembly_callable(schur_comp,
tensor=self.S,
bcs=trace_bcs,
form_compiler_parameters=self.ctx.fc_params,
mat_type=mat_type)
with timed_region("HybridOperatorAssembly"):
self._assemble_S()
Smat = self.S.petscmat
nullspace = self.ctx.appctx.get("trace_nullspace", None)
if nullspace is not None:
nsp = nullspace(TraceSpace)
Smat.setNullSpace(nsp.nullspace(comm=pc.comm))
# Create a SNESContext for the DM associated with the trace problem
self._ctx_ref = self.new_snes_ctx(pc,
schur_comp,
trace_bcs,
mat_type,
self.ctx.fc_params)
# dm associated with the trace problem
trace_dm = TraceSpace.dm
# KSP for the system of Lagrange multipliers
trace_ksp = PETSc.KSP().create(comm=pc.comm)
trace_ksp.incrementTabLevel(1, parent=pc)
# Set the dm for the trace solver
trace_ksp.setDM(trace_dm)
trace_ksp.setDMActive(False)
trace_ksp.setOptionsPrefix(prefix)
trace_ksp.setOperators(Smat, Smat)
self.trace_ksp = trace_ksp
with dmhooks.add_hooks(trace_dm, self,
appctx=self._ctx_ref,
save=False):
trace_ksp.setFromOptions()
split_mixed_op = dict(split_form(Atilde.form))
split_trace_op = dict(split_form(K.form))
# Generate reconstruction calls
self._reconstruction_calls(split_mixed_op, split_trace_op)
