はじめに

最近Qiitaでpython関係の記事を検索して読んでいるのですが、割とよくnumpyやscipyの良さを活かしきれていないコードを見かけます。おそらく私と同じくコンパイラ型言語のFortran, C, C++を使ってプログラミングを始めたために、インタープリタ型言語のPythonを使っても同じようなコードを書いてしまっているのだろうと思いますが、非常にもったいないです。そこで数値計算に特化して、どうやって速いコードを書くかを実例を示しながら説明していこうと思います。

計算を高速に実行するためのポイント

forループを書かない

一番大事なことはforループを可能な限り避けることです。配列に対して計算を行うとき、Fortran, C, C++ではforループを使って各要素ごとに計算を行います。

const int n = 1000;
auto v = std::vector(n);
auto v1 = std::vector(n);
auto v2 = std::vector(n);

initialize(v1, v2);

for (int i = 0; i < v.size(); ++i) {
    v[i] = v1[i] * v2[i];
}

pythonでは、このようなforループを使った計算が 非常に遅いです 。その代わりとして、numpyに各種配列計算の機能・関数が用意されているので、可能な限りそれらを使うようにすることが大事です。

例えば上のベクトル要素の掛け算はnumpy.arrayを使うと簡単に表せてかつforループで書くよりも圧倒的に高速です。

# v1 and v2 are of type numpy.array
v = v1 * v2

またこのようにforループを使わないことでコードの可読性も向上します。

例えば次の記事の問題を考えてみましょう。
Pythonによる数値計算-1次元電子系のパイエルス転移
この記事ではnumpy.arrayではなくリストを用意し、forループを回して計算しています。

for i in range(step):
    y[i] = (math.log(abs((x[i] + 1)/(x[i] - 1)))/x[i]) / (math.log(abs((x[0] + 1)/(x[0] - 1)))/x[0])

この部分をnumpy.arrayを使って書き直したものがこちらになります。
（簡単のため発散を避ける部分はあえて除外しています。）

import numpy as np

def calc_chi(x):
    return np.log(np.fabs((x + 1.0)/(x - 1.0))/x)

x0 = 0.0
x1 = 2.0
n = 100
x = np.linespace(x0, x1, n)

# this line corresponds to the above for-loop
chi = calc_chi(x)/calc_chi(x0)

この問題は大した計算ではないため高速化のメリットはほとんどありませんが、コードの可読性が大きく向上しており、この点のメリットが大きいことがわかります。

一次元偏微分方程式

偏微分方程式を高速に解くには、線形な偏微分方程式を離散化すると行列の和・積の形に帰着するということを利用します。例えば一次元拡散方程式をFTCS法で解くことを考えましょう。

\frac{\partial u}{\partial t} = \alpha \frac{\partial^2 u}{\partial x^2}

時間について前進差分、空間について中心差分を取ると

u^{n+1}_i = u^n_i + r \, (u^n_{i+1} - 2u^n_i + u^n_{i-1})

となります。ここで$r=\frac{\alpha \Delta t}{ \Delta x^2}$です。これを単純に実装すると

for i in range(n):
    u_new[i] = u[i] + r*(u[i + 1] - 2.0*u[i] + u[i - 1])

となります。当然ながらこれは遅いです。

もう一度離散化した式を見てみましょう。実はこの式は行列を使って表すことができます。

\mathbf{u}^{n+1} = \mathbf{A} \mathbf{u}^n

ここで$\mathbf{u} = [u_0, u_1, \cdots, u_{N-1}]^T$、$\mathbf{A} = \mathbf{E} + r\mathbf{B}$です。$\mathbf{E}$は単位行列、$\mathbf{B}$は次式で定義される帯行列です。

\mathbf{B} =
\left[
\begin{matrix}
-2 & 1 &   &  &  \\
 1 & -2 & 1 &  &  \\
   & & \ddots & & \\
   & & 1 & -2 & 1 \\
   & & & 1 & -2
\end{matrix}
\right]

（簡単のため境界条件については無視しています。）
したがって、行列$\mathbf{A}$を一度計算してしまえば、forループを使わずに時間発展を計算することが可能となります。行列$\mathbf{A}$は疎行列であり、scipyを利用すれば疎行列とベクトルの掛け算を高速に計算することができます。

import numpy as np
from scipy.sparse import dia_matrix identity

# construct matrix A
data = np.array([np.ones(n), -2.0*np.ones(n), np.ones(n)])
offsets = np.array([-1, 0, 1])
B = dia_matrix((data, offsets), shape=(n, n))
E = identity(n)
A = E + r * B

行列計算を使って修正したものが下記コードになります。

# calculate u at new time step
u_new = A.dot(u)

追記：
numpy.diffを使うと一次元配列の隣接要素の差分を計算することができます。

import numpy as np
u = ...              # u is of type numpy.array with n elements
du1 = np.diff(u)     # du1[i] = u[i+1] - u[i], i = 0,...,n-2
du2 = np.diff(du1)   # du2[i] = du1[i+1] - du1[i] = u[i+2] - 2*u[i+1] + u[i], i = 0,...,n-3

二次元偏微分方程式

二次元偏微分方程式も、一次元偏微分方程式と同様に離散化した式が行列で表されることを利用して高速に解くことができます。

例えばこちらの記事の問題を考えてみましょう。
[Pythonによる科学・技術計算] 静電位に対する2次元ラプラス・ポアソン方程式のヤコビ法による数値解法，楕円型偏微分方程式，境界値問題
ヤコビ法を用いたときの離散化式は次のようになります。

\phi_{i,j}^{n+1} = \frac{1}{4}(\phi_{i+1,j}^n + \phi_{i-1,j}^n + \phi_{i,j+1}^n + \phi_{i,j-1}^n)

この式を素直にforループを使って実装すると、下のようになります。

for i in range(nx):
    for j in range(ny):
        if i == 0 or i == nx - 1 or j == 0 or j == ny - 1:
            # for boundary condition
            phi_new[i,j] = phi[i,j]
        else: 
            # for Jacobi method
            phi_new[i,j] = 0.25 * (phi[i + 1, j] + phi[i - 1, j] + phi[i, j + 1] + phi[i, j - 1])

毎反復ごとにこのループを回しつつ、さらにそのループの中で境界条件のための場合分けまで行っているため非常に遅くなります。（例えコンパイラ型言語だったとしても、速度が重要な場合はforループ内で場合分けするのは避けるべきです。）

さてもう一度離散化式を見てみましょう。この場合$\phi$が二次元配列で表されているため行列表現に書き換えるのは難しそうに見えます。しかし二次元配列は一次元配列に変換が可能です（$(i,j) \leftrightarrow k$とマップさせれば、$\phi_{i,j} \leftrightarrow \phi_k$となる）。よってこの式も

\mathbf{\phi}^{n+1} = \mathbf{A} \mathbf{\phi}^n

と行列で表すことが可能です。$\mathbf{A}$は疎行列です。

forループを使った上記記事のコードだと、私のPCではwhile文の部分の計算に約９秒かかりました。しかし$\phi$を一次元配列に変換し、forループ部分の計算を疎行列の掛け算に変換したコードの場合、while文の部分の計算は約0.05秒で終わりました。約180倍高速になったわけです。forループを行列計算に変換した部分を下に示します。
（リファクタリングしたスクリプトはGitHubにアップしました。）

laplace_refactored.py

# update field using Jacobi method
# A is sparse matrix
phi_new = A.dot(phi)

ライブラリを使用する

自分が行いたい処理は、大抵の場合既にライブラリとして提供されています。スクリプトを書く前にググりましょう。
（その処理について勉強するためにあえてスクリプトを書く場合はこの限りではありませんが…）

ググるコツは

その処理の一般的な名前で検索する
その処理に関連する単語を加える
日本語でヒットしなければ英語で検索する

とかでしょうか。私の場合ほとんど英語で検索をかけるため、大体は本家レファレンスかStack Overflowが引っ掛かり答えを見つけることができます。

例えばヒストグラムを計算しようとしている記事
情報処理の問題をfortranとpythonで書く！#08 0~1まで0.1区分でそれぞれの個数を求める
でコメントしたのですが、ヒストグラムを計算する関数numpy.histogramを使うとforループを使わず、わずか４行で処理が終わります。

import numpy as np
data = np.loadtxt('data.txt')
bin_edges = np.arange(11) / 10.0
hist, bin_edges = np.histogram(data, bins=bin_edges)

またJuliaの速さを体感するという記事では、常微分方程式をforループを使って解いたときのJuliaとPythonの速度比較を行っています。ただ、scipy.integrate.odeintを使えばpythonでも容易にJulia並の速度を出すことができるので、forループを使った比較は実用上あまり意味がありません。

import numpy as np
from scipy.integrate import odeint

# discretization
t = np.linspace(0.0, 4.0, 100)
# initial value
y0 = 1.0
# Differential equation: dy/dt = -y
# function definition which returns the right-hand side
def f(y, t): return -y
y = odeint(f, y0, t)

数値計算に関するライブラリを予めざっと読み込んでおき、どんな関数があるのか把握しておくのも大切です。

CythonやNumbaを使って高速化する

JITコンパイラのNumbaを使うと簡単に高速化できますが、コードによってはエラーが出て使用できないこともあります。一方、CythonはCコンパイラでコードをコンパイルしているので、実質的にC言語でコードを書いて呼び出すのと変わらない速度でコードを実行できます。

Numba
Cython

こちらの記事でNumbaによりどのくらい高速化されるかｇ比較されています。

Cythonについてはこちらの記事がおすすめです。

Pythonで書くと遅くなる関数・クラスのみをFortran, C, C++で書く

非常に複雑な処理が必要で、どうしてもforループをなくすことができない場合、Fortran, C, C++でその処理を書いてコンパイルし、pythonから呼び出すことで高速化できます。Fortran, C, C++のpython用ラッパーは下記に示す通りです。

F2PY for Fortran
- scipy.org: three ways to wrap
SWIG for C, C++
CFFI for C, C++
boost.python for C++
pybind11 for C++11

C++を使っているならpybind11をおすすめします。