趣味で計算流砂水理

趣味で計算流砂水理 Computational Sediment Hydraulics for Fun Learning

数値計算とか河川工学とかプログラミングのことを書いています

MENU

pythonの高速化の基礎：二次元配列の処理を例に

R

スポンサーリンク

結論:何も考えずnumbaを使え。

モチベーション

pythonの具体的な高速化のテクニックを備忘録的に整理しておこうと思ったので。

テストケース

標高zの二次元配列1000x1000から最急勾配を求める以下の式を対象としました。

\begin{align} {\rm slope} = \arctan{ \left( \sqrt{ \left( \frac{\partial z}{\partial x} \right)^2 + \left( \frac{\partial z}{\partial y} \right)^2}\right) } \end{align}

離散化は以下記事を参照。

computational-sediment-hyd.hatenablog.jp

計算結果一覧

PCのスペックは以下のとおり

computational-sediment-hyd.hatenablog.jp

基本ケース

基本ケースのコードは下のとおり。

import numpy as np
import xarray as xr
from numba import jit

z = np.random.rand(1000,1000)

def Slope(z):
    dx, dy = float(1), float(1)
    ny, nx = z.shape
    
    Slope = np.full_like(z, float(-9999))
    for j in range(1, ny-1):
        for i in range(1, nx-1):
            dzdx = 0.25*(
                0.5*(z[j+1,i+1] - z[j+1,i-1])/dx
              + 1.0*(z[j  ,i+1] - z[j  ,i-1])/dx
              + 0.5*(z[j-1,i+1] - z[j-1,i-1])/dx)
                
            dzdy = 0.25*(
                0.5*(z[j+1,i+1] - z[j-1,i+1])/dy
              + 1.0*(z[j+1,i  ] - z[j-1,i  ])/dy
              + 0.5*(z[j+1,i-1] - z[j-1,i-1])/dy)
            
            Slope[j,i] = np.arctan( np.sqrt(dzdx**2 + dzdy**2) )
            
    return Slope

%%timeit
d = Slope(z)

5.54 s ± 75.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

リスト内包表記

コードの可読性が一気に下がります。その上全然早くなりません。

def SlopeListComp(z):
    dx, dy = float(1), float(1)
    ny, nx = z.shape
    
    Slope = np.full_like(z, float(-9999))
    
    t =np.array([[ np.arctan( np.sqrt(
             (0.25*(
                0.5*(z[j+1,i+1] - z[j+1,i-1])/dx
              + 1.0*(z[j  ,i+1] - z[j  ,i-1])/dx
              + 0.5*(z[j-1,i+1] - z[j-1,i-1])/dx))**2
            +(0.25*(
                0.5*(z[j+1,i+1] - z[j-1,i+1])/dy
              + 1.0*(z[j+1,i  ] - z[j-1,i  ])/dy
              + 0.5*(z[j+1,i-1] - z[j-1,i-1])/dy))**2)
                          ) for i in range(1, nx-1) ] for j in range(1, ny-1)])
    
    Slope[1:-1,1:-1] = t
    
    return Slope

%%timeit
d = SlopeListComp(z)

5.76 s ± 154 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

スライス

高速化の基本です。爆速になります。

def SlopeSlice(z):
    dx, dy = float(1), float(1)
    ny, nx = z.shape
    
    Slope = np.full_like(z, float(-9999))
    
    dzdx = 0.25*(
        0.5*(z[2:  ,2:] - z[2:  , :-2])/dx
      + 1.0*(z[1:-1,2:] - z[1:-1, :-2])/dx
      + 0.5*(z[:-2 ,2:] - z[:-2 , :-2])/dx)
    
    dzdy = 0.25*(
        0.5*(z[2:, 2:  ] - z[:-2, 2:  ])/dy
      + 1.0*(z[2:, 1:-1] - z[:-2, 1:-1])/dy
      + 0.5*(z[2:,  :-2] - z[:-2,  :-2])/dy)

    Slope[1:-1, 1:-1] = np.arctan( np.sqrt(dzdx**2 + dzdy**2) )

    return Slope

%%timeit
d = SlopeSlice(z)

60.4 ms ± 863 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

スライス ケース2

可読性は下がりますが、ワンライナーで書きます。結果速度は全く変わりません。

def SlopeSlice2(z):
    dx, dy = float(1), float(1)
    ny, nx = z.shape
    
    Slope = np.full_like(z, float(-9999))
    
    Slope[1:-1, 1:-1] = np.arctan( np.sqrt(
         (0.25*(
        0.5*(z[2:  ,2:] - z[2:  , :-2])/dx
      + 1.0*(z[1:-1,2:] - z[1:-1, :-2])/dx
      + 0.5*(z[:-2 ,2:] - z[:-2 , :-2])/dx))**2
       + (0.25*(
        0.5*(z[2:, 2:  ] - z[:-2, 2:  ])/dy
      + 1.0*(z[2:, 1:-1] - z[:-2, 1:-1])/dy
      + 0.5*(z[2:,  :-2] - z[:-2,  :-2])/dy))**2
        ) )

    return Slope

%%timeit
d = SlopeSlice2(z)

    59.2 ms ± 2.46 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

numba

コーディング方法は以下の記事を参照。

computational-sediment-hyd.hatenablog.jp

本命です。基本コードの前に1行付け加えるだけですが当然最速です。

@jit(nopython=True)
def SlopeNumba(z):
    dx, dy = float(1), float(1)
    ny, nx = z.shape
    
    Slope = np.full_like(z, float(-9999))
    for j in range(1, ny-1):
        for i in range(1, nx-1):
            dzdx = 0.25*(
                0.5*(z[j+1,i+1] - z[j+1,i-1])/dx
              + 1.0*(z[j  ,i+1] - z[j  ,i-1])/dx
              + 0.5*(z[j-1,i+1] - z[j-1,i-1])/dx)
                
            dzdy = 0.25*(
                0.5*(z[j+1,i+1] - z[j-1,i+1])/dy
              + 1.0*(z[j+1,i  ] - z[j-1,i  ])/dy
              + 0.5*(z[j+1,i-1] - z[j-1,i-1])/dy)
            
            Slope[j,i] = np.arctan( np.sqrt(dzdx**2 + dzdy**2) )
            
    return Slope

%%timeit
d = SlopeNumba(z)

11.8 ms ± 246 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

numba + リスト内包表記

エラーが出て回りませんでした。一次元配列なら回りますが基本ケースのnumbaと同速度です。

@jit(nopython=True)
def SlopeListCompNumba(z):
    dx, dy = float(1), float(1)
    ny, nx = z.shape
    
    Slope = np.full_like(z, float(-9999))
    
    t =np.array([[ np.arctan( np.sqrt(
             (0.25*(
                0.5*(z[j+1,i+1] - z[j+1,i-1])/dx
              + 1.0*(z[j  ,i+1] - z[j  ,i-1])/dx
              + 0.5*(z[j-1,i+1] - z[j-1,i-1])/dx))**2
            +(0.25*(
                0.5*(z[j+1,i+1] - z[j-1,i+1])/dy
              + 1.0*(z[j+1,i  ] - z[j-1,i  ])/dy
              + 0.5*(z[j+1,i-1] - z[j-1,i-1])/dy))**2)
                          ) for i in range(1, nx-1) ] for j in range(1, ny-1)])
    
    Slope[1:-1,1:-1] = t
    
    return Slope

%%timeit
d = SlopeListCompNumba(z)

numba + スライス

基本ケースのnumbaと同速度です。なので、スライスで書く理由はないです。

@jit(nopython=True)
def SlopeSliceNumba(z):
    dx, dy = float(1), float(1)
    ny, nx = z.shape
    
    Slope = np.full_like(z, float(-9999))
    
    dzdx = 0.25*(
        0.5*(z[2:  ,2:] - z[2:  , :-2])/dx
      + 1.0*(z[1:-1,2:] - z[1:-1, :-2])/dx
      + 0.5*(z[:-2 ,2:] - z[:-2 , :-2])/dx)
    
    dzdy = 0.25*(
        0.5*(z[2:, 2:  ] - z[:-2, 2:  ])/dy
      + 1.0*(z[2:, 1:-1] - z[:-2, 1:-1])/dy
      + 0.5*(z[2:,  :-2] - z[:-2,  :-2])/dy)

    Slope[1:-1, 1:-1] = np.arctan( np.sqrt(dzdx**2 + dzdy**2) )

    return Slope

%%timeit
d = SlopeSliceNumba(z)

21.1 ms ± 297 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

おまけ

おまけ１：配列へのアクセスは結構時間がかかる

配列へのアクセス速度を確認するため、計算結果を配列に格納する行「Slope[j,i]…」をコメントアウトして計算すると計算時間が

高速になります。

つまり、計算時間の多くをメモリ上の配列へのアクセスに使っていることになります。 この部分を工夫することにより高速化できそうです。

def Slopetmp(z):
    dx, dy = float(1), float(1)
    ny, nx = z.shape
    
    Slope = np.full_like(z, float(-9999))
    for j in range(1, ny-1):
        for i in range(1, nx-1):
            dzdx = 0.25*(
                0.5*(z[j+1,i+1] - z[j+1,i-1])/dx
              + 1.0*(z[j  ,i+1] - z[j  ,i-1])/dx
              + 0.5*(z[j-1,i+1] - z[j-1,i-1])/dx)
                
            dzdy = 0.25*(
                0.5*(z[j+1,i+1] - z[j-1,i+1])/dy
              + 1.0*(z[j+1,i  ] - z[j-1,i  ])/dy
              + 0.5*(z[j+1,i-1] - z[j-1,i-1])/dy)
            
#             Slope[j,i] = np.arctan( np.sqrt(dzdx**2 + dzdy**2) )
            
    return Slope

%%timeit
d = Slopetmp(z)

2.98 s ± 91.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

@jit(nopython=True)
def SlopetmpNumba(z):
    dx, dy = float(1), float(1)
    ny, nx = z.shape
    
    Slope = np.full_like(z, float(-9999))
    for j in range(1, ny-1):
        for i in range(1, nx-1):
            dzdx = 0.25*(
                0.5*(z[j+1,i+1] - z[j+1,i-1])/dx
              + 1.0*(z[j  ,i+1] - z[j  ,i-1])/dx
              + 0.5*(z[j-1,i+1] - z[j-1,i-1])/dx)
                
            dzdy = 0.25*(
                0.5*(z[j+1,i+1] - z[j-1,i+1])/dy
              + 1.0*(z[j+1,i  ] - z[j-1,i  ])/dy
              + 0.5*(z[j+1,i-1] - z[j-1,i-1])/dy)
            
#             Slope[j,i] = np.arctan( np.sqrt(dzdx**2 + dzdy**2) )
            
    return Slope

%%timeit
d = SlopetmpNumba(z)

1.98 ms ± 332 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

おまけ2：ループの順を内側→外側

プログラミングの基本ではありますが、ループの順を内側→外側にすると計算が遅くなるはずです。 ところが、計算結果みると

  • 基本ケース場合、5.55s:外側→内側のケース5.54s
  • numbaの場合、15.9ms:外側→内側のケース11.8 ms

であり、numbaの場合に若干早くなる程度であまり変わらない。

なので、pythonの場合はループの順はあまり意識しなくても良いと思います。

def Slopetmp2(z):
    dx, dy = float(1), float(1)
    ny, nx = z.shape
    
    Slope = np.full_like(z, float(-9999))
    for i in range(1, nx-1):
        for j in range(1, ny-1):
            dzdx = 0.25*(
                0.5*(z[j+1,i+1] - z[j+1,i-1])/dx
              + 1.0*(z[j  ,i+1] - z[j  ,i-1])/dx
              + 0.5*(z[j-1,i+1] - z[j-1,i-1])/dx)
                
            dzdy = 0.25*(
                0.5*(z[j+1,i+1] - z[j-1,i+1])/dy
              + 1.0*(z[j+1,i  ] - z[j-1,i  ])/dy
              + 0.5*(z[j+1,i-1] - z[j-1,i-1])/dy)
            
            Slope[j,i] = np.arctan( np.sqrt(dzdx**2 + dzdy**2) )
            
    return Slope

%%timeit
d = Slopetmp2(z)

5.55 s ± 192 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

@jit(nopython=True)
def Slopetmp2Numba(z):
    dx, dy = float(1), float(1)
    ny, nx = z.shape
    
    Slope = np.full_like(z, float(-9999))
    for i in range(1, nx-1):
        for j in range(1, ny-1):
            dzdx = 0.25*(
                0.5*(z[j+1,i+1] - z[j+1,i-1])/dx
              + 1.0*(z[j  ,i+1] - z[j  ,i-1])/dx
              + 0.5*(z[j-1,i+1] - z[j-1,i-1])/dx)
                
            dzdy = 0.25*(
                0.5*(z[j+1,i+1] - z[j-1,i+1])/dy
              + 1.0*(z[j+1,i  ] - z[j-1,i  ])/dy
              + 0.5*(z[j+1,i-1] - z[j-1,i-1])/dy)
            
            Slope[j,i] = np.arctan( np.sqrt(dzdx**2 + dzdy**2) )
            
    return Slope

%%timeit
d = Slopetmp2Numba(z)

15.9 ms ± 313 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

おわりに

改めてnumba最強

このブログではpythonによる科学技術計算を強く推奨しています。

pythonを使用したことない方はぜひ下記を一読下さい。

computational-sediment-hyd.hatenablog.jp