PythonからFortranサブルーチンを呼び出す時の多次元配列のアクセス順（ctypeslibとf2py）

はじめに

Pythonで多次元配列に対してfor文を回して複雑な処理をしたい事がありますが（大気海洋分野ではデータが多次元配列で表されることが多いです）、そうした場合しばしば実行時間が長くなり困ることがあります。これを解決する高速化方法は様々ですが、流体計算でよく使われるFortranでサブルーチンを書いて、それをPythonから呼び出すことは一つの解決方法です。PythonからFortranのサブルーチンを呼び出す方法には、numpy.ctypeslibを使う方法とf2pyを使う方法の大きく2つがあります。両者で多次元配列のアクセス順が異なりますが、日本語でまとめられているページが見つからなかったので、記事を書きました。

実行環境は以下の通りです。

• macOS 10.15.7 Catalina
• Python 3.7.7 (conda 4.9.0)
• NumPy 1.19.2
• f2py version 2
• gfortran (gcc) 10.2.0 (Homebrew)

PythonとFortranでの多次元配列の扱いの違い

基礎的な話なので、既知の場合はこの節を読み飛ばして大丈夫です。

プログラミング言語で多次元配列が扱われますが、メモリ上では1次元方向に並べて配置されます。その際、どのように格納されるかは言語によって異なります。以下のコードでは、メモリ上に配置された順番で配列の要素に値0を書き込んでいます。

loop.py

import numpy as np

nk = 2
nj = 3
ni = 4
array = np.empty((nk, nj, ni), dtype=np.float32, order="C")
for k in range(nk):
    for j in range(nj):
        for i in range(ni):
            array[k, j, i] = 0    

loop.f90

integer(4), parameter :: nk = 2, nj = 3, ni = 4
real(4), dimension(ni, nj, nk) :: array

do k = 1, nk
    do j = 1, nj
        do i = 1, ni
            array(i, j, k) = 0
        end do
    end do
end do

Python（正確にはnumpyの配列のorder="C"の場合）では[k, j, i]、Fortranでは(i, j, k)の順になっている点が違います（それぞれ、行優先と列優先と呼ばれたりします）。メモリ上の格納順にアクセスするか飛び飛びにアクセスするかは、特に配列のサイズが大きくなると実行時間に影響します。

numpy.ctypeslibを使う場合

例として、numpyの2次元配列をFortranにサブルーチンに渡し、それを2倍したものを返すコードを書いてみます（3次元以上の場合についても同様です）。コードの詳細についてはここでは解説しないので、知りたい場合は例えば Fortran, C言語 との連携 など他の記事を見てください。

doubling1.py

from ctypes import POINTER, c_int32, c_void_p
import numpy as np

libDoubling = np.ctypeslib.load_library("doubling1.so", ".")
libDoubling.doubling.argtypes = [
    np.ctypeslib.ndpointer(dtype=np.float32),
    np.ctypeslib.ndpointer(dtype=np.float32),
    POINTER(c_int32),
    POINTER(c_int32),
]
libDoubling.doubling.restype = c_void_p

nj = 2
ni = 3

vin = np.empty((nj, ni), dtype=np.float32, order="C")
vin[0, 0] = 1
vin[0, 1] = 2
vin[0, 2] = 3
vin[1, 0] = 4
vin[1, 1] = 5
vin[1, 2] = 6

vout = np.empty((nj, ni), dtype=np.float32, order="C")

_nj = POINTER((c_int32))(c_int32(nj))
_ni = POINTER((c_int32))(c_int32(ni))

libDoubling.doubling(vin, vout, _ni, _nj)

print(vin)
# [[1. 2. 3.]
#  [4. 5. 6.]]
print(vout)
# [[ 2.  4.  6.]
#  [ 8. 10. 12.]]

doubling1.f90

! gfortran -Wall -cpp -fPIC -shared doubling1.f90 -o doubling1.so
subroutine doubling(vin, vout, ni, nj) bind(C)
    use, intrinsic :: ISO_C_BINDING
    implicit none
    integer(c_int32_t), intent(in) :: ni, nj
    real(c_float), dimension(ni, nj), intent(in) :: vin
    real(c_float), dimension(ni, nj), intent(out) :: vout
    integer(c_int32_t) :: i, j

    do j = 1, nj
        do i = 1, ni
            vout(i, j) = vin(i, j) * 2
        end do
    end do
end subroutine doubling

Fortranのサブルーチン（共有ライブラリ）の引数にはポインタが渡されます。PythonとFortranでインデックスが[j, i]と(i, j)と逆になりますが、メモリ上にデータがどういう順番で格納されているかを考えれば、素直に理解できるはずです。

f2pyを使う場合

次にf2pyを使い、先ほどと同様のコードを書いてみます。これについても詳細は解説しませんが、例えば f2pyでnumpyの配列を扱う といった記事が参考になると思います。

doubling2.f90

! f2py -c --fcompiler=gfortran -m doubling2 doubling2.f90
subroutine doubling2(vin, vout, ni, nj)
    implicit none
    integer(4), intent(in) :: ni, nj
    real(4), dimension(nj, ni), intent(in) :: vin
    real(4), dimension(nj, ni), intent(out) :: vout
    integer(4) :: i, j

    do j = 1, nj
        do i = 1, ni
            vout(j, i) = vin(j, i) * 2
        end do
    end do
end subroutine doubling2

doubling2.py

import numpy as np
import doubling2

nj = 2
ni = 3
vin = np.empty((nj, ni), dtype=np.float32, order="C")
vin[0, 0] = 1
vin[0, 1] = 2
vin[0, 2] = 3
vin[1, 0] = 4
vin[1, 1] = 5
vin[1, 2] = 6

vout = np.empty((nj, ni), dtype=np.float32, order="C")

vout = doubling2.doubling2(vin, ni, nj)
print(vin)
# [[1. 2. 3.]
#  [4. 5. 6.]]
print(vout)
# [[ 2.  4.  6.]
#  [ 8. 10. 12.]]

f2pyを使う場合は、Pythonの[j, i]はそのままFortranでも(j, i)とすればOKです。コード上の見た目が同じになるメリットがありますが、doubling2.f90の中のdo文では、配列のアクセス順がメモリ上のデータ格納順と対応しなくなるというデメリットがあります。

これを解決するにはどうすれば良いでしょうか？　つまりFortran側のコードをdoubling1.f90と同じく、メモリ上のデータ格納順にアクセスするようにしたいわけです。

doubling3.f90

! f2py -c --fcompiler=gfortran -m doubling3 doubling3.f90
subroutine doubling3(vin, vout, ni, nj)
    implicit none
    integer(4), intent(in) :: ni, nj
    real(4), dimension(ni, nj), intent(in) :: vin
    real(4), dimension(ni, nj), intent(out) :: vout
    integer(4) :: i, j

    do j = 1, nj
        do i = 1, ni
            vout(i, j) = vin(i, j) * 2
        end do
    end do
end subroutine doubling3

numpyの配列のorder引数をデフォルトの"C"ではなく"F"にすることで、メモリ上の格納順を逆にすることが出来ることを思いつくかもしれません。しかし、これではこの問題は解決しません。すなわち、doubling2.pyでorder="C"をorder="F"に変えたとしても、対応するFortranコードは依然としてdoubling2.f90であり、doubling3.f90を使うとエラーになるということです。

解決策は、f2py — prevent array reordering - Stack Overflowにあるとおり、T属性を用いて配列を転置することです。

doubling3.py

import numpy as np
import doubling3

nj = 2
ni = 3
vin = np.empty((nj, ni), dtype=np.float32, order="C")
vin[0, 0] = 1
vin[0, 1] = 2
vin[0, 2] = 3
vin[1, 0] = 4
vin[1, 1] = 5
vin[1, 2] = 6

_vout = np.empty((nj, ni), dtype=np.float32, order="C")

_vout = doubling3.doubling3(vin.T, ni, nj)
vout = _vout.T
print(vin)
# [[1. 2. 3.]
#  [4. 5. 6.]]
print(vout)
# [[ 2.  4.  6.]
#  [ 8. 10. 12.]]

例えば「netCDF4ライブラリを用いてGPVを読み込み、それをFortranのサブルーチンで効率良く処理した後、matplotlib+cartopyなどで描画する」といった用途の場合、この方法がよいと思います。"お節介"を回避するために転置するのは面倒ですが、それを含めてもnumpy.ctypeslibを使う場合と比べると記述量は少なく、f2pyのほうが使い勝手がよいように思います。