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 __init__(self, state, iterations=1):
super().__init__(state)
self.iterations = iterations
# obtain our fields
self.theta = state.fields('theta')
self.water_v = state.fields('water_v')
self.water_c = state.fields('water_c')
rho = state.fields('rho')
try:
rain = state.fields('rain')
water_l = self.water_c + rain
except NotImplementedError:
water_l = self.water_c
# declare function space
Vt = self.theta.function_space()
# make rho variables
# we recover rho into theta space
if state.vertical_degree == 0 and state.horizontal_degree == 0:
boundary_method = Boundary_Method.physics
else:
boundary_method = None
Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element()))
rho_averaged = Function(Vt)
self.rho_recoverer = Recoverer(rho,
rho_averaged,
VDG=Vt_broken,
boundary_method=boundary_method)
# define some parameters as attributes
dt = state.timestepping.dt
R_d = state.parameters.R_d
cp = state.parameters.cp
cv = state.parameters.cv
c_pv = state.parameters.c_pv
c_pl = state.parameters.c_pl
c_vv = state.parameters.c_vv
R_v = state.parameters.R_v
# make useful fields
Pi = thermodynamics.pi(state.parameters, rho_averaged, self.theta)
T = thermodynamics.T(state.parameters,
self.theta,
Pi,
r_v=self.water_v)
p = thermodynamics.p(state.parameters, Pi)
L_v = thermodynamics.Lv(state.parameters, T)
R_m = R_d + R_v * self.water_v
c_pml = cp + c_pv * self.water_v + c_pl * water_l
c_vml = cv + c_vv * self.water_v + c_pl * water_l
# use Teten's formula to calculate w_sat
w_sat = thermodynamics.r_sat(state.parameters, T, p)
# make appropriate condensation rate
dot_r_cond = ((self.water_v - w_sat) / (dt * (1.0 +
((L_v**2.0 * w_sat) /
(cp * R_v * T**2.0)))))
# make cond_rate function, that needs to be the same for all updates in one time step
cond_rate = Function(Vt)
# adjust cond rate so negative concentrations don't occur
self.lim_cond_rate = Interpolator(
conditional(dot_r_cond < 0,
max_value(dot_r_cond, -self.water_c / dt),
min_value(dot_r_cond, self.water_v / dt)), cond_rate)
# tell the prognostic fields what to update to
self.water_v_new = Interpolator(self.water_v - dt * cond_rate, Vt)
self.water_c_new = Interpolator(self.water_c + dt * cond_rate, Vt)
self.theta_new = Interpolator(
self.theta * (1.0 + dt * cond_rate *
(cv * L_v / (c_vml * cp * T) - R_v * cv * c_pml /
(R_m * cp * c_vml))), Vt)
# find perturbed water_v
w_v = Function(Vt)
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")
def __init__(self, state):
super().__init__(state)
# obtain our fields
self.theta = state.fields('theta')
self.water_v = state.fields('water_v')
self.rain = state.fields('rain')
rho = state.fields('rho')
try:
water_c = state.fields('water_c')
water_l = self.rain + water_c
except NotImplementedError:
water_l = self.rain
# declare function space
Vt = self.theta.function_space()
# make rho variables
# we recover rho into theta space
if state.vertical_degree == 0 and state.horizontal_degree == 0:
boundary_method = Boundary_Method.physics
else:
boundary_method = None
Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element()))
rho_averaged = Function(Vt)
self.rho_recoverer = Recoverer(rho,
rho_averaged,
VDG=Vt_broken,
boundary_method=boundary_method)
# define some parameters as attributes
dt = state.timestepping.dt
R_d = state.parameters.R_d
cp = state.parameters.cp
cv = state.parameters.cv
c_pv = state.parameters.c_pv
c_pl = state.parameters.c_pl
c_vv = state.parameters.c_vv
R_v = state.parameters.R_v
# make useful fields
Pi = thermodynamics.pi(state.parameters, rho_averaged, self.theta)
T = thermodynamics.T(state.parameters,
self.theta,
Pi,
r_v=self.water_v)
p = thermodynamics.p(state.parameters, Pi)
L_v = thermodynamics.Lv(state.parameters, T)
R_m = R_d + R_v * self.water_v
c_pml = cp + c_pv * self.water_v + c_pl * water_l
c_vml = cv + c_vv * self.water_v + c_pl * water_l
# use Teten's formula to calculate w_sat
w_sat = thermodynamics.r_sat(state.parameters, T, p)
# expression for ventilation factor
a = Constant(1.6)
b = Constant(124.9)
c = Constant(0.2046)
C = a + b * (rho_averaged * self.rain)**c
# make appropriate condensation rate
f = Constant(5.4e5)
g = Constant(2.55e6)
h = Constant(0.525)
dot_r_evap = (((1 - self.water_v / w_sat) * C *
(rho_averaged * self.rain)**h) /
(rho_averaged * (f + g / (p * w_sat))))
# make evap_rate function, needs to be the same for all updates in one time step
evap_rate = Function(Vt)
# adjust evap rate so negative rain doesn't occur
self.lim_evap_rate = Interpolator(
conditional(
dot_r_evap < 0, 0.0,
conditional(self.rain < 0.0, 0.0,
min_value(dot_r_evap, self.rain / dt))), evap_rate)
# tell the prognostic fields what to update to
self.water_v_new = Interpolator(self.water_v + dt * evap_rate, Vt)
self.rain_new = Interpolator(self.rain - dt * evap_rate, Vt)
self.theta_new = Interpolator(
self.theta * (1.0 - dt * evap_rate *
(cv * L_v / (c_vml * cp * T) - R_v * cv * c_pml /
(R_m * cp * c_vml))), Vt)
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)