CODEWITHSUNDEEP
pythonnumpy
gutelfuldead
gutelfuldead

Reputation: 602

Return multiple matrices based on user input to function

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

Answers (2)

Divakar
Divakar

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

Alfe
Alfe

Reputation: 59616

Sure:

return tuple(a[i,:,:] for i in range(k))

Python's great! 8-))

Upvotes: 2

Related Questions