Finite difference Jacobians of numeric functions

jacobian_fd.py

'''Finite difference Jacobians

Author: James Gregson 
Date: 2020-10-04

Utility function to differentiate numeric functions with finite differences.
Useful for checking analytically derived Jacobian & gradient computations.

See bottom of file for example usage.

'''
import numpy as np

def dirac( shape, on, val=1.0 ):
    '''Multidimensional dirac function returning a single on value
    at specified location with specified value
    '''
    V = np.zeros(shape)
    V[on] = val
    return V

def jacobian_fd( func, *args, jacfd_epsilon=1e-4, **kwargs ):
    '''Numerically differentiate function with respect to it's positional arguments

    Args:
        - func (function): function to numerically differentiate, takes
            ndarray positional arguments and optional keyword arguments
            that are not differentiated. Returns a single ndarray return value

        - args (ndarray list): ndarray arguments to func

        - jacfd_epsilon (float): finite-difference epsilon

        - kwargs (keyword args): optional keyword arguments for func,
            passed to func unmodified

    Returns:
        - f (ndarray): return value of function

        - dF (list of ndarrays): partial derivatives of function w.r.t.
            each argument in args. Shape of the i'th return value will be
            [*f.shape,*args[i].shape]

    '''
    f = func( *args, **kwargs )
    if not isinstance( f, np.ndarray ):
        raise RuntimeError('Return of func must be an ndarray')
    dF = []
    args = [ a.copy() for a in args ]
    for aid in range(len(args)):
        if not isinstance( args[aid], np.ndarray ):
            raise RuntimeError('All positional arguments must be numpy ndarrays')
        shape = args[aid].shape
        count = args[aid].size
        darg = np.zeros((*f.shape,count))
        for i in range(count):
            tmp = args[aid]
            args[aid] = tmp + dirac(count,i,jacfd_epsilon).reshape(shape)
            f1 = func( *args, **kwargs )
            args[aid] = tmp - dirac(count,i,jacfd_epsilon).reshape(shape)
            f0 = func( *args, **kwargs )
            args[aid] = tmp
            darg[...,i] = (f1-f0)/(2.0*jacfd_epsilon)
        dF.append( darg.reshape((*f.shape,*shape)) )
    return f, dF

if __name__ == '__main__':

    def func( A, B ):
        return A@B

    A = np.random.standard_normal((2,4))
    B = np.random.standard_normal((4,2))

    f, (dfdA,dfdB) = jacobian_fd( func, A, B )

    dB = np.random.standard_normal((4,2))*1e-5
    f2_an = func( A, B+dB )
    f2_fd = f + np.sum( dfdB*dB[None,None,...], axis=(-2,-1) )
    print( 'analytic vs fd: {}'.format( np.linalg.norm(f2_an-f2_fd) ) )

    dA = np.random.standard_normal((2,4))*1e-5
    f2_an = func( A+dA, B )
    f2_fd = f + np.sum( dfdA*dA[None,None,...], axis=(-2,-1) )
    print( 'analytic vs fd: {}'.format( np.linalg.norm(f2_an-f2_fd) ) )