def evaluate_adj_component(self,
inputs,
adj_inputs,
block_variable,
idx,
prepared=None):
if not self.linear and self.func == block_variable.output:
# We are not able to calculate derivatives wrt initial guess.
return None
F_form = prepared["form"]
adj_sol = prepared["adj_sol"]
adj_sol_bdy = prepared["adj_sol_bdy"]
c = block_variable.output
c_rep = block_variable.saved_output
if isinstance(c, firedrake.Function):
trial_function = firedrake.TrialFunction(c.function_space())
elif isinstance(c, firedrake.Constant):
mesh = self.compat.extract_mesh_from_form(F_form)
trial_function = firedrake.TrialFunction(
c._ad_function_space(mesh))
elif isinstance(c, firedrake.DirichletBC):
tmp_bc = self.compat.create_bc(
c,
value=self.compat.extract_subfunction(adj_sol_bdy,
c.function_space()))
return [tmp_bc]
elif isinstance(c, self.compat.MeshType):
# Using CoordianteDerivative requires us to do action before
# differentiating, might change in the future.
F_form_tmp = firedrake.action(F_form, adj_sol)
X = firedrake.SpatialCoordinate(c_rep)
dFdm = firedrake.derivative(
-F_form_tmp, X, firedrake.TestFunction(c._ad_function_space()))
dFdm = self.compat.assemble_adjoint_value(dFdm,
**self.assemble_kwargs)
return dFdm
# dFdm_cache works with original variables, not block saved outputs.
if c in self._dFdm_cache:
dFdm = self._dFdm_cache[c]
else:
dFdm = -firedrake.derivative(self.lhs, c, trial_function)
dFdm = firedrake.adjoint(dFdm)
self._dFdm_cache[c] = dFdm
# Replace the form coefficients with checkpointed values.
replace_map = self._replace_map(dFdm)
replace_map[self.func] = self.get_outputs()[0].saved_output
dFdm = replace(dFdm, replace_map)
dFdm = dFdm * adj_sol
dFdm = self.compat.assemble_adjoint_value(dFdm, **self.assemble_kwargs)
return dFdm
def get_form(self, V):
u = fd.TrialFunction(V)
v = fd.TestFunction(V)
base = self.base_form.get_form(V)
n = fd.FacetNormal(V.ufl_domain())
dim = V.mesh().topological_dimension()
def h(domain):
if dim == 3:
return (fd.FacetArea(domain) * 2)**(1. / 2)
else:
return fd.FacetArea(domain)
h = h(V.mesh())
alpha = fd.Constant(10)
from firedrake import div, grad, dS, dx, inner, jump, avg, ds
def form(u, v):
return inner(div(grad(u)), div(grad(v)))*dx \
- inner(avg(div(grad(u))), jump(grad(v), n))*dS \
- inner(jump(grad(u), n), avg(div(grad(v))))*dS \
+ alpha/h*inner(jump(grad(u), n), jump(grad(v), n))*dS \
- inner(div(grad(u)), inner(grad(v), n))*ds \
- inner(inner(grad(u), n), div(grad(v)))*ds \
+ alpha/h*inner(grad(u), grad(v))*ds
form = base + self.mu * (sum(form(u[i], v[i]) for i in range(dim)))
return form
def _advect(self, dt, E, u, w, h, s, E_inflow, E_surface):
Q = E.function_space()
φ, ψ = firedrake.TrialFunction(Q), firedrake.TestFunction(Q)
U = firedrake.as_vector((u[0], u[1], w))
flux_cells = -φ * inner(U, grad(ψ)) * h * dx
mesh = Q.mesh()
ν = facet_normal_2(mesh)
outflow = firedrake.max_value(inner(u, ν), 0)
inflow = firedrake.min_value(inner(u, ν), 0)
flux_outflow = φ * ψ * outflow * h * ds_v + \
φ * ψ * firedrake.max_value(-w, 0) * h * ds_b + \
φ * ψ * firedrake.max_value(+w, 0) * h * ds_t
F = φ * ψ * h * dx + dt * (flux_cells + flux_outflow)
flux_inflow = -E_inflow * ψ * inflow * h * ds_v \
-E_surface * ψ * firedrake.min_value(-w, 0) * h * ds_b \
-E_surface * ψ * firedrake.min_value(+w, 0) * h * ds_t
A = E * ψ * h * dx + dt * flux_inflow
solver_parameters = {'ksp_type': 'preonly', 'pc_type': 'lu'}
degree_E = E.ufl_element().degree()
degree_u = u.ufl_element().degree()
degree = (3 * degree_E[0] + degree_u[0], 2 * degree_E[1] + degree_u[1])
form_compiler_parameters = {'quadrature_degree': degree}
firedrake.solve(F == A,
E,
solver_parameters=solver_parameters,
form_compiler_parameters=form_compiler_parameters)
def create_form(form_string, family, degree, dimension, operation):
from firedrake import dx, inner, grad
import firedrake as f
f.parameters['coffee']['optlevel'] = 'O3'
f.parameters['pyop2_options']['opt_level'] = 'O3'
f.parameters['pyop2_options']['simd_isa'] = 'avx'
f.parameters["pyop2_options"]["lazy_evaluation"] = False
m = f.UnitSquareMesh(2, 2, quadrilateral=True)
if dimension == 3:
m = f.ExtrudedMesh(m, 2)
fs = f.FunctionSpace(m, family, degree)
v = f.TestFunction(fs)
if operation == "matrix":
u = f.TrialFunction(fs)
elif operation == "action":
u = f.Function(fs)
else:
raise ValueError("Unknown operation: %s" % operation)
if form_string == "u * v * dx":
return u * v * dx
elif form_string == "inner(grad(u), grad(v)) * dx":
return inner(grad(u), grad(v)) * dx
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, 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)
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 _initialise_problem(self):
"""
Set up the Firedrake machinery for solving the problem.
"""
# Define trial and test functions on the space
self._u = fd.TrialFunction(self.V)
self._v = fd.TestFunction(self.V)
# Define sesquilinear form and antilinear functional
self._a = self._define_form(self._A, self._n)
self._set_L()
self._set_pre()
# Define problem and solver (following code courtesy of Lawrence
# Mitchell, via Slack)
problem = fd.LinearVariationalProblem(self._a,
self._L,
self.u_h,
aP=self._a_pre,
constant_jacobian=False)
self._solver = fd.LinearVariationalSolver(
problem, solver_parameters=self._solver_parameters)
self._initialised = True
def get_facet_areas(mesh):
"""
Compute area of each facet of `mesh`.
The facet areas are stored as a HDiv
trace field.
Note that the plus sign is arbitrary
and could equally well be chosen as minus.
:arg mesh: the input mesh to do computations on
:rtype: firedrake.function.Function facet_areas with
facet area data
"""
HDivTrace = firedrake.FunctionSpace(mesh, "HDiv Trace", 0)
v = firedrake.TestFunction(HDivTrace)
u = firedrake.TrialFunction(HDivTrace)
facet_areas = firedrake.Function(HDivTrace, name="Facet areas")
mass_term = v("+") * u("+") * ufl.dS + v * u * ufl.ds
rhs = v("+") * ufl.FacetArea(mesh) * ufl.dS + v * ufl.FacetArea(mesh) * ufl.ds
sp = {
"snes_type": "ksponly",
"ksp_type": "preonly",
"pc_type": "jacobi",
}
firedrake.solve(mass_term == rhs, facet_areas, solver_parameters=sp)
return facet_areas
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 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 update_search_direction(self):
r"""Set the search direction to be the inverse of the mass matrix times
the gradient of the objective"""
q, dJ = self.search_direction, self.gradient
Q = q.function_space()
M = firedrake.TrialFunction(Q) * firedrake.TestFunction(Q) * dx
firedrake.solve(M == -dJ,
q,
solver_parameters=self._solver_params,
form_compiler_parameters=self._fc_params)
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 __init__(self, Q, free_bids=["on_boundary"]):
(V, I_interp) = Q.get_space_for_inner()
self.free_bids = free_bids
u = fd.TrialFunction(V)
v = fd.TestFunction(V)
n = fd.FacetNormal(V.mesh())
def surf_grad(u):
return fd.sym(fd.grad(u) - fd.outer(fd.grad(u) * n, n))
a = (fd.inner(surf_grad(u), surf_grad(v)) + fd.inner(u, v)) * fd.ds
# petsc doesn't like matrices with zero rows
a += 1e-10 * fd.inner(u, v) * fd.dx
A = fd.assemble(a, mat_type="aij")
A = A.petscmat
tdim = V.mesh().topological_dimension()
lsize = fd.Function(V).vector().local_size()
def get_nodes_bc(bc):
nodes = bc.nodes
return nodes[nodes < lsize]
def get_nodes_bid(bid):
bc = fd.DirichletBC(V, fd.Constant(tdim * (0, )), bid)
return get_nodes_bc(bc)
free_nodes = np.concatenate(
[get_nodes_bid(bid) for bid in self.free_bids])
free_dofs = np.concatenate(
[tdim * free_nodes + i for i in range(tdim)])
free_dofs = np.unique(np.sort(free_dofs))
self.free_is = PETSc.IS().createGeneral(free_dofs)
lgr, lgc = A.getLGMap()
self.global_free_is_row = lgr.applyIS(self.free_is)
self.global_free_is_col = lgc.applyIS(self.free_is)
A = A.createSubMatrix(self.global_free_is_row, self.global_free_is_col)
# A.view()
A.assemble()
self.A = A
Aksp = PETSc.KSP().create()
Aksp.setOperators(self.A)
Aksp.setOptionsPrefix("A_")
opts = PETSc.Options()
opts["A_ksp_type"] = "cg"
opts["A_pc_type"] = "hypre"
opts["A_ksp_atol"] = 1e-10
opts["A_ksp_rtol"] = 1e-10
Aksp.setUp()
Aksp.setFromOptions()
self.Aksp = Aksp
def transfer_form(F, newmesh, transfer=firedrake.prolong, replace_map={}):
"""
Given a form defined on some mesh, generate a new form with all
the same components, but transferred onto a different mesh.
:arg F: the form to be transferred
:arg newmesh: the mesh to transfer the form to
:kwarg transfer: the transfer operator to use
:kwarg replace_map: user-provided replace map
"""
f = 0
# We replace at least the coordinate map
replace_map = {
ufl.SpatialCoordinate(F.ufl_domain()): ufl.SpatialCoordinate(newmesh)
}
# Test and trial functions are also replaced
if len(F.arguments()) > 0:
Vold = F.arguments()[0].function_space()
Vnew = firedrake.FunctionSpace(newmesh, Vold.ufl_element())
replace_map[firedrake.TestFunction(Vold)] = firedrake.TestFunction(
Vnew)
replace_map[firedrake.TrialFunction(Vold)] = firedrake.TrialFunction(
Vnew)
# As well as any spatially varying coefficients
for c in F.coefficients():
if isinstance(c, firedrake.Function) and c not in replace_map:
replace_map[c] = firedrake.Function(
firedrake.FunctionSpace(newmesh, c.ufl_element()))
transfer(c, replace_map[c])
# The form is reconstructed by cell type
for cell_type, dX in zip(("cell", "exterior_facet", "interior_facet"),
(ufl.dx, ufl.ds, ufl.dS)):
for integral in F.integrals_by_type(cell_type):
differential = dX(integral.subdomain_id(), domain=newmesh)
f += ufl.replace(integral.integrand(), replace_map) * differential
return f
def solve(self, dt, h0, a, u, h_inflow=None):
r"""Propagate the thickness forward by one timestep
This function uses the implicit Euler timestepping scheme to avoid
the stability issues associated to using continuous finite elements
for advection-type equations. The implicit Euler scheme is stable
for any timestep; you do not need to ensure that the CFL condition
is satisfied in order to get an answer. Nonetheless, keeping the
timestep within the CFL bound is a good idea for accuracy.
Parameters
----------
dt : float
Timestep
h0 : firedrake.Function
Initial ice thickness
a : firedrake.Function
Sum of accumulation and melt rates
u : firedrake.Function
Ice velocity
h_inflow : firedrake.Function
Thickness of the upstream ice that advects into the domain
Returns
-------
h : firedrake.Function
Ice thickness at `t + dt`
"""
grad, ds = self.grad, self.ds
h_inflow = h_inflow if h_inflow is not None else h0
Q = h0.function_space()
h, φ = firedrake.TrialFunction(Q), firedrake.TestFunction(Q)
n = self.facet_normal(Q.mesh())
outflow = firedrake.max_value(inner(u, n), 0)
inflow = firedrake.min_value(inner(u, n), 0)
flux_cells = -h * inner(u, grad(φ)) * dx
flux_out = h * φ * outflow * ds
F = h * φ * dx + dt * (flux_cells + flux_out)
accumulation = a * φ * dx
flux_in = -h_inflow * φ * inflow * ds
A = h0 * φ * dx + dt * (accumulation + flux_in)
h = h0.copy(deepcopy=True)
solver_parameters = {'ksp_type': 'preonly', 'pc_type': 'lu'}
firedrake.solve(F == A, h, solver_parameters=solver_parameters)
return h
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 get_weak_form(self, V):
"""
mu is a spatially varying coefficient in the weak form
of the elasticity equations. The idea is to make the mesh stiff near
the boundary that is being deformed.
"""
if self.fixed_bids is not None and len(self.fixed_bids) > 0:
mu = self.get_mu(V)
else:
mu = fd.Constant(1.0)
u = fd.TrialFunction(V)
v = fd.TestFunction(V)
return mu * fd.inner(fd.sym(fd.grad(u)), fd.sym(fd.grad(v))) * fd.dx
def DG_inv_mass(self, DG):
"""
Inverse DG mass matrix
:arg DG: the DG space
:returns: A PETSc Mat.
"""
cache = self.caches[DG.ufl_element()]
key = DG.dim()
try:
return cache._DG_inv_mass[key]
except KeyError:
M = firedrake.assemble(firedrake.Tensor(firedrake.inner(firedrake.TrialFunction(DG),
firedrake.TestFunction(DG))*firedrake.dx).inv)
return cache._DG_inv_mass.setdefault(key, M.petscmat)
def solve(self, dt, h0, a, u, h_inflow=None):
r"""Propagate the thickness forward by one timestep
This function uses an implicit second-order Taylor-Galerkin (also
known as Lax-Wendroff) scheme to solve the conservative advection
equation for ice thickness.
Parameters
----------
dt : float
Timestep
h0 : firedrake.Function
Initial ice thickness
a : firedrake.Function
Sum of accumulation and melt rates
u : firedrake.Function
Ice velocity
h_inflow : firedrake.Function
Thickness of the upstream ice that advects into the domain
Returns
-------
h : firedrake.Function
Ice thickness at `t + dt`
"""
grad, div, ds = self.grad, self.div, self.ds
h_inflow = h_inflow if h_inflow is not None else h0
Q = h0.function_space()
h, φ = firedrake.TrialFunction(Q), firedrake.TestFunction(Q)
n = self.facet_normal(Q.mesh())
outflow = firedrake.max_value(inner(u, n), 0)
inflow = firedrake.min_value(inner(u, n), 0)
flux_cells = -h * inner(u, grad(φ)) * dx
flux_cells_lax = 0.5 * dt * div(h * u) * inner(u, grad(φ)) * dx
flux_out = (h - 0.5 * dt * div(h * u)) * φ * outflow * ds
F = h * φ * dx + dt * (flux_cells + flux_cells_lax + flux_out)
accumulation = a * φ * dx
flux_in = -(h_inflow - 0.5 * dt * div(h0 * u)) * φ * inflow * ds
A = h0 * φ * dx + dt * (accumulation + flux_in)
h = h0.copy(deepcopy=True)
solver_parameters = {'ksp_type': 'preonly', 'pc_type': 'lu'}
firedrake.solve(F == A, h, solver_parameters=solver_parameters)
return h
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 __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 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)
return self._V_DG_mass.setdefault(key, M.petscmat)
def __init__(self, problem, callback=(lambda s: None)):
self._setup(problem, callback)
self.update_state()
self.update_adjoint_state()
q, dJ = self.search_direction, self.gradient
Q = q.function_space()
M = firedrake.TrialFunction(Q) * firedrake.TestFunction(Q) * dx
problem = firedrake.LinearVariationalProblem(
M, -dJ, q, form_compiler_parameters=self._fc_params)
self._search_direction_solver = firedrake.LinearVariationalSolver(
problem, solver_parameters=self._solver_params)
self.update_search_direction()
self._callback(self)
def form2indicator(F):
"""
Multiply throughout in a form and
assemble as a cellwise error
indicator.
:arg F: the form
"""
mesh = F.ufl_domain()
P0 = firedrake.FunctionSpace(mesh, "DG", 0)
p0test = firedrake.TestFunction(P0)
indicator = firedrake.Function(P0)
# Contributions from surface integrals
flux_terms = 0
integrals = F.integrals_by_type("exterior_facet")
if len(integrals) > 0:
for integral in integrals:
ds = firedrake.ds(integral.subdomain_id())
flux_terms += p0test * integral.integrand() * ds
integrals = F.integrals_by_type("interior_facet")
if len(integrals) > 0:
for integral in integrals:
dS = firedrake.dS(integral.subdomain_id())
flux_terms += p0test("+") * integral.integrand() * dS
flux_terms += p0test("-") * integral.integrand() * dS
if flux_terms != 0:
dx = firedrake.dx
mass_term = firedrake.TrialFunction(P0) * p0test * dx
sp = {
"snes_type": "ksponly",
"ksp_type": "preonly",
"pc_type": "jacobi",
}
firedrake.solve(mass_term == flux_terms,
indicator,
solver_parameters=sp)
# Contributions from volume integrals
cell_terms = 0
integrals = F.integrals_by_type("cell")
if len(integrals) > 0:
for integral in integrals:
dx = firedrake.dx(integral.subdomain_id())
cell_terms += p0test * integral.integrand() * dx
indicator += firedrake.assemble(cell_terms)
return indicator
def get_mu(self, V):
W = fd.FunctionSpace(V.mesh(), "CG", 1)
bcs = []
if len(self.fixed_bids):
bcs.append(fd.DirichletBC(W, 1, self.fixed_bids))
if len(self.free_bids):
bcs.append(fd.DirichletBC(W, 10, self.free_bids))
if len(bcs) == 0:
bcs = None
u = fd.TrialFunction(W)
v = fd.TestFunction(W)
a = fd.inner(fd.grad(u), fd.grad(v)) * fd.dx
b = fd.inner(fd.Constant(0.), v) * fd.dx
mu = fd.Function(W)
fd.solve(a == b, mu, bcs=bcs)
return mu
def get_weak_form(self, V):
"""
mu is a spatially varying coefficient in the weak form
of the elasticity equations. The idea is to make the mesh stiff near
the boundary that is being deformed.
"""
a = self.inner.get_weak_form(V)
u = fd.TrialFunction(V)
v = fd.TestFunction(V)
(u1, u2) = fd.split(u)
(v1, v2) = fd.split(v)
a_cr = (fd.inner(
fd.grad(u1)[0] - fd.grad(u2)[1],
fd.grad(v1)[0] - fd.grad(v2)[1]) + fd.inner(
fd.grad(u1)[1] + fd.grad(u2)[0],
fd.grad(v1)[1] + fd.grad(v2)[0])) * self.mu_cr * fd.dx
return a + a_cr
def assemble_mass_matrix(space, norm_type="L2"):
"""
Assemble the ``norm_type`` mass matrix
associated with some finite element ``space``.
"""
trial = firedrake.TrialFunction(space)
test = firedrake.TestFunction(space)
if norm_type == "L2":
lhs = ufl.inner(trial, test) * ufl.dx
elif norm_type == "H1":
lhs = (
ufl.inner(trial, test) * ufl.dx
+ ufl.inner(ufl.grad(trial), ufl.grad(test)) * ufl.dx
)
else:
raise ValueError(f"Norm type {norm_type} not recognised.")
return firedrake.assemble(lhs).petscmat
def _diffuse(self, dt, E, h, q, q_bed, E_surface):
Q = E.function_space()
degree = Q.ufl_element().degree()[1]
φ, ψ = firedrake.TrialFunction(Q), firedrake.TestFunction(Q)
a = (h * φ * ψ + dt * α * φ.dx(2) * ψ.dx(2) / h) * dx \
+ degree**2 * dt * α * φ * ψ / h * ds_t
f = E * ψ * h * dx \
+ dt * q * ψ * h * dx \
+ dt * q_bed * ψ * ds_b \
+ degree**2 * dt * α * E_surface * ψ / h * ds_t
degree_E = E.ufl_element().degree()
degree = (3 * degree_E[0], 2 * degree_E[1])
form_compiler_parameters = {'quadrature_degree': degree}
firedrake.solve(a == f,
E,
form_compiler_parameters=form_compiler_parameters)
