[Numpy-discussion] [Matplotlib-users] Vectorization

[Keith Goodman]

On Fri, Jul 2, 2010 at 11:45 AM, Keith Goodman <kwgoodman at gmail.com> wrote:
> On Fri, Jul 2, 2010 at 11:33 AM, Benjamin Root <ben.root at ou.edu> wrote:
>> I am moving this over to numpy-discussion maillist...
>>
>> I don't have a firm answer for you, but I did notice one issue in your
>> code.  You call arange(len(dx) - 1) for your loops, but you probably really
>> need arange(1, len(dx) - 1) because you are accessing elements both after
>> *and* before the current index.  An index of -1 is actually valid because
>> that means the last element of the array, and may not be what you intended.
>>
>> Ben Root
>>
>> On Fri, Jul 2, 2010 at 1: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?
>
> If no one knows how to vectorize it then one way to go is cython. If
> you convert your arrays to lists then it is very easy to convert the
> loop to cython. Fast too.

Some more thoughts: you pull phi[k,j,i] four times per loop. Setting
phikji = phi[k,j,i] might give a little more speed. Might also help to
do stuff like phi_i = phi[:,:,i], phi_ip1 = phi[:,:,i+1], etc right
after "for i in range()". Also everything gets multiplied by "-0.5 * m
* phi" so you should do that outside the loop. You can square the d's
(dx, dy, and dz) outside the loop. Doing all these and then using
cython should makes things very fast.