f2pyを使って自作したfortranのモジュールをpythonで実装してみる（テンソル積）

概要

pythonはライブラリが豊富で可読性が高い言語であるが、計算速度の面ではfortranに劣る。そこで、pythonにfortranで書いたモジュールを実装してみよう。今回はf2pyを使ってfortranのモジュールをコンパイルする方法を紹介する。

使い方

f2pyのインストール

f2pyはpythonのライブラリであるnumpyに含まれている。よって、コマンドプロンプトやターミナルでpipを使ってnumpyをインストールしよう。

pip install numpy

コンパイル

f2pyを使ってfortranのモジュールをコンパイルしよう。

f2py -m モジュール名 -c ファイル名.f90

コンパイルが成功していたら、モジュール名.soのようなファイルが作られているはず。
モジュール名とファイル名は好きな名前を付けよう。

fortranのモジュールの中身

実装したい関数をsubroutineで書いてみよう。
今回は、テンソル（多次元配列）と行列の積の計算を行うモジュールを実装してみる。テンソルについて詳しく知りたい方はこちらの論文を参考にしてほしい。
以下のコードは3次元テンソルと行列の積（mode-n積）を定義通りに書いたもの。計算効率などは全く無視してるので参考程度に。

mode_n_product.f90

subroutine mode_n_product(x, m, mode, result, n1, n2, n3)
    use iso_c_binding

    implicit none
    integer(kind=c_int), intent(in) :: mode, n1, n2, n3
    real(kind=8), intent(in) :: x(:,:,:), m(:,:)
    real(kind=8), intent(inout) :: result(n1,n2,n3)
    integer :: i, j, k, l

    if (size(x,mode) /= size(m,2)) stop

    select case(mode)
    case(1)
        result = 0.0
        do l=1, size(m,2)
            do k=1, size(x,3)
                do j=1, size(x,2)
                    do i=1, size(m,1)
                        result(i,j,k) = result(i,j,k) + x(l,j,k) * m(i,l)
                    end do
                end do
            end do
        end do
    case(2)            
        result = 0.0
        do l=1, size(m,2)
            do k=1, size(x,3)
                do j=1, size(m,1)
                    do i=1, size(x,1)
                        result(i,j,k) = result(i,j,k) + x(i,l,k) * m(j,l)
                    end do
                end do
            end do
        end do
    case(3)
        result = 0.0
        do l=1, size(m,2)
            do k=1, size(m,1)
                do j=1, size(x,2)
                    do i=1, size(x,1)
                        result(i,j,k) = result(i,j,k) + x(i,j,l) * m(k,l)
                    end do
                end do
            end do
        end do
    case default
        print *, 'Invalid mode'
        stop
    end select
end subroutine mode_n_product

引数はintent(in)で、出力はintent(inout)またはintent(out)で定義する。
intent(inout)で出力する配列を定義するとき、result(n1,n2,n3)のように配列のサイズを指定しないとエラーになるので注意。つまり、result(:,:,:)ではダメ。

pythonの中身

さて、fortranで書いたモジュールを使ってみよう。import モジュール名で作ったモジュールをインポートしよう。

main.py

import numpy as np
import mode_n_product_module

x = np.random.rand(4, 5, 2)
m = np.random.rand(3, 2)

mode = 3

if mode == 1:
    result = np.zeros((m.shape[0], x.shape[1], x.shape[2]))
elif mode == 2:
    result = np.zeros((x.shape[0], m.shape[0], x.shape[2]))
elif mode == 3:
    result = np.zeros((x.shape[0], x.shape[1], m.shape[0]))

tensor = np.asfortranarray(x)
matrix = np.asfortranarray(m)
result = np.asfortranarray(result)

mode_n_product_module.mode_n_product(tensor, matrix, mode, result, *result.shape)

result = np.ascontiguousarray(result)

for i in range(result.shape[2]):
    for j in range(result.shape[0]):
        print(result[j, :, i])
    print()

モジュールを使って出力する配列のサイズは、result = np.zeros((m.shape[0], x.shape[1], x.shape[2]))のように事前に定義しておく必要がある。また、numpyで作った配列をfortranに渡す用の配列に直すためにnp.asfortranarray()を使う。

実行

fortranのモジュールをf2pyでコンパイルしたら、あとはpythonを実行するだけ。

python main.py

実行してみると以下のような結果が出力されるはず。

結果

[0.13324186 0.80214173 0.63156802 0.34799383 0.57245455]
[0.26886559 0.69130153 0.27121069 0.74469734 0.15505326]
[0.32695887 0.26361047 0.59345141 0.74017049 0.74162445]
[0.57436783 0.45527635 0.64633685 0.2854053  0.5319504 ]

[0.2612949  0.341772   0.27374836 0.07839377 0.18997   ]
[0.38773648 0.16981612 0.43738313 0.4340993  0.2645207 ]
[0.35061437 0.34622163 0.4684826  0.37916372 0.47211926]
[0.48303257 0.46443669 0.31738594 0.44506599 0.28343468]

[0.20930631 0.63111042 0.49928321 0.23810066 0.42284428]
[0.35108303 0.48018949 0.377778   0.64557898 0.22336776]
[0.36518347 0.32688571 0.57706297 0.61494062 0.66325127]
[0.57363416 0.49635582 0.52997944 0.38978119 0.44753513]

終わりに

今回はf2pyを使ってfortranのモジュールをpythonに実装してみた。これを応用すればテンソルや行列についての様々なモジュールを自分で作って、pythonを高速化できるはず。
何かご指摘があれば、ぜひコメントしてほしい。