[Numpy-discussion] [Matplotlib-users] Vectorization

[Bruce Southey]

On 07/02/2010 01:45 PM, Keith Goodman 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.
> _______________________________________________
>    
Since things do not depend on previous results. Without thinking much 
can you just replace all your phi[] references to being phi[].sum()?

Probably wrong but hopefully you can figure out what I mean.



My reasoning is that over three loops then

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]

is just sum of all the entries except the last dimension in each axis of phi ie

K= -0.5 * m * phi[0:len(dz)-1, 0:len(dz)-1,0:len(dx)-1].sum()

Similarity for the remaining places you have to sum the appropriate sections of the phi array. So phi[k,j,i-1] becomes something like:
phi[:,:,0:i-1].sum() except you have to address the division. Therefore you have to sum over the appropriate access

(phi[:,:,0:i-1].sum(axis=2)/dx[0:i-1]**2).sum() # or
numpy.dot((phi[:,:,0:i-1].sum(axis=2), 1/(dx[0:i-1]**2))

Then I got confused when say i=0 because the reference phi[k,j,i-1] becomes phi[k,j,-1] - which is the last index of that axis. Is that what you meant to happen?


Bruce