def __init__(self, mesh, conditions, timestepping, params, output, solver_params):
self.timestepping = timestepping
self.timestep = timestepping.timestep
self.timescale = timestepping.timescale
self.params = params
if output is None:
raise RuntimeError("You must provide a directory name for dumping results")
else:
self.output = output
self.outfile = File(output.dirname)
self.dump_count = 0
self.dump_freq = output.dumpfreq
self.solver_params = solver_params
self.mesh = mesh
self.conditions = conditions
if conditions.steady_state == True:
self.ind = 1
else:
self.ind = 1
family = conditions.family
self.x, self.y = SpatialCoordinate(mesh)
self.n = FacetNormal(mesh)
self.V = VectorFunctionSpace(mesh, family, conditions.order + 1)
self.U = FunctionSpace(mesh, family, conditions.order + 1)
self.U1 = FunctionSpace(mesh, 'DG', conditions.order)
self.S = TensorFunctionSpace(mesh, 'DG', conditions.order)
self.D = FunctionSpace(mesh, 'DG', 0)
self.W1 = MixedFunctionSpace([self.V, self.S])
self.W2 = MixedFunctionSpace([self.V, self.U1, self.U1])
self.W3 = MixedFunctionSpace([self.V, self.S, self.U1, self.U1])
def _build_spaces(self, mesh, vertical_degree, horizontal_degree, family):
"""
Build:
velocity space self.V2,
pressure space self.V3,
temperature space self.Vt,
mixed function space self.W = (V2,V3,Vt)
"""
self.spaces = SpaceCreator()
if vertical_degree is not None:
# horizontal base spaces
cell = mesh._base_mesh.ufl_cell().cellname()
S1 = FiniteElement(family, cell, horizontal_degree+1)
S2 = FiniteElement("DG", cell, horizontal_degree)
# vertical base spaces
T0 = FiniteElement("CG", interval, vertical_degree+1)
T1 = FiniteElement("DG", interval, vertical_degree)
# build spaces V2, V3, Vt
V2h_elt = HDiv(TensorProductElement(S1, T1))
V2t_elt = TensorProductElement(S2, T0)
V3_elt = TensorProductElement(S2, T1)
V2v_elt = HDiv(V2t_elt)
V2_elt = V2h_elt + V2v_elt
V0 = self.spaces("HDiv", mesh, V2_elt)
V1 = self.spaces("DG", mesh, V3_elt)
V2 = self.spaces("HDiv_v", mesh, V2t_elt)
self.Vv = self.spaces("Vv", mesh, V2v_elt)
self.W = MixedFunctionSpace((V0, V1, V2))
else:
cell = mesh.ufl_cell().cellname()
V1_elt = FiniteElement(family, cell, horizontal_degree+1)
V0 = self.spaces("HDiv", mesh, V1_elt)
V1 = self.spaces("DG", mesh, "DG", horizontal_degree)
self.W = MixedFunctionSpace((V0, V1))
def argument(self, o):
from ufl import split
from firedrake import MixedFunctionSpace, FunctionSpace
V = o.function_space()
if len(V) == 1:
# Not on a mixed space, just return ourselves.
return o
if o in self._arg_cache:
return self._arg_cache[o]
V_is = V.split()
indices = self.blocks[o.number()]
try:
indices = tuple(indices)
nidx = len(indices)
except TypeError:
# Only one index provided.
indices = (indices, )
nidx = 1
if nidx == 1:
W = V_is[indices[0]]
W = FunctionSpace(W.mesh(), W.ufl_element())
a = (Argument(W, o.number(), part=o.part()), )
else:
W = MixedFunctionSpace([V_is[i] for i in indices])
a = split(Argument(W, o.number(), part=o.part()))
args = []
for i in range(len(V_is)):
if i in indices:
c = indices.index(i)
a_ = a[c]
if len(a_.ufl_shape) == 0:
args += [a_]
else:
args += [a_[j] for j in numpy.ndindex(a_.ufl_shape)]
else:
args += [
Zero()
for j in numpy.ndindex(V_is[i].ufl_element().value_shape())
]
return self._arg_cache.setdefault(o, as_vector(args))
def __init__(self,
field_names,
state,
family,
degree,
terms_to_linearise=None,
no_normal_flow_bc_ids=None,
active_tracers=None):
self.field_names = field_names
self.active_tracers = active_tracers
self.terms_to_linearise = {} if terms_to_linearise is None else terms_to_linearise
# Build finite element spaces
self.spaces = [
space for space in self._build_spaces(state, family, degree)
]
# Add active tracers to the list of prognostics
if active_tracers is None:
active_tracers = []
self.add_tracers_to_prognostics(state, active_tracers)
# Make the full mixed function space
W = MixedFunctionSpace(self.spaces)
# Can now call the underlying PrognosticEquation
full_field_name = "_".join(self.field_names)
super().__init__(state, W, full_field_name)
# Set up test functions, trials and prognostics
self.tests = TestFunctions(W)
self.trials = TrialFunctions(W)
self.X = Function(W)
# Set up no-normal-flow boundary conditions
if no_normal_flow_bc_ids is None:
no_normal_flow_bc_ids = []
self.set_no_normal_flow_bcs(state, no_normal_flow_bc_ids)
def _setup_solver(self):
state = self.state # just cutting down line length a bit
dt = state.timestepping.dt
beta = dt*state.timestepping.alpha
mu = state.mu
# Split up the rhs vector (symbolically)
u_in, p_in, b_in = split(state.xrhs)
# Build the reduced function space for u,p
M = MixedFunctionSpace((state.V[0], state.V[1]))
w, phi = TestFunctions(M)
u, p = TrialFunctions(M)
# Get background fields
bbar = state.bbar
# Analytical (approximate) elimination of theta
k = state.k # Upward pointing unit vector
b = -dot(k,u)*dot(k,grad(bbar))*beta + b_in
# vertical projection
def V(u):
return k*inner(u,k)
eqn = (
inner(w, (u - u_in))*dx
- beta*div(w)*p*dx
- beta*inner(w,k)*b*dx
+ phi*div(u)*dx
)
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 p solver
self.up = Function(M)
# Boundary conditions (assumes extruded mesh)
dim = M.sub(0).ufl_element().value_shape()[0]
bc = ("0.0",)*dim
bcs = [DirichletBC(M.sub(0), Expression(bc), "bottom"),
DirichletBC(M.sub(0), Expression(bc), "top")]
# preconditioner equation
L = self.L
Ap = (
inner(w,u) + L*L*div(w)*div(u) +
phi*p/L/L
)*dx
# Solver for u, p
up_problem = LinearVariationalProblem(
aeqn, Leqn, self.up, bcs=bcs, aP=Ap)
nullspace = MixedVectorSpaceBasis(M,
[M.sub(0),
VectorSpaceBasis(constant=True)])
self.up_solver = LinearVariationalSolver(up_problem,
solver_parameters=self.params,
nullspace=nullspace)
# Reconstruction of b
b = TrialFunction(state.V[2])
gamma = TestFunction(state.V[2])
u, p = self.up.split()
self.b = Function(state.V[2])
b_eqn = gamma*(b - b_in +
dot(k,u)*dot(k,grad(bbar))*beta)*dx
b_problem = LinearVariationalProblem(lhs(b_eqn),
rhs(b_eqn),
self.b)
self.b_solver = LinearVariationalSolver(b_problem)
cartesian_u_expr = -u_zonal * sin(lamda) - u_merid * sin(theta) * cos(
lamda)
cartesian_v_expr = u_zonal * cos(lamda) - u_merid * sin(theta) * sin(lamda)
cartesian_w_expr = u_merid * cos(theta)
return as_vector((cartesian_u_expr, cartesian_v_expr, cartesian_w_expr))
# Build function spaces
degree = 2
family = ("DG", "BDM", "CG")
W0 = FunctionSpace(mesh, family[0], degree - 1, family[0])
W1_elt = FiniteElement(family[1], triangle, degree)
W1 = FunctionSpace(mesh, W1_elt, name="HDiv")
W2 = FunctionSpace(mesh, family[2], degree + 1)
M = MixedFunctionSpace((W1, W0))
# Set up initial condition parameters
u_max, D_bump, D_mean = 80, 120, 10000
theta_0, theta_1, theta_2 = pi / 7., 5 * pi / 14., pi / 4.
a, b = 1 / 3., 1 / 15.
e_n = exp(-4 / (theta_1 - theta_0)**2)
# uexpr
u_zonal_expr = (u_max / e_n) * exp(1 / ((theta - theta_0) * (theta - theta_1)))
u_zonal = conditional(ge(theta, theta_0),
conditional(le(theta, theta_1), u_zonal_expr, 0.), 0.)
u_merid = 0.0
# get cartesian components of velocity
uexpr = sphere_to_cartesian(mesh, u_zonal, u_merid)
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),
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.
"""
# Calculate hydrostatic Pi
W = MixedFunctionSpace((state.Vv,state.V[1]))
v, pi = TrialFunctions(W)
dv, dpi = TestFunctions(W)
n = FacetNormal(state.mesh)
cp = state.parameters.cp
alhs = (
(cp*inner(v,dv) - cp*div(dv*theta0)*pi)*dx
+ dpi*div(theta0*v)*dx
)
if top:
bmeasure = ds_t
bstring = "bottom"
else:
bmeasure = ds_b
bstring = "top"
arhs = -cp*inner(dv,n)*theta0*pi_boundary*bmeasure
if state.parameters.geopotential:
Phi = state.Phi
arhs += div(dv)*Phi*dx - inner(dv,n)*Phi*bmeasure
else:
g = state.parameters.g
arhs -= g*inner(dv,state.k)*dx
if(state.mesh.geometric_dimension() == 2):
bcs = [DirichletBC(W.sub(0), Expression(("0.", "0.")), bstring)]
elif(state.mesh.geometric_dimension() == 3):
bcs = [DirichletBC(W.sub(0), Expression(("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_estimate_eigenvalues': True,
'fieldsplit_1_mg_levels_ksp_chebyshev_estimate_eigenvalues_random': 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)
kappa = state.parameters.kappa
R_d = state.parameters.R_d
p_0 = state.parameters.p_0
if solve_for_rho:
w1 = Function(W)
v, rho = w1.split()
rho.interpolate(p_0*(Pi**((1-kappa)/kappa))/R_d/theta0)
v, rho = split(w1)
dv, dpi = TestFunctions(W)
pi = ((R_d/p_0)*rho*theta0)**(kappa/(1.-kappa))
F = (
(cp*inner(v,dv) - cp*div(dv*theta0)*pi)*dx
+ dpi*div(theta0*v)*dx
+ cp*inner(dv,n)*theta0*pi_boundary*bmeasure
)
if state.parameters.geopotential:
F += - div(dv)*Phi*dx + inner(dv,n)*Phi*bmeasure
else:
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(p_0*(Pi**((1-kappa)/kappa))/R_d/theta0)
def _setup_solver(self):
state = self.state # just cutting down line length a bit
dt = state.dt
beta_ = dt * self.alpha
Vu = state.spaces("HDiv")
Vb = state.spaces("theta")
Vp = state.spaces("DG")
# Store time-stepping coefficients as UFL Constants
beta = Constant(beta_)
# Split up the rhs vector (symbolically)
self.xrhs = Function(self.equations.function_space)
u_in, p_in, b_in = split(self.xrhs)
# Build the reduced function space for u,p
M = MixedFunctionSpace((Vu, Vp))
w, phi = TestFunctions(M)
u, p = TrialFunctions(M)
# Get background fields
bbar = state.fields("bbar")
# Analytical (approximate) elimination of theta
k = state.k # Upward pointing unit vector
b = -dot(k, u) * dot(k, grad(bbar)) * beta + b_in
# vertical projection
def V(u):
return k * inner(u, k)
eqn = (inner(w, (u - u_in)) * dx - beta * div(w) * p * dx -
beta * inner(w, k) * b * dx + phi * div(u) * dx)
if hasattr(self.equations, "mu"):
eqn += dt * self.equations.mu * inner(w, k) * inner(u, k) * dx
aeqn = lhs(eqn)
Leqn = rhs(eqn)
# Place to put result of u p solver
self.up = Function(M)
# Boundary conditions (assumes extruded mesh)
# BCs are declared for the plain velocity space. As we need them in
# a mixed problem, we replicate the BCs but for subspace of M
bcs = [
DirichletBC(M.sub(0), bc.function_arg, bc.sub_domain)
for bc in self.equations.bcs['u']
]
# Solver for u, p
up_problem = LinearVariationalProblem(aeqn, Leqn, self.up, bcs=bcs)
# Provide callback for the nullspace of the trace system
def trace_nullsp(T):
return VectorSpaceBasis(constant=True)
appctx = {"trace_nullspace": trace_nullsp}
self.up_solver = LinearVariationalSolver(
up_problem,
solver_parameters=self.solver_parameters,
appctx=appctx)
# Reconstruction of b
b = TrialFunction(Vb)
gamma = TestFunction(Vb)
u, p = self.up.split()
self.b = Function(Vb)
b_eqn = gamma * (b - b_in + dot(k, u) * dot(k, grad(bbar)) * beta) * dx
b_problem = LinearVariationalProblem(lhs(b_eqn), rhs(b_eqn), self.b)
self.b_solver = LinearVariationalSolver(b_problem)
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 compressible_hydrostatic_balance(state,
theta0,
rho0,
exner0=None,
top=False,
exner_boundary=Constant(1.0),
mr_t=None,
solve_for_rho=False,
params=None):
"""
Compute a hydrostatically balanced density given a potential temperature
profile. By default, this uses a vertically-oriented hybridization
procedure for solving the resulting discrete systems.
:arg state: The :class:`State` object.
:arg theta0: :class:`.Function`containing the potential temperature.
:arg rho0: :class:`.Function` to write the initial density into.
:arg top: If True, set a boundary condition at the top. Otherwise, set
it at the bottom.
:arg exner_boundary: a field or expression to use as boundary data for exner
on the top or bottom as specified.
:arg mr_t: the initial total water mixing ratio field.
"""
# Calculate hydrostatic Pi
VDG = state.spaces("DG")
Vu = state.spaces("HDiv")
Vv = FunctionSpace(state.mesh, Vu.ufl_element()._elements[-1])
W = MixedFunctionSpace((Vv, VDG))
v, exner = TrialFunctions(W)
dv, dexner = TestFunctions(W)
n = FacetNormal(state.mesh)
cp = state.parameters.cp
# add effect of density of water upon theta
theta = theta0
if mr_t is not None:
theta = theta0 / (1 + mr_t)
alhs = ((cp * inner(v, dv) - cp * div(dv * theta) * exner) * dx +
dexner * div(theta * v) * dx)
if top:
bmeasure = ds_t
bstring = "bottom"
else:
bmeasure = ds_b
bstring = "top"
arhs = -cp * inner(dv, n) * theta * exner_boundary * bmeasure
# Possibly make g vary with spatial coordinates?
g = state.parameters.g
arhs -= g * inner(dv, state.k) * dx
bcs = [DirichletBC(W.sub(0), zero(), bstring)]
w = Function(W)
exner_problem = LinearVariationalProblem(alhs, arhs, w, bcs=bcs)
if params is None:
params = {
'ksp_type': 'preonly',
'pc_type': 'python',
'mat_type': 'matfree',
'pc_python_type': 'gusto.VerticalHybridizationPC',
# Vertical trace system is only coupled vertically in columns
# block ILU is a direct solver!
'vert_hybridization': {
'ksp_type': 'preonly',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
}
}
exner_solver = LinearVariationalSolver(exner_problem,
solver_parameters=params,
options_prefix="exner_solver")
exner_solver.solve()
v, exner = w.split()
if exner0 is not None:
exner0.assign(exner)
if solve_for_rho:
w1 = Function(W)
v, rho = w1.split()
rho.interpolate(thermodynamics.rho(state.parameters, theta0, exner))
v, rho = split(w1)
dv, dexner = TestFunctions(W)
exner = thermodynamics.exner_pressure(state.parameters, rho, theta0)
F = ((cp * inner(v, dv) - cp * div(dv * theta) * exner) * dx +
dexner * div(theta0 * v) * dx +
cp * inner(dv, n) * theta * exner_boundary * bmeasure)
F += g * inner(dv, state.k) * dx
rhoproblem = NonlinearVariationalProblem(F, w1, bcs=bcs)
rhosolver = NonlinearVariationalSolver(rhoproblem,
solver_parameters=params,
options_prefix="rhosolver")
rhosolver.solve()
v, rho_ = w1.split()
rho0.assign(rho_)
else:
rho0.interpolate(thermodynamics.rho(state.parameters, theta0, exner))
def test_replace_subject(subject_type, replacement_type):
# ------------------------------------------------------------------------ #
# Only certain combinations of options are valid
# ------------------------------------------------------------------------ #
if subject_type == 'vector' and replacement_type != 'vector':
return True
elif replacement_type == 'vector' and subject_type != 'vector':
return True
# ------------------------------------------------------------------------ #
# Set up
# ------------------------------------------------------------------------ #
# Some basic labels
foo_label = Label("foo")
bar_label = Label("bar")
# Create mesh, function space and forms
n = 3
mesh = UnitSquareMesh(n, n)
V0 = FunctionSpace(mesh, "DG", 0)
V1 = FunctionSpace(mesh, "CG", 1)
V2 = VectorFunctionSpace(mesh, "DG", 0)
Vmixed = MixedFunctionSpace((V0, V1))
idx = None
# ------------------------------------------------------------------------ #
# Choose subject
# ------------------------------------------------------------------------ #
if subject_type == 'normal':
V = V0
elif subject_type == 'mixed':
V = Vmixed
if replacement_type == 'normal':
idx = 0
elif subject_type == 'vector':
V = V2
else:
raise ValueError
the_subject = Function(V)
not_subject = Function(V)
test = TestFunction(V)
form_1 = inner(the_subject, test)*dx
form_2 = inner(not_subject, test)*dx
term_1 = foo_label(subject(form_1, the_subject))
term_2 = bar_label(form_2)
labelled_form = term_1 + term_2
# ------------------------------------------------------------------------ #
# Choose replacement
# ------------------------------------------------------------------------ #
if replacement_type == 'normal':
V = V1
elif replacement_type == 'mixed':
V = Vmixed
if subject_type != 'mixed':
idx = 0
elif replacement_type == 'vector':
V = V2
elif replacement_type == 'tuple':
V = Vmixed
else:
raise ValueError
the_replacement = Function(V)
if replacement_type == 'tuple':
the_replacement = TrialFunctions(Vmixed)
if subject_type == 'normal':
idx = 0
# ------------------------------------------------------------------------ #
# Test replace_subject
# ------------------------------------------------------------------------ #
labelled_form = labelled_form.label_map(
lambda t: t.has_label(subject),
map_if_true=replace_subject(the_replacement, idx=idx)
)
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
# 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((state.V[0], state.V[1]))
w, phi = TestFunctions(M)
u, rho = TrialFunctions(M)
n = FacetNormal(state.mesh)
# Get background fields
thetabar = state.thetabar
rhobar = state.rhobar
pibar = exner(thetabar, rhobar, state)
pibar_rho = exner_rho(thetabar, rhobar, state)
pibar_theta = exner_theta(thetabar, rhobar, state)
# 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)
eqn = (
inner(w, (u - u_in))*dx
- beta*cp*div(theta*V(w))*pibar*dx
# 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)*pi*dx
+ beta*cp*jump(thetabar*w,n)*avg(pi)*dS_v
+ (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)
dim = M.sub(0).ufl_element().value_shape()[0]
bc = ("0.0",)*dim
bcs = [DirichletBC(M.sub(0), Expression(bc), "bottom"),
DirichletBC(M.sub(0), Expression(bc), "top")]
# Solver for u, rho
urho_problem = LinearVariationalProblem(
aeqn, Leqn, self.urho, bcs=bcs)
self.urho_solver = LinearVariationalSolver(urho_problem,
solver_parameters=self.params,
options_prefix='ImplicitSolver')
# Reconstruction of theta
theta = TrialFunction(state.V[2])
gamma = TestFunction(state.V[2])
u, rho = self.urho.split()
self.theta = Function(state.V[2])
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))
def moist_hydrostatic_balance(state, theta_e, water_t, pi_boundary=Constant(1.0)):
"""
Given a wet equivalent potential temperature, theta_e, and the total moisture
content, water_t, compute a hydrostatically balance virtual potential temperature,
dry density and water vapour profile.
:arg state: The :class:`State` object.
:arg theta_e: The initial wet equivalent potential temperature profile.
:arg water_t: The total water pseudo-mixing ratio profile.
:arg pi_boundary: the value of pi on the lower boundary of the domain.
"""
theta0 = state.fields('theta')
rho0 = state.fields('rho')
water_v0 = state.fields('water_v')
# Calculate hydrostatic Pi
Vt = theta0.function_space()
Vr = rho0.function_space()
Vv = state.fields('u').function_space()
n = FacetNormal(state.mesh)
g = state.parameters.g
cp = state.parameters.cp
R_d = state.parameters.R_d
p_0 = state.parameters.p_0
VDG = state.spaces("DG")
if any(deg > 2 for deg in VDG.ufl_element().degree()):
state.logger.warning("default quadrature degree most likely not sufficient for this degree element")
quadrature_degree = (5, 5)
params = {'ksp_type': 'preonly',
'ksp_monitor_true_residual': True,
'ksp_converged_reason': True,
'snes_converged_reason': True,
'ksp_max_it': 100,
'mat_type': 'aij',
'pc_type': 'lu',
'pc_factor_mat_solver_type': 'mumps'}
theta0.interpolate(theta_e)
water_v0.interpolate(water_t)
Pi = Function(Vr)
epsilon = 0.9 # relaxation constant
# set up mixed space
Z = MixedFunctionSpace((Vt, Vt))
z = Function(Z)
gamma, phi = TestFunctions(Z)
theta_v, w_v = z.split()
# give first guesses for trial functions
theta_v.assign(theta0)
w_v.assign(water_v0)
theta_v, w_v = split(z)
# define variables
T = thermodynamics.T(state.parameters, theta_v, Pi, r_v=w_v)
p = thermodynamics.p(state.parameters, Pi)
w_sat = thermodynamics.r_sat(state.parameters, T, p)
dxp = dx(degree=(quadrature_degree))
# set up weak form of theta_e and w_sat equations
F = (-gamma * theta_e * dxp
+ gamma * thermodynamics.theta_e(state.parameters, T, p, w_v, water_t) * dxp
- phi * w_v * dxp
+ phi * w_sat * dxp)
problem = NonlinearVariationalProblem(F, z)
solver = NonlinearVariationalSolver(problem, solver_parameters=params)
theta_v, w_v = z.split()
Pi_h = Function(Vr).interpolate((p / p_0) ** (R_d / cp))
# solve for Pi with theta_v and w_v constant
# then solve for theta_v and w_v with Pi constant
for i in range(5):
compressible_hydrostatic_balance(state, theta0, rho0, pi0=Pi_h, water_t=water_t)
Pi.assign(Pi * (1 - epsilon) + epsilon * Pi_h)
solver.solve()
theta0.assign(theta0 * (1 - epsilon) + epsilon * theta_v)
water_v0.assign(water_v0 * (1 - epsilon) + epsilon * w_v)
# now begin on Newton solver, setup up new mixed space
Z = MixedFunctionSpace((Vt, Vt, Vr, Vv))
z = Function(Z)
gamma, phi, psi, w = TestFunctions(Z)
theta_v, w_v, pi, v = z.split()
# use previous values as first guesses for newton solver
theta_v.assign(theta0)
w_v.assign(water_v0)
pi.assign(Pi)
theta_v, w_v, pi, v = split(z)
# define variables
T = thermodynamics.T(state.parameters, theta_v, pi, r_v=w_v)
p = thermodynamics.p(state.parameters, pi)
w_sat = thermodynamics.r_sat(state.parameters, T, p)
F = (-gamma * theta_e * dxp
+ gamma * thermodynamics.theta_e(state.parameters, T, p, w_v, water_t) * dxp
- phi * w_v * dxp
+ phi * w_sat * dxp
+ cp * inner(v, w) * dxp
- cp * div(w * theta_v / (1.0 + water_t)) * pi * dxp
+ psi * div(theta_v * v / (1.0 + water_t)) * dxp
+ cp * inner(w, n) * pi_boundary * theta_v / (1.0 + water_t) * ds_b
+ g * inner(w, state.k) * dxp)
bcs = [DirichletBC(Z.sub(3), 0.0, "top")]
problem = NonlinearVariationalProblem(F, z, bcs=bcs)
solver = NonlinearVariationalSolver(problem, solver_parameters=params)
solver.solve()
theta_v, w_v, pi, v = z.split()
# assign final values
theta0.assign(theta_v)
water_v0.assign(w_v)
# find rho
compressible_hydrostatic_balance(state, theta0, rho0, water_t=water_t, solve_for_rho=True)
def __init__(self, prognostic_variables, simulation_parameters):
mesh = simulation_parameters['mesh'][-1]
self.scheme = simulation_parameters['scheme'][-1]
self.timestepping = simulation_parameters['timestepping'][-1]
alphasq = simulation_parameters['alphasq'][-1]
c0 = simulation_parameters['c0'][-1]
gamma = simulation_parameters['gamma'][-1]
Dt = Constant(simulation_parameters['dt'][-1])
self.solvers = []
if alphasq.values()[0] > 0.0 and gamma.values()[0] == 0.0:
self.setup = 'ch'
if self.scheme == 'upwind' and self.timestepping == 'ssprk3':
Vm = prognostic_variables.Vm
Vu = prognostic_variables.Vu
self.m = prognostic_variables.m
self.u = prognostic_variables.u
self.Xi = prognostic_variables.dXi
self.m0 = Function(Vm).assign(self.m)
# now make problem for the actual problem
psi = TestFunction(Vm)
self.m_trial = Function(Vm)
self.dm = Function(
Vm
) # introduce this as the advection operator for a single step
us = Dt * self.u + self.Xi
nhat = FacetNormal(mesh)
un = 0.5 * (dot(us, nhat) + abs(dot(us, nhat)))
ones = Function(Vu).project(as_vector([Constant(1.)]))
Lm = (psi * self.dm * dx -
psi.dx(0) * self.m_trial * dot(ones, us) * dx +
psi * self.m_trial * dot(ones, us.dx(0)) * dx +
jump(psi) * (un('+') * self.m_trial('+') -
un('-') * self.m_trial('-')) * dS)
mprob = NonlinearVariationalProblem(Lm, self.dm)
self.msolver = NonlinearVariationalSolver(mprob,
solver_parameters={
'ksp_type':
'preonly',
'pc_type':
'bjacobi',
'sub_pc_type':
'ilu'
})
phi = TestFunction(Vu)
Lu = (dot(phi, ones) * self.m * dx - dot(phi, self.u) * dx -
alphasq * dot(self.u.dx(0), phi.dx(0)) * dx)
uprob = NonlinearVariationalProblem(Lu, self.u)
self.usolver = NonlinearVariationalSolver(uprob,
solver_parameters={
'ksp_type':
'preonly',
'pc_type': 'lu'
})
elif self.scheme == 'hydrodynamic' and self.timestepping == 'midpoint':
Vu = prognostic_variables.Vu
self.u = prognostic_variables.u
W = MixedFunctionSpace((Vu, ) * 3)
psi, phi, zeta = TestFunctions(W)
w1 = Function(W)
self.u1, dFh, dGh = split(w1)
uh = (self.u1 + self.u) / 2
dXi = prognostic_variables.dXi
dXi_x = prognostic_variables.dXi_x
dXi_xx = prognostic_variables.dXi_xx
dvh = Dt * uh + dXi
Lu = (psi * (self.u1 - self.u) * dx +
psi * uh.dx(0) * dvh * dx - psi.dx(0) * dFh * dx +
psi * dGh * dx + phi * dFh * dx +
alphasq * phi.dx(0) * dFh.dx(0) * dx -
phi * uh * uh * Dt * dx -
0.5 * alphasq * phi * uh.dx(0) * uh.dx(0) * Dt * dx +
zeta * dGh * dx + alphasq * zeta.dx(0) * dGh.dx(0) * dx -
2 * zeta * uh * dXi_x * dx -
alphasq * zeta * uh.dx(0) * dXi_xx * dx)
self.u1, dFh, dGh = w1.split()
uprob = NonlinearVariationalProblem(Lu, w1)
self.usolver = NonlinearVariationalSolver(uprob,
solver_parameters={
'mat_type':
'aij',
'ksp_type':
'preonly',
'pc_type': 'lu'
})
elif self.scheme == 'no_gradient' and self.timestepping == 'midpoint':
# a version of the hydrodynamic form but without exploiting the gradient
Vu = prognostic_variables.Vu
self.u = prognostic_variables.u
W = MixedFunctionSpace((Vu, ) * 3)
psi, phi, zeta = TestFunctions(W)
w1 = Function(W)
self.u1, dFh, dGh = split(w1)
uh = (self.u1 + self.u) / 2
dXi = prognostic_variables.dXi
dXi_x = prognostic_variables.dXi_x
dXi_xx = prognostic_variables.dXi_xx
dvh = Dt * uh + dXi
Lu = (psi * (self.u1 - self.u) * dx +
psi * uh.dx(0) * dvh * dx + psi * dFh.dx(0) * dx +
psi * dGh * dx + phi * dFh * dx +
alphasq * phi.dx(0) * dFh.dx(0) * dx -
phi * uh * uh * Dt * dx -
0.5 * alphasq * phi * uh.dx(0) * uh.dx(0) * Dt * dx +
zeta * dGh * dx + alphasq * zeta.dx(0) * dGh.dx(0) * dx -
2 * zeta * uh * dXi_x * dx -
alphasq * zeta * uh.dx(0) * dXi_xx * dx)
self.u1, dFh, dGh = w1.split()
uprob = NonlinearVariationalProblem(Lu, w1)
self.usolver = NonlinearVariationalSolver(uprob,
solver_parameters={
'mat_type':
'aij',
'ksp_type':
'preonly',
'pc_type': 'lu'
})
elif self.scheme == 'test' and self.timestepping == 'midpoint':
self.u = prognostic_variables.u
Vu = prognostic_variables.Vu
psi = TestFunction(Vu)
self.u1 = Function(Vu)
uh = (self.u1 + self.u) / 2
dvh = Dt * uh + prognostic_variables.dXi
eqn = (psi * (self.u1 - self.u) * dx -
psi * uh * dvh.dx(0) * dx)
prob = NonlinearVariationalProblem(eqn, self.u1)
self.usolver = NonlinearVariationalSolver(prob,
solver_parameters={
'mat_type':
'aij',
'ksp_type':
'preonly',
'pc_type': 'lu'
})
else:
raise ValueError(
'Scheme %s and timestepping %s either not compatible or not recognised.'
% (self.scheme, self.timestepping))
elif alphasq.values()[0] == 0.0 and gamma.values()[0] > 0.0:
self.setup = 'kdv'
if self.scheme == 'upwind' and self.timestepping == 'ssprk3':
raise NotImplementedError(
'Scheme %s and timestepping %s not yet implemented.' %
(self.scheme, self.timestepping))
elif self.scheme == 'upwind' and self.timestepping == 'midpoint':
raise NotImplementedError(
'Scheme %s and timestepping %s not yet implemented.' %
(self.scheme, self.timestepping))
elif self.scheme == 'hydrodynamic' and self.timestepping == 'midpoint':
raise NotImplementedError(
'Scheme %s and timestepping %s not yet implemented.' %
(self.scheme, self.timestepping))
else:
raise ValueError(
'Scheme %s and timestepping %s either not compatible or not recognised.'
% (self.scheme, self.timestepping))
else:
raise NotImplementedError(
'Schemes for your values of alpha squared %.3f and gamma %.3f are not yet implemented.'
% (alphasq, gamma))
f.interpolate(2 * Omega * x[2] / R)
g, H = 9.810616, 5960.
u_0 = 2 * pi * R / (12 * 24 * 60 * 60.)
uexpr = u_0 * as_vector([-x[1], x[0], 0.0]) / R
Dexpr = H - (R * Omega * u_0 + u_0**2 / 2.) * x[2]**2 / (g * R**2)
bexpr = Constant(0.)
# Build function spaces
degree = 2
family = ("DG", "BDM", "CG")
W0 = FunctionSpace(mesh, family[0], degree - 1, family[0])
W1_elt = FiniteElement(family[1], triangle, degree)
W1 = FunctionSpace(mesh, W1_elt, name="HDiv")
W2 = FunctionSpace(mesh, family[2], degree + 1)
M = MixedFunctionSpace((W1, W0))
# Set up functions
xn = Function(M)
un, Dn = xn.split()
un.rename('u')
Dn.rename('D')
fields = {'u': un, 'D': Dn}
# Energy field for output
En = Function(W0, name='Energy')
e_form = lambda u, D: 0.5 * (D * inner(u, u) + g * (D + b)**2)
# Vorticity field
qn = Function(W2, name='potential vorticity')
fields['q'] = qn
def _setup_solver(self):
state = self.state # just cutting down line length a bit
Dt = state.timestepping.dt
beta_ = Dt*state.timestepping.alpha
mu = state.mu
Vu = state.spaces("HDiv")
Vb = state.spaces("HDiv_v")
Vp = state.spaces("DG")
# Store time-stepping coefficients as UFL Constants
dt = Constant(Dt)
beta = Constant(beta_)
# Split up the rhs vector (symbolically)
u_in, p_in, b_in = split(state.xrhs)
# Build the reduced function space for u,p
M = MixedFunctionSpace((Vu, Vp))
w, phi = TestFunctions(M)
u, p = TrialFunctions(M)
# Get background fields
bbar = state.fields("bbar")
# Analytical (approximate) elimination of theta
k = state.k # Upward pointing unit vector
b = -dot(k, u)*dot(k, grad(bbar))*beta + b_in
# vertical projection
def V(u):
return k*inner(u, k)
eqn = (
inner(w, (u - u_in))*dx
- beta*div(w)*p*dx
- beta*inner(w, k)*b*dx
+ phi*div(u)*dx
)
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 p solver
self.up = Function(M)
# Boundary conditions (assumes extruded mesh)
bcs = None if len(self.state.bcs) == 0 else self.state.bcs
# Solver for u, p
up_problem = LinearVariationalProblem(aeqn, Leqn, self.up, bcs=bcs)
# Provide callback for the nullspace of the trace system
def trace_nullsp(T):
return VectorSpaceBasis(constant=True)
appctx = {"trace_nullspace": trace_nullsp}
self.up_solver = LinearVariationalSolver(up_problem,
solver_parameters=self.solver_parameters,
appctx=appctx)
# Reconstruction of b
b = TrialFunction(Vb)
gamma = TestFunction(Vb)
u, p = self.up.split()
self.b = Function(Vb)
b_eqn = gamma*(b - b_in
+ dot(k, u)*dot(k, grad(bbar))*beta)*dx
b_problem = LinearVariationalProblem(lhs(b_eqn),
rhs(b_eqn),
self.b)
self.b_solver = LinearVariationalSolver(b_problem)
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")]
