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