Modify element and neighbor in numpy 2d array

I am trying find a more efficient way to modify every element of an array, and its lowest neighbor, by some amount. In the example code below, I modify them by their difference, but the change() function could be anything.

After searching, scipy.ndimage.generic_filter() seemed to be the ideal method to use, since it allows for easy comparison between elements and their neighbors. After getting the offset from ndimage.generic_filter(), I am feeding that to numpy.roll() to modify each element's chosen neighbor.

The problem is that, with very large arrays and multiple iterations, the inefficiency of looping through np.roll() ndimage.generic_filter() eats into performance. With a 10000x10000 array, my execution time of the code below is 5m42s. Is there a more efficient way of doing this using scipy or numpy?

import numpy as np
from scipy import ndimage

dic = {0: (-1, -1), 1: (-1, 0), 2: (-1, 1),
       3: (0, -1),  4: (0, 0),  5: (0, 1),
       6: (1, -1),  7: (1, 0),  8: (1, 1)}

def change(a):
    return (a[4] - min(a)) / 2

def movement(a):
    return np.argmin(a)

P = np.random.uniform(10, size=(5,5))

# determine element change and direction
chng = np.zeros((5, 5)) 
ndimage.generic_filter(P, change, size=3, output=chng)
move = np.random.randint(10, size=(5, 5))
ndimage.generic_filter(P, movement, size=3, output=move)
P -= chng

# determine neighbor change
chng2 = np.zeros((5, 5))
for j in range(9):
    if j == 4:
        continue
    p = np.where(move == j, chng, 0)
    p = np.roll(p, dic[j], axis=(0, 1))
    chng2 += p
P += chng2

EDIT: Below is a more efficient solution. Many thanks @Paul Panzer.

import numpy as np

P = np.random.uniform(10, size=(1000, 1000))

# determine element change and direction
PP = np.bmat([[P[:, -1:], P, P[:, :1]]])
PP = np.bmat([[PP[-1:]], [PP], [PP[:1]]])
PPP = np.lib.stride_tricks.as_strided(PP, (1000, 1000, 3, 3), 2 * PP.strides)
am1 = np.argmin(PPP, axis=3)
i, j, k = np.ogrid[(*map(slice, PPP.shape[:-1]),)]
am0 = np.argmin(PPP[i, j, k, am1], axis=2)
i, j = np.ogrid[(*map(slice, PPP.shape[:-2]),)]
am1 = am1[i, j, am0]
mn = PPP[i, j, am0, am1]
change = (P - mn) / 2
P -= change

# determine neighbor change
am1 -= 1
am0 -= 1
i, j = np.ogrid[(*map(slice, P.shape),)]
np.add.at(P, ((am0 + i) % P.shape[0], (am1 + j) % P.shape[1]), change)

2 answers

  • answered 2018-02-13 03:01 David Zarebski

    This might be what you are looking for https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.convolve2d.html

    in2 (Cf documentation) would be a matrix corresponding to what you wrote as a dict

    dic = {0: (-1, -1), 1: (-1, 0), 2: (-1, 1),
       3: (0, -1),  4: (0, 0),  5: (0, 1),
       6: (1, -1),  7: (1, 0),  8: (1, 1)}
    

    Hope it helps

  • answered 2018-02-13 04:49 Paul Panzer

    You could use np.add.at. The following snippet picks up after your P -= chng:

    >>> P_pp = P.copy()
    >>> dic_i, dic_j = np.indices((3, 3)).reshape(2, 9) - 1
    >>> i, j = np.ogrid[(*map(slice, P_pp.shape),)]
    >>> np.add.at(P_pp, ((dic_i[move] + i) % P_pp.shape[0], (dic_j[move] + j) % P_pp.shape[1]), chng)
    

    Since we worked on a copy of P we can now execute the rest of your code and then:

    # Tada!
    >>> np.allclose(P_pp, P)
    True
    

    Update: Here is a method to calculate local argmin without using ndimage. One potential advantage is that we get the corresponding minima cheaply once we have the argminima. Note that the argmin is already 2D first component is in am0, second in am1. Each ranges between 0 and 2, so the center is at 1,1, The minima are in mn

    >>> P = np.random.uniform(10, size=(5,5))
    >>> PP = np.bmat([[P[:,-1:], P, P[:, :1]]])
    >>> PP = np.bmat([[PP[-1:]], [PP], [PP[:1]]])
    >>> PPP = np.lib.stride_tricks.as_strided(PP, (5, 5, 3, 3), 2 * PP.strides)
    >>> am1 = np.argmin(PPP, axis=3)
    >>> i, j, k = np.ogrid[(*map(slice, PPP.shape[:-1]),)]
    >>> am0 = np.argmin(PPP[i, j, k, am1], axis=2)
    >>> i, j = np.ogrid[(*map(slice, PPP.shape[:-2]),)]
    >>> am1 = am1[i, j, am0]
    >>> mn = PPP[i, j, am0, am1]