tlm_adjoint.equation

Module Contents

class tlm_adjoint.equation.Equation(X, deps, nl_deps=None, *, ic_deps=None, ic=None, adj_ic_deps=None, adj_ic=None, adj_type='conjugate_dual')

Core equation class. Defines a differentiable operation for use as an adjoint tape record.

The equation is defined via a residual function \(\mathcal{F}\). The forward solution is defined implicitly as the value \(x\) for which

\[\mathcal{F} \left( x, y_0, y_1, \ldots \right) = 0,\]

where \(y_i\) are dependencies.

This is an abstract base class. Information required to solve forward equations, perform adjoint calculations, and define tangent-linear equations, is provided by overloading abstract methods. This class does not inherit from abc.ABC, so that methods may be implemented as needed.

Parameters:

• X – A variable, or a Sequence of variables, defining the forward solution.

• deps – A Sequence of variables defining dependencies. Must define a superset of X.

• nl_deps – A Sequence of variables defining non-linear dependencies. Must define a subset of deps. Defaults to deps.

• ic_deps – A Sequence of variables whose value must be available prior to computing the forward solution. Intended for iterative methods with non-zero initial guesses. Must define a subset of X. Can be overridden by ic.

• ic – Whether ic_deps should be set equal to X. Defaults to True if ic_deps is not supplied, and to False otherwise.

• adj_ic_deps – A Sequence of variables whose value must be available prior to computing the adjoint solution. Intended for iterative methods with non-zero initial guesses. Must define a subset of X. Can be overridden by adj_ic.

• adj_ic – Whether adj_ic_deps should be set equal to X. Defaults to True if adj_ic_deps is not supplied, and to False otherwise.

• adj_type – The space type relative to X of adjoint variables. ‘primal’ or ‘conjugate_dual’, or a Sequence of these.

drop_references()

Drop references to variables which store values.

x()

Return the forward solution variable, assuming the forward solution has one component.

Returns:

A variable defining the forward solution.

X(m=None)

Return forward solution variables.

Returns:

If m is supplied, a variable defining the m th component of the forward solution. If m is not supplied, a tuple of variables defining the forward solution.

dependencies()

Return dependencies.

Returns:

A tuple of variables defining dependencies.

nonlinear_dependencies()

Return non-linear dependencies.

Returns:

A tuple of variables defining non-linear dependencies.

initial_condition_dependencies()

Return ‘initial condition’ dependencies – dependencies whose value is needed prior to computing the forward solution.

Returns:

A tuple of variables defining initial condition dependencies.

adjoint_initial_condition_dependencies()

Return adjoint ‘initial condition’ dependencies – dependencies whose value is needed prior to computing the adjoint solution.

Returns:

A tuple of variables defining adjoint initial condition dependencies.

adj_x_type()

Return the space type for the adjoint solution, relative to the forward solution, assuming the forward solution has exactly one component.

Returns:

One of ‘primal’ or ‘conjugate_dual’.

adj_X_type(m=None)

Return the space type for the adjoint solution, relative to the forward solution.

Returns:

If m is supplied, one of ‘primal’ or ‘conjugate_dual’ defining the relative space type for the m th component of the adjoint solution. If m is not supplied, a tuple whose elements are ‘primal’ or ‘conjugate_dual’, defining the relative space type of the adjoint solution.

new_adj_x()

Return a new variable suitable for storing the adjoint solution, assuming the forward solution has exactly one component.

Returns:

A variable suitable for storing the adjoint solution.

new_adj_X(m=None)

Return new variables suitable for storing the adjoint solution.

Returns:

If m is supplied, a variable suitable for storing the m th component of the adjoint solution. If m is not supplied, a tuple of variables suitable for storing the adjoint solution.

solve(*, annotate=None, tlm=None)

Compute the forward solution.

Parameters:

• annotate – Whether the EquationManager should record the solution of equations.

• tlm – Whether tangent-linear equations should be solved.

forward(X, deps=None)

Wraps Equation.forward_solve() to handle cache invalidation.

abstract forward_solve(X, deps=None)

Compute the forward solution.

Can assume that the currently active EquationManager is paused.

Parameters:

• X – A variable if the forward solution has a single component, otherwise a Sequence of variables. May define an initial guess, and should be set by this method. Subclasses may replace this argument with x if the forward solution has a single component.

• deps – A tuple of variables, defining values for dependencies. Only the elements corresponding to X may be modified. self.dependencies() should be used if not supplied.

adjoint(adj_X, nl_deps, B, dep_Bs)

Compute the adjoint solution, and subtract terms from other adjoint right-hand-sides.

Parameters:

• adj_X – Either None, or a Sequence of variables defining the initial guess for an iterative solve. May be modified or returned.

• nl_deps – A Sequence of variables defining values for non-linear dependencies. Should not be modified.

• B – A sequence of variables defining the right-hand-side of the adjoint equation. May be modified or returned.

• dep_Bs – A Mapping whose items are (dep_index, dep_B). Each dep_B is an AdjointRHS which should be updated by subtracting adjoint derivative information computed by differentiating with respect to self.dependencies()[dep_index].

Returns:

A tuple of variables defining the adjoint solution, or None to indicate that the solution is zero.

adjoint_cached(adj_X, nl_deps, dep_Bs)

Subtract terms from other adjoint right-hand-sides.

Parameters:

• adj_X – A Sequence of variables defining the adjoint solution. Should not be modified.

• nl_deps – A Sequence of variables defining values for non-linear dependencies. Should not be modified.

• dep_Bs – A Mapping whose items are (dep_index, dep_B). Each dep_B is an AdjointRHS which should be updated by subtracting adjoint derivative information computed by differentiating with respect to self.dependencies()[dep_index].

abstract adjoint_derivative_action(nl_deps, dep_index, adj_X)

Return the action of the adjoint of a derivative of the forward residual on the adjoint solution. This is the negative of an adjoint right-hand-side term.

Parameters:

• nl_deps – A Sequence of variables defining values for non-linear dependencies. Should not be modified.

• dep_index – An int. The derivative is defined by differentiation of the forward residual with respect to self.dependencies()[dep_index].

• adj_X – The adjoint solution. A variable if the adjoint solution has a single component, otherwise a Sequence of variables. Should not be modified. Subclasses may replace this argument with adj_x if the adjoint solution has a single component.

Returns:

The action of the adjoint of a derivative on the adjoint solution. Will be passed to subtract_adjoint_derivative_action(), and valid types depend upon the adjoint variable type. Typically this will be a variable, or a two element tuple (alpha, F), where alpha is a numbers.Complex and F a variable, with the value defined by the product of alpha and F.

subtract_adjoint_derivative_actions(adj_X, nl_deps, dep_Bs)

Subtract terms from other adjoint right-hand-sides.

Can be overridden for an optimized implementation, but otherwise uses Equation.adjoint_derivative_action().

Parameters:

• adj_X – The adjoint solution. A variable if the adjoint solution has a single component, otherwise a Sequence of variables. Should not be modified. Subclasses may replace this argument with adj_x if the adjoint solution has a single component.

• nl_deps – A Sequence of variables defining values for non-linear dependencies. Should not be modified.

• dep_Bs – A Mapping whose items are (dep_index, dep_B). Each dep_B is an AdjointRHS which should be updated by subtracting adjoint derivative information computed by differentiating with respect to self.dependencies()[dep_index].

abstract adjoint_jacobian_solve(adj_X, nl_deps, B)

Compute an adjoint solution.

Parameters:

• adj_X – Either None, or a variable (if the adjoint solution has a single component) or Sequence of variables (otherwise) defining the initial guess for an iterative solve. May be modified or returned. Subclasses may replace this argument with adj_x if the adjoint solution has a single component.

• nl_deps – A Sequence of variables defining values for non-linear dependencies. Should not be modified.

• B – The right-hand-side. A variable (if the adjoint solution has a single component) or Sequence of variables (otherwise) storing the value of the right-hand-side. May be modified or returned. Subclasses may replace this argument with b if the adjoint solution has a single component.

Returns:

A variable or Sequence of variables storing the value of the adjoint solution. May return None to indicate a value of zero.

abstract tangent_linear(tlm_map)

Derive an Equation corresponding to a tangent-linear operation.

Parameters:

tlm_map – A TangentLinearMap storing values for tangent-linear variables.

Returns:

An Equation, corresponding to the tangent-linear operation.

class tlm_adjoint.equation.ZeroAssignment(X)

Represents an assignment

\[x = 0.\]

The forward residual is defined

\[\mathcal{F} \left( x \right) = x.\]

Parameters:

X – A variable or a Sequence of variables defining the forward solution \(x\).

forward_solve(X, deps=None)

Compute the forward solution.

Can assume that the currently active EquationManager is paused.

Parameters:

• X – A variable if the forward solution has a single component, otherwise a Sequence of variables. May define an initial guess, and should be set by this method. Subclasses may replace this argument with x if the forward solution has a single component.

• deps – A tuple of variables, defining values for dependencies. Only the elements corresponding to X may be modified. self.dependencies() should be used if not supplied.

adjoint_jacobian_solve(adj_X, nl_deps, B)

Compute an adjoint solution.

Parameters:

• adj_X – Either None, or a variable (if the adjoint solution has a single component) or Sequence of variables (otherwise) defining the initial guess for an iterative solve. May be modified or returned. Subclasses may replace this argument with adj_x if the adjoint solution has a single component.

• nl_deps – A Sequence of variables defining values for non-linear dependencies. Should not be modified.

• B – The right-hand-side. A variable (if the adjoint solution has a single component) or Sequence of variables (otherwise) storing the value of the right-hand-side. May be modified or returned. Subclasses may replace this argument with b if the adjoint solution has a single component.

Returns:

A variable or Sequence of variables storing the value of the adjoint solution. May return None to indicate a value of zero.

tangent_linear(tlm_map)

Derive an Equation corresponding to a tangent-linear operation.

Parameters:

tlm_map – A TangentLinearMap storing values for tangent-linear variables.

Returns:

An Equation, corresponding to the tangent-linear operation.

class tlm_adjoint.equation.NullSolver(X)

Represents an assignment

\[x = 0.\]

The forward residual is defined

\[\mathcal{F} \left( x \right) = x.\]

Parameters:

X – A variable or a Sequence of variables defining the forward solution \(x\).