def test_advection():
E_initial = firedrake.interpolate(E_surface + q_bed / α * h * (1 - ζ), Q)
E = E_initial.copy(deepcopy=True)
u0 = 100.0
du = 100.0
u_expr = as_vector((u0 + du * x / Lx, 0))
u = firedrake.interpolate(u_expr, V)
w = firedrake.interpolate((-du / Lx + dh / Lx / h * u[0]) * ζ, W)
dt = 10.0
final_time = Lx / u0
num_steps = int(final_time / dt) + 1
model = icepack.models.HeatTransport3D()
for step in range(num_steps):
model._advect(dt,
E=E,
u=u,
w=w,
h=h,
s=s,
E_inflow=E_initial,
E_surface=Constant(E_surface))
error_surface = assemble((E - E_surface)**2 * ds_t)
assert error_surface / assemble(E_surface**2 * ds_t(mesh)) < 1e-2
error_bed = assemble((E - E_initial)**2 * ds_b)
assert error_bed / assemble(E_initial**2 * ds_b(mesh)) < 1e-2
def test_advection():
E_initial = firedrake.interpolate(E_surface + q_bed / α * h * (1 - ζ), Q)
E = E_initial.copy(deepcopy=True)
# Subclass the heat transport model and turn off diffusion so that we can
# test advection by itself
class AdvectionTransportModel(icepack.models.HeatTransport3D):
def __init__(self):
super(AdvectionTransportModel, self).__init__()
def diffusive_flux(self, **kwargs):
E = kwargs["energy"]
h = kwargs["thickness"]
Q = E.function_space()
ψ = firedrake.TestFunction(Q)
return Constant(0) * ψ * h * dx
model = AdvectionTransportModel()
solver = icepack.solvers.HeatTransportSolver(model)
u0 = 100.0
du = 100.0
u_expr = as_vector((u0 + du * x / Lx, 0))
u = firedrake.interpolate(u_expr, V)
w = firedrake.interpolate((-du / Lx + dh / Lx / h * u[0]) * ζ, W)
dt = 10.0
final_time = Lx / u0
num_steps = int(final_time / dt) + 1
for step in range(num_steps):
E = solver.solve(
dt,
energy=E,
velocity=u,
vertical_velocity=w,
thickness=h,
surface=s,
heat=Constant(0),
heat_bed=Constant(q_bed),
energy_inflow=E_initial,
energy_surface=Constant(E_surface),
)
error_surface = assemble((E - E_surface)**2 * ds_t)
assert error_surface / assemble(E_surface**2 * ds_t(mesh)) < 1e-2
error_bed = assemble((E - E_initial)**2 * ds_b)
assert error_bed / assemble(E_initial**2 * ds_b(mesh)) < 1e-2
def __init__(self, domain, degree):
self.ds_v = ds_v(domain=domain, degree=degree)
self.ds_t = ds_t(domain=domain, degree=degree)
self.ds_b = ds_b(domain=domain, degree=degree)
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 _setup_solver(self):
from firedrake.assemble import create_assembly_callable
import numpy as np
state = self.state
dt = state.timestepping.dt
beta = dt*state.timestepping.alpha
cp = state.parameters.cp
mu = state.mu
Vu = state.spaces("HDiv")
Vu_broken = FunctionSpace(state.mesh, BrokenElement(Vu.ufl_element()))
Vtheta = state.spaces("HDiv_v")
Vrho = state.spaces("DG")
h_deg = state.horizontal_degree
v_deg = state.vertical_degree
Vtrace = FunctionSpace(state.mesh, "HDiv Trace", degree=(h_deg, v_deg))
# Split up the rhs vector (symbolically)
u_in, rho_in, theta_in = split(state.xrhs)
# Build the function space for "broken" u and rho
# and add the trace variable
M = MixedFunctionSpace((Vu_broken, Vrho))
w, phi = TestFunctions(M)
u, rho = TrialFunctions(M)
l0 = TrialFunction(Vtrace)
dl = TestFunction(Vtrace)
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))
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))
# Mass matrix for the trace space
tM = assemble(dl('+')*l0('+')*(dS_v + dS_h)
+ dl*l0*ds_v + dl*l0*(ds_t + ds_b))
Lrhobar = Function(Vtrace)
Lpibar = Function(Vtrace)
rhopi_solver = LinearSolver(tM, solver_parameters={'ksp_type': 'cg',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'},
options_prefix='rhobarpibar_solver')
rhobar_avg = Function(Vtrace)
pibar_avg = Function(Vtrace)
def _traceRHS(f):
return (dl('+')*avg(f)*(dS_v + dS_h)
+ dl*f*ds_v + dl*f*(ds_t + ds_b))
assemble(_traceRHS(rhobar), tensor=Lrhobar)
assemble(_traceRHS(pibar), tensor=Lpibar)
# Project averages of coefficients into the trace space
with timed_region("Gusto:HybridProjectRhobar"):
rhopi_solver.solve(rhobar_avg, Lrhobar)
with timed_region("Gusto:HybridProjectPibar"):
rhopi_solver.solve(pibar_avg, Lpibar)
# 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
# "broken" u and rho system
Aeqn = (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*dot(theta_w*V(w), n)*pibar_avg('+')*dS_vp
+ beta*cp*dot(theta_w*V(w), n)*pibar_avg('+')*dS_hp
+ beta*cp*dot(theta_w*V(w), n)*pibar_avg*ds_tbp
- beta*cp*div(thetabar_w*w)*pi*dxp
+ (phi*(rho - rho_in) - beta*inner(grad(phi), u)*rhobar)*dx
+ beta*dot(phi*u, n)*rhobar_avg('+')*(dS_v + dS_h))
if mu is not None:
Aeqn += dt*mu*inner(w, k)*inner(u, k)*dx
# Form the mixed operators using Slate
# (A K)(X) = (X_r)
# (K.T 0)(l) (0 )
# where X = ("broken" u, rho)
A = Tensor(lhs(Aeqn))
X_r = Tensor(rhs(Aeqn))
# Off-diagonal block matrices containing the contributions
# of the Lagrange multipliers (surface terms in the momentum equation)
K = Tensor(beta*cp*dot(thetabar_w*w, n)*l0('+')*(dS_vp + dS_hp)
+ beta*cp*dot(thetabar_w*w, n)*l0*ds_vp
+ beta*cp*dot(thetabar_w*w, n)*l0*ds_tbp)
# X = A.inv * (X_r - K * l),
# 0 = K.T * X = -(K.T * A.inv * K) * l + K.T * A.inv * X_r,
# so (K.T * A.inv * K) * l = K.T * A.inv * X_r
# is the reduced equation for the Lagrange multipliers.
# Right-hand side expression: (Forward substitution)
Rexp = K.T * A.inv * X_r
self.R = Function(Vtrace)
# We need to rebuild R everytime data changes
self._assemble_Rexp = create_assembly_callable(Rexp, tensor=self.R)
# Schur complement operator:
Smatexp = K.T * A.inv * K
with timed_region("Gusto:HybridAssembleTraceOp"):
S = assemble(Smatexp)
S.force_evaluation()
# Set up the Linear solver for the system of Lagrange multipliers
self.lSolver = LinearSolver(S, solver_parameters=self.solver_parameters,
options_prefix='lambda_solve')
# Result function for the multiplier solution
self.lambdar = Function(Vtrace)
# Place to put result of u rho reconstruction
self.urho = Function(M)
# Reconstruction of broken u and rho
u_, rho_ = self.urho.split()
# Split operators for two-stage reconstruction
_A = A.blocks
_K = K.blocks
_Xr = X_r.blocks
A00 = _A[0, 0]
A01 = _A[0, 1]
A10 = _A[1, 0]
A11 = _A[1, 1]
K0 = _K[0, 0]
Ru = _Xr[0]
Rrho = _Xr[1]
lambda_vec = AssembledVector(self.lambdar)
# rho reconstruction
Srho = A11 - A10 * A00.inv * A01
rho_expr = Srho.solve(Rrho - A10 * A00.inv * (Ru - K0 * lambda_vec),
decomposition="PartialPivLU")
self._assemble_rho = create_assembly_callable(rho_expr, tensor=rho_)
# "broken" u reconstruction
rho_vec = AssembledVector(rho_)
u_expr = A00.solve(Ru - A01 * rho_vec - K0 * lambda_vec,
decomposition="PartialPivLU")
self._assemble_u = create_assembly_callable(u_expr, tensor=u_)
# Project broken u into the HDiv space using facet averaging.
# Weight function counting the dofs of the HDiv element:
shapes = (Vu.finat_element.space_dimension(), np.prod(Vu.shape))
weight_kernel = """
for (int i=0; i<%d; ++i) {
for (int j=0; j<%d; ++j) {
w[i][j] += 1.0;
}}""" % shapes
self._weight = Function(Vu)
par_loop(weight_kernel, dx, {"w": (self._weight, INC)})
# Averaging kernel
self._average_kernel = """
for (int i=0; i<%d; ++i) {
for (int j=0; j<%d; ++j) {
vec_out[i][j] += vec_in[i][j]/w[i][j];
}}""" % 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={'ksp_type': 'cg',
'pc_type': 'bjacobi',
'pc_sub_type': 'ilu'},
options_prefix='thetabacksubstitution')
self.bcs = [DirichletBC(Vu, 0.0, "bottom"),
DirichletBC(Vu, 0.0, "top")]