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 __init__(self, *args, **kwargs):
annotate = kwargs.pop("annotate", None)
to_annotate = utils.to_annotate(annotate)
if "name" in kwargs:
self.adj_name = kwargs["name"]
#if self.adj_name in adjglobals.function_names and to_annotate:
# backend.info_red("Warning: got duplicate function name %s" % self.adj_name)
adjglobals.function_names.add(self.adj_name)
del kwargs["name"]
backend.Function.__init__(self, *args, **kwargs)
if hasattr(self, 'adj_name'):
if backend.__name__ == "dolfin":
self.rename(self.adj_name, "a Function from dolfin-adjoint")
else:
self.name = self.__str__
if to_annotate:
if not isinstance(args[0], compatibility.function_space_type):
if isinstance(args[0], backend.Function):
known = adjglobals.adjointer.variable_known(adjglobals.adj_variables[args[0]])
else:
known = True
if known or (annotate is True):
assignment.register_assign(self, args[0])
else:
adjglobals.adj_variables.forget(args[0])
def project_firedrake(v, V, **kwargs):
annotate = kwargs.pop("annotate", None)
to_annotate = utils.to_annotate(annotate)
if isinstance(v, backend.Expression) and (annotate is not True):
to_annotate = False
if isinstance(v, backend.Constant) and (annotate is not True):
to_annotate = False
if isinstance(V, backend.functionspaceimpl.WithGeometry):
result = utils.function_to_da_function(backend.Function(V, name=None))
elif isinstance(V, backend.function.Function):
result = utils.function_to_da_function(V)
else:
raise ValueError("Can't project into a '%r'" % V)
name = kwargs.pop("name", None)
if name is not None:
result.adj_name = name
result.rename(name, "a Function from dolfin-adjoint")
with misc.annotations(to_annotate):
result = backend.project(v, result, **kwargs)
return result
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 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:
factory = args[0]
vec = args[1]
b = dolfin.as_backend_type(vec).__class__()
factory.F(b=b, x=vec)
F = b.form
bcs = b.bcs
u = vec.function
var = adjglobals.adj_variables[u]
solving.annotate(F == 0, u, bcs, solver_parameters={"newton_solver": self.parameters.to_dict()})
newargs = [self] + list(args)
out = dolfin.NewtonSolver.solve(*newargs, **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, annotate=None):
"""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)."""
annotate = utils.to_annotate(annotate)
if annotate:
problem = self.problem
solving.annotate(
problem.F == 0,
problem.u,
problem.bcs,
J=problem.J,
solver_parameters=compatibility.to_dict(self.parameters),
)
out = backend.NonlinearVariationalSolver.solve(self)
if annotate and backend.parameters["adjoint"]["record_all"]:
adjglobals.adjointer.record_variable(
adjglobals.adj_variables[self.problem.u], libadjoint.MemoryStorage(adjlinalg.Vector(self.problem.u))
)
return out
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 project(self, other, annotate=None, *args, **kwargs):
'''To disable the annotation, just pass :py:data:`annotate=False` to this routine, and it acts exactly like the
Firedrake project call.'''
to_annotate = utils.to_annotate(annotate)
if not to_annotate:
flag = misc.pause_annotation()
res = firedrake_project(self, other, *args, **kwargs)
if not to_annotate:
misc.continue_annotation(flag)
return res
def project_dolfin(v, V=None, bcs=None, mesh=None, solver_type="cg", preconditioner_type="default", form_compiler_parameters=None, annotate=None, name=None):
'''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 libadjoint.
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).'''
to_annotate = utils.to_annotate(annotate)
if isinstance(v, backend.Expression) and (annotate is not True):
to_annotate = False
if isinstance(v, backend.Constant) and (annotate is not True):
to_annotate = False
out = backend.project(v=v, V=V, bcs=bcs, mesh=mesh, solver_type=solver_type, preconditioner_type=preconditioner_type, form_compiler_parameters=form_compiler_parameters)
out = utils.function_to_da_function(out)
if name is not None:
out.adj_name = name
out.rename(name, "a Function from dolfin-adjoint")
if to_annotate:
# reproduce the logic from project. This probably isn't future-safe, but anyway
if V is None:
V = backend.fem.projection._extract_function_space(v, mesh)
if mesh is None:
mesh = V.mesh()
# Define variational problem for projection
w = backend.TestFunction(V)
Pv = backend.TrialFunction(V)
a = backend.inner(w, Pv)*backend.dx(domain=mesh)
L = backend.inner(w, v)*backend.dx(domain=mesh)
solving.annotate(a == L, out, bcs, solver_parameters={"linear_solver": solver_type, "preconditioner": preconditioner_type, "symmetric": True})
if backend.parameters["adjoint"]["record_all"]:
adjglobals.adjointer.record_variable(adjglobals.adj_variables[out], libadjoint.MemoryStorage(adjlinalg.Vector(out)))
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 libadjoint.
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)."""
# First, decide if we should annotate or not.
to_annotate = utils.to_annotate(kwargs.pop("annotate", None))
if to_annotate:
linear = annotate(*args, **kwargs)
# Avoid recursive annotation
flag = misc.pause_annotation()
try:
ret = backend.solve(*args, **kwargs)
except:
raise
finally:
misc.continue_annotation(flag)
if to_annotate:
# Finally, if we want to record all of the solutions of the real forward model
# (for comparison with a libadjoint replay later),
# then we should record the value of the variable we just solved for.
if backend.parameters["adjoint"]["record_all"]:
if isinstance(args[0], ufl.classes.Equation):
unpacked_args = compatibility._extract_args(*args, **kwargs)
u = unpacked_args[1]
adjglobals.adjointer.record_variable(
adjglobals.adj_variables[u], libadjoint.MemoryStorage(adjlinalg.Vector(u))
)
elif isinstance(args[0], compatibility.matrix_types()):
u = args[1].function
adjglobals.adjointer.record_variable(
adjglobals.adj_variables[u], libadjoint.MemoryStorage(adjlinalg.Vector(u))
)
else:
raise libadjoint.exceptions.LibadjointErrorInvalidInputs("Don't know how to record, sorry")
return ret
def copy(self, *args, **kwargs):
name = kwargs.pop("name", None)
annotate = kwargs.pop("annotate", None)
to_annotate = utils.to_annotate(annotate)
with misc.annotations(False):
copy = backend.Function.copy(self, *args, **kwargs)
copy = utils.function_to_da_function(copy)
if name is not None:
copy.adj_name = name
copy.rename(name, "a Function from dolfin-adjoint")
adjglobals.function_names.add(name)
if to_annotate:
assignment.register_assign(copy, self)
return copy
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, dolfin_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)`.
"""
form = args[0]
caching.assembled_fwd_forms.add(form)
cache = kwargs.pop("cache", False)
to_annotate = utils.to_annotate(kwargs.pop("annotate", None))
output = backend_assemble(*args, **kwargs)
if not isinstance(output, float) and to_annotate:
output.form = form
output.assemble_system = False
if cache:
caching.assembled_adj_forms[form] = output
return output
def assemble_system(*args, **kwargs):
"""When a form is assembled, the information about its nonlinear dependencies is lost,
and it is no longer easy to manipulate. Therefore, dolfin_adjoint overloads the :py:func:`dolfin.assemble_system`
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)`.
"""
lhs = args[0]
rhs = args[1]
caching.assembled_fwd_forms.add(lhs)
caching.assembled_fwd_forms.add(rhs)
cache = kwargs.pop("cache", False)
if "bcs" in kwargs:
bcs = kwargs["bcs"]
elif len(args) > 2:
bcs = args[2]
else:
bcs = []
if not isinstance(bcs, list):
bcs = [bcs]
(lhs_out, rhs_out) = backend.assemble_system(*args, **kwargs)
to_annotate = utils.to_annotate(kwargs.pop("annotate", None))
if to_annotate:
lhs_out.form = lhs
lhs_out.bcs = bcs
lhs_out.assemble_system = True
rhs_out.form = rhs
rhs_out.bcs = bcs
rhs_out.assemble_system = True
if cache:
caching.assembled_adj_forms[lhs] = lhs_out
caching.assembled_adj_forms[rhs] = rhs_out
return (lhs_out, rhs_out)
def solve_local(self, x_vec, b_vec, b_dofmap, **kwargs):
# Figure out whether to annotate or not
to_annotate = utils.to_annotate(kwargs.pop("annotate", None))
x = x_vec.function
if to_annotate:
L = b_vec.form
# Set Matrix class for solving the adjoint systems
solving.annotate(self.a == L, x, \
solver_parameters={"solver_type": self.solver_type, "factorize" : self.adjoint_factorize}, \
matrix_class=LocalSolverMatrix)
# Use standard local solver
out = dolfin.LocalSolver.solve_local(self, x_vec, b_vec, b_dofmap)
if to_annotate:
# checkpointing
if dolfin.parameters["adjoint"]["record_all"]:
adjglobals.adjointer.record_variable(adjglobals.adj_variables[x],
libadjoint.MemoryStorage(adjlinalg.Vector(x)))
return out
def project_firedrake(*args, **kwargs):
try:
annotate = kwargs["annotate"]
kwargs.pop("annotate")
except KeyError:
annotate = None
to_annotate = utils.to_annotate(annotate)
if isinstance(args[0], backend.Expression) and (annotate is not True):
to_annotate = False
if isinstance(args[0], backend.Constant) and (annotate is not True):
to_annotate = False
if to_annotate:
result = backend.project(*args, **kwargs)
else:
flag = misc.pause_annotation()
result = backend.project(*args, **kwargs)
misc.continue_annotation(flag)
return result
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 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
