def compute_tlm(parameter, forget=False): if isinstance(parameter, (list, tuple)): parameter = ListControl(parameter) for i in range(adjglobals.adjointer.equation_count): (tlm_var, output) = adjglobals.adjointer.get_tlm_solution(i, parameter) if output.data: output.data.rename(str(tlm_var), "a Function from dolfin-adjoint") storage = libadjoint.MemoryStorage(output) storage.set_overwrite(True) adjglobals.adjointer.record_variable(tlm_var, storage) yield (output.data, tlm_var) # forget is None: forget *nothing*. # forget is True: forget everything we can, forward and adjoint # forget is False: forget only unnecessary tlm values if forget is None: pass elif forget: adjglobals.adjointer.forget_tlm_equation(i) else: adjglobals.adjointer.forget_tlm_values(i)
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 __call__(self, infunc, annotate=None): # Perform the forward operation out = Function(infunc.function_space()) out.vector()[:] = out.vector().get_local()**2 # Annotate the operation on the dolfin-adjoint tape if utils.to_annotate(annotate): rhs = CustomDolfinAdjointFunctionRhs(coeffs=[infunc]) out_dep = adjglobals.adj_variables.next(out) solving.register_initial_conditions(zip(rhs.coefficients(), rhs.dependencies()), linear=False) if parameters["adjoint"]["record_all"]: adjglobals.adjointer.record_variable( out_dep, libadjoint.MemoryStorage(adjlinalg.Vector(infunc))) identity = utils.get_identity_block(out.function_space()) eq = libadjoint.Equation(out_dep, blocks=[identity], targets=[out_dep], rhs=rhs) cs = adjglobals.adjointer.register_equation(eq) solving.do_checkpoint(cs, out_dep, rhs) return out
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 dolfin_adjoint_interpolate(self, other, annotate=None): out = dolfin_interpolate(self, other) if annotate is True: assignment.register_assign(self, other, op=backend.interpolate) adjglobals.adjointer.record_variable( adjglobals.adj_variables[self], libadjoint.MemoryStorage(adjlinalg.Vector(self))) return out
def solve(self, *args, **kwargs): # Annotate an assignment solve record = da_annotate_equation_solve(self) annotate = not dolfin.parameters["adjoint"]["stop_annotating"] dolfin.parameters["adjoint"]["stop_annotating"] = True ret = solve_orig(self, *args, **kwargs) dolfin.parameters["adjoint"]["stop_annotating"] = not annotate if record: x = unwrap_fns(self.x()) adjglobals.adjointer.record_variable( adjglobals.adj_variables[x], libadjoint.MemoryStorage(adjlinalg.Vector(x))) return ret
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) out = utils.function_to_da_function(out) if name is not None: out.adj_name = name to_annotate = utils.to_annotate(annotate) if to_annotate: if isinstance(v, backend.Function): 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) elif annotate is not None: outvar = adjglobals.adj_variables[out] solving.register_initial_conditions([ [out, outvar], ], linear=True) return out
def tlm_dolfin(control, forget=False): for i in range(adjglobals.adjointer.equation_count): (tlm_var, output) = adjglobals.adjointer.get_tlm_solution(i, control) storage = libadjoint.MemoryStorage(output) storage.set_overwrite(True) adjglobals.adjointer.record_variable(tlm_var, storage) if forget is None: pass elif forget: adjglobals.adjointer.forget_tlm_equation(i) else: adjglobals.adjointer.forget_tlm_values(i) return output
def compute_adjoint(functional, forget=True, ignore=[]): ignorelist = [] for fn in ignore: if isinstance(fn, backend.Function): ignorelist.append(adjglobals.adj_variables[fn]) elif isinstance(fn, str): ignorelist.append(libadjoint.Variable(fn, 0, 0)) else: ignorelist.append(fn) for i in range(adjglobals.adjointer.timestep_count): adjglobals.adjointer.set_functional_dependencies(functional, i) for i in range(adjglobals.adjointer.equation_count)[::-1]: fwd_var = adjglobals.adjointer.get_forward_variable(i) if fwd_var in ignorelist: info("Ignoring the adjoint equation for %s" % fwd_var) continue (adj_var, output) = adjglobals.adjointer.get_adjoint_solution(i, functional) if output.data: if backend.__name__ == "dolfin": output.data.rename(str(adj_var), "a Function from dolfin-adjoint") else: output.data.name = str(adj_var) storage = libadjoint.MemoryStorage(output) storage.set_overwrite(True) adjglobals.adjointer.record_variable(adj_var, storage) # forget is None: forget *nothing*. # forget is True: forget everything we can, forward and adjoint # forget is False: forget only unnecessary adjoint values if forget is None: pass elif forget: adjglobals.adjointer.forget_adjoint_equation(i) else: adjglobals.adjointer.forget_adjoint_values(i) yield (output.data, adj_var)
def annotate_slope_limit(f): # First annotate the equation adj_var = adjglobals.adj_variables[f] rhs = SlopeRHS(f) adj_var_next = adjglobals.adj_variables.next(f) identity_block = solving.get_identity(f.function_space()) eq = libadjoint.Equation(adj_var_next, blocks=[identity_block], targets=[adj_var_next], rhs=rhs) cs = adjglobals.adjointer.register_equation(eq) # Record the result adjglobals.adjointer.record_variable( adjglobals.adj_variables[f], libadjoint.MemoryStorage(adjlinalg.Vector(f)))
def solve(self, *args, **kwargs): # Annotate an equation solve record = da_annotate_equation_solve(self) # The EquationSolver solve method could do a lot of work, including # further solves. Temporarily disable annotation during the solve ... annotate = not dolfin.parameters["adjoint"]["stop_annotating"] dolfin.parameters["adjoint"]["stop_annotating"] = True ret = solve_orig(self, *args, **kwargs) # ... and then restore the previous annotation status. dolfin.parameters["adjoint"]["stop_annotating"] = not annotate if record: # We still want to wrap functions with WrappedFunction so that we can for # example write separate equations for u[0] and u[n]. However # dolfin-adjoint needs to treat these as the same Function, so # strategically unwrap WrappedFunction s. x = unwrap_fns(self.x()) adjglobals.adjointer.record_variable( adjglobals.adj_variables[x], libadjoint.MemoryStorage(adjlinalg.Vector(x))) return ret
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 register_assign(new, old, op=None): if not isinstance(old, backend.Function): assert op is not None fn_space = new.function_space() identity_block = utils.get_identity_block(fn_space) dep = adjglobals.adj_variables.next(new) if backend.parameters["adjoint"]["record_all"] and isinstance( old, backend.Function): adjglobals.adjointer.record_variable( dep, libadjoint.MemoryStorage(adjlinalg.Vector(old))) rhs = IdentityRHS(old, fn_space, op) 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)
def replay_dolfin(forget=False, tol=0.0, stop=False): with misc.annotations(False): if not backend.parameters["adjoint"]["record_all"]: info_red( "Warning: your replay test will be much more effective with dolfin.parameters['adjoint']['record_all'] = True." ) success = True for i in range(adjglobals.adjointer.equation_count): (fwd_var, output) = adjglobals.adjointer.get_forward_solution(i) storage = libadjoint.MemoryStorage(output) storage.set_compare(tol=tol) storage.set_overwrite(True) out = adjglobals.adjointer.record_variable(fwd_var, storage) success = success and out if forget: adjglobals.adjointer.forget_forward_equation(i) if not out and stop: break return success
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 solve(self, *args, **kwargs): timer = backend.Timer("Matrix-free solver") annotate = True if "annotate" in kwargs: annotate = kwargs["annotate"] del kwargs["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] if annotate: if not isinstance(A, AdjointKrylovMatrix): try: A = AdjointKrylovMatrix(A.form) except AttributeError: raise libadjoint.exceptions.LibadjointErrorInvalidInputs( "Your A has to either be an AdjointKrylovMatrix or have been assembled after backend_adjoint was imported." ) if not hasattr(x, 'function'): raise libadjoint.exceptions.LibadjointErrorInvalidInputs( "Your x has to come from code like down_cast(my_function.vector())." ) if not hasattr(b, 'form'): raise libadjoint.exceptions.LibadjointErrorInvalidInputs( "Your b has to have the .form attribute: was it assembled with from backend_adjoint import *?" ) if not hasattr(A, 'dependencies'): backend.info_red( "A has no .dependencies method; assuming no nonlinear dependencies of the matrix-free operator." ) coeffs = [] dependencies = [] else: coeffs = [ coeff for coeff in A.dependencies() if hasattr(coeff, 'function_space') ] dependencies = [ adjglobals.adj_variables[coeff] for coeff in coeffs ] if len(dependencies) > 0: assert hasattr( A, "set_dependencies" ), "Need a set_dependencies method to replace your values, if you have nonlinear dependencies ... " rhs = adjrhs.RHS(b.form) key = '{}{}'.format(hash(A), random.random()).encode('utf8') diag_name = hashlib.md5(key).hexdigest() diag_block = libadjoint.Block( diag_name, dependencies=dependencies, test_hermitian=backend.parameters["adjoint"]["test_hermitian"], test_derivative=backend.parameters["adjoint"] ["test_derivative"]) solving.register_initial_conditions( zip(rhs.coefficients(), rhs.dependencies()) + zip(coeffs, dependencies), linear=False, var=None) var = adjglobals.adj_variables.next(x.function) frozen_expressions_dict = expressions.freeze_dict() frozen_parameters = self.parameters.to_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) if len(dependencies) > 0: A.set_dependencies(dependencies, [val.data for val in values]) if hermitian: A_transpose = A.hermitian() return (MatrixFree( A_transpose, fn_space=x.function.function_space(), bcs=A_transpose.bcs, solver_parameters=self.solver_parameters, operators=transpose_operators(self.operators), parameters=frozen_parameters), adjlinalg.Vector( None, fn_space=x.function.function_space())) else: return (MatrixFree( A, fn_space=x.function.function_space(), bcs=b.bcs, solver_parameters=self.solver_parameters, operators=self.operators, parameters=frozen_parameters), adjlinalg.Vector( None, fn_space=x.function.function_space())) diag_block.assemble = diag_assembly_cb def diag_action_cb(dependencies, values, hermitian, coefficient, input, context): expressions.update_expressions(frozen_expressions_dict) A.set_dependencies(dependencies, [val.data for val in values]) if hermitian: acting_mat = A.transpose() else: acting_mat = A output_fn = backend.Function(input.data.function_space()) acting_mat.mult(input.data.vector(), output_fn.vector()) vec = output_fn.vector() for i in range(len(vec)): vec[i] = coefficient * vec[i] return adjlinalg.Vector(output_fn) diag_block.action = diag_action_cb if len(dependencies) > 0: def derivative_action(dependencies, values, variable, contraction_vector, hermitian, input, coefficient, context): expressions.update_expressions(frozen_expressions_dict) A.set_dependencies(dependencies, [val.data for val in values]) action = A.derivative_action( values[dependencies.index(variable)].data, contraction_vector.data, hermitian, input.data, coefficient) return adjlinalg.Vector(action) diag_block.derivative_action = derivative_action eqn = libadjoint.Equation(var, blocks=[diag_block], targets=[var], rhs=rhs) cs = adjglobals.adjointer.register_equation(eqn) solving.do_checkpoint(cs, var, rhs) out = backend.PETScKrylovSolver.solve(self, *args) if annotate: if backend.parameters["adjoint"]["record_all"]: adjglobals.adjointer.record_variable( var, libadjoint.MemoryStorage(adjlinalg.Vector(x.function))) timer.stop() 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) == 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 KrylovSolver 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(krylov_solvers) krylov_solvers.append(self) adj_krylov_solvers.append(None) idx = self.__global_list_idx__ class KrylovSolverMatrix(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 = krylov_solvers[idx] need_to_set_operator = self._need_to_reset_operator else: if adj_krylov_solvers[idx] is None: need_to_set_operator = True adj_krylov_solvers[idx] = KrylovSolver( *solver_parameters) else: need_to_set_operator = self._need_to_reset_operator solver = adj_krylov_solvers[idx] 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) # 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) solver.solve(x.vector(), rhs) return adjlinalg.Vector(x) solving.annotate(A == b, u, bcs, matrix_class=KrylovSolverMatrix, initial_guess=parameters['nonzero_initial_guess'], replace_map=True) out = dolfin.KrylovSolver.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 __call__(self, value): """ Evaluates the reduced functional for the given control value. Args: value: The point in control space where to perform the Taylor test. Must be of the same type as the Control (e.g. Function, Constant or lists of latter). Returns: float: The functional value. """ # Make sure we do not annotate # Reset any cached data in dolfin-adjoint adj_reset_cache() #: The control values at which the reduced functional is to be evaluated. value = enlist(value) # Call callback self.eval_cb_pre(delist(value, list_type=self.controls)) # Update the control values on the tape ListControl(self.controls).update(value) # Check if the result is already cached if self.cache: hash = value_hash(value) if hash in self._cache["functional_cache"]: # Found a cache info_green("Got a functional cache hit") return self._cache["functional_cache"][hash] # Replay the annotation and evaluate the functional func_value = 0. for i in range(adjointer.equation_count): (fwd_var, output) = adjointer.get_forward_solution(i) if isinstance(output.data, Function): output.data.rename(str(fwd_var), "a Function from dolfin-adjoint") # Call callback self.replay_cb(fwd_var, output.data, delist(value, list_type=self.controls)) # Check if we checkpointing is active and if yes # record the exact same checkpoint variables as # in the initial forward run if adjointer.get_checkpoint_strategy() != None: if str(fwd_var) in mem_checkpoints: storage = libadjoint.MemoryStorage(output, cs=True) storage.set_overwrite(True) adjointer.record_variable(fwd_var, storage) if str(fwd_var) in disk_checkpoints: storage = libadjoint.MemoryStorage(output) adjointer.record_variable(fwd_var, storage) storage = libadjoint.DiskStorage(output, cs=True) storage.set_overwrite(True) adjointer.record_variable(fwd_var, storage) if not str(fwd_var) in mem_checkpoints and not str( fwd_var) in disk_checkpoints: storage = libadjoint.MemoryStorage(output) storage.set_overwrite(True) adjointer.record_variable(fwd_var, storage) # No checkpointing, so we record everything else: storage = libadjoint.MemoryStorage(output) storage.set_overwrite(True) adjointer.record_variable(fwd_var, storage) if i == adjointer.timestep_end_equation(fwd_var.timestep): func_value += adjointer.evaluate_functional( self.functional, fwd_var.timestep) if adjointer.get_checkpoint_strategy() != None: adjointer.forget_forward_equation(i) self.current_func_value = func_value # Call callback self.eval_cb_post(self.scale * func_value, delist(value, list_type=self.controls)) if self.cache: # Add result to cache info_red("Got a functional cache miss") self._cache["functional_cache"][hash] = self.scale * func_value return self.scale * func_value
def project_dolfin(v, V=None, bcs=None, mesh=None, solver_type="lu", 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(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 __call__(self, m_dot, project=False): flag = misc.pause_annotation() hess_action_timer = backend.Timer("Hessian action") m_p = self.m.set_perturbation(m_dot) last_timestep = adjglobals.adjointer.timestep_count m_dot = enlist(m_dot) Hm = [] for m_dot_cmp in m_dot: if hasattr(m_dot_cmp, 'function_space'): Hm.append(backend.Function(m_dot_cmp.function_space())) elif isinstance(m_dot_cmp, float): Hm.append(0.0) else: raise NotImplementedError( "Sorry, don't know how to handle this") tlm_timer = backend.Timer("Hessian action (TLM)") # run the tangent linear model for (tlm, tlm_var) in compute_tlm(m_p, forget=None): pass tlm_timer.stop() # run the adjoint and second-order adjoint equations. for i in range(adjglobals.adjointer.equation_count)[::-1]: adj_var = adjglobals.adjointer.get_forward_variable(i).to_adjoint( self.J) # Only recompute the adjoint variable if we do not have it yet try: adj = adjglobals.adjointer.get_variable_value(adj_var) except (libadjoint.exceptions.LibadjointErrorHashFailed, libadjoint.exceptions.LibadjointErrorNeedValue): adj_timer = backend.Timer("Hessian action (ADM)") adj = adjglobals.adjointer.get_adjoint_solution(i, self.J)[1] adj_timer.stop() storage = libadjoint.MemoryStorage(adj) adjglobals.adjointer.record_variable(adj_var, storage) adj = adj.data soa_timer = backend.Timer("Hessian action (SOA)") (soa_var, soa_vec) = adjglobals.adjointer.get_soa_solution(i, self.J, m_p) soa_timer.stop() soa = soa_vec.data def hess_inner(Hm, out): assert len(out) == len(Hm) for i in range(len(out)): if out[i] is not None: if isinstance(Hm[i], backend.Function): Hm[i].vector().axpy(1.0, out[i].vector()) elif isinstance(Hm[i], float): Hm[i] += out[i] else: raise ValueError( "Do not know what to do with this") return Hm func_timer = backend.Timer("Hessian action (derivative formula)") # now implement the Hessian action formula. out = self.m.equation_partial_derivative(adjglobals.adjointer, soa, i, soa_var.to_forward()) Hm = hess_inner(Hm, out) out = self.m.equation_partial_second_derivative( adjglobals.adjointer, adj, i, soa_var.to_forward(), m_dot) Hm = hess_inner(Hm, out) if last_timestep > adj_var.timestep: # We have hit a new timestep, and need to compute this timesteps' \partial^2 J/\partial m^2 contribution last_timestep = adj_var.timestep out = self.m.functional_partial_second_derivative( adjglobals.adjointer, self.J, adj_var.timestep, m_dot) Hm = hess_inner(Hm, out) func_timer.stop() storage = libadjoint.MemoryStorage(soa_vec) storage.set_overwrite(True) adjglobals.adjointer.record_variable(soa_var, storage) for Hm_cmp in Hm: if isinstance(Hm_cmp, backend.Function): Hm_cmp.rename("d^2(%s)/d(%s)^2" % (str(self.J), str(self.m)), "a Function from dolfin-adjoint") misc.continue_annotation(flag) return postprocess(Hm, project, list_type=self.enlisted_controls)
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
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 perturbed_replay(parameter, perturbation, perturbation_scale, observation, perturbation_norm="mass", observation_norm="mass", callback=None, forget=False): r"""Perturb the forward run and compute .. math:: \frac{ \left|\left| \delta \mathrm{observation} \right|\right| }{ \left|\left| \delta \mathrm{input} \right| \right| } as a function of time. :py:data:`parameter` -- an FunctionControl to say what variable should be perturbed (e.g. FunctionControl('InitialConcentration')) :py:data:`perturbation` -- a Function to give the perturbation direction (from a GST analysis, for example) :py:data:`perturbation_norm` -- a bilinear Form which induces a norm on the space of perturbation inputs :py:data:`perturbation_scale` -- how big the norm of the initial perturbation should be :py:data:`observation` -- the variable to observe (e.g. 'Concentration') :py:data:`observation_norm` -- a bilinear Form which induces a norm on the space of perturbation outputs :py:data:`callback` -- a function f(var, perturbed, unperturbed) that the user can supply (e.g. to dump out variables during the perturbed replay) """ if not backend.parameters["adjoint"]["record_all"]: info_red( "Warning: your replay test will be much more effective with backend.parameters['adjoint']['record_all'] = True." ) assert isinstance(parameter, controls.FunctionControl) if perturbation_norm == "mass": p_fnsp = perturbation.function_space() u = backend.TrialFunction(p_fnsp) v = backend.TestFunction(p_fnsp) p_mass = backend.inner(u, v) * backend.dx perturbation_norm = p_mass if not isinstance(perturbation_norm, backend.GenericMatrix): perturbation_norm = backend.assemble(perturbation_norm) if not isinstance(observation_norm, backend.GenericMatrix) and observation_norm != "mass": observation_norm = backend.assemble(observation_norm) def compute_norm(perturbation, norm): # Need to compute <x, Ax> and then take its sqrt # where x is perturbation, A is norm try: vec = perturbation.vector() except: vec = perturbation Ax = vec.copy() norm.mult(vec, Ax) xAx = vec.inner(Ax) return math.sqrt(xAx) growths = [] for i in range(adjglobals.adjointer.equation_count): (fwd_var, output) = adjglobals.adjointer.get_forward_solution(i) if fwd_var == parameter.var: # we've hit the initial condition we want to perturb current_norm = compute_norm(perturbation, perturbation_norm) output.data.vector()[:] += (perturbation_scale / current_norm) * perturbation.vector() unperturbed = adjglobals.adjointer.get_variable_value(fwd_var).data if fwd_var.name == observation: # we've hit something we want to observe # Fetch the unperturbed result from the record if observation_norm == "mass": # we can't do this earlier, because we don't have the observation function space yet o_fnsp = output.data.function_space() u = backend.TrialFunction(o_fnsp) v = backend.TestFunction(o_fnsp) o_mass = backend.inner(u, v) * backend.dx observation_norm = backend.assemble(o_mass) diff = output.data.vector() - unperturbed.vector() growths.append( compute_norm(diff, observation_norm) / perturbation_scale) # <--- the action line if callback is not None: callback(fwd_var, output.data, unperturbed) storage = libadjoint.MemoryStorage(output) storage.set_compare(tol=None) storage.set_overwrite(True) out = adjglobals.adjointer.record_variable(fwd_var, storage) if forget: adjglobals.adjointer.forget_forward_equation(i) # can happen if we initialised a nonlinear solve with a constant zero guess if growths[0] == 0.0: return growths[1:] else: return growths
def compute_gradient(J, param, forget=True, ignore=[], callback=lambda var, output: None, project=False): if not isinstance(J, Functional): raise ValueError("J must be of type dolfin_adjoint.Functional.") flag = misc.pause_annotation() enlisted_controls = enlist(param) param = ListControl(enlisted_controls) if backend.parameters["adjoint"]["allow_zero_derivatives"]: dJ_init = [] for c in enlisted_controls: if isinstance(c.data(), backend.Constant): dJ_init.append(backend.Constant(0)) elif isinstance(c.data(), backend.Function): space = c.data().function_space() dJ_init.append(backend.Function(space)) else: dJ_init = [None] * len(enlisted_controls) dJdparam = enlisted_controls.__class__(dJ_init) last_timestep = adjglobals.adjointer.timestep_count ignorelist = [] for fn in ignore: if isinstance(fn, backend.Function): ignorelist.append(adjglobals.adj_variables[fn]) elif isinstance(fn, str): ignorelist.append(libadjoint.Variable(fn, 0, 0)) else: ignorelist.append(fn) for i in range(adjglobals.adjointer.timestep_count): adjglobals.adjointer.set_functional_dependencies(J, i) for i in range(adjglobals.adjointer.equation_count)[::-1]: fwd_var = adjglobals.adjointer.get_forward_variable(i) if fwd_var in ignorelist: info("Ignoring the adjoint equation for %s" % fwd_var) continue (adj_var, output) = adjglobals.adjointer.get_adjoint_solution(i, J) callback(adj_var, output.data) storage = libadjoint.MemoryStorage(output) storage.set_overwrite(True) adjglobals.adjointer.record_variable(adj_var, storage) fwd_var = libadjoint.Variable(adj_var.name, adj_var.timestep, adj_var.iteration) out = param.equation_partial_derivative(adjglobals.adjointer, output.data, i, fwd_var) dJdparam = _add(dJdparam, out) if last_timestep > adj_var.timestep: # We have hit a new timestep, and need to compute this timesteps' \partial J/\partial m contribution out = param.functional_partial_derivative(adjglobals.adjointer, J, adj_var.timestep) dJdparam = _add(dJdparam, out) last_timestep = adj_var.timestep if forget is None: pass elif forget: adjglobals.adjointer.forget_adjoint_equation(i) else: adjglobals.adjointer.forget_adjoint_values(i) rename(J, dJdparam, param) misc.continue_annotation(flag) return postprocess(dJdparam, project, list_type=enlisted_controls)
J = Functional(inner(sigma0[2], sigma0[2]) * dx * dtm[FINISH_TIME]) param = ScalarParameter(amplitude) adjointer = solving.adjointer # Copy the code from compute_gradient: #dJdp = compute_gradient(J, ScalarParameter(amplitude)) # as we want to also dump out norms of the adjoint solution. # -------------------------------------------------------------- norm0s = [] norm1s = [] dJdparam = None for i in range(adjointer.equation_count)[::-1]: (adj_var, output) = adjointer.get_adjoint_solution(i, J) storage = libadjoint.MemoryStorage(output) adjointer.record_variable(adj_var, storage) fwd_var = libadjoint.Variable(adj_var.name, adj_var.timestep, adj_var.iteration) afile = File("%s/adjoint_%s_%d.xml" % (dirname, adj_var.name, i)) afile << output.data out = param.inner_adjoint(adjointer, output.data, i, fwd_var) if dJdparam is None: dJdparam = out elif dJdparam is not None and out is not None: dJdparam += out if adj_var.name == "w_{10}": (adj_sigma0, adj_sigma1, adj_v, adj_gamma) = output.data.split()