Zenn
🗂

[Julia] CartesianIndices で、多次元配列処理をエレガントに操ろう

2022/12/24に公開
Julia
配列
julialang
cartesianindex
tech

Julia Advent Calendar 2022 20日目ですが、クリスマスイブの日が暮れる前に公開できた（はずです）。

Cartesian indicesは、多次元配列の要素アクセスを行う方法の一つです。
これと、そのイテレータ CartesianIndices を用いると、多次元配列の一般的なアルゴリズムを極めてエレガントに操れます。ただし、Julia マニュアルの対応する項目を読んだだけでは、以下に紹介するようなエレガントなコードは、なかなか思いつかないでしょう。
以下の記事は、Tim Holy先生による Julia blogger記事"Multidimensional algorithms and iteration"の翻案です。翻案ということで、原文とは説明の順番を変え、私なりに膨らませた解説を行ってみます。
いまは優秀な翻訳ツールもありますから、原記事もぜひ読んでみるとよいでしょう。

CartesianIndex と整数とを合体する

CartesianIndex は、多次元配列の各軸の添字をまとめたオブジェクトです。
例えば、4次元配列 A に対する要素アクセス A[i,j,k,l] を A[CartesianIndex(i,j,k,l)] と書けます。
CartesianIndex と整数を混ぜることもできます。例えば、A[CartesianIndex(i,j),k,CartesianIndex(l))] と書くこともできます。
CartesianIndex のヘルプにも "One can freely mix integer and CartesianIndex indices ..." という下りで説明されています（が、簡潔すぎて分かりませんよね）。

help?> CartesianIndex
  (中略)
  One can freely mix integer and CartesianIndex indices; for example, `A[Ipre,i, Ipost]` (where `Ipre` and `Ipost` are CartesianIndex indices and `i` is an `Int`) can be a useful expression when writing algorithms that work along a single dimension of an array of arbitrary dimensionality.

複数の CartesianIndex と整数とを合体できます。空の CartesianIndex() は単純に無視されますから、 A[CartesianIndex(),CartesianIndex(i,j),k,CartesianIndex(l))] と書くこともできます。空の CartesianIndex() は、次の例で用いられます。

例1. 多次元配列の指数平滑法

(原文の最後のコード例です)

多次元配列の特定の次元に着目した操作に、CartesianIndex を使ってみましょう。

例として、指数平滑法 （Exponential Smoothing) を実装しましょう。
1次元配列 x[i] の指数平滑 s[i] は、以下のように計算されます。重み \alpha を付けた移動平均 (moving average) ですね。

\begin{cases} s \left\lbrack 1 \right\rbrack & = x\left\lbrack 1 \right\rbrack & \\ s \left\lbrack i \right\rbrack & = \alpha x\left\lbrack i \right\rbrack + (1-\alpha) s\left\lbrack i-1 \right\rbrack & \text{for}\; i \ge 2 \end{cases}

これを計算するJuliaコード expfilt1! は、ほぼ数式通りです。引数を書き換える関数の名前は、末尾に ! を付けるのが、Juliaの命名法のお約束です。

function expfilt1!(s, x, α)
    0 < α <= 1 || error("α must be between 0 and 1")
    s[1] = x[1]
    for i = 2:length(x)
        s[i] = α*x[i] + (1-α)*s[i-1]
    end
    s
end

では、多次元配列 x\left\lbrack\right\rbrack の、第dim軸に着目して指数平滑 s_i を計算しましょう。

\begin{cases} s\left\lbrack\ldots,1,\ldots\right\rbrack & = x\left\lbrack\ldots,1,\ldots\right\rbrack & \\ s\left\lbrack\ldots,i,\ldots\right\rbrack & = \alpha x\left\lbrack\ldots,i,\ldots\right\rbrack + (1-\alpha) s\left\lbrack\ldots,i-1,\ldots\right\rbrack & \text{for}\; i \ge 2 \end{cases}

指数平滑を実行する関数 _expfilt! は、以下のように書けます。
着目する軸を指定する引数として、次の3つを受け取ることにします。

  • x の第1軸から(dim-1)軸までの全ての添字を巡回するイテレータ Rpre
  • 第dim軸の大きさ n （添字は 1 から始まると仮定します）
  • 第(dim+1)軸から最後の軸までの全ての添字を巡回するイテレータ Rpost 
function _expfilt!(s, x, α, Rpre, n, Rpost)
    for Ipost in Rpost
        for Ipre in Rpre
            s[Ipre, 1, Ipost] = x[Ipre, 1, Ipost]
        end
        for i = 2:n
            for Ipre in Rpre
                s[Ipre, i, Ipost] = α*x[Ipre, i, Ipost] + (1-α)*s[Ipre, i-1, Ipost]
            end
        end
    end
    s
end

第dim軸に対する繰り返しは、1次元配列のコードの自然な拡張になっています。
第dim軸の前後の添字を、イテレータ RpreとRpost を用いて全て巡回します。つまり、配列要素を [Ipre, i, Ipost] のようにアクセスするわけです。ここで、IpreはRpreの要素 （1軸からdim-1軸の添字） 、iは第n軸の添字、Ipostは、Rpostの要素 （dim+1軸からndim軸の添字）です。

Rpre のループを最内側で回すのは、第1軸がもっとも早く動くように多次元配列が配置されるからですね（原文 Cache efficiencyを参照してください）。

では、_expfilt! を呼び出す関数 expfiltdim を定義しましょう。
引数として、着目する軸 dim (整数)を受け取ります。

size(x) は、xの各軸の寸法のタプルを返しますから、

  • これを [1:dim-1] の範囲で切り出してCartesianIndicesに渡し、イテレータ Ipreを、
  • これを [dim+1:end] の範囲で切り出して、同様にイテレータ Ipost を

作るわけです。
dim=1 の場合は Rpre が空のイテレータ、dim=ndims(x) (xの最大次元数)の場合は Rpost が空のイテレータになりますが、例外は発生しません。

function expfiltdim(x, dim::Integer, α)
    s = similar(x)
    Rpre = CartesianIndices(size(x)[1:dim-1])
    Rpost = CartesianIndices(size(x)[dim+1:end])
    _expfilt!(s, x, α, Rpre, size(x, dim), Rpost)
end

合わせて、20行に満たないコードで、着目する軸を指定するだけで計算するプログラムを実装できました。

CartesianIndex のようなツールがなければ、多次元配列を1次元配列とみなして、デカルト添字と線形添字を写像するような関数を書いていたところです。私自身 FORTRAN ではよく書いていたのですが、もう書きたくありません．

CartesianIndex に対する演算

以下のコードで用いる CartesianIndex の演算をいくつか紹介しましょう。

以下、CartesianIndex のオブジェクト I を作っておきましょう。(I には、線形恒等写像 LinearAlgebra.UniformScalingの定義もありますので、using LinearAlgebra しているときには、使わない方がよいかもしれません）

julia> I = CartesianIndex(3,4,5)
CartesianIndex(3, 4, 5)

Tuple に作用させると、タプルが得られます。

julia> Tuple(I)
(3, 4, 5)

要素をスカラー倍できます。

julia> 2I
CartesianIndex(6, 8, 10)

oneunit(x::T) は T(one(x)) と同じです。oneunit に CartesianIndex を渡すと、各軸が 1 になった CartesisnIndex が得られます。

julia> I1 = oneunit(I)
CartesianIndex(1, 1, 1)

要素同士で加減できます。
上の I1 を加減すると、各軸の添字を同時にインクリメント・デクリメントできますね。

julia> I+I1
CartesianIndex(4, 5, 6)

julia> I-I1
CartesianIndex(2, 3, 4)

複数の CaretesinIndex を max, min 関数に与えると、各軸の最大または最小の添字からなる CartesianIndex が得られます。

julia> max(CartesianIndex(1,2,3),CartesianIndex(4,3,2))
CartesianIndex(4, 3, 3)

julia> min(CartesianIndex(1,2,3),CartesianIndex(4,3,2))
CartesianIndex(1, 2, 2)

例2. ボックスカー平均

(原文の最初のコード例です)

Boxcar average とは、配列の各要素に着目して、各軸の添字を \pm 1 加減した添字でアクセスできる全ての要素を対象とした平均値を計算するものです。存在しない配列要素は、除外します。
例えば、3次元配列 A の要素 A[i,j,k] に着目すると、A[i-1,j-1,k-1] からA[i+1,j+1,k+1] までの全ての要素を対象に平均値を計算するわけです。各軸の添字が、1 を下回ったり、各軸の寸法 size(A,d) を越えたりする要素は存在しないので、除外します。

単純に、以下のような雛形を、思いつきます（一部は擬似コード）。

function boxcar_mediocre1(A::AbstractArray)
    out = similar(A)
    R = CartesianIndices(A)
    Ifirst = first(R)
    I1 = oneunit(Ifirst)
    for I in R
        n, s = 0, zero(eltype(out))
        for J in I-I1:I+I1
            if (`J` が存在する範囲なら)
                s += A[J]
                n += 1
            end
        end
        out[I] = s/n
    end
    out
end
  • out は、boxcar average の計算結果を入れる配列です
  • R は、配列A 全体の要素を巡回するイテレータです
  • first(R) は、イテレータ Rの最初の要素の添字を表す CartesianIndex です
  • I-I1:I+I1 は、I を中心に各軸の添字を 1 だけ加減した範囲のイテレータです

擬似コードとしていた「添字 J について、各軸の添字が有効な範囲にあることの検査」を書き加えましょう。

function boxcar_mediocre2(A::AbstractArray)
    out = similar(A)
    R = CartesianIndices(A)
    Ifirst = first(R)
    I1 = oneunit(Ifirst)
    for I in R
        n, s = 0, zero(eltype(out))
        for J in I-I1:I+I1
            # `J` が存在する範囲なら
            if all( 1 <= Tuple(J)[d] <= size(A,d) for d in 1:ndims(A) ) 
                s += A[J]
                n += 1
            end
        end
        out[I] = s/n
    end
    out
end

上のプログラムでちゃんと動きますが、内側ループの中で、条件判断を毎回行うはかっこ悪いですね。この条件判断は、各軸の添字が始点・終点に達する場合だけ成立しません。すなわち、多くの場合、この条件判断は成り立ちますから、検査しないに越したことはありません。

各軸の添字が端に達しない多次元添字の集合だけで、イテレータを作ることができます。Tim Holy先生の原記事は、以下のようなエレガントなコードを示しています。

function boxcar_elegant(A::AbstractArray)
    out = similar(A)
    R = CartesianIndices(A)
    Ifirst, Ilast = first(R), last(R)
    I1 = oneunit(Ifirst)
    for I in R
        n, s = 0, zero(eltype(out))
        for J in max(Ifirst, I-I1):min(Ilast, I+I1)
            s += A[J]
            n += 1
        end
        out[I] = s/n
    end
    out
end
  • last(R) は、Rの最後の要素の添字を表します
  • max(Ifirst, I-I1) は、着目した範囲で、各軸の添字の最小値を表します
  • min(Ilast, I+I1)は、着目した範囲で、各軸の添字の最大値を表します

つまり、max(Ifirst, I-I1):min(Ilast, I+I1) というイテレータは、範囲内の添字だけを巡回します。範囲の検査を行う必要はないので、実行性能も上がりますね。

Juliaの配列の添字は、1 以外から開始することも可能です（→ [Julia 0.5] 配列添字を0から数える）が、そのような場合でも、上のコードは修正なく動作します。

例3. 多次元配列の reduce 演算

(原文の2番目のコード例です)

多次元配列の特定の次元に沿った和を求める関数を書いてみましょう。

元データは多次元配列 A にあります。結果を格納する多次元配列 B は、和を求める軸のみ寸法を 1 とし、それ以外の軸の大きさは A と同じにしておきます。

b\left\lbrack\ldots,1,\ldots\right\rbrack = \sum_{i} a\left\lbrack\ldots,i,\ldots\right\rbrack

例えば、寸法 [l, m, n] の配列 A の第2軸に沿った和をとる場合は、寸法 [l, 1, n] の配列Bを用意しておくわけです。

B の要素は 0 に初期化されているとします。
どのように書きましょうか？

例1. で示した CartesianIndex と整数の合成による配列アクセスが考えられます。

つまり、B の寸法が 1 となる軸 dim を見つけて、3つのイテレータ

  • A の第1軸から(dim-1)軸までの全ての添字を巡回する Rpre
  • 第dim軸の全ての添字を巡回する Rdim
  • 第(d+1)軸から最後の軸までの全ての添字を巡回する Rpost

を作って、それぞれで巡回するわけです。

function sumalongdims_mediocre!(B, A)
    dim = findfirst(isequal(1), size(B))
    Rpre =  CartesianIndices(size(A)[1:dim-1])
    Rdim =  axes(A,dim)
    Rpost = CartesianIndices(size(A)[dim+1:end])

    for Ipost in Rpost
        for dim in Rdim 
            for Ipre in Rpre
                B[Ipre,1,Ipost] += A[Ipre,dim,Ipost] 
            end
        end
    end
    B
end

上のコードで正しく動くのですが、3重ループを作ったり大げさですね。

さて、Tim Holy先生のエレガントなコードは、こんなに簡潔です。

function sumalongdims!(B, A)
    fill!(B, 0)
    Bmax = last(CartesianIndices(B))
    for I in CartesianIndices(A)
        B[min(Bmax,I)] += A[I]
    end
    B
end

I は、Aの各要素を表すデカルト添字です。Bmax は、B のデカルト添字のうち、各軸の最大寸法からなります。和を取るべき軸では Bmax が 1 ですから、min(Bmax,I) は、その軸の添字だけ 1 にした添字です。

例えば、Aの寸法が [l,m,n]、Bの寸法が [l,1,n] の場合、B[min(Bmax,I)] += A[I] は、B[i,1,k] += A[i,j,k]となるわけです。

この形式であれば、複数の軸に対して輪を取るように拡張しても、同じコードで動作します。

b\left\lbrack\ldots,1,\ldots,1,\ldots\right\rbrack = \sum_{i,j} a\left\lbrack\ldots,i,\ldots,j,\ldots\right\rbrack

次に、sumalongdims! を呼び出す関数を定義しましょう。
和を取るべき軸をベクトル dims で与えて、結果を格納する多次元配列 B を確保しましょう。

function sumalongdims_noninferrable(A, dims)
    sz = [size(A)...]
    sz[[dims...]] .= 1
    B = Array{eltype(A)}(undef, sz...)
    sumalongdims!(B, A)
end
  • sz に、A の寸法を用意します
  • sz の dims で示される要素を 1 に置換します
  • B = Array{eltype(A)}(undef, sz...) で B を確保します。

上のコードは動くのですが、B の寸法のヒントがないので、B の型が正しく推論できません。つまり、実行性能が悪化する恐れがあります。
実際 @code_warntype を使って、局所関数 B の型を調べましょう。

julia> @code_warntype sumalongdims_noninferrable(rand(5,4,3),[2])
MethodInstance for sumalongdims_noninferrable(::Array{Float64, 3}, ::Vector{Int64})
  from sumalongdims_noninferrable(A, dims) in Main
Arguments
  #self#::Core.Const(sumalongdims_noninferrable)
  A::Array{Float64, 3}
  dims::Vector{Int64}
Locals
  B::Any　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　# <== 注目
  sz::Vector{Int64}
(以下、略)

Locals の次の行は、B が Any 型であることを示します。

これを解決する一つの案は、B の寸法・型を明示することです。

function sumalongdims(A, dims)
    sz = [size(A)...]
    sz[[dims...]] .= 1
    B = Array{eltype(A)}(undef, sz...)::Array{eltype(A),ndims(A)}
    sumalongdims!(B, A)
end

@code_warntype を使って確認しましょう。

julia> @code_warntype sumalongdims(rand(5,4,3),[2])
MethodInstance for sumalongdims(::Array{Float64, 3}, ::Vector{Int64})
  from sumalongdims(A, dims) in Main
Arguments
  #self#::Core.Const(sumalongdims)
  A::Array{Float64, 3}
  dims::Vector{Int64}
Locals
  B::Array{Float64, 3}    # <== 注目
  sz::Vector{Int64}
(以下、略)

B の型は Array{Float64, 3} になりました。

あるいは、次のコードのように、sz として Aの次元を持った ntuple を用意すれば、B の型も推論可能となります。@code_warntype を使った確認は省略しましょう。

function sumalongdims(A, dims)
    sz = ntuple(i->i ∈ dims ? 1 : size(A, i), Val(ndims(A)))
    B = Array{eltype(A)}(undef, sz...)
    sumalongdims!(B, A)
end

関数障壁を設ける

例1と例3では、実際の計算を行う関数( _expfilt! や sumalongdims! )と、結果を格納する配列を確保したりする前処理の関数 (expfiltdim や sumalongdims) を分けています。
これらをまとめて一つの関数にすることもできますが、分けたままにしておくのが Julia のパターンの一つで、「関数障壁を設ける (introduction of function barriers)」と呼ばれています。→ Separate kernel functions (aka, function barriers)

例3であれば、前処理関数 sumalongdims は、型不安定になりえます（A の型に応じて、B や戻り値の型が変わります）。実際の計算を行う関数 sumalongdims! は、引数 Bの型が不明でも、返り値の型が決まります。例1と例3では、前者の末尾で後者を呼び出しているので、インライン展開されたりします。このように、実際の計算を行う関数が高速に実行されることが期待できます。

まとめ

ここまで、たどたどしい文章でしたが、次のようなことを説明しました。

  • CartesianIndex と整数を合体する要素アクセス
  • CartesianIndex に対する演算: oneunit, min, max
  • CartesianIndex 同士の加減算: +, -
  • 一部の軸に対する CartesianIndices の生成
  • @code_warntype を用いた変数の型の調査
  • 「関数障壁を設ける」パターン

各軸の添字個別ではなく、デカルト添字ひとまとめに対する操作を考えるとエレガントなコードにたどり着けそうですね。Julia愛好家のみなさんの参考になれば幸いです。

Discussion