配列差分計算のfor文をスライスに置き換え

# はじめに
最近，数値計算の勉強を始めました．試しにキャビティ流れを計算してみたのですが，巷で言われるようにPythonでfor文を多用すると処理時間がとても長くなりました．備忘録を兼ねて，for文を極力使わないで書いたプログラムと実行速度を比較します．

# 参考
キャビティ流れの計算は名古屋工業大学の講義資料を参考にしました．
[H21年度・４日目 （ ２００９年８月２９日（土） ）
流体数値計算法](http://fluid.web.nitech.ac.jp/Gotoh_Home_page/Edu/Public_course/Text/CAE_4th.pdf)

for文を使わないで差分を計算する方法は次を参考にしました．
[Pythonで数値計算 2. Numpy array (配列) のコツ．](https://japanichaos.appspot.com/NumpyArray.html)

# プログラム
まずは教科書通りにfor文を多用したプログラムを示します．
こちらのプログラムをBaselineとします．

```python:HSMAC_Baseline
import numpy as np

def update(km, xn, yn, dx, dy, dt, eps, rho, mu, p, u, v):
    #速度と圧力の修正
    for k in range(km):
        flag = True #収束判定用
        for i in range(1, xn):
            for j in range(1, yn):
                #p
                dp = eps * ((u[i+1,j] - u[i,j]) / dx + (v[i,j+1] - v[i,j]) / dy)
                p[i,j] = p[i,j] + dp

                #u
                u[i+1,j] = u[i+1,j] + dt * dp / dx / rho
                u[i,j] = u[i,j] - dt * dp / dx / rho

                #v
                v[i,j+1] = v[i,j+1] + dt * dp / dy / rho
                v[i,j] = v[i,j] - dt * dp / dy / rho
                
                divv = (u[i+1,j] - u[i,j]) / dx + (v[i,j+1] - v[i,j]) / dy
                if divv >= 1.0e-6 : flag = False
        if flag : break

def velocity(xn, yn, u, v, dx, dy, dt, rho, mu, p, uwall):
    #右と左の境界条件
    for j in range(yn+1):
        u[1,j] = 0.0
        u[0,j] = u[2,j]
        v[0,j] = -v[1,j]

        u[xn,j] = 0
        u[xn+1,j] = u[xn-1,j]
        v[xn,j] = -v[xn-1,j]

    #上と下の境界条件
    for i in range(xn+1):
        v[i,1] = 0.0
        v[i,0] = v[i,2]
        u[i,0] = -u[i,1]

        v[i,yn] = 0.0
        v[i,yn+1] = v[i,yn-1]
        u[i,yn] = 2.0*uwall - u[i,yn-1]

    u[xn+1,0] = -u[xn,0]
    u[xn+1,yn] = -u[xn,yn]

    #u
    for i in range(2, xn):
        for j in range(1,yn):
            vhat = (v[i,j] + v[i,j+1] + v[i-1,j+1] + v[i-1,j]) / 4.0
            u1 = -u[i,j] * (u[i+1,j] - u[i-1,j]) / (2.0*dx) - vhat * (u[i,j+1] - u[i,j-1]) / (2.0*dy)
            u2 = (u[i+1,j] - 2.0*u[i,j] + u[i-1,j]) / dx**2 + (u[i,j+1] - 2.0*u[i,j] + u[i,j-1]) / dy**2
            u[i,j] = u[i,j] + dt * (u1 + mu/rho * u2)

    #v
    for i in range(1, xn):
        for j in range(2, yn):
            uhat = (u[i,j] + u[i+1,j] + u[i+1,j-1] + u[i,j-1]) / 4.0
            v1 = -uhat * (v[i+1,j] - v[i-1,j]) / (2.0*dx) - v[i,j] * (v[i,j+1] - v[i,j-1]) / (2.0*dy)
            v2 = (v[i+1,j] - 2.0*v[i,j] + v[i-1,j]) / dx**2 + (v[i,j+1] - 2.0*v[i,j] + v[i,j-1]) / dy**2
            v[i,j] = v[i,j] + dt * (v1 + mu/rho * v2)
            
def post_process(dx, dy, xn, yn, p, u, v):
    x = np.arange(start=dx, step=dx, stop=xn*dx) - 0.5*dx
    y = np.arange(start=dy, step=dy, stop=yn*dy) - 0.5*dy
    y, x = np.meshgrid(x, y)

    u = (u[1:-2,1:-1] + u[2:-1,1:-1]) / 2.0
    v = (v[1:-1,1:-2] + v[1:-1,2:-1]) / 2.0
    
    return x, y, p, u, v

def mac(uwall):
    xn = 21
    yn = 21
    L = 10.0e-3 #幅[m]
    D = 10.0e-3 #深さ[m]
    u = np.zeros((xn+2, yn+1))
    v = np.zeros((xn+1, yn+2))
    p = np.zeros((xn+1, yn+1))
    dx = L / (xn-1) #x方向刻み幅
    dy = D / (yn-1) #y方向刻み幅
    dt = 1.0e-3 #時間刻み
    mu = 0.01 #粘性係数[m*s]
    rho = 1000.0 #密度[kg/m**3]
    eps = -0.4 * rho / dt / (1.0 / dx**2 + 1.0 / dy**2)
    lm = 1000
    km = 2000

    #time step
    for l in range(lm):
        velocity(xn, yn, u, v, dx, dy, dt, rho, mu, p, uwall)
        update(km, xn, yn, dx, dy, dt, eps, rho, mu, p, u, v)

    return post_process(dx, dy, xn, yn, p, u, v)
```

次にnumpyのスライス機能を使って，for文を極力減らしたプログラムを示します．共通部分は省きました．こちらのほうが可読性が高いように思います．
以後，このプログラムをSliceと表記します．

```python:HSMAC_Slice
def update(km, xn, yn, dx, dy, dt, eps, rho, mu, p, u, v):
    #速度と圧力の修正
    for k in range(km):      
        #p
        dp = eps * ((u[1+1:xn+1,1:yn] - u[1:xn,1:yn]) / dx + (v[1:xn,1+1:yn+1] - v[1:xn,1:yn]) / dy)
        p[1:xn,1:yn] = p[1:xn,1:yn] + dp
        
        #u
        u[1+1:xn+1,1:yn] = u[1+1:xn+1,1:yn] + dt * dp / dx / rho
        u[1:xn,1:yn] = u[1:xn,1:yn] - dt * dp / dx / rho
        
        #v
        v[1:xn,1+1:yn+1] = v[1:xn,1+1:yn+1] + dt * dp / dy / rho
        v[1:xn,1:yn] = v[1:xn,1:yn] - dt * dp / dy / rho
        
        divv = (u[1+1:xn+1,1:yn] - u[1:xn,1:yn]) / dx + (v[1:xn,1+1:yn+1] - v[1:xn,1:yn]) / dy
        if np.amax(divv) <= 1.0e-6 : break

def velocity(xn, yn, u, v, dx, dy, dt, rho, mu, p, uwall):
    #右と左の境界条件
    u[1,:] = 0.0
    u[0,:] = u[2,:]
    v[0,:] = -v[1,:]
    
    u[xn,:] = 0.0
    u[xn+1,:] = u[xn-1,:]
    v[xn,:] = -v[xn-1,:]

    #上と下の境界条件
    v[:,1] = 0.0
    v[:,0] = v[:,2]
    u[:,0] = -u[:,1]
    
    v[:,yn] = 0.0
    v[:,yn+1] = v[:,yn-1]
    u[:,yn] = 2.0*uwall - u[:,yn-1]

    u[xn+1,0] = -u[xn,0]
    u[xn+1,yn] = -u[xn,yn]

    #u
    vhat = (v[2:xn,1:yn] + v[2:xn,1+1:yn+1] + v[2-1:xn-1,1+1:yn+1] + v[2-1:xn-1,1:yn]) / 4.0
    u1 = -u[2:xn,1:yn] * (u[2+1:xn+1,1:yn] - u[2-1:xn-1,1:yn]) / (2.0*dx) - vhat * (u[2:xn,1+1:yn+1] - u[2:xn,1-1:yn-1]) / (2.0*dy)
    u2 = (u[2+1:xn+1,1:yn] - 2.0*u[2:xn,1:yn] + u[2-1:xn-1,1:yn]) / dx**2 + (u[2:xn,1+1:yn+1] - 2.0*u[2:xn,1:yn] + u[2:xn,1-1:yn-1]) / dy**2
    u[2:xn,1:yn] = u[2:xn,1:yn] + dt * (u1 + mu/rho * u2)

    #v     
    uhat = (u[1:xn,2:yn] + u[1+1:xn+1,2:yn] + u[1+1:xn+1,2-1:yn-1] + u[1:xn,2-1:yn-1]) / 4.0
    v1 = -uhat * (v[1+1:xn+1,2:yn] - v[1-1:xn-1,2:yn]) / (2.0*dx) - v[1:xn,2:yn] * (v[1:xn,2+1:yn+1] - v[1:xn,2-1:yn-1]) / (2.0*dy)
    v2 = (v[1+1:xn+1,2:yn] - 2.0*v[1:xn,2:yn] + v[1-1:xn-1,2:yn]) / dx**2 + (v[1:xn,2+1:yn+1] - 2.0*v[1:xn,2:yn] + v[1:xn,2-1:yn-1]) / dy**2
    v[1:xn,2:yn] = v[1:xn,2:yn] + dt * (v1 + mu/rho * v2)
```

# 結果 キャビティ流れ
まずは，キャビティ流れが正しく計算されているか確認します．

```python:result
import matplotlib.pyplot as plt

%time x, y, p, u, v = mac(10.0e-3)
fig, axes = plt.subplots()
qu = axes.quiver(x*1000, y*1000, u, v)
axes.set_title("10 mm/s")
axes.set_xlabel("x mm")
axes.set_ylabel("y mm")
axes.set_aspect("equal")
fig.show()
```

![ダウンロード.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/455608/b2fc2b5c-20d3-85b5-1b2b-628962a69b81.png)



# 結果 処理時間
Baseline，SliceとそれぞれにNumbaを用いたものの処理時間を計測しました．
Baseline以外は3回の平均値を用いています．

|  | Baseline | Baseline + Numba | Slice | Slice + Numba |
|:-:|:-:|:-:|:-:|:-:|
| Time[s] | 603 | 2.50 | 34.6 | 7.17 |
| Gain[-] | 1.00 | 241 | 17.4 | 84.1 |


#まとめ
Numbaすごい
numpyのスライスを使ってfor文を減らすよりもNumbaの影響のほうが大きいですね．

## 追記
shiracamusさん，ご指摘ありがとうございます．
Python標準機能だけでは実現できないので，タグにnumpyを追加しました．