Modern fortranでルンゲクッタ法

#はじめに
微分方程式を導いたら、次は気軽に解を導出したい。そんなとき、プログラムが煩雑だとやる気が出ないので、サクッと使えるルーチンをモダンフォートランで実装した。

#微分方程式の初期値問題
微分方程式の初期値問題は一般的に次の形をしている。

\frac{d\boldsymbol{x}}{dt}=f(t,\boldsymbol{x})

##ルンゲクッタ法
時刻 $t=0$ 秒における初期値を $\boldsymbol{x}_0$とする。
$h_i$を微小な時間刻み幅とし、

$t_n=t_{n-1}+h_n$ 秒とする。

このとき、$\boldsymbol{x}_n$を

$\boldsymbol{x}_{n-1}$と$h_n$から順次計算する手法の一つがルンゲクッタ法である。

\boldsymbol{k}_1=\boldsymbol{f}(t_{n-1},\boldsymbol{x}_{n-1})\\
\boldsymbol{k}_2=\boldsymbol{f}(t_{n-1}+\frac{h}{2},\boldsymbol{x}_{n-1}+\boldsymbol{k}_1 \frac{h}{2})\\
\boldsymbol{k}_3=\boldsymbol{f}(t_{n-1}+\frac{h}{2},\boldsymbol{x}_{n-1}+\boldsymbol{k}_2 \frac{h}{2})\\
\boldsymbol{k}_4=\boldsymbol{f}(t_{n-1}+h,\boldsymbol{x}_{n-1}+\boldsymbol{k}_3 h)\\
\boldsymbol{x}_n=\boldsymbol{x}_{n-1}+\frac{h}{6}
(\boldsymbol{k}_1+2\boldsymbol{k}_2+2\boldsymbol{k}_3+\boldsymbol{k}_4)

#fortranメインプログラム
GSLやRKSUITE、ODEPACKなどのライブラリと同じく、微分方程式の右辺はユーザーが定義する必要がある。ここで、Modern fortranの流儀でソルバを作ることで、メインプログラムを至極簡単に書くことができる。例は以下の通り。

program rkmain
  use rungekutta !ルンゲクッタ解法をまとめたモジュール
  implicit none
  real(kind=dpkind),dimension(3) :: x !状態変数
  type(rk) :: rktype  !ルンゲクッタ法ソルバに関連する型

  !微分方程式の右辺を計算する関数、解法、時間刻みの指定
  call rktype%setup(rhside,"rk4",1d-5) 
  !初期値の指定
  x=[0.5e0,0.5e0,0.5e0]
  !時刻0〜50秒まで、0.01秒毎に標準出力へ書き出し。
  !ただし、内部的には刻み幅1e-5秒としている。
  do 
    if(rktype%time>5e1)exit
    call rktype%solver(rktype%time+1e-2,x) !積分実行
    print*,rktype%time,x
  enddo
  !========================================
  contains
    !微分方程式の右辺を定義
    function rhside(self,time,wk)
      class(rk),intent(inout) :: self !parameter
      real(kind=dpkind) :: time
      real(kind=dpkind),dimension(:) :: wk
      real(kind=dpkind),dimension(size(wk)) :: rhside
      real(kind=dpkind),parameter :: a=10e0,b=28e0,c=8e0/3e0
      rhside(1)=a*(wk(2)-wk(1))
      rhside(2)=b*wk(1)-wk(2)-wk(1)*wk(3)
      rhside(3)=wk(1)*wk(2)-c*wk(3)
    end function
end program

実行結果をプロットすると次のようになる。

#ソルバを定義したモジュール
以下では、最も単純な固定刻み幅の解法を実装した。
刻み幅可変のルンゲクッタFehlberg法やAdams法、ギア法などを実装すれば、
setupルーチン引数で解法を指定するだけで切り替えることができるようにする設計思想にしている。（f77のODEPACKなどをラッピングしてもよいと思う。）

（2020/10/24追記）下記のソースコードは、とある工夫をすることで、大きく計算速度が向上します。

module rungekutta
  use,intrinsic :: iso_fortran_env
  private
  !---
  integer,parameter,public :: dpkind=real64 
  real(kind=dpkind),public,parameter :: PI=3.1415926535897932384626433_dpkind
  !---
  type,public :: rk
    real(kind=dpkind) :: time=0e0 !<計算開始からの総経過時間
    real(kind=dpkind) :: dt=1e-4   !<時間刻み
    procedure(rhside),pass,pointer :: fun => null() 
                !<微分方程式の右辺を計算する関数 
    procedure(rksol),pass,pointer :: solver => null()      
                !<微分方程式の解法
    contains
      final :: final_rk !<構造体のファイナライズ処理
      procedure :: setup !<数値積分の初期値、解法をセットする手続き
  end type
  !-------------
  interface
    !---
    ! Right hand side function interface
    function rhside(self,time,wk)
      import :: dpkind,rk
      class(rk),intent(inout) :: self !parameter
      real(kind=dpkind) :: time
      real(kind=dpkind),dimension(:) :: wk
      real(kind=dpkind),dimension(size(wk)) :: rhside
    end function
    !---
    ! Solver interface
    subroutine rksol(self,tnext,wk)
      import :: dpkind,rk
      class(rk),intent(inout) :: self
      real(kind=dpkind),intent(in) :: tnext !output time
      real(kind=dpkind),dimension(:) :: wk
    end subroutine
  end interface
  !============================--
  contains
    !---
    pure elemental subroutine final_rk(self)
      type(rk),intent(inout) :: self
      if(associated(self%fun))nullify(self%fun)
      if(associated(self%solver))nullify(self%solver)
    end subroutine
    !---
    subroutine setup(self,fun,solver,dt)
      class(rk),intent(inout) :: self
      procedure(rhside) :: fun  !微分方程式の右辺を計算する関数
      character(len=*),intent(in) :: solver !ソルバーの選択
      real(kind=dpkind),intent(in),optional :: dt !時間刻み幅 
      !時間刻みのセット
      if(present(dt))then
        self%dt=dt
      endif
      self%fun=>fun
      !ソルバーの選択
      select case(solver)
      case("rk4")
        !ルンゲクッタ法
        self%solver=>rk4
      case("rkf45")
        !ルンゲクッタ・フェルベルグ法(刻み幅自動調整型)
        !self%solver=>rkf45
      case default
        write(error_unit,*) "solver must one of the 'rk4, rkf45'"
      endselect
    end subroutine
    !---
    subroutine rk4(self,tn,wk)
      class(rk),intent(inout) :: self
      real(kind=dpkind),intent(in) :: tn !次に出力したい時間
      real(kind=dpkind),dimension(:) :: wk
      real(kind=dpkind),dimension(size(wk)) :: k1,k2,k3,k4
      real(kind=dpkind) :: dt_tmp
      real(kind=dpkind) :: dthalf
      associate(time=>self%time,dt=>self%dt)
        dt_tmp=dt
        do
          if(time+dt>tn)dt_tmp=tn-time
          dthalf=dt_tmp/2e0
          !!$OMP workshare
          k1=self%fun(time,wk)
          k2=self%fun(time+dthalf,wk+k1*dthalf)
          k3=self%fun(time+dthalf,wk+k2*dthalf)
          k4=self%fun(time+dt_tmp,wk+k3*dt_tmp)
          !状態変数を更新
          wk=wk+dt_tmp*(k1+2e0*k2+2e0*k3+k4)/6e0
          !!$OMP end workshare
          !1ステップ次の解が計算できたので、時間を進める
          time=time+dt_tmp 
          !時間の計算範囲がtnを超えたら計算終了
          if(time>=tn)then
            exit
          endif
        enddo
      end associate
    end subroutine
end module

#パラメータスタディなどを考える場合
微分方程式のパラメータを最適化する場合などは、パラメータを変更して何度も同じ計算を実行する必要があるかもしれない。そのようなときは、rk型を拡張して変数 $\boldsymbol{x}$ を型に入れ込むなどの対処をすると、拡張型の配列を用いて並列計算なども可能になる。

#おわりに
本稿で記した方法は、型とその拡張型を利用することによって応用範囲が無限に広がる。また、f77のようにサブルーチン引数を延々と意識する必要もないので、Modern Fortran はとても便利だと思う。