Vectorizing for loop using splicing in NumPy

Question

I have this for loop:

blockSize = 5
ds = np.arange(20)
ds = np.reshape(ds, (1, len(ds))
counts = np.zeros(len(ds[0]/blockSize))

for i in range(len(counts[0])):
    counts[0, i] = np.floor(np.sum(ds[0, i*blockSize:i*blockSize+blockSize]))

I am trying to vectorize it, doing something like this:

countIndices = np.arange(len(counts[0]))
counts[0, countsIndices] = np.floor(np.sum(ds[0, countIndices*blockSize:countIndices*blockSize + blockSize]))

However, this does not work and gives this error:

counts[0, countIndices] = np.floor(np.sum(ds[0, countIndices*blockSize:countIndices*blockSize + blockSize]))
TypeError: only integer scalar arrays can be converted to a scalar index

I know that something like this works:

counts[0, countIndices] = np.floor(ds[0, countIndices*blockSize]
                         + ds[0, countIndices*blockSize + 2] +
                         ... ds[0, countIndices*blockSize + blockSize])

The issue is that for large values of blocksize (which blocksize is very large in my actual code), this is not feasible to implement. I am confused as to how to accomplish what I want. Any help is greatly appreciated.

Show any and all errors with full traceback. Did you mean `counts = np.zeros(len(ds[0])//blockSize)` instead of `counts = np.zeros(len(ds[0]/blockSize))`? — Mad Physicist, Jul 20 '21 at 18:47

score 1 · Accepted Answer · answered Jul 20 '21 at 18:53

You don't need to do floor if you store the result to an integer array. You can also create a fake new axis of size blockSize to fully vectorize your operation.

block_size = 5
ds = np.arange(80.0).reshape(4, -1)   # Shape (4, 20)
counts = np.empty((ds.shape[0], ds.shape[1] // block_size), dtype=int)

To introduce the fake dimension and sum:

ds.reshape(ds.shape[0], -1, block_size).sum(axis=-1, out=counts)

Reshaping does not copy the data, so the operation ds.reshape(ds.shape[0], -1, block_size) is extremely cheap in both time and space.

You can use -1 for one of the reshape dimensions to avoid computing/writing out long division expressions.

Vectorizing for loop using splicing in NumPy

1 Answers1