How to vectorize to remove for loops in Python?

Question

I want to remove the for loops in the following Python code by vectorizing it. I searched the Stack overflow and all over the internet to find a solution to no avail.

for v in range(height):
  for u in range(width):
    start[v,u,0] = -0.5 + u / (width-1)
    start[v,u,1] = (-0.5 + v / (height-1)) * height / width
    start[v,u,2] = 0

I think we should use NumPy. Perhaps numpy.linspace() is to be used? So, what do you think?

What are all these variables to begin with? ``start[v,u,0]`` suggests you are *already* using numpy. — MisterMiyagi, Oct 13 '20 at 15:58
Welcome to StackOverflow! Please take the time to read this post on how to [How do I ask a good question?](https://stackoverflow.com/help/how-to-ask) as well as how to provide a [minimal, Complete, and Verifiable example](https://stackoverflow.com/help/mcve) and revise your question accordingly — yatu, Oct 13 '20 at 16:00
It seems that you can vectorize that indexing using `ax0, ax1 = np.ogrid[:height, :width`] — yatu, Oct 13 '20 at 16:10

mathfux · Answer 1 · 2020-10-13T16:47:02.397

The following:

start[v,u,0] = -0.5 + u / (width-1)
start[v,u,1] = (-0.5 + v / (height-1)) * height / width
start[v,u,2] = 0

works if v and u are fancy indices:

start = np.random.randint(99, size=(5,7,3)).astype(float)
height, width = start.shape[:2]
v, u = np.indices((height, width))

Sample run (values are rounded):

>>> v
[[0 0 0 0 0 0 0]
 [1 1 1 1 1 1 1]
 [2 2 2 2 2 2 2]
 [3 3 3 3 3 3 3]
 [4 4 4 4 4 4 4]]
>>> u
[[0 1 2 3 4 5 6]
 [0 1 2 3 4 5 6]
 [0 1 2 3 4 5 6]
 [0 1 2 3 4 5 6]
 [0 1 2 3 4 5 6]]
>>> start[v,u,0]
[[-0.5  -0.33 -0.17  0.    0.17  0.33  0.5 ]
 [-0.5  -0.33 -0.17  0.    0.17  0.33  0.5 ]
 [-0.5  -0.33 -0.17  0.    0.17  0.33  0.5 ]
 [-0.5  -0.33 -0.17  0.    0.17  0.33  0.5 ]
 [-0.5  -0.33 -0.17  0.    0.17  0.33  0.5 ]]
>>> start[v,u,1]
[[-0.36 -0.36 -0.36 -0.36 -0.36 -0.36 -0.36]
 [-0.18 -0.18 -0.18 -0.18 -0.18 -0.18 -0.18]
 [ 0.    0.    0.    0.    0.    0.    0.  ]
 [ 0.18  0.18  0.18  0.18  0.18  0.18  0.18]
 [ 0.36  0.36  0.36  0.36  0.36  0.36  0.36]]

A quick preview for better understanding:

score 0 · Answer 2 · answered Oct 13 '20 at 16:48

It's hard to answer your question without more info (do you want to create the array? Is it already created?? why are your v, u indicies 'backwards'? ...

But to answer your question, yes - there are many ways to do what you want with numpy.

Have a look below, and then go and read for an hour or two some numpy tutorials etc:

import numpy as np

rows = 10
cols = 5

values = np.zeros((rows, cols, 3))

values[:,:,0] = np.fromfunction(
    lambda r, c: -0.5 + c / (cols - 1),
    (rows, cols))

values[:,:,1] = np.fromfunction(
    lambda r, c: (-0.5 + r / (rows - 1) * rows / cols),
    (rows, cols))

Note there are many ways to do what you want with numpy.

The above is just one, and written in a more expanded, easy to follow way. You could do the same thing in one line probably, but it would be harder to understand if you're new to numpy.

How to vectorize to remove for loops in Python?

2 Answers2