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 _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 compressible_hydrostatic_balance(state, theta0, rho0, pi0=None,
top=False, pi_boundary=Constant(1.0),
water_t=None,
solve_for_rho=False,
params=None):
"""
Compute a hydrostatically balanced density given a potential temperature
profile. By default, this uses a vertically-oriented hybridization
procedure for solving the resulting discrete systems.
: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.
:arg water_t: the initial total water mixing ratio field.
"""
# Calculate hydrostatic Pi
VDG = state.spaces("DG")
Vv = state.spaces("Vv")
W = MixedFunctionSpace((Vv, VDG))
v, pi = TrialFunctions(W)
dv, dpi = TestFunctions(W)
n = FacetNormal(state.mesh)
cp = state.parameters.cp
# add effect of density of water upon theta
theta = theta0
if water_t is not None:
theta = theta0 / (1 + water_t)
alhs = (
(cp*inner(v, dv) - cp*div(dv*theta)*pi)*dx
+ dpi*div(theta*v)*dx
)
if top:
bmeasure = ds_t
bstring = "bottom"
else:
bmeasure = ds_b
bstring = "top"
arhs = -cp*inner(dv, n)*theta*pi_boundary*bmeasure
# Possibly make g vary with spatial coordinates?
g = state.parameters.g
arhs -= g*inner(dv, state.k)*dx
bcs = [DirichletBC(W.sub(0), Constant(0.0), bstring)]
w = Function(W)
PiProblem = LinearVariationalProblem(alhs, arhs, w, bcs=bcs)
if params is None:
params = {'ksp_type': 'preonly',
'pc_type': 'python',
'mat_type': 'matfree',
'pc_python_type': 'gusto.VerticalHybridizationPC',
# Vertical trace system is only coupled vertically in columns
# block ILU is a direct solver!
'vert_hybridization': {'ksp_type': 'preonly',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}}
PiSolver = LinearVariationalSolver(PiProblem,
solver_parameters=params,
options_prefix="pisolver")
PiSolver.solve()
v, Pi = w.split()
if pi0 is not None:
pi0.assign(Pi)
if solve_for_rho:
w1 = Function(W)
v, rho = w1.split()
rho.interpolate(thermodynamics.rho(state.parameters, theta0, Pi))
v, rho = split(w1)
dv, dpi = TestFunctions(W)
pi = thermodynamics.pi(state.parameters, rho, theta0)
F = (
(cp*inner(v, dv) - cp*div(dv*theta)*pi)*dx
+ dpi*div(theta0*v)*dx
+ cp*inner(dv, n)*theta*pi_boundary*bmeasure
)
F += g*inner(dv, state.k)*dx
rhoproblem = NonlinearVariationalProblem(F, w1, bcs=bcs)
rhosolver = NonlinearVariationalSolver(rhoproblem, solver_parameters=params,
options_prefix="rhosolver")
rhosolver.solve()
v, rho_ = w1.split()
rho0.assign(rho_)
else:
rho0.interpolate(thermodynamics.rho(state.parameters, theta0, Pi))
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.
shape = Vv.finat_element.space_dimension() * np.prod(Vv.shape)
weight_kernel = """
for (int i=0; i<%d; ++i) {
w[i] += 1.0;
}""" % shape
self.weight = Function(Vv)
par_loop(weight_kernel, dx, {"w": (self.weight, INC)})
self.average_kernel = """
for (int i=0; i<%d; ++i) {
vec_out[i] += vec_in[i]/w[i];
}""" % shape
# 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()
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 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)
B_x = B_0 + B_2 * X**2 + B_4 * X**4 + B_6 * X**6
f_c = 4e3
d_c = 500
w_c = 24e3
B_y = d_c * (1 / (1 + exp(-2 * (y - Ly / 2 - w_c) / f_c)) + 1 /
(1 + exp(+2 * (y - Ly / 2 + w_c) / f_c)))
z_deep = -720
z_b = interpolate(max_value(B_x + B_y, z_deep), Q)
# We use units of megapascals-meters-years, whereas in the paper the parameters
# are reported in units of pascals, so these values include some conversion
# factors.
A = Constant(20)
C = Constant(1e-2)
# Create the bed friction dissipation functional
from icepack.constants import (ice_density as ρ_I, water_density as ρ_W,
gravity as g, weertman_sliding_law as m)
def friction(**kwargs):
keys = ('velocity', 'thickness', 'surface', 'friction')
u, h, s, C = map(kwargs.get, keys)
p_W = ρ_W * g * max_value(0, -(s - h))
p_I = ρ_I * g * h
N = p_I - p_W
τ_c = N / 2
def initialize(self, pc):
from firedrake import TrialFunction, TestFunction, dx, \
assemble, inner, grad, split, Constant, parameters
from firedrake.assemble import allocate_matrix
if pc.getType() != "python":
raise ValueError("Expecting PC type python")
prefix = pc.getOptionsPrefix() + "pcd_"
# we assume P has things stuffed inside of it
_, P = pc.getOperators()
context = P.getPythonContext()
test, trial = context.a.arguments()
if test.function_space() != trial.function_space():
raise ValueError("Pressure space test and trial space differ")
Q = test.function_space()
p = TrialFunction(Q)
q = TestFunction(Q)
mass = p*q*dx
# Regularisation to avoid having to think about nullspaces.
stiffness = inner(grad(p), grad(q))*dx + Constant(1e-6)*p*q*dx
opts = PETSc.Options()
# we're inverting Mp and Kp, so default them to assembled.
# Fp only needs its action, so default it to mat-free.
# These can of course be overridden.
# only Fp is referred to in update, so that's the only
# one we stash.
default = parameters["default_matrix_type"]
Mp_mat_type = opts.getString(prefix+"Mp_mat_type", default)
Kp_mat_type = opts.getString(prefix+"Kp_mat_type", default)
self.Fp_mat_type = opts.getString(prefix+"Fp_mat_type", "matfree")
Mp = assemble(mass, form_compiler_parameters=context.fc_params,
mat_type=Mp_mat_type,
options_prefix=prefix + "Mp_")
Kp = assemble(stiffness, form_compiler_parameters=context.fc_params,
mat_type=Kp_mat_type,
options_prefix=prefix + "Kp_")
# FIXME: Should we transfer nullspaces over. I think not.
Mksp = PETSc.KSP().create(comm=pc.comm)
Mksp.incrementTabLevel(1, parent=pc)
Mksp.setOptionsPrefix(prefix + "Mp_")
Mksp.setOperators(Mp.petscmat)
Mksp.setUp()
Mksp.setFromOptions()
self.Mksp = Mksp
Kksp = PETSc.KSP().create(comm=pc.comm)
Kksp.incrementTabLevel(1, parent=pc)
Kksp.setOptionsPrefix(prefix + "Kp_")
Kksp.setOperators(Kp.petscmat)
Kksp.setUp()
Kksp.setFromOptions()
self.Kksp = Kksp
state = context.appctx["state"]
Re = context.appctx.get("Re", 1.0)
velid = context.appctx["velocity_space"]
u0 = split(state)[velid]
fp = 1.0/Re * inner(grad(p), grad(q))*dx + inner(u0, grad(p))*q*dx
self.Re = Re
self.Fp = allocate_matrix(fp, form_compiler_parameters=context.fc_params,
mat_type=self.Fp_mat_type,
options_prefix=prefix + "Fp_")
self._assemble_Fp = functools.partial(assemble,
fp,
tensor=self.Fp,
form_compiler_parameters=context.fc_params,
mat_type=self.Fp_mat_type,
assembly_type="residual")
self._assemble_Fp()
Fpmat = self.Fp.petscmat
self.workspace = [Fpmat.createVecLeft() for i in (0, 1)]
def initialize(self, pc):
"""Set up the problem context. Take the original
mixed problem and reformulate the problem as a
hybridized mixed system.
A KSP is created for the Lagrange multiplier system.
"""
from firedrake import (FunctionSpace, Function, Constant,
TrialFunction, TrialFunctions, TestFunction,
DirichletBC, assemble)
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.formmanipulation import split_form
from ufl.algorithms.replace import replace
# Extract the problem context
prefix = pc.getOptionsPrefix() + "hybridization_"
_, P = pc.getOperators()
self.cxt = P.getPythonContext()
if not isinstance(self.cxt, ImplicitMatrixContext):
raise ValueError(
"The python context must be an ImplicitMatrixContext")
test, trial = self.cxt.a.arguments()
V = test.function_space()
mesh = V.mesh()
if len(V) != 2:
raise ValueError("Expecting two function spaces.")
if all(Vi.ufl_element().value_shape() for Vi in V):
raise ValueError("Expecting an H(div) x L2 pair of spaces.")
# Automagically determine which spaces are vector and scalar
for i, Vi in enumerate(V):
if Vi.ufl_element().sobolev_space().name == "HDiv":
self.vidx = i
else:
assert Vi.ufl_element().sobolev_space().name == "L2"
self.pidx = i
# Create the space of approximate traces.
W = V[self.vidx]
if W.ufl_element().family() == "Brezzi-Douglas-Marini":
tdegree = W.ufl_element().degree()
else:
try:
# If we have a tensor product element
h_deg, v_deg = W.ufl_element().degree()
tdegree = (h_deg - 1, v_deg - 1)
except TypeError:
tdegree = W.ufl_element().degree() - 1
TraceSpace = FunctionSpace(mesh, "HDiv Trace", tdegree)
# Break the function spaces and define fully discontinuous spaces
broken_elements = ufl.MixedElement(
[ufl.BrokenElement(Vi.ufl_element()) for Vi in V])
V_d = FunctionSpace(mesh, broken_elements)
# Set up the functions for the original, hybridized
# and schur complement systems
self.broken_solution = Function(V_d)
self.broken_residual = Function(V_d)
self.trace_solution = Function(TraceSpace)
self.unbroken_solution = Function(V)
self.unbroken_residual = Function(V)
# Set up the KSP for the hdiv residual projection
hdiv_mass_ksp = PETSc.KSP().create(comm=pc.comm)
hdiv_mass_ksp.setOptionsPrefix(prefix + "hdiv_residual_")
# HDiv mass operator
p = TrialFunction(V[self.vidx])
q = TestFunction(V[self.vidx])
mass = ufl.dot(p, q) * ufl.dx
# TODO: Bcs?
M = assemble(mass,
bcs=None,
form_compiler_parameters=self.cxt.fc_params)
M.force_evaluation()
Mmat = M.petscmat
hdiv_mass_ksp.setOperators(Mmat)
hdiv_mass_ksp.setUp()
hdiv_mass_ksp.setFromOptions()
self.hdiv_mass_ksp = hdiv_mass_ksp
# Storing the result of A.inv * r, where A is the HDiv
# mass matrix and r is the HDiv residual
self._primal_r = Function(V[self.vidx])
tau = TestFunction(V_d[self.vidx])
self._assemble_broken_r = create_assembly_callable(
ufl.dot(self._primal_r, tau) * ufl.dx,
tensor=self.broken_residual.split()[self.vidx],
form_compiler_parameters=self.cxt.fc_params)
# Create the symbolic Schur-reduction:
# Original mixed operator replaced with "broken"
# arguments
arg_map = {test: TestFunction(V_d), trial: TrialFunction(V_d)}
Atilde = Tensor(replace(self.cxt.a, arg_map))
gammar = TestFunction(TraceSpace)
n = ufl.FacetNormal(mesh)
sigma = TrialFunctions(V_d)[self.vidx]
# We zero out the contribution of the trace variables on the exterior
# boundary. Extruded cells will have both horizontal and vertical
# facets
if mesh.cell_set._extruded:
trace_bcs = [
DirichletBC(TraceSpace, Constant(0.0), "on_boundary"),
DirichletBC(TraceSpace, Constant(0.0), "bottom"),
DirichletBC(TraceSpace, Constant(0.0), "top")
]
K = Tensor(
gammar('+') * ufl.dot(sigma, n) * ufl.dS_h +
gammar('+') * ufl.dot(sigma, n) * ufl.dS_v)
else:
trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), "on_boundary")]
K = Tensor(gammar('+') * ufl.dot(sigma, n) * ufl.dS)
# If boundary conditions are contained in the ImplicitMatrixContext:
if self.cxt.row_bcs:
raise NotImplementedError(
"Strong BCs not currently handled. Try imposing them weakly.")
# Assemble the Schur complement operator and right-hand side
self.schur_rhs = Function(TraceSpace)
self._assemble_Srhs = create_assembly_callable(
K * Atilde.inv * self.broken_residual,
tensor=self.schur_rhs,
form_compiler_parameters=self.cxt.fc_params)
schur_comp = K * Atilde.inv * K.T
self.S = allocate_matrix(schur_comp,
bcs=trace_bcs,
form_compiler_parameters=self.cxt.fc_params)
self._assemble_S = create_assembly_callable(
schur_comp,
tensor=self.S,
bcs=trace_bcs,
form_compiler_parameters=self.cxt.fc_params)
self._assemble_S()
self.S.force_evaluation()
Smat = self.S.petscmat
# Nullspace for the multiplier problem
nullspace = create_schur_nullspace(P, -K * Atilde, V, V_d, TraceSpace,
pc.comm)
if nullspace:
Smat.setNullSpace(nullspace)
# Set up the KSP for the system of Lagrange multipliers
trace_ksp = PETSc.KSP().create(comm=pc.comm)
trace_ksp.setOptionsPrefix(prefix)
trace_ksp.setOperators(Smat)
trace_ksp.setUp()
trace_ksp.setFromOptions()
self.trace_ksp = trace_ksp
split_mixed_op = dict(split_form(Atilde.form))
split_trace_op = dict(split_form(K.form))
# Generate reconstruction calls
self._reconstruction_calls(split_mixed_op, split_trace_op)
# NOTE: The projection stage *might* be replaced by a Fortin
# operator. We may want to allow the user to specify if they
# wish to use a Fortin operator over a projection, or vice-versa.
# In a future add-on, we can add a switch which chooses either
# the Fortin reconstruction or the usual KSP projection.
# Set up the projection KSP
hdiv_projection_ksp = PETSc.KSP().create(comm=pc.comm)
hdiv_projection_ksp.setOptionsPrefix(prefix + 'hdiv_projection_')
# Reuse the mass operator from the hdiv_mass_ksp
hdiv_projection_ksp.setOperators(Mmat)
# Construct the RHS for the projection stage
self._projection_rhs = Function(V[self.vidx])
self._assemble_projection_rhs = create_assembly_callable(
ufl.dot(self.broken_solution.split()[self.vidx], q) * ufl.dx,
tensor=self._projection_rhs,
form_compiler_parameters=self.cxt.fc_params)
# Finalize ksp setup
hdiv_projection_ksp.setUp()
hdiv_projection_ksp.setFromOptions()
self.hdiv_projection_ksp = hdiv_projection_ksp
stabilisation = args.stabilisation or 'lax_friedrichs'
if stabilisation == 'none' or family == 'cg-cg' or not nonlinear:
stabilisation = None
approach = args.approach or 'hessian'
kwargs = {
'approach': approach,
'inversion_level': 2, # TODO: Avoid this hard-code
# Space-time domain
'level': int(args.level or 0),
'end_time': float(args.end_time or 1440.0),
'num_meshes': int(args.num_meshes or 24),
# Physics
'bathymetry_cap': 30.0,
'base_viscosity': Constant(args.base_viscosity or 1.0e-03),
# Solver
'family': family,
'stabilisation': stabilisation,
'use_wetting_and_drying': False,
# Mesh adaptation
'adapt_field': args.adapt_field or 'all_avg',
'hessian_time_combination': args.time_combine or 'integrate',
'hessian_timestep_lag': int(args.hessian_lag
or 1), # NOTE: Not used in weighted Hessian
'normalisation': args.normalisation or 'complexity',
'norm_order': 1 if p is None else None if p == 'inf' else float(p),
'target': float(args.target or 5.0e+03),
'h_min': float(args.h_min or 1.0e+02),
diagnostics=diagnostics,
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()
# Isentropic background state
Tsurf = Constant(300.)
theta_b = Function(Vt).interpolate(Tsurf)
rho_b = Function(Vr)
# Calculate hydrostatic Pi
compressible_hydrostatic_balance(state, theta_b, rho_b, solve_for_rho=True)
x = SpatialCoordinate(mesh)
a = 5.0e3
deltaTheta = 1.0e-2
xc = 0.5 * L
xr = 4000.
zc = 3000.
zr = 2000.
r = sqrt(((x[0] - xc) / xr)**2 + ((x[1] - zc) / zr)**2)
"| Diagnostics:", tuple(nc_diag.keys()))
print("Starting Initial condition, and function setup at", ctime())
R, Omega, = 6371220., 7.292e-5
mesh = IcosahedralSphereMesh(radius=R, refinement_level=ref_level, degree=2)
mesh.init_cell_orientations(SpatialCoordinate(mesh))
x = SpatialCoordinate(mesh)
f = Function(FunctionSpace(mesh, "CG", 1))
f.interpolate(2 * Omega * x[2] / R)
g, H = 9.810616, 5960.
u_0 = 2 * pi * R / (12 * 24 * 60 * 60.)
uexpr = u_0 * as_vector([-x[1], x[0], 0.0]) / R
Dexpr = H - (R * Omega * u_0 + u_0**2 / 2.) * x[2]**2 / (g * R**2)
bexpr = Constant(0.)
# Build function spaces
degree = 2
family = ("DG", "BDM", "CG")
W0 = FunctionSpace(mesh, family[0], degree - 1, family[0])
W1_elt = FiniteElement(family[1], triangle, degree)
W1 = FunctionSpace(mesh, W1_elt, name="HDiv")
W2 = FunctionSpace(mesh, family[2], degree + 1)
M = MixedFunctionSpace((W1, W0))
# Set up functions
xn = Function(M)
un, Dn = xn.split()
un.rename('u')
Dn.rename('D')
# Define constant theta_e and water_t
Tsurf = 283.0
psurf = 85000.
pi_surf = (psurf / state.parameters.p_0)**state.parameters.kappa
humidity = 0.2
S = 1.3e-5
theta_surf = thermodynamics.theta(state.parameters, Tsurf, psurf)
theta_d = Function(Vt).interpolate(theta_surf * exp(S * z))
H = Function(Vt).assign(humidity)
# Calculate hydrostatic fields
unsaturated_hydrostatic_balance(state,
theta_d,
H,
pi_boundary=Constant(pi_surf))
# make mean fields
theta_b = Function(Vt).assign(theta0)
rho_b = Function(Vr).assign(rho0)
water_vb = Function(Vt).assign(water_v0)
# define perturbation to RH
xc = L / 2
zc = 800.
r1 = 300.
r2 = 200.
r = sqrt((x - xc)**2 + (z - zc)**2)
H_expr = conditional(
r > r1, 0.0,
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))
weight_kernel = """
for (int i=0; i<%d; ++i) {
for (int j=0; j<%d; ++j) {
w[i][j] += 1.0;
}}""" % shapes
self.weight = Function(V[self.vidx])
par_loop(weight_kernel, ufl.dx, {"w": (self.weight, INC)})
self.average_kernel = """
for (int i=0; i<%d; ++i) {
for (int j=0; j<%d; ++j) {
vec_out[i][j] += vec_in[i][j]/w[i][j];
}}""" % shapes
# 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 test_ice_shelf_prognostic_solver(solver_type):
ρ = ρ_I * (1 - ρ_I / ρ_W)
Lx, Ly = 20.0e3, 20.0e3
h0 = 500.0
u0 = 100.0
T = 254.15
model = icepack.models.IceShelf()
opts = {
"dirichlet_ids": [1],
"side_wall_ids": [3, 4],
"prognostic_solver_type": solver_type,
}
delta_x, error = [], []
for N in range(16, 65, 4):
delta_x.append(Lx / N)
mesh = firedrake.RectangleMesh(N, N, Lx, Ly)
x, y = firedrake.SpatialCoordinate(mesh)
degree = 2
V = firedrake.VectorFunctionSpace(mesh, "CG", degree)
Q = firedrake.FunctionSpace(mesh, "CG", degree)
q = (n + 1) * (ρ * g * h0 * u0 / 4)**n * icepack.rate_factor(T)
ux = (u0**(n + 1) + q * x)**(1 / (n + 1))
h = interpolate(h0 * u0 / ux, Q)
h_initial = h.copy(deepcopy=True)
A = Constant(icepack.rate_factor(T))
a = Constant(0)
solver = icepack.solvers.FlowSolver(model, **opts)
u_guess = interpolate(firedrake.as_vector((ux, 0)), V)
u = solver.diagnostic_solve(velocity=u_guess,
thickness=h,
fluidity=A,
strain_rate_min=Constant(0.0))
final_time, dt = 1.0, 1.0 / 12
num_timesteps = int(final_time / dt)
for k in range(num_timesteps):
h = solver.prognostic_solve(dt,
thickness=h,
velocity=u,
accumulation=a,
thickness_inflow=h_initial)
u = solver.diagnostic_solve(velocity=u,
thickness=h,
fluidity=A,
strain_rate_min=Constant(0.0))
error.append(norm(h - h_initial) / norm(h_initial))
print(delta_x[-1], error[-1])
log_delta_x = np.log2(np.array(delta_x))
log_error = np.log2(np.array(error))
slope, intercept = np.polyfit(log_delta_x, log_error, 1)
print(f"log(error) ~= {slope:g} * log(dx) + {intercept:g}")
assert slope > degree - 0.05
def test_hybrid_prognostic_solve(solver_type):
Lx, Ly = 20e3, 20e3
h0, dh = 500.0, 100.0
T = 254.15
u_in = 100.0
model = icepack.models.HybridModel()
opts = {
"dirichlet_ids": [1],
"side_wall_ids": [3, 4],
"prognostic_solver_type": solver_type,
}
Nx, Ny = 32, 32
mesh2d = firedrake.RectangleMesh(Nx, Ny, Lx, Ly)
mesh = firedrake.ExtrudedMesh(mesh2d, layers=1)
V = firedrake.VectorFunctionSpace(mesh,
"CG",
2,
vfamily="GL",
vdegree=1,
dim=2)
Q = firedrake.FunctionSpace(mesh, "CG", 2, vfamily="DG", vdegree=0)
x, y, ζ = firedrake.SpatialCoordinate(mesh)
height_above_flotation = 10.0
d = -ρ_I / ρ_W * (h0 - dh) + height_above_flotation
ρ = ρ_I - ρ_W * d**2 / (h0 - dh)**2
Z = icepack.rate_factor(T) * (ρ * g * h0 / 4)**n
q = 1 - (1 - (dh / h0) * (x / Lx))**(n + 1)
ux = u_in + Z * q * Lx * (h0 / dh) / (n + 1)
u0 = interpolate(firedrake.as_vector((ux, 0)), V)
thickness = h0 - dh * x / Lx
β = 1 / 2
α = β * ρ / ρ_I * dh / Lx
h = interpolate(h0 - dh * x / Lx, Q)
h_inflow = h.copy(deepcopy=True)
ds = (1 + β) * ρ / ρ_I * dh
s = interpolate(d + h0 - dh + ds * (1 - x / Lx), Q)
b = interpolate(s - h, Q)
C = interpolate(α * (ρ_I * g * thickness) * ux**(-1 / m), Q)
A = Constant(icepack.rate_factor(T))
final_time, dt = 1.0, 1.0 / 12
num_timesteps = int(final_time / dt)
solver = icepack.solvers.FlowSolver(model, **opts)
u = solver.diagnostic_solve(
velocity=u0,
thickness=h,
surface=s,
fluidity=A,
friction=C,
strain_rate_min=Constant(0.0),
)
a = firedrake.Function(Q)
h_n = solver.prognostic_solve(dt,
thickness=h,
velocity=u,
accumulation=a,
thickness_inflow=h_inflow)
a.interpolate((h_n - h) / dt)
for k in range(num_timesteps):
h = solver.prognostic_solve(dt,
thickness=h,
velocity=u,
accumulation=a,
thickness_inflow=h_inflow)
s = icepack.compute_surface(thickness=h, bed=b)
u = solver.diagnostic_solve(
velocity=u,
thickness=h,
surface=s,
fluidity=A,
friction=C,
strain_rate_min=Constant(0.0),
)
assert icepack.norm(h, norm_type="Linfty") < np.inf
}
discrete_turbines = True
# discrete_turbines = False
op = TurbineArrayOptions(**kwargs)
op.update({
'spun': False,
'di': create_directory(os.path.join(op.di, 'unsteady')),
# Extend to time-dependent case
'timestepper': 'CrankNicolson',
'dt': 5.0,
'dt_per_export': 1,
'end_time': 600.0,
# Crank down viscosity and plot vorticity
'base_viscosity': Constant(1.0e-05),
'characteristic_velocity': Constant(op.inflow_velocity),
'grad_depth_viscosity': True,
'max_reynolds_number': 5.0e+05,
'recover_vorticity': True,
# Only consider the first turbine
'region_of_interest': [op.region_of_interest[0]],
'array_ids': np.array([2]),
'farm_ids': (2, ),
})
# Solve forward problem
tp = AdaptiveTurbineProblem(op,
discrete_turbines=discrete_turbines,
ramp_dir=op.di)
def setup_tracer(dirname):
# declare grid shape, with length L and height H
L = 1000.
H = 1000.
nlayers = int(H / 100.)
ncolumns = int(L / 100.)
# make mesh
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=10.0, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname + "/tracer",
dumpfreq=1,
dumplist=['u'],
perturbation_fields=['theta', 'rho'])
parameters = CompressibleParameters()
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist,
diagnostic_fields=[Difference('theta', 'tracer')])
# declare initial fields
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()
# declare tracer field and a background field
tracer0 = state.fields("tracer", Vt)
# Isentropic background state
Tsurf = Constant(300.)
theta_b = Function(Vt).interpolate(Tsurf)
rho_b = Function(Vr)
# Calculate initial rho
compressible_hydrostatic_balance(state, theta_b, rho_b, solve_for_rho=True)
# set up perturbation to theta
xc = 500.
zc = 350.
rc = 250.
x = SpatialCoordinate(mesh)
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)
tracer0.interpolate(theta0)
state.initialise([('u', u0), ('rho', rho0), ('theta', theta0),
('tracer', tracer0)])
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, equation_form="advective")
# build advection dictionary
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)))
advected_fields.append(("tracer", SSPRK3(state, tracer0, thetaeqn)))
# Set up linear solver
linear_solver = CompressibleSolver(state)
compressible_forcing = CompressibleForcing(state)
# build time stepper
stepper = CrankNicolson(state, advected_fields, linear_solver,
compressible_forcing)
return stepper, 100.0
def __init__(self, mesh, vertical_degree=None, horizontal_degree=1,
family="RT",
Coriolis=None, sponge_function=None,
hydrostatic=None,
timestepping=None,
output=None,
parameters=None,
diagnostics=None,
fieldlist=None,
diagnostic_fields=None,
u_bc_ids=None):
self.family = family
self.vertical_degree = vertical_degree
self.horizontal_degree = horizontal_degree
self.Omega = Coriolis
self.mu = sponge_function
self.hydrostatic = hydrostatic
self.timestepping = timestepping
if output is None:
raise RuntimeError("You must provide a directory name for dumping results")
else:
self.output = output
self.parameters = parameters
if fieldlist is None:
raise RuntimeError("You must provide a fieldlist containing the names of the prognostic fields")
else:
self.fieldlist = fieldlist
if diagnostics is not None:
self.diagnostics = diagnostics
else:
self.diagnostics = Diagnostics(*fieldlist)
if diagnostic_fields is not None:
self.diagnostic_fields = diagnostic_fields
else:
self.diagnostic_fields = []
if u_bc_ids is not None:
self.u_bc_ids = u_bc_ids
else:
self.u_bc_ids = []
# The mesh
self.mesh = mesh
# Build the spaces
self._build_spaces(mesh, vertical_degree, horizontal_degree, family)
# Allocate state
self._allocate_state()
if self.output.dumplist is None:
self.output.dumplist = fieldlist
self.fields = FieldCreator(fieldlist, self.xn, self.output.dumplist)
# set up bcs
V = self.fields('u').function_space()
self.bcs = []
if V.extruded:
self.bcs.append(DirichletBC(V, 0.0, "bottom"))
self.bcs.append(DirichletBC(V, 0.0, "top"))
for id in self.u_bc_ids:
self.bcs.append(DirichletBC(V, 0.0, id))
self.dumpfile = None
# figure out if we're on a sphere
try:
self.on_sphere = (mesh._base_mesh.geometric_dimension() == 3 and mesh._base_mesh.topological_dimension() == 2)
except AttributeError:
self.on_sphere = (mesh.geometric_dimension() == 3 and mesh.topological_dimension() == 2)
# build the vertical normal and define perp for 2d geometries
dim = mesh.topological_dimension()
if self.on_sphere:
x = SpatialCoordinate(mesh)
R = sqrt(inner(x, x))
self.k = interpolate(x/R, mesh.coordinates.function_space())
if dim == 2:
outward_normals = CellNormal(mesh)
self.perp = lambda u: cross(outward_normals, u)
else:
kvec = [0.0]*dim
kvec[dim-1] = 1.0
self.k = Constant(kvec)
if dim == 2:
self.perp = lambda u: as_vector([-u[1], u[0]])
# project test function for hydrostatic case
if self.hydrostatic:
self.h_project = lambda u: u - self.k*inner(u, self.k)
else:
self.h_project = lambda u: u
# Constant to hold current time
self.t = Constant(0.0)
# setup logger
logger.setLevel(output.log_level)
set_log_handler(mesh.comm)
logger.info("Timestepping parameters that take non-default values:")
logger.info(", ".join("%s: %s" % item for item in vars(timestepping).items()))
if parameters is not None:
logger.info("Physical parameters that take non-default values:")
logger.info(", ".join("%s: %s" % item for item in vars(parameters).items()))
def streamplot(function,
resolution=None,
min_length=None,
max_time=None,
start_width=0.5,
end_width=1.5,
tolerance=3e-3,
loc_tolerance=1e-10,
seed=None,
**kwargs):
r"""Create a streamline plot of a vector field
Similar to matplotlib :func:`streamplot <matplotlib.pyplot.streamplot>`
:arg function: the Firedrake :class:`~.Function` to plot
:arg resolution: minimum spacing between streamlines (defaults to domain size / 20)
:arg min_length: minimum length of a streamline (defaults to 4x resolution)
:arg max_time: maximum time to integrate a streamline
:arg start_width: line width at beginning of streamline
:arg end_width: line width at end of streamline, to convey direction
:arg tolerance: dimensionless tolerance for adaptive ODE integration
:arg loc_tolerance: point location tolerance for :meth:`~firedrake.functions.Function.at`
:kwarg kwargs: same as for matplotlib :class:`~matplotlib.collections.LineCollection`
"""
import randomgen
if function.ufl_shape != (2, ):
raise ValueError("Streamplot only defined for 2D vector fields!")
axes = kwargs.pop("axes", None)
if axes is None:
figure = plt.figure()
axes = figure.add_subplot(111)
mesh = function.ufl_domain()
if resolution is None:
coords = mesh.coordinates.dat.data_ro
resolution = (coords.max(axis=0) - coords.min(axis=0)).max() / 20
if min_length is None:
min_length = 4 * resolution
if max_time is None:
area = assemble(Constant(1) * dx(mesh))
average_speed = np.sqrt(
assemble(inner(function, function) * dx) / area)
max_time = 50 * min_length / average_speed
streamplotter = Streamplotter(function, resolution, min_length, max_time,
tolerance, loc_tolerance)
# TODO: better way of seeding start points
shape = streamplotter._grid.shape
xmin = streamplotter._grid_point((0, 0))
xmax = streamplotter._grid_point((shape[0] - 2, shape[1] - 2))
X, Y = np.meshgrid(np.linspace(xmin[0], xmax[0], shape[0] - 2),
np.linspace(xmin[1], xmax[1], shape[1] - 2))
start_points = np.vstack((X.ravel(), Y.ravel())).T
# Randomly shuffle the start points
generator = randomgen.MT19937(seed).generator
for x in generator.permutation(np.array(start_points)):
streamplotter.add_streamline(x)
# Colors are determined by the speed, thicknesses by arc length
speeds = []
widths = []
for streamline in streamplotter.streamlines:
velocity = np.array(function.at(streamline, tolerance=loc_tolerance))
speed = np.sqrt(np.sum(velocity**2, axis=1))
speeds.extend(speed[:-1])
delta = np.sqrt(np.sum(np.diff(streamline, axis=0)**2, axis=1))
arc_length = np.cumsum(delta)
length = arc_length[-1]
s = arc_length / length
linewidth = (1 - s) * start_width + s * end_width
widths.extend(linewidth)
points = []
for streamline in streamplotter.streamlines:
pts = streamline.reshape(-1, 1, 2)
points.extend(np.hstack((pts[:-1], pts[1:])))
speeds = np.array(speeds)
widths = np.array(widths)
vmin = kwargs.pop("vmin", speeds.min())
vmax = kwargs.pop("vmax", speeds.max())
norm = kwargs.pop("norm", matplotlib.colors.Normalize(vmin=vmin,
vmax=vmax))
cmap = plt.get_cmap(kwargs.pop("cmap", None))
collection = LineCollection(points, cmap=cmap, norm=norm, linewidth=widths)
collection.set_array(speeds)
axes.add_collection(collection)
_autoscale_view(axes, function.ufl_domain().coordinates.dat.data_ro)
return collection
('rho', SSPRK3(state, rho0, rhoeqn)),
('theta', SSPRK3(state, theta0, thetaeqn, limiter=limiter)),
('water_v', SSPRK3(state, water_v0, thetaeqn, limiter=limiter)),
('water_c', SSPRK3(state, water_c0, thetaeqn, limiter=limiter))]
# Set up linear solver
linear_solver = CompressibleSolver(state, moisture=moisture)
# Set up forcing
compressible_forcing = CompressibleForcing(state, moisture=moisture, euler_poincare=euler_poincare)
# diffusion
bcs = [DirichletBC(Vu, 0.0, "bottom"),
DirichletBC(Vu, 0.0, "top")]
diffused_fields = []
if diffusion:
diffused_fields.append(('u', InteriorPenalty(state, Vu, kappa=Constant(60.),
mu=Constant(10./deltax), bcs=bcs)))
# define condensation
physics_list = [Condensation(state)]
# build time stepper
stepper = CrankNicolson(state, advected_fields, linear_solver,
compressible_forcing, physics_list=physics_list,
diffused_fields=diffused_fields)
stepper.run(t=0, tmax=tmax)
else:
# Fixed quadrature with 64 points gives absolute errors below 1e-13
# for a quantity of order 1e-3.
v, _ = integrate.fixed_quad(D_integrand, theta_0, angles[ii], n=64)
val[ii] = -(R / g) * v
return val
# Get coordinates to pass to Dval function
D0 = Function(W0)
W = VectorFunctionSpace(mesh, D0.ufl_element())
X = interpolate(mesh.coordinates, W)
D0.dat.data[:] = Dval(X.dat.data_ro)
# Adjust mean value of initial D
C = Function(D0.function_space()).assign(Constant(1.0))
area = assemble(C * dx)
Dmean = assemble(D0 * dx) / area
D0 -= Dmean
D0 += Constant(D_mean)
# Dpert
D_p = Function(W0)
Dpert = D_bump * cos(theta) * exp(-(lamda / a)**2) * exp(-(
(theta_2 - theta) / b)**2)
D_p.interpolate(Dpert)
# Dexpr = Dbar + Dpert
Dexpr = D0 + D_p
bexpr = Constant(0)
# Set up functions
# Set up linear solver
linear_solver = CompressibleSolver(state, moisture=moisture)
# Set up forcing
compressible_forcing = CompressibleForcing(state, moisture=moisture)
# diffusion
bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")]
diffused_fields = []
if diffusion:
diffused_fields.append(('u',
InteriorPenalty(state,
Vu,
kappa=Constant(60.),
mu=Constant(10. / deltax),
bcs=bcs)))
# define condensation
physics_list = [Condensation(state)]
# build time stepper
stepper = CrankNicolson(state,
advected_fields,
linear_solver,
compressible_forcing,
physics_list=physics_list,
diffused_fields=diffused_fields)
stepper.run(t=0, tmax=tmax)
def getForm(F, butch, t, dt, u0, bcs=None):
"""Given a time-dependent variational form and a
:class:`ButcherTableau`, produce UFL for the s-stage RK method.
:arg F: UFL form for the semidiscrete ODE/DAE
:arg butch: the :class:`ButcherTableau` for the RK method being used to
advance in time.
:arg t: a :class:`Constant` referring to the current time level.
Any explicit time-dependence in F is included
:arg dt: a :class:`Constant` referring to the size of the current
time step.
:arg u0: a :class:`Function` referring to the state of
the PDE system at time `t`
:arg bcs: optionally, a :class:`DirichletBC` object (or iterable thereof)
containing (possible time-dependent) boundary conditions imposed
on the system.
On output, we return a tuple consisting of four parts:
- Fnew, the :class:`Form`
- k, the :class:`firedrake.Function` holding all the stages.
It lives in a :class:`firedrake.FunctionSpace` corresponding to the
s-way tensor product of the space on which the semidiscrete
form lives.
- `bcnew`, a list of :class:`firedrake.DirichletBC` objects to be posed
on the stages,
- `gblah`, a list of pairs of the form (f, expr), where f is
a :class:`firedrake.Function` and expr is a :class:`ufl.Expr`.
at each time step, each expr needs to be re-interpolated/projected
onto the corresponding f in order for Firedrake to pick up that
time-dependent boundary conditions need to be re-applied.
"""
v = F.arguments()[0]
V = v.function_space()
assert V == u0.function_space()
A = numpy.array([[Constant(aa) for aa in arow] for arow in butch.A])
c = numpy.array([Constant(ci) for ci in butch.c])
num_stages = len(c)
num_fields = len(V)
Vbig = numpy.prod([V for i in range(num_stages)])
# Silence a warning about transfer managers when we
# coarsen coefficients in V
push_parent(V.dm, Vbig.dm)
vnew = TestFunction(Vbig)
k = Function(Vbig)
if len(V) == 1:
u0bits = [u0]
vbits = [v]
if num_stages == 1:
vbigbits = [vnew]
kbits = [k]
else:
vbigbits = split(vnew)
kbits = split(k)
else:
u0bits = split(u0)
vbits = split(v)
vbigbits = split(vnew)
kbits = split(k)
kbits_np = numpy.zeros((num_stages, num_fields), dtype="object")
for i in range(num_stages):
for j in range(num_fields):
kbits_np[i, j] = kbits[i * num_fields + j]
Ak = A @ kbits_np
Fnew = Zero()
for i in range(num_stages):
repl = {t: t + c[i] * dt}
for j, (ubit, vbit, kbit) in enumerate(zip(u0bits, vbits, kbits)):
repl[ubit] = ubit + dt * Ak[i, j]
repl[vbit] = vbigbits[num_fields * i + j]
repl[TimeDerivative(ubit)] = kbits_np[i, j]
if (len(ubit.ufl_shape) == 1):
for kk, kbitbit in enumerate(kbits_np[i, j]):
repl[TimeDerivative(ubit[kk])] = kbitbit
repl[ubit[kk]] = repl[ubit][kk]
repl[vbit[kk]] = repl[vbit][kk]
Fnew += replace(F, repl)
bcnew = []
gblah = []
if bcs is None:
bcs = []
for bc in bcs:
if isinstance(bc.domain_args[0], str):
boundary = bc.domain_args[0]
else:
boundary = ()
try:
for j in bc.sub_domain:
boundary += j
except TypeError:
boundary = (bc.sub_domain, )
gfoo = expand_derivatives(diff(bc._original_arg, t))
if len(V) == 1:
for i in range(num_stages):
gcur = replace(gfoo, {t: t + c[i] * dt})
try:
gdat = interpolate(gcur, V)
except NotImplementedError:
gdat = project(gcur, V)
gblah.append((gdat, gcur))
bcnew.append(DirichletBC(Vbig[i], gdat, boundary))
else:
sub = bc.function_space_index()
for i in range(num_stages):
gcur = replace(gfoo, {t: t + butch.c[i] * dt})
try:
gdat = interpolate(gcur, V.sub(sub))
except NotImplementedError:
gdat = project(gcur, V)
gblah.append((gdat, gcur))
bcnew.append(
DirichletBC(Vbig[sub + num_fields * i], gdat, boundary))
return Fnew, k, bcnew, gblah
# Define constant theta_e and water_t
Tsurf = 283.0
psurf = 85000.
exner_surf = (psurf / state.parameters.p_0)**state.parameters.kappa
humidity = 0.2
S = 1.3e-5
theta_surf = thermodynamics.theta(state.parameters, Tsurf, psurf)
theta_d = Function(Vt).interpolate(theta_surf * exp(S * z))
H = Function(Vt).assign(humidity)
# Calculate hydrostatic fields
unsaturated_hydrostatic_balance(state,
theta_d,
H,
exner_boundary=Constant(exner_surf))
# make mean fields
theta_b = Function(Vt).assign(theta0)
rho_b = Function(Vr).assign(rho0)
water_vb = Function(Vt).assign(water_v0)
# define perturbation to RH
xc = L / 2
zc = 800.
r1 = 300.
r2 = 200.
r = sqrt((x - xc)**2 + (z - zc)**2)
H_expr = conditional(
r > r1, 0.0,
diagnostics=diagnostics,
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()
# Isentropic background state
Tsurf = Constant(300.)
theta_b = Function(Vt).interpolate(Tsurf)
rho_b = Function(Vr)
# Calculate hydrostatic Pi
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)
def setup(self, state):
space = state.spaces("HDiv")
super(SawyerEliassenU, self).setup(state, space=space)
u = state.fields("u")
b = state.fields("b")
v = inner(u, as_vector([0., 1., 0.]))
# spaces
V0 = FunctionSpace(state.mesh, "CG", 2)
Vu = u.function_space()
# project b to V0
self.b_v0 = Function(V0)
btri = TrialFunction(V0)
btes = TestFunction(V0)
a = inner(btes, btri) * dx
L = inner(btes, b) * dx
projectbproblem = LinearVariationalProblem(a, L, self.b_v0)
self.project_b_solver = LinearVariationalSolver(
projectbproblem, solver_parameters={'ksp_type': 'cg'})
# project v to V0
self.v_v0 = Function(V0)
vtri = TrialFunction(V0)
vtes = TestFunction(V0)
a = inner(vtes, vtri) * dx
L = inner(vtes, v) * dx
projectvproblem = LinearVariationalProblem(a, L, self.v_v0)
self.project_v_solver = LinearVariationalSolver(
projectvproblem, solver_parameters={'ksp_type': 'cg'})
# stm/psi is a stream function
self.stm = Function(V0)
psi = TrialFunction(V0)
xsi = TestFunction(V0)
f = state.parameters.f
H = state.parameters.H
L = state.parameters.L
dbdy = state.parameters.dbdy
x, y, z = SpatialCoordinate(state.mesh)
bcs = [DirichletBC(V0, 0., "bottom"),
DirichletBC(V0, 0., "top")]
Mat = as_matrix([[b.dx(2), 0., -f*self.v_v0.dx(2)],
[0., 0., 0.],
[-self.b_v0.dx(0), 0., f**2+f*self.v_v0.dx(0)]])
Equ = (
inner(grad(xsi), dot(Mat, grad(psi)))
- dbdy*inner(grad(xsi), as_vector([-v, 0., f*(z-H/2)]))
)*dx
# fourth-order terms
if state.parameters.fourthorder:
eps = Constant(0.0001)
brennersigma = Constant(10.0)
n = FacetNormal(state.mesh)
deltax = Constant(state.parameters.deltax)
deltaz = Constant(state.parameters.deltaz)
nn = as_matrix([[sqrt(brennersigma/Constant(deltax)), 0., 0.],
[0., 0., 0.],
[0., 0., sqrt(brennersigma/Constant(deltaz))]])
mu = as_matrix([[1., 0., 0.],
[0., 0., 0.],
[0., 0., H/L]])
# anisotropic form
Equ += eps*(
div(dot(mu, grad(psi)))*div(dot(mu, grad(xsi)))*dx
- (
avg(dot(dot(grad(grad(psi)), n), n))*jump(grad(xsi), n=n)
+ avg(dot(dot(grad(grad(xsi)), n), n))*jump(grad(psi), n=n)
- jump(nn*grad(psi), n=n)*jump(nn*grad(xsi), n=n)
)*(dS_h + dS_v)
)
Au = lhs(Equ)
Lu = rhs(Equ)
stmproblem = LinearVariationalProblem(Au, Lu, self.stm, bcs=bcs)
self.stream_function_solver = LinearVariationalSolver(
stmproblem, solver_parameters={'ksp_type': 'cg'})
# solve for sawyer_eliassen u
self.u = Function(Vu)
utrial = TrialFunction(Vu)
w = TestFunction(Vu)
a = inner(w, utrial)*dx
L = (w[0]*(-self.stm.dx(2))+w[2]*(self.stm.dx(0)))*dx
ugproblem = LinearVariationalProblem(a, L, self.u)
self.sawyer_eliassen_u_solver = LinearVariationalSolver(
ugproblem, solver_parameters={'ksp_type': 'cg'})
def _ad_assign_numpy(dst, src, offset):
dst.assign(Constant(numpy.reshape(src[offset:offset + dst.value_size()], dst.ufl_shape)))
offset += dst.value_size()
return dst, offset
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
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 saturated_hydrostatic_balance(state, theta_e, water_t, pi0=None,
top=False, pi_boundary=Constant(1.0),
max_outer_solve_count=40,
max_theta_solve_count=5,
max_inner_solve_count=3):
"""
Given a wet equivalent potential temperature, theta_e, and the total moisture
content, water_t, compute a hydrostatically balance virtual potential temperature,
dry density and water vapour profile.
The general strategy is to split up the solving into two steps:
1) finding rho to balance the theta profile
2) finding theta_v and r_v to get back theta_e and saturation
We iteratively solve these steps until we (hopefully)
converge to a solution.
:arg state: The :class:`State` object.
:arg theta_e: The initial wet equivalent potential temperature profile.
:arg water_t: The total water pseudo-mixing ratio profile.
:arg pi0: Optional function to put exner pressure into.
:arg top: If True, set a boundary condition at the top, otherwise
it will be at the bottom.
:arg pi_boundary: The value of pi on the specified boundary.
:arg max_outer_solve_count: Max number of outer iterations for balance solver.
:arg max_theta_solve_count: Max number of iterations for theta solver (middle part of solve).
:arg max_inner_solve_count: Max number of iterations on the inner most
loop for the water vapour solver.
"""
theta0 = state.fields('theta')
rho0 = state.fields('rho')
water_v0 = state.fields('water_v')
# Calculate hydrostatic Pi
Vt = theta0.function_space()
Vr = rho0.function_space()
VDG = state.spaces("DG")
if any(deg > 2 for deg in VDG.ufl_element().degree()):
logger.warning("default quadrature degree most likely not sufficient for this degree element")
theta0.interpolate(theta_e)
water_v0.interpolate(water_t)
if state.vertical_degree == 0:
boundary_method = Boundary_Method.physics
else:
boundary_method = None
rho_h = Function(Vr)
Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element()))
rho_averaged = Function(Vt)
rho_recoverer = Recoverer(rho0, rho_averaged, VDG=Vt_broken, boundary_method=boundary_method)
w_h = Function(Vt)
theta_h = Function(Vt)
theta_e_test = Function(Vt)
delta = 0.8
# expressions for finding theta0 and water_v0 from theta_e and water_t
pie = thermodynamics.pi(state.parameters, rho_averaged, theta0)
p = thermodynamics.p(state.parameters, pie)
T = thermodynamics.T(state.parameters, theta0, pie, water_v0)
r_v_expr = thermodynamics.r_sat(state.parameters, T, p)
theta_e_expr = thermodynamics.theta_e(state.parameters, T, p, water_v0, water_t)
for i in range(max_outer_solve_count):
# solve for rho with theta_vd and w_v guesses
compressible_hydrostatic_balance(state, theta0, rho_h, top=top,
pi_boundary=pi_boundary, water_t=water_t,
solve_for_rho=True)
# damp solution
rho0.assign(rho0 * (1 - delta) + delta * rho_h)
theta_e_test.assign(theta_e_expr)
if errornorm(theta_e_test, theta_e) < 1e-8:
break
# calculate averaged rho
rho_recoverer.project()
# now solve for r_v
for j in range(max_theta_solve_count):
theta_h.interpolate(theta_e / theta_e_expr * theta0)
theta0.assign(theta0 * (1 - delta) + delta * theta_h)
# break when close enough
if errornorm(theta_e_test, theta_e) < 1e-6:
break
for k in range(max_inner_solve_count):
w_h.interpolate(r_v_expr)
water_v0.assign(water_v0 * (1 - delta) + delta * w_h)
# break when close enough
theta_e_test.assign(theta_e_expr)
if errornorm(theta_e_test, theta_e) < 1e-6:
break
if i == max_outer_solve_count:
raise RuntimeError('Hydrostatic balance solve has not converged within %i' % i, 'iterations')
if pi0 is not None:
pie = thermodynamics.pi(state.parameters, rho0, theta0)
pi0.interpolate(pie)
# do one extra solve for rho
compressible_hydrostatic_balance(state, theta0, rho0, top=top,
pi_boundary=pi_boundary,
water_t=water_t, solve_for_rho=True)
def initialize(self, pc):
"""Set up the problem context. Take the original
mixed problem and reformulate the problem as a
hybridized mixed system.
A KSP is created for the Lagrange multiplier system.
"""
from firedrake import (FunctionSpace, Function, Constant,
TrialFunction, TrialFunctions, TestFunction,
DirichletBC, assemble)
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.formmanipulation import split_form
from ufl.algorithms.replace import replace
# Extract the problem context
prefix = pc.getOptionsPrefix() + "hybridization_"
_, P = pc.getOperators()
self.cxt = P.getPythonContext()
if not isinstance(self.cxt, ImplicitMatrixContext):
raise ValueError("The python context must be an ImplicitMatrixContext")
test, trial = self.cxt.a.arguments()
V = test.function_space()
mesh = V.mesh()
if len(V) != 2:
raise ValueError("Expecting two function spaces.")
if all(Vi.ufl_element().value_shape() for Vi in V):
raise ValueError("Expecting an H(div) x L2 pair of spaces.")
# Automagically determine which spaces are vector and scalar
for i, Vi in enumerate(V):
if Vi.ufl_element().sobolev_space().name == "HDiv":
self.vidx = i
else:
assert Vi.ufl_element().sobolev_space().name == "L2"
self.pidx = i
# Create the space of approximate traces.
W = V[self.vidx]
if W.ufl_element().family() == "Brezzi-Douglas-Marini":
tdegree = W.ufl_element().degree()
else:
try:
# If we have a tensor product element
h_deg, v_deg = W.ufl_element().degree()
tdegree = (h_deg - 1, v_deg - 1)
except TypeError:
tdegree = W.ufl_element().degree() - 1
TraceSpace = FunctionSpace(mesh, "HDiv Trace", tdegree)
# Break the function spaces and define fully discontinuous spaces
broken_elements = ufl.MixedElement([ufl.BrokenElement(Vi.ufl_element()) for Vi in V])
V_d = FunctionSpace(mesh, broken_elements)
# Set up the functions for the original, hybridized
# and schur complement systems
self.broken_solution = Function(V_d)
self.broken_residual = Function(V_d)
self.trace_solution = Function(TraceSpace)
self.unbroken_solution = Function(V)
self.unbroken_residual = Function(V)
# Set up the KSP for the hdiv residual projection
hdiv_mass_ksp = PETSc.KSP().create(comm=pc.comm)
hdiv_mass_ksp.setOptionsPrefix(prefix + "hdiv_residual_")
# HDiv mass operator
p = TrialFunction(V[self.vidx])
q = TestFunction(V[self.vidx])
mass = ufl.dot(p, q)*ufl.dx
# TODO: Bcs?
M = assemble(mass, bcs=None, form_compiler_parameters=self.cxt.fc_params)
M.force_evaluation()
Mmat = M.petscmat
hdiv_mass_ksp.setOperators(Mmat)
hdiv_mass_ksp.setUp()
hdiv_mass_ksp.setFromOptions()
self.hdiv_mass_ksp = hdiv_mass_ksp
# Storing the result of A.inv * r, where A is the HDiv
# mass matrix and r is the HDiv residual
self._primal_r = Function(V[self.vidx])
tau = TestFunction(V_d[self.vidx])
self._assemble_broken_r = create_assembly_callable(
ufl.dot(self._primal_r, tau)*ufl.dx,
tensor=self.broken_residual.split()[self.vidx],
form_compiler_parameters=self.cxt.fc_params)
# Create the symbolic Schur-reduction:
# Original mixed operator replaced with "broken"
# arguments
arg_map = {test: TestFunction(V_d),
trial: TrialFunction(V_d)}
Atilde = Tensor(replace(self.cxt.a, arg_map))
gammar = TestFunction(TraceSpace)
n = ufl.FacetNormal(mesh)
sigma = TrialFunctions(V_d)[self.vidx]
if mesh.cell_set._extruded:
Kform = (gammar('+') * ufl.dot(sigma, n) * ufl.dS_h +
gammar('+') * ufl.dot(sigma, n) * ufl.dS_v)
else:
Kform = (gammar('+') * ufl.dot(sigma, n) * ufl.dS)
# Here we deal with boundaries. If there are Neumann
# conditions (which should be enforced strongly for
# H(div)xL^2) then we need to add jump terms on the exterior
# facets. If there are Dirichlet conditions (which should be
# enforced weakly) then we need to zero out the trace
# variables there as they are not active (otherwise the hybrid
# problem is not well-posed).
# If boundary conditions are contained in the ImplicitMatrixContext:
if self.cxt.row_bcs:
# Find all the subdomains with neumann BCS
# These are Dirichlet BCs on the vidx space
neumann_subdomains = set()
for bc in self.cxt.row_bcs:
if bc.function_space().index == self.pidx:
raise NotImplementedError("Dirichlet conditions for scalar variable not supported. Use a weak bc")
if bc.function_space().index != self.vidx:
raise NotImplementedError("Dirichlet bc set on unsupported space.")
# append the set of sub domains
subdom = bc.sub_domain
if isinstance(subdom, str):
neumann_subdomains |= set([subdom])
else:
neumann_subdomains |= set(as_tuple(subdom, int))
# separate out the top and bottom bcs
extruded_neumann_subdomains = neumann_subdomains & {"top", "bottom"}
neumann_subdomains = neumann_subdomains.difference(extruded_neumann_subdomains)
integrand = gammar * ufl.dot(sigma, n)
measures = []
trace_subdomains = []
if mesh.cell_set._extruded:
ds = ufl.ds_v
for subdomain in extruded_neumann_subdomains:
measures.append({"top": ufl.ds_t, "bottom": ufl.ds_b}[subdomain])
trace_subdomains.extend(sorted({"top", "bottom"} - extruded_neumann_subdomains))
else:
ds = ufl.ds
if "on_boundary" in neumann_subdomains:
measures.append(ds)
else:
measures.append(ds(tuple(neumann_subdomains)))
dirichlet_subdomains = set(mesh.exterior_facets.unique_markers) - neumann_subdomains
trace_subdomains.append(sorted(dirichlet_subdomains))
for measure in measures:
Kform += integrand*measure
trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), subdomain) for subdomain in trace_subdomains]
else:
# No bcs were provided, we assume weak Dirichlet conditions.
# We zero out the contribution of the trace variables on
# the exterior boundary. Extruded cells will have both
# horizontal and vertical facets
trace_subdomains = ["on_boundary"]
if mesh.cell_set._extruded:
trace_subdomains.extend(["bottom", "top"])
trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), subdomain) for subdomain in trace_subdomains]
# Make a SLATE tensor from Kform
K = Tensor(Kform)
# Assemble the Schur complement operator and right-hand side
self.schur_rhs = Function(TraceSpace)
self._assemble_Srhs = create_assembly_callable(
K * Atilde.inv * AssembledVector(self.broken_residual),
tensor=self.schur_rhs,
form_compiler_parameters=self.cxt.fc_params)
schur_comp = K * Atilde.inv * K.T
self.S = allocate_matrix(schur_comp, bcs=trace_bcs,
form_compiler_parameters=self.cxt.fc_params)
self._assemble_S = create_assembly_callable(schur_comp,
tensor=self.S,
bcs=trace_bcs,
form_compiler_parameters=self.cxt.fc_params)
self._assemble_S()
self.S.force_evaluation()
Smat = self.S.petscmat
# Nullspace for the multiplier problem
nullspace = create_schur_nullspace(P, -K * Atilde,
V, V_d, TraceSpace,
pc.comm)
if nullspace:
Smat.setNullSpace(nullspace)
# Set up the KSP for the system of Lagrange multipliers
trace_ksp = PETSc.KSP().create(comm=pc.comm)
trace_ksp.setOptionsPrefix(prefix)
trace_ksp.setOperators(Smat)
trace_ksp.setUp()
trace_ksp.setFromOptions()
self.trace_ksp = trace_ksp
split_mixed_op = dict(split_form(Atilde.form))
split_trace_op = dict(split_form(K.form))
# Generate reconstruction calls
self._reconstruction_calls(split_mixed_op, split_trace_op)
# NOTE: The projection stage *might* be replaced by a Fortin
# operator. We may want to allow the user to specify if they
# wish to use a Fortin operator over a projection, or vice-versa.
# In a future add-on, we can add a switch which chooses either
# the Fortin reconstruction or the usual KSP projection.
# Set up the projection KSP
hdiv_projection_ksp = PETSc.KSP().create(comm=pc.comm)
hdiv_projection_ksp.setOptionsPrefix(prefix + 'hdiv_projection_')
# Reuse the mass operator from the hdiv_mass_ksp
hdiv_projection_ksp.setOperators(Mmat)
# Construct the RHS for the projection stage
self._projection_rhs = Function(V[self.vidx])
self._assemble_projection_rhs = create_assembly_callable(
ufl.dot(self.broken_solution.split()[self.vidx], q)*ufl.dx,
tensor=self._projection_rhs,
form_compiler_parameters=self.cxt.fc_params)
# Finalize ksp setup
hdiv_projection_ksp.setUp()
hdiv_projection_ksp.setFromOptions()
self.hdiv_projection_ksp = hdiv_projection_ksp
def unsaturated_hydrostatic_balance(state, theta_d, H, pi0=None,
top=False, pi_boundary=Constant(1.0),
max_outer_solve_count=40,
max_inner_solve_count=20):
"""
Given vertical profiles for dry potential temperature
and relative humidity compute hydrostatically balanced
virtual potential temperature, dry density and water vapour profiles.
The general strategy is to split up the solving into two steps:
1) finding rho to balance the theta profile
2) finding theta_v and r_v to get back theta_d and H
We iteratively solve these steps until we (hopefully)
converge to a solution.
:arg state: The :class:`State` object.
:arg theta_d: The initial dry potential temperature profile.
:arg H: The relative humidity profile.
:arg pi0: Optional function to put exner pressure into.
:arg top: If True, set a boundary condition at the top, otherwise
it will be at the bottom.
:arg pi_boundary: The value of pi on the specified boundary.
:arg max_outer_solve_count: Max number of iterations for outer loop of balance solver.
:arg max_inner_solve_count: Max number of iterations for inner loop of balanace solver.
"""
theta0 = state.fields('theta')
rho0 = state.fields('rho')
water_v0 = state.fields('water_v')
# Calculate hydrostatic Pi
Vt = theta0.function_space()
Vr = rho0.function_space()
R_d = state.parameters.R_d
R_v = state.parameters.R_v
epsilon = R_d / R_v
VDG = state.spaces("DG")
if any(deg > 2 for deg in VDG.ufl_element().degree()):
logger.warning("default quadrature degree most likely not sufficient for this degree element")
# apply first guesses
theta0.assign(theta_d * 1.01)
water_v0.assign(0.01)
if state.vertical_degree == 0:
method = Boundary_Method.physics
else:
method = None
rho_h = Function(Vr)
rho_averaged = Function(Vt)
Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element()))
rho_recoverer = Recoverer(rho0, rho_averaged, VDG=Vt_broken, boundary_method=method)
w_h = Function(Vt)
delta = 1.0
# make expressions for determining water_v0
pie = thermodynamics.pi(state.parameters, rho_averaged, theta0)
p = thermodynamics.p(state.parameters, pie)
T = thermodynamics.T(state.parameters, theta0, pie, water_v0)
r_v_expr = thermodynamics.r_v(state.parameters, H, T, p)
# make expressions to evaluate residual
pi_ev = thermodynamics.pi(state.parameters, rho_averaged, theta0)
p_ev = thermodynamics.p(state.parameters, pi_ev)
T_ev = thermodynamics.T(state.parameters, theta0, pi_ev, water_v0)
RH_ev = thermodynamics.RH(state.parameters, water_v0, T_ev, p_ev)
RH = Function(Vt)
for i in range(max_outer_solve_count):
# solve for rho with theta_vd and w_v guesses
compressible_hydrostatic_balance(state, theta0, rho_h, top=top,
pi_boundary=pi_boundary, water_t=water_v0,
solve_for_rho=True)
# damp solution
rho0.assign(rho0 * (1 - delta) + delta * rho_h)
# calculate averaged rho
rho_recoverer.project()
RH.assign(RH_ev)
if errornorm(RH, H) < 1e-10:
break
# now solve for r_v
for j in range(max_inner_solve_count):
w_h.interpolate(r_v_expr)
water_v0.assign(water_v0 * (1 - delta) + delta * w_h)
# compute theta_vd
theta0.assign(theta_d * (1 + water_v0 / epsilon))
# test quality of solution by re-evaluating expression
RH.assign(RH_ev)
if errornorm(RH, H) < 1e-10:
break
if i == max_outer_solve_count:
raise RuntimeError('Hydrostatic balance solve has not converged within %i' % i, 'iterations')
if pi0 is not None:
pie = thermodynamics.pi(state.parameters, rho0, theta0)
pi0.interpolate(pie)
# do one extra solve for rho
compressible_hydrostatic_balance(state, theta0, rho0, top=top,
pi_boundary=pi_boundary,
water_t=water_v0, solve_for_rho=True)
def transmission(
mesh,
scatterer_bdy_id,
outer_bdy_id,
wave_number,
options_prefix=None,
solver_parameters=None,
fspace=None,
true_sol_grad_expr=None,
):
r"""
preconditioner_gamma and preconditioner_lambda are used to precondition
with the following equation:
\Delta u - \kappa^2 \gamma u = 0
(\partial_n - i\kappa\beta) u |_\Sigma = 0
"""
# need as tuple so can use integral measure
if isinstance(outer_bdy_id, int):
outer_bdy_id = [outer_bdy_id]
outer_bdy_id = tuple(outer_bdy_id)
u = TrialFunction(fspace)
v = TestFunction(fspace)
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)
n = FacetNormal(mesh)
metadata = {'quadrature_degree': 2 * fspace.ufl_element().degree()}
L = inner(inner(true_sol_grad_expr, n), v) * ds(scatterer_bdy_id,
metadata=metadata)
solution = Function(fspace)
# {{{ Used for preconditioning
if 'gamma' in solver_parameters or 'beta' in solver_parameters:
solver_params = dict(solver_parameters)
gamma = complex(solver_parameters.pop('gamma', 1.0))
import cmath
beta = complex(solver_parameters.pop('beta', cmath.sqrt(gamma)))
aP = 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)
else:
aP = None
solver_params = solver_parameters
# }}}
# prepare to set up pyamg preconditioner if using it
using_pyamg = solver_params['pc_type'] == 'pyamg'
if using_pyamg:
pyamg_tol = solver_parameters.get('pyamg_tol', None)
if pyamg_tol is not None:
pyamg_tol = float(pyamg_tol)
pyamg_maxiter = solver_params.get('pyamg_maxiter', None)
if pyamg_maxiter is not None:
pyamg_maxiter = int(pyamg_maxiter)
del solver_params['pc_type']
# Create a solver and return the KSP object with the solution so that can get
# PETSc information
# Create problem
problem = vs.LinearVariationalProblem(a, L, solution, aP=aP)
# Create solver and call solve
solver = vs.LinearVariationalSolver(problem,
solver_parameters=solver_params,
options_prefix=options_prefix)
# prepare to set up pyamg preconditioner if using it
if using_pyamg:
A = assemble(a).M.handle
pc = solver.snes.getKSP().pc
pc.setType(pc.Type.PYTHON)
pc.setPythonContext(
AMGTransmissionPreconditioner(wave_number,
fspace,
A,
tol=pyamg_tol,
maxiter=pyamg_maxiter,
use_plane_waves=True))
# If using pyamg as preconditioner, use it!
try:
solver.solve()
except ConvergenceError:
pass
return solver.snes, solution
def simulation(simulation_parameters,
diagnostic_values=None,
fields_to_output=None,
expected_u=False):
"""
This function sets up and runs a basic simulation of the stochastic
Camassa-Holm equation.
:arg simulation_parameters: a dictionary containing the parameters for
the simulation. Each value should be a tuple, which will normally
have only one element. Two-element tuples will also contain an index
for outputting -- these are variables being looped over by the experiment.
:arg diagnostic_values: a list of diagnostic numbers to output.
:arg fields_to_output: a list of diagnostic fields to be dumped.
"""
# ensure simulation parameters use ufl Constant
for key in ['alphasq', 'c0', 'gamma']:
if not isinstance(simulation_parameters[key][-1], Constant):
if len(simulation_parameters[key]) == 2:
simulation_parameters[key] = (
simulation_parameters[key][0],
Constant(simulation_parameters[key][1]))
elif len(simulation_parameters[key]) == 1:
simulation_parameters[key] = (Constant(
simulation_parameters[key][-1]), )
else:
raise ValueError(
'The value of simulation_parameters %s seems to be of length %d. It should be either 1 or 2.'
% (key, len(simulation_parameters[key])))
mesh = simulation_parameters['mesh'][-1]
dt = simulation_parameters['dt'][-1]
tmax = simulation_parameters['tmax'][-1]
ndump = simulation_parameters['ndump'][-1]
field_ndump = simulation_parameters['field_ndump'][-1]
file_name = simulation_parameters['file_name'][-1]
scheme = simulation_parameters['scheme'][-1]
allow_fail = simulation_parameters['allow_fail'][-1]
only_peakons = simulation_parameters['only_peakons'][-1]
peakon_speed = simulation_parameters['peakon_speed'][-1]
peakon_method = simulation_parameters['peakon_method'][-1]
num_Xis = simulation_parameters['num_Xis'][-1]
allow_fail = True if only_peakons else allow_fail
prognostic_variables = PrognosticVariables(scheme, mesh, num_Xis)
diagnostic_variables = DiagnosticVariables(
prognostic_variables,
fields_to_output,
diagnostic_values=diagnostic_values)
build_initial_conditions(prognostic_variables, simulation_parameters)
Xis = StochasticFunctions(prognostic_variables, simulation_parameters)
equations = Equations(prognostic_variables, simulation_parameters)
# set up peakon equations
if simulation_parameters['peakon_equations'][-1] == True:
peakon_equations = PeakonEquations(prognostic_variables,
simulation_parameters, peakon_speed,
peakon_method)
elif simulation_parameters['peakon_equations'][-1] == False:
peakon_equations = None
else:
raise ValueError('peakon_equations must be True or False, not %s' %
simulation_parameters['peakon_equations'][-1])
outputting = Outputting(prognostic_variables,
diagnostic_variables,
simulation_parameters,
peakon_equations=peakon_equations,
diagnostic_values=diagnostic_values)
diagnostic_equations = DiagnosticEquations(diagnostic_variables,
prognostic_variables,
outputting,
simulation_parameters)
diagnostic_equations.solve()
dumpn = 0
field_dumpn = 0
# want the update to always happen on first time step (so results are consistent with earlier results)
t = 0
failed = True if only_peakons else False
failed_time = np.nan
outputting.dump_diagnostics(t)
outputting.dump_fields(t)
# run simulation
while (t < tmax - 0.5 * dt):
t += dt
# update stochastic basis functions
Xis.update(t)
# solve problems
if allow_fail:
if not failed:
try:
equations.solve()
except ConvergenceError:
failed = True
failed_time = t
print("Solver failed at t = %f" % t)
else:
equations.solve()
if peakon_equations is not None:
peakon_equations.update()
# output if necessary
dumpn += 1
field_dumpn += 1
# do interpolations/projections for any outputing
if dumpn == ndump or field_dumpn == field_ndump:
if not failed:
diagnostic_equations.solve()
# output diagnostic values
if dumpn == ndump:
outputting.dump_diagnostics(t, failed=failed)
dumpn -= ndump
# output fields
if field_dumpn == field_ndump:
if not failed:
outputting.dump_fields(t)
field_dumpn -= field_ndump
outputting.store_times(failed_time)
outputting.data_file.close()
if expected_u:
return prognostic_variables.u
