Выберите NumPy значений вокруг индекса

Question

У меня есть два NumPy массива:

import numpy as np

m = 3
x = np.array([1, 0, 0, np.inf, 0, 0, 1, 1, 2, np.inf, np.inf, 2, 3, 4, 5, 6, 7, 8, 9, 10])
y = np.arange(x.shape[0]-m+1)

Допустим, что когда бы ни было np.inf в x, эта позиция индекса называется i. Для каждого i я хочу установить значения y[i-m+1:i+m] = np.inf. Таким образом, после замены y должно выглядеть следующим образом:

array([0, np.inf, np.inf, np.inf, np.inf, np.inf, 6, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, 13, 14, 15, 16, 17])

Это также должно работать, когда значение m увеличивается или уменьшается.

Answer 1

Самым простым решением, которое я могу придумать, является использование функции np.convolve для расширения маски. Это можно сделать следующим образом:

mask = np.convolve(x==np.inf, [True]*(m*2-1), mode='same')
y[mask[:-m+1]] = np.inf

Answer 2

Вот один подход, определяющий целочисленную маску для индексации y и использующий broadcasting:

m = 3
x = np.array([1, 0, 0, np.inf, 0, 0, 1, 1, 2, np.inf, np.inf, 2, 3, 4, 5, 6, 7, 8, 9, 10])
y = np.arange(x.shape[0]-m+1).astype(float)

---

i = np.flatnonzero(x == np.inf)
y[(i + np.arange(-m+1,m)[:,None]).ravel('F')] = np.inf

print(y)
array([ 0., inf, inf, inf, inf, inf,  6., inf, inf, inf, inf, inf, inf,
       13., 14., 15., 16., 17.])

Answer 3

Так что у вас уже есть решение,

import numpy as np
m = 3
x = np.array([1, 0, 0, np.inf, 0, 0, 1, 1, 2, np.inf, np.inf, 2, 3, 4, 5, 6, 7, 8, 9, 10])
y = np.arange(x.shape[0]-m+1)

idx = np.where(x == np.inf)[0] # grabs all i's
for i in idx: # applies the operation you described at each i
    y[(i-m+1):(i+m)] = np.inf # slightly changed indices

Единственное, что я думаю, что вы ошибаетесь, это то, что я считаю, что оно должно быть от i-m+1 до i+m. Надеюсь, это поможет!

Answer 4

Использование np.meshgrid

m = 3
x = np.array([1, 0, 0, np.inf, 0, 0, 1, 1, 2, np.inf, np.inf, 2, 3, 4, 5, 6, 7, 8, 9, 10]) 
y = np.arange(x.shape[0]-m+1, dtype=float)

inf_ndx = np.flatnonzero(np.isinf(x))  # where the infs are
a, b = np.meshgrid(inf_ndx, np.arange(-m+1,m))  # expand the indexes with the ranges
y[a+b] = np.inf  # set all such locations to inf

print(inf_ndx)    
print(a)
print(b)
print(y)

производит

[ 3  9 10]
[[ 3  9 10]
 [ 3  9 10]
 [ 3  9 10]
 [ 3  9 10]
 [ 3  9 10]]
[[-2 -2 -2]
 [-1 -1 -1]
 [ 0  0  0]
 [ 1  1  1]
 [ 2  2  2]]
[ 0. inf inf inf inf inf  6. inf inf inf inf inf inf 13. 14. 15. 16. 17.]