Automatic differentiation with the autodiff sub-package

Automatic differentiaton (AD) refers to a diverse set of techniques used to numerically evaluate the derivatives of arbitrary (analytical) functions. NITROGEN requires access to such derivatives for a variety of tasks. Although the DFun class defines an interface to differentiable functions and some high-level operations such as composition, it does not prescribe an implementation. In many built-in NITROGEN objects, derivatives are computed via “forward accumulation” AD module (forward) of the nitrogen.autodiff sub-package. This chapter provides a guide to the design and usage of this module.

A brief introduction to AD

Forward AD relies on the accumulative use of the chain rule, product rule, and Taylor series to evaluate the derivatives of analytical expressions in terms of the derivatives of their arguments. Given an expression \(f\) with derivatives \(\partial^\alpha f\) (using multi-index notation), we begin with some simple rules from the linearity of differentiation:

\[ \begin{align}\begin{aligned}w &= af &\rightarrow \partial^\alpha w &= a \left(\partial^\alpha f\right)\\w &= f + g &\rightarrow \partial^\alpha w &= \partial^\alpha f + \partial^\alpha g,\end{aligned}\end{align} \]

where \(a\) is a scalar constant. It is convenient to introduce a scaled derivative,

\[f^{(\alpha)} = \frac{1}{\alpha !} \partial^\alpha f,\]

which renders the above rules unchanged,

\[ \begin{align}\begin{aligned}(af)^{(\alpha)}&= a f^{(\alpha)}\\(f + g)^{(\alpha)} &= f^{(\alpha)} + g^{(\alpha)}.\end{aligned}\end{align} \]

Another fundamental result is the generalized product rule, or Leibniz formula, for multi-variable functions,

\[(fg)^{(\alpha)} = \sum_{\beta \leq \alpha} f^{(\beta)} g^{(\alpha-\beta)},\]

where \(\beta \leq \alpha\) if and only if \(\beta_i \leq \alpha_i\) for all \(i\). (The use of scaled derivatives removes multi-index binomial coefficients otherwise present in this formula.) The Leibniz formula implies that product derivatives up to a given order \(k = \vert \alpha \vert\) only require the derivatives of the factors up to the same order.

Finally, we also need to consider chain rule or function composition expressions, \(h(x) = f(g(x))\). Explicit formulae, such as that of Faa di Bruno, provide formal expressions relating the derivatives of \(h\) to those of \(f\) and \(g\) with their respective arguments. These expressions, however, are somewhat cumbersome to implement directly. Instead, we rely on a Taylor series approach. The multi-variate function \(f(y)\) can be expanded about the value \(y_0\) as

\[f(y) = \sum_\beta f^{(\beta)} \vert_{y=y_0} (y - y_0)^\beta\]

If we define \(h(x) = f(g(x))\), then we can expand \(h\) about \(x_0\) as

\[h(x) = \sum_\beta f^{(\beta)} \vert_{y=g(x_0)} (g - g_0)^\beta\]

Assuming all functions are sufficiently well behaved, it can be shown that the truncated series

\[h_k(x) \equiv \sum_{\vert \beta \vert \leq k} f^{(\beta)} \vert_{y=g(x_0)} (g - g_0)^\beta\]

has the same derivatives as \(h(x)\) at \(x = x_0\) up to the expansion order, i.e.

\[h_k^{(\alpha)} \vert_{x_0} = h^{(\alpha)} \vert_{x_0} \text{ for } \vert \alpha \vert \leq k.\]

The derivatives of \(f\) and \(g\) are assumed to be known, so repeated application of the Leibniz formula allows one to straightforwardly compute the truncated Taylor series. This procedure is how autodiff implements most mathematical functions, including trigonometric, exponential, and logarithmic functions.

Working with adarray objects

NITROGEN implements the concepts introduced above with the adarray class. An instance of this class stores the derivatives for an expression with respect to a given number of variables up to a fixed order. To illustrate how to use this class, let’s construct the expression \(f = 3 + 2 x^2 - 4(x-y)\), where \(x\) and \(y\) are the two independent variables.

First, we create the adarray objects for the two independent variables (“symbols”) with the sym() function:

>>> import nitrogen.autodiff.forward as adf
>>> x = adf.sym(1.0, 0, 3, 2) # x <-- 1.0
>>> y = adf.sym(2.5, 1, 3, 2) # y <-- 2.5

Note that the value of the symbol, the symbol index, the maximum derivative order, and the total number of symbols must all be specified upon construction and cannot be changed. The basic arithmetic operators +, -, *, and / are overloaded for adarray operands, so we can evaluate \(f\) simply as

>>> f = 3 + 2 * x * x - 4 * (x - y)

The numerical values of the derivatives are stored in an ndarray referred to by the d attribute of the adarray. The derivatives are stored with their scaled values using the same lexical ordering as DFun derivative arrays. Let’s inspect the contents of the derivative arrays we created:

>>> x.d
array([1., 1., 0., 0., 0., 0., 0., 0., 0., 0.])
>>> y.d
array([2.5, 0. , 1. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ])
>>> f.d
array([11.,  0.,  4.,  2.,  0.,  0.,  0.,  0.,  0.,  0.])

The derivative arrays for symbols are very simple. The first element is the value of the symbol. The only other non-zero element is the first derivative with respect to itself, which is always 1. By inspection of f.d we can see that it contains the correct (scaled) derivatives for \(f\).

As noted above, the symbol values must be specified when their adarray objects are created. Instead of a single value, an arbitrarily shaped ndarray (or array-like object) can be passed for vectorized processing. This “base-shape” must be the same for each adarray used together.

>>> x = adf.sym(np.linspace(0,1,5), 0, 3, 2) # x <-- [0, 0.25, 0.50, 0.75, 1.0]
>>> y = adf.sym(np.linspace(0,2,5), 1, 3, 2) # y <-- [0, 0.5,  1.0,  1.5,  2.0]
>>> x.d # shape = (nd,) + base_shape
array([[0.  , 0.25, 0.5 , 0.75, 1.  ],
       [1.  , 1.  , 1.  , 1.  , 1.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  ]])
>>> g = x*x*y*y
>>> g.d
array([[0.      , 0.015625, 0.25    , 1.265625, 4.      ],
       [0.      , 0.125   , 1.      , 3.375   , 8.      ],
       [0.      , 0.0625  , 0.5     , 1.6875  , 4.      ],
       [0.      , 0.25    , 1.      , 2.25    , 4.      ],
       [0.      , 0.5     , 2.      , 4.5     , 8.      ],
       [0.      , 0.0625  , 0.25    , 0.5625  , 1.      ],
       [0.      , 0.      , 0.      , 0.      , 0.      ],
       [0.      , 1.      , 2.      , 3.      , 4.      ],
       [0.      , 0.5     , 1.      , 1.5     , 2.      ],
       [0.      , 0.      , 0.      , 0.      , 0.      ]])

We can also create adarray objects for constant expressions

>>> c = adf.const(np.linspace(1,2,5), 3, 2) # c <-- [1.0, 1.2, 1.4, 1.6, 1.8, 2.0] (constant)
>>> c.d
array([[1.  , 1.25, 1.5 , 1.75, 2.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  ]])

Because c is a constant, its derivatives are identically zero. Other constructor functions include array(), empty_like(), and const_like().

Mathematical functions

The forward module implements the following mathematical and arithmetical functions:

add()

Addition (used by +)

subtract()

Subtraction (used by -)

mul()

Multiplication (used by *)

div()

Division (used by /)

sin()

Sine.

cos()

Cosine.

exp()

The exponential function, \(e^x\).

log()

Natural logarithm.

powi()

Integer exponentiation, \(x^i\).

powf()

Real or complex exponentiatoin, \(x^p\).

sqrt()

Square root, \(\sqrt{x}\).

In general, complex arithmetic is supported whenever possible. For certain functions, like sqrt() or powf(), different conventions for branch cuts are possible. These issues are discussed in more detail in the function documentation.

Linear algebra

The forward module also supports differentiation of some basic linear algebra algorithms, including:

chol_sp()

Cholesky decomposition of a symmetric packed matrix.

inv_tp()

Inversion of a triangular matrix.

llt_tp()

L @ L.T of a lower triangular matrix.

ltl_tp()

L.T @ L of a lower triangular matrix.