[Numpy-discussion] any interest in including a second-order gradient?
Andrew Hawryluk
HAWRYLA at novachem.com
Mon Oct 27 11:29:41 EDT 2008
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
More information about the NumPy-Discussion
mailing list