def __init__(self, mesh, conditions, timestepping, params, output, solver_params):
super().__init__(mesh, conditions, timestepping, params, output, solver_params)
self.u0 = Function(self.V)
self.u1 = Function(self.V)
self.sigma0 = Function(self.S)
self.sigma1 = Function(self.S)
theta = conditions.theta
uh = (1-theta) * self.u0 + theta * self.u1
a = Function(self.U)
h = Function(self.U)
p = TestFunction(self.V)
q = TestFunction(self.S)
self.initial_condition((self.u0, conditions.ic['u']), (a, conditions.ic['a']), (h, conditions.ic['h']))
ep_dot = self.strain(grad(uh))
zeta = self.zeta(h, a, self.delta(uh))
eta = zeta * params.e ** (-2)
rheology = 2 * eta * ep_dot + (zeta - eta) * tr(ep_dot) * Identity(2) - 0.5 * self.Ice_Strength(h, a) * Identity(2)
self.initial_condition((self.sigma0, rheology),(self.sigma1, self.sigma0))
def sigma_next(timestep, zeta, ep_dot, sigma, P):
A = 1 + 0.25 * (timestep * params.e ** 2) / params.T
B = timestep * 0.125 * (1 - params.e ** 2) / params.T
rhs = (1 - (timestep * params.e ** 2) / (4 * params.T)) * sigma - timestep / params.T * (
0.125 * (1 - params.e ** 2) * tr(sigma) * Identity(2) - 0.25 * P * Identity(2) + zeta * ep_dot)
C = (rhs[0, 0] - rhs[1, 1]) / A
D = (rhs[0, 0] + rhs[1, 1]) / (A + 2 * B)
sigma00 = 0.5 * (C + D)
sigma11 = 0.5 * (D - C)
sigma01 = rhs[0, 1]
sigma = as_matrix([[sigma00, sigma01], [sigma01, sigma11]])
return sigma
s = sigma_next(self.timestep, zeta, ep_dot, self.sigma0, self.Ice_Strength(h, a))
sh = (1-theta) * s + theta * self.sigma0
eqn = self.momentum_equation(h, self.u1, self.u0, p, sh, params.rho, uh, conditions.ocean_curr,
params.rho_a, params.C_a, params.rho_w, params.C_w, conditions.geo_wind, params.cor, self.timestep, ind=self.ind)
tensor_eqn = inner(self.sigma1-s, q) * dx
if conditions.stabilised['state']:
alpha = conditions.stabilised['alpha']
eqn += stabilisation_term(alpha=alpha, zeta=avg(zeta), mesh=mesh, v=uh, test=p)
bcs = DirichletBC(self.V, conditions.bc['u'], "on_boundary")
uprob = NonlinearVariationalProblem(eqn, self.u1, bcs)
self.usolver = NonlinearVariationalSolver(uprob, solver_parameters=solver_params.bt_params)
sprob = NonlinearVariationalProblem(tensor_eqn, self.sigma1)
self.ssolver = NonlinearVariationalSolver(sprob, solver_parameters=solver_params.bt_params)
def split(self, fields):
from ufl import as_vector, replace
from firedrake import NonlinearVariationalProblem as NLVP, FunctionSpace
splits = self._splits.get(tuple(fields))
if splits is not None:
return splits
splits = []
problem = self._problem
splitter = ExtractSubBlock()
for field in fields:
try:
if len(field) > 1:
raise NotImplementedError("Can't split into subblock")
except TypeError:
# Just a single field, we can handle that
pass
F = splitter.split(problem.F, argument_indices=(field, ))
J = splitter.split(problem.J, argument_indices=(field, field))
us = problem.u.split()
subu = us[field]
vec = []
for i, u in enumerate(us):
for idx in numpy.ndindex(u.ufl_shape):
vec.append(u[idx])
u = as_vector(vec)
F = replace(F, {problem.u: u})
J = replace(J, {problem.u: u})
if problem.Jp is not None:
Jp = splitter.split(problem.Jp,
argument_indices=(field, field))
Jp = replace(Jp, {problem.u: u})
else:
Jp = None
bcs = []
for bc in problem.bcs:
if bc.function_space().index == field:
V = FunctionSpace(subu.ufl_domain(), subu.ufl_element())
bcs.append(
type(bc)(V,
bc.function_arg,
bc.sub_domain,
method=bc.method))
new_problem = NLVP(
F,
subu,
bcs=bcs,
J=J,
Jp=None,
form_compiler_parameters=problem.form_compiler_parameters)
new_problem._constant_jacobian = problem._constant_jacobian
splits.append(
type(self)(new_problem,
mat_type=self.mat_type,
pmat_type=self.pmat_type,
appctx=self.appctx))
return self._splits.setdefault(tuple(fields), splits)
def __init__(self,
F,
butcher_tableau,
t,
dt,
u0,
bcs=None,
solver_parameters=None,
update_solver_parameters=None,
splitting=AI,
nullspace=None,
appctx=None):
self.u0 = u0
self.t = t
self.dt = dt
self.num_fields = len(u0.function_space())
self.num_stages = len(butcher_tableau.b)
self.butcher_tableau = butcher_tableau
Fbig, update_stuff, UU, bigBCs, gblah, nsp = getFormStage(
F, butcher_tableau, u0, t, dt, bcs, splitting=splitting)
self.UU = UU
self.bigBCs = bigBCs
self.bcdat = gblah
self.update_stuff = update_stuff
self.prob = NonlinearVariationalProblem(Fbig, UU, bigBCs)
appctx_irksome = {
"F": F,
"butcher_tableau": butcher_tableau,
"t": t,
"dt": dt,
"u0": u0,
"bcs": bcs,
"nullspace": nullspace
}
if appctx is None:
appctx = appctx_irksome
else:
appctx = {**appctx, **appctx_irksome}
self.solver = NonlinearVariationalSolver(
self.prob,
appctx=appctx,
nullspace=nsp,
solver_parameters=solver_parameters)
unew, Fupdate, update_bcs, update_bcs_gblah = self.update_stuff
self.update_problem = NonlinearVariationalProblem(
Fupdate, unew, update_bcs)
self.update_solver = NonlinearVariationalSolver(
self.update_problem, solver_parameters=update_solver_parameters)
self._update = self._update_general
def split(self, fields):
from ufl import as_vector, replace
from firedrake import NonlinearVariationalProblem as NLVP, FunctionSpace
splits = self._splits.get(tuple(fields))
if splits is not None:
return splits
splits = []
problem = self._problem
splitter = ExtractSubBlock()
for field in fields:
try:
if len(field) > 1:
raise NotImplementedError("Can't split into subblock")
except TypeError:
# Just a single field, we can handle that
pass
F = splitter.split(problem.F, argument_indices=(field, ))
J = splitter.split(problem.J, argument_indices=(field, field))
us = problem.u.split()
subu = us[field]
vec = []
for i, u in enumerate(us):
for idx in numpy.ndindex(u.ufl_shape):
vec.append(u[idx])
u = as_vector(vec)
F = replace(F, {problem.u: u})
J = replace(J, {problem.u: u})
if problem.Jp is not None:
Jp = splitter.split(problem.Jp, argument_indices=(field, field))
Jp = replace(Jp, {problem.u: u})
else:
Jp = None
bcs = []
for bc in problem.bcs:
if bc.function_space().index == field:
V = FunctionSpace(subu.ufl_domain(), subu.ufl_element())
bcs.append(type(bc)(V,
bc.function_arg,
bc.sub_domain,
method=bc.method))
new_problem = NLVP(F, subu, bcs=bcs, J=J, Jp=None,
form_compiler_parameters=problem.form_compiler_parameters)
new_problem._constant_jacobian = problem._constant_jacobian
splits.append(type(self)(new_problem, mat_type=self.mat_type, pmat_type=self.pmat_type,
appctx=self.appctx))
return self._splits.setdefault(tuple(fields), splits)
def setup_solver(self, up_init=None):
""" Setup the solvers
"""
self.up0 = Function(self.W)
if up_init is not None:
chk_in = checkpointing.HDF5File(up_init, file_mode='r')
chk_in.read(self.up0, "/up")
chk_in.close()
self.u0, self.p0 = split(self.up0)
self.up = Function(self.W)
if up_init is not None:
chk_in = checkpointing.HDF5File(up_init, file_mode='r')
chk_in.read(self.up, "/up")
chk_in.close()
self.u1, self.p1 = split(self.up)
self.up.sub(0).rename("velocity")
self.up.sub(1).rename("pressure")
v, q = TestFunctions(self.W)
h = CellVolume(self.mesh)
u_norm = sqrt(dot(self.u0, self.u0))
if self.has_nullspace:
nullspace = MixedVectorSpaceBasis(
self.W,
[self.W.sub(0), VectorSpaceBasis(constant=True)])
else:
nullspace = None
tau = ((2.0 / self.dt)**2 + (2.0 * u_norm / h)**2 +
(4.0 * self.nu / h**2)**2)**(-0.5)
# temporal discretization
F = (1.0 / self.dt) * inner(self.u1 - self.u0, v) * dx
# weak form
F += (+inner(dot(self.u0, nabla_grad(self.u1)), v) * dx +
self.nu * inner(grad(self.u1), grad(v)) * dx -
(1.0 / self.rho) * self.p1 * div(v) * dx +
div(self.u1) * q * dx - inner(self.forcing, v) * dx)
# residual form
R = (+(1.0 / self.dt) * (self.u1 - self.u0) +
dot(self.u0, nabla_grad(self.u1)) - self.nu * div(grad(self.u1)) +
(1.0 / self.rho) * grad(self.p1) - self.forcing)
# GLS
F += tau * inner(
+dot(self.u0, nabla_grad(v)) - self.nu * div(grad(v)) +
(1.0 / self.rho) * grad(q), R) * dx
self.problem = NonlinearVariationalProblem(F, self.up, self.bcs)
self.solver = NonlinearVariationalSolver(
self.problem,
nullspace=nullspace,
solver_parameters=self.solver_parameters)
def solve(self, file_path='ttip_result/solution.pvd'):
"""
Setup and solve the nonlinear problem.
Save value to file given.
Any additional keyword arguments are passed to the iteration method.
Args:
file_path (string, optional):
The path to save the pvd file to.
vtk files will be generated in the same directory as the pvd.
It is recommended that this is a separate drectory per run.
Defaults to 'TTiP_result/solution.pvd'.
"""
F = self.problem.a - self.problem.L
steady_state = self.is_steady_state()
if isinstance(self.problem, BoundaryMixin):
var_prob = NonlinearVariationalProblem(F,
self.u,
bcs=self.problem.bcs)
else:
var_prob = NonlinearVariationalProblem(F, self.u)
solver = NonlinearVariationalSolver(problem=var_prob,
solver_parameters=self.params)
outfile = File(file_path)
outfile.write(self.u, target_degree=1, target_continuity=H1)
if steady_state:
solver.solve()
outfile.write(self.u, target_degree=1, target_continuity=H1)
else:
self.problem.T_.assign(self.u)
last_perc = 0
for i in range(self.problem.steps):
solver.solve()
perc = int(100 * (i + 1) / self.problem.steps)
if perc > last_perc:
print(f'{perc}%')
last_perc = perc
self.problem.T_.assign(self.u)
outfile.write(self.u, target_degree=1, target_continuity=H1)
def _ad_problem_clone(self, problem, dependencies):
"""Replaces every coefficient in the residual and jacobian with a deepcopy to return
a clone of the original NonlinearVariationalProblem instance. We'll be modifying the
numerical values of the coefficients in the residual and jacobian, so in order not to
affect the user-defined self._ad_problem.F, self._ad_problem.J and self._ad_problem.u
expressions, we'll instead create clones of them.
"""
from firedrake import NonlinearVariationalProblem
F_replace_map = {}
J_replace_map = {}
F_coefficients = problem.F.coefficients()
J_coefficients = problem.J.coefficients()
_ad_count_map = {}
for block_variable in dependencies:
coeff = block_variable.output
if coeff in F_coefficients and coeff not in F_replace_map:
if isinstance(coeff, Constant):
F_replace_map[coeff] = copy.deepcopy(coeff)
else:
F_replace_map[coeff] = coeff.copy(deepcopy=True)
_ad_count_map[F_replace_map[coeff]] = coeff.count()
if coeff in J_coefficients and coeff not in J_replace_map:
if coeff in F_replace_map:
J_replace_map[coeff] = F_replace_map[coeff]
elif isinstance(coeff, Constant):
J_replace_map[coeff] = copy.deepcopy(coeff)
else:
J_replace_map[coeff] = coeff.copy()
_ad_count_map[J_replace_map[coeff]] = coeff.count()
nlvp = NonlinearVariationalProblem(replace(problem.F, F_replace_map),
F_replace_map[problem.u],
bcs=problem.bcs,
J=replace(problem.J, J_replace_map))
nlvp._ad_count_map_update(_ad_count_map)
return nlvp
def __init__(self, mesh, conditions, timestepping, params, output, solver_params):
super().__init__(mesh, conditions, timestepping, params, output, solver_params)
self.w0 = Function(self.W3)
self.w1 = Function(self.W3)
u0, s0, h0, a0 = self.w0.split()
p, q, r, m = TestFunctions(self.W3)
self.initial_condition((u0, conditions.ic['u']), (s0, conditions.ic['s']),
(a0, conditions.ic['a']), (h0, conditions.ic['h']))
self.w1.assign(self.w0)
u1, s1, h1, a1 = split(self.w1)
u0, s0, h0, a0 = split(self.w0)
theta = conditions.theta
uh = (1-theta) * u0 + theta * u1
sh = (1-theta) * s0 + theta * s1
hh = (1-theta) * h0 + theta * h1
ah = (1-theta) * a0 + theta * a1
ep_dot = self.strain(grad(uh))
zeta = self.zeta(hh, ah, self.delta(uh))
rheology = params.e ** 2 * sh + Identity(2) * 0.5 * ((1 - params.e ** 2) * tr(sh) + self.Ice_Strength(hh, ah))
eqn = self.momentum_equation(hh, u1, u0, p, sh, params.rho, uh, conditions.ocean_curr, params.rho_a,
params.C_a, params.rho_w, params.C_w, conditions.geo_wind, params.cor, self.timestep, ind=self.ind)
eqn += self.transport_equation(uh, hh, ah, h1, h0, a1, a0, r, m, self.n, self.timestep)
eqn += inner(self.ind * (s1 - s0) + 0.5 * self.timestep * rheology / params.T, q) * dx
eqn -= inner(q * zeta * self.timestep / params.T, ep_dot) * dx
if conditions.stabilised['state']:
alpha = conditions.stabilised['alpha']
eqn += self.stabilisation_term(alpha=alpha, zeta=avg(zeta), mesh=mesh, v=uh, test=p)
bcs = DirichletBC(self.W3.sub(0), conditions.bc['u'], "on_boundary")
uprob = NonlinearVariationalProblem(eqn, self.w1, bcs)
self.usolver = NonlinearVariationalSolver(uprob, solver_parameters=solver_params.bt_params)
self.u1, self.s0, self.h1, self.a1 = self.w1.split()
def __init__(self, mesh, conditions, timestepping, params, output, solver_params):
super().__init__(mesh, conditions, timestepping, params, output, solver_params)
self.w0 = Function(self.W2)
self.w1 = Function(self.W2)
u0, h0, a0 = self.w0.split()
p, q, r = TestFunctions(self.W2)
self.initial_condition((u0, conditions.ic['u']), (h0, conditions.ic['h']),
(a0, conditions.ic['a']))
self.w1.assign(self.w0)
u1, h1, a1 = split(self.w1)
u0, h0, a0 = split(self.w0)
theta = conditions.theta
uh = (1-theta) * u0 + theta * u1
ah = (1-theta) * a0 + theta * a1
hh = (1-theta) * h0 + theta * h1
ep_dot = self.strain(grad(uh))
zeta = self.zeta(hh, ah, self.delta(uh))
eta = zeta * params.e ** (-2)
sigma = 2 * eta * ep_dot + (zeta - eta) * tr(ep_dot) * Identity(2) - self.Ice_Strength(hh, ah) * 0.5 * Identity(
2)
eqn = self.momentum_equation(hh, u1, u0, p, sigma, params.rho, uh, conditions.ocean_curr, params.rho_a,
params.C_a, params.rho_w, params.C_w, conditions.geo_wind, params.cor, self.timestep)
eqn += self.transport_equation(uh, hh, ah, h1, h0, a1, a0, q, r, self.n, self.timestep)
if conditions.stabilised['state']:
alpha = conditions.stabilised['alpha']
eqn += self.stabilisation_term(alpha=alpha, zeta=avg(zeta), mesh=mesh, v=uh, test=p)
bcs = DirichletBC(self.W2.sub(0), conditions.bc['u'], "on_boundary")
uprob = NonlinearVariationalProblem(eqn, self.w1, bcs)
self.usolver = NonlinearVariationalSolver(uprob, solver_parameters=solver_params.bt_params)
self.u1, self.h1, self.a1 = self.w1.split()
def run(steady=False):
"""
solve CdT/dt = S + div(k*grad(T))
=> C*v*(dT/dt)/k*dx - S*v/k*dx + grad(v)*grad(T)*dx = v*dot(grad(T), n)*ds
"""
steps = 250
dt = 1e-10
timescale = (0, steps * dt)
if steady:
print('Running steady state.')
else:
print(f'Running with time step {dt:.2g}s on time interval: '
f'{timescale[0]:.2g}s - {timescale[1]:.2g}s')
dt_invc = Constant(1 / dt)
extent = [40e-6, 40e-6, 40e-6]
mesh = BoxMesh(20, 20, 20, *extent)
V = FunctionSpace(mesh, 'CG', 1)
print(V.dim())
T = Function(V) # temperature at time i+1 (electron for now)
T_ = Function(V) # temperature at time i
v = TestFunction(V) # test function
S = create_S(mesh, V, extent)
C = create_heat_capacity(mesh, V, extent)
k = create_conductivity(mesh, V, T)
set_initial_value(mesh, T_, extent)
# Mass matrix section
M = C * T * dt_invc * v * dx
M_ = C * T_ * dt_invc * v * dx
# Stiffness matrix section
A = k * dot(grad(T), grad(v)) * dx
# function section
f = S * v * dx
# boundaries
bcs, R, b = create_dirichlet_bounds(mesh,
V,
T,
v,
k,
g=100,
boundary=[1, 2, 3, 4, 5, 6])
# bcs += create_dirichlet_bounds(mesh, V, T, v, k, 500, [6])[0]
# bcs, R, b = create_robin_bounds(mesh, T, v, k, 1e8/(100), 1e8)
if steady:
steps = 1
a = A + R
L = f + b
else:
a = M + A + R
L = M_ + f + b
prob = NonlinearVariationalProblem(a - L, T, bcs=bcs)
solver = NonlinearVariationalSolver(prob, solver_parameters=SOLVE_PARAMS)
T.assign(T_)
timestamp = datetime.now().strftime("%d-%b-%Y-%H-%M-%S")
outfile = File(f'{timestamp}/first_output.pvd')
outfile.write(T_, target_degree=1, target_continuity=H1)
last_perc = 0
for i in range(steps):
solver.solve()
perc = int(100 * (i + 1) / steps)
if perc > last_perc:
print(f'{perc}%')
last_perc = perc
T_.assign(T)
outfile.write(T_, target_degree=1, target_continuity=H1)
def split(self, fields):
from firedrake import replace, as_vector, split
from firedrake import NonlinearVariationalProblem as NLVP
fields = tuple(tuple(f) for f in fields)
splits = self._splits.get(tuple(fields))
if splits is not None:
return splits
splits = []
problem = self._problem
splitter = ExtractSubBlock()
for field in fields:
F = splitter.split(problem.F, argument_indices=(field, ))
J = splitter.split(problem.J, argument_indices=(field, field))
us = problem.u.split()
V = F.arguments()[0].function_space()
# Exposition:
# We are going to make a new solution Function on the sub
# mixed space defined by the relevant fields.
# But the form may refer to the rest of the solution
# anyway.
# So we pull it apart and will make a new function on the
# subspace that shares data.
pieces = [us[i].dat for i in field]
if len(pieces) == 1:
val, = pieces
subu = function.Function(V, val=val)
subsplit = (subu, )
else:
val = op2.MixedDat(pieces)
subu = function.Function(V, val=val)
# Split it apart to shove in the form.
subsplit = split(subu)
# Permutation from field indexing to indexing of pieces
field_renumbering = dict([f, i] for i, f in enumerate(field))
vec = []
for i, u in enumerate(us):
if i in field:
# If this is a field we're keeping, get it from
# the new function. Otherwise just point to the
# old data.
u = subsplit[field_renumbering[i]]
if u.ufl_shape == ():
vec.append(u)
else:
for idx in numpy.ndindex(u.ufl_shape):
vec.append(u[idx])
# So now we have a new representation for the solution
# vector in the old problem. For the fields we're going
# to solve for, it points to a new Function (which wraps
# the original pieces). For the rest, it points to the
# pieces from the original Function.
# IOW, we've reinterpreted our original mixed solution
# function as being made up of some spaces we're still
# solving for, and some spaces that have just become
# coefficients in the new form.
u = as_vector(vec)
F = replace(F, {problem.u: u})
J = replace(J, {problem.u: u})
if problem.Jp is not None:
Jp = splitter.split(problem.Jp, argument_indices=(field, field))
Jp = replace(Jp, {problem.u: u})
else:
Jp = None
bcs = []
for bc in problem.bcs:
Vbc = bc.function_space()
if Vbc.parent is not None and isinstance(Vbc.parent.ufl_element(), VectorElement):
index = Vbc.parent.index
else:
index = Vbc.index
cmpt = Vbc.component
# TODO: need to test this logic
if index in field:
if len(field) == 1:
W = V
else:
W = V.sub(field_renumbering[index])
if cmpt is not None:
W = W.sub(cmpt)
bcs.append(type(bc)(W,
bc.function_arg,
bc.sub_domain,
method=bc.method))
new_problem = NLVP(F, subu, bcs=bcs, J=J, Jp=Jp,
form_compiler_parameters=problem.form_compiler_parameters)
new_problem._constant_jacobian = problem._constant_jacobian
splits.append(type(self)(new_problem, mat_type=self.mat_type, pmat_type=self.pmat_type,
appctx=self.appctx))
return self._splits.setdefault(tuple(fields), splits)
def build_initial_conditions(prognostic_variables, simulation_parameters):
"""
Initialises the prognostic variables based on the
initial condition string.
:arg prognostic_variables: a PrognosticVariables object.
:arg simulation_parameters: a dictionary containing the simulation parameters.
"""
mesh = simulation_parameters['mesh'][-1]
ic = simulation_parameters['ic'][-1]
alphasq = simulation_parameters['alphasq'][-1]
c0 = simulation_parameters['c0'][-1]
gamma = simulation_parameters['gamma'][-1]
x, = SpatialCoordinate(mesh)
Ld = simulation_parameters['Ld'][-1]
deltax = Ld / simulation_parameters['resolution'][-1]
w = simulation_parameters['peak_width'][-1]
epsilon = 1
ic_dict = {
'two_peaks':
(0.2 * 2 /
(exp(x - 403. / 15. * 40. / Ld) + exp(-x + 403. / 15. * 40. / Ld)) +
0.5 * 2 /
(exp(x - 203. / 15. * 40. / Ld) + exp(-x + 203. / 15. * 40. / Ld))),
'gaussian':
0.5 * exp(-((x - 10.) / 2.)**2),
'gaussian_narrow':
0.5 * exp(-((x - 10.) / 1.)**2),
'gaussian_wide':
0.5 * exp(-((x - 10.) / 3.)**2),
'peakon':
conditional(x < Ld / 2., exp((x - Ld / 2) / sqrt(alphasq)),
exp(-(x - Ld / 2) / sqrt(alphasq))),
'one_peak':
0.5 * 2 /
(exp(x - 203. / 15. * 40. / Ld) + exp(-x + 203. / 15. * 40. / Ld)),
'proper_peak':
0.5 * 2 / (exp(x - Ld / 4) + exp(-x + Ld / 4)),
'new_peak':
0.5 * 2 / (exp((x - Ld / 4) / w) + exp((-x + Ld / 4) / w)),
'flat':
Constant(2 * pi**2 / (9 * 40**2)),
'fast_flat':
Constant(0.1),
'coshes':
Constant(2000) * cosh((2000**0.5 / 2) * (x - 0.75))**(-2) +
Constant(1000) * cosh(1000**0.5 / 2 * (x - 0.25))**(-2),
'd_peakon':
exp(-sqrt((x - Ld / 2)**2 + epsilon * deltax**2) / sqrt(alphasq)),
'zero':
Constant(0.0),
'two_peakons':
conditional(
x < Ld / 4,
exp((x - Ld / 4) / sqrt(alphasq)) -
exp(-(x + Ld / 4) / sqrt(alphasq)),
conditional(
x < 3 * Ld / 4,
exp(-(x - Ld / 4) / sqrt(alphasq)) - exp(
(x - 3 * Ld / 4) / sqrt(alphasq)),
exp((x - 5 * Ld / 4) / sqrt(alphasq)) -
exp(-(x - 3 * Ld / 4) / sqrt(alphasq)))),
'twin_peakons':
conditional(
x < Ld / 4,
exp((x - Ld / 4) / sqrt(alphasq)) + 0.5 * exp(
(x - Ld / 2) / sqrt(alphasq)),
conditional(
x < Ld / 2,
exp(-(x - Ld / 4) / sqrt(alphasq)) + 0.5 * exp(
(x - Ld / 2) / sqrt(alphasq)),
conditional(
x < 3 * Ld / 4,
exp(-(x - Ld / 4) / sqrt(alphasq)) +
0.5 * exp(-(x - Ld / 2) / sqrt(alphasq)),
exp((x - 5 * Ld / 4) / sqrt(alphasq)) +
0.5 * exp(-(x - Ld / 2) / sqrt(alphasq))))),
'periodic_peakon': (conditional(
x < Ld / 2, 0.5 / (1 - exp(-Ld / sqrt(alphasq))) *
(exp((x - Ld / 2) / sqrt(alphasq)) +
exp(-Ld / sqrt(alphasq)) * exp(-(x - Ld / 2) / sqrt(alphasq))),
0.5 / (1 - exp(-Ld / sqrt(alphasq))) *
(exp(-(x - Ld / 2) / sqrt(alphasq)) +
exp(-Ld / sqrt(alphasq)) * exp((x - Ld / 2) / sqrt(alphasq))))),
'cos_bell':
conditional(x < Ld / 4, (cos(pi * (x - Ld / 8) / (2 * Ld / 8)))**2,
0.0),
'antisymmetric':
1 / (exp((x - Ld / 4) / Ld) + exp((-x + Ld / 4) / Ld)) - 1 / (exp(
(Ld - x - Ld / 4) / Ld) + exp((Ld + x + Ld / 4) / Ld))
}
ic_expr = ic_dict[ic]
if prognostic_variables.scheme in ['upwind', 'LASCH']:
VCG5 = FunctionSpace(mesh, "CG", 5)
smooth_condition = Function(VCG5).interpolate(ic_expr)
prognostic_variables.u.project(as_vector([smooth_condition]))
# need to find initial m by solving helmholtz problem
CG1 = FunctionSpace(mesh, "CG", 1)
u0 = prognostic_variables.u
p = TestFunction(CG1)
m_CG = Function(CG1)
ones = Function(prognostic_variables.Vu).project(
as_vector([Constant(1.)]))
Lm = (p * m_CG - p * dot(ones, u0) -
alphasq * p.dx(0) * dot(ones, u0.dx(0))) * dx
mprob0 = NonlinearVariationalProblem(Lm, m_CG)
msolver0 = NonlinearVariationalSolver(mprob0,
solver_parameters={
'ksp_type': 'preonly',
'pc_type': 'lu'
})
msolver0.solve()
prognostic_variables.m.interpolate(m_CG)
if prognostic_variables.scheme == 'LASCH':
prognostic_variables.Eu.assign(prognostic_variables.u)
prognostic_variables.Em.assign(prognostic_variables.m)
elif prognostic_variables.scheme in ('conforming', 'hydrodynamic', 'test',
'LASCH_hydrodynamic',
'LASCH_hydrodynamic_m',
'no_gradient'):
if ic == 'peakon':
Vu = prognostic_variables.Vu
# delta = Function(Vu)
# middle_index = int(len(delta.dat.data[:]) / 2)
# delta.dat.data[middle_index] = 1
# u0 = prognostic_variables.u
# phi = TestFunction(Vu)
#
# eqn = phi * u0 * dx + alphasq * phi.dx(0) * u0.dx(0) * dx - phi * delta * dx
# prob = NonlinearVariationalProblem(eqn, u0)
# solver = NonlinearVariationalSolver(prob)
# solver.solve()
# W = MixedFunctionSpace((Vu, Vu))
# psi, phi = TestFunctions(W)
# w = Function(W)
# u, F = w.split()
# u.interpolate(ic_expr)
# u, F = split(w)
#
# eqn = (psi * u * dx - psi * (0.5 * u * u + F) * dx
# + phi * F * dx + alphasq * phi.dx(0) * F.dx(0) * dx
# - phi * u * u * dx - 0.5 * alphasq * phi * u.dx(0) * u.dx(0) * dx)
#
# u, F = w.split()
#
# prob = NonlinearVariationalProblem(eqn, w)
# solver = NonlinearVariationalSolver(prob)
# solver.solve()
# prognostic_variables.u.assign(u)
prognostic_variables.u.project(ic_expr)
# prognostic_variables.u.interpolate(ic_expr)
else:
VCG5 = FunctionSpace(mesh, "CG", 5)
smooth_condition = Function(VCG5).interpolate(ic_expr)
prognostic_variables.u.project(smooth_condition)
if prognostic_variables.scheme in [
'LASCH_hydrodynamic', 'LASCH_hydrodynamic_m'
]:
prognostic_variables.Eu.assign(prognostic_variables.u)
else:
raise NotImplementedError('Other schemes not yet implemented.')
def __init__(self, prognostic_variables, simulation_parameters):
mesh = simulation_parameters['mesh'][-1]
self.scheme = simulation_parameters['scheme'][-1]
self.timestepping = simulation_parameters['timestepping'][-1]
alphasq = simulation_parameters['alphasq'][-1]
c0 = simulation_parameters['c0'][-1]
gamma = simulation_parameters['gamma'][-1]
Dt = Constant(simulation_parameters['dt'][-1])
self.solvers = []
if alphasq.values()[0] > 0.0 and gamma.values()[0] == 0.0:
self.setup = 'ch'
if self.scheme == 'upwind' and self.timestepping == 'ssprk3':
Vm = prognostic_variables.Vm
Vu = prognostic_variables.Vu
self.m = prognostic_variables.m
self.u = prognostic_variables.u
self.Xi = prognostic_variables.dXi
self.m0 = Function(Vm).assign(self.m)
# now make problem for the actual problem
psi = TestFunction(Vm)
self.m_trial = Function(Vm)
self.dm = Function(
Vm
) # introduce this as the advection operator for a single step
us = Dt * self.u + self.Xi
nhat = FacetNormal(mesh)
un = 0.5 * (dot(us, nhat) + abs(dot(us, nhat)))
ones = Function(Vu).project(as_vector([Constant(1.)]))
Lm = (psi * self.dm * dx -
psi.dx(0) * self.m_trial * dot(ones, us) * dx +
psi * self.m_trial * dot(ones, us.dx(0)) * dx +
jump(psi) * (un('+') * self.m_trial('+') -
un('-') * self.m_trial('-')) * dS)
mprob = NonlinearVariationalProblem(Lm, self.dm)
self.msolver = NonlinearVariationalSolver(mprob,
solver_parameters={
'ksp_type':
'preonly',
'pc_type':
'bjacobi',
'sub_pc_type':
'ilu'
})
phi = TestFunction(Vu)
Lu = (dot(phi, ones) * self.m * dx - dot(phi, self.u) * dx -
alphasq * dot(self.u.dx(0), phi.dx(0)) * dx)
uprob = NonlinearVariationalProblem(Lu, self.u)
self.usolver = NonlinearVariationalSolver(uprob,
solver_parameters={
'ksp_type':
'preonly',
'pc_type': 'lu'
})
elif self.scheme == 'hydrodynamic' and self.timestepping == 'midpoint':
Vu = prognostic_variables.Vu
self.u = prognostic_variables.u
W = MixedFunctionSpace((Vu, ) * 3)
psi, phi, zeta = TestFunctions(W)
w1 = Function(W)
self.u1, dFh, dGh = split(w1)
uh = (self.u1 + self.u) / 2
dXi = prognostic_variables.dXi
dXi_x = prognostic_variables.dXi_x
dXi_xx = prognostic_variables.dXi_xx
dvh = Dt * uh + dXi
Lu = (psi * (self.u1 - self.u) * dx +
psi * uh.dx(0) * dvh * dx - psi.dx(0) * dFh * dx +
psi * dGh * dx + phi * dFh * dx +
alphasq * phi.dx(0) * dFh.dx(0) * dx -
phi * uh * uh * Dt * dx -
0.5 * alphasq * phi * uh.dx(0) * uh.dx(0) * Dt * dx +
zeta * dGh * dx + alphasq * zeta.dx(0) * dGh.dx(0) * dx -
2 * zeta * uh * dXi_x * dx -
alphasq * zeta * uh.dx(0) * dXi_xx * dx)
self.u1, dFh, dGh = w1.split()
uprob = NonlinearVariationalProblem(Lu, w1)
self.usolver = NonlinearVariationalSolver(uprob,
solver_parameters={
'mat_type':
'aij',
'ksp_type':
'preonly',
'pc_type': 'lu'
})
elif self.scheme == 'no_gradient' and self.timestepping == 'midpoint':
# a version of the hydrodynamic form but without exploiting the gradient
Vu = prognostic_variables.Vu
self.u = prognostic_variables.u
W = MixedFunctionSpace((Vu, ) * 3)
psi, phi, zeta = TestFunctions(W)
w1 = Function(W)
self.u1, dFh, dGh = split(w1)
uh = (self.u1 + self.u) / 2
dXi = prognostic_variables.dXi
dXi_x = prognostic_variables.dXi_x
dXi_xx = prognostic_variables.dXi_xx
dvh = Dt * uh + dXi
Lu = (psi * (self.u1 - self.u) * dx +
psi * uh.dx(0) * dvh * dx + psi * dFh.dx(0) * dx +
psi * dGh * dx + phi * dFh * dx +
alphasq * phi.dx(0) * dFh.dx(0) * dx -
phi * uh * uh * Dt * dx -
0.5 * alphasq * phi * uh.dx(0) * uh.dx(0) * Dt * dx +
zeta * dGh * dx + alphasq * zeta.dx(0) * dGh.dx(0) * dx -
2 * zeta * uh * dXi_x * dx -
alphasq * zeta * uh.dx(0) * dXi_xx * dx)
self.u1, dFh, dGh = w1.split()
uprob = NonlinearVariationalProblem(Lu, w1)
self.usolver = NonlinearVariationalSolver(uprob,
solver_parameters={
'mat_type':
'aij',
'ksp_type':
'preonly',
'pc_type': 'lu'
})
elif self.scheme == 'test' and self.timestepping == 'midpoint':
self.u = prognostic_variables.u
Vu = prognostic_variables.Vu
psi = TestFunction(Vu)
self.u1 = Function(Vu)
uh = (self.u1 + self.u) / 2
dvh = Dt * uh + prognostic_variables.dXi
eqn = (psi * (self.u1 - self.u) * dx -
psi * uh * dvh.dx(0) * dx)
prob = NonlinearVariationalProblem(eqn, self.u1)
self.usolver = NonlinearVariationalSolver(prob,
solver_parameters={
'mat_type':
'aij',
'ksp_type':
'preonly',
'pc_type': 'lu'
})
else:
raise ValueError(
'Scheme %s and timestepping %s either not compatible or not recognised.'
% (self.scheme, self.timestepping))
elif alphasq.values()[0] == 0.0 and gamma.values()[0] > 0.0:
self.setup = 'kdv'
if self.scheme == 'upwind' and self.timestepping == 'ssprk3':
raise NotImplementedError(
'Scheme %s and timestepping %s not yet implemented.' %
(self.scheme, self.timestepping))
elif self.scheme == 'upwind' and self.timestepping == 'midpoint':
raise NotImplementedError(
'Scheme %s and timestepping %s not yet implemented.' %
(self.scheme, self.timestepping))
elif self.scheme == 'hydrodynamic' and self.timestepping == 'midpoint':
raise NotImplementedError(
'Scheme %s and timestepping %s not yet implemented.' %
(self.scheme, self.timestepping))
else:
raise ValueError(
'Scheme %s and timestepping %s either not compatible or not recognised.'
% (self.scheme, self.timestepping))
else:
raise NotImplementedError(
'Schemes for your values of alpha squared %.3f and gamma %.3f are not yet implemented.'
% (alphasq, gamma))
def compressible_hydrostatic_balance(state,
theta0,
rho0,
exner0=None,
top=False,
exner_boundary=Constant(1.0),
mr_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 exner_boundary: a field or expression to use as boundary data for exner
on the top or bottom as specified.
:arg mr_t: the initial total water mixing ratio field.
"""
# Calculate hydrostatic Pi
VDG = state.spaces("DG")
Vu = state.spaces("HDiv")
Vv = FunctionSpace(state.mesh, Vu.ufl_element()._elements[-1])
W = MixedFunctionSpace((Vv, VDG))
v, exner = TrialFunctions(W)
dv, dexner = TestFunctions(W)
n = FacetNormal(state.mesh)
cp = state.parameters.cp
# add effect of density of water upon theta
theta = theta0
if mr_t is not None:
theta = theta0 / (1 + mr_t)
alhs = ((cp * inner(v, dv) - cp * div(dv * theta) * exner) * dx +
dexner * div(theta * v) * dx)
if top:
bmeasure = ds_t
bstring = "bottom"
else:
bmeasure = ds_b
bstring = "top"
arhs = -cp * inner(dv, n) * theta * exner_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), zero(), bstring)]
w = Function(W)
exner_problem = 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'
}
}
exner_solver = LinearVariationalSolver(exner_problem,
solver_parameters=params,
options_prefix="exner_solver")
exner_solver.solve()
v, exner = w.split()
if exner0 is not None:
exner0.assign(exner)
if solve_for_rho:
w1 = Function(W)
v, rho = w1.split()
rho.interpolate(thermodynamics.rho(state.parameters, theta0, exner))
v, rho = split(w1)
dv, dexner = TestFunctions(W)
exner = thermodynamics.exner_pressure(state.parameters, rho, theta0)
F = ((cp * inner(v, dv) - cp * div(dv * theta) * exner) * dx +
dexner * div(theta0 * v) * dx +
cp * inner(dv, n) * theta * exner_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, exner))
def __init__(self, prognostic_variables, simulation_parameters):
mesh = simulation_parameters['mesh'][-1]
x, = SpatialCoordinate(mesh)
Ld = simulation_parameters['Ld'][-1]
self.scheme = simulation_parameters['scheme'][-1]
self.dt = simulation_parameters['dt'][-1]
self.num_Xis = simulation_parameters['num_Xis'][-1]
self.Xi_family = simulation_parameters['Xi_family'][-1]
self.dXi = prognostic_variables.dXi
self.dWs = [Constant(0.0) for dw in range(self.num_Xis)]
self.dW_nums = prognostic_variables.dW_nums
self.Xi_functions = []
self.nXi_updates = simulation_parameters['nXi_updates'][-1]
self.smooth_t = simulation_parameters['smooth_t'][-1]
self.fixed_dW = simulation_parameters['fixed_dW'][-1]
if self.smooth_t is not None and self.nXi_updates > 1:
raise ValueError('Prescribing forcing and including multiple Xi updates are not compatible.')
if self.smooth_t is not None or self.fixed_dW is not None:
print('WARNING: Remember to change sigma to sigma * sqrt(dt) with the prescribed forcing option or the fixed_dW option.')
seed = simulation_parameters['seed'][-1]
np.random.seed(seed)
# make sure sigma is a Constant
if self.num_Xis != 0:
if isinstance(simulation_parameters['sigma'][-1], Constant):
self.sigma = simulation_parameters['sigma'][-1]
else:
self.sigma = Constant(simulation_parameters['sigma'][-1])
else:
self.sigma = Constant(0.0)
self.pure_xi_list = prognostic_variables.pure_xi_list
self.pure_xi_x_list = prognostic_variables.pure_xi_x_list
self.pure_xi_xx_list = prognostic_variables.pure_xi_xx_list
self.pure_xi_xxx_list = prognostic_variables.pure_xi_xxx_list
self.pure_xi_xxxx_list = prognostic_variables.pure_xi_xxxx_list
for xi in range(self.num_Xis):
self.pure_xi_list.append(Function(self.dXi.function_space()))
self.pure_xi_x_list.append(Function(self.dXi.function_space()))
self.pure_xi_xx_list.append(Function(self.dXi.function_space()))
self.pure_xi_xxx_list.append(Function(self.dXi.function_space()))
self.pure_xi_xxxx_list.append(Function(self.dXi.function_space()))
if self.Xi_family == 'sines':
for n in range(self.num_Xis):
if (n+1) % 2 == 1:
self.Xi_functions.append(self.sigma * sin(2*(n+1)*pi*x/Ld))
else:
self.Xi_functions.append(self.sigma * cos(2*(n+1)*pi*x/Ld))
elif self.Xi_family == 'double_sines':
for n in range(self.num_Xis):
if (n+1) % 2 == 1:
self.Xi_functions.append(self.sigma * sin(4*(n+1)*pi*x/Ld))
else:
self.Xi_functions.append(self.sigma * cos(4*(n+1)*pi*x/Ld))
elif self.Xi_family == 'high_freq_sines':
for n in range(self.num_Xis):
if (n+1) % 2 == 1:
self.Xi_functions.append(self.sigma * sin((2*(n+1)+10)*pi*x/Ld))
else:
self.Xi_functions.append(self.sigma * cos((2*(n+1)+10)*pi*x/Ld))
elif self.Xi_family == 'gaussians':
for n in range(self.num_Xis):
self.Xi_functions.append(self.sigma * 0.5*self.num_Xis*exp(-((x-Ld*(n+1)/(self.num_Xis +1.0))/2.)**2))
elif self.Xi_family == 'quadratic':
if self.num_Xis > 1:
raise NotImplementedError('Quadratic Xi not yet implemented for more than one Xi')
else:
self.Xi_functions.append(32/(Ld*Ld)*conditional(x > Ld/4,
conditional(x > 3*Ld/8,
conditional(x > 5*Ld/8,
conditional(x < 3*Ld/4,
self.sigma * (x - 3*Ld/4)**2,
0.0),
(x-Ld/2)**2+Ld**2/32),
(x-Ld/4)**2),
0.0))
elif self.Xi_family == 'proper_peak':
if self.num_Xis > 1:
raise NotImplementedError('Quadratic Xi not yet implemented for more than one Xi')
else:
self.Xi_functions.append(self.sigma * 0.5*2/(exp(x-Ld/2)+exp(-x+Ld/2)))
elif self.Xi_family == 'constant':
if self.num_Xis > 1:
raise NotImplementedError('Constant Xi not yet implemented for more than one Xi')
else:
self.Xi_functions.append(self.sigma * (sin(0*pi*x/Ld)+1))
else:
raise NotImplementedError('Xi_family %s not implemented' % self.Xi_family)
# make lists of functions for xi_x, xi_xx and xi_xxx
if self.scheme in ['hydrodynamic', 'LASCH_hydrodynamic']:
self.dXi_x = prognostic_variables.dXi_x
self.dXi_xx = prognostic_variables.dXi_xx
self.Xi_x_functions = []
self.Xi_xx_functions = []
for Xi_expr in self.Xi_functions:
Xi_x_function = Function(self.dXi_x.function_space())
Xi_xx_function = Function(self.dXi_xx.function_space())
phi_x = TestFunction(self.dXi_x.function_space())
phi_xx = TestFunction(self.dXi_xx.function_space())
Xi_x_eqn = phi_x * Xi_x_function * dx + phi_x.dx(0) * Xi_expr * dx
Xi_xx_eqn = phi_xx * Xi_xx_function * dx + phi_xx.dx(0) * Xi_x_function * dx
Xi_x_problem = NonlinearVariationalProblem(Xi_x_eqn, Xi_x_function)
Xi_xx_problem = NonlinearVariationalProblem(Xi_xx_eqn, Xi_xx_function)
Xi_x_solver = NonlinearVariationalSolver(Xi_x_problem)
Xi_xx_solver = NonlinearVariationalSolver(Xi_xx_problem)
# for some reason these solvers don't work for constant Xi functions
# so just manually make the derivatives be zero
if self.Xi_family == 'constant':
Xi_x_function.interpolate(0.0*x)
Xi_xx_function.interpolate(0.0*x)
else:
Xi_x_solver.solve()
Xi_xx_solver.solve()
self.Xi_x_functions.append(Xi_x_function)
self.Xi_xx_functions.append(Xi_xx_function)
# now make a master xi
Xi_expr = 0.0*x
for dW, Xi_function, pure_xi, pure_xi_x, pure_xi_xx, pure_xi_xxx, pure_xi_xxxx in zip(self.dWs, self.Xi_functions, self.pure_xi_list, self.pure_xi_x_list, self.pure_xi_xx_list, self.pure_xi_xxx_list, self.pure_xi_xxxx_list):
Xi_expr += dW * Xi_function
if self.scheme in ['upwind', 'LASCH']:
pure_xi.interpolate(as_vector([Xi_function]))
pure_xi_x.project(as_vector([Xi_function.dx(0)]))
CG1 = FunctionSpace(mesh, "CG", 1)
psi = TestFunction(CG1)
xixx_scalar = Function(CG1)
xixx_eqn = psi * xixx_scalar * dx + psi.dx(0) * Xi_function.dx(0) * dx
prob = NonlinearVariationalProblem(xixx_eqn, xixx_scalar)
solver = NonlinearVariationalSolver(prob)
solver.solve()
pure_xi_xx.interpolate(as_vector([xixx_scalar]))
else:
pure_xi.interpolate(Xi_function)
# I guess we can't take the gradient of constants
if self.Xi_family != 'constant':
pure_xi_x.project(Xi_function.dx(0))
pure_xi_xx.project(pure_xi_x.dx(0))
pure_xi_xxx.project(pure_xi_xx.dx(0))
pure_xi_xxxx.project(pure_xi_xxx.dx(0))
if self.scheme in ['upwind', 'LASCH']:
self.dXi_interpolator = Interpolator(as_vector([Xi_expr]), self.dXi)
else:
self.dXi_interpolator = Interpolator(Xi_expr, self.dXi)
if self.scheme in ['hydrodynamic', 'LASCH_hydrodynamic']:
# initialise blank expressions
Xi_x_expr = 0.0*x
Xi_xx_expr = 0.0*x
# make full expressions by adding all dW * Xi_xs
for dW, Xi_x_function, Xi_xx_function in zip(self.dWs, self.Xi_x_functions, self.Xi_xx_functions):
Xi_x_expr += dW * Xi_x_function
Xi_xx_expr += dW * Xi_xx_function
self.dXi_x_interpolator = Interpolator(Xi_x_expr, self.dXi_x)
self.dXi_xx_interpolator = Interpolator(Xi_xx_expr, self.dXi_xx)
def moist_hydrostatic_balance(state, theta_e, water_t, pi_boundary=Constant(1.0)):
"""
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.
: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 pi_boundary: the value of pi on the lower boundary of the domain.
"""
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()
Vv = state.fields('u').function_space()
n = FacetNormal(state.mesh)
g = state.parameters.g
cp = state.parameters.cp
R_d = state.parameters.R_d
p_0 = state.parameters.p_0
VDG = state.spaces("DG")
if any(deg > 2 for deg in VDG.ufl_element().degree()):
state.logger.warning("default quadrature degree most likely not sufficient for this degree element")
quadrature_degree = (5, 5)
params = {'ksp_type': 'preonly',
'ksp_monitor_true_residual': True,
'ksp_converged_reason': True,
'snes_converged_reason': True,
'ksp_max_it': 100,
'mat_type': 'aij',
'pc_type': 'lu',
'pc_factor_mat_solver_type': 'mumps'}
theta0.interpolate(theta_e)
water_v0.interpolate(water_t)
Pi = Function(Vr)
epsilon = 0.9 # relaxation constant
# set up mixed space
Z = MixedFunctionSpace((Vt, Vt))
z = Function(Z)
gamma, phi = TestFunctions(Z)
theta_v, w_v = z.split()
# give first guesses for trial functions
theta_v.assign(theta0)
w_v.assign(water_v0)
theta_v, w_v = split(z)
# define variables
T = thermodynamics.T(state.parameters, theta_v, Pi, r_v=w_v)
p = thermodynamics.p(state.parameters, Pi)
w_sat = thermodynamics.r_sat(state.parameters, T, p)
dxp = dx(degree=(quadrature_degree))
# set up weak form of theta_e and w_sat equations
F = (-gamma * theta_e * dxp
+ gamma * thermodynamics.theta_e(state.parameters, T, p, w_v, water_t) * dxp
- phi * w_v * dxp
+ phi * w_sat * dxp)
problem = NonlinearVariationalProblem(F, z)
solver = NonlinearVariationalSolver(problem, solver_parameters=params)
theta_v, w_v = z.split()
Pi_h = Function(Vr).interpolate((p / p_0) ** (R_d / cp))
# solve for Pi with theta_v and w_v constant
# then solve for theta_v and w_v with Pi constant
for i in range(5):
compressible_hydrostatic_balance(state, theta0, rho0, pi0=Pi_h, water_t=water_t)
Pi.assign(Pi * (1 - epsilon) + epsilon * Pi_h)
solver.solve()
theta0.assign(theta0 * (1 - epsilon) + epsilon * theta_v)
water_v0.assign(water_v0 * (1 - epsilon) + epsilon * w_v)
# now begin on Newton solver, setup up new mixed space
Z = MixedFunctionSpace((Vt, Vt, Vr, Vv))
z = Function(Z)
gamma, phi, psi, w = TestFunctions(Z)
theta_v, w_v, pi, v = z.split()
# use previous values as first guesses for newton solver
theta_v.assign(theta0)
w_v.assign(water_v0)
pi.assign(Pi)
theta_v, w_v, pi, v = split(z)
# define variables
T = thermodynamics.T(state.parameters, theta_v, pi, r_v=w_v)
p = thermodynamics.p(state.parameters, pi)
w_sat = thermodynamics.r_sat(state.parameters, T, p)
F = (-gamma * theta_e * dxp
+ gamma * thermodynamics.theta_e(state.parameters, T, p, w_v, water_t) * dxp
- phi * w_v * dxp
+ phi * w_sat * dxp
+ cp * inner(v, w) * dxp
- cp * div(w * theta_v / (1.0 + water_t)) * pi * dxp
+ psi * div(theta_v * v / (1.0 + water_t)) * dxp
+ cp * inner(w, n) * pi_boundary * theta_v / (1.0 + water_t) * ds_b
+ g * inner(w, state.k) * dxp)
bcs = [DirichletBC(Z.sub(3), 0.0, "top")]
problem = NonlinearVariationalProblem(F, z, bcs=bcs)
solver = NonlinearVariationalSolver(problem, solver_parameters=params)
solver.solve()
theta_v, w_v, pi, v = z.split()
# assign final values
theta0.assign(theta_v)
water_v0.assign(w_v)
# find rho
compressible_hydrostatic_balance(state, theta0, rho0, water_t=water_t, solve_for_rho=True)
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)
def assemble(self, eqn, func, bcs, params):
uprob = NonlinearVariationalProblem(eqn, func, bcs)
self.usolver = NonlinearVariationalSolver(uprob, solver_parameters=params)
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.
: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
g = state.parameters.g
arhs -= g*inner(dv, state.k)*dx
bcs = [DirichletBC(W.sub(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_esteig': 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)
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)
rhosolver.solve()
v, rho_ = w1.split()
rho0.assign(rho_)
else:
rho0.interpolate(thermodynamics.rho(state.parameters, theta0, Pi))
phi = TestFunction(Vt)
rho_averaged = Function(Vt)
rho_recoverer = Recoverer(rho0,
rho_averaged,
VDG=FunctionSpace(mesh,
BrokenElement(Vt.ufl_element())),
boundary_method=physics_boundary_method)
rho_recoverer.project()
exner = thermodynamics.exner_pressure(state.parameters, rho_averaged, theta0)
p = thermodynamics.p(state.parameters, exner)
T = thermodynamics.T(state.parameters, theta0, exner, r_v=w_v)
w_sat = thermodynamics.r_sat(state.parameters, T, p)
w_functional = (phi * w_v * dxp - phi * w_sat * dxp)
w_problem = NonlinearVariationalProblem(w_functional, w_v)
w_solver = NonlinearVariationalSolver(w_problem)
w_solver.solve()
water_v0.assign(w_v)
water_c0.assign(water_t - water_v0)
state.set_reference_profiles([('rho', rho_b), ('theta', theta_b),
('vapour_mixing_ratio', water_vb)])
rho_opts = None
theta_opts = EmbeddedDGOptions()
u_transport = ImplicitMidpoint(state, "u")
transported_fields = [
SSPRK3(state, "rho", options=rho_opts),
def __init__(self, diagnostic_variables, prognostic_variables, outputting,
simulation_parameters):
self.diagnostic_variables = diagnostic_variables
self.prognostic_variables = prognostic_variables
self.outputting = outputting
self.simulation_parameters = simulation_parameters
Dt = Constant(simulation_parameters['dt'][-1])
Ld = simulation_parameters['Ld'][-1]
u = self.prognostic_variables.u
Xi = self.prognostic_variables.dXi
Vu = u.function_space()
vector_u = True if Vu.ufl_element() == VectorElement else False
ones = Function(
VectorFunctionSpace(self.prognostic_variables.mesh, "CG",
1)).project(as_vector([Constant(1.0)]))
self.to_update_constants = False
self.interpolators = []
self.projectors = []
self.solvers = []
mesh = u.function_space().mesh()
x, = SpatialCoordinate(mesh)
alphasq = simulation_parameters['alphasq'][-1]
periodic = simulation_parameters['periodic'][-1]
# do peakon data checks here
true_peakon_data = simulation_parameters['true_peakon_data'][-1]
if true_peakon_data is not None:
self.true_peakon_file = Dataset(
'results/' + true_peakon_data + '/data.nc', 'r')
# check length of file is correct
ndump = simulation_parameters['ndump'][-1]
tmax = simulation_parameters['tmax'][-1]
dt = simulation_parameters['dt'][-1]
if len(self.true_peakon_file['time'][:]) != int(tmax /
(ndump * dt)) + 1:
raise ValueError(
'If reading in true peakon data, the dump frequency must be the same as that used for the true peakon data.'
+
' Length of true peakon data as %i, but proposed length is %i'
% (len(self.true_peakon_file['time'][:]),
int(tmax / (ndump * dt)) + 1))
if self.true_peakon_file['p'][:].shape != (int(tmax /
(ndump * dt)) +
1, ):
raise ValueError(
'True peakon data shape %i must be the same shape as proposed data %i'
% ((int(tmax / (ndump * dt)) + 1, ),
self.true_peakon_file['p'][:].shape))
# do peakon data checks here
true_mean_peakon_data = simulation_parameters['true_mean_peakon_data'][
-1]
if true_mean_peakon_data is not None:
self.true_mean_peakon_file = Dataset(
'results/' + true_mean_peakon_data + '/data.nc', 'r')
# check length of file is correct
ndump = simulation_parameters['ndump'][-1]
tmax = simulation_parameters['tmax'][-1]
dt = simulation_parameters['dt'][-1]
if len(self.true_mean_peakon_file['time'][:]) != int(tmax /
(ndump * dt)):
raise ValueError(
'If reading in true peakon data, the dump frequency must be the same as that used for the true peakon data.'
)
if self.true_mean_peakon_file['p'][:].shape != (int(
tmax / (ndump * dt)), ):
raise ValueError(
'True peakon data must have same shape as proposed data!')
for key, value in self.diagnostic_variables.fields.items():
if key == 'uscalar':
uscalar = self.diagnostic_variables.fields['uscalar']
u_interpolator = Interpolator(dot(ones, u), uscalar)
self.interpolators.append(u_interpolator)
elif key == 'Euscalar':
Eu = self.prognostic_variables.Eu
Euscalar = self.diagnostic_variables.fields['Euscalar']
Eu_interpolator = Interpolator(dot(ones, Eu), Euscalar)
self.interpolators.append(Eu_interpolator)
elif key == 'Xiscalar':
Xi = self.prognostic_variables.dXi
Xiscalar = self.diagnostic_variables.fields['Xiscalar']
Xi_interpolator = Interpolator(dot(ones, Xi), Xiscalar)
self.interpolators.append(Xi_interpolator)
elif key == 'du':
if type(u.function_space().ufl_element()) == VectorElement:
u_to_project = self.diagnostic_variables.fields['uscalar']
else:
u_to_project = u
du = self.diagnostic_variables.fields['du']
du_projector = Projector(u_to_project.dx(0), du)
self.projectors.append(du_projector)
elif key == 'jump_du':
du = self.diagnostic_variables.fields['du']
jump_du = self.diagnostic_variables.fields['jump_du']
V = jump_du.function_space()
jtrial = TrialFunction(V)
psi = TestFunction(V)
Lj = psi('+') * abs(jump(du)) * dS
aj = psi('+') * jtrial('+') * dS
jprob = LinearVariationalProblem(aj, Lj, jump_du)
jsolver = LinearVariationalSolver(jprob)
self.solvers.append(jsolver)
elif key == 'du_smooth':
du = self.diagnostic_variables.fields['du']
du_smooth = self.diagnostic_variables.fields['du_smooth']
projector = Projector(du, du_smooth)
self.projectors.append(projector)
elif key == 'u2_flux':
gamma = simulation_parameters['gamma'][-1]
u2_flux = self.diagnostic_variables.fields['u2_flux']
xis = self.prognostic_variables.pure_xi_list
xis_x = []
xis_xxx = []
CG1 = FunctionSpace(mesh, "CG", 1)
psi = TestFunction(CG1)
for xi in xis:
xis_x.append(Function(CG1).project(xi.dx(0)))
for xi_x in xis_x:
xi_xxx = Function(CG1)
form = (psi * xi_xxx + psi.dx(0) * xi_x.dx(0)) * dx
prob = NonlinearVariationalProblem(form, xi_xxx)
solver = NonlinearVariationalSolver(prob)
solver.solve()
xis_xxx.append(xi_xxx)
flux_expr = 0.0 * x
for xi, xi_x, xi_xxx in zip(xis, xis_x, xis_xxx):
flux_expr += (6 * u.dx(0) * xi + 12 * u * xi_x + gamma *
xi_xxx) * (6 * u.dx(0) * xi + 24 * u * xi_x +
gamma * xi_xxx)
projector = Projector(flux_expr, u2_flux)
self.projectors.append(projector)
elif key == 'a':
# find 6 * u_x * Xi + gamma * Xi_xxx
mesh = u.function_space().mesh()
gamma = simulation_parameters['gamma'][-1]
a_flux = self.diagnostic_variables.fields['a']
xis = self.prognostic_variables.pure_xis
xis_x = []
xis_xxx = []
CG1 = FunctionSpace(mesh, "CG", 1)
psi = TestFunction(CG1)
for xi in xis:
xis_x.append(Function(CG1).project(xi.dx(0)))
for xi_x in xis_x:
xi_xxx = Function(CG1)
form = (psi * xi_xxx + psi.dx(0) * xi_x.dx(0)) * dx
prob = NonlinearVariationalProblem(form, xi_xxx)
solver = NonlinearVariationalSolver(prob)
solver.solve()
xis_xxx.append(xi_xxx)
x, = SpatialCoordinate(mesh)
a_expr = 0.0 * x
for xi, xi_x, xi_xxx in zip(xis, xis_x, xis_xxx):
a_expr += 6 * u.dx(0) * xi + gamma * xi_xxx
projector = Projector(a_expr, a_flux)
self.projectors.append(projector)
elif key == 'b':
# find 12 * u * Xi_x
mesh = u.function_space().mesh()
gamma = simulation_parameters['gamma'][-1]
b_flux = self.diagnostic_variables.fields['b']
xis = self.prognostic_variables.pure_xis
x, = SpatialCoordinate(mesh)
b_expr = 0.0 * x
for xi, xi_x, xi_xxx in zip(xis, xis_x, xis_xxx):
b_expr += 12 * u * xi.dx(0)
projector = Projector(b_expr, b_flux)
self.projectors.append(projector)
elif key == 'kdv_1':
# find the first part of the kdv form
u0 = prognostic_variables.u0
uh = (u + u0) / 2
us = Dt * uh + sqrt(Dt) * Xi
psi = TestFunction(Vu)
du_1 = self.diagnostic_variables.fields['kdv_1']
eqn = psi * du_1 * dx - 6 * psi.dx(0) * uh * us * dx
prob = NonlinearVariationalProblem(eqn, du_1)
solver = NonlinearVariationalSolver(prob)
self.solvers.append(solver)
elif key == 'kdv_2':
# find the second part of the kdv form
u0 = prognostic_variables.u0
uh = (u + u0) / 2
us = Dt * uh + sqrt(Dt) * Xi
psi = TestFunction(Vu)
du_2 = self.diagnostic_variables.fields['kdv_2']
eqn = psi * du_2 * dx + 6 * psi * uh * us.dx(0) * dx
prob = NonlinearVariationalProblem(eqn, du_2)
solver = NonlinearVariationalSolver(prob)
self.solvers.append(solver)
elif key == 'kdv_3':
# find the third part of the kdv form
u0 = prognostic_variables.u0
uh = (u + u0) / 2
us = Dt * uh + sqrt(Dt) * Xi
du_3 = self.diagnostic_variables.fields['kdv_3']
gamma = simulation_parameters['gamma'][-1]
phi = TestFunction(Vu)
F = Function(Vu)
eqn = (phi * F * dx + phi.dx(0) * us.dx(0) * dx)
prob = NonlinearVariationalProblem(eqn, F)
solver = NonlinearVariationalSolver(prob)
self.solvers.append(solver)
self.projectors.append(Projector(-gamma * F.dx(0), du_3))
# nu = TestFunction(Vu)
# back_eqn = nu * du_3 * dx - gamma * nu.dx(0) * F * dx
# back_prob = NonlinearVariationalProblem(back_eqn, du_3)
# back_solver = NonlinearVariationalSolver(back_prob)
# self.solvers.append(solver)
elif key == 'm':
m = self.diagnostic_variables.fields['m']
phi = TestFunction(Vu)
eqn = phi * m * dx - phi * u * dx - alphasq * phi.dx(0) * u.dx(
0) * dx
prob = NonlinearVariationalProblem(eqn, m)
solver = NonlinearVariationalSolver(prob)
self.solvers.append(solver)
elif key == 'u_xx':
u_xx = self.diagnostic_variables.fields['u_xx']
phi = TestFunction(Vu)
eqn = phi * u_xx * dx + phi.dx(0) * u_xx.dx(0) * dx
prob = NonlinearVariationalProblem(eqn, u_xx)
solver = NonlinearVariationalSolver(prob)
self.solvers.append(solver)
elif key == 'u_sde':
self.to_update_constants = True
self.Ld = Ld
self.alphasq = alphasq
self.p = Constant(1.0 * 0.5 * (1 + exp(-Ld / sqrt(alphasq))) /
(1 - exp(-Ld / sqrt(alphasq))))
self.q = Constant(Ld / 2)
u_sde = self.diagnostic_variables.fields['u_sde']
if periodic:
expr = conditional(
x < self.q - Ld / 2,
self.p * ((exp(-(x - self.q + Ld) / sqrt(alphasq)) +
exp(-Ld / sqrt(alphasq)) * exp(
(x - self.q + Ld) / sqrt(alphasq))) /
(1 - exp(-Ld / sqrt(alphasq)))),
conditional(
x < self.q + Ld / 2,
self.p * ((exp(-sqrt((self.q - x)**2 / alphasq)) +
exp(-Ld / sqrt(alphasq)) *
exp(sqrt((self.q - x)**2 / alphasq))) /
(1 - exp(-Ld / sqrt(alphasq)))),
self.p *
((exp(-(self.q + Ld - x) / sqrt(alphasq)) +
exp(-Ld / sqrt(alphasq) * exp(
(self.q + Ld - x) / sqrt(alphasq)))) /
(1 - exp(-Ld / sqrt(alphasq))))))
else:
expr = conditional(
x < self.q - Ld / 2,
self.p * exp(-(x - self.q + Ld) / sqrt(alphasq)),
conditional(
x < self.q + Ld / 2,
self.p * exp(-sqrt((self.q - x)**2 / alphasq)),
self.p * exp(-(self.q + Ld - x) / sqrt(alphasq))))
self.interpolators.append(Interpolator(expr, u_sde))
elif key == 'u_sde_weak':
u_sde = self.diagnostic_variables.fields['u_sde']
u_sde_weak = self.diagnostic_variables.fields['u_sde_weak']
psi = TestFunction(Vu)
eqn = psi * u_sde_weak * dx - psi * (u - u_sde) * dx
prob = NonlinearVariationalProblem(eqn, u_sde_weak)
solver = NonlinearVariationalSolver(prob)
self.solvers.append(solver)
elif key == 'u_sde_mean':
self.to_update_constants = True
self.p = Constant(1.0)
self.q = Constant(Ld / 2)
if periodic:
raise NotImplementedError(
'u_sde_mean not yet implemented for periodic peakon')
u_sde = self.diagnostic_variables.fields['u_sde_mean']
expr = conditional(
x < self.q - Ld / 2,
self.p * exp(-(x - self.q + Ld) / sqrt(alphasq)),
conditional(
x < self.q + Ld / 2,
self.p * exp(-sqrt((self.q - x)**2 / alphasq)),
self.p * exp(-(self.q + Ld - x) / sqrt(alphasq))))
self.interpolators.append(Interpolator(expr, u_sde))
elif key == 'u_sde_weak_mean':
u_sde = self.diagnostic_variables.fields['u_sde_mean']
u_sde_weak = self.diagnostic_variables.fields[
'u_sde_weak_mean']
psi = TestFunction(Vu)
eqn = psi * u_sde_weak * dx - psi * (u - u_sde) * dx
prob = NonlinearVariationalProblem(eqn, u_sde_weak)
solver = NonlinearVariationalSolver(prob)
self.solvers.append(solver)
elif key == 'pure_xi':
pure_xi = 0.0 * x
for xi in self.prognostic_variables.pure_xi_list:
if vector_u:
pure_xi += dot(ones, xi)
else:
pure_xi += xi
Xiscalar = self.diagnostic_variables.fields['pure_xi']
Xi_interpolator = Interpolator(pure_xi, Xiscalar)
self.interpolators.append(Xi_interpolator)
elif key == 'pure_xi_x':
pure_xi_x = 0.0 * x
for xix in self.prognostic_variables.pure_xi_x_list:
if vector_u:
pure_xi_x += dot(ones, xix)
else:
pure_xi_x += xix
Xiscalar = self.diagnostic_variables.fields['pure_xi_x']
Xi_interpolator = Interpolator(pure_xi_x, Xiscalar)
self.interpolators.append(Xi_interpolator)
elif key == 'pure_xi_xx':
pure_xi_xx = 0.0 * x
for xixx in self.prognostic_variables.pure_xi_xx_list:
if vector_u:
pure_xi_xx += dot(ones, xixx)
else:
pure_xi_xx += xixx
Xiscalar = self.diagnostic_variables.fields['pure_xi_xx']
Xi_interpolator = Interpolator(pure_xi_xx, Xiscalar)
self.interpolators.append(Xi_interpolator)
elif key == 'pure_xi_xxx':
pure_xi_xxx = 0.0 * x
for xixxx in self.prognostic_variables.pure_xi_xxx_list:
if vector_u:
pure_xi_xxx += dot(ones, xixxx)
else:
pure_xi_xxx += xixxx
Xiscalar = self.diagnostic_variables.fields['pure_xi_xxx']
Xi_interpolator = Interpolator(pure_xi_xxx, Xiscalar)
self.interpolators.append(Xi_interpolator)
elif key == 'pure_xi_xxxx':
pure_xi_xxxx = 0.0 * x
for xixxxx in self.prognostic_variables.pure_xi_xx_list:
if vector_u:
pure_xi_xxxx += dot(ones, xixxxx)
else:
pure_xi_xxxx += xixxxx
Xiscalar = self.diagnostic_variables.fields['pure_xi_xxxx']
Xi_interpolator = Interpolator(pure_xi_xxxx, Xiscalar)
self.interpolators.append(Xi_interpolator)
else:
raise NotImplementedError('Diagnostic %s not yet implemented' %
key)