Nested list comprehension in Python

Question

I've got a list comprehension I'm trying to get my head around and I just can't seem to get what I'm after and thought I'd see if anybody else knew how!

My basic data structure is this:

structure = [[np.array([[1,2,3],[4,5,6]]), np.array([[7,8,9],[10,11,12]])], [np.array([[13,14,15],[16,17,18]]), np.array([[19,20,21],[22,23,24]])]]

So I've got an overall list containing sublists of numpy arrays and my desired output is some sort of grouping (don't care if it's a list or an array) with the following elements paired:

[1, 13]
[4, 16]
[2, 14]
[5, 17]
[3, 15]
[6, 18]

I thought I'd got it with the following style construct:

output = [structure[i][0][j] for j in range(9) for i in range(len(structure))] but alas, no joy.

I don't really mind if it needs more than one stage - just want to get those elements grouped together!

(as a bit of background - I've got lists of probabilities outputted from various models and within those models I've got a training list and a validation list:

[[model_1], [model_2], ..., [model_n]]

where [model_1] is [[training_set], [validation_set], [test_set]]

and [training_set] is np.array([p_1, p_2, ..., p_n],[p_1, p_2, ..., p_n],...])

I'd like to group together the prediction for item 1 for each of the models and create a training vector out of it of length equal to the number of models I've got. I'd then like to do the same but for the second row of [training_set].

If that doesn't make sense let me know!

Answer 1

Since all the arrays (and sublists) in structure are the same size you can turn it into one higher dimensional array:

In [189]: A=np.array(structure)
Out[189]: 
array([[[[ 1,  2,  3],
         [ 4,  5,  6]],

        [[ 7,  8,  9],
         [10, 11, 12]]],


       [[[13, 14, 15],
         [16, 17, 18]],

        [[19, 20, 21],
         [22, 23, 24]]]])

In [190]: A.shape
Out[190]: (2, 2, 2, 3)

Reshaping and swapaxes can give you all kinds of combinations.

For example, the values in your sample sublist can be selected with:

In [194]: A[:,0,:,:]
Out[194]: 
array([[[ 1,  2,  3],
        [ 4,  5,  6]],

       [[13, 14, 15],
        [16, 17, 18]]])

and reshape to get

In [197]: A[:,0,:,:].reshape(2,6)
Out[197]: 
array([[ 1,  2,  3,  4,  5,  6],
       [13, 14, 15, 16, 17, 18]])

and transpose to get the 6 rows of pairs:

In [198]: A[:,0,:,:].reshape(2,6).T
Out[198]: 
array([[ 1, 13],
       [ 2, 14],
       [ 3, 15],
       [ 4, 16],
       [ 5, 17],
       [ 6, 18]])

To get them in the 1,4,2,5.. order I can transpose first

In [208]: A[:,0,:,:].T.reshape(6,2)
Out[208]: 
array([[ 1, 13],
       [ 4, 16],
       [ 2, 14],
       [ 5, 17],
       [ 3, 15],
       [ 6, 18]])