class Gradient(DiagnosticField):
def __init__(self, name):
super().__init__()
self.fname = name
@property
def name(self):
return self.fname+"_gradient"
def setup(self, state):
if not self._initialised:
mesh_dim = state.mesh.geometric_dimension()
try:
field_dim = state.fields(self.fname).ufl_shape[0]
except IndexError:
field_dim = 1
shape = (mesh_dim, ) * field_dim
space = TensorFunctionSpace(state.mesh, "CG", 1, shape=shape)
super().setup(state, space=space)
f = state.fields(self.fname)
test = TestFunction(space)
trial = TrialFunction(space)
n = FacetNormal(state.mesh)
a = inner(test, trial)*dx
L = -inner(div(test), f)*dx
if space.extruded:
L += dot(dot(test, n), f)*(ds_t + ds_b)
prob = LinearVariationalProblem(a, L, self.field)
self.solver = LinearVariationalSolver(prob)
def compute(self, state):
self.solver.solve()
return self.field
class HydrostaticImbalance(DiagnosticField):
name = "HydrostaticImbalance"
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 compute(self, state):
self.imbalance_solver.solve()
return self.field[1]
class Precipitation(DiagnosticField):
name = "Precipitation"
def setup(self, state):
if not self._initialised:
space = state.spaces("DG0", state.mesh, "DG", 0)
super().setup(state, space=space)
rain = state.fields('rain')
rho = state.fields('rho')
v = state.fields('rainfall_velocity')
self.phi = TestFunction(space)
flux = TrialFunction(space)
n = FacetNormal(state.mesh)
un = 0.5 * (dot(v, n) + abs(dot(v, n)))
self.flux = Function(space)
a = self.phi * flux * dx
L = self.phi * rain * un * rho
if space.extruded:
L = L * (ds_b + ds_t + ds_v)
else:
L = L * ds
# setup solver
problem = LinearVariationalProblem(a, L, self.flux)
self.solver = LinearVariationalSolver(problem)
def compute(self, state):
self.solver.solve()
self.field.assign(self.field + assemble(self.flux * self.phi * dx))
return self.field
class TrueResidualV(DiagnosticField):
name = "TrueResidualV"
def setup(self, state):
super(TrueResidualV, self).setup(state)
unew, pnew, bnew = state.xn.split()
uold, pold, bold = state.xb.split()
ubar = 0.5 * (unew + uold)
H = state.parameters.H
f = state.parameters.f
dbdy = state.parameters.dbdy
dt = state.timestepping.dt
x, y, z = SpatialCoordinate(state.mesh)
V = FunctionSpace(state.mesh, "DG", 0)
wv = TestFunction(V)
v = TrialFunction(V)
vlhs = wv * v * dx
vrhs = wv * ((unew[1] - uold[1]) / dt + ubar[0] * ubar[1].dx(0) +
ubar[2] * ubar[1].dx(2) + f * ubar[0] + dbdy *
(z - H / 2)) * dx
self.vtres = Function(V)
vtresproblem = LinearVariationalProblem(vlhs, vrhs, self.vtres)
self.v_residual_solver = LinearVariationalSolver(
vtresproblem, solver_parameters={'ksp_type': 'cg'})
def compute(self, state):
self.v_residual_solver.solve()
v_residual = self.vtres
return self.field.interpolate(v_residual)
class GeostrophicImbalance(DiagnosticField):
name = "GeostrophicImbalance"
def setup(self, state):
super(GeostrophicImbalance, self).setup(state)
u = state.fields("u")
b = state.fields("b")
p = state.fields("p")
f = state.parameters.f
Vu = u.function_space()
v = TrialFunction(Vu)
w = TestFunction(Vu)
a = inner(w, v) * dx
L = (div(w) * p + inner(w, as_vector([f * u[1], 0.0, b]))) * dx
bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")]
self.imbalance = Function(Vu)
imbalanceproblem = LinearVariationalProblem(a,
L,
self.imbalance,
bcs=bcs)
self.imbalance_solver = LinearVariationalSolver(
imbalanceproblem, solver_parameters={'ksp_type': 'cg'})
def compute(self, state):
f = state.parameters.f
self.imbalance_solver.solve()
geostrophic_imbalance = self.imbalance[0] / f
return self.field.interpolate(geostrophic_imbalance)
class CompressibleEadyForcing(CompressibleForcing):
"""
Forcing class for compressible Eady equations.
"""
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 _build_forcing_solvers(self):
super(CompressibleEadyForcing, self)._build_forcing_solvers()
# theta_forcing
dthetady = self.state.parameters.dthetady
Vt = self.state.spaces("HDiv_v")
F = TrialFunction(Vt)
gamma = TestFunction(Vt)
self.thetaF = Function(Vt)
u0, _, _ = split(self.x0)
a = gamma * F * dx
L = -self.scaling * gamma * (dthetady *
inner(u0, as_vector([0., 1., 0.]))) * dx
theta_forcing_problem = LinearVariationalProblem(a, L, self.thetaF)
solver_parameters = {}
if logger.isEnabledFor(DEBUG):
solver_parameters["ksp_monitor_true_residual"] = None
self.theta_forcing_solver = LinearVariationalSolver(
theta_forcing_problem,
solver_parameters=solver_parameters,
options_prefix="ThetaForcingSolver")
def apply(self, scaling, x_in, x_nl, x_out, **kwargs):
Forcing.apply(self, scaling, x_in, x_nl, x_out, **kwargs)
self.theta_forcing_solver.solve() # places forcing in self.thetaF
_, _, theta_out = x_out.split()
theta_out += self.thetaF
class EadyForcing(IncompressibleForcing):
"""
Forcing class for Eady Boussinesq equations.
"""
def forcing_term(self):
L = Forcing.forcing_term(self)
dbdy = self.state.parameters.dbdy
H = self.state.parameters.H
Vp = self.state.spaces("DG")
_, _, z = SpatialCoordinate(self.state.mesh)
eady_exp = Function(Vp).interpolate(z - H / 2.)
L -= self.scaling * dbdy * eady_exp * inner(
self.test, as_vector([0., 1., 0.])) * dx
return L
def _build_forcing_solvers(self):
super(EadyForcing, self)._build_forcing_solvers()
# b_forcing
dbdy = self.state.parameters.dbdy
Vb = self.state.spaces("HDiv_v")
F = TrialFunction(Vb)
gamma = TestFunction(Vb)
self.bF = Function(Vb)
u0, _, b0 = split(self.x0)
a = gamma * F * dx
L = -self.scaling * gamma * (dbdy *
inner(u0, as_vector([0., 1., 0.]))) * dx
b_forcing_problem = LinearVariationalProblem(a, L, self.bF)
solver_parameters = {}
if logger.isEnabledFor(DEBUG):
solver_parameters["ksp_monitor_true_residual"] = None
self.b_forcing_solver = LinearVariationalSolver(
b_forcing_problem,
solver_parameters=solver_parameters,
options_prefix="BForcingSolver")
def apply(self, scaling, x_in, x_nl, x_out, **kwargs):
super(EadyForcing, self).apply(scaling, x_in, x_nl, x_out, **kwargs)
self.b_forcing_solver.solve() # places forcing in self.bF
_, _, b_out = x_out.split()
b_out += self.bF
class Vorticity(DiagnosticField):
"""Base diagnostic field class for vorticity."""
def setup(self, state, vorticity_type=None):
"""Solver for vorticity.
:arg state: The state containing model.
:arg vorticity_type: must be "relative", "absolute" or "potential"
"""
if not self._initialised:
vorticity_types = ["relative", "absolute", "potential"]
if vorticity_type not in vorticity_types:
raise ValueError("vorticity type must be one of %s, not %s" % (vorticity_types, vorticity_type))
try:
space = state.spaces("CG")
except AttributeError:
dgspace = state.spaces("DG")
cg_degree = dgspace.ufl_element().degree() + 2
space = FunctionSpace(state.mesh, "CG", cg_degree)
super().setup(state, space=space)
u = state.fields("u")
gamma = TestFunction(space)
q = TrialFunction(space)
if vorticity_type == "potential":
D = state.fields("D")
a = q*gamma*D*dx
else:
a = q*gamma*dx
if state.on_sphere:
cell_normals = CellNormal(state.mesh)
gradperp = lambda psi: cross(cell_normals, grad(psi))
L = (- inner(gradperp(gamma), u))*dx
else:
raise NotImplementedError("The vorticity diagnostics have only been implemented for 2D spherical geometries.")
if vorticity_type != "relative":
f = state.fields("coriolis")
L += gamma*f*dx
problem = LinearVariationalProblem(a, L, self.field)
self.solver = LinearVariationalSolver(problem, solver_parameters={"ksp_type": "cg"})
def compute(self, state):
"""Computes the vorticity.
"""
self.solver.solve()
return self.field
class ShallowWaterForcing(Forcing):
def __init__(self, state, linear=False):
self.state = state
g = state.parameters.g
f = state.f
Vu = state.V[0]
W = state.W
self.x0 = Function(W) # copy x to here
u0, D0 = split(self.x0)
n = FacetNormal(state.mesh)
un = 0.5 * (dot(u0, n) + abs(dot(u0, n)))
F = TrialFunction(Vu)
w = TestFunction(Vu)
self.uF = Function(Vu)
outward_normals = CellNormal(state.mesh)
perp = lambda u: cross(outward_normals, u)
a = inner(w, F) * dx
L = (-f * inner(w, perp(u0)) + g * div(w) * D0) * dx - g * inner(
jump(w, n), un("+") * D0("+") - un("-") * D0("-")
) * dS
if not linear:
L -= 0.5 * div(w) * inner(u0, u0) * dx
u_forcing_problem = LinearVariationalProblem(a, L, self.uF)
self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem)
def apply(self, scaling, x_in, x_nl, x_out, **kwargs):
self.x0.assign(x_nl)
self.u_forcing_solver.solve() # places forcing in self.uF
self.uF *= scaling
uF, _ = x_out.split()
x_out.assign(x_in)
uF += self.uF
class IncompressibleForcing(Forcing):
"""
Forcing class for incompressible Euler Boussinesq equations.
"""
def pressure_gradient_term(self):
_, p0, _ = split(self.x0)
L = div(self.test) * p0 * dx
return L
def gravity_term(self):
_, _, b0 = split(self.x0)
L = b0 * inner(self.test, self.state.k) * dx
return L
def _build_forcing_solvers(self):
super(IncompressibleForcing, self)._build_forcing_solvers()
Vp = self.state.spaces("DG")
p = TrialFunction(Vp)
q = TestFunction(Vp)
self.divu = Function(Vp)
u0, _, _ = split(self.x0)
a = p * q * dx
L = q * div(u0) * dx
divergence_problem = LinearVariationalProblem(a, L, self.divu)
solver_parameters = {}
if logger.isEnabledFor(DEBUG):
solver_parameters["ksp_monitor_true_residual"] = None
self.divergence_solver = LinearVariationalSolver(
divergence_problem,
solver_parameters=solver_parameters,
options_prefix="DivergenceSolver")
def apply(self, scaling, x_in, x_nl, x_out, **kwargs):
super(IncompressibleForcing, self).apply(scaling, x_in, x_nl, x_out,
**kwargs)
if 'incompressible' in kwargs and kwargs['incompressible']:
_, p_out, _ = x_out.split()
self.divergence_solver.solve()
p_out.assign(self.divu)
class InteriorPenalty(Diffusion):
"""
Interior penalty diffusion method
:arg state: :class:`.State` object.
:arg V: Function space of diffused field
:arg direction: list containing directions in which function space
:arg: mu: the penalty weighting function, which is
:recommended to be proportional to 1/dx
:arg: kappa: strength of diffusion
:arg: bcs: (optional) a list of boundary conditions to apply
"""
def __init__(self, state, V, kappa, mu, bcs=None):
super(InteriorPenalty, self).__init__(state)
dt = state.timestepping.dt
gamma = TestFunction(V)
phi = TrialFunction(V)
self.phi1 = Function(V)
n = FacetNormal(state.mesh)
a = inner(gamma, phi) * dx + dt * inner(grad(gamma),
grad(phi) * kappa) * dx
def get_flux_form(dS, M):
fluxes = (-inner(2 * avg(outer(phi, n)), avg(grad(gamma) * M)) -
inner(avg(grad(phi) * M), 2 * avg(outer(gamma, n))) +
mu * inner(2 * avg(outer(phi, n)),
2 * avg(outer(gamma, n) * kappa))) * dS
return fluxes
a += dt * get_flux_form(dS_v, kappa)
a += dt * get_flux_form(dS_h, kappa)
L = inner(gamma, phi) * dx
problem = LinearVariationalProblem(a,
action(L, self.phi1),
self.phi1,
bcs=bcs)
self.solver = LinearVariationalSolver(problem)
def apply(self, x_in, x_out):
self.phi1.assign(x_in)
self.solver.solve()
x_out.assign(self.phi1)
class InteriorPenalty(Diffusion):
"""
Interior penalty diffusion method
:arg state: :class:`.State` object.
:arg V: Function space of diffused field
:arg direction: list containing directions in which function space
is discontinuous: 1 corresponds to vertical, 2 to horizontal.
:arg params: dictionary containing the interior penalty parameters
:mu and kappa where mu is the penalty weighting function, which is
:recommended to be proportional to 1/dx
"""
def __init__(self, state, V, direction=[1,2], params=None):
super(InteriorPenalty, self).__init__(state)
dt = state.timestepping.dt
kappa = params['kappa']
mu = params['mu']
gamma = TestFunction(V)
phi = TrialFunction(V)
self.phi1 = Function(V)
n = FacetNormal(state.mesh)
a = inner(gamma,phi)*dx + dt*inner(grad(gamma), grad(phi)*kappa)*dx
def get_flux_form(dS, M):
fluxes = (-inner(2*avg(outer(phi, n)), avg(grad(gamma)*M))
- inner(avg(grad(phi)*M), 2*avg(outer(gamma, n)))
+ mu*inner(2*avg(outer(phi, n)), 2*avg(outer(gamma, n)*kappa)))*dS
return fluxes
if 1 in direction:
a += dt*get_flux_form(dS_v, kappa)
if 2 in direction:
a += dt*get_flux_form(dS_h, kappa)
L = inner(gamma,phi)*dx
problem = LinearVariationalProblem(a, action(L,self.phi1), self.phi1)
self.solver = LinearVariationalSolver(problem)
def apply(self, x_in, x_out):
self.phi1.assign(x_in)
self.solver.solve()
x_out.assign(self.phi1)
class ShallowWaterSolver(TimesteppingSolver):
def _setup_solver(self):
state = self.state
H = state.parameters.H
g = state.parameters.g
beta = state.timestepping.dt*state.timestepping.alpha
# Split up the rhs vector (symbolically)
u_in, D_in = split(state.xrhs)
W = state.W
w, phi = TestFunctions(W)
u, D = TrialFunctions(W)
eqn = (
inner(w, u) - beta*g*div(w)*D
- inner(w, u_in)
+ phi*D + beta*H*phi*div(u)
- phi*D_in
)*dx
aeqn = lhs(eqn)
Leqn = rhs(eqn)
# Place to put result of u rho solver
self.uD = Function(W)
# Solver for u, D
uD_problem = LinearVariationalProblem(
aeqn, Leqn, self.state.dy)
self.uD_solver = LinearVariationalSolver(uD_problem,
solver_parameters=self.params,
options_prefix='SWimplicit')
def solve(self):
"""
Apply the solver with rhs state.xrhs and result state.dy.
"""
self.uD_solver.solve()
class LinearAdvection_V3(Advection):
"""
An advection scheme that uses the linearised background state
in evaluation of the advection term for a DG space.
:arg state: :class:`.State` object.
:arg V:class:`.FunctionSpace` object. The DG Function space.
:arg qbar: :class:`.Function` object. The reference function that we
are linearising around.
:arg options: a PETSc options dictionary
"""
def __init__(self, state, V, qbar, options=None):
super(LinearAdvection_V3, self).__init__(state)
p = TestFunction(V)
q = TrialFunction(V)
self.dq = Function(V)
n = FacetNormal(state.mesh)
a = p*q*dx
L = (dot(grad(p), self.ubar)*qbar*dx -
jump(self.ubar*p, n)*avg(qbar)*(dS_v + dS_h))
aProblem = LinearVariationalProblem(a,L,self.dq)
if options is None:
options = {'ksp_type':'cg',
'pc_type':'bjacobi',
'sub_pc_type':'ilu'}
self.solver = LinearVariationalSolver(aProblem,
solver_parameters=options,
options_prefix='LinearAdvectionV3')
def apply(self, x_in, x_out):
dt = self.state.timestepping.dt
self.solver.solve()
x_out.assign(x_in + dt*self.dq)
class LinearAdvection_Vt(Advection):
"""
An advection scheme that uses the linearised background state
in evaluation of the advection term for the Vt temperature space.
:arg state: :class:`.State` object.
:arg V:class:`.FunctionSpace` object. The Function space for temperature.
:arg qbar: :class:`.Function` object. The reference function that we
are linearising around.
:arg options: a PETSc options dictionary
"""
def __init__(self, state, V, qbar, options=None):
super(LinearAdvection_Vt, self).__init__(state)
p = TestFunction(V)
q = TrialFunction(V)
self.dq = Function(V)
a = p*q*dx
k = state.k # Upward pointing unit vector
L = -p*dot(self.ubar,k)*dot(k,grad(qbar))*dx
aProblem = LinearVariationalProblem(a,L,self.dq)
if options is None:
options = {'ksp_type':'cg',
'pc_type':'bjacobi',
'sub_pc_type':'ilu'}
self.solver = LinearVariationalSolver(aProblem,
solver_parameters=options,
options_prefix='LinearAdvectionVt')
def apply(self, x_in, x_out):
dt = self.state.timestepping.dt
self.solver.solve()
x_out.assign(x_in + dt*self.dq)
r = sqrt((x - xc)**2 + (z - zc)**2)
theta_pert = Function(Vt).interpolate(
conditional(r > rc, 0.0,
Tdash * (cos(pi * r / (2.0 * rc)))**2))
# define initial theta
theta0.assign(theta_b * (theta_pert / 300.0 + 1.0))
# find perturbed rho
gamma = TestFunction(Vr)
rho_trial = TrialFunction(Vr)
lhs = gamma * rho_trial * dx
rhs = gamma * (rho_b * theta_b / theta0) * dx
rho_problem = LinearVariationalProblem(lhs, rhs, rho0)
rho_solver = LinearVariationalSolver(rho_problem)
rho_solver.solve()
state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])
# Set up transport schemes
VDG1 = state.spaces("DG1_equispaced")
VCG1 = FunctionSpace(mesh, "CG", 1)
Vt_brok = FunctionSpace(mesh, BrokenElement(Vt.ufl_element()))
Vu_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element())
Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1)
Vu_brok = FunctionSpace(mesh, BrokenElement(Vu.ufl_element()))
u_opts = RecoveredOptions(embedding_space=Vu_DG1,
recovered_space=Vu_CG1,
broken_space=Vu_brok,
boundary_method=Boundary_Method.dynamics)
class ShallowWaterSolver(TimesteppingSolver):
"""
Timestepping linear solver object for the nonlinear shallow water
equations with prognostic variables u and D. The linearized system
is solved using a hybridized-mixed method.
"""
solver_parameters = {
'ksp_type': 'preonly',
'mat_type': 'matfree',
'pc_type': 'python',
'pc_python_type': 'firedrake.HybridizationPC',
'hybridization': {'ksp_type': 'cg',
'pc_type': 'gamg',
'ksp_rtol': 1e-8,
'mg_levels': {'ksp_type': 'chebyshev',
'ksp_max_it': 2,
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}}
}
@timed_function("Gusto:SolverSetup")
def _setup_solver(self):
state = self.state
H_ = state.parameters.H
g_ = state.parameters.g
beta_ = state.timestepping.dt*state.timestepping.alpha
# Store time-stepping coefficients as UFL Constants
beta = Constant(beta_)
H = Constant(H_)
g = Constant(g_)
# Split up the rhs vector (symbolically)
u_in, D_in = split(state.xrhs)
W = state.W
w, phi = TestFunctions(W)
u, D = TrialFunctions(W)
eqn = (
inner(w, u) - beta*g*div(w)*D
- inner(w, u_in)
+ phi*D + beta*H*phi*div(u)
- phi*D_in
)*dx
aeqn = lhs(eqn)
Leqn = rhs(eqn)
# Place to put result of u rho solver
self.uD = Function(W)
# Solver for u, D
uD_problem = LinearVariationalProblem(
aeqn, Leqn, self.state.dy)
self.uD_solver = LinearVariationalSolver(uD_problem,
solver_parameters=self.solver_parameters,
options_prefix='SWimplicit')
@timed_function("Gusto:LinearSolve")
def solve(self):
"""
Apply the solver with rhs state.xrhs and result state.dy.
"""
self.uD_solver.solve()
class OldCompressibleSolver(TimesteppingSolver):
"""
Timestepping linear solver object for the compressible equations
in theta-pi formulation with prognostic variables u,rho,theta.
This solver follows the following strategy:
(1) Analytically eliminate theta (introduces error near topography)
(2) Solve resulting system for (u,rho) using a Schur preconditioner
(3) Reconstruct theta
:arg state: a :class:`.State` object containing everything else.
:arg quadrature degree: tuple (q_h, q_v) where q_h is the required
quadrature degree in the horizontal direction and q_v is that in
the vertical direction
:arg solver_parameters (optional): solver parameters
:arg overwrite_solver_parameters: boolean, if True use only the
solver_parameters that have been passed in, if False then update
the default solver parameters with the solver_parameters passed in.
:arg moisture (optional): list of names of moisture fields.
"""
solver_parameters = {
'pc_type': 'fieldsplit',
'pc_fieldsplit_type': 'schur',
'ksp_type': 'gcr',
'ksp_monitor_true_residual': None,
'ksp_max_it': 100,
'pc_fieldsplit_schur_fact_type': 'FULL',
'pc_fieldsplit_schur_precondition': 'selfp',
'fieldsplit_0': {'ksp_type': 'preonly',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'},
'fieldsplit_1': {'ksp_type': 'fgmres',
'ksp_monitor_true_residual': None,
'ksp_rtol': 1.0e-8,
'ksp_atol': 1.0e-8,
'ksp_max_it': 100,
'pc_type': 'gamg',
'pc_gamg_sym_graph': None,
'mg_levels': {'ksp_type': 'gmres',
'ksp_max_it': 5,
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}}
}
def __init__(self, state, quadrature_degree=None, solver_parameters=None,
overwrite_solver_parameters=False, moisture=None):
self.moisture = moisture
if quadrature_degree is not None:
self.quadrature_degree = quadrature_degree
else:
dgspace = state.spaces("DG")
if any(deg > 2 for deg in dgspace.ufl_element().degree()):
logger.warning("default quadrature degree most likely not sufficient for this degree element")
self.quadrature_degree = (5, 5)
super().__init__(state, solver_parameters, overwrite_solver_parameters)
@timed_function("Gusto:SolverSetup")
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')
@timed_function("Gusto:SchurCompLinearSolve")
def solve(self):
"""
Apply the solver with rhs state.xrhs and result state.dy.
"""
with timed_region("Gusto:VelocityDensitySolve"):
self.urho_solver.solve()
u1, rho1 = self.urho.split()
u, rho, theta = self.state.dy.split()
u.assign(u1)
rho.assign(rho1)
with timed_region("Gusto:ThetaRecon"):
self.theta_solver.solve()
theta.assign(self.theta)
perp = lambda u: cross(CellNormal(mesh), u)
out_n = CellNormal(mesh)
p_uw = lambda u, v: (s(u)('+') * cross(out_n('+'), v('+')) + s(u)
('-') * cross(out_n('-'), v('-')))
# Initial solve for vorticity for output
eta = TestFunction(W2)
q_ = TrialFunction(W2)
q_eqn = eta * q_ * Dn * dx + inner(perp(grad(eta)), un) * dx - eta * f * dx
q_p = LinearVariationalProblem(lhs(q_eqn), rhs(q_eqn), qn)
q_solver = LinearVariationalSolver(q_p,
solver_parameters={
"ksp_type": "preonly",
"pc_type": "lu"
})
q_solver.solve()
# Build advection, forcing forms
K = inner(un, un) / 3. + inner(un, unk) / 3. + inner(unk, unk) / 3.
Prhs = g * (0.5 * (Dn + Dp) + b) + 0.5 * K
# D advection solver
v, psi = TestFunctions(M)
F_, D_ = TrialFunctions(M)
D_bar = 0.5 * (Dn + D_)
Frhs = unk * D_ / 3. + un * D_ / 6. + unk * Dn / 6. + un * Dn / 3.
D_ad = -psi * div(F_) * dx
D_eqn = (D_ - Dn) * psi * dx - dt * D_ad + inner(v, F_ - Frhs) * dx
Dad_p = LinearVariationalProblem(lhs(D_eqn), rhs(D_eqn), xp)
D_ad_solver = LinearVariationalSolver(Dad_p,
perp = lambda u: cross(CellNormal(mesh), u)
out_n = CellNormal(mesh)
p_uw = lambda u, v: (s(u)('+') * cross(out_n('+'), v('+')) + s(u)
('-') * cross(out_n('-'), v('-')))
# Initial solve for vorticity for output
eta = TestFunction(W2)
q_ = TrialFunction(W2)
q_eqn = eta * q_ * Dn * dx + inner(perp(grad(eta)), un) * dx - eta * f * dx
q_p = LinearVariationalProblem(lhs(q_eqn), rhs(q_eqn), qn)
q_solver = LinearVariationalSolver(q_p,
solver_parameters={
"ksp_type": "preonly",
"pc_type": "lu"
})
q_solver.solve()
q0.assign(qn)
# Build advection, forcing forms
M_ext = MixedFunctionSpace((W1, W0, W0))
xp_ext = Function(M_ext)
u_, D_, P_ = TrialFunctions(M_ext)
v, psi, chi = TestFunctions(M_ext)
D_bar = 0.5 * (Dn + D_)
u_bar = 0.5 * (un + u_)
Frhs = unk * Dnk / 3. + un * Dnk / 6. + unk * Dn / 6. + un * Dn / 3.
K = inner(un, un) / 3. + inner(un, unk) / 3. + inner(unk, unk) / 3.
Prhs = g * (D_bar + b) + 0.5 * K
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)
class DGAdvection(Advection):
"""
DG 3 step SSPRK advection scheme that can be applied to a scalar
or vector field
:arg state: :class:`.State` object.
:arg V: function space of advected field - should be DG
:arg continuity: optional boolean.
If ``True``, the advection equation is of the form:
:math: `D_t +\nabla \cdot(uD) = 0`.
If ``False``, the advection equation is of the form:
:math: `D_t + (u \cdot \nabla)D = 0`.
"""
def __init__(self, state, V, continuity=False):
super(DGAdvection, self).__init__(state)
element = V.fiat_element
assert element.entity_dofs() == element.entity_closure_dofs(), "Provided space is not discontinuous"
dt = state.timestepping.dt
if V.extruded:
surface_measure = (dS_h + dS_v)
else:
surface_measure = dS
phi = TestFunction(V)
D = TrialFunction(V)
self.D1 = Function(V)
self.dD = Function(V)
n = FacetNormal(state.mesh)
# ( dot(v, n) + |dot(v, n)| )/2.0
un = 0.5*(dot(self.ubar, n) + abs(dot(self.ubar, n)))
a_mass = inner(phi,D)*dx
if continuity:
a_int = -inner(grad(phi), outer(D, self.ubar))*dx
else:
a_int = -inner(div(outer(phi,self.ubar)),D)*dx
a_flux = (dot(jump(phi), un('+')*D('+') - un('-')*D('-')))*surface_measure
arhs = a_mass - dt*(a_int + a_flux)
DGproblem = LinearVariationalProblem(a_mass, action(arhs,self.D1),
self.dD)
self.DGsolver = LinearVariationalSolver(DGproblem,
solver_parameters={
'ksp_type':'preonly',
'pc_type':'bjacobi',
'sub_pc_type': 'ilu'},
options_prefix='DGAdvection')
def apply(self, x_in, x_out):
# SSPRK Stage 1
self.D1.assign(x_in)
self.DGsolver.solve()
self.D1.assign(self.dD)
# SSPRK Stage 2
self.DGsolver.solve()
self.D1.assign(0.75*x_in + 0.25*self.dD)
# SSPRK Stage 3
self.DGsolver.solve()
x_out.assign((1.0/3.0)*x_in + (2.0/3.0)*self.dD)
class SawyerEliassenU(DiagnosticField):
name = "SawyerEliassenU"
def setup(self, state):
space = state.spaces("HDiv")
super(SawyerEliassenU, self).setup(state, space=space)
u = state.fields("u")
b = state.fields("b")
v = inner(u, as_vector([0., 1., 0.]))
# spaces
V0 = FunctionSpace(state.mesh, "CG", 2)
Vu = u.function_space()
# project b to V0
self.b_v0 = Function(V0)
btri = TrialFunction(V0)
btes = TestFunction(V0)
a = inner(btes, btri) * dx
L = inner(btes, b) * dx
projectbproblem = LinearVariationalProblem(a, L, self.b_v0)
self.project_b_solver = LinearVariationalSolver(
projectbproblem, solver_parameters={'ksp_type': 'cg'})
# project v to V0
self.v_v0 = Function(V0)
vtri = TrialFunction(V0)
vtes = TestFunction(V0)
a = inner(vtes, vtri) * dx
L = inner(vtes, v) * dx
projectvproblem = LinearVariationalProblem(a, L, self.v_v0)
self.project_v_solver = LinearVariationalSolver(
projectvproblem, solver_parameters={'ksp_type': 'cg'})
# stm/psi is a stream function
self.stm = Function(V0)
psi = TrialFunction(V0)
xsi = TestFunction(V0)
f = state.parameters.f
H = state.parameters.H
L = state.parameters.L
dbdy = state.parameters.dbdy
x, y, z = SpatialCoordinate(state.mesh)
bcs = [DirichletBC(V0, 0., "bottom"), DirichletBC(V0, 0., "top")]
Mat = as_matrix([[b.dx(2), 0., -f * self.v_v0.dx(2)], [0., 0., 0.],
[-self.b_v0.dx(0), 0., f**2 + f * self.v_v0.dx(0)]])
Equ = (inner(grad(xsi), dot(Mat, grad(psi))) -
dbdy * inner(grad(xsi), as_vector([-v, 0., f *
(z - H / 2)]))) * dx
# fourth-order terms
if state.parameters.fourthorder:
eps = Constant(0.0001)
brennersigma = Constant(10.0)
n = FacetNormal(state.mesh)
deltax = Constant(state.parameters.deltax)
deltaz = Constant(state.parameters.deltaz)
nn = as_matrix([[sqrt(brennersigma / Constant(deltax)), 0., 0.],
[0., 0., 0.],
[0., 0.,
sqrt(brennersigma / Constant(deltaz))]])
mu = as_matrix([[1., 0., 0.], [0., 0., 0.], [0., 0., H / L]])
# anisotropic form
Equ += eps * (
div(dot(mu, grad(psi))) * div(dot(mu, grad(xsi))) * dx -
(avg(dot(dot(grad(grad(psi)), n), n)) * jump(grad(xsi), n=n) +
avg(dot(dot(grad(grad(xsi)), n), n)) * jump(grad(psi), n=n) -
jump(nn * grad(psi), n=n) * jump(nn * grad(xsi), n=n)) *
(dS_h + dS_v))
Au = lhs(Equ)
Lu = rhs(Equ)
stmproblem = LinearVariationalProblem(Au, Lu, self.stm, bcs=bcs)
self.stream_function_solver = LinearVariationalSolver(
stmproblem, solver_parameters={'ksp_type': 'cg'})
# solve for sawyer_eliassen u
self.u = Function(Vu)
utrial = TrialFunction(Vu)
w = TestFunction(Vu)
a = inner(w, utrial) * dx
L = (w[0] * (-self.stm.dx(2)) + w[2] * (self.stm.dx(0))) * dx
ugproblem = LinearVariationalProblem(a, L, self.u)
self.sawyer_eliassen_u_solver = LinearVariationalSolver(
ugproblem, solver_parameters={'ksp_type': 'cg'})
def compute(self, state):
self.project_b_solver.solve()
self.project_v_solver.solve()
self.stream_function_solver.solve()
self.sawyer_eliassen_u_solver.solve()
sawyer_eliassen_u = self.u
return self.field.project(sawyer_eliassen_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))
class CompressibleSolver(TimesteppingSolver):
"""
Timestepping linear solver object for the compressible equations
in theta-exner formulation with prognostic variables u, rho, and theta.
This solver follows the following strategy:
(1) Analytically eliminate theta (introduces error near topography)
(2a) Formulate the resulting mixed system for u and rho using a
hybridized mixed method. This breaks continuity in the
linear perturbations of u, and introduces a new unknown on the
mesh interfaces approximating the average of the Exner pressure
perturbations. These trace unknowns also act as Lagrange
multipliers enforcing normal continuity of the "broken" u variable.
(2b) Statically condense the block-sparse system into a single system
for the Lagrange multipliers. This is the only globally coupled
system requiring a linear solver.
(2c) Using the computed trace variables, we locally recover the
broken velocity and density perturbations. This is accomplished
in two stages:
(i): Recover rho locally using the multipliers.
(ii): Recover "broken" u locally using rho and the multipliers.
(2d) Project the "broken" velocity field into the HDiv-conforming
space using local averaging.
(3) Reconstruct theta
:arg state: a :class:`.State` object containing everything else.
:arg quadrature degree: tuple (q_h, q_v) where q_h is the required
quadrature degree in the horizontal direction and q_v is that in
the vertical direction.
:arg solver_parameters (optional): solver parameters for the
trace system.
:arg overwrite_solver_parameters: boolean, if True use only the
solver_parameters that have been passed in, if False then update.
the default solver parameters with the solver_parameters passed in.
:arg moisture (optional): list of names of moisture fields.
"""
solver_parameters = {
'mat_type': 'matfree',
'ksp_type': 'preonly',
'pc_type': 'python',
'pc_python_type': 'firedrake.SCPC',
'pc_sc_eliminate_fields': '0, 1',
# The reduced operator is not symmetric
'condensed_field': {
'ksp_type': 'fgmres',
'ksp_rtol': 1.0e-8,
'ksp_atol': 1.0e-8,
'ksp_max_it': 100,
'pc_type': 'gamg',
'pc_gamg_sym_graph': None,
'mg_levels': {
'ksp_type': 'gmres',
'ksp_max_it': 5,
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
}
}
}
def __init__(self,
state,
equations,
alpha=0.5,
quadrature_degree=None,
solver_parameters=None,
overwrite_solver_parameters=False,
moisture=None):
self.moisture = moisture
if quadrature_degree is not None:
self.quadrature_degree = quadrature_degree
else:
dgspace = state.spaces("DG")
if any(deg > 2 for deg in dgspace.ufl_element().degree()):
logger.warning(
"default quadrature degree most likely not sufficient for this degree element"
)
self.quadrature_degree = (5, 5)
if logger.isEnabledFor(DEBUG):
# Set outer solver to FGMRES and turn on KSP monitor for the outer system
self.solver_parameters["ksp_type"] = "fgmres"
self.solver_parameters["mat_type"] = "aij"
self.solver_parameters["pmat_type"] = "matfree"
self.solver_parameters["ksp_monitor_true_residual"] = None
# Turn monitor on for the trace system
self.solver_parameters["condensed_field"][
"ksp_monitor_true_residual"] = None
super().__init__(state, equations, alpha, solver_parameters,
overwrite_solver_parameters)
@timed_function("Gusto:SolverSetup")
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']
@timed_function("Gusto:LinearSolve")
def solve(self, xrhs, dy):
"""
Apply the solver with rhs xrhs and result dy.
"""
self.xrhs.assign(xrhs)
# Solve the hybridized system
self.hybridized_solver.solve()
broken_u, rho1, _ = self.urhol0.split()
u1 = self.u_hdiv
# Project broken_u into the HDiv space
u1.assign(0.0)
with timed_region("Gusto:HybridProjectHDiv"):
par_loop(
self._average_kernel, dx, {
"w": (self._weight, READ),
"vec_in": (broken_u, READ),
"vec_out": (u1, INC)
})
# Reapply bcs to ensure they are satisfied
for bc in self.bcs:
bc.apply(u1)
# Copy back into u and rho cpts of dy
u, rho, theta = dy.split()[0:3]
u.assign(u1)
rho.assign(rho1)
# Reconstruct theta
with timed_region("Gusto:ThetaRecon"):
self.theta_solver.solve()
# Copy into theta cpt of dy
theta.assign(self.theta)
class IncompressibleForcing(Forcing):
"""
Forcing class for incompressible Euler Boussinesq equations.
"""
def __init__(self, state, linear=False):
self.state = state
self._build_forcing_solver(linear)
def _build_forcing_solver(self, linear):
"""
Only put forcing terms into the u equation.
"""
state = self.state
self.scaling = Constant(1.0)
Vu = state.V[0]
W = state.W
self.x0 = Function(W) # copy x to here
u0, p0, b0 = split(self.x0)
F = TrialFunction(Vu)
w = TestFunction(Vu)
self.uF = Function(Vu)
Omega = state.Omega
mu = state.mu
a = inner(w, F) * dx
L = (
self.scaling * div(w) * p0 * dx # pressure gradient
+ self.scaling * b0 * inner(w, state.k) * dx # gravity term
)
if not linear:
L -= self.scaling * 0.5 * div(w) * inner(u0, u0) * dx
if Omega is not None:
L -= self.scaling * inner(w, cross(2 * Omega, u0)) * dx # Coriolis term
if mu is not None:
self.mu_scaling = Constant(1.0)
L -= self.mu_scaling * mu * inner(w, state.k) * inner(u0, state.k) * dx
bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")]
u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs)
self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem)
Vp = state.V[1]
p = TrialFunction(Vp)
q = TestFunction(Vp)
self.divu = Function(Vp)
a = p * q * dx
L = q * div(u0) * dx
divergence_problem = LinearVariationalProblem(a, L, self.divu)
self.divergence_solver = LinearVariationalSolver(divergence_problem)
def apply(self, scaling, x_in, x_nl, x_out, **kwargs):
self.x0.assign(x_nl)
self.scaling.assign(scaling)
if "mu_alpha" in kwargs and kwargs["mu_alpha"] is not None:
self.mu_scaling.assign(kwargs["mu_alpha"])
self.u_forcing_solver.solve() # places forcing in self.uF
u_out, p_out, _ = x_out.split()
x_out.assign(x_in)
u_out += self.uF
if "incompressible" in kwargs and kwargs["incompressible"]:
self.divergence_solver.solve()
p_out.assign(self.divu)
class CompressibleForcing(Forcing):
"""
Forcing class for compressible Euler equations.
"""
def __init__(self, state, linear=False):
self.state = state
self._build_forcing_solver(linear)
def _build_forcing_solver(self, linear):
"""
Only put forcing terms into the u equation.
"""
state = self.state
self.scaling = Constant(1.0)
Vu = state.V[0]
W = state.W
self.x0 = Function(W) # copy x to here
u0, rho0, theta0 = split(self.x0)
F = TrialFunction(Vu)
w = TestFunction(Vu)
self.uF = Function(Vu)
Omega = state.Omega
cp = state.parameters.cp
mu = state.mu
n = FacetNormal(state.mesh)
pi = exner(theta0, rho0, state)
a = inner(w, F) * dx
L = self.scaling * (
+cp * div(theta0 * w) * pi * dx # pressure gradient [volume]
- cp * jump(w * theta0, n) * avg(pi) * dS_v # pressure gradient [surface]
)
if state.parameters.geopotential:
Phi = state.Phi
L += self.scaling * div(w) * Phi * dx # gravity term
else:
g = state.parameters.g
L -= self.scaling * g * inner(w, state.k) * dx # gravity term
if not linear:
L -= self.scaling * 0.5 * div(w) * inner(u0, u0) * dx
if Omega is not None:
L -= self.scaling * inner(w, cross(2 * Omega, u0)) * dx # Coriolis term
if mu is not None:
self.mu_scaling = Constant(1.0)
L -= self.mu_scaling * mu * inner(w, state.k) * inner(u0, state.k) * dx
bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")]
u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs)
self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem)
def apply(self, scaling, x_in, x_nl, x_out, **kwargs):
self.x0.assign(x_nl)
self.scaling.assign(scaling)
if "mu_alpha" in kwargs and kwargs["mu_alpha"] is not None:
self.mu_scaling.assign(kwargs["mu_alpha"])
self.u_forcing_solver.solve() # places forcing in self.uF
u_out, _, _ = x_out.split()
x_out.assign(x_in)
u_out += self.uF
class EulerPoincareForm(Advection):
def __init__(self, state, V):
super(EulerPoincareForm, self).__init__(state)
dt = state.timestepping.dt
w = TestFunction(V)
u = TrialFunction(V)
self.u0 = Function(V)
ustar = 0.5*(self.u0 + u)
n = FacetNormal(state.mesh)
Upwind = 0.5*(sign(dot(self.ubar, n))+1)
if state.mesh.geometric_dimension() == 3:
if V.extruded:
surface_measure = (dS_h + dS_v)
else:
surface_measure = dS
# <w,curl(u) cross ubar + grad( u.ubar)>
# =<curl(u),ubar cross w> - <div(w), u.ubar>
# =<u,curl(ubar cross w)> -
# <<u_upwind, [[n cross(ubar cross w)cross]]>>
both = lambda u: 2*avg(u)
Eqn = (
inner(w, u-self.u0)*dx
+ dt*inner(ustar, curl(cross(self.ubar, w)))*dx
- dt*inner(both(Upwind*ustar),
both(cross(n, cross(self.ubar, w))))*surface_measure
- dt*div(w)*inner(ustar, self.ubar)*dx
)
# define surface measure and terms involving perp differently
# for slice (i.e. if V.extruded is True) and shallow water
# (V.extruded is False)
else:
if V.extruded:
surface_measure = (dS_h + dS_v)
perp = lambda u: as_vector([-u[1], u[0]])
perp_u_upwind = Upwind('+')*perp(ustar('+')) + Upwind('-')*perp(ustar('-'))
else:
surface_measure = dS
outward_normals = CellNormal(state.mesh)
perp = lambda u: cross(outward_normals, u)
perp_u_upwind = Upwind('+')*cross(outward_normals('+'),ustar('+')) + Upwind('-')*cross(outward_normals('-'),ustar('-'))
Eqn = (
(inner(w, u-self.u0)
- dt*inner(w, div(perp(ustar))*perp(self.ubar))
- dt*div(w)*inner(ustar, self.ubar))*dx
- dt*inner(jump(inner(w, perp(self.ubar)), n),
perp_u_upwind)*surface_measure
+ dt*jump(inner(w,
perp(self.ubar))*perp(ustar), n)*surface_measure
)
a = lhs(Eqn)
L = rhs(Eqn)
self.u1 = Function(V)
uproblem = LinearVariationalProblem(a, L, self.u1)
self.usolver = LinearVariationalSolver(uproblem,
options_prefix='EPAdvection')
def apply(self, x_in, x_out):
self.u0.assign(x_in)
self.usolver.solve()
x_out.assign(self.u1)
class Forcing(object, metaclass=ABCMeta):
"""
Base class for forcing terms for Gusto.
:arg state: x :class:`.State` object.
:arg euler_poincare: if True then the momentum equation is in Euler
Poincare form and we need to add 0.5*grad(u^2) to the forcing term.
If False then this term is not added.
:arg linear: if True then we are solving a linear equation so nonlinear
terms (namely the Euler Poincare term) should not be added.
:arg extra_terms: extra terms to add to the u component of the forcing
term - these will be multiplied by the appropriate test function.
"""
def __init__(self,
state,
euler_poincare=True,
linear=False,
extra_terms=None,
moisture=None):
self.state = state
if linear:
self.euler_poincare = False
logger.warning(
'Setting euler_poincare to False because you have set linear=True'
)
else:
self.euler_poincare = euler_poincare
# set up functions
self.Vu = state.spaces("HDiv")
# this is the function that the forcing term is applied to
self.x0 = Function(state.W)
self.test = TestFunction(self.Vu)
self.trial = TrialFunction(self.Vu)
# this is the function that contains the result of solving
# <test, trial> = <test, F(x0)>, where F is the forcing term
self.uF = Function(self.Vu)
# find out which terms we need
self.extruded = self.Vu.extruded
self.coriolis = state.Omega is not None or hasattr(
state.fields, "coriolis")
self.sponge = state.mu is not None
self.hydrostatic = state.hydrostatic
self.topography = hasattr(state.fields, "topography")
self.extra_terms = extra_terms
self.moisture = moisture
# some constants to use for scaling terms
self.scaling = Constant(1.)
self.impl = Constant(1.)
self._build_forcing_solvers()
def mass_term(self):
return inner(self.test, self.trial) * dx
def coriolis_term(self):
u0 = split(self.x0)[0]
return -inner(self.test, cross(2 * self.state.Omega, u0)) * dx
def sponge_term(self):
u0 = split(self.x0)[0]
return self.state.mu * inner(self.test, self.state.k) * inner(
u0, self.state.k) * dx
def euler_poincare_term(self):
u0 = split(self.x0)[0]
return -0.5 * div(self.test) * inner(self.state.h_project(u0), u0) * dx
def hydrostatic_term(self):
u0 = split(self.x0)[0]
return inner(u0, self.state.k) * inner(self.test, self.state.k) * dx
@abstractmethod
def pressure_gradient_term(self):
pass
def forcing_term(self):
L = self.pressure_gradient_term()
if self.extruded:
L += self.gravity_term()
if self.coriolis:
L += self.coriolis_term()
if self.euler_poincare:
L += self.euler_poincare_term()
if self.topography:
L += self.topography_term()
if self.extra_terms is not None:
L += inner(self.test, self.extra_terms) * dx
# scale L
L = self.scaling * L
# sponge term has a separate scaling factor as it is always implicit
if self.sponge:
L -= self.impl * self.state.timestepping.dt * self.sponge_term()
# hydrostatic term has no scaling factor
if self.hydrostatic:
L += (2 * self.impl - 1) * self.hydrostatic_term()
return L
def _build_forcing_solvers(self):
a = self.mass_term()
L = self.forcing_term()
bcs = None if len(self.state.bcs) == 0 else self.state.bcs
u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs)
solver_parameters = {}
if logger.isEnabledFor(DEBUG):
solver_parameters["ksp_monitor_true_residual"] = None
self.u_forcing_solver = LinearVariationalSolver(
u_forcing_problem,
solver_parameters=solver_parameters,
options_prefix="UForcingSolver")
def apply(self, scaling, x_in, x_nl, x_out, **kwargs):
"""
Function takes x as input, computes F(x_nl) and returns
x_out = x + scale*F(x_nl)
as output.
:arg x_in: :class:`.Function` object
:arg x_nl: :class:`.Function` object
:arg x_out: :class:`.Function` object
:arg implicit: forcing stage for sponge and hydrostatic terms, if present
"""
self.scaling.assign(scaling)
self.x0.assign(x_nl)
implicit = kwargs.get("implicit")
if implicit is not None:
self.impl.assign(int(implicit))
self.u_forcing_solver.solve() # places forcing in self.uF
uF = x_out.split()[0]
x_out.assign(x_in)
uF += self.uF
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))
class CompressibleForcing(Forcing):
"""
Forcing class for compressible Euler equations.
"""
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 gravity_term(self):
g = self.state.parameters.g
L = -g * inner(self.test, self.state.k) * dx
return L
def theta_forcing(self):
cv = self.state.parameters.cv
cp = self.state.parameters.cp
c_vv = self.state.parameters.c_vv
c_pv = self.state.parameters.c_pv
c_pl = self.state.parameters.c_pl
R_d = self.state.parameters.R_d
R_v = self.state.parameters.R_v
u0, _, theta0 = split(self.x0)
water_v = self.state.fields('water_v')
water_c = self.state.fields('water_c')
c_vml = cv + water_v * c_vv + water_c * c_pl
c_pml = cp + water_v * c_pv + water_c * c_pl
R_m = R_d + water_v * R_v
L = -theta0 * (R_m / c_vml - (R_d * c_pml) / (cp * c_vml)) * div(u0)
return self.scaling * L
def _build_forcing_solvers(self):
super(CompressibleForcing, self)._build_forcing_solvers()
# build forcing for theta equation
if self.moisture is not None:
_, _, theta0 = split(self.x0)
Vt = self.state.spaces("HDiv_v")
p = TrialFunction(Vt)
q = TestFunction(Vt)
self.thetaF = Function(Vt)
a = p * q * dx
L = self.theta_forcing()
L = q * L * dx
theta_problem = LinearVariationalProblem(a, L, self.thetaF)
solver_parameters = {}
if logger.isEnabledFor(DEBUG):
solver_parameters["ksp_monitor_true_residual"] = None
self.theta_solver = LinearVariationalSolver(
theta_problem,
solver_parameters=solver_parameters,
option_prefix="ThetaForcingSolver")
def apply(self, scaling, x_in, x_nl, x_out, **kwargs):
super(CompressibleForcing, self).apply(scaling, x_in, x_nl, x_out,
**kwargs)
if self.moisture is not None:
self.theta_solver.solve()
_, _, theta_out = x_out.split()
theta_out += self.thetaF
class SUPGAdvection(Advection):
"""
An SUPG advection scheme that can apply DG upwinding (in the direction
specified by the direction arg) if the function space is only
partially continuous.
:arg state: :class:`.State` object.
:arg V:class:`.FunctionSpace` object. The advected field function space.
:arg direction: list containing the directions in which the function
space is discontinuous. 1 corresponds to the vertical direction, 2 to
the horizontal direction
:arg supg_params: dictionary containing SUPG parameters tau for each
direction. If not supplied tau is set to dt/sqrt(15.)
"""
def __init__(self, state, V, direction=[], supg_params=None):
super(SUPGAdvection, self).__init__(state)
dt = state.timestepping.dt
params = supg_params.copy() if supg_params else {}
params.setdefault('a0', dt/sqrt(15.))
params.setdefault('a1', dt/sqrt(15.))
gamma = TestFunction(V)
theta = TrialFunction(V)
self.theta0 = Function(V)
# make SUPG test function
taus = [params["a0"], params["a1"]]
for i in direction:
taus[i] = 0.0
tau = Constant(((taus[0], 0.), (0., taus[1])))
dgamma = dot(dot(self.ubar, tau), grad(gamma))
gammaSU = gamma + dgamma
n = FacetNormal(state.mesh)
un = 0.5*(dot(self.ubar, n) + abs(dot(self.ubar, n)))
a_mass = gammaSU*theta*dx
arhs = a_mass - dt*gammaSU*dot(self.ubar, grad(theta))*dx
if 1 in direction:
arhs -= (
dt*dot(jump(gammaSU), (un('+')*theta('+')
- un('-')*theta('-')))*dS_v
- dt*(gammaSU('+')*dot(self.ubar('+'), n('+'))*theta('+')
+ gammaSU('-')*dot(self.ubar('-'), n('-'))*theta('-'))*dS_v
)
if 2 in direction:
arhs -= (
dt*dot(jump(gammaSU), (un('+')*theta('+')
- un('-')*theta('-')))*dS_h
- dt*(gammaSU('+')*dot(self.ubar('+'), n('+'))*theta('+')
+ gammaSU('-')*dot(self.ubar('-'), n('-'))*theta('-'))*dS_h
)
self.theta1 = Function(V)
self.dtheta = Function(V)
problem = LinearVariationalProblem(a_mass, action(arhs,self.theta1), self.dtheta)
self.solver = LinearVariationalSolver(problem,
options_prefix='SUPGAdvection')
def apply(self, x_in, x_out):
# SSPRK Stage 1
self.theta1.assign(x_in)
self.solver.solve()
self.theta1.assign(self.dtheta)
# SSPRK Stage 2
self.solver.solve()
self.theta1.assign(0.75*x_in + 0.25*self.dtheta)
# SSPRK Stage 3
self.solver.solve()
x_out.assign((1.0/3.0)*x_in + (2.0/3.0)*self.dtheta)
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']
class LinearTimesteppingSolver(object):
"""
Timestepping linear solver object for the nonlinear shallow water
equations with prognostic variables u and D. The linearized system
is solved using a hybridized-mixed method.
"""
solver_parameters = {
'ksp_type': 'preonly',
'mat_type': 'matfree',
'pc_type': 'python',
'pc_python_type': 'firedrake.HybridizationPC',
'hybridization': {
'ksp_type': 'cg',
'pc_type': 'gamg',
'ksp_rtol': 1e-8,
'mg_levels': {
'ksp_type': 'chebyshev',
'ksp_max_it': 2,
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
}
}
}
def __init__(self, equation, alpha):
residual = equation.residual.label_map(
lambda t: t.has_label(linearisation),
lambda t: Term(t.get(linearisation).form, t.labels), drop)
dt = equation.state.dt
W = equation.function_space
beta = dt * alpha
# Split up the rhs vector (symbolically)
self.xrhs = Function(W)
aeqn = residual.label_map(
lambda t:
(t.has_label(time_derivative) and t.has_label(linearisation)),
map_if_false=lambda t: beta * t)
Leqn = residual.label_map(
lambda t:
(t.has_label(time_derivative) and t.has_label(linearisation)),
map_if_false=drop)
# Place to put result of solver
self.dy = Function(W)
# Solver
bcs = equation.bcs['u']
problem = LinearVariationalProblem(aeqn.form,
action(Leqn.form, self.xrhs),
self.dy,
bcs=bcs)
self.solver = LinearVariationalSolver(
problem,
solver_parameters=self.solver_parameters,
options_prefix='linear_solver')
@timed_function("Gusto:LinearSolve")
def solve(self, xrhs, dy):
"""
Apply the solver with rhs xrhs and result dy.
"""
self.xrhs.assign(xrhs)
self.solver.solve()
dy.assign(self.dy)
class IncompressibleSolver(TimesteppingSolver):
"""Timestepping linear solver object for the incompressible
Boussinesq equations with prognostic variables u, p, b.
This solver follows the following strategy:
(1) Analytically eliminate b (introduces error near topography)
(2) Solve resulting system for (u,p) using a block Hdiv preconditioner
(3) Reconstruct b
This currently requires a (parallel) direct solver so is probably
a bit memory-hungry, we'll improve this with a hybridised solver
soon.
:arg state: a :class:`.State` object containing everything else.
:arg L: the width of the domain, used in the preconditioner.
:arg params: Solver parameters.
"""
def __init__(self, state, L, params=None):
self.state = state
if params is None:
self.params = {'ksp_type':'gmres',
'pc_type':'fieldsplit',
'pc_fieldsplit_type':'additive',
'fieldsplit_0_pc_type':'lu',
'fieldsplit_1_pc_type':'lu',
'fieldsplit_0_pc_factor_mat_solver_package': 'mumps',
'fieldsplit_0_pc_factor_mat_solver_package': 'mumps',
'fieldsplit_0_ksp_type':'preonly',
'fieldsplit_1_ksp_type':'preonly'}
else:
self.params = params
self.L = L
# setup the solver
self._setup_solver()
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)
def solve(self):
"""
Apply the solver with rhs state.xrhs and result state.dy.
"""
self.up_solver.solve()
u1, p1 = self.up.split()
u, p, b = self.state.dy.split()
u.assign(u1)
p.assign(p1)
self.b_solver.solve()
b.assign(self.b)
class IncompressibleSolver(TimesteppingSolver):
"""Timestepping linear solver object for the incompressible
Boussinesq equations with prognostic variables u, p, b.
This solver follows the following strategy:
(1) Analytically eliminate b (introduces error near topography)
(2) Solve resulting system for (u,p) using a hybrid-mixed method
(3) Reconstruct b
:arg state: a :class:`.State` object containing everything else.
:arg solver_parameters: (optional) Solver parameters.
:arg overwrite_solver_parameters: boolean, if True use only the
solver_parameters that have been passed in, if False then
update the default solver parameters with the solver_parameters
passed in.
"""
solver_parameters = {
'ksp_type': 'preonly',
'mat_type': 'matfree',
'pc_type': 'python',
'pc_python_type': 'firedrake.HybridizationPC',
'hybridization': {
'ksp_type': 'cg',
'pc_type': 'gamg',
'ksp_rtol': 1e-8,
'mg_levels': {
'ksp_type': 'chebyshev',
'ksp_max_it': 2,
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
}
}
}
@timed_function("Gusto:SolverSetup")
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)
@timed_function("Gusto:LinearSolve")
def solve(self, xrhs, dy):
"""
Apply the solver with rhs xrhs and result dy.
"""
self.xrhs.assign(xrhs)
with timed_region("Gusto:VelocityPressureSolve"):
self.up_solver.solve()
u1, p1 = self.up.split()
u, p, b = dy.split()
u.assign(u1)
p.assign(p1)
with timed_region("Gusto:BuoyancyRecon"):
self.b_solver.solve()
b.assign(self.b)
class CompressibleSolver(TimesteppingSolver):
"""
Timestepping linear solver object for the compressible equations
in theta-pi formulation with prognostic variables u,rho,theta.
This solver follows the following strategy:
(1) Analytically eliminate theta (introduces error near topography)
(2) Solve resulting system for (u,rho) using a Schur preconditioner
(3) Reconstruct theta
:arg state: a :class:`.State` object containing everything else.
"""
def __init__(self, state, params=None):
self.state = state
if params is None:
self.params = {'pc_type': 'fieldsplit',
'pc_fieldsplit_type': 'schur',
'ksp_type': 'gmres',
'ksp_max_it': 100,
'ksp_gmres_restart': 50,
'pc_fieldsplit_schur_fact_type': 'FULL',
'pc_fieldsplit_schur_precondition': 'selfp',
'fieldsplit_0_ksp_type': 'preonly',
'fieldsplit_0_pc_type': 'bjacobi',
'fieldsplit_0_sub_pc_type': 'ilu',
'fieldsplit_1_ksp_type': 'preonly',
'fieldsplit_1_pc_type': 'gamg',
'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': 1,
'fieldsplit_1_mg_levels_pc_type': 'bjacobi',
'fieldsplit_1_mg_levels_sub_pc_type': 'ilu'}
else:
self.params = params
# setup the solver
self._setup_solver()
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 solve(self):
"""
Apply the solver with rhs state.xrhs and result state.dy.
"""
self.urho_solver.solve()
u1, rho1 = self.urho.split()
u, rho, theta = self.state.dy.split()
u.assign(u1)
rho.assign(rho1)
self.theta_solver.solve()
theta.assign(self.theta)
class HybridizedCompressibleSolver(TimesteppingSolver):
"""
Timestepping linear solver object for the compressible equations
in theta-pi formulation with prognostic variables u, rho, and theta.
This solver follows the following strategy:
(1) Analytically eliminate theta (introduces error near topography)
(2a) Formulate the resulting mixed system for u and rho using a
hybridized mixed method. This breaks continuity in the
linear perturbations of u, and introduces a new unknown on the
mesh interfaces approximating the average of the Exner pressure
perturbations. These trace unknowns also act as Lagrange
multipliers enforcing normal continuity of the "broken" u variable.
(2b) Statically condense the block-sparse system into a single system
for the Lagrange multipliers. This is the only globally coupled
system requiring a linear solver.
(2c) Using the computed trace variables, we locally recover the
broken velocity and density perturbations. This is accomplished
in two stages:
(i): Recover rho locally using the multipliers.
(ii): Recover "broken" u locally using rho and the multipliers.
(2d) Project the "broken" velocity field into the HDiv-conforming
space using local averaging.
(3) Reconstruct theta
:arg state: a :class:`.State` object containing everything else.
:arg quadrature degree: tuple (q_h, q_v) where q_h is the required
quadrature degree in the horizontal direction and q_v is that in
the vertical direction.
:arg solver_parameters (optional): solver parameters for the
trace system.
:arg overwrite_solver_parameters: boolean, if True use only the
solver_parameters that have been passed in, if False then update.
the default solver parameters with the solver_parameters passed in.
:arg moisture (optional): list of names of moisture fields.
"""
# Solver parameters for the Lagrange multiplier system
# NOTE: The reduced operator is not symmetric
solver_parameters = {'ksp_type': 'gmres',
'pc_type': 'gamg',
'ksp_rtol': 1.0e-8,
'mg_levels': {'ksp_type': 'richardson',
'ksp_max_it': 2,
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'}}
def __init__(self, state, quadrature_degree=None, solver_parameters=None,
overwrite_solver_parameters=False, moisture=None):
self.moisture = moisture
self.state = state
if quadrature_degree is not None:
self.quadrature_degree = quadrature_degree
else:
dgspace = state.spaces("DG")
if any(deg > 2 for deg in dgspace.ufl_element().degree()):
state.logger.warning("default quadrature degree most likely not sufficient for this degree element")
self.quadrature_degree = (5, 5)
super().__init__(state, solver_parameters, overwrite_solver_parameters)
@timed_function("Gusto:SolverSetup")
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")]
@timed_function("Gusto:LinearSolve")
def solve(self):
"""
Apply the solver with rhs state.xrhs and result state.dy.
"""
# Solve the velocity-density system
with timed_region("Gusto:VelocityDensitySolve"):
# Assemble the RHS for lambda into self.R
with timed_region("Gusto:HybridRHS"):
self._assemble_Rexp()
# Solve for lambda
with timed_region("Gusto:HybridTraceSolve"):
self.lSolver.solve(self.lambdar, self.R)
# Reconstruct broken u and rho
with timed_region("Gusto:HybridRecon"):
self._assemble_rho()
self._assemble_u()
broken_u, rho1 = self.urho.split()
u1 = self.u_hdiv
# Project broken_u into the HDiv space
u1.assign(0.0)
with timed_region("Gusto:HybridProjectHDiv"):
par_loop(self._average_kernel, dx,
{"w": (self._weight, READ),
"vec_in": (broken_u, READ),
"vec_out": (u1, INC)})
# Reapply bcs to ensure they are satisfied
for bc in self.bcs:
bc.apply(u1)
# Copy back into u and rho cpts of dy
u, rho, theta = self.state.dy.split()
u.assign(u1)
rho.assign(rho1)
# Reconstruct theta
with timed_region("Gusto:ThetaRecon"):
self.theta_solver.solve()
# Copy into theta cpt of dy
theta.assign(self.theta)
