def initialize(self, pc):
_, K = pc.getOperators()
self.ctx = K.getPythonContext()
self.prefix = pc.getOptionsPrefix()
self.assembleK()
# Reassign K
K = self.K.M.handle
self.Z = self.ctx.appctx.get("projection_mat", None)
assert self.Z != None
# Project the jacobian
Kp = K.PtAP(self.Z)
KpInv = PETSc.KSP().create()
KpInv.setOperators(Kp)
# KpInv.setType("preonly")
# KpInvPC = KpInv.getPC()
# KpInvPC.setFactorSolverType("mumps")
KpInv.setUp()
self.KpInv = KpInv
# Create reduced space vectors
self.Xp, _ = self.Z.createVecs()
self.Yp, _ = self.Z.createVecs()
def setup_ksp(self, A, pc):
ksp = PETSc.KSP().create(comm=pc.comm)
ksp.incrementTabLevel(1, parent=pc)
ksp.setOperators(A)
ksp.setOptionsPrefix(self.options_prefix)
ksp.setFromOptions()
ksp.setUp()
return ksp
def __init__(self, Q, free_bids=["on_boundary"]):
(V, I_interp) = Q.get_space_for_inner()
self.free_bids = free_bids
u = fd.TrialFunction(V)
v = fd.TestFunction(V)
n = fd.FacetNormal(V.mesh())
def surf_grad(u):
return fd.sym(fd.grad(u) - fd.outer(fd.grad(u) * n, n))
a = (fd.inner(surf_grad(u), surf_grad(v)) + fd.inner(u, v)) * fd.ds
# petsc doesn't like matrices with zero rows
a += 1e-10 * fd.inner(u, v) * fd.dx
A = fd.assemble(a, mat_type="aij")
A = A.petscmat
tdim = V.mesh().topological_dimension()
lsize = fd.Function(V).vector().local_size()
def get_nodes_bc(bc):
nodes = bc.nodes
return nodes[nodes < lsize]
def get_nodes_bid(bid):
bc = fd.DirichletBC(V, fd.Constant(tdim * (0, )), bid)
return get_nodes_bc(bc)
free_nodes = np.concatenate(
[get_nodes_bid(bid) for bid in self.free_bids])
free_dofs = np.concatenate(
[tdim * free_nodes + i for i in range(tdim)])
free_dofs = np.unique(np.sort(free_dofs))
self.free_is = PETSc.IS().createGeneral(free_dofs)
lgr, lgc = A.getLGMap()
self.global_free_is_row = lgr.applyIS(self.free_is)
self.global_free_is_col = lgc.applyIS(self.free_is)
A = A.createSubMatrix(self.global_free_is_row, self.global_free_is_col)
# A.view()
A.assemble()
self.A = A
Aksp = PETSc.KSP().create()
Aksp.setOperators(self.A)
Aksp.setOptionsPrefix("A_")
opts = PETSc.Options()
opts["A_ksp_type"] = "cg"
opts["A_pc_type"] = "hypre"
opts["A_ksp_atol"] = 1e-10
opts["A_ksp_rtol"] = 1e-10
Aksp.setUp()
Aksp.setFromOptions()
self.Aksp = Aksp
def initialize(self, pc):
from firedrake import TrialFunction, TestFunction, dx, assemble, inner, parameters
if pc.getType() != "python":
raise ValueError("Expecting PC type python")
prefix = pc.getOptionsPrefix()
options_prefix = prefix + "Mp_"
# we assume P has things stuffed inside of it
_, P = pc.getOperators()
context = P.getPythonContext()
test, trial = context.a.arguments()
if test.function_space() != trial.function_space():
raise ValueError(
"MassInvPC only makes sense if test and trial space are the same"
)
V = test.function_space()
mu = context.appctx.get("mu", 1.0)
u = TrialFunction(V)
v = TestFunction(V)
# Handle vector and tensor-valued spaces.
# 1/mu goes into the inner product in case it varies spatially.
a = inner(1 / mu * u, v) * dx
opts = PETSc.Options()
mat_type = opts.getString(options_prefix + "mat_type",
parameters["default_matrix_type"])
A = assemble(a,
form_compiler_parameters=context.fc_params,
mat_type=mat_type,
options_prefix=options_prefix)
A.force_evaluation()
Pmat = A.petscmat
Pmat.setNullSpace(P.getNullSpace())
tnullsp = P.getTransposeNullSpace()
if tnullsp.handle != 0:
Pmat.setTransposeNullSpace(tnullsp)
ksp = PETSc.KSP().create(comm=pc.comm)
ksp.incrementTabLevel(1, parent=pc)
ksp.setOperators(Pmat)
ksp.setOptionsPrefix(options_prefix)
ksp.setFromOptions()
ksp.setUp()
self.ksp = ksp
def GetPETScBFBTApproximateToSchurInverse():
Fass = assemble(lhs(F))
Gass = assemble(inner(u, v) * dx)
# bcin.apply(Fass)
# bcnoslip.apply(Fass)
# bcin.apply(Gass)
# bcnoslip.apply(Gass)
A = Fass.M[0,0].handle
BT = Fass.M[0,1].handle
B = Fass.M[1,0].handle
D = Fass.M[1,1].handle
VM = Gass.M[0, 0].handle
# Adiag = A.getDiagonal()
Adiag = VM.getRowSum()
invAdiag = Adiag.duplicate()
Adiag.copy(invAdiag)
invAdiag.reciprocal()
Einv = A.duplicate()
Einv.setDiagonal(invAdiag)
# Einv.convert(PETSc.Mat.Type.SEQAIJ)
BEBT = B.matMult(Einv.matMult(BT))
BEBT_ksp = PETSc.KSP().create()
BEBT_ksp.setOperators(BEBT)
opts = PETSc.Options()
BEBT_ksp.setOptionsPrefix("bebt_")
opts['bebt_ksp_type'] = 'richardson'
opts['bebt_pc_type'] = 'hypre'
opts['bebt_ksp_max_it'] = 1
opts['bebt_ksp_atol'] = 1.0e-9
opts['bebt_ksp_rtol'] = 1.0e-9
BEBT_ksp.setUp()
BEBT_ksp.setFromOptions()
MIDDLE = B.matMult(Einv.matMult(A.matMult(Einv.matMult(BT))))
MIDDLE.scale(-1)
class SchurInvApprox(object):
def mult(self, mat, x, y):
y1 = y.duplicate()
BEBT_ksp.solve(x, y1)
y2 = y.duplicate()
MIDDLE.mult(y1, y2)
BEBT_ksp.solve(y2, y)
schur = PETSc.Mat()
schur.createPython(D.getSizes(), SchurInvApprox())
schur.setUp()
return schur
def solve(ctx, state):
out_file = File(solution_out) if solution_out else None
mass = state["mass"]
hats = state["hats"]
u = ctx[2]["u"]
u0 = ctx[2]["u0"]
solver = ctx[1]
from firedrake.petsc import PETSc
ksp_hats = PETSc.KSP()
ksp_hats.create()
ksp_hats.setOperators(hats)
opts = PETSc.Options()
opts['ksp_type'] = inner_ksp
opts['ksp_max_it'] = max_iterations
opts['pc_type'] = 'hypre'
ksp_hats.setFromOptions()
class SchurInv(object):
def mult(self, mat, x, y):
tmp1 = y.duplicate()
tmp2 = y.duplicate()
ksp_hats.solve(x, tmp1)
mass.mult(tmp1, tmp2)
ksp_hats.solve(tmp2, y)
pc_schur = PETSc.Mat()
pc_schur.createPython(mass.getSizes(), SchurInv())
pc_schur.setUp()
pc = solver.snes.ksp.pc
pc.setFieldSplitSchurPreType(PETSc.PC.SchurPreType.USER, pc_schur)
for step in range(steps):
casper.invoke_task(ctx, "assign", state)
solver.solve()
if out_file is not None:
out_file.write(u.split()[0], time=step)
if compute_norms:
nu = norm(u)
if comm.rank == 0:
print(step, 'L2(u):', nu)
def V_inv_mass_ksp(self, V):
"""
A KSP inverting a mass matrix
:arg V: a function space.
:returns: A PETSc KSP for inverting (V, V).
"""
key = V.dim()
try:
return self._V_inv_mass_ksp[key]
except KeyError:
M = firedrake.assemble(firedrake.inner(firedrake.TestFunction(V),
firedrake.TrialFunction(V))*firedrake.dx)
M.force_evaluation()
ksp = PETSc.KSP().create(comm=V.comm)
ksp.setOperators(M.petscmat)
ksp.setOptionsPrefix("{}_prolongation_mass_".format(V.ufl_element()._short_name))
ksp.setType("preonly")
ksp.pc.setType("cholesky")
ksp.setFromOptions()
ksp.setUp()
return self._V_inv_mass_ksp.setdefault(key, ksp)
def test_duplicate(a, bcs):
test, trial = a.arguments()
if test.function_space().shape == ():
rhs_form = inner(Constant(1), test)*dx
elif test.function_space().shape == (2, ):
rhs_form = inner(Constant((1, 1)), test)*dx
if bcs is not None:
Af = assemble(a, mat_type="matfree", bcs=bcs)
rhs = assemble(rhs_form, bcs=bcs)
else:
Af = assemble(a, mat_type="matfree")
rhs = assemble(rhs_form)
# matrix-free duplicate creates a matrix-free copy of Af
# we have not implemented the default copy = False
B_petsc = Af.petscmat.duplicate(copy=True)
ksp = PETSc.KSP().create()
ksp.setOperators(Af.petscmat)
ksp.setFromOptions()
solution1 = Function(test.function_space())
solution2 = Function(test.function_space())
# Solve system with original matrix A
with rhs.dat.vec_ro as b, solution1.dat.vec as x:
ksp.solve(b, x)
# Multiply with copied matrix B
with solution1.dat.vec_ro as x, solution2.dat.vec_ro as y:
B_petsc.mult(x, y)
# Check if original rhs is equal to BA^-1 (rhs)
assert np.allclose(rhs.vector().array(), solution2.vector().array())
def __init__(self,
A,
P=None,
solver_parameters=None,
nullspace=None,
transpose_nullspace=None,
near_nullspace=None,
options_prefix=None):
"""A linear solver for assembled systems (Ax = b).
:arg A: a :class:`~.MatrixBase` (the operator).
:arg P: an optional :class:`~.MatrixBase` to construct any
preconditioner from; if none is supplied ``A`` is
used to construct the preconditioner.
:kwarg parameters: (optional) dict of 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 set
the near nullpace.
: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::
Any boundary conditions for this solve *must* have been
applied when assembling the operator.
"""
if not isinstance(A, matrix.MatrixBase):
raise TypeError("Provided operator is a '%s', not a MatrixBase" %
type(A).__name__)
if P is not None and not isinstance(P, matrix.MatrixBase):
raise TypeError(
"Provided preconditioner is a '%s', not a MatrixBase" %
type(P).__name__)
self.A = A
self.comm = A.comm
self.P = P if P is not None else A
# Set up parameters mixin
super(LinearSolver, self).__init__(solver_parameters, options_prefix)
self.A.petscmat.setOptionsPrefix(self.options_prefix)
self.P.petscmat.setOptionsPrefix(self.options_prefix)
# Set some defaults
self.set_default_parameter("ksp_rtol", "1e-7")
# If preconditioning matrix is matrix-free, then default to no
# preconditioning.
if isinstance(self.P, matrix.ImplicitMatrix):
self.set_default_parameter("pc_type", "none")
elif self.P.block_shape != (1, 1):
# Otherwise, mixed problems default to jacobi.
self.set_default_parameter("pc_type", "jacobi")
self.ksp = PETSc.KSP().create(comm=self.comm)
W = self.test_space
# DM provides fieldsplits (but not operators)
self.ksp.setDM(W.dm)
self.ksp.setDMActive(False)
if nullspace is not None:
nullspace._apply(self.A)
if P is not None:
nullspace._apply(self.P)
if transpose_nullspace is not None:
transpose_nullspace._apply(self.A, transpose=True)
if P is not None:
transpose_nullspace._apply(self.P, transpose=True)
if near_nullspace is not None:
near_nullspace._apply(self.A, near=True)
if P is not None:
near_nullspace._apply(self.P, near=True)
self.nullspace = nullspace
self.transpose_nullspace = transpose_nullspace
self.near_nullspace = near_nullspace
# Operator setting must come after null space has been
# applied
# Force evaluation here
self.A.force_evaluation()
self.P.force_evaluation()
self.ksp.setOperators(A=self.A.petscmat, P=self.P.petscmat)
# Set from options now (we're not allowed to change parameters
# anyway).
self.set_from_options(self.ksp)
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 ufl.algorithms.map_integrands import map_integrand_dags
from firedrake import (FunctionSpace, TrialFunction, TrialFunctions,
TestFunction, Function, BrokenElement,
MixedElement, FacetNormal, Constant,
DirichletBC, Projector)
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.formmanipulation import ArgumentReplacer, split_form
# Extract the problem context
prefix = pc.getOptionsPrefix()
_, P = pc.getOperators()
context = P.getPythonContext()
test, trial = context.a.arguments()
V = test.function_space()
if V.mesh().cell_set._extruded:
# TODO: Merge FIAT branch to support TPC trace elements
raise NotImplementedError("Not implemented on extruded meshes.")
# Break the function spaces and define fully discontinuous spaces
broken_elements = [BrokenElement(Vi.ufl_element()) for Vi in V]
elem = MixedElement(broken_elements)
V_d = FunctionSpace(V.mesh(), elem)
arg_map = {test: TestFunction(V_d), trial: TrialFunction(V_d)}
# Replace the problems arguments with arguments defined
# on the new discontinuous spaces
replacer = ArgumentReplacer(arg_map)
new_form = map_integrand_dags(replacer, context.a)
# Create the space of approximate traces.
# The vector function space will have a non-empty value_shape
W = next(v for v in V if bool(v.ufl_element().value_shape()))
if W.ufl_element().family() in ["Raviart-Thomas", "RTCF"]:
tdegree = W.ufl_element().degree() - 1
else:
tdegree = W.ufl_element().degree()
# NOTE: Once extruded is ready, we will need to be aware of this
# and construct the appropriate trace space for the HDiv element
TraceSpace = FunctionSpace(V.mesh(), "HDiv Trace", tdegree)
# NOTE: For extruded, we will need to add "on_top" and "on_bottom"
trace_conditions = [
DirichletBC(TraceSpace, Constant(0.0), "on_boundary")
]
# Set up the functions for the original, hybridized
# and schur complement systems
self.broken_solution = Function(V_d)
self.broken_rhs = Function(V_d)
self.trace_solution = Function(TraceSpace)
self.unbroken_solution = Function(V)
self.unbroken_rhs = Function(V)
# Create the symbolic Schur-reduction
Atilde = Tensor(new_form)
gammar = TestFunction(TraceSpace)
n = FacetNormal(V.mesh())
# Vector trial function will have a non-empty ufl_shape
sigma = next(f for f in TrialFunctions(V_d) if bool(f.ufl_shape))
# NOTE: Once extruded is ready, this will change slightly
# to include both horizontal and vertical interior facets
K = Tensor(gammar('+') * ufl.dot(sigma, n) * ufl.dS)
# Assemble the Schur complement operator and right-hand side
self.schur_rhs = Function(TraceSpace)
self._assemble_Srhs = create_assembly_callable(
K * Atilde.inv * self.broken_rhs,
tensor=self.schur_rhs,
form_compiler_parameters=context.fc_params)
schur_comp = K * Atilde.inv * K.T
self.S = allocate_matrix(schur_comp,
bcs=trace_conditions,
form_compiler_parameters=context.fc_params)
self._assemble_S = create_assembly_callable(
schur_comp,
tensor=self.S,
bcs=trace_conditions,
form_compiler_parameters=context.fc_params)
self._assemble_S()
self.S.force_evaluation()
Smat = self.S.petscmat
# Nullspace for the multiplier problem
nullsp = P.getNullSpace()
if nullsp.handle != 0:
new_vecs = get_trace_nullspace_vecs(K * Atilde.inv, nullsp, V, V_d,
TraceSpace)
tr_nullsp = PETSc.NullSpace().create(vectors=new_vecs,
comm=pc.comm)
Smat.setNullSpace(tr_nullsp)
# Set up the KSP for the system of Lagrange multipliers
ksp = PETSc.KSP().create(comm=pc.comm)
ksp.setOptionsPrefix(prefix + "trace_")
ksp.setTolerances(rtol=1e-13)
ksp.setOperators(Smat)
ksp.setUp()
ksp.setFromOptions()
self.ksp = ksp
# Now we construct the local tensors for the reconstruction stage
# TODO: Add support for mixed tensors and these variables
# become unnecessary
split_forms = split_form(new_form)
A = Tensor(next(sf.form for sf in split_forms if sf.indices == (0, 0)))
B = Tensor(next(sf.form for sf in split_forms if sf.indices == (1, 0)))
C = Tensor(next(sf.form for sf in split_forms if sf.indices == (1, 1)))
trial = TrialFunction(
FunctionSpace(V.mesh(), BrokenElement(W.ufl_element())))
K_local = Tensor(gammar('+') * ufl.dot(trial, n) * ufl.dS)
# Split functions and reconstruct each bit separately
sigma_h, u_h = self.broken_solution.split()
g, f = self.broken_rhs.split()
# Pressure reconstruction
M = B * A.inv * B.T + C
u_sol = M.inv * f + M.inv * (
B * A.inv * K_local.T * self.trace_solution - B * A.inv * g)
self._assemble_pressure = create_assembly_callable(
u_sol, tensor=u_h, form_compiler_parameters=context.fc_params)
# Velocity reconstruction
sigma_sol = A.inv * g + A.inv * (B.T * u_h -
K_local.T * self.trace_solution)
self._assemble_velocity = create_assembly_callable(
sigma_sol,
tensor=sigma_h,
form_compiler_parameters=context.fc_params)
# Set up the projector for projecting the broken solution
# into the unbroken finite element spaces
# NOTE: Tolerance here matters!
sigma_b, _ = self.broken_solution.split()
sigma_u, _ = self.unbroken_solution.split()
self.projector = Projector(sigma_b,
sigma_u,
solver_parameters={
"ksp_type": "cg",
"ksp_rtol": 1e-13
})
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 nonlocal_integral_eq(
mesh,
scatterer_bdy_id,
outer_bdy_id,
wave_number,
options_prefix=None,
solver_parameters=None,
fspace=None,
vfspace=None,
true_sol_grad=None,
queue=None,
fspace_analog=None,
qbx_kwargs=None,
):
r"""
see run_method for descriptions of unlisted args
args:
:arg queue: A command queue for the computing context
gamma and beta are used to precondition
with the following equation:
\Delta u - \kappa^2 \gamma u = 0
(\partial_n - i\kappa\beta) u |_\Sigma = 0
"""
with_refinement = True
# away from the excluded region, but firedrake and meshmode point
# into
pyt_inner_normal_sign = -1
ambient_dim = mesh.geometric_dimension()
# {{{ Create operator
from pytential import sym
r"""
..math:
x \in \Sigma
grad_op(x) =
\nabla(
\int_\Gamma(
u(y) \partial_n H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
)
"""
grad_op = pyt_inner_normal_sign * sym.grad(
ambient_dim,
sym.D(HelmholtzKernel(ambient_dim),
sym.var("u"),
k=sym.var("k"),
qbx_forced_limit=None))
r"""
..math:
x \in \Sigma
op(x) =
i \kappa \cdot
\int_\Gamma(
u(y) \partial_n H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
"""
op = pyt_inner_normal_sign * 1j * sym.var("k") * (sym.D(
HelmholtzKernel(ambient_dim),
sym.var("u"),
k=sym.var("k"),
qbx_forced_limit=None))
pyt_grad_op = fd2mm.fd_bind(
queue.context,
fspace_analog,
grad_op,
source=(fspace, scatterer_bdy_id),
target=(vfspace, outer_bdy_id),
with_refinement=with_refinement,
qbx_kwargs=qbx_kwargs,
)
pyt_op = fd2mm.fd_bind(
queue.context,
fspace_analog,
op,
source=(fspace, scatterer_bdy_id),
target=(fspace, outer_bdy_id),
with_refinement=with_refinement,
qbx_kwargs=qbx_kwargs,
)
# }}}
class MatrixFreeB(object):
def __init__(self, A, pyt_grad_op, pyt_op, queue, kappa):
"""
:arg kappa: The wave number
"""
self.queue = queue
self.k = kappa
self.pyt_op = pyt_op
self.pyt_grad_op = pyt_grad_op
self.A = A
# {{{ Create some functions needed for multing
self.x_fntn = Function(fspace)
self.potential_int = Function(fspace)
self.potential_int.dat.data[:] = 0.0
self.grad_potential_int = Function(vfspace)
self.grad_potential_int.dat.data[:] = 0.0
self.pyt_result = Function(fspace)
self.n = FacetNormal(mesh)
self.v = TestFunction(fspace)
# }}}
def mult(self, mat, x, y):
# Perform pytential operation
self.x_fntn.dat.data[:] = x[:]
self.pyt_op(self.queue,
self.potential_int,
u=self.x_fntn,
k=self.k)
self.pyt_grad_op(self.queue,
self.grad_potential_int,
u=self.x_fntn,
k=self.k)
# Integrate the potential
r"""
Compute the inner products using firedrake. Note this
will be subtracted later, hence appears off by a sign.
.. math::
\langle
n(x) \cdot \nabla(
\int_\Gamma(
u(y) \partial_n H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
), v
\rangle_\Sigma
- \langle
i \kappa \cdot
\int_\Gamma(
u(y) \partial_n H_0^{(1)}(\kappa |x - y|)
)d\gamma(y), v
\rangle_\Sigma
"""
self.pyt_result = assemble(
inner(inner(self.grad_potential_int, self.n), self.v) *
ds(outer_bdy_id) -
inner(self.potential_int, self.v) * ds(outer_bdy_id))
# y <- Ax - evaluated potential
self.A.mult(x, y)
with self.pyt_result.dat.vec_ro as ep:
y.axpy(-1, ep)
# {{{ Compute normal helmholtz operator
u = TrialFunction(fspace)
v = TestFunction(fspace)
r"""
.. math::
\langle
\nabla u, \nabla v
\rangle
- \kappa^2 \cdot \langle
u, v
\rangle
- i \kappa \langle
u, v
\rangle_\Sigma
"""
a = inner(grad(u), grad(v)) * dx \
- Constant(wave_number**2) * inner(u, v) * dx \
- Constant(1j * wave_number) * inner(u, v) * ds(outer_bdy_id)
# get the concrete matrix from a general bilinear form
A = assemble(a).M.handle
# }}}
# {{{ Setup Python matrix
B = PETSc.Mat().create()
# build matrix context
Bctx = MatrixFreeB(A, pyt_grad_op, pyt_op, queue, wave_number)
# set up B as same size as A
B.setSizes(*A.getSizes())
B.setType(B.Type.PYTHON)
B.setPythonContext(Bctx)
B.setUp()
# }}}
# {{{ Create rhs
# Remember f is \partial_n(true_sol)|_\Gamma
# so we just need to compute \int_\Gamma\partial_n(true_sol) H(x-y)
from pytential import sym
sigma = sym.make_sym_vector("sigma", ambient_dim)
r"""
..math:
x \in \Sigma
grad_op(x) =
\nabla(
\int_\Gamma(
f(y) H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
)
"""
grad_op = pyt_inner_normal_sign * \
sym.grad(ambient_dim, sym.S(HelmholtzKernel(ambient_dim),
sym.n_dot(sigma),
k=sym.var("k"), qbx_forced_limit=None))
r"""
..math:
x \in \Sigma
op(x) =
i \kappa \cdot
\int_\Gamma(
f(y) H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
)
"""
op = 1j * sym.var("k") * pyt_inner_normal_sign * \
sym.S(HelmholtzKernel(ambient_dim),
sym.n_dot(sigma),
k=sym.var("k"),
qbx_forced_limit=None)
rhs_grad_op = fd2mm.fd_bind(
queue.context,
fspace_analog,
grad_op,
source=(vfspace, scatterer_bdy_id),
target=(vfspace, outer_bdy_id),
with_refinement=with_refinement,
qbx_kwargs=qbx_kwargs,
)
rhs_op = fd2mm.fd_bind(
queue.context,
fspace_analog,
op,
source=(vfspace, scatterer_bdy_id),
target=(fspace, outer_bdy_id),
with_refinement=with_refinement,
qbx_kwargs=qbx_kwargs,
)
f_grad_convoluted = Function(vfspace)
f_convoluted = Function(fspace)
rhs_grad_op(queue, f_grad_convoluted, sigma=true_sol_grad, k=wave_number)
rhs_op(queue, f_convoluted, sigma=true_sol_grad, k=wave_number)
r"""
\langle
f, v
\rangle_\Gamma
+ \langle
i \kappa \cdot \int_\Gamma(
f(y) H_0^{(1)}(\kappa |x - y|)
)d\gamma(y), v
\rangle_\Sigma
- \langle
n(x) \cdot \nabla(
\int_\Gamma(
f(y) H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
), v
\rangle_\Sigma
"""
rhs_form = inner(inner(true_sol_grad, FacetNormal(mesh)),
v) * ds(scatterer_bdy_id) \
+ inner(f_convoluted, v) * ds(outer_bdy_id) \
- inner(inner(f_grad_convoluted, FacetNormal(mesh)),
v) * ds(outer_bdy_id)
rhs = assemble(rhs_form)
# {{{ set up a solver:
solution = Function(fspace, name="Computed Solution")
# {{{ Used for preconditioning
if 'gamma' in solver_parameters or 'beta' in solver_parameters:
gamma = complex(solver_parameters.pop('gamma', 1.0))
import cmath
beta = complex(solver_parameters.pop('beta', cmath.sqrt(gamma)))
p = inner(grad(u), grad(v)) * dx \
- Constant(wave_number**2 * gamma) * inner(u, v) * dx \
- Constant(1j * wave_number * beta) * inner(u, v) * ds(outer_bdy_id)
P = assemble(p).M.handle
else:
P = A
# }}}
# Set up options to contain solver parameters:
ksp = PETSc.KSP().create()
if solver_parameters['pc_type'] == 'pyamg':
del solver_parameters['pc_type'] # We are using the AMG preconditioner
pyamg_tol = solver_parameters.get('pyamg_tol', None)
if pyamg_tol is not None:
pyamg_tol = float(pyamg_tol)
pyamg_maxiter = solver_parameters.get('pyamg_maxiter', None)
if pyamg_maxiter is not None:
pyamg_maxiter = int(pyamg_maxiter)
ksp.setOperators(B)
ksp.setUp()
pc = ksp.pc
pc.setType(pc.Type.PYTHON)
pc.setPythonContext(
AMGTransmissionPreconditioner(wave_number,
fspace,
A,
tol=pyamg_tol,
maxiter=pyamg_maxiter,
use_plane_waves=True))
# Otherwise use regular preconditioner
else:
ksp.setOperators(B, P)
options_manager = OptionsManager(solver_parameters, options_prefix)
options_manager.set_from_options(ksp)
with rhs.dat.vec_ro as b:
with solution.dat.vec as x:
ksp.solve(b, x)
# }}}
return ksp, solution
def initialize(self, pc):
"""Set up the problem context. Take the original
H1-problem and partition the spaces/functions
into 'interior' and 'facet' parts.
A KSP is created for the reduced system after
static condensation is applied.
"""
from firedrake import (FunctionSpace, Function, TrialFunction,
TestFunction)
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from ufl.algorithms.replace import replace
# Extract python context
prefix = pc.getOptionsPrefix() + "static_condensation_"
_, P = pc.getOperators()
self.cxt = P.getPythonContext()
if not isinstance(self.cxt, ImplicitMatrixContext):
raise ValueError("Context must be an ImplicitMatrixContext")
test, trial = self.cxt.a.arguments()
V = test.function_space()
mesh = V.mesh()
if len(V) > 1:
raise ValueError("Cannot use this PC for mixed problems.")
if V.ufl_element().sobolev_space().name != "H1":
raise ValueError("Expecting an H1-conforming element.")
if not V.ufl_element().cell().is_simplex():
raise NotImplementedError("Only simplex meshes are implemented.")
top_dim = V.finat_element._element.ref_el.get_dimension()
if not V.finat_element.entity_dofs()[top_dim][0]:
raise RuntimeError("There are no interior dofs to eliminate.")
# We decompose the space into an interior part and facet part
interior_element = V.ufl_element()["interior"]
facet_element = V.ufl_element()["facet"]
V_int = FunctionSpace(mesh, interior_element)
V_facet = FunctionSpace(mesh, facet_element)
# Get transfer kernel for moving data
self._transfer_kernel = get_transfer_kernels({
'h1-space': V,
'interior-space': V_int,
'facet-space': V_facet
})
# Set up functions for the H1 functions and the interior/trace parts
self.trace_solution = Function(V_facet)
self.interior_solution = Function(V_int)
self.h1_solution = Function(V)
self.h1_residual = Function(V)
self.interior_residual = Function(V_int)
self.trace_residual = Function(V_facet)
# TODO: Handle strong bcs in Slate
if self.cxt.row_bcs:
raise NotImplementedError("Strong bcs not implemented yet")
self.bcs = None
A00 = Tensor(
replace(self.cxt.a, {
test: TestFunction(V_int),
trial: TrialFunction(V_int)
}))
A01 = Tensor(
replace(self.cxt.a, {
test: TestFunction(V_int),
trial: TrialFunction(V_facet)
}))
A10 = Tensor(
replace(self.cxt.a, {
test: TestFunction(V_facet),
trial: TrialFunction(V_int)
}))
A11 = Tensor(
replace(self.cxt.a, {
test: TestFunction(V_facet),
trial: TrialFunction(V_facet)
}))
# Schur complement operator
S = A11 - A10 * A00.inv * A01
self.S = allocate_matrix(S,
bcs=self.bcs,
form_compiler_parameters=self.cxt.fc_params)
self._assemble_S = create_assembly_callable(
S,
tensor=self.S,
bcs=self.bcs,
form_compiler_parameters=self.cxt.fc_params)
self._assemble_S()
Smat = self.S.petscmat
# Nullspace for the reduced system
nullspace = create_sc_nullspace(P, V, V_facet, pc.comm)
if nullspace:
Smat.setNullSpace(nullspace)
# Set up KSP for the reduced problem
sc_ksp = PETSc.KSP().create(comm=pc.comm)
sc_ksp.setOptionsPrefix(prefix)
sc_ksp.setOperators(Smat)
sc_ksp.setUp()
sc_ksp.setFromOptions()
self.sc_ksp = sc_ksp
# Set up rhs for the reduced problem
F0 = AssembledVector(self.interior_residual)
self.sc_rhs = Function(V_facet)
self.sc_rhs_thunk = Function(V_facet)
self._assemble_sc_rhs_thunk = create_assembly_callable(
-A10 * A00.inv * F0,
tensor=self.sc_rhs_thunk,
form_compiler_parameters=self.cxt.fc_params)
# Reconstruction calls
u_facet = AssembledVector(self.trace_solution)
self._assemble_interior_u = create_assembly_callable(
A00.inv * (F0 - A01 * u_facet),
tensor=self.interior_solution,
form_compiler_parameters=self.cxt.fc_params)
def nonlocal_integral_eq(
mesh,
scatterer_bdy_id,
outer_bdy_id,
wave_number,
options_prefix=None,
solver_parameters=None,
fspace=None,
vfspace=None,
true_sol_grad_expr=None,
actx=None,
dgfspace=None,
dgvfspace=None,
meshmode_src_connection=None,
qbx_kwargs=None,
):
r"""
see run_method for descriptions of unlisted args
args:
gamma and beta are used to precondition
with the following equation:
\Delta u - \kappa^2 \gamma u = 0
(\partial_n - i\kappa\beta) u |_\Sigma = 0
"""
# make sure we get outer bdy id as tuple in case it consists of multiple ids
if isinstance(outer_bdy_id, int):
outer_bdy_id = [outer_bdy_id]
outer_bdy_id = tuple(outer_bdy_id)
# away from the excluded region, but firedrake and meshmode point
# into
pyt_inner_normal_sign = -1
ambient_dim = mesh.geometric_dimension()
# {{{ Build src and tgt
# build connection meshmode near src boundary -> src boundary inside meshmode
from meshmode.discretization.poly_element import \
InterpolatoryQuadratureSimplexGroupFactory
from meshmode.discretization.connection import make_face_restriction
factory = InterpolatoryQuadratureSimplexGroupFactory(
dgfspace.finat_element.degree)
src_bdy_connection = make_face_restriction(actx,
meshmode_src_connection.discr,
factory, scatterer_bdy_id)
# source is a qbx layer potential
from pytential.qbx import QBXLayerPotentialSource
disable_refinement = (fspace.mesh().geometric_dimension() == 3)
qbx = QBXLayerPotentialSource(src_bdy_connection.to_discr,
**qbx_kwargs,
_disable_refinement=disable_refinement)
# get target indices and point-set
target_indices, target = get_target_points_and_indices(
fspace, outer_bdy_id)
# }}}
# build the operations
from pytential import bind, sym
r"""
..math:
x \in \Sigma
grad_op(x) =
\nabla(
\int_\Gamma(
u(y) \partial_n H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
)
"""
grad_op = pyt_inner_normal_sign * sym.grad(
ambient_dim,
sym.D(HelmholtzKernel(ambient_dim),
sym.var("u"),
k=sym.var("k"),
qbx_forced_limit=None))
r"""
..math:
x \in \Sigma
op(x) =
i \kappa \cdot
\int_\Gamma(
u(y) \partial_n H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
"""
op = pyt_inner_normal_sign * 1j * sym.var("k") * (sym.D(
HelmholtzKernel(ambient_dim),
sym.var("u"),
k=sym.var("k"),
qbx_forced_limit=None))
# bind the operations
pyt_grad_op = bind((qbx, target), grad_op)
pyt_op = bind((qbx, target), op)
# }}}
class MatrixFreeB(object):
def __init__(self, A, pyt_grad_op, pyt_op, actx, kappa):
"""
:arg kappa: The wave number
"""
self.actx = actx
self.k = kappa
self.pyt_op = pyt_op
self.pyt_grad_op = pyt_grad_op
self.A = A
self.meshmode_src_connection = meshmode_src_connection
# {{{ Create some functions needed for multing
self.x_fntn = Function(fspace)
# CG
self.potential_int = Function(fspace)
self.potential_int.dat.data[:] = 0.0
self.grad_potential_int = Function(vfspace)
self.grad_potential_int.dat.data[:] = 0.0
self.pyt_result = Function(fspace)
self.n = FacetNormal(mesh)
self.v = TestFunction(fspace)
# some meshmode ones
self.x_mm_fntn = self.meshmode_src_connection.discr.empty(
self.actx, dtype='c')
# }}}
def mult(self, mat, x, y):
# Copy function data into the fivredrake function
self.x_fntn.dat.data[:] = x[:]
# Transfer the function to meshmode
self.meshmode_src_connection.from_firedrake(project(
self.x_fntn, dgfspace),
out=self.x_mm_fntn)
# Restrict to boundary
x_mm_fntn_on_bdy = src_bdy_connection(self.x_mm_fntn)
# Apply the operation
potential_int_mm = self.pyt_op(self.actx,
u=x_mm_fntn_on_bdy,
k=self.k)
grad_potential_int_mm = self.pyt_grad_op(self.actx,
u=x_mm_fntn_on_bdy,
k=self.k)
# Store in firedrake
self.potential_int.dat.data[target_indices] = potential_int_mm.get(
)
for dim in range(grad_potential_int_mm.shape[0]):
self.grad_potential_int.dat.data[
target_indices, dim] = grad_potential_int_mm[dim].get()
# Integrate the potential
r"""
Compute the inner products using firedrake. Note this
will be subtracted later, hence appears off by a sign.
.. math::
\langle
n(x) \cdot \nabla(
\int_\Gamma(
u(y) \partial_n H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
), v
\rangle_\Sigma
- \langle
i \kappa \cdot
\int_\Gamma(
u(y) \partial_n H_0^{(1)}(\kappa |x - y|)
)d\gamma(y), v
\rangle_\Sigma
"""
self.pyt_result = assemble(
inner(inner(self.grad_potential_int, self.n), self.v) *
ds(outer_bdy_id) -
inner(self.potential_int, self.v) * ds(outer_bdy_id))
# y <- Ax - evaluated potential
self.A.mult(x, y)
with self.pyt_result.dat.vec_ro as ep:
y.axpy(-1, ep)
# {{{ Compute normal helmholtz operator
u = TrialFunction(fspace)
v = TestFunction(fspace)
r"""
.. math::
\langle
\nabla u, \nabla v
\rangle
- \kappa^2 \cdot \langle
u, v
\rangle
- i \kappa \langle
u, v
\rangle_\Sigma
"""
a = inner(grad(u), grad(v)) * dx \
- Constant(wave_number**2) * inner(u, v) * dx \
- Constant(1j * wave_number) * inner(u, v) * ds(outer_bdy_id)
# get the concrete matrix from a general bilinear form
A = assemble(a).M.handle
# }}}
# {{{ Setup Python matrix
B = PETSc.Mat().create()
# build matrix context
Bctx = MatrixFreeB(A, pyt_grad_op, pyt_op, actx, wave_number)
# set up B as same size as A
B.setSizes(*A.getSizes())
B.setType(B.Type.PYTHON)
B.setPythonContext(Bctx)
B.setUp()
# }}}
# {{{ Create rhs
# Remember f is \partial_n(true_sol)|_\Gamma
# so we just need to compute \int_\Gamma\partial_n(true_sol) H(x-y)
sigma = sym.make_sym_vector("sigma", ambient_dim)
r"""
..math:
x \in \Sigma
grad_op(x) =
\nabla(
\int_\Gamma(
f(y) H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
)
"""
grad_op = pyt_inner_normal_sign * \
sym.grad(ambient_dim, sym.S(HelmholtzKernel(ambient_dim),
sym.n_dot(sigma),
k=sym.var("k"), qbx_forced_limit=None))
r"""
..math:
x \in \Sigma
op(x) =
i \kappa \cdot
\int_\Gamma(
f(y) H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
)
"""
op = 1j * sym.var("k") * pyt_inner_normal_sign * \
sym.S(HelmholtzKernel(ambient_dim),
sym.n_dot(sigma),
k=sym.var("k"),
qbx_forced_limit=None)
rhs_grad_op = bind((qbx, target), grad_op)
rhs_op = bind((qbx, target), op)
# Transfer to meshmode
metadata = {'quadrature_degree': 2 * fspace.ufl_element().degree()}
dg_true_sol_grad = project(true_sol_grad_expr,
dgvfspace,
form_compiler_parameters=metadata)
true_sol_grad_mm = meshmode_src_connection.from_firedrake(dg_true_sol_grad,
actx=actx)
true_sol_grad_mm = src_bdy_connection(true_sol_grad_mm)
# Apply the operations
f_grad_convoluted_mm = rhs_grad_op(actx,
sigma=true_sol_grad_mm,
k=wave_number)
f_convoluted_mm = rhs_op(actx, sigma=true_sol_grad_mm, k=wave_number)
# Transfer function back to firedrake
f_grad_convoluted = Function(vfspace)
f_convoluted = Function(fspace)
f_grad_convoluted.dat.data[:] = 0.0
f_convoluted.dat.data[:] = 0.0
for dim in range(f_grad_convoluted_mm.shape[0]):
f_grad_convoluted.dat.data[target_indices,
dim] = f_grad_convoluted_mm[dim].get()
f_convoluted.dat.data[target_indices] = f_convoluted_mm.get()
r"""
\langle
f, v
\rangle_\Gamma
+ \langle
i \kappa \cdot \int_\Gamma(
f(y) H_0^{(1)}(\kappa |x - y|)
)d\gamma(y), v
\rangle_\Sigma
- \langle
n(x) \cdot \nabla(
\int_\Gamma(
f(y) H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
), v
\rangle_\Sigma
"""
rhs_form = inner(inner(true_sol_grad_expr, FacetNormal(mesh)),
v) * ds(scatterer_bdy_id, metadata=metadata) \
+ inner(f_convoluted, v) * ds(outer_bdy_id) \
- inner(inner(f_grad_convoluted, FacetNormal(mesh)),
v) * ds(outer_bdy_id)
rhs = assemble(rhs_form)
# {{{ set up a solver:
solution = Function(fspace, name="Computed Solution")
# {{{ Used for preconditioning
if 'gamma' in solver_parameters or 'beta' in solver_parameters:
gamma = complex(solver_parameters.pop('gamma', 1.0))
import cmath
beta = complex(solver_parameters.pop('beta', cmath.sqrt(gamma)))
p = inner(grad(u), grad(v)) * dx \
- Constant(wave_number**2 * gamma) * inner(u, v) * dx \
- Constant(1j * wave_number * beta) * inner(u, v) * ds(outer_bdy_id)
P = assemble(p).M.handle
else:
P = A
# }}}
# Set up options to contain solver parameters:
ksp = PETSc.KSP().create()
if solver_parameters['pc_type'] == 'pyamg':
del solver_parameters['pc_type'] # We are using the AMG preconditioner
pyamg_tol = solver_parameters.get('pyamg_tol', None)
if pyamg_tol is not None:
pyamg_tol = float(pyamg_tol)
pyamg_maxiter = solver_parameters.get('pyamg_maxiter', None)
if pyamg_maxiter is not None:
pyamg_maxiter = int(pyamg_maxiter)
ksp.setOperators(B)
ksp.setUp()
pc = ksp.pc
pc.setType(pc.Type.PYTHON)
pc.setPythonContext(
AMGTransmissionPreconditioner(wave_number,
fspace,
A,
tol=pyamg_tol,
maxiter=pyamg_maxiter,
use_plane_waves=True))
# Otherwise use regular preconditioner
else:
ksp.setOperators(B, P)
options_manager = OptionsManager(solver_parameters, options_prefix)
options_manager.set_from_options(ksp)
import petsc4py.PETSc
petsc4py.PETSc.Sys.popErrorHandler()
with rhs.dat.vec_ro as b:
with solution.dat.vec as x:
ksp.solve(b, x)
# }}}
return ksp, solution
def initialize(self, pc):
from firedrake import TrialFunction, TestFunction, dx, \
assemble, inner, grad, split, Constant, parameters
from firedrake.assemble import allocate_matrix, create_assembly_callable
if pc.getType() != "python":
raise ValueError("Expecting PC type python")
prefix = pc.getOptionsPrefix() + "pcd_"
# we assume P has things stuffed inside of it
_, P = pc.getOperators()
context = P.getPythonContext()
test, trial = context.a.arguments()
if test.function_space() != trial.function_space():
raise ValueError("Pressure space test and trial space differ")
Q = test.function_space()
p = TrialFunction(Q)
q = TestFunction(Q)
mass = p * q * dx
# Regularisation to avoid having to think about nullspaces.
stiffness = inner(grad(p), grad(q)) * dx + Constant(1e-6) * p * q * dx
opts = PETSc.Options()
# we're inverting Mp and Kp, so default them to assembled.
# Fp only needs its action, so default it to mat-free.
# These can of course be overridden.
# only Fp is referred to in update, so that's the only
# one we stash.
default = parameters["default_matrix_type"]
Mp_mat_type = opts.getString(prefix + "Mp_mat_type", default)
Kp_mat_type = opts.getString(prefix + "Kp_mat_type", default)
self.Fp_mat_type = opts.getString(prefix + "Fp_mat_type", "matfree")
Mp = assemble(mass,
form_compiler_parameters=context.fc_params,
mat_type=Mp_mat_type,
options_prefix=prefix + "Mp_")
Kp = assemble(stiffness,
form_compiler_parameters=context.fc_params,
mat_type=Kp_mat_type,
options_prefix=prefix + "Kp_")
Mp.force_evaluation()
Kp.force_evaluation()
# FIXME: Should we transfer nullspaces over. I think not.
Mksp = PETSc.KSP().create(comm=pc.comm)
Mksp.incrementTabLevel(1, parent=pc)
Mksp.setOptionsPrefix(prefix + "Mp_")
Mksp.setOperators(Mp.petscmat)
Mksp.setUp()
Mksp.setFromOptions()
self.Mksp = Mksp
Kksp = PETSc.KSP().create(comm=pc.comm)
Kksp.incrementTabLevel(1, parent=pc)
Kksp.setOptionsPrefix(prefix + "Kp_")
Kksp.setOperators(Kp.petscmat)
Kksp.setUp()
Kksp.setFromOptions()
self.Kksp = Kksp
state = context.appctx["state"]
Re = context.appctx.get("Re", 1.0)
velid = context.appctx["velocity_space"]
u0 = split(state)[velid]
fp = 1.0 / Re * inner(grad(p), grad(q)) * dx + inner(u0,
grad(p)) * q * dx
self.Re = Re
self.Fp = allocate_matrix(fp,
form_compiler_parameters=context.fc_params,
mat_type=self.Fp_mat_type,
options_prefix=prefix + "Fp_")
self._assemble_Fp = create_assembly_callable(
fp,
tensor=self.Fp,
form_compiler_parameters=context.fc_params,
mat_type=self.Fp_mat_type)
self._assemble_Fp()
self.Fp.force_evaluation()
Fpmat = self.Fp.petscmat
self.workspace = [Fpmat.createVecLeft() for i in (0, 1)]
def __init__(self, Q, fixed_bids=[], extra_bcs=[], direct_solve=False):
if isinstance(extra_bcs, fd.DirichletBC):
extra_bcs = [extra_bcs]
self.direct_solve = direct_solve
self.fixed_bids = fixed_bids # fixed parts of bdry
self.params = self.get_params() # solver parameters
self.Q = Q
"""
V: type fd.FunctionSpace
I: type PETSc.Mat, interpolation matrix between V and ControlSpace
"""
(V, I_interp) = Q.get_space_for_inner()
free_bids = list(V.mesh().topology.exterior_facets.unique_markers)
self.free_bids = [int(i) for i in free_bids] # np.int->int
for bid in self.fixed_bids:
self.free_bids.remove(bid)
# Some weak forms have a nullspace. We import the nullspace if no
# parts of the bdry are fixed (we assume that a DirichletBC is
# sufficient to empty the nullspace).
nsp = None
if len(self.fixed_bids) == 0:
nsp_functions = self.get_nullspace(V)
if nsp_functions is not None:
nsp = fd.VectorSpaceBasis(nsp_functions)
nsp.orthonormalize()
bcs = []
# impose homogeneous Dirichlet bcs on bdry parts that are fixed.
if len(self.fixed_bids) > 0:
dim = V.value_size
if dim == 2:
zerovector = fd.Constant((0, 0))
elif dim == 3:
zerovector = fd.Constant((0, 0, 0))
else:
raise NotImplementedError
bcs.append(fd.DirichletBC(V, zerovector, self.fixed_bids))
if len(extra_bcs) > 0:
bcs += extra_bcs
if len(bcs) == 0:
bcs = None
a = self.get_weak_form(V)
A = fd.assemble(a, mat_type='aij', bcs=bcs)
ls = fd.LinearSolver(A,
solver_parameters=self.params,
nullspace=nsp,
transpose_nullspace=nsp)
self.ls = ls
self.A = A.petscmat
self.interpolated = False
# If the matrix I is passed, replace A with transpose(I)*A*I
# and set up a ksp solver for self.riesz_map
if I_interp is not None:
self.interpolated = True
ITAI = self.A.PtAP(I_interp)
from firedrake.petsc import PETSc
import numpy as np
zero_rows = []
# if there are zero-rows, replace them with rows that
# have 1 on the diagonal entry
for row in range(*ITAI.getOwnershipRange()):
(cols, vals) = ITAI.getRow(row)
valnorm = np.linalg.norm(vals)
if valnorm < 1e-13:
zero_rows.append(row)
for row in zero_rows:
ITAI.setValue(row, row, 1.0)
ITAI.assemble()
# overwrite the self.A created by get_impl
self.A = ITAI
# create ksp solver for self.riesz_map
Aksp = PETSc.KSP().create(comm=V.comm)
Aksp.setOperators(ITAI)
Aksp.setType("preonly")
Aksp.pc.setType("cholesky")
Aksp.pc.setFactorSolverType("mumps")
Aksp.setFromOptions()
Aksp.setUp()
self.Aksp = Aksp
if pc in ["fieldsplit", "ilu"]:
sigma = 100
# PC for the Schur complement solve
trial = TrialFunction(V)
test = TestFunction(V)
mass = assemble(inner(trial, test) * dx).M.handle
a = 1
c = (dt * lmbda) / (1 + dt * sigma)
hats = assemble(
sqrt(a) * inner(trial, test) * dx +
sqrt(c) * inner(grad(trial), grad(test)) * dx).M.handle
from firedrake.petsc import PETSc
ksp_hats = PETSc.KSP()
ksp_hats.create()
ksp_hats.setOperators(hats)
opts = PETSc.Options()
opts["ksp_type"] = inner_ksp
opts["ksp_max_it"] = maxit
opts["pc_type"] = "hypre"
ksp_hats.setFromOptions()
class SchurInv(object):
def mult(self, mat, x, y):
tmp1 = y.duplicate()
tmp2 = y.duplicate()
ksp_hats.solve(x, tmp1)
mass.mult(tmp1, tmp2)
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, assemble)
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.cxt = P.getPythonContext()
if not isinstance(self.cxt, ImplicitMatrixContext):
raise ValueError("The python context must be an ImplicitMatrixContext")
test, trial = self.cxt.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)
# Set up the KSP for the hdiv residual projection
hdiv_mass_ksp = PETSc.KSP().create(comm=pc.comm)
hdiv_mass_ksp.setOptionsPrefix(prefix + "hdiv_residual_")
# HDiv mass operator
p = TrialFunction(V[self.vidx])
q = TestFunction(V[self.vidx])
mass = ufl.dot(p, q)*ufl.dx
# TODO: Bcs?
M = assemble(mass, bcs=None, form_compiler_parameters=self.cxt.fc_params)
M.force_evaluation()
Mmat = M.petscmat
hdiv_mass_ksp.setOperators(Mmat)
hdiv_mass_ksp.setUp()
hdiv_mass_ksp.setFromOptions()
self.hdiv_mass_ksp = hdiv_mass_ksp
# Storing the result of A.inv * r, where A is the HDiv
# mass matrix and r is the HDiv residual
self._primal_r = Function(V[self.vidx])
tau = TestFunction(V_d[self.vidx])
self._assemble_broken_r = create_assembly_callable(
ufl.dot(self._primal_r, tau)*ufl.dx,
tensor=self.broken_residual.split()[self.vidx],
form_compiler_parameters=self.cxt.fc_params)
# 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.cxt.a, arg_map))
gammar = TestFunction(TraceSpace)
n = ufl.FacetNormal(mesh)
sigma = TrialFunctions(V_d)[self.vidx]
if mesh.cell_set._extruded:
Kform = (gammar('+') * ufl.dot(sigma, n) * ufl.dS_h +
gammar('+') * ufl.dot(sigma, n) * ufl.dS_v)
else:
Kform = (gammar('+') * ufl.dot(sigma, 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.cxt.row_bcs:
# Find all the subdomains with neumann BCS
# These are Dirichlet BCs on the vidx space
neumann_subdomains = set()
for bc in self.cxt.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, int))
# separate out the top and bottom bcs
extruded_neumann_subdomains = neumann_subdomains & {"top", "bottom"}
neumann_subdomains = neumann_subdomains.difference(extruded_neumann_subdomains)
integrand = gammar * ufl.dot(sigma, n)
measures = []
trace_subdomains = []
if mesh.cell_set._extruded:
ds = ufl.ds_v
for subdomain in 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.append(ds(tuple(neumann_subdomains)))
dirichlet_subdomains = set(mesh.exterior_facets.unique_markers) - neumann_subdomains
trace_subdomains.append(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.cxt.fc_params)
schur_comp = K * Atilde.inv * K.T
self.S = allocate_matrix(schur_comp, bcs=trace_bcs,
form_compiler_parameters=self.cxt.fc_params)
self._assemble_S = create_assembly_callable(schur_comp,
tensor=self.S,
bcs=trace_bcs,
form_compiler_parameters=self.cxt.fc_params)
self._assemble_S()
self.S.force_evaluation()
Smat = self.S.petscmat
# Nullspace for the multiplier problem
nullspace = create_schur_nullspace(P, -K * Atilde,
V, V_d, TraceSpace,
pc.comm)
if nullspace:
Smat.setNullSpace(nullspace)
# Set up the KSP for the system of Lagrange multipliers
trace_ksp = PETSc.KSP().create(comm=pc.comm)
trace_ksp.setOptionsPrefix(prefix)
trace_ksp.setOperators(Smat)
trace_ksp.setUp()
trace_ksp.setFromOptions()
self.trace_ksp = trace_ksp
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)
# NOTE: The projection stage *might* be replaced by a Fortin
# operator. We may want to allow the user to specify if they
# wish to use a Fortin operator over a projection, or vice-versa.
# In a future add-on, we can add a switch which chooses either
# the Fortin reconstruction or the usual KSP projection.
# Set up the projection KSP
hdiv_projection_ksp = PETSc.KSP().create(comm=pc.comm)
hdiv_projection_ksp.setOptionsPrefix(prefix + 'hdiv_projection_')
# Reuse the mass operator from the hdiv_mass_ksp
hdiv_projection_ksp.setOperators(Mmat)
# Construct the RHS for the projection stage
self._projection_rhs = Function(V[self.vidx])
self._assemble_projection_rhs = create_assembly_callable(
ufl.dot(self.broken_solution.split()[self.vidx], q)*ufl.dx,
tensor=self._projection_rhs,
form_compiler_parameters=self.cxt.fc_params)
# Finalize ksp setup
hdiv_projection_ksp.setUp()
hdiv_projection_ksp.setFromOptions()
self.hdiv_projection_ksp = hdiv_projection_ksp
def initialize(self, pc):
"""Set up the problem context. This takes the incoming
three-field hybridized system and constructs the static
condensation operators using Slate expressions.
A KSP is created for the reduced system for the Lagrange
multipliers. The scalar and flux fields are reconstructed
locally.
"""
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() + "hybrid_sc_"
_, P = pc.getOperators()
self.cxt = P.getPythonContext()
if not isinstance(self.cxt, ImplicitMatrixContext):
raise ValueError("Context must be an ImplicitMatrixContext")
# Retrieve the mixed function space, which is expected to
# be of the form: W = (DG_k)^n \times DG_k \times DG_trace
W = self.cxt.a.arguments()[0].function_space()
if len(W) != 3:
raise RuntimeError("Expecting three function spaces.")
# Assert a specific ordering of the spaces
# TODO: Clean this up
assert W[2].ufl_element().family() == "HDiv Trace"
# Extract trace space
T = W[2]
# Need to duplicate a trace space which is NOT
# associated with a subspace of a mixed space.
Tr = FunctionSpace(T.mesh(), T.ufl_element())
bcs = []
cxt_bcs = self.cxt.row_bcs
for bc in cxt_bcs:
assert bc.function_space() == T, (
"BCs should be imposing vanishing conditions on traces")
if isinstance(bc.function_arg, Function):
bc_arg = interpolate(bc.function_arg, Tr)
else:
# Constants don't need to be interpolated
bc_arg = bc.function_arg
bcs.append(DirichletBC(Tr, bc_arg, bc.sub_domain))
mat_type = PETSc.Options().getString(prefix + "mat_type", "aij")
self.r_lambda = Function(T)
self.residual = Function(W)
self.solution = Function(W)
# Perform symbolics only once
S_expr, r_lambda_expr, u_h_expr, q_h_expr = self._slate_expressions
self.S = allocate_matrix(S_expr,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type)
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
# Set up ksp for the trace problem
trace_ksp = PETSc.KSP().create(comm=pc.comm)
trace_ksp.incrementTabLevel(1, parent=pc)
trace_ksp.setOptionsPrefix(prefix)
trace_ksp.setOperators(Smat)
trace_ksp.setUp()
trace_ksp.setFromOptions()
self.trace_ksp = trace_ksp
self._assemble_Srhs = create_assembly_callable(
r_lambda_expr,
tensor=self.r_lambda,
form_compiler_parameters=self.cxt.fc_params)
q_h, u_h, lambda_h = self.solution.split()
# Assemble u_h using lambda_h
self._assemble_u = create_assembly_callable(
u_h_expr, tensor=u_h, form_compiler_parameters=self.cxt.fc_params)
# Recover q_h using both u_h and lambda_h
self._assemble_q = create_assembly_callable(
q_h_expr, tensor=q_h, form_compiler_parameters=self.cxt.fc_params)
def initialize(self, pc):
from firedrake import TrialFunction, TestFunction, Function, DirichletBC, dx, \
assemble, Mesh, inner, grad, split, Constant, parameters
from firedrake.assemble import allocate_matrix, create_assembly_callable
prefix = pc.getOptionsPrefix() + "pcd_"
_, P = pc.getOperators()
context = P.getPythonContext()
test, trial = context.a.arguments()
Q = test.function_space()
p = TrialFunction(Q)
q = TestFunction(Q)
nu = context.appctx["nu"]
dt = context.appctx["dt"]
mass = (1.0 / nu) * p * q * dx
stiffness = inner(grad(p), grad(q)) * dx
velid = context.appctx["velocity_space"]
self.bcs = context.appctx["PCDbc"]
opts = PETSc.Options()
default = parameters["default_matrix_type"]
Mp_mat_type = opts.getString(prefix + "Mp_mat_type", default)
Kp_mat_type = opts.getString(prefix + "Kp_mat_type", default)
Fp_mat_type = opts.getString(prefix + "Fp_mat_type", "matfree")
Mp = assemble(mass,
form_compiler_parameters=context.fc_params,
mat_type=Mp_mat_type,
options_prefix=prefix + "Mp_")
Kp = assemble(stiffness,
bcs=self.bcs,
form_compiler_parameters=context.fc_params,
mat_type=Kp_mat_type,
options_prefix=prefix + "Kp_")
Mksp = PETSc.KSP().create(comm=pc.comm)
Mksp.incrementTabLevel(1, parent=pc)
Mksp.setOptionsPrefix(prefix + "Mp_")
Mksp.setOperators(Mp.petscmat)
Mksp.setUp()
Mksp.setFromOptions()
self.Mksp = Mksp
Kksp = PETSc.KSP().create(comm=pc.comm)
Kksp.incrementTabLevel(1, parent=pc)
Kksp.setOptionsPrefix(prefix + "Kp_")
Kksp.setOperators(Kp.petscmat)
Kksp.setUp()
Kksp.setFromOptions()
self.Kksp = Kksp
u = context.appctx["u"]
fp = (1.0 / nu) * ((1.0 / dt) * p + dot(u, grad(p))) * q * dx
Fp = allocate_matrix(fp,
form_compiler_parameters=context.fc_params,
mat_type=Fp_mat_type,
options_prefix=prefix + "Fp_")
self._assemble_Fp = create_assembly_callable(
fp,
tensor=Fp,
form_compiler_parameters=context.fc_params,
mat_type=Fp_mat_type)
self.Fp = Fp
self._assemble_Fp()
self.workspace = [Fp.petscmat.createVecLeft() for i in (0, 1)]
self.tmp = Function(Q)
def __init__(self,
A,
P=None,
solver_parameters=None,
nullspace=None,
options_prefix=None):
"""A linear solver for assembled systems (Ax = b).
:arg A: a :class:`~.Matrix` (the operator).
:arg P: an optional :class:`~.Matrix` to construct any
preconditioner from; if none is supplied :data:`A` is
used to construct the preconditioner.
:kwarg parameters: (optional) dict of solver parameters.
:kwarg nullspace: an optional :class:`~.VectorSpaceBasis` (or
:class:`~.MixedVectorSpaceBasis` spanning the null space
of the operator.
: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 :data:`solver_parameters` dict.
.. note::
Any boundary conditions for this solve *must* have been
applied when assembling the operator.
"""
# Do this first so __del__ doesn't barf horribly if we get an
# error in __init__
if options_prefix is not None:
self._opt_prefix = options_prefix
self._auto_prefix = False
else:
self._opt_prefix = "firedrake_ksp_%d_" % LinearSolver._id
self._auto_prefix = True
LinearSolver._id += 1
self.A = A
self.P = P if P is not None else A
parameters = solver_parameters.copy(
) if solver_parameters is not None else {}
parameters.setdefault("ksp_rtol", "1e-7")
if self.P._M.sparsity.shape != (1, 1):
parameters.setdefault('pc_type', 'jacobi')
self.ksp = PETSc.KSP().create()
self.ksp.setOptionsPrefix(self._opt_prefix)
# Allow command-line arguments to override dict parameters
opts = PETSc.Options()
for k, v in opts.getAll().iteritems():
if k.startswith(self._opt_prefix):
parameters[k[len(self._opt_prefix):]] = v
self.parameters = parameters
W = self.A.a.arguments()[0].function_space()
# DM provides fieldsplits (but not operators)
self.ksp.setDM(W._dm)
self.ksp.setDMActive(False)
if nullspace is not None:
nullspace._apply(self.A._M)
if P is not None:
nullspace._apply(self.P._M)
self.nullspace = nullspace
self._W = W
# Operator setting must come after null space has been
# applied
# Force evaluation here
self.ksp.setOperators(A=self.A.M.handle, P=self.P.M.handle)
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)
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))
# Set up the KSP for the system of Lagrange multipliers
trace_ksp = PETSc.KSP().create(comm=pc.comm)
trace_ksp.setOptionsPrefix(prefix)
trace_ksp.setOperators(Smat)
trace_ksp.setUp()
trace_ksp.setFromOptions()
self.trace_ksp = trace_ksp
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)
def initialize(self, pc):
""" Set up the problem context. Takes the original
mixed problem and transforms it into the equivalent
hybrid-mixed system.
A KSP object is created for the Lagrange multipliers
on the top/bottom faces of the mesh cells.
"""
from firedrake import (FunctionSpace, Function, Constant,
FiniteElement, TensorProductElement,
TrialFunction, TrialFunctions, TestFunction,
DirichletBC, interval, MixedElement,
BrokenElement)
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.formmanipulation import split_form
from ufl.algorithms.replace import replace
from ufl.cell import TensorProductCell
# Extract PC context
prefix = pc.getOptionsPrefix() + "vert_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()
# Magically determine which spaces are vector and scalar valued
for i, Vi in enumerate(V):
# Vector-valued spaces will have a non-empty value_shape
if Vi.ufl_element().value_shape():
self.vidx = i
else:
self.pidx = i
Vv = V[self.vidx]
Vp = V[self.pidx]
# Create the space of approximate traces in the vertical.
# NOTE: Technically a hack since the resulting space is technically
# defined in cell interiors, however the degrees of freedom will only
# be geometrically defined on edges. Arguments will only be used in
# surface integrals
deg, _ = Vv.ufl_element().degree()
# Assumes a tensor product cell (quads, triangular-prisms, cubes)
if not isinstance(Vp.ufl_element().cell(), TensorProductCell):
raise NotImplementedError(
"Currently only implemented for tensor product discretizations"
)
# Only want the horizontal cell
cell, _ = Vp.ufl_element().cell()._cells
DG = FiniteElement("DG", cell, deg)
CG = FiniteElement("CG", interval, 1)
Vv_tr_element = TensorProductElement(DG, CG)
Vv_tr = FunctionSpace(mesh, Vv_tr_element)
# Break the spaces
broken_elements = MixedElement(
[BrokenElement(Vi.ufl_element()) for Vi in V])
V_d = FunctionSpace(mesh, broken_elements)
# Set up relevant functions
self.broken_solution = Function(V_d)
self.broken_residual = Function(V_d)
self.trace_solution = Function(Vv_tr)
self.unbroken_solution = Function(V)
self.unbroken_residual = Function(V)
# Set up transfer kernels to and from the broken velocity space
# NOTE: Since this snippet of code is used in a couple places in
# in Gusto, might be worth creating a utility function that is
# is importable and just called where needed.
shapes = {
"i": Vv.finat_element.space_dimension(),
"j": np.prod(Vv.shape, dtype=int)
}
weight_kernel = """
for (int i=0; i<{i}; ++i)
for (int j=0; j<{j}; ++j)
w[i*{j} + j] += 1.0;
""".format(**shapes)
self.weight = Function(Vv)
par_loop(weight_kernel, dx, {"w": (self.weight, INC)})
# Averaging kernel
self.average_kernel = """
for (int i=0; i<{i}; ++i)
for (int j=0; j<{j}; ++j)
vec_out[i*{j} + j] += vec_in[i*{j} + j]/w[i*{j} + j];
""".format(**shapes)
# 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(Vv_tr)
n = FacetNormal(mesh)
sigma = TrialFunctions(V_d)[self.vidx]
# Again, assumes tensor product structure. Why use this if you
# don't have some form of vertical extrusion?
Kform = gammar('+') * jump(sigma, n=n) * dS_h
# Here we deal with boundary conditions
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, int))
# separate out the top and bottom bcs
extruded_neumann_subdomains = neumann_subdomains & {
"top", "bottom"
}
neumann_subdomains = neumann_subdomains - extruded_neumann_subdomains
integrand = gammar * dot(sigma, n)
measures = []
trace_subdomains = []
for subdomain in sorted(extruded_neumann_subdomains):
measures.append({"top": ds_t, "bottom": ds_b}[subdomain])
trace_subdomains.extend(
sorted({"top", "bottom"} - extruded_neumann_subdomains))
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
else:
trace_subdomains = ["top", "bottom"]
trace_bcs = [
DirichletBC(Vv_tr, 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(Vv_tr)
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)
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)
self._assemble_S()
self.S.force_evaluation()
Smat = self.S.petscmat
nullspace = self.ctx.appctx.get("vert_trace_nullspace", None)
if nullspace is not None:
nsp = nullspace(Vv_tr)
Smat.setNullSpace(nsp.nullspace(comm=pc.comm))
# Set up the KSP for the system of Lagrange multipliers
trace_ksp = PETSc.KSP().create(comm=pc.comm)
trace_ksp.setOptionsPrefix(prefix)
trace_ksp.setOperators(Smat)
trace_ksp.setUp()
trace_ksp.setFromOptions()
self.trace_ksp = trace_ksp
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)
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, assemble)
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.cxt = P.getPythonContext()
if not isinstance(self.cxt, ImplicitMatrixContext):
raise ValueError("The python context must be an ImplicitMatrixContext")
test, trial = self.cxt.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)
# Set up the KSP for the hdiv residual projection
hdiv_mass_ksp = PETSc.KSP().create(comm=pc.comm)
hdiv_mass_ksp.setOptionsPrefix(prefix + "hdiv_residual_")
# HDiv mass operator
p = TrialFunction(V[self.vidx])
q = TestFunction(V[self.vidx])
mass = ufl.dot(p, q)*ufl.dx
# TODO: Bcs?
M = assemble(mass, bcs=None, form_compiler_parameters=self.cxt.fc_params)
M.force_evaluation()
Mmat = M.petscmat
hdiv_mass_ksp.setOperators(Mmat)
hdiv_mass_ksp.setUp()
hdiv_mass_ksp.setFromOptions()
self.hdiv_mass_ksp = hdiv_mass_ksp
# Storing the result of A.inv * r, where A is the HDiv
# mass matrix and r is the HDiv residual
self._primal_r = Function(V[self.vidx])
tau = TestFunction(V_d[self.vidx])
self._assemble_broken_r = create_assembly_callable(
ufl.dot(self._primal_r, tau)*ufl.dx,
tensor=self.broken_residual.split()[self.vidx],
form_compiler_parameters=self.cxt.fc_params)
# 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.cxt.a, arg_map))
gammar = TestFunction(TraceSpace)
n = ufl.FacetNormal(mesh)
sigma = TrialFunctions(V_d)[self.vidx]
# We zero out the contribution of the trace variables on the exterior
# boundary. Extruded cells will have both horizontal and vertical
# facets
if mesh.cell_set._extruded:
trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), "on_boundary"),
DirichletBC(TraceSpace, Constant(0.0), "bottom"),
DirichletBC(TraceSpace, Constant(0.0), "top")]
K = Tensor(gammar('+') * ufl.dot(sigma, n) * ufl.dS_h +
gammar('+') * ufl.dot(sigma, n) * ufl.dS_v)
else:
trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), "on_boundary")]
K = Tensor(gammar('+') * ufl.dot(sigma, n) * ufl.dS)
# If boundary conditions are contained in the ImplicitMatrixContext:
if self.cxt.row_bcs:
raise NotImplementedError("Strong BCs not currently handled. Try imposing them weakly.")
# Assemble the Schur complement operator and right-hand side
self.schur_rhs = Function(TraceSpace)
self._assemble_Srhs = create_assembly_callable(
K * Atilde.inv * self.broken_residual,
tensor=self.schur_rhs,
form_compiler_parameters=self.cxt.fc_params)
schur_comp = K * Atilde.inv * K.T
self.S = allocate_matrix(schur_comp, bcs=trace_bcs,
form_compiler_parameters=self.cxt.fc_params)
self._assemble_S = create_assembly_callable(schur_comp,
tensor=self.S,
bcs=trace_bcs,
form_compiler_parameters=self.cxt.fc_params)
self._assemble_S()
self.S.force_evaluation()
Smat = self.S.petscmat
# Nullspace for the multiplier problem
nullspace = create_schur_nullspace(P, -K * Atilde,
V, V_d, TraceSpace,
pc.comm)
if nullspace:
Smat.setNullSpace(nullspace)
# Set up the KSP for the system of Lagrange multipliers
trace_ksp = PETSc.KSP().create(comm=pc.comm)
trace_ksp.setOptionsPrefix(prefix)
trace_ksp.setOperators(Smat)
trace_ksp.setUp()
trace_ksp.setFromOptions()
self.trace_ksp = trace_ksp
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)
# NOTE: The projection stage *might* be replaced by a Fortin
# operator. We may want to allow the user to specify if they
# wish to use a Fortin operator over a projection, or vice-versa.
# In a future add-on, we can add a switch which chooses either
# the Fortin reconstruction or the usual KSP projection.
# Set up the projection KSP
hdiv_projection_ksp = PETSc.KSP().create(comm=pc.comm)
hdiv_projection_ksp.setOptionsPrefix(prefix + 'hdiv_projection_')
# Reuse the mass operator from the hdiv_mass_ksp
hdiv_projection_ksp.setOperators(Mmat)
# Construct the RHS for the projection stage
self._projection_rhs = Function(V[self.vidx])
self._assemble_projection_rhs = create_assembly_callable(
ufl.dot(self.broken_solution.split()[self.vidx], q)*ufl.dx,
tensor=self._projection_rhs,
form_compiler_parameters=self.cxt.fc_params)
# Finalize ksp setup
hdiv_projection_ksp.setUp()
hdiv_projection_ksp.setFromOptions()
self.hdiv_projection_ksp = hdiv_projection_ksp
def do_setup(mesh,
pc,
degree=1,
theta=0.5,
dt=5.0e-06,
lmbda=1.0e-02,
inner_ksp='preonly',
maxit=1,
params={}):
V = FunctionSpace(mesh, "Lagrange", degree)
ME = V * V
# Define trial and test functions
du = TrialFunction(ME)
q, v = TestFunctions(ME)
# Define functions
u = Function(ME) # current solution
u0 = Function(ME) # solution from previous converged step
# Split mixed functions
dc, dmu = split(du)
c, mu = split(u)
c0, mu0 = split(u0)
# Create intial conditions and interpolate
init_code = "A[0] = 0.63 + 0.02*(0.5 - (double)random()/RAND_MAX);"
user_code = """int __rank;
MPI_Comm_rank(MPI_COMM_WORLD, &__rank);
srandom(2 + __rank);"""
par_loop(init_code,
direct, {'A': (u[0], WRITE)},
headers=["#include <stdlib.h>"],
user_code=user_code)
u.dat.data_ro
# Compute the chemical potential df/dc
c = variable(c)
f = 100 * c**2 * (1 - c)**2
dfdc = diff(f, c)
mu_mid = (1.0 - theta) * mu0 + theta * mu
# Weak statement of the equations
F0 = c * q * dx - c0 * q * dx + dt * dot(grad(mu_mid), grad(q)) * dx
F1 = mu * v * dx - dfdc * v * dx - lmbda * dot(grad(c), grad(v)) * dx
F = F0 + F1
# Compute directional derivative about u in the direction of du (Jacobian)
J = derivative(F, u, du)
problem = NonlinearVariationalProblem(F, u, J=J)
solver = NonlinearVariationalSolver(problem, solver_parameters=params)
if pc in ['fieldsplit', 'ilu']:
sigma = 100
# PC for the Schur complement solve
trial = TrialFunction(V)
test = TestFunction(V)
mass = assemble(inner(trial, test) * dx).M.handle
a = 1
c = (dt * lmbda) / (1 + dt * sigma)
hats = assemble(
sqrt(a) * inner(trial, test) * dx +
sqrt(c) * inner(grad(trial), grad(test)) * dx).M.handle
from firedrake.petsc import PETSc
ksp_hats = PETSc.KSP()
ksp_hats.create()
ksp_hats.setOperators(hats)
opts = PETSc.Options()
opts['ksp_type'] = inner_ksp
opts['ksp_max_it'] = maxit
opts['pc_type'] = 'hypre'
ksp_hats.setFromOptions()
class SchurInv(object):
def mult(self, mat, x, y):
tmp1 = y.duplicate()
tmp2 = y.duplicate()
ksp_hats.solve(x, tmp1)
mass.mult(tmp1, tmp2)
ksp_hats.solve(tmp2, y)
pc_schur = PETSc.Mat()
pc_schur.createPython(mass.getSizes(), SchurInv())
pc_schur.setUp()
pc = solver.snes.ksp.pc
pc.setFieldSplitSchurPreType(PETSc.PC.SchurPreType.USER, pc_schur)
return u, u0, solver
