Reputation: 602
I have a function that accepts user parameter k and generates a matrix of the form:
a = np.array((k,2,5))
Leaving the function I would like to return k 2x5 matrices. Essentially,
if k == 2:
return a[0,:,:], a[1,:,:]
if k == 3:
return a[0,:,:], a[1,:,:], a[2,:,:]
Is there a way to do this dynamically without having to hard-code each of the possible values of k?
Thanks.
Upvotes: 1
Views: 65
Reputation: 221774
Here's an approach using np.array_split
-
np.array_split(a[:k].reshape(-1,a.shape[-1]),k)shorter
Even shorter with np.vsplit
-
np.vsplit(a[:k].reshape(-1,a.shape[-1]), k)
Please note that if you are assigning the output from the relevant function to k
number of outputs, you don't need all this mess and can simply return the slice . i.e. do : return a[:k]
.
As an example :
In [595]: a = np.random.rand(4,2,5)
In [596]: k = 3
In [597]: a0,a1,a2 = a[:k]
In [598]: np.allclose(a0, a[0,:,:]) # Verify for all three outputs
Out[598]: True
In [599]: np.allclose(a1, a[1,:,:])
Out[599]: True
In [600]: np.allclose(a2, a[2,:,:])
Out[600]: True
Upvotes: 2