np.column_stack could be much faster? #10321

Activity

1. eric-wieser commented on Jan 3, 2018

   

   eric-wieser

   on Jan 3, 2018

   Last edited by eric-wieser

   Member

   More actions

   For reference, column_stack(tup) is implemented as

   numpy/numpy/lib/shape_base.py

   Lines 363 to 369 in c08d264

   	arrays = []
   	for v in tup:
   	arr = array(v, copy=False, subok=True)
   	if arr.ndim < 2:
   	arr = array(arr, copy=False, subok=True, ndmin=2).T
   	arrays.append(arr)
   	return _nx.concatenate(arrays, 1)
2. eric-wieser commented on Jan 3, 2018

   

   eric-wieser

   on Jan 3, 2018

   Member

   More actions

   Can you give your exact benchmark code and results?
3. toobaz commented on Jan 4, 2018

   

   toobaz

   on Jan 4, 2018

   Author

   More actions

   I had only made sparse attempts, but here is a more systematic one (showing that the ~25% slowdown was optimistic - it can reach 60% - but the ~4x gain was pessimistic - it can reach 7x).

   In [1]: import numpy as np
   
   In [2]: def my_column_stack(tup):
      ...:     arrays = [np.asarray(arr) for arr in tup]
      ...:     shape = arrays[0].shape
      ...:     for v in arrays:
      ...:         if v.shape != shape:
      ...:             raise ValueError("All input arrays must have the same shape")
      ...:     res = np.array(arrays).T
      ...:     if len(shape) > 1:
      ...:         return res[0]
      ...:     return res
      ...: 
   
   In [3]: for N in 3, 10, 20, 100, 1000, 10000:
      ...:     print("N =", N)
      ...:     print("1d")
      ...:     l = [np.arange(N)]*N
      ...:     %timeit my_column_stack(l)
      ...:     %timeit np.column_stack(l)
      ...: 
      ...:     print("2d")
      ...:     l = [np.array(np.arange(N), ndmin=2).T]*N
      ...:     %timeit my_column_stack(l)
      ...:     %timeit np.column_stack(l)
      ...:     
      ...:     print()
      ...:     
   N = 3
   1d
   4.95 µs ± 48.6 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
   5.32 µs ± 28.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
   2d
   5.39 µs ± 74.6 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
   3.31 µs ± 27.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
   
   N = 10
   1d
   11.6 µs ± 183 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
   16.3 µs ± 1.81 µs per loop (mean ± std. dev. of 7 runs, 100000 loops each)
   2d
   13 µs ± 261 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
   8.91 µs ± 222 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
   
   N = 20
   1d
   21.1 µs ± 367 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
   31.3 µs ± 213 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
   2d
   23.8 µs ± 290 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
   17 µs ± 258 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
   
   N = 100
   1d
   96.5 µs ± 340 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
   157 µs ± 1.14 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
   2d
   109 µs ± 979 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
   83.9 µs ± 553 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
   
   N = 1000
   1d
   1.61 ms ± 39.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
   6.27 ms ± 107 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
   2d
   1.66 ms ± 14.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
   5.33 ms ± 35.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
   
   N = 10000
   1d
   343 ms ± 7.52 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
   2.39 s ± 21.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
   2d
   345 ms ± 2.75 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
   2.39 s ± 34.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
   
4. eric-wieser commented on Jan 4, 2018

   

   eric-wieser

   on Jan 4, 2018

   Member

   More actions

   I don't think your implementation does the same thing - for better or worse, column_stack accepts a mixture of 1d and 2d arrays
5. toobaz commented on Jan 4, 2018

   

   toobaz

   on Jan 4, 2018

   Author

   More actions

   You are right. The fact that np.column_stack([[1,2], [3,4,5]]) raises ValueError: all the input array dimensions except for the concatenation axis must match exactly made me think this case was supported by mistake. But the docs explicitly say "All of them must have the same first dimension.", so the behavior you mention is even documented. Sorry for the noise.
6. toobaz

   closed this as completedon Jan 4, 2018