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
L = (- inner(state.perp(grad(gamma)), u))*dx
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"})
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 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
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")
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)
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.solver_parameters,
options_prefix='SWimplicit')
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 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 _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, **kwargs):
for name, field in kwargs.items():
if name in self._fields.keys():
self._fields[name].assign(field)
else:
if isinstance(field, firedrake.Constant):
self._fields[name] = firedrake.Constant(field)
elif isinstance(field, firedrake.Function):
self._fields[name] = field.copy(deepcopy=True)
else:
raise TypeError(
"Input %s field has type %s, must be Constant or Function!"
% (name, type(field))
)
# Create symbolic representations of the flux and sources of damage
dt = firedrake.Constant(1.0)
flux = self.model.flux(**self.fields)
# Create the finite element mass matrix
D = self.fields["damage"]
Q = D.function_space()
φ, ψ = firedrake.TrialFunction(Q), firedrake.TestFunction(Q)
M = φ * ψ * dx
L1 = -dt * flux
D1 = firedrake.Function(Q)
D2 = firedrake.Function(Q)
L2 = firedrake.replace(L1, {D: D1})
L3 = firedrake.replace(L1, {D: D2})
dD = firedrake.Function(Q)
parameters = {
"solver_parameters": {
"ksp_type": "preonly",
"pc_type": "bjacobi",
"sub_pc_type": "ilu",
}
}
problem1 = LinearVariationalProblem(M, L1, dD)
problem2 = LinearVariationalProblem(M, L2, dD)
problem3 = LinearVariationalProblem(M, L3, dD)
solver1 = LinearVariationalSolver(problem1, **parameters)
solver2 = LinearVariationalSolver(problem2, **parameters)
solver3 = LinearVariationalSolver(problem3, **parameters)
self._solvers = [solver1, solver2, solver3]
self._stages = [D1, D2]
self._damage_change = dD
self._timestep = dt
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
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")
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
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 __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 __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)
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 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 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 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)
def solver(self):
# setup solver using lhs and rhs defined in derived class
problem = LinearVariationalProblem(self.lhs, self.rhs, self.dq)
solver_name = self.field.name(
) + self.equation.__class__.__name__ + self.__class__.__name__
return LinearVariationalSolver(
problem,
solver_parameters=self.solver_parameters,
options_prefix=solver_name)
def _setup(self, **kwargs):
for name, field in kwargs.items():
if name in self.fields.keys():
self.fields[name].assign(field)
else:
self.fields[name] = utilities.copy(field)
# Create symbolic representations of the flux and sources of damage
dt = firedrake.Constant(1.)
flux = self.model.flux(**self.fields)
# Create the finite element mass matrix
D = self.fields.get('damage', self.fields.get('D'))
Q = D.function_space()
φ, ψ = firedrake.TrialFunction(Q), firedrake.TestFunction(Q)
M = φ * ψ * dx
L1 = -dt * flux
D1 = firedrake.Function(Q)
D2 = firedrake.Function(Q)
L2 = firedrake.replace(L1, {D: D1})
L3 = firedrake.replace(L1, {D: D2})
dD = firedrake.Function(Q)
parameters = {
'solver_parameters': {
'ksp_type': 'preonly',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
}
}
problem1 = LinearVariationalProblem(M, L1, dD)
problem2 = LinearVariationalProblem(M, L2, dD)
problem3 = LinearVariationalProblem(M, L3, dD)
solver1 = LinearVariationalSolver(problem1, **parameters)
solver2 = LinearVariationalSolver(problem2, **parameters)
solver3 = LinearVariationalSolver(problem3, **parameters)
self._solvers = [solver1, solver2, solver3]
self._stages = [D1, D2]
self._damage_change = dD
self._timestep = dt
def _build_forcing_solvers(self):
a = self.mass_term()
L = self.forcing_term()
if self.Vu.extruded:
bcs = [
DirichletBC(self.Vu, 0.0, "bottom"),
DirichletBC(self.Vu, 0.0, "top")
]
else:
bcs = None
u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs)
solver_parameters = {}
if self.state.output.log_level == DEBUG:
solver_parameters["ksp_monitor_true_residual"] = True
self.u_forcing_solver = LinearVariationalSolver(
u_forcing_problem,
solver_parameters=solver_parameters,
options_prefix="UForcingSolver")
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 _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)
class HydrostaticImbalance(DiagnosticField):
name = "HydrostaticImbalance"
def setup(self, state):
if not self._initialised:
Vu = state.spaces("HDiv")
space = FunctionSpace(state.mesh, Vu.ufl_element()._elements[-1])
super().setup(state, space=space)
rho = state.fields("rho")
rhobar = state.fields("rhobar")
theta = state.fields("theta")
thetabar = state.fields("thetabar")
exner = thermodynamics.exner_pressure(state.parameters, rho, theta)
exnerbar = thermodynamics.exner_pressure(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) * exnerbar * dx + cp * jump(
(theta - thetabar) * w, n) * avg(exnerbar) * dS_v -
cp * div(thetabar * w) * (exner - exnerbar) * dx +
cp * jump(thetabar * w, n) * avg(exner - exnerbar) * 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]
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 _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 self.state.output.log_level == DEBUG:
solver_parameters["ksp_monitor_true_residual"] = True
self.divergence_solver = LinearVariationalSolver(
divergence_problem,
solver_parameters=solver_parameters,
options_prefix="DivergenceSolver")
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)
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 _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 __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')
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 __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 __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 __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 __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')
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)
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)
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)
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")
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 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)
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 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 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 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
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)
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 __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')
