def lhs(self):
l = self.residual.label_map(
lambda t: t.has_label(time_derivative),
map_if_true=replace_subject(self.dq, self.idx),
map_if_false=drop)
return l.form
def rhs(self):
r = self.residual.label_map(
all_terms,
map_if_true=replace_subject(self.q1, self.idx))
r = r.label_map(lambda t: t.has_label(time_derivative),
map_if_false=lambda t: -(1-self.theta)*self.dt*t)
return r.form
def lhs(self):
l = self.residual.label_map(
all_terms,
map_if_true=replace_subject(self.dq, self.idx))
l = l.label_map(lambda t: t.has_label(time_derivative),
map_if_false=lambda t: self.theta*self.dt*t)
return l.form
def hydrostatic_projection(self, t):
# TODO: make this more general, i.e. should work on the sphere
assert not self.state.on_sphere, "the hydrostatic projection is not yet implemented for spherical geometry"
k = Constant((*self.state.k, 0, 0))
X = t.get(subject)
new_subj = X - k * inner(X, k)
return replace_subject(new_subj)(t)
def generate_linear_terms(self, residual, terms_to_linearise):
"""
Generates linear forms for each of the terms in the equation set
(unless specified otherwise). The linear forms are then added to the
terms through a `linearisation` :class:`Label`.
Linear forms are currently generated by replacing the `subject` with its
corresponding `TrialFunction`, however the intention in the future is to
use the `ufl.derivative` to obtain the linear forms.
Terms that already have a `linearisation` label will be left.
:arg residual: The residual of the equation set. A
:class:`LabelledForm` containing all the terms
of the equation set.
:arg terms_to_linearise: A dictionary describing the terms which should
be linearised.
"""
# TODO: Neaten up the `terms_to_linearise` variable. This should not be
# a dictionary, it should be a filter of some sort
# Loop through prognostic variables and linearise any specified terms
for field in self.field_names:
residual = residual.label_map(
# Extract all terms to be linearised for this field's equation that haven't already been
lambda t: (not t.has_label(linearisation) and
(t.get(prognostic) == field) and any(
t.has_label(*terms_to_linearise[field],
return_tuple=True))),
lambda t: linearisation(
t,
LabelledForm(t).label_map(
all_terms, replace_subject(self.trials)).terms[0]))
return residual
def __init__(self, equation, alpha):
self.field_name = equation.field_name
implicit_terms = ["incompressibility", "sponge"]
dt = equation.state.dt
W = equation.function_space
self.x0 = Function(W)
self.xF = Function(W)
# set up boundary conditions on the u subspace of W
bcs = [
DirichletBC(W.sub(0), bc.function_arg, bc.sub_domain)
for bc in equation.bcs['u']
]
# drop terms relating to transport and diffusion
residual = equation.residual.label_map(
lambda t: any(t.has_label(transport, diffusion, return_tuple=True)
), drop)
# the lhs of both of the explicit and implicit solvers is just
# the time derivative form
trials = TrialFunctions(W)
a = residual.label_map(lambda t: t.has_label(time_derivative),
replace_subject(trials),
map_if_false=drop)
# the explicit forms are multiplied by (1-alpha) and moved to the rhs
L_explicit = -(1 - alpha) * dt * residual.label_map(
lambda t: t.has_label(time_derivative) or t.get(name) in
implicit_terms or t.get(name) == "hydrostatic_form", drop,
replace_subject(self.x0))
# the implicit forms are multiplied by alpha and moved to the rhs
L_implicit = -alpha * dt * residual.label_map(
lambda t: t.has_label(time_derivative) or t.get(name) in
implicit_terms or t.get(name) == "hydrostatic_form", drop,
replace_subject(self.x0))
# now add the terms that are always fully implicit
if any(t.get(name) in implicit_terms for t in residual):
L_implicit -= dt * residual.label_map(
lambda t: t.get(name) in implicit_terms,
replace_subject(self.x0), drop)
# the hydrostatic equations require some additional forms:
if any([t.has_label(hydrostatic) for t in residual]):
L_explicit += residual.label_map(
lambda t: t.get(name) == "hydrostatic_form",
replace_subject(self.x0), drop)
L_implicit -= residual.label_map(
lambda t: t.get(name) == "hydrostatic_form",
replace_subject(self.x0), drop)
# now we can set up the explicit and implicit problems
explicit_forcing_problem = LinearVariationalProblem(a.form,
L_explicit.form,
self.xF,
bcs=bcs)
implicit_forcing_problem = LinearVariationalProblem(a.form,
L_implicit.form,
self.xF,
bcs=bcs)
solver_parameters = {}
if logger.isEnabledFor(DEBUG):
solver_parameters["ksp_monitor_true_residual"] = None
self.solvers = {}
self.solvers["explicit"] = LinearVariationalSolver(
explicit_forcing_problem,
solver_parameters=solver_parameters,
options_prefix="ExplicitForcingSolver")
self.solvers["implicit"] = LinearVariationalSolver(
implicit_forcing_problem,
solver_parameters=solver_parameters,
options_prefix="ImplicitForcingSolver")