自分で作ったtypeをベクトルとみなす in Julia

この記事では、自分で作ったtypeをAbstractVectorのsubtypeとみなす方法について述べます。
ベクトルに対して色々できるパッケージがあるので、自前の型もベクトルとしてしまおう、ということです。

バージョン

Julia 1.1.0

線形演算子とそれを作用させる"ベクトル"

Juliaでは有力な線形計算パッケージがあります。
例えば、連立方程式を解きたい場合には、IterativeSolvers.jl、少数の固有値を得たい場合には、Arpack.jlがあります。
IterativeSolvers.jlにおいては、連立方程式Ax=bの解を求めるためには、共役勾配法(CG法)を使います。この時、逐次演算では行列ベクトル積を何回も演算させます。
しかし、行列Aがあまりにも大きかった場合、メモリにAが入らずに計算ができない！というような場合があります。一つの方法としては、Aが疎行列であれば、行列要素のうち0ではない部分だけを格納することで、メモリを節約することができます。Juliaではその場合にはSparseArrays.jlというパッケージがあります。
一方、非ゼロ要素が沢山ある場合でも、行列自体をメモリに保持するのではなく、あるベクトルv
に対してAvのみを返すようにすれば、計算が可能です。
そのようなケースに対応するパッケージが、LinearMaps.jlというパッケージです。
https://github.com/Jutho/LinearMaps.jl
これを使えば、線形な演算子Aが与えられているときにAvを計算するというようなより抽象的な演算も行う事ができます。
これについては、
「JuliaでLinearMaps.jlを使ってみる：行列を定義せずに線形演算」
https://qiita.com/cometscome_phys/items/9f9cbd4f2ffa75e09f5e
においてまとめました。

しかし、これではまだ、複雑な問題を解くためには不十分かもしれません。
なぜならば、IterativeSolvers.jlにおいて線形演算子を作用させるものは、ベクトルでなければならないからです。つまり、自前で作ったtypeをそのまま素朴にIterativeSolvers.jlに突っ込むことはできません。
そこで、「自分で作ったtypeをAbstractVectorのsubtypeとみなす」ことで、これを実現させます。

参考url
https://docs.julialang.org/en/v1/manual/interfaces/#man-interface-array-1

例

このようなtypeを作ったとしましょう。

    mutable struct Fermions 
        fields::Array{ComplexF64,1}
        N::Int64
    end

このFermionsタイプには、fieldsというものがあり、

a = Fermions(rand(5),5)
println(a.fields[2])

のようにすればfieldsというフィールドの値を取り出すことができます。
しかし、このtypeはそのままIterativeSolvers.jlに突っ込むことはできません。なぜなら、ベクトルではないからです。

AbstractVectorのsubtype化

そこで、このtypeをAbstractVectorのsubtypeとみなしてしまいましょう。
それによって、引数がAbstractVectorである関数は全て使えるようになります。AbstractVectorとは、一次元配列のことを意味しています。一次元配列であれば、A[1]のような形で要素を取り出すことができなければなりません。そこで、

    mutable struct Fermions <: AbstractVector{ComplexF64}
        fields::Array{ComplexF64,1}
        N::Int64
    end
    Base.size(A::Fermions) = (A.N,)
    Base.getindex(A::Fermions, i::Int) = A.fields[i]
    function Base.setindex!(A::Fermions, v, i::Int) 
        A.fields[i] = v
    end

とします。ここで、getindex(A,i)は、A[i]のことです。このように、オリジナルの型Fermionsが引数に入った場合に示す挙動を入れておけば、Fermionsをベクトルとみなすことが可能となります。

IterativeSolvers.jlで使うために、以下のように関数を定義しました。

module Fermionfields   
    using LinearAlgebra
    export Fermions

    mutable struct Fermions <: AbstractVector{ComplexF64}
        fields::Array{ComplexF64,1}
        N::Int64
    end
    Fermions(N)  = Fermions(rand(N),N)

    Base.size(A::Fermions) = (A.N,)
    Base.getindex(A::Fermions, i::Int) = A.fields[i]
    function Base.setindex!(A::Fermions, v, i::Int) 
        A.fields[i] = v
    end

    Base.similar(A::Fermions) = Fermions(similar(A.fields),A.N)
    Base.similar(A::Fermions, ::Type{<:ComplexF64}, dims::Union{Integer, AbstractUnitRange})=Fermions(similar(A,dims),dims)

end

ベクトルとして扱った時に使いそうな関数をそれぞれ定義しました。これで、引数の型にFermionsが来た時には、自動的に定義した関数が呼ばれます。これが多重ディスパッチですね。

連立方程式を解く

このFermions型を使って、連立方程式を解いてみましょう。
行列は以前の記事と同じようにLinearMaps.jlを使います。
全体のコードは

module Fermionfields   
    using LinearAlgebra
    export Fermions

    mutable struct Fermions <: AbstractVector{ComplexF64}
        fields::Array{ComplexF64,1}
        N::Int64
    end
    Fermions(N)  = Fermions(rand(N),N)

    Base.size(A::Fermions) = (A.N,)
    Base.getindex(A::Fermions, i::Int) = A.fields[i]
    function Base.setindex!(A::Fermions, v, i::Int) 
        A.fields[i] = v
    end

    Base.similar(A::Fermions) = Fermions(similar(A.fields),A.N)
    Base.similar(A::Fermions, ::Type{<:ComplexF64}, dims::Union{Integer, AbstractUnitRange})=Fermions(similar(A,dims),dims)

    LinearAlgebra.norm(A::Fermions) = norm(A.fields)
end

function set_diff(v) 
    function calc_diff!(y::AbstractVector, x::AbstractVector)
        n = length(x)
        length(y) == n || throw(DimensionMismatch())
        for i=1:n
            dx = 1
            jp = i+dx
            jp += ifelse(jp > n,-n,0) #+1方向
            dx = -1
            jm = i+dx
            jm += ifelse(jm < 1,n,0) #-1方向
            y[i] = v*(x[jp]+x[jm])+0.1*x[i]
        end
        return y
    end
    (y,x) -> calc_diff!(y,x)
end

using .Fermionfields
using LinearAlgebra
using Random
Random.seed!(123)
n = 24
phi = Fermions(n)


A = set_diff(-1.0)
using LinearMaps
using IterativeSolvers
tol = 1e-12
D = LinearMap(A, n; ismutating=true,issymmetric=true)
a = rand(2)
println(zero(phi))

phi = phi ./norm(phi)
println(norm(phi))
println(typeof(phi))
x = cg(D,phi,tol=tol)
println(norm(x))
println(norm(D*x))

println(norm(D*x-phi))

となります。