numpy(行列計算)高速化のためのnumbaの型指定について

Question

前提・実現したいこと

numbaを用いて行列計算を高速にしたいと考えています。
numba.jitにおいて型指定を行うとより早くなるそうなので試しています。

発生している問題・エラーメッセージ

型推定に失敗？しているようです。

Compilation is falling back to object mode WITH looplifting enabled because Function "compute" failed type inference due to: Invalid use of Function(<built-in function zeros>) with argument(s) of type(s): ((Literal[int](10), array(float64, 1d, A)), dtype=Literal[str](f8))
 * parameterized
In definition 0:
    All templates rejected with literals.
In definition 1:
    All templates rejected without literals.
This error is usually caused by passing an argument of a type that is unsupported by the named function.
[1] During: resolving callee type: Function(<built-in function zeros>)
[2] During: typing of call at <ipython-input-45-5c375a54c6a4> (6)


File "<ipython-input-45-5c375a54c6a4>", line 6:
def compute(T):
    <source elided>
    dt = 0.001
    X = np.zeros((N,T),dtype='f8')
    ^

  @numba.jit('f8(f8[:])')
<ipython-input-45-5c375a54c6a4>:3: NumbaWarning: 
Compilation is falling back to object mode WITHOUT looplifting enabled because Function "compute" failed type inference due to: cannot determine Numba type of <class 'numba.dispatcher.LiftedLoop'>

File "<ipython-input-45-5c375a54c6a4>", line 12:
def compute(T):
    <source elided>
    
    for t in range(0,T-1,1):
    ^

  @numba.jit('f8(f8[:])')
C:\Users\takuy\Anaconda3\envs\keras_env\lib\site-packages\numba\object_mode_passes.py:178: NumbaWarning: Function "compute" was compiled in object mode without forceobj=True, but has lifted loops.

File "<ipython-input-45-5c375a54c6a4>", line 4:
@numba.jit('f8(f8[:])')
def compute(T):
^

  state.func_ir.loc))
C:\Users\takuy\Anaconda3\envs\keras_env\lib\site-packages\numba\object_mode_passes.py:187: NumbaDeprecationWarning: 
Fall-back from the nopython compilation path to the object mode compilation path has been detected, this is deprecated behaviour.

For more information visit http://numba.pydata.org/numba-doc/latest/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit

File "<ipython-input-45-5c375a54c6a4>", line 4:
@numba.jit('f8(f8[:])')
def compute(T):
^

  warnings.warn(errors.NumbaDeprecationWarning(msg, state.func_ir.loc))
<ipython-input-45-5c375a54c6a4>:16: NumbaPerformanceWarning: '@' is faster on contiguous arrays, called on (array(float64, 2d, C), array(float64, 1d, A))
  X_2[:,t+1] = -M_inv@K@X[:,t+1]

該当のソースコード

import numpy as np
from scipy import linalg

m=1
g=9.8
k=100
N=10

#剛性マトリックスの作成
K=np.eye(N,dtype=float)*2-np.eye(N,k=1)-np.eye(N,k=-1)
K[0,0]=2;K[N-1,N-1]=1
M=np.eye(N,dtype=float)
K_inv=np.linalg.inv(K)
M_inv=np.linalg.inv(M)


import numba

#問題の部分
@numba.jit('f8(f8[:])')
def compute(T):
    dt = 0.001
    X = np.zeros((N,T),dtype='f8')
    V = np.zeros((N,T),dtype='f8')
    V[-1,0] = 1.0
    X_2 = np.zeros((N,T),dtype='f8')
    print(V)
    
    for t in range(0,T-1,1):

        X[:,t+1] = X[:,t] + dt*V[:,t]
        V[:,t+1] = V[:,t] + dt*X_2[:,t]
        X_2[:,t+1] = -M_inv@K@X[:,t+1]
        
    
    return (X)

%time X_=compute(int(1e4))

試したこと

すべてのarrayををfloat64に型指定しましたが,推定に失敗するようです

補足情報（FW/ツールのバージョンなど）

jupyter notebook　使用

Answer 1

全部を説明できるほど私も完全には理解していないため端的に書きます。

いくつかWarningが出ていますが、最初のはnp.zeros()の引数が原因です。
numbaはnumpyの関数を独自実装のコードと置き換えてコンパイルするのですが、numba実装のzerosは文字列でのdtypeを受け付けないようです。(ドキュメントは見つかりませんでしたがWarningからの推測)

参考: https://numba.pydata.org/numba-doc/dev/reference/numpysupported.html

2番目と3番目のWarningは型情報の間違いによるものです。
compute()は実装をみるとintをひとつ引数にとって、float64の2D-arrayを返す関数になっています。
ところが、デコレータで与えている型情報はf8(f8[:])となっており、これはfloat64の1D-arrayを引数にとって、float64のスカラを返す関数という意味になります。
そのためコンパイラが混乱しているようです。
正しくはf8[:,:](i8)になるのではないでしょうか。

その次のWarningはDeprecationWarningで、上記の原因でコンパイラがnopythonモード(はやい)でのコンパイルに失敗して、objectモード(はやくない)でのコンパイルに切り替えているのですが、この、nopythonモードでのコンパイルに失敗した場合にobjectモードに切り替わる機能が将来的にはなくなりますよという警告です。

一番最後のはパフォーマンスに関する警告で@演算子は連続領域に対する演算のほうが高速ですということを教えてくれているのですが、とりあえず今回は無視するしかなさそうです。

下記にすれば、一番最後のWarning以外は消すことができます。

@numba.jit('f8[:,:](i8)')
def compute(T):
    dt = 0.001
    X = np.zeros((N,T),dtype=np.float64)
    V = np.zeros((N,T),dtype=np.float64)
    V[-1,0] = 1.0
    X_2 = np.zeros((N,T),dtype=np.float64)
    print(V)

    for t in range(0,T-1,1):

        X[:,t+1] = X[:,t] + dt*V[:,t]
        V[:,t+1] = V[:,t] + dt*X_2[:,t]
        X_2[:,t+1] = -M_inv@K@X[:,t+1]


    return (X)

ただし、元のままでもint(1e4)を引数にしての最初の実行で、numbaがいい感じに型を推定してちゃんとコンパイルして動作しているはずです。