def setup(self, state):
if not self._initialised:
space = state.spaces("Vv")
super().setup(state, space=space)
rho = state.fields("rho")
rhobar = state.fields("rhobar")
theta = state.fields("theta")
thetabar = state.fields("thetabar")
pi = thermodynamics.pi(state.parameters, rho, theta)
pibar = thermodynamics.pi(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)*pibar*dx
+ cp*jump((theta-thetabar)*w, n)*avg(pibar)*dS_v
- cp*div(thetabar*w)*(pi-pibar)*dx
+ cp*jump(thetabar*w, n)*avg(pi-pibar)*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_Pi0(state, theta0, rho0):
# exner function
Vr = rho0.function_space()
Pi = Function(Vr).interpolate(thermodynamics.pi(state.parameters, rho0, theta0))
Pi0 = assemble(Pi*dx)/assemble(Constant(1)*dx(domain=state.mesh))
return Pi0
def setup(self, state):
if not self._initialised:
space = state.fields("theta").function_space()
broken_space = FunctionSpace(state.mesh, BrokenElement(space.ufl_element()))
boundary_method = Boundary_Method.physics if (state.vertical_degree == 0 and state.horizontal_degree == 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("water_v")
except NotImplementedError:
self.r_v = Constant(0.0)
try:
self.r_c = state.fields("water_c")
except NotImplementedError:
self.r_c = Constant(0.0)
try:
self.rain = state.fields("rain")
except NotImplementedError:
self.rain = Constant(0.0)
# now let's store the most common expressions
self.pi = thermodynamics.pi(state.parameters, self.rho_averaged, self.theta)
self.T = thermodynamics.T(state.parameters, self.theta, self.pi, r_v=self.r_v)
self.p = thermodynamics.p(state.parameters, self.pi)
self.r_l = self.r_c + self.rain
self.r_t = self.r_v + self.r_c + self.rain
def compressible_eady_initial_v(state, theta0, rho0, v):
f = state.parameters.f
cp = state.parameters.cp
# exner function
Vr = rho0.function_space()
Pi = Function(Vr).interpolate(thermodynamics.pi(state.parameters, rho0, theta0))
# get Pi gradient
Vu = state.spaces("HDiv")
g = TrialFunction(Vu)
wg = TestFunction(Vu)
n = FacetNormal(state.mesh)
a = inner(wg, g)*dx
L = -div(wg)*Pi*dx + inner(wg, n)*Pi*ds_tb
pgrad = Function(Vu)
solve(a == L, pgrad)
# get initial v
m = TrialFunction(Vr)
phi = TestFunction(Vr)
a = phi*f*m*dx
L = phi*cp*theta0*pgrad[0]*dx
solve(a == L, v)
return v
def compute(self, state):
theta = state.fields('theta')
rho = state.fields('rho')
w_v = state.fields('water_v')
w_c = state.fields('water_c')
pi = thermodynamics.pi(state.parameters, rho, theta)
T = thermodynamics.T(state.parameters, theta, pi, r_v=w_v)
return self.field.interpolate(thermodynamics.internal_energy(state.parameters, rho, T, r_v=w_v, r_l=w_c))
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):
x, y, z = SpatialCoordinate(state.mesh)
g = state.parameters.g
cp = state.parameters.cp
cv = state.parameters.cv
Pi0 = state.parameters.Pi0
rho = state.fields("rho")
theta = state.fields("theta")
Pi = thermodynamics.pi(state.parameters, rho, theta)
potential = rho * (g * z + cv * Pi * theta - cp * Pi0 * theta)
return self.field.interpolate(potential)
def forcing_term(self):
# L = super(EadyForcing, self).forcing_term()
L = Forcing.forcing_term(self)
dthetady = self.state.parameters.dthetady
Pi0 = self.state.parameters.Pi0
cp = self.state.parameters.cp
_, rho0, theta0 = split(self.x0)
Pi = thermodynamics.pi(self.state.parameters, rho0, theta0)
Pi_0 = Constant(Pi0)
L += self.scaling * cp * dthetady * (Pi - Pi_0) * inner(
self.test, as_vector([0., 1., 0.])) * dx # Eady forcing
return L
def pressure_gradient_term(self):
u0, rho0, theta0 = split(self.x0)
cp = self.state.parameters.cp
n = FacetNormal(self.state.mesh)
Vtheta = self.state.spaces("HDiv_v")
# introduce new theta so it can be changed by moisture
theta = theta0
# 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 = theta / (1 + water_t)
pi = thermodynamics.pi(self.state.parameters, rho0, theta0)
L = (+cp * div(theta * self.test) * pi * dx -
cp * jump(self.test * theta, n) * avg(pi) * dS_v)
return L
def compute(self, state):
rho = state.fields(self.rho_name)
theta = state.fields(self.theta_name)
Pi = thermodynamics.pi(state.parameters, rho, theta)
return self.field.interpolate(Pi)
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)
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 _setup_solver(self):
import numpy as np
state = self.state
Dt = state.timestepping.dt
beta_ = Dt*state.timestepping.alpha
cp = state.parameters.cp
mu = state.mu
Vu = state.spaces("HDiv")
Vu_broken = FunctionSpace(state.mesh, BrokenElement(Vu.ufl_element()))
Vtheta = state.spaces("HDiv_v")
Vrho = state.spaces("DG")
# Store time-stepping coefficients as UFL Constants
dt = Constant(Dt)
beta = Constant(beta_)
beta_cp = Constant(beta_ * cp)
h_deg = state.horizontal_degree
v_deg = state.vertical_degree
Vtrace = FunctionSpace(state.mesh, "HDiv Trace", degree=(h_deg, v_deg))
# Split up the rhs vector (symbolically)
u_in, rho_in, theta_in = split(state.xrhs)
# 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")
pibar = thermodynamics.pi(state.parameters, rhobar, thetabar)
pibar_rho = thermodynamics.pi_rho(state.parameters, rhobar, thetabar)
pibar_theta = thermodynamics.pi_theta(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 pi') in the vertical
# component of the gradient
# The pi prime term (here, bars are for mean and no bars are
# for linear perturbations)
pi = pibar_theta*theta + pibar_rho*rho
# Vertical projection
def V(u):
return 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)
pibar_avg = Function(Vtrace)
rho_avg_prb = LinearVariationalProblem(a_tr, L_tr(rhobar), rhobar_avg)
pi_avg_prb = LinearVariationalProblem(a_tr, L_tr(pibar), pibar_avg)
rho_avg_solver = LinearVariationalSolver(rho_avg_prb,
solver_parameters=cg_ilu_parameters,
options_prefix='rhobar_avg_solver')
pi_avg_solver = LinearVariationalSolver(pi_avg_prb,
solver_parameters=cg_ilu_parameters,
options_prefix='pibar_avg_solver')
with timed_region("Gusto:HybridProjectRhobar"):
rho_avg_solver.solve()
with timed_region("Gusto:HybridProjectPibar"):
pi_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.
eqn = (
# momentum equation
inner(w, (state.h_project(u) - u_in))*dx
- beta_cp*div(theta_w*V(w))*pibar*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)*pibar_avg('+')*dS_vp
+ beta_cp*jump(theta_w*V(w), n=n)*pibar_avg('+')*dS_hp
+ beta_cp*dot(theta_w*V(w), n)*pibar_avg*ds_tbp
- beta_cp*div(thetabar_w*w)*pi*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 mu is not None:
eqn += dt*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.state.bcs
# 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)
pibar = thermodynamics.pi(state.parameters, rho_b, theta_b)
Tb = thermodynamics.T(state.parameters, theta_b, pibar, r_v=water_vb)
Ibar = state.fields("InternalEnergybar", Vt, dump=False)
Ibar.interpolate(
thermodynamics.internal_energy(state.parameters,
rho_b,
Tb,
r_v=water_vb,
r_l=water_cb))
# define perturbation
xc = L / 2
zc = 2000.
rc = 2000.
Tdash = 2.0
r = sqrt((x - xc)**2 + (z - zc)**2)
def compressible_hydrostatic_balance(state, theta0, rho0, pi0=None,
top=False, pi_boundary=Constant(1.0),
water_t=None,
solve_for_rho=False,
params=None):
"""
Compute a hydrostatically balanced density given a potential temperature
profile. By default, this uses a vertically-oriented hybridization
procedure for solving the resulting discrete systems.
:arg state: The :class:`State` object.
:arg theta0: :class:`.Function`containing the potential temperature.
:arg rho0: :class:`.Function` to write the initial density into.
:arg top: If True, set a boundary condition at the top. Otherwise, set
it at the bottom.
:arg pi_boundary: a field or expression to use as boundary data for pi on
the top or bottom as specified.
:arg water_t: the initial total water mixing ratio field.
"""
# Calculate hydrostatic Pi
VDG = state.spaces("DG")
Vv = state.spaces("Vv")
W = MixedFunctionSpace((Vv, VDG))
v, pi = TrialFunctions(W)
dv, dpi = TestFunctions(W)
n = FacetNormal(state.mesh)
cp = state.parameters.cp
# add effect of density of water upon theta
theta = theta0
if water_t is not None:
theta = theta0 / (1 + water_t)
alhs = (
(cp*inner(v, dv) - cp*div(dv*theta)*pi)*dx
+ dpi*div(theta*v)*dx
)
if top:
bmeasure = ds_t
bstring = "bottom"
else:
bmeasure = ds_b
bstring = "top"
arhs = -cp*inner(dv, n)*theta*pi_boundary*bmeasure
# Possibly make g vary with spatial coordinates?
g = state.parameters.g
arhs -= g*inner(dv, state.k)*dx
bcs = [DirichletBC(W.sub(0), Constant(0.0), bstring)]
w = Function(W)
PiProblem = LinearVariationalProblem(alhs, arhs, w, bcs=bcs)
if params is None:
params = {'ksp_type': 'preonly',
'pc_type': 'python',
'mat_type': 'matfree',
'pc_python_type': 'gusto.VerticalHybridizationPC',
# Vertical trace system is only coupled vertically in columns
# block ILU is a direct solver!
'vert_hybridization': {'ksp_type': 'preonly',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}}
PiSolver = LinearVariationalSolver(PiProblem,
solver_parameters=params,
options_prefix="pisolver")
PiSolver.solve()
v, Pi = w.split()
if pi0 is not None:
pi0.assign(Pi)
if solve_for_rho:
w1 = Function(W)
v, rho = w1.split()
rho.interpolate(thermodynamics.rho(state.parameters, theta0, Pi))
v, rho = split(w1)
dv, dpi = TestFunctions(W)
pi = thermodynamics.pi(state.parameters, rho, theta0)
F = (
(cp*inner(v, dv) - cp*div(dv*theta)*pi)*dx
+ dpi*div(theta0*v)*dx
+ cp*inner(dv, n)*theta*pi_boundary*bmeasure
)
F += g*inner(dv, state.k)*dx
rhoproblem = NonlinearVariationalProblem(F, w1, bcs=bcs)
rhosolver = NonlinearVariationalSolver(rhoproblem, solver_parameters=params,
options_prefix="rhosolver")
rhosolver.solve()
v, rho_ = w1.split()
rho0.assign(rho_)
else:
rho0.interpolate(thermodynamics.rho(state.parameters, theta0, Pi))
def unsaturated_hydrostatic_balance(state, theta_d, H, pi0=None,
top=False, pi_boundary=Constant(1.0),
max_outer_solve_count=40,
max_inner_solve_count=20):
"""
Given vertical profiles for dry potential temperature
and relative humidity compute hydrostatically balanced
virtual potential temperature, dry density and water vapour profiles.
The general strategy is to split up the solving into two steps:
1) finding rho to balance the theta profile
2) finding theta_v and r_v to get back theta_d and H
We iteratively solve these steps until we (hopefully)
converge to a solution.
:arg state: The :class:`State` object.
:arg theta_d: The initial dry potential temperature profile.
:arg H: The relative humidity profile.
:arg pi0: Optional function to put exner pressure into.
:arg top: If True, set a boundary condition at the top, otherwise
it will be at the bottom.
:arg pi_boundary: The value of pi on the specified boundary.
:arg max_outer_solve_count: Max number of iterations for outer loop of balance solver.
:arg max_inner_solve_count: Max number of iterations for inner loop of balanace solver.
"""
theta0 = state.fields('theta')
rho0 = state.fields('rho')
water_v0 = state.fields('water_v')
# Calculate hydrostatic Pi
Vt = theta0.function_space()
Vr = rho0.function_space()
R_d = state.parameters.R_d
R_v = state.parameters.R_v
epsilon = R_d / R_v
VDG = state.spaces("DG")
if any(deg > 2 for deg in VDG.ufl_element().degree()):
logger.warning("default quadrature degree most likely not sufficient for this degree element")
# apply first guesses
theta0.assign(theta_d * 1.01)
water_v0.assign(0.01)
if state.vertical_degree == 0:
method = Boundary_Method.physics
else:
method = None
rho_h = Function(Vr)
rho_averaged = Function(Vt)
Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element()))
rho_recoverer = Recoverer(rho0, rho_averaged, VDG=Vt_broken, boundary_method=method)
w_h = Function(Vt)
delta = 1.0
# make expressions for determining water_v0
pie = thermodynamics.pi(state.parameters, rho_averaged, theta0)
p = thermodynamics.p(state.parameters, pie)
T = thermodynamics.T(state.parameters, theta0, pie, water_v0)
r_v_expr = thermodynamics.r_v(state.parameters, H, T, p)
# make expressions to evaluate residual
pi_ev = thermodynamics.pi(state.parameters, rho_averaged, theta0)
p_ev = thermodynamics.p(state.parameters, pi_ev)
T_ev = thermodynamics.T(state.parameters, theta0, pi_ev, water_v0)
RH_ev = thermodynamics.RH(state.parameters, water_v0, T_ev, p_ev)
RH = Function(Vt)
for i in range(max_outer_solve_count):
# solve for rho with theta_vd and w_v guesses
compressible_hydrostatic_balance(state, theta0, rho_h, top=top,
pi_boundary=pi_boundary, water_t=water_v0,
solve_for_rho=True)
# damp solution
rho0.assign(rho0 * (1 - delta) + delta * rho_h)
# calculate averaged rho
rho_recoverer.project()
RH.assign(RH_ev)
if errornorm(RH, H) < 1e-10:
break
# now solve for r_v
for j in range(max_inner_solve_count):
w_h.interpolate(r_v_expr)
water_v0.assign(water_v0 * (1 - delta) + delta * w_h)
# compute theta_vd
theta0.assign(theta_d * (1 + water_v0 / epsilon))
# test quality of solution by re-evaluating expression
RH.assign(RH_ev)
if errornorm(RH, H) < 1e-10:
break
if i == max_outer_solve_count:
raise RuntimeError('Hydrostatic balance solve has not converged within %i' % i, 'iterations')
if pi0 is not None:
pie = thermodynamics.pi(state.parameters, rho0, theta0)
pi0.interpolate(pie)
# do one extra solve for rho
compressible_hydrostatic_balance(state, theta0, rho0, top=top,
pi_boundary=pi_boundary,
water_t=water_v0, solve_for_rho=True)
def saturated_hydrostatic_balance(state, theta_e, water_t, pi0=None,
top=False, pi_boundary=Constant(1.0),
max_outer_solve_count=40,
max_theta_solve_count=5,
max_inner_solve_count=3):
"""
Given a wet equivalent potential temperature, theta_e, and the total moisture
content, water_t, compute a hydrostatically balance virtual potential temperature,
dry density and water vapour profile.
The general strategy is to split up the solving into two steps:
1) finding rho to balance the theta profile
2) finding theta_v and r_v to get back theta_e and saturation
We iteratively solve these steps until we (hopefully)
converge to a solution.
:arg state: The :class:`State` object.
:arg theta_e: The initial wet equivalent potential temperature profile.
:arg water_t: The total water pseudo-mixing ratio profile.
:arg pi0: Optional function to put exner pressure into.
:arg top: If True, set a boundary condition at the top, otherwise
it will be at the bottom.
:arg pi_boundary: The value of pi on the specified boundary.
:arg max_outer_solve_count: Max number of outer iterations for balance solver.
:arg max_theta_solve_count: Max number of iterations for theta solver (middle part of solve).
:arg max_inner_solve_count: Max number of iterations on the inner most
loop for the water vapour solver.
"""
theta0 = state.fields('theta')
rho0 = state.fields('rho')
water_v0 = state.fields('water_v')
# Calculate hydrostatic Pi
Vt = theta0.function_space()
Vr = rho0.function_space()
VDG = state.spaces("DG")
if any(deg > 2 for deg in VDG.ufl_element().degree()):
logger.warning("default quadrature degree most likely not sufficient for this degree element")
theta0.interpolate(theta_e)
water_v0.interpolate(water_t)
if state.vertical_degree == 0:
boundary_method = Boundary_Method.physics
else:
boundary_method = None
rho_h = Function(Vr)
Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element()))
rho_averaged = Function(Vt)
rho_recoverer = Recoverer(rho0, rho_averaged, VDG=Vt_broken, boundary_method=boundary_method)
w_h = Function(Vt)
theta_h = Function(Vt)
theta_e_test = Function(Vt)
delta = 0.8
# expressions for finding theta0 and water_v0 from theta_e and water_t
pie = thermodynamics.pi(state.parameters, rho_averaged, theta0)
p = thermodynamics.p(state.parameters, pie)
T = thermodynamics.T(state.parameters, theta0, pie, water_v0)
r_v_expr = thermodynamics.r_sat(state.parameters, T, p)
theta_e_expr = thermodynamics.theta_e(state.parameters, T, p, water_v0, water_t)
for i in range(max_outer_solve_count):
# solve for rho with theta_vd and w_v guesses
compressible_hydrostatic_balance(state, theta0, rho_h, top=top,
pi_boundary=pi_boundary, water_t=water_t,
solve_for_rho=True)
# damp solution
rho0.assign(rho0 * (1 - delta) + delta * rho_h)
theta_e_test.assign(theta_e_expr)
if errornorm(theta_e_test, theta_e) < 1e-8:
break
# calculate averaged rho
rho_recoverer.project()
# now solve for r_v
for j in range(max_theta_solve_count):
theta_h.interpolate(theta_e / theta_e_expr * theta0)
theta0.assign(theta0 * (1 - delta) + delta * theta_h)
# break when close enough
if errornorm(theta_e_test, theta_e) < 1e-6:
break
for k in range(max_inner_solve_count):
w_h.interpolate(r_v_expr)
water_v0.assign(water_v0 * (1 - delta) + delta * w_h)
# break when close enough
theta_e_test.assign(theta_e_expr)
if errornorm(theta_e_test, theta_e) < 1e-6:
break
if i == max_outer_solve_count:
raise RuntimeError('Hydrostatic balance solve has not converged within %i' % i, 'iterations')
if pi0 is not None:
pie = thermodynamics.pi(state.parameters, rho0, theta0)
pi0.interpolate(pie)
# do one extra solve for rho
compressible_hydrostatic_balance(state, theta0, rho0, top=top,
pi_boundary=pi_boundary,
water_t=water_t, solve_for_rho=True)
def _setup_solver(self):
state = self.state # just cutting down line length a bit
Dt = state.timestepping.dt
beta_ = Dt*state.timestepping.alpha
cp = state.parameters.cp
mu = state.mu
Vu = state.spaces("HDiv")
Vtheta = state.spaces("HDiv_v")
Vrho = state.spaces("DG")
# Store time-stepping coefficients as UFL Constants
dt = Constant(Dt)
beta = Constant(beta_)
beta_cp = Constant(beta_ * cp)
# Split up the rhs vector (symbolically)
u_in, rho_in, theta_in = split(state.xrhs)
# Build the reduced function space for u,rho
M = MixedFunctionSpace((Vu, Vrho))
w, phi = TestFunctions(M)
u, rho = TrialFunctions(M)
n = FacetNormal(state.mesh)
# Get background fields
thetabar = state.fields("thetabar")
rhobar = state.fields("rhobar")
pibar = thermodynamics.pi(state.parameters, rhobar, thetabar)
pibar_rho = thermodynamics.pi_rho(state.parameters, rhobar, thetabar)
pibar_theta = thermodynamics.pi_theta(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 pi') in the vertical
# component of the gradient
# the pi prime term (here, bars are for mean and no bars are
# for linear perturbations)
pi = pibar_theta*theta + pibar_rho*rho
# vertical projection
def V(u):
return 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))
# 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
eqn = (
inner(w, (state.h_project(u) - u_in))*dx
- beta_cp*div(theta_w*V(w))*pibar*dxp
# following does nothing but is preserved in the comments
# to remind us why (because V(w) is purely vertical).
# + beta_cp*jump(theta*V(w), n)*avg(pibar)*dS_v
- beta_cp*div(thetabar_w*w)*pi*dxp
+ beta_cp*jump(thetabar_w*w, n)*avg(pi)*dS_vp
+ (phi*(rho - rho_in) - beta*inner(grad(phi), u)*rhobar)*dx
+ beta*jump(phi*u, n)*avg(rhobar)*(dS_v + dS_h)
)
if mu is not None:
eqn += dt*mu*inner(w, k)*inner(u, k)*dx
aeqn = lhs(eqn)
Leqn = rhs(eqn)
# Place to put result of u rho solver
self.urho = Function(M)
# Boundary conditions (assumes extruded mesh)
bcs = [DirichletBC(M.sub(0), 0.0, "bottom"),
DirichletBC(M.sub(0), 0.0, "top")]
# Solver for u, rho
urho_problem = LinearVariationalProblem(
aeqn, Leqn, self.urho, bcs=bcs)
self.urho_solver = LinearVariationalSolver(urho_problem,
solver_parameters=self.solver_parameters,
options_prefix='ImplicitSolver')
# Reconstruction of theta
theta = TrialFunction(Vtheta)
gamma = TestFunction(Vtheta)
u, rho = self.urho.split()
self.theta = Function(Vtheta)
theta_eqn = gamma*(theta - theta_in
+ dot(k, u)*dot(k, grad(thetabar))*beta)*dx
theta_problem = LinearVariationalProblem(lhs(theta_eqn),
rhs(theta_eqn),
self.theta)
self.theta_solver = LinearVariationalSolver(theta_problem,
options_prefix='thetabacksubstitution')
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))
# set theta0
theta0.interpolate(theta_b + theta_pert)
# calculate hydrostatic Pi
rho_b = Function(Vr)
compressible_hydrostatic_balance(state, theta_b, rho_b)
compressible_hydrostatic_balance(state, theta0, rho0)
# set Pi0
Pi0 = calculate_Pi0(state, theta0, rho0)
state.parameters.Pi0 = Pi0
# set x component of velocity
cp = state.parameters.cp
dthetady = state.parameters.dthetady
Pi = thermodynamics.pi(state.parameters, rho0, theta0)
u = cp * dthetady / f * (Pi - Pi0)
# set y component of velocity
v = Function(Vr).assign(0.)
compressible_eady_initial_v(state, theta0, rho0, v)
# set initial u
u_exp = as_vector([u, v, 0.])
u0.project(u_exp)
# pass these initial conditions to the state.initialise method
state.initialise([('u', u0), ('rho', rho0), ('theta', theta0)])
# set the background profiles
state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])
def _setup_solver(self):
from firedrake.assemble import create_assembly_callable
import numpy as np
state = self.state
dt = state.timestepping.dt
beta = dt*state.timestepping.alpha
cp = state.parameters.cp
mu = state.mu
Vu = state.spaces("HDiv")
Vu_broken = FunctionSpace(state.mesh, BrokenElement(Vu.ufl_element()))
Vtheta = state.spaces("HDiv_v")
Vrho = state.spaces("DG")
h_deg = state.horizontal_degree
v_deg = state.vertical_degree
Vtrace = FunctionSpace(state.mesh, "HDiv Trace", degree=(h_deg, v_deg))
# Split up the rhs vector (symbolically)
u_in, rho_in, theta_in = split(state.xrhs)
# Build the function space for "broken" u and rho
# and add the trace variable
M = MixedFunctionSpace((Vu_broken, Vrho))
w, phi = TestFunctions(M)
u, rho = TrialFunctions(M)
l0 = TrialFunction(Vtrace)
dl = TestFunction(Vtrace)
n = FacetNormal(state.mesh)
# Get background fields
thetabar = state.fields("thetabar")
rhobar = state.fields("rhobar")
pibar = thermodynamics.pi(state.parameters, rhobar, thetabar)
pibar_rho = thermodynamics.pi_rho(state.parameters, rhobar, thetabar)
pibar_theta = thermodynamics.pi_theta(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 pi') in the vertical
# component of the gradient
# The pi prime term (here, bars are for mean and no bars are
# for linear perturbations)
pi = pibar_theta*theta + pibar_rho*rho
# Vertical projection
def V(u):
return 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))
# Mass matrix for the trace space
tM = assemble(dl('+')*l0('+')*(dS_v + dS_h)
+ dl*l0*ds_v + dl*l0*(ds_t + ds_b))
Lrhobar = Function(Vtrace)
Lpibar = Function(Vtrace)
rhopi_solver = LinearSolver(tM, solver_parameters={'ksp_type': 'cg',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'},
options_prefix='rhobarpibar_solver')
rhobar_avg = Function(Vtrace)
pibar_avg = Function(Vtrace)
def _traceRHS(f):
return (dl('+')*avg(f)*(dS_v + dS_h)
+ dl*f*ds_v + dl*f*(ds_t + ds_b))
assemble(_traceRHS(rhobar), tensor=Lrhobar)
assemble(_traceRHS(pibar), tensor=Lpibar)
# Project averages of coefficients into the trace space
with timed_region("Gusto:HybridProjectRhobar"):
rhopi_solver.solve(rhobar_avg, Lrhobar)
with timed_region("Gusto:HybridProjectPibar"):
rhopi_solver.solve(pibar_avg, Lpibar)
# 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
# "broken" u and rho system
Aeqn = (inner(w, (state.h_project(u) - u_in))*dx
- beta*cp*div(theta_w*V(w))*pibar*dxp
# following does nothing but is preserved in the comments
# to remind us why (because V(w) is purely vertical).
# + beta*cp*dot(theta_w*V(w), n)*pibar_avg('+')*dS_vp
+ beta*cp*dot(theta_w*V(w), n)*pibar_avg('+')*dS_hp
+ beta*cp*dot(theta_w*V(w), n)*pibar_avg*ds_tbp
- beta*cp*div(thetabar_w*w)*pi*dxp
+ (phi*(rho - rho_in) - beta*inner(grad(phi), u)*rhobar)*dx
+ beta*dot(phi*u, n)*rhobar_avg('+')*(dS_v + dS_h))
if mu is not None:
Aeqn += dt*mu*inner(w, k)*inner(u, k)*dx
# Form the mixed operators using Slate
# (A K)(X) = (X_r)
# (K.T 0)(l) (0 )
# where X = ("broken" u, rho)
A = Tensor(lhs(Aeqn))
X_r = Tensor(rhs(Aeqn))
# Off-diagonal block matrices containing the contributions
# of the Lagrange multipliers (surface terms in the momentum equation)
K = Tensor(beta*cp*dot(thetabar_w*w, n)*l0('+')*(dS_vp + dS_hp)
+ beta*cp*dot(thetabar_w*w, n)*l0*ds_vp
+ beta*cp*dot(thetabar_w*w, n)*l0*ds_tbp)
# X = A.inv * (X_r - K * l),
# 0 = K.T * X = -(K.T * A.inv * K) * l + K.T * A.inv * X_r,
# so (K.T * A.inv * K) * l = K.T * A.inv * X_r
# is the reduced equation for the Lagrange multipliers.
# Right-hand side expression: (Forward substitution)
Rexp = K.T * A.inv * X_r
self.R = Function(Vtrace)
# We need to rebuild R everytime data changes
self._assemble_Rexp = create_assembly_callable(Rexp, tensor=self.R)
# Schur complement operator:
Smatexp = K.T * A.inv * K
with timed_region("Gusto:HybridAssembleTraceOp"):
S = assemble(Smatexp)
S.force_evaluation()
# Set up the Linear solver for the system of Lagrange multipliers
self.lSolver = LinearSolver(S, solver_parameters=self.solver_parameters,
options_prefix='lambda_solve')
# Result function for the multiplier solution
self.lambdar = Function(Vtrace)
# Place to put result of u rho reconstruction
self.urho = Function(M)
# Reconstruction of broken u and rho
u_, rho_ = self.urho.split()
# Split operators for two-stage reconstruction
_A = A.blocks
_K = K.blocks
_Xr = X_r.blocks
A00 = _A[0, 0]
A01 = _A[0, 1]
A10 = _A[1, 0]
A11 = _A[1, 1]
K0 = _K[0, 0]
Ru = _Xr[0]
Rrho = _Xr[1]
lambda_vec = AssembledVector(self.lambdar)
# rho reconstruction
Srho = A11 - A10 * A00.inv * A01
rho_expr = Srho.solve(Rrho - A10 * A00.inv * (Ru - K0 * lambda_vec),
decomposition="PartialPivLU")
self._assemble_rho = create_assembly_callable(rho_expr, tensor=rho_)
# "broken" u reconstruction
rho_vec = AssembledVector(rho_)
u_expr = A00.solve(Ru - A01 * rho_vec - K0 * lambda_vec,
decomposition="PartialPivLU")
self._assemble_u = create_assembly_callable(u_expr, tensor=u_)
# Project broken u into the HDiv space using facet averaging.
# Weight function counting the dofs of the HDiv element:
shapes = (Vu.finat_element.space_dimension(), np.prod(Vu.shape))
weight_kernel = """
for (int i=0; i<%d; ++i) {
for (int j=0; j<%d; ++j) {
w[i][j] += 1.0;
}}""" % shapes
self._weight = Function(Vu)
par_loop(weight_kernel, dx, {"w": (self._weight, INC)})
# Averaging kernel
self._average_kernel = """
for (int i=0; i<%d; ++i) {
for (int j=0; j<%d; ++j) {
vec_out[i][j] += vec_in[i][j]/w[i][j];
}}""" % shapes
# 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={'ksp_type': 'cg',
'pc_type': 'bjacobi',
'pc_sub_type': 'ilu'},
options_prefix='thetabacksubstitution')
self.bcs = [DirichletBC(Vu, 0.0, "bottom"),
DirichletBC(Vu, 0.0, "top")]
