Is it possible to speed up calculation of very large arrays in numpy?

Question

I have looked around, but could not find an exact answer.

I have several matrices of 1200x2000 that I want to save in a way that I can read them later in some sort of loop. When I want to save 100 of these matrices using the method below, the process is really fast. But if I want to save 1000, the time necessary for each iteration gets really long.

Here is an example of the problem (here the matrices are only 1s because the problem already appears in this simple case):

reps = 1000
ret100 = np.zeros([1200, 2000, reps])
for k in range(reps):
    ret100[:,:,k] = 1
    print(k)

compared to the same thing, but reps=100:

reps = 100
ret100 = np.zeros([1200, 2000, reps])
for k in range(reps):
    ret100[:,:,k] = 1
    print(k)

I have two questions:

1. Why does each step in the iteration above take longer when reps = 1000, compared to 100? Is this the expected behavior?
2. Is there a better way to store a bunch of matrices than one presented above?

Answer 1

You are accessing cells in memory in a very inefficient way. Here is a simple test to see why:

a = np.zeros([1200, 2000, 100])
a[1,:,:] = 1 # time: 97.3 µs
a[:,1,:] = 1 # time: 345 µs
a[:,:,1] = 1 # time: 16 ms

The access pattern to the values of the first dimensions is not efficient. This is because the cells of the last dimension is stored contiguously in memory as opposed to the others. Accessing memory in a non-contiguous way is generally much slower (the bigger the stride, the slower it is).

So consider swapping the order of the dimensions like this:

reps = 100
ret100 = np.zeros([reps, 2000, 1200])
for k in range(reps):
    ret100[k,:,:] = 1
    print(k)

This is more than 10 times faster on my machine. The speed-up is be even bigger on larger arrays (eg. with reps set to 1000).

Answer 2

I am not sure if this is what you want, please correct me if I misunderstand.

import numpy as np
reps = 1000
ret100 = np.zeros([1200, 2000, reps])
ret100[:,:,0:reps] = 1