Looking for a way to create a particular matrix in python in a less amount of time

Question

I want to create a matrix like this in Python: $$ \begin{pmatrix} 0 & 0 & 0 & 0 & 1 & 1 & 1 & 1 & 2 & 2 & 2 & 2 & 3 & 3 & 3 & 3\\ 0 & 1 & 2 & 3 & 0 & 1 & 2 & 3 & 0 & 1 & 2 & 3 & 0 & 1 & 2 & 3\\ 1 & 1 & 1 & 1 & 1 & 1 & 1 & 1 & 1 & 1 & 1 & 1 & 1 & 1 & 1 & 1 \end{pmatrix} $$

I have written a code that works, but takes a long time to execute:

row = 3000
col = 4000

row_1 = np.asarray((col)*[0])
for i in range(1,row):
    row_1 = np.append(row_1, np.asarray((col)*[i]))
row_2 = np.asarray(row*[x for x in range(0,col)])
row_3 = np.asarray(col*(row)*[1])

output = np.vstack([row_1,row_2,row_3])

It takes between 55 to 65 seconds to run this code. Is it possible to create the same matrix in a way that is more efficient?

Answer 1

The NumPy Reference should be the first place you look when you have a problem like this. The operations you need are nearly always in there somewhere. And functions that you find while browsing the reference are sure to come in useful later in your NumPy career.

For the first row, you could use numpy.arange and numpy.repeat:

>>> np.repeat(np.arange(4), 4)
array([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3])

For the second row, you could use numpy.tile:

>>> np.tile(np.arange(4), 4)
array([0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3])

For the third row, you could use numpy.ones:

>>> np.ones(4 * 4, dtype=int)
array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])

Putting that together (using numpy.ones_like for row3 so that it has the same shape and data type as row2):

row1 = np.repeat(np.arange(row), col)
row2 = np.tile(np.arange(col), row)
row3 = np.ones_like(row2)
output = np.vstack([row1, row2, row3])

This is about 100 times as fast as the code in the post.

But you can do better than that! There is a hint in the documentation for numpy.tile:

Note: Although tile may be used for broadcasting, it is strongly recommended to use numpy’s broadcasting operations and functions.

Broadcasting is the way that NumPy adapts the inputs to an operation so that they match in shape. Using this we can write the whole function using array indexing and reshaping operations, like this:

output = np.zeros((3, row, col), dtype=int)
output[0] = np.arange(row).reshape(-1, 1)
output[1] = np.arange(col).reshape(1, -1)
output[2] = 1
output = output.reshape(3, -1)

This is about 160 times as fast as that code in the post.

A slightly clearer way to write this is to use numpy.meshgrid to generate the coordinate arrays:

output = np.zeros((3, row, col), dtype=int)
i, j = np.meshgrid(np.arange(row), np.arange(col), indexing='ij', sparse=True)
output[0] = i
output[1] = j
output[2] = 1
output = output.reshape(3, -1)

This isn't any faster than the previous version but the use of numpy.meshgrid gives more of a clue to the reader as to what is happening.

(You might be able to do even better than that. Do you really need to construct the whole matrix? Could you combine broadcasting and numpy.meshgrid to get the results you want instead? I can't answer this since you didn't explain what you are going to use this code for. But maybe you can ask another question showing us more of your program.)