What is the best way to stack the same row to stacked 2D matrices in numpy?

user_14559854 0
Viewed 93

I am looking for a way to efficiently stack a the same row to several stacked 2D matrices.

Concretely, I am interested in stacking matrices of shape (3,4) which have float elements. For example, if there were 2 matrices to be stacked:

import numpy as np
np.arange(24.).reshape(2,3,4)

array([[[ 0.,  1.,  2.,  3.],
        [ 4.,  5.,  6.,  7.],
        [ 8.,  9., 10., 11.]],

       [[12., 13., 14., 15.],
        [16., 17., 18., 19.],
        [20., 21., 22., 23.]]])

and having a row of shape (1,1,4) to be stacked:

row = np.array([[[101.,102.,103.,104.]]])

The final result would look like this (stacked (4,4) matrices):

array([[[  0.,   1.,   2.,   3.],
        [  4.,   5.,   6.,   7.],
        [  8.,   9.,  10.,  11.],
        [101., 102., 103., 104.]],

       [[ 12.,  13.,  14.,  15.],
        [ 16.,  17.,  18.,  19.],
        [ 20.,  21.,  22.,  23.],
        [101., 102., 103., 104.]]])

Until know, the best attempt that I did is by using np.tile:

import numpy as np
M_stacked = np.arange(24.).reshape(2,3,4)
row = np.array([[[101.,102.,103.,104.]]])
np.concatenate((M_stacked, np.tile(row, (len(M_stacked),1,1))), axis=1)

But I believe this may not be the most efficient solution, specially when the number of stacked matrices increases. Is there a better approach?

Thanks in advance!

As a reference, these are the timings that I am getting:

If there are 2 stacked matrices:

M_stacked = np.arange(2.*3.*4.).reshape(2,3,4)
%timeit np.concatenate((M_stacked, np.tile(row, (len(M_stacked),1,1))), axis=1)
7.2 µs ± 85 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each))

If there are 1000 stacked matrices:

M_stacked = np.arange(1000.*3.*4.).reshape(1000,3,4)
%timeit np.concatenate((M_stacked, np.tile(row, (len(M_stacked),1,1))), axis=1)
28.8 µs ± 108 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
multidimensional-array numpy python
1 Answers

To check how good the operation is, you can perform an analytical memory throughput analysis. 1 matrix of size (1000, 3, 4) need to be read and 1 matrix of size (1000, 4, 4) need to be written. The size of the double-precision values is of 8 bytes (on standard IEEE-754 compliant systems). Since the computation is done in 28.8 us, the memory throughput is 8*(1000*3*4 + 1000*4*4) / 28.8e-6 / 1024**3 = 7.2 GiB/s which is relatively good but not great.

Since the computation is very quick, a significant part of the time spent is just the overhead of calling Numpy functions, performing internal array checks, allocating temporary Python objects (eg. tuples, Numpy views, Python integers), etc.

You can reduce the overhead using the Numba's JIT with a carefully tune configuration. Numba will generate a function that will do less checks/allocations a do them only once.

One solution do speed up the code a lot is to manually fix the size of the stacked 2D arrays with Numba. Indeed, working on variable-sized arrays is much more expensive on very small arrays because Numpy/Numba needs no iterate over the two last small dimension using inefficient loops. With a manually fix sized, Numba can unroll the loop and generate a much more efficient code (eg. using SIMD instructions). One simple and elegant way to tell to Numba the size of the matrix is to use assertions. Assertions are also useful to prevent the execution of a wrong code if the fixed size is incorrect.

Here is the resulting code:

import numpy as np
import numba as nb

@nb.njit('f8[:,:,::1](f8[:,:,::1], f8[:,:,::1])')
def compute(M_stacked, row):
    m, n, o = M_stacked.shape
    assert o == 4 # The assertion help numba to generate a much faster code
    assert n == 3
    assert row.shape == (1, 1, o)
    res = np.empty((m, n+1, o), dtype=np.float64)
    for i in range(m):
        # With the above assertions, the following loops will 
        # be unrolled by Numba to produce a very fast code.
        for j in range(n):
            for k in range(o):
                res[i, j, k] = M_stacked[i, j, k]
        for k in range(o):
            res[i, n, k] = row[0, 0, k]
    return res

row = np.array([[[101.,102.,103.,104.]]])
M_stacked = np.arange(1000.*3.*4.).reshape(1000,3,4)
%timeit compute(M_stacked, row)

Here is the results on my machine:

Numpy reference code:           19.5 us (10.7 GiB/s)
Numba code without assertions:  14.9 us (14.0 GiB/s)
Numba code with the assertions:  5.7 us (36.6 GiB/s)

The last implementation is 3.4 faster then the initial implementation. The throughput is very good for a sequential code executed in a Python interpreter.

Note that the bandwidth of my RAM is about 40 GiB/s although the matrices are probably stored in the CPU caches with a higher bandwidth. Moreover, the maximum practical throughput of a sequential Numpy code (working CPU caches) on my machine is 56 GiB/s.

user_129395570
Related