for文の繰り返しをdef関数を用いて、なるべく避けたい。

以前の質問で境界条件を考慮したランダムウォークの続きなのですが、以下のメイン計算を100000万回繰り返したいとなると、

python
1import random
2import numpy as np
3import itertools
4n = 3
5itr = 5
6step = [i for i in range(0,itr+1)]
7r_list = [i for i in range(0,n)]
8walk = np.zeros([n,n,n],dtype=np.uint8)
9walk[0,0,0] = 1
10count =0
11time=[]
12for k in range(0,10000):
13    for t in step:
14        if t == 0:
15            pass
16        else:
17            number = random.randint(1,6)
18            next_walk = np.zeros([n,n,n],dtype=np.uint8)
19            for i in itertools.product(r_list,r_list,r_list):
20                x = i[0]
21                y = i[1]
22                z = i[2]
23        #Boundary condition
24                x1 = (x-1 + n) % n
25                x2 = (x+1) % n
26                y1 = (y-1 + n) % n
27                y2 = (y+1) % n
28                z1 = (z-1 + n) % n
29                z2 = (z+1) % n
30                if  stage[i]== 0:
31                    continue
32                else:
33                    if walk[i]==0:
34                        continue
35                    else:
36                        if number == 1:
37                            if stage[x2,y,z]==1:
38                                next_walk[x2,y,z]=walk[i]
39                            else:
40                                next_walk[x,y,z]=walk[i]
41                        elif number == 2:
42                            if stage[x,y2,z]==1:
43                                next_walk[x,y2,z]=walk[i]
44                            else:
45                                next_walk[x,y,z]=walk[i]
46                        elif number == 3:
47                            if stage[x,y,z2]==1:
48                                next_walk[x,y,z2]=walk[i]
49                            else:
50                                next_walk[x,y,z]=walk[i]
51                        elif number == 4:
52                            if stage[x1,y,z]==1:
53                                next_walk[x1,y,z]=walk[i]
54                            else:
55                                next_walk[x,y,z]=walk[i]
56                        elif number == 5:
57                            if stage[x,y1,z]==1:
58                                next_walk[x,y1,z]=walk[i]
59                            else:
60                                next_walk[x,y,z]=walk[i]
61                        elif number == 6:
62                            if stage[x,y,z1]==1:
63                                next_walk[x,y,z1]=walk[i]
64                            else:
65                                next_walk[x,y,z]=walk[i]
66            walk = np.copy(next_walk)

となってしまい、for文の分岐が多くなってしまい、処理がものすごく遅くなってしまい、これよりもう少し早くしたいです。とりあえず、メイン計算と繰り返し回数を分けてコードを書きたいのですが、

def move():
    for i in itertools.product(r_list,r_list,r_list):
        x = i[0]
        y = i[1]
        z = i[2]
        #Boundary condition
        x1 = (x-1 + n) % n
        x2 = (x+1) % n
        y1 = (y-1 + n) % n
        y2 = (y+1) % n
        z1 = (z-1 + n) % n
        z2 = (z+1) % n
        if stage[i] == 0:
            continue
        else:
            if walk[i]==1:
                if number == 1:
                    if stage[x2,y,z]==1:
                        next_walk[x2,y,z]=walk[i]
                    else:
                        next_walk[x,y,z]=walk[i]
                elif number == 2:
                    if stage[x,y2,z]==1:
                        next_walk[x,y2,z]=walk[i]
                    else:
                        next_walk[x,y,z]=walk[i]
                elif number == 3:
                    if stage[x,y,z2]==1:
                        next_walk[x,y,z2]=walk[i]
                    else:
                        next_walk[x,y,z]=walk[i]
                elif number == 4:
                    if stage[x1,y,z]==1:
                        next_walk[x1,y,z]=walk[i]
                    else:
                        next_walk[x,y,z]=walk[i]
                elif number == 5:
                    if stage[x,y1,z]==1:
                        next_walk[x,y1,z]=walk[i]
                    else:
                        next_walk[x,y,z]=walk[i]
                elif number == 6:
                    if stage[x,y,z1]==1:
                        next_walk[x,y,z1]=walk[i]
                    else:
                        next_walk[x,y,z]=walk[i]
                else:
                    continue
    walk = np.copy(next_walk)
    return(walk)

def experiment(t,N):
    for i in range(0,N):
        for j in range(0,t):
             move()
のようになるべくfor文の分岐を減らして書きたいのですが、それと同時に愚直に書いた時よりも処理は軽くなるのでしょうか？

エラー

TypeError                                 Traceback (most recent call last)
<ipython-input-21-96add303636d> in <module>()
     56     walk = np.copy(next_walk)
     57     return(walk)
---> 58 print(move())

<ipython-input-21-96add303636d> in move()
     18         z1 = (z-1 + n) % n
     19         z2 = (z+1) % n
---> 20         if stage[i] == 0:
     21             continue
     22         else:

TypeError: list indices must be integers or slices, not tuple

正直、def関数がわかっているようで、全くわかっていない。。。