cshiftの使いどころ

概要

配列要素を循環シフトするcshift関数は，境界条件によっては格子ボルツマン法のstreaming stepで使えそうです．

格子ボルツマン法

格子ボルツマン法は，流れの数値計算法の一つです．流れの数値計算法の多くは，流体を巨視的に捉えることによって導出されるNavier-Stokes方程式を離散化しており，これは偏微分や時間積分の数値計算法に，上流化など流体現象独自の取り扱いを導入していると考えられます．一方で，流れを微視的に見ると，途方もない数の水分子の集合として表され，分子の相互作用を直接計算するような手法も存在ます．

格子ボルツマン法は，どちらかというと流体を微視的に捉えていますが，分子そのものではなくそれらの分布という少し大きめのスケールで捉えています．この分布を，非常に強い制約の下で離散化することで，非常に簡単なアルゴリズムで流体現象の計算に成功しています．

格子ボルツマン法の支配方程式は，分布関数$f$の時間発展方程式です．

\frac{\partial f_k(t,\boldsymbol{x})}{\partial t} + \boldsymbol{c}_k\cdot\nabla f_k(t,\boldsymbol{x}) = \Omega_k(f)

格子ボルツマン法では，分布関数$f$を空間上で離散化して中間スケールの粒子とみなすだけでなく，その粒子が動ける方向も離散化します．これは，移動速度のベクトル$\boldsymbol{c}_k$が取り得る値を制限するということです．空間のある地点で離散化された中間スケールの粒子は，電車のように空間上に引かれたレールの上を隣の地点まで移動し，その過程で他の粒子と相互作用します．その相互作用を表すのが，$\Omega$で表された衝突項です．

簡単な相互作用のモデル（単一緩和時間の衝突則）を導入すると，粒子分布の時間発展方程式は次のようになります．

f_k(t+\varDelta t,\boldsymbol{x}+\boldsymbol{c}_k\varDelta t) = f_k(t,\boldsymbol{x}) - \frac{1}{\tau}\left[ f_k(t,\boldsymbol{x}) -  f_k^{\text{eq}}(t,\boldsymbol{x})\right]

$f_k^{\text{eq}}(t,\boldsymbol{x})$は平衡分布関数，$\tau$は緩和時間とよばれます．格子ボルツマン法では，上の式を，二つの段階に分けて計算します．

\begin{align}
f_k^*(t+\varDelta t,\boldsymbol{x}) &= f_k(t,\boldsymbol{x}) - \frac{1}{\tau}\left[ f_k(t,\boldsymbol{x}) -  f_k^{\text{eq}}(t,\boldsymbol{x})\right]\\
f_k(t+\varDelta t,\boldsymbol{x}+\boldsymbol{c}_k\varDelta t) &= f_k^*(t+\varDelta t,\boldsymbol{x})
\end{align}

まず衝突項を計算し，次に衝突項の影響が反映された分布関数$f_k^*$を速度$\boldsymbol{c}_k$で移動させます．最初の段階をcollision step，次の段階をstreaming stepとよびます．

2次元において，速度を9方向に離散化する2次元9速度モデルの概念図を示します．ある点$(i,j)$に存在する，粒子分布を値に持つ擬似的な粒子が，時刻$\varDelta t$秒後に9方向へ移動します．

粒子は9個ありますが，私の理解では，ある点の分布関数が9成分あるのではなく，9方向に移動することを表現するために，便宜上9個の粒子を用いているということかと思います．

この過程をプログラムで書くと，以下のような感じになります．dir_??は粒子番号をもつ列挙子で，pが当該軸の正方向への移動，mが負方向への移動です．

! f(i  , j  , dir_00) = f_upstream(i, j, dir_00) ! 移動しない粒子
f(i+1, j  , dir_p0) = f_upstream(i, j, dir_p0)
f(i  , j+1, dir_0p) = f_upstream(i, j, dir_0p)
f(i-1, j  , dir_m0) = f_upstream(i, j, dir_m0)
f(i  , j-1, dir_0m) = f_upstream(i, j, dir_0m)
f(i+1, j+1, dir_pp) = f_upstream(i, j, dir_pp)
f(i-1, j+1, dir_mp) = f_upstream(i, j, dir_mp)
f(i-1, j-1, dir_mm) = f_upstream(i, j, dir_mm)
f(i+1, j-1, dir_pm) = f_upstream(i, j, dir_pm)

Fortranでは，配列演算を用いてf(2:Nx, : , dir_p0) = f_upstream(1:Nx-1, :, dir_p0)などとも書けるのですが，どうしても配列の端（物理的には領域の境界）での取り扱いに工夫が必要になります．やり方としては，moduleを使って範囲から出た値を折りたたむとか，

integer(int32) :: i, j, ip1, im1, jp1, jm1
do concurrent(i=1:Nx, j=1:Ny)
    ip1 = modulo((i+1)-1, Nx) + 1
    im1 = modulo((i-1)-1, Nx) + 1
    jp1 = modulo((j+1)-1, Ny) + 1
    jm1 = modulo((j-1)-1, Ny) + 1

    f(ip1, j  , dir_p0) = f_upstream(i, j, dir_p0)
    f(i  , jp1, dir_0p) = f_upstream(i, j, dir_0p)
    f(im1, j  , dir_m0) = f_upstream(i, j, dir_m0)
    f(i  , jm1, dir_0m) = f_upstream(i, j, dir_0m)
    f(ip1, jp1, dir_pp) = f_upstream(i, j, dir_pp)
    f(im1, jp1, dir_mp) = f_upstream(i, j, dir_mp)
    f(im1, jm1, dir_mm) = f_upstream(i, j, dir_mm)
    f(ip1, jm1, dir_pm) = f_upstream(i, j, dir_pm)
end do

あるいは部分配列の記述を使って頑張るとかになります．

f(2:Nx  , 1:Ny  , dir_p0) = f_upstream(1:Nx-1, 1:Ny  , dir_p0)
f(1:Nx  , 2:Ny  , dir_0p) = f_upstream(1:Nx  , 1:Ny-1, dir_0p)
f(1:Nx-1, 1:Ny  , dir_m0) = f_upstream(2:Nx  , 1:Ny  , dir_m0)
f(1:Nx  , 1:Ny-1, dir_0m) = f_upstream(1:Nx  , 2:Ny  , dir_0m)
f(2:Nx  , 2:Ny  , dir_pp) = f_upstream(1:Nx-1, 1:Ny-1, dir_pp)
f(1:Nx-1, 2:Ny  , dir_mp) = f_upstream(2:Nx  , 1:Ny-1, dir_mp)
f(1:Nx-1, 1:Ny-1, dir_mm) = f_upstream(2:Nx  , 2:Ny  , dir_mm)
f(2:Nx  , 1:Ny-1, dir_pm) = f_upstream(1:Nx-1, 2:Ny  , dir_pm)
! 境界条件は以下で書く
f(1  , 1:Ny  , dir_p0) = f_upstream(Nx, 1:Ny  , dir_p0)
...

すべての境界が周期境界条件の場合，配列要素を循環シフトするcshiftを利用すると，この手間から解放されます．

循環シフト

cshiftはFortran 90から追加された組み込み手続（関数）で，引数にシフトする配列，シフト量，シフトする配列次元をとり，シフトされた配列を返します．シフト量はスカラでも配列（シフトする次元と同じ要素数）でも指定できます．

    integer(int32) :: a(3, 2)
    a(1:3, 1) = [1, 2, 3]
    a(1:3, 2) = [4, 5, 6]

    print '(3(I2))', cshift(a, 1, dim=1) ! 左方向に1要素分循環シフトする
    ! 2 3 1
    ! 5 6 4
    print '(3(I2))', cshift(a, [2, 1], dim=1) ! a(1:3, 1)を左方向に2要素，a(1:3, 2)を左方向に1要素循環シフトする
    ! 3 1 2
    ! 5 6 4

このcshiftを使えば，先ほどのstreaming stepは，次のように書けます．

f(:, :, dir_p0) = cshift(f_upstream(:, :, dir_p0), shift= 1, dim=dir_x)
f(:, :, dir_0p) = cshift(f_upstream(:, :, dir_0p), shift= 1, dim=dir_y)
f(:, :, dir_m0) = cshift(f_upstream(:, :, dir_m0), shift=-1, dim=dir_x)
f(:, :, dir_0m) = cshift(f_upstream(:, :, dir_0m), shift=-1, dim=dir_y)

dir_x, dir_yは，配列次元と空間次元を対応づける列挙子です．残念なことにcshiftはどれか一つの次元方向にしかシフトできないので，二つの次元をシフトするには，cshiftを2回用いる必要があります．

f(:, :, dir_pp) = cshift( &
                  cshift(f_upstream(:, :, dir_pp), shift=1, dim=dir_x), &
                                                   shift=1, dim=dir_y)

これは仕方ないこととして目をつむっても，毎回書くのは面倒なので，関数にしてしまいます．dimを配列として受け取れるようにしておきます．このとき，シフト量もシフトする次元も，シフトする配列arrayのランクと同じに制限しておきます．格子ボルツマン法のstreaming stepに用いる用途では，配列のシフト量は均一なので，これで問題ありません．

module function_op_unary_shift_circular
    use, intrinsic :: iso_fortran_env
    implicit none
    private
    public :: cshift

    interface cshift
        procedure :: cshift_2d_r32
        procedure :: cshift_2d_r64
    end interface
contains 
    pure function cshift_2d_r32(array, shift, dim) result(new_array)
        implicit none
        real(real32), intent(in) :: array(:, :)
        integer(int32), intent(in) :: shift(rank(array))
        integer(int32), intent(in) :: dim(rank(array))

        real(real32), allocatable :: new_array(:, :)

        new_array = cshift( &
                    cshift(array, shift=shift(1), dim=dim(1)), &
                    shift=shift(2), dim=dim(2))
    end function cshift_2d_r32

    pure function cshift_2d_r64(array, shift, dim) result(new_array)
        implicit none
        real(real64), intent(in) :: array(:, :)
        integer(int32), intent(in) :: shift(rank(array))
        integer(int32), intent(in) :: dim(rank(array))

        real(real64), allocatable :: new_array(:, :)

        new_array = cshift( &
                    cshift(array, shift=shift(1), dim=dim(1)), &
                    shift=shift(2), dim=dim(2))
    end function cshift_2d_r64

    ! 必要なら他の型に対する手続を追加する
end module function_op_unary_shift_circular

最終的に，streaming stepは次のように書けるようになりました．

use :: function_op_unary_shift_circular
f_upstream(:, :, :) = f(:, :, :)
f(:, :, dir_p0) = cshift(f_upstream(:, :, dir_p0), shift= 1, dim=dir_x) !&
f(:, :, dir_0p) = cshift(f_upstream(:, :, dir_0p), shift= 1, dim=dir_y) !&
f(:, :, dir_m0) = cshift(f_upstream(:, :, dir_m0), shift=-1, dim=dir_x) !&
f(:, :, dir_0m) = cshift(f_upstream(:, :, dir_0m), shift=-1, dim=dir_y) !&
f(:, :, dir_pp) = cshift(f_upstream(:, :, dir_pp), shift=[ 1,  1], dim=[dir_x, dir_y]) !&
f(:, :, dir_mp) = cshift(f_upstream(:, :, dir_mp), shift=[-1,  1], dim=[dir_x, dir_y]) !&
f(:, :, dir_mm) = cshift(f_upstream(:, :, dir_mm), shift=[-1, -1], dim=[dir_x, dir_y]) !&
f(:, :, dir_pm) = cshift(f_upstream(:, :, dir_pm), shift=[ 1, -1], dim=[dir_x, dir_y]) !&

格子ボルツマン法の入門書¹では，最初の問題としてTaylor-Green渦が取り扱われています．この問題は全方向が周期境界条件なので，cshiftを使えばすっきりと書けます．

まとめ

最初cshiftを知ったときはその利便性がイマイチ理解できなかったのですが，意外な使いどころがありました．知らないけど便利に使える手続がまだまだありそうな気がします．

1. 瀬田，格子ボルツマン法，森北出版，東京，2021． ↩