TI-Nspire & Lua / 特殊な方程式に対する補外法

参考: パソコンで見る天体の動き, pp.131-142

function gragg(funcs, t0, inits, initsDot, h)
   local unpack = unpack or table.unpack
   
   local t0       = t0
   local inits    = inits
   local initsDot = initsDot
   
   local function list2oneColMat(list)
            local oneColMat = {}
            for i = 1, #list do
               oneColMat[i] = {}
               oneColMat[i][1] = list[i]
            end
            return oneColMat
         end
   local function newMat(numRows, numCols, val)
            local mat = {}
            for i = 1, numRows do
               mat[i] = {} 
               for j = 1, numCols do
                  mat[i][j] = val or nil
               end
            end
            return mat
         end
   local function euler(funcs, t0, inits, initsDot, h)
            local S    = 7 -- テキストどおり 7 通り計算する。
            local temp = {}
            
            local x  = newMat(#funcs, 1, _) -- 一種のオイラー法で求めた値を入れておくための 1 列行列を用意しておく。
            local p  = newMat(#funcs, 1, _) -- 一種のオイラー法で求めた値を入れておくための 1 列行列を用意しておく。
            local p0 = {}                   -- 一種のオイラー法で求めた値を入れておくためのリストを用意しておく。
            
            local X = newMat(#funcs, 1, _) -- 補外にかける値を入れておくための 1 列行列を用意しておく。
            local U = newMat(#funcs, 1, _) -- 補外にかける値を入れておくための 1 列行列を用意しておく。
            
            for s = 1, S do
               local N  = 2^s
               local hs = h/N -- 内部分割する刻み幅の設定はテキストどおり h/2^s とする。
               
               for n = 0, N do
                  for i = 1, #funcs do
                     p0[i]     = initsDot[i] + (hs/2) * funcs[i](unpack(inits))
                     x[i][n+1] = (x[i][n] or inits[i]) + hs * (p[i][n] or p0[i])  
                     temp[i]   = x[i][n+1]  
                  end 
                  for i = 1, #funcs do
                     p[i][n+1] = (p[i][n] or p0[i]) + hs * funcs[i](unpack(temp))   
                  end
               end
               for i = 1, #funcs do
                  X[i][s] = x[i][N]
                  temp[i] = X[i][s] 
               end
               for i = 1, #funcs do
                  U[i][s] = p[i][N] - (hs/2) * funcs[i](unpack(temp))
               end
            end
            return X, U
         end
   local function neville(listOfLists)
            local numOfLists = #listOfLists
            local numOfHs    = #listOfLists[1]
            local X          = {}
            local extrapo    = {}
            for n = 1, numOfLists do
               X[n] = list2oneColMat(listOfLists[n]) -- 既知の点列を 1 列行列に変換する。
               for i = 2, numOfHs do
               for j = 2, i  do
                  X[n][i][j] = X[n][i][j-1] + (X[n][i][j-1] - X[n][i-1][j-1]) / (4^(j-1) - 1) --順次 1/2 刻みで計算する場合の公式。
               end
               end
               extrapo[n] = X[n][numOfHs][numOfHs] -- 補外値だけ返す。
            end
            return extrapo
         end
   return function()
             local temp1, temp2 = euler(funcs, t0, inits, initsDot, h)
             inits, initsDot    = neville(temp1), neville(temp2)  
             t0                 = t0 + h
             return t0, inits, initsDot  
          end
end


-- 確かめる。
do

-- この聯立微分方程式を補外法で解く。
local function xDotDot(x, y) return -x/((math.sqrt(x^2 + y^2))^3) end
local function yDotDot(x, y) return -y/((math.sqrt(x^2 + y^2))^3) end
local inits    = {1, 0}
local initsDot = {0, 1}

-- 1 ステップだけ計算する函数を作って、
local a = gragg({xDotDot, yDotDot}, 0, inits, initsDot, 1)

-- 初期値を更新しながら必要な回数だけ繰り返す。
for i = 1, 5 do
   local t0, inits, initsDot = a()
   print(t0, table.concat(inits, ", "), table.concat(initsDot, ", "))
end

end

これで、『パソコンで見る天体の動き』に解説されている常微分方程式の数値解法はオイラー法以外全部 Lua で実装できた。