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
value_coeffs = [v.data for v in values]
expressions.update_expressions(frozen_expressions)
constant.update_constants(frozen_constants)
eq_l = backend.replace(eq_lhs, dict(zip(diag_coeffs, value_coeffs)))
kwargs = {"cache": eq_l in caching.assembled_fwd_forms} # should we cache our matrices on the way backwards?
if hermitian:
# Homogenise the adjoint boundary conditions. This creates the adjoint
# solution associated with the lifted discrete system that is actually solved.
adjoint_bcs = [utils.homogenize(bc) for bc in eq_bcs if isinstance(bc, backend.DirichletBC)] + [
bc for bc in eq_bcs if not isinstance(bc, backend.DirichletBC)
]
if len(adjoint_bcs) == 0:
adjoint_bcs = None
else:
adjoint_bcs = misc.uniq(adjoint_bcs)
kwargs["bcs"] = adjoint_bcs
kwargs["solver_parameters"] = solver_parameters
kwargs["adjoint"] = True
if initial_guess:
kwargs["initial_guess"] = value_coeffs[dependencies.index(initial_guess_var)]
if replace_map:
kwargs["replace_map"] = dict(zip(diag_coeffs, value_coeffs))
return (
matrix_class(
backend.adjoint(eq_l, reordered_arguments=ufl.algorithms.extract_arguments(eq_l)), **kwargs
),
adjlinalg.Vector(None, fn_space=u.function_space()),
)
else:
kwargs["bcs"] = misc.uniq(eq_bcs)
kwargs["solver_parameters"] = solver_parameters
kwargs["adjoint"] = False
if initial_guess:
kwargs["initial_guess"] = value_coeffs[dependencies.index(initial_guess_var)]
if replace_map:
kwargs["replace_map"] = dict(zip(diag_coeffs, value_coeffs))
return (matrix_class(eq_l, **kwargs), adjlinalg.Vector(None, fn_space=u.function_space()))
def axpy(self, alpha, x):
assert isinstance(x.data, ufl.Form)
assert isinstance(self.data, ufl.Form)
# Let's do a bit of argument shuffling, shall we?
xargs = ufl.algorithms.extract_arguments(x.data)
sargs = ufl.algorithms.extract_arguments(self.data)
if xargs != sargs:
# OK, let's check that all of the function spaces are happy and so on.
for i in range(len(xargs)):
assert xargs[i].element() == sargs[i].element()
assert xargs[i].function_space() == sargs[i].function_space()
# Now that we are happy, let's replace the xargs with the sargs ones.
x_form = backend.replace(x.data, dict(zip(xargs, sargs)))
else:
x_form = x.data
self.data+=alpha*x_form
self.bcs += x.bcs # Err, I hope they are compatible ...
self.bcs = misc.uniq(self.bcs)
def solve(self, *args, **kwargs):
'''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).'''
to_annotate = utils.to_annotate(kwargs.pop("annotate", None))
if to_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]
bcs = []
if hasattr(A, 'bcs'):
bcs += A.bcs
if hasattr(b, 'bcs'):
bcs += b.bcs
bcs = misc.uniq(bcs)
assemble_system = A.assemble_system
A = A.form
u = x.function
b = b.form
if self.operators[1] is not None:
P = self.operators[1].form
else:
P = None
solver_parameters = self.solver_parameters
parameters = self.parameters.to_dict()
fn_space = u.function_space()
has_preconditioner = P is not None
nsp = self.nsp
tnsp = self.tnsp
if nsp is not None:
msg = """
The transpose nullspace is not set.
The nullspace of the PETScKrylovSolver is set. In this case,
the transpose nullspace must also be set, use:
solver.set_transpose_nullspace(nullspace)
"""
assert tnsp is not None, msg
if self.__global_list_idx__ is None:
self.__global_list_idx__ = len(petsc_krylov_solvers)
petsc_krylov_solvers.append(self)
adj_petsc_krylov_solvers.append(None)
idx = self.__global_list_idx__
class PETScKrylovSolverMatrix(adjlinalg.Matrix):
def __init__(selfmat, *args, **kwargs):
if 'initial_guess' in kwargs:
selfmat.initial_guess = kwargs['initial_guess']
del kwargs['initial_guess']
else:
selfmat.initial_guess = None
replace_map = kwargs['replace_map']
del kwargs['replace_map']
adjlinalg.Matrix.__init__(selfmat, *args, **kwargs)
selfmat.adjoint = kwargs['adjoint']
if P is None:
selfmat.operators = (dolfin.replace(A, replace_map), None)
else:
selfmat.operators = (dolfin.replace(A, replace_map), dolfin.replace(P, replace_map))
def axpy(selfmat, alpha, x):
raise libadjoint.exceptions.LibadjointErrorNotImplemented("Shouldn't ever get here")
def solve(selfmat, var, b):
if selfmat.adjoint:
operators = transpose_operators(selfmat.operators)
else:
operators = selfmat.operators
# Fetch/construct the solver
if var.type in ['ADJ_FORWARD', 'ADJ_TLM']:
solver = petsc_krylov_solvers[idx]
need_to_set_operator = self._need_to_reset_operator
else:
if adj_petsc_krylov_solvers[idx] is None:
need_to_set_operator = True
adj_petsc_krylov_solvers[idx] = PETScKrylovSolver(*solver_parameters)
adj_ksp = adj_petsc_krylov_solvers[idx].ksp()
fwd_ksp = petsc_krylov_solvers[idx].ksp()
adj_ksp.setOptionsPrefix(fwd_ksp.getOptionsPrefix())
adj_ksp.setType(fwd_ksp.getType())
adj_ksp.pc.setType(fwd_ksp.pc.getType())
adj_ksp.setFromOptions()
else:
need_to_set_operator = self._need_to_reset_operator
solver = adj_petsc_krylov_solvers[idx]
# FIXME: work around DOLFIN bug #583
try:
solver.parameters.convergence_norm_type
except:
solver.parameters.convergence_norm_type = "preconditioned"
# end FIXME
solver.parameters.update(parameters)
self._need_to_reset_operator = False
if selfmat.adjoint:
(nsp_, tnsp_) = (tnsp, nsp)
else:
(nsp_, tnsp_) = (nsp, tnsp)
x = dolfin.Function(fn_space)
if selfmat.initial_guess is not None and var.type == 'ADJ_FORWARD':
x.vector()[:] = selfmat.initial_guess.vector()
if b.data is None:
dolfin.info_red("Warning: got zero RHS for the solve associated with variable %s" % var)
return adjlinalg.Vector(x)
if var.type in ['ADJ_TLM', 'ADJ_ADJOINT']:
selfmat.bcs = [utils.homogenize(bc) for bc in selfmat.bcs if isinstance(bc, dolfin.cpp.DirichletBC)] + [bc for bc in selfmat.bcs if not isinstance(bc, dolfin.cpp.DirichletBC)]
# This is really hideous. Sorry.
if isinstance(b.data, dolfin.Function):
rhs = b.data.vector().copy()
[bc.apply(rhs) for bc in selfmat.bcs]
if need_to_set_operator:
if assemble_system: # if we called assemble_system, rather than assemble
v = dolfin.TestFunction(fn_space)
(A, rhstmp) = dolfin.assemble_system(operators[0], dolfin.inner(b.data, v)*dolfin.dx, selfmat.bcs)
if has_preconditioner:
(P, rhstmp) = dolfin.assemble_system(operators[1], dolfin.inner(b.data, v)*dolfin.dx, selfmat.bcs)
solver.set_operators(A, P)
else:
solver.set_operator(A)
else: # we called assemble
A = dolfin.assemble(operators[0])
[bc.apply(A) for bc in selfmat.bcs]
if has_preconditioner:
P = dolfin.assemble(operators[1])
[bc.apply(P) for bc in selfmat.bcs]
solver.set_operators(A, P)
else:
solver.set_operator(A)
else:
if assemble_system: # if we called assemble_system, rather than assemble
(A, rhs) = dolfin.assemble_system(operators[0], b.data, selfmat.bcs)
if need_to_set_operator:
if has_preconditioner:
(P, rhstmp) = dolfin.assemble_system(operators[1], b.data, selfmat.bcs)
solver.set_operators(A, P)
else:
solver.set_operator(A)
else: # we called assemble
A = dolfin.assemble(operators[0])
rhs = dolfin.assemble(b.data)
[bc.apply(A) for bc in selfmat.bcs]
[bc.apply(rhs) for bc in selfmat.bcs]
if need_to_set_operator:
if has_preconditioner:
P = dolfin.assemble(operators[1])
[bc.apply(P) for bc in selfmat.bcs]
solver.set_operators(A, P)
else:
solver.set_operator(A)
if need_to_set_operator:
print "|A|: %.6e" % A.norm("frobenius")
# Set the nullspace for the linear operator
if nsp_ is not None and need_to_set_operator:
dolfin.as_backend_type(A).set_nullspace(nsp_)
# (Possibly override the user in) orthogonalize
# the right-hand-side
if tnsp_ is not None:
tnsp_.orthogonalize(rhs)
print "%s: |b|: %.6e" % (var, rhs.norm("l2"))
solver.solve(x.vector(), rhs)
return adjlinalg.Vector(x)
solving.annotate(A == b, u, bcs, matrix_class=PETScKrylovSolverMatrix, initial_guess=parameters['nonzero_initial_guess'], replace_map=True)
out = dolfin.PETScKrylovSolver.solve(self, *args, **kwargs)
if to_annotate and dolfin.parameters["adjoint"]["record_all"]:
adjglobals.adjointer.record_variable(adjglobals.adj_variables[u], libadjoint.MemoryStorage(adjlinalg.Vector(u)))
return out
def solve(self, *args, **kwargs):
'''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).'''
to_annotate = utils.to_annotate(kwargs.pop("annotate", None))
if to_annotate:
if len(args) == 2:
try:
A = self.matrix.form
except AttributeError:
raise libadjoint.exceptions.LibadjointErrorInvalidInputs("Your matrix A has to have the .form attribute: was it assembled after from dolfin_adjoint import *?")
try:
self.op_bcs = self.matrix.bcs
except AttributeError:
self.op_bcs = []
try:
x = args[0].function
except AttributeError:
raise libadjoint.exceptions.LibadjointErrorInvalidInputs("Your solution x has to have a .function attribute; is it the .vector() of a Function?")
try:
b = args[1].form
except AttributeError:
raise libadjoint.exceptions.LibadjointErrorInvalidInputs("Your RHS b has to have the .form attribute: was it assembled after from dolfin_adjoint import *?")
try:
eq_bcs = misc.uniq(self.op_bcs + args[1].bcs)
except AttributeError:
eq_bcs = self.op_bcs
elif len(args) == 3:
A = args[0].form
try:
x = args[1].function
except AttributeError:
raise libadjoint.exceptions.LibadjointErrorInvalidInputs("Your solution x has to have a .function attribute; is it the .vector() of a Function?")
try:
self.op_bcs = A.bcs
except AttributeError:
self.op_bcs = []
try:
b = args[2].form
except AttributeError:
raise libadjoint.exceptions.LibadjointErrorInvalidInputs("Your RHS b has to have the .form attribute: was it assembled after from dolfin_adjoint import *?")
try:
eq_bcs = misc.uniq(self.op_bcs + args[2].bcs)
except AttributeError:
eq_bcs = self.op_bcs
else:
raise libadjoint.exceptions.LibadjointErrorInvalidInputs("LUSolver.solve() must be called with either (A, x, b) or (x, b).")
if self.parameters["reuse_factorization"] and self.__global_list_idx__ is None:
self.__global_list_idx__ = len(lu_solvers)
lu_solvers.append(self)
adj_lu_solvers.append(None)
solving.annotate(A == b, x, eq_bcs, solver_parameters={"linear_solver": "lu"}, matrix_class=make_LUSolverMatrix(self.__global_list_idx__, self.parameters["reuse_factorization"]))
out = dolfin.LUSolver.solve(self, *args, **kwargs)
if to_annotate:
if dolfin.parameters["adjoint"]["record_all"]:
adjglobals.adjointer.record_variable(adjglobals.adj_variables[x], libadjoint.MemoryStorage(adjlinalg.Vector(x)))
return out
def solve(self, *args, **kwargs):
'''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).'''
to_annotate = utils.to_annotate(kwargs.pop("annotate", None))
nsp = self.nsp
if to_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]
bcs = []
if hasattr(A, 'bcs'):
bcs += A.bcs
if hasattr(b, 'bcs'):
bcs += b.bcs
bcs = misc.uniq(bcs)
assemble_system = A.assemble_system
A = A.form
u = x.function
b = b.form
if self.operators[1] is not None:
P = self.operators[1].form
else:
P = None
solver_parameters = self.solver_parameters
parameters = self.parameters.to_dict()
fn_space = u.function_space()
has_preconditioner = P is not None
class LinearSolverMatrix(adjlinalg.Matrix):
def __init__(self, *args, **kwargs):
if 'initial_guess' in kwargs:
self.initial_guess = kwargs['initial_guess']
del kwargs['initial_guess']
else:
self.initial_guess = None
replace_map = kwargs['replace_map']
del kwargs['replace_map']
adjlinalg.Matrix.__init__(self, *args, **kwargs)
self.adjoint = kwargs['adjoint']
if P is None:
self.operators = (dolfin.replace(A, replace_map), None)
else:
self.operators = (dolfin.replace(A, replace_map), dolfin.replace(P, replace_map))
def axpy(self, alpha, x):
raise libadjoint.exceptions.LibadjointErrorNotImplemented("Shouldn't ever get here")
def solve(self, var, b):
if self.adjoint:
operators = transpose_operators(self.operators)
else:
operators = self.operators
solver = dolfin.LinearSolver(*solver_parameters)
solver.parameters.update(parameters)
x = dolfin.Function(fn_space)
if self.initial_guess is not None and var.type == 'ADJ_FORWARD':
x.vector()[:] = self.initial_guess.vector()
if b.data is None:
dolfin.info_red("Warning: got zero RHS for the solve associated with variable %s" % var)
return adjlinalg.Vector(x)
if var.type in ['ADJ_TLM', 'ADJ_ADJOINT']:
self.bcs = [utils.homogenize(bc) for bc in self.bcs if isinstance(bc, dolfin.cpp.DirichletBC)] + [bc for bc in self.bcs if not isinstance(bc, dolfin.cpp.DirichletBC)]
# This is really hideous. Sorry.
if isinstance(b.data, dolfin.Function):
rhs = b.data.vector().copy()
[bc.apply(rhs) for bc in self.bcs]
if assemble_system: # if we called assemble_system, rather than assemble
v = dolfin.TestFunction(fn_space)
(A, rhstmp) = dolfin.assemble_system(operators[0], dolfin.inner(b.data, v)*dolfin.dx, self.bcs)
if has_preconditioner:
(P, rhstmp) = dolfin.assemble_system(operators[1], dolfin.inner(b.data, v)*dolfin.dx, self.bcs)
solver.set_operators(A, P)
else:
solver.set_operator(A)
else: # we called assemble
A = dolfin.assemble(operators[0])
[bc.apply(A) for bc in self.bcs]
# Set nullspace
if nsp:
dolfin.as_backend_type(A).set_nullspace(nsp)
nsp.orthogonalize(b);
if has_preconditioner:
P = dolfin.assemble(operators[1])
[bc.apply(P) for bc in self.bcs]
solver.set_operators(A, P)
else:
solver.set_operator(A)
else:
if assemble_system: # if we called assemble_system, rather than assemble
(A, rhs) = dolfin.assemble_system(operators[0], b.data, self.bcs)
if has_preconditioner:
(P, rhstmp) = dolfin.assemble_system(operators[1], b.data, self.bcs)
solver.set_operators(A, P)
else:
solver.set_operator(A)
else: # we called assemble
A = dolfin.assemble(operators[0])
rhs = dolfin.assemble(b.data)
[bc.apply(A) for bc in self.bcs]
[bc.apply(rhs) for bc in self.bcs]
# Set nullspace
if nsp:
dolfin.as_backend_type(A).set_nullspace(nsp)
nsp.orthogonalize(rhs);
if has_preconditioner:
P = dolfin.assemble(operators[1])
[bc.apply(P) for bc in self.bcs]
solver.set_operators(A, P)
else:
solver.set_operator(A)
solver.solve(x.vector(), rhs)
return adjlinalg.Vector(x)
nonzero_initial_guess = parameters['nonzero_initial_guess'] if 'nonzero_initial_guess' in parameters.keys() else False
solving.annotate(A == b, u, bcs, matrix_class=LinearSolverMatrix, initial_guess=nonzero_initial_guess, replace_map=True)
out = dolfin.LinearSolver.solve(self, *args, **kwargs)
if to_annotate and dolfin.parameters["adjoint"]["record_all"]:
adjglobals.adjointer.record_variable(adjglobals.adj_variables[u], libadjoint.MemoryStorage(adjlinalg.Vector(u)))
return out
def _generateChecksForAllModels(self):
l = [m.members for m in self.models]
invocs = uniq(flatten(l))
self.contentProvider.generateChecksForInvocations(invocs)
def annotate(*args, **kwargs):
"""This routine handles all of the annotation, recording the solves as they
happen so that libadjoint can rewind them later."""
if "matrix_class" in kwargs:
matrix_class = kwargs["matrix_class"]
del kwargs["matrix_class"]
else:
matrix_class = adjlinalg.Matrix
if "initial_guess" in kwargs:
initial_guess = kwargs["initial_guess"]
del kwargs["initial_guess"]
else:
initial_guess = False
replace_map = False
if "replace_map" in kwargs:
replace_map = kwargs["replace_map"]
del kwargs["replace_map"]
if isinstance(args[0], ufl.classes.Equation):
# annotate !
# Unpack the arguments, using the same routine as the real Dolfin solve call
unpacked_args = compatibility._extract_args(*args, **kwargs)
eq = unpacked_args[0]
u = unpacked_args[1]
bcs = unpacked_args[2]
J = unpacked_args[3]
# create a deep copy of the parameters. They can be of type
# backend.Parameters or just a list
if type(unpacked_args[7]) == backend.Parameters:
solver_parameters = backend.Parameters(unpacked_args[7])
else:
solver_parameters = copy.deepcopy(unpacked_args[7])
if isinstance(eq.lhs, ufl.Form) and isinstance(eq.rhs, ufl.Form):
eq_lhs = eq.lhs
eq_rhs = eq.rhs
eq_bcs = bcs
linear = True
else:
eq_lhs, eq_rhs = define_nonlinear_equation(eq.lhs, u)
F = eq.lhs
eq_bcs = []
linear = False
elif isinstance(args[0], compatibility.matrix_types()):
linear = True
try:
eq_lhs = args[0].form
except (KeyError, AttributeError) as e:
raise libadjoint.exceptions.LibadjointErrorInvalidInputs(
"dolfin_adjoint did not assemble your form, and so does not recognise your matrix. Did you from dolfin_adjoint import *?"
)
try:
eq_rhs = args[2].form
except (KeyError, AttributeError) as e:
raise libadjoint.exceptions.LibadjointErrorInvalidInputs(
"dolfin_adjoint did not assemble your form, and so does not recognise your right-hand side. Did you from dolfin_adjoint import *?"
)
u = args[1]
u = u.function
solver_parameters = {}
try:
solver_parameters["linear_solver"] = args[3]
except IndexError:
pass
try:
solver_parameters["preconditioner"] = args[4]
except IndexError:
pass
try:
eq_bcs = misc.uniq(args[0].bcs + args[2].bcs)
except AttributeError:
assert not hasattr(args[0], "bcs") and not hasattr(args[2], "bcs")
eq_bcs = []
else:
print "args[0].__class__: ", args[0].__class__
raise libadjoint.exceptions.LibadjointErrorNotImplemented("Don't know how to annotate your equation, sorry!")
# Suppose we are solving for a variable w, and that variable shows up in the
# coefficients of eq_lhs/eq_rhs.
# Does that mean:
# a) the /previous value/ of that variable, and you want to timestep?
# b) the /value to be solved for/ in this solve?
# i.e. is it timelevel n-1, or n?
# if Dolfin is doing a linear solve, we want case a);
# if Dolfin is doing a nonlinear solve, we want case b).
# so if we are doing a nonlinear solve, we bump the timestep number here
# /before/ we map the coefficients -> dependencies,
# so that libadjoint records the dependencies with the right timestep number.
if not linear:
# Register the initial condition before the first nonlinear solve
register_initial_conditions([[u, adjglobals.adj_variables[u]]], linear=False)
var = adjglobals.adj_variables.next(u)
else:
var = None
# Set up the data associated with the matrix on the left-hand side. This goes on the diagonal
# of the 'large' system that incorporates all of the timelevels, which is why it is prefixed
# with diag.
diag_name = hashlib.md5(
str(eq_lhs) + str(eq_rhs) + str(u) + str(random.random())
).hexdigest() # we don't have a useful human-readable name, so take the md5sum of the string representation of the forms
diag_deps = [
adjglobals.adj_variables[coeff]
for coeff in ufl.algorithms.extract_coefficients(eq_lhs)
if hasattr(coeff, "function_space")
]
diag_coeffs = [coeff for coeff in ufl.algorithms.extract_coefficients(eq_lhs) if hasattr(coeff, "function_space")]
if (
initial_guess and linear
): # if the initial guess matters, we're going to have to add this in as a dependency of the system
initial_guess_var = adjglobals.adj_variables[u]
diag_deps.append(initial_guess_var)
diag_coeffs.append(u)
diag_block = libadjoint.Block(
diag_name,
dependencies=diag_deps,
test_hermitian=backend.parameters["adjoint"]["test_hermitian"],
test_derivative=backend.parameters["adjoint"]["test_derivative"],
)
# Similarly, create the object associated with the right-hand side data.
if linear:
rhs = adjrhs.RHS(eq_rhs)
else:
rhs = adjrhs.NonlinearRHS(eq_rhs, F, u, bcs, mass=eq_lhs, solver_parameters=solver_parameters, J=J)
# We need to check if this is the first equation,
# so that we can register the appropriate initial conditions.
# These equations are necessary so that libadjoint can assemble the
# relevant adjoint equations for the adjoint variables associated with
# the initial conditions.
assert len(rhs.coefficients()) == len(rhs.dependencies())
register_initial_conditions(
zip(rhs.coefficients(), rhs.dependencies()) + zip(diag_coeffs, diag_deps), linear=linear, var=var
)
# c.f. the discussion above. In the linear case, we want to bump the
# timestep number /after/ all of the dependencies' timesteps have been
# computed for libadjoint.
if linear:
var = adjglobals.adj_variables.next(u)
# With the initial conditions out of the way, let us now define the callbacks that
# define the actions of the operator the user has passed in on the lhs of this equation.
# Our equation may depend on Expressions, and those Expressions may have parameters
# (e.g. for time-dependent boundary conditions).
# In order to successfully replay the forward solve, we need to keep those parameters around.
# In expressions.py, we overloaded the Expression class to record all of the parameters
# as they are set. We're now going to copy that dictionary as it is at the annotation time,
# so that we can get back to this exact state:
frozen_expressions = expressions.freeze_dict()
frozen_constants = constant.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
value_coeffs = [v.data for v in values]
expressions.update_expressions(frozen_expressions)
constant.update_constants(frozen_constants)
eq_l = backend.replace(eq_lhs, dict(zip(diag_coeffs, value_coeffs)))
kwargs = {"cache": eq_l in caching.assembled_fwd_forms} # should we cache our matrices on the way backwards?
if hermitian:
# Homogenise the adjoint boundary conditions. This creates the adjoint
# solution associated with the lifted discrete system that is actually solved.
adjoint_bcs = [utils.homogenize(bc) for bc in eq_bcs if isinstance(bc, backend.DirichletBC)] + [
bc for bc in eq_bcs if not isinstance(bc, backend.DirichletBC)
]
if len(adjoint_bcs) == 0:
adjoint_bcs = None
else:
adjoint_bcs = misc.uniq(adjoint_bcs)
kwargs["bcs"] = adjoint_bcs
kwargs["solver_parameters"] = solver_parameters
kwargs["adjoint"] = True
if initial_guess:
kwargs["initial_guess"] = value_coeffs[dependencies.index(initial_guess_var)]
if replace_map:
kwargs["replace_map"] = dict(zip(diag_coeffs, value_coeffs))
return (
matrix_class(
backend.adjoint(eq_l, reordered_arguments=ufl.algorithms.extract_arguments(eq_l)), **kwargs
),
adjlinalg.Vector(None, fn_space=u.function_space()),
)
else:
kwargs["bcs"] = misc.uniq(eq_bcs)
kwargs["solver_parameters"] = solver_parameters
kwargs["adjoint"] = False
if initial_guess:
kwargs["initial_guess"] = value_coeffs[dependencies.index(initial_guess_var)]
if replace_map:
kwargs["replace_map"] = dict(zip(diag_coeffs, value_coeffs))
return (matrix_class(eq_l, **kwargs), adjlinalg.Vector(None, fn_space=u.function_space()))
diag_block.assemble = diag_assembly_cb
def diag_action_cb(dependencies, values, hermitian, coefficient, input, context):
value_coeffs = [v.data for v in values]
expressions.update_expressions(frozen_expressions)
constant.update_constants(frozen_constants)
eq_l = backend.replace(eq_lhs, dict(zip(diag_coeffs, value_coeffs)))
if hermitian:
eq_l = backend.adjoint(eq_l)
output = coefficient * backend.action(eq_l, input.data)
return adjlinalg.Vector(output)
diag_block.action = diag_action_cb
if len(diag_deps) > 0:
# If this block is nonlinear (the entries of the matrix on the LHS
# depend on any variable previously computed), then that will induce
# derivative terms in the adjoint equations. Here, we define the
# callback libadjoint will need to compute such terms.
def derivative_action(
dependencies, values, variable, contraction_vector, hermitian, input, coefficient, context
):
dolfin_variable = values[dependencies.index(variable)].data
dolfin_values = [val.data for val in values]
expressions.update_expressions(frozen_expressions)
constant.update_constants(frozen_constants)
current_form = backend.replace(eq_lhs, dict(zip(diag_coeffs, dolfin_values)))
deriv = backend.derivative(current_form, dolfin_variable)
args = ufl.algorithms.extract_arguments(deriv)
deriv = backend.replace(deriv, {args[1]: contraction_vector.data}) # contract over the middle index
# Assemble the G-matrix now, so that we can apply the Dirichlet BCs to it
if len(ufl.algorithms.extract_arguments(ufl.algorithms.expand_derivatives(coefficient * deriv))) == 0:
return adjlinalg.Vector(None)
G = coefficient * deriv
if hermitian:
output = backend.action(backend.adjoint(G), input.data)
else:
output = backend.action(G, input.data)
return adjlinalg.Vector(output)
diag_block.derivative_action = derivative_action
def derivative_outer_action(
dependencies, values, variable, contraction_vector, hermitian, input, coefficient, context
):
dolfin_variable = values[dependencies.index(variable)].data
dolfin_values = [val.data for val in values]
expressions.update_expressions(frozen_expressions)
constant.update_constants(frozen_constants)
current_form = backend.replace(eq_lhs, dict(zip(diag_coeffs, dolfin_values)))
deriv = backend.derivative(current_form, dolfin_variable)
args = ufl.algorithms.extract_arguments(deriv)
deriv = backend.replace(deriv, {args[2]: contraction_vector.data}) # contract over the outer index
# Assemble the G-matrix now, so that we can apply the Dirichlet BCs to it
if len(ufl.algorithms.extract_arguments(ufl.algorithms.expand_derivatives(coefficient * deriv))) == 0:
return adjlinalg.Vector(None)
G = coefficient * deriv
if hermitian:
output = backend.action(backend.adjoint(G), input.data)
else:
output = backend.action(G, input.data)
return adjlinalg.Vector(output)
diag_block.derivative_outer_action = derivative_outer_action
def second_derivative_action(
dependencies,
values,
inner_variable,
inner_contraction_vector,
outer_variable,
outer_contraction_vector,
hermitian,
input,
coefficient,
context,
):
dolfin_inner_variable = values[dependencies.index(inner_variable)].data
dolfin_outer_variable = values[dependencies.index(outer_variable)].data
dolfin_values = [val.data for val in values]
expressions.update_expressions(frozen_expressions)
constant.update_constants(frozen_constants)
current_form = backend.replace(eq_lhs, dict(zip(diag_coeffs, dolfin_values)))
deriv = backend.derivative(current_form, dolfin_inner_variable)
args = ufl.algorithms.extract_arguments(deriv)
deriv = backend.replace(deriv, {args[1]: inner_contraction_vector.data}) # contract over the middle index
deriv = backend.derivative(deriv, dolfin_outer_variable)
args = ufl.algorithms.extract_arguments(deriv)
deriv = backend.replace(deriv, {args[1]: outer_contraction_vector.data}) # contract over the middle index
# Assemble the G-matrix now, so that we can apply the Dirichlet BCs to it
if len(ufl.algorithms.extract_arguments(ufl.algorithms.expand_derivatives(coefficient * deriv))) == 0:
return adjlinalg.Vector(None)
G = coefficient * deriv
if hermitian:
output = backend.action(backend.adjoint(G), input.data)
else:
output = backend.action(G, input.data)
return adjlinalg.Vector(output)
diag_block.second_derivative_action = second_derivative_action
eqn = libadjoint.Equation(var, blocks=[diag_block], targets=[var], rhs=rhs)
cs = adjglobals.adjointer.register_equation(eqn)
do_checkpoint(cs, var, rhs)
return linear
def _generateChecksForAllModels(self):
l = [m.members for m in self.models]
invocs = uniq(flatten(l))
self.contentProvider.generateChecksForInvocations(invocs)
