getting elements of array in python with a single line

Question

I need to work with this beautiful np array

import numpy as np
train_predicteds = np.asarray([ 
 [[0.1, 0.2, 0.3], [0.5, 0.6, 0.7], [0.7, 0.8, 0.9]],
 [[0.3, 0.1, 0.4], [0.4, 0.5, 0.6], [0.5, 0.6, 0.1]]])

Now I want to get the elements in this fashion:

[[0.1, 0.3], [0.2, 0.1], [0.3, 0.4],
 [0.5, 0.4], [0.6, 0.5], [0.7, 0.6],
 [0.7, 0.5], [0.8, 0.6], [0.9, 0.1]]

Some sort of solution I found was using these two lines:

aux = [item[0] for item in train_predicteds]
x = [item[0] for item in aux]

Which produces me x equals to

[0.10000000000000001, 0.30000000000000001]

But I can't merge these two lines in a single one, is it possible? Or is there a better pythonic solution?

Thanks guys

Answer 1

Starting with:

In [17]: arr = np.asarray([  
    ...:  [[0.1, 0.2, 0.3], [0.5, 0.6, 0.7], [0.7, 0.8, 0.9]], 
    ...:  [[0.3, 0.1, 0.4], [0.4, 0.5, 0.6], [0.5, 0.6, 0.1]]])                 
In [18]: arr                                                                    
Out[18]: 
array([[[0.1, 0.2, 0.3],
        [0.5, 0.6, 0.7],
        [0.7, 0.8, 0.9]],

       [[0.3, 0.1, 0.4],
        [0.4, 0.5, 0.6],
        [0.5, 0.6, 0.1]]])
In [19]: arr.shape                                                              
Out[19]: (2, 3, 3)

After trying several transpose orders I got:

In [26]: arr.transpose(1,2,0)        # shape (3,3,2) moves 1st dim to end                                          
Out[26]: 
array([[[0.1, 0.3],
        [0.2, 0.1],
        [0.3, 0.4]],

       [[0.5, 0.4],
        [0.6, 0.5],
        [0.7, 0.6]],

       [[0.7, 0.5],
        [0.8, 0.6],
        [0.9, 0.1]]])

The first two dimensions can be merged with reshape:

In [27]: arr.transpose(1,2,0).reshape(9,2)                                      
Out[27]: 
array([[0.1, 0.3],
       [0.2, 0.1],
       [0.3, 0.4],
       [0.5, 0.4],
       [0.6, 0.5],
       [0.7, 0.6],
       [0.7, 0.5],
       [0.8, 0.6],
       [0.9, 0.1]])

Answer 2

The better Pythonic solution

>>> train_predicteds[:,0,0]
array([0.1, 0.3])