CODEWITHSUNDEEP
pythonperformancenumpyindexingvectorization
Canuck
Canuck

Reputation: 585

Creating an ordered matrix through indexing

I have a problem related to a detector that reads out the amount of photons that get fed into channels and the times they enter the detector, Lets say its channels 0 to 6 for simplicity sake. Array A will hold the channels, basically the index list, and while I am fine with counting the photons I am having troubles storing the times in a sensible container without looping (The data file is huge). So think of the array A as the index list and B as the times.

A=np.array([3,0,4,2,4,1,6])
#so this just says channel 3 got one photon, channel 0 got one, 
#channel 4 got two, 2 got one, 1 got one, channel 5 never got any so 
#it doesn't show up, and 6 got one.
B=np.array([1.2,1.6,3.,.7,.1,.05,9.])
#so here B are the times and they say (by referencing A) that channel 
#1 got a photon at .05s, channel 0 got its photon at 1.6s, channel 4 
#got a photon at 3s and another at .1s etc.
#I would like to somehow store these times in a coo sparse array or
# perhaps just a regular array that would look like:
C=np.array([[1.6,0],[.05,0],[.7,0],[1.2,0],[.1,3.0],[0,0],[.9,0]])
#the zeros could be nans of course. It would be helpful if each row 
# was ordered from earliest times to latest. This final array is
#of course ordered properly from 0 to 6 in terms of channels down
#the first axis (not in the random order that the index list was)

Not a difficult problem if you don't care about speed but unfortunately everything I do lately need to be fast. Thanks guys

Upvotes: 2

Views: 70

Answers (1)

Divakar
Divakar

Reputation: 221524

Here's a vectorized approach -

from scipy.sparse import coo_matrix

# Get sorting indices for A
n = len(A)
sidx = A.argsort()

# Use those indices to get sorted A
sA = A[sidx]

# Get shifts going from one group of identical sorted A values to another
shift_mask = np.concatenate(( [True], sA[1:] != sA[:-1] ))

# Get row indices for output array assigning
row_ids = np.zeros(n,dtype=int)
row_ids[shift_mask] = sA[shift_mask]
np.maximum.accumulate(row_ids, out=row_ids)

# Get col indices for output array assigning by using shifting mask
col_ids = intervaled_cumsum(shift_mask,trigger_val=1,start_val=0)

# Setup output sparse matrix and assign values from sorted array B
out = coo_matrix((B[sidx], (row_ids, col_ids)))

Function intervaled_cumsum is taken from here.

Sample run (on a more generic one) -

In [173]: A
Out[173]: array([3, 0, 4, 2, 4, 1, 6, 4, 2, 6])

In [174]: B
Out[174]: array([ 1.2 , 1.6 , 3.  , 0.7 , 0.1 , 0.05, 9.  , 1.5 , 2.9 , 3.1 ])

In [175]: out.toarray()
Out[175]: 
array([[ 1.6 ,  0.  ,  0.  ],
       [ 0.05,  0.  ,  0.  ],
       [ 0.7 ,  2.9 ,  0.  ],
       [ 1.2 ,  0.  ,  0.  ],
       [ 3.  ,  0.1 ,  1.5 ],
       [ 0.  ,  0.  ,  0.  ],
       [ 9.  ,  3.1 ,  0.  ]])

To explain the part that computes those shifts for the sorted A, we are making use of one-shifted slices of the sorted A to get a mask that represents shifts -

In [223]: sA # sorted A
Out[223]: array([0, 1, 2, 2, 3, 4, 4, 4, 6, 6])

In [224]: sA[1:] != sA[:-1]
Out[224]: array([ True,  True, False,  True,  True, False, False,  True, False], dtype=bool)

In [225]: np.concatenate(( [True], sA[1:] != sA[:-1] ))
Out[225]: array([ True,  True,  True, False,  True,  True, False, False,  True, False], dtype=bool)

So, correlating this output mask to sorted A, its basically all 1s except where the indices repeat.

Upvotes: 1

Related Questions