Python による常微分方程式の数値解法 / 古典的 Runge-Kutta 法

RK4.py

# Python による常微分方程式の数値解法 / 古典的 Runge-Kutta 法
class RK4:
    def __init__(self, funcs, t0, inits, h, numOfDiv=1):
        self.funcs    = funcs
        self.t0       = t0
        self.inits    = inits
        self.dim      = len(funcs)
        self.numOfDiv = numOfDiv
        self.h        = h / self.numOfDiv
        self.f        = [[None for i in range(self.dim)] for i in range(4)]
        self.temp     = [[None for i in range(self.dim)] for i in range(3)]
    def update(self):
        for i in range(self.numOfDiv):
            for j in range(self.dim):
                self.f[0][j] = self.funcs[j](self.t0 , *self.inits)
                self.temp[0][j] = self.inits[j] + self.h/2 * self.f[0][j]
            for j in range(self.dim):
                self.f[1][j] = self.funcs[j](self.t0 + self.h/2 , *self.temp[0])
                self.temp[1][j] = self.inits[j] + self.h/2 * self.f[1][j]
            for j in range(self.dim):
                self.f[2][j] = self.funcs[j](self.t0 + self.h/2 , *self.temp[1])
                self.temp[2][j] = self.inits[j] + self.h * self.f[2][j]
            for j in range(self.dim):
                self.f[3][j] = self.funcs[j](self.t0 + self.h , *self.temp[2])
            for j in range(self.dim):
                self.inits[j] += self.h/6 * (self.f[0][j] + 2 * (self.f[1][j] + self.f[2][j]) + self.f[3][j])
            self.t0 += self.h
        return self
    def print(self):
        print(self.t0, self.inits)
        return self


########
# test #
########
if __name__ == "__main__":
    # 解くべき聯立微分方程式を定義する。リストで括っておく。
    def xDot(t, x, y):  # x'(t) = y(t)
        return y
    def yDot(t, x, y):  # y'(t) = t - x(t)
        return t - x
    funcs = [xDot, yDot]

    # 独立変数の開始値と終了値とを指定する。
    t0 = 0
    tMax = 7.5

    # 従属変数の初期値を指定する。リストで括っておく。
    x0 = 0
    y0 = 0
    inits = [x0, y0]

    # 刻み幅を指定する。2 の整数乗分の 1 にすることが望ましい。
    h = 1/2**2

    # 1 刻みの内部分割数を指定する場合は 2 の整数乗にすることが望ましい。
    # numOfDiv = 2**4

    # 1 刻みだけ計算する函数を実体化して、
    sol = RK4(funcs, t0, inits, h)

    # 初期値を更新しながら必要な回数だけ実行を繰り返す。
    while sol.t0 < tMax:
        sol.update().print()

RK4_modified.py

# Python による常微分方程式の数値解法_改良版 / 古典的 Runge-Kutta 法
class RK4:
  def __init__(self, funcs, inits, h, numOfDiv=1):
    self.funcs = funcs
    self.t0 = 0
    self.inits = inits
    self.h = h / numOfDiv
    self.half_h = self.h / 2
    self.sixth_h = self.h / 6
    self.range_dim = range(len(funcs))
    self.range_div = range(numOfDiv)
    self.f = [[0 for i in self.range_dim] for i in range(4)]
    self.temp = [[0 for i in self.range_dim] for i in range(3)]

  def update(self):
    for i in self.range_div:
      self.f[0] = [self.funcs[j](self.t0, *self.inits) for j in self.range_dim]
      self.temp[0] = [self.inits[j] + self.half_h * self.f[0][j] for j in self.range_dim]
      self.f[1] = [self.funcs[j](self.t0 + self.half_h, *self.temp[0]) for j in self.range_dim]
      self.temp[1] = [self.inits[j] + self.half_h * self.f[1][j] for j in self.range_dim]
      self.f[2] = [self.funcs[j](self.t0 + self.half_h, *self.temp[1]) for j in self.range_dim]
      self.temp[2] = [self.inits[j] + self.h * self.f[2][j] for j in self.range_dim]
      self.f[3] = [self.funcs[j](self.t0 + self.h , *self.temp[2]) for j in self.range_dim]
      self.inits = [self.inits[j] + self.sixth_h * (self.f[0][j] + 2 * (self.f[1][j] + self.f[2][j]) + self.f[3][j]) for j in self.range_dim]
      self.t0 += self.h
    return self

  def print(self):
    print(self.t0, self.inits)
    return self

########
# test #
########
if __name__ == "__main__":
#def main():
  # 解くべき連立常微分方程式を定義する。
  # 確認のため、厳密解の存在する方程式を定義する。単振り子とは無関係。
  def xDot(t, x, y): # x'(t) = y(t)
    return y
  def yDot(t, x, y): # y'(t) = t - x(t)
    return t - x
  funcs = [xDot, yDot]

  # 従属変数の初期値を指定する。
  x0 = 0
  y0 = 0
  inits = [x0, y0]

  # ステップ幅を指定する。
  h = 0.25 # =1/2**2
  
  # 1ステップだけ計算する函数を実体化して、
  sol = RK4(funcs, inits, h)

  # 初期値を何度か更新して確認する。
  while sol.t0 < 7.5:
    sol.update().print()