def compute(self, state):
theta = state.fields('theta')
rho = state.fields('rho')
w_v = state.fields('water_v')
w_c = state.fields('water_c')
w_t = w_c + w_v
pi = thermodynamics.pi(state.parameters, rho, theta)
p = thermodynamics.p(state.parameters, pi)
T = thermodynamics.T(state.parameters, theta, pi, r_v=w_v)
return self.field.interpolate(thermodynamics.theta_e(state.parameters, T, p, w_v, w_t))
def compute(self, state):
super().compute(state)
return self.field.assign(thermodynamics.theta_e(state.parameters, self.T, self.p, self.r_v, self.r_t))
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 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)