def __init__(self, state, V, direction=[1,2], params=None):
super(InteriorPenalty, self).__init__(state)
dt = state.timestepping.dt
kappa = params['kappa']
mu = params['mu']
gamma = TestFunction(V)
phi = TrialFunction(V)
self.phi1 = Function(V)
n = FacetNormal(state.mesh)
a = inner(gamma,phi)*dx + dt*inner(grad(gamma), grad(phi)*kappa)*dx
def get_flux_form(dS, M):
fluxes = (-inner(2*avg(outer(phi, n)), avg(grad(gamma)*M))
- inner(avg(grad(phi)*M), 2*avg(outer(gamma, n)))
+ mu*inner(2*avg(outer(phi, n)), 2*avg(outer(gamma, n)*kappa)))*dS
return fluxes
if 1 in direction:
a += dt*get_flux_form(dS_v, kappa)
if 2 in direction:
a += dt*get_flux_form(dS_h, kappa)
L = inner(gamma,phi)*dx
problem = LinearVariationalProblem(a, action(L,self.phi1), self.phi1)
self.solver = LinearVariationalSolver(problem)
def _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 solvevdisp(mesh,bdryids,deltah):
P1 = FunctionSpace(mesh, "CG", 1)
r = TrialFunction(P1)
s = TestFunction(P1)
a = inner(grad(r), grad(s)) * dx # note natural b.c. on outflow
L = inner(Constant(0.0), s) * dx
# WARNING: top must go *first* so closed top gets zero; is this documented behavior?
bcs = [ DirichletBC(P1, deltah, bdryids['top']),
DirichletBC(P1, Constant(0.0), (bdryids['base'],bdryids['inflow'])) ]
rsoln = Function(P1)
solve(a == L, rsoln, bcs=bcs, options_prefix='vd', solver_parameters={})
return rsoln
def V_approx_inv_mass(self, V, DG):
"""
Approximate inverse mass. Computes (cellwise) (V, V)^{-1} (V, DG).
:arg V: a function space
:arg DG: the DG space
:returns: A PETSc Mat mapping from V -> DG.
"""
key = V.dim()
try:
return self._V_approx_inv_mass[key]
except KeyError:
a = firedrake.Tensor(firedrake.inner(firedrake.TestFunction(V),
firedrake.TrialFunction(V))*firedrake.dx)
b = firedrake.Tensor(firedrake.inner(firedrake.TestFunction(V),
firedrake.TrialFunction(DG))*firedrake.dx)
M = firedrake.assemble(a.inv * b)
M.force_evaluation()
return self._V_approx_inv_mass.setdefault(key, M.petscmat)
def A(self):
u = firedrake.TrialFunction(self.target.function_space())
v = firedrake.TestFunction(self.target.function_space())
a = firedrake.inner(u, v)*firedrake.dx
if self.use_slate_for_inverse:
a = firedrake.Tensor(a).inv
A = firedrake.assemble(a, bcs=self.bcs,
form_compiler_parameters=self.form_compiler_parameters)
return A
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 DG_inv_mass(self, DG):
"""
Inverse DG mass matrix
:arg DG: the DG space
:returns: A PETSc Mat.
"""
key = DG.dim()
try:
return self._DG_inv_mass[key]
except KeyError:
M = firedrake.assemble(firedrake.Tensor(firedrake.inner(firedrake.TestFunction(DG),
firedrake.TrialFunction(DG))*firedrake.dx).inv)
M.force_evaluation()
return self._DG_inv_mass.setdefault(key, M.petscmat)
def V_DG_mass(self, V, DG):
"""
Mass matrix from between V and DG spaces.
:arg V: a function space
:arg DG: the DG space
:returns: A PETSc Mat mapping from V -> DG
"""
key = V.dim()
try:
return self._V_DG_mass[key]
except KeyError:
M = firedrake.assemble(firedrake.inner(firedrake.TestFunction(DG),
firedrake.TrialFunction(V))*firedrake.dx)
M.force_evaluation()
return self._V_DG_mass.setdefault(key, M.petscmat)
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 _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 initialize(self, pc):
from firedrake import TrialFunction, TestFunction, dx, assemble, inner, parameters
prefix = pc.getOptionsPrefix()
options_prefix = prefix + "Mp_"
# we assume P has things stuffed inside of it
_, P = pc.getOperators()
context = P.getPythonContext()
test, trial = context.a.arguments()
if test.function_space() != trial.function_space():
raise ValueError("MassInvPC only makes sense if test and trial space are the same")
V = test.function_space()
mu = context.appctx.get("mu", 1.0)
u = TrialFunction(V)
v = TestFunction(V)
# Handle vector and tensor-valued spaces.
# 1/mu goes into the inner product in case it varies spatially.
a = inner(1/mu * u, v)*dx
opts = PETSc.Options()
mat_type = opts.getString(options_prefix + "mat_type",
parameters["default_matrix_type"])
A = assemble(a, form_compiler_parameters=context.fc_params,
mat_type=mat_type, options_prefix=options_prefix)
A.force_evaluation()
Pmat = A.petscmat
Pmat.setNullSpace(P.getNullSpace())
tnullsp = P.getTransposeNullSpace()
if tnullsp.handle != 0:
Pmat.setTransposeNullSpace(tnullsp)
ksp = PETSc.KSP().create(comm=pc.comm)
ksp.incrementTabLevel(1, parent=pc)
ksp.setOperators(Pmat)
ksp.setOptionsPrefix(options_prefix)
ksp.setFromOptions()
ksp.setUp()
self.ksp = ksp
def V_inv_mass_ksp(self, V):
"""
A KSP inverting a mass matrix
:arg V: a function space.
:returns: A PETSc KSP for inverting (V, V).
"""
key = V.dim()
try:
return self._V_inv_mass_ksp[key]
except KeyError:
M = firedrake.assemble(firedrake.inner(firedrake.TestFunction(V),
firedrake.TrialFunction(V))*firedrake.dx)
M.force_evaluation()
ksp = PETSc.KSP().create(comm=V.comm)
ksp.setOperators(M.petscmat)
ksp.setOptionsPrefix("{}_prolongation_mass_".format(V.ufl_element()._short_name))
ksp.setType("preonly")
ksp.pc.setType("cholesky")
ksp.setFromOptions()
ksp.setUp()
return self._V_inv_mass_ksp.setdefault(key, ksp)
def momentum_term():
return inner(rho * hh * (u1 - u0), p) * dx
def advection_term(self, q):
L = super().advection_term(q)
L -= 0.5 * div(self.test) * inner(q, self.ubar) * dx
return L
def mass_term(self, q):
return inner(self.test, q) * dx
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
Vu = state.spaces("HDiv")
Vb = state.spaces("HDiv_v")
Vp = state.spaces("DG")
# Store time-stepping coefficients as UFL Constants
dt = Constant(Dt)
beta = Constant(beta_)
# 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((Vu, Vp))
w, phi = TestFunctions(M)
u, p = TrialFunctions(M)
# Get background fields
bbar = state.fields("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)
bcs = None if len(self.state.bcs) == 0 else self.state.bcs
# Solver for u, p
up_problem = LinearVariationalProblem(aeqn, Leqn, self.up, bcs=bcs)
# Provide callback for the nullspace of the trace system
def trace_nullsp(T):
return VectorSpaceBasis(constant=True)
appctx = {"trace_nullspace": trace_nullsp}
self.up_solver = LinearVariationalSolver(up_problem,
solver_parameters=self.solver_parameters,
appctx=appctx)
# Reconstruction of b
b = TrialFunction(Vb)
gamma = TestFunction(Vb)
u, p = self.up.split()
self.b = Function(Vb)
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 compliance_optimization(n_iters=200):
output_dir = "cantilever/"
path = os.path.abspath(__file__)
dir_path = os.path.dirname(path)
m = fd.Mesh(f"{dir_path}/mesh_cantilever.msh")
mesh = fd.MeshHierarchy(m, 0)[-1]
# 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)
lx = 2.0
ly = 1.0
phi_expr = (
-cos(6.0 / lx * pi * x) * cos(4.0 * pi * y)
- 0.6
+ max_value(200.0 * (0.01 - x ** 2 - (y - ly / 2) ** 2), 0.0)
+ max_value(100.0 * (x + y - lx - ly + 0.1), 0.0)
+ max_value(100.0 * (x - y - lx + 0.1), 0.0)
)
# 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. Elasticity
rho_min = 1e-5
beta = fd.Constant(200.0)
def hs(phi, beta):
return fd.Constant(1.0) / (
fd.Constant(1.0) + exp(-beta * phi)
) + fd.Constant(rho_min)
H1_elem = fd.VectorElement("CG", mesh.ufl_cell(), 1)
W = fd.FunctionSpace(mesh, H1_elem)
u = fd.TrialFunction(W)
v = fd.TestFunction(W)
# Elasticity parameters
E, nu = 1.0, 0.3
mu, lmbda = fd.Constant(E / (2 * (1 + nu))), fd.Constant(
E * nu / ((1 + nu) * (1 - 2 * nu))
)
def epsilon(u):
return sym(nabla_grad(u))
def sigma(v):
return 2.0 * mu * epsilon(v) + lmbda * tr(epsilon(v)) * Identity(2)
a = inner(hs(-phi, beta) * sigma(u), nabla_grad(v)) * dx
t = fd.Constant((0.0, -75.0))
L = inner(t, v) * ds(2)
bc = fd.DirichletBC(W, fd.Constant((0.0, 0.0)), 1)
parameters = {
"ksp_type": "preonly",
"pc_type": "lu",
"mat_type": "aij",
"ksp_converged_reason": None,
"pc_factor_mat_solver_type": "mumps",
}
u_sol = fd.Function(W)
F = fd.action(a, u_sol) - L
problem = fd.NonlinearVariationalProblem(F, u_sol, bcs=bc)
solver = fd.NonlinearVariationalSolver(
problem, solver_parameters=parameters
)
solver.solve()
# fd.solve(
# a == L, u_sol, bcs=[bc], solver_parameters=parameters
# ) # , nullspace=nullspace)
with fda.stop_annotating():
fd.File("u_sol.pvd").write(u_sol)
# Cost function: Compliance
J = fd.assemble(
fd.Constant(1e-2)
* inner(hs(-phi, beta) * sigma(u_sol), epsilon(u_sol))
* dx
)
# Constraint: Volume
with fda.stop_annotating():
total_volume = fd.assemble(fd.Constant(1.0) * dx(domain=mesh))
VolPen = fd.assemble(hs(-phi, beta) * dx)
# Needed to track the value of the volume
VolControl = fda.Control(VolPen)
Vval = total_volume / 2.0
phi_pvd = fd.File("phi_evolution.pvd", target_continuity=fd.H1)
def deriv_cb(phi):
with fda.stop_annotating():
phi_pvd.write(phi[0])
c = fda.Control(s)
Jhat = LevelSetFunctional(J, c, phi, derivative_cb_pre=deriv_cb)
Vhat = LevelSetFunctional(VolPen, c, phi)
beta_param = 0.1
# Boundary conditions for the shape derivatives.
# They must be zero at the boundary conditions.
bcs_vel = fd.DirichletBC(S, fd.Constant((0.0, 0.0)), (1, 2))
# Regularize the shape derivatives
reg_solver = RegularizationSolver(
S,
mesh,
beta=beta_param,
gamma=1.0e5,
dx=dx,
bcs=bcs_vel,
output_dir=None,
)
# Hamilton-Jacobi equation to advect the level set
dt = 0.05
tol = 1e-5
# Optimization problem
vol_constraint = Constraint(Vhat, Vval, VolControl)
problem = InfDimProblem(Jhat, reg_solver, ineqconstraints=vol_constraint)
parameters = {
"ksp_type": "preonly",
"pc_type": "lu",
"mat_type": "aij",
"ksp_converged_reason": None,
"pc_factor_mat_solver_type": "mumps",
}
params = {
"alphaC": 3.0,
"K": 0.1,
"debug": 5,
"alphaJ": 1.0,
"dt": dt,
"maxtrials": 10,
"maxit": n_iters,
"itnormalisation": 50,
"tol": tol,
}
results = nlspace_solve(problem, params)
return results
h_undamaged = firedrake.Function(Q)
u_undamaged = firedrake.Function(V)
chk.load(h_undamaged, name='h')
chk.load(u_undamaged, name='u')
with firedrake.DumbCheckpoint(args.damaged, mode=firedrake.FILE_READ) as chk:
h_damaged = firedrake.Function(Q)
u_damaged = firedrake.Function(V)
D = firedrake.Function(Δ)
chk.load(h_damaged, name='h')
chk.load(u_damaged, name='u')
chk.load(D, name='D')
δh = firedrake.interpolate(h_damaged - h_undamaged, Q)
speed = sqrt(inner(u_undamaged, u_undamaged))
v = u_undamaged / speed
δu = firedrake.interpolate(inner(u_damaged - u_undamaged, v), Q)
fig, axes = icepack.plot.subplots(nrows=2, ncols=2, sharex=True, sharey=True)
axes[0, 0].set_ylabel('distance (m)')
umax = icepack.norm(u_undamaged, norm_type='Linfty')
umax = (int(umax // 50) + 1) * 50
streamlines = icepack.plot.streamplot(u_undamaged, precision=1e3, density=5e3, axes=axes[0, 0])
colorbar(fig, axes[0, 0], streamlines, label='velocity (m/a)')
dumax = icepack.norm(δu, norm_type='Linfty')
dumax = int(dumax) + 1
levels = np.linspace(-dumax, +dumax, 2 * dumax + 1)
contours = icepack.plot.tricontourf(δu, levels=levels, cmap='twilight', axes=axes[0, 1])
def rhs_form(self):
v = firedrake.TestFunction(self.target.function_space())
form = firedrake.inner(self.source, v) * firedrake.dx
return form
def delta(self, u):
return sqrt(self.params.Delta_min ** 2 + 2 * self.params.e ** (-2) * inner(dev(self.strain(grad(u))),
dev(self.strain(grad(u)))) + tr(
self.strain(grad(u))) ** 2)
def rheology_term():
return inner(sigma, grad(p)) * dx
def compute_basis(self, n_basis, inner_product="L2", time_scaling=False, delta_t=None):
"""
:arg n_basis: Number of basis.
:arg inner_product: Type of inner product (L2 or H1).
:arg time_scaling: Use time scaling.
:arg delta_t: :class:`numpy.ndarray` with used timesteps to scale.
:return: Estimated error.
"""
# Build inner product matrix
V = self.snaps[-1].function_space()
if inner_product == "L2":
ip_form = inner(TrialFunction(V), TestFunction(V)) * dx
elif inner_product == "H1":
ip_form = inner(TrialFunction(V), TestFunction(V)) * dx + inner(grad(TrialFunction(V)),
grad(TestFunction(V))) * dx
ip_mat = assemble(ip_form, mat_type="aij").M.handle
M = self.snapshots_to_matrix()
# This matrix is symmetric positive semidefinite
corr_mat = np.matmul(self.snap_mat.transpose(), petsc2sp(ip_mat).dot(self.snap_mat))
# Build time scaling diagonal matrix
if time_scaling is True and delta_t is not None:
D = np.zeros((len(self.snaps), len(self.snaps)))
for i in range(len(self.snaps)):
D[i, i] = np.sqrt(delta_t[i])
D[0, 0] = np.sqrt(delta_t[0] / 2.0)
D[-1, -1] = np.sqrt(delta_t[-1] / 2.0)
# D'MD
corr_mat = np.matmul(D.transpose(), np.matmul(corr_mat, D))
self.n_basis = n_basis
# Compute eigenvalues, all real and non negative
w, v = np.linalg.eigh(corr_mat)
idx = np.argsort(w)[::-1]
w = w[idx]
v = v[:, idx]
# Skip negative entries
idx_neg = np.argwhere(w < 0)
if len(idx_neg) > 0:
# Reduce number of basis to min(n_basis, first_negative_eigenvalue)
n_basis = np.minimum(n_basis, idx_neg[0][0])
psi_mat = np.zeros((self.snaps[-1].function_space().dof_count, n_basis))
for i in range(n_basis):
psi_mat[:, i] = self.snap_mat.dot(v[:, i]) / np.sqrt(w[i])
ratio = np.sum(w[:n_basis]) / np.sum(w[:M])
self.basis_mat = PETSc.Mat().create(PETSc.COMM_WORLD)
self.basis_mat.setType('dense')
self.basis_mat.setSizes([self.snap_mat.shape[0], n_basis])
self.basis_mat.setUp()
self.basis_mat.setValues(range(self.snap_mat.shape[0]),
range(n_basis),
psi_mat.reshape((self.snap_mat.shape[0], n_basis)))
self.basis_mat.assemble()
return ratio
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 _effective_strain_rate(ε, ε_min):
return sqrt((inner(ε, ε) + trace(ε) ** 2 + ε_min ** 2) / 2)
def forcing_term():
return inner(rho * hh * cor * perp(ocean_curr - uh), p) * dx
def inner_product(self, x, y):
return fd.assemble(inner(x, y) * dx)
def stress_term(density, drag, func):
return inner(density * drag * sqrt(dot(func, func)) * func, p) * dx(degree=3)
def initialize(self, pc):
from firedrake import TrialFunction, TestFunction, dx, \
assemble, inner, grad, split, Constant, parameters
from firedrake.assemble import allocate_matrix, create_assembly_callable
prefix = pc.getOptionsPrefix() + "pcd_"
# we assume P has things stuffed inside of it
_, P = pc.getOperators()
context = P.getPythonContext()
test, trial = context.a.arguments()
if test.function_space() != trial.function_space():
raise ValueError("Pressure space test and trial space differ")
Q = test.function_space()
p = TrialFunction(Q)
q = TestFunction(Q)
mass = p*q*dx
# Regularisation to avoid having to think about nullspaces.
stiffness = inner(grad(p), grad(q))*dx + Constant(1e-6)*p*q*dx
opts = PETSc.Options()
# we're inverting Mp and Kp, so default them to assembled.
# Fp only needs its action, so default it to mat-free.
# These can of course be overridden.
# only Fp is referred to in update, so that's the only
# one we stash.
default = parameters["default_matrix_type"]
Mp_mat_type = opts.getString(prefix+"Mp_mat_type", default)
Kp_mat_type = opts.getString(prefix+"Kp_mat_type", default)
self.Fp_mat_type = opts.getString(prefix+"Fp_mat_type", "matfree")
Mp = assemble(mass, form_compiler_parameters=context.fc_params,
mat_type=Mp_mat_type)
Kp = assemble(stiffness, form_compiler_parameters=context.fc_params,
mat_type=Kp_mat_type)
Mp.force_evaluation()
Kp.force_evaluation()
# FIXME: Should we transfer nullspaces over. I think not.
Mksp = PETSc.KSP().create()
Mksp.setOptionsPrefix(prefix + "Mp_")
Mksp.setOperators(Mp.petscmat)
Mksp.setUp()
Mksp.setFromOptions()
self.Mksp = Mksp
Kksp = PETSc.KSP().create()
Kksp.setOptionsPrefix(prefix + "Kp_")
Kksp.setOperators(Kp.petscmat)
Kksp.setUp()
Kksp.setFromOptions()
self.Kksp = Kksp
state = context.appctx["state"]
Re = context.appctx.get("Re", 1.0)
velid = context.appctx["velocity_space"]
u0 = split(state)[velid]
fp = 1.0/Re * inner(grad(p), grad(q))*dx + inner(u0, grad(p))*q*dx
self.Re = Re
self.Fp = allocate_matrix(fp, form_compiler_parameters=context.fc_params,
mat_type=self.Fp_mat_type)
self._assemble_Fp = create_assembly_callable(fp, tensor=self.Fp,
form_compiler_parameters=context.fc_params,
mat_type=self.Fp_mat_type)
self._assemble_Fp()
self.Fp.force_evaluation()
Fpmat = self.Fp.petscmat
self.workspace = [Fpmat.createVecLeft() for i in (0, 1)]
inlet_angles = π * np.array([-3/4, -1/2, -1/3, -1/6])
inlet_widths = π * np.array([1/8, 1/12, 1/24, 1/12])
x = firedrake.SpatialCoordinate(mesh)
R = 200e3
u_in = 300
h_in = 350
hb = 100
dh, du = 400, 250
hs, us = [], []
for θ, ϕ in zip(inlet_angles, inlet_widths):
x0 = R * as_vector((np.cos(θ), np.sin(θ)))
v = -as_vector((np.cos(θ), np.sin(θ)))
L = inner(x - x0, v)
W = x - x0 - L * v
Rn = 2 * ϕ / π * R
q = firedrake.max_value(1 - (W / Rn)**2, 0)
hs.append(hb + q * ((h_in - hb) - dh * L /R))
us.append(firedrake.exp(-4 * (W/R)**2) * (u_in + du * L / R) * v)
h_expr = firedrake.Constant(hb)
for h in hs:
h_expr = firedrake.max_value(h, h_expr)
u_expr = sum(us)
# Interpolate the initial data to finite element spaces
h0 = firedrake.interpolate(h_expr, Q)
u0 = firedrake.interpolate(u_expr, V)
def _setup_solver(self):
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")
# Store time-stepping coefficients as UFL Constants
dt = Constant(Dt)
beta = Constant(beta_)
beta_cp = Constant(beta_ * cp)
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, 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")
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)))
# 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)
pibar_avg = Function(Vtrace)
rho_avg_prb = LinearVariationalProblem(a_tr, L_tr(rhobar), rhobar_avg)
pi_avg_prb = LinearVariationalProblem(a_tr, L_tr(pibar), pibar_avg)
rho_avg_solver = LinearVariationalSolver(rho_avg_prb,
solver_parameters=cg_ilu_parameters,
options_prefix='rhobar_avg_solver')
pi_avg_solver = LinearVariationalSolver(pi_avg_prb,
solver_parameters=cg_ilu_parameters,
options_prefix='pibar_avg_solver')
with timed_region("Gusto:HybridProjectRhobar"):
rho_avg_solver.solve()
with timed_region("Gusto:HybridProjectPibar"):
pi_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.
eqn = (
# momentum equation
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_w*V(w), n=n)*pibar_avg('+')*dS_vp
+ beta_cp*jump(theta_w*V(w), n=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
# 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 mu is not None:
eqn += dt*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.state.bcs
u_exact = x * y * (1 - fd.exp(c[0] *
(x - 1))) * (1 - fd.exp(c[1] * (y - 1))) / (
(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
def V(u):
return k*inner(u, k)
def rhs_form(self):
v = firedrake.TestFunction(self.target.function_space())
form = firedrake.inner(self.source, v)*firedrake.dx
return form
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 mass_term(self, test, trial):
r"""Return the UFL for the mass term \int test * trial * dx typically used in the time term."""
return firedrake.inner(test, trial) * self.dx
def mass(mesh, degree):
V = FunctionSpace(mesh, 'CG', degree)
u = TrialFunction(V)
v = TestFunction(V)
return inner(u, v) * dx
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)
timesteps, indices = chk.get_timesteps()
chk.set_timestep(timesteps[-1], idx=indices[-1])
chk.load(h, name='h')
chk.load(u, name='u')
x, y, ζ = firedrake.SpatialCoordinate(mesh)
P_0 = 1
P_1 = np.sqrt(3) * (2 * ζ - 1)
P_2 = np.sqrt(5) * (6 * ζ**2 - 6 * ζ + 1)
weight = P_0 - np.sqrt(3) * P_1 + np.sqrt(5) * P_2
u_b = icepack.depth_average(u, weight=weight)
weight = P_0 + np.sqrt(3) * P_1 + np.sqrt(5) * P_2
u_s = icepack.depth_average(u, weight=weight)
U_b = sqrt(inner(u_b, u_b))
U_s = sqrt(inner(u_s, u_s))
Q2d = firedrake.FunctionSpace(mesh2d, 'CG', 1)
ratio = firedrake.project(U_b / U_s, Q2d)
fig, axes = icepack.plot.subplots()
axes.set_xlim(0, 640e3)
axes.set_ylim(0, 80e3)
axes.get_yaxis().set_visible(False)
colors = icepack.plot.tripcolor(ratio, vmin=0.8, vmax=1.0, axes=axes)
fig.colorbar(colors, ax=axes, fraction=0.0075, pad=0.04)
fig.savefig(args.output, dpi=300, bbox_inches='tight')
def exact_u(h_expr, Q):
h = interpolate(h_expr,Q)
u_exact = -A0 * h**(n + 1) * inner(grad(h), grad(h)) * grad(h)
return u_exact
def initialize(self, pc):
from firedrake import TrialFunction, TestFunction, Function, DirichletBC, dx, \
assemble, Mesh, inner, grad, split, Constant, parameters
from firedrake.assemble import allocate_matrix, create_assembly_callable
prefix = pc.getOptionsPrefix() + "pcd_"
_, P = pc.getOperators()
context = P.getPythonContext()
test, trial = context.a.arguments()
Q = test.function_space()
p = TrialFunction(Q)
q = TestFunction(Q)
nu = context.appctx["nu"]
dt = context.appctx["dt"]
mass = (1.0 / nu) * p * q * dx
stiffness = inner(grad(p), grad(q)) * dx
velid = context.appctx["velocity_space"]
self.bcs = context.appctx["PCDbc"]
opts = PETSc.Options()
default = parameters["default_matrix_type"]
Mp_mat_type = opts.getString(prefix + "Mp_mat_type", default)
Kp_mat_type = opts.getString(prefix + "Kp_mat_type", default)
Fp_mat_type = opts.getString(prefix + "Fp_mat_type", "matfree")
Mp = assemble(mass,
form_compiler_parameters=context.fc_params,
mat_type=Mp_mat_type,
options_prefix=prefix + "Mp_")
Kp = assemble(stiffness,
bcs=self.bcs,
form_compiler_parameters=context.fc_params,
mat_type=Kp_mat_type,
options_prefix=prefix + "Kp_")
Mksp = PETSc.KSP().create(comm=pc.comm)
Mksp.incrementTabLevel(1, parent=pc)
Mksp.setOptionsPrefix(prefix + "Mp_")
Mksp.setOperators(Mp.petscmat)
Mksp.setUp()
Mksp.setFromOptions()
self.Mksp = Mksp
Kksp = PETSc.KSP().create(comm=pc.comm)
Kksp.incrementTabLevel(1, parent=pc)
Kksp.setOptionsPrefix(prefix + "Kp_")
Kksp.setOperators(Kp.petscmat)
Kksp.setUp()
Kksp.setFromOptions()
self.Kksp = Kksp
u = context.appctx["u"]
fp = (1.0 / nu) * ((1.0 / dt) * p + dot(u, grad(p))) * q * dx
Fp = allocate_matrix(fp,
form_compiler_parameters=context.fc_params,
mat_type=Fp_mat_type,
options_prefix=prefix + "Fp_")
self._assemble_Fp = create_assembly_callable(
fp,
tensor=Fp,
form_compiler_parameters=context.fc_params,
mat_type=Fp_mat_type)
self.Fp = Fp
self._assemble_Fp()
self.workspace = [Fp.petscmat.createVecLeft() for i in (0, 1)]
self.tmp = Function(Q)
# tau_a = vn
# numerical flux
chat = lmbd_h
###########################################
qhat = qh + tau_d * (ch - chat) * n + velocity * chat + tau_a * (ch - chat) * n
# qhat = qh + tau*(ch - chat)*n + velocity * cupw
# qhat_n = fd.dot(qh, n) + tau*(ch - chat) + chat*vn
# ###########################################
# # Lax-Friedrichs/Roe form
# ch_a = 0.5*(ch('+') + ch('-'))
# qhat = qh + tau_d*(ch - chat)*n + velocity * ch_a + tau_a*(ch - ch_a)*n
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
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 norm(u):
return fd.inner(u, u) * fd.dx
def get_flux_form(dS, M):
fluxes = (-inner(2*avg(outer(phi, n)), avg(grad(gamma)*M))
- inner(avg(grad(phi)*M), 2*avg(outer(gamma, n)))
+ mu*inner(2*avg(outer(phi, n)), 2*avg(outer(gamma, n)*kappa)))*dS
return fluxes
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 V(u):
return k*inner(u,k)
Q,
fixed_bids=fixed_bids,
direct_solve=True)
# import IPython; IPython.embed()
res = [0, 1, 10, 50, 100, 150, 200, 250, 300, 400, 499, 625, 750, 875, 999]
res = [r for r in res if r <= optre - 1]
if res[-1] != optre - 1:
res.append(optre - 1)
results = run_solver(solver, res, args)
# import sys; sys.exit()
u, _ = fd.split(solver.z)
nu = solver.nu
objective_form = nu * fd.inner(fd.sym(fd.grad(u)), fd.sym(fd.grad(u))) * fd.dx
solver.setup_adjoint(objective_form)
solver.solver_adjoint.solve()
# import sys; sys.exit()
class Constraint(fs.PdeConstraint):
def solve(self):
super().solve()
solver.solve(optre)
class Objective(fs.ShapeObjective):
def __init__(self, *args_, **kwargs_):
super().__init__(*args_, **kwargs_)
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 value_form(self):
"""Evaluate misfit functional."""
u = fd.split(self.pde_solver.solution)[0]
nu = self.pde_solver.nu
return 0.5 * nu * fd.inner(fd.grad(u), fd.grad(u)) * fd.dx
