从numpy数组列表创建numpy数组的Pythonic方法

Question

我在一个循环里生成了一系列一维的numpy数组，之后再把这些数组放到一个列表里，最后把这个列表转换成一个二维的numpy数组。如果我提前知道要生成多少个数组，我可以直接创建一个二维的numpy数组，但我不知道，所以只能先把它们放在列表里。

下面是一个示例：

>>> list_of_arrays = map(lambda x: x*ones(2), range(5))
>>> list_of_arrays
[array([ 0.,  0.]), array([ 1.,  1.]), array([ 2.,  2.]), array([ 3.,  3.]), array([ 4.,  4.])]
>>> arr = array(list_of_arrays)
>>> arr
array([[ 0.,  0.],
       [ 1.,  1.],
       [ 2.,  2.],
       [ 3.,  3.],
       [ 4.,  4.]])

我想问的是：

有没有比把这些数组放在列表里再转换成numpy数组更好的方法（从性能上考虑）来收集连续的数字数据（在我的例子中是numpy数组）？有没有什么“可扩展”的矩阵数据结构可以用在经过验证的模块里？

我二维矩阵的典型大小大约在100x10到5000x10之间，都是浮点数。

补充：在这个例子中我使用了map，但在我的实际应用中我用的是for循环。

Answer 1

比@Gill Bates的回答还简单，这里有一行代码：

np.stack(list_of_arrays, axis=0)

Answer 2

这里有一个方便的方法，可以使用 numpy.concatenate。我觉得这个方法的速度也比@unutbu的回答要快：

In [32]: import numpy as np 

In [33]: list_of_arrays = list(map(lambda x: x * np.ones(2), range(5)))

In [34]: list_of_arrays
Out[34]: 
[array([ 0.,  0.]),
 array([ 1.,  1.]),
 array([ 2.,  2.]),
 array([ 3.,  3.]),
 array([ 4.,  4.])]

In [37]: shape = list(list_of_arrays[0].shape)

In [38]: shape
Out[38]: [2]

In [39]: shape[:0] = [len(list_of_arrays)]

In [40]: shape
Out[40]: [5, 2]

In [41]: arr = np.concatenate(list_of_arrays).reshape(shape)

In [42]: arr
Out[42]: 
array([[ 0.,  0.],
       [ 1.,  1.],
       [ 2.,  2.],
       [ 3.,  3.],
       [ 4.,  4.]])

Answer 3

假设你知道最终的数组 arr 最大不会超过 5000x10 的大小。这样的话，你可以提前分配一个最大大小的数组，然后在循环中逐步填充数据，最后在循环结束后用 arr.resize 来调整数组的大小，去掉多余的部分。

下面的测试表明，这种做法比起在构建过程中使用中间的 Python 列表，无论最终数组的大小如何，速度都会稍微快一些。

另外，arr.resize 会释放掉未使用的内存，所以最终的内存占用（虽然可能不是中间的）会比 python_lists_to_array 使用的要小。

这表明 numpy_all_the_way 的速度更快：

% python -mtimeit -s"import test" "test.numpy_all_the_way(100)"
100 loops, best of 3: 1.78 msec per loop
% python -mtimeit -s"import test" "test.numpy_all_the_way(1000)"
100 loops, best of 3: 18.1 msec per loop
% python -mtimeit -s"import test" "test.numpy_all_the_way(5000)"
10 loops, best of 3: 90.4 msec per loop

% python -mtimeit -s"import test" "test.python_lists_to_array(100)"
1000 loops, best of 3: 1.97 msec per loop
% python -mtimeit -s"import test" "test.python_lists_to_array(1000)"
10 loops, best of 3: 20.3 msec per loop
% python -mtimeit -s"import test" "test.python_lists_to_array(5000)"
10 loops, best of 3: 101 msec per loop

这表明 numpy_all_the_way 使用的内存更少：

% test.py
Initial memory usage: 19788
After python_lists_to_array: 20976
After numpy_all_the_way: 20348

test.py:

import numpy as np
import os


def memory_usage():
    pid = os.getpid()
    return next(line for line in open('/proc/%s/status' % pid).read().splitlines()
                if line.startswith('VmSize')).split()[-2]

N, M = 5000, 10


def python_lists_to_array(k):
    list_of_arrays = list(map(lambda x: x * np.ones(M), range(k)))
    arr = np.array(list_of_arrays)
    return arr


def numpy_all_the_way(k):
    arr = np.empty((N, M))
    for x in range(k):
        arr[x] = x * np.ones(M)
    arr.resize((k, M))
    return arr

if __name__ == '__main__':
    print('Initial memory usage: %s' % memory_usage())
    arr = python_lists_to_array(5000)
    print('After python_lists_to_array: %s' % memory_usage())
    arr = numpy_all_the_way(5000)
    print('After numpy_all_the_way: %s' % memory_usage())