Background
I am trying to calculate the "social traveler" distance between a set of points. For two points a and b, we define ST(a,b) as the cumulative sum of euclidean distances between a and all points closer to a than b, up to and including b.
(Imagine a is making various trips to see people. They first see the nearest person, then the next nearest, and so on. The social-traveler distance is how much they will have traveled (counting only the outgoing distance) by the time they reach b.)
Social traveler distance is defined up to some arbitrary radius; beyond that the "visited" point is defined as unreachable. This is because reasons. ;)
What I have done so far
I have a numpy array with points. A toy example:
>>> import numpy as np
>>> x = np.array([(0,0), (-1,2), (-2,-2), (6,-2), (4,0)])
>>> x
array([[ 0, 0],
[-1, 2],
[-2, -2],
[ 6, -2],
[ 4, 0]])
I generate an array with the pairwise distances between each point. For this example, distance 4 is the upper limit I consider.
>>> from sklearn.metrics import pairwise_distances
>>> y = pairwise_distances(x)
>>> y
array([[0. , 2.24, 2.83, 6.32, 4. ],
[2.24, 0. , 4.12, 8.06, 5.39],
[2.83, 4.12, 0. , 8. , 6.32],
[6.32, 8.06, 8. , 0. , 2.83],
[4. , 5.39, 6.32, 2.83, 0. ]])
>>> import numpy.ma as ma
>>> ym = ma.masked_greater(y, 4)
>>> ym
masked_array(
data=[[0.0 , 2.24, 2.83, --, 4.0 ],
[2.24, 0.0 , --, --, --],
[2.83, --, 0.0 , --, --],
[ --, --, --, 0.0 , 2.83],
[4.0 , --, --, 2.83, 0.0 ]],
mask=[[False, False, False, True, False],
[False, False, True, True, True],
[False, True, False, True, True],
[ True, True, True, False, False],
[False, True, True, False, False]],
fill_value=1e+20)
(I have truncated numbers for everyone's sanity.)
My goal here is to replace each row element with the cumulative sum of the row's non-missing values up to and including that element. (For simplicity, ignore the issue of ties.) That is, I want to get this array:
>>> hypothetical_new_y
array([[0. , 2.24, 5.06, nan, 9.06],
[2.24, 0. , nan, nan, nan],
[2.83, nan, 0. , nan, nan],
[nan , nan, nan, 0. , 2.83],
[6.83, nan, nan, 2.83, 0. ]])
I know how to get the cumulative sum along the rows:
>>> np.cumsum(ym, axis=1)
masked_array(
data=[[0.0 , 2.24, 5.06, --, 9.06],
[2.24, 2.24, --, --, --],
[2.83, --, 2.83, --, --],
[ --, --, --, 0.0 , 2.83],
[4.0 , --, --, 6.83, 6.83]],
...
Because this is adding based on the row order rather than sorting, it gives incorrect values (except by coincidence on the first row here). I can do the same, first sorting the rows:
>>> np.cumsum(np.sort(ym, axis=1), axis=1)
masked_array(
data=[[0.0 , 2.24, 5.06, 9.06, --],
[0.0 , 2.24, --, --, --],
[0.0 , 2.83, --, --, --],
[0.0 , 2.83, --, --, --],
[0.0 , 2.83, 6.83, --, --]],
mask=[[False, False, False, False, True],
[False, False, True, True, True],
[False, False, True, True, True],
[False, False, True, True, True],
[False, False, False, True, True]],
fill_value=1e+20)
This gives me the correct values, but they are sorted in ascending value rather than in the order I want. I want these cumulative sums positioned in accordance with the original array's argsort values:
>>> np.argsort(ym)
array([[0, 1, 2, 4, 3],
[1, 0, 2, 3, 4],
[2, 0, 1, 3, 4],
[3, 4, 0, 1, 2],
[4, 3, 0, 1, 2]])
It seems I need a way to sort these rows, do the cumulative sums, and then return the cumulative sums sorted by those argsort vectors. Notice that if you present the results of np.cumsum(np.sort(ym, axis=1), axis=1) ordered by the vectors in np.argsort(ym), you have my hypothethical_new_y array. But it's this last step where I'm falling on my face.
I feel like I have two-thirds of a Venn diagram here. Any thoughts on how I can take this last step? My hope is that this is a relatively simple gap in my numpy-syntax knowledge that I've uncovered.
EDIT: After posting this and continuing to dig, I thought I had found the "obvious" answer with np.take_along_axis()--but no. Consider this code:
>>> foo = np.argsort(ym)
>>> bar = np.cumsum(np.sort(ym, axis=1), axis=1)
>>> np.take_along_axis(bar, foo, axis=1)
masked_array(
data=[[0.0 , 2.24, 5.06, --, 9.06],
[2.24, 0.0 , --, --, --],
[ --, 0.0 , 2.83, --, --],
[ --, --, 0.0 , 2.83, --],
[ --, --, 0.0 , 2.83, 6.83]],
...
If you look at the last row, for example, it is clear that the command is taking the 4th element of the last row of bar and putting it in the first spot, the 3rd element of the last row of bar and putting it in the second spot, etc. (Recall that the last row of foo is [4, 3, 0, 1, 2].) What I want instead is to take the 1st element of the last row of bar and put it in the fourth spot, the 2nd element of bar and put it in the third spot, etc. And np.put_along_axis isn't my jam, either.