Flatten numpy array with sub-arrays of different dimensions

Question

This seems like a simple enough task, but I haven't found how to do it using numpy. Consider the example array:

import numpy as np
aa = np.array([np.array([13.16]), np.array([1.58 , 1.2]), np.array([13.1]), np.array([1. , 2.6])], dtype=object)

I need a general way to flatten that array into a single array of N elements, with N=every float in all the sub-arrays. In this case it would be:

aa = np.array([13.16, 1.58 , 1.2, 13.1, 1. , 2.6])

I've tried np.ndarray.flatten() (tried all the 'order' options)) but I get back the same unchanged aa array.

Why is np.ndarray.flatten() not working and how can I accomplish this?

The solution should be as general as possible since the example aa array I'm using here will actually be filled with sub-arrays of different lengths in my real code.

Answer 1

You can use numpy.hstack

>>> np.hstack(aa)
array([13.16,  1.58,  1.2 , 13.1 ,  1.  ,  2.6 ])

Answer 2

If all sub-arrays are 1D, you can also use np.concatenate to save a bit of time. As far as I understand np.hstack is a wrapper of np.concatenate that reduces to the same operation if the input is 1D.

np.concatenate(aa)
# array([13.16,  1.58,  1.2 , 13.1 ,  1.  ,  2.6 ])

np.allclose(np.concatenate(aa), np.hstack(aa))
# True

%timeit np.hstack(aa)
5.53 µs ± 28.2 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
%timeit np.concatenate(aa)
2.2 µs ± 46.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)