def __init__(self, source, target_space, target, bcs=[], **kwargs):
super(SupermeshProjectBlock,
self).__init__(ad_block_tag=kwargs.pop("ad_block_tag", None))
import firedrake.supermeshing as supermesh
# Process args and kwargs
if not isinstance(source, self.backend.Function):
raise NotImplementedError(
f"Source function must be a Function, not {type(source)}.")
if bcs != []:
raise NotImplementedError(
"Boundary conditions not yet considered.")
# Store spaces
mesh = kwargs.pop("mesh", None)
if mesh is None:
mesh = target_space.mesh()
self.source_space = source.function_space()
self.target_space = target_space
self.projector = firedrake.Projector(source, target_space, **kwargs)
# Assemble mixed mass matrix
self.mixed_mass = supermesh.assemble_mixed_mass_matrix(
source.function_space(), target_space)
# Add dependencies
self.add_dependency(source, no_duplicates=True)
for bc in bcs:
self.add_dependency(bc, no_duplicates=True)
def adjoint_supermesh_project(tgt_b,
src_b,
mixed_mass_matrix=None,
solver=None):
"""
Hand-coded adjoint of a supermesh projection.
:arg tgt_b: seed vector in target space.
:arg src_b: adjoint supermesh projection into source space.
"""
source_space = src_b.function_space()
target_space = tgt_b.function_space()
# Adjoint of step 2: Solve the linear system for the target:
# Mt^T * sol = tgt_b
ksp = solver or PETSc.KSP().create()
if solver is None:
Mt = assemble(
inner(TrialFunction(target_space), TestFunction(target_space)) *
dx).M.handle
ksp.setOperators(Mt.transpose())
ksp.setFromOptions()
with tgt_b.dat.vec_ro as rhs:
sol = rhs.copy()
ksp.solve(rhs, sol)
# Adjoint of 1: Multiply with the tranpose mass matrix
# src_b := Mst^T * sol
Mst = mixed_mass_matrix or supermesh.assemble_mixed_mass_matrix(
source_space, target_space)
with src_b.dat.vec_ro as vs:
Mst.multTranspose(sol, vs)
return src_b
def mesh2mesh_project_adjoint(target_b, source_b, **kwargs):
"""
Apply the adjoint of a mesh-to-mesh conservative
projection to some seed ``target_b``, mapping to
``source_b``.
The notation used here is in terms of the adjoint of
``mesh2mesh_project``. However, this function may also
be interpreted as a projector in its own right,
mapping ``target_b`` to ``source_b``.
Extra keyword arguments are passed to Firedrake's
``project`` function.
:arg target_b: seed :class:`Function` from the target
space of the forward projection
:arg source_b: the :class:`Function` from the source
space of the forward projection
"""
from firedrake.supermeshing import assemble_mixed_mass_matrix
target_space = target_b.function_space()
assert isinstance(source_b, firedrake.Function)
source_space = source_b.function_space()
# Get subspaces
if source_space == target_space:
source_b.assign(target_b)
return source_b
elif hasattr(source_space, "num_sub_spaces"):
if not hasattr(target_space, "num_sub_spaces"):
raise ValueError(f"Incompatible spaces {source_space} and {target_space}")
if not target_space.num_sub_spaces() == source_space.num_sub_spaces():
raise ValueError(f"Incompatible spaces {source_space} and {target_space}")
target_b_split = target_b.split()
source_b_split = source_b.split()
elif hasattr(target_space, "num_sub_spaces"):
raise ValueError(f"Incompatible spaces {source_space} and {target_space}")
else:
target_b_split = [target_b]
source_b_split = [source_b]
# Apply adjoint projection operator to each component
for i, (t_b, s_b) in enumerate(zip(target_b_split, source_b_split)):
ksp = petsc4py.KSP().create()
ksp.setOperators(assemble_mass_matrix(t_b.function_space()))
mixed_mass = assemble_mixed_mass_matrix(
t_b.function_space(), s_b.function_space()
)
with t_b.dat.vec_ro as tb, s_b.dat.vec_wo as sb:
residual = tb.copy()
ksp.solveTranspose(tb, residual)
mixed_mass.mult(residual, sb) # NOTE: mixed mass already transposed
return source_b
def test_supermesh_projection(plot=False):
"""
Ping-pong test for hand-coded supermesh projection.
"""
# Setup operators
s_init, s, t = setup()
Vs, Vt = s.function_space(), t.function_space()
M_st = supermesh.assemble_mixed_mass_matrix(Vs, Vt)
M_ts = supermesh.assemble_mixed_mass_matrix(Vt, Vs)
M_s = assemble(inner(TrialFunction(Vs), TestFunction(Vs)) * dx).M.handle
M_t = assemble(inner(TrialFunction(Vt), TestFunction(Vt)) * dx).M.handle
solver_s = PETSc.KSP().create()
solver_s.setOperators(M_s)
solver_s.setFromOptions()
solver_t = PETSc.KSP().create()
solver_t.setOperators(M_t)
solver_t.setFromOptions()
M_ts = M_st.copy()
M_ts.transpose()
# Ping pong test
N = 100
mass_init = assemble(s_init * dx)
l2_norm_init = norm(s_init)
l2_error = []
mass_error = []
s.assign(s_init)
for i in range(N):
supermesh_project(s, t, mixed_mass_matrix=M_st, solver=solver_t)
supermesh_project(t, s, mixed_mass_matrix=M_ts, solver=solver_s)
l2_error.append(errornorm(s, s_init) / l2_norm_init)
mass_error.append(abs(assemble(s * dx) - mass_init) / abs(mass_init))
assert np.allclose(mass_error, 0.0, atol=1.0e-04)
if plot:
plot_error(l2_error, "l2", interp="projection")
plot_error(mass_error, "mass", interp="projection")
levels = np.linspace(-0.15, 1.05, 7)
plot_field(s_init, "source_for_projection", levels, levels)
plot_field(s, "after_100_projections", levels, levels)
def mixed_mass(self, V_A, V_B):
"""
Compute the mixed mass matrix of two function spaces.
:arg V_A: the donor space
:arg V_B: the target space
:returns: A PETSc Mat.
"""
from firedrake.supermeshing import assemble_mixed_mass_matrix
key = (V_A.dim(), V_B.dim())
try:
return self._mixed_mass[key]
except KeyError:
M = assemble_mixed_mass_matrix(V_A, V_B)
return self._mixed_mass.setdefault(key, M)
def supermesh_project(src,
tgt,
check_mass=False,
mixed_mass_matrix=None,
solver=None):
"""
Hand-coded supermesh projection.
:arg src: source field.
:arg tgt: target field.
"""
source_space = src.function_space()
target_space = tgt.function_space()
# Step 1: Form the RHS:
# rhs := Mst * src
Mst = mixed_mass_matrix or supermesh.assemble_mixed_mass_matrix(
source_space, target_space)
with tgt.dat.vec_ro as vt:
rhs = vt.copy()
with src.dat.vec_ro as vs:
Mst.mult(vs, rhs)
# Step 2: Solve the linear system for the target:
# Mt * tgt = rhs
ksp = solver or PETSc.KSP().create()
if solver is None:
Mt = assemble(
inner(TrialFunction(target_space), TestFunction(target_space)) *
dx).M.handle
ksp.setOperators(Mt)
ksp.setFromOptions()
with tgt.dat.vec as vt:
ksp.solve(rhs, vt)
if check_mass:
assert np.allclose(assemble(src * dx), assemble(tgt * dx))
return tgt
def mixed_mass(self):
from firedrake.supermeshing import assemble_mixed_mass_matrix
return assemble_mixed_mass_matrix(self.source.function_space(),
self.target.function_space())
def mixed_mass(self):
from firedrake.supermeshing import assemble_mixed_mass_matrix
return assemble_mixed_mass_matrix(self.source.function_space(),
self.target.function_space())