Mr.Owl
Mr.Owl

Reputation: 44

Numpy - Fastest way to create a integer index matrix

Problem:

I have an array that represents products, lets say 3 for example

prod = [1,2,3]

then I have a correlation matrix for those products (just a number that represents something between two products, lets call c_ij for simplicity), in this case a 3x3 matrix

corr = [[c_11,c_12,c_13],
        [c_21,c_22,c_23],
        [c_31,c_32,c_33]]

The problem is that a need to shuffle the prod array, then I need to shuffle the corr matrix in a way that corr[i,j] still represent the correlation between prod[i] and prod[j]

My solution:

I know I can use a integer array as index to shuffle multiple array in the same way, like this:

order = [0,1,2]
new_order = np.random.permutation(order) # [2,0,1] for example
shuf_prod = prod[new_order]

Looking in the web I find that to make this work in a matrix I need to transform the order array in a matrix like

new_order = [2,0,1]
new_index = [ [[2,2,2],[0,0,0],[1,1,1]],
              [[2,0,1],[2,0,1],[2,0,1]] ]

new_corr = corr[tuple(new_index)]
# this output what I want that is:
# [[c_33,c_31,c_32],
#  [c_13,c_11,c_12],
#  [c_23,c_21,c_22]]

Question:

  1. The entire solution of shuffling look chunky and not efficient, this is a performance critical application so there is a faster way to do this? (I don't really care for simplicity of code, just performance)
  2. If this is a good way of doing this, how I can create the new_index matrix from new_order array?

EDIT: Michael Szczesny solved the problem

new_corr = corr[new_order].T[new_order].T

Upvotes: 0

Views: 137

Answers (1)

Alain T.
Alain T.

Reputation: 42133

you can use the indices directly as subscripts to the matrix as long as you provide the right shape for the second axis:

import numpy as np

mat = np.array([[3,4,5],
                [4,8,9],
                [5,9,7]])

order = np.array([2,0,1])

mat[order,order[:,None]]
array([[7, 5, 9],
       [5, 3, 4],
       [9, 4, 8]])

Upvotes: 1

Related Questions