Is there unstack in NumPy?

user_941531 0
Viewed 6914
numpy python
4 Answers

Coming across this late, here is a much simpler answer:

def unstack(a, axis=0):
    return np.moveaxis(a, axis, 0)
#    return list(np.moveaxis(a, axis, 0))

As a bonus, the result is still a numpy array. The unwrapping happens if you just python-unwrap it:

A, B, = unstack([[1, 2], [3, 4]], axis=1)
assert list(A) == [1, 3]
assert list(B) == [2, 4]

Unsurprisingly, it is also the fastest:

# np.squeeze
❯ python -m timeit -s "import numpy as np; a=np.array(np.meshgrid(np.arange(1000), np.arange(1000)));" "C = [np.squeeze(e, 1) for e in np.split(a, a.shape[1], axis = 1)]"
100 loops, best of 5: 2.64 msec per loop
    
# np.take
❯ python -m timeit -s "import numpy as np; a=np.array(np.meshgrid(np.arange(1000), np.arange(1000)));" "C = [np.take(a, i, axis = 1) for i in range(a.shape[1])]"       
50 loops, best of 5: 5.08 msec per loop

# np.moveaxis
❯ python -m timeit -s "import numpy as np; a=np.array(np.meshgrid(np.arange(1000), np.arange(1000)));" "C = np.moveaxis(a, 1, 0)"
100000 loops, best of 5: 3.89 usec per loop

# list(np.moveaxis)
❯ python -m timeit -s "import numpy as np; a=np.array(np.meshgrid(np.arange(1000), np.arange(1000)));" "C = list(np.moveaxis(a, 1, 0))"
1000 loops, best of 5: 205 usec per loop
user_5038220

Thanks to suggestion of @hpaulj solved task efficiently using np.split.

Try it online!

import numpy as np

def unstack(a, axis = 0):
    return [np.squeeze(e, axis) for e in np.split(a, a.shape[axis], axis = axis)]

a = [np.array([[1,2,3],[4,5,6]]), np.array([[7,8,9],[10,11,12]])]

for axis in range(len(a[0].shape) + 1):
    b = np.stack(a, axis)
    c = unstack(b, axis)
    # Check that we have same "c" as input "a"
    assert len(c) == len(a) and all(np.all(sc == sa) for sc, sa in zip(c, a)), (c, a)
user_9415310

it's simple like:

list(array)

it will iterate over the first axis.

an example:

a = np.array([1,2,3])
b = np.array([4,5,6])
c = np.stack([a,b])
d = list(c)

Output of c

array([[1, 2, 3],
       [4, 5, 6]])

Output of d

[array([1,2,3]), array([4,5,6])]
user_144080780

Can also be solved using np.take.

It is a simpler than this solution but could be a bit less efficient, needs timing measurement for different inputs.

Try it online!

import numpy as np

def unstack(a, axis = 0):
    return [np.take(a, i, axis = axis) for i in range(a.shape[axis])]

a = [np.array([[1,2,3],[4,5,6]]), np.array([[7,8,9],[10,11,12]])]

for axis in range(len(a[0].shape) + 1):
    b = np.stack(a, axis)
    c = unstack(b, axis)
    # Check that we have same "c" as input "a"
    assert len(c) == len(a) and all(np.all(sc == sa) for sc, sa in zip(c, a)), (c, a)
user_9415310
Related