[Numpy-discussion] [Matplotlib-users] Vectorization
Geoffrey Ely
gely at usc.edu
Fri Jul 2 15:47:22 EDT 2010
On Jul 2, 2010, at 11:33 AM, Benjamin Root wrote:
> I want to do the same for the calculation of the kinetic energy: <phi|p^2|phi>/2m. There is a laplacian in the volume integral which complicates things:
> K = 0.0
> for i in numpy.arange(len(dx)-1):
> for j in numpy.arange(len(dy)-1):
> for k in numpy.arange(len(dz)-1):
> K += -0.5 * m * phi[k,j,i] * (
> (phi[k,j,i-1] - 2.0*phi[k,j,i] + phi[k,j,i+1]) / dx[i]**2
> + (phi[k,j-1,i] - 2.0*phi[k,j,i] + phi[k,j+1,i]) / dy[j]**2
> + (phi[k-1,j,i] - 2.0*phi[k,j,i] + phi[k+1,j,i]) / dz[k]**2
> )
>
> My question is, how would I vectorize such loops? I don't know how I would manage the "numpy.newaxis" code-foo with neighbours dependency... Any idea?
How about:
K = -0.5 * m * (
phi[1:-1,1:-1,1:-1] * (
np.diff(phi[1:-1,1:-1,:], 2, 2) / dx[None,None,1:-1]**2
+ np.diff(phi[1:-1,:,1:-1], 2, 1) / dy[None,1:-1,None]**2
+ np.diff(phi[:,1:-1,1:-1], 2, 0) / dz[1:-1,None,None]**2
)
).sum()
(Not tested)
More information about the NumPy-Discussion
mailing list