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 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 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 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 adjoint_genericmatrix_mul(self, other):
out = backend_genericmatrix_mul(self, other)
if hasattr(self, 'form') and isinstance(other, backend.cpp.la.GenericVector):
if hasattr(other, 'form'):
out.form = backend.action(self.form, other.form)
elif hasattr(other, 'function'):
if hasattr(other, 'function_factor'):
out.form = backend.action(other.function_factor * self.form, other.function)
else:
out.form = backend.action(self.form, other.function)
return out
def get_residual(i):
from .adjrhs import adj_get_forward_equation
(fwd_var, lhs, rhs) = adj_get_forward_equation(i)
if isinstance(lhs, adjlinalg.IdentityMatrix):
return None
fn_space = ufl.algorithms.extract_arguments(lhs)[0].function_space()
x = backend.Function(fn_space)
if rhs == 0:
form = lhs
x = fwd_var.nonlinear_u
else:
form = backend.action(lhs, x) - rhs
try:
y = adjglobals.adjointer.get_variable_value(fwd_var).data
except libadjoint.exceptions.LibadjointErrorNeedValue:
info_red(
"Warning: recomputing forward solution; please report this script on launchpad"
)
y = adjglobals.adjointer.get_forward_solution(i)[1].data
form = backend.replace(form, {x: y})
return form
def _assemble_and_solve_adj_eq(self,
dFdu_adj_form,
dJdu,
compute_bdy=True):
dJdu_copy = dJdu.copy()
dFdu = compat.assemble_adjoint_value(dFdu_adj_form,
**self.assemble_kwargs)
bcs = self._homogenize_bcs()
# Apply boundary conditions on adj_dFdu and dJdu.
for bc in bcs:
bc.apply(dFdu, dJdu)
adj_sol = compat.create_function(self.function_space)
lu_solver_methods = backend.lu_solver_methods()
solver_method = self.adj_args[0] if len(
self.adj_args) >= 1 else "default"
solver_method = "default" if solver_method == "lu" else solver_method
if solver_method in lu_solver_methods:
solver = backend.LUSolver(solver_method)
solver_parameters = self.adj_kwargs.get("lu_solver", {})
else:
solver = backend.KrylovSolver(*self.adj_args)
solver_parameters = self.adj_kwargs.get("krylov_solver", {})
solver.parameters.update(solver_parameters)
solver.solve(dFdu, adj_sol.vector(), dJdu)
adj_sol_bdy = None
if compute_bdy:
adj_sol_bdy = compat.function_from_vector(
self.function_space, dJdu_copy - compat.assemble_adjoint_value(
backend.action(dFdu_adj_form, adj_sol)))
return adj_sol, adj_sol_bdy
def action(self, x, y):
assert isinstance(x.data, backend.Function)
assert isinstance(y.data, backend.Function)
action_form = backend.action(self.data, x.data)
action_vec = backend.assemble(action_form)
y.data.vector()[:] = action_vec
def get_residual(i):
from adjrhs import adj_get_forward_equation
(fwd_var, lhs, rhs) = adj_get_forward_equation(i)
if isinstance(lhs, adjlinalg.IdentityMatrix):
return None
fn_space = ufl.algorithms.extract_arguments(lhs)[0].function_space()
x = backend.Function(fn_space)
if rhs == 0:
form = lhs
x = fwd_var.nonlinear_u
else:
form = backend.action(lhs, x) - rhs
try:
y = adjglobals.adjointer.get_variable_value(fwd_var).data
except libadjoint.exceptions.LibadjointErrorNeedValue:
info_red("Warning: recomputing forward solution; please report this script on launchpad")
y = adjglobals.adjointer.get_forward_solution(i)[1].data
form = backend.replace(form, {x: y})
return form
def _assemble_and_solve_adj_eq(self, dFdu_adj_form, dJdu, compute_bdy):
dJdu_copy = dJdu.copy()
bcs = self._homogenize_bcs()
solver = self.block_helper.adjoint_solver
if solver is None:
if self.assemble_system:
rhs_bcs_form = backend.inner(
backend.Function(self.function_space),
dFdu_adj_form.arguments()[0]) * backend.dx
A, _ = backend.assemble_system(dFdu_adj_form, rhs_bcs_form,
bcs)
else:
A = compat.assemble_adjoint_value(dFdu_adj_form)
[bc.apply(A) for bc in bcs]
solver = backend.LUSolver(A, self.method)
self.block_helper.adjoint_solver = solver
solver.parameters.update(self.lu_solver_parameters)
[bc.apply(dJdu) for bc in bcs]
adj_sol = backend.Function(self.function_space)
solver.solve(adj_sol.vector(), dJdu)
adj_sol_bdy = None
if compute_bdy:
adj_sol_bdy = compat.function_from_vector(
self.function_space, dJdu_copy - compat.assemble_adjoint_value(
backend.action(dFdu_adj_form, adj_sol)))
return adj_sol, adj_sol_bdy
def _assemble_and_solve_adj_eq(self, dFdu_adj_form, dJdu, compute_bdy):
dJdu_copy = dJdu.copy()
bcs = self._homogenize_bcs()
if self.assemble_system:
rhs_bcs_form = backend.inner(
backend.Function(self.function_space),
dFdu_adj_form.arguments()[0]) * backend.dx
A, _ = backend.assemble_system(dFdu_adj_form, rhs_bcs_form, bcs)
else:
A = backend.assemble(dFdu_adj_form)
[bc.apply(A) for bc in bcs]
[bc.apply(dJdu) for bc in bcs]
adj_sol = compat.create_function(self.function_space)
compat.linalg_solve(A, adj_sol.vector(), dJdu, *self.adj_args,
**self.adj_kwargs)
adj_sol_bdy = None
if compute_bdy:
adj_sol_bdy = compat.function_from_vector(
self.function_space, dJdu_copy - compat.assemble_adjoint_value(
backend.action(dFdu_adj_form, adj_sol)))
return adj_sol, adj_sol_bdy
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 _assemble_and_solve_adj_eq(self, dFdu_adj_form, dJdu, compute_bdy):
dJdu_copy = dJdu.copy()
bcs = self._homogenize_bcs()
solver = self.block_helper.adjoint_solver
if solver is None:
solver = backend.PETScKrylovSolver(self.method, self.preconditioner)
solver.ksp().setOptionsPrefix(self.ksp_options_prefix)
solver.set_from_options()
if self.assemble_system:
rhs_bcs_form = backend.inner(backend.Function(self.function_space),
dFdu_adj_form.arguments()[0]) * backend.dx
A, _ = backend.assemble_system(dFdu_adj_form, rhs_bcs_form, bcs)
if self._ad_nullspace is not None:
as_backend_type(A).set_nullspace(self._ad_nullspace)
if self.pc_operator is not None:
P = self._replace_form(self.pc_operator)
P, _ = backend.assemble_system(P, rhs_bcs_form, bcs)
solver.set_operators(A, P)
else:
solver.set_operator(A)
else:
A = compat.assemble_adjoint_value(dFdu_adj_form)
[bc.apply(A) for bc in bcs]
if self._ad_nullspace is not None:
as_backend_type(A).set_nullspace(self._ad_nullspace)
if self.pc_operator is not None:
P = self._replace_form(self.pc_operator)
P = compat.assemble_adjoint_value(P)
[bc.apply(P) for bc in bcs]
solver.set_operators(A, P)
else:
solver.set_operator(A)
self.block_helper.adjoint_solver = solver
solver.parameters.update(self.krylov_solver_parameters)
[bc.apply(dJdu) for bc in bcs]
if self._ad_nullspace is not None:
if self._ad_nullspace._ad_orthogonalized:
self._ad_nullspace.orthogonalize(dJdu)
adj_sol = backend.Function(self.function_space)
solver.solve(adj_sol.vector(), dJdu)
adj_sol_bdy = None
if compute_bdy:
adj_sol_bdy = compat.function_from_vector(self.function_space, dJdu_copy - compat.assemble_adjoint_value(
backend.action(dFdu_adj_form, adj_sol)))
return adj_sol, adj_sol_bdy
def jacobian_action(self, m, dm, result):
"""Computes the Jacobian action of c(m) in direction dm and stores the result in result. """
if isinstance(m, list):
assert len(m) == 1
m = m[0]
self.update_control(m)
form = backend.action(self.dform, dm)
result.assign(backend.assemble(form))
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 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
diag_name = "Split:%s:" % idx + hashlib.md5(str(eq_lhs) + "split" + str(smallfn) + str(bigfn) + str(idx) + str(random.random())).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 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 _create_F_form(self):
# Process the equation forms, replacing values with checkpoints,
# and gathering lhs and rhs in one single form.
if self.linear:
tmp_u = compat.create_function(self.function_space)
F_form = backend.action(self.lhs, tmp_u) - self.rhs
else:
tmp_u = self.func
F_form = self.lhs
replace_map = self._replace_map(F_form)
replace_map[tmp_u] = self.get_outputs()[0].saved_output
return ufl.replace(F_form, replace_map)
def define_nonlinear_equation(F, u):
# Given an F := 0,
# we write the equation for libadjoint's annotation purposes as
# M.u = M.u + F(u)
# as we need to have something on the diagonal in our big time system
fn_space = u.function_space()
test = backend.TestFunction(fn_space)
trial = backend.TrialFunction(fn_space)
mass = backend.inner(test, trial) * backend.dx
return (mass, backend.action(mass, u) - F)
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 mult(self, *args):
shapes = self.shape(self.current_form)
y = backend.PETScVector(shapes[0])
action_fn = backend.Function(ufl.algorithms.extract_arguments(self.current_form)[-1].function_space())
action_vec = action_fn.vector()
for i in range(len(args[0])):
action_vec[i] = args[0][i]
action_form = backend.action(self.current_form, action_fn)
backend.assemble(action_form, tensor=y)
for bc in self.bcs:
bcvals = bc.get_boundary_values()
for idx in bcvals:
y[idx] = action_vec[idx]
args[1].set_local(y.array())
def mult(self, *args):
shapes = self.shape(self.current_form)
y = backend.PETScVector(shapes[0])
action_fn = backend.Function(
ufl.algorithms.extract_arguments(
self.current_form)[-1].function_space())
action_vec = action_fn.vector()
for i in range(len(args[0])):
action_vec[i] = args[0][i]
action_form = backend.action(self.current_form, action_fn)
backend.assemble(action_form, tensor=y)
for bc in self.bcs:
bcvals = bc.get_boundary_values()
for idx in bcvals:
y[idx] = action_vec[idx]
args[1].set_local(y.array())
def _assemble_and_solve_adj_eq(self, dFdu_adj_form, dJdu, compute_bdy):
dJdu_copy = dJdu.copy()
kwargs = self.assemble_kwargs.copy()
# Homogenize and apply boundary conditions on adj_dFdu and dJdu.
bcs = self._homogenize_bcs()
kwargs["bcs"] = bcs
dFdu = compat.assemble_adjoint_value(dFdu_adj_form, **kwargs)
for bc in bcs:
bc.apply(dJdu)
adj_sol = compat.create_function(self.function_space)
compat.linalg_solve(dFdu, adj_sol.vector(), dJdu, *self.adj_args,
**self.adj_kwargs)
adj_sol_bdy = None
if compute_bdy:
adj_sol_bdy = compat.function_from_vector(
self.function_space, dJdu_copy - compat.assemble_adjoint_value(
backend.action(dFdu_adj_form, adj_sol)))
return adj_sol, adj_sol_bdy
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 evaluate_hessian_component(self,
inputs,
hessian_inputs,
adj_inputs,
block_variable,
idx,
relevant_dependencies,
prepared=None):
c = block_variable.output
if c == self.func and not self.linear:
return None
adj_sol2 = prepared["adj_sol2"]
adj_sol2_bdy = prepared["adj_sol2_bdy"]
F_form = prepared["form"]
adj_sol = prepared["adj_sol"]
fwd_block_variable = self.get_outputs()[0]
tlm_output = fwd_block_variable.tlm_value
c_rep = block_variable.saved_output
# If m = DirichletBC then d^2F(u,m)/dm^2 = 0 and d^2F(u,m)/dudm = 0,
# so we only have the term dF(u,m)/dm * adj_sol2
if isinstance(c, backend.DirichletBC):
tmp_bc = compat.create_bc(c,
value=extract_subfunction(
adj_sol2_bdy, c.function_space()))
return [tmp_bc]
if isinstance(c_rep, backend.Constant):
mesh = compat.extract_mesh_from_form(F_form)
W = c._ad_function_space(mesh)
elif isinstance(c, compat.ExpressionType):
mesh = F_form.ufl_domain().ufl_cargo()
W = c._ad_function_space(mesh)
elif isinstance(c, compat.MeshType):
X = backend.SpatialCoordinate(c)
element = X.ufl_domain().ufl_coordinate_element()
W = backend.FunctionSpace(c, element)
else:
W = c.function_space()
dc = backend.TestFunction(W)
form_adj = backend.action(F_form, adj_sol)
form_adj2 = backend.action(F_form, adj_sol2)
if isinstance(c, compat.MeshType):
dFdm_adj = backend.derivative(form_adj, X, dc)
dFdm_adj2 = backend.derivative(form_adj2, X, dc)
else:
dFdm_adj = backend.derivative(form_adj, c_rep, dc)
dFdm_adj2 = backend.derivative(form_adj2, c_rep, dc)
# TODO: Old comment claims this might break on split. Confirm if true or not.
d2Fdudm = ufl.algorithms.expand_derivatives(
backend.derivative(dFdm_adj, fwd_block_variable.saved_output,
tlm_output))
hessian_output = 0
# We need to add terms from every other dependency
# i.e. the terms d^2F/dm_1dm_2
for _, bv in relevant_dependencies:
c2 = bv.output
c2_rep = bv.saved_output
if isinstance(c2, backend.DirichletBC):
continue
tlm_input = bv.tlm_value
if tlm_input is None:
continue
if c2 == self.func and not self.linear:
continue
# TODO: If tlm_input is a Sum, this crashes in some instances?
if isinstance(c2_rep, compat.MeshType):
X = backend.SpatialCoordinate(c2_rep)
d2Fdm2 = ufl.algorithms.expand_derivatives(
backend.derivative(dFdm_adj, X, tlm_input))
else:
d2Fdm2 = ufl.algorithms.expand_derivatives(
backend.derivative(dFdm_adj, c2_rep, tlm_input))
if d2Fdm2.empty():
continue
hessian_output -= compat.assemble_adjoint_value(d2Fdm2)
if not d2Fdudm.empty():
# FIXME: This can be empty in the multimesh case, ask sebastian
hessian_output -= compat.assemble_adjoint_value(d2Fdudm)
hessian_output -= compat.assemble_adjoint_value(dFdm_adj2)
if isinstance(c, compat.ExpressionType):
return [(hessian_output, W)]
else:
return hessian_output