How to fill a numpy array with a gradient with a vectorized code?

Question

Two numpy arrays have to be filled as follow:

[0 1 2]
 [0 1 2]
 [0 1 2]
 [0 1 2]
 [0 1 2]]

[[0 0 0]
 [1 1 1]
 [2 2 2]
 [3 3 3]
 [4 4 4]]

Or enter image description here

A first idea was:

ax=np.zeros((5,3),np.int16)
ay=np.zeros((5,3),np.int16)

for j in range(0,3):
    for i in range(0,5):
        ax[i,j]=j#filling ax with x=col
        ay[i,j]=i#filling ay with y values y=li

A second idea was:

bx = np.zeros((5,3),np.int16)
by = np.zeros((5,3),np.int16)

for j in range(3):
    bx[:,j]=j

for i in range(5):
    by[i,:]=i

I am sure there's a better way, which one? Thanks jp

zhangxaochen · Accepted Answer · 2014-01-31 15:39:24Z

2

I think using numpy.tile might be better:

In [422]: np.tile((0,1,2), (5,1))
Out[422]: 
array([[0, 1, 2],
       [0, 1, 2],
       [0, 1, 2],
       [0, 1, 2],
       [0, 1, 2]])

In [473]: tile(arange(5)[:,None], 3)
Out[473]: 
array([[0, 0, 0],
       [1, 1, 1],
       [2, 2, 2],
       [3, 3, 3],
       [4, 4, 4]])

Time efficiency:

for a small matrix of shape (5,3), the for-loop is faster:

In [490]: timeit np.tile((0,1,2), (5,1))
10000 loops, best of 3: 38.3 us per loop

In [491]: %%timeit
     ...: bx = np.zeros((5,3),np.int16)
     ...: for j in range(3):
     ...:     bx[:,j]=j
     ...: 
100000 loops, best of 3: 16.5 us per loop

but for a large matrix of shape (5, 1000), tile is faster:

In [488]: timeit n=1000; tile(xrange(n), (5,1))
1000 loops, best of 3: 313 us per loop

In [489]: %%timeit
     ...: n=1000
     ...: bx=zeros((5, n))
     ...: for j in range(n):
     ...:     bx[:,j]=j
     ...: 
100 loops, best of 3: 3.97 ms per loop

Anyway, tile makes code cleaner.

edited Jan 31, 2014 at 15:39

answered Jan 31, 2014 at 14:44

zhangxaochen

34.2k15 gold badges82 silver badges115 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Jean-Pat Over a year ago

Great! How to timeit ?

Collectives™ on Stack Overflow

How to fill a numpy array with a gradient with a vectorized code?

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related