配列の添字を自由にする(構造保存解法コードサンプル有り)

Julia では配列の添字は 1 から始まる ¹ ことになっていて，しかも変えられない． 1から始まるこのルールは線形代数を始めとする多くの数学の場面で使い勝手の良いものだが，かといってそういう場面ばかりでもないよなあ．

これが自由に変えられれば便利なのになあ，ということは皆が思うところで，そのあたりを解消しようというパッケージもいくつかある．特に柔軟そうなものとして FlexibleArrays というのがあって，仕様や source をみるとかなり力が入っている． ただまだ開発途中で，最小限の機能もまだ実装されていないし ² ， 結構大きなバグ ³ もある．

こんな状況なので，言語仕様やパッケージに頼れるようになるのは少し先になりそう．

仕方ないので急場しのぎを考える

てきとー その1: マクロで数字を書き換える

以前のコードで使っている方法なんだけど，見えている数字を実際は違う数字に変えてしまうようなマクロを使うという手がある．

Julia のマクロはコードの(コンパイル前の) parse 時に実行されるので，コンパイル時には意図したものにコードが直っているという意味では筋がいいし，コードそのものが書き換わるので代入などの操作の左辺にも使えて便利だ． しかも Julia のマクロは文字列変換ではなく式変換なので，ほんとうの意味でコードを変換できる．

問題はコードの見た目がうるさくなることと，一部分でマクロの利用を忘れてもコードが動きかねない(わかりにくいバグの温床だぜ)というあたりと，それになにより Julia のマクロはわかりにくい ことだ． 上にも書いたように，Julia のマクロは式変換なので，まるで Lisp のような式を書くことになる のだ．

… とまあ，バランスを考えるとあまり複雑なマクロは嫌だなあ，というのが正直なトコロ．

まあ，実際の例を見てみよう． Julia のコードの中で

# 数式での離散関数の添字が -1 からなため，このズレに基づくトラブルに悩まさ
# れないように、離散関数 u(k) = (プログラム) @U(u,k) = 配列 u[k+2] となる
# ように マクロ @U を用意する．
# この @U(u,k) は実行前に u[k+2] そのものに変換され，「代入も可能」である!

macro U(v, k)
  Expr(:ref, v, Expr(:call, +, k, 2) )
end

とマクロを定義しておいて，

# 初期値代入
for k in -1:N+1
  @U(u, k) = initial_func(k)
end

なんて感じでコードを書いていくわけだ(マクロは利用時に @ をつけることになってる)．

ちなみに，数字の変換式だけをマクロにしないで配列名までマクロに巻き込んでいるのは，マクロの利用を忘れてバグを作りこむことを少しでも防ぐためのささやかな工夫である．泣かせるなあ．というか泣きたい．

てきとーその2: ベタでいいから型を用意しようぜ

マクロの利用は簡単だけど嫌な予感しまくりなので避けたい． というわけでベタでも何でも，使える「型」を作ろう．

まあ，一番簡単なのは，添字の最小値と最大値，そして配列値そのものを含めた型を作ることで，ここでもそうしてみるわけだ． 1 次元版ができれば多次元はなんとかなるだろうから，まずは 1次元に絞って話をしよう．

で，C で言う struct 型は Julia では composite type って呼ばれててこれを作るのは簡単だが，文法云々以前にちょいと気をつけないといけない点がある． というのは，Julia が multiple dispatch 言語だ，ってことだ． multiple dispatch 言語ってのは，

• object 指向言語から継承関連の機能を(見かけ上)除去して
• object に貼り付けてあった method/function を型定義の外で定義する

とする代わりに，同じ名前の function の定義を入力変数の型毎に別個に行う( = multiple dispatch )という言語，と思えばいい ⁴ ．

だから，Julia では型を作ったら，それにあわせて自分の使う(既にある)関数の新しい定義を与えないと不便だ．まあ，与えなくても十分使えるけど，せっかく作った型を自然に使いたかったらそうした方がいいのだ．これが Julia における，「新しく型を作る時の」ちょっとした注意点．

さて，ぐだぐだと書いたけど，まあ現物を見てみよう．

# 例えば u = Array1D(-3,5) とすると，これは添字が
# -3,-2,...,4,5 である実数配列となる．
type Array1D
    # 下記の変数の type をきちんと指定すると ForwardDiff 
    # などのパッケージがうまく動かない．よって曖昧にしておく．
    firstidx # 添字範囲の最小値．Int64 を想定．
    endidx   # 添字範囲の最大値．Int64 を想定．
    inarr    # Array{Float64,1}( endidx - firstidx +1 ) を想定している.
    # 全部具体的に与える constructor
    Array1D(firstidx,endidx, v) = new(firstidx,endidx,copy(v)) 
    # 範囲だけ与える contructor
    Array1D(firstidx,endidx) = new(firstidx,endidx,Array{Float64,1}(endidx-firstidx+1)) 
end

### 添字を指定して中身を見たり変更できるように．
Base.getindex(T::Array1D, key) = T.inarr[key-T.firstidx+1]
Base.setindex!(T::Array1D, value, key) = setindex!(T.inarr, value, key-T.firstidx+1)
### iterative 処理等も受けられるようにする. 全部は面倒なので，使う範囲だけ．
Base.start(S::Array1D) = S.firstidx
Base.next(S::Array1D,state) = (S[state], state+1)
Base.done(S::Array1D,state) = state > S.endidx
Base.eltype(::Type{Array1D})= Float64
Base.length(S::Array1D) = S.endidx-S.firstidx+1

Base.copy(S::Array1D) = Array1D(S.firstidx,S.endidx,copy(S.inarr))
Base.similar(S::Array1D) = Array1D(S.firstidx,S.endidx)
Base.map(f::Function, S::Array1D) = Array1D(S.firstidx,S.endidx,map(f,S.inarr))

最初の type で型を作っている． 中で，定義しておくと何かと便利な constructor も定義してある．

その後は，様々な通常の関数の，この型での動作の定義をしている． 配列周りの関数はほとんどは Julia の Base で定義されているので，これをこのように重ね書き(multiple dispatch)することになる．

で，これを使うと，例えば先の構造保存数値解法のプログラムは以下のように書き換えられる． 添字に変な変換を必要とすることもなく，ForwardDiff パッケージを利用する部分(Newton! 関数)以外では, sum とか copy とか map とかが新しい型にも使えて，無理なくコードが書けるようになる ⁵ のがいい感じだ．

あと，この型とは関係ないけど，台形則を毎回コードするのはバカみたいだったので Tsum という関数を定義しておいた．この関数の定義は，もとからある sum の定義に沿ったもので， for などを使うよりも汎用的なため，こうしておくと配列以外の変数にも Tsum が使えるようになるのだ．

コード全体を通して他の部分も少しいじっているが，まあ本質的にはおんなじで，もちろん計算結果も全く同じだった．

### 自動微分 パッケージ
using ForwardDiff

# 例えば u = Array1D(-3,5) とすると，これは添字が
# -3,-2,...,4,5 である実数配列となる．
type Array1D
    # 下記の変数の type をきちんと指定すると ForwardDiff 
    # などのパッケージがうまく動かない．よって曖昧にしておく．
    firstidx # 添字範囲の最小値．Int64 を想定．
    endidx   # 添字範囲の最大値．Int64 を想定．
    inarr    # Array{Float64,1}( endidx - firstidx +1 ) を想定している.
    # 全部具体的に与える constructor
    Array1D(firstidx,endidx, v) = new(firstidx,endidx,copy(v)) 
    # 範囲だけ与える contructor
    Array1D(firstidx,endidx) = new(firstidx,endidx,Array{Float64,1}(endidx-firstidx+1)) 
end

### 添字を指定して中身を見たり変更できるように．
Base.getindex(T::Array1D, key) = T.inarr[key-T.firstidx+1]
Base.setindex!(T::Array1D, value, key) = setindex!(T.inarr, value, key-T.firstidx+1)
### iterative 処理等も受けられるようにする. 全部は面倒なので，使う範囲だけ．
Base.start(S::Array1D) = S.firstidx
Base.next(S::Array1D,state) = (S[state], state+1)
Base.done(S::Array1D,state) = state > S.endidx
Base.eltype(::Type{Array1D})= Float64
Base.length(S::Array1D) = S.endidx-S.firstidx+1

Base.copy(S::Array1D) = Array1D(S.firstidx,S.endidx,copy(S.inarr))
Base.similar(S::Array1D) = Array1D(S.firstidx,S.endidx)
Base.map(f::Function, S::Array1D) = Array1D(S.firstidx,S.endidx,map(f,S.inarr))

### 一般の collection に使える台形則. 
### 純粋な和 sum のコードを少し変更したもの．最後の足し算1回は本来減らせるが，まあいいか．
function Tsum(S)

  state = start(S)
  (val, state) = next(S, state)
  result = val/2.0

  while !done(S, state)
    (val, state) = next(S, state)
    result += val
  end
  result -= val/2.0

  return result
end


### 各種計算パラメータ
const L = 2.0
const N = 200
const Dx = L/N

const T = 1.0
const Dt = 10.0^(-3)
const N_time = Int(round(T/Dt))
const skip_time = div(N_time,10)

const gamma_r = Dt/(Dx^2)

const data_fname = "data.dat"
const I_fname    = "I.dat"
const M_fname    = "M.dat"

const Newton_Criterion = 10.0^(-6)

### 初期値を決める関数．境界条件に合っている．
function u0(x)
  A0 = 1.0; A1 = 2.0; B2 = 0.3;
  A0 + A1 - A1 * cos(pi*x/L) + B2*cos(2*pi*2x/L)
  end

### 近似解の値の出力. こういうのは用意しておくと良いねえ．
function result_out(v::Array1D, time, data_f, I_f, M_f)
  println(data_f, "\n## time = $(time)")

  for k in 0:N
    println(data_f, v[k] )
  end

  println(I_f, time, " ", I(v) )
  println(M_f, time, " ", M(v) )

end

### 散逸量 I[u], 保存量 M[u] を計算する．
I(v) = Dx * Tsum(v)  # 単なる，台形則による値の離散和
M(v) = Dx * Tsum(map(log, v)) # log を適用してから台形則．

### exp の差分商式
dexp(a,b) = ifelse( a != b, ( exp(a)-exp(b) )/( a - b ), exp(a) )

### DVDMスキームそのものの式．現在値と次ステップの値をもらって、本来ゼロになるべき量．
function scheme(vp,v) 

  # Array1D の similar を使わずに, 実際に変化する array のサイズを julia に認識させる．
  # そうしないと，ForwardDiff package が正しく動作しない．まあある種のバグとも言えるが…
  fia = copy(vp.inarr) 
  f  = Array1D(v.firstidx, v.endidx, fia)

  # 本来の方程式+BC の構成は DAE(代数微分方程式)に近い．数値的には DAE は解きにくいことで有名で厄介だ．
  # よって、解くときは方程式に境界条件を繰り込んで、純粋な微分方程式に近い形に直したほうが、数値的には安定して解きやすい．
  # その結果が以下のとおりで、最初と最後の式が境界条件を受けて変形する．
  f[0] -= v[0] + 2 * gamma_r * ( dexp( vp[1], v[1] ) - dexp( vp[0], v[0] ) )
  f[N] -= v[N] + 2 * gamma_r * ( dexp( vp[N-1], v[N-1] ) - dexp( vp[N], v[N] ) )

  # 残りの式
  for k in 1:N-1
    f[k] -= v[k] + gamma_r * ( dexp( vp[k+1], v[k+1] ) - 2 * dexp( vp[k], v[k] ) + dexp( vp[k-1], v[k-1] ))
  end

  return f.inarr

end

### Newton 法でスキームの解を求める．Jacobian を求めるのに自動微分を使っている．なんて楽なんだ…
function Newton!(vp, v)

  # input/output for jacobian should be a normal array since of the ForwardDiff package
  func(w) = scheme(Array1D(vp.firstidx, vp.endidx, w), v)

  local i = 1
  while true

    # 十分にベクトル f がゼロに近ければ終了
    if sum(abs, func(vp.inarr)) < Newton_Criterion
      println("Final Newton loop number: ", i) # 最終的な Newton 反復の回数をみておこう．
      return vp # 本来は返さなくてももとの vector が書き換わらないといけないんだが．
    end

    # Newton 反復．これを教科書通りに 1行でプログラムにできる日が来るとは感動的だなあ．
    vp.inarr -= jacobian(func, vp.inarr) \ func(vp.inarr)
    i += 1

  end

end

### 以下、メイン文
println("Program Started.")

# 時間測定開始!
tic() 

# 現在ステップ、次ステップの近似解の入れ物を用意
u = Array1D(0,N) # 添字が 0:N になる配列を用意する(添字範囲は自由に取れる)．
# 初期値代入
for k in 0:N
  u[k] = u0( k*Dx )
  end

# up = similar(u)
up = copy(u) # いったん中身をコピーしてから
for k in 0:N
  up[k] += 0.0001 * randn() # 差分商の微分を考えると, 全く同じはよくないので少しずらす．
end

# 計算は v = log(u) 相当でやるので, それらを用意する．
lu = map(log,u)
lup = map(log,up)

# ファイル open
data_file = open(data_fname, "w")
I_file = open(I_fname, "w")
M_file = open(M_fname, "w")

result_out(u, 0.0, data_file, I_file, M_file)

### 時間発展
for loop = 1:N_time

  # Newton 法で 1step 計算する．
  lup = Newton!(lup, lu)

  # 毎回出力するとデータが多すぎるので、skip_time ステップごとに出力する．
  if mod(loop, skip_time) == 0
    up = map(exp,lup) # 必要なときだけ u = exp(v) に戻す
    result_out(up, loop*Dt, data_file, I_file, M_file)
    end

  # u と up の中身交換(正確にはポインタを交換してるだけ)
  lu, lup = lup, lu 

end

# ファイル close
close(data_file)
close(I_file)
close(M_file)

# 時間測定終了! 実際に計算してみると、さすがに Euler 法より遅いけど，
# Crank-Nicolson 法とほぼ同じ計算時間になる．
s = toq()
# 計算にかかった時間を報告
println("Program ended.")
println("# elapsed time: ", s, " seconds")