current community

your communities

more stack exchange communities

Stack Exchange
tour help
Take the 2-minute tour ×
Code Review Stack Exchange is a question and answer site for peer programmer code reviews. It's 100% free, no registration required.
up vote 1 down vote favorite

I have the following code:

for j in range(rows):
    for i in range(cols):
        k = j * rows + i

        k1 = k + 1 if i < cols - 1 else k - 1
        k2 = k - 1 if i > 0 else k + 1
        k3 = k + cols if j < rows - 1 else k - cols
        k4 = k - cols if j > 0 else k + cols


        w1 = U[k]
        w2 = U[k-1] if i > 0 else U[k]
        w3 = V[k]
        w4 = V[k-cols] if j > 0 else V[k]

        zarray[k] = (w1 + w2 + w3 + w4) * parray[k] - (w1 * parray[k1] + w2 * parray[k2] + w3 * parray[k3] + w4 * parray[k4])

I want to know if there is a way to vectorize this loops, because I think that exists a kind of "convolution" for zarray.

U and V are arrays representing 2D matrices with 512x512 elements, and parray is also a 2D representation with 512x512 elements.

In a previous post vectorization was recommended, but now I can not figure out how to vectorize when I have different indices operations.

share|improve this question

1 Answer 1

up vote 2 down vote accepted

I was the one who recommended vectorizing. I think the first approach would be to convert your code to using 2D indexes. I think this will make the vectorizing clearer:

for j in range(rows):
    for i in range(cols):
        i1 = i+1 if i<cols-1 else i-1
        i2 = i-1 if i>0 else i+1
        j1 = j+1 if j<rows-1 else j-1
        j2 = j-1 if j>0 else j+1

        w1 = U[j, i]
        w2 = U[j, i-1] if i > 0 else U[j, i]
        w3 = V[j, i]
        w4 = V[j-1, i] if j > 0 else V[j, i]

        p = parray[j, i]
        p1 = parray[j, i1]
        p2 = parray[j, i2]
        p3 = parray[j1, i]
        p4 = parray[j2, i]

        zarray[j, i] = (w1 + w2 + w3 + w4)*p - (w1*p1 + w2*p2 + w3*p3 + w4*p4)

So, if I am reading this code right, you are shifting some rows and columns around. So the next step is to re-implement this by making copies of the arrays that follow the same patterns:

U1 = np.empty_like(U)
V1 = np.empty_like(V)
parray1 = np.empty_like(parray)
parray2 = np.empty_like(parray)
parray3 = np.empty_like(parray)
parray4 = np.empty_like(parray)

U1[:, 1:] = U[:, :-1]  # this is w2
V1[1:, :] = V[:-1, :]  # this is w4
U1[:, 0] = U[:, 0]
V1[0, :] = V[0, :]

parray1[:, :-1] = parray[:, 1:]  # this is the result of k1
parray2[:, 1:] = parray[:, :-1]  # this is the result of k2
parray3[:-1, :] = parray[1:, :]  # this is the result of k3
parray4[1:, :] = parray[:-1, :]  # this is the result of k4

parray1[:, -1] = parray[:, -2]
parray2[:, 0] = parray[:, 1]
parray3[-1, :] = parray[-2, :]
parray4[0, :] = parray[1, :]

zarray = (U+U1+V+V1)*parray - (U*parray1 + U1*parray2 + V*parray3 + V1*parray4)

You could also use np.hstack and np.vstack instead of slices:

U1 = np.hstack([U[:, :1], U[:, :-1]])
V1 = np.vstack([V[:1, :], V[:-1, :]])

parray1 = np.hstack([parray[:, 1:], parray[:, -2:-1]])
parray2 = np.hstack([parray[:, 1:2], parray[:, :-1]])
parray3 = np.vstack([parray[1:, :], parray[-2:-1, :]])
parray4 = np.vstack([parray[1:2, :], parray[:-1, :]])

zarray = (U+U1+V+V1)*parray - (U*parray1 + U1*parray2 + V*parray3 + V1*parray4)
share|improve this answer
    
Yes, this is the way I just solved! Thank you a lot! I thought (w1+w2+w3+w4) would be a inconvenience, but it just was the same as de previuos one. –  FacundoGFlores Jun 23 at 19:46
    
Just an observation: which is a better way for copying? a[:] = v or a = v.copy()? –  FacundoGFlores Jun 23 at 19:48
1  
a[:] = v only works if a is already defined, so a = v.copy() is better for copying the whole thing. I guess you could pre-define the arrays using np.empty_like rather than by making copies, then slice those. I am not sure whether that would save much time or not. –  TheBlackCat Jun 23 at 19:56
1  
It turns out using np.empty_like is orders of magnitude faster, so I updated my answer to use it instead. –  TheBlackCat Jun 23 at 20:01
1  
I have also added an approach using hstack and vstack. I can't tell which is faster, my benchmarks are all over the place. –  TheBlackCat Jun 23 at 20:09

Your Answer

 
discard

By posting your answer, you agree to the privacy policy and terms of service.

Not the answer you're looking for? Browse other questions tagged or ask your own question.