def _init_params(self, args, kwargs, varform):
if len(self.forward_args) <= 0:
self.forward_args = args
if len(self.forward_kwargs) <= 0:
self.forward_kwargs = kwargs.copy()
if len(self.adj_args) <= 0:
self.adj_args = self.forward_args
if len(self.adj_kwargs) <= 0:
self.adj_kwargs = self.forward_kwargs.copy()
if varform:
if "J" in self.forward_kwargs:
self.adj_kwargs["J"] = backend.adjoint(
self.forward_kwargs["J"])
if "Jp" in self.forward_kwargs:
self.adj_kwargs["Jp"] = backend.adjoint(
self.forward_kwargs["Jp"])
if "M" in self.forward_kwargs:
raise NotImplementedError(
"Annotation of adaptive solves not implemented.")
self.adj_kwargs.pop("appctx", None)
if "solver_parameters" in kwargs and "mat_type" in kwargs[
"solver_parameters"]:
self.assemble_kwargs["mat_type"] = kwargs["solver_parameters"][
"mat_type"]
if varform:
if "appctx" in kwargs:
self.assemble_kwargs["appctx"] = kwargs["appctx"]
def _assemble_soa_eq_rhs(self, dFdu_form, adj_sol, hessian_input, d2Fdu2):
# Start piecing together the rhs of the soa equation
b = hessian_input.copy()
b_form = d2Fdu2
for bo in self.get_dependencies():
c = bo.output
c_rep = bo.saved_output
tlm_input = bo.tlm_value
if (c == self.func and not self.linear) or tlm_input is None:
continue
if isinstance(c, compat.MeshType):
X = backend.SpatialCoordinate(c)
dFdu_adj = backend.action(backend.adjoint(dFdu_form), adj_sol)
d2Fdudm = ufl.algorithms.expand_derivatives(
backend.derivative(dFdu_adj, X, tlm_input))
if len(d2Fdudm.integrals()) > 0:
b -= compat.assemble_adjoint_value(d2Fdudm)
elif not isinstance(c, backend.DirichletBC):
b_form += backend.derivative(dFdu_form, c_rep, tlm_input)
b_form = ufl.algorithms.expand_derivatives(b_form)
if len(b_form.integrals()) > 0:
b_form = backend.adjoint(b_form)
b -= compat.assemble_adjoint_value(backend.action(b_form, adj_sol))
return b
def transpose_operators(operators):
out = [None, None]
for i in range(2):
op = operators[i]
if op is None:
out[i] = None
elif isinstance(op, backend.cpp.GenericMatrix):
out[i] = op.__class__()
backend.assemble(backend.adjoint(op.form), tensor=out[i])
if hasattr(op, 'bcs'):
adjoint_bcs = [backend.homogenize(bc) for bc in op.bcs if isinstance(bc, backend.cpp.DirichletBC)] + [bc for bc in op.bcs if not isinstance(bc, backend.DirichletBC)]
[bc.apply(out[i]) for bc in adjoint_bcs]
elif isinstance(op, backend.Form) or isinstance(op, ufl.form.Form):
out[i] = backend.adjoint(op)
if hasattr(op, 'bcs'):
out[i].bcs = [backend.homogenize(bc) for bc in op.bcs if isinstance(bc, backend.cpp.DirichletBC)] + [bc for bc in op.bcs if not isinstance(bc, backend.DirichletBC)]
elif isinstance(op, AdjointKrylovMatrix):
pass
else:
print "op.__class__: ", op.__class__
raise libadjoint.exceptions.LibadjointErrorNotImplemented("Don't know how to transpose anything else!")
return out
def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
fwd_block_variable = self.get_outputs()[0]
u = fwd_block_variable.output
dJdu = adj_inputs[0]
F_form = self._create_F_form()
dFdu = backend.derivative(F_form, fwd_block_variable.saved_output,
backend.TrialFunction(u.function_space()))
dFdu_form = backend.adjoint(dFdu)
dJdu = dJdu.copy()
compute_bdy = self._should_compute_boundary_adjoint(
relevant_dependencies)
adj_sol, adj_sol_bdy = self._assemble_and_solve_adj_eq(
dFdu_form, dJdu, compute_bdy)
self.adj_sol = adj_sol
if self.adj_cb is not None:
self.adj_cb(adj_sol)
if self.adj_bdy_cb is not None and compute_bdy:
self.adj_bdy_cb(adj_sol_bdy)
r = {}
r["form"] = F_form
r["adj_sol"] = adj_sol
r["adj_sol_bdy"] = adj_sol_bdy
return r
def derivative_assembly(self, dependencies, values, variable, hermitian):
replace_map = {}
for i in range(len(self.deps)):
if self.deps[i] == self.ic_var: continue
j = dependencies.index(self.deps[i])
replace_map[self.coeffs[i]] = values[j].data
diff_var = values[dependencies.index(variable)].data
current_form = backend.replace(self.form, replace_map)
deriv = backend.derivative(current_form, diff_var)
if hermitian:
deriv = backend.adjoint(deriv)
bcs = [
utils.homogenize(bc)
for bc in self.bcs if isinstance(bc, backend.DirichletBC)
] + [
bc
for bc in self.bcs if not isinstance(bc, backend.DirichletBC)
]
else:
bcs = self.bcs
return adjlinalg.Matrix(deriv, bcs=bcs)
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)
def hermitian(self):
adjoint_bcs = [
utils.homogenize(bc)
for bc in self.bcs if isinstance(bc, backend.cpp.DirichletBC)
] + [bc for bc in self.bcs if not isinstance(bc, backend.DirichletBC)]
return AdjointKrylovMatrix(backend.adjoint(self.original_form),
bcs=adjoint_bcs)
def second_derivative_action(self, dependencies, values, inner_variable, inner_contraction_vector, outer_variable, hermitian, action_vector):
if isinstance(self.form, ufl.form.Form):
# Find the dolfin Function corresponding to variable.
dolfin_inner_variable = values[dependencies.index(inner_variable)].data
dolfin_outer_variable = values[dependencies.index(outer_variable)].data
dolfin_dependencies = [dep for dep in _extract_function_coeffs(self.form)]
dolfin_values = [val.data for val in values]
current_form = backend.replace(self.form, dict(zip(dolfin_dependencies, dolfin_values)))
trial = backend.TrialFunction(dolfin_outer_variable.function_space())
d_rhs = backend.derivative(current_form, dolfin_inner_variable, inner_contraction_vector.data)
d_rhs = ufl.algorithms.expand_derivatives(d_rhs)
if len(d_rhs.integrals()) == 0:
return None
d_rhs = backend.derivative(d_rhs, dolfin_outer_variable, trial)
d_rhs = ufl.algorithms.expand_derivatives(d_rhs)
if len(d_rhs.integrals()) == 0:
return None
if hermitian:
action = backend.action(backend.adjoint(d_rhs), action_vector.data)
else:
action = backend.action(d_rhs, action_vector.data)
return adjlinalg.Vector(action)
else:
# RHS is a adjlinalg.Vector. Its derivative is therefore zero.
raise exceptions.LibadjointErrorNotImplemented("No derivative method for constant RHS.")
def derivative_action(self, dependencies, values, variable,
contraction_vector, hermitian):
if contraction_vector.data is None:
return adjlinalg.Vector(None)
if isinstance(self.form, ufl.form.Form):
# Find the dolfin Function corresponding to variable.
dolfin_variable = values[dependencies.index(variable)].data
dolfin_dependencies = [
dep for dep in _extract_function_coeffs(self.form)
]
dolfin_values = [val.data for val in values]
current_form = backend.replace(
self.form, dict(zip(dolfin_dependencies, dolfin_values)))
trial = backend.TrialFunction(dolfin_variable.function_space())
d_rhs = backend.derivative(current_form, dolfin_variable, trial)
if hermitian:
action = backend.action(backend.adjoint(d_rhs),
contraction_vector.data)
else:
action = backend.action(d_rhs, contraction_vector.data)
return adjlinalg.Vector(action)
else:
# RHS is a adjlinalg.Vector. Its derivative is therefore zero.
return adjlinalg.Vector(None)
def derivative_action(self, dependencies, values, variable, contraction_vector, hermitian):
if contraction_vector.data is None:
return adjlinalg.Vector(None)
if isinstance(self.form, ufl.form.Form):
# Find the dolfin Function corresponding to variable.
dolfin_variable = values[dependencies.index(variable)].data
dolfin_dependencies = [dep for dep in _extract_function_coeffs(self.form)]
dolfin_values = [val.data for val in values]
current_form = backend.replace(self.form, dict(zip(dolfin_dependencies, dolfin_values)))
trial = backend.TrialFunction(dolfin_variable.function_space())
d_rhs = backend.derivative(current_form, dolfin_variable, trial)
if hermitian:
action = backend.action(backend.adjoint(d_rhs), contraction_vector.data)
else:
action = backend.action(d_rhs, contraction_vector.data)
return adjlinalg.Vector(action)
else:
# RHS is a adjlinalg.Vector. Its derivative is therefore zero.
return adjlinalg.Vector(None)
def transpose_operators(operators):
out = [None, None]
for i in range(2):
op = operators[i]
if op is None:
out[i] = None
elif isinstance(op, backend.cpp.GenericMatrix):
out[i] = op.__class__()
backend.assemble(backend.adjoint(op.form), tensor=out[i])
if hasattr(op, 'bcs'):
adjoint_bcs = [
utils.homogenize(bc)
for bc in op.bcs if isinstance(bc, backend.cpp.DirichletBC)
] + [
bc
for bc in op.bcs if not isinstance(bc, backend.DirichletBC)
]
[bc.apply(out[i]) for bc in adjoint_bcs]
elif isinstance(op, backend.Form) or isinstance(op, ufl.form.Form):
out[i] = backend.adjoint(op)
if hasattr(op, 'bcs'):
out[i].bcs = [
utils.homogenize(bc)
for bc in op.bcs if isinstance(bc, backend.cpp.DirichletBC)
] + [
bc
for bc in op.bcs if not isinstance(bc, backend.DirichletBC)
]
elif isinstance(op, AdjointKrylovMatrix):
pass
else:
print("op.__class__: ", op.__class__)
raise libadjoint.exceptions.LibadjointErrorNotImplemented(
"Don't know how to transpose anything else!")
return out
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 derivative_action(self, variable, contraction_vector, hermitian, input, coefficient):
deriv = backend.derivative(self.current_form, variable)
args = ufl.algorithms.extract_arguments(deriv)
deriv = backend.replace(deriv, {args[1]: contraction_vector})
if hermitian:
deriv = backend.adjoint(deriv)
action = coefficient * backend.action(deriv, input)
return action
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)
def derivative_action(self, variable, contraction_vector, hermitian, input,
coefficient):
deriv = backend.derivative(self.current_form, variable)
args = ufl.algorithms.extract_arguments(deriv)
deriv = backend.replace(deriv, {args[1]: contraction_vector})
if hermitian:
deriv = backend.adjoint(deriv)
action = coefficient * backend.action(deriv, input)
return action
def _assemble_and_solve_soa_eq(self, dFdu_form, adj_sol, hessian_input,
d2Fdu2, compute_bdy):
b = self._assemble_soa_eq_rhs(dFdu_form, adj_sol, hessian_input,
d2Fdu2)
dFdu_form = backend.adjoint(dFdu_form)
adj_sol2, adj_sol2_bdy = self._assemble_and_solve_adj_eq(
dFdu_form, b, compute_bdy)
if self.adj2_cb is not None:
self.adj2_cb(adj_sol2)
if self.adj2_bdy_cb is not None and compute_bdy:
self.adj2_bdy_cb(adj_sol2_bdy)
return adj_sol2, adj_sol2_bdy
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))
def evaluate_adj_component(self,
inputs,
adj_inputs,
block_variable,
idx,
prepared=None):
if not self.linear and self.func == block_variable.output:
# We are not able to calculate derivatives wrt initial guess.
return None
F_form = prepared["form"]
adj_sol = prepared["adj_sol"]
adj_sol_bdy = prepared["adj_sol_bdy"]
c = block_variable.output
c_rep = block_variable.saved_output
if isinstance(c, backend.Function):
trial_function = backend.TrialFunction(c.function_space())
elif isinstance(c, backend.Constant):
mesh = compat.extract_mesh_from_form(F_form)
trial_function = backend.TrialFunction(c._ad_function_space(mesh))
elif isinstance(c, compat.ExpressionType):
mesh = F_form.ufl_domain().ufl_cargo()
c_fs = c._ad_function_space(mesh)
trial_function = backend.TrialFunction(c_fs)
elif isinstance(c, backend.DirichletBC):
tmp_bc = compat.create_bc(c,
value=extract_subfunction(
adj_sol_bdy, c.function_space()))
return [tmp_bc]
elif isinstance(c, compat.MeshType):
# Using CoordianteDerivative requires us to do action before
# differentiating, might change in the future.
F_form_tmp = backend.action(F_form, adj_sol)
X = backend.SpatialCoordinate(c_rep)
dFdm = backend.derivative(-F_form_tmp, X)
dFdm = compat.assemble_adjoint_value(dFdm, **self.assemble_kwargs)
return dFdm
dFdm = -backend.derivative(F_form, c_rep, trial_function)
dFdm = backend.adjoint(dFdm)
dFdm = dFdm * adj_sol
dFdm = compat.assemble_adjoint_value(dFdm, **self.assemble_kwargs)
if isinstance(c, compat.ExpressionType):
return [[dFdm, c_fs]]
else:
return dFdm
def second_derivative_action(self, dependencies, values, inner_variable,
inner_contraction_vector, outer_variable,
hermitian, action_vector):
if isinstance(self.form, ufl.form.Form):
# Find the dolfin Function corresponding to variable.
dolfin_inner_variable = values[dependencies.index(
inner_variable)].data
dolfin_outer_variable = values[dependencies.index(
outer_variable)].data
dolfin_dependencies = [
dep for dep in _extract_function_coeffs(self.form)
]
dolfin_values = [val.data for val in values]
current_form = backend.replace(
self.form, dict(zip(dolfin_dependencies, dolfin_values)))
trial = backend.TrialFunction(
dolfin_outer_variable.function_space())
d_rhs = backend.derivative(current_form, dolfin_inner_variable,
inner_contraction_vector.data)
d_rhs = ufl.algorithms.expand_derivatives(d_rhs)
if len(d_rhs.integrals()) == 0:
return None
d_rhs = backend.derivative(d_rhs, dolfin_outer_variable, trial)
d_rhs = ufl.algorithms.expand_derivatives(d_rhs)
if len(d_rhs.integrals()) == 0:
return None
if hermitian:
action = backend.action(backend.adjoint(d_rhs),
action_vector.data)
else:
action = backend.action(d_rhs, action_vector.data)
return adjlinalg.Vector(action)
else:
# RHS is a adjlinalg.Vector. Its derivative is therefore zero.
raise libadjoint.exceptions.LibadjointErrorNotImplemented(
"No derivative method for constant RHS.")
def derivative_assembly(self, dependencies, values, variable, hermitian):
replace_map = {}
for i in range(len(self.deps)):
if self.deps[i] == self.ic_var: continue
j = dependencies.index(self.deps[i])
replace_map[self.coeffs[i]] = values[j].data
diff_var = values[dependencies.index(variable)].data
current_form = backend.replace(self.form, replace_map)
deriv = backend.derivative(current_form, diff_var)
if hermitian:
deriv = backend.adjoint(deriv)
bcs = [utils.homogenize(bc) for bc in self.bcs if isinstance(bc, backend.DirichletBC)] + [bc for bc in self.bcs if not isinstance(bc, backend.DirichletBC)]
else:
bcs = self.bcs
return adjlinalg.Matrix(deriv, bcs=bcs)
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))
def hermitian(self): adjoint_bcs = [backend.homogenize(bc) for bc in self.bcs if isinstance(bc, backend.cpp.DirichletBC)] + [bc for bc in self.bcs if not isinstance(bc, backend.DirichletBC)] return AdjointKrylovMatrix(backend.adjoint(self.original_form), bcs=adjoint_bcs)