NumPy: Create a multidimensional array from an iterable

Viewed 631

I have an iterable of tuples, and I'd like to build an ndarray from it. Say that the shape would be (12345, 67890). What would be an efficient and elegant way to do so?

Here are a few options, and why I ruled them out:

  1. np.array(my_tuples) starts allocating the array before it knows the size, which requires inefficient relocations according to NumPy's documentation.

  2. Create an array with uninitialized content using np.ndarray((12345, 67890)) and then do a loop that populates it with data. It works and it's efficient, but a bit inelegant because it requires multiple statements.

  3. Use np.fromiter which appears to be geared towards 1-dimensional arrays only.

Does anyone have a better solution?

(I've seen this question, but I'm not seeing any promising answers there.)

3 Answers

Define a generator:

def foo(m,n):
    for i in range(m):
        yield list(range(i,i+n))

timing several alternatives:

In [93]: timeit np.array(list(foo(3000,4000)))                                  
1.74 s ± 17.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [94]: timeit list(foo(3000,4000))                                            
663 ms ± 3.84 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [95]: timeit np.stack([np.array(row) for row in foo(3000,4000)])             
1.32 s ± 2.04 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [96]: timeit np.concatenate([np.array(row, ndmin=2) for row in foo(3000,4000)
    ...: ])                                                                     
1.33 s ± 23.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [97]: %%timeit  
    ...: arr = np.empty((3000,4000),int) 
    ...: for i,row in enumerate(foo(3000,4000)): 
    ...:     arr[i] = row 
    ...:                                                                        
1.29 s ± 3.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

and with a flat generator:

def foo1(m,n):
    for i in range(m):
        for j in range(n):
            yield i+j
In [104]: timeit np.fromiter(foo1(3000,4000),int).reshape(3000,4000)            
1.54 s ± 5.64 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Use fromiter() with .reshape(). Reshaping does not require more memory or processing.

I suspect you'll find this not elegant enough, but fast it is:

from timeit import timeit
import itertools as it 

def x():
   for i in range(3000):
       yield list(range(i,i+4000))

timeit(lambda:np.fromiter(it.chain.from_iterable(x()),int,12000000).reshape(3000,4000),number=10)
# 5.048861996969208

Compare that to, for example

timeit(lambda:np.concatenate(list(x()),0),number=10)
# 12.466914481949061

Btw. if you do not know the total number of elements in advance, no big deal:

timeit(lambda:np.fromiter(it.chain.from_iterable(x()),int).reshape(3000,-1),number=10)
# 5.331893905065954
Related