CODEWITHSUNDEEP
pythonnumpyvectorization
FooBar
FooBar

Reputation: 16508

Cumsum and vectorized slicing

I have a matrix J of size (j,v) and a vector JTildeIDX of size v. The vector contains the start for a cumsum operation over J. That is, given

>>> JTildeIDX
array([0, 0, 9, 9, 9])
>>> J
array([[ 1.        ,  1.        ,  1.        ,  1.        ,  1.        ],
       [ 1.11111111,  1.11111111,  1.11111111,  1.11111111,  1.11111111],
       [ 1.22222222,  1.22222222,  1.22222222,  1.22222222,  1.22222222],
       [ 1.33333333,  1.33333333,  1.33333333,  1.33333333,  1.33333333],
       [ 1.44444444,  1.44444444,  1.44444444,  1.44444444,  1.44444444],
       [ 1.55555556,  1.55555556,  1.55555556,  1.55555556,  1.55555556],
       [ 1.66666667,  1.66666667,  1.66666667,  1.66666667,  1.66666667],
       [ 1.77777778,  1.77777778,  1.77777778,  1.77777778,  1.77777778],
       [ 1.88888889,  1.88888889,  1.88888889,  1.88888889,  1.88888889],
       [ 2.        ,  2.        ,  2.        ,  2.        ,  2.        ]])

I want to get an output of size v == 5, which contains the following cumsum:

np.array([1+1.1+....+2, 1+1.1+....+2, 2, 2, 2])

How would I achieve this?

Upvotes: 1

Views: 148

Answers (1)

Divakar
Divakar

Reputation: 221614

One approach with braodcasting and boolean indexing -

import numpy as np

# Mask of elements from J to be summed up
mask = np.arange(J.shape[0])[:,None] >= JTildeIDX

# Map the mask onto input array and sum along columns for final output
out = (J*mask).sum(0)

Sample run -

In [154]: J
Out[154]: 
array([[0, 7, 2, 1, 3],
       [0, 8, 4, 3, 1],
       [8, 5, 1, 8, 7],
       [6, 6, 0, 3, 8],
       [6, 5, 6, 4, 2],
       [8, 0, 1, 5, 2],
       [6, 3, 0, 6, 4]])

In [155]: JTildeIDX
Out[155]: array([0, 0, 4, 6, 5])

In [156]: mask = np.arange(J.shape[0])[:,None] >= JTildeIDX

In [157]: mask
Out[157]: 
array([[ True,  True, False, False, False],
       [ True,  True, False, False, False],
       [ True,  True, False, False, False],
       [ True,  True, False, False, False],
       [ True,  True,  True, False, False],
       [ True,  True,  True, False,  True],
       [ True,  True,  True,  True,  True]], dtype=bool)

In [158]: out = (J*mask).sum(0)

In [159]: out
Out[159]: array([34, 34,  7,  6,  6])

Upvotes: 2

Related Questions