def handle_exit_annotation():
yield
# Since importing firedrake_adjoint modifies a global variable, we need to
# pause annotations at the end of the module
annotate = annotate_tape()
if annotate:
pause_annotation()
def wrapper(*args, **kwargs):
"""When a form is assembled, the information about its nonlinear dependencies is lost,
and it is no longer easy to manipulate. Therefore, we decorate :func:`.assemble`
to *attach the form to the assembled object*. This lets the automatic annotation work,
even when the user calls the lower-level :py:data:`solve(A, x, b)`.
"""
ad_block_tag = kwargs.pop("ad_block_tag", None)
annotate = annotate_tape(kwargs)
with stop_annotating():
output = assemble(*args, **kwargs)
form = args[0]
if isinstance(output, numbers.Complex):
if not annotate:
return output
if not isinstance(output, float):
raise NotImplementedError(
"Taping for complex-valued 0-forms not yet done!")
output = create_overloaded_object(output)
block = AssembleBlock(form, ad_block_tag=ad_block_tag)
tape = get_working_tape()
tape.add_block(block)
block.add_output(output.block_variable)
else:
# Assembled a vector or matrix
output.form = form
return output
def wrapper(*args, **kwargs):
"""The project call performs an equation solve, and so it too must be annotated so that the
adjoint and tangent linear models may be constructed automatically by pyadjoint.
To disable the annotation of this function, just pass :py:data:`annotate=False`. This is useful in
cases where the solve is known to be irrelevant or diagnostic for the purposes of the adjoint
computation (such as projecting fields to other function spaces for the purposes of
visualisation)."""
annotate = annotate_tape(kwargs)
with stop_annotating():
output = project(*args, **kwargs)
output = create_overloaded_object(output)
if annotate:
bcs = kwargs.pop("bcs", [])
sb_kwargs = ProjectBlock.pop_kwargs(kwargs)
block = ProjectBlock(args[0], args[1], output, bcs, **sb_kwargs)
tape = get_working_tape()
tape.add_block(block)
block.add_output(output.block_variable)
return output
def wrapper(self, b, *args, **kwargs):
ad_block_tag = kwargs.pop("ad_block_tag", None)
annotate = annotate_tape(kwargs)
if annotate:
bcs = kwargs.get("bcs", [])
if isinstance(
b, firedrake.Function
) and b.ufl_domain() != self.function_space().mesh():
block = SupermeshProjectBlock(b,
self.function_space(),
self,
bcs,
ad_block_tag=ad_block_tag)
else:
block = ProjectBlock(b,
self.function_space(),
self,
bcs,
ad_block_tag=ad_block_tag)
tape = get_working_tape()
tape.add_block(block)
with stop_annotating():
output = project(self, b, *args, **kwargs)
if annotate:
block.add_output(output.create_block_variable())
return output
def wrapper(interpolator, *function, **kwargs):
"""To disable the annotation, just pass :py:data:`annotate=False` to this routine, and it acts exactly like the
Firedrake interpolate call."""
ad_block_tag = kwargs.pop("ad_block_tag", None)
annotate = annotate_tape(kwargs)
if annotate:
sb_kwargs = InterpolateBlock.pop_kwargs(kwargs)
sb_kwargs.update(kwargs)
block = InterpolateBlock(interpolator,
*function,
ad_block_tag=ad_block_tag,
**sb_kwargs)
tape = get_working_tape()
tape.add_block(block)
with stop_annotating():
output = interpolate(interpolator, *function, **kwargs)
if annotate:
from firedrake import Function
if isinstance(interpolator.V, Function):
block.add_output(output.create_block_variable())
else:
block.add_output(output.block_variable)
return output
def solve(self, *args, **kwargs):
annotate = annotate_tape(kwargs)
if annotate:
if len(args) == 3:
block_helper = LUSolveBlockHelper()
A = args[0]
x = args[1]
b = args[2]
elif len(args) == 2:
block_helper = self.block_helper
A = self.operator
x = args[0]
b = args[1]
u = x.function
parameters = self.parameters.copy()
tape = get_working_tape()
sb_kwargs = LUSolveBlock.pop_kwargs(kwargs)
block = LUSolveBlock(A,
x,
b,
lu_solver_parameters=parameters,
block_helper=block_helper,
lu_solver_method=self.method,
**sb_kwargs)
tape.add_block(block)
out = backend.LUSolver.solve(self, *args, **kwargs)
if annotate:
block.add_output(u.create_block_variable())
return out
def wrapper(self, **kwargs):
"""To disable the annotation, just pass :py:data:`annotate=False` to this routine, and it acts exactly like the
Firedrake solve call. This is useful in cases where the solve is known to be irrelevant or diagnostic
for the purposes of the adjoint computation (such as projecting fields to other function spaces
for the purposes of visualisation)."""
annotate = annotate_tape(kwargs)
if annotate:
tape = get_working_tape()
problem = self._ad_problem
sb_kwargs = NonlinearVariationalSolveBlock.pop_kwargs(kwargs)
sb_kwargs.update(kwargs)
block = NonlinearVariationalSolveBlock(
problem._ad_F == 0,
problem._ad_u,
problem._ad_bcs,
problem_J=problem._ad_J,
solver_params=self.parameters,
solver_kwargs=self._ad_kwargs,
**sb_kwargs)
tape.add_block(block)
with stop_annotating():
out = solve(self, **kwargs)
if annotate:
block.add_output(
self._ad_problem._ad_u.create_block_variable())
return out
def wrapper(*args, **kwargs):
ad_block_tag = kwargs.pop("ad_block_tag", None)
annotate = annotate_tape(kwargs)
if annotate:
tape = get_working_tape()
solve_block_type = SolveVarFormBlock
if not isinstance(args[0], ufl.equation.Equation):
solve_block_type = SolveLinearSystemBlock
sb_kwargs = solve_block_type.pop_kwargs(kwargs)
sb_kwargs.update(kwargs)
block = solve_block_type(*args,
ad_block_tag=ad_block_tag,
**sb_kwargs)
tape.add_block(block)
with stop_annotating():
output = solve(*args, **kwargs)
if annotate:
if hasattr(args[1], "create_block_variable"):
block_variable = args[1].create_block_variable()
else:
block_variable = args[1].function.create_block_variable()
block.add_output(block_variable)
return output
def __init__(self, *args, **kwargs):
annotate = annotate_tape(kwargs)
BaseExpression.__init__(self, *args, **kwargs, annotate=annotate)
backend.UserExpression.__init__(self, *args, **kwargs)
for k, v in kwargs.items():
self.save_attribute(k, v)
def solve(self, *args, **kwargs):
annotate = annotate_tape(kwargs)
if annotate:
tape = get_working_tape()
factory = args[0]
vec = args[1]
b = backend.as_backend_type(vec).__class__()
factory.F(b=b, x=vec)
F = b.form
bcs = b.bcs
u = vec.function
sb_kwargs = SolveVarFormBlock.pop_kwargs(kwargs)
block = SolveVarFormBlock(
F == 0,
u,
bcs,
solver_parameters={"newton_solver": self.parameters.copy()},
**sb_kwargs)
tape.add_block(block)
newargs = [self] + list(args)
out = backend.NewtonSolver.solve(*newargs, **kwargs)
if annotate:
block.add_output(u.create_block_variable())
return out
def assemble(*args, **kwargs):
"""When a form is assembled, the information about its nonlinear dependencies is lost,
and it is no longer easy to manipulate. Therefore, fenics_adjoint overloads the :py:func:`dolfin.assemble`
function to *attach the form to the assembled object*. This lets the automatic annotation work,
even when the user calls the lower-level :py:data:`solve(A, x, b)`.
"""
annotate = annotate_tape(kwargs)
with stop_annotating():
output = backend.assemble(*args, **kwargs)
form = args[0]
if isinstance(output, float):
output = create_overloaded_object(output)
if annotate:
block = AssembleBlock(form)
tape = get_working_tape()
tape.add_block(block)
block.add_output(output.block_variable)
else:
# Assembled a vector or matrix
output.form = form
return output
def assign(self, *args, **kwargs):
annotate = annotate_tape(kwargs)
outputs = Enlist(args[0])
inputs = Enlist(args[1])
if annotate:
for i, o in enumerate(outputs):
if not isinstance(o, OverloadedType):
outputs[i] = create_overloaded_object(o)
for j, i in enumerate(outputs):
if not isinstance(i, OverloadedType):
inputs[j] = create_overloaded_object(i)
block = FunctionAssignerBlock(self, inputs)
tape = get_working_tape()
tape.add_block(block)
with stop_annotating():
ret = backend.FunctionAssigner.assign(self, outputs.delist(), inputs.delist(), **kwargs)
if annotate:
for output in outputs:
block.add_output(output.block_variable)
return ret
def solve(self, **kwargs):
annotate = annotate_tape()
if annotate:
block_helper = BlockSolveBlockHelper()
tape = get_working_tape()
problem = self._ad_problem
# sb_kwargs = SolveBlock.pop_kwargs(kwargs)
block = NonlinearBlockSolveBlock(
problem._ad_b == 0,
problem._ad_u,
problem._ad_bcs,
block_helper=block_helper,
problem_J=problem._ad_A,
block_field=self._ad_problem.block_field,
block_split=self._ad_problem.block_split)
tape.add_block(block)
with stop_annotating():
out = super(NonlinearBlockSolver, self).solve()
if annotate:
block.add_output(self._ad_problem._ad_u.create_block_variable())
return out
def wrapper(*args, **kwargs):
"""The project call performs an equation solve, and so it too must be annotated so that the
adjoint and tangent linear models may be constructed automatically by pyadjoint.
To disable the annotation of this function, just pass :py:data:`annotate=False`. This is useful in
cases where the solve is known to be irrelevant or diagnostic for the purposes of the adjoint
computation (such as projecting fields to other function spaces for the purposes of
visualisation)."""
annotate = annotate_tape(kwargs)
if annotate:
bcs = kwargs.get("bcs", [])
sb_kwargs = ProjectBlock.pop_kwargs(kwargs)
if isinstance(args[1], function.Function):
# block should be created before project because output might also be an input that needs checkpointing
output = args[1]
V = output.function_space()
block = ProjectBlock(args[0], V, output, bcs, **sb_kwargs)
with stop_annotating():
output = project(*args, **kwargs)
if annotate:
tape = get_working_tape()
if not isinstance(args[1], function.Function):
block = ProjectBlock(args[0], args[1], output, bcs,
**sb_kwargs)
tape.add_block(block)
block.add_output(output.create_block_variable())
return output
def sub(self, i, deepcopy=False, **kwargs):
from .function_assigner import FunctionAssigner, FunctionAssignerBlock
annotate = annotate_tape(kwargs)
if deepcopy:
ret = create_overloaded_object(
backend.Function.sub(self, i, deepcopy, **kwargs))
if annotate:
fa = FunctionAssigner(ret.function_space(),
self.function_space())
block = FunctionAssignerBlock(fa, Enlist(self))
tape = get_working_tape()
tape.add_block(block)
block.add_output(ret.block_variable)
else:
extra_kwargs = {}
if annotate:
extra_kwargs = {
"block_class": FunctionSplitBlock,
"_ad_floating_active": True,
"_ad_args": [self, i],
"_ad_output_args": [i],
"output_block_class": FunctionMergeBlock,
"_ad_outputs": [self],
}
ret = compat.create_function(self, i, **extra_kwargs)
return ret
def update(self, x, flag, iteration):
"""Update domain and solution to state and adjoint equation."""
if self.Q.update_domain(x):
try:
# We use pyadjoint to calculate adjoint and shape derivatives,
# in order to do this we need to "record a tape of the forward
# solve", pyadjoint will then figure out all necessary
# adjoints.
import firedrake_adjoint as fda
tape = fda.get_working_tape()
tape.clear_tape()
# ensure we are annotating
from pyadjoint.tape import annotate_tape
safety_counter = 0
while not annotate_tape():
safety_counter += 1
fda.continue_annotation()
if safety_counter > 1e2:
import sys
sys.exit('Cannot annotate even after 100 attempts.')
mesh_m = self.J.Q.mesh_m
s = fd.Function(self.J.V_m)
mesh_m.coordinates.assign(mesh_m.coordinates + s)
self.s = s
self.c = fda.Control(s)
self.e.solve()
Jpyadj = fd.assemble(self.J.value_form())
self.Jred = fda.ReducedFunctional(Jpyadj, self.c)
fda.pause_annotation()
except fd.ConvergenceError:
if self.cb is not None:
self.cb()
raise
if iteration >= 0 and self.cb is not None:
self.cb()
def __init__(self, *args, **kwargs):
annotate = annotate_tape(kwargs)
BaseExpression.__init__(self, *args, **kwargs, annotate=annotate)
backend.Expression.__init__(self, *args, **kwargs)
for k, v in kwargs.items():
if k not in _IGNORED_EXPRESSION_ATTRIBUTES:
self.save_attribute(k, v)
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# stop any annotation that might be ongoing as we only need to record
# what happens in self.solvePDE()
from pyadjoint.tape import pause_annotation, annotate_tape
if annotate_tape():
pause_annotation()
def XDMFFile_read_checkpoint(self, *args, **kwargs):
annotate = annotate_tape(kwargs)
output = __XDMFFile_read_checkpoint__(self, *args, **kwargs)
if annotate:
func = args[0]
if isinstance(func, OverloadedType):
func.create_block_variable()
return output
def wrapper(self, *args, **kwargs):
annotate = annotate_tape(kwargs)
if annotate:
for arg in args:
if not hasattr(arg, "bcs"):
arg.bcs = []
arg.bcs.append(self)
with stop_annotating():
ret = apply(self, *args, **kwargs)
return ret
def wrapper(self, other, **kwargs):
annotate = annotate_tape(kwargs)
func = __idiv__(self, other, **kwargs)
if annotate:
block = FunctionAssignBlock(func, self / other)
tape = get_working_tape()
tape.add_block(block)
block.add_output(func.create_block_variable())
return func
def HDF5File_read(self, *args, **kwargs):
annotate = annotate_tape(kwargs)
output = __HDF5File_read__(self, *args, **kwargs)
if annotate:
func = args[0]
if isinstance(func, backend.Mesh):
func.org_mesh_coords = func.coordinates().copy()
if isinstance(func, OverloadedType):
func.create_block_variable()
return output
def move(mesh, vector, **kwargs):
annotate = annotate_tape(kwargs)
if annotate:
assert isinstance(mesh, OverloadedType)
assert isinstance(vector, OverloadedType)
tape = get_working_tape()
block = ALEMoveBlock(mesh, vector, **kwargs)
tape.add_block(block)
with stop_annotating():
output = __backend_ALE_move(mesh, vector)
if annotate:
block.add_output(mesh.create_block_variable())
return output
def wrapper(self, other, **kwargs):
ad_block_tag = kwargs.pop("ad_block_tag", None)
annotate = annotate_tape(kwargs)
func = __imul__(self, other, **kwargs)
if annotate:
block = FunctionAssignBlock(func,
self * other,
ad_block_tag=ad_block_tag)
tape = get_working_tape()
tape.add_block(block)
block.add_output(func.create_block_variable())
return func
def __getitem__(self, item):
annotate = annotate_tape()
if annotate:
block = NumpyArraySliceBlock(self, item)
tape = get_working_tape()
tape.add_block(block)
with stop_annotating():
out = numpy.ndarray.__getitem__(self, item)
if annotate:
out = create_overloaded_object(out)
block.add_output(out.create_block_variable())
return out
def solve(*args, **kwargs):
"""This solve routine wraps the real Dolfin solve call. Its purpose is to annotate the model,
recording what solves occur and what forms are involved, so that the adjoint and tangent linear models may be
constructed automatically by pyadjoint.
To disable the annotation, just pass :py:data:`annotate=False` to this routine, and it acts exactly like the
Dolfin solve call. This is useful in cases where the solve is known to be irrelevant or diagnostic
for the purposes of the adjoint computation (such as projecting fields to other function spaces
for the purposes of visualisation).
The overloaded solve takes optional callback functions to extract adjoint solutions.
All of the callback functions follow the same signature, taking a single argument of type Function.
Keyword Args:
adj_cb (function, optional): callback function supplying the adjoint solution in the interior.
The boundary values are zero.
adj_bdy_cb (function, optional): callback function supplying the adjoint solution on the boundary.
The interior values are not guaranteed to be zero.
adj2_cb (function, optional): callback function supplying the second-order adjoint solution in the interior.
The boundary values are zero.
adj2_bdy_cb (function, optional): callback function supplying the second-order adjoint solution on
the boundary. The interior values are not guaranteed to be zero.
"""
ad_block_tag = kwargs.pop("ad_block_tag", None)
annotate = annotate_tape(kwargs)
if annotate:
tape = get_working_tape()
solve_block_type = SolveVarFormBlock
if not isinstance(args[0], ufl.equation.Equation):
solve_block_type = SolveLinearSystemBlock
sb_kwargs = solve_block_type.pop_kwargs(kwargs)
sb_kwargs.update(kwargs)
block = solve_block_type(*args, ad_block_tag=ad_block_tag, **sb_kwargs)
tape.add_block(block)
with stop_annotating():
output = backend.solve(*args, **kwargs)
if annotate:
if hasattr(args[1], "create_block_variable"):
block_variable = args[1].create_block_variable()
else:
block_variable = args[1].function.create_block_variable()
block.add_output(block_variable)
return output
def wrapper(self, *args, **kwargs):
annotate = annotate_tape(kwargs)
func = copy(self, *args, **kwargs)
if annotate:
if kwargs.pop("deepcopy", False):
block = FunctionAssignBlock(func, self)
tape = get_working_tape()
tape.add_block(block)
block.add_output(func.create_block_variable())
else:
# TODO: Implement. Here we would need to use floating types.
raise NotImplementedError("Currently kwargs['deepcopy'] must be set True")
return func
def __init__(self, *args, **kwargs):
annotate = annotate_tape(kwargs)
self._ad_initialized = False
self._ad_attributes_dict = {}
self.ad_ignored_attributes = []
self.user_defined_derivatives = {}
FloatingType.__init__(self,
*args,
block_class=ExpressionBlock,
annotate=annotate,
_ad_args=[self],
_ad_floating_active=True,
**kwargs)
self._ad_initialized = True
self._cached_fs = {}
def __init__(self, *args, **kwargs):
annotate = annotate_tape(kwargs)
BaseExpression.__init__(self, *args, **kwargs, annotate=annotate)
kwargs_copy = kwargs.copy()
for k in kwargs:
v = kwargs[k]
if isinstance(v, OverloadedType):
if hasattr(v, "_cpp_object"):
kwargs[k] = v._cpp_object
backend.CompiledExpression.__init__(self, *args, **kwargs)
for k, val in kwargs_copy.items():
if hasattr(self._cpp_object, k):
self.save_attribute(k, val)
def project(self, b, *args, **kwargs):
annotate = annotate_tape(kwargs)
with stop_annotating():
output = super(Function, self).project(b, *args, **kwargs)
output = create_overloaded_object(output)
if annotate:
bcs = kwargs.pop("bcs", [])
block = ProjectBlock(b, self.function_space(), output, bcs)
tape = get_working_tape()
tape.add_block(block)
block.add_output(output.create_block_variable())
return output