[Numpy-discussion] [Matplotlib-users] Vectorization
Keith Goodman
kwgoodman@gmail....
Fri Jul 2 13:45:34 CDT 2010
On Fri, Jul 2, 2010 at 11:33 AM, Benjamin Root <ben.root@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@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.
More information about the NumPy-Discussion
mailing list