def setup(self, state):
if not self._initialised:
Vu = state.spaces("HDiv")
space = FunctionSpace(state.mesh, Vu.ufl_element()._elements[-1])
super().setup(state, space=space)
rho = state.fields("rho")
rhobar = state.fields("rhobar")
theta = state.fields("theta")
thetabar = state.fields("thetabar")
exner = thermodynamics.exner_pressure(state.parameters, rho, theta)
exnerbar = thermodynamics.exner_pressure(state.parameters, rhobar,
thetabar)
cp = Constant(state.parameters.cp)
n = FacetNormal(state.mesh)
F = TrialFunction(space)
w = TestFunction(space)
a = inner(w, F) * dx
L = (-cp * div((theta - thetabar) * w) * exnerbar * dx + cp * jump(
(theta - thetabar) * w, n) * avg(exnerbar) * dS_v -
cp * div(thetabar * w) * (exner - exnerbar) * dx +
cp * jump(thetabar * w, n) * avg(exner - exnerbar) * dS_v)
bcs = [
DirichletBC(space, 0.0, "bottom"),
DirichletBC(space, 0.0, "top")
]
imbalanceproblem = LinearVariationalProblem(a,
L,
self.field,
bcs=bcs)
self.imbalance_solver = LinearVariationalSolver(imbalanceproblem)
def calculate_exner0(state, theta0, rho0):
# exner function
Vr = rho0.function_space()
exner = Function(Vr).interpolate(
thermodynamics.exner_pressure(state.parameters, rho0, theta0))
exner0 = assemble(exner * dx) / assemble(
Constant(1) * dx(domain=state.mesh))
return exner0
def setup(self, state):
if not self._initialised:
space = state.fields("theta").function_space()
broken_space = FunctionSpace(state.mesh,
BrokenElement(space.ufl_element()))
h_deg = space.ufl_element().degree()[0]
v_deg = space.ufl_element().degree()[1] - 1
boundary_method = Boundary_Method.physics if (
v_deg == 0 and h_deg == 0) else None
super().setup(state, space=space)
# now let's attach all of our fields
self.u = state.fields("u")
self.rho = state.fields("rho")
self.theta = state.fields("theta")
self.rho_averaged = Function(space)
self.recoverer = Recoverer(self.rho,
self.rho_averaged,
VDG=broken_space,
boundary_method=boundary_method)
try:
self.r_v = state.fields("vapour_mixing_ratio")
except NotImplementedError:
self.r_v = Constant(0.0)
try:
self.r_c = state.fields("cloud_liquid_mixing_ratio")
except NotImplementedError:
self.r_c = Constant(0.0)
try:
self.rain = state.fields("rain_mixing_ratio")
except NotImplementedError:
self.rain = Constant(0.0)
# now let's store the most common expressions
self.exner = thermodynamics.exner_pressure(state.parameters,
self.rho_averaged,
self.theta)
self.T = thermodynamics.T(state.parameters,
self.theta,
self.exner,
r_v=self.r_v)
self.p = thermodynamics.p(state.parameters, self.exner)
self.r_l = self.r_c + self.rain
self.r_t = self.r_v + self.r_c + self.rain
def _setup_solver(self):
import numpy as np
state = self.state
dt = state.dt
beta_ = dt * self.alpha
cp = state.parameters.cp
Vu = state.spaces("HDiv")
Vu_broken = FunctionSpace(state.mesh, BrokenElement(Vu.ufl_element()))
Vtheta = state.spaces("theta")
Vrho = state.spaces("DG")
# Store time-stepping coefficients as UFL Constants
beta = Constant(beta_)
beta_cp = Constant(beta_ * cp)
h_deg = Vrho.ufl_element().degree()[0]
v_deg = Vrho.ufl_element().degree()[1]
Vtrace = FunctionSpace(state.mesh, "HDiv Trace", degree=(h_deg, v_deg))
# Split up the rhs vector (symbolically)
self.xrhs = Function(self.equations.function_space)
u_in, rho_in, theta_in = split(self.xrhs)[0:3]
# Build the function space for "broken" u, rho, and pressure trace
M = MixedFunctionSpace((Vu_broken, Vrho, Vtrace))
w, phi, dl = TestFunctions(M)
u, rho, l0 = TrialFunctions(M)
n = FacetNormal(state.mesh)
# Get background fields
thetabar = state.fields("thetabar")
rhobar = state.fields("rhobar")
exnerbar = thermodynamics.exner_pressure(state.parameters, rhobar,
thetabar)
exnerbar_rho = thermodynamics.dexner_drho(state.parameters, rhobar,
thetabar)
exnerbar_theta = thermodynamics.dexner_dtheta(state.parameters, rhobar,
thetabar)
# Analytical (approximate) elimination of theta
k = state.k # Upward pointing unit vector
theta = -dot(k, u) * dot(k, grad(thetabar)) * beta + theta_in
# Only include theta' (rather than exner') in the vertical
# component of the gradient
# The exner prime term (here, bars are for mean and no bars are
# for linear perturbations)
exner = exnerbar_theta * theta + exnerbar_rho * rho
# Vertical projection
def V(u):
return k * inner(u, k)
# hydrostatic projection
h_project = lambda u: u - k * inner(u, k)
# Specify degree for some terms as estimated degree is too large
dxp = dx(degree=(self.quadrature_degree))
dS_vp = dS_v(degree=(self.quadrature_degree))
dS_hp = dS_h(degree=(self.quadrature_degree))
ds_vp = ds_v(degree=(self.quadrature_degree))
ds_tbp = (ds_t(degree=(self.quadrature_degree)) +
ds_b(degree=(self.quadrature_degree)))
# Add effect of density of water upon theta
if self.moisture is not None:
water_t = Function(Vtheta).assign(0.0)
for water in self.moisture:
water_t += self.state.fields(water)
theta_w = theta / (1 + water_t)
thetabar_w = thetabar / (1 + water_t)
else:
theta_w = theta
thetabar_w = thetabar
_l0 = TrialFunction(Vtrace)
_dl = TestFunction(Vtrace)
a_tr = _dl('+') * _l0('+') * (
dS_vp + dS_hp) + _dl * _l0 * ds_vp + _dl * _l0 * ds_tbp
def L_tr(f):
return _dl('+') * avg(f) * (
dS_vp + dS_hp) + _dl * f * ds_vp + _dl * f * ds_tbp
cg_ilu_parameters = {
'ksp_type': 'cg',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
}
# Project field averages into functions on the trace space
rhobar_avg = Function(Vtrace)
exnerbar_avg = Function(Vtrace)
rho_avg_prb = LinearVariationalProblem(a_tr, L_tr(rhobar), rhobar_avg)
exner_avg_prb = LinearVariationalProblem(a_tr, L_tr(exnerbar),
exnerbar_avg)
rho_avg_solver = LinearVariationalSolver(
rho_avg_prb,
solver_parameters=cg_ilu_parameters,
options_prefix='rhobar_avg_solver')
exner_avg_solver = LinearVariationalSolver(
exner_avg_prb,
solver_parameters=cg_ilu_parameters,
options_prefix='exnerbar_avg_solver')
with timed_region("Gusto:HybridProjectRhobar"):
rho_avg_solver.solve()
with timed_region("Gusto:HybridProjectExnerbar"):
exner_avg_solver.solve()
# "broken" u, rho, and trace system
# NOTE: no ds_v integrals since equations are defined on
# a periodic (or sphere) base mesh.
if any([t.has_label(hydrostatic) for t in self.equations.residual]):
u_mass = inner(w, (h_project(u) - u_in)) * dx
else:
u_mass = inner(w, (u - u_in)) * dx
eqn = (
# momentum equation
u_mass - beta_cp * div(theta_w * V(w)) * exnerbar * dxp
# following does nothing but is preserved in the comments
# to remind us why (because V(w) is purely vertical).
# + beta_cp*jump(theta_w*V(w), n=n)*exnerbar_avg('+')*dS_vp
+ beta_cp * jump(theta_w * V(w), n=n) * exnerbar_avg('+') * dS_hp +
beta_cp * dot(theta_w * V(w), n) * exnerbar_avg * ds_tbp -
beta_cp * div(thetabar_w * w) * exner * dxp
# trace terms appearing after integrating momentum equation
+ beta_cp * jump(thetabar_w * w, n=n) * l0('+') * (dS_vp + dS_hp) +
beta_cp * dot(thetabar_w * w, n) * l0 * (ds_tbp + ds_vp)
# mass continuity equation
+ (phi *
(rho - rho_in) - beta * inner(grad(phi), u) * rhobar) * dx +
beta * jump(phi * u, n=n) * rhobar_avg('+') * (dS_v + dS_h)
# term added because u.n=0 is enforced weakly via the traces
+ beta * phi * dot(u, n) * rhobar_avg * (ds_tb + ds_v)
# constraint equation to enforce continuity of the velocity
# through the interior facets and weakly impose the no-slip
# condition
+ dl('+') * jump(u, n=n) * (dS_vp + dS_hp) + dl * dot(u, n) *
(ds_tbp + ds_vp))
# contribution of the sponge term
if hasattr(self.equations, "mu"):
eqn += dt * self.equations.mu * inner(w, k) * inner(u, k) * dx
aeqn = lhs(eqn)
Leqn = rhs(eqn)
# Function for the hybridized solutions
self.urhol0 = Function(M)
hybridized_prb = LinearVariationalProblem(aeqn, Leqn, self.urhol0)
hybridized_solver = LinearVariationalSolver(
hybridized_prb,
solver_parameters=self.solver_parameters,
options_prefix='ImplicitSolver')
self.hybridized_solver = hybridized_solver
# Project broken u into the HDiv space using facet averaging.
# Weight function counting the dofs of the HDiv element:
shapes = {
"i": Vu.finat_element.space_dimension(),
"j": np.prod(Vu.shape, dtype=int)
}
weight_kernel = """
for (int i=0; i<{i}; ++i)
for (int j=0; j<{j}; ++j)
w[i*{j} + j] += 1.0;
""".format(**shapes)
self._weight = Function(Vu)
par_loop(weight_kernel, dx, {"w": (self._weight, INC)})
# Averaging kernel
self._average_kernel = """
for (int i=0; i<{i}; ++i)
for (int j=0; j<{j}; ++j)
vec_out[i*{j} + j] += vec_in[i*{j} + j]/w[i*{j} + j];
""".format(**shapes)
# HDiv-conforming velocity
self.u_hdiv = Function(Vu)
# Reconstruction of theta
theta = TrialFunction(Vtheta)
gamma = TestFunction(Vtheta)
self.theta = Function(Vtheta)
theta_eqn = gamma * (theta - theta_in + dot(k, self.u_hdiv) *
dot(k, grad(thetabar)) * beta) * dx
theta_problem = LinearVariationalProblem(lhs(theta_eqn),
rhs(theta_eqn), self.theta)
self.theta_solver = LinearVariationalSolver(
theta_problem,
solver_parameters=cg_ilu_parameters,
options_prefix='thetabacksubstitution')
# Store boundary conditions for the div-conforming velocity to apply
# post-solve
self.bcs = self.equations.bcs['u']
def compute(self, state):
rho = state.fields(self.rho_name)
theta = state.fields(self.theta_name)
exner = thermodynamics.exner_pressure(state.parameters, rho, theta)
return self.field.interpolate(exner)
def unsaturated_hydrostatic_balance(state,
theta_d,
H,
exner0=None,
top=False,
exner_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 exner0: 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 exner_boundary: The value of exner 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')
mr_v0 = state.fields('vapour_mixing_ratio')
# Calculate hydrostatic exner pressure
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)
mr_v0.assign(0.01)
v_deg = Vr.ufl_element().degree()[1]
if v_deg == 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 mr_v0
exner = thermodynamics.exner_pressure(state.parameters, rho_averaged,
theta0)
p = thermodynamics.p(state.parameters, exner)
T = thermodynamics.T(state.parameters, theta0, exner, mr_v0)
r_v_expr = thermodynamics.r_v(state.parameters, H, T, p)
# make expressions to evaluate residual
exner_ev = thermodynamics.exner_pressure(state.parameters, rho_averaged,
theta0)
p_ev = thermodynamics.p(state.parameters, exner_ev)
T_ev = thermodynamics.T(state.parameters, theta0, exner_ev, mr_v0)
RH_ev = thermodynamics.RH(state.parameters, mr_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,
exner_boundary=exner_boundary,
mr_t=mr_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)
mr_v0.assign(mr_v0 * (1 - delta) + delta * w_h)
# compute theta_vd
theta0.assign(theta_d * (1 + mr_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 exner0 is not None:
exner = thermodynamics.exner_pressure(state.parameters, rho0, theta0)
exner0.interpolate(exner)
# do one extra solve for rho
compressible_hydrostatic_balance(state,
theta0,
rho0,
top=top,
exner_boundary=exner_boundary,
mr_t=mr_v0,
solve_for_rho=True)
def saturated_hydrostatic_balance(state,
theta_e,
mr_t,
exner0=None,
top=False,
exner_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, mr_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 mr_t: The total water pseudo-mixing ratio profile.
:arg exner0: 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 exner_boundary: The value of exner 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')
mr_v0 = state.fields('vapour_mixing_ratio')
# Calculate hydrostatic exner pressure
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)
mr_v0.interpolate(mr_t)
v_deg = Vr.ufl_element().degree()[1]
if v_deg == 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 mr_v0 from theta_e and mr_t
exner = thermodynamics.exner_pressure(state.parameters, rho_averaged,
theta0)
p = thermodynamics.p(state.parameters, exner)
T = thermodynamics.T(state.parameters, theta0, exner, mr_v0)
r_v_expr = thermodynamics.r_sat(state.parameters, T, p)
theta_e_expr = thermodynamics.theta_e(state.parameters, T, p, mr_v0, mr_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,
exner_boundary=exner_boundary,
mr_t=mr_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)
mr_v0.assign(mr_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 exner0 is not None:
exner = thermodynamics.exner(state.parameters, rho0, theta0)
exner0.interpolate(exner)
# do one extra solve for rho
compressible_hydrostatic_balance(state,
theta0,
rho0,
top=top,
exner_boundary=exner_boundary,
mr_t=mr_t,
solve_for_rho=True)
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))
# Define constant theta_e and water_t
Tsurf = 320.0
total_water = 0.02
theta_e = Function(Vt).assign(Tsurf)
water_t = Function(Vt).assign(total_water)
# Calculate hydrostatic fields
saturated_hydrostatic_balance(state, theta_e, water_t)
# make mean fields
theta_b = Function(Vt).assign(theta0)
rho_b = Function(Vr).assign(rho0)
water_vb = Function(Vt).assign(water_v0)
water_cb = Function(Vt).assign(water_t - water_vb)
exner_b = thermodynamics.exner_pressure(state.parameters, rho_b, theta_b)
Tb = thermodynamics.T(state.parameters, theta_b, exner_b, r_v=water_vb)
# 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,
Tdash * (cos(pi * r / (2.0 * rc)))**2))
# define initial theta
theta0.assign(theta_b * (theta_pert / 300.0 + 1.0))
def __init__(self, state, iterations=1):
super().__init__(state)
self.iterations = iterations
# obtain our fields
self.theta = state.fields('theta')
self.water_v = state.fields('vapour_mixing_ratio')
self.water_c = state.fields('cloud_liquid_mixing_ratio')
rho = state.fields('rho')
try:
# TODO: use the phase flag for the tracers here
rain = state.fields('rain_mixing_ratio')
water_l = self.water_c + rain
except NotImplementedError:
water_l = self.water_c
# declare function space
Vt = self.theta.function_space()
# make rho variables
# we recover rho into theta space
h_deg = rho.function_space().ufl_element().degree()[0]
v_deg = rho.function_space().ufl_element().degree()[1]
if v_deg == 0 and h_deg == 0:
boundary_method = Boundary_Method.physics
else:
boundary_method = None
Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element()))
rho_averaged = Function(Vt)
self.rho_recoverer = Recoverer(rho, rho_averaged, VDG=Vt_broken, boundary_method=boundary_method)
# define some parameters as attributes
dt = state.dt
R_d = state.parameters.R_d
cp = state.parameters.cp
cv = state.parameters.cv
c_pv = state.parameters.c_pv
c_pl = state.parameters.c_pl
c_vv = state.parameters.c_vv
R_v = state.parameters.R_v
# make useful fields
exner = thermodynamics.exner_pressure(state.parameters, rho_averaged, self.theta)
T = thermodynamics.T(state.parameters, self.theta, exner, r_v=self.water_v)
p = thermodynamics.p(state.parameters, exner)
L_v = thermodynamics.Lv(state.parameters, T)
R_m = R_d + R_v * self.water_v
c_pml = cp + c_pv * self.water_v + c_pl * water_l
c_vml = cv + c_vv * self.water_v + c_pl * water_l
# use Teten's formula to calculate w_sat
w_sat = thermodynamics.r_sat(state.parameters, T, p)
# make appropriate condensation rate
dot_r_cond = ((self.water_v - w_sat)
/ (dt * (1.0 + ((L_v ** 2.0 * w_sat)
/ (cp * R_v * T ** 2.0)))))
# make cond_rate function, that needs to be the same for all updates in one time step
cond_rate = Function(Vt)
# adjust cond rate so negative concentrations don't occur
self.lim_cond_rate = Interpolator(conditional(dot_r_cond < 0,
max_value(dot_r_cond, - self.water_c / dt),
min_value(dot_r_cond, self.water_v / dt)), cond_rate)
# tell the prognostic fields what to update to
self.water_v_new = Interpolator(self.water_v - dt * cond_rate, Vt)
self.water_c_new = Interpolator(self.water_c + dt * cond_rate, Vt)
self.theta_new = Interpolator(self.theta
* (1.0 + dt * cond_rate
* (cv * L_v / (c_vml * cp * T)
- R_v * cv * c_pml / (R_m * cp * c_vml))), Vt)
def __init__(self, state):
super().__init__(state)
# obtain our fields
self.theta = state.fields('theta')
self.water_v = state.fields('vapour_mixing_ratio')
self.rain = state.fields('rain_mixing_ratio')
rho = state.fields('rho')
try:
water_c = state.fields('cloud_liquid_mixing_ratio')
water_l = self.rain + water_c
except NotImplementedError:
water_l = self.rain
# declare function space
Vt = self.theta.function_space()
# make rho variables
# we recover rho into theta space
h_deg = rho.function_space().ufl_element().degree()[0]
v_deg = rho.function_space().ufl_element().degree()[1]
if v_deg == 0 and h_deg == 0:
boundary_method = Boundary_Method.physics
else:
boundary_method = None
Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element()))
rho_averaged = Function(Vt)
self.rho_recoverer = Recoverer(rho, rho_averaged, VDG=Vt_broken, boundary_method=boundary_method)
# define some parameters as attributes
dt = state.dt
R_d = state.parameters.R_d
cp = state.parameters.cp
cv = state.parameters.cv
c_pv = state.parameters.c_pv
c_pl = state.parameters.c_pl
c_vv = state.parameters.c_vv
R_v = state.parameters.R_v
# make useful fields
exner = thermodynamics.exner_pressure(state.parameters, rho_averaged, self.theta)
T = thermodynamics.T(state.parameters, self.theta, exner, r_v=self.water_v)
p = thermodynamics.p(state.parameters, exner)
L_v = thermodynamics.Lv(state.parameters, T)
R_m = R_d + R_v * self.water_v
c_pml = cp + c_pv * self.water_v + c_pl * water_l
c_vml = cv + c_vv * self.water_v + c_pl * water_l
# use Teten's formula to calculate w_sat
w_sat = thermodynamics.r_sat(state.parameters, T, p)
# expression for ventilation factor
a = Constant(1.6)
b = Constant(124.9)
c = Constant(0.2046)
C = a + b * (rho_averaged * self.rain) ** c
# make appropriate condensation rate
f = Constant(5.4e5)
g = Constant(2.55e6)
h = Constant(0.525)
dot_r_evap = (((1 - self.water_v / w_sat) * C * (rho_averaged * self.rain) ** h)
/ (rho_averaged * (f + g / (p * w_sat))))
# make evap_rate function, needs to be the same for all updates in one time step
evap_rate = Function(Vt)
# adjust evap rate so negative rain doesn't occur
self.lim_evap_rate = Interpolator(conditional(dot_r_evap < 0,
0.0,
conditional(self.rain < 0.0,
0.0,
min_value(dot_r_evap, self.rain / dt))),
evap_rate)
# tell the prognostic fields what to update to
self.water_v_new = Interpolator(self.water_v + dt * evap_rate, Vt)
self.rain_new = Interpolator(self.rain - dt * evap_rate, Vt)
self.theta_new = Interpolator(self.theta
* (1.0 - dt * evap_rate
* (cv * L_v / (c_vml * cp * T)
- R_v * cv * c_pml / (R_m * cp * c_vml))), Vt)
def __init__(
self,
state,
family,
degree,
Omega=None,
sponge=None,
extra_terms=None,
terms_to_linearise={
'u': [time_derivative],
'rho': [time_derivative, transport],
'theta': [time_derivative, transport]
},
u_transport_option="vector_invariant_form",
diffusion_options=None,
no_normal_flow_bc_ids=None,
active_tracers=None):
field_names = ['u', 'rho', 'theta']
if active_tracers is None:
active_tracers = []
super().__init__(field_names,
state,
family,
degree,
terms_to_linearise=terms_to_linearise,
no_normal_flow_bc_ids=no_normal_flow_bc_ids,
active_tracers=active_tracers)
g = state.parameters.g
cp = state.parameters.cp
w, phi, gamma = self.tests[0:3]
u, rho, theta = split(self.X)[0:3]
rhobar = state.fields("rhobar", space=state.spaces("DG"), dump=False)
thetabar = state.fields("thetabar",
space=state.spaces("theta"),
dump=False)
zero_expr = Constant(0.0) * theta
exner = exner_pressure(state.parameters, rho, theta)
n = FacetNormal(state.mesh)
# -------------------------------------------------------------------- #
# Time Derivative Terms
# -------------------------------------------------------------------- #
mass_form = self.generate_mass_terms()
# -------------------------------------------------------------------- #
# Transport Terms
# -------------------------------------------------------------------- #
# Velocity transport term -- depends on formulation
if u_transport_option == "vector_invariant_form":
u_adv = prognostic(vector_invariant_form(state, w, u), "u")
elif u_transport_option == "vector_advection_form":
u_adv = prognostic(advection_form(state, w, u), "u")
elif u_transport_option == "vector_manifold_advection_form":
u_adv = prognostic(vector_manifold_advection_form(state, w, u),
"u")
elif u_transport_option == "circulation_form":
ke_form = kinetic_energy_form(state, w, u)
ke_form = transport.remove(ke_form)
ke_form = ke_form.label_map(
lambda t: t.has_label(transporting_velocity), lambda t: Term(
ufl.replace(t.form, {t.get(transporting_velocity): u}), t.
labels))
ke_form = transporting_velocity.remove(ke_form)
u_adv = advection_equation_circulation_form(state, w, u) + ke_form
else:
raise ValueError("Invalid u_transport_option: %s" %
u_transport_option)
# Density transport (conservative form)
rho_adv = prognostic(continuity_form(state, phi, rho), "rho")
# Transport term needs special linearisation
if transport in terms_to_linearise['rho']:
linear_rho_adv = linear_continuity_form(
state, phi, rhobar).label_map(
lambda t: t.has_label(transporting_velocity),
lambda t: Term(
ufl.replace(t.form, {
t.get(transporting_velocity):
self.trials[0]
}), t.labels))
rho_adv = linearisation(rho_adv, linear_rho_adv)
# Potential temperature transport (advective form)
theta_adv = prognostic(advection_form(state, gamma, theta), "theta")
# Transport term needs special linearisation
if transport in terms_to_linearise['theta']:
linear_theta_adv = linear_advection_form(
state, gamma, thetabar).label_map(
lambda t: t.has_label(transporting_velocity),
lambda t: Term(
ufl.replace(t.form, {
t.get(transporting_velocity):
self.trials[0]
}), t.labels))
theta_adv = linearisation(theta_adv, linear_theta_adv)
adv_form = subject(u_adv + rho_adv + theta_adv, self.X)
# Add transport of tracers
if len(active_tracers) > 0:
adv_form += self.generate_tracer_transport_terms(
state, active_tracers)
# -------------------------------------------------------------------- #
# Pressure Gradient Term
# -------------------------------------------------------------------- #
# First get total mass of tracers
tracer_mr_total = zero_expr
for tracer in active_tracers:
if tracer.variable_type == TracerVariableType.mixing_ratio:
idx = self.field_names.index(tracer.name)
tracer_mr_total += split(self.X)[idx]
else:
raise NotImplementedError(
'Only mixing ratio tracers are implemented')
theta_v = theta / (Constant(1.0) + tracer_mr_total)
pressure_gradient_form = name(
subject(
prognostic(
cp * (-div(theta_v * w) * exner * dx +
jump(theta_v * w, n) * avg(exner) * dS_v), "u"),
self.X), "pressure_gradient")
# -------------------------------------------------------------------- #
# Gravitational Term
# -------------------------------------------------------------------- #
gravity_form = subject(
prognostic(Term(g * inner(state.k, w) * dx), "u"), self.X)
residual = (mass_form + adv_form + pressure_gradient_form +
gravity_form)
# -------------------------------------------------------------------- #
# Moist Thermodynamic Divergence Term
# -------------------------------------------------------------------- #
if len(active_tracers) > 0:
cv = state.parameters.cv
c_vv = state.parameters.c_vv
c_pv = state.parameters.c_pv
c_pl = state.parameters.c_pl
R_d = state.parameters.R_d
R_v = state.parameters.R_v
# Get gas and liquid moisture mixing ratios
mr_l = zero_expr
mr_v = zero_expr
for tracer in active_tracers:
if tracer.is_moisture:
if tracer.variable_type == TracerVariableType.mixing_ratio:
idx = self.field_names.index(tracer.name)
if tracer.phase == Phases.gas:
mr_v += split(self.X)[idx]
elif tracer.phase == Phases.liquid:
mr_l += split(self.X)[idx]
else:
raise NotImplementedError(
'Only mixing ratio tracers are implemented')
c_vml = cv + mr_v * c_vv + mr_l * c_pl
c_pml = cp + mr_v * c_pv + mr_l * c_pl
R_m = R_d + mr_v * R_v
residual += subject(
prognostic(
gamma * theta * div(u) * (R_m / c_vml - (R_d * c_pml) /
(cp * c_vml)) * dx, "theta"),
self.X)
# -------------------------------------------------------------------- #
# Extra Terms (Coriolis, Sponge, Diffusion and others)
# -------------------------------------------------------------------- #
if Omega is not None:
residual += subject(
prognostic(inner(w, cross(2 * Omega, u)) * dx, "u"), self.X)
if sponge is not None:
W_DG = FunctionSpace(state.mesh, "DG", 2)
x = SpatialCoordinate(state.mesh)
z = x[len(x) - 1]
H = sponge.H
zc = sponge.z_level
assert float(zc) < float(
H
), "you have set the sponge level above the height of your domain"
mubar = sponge.mubar
muexpr = conditional(
z <= zc, 0.0,
mubar * sin((pi / 2.) * (z - zc) / (H - zc))**2)
self.mu = Function(W_DG).interpolate(muexpr)
residual += name(
subject(
prognostic(
self.mu * inner(w, state.k) * inner(u, state.k) * dx,
"u"), self.X), "sponge")
if diffusion_options is not None:
for field, diffusion in diffusion_options:
idx = self.field_names.index(field)
test = self.tests[idx]
fn = split(self.X)[idx]
residual += subject(
prognostic(
interior_penalty_diffusion_form(
state, test, fn, diffusion), field), self.X)
if extra_terms is not None:
for field, term in extra_terms:
idx = self.field_names.index(field)
test = self.tests[idx]
residual += subject(prognostic(inner(test, term) * dx, field),
self.X)
# -------------------------------------------------------------------- #
# Linearise equations
# -------------------------------------------------------------------- #
# TODO: add linearisation states for variables
# Add linearisations to equations
self.residual = self.generate_linear_terms(residual,
self.terms_to_linearise)