def __call__(self, infunc, annotate=None):
# Perform the forward operation
out = Function(infunc.function_space())
out.vector()[:] = out.vector().get_local()**2
# Annotate the operation on the dolfin-adjoint tape
if utils.to_annotate(annotate):
rhs = CustomDolfinAdjointFunctionRhs(coeffs=[infunc])
out_dep = adjglobals.adj_variables.next(out)
solving.register_initial_conditions(zip(rhs.coefficients(),
rhs.dependencies()),
linear=False)
if parameters["adjoint"]["record_all"]:
adjglobals.adjointer.record_variable(
out_dep,
libadjoint.MemoryStorage(adjlinalg.Vector(infunc)))
identity = utils.get_identity_block(out.function_space())
eq = libadjoint.Equation(out_dep,
blocks=[identity],
targets=[out_dep],
rhs=rhs)
cs = adjglobals.adjointer.register_equation(eq)
solving.do_checkpoint(cs, out_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)
out = utils.function_to_da_function(out)
if name is not None:
out.adj_name = name
to_annotate = utils.to_annotate(annotate)
if to_annotate:
if isinstance(v, backend.Function):
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)
elif annotate is not None:
outvar = adjglobals.adj_variables[out]
solving.register_initial_conditions([
[out, outvar],
], linear=True)
return out
def annotate_slope_limit(f):
# First annotate the equation
adj_var = adjglobals.adj_variables[f]
rhs = SlopeRHS(f)
adj_var_next = adjglobals.adj_variables.next(f)
identity_block = solving.get_identity(f.function_space())
eq = libadjoint.Equation(adj_var_next,
blocks=[identity_block],
targets=[adj_var_next],
rhs=rhs)
cs = adjglobals.adjointer.register_equation(eq)
# Record the result
adjglobals.adjointer.record_variable(
adjglobals.adj_variables[f],
libadjoint.MemoryStorage(adjlinalg.Vector(f)))
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 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 annotate_split(bigfn, idx, smallfn, bcs):
fn_space = smallfn.function_space().collapse()
test = backend.TestFunction(fn_space)
trial = backend.TrialFunction(fn_space)
eq_lhs = backend.inner(test, trial) * backend.dx
key = "{}split{}{}{}{}".format(eq_lhs, smallfn, bigfn, idx,
random.random()).encode('utf8')
diag_name = "Split:%s:" % idx + hashlib.md5(key).hexdigest()
diag_deps = []
diag_block = libadjoint.Block(
diag_name,
dependencies=diag_deps,
test_hermitian=backend.parameters["adjoint"]["test_hermitian"],
test_derivative=backend.parameters["adjoint"]["test_derivative"])
solving.register_initial_conditions(
[(bigfn, adjglobals.adj_variables[bigfn])], linear=True, var=None)
var = adjglobals.adj_variables.next(smallfn)
frozen_expressions_dict = expressions.freeze_dict()
def diag_assembly_cb(dependencies, values, hermitian, coefficient,
context):
'''This callback must conform to the libadjoint Python block assembly
interface. It returns either the form or its transpose, depending on
the value of the logical hermitian.'''
assert coefficient == 1
expressions.update_expressions(frozen_expressions_dict)
value_coeffs = [v.data for v in values]
eq_l = eq_lhs
if hermitian:
adjoint_bcs = [
utils.homogenize(bc)
for bc in bcs if isinstance(bc, backend.DirichletBC)
] + [bc for bc in bcs if not isinstance(bc, backend.DirichletBC)]
if len(adjoint_bcs) == 0: adjoint_bcs = None
return (adjlinalg.Matrix(backend.adjoint(eq_l), bcs=adjoint_bcs),
adjlinalg.Vector(None, fn_space=fn_space))
else:
return (adjlinalg.Matrix(eq_l, bcs=bcs),
adjlinalg.Vector(None, fn_space=fn_space))
diag_block.assemble = diag_assembly_cb
rhs = SplitRHS(test, bigfn, idx)
eqn = libadjoint.Equation(var, blocks=[diag_block], targets=[var], rhs=rhs)
cs = adjglobals.adjointer.register_equation(eqn)
solving.do_checkpoint(cs, var, rhs)
if backend.parameters["adjoint"]["fussy_replay"]:
mass = eq_lhs
smallfn_massed = backend.Function(fn_space)
backend.solve(mass == backend.action(mass, smallfn), smallfn_massed)
assert False, "No idea how to assign to a subfunction yet .. "
#assignment.dolfin_assign(bigfn, smallfn_massed)
if backend.parameters["adjoint"]["record_all"]:
smallfn_record = backend.Function(fn_space)
assignment.dolfin_assign(smallfn_record, smallfn)
adjglobals.adjointer.record_variable(
var, libadjoint.MemoryStorage(adjlinalg.Vector(smallfn_record)))
def solve(self, *args, **kwargs):
timer = backend.Timer("Matrix-free solver")
annotate = True
if "annotate" in kwargs:
annotate = kwargs["annotate"]
del kwargs["annotate"]
if len(args) == 3:
A = args[0]
x = args[1]
b = args[2]
elif len(args) == 2:
A = self.operators[0]
x = args[0]
b = args[1]
if annotate:
if not isinstance(A, AdjointKrylovMatrix):
try:
A = AdjointKrylovMatrix(A.form)
except AttributeError:
raise libadjoint.exceptions.LibadjointErrorInvalidInputs(
"Your A has to either be an AdjointKrylovMatrix or have been assembled after backend_adjoint was imported."
)
if not hasattr(x, 'function'):
raise libadjoint.exceptions.LibadjointErrorInvalidInputs(
"Your x has to come from code like down_cast(my_function.vector())."
)
if not hasattr(b, 'form'):
raise libadjoint.exceptions.LibadjointErrorInvalidInputs(
"Your b has to have the .form attribute: was it assembled with from backend_adjoint import *?"
)
if not hasattr(A, 'dependencies'):
backend.info_red(
"A has no .dependencies method; assuming no nonlinear dependencies of the matrix-free operator."
)
coeffs = []
dependencies = []
else:
coeffs = [
coeff for coeff in A.dependencies()
if hasattr(coeff, 'function_space')
]
dependencies = [
adjglobals.adj_variables[coeff] for coeff in coeffs
]
if len(dependencies) > 0:
assert hasattr(
A, "set_dependencies"
), "Need a set_dependencies method to replace your values, if you have nonlinear dependencies ... "
rhs = adjrhs.RHS(b.form)
key = '{}{}'.format(hash(A), random.random()).encode('utf8')
diag_name = hashlib.md5(key).hexdigest()
diag_block = libadjoint.Block(
diag_name,
dependencies=dependencies,
test_hermitian=backend.parameters["adjoint"]["test_hermitian"],
test_derivative=backend.parameters["adjoint"]
["test_derivative"])
solving.register_initial_conditions(
zip(rhs.coefficients(), rhs.dependencies()) +
zip(coeffs, dependencies),
linear=False,
var=None)
var = adjglobals.adj_variables.next(x.function)
frozen_expressions_dict = expressions.freeze_dict()
frozen_parameters = self.parameters.to_dict()
def diag_assembly_cb(dependencies, values, hermitian, coefficient,
context):
'''This callback must conform to the libadjoint Python block assembly
interface. It returns either the form or its transpose, depending on
the value of the logical hermitian.'''
assert coefficient == 1
expressions.update_expressions(frozen_expressions_dict)
if len(dependencies) > 0:
A.set_dependencies(dependencies,
[val.data for val in values])
if hermitian:
A_transpose = A.hermitian()
return (MatrixFree(
A_transpose,
fn_space=x.function.function_space(),
bcs=A_transpose.bcs,
solver_parameters=self.solver_parameters,
operators=transpose_operators(self.operators),
parameters=frozen_parameters),
adjlinalg.Vector(
None, fn_space=x.function.function_space()))
else:
return (MatrixFree(
A,
fn_space=x.function.function_space(),
bcs=b.bcs,
solver_parameters=self.solver_parameters,
operators=self.operators,
parameters=frozen_parameters),
adjlinalg.Vector(
None, fn_space=x.function.function_space()))
diag_block.assemble = diag_assembly_cb
def diag_action_cb(dependencies, values, hermitian, coefficient,
input, context):
expressions.update_expressions(frozen_expressions_dict)
A.set_dependencies(dependencies, [val.data for val in values])
if hermitian:
acting_mat = A.transpose()
else:
acting_mat = A
output_fn = backend.Function(input.data.function_space())
acting_mat.mult(input.data.vector(), output_fn.vector())
vec = output_fn.vector()
for i in range(len(vec)):
vec[i] = coefficient * vec[i]
return adjlinalg.Vector(output_fn)
diag_block.action = diag_action_cb
if len(dependencies) > 0:
def derivative_action(dependencies, values, variable,
contraction_vector, hermitian, input,
coefficient, context):
expressions.update_expressions(frozen_expressions_dict)
A.set_dependencies(dependencies,
[val.data for val in values])
action = A.derivative_action(
values[dependencies.index(variable)].data,
contraction_vector.data, hermitian, input.data,
coefficient)
return adjlinalg.Vector(action)
diag_block.derivative_action = derivative_action
eqn = libadjoint.Equation(var,
blocks=[diag_block],
targets=[var],
rhs=rhs)
cs = adjglobals.adjointer.register_equation(eqn)
solving.do_checkpoint(cs, var, rhs)
out = backend.PETScKrylovSolver.solve(self, *args)
if annotate:
if backend.parameters["adjoint"]["record_all"]:
adjglobals.adjointer.record_variable(
var,
libadjoint.MemoryStorage(adjlinalg.Vector(x.function)))
timer.stop()
return out
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