hrithik mahesh
hrithik mahesh

Reputation: 113

vectorize image generation with two for loops

I'm trying to generate sinusiod image but i see no alternative without using two for loops. Can someone help me vectorize this code?

import numpy as np
M, N = 1001, 1001
u = 100
v = 200
sin_img = np.zeros((M, N))
for m in range(M):
  for n in range(N):
    sin_img[m, n] = np.sin((2 * np.pi * u * m / M) + (2 * np.pi * v * n / N))

Upvotes: 2

Views: 94

Answers (2)

Akshay Sehgal
Akshay Sehgal

Reputation: 19322

Here is the proper vectorized way of doing it without any loops what so ever.

  1. You need to create a matrix that holds the indexes of its own positions as tuples. So if the matrix is 5,5, then this matrix should hold tuples of indexes for each of its positions and therefore be a 5,5,2 matrix. So 0th row, 0th column should hold (0,0), 4th row and 2nd column should hold (4,2) and so on.
  2. Once you have that 1001,1001,2 created using np.indices, you can use that in a vectorized way in the numpy functions with 1001, 1001, 1 in each of the 2 parts.
  3. Post that use your np.sin on the sum of the 2 parts.
idx_sin_img = np.stack(np.indices((M,N)), axis=-1) #(1001,1001,2)

out = np.sin((2 * np.pi * u * idx_sin_img[...,0] / M) + 
             (2 * np.pi * v * idx_sin_img[...,1] / N))

out.shape
(1001, 1001)

Just to confirm that it gives the same output -

np.all(sin_img == out) #sin_img is from OP's own code
#True

Benchmarks -

  1. Original method posted by OP - 1.25 s ± 28 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
  2. List comprehension method by Roy - 1.44 s ± 68.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
  3. Purely vectorized method by Akshay - 26.9 ms ± 1.93 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Upvotes: 2

Roy Cohen
Roy Cohen

Reputation: 1570

I don't know what "vectorize" means but if you want to get rid of the two for loops you can use a list comprehension:

sin_img = np.array([[np.sin((2 * np.pi * u * m / M) + (2 * np.pi * v * n / N))
    for n in range(N)]
    for m in range(M)])

Explanation:
The code is made up from a list comprehension inside a list comprehension. Each list comprehension is made like this:

[some_expression for NAME in some_iterable]

where the some_expression part can be another list comprehension.

Upvotes: 2

Related Questions