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 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