def advection_term(self, q):
if self.state.mesh.topological_dimension() == 3:
# <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)
L = (inner(q, curl(cross(self.ubar, self.test))) * dx -
inner(both(self.Upwind * q),
both(cross(self.n, cross(self.ubar, self.test)))) *
self.dS)
else:
if self.ibp == "once":
L = (-inner(
self.gradperp(inner(self.test, self.perp(self.ubar))), q) *
dx - inner(
jump(inner(self.test, self.perp(self.ubar)), self.n),
self.perp_u_upwind(q)) * self.dS)
else:
L = (
(-inner(self.test,
div(self.perp(q)) * self.perp(self.ubar))) * dx -
inner(jump(inner(self.test, self.perp(self.ubar)), self.n),
self.perp_u_upwind(q)) * self.dS + jump(
inner(self.test, self.perp(self.ubar)) *
self.perp(q), self.n) * self.dS)
L -= 0.5 * div(self.test) * inner(q, self.ubar) * dx
return L
def __init__(self, state, V, *, ibp="once", solver_params=None):
super().__init__(state=state,
V=V,
ibp=ibp,
solver_params=solver_params)
self.Upwind = 0.5 * (sign(dot(self.ubar, self.n)) + 1)
if self.state.mesh.topological_dimension() == 2:
self.perp = state.perp
if state.on_sphere:
outward_normals = CellNormal(state.mesh)
self.perp_u_upwind = lambda q: self.Upwind('+') * cross(
outward_normals('+'), q('+')) + self.Upwind('-') * cross(
outward_normals('-'), q('-'))
else:
self.perp_u_upwind = lambda q: self.Upwind('+') * state.perp(
q('+')) + self.Upwind('-') * state.perp(q('-'))
self.gradperp = lambda u: state.perp(grad(u))
elif self.state.mesh.topological_dimension() == 3:
if self.ibp == "twice":
raise NotImplementedError(
"ibp=twice is not implemented for 3d problems")
else:
raise RuntimeError("topological mesh dimension must be 2 or 3")
def vector_invariant_form(state, test, q, ibp=IntegrateByParts.ONCE):
Vu = state.spaces("HDiv")
dS_ = (dS_v + dS_h) if Vu.extruded else dS
ubar = Function(Vu)
n = FacetNormal(state.mesh)
Upwind = 0.5 * (sign(dot(ubar, n)) + 1)
if state.mesh.topological_dimension() == 3:
if ibp != IntegrateByParts.ONCE:
raise NotImplementedError
# <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)
L = (inner(q, curl(cross(ubar, test))) * dx -
inner(both(Upwind * q), both(cross(n, cross(ubar, test)))) * dS_)
else:
perp = state.perp
if state.on_sphere:
outward_normals = CellNormal(state.mesh)
perp_u_upwind = lambda q: Upwind('+') * cross(
outward_normals('+'), q('+')) + Upwind('-') * cross(
outward_normals('-'), q('-'))
else:
perp_u_upwind = lambda q: Upwind('+') * perp(q('+')) + Upwind(
'-') * perp(q('-'))
if ibp == IntegrateByParts.ONCE:
L = (-inner(perp(grad(inner(test, perp(ubar)))), q) * dx -
inner(jump(inner(test, perp(ubar)), n), perp_u_upwind(q)) *
dS_)
else:
L = ((-inner(test,
div(perp(q)) * perp(ubar))) * dx -
inner(jump(inner(test, perp(ubar)), n), perp_u_upwind(q)) *
dS_ + jump(inner(test, perp(ubar)) * perp(q), n) * dS_)
L -= 0.5 * div(test) * inner(q, ubar) * dx
form = transporting_velocity(L, ubar)
return transport(form, TransportEquationType.vector_invariant)
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 advection_term(self, q):
n = FacetNormal(self.state.mesh)
Upwind = 0.5 * (sign(dot(self.ubar, n)) + 1)
if self.state.mesh.topological_dimension() == 3:
# <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)
L = (inner(q, curl(cross(self.ubar, self.test))) * dx -
inner(both(Upwind * q),
both(cross(n, cross(self.ubar, self.test)))) * self.dS)
else:
perp = self.state.perp
if self.state.on_sphere:
outward_normals = CellNormal(self.state.mesh)
perp_u_upwind = lambda q: Upwind('+') * cross(
outward_normals('+'), q('+')) + Upwind('-') * cross(
outward_normals('-'), q('-'))
else:
perp_u_upwind = lambda q: Upwind('+') * perp(q('+')) + Upwind(
'-') * perp(q('-'))
if self.ibp == IntegrateByParts.ONCE:
L = (-inner(perp(grad(inner(self.test, perp(self.ubar)))), q) *
dx - inner(jump(inner(self.test, perp(self.ubar)), n),
perp_u_upwind(q)) * self.dS)
else:
L = ((-inner(self.test,
div(perp(q)) * perp(self.ubar))) * dx -
inner(jump(inner(self.test, perp(self.ubar)), n),
perp_u_upwind(q)) * self.dS +
jump(inner(self.test, perp(self.ubar)) * perp(q), n) *
self.dS)
L -= 0.5 * div(self.test) * inner(q, self.ubar) * dx
return L
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 _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 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 coriolis_term(self):
u0 = split(self.x0)[0]
return -inner(self.test, cross(2 * self.state.Omega, u0)) * dx
def __init__(self,
mesh,
dt,
output=None,
parameters=None,
diagnostics=None,
diagnostic_fields=None):
if output is None:
raise RuntimeError(
"You must provide a directory name for dumping results")
else:
self.output = output
self.parameters = parameters
if diagnostics is not None:
self.diagnostics = diagnostics
else:
self.diagnostics = Diagnostics()
if diagnostic_fields is not None:
self.diagnostic_fields = diagnostic_fields
else:
self.diagnostic_fields = []
# The mesh
self.mesh = mesh
self.spaces = SpaceCreator(mesh)
if self.output.dumplist is None:
self.output.dumplist = []
self.fields = StateFields(*self.output.dumplist)
self.dumpdir = None
self.dumpfile = None
self.to_pickup = None
# figure out if we're on a sphere
try:
self.on_sphere = (mesh._base_mesh.geometric_dimension() == 3
and mesh._base_mesh.topological_dimension() == 2)
except AttributeError:
self.on_sphere = (mesh.geometric_dimension() == 3
and mesh.topological_dimension() == 2)
# build the vertical normal and define perp for 2d geometries
dim = mesh.topological_dimension()
if self.on_sphere:
x = SpatialCoordinate(mesh)
R = sqrt(inner(x, x))
self.k = interpolate(x / R, mesh.coordinates.function_space())
if dim == 2:
outward_normals = CellNormal(mesh)
self.perp = lambda u: cross(outward_normals, u)
else:
kvec = [0.0] * dim
kvec[dim - 1] = 1.0
self.k = Constant(kvec)
if dim == 2:
self.perp = lambda u: as_vector([-u[1], u[0]])
# setup logger
logger.setLevel(output.log_level)
set_log_handler(mesh.comm)
if parameters is not None:
logger.info("Physical parameters that take non-default values:")
logger.info(", ".join("%s: %s" % (k, float(v))
for (k, v) in vars(parameters).items()))
# Constant to hold current time
self.t = Constant(0.0)
if type(dt) is Constant:
self.dt = dt
elif type(dt) in (float, int):
self.dt = Constant(dt)
else:
raise TypeError(
f'dt must be a Constant, float or int, not {type(dt)}')
Dn.interpolate(Dexpr)
b.interpolate(bexpr)
if COMM_WORLD.Get_rank() == 0:
print("Finished setting up functions at", ctime())
# Build perp and upwind perp (including 2D version for test purposes)
n = FacetNormal(mesh)
s = lambda u: 0.5 * (sign(dot(u, n)) + 1)
uw = lambda u, v: (s(u)('+') * v('+') + s(u)('-') * v('-'))
if mesh.geometric_dimension() == 2:
perp = lambda u: as_vector([-u[1], u[0]])
p_uw = lambda u, v: perp(uw(u, v))
else:
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()
def perp(u):
return fd.cross(outward_normals, u)
def __init__(self,
mesh,
vertical_degree=None,
horizontal_degree=1,
family="RT",
Coriolis=None,
sponge_function=None,
hydrostatic=None,
timestepping=None,
output=None,
parameters=None,
diagnostics=None,
fieldlist=None,
diagnostic_fields=None,
u_bc_ids=None):
self.family = family
self.vertical_degree = vertical_degree
self.horizontal_degree = horizontal_degree
self.Omega = Coriolis
self.mu = sponge_function
self.hydrostatic = hydrostatic
self.timestepping = timestepping
if output is None:
raise RuntimeError(
"You must provide a directory name for dumping results")
else:
self.output = output
self.parameters = parameters
if fieldlist is None:
raise RuntimeError(
"You must provide a fieldlist containing the names of the prognostic fields"
)
else:
self.fieldlist = fieldlist
if diagnostics is not None:
self.diagnostics = diagnostics
else:
self.diagnostics = Diagnostics(*fieldlist)
if diagnostic_fields is not None:
self.diagnostic_fields = diagnostic_fields
else:
self.diagnostic_fields = []
if u_bc_ids is not None:
self.u_bc_ids = u_bc_ids
else:
self.u_bc_ids = []
# The mesh
self.mesh = mesh
# Build the spaces
self._build_spaces(mesh, vertical_degree, horizontal_degree, family)
# Allocate state
self._allocate_state()
if self.output.dumplist is None:
self.output.dumplist = fieldlist
self.fields = FieldCreator(fieldlist, self.xn, self.output.dumplist)
# set up bcs
V = self.fields('u').function_space()
self.bcs = []
if V.extruded:
self.bcs.append(DirichletBC(V, 0.0, "bottom"))
self.bcs.append(DirichletBC(V, 0.0, "top"))
for id in self.u_bc_ids:
self.bcs.append(DirichletBC(V, 0.0, id))
self.dumpfile = None
# figure out if we're on a sphere
try:
self.on_sphere = (mesh._base_mesh.geometric_dimension() == 3
and mesh._base_mesh.topological_dimension() == 2)
except AttributeError:
self.on_sphere = (mesh.geometric_dimension() == 3
and mesh.topological_dimension() == 2)
# build the vertical normal and define perp for 2d geometries
dim = mesh.topological_dimension()
if self.on_sphere:
x = SpatialCoordinate(mesh)
R = sqrt(inner(x, x))
self.k = interpolate(x / R, mesh.coordinates.function_space())
if dim == 2:
outward_normals = CellNormal(mesh)
self.perp = lambda u: cross(outward_normals, u)
else:
kvec = [0.0] * dim
kvec[dim - 1] = 1.0
self.k = Constant(kvec)
if dim == 2:
self.perp = lambda u: as_vector([-u[1], u[0]])
# project test function for hydrostatic case
if self.hydrostatic:
self.h_project = lambda u: u - self.k * inner(u, self.k)
else:
self.h_project = lambda u: u
# Constant to hold current time
self.t = Constant(0.0)
# setup logger
logger.setLevel(output.log_level)
set_log_handler(mesh.comm)
logger.info("Timestepping parameters that take non-default values:")
logger.info(", ".join("%s: %s" % item
for item in vars(timestepping).items()))
if parameters is not None:
logger.info("Physical parameters that take non-default values:")
logger.info(", ".join("%s: %s" % item
for item in vars(parameters).items()))
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 __init__(
self,
state,
family,
degree,
Omega=None,
terms_to_linearise={
'u': [time_derivative],
'p': [time_derivative],
'b': [time_derivative, transport]
},
u_transport_option="vector_invariant_form",
no_normal_flow_bc_ids=None,
active_tracers=None):
field_names = ['u', 'p', 'b']
if active_tracers is not None:
raise NotImplementedError(
'Tracers not implemented for Boussinesq equations')
if active_tracers is None:
active_tracers = []
super().__init__(field_names,
state,
family,
degree,
terms_to_linearise=terms_to_linearise,
no_normal_flow_bc_ids=no_normal_flow_bc_ids,
active_tracers=active_tracers)
w, phi, gamma = self.tests[0:3]
u, p, b = split(self.X)
bbar = state.fields("bbar", space=state.spaces("theta"), dump=False)
bbar = state.fields("pbar", space=state.spaces("DG"), dump=False)
# -------------------------------------------------------------------- #
# Time Derivative Terms
# -------------------------------------------------------------------- #
mass_form = self.generate_mass_terms()
# -------------------------------------------------------------------- #
# Transport Terms
# -------------------------------------------------------------------- #
# Velocity transport term -- depends on formulation
if u_transport_option == "vector_invariant_form":
u_adv = prognostic(vector_invariant_form(state, w, u), "u")
elif u_transport_option == "vector_advection_form":
u_adv = prognostic(advection_form(state, w, u), "u")
elif u_transport_option == "vector_manifold_advection_form":
u_adv = prognostic(vector_manifold_advection_form(state, w, u),
"u")
elif u_transport_option == "circulation_form":
ke_form = kinetic_energy_form(state, w, u)
ke_form = transport.remove(ke_form)
ke_form = ke_form.label_map(
lambda t: t.has_label(transporting_velocity), lambda t: Term(
ufl.replace(t.form, {t.get(transporting_velocity): u}), t.
labels))
ke_form = transporting_velocity.remove(ke_form)
u_adv = advection_equation_circulation_form(state, w, u) + ke_form
else:
raise ValueError("Invalid u_transport_option: %s" %
u_transport_option)
# Buoyancy transport
b_adv = prognostic(advection_form(state, gamma, b), "b")
linear_b_adv = linear_advection_form(state, gamma, bbar).label_map(
lambda t: t.has_label(transporting_velocity), lambda t: Term(
ufl.replace(t.form, {
t.get(transporting_velocity): self.trials[0]
}), t.labels))
b_adv = linearisation(b_adv, linear_b_adv)
adv_form = subject(u_adv + b_adv, self.X)
# Add transport of tracers
if len(active_tracers) > 0:
adv_form += self.generate_tracer_transport_terms(
state, active_tracers)
# -------------------------------------------------------------------- #
# Pressure Gradient Term
# -------------------------------------------------------------------- #
pressure_gradient_form = subject(prognostic(-div(w) * p * dx, "u"),
self.X)
# -------------------------------------------------------------------- #
# Gravitational Term
# -------------------------------------------------------------------- #
gravity_form = subject(prognostic(-b * inner(w, state.k) * dx, "u"),
self.X)
# -------------------------------------------------------------------- #
# Divergence Term
# -------------------------------------------------------------------- #
# This enforces that div(u) = 0
# The p features here so that the div(u) evaluated in the "forcing" step
# replaces the whole pressure field, rather than merely providing an
# increment to it.
divergence_form = name(
subject(prognostic(phi * (p - div(u)) * dx, "p"), self.X),
"incompressibility")
residual = (mass_form + adv_form + divergence_form +
pressure_gradient_form + gravity_form)
# -------------------------------------------------------------------- #
# Extra Terms (Coriolis)
# -------------------------------------------------------------------- #
if Omega is not None:
residual += subject(
prognostic(inner(w, cross(2 * Omega, u)) * dx, "u"), self.X)
# -------------------------------------------------------------------- #
# Linearise equations
# -------------------------------------------------------------------- #
# TODO: add linearisation states for variables
# Add linearisations to equations
self.residual = self.generate_linear_terms(residual,
self.terms_to_linearise)
def __init__(
self,
state,
family,
degree,
Omega=None,
sponge=None,
extra_terms=None,
terms_to_linearise={
'u': [time_derivative],
'rho': [time_derivative, transport],
'theta': [time_derivative, transport]
},
u_transport_option="vector_invariant_form",
diffusion_options=None,
no_normal_flow_bc_ids=None,
active_tracers=None):
field_names = ['u', 'rho', 'theta']
if active_tracers is None:
active_tracers = []
super().__init__(field_names,
state,
family,
degree,
terms_to_linearise=terms_to_linearise,
no_normal_flow_bc_ids=no_normal_flow_bc_ids,
active_tracers=active_tracers)
g = state.parameters.g
cp = state.parameters.cp
w, phi, gamma = self.tests[0:3]
u, rho, theta = split(self.X)[0:3]
rhobar = state.fields("rhobar", space=state.spaces("DG"), dump=False)
thetabar = state.fields("thetabar",
space=state.spaces("theta"),
dump=False)
zero_expr = Constant(0.0) * theta
exner = exner_pressure(state.parameters, rho, theta)
n = FacetNormal(state.mesh)
# -------------------------------------------------------------------- #
# Time Derivative Terms
# -------------------------------------------------------------------- #
mass_form = self.generate_mass_terms()
# -------------------------------------------------------------------- #
# Transport Terms
# -------------------------------------------------------------------- #
# Velocity transport term -- depends on formulation
if u_transport_option == "vector_invariant_form":
u_adv = prognostic(vector_invariant_form(state, w, u), "u")
elif u_transport_option == "vector_advection_form":
u_adv = prognostic(advection_form(state, w, u), "u")
elif u_transport_option == "vector_manifold_advection_form":
u_adv = prognostic(vector_manifold_advection_form(state, w, u),
"u")
elif u_transport_option == "circulation_form":
ke_form = kinetic_energy_form(state, w, u)
ke_form = transport.remove(ke_form)
ke_form = ke_form.label_map(
lambda t: t.has_label(transporting_velocity), lambda t: Term(
ufl.replace(t.form, {t.get(transporting_velocity): u}), t.
labels))
ke_form = transporting_velocity.remove(ke_form)
u_adv = advection_equation_circulation_form(state, w, u) + ke_form
else:
raise ValueError("Invalid u_transport_option: %s" %
u_transport_option)
# Density transport (conservative form)
rho_adv = prognostic(continuity_form(state, phi, rho), "rho")
# Transport term needs special linearisation
if transport in terms_to_linearise['rho']:
linear_rho_adv = linear_continuity_form(
state, phi, rhobar).label_map(
lambda t: t.has_label(transporting_velocity),
lambda t: Term(
ufl.replace(t.form, {
t.get(transporting_velocity):
self.trials[0]
}), t.labels))
rho_adv = linearisation(rho_adv, linear_rho_adv)
# Potential temperature transport (advective form)
theta_adv = prognostic(advection_form(state, gamma, theta), "theta")
# Transport term needs special linearisation
if transport in terms_to_linearise['theta']:
linear_theta_adv = linear_advection_form(
state, gamma, thetabar).label_map(
lambda t: t.has_label(transporting_velocity),
lambda t: Term(
ufl.replace(t.form, {
t.get(transporting_velocity):
self.trials[0]
}), t.labels))
theta_adv = linearisation(theta_adv, linear_theta_adv)
adv_form = subject(u_adv + rho_adv + theta_adv, self.X)
# Add transport of tracers
if len(active_tracers) > 0:
adv_form += self.generate_tracer_transport_terms(
state, active_tracers)
# -------------------------------------------------------------------- #
# Pressure Gradient Term
# -------------------------------------------------------------------- #
# First get total mass of tracers
tracer_mr_total = zero_expr
for tracer in active_tracers:
if tracer.variable_type == TracerVariableType.mixing_ratio:
idx = self.field_names.index(tracer.name)
tracer_mr_total += split(self.X)[idx]
else:
raise NotImplementedError(
'Only mixing ratio tracers are implemented')
theta_v = theta / (Constant(1.0) + tracer_mr_total)
pressure_gradient_form = name(
subject(
prognostic(
cp * (-div(theta_v * w) * exner * dx +
jump(theta_v * w, n) * avg(exner) * dS_v), "u"),
self.X), "pressure_gradient")
# -------------------------------------------------------------------- #
# Gravitational Term
# -------------------------------------------------------------------- #
gravity_form = subject(
prognostic(Term(g * inner(state.k, w) * dx), "u"), self.X)
residual = (mass_form + adv_form + pressure_gradient_form +
gravity_form)
# -------------------------------------------------------------------- #
# Moist Thermodynamic Divergence Term
# -------------------------------------------------------------------- #
if len(active_tracers) > 0:
cv = state.parameters.cv
c_vv = state.parameters.c_vv
c_pv = state.parameters.c_pv
c_pl = state.parameters.c_pl
R_d = state.parameters.R_d
R_v = state.parameters.R_v
# Get gas and liquid moisture mixing ratios
mr_l = zero_expr
mr_v = zero_expr
for tracer in active_tracers:
if tracer.is_moisture:
if tracer.variable_type == TracerVariableType.mixing_ratio:
idx = self.field_names.index(tracer.name)
if tracer.phase == Phases.gas:
mr_v += split(self.X)[idx]
elif tracer.phase == Phases.liquid:
mr_l += split(self.X)[idx]
else:
raise NotImplementedError(
'Only mixing ratio tracers are implemented')
c_vml = cv + mr_v * c_vv + mr_l * c_pl
c_pml = cp + mr_v * c_pv + mr_l * c_pl
R_m = R_d + mr_v * R_v
residual += subject(
prognostic(
gamma * theta * div(u) * (R_m / c_vml - (R_d * c_pml) /
(cp * c_vml)) * dx, "theta"),
self.X)
# -------------------------------------------------------------------- #
# Extra Terms (Coriolis, Sponge, Diffusion and others)
# -------------------------------------------------------------------- #
if Omega is not None:
residual += subject(
prognostic(inner(w, cross(2 * Omega, u)) * dx, "u"), self.X)
if sponge is not None:
W_DG = FunctionSpace(state.mesh, "DG", 2)
x = SpatialCoordinate(state.mesh)
z = x[len(x) - 1]
H = sponge.H
zc = sponge.z_level
assert float(zc) < float(
H
), "you have set the sponge level above the height of your domain"
mubar = sponge.mubar
muexpr = conditional(
z <= zc, 0.0,
mubar * sin((pi / 2.) * (z - zc) / (H - zc))**2)
self.mu = Function(W_DG).interpolate(muexpr)
residual += name(
subject(
prognostic(
self.mu * inner(w, state.k) * inner(u, state.k) * dx,
"u"), self.X), "sponge")
if diffusion_options is not None:
for field, diffusion in diffusion_options:
idx = self.field_names.index(field)
test = self.tests[idx]
fn = split(self.X)[idx]
residual += subject(
prognostic(
interior_penalty_diffusion_form(
state, test, fn, diffusion), field), self.X)
if extra_terms is not None:
for field, term in extra_terms:
idx = self.field_names.index(field)
test = self.tests[idx]
residual += subject(prognostic(inner(test, term) * dx, field),
self.X)
# -------------------------------------------------------------------- #
# Linearise equations
# -------------------------------------------------------------------- #
# TODO: add linearisation states for variables
# Add linearisations to equations
self.residual = self.generate_linear_terms(residual,
self.terms_to_linearise)
