"""
Driver functions for derivative recovery.
"""
import os
import firedrake
import ufl
from firedrake.__future__ import interpolate
from firedrake.petsc import PETSc
from pyop2 import op2
from .interpolation import clement_interpolant
from .math import construct_basis
from .quality import QualityMeasure, include_dir
__all__ = ["recover_gradient_l2", "recover_hessian_clement", "recover_boundary_hessian"]
def get_metric_kernel(func: str, dim: int) -> op2.Kernel:
"""
Helper function to easily pass Eigen kernels
for metric utilities to Firedrake via PyOP2.
:arg func: function name
:arg dim: spatial dimension
"""
pwd = os.path.abspath(os.path.join(os.path.dirname(__file__), "cxx"))
with open(os.path.join(pwd, f"metric{dim}d.cxx"), "r") as code:
return op2.Kernel(code.read(), func, cpp=True, include_dirs=include_dir)
[docs]
@PETSc.Log.EventDecorator()
def recover_gradient_l2(f, target_space=None):
r"""
Recover the gradient of a scalar or vector field using :math:`L^2` projection.
:arg f: the scalar field whose derivatives we seek to recover
:kwarg mesh: the underlying mesh
:kwarg target_space: the :func:`firedrake.functionspace.FunctionSpace`
recovered gradient should live in
"""
if target_space is None:
if not isinstance(f, firedrake.Function):
raise ValueError(
"If a target space is not provided then the input must be a Function."
)
degree = max(1, f.ufl_element().degree() - 1)
mesh = f.function_space().mesh()
rank = len(f.ufl_element().value_shape)
if rank == 0:
target_space = firedrake.VectorFunctionSpace(mesh, "CG", degree)
elif rank == 1:
target_space = firedrake.TensorFunctionSpace(mesh, "CG", degree)
else:
raise ValueError(
"L2 projection can only be used to compute gradients of scalar or"
f" vector Functions, not Functions of rank {rank}."
)
return firedrake.project(ufl.grad(f), target_space)
[docs]
@PETSc.Log.EventDecorator()
def recover_hessian_clement(f):
r"""
Recover the gradient and Hessian of a scalar field using two applications of
Clement interpolation.
Note that if the field is of degree 2 then projection will be used to obtain the
gradient. If the field is of degree 3 or greater then projection will be used
for the Hessian recovery, too.
:arg f: the scalar field whose derivatives we seek to recover
"""
if not isinstance(f, firedrake.Function):
raise ValueError(
"Clement interpolation can only be used to compute gradients of"
" Lagrange Functions of degree > 0."
)
family = f.ufl_element().family()
degree = f.ufl_element().degree()
if family not in ("Lagrange", "Discontinuous Lagrange") or degree == 0:
raise ValueError(
"Clement interpolation can only be used to compute gradients of"
" Lagrange Functions of degree > 0."
)
mesh = f.function_space().mesh()
# Recover gradient
if degree <= 1:
V = firedrake.VectorFunctionSpace(mesh, "DG", 0)
g = clement_interpolant(firedrake.project(ufl.grad(f), V))
else:
V = firedrake.VectorFunctionSpace(mesh, "DG", degree - 1)
g = recover_gradient_l2(f, target_space=V)
# Recover Hessian
if degree <= 2:
W = firedrake.TensorFunctionSpace(mesh, "DG", 0)
H = clement_interpolant(firedrake.project(ufl.grad(g), W))
else:
W = firedrake.TensorFunctionSpace(mesh, "DG", degree - 2)
H = recover_gradient_l2(g, target_space=W)
return g, H
[docs]
@PETSc.Log.EventDecorator()
def recover_boundary_hessian(f, method="Clement", target_space=None, **kwargs):
"""
Recover the Hessian of a scalar field on the domain boundary.
:arg f: field to recover over the domain boundary
:kwarg method: choose from 'mixed_L2' and 'Clement'
:kwarg target_space: the :func:`firedrake.functionspace.TensorFunctionSpace`
in which the metric will exist
"""
mesh = ufl.domain.extract_unique_domain(f)
d = mesh.topological_dimension()
assert d in (2, 3)
# Apply Gram-Schmidt to get tangent vectors
n = ufl.FacetNormal(mesh)
ns = construct_basis(n)
s = ns[1:]
ns = ufl.as_vector(ns)
# Setup
P1 = firedrake.FunctionSpace(mesh, "CG", 1)
P1_ten = target_space or firedrake.TensorFunctionSpace(mesh, "CG", 1)
assert P1_ten.ufl_element().family() == "Lagrange"
assert P1_ten.ufl_element().degree() == 1
boundary_tag = kwargs.get("boundary_tag", "on_boundary")
Hs = firedrake.TrialFunction(P1)
v = firedrake.TestFunction(P1)
l2_proj = [[firedrake.Function(P1) for i in range(d - 1)] for j in range(d - 1)]
h = firedrake.assemble(
interpolate(ufl.CellDiameter(mesh), firedrake.FunctionSpace(mesh, "DG", 0))
)
h = firedrake.Constant(1 / h.vector().gather().max() ** 2)
sp = {
"ksp_type": "gmres",
"ksp_gmres_restart": 20,
"pc_type": "ilu",
}
if method == "mixed_L2":
# Arbitrary value on domain interior
a = v * Hs * ufl.dx
L = v * h * ufl.dx
# Hessian on boundary
nullspace = firedrake.VectorSpaceBasis(constant=True)
for j, s1 in enumerate(s):
for i, s0 in enumerate(s):
bcs = []
for tag in mesh.exterior_facets.unique_markers:
a_bc = v * Hs * ufl.ds(tag)
L_bc = (
-ufl.dot(s0, ufl.grad(v))
* ufl.dot(s1, ufl.grad(f))
* ufl.ds(tag)
)
bcs.append(firedrake.EquationBC(a_bc == L_bc, l2_proj[i][j], tag))
firedrake.solve(
a == L,
l2_proj[i][j],
bcs=bcs,
nullspace=nullspace,
solver_parameters=sp,
)
elif method == "Clement":
P0_vec = firedrake.VectorFunctionSpace(mesh, "DG", 0)
P0_ten = firedrake.TensorFunctionSpace(mesh, "DG", 0)
P1_vec = firedrake.VectorFunctionSpace(mesh, "CG", 1)
H = firedrake.Function(P1_ten)
p0test = firedrake.TestFunction(P0_vec)
p1test = firedrake.TestFunction(P1)
fa = QualityMeasure(mesh, python=True)("facet_area")
source = firedrake.assemble(ufl.inner(p0test, ufl.grad(f)) / fa * ufl.ds)
# Recover gradient
c = clement_interpolant(source, boundary=True, target_space=P1_vec)
# Recover Hessian
H += clement_interpolant(
firedrake.assemble(interpolate(ufl.grad(c), P0_ten)),
boundary=True,
target_space=P1_ten,
)
# Compute tangential components
for j, s1 in enumerate(s):
for i, s0 in enumerate(s):
l2_proj[i][j] = firedrake.Function(P1)
l2_proj[i][j].dat.data_with_halos[:] = firedrake.assemble(
p1test * ufl.dot(ufl.dot(s0, H), s1) / fa * ufl.ds
).dat.data_with_halos
# TODO: Avoid accessing .dat.data_with_halos (#131)
else:
raise ValueError(
f"Recovery method '{method}' not supported for Hessians on the boundary."
)
# Construct tensor field
Hbar = firedrake.Function(P1_ten)
if d == 2:
Hsub = firedrake.assemble(interpolate(abs(l2_proj[0][0]), P1))
H = ufl.as_matrix([[h, 0], [0, Hsub]])
else:
fs = firedrake.TensorFunctionSpace(mesh, "CG", 1, shape=(2, 2))
Hsub = firedrake.Function(fs)
Hsub.interpolate(
ufl.as_matrix(
[[l2_proj[0][0], l2_proj[0][1]], [l2_proj[1][0], l2_proj[1][1]]]
)
)
# Enforce SPD
metric = firedrake.Function(fs)
op2.par_loop(
get_metric_kernel("metric_from_hessian", 2),
fs.node_set,
metric.dat(op2.RW),
Hsub.dat(op2.READ),
)
Hsub.assign(metric)
# TODO: Could this be supported using RiemannianMetric.enforce_spd? (#131)
# Construct Hessian
H = ufl.as_matrix(
[[h, 0, 0], [0, Hsub[0, 0], Hsub[0, 1]], [0, Hsub[1, 0], Hsub[1, 1]]]
)
# Arbitrary value on domain interior
sigma = firedrake.TrialFunction(P1_ten)
tau = firedrake.TestFunction(P1_ten)
a = ufl.inner(tau, sigma) * ufl.dx
L = ufl.inner(tau, h * ufl.Identity(d)) * ufl.dx
# Boundary values imposed as in [Loseille et al. 2011]
a_bc = ufl.inner(tau, sigma) * ufl.ds
L_bc = ufl.inner(tau, ufl.dot(ufl.transpose(ns), ufl.dot(H, ns))) * ufl.ds
bcs = firedrake.EquationBC(a_bc == L_bc, Hbar, boundary_tag)
firedrake.solve(a == L, Hbar, bcs=bcs, solver_parameters=sp)
return Hbar