Python による常微分方程式の数値解法 / Fehlberg 法

Fehlberg.py

# Python による常微分方程式の数値解法 / Fehlberg 法
from math import *

class Fehlberg:
    def __init__(self, funcs, t0, inits, h, tol=1e-3):
        self.funcs    = funcs
        self.t0       = t0
        self.inits    = inits
        self.h        = h
        self.numOfDiv = None
        self.dim      = len(funcs)
        self.tol      = tol
        self.f        = [[None for i in range(self.dim)] for i in range(6)]
        self.temp     = [[None for i in range(self.dim)] for i in range(5)]
        self.temp2    = [None for i in range(self.dim)]
        self.inits4   = [None for i in range(self.dim)]
    def __floorB(self, num):
        if 0 < num < 1:
            return 2**floor(log(num)/log(2))
        else:
            return floor(num)
    def __hNew(self, h, err, tol):
        if err > tol:
            return 0.9 * h * (tol * h/(h * err))**(1/4)
        else:
            return h
    def __maxOfErr(self, listA, listB):
        for i in range(self.dim):
            self.temp2[i] = abs(listA[i] - listB[i])
        return max(self.temp2)
    def __preCalc(self, numOfDiv):
        preInits = [*self.inits]
        preT0    = self.t0
        preH     = self.h/numOfDiv
        for i in range(numOfDiv):
            for j in range(self.dim):
                self.f[0][j] = self.funcs[j](preT0 , *preInits)
                self.temp[0][j] = preInits[j] + preH * (1/4) * self.f[0][j]
            for j in range(self.dim):
                self.f[1][j] = self.funcs[j](preT0 + preH * (1/4) , *self.temp[0])
                self.temp[1][j] = preInits[j] + preH * (1/32) * (3 * self.f[0][j] + 9 * self.f[1][j])
            for j in range(self.dim):
                self.f[2][j] = self.funcs[j](preT0 + preH * (3/8) , *self.temp[1])
                self.temp[2][j] = preInits[j] + preH * (1/2197) * (1932 * self.f[0][j] - 7200 * self.f[1][j] + 7296 * self.f[2][j])
            for j in range(self.dim):
                self.f[3][j] = self.funcs[j](preT0 + preH * (12/13), *(self.temp[2]))
                self.temp[3][j] = preInits[j] + preH * ((439/216) * self.f[0][j] - 8 * self.f[1][j] + (3680/513) * self.f[2][j] - (845/4104) * self.f[3][j])
            for j in range(self.dim):
                self.f[4][j] = self.funcs[j](preT0 + preH , *self.temp[3])
                self.temp[4][j] = preInits[j] + preH * ((-8/27) * self.f[0][j] + 2 * self.f[1][j] - (3544/2565) * self.f[2][j] + (1859/4104) * self.f[3][j] - (11/40) * self.f[4][j])
            for j in range(self.dim):
                self.f[5][j] = self.funcs[j](preT0 + preH * (1/2) , *self.temp[4])
            if numOfDiv == 1:
                for j in range(self.dim):
                    self.inits4[j] = preInits[j] + preH * ((25/216) * self.f[0][j] + (1408/2565) * self.f[2][j] + (2197/4104) * self.f[3][j] - (1/5) * self.f[4][j])
            for j in range(self.dim):
                preInits[j] = preInits[j] + preH * ((16/135) * self.f[0][j] + (6656/12825) * self.f[2][j] + (28561/56430) * self.f[3][j] - (9/50) * self.f[4][j] + (2/55) * self.f[5][j])
            preT0 = preT0 + preH
        return preT0, self.inits4, preInits
    def update(self):
        self.numOfDiv = 1
        t, a4, a5     = self.__preCalc(self.numOfDiv)
        err           = self.__maxOfErr(a4, a5)
        if err < self.tol:
            self.t0, self.inits = t, a5
        else:
            self.numOfDiv = int(self.h/self.__floorB(self.__hNew(self.h, err, self.tol)))
            t, a4, a5     = self.__preCalc(self.numOfDiv)
            self.t0, self.inits = t, a5
        return self
    def print(self):
        print(self.t0, self.inits, self.numOfDiv)
        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

    # 許容誤差を指定する。デフォルトは 1e-3。
    tol = 1e-9

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

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