CODEWITHSUNDEEP
pythonnumpy
slaw
slaw

Reputation: 6869

Select NumPy Values Around Index

I have two NumPy arrays:

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)

Let's say that where ever there is an np.inf in x, that index position is called i. For each i, I want to set the values of y[i-m+1:i+m] = np.inf. So, after the replacement, y should look like:

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])

This should also work when the value of m is either increased or decreased.

Upvotes: 5

Views: 256

Answers (4)

Sholto Armstrong
Sholto Armstrong

Reputation: 273

The easiest solution I can think of is to use the np.convolve function to dilate a mask. This can be done as follows:

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

Upvotes: 3

Pynchia
Pynchia

Reputation: 11580

Using 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)

produces

[ 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.]

Upvotes: 0

yatu
yatu

Reputation: 88226

Here's one approach defineing an integer mask to index y and using 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.])

Upvotes: 2

cmxu
cmxu

Reputation: 972

So you basically have the solution already,

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

The only thing I think you have wrong is that I believe it should be from i-m+1 to i+m. Hope that helps!

Upvotes: 1

Related Questions