def _split_mixed_operator(self):
split_mixed_op = dict(split_form(self.Atilde.form))
id0, id1 = (self.vidx, self.pidx)
A00 = Tensor(split_mixed_op[(id0, id0)])
A01 = Tensor(split_mixed_op[(id0, id1)])
A10 = Tensor(split_mixed_op[(id1, id0)])
A11 = Tensor(split_mixed_op[(id1, id1)])
self.list_split_mixed_ops = [A00, A01, A10, A11]
split_trace_op = dict(split_form(self.K.form))
K0 = Tensor(split_trace_op[(0, id0)])
K1 = Tensor(split_trace_op[(0, id1)])
self.list_split_trace_ops = [K0, K1]
def test_matrix_subblocks(mesh):
if mesh.ufl_cell() == quadrilateral:
U = FunctionSpace(mesh, "RTCF", 1)
else:
U = FunctionSpace(mesh, "RT", 1)
V = FunctionSpace(mesh, "DG", 0)
T = FunctionSpace(mesh, "HDiv Trace", 0)
n = FacetNormal(mesh)
W = U * V * T
u, p, lambdar = TrialFunctions(W)
w, q, gammar = TestFunctions(W)
A = Tensor(inner(u, w)*dx + p*q*dx - div(w)*p*dx + q*div(u)*dx
+ lambdar('+')*jump(w, n=n)*dS + gammar('+')*jump(u, n=n)*dS
+ lambdar*gammar*ds)
# Test individual blocks
indices = [(0, 0), (0, 1), (1, 0), (1, 1), (1, 2), (2, 1), (2, 2)]
refs = dict(split_form(A.form))
_A = A.blocks
for x, y in indices:
ref = assemble(refs[x, y]).M.values
block = _A[x, y]
assert np.allclose(assemble(block).M.values, ref, rtol=1e-14)
# Mixed blocks
A0101 = _A[:2, :2]
A1212 = _A[1:3, 1:3]
_A0101 = A0101.blocks
_A1212 = A1212.blocks
# Block of blocks
A0101_00 = _A0101[0, 0]
A0101_11 = _A0101[1, 1]
A0101_01 = _A0101[0, 1]
A0101_10 = _A0101[1, 0]
A1212_00 = _A1212[0, 0]
A1212_11 = _A1212[1, 1]
A1212_01 = _A1212[0, 1]
A1212_10 = _A1212[1, 0]
items = [(A0101_00, refs[(0, 0)]),
(A0101_11, refs[(1, 1)]),
(A0101_01, refs[(0, 1)]),
(A0101_10, refs[(1, 0)]),
(A1212_00, refs[(1, 1)]),
(A1212_11, refs[(2, 2)]),
(A1212_01, refs[(1, 2)]),
(A1212_10, refs[(2, 1)])]
# Test assembly of blocks of mixed blocks
for tensor, form in items:
ref = assemble(form).M.values
assert np.allclose(assemble(tensor).M.values, ref, rtol=1e-14)
def test_matrix_subblocks(mesh):
if mesh.ufl_cell() == quadrilateral:
U = FunctionSpace(mesh, "RTCF", 1)
else:
U = FunctionSpace(mesh, "RT", 1)
V = FunctionSpace(mesh, "DG", 0)
T = FunctionSpace(mesh, "HDiv Trace", 0)
n = FacetNormal(mesh)
W = U * V * T
u, p, lambdar = TrialFunctions(W)
w, q, gammar = TestFunctions(W)
A = Tensor(inner(u, w)*dx + p*q*dx - div(w)*p*dx + q*div(u)*dx +
lambdar('+')*jump(w, n=n)*dS + gammar('+')*jump(u, n=n)*dS +
lambdar*gammar*ds)
# Test individual blocks
indices = [(0, 0), (0, 1), (1, 0), (1, 1), (1, 2), (2, 1), (2, 2)]
refs = dict(split_form(A.form))
_A = A.blocks
for x, y in indices:
ref = assemble(refs[x, y]).M.values
block = _A[x, y]
assert np.allclose(assemble(block).M.values, ref, rtol=1e-14)
# Mixed blocks
A0101 = _A[:2, :2]
A1212 = _A[1:3, 1:3]
_A0101 = A0101.blocks
_A1212 = A1212.blocks
# Block of blocks
A0101_00 = _A0101[0, 0]
A0101_11 = _A0101[1, 1]
A0101_01 = _A0101[0, 1]
A0101_10 = _A0101[1, 0]
A1212_00 = _A1212[0, 0]
A1212_11 = _A1212[1, 1]
A1212_01 = _A1212[0, 1]
A1212_10 = _A1212[1, 0]
items = [(A0101_00, refs[(0, 0)]),
(A0101_11, refs[(1, 1)]),
(A0101_01, refs[(0, 1)]),
(A0101_10, refs[(1, 0)]),
(A1212_00, refs[(1, 1)]),
(A1212_11, refs[(2, 2)]),
(A1212_01, refs[(1, 2)]),
(A1212_10, refs[(2, 1)])]
# Test assembly of blocks of mixed blocks
for tensor, form in items:
ref = assemble(form).M.values
assert np.allclose(assemble(tensor).M.values, ref, rtol=1e-14)
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 compile_form(form, name, parameters=None, inverse=False, split=True):
"""Compile a form using TSFC.
:arg form: the :class:`~ufl.classes.Form` to compile.
:arg name: a prefix for the generated kernel functions.
:arg parameters: optional dict of parameters to pass to the form
compiler. If not provided, parameters are read from the
``form_compiler`` slot of the Firedrake
:data:`~.parameters` dictionary (which see).
:arg inverse: If True then assemble the inverse of the local tensor.
:arg split: If ``False``, then don't split mixed forms.
Returns a tuple of tuples of
(index, integral type, subdomain id, coordinates, coefficients, needs_orientations, :class:`Kernels <pyop2.op2.Kernel>`).
``needs_orientations`` indicates whether the form requires cell
orientation information (for correctly pulling back to reference
elements on embedded manifolds).
The coordinates are extracted from the domain of the integral (a
:func:`~.Mesh`)
"""
# Check that we get a Form
if not isinstance(form, Form):
raise RuntimeError("Unable to convert object to a UFL form: %s" %
repr(form))
if parameters is None:
parameters = default_parameters["form_compiler"].copy()
else:
# Override defaults with user-specified values
_ = parameters
parameters = default_parameters["form_compiler"].copy()
parameters.update(_)
# We stash the compiled kernels on the form so we don't have to recompile
# if we assemble the same form again with the same optimisations
cache = form._cache.setdefault("firedrake_kernels", {})
def tuplify(params):
return tuple((k, params[k]) for k in sorted(params))
key = (tuplify(default_parameters["coffee"]), name, tuplify(parameters),
split)
try:
return cache[key]
except KeyError:
pass
kernels = []
# A map from all form coefficients to their number.
coefficient_numbers = dict(
(c, n) for (n, c) in enumerate(form.coefficients()))
if split:
iterable = split_form(form)
else:
iterable = ([(0, ) * len(form.arguments()), form], )
for idx, f in iterable:
f = _real_mangle(f)
# Map local coefficient numbers (as seen inside the
# compiler) to the global coefficient numbers
number_map = dict((n, coefficient_numbers[c])
for (n, c) in enumerate(f.coefficients()))
kinfos = TSFCKernel(f, name + "".join(map(str, idx)), parameters,
number_map).kernels
for kinfo in kinfos:
kernels.append(SplitKernel(idx, kinfo))
kernels = tuple(kernels)
return cache.setdefault(key, kernels)
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. 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 compile_form(form, name, parameters=None, inverse=False):
"""Compile a form using TSFC.
:arg form: the :class:`~ufl.classes.Form` to compile.
:arg name: a prefix for the generated kernel functions.
:arg parameters: optional dict of parameters to pass to the form
compiler. If not provided, parameters are read from the
``form_compiler`` slot of the Firedrake
:data:`~.parameters` dictionary (which see).
:arg inverse: If True then assemble the inverse of the local tensor.
Returns a tuple of tuples of
(index, integral type, subdomain id, coordinates, coefficients, needs_orientations, :class:`Kernels <pyop2.op2.Kernel>`).
``needs_orientations`` indicates whether the form requires cell
orientation information (for correctly pulling back to reference
elements on embedded manifolds).
The coordinates are extracted from the domain of the integral (a
:func:`~.Mesh`)
"""
# Check that we get a Form
if not isinstance(form, Form):
raise RuntimeError("Unable to convert object to a UFL form: %s" %
repr(form))
if parameters is None:
parameters = default_parameters["form_compiler"].copy()
else:
# Override defaults with user-specified values
_ = parameters
parameters = default_parameters["form_compiler"].copy()
parameters.update(_)
# We stash the compiled kernels on the form so we don't have to recompile
# if we assemble the same form again with the same optimisations
if "firedrake_kernels" in form._cache:
# Save both kernels and TSFC params so we can tell if this
# cached version is valid (the TSFC parameters might have changed)
kernels, coffee_params, old_name, params = form._cache[
"firedrake_kernels"]
if coffee_params == default_parameters["coffee"] and \
name == old_name and \
params == parameters:
return kernels
kernels = []
# A map from all form coefficients to their number.
coefficient_numbers = dict(
(c, n) for (n, c) in enumerate(form.coefficients()))
for idx, f in split_form(form):
f = _real_mangle(f)
# Map local coefficient numbers (as seen inside the
# compiler) to the global coefficient numbers
number_map = dict((n, coefficient_numbers[c])
for (n, c) in enumerate(f.coefficients()))
kinfos = TSFCKernel(f, name + "".join(map(str, idx)), parameters,
number_map).kernels
for kinfo in kinfos:
kernels.append(SplitKernel(idx, kinfo))
kernels = tuple(kernels)
form._cache["firedrake_kernels"] = (kernels,
default_parameters["coffee"].copy(),
name, parameters)
return kernels
def compile_form(form, name, parameters=None, inverse=False):
"""Compile a form using TSFC.
:arg form: the :class:`~ufl.classes.Form` to compile.
:arg name: a prefix for the generated kernel functions.
:arg parameters: optional dict of parameters to pass to the form
compiler. If not provided, parameters are read from the
``form_compiler`` slot of the Firedrake
:data:`~.parameters` dictionary (which see).
:arg inverse: If True then assemble the inverse of the local tensor.
Returns a tuple of tuples of
(index, integral type, subdomain id, coordinates, coefficients, needs_orientations, :class:`Kernels <pyop2.op2.Kernel>`).
``needs_orientations`` indicates whether the form requires cell
orientation information (for correctly pulling back to reference
elements on embedded manifolds).
The coordinates are extracted from the domain of the integral (a
:func:`~.Mesh`)
"""
# Check that we get a Form
if not isinstance(form, Form):
raise RuntimeError("Unable to convert object to a UFL form: %s" % repr(form))
if parameters is None:
parameters = default_parameters["form_compiler"].copy()
else:
# Override defaults with user-specified values
_ = parameters
parameters = default_parameters["form_compiler"].copy()
parameters.update(_)
# We stash the compiled kernels on the form so we don't have to recompile
# if we assemble the same form again with the same optimisations
if "firedrake_kernels" in form._cache:
# Save both kernels and TSFC params so we can tell if this
# cached version is valid (the TSFC parameters might have changed)
kernels, coffee_params, old_name, params = form._cache["firedrake_kernels"]
if coffee_params == default_parameters["coffee"] and \
name == old_name and \
params == parameters:
return kernels
kernels = []
# A map from all form coefficients to their number.
coefficient_numbers = dict((c, n)
for (n, c) in enumerate(form.coefficients()))
for idx, f in split_form(form):
# Map local coefficient numbers (as seen inside the
# compiler) to the global coefficient numbers
number_map = dict((n, coefficient_numbers[c])
for (n, c) in enumerate(f.coefficients()))
kinfos = TSFCKernel(f, name + "".join(map(str, idx)), parameters,
number_map).kernels
for kinfo in kinfos:
kernels.append(SplitKernel(idx, kinfo))
kernels = tuple(kernels)
form._cache["firedrake_kernels"] = (kernels, default_parameters["coffee"].copy(),
name, parameters)
return kernels
def compile_form(form, name, parameters=None, inverse=False, split=True, interface=None, coffee=False):
"""Compile a form using TSFC.
:arg form: the :class:`~ufl.classes.Form` to compile.
:arg name: a prefix for the generated kernel functions.
:arg parameters: optional dict of parameters to pass to the form
compiler. If not provided, parameters are read from the
``form_compiler`` slot of the Firedrake
:data:`~.parameters` dictionary (which see).
:arg inverse: If True then assemble the inverse of the local tensor.
:arg split: If ``False``, then don't split mixed forms.
:arg coffee: compile coffee kernel instead of loopy kernel
Returns a tuple of tuples of
(index, integral type, subdomain id, coordinates, coefficients, needs_orientations, :class:`Kernels <pyop2.op2.Kernel>`).
``needs_orientations`` indicates whether the form requires cell
orientation information (for correctly pulling back to reference
elements on embedded manifolds).
The coordinates are extracted from the domain of the integral (a
:func:`~.Mesh`)
"""
# Check that we get a Form
if not isinstance(form, Form):
raise RuntimeError("Unable to convert object to a UFL form: %s" % repr(form))
if parameters is None:
parameters = default_parameters["form_compiler"].copy()
else:
# Override defaults with user-specified values
_ = parameters
parameters = default_parameters["form_compiler"].copy()
parameters.update(_)
# We stash the compiled kernels on the form so we don't have to recompile
# if we assemble the same form again with the same optimisations
cache = form._cache.setdefault("firedrake_kernels", {})
def tuplify(params):
return tuple((k, params[k]) for k in sorted(params))
key = (tuplify(default_parameters["coffee"]), name, tuplify(parameters), split)
try:
return cache[key]
except KeyError:
pass
kernels = []
# A map from all form coefficients to their number.
coefficient_numbers = dict((c, n)
for (n, c) in enumerate(form.coefficients()))
if split:
iterable = split_form(form)
else:
iterable = ([(0, )*len(form.arguments()), form], )
for idx, f in iterable:
f = _real_mangle(f)
# Map local coefficient numbers (as seen inside the
# compiler) to the global coefficient numbers
number_map = dict((n, coefficient_numbers[c])
for (n, c) in enumerate(f.coefficients()))
kinfos = TSFCKernel(f, name + "".join(map(str, idx)), parameters,
number_map, interface, coffee).kernels
for kinfo in kinfos:
kernels.append(SplitKernel(idx, kinfo))
kernels = tuple(kernels)
return cache.setdefault(key, kernels)
def setup(self, equation, uadv=None, apply_bcs=True, *active_labels):
self.residual = equation.residual
if self.field_name is not None:
self.idx = equation.field_names.index(self.field_name)
self.fs = self.state.fields(self.field_name).function_space()
self.residual = self.residual.label_map(
lambda t: t.get(prognostic) == self.field_name,
lambda t: Term(
split_form(t.form)[self.idx].form,
t.labels),
drop)
bcs = equation.bcs[self.field_name]
else:
self.field_name = equation.field_name
self.fs = equation.function_space
self.idx = None
if type(self.fs.ufl_element()) is MixedElement:
bcs = [bc for _, bcs in equation.bcs.items() for bc in bcs]
else:
bcs = equation.bcs[self.field_name]
if len(active_labels) > 0:
self.residual = self.residual.label_map(
lambda t: any(t.has_label(time_derivative, *active_labels)),
map_if_false=drop)
options = self.options
# -------------------------------------------------------------------- #
# Routines relating to transport
# -------------------------------------------------------------------- #
if hasattr(self.options, 'ibp'):
self.replace_transport_term()
self.replace_transporting_velocity(uadv)
# -------------------------------------------------------------------- #
# Wrappers for embedded / recovery methods
# -------------------------------------------------------------------- #
if self.discretisation_option in ["embedded_dg", "recovered"]:
# construct the embedding space if not specified
if options.embedding_space is None:
V_elt = BrokenElement(self.fs.ufl_element())
self.fs = FunctionSpace(self.state.mesh, V_elt)
else:
self.fs = options.embedding_space
self.xdg_in = Function(self.fs)
self.xdg_out = Function(self.fs)
if self.idx is None:
self.x_projected = Function(equation.function_space)
else:
self.x_projected = Function(self.state.fields(self.field_name).function_space())
new_test = TestFunction(self.fs)
parameters = {'ksp_type': 'cg',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}
# -------------------------------------------------------------------- #
# Make boundary conditions
# -------------------------------------------------------------------- #
if not apply_bcs:
self.bcs = None
elif self.discretisation_option in ["embedded_dg", "recovered"]:
# Transfer boundary conditions onto test function space
self.bcs = [DirichletBC(self.fs, bc.function_arg, bc.sub_domain) for bc in bcs]
else:
self.bcs = bcs
# -------------------------------------------------------------------- #
# Modify test function for SUPG methods
# -------------------------------------------------------------------- #
if self.discretisation_option == "supg":
# construct tau, if it is not specified
dim = self.state.mesh.topological_dimension()
if options.tau is not None:
# if tau is provided, check that is has the right size
tau = options.tau
assert as_ufl(tau).ufl_shape == (dim, dim), "Provided tau has incorrect shape!"
else:
# create tuple of default values of size dim
default_vals = [options.default*self.dt]*dim
# check for directions is which the space is discontinuous
# so that we don't apply supg in that direction
if is_cg(self.fs):
vals = default_vals
else:
space = self.fs.ufl_element().sobolev_space()
if space.name in ["HDiv", "DirectionalH"]:
vals = [default_vals[i] if space[i].name == "H1"
else 0. for i in range(dim)]
else:
raise ValueError("I don't know what to do with space %s" % space)
tau = Constant(tuple([
tuple(
[vals[j] if i == j else 0. for i, v in enumerate(vals)]
) for j in range(dim)])
)
self.solver_parameters = {'ksp_type': 'gmres',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}
test = TestFunction(self.fs)
new_test = test + dot(dot(uadv, tau), grad(test))
if self.discretisation_option is not None:
# replace the original test function with one defined on
# the embedding space, as this is the space where the
# the problem will be solved
self.residual = self.residual.label_map(
all_terms,
map_if_true=replace_test_function(new_test))
if self.discretisation_option == "embedded_dg":
if self.limiter is None:
self.x_out_projector = Projector(self.xdg_out, self.x_projected,
solver_parameters=parameters)
else:
self.x_out_projector = Recoverer(self.xdg_out, self.x_projected)
if self.discretisation_option == "recovered":
# set up the necessary functions
self.x_in = Function(self.state.fields(self.field_name).function_space())
x_rec = Function(options.recovered_space)
x_brok = Function(options.broken_space)
# set up interpolators and projectors
self.x_rec_projector = Recoverer(self.x_in, x_rec, VDG=self.fs, boundary_method=options.boundary_method) # recovered function
self.x_brok_projector = Projector(x_rec, x_brok) # function projected back
self.xdg_interpolator = Interpolator(self.x_in + x_rec - x_brok, self.xdg_in)
if self.limiter is not None:
self.x_brok_interpolator = Interpolator(self.xdg_out, x_brok)
self.x_out_projector = Recoverer(x_brok, self.x_projected)
else:
self.x_out_projector = Projector(self.xdg_out, self.x_projected)
# setup required functions
self.dq = Function(self.fs)
self.q1 = Function(self.fs)
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]
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. 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, int))
# 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)
self._assemble_S()
self.S.force_evaluation()
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)
