Эффективный numpy.cumsum и numpy.digitize

Question

Учитывая матрицу значений, которые представляют вероятности, я пытаюсь написать эффективный процесс, который возвращает корзину, которой принадлежит значение. Например:

sample = 0.5
x = np.array([0.1]*10)
np.digitize( sample, np.cumsum(x))-1
#returns 5

это результат, который я ищу. В соответствии с timeit за x для массивов с несколькими элементами более эффективно сделать это следующим образом:

cdf = 0
for key,val in enumerate(x):
    cdf += val
    if sample<=cdf:
        print key
        break

пока для большего x Массивы решение NumPy быстрее. Вопрос:

1. Есть ли способ еще больше ускорить его, например, функцию, которая объединяет шаги?
2. Можем ли мы векторизовать процесс для случая, когда sample список, каждый элемент которого связан со своим x массив (x тогда будет 2-й)?

В приложении x содержит предельные вероятности; это способ, которым я должен уменьшить результаты np.digitize

Answer 1

Вы могли бы использовать некоторые broadcasting магия там -

(x.cumsum(1) > sample[:,None]).argmax(1)-1

Шаги вовлечены:

I. Произведите слияние по каждому ряду.

II. Используйте широковещательное сравнение для каждой строки совокупности с каждым значением выборки и ищите, чтобы первое вхождение выборки было меньше значений совокупности, сигнализируя, что элемент перед этим в x это индекс, который мы ищем.

Пошаговый прогон -

In [64]: x
Out[64]: 
array([[ 0.1 ,  0.1 ,  0.1 ,  0.1 ,  0.1 ,  0.1 ,  0.1 ],
       [ 0.8 ,  0.96,  0.88,  0.36,  0.5 ,  0.68,  0.71],
       [ 0.37,  0.56,  0.5 ,  0.01,  0.77,  0.88,  0.36],
       [ 0.62,  0.08,  0.37,  0.93,  0.65,  0.4 ,  0.79]])

In [65]: sample # one elem per row of x
Out[65]: array([ 0.5,  2.2,  1.9,  2.2])

In [78]: x.cumsum(1)
Out[78]: 
array([[ 0.1 ,  0.2 ,  0.3 ,  0.4 ,  0.5 ,  0.6 ,  0.7 ],
       [ 0.8 ,  1.76,  2.64,  2.99,  3.49,  4.18,  4.89],
       [ 0.37,  0.93,  1.43,  1.45,  2.22,  3.1 ,  3.47],
       [ 0.62,  0.69,  1.06,  1.99,  2.64,  3.04,  3.83]])

In [79]: x.cumsum(1) > sample[:,None]
Out[79]: 
array([[False, False, False, False, False,  True,  True],
       [False, False,  True,  True,  True,  True,  True],
       [False, False, False, False,  True,  True,  True],
       [False, False, False, False,  True,  True,  True]], dtype=bool)

In [80]: (x.cumsum(1) > sample[:,None]).argmax(1)-1
Out[80]: array([4, 1, 3, 3])

# A loopy solution to verify results against
In [81]: [np.digitize( sample[i], np.cumsum(x[i]))-1 for i in range(x.shape[0])]
Out[81]: [4, 1, 3, 3]

---

Граничные случаи:

Предлагаемое решение автоматически обрабатывает случаи, когда sample значения меньше наименьшего из совокупных суммированных значений -

In [113]: sample[0] = 0.08  # editing first sample to be lesser than 0.1

In [114]: [np.digitize( sample[i], np.cumsum(x[i]))-1 for i in range(x.shape[0])]
Out[114]: [-1, 1, 3, 3]

In [115]: (x.cumsum(1) > sample[:,None]).argmax(1)-1
Out[115]: array([-1,  1,  3,  3])

Для случаев, когда sample значение больше, чем наибольшее из совокупных суммированных значений, нам нужно сделать один дополнительный шаг -

In [116]: sample[0] = 0.8  # editing first sample to be greater than 0.7

In [121]: mask = (x.cumsum(1) > sample[:,None])

In [122]: idx = mask.argmax(1)-1

In [123]: np.where(mask.any(1),idx,x.shape[1]-1)
Out[123]: array([6, 1, 3, 3])

In [124]: [np.digitize( sample[i], np.cumsum(x[i]))-1 for i in range(x.shape[0])]
Out[124]: [6, 1, 3, 3]