# Vectorizing in Python [closed]

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.

• This question lacks any indication of what the code is intended to achieve. To help reviewers give you better answers, please add sufficient context to your question, including a title that summarises the purpose of the code. We want to know why much more than how. The more you tell us about what your code is for, the easier it will be for reviewers to help you. The title needs an edit to simply state the task. Commented Feb 27, 2019 at 9:23

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)

• 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. Commented Jun 23, 2015 at 19:46
• Just an observation: which is a better way for copying? a[:] = v or a = v.copy()? Commented Jun 23, 2015 at 19:48
• 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. Commented Jun 23, 2015 at 19:56
• It turns out using np.empty_like is orders of magnitude faster, so I updated my answer to use it instead. Commented Jun 23, 2015 at 20:01
• I have also added an approach using hstack and vstack. I can't tell which is faster, my benchmarks are all over the place. Commented Jun 23, 2015 at 20:09