def assign(self, receiving, giving, annotate=None):
out = backend.FunctionAssigner.assign(self, receiving, giving)
to_annotate = utils.to_annotate(annotate)
if to_annotate:
# Receiving is always a single Function, or a single Function.sub(foo).sub(bar)
# If the user passes in v.sub(0) to this function, we need to get a handle on v;
# fetch that now
receiving_super = get_super_function(receiving)
receiving_fnspace = receiving_super.function_space()
receiving_identity = utils.get_identity_block(receiving_fnspace)
receiving_idx = get_super_idx(receiving)
rhs = FunctionAssignerRHS(self, self.adj_function_assigner, receiving_super, receiving_idx, giving)
receiving_dep = adjglobals.adj_variables.next(receiving_super)
solving.register_initial_conditions(zip(rhs.coefficients(),rhs.dependencies()), linear=True)
if backend.parameters["adjoint"]["record_all"]:
adjglobals.adjointer.record_variable(receiving_dep, libadjoint.MemoryStorage(adjlinalg.Vector(receiving_super)))
eq = libadjoint.Equation(receiving_dep, blocks=[receiving_identity], targets=[receiving_dep], rhs=rhs)
cs = adjglobals.adjointer.register_equation(eq)
solving.do_checkpoint(cs, receiving_dep, rhs)
return out
def step(self, dt, annotate=None):
to_annotate = utils.to_annotate(annotate)
if to_annotate:
scheme = self.scheme()
var = scheme.solution()
fn_space = var.function_space()
current_var = adjglobals.adj_variables[var]
if not adjglobals.adjointer.variable_known(current_var):
solving.register_initial_conditions([(var, current_var)], linear=True)
identity_block = utils.get_identity_block(fn_space)
frozen_expressions = expressions.freeze_dict()
frozen_constants = constant.freeze_dict()
rhs = PointIntegralRHS(self, dt, current_var, frozen_expressions, frozen_constants)
next_var = adjglobals.adj_variables.next(var)
eqn = libadjoint.Equation(next_var, blocks=[identity_block], targets=[next_var], rhs=rhs)
cs = adjglobals.adjointer.register_equation(eqn)
super(PointIntegralSolver, self).step(dt)
if to_annotate:
curtime = float(scheme.t())
scheme.t().assign(curtime) # so that d-a sees the time update, which is implict in step
solving.do_checkpoint(cs, next_var, rhs)
if dolfin.parameters["adjoint"]["record_all"]:
adjglobals.adjointer.record_variable(next_var, libadjoint.MemoryStorage(adjlinalg.Vector(var)))
def interpolate(v, V, annotate=None, name=None):
'''The interpolate call changes Function data, and so it too must be annotated so that the
adjoint and tangent linear models may be constructed automatically by libadjoint.
To disable the annotation of this function, just pass :py:data:`annotate=False`. This is useful in
cases where the interpolation is known to be irrelevant or diagnostic for the purposes of the adjoint
computation (such as interpolating fields to other function spaces for the purposes of
visualisation).'''
out = backend.interpolate(v, V)
if name is not None:
out.adj_name = name
to_annotate = utils.to_annotate(annotate)
if isinstance(v, backend.Function) and to_annotate:
rhsdep = adjglobals.adj_variables[v]
if adjglobals.adjointer.variable_known(rhsdep):
rhs = InterpolateRHS(v, V)
identity_block = utils.get_identity_block(V)
solving.register_initial_conditions(zip(rhs.coefficients(),rhs.dependencies()), linear=True)
dep = adjglobals.adj_variables.next(out)
if backend.parameters["adjoint"]["record_all"]:
adjglobals.adjointer.record_variable(dep, libadjoint.MemoryStorage(adjlinalg.Vector(out)))
initial_eq = libadjoint.Equation(dep, blocks=[identity_block], targets=[dep], rhs=rhs)
cs = adjglobals.adjointer.register_equation(initial_eq)
solving.do_checkpoint(cs, dep, rhs)
return out
def register_initial_condition(coeff, dep):
fn_space = coeff.function_space()
identity_block = utils.get_identity_block(fn_space)
if backend.parameters["adjoint"]["record_all"]:
adjglobals.adjointer.record_variable(dep, libadjoint.MemoryStorage(adjlinalg.Vector(coeff)))
init_rhs=adjlinalg.Vector(coeff).duplicate()
init_rhs.axpy(1.0,adjlinalg.Vector(coeff))
rhs = adjrhs.RHS(init_rhs)
initial_eq = libadjoint.Equation(dep, blocks=[identity_block], targets=[dep], rhs=rhs)
cs = adjglobals.adjointer.register_equation(initial_eq)
assert adjglobals.adjointer.variable_known(dep)
do_checkpoint(cs, dep, rhs)
def register_assign(new, old, op=None):
if not isinstance(old, backend.Function):
assert op is not None
fn_space = new.function_space()
identity_block = utils.get_identity_block(fn_space)
dep = adjglobals.adj_variables.next(new)
if backend.parameters["adjoint"]["record_all"] and isinstance(old, backend.Function):
adjglobals.adjointer.record_variable(dep, libadjoint.MemoryStorage(adjlinalg.Vector(old)))
rhs = IdentityRHS(old, fn_space, op)
register_initial_conditions(zip(rhs.coefficients(),rhs.dependencies()), linear=True)
initial_eq = libadjoint.Equation(dep, blocks=[identity_block], targets=[dep], rhs=rhs)
cs = adjglobals.adjointer.register_equation(initial_eq)
do_checkpoint(cs, dep, rhs)
def dolfin_adjoint_assign(self, other, annotate=None, *args, **kwargs):
'''We also need to monkeypatch the Function.assign method, as it is often used inside
the main time loop, and not annotating it means you get the adjoint wrong for totally
nonobvious reasons. If anyone objects to me monkeypatching your objects, my apologies
in advance.'''
if self is other:
return
to_annotate = utils.to_annotate(annotate)
# if we shouldn't annotate, just assign
if not to_annotate:
return dolfin_assign(self, other, *args, **kwargs)
if isinstance(other, ufl.algebra.Sum) or isinstance(other, ufl.algebra.Product):
if backend.__name__ != 'dolfin':
errmsg = '''Cannot use Function.assign(linear combination of other Functions) yet.'''
raise libadjoint.exceptions.LibadjointErrorNotImplemented(errmsg)
else:
lincom = _check_and_contract_linear_comb(other, self)
else:
lincom = [(other, 1.0)]
# ignore anything not a backend.Function, unless the user insists
if not isinstance(other, backend.Function) and (annotate is not True):
return dolfin_assign(self, other, *args, **kwargs)
# ignore anything that is an interpolation, rather than a straight assignment
if hasattr(self, "function_space") and hasattr(other, "function_space"):
if str(self.function_space()) != str(other.function_space()):
return dolfin_assign(self, other, *args, **kwargs)
functions, weights = zip(*lincom)
self_var = adjglobals.adj_variables[self]
function_vars = [adjglobals.adj_variables[function] for function in functions]
# ignore any functions we haven't seen before -- we DON'T want to
# annotate the assignment of initial conditions here. That happens
# in the main solve wrapper.
for function_var in function_vars:
if not adjglobals.adjointer.variable_known(function_var) and not adjglobals.adjointer.variable_known(self_var) and (annotate is not True):
[adjglobals.adj_variables.forget(function) for function in functions]
adjglobals.adj_variables.forget(self)
return dolfin_assign(self, other, *args, **kwargs)
# OK, so we have a variable we've seen before. Beautiful.
if not adjglobals.adjointer.variable_known(self_var):
adjglobals.adj_variables.forget(self)
out = dolfin_assign(self, other, *args, **kwargs)
fn_space = self.function_space()
identity_block = utils.get_identity_block(fn_space)
dep = adjglobals.adj_variables.next(self)
if backend.parameters["adjoint"]["record_all"]:
adjglobals.adjointer.record_variable(dep, libadjoint.MemoryStorage(adjlinalg.Vector(self)))
rhs = LinComRHS(functions, weights, fn_space)
register_initial_conditions(zip(rhs.coefficients(),rhs.dependencies()), linear=True)
initial_eq = libadjoint.Equation(dep, blocks=[identity_block], targets=[dep], rhs=rhs)
cs = adjglobals.adjointer.register_equation(initial_eq)
do_checkpoint(cs, dep, rhs)
return out