[Numpy-discussion] Little vectorization help...

[Andrea Gavana]

Hi All,

    I am tryin to "vectorize" 3 nested for loops but I am not having
much success. Here is the code I use:

import numpy
import numpy.ma as masked

grid = numpy.zeros((nx, ny), dtype=numpy.float32)
xOut = numpy.zeros((nx, ny), dtype=numpy.float32)
yOut = numpy.zeros((nx, ny), dtype=numpy.float32)

z = GetCentroids() # Some vector z values
prop = GetValue()  # Some other vector values

NaN = numpy.NaN

for I in xrange(1, nx+1):

    for J in xrange(1, ny+1):

        theSum = []

        for K in xrange(1, nz+1):

            cellPos = I-1 + nx*(J-1) + nx*ny*(K-1)
            centroid = z[cellPos]

            if low <= centroid <= high and actnum[cellPos] > 0:
                theSum.append(prop[cellPos])

        if theSum:
            grid[I-1, J-1] = sum(theSum)/len(theSum)
        else:
            grid[I-1, J-1] = NaN

        xOut[I-1, J-1], yOut[I-1, J-1] = x[cellPos], y[cellPos]

grid = masked.masked_where(numpy.isnan(grid), grid)


Some explanation:

1) "z" is a vector of nx*ny*nz components, where nx = 100, ny = 73, nz
= 23, which represents 3D hexahedron cell centroids;
2) "prop" is a vector like z, with the same shape, with some floating
point values in it;
3) "actnum" is a vector of integers (0 or 1) with the same shape as z,
and indicates if a cell should be considered in the loop or not;
4) low and high are 2 floating point values with low < high: if the
cell centroid fall between low and high and the cell is active (as
stated in "actnum"), then I take the value of "prop" in that cell and
I append it to the "theSum" list;
5) At the end of the K loop, I just take an arithmetic mean of the
values in "theSum" list.

I think I may be able to figure out how to vectorize the part
regarding the "grid" variable, but I have no idea on what to do for
the xOut and yOut variables, and I need them because I use them in a
later call to matplotlib.contourf.

If you could drop some hint on how to proceed, ot would be very
appreciated. Thank you for your suggestions.

Andrea.

"Imagination Is The Only Weapon In The War Against Reality."
http://xoomer.alice.it/infinity77/