NaysMiniを使ってみたらサクッと回ったけど遅かったのでnumbaで爆速化した - 趣味で計算流砂水理 Computational Sediment Hydraulics for Fun Learning

########### Main ############
# 追記
fig, ax = plt.subplots(figsize = (xsize, ysize))
# ----
while time<= etime:
    usta,ep,ep_x=fric.us_cal(usta,ep,ep_x,u,v,hs,nx,ny,snm,g_sqrt,hmin,ep_alpha)
    if icount%it_out==0:
        print('time=',np.round(time,3),l)
---
---
---
# 削除
#        fig, ax = plt.subplots(figsize = (xsize, ysize))
---
---
---
# 削除
#        plt.clf()
#        plt.close()
# 追記
        cb.remove() 
        plt.cla()

参考サイト

修正したコードはこちら

速度計測

jupyter notebook上の%%timeコマンドで計測しました。
ただし、ffmpegによるgif作成の部分は含んでおりません。
コード全体とmatplotlibによる図化部をコメントアウトした場合の2ケースで計測しました。
PCのスペックはこちら
オリジナルコードの計算時間は、
- コード全体 : Wall time: 9min
- 図化なし：Wall time: 6min 47s
- 図化時間：9min - 6min 47s = 2min13s

私の感覚では結構かかるなという感じです。

numbaで爆速化

ではnumbaで高速化していきます。numbaをご存じない方は以下を参照ください。

作業は非常に簡単で、whileループの中の関数をnumbaにします。（一部オリジナルでnumbaを使用している。）

例えば、fric.pyのus_cal関数を例に取ると、import文と関数の前に@jit--を記述するだけです。

import numpy as np
from numba import jit #追記

@jit(nopython=True) #追記
def us_cal(usta,ep,ep_x,u,v,hs,nx,ny,snm,g_sqrt,hmin,ep_alpha):
    for i in np.arange(1,nx+1):
        for j in np.arange(1,ny+1):
            ux=(u[i,j]+u[i-1,j])*.5
            vx=(v[i,j]+v[i,j-1])*.5
---
---
---

同様に、boundary.py ,stgg.py ,cfxy.py ,rhs.py ,diffusion.py ,newgrd.py ,cip2d.py ,uxvx.py ,fric.py ,mkzero.py ,hcal.pyにも書き加えます。

修正したコードはこちら

修正したコードの計算時間は、

コード全体 : Wall time: 3min17s（3.7倍高速化）
図化なし：5.85 s （69.6倍高速化）
図化時間：3min17s - 5.85s = 2min11s

となり、計算部分は爆速化しました。でも、図化部分は変わらないので、結果的には3.7倍程度です。

この程度の計算で3minというのはちょっと時間がかかり過ぎな気がします。

図化部分を分ける

MVCやMVVMの観点からも、計算と図化は分けたほうが良いので計算結果はファイルに出力する形にします。

出力形式はxarrayを使ったNetCDF形式とします。

xarrayのインストール方法は以下を参照。conda推奨

https://xarray.pydata.org/en/stable/getting-started-guide/installing.html#instructions

import xarray as xr
---
---
---
########### Main ############
fig, ax = plt.subplots(figsize = (xsize, ysize))
while time<= etime:
    usta,ep,ep_x=fric.us_cal(usta,ep,ep_x,u,v,hs,nx,ny,snm,g_sqrt,hmin,ep_alpha)
    if icount%it_out==0:
        print('time=',np.round(time,3),l)
        nfile=nfile+1
        ux,vx,uv2,hx,hsx=uxvx.uv(ux,vx,uv2,hx,hsx,u,v,h,hs,nx,ny)
        vor=uxvx.vortex(vor,ux,vx,nx,ny,dx,dy)
        
# 追加--------------
        dss = xr.Dataset({'vortex': (['t','x','y'], np.array([vor]))
                          , 'u': (['t','x','y'], np.array([ux]))
                          , 'v': (['t','x','y'], np.array([vx]))}
                          , coords={'t':[round(time,3)] , 'x': x, 'y':y}
                           )
        out = dss.to_netcdf('out' + str(nfile).zfill(8) + '.nc')
        out = dss.close()
#------------------
---
---
---

修正したコードはこちら

このときの計算時間は、Wall time: 12.9 sであり実用的な時間になりました。

図化

図化にはこのブログではおなじみのholovizファミリーを使います。

computational-sediment-hyd.hatenablog.jp

condaでの環境構築は下のとおりです。（2022/1/2時点）

conda create -y -n holo python=3.7
conda activate holo
# https://xarray.pydata.org/en/stable/getting-started-guide/installing.html#instructions
conda install -y -c conda-forge xarray dask netCDF4 bottleneck
conda install -y -c conda-forge hvplot
conda install -y -c conda-forge selenium
conda install -y -c conda-forge firefox geckodriver
conda install -y -c conda-forge jupyter

元サイトと同様に渦度と流速ベクトル図を書きます。

import numpy as np
import xarray as xr
import hvplot.xarray
import glob
import holoviews as hv
hv.extension('bokeh')

fs = glob.glob('out0*.nc')
fs.sort()
dd = []
for f in fs:
    ds = xr.open_dataset(f)
    U = ds['u'].values
    V = ds['v'].values
    Vmag = np.sqrt( U**2 + V**2) + 0.00000001
    angle = (np.pi/2.) - np.arctan2(U/Vmag, V/Vmag)
    ds['Vmag'] = (['t','x','y'], Vmag)
    ds['angle'] =(['t','x','y'], angle)
    ds = ds.drop(['u','v'])
    dd.append(ds)
    
dss = xr.concat(dd, dim="t")

fVec = dss.hvplot.vectorfield(x='x', y='y', groupby='t', angle='angle', mag='Vmag',hover=False)
fVor = dss['vortex'].hvplot(frame_height=150, frame_width=600, x='x', y='y', cmap='bwr', clim=(-10,10)) 

with open('obst.dat', mode='r', encoding='utf-8') as response:
    l = next(response)
    l = next(response)

ll = l.replace('\n','').split(',')
i1,i2,j1,j2 = np.array(ll, dtype=int)

x1 = dss['x'].values[i1]
x2 = dss['x'].values[i2]
y1 = dss['y'].values[j1]
y2 = dss['y'].values[j2]

fPoly = hv.Polygons(np.array([[x1,y1],[x2,y1],[x2,y2],[x1,y2]])).options(fill_color='green')

g = fVor * fVec * fPoly

d = hvplot.save(g,'out.html')
del d

出力されたhtmlファイルは、以下のような感じでユーザー側でコントロールするものになりwebフレンドリーです。

f:id:SedimentHydraulics:20211226184135g:plain

GitHubにインタラクティブなグラフをあげました。

https://computational-sediment-hyd.github.io/2DH_Python/out.html

上記をgif化する場合、結構時間がかかりますが次のように書きます。

from PIL import Image

for i, t in enumerate( dss['t'].values ):
     gp = g[t].options(title=str(np.round(t,3)) + ' sec', toolbar=None)
     d = hvplot.save(gp,'png'+str(i).zfill(8) +'.png')
del d

fs = glob.glob("*.png")
imgs = [Image.open(f) for f in fs]
imgs = imgs[0:501:2]
# appendした画像配列をGIFにする。durationで持続時間、loopでループ数を指定可能。
d= imgs[0].save('out.gif',
             save_all=True, append_images=imgs[1:], optimize=False, duration=20, loop=0)
del d

出力した図はこちら

f:id:SedimentHydraulics:20211226184429g:plain

コード全体はこちら

まとめ

計算部分はnumbaで爆速化（約70倍）
計算部分と図化部分は分けたほうが良い。
図化ライブラリはHolovizがおすすめ。

pythonは遅いと言われますが少し工夫するだけで爆速化するので科学技術計算でも十分に使えると思います。

むしろfortranの開発コストを考えれば、小～中規模のソフトウェアであればpythonのほうが優位と考えています。

おまけ

計算結果の渦度ですが左右対称になっていない気がします。ちょっとチェックしてみます。
元コードはループをバラバラに記述しているため、工夫次第でまだまだ早くなります。多分、処理をわかりやすくするためにそのような書き方にしていると思いますが、コード全体を抽象化することにより可読性と速度を両立したコードになると思います。
リクエストがあれば改良版のコードも公開します。⇒以下で公開しています。

GitHub

元のものをforkしてつくりました。

GitHub - computational-sediment-hyd/2DH_Python

リンク