def test_diffusion(params):
E_true = firedrake.interpolate(E_surface + q_bed / α * h * (1 - ζ), Q)
E = firedrake.interpolate(Constant(480), Q)
# Subclass the heat transport model and turn off advection so that we can
# test diffusion by itself
class DiffusionTransportModel(icepack.models.HeatTransport3D):
def __init__(self):
super(DiffusionTransportModel, self).__init__()
def advective_flux(self, **kwargs):
E = kwargs['energy']
h = kwargs['thickness']
Q = E.function_space()
ψ = firedrake.TestFunction(Q)
return firedrake.Constant(0) * ψ * h * dx
model = DiffusionTransportModel()
solver = icepack.solvers.HeatTransportSolver(model,
solver_parameters=params)
dt = 250.0
final_time = 6000
num_steps = int(final_time / dt) + 1
for step in range(num_steps):
E = solver.solve(dt,
energy=E,
thickness=h,
energy_surface=Constant(E_surface),
heat=Constant(0),
heat_bed=Constant(q_bed))
assert assemble((E - E_true)**2 * ds_t) / assemble(E_true**2 * ds_t) < 1e-3
assert assemble((E - E_true)**2 * ds_b) / assemble(E_true**2 * ds_b) < 1e-3
def compute_inf_sup_constant(spaces, b, inner_products):
V, Q = spaces
m, n = inner_products
Z = V * Q
u, p = firedrake.TrialFunctions(Z)
v, q = firedrake.TestFunctions(Z)
A = assemble(-(m(u, v) + b(v, p) + b(u, q)), mat_type='aij').M.handle
M = assemble(n(p, q), mat_type='aij').M.handle
opts = PETSc.Options()
opts.setValue('eps_gen_hermitian', None)
opts.setValue('eps_target_real', None)
opts.setValue('eps_smallest_magnitude', None)
opts.setValue('st_type', 'sinvert')
opts.setValue('st_ksp_type', 'preonly')
opts.setValue('st_pc_type', 'lu')
opts.setValue('st_pc_factor_mat_solver_type', 'mumps')
opts.setValue('eps_tol', 1e-8)
num_values = 1
eigensolver = SLEPc.EPS().create(comm=firedrake.COMM_WORLD)
eigensolver.setDimensions(num_values)
eigensolver.setOperators(A, M)
eigensolver.setFromOptions()
eigensolver.solve()
Vr, Vi = A.getVecs()
λ = eigensolver.getEigenpair(0, Vr, Vi)
return λ
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 form_function(ts, t, X, F):
r"""Form the residual for this problem
:arg ts: a PETSc TS object
:arg t: the time at step/stage being solved
:arg X: state vector
:arg F: function vector
"""
dm = ts.getDM()
ctx = dmhooks.get_appctx(dm)
# X may not be the same vector as the vec behind self._problem.phi_n, so
# copy guess in from X.
with ctx._problem.phi_n.dat.vec_wo as v:
X.copy(v)
b1 = assemble(ctx._problem.L1)
ctx._problem.solver1.solve(ctx._problem.p1, b1)
b2 = assemble(ctx._problem.L2)
ctx._problem.solver2.solve(ctx._problem.p2, b2)
b3 = assemble(ctx._problem.Lb)
ctx._problem.solver_b.solve(ctx._F, b3)
# F may not be the same vector as self._F, so copy
# residual out to F.
with ctx._F.dat.vec_ro as v:
v.copy(F)
def _newton_solve(z, E, scale, tolerance=1e-6, armijo=1e-4, max_iterations=50):
F = derivative(E, z)
H = derivative(F, z)
Q = z.function_space()
bc = firedrake.DirichletBC(Q, 0, 'on_boundary')
p = firedrake.Function(Q)
for iteration in range(max_iterations):
firedrake.solve(H == -F, p, bc,
solver_parameters={'ksp_type': 'preonly', 'pc_type': 'lu'})
dE_dp = assemble(action(F, p))
α = 1.0
E0 = assemble(E)
Ez = assemble(replace(E, {z: z + firedrake.Constant(α) * p}))
while (Ez > E0 + armijo * α * dE_dp) or np.isnan(Ez):
α /= 2
Ez = assemble(replace(E, {z: z + firedrake.Constant(α) * p}))
z.assign(z + firedrake.Constant(α) * p)
if abs(dE_dp) < tolerance * assemble(scale):
return z
raise ValueError("Newton solver failed to converge after {0} iterations"
.format(max_iterations))
def _ad_convert_type(self, value, options=None):
from firedrake import Function, TrialFunction, TestFunction, assemble
options = {} if options is None else options
riesz_representation = options.get("riesz_representation", "l2")
if riesz_representation == "l2":
return Function(self.function_space(), val=value)
elif riesz_representation == "L2":
ret = Function(self.function_space())
u = TrialFunction(self.function_space())
v = TestFunction(self.function_space())
M = assemble(firedrake.inner(u, v) * firedrake.dx)
firedrake.solve(M, ret, value)
return ret
elif riesz_representation == "H1":
ret = Function(self.function_space())
u = TrialFunction(self.function_space())
v = TestFunction(self.function_space())
M = assemble(
firedrake.inner(u, v) * firedrake.dx +
firedrake.inner(firedrake.grad(u), firedrake.grad(v)) *
firedrake.dx)
firedrake.solve(M, ret, value)
return ret
elif callable(riesz_representation):
return riesz_representation(value)
else:
raise NotImplementedError("Unknown Riesz representation %s" %
riesz_representation)
def __init__(self, V, fixed_dims=[], direct_solve=False):
if isinstance(fixed_dims, int):
fixed_dims = [fixed_dims]
self.V = V
self.fixed_dims = fixed_dims
self.direct_solve = direct_solve
self.zero = fd.Constant(V.mesh().topological_dimension() * (0, ))
u = fd.TrialFunction(V)
v = fd.TestFunction(V)
self.zero_fun = fd.Function(V)
self.a = 1e-2 * \
fd.inner(u, v) * fd.dx + fd.inner(fd.sym(fd.grad(u)),
fd.sym(fd.grad(v))) * fd.dx
self.bc_fun = fd.Function(V)
if len(self.fixed_dims) == 0:
bcs = [fd.DirichletBC(self.V, self.bc_fun, "on_boundary")]
else:
bcs = []
for i in range(self.V.mesh().topological_dimension()):
if i in self.fixed_dims:
bcs.append(fd.DirichletBC(self.V.sub(i), 0, "on_boundary"))
else:
bcs.append(
fd.DirichletBC(self.V.sub(i), self.bc_fun.sub(i),
"on_boundary"))
self.A_ext = fd.assemble(self.a, bcs=bcs, mat_type="aij")
self.ls_ext = fd.LinearSolver(self.A_ext,
solver_parameters=self.get_params())
self.A_adj = fd.assemble(self.a,
bcs=fd.DirichletBC(self.V, self.zero,
"on_boundary"),
mat_type="aij")
self.ls_adj = fd.LinearSolver(self.A_adj,
solver_parameters=self.get_params())
def _update_centroids(self, field):
"""
Update centroid values
"""
assemble(inner(field, TestFunction(self.P0)) * dx,
tensor=self.centroids_rhs)
self.centroid_solver.solve(self.centroids, self.centroids_rhs)
def area(self, mesh):
if not hasattr(self, "_area"):
V = FunctionSpace(mesh, "DG", 0)
self.expr = TestFunction(V)*dx
self._area = Function(V)
assemble(self.expr, tensor=self._area)
return self._area
def test_advection_diffusion():
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()
solver = icepack.solvers.HeatTransportSolver(model)
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),
)
rms = np.sqrt(assemble(E**2 * h * dx) / assemble(h * dx))
assert (E_surface - 5 < rms) and (rms < E_surface + 5 + q_bed / α * h0)
def callback(inverse_solver):
E, R = inverse_solver.objective, inverse_solver.regularization
misfit = firedrake.assemble(E) / area
regularization = firedrake.assemble(R) / area
q = inverse_solver.parameter
error = firedrake.norm(q - q_true) / np.sqrt(area)
print(misfit, regularization, error, flush=True)
def calculate_Pi0(state, theta0, rho0):
# exner function
Vr = rho0.function_space()
Pi = Function(Vr).interpolate(thermodynamics.pi(state.parameters, rho0, theta0))
Pi0 = assemble(Pi*dx)/assemble(Constant(1)*dx(domain=state.mesh))
return Pi0
def get_trace_nullspace_vecs(forward, nullspace, V, V_d, TraceSpace):
"""Gets the nullspace vectors corresponding to the Schur complement
system for the multipliers.
:arg forward: A Slate expression denoting the forward elimination
operator.
:arg nullspace: A nullspace for the original mixed problem
:arg V: The original "unbroken" space.
:arg V_d: The broken space.
:arg TraceSpace: The space of approximate traces.
Returns: A list of vectors describing the nullspace of the multiplier
system
"""
from firedrake import project, assemble, Function
vecs = nullspace.getVecs()
tmp = Function(V)
tmp_b = Function(V_d)
tnsp_tmp = Function(TraceSpace)
forward_action = forward * tmp_b
new_vecs = []
for v in vecs:
with tmp.dat.vec as t:
v.copy(t)
project(tmp, tmp_b)
assemble(forward_action, tensor=tnsp_tmp)
with tnsp_tmp.dat.vec_ro as v:
new_vecs.append(v.copy())
return new_vecs
def L2_error(solution, true_solution, integration_measure):
dx = integration_measure
w_h = solution.split()
if len(w_h) == 1:
u_h = w_h[0]
u = true_solution
e = math.sqrt(fe.assemble(fe.inner(u_h - u, u_h - u) * dx))
else:
e = 0.
for u_h, u in zip(w_h, true_solution):
e += fe.assemble(fe.inner(u_h - u, u_h - u) * dx)
e = math.sqrt(e)
return e
def test_advection_diffusion():
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):
E = model.solve(dt,
E0=E,
u=u,
w=w,
h=h,
s=s,
q=Constant(0),
q_bed=q_bed,
E_inflow=E_initial,
E_surface=Constant(E_surface))
rms = np.sqrt(assemble(E**2 * h * dx) / assemble(h * dx))
assert (E_surface - 5 < rms) and (rms < E_surface + 5 + q_bed / α * h0)
def qoi_eval(prob,this_qoi,comm):
"""Helper function that evaluates qois.
prob - Helmholtz problem (or, for testing purposes only, a float)
this_qoi - string, one of ['testing','integral','origin']
comm - the communicator for spatial parallelism.
output - the value of the qoi for this realisation of the
problem. None if this_qoi is not in the list above.
"""
if this_qoi == 'testing':
output = prob
elif this_qoi == 'integral':
# This is currently a bit of a hack, because there's a bug
# in complex firedrake.
# It's also non-obvious why this works in parallel....
V = prob.u_h.function_space()
func_real = fd.Function(V)
func_imag = fd.Function(V)
func_real.dat.data[:] = np.real(prob.u_h.dat.data)
func_imag.dat.data[:] = np.imag(prob.u_h.dat.data)
output = fd.assemble(func_real * fd.dx) + 1j * fd.assemble(func_imag * fd.dx)
elif this_qoi == 'origin':
# This gives the value of the function at (0,0).
output = eval_at_mesh_point(prob.u_h,np.array([0.0,0.0]),comm)
elif this_qoi == 'top_right':
# This gives the value of the function at (1,1).
output = eval_at_mesh_point(prob.u_h,np.array([1.0,1.0]),comm)
elif this_qoi == 'gradient_top_right':
# This gives the gradient of the solution at the
# top-right-hand corner of the domain.
gradient = fd.grad(prob.u_h)
DG_spaces = [fd.FunctionSpace(prob.V.mesh(),"DG",1) for ii in range(len(gradient))]
DG_functions = [fd.Function(DG_space) for DG_space in DG_spaces]
for ii in range(len(DG_functions)):
DG_functions[ii].interpolate(gradient[ii])
point = tuple([1.0 for ii in range(len(gradient))])
# A bit funny because output needs to be a column vector
#output = np.array([eval_at_mesh_point(DG_fun,point,comm) for DG_fun in DG_functions],ndmin=2).transpose()
# For now, set the output to be the first component of the gradient
output = eval_at_mesh_point(DG_functions[0],point,comm)
else:
output = None
return output
def setup(self, state):
if not self._initialised:
super(CourantNumber, self).setup(state)
# set up area computation
V = state.spaces("DG0")
test = TestFunction(V)
self.area = Function(V)
assemble(test*dx, tensor=self.area)
def calculate_exner0(state, theta0, rho0):
# exner function
Vr = rho0.function_space()
exner = Function(Vr).interpolate(
thermodynamics.exner_pressure(state.parameters, rho0, theta0))
exner0 = assemble(exner * dx) / assemble(
Constant(1) * dx(domain=state.mesh))
return exner0
def callback(solver):
q = solver.search_direction
dJ = solver.gradient
dJ_dq = firedrake.action(dJ, q)
decrement = firedrake.assemble(dJ_dq, **fc_params) / area
error = firedrake.assemble(solver.objective) / area
penalty = firedrake.assemble(solver.regularization) / area
print(f'{error:10.4g} | {penalty:10.4g} | {decrement:10.4g}')
def solve_something(mesh):
V = fd.FunctionSpace(mesh, "CG", 1)
u = fd.Function(V)
v = fd.TestFunction(V)
x, y = fd.SpatialCoordinate(mesh)
# f = fd.sin(x) * fd.sin(y) + x**2 + y**2
# uex = x**4 * y**4
uex = fd.sin(x) * fd.sin(y) #*(x*y)**3
# def source(xs, ys):
# return 1/((x-xs)**2+(y-ys)**2 + 0.1)
# uex = source(0, 0)
uex = uex - fd.assemble(uex * fd.dx) / fd.assemble(1 * fd.dx(domain=mesh))
# f = fd.conditional(fd.ge(abs(x)-abs(y), 0), 1, 0)
from firedrake import inner, grad, dx, ds, div, sym
eps = fd.Constant(0.0)
f = uex - div(grad(uex)) + eps * div(grad(div(grad(uex))))
n = fd.FacetNormal(mesh)
g = inner(grad(uex), n)
g1 = inner(grad(div(grad(uex))), n)
g2 = div(grad(uex))
# F = 0.1 * inner(u, v) * dx + inner(grad(u), grad(v)) * dx + inner(grad(grad(u)), grad(grad(v))) * dx - f * v * dx - g * v * ds
F = inner(u, v) * dx + inner(grad(u),
grad(v)) * dx - f * v * dx - g * v * ds
F += eps * inner(div(grad(u)), div(grad(v))) * dx
F += eps * g1 * v * ds
F -= eps * g2 * inner(grad(v), n) * ds
# f = -div(grad(uex))
# F = inner(grad(u), grad(v)) * dx - f * v * dx
# bc = fd.DirichletBC(V, uex, "on_boundary")
bc = None
fd.solve(F == 0,
u,
bcs=bc,
solver_parameters={
"ksp_type": "cg",
"ksp_atol": 1e-13,
"ksp_rtol": 1e-13,
"ksp_dtol": 1e-13,
"ksp_stol": 1e-13,
"pc_type": "jacobi",
"pc_factor_mat_solver_type": "mumps",
"snes_type": "ksponly",
"ksp_converged_reason": None
})
print("||u-uex|| =", fd.norm(u - uex))
print("||grad(u-uex)|| =", fd.norm(grad(u - uex)))
print("||grad(grad(u-uex))|| =", fd.norm(grad(grad(u - uex))))
err = fd.Function(
fd.TensorFunctionSpace(mesh, "DG",
V.ufl_element().degree() - 2)).interpolate(
grad(grad(u - uex)))
# err = fd.Function(fd.FunctionSpace(mesh, "DG", V.ufl_element().degree())).interpolate(u-uex)
fd.File(outdir + "sln.pvd").write(u)
fd.File(outdir + "err.pvd").write(err)
def run(problem, tensor, size_factor, degree):
formname = problem.__name__
cellname = 'cube' if tensor else 'simplex'
PETSc.Sys.Print("%s: %s, degree=%d" % (formname, cellname, degree))
num_cells = COMM_WORLD.size * max(1, 1e7 * size_factor / (degree + 1)**{'spectral': 4, 'coffee': 6}[args.mode])
h = int(floor(cbrt(num_cells / COMM_WORLD.size)))
w = int(floor(sqrt(num_cells / h)))
d = int(round(num_cells / (w * h)))
num_cells = w * d * h
if tensor:
mesh = ExtrudedMesh(UnitSquareMesh(w, d, quadrilateral=True), h)
else:
mesh = UnitCubeMesh(w, d, h)
comm = mesh.comm
J = problem(mesh, int(degree))
# Warmup and allocate
A = assemble(J, mat_type="matfree", form_compiler_parameters={'mode': args.mode})
A.force_evaluation()
Ap = A.petscmat
x, y = Ap.createVecs()
assert x.size == y.size
num_dofs = x.size
Ap.mult(x, y)
stage = PETSc.Log.Stage("%s(%d) %s" % (formname, degree, cellname))
with stage:
assemble(J, mat_type="matfree", form_compiler_parameters={'mode': args.mode}, tensor=A)
A.force_evaluation()
Ap = A.petscmat
for _ in range(num_matvecs):
Ap.mult(x, y)
parloop = parloop_event.getPerfInfo()
matmult = matmult_event.getPerfInfo()
assert parloop["count"] == 2 * num_matvecs
assert matmult["count"] == num_matvecs
parloop_time = comm.allreduce(parloop["time"], op=MPI.SUM) / (comm.size * num_matvecs)
matmult_time = comm.allreduce(matmult["time"], op=MPI.SUM) / (comm.size * num_matvecs)
matfree_overhead = (1 - parloop_time / matmult_time)
PETSc.Sys.Print("Matrix-free action overhead: %.1f%%" % (matfree_overhead * 100,))
if COMM_WORLD.rank == 0:
header = not os.path.exists(filepath)
data = {"num_cells": num_cells,
"num_dofs": num_dofs,
"num_procs": comm.size,
"tsfc_mode": args.mode,
"problem": formname,
"cell_type": cellname,
"degree": degree,
"matmult_time": matmult_time,
"parloop_time": parloop_time}
df = pandas.DataFrame(data, index=[0])
df.to_csv(filepath, index=False, mode='a', header=header)
def derivative(self, out):
"""
Assemble partial directional derivative wrt ControlSpace perturbations.
"""
v = fd.TestFunction(self.V_control)
fd.assemble(self.derivative_form(v),
tensor=self.deriv_r,
form_compiler_parameters=self.params)
out.fun.assign(self.deriv_r)
out.scale(self.scale)
def estimate_timestep(mesh, V, c, estimate_max_eigenvalue=True):
"""Estimate the maximum stable timestep based on the spectral radius
using optionally the Gershgorin Circle Theorem to estimate the
maximum generalized eigenvalue. Otherwise computes the maximum
generalized eigenvalue exactly
ONLY WORKS WITH KMV ELEMENTS
"""
u, v = fire.TrialFunction(V), fire.TestFunction(V)
quad_rule = finat.quadrature.make_quadrature(V.finat_element.cell,
V.ufl_element().degree(),
"KMV")
dxlump = fire.dx(rule=quad_rule)
A = fire.assemble(u * v * dxlump)
ai, aj, av = A.petscmat.getValuesCSR()
av_inv = []
for value in av:
if value == 0:
av_inv.append(0.0)
else:
av_inv.append(1 / value)
Asp = scipy.sparse.csr_matrix((av, aj, ai))
Asp_inv = scipy.sparse.csr_matrix((av_inv, aj, ai))
K = fire.assemble(c * c * dot(grad(u), grad(v)) * dxlump)
ai, aj, av = K.petscmat.getValuesCSR()
Ksp = scipy.sparse.csr_matrix((av, aj, ai))
# operator
Lsp = Asp_inv.multiply(Ksp)
if estimate_max_eigenvalue:
# absolute maximum of diagonals
max_eigval = np.amax(np.abs(Lsp.diagonal()))
else:
print(
"Computing exact eigenvalues is extremely computationally demanding!",
flush=True,
)
max_eigval = scipy.sparse.linalg.eigs(Ksp,
M=Asp,
k=1,
which="LM",
return_eigenvectors=False)[0]
# print(max_eigval)
if np.sqrt(max_eigval) > 0.0:
max_dt = np.float(2 / np.sqrt(max_eigval))
else:
max_dt = 100000000
#print(
# f"Maximum stable timestep should be about: {np.float(2 / np.sqrt(max_eigval))} seconds",
# flush=True,
#)
return max_dt
def apply_massinv(self):
if not self.constant_jacobian:
firedrake.assemble(self.A.a, tensor=self.A, bcs=self.bcs,
form_compiler_parameters=self.form_compiler_parameters)
if self.use_slate_for_inverse:
def solve(x, b):
with x.dat.vec_wo as x_, b.dat.vec_ro as b_:
self.A.petscmat.mult(b_, x_)
return solve
else:
return self.solver.solve
def apply_massinv(self):
if not self.constant_jacobian:
firedrake.assemble(self.A.a, tensor=self.A, bcs=self.bcs,
form_compiler_parameters=self.form_compiler_parameters)
if self.use_slate_for_inverse:
def solve(x, b):
self.A.force_evaluation()
with x.dat.vec_wo as x_, b.dat.vec_ro as b_:
self.A.petscmat.mult(b_, x_)
return solve
else:
return self.solver.solve
def stepInfo(solver):
"""Printer summary information from solver
Parameters
----------
solver : firedrake solver
solver function
"""
E = firedrake.assemble(solver.objective)
R = firedrake.assemble(solver.regularization)
area = firedrake.assemble(firedrake.Constant(1) * firedrake.dx(mesh))
Print(f'{E/area:g}, {R/area:g} '
f'{datetime.now().strftime("%H:%M:%S")} {area:10.3e}')
def parameterInfo(solver, area, message=''):
"""
Print various statistics
"""
floating = solver.problem.diagnostic_solve_kwargs['floating']
grounded = solver.problem.diagnostic_solve_kwargs['grounded']
areaFloating = firedrake.assemble(floating * firedrake.dx)
areaGrounded = firedrake.assemble(grounded * firedrake.dx)
avgFloat = firedrake.assemble(solver.parameter * floating *
firedrake.dx) / areaFloating
avgGrounded = firedrake.assemble(solver.parameter * grounded *
firedrake.dx) / areaGrounded
Print(f'{message} grounded {avgGrounded:9.2e} floating {avgFloat:9.2e}')
def solvePDE(self):
"""Solve the heat equation and evaluate the objective function."""
self.J = 0
t = 0
self.u.assign(self.u0)
self.J += fd.assemble(self.dt * (self.u - self.u_t(t))**2 * self.dx)
for ii in range(10):
self.u_old.assign(self.u)
fd.solve(self.F(t, self.u, self.u_old) == 0, self.u, bcs=self.bcs)
t += self.dt
self.J += fd.assemble(self.dt * (self.u - self.u_t(t))**2 *
self.dx)
def postprocess(self):
p, u, T = self.solution.split()
div = fe.div
self.velocity_divergence = fe.assemble(div(u) * self.dx)
phil = self.liquid_volume_fraction(temperature=T)
self.liquid_area = fe.assemble(phil * self.dx)
return self
def write_output(diag, t, counter, dfr, nc_h5_dfr, outf, outf_ll, fld_out,
fld_out_ll):
"""Function to write vtu, txt output"""
# Update and output diagnostics
energy = e_form(un, Dn)
En.interpolate(energy)
if 'Energy' in diag.keys():
diag['Energy'] = assemble(energy * dx)
if 'Mass' in diag.keys():
diag['Mass'] = assemble(Dn * dx)
if 'Enstrophy' in diag.keys():
diag['Enstrophy'] = assemble(qn**2 * Dn * dx)
if 'L2_error_D' in diag.keys():
diag['L2_error_D'] = sqrt(assemble((Dn - D0) * (Dn - D0) * dx))
if 'L2_error_u' in diag.keys():
diag['L2_error_u'] = sqrt(assemble(inner(un - u0, un - u0) * dx))
_c = 0
for k in diag:
nc_temp[counter % nc_h5_dfr][_c] = diag[k]
_c += 1
# Save I/O time by not writing at each timestep
if (counter % nc_h5_dfr) == nc_h5_dfr - 1:
if COMM_WORLD.Get_rank() == 0:
print("Timestep nr", counter, ": Writing nc files at", ctime())
write_to_nc_file(t, nc_temp)
if nc_safe:
write_to_nc_file(t, nc_temp, '_bckp')
# Write to checkpoint
chkpt.store(xn)
chkpt.write_attribute("/", "time", t)
if h5_safe:
chkpt_bckp.store(xn)
chkpt_bckp.write_attribute("/", "time", t)
# Output vtu file
if t in field_dumptimes:
#q2Dn.project(qn**2*Dn)
#eta_out.interpolate(Dn + b)
#vortsolver.solve()
u_err.assign(un - u0)
D_err.assign(Dn - D0)
q_err.assign(qn - q0)
outf.write(*fld_out)
if write_latlon:
outf_ll.write(*fld_out_ll)
def _ad_dot(self, other, options=None):
from firedrake import assemble
options = {} if options is None else options
riesz_representation = options.get("riesz_representation", "l2")
if riesz_representation == "l2":
return self.vector().inner(other.vector())
elif riesz_representation == "L2":
return assemble(firedrake.inner(self, other)*firedrake.dx)
elif riesz_representation == "H1":
return assemble((firedrake.inner(self, other)
+ firedrake.inner(firedrake.grad(self), other))*firedrake.dx)
else:
raise NotImplementedError(
"Unknown Riesz representation %s" % riesz_representation)
def derivative(self, out):
"""
Assemble partial directional derivative wrt ControlSpace perturbations.
First, assemble directional derivative (wrt FEspace V_m) and
store it in self.deriv_m. This automatically updates self.deriv_r,
which is then converted to the directional derivative wrt
ControSpace perturbations restrict.
"""
v = fd.TestFunction(self.V_m)
fd.assemble(self.derivative_form(v),
tensor=self.deriv_m,
form_compiler_parameters=self.params)
out.from_first_derivative(self.deriv_r)
out.scale(self.scale)
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 _construct_centroid_solver(self):
"""
Constructs a linear problem for computing the centroids
:return: LinearSolver instance
"""
u = TrialFunction(self.P0)
v = TestFunction(self.P0)
a = assemble(u * v * dx)
return LinearSolver(a, solver_parameters={'ksp_type': 'preonly',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'})
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 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_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 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)
print("END: Read in reservoir fields")
# Permeability field harmonic interpolation to facets
Kx_facet = fd.conditional(fd.gt(fd.avg(Kx), 0.0), Kx('+')*Kx('-') / fd.avg(Kx), 0.0)
Ky_facet = fd.conditional(fd.gt(fd.avg(Ky), 0.0), Ky('+')*Ky('-') / fd.avg(Ky), 0.0)
Kz_facet = fd.conditional(fd.gt(fd.avg(Kz), 0.0), Kz('+')*Kz('-') / fd.avg(Kz), 0.0)
# We can now define the bilinear and linear forms for the left and right
dx = fd.dx
KdivU = fd.as_vector((Kx_facet*u.dx(0), Ky_facet*u.dx(1), Kz_facet*u.dx(2)))
a = (fd.dot(KdivU, fd.grad(v))) * dx
m = u * v * phi * dx
# Defining the eigenvalue problem
petsc_a = fd.assemble(a).M.handle
petsc_m = fd.assemble(m).M.handle
num_eigenvalues = 3
# Set solver options
opts = PETSc.Options()
opts.setValue("eps_gen_hermitian", None)
#opts.setValue("st_pc_factor_shift_type", "NONZERO")
opts.setValue("eps_type", "krylovschur")
#opts.setValue("eps_smallest_magnitude", None)
opts.setValue("eps_target_magnitude", None)
opts.setValue("eps_target", 1e-8)
opts.setValue("eps_tol", 1e-10)
# Solve for eigenvalues
def _update_centroids(self, field):
"""
Update centroid values
"""
assemble(TestFunction(self.P0) * field * dx, tensor=self.centroids_rhs)
self.centroid_solver.solve(self.centroids, self.centroids_rhs)
def _update_centroids(self, field):
"""
Update centroid values
"""
b = assemble(TestFunction(self.P0) * field * dx)
self.centroid_solver.solve(self.centroids, b)
def initialize(self, pc):
_, P = pc.getOperators()
assert P.type == "python"
context = P.getPythonContext()
(self.J, self.bcs) = (context.a, context.row_bcs)
test, trial = self.J.arguments()
if test.function_space() != trial.function_space():
raise NotImplementedError("test and trial spaces must be the same")
Pk = test.function_space()
element = Pk.ufl_element()
shape = element.value_shape()
mesh = Pk.ufl_domain()
if len(shape) == 0:
P1 = firedrake.FunctionSpace(mesh, "CG", 1)
elif len(shape) == 1:
P1 = firedrake.VectorFunctionSpace(mesh, "CG", 1, dim=shape[0])
else:
P1 = firedrake.TensorFunctionSpace(mesh, "CG", 1, shape=shape,
symmetry=element.symmetry())
# TODO: A smarter low-order operator would also interpolate
# any coefficients to the coarse space.
mapper = ArgumentReplacer({test: firedrake.TestFunction(P1),
trial: firedrake.TrialFunction(P1)})
self.lo_J = map_integrands.map_integrand_dags(mapper, self.J)
lo_bcs = []
for bc in self.bcs:
# Don't actually need the value, since it's only used for
# killing parts of the restriction matrix.
lo_bcs.append(firedrake.DirichletBC(P1, firedrake.zero(P1.shape),
bc.sub_domain,
method=bc.method))
self.lo_bcs = tuple(lo_bcs)
mat_type = PETSc.Options().getString(pc.getOptionsPrefix() + "lo_mat_type",
firedrake.parameters["default_matrix_type"])
self.lo_op = firedrake.assemble(self.lo_J, bcs=self.lo_bcs,
mat_type=mat_type)
self.lo_op.force_evaluation()
A, P = pc.getOperators()
nearnullsp = P.getNearNullSpace()
if nearnullsp.handle != 0:
# Actually have a near nullspace
tmp = firedrake.Function(Pk)
low = firedrake.Function(P1)
vecs = []
for vec in nearnullsp.getVecs():
with tmp.dat.vec as v:
vec.copy(v)
low.interpolate(tmp)
with low.dat.vec_ro as v:
vecs.append(v.copy())
nullsp = PETSc.NullSpace().create(vectors=vecs, comm=pc.comm)
self.lo_op.petscmat.setNearNullSpace(nullsp)
lo = PETSc.PC().create(comm=pc.comm)
lo.incrementTabLevel(1, parent=pc)
lo.setOperators(self.lo_op.petscmat, self.lo_op.petscmat)
lo.setOptionsPrefix(pc.getOptionsPrefix() + "lo_")
lo.setFromOptions()
self.lo = lo
self.restriction = restriction_matrix(Pk, P1, self.bcs, self.lo_bcs)
self.work = self.lo_op.petscmat.createVecs()
if len(self.bcs) > 0:
bc_nodes = numpy.unique(numpy.concatenate([bc.nodes for bc in self.bcs]))
bc_nodes = bc_nodes[bc_nodes < Pk.dof_dset.size]
bc_iset = PETSc.IS().createBlock(numpy.prod(shape), bc_nodes,
comm=PETSc.COMM_SELF)
self.bc_indices = bc_iset.getIndices()
bc_iset.destroy()
else:
self.bc_indices = numpy.empty(0, dtype=numpy.int32)
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)]
def update(self, pc):
firedrake.assemble(self.lo_J, bcs=self.lo_bcs, tensor=self.lo_op)
self.lo_op.force_evaluation()
def l2(f):
return sqrt(assemble(dot(f, f)*dx))
