Group and summarize matrix by value

Question

I have two matrices, prob and totalHigh both of shape axbxcxd. a and b are coordinates. Here are two samples:

In [77]: prob[1,1,:]
Out[77]: 
array([[ 0.09,  0.01,  0.  ,  0.  ,  0.  ],
       [ 0.81,  0.09,  0.  ,  0.  ,  0.  ],
       [ 0.  ,  0.  ,  0.  ,  0.  ,  0.  ]])

In [78]: totalHigh[1,1,:]
Out[78]: 
array([[0, 1, 2, 3, 4],
       [1, 2, 3, 4, 5],
       [2, 3, 4, 5, 6]])

totalHigh contains information about outcomes, unfortunately on two dimensions. Correspondingly, prob contains probabilities of these outcomes. For example, the total probability of outcome 1, at coordinates 1,1 is 0.01+0.81.

How can I remove the redundant dimension?

Expected Outcome

simplifiedHigh[1,1,:]
array([0, 1, 2, 3, 4, 5, 6])
simplifiedProb[1,1,:]
array([0.09, 0.82, 0.09, 0, 0, 0, 0])

How do I get that in the most efficient way?

Answer 1

You can use np.bincount and np.unique -

IDs = np.unique(totalHigh_sliced)
counts = np.bincount(totalHigh_sliced.ravel(),prob_sliced.ravel())

Sample run -

In [215]: prob_sliced
Out[215]: 
array([[ 0.09,  0.01,  0.  ,  0.  ,  0.  ],
       [ 0.81,  0.09,  0.  ,  0.  ,  0.  ],
       [ 0.  ,  0.  ,  0.  ,  0.  ,  0.  ]])

In [216]: totalHigh_sliced
Out[216]: 
array([[0, 1, 2, 3, 4],
       [1, 2, 3, 4, 5],
       [2, 3, 4, 5, 6]])

In [217]: IDs = np.unique(totalHigh_sliced)
     ...: counts = np.bincount(totalHigh_sliced.ravel(),prob_sliced.ravel())
     ...: 

In [218]: IDs
Out[218]: array([0, 1, 2, 3, 4, 5, 6])

In [219]: counts
Out[219]: array([ 0.09,  0.82,  0.09,  0.  ,  0.  ,  0.  ,  0.  ])