def backward_substitution(self, pc, y):
"""
"""
# We assemble the unknown which is an expression
# of the first eliminated variable.
self._sub_unknown()
# Recover the eliminated unknown
self._elim_unknown()
with timed_region("HybridProject"):
# Project the broken solution into non-broken spaces
broken_pressure = self.broken_solution.split()[self.pidx]
unbroken_pressure = self.unbroken_solution.split()[self.pidx]
broken_pressure.dat.copy(unbroken_pressure.dat)
# Compute the hdiv projection of the broken hdiv solution
broken_hdiv = self.broken_solution.split()[self.vidx]
unbroken_hdiv = self.unbroken_solution.split()[self.vidx]
unbroken_hdiv.assign(0)
par_loop(self.average_kernel, ufl.dx,
{"w": (self.weight, READ),
"vec_in": (broken_hdiv, READ),
"vec_out": (unbroken_hdiv, INC)})
with self.unbroken_solution.dat.vec_ro as v:
v.copy(y)
def __init__(self, equation):
"""
Initialise limiter
:param space : equation, as we need the broken space attached to it
"""
self.Vt = equation.space
# check this is the right space, only currently working for 2D extruded mesh
if self.Vt.extruded and self.Vt.mesh().topological_dimension() == 2:
# check that horizontal degree is 1 and vertical degree is 2
if self.Vt.ufl_element().degree()[0] is not 1 or \
self.Vt.ufl_element().degree()[1] is not 2:
raise ValueError('This is not the right limiter for this space.')
# check that continuity of the spaces is correct
# this will fail if the space does not use broken elements
if self.Vt.ufl_element()._element.sobolev_space()[0].name is not 'L2' or \
self.Vt.ufl_element()._element.sobolev_space()[1].name is not 'H1':
raise ValueError('This is not the right limiter for this space.')
else:
logger.warning('This limiter may not work for the space you are using.')
self.Q1DG = FunctionSpace(self.Vt.mesh(), 'DG', 1) # space with only vertex DOFs
self.vertex_limiter = VertexBasedLimiter(self.Q1DG)
self.theta_hat = Function(self.Q1DG) # theta function with only vertex DOFs
self.w = Function(self.Vt)
self.result = Function(self.Vt)
par_loop(_weight_kernel, dx, {"weight": (self.w, INC)})
def solve(self):
"""
Apply the solver with rhs state.xrhs and result state.dy.
"""
# Solve the hybridized system
self.hybridized_solver.solve()
broken_u, rho1, _ = self.urhol0.split()
u1 = self.u_hdiv
# Project broken_u into the HDiv space
u1.assign(0.0)
with timed_region("Gusto:HybridProjectHDiv"):
par_loop(self._average_kernel, dx,
{"w": (self._weight, READ),
"vec_in": (broken_u, READ),
"vec_out": (u1, INC)})
# Reapply bcs to ensure they are satisfied
for bc in self.bcs:
bc.apply(u1)
# Copy back into u and rho cpts of dy
u, rho, theta = self.state.dy.split()
u.assign(u1)
rho.assign(rho1)
# Reconstruct theta
with timed_region("Gusto:ThetaRecon"):
self.theta_solver.solve()
# Copy into theta cpt of dy
theta.assign(self.theta)
def forward_elimination(self, pc, x):
"""
"""
with timed_region("HybridBreak"):
with self.unbroken_residual.dat.vec_wo as v:
x.copy(v)
# Transfer unbroken_rhs into broken_rhs
# NOTE: Scalar space is already "broken" so no need for
# any projections
unbroken_scalar_data = self.unbroken_residual.split()[self.pidx]
broken_scalar_data = self.broken_residual.split()[self.pidx]
unbroken_scalar_data.dat.copy(broken_scalar_data.dat)
# Assemble the new "broken" hdiv residual
# We need a residual R' in the broken space that
# gives R'[w] = R[w] when w is in the unbroken space.
# We do this by splitting the residual equally between
# basis functions that add together to give unbroken
# basis functions.
unbroken_res_hdiv = self.unbroken_residual.split()[self.vidx]
broken_res_hdiv = self.broken_residual.split()[self.vidx]
broken_res_hdiv.assign(0)
par_loop(self.average_kernel, ufl.dx,
{"w": (self.weight, READ),
"vec_in": (unbroken_res_hdiv, READ),
"vec_out": (broken_res_hdiv, INC)})
with timed_region("HybridRHS"):
# Compute the rhs for the multiplier system
self._assemble_Srhs()
def backward_substitution(self, pc, y):
"""Perform the backwards recovery of eliminated fields.
:arg pc: a Preconditioner instance.
:arg y: a PETSc vector for placing the resulting fields.
"""
# We assemble the unknown which is an expression
# of the first eliminated variable.
self._sub_unknown()
# Recover the eliminated unknown
self._elim_unknown()
with timed_region("HybridProject"):
# Project the broken solution into non-broken spaces
broken_pressure = self.broken_solution.split()[self.pidx]
unbroken_pressure = self.unbroken_solution.split()[self.pidx]
broken_pressure.dat.copy(unbroken_pressure.dat)
# Compute the hdiv projection of the broken hdiv solution
broken_hdiv = self.broken_solution.split()[self.vidx]
unbroken_hdiv = self.unbroken_solution.split()[self.vidx]
unbroken_hdiv.assign(0)
par_loop(self.average_kernel, ufl.dx,
{"w": (self.weight, READ),
"vec_in": (broken_hdiv, READ),
"vec_out": (unbroken_hdiv, INC)},
is_loopy_kernel=True)
with self.unbroken_solution.dat.vec_ro as v:
v.copy(y)
def backward_substitution(self, pc, y):
"""Perform the backwards recovery of eliminated fields.
:arg pc: a Preconditioner instance.
:arg y: a PETSc vector for placing the resulting fields.
"""
# We assemble the unknown which is an expression
# of the first eliminated variable.
self._sub_unknown()
# Recover the eliminated unknown
self._elim_unknown()
with timed_region("HybridProject"):
# Project the broken solution into non-broken spaces
broken_pressure = self.broken_solution.split()[self.pidx]
unbroken_pressure = self.unbroken_solution.split()[self.pidx]
broken_pressure.dat.copy(unbroken_pressure.dat)
# Compute the hdiv projection of the broken hdiv solution
broken_hdiv = self.broken_solution.split()[self.vidx]
unbroken_hdiv = self.unbroken_solution.split()[self.vidx]
unbroken_hdiv.assign(0)
par_loop(self.average_kernel,
ufl.dx, {
"w": (self.weight, READ),
"vec_in": (broken_hdiv, READ),
"vec_out": (unbroken_hdiv, INC)
},
is_loopy_kernel=True)
with self.unbroken_solution.dat.vec_ro as v:
v.copy(y)
def copy_vertex_values(self, field):
"""
Copies the vertex values from temperature space to
Q1DG space which only has vertices.
"""
par_loop(_copy_into_Q1DG_loop, dx,
{"theta": (field, READ),
"theta_hat": (self.theta_hat, RW)})
def apply(self, w):
"""
Perform the par loop for calculating the weightings for the Averager.
:arg w: the field to store the weights in.
"""
par_loop(self._kernel, dx, {"w": (w, INC)}, is_loopy_kernel=True)
def check_midpoint_values(self, field):
"""
Checks the midpoint field values are less than the maximum
and more than the minimum values. Amends them to the average
if they are not.
"""
par_loop(_check_midpoint_values_loop, dx,
{"theta": (field, RW)})
def copy_vertex_values_back(self, field):
"""
Copies the vertex values back from the Q1DG space to
the original temperature space.
"""
par_loop(_copy_from_Q1DG_loop, dx,
{"theta": (field, RW),
"theta_hat": (self.theta_hat, READ)})
def apply_limiter(self, field):
"""
Only applies limiting loop on the given field
"""
par_loop(self._limit_kernel, dx,
{"qbar": (self.centroids, READ),
"q": (field, RW),
"qmax": (self.max_field, READ),
"qmin": (self.min_field, READ)})
def remap_to_embedded_space(self, field):
"""
Remap from DG space to embedded DG space.
"""
self.result.assign(0.)
par_loop(_average_kernel, dx, {"vrec": (self.result, INC),
"v_b": (field, READ),
"weight": (self.w, READ)})
field.assign(self.result)
def copy_vertex_values(self, field):
"""
Copies the vertex values from temperature space to
DG1 space which only has vertices.
"""
par_loop(self._copy_into_DG1_kernel,
dx, {
"theta": (field, READ),
"theta_hat": (self.theta_hat, WRITE)
},
is_loopy_kernel=True)
def apply_limiter(self, field):
"""
Only applies limiting loop on the given field
"""
par_loop(
self._limit_kernel, dx, {
"qbar": (self.centroids, READ),
"q": (field, RW),
"qmax": (self.max_field, READ),
"qmin": (self.min_field, READ)
})
def project(self):
"""
Apply the recovery.
"""
# Ensure that the function being populated is zeroed out
self.v_out.dat.zero()
par_loop(self._average_kernel, ufl.dx, {"vo": (self.v_out, INC),
"w": (self._weighting, READ),
"v": (self.v, READ)})
return self.v_out
def create_sc_nullspace(P, V, V_facet, comm):
"""Gets the nullspace vectors corresponding to the Schur complement
system.
:arg P: The H1 operator from the ImplicitMatrixContext.
:arg V: The H1 finite element space.
:arg V_facet: The finite element space of H1 basis functions
restricted to the mesh skeleton.
Returns: A nullspace (if there is one) for the Schur-complement system.
"""
from firedrake import Function
nullspace = P.getNullSpace()
if nullspace.handle == 0:
# No nullspace
return None
vecs = nullspace.getVecs()
tmp = Function(V)
scsp_tmp = Function(V_facet)
new_vecs = []
# Transfer the trace bit (the nullspace vector restricted
# to facet nodes is the nullspace for the condensed system)
kernel = """
for (int i=0; i<%(d)d; ++i){
for (int j=0; j<%(s)d; ++j){
x_facet[i*%(s)d + j] = x_h[i*%(s)d + j];
}
}""" % {
"d": V_facet.finat_element.space_dimension(),
"s": np.prod(V_facet.shape)
}
for v in vecs:
with tmp.dat.vec_wo as t:
v.copy(t)
par_loop(kernel, ufl.dx, {
"x_facet": (scsp_tmp, WRITE),
"x_h": (tmp, READ)
})
# Map vecs to the facet space
with scsp_tmp.dat.vec_ro as v:
new_vecs.append(v.copy())
# Normalize
for v in new_vecs:
v.normalize()
sc_nullspace = PETSc.NullSpace().create(vectors=new_vecs, comm=comm)
return sc_nullspace
def _partition_residual(self):
"""Partition the incoming right-hand side residual
into 'interior' and 'facet' sections.
"""
r_int = self.interior_residual
r_facet = self.trace_residual
par_loop(
self._transfer_kernel.partition, ufl.dx, {
"x_int": (r_int, WRITE),
"x_facet": (r_facet, WRITE),
"x": (self.h1_residual, READ)
})
def _weighting(self):
"""
Generates a weight function for computing a projection via averaging.
"""
w = Function(self.V)
weight_kernel = """
for (int i=0; i<%d; ++i) {
for (int j=0; j<%d; ++j) {
w[i][j] += 1.0;
}}""" % self._shapes
par_loop(weight_kernel, ufl.dx, {"w": (w, INC)})
return w
def compute_bounds(self, field):
"""
Only computes min and max bounds of neighbouring cells
"""
self._update_centroids(field)
self.max_field.assign(-1.0e10) # small number
self.min_field.assign(1.0e10) # big number
par_loop(self._min_max_loop,
dx,
{"maxq": (self.max_field, RW),
"minq": (self.min_field, RW),
"q": (self.centroids, READ)})
def compute_bounds(self, field):
"""
Only computes min and max bounds of neighbouring cells
"""
self._update_centroids(field)
self.max_field.assign(-1.0e10) # small number
self.min_field.assign(1.0e10) # big number
par_loop(
self._min_max_loop, dx, {
"maxq": (self.max_field, MAX),
"minq": (self.min_field, MIN),
"q": (self.centroids, READ)
})
def remap_to_embedded_space(self, field):
"""
Remap from DG space to embedded DG space.
"""
self.result.assign(0.)
par_loop(self._average_kernel,
dx, {
"vo": (self.result, INC),
"v": (field, READ),
"w": (self.w, READ)
},
is_loopy_kernel=True)
field.assign(self.result)
def solve(self):
"""
Apply the solver with rhs state.xrhs and result state.dy.
"""
# Solve the velocity-density system
with timed_region("Gusto:VelocityDensitySolve"):
# Assemble the RHS for lambda into self.R
with timed_region("Gusto:HybridRHS"):
self._assemble_Rexp()
# Solve for lambda
with timed_region("Gusto:HybridTraceSolve"):
self.lSolver.solve(self.lambdar, self.R)
# Reconstruct broken u and rho
with timed_region("Gusto:HybridRecon"):
self._assemble_rho()
self._assemble_u()
broken_u, rho1 = self.urho.split()
u1 = self.u_hdiv
# Project broken_u into the HDiv space
u1.assign(0.0)
with timed_region("Gusto:HybridProjectHDiv"):
par_loop(self._average_kernel, dx,
{"w": (self._weight, READ),
"vec_in": (broken_u, READ),
"vec_out": (u1, INC)})
# Reapply bcs to ensure they are satisfied
for bc in self.bcs:
bc.apply(u1)
# Copy back into u and rho cpts of dy
u, rho, theta = self.state.dy.split()
u.assign(u1)
rho.assign(rho1)
# Reconstruct theta
with timed_region("Gusto:ThetaRecon"):
self.theta_solver.solve()
# Copy into theta cpt of dy
theta.assign(self.theta)
def copy_vertex_values_back(self, field):
"""
Copies the vertex values back from the DG1 space to
the original temperature space, and checks that the
midpoint values are within the minimum and maximum
at the adjacent vertices.
If outside of the minimum and maximum, correct the values
to be the average.
"""
par_loop(self._copy_from_DG1_kernel,
dx, {
"theta": (field, WRITE),
"theta_hat": (self.theta_hat, READ),
"theta_old": (self.theta_old, READ)
},
is_loopy_kernel=True)
def apply(self, v_DG1, v_CG1):
"""
Performs the par loop.
:arg v_DG1: the target field to correct.
:arg v_CG1: the initially recovered uncorrected field.
"""
par_loop(self._kernel,
dx,
args={
"DG1": (v_DG1, WRITE),
"CG1": (v_CG1, READ)
},
is_loopy_kernel=True,
iterate=ON_BOTTOM)
def apply(self, v_out, weighting, v_in):
"""
Perform the averaging par loop.
:arg v_out: the continuous target field.
:arg weighting: the weights to be used for the averaging.
:arg v_in: the input field.
"""
par_loop(self._kernel,
dx, {
"vo": (v_out, INC),
"w": (weighting, READ),
"v": (v_in, READ)
},
is_loopy_kernel=True)
def apply(self, field_hat, field_DG1, field_old):
"""
Performs the par loop.
:arg field_hat: The field to write to in the broken temperature space.
:arg field_DG1: A field in the equispaced DG1 space whose vertex values
have been limited.
:arg field_old: The original un-limited field in the broken temperature
space.
"""
par_loop(self._kernel,
dx, {
"field_hat": (field_hat, WRITE),
"field_DG1": (field_DG1, READ),
"field_old": (field_old, READ)
},
is_loopy_kernel=True)
def spatial_index(self):
"""Spatial index to quickly find which cell contains a given point."""
from firedrake import function, functionspace
from firedrake.parloops import par_loop, READ, RW
gdim = self.ufl_cell().geometric_dimension()
if gdim <= 1:
info_red("libspatialindex does not support 1-dimension, falling back on brute force.")
return None
# Calculate the bounding boxes for all cells by running a kernel
V = functionspace.VectorFunctionSpace(self, "DG", 0, dim=gdim)
coords_min = function.Function(V)
coords_max = function.Function(V)
coords_min.dat.data.fill(np.inf)
coords_max.dat.data.fill(-np.inf)
kernel = """
for (int d = 0; d < gdim; d++) {
for (int i = 0; i < nodes_per_cell; i++) {
f_min[0][d] = fmin(f_min[0][d], f[i][d]);
f_max[0][d] = fmax(f_max[0][d], f[i][d]);
}
}
"""
cell_node_list = self.coordinates.function_space().cell_node_list
nodes_per_cell = len(cell_node_list[0])
kernel = kernel.replace("gdim", str(gdim))
kernel = kernel.replace("nodes_per_cell", str(nodes_per_cell))
par_loop(kernel, ufl.dx, {'f': (self.coordinates, READ),
'f_min': (coords_min, RW),
'f_max': (coords_max, RW)})
# Reorder bounding boxes according to the cell indices we use
column_list = V.cell_node_list.reshape(-1)
coords_min = self._order_data_by_cell_index(column_list, coords_min.dat.data_ro_with_halos)
coords_max = self._order_data_by_cell_index(column_list, coords_max.dat.data_ro_with_halos)
# Build spatial index
return spatialindex.from_regions(coords_min, coords_max)
def _weighting(self):
"""
Generates a weight function for computing a projection via averaging.
"""
w = Function(self.V)
weight_domain = "{{[i, j]: 0 <= i < {nDOFs} and 0 <= j < {dim}}}".format(
**self.shapes)
weight_instructions = ("""
for i
for j
w[i,j] = w[i,j] + 1.0
end
end
""")
_weight_kernel = (weight_domain, weight_instructions)
par_loop(_weight_kernel, dx, {"w": (w, INC)}, is_loopy_kernel=True)
return w
def _reconstruct(self):
"""Locally solve for the interior degrees of
freedom using the computed unknowns for the facets.
A transfer kernel is used to join the interior and
facet solutions together.
"""
with timed_region("SCAssembleInterior"):
self._assemble_interior_u()
u_int = self.interior_solution
u_facet = self.trace_solution
with timed_region("SCReconSolution"):
par_loop(
self._transfer_kernel.join, ufl.dx, {
"x": (self.h1_solution, WRITE),
"x_int": (u_int, READ),
"x_facet": (u_facet, READ)
})
def apply(self, v_DG1_old, v_DG1, act_coords, eff_coords, num_ext):
"""
Performs the par loop for the Gaussian elimination kernel.
:arg v_DG1_old: the originally recovered field in DG1.
:arg v_DG1: the new target field in DG1.
:arg act_coords: the actual coordinates in vec DG1.
:arg eff_coords: the effective coordinates of the recovery in vec DG1.
:arg num_ext: the number of exterior DOFs in the cell, in DG0.
"""
par_loop(self._kernel,
dx, {
"DG1_OLD": (v_DG1_old, READ),
"DG1": (v_DG1, WRITE),
"ACT_COORDS": (act_coords, READ),
"EFF_COORDS": (eff_coords, READ),
"NUM_EXT": (num_ext, READ)
},
is_loopy_kernel=True)
def forward_elimination(self, pc, x):
"""Perform the forward elimination of fields and
provide the reduced right-hand side for the condensed
system.
:arg pc: a Preconditioner instance.
:arg x: a PETSc vector containing the incoming right-hand side.
"""
with timed_region("HybridBreak"):
with self.unbroken_residual.dat.vec_wo as v:
x.copy(v)
# Transfer unbroken_rhs into broken_rhs
# NOTE: Scalar space is already "broken" so no need for
# any projections
unbroken_scalar_data = self.unbroken_residual.split()[self.pidx]
broken_scalar_data = self.broken_residual.split()[self.pidx]
unbroken_scalar_data.dat.copy(broken_scalar_data.dat)
# Assemble the new "broken" hdiv residual
# We need a residual R' in the broken space that
# gives R'[w] = R[w] when w is in the unbroken space.
# We do this by splitting the residual equally between
# basis functions that add together to give unbroken
# basis functions.
unbroken_res_hdiv = self.unbroken_residual.split()[self.vidx]
broken_res_hdiv = self.broken_residual.split()[self.vidx]
broken_res_hdiv.assign(0)
par_loop(self.average_kernel,
ufl.dx, {
"w": (self.weight, READ),
"vec_in": (unbroken_res_hdiv, READ),
"vec_out": (broken_res_hdiv, INC)
},
is_loopy_kernel=True)
with timed_region("HybridRHS"):
# Compute the rhs for the multiplier system
self._assemble_Srhs()
def apply(self):
self.interpolator.interpolate()
if self.method == Boundary_Method.physics:
par_loop(self._bottom_kernel,
dx,
args={
"DG1": (self.v_DG1, WRITE),
"CG1": (self.v_CG1, READ)
},
is_loopy_kernel=True,
iterate=ON_BOTTOM)
par_loop(self._top_kernel,
dx,
args={
"DG1": (self.v_DG1, WRITE),
"CG1": (self.v_CG1, READ)
},
is_loopy_kernel=True,
iterate=ON_TOP)
else:
self.v_DG1_old.assign(self.v_DG1)
par_loop(self._gaussian_elimination_kernel,
dx, {
"DG1_OLD": (self.v_DG1_old, READ),
"DG1": (self.v_DG1, WRITE),
"ACT_COORDS": (self.act_coords, READ),
"EFF_COORDS": (self.eff_coords, READ),
"NUM_EXT": (self.num_ext, READ)
},
is_loopy_kernel=True)
def forward_elimination(self, pc, x):
"""Perform the forward elimination of fields and
provide the reduced right-hand side for the condensed
system.
:arg pc: a Preconditioner instance.
:arg x: a PETSc vector containing the incoming right-hand side.
"""
with timed_region("HybridBreak"):
with self.unbroken_residual.dat.vec_wo as v:
x.copy(v)
# Transfer unbroken_rhs into broken_rhs
# NOTE: Scalar space is already "broken" so no need for
# any projections
unbroken_scalar_data = self.unbroken_residual.split()[self.pidx]
broken_scalar_data = self.broken_residual.split()[self.pidx]
unbroken_scalar_data.dat.copy(broken_scalar_data.dat)
# Assemble the new "broken" hdiv residual
# We need a residual R' in the broken space that
# gives R'[w] = R[w] when w is in the unbroken space.
# We do this by splitting the residual equally between
# basis functions that add together to give unbroken
# basis functions.
unbroken_res_hdiv = self.unbroken_residual.split()[self.vidx]
broken_res_hdiv = self.broken_residual.split()[self.vidx]
broken_res_hdiv.assign(0)
par_loop(self.average_kernel, ufl.dx,
{"w": (self.weight, READ),
"vec_in": (unbroken_res_hdiv, READ),
"vec_out": (broken_res_hdiv, INC)},
is_loopy_kernel=True)
with timed_region("HybridRHS"):
# Compute the rhs for the multiplier system
self._assemble_Srhs()
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 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)