class IncompressibleState(BaroclinicState):
def __init__(self, mesh, vertical_degree=1, horizontal_degree=1,
family="RT", z=None, k=None, Omega=None, mu=None,
timestepping=None,
output=None,
parameters=None,
diagnostics=None,
fieldlist=None,
diagnostic_fields=[],
on_sphere=False):
super(IncompressibleState, self).__init__(mesh=mesh,
vertical_degree=vertical_degree,
horizontal_degree=horizontal_degree,
family=family,
z=z, k=k, Omega=Omega, mu=mu,
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist,
diagnostic_fields=diagnostic_fields)
if parameters.geopotential:
raise RuntimeError("geopotential formulation is not compatible with incompressible Boussinesq")
def set_reference_profiles(self, b_ref):
"""
Initialise reference profiles
:arg b_ref: :class:`.Function` object, reference bouyancy
"""
self.bbar = Function(self.V[2])
self.bbar.project(b_ref)
class EmbeddedDGAdvection(DGAdvection):
def __init__(self, state, V, Vdg=None, continuity=False):
if Vdg is None:
Vdg_elt = BrokenElement(V.ufl_element())
Vdg = FunctionSpace(state.mesh, Vdg_elt)
super(EmbeddedDGAdvection, self).__init__(state, Vdg, continuity)
self.xdg_in = Function(Vdg)
self.xdg_out = Function(Vdg)
self.x_projected = Function(V)
pparameters = {'ksp_type':'cg',
'pc_type':'bjacobi',
'sub_pc_type':'ilu'}
self.Projector = Projector(self.xdg_out, self.x_projected,
solver_parameters=pparameters)
def apply(self, x_in, x_out):
self.xdg_in.interpolate(x_in)
super(EmbeddedDGAdvection, self).apply(self.xdg_in, self.xdg_out)
self.Projector.project()
x_out.assign(self.x_projected)
def set_reference_profiles(self, rho_ref, theta_ref):
"""
Initialise reference profiles
:arg rho_ref: :class:`.Function` object, initial rho
:arg theta_ref: :class:`.Function` object, initial theta
"""
self.rhobar = Function(self.V[1])
self.thetabar = Function(self.V[2])
self.rhobar.project(rho_ref)
self.thetabar.project(theta_ref)
def __init__(self, state, linear=False):
self.state = state
g = state.parameters.g
f = state.f
Vu = state.V[0]
W = state.W
self.x0 = Function(W) # copy x to here
u0, D0 = split(self.x0)
n = FacetNormal(state.mesh)
un = 0.5 * (dot(u0, n) + abs(dot(u0, n)))
F = TrialFunction(Vu)
w = TestFunction(Vu)
self.uF = Function(Vu)
outward_normals = CellNormal(state.mesh)
perp = lambda u: cross(outward_normals, u)
a = inner(w, F) * dx
L = (-f * inner(w, perp(u0)) + g * div(w) * D0) * dx - g * inner(
jump(w, n), un("+") * D0("+") - un("-") * D0("-")
) * dS
if not linear:
L -= 0.5 * div(w) * inner(u0, u0) * dx
u_forcing_problem = LinearVariationalProblem(a, L, self.uF)
self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem)
def __init__(self, state, V, direction=[1,2], params=None):
super(InteriorPenalty, self).__init__(state)
dt = state.timestepping.dt
kappa = params['kappa']
mu = params['mu']
gamma = TestFunction(V)
phi = TrialFunction(V)
self.phi1 = Function(V)
n = FacetNormal(state.mesh)
a = inner(gamma,phi)*dx + dt*inner(grad(gamma), grad(phi)*kappa)*dx
def get_flux_form(dS, M):
fluxes = (-inner(2*avg(outer(phi, n)), avg(grad(gamma)*M))
- inner(avg(grad(phi)*M), 2*avg(outer(gamma, n)))
+ mu*inner(2*avg(outer(phi, n)), 2*avg(outer(gamma, n)*kappa)))*dS
return fluxes
if 1 in direction:
a += dt*get_flux_form(dS_v, kappa)
if 2 in direction:
a += dt*get_flux_form(dS_h, kappa)
L = inner(gamma,phi)*dx
problem = LinearVariationalProblem(a, action(L,self.phi1), self.phi1)
self.solver = LinearVariationalSolver(problem)
def get_latlon_mesh(mesh):
coords_orig = mesh.coordinates
mesh_dg_fs = VectorFunctionSpace(mesh, "DG", 1)
coords_dg = Function(mesh_dg_fs)
coords_latlon = Function(mesh_dg_fs)
par_loop("""
for (int i=0; i<3; i++) {
for (int j=0; j<3; j++) {
dg[i][j] = cg[i][j];
}
}
""", dx, {'dg': (coords_dg, WRITE),
'cg': (coords_orig, READ)})
# lat-lon 'x' = atan2(y, x)
coords_latlon.dat.data[:,0] = np.arctan2(coords_dg.dat.data[:,1], coords_dg.dat.data[:,0])
# lat-lon 'y' = asin(z/sqrt(x^2 + y^2 + z^2))
coords_latlon.dat.data[:,1] = np.arcsin(coords_dg.dat.data[:,2]/np.sqrt(coords_dg.dat.data[:,0]**2 + coords_dg.dat.data[:,1]**2 + coords_dg.dat.data[:,2]**2))
coords_latlon.dat.data[:,2] = 0.0
kernel = op2.Kernel("""
#define PI 3.141592653589793
#define TWO_PI 6.283185307179586
void splat_coords(double **coords) {
double diff0 = (coords[0][0] - coords[1][0]);
double diff1 = (coords[0][0] - coords[2][0]);
double diff2 = (coords[1][0] - coords[2][0]);
if (fabs(diff0) > PI || fabs(diff1) > PI || fabs(diff2) > PI) {
const int sign0 = coords[0][0] < 0 ? -1 : 1;
const int sign1 = coords[1][0] < 0 ? -1 : 1;
const int sign2 = coords[2][0] < 0 ? -1 : 1;
if (sign0 < 0) {
coords[0][0] += TWO_PI;
}
if (sign1 < 0) {
coords[1][0] += TWO_PI;
}
if (sign2 < 0) {
coords[2][0] += TWO_PI;
}
}
}""", "splat_coords")
op2.par_loop(kernel, coords_latlon.cell_set,
coords_latlon.dat(op2.RW, coords_latlon.cell_node_map()))
return Mesh(coords_latlon)
class ShallowWaterForcing(Forcing):
def __init__(self, state, linear=False):
self.state = state
g = state.parameters.g
f = state.f
Vu = state.V[0]
W = state.W
self.x0 = Function(W) # copy x to here
u0, D0 = split(self.x0)
n = FacetNormal(state.mesh)
un = 0.5 * (dot(u0, n) + abs(dot(u0, n)))
F = TrialFunction(Vu)
w = TestFunction(Vu)
self.uF = Function(Vu)
outward_normals = CellNormal(state.mesh)
perp = lambda u: cross(outward_normals, u)
a = inner(w, F) * dx
L = (-f * inner(w, perp(u0)) + g * div(w) * D0) * dx - g * inner(
jump(w, n), un("+") * D0("+") - un("-") * D0("-")
) * dS
if not linear:
L -= 0.5 * div(w) * inner(u0, u0) * dx
u_forcing_problem = LinearVariationalProblem(a, L, self.uF)
self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem)
def apply(self, scaling, x_in, x_nl, x_out, **kwargs):
self.x0.assign(x_nl)
self.u_forcing_solver.solve() # places forcing in self.uF
self.uF *= scaling
uF, _ = x_out.split()
x_out.assign(x_in)
uF += self.uF
class InteriorPenalty(Diffusion):
"""
Interior penalty diffusion method
:arg state: :class:`.State` object.
:arg V: Function space of diffused field
:arg direction: list containing directions in which function space
is discontinuous: 1 corresponds to vertical, 2 to horizontal.
:arg params: dictionary containing the interior penalty parameters
:mu and kappa where mu is the penalty weighting function, which is
:recommended to be proportional to 1/dx
"""
def __init__(self, state, V, direction=[1,2], params=None):
super(InteriorPenalty, self).__init__(state)
dt = state.timestepping.dt
kappa = params['kappa']
mu = params['mu']
gamma = TestFunction(V)
phi = TrialFunction(V)
self.phi1 = Function(V)
n = FacetNormal(state.mesh)
a = inner(gamma,phi)*dx + dt*inner(grad(gamma), grad(phi)*kappa)*dx
def get_flux_form(dS, M):
fluxes = (-inner(2*avg(outer(phi, n)), avg(grad(gamma)*M))
- inner(avg(grad(phi)*M), 2*avg(outer(gamma, n)))
+ mu*inner(2*avg(outer(phi, n)), 2*avg(outer(gamma, n)*kappa)))*dS
return fluxes
if 1 in direction:
a += dt*get_flux_form(dS_v, kappa)
if 2 in direction:
a += dt*get_flux_form(dS_h, kappa)
L = inner(gamma,phi)*dx
problem = LinearVariationalProblem(a, action(L,self.phi1), self.phi1)
self.solver = LinearVariationalSolver(problem)
def apply(self, x_in, x_out):
self.phi1.assign(x_in)
self.solver.solve()
x_out.assign(self.phi1)
def set_reference_profiles(self, b_ref):
"""
Initialise reference profiles
:arg b_ref: :class:`.Function` object, reference bouyancy
"""
self.bbar = Function(self.V[2])
self.bbar.project(b_ref)
def _build_forcing_solver(self, linear):
"""
Only put forcing terms into the u equation.
"""
state = self.state
self.scaling = Constant(1.0)
Vu = state.V[0]
W = state.W
self.x0 = Function(W) # copy x to here
u0, rho0, theta0 = split(self.x0)
F = TrialFunction(Vu)
w = TestFunction(Vu)
self.uF = Function(Vu)
Omega = state.Omega
cp = state.parameters.cp
mu = state.mu
n = FacetNormal(state.mesh)
pi = exner(theta0, rho0, state)
a = inner(w, F) * dx
L = self.scaling * (
+cp * div(theta0 * w) * pi * dx # pressure gradient [volume]
- cp * jump(w * theta0, n) * avg(pi) * dS_v # pressure gradient [surface]
)
if state.parameters.geopotential:
Phi = state.Phi
L += self.scaling * div(w) * Phi * dx # gravity term
else:
g = state.parameters.g
L -= self.scaling * g * inner(w, state.k) * dx # gravity term
if not linear:
L -= self.scaling * 0.5 * div(w) * inner(u0, u0) * dx
if Omega is not None:
L -= self.scaling * inner(w, cross(2 * Omega, u0)) * dx # Coriolis term
if mu is not None:
self.mu_scaling = Constant(1.0)
L -= self.mu_scaling * mu * inner(w, state.k) * inner(u0, state.k) * dx
bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")]
u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs)
self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem)
def _build_forcing_solver(self, linear):
"""
Only put forcing terms into the u equation.
"""
state = self.state
self.scaling = Constant(1.0)
Vu = state.V[0]
W = state.W
self.x0 = Function(W) # copy x to here
u0, p0, b0 = split(self.x0)
F = TrialFunction(Vu)
w = TestFunction(Vu)
self.uF = Function(Vu)
Omega = state.Omega
mu = state.mu
a = inner(w, F) * dx
L = (
self.scaling * div(w) * p0 * dx # pressure gradient
+ self.scaling * b0 * inner(w, state.k) * dx # gravity term
)
if not linear:
L -= self.scaling * 0.5 * div(w) * inner(u0, u0) * dx
if Omega is not None:
L -= self.scaling * inner(w, cross(2 * Omega, u0)) * dx # Coriolis term
if mu is not None:
self.mu_scaling = Constant(1.0)
L -= self.mu_scaling * mu * inner(w, state.k) * inner(u0, state.k) * dx
bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")]
u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs)
self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem)
Vp = state.V[1]
p = TrialFunction(Vp)
q = TestFunction(Vp)
self.divu = Function(Vp)
a = p * q * dx
L = q * div(u0) * dx
divergence_problem = LinearVariationalProblem(a, L, self.divu)
self.divergence_solver = LinearVariationalSolver(divergence_problem)
def bezier_plot(function, axes=None, **kwargs):
"""Plot a 1D function on a function space with order no more than 4 using
Bezier curve within each cell, return a matplotlib axes
:arg function: 1D function for plotting
:arg axes: Axes for plotting, if None, a new one will be created
:arg kwargs: additional key work arguments to plot
"""
try:
import matplotlib.pyplot as plt
from matplotlib.path import Path
import matplotlib.patches as patches
except ImportError:
raise RuntimeError("Matplotlib not importable, is it installed?")
deg = function.function_space().ufl_element().degree()
mesh = function.function_space().mesh()
if deg == 0:
V = FunctionSpace(mesh, "DG", 1)
func = Function(V).interpolate(function)
return bezier_plot(func, axes, **kwargs)
y_vals = _bezier_calculate_points(function)
x = SpatialCoordinate(mesh)
coords = Function(FunctionSpace(mesh, 'DG', deg))
coords.interpolate(x[0])
x_vals = _bezier_calculate_points(coords)
vals = np.dstack((x_vals, y_vals))
if axes is None:
figure = plt.figure()
axes = figure.add_subplot(111)
codes = {1: [Path.MOVETO, Path.LINETO],
2: [Path.MOVETO, Path.CURVE3, Path.CURVE3],
3: [Path.MOVETO, Path.CURVE4, Path.CURVE4, Path.CURVE4]}
vertices = vals.reshape(-1, 2)
path = Path(vertices, np.tile(codes[deg],
function.function_space().cell_node_list.shape[0]))
patch = patches.PathPatch(path, facecolor='none', lw=2)
axes.add_patch(patch)
axes.plot(**kwargs)
return axes
def getranks(mesh):
element_rank = Function(FunctionSpace(mesh,'DG',0))
element_rank.dat.data[:] = mesh.comm.rank
element_rank.rename('element_rank')
vertex_rank = Function(FunctionSpace(mesh,'CG',1))
vertex_rank.dat.data[:] = mesh.comm.rank
vertex_rank.rename('vertex_rank')
return (element_rank,vertex_rank)
def __init__(self, state, V, direction=[], supg_params=None):
super(SUPGAdvection, self).__init__(state)
dt = state.timestepping.dt
params = supg_params.copy() if supg_params else {}
params.setdefault('a0', dt/sqrt(15.))
params.setdefault('a1', dt/sqrt(15.))
gamma = TestFunction(V)
theta = TrialFunction(V)
self.theta0 = Function(V)
# make SUPG test function
taus = [params["a0"], params["a1"]]
for i in direction:
taus[i] = 0.0
tau = Constant(((taus[0], 0.), (0., taus[1])))
dgamma = dot(dot(self.ubar, tau), grad(gamma))
gammaSU = gamma + dgamma
n = FacetNormal(state.mesh)
un = 0.5*(dot(self.ubar, n) + abs(dot(self.ubar, n)))
a_mass = gammaSU*theta*dx
arhs = a_mass - dt*gammaSU*dot(self.ubar, grad(theta))*dx
if 1 in direction:
arhs -= (
dt*dot(jump(gammaSU), (un('+')*theta('+')
- un('-')*theta('-')))*dS_v
- dt*(gammaSU('+')*dot(self.ubar('+'), n('+'))*theta('+')
+ gammaSU('-')*dot(self.ubar('-'), n('-'))*theta('-'))*dS_v
)
if 2 in direction:
arhs -= (
dt*dot(jump(gammaSU), (un('+')*theta('+')
- un('-')*theta('-')))*dS_h
- dt*(gammaSU('+')*dot(self.ubar('+'), n('+'))*theta('+')
+ gammaSU('-')*dot(self.ubar('-'), n('-'))*theta('-'))*dS_h
)
self.theta1 = Function(V)
self.dtheta = Function(V)
problem = LinearVariationalProblem(a_mass, action(arhs,self.theta1), self.dtheta)
self.solver = LinearVariationalSolver(problem,
options_prefix='SUPGAdvection')
def __init__(self, state, V, continuity=False):
super(DGAdvection, self).__init__(state)
element = V.fiat_element
assert element.entity_dofs() == element.entity_closure_dofs(), "Provided space is not discontinuous"
dt = state.timestepping.dt
if V.extruded:
surface_measure = (dS_h + dS_v)
else:
surface_measure = dS
phi = TestFunction(V)
D = TrialFunction(V)
self.D1 = Function(V)
self.dD = Function(V)
n = FacetNormal(state.mesh)
# ( dot(v, n) + |dot(v, n)| )/2.0
un = 0.5*(dot(self.ubar, n) + abs(dot(self.ubar, n)))
a_mass = inner(phi,D)*dx
if continuity:
a_int = -inner(grad(phi), outer(D, self.ubar))*dx
else:
a_int = -inner(div(outer(phi,self.ubar)),D)*dx
a_flux = (dot(jump(phi), un('+')*D('+') - un('-')*D('-')))*surface_measure
arhs = a_mass - dt*(a_int + a_flux)
DGproblem = LinearVariationalProblem(a_mass, action(arhs,self.D1),
self.dD)
self.DGsolver = LinearVariationalSolver(DGproblem,
solver_parameters={
'ksp_type':'preonly',
'pc_type':'bjacobi',
'sub_pc_type': 'ilu'},
options_prefix='DGAdvection')
def PeriodicRectangleMesh(nx, ny, Lx, Ly, direction="both",
quadrilateral=False, reorder=None, comm=COMM_WORLD):
"""Generate a periodic rectangular mesh
:arg nx: The number of cells in the x direction
:arg ny: The number of cells in the y direction
:arg Lx: The extent in the x direction
:arg Ly: The extent in the y direction
:arg direction: The direction of the periodicity, one of
``"both"``, ``"x"`` or ``"y"``.
:kwarg quadrilateral: (optional), creates quadrilateral mesh, defaults to False
:kwarg reorder: (optional), should the mesh be reordered
:kwarg comm: Optional communicator to build the mesh on (defaults to
COMM_WORLD).
If direction == "x" the boundary edges in this mesh are numbered as follows:
* 1: plane y == 0
* 2: plane y == Ly
If direction == "y" the boundary edges are:
* 1: plane x == 0
* 2: plane x == Lx
"""
if direction == "both" and ny == 1 and quadrilateral:
return OneElementThickMesh(nx, Lx, Ly)
if direction not in ("both", "x", "y"):
raise ValueError("Cannot have a periodic mesh with periodicity '%s'" % direction)
if direction != "both":
return PartiallyPeriodicRectangleMesh(nx, ny, Lx, Ly, direction=direction,
quadrilateral=quadrilateral,
reorder=reorder,
comm=comm)
if nx < 3 or ny < 3:
raise ValueError("2D periodic meshes with fewer than 3 \
cells in each direction are not currently supported")
m = TorusMesh(nx, ny, 1.0, 0.5, quadrilateral=quadrilateral, reorder=reorder,
comm=comm)
coord_fs = VectorFunctionSpace(m, 'DG', 1, dim=2)
old_coordinates = m.coordinates
new_coordinates = Function(coord_fs)
periodic_kernel = """
double pi = 3.141592653589793;
double eps = 1e-12;
double bigeps = 1e-1;
double phi, theta, Y, Z;
Y = 0.0;
Z = 0.0;
for(int i=0; i<old_coords.dofs; i++) {
Y += old_coords[i][1];
Z += old_coords[i][2];
}
for(int i=0; i<new_coords.dofs; i++) {
phi = atan2(old_coords[i][1], old_coords[i][0]);
if (fabs(sin(phi)) > bigeps)
theta = atan2(old_coords[i][2], old_coords[i][1]/sin(phi) - 1.0);
else
theta = atan2(old_coords[i][2], old_coords[i][0]/cos(phi) - 1.0);
new_coords[i][0] = phi/(2.0*pi);
if(new_coords[i][0] < -eps) {
new_coords[i][0] += 1.0;
}
if(fabs(new_coords[i][0]) < eps && Y < 0.0) {
new_coords[i][0] = 1.0;
}
new_coords[i][1] = theta/(2.0*pi);
if(new_coords[i][1] < -eps) {
new_coords[i][1] += 1.0;
}
if(fabs(new_coords[i][1]) < eps && Z < 0.0) {
new_coords[i][1] = 1.0;
}
new_coords[i][0] *= Lx[0];
new_coords[i][1] *= Ly[0];
}
"""
cLx = Constant(Lx)
cLy = Constant(Ly)
par_loop(periodic_kernel, dx,
{"new_coords": (new_coordinates, WRITE),
"old_coords": (old_coordinates, READ),
"Lx": (cLx, READ),
"Ly": (cLy, READ)})
return mesh.Mesh(new_coordinates)
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 OneElementThickMesh(ncells, Lx, Ly, comm=COMM_WORLD):
"""
Generate a rectangular mesh in the domain with corners [0,0]
and [Lx, Ly] with ncells, that is periodic in the x-direction.
:arg ncells: The number of cells in the mesh.
:arg Lx: The width of the domain in the x-direction.
:arg Ly: The width of the domain in the y-direction.
:kwarg comm: Optional communicator to build the mesh on (defaults to
COMM_WORLD).
"""
left = np.arange(ncells, dtype=np.int32)
right = np.roll(left, -1)
cells = np.array([left, left, right, right]).T
dx = Lx/ncells
X = np.arange(1.0*ncells, dtype=np.double)*dx
Y = 0.*X
coords = np.array([X, Y]).T
# a line of coordinates, with a looped topology
plex = mesh._from_cell_list(2, cells, coords, comm)
mesh1 = mesh.Mesh(plex)
mesh1.topology.init()
cell_numbering = mesh1._cell_numbering
cell_range = plex.getHeightStratum(0)
cell_closure = np.zeros((cell_range[1], 9), dtype=IntType)
# Get the coordinates for this process
coords = plex.getCoordinatesLocal().array_r
# get the PETSc section
coords_sec = plex.getCoordinateSection()
for e in range(*cell_range):
closure, orient = plex.getTransitiveClosure(e)
# get the row for this cell
row = cell_numbering.getOffset(e)
# run some checks
assert(closure[0] == e)
assert len(closure) == 7, closure
edge_range = plex.getHeightStratum(1)
assert(all(closure[1:5] >= edge_range[0]))
assert(all(closure[1:5] < edge_range[1]))
vertex_range = plex.getHeightStratum(2)
assert(all(closure[5:] >= vertex_range[0]))
assert(all(closure[5:] < vertex_range[1]))
# enter the cell number
cell_closure[row][8] = e
# Get a list of unique edges
edge_set = list(set(closure[1:5]))
# there are two vertices in the cell
cell_vertices = closure[5:]
cell_X = np.array([0., 0.])
for i, v in enumerate(cell_vertices):
cell_X[i] = coords[coords_sec.getOffset(v)]
# Add in the edges
for i in range(3):
# count up how many times each edge is repeated
repeats = list(closure[1:5]).count(edge_set[i])
if repeats == 2:
# we have a y-periodic edge
cell_closure[row][6] = edge_set[i]
cell_closure[row][7] = edge_set[i]
elif repeats == 1:
# in this code we check if it is a right edge, or a left edge
# by inspecting the x coordinates of the edge vertex (1)
# and comparing with the x coordinates of the cell vertices (2)
# there is only one vertex on the edge in this case
edge_vertex = plex.getCone(edge_set[i])[0]
# get X coordinate for this edge
edge_X = coords[coords_sec.getOffset(edge_vertex)]
# get X coordinates for this cell
if(cell_X.min() < dx/2):
if cell_X.max() < 3*dx/2:
# We are in the first cell
if(edge_X.min() < dx/2):
# we are on left hand edge
cell_closure[row][4] = edge_set[i]
else:
# we are on right hand edge
cell_closure[row][5] = edge_set[i]
else:
# We are in the last cell
if(edge_X.min() < dx/2):
# we are on right hand edge
cell_closure[row][5] = edge_set[i]
else:
# we are on left hand edge
cell_closure[row][4] = edge_set[i]
else:
if(abs(cell_X.min()-edge_X.min()) < dx/2):
# we are on left hand edge
cell_closure[row][4] = edge_set[i]
else:
# we are on right hand edge
cell_closure[row][5] = edge_set[i]
# Add in the vertices
vertices = closure[5:]
v1 = vertices[0]
v2 = vertices[1]
x1 = coords[coords_sec.getOffset(v1)]
x2 = coords[coords_sec.getOffset(v2)]
# Fix orientations
if(x1 > x2):
if(x1 - x2 < dx*1.5):
# we are not on the rightmost cell and need to swap
v1, v2 = v2, v1
elif(x2 - x1 > dx*1.5):
# we are on the rightmost cell and need to swap
v1, v2 = v2, v1
cell_closure[row][0:4] = [v1, v1, v2, v2]
mesh1.topology.cell_closure = np.array(cell_closure, dtype=IntType)
mesh1.init()
Vc = VectorFunctionSpace(mesh1, 'DQ', 1)
fc = Function(Vc).interpolate(mesh1.coordinates)
mash = mesh.Mesh(fc)
topverts = Vc.cell_node_list[:, 1::2].flatten()
mash.coordinates.dat.data_with_halos[topverts, 1] = Ly
# search for the last cell
mcoords_ro = mash.coordinates.dat.data_ro_with_halos
mcoords = mash.coordinates.dat.data_with_halos
for e in range(*cell_range):
cell = cell_numbering.getOffset(e)
cell_nodes = Vc.cell_node_list[cell, :]
Xvals = mcoords_ro[cell_nodes, 0]
if(Xvals.max() - Xvals.min() > Lx/2):
mcoords[cell_nodes[2:], 0] = Lx
else:
mcoords
local_facet_dat = mash.topology.interior_facets.local_facet_dat
local_facet_number = mash.topology.interior_facets.local_facet_number
lfd_ro = local_facet_dat.data_ro
for i in range(lfd_ro.shape[0]):
if all(lfd_ro[i, :] == np.array([3, 3])):
local_facet_dat.data[i, :] = [2, 3]
local_facet_number[i, :] = [2, 3]
return mash
def setup_sk(dirname):
nlayers = 10 # horizontal layers
columns = 30 # number of columns
L = 1.e5
m = PeriodicIntervalMesh(columns, L)
dt = 6.0
# build volume mesh
H = 1.0e4 # Height position of the model top
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers)
# Set up points for output at the centre of the domain, edges and corners.
# The point at x=L*(13.0/30) is in the halo region for a two-way MPI decomposition
points_x = [0.0, L * (13.0 / 30), L / 2.0, L]
points_z = [0.0, H / 2.0, H]
points = np.array([p for p in itertools.product(points_x, points_z)])
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=dt)
output = OutputParameters(dirname=dirname + "/sk_nonlinear",
dumplist=['u'],
dumpfreq=5,
log_level=INFO,
point_data=[('rho', points), ('u', points)])
parameters = CompressibleParameters()
diagnostic_fields = [CourantNumber()]
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist,
diagnostic_fields=diagnostic_fields)
# Initial conditions
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()
# Thermodynamic constants required for setting initial conditions
# and reference profiles
g = parameters.g
N = parameters.N
# N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz)
x, z = SpatialCoordinate(mesh)
Tsurf = 300.
thetab = Tsurf * exp(N**2 * z / g)
theta_b = Function(Vt).interpolate(thetab)
rho_b = Function(Vr)
# Calculate hydrostatic Pi
compressible_hydrostatic_balance(state, theta_b, rho_b)
a = 5.0e3
deltaTheta = 1.0e-2
theta_pert = deltaTheta * sin(np.pi * z / H) / (1 + (x - L / 2)**2 / a**2)
theta0.interpolate(theta_b + theta_pert)
rho0.assign(rho_b)
u0.project(as_vector([20.0, 0.0]))
state.initialise([('u', u0), ('rho', rho0), ('theta', theta0)])
state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])
# Set up advection schemes
ueqn = EulerPoincare(state, Vu)
rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity")
thetaeqn = SUPGAdvection(state, Vt)
advected_fields = []
advected_fields.append(("u", ThetaMethod(state, u0, ueqn)))
advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn)))
advected_fields.append(("theta", SSPRK3(state, theta0, thetaeqn)))
# Set up linear solver
linear_solver = CompressibleSolver(state)
# Set up forcing
compressible_forcing = CompressibleForcing(state)
# build time stepper
stepper = CrankNicolson(state, advected_fields, linear_solver,
compressible_forcing)
return stepper, 2 * dt
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()
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)
p0 = state.fields("p")
# spaces
Vu = u0.function_space()
Vb = b0.function_space()
x, z = SpatialCoordinate(mesh)
# first setup the background buoyancy profile
# z.grad(bref) = N**2
# the following is symbolic algebra, using the default buoyancy frequency
# from the parameters class.
N = parameters.N
bref = z * (N**2)
# interpolate the expression to the function
b_b = Function(Vb).interpolate(bref)
# setup constants
a = 5.0e3
deltab = 1.0e-2
b_pert = deltab * sin(np.pi * z / H) / (1 + (x - L / 2)**2 / a**2)
# interpolate the expression to the function
b0.interpolate(b_b + b_pert)
incompressible_hydrostatic_balance(state, b_b, p0)
# interpolate velocity to vector valued function space
W_VectorCG1 = VectorFunctionSpace(mesh, "CG", 1)
uinit = Function(W_VectorCG1).interpolate(as_vector([20.0, 0.0]))
# project to the function space we actually want to use
# this step is purely because it is not yet possible to interpolate to the
def get_latlon_mesh(mesh):
coords_orig = mesh.coordinates
coords_fs = coords_orig.function_space()
if coords_fs.extruded:
cell = mesh._base_mesh.ufl_cell().cellname()
DG1_hori_elt = FiniteElement("DG", cell, 1, variant="equispaced")
DG1_vert_elt = FiniteElement("DG", interval, 1, variant="equispaced")
DG1_elt = TensorProductElement(DG1_hori_elt, DG1_vert_elt)
else:
cell = mesh.ufl_cell().cellname()
DG1_elt = FiniteElement("DG", cell, 1, variant="equispaced")
vec_DG1 = VectorFunctionSpace(mesh, DG1_elt)
coords_dg = Function(vec_DG1).interpolate(coords_orig)
coords_latlon = Function(vec_DG1)
shapes = {"nDOFs": vec_DG1.finat_element.space_dimension(), 'dim': 3}
radius = np.min(
np.sqrt(coords_dg.dat.data[:, 0]**2 + coords_dg.dat.data[:, 1]**2 +
coords_dg.dat.data[:, 2]**2))
# lat-lon 'x' = atan2(y, x)
coords_latlon.dat.data[:, 0] = np.arctan2(coords_dg.dat.data[:, 1],
coords_dg.dat.data[:, 0])
# lat-lon 'y' = asin(z/sqrt(x^2 + y^2 + z^2))
coords_latlon.dat.data[:, 1] = np.arcsin(
coords_dg.dat.data[:, 2] /
np.sqrt(coords_dg.dat.data[:, 0]**2 + coords_dg.dat.data[:, 1]**2 +
coords_dg.dat.data[:, 2]**2))
# our vertical coordinate is radius - the minimum radius
coords_latlon.dat.data[:,
2] = np.sqrt(coords_dg.dat.data[:, 0]**2 +
coords_dg.dat.data[:, 1]**2 +
coords_dg.dat.data[:, 2]**2) - radius
# We need to ensure that all points in a cell are on the same side of the branch cut in longitude coords
# This kernel amends the longitude coords so that all longitudes in one cell are close together
kernel = op2.Kernel(
"""
#define PI 3.141592653589793
#define TWO_PI 6.283185307179586
void splat_coords(double *coords) {{
double max_diff = 0.0;
double diff = 0.0;
for (int i=0; i<{nDOFs}; i++) {{
for (int j=0; j<{nDOFs}; j++) {{
diff = coords[i*{dim}] - coords[j*{dim}];
if (fabs(diff) > max_diff) {{
max_diff = diff;
}}
}}
}}
if (max_diff > PI) {{
for (int i=0; i<{nDOFs}; i++) {{
if (coords[i*{dim}] < 0) {{
coords[i*{dim}] += TWO_PI;
}}
}}
}}
}}
""".format(**shapes), "splat_coords")
op2.par_loop(kernel, coords_latlon.cell_set,
coords_latlon.dat(op2.RW, coords_latlon.cell_node_map()))
return Mesh(coords_latlon)
def setup(self, equation, uadv=None, apply_bcs=True, *active_labels):
self.residual = equation.residual
if self.field_name is not None:
self.idx = equation.field_names.index(self.field_name)
self.fs = self.state.fields(self.field_name).function_space()
self.residual = self.residual.label_map(
lambda t: t.get(prognostic) == self.field_name,
lambda t: Term(
split_form(t.form)[self.idx].form,
t.labels),
drop)
bcs = equation.bcs[self.field_name]
else:
self.field_name = equation.field_name
self.fs = equation.function_space
self.idx = None
if type(self.fs.ufl_element()) is MixedElement:
bcs = [bc for _, bcs in equation.bcs.items() for bc in bcs]
else:
bcs = equation.bcs[self.field_name]
if len(active_labels) > 0:
self.residual = self.residual.label_map(
lambda t: any(t.has_label(time_derivative, *active_labels)),
map_if_false=drop)
options = self.options
# -------------------------------------------------------------------- #
# Routines relating to transport
# -------------------------------------------------------------------- #
if hasattr(self.options, 'ibp'):
self.replace_transport_term()
self.replace_transporting_velocity(uadv)
# -------------------------------------------------------------------- #
# Wrappers for embedded / recovery methods
# -------------------------------------------------------------------- #
if self.discretisation_option in ["embedded_dg", "recovered"]:
# construct the embedding space if not specified
if options.embedding_space is None:
V_elt = BrokenElement(self.fs.ufl_element())
self.fs = FunctionSpace(self.state.mesh, V_elt)
else:
self.fs = options.embedding_space
self.xdg_in = Function(self.fs)
self.xdg_out = Function(self.fs)
if self.idx is None:
self.x_projected = Function(equation.function_space)
else:
self.x_projected = Function(self.state.fields(self.field_name).function_space())
new_test = TestFunction(self.fs)
parameters = {'ksp_type': 'cg',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}
# -------------------------------------------------------------------- #
# Make boundary conditions
# -------------------------------------------------------------------- #
if not apply_bcs:
self.bcs = None
elif self.discretisation_option in ["embedded_dg", "recovered"]:
# Transfer boundary conditions onto test function space
self.bcs = [DirichletBC(self.fs, bc.function_arg, bc.sub_domain) for bc in bcs]
else:
self.bcs = bcs
# -------------------------------------------------------------------- #
# Modify test function for SUPG methods
# -------------------------------------------------------------------- #
if self.discretisation_option == "supg":
# construct tau, if it is not specified
dim = self.state.mesh.topological_dimension()
if options.tau is not None:
# if tau is provided, check that is has the right size
tau = options.tau
assert as_ufl(tau).ufl_shape == (dim, dim), "Provided tau has incorrect shape!"
else:
# create tuple of default values of size dim
default_vals = [options.default*self.dt]*dim
# check for directions is which the space is discontinuous
# so that we don't apply supg in that direction
if is_cg(self.fs):
vals = default_vals
else:
space = self.fs.ufl_element().sobolev_space()
if space.name in ["HDiv", "DirectionalH"]:
vals = [default_vals[i] if space[i].name == "H1"
else 0. for i in range(dim)]
else:
raise ValueError("I don't know what to do with space %s" % space)
tau = Constant(tuple([
tuple(
[vals[j] if i == j else 0. for i, v in enumerate(vals)]
) for j in range(dim)])
)
self.solver_parameters = {'ksp_type': 'gmres',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}
test = TestFunction(self.fs)
new_test = test + dot(dot(uadv, tau), grad(test))
if self.discretisation_option is not None:
# replace the original test function with one defined on
# the embedding space, as this is the space where the
# the problem will be solved
self.residual = self.residual.label_map(
all_terms,
map_if_true=replace_test_function(new_test))
if self.discretisation_option == "embedded_dg":
if self.limiter is None:
self.x_out_projector = Projector(self.xdg_out, self.x_projected,
solver_parameters=parameters)
else:
self.x_out_projector = Recoverer(self.xdg_out, self.x_projected)
if self.discretisation_option == "recovered":
# set up the necessary functions
self.x_in = Function(self.state.fields(self.field_name).function_space())
x_rec = Function(options.recovered_space)
x_brok = Function(options.broken_space)
# set up interpolators and projectors
self.x_rec_projector = Recoverer(self.x_in, x_rec, VDG=self.fs, boundary_method=options.boundary_method) # recovered function
self.x_brok_projector = Projector(x_rec, x_brok) # function projected back
self.xdg_interpolator = Interpolator(self.x_in + x_rec - x_brok, self.xdg_in)
if self.limiter is not None:
self.x_brok_interpolator = Interpolator(self.xdg_out, x_brok)
self.x_out_projector = Recoverer(x_brok, self.x_projected)
else:
self.x_out_projector = Projector(self.xdg_out, self.x_projected)
# setup required functions
self.dq = Function(self.fs)
self.q1 = Function(self.fs)
def __init__(self, state, V):
super(EulerPoincareForm, self).__init__(state)
dt = state.timestepping.dt
w = TestFunction(V)
u = TrialFunction(V)
self.u0 = Function(V)
ustar = 0.5*(self.u0 + u)
n = FacetNormal(state.mesh)
Upwind = 0.5*(sign(dot(self.ubar, n))+1)
if state.mesh.geometric_dimension() == 3:
if V.extruded:
surface_measure = (dS_h + dS_v)
else:
surface_measure = dS
# <w,curl(u) cross ubar + grad( u.ubar)>
# =<curl(u),ubar cross w> - <div(w), u.ubar>
# =<u,curl(ubar cross w)> -
# <<u_upwind, [[n cross(ubar cross w)cross]]>>
both = lambda u: 2*avg(u)
Eqn = (
inner(w, u-self.u0)*dx
+ dt*inner(ustar, curl(cross(self.ubar, w)))*dx
- dt*inner(both(Upwind*ustar),
both(cross(n, cross(self.ubar, w))))*surface_measure
- dt*div(w)*inner(ustar, self.ubar)*dx
)
# define surface measure and terms involving perp differently
# for slice (i.e. if V.extruded is True) and shallow water
# (V.extruded is False)
else:
if V.extruded:
surface_measure = (dS_h + dS_v)
perp = lambda u: as_vector([-u[1], u[0]])
perp_u_upwind = Upwind('+')*perp(ustar('+')) + Upwind('-')*perp(ustar('-'))
else:
surface_measure = dS
outward_normals = CellNormal(state.mesh)
perp = lambda u: cross(outward_normals, u)
perp_u_upwind = Upwind('+')*cross(outward_normals('+'),ustar('+')) + Upwind('-')*cross(outward_normals('-'),ustar('-'))
Eqn = (
(inner(w, u-self.u0)
- dt*inner(w, div(perp(ustar))*perp(self.ubar))
- dt*div(w)*inner(ustar, self.ubar))*dx
- dt*inner(jump(inner(w, perp(self.ubar)), n),
perp_u_upwind)*surface_measure
+ dt*jump(inner(w,
perp(self.ubar))*perp(ustar), n)*surface_measure
)
a = lhs(Eqn)
L = rhs(Eqn)
self.u1 = Function(V)
uproblem = LinearVariationalProblem(a, L, self.u1)
self.usolver = LinearVariationalSolver(uproblem,
options_prefix='EPAdvection')
class RK4(ExplicitTimeDiscretisation):
"""
Class to implement the 4-stage Runge-Kutta timestepping method:
k1 = f(y_n)
k2 = f(y_n + 1/2*dt*k1)
k3 = f(y_n + 1/2*dt*k2)
k4 = f(y_n + dt*k3)
y_(n+1) = y_n + (1/6) * dt * (k1 + 2*k2 + 2*k3 + k4)
where subscripts indicate the timelevel, superscripts indicate the stage
number and f is the RHS.
"""
def setup(self, equation, uadv, *active_labels):
super().setup(equation, uadv, *active_labels)
self.k1 = Function(self.fs)
self.k2 = Function(self.fs)
self.k3 = Function(self.fs)
self.k4 = Function(self.fs)
@cached_property
def lhs(self):
l = self.residual.label_map(
lambda t: t.has_label(time_derivative),
map_if_true=replace_subject(self.dq, self.idx),
map_if_false=drop)
return l.form
@cached_property
def rhs(self):
r = self.residual.label_map(
all_terms,
map_if_true=replace_subject(self.q1, self.idx))
r = r.label_map(
lambda t: t.has_label(time_derivative),
map_if_true=drop,
map_if_false=lambda t: -1*t)
return r.form
def solve_stage(self, x_in, stage):
if stage == 0:
self.solver.solve()
self.k1.assign(self.dq)
self.q1.assign(x_in + 0.5 * self.dt * self.k1)
elif stage == 1:
self.solver.solve()
self.k2.assign(self.dq)
self.q1.assign(x_in + 0.5 * self.dt * self.k2)
elif stage == 2:
self.solver.solve()
self.k3.assign(self.dq)
self.q1.assign(x_in + self.dt * self.k3)
elif stage == 3:
self.solver.solve()
self.k4.assign(self.dq)
self.q1.assign(x_in + 1/6 * self.dt * (self.k1 + 2*self.k2 + 2*self.k3 + self.k4))
def apply_cycle(self, x_in, x_out):
if self.limiter is not None:
self.limiter.apply(x_in)
self.q1.assign(x_in)
for i in range(4):
self.solve_stage(x_in, i)
x_out.assign(self.q1)
class TimeDiscretisation(object, metaclass=ABCMeta):
"""
Base class for time discretisation schemes.
:arg state: :class:`.State` object.
:arg field: field to be evolved
:arg equation: :class:`.Equation` object, specifying the equation
that field satisfies
:arg solver_parameters: solver_parameters
:arg limiter: :class:`.Limiter` object.
:arg options: :class:`.DiscretisationOptions` object
"""
def __init__(self, state, field_name=None, solver_parameters=None,
limiter=None, options=None):
self.state = state
self.field_name = field_name
self.dt = self.state.dt
self.limiter = limiter
self.options = options
if options is not None:
self.discretisation_option = options.name
else:
self.discretisation_option = None
# get default solver options if none passed in
if solver_parameters is None:
self.solver_parameters = {'ksp_type': 'cg',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}
else:
self.solver_parameters = solver_parameters
if logger.isEnabledFor(DEBUG):
self.solver_parameters["ksp_monitor_true_residual"] = None
def setup(self, equation, uadv=None, apply_bcs=True, *active_labels):
self.residual = equation.residual
if self.field_name is not None:
self.idx = equation.field_names.index(self.field_name)
self.fs = self.state.fields(self.field_name).function_space()
self.residual = self.residual.label_map(
lambda t: t.get(prognostic) == self.field_name,
lambda t: Term(
split_form(t.form)[self.idx].form,
t.labels),
drop)
bcs = equation.bcs[self.field_name]
else:
self.field_name = equation.field_name
self.fs = equation.function_space
self.idx = None
if type(self.fs.ufl_element()) is MixedElement:
bcs = [bc for _, bcs in equation.bcs.items() for bc in bcs]
else:
bcs = equation.bcs[self.field_name]
if len(active_labels) > 0:
self.residual = self.residual.label_map(
lambda t: any(t.has_label(time_derivative, *active_labels)),
map_if_false=drop)
options = self.options
# -------------------------------------------------------------------- #
# Routines relating to transport
# -------------------------------------------------------------------- #
if hasattr(self.options, 'ibp'):
self.replace_transport_term()
self.replace_transporting_velocity(uadv)
# -------------------------------------------------------------------- #
# Wrappers for embedded / recovery methods
# -------------------------------------------------------------------- #
if self.discretisation_option in ["embedded_dg", "recovered"]:
# construct the embedding space if not specified
if options.embedding_space is None:
V_elt = BrokenElement(self.fs.ufl_element())
self.fs = FunctionSpace(self.state.mesh, V_elt)
else:
self.fs = options.embedding_space
self.xdg_in = Function(self.fs)
self.xdg_out = Function(self.fs)
if self.idx is None:
self.x_projected = Function(equation.function_space)
else:
self.x_projected = Function(self.state.fields(self.field_name).function_space())
new_test = TestFunction(self.fs)
parameters = {'ksp_type': 'cg',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}
# -------------------------------------------------------------------- #
# Make boundary conditions
# -------------------------------------------------------------------- #
if not apply_bcs:
self.bcs = None
elif self.discretisation_option in ["embedded_dg", "recovered"]:
# Transfer boundary conditions onto test function space
self.bcs = [DirichletBC(self.fs, bc.function_arg, bc.sub_domain) for bc in bcs]
else:
self.bcs = bcs
# -------------------------------------------------------------------- #
# Modify test function for SUPG methods
# -------------------------------------------------------------------- #
if self.discretisation_option == "supg":
# construct tau, if it is not specified
dim = self.state.mesh.topological_dimension()
if options.tau is not None:
# if tau is provided, check that is has the right size
tau = options.tau
assert as_ufl(tau).ufl_shape == (dim, dim), "Provided tau has incorrect shape!"
else:
# create tuple of default values of size dim
default_vals = [options.default*self.dt]*dim
# check for directions is which the space is discontinuous
# so that we don't apply supg in that direction
if is_cg(self.fs):
vals = default_vals
else:
space = self.fs.ufl_element().sobolev_space()
if space.name in ["HDiv", "DirectionalH"]:
vals = [default_vals[i] if space[i].name == "H1"
else 0. for i in range(dim)]
else:
raise ValueError("I don't know what to do with space %s" % space)
tau = Constant(tuple([
tuple(
[vals[j] if i == j else 0. for i, v in enumerate(vals)]
) for j in range(dim)])
)
self.solver_parameters = {'ksp_type': 'gmres',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}
test = TestFunction(self.fs)
new_test = test + dot(dot(uadv, tau), grad(test))
if self.discretisation_option is not None:
# replace the original test function with one defined on
# the embedding space, as this is the space where the
# the problem will be solved
self.residual = self.residual.label_map(
all_terms,
map_if_true=replace_test_function(new_test))
if self.discretisation_option == "embedded_dg":
if self.limiter is None:
self.x_out_projector = Projector(self.xdg_out, self.x_projected,
solver_parameters=parameters)
else:
self.x_out_projector = Recoverer(self.xdg_out, self.x_projected)
if self.discretisation_option == "recovered":
# set up the necessary functions
self.x_in = Function(self.state.fields(self.field_name).function_space())
x_rec = Function(options.recovered_space)
x_brok = Function(options.broken_space)
# set up interpolators and projectors
self.x_rec_projector = Recoverer(self.x_in, x_rec, VDG=self.fs, boundary_method=options.boundary_method) # recovered function
self.x_brok_projector = Projector(x_rec, x_brok) # function projected back
self.xdg_interpolator = Interpolator(self.x_in + x_rec - x_brok, self.xdg_in)
if self.limiter is not None:
self.x_brok_interpolator = Interpolator(self.xdg_out, x_brok)
self.x_out_projector = Recoverer(x_brok, self.x_projected)
else:
self.x_out_projector = Projector(self.xdg_out, self.x_projected)
# setup required functions
self.dq = Function(self.fs)
self.q1 = Function(self.fs)
def pre_apply(self, x_in, discretisation_option):
"""
Extra steps to discretisation if using an embedded method,
which might be either the plain embedded method or the
recovered space scheme.
:arg x_in: the input set of prognostic fields.
:arg discretisation option: string specifying which scheme to use.
"""
if discretisation_option == "embedded_dg":
try:
self.xdg_in.interpolate(x_in)
except NotImplementedError:
self.xdg_in.project(x_in)
elif discretisation_option == "recovered":
self.x_in.assign(x_in)
self.x_rec_projector.project()
self.x_brok_projector.project()
self.xdg_interpolator.interpolate()
def post_apply(self, x_out, discretisation_option):
"""
The projection steps, returning a field to its original space
for an embedded DG scheme. For the case of the
recovered scheme, there are two options dependent on whether
the scheme is limited or not.
:arg x_out: the outgoing field.
:arg discretisation_option: string specifying which option to use.
"""
if discretisation_option == "recovered" and self.limiter is not None:
self.x_brok_interpolator.interpolate()
self.x_out_projector.project()
x_out.assign(self.x_projected)
@abstractproperty
def lhs(self):
l = self.residual.label_map(
lambda t: t.has_label(time_derivative),
map_if_true=replace_subject(self.dq, self.idx),
map_if_false=drop)
return l.form
@abstractproperty
def rhs(self):
r = self.residual.label_map(
all_terms,
map_if_true=replace_subject(self.q1, self.idx))
r = r.label_map(
lambda t: t.has_label(time_derivative),
map_if_false=lambda t: -self.dt*t)
return r.form
def replace_transport_term(self):
"""
This routine allows the default transport term to be replaced with a
different one, specified through the transport options.
This is necessary because when the prognostic equations are declared,
the whole transport
"""
# Extract transport term of equation
old_transport_term_list = self.residual.label_map(
lambda t: t.has_label(transport), map_if_false=drop)
# If there are more transport terms, extract only the one for this variable
if len(old_transport_term_list.terms) > 1:
raise NotImplementedError('Cannot replace transport terms when there are more than one')
# Then we should only have one transport term
old_transport_term = old_transport_term_list.terms[0]
# If the transport term has an ibp label, then it could be replaced
if old_transport_term.has_label(ibp_label) and hasattr(self.options, 'ibp'):
# Do the options specify a different ibp to the old transport term?
if old_transport_term.labels['ibp'] != self.options.ibp:
# Set up a new transport term
field = self.state.fields(self.field_name)
test = TestFunction(self.fs)
# Set up new transport term (depending on the type of transport equation)
if old_transport_term.labels['transport'] == TransportEquationType.advective:
new_transport_term = advection_form(self.state, test, field, ibp=self.options.ibp)
elif old_transport_term.labels['transport'] == TransportEquationType.conservative:
new_transport_term = continuity_form(self.state, test, field, ibp=self.options.ibp)
else:
raise NotImplementedError(f'Replacement of transport term not implemented yet for {old_transport_term.labels["transport"]}')
# Finally, drop the old transport term and add the new one
self.residual = self.residual.label_map(
lambda t: t.has_label(transport), map_if_true=drop)
self.residual += subject(new_transport_term, field)
def replace_transporting_velocity(self, uadv):
# replace the transporting velocity in any terms that contain it
if any([t.has_label(transporting_velocity) for t in self.residual]):
assert uadv is not None
if uadv == "prognostic":
self.residual = self.residual.label_map(
lambda t: t.has_label(transporting_velocity),
map_if_true=lambda t: Term(ufl.replace(
t.form, {t.get(transporting_velocity): split(t.get(subject))[0]}), t.labels)
)
else:
self.residual = self.residual.label_map(
lambda t: t.has_label(transporting_velocity),
map_if_true=lambda t: Term(ufl.replace(
t.form, {t.get(transporting_velocity): uadv}), t.labels)
)
self.residual = transporting_velocity.update_value(self.residual, uadv)
@cached_property
def solver(self):
# setup solver using lhs and rhs defined in derived class
problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.dq, bcs=self.bcs)
solver_name = self.field_name+self.__class__.__name__
return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name)
@abstractmethod
def apply(self, x_in, x_out):
"""
Function takes x as input, computes L(x) as defined by the equation,
and returns x_out as output.
:arg x: :class:`.Function` object, the input Function.
:arg x_out: :class:`.Function` object, the output Function.
"""
pass
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)
class VerticalHybridizationPC(PCBase):
""" A Slate-based python preconditioner for solving
the hydrostatic pressure equation (after rewriting as
a mixed vertical HDiv x L2 system). This preconditioner
hybridizes a mixed system in the vertical direction. This
means that velocities are rendered discontinuous in the
vertical and Lagrange multipliers are introduced
on the top/bottom facets to weakly enforce continuity
through the top/bottom faces of each cell.
This PC assembles a statically condensed problem for the
multipliers and inverts the resulting system using the provided
solver options. The original unknowns are recovered element-wise
by solving local linear systems.
All elimination and recovery kernels are generated using
the Slate DSL in Firedrake.
"""
@timed_function("VertHybridInit")
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 _reconstruction_calls(self, split_mixed_op, split_trace_op):
"""This generates the reconstruction calls for the unknowns using the
Lagrange multipliers.
:arg split_mixed_op: a ``dict`` of split forms that make up the broken
mixed operator from the original problem.
:arg split_trace_op: a ``dict`` of split forms that make up the trace
contribution in the hybridized mixed system.
"""
from firedrake.assemble import create_assembly_callable
# We always eliminate the velocity block first
id0, id1 = (self.vidx, self.pidx)
# TODO: When PyOP2 is able to write into mixed dats,
# the reconstruction expressions can simplify into
# one clean expression.
A = Tensor(split_mixed_op[(id0, id0)])
B = Tensor(split_mixed_op[(id0, id1)])
C = Tensor(split_mixed_op[(id1, id0)])
D = Tensor(split_mixed_op[(id1, id1)])
K_0 = Tensor(split_trace_op[(0, id0)])
K_1 = Tensor(split_trace_op[(0, id1)])
# Split functions and reconstruct each bit separately
split_residual = self.broken_residual.split()
split_sol = self.broken_solution.split()
g = AssembledVector(split_residual[id0])
f = AssembledVector(split_residual[id1])
sigma = split_sol[id0]
u = split_sol[id1]
lambdar = AssembledVector(self.trace_solution)
M = D - C * A.inv * B
R = K_1.T - C * A.inv * K_0.T
u_rec = M.solve(f - C * A.inv * g - R * lambdar,
decomposition="PartialPivLU")
self._sub_unknown = create_assembly_callable(
u_rec, tensor=u, form_compiler_parameters=self.ctx.fc_params)
sigma_rec = A.solve(g - B * AssembledVector(u) - K_0.T * lambdar,
decomposition="PartialPivLU")
self._elim_unknown = create_assembly_callable(
sigma_rec,
tensor=sigma,
form_compiler_parameters=self.ctx.fc_params)
@timed_function("VertHybridRecon")
def _reconstruct(self):
"""Reconstructs the system unknowns using the multipliers.
Note that the reconstruction calls are assumed to be
initialized at this point.
"""
# We assemble the unknown which is an expression
# of the first eliminated variable.
self._sub_unknown()
# Recover the eliminated unknown
self._elim_unknown()
@timed_function("VertHybridUpdate")
def update(self, pc):
"""Update by assembling into the operator. No need to
reconstruct symbolic objects.
"""
self._assemble_S()
self.S.force_evaluation()
def apply(self, pc, x, y):
"""We solve the forward eliminated problem for the
approximate traces of the scalar solution (the multipliers)
and reconstruct the "broken flux and scalar variable."
Lastly, we project the broken solutions into the mimetic
non-broken finite element space.
"""
with timed_region("VertHybridBreak"):
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)
with timed_region("VertHybridRHS"):
# 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, dx, {
"w": (self.weight, READ),
"vec_in": (unbroken_res_hdiv, READ),
"vec_out": (broken_res_hdiv, INC)
})
# Compute the rhs for the multiplier system
self._assemble_Srhs()
with timed_region("VertHybridSolve"):
# Solve the system for the Lagrange multipliers
with self.schur_rhs.dat.vec_ro as b:
if self.trace_ksp.getInitialGuessNonzero():
acc = self.trace_solution.dat.vec
else:
acc = self.trace_solution.dat.vec_wo
with acc as x_trace:
self.trace_ksp.solve(b, x_trace)
# Reconstruct the unknowns
self._reconstruct()
with timed_region("VertHybridRecover"):
# 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, 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 applyTranspose(self, pc, x, y):
"""Apply the transpose of the preconditioner."""
raise NotImplementedError(
"The transpose application of the PC is not implemented.")
def view(self, pc, viewer=None):
"""Viewer calls for the various configurable objects in this PC."""
super(VerticalHybridizationPC, self).view(pc, viewer)
viewer.pushASCIITab()
viewer.printfASCII("Solves K * P^-1 * K.T using local eliminations.\n")
viewer.printfASCII("KSP solver for the multipliers:\n")
viewer.pushASCIITab()
self.trace_ksp.view(viewer)
viewer.popASCIITab()
viewer.printfASCII(
"Locally reconstructing the broken solutions from the multipliers.\n"
)
viewer.pushASCIITab()
viewer.printfASCII(
"Project the broken hdiv solution into the HDiv space.\n")
viewer.popASCIITab()
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()
# first setup the background buoyancy profile
# z.grad(bref) = N**2
# the following is symbolic algebra, using the default buoyancy frequency
# from the parameters class.
x, y, z = SpatialCoordinate(mesh)
g = parameters.g
Nsq = parameters.Nsq
theta_surf = parameters.theta_surf
# N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz)
theta_ref = theta_surf * exp(Nsq * (z - H / 2) / g)
theta_b = Function(Vt).interpolate(theta_ref)
# set theta_pert
def coth(x):
return cosh(x) / sinh(x)
def Z(z):
return Bu * ((z / H) - 0.5)
def n():
return Bu**(-1) * sqrt(
(Bu * 0.5 - tanh(Bu * 0.5)) * (coth(Bu * 0.5) - Bu * 0.5))
Test run: %s.\n
""" % (dt, deltax, deltaz, bool(args.test)))
PETSc.Sys.Print("Initializing problem with dt: %s and tmax: %s.\n" %
(dt, tmax))
PETSc.Sys.Print("Creating mesh with %s columns and %s layers...\n" %
(columns, nlayers))
m = PeriodicIntervalMesh(columns, L)
# build volume mesh
ext_mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers)
Vc = VectorFunctionSpace(ext_mesh, "DG", 2)
coord = SpatialCoordinate(ext_mesh)
x = Function(Vc).interpolate(as_vector([coord[0], coord[1]]))
a = 10000.
xc = L / 2.
x, z = SpatialCoordinate(ext_mesh)
hm = 1.
zs = hm * a**2 / ((x - xc)**2 + a**2)
smooth_z = True
dirname = 'h_mountain'
if smooth_z:
dirname += '_smootherz'
zh = 5000.
xexpr = as_vector(
[x, conditional(z < zh, z + cos(0.5 * pi * z / zh)**6 * zs, z)])
else:
xexpr = as_vector([x, z + ((H - z) / H) * zs])
def __init__(self, state, iterations=1):
super().__init__(state)
self.iterations = iterations
# obtain our fields
self.theta = state.fields('theta')
self.water_v = state.fields('water_v')
self.water_c = state.fields('water_c')
rho = state.fields('rho')
try:
rain = state.fields('rain')
water_l = self.water_c + rain
except NotImplementedError:
water_l = self.water_c
# declare function space
Vt = self.theta.function_space()
# make rho variables
# we recover rho into theta space
if state.vertical_degree == 0 and state.horizontal_degree == 0:
boundary_method = Boundary_Method.physics
else:
boundary_method = None
Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element()))
rho_averaged = Function(Vt)
self.rho_recoverer = Recoverer(rho,
rho_averaged,
VDG=Vt_broken,
boundary_method=boundary_method)
# define some parameters as attributes
dt = state.timestepping.dt
R_d = state.parameters.R_d
cp = state.parameters.cp
cv = state.parameters.cv
c_pv = state.parameters.c_pv
c_pl = state.parameters.c_pl
c_vv = state.parameters.c_vv
R_v = state.parameters.R_v
# make useful fields
Pi = thermodynamics.pi(state.parameters, rho_averaged, self.theta)
T = thermodynamics.T(state.parameters,
self.theta,
Pi,
r_v=self.water_v)
p = thermodynamics.p(state.parameters, Pi)
L_v = thermodynamics.Lv(state.parameters, T)
R_m = R_d + R_v * self.water_v
c_pml = cp + c_pv * self.water_v + c_pl * water_l
c_vml = cv + c_vv * self.water_v + c_pl * water_l
# use Teten's formula to calculate w_sat
w_sat = thermodynamics.r_sat(state.parameters, T, p)
# make appropriate condensation rate
dot_r_cond = ((self.water_v - w_sat) / (dt * (1.0 +
((L_v**2.0 * w_sat) /
(cp * R_v * T**2.0)))))
# make cond_rate function, that needs to be the same for all updates in one time step
cond_rate = Function(Vt)
# adjust cond rate so negative concentrations don't occur
self.lim_cond_rate = Interpolator(
conditional(dot_r_cond < 0,
max_value(dot_r_cond, -self.water_c / dt),
min_value(dot_r_cond, self.water_v / dt)), cond_rate)
# tell the prognostic fields what to update to
self.water_v_new = Interpolator(self.water_v - dt * cond_rate, Vt)
self.water_c_new = Interpolator(self.water_c + dt * cond_rate, Vt)
self.theta_new = Interpolator(
self.theta * (1.0 + dt * cond_rate *
(cv * L_v / (c_vml * cp * T) - R_v * cv * c_pml /
(R_m * cp * c_vml))), Vt)
def nonlocal_integral_eq(
mesh,
scatterer_bdy_id,
outer_bdy_id,
wave_number,
options_prefix=None,
solver_parameters=None,
fspace=None,
vfspace=None,
true_sol_grad=None,
queue=None,
fspace_analog=None,
qbx_kwargs=None,
):
r"""
see run_method for descriptions of unlisted args
args:
:arg queue: A command queue for the computing context
gamma and beta are used to precondition
with the following equation:
\Delta u - \kappa^2 \gamma u = 0
(\partial_n - i\kappa\beta) u |_\Sigma = 0
"""
with_refinement = True
# away from the excluded region, but firedrake and meshmode point
# into
pyt_inner_normal_sign = -1
ambient_dim = mesh.geometric_dimension()
# {{{ Create operator
from pytential import sym
r"""
..math:
x \in \Sigma
grad_op(x) =
\nabla(
\int_\Gamma(
u(y) \partial_n H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
)
"""
grad_op = pyt_inner_normal_sign * sym.grad(
ambient_dim,
sym.D(HelmholtzKernel(ambient_dim),
sym.var("u"),
k=sym.var("k"),
qbx_forced_limit=None))
r"""
..math:
x \in \Sigma
op(x) =
i \kappa \cdot
\int_\Gamma(
u(y) \partial_n H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
"""
op = pyt_inner_normal_sign * 1j * sym.var("k") * (sym.D(
HelmholtzKernel(ambient_dim),
sym.var("u"),
k=sym.var("k"),
qbx_forced_limit=None))
pyt_grad_op = fd2mm.fd_bind(
queue.context,
fspace_analog,
grad_op,
source=(fspace, scatterer_bdy_id),
target=(vfspace, outer_bdy_id),
with_refinement=with_refinement,
qbx_kwargs=qbx_kwargs,
)
pyt_op = fd2mm.fd_bind(
queue.context,
fspace_analog,
op,
source=(fspace, scatterer_bdy_id),
target=(fspace, outer_bdy_id),
with_refinement=with_refinement,
qbx_kwargs=qbx_kwargs,
)
# }}}
class MatrixFreeB(object):
def __init__(self, A, pyt_grad_op, pyt_op, queue, kappa):
"""
:arg kappa: The wave number
"""
self.queue = queue
self.k = kappa
self.pyt_op = pyt_op
self.pyt_grad_op = pyt_grad_op
self.A = A
# {{{ Create some functions needed for multing
self.x_fntn = Function(fspace)
self.potential_int = Function(fspace)
self.potential_int.dat.data[:] = 0.0
self.grad_potential_int = Function(vfspace)
self.grad_potential_int.dat.data[:] = 0.0
self.pyt_result = Function(fspace)
self.n = FacetNormal(mesh)
self.v = TestFunction(fspace)
# }}}
def mult(self, mat, x, y):
# Perform pytential operation
self.x_fntn.dat.data[:] = x[:]
self.pyt_op(self.queue,
self.potential_int,
u=self.x_fntn,
k=self.k)
self.pyt_grad_op(self.queue,
self.grad_potential_int,
u=self.x_fntn,
k=self.k)
# Integrate the potential
r"""
Compute the inner products using firedrake. Note this
will be subtracted later, hence appears off by a sign.
.. math::
\langle
n(x) \cdot \nabla(
\int_\Gamma(
u(y) \partial_n H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
), v
\rangle_\Sigma
- \langle
i \kappa \cdot
\int_\Gamma(
u(y) \partial_n H_0^{(1)}(\kappa |x - y|)
)d\gamma(y), v
\rangle_\Sigma
"""
self.pyt_result = assemble(
inner(inner(self.grad_potential_int, self.n), self.v) *
ds(outer_bdy_id) -
inner(self.potential_int, self.v) * ds(outer_bdy_id))
# y <- Ax - evaluated potential
self.A.mult(x, y)
with self.pyt_result.dat.vec_ro as ep:
y.axpy(-1, ep)
# {{{ Compute normal helmholtz operator
u = TrialFunction(fspace)
v = TestFunction(fspace)
r"""
.. math::
\langle
\nabla u, \nabla v
\rangle
- \kappa^2 \cdot \langle
u, v
\rangle
- i \kappa \langle
u, v
\rangle_\Sigma
"""
a = inner(grad(u), grad(v)) * dx \
- Constant(wave_number**2) * inner(u, v) * dx \
- Constant(1j * wave_number) * inner(u, v) * ds(outer_bdy_id)
# get the concrete matrix from a general bilinear form
A = assemble(a).M.handle
# }}}
# {{{ Setup Python matrix
B = PETSc.Mat().create()
# build matrix context
Bctx = MatrixFreeB(A, pyt_grad_op, pyt_op, queue, wave_number)
# set up B as same size as A
B.setSizes(*A.getSizes())
B.setType(B.Type.PYTHON)
B.setPythonContext(Bctx)
B.setUp()
# }}}
# {{{ Create rhs
# Remember f is \partial_n(true_sol)|_\Gamma
# so we just need to compute \int_\Gamma\partial_n(true_sol) H(x-y)
from pytential import sym
sigma = sym.make_sym_vector("sigma", ambient_dim)
r"""
..math:
x \in \Sigma
grad_op(x) =
\nabla(
\int_\Gamma(
f(y) H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
)
"""
grad_op = pyt_inner_normal_sign * \
sym.grad(ambient_dim, sym.S(HelmholtzKernel(ambient_dim),
sym.n_dot(sigma),
k=sym.var("k"), qbx_forced_limit=None))
r"""
..math:
x \in \Sigma
op(x) =
i \kappa \cdot
\int_\Gamma(
f(y) H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
)
"""
op = 1j * sym.var("k") * pyt_inner_normal_sign * \
sym.S(HelmholtzKernel(ambient_dim),
sym.n_dot(sigma),
k=sym.var("k"),
qbx_forced_limit=None)
rhs_grad_op = fd2mm.fd_bind(
queue.context,
fspace_analog,
grad_op,
source=(vfspace, scatterer_bdy_id),
target=(vfspace, outer_bdy_id),
with_refinement=with_refinement,
qbx_kwargs=qbx_kwargs,
)
rhs_op = fd2mm.fd_bind(
queue.context,
fspace_analog,
op,
source=(vfspace, scatterer_bdy_id),
target=(fspace, outer_bdy_id),
with_refinement=with_refinement,
qbx_kwargs=qbx_kwargs,
)
f_grad_convoluted = Function(vfspace)
f_convoluted = Function(fspace)
rhs_grad_op(queue, f_grad_convoluted, sigma=true_sol_grad, k=wave_number)
rhs_op(queue, f_convoluted, sigma=true_sol_grad, k=wave_number)
r"""
\langle
f, v
\rangle_\Gamma
+ \langle
i \kappa \cdot \int_\Gamma(
f(y) H_0^{(1)}(\kappa |x - y|)
)d\gamma(y), v
\rangle_\Sigma
- \langle
n(x) \cdot \nabla(
\int_\Gamma(
f(y) H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
), v
\rangle_\Sigma
"""
rhs_form = inner(inner(true_sol_grad, FacetNormal(mesh)),
v) * ds(scatterer_bdy_id) \
+ inner(f_convoluted, v) * ds(outer_bdy_id) \
- inner(inner(f_grad_convoluted, FacetNormal(mesh)),
v) * ds(outer_bdy_id)
rhs = assemble(rhs_form)
# {{{ set up a solver:
solution = Function(fspace, name="Computed Solution")
# {{{ Used for preconditioning
if 'gamma' in solver_parameters or 'beta' in solver_parameters:
gamma = complex(solver_parameters.pop('gamma', 1.0))
import cmath
beta = complex(solver_parameters.pop('beta', cmath.sqrt(gamma)))
p = inner(grad(u), grad(v)) * dx \
- Constant(wave_number**2 * gamma) * inner(u, v) * dx \
- Constant(1j * wave_number * beta) * inner(u, v) * ds(outer_bdy_id)
P = assemble(p).M.handle
else:
P = A
# }}}
# Set up options to contain solver parameters:
ksp = PETSc.KSP().create()
if solver_parameters['pc_type'] == 'pyamg':
del solver_parameters['pc_type'] # We are using the AMG preconditioner
pyamg_tol = solver_parameters.get('pyamg_tol', None)
if pyamg_tol is not None:
pyamg_tol = float(pyamg_tol)
pyamg_maxiter = solver_parameters.get('pyamg_maxiter', None)
if pyamg_maxiter is not None:
pyamg_maxiter = int(pyamg_maxiter)
ksp.setOperators(B)
ksp.setUp()
pc = ksp.pc
pc.setType(pc.Type.PYTHON)
pc.setPythonContext(
AMGTransmissionPreconditioner(wave_number,
fspace,
A,
tol=pyamg_tol,
maxiter=pyamg_maxiter,
use_plane_waves=True))
# Otherwise use regular preconditioner
else:
ksp.setOperators(B, P)
options_manager = OptionsManager(solver_parameters, options_prefix)
options_manager.set_from_options(ksp)
with rhs.dat.vec_ro as b:
with solution.dat.vec as x:
ksp.solve(b, x)
# }}}
return ksp, solution
def setup_3d_recovery(dirname):
L = 3.
H = 3.
W = 3.
deltax = L / 3.
deltay = W / 3.
deltaz = H / 3.
nlayers = int(H / deltaz)
ncolumnsx = int(L / deltax)
ncolumnsy = int(W / deltay)
m = PeriodicRectangleMesh(ncolumnsx,
ncolumnsy,
L,
W,
direction='y',
quadrilateral=True)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers)
x, y, z = SpatialCoordinate(mesh)
# horizontal base spaces
cell = mesh._base_mesh.ufl_cell().cellname()
u_hori = FiniteElement("RTCF", cell, 1)
w_hori = FiniteElement("DG", cell, 0)
# vertical base spaces
u_vert = FiniteElement("DG", interval, 0)
w_vert = FiniteElement("CG", interval, 1)
# build elements
u_element = HDiv(TensorProductElement(u_hori, u_vert))
w_element = HDiv(TensorProductElement(w_hori, w_vert))
theta_element = TensorProductElement(w_hori, w_vert)
v_element = u_element + w_element
# spaces
VDG0 = FunctionSpace(mesh, "DG", 0)
VCG1 = FunctionSpace(mesh, "CG", 1)
VDG1 = FunctionSpace(mesh, "DG", 1)
Vt = FunctionSpace(mesh, theta_element)
Vt_brok = FunctionSpace(mesh, BrokenElement(theta_element))
Vu = FunctionSpace(mesh, v_element)
VuCG1 = VectorFunctionSpace(mesh, "CG", 1)
VuDG1 = VectorFunctionSpace(mesh, "DG", 1)
# set up initial conditions
np.random.seed(0)
expr = np.random.randn() + np.random.randn() * x + np.random.randn(
) * z + np.random.randn() * x * z
# our actual theta and rho and v
rho_CG1_true = Function(VCG1).interpolate(expr)
theta_CG1_true = Function(VCG1).interpolate(expr)
v_CG1_true = Function(VuCG1).interpolate(as_vector([expr, expr, expr]))
rho_Vt_true = Function(Vt).interpolate(expr)
# make the initial fields by projecting expressions into the lowest order spaces
rho_DG0 = Function(VDG0).interpolate(expr)
rho_CG1 = Function(VCG1)
theta_Vt = Function(Vt).interpolate(expr)
theta_CG1 = Function(VCG1)
v_Vu = Function(Vu).project(as_vector([expr, expr, expr]))
v_CG1 = Function(VuCG1)
rho_Vt = Function(Vt)
# make the recoverers and do the recovery
rho_recoverer = Recoverer(rho_DG0,
rho_CG1,
VDG=VDG1,
boundary_method=Boundary_Method.dynamics)
theta_recoverer = Recoverer(theta_Vt,
theta_CG1,
VDG=VDG1,
boundary_method=Boundary_Method.dynamics)
v_recoverer = Recoverer(v_Vu,
v_CG1,
VDG=VuDG1,
boundary_method=Boundary_Method.dynamics)
rho_Vt_recoverer = Recoverer(rho_DG0,
rho_Vt,
VDG=Vt_brok,
boundary_method=Boundary_Method.physics)
rho_recoverer.project()
theta_recoverer.project()
v_recoverer.project()
rho_Vt_recoverer.project()
rho_diff = errornorm(rho_CG1, rho_CG1_true) / norm(rho_CG1_true)
theta_diff = errornorm(theta_CG1, theta_CG1_true) / norm(theta_CG1_true)
v_diff = errornorm(v_CG1, v_CG1_true) / norm(v_CG1_true)
rho_Vt_diff = errornorm(rho_Vt, rho_Vt_true) / norm(rho_Vt_true)
return (rho_diff, theta_diff, v_diff, rho_Vt_diff)
def setup_dump(self, t, tmax, pickup=False):
"""
Setup dump files
Check for existence of directory so as not to overwrite
output files
Setup checkpoint file
:arg tmax: model stop time
:arg pickup: recover state from the checkpointing file if true,
otherwise dump and checkpoint to disk. (default is False).
"""
if any([
self.output.dump_vtus, self.output.dumplist_latlon,
self.output.dump_diagnostics, self.output.point_data,
self.output.checkpoint and not pickup
]):
# setup output directory and check that it does not already exist
self.dumpdir = path.join("results", self.output.dirname)
running_tests = '--running-tests' in sys.argv or "pytest" in self.output.dirname
if self.mesh.comm.rank == 0:
if not running_tests and path.exists(
self.dumpdir) and not pickup:
raise IOError("results directory '%s' already exists" %
self.dumpdir)
else:
if not running_tests:
makedirs(self.dumpdir)
if self.output.dump_vtus:
# setup pvd output file
outfile = path.join(self.dumpdir, "field_output.pvd")
self.dumpfile = File(outfile,
project_output=self.output.project_fields,
comm=self.mesh.comm)
# make list of fields to dump
self.to_dump = [field for field in self.fields if field.dump]
# make dump counter
self.dumpcount = itertools.count()
# if there are fields to be dumped in latlon coordinates,
# setup the latlon coordinate mesh and make output file
if len(self.output.dumplist_latlon) > 0:
mesh_ll = get_latlon_mesh(self.mesh)
outfile_ll = path.join(self.dumpdir, "field_output_latlon.pvd")
self.dumpfile_ll = File(outfile_ll,
project_output=self.output.project_fields,
comm=self.mesh.comm)
# make functions on latlon mesh, as specified by dumplist_latlon
self.to_dump_latlon = []
for name in self.output.dumplist_latlon:
f = self.fields(name)
field = Function(functionspaceimpl.WithGeometry(
f.function_space(), mesh_ll),
val=f.topological,
name=name + '_ll')
self.to_dump_latlon.append(field)
# we create new netcdf files to write to, unless pickup=True, in
# which case we just need the filenames
if self.output.dump_diagnostics:
diagnostics_filename = self.dumpdir + "/diagnostics.nc"
self.diagnostic_output = DiagnosticsOutput(diagnostics_filename,
self.diagnostics,
self.output.dirname,
self.mesh.comm,
create=not pickup)
if len(self.output.point_data) > 0:
pointdata_filename = self.dumpdir + "/point_data.nc"
ndt = int(tmax / self.timestepping.dt)
self.pointdata_output = PointDataOutput(pointdata_filename,
ndt,
self.output.point_data,
self.output.dirname,
self.fields,
self.mesh.comm,
create=not pickup)
# if we want to checkpoint and are not picking up from a previous
# checkpoint file, setup the dumb checkpointing
if self.output.checkpoint and not pickup:
self.chkpt = DumbCheckpoint(path.join(self.dumpdir, "chkpt"),
mode=FILE_CREATE)
# make list of fields to pickup (this doesn't include
# diagnostic fields)
self.to_pickup = [field for field in self.fields if field.pickup]
# if we want to checkpoint then make a checkpoint counter
if self.output.checkpoint:
self.chkptcount = itertools.count()
# dump initial fields
self.dump(t)
class Advection(object, metaclass=ABCMeta):
"""
Base class for advection schemes.
:arg state: :class:`.State` object.
:arg field: field to be advected
:arg equation: :class:`.Equation` object, specifying the equation
that field satisfies
:arg solver_parameters: solver_parameters
:arg limiter: :class:`.Limiter` object.
:arg options: :class:`.AdvectionOptions` object
"""
def __init__(self, state, field, equation=None, *, solver_parameters=None,
limiter=None):
if equation is not None:
self.state = state
self.field = field
self.equation = equation
# get ubar from the equation class
self.ubar = self.equation.ubar
self.dt = self.state.timestepping.dt
# get default solver options if none passed in
if solver_parameters is None:
self.solver_parameters = equation.solver_parameters
else:
self.solver_parameters = solver_parameters
if logger.isEnabledFor(DEBUG):
self.solver_parameters["ksp_monitor_true_residual"] = True
self.limiter = limiter
if hasattr(equation, "options"):
self.discretisation_option = equation.options.name
self._setup(state, field, equation.options)
else:
self.discretisation_option = None
self.fs = field.function_space()
# setup required functions
self.dq = Function(self.fs)
self.q1 = Function(self.fs)
def _setup(self, state, field, options):
if options.name in ["embedded_dg", "recovered"]:
self.fs = options.embedding_space
self.xdg_in = Function(self.fs)
self.xdg_out = Function(self.fs)
self.x_projected = Function(field.function_space())
parameters = {'ksp_type': 'cg',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}
self.Projector = Projector(self.xdg_out, self.x_projected,
solver_parameters=parameters)
if options.name == "recovered":
# set up the necessary functions
self.x_in = Function(field.function_space())
x_rec = Function(options.recovered_space)
x_brok = Function(options.broken_space)
# set up interpolators and projectors
self.x_rec_projector = Recoverer(self.x_in, x_rec, VDG=self.fs, boundary_method=options.boundary_method) # recovered function
self.x_brok_projector = Projector(x_rec, x_brok) # function projected back
self.xdg_interpolator = Interpolator(self.x_in + x_rec - x_brok, self.xdg_in)
if self.limiter is not None:
self.x_brok_interpolator = Interpolator(self.xdg_out, x_brok)
self.x_out_projector = Recoverer(x_brok, self.x_projected)
def pre_apply(self, x_in, discretisation_option):
"""
Extra steps to advection if using an embedded method,
which might be either the plain embedded method or the
recovered space advection scheme.
:arg x_in: the input set of prognostic fields.
:arg discretisation option: string specifying which scheme to use.
"""
if discretisation_option == "embedded_dg":
try:
self.xdg_in.interpolate(x_in)
except NotImplementedError:
self.xdg_in.project(x_in)
elif discretisation_option == "recovered":
self.x_in.assign(x_in)
self.x_rec_projector.project()
self.x_brok_projector.project()
self.xdg_interpolator.interpolate()
def post_apply(self, x_out, discretisation_option):
"""
The projection steps, returning a field to its original space
for an embedded DG advection scheme. For the case of the
recovered scheme, there are two options dependent on whether
the scheme is limited or not.
:arg x_out: the outgoing field.
:arg discretisation_option: string specifying which option to use.
"""
if discretisation_option == "embedded_dg":
self.Projector.project()
elif discretisation_option == "recovered":
if self.limiter is not None:
self.x_brok_interpolator.interpolate()
self.x_out_projector.project()
else:
self.Projector.project()
x_out.assign(self.x_projected)
@abstractproperty
def lhs(self):
return self.equation.mass_term(self.equation.trial)
@abstractproperty
def rhs(self):
return self.equation.mass_term(self.q1) - self.dt*self.equation.advection_term(self.q1)
def update_ubar(self, xn, xnp1, alpha):
un = xn.split()[0]
unp1 = xnp1.split()[0]
self.ubar.assign(un + alpha*(unp1-un))
@cached_property
def solver(self):
# setup solver using lhs and rhs defined in derived class
problem = LinearVariationalProblem(self.lhs, self.rhs, self.dq)
solver_name = self.field.name()+self.equation.__class__.__name__+self.__class__.__name__
return LinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name)
@abstractmethod
def apply(self, x_in, x_out):
"""
Function takes x as input, computes L(x) as defined by the equation,
and returns x_out as output.
:arg x: :class:`.Function` object, the input Function.
:arg x_out: :class:`.Function` object, the output Function.
"""
pass
def __init__(self, state):
super().__init__(state)
# obtain our fields
self.theta = state.fields('theta')
self.water_v = state.fields('water_v')
self.rain = state.fields('rain')
rho = state.fields('rho')
try:
water_c = state.fields('water_c')
water_l = self.rain + water_c
except NotImplementedError:
water_l = self.rain
# declare function space
Vt = self.theta.function_space()
# make rho variables
# we recover rho into theta space
if state.vertical_degree == 0 and state.horizontal_degree == 0:
boundary_method = Boundary_Method.physics
else:
boundary_method = None
Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element()))
rho_averaged = Function(Vt)
self.rho_recoverer = Recoverer(rho,
rho_averaged,
VDG=Vt_broken,
boundary_method=boundary_method)
# define some parameters as attributes
dt = state.timestepping.dt
R_d = state.parameters.R_d
cp = state.parameters.cp
cv = state.parameters.cv
c_pv = state.parameters.c_pv
c_pl = state.parameters.c_pl
c_vv = state.parameters.c_vv
R_v = state.parameters.R_v
# make useful fields
Pi = thermodynamics.pi(state.parameters, rho_averaged, self.theta)
T = thermodynamics.T(state.parameters,
self.theta,
Pi,
r_v=self.water_v)
p = thermodynamics.p(state.parameters, Pi)
L_v = thermodynamics.Lv(state.parameters, T)
R_m = R_d + R_v * self.water_v
c_pml = cp + c_pv * self.water_v + c_pl * water_l
c_vml = cv + c_vv * self.water_v + c_pl * water_l
# use Teten's formula to calculate w_sat
w_sat = thermodynamics.r_sat(state.parameters, T, p)
# expression for ventilation factor
a = Constant(1.6)
b = Constant(124.9)
c = Constant(0.2046)
C = a + b * (rho_averaged * self.rain)**c
# make appropriate condensation rate
f = Constant(5.4e5)
g = Constant(2.55e6)
h = Constant(0.525)
dot_r_evap = (((1 - self.water_v / w_sat) * C *
(rho_averaged * self.rain)**h) /
(rho_averaged * (f + g / (p * w_sat))))
# make evap_rate function, needs to be the same for all updates in one time step
evap_rate = Function(Vt)
# adjust evap rate so negative rain doesn't occur
self.lim_evap_rate = Interpolator(
conditional(
dot_r_evap < 0, 0.0,
conditional(self.rain < 0.0, 0.0,
min_value(dot_r_evap, self.rain / dt))), evap_rate)
# tell the prognostic fields what to update to
self.water_v_new = Interpolator(self.water_v + dt * evap_rate, Vt)
self.rain_new = Interpolator(self.rain - dt * evap_rate, Vt)
self.theta_new = Interpolator(
self.theta * (1.0 - dt * evap_rate *
(cv * L_v / (c_vml * cp * T) - R_v * cv * c_pml /
(R_m * cp * c_vml))), Vt)
u_bc_ids=[1, 2])
# Initial conditions
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()
x, z = SpatialCoordinate(mesh)
# Define constant theta_e and water_t
Tsurf = 300.0
theta_b = Function(Vt).interpolate(Constant(Tsurf))
# Calculate hydrostatic fields
compressible_hydrostatic_balance(state, theta_b, rho0, solve_for_rho=True)
# make mean fields
rho_b = Function(Vr).assign(rho0)
# define perturbation
xc = L / 2
zc = 2000.
rc = 2000.
Tdash = 2.0
r = sqrt((x - xc)**2 + (z - zc)**2)
theta_pert = Function(Vt).interpolate(
conditional(r > rc, 0.0,
def setup_vert_limiters(dirname):
# declare grid shape
L = 200.
H = 800.
ncolumns = int(L / 20.)
nlayers = int(H / 20.)
# make mesh
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
x, z = SpatialCoordinate(mesh)
fieldlist = ['u']
timestepping = TimesteppingParameters(dt=1.0, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname + "/limiting_vert",
dumpfreq=5,
dumplist=['u'],
perturbation_fields=['theta0', 'theta1'])
parameters = CompressibleParameters()
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist)
# make elements
# v is continuous in vertical, h is horizontal
cell = mesh._base_mesh.ufl_cell().cellname()
DG0_element = FiniteElement("DG", cell, 0)
CG1_element = FiniteElement("CG", cell, 1)
DG1_element = FiniteElement("DG", cell, 1)
CG2_element = FiniteElement("CG", cell, 2)
V0_element = TensorProductElement(DG0_element, CG1_element)
V1_element = TensorProductElement(DG1_element, CG2_element)
# spaces
VDG1 = FunctionSpace(mesh, "DG", 1)
VCG1 = FunctionSpace(mesh, "CG", 1)
V0 = FunctionSpace(mesh, V0_element)
V1 = FunctionSpace(mesh, V1_element)
V0_brok = FunctionSpace(mesh, BrokenElement(V0.ufl_element()))
V0_spaces = (VDG1, VCG1, V0_brok)
# declare initial fields
u0 = state.fields("u")
theta0 = state.fields("theta0", V0)
theta1 = state.fields("theta1", V1)
# Isentropic background state
Tsurf = 0.
thetab = Constant(Tsurf)
theta_b1 = Function(V1).interpolate(thetab)
theta_b0 = Function(V0).interpolate(thetab)
# set up bubble
xc = L / 2
zc = 700.
rc = 80.
theta_expr = conditional(
sqrt((x - xc)**2.0) < rc,
conditional(sqrt((z - zc)**2.0) < rc, Constant(2.0), Constant(0.0)),
Constant(0.0))
theta_pert1 = Function(V1).interpolate(theta_expr)
theta_pert0 = Function(V0).interpolate(theta_expr)
# set up velocity field
u_max = Constant(5.0)
u0.project(as_vector([0, -u_max]))
theta0.interpolate(theta_b0 + theta_pert0)
theta1.interpolate(theta_b1 + theta_pert1)
state.initialise([('u', u0), ('theta1', theta1), ('theta0', theta0)])
state.set_reference_profiles([('theta1', theta_b1), ('theta0', theta_b0)])
# set up advection schemes
thetaeqn1 = EmbeddedDGAdvection(state, V1, equation_form="advective")
thetaeqn0 = EmbeddedDGAdvection(state,
V0,
equation_form="advective",
recovered_spaces=V0_spaces)
# build advection dictionary
advected_fields = []
advected_fields.append(('u', NoAdvection(state, u0, None)))
advected_fields.append(('theta1',
SSPRK3(state,
theta1,
thetaeqn1,
limiter=ThetaLimiter(thetaeqn1))))
advected_fields.append(('theta0',
SSPRK3(state,
theta0,
thetaeqn0,
limiter=VertexBasedLimiter(VDG1))))
# build time stepper
stepper = AdvectionDiffusion(state, advected_fields)
return stepper, 40.0
def _diagonal(self):
from firedrake import Function
assert self.on_diag
return Function(self._x.function_space())
eqns = CompressibleEulerEquations(state, "CG", 1)
# Initial conditions
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# spaces
Vu = state.spaces("HDiv")
Vt = state.spaces("theta")
Vr = state.spaces("DG")
# Isentropic background state
Tsurf = Constant(300.)
theta_b = Function(Vt).interpolate(Tsurf)
rho_b = Function(Vr)
# Calculate hydrostatic exner
compressible_hydrostatic_balance(state, theta_b, rho_b, solve_for_rho=True)
x = SpatialCoordinate(mesh)
xc = 500.
zc = 350.
rc = 250.
r = sqrt((x[0]-xc)**2 + (x[1]-zc)**2)
theta_pert = conditional(r > rc, 0., 0.25*(1. + cos((pi/rc)*r)))
theta0.interpolate(theta_b + theta_pert)
rho0.interpolate(rho_b)
class DGAdvection(Advection):
"""
DG 3 step SSPRK advection scheme that can be applied to a scalar
or vector field
:arg state: :class:`.State` object.
:arg V: function space of advected field - should be DG
:arg continuity: optional boolean.
If ``True``, the advection equation is of the form:
:math: `D_t +\nabla \cdot(uD) = 0`.
If ``False``, the advection equation is of the form:
:math: `D_t + (u \cdot \nabla)D = 0`.
"""
def __init__(self, state, V, continuity=False):
super(DGAdvection, self).__init__(state)
element = V.fiat_element
assert element.entity_dofs() == element.entity_closure_dofs(), "Provided space is not discontinuous"
dt = state.timestepping.dt
if V.extruded:
surface_measure = (dS_h + dS_v)
else:
surface_measure = dS
phi = TestFunction(V)
D = TrialFunction(V)
self.D1 = Function(V)
self.dD = Function(V)
n = FacetNormal(state.mesh)
# ( dot(v, n) + |dot(v, n)| )/2.0
un = 0.5*(dot(self.ubar, n) + abs(dot(self.ubar, n)))
a_mass = inner(phi,D)*dx
if continuity:
a_int = -inner(grad(phi), outer(D, self.ubar))*dx
else:
a_int = -inner(div(outer(phi,self.ubar)),D)*dx
a_flux = (dot(jump(phi), un('+')*D('+') - un('-')*D('-')))*surface_measure
arhs = a_mass - dt*(a_int + a_flux)
DGproblem = LinearVariationalProblem(a_mass, action(arhs,self.D1),
self.dD)
self.DGsolver = LinearVariationalSolver(DGproblem,
solver_parameters={
'ksp_type':'preonly',
'pc_type':'bjacobi',
'sub_pc_type': 'ilu'},
options_prefix='DGAdvection')
def apply(self, x_in, x_out):
# SSPRK Stage 1
self.D1.assign(x_in)
self.DGsolver.solve()
self.D1.assign(self.dD)
# SSPRK Stage 2
self.DGsolver.solve()
self.D1.assign(0.75*x_in + 0.25*self.dD)
# SSPRK Stage 3
self.DGsolver.solve()
x_out.assign((1.0/3.0)*x_in + (2.0/3.0)*self.dD)
def __init__(self, state, field, equation=None, *, solver_parameters=None,
limiter=None):
if equation is not None:
self.state = state
self.field = field
self.equation = equation
# get ubar from the equation class
self.ubar = self.equation.ubar
self.dt = self.state.timestepping.dt
# get default solver options if none passed in
if solver_parameters is None:
self.solver_parameters = equation.solver_parameters
else:
self.solver_parameters = solver_parameters
if state.output.log_level == DEBUG:
self.solver_parameters["ksp_monitor_true_residual"] = True
self.limiter = limiter
# check to see if we are using an embedded DG method - if we are then
# the projector and output function will have been set up in the
# equation class and we can get the correct function space from
# the output function.
if isinstance(equation, EmbeddedDGAdvection):
# check that the field and the equation are compatible
if equation.V0 != field.function_space():
raise ValueError('The field to be advected is not compatible with the equation used.')
self.embedded_dg = True
fs = equation.space
self.xdg_in = Function(fs)
self.xdg_out = Function(fs)
self.x_projected = Function(field.function_space())
parameters = {'ksp_type': 'cg',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}
self.Projector = Projector(self.xdg_out, self.x_projected,
solver_parameters=parameters)
self.recovered = equation.recovered
if self.recovered:
# set up the necessary functions
self.x_in = Function(field.function_space())
x_adv = Function(fs)
x_rec = Function(equation.V_rec)
x_brok = Function(equation.V_brok)
# set up interpolators and projectors
self.x_adv_interpolator = Interpolator(self.x_in, x_adv) # interpolate before recovery
self.x_rec_projector = Recoverer(x_adv, x_rec) # recovered function
# when the "average" method comes into firedrake master, this will be
# self.x_rec_projector = Projector(self.x_in, equation.Vrec, method="average")
self.x_brok_projector = Projector(x_rec, x_brok) # function projected back
self.xdg_interpolator = Interpolator(self.x_in + x_rec - x_brok, self.xdg_in)
if self.limiter is not None:
self.x_brok_interpolator = Interpolator(self.xdg_out, x_brok)
self.x_out_projector = Recoverer(x_brok, self.x_projected)
# when the "average" method comes into firedrake master, this will be
# self.x_out_projector = Projector(x_brok, self.x_projected, method="average")
else:
self.embedded_dg = False
fs = field.function_space()
# setup required functions
self.fs = fs
self.dq = Function(fs)
self.q1 = Function(fs)
class SUPGAdvection(Advection):
"""
An SUPG advection scheme that can apply DG upwinding (in the direction
specified by the direction arg) if the function space is only
partially continuous.
:arg state: :class:`.State` object.
:arg V:class:`.FunctionSpace` object. The advected field function space.
:arg direction: list containing the directions in which the function
space is discontinuous. 1 corresponds to the vertical direction, 2 to
the horizontal direction
:arg supg_params: dictionary containing SUPG parameters tau for each
direction. If not supplied tau is set to dt/sqrt(15.)
"""
def __init__(self, state, V, direction=[], supg_params=None):
super(SUPGAdvection, self).__init__(state)
dt = state.timestepping.dt
params = supg_params.copy() if supg_params else {}
params.setdefault('a0', dt/sqrt(15.))
params.setdefault('a1', dt/sqrt(15.))
gamma = TestFunction(V)
theta = TrialFunction(V)
self.theta0 = Function(V)
# make SUPG test function
taus = [params["a0"], params["a1"]]
for i in direction:
taus[i] = 0.0
tau = Constant(((taus[0], 0.), (0., taus[1])))
dgamma = dot(dot(self.ubar, tau), grad(gamma))
gammaSU = gamma + dgamma
n = FacetNormal(state.mesh)
un = 0.5*(dot(self.ubar, n) + abs(dot(self.ubar, n)))
a_mass = gammaSU*theta*dx
arhs = a_mass - dt*gammaSU*dot(self.ubar, grad(theta))*dx
if 1 in direction:
arhs -= (
dt*dot(jump(gammaSU), (un('+')*theta('+')
- un('-')*theta('-')))*dS_v
- dt*(gammaSU('+')*dot(self.ubar('+'), n('+'))*theta('+')
+ gammaSU('-')*dot(self.ubar('-'), n('-'))*theta('-'))*dS_v
)
if 2 in direction:
arhs -= (
dt*dot(jump(gammaSU), (un('+')*theta('+')
- un('-')*theta('-')))*dS_h
- dt*(gammaSU('+')*dot(self.ubar('+'), n('+'))*theta('+')
+ gammaSU('-')*dot(self.ubar('-'), n('-'))*theta('-'))*dS_h
)
self.theta1 = Function(V)
self.dtheta = Function(V)
problem = LinearVariationalProblem(a_mass, action(arhs,self.theta1), self.dtheta)
self.solver = LinearVariationalSolver(problem,
options_prefix='SUPGAdvection')
def apply(self, x_in, x_out):
# SSPRK Stage 1
self.theta1.assign(x_in)
self.solver.solve()
self.theta1.assign(self.dtheta)
# SSPRK Stage 2
self.solver.solve()
self.theta1.assign(0.75*x_in + 0.25*self.dtheta)
# SSPRK Stage 3
self.solver.solve()
x_out.assign((1.0/3.0)*x_in + (2.0/3.0)*self.dtheta)
class IncompressibleForcing(Forcing):
"""
Forcing class for incompressible Euler Boussinesq equations.
"""
def __init__(self, state, linear=False):
self.state = state
self._build_forcing_solver(linear)
def _build_forcing_solver(self, linear):
"""
Only put forcing terms into the u equation.
"""
state = self.state
self.scaling = Constant(1.0)
Vu = state.V[0]
W = state.W
self.x0 = Function(W) # copy x to here
u0, p0, b0 = split(self.x0)
F = TrialFunction(Vu)
w = TestFunction(Vu)
self.uF = Function(Vu)
Omega = state.Omega
mu = state.mu
a = inner(w, F) * dx
L = (
self.scaling * div(w) * p0 * dx # pressure gradient
+ self.scaling * b0 * inner(w, state.k) * dx # gravity term
)
if not linear:
L -= self.scaling * 0.5 * div(w) * inner(u0, u0) * dx
if Omega is not None:
L -= self.scaling * inner(w, cross(2 * Omega, u0)) * dx # Coriolis term
if mu is not None:
self.mu_scaling = Constant(1.0)
L -= self.mu_scaling * mu * inner(w, state.k) * inner(u0, state.k) * dx
bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")]
u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs)
self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem)
Vp = state.V[1]
p = TrialFunction(Vp)
q = TestFunction(Vp)
self.divu = Function(Vp)
a = p * q * dx
L = q * div(u0) * dx
divergence_problem = LinearVariationalProblem(a, L, self.divu)
self.divergence_solver = LinearVariationalSolver(divergence_problem)
def apply(self, scaling, x_in, x_nl, x_out, **kwargs):
self.x0.assign(x_nl)
self.scaling.assign(scaling)
if "mu_alpha" in kwargs and kwargs["mu_alpha"] is not None:
self.mu_scaling.assign(kwargs["mu_alpha"])
self.u_forcing_solver.solve() # places forcing in self.uF
u_out, p_out, _ = x_out.split()
x_out.assign(x_in)
u_out += self.uF
if "incompressible" in kwargs and kwargs["incompressible"]:
self.divergence_solver.solve()
p_out.assign(self.divu)
class IncompressibleSolver(TimesteppingSolver):
"""Timestepping linear solver object for the incompressible
Boussinesq equations with prognostic variables u, p, b.
This solver follows the following strategy:
(1) Analytically eliminate b (introduces error near topography)
(2) Solve resulting system for (u,p) using a block Hdiv preconditioner
(3) Reconstruct b
This currently requires a (parallel) direct solver so is probably
a bit memory-hungry, we'll improve this with a hybridised solver
soon.
:arg state: a :class:`.State` object containing everything else.
:arg L: the width of the domain, used in the preconditioner.
:arg params: Solver parameters.
"""
def __init__(self, state, L, params=None):
self.state = state
if params is None:
self.params = {'ksp_type':'gmres',
'pc_type':'fieldsplit',
'pc_fieldsplit_type':'additive',
'fieldsplit_0_pc_type':'lu',
'fieldsplit_1_pc_type':'lu',
'fieldsplit_0_pc_factor_mat_solver_package': 'mumps',
'fieldsplit_0_pc_factor_mat_solver_package': 'mumps',
'fieldsplit_0_ksp_type':'preonly',
'fieldsplit_1_ksp_type':'preonly'}
else:
self.params = params
self.L = L
# setup the solver
self._setup_solver()
def _setup_solver(self):
state = self.state # just cutting down line length a bit
dt = state.timestepping.dt
beta = dt*state.timestepping.alpha
mu = state.mu
# Split up the rhs vector (symbolically)
u_in, p_in, b_in = split(state.xrhs)
# Build the reduced function space for u,p
M = MixedFunctionSpace((state.V[0], state.V[1]))
w, phi = TestFunctions(M)
u, p = TrialFunctions(M)
# Get background fields
bbar = state.bbar
# Analytical (approximate) elimination of theta
k = state.k # Upward pointing unit vector
b = -dot(k,u)*dot(k,grad(bbar))*beta + b_in
# vertical projection
def V(u):
return k*inner(u,k)
eqn = (
inner(w, (u - u_in))*dx
- beta*div(w)*p*dx
- beta*inner(w,k)*b*dx
+ phi*div(u)*dx
)
if mu is not None:
eqn += dt*mu*inner(w,k)*inner(u,k)*dx
aeqn = lhs(eqn)
Leqn = rhs(eqn)
# Place to put result of u p solver
self.up = Function(M)
# Boundary conditions (assumes extruded mesh)
dim = M.sub(0).ufl_element().value_shape()[0]
bc = ("0.0",)*dim
bcs = [DirichletBC(M.sub(0), Expression(bc), "bottom"),
DirichletBC(M.sub(0), Expression(bc), "top")]
# preconditioner equation
L = self.L
Ap = (
inner(w,u) + L*L*div(w)*div(u) +
phi*p/L/L
)*dx
# Solver for u, p
up_problem = LinearVariationalProblem(
aeqn, Leqn, self.up, bcs=bcs, aP=Ap)
nullspace = MixedVectorSpaceBasis(M,
[M.sub(0),
VectorSpaceBasis(constant=True)])
self.up_solver = LinearVariationalSolver(up_problem,
solver_parameters=self.params,
nullspace=nullspace)
# Reconstruction of b
b = TrialFunction(state.V[2])
gamma = TestFunction(state.V[2])
u, p = self.up.split()
self.b = Function(state.V[2])
b_eqn = gamma*(b - b_in +
dot(k,u)*dot(k,grad(bbar))*beta)*dx
b_problem = LinearVariationalProblem(lhs(b_eqn),
rhs(b_eqn),
self.b)
self.b_solver = LinearVariationalSolver(b_problem)
def solve(self):
"""
Apply the solver with rhs state.xrhs and result state.dy.
"""
self.up_solver.solve()
u1, p1 = self.up.split()
u, p, b = self.state.dy.split()
u.assign(u1)
p.assign(p1)
self.b_solver.solve()
b.assign(self.b)
def _setup_solver(self):
import numpy as np
state = self.state
Dt = state.timestepping.dt
beta_ = Dt*state.timestepping.alpha
cp = state.parameters.cp
mu = state.mu
Vu = state.spaces("HDiv")
Vu_broken = FunctionSpace(state.mesh, BrokenElement(Vu.ufl_element()))
Vtheta = state.spaces("HDiv_v")
Vrho = state.spaces("DG")
# Store time-stepping coefficients as UFL Constants
dt = Constant(Dt)
beta = Constant(beta_)
beta_cp = Constant(beta_ * cp)
h_deg = state.horizontal_degree
v_deg = state.vertical_degree
Vtrace = FunctionSpace(state.mesh, "HDiv Trace", degree=(h_deg, v_deg))
# Split up the rhs vector (symbolically)
u_in, rho_in, theta_in = split(state.xrhs)
# Build the function space for "broken" u, rho, and pressure trace
M = MixedFunctionSpace((Vu_broken, Vrho, Vtrace))
w, phi, dl = TestFunctions(M)
u, rho, l0 = TrialFunctions(M)
n = FacetNormal(state.mesh)
# Get background fields
thetabar = state.fields("thetabar")
rhobar = state.fields("rhobar")
pibar = thermodynamics.pi(state.parameters, rhobar, thetabar)
pibar_rho = thermodynamics.pi_rho(state.parameters, rhobar, thetabar)
pibar_theta = thermodynamics.pi_theta(state.parameters, rhobar, thetabar)
# Analytical (approximate) elimination of theta
k = state.k # Upward pointing unit vector
theta = -dot(k, u)*dot(k, grad(thetabar))*beta + theta_in
# Only include theta' (rather than pi') in the vertical
# component of the gradient
# The pi prime term (here, bars are for mean and no bars are
# for linear perturbations)
pi = pibar_theta*theta + pibar_rho*rho
# Vertical projection
def V(u):
return k*inner(u, k)
# Specify degree for some terms as estimated degree is too large
dxp = dx(degree=(self.quadrature_degree))
dS_vp = dS_v(degree=(self.quadrature_degree))
dS_hp = dS_h(degree=(self.quadrature_degree))
ds_vp = ds_v(degree=(self.quadrature_degree))
ds_tbp = (ds_t(degree=(self.quadrature_degree))
+ ds_b(degree=(self.quadrature_degree)))
# Add effect of density of water upon theta
if self.moisture is not None:
water_t = Function(Vtheta).assign(0.0)
for water in self.moisture:
water_t += self.state.fields(water)
theta_w = theta / (1 + water_t)
thetabar_w = thetabar / (1 + water_t)
else:
theta_w = theta
thetabar_w = thetabar
_l0 = TrialFunction(Vtrace)
_dl = TestFunction(Vtrace)
a_tr = _dl('+')*_l0('+')*(dS_vp + dS_hp) + _dl*_l0*ds_vp + _dl*_l0*ds_tbp
def L_tr(f):
return _dl('+')*avg(f)*(dS_vp + dS_hp) + _dl*f*ds_vp + _dl*f*ds_tbp
cg_ilu_parameters = {'ksp_type': 'cg',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}
# Project field averages into functions on the trace space
rhobar_avg = Function(Vtrace)
pibar_avg = Function(Vtrace)
rho_avg_prb = LinearVariationalProblem(a_tr, L_tr(rhobar), rhobar_avg)
pi_avg_prb = LinearVariationalProblem(a_tr, L_tr(pibar), pibar_avg)
rho_avg_solver = LinearVariationalSolver(rho_avg_prb,
solver_parameters=cg_ilu_parameters,
options_prefix='rhobar_avg_solver')
pi_avg_solver = LinearVariationalSolver(pi_avg_prb,
solver_parameters=cg_ilu_parameters,
options_prefix='pibar_avg_solver')
with timed_region("Gusto:HybridProjectRhobar"):
rho_avg_solver.solve()
with timed_region("Gusto:HybridProjectPibar"):
pi_avg_solver.solve()
# "broken" u, rho, and trace system
# NOTE: no ds_v integrals since equations are defined on
# a periodic (or sphere) base mesh.
eqn = (
# momentum equation
inner(w, (state.h_project(u) - u_in))*dx
- beta_cp*div(theta_w*V(w))*pibar*dxp
# following does nothing but is preserved in the comments
# to remind us why (because V(w) is purely vertical).
# + beta_cp*jump(theta_w*V(w), n=n)*pibar_avg('+')*dS_vp
+ beta_cp*jump(theta_w*V(w), n=n)*pibar_avg('+')*dS_hp
+ beta_cp*dot(theta_w*V(w), n)*pibar_avg*ds_tbp
- beta_cp*div(thetabar_w*w)*pi*dxp
# trace terms appearing after integrating momentum equation
+ beta_cp*jump(thetabar_w*w, n=n)*l0('+')*(dS_vp + dS_hp)
+ beta_cp*dot(thetabar_w*w, n)*l0*(ds_tbp + ds_vp)
# mass continuity equation
+ (phi*(rho - rho_in) - beta*inner(grad(phi), u)*rhobar)*dx
+ beta*jump(phi*u, n=n)*rhobar_avg('+')*(dS_v + dS_h)
# term added because u.n=0 is enforced weakly via the traces
+ beta*phi*dot(u, n)*rhobar_avg*(ds_tb + ds_v)
# constraint equation to enforce continuity of the velocity
# through the interior facets and weakly impose the no-slip
# condition
+ dl('+')*jump(u, n=n)*(dS_vp + dS_hp)
+ dl*dot(u, n)*(ds_tbp + ds_vp)
)
# contribution of the sponge term
if mu is not None:
eqn += dt*mu*inner(w, k)*inner(u, k)*dx
aeqn = lhs(eqn)
Leqn = rhs(eqn)
# Function for the hybridized solutions
self.urhol0 = Function(M)
hybridized_prb = LinearVariationalProblem(aeqn, Leqn, self.urhol0)
hybridized_solver = LinearVariationalSolver(hybridized_prb,
solver_parameters=self.solver_parameters,
options_prefix='ImplicitSolver')
self.hybridized_solver = hybridized_solver
# Project broken u into the HDiv space using facet averaging.
# Weight function counting the dofs of the HDiv element:
shapes = {"i": Vu.finat_element.space_dimension(),
"j": np.prod(Vu.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(Vu)
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)
# HDiv-conforming velocity
self.u_hdiv = Function(Vu)
# Reconstruction of theta
theta = TrialFunction(Vtheta)
gamma = TestFunction(Vtheta)
self.theta = Function(Vtheta)
theta_eqn = gamma*(theta - theta_in
+ dot(k, self.u_hdiv)*dot(k, grad(thetabar))*beta)*dx
theta_problem = LinearVariationalProblem(lhs(theta_eqn), rhs(theta_eqn), self.theta)
self.theta_solver = LinearVariationalSolver(theta_problem,
solver_parameters=cg_ilu_parameters,
options_prefix='thetabacksubstitution')
# Store boundary conditions for the div-conforming velocity to apply
# post-solve
self.bcs = self.state.bcs
class CompressibleSolver(TimesteppingSolver):
"""
Timestepping linear solver object for the compressible equations
in theta-pi formulation with prognostic variables u,rho,theta.
This solver follows the following strategy:
(1) Analytically eliminate theta (introduces error near topography)
(2) Solve resulting system for (u,rho) using a Schur preconditioner
(3) Reconstruct theta
:arg state: a :class:`.State` object containing everything else.
"""
def __init__(self, state, params=None):
self.state = state
if params is None:
self.params = {'pc_type': 'fieldsplit',
'pc_fieldsplit_type': 'schur',
'ksp_type': 'gmres',
'ksp_max_it': 100,
'ksp_gmres_restart': 50,
'pc_fieldsplit_schur_fact_type': 'FULL',
'pc_fieldsplit_schur_precondition': 'selfp',
'fieldsplit_0_ksp_type': 'preonly',
'fieldsplit_0_pc_type': 'bjacobi',
'fieldsplit_0_sub_pc_type': 'ilu',
'fieldsplit_1_ksp_type': 'preonly',
'fieldsplit_1_pc_type': 'gamg',
'fieldsplit_1_mg_levels_ksp_type': 'chebyshev',
'fieldsplit_1_mg_levels_ksp_chebyshev_estimate_eigenvalues': True,
'fieldsplit_1_mg_levels_ksp_chebyshev_estimate_eigenvalues_random': True,
'fieldsplit_1_mg_levels_ksp_max_it': 1,
'fieldsplit_1_mg_levels_pc_type': 'bjacobi',
'fieldsplit_1_mg_levels_sub_pc_type': 'ilu'}
else:
self.params = params
# setup the solver
self._setup_solver()
def _setup_solver(self):
state = self.state # just cutting down line length a bit
dt = state.timestepping.dt
beta = dt*state.timestepping.alpha
cp = state.parameters.cp
mu = state.mu
# Split up the rhs vector (symbolically)
u_in, rho_in, theta_in = split(state.xrhs)
# Build the reduced function space for u,rho
M = MixedFunctionSpace((state.V[0], state.V[1]))
w, phi = TestFunctions(M)
u, rho = TrialFunctions(M)
n = FacetNormal(state.mesh)
# Get background fields
thetabar = state.thetabar
rhobar = state.rhobar
pibar = exner(thetabar, rhobar, state)
pibar_rho = exner_rho(thetabar, rhobar, state)
pibar_theta = exner_theta(thetabar, rhobar, state)
# Analytical (approximate) elimination of theta
k = state.k # Upward pointing unit vector
theta = -dot(k,u)*dot(k,grad(thetabar))*beta + theta_in
# Only include theta' (rather than pi') in the vertical
# component of the gradient
# the pi prime term (here, bars are for mean and no bars are
# for linear perturbations)
pi = pibar_theta*theta + pibar_rho*rho
# vertical projection
def V(u):
return k*inner(u,k)
eqn = (
inner(w, (u - u_in))*dx
- beta*cp*div(theta*V(w))*pibar*dx
# following does nothing but is preserved in the comments
# to remind us why (because V(w) is purely vertical.
# + beta*cp*jump(theta*V(w),n)*avg(pibar)*dS_v
- beta*cp*div(thetabar*w)*pi*dx
+ beta*cp*jump(thetabar*w,n)*avg(pi)*dS_v
+ (phi*(rho - rho_in) - beta*inner(grad(phi), u)*rhobar)*dx
+ beta*jump(phi*u, n)*avg(rhobar)*(dS_v + dS_h)
)
if mu is not None:
eqn += dt*mu*inner(w,k)*inner(u,k)*dx
aeqn = lhs(eqn)
Leqn = rhs(eqn)
# Place to put result of u rho solver
self.urho = Function(M)
# Boundary conditions (assumes extruded mesh)
dim = M.sub(0).ufl_element().value_shape()[0]
bc = ("0.0",)*dim
bcs = [DirichletBC(M.sub(0), Expression(bc), "bottom"),
DirichletBC(M.sub(0), Expression(bc), "top")]
# Solver for u, rho
urho_problem = LinearVariationalProblem(
aeqn, Leqn, self.urho, bcs=bcs)
self.urho_solver = LinearVariationalSolver(urho_problem,
solver_parameters=self.params,
options_prefix='ImplicitSolver')
# Reconstruction of theta
theta = TrialFunction(state.V[2])
gamma = TestFunction(state.V[2])
u, rho = self.urho.split()
self.theta = Function(state.V[2])
theta_eqn = gamma*(theta - theta_in +
dot(k,u)*dot(k,grad(thetabar))*beta)*dx
theta_problem = LinearVariationalProblem(lhs(theta_eqn),
rhs(theta_eqn),
self.theta)
self.theta_solver = LinearVariationalSolver(theta_problem,
options_prefix='thetabacksubstitution')
def solve(self):
"""
Apply the solver with rhs state.xrhs and result state.dy.
"""
self.urho_solver.solve()
u1, rho1 = self.urho.split()
u, rho, theta = self.state.dy.split()
u.assign(u1)
rho.assign(rho1)
self.theta_solver.solve()
theta.assign(self.theta)
class CompressibleForcing(Forcing):
"""
Forcing class for compressible Euler equations.
"""
def __init__(self, state, linear=False):
self.state = state
self._build_forcing_solver(linear)
def _build_forcing_solver(self, linear):
"""
Only put forcing terms into the u equation.
"""
state = self.state
self.scaling = Constant(1.0)
Vu = state.V[0]
W = state.W
self.x0 = Function(W) # copy x to here
u0, rho0, theta0 = split(self.x0)
F = TrialFunction(Vu)
w = TestFunction(Vu)
self.uF = Function(Vu)
Omega = state.Omega
cp = state.parameters.cp
mu = state.mu
n = FacetNormal(state.mesh)
pi = exner(theta0, rho0, state)
a = inner(w, F) * dx
L = self.scaling * (
+cp * div(theta0 * w) * pi * dx # pressure gradient [volume]
- cp * jump(w * theta0, n) * avg(pi) * dS_v # pressure gradient [surface]
)
if state.parameters.geopotential:
Phi = state.Phi
L += self.scaling * div(w) * Phi * dx # gravity term
else:
g = state.parameters.g
L -= self.scaling * g * inner(w, state.k) * dx # gravity term
if not linear:
L -= self.scaling * 0.5 * div(w) * inner(u0, u0) * dx
if Omega is not None:
L -= self.scaling * inner(w, cross(2 * Omega, u0)) * dx # Coriolis term
if mu is not None:
self.mu_scaling = Constant(1.0)
L -= self.mu_scaling * mu * inner(w, state.k) * inner(u0, state.k) * dx
bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")]
u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs)
self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem)
def apply(self, scaling, x_in, x_nl, x_out, **kwargs):
self.x0.assign(x_nl)
self.scaling.assign(scaling)
if "mu_alpha" in kwargs and kwargs["mu_alpha"] is not None:
self.mu_scaling.assign(kwargs["mu_alpha"])
self.u_forcing_solver.solve() # places forcing in self.uF
u_out, _, _ = x_out.split()
x_out.assign(x_in)
u_out += self.uF
class IncompressibleSolver(TimesteppingSolver):
"""Timestepping linear solver object for the incompressible
Boussinesq equations with prognostic variables u, p, b.
This solver follows the following strategy:
(1) Analytically eliminate b (introduces error near topography)
(2) Solve resulting system for (u,p) using a hybrid-mixed method
(3) Reconstruct b
:arg state: a :class:`.State` object containing everything else.
:arg solver_parameters: (optional) Solver parameters.
:arg overwrite_solver_parameters: boolean, if True use only the
solver_parameters that have been passed in, if False then
update the default solver parameters with the solver_parameters
passed in.
"""
solver_parameters = {
'ksp_type': 'preonly',
'mat_type': 'matfree',
'pc_type': 'python',
'pc_python_type': 'firedrake.HybridizationPC',
'hybridization': {'ksp_type': 'cg',
'pc_type': 'gamg',
'ksp_rtol': 1e-8,
'mg_levels': {'ksp_type': 'chebyshev',
'ksp_max_it': 2,
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}}
}
def __init__(self, state, solver_parameters=None,
overwrite_solver_parameters=False):
super().__init__(state, solver_parameters, overwrite_solver_parameters)
@timed_function("Gusto:SolverSetup")
def _setup_solver(self):
state = self.state # just cutting down line length a bit
Dt = state.timestepping.dt
beta_ = Dt*state.timestepping.alpha
mu = state.mu
Vu = state.spaces("HDiv")
Vb = state.spaces("HDiv_v")
Vp = state.spaces("DG")
# Store time-stepping coefficients as UFL Constants
dt = Constant(Dt)
beta = Constant(beta_)
# Split up the rhs vector (symbolically)
u_in, p_in, b_in = split(state.xrhs)
# Build the reduced function space for u,p
M = MixedFunctionSpace((Vu, Vp))
w, phi = TestFunctions(M)
u, p = TrialFunctions(M)
# Get background fields
bbar = state.fields("bbar")
# Analytical (approximate) elimination of theta
k = state.k # Upward pointing unit vector
b = -dot(k, u)*dot(k, grad(bbar))*beta + b_in
# vertical projection
def V(u):
return k*inner(u, k)
eqn = (
inner(w, (u - u_in))*dx
- beta*div(w)*p*dx
- beta*inner(w, k)*b*dx
+ phi*div(u)*dx
)
if mu is not None:
eqn += dt*mu*inner(w, k)*inner(u, k)*dx
aeqn = lhs(eqn)
Leqn = rhs(eqn)
# Place to put result of u p solver
self.up = Function(M)
# Boundary conditions (assumes extruded mesh)
bcs = None if len(self.state.bcs) == 0 else self.state.bcs
# Solver for u, p
up_problem = LinearVariationalProblem(aeqn, Leqn, self.up, bcs=bcs)
# Provide callback for the nullspace of the trace system
def trace_nullsp(T):
return VectorSpaceBasis(constant=True)
appctx = {"trace_nullspace": trace_nullsp}
self.up_solver = LinearVariationalSolver(up_problem,
solver_parameters=self.solver_parameters,
appctx=appctx)
# Reconstruction of b
b = TrialFunction(Vb)
gamma = TestFunction(Vb)
u, p = self.up.split()
self.b = Function(Vb)
b_eqn = gamma*(b - b_in
+ dot(k, u)*dot(k, grad(bbar))*beta)*dx
b_problem = LinearVariationalProblem(lhs(b_eqn),
rhs(b_eqn),
self.b)
self.b_solver = LinearVariationalSolver(b_problem)
@timed_function("Gusto:LinearSolve")
def solve(self):
"""
Apply the solver with rhs state.xrhs and result state.dy.
"""
with timed_region("Gusto:VelocityPressureSolve"):
self.up_solver.solve()
u1, p1 = self.up.split()
u, p, b = self.state.dy.split()
u.assign(u1)
p.assign(p1)
with timed_region("Gusto:BuoyancyRecon"):
self.b_solver.solve()
b.assign(self.b)
def compressible_hydrostatic_balance(state, theta0, rho0, pi0=None,
top=False, pi_boundary=Constant(1.0),
solve_for_rho=False,
params=None):
"""
Compute a hydrostatically balanced density given a potential temperature
profile.
:arg state: The :class:`State` object.
:arg theta0: :class:`.Function`containing the potential temperature.
:arg rho0: :class:`.Function` to write the initial density into.
:arg top: If True, set a boundary condition at the top. Otherwise, set
it at the bottom.
:arg pi_boundary: a field or expression to use as boundary data for pi on
the top or bottom as specified.
"""
# Calculate hydrostatic Pi
W = MixedFunctionSpace((state.Vv,state.V[1]))
v, pi = TrialFunctions(W)
dv, dpi = TestFunctions(W)
n = FacetNormal(state.mesh)
cp = state.parameters.cp
alhs = (
(cp*inner(v,dv) - cp*div(dv*theta0)*pi)*dx
+ dpi*div(theta0*v)*dx
)
if top:
bmeasure = ds_t
bstring = "bottom"
else:
bmeasure = ds_b
bstring = "top"
arhs = -cp*inner(dv,n)*theta0*pi_boundary*bmeasure
if state.parameters.geopotential:
Phi = state.Phi
arhs += div(dv)*Phi*dx - inner(dv,n)*Phi*bmeasure
else:
g = state.parameters.g
arhs -= g*inner(dv,state.k)*dx
if(state.mesh.geometric_dimension() == 2):
bcs = [DirichletBC(W.sub(0), Expression(("0.", "0.")), bstring)]
elif(state.mesh.geometric_dimension() == 3):
bcs = [DirichletBC(W.sub(0), Expression(("0.", "0.", "0.")), bstring)]
w = Function(W)
PiProblem = LinearVariationalProblem(alhs, arhs, w, bcs=bcs)
if(params is None):
params = {'pc_type': 'fieldsplit',
'pc_fieldsplit_type': 'schur',
'ksp_type': 'gmres',
'ksp_monitor_true_residual': True,
'ksp_max_it': 100,
'ksp_gmres_restart': 50,
'pc_fieldsplit_schur_fact_type': 'FULL',
'pc_fieldsplit_schur_precondition': 'selfp',
'fieldsplit_0_ksp_type': 'richardson',
'fieldsplit_0_ksp_max_it': 5,
'fieldsplit_0_pc_type': 'gamg',
'fieldsplit_1_pc_gamg_sym_graph': True,
'fieldsplit_1_mg_levels_ksp_type': 'chebyshev',
'fieldsplit_1_mg_levels_ksp_chebyshev_estimate_eigenvalues': True,
'fieldsplit_1_mg_levels_ksp_chebyshev_estimate_eigenvalues_random': True,
'fieldsplit_1_mg_levels_ksp_max_it': 5,
'fieldsplit_1_mg_levels_pc_type': 'bjacobi',
'fieldsplit_1_mg_levels_sub_pc_type': 'ilu'}
PiSolver = LinearVariationalSolver(PiProblem,
solver_parameters=params)
PiSolver.solve()
v, Pi = w.split()
if pi0 is not None:
pi0.assign(Pi)
kappa = state.parameters.kappa
R_d = state.parameters.R_d
p_0 = state.parameters.p_0
if solve_for_rho:
w1 = Function(W)
v, rho = w1.split()
rho.interpolate(p_0*(Pi**((1-kappa)/kappa))/R_d/theta0)
v, rho = split(w1)
dv, dpi = TestFunctions(W)
pi = ((R_d/p_0)*rho*theta0)**(kappa/(1.-kappa))
F = (
(cp*inner(v,dv) - cp*div(dv*theta0)*pi)*dx
+ dpi*div(theta0*v)*dx
+ cp*inner(dv,n)*theta0*pi_boundary*bmeasure
)
if state.parameters.geopotential:
F += - div(dv)*Phi*dx + inner(dv,n)*Phi*bmeasure
else:
F += g*inner(dv,state.k)*dx
rhoproblem = NonlinearVariationalProblem(F, w1, bcs=bcs)
rhosolver = NonlinearVariationalSolver(rhoproblem, solver_parameters=params)
rhosolver.solve()
v, rho_ = w1.split()
rho0.assign(rho_)
else:
rho0.interpolate(p_0*(Pi**((1-kappa)/kappa))/R_d/theta0)
def _setup_solver(self):
state = self.state # just cutting down line length a bit
Dt = state.timestepping.dt
beta_ = Dt*state.timestepping.alpha
mu = state.mu
Vu = state.spaces("HDiv")
Vb = state.spaces("HDiv_v")
Vp = state.spaces("DG")
# Store time-stepping coefficients as UFL Constants
dt = Constant(Dt)
beta = Constant(beta_)
# Split up the rhs vector (symbolically)
u_in, p_in, b_in = split(state.xrhs)
# Build the reduced function space for u,p
M = MixedFunctionSpace((Vu, Vp))
w, phi = TestFunctions(M)
u, p = TrialFunctions(M)
# Get background fields
bbar = state.fields("bbar")
# Analytical (approximate) elimination of theta
k = state.k # Upward pointing unit vector
b = -dot(k, u)*dot(k, grad(bbar))*beta + b_in
# vertical projection
def V(u):
return k*inner(u, k)
eqn = (
inner(w, (u - u_in))*dx
- beta*div(w)*p*dx
- beta*inner(w, k)*b*dx
+ phi*div(u)*dx
)
if mu is not None:
eqn += dt*mu*inner(w, k)*inner(u, k)*dx
aeqn = lhs(eqn)
Leqn = rhs(eqn)
# Place to put result of u p solver
self.up = Function(M)
# Boundary conditions (assumes extruded mesh)
bcs = None if len(self.state.bcs) == 0 else self.state.bcs
# Solver for u, p
up_problem = LinearVariationalProblem(aeqn, Leqn, self.up, bcs=bcs)
# Provide callback for the nullspace of the trace system
def trace_nullsp(T):
return VectorSpaceBasis(constant=True)
appctx = {"trace_nullspace": trace_nullsp}
self.up_solver = LinearVariationalSolver(up_problem,
solver_parameters=self.solver_parameters,
appctx=appctx)
# Reconstruction of b
b = TrialFunction(Vb)
gamma = TestFunction(Vb)
u, p = self.up.split()
self.b = Function(Vb)
b_eqn = gamma*(b - b_in
+ dot(k, u)*dot(k, grad(bbar))*beta)*dx
b_problem = LinearVariationalProblem(lhs(b_eqn),
rhs(b_eqn),
self.b)
self.b_solver = LinearVariationalSolver(b_problem)
class CompressibleSolver(TimesteppingSolver):
"""
Timestepping linear solver object for the compressible equations
in theta-pi formulation with prognostic variables u, rho, and theta.
This solver follows the following strategy:
(1) Analytically eliminate theta (introduces error near topography)
(2a) Formulate the resulting mixed system for u and rho using a
hybridized mixed method. This breaks continuity in the
linear perturbations of u, and introduces a new unknown on the
mesh interfaces approximating the average of the Exner pressure
perturbations. These trace unknowns also act as Lagrange
multipliers enforcing normal continuity of the "broken" u variable.
(2b) Statically condense the block-sparse system into a single system
for the Lagrange multipliers. This is the only globally coupled
system requiring a linear solver.
(2c) Using the computed trace variables, we locally recover the
broken velocity and density perturbations. This is accomplished
in two stages:
(i): Recover rho locally using the multipliers.
(ii): Recover "broken" u locally using rho and the multipliers.
(2d) Project the "broken" velocity field into the HDiv-conforming
space using local averaging.
(3) Reconstruct theta
:arg state: a :class:`.State` object containing everything else.
:arg quadrature degree: tuple (q_h, q_v) where q_h is the required
quadrature degree in the horizontal direction and q_v is that in
the vertical direction.
:arg solver_parameters (optional): solver parameters for the
trace system.
:arg overwrite_solver_parameters: boolean, if True use only the
solver_parameters that have been passed in, if False then update.
the default solver parameters with the solver_parameters passed in.
:arg moisture (optional): list of names of moisture fields.
"""
solver_parameters = {'mat_type': 'matfree',
'ksp_type': 'preonly',
'pc_type': 'python',
'pc_python_type': 'firedrake.SCPC',
'pc_sc_eliminate_fields': '0, 1',
# The reduced operator is not symmetric
'condensed_field': {'ksp_type': 'fgmres',
'ksp_rtol': 1.0e-8,
'ksp_atol': 1.0e-8,
'ksp_max_it': 100,
'pc_type': 'gamg',
'pc_gamg_sym_graph': None,
'mg_levels': {'ksp_type': 'gmres',
'ksp_max_it': 5,
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}}}
def __init__(self, state, quadrature_degree=None, solver_parameters=None,
overwrite_solver_parameters=False, moisture=None):
self.moisture = moisture
self.state = state
if quadrature_degree is not None:
self.quadrature_degree = quadrature_degree
else:
dgspace = state.spaces("DG")
if any(deg > 2 for deg in dgspace.ufl_element().degree()):
logger.warning("default quadrature degree most likely not sufficient for this degree element")
self.quadrature_degree = (5, 5)
if logger.isEnabledFor(DEBUG):
# Set outer solver to FGMRES and turn on KSP monitor for the outer system
self.solver_parameters["ksp_type"] = "fgmres"
self.solver_parameters["mat_type"] = "aij"
self.solver_parameters["pmat_type"] = "matfree"
self.solver_parameters["ksp_monitor_true_residual"] = None
# Turn monitor on for the trace system
self.solver_parameters["condensed_field"]["ksp_monitor_true_residual"] = None
super().__init__(state, solver_parameters, overwrite_solver_parameters)
@timed_function("Gusto:SolverSetup")
def _setup_solver(self):
import numpy as np
state = self.state
Dt = state.timestepping.dt
beta_ = Dt*state.timestepping.alpha
cp = state.parameters.cp
mu = state.mu
Vu = state.spaces("HDiv")
Vu_broken = FunctionSpace(state.mesh, BrokenElement(Vu.ufl_element()))
Vtheta = state.spaces("HDiv_v")
Vrho = state.spaces("DG")
# Store time-stepping coefficients as UFL Constants
dt = Constant(Dt)
beta = Constant(beta_)
beta_cp = Constant(beta_ * cp)
h_deg = state.horizontal_degree
v_deg = state.vertical_degree
Vtrace = FunctionSpace(state.mesh, "HDiv Trace", degree=(h_deg, v_deg))
# Split up the rhs vector (symbolically)
u_in, rho_in, theta_in = split(state.xrhs)
# Build the function space for "broken" u, rho, and pressure trace
M = MixedFunctionSpace((Vu_broken, Vrho, Vtrace))
w, phi, dl = TestFunctions(M)
u, rho, l0 = TrialFunctions(M)
n = FacetNormal(state.mesh)
# Get background fields
thetabar = state.fields("thetabar")
rhobar = state.fields("rhobar")
pibar = thermodynamics.pi(state.parameters, rhobar, thetabar)
pibar_rho = thermodynamics.pi_rho(state.parameters, rhobar, thetabar)
pibar_theta = thermodynamics.pi_theta(state.parameters, rhobar, thetabar)
# Analytical (approximate) elimination of theta
k = state.k # Upward pointing unit vector
theta = -dot(k, u)*dot(k, grad(thetabar))*beta + theta_in
# Only include theta' (rather than pi') in the vertical
# component of the gradient
# The pi prime term (here, bars are for mean and no bars are
# for linear perturbations)
pi = pibar_theta*theta + pibar_rho*rho
# Vertical projection
def V(u):
return k*inner(u, k)
# Specify degree for some terms as estimated degree is too large
dxp = dx(degree=(self.quadrature_degree))
dS_vp = dS_v(degree=(self.quadrature_degree))
dS_hp = dS_h(degree=(self.quadrature_degree))
ds_vp = ds_v(degree=(self.quadrature_degree))
ds_tbp = (ds_t(degree=(self.quadrature_degree))
+ ds_b(degree=(self.quadrature_degree)))
# Add effect of density of water upon theta
if self.moisture is not None:
water_t = Function(Vtheta).assign(0.0)
for water in self.moisture:
water_t += self.state.fields(water)
theta_w = theta / (1 + water_t)
thetabar_w = thetabar / (1 + water_t)
else:
theta_w = theta
thetabar_w = thetabar
_l0 = TrialFunction(Vtrace)
_dl = TestFunction(Vtrace)
a_tr = _dl('+')*_l0('+')*(dS_vp + dS_hp) + _dl*_l0*ds_vp + _dl*_l0*ds_tbp
def L_tr(f):
return _dl('+')*avg(f)*(dS_vp + dS_hp) + _dl*f*ds_vp + _dl*f*ds_tbp
cg_ilu_parameters = {'ksp_type': 'cg',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}
# Project field averages into functions on the trace space
rhobar_avg = Function(Vtrace)
pibar_avg = Function(Vtrace)
rho_avg_prb = LinearVariationalProblem(a_tr, L_tr(rhobar), rhobar_avg)
pi_avg_prb = LinearVariationalProblem(a_tr, L_tr(pibar), pibar_avg)
rho_avg_solver = LinearVariationalSolver(rho_avg_prb,
solver_parameters=cg_ilu_parameters,
options_prefix='rhobar_avg_solver')
pi_avg_solver = LinearVariationalSolver(pi_avg_prb,
solver_parameters=cg_ilu_parameters,
options_prefix='pibar_avg_solver')
with timed_region("Gusto:HybridProjectRhobar"):
rho_avg_solver.solve()
with timed_region("Gusto:HybridProjectPibar"):
pi_avg_solver.solve()
# "broken" u, rho, and trace system
# NOTE: no ds_v integrals since equations are defined on
# a periodic (or sphere) base mesh.
eqn = (
# momentum equation
inner(w, (state.h_project(u) - u_in))*dx
- beta_cp*div(theta_w*V(w))*pibar*dxp
# following does nothing but is preserved in the comments
# to remind us why (because V(w) is purely vertical).
# + beta_cp*jump(theta_w*V(w), n=n)*pibar_avg('+')*dS_vp
+ beta_cp*jump(theta_w*V(w), n=n)*pibar_avg('+')*dS_hp
+ beta_cp*dot(theta_w*V(w), n)*pibar_avg*ds_tbp
- beta_cp*div(thetabar_w*w)*pi*dxp
# trace terms appearing after integrating momentum equation
+ beta_cp*jump(thetabar_w*w, n=n)*l0('+')*(dS_vp + dS_hp)
+ beta_cp*dot(thetabar_w*w, n)*l0*(ds_tbp + ds_vp)
# mass continuity equation
+ (phi*(rho - rho_in) - beta*inner(grad(phi), u)*rhobar)*dx
+ beta*jump(phi*u, n=n)*rhobar_avg('+')*(dS_v + dS_h)
# term added because u.n=0 is enforced weakly via the traces
+ beta*phi*dot(u, n)*rhobar_avg*(ds_tb + ds_v)
# constraint equation to enforce continuity of the velocity
# through the interior facets and weakly impose the no-slip
# condition
+ dl('+')*jump(u, n=n)*(dS_vp + dS_hp)
+ dl*dot(u, n)*(ds_tbp + ds_vp)
)
# contribution of the sponge term
if mu is not None:
eqn += dt*mu*inner(w, k)*inner(u, k)*dx
aeqn = lhs(eqn)
Leqn = rhs(eqn)
# Function for the hybridized solutions
self.urhol0 = Function(M)
hybridized_prb = LinearVariationalProblem(aeqn, Leqn, self.urhol0)
hybridized_solver = LinearVariationalSolver(hybridized_prb,
solver_parameters=self.solver_parameters,
options_prefix='ImplicitSolver')
self.hybridized_solver = hybridized_solver
# Project broken u into the HDiv space using facet averaging.
# Weight function counting the dofs of the HDiv element:
shapes = {"i": Vu.finat_element.space_dimension(),
"j": np.prod(Vu.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(Vu)
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)
# HDiv-conforming velocity
self.u_hdiv = Function(Vu)
# Reconstruction of theta
theta = TrialFunction(Vtheta)
gamma = TestFunction(Vtheta)
self.theta = Function(Vtheta)
theta_eqn = gamma*(theta - theta_in
+ dot(k, self.u_hdiv)*dot(k, grad(thetabar))*beta)*dx
theta_problem = LinearVariationalProblem(lhs(theta_eqn), rhs(theta_eqn), self.theta)
self.theta_solver = LinearVariationalSolver(theta_problem,
solver_parameters=cg_ilu_parameters,
options_prefix='thetabacksubstitution')
# Store boundary conditions for the div-conforming velocity to apply
# post-solve
self.bcs = self.state.bcs
@timed_function("Gusto:LinearSolve")
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 __init__(self, state):
super(Condensation, self).__init__(state)
# obtain our fields
self.theta = state.fields('theta')
self.water_v = state.fields('water_v')
self.water_c = state.fields('water_c')
rho = state.fields('rho')
# declare function space
Vt = self.theta.function_space()
param = self.state.parameters
# define some parameters as attributes
dt = self.state.timestepping.dt
R_d = param.R_d
p_0 = param.p_0
kappa = param.kappa
cp = param.cp
cv = param.cv
c_pv = param.c_pv
c_pl = param.c_pl
c_vv = param.c_vv
R_v = param.R_v
L_v0 = param.L_v0
T_0 = param.T_0
w_sat1 = param.w_sat1
w_sat2 = param.w_sat2
w_sat3 = param.w_sat3
w_sat4 = param.w_sat4
# make useful fields
Pi = ((R_d * rho * self.theta / p_0)**(kappa / (1.0 - kappa)))
T = Pi * self.theta * R_d / (R_d + self.water_v * R_v)
p = p_0 * Pi**(1.0 / kappa)
L_v = L_v0 - (c_pl - c_pv) * (T - T_0)
R_m = R_d + R_v * self.water_v
c_pml = cp + c_pv * self.water_v + c_pl * self.water_c
c_vml = cv + c_vv * self.water_v + c_pl * self.water_c
# use Teten's formula to calculate w_sat
w_sat = (w_sat1 / (p * exp(w_sat2 * (T - T_0) /
(T - w_sat3)) - w_sat4))
# make appropriate condensation rate
dot_r_cond = ((self.water_v - w_sat) / (dt * (1.0 +
((L_v**2.0 * w_sat) /
(cp * R_v * T**2.0)))))
# make cond_rate function, that needs to be the same for all updates in one time step
self.cond_rate = Function(Vt)
# adjust cond rate so negative concentrations don't occur
self.lim_cond_rate = Interpolator(
conditional(dot_r_cond < 0,
max_value(dot_r_cond, -self.water_c / dt),
min_value(dot_r_cond, self.water_v / dt)),
self.cond_rate)
# tell the prognostic fields what to update to
self.water_v_new = Interpolator(self.water_v - dt * self.cond_rate, Vt)
self.water_c_new = Interpolator(self.water_c + dt * self.cond_rate, Vt)
self.theta_new = Interpolator(
self.theta * (1.0 + dt * self.cond_rate *
(cv * L_v / (c_vml * cp * T) - R_v * cv * c_pml /
(R_m * cp * c_vml))), Vt)
def _setup_solver(self):
state = self.state # just cutting down line length a bit
dt = state.timestepping.dt
beta = dt*state.timestepping.alpha
mu = state.mu
# Split up the rhs vector (symbolically)
u_in, p_in, b_in = split(state.xrhs)
# Build the reduced function space for u,p
M = MixedFunctionSpace((state.V[0], state.V[1]))
w, phi = TestFunctions(M)
u, p = TrialFunctions(M)
# Get background fields
bbar = state.bbar
# Analytical (approximate) elimination of theta
k = state.k # Upward pointing unit vector
b = -dot(k,u)*dot(k,grad(bbar))*beta + b_in
# vertical projection
def V(u):
return k*inner(u,k)
eqn = (
inner(w, (u - u_in))*dx
- beta*div(w)*p*dx
- beta*inner(w,k)*b*dx
+ phi*div(u)*dx
)
if mu is not None:
eqn += dt*mu*inner(w,k)*inner(u,k)*dx
aeqn = lhs(eqn)
Leqn = rhs(eqn)
# Place to put result of u p solver
self.up = Function(M)
# Boundary conditions (assumes extruded mesh)
dim = M.sub(0).ufl_element().value_shape()[0]
bc = ("0.0",)*dim
bcs = [DirichletBC(M.sub(0), Expression(bc), "bottom"),
DirichletBC(M.sub(0), Expression(bc), "top")]
# preconditioner equation
L = self.L
Ap = (
inner(w,u) + L*L*div(w)*div(u) +
phi*p/L/L
)*dx
# Solver for u, p
up_problem = LinearVariationalProblem(
aeqn, Leqn, self.up, bcs=bcs, aP=Ap)
nullspace = MixedVectorSpaceBasis(M,
[M.sub(0),
VectorSpaceBasis(constant=True)])
self.up_solver = LinearVariationalSolver(up_problem,
solver_parameters=self.params,
nullspace=nullspace)
# Reconstruction of b
b = TrialFunction(state.V[2])
gamma = TestFunction(state.V[2])
u, p = self.up.split()
self.b = Function(state.V[2])
b_eqn = gamma*(b - b_in +
dot(k,u)*dot(k,grad(bbar))*beta)*dx
b_problem = LinearVariationalProblem(lhs(b_eqn),
rhs(b_eqn),
self.b)
self.b_solver = LinearVariationalSolver(b_problem)
def PartiallyPeriodicRectangleMesh(nx, ny, Lx, Ly, direction="x", quadrilateral=False, reorder=None, comm=COMM_WORLD):
"""Generates RectangleMesh that is periodic in the x or y direction.
:arg nx: The number of cells in the x direction
:arg ny: The number of cells in the y direction
:arg Lx: The extent in the x direction
:arg Ly: The extent in the y direction
:kwarg direction: The direction of the periodicity (default x).
:kwarg quadrilateral: (optional), creates quadrilateral mesh, defaults to False
:kwarg reorder: (optional), should the mesh be reordered
:kwarg comm: Optional communicator to build the mesh on (defaults to
COMM_WORLD).
If direction == "x" the boundary edges in this mesh are numbered as follows:
* 1: plane y == 0
* 2: plane y == Ly
If direction == "y" the boundary edges are:
* 1: plane x == 0
* 2: plane x == Lx
"""
if direction not in ("x", "y"):
raise ValueError("Unsupported periodic direction '%s'" % direction)
# handle x/y directions: na, La are for the periodic axis
na, nb, La, Lb = nx, ny, Lx, Ly
if direction == "y":
na, nb, La, Lb = ny, nx, Ly, Lx
if na < 3:
raise ValueError("2D periodic meshes with fewer than 3 \
cells in each direction are not currently supported")
m = CylinderMesh(na, nb, 1.0, 1.0, longitudinal_direction="z",
quadrilateral=quadrilateral, reorder=reorder, comm=comm)
coord_fs = VectorFunctionSpace(m, 'DG', 1, dim=2)
old_coordinates = m.coordinates
new_coordinates = Function(coord_fs)
# make x-periodic mesh
# unravel x coordinates like in periodic interval
# set y coordinates to z coordinates
periodic_kernel = """double Y,pi;
Y = 0.0;
for(int i=0; i<old_coords.dofs; i++) {
Y += old_coords[i][1];
}
pi=3.141592653589793;
for(int i=0;i<new_coords.dofs;i++){
new_coords[i][0] = atan2(old_coords[i][1],old_coords[i][0])/pi/2;
if(new_coords[i][0]<0.) new_coords[i][0] += 1;
if(new_coords[i][0]==0 && Y<0.) new_coords[i][0] = 1.0;
new_coords[i][0] *= Lx[0];
new_coords[i][1] = old_coords[i][2]*Ly[0];
}"""
cLx = Constant(La)
cLy = Constant(Lb)
par_loop(periodic_kernel, dx,
{"new_coords": (new_coordinates, WRITE),
"old_coords": (old_coordinates, READ),
"Lx": (cLx, READ),
"Ly": (cLy, READ)})
if direction == "y":
# flip x and y coordinates
operator = np.asarray([[0, 1],
[1, 0]])
new_coordinates.dat.data[:] = np.dot(new_coordinates.dat.data, operator.T)
return mesh.Mesh(new_coordinates)
def _setup_solver(self):
state = self.state # just cutting down line length a bit
dt = state.timestepping.dt
beta = dt*state.timestepping.alpha
cp = state.parameters.cp
mu = state.mu
# Split up the rhs vector (symbolically)
u_in, rho_in, theta_in = split(state.xrhs)
# Build the reduced function space for u,rho
M = MixedFunctionSpace((state.V[0], state.V[1]))
w, phi = TestFunctions(M)
u, rho = TrialFunctions(M)
n = FacetNormal(state.mesh)
# Get background fields
thetabar = state.thetabar
rhobar = state.rhobar
pibar = exner(thetabar, rhobar, state)
pibar_rho = exner_rho(thetabar, rhobar, state)
pibar_theta = exner_theta(thetabar, rhobar, state)
# Analytical (approximate) elimination of theta
k = state.k # Upward pointing unit vector
theta = -dot(k,u)*dot(k,grad(thetabar))*beta + theta_in
# Only include theta' (rather than pi') in the vertical
# component of the gradient
# the pi prime term (here, bars are for mean and no bars are
# for linear perturbations)
pi = pibar_theta*theta + pibar_rho*rho
# vertical projection
def V(u):
return k*inner(u,k)
eqn = (
inner(w, (u - u_in))*dx
- beta*cp*div(theta*V(w))*pibar*dx
# following does nothing but is preserved in the comments
# to remind us why (because V(w) is purely vertical.
# + beta*cp*jump(theta*V(w),n)*avg(pibar)*dS_v
- beta*cp*div(thetabar*w)*pi*dx
+ beta*cp*jump(thetabar*w,n)*avg(pi)*dS_v
+ (phi*(rho - rho_in) - beta*inner(grad(phi), u)*rhobar)*dx
+ beta*jump(phi*u, n)*avg(rhobar)*(dS_v + dS_h)
)
if mu is not None:
eqn += dt*mu*inner(w,k)*inner(u,k)*dx
aeqn = lhs(eqn)
Leqn = rhs(eqn)
# Place to put result of u rho solver
self.urho = Function(M)
# Boundary conditions (assumes extruded mesh)
dim = M.sub(0).ufl_element().value_shape()[0]
bc = ("0.0",)*dim
bcs = [DirichletBC(M.sub(0), Expression(bc), "bottom"),
DirichletBC(M.sub(0), Expression(bc), "top")]
# Solver for u, rho
urho_problem = LinearVariationalProblem(
aeqn, Leqn, self.urho, bcs=bcs)
self.urho_solver = LinearVariationalSolver(urho_problem,
solver_parameters=self.params,
options_prefix='ImplicitSolver')
# Reconstruction of theta
theta = TrialFunction(state.V[2])
gamma = TestFunction(state.V[2])
u, rho = self.urho.split()
self.theta = Function(state.V[2])
theta_eqn = gamma*(theta - theta_in +
dot(k,u)*dot(k,grad(thetabar))*beta)*dx
theta_problem = LinearVariationalProblem(lhs(theta_eqn),
rhs(theta_eqn),
self.theta)
self.theta_solver = LinearVariationalSolver(theta_problem,
options_prefix='thetabacksubstitution')
class HybridizationPC(SCBase):
needs_python_pmat = True
"""A Slate-based python preconditioner that solves a
mixed H(div)-conforming problem using hybridization.
Currently, this preconditioner supports the hybridization
of the RT and BDM mixed methods of arbitrary degree.
The forward eliminations and backwards reconstructions
are performed element-local using the Slate language.
"""
@timed_function("HybridInit")
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 _reconstruction_calls(self, split_mixed_op, split_trace_op):
"""This generates the reconstruction calls for the unknowns using the
Lagrange multipliers.
:arg split_mixed_op: a ``dict`` of split forms that make up the broken
mixed operator from the original problem.
:arg split_trace_op: a ``dict`` of split forms that make up the trace
contribution in the hybridized mixed system.
"""
from firedrake.assemble import create_assembly_callable
# We always eliminate the velocity block first
id0, id1 = (self.vidx, self.pidx)
# TODO: When PyOP2 is able to write into mixed dats,
# the reconstruction expressions can simplify into
# one clean expression.
A = Tensor(split_mixed_op[(id0, id0)])
B = Tensor(split_mixed_op[(id0, id1)])
C = Tensor(split_mixed_op[(id1, id0)])
D = Tensor(split_mixed_op[(id1, id1)])
K_0 = Tensor(split_trace_op[(0, id0)])
K_1 = Tensor(split_trace_op[(0, id1)])
# Split functions and reconstruct each bit separately
split_residual = self.broken_residual.split()
split_sol = self.broken_solution.split()
g = AssembledVector(split_residual[id0])
f = AssembledVector(split_residual[id1])
sigma = split_sol[id0]
u = split_sol[id1]
lambdar = AssembledVector(self.trace_solution)
M = D - C * A.inv * B
R = K_1.T - C * A.inv * K_0.T
u_rec = M.solve(f - C * A.inv * g - R * lambdar,
decomposition="PartialPivLU")
self._sub_unknown = create_assembly_callable(u_rec,
tensor=u,
form_compiler_parameters=self.ctx.fc_params)
sigma_rec = A.solve(g - B * AssembledVector(u) - K_0.T * lambdar,
decomposition="PartialPivLU")
self._elim_unknown = create_assembly_callable(sigma_rec,
tensor=sigma,
form_compiler_parameters=self.ctx.fc_params)
@timed_function("HybridUpdate")
def update(self, pc):
"""Update by assembling into the operator. No need to
reconstruct symbolic objects.
"""
self._assemble_S()
self.S.force_evaluation()
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 sc_solve(self, pc):
"""Solve the condensed linear system for the
condensed field.
:arg pc: a Preconditioner instance.
"""
# Solve the system for the Lagrange multipliers
with self.schur_rhs.dat.vec_ro as b:
if self.trace_ksp.getInitialGuessNonzero():
acc = self.trace_solution.dat.vec
else:
acc = self.trace_solution.dat.vec_wo
with acc as x_trace:
self.trace_ksp.solve(b, x_trace)
def 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 view(self, pc, viewer=None):
"""Viewer calls for the various configurable objects in this PC."""
super(HybridizationPC, self).view(pc, viewer)
viewer.pushASCIITab()
viewer.printfASCII("Solves K * P^-1 * K.T using local eliminations.\n")
viewer.printfASCII("KSP solver for the multipliers:\n")
viewer.pushASCIITab()
self.trace_ksp.view(viewer)
viewer.popASCIITab()
viewer.printfASCII("Locally reconstructing the broken solutions from the multipliers.\n")
viewer.pushASCIITab()
viewer.printfASCII("Project the broken hdiv solution into the HDiv space.\n")
viewer.popASCIITab()
