Juliaで数値計算 その5:コードサンプル〜Typeに対する計算を切り替える方法

Juliaで数値計算　その1：コードサンプル〜基本的計算編〜
https://qiita.com/cometscome_phys/items/31d0b811345a3e12fcef
Juliaで数値計算　その2：コードサンプル〜わかりやすい書き方編〜
https://qiita.com/cometscome_phys/items/dfd63f1f7b51bdda4ca4
Juliaで数値計算　その3：コードサンプル〜2次元プロット・可視化編〜
https://qiita.com/cometscome_phys/items/19b041b2f3364705f428
Juliaで数値計算　その4:コードサンプル〜MPI並列計算編
https://qiita.com/cometscome_phys/items/f1971f8590a6d69d472b
に引き続いて、FortranやCからやってきた人のためにJuliaの書き方を紹介する。
この記事はその5である。

今回は、Fortranでは実装しにくい方法を使って、Juliaでの面白い書き方を見てみることにする。
新しいTypeを作って、そのTypeの要素にmoduleを入れることで、「そのTypeに対する計算を切り替える」ということをしてみる。

バージョン

Julia: 1.1.0

ModuleをフィールドにもつType

以下のようなTypeを考える。

module atoms
    export Atoms,get_energy
    mutable struct Atom
        rx::Float64
        ry::Float64
        rz::Float64
        px::Float64
        py::Float64
        pz::Float64        
    end

    mutable struct Atoms
        atoms::Array{Atom,1}
        num::Int64
        method::Module
    end

    module dummy
        using ..atoms
        function get_energy(atoms::Atoms)
            println("error! This is dummy! No method exist in atoms")
        end
    end

    function Atoms()
        Atoms([],0,dummy)
    end

    function get_energy(atoms::Atoms)
        atoms.method.get_energy(atoms)
    end

end

このmodule atomsには、Atomsというtypeがあり、このtypeはフィールドとして、atoms,num,methodを持つ。atomsは要素がtype Atomの配列、numは整数、methodはmoduleである。
type Atomは、x,y,z,px,py,pzというフィールドを持つ。想定としては、原子の座標と運動量を保持するtypeがAtomで、全体の情報を持っているのがAtomsとなる。
なお、初期化としてAtoms()を用意してある。

実は自分もこれまで気がついていなかったのだが、typeのフィールドとして、Moduleを持つことができる。初期化の時には、dummyというModuleを入れている。

このAtomsの"エネルギー"を求める関数をget_energyとしよう。
コードとして、Atoms型の引数を入れる関数get_energy(atoms)を呼べばアウトプットとしてエネルギーを出すようにしたい。そして、エネルギーを求める関数は必要に応じて切り替えられるようにするとする。一つの方法としては、get_energy(atoms,func)として、使う関数funcを引数とすることができる。しかし、エネルギーを求める関数以外にも例えばその微分を求める関数func2とかも切り替えて使いたいとすると、get_forces(atoms,func2)みたいにしなければならない。そして、今どの手法を使って計算しているかを意識しておかないと、funcとfunc2に何を入れるかがわからなくなる。
これを解決するために、methodをフィールドに入れた。例えば、

module testmodule
    using ..atoms
    function get_energy(atoms::Atoms)
        e = 4
    end
end

using .atoms
using .testmodule

testatom = Atoms()
println(testatom)
get_energy(testatom)
testatom.method = testmodule
e = get_energy(testatom)
println(e)

と新しいmoduleであるtestmoduleを定義して、methodフィールドにこれを代入testatom.method = testmoduleしてみる。すると、get_energy(testatom)をすればちゃんとtestmodule内のget_energyを使うことができる。
同様に、新しい関数get_forcesみたいなものを定義しておけば、methodを異なるmoduleにするだけで、手法を切り替えることができる。
もし、新しい手法を追加したい場合には、get_energyを新しいmoduleに定義してmethodを切り替えるだけでコードの他の部分を一切変える必要がなくなる。