[Numpy-discussion] any interest in including a second-order gradient?

[Andrew Hawryluk]

We wrote a simple variation on the gradient() function to calculate the
second derivatives. Would there be any interest in including a
gradient2() in numpy?

Andrew



def gradient2(f, *varargs):
    """Calculate the second-order gradient of an N-dimensional scalar
function.

    Uses central differences on the interior and first differences on
boundaries
    to give the same shape.

    Inputs:

      f -- An N-dimensional array giving samples of a scalar function

      varargs -- 0, 1, or N scalars giving the sample distances in each
direction

    Outputs:

      N arrays of the same shape as f giving the derivative of f with
respect
      to each dimension.

    """
    N = len(f.shape)  # number of dimensions
    n = len(varargs)
    if n == 0:
        dx = [1.0]*N
    elif n == 1:
        dx = [varargs[0]]*N
    elif n == N:
        dx = list(varargs)
    else:
        raise SyntaxError, "invalid number of arguments"

    # use central differences on interior and first differences on
endpoints

    outvals = []

    # create slice objects --- initially all are [:, :, ..., :]
    slice1 = [slice(None)]*N
    slice2 = [slice(None)]*N
    slice3 = [slice(None)]*N
    

    otype = f.dtype.char
    if otype not in ['f', 'd', 'F', 'D']:
        otype = 'd'

    for axis in range(N):
        # select out appropriate parts for this dimension
        out = zeros(f.shape, f.dtype.char)
        slice1[axis] = slice(1, -1)
        slice2[axis] = slice(2, None)
        slice3[axis] = slice(None, -2)
        # 1D equivalent -- out[1:-1] = (f[2:] - 2*f[1:-1] + f[:-2])
        out[slice1] = (f[slice2] - 2*f[slice1] + f[slice3])
        slice1[axis] = 0
        slice2[axis] = 1
        slice3[axis] = 2
        # 1D equivalent -- out[0] = (f[2] - 2*f[1] + f[0])
        out[slice1] = (f[slice3] - 2*f[slice2] + f[slice1])
        slice1[axis] = -1
        slice2[axis] = -2
        slice3[axis] = -3
        # 1D equivalent -- out[-1] = (f[-1] - 2*f{-2] + f[-3])
        out[slice1] = (f[slice1] - 2*f[slice2] + f[slice3])

        # divide by the squared step size
        outvals.append(out / dx[axis] / dx[axis])

        # reset the slice object in this dimension to ":"
        slice1[axis] = slice(None)
        slice2[axis] = slice(None)
        slice3[axis] = slice(None)

    if N == 1:
        return outvals[0]
    else:
        return outvals