def compute_form_action(form, coefficient):
"""Compute the action of a form on a Coefficient.
This works simply by replacing the last Argument
with a Coefficient on the same function space (element).
The form returned will thus have one Argument less
and one additional Coefficient at the end if no
Coefficient has been provided.
"""
# TODO: Check whatever makes sense for coefficient
# Extract all arguments
arguments = form.arguments()
parts = [arg.part() for arg in arguments]
if set(parts) - {None}:
error("compute_form_action cannot handle parts.")
# Pick last argument (will be replaced)
u = arguments[-1]
fs = u.ufl_function_space()
if coefficient is None:
coefficient = Coefficient(fs)
elif coefficient.ufl_function_space() != fs:
debug("Computing action of form on a coefficient in a different function space.")
return replace(form, {u: coefficient})
def compute_energy_norm(form, coefficient):
"""Compute the a-norm of a Coefficient given a form a.
This works simply by replacing the two Arguments
with a Coefficient on the same function space (element).
The Form returned will thus be a functional with no
Arguments, and one additional Coefficient at the
end if no coefficient has been provided.
"""
arguments = form.arguments()
parts = [arg.part() for arg in arguments]
if set(parts) - {None}:
error("compute_energy_norm cannot handle parts.")
if len(arguments) != 2:
error("Expecting bilinear form.")
v, u = arguments
U = u.ufl_function_space()
V = v.ufl_function_space()
if U != V:
error("Expecting equal finite elements for test and trial functions, got '%s' and '%s'." % (U, V))
if coefficient is None:
coefficient = Coefficient(V)
else:
if coefficient.ufl_function_space() != U:
error("Trying to compute action of form on a "
"coefficient in an incompatible element space.")
return replace(form, {u: coefficient, v: coefficient})
def compute_form_action(form, coefficient):
"""Compute the action of a form on a Coefficient.
This works simply by replacing the last Argument
with a Coefficient on the same function space (element).
The form returned will thus have one Argument less
and one additional Coefficient at the end if no
Coefficient has been provided.
"""
# TODO: Check whatever makes sense for coefficient
# Extract all arguments
arguments = form.arguments()
parts = [arg.part() for arg in arguments]
if set(parts) - {None}:
error("compute_form_action cannot handle parts.")
# Pick last argument (will be replaced)
u = arguments[-1]
fs = u.ufl_function_space()
if coefficient is None:
coefficient = Coefficient(fs)
elif coefficient.ufl_function_space() != fs:
debug("Computing action of form on a coefficient in a different function space.")
return replace(form, {u: coefficient})
def compute_form_action(form, coefficient):
"""Compute the action of a form on a Coefficient.
This works simply by replacing the last Argument
with a Coefficient on the same function space (element).
The form returned will thus have one Argument less
and one additional Coefficient at the end if no
Coefficient has been provided.
"""
# TODO: Check whatever makes sense for coefficient
# Extract all arguments
arguments = extract_arguments(form)
# Pick last argument (will be replaced)
u = arguments[-1]
e = u.element()
if coefficient is None:
coefficient = Coefficient(e)
else:
#ufl_assert(coefficient.element() == e, \
if coefficient.element() != e:
debug("Computing action of form on a coefficient in a different element space.")
return replace(form, { u: coefficient })
def compute_form_action(form, coefficient):
"""Compute the action of a form on a Coefficient.
This works simply by replacing the last Argument
with a Coefficient on the same function space (element).
The form returned will thus have one Argument less
and one additional Coefficient at the end if no
Coefficient has been provided.
"""
# TODO: Check whatever makes sense for coefficient
# Extract all arguments
arguments = extract_arguments(form)
# Pick last argument (will be replaced)
u = arguments[-1]
e = u.element()
if coefficient is None:
coefficient = Coefficient(e)
else:
#ufl_assert(coefficient.element() == e, \
if coefficient.element() != e:
debug(
"Computing action of form on a coefficient in a different element space."
)
return replace(form, {u: coefficient})
def compute_energy_norm(form, coefficient):
"""Compute the a-norm of a Coefficient given a form a.
This works simply by replacing the two Arguments
with a Coefficient on the same function space (element).
The Form returned will thus be a functional with no
Arguments, and one additional Coefficient at the
end if no coefficient has been provided.
"""
arguments = form.arguments()
parts = [arg.part() for arg in arguments]
if set(parts) - {None}:
error("compute_energy_norm cannot handle parts.")
if len(arguments) != 2:
error("Expecting bilinear form.")
v, u = arguments
U = u.ufl_function_space()
V = v.ufl_function_space()
if U != V:
error("Expecting equal finite elements for test and trial functions, got '%s' and '%s'." % (U, V))
if coefficient is None:
coefficient = Coefficient(V)
else:
if coefficient.ufl_function_space() != U:
error("Trying to compute action of form on a "
"coefficient in an incompatible element space.")
return replace(form, {u: coefficient, v: coefficient})
def compute_energy_norm(form, coefficient):
"""Compute the a-norm of a Coefficient given a form a.
This works simply by replacing the two Arguments
with a Coefficient on the same function space (element).
The Form returned will thus be a functional with no
Arguments, and one additional Coefficient at the
end if no coefficient has been provided.
"""
arguments = extract_arguments(form)
ufl_assert(len(arguments) == 2, "Expecting bilinear form.")
v, u = arguments
e = u.element()
e2 = v.element()
ufl_assert(
e == e2,
"Expecting equal finite elements for test and trial functions, got '%s' and '%s'."
% (str(e), str(e2)))
if coefficient is None:
coefficient = Coefficient(e)
else:
ufl_assert(coefficient.element() == e, \
"Trying to compute action of form on a "\
"coefficient in an incompatible element space.")
return replace(form, {u: coefficient, v: coefficient})
def compute_form_adjoint(form, reordered_arguments=None):
"""Compute the adjoint of a bilinear form.
This works simply by swapping the number and part of the two arguments,
but keeping their elements and places in the integrand expressions.
"""
arguments = form.arguments()
parts = [arg.part() for arg in arguments]
if set(parts) - {None}:
error("compute_form_adjoint cannot handle parts.")
if len(arguments) != 2:
error("Expecting bilinear form.")
v, u = arguments
if v.number() >= u.number():
error("Mistaken assumption in code!")
if reordered_arguments is None:
reordered_u = Argument(u.ufl_function_space(),
number=v.number(),
part=v.part())
reordered_v = Argument(v.ufl_function_space(),
number=u.number(),
part=u.part())
else:
reordered_u, reordered_v = reordered_arguments
if reordered_u.number() >= reordered_v.number():
error("Ordering of new arguments is the same as the old arguments!")
if reordered_u.part() != v.part():
error("Ordering of new arguments is the same as the old arguments!")
if reordered_v.part() != u.part():
error("Ordering of new arguments is the same as the old arguments!")
if reordered_u.ufl_function_space() != u.ufl_function_space():
error(
"Element mismatch between new and old arguments (trial functions)."
)
if reordered_v.ufl_function_space() != v.ufl_function_space():
error(
"Element mismatch between new and old arguments (test functions).")
return map_integrands(Conj, replace(form, {
v: reordered_v,
u: reordered_u
}))
def replace_terminal_data(o, mapping):
"""Return a new form where the terminals have been replaced using the
provided mapping.
:arg o: The object to have its terminals replaced. This must either be a
:class:`~.Form` or :class:`~.Integral`.
:arg mapping: A mapping suitable for reconstructing the form such as the one
returned by :func:`strip_terminal_data`.
:returns: The new form.
"""
if isinstance(o, Form):
return Form(
[replace_terminal_data(itg, mapping) for itg in o.integrals()])
elif isinstance(o, Integral):
expr_map, domain_map = mapping
integrand = replace(o.integrand(), expr_map)
return o.reconstruct(integrand, domain=domain_map[o.ufl_domain()])
else:
raise ValueError("Only Form or Integral inputs expected")
def compute_form_adjoint(form, reordered_arguments=None):
"""Compute the adjoint of a bilinear form.
This works simply by swapping the number and part of the two arguments,
but keeping their elements and places in the integrand expressions.
"""
arguments = form.arguments()
parts = [arg.part() for arg in arguments]
if set(parts) - {None}:
error("compute_form_adjoint cannot handle parts.")
if len(arguments) != 2:
error("Expecting bilinear form.")
v, u = arguments
if v.number() >= u.number():
error("Mistaken assumption in code!")
if reordered_arguments is None:
reordered_u = Argument(u.ufl_function_space(), number=v.number(),
part=v.part())
reordered_v = Argument(v.ufl_function_space(), number=u.number(),
part=u.part())
else:
reordered_u, reordered_v = reordered_arguments
if reordered_u.number() >= reordered_v.number():
error("Ordering of new arguments is the same as the old arguments!")
if reordered_u.part() != v.part():
error("Ordering of new arguments is the same as the old arguments!")
if reordered_v.part() != u.part():
error("Ordering of new arguments is the same as the old arguments!")
if reordered_u.ufl_function_space() != u.ufl_function_space():
error("Element mismatch between new and old arguments (trial functions).")
if reordered_v.ufl_function_space() != v.ufl_function_space():
error("Element mismatch between new and old arguments (test functions).")
return map_integrands(Conj, replace(form, {v: reordered_v, u: reordered_u}))
def compute_form_adjoint(form, reordered_arguments=None):
"""Compute the adjoint of a bilinear form.
This works simply by changing the ordering (count) of the two arguments.
"""
arguments = extract_arguments(form)
ufl_assert(len(arguments) == 2, "Expecting bilinear form.")
v, u = arguments
ufl_assert(v.count() < u.count(), "Mistaken assumption in code!")
if reordered_arguments is None:
reordered_arguments = (Argument(u.element()), Argument(v.element()))
reordered_u, reordered_v = reordered_arguments
ufl_assert(reordered_u.count() < reordered_v.count(),
"Ordering of new arguments is the same as the old arguments!")
ufl_assert(reordered_u.element() == u.element(),
"Element mismatch between new and old arguments (trial functions).")
ufl_assert(reordered_v.element() == v.element(),
"Element mismatch between new and old arguments (test functions).")
return replace(form, {v: reordered_v, u: reordered_u})
def compute_energy_norm(form, coefficient):
"""Compute the a-norm of a Coefficient given a form a.
This works simply by replacing the two Arguments
with a Coefficient on the same function space (element).
The Form returned will thus be a functional with no
Arguments, and one additional Coefficient at the
end if no coefficient has been provided.
"""
arguments = extract_arguments(form)
ufl_assert(len(arguments) == 2, "Expecting bilinear form.")
v, u = arguments
e = u.element()
e2 = v.element()
ufl_assert(e == e2, "Expecting equal finite elements for test and trial functions, got '%s' and '%s'." % (str(e), str(e2)))
if coefficient is None:
coefficient = Coefficient(e)
else:
ufl_assert(coefficient.element() == e, \
"Trying to compute action of form on a "\
"coefficient in an incompatible element space.")
return replace(form, { u: coefficient, v: coefficient })
def compute_form_adjoint(form, reordered_arguments=None):
"""Compute the adjoint of a bilinear form.
This works simply by changing the ordering (count) of the two arguments.
"""
arguments = extract_arguments(form)
ufl_assert(len(arguments) == 2, "Expecting bilinear form.")
v, u = arguments
ufl_assert(v.count() < u.count(), "Mistaken assumption in code!")
if reordered_arguments is None:
reordered_arguments = (Argument(u.element()), Argument(v.element()))
reordered_u, reordered_v = reordered_arguments
ufl_assert(reordered_u.count() < reordered_v.count(),
"Ordering of new arguments is the same as the old arguments!")
ufl_assert(
reordered_u.element() == u.element(),
"Element mismatch between new and old arguments (trial functions).")
ufl_assert(
reordered_v.element() == v.element(),
"Element mismatch between new and old arguments (test functions).")
return replace(form, {v: reordered_v, u: reordered_u})
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 preprocess_expression(expr, object_names=None, common_cell=None, element_mapping=None):
"""
Preprocess raw input expression to obtain expression metadata,
including a modified (preprocessed) expression more easily
manipulated by expression compilers. The original expression
is left untouched. Currently, the following transformations
are made to the preprocessed form:
expand_compounds (side effect of calling expand_derivatives)
expand_derivatives
renumber arguments and coefficients and apply evt. element mapping
"""
tic = Timer('preprocess_expression')
# Check that we get an expression
ufl_assert(isinstance(expr, Expr), "Expecting Expr.")
# Object names is empty if not given
object_names = object_names or {}
# Element mapping is empty if not given
element_mapping = element_mapping or {}
# Create empty expression data
expr_data = ExprData()
# Store original expression
expr_data.original_expr = expr
# Store name of expr if given, otherwise empty string
# such that automatic names can be assigned externally
expr_data.name = object_names.get(id(expr), "") # TODO: Or default to 'expr'?
# Extract common cell
common_cell = extract_common_cell(expr, common_cell)
# TODO: Split out expand_compounds from expand_derivatives
# Expand derivatives
tic('expand_derivatives')
expr = expand_derivatives(expr, common_cell.geometric_dimension())
# Renumber indices
#expr = renumber_indices(expr) # TODO: No longer needed?
# Replace arguments and coefficients with new renumbered objects
tic('extract_arguments_and_coefficients')
original_arguments, original_coefficients = \
extract_arguments_and_coefficients(expr)
tic('build_element_mapping')
element_mapping = build_element_mapping(element_mapping,
common_cell,
original_arguments,
original_coefficients)
tic('build_argument_replace_map')
replace_map, renumbered_arguments, renumbered_coefficients = \
build_argument_replace_map(original_arguments,
original_coefficients,
element_mapping)
expr = replace(expr, replace_map)
# Build mapping to original arguments and coefficients, which is
# useful if the original arguments have data attached to them
inv_replace_map = dict((w,v) for (v,w) in replace_map.iteritems())
original_arguments = [inv_replace_map[v] for v in renumbered_arguments]
original_coefficients = [inv_replace_map[w] for w in renumbered_coefficients]
# Store data extracted by preprocessing
expr_data.arguments = renumbered_arguments # TODO: Needed?
expr_data.coefficients = renumbered_coefficients # TODO: Needed?
expr_data.original_arguments = original_arguments
expr_data.original_coefficients = original_coefficients
expr_data.renumbered_arguments = renumbered_arguments
expr_data.renumbered_coefficients = renumbered_coefficients
tic('replace')
# Mappings from elements and functions (coefficients and arguments)
# that reside in expr to objects with canonical numbering as well as
# completed cells and elements
expr_data.element_replace_map = element_mapping
expr_data.function_replace_map = replace_map
# Store signature of form
tic('signature')
expr_data.signature = compute_expression_signature(expr, expr_data.function_replace_map)
# Store elements, sub elements and element map
tic('extract_elements')
expr_data.elements = tuple(f.element() for f in
chain(renumbered_arguments,
renumbered_coefficients))
expr_data.unique_elements = unique_tuple(expr_data.elements)
expr_data.sub_elements = extract_sub_elements(expr_data.elements)
expr_data.unique_sub_elements = unique_tuple(expr_data.sub_elements)
# Store element domains (NB! This is likely to change!)
# FIXME: DOMAINS: What is a sensible way to store domains for a expr?
expr_data.domains = tuple(sorted(set(element.domain()
for element in expr_data.unique_elements)))
# Store toplevel domains (NB! This is likely to change!)
# FIXME: DOMAINS: What is a sensible way to store domains for a expr?
expr_data.top_domains = tuple(sorted(set(domain.top_domain()
for domain in expr_data.domains)))
# Store common cell
expr_data.cell = common_cell
# Store data related to cell
expr_data.geometric_dimension = expr_data.cell.geometric_dimension()
expr_data.topological_dimension = expr_data.cell.topological_dimension()
# Store some useful dimensions
expr_data.rank = len(expr_data.arguments) # TODO: Is this useful for expr?
expr_data.num_coefficients = len(expr_data.coefficients)
# Store argument names # TODO: Is this useful for expr?
expr_data.argument_names = \
[object_names.get(id(expr_data.original_arguments[i]), "v%d" % i)
for i in range(expr_data.rank)]
# Store coefficient names
expr_data.coefficient_names = \
[object_names.get(id(expr_data.original_coefficients[i]), "w%d" % i)
for i in range(expr_data.num_coefficients)]
# Store preprocessed expression
expr_data.preprocessed_expr = expr
tic.end()
# A coarse profiling implementation
# TODO: Add counting of nodes
# TODO: Add memory usage
if preprocess_expression.enable_profiling:
print tic
return expr_data
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. 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 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)
def preprocess_expression(expr,
object_names=None,
common_cell=None,
element_mapping=None):
"""
Preprocess raw input expression to obtain expression metadata,
including a modified (preprocessed) expression more easily
manipulated by expression compilers. The original expression
is left untouched. Currently, the following transformations
are made to the preprocessed form:
expand_compounds (side effect of calling expand_derivatives)
expand_derivatives
renumber arguments and coefficients and apply evt. element mapping
"""
tic = Timer('preprocess_expression')
# Check that we get an expression
ufl_assert(isinstance(expr, Expr), "Expecting Expr.")
# Object names is empty if not given
object_names = object_names or {}
# Element mapping is empty if not given
element_mapping = element_mapping or {}
# Create empty expression data
expr_data = ExprData()
# Store original expression
expr_data.original_expr = expr
# Store name of expr if given, otherwise empty string
# such that automatic names can be assigned externally
expr_data.name = object_names.get(id(expr),
"") # TODO: Or default to 'expr'?
# Extract common cell
common_cell = extract_common_cell(expr, common_cell)
# TODO: Split out expand_compounds from expand_derivatives
# Expand derivatives
tic('expand_derivatives')
expr = expand_derivatives(expr, common_cell.geometric_dimension())
# Renumber indices
#expr = renumber_indices(expr) # TODO: No longer needed?
# Replace arguments and coefficients with new renumbered objects
tic('extract_arguments_and_coefficients')
original_arguments, original_coefficients = \
extract_arguments_and_coefficients(expr)
tic('build_element_mapping')
element_mapping = build_element_mapping(element_mapping, common_cell,
original_arguments,
original_coefficients)
tic('build_argument_replace_map')
replace_map, renumbered_arguments, renumbered_coefficients = \
build_argument_replace_map(original_arguments,
original_coefficients,
element_mapping)
expr = replace(expr, replace_map)
# Build mapping to original arguments and coefficients, which is
# useful if the original arguments have data attached to them
inv_replace_map = dict((w, v) for (v, w) in replace_map.iteritems())
original_arguments = [inv_replace_map[v] for v in renumbered_arguments]
original_coefficients = [
inv_replace_map[w] for w in renumbered_coefficients
]
# Store data extracted by preprocessing
expr_data.arguments = renumbered_arguments # TODO: Needed?
expr_data.coefficients = renumbered_coefficients # TODO: Needed?
expr_data.original_arguments = original_arguments
expr_data.original_coefficients = original_coefficients
expr_data.renumbered_arguments = renumbered_arguments
expr_data.renumbered_coefficients = renumbered_coefficients
tic('replace')
# Mappings from elements and functions (coefficients and arguments)
# that reside in expr to objects with canonical numbering as well as
# completed cells and elements
expr_data.element_replace_map = element_mapping
expr_data.function_replace_map = replace_map
# Store signature of form
tic('signature')
expr_data.signature = compute_expression_signature(
expr, expr_data.function_replace_map)
# Store elements, sub elements and element map
tic('extract_elements')
expr_data.elements = tuple(
f.element()
for f in chain(renumbered_arguments, renumbered_coefficients))
expr_data.unique_elements = unique_tuple(expr_data.elements)
expr_data.sub_elements = extract_sub_elements(expr_data.elements)
expr_data.unique_sub_elements = unique_tuple(expr_data.sub_elements)
# Store element domains (NB! This is likely to change!)
# FIXME: DOMAINS: What is a sensible way to store domains for a expr?
expr_data.domains = tuple(
sorted(set(element.domain() for element in expr_data.unique_elements)))
# Store toplevel domains (NB! This is likely to change!)
# FIXME: DOMAINS: What is a sensible way to store domains for a expr?
expr_data.top_domains = tuple(
sorted(set(domain.top_domain() for domain in expr_data.domains)))
# Store common cell
expr_data.cell = common_cell
# Store data related to cell
expr_data.geometric_dimension = expr_data.cell.geometric_dimension()
expr_data.topological_dimension = expr_data.cell.topological_dimension()
# Store some useful dimensions
expr_data.rank = len(expr_data.arguments) # TODO: Is this useful for expr?
expr_data.num_coefficients = len(expr_data.coefficients)
# Store argument names # TODO: Is this useful for expr?
expr_data.argument_names = \
[object_names.get(id(expr_data.original_arguments[i]), "v%d" % i)
for i in range(expr_data.rank)]
# Store coefficient names
expr_data.coefficient_names = \
[object_names.get(id(expr_data.original_coefficients[i]), "w%d" % i)
for i in range(expr_data.num_coefficients)]
# Store preprocessed expression
expr_data.preprocessed_expr = expr
tic.end()
# A coarse profiling implementation
# TODO: Add counting of nodes
# TODO: Add memory usage
if preprocess_expression.enable_profiling:
print tic
return expr_data
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)
