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(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 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 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 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 __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 __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 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 advection_term(self, q):
if self.continuity:
n = FacetNormal(self.state.mesh)
L = (-dot(grad(self.test), self.ubar) * self.qbar * dx +
jump(self.ubar * self.test, n) * avg(self.qbar) * self.dS)
else:
L = self.test * dot(self.ubar, self.state.k) * dot(
self.state.k, grad(self.qbar)) * dx
return L
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 _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 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 linear_continuity_form(state, test, qbar, facet_term=False):
Vu = state.spaces("HDiv")
ubar = Function(Vu)
L = qbar * test * div(ubar) * dx
if facet_term:
n = FacetNormal(state.mesh)
Vu = state.spaces("HDiv")
dS_ = (dS_v + dS_h) if Vu.extruded else dS
L += jump(ubar * test, n) * avg(qbar) * dS_
form = transporting_velocity(L, ubar)
return transport(form, TransportEquationType.conservative)
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 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 __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')
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 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
(1 - fd.exp(-c[0])) * (1 - fd.exp(-c[1])))
source = fd.div(velocity * u_exact) - fd.div(Dm * fd.grad(u_exact))
# stability
tau_d = fd.Constant(max([Dm, tau_e])) / h
tau_a = abs(fd.dot(velocity, n))
# tau = fd.Constant(1.0) / h + vn
# numerical flux
qhat = qh + tau_d * (ch - lmbd_h) * n
fhat = velocity * lmbd_h + tau_a * (ch - lmbd_h) * n
a_u = (
fd.inner(fd.inv(Dm) * qh, vh) * fd.dx - ch * fd.div(vh) * fd.dx +
# internal faces
fd.jump(lmbd_h * vh, n) * fd.dS +
# other faces
lmbd_h * fd.inner(vh, n) * fd.ds)
# Dirichlet faces
L_u = 0
a_c = (-fd.inner(fd.grad(wh), qh + ch * velocity) * fd.dx +
wh('+') * fd.jump(qhat + fhat, n) * fd.dS +
wh * fd.inner(qhat + fhat, n) * fd.ds - wh * source * fd.dx)
L_c = 0
# transmission boundary condition
F_q = (
mu_h('+')*fd.jump(qhat + fhat, n)*fd.dS + \
# Diff = fd.Identity(mesh.geometric_dimension())*(Dm + d_t*vnorm) + \
# fd.Constant(d_l-d_t)*fd.outer(velocity, velocity)/vnorm
# stability
# tau = fd.Constant(1.0) / h + abs(fd.dot(velocity, n))
tau = fd.Constant(5) / h + vn
# numerical flux
chat = lmbd_h
qhat = qh + tau * (ch - chat) * n + velocity * chat
# qhat_n = fd.dot(qh, n) + tau*(ch - chat) + chat*vn
a_u = (
fd.inner(fd.inv(Diff) * qh, vh) * fd.dx - ch * fd.div(vh) * fd.dx +
# internal faces
fd.jump(lmbd_h * vh, n) * fd.dS +
# other faces
lmbd_h * fd.inner(vh, n) * fd.ds(outflow) +
lmbd_h * fd.inner(vh, n) * fd.ds(TOP) +
lmbd_h * fd.inner(vh, n) * fd.ds(BOTTOM))
# Dirichlet faces
L_u = -fd.Constant(cIn) * fd.inner(vh, n) * fd.ds(inlet)
a_c = (wh * (ch - c0) / dtc * fd.dx -
fd.inner(fd.grad(wh), qh + ch * velocity) * fd.dx +
wh("+") * fd.jump(qhat, n) * fd.dS + wh * fd.inner(qhat, n) * fd.ds)
L_c = 0
# transmission boundary condition
D_eqn = (D_ - Dn) * psi * dx - dt * D_ad + inner(v, F_ - Frhs) * dx
Dad_p = LinearVariationalProblem(lhs(D_eqn), rhs(D_eqn), xp)
D_ad_solver = LinearVariationalSolver(Dad_p,
solver_parameters={
"mat_type": "aij",
"snes_linesearch_type": "basic",
"ksp_type": "preonly",
"pc_type": "lu"
})
# u advection solver
w = TestFunction(W1)
u_ = TrialFunction(W1)
u_bar = 0.5 * (un + u_)
u_ad = (inner(perp(grad(inner(Dbar * w, perp(u_rec)))), u_bar) * dx +
inner(jump(inner(Dbar * w, perp(u_rec)), n), p_uw(ubar, u_bar)) * dS -
f * inner(perp(u_rec), Dbar * w) * dx)
u_eqn = inner(u_ - un, Dbar * w) * dx - dt * u_ad
uad_p = LinearVariationalProblem(lhs(u_eqn), rhs(u_eqn), uad)
u_ad_solver = LinearVariationalSolver(uad_p)
# u forcing solver
u_f = div(w) * Prhs * dx
f_eqn = inner(u_, w) * dx - dt * u_f
uf_p = LinearVariationalProblem(lhs(f_eqn), rhs(f_eqn), uf)
f_u_solver = LinearVariationalSolver(uf_p)
# Auxiliary solvers
Frhs = unk * Dnk / 3. + un * Dnk / 6. + unk * Dn / 6. + un * Dn / 3.
u_rec_eqn = inner(w, Dbar * u_ - Frhs) * dx
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
Vu = state.spaces("HDiv")
Vtheta = state.spaces("HDiv_v")
Vrho = state.spaces("DG")
# Store time-stepping coefficients as UFL Constants
dt = Constant(Dt)
beta = Constant(beta_)
beta_cp = Constant(beta_ * cp)
# 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((Vu, Vrho))
w, phi = TestFunctions(M)
u, rho = TrialFunctions(M)
n = FacetNormal(state.mesh)
# Get background fields
thetabar = state.fields("thetabar")
rhobar = state.fields("rhobar")
pibar = thermodynamics.pi(state.parameters, rhobar, thetabar)
pibar_rho = thermodynamics.pi_rho(state.parameters, rhobar, thetabar)
pibar_theta = thermodynamics.pi_theta(state.parameters, rhobar, thetabar)
# 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)
# specify degree for some terms as estimated degree is too large
dxp = dx(degree=(self.quadrature_degree))
dS_vp = dS_v(degree=(self.quadrature_degree))
# 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_w = theta / (1 + water_t)
thetabar_w = thetabar / (1 + water_t)
else:
theta_w = theta
thetabar_w = thetabar
eqn = (
inner(w, (state.h_project(u) - u_in))*dx
- beta_cp*div(theta_w*V(w))*pibar*dxp
# 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*w)*pi*dxp
+ beta_cp*jump(thetabar_w*w, n)*avg(pi)*dS_vp
+ (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)
bcs = [DirichletBC(M.sub(0), 0.0, "bottom"),
DirichletBC(M.sub(0), 0.0, "top")]
# Solver for u, rho
urho_problem = LinearVariationalProblem(
aeqn, Leqn, self.urho, bcs=bcs)
self.urho_solver = LinearVariationalSolver(urho_problem,
solver_parameters=self.solver_parameters,
options_prefix='ImplicitSolver')
# Reconstruction of theta
theta = TrialFunction(Vtheta)
gamma = TestFunction(Vtheta)
u, rho = self.urho.split()
self.theta = Function(Vtheta)
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, prognostic_variables, simulation_parameters):
mesh = simulation_parameters['mesh'][-1]
self.scheme = simulation_parameters['scheme'][-1]
self.timestepping = simulation_parameters['timestepping'][-1]
alphasq = simulation_parameters['alphasq'][-1]
c0 = simulation_parameters['c0'][-1]
gamma = simulation_parameters['gamma'][-1]
Dt = Constant(simulation_parameters['dt'][-1])
self.solvers = []
if alphasq.values()[0] > 0.0 and gamma.values()[0] == 0.0:
self.setup = 'ch'
if self.scheme == 'upwind' and self.timestepping == 'ssprk3':
Vm = prognostic_variables.Vm
Vu = prognostic_variables.Vu
self.m = prognostic_variables.m
self.u = prognostic_variables.u
self.Xi = prognostic_variables.dXi
self.m0 = Function(Vm).assign(self.m)
# now make problem for the actual problem
psi = TestFunction(Vm)
self.m_trial = Function(Vm)
self.dm = Function(
Vm
) # introduce this as the advection operator for a single step
us = Dt * self.u + self.Xi
nhat = FacetNormal(mesh)
un = 0.5 * (dot(us, nhat) + abs(dot(us, nhat)))
ones = Function(Vu).project(as_vector([Constant(1.)]))
Lm = (psi * self.dm * dx -
psi.dx(0) * self.m_trial * dot(ones, us) * dx +
psi * self.m_trial * dot(ones, us.dx(0)) * dx +
jump(psi) * (un('+') * self.m_trial('+') -
un('-') * self.m_trial('-')) * dS)
mprob = NonlinearVariationalProblem(Lm, self.dm)
self.msolver = NonlinearVariationalSolver(mprob,
solver_parameters={
'ksp_type':
'preonly',
'pc_type':
'bjacobi',
'sub_pc_type':
'ilu'
})
phi = TestFunction(Vu)
Lu = (dot(phi, ones) * self.m * dx - dot(phi, self.u) * dx -
alphasq * dot(self.u.dx(0), phi.dx(0)) * dx)
uprob = NonlinearVariationalProblem(Lu, self.u)
self.usolver = NonlinearVariationalSolver(uprob,
solver_parameters={
'ksp_type':
'preonly',
'pc_type': 'lu'
})
elif self.scheme == 'hydrodynamic' and self.timestepping == 'midpoint':
Vu = prognostic_variables.Vu
self.u = prognostic_variables.u
W = MixedFunctionSpace((Vu, ) * 3)
psi, phi, zeta = TestFunctions(W)
w1 = Function(W)
self.u1, dFh, dGh = split(w1)
uh = (self.u1 + self.u) / 2
dXi = prognostic_variables.dXi
dXi_x = prognostic_variables.dXi_x
dXi_xx = prognostic_variables.dXi_xx
dvh = Dt * uh + dXi
Lu = (psi * (self.u1 - self.u) * dx +
psi * uh.dx(0) * dvh * dx - psi.dx(0) * dFh * dx +
psi * dGh * dx + phi * dFh * dx +
alphasq * phi.dx(0) * dFh.dx(0) * dx -
phi * uh * uh * Dt * dx -
0.5 * alphasq * phi * uh.dx(0) * uh.dx(0) * Dt * dx +
zeta * dGh * dx + alphasq * zeta.dx(0) * dGh.dx(0) * dx -
2 * zeta * uh * dXi_x * dx -
alphasq * zeta * uh.dx(0) * dXi_xx * dx)
self.u1, dFh, dGh = w1.split()
uprob = NonlinearVariationalProblem(Lu, w1)
self.usolver = NonlinearVariationalSolver(uprob,
solver_parameters={
'mat_type':
'aij',
'ksp_type':
'preonly',
'pc_type': 'lu'
})
elif self.scheme == 'no_gradient' and self.timestepping == 'midpoint':
# a version of the hydrodynamic form but without exploiting the gradient
Vu = prognostic_variables.Vu
self.u = prognostic_variables.u
W = MixedFunctionSpace((Vu, ) * 3)
psi, phi, zeta = TestFunctions(W)
w1 = Function(W)
self.u1, dFh, dGh = split(w1)
uh = (self.u1 + self.u) / 2
dXi = prognostic_variables.dXi
dXi_x = prognostic_variables.dXi_x
dXi_xx = prognostic_variables.dXi_xx
dvh = Dt * uh + dXi
Lu = (psi * (self.u1 - self.u) * dx +
psi * uh.dx(0) * dvh * dx + psi * dFh.dx(0) * dx +
psi * dGh * dx + phi * dFh * dx +
alphasq * phi.dx(0) * dFh.dx(0) * dx -
phi * uh * uh * Dt * dx -
0.5 * alphasq * phi * uh.dx(0) * uh.dx(0) * Dt * dx +
zeta * dGh * dx + alphasq * zeta.dx(0) * dGh.dx(0) * dx -
2 * zeta * uh * dXi_x * dx -
alphasq * zeta * uh.dx(0) * dXi_xx * dx)
self.u1, dFh, dGh = w1.split()
uprob = NonlinearVariationalProblem(Lu, w1)
self.usolver = NonlinearVariationalSolver(uprob,
solver_parameters={
'mat_type':
'aij',
'ksp_type':
'preonly',
'pc_type': 'lu'
})
elif self.scheme == 'test' and self.timestepping == 'midpoint':
self.u = prognostic_variables.u
Vu = prognostic_variables.Vu
psi = TestFunction(Vu)
self.u1 = Function(Vu)
uh = (self.u1 + self.u) / 2
dvh = Dt * uh + prognostic_variables.dXi
eqn = (psi * (self.u1 - self.u) * dx -
psi * uh * dvh.dx(0) * dx)
prob = NonlinearVariationalProblem(eqn, self.u1)
self.usolver = NonlinearVariationalSolver(prob,
solver_parameters={
'mat_type':
'aij',
'ksp_type':
'preonly',
'pc_type': 'lu'
})
else:
raise ValueError(
'Scheme %s and timestepping %s either not compatible or not recognised.'
% (self.scheme, self.timestepping))
elif alphasq.values()[0] == 0.0 and gamma.values()[0] > 0.0:
self.setup = 'kdv'
if self.scheme == 'upwind' and self.timestepping == 'ssprk3':
raise NotImplementedError(
'Scheme %s and timestepping %s not yet implemented.' %
(self.scheme, self.timestepping))
elif self.scheme == 'upwind' and self.timestepping == 'midpoint':
raise NotImplementedError(
'Scheme %s and timestepping %s not yet implemented.' %
(self.scheme, self.timestepping))
elif self.scheme == 'hydrodynamic' and self.timestepping == 'midpoint':
raise NotImplementedError(
'Scheme %s and timestepping %s not yet implemented.' %
(self.scheme, self.timestepping))
else:
raise ValueError(
'Scheme %s and timestepping %s either not compatible or not recognised.'
% (self.scheme, self.timestepping))
else:
raise NotImplementedError(
'Schemes for your values of alpha squared %.3f and gamma %.3f are not yet implemented.'
% (alphasq, gamma))
{"ksp_type":"preonly",
"pc_type":"lu"})
q_solver.solve()
# Build advection, forcing forms
Frhs = unk*Dnk/3. + un*Dnk/6. + unk*Dn/6. + un*Dn/3.
K = inner(un, un)/3. + inner(un, unk)/3. + inner(unk, unk)/3.
Prhs = g*(0.5*(Dn + Dp) + b) + 0.5*K
# D advection solver
phi = TestFunction(W0)
D_ = TrialFunction(W0)
D_bar = 0.5*(Dn + D_)
D_ad = (inner(grad(phi), D_bar*u_rec)*dx
- jump(phi*u_rec, n)*uw(ubar, D_bar)*dS)
D_eqn = (D_ - Dn)*phi*dx - dt*D_ad
Dad_p = LinearVariationalProblem(lhs(D_eqn), rhs(D_eqn), Dp)
D_ad_solver = LinearVariationalSolver(Dad_p)
# u advection solver
w = TestFunction(W1)
u_ = TrialFunction(W1)
u_bar = 0.5*(un + u_)
u_ad = (inner(perp(grad(inner(Dbar*w, perp(u_rec)))), u_bar)*dx
+ inner(jump(inner(Dbar*w, perp(u_rec)), n),
p_uw(ubar, u_bar))*dS)
u_eqn = inner(u_ - un, Dbar*w)*dx - dt*u_ad
uad_p = LinearVariationalProblem(lhs(u_eqn), rhs(u_eqn), uad)
u_ad_solver = LinearVariationalSolver(uad_p)
def initialize(self, pc):
""" Set up the problem context. Takes the original
mixed problem and transforms it into the equivalent
hybrid-mixed system.
A KSP object is created for the Lagrange multipliers
on the top/bottom faces of the mesh cells.
"""
from firedrake import (FunctionSpace, Function, Constant,
FiniteElement, TensorProductElement,
TrialFunction, TrialFunctions, TestFunction,
DirichletBC, interval, MixedElement,
BrokenElement)
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.formmanipulation import split_form
from ufl.algorithms.replace import replace
from ufl.cell import TensorProductCell
# Extract PC context
prefix = pc.getOptionsPrefix() + "vert_hybridization_"
_, P = pc.getOperators()
self.ctx = P.getPythonContext()
if not isinstance(self.ctx, ImplicitMatrixContext):
raise ValueError(
"The python context must be an ImplicitMatrixContext")
test, trial = self.ctx.a.arguments()
V = test.function_space()
mesh = V.mesh()
# Magically determine which spaces are vector and scalar valued
for i, Vi in enumerate(V):
# Vector-valued spaces will have a non-empty value_shape
if Vi.ufl_element().value_shape():
self.vidx = i
else:
self.pidx = i
Vv = V[self.vidx]
Vp = V[self.pidx]
# Create the space of approximate traces in the vertical.
# NOTE: Technically a hack since the resulting space is technically
# defined in cell interiors, however the degrees of freedom will only
# be geometrically defined on edges. Arguments will only be used in
# surface integrals
deg, _ = Vv.ufl_element().degree()
# Assumes a tensor product cell (quads, triangular-prisms, cubes)
if not isinstance(Vp.ufl_element().cell(), TensorProductCell):
raise NotImplementedError(
"Currently only implemented for tensor product discretizations"
)
# Only want the horizontal cell
cell, _ = Vp.ufl_element().cell()._cells
DG = FiniteElement("DG", cell, deg)
CG = FiniteElement("CG", interval, 1)
Vv_tr_element = TensorProductElement(DG, CG)
Vv_tr = FunctionSpace(mesh, Vv_tr_element)
# Break the spaces
broken_elements = MixedElement(
[BrokenElement(Vi.ufl_element()) for Vi in V])
V_d = FunctionSpace(mesh, broken_elements)
# Set up relevant functions
self.broken_solution = Function(V_d)
self.broken_residual = Function(V_d)
self.trace_solution = Function(Vv_tr)
self.unbroken_solution = Function(V)
self.unbroken_residual = Function(V)
# Set up transfer kernels to and from the broken velocity space
# NOTE: Since this snippet of code is used in a couple places in
# in Gusto, might be worth creating a utility function that is
# is importable and just called where needed.
shapes = {
"i": Vv.finat_element.space_dimension(),
"j": np.prod(Vv.shape, dtype=int)
}
weight_kernel = """
for (int i=0; i<{i}; ++i)
for (int j=0; j<{j}; ++j)
w[i*{j} + j] += 1.0;
""".format(**shapes)
self.weight = Function(Vv)
par_loop(weight_kernel, dx, {"w": (self.weight, INC)})
# Averaging kernel
self.average_kernel = """
for (int i=0; i<{i}; ++i)
for (int j=0; j<{j}; ++j)
vec_out[i*{j} + j] += vec_in[i*{j} + j]/w[i*{j} + j];
""".format(**shapes)
# Original mixed operator replaced with "broken" arguments
arg_map = {test: TestFunction(V_d), trial: TrialFunction(V_d)}
Atilde = Tensor(replace(self.ctx.a, arg_map))
gammar = TestFunction(Vv_tr)
n = FacetNormal(mesh)
sigma = TrialFunctions(V_d)[self.vidx]
# Again, assumes tensor product structure. Why use this if you
# don't have some form of vertical extrusion?
Kform = gammar('+') * jump(sigma, n=n) * dS_h
# Here we deal with boundary conditions
if self.ctx.row_bcs:
# Find all the subdomains with neumann BCS
# These are Dirichlet BCs on the vidx space
neumann_subdomains = set()
for bc in self.ctx.row_bcs:
if bc.function_space().index == self.pidx:
raise NotImplementedError(
"Dirichlet conditions for scalar variable not supported. Use a weak bc."
)
if bc.function_space().index != self.vidx:
raise NotImplementedError(
"Dirichlet bc set on unsupported space.")
# append the set of sub domains
subdom = bc.sub_domain
if isinstance(subdom, str):
neumann_subdomains |= set([subdom])
else:
neumann_subdomains |= set(as_tuple(subdom, int))
# separate out the top and bottom bcs
extruded_neumann_subdomains = neumann_subdomains & {
"top", "bottom"
}
neumann_subdomains = neumann_subdomains - extruded_neumann_subdomains
integrand = gammar * dot(sigma, n)
measures = []
trace_subdomains = []
for subdomain in sorted(extruded_neumann_subdomains):
measures.append({"top": ds_t, "bottom": ds_b}[subdomain])
trace_subdomains.extend(
sorted({"top", "bottom"} - extruded_neumann_subdomains))
measures.extend((ds(sd) for sd in sorted(neumann_subdomains)))
markers = [int(x) for x in mesh.exterior_facets.unique_markers]
dirichlet_subdomains = set(markers) - neumann_subdomains
trace_subdomains.extend(sorted(dirichlet_subdomains))
for measure in measures:
Kform += integrand * measure
else:
trace_subdomains = ["top", "bottom"]
trace_bcs = [
DirichletBC(Vv_tr, Constant(0.0), subdomain)
for subdomain in trace_subdomains
]
# Make a SLATE tensor from Kform
K = Tensor(Kform)
# Assemble the Schur complement operator and right-hand side
self.schur_rhs = Function(Vv_tr)
self._assemble_Srhs = create_assembly_callable(
K * Atilde.inv * AssembledVector(self.broken_residual),
tensor=self.schur_rhs,
form_compiler_parameters=self.ctx.fc_params)
mat_type = PETSc.Options().getString(prefix + "mat_type", "aij")
schur_comp = K * Atilde.inv * K.T
self.S = allocate_matrix(schur_comp,
bcs=trace_bcs,
form_compiler_parameters=self.ctx.fc_params,
mat_type=mat_type,
options_prefix=prefix)
self._assemble_S = create_assembly_callable(
schur_comp,
tensor=self.S,
bcs=trace_bcs,
form_compiler_parameters=self.ctx.fc_params,
mat_type=mat_type)
self._assemble_S()
self.S.force_evaluation()
Smat = self.S.petscmat
nullspace = self.ctx.appctx.get("vert_trace_nullspace", None)
if nullspace is not None:
nsp = nullspace(Vv_tr)
Smat.setNullSpace(nsp.nullspace(comm=pc.comm))
# Set up the KSP for the system of Lagrange multipliers
trace_ksp = PETSc.KSP().create(comm=pc.comm)
trace_ksp.setOptionsPrefix(prefix)
trace_ksp.setOperators(Smat)
trace_ksp.setUp()
trace_ksp.setFromOptions()
self.trace_ksp = trace_ksp
split_mixed_op = dict(split_form(Atilde.form))
split_trace_op = dict(split_form(K.form))
# Generate reconstruction calls
self._reconstruction_calls(split_mixed_op, split_trace_op)
def _setup_solver(self):
import numpy as np
state = self.state
dt = state.dt
beta_ = dt * self.alpha
cp = state.parameters.cp
Vu = state.spaces("HDiv")
Vu_broken = FunctionSpace(state.mesh, BrokenElement(Vu.ufl_element()))
Vtheta = state.spaces("theta")
Vrho = state.spaces("DG")
# Store time-stepping coefficients as UFL Constants
beta = Constant(beta_)
beta_cp = Constant(beta_ * cp)
h_deg = Vrho.ufl_element().degree()[0]
v_deg = Vrho.ufl_element().degree()[1]
Vtrace = FunctionSpace(state.mesh, "HDiv Trace", degree=(h_deg, v_deg))
# Split up the rhs vector (symbolically)
self.xrhs = Function(self.equations.function_space)
u_in, rho_in, theta_in = split(self.xrhs)[0:3]
# Build the function space for "broken" u, rho, and pressure trace
M = MixedFunctionSpace((Vu_broken, Vrho, Vtrace))
w, phi, dl = TestFunctions(M)
u, rho, l0 = TrialFunctions(M)
n = FacetNormal(state.mesh)
# Get background fields
thetabar = state.fields("thetabar")
rhobar = state.fields("rhobar")
exnerbar = thermodynamics.exner_pressure(state.parameters, rhobar,
thetabar)
exnerbar_rho = thermodynamics.dexner_drho(state.parameters, rhobar,
thetabar)
exnerbar_theta = thermodynamics.dexner_dtheta(state.parameters, rhobar,
thetabar)
# 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 exner') in the vertical
# component of the gradient
# The exner prime term (here, bars are for mean and no bars are
# for linear perturbations)
exner = exnerbar_theta * theta + exnerbar_rho * rho
# Vertical projection
def V(u):
return k * inner(u, k)
# hydrostatic projection
h_project = lambda u: u - k * inner(u, k)
# Specify degree for some terms as estimated degree is too large
dxp = dx(degree=(self.quadrature_degree))
dS_vp = dS_v(degree=(self.quadrature_degree))
dS_hp = dS_h(degree=(self.quadrature_degree))
ds_vp = ds_v(degree=(self.quadrature_degree))
ds_tbp = (ds_t(degree=(self.quadrature_degree)) +
ds_b(degree=(self.quadrature_degree)))
# 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_w = theta / (1 + water_t)
thetabar_w = thetabar / (1 + water_t)
else:
theta_w = theta
thetabar_w = thetabar
_l0 = TrialFunction(Vtrace)
_dl = TestFunction(Vtrace)
a_tr = _dl('+') * _l0('+') * (
dS_vp + dS_hp) + _dl * _l0 * ds_vp + _dl * _l0 * ds_tbp
def L_tr(f):
return _dl('+') * avg(f) * (
dS_vp + dS_hp) + _dl * f * ds_vp + _dl * f * ds_tbp
cg_ilu_parameters = {
'ksp_type': 'cg',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
}
# Project field averages into functions on the trace space
rhobar_avg = Function(Vtrace)
exnerbar_avg = Function(Vtrace)
rho_avg_prb = LinearVariationalProblem(a_tr, L_tr(rhobar), rhobar_avg)
exner_avg_prb = LinearVariationalProblem(a_tr, L_tr(exnerbar),
exnerbar_avg)
rho_avg_solver = LinearVariationalSolver(
rho_avg_prb,
solver_parameters=cg_ilu_parameters,
options_prefix='rhobar_avg_solver')
exner_avg_solver = LinearVariationalSolver(
exner_avg_prb,
solver_parameters=cg_ilu_parameters,
options_prefix='exnerbar_avg_solver')
with timed_region("Gusto:HybridProjectRhobar"):
rho_avg_solver.solve()
with timed_region("Gusto:HybridProjectExnerbar"):
exner_avg_solver.solve()
# "broken" u, rho, and trace system
# NOTE: no ds_v integrals since equations are defined on
# a periodic (or sphere) base mesh.
if any([t.has_label(hydrostatic) for t in self.equations.residual]):
u_mass = inner(w, (h_project(u) - u_in)) * dx
else:
u_mass = inner(w, (u - u_in)) * dx
eqn = (
# momentum equation
u_mass - beta_cp * div(theta_w * V(w)) * exnerbar * dxp
# following does nothing but is preserved in the comments
# to remind us why (because V(w) is purely vertical).
# + beta_cp*jump(theta_w*V(w), n=n)*exnerbar_avg('+')*dS_vp
+ beta_cp * jump(theta_w * V(w), n=n) * exnerbar_avg('+') * dS_hp +
beta_cp * dot(theta_w * V(w), n) * exnerbar_avg * ds_tbp -
beta_cp * div(thetabar_w * w) * exner * dxp
# trace terms appearing after integrating momentum equation
+ beta_cp * jump(thetabar_w * w, n=n) * l0('+') * (dS_vp + dS_hp) +
beta_cp * dot(thetabar_w * w, n) * l0 * (ds_tbp + ds_vp)
# mass continuity equation
+ (phi *
(rho - rho_in) - beta * inner(grad(phi), u) * rhobar) * dx +
beta * jump(phi * u, n=n) * rhobar_avg('+') * (dS_v + dS_h)
# term added because u.n=0 is enforced weakly via the traces
+ beta * phi * dot(u, n) * rhobar_avg * (ds_tb + ds_v)
# constraint equation to enforce continuity of the velocity
# through the interior facets and weakly impose the no-slip
# condition
+ dl('+') * jump(u, n=n) * (dS_vp + dS_hp) + dl * dot(u, n) *
(ds_tbp + ds_vp))
# contribution of the sponge term
if hasattr(self.equations, "mu"):
eqn += dt * self.equations.mu * inner(w, k) * inner(u, k) * dx
aeqn = lhs(eqn)
Leqn = rhs(eqn)
# Function for the hybridized solutions
self.urhol0 = Function(M)
hybridized_prb = LinearVariationalProblem(aeqn, Leqn, self.urhol0)
hybridized_solver = LinearVariationalSolver(
hybridized_prb,
solver_parameters=self.solver_parameters,
options_prefix='ImplicitSolver')
self.hybridized_solver = hybridized_solver
# Project broken u into the HDiv space using facet averaging.
# Weight function counting the dofs of the HDiv element:
shapes = {
"i": Vu.finat_element.space_dimension(),
"j": np.prod(Vu.shape, dtype=int)
}
weight_kernel = """
for (int i=0; i<{i}; ++i)
for (int j=0; j<{j}; ++j)
w[i*{j} + j] += 1.0;
""".format(**shapes)
self._weight = Function(Vu)
par_loop(weight_kernel, dx, {"w": (self._weight, INC)})
# Averaging kernel
self._average_kernel = """
for (int i=0; i<{i}; ++i)
for (int j=0; j<{j}; ++j)
vec_out[i*{j} + j] += vec_in[i*{j} + j]/w[i*{j} + j];
""".format(**shapes)
# HDiv-conforming velocity
self.u_hdiv = Function(Vu)
# Reconstruction of theta
theta = TrialFunction(Vtheta)
gamma = TestFunction(Vtheta)
self.theta = Function(Vtheta)
theta_eqn = gamma * (theta - theta_in + dot(k, self.u_hdiv) *
dot(k, grad(thetabar)) * beta) * dx
theta_problem = LinearVariationalProblem(lhs(theta_eqn),
rhs(theta_eqn), self.theta)
self.theta_solver = LinearVariationalSolver(
theta_problem,
solver_parameters=cg_ilu_parameters,
options_prefix='thetabacksubstitution')
# Store boundary conditions for the div-conforming velocity to apply
# post-solve
self.bcs = self.equations.bcs['u']
def stabilisation_term(self, alpha, zeta, mesh, v, test):
e = avg(CellVolume(mesh)) / FacetArea(mesh)
return 2 * alpha * zeta / e * (dot(jump(v), jump(test))) * dS
def setup(self, state):
space = state.spaces("HDiv")
super(SawyerEliassenU, self).setup(state, space=space)
u = state.fields("u")
b = state.fields("b")
v = inner(u, as_vector([0., 1., 0.]))
# spaces
V0 = FunctionSpace(state.mesh, "CG", 2)
Vu = u.function_space()
# project b to V0
self.b_v0 = Function(V0)
btri = TrialFunction(V0)
btes = TestFunction(V0)
a = inner(btes, btri) * dx
L = inner(btes, b) * dx
projectbproblem = LinearVariationalProblem(a, L, self.b_v0)
self.project_b_solver = LinearVariationalSolver(
projectbproblem, solver_parameters={'ksp_type': 'cg'})
# project v to V0
self.v_v0 = Function(V0)
vtri = TrialFunction(V0)
vtes = TestFunction(V0)
a = inner(vtes, vtri) * dx
L = inner(vtes, v) * dx
projectvproblem = LinearVariationalProblem(a, L, self.v_v0)
self.project_v_solver = LinearVariationalSolver(
projectvproblem, solver_parameters={'ksp_type': 'cg'})
# stm/psi is a stream function
self.stm = Function(V0)
psi = TrialFunction(V0)
xsi = TestFunction(V0)
f = state.parameters.f
H = state.parameters.H
L = state.parameters.L
dbdy = state.parameters.dbdy
x, y, z = SpatialCoordinate(state.mesh)
bcs = [DirichletBC(V0, 0., "bottom"), DirichletBC(V0, 0., "top")]
Mat = as_matrix([[b.dx(2), 0., -f * self.v_v0.dx(2)], [0., 0., 0.],
[-self.b_v0.dx(0), 0., f**2 + f * self.v_v0.dx(0)]])
Equ = (inner(grad(xsi), dot(Mat, grad(psi))) -
dbdy * inner(grad(xsi), as_vector([-v, 0., f *
(z - H / 2)]))) * dx
# fourth-order terms
if state.parameters.fourthorder:
eps = Constant(0.0001)
brennersigma = Constant(10.0)
n = FacetNormal(state.mesh)
deltax = Constant(state.parameters.deltax)
deltaz = Constant(state.parameters.deltaz)
nn = as_matrix([[sqrt(brennersigma / Constant(deltax)), 0., 0.],
[0., 0., 0.],
[0., 0.,
sqrt(brennersigma / Constant(deltaz))]])
mu = as_matrix([[1., 0., 0.], [0., 0., 0.], [0., 0., H / L]])
# anisotropic form
Equ += eps * (
div(dot(mu, grad(psi))) * div(dot(mu, grad(xsi))) * dx -
(avg(dot(dot(grad(grad(psi)), n), n)) * jump(grad(xsi), n=n) +
avg(dot(dot(grad(grad(xsi)), n), n)) * jump(grad(psi), n=n) -
jump(nn * grad(psi), n=n) * jump(nn * grad(xsi), n=n)) *
(dS_h + dS_v))
Au = lhs(Equ)
Lu = rhs(Equ)
stmproblem = LinearVariationalProblem(Au, Lu, self.stm, bcs=bcs)
self.stream_function_solver = LinearVariationalSolver(
stmproblem, solver_parameters={'ksp_type': 'cg'})
# solve for sawyer_eliassen u
self.u = Function(Vu)
utrial = TrialFunction(Vu)
w = TestFunction(Vu)
a = inner(w, utrial) * dx
L = (w[0] * (-self.stm.dx(2)) + w[2] * (self.stm.dx(0))) * dx
ugproblem = LinearVariationalProblem(a, L, self.u)
self.sawyer_eliassen_u_solver = LinearVariationalSolver(
ugproblem, solver_parameters={'ksp_type': 'cg'})
def __init__(self, diagnostic_variables, prognostic_variables, outputting,
simulation_parameters):
self.diagnostic_variables = diagnostic_variables
self.prognostic_variables = prognostic_variables
self.outputting = outputting
self.simulation_parameters = simulation_parameters
Dt = Constant(simulation_parameters['dt'][-1])
Ld = simulation_parameters['Ld'][-1]
u = self.prognostic_variables.u
Xi = self.prognostic_variables.dXi
Vu = u.function_space()
vector_u = True if Vu.ufl_element() == VectorElement else False
ones = Function(
VectorFunctionSpace(self.prognostic_variables.mesh, "CG",
1)).project(as_vector([Constant(1.0)]))
self.to_update_constants = False
self.interpolators = []
self.projectors = []
self.solvers = []
mesh = u.function_space().mesh()
x, = SpatialCoordinate(mesh)
alphasq = simulation_parameters['alphasq'][-1]
periodic = simulation_parameters['periodic'][-1]
# do peakon data checks here
true_peakon_data = simulation_parameters['true_peakon_data'][-1]
if true_peakon_data is not None:
self.true_peakon_file = Dataset(
'results/' + true_peakon_data + '/data.nc', 'r')
# check length of file is correct
ndump = simulation_parameters['ndump'][-1]
tmax = simulation_parameters['tmax'][-1]
dt = simulation_parameters['dt'][-1]
if len(self.true_peakon_file['time'][:]) != int(tmax /
(ndump * dt)) + 1:
raise ValueError(
'If reading in true peakon data, the dump frequency must be the same as that used for the true peakon data.'
+
' Length of true peakon data as %i, but proposed length is %i'
% (len(self.true_peakon_file['time'][:]),
int(tmax / (ndump * dt)) + 1))
if self.true_peakon_file['p'][:].shape != (int(tmax /
(ndump * dt)) +
1, ):
raise ValueError(
'True peakon data shape %i must be the same shape as proposed data %i'
% ((int(tmax / (ndump * dt)) + 1, ),
self.true_peakon_file['p'][:].shape))
# do peakon data checks here
true_mean_peakon_data = simulation_parameters['true_mean_peakon_data'][
-1]
if true_mean_peakon_data is not None:
self.true_mean_peakon_file = Dataset(
'results/' + true_mean_peakon_data + '/data.nc', 'r')
# check length of file is correct
ndump = simulation_parameters['ndump'][-1]
tmax = simulation_parameters['tmax'][-1]
dt = simulation_parameters['dt'][-1]
if len(self.true_mean_peakon_file['time'][:]) != int(tmax /
(ndump * dt)):
raise ValueError(
'If reading in true peakon data, the dump frequency must be the same as that used for the true peakon data.'
)
if self.true_mean_peakon_file['p'][:].shape != (int(
tmax / (ndump * dt)), ):
raise ValueError(
'True peakon data must have same shape as proposed data!')
for key, value in self.diagnostic_variables.fields.items():
if key == 'uscalar':
uscalar = self.diagnostic_variables.fields['uscalar']
u_interpolator = Interpolator(dot(ones, u), uscalar)
self.interpolators.append(u_interpolator)
elif key == 'Euscalar':
Eu = self.prognostic_variables.Eu
Euscalar = self.diagnostic_variables.fields['Euscalar']
Eu_interpolator = Interpolator(dot(ones, Eu), Euscalar)
self.interpolators.append(Eu_interpolator)
elif key == 'Xiscalar':
Xi = self.prognostic_variables.dXi
Xiscalar = self.diagnostic_variables.fields['Xiscalar']
Xi_interpolator = Interpolator(dot(ones, Xi), Xiscalar)
self.interpolators.append(Xi_interpolator)
elif key == 'du':
if type(u.function_space().ufl_element()) == VectorElement:
u_to_project = self.diagnostic_variables.fields['uscalar']
else:
u_to_project = u
du = self.diagnostic_variables.fields['du']
du_projector = Projector(u_to_project.dx(0), du)
self.projectors.append(du_projector)
elif key == 'jump_du':
du = self.diagnostic_variables.fields['du']
jump_du = self.diagnostic_variables.fields['jump_du']
V = jump_du.function_space()
jtrial = TrialFunction(V)
psi = TestFunction(V)
Lj = psi('+') * abs(jump(du)) * dS
aj = psi('+') * jtrial('+') * dS
jprob = LinearVariationalProblem(aj, Lj, jump_du)
jsolver = LinearVariationalSolver(jprob)
self.solvers.append(jsolver)
elif key == 'du_smooth':
du = self.diagnostic_variables.fields['du']
du_smooth = self.diagnostic_variables.fields['du_smooth']
projector = Projector(du, du_smooth)
self.projectors.append(projector)
elif key == 'u2_flux':
gamma = simulation_parameters['gamma'][-1]
u2_flux = self.diagnostic_variables.fields['u2_flux']
xis = self.prognostic_variables.pure_xi_list
xis_x = []
xis_xxx = []
CG1 = FunctionSpace(mesh, "CG", 1)
psi = TestFunction(CG1)
for xi in xis:
xis_x.append(Function(CG1).project(xi.dx(0)))
for xi_x in xis_x:
xi_xxx = Function(CG1)
form = (psi * xi_xxx + psi.dx(0) * xi_x.dx(0)) * dx
prob = NonlinearVariationalProblem(form, xi_xxx)
solver = NonlinearVariationalSolver(prob)
solver.solve()
xis_xxx.append(xi_xxx)
flux_expr = 0.0 * x
for xi, xi_x, xi_xxx in zip(xis, xis_x, xis_xxx):
flux_expr += (6 * u.dx(0) * xi + 12 * u * xi_x + gamma *
xi_xxx) * (6 * u.dx(0) * xi + 24 * u * xi_x +
gamma * xi_xxx)
projector = Projector(flux_expr, u2_flux)
self.projectors.append(projector)
elif key == 'a':
# find 6 * u_x * Xi + gamma * Xi_xxx
mesh = u.function_space().mesh()
gamma = simulation_parameters['gamma'][-1]
a_flux = self.diagnostic_variables.fields['a']
xis = self.prognostic_variables.pure_xis
xis_x = []
xis_xxx = []
CG1 = FunctionSpace(mesh, "CG", 1)
psi = TestFunction(CG1)
for xi in xis:
xis_x.append(Function(CG1).project(xi.dx(0)))
for xi_x in xis_x:
xi_xxx = Function(CG1)
form = (psi * xi_xxx + psi.dx(0) * xi_x.dx(0)) * dx
prob = NonlinearVariationalProblem(form, xi_xxx)
solver = NonlinearVariationalSolver(prob)
solver.solve()
xis_xxx.append(xi_xxx)
x, = SpatialCoordinate(mesh)
a_expr = 0.0 * x
for xi, xi_x, xi_xxx in zip(xis, xis_x, xis_xxx):
a_expr += 6 * u.dx(0) * xi + gamma * xi_xxx
projector = Projector(a_expr, a_flux)
self.projectors.append(projector)
elif key == 'b':
# find 12 * u * Xi_x
mesh = u.function_space().mesh()
gamma = simulation_parameters['gamma'][-1]
b_flux = self.diagnostic_variables.fields['b']
xis = self.prognostic_variables.pure_xis
x, = SpatialCoordinate(mesh)
b_expr = 0.0 * x
for xi, xi_x, xi_xxx in zip(xis, xis_x, xis_xxx):
b_expr += 12 * u * xi.dx(0)
projector = Projector(b_expr, b_flux)
self.projectors.append(projector)
elif key == 'kdv_1':
# find the first part of the kdv form
u0 = prognostic_variables.u0
uh = (u + u0) / 2
us = Dt * uh + sqrt(Dt) * Xi
psi = TestFunction(Vu)
du_1 = self.diagnostic_variables.fields['kdv_1']
eqn = psi * du_1 * dx - 6 * psi.dx(0) * uh * us * dx
prob = NonlinearVariationalProblem(eqn, du_1)
solver = NonlinearVariationalSolver(prob)
self.solvers.append(solver)
elif key == 'kdv_2':
# find the second part of the kdv form
u0 = prognostic_variables.u0
uh = (u + u0) / 2
us = Dt * uh + sqrt(Dt) * Xi
psi = TestFunction(Vu)
du_2 = self.diagnostic_variables.fields['kdv_2']
eqn = psi * du_2 * dx + 6 * psi * uh * us.dx(0) * dx
prob = NonlinearVariationalProblem(eqn, du_2)
solver = NonlinearVariationalSolver(prob)
self.solvers.append(solver)
elif key == 'kdv_3':
# find the third part of the kdv form
u0 = prognostic_variables.u0
uh = (u + u0) / 2
us = Dt * uh + sqrt(Dt) * Xi
du_3 = self.diagnostic_variables.fields['kdv_3']
gamma = simulation_parameters['gamma'][-1]
phi = TestFunction(Vu)
F = Function(Vu)
eqn = (phi * F * dx + phi.dx(0) * us.dx(0) * dx)
prob = NonlinearVariationalProblem(eqn, F)
solver = NonlinearVariationalSolver(prob)
self.solvers.append(solver)
self.projectors.append(Projector(-gamma * F.dx(0), du_3))
# nu = TestFunction(Vu)
# back_eqn = nu * du_3 * dx - gamma * nu.dx(0) * F * dx
# back_prob = NonlinearVariationalProblem(back_eqn, du_3)
# back_solver = NonlinearVariationalSolver(back_prob)
# self.solvers.append(solver)
elif key == 'm':
m = self.diagnostic_variables.fields['m']
phi = TestFunction(Vu)
eqn = phi * m * dx - phi * u * dx - alphasq * phi.dx(0) * u.dx(
0) * dx
prob = NonlinearVariationalProblem(eqn, m)
solver = NonlinearVariationalSolver(prob)
self.solvers.append(solver)
elif key == 'u_xx':
u_xx = self.diagnostic_variables.fields['u_xx']
phi = TestFunction(Vu)
eqn = phi * u_xx * dx + phi.dx(0) * u_xx.dx(0) * dx
prob = NonlinearVariationalProblem(eqn, u_xx)
solver = NonlinearVariationalSolver(prob)
self.solvers.append(solver)
elif key == 'u_sde':
self.to_update_constants = True
self.Ld = Ld
self.alphasq = alphasq
self.p = Constant(1.0 * 0.5 * (1 + exp(-Ld / sqrt(alphasq))) /
(1 - exp(-Ld / sqrt(alphasq))))
self.q = Constant(Ld / 2)
u_sde = self.diagnostic_variables.fields['u_sde']
if periodic:
expr = conditional(
x < self.q - Ld / 2,
self.p * ((exp(-(x - self.q + Ld) / sqrt(alphasq)) +
exp(-Ld / sqrt(alphasq)) * exp(
(x - self.q + Ld) / sqrt(alphasq))) /
(1 - exp(-Ld / sqrt(alphasq)))),
conditional(
x < self.q + Ld / 2,
self.p * ((exp(-sqrt((self.q - x)**2 / alphasq)) +
exp(-Ld / sqrt(alphasq)) *
exp(sqrt((self.q - x)**2 / alphasq))) /
(1 - exp(-Ld / sqrt(alphasq)))),
self.p *
((exp(-(self.q + Ld - x) / sqrt(alphasq)) +
exp(-Ld / sqrt(alphasq) * exp(
(self.q + Ld - x) / sqrt(alphasq)))) /
(1 - exp(-Ld / sqrt(alphasq))))))
else:
expr = conditional(
x < self.q - Ld / 2,
self.p * exp(-(x - self.q + Ld) / sqrt(alphasq)),
conditional(
x < self.q + Ld / 2,
self.p * exp(-sqrt((self.q - x)**2 / alphasq)),
self.p * exp(-(self.q + Ld - x) / sqrt(alphasq))))
self.interpolators.append(Interpolator(expr, u_sde))
elif key == 'u_sde_weak':
u_sde = self.diagnostic_variables.fields['u_sde']
u_sde_weak = self.diagnostic_variables.fields['u_sde_weak']
psi = TestFunction(Vu)
eqn = psi * u_sde_weak * dx - psi * (u - u_sde) * dx
prob = NonlinearVariationalProblem(eqn, u_sde_weak)
solver = NonlinearVariationalSolver(prob)
self.solvers.append(solver)
elif key == 'u_sde_mean':
self.to_update_constants = True
self.p = Constant(1.0)
self.q = Constant(Ld / 2)
if periodic:
raise NotImplementedError(
'u_sde_mean not yet implemented for periodic peakon')
u_sde = self.diagnostic_variables.fields['u_sde_mean']
expr = conditional(
x < self.q - Ld / 2,
self.p * exp(-(x - self.q + Ld) / sqrt(alphasq)),
conditional(
x < self.q + Ld / 2,
self.p * exp(-sqrt((self.q - x)**2 / alphasq)),
self.p * exp(-(self.q + Ld - x) / sqrt(alphasq))))
self.interpolators.append(Interpolator(expr, u_sde))
elif key == 'u_sde_weak_mean':
u_sde = self.diagnostic_variables.fields['u_sde_mean']
u_sde_weak = self.diagnostic_variables.fields[
'u_sde_weak_mean']
psi = TestFunction(Vu)
eqn = psi * u_sde_weak * dx - psi * (u - u_sde) * dx
prob = NonlinearVariationalProblem(eqn, u_sde_weak)
solver = NonlinearVariationalSolver(prob)
self.solvers.append(solver)
elif key == 'pure_xi':
pure_xi = 0.0 * x
for xi in self.prognostic_variables.pure_xi_list:
if vector_u:
pure_xi += dot(ones, xi)
else:
pure_xi += xi
Xiscalar = self.diagnostic_variables.fields['pure_xi']
Xi_interpolator = Interpolator(pure_xi, Xiscalar)
self.interpolators.append(Xi_interpolator)
elif key == 'pure_xi_x':
pure_xi_x = 0.0 * x
for xix in self.prognostic_variables.pure_xi_x_list:
if vector_u:
pure_xi_x += dot(ones, xix)
else:
pure_xi_x += xix
Xiscalar = self.diagnostic_variables.fields['pure_xi_x']
Xi_interpolator = Interpolator(pure_xi_x, Xiscalar)
self.interpolators.append(Xi_interpolator)
elif key == 'pure_xi_xx':
pure_xi_xx = 0.0 * x
for xixx in self.prognostic_variables.pure_xi_xx_list:
if vector_u:
pure_xi_xx += dot(ones, xixx)
else:
pure_xi_xx += xixx
Xiscalar = self.diagnostic_variables.fields['pure_xi_xx']
Xi_interpolator = Interpolator(pure_xi_xx, Xiscalar)
self.interpolators.append(Xi_interpolator)
elif key == 'pure_xi_xxx':
pure_xi_xxx = 0.0 * x
for xixxx in self.prognostic_variables.pure_xi_xxx_list:
if vector_u:
pure_xi_xxx += dot(ones, xixxx)
else:
pure_xi_xxx += xixxx
Xiscalar = self.diagnostic_variables.fields['pure_xi_xxx']
Xi_interpolator = Interpolator(pure_xi_xxx, Xiscalar)
self.interpolators.append(Xi_interpolator)
elif key == 'pure_xi_xxxx':
pure_xi_xxxx = 0.0 * x
for xixxxx in self.prognostic_variables.pure_xi_xx_list:
if vector_u:
pure_xi_xxxx += dot(ones, xixxxx)
else:
pure_xi_xxxx += xixxxx
Xiscalar = self.diagnostic_variables.fields['pure_xi_xxxx']
Xi_interpolator = Interpolator(pure_xi_xxxx, Xiscalar)
self.interpolators.append(Xi_interpolator)
else:
raise NotImplementedError('Diagnostic %s not yet implemented' %
key)
def heat_exchanger_optimization(mu=0.03, n_iters=1000):
output_dir = "2D/"
path = os.path.abspath(__file__)
dir_path = os.path.dirname(path)
mesh = fd.Mesh(f"{dir_path}/2D_mesh.msh")
# Perturb the mesh coordinates. Necessary to calculate shape derivatives
S = fd.VectorFunctionSpace(mesh, "CG", 1)
s = fd.Function(S, name="deform")
mesh.coordinates.assign(mesh.coordinates + s)
# Initial level set function
x, y = fd.SpatialCoordinate(mesh)
PHI = fd.FunctionSpace(mesh, "CG", 1)
phi_expr = sin(y * pi / 0.2) * cos(x * pi / 0.2) - fd.Constant(0.8)
# Avoid recording the operation interpolate into the tape.
# Otherwise, the shape derivatives will not be correct
with fda.stop_annotating():
phi = fd.interpolate(phi_expr, PHI)
phi.rename("LevelSet")
fd.File(output_dir + "phi_initial.pvd").write(phi)
# Physics
mu = fd.Constant(mu) # viscosity
alphamin = 1e-12
alphamax = 2.5 / (2e-4)
parameters = {
"mat_type": "aij",
"ksp_type": "preonly",
"ksp_converged_reason": None,
"pc_type": "lu",
"pc_factor_mat_solver_type": "mumps",
}
stokes_parameters = parameters
temperature_parameters = parameters
u_inflow = 2e-3
tin1 = fd.Constant(10.0)
tin2 = fd.Constant(100.0)
P2 = fd.VectorElement("CG", mesh.ufl_cell(), 2)
P1 = fd.FiniteElement("CG", mesh.ufl_cell(), 1)
TH = P2 * P1
W = fd.FunctionSpace(mesh, TH)
U = fd.TrialFunction(W)
u, p = fd.split(U)
V = fd.TestFunction(W)
v, q = fd.split(V)
epsilon = fd.Constant(10000.0)
def hs(phi, epsilon):
return fd.Constant(alphamax) * fd.Constant(1.0) / (
fd.Constant(1.0) + exp(-epsilon * phi)) + fd.Constant(alphamin)
def stokes(phi, BLOCK_INLET_MOUTH, BLOCK_OUTLET_MOUTH):
a_fluid = mu * inner(grad(u), grad(v)) - div(v) * p - q * div(u)
darcy_term = inner(u, v)
return (a_fluid * dx + hs(phi, epsilon) * darcy_term * dx(0) +
alphamax * darcy_term *
(dx(BLOCK_INLET_MOUTH) + dx(BLOCK_OUTLET_MOUTH)))
# Dirichlet boundary conditions
inflow1 = fd.as_vector([
u_inflow * sin(
((y - (line_sep -
(dist_center + inlet_width))) * pi) / inlet_width),
0.0,
])
inflow2 = fd.as_vector([
u_inflow * sin(((y - (line_sep + dist_center)) * pi) / inlet_width),
0.0,
])
noslip = fd.Constant((0.0, 0.0))
# Stokes 1
bcs1_1 = fd.DirichletBC(W.sub(0), noslip, WALLS)
bcs1_2 = fd.DirichletBC(W.sub(0), inflow1, INLET1)
bcs1_3 = fd.DirichletBC(W.sub(1), fd.Constant(0.0), OUTLET1)
bcs1_4 = fd.DirichletBC(W.sub(0), noslip, INLET2)
bcs1_5 = fd.DirichletBC(W.sub(0), noslip, OUTLET2)
bcs1 = [bcs1_1, bcs1_2, bcs1_3, bcs1_4, bcs1_5]
# Stokes 2
bcs2_1 = fd.DirichletBC(W.sub(0), noslip, WALLS)
bcs2_2 = fd.DirichletBC(W.sub(0), inflow2, INLET2)
bcs2_3 = fd.DirichletBC(W.sub(1), fd.Constant(0.0), OUTLET2)
bcs2_4 = fd.DirichletBC(W.sub(0), noslip, INLET1)
bcs2_5 = fd.DirichletBC(W.sub(0), noslip, OUTLET1)
bcs2 = [bcs2_1, bcs2_2, bcs2_3, bcs2_4, bcs2_5]
# Forward problems
U1, U2 = fd.Function(W), fd.Function(W)
L = inner(fd.Constant((0.0, 0.0, 0.0)), V) * dx
problem = fd.LinearVariationalProblem(stokes(-phi, INMOUTH2, OUTMOUTH2),
L,
U1,
bcs=bcs1)
solver_stokes1 = fd.LinearVariationalSolver(
problem,
solver_parameters=stokes_parameters,
options_prefix="stokes_1")
solver_stokes1.solve()
problem = fd.LinearVariationalProblem(stokes(phi, INMOUTH1, OUTMOUTH1),
L,
U2,
bcs=bcs2)
solver_stokes2 = fd.LinearVariationalSolver(
problem,
solver_parameters=stokes_parameters,
options_prefix="stokes_2")
solver_stokes2.solve()
# Convection difussion equation
ks = fd.Constant(1e0)
cp_value = 5.0e5
cp = fd.Constant(cp_value)
T = fd.FunctionSpace(mesh, "DG", 1)
t = fd.Function(T, name="Temperature")
w = fd.TestFunction(T)
# Mesh-related functions
n = fd.FacetNormal(mesh)
h = fd.CellDiameter(mesh)
u1, p1 = fd.split(U1)
u2, p2 = fd.split(U2)
def upwind(u):
return (dot(u, n) + abs(dot(u, n))) / 2.0
u1n = upwind(u1)
u2n = upwind(u2)
# Penalty term
alpha = fd.Constant(500.0)
# Bilinear form
a_int = dot(grad(w), ks * grad(t) - cp * (u1 + u2) * t) * dx
a_fac = (fd.Constant(-1.0) * ks * dot(avg(grad(w)), jump(t, n)) * dS +
fd.Constant(-1.0) * ks * dot(jump(w, n), avg(grad(t))) * dS +
ks("+") *
(alpha("+") / avg(h)) * dot(jump(w, n), jump(t, n)) * dS)
a_vel = (dot(
jump(w),
cp * (u1n("+") + u2n("+")) * t("+") - cp *
(u1n("-") + u2n("-")) * t("-"),
) * dS + dot(w,
cp * (u1n + u2n) * t) * ds)
a_bnd = (dot(w,
cp * dot(u1 + u2, n) * t) * (ds(INLET1) + ds(INLET2)) +
w * t * (ds(INLET1) + ds(INLET2)) - w * tin1 * ds(INLET1) -
w * tin2 * ds(INLET2) + alpha / h * ks * w * t *
(ds(INLET1) + ds(INLET2)) - ks * dot(grad(w), t * n) *
(ds(INLET1) + ds(INLET2)) - ks * dot(grad(t), w * n) *
(ds(INLET1) + ds(INLET2)))
aT = a_int + a_fac + a_vel + a_bnd
LT_bnd = (alpha / h * ks * tin1 * w * ds(INLET1) +
alpha / h * ks * tin2 * w * ds(INLET2) -
tin1 * ks * dot(grad(w), n) * ds(INLET1) -
tin2 * ks * dot(grad(w), n) * ds(INLET2))
problem = fd.LinearVariationalProblem(derivative(aT, t), LT_bnd, t)
solver_temp = fd.LinearVariationalSolver(
problem,
solver_parameters=temperature_parameters,
options_prefix="temperature",
)
solver_temp.solve()
# fd.solve(eT == 0, t, solver_parameters=temperature_parameters)
# Cost function: Flux at the cold outlet
scale_factor = 4e-4
Jform = fd.assemble(
fd.Constant(-scale_factor * cp_value) * inner(t * u1, n) * ds(OUTLET1))
# Constraints: Pressure drop on each fluid
power_drop = 1e-2
Power1 = fd.assemble(p1 / power_drop * ds(INLET1))
Power2 = fd.assemble(p2 / power_drop * ds(INLET2))
phi_pvd = fd.File("phi_evolution.pvd")
def deriv_cb(phi):
with stop_annotating():
phi_pvd.write(phi[0])
c = fda.Control(s)
# Reduced Functionals
Jhat = LevelSetFunctional(Jform, c, phi, derivative_cb_pre=deriv_cb)
P1hat = LevelSetFunctional(Power1, c, phi)
P1control = fda.Control(Power1)
P2hat = LevelSetFunctional(Power2, c, phi)
P2control = fda.Control(Power2)
Jhat_v = Jhat(phi)
print("Initial cost function value {:.5f}".format(Jhat_v), flush=True)
print("Power drop 1 {:.5f}".format(Power1), flush=True)
print("Power drop 2 {:.5f}".format(Power2), flush=True)
beta_param = 0.08
# Regularize the shape derivatives only in the domain marked with 0
reg_solver = RegularizationSolver(S,
mesh,
beta=beta_param,
gamma=1e5,
dx=dx,
design_domain=0)
tol = 1e-5
dt = 0.05
params = {
"alphaC": 1.0,
"debug": 5,
"alphaJ": 1.0,
"dt": dt,
"K": 1e-3,
"maxit": n_iters,
"maxtrials": 5,
"itnormalisation": 10,
"tol_merit":
5e-3, # new merit can be within 0.5% of the previous merit
# "normalize_tol" : -1,
"tol": tol,
}
solver_parameters = {
"reinit_solver": {
"h_factor": 2.0,
}
}
# Optimization problem
problem = InfDimProblem(
Jhat,
reg_solver,
ineqconstraints=[
Constraint(P1hat, 1.0, P1control),
Constraint(P2hat, 1.0, P2control),
],
solver_parameters=solver_parameters,
)
results = nlspace_solve(problem, params)
return results
mu = fd.Constant(mu)
rho = fd.Constant(rho)
n = fd.FacetNormal(mesh)
mul = fd.Function(DG0)
mul.interpolate(fd.conditional(logical, mueff[0], mueff[1]))
def D(x):
return 2*fd.sym(fd.grad(x))
# weak form -> SB
F_sb = 2*mu*fd.inner(D(v), D(u))*fd.dx - p*fd.div(v)*fd.dx + \
mu*ikm*fd.inner(u, v)*fd.dx + q*fd.div(u)*fd.dx - \
2*mu*fd.inner(fd.avg(D(u)*n), fd.jump(v))*fd.dS - \
2*mu*Beta*fd.inner(fd.avg(D(v)*n), fd.jump(u))*fd.dS -\
fd.inner(f, v)*fd.dx - p_out*fd.inner(v, n)*fd.ds(outlet)
# -------------------------------------------------------------------------
# transport
# coefficients
Dt = fd.Constant(np.sqrt(tol))
c_mid = 0.5 * (c0 + c) # Crank-Nicolson timestepping
# stability
gamma = fd.Constant(2*(mesh.geometric_dimension()+1) /
mesh.geometric_dimension())
h_E = fd.sqrt(2) * fd.CellVolume(mesh) / fd.CellDiameter(mesh)
Ty('+') * Ty('-') / fd.avg(Ty), 0.0)
Tz_facet = fd.conditional(fd.gt(fd.avg(Tz), 0.0),
Tz('+') * Tz('-') / fd.avg(Tz), 0.0)
T_facet = (Tx_facet * (abs(n[0]('+')) + abs(n[0]('-'))) / 2 + Ty_facet *
(abs(n[1]('+')) + abs(n[1]('-'))) / 2 + Tz_facet *
(abs(n[2]('+')) + abs(n[2]('-'))) / 2)
# We can now define the bilinear and linear forms for the left and right
x, y, z = mesh.coordinates
x_func = fd.interpolate(x, V)
y_func = fd.interpolate(y, V)
z_func = fd.interpolate(z, V)
Delta_x = fd.jump(x_func)
Delta_y = fd.jump(y_func)
Delta_z = fd.jump(z_func)
Delta = fd.sqrt(Delta_x**2 + Delta_y**2 + Delta_z**2)
dx = fd.dx
alpha = 1
gamma = 1
m = u * v * w * dx
a = T_facet/Delta*fd.dot(fd.jump(v,n), fd.jump(u,n))*fd.dS_h \
+ T_facet/Delta*fd.dot(fd.jump(v,n), fd.jump(u,n))*fd.dS_v
# Defining the eigenvalue problem
def __init__(self, state, V):
super(EulerPoincareForm, self).__init__(state)
dt = state.timestepping.dt
w = TestFunction(V)
u = TrialFunction(V)
self.u0 = Function(V)
ustar = 0.5*(self.u0 + u)
n = FacetNormal(state.mesh)
Upwind = 0.5*(sign(dot(self.ubar, n))+1)
if state.mesh.geometric_dimension() == 3:
if V.extruded:
surface_measure = (dS_h + dS_v)
else:
surface_measure = dS
# <w,curl(u) cross ubar + grad( u.ubar)>
# =<curl(u),ubar cross w> - <div(w), u.ubar>
# =<u,curl(ubar cross w)> -
# <<u_upwind, [[n cross(ubar cross w)cross]]>>
both = lambda u: 2*avg(u)
Eqn = (
inner(w, u-self.u0)*dx
+ dt*inner(ustar, curl(cross(self.ubar, w)))*dx
- dt*inner(both(Upwind*ustar),
both(cross(n, cross(self.ubar, w))))*surface_measure
- dt*div(w)*inner(ustar, self.ubar)*dx
)
# define surface measure and terms involving perp differently
# for slice (i.e. if V.extruded is True) and shallow water
# (V.extruded is False)
else:
if V.extruded:
surface_measure = (dS_h + dS_v)
perp = lambda u: as_vector([-u[1], u[0]])
perp_u_upwind = Upwind('+')*perp(ustar('+')) + Upwind('-')*perp(ustar('-'))
else:
surface_measure = dS
outward_normals = CellNormal(state.mesh)
perp = lambda u: cross(outward_normals, u)
perp_u_upwind = Upwind('+')*cross(outward_normals('+'),ustar('+')) + Upwind('-')*cross(outward_normals('-'),ustar('-'))
Eqn = (
(inner(w, u-self.u0)
- dt*inner(w, div(perp(ustar))*perp(self.ubar))
- dt*div(w)*inner(ustar, self.ubar))*dx
- dt*inner(jump(inner(w, perp(self.ubar)), n),
perp_u_upwind)*surface_measure
+ dt*jump(inner(w,
perp(self.ubar))*perp(ustar), n)*surface_measure
)
a = lhs(Eqn)
L = rhs(Eqn)
self.u1 = Function(V)
uproblem = LinearVariationalProblem(a, L, self.u1)
self.usolver = LinearVariationalSolver(uproblem,
options_prefix='EPAdvection')
def _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')
"pc_type": "lu"
})
q_solver.solve()
# Build advection, forcing forms
Frhs = unk * Dnk / 3. + un * Dnk / 6. + unk * Dn / 6. + un * Dn / 3.
K = inner(un, un) / 3. + inner(un, unk) / 3. + inner(unk, unk) / 3.
Prhs = g * (0.5 * (Dn + Dp) + b) + 0.5 * K
# D advection solver
phi = TestFunction(W0)
D_ = TrialFunction(W0)
D_bar = 0.5 * (Dn + D_)
D_ad = (inner(grad(phi), D_bar * u_rec) * dx -
jump(phi * u_rec, n) * uw(ubar, D_bar) * dS)
D_eqn = (D_ - Dn) * phi * dx - dt * D_ad
Dad_p = LinearVariationalProblem(lhs(D_eqn), rhs(D_eqn), Dp)
D_ad_solver = LinearVariationalSolver(Dad_p)
# u advection solver
w = TestFunction(W1)
u_ = TrialFunction(W1)
u_bar = 0.5 * (un + u_)
u_ad = (inner(perp(grad(inner(w, perp(u_rec)))), u_bar) * dx +
inner(jump(inner(w, perp(u_rec)), n), p_uw(ubar, u_bar)) * dS)
u_eqn = inner(u_ - un, w) * dx - dt * u_ad
uad_p = LinearVariationalProblem(lhs(u_eqn), rhs(u_eqn), uad)
u_ad_solver = LinearVariationalSolver(uad_p)