[Numpy-discussion] [Matplotlib-users] Vectorization

[Nicolas Bigaouette]

2010/7/2 Charles R Harris <charlesr.harris at gmail.com>

>
>
> On Fri, Jul 2, 2010 at 12:15 PM, Nicolas Bigaouette <nbigaouette at gmail.com
> > wrote:
>
>> Hi all,
>>
>> I don't really know where to ask, so here it is.
>>
>> I was able to vectorize the normalization calculation in quantum
>> mechanics: <phi|phi>. Basically it's a volume integral of a scalar field.
>> Using:
>>
>>> norm = 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):
>>>             norm += psi[k,j,i]**2 * dx[i] * dy[j] * dz[k]
>>>
>> if dead slow. I replaced that with:
>>
>>> norm = (psi**2 *
>>> dx*dy[:,numpy.newaxis]*dz[:,numpy.newaxis,numpy.newaxis]).sum()
>>>
>> which is almost instantanious.
>>
>> 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?
>>
>> Thanx!
>>
>>
> Wouldn't the Laplacian look better in momentum space? That is, do a 3d
> Fourier transform, multiply by the appropriate weight, and do the integral
> in the transform space.
>
> Chuck
>

Thanx all for suggestions. The best approach was the convolution. I couldn't
do a single convolution since the cell sizes are the same in x, y and z. I
had to separate the laplacian into a sum over the three dimensions and
convolve each individually. There is also terms appearing because of the
non-uniform grid.

Not that I work in atomic units, so m=1.

I'm pasting the code I've finally written:

> def Energy_Kinetic(phi, J1x, J1y, J1z, J2x, J2y, J2z):
>
>     kernel_x1 = numpy.array([[[0,0,0], [0,0,0],  [0,0,0]],
>                              [[0,0,0], [1,-2,1], [0,0,0]],
>                              [[0,0,0], [0,0,0],  [0,0,0]]])
>     kernel_y1 = numpy.array([[[0,0,0], [0,0,0],  [0,0,0]],
>                              [[0,1,0], [0,-2,0], [0,1,0]],
>                              [[0,0,0], [0,0,0],  [0,0,0]]])
>     kernel_z1 = numpy.array([[[0,0,0], [0,1,0],  [0,0,0]],
>                              [[0,0,0], [0,-2,0], [0,0,0]],
>                              [[0,0,0], [0,1,0],  [0,0,0]]])
>
>     kernel_x2 = numpy.array([[[0,0,0], [0,0,0],  [0,0,0]],
>                              [[0,0,0], [1,0,-1], [0,0,0]],
>                              [[0,0,0], [0,0,0],  [0,0,0]]])
>     kernel_y2 = numpy.array([[[0,0,0], [0,0,0],  [0,0,0]],
>                              [[0,1,0], [0,0,0],  [0,-1,0]],
>                              [[0,0,0], [0,0,0],  [0,0,0]]])
>     kernel_z2 = numpy.array([[[0,0,0], [0,1,0],  [0,0,0]],
>                              [[0,0,0], [0,0,0],  [0,0,0]],
>                              [[0,0,0], [0,-1,0], [0,0,0]]])
>
>     convolution = ( \
>         scipy.signal.convolve(phi, kernel_x1, mode='same') / J1x**2

      + scipy.signal.convolve(phi, kernel_y1, mode='same') /
> J1y[:,numpy.newaxis]**2
>       + scipy.signal.convolve(phi, kernel_z1, mode='same') /
> J1z[:,numpy.newaxis,numpy.newaxis]**2
>
      + 0.5*(J2x / J1x**3)*scipy.signal.convolve(phi, kernel_x2,
> mode='same') \
>       + 0.5*(J2y[:,numpy.newaxis] /
> J1y[:,numpy.newaxis]**3)*scipy.signal.convolve(phi, kernel_y2, mode='same')
> \
>       + 0.5*(J2z[:,numpy.newaxis,numpy.newaxis] /
> J1z[:,numpy.newaxis,numpy.newaxis]**3)*scipy.signal.convolve(phi,kernel_z2,
> mode='same') \
>     )
>
>     K = -0.5 * (numpy.conjugate(phi) * convolution * J1x *
> J1y[:,numpy.newaxis] * J1z[:,numpy.newaxis,numpy.newaxis]).real.sum()
>
>     return K
> #
>
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.python.org/pipermail/numpy-discussion/attachments/20100730/f7ca4474/attachment.html>