Как посчитать ненулевые значения, используя binned_statistic

Question

Мне нужно эффективно обрабатывать очень большие одномерные массивы, извлекая некоторую статистику для каждого бина, и я нашел очень полезной функцию binned_statistic из scipy.stats, так как она включает аргумент "статистика", который работает довольно эффективно.

Я хотел бы выполнить функцию "подсчета", но без учета нулевых значений.

Я работаю параллельно со скользящими окнами (функция скручивания панд) над одними и теми же массивами, и это прекрасно работает для замены нулей на NaN, но это поведение не распространяется на мой случай.

Это игрушечный пример того, что я делаю:

import numpy as np
import pandas as pd
from scipy.stats import binned_statistic

# As example with sliding windows, this returns just the length of each window:
a = np.array([1., 0., 0., 1.])
pd.Series(a).rolling(2).count() # Returns [1.,2.,2.,2.]

# You can make the count to do it only if not zero:
nonzero_a = a.copy()
nonzero_a[nonzero_a==0.0]='nan'
pd.Series(nonzero_a).rolling(2).count()   # Returns [1.,1.,0.,1.]

# However, with binned_statistic I am not able to do anything similar:
binned_statistic(range(4), a, bins=2, statistic='count')[0] 
binned_statistic(range(4), nonzero_a, bins=2, statistic='count')[0]
binned_statistic(range(4), np.array([1., False, None, 1.], bins=2, statistic='count')[0]

Все предыдущие прогоны дают одинаковый результат: [2., 2.], но я ожидаю [1., 1.].

Единственный найденный вариант - передать пользовательскую функцию, но она работает значительно хуже, чем реализованные функции в реальных случаях.

binned_statistic(range(4), a, bins=2, statistic=np.count_nonzero)

Answer 1

Я нашел и простой способ скопировать ненулевой счетчик, преобразовав массив в 0-1 и применяя сумму:

 # Transform all non-zero to 1s
 a = np.array([1., 0., 0., 2.])
 nonzero_a = a.copy()
 nonzero_a[nonzero_a>0.0]=1.0     # nonzero_a = [1., 0., 0., 1.]

 binned_statistic(np.arange(len(nonzero_a)), nonzero_a, bins=bins, statistic='sum')[0]   # Returns [1.0, 1.0]