How can I use Numpy functions instead of For-loops to access and manipulate adjacent array entries?

Question

I'm trying to simulate a stochastic system on a 2D grid. Basically, everytime a random clock rings, the color of a pixel changes to some other color, possibly changing the color of some random nearest neighbor too. The following code gives the desired result but not the desired speed (at least for a 1000x1000 grid):

import numpy as np
import random as rd
import matplotlib.pyplot as plt
from matplotlib import colors
from matplotlib.animation import FuncAnimation

#define colormap
color= ['aquamarine','cyan','deepskyblue','dodgerblue',
        'fuchsia','mediumorchid','blueviolet','darkslateblue']
cmap=colors.ListedColormap(color)

#define transition rates: 0->1 at rate[0], 1->2 at rate[1], 2->3 at rate [2],
#3->0 at rate[3], 4->5 at rate[4], 5->6 at rate[5], 6->7 at rate[6], 7->4 at rate[7]
rate = np.array([0.21, 1, 1, 1, 0.25, 1, 1, 1])

#set time interval                    
T=100
time =np.arange(T)

#define initial configuration
N=1000
config = np.zeros((N,N), dtype=int)
for i in range(450,550):
    for j in range(450,550):
        config[i][j]=4

#change configuration according to some rule
def step(config, n):
        r = rate[config]                #get only those entries
        k = np.random.poisson(lam=r)    #that are randomly 
        x = np.transpose(np.nonzero(k)) #activated and
        np.random.shuffle(x)            #traverse them in
        for ind in x:                   #random order
            i=ind[0]
            j=ind[1]
            if config[i,j]==3:          #if in state 3
                vec = rd.choice([[-1,0],[1,0],[0,1],[0,-1]])#choose
                x =(i+vec[0])%n         #a random neighbour
                y =(j+vec[1])%n
                config[x,y] = 0         #and send them
                config[i,j] = 0         #to state 0
            elif config[i,j]==7:        #same here, but send them to state 4
                vec = rd.choice([[-1,0],[1,0],[0,1],[0,-1]])
                x =(i+vec[0])%n
                y =(j+vec[1])%n
                config[x,y] = 4
                config[i,j] = 4
            else:
                config[i,j] = config[i,j]+1#othwise move to next state

#set up plot        
fig, ax = plt.subplots(figsize=(5,5))
im=ax.imshow(config[::-1],cmap=cmap, vmin=0, vmax=cmap.N)
time_text=ax.text(2*N/3,N-1,"time="+str(0))

def update(t):
    step(config, N)
    im.set_data(config[::-1])
    time_text.set_text("time="+str((t+1)/T))
    return im, time_text,

ani = FuncAnimation(fig, update, frames=time,
                    interval=5, blit=True, repeat=False)

plt.show()

(For reference: A 100x100 grid takes roughly 7 seconds, a 500x500 grid 45 s) I suspect that the step() function is rather slow due to looping over a possibly large array. Question 1: (How) can Numpy functions be uitilized to avoid the For-loop? Question 2: Is there any other (easy) option to increase the simulation speed?

At first I thought that np.where() could be employed to circumvent the For-loop and its inner If-statement but I don't kow how to access and modify (random) adjacent elements. Besides that I've tried to use Numba with zero success (but lots of warnings and errors). Is it worth trying to make it work?

Thank you for your help!

Answer 1

Main idea is that most pixels are not 3 or 7, and so their changes can be vectorized. Then we do the neighbor changes from 3's and 7's one by one. For competing influences, we keep track of it in spread_changes and then randomly pick one as the one that "applied last." The neighbor influence is definitely the slow step here, but I haven't figured out a good way to vectorize it.

import random as rd
import matplotlib.pyplot as plt
from matplotlib import colors
from matplotlib.animation import FuncAnimation

#define colormap
color= ['aquamarine','cyan','deepskyblue','dodgerblue',
        'fuchsia','mediumorchid','blueviolet','darkslateblue']
cmap=colors.ListedColormap(color)

#define transition rates: 0->1 at rate[0], 1->2 at rate[1], 2->3 at rate [2],
#3->0 at rate[3], 4->5 at rate[4], 5->6 at rate[5], 6->7 at rate[6], 7->4 at rate[7]
rate = np.array([0.21, 1, 1, 1, 0.25, 1, 1, 1])

#set time interval                    
T=100
time =np.arange(T)

#define initial configuration
N=1000
config = np.zeros((N,N), dtype=int)
for i in range(450,550):
    for j in range(450,550):
        config[i][j]=4


def step(n):
    global config
    k = np.random.poisson(lam = rate[config])
    change_mask = np.array(k > 0, dtype = 'int')
    three_mask = np.array(config == 3, dtype = 'int')
    seven_mask = np.array(config == 7, dtype = 'int')
    
    plus_one = (config+1)*change_mask*(1-three_mask)*(1-seven_mask)
    no_change = config*(1-change_mask)
    # three_to_zero = three_mask*change_mask*0  #Don't need this, it's always all zeros
    seven_to_four = seven_mask*change_mask*4
    
    config = plus_one + no_change + seven_to_four
    
    
    spread_changes = defaultdict(list)
    
    for y,x in zip(*np.nonzero(seven_mask*change_mask)):
        dy, dx = random.choice([[1, 0], [-1, 0], [0, 1], [0, -1]])
    
        spread_changes[(y+dy)%N, (x+dx)%N].append(4)
    
    for y,x in zip(*np.nonzero(three_mask*change_mask)):
        dy, dx = random.choice([[1, 0], [-1, 0], [0, 1], [0, -1]])
    
        spread_changes[(y+dy)%N, (x+dx)%N].append(0)

    for i,j in spread_changes:
  
        choice = random.randint(0, len(spread_changes[i,j]))

        if choice < len(spread_changes[i,j]):

            config[i,j] = spread_changes[i,j][choice]


#set up plot        
fig, ax = plt.subplots(figsize=(5,5))
im=ax.imshow(config[::-1],cmap=cmap, vmin=0, vmax=cmap.N)
time_text=ax.text(2*N/3,N-1,"time="+str(0))

def update(t):
    global config
    step(N)
    im.set_data(config[::-1])
    time_text.set_text("time="+str((t+1)/T))
    return im, time_text,

ani = FuncAnimation(fig, update, frames=time,
                    interval=5, blit=True, repeat=False)

plt.show()