Consider a numpy array A of shape (3,). Then the following line
x,y,z = A
assigns x to A[0], y to A[1] and z to A[2]. Supoose now that A instead has shape s + (3,) for some arbitrary shape s. I would like to similarly assign x to A[...,0], y to A[...,1] and z to A[...,2]. The above line
x,y,z = A
does not work and gives a ValueError: not enough values to unpack (expected 3, got 2) [when A has shape (2,3)]. How can I make the desired assignment in a clean way? Obviously the following
x,y,z = A[...,0], A[...,1], A[...,2]
works but is a bit tedious if 3 is replaced by some large number.
You can use numpy.rollaxis:
x, y, z = np.rollaxis(A, -1)
Assuming this input:
array([[[0, 1, 2],
[3, 4, 5]]])
Output:
# x, y, z
(array([[0, 3]]), array([[1, 4]]), array([[2, 5]]))
This works with any position, just specify the dimension to use as the second parameter of rollaxis:
x, y, z = np.rollaxis(A, 2)
on shape = (7,6,8,2,1,4,5,9,3) ; A = np.arange(np.prod(shape)).reshape(shape) as input.
%timeit x, y, z = np.rollaxis(A, -1)
# 1.96 µs ± 59 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)
%timeit x, y, z = np.moveaxis(A, -1, 0)
# 3.78 µs ± 298 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
# credit to @QuangHoang
%timeit x, y, z = [A[...,i] for i in range(3)]
# 691 ns ± 8.24 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)