Computing derivatives using numpy

Question

I'm trying to implement a differential in python via numpy that can accept a scalar, a vector, or a matrix.

import numpy as np

def foo_scalar(x):
    f = x * x
    df = 2 * x

    return f, df

def foo_vector(x):
    f = x * x
    n = x.size
    df = np.zeros((n, n))
    for mu in range(n):
        for i in range(n):
            if mu == i:
                df[mu, i] = 2 * x[i]

    return f, df

def foo_matrix(x):
    f = x * x
    m, n = x.shape
    df = np.zeros((m, n, m, n))
    for mu in range(m):
        for nu in range(n):
            for i in range(m):
                for j in range(n):
                    if (mu == i) and (nu == j):
                        df[mu, nu, i, j] = 2 * x[i, j]

    return f, df

This works fine, but it seems like there should be a way to do this in a single function, and let numpy "figure out" the correct dimensions. I could force everything into a 2-D array form with something like

x = np.array(x)
if len(x.shape) == 0:
    x = x.reshape(1, 1)
elif len(x.shape) == 1:
    x = x.reshape(-1, 1)
if len(f.shape) == 0:
   f = f.reshape(1, 1)
elif len(f.shape) == 1:
   f = f.reshape(-1, 1)

and always have 4 nested for loops, but this doesn't scale if I need to generalize to higher-order tensors.

Is what I'm trying to do possible, and if so, how?

Answer 1

I highly doubt there is a function to generate the second parameter returned by the function in Numpy. That being said you can play with the feature of Numpy and Python so to vectorize this and make the function faster. You first need to generate the indices and, then generate the target matrix and set it. Note that operating with N-dimensional generic arrays tends to be slow and tricky in non-trivial cases. The magic * unrolling operator is used to generate N parameters.

def foo_generic(x):
    f = x ** 2
    idx = np.stack(np.meshgrid(*[np.arange(e) for e in x.shape], indexing='ij'))
    idx = tuple(np.concatenate((idx, idx)).reshape(2*x.ndim, -1))
    df = np.zeros([*x.shape, *x.shape])
    df[idx] = 2 * x.ravel()
    return f, df

Note that foo_generic does not support scalar and it would be very inefficient to use it for that anyway, but you can add a condition in it to support this special case apart.

The df matrix will very quickly be huge for higher order so I strongly advise you not to use dense matrices for that since the number of zeros is huge compared to the number of values in the matrix case already. Sparse matrices fix this. In fact, for a 5x5 matrix, there are >95% of zeros. Not to mention the matrix becomes quickly huge and willing a huge matrix full of zeros is not efficient.