def __init__(self,
state,
function_space,
field_name,
ufamily=None,
udegree=None,
Vu=None,
diffusion_parameters=None,
**kwargs):
super().__init__(state, function_space, field_name)
if not hasattr(state.fields, "u"):
if Vu is not None:
V = state.spaces("HDiv", V=Vu)
else:
assert ufamily is not None, "Specify the family for u"
assert udegree is not None, "Specify the degree of the u space"
V = state.spaces("HDiv", ufamily, udegree)
state.fields("u", V)
test = TestFunction(function_space)
q = Function(function_space)
mass_form = time_derivative(inner(q, test) * dx)
self.residual = subject(
mass_form + advection_form(state, test, q, **kwargs) +
interior_penalty_diffusion_form(state, test, q,
diffusion_parameters), q)
def __init__(self, state, function_space, field_name,
diffusion_parameters):
super().__init__(state, function_space, field_name)
test = TestFunction(function_space)
q = Function(function_space)
mass_form = time_derivative(inner(q, test) * dx)
self.residual = subject(
mass_form + interior_penalty_diffusion_form(
state, test, q, diffusion_parameters), q)
def generate_tracer_transport_terms(self, state, active_tracers):
"""
Adds the transport forms for the active tracers to the equation.
:arg state: The `State` object.
:arg active_tracers: A list of `ActiveTracer` objects that encode the
metadata for the active tracers.
"""
# By default return None if no tracers are to be transported
adv_form = None
no_tracer_transported = True
for i, tracer in enumerate(active_tracers):
if tracer.transport_flag:
idx = self.field_names.index(tracer.name)
tracer_prog = split(self.X)[idx]
tracer_test = self.tests[idx]
if tracer.transport_eqn == TransportEquationType.advective:
tracer_adv = prognostic(
advection_form(state, tracer_test, tracer_prog),
tracer.name)
elif tracer.transport_eqn == TransportEquationType.conservative:
tracer_adv = prognostic(
continuity_form(state, tracer_test, tracer_prog),
tracer.name)
else:
raise ValueError(
f'Transport eqn {tracer.transport_eqn} not recognised')
if no_tracer_transported:
# We arrive here for the first tracer to be transported
adv_form = subject(tracer_adv, self.X)
no_tracer_transported = False
else:
adv_form += subject(tracer_adv, self.X)
return adv_form
def generate_mass_terms(self):
"""
Builds the weak time derivative terms ("mass terms") for all the
prognostic variables of the equation set. Returns a
:class:`LabelledForm` containing all of these terms.
"""
for i, (test,
field_name) in enumerate(zip(self.tests, self.field_names)):
prog = split(self.X)[i]
mass = subject(prognostic(inner(prog, test) * dx, field_name),
self.X)
if i == 0:
mass_form = time_derivative(mass)
else:
mass_form += time_derivative(mass)
return mass_form
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):
super().__init__(state,
family,
degree,
Omega=Omega,
sponge=sponge,
extra_terms=extra_terms,
terms_to_linearise=terms_to_linearise,
u_transport_option=u_transport_option,
diffusion_options=diffusion_options,
no_normal_flow_bc_ids=no_normal_flow_bc_ids,
active_tracers=active_tracers)
self.residual = self.residual.label_map(
lambda t: t.has_label(time_derivative),
map_if_true=lambda t: hydrostatic(t, self.hydrostatic_projection(t)
))
k = self.state.k
u = split(self.X)[0]
self.residual += name(
subject(
prognostic(-inner(k, self.tests[0]) * inner(k, u) * dx, "u"),
self.X), "hydrostatic_form")
def replace_transport_term(self):
"""
This routine allows the default transport term to be replaced with a
different one, specified through the transport options.
This is necessary because when the prognostic equations are declared,
the whole transport
"""
# Extract transport term of equation
old_transport_term_list = self.residual.label_map(
lambda t: t.has_label(transport), map_if_false=drop)
# If there are more transport terms, extract only the one for this variable
if len(old_transport_term_list.terms) > 1:
raise NotImplementedError('Cannot replace transport terms when there are more than one')
# Then we should only have one transport term
old_transport_term = old_transport_term_list.terms[0]
# If the transport term has an ibp label, then it could be replaced
if old_transport_term.has_label(ibp_label) and hasattr(self.options, 'ibp'):
# Do the options specify a different ibp to the old transport term?
if old_transport_term.labels['ibp'] != self.options.ibp:
# Set up a new transport term
field = self.state.fields(self.field_name)
test = TestFunction(self.fs)
# Set up new transport term (depending on the type of transport equation)
if old_transport_term.labels['transport'] == TransportEquationType.advective:
new_transport_term = advection_form(self.state, test, field, ibp=self.options.ibp)
elif old_transport_term.labels['transport'] == TransportEquationType.conservative:
new_transport_term = continuity_form(self.state, test, field, ibp=self.options.ibp)
else:
raise NotImplementedError(f'Replacement of transport term not implemented yet for {old_transport_term.labels["transport"]}')
# Finally, drop the old transport term and add the new one
self.residual = self.residual.label_map(
lambda t: t.has_label(transport), map_if_true=drop)
self.residual += subject(new_transport_term, field)
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)
def __init__(
self,
state,
family,
degree,
fexpr=None,
bexpr=None,
terms_to_linearise={
'D': [time_derivative, transport],
'u': [time_derivative, pressure_gradient]
},
u_transport_option="vector_invariant_form",
no_normal_flow_bc_ids=None,
active_tracers=None):
field_names = ["u", "D"]
if active_tracers is not None:
raise NotImplementedError(
'Tracers not implemented for shallow water 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)
g = state.parameters.g
H = state.parameters.H
w, phi = self.tests
u, D = split(self.X)
# -------------------------------------------------------------------- #
# 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)
# Depth transport term
D_adv = prognostic(continuity_form(state, phi, D), "D")
# Transport term needs special linearisation
if transport in terms_to_linearise['D']:
linear_D_adv = linear_continuity_form(state, phi, H).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))
# Add linearisation to D_adv
D_adv = linearisation(D_adv, linear_D_adv)
adv_form = subject(u_adv + D_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 = pressure_gradient(
subject(prognostic(-g * div(w) * D * dx, "u"), self.X))
residual = (mass_form + adv_form + pressure_gradient_form)
# -------------------------------------------------------------------- #
# Extra Terms (Coriolis and Topography)
# -------------------------------------------------------------------- #
if fexpr is not None:
V = FunctionSpace(state.mesh, "CG", 1)
f = state.fields("coriolis", space=V)
f.interpolate(fexpr)
coriolis_form = coriolis(
subject(prognostic(f * inner(state.perp(u), w) * dx, "u"),
self.X))
residual += coriolis_form
if bexpr is not None:
b = state.fields("topography", state.spaces("DG"))
b.interpolate(bexpr)
topography_form = subject(prognostic(-g * div(w) * b * dx, "u"),
self.X)
residual += topography_form
# -------------------------------------------------------------------- #
# Linearise equations
# -------------------------------------------------------------------- #
u, D = self.X.split()
# Linearise about D = H
# TODO: add linearisation state for u
D.assign(Constant(H))
u, D = split(self.X)
# Add linearisations to equations
self.residual = self.generate_linear_terms(residual,
self.terms_to_linearise)
