How to extract columns from an indexed matrix?

Question

I have the following matrix:

M = np.matrix([[1,2,3,4,5,6,7,8,9,10], 
               [11,12,13,14,15,16,17,18,19,20], 
               [21,22,23,24,25,26,27,28,29,30]])

And I receive a vector indexing the columns of the matrix:

index = np.array([1,1,2,2,2,2,3,4,4,4])

This vector has 4 different values, so my objective is to create a list containing four new matrices so that the first matrix is made by the first two columns of M, the second matrix is made by columns 3 to 6 and so on:

M1 = np.matrix([[1,2],[11,12],[21,22]])
M2 = np.matrix([[3,4,5,6],[13,14,15,16],[23,24,25,26]])
M3 = np.matrix([[7],[17],[27]])
M4 = np.matrix([[8,9,10],[18,19,20],[28,29,30]])
l = list(M1,M2,M3,M4)

I need to do this in a automated way, since the number of rows and columns of M as well as the indexing scheme are not fixed. How can I do this?

score 1 · Accepted Answer · answered Apr 07 '18 at 12:21

There are 3 points to note:

For a variable number of variables, as in this case, the recommended solution is to use a dictionary.
You can use simple numpy indexing for the individual case.
Unless you have a very specific reason, use numpy.array instead of numpy.matrix.

Combining these points, you can use a dictionary comprehension:

d = {k: np.array(M[:, np.where(index==k)[0]]) for k in np.unique(index)}

Result:

{1: array([[ 1,  2],
           [11, 12],
           [21, 22]]),
 2: array([[ 3,  4,  5,  6],
           [13, 14, 15, 16],
           [23, 24, 25, 26]]),
 3: array([[ 7],
           [17],
           [27]]),
 4: array([[ 8,  9, 10],
           [18, 19, 20],
           [28, 29, 30]])}

score 0 · Answer 2 · answered Apr 07 '18 at 11:30

import numpy as np

M = np.matrix([[1,2,3,4,5,6,7,8,9,10], 
               [11,12,13,14,15,16,17,18,19,20], 
               [21,22,23,24,25,26,27,28,29,30]])

index = np.array([1,1,2,2,2,2,3,4,4,4])

m = [[],[],[],[]]
for i,c in enumerate(index):
  m[k-1].append(c)

for idx in m:
  print M[:,idx]

this is a little hard coded, I assumed you will always want 4 matrixes and such.. you can change it for more generalisation

How to extract columns from an indexed matrix?

2 Answers2