def AdvectionDiffusionGLS(
V: fd.FunctionSpace,
theta: fd.Function,
phi: fd.Function,
PeInv: float = 1e-4,
phi_t: fd.Function = None,
):
PeInv_ct = fd.Constant(PeInv)
rho = fd.TestFunction(V)
F = (inner(theta, grad(phi)) * rho +
PeInv_ct * inner(grad(phi), grad(rho))) * dx
if phi_t:
F += phi_t * rho * dx
h = fd.CellDiameter(V.ufl_domain())
R_U = dot(theta, grad(phi)) - PeInv_ct * div(grad(phi))
if phi_t:
R_U += phi_t
beta_gls = 0.9
tau_gls = beta_gls * ((4.0 * dot(theta, theta) / h**2) + 9.0 *
(4.0 * PeInv_ct / h**2)**2)**(-0.5)
theta_U = dot(theta, grad(rho)) - PeInv_ct * div(grad(rho))
F += tau_gls * inner(R_U, theta_U) * dx()
return F
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 _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, 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 _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 advection_term(self, q):
if self.state.mesh.topological_dimension() == 3:
# <w,curl(u) cross ubar + grad( u.ubar)>
# =<curl(u),ubar cross w> - <div(w), u.ubar>
# =<u,curl(ubar cross w)> -
# <<u_upwind, [[n cross(ubar cross w)cross]]>>
both = lambda u: 2 * avg(u)
L = (inner(q, curl(cross(self.ubar, self.test))) * dx -
inner(both(self.Upwind * q),
both(cross(self.n, cross(self.ubar, self.test)))) *
self.dS)
else:
if self.ibp == "once":
L = (-inner(
self.gradperp(inner(self.test, self.perp(self.ubar))), q) *
dx - inner(
jump(inner(self.test, self.perp(self.ubar)), self.n),
self.perp_u_upwind(q)) * self.dS)
else:
L = (
(-inner(self.test,
div(self.perp(q)) * self.perp(self.ubar))) * dx -
inner(jump(inner(self.test, self.perp(self.ubar)), self.n),
self.perp_u_upwind(q)) * self.dS + jump(
inner(self.test, self.perp(self.ubar)) *
self.perp(q), self.n) * self.dS)
L -= 0.5 * div(self.test) * inner(q, self.ubar) * dx
return L
def setup_solver(self, up_init=None):
""" Setup the solvers
"""
self.up0 = Function(self.W)
if up_init is not None:
chk_in = checkpointing.HDF5File(up_init, file_mode='r')
chk_in.read(self.up0, "/up")
chk_in.close()
self.u0, self.p0 = split(self.up0)
self.up = Function(self.W)
if up_init is not None:
chk_in = checkpointing.HDF5File(up_init, file_mode='r')
chk_in.read(self.up, "/up")
chk_in.close()
self.u1, self.p1 = split(self.up)
self.up.sub(0).rename("velocity")
self.up.sub(1).rename("pressure")
v, q = TestFunctions(self.W)
h = CellVolume(self.mesh)
u_norm = sqrt(dot(self.u0, self.u0))
if self.has_nullspace:
nullspace = MixedVectorSpaceBasis(
self.W,
[self.W.sub(0), VectorSpaceBasis(constant=True)])
else:
nullspace = None
tau = ((2.0 / self.dt)**2 + (2.0 * u_norm / h)**2 +
(4.0 * self.nu / h**2)**2)**(-0.5)
# temporal discretization
F = (1.0 / self.dt) * inner(self.u1 - self.u0, v) * dx
# weak form
F += (+inner(dot(self.u0, nabla_grad(self.u1)), v) * dx +
self.nu * inner(grad(self.u1), grad(v)) * dx -
(1.0 / self.rho) * self.p1 * div(v) * dx +
div(self.u1) * q * dx - inner(self.forcing, v) * dx)
# residual form
R = (+(1.0 / self.dt) * (self.u1 - self.u0) +
dot(self.u0, nabla_grad(self.u1)) - self.nu * div(grad(self.u1)) +
(1.0 / self.rho) * grad(self.p1) - self.forcing)
# GLS
F += tau * inner(
+dot(self.u0, nabla_grad(v)) - self.nu * div(grad(v)) +
(1.0 / self.rho) * grad(q), R) * dx
self.problem = NonlinearVariationalProblem(F, self.up, self.bcs)
self.solver = NonlinearVariationalSolver(
self.problem,
nullspace=nullspace,
solver_parameters=self.solver_parameters)
def get_weak_form(self):
(v, q) = fd.TestFunctions(self.V)
(u, p) = fd.split(self.solution)
F = self.nu * fd.inner(fd.grad(u), fd.grad(v)) * fd.dx \
- p * fd.div(v) * fd.dx \
+ fd.div(u) * q * fd.dx \
+ fd.inner(fda.Constant((0., 0.)), v) * fd.dx
return F
def incompressible_hydrostatic_balance(state, b0, p0, top=False, params=None):
# get F
Vu = state.spaces("HDiv")
Vv = FunctionSpace(state.mesh, Vu.ufl_element()._elements[-1])
v = TrialFunction(Vv)
w = TestFunction(Vv)
if top:
bstring = "top"
else:
bstring = "bottom"
bcs = [DirichletBC(Vv, 0.0, bstring)]
a = inner(w, v) * dx
L = inner(state.k, w) * b0 * dx
F = Function(Vv)
solve(a == L, F, bcs=bcs)
# define mixed function space
VDG = state.spaces("DG")
WV = (Vv) * (VDG)
# get pprime
v, pprime = TrialFunctions(WV)
w, phi = TestFunctions(WV)
bcs = [DirichletBC(WV[0], zero(), bstring)]
a = (inner(w, v) + div(w) * pprime + div(v) * phi) * dx
L = phi * div(F) * dx
w1 = Function(WV)
if params is None:
params = {
'ksp_type': 'gmres',
'pc_type': 'fieldsplit',
'pc_fieldsplit_type': 'schur',
'pc_fieldsplit_schur_fact_type': 'full',
'pc_fieldsplit_schur_precondition': 'selfp',
'fieldsplit_1_ksp_type': 'preonly',
'fieldsplit_1_pc_type': 'gamg',
'fieldsplit_1_mg_levels_pc_type': 'bjacobi',
'fieldsplit_1_mg_levels_sub_pc_type': 'ilu',
'fieldsplit_0_ksp_type': 'richardson',
'fieldsplit_0_ksp_max_it': 4,
'ksp_atol': 1.e-08,
'ksp_rtol': 1.e-08
}
solve(a == L, w1, bcs=bcs, solver_parameters=params)
v, pprime = w1.split()
p0.project(pprime)
def stokes(self):
v_ = self.v
if self.opt['form'] == 'linear' or self.opt['form'] == 'nonlinear':
u_ = self.sol.split()
elif self.opt['form'] == 'bilinear':
u_ = self.u
l = self.nu * inner(grad(u_[0]), grad(v_[0])) * dx
b = -div(v_[0]) * u_[1] * dx
bt = -div(u_[0]) * v_[1] * dx
return l + b + bt
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 _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 dgls_form(self, problem, mesh, bcs_p):
rho = problem.rho
mu = problem.mu
k = problem.k
f = problem.f
q, p = fire.TrialFunctions(self._W)
w, v = fire.TestFunctions(self._W)
n = fire.FacetNormal(mesh)
h = fire.CellDiameter(mesh)
# Stabilizing parameters
has_mesh_characteristic_length = True
delta_0 = fire.Constant(1)
delta_1 = fire.Constant(-1 / 2)
delta_2 = fire.Constant(1 / 2)
delta_3 = fire.Constant(1 / 2)
eta_p = fire.Constant(100)
eta_q = fire.Constant(100)
h_avg = (h("+") + h("-")) / 2.0
if has_mesh_characteristic_length:
delta_2 = delta_2 * h * h
delta_3 = delta_3 * h * h
kappa = rho * k / mu
inv_kappa = 1.0 / kappa
# Classical mixed terms
a = (dot(inv_kappa * q, w) - div(w) * p - delta_0 * v * div(q)) * dx
L = -delta_0 * f * v * dx
# DG terms
a += jump(w, n) * avg(p) * dS - avg(v) * jump(q, n) * dS
# Edge stabilizing terms
a += (eta_q * h_avg) * avg(inv_kappa) * (
jump(q, n) * jump(w, n)) * dS + (eta_p / h_avg) * avg(kappa) * dot(
jump(v, n), jump(p, n)) * dS
# Add the contributions of the pressure boundary conditions to L
for pboundary, iboundary in bcs_p:
L -= pboundary * dot(w, n) * ds(iboundary)
# Stabilizing terms
a += (delta_1 * inner(kappa * (inv_kappa * q + grad(p)),
delta_0 * inv_kappa * w + grad(v)) * dx)
a += delta_2 * inv_kappa * div(q) * div(w) * dx
a += delta_3 * inner(kappa * curl(inv_kappa * q), curl(
inv_kappa * w)) * dx
L += delta_2 * inv_kappa * f * div(w) * dx
return a, L
def vector_invariant_form(state, test, q, ibp=IntegrateByParts.ONCE):
Vu = state.spaces("HDiv")
dS_ = (dS_v + dS_h) if Vu.extruded else dS
ubar = Function(Vu)
n = FacetNormal(state.mesh)
Upwind = 0.5 * (sign(dot(ubar, n)) + 1)
if state.mesh.topological_dimension() == 3:
if ibp != IntegrateByParts.ONCE:
raise NotImplementedError
# <w,curl(u) cross ubar + grad( u.ubar)>
# =<curl(u),ubar cross w> - <div(w), u.ubar>
# =<u,curl(ubar cross w)> -
# <<u_upwind, [[n cross(ubar cross w)cross]]>>
both = lambda u: 2 * avg(u)
L = (inner(q, curl(cross(ubar, test))) * dx -
inner(both(Upwind * q), both(cross(n, cross(ubar, test)))) * dS_)
else:
perp = state.perp
if state.on_sphere:
outward_normals = CellNormal(state.mesh)
perp_u_upwind = lambda q: Upwind('+') * cross(
outward_normals('+'), q('+')) + Upwind('-') * cross(
outward_normals('-'), q('-'))
else:
perp_u_upwind = lambda q: Upwind('+') * perp(q('+')) + Upwind(
'-') * perp(q('-'))
if ibp == IntegrateByParts.ONCE:
L = (-inner(perp(grad(inner(test, perp(ubar)))), q) * dx -
inner(jump(inner(test, perp(ubar)), n), perp_u_upwind(q)) *
dS_)
else:
L = ((-inner(test,
div(perp(q)) * perp(ubar))) * dx -
inner(jump(inner(test, perp(ubar)), n), perp_u_upwind(q)) *
dS_ + jump(inner(test, perp(ubar)) * perp(q), n) * dS_)
L -= 0.5 * div(test) * inner(q, ubar) * dx
form = transporting_velocity(L, ubar)
return transport(form, TransportEquationType.vector_invariant)
def __init__(self, mesh_m, viscosity):
super().__init__()
self.mesh_m = mesh_m
self.failed_to_solve = False # when self.solver.solve() fail
# Setup problem, Taylor-Hood finite elements
self.V = fd.VectorFunctionSpace(self.mesh_m, "CG", 2) \
* fd.FunctionSpace(self.mesh_m, "CG", 1)
# Preallocate solution variables for state equation
self.solution = fd.Function(self.V, name="State")
self.testfunction = fd.TestFunction(self.V)
# Define viscosity parameter
self.viscosity = viscosity
# Weak form of incompressible Navier-Stokes equations
z = self.solution
u, p = fd.split(z)
test = self.testfunction
v, q = fd.split(test)
nu = self.viscosity # shorten notation
self.F = nu*fd.inner(fd.grad(u), fd.grad(v))*fd.dx - p*fd.div(v)*fd.dx\
+ fd.inner(fd.dot(fd.grad(u), u), v)*fd.dx + fd.div(u)*q*fd.dx
# Dirichlet Boundary conditions
X = fd.SpatialCoordinate(self.mesh_m)
dim = self.mesh_m.topological_dimension()
if dim == 2:
uin = 4 * fd.as_vector([(1 - X[1]) * X[1], 0])
elif dim == 3:
rsq = X[0]**2 + X[1]**2 # squared radius = 0.5**2 = 1/4
uin = fd.as_vector([0, 0, 1 - 4 * rsq])
else:
raise NotImplementedError
self.bcs = [
fd.DirichletBC(self.V.sub(0), 0., [12, 13]),
fd.DirichletBC(self.V.sub(0), uin, [10])
]
# PDE-solver parameters
self.nsp = None
self.params = {
"snes_max_it": 10,
"mat_type": "aij",
"pc_type": "lu",
"pc_factor_mat_solver_type": "superlu_dist",
# "snes_monitor": None, "ksp_monitor": None,
}
def continuity_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False):
if outflow and ibp == IntegrateByParts.NEVER:
raise ValueError(
"outflow is True and ibp is None are incompatible options")
Vu = state.spaces("HDiv")
dS_ = (dS_v + dS_h) if Vu.extruded else dS
ubar = Function(Vu)
if ibp == IntegrateByParts.ONCE:
L = -inner(grad(test), outer(q, ubar)) * dx
else:
L = inner(test, div(outer(q, ubar))) * dx
if ibp != IntegrateByParts.NEVER:
n = FacetNormal(state.mesh)
un = 0.5 * (dot(ubar, n) + abs(dot(ubar, n)))
L += dot(jump(test), (un('+') * q('+') - un('-') * q('-'))) * dS_
if ibp == IntegrateByParts.TWICE:
L -= (inner(test('+'),
dot(ubar('+'), n('+')) * q('+')) +
inner(test('-'),
dot(ubar('-'), n('-')) * q('-'))) * dS_
if outflow:
n = FacetNormal(state.mesh)
un = 0.5 * (dot(ubar, n) + abs(dot(ubar, n)))
L += test * un * q * (ds_v + ds_t + ds_b)
form = transporting_velocity(L, ubar)
return ibp_label(transport(form, TransportEquationType.conservative), ibp)
def residual(self, test, trial, trial_lagged, fields, bcs):
u_adv = trial_lagged
phi = test
n = self.n
u = trial
F = -dot(u, div(outer(phi, u_adv)))*self.dx
for id, bc in bcs.items():
if 'u' in bc:
u_in = bc['u']
elif 'un' in bc:
u_in = bc['un'] * n # this implies u_t=0 on the inflow
else:
u_in = zero(self.dim)
F += conditional(dot(u_adv, n) < 0,
dot(phi, u_in)*dot(u_adv, n),
dot(phi, u)*dot(u_adv, n)) * self.ds(id)
if not (is_continuous(self.trial_space) and normal_is_continuous(u_adv)):
# s=0: u.n(-)<0 => flow goes from '+' to '-' => '+' is upwind
# s=1: u.n(-)>0 => flow goes from '-' to '+' => '-' is upwind
s = 0.5*(sign(dot(avg(u), n('-'))) + 1.0)
u_up = u('-')*s + u('+')*(1-s)
F += dot(u_up, (dot(u_adv('+'), n('+'))*phi('+') + dot(u_adv('-'), n('-'))*phi('-'))) * self.dS
return -F
def eady_initial_v(state, p0, v):
f = state.parameters.f
x, y, z = SpatialCoordinate(state.mesh)
# get pressure gradient
Vu = state.spaces("HDiv")
g = TrialFunction(Vu)
wg = TestFunction(Vu)
n = FacetNormal(state.mesh)
a = inner(wg, g)*dx
L = -div(wg)*p0*dx + inner(wg, n)*p0*ds_tb
pgrad = Function(Vu)
solve(a == L, pgrad)
# get initial v
Vp = p0.function_space()
phi = TestFunction(Vp)
m = TrialFunction(Vp)
a = f*phi*m*dx
L = phi*pgrad[0]*dx
solve(a == L, v)
return v
def compressible_eady_initial_v(state, theta0, rho0, v):
f = state.parameters.f
cp = state.parameters.cp
# exner function
Vr = rho0.function_space()
Pi = Function(Vr).interpolate(thermodynamics.pi(state.parameters, rho0, theta0))
# get Pi gradient
Vu = state.spaces("HDiv")
g = TrialFunction(Vu)
wg = TestFunction(Vu)
n = FacetNormal(state.mesh)
a = inner(wg, g)*dx
L = -div(wg)*Pi*dx + inner(wg, n)*Pi*ds_tb
pgrad = Function(Vu)
solve(a == L, pgrad)
# get initial v
m = TrialFunction(Vr)
phi = TestFunction(Vr)
a = phi*f*m*dx
L = phi*cp*theta0*pgrad[0]*dx
solve(a == L, v)
return v
def divMelt(h, floating, meltParams, u, Q):
""" Melt function that is a scaled version of the flux divergence
h : firedrake function
ice thickness
u : firedrake vector function
surface elevation
floating : firedrake function
floating mask
V : firedrake vector space
vector space for velocity
meltParams : dict
parameters for melt function
Returns
-------
firedrake function
melt rates
"""
flux = u * h
fluxDiv = icepack.interpolate(firedrake.div(flux), Q)
fluxDivS = firedrakeSmooth(fluxDiv, alpha=8000)
fluxDivS = firedrake.min_value(
fluxDivS * floating * meltParams['meltMask'], 0)
intFluxDiv = firedrake.assemble(fluxDivS * firedrake.dx)
scale = -1.0 * float(meltParams['intMelt']) / float(intFluxDiv)
scale = firedrake.Constant(scale)
melt = icepack.interpolate(
firedrake.min_value(fluxDivS * scale, meltParams['maxMelt']), Q)
return melt
def div(u):
r"""Compute the horizontal divergence of a velocity field"""
axes = get_mesh_axes(u.ufl_domain())
if axes == "xy":
return firedrake.div(u)
if axes == "xyz":
return u[0].dx(0) + u[1].dx(1)
return u.dx(0)
def kinetic_energy_form(state, test, q):
ubar = Function(state.spaces("HDiv"))
L = 0.5 * div(test) * inner(q, ubar) * dx
form = transporting_velocity(L, ubar)
return transport(form, TransportEquationType.vector_invariant)
def test_solver_no_flow_region():
mesh = fd.Mesh("./2D_mesh.msh")
no_flow = [2]
no_flow_markers = [1]
mesh = mark_no_flow_regions(mesh, no_flow, no_flow_markers)
P2 = fd.VectorElement("CG", mesh.ufl_cell(), 1)
P1 = fd.FiniteElement("CG", mesh.ufl_cell(), 1)
TH = P2 * P1
W = fd.FunctionSpace(mesh, TH)
(v, q) = fd.TestFunctions(W)
# Stokes 1
w_sol1 = fd.Function(W)
nu = fd.Constant(0.05)
F = NavierStokesBrinkmannForm(W, w_sol1, nu, beta_gls=2.0)
x, y = fd.SpatialCoordinate(mesh)
u_mms = fd.as_vector(
[sin(2.0 * pi * x) * sin(pi * y),
sin(pi * x) * sin(2.0 * pi * y)])
p_mms = -0.5 * (u_mms[0]**2 + u_mms[1]**2)
f_mms_u = (grad(u_mms) * u_mms + grad(p_mms) -
2.0 * nu * div(sym(grad(u_mms))))
f_mms_p = div(u_mms)
F += -inner(f_mms_u, v) * dx - f_mms_p * q * dx
bc1 = fd.DirichletBC(W.sub(0), u_mms, "on_boundary")
bc2 = fd.DirichletBC(W.sub(1), p_mms, "on_boundary")
bc_no_flow = InteriorBC(W.sub(0), fd.Constant((0.0, 0.0)), no_flow_markers)
solver_parameters = {"ksp_max_it": 500, "ksp_monitor": None}
problem1 = fd.NonlinearVariationalProblem(F,
w_sol1,
bcs=[bc1, bc2, bc_no_flow])
solver1 = NavierStokesBrinkmannSolver(
problem1,
options_prefix="navier_stokes",
solver_parameters=solver_parameters,
)
solver1.solve()
u_sol, _ = w_sol1.split()
u_mms_func = fd.interpolate(u_mms, W.sub(0))
error = fd.errornorm(u_sol, u_mms_func)
assert error < 0.07
def advection_term(self, q):
n = FacetNormal(self.state.mesh)
Upwind = 0.5 * (sign(dot(self.ubar, n)) + 1)
if self.state.mesh.topological_dimension() == 3:
# <w,curl(u) cross ubar + grad( u.ubar)>
# =<curl(u),ubar cross w> - <div(w), u.ubar>
# =<u,curl(ubar cross w)> -
# <<u_upwind, [[n cross(ubar cross w)cross]]>>
both = lambda u: 2 * avg(u)
L = (inner(q, curl(cross(self.ubar, self.test))) * dx -
inner(both(Upwind * q),
both(cross(n, cross(self.ubar, self.test)))) * self.dS)
else:
perp = self.state.perp
if self.state.on_sphere:
outward_normals = CellNormal(self.state.mesh)
perp_u_upwind = lambda q: Upwind('+') * cross(
outward_normals('+'), q('+')) + Upwind('-') * cross(
outward_normals('-'), q('-'))
else:
perp_u_upwind = lambda q: Upwind('+') * perp(q('+')) + Upwind(
'-') * perp(q('-'))
if self.ibp == IntegrateByParts.ONCE:
L = (-inner(perp(grad(inner(self.test, perp(self.ubar)))), q) *
dx - inner(jump(inner(self.test, perp(self.ubar)), n),
perp_u_upwind(q)) * self.dS)
else:
L = ((-inner(self.test,
div(perp(q)) * perp(self.ubar))) * dx -
inner(jump(inner(self.test, perp(self.ubar)), n),
perp_u_upwind(q)) * self.dS +
jump(inner(self.test, perp(self.ubar)) * perp(q), n) *
self.dS)
L -= 0.5 * div(self.test) * inner(q, self.ubar) * dx
return L
def run_solver(r):
mesh = fd.UnitSquareMesh(2**r, 2**r)
P2 = fd.VectorElement("CG", mesh.ufl_cell(), 1)
P1 = fd.FiniteElement("CG", mesh.ufl_cell(), 1)
TH = P2 * P1
W = fd.FunctionSpace(mesh, TH)
(v, q) = fd.TestFunctions(W)
# Stokes 1
w_sol1 = fd.Function(W)
nu = fd.Constant(0.05)
F = NavierStokesBrinkmannForm(W, w_sol1, nu, beta_gls=2.0)
from firedrake import sin, grad, pi, sym, div, inner
x, y = fd.SpatialCoordinate(mesh)
u_mms = fd.as_vector(
[sin(2.0 * pi * x) * sin(pi * y),
sin(pi * x) * sin(2.0 * pi * y)])
p_mms = -0.5 * (u_mms[0]**2 + u_mms[1]**2)
f_mms_u = (grad(u_mms) * u_mms + grad(p_mms) -
2.0 * nu * div(sym(grad(u_mms))))
f_mms_p = div(u_mms)
F += -inner(f_mms_u, v) * dx - f_mms_p * q * dx
bc1 = fd.DirichletBC(W.sub(0), u_mms, "on_boundary")
bc2 = fd.DirichletBC(W.sub(1), p_mms, "on_boundary")
solver_parameters = {"ksp_max_it": 200}
problem1 = fd.NonlinearVariationalProblem(F, w_sol1, bcs=[bc1, bc2])
solver1 = NavierStokesBrinkmannSolver(
problem1,
options_prefix="navier_stokes",
solver_parameters=solver_parameters,
)
solver1.solve()
u_sol, _ = w_sol1.split()
fd.File("test_u_sol.pvd").write(u_sol)
u_mms_func = fd.interpolate(u_mms, W.sub(0))
error = fd.errornorm(u_sol, u_mms_func)
print(f"Error: {error}")
return error
def topography_term(self):
g = self.state.parameters.g
u0, _ = split(self.x0)
b = self.state.fields("topography")
n = FacetNormal(self.state.mesh)
un = 0.5 * (dot(u0, n) + abs(dot(u0, n)))
L = g * div(self.test) * b * dx - g * inner(
jump(self.test, n),
un('+') * b('+') - un('-') * b('-')) * dS
return L
def pressure_gradient_term(self):
g = self.state.parameters.g
u0, D0 = split(self.x0)
n = FacetNormal(self.state.mesh)
un = 0.5 * (dot(u0, n) + abs(dot(u0, n)))
L = g * (div(self.test) * D0 * dx -
inner(jump(self.test, n),
un('+') * D0('+') - un('-') * D0('-')) * dS)
return L
def solve_something(mesh):
V = fd.FunctionSpace(mesh, "CG", 1)
u = fd.Function(V)
v = fd.TestFunction(V)
x, y = fd.SpatialCoordinate(mesh)
# f = fd.sin(x) * fd.sin(y) + x**2 + y**2
# uex = x**4 * y**4
uex = fd.sin(x) * fd.sin(y) #*(x*y)**3
# def source(xs, ys):
# return 1/((x-xs)**2+(y-ys)**2 + 0.1)
# uex = source(0, 0)
uex = uex - fd.assemble(uex * fd.dx) / fd.assemble(1 * fd.dx(domain=mesh))
# f = fd.conditional(fd.ge(abs(x)-abs(y), 0), 1, 0)
from firedrake import inner, grad, dx, ds, div, sym
eps = fd.Constant(0.0)
f = uex - div(grad(uex)) + eps * div(grad(div(grad(uex))))
n = fd.FacetNormal(mesh)
g = inner(grad(uex), n)
g1 = inner(grad(div(grad(uex))), n)
g2 = div(grad(uex))
# F = 0.1 * inner(u, v) * dx + inner(grad(u), grad(v)) * dx + inner(grad(grad(u)), grad(grad(v))) * dx - f * v * dx - g * v * ds
F = inner(u, v) * dx + inner(grad(u),
grad(v)) * dx - f * v * dx - g * v * ds
F += eps * inner(div(grad(u)), div(grad(v))) * dx
F += eps * g1 * v * ds
F -= eps * g2 * inner(grad(v), n) * ds
# f = -div(grad(uex))
# F = inner(grad(u), grad(v)) * dx - f * v * dx
# bc = fd.DirichletBC(V, uex, "on_boundary")
bc = None
fd.solve(F == 0,
u,
bcs=bc,
solver_parameters={
"ksp_type": "cg",
"ksp_atol": 1e-13,
"ksp_rtol": 1e-13,
"ksp_dtol": 1e-13,
"ksp_stol": 1e-13,
"pc_type": "jacobi",
"pc_factor_mat_solver_type": "mumps",
"snes_type": "ksponly",
"ksp_converged_reason": None
})
print("||u-uex|| =", fd.norm(u - uex))
print("||grad(u-uex)|| =", fd.norm(grad(u - uex)))
print("||grad(grad(u-uex))|| =", fd.norm(grad(grad(u - uex))))
err = fd.Function(
fd.TensorFunctionSpace(mesh, "DG",
V.ufl_element().degree() - 2)).interpolate(
grad(grad(u - uex)))
# err = fd.Function(fd.FunctionSpace(mesh, "DG", V.ufl_element().degree())).interpolate(u-uex)
fd.File(outdir + "sln.pvd").write(u)
fd.File(outdir + "err.pvd").write(err)
def advection_term(self, q):
if self.continuity:
if self.ibp == IntegrateByParts.ONCE:
L = -inner(grad(self.test), outer(q, self.ubar)) * dx
else:
L = inner(self.test, div(outer(q, self.ubar))) * dx
else:
if self.ibp == IntegrateByParts.ONCE:
L = -inner(div(outer(self.test, self.ubar)), q) * dx
else:
L = inner(outer(self.test, self.ubar), grad(q)) * dx
if self.dS is not None and self.ibp != IntegrateByParts.NEVER:
n = FacetNormal(self.state.mesh)
un = 0.5 * (dot(self.ubar, n) + abs(dot(self.ubar, n)))
L += dot(jump(self.test),
(un('+') * q('+') - un('-') * q('-'))) * self.dS
if self.ibp == IntegrateByParts.TWICE:
L -= (inner(self.test('+'),
dot(self.ubar('+'), n('+')) * q('+')) +
inner(self.test('-'),
dot(self.ubar('-'), n('-')) * q('-'))) * self.dS
if self.outflow:
n = FacetNormal(self.state.mesh)
un = 0.5 * (dot(self.ubar, n) + abs(dot(self.ubar, n)))
L += self.test * un * q * (ds_v + ds_t + ds_b)
if self.vector_manifold:
n = FacetNormal(self.state.mesh)
w = self.test
dS = self.dS
u = q
L += un('+') * inner(w('-'),
n('+') + n('-')) * inner(u('+'), n('+')) * dS
L += un('-') * inner(w('+'),
n('+') + n('-')) * inner(u('-'), n('-')) * dS
return L
def form_function(u, h, v, q):
K = 0.5 * fd.inner(u, u)
n = fd.FacetNormal(mesh)
uup = 0.5 * (fd.dot(u, n) + abs(fd.dot(u, n)))
Upwind = 0.5 * (fd.sign(fd.dot(u, n)) + 1)
eqn = (fd.inner(v, f * perp(u)) * fd.dx -
fd.inner(perp(fd.grad(fd.inner(v, perp(u)))), u) * fd.dx +
fd.inner(both(perp(n) * fd.inner(v, perp(u))), both(Upwind * u)) *
fd.dS - fd.div(v) * (g * (h + b) + K) * fd.dx -
fd.inner(fd.grad(q), u) * h * fd.dx + fd.jump(q) *
(uup('+') * h('+') - uup('-') * h('-')) * fd.dS)
return eqn
def form(u, v):
return inner(div(grad(u)), div(grad(v)))*dx \
- inner(avg(div(grad(u))), jump(grad(v), n))*dS \
- inner(jump(grad(u), n), avg(div(grad(v))))*dS \
+ alpha/h*inner(jump(grad(u), n), jump(grad(v), n))*dS \
- inner(div(grad(u)), inner(grad(v), n))*ds \
- inner(inner(grad(u), n), div(grad(v)))*ds \
+ alpha/h*inner(grad(u), grad(v))*ds
def advection_term(self, q):
if self.continuity:
if self.ibp == "once":
L = -inner(grad(self.test), outer(q, self.ubar)) * dx
else:
L = inner(self.test, div(outer(q, self.ubar))) * dx
else:
if self.ibp == "once":
L = -inner(div(outer(self.test, self.ubar)), q) * dx
else:
L = inner(outer(self.test, self.ubar), grad(q)) * dx
if self.dS is not None and self.ibp is not None:
L += dot(jump(self.test),
(self.un('+') * q('+') - self.un('-') * q('-'))) * self.dS
if self.ibp == "twice":
L -= (
inner(self.test('+'),
dot(self.ubar('+'), self.n('+')) * q('+')) +
inner(self.test('-'),
dot(self.ubar('-'), self.n('-')) * q('-'))) * self.dS
if self.outflow:
L += self.test * self.un * q * self.ds
if self.vector_manifold:
un = self.un
w = self.test
u = q
n = self.n
dS = self.dS
L += un('+') * inner(w('-'),
n('+') + n('-')) * inner(u('+'), n('+')) * dS
L += un('-') * inner(w('+'),
n('+') + n('-')) * inner(u('-'), n('-')) * dS
return L
def derivative(self, out):
super().derivative(out)
if args.discretisation != "pkp0":
return
w = fd.TestFunction(self.V_m)
u = solver.z.split()[0]
v = solver.z_adj.split()[0]
from firedrake import div, cell_avg, dx, tr, grad
gamma = solver.gamma
deriv = gamma * div(w) * cell_avg(div(u)) * div(v) * dx \
+ gamma * (cell_avg(div(u) * div(w) - tr(grad(u)*grad(w)))
- cell_avg(div(u)) * cell_avg(div(w))) * div(v) * dx \
- gamma * cell_avg(div(u)) * tr(grad(v)*grad(w)) * dx
fd.assemble(deriv,
tensor=self.deriv_m,
form_compiler_parameters=self.params)
outcopy = out.clone()
outcopy.from_first_derivative(self.deriv_r)
fd.warning(fd.RED % ("norm of extra term %e" % outcopy.norm()))
out.plus(outcopy)
def heat(n, deg, time_stages, stage_type="deriv", splitting=IA):
N = 2**n
msh = UnitIntervalMesh(N)
params = {
"snes_type": "ksponly",
"ksp_type": "preonly",
"mat_type": "aij",
"pc_type": "lu"
}
V = FunctionSpace(msh, "CG", deg)
x, = SpatialCoordinate(msh)
t = Constant(0.0)
dt = Constant(2.0 / N)
uexact = exp(-t) * sin(pi * x)
rhs = expand_derivatives(diff(uexact, t)) - div(grad(uexact))
butcher_tableau = GaussLegendre(time_stages)
u = project(uexact, V)
v = TestFunction(V)
F = (inner(Dt(u), v) * dx + inner(grad(u), grad(v)) * dx -
inner(rhs, v) * dx)
bc = DirichletBC(V, Constant(0), "on_boundary")
stepper = TimeStepper(F,
butcher_tableau,
t,
dt,
u,
bcs=bc,
solver_parameters=params,
stage_type=stage_type,
splitting=splitting)
while (float(t) < 1.0):
if (float(t) + float(dt) > 1.0):
dt.assign(1.0 - float(t))
stepper.advance()
t.assign(float(t) + float(dt))
return errornorm(uexact, u) / norm(uexact)
def __init__(self, mesh, vertical_degree=1, horizontal_degree=1,
family="RT", z=None, k=None, Omega=None, mu=None,
timestepping=None,
output=None,
parameters=None,
diagnostics=None,
fieldlist=None,
diagnostic_fields=[],
on_sphere=False):
super(BaroclinicState, self).__init__(mesh=mesh,
vertical_degree=vertical_degree,
horizontal_degree=horizontal_degree,
family=family,
z=z, k=k, Omega=Omega, mu=mu,
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist,
diagnostic_fields=diagnostic_fields)
# build the geopotential
if parameters.geopotential:
V = FunctionSpace(mesh, "CG", 1)
if on_sphere:
self.Phi = Function(V).interpolate(Expression("pow(x[0]*x[0]+x[1]*x[1]+x[2]*x[2],0.5)"))
else:
self.Phi = Function(V).interpolate(Expression("x[1]"))
self.Phi *= parameters.g
if self.k is None:
# build the vertical normal
w = TestFunction(self.Vv)
u = TrialFunction(self.Vv)
self.k = Function(self.Vv)
n = FacetNormal(self.mesh)
krhs = -div(w)*self.z*dx + inner(w,n)*self.z*ds_tb
klhs = inner(w,u)*dx
solve(klhs == krhs, self.k)
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 _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 _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 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)
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')