def __init__(self, function_space, model_parameters):
super(AllenCahnHeatModel, self).__init__()
self._solution_function = fd.Function(function_space)
self._solutiondot_function = fd.Function(function_space)
test_function = fd.TestFunction(function_space)
# Split mixed FE function into separate functions and put them into namedtuple
FiniteElementFunction = namedtuple("FiniteElementFunction",
["phi", "T"])
split_solution_function = FiniteElementFunction(
*ufl.split(self._solution_function))
split_solutiondot_function = FiniteElementFunction(
*ufl.split(self._solutiondot_function))
split_test_function = FiniteElementFunction(*ufl.split(test_function))
finite_element_functions = (
split_solution_function,
split_solutiondot_function,
split_test_function,
)
F_AC = get_allen_cahn_weak_form(*finite_element_functions,
model_parameters)
F_heat = get_heat_weak_form(*finite_element_functions,
model_parameters)
F = F_AC + F_heat
self._residual = F
def get_zero_vec(self):
if self.is_DG:
fun = fd.Function(self.V_c)
else:
fun = fd.Function(self.V_r)
fun *= 0.
return fun
def create_function_marker(PHI, W, xlimits, ylimits):
x, y, z = fd.SpatialCoordinate(PHI.ufl_domain())
x_func, y_func, z_func = (
fd.Function(PHI),
fd.Function(PHI),
fd.Function(PHI),
)
with fda.stop_annotating():
x_func.interpolate(x)
y_func.interpolate(y)
z_func.interpolate(z)
domain = "{[i, j]: 0 <= i < f.dofs and 0<= j <= 3}"
instruction = f"""
f[i, j] = 1.0 if (x[i, 0] < {xlimits[1]} and x[i, 0] > {xlimits[0]}) and (y[i, 0] < {ylimits[0]} or y[i, 0] > {ylimits[1]}) and z[i, 0] < 1e-7 else 0.0
"""
I_BC = fd.Function(W)
fd.par_loop(
(domain, instruction),
dx,
{
"f": (I_BC, fd.RW),
"x": (x_func, fd.READ),
"y": (y_func, fd.READ),
"z": (z_func, fd.READ),
},
is_loopy_kernel=True,
)
return I_BC
def get_restriction_weights(coarse, fine):
mesh = coarse.mesh()
assert hasattr(mesh, "_shared_data_cache")
cache = mesh._shared_data_cache["hierarchy_restriction_weights"]
key = entity_dofs_key(coarse.finat_element.entity_dofs())
try:
return cache[key]
except KeyError:
# We hit each fine dof more than once since we loop
# elementwise over the coarse cells. So we need a count of
# how many times we did this to weight the final contribution
# appropriately.
if not (coarse.ufl_element() == fine.ufl_element()):
raise ValueError("Can't transfer between different spaces")
if coarse.finat_element.entity_dofs(
) == coarse.finat_element.entity_closure_dofs():
return cache.setdefault(key, None)
ele = coarse.ufl_element()
if isinstance(ele, ufl.VectorElement):
ele = ele.sub_elements()[0]
weights = firedrake.Function(
firedrake.FunctionSpace(fine.mesh(), ele))
else:
weights = firedrake.Function(fine)
c2f_map = coarse_to_fine_node_map(coarse, fine)
kernel = get_count_kernel(c2f_map.arity)
op2.par_loop(kernel, c2f_map.iterset,
weights.dat(op2.INC, c2f_map[op2.i[0]]))
weights.assign(1 / weights)
return cache.setdefault(key, weights)
def referencemesh(mesh, b, hinitial, Href):
'''In-place modification of an extruded mesh to create the reference mesh.
Changes the top surface to lambda = b + sqrt(Href^2 + (Hinitial - b)^2).
Assumes b,hinitial are Functions defined on the base mesh. Assumes the
input mesh is extruded 3D mesh with 0 <= z <= 1.'''
if not _admissible(b, hinitial):
assert ValueError('input hinitial not admissible')
P1base = fd.FunctionSpace(mesh._base_mesh, 'P', 1)
HH = fd.Function(P1base).interpolate(hinitial - b)
# alternative:
#lambase = fd.Function(P1base).interpolate(b + fd.max_value(Href, Hstart))
lambase = fd.Function(P1base).interpolate(b + fd.sqrt(HH**2 + Href**2))
lam = extend(mesh, lambase)
Vcoord = mesh.coordinates.function_space()
if mesh._base_mesh.cell_dimension() == 1:
x, z = fd.SpatialCoordinate(mesh)
XX = fd.Function(Vcoord).interpolate(fd.as_vector([x, lam * z]))
elif mesh._base_mesh.cell_dimension() == 2:
x, y, z = fd.SpatialCoordinate(mesh)
XX = fd.Function(Vcoord).interpolate(fd.as_vector([x, y, lam * z]))
else:
raise ValueError('only 2D and 3D reference meshes are generated')
mesh.coordinates.assign(XX)
return 0
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 __init__(self, Vf, Vc, Vf_bcs, Vc_bcs):
self.Vf = Vf
self.Vc = Vc
self.Vf_bcs = Vf_bcs
self.Vc_bcs = Vc_bcs
self.uc = firedrake.Function(Vc)
self.uf = firedrake.Function(Vf)
self.mesh = Vf.mesh()
self.weight = self.multiplicity(Vf)
with self.weight.dat.vec as w:
w.reciprocal()
tf, _, _, _, _ = tensor_product_space_query(Vf)
tc, _, _, _, _ = tensor_product_space_query(Vc)
mf = Vf.ufl_element().mapping().lower()
mc = Vc.ufl_element().mapping().lower()
if tf and tc and mf == mc:
self.Vf_map = get_permuted_map(Vf)
self.Vc_map = get_permuted_map(Vc)
self.prolong_kernel, self.restrict_kernel = self.make_blas_kernels(
Vf, Vc)
else:
self.Vf_map = Vf.cell_node_map()
self.Vc_map = Vc.cell_node_map()
self.prolong_kernel, self.restrict_kernel = self.make_kernels(
Vf, Vc)
def define_functions(self):
self.X = fd.Function(self.V, name="X") # displacement, not position
self.U = fd.Function(self.V, name="U") # velocity
self.trial = fd.TrialFunction(self.V)
# self.phi_vect = fd.Function(self.V, name="Surface forcing")
#Fgr = Function(V, name="Gravity")
self.v = fd.TestFunction(self.V)
def __init__(self, mesh_r, refinements=1, order=1):
mh = fd.MeshHierarchy(mesh_r, refinements)
self.mesh_hierarchy = mh
# Control space on coarsest mesh
self.mesh_r_coarse = self.mesh_hierarchy[0]
self.V_r_coarse = fd.VectorFunctionSpace(self.mesh_r_coarse, "CG",
order)
# Create self.id and self.T on refined mesh.
element = self.V_r_coarse.ufl_element()
self.intermediate_Ts = []
for i in range(refinements - 1):
mesh = self.mesh_hierarchy[i + 1]
V = fd.FunctionSpace(mesh, element)
self.intermediate_Ts.append(fd.Function(V))
self.mesh_r = self.mesh_hierarchy[-1]
element = self.V_r_coarse.ufl_element()
self.V_r = fd.FunctionSpace(self.mesh_r, element)
X = fd.SpatialCoordinate(self.mesh_r)
self.id = fd.Function(self.V_r).interpolate(X)
self.T = fd.Function(self.V_r, name="T")
self.T.assign(self.id)
self.mesh_m = fda.Mesh(self.T)
self.V_m = fd.FunctionSpace(self.mesh_m, element)
def create_interpolation(dmc, dmf):
cctx = firedrake.dmhooks.get_appctx(dmc)
fctx = firedrake.dmhooks.get_appctx(dmf)
manager = firedrake.dmhooks.get_transfer_manager(dmf)
V_c = cctx._problem.u.function_space()
V_f = fctx._problem.u.function_space()
row_size = V_f.dof_dset.layout_vec.getSizes()
col_size = V_c.dof_dset.layout_vec.getSizes()
cfn = firedrake.Function(V_c)
ffn = firedrake.Function(V_f)
cbcs = cctx._problem.bcs
fbcs = fctx._problem.bcs
ctx = Interpolation(cfn, ffn, manager, cbcs, fbcs)
mat = PETSc.Mat().create(comm=dmc.comm)
mat.setSizes((row_size, col_size))
mat.setType(mat.Type.PYTHON)
mat.setPythonContext(ctx)
mat.setUp()
return mat, None
def function_arg(self, g):
'''Set the value of this boundary condition.'''
if isinstance(g, firedrake.Function):
if g.function_space() != self.function_space():
raise RuntimeError("%r is defined on incompatible FunctionSpace!" % g)
self._function_arg = g
elif isinstance(g, ufl.classes.Zero):
if g.ufl_shape and g.ufl_shape != self.function_space().ufl_element().value_shape():
raise ValueError(f"Provided boundary value {g} does not match shape of space")
# Special case. Scalar zero for direct Function.assign.
self._function_arg = ufl.zero()
elif isinstance(g, ufl.classes.Expr):
if g.ufl_shape != self.function_space().ufl_element().value_shape():
raise RuntimeError(f"Provided boundary value {g} does not match shape of space")
try:
self._function_arg = firedrake.Function(self.function_space())
self._function_arg_update = firedrake.Interpolator(g, self._function_arg).interpolate
except (NotImplementedError, AttributeError):
# Element doesn't implement interpolation
self._function_arg = firedrake.Function(self.function_space()).project(g)
self._function_arg_update = firedrake.Projector(g, self._function_arg).project
else:
try:
g = as_ufl(g)
self._function_arg = g
except UFLException:
try:
# Recurse to handle this through interpolation.
self.function_arg = as_ufl(as_tensor(g))
except UFLException:
raise ValueError(f"{g} is not a valid DirichletBC expression")
def __init__(self,
mesh,
element,
variational_form_residual,
dirichlet_boundary_conditions,
initial_values,
quadrature_degree=None,
time_dependent=True,
time_stencil_size=2,
output_directory_path="output/"):
self.mesh = mesh
self.element = element
self.function_space = fe.FunctionSpace(mesh, element)
self.quadrature_degree = quadrature_degree
self.solutions = [
fe.Function(self.function_space) for i in range(time_stencil_size)
]
self.solution = self.solutions[0]
self.backup_solution = fe.Function(self.solution)
if time_dependent:
assert (time_stencil_size > 1)
self.time = fe.Constant(0.)
self.timestep_size = fe.Constant(1.)
else:
self.time = None
self.timestep_size = None
self.output_directory_path = pathlib.Path(output_directory_path)
self.solution_file = None
self.plotvars = None
self.initial_values = initial_values(sim=self)
for solution in self.solutions:
solution.assign(self.initial_values)
self.variational_form_residual = variational_form_residual(
sim=self, solution=self.solution)
self.dirichlet_boundary_conditions = \
dirichlet_boundary_conditions(sim = self)
self.snes_iteration_count = 0
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 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 test_poisson_inverse(solver_type):
Nx, Ny = 32, 32
mesh = firedrake.UnitSquareMesh(Nx, Ny)
degree = 2
Q = firedrake.FunctionSpace(mesh, "CG", degree)
x, y = firedrake.SpatialCoordinate(mesh)
q_true = interpolate(-4 * ((x - 0.5)**2 + (y - 0.5)**2), Q)
f = interpolate(firedrake.Constant(1), Q)
dirichlet_ids = [1, 2, 3, 4]
model = PoissonModel()
poisson_solver = PoissonSolver(model, dirichlet_ids=dirichlet_ids)
u_bdry = firedrake.Function(Q)
u_obs = poisson_solver.diagnostic_solve(u=u_bdry, q=q_true, f=f)
q0 = firedrake.Function(Q)
u0 = poisson_solver.diagnostic_solve(u=u_bdry, q=q0, f=f)
def callback(inverse_solver):
misfit = firedrake.assemble(inverse_solver.objective)
regularization = firedrake.assemble(inverse_solver.regularization)
q = inverse_solver.parameter
error = firedrake.norm(q - q_true)
print(misfit, regularization, error)
L = firedrake.Constant(1e-4)
problem = icepack.inverse.InverseProblem(
model=model,
objective=lambda u: 0.5 * (u - u_obs)**2 * dx,
regularization=lambda q: 0.5 * L**2 * inner(grad(q), grad(q)) * dx,
state_name="u",
state=u0,
parameter_name="q",
parameter=q0,
solver_type=PoissonSolver,
solver_kwargs={"dirichlet_ids": dirichlet_ids},
diagnostic_solve_kwargs={"f": f},
)
solver = solver_type(problem, callback)
assert solver.state is not None
assert icepack.norm(solver.state) > 0
assert icepack.norm(solver.adjoint_state) > 0
assert icepack.norm(solver.search_direction) > 0
max_iterations = 1000
iterations = solver.solve(rtol=2.5e-2,
atol=1e-8,
max_iterations=max_iterations)
print(f"Number of iterations: {iterations}")
assert iterations < max_iterations
q = solver.parameter
assert icepack.norm(q - q_true) < 0.25
def stresses(mesh, icemodel, u):
Q1 = fd.FunctionSpace(mesh, 'Q', 1)
TQ1 = fd.TensorFunctionSpace(mesh, 'Q', 1)
Du = fd.Function(TQ1).interpolate(0.5 * (fd.grad(u) + fd.grad(u).T))
Du2 = fd.Function(Q1).interpolate(0.5 * fd.inner(Du, Du) +
icemodel.eps * icemodel.Dtyp**2.0)
nu = fd.Function(Q1).interpolate(0.5 * Bn * Du2**(-1.0 / n))
nu.rename('effective viscosity')
tau = fd.Function(TQ1).interpolate(2.0 * nu * Du)
tau.rename('tau')
return tau, nu
def _setup(self, **kwargs):
for name, field in kwargs.items():
if name in self._fields.keys():
self._fields[name].assign(field)
else:
if isinstance(field, firedrake.Constant):
self._fields[name] = firedrake.Constant(field)
elif isinstance(field, firedrake.Function):
self._fields[name] = field.copy(deepcopy=True)
else:
raise TypeError(
"Input %s field has type %s, must be Constant or Function!"
% (name, type(field))
)
# Create symbolic representations of the flux and sources of damage
dt = firedrake.Constant(1.0)
flux = self.model.flux(**self.fields)
# Create the finite element mass matrix
D = self.fields["damage"]
Q = D.function_space()
φ, ψ = firedrake.TrialFunction(Q), firedrake.TestFunction(Q)
M = φ * ψ * dx
L1 = -dt * flux
D1 = firedrake.Function(Q)
D2 = firedrake.Function(Q)
L2 = firedrake.replace(L1, {D: D1})
L3 = firedrake.replace(L1, {D: D2})
dD = firedrake.Function(Q)
parameters = {
"solver_parameters": {
"ksp_type": "preonly",
"pc_type": "bjacobi",
"sub_pc_type": "ilu",
}
}
problem1 = LinearVariationalProblem(M, L1, dD)
problem2 = LinearVariationalProblem(M, L2, dD)
problem3 = LinearVariationalProblem(M, L3, dD)
solver1 = LinearVariationalSolver(problem1, **parameters)
solver2 = LinearVariationalSolver(problem2, **parameters)
solver3 = LinearVariationalSolver(problem3, **parameters)
self._solvers = [solver1, solver2, solver3]
self._stages = [D1, D2]
self._damage_change = dD
self._timestep = dt
def update_search_direction(self):
r"""Solve the Gauss-Newton system for the new search direction using
the preconditioned conjugate gradient method"""
p, q, dJ = self.parameter, self.search_direction, self.gradient
dR = derivative(self.regularization, self.parameter)
Q = q.function_space()
M = firedrake.TrialFunction(Q) * firedrake.TestFunction(Q) * dx + \
derivative(dR, p)
# Compute the preconditioned residual
z = firedrake.Function(Q)
firedrake.solve(M == -dJ,
z,
solver_parameters=self._solver_params,
form_compiler_parameters=self._fc_params)
# This variable is a search direction for a search direction, which
# is definitely not confusing at all.
s = z.copy(deepcopy=True)
q *= 0.0
old_cost = np.inf
while True:
z_mnorm = self._assemble(firedrake.energy_norm(M, z))
s_hnorm = self.gauss_newton_energy_norm(s)
α = z_mnorm / s_hnorm
δz = firedrake.Function(Q)
g = self.gauss_newton_mult(s)
firedrake.solve(M == g,
δz,
solver_parameters=self._solver_params,
form_compiler_parameters=self._fc_params)
q += α * s
z -= α * δz
β = self._assemble(firedrake.energy_norm(M, z)) / z_mnorm
s *= β
s += z
energy_norm = self.gauss_newton_energy_norm(q)
cost = 0.5 * energy_norm + self._assemble(action(dJ, q))
if (abs(old_cost - cost) /
(0.5 * energy_norm) < self._search_tolerance):
return
old_cost = cost
def _setup(self, problem, callback=(lambda s: None)):
self._problem = problem
self._callback = callback
self._p = problem.parameter.copy(deepcopy=True)
self._u = problem.state.copy(deepcopy=True)
self._model_args = dict(**problem.model_args,
dirichlet_ids=problem.dirichlet_ids)
u_name, p_name = problem.state_name, problem.parameter_name
args = dict(**self._model_args, **{u_name: self._u, p_name: self._p})
# Make the form compiler use a reasonable number of quadrature points
degree = problem.model.quadrature_degree(**args)
self._fc_params = {'quadrature_degree': degree}
# Create the error, regularization, and barrier functionals
self._E = problem.objective(self._u)
self._R = problem.regularization(self._p)
self._J = self._E + self._R
# Create the weak form of the forward model, the adjoint state, and
# the derivative of the objective functional
self._F = derivative(problem.model.action(**args), self._u)
self._dF_du = derivative(self._F, self._u)
# Create a search direction
dR = derivative(self._R, self._p)
self._solver_params = {'ksp_type': 'preonly', 'pc_type': 'lu'}
Q = self._p.function_space()
self._q = firedrake.Function(Q)
# Create the adjoint state variable
V = self.state.function_space()
self._λ = firedrake.Function(V)
dF_dp = derivative(self._F, self._p)
# Create Dirichlet BCs where they apply for the adjoint solve
rank = self._λ.ufl_element().num_sub_elements()
if rank == 0:
zero = firedrake.Constant(0)
else:
zero = firedrake.as_vector((0, ) * rank)
self._bc = firedrake.DirichletBC(V, zero, problem.dirichlet_ids)
# Create the derivative of the objective functional
self._dE = derivative(self._E, self._u)
dR = derivative(self._R, self._p)
self._dJ = (action(adjoint(dF_dp), self._λ) + dR)
def __init__(self,
equation,
solution,
fields,
dt,
bnd_conditions=None,
solver_parameters={}):
"""
:arg equation: the equation to solve
:type equation: :class:`Equation` object
:arg solution: :class:`Function` where solution will be stored
:arg fields: Dictionary of fields that are passed to the equation
:type fields: dict of :class:`Function` or :class:`Constant` objects
:arg float dt: time step in seconds
:kwarg dict bnd_conditions: Dictionary of boundary conditions passed to the equation
:kwarg dict solver_parameters: PETSc solver options
"""
super(ERKGeneric, self).__init__(equation, solution, fields, dt,
solver_parameters)
self._initialized = False
V = solution.function_space()
assert V == equation.trial_space
self.solution_old = firedrake.Function(V, name='old solution')
self.tendency = []
for i in range(self.n_stages):
k = firedrake.Function(V, name='tendency{:}'.format(i))
self.tendency.append(k)
# fully explicit evaluation
trial = firedrake.TrialFunction(V)
self.a_rk = self.equation.mass_term(self.test, trial)
self.l_rk = self.dt_const * self.equation.residual(
self.test, self.solution, self.solution, self.fields,
bnd_conditions)
self._nontrivial = self.l_rk != 0
# construct expressions for stage solutions
if self._nontrivial:
self.sol_expressions = []
for i_stage in range(self.n_stages):
sol_expr = sum(
map(operator.mul, self.tendency[:i_stage],
self.a[i_stage][:i_stage]))
self.sol_expressions.append(sol_expr)
self.final_sol_expr = sum(map(operator.mul, self.tendency, self.b))
self.update_solver()
def __init__(self,
c,
bids,
*args,
lower_bound=None,
upper_bound=None,
**kwargs):
super().__init__(*args, **kwargs)
self.T = self.Q.T
self.lam = fd.Function(self.T.function_space())
self.lower_bound = lower_bound
self.upper_bound = upper_bound
self.bids = bids
self.viol = fd.Function(self.T.function_space())
self.c = c
def get_nullspace(self, V):
"""This nullspace contains constant functions."""
dim = V.value_size
if dim == 2:
n1 = fd.Function(V).interpolate(fd.Constant((1.0, 0.0)))
n2 = fd.Function(V).interpolate(fd.Constant((0.0, 1.0)))
res = [n1, n2]
elif dim == 3:
n1 = fd.Function(V).interpolate(fd.Constant((1.0, 0.0, 0.0)))
n2 = fd.Function(V).interpolate(fd.Constant((0.0, 1.0, 0.0)))
n3 = fd.Function(V).interpolate(fd.Constant((0.0, 0.0, 1.0)))
res = [n1, n2, n3]
else:
raise NotImplementedError
return res
def __init__(self, Vf, Vc, Vf_bcs, Vc_bcs):
self.Vf = Vf
self.Vc = Vc
self.Vf_bcs = Vf_bcs
self.Vc_bcs = Vc_bcs
self.uc = firedrake.Function(Vc)
self.uf = firedrake.Function(Vf)
self.prolong_kernel = self.prolongation_transfer_kernel_action(
Vf, self.uc)
matrix_kernel = self.prolongation_transfer_kernel_action(
Vf, firedrake.TestFunction(Vc))
element_kernel = loopy.generate_code_v2(
matrix_kernel.code).device_code()
element_kernel = element_kernel.replace(
"void expression_kernel", "static void expression_kernel")
dimc = Vc.finat_element.space_dimension() * Vc.value_size
dimf = Vf.finat_element.space_dimension() * Vf.value_size
self.restrict_code = f"""
{element_kernel}
void restriction(double *restrict Rc, const double *restrict Rf, const double *restrict w)
{{
double Afc[{dimf}*{dimc}] = {{0}};
expression_kernel(Afc);
for (int32_t i = 0; i < {dimf}; i++)
for (int32_t j = 0; j < {dimc}; j++)
Rc[j] += Afc[i*{dimc} + j] * Rf[i] / w[i];
}}
"""
self.restrict_kernel = op2.Kernel(self.restrict_code, "restriction")
self.mesh = Vf.mesh()
# Lawrence's magic code for calculating dof multiplicities
shapes = (Vf.finat_element.space_dimension(), np.prod(Vf.shape))
domain = "{[i,j]: 0 <= i < %d and 0 <= j < %d}" % shapes
instructions = """
for i, j
w[i,j] = w[i,j] + 1
end
"""
self.weight = firedrake.Function(Vf)
firedrake.par_loop((domain, instructions),
firedrake.dx, {"w": (self.weight, op2.INC)},
is_loopy_kernel=True)
def gauss_newton_energy_norm(self, q):
r"""Compute the energy norm of a field w.r.t. the Gauss-Newton operator
The energy norm of a field :math:`q` w.r.t. the Gauss-Newton operator
:math:`H` can be computed using one fewer linear solve than if we were
to calculate the action of :math:`H\cdot q` on :math:`q`. This saves
computation when using the conjugate gradient method to solve for the
search direction.
"""
u, p = self.state, self.parameter
dE = derivative(self._E, u)
dR = derivative(self._R, p)
dF_du, dF_dp = self._dF_du, derivative(self._F, p)
v = firedrake.Function(u.function_space())
firedrake.solve(dF_du == action(dF_dp, q),
v,
self._bc,
solver_parameters=self._solver_params,
form_compiler_parameters=self._fc_params)
return self._assemble(
firedrake.energy_norm(derivative(dE, u), v) +
firedrake.energy_norm(derivative(dR, p), q))
def __init__(self,
*args,
rayleigh_number=7.e5,
prandtl_number=0.0216,
stefan_number=0.046,
liquidus_temperature=0.,
hotwall_temperature=1.,
initial_temperature=-0.1546,
cutoff_length=0.5,
element_degrees=(1, 2, 2),
mesh_dimensions=(20, 40),
**kwargs):
if "solution" not in kwargs:
mesh = fe.RectangleMesh(nx=mesh_dimensions[0],
ny=mesh_dimensions[1],
Lx=cutoff_length,
Ly=1.)
element = sapphire.simulations.navier_stokes_boussinesq.element(
cell=mesh.ufl_cell(), degrees=element_degrees)
kwargs["solution"] = fe.Function(fe.FunctionSpace(mesh, element))
super().__init__(*args,
reynolds_number=1. / prandtl_number,
rayleigh_number=rayleigh_number,
prandtl_number=prandtl_number,
stefan_number=stefan_number,
liquidus_temperature=liquidus_temperature,
hotwall_temperature=hotwall_temperature,
initial_temperature=initial_temperature,
**kwargs)
def test_diagnostic_solver_convergence():
shallow_ice = icepack.models.ShallowIce()
for degree in range(1, 4):
delta_x, error = [], []
for N in range(10, 110 - 20 * (degree - 1), 10):
mesh = make_mesh(R_mesh, R / N)
Q = firedrake.FunctionSpace(mesh, 'CG', degree)
V = firedrake.VectorFunctionSpace(mesh, 'CG', degree)
h_expr = Bueler_profile(mesh, R)
u_exact = interpolate(exact_u(h_expr, Q), V)
h = interpolate(h_expr, Q)
s = interpolate(h_expr, Q)
u = firedrake.Function(V)
u_num = shallow_ice.diagnostic_solve(u0=u, h=h, s=s, A=A)
error.append(norm(u_exact - u_num) / norm(u_exact))
delta_x.append(R / N)
print(delta_x[-1], error[-1])
assert assemble(shallow_ice.scale(u=u_num)) > 0
log_delta_x = np.log2(np.array(delta_x))
log_error = np.log2(np.array(error))
slope, intercept = np.polyfit(log_delta_x, log_error, 1)
print('log(error) ~= {:g} * log(dx) + {:g}'.format(slope, intercept))
assert slope > 0.9
def lift3d(q2d, Q3D):
r"""Return a 3D function that extends the given 2D function as a constant
in the vertical
This is the reverse operation of depth averaging -- it takes a 2D function
`q2d` and returns a function `q3d` defined over a 3D function space such
`q3d(x, y, z) == q2d(x, y)` for any `x, y`. The space `Q3D` of the result
must have the same horizontal element as the input function and a vertical
degree of 0.
Parameters
----------
q2d : firedrake.Function
A function defined on a 2D footprint mesh
Q3D : firedrake.Function
A function space defined on a 3D mesh extruded from the footprint mesh;
the function space must only go up to degree 0 in the vertical.
Returns
-------
q3d : firedrake.Function
The 3D-lifted input field
"""
q3d = firedrake.Function(Q3D)
assert q3d.dat.data_ro.shape == q2d.dat.data_ro.shape
q3d.dat.data[:] = q2d.dat.data_ro[:]
return q3d
def __init__(self, mesh_m):
super().__init__()
self.mesh_m = mesh_m
# Setup problem
self.V = fd.FunctionSpace(self.mesh_m, "CG", 1)
# Preallocate solution variables for state and adjoint equations
self.solution = fd.Function(self.V, name="State")
# Weak form of Poisson problem
u = self.solution
v = fd.TestFunction(self.V)
self.f = fd.Constant(4.)
self.F = (fd.inner(fd.grad(u), fd.grad(v)) - self.f * v) * fd.dx
self.bcs = fd.DirichletBC(self.V, 0., "on_boundary")
# PDE-solver parameters
self.params = {
"ksp_type": "cg",
"mat_type": "aij",
"pc_type": "hypre",
"pc_factor_mat_solver_package": "boomerang",
"ksp_rtol": 1e-11,
"ksp_atol": 1e-11,
"ksp_stol": 1e-15,
}
stateproblem = fd.NonlinearVariationalProblem(self.F,
self.solution,
bcs=self.bcs)
self.solver = fd.NonlinearVariationalSolver(
stateproblem, solver_parameters=self.params)
def getCheckPointVars(checkFile, varNames, Q, t=None):
""" Read a variable from a firedrake checkpoint file
Parameters
----------
checkFile : str
checkfile name sans .h5
varNames : str or list of str
Names of variables to extract
Q : firedrake function space
firedrake function space can be a vector space, V, but not mixed
Returns
-------
myVars: dict
{'myVar':}
"""
# Ensure a list since a single str is allowed
if type(varNames) is not list:
varNames = [varNames]
# open checkpoint
myVars = {}
if not os.path.exists(f'{checkFile}.h5'):
myerror(f'getCheckPointVar: file {checkFile}.h5 does not exist')
with firedrake.DumbCheckpoint(checkFile, mode=firedrake.FILE_READ) as chk:
if t is not None:
print(t)
chk.set_timestep(t)
for varName in varNames:
myVar = firedrake.Function(Q, name=varName)
chk.load(myVar, name=varName)
myVars[varName] = myVar
return myVars
def __init__(self, problem, callback=(lambda s: None), memory=5):
self._setup(problem, callback)
self.update_state()
self.update_adjoint_state()
Q = self.parameter.function_space()
self._memory = memory
q, dJ = self.search_direction, self.gradient
M = firedrake.TrialFunction(Q) * firedrake.TestFunction(Q) * dx
f = firedrake.Function(Q)
problem = firedrake.LinearVariationalProblem(
M, dJ, f, form_compiler_parameters=self._fc_params)
self._search_direction_solver = firedrake.LinearVariationalSolver(
problem, solver_parameters=self._solver_params)
self._search_direction_solver.solve()
q.assign(-f)
self._f = f
self._rho = []
self._ps = [self.parameter.copy(deepcopy=True)]
self._fs = [q.copy(deepcopy=True)]
self._fs[-1] *= -1
self._callback(self)
