【Python】Heat Kernel行列をfor文を使わずに求める

ふと、過去の記事を眺めていたら、ESFSの記事でこんな記述があるのを発見。

高速化のためには、for文で計算している部分を行列計算になおすなど、
繰り返し構文をなるべく使わない工夫が必要になります。
実際にどのように実装したかは、気が向いたら書こうと思います。

Pythonのようなスクリプト言語の場合、繰り返し構文で計算するよりも、
numpyなどによる行列演算を用いて、一括で計算した方が速くなることが多々あります。
これはMATLABでも同じことが言えますし（最近のはJIT-Acceleratorがあるので
そこまで変わらないかも）、Rなんかも同様のようです。
ということで、今回はPythonでESFSやLE、LPP等で必要になるHeat Kernel行列を
（一部強引に）繰り返し構文を使わず実装したコードを紹介します。
まず、ここでいうHeat Kernel行列は以下のようなものを指します。

$N$ 個の $d$ 次元ベクトル ${\bf x}_1 , {\bf x}_2 , \cdots , {\bf x}_N$ があり、 ${\bf x}_i$ が ${\bf x}_j$ のn近傍、または、 ${\bf x}_j$ が ${\bf x}_i$ のn近傍の場合、
$W_{ij} = e^{-\frac{\parallel {\bf x}_i - {\bf x}_j \parallel_2^2}{t}}$
それ以外の場合は、 $W_{ij} = 0$ となる $W_{ij}$ を各要素として持つ $N \times N$ の対称行列 ${\bf W}$ 。

ということで、この行列をつくるためにやらければならないことは、

各データ間の距離の2乗を計算
その距離に基づいて、各データのn近傍のサンプルを導出
近傍関係から接続行列（n近傍なら接続）を導出
データ間の距離と接続行列から、Heat Kernel行列を計算

この4つです。
これらの計算にfor文を一切使わないことで、高速化に努めました。

前提

下記のソースコードは、以下のようにインポートしていることを前提とします。

import numpy as np
import scipy as sci
import scipy.linalg

各データ間の距離の2乗を計算

ここで求めるのは次のような距離の2乗を要素として持つ行列です。
${\bf {\rm DistMat}} = \begin{equation} \begin{pmatrix} \parallel {\bf x}_1 - {\bf x}_1 \parallel_2^2 & \parallel {\bf x}_1 - {\bf x}_2 \parallel_2^2 & \cdots & \parallel {\bf x}_1 - {\bf x}_N \parallel_2^2 \\ \parallel {\bf x}_2 - {\bf x}_1 \parallel_2^2 & \parallel {\bf x}_2 - {\bf x}_2 \parallel_2^2 & \cdots & \parallel {\bf x}_2 - {\bf x}_N \parallel_2^2 \\ \vdots & \vdots & \ddots & \vdots \\ \parallel {\bf x}_N - {\bf x}_1 \parallel_2^2 & \parallel {\bf x}_N - {\bf x}_2 \parallel_2^2 & \cdots & \parallel {\bf x}_N - {\bf x}_N \parallel_2^2 \end{pmatrix} \end{equation}$
単純にループを回して、それぞれで距離を計算すれば、各データ間の距離は得られますが、
その実装方法だと $O(N^2)$ 回の繰り返しなので、 $N$ が大きくなればかなり遅くなります。
そのため、ここがHeat Kernel行列の計算の実装において、最も工夫すべき箇所だといえます。
さて、上記の行列ですが、次のように分解できます。
$\begin{eqnarray} {\bf {\rm DistMat}} &=& \begin{pmatrix} \parallel {\bf x}_1 \parallel_2^2 - 2 {\bf x}_1^T {\bf x}_1 + \parallel {\bf x}_1 \parallel_2^2 & \cdots & \parallel {\bf x}_1 \parallel_2^2 - 2 {\bf x}_1^T {\bf x}_N + \parallel {\bf x}_N \parallel_2^2 \\ \parallel {\bf x}_2 \parallel_2^2 - 2 {\bf x}_2^T {\bf x}_1 + \parallel {\bf x}_1 \parallel_2^2 & \cdots & \parallel {\bf x}_2 \parallel_2^2 - 2 {\bf x}_2^T {\bf x}_N + \parallel {\bf x}_N \parallel_2^2 \\ \vdots & \ddots & \vdots \\ \parallel {\bf x}_N \parallel_2^2 - 2 {\bf x}_N^T {\bf x}_1 + \parallel {\bf x}_1 \parallel_2^2 & \cdots & \parallel {\bf x}_N \parallel_2^2 - 2 {\bf x}_N^T {\bf x}_N + \parallel {\bf x}_N \parallel_2^2 \end{pmatrix} \\ & = & \begin{pmatrix}\parallel {\bf x}_1 \parallel_2^2 & \cdots & \parallel {\bf x}_1 \parallel_2^2 \\ \parallel {\bf x}_2 \parallel_2^2 & \cdots & \parallel {\bf x}_2 \parallel_2^2 \\ \vdots & \ddots & \vdots \\ \parallel {\bf x}_N \parallel_2^2 & \cdots & \parallel {\bf x}_N \parallel_2^2 \end{pmatrix} - 2 \begin{pmatrix} {\bf x}_1^T {\bf x}_1 & \cdots & {\bf x}_1^T {\bf x}_N \\ {\bf x}_2^T {\bf x}_1 & \cdots & {\bf x}_2^T {\bf x}_N \\ \vdots & \ddots & \vdots \\ {\bf x}_N^T {\bf x}_1 & \cdots & {\bf x}_N^T {\bf x}_N \end{pmatrix} + \begin{pmatrix}\parallel {\bf x}_1 \parallel_2^2 & \cdots & \parallel {\bf x}_N \parallel_2^2 \\ \parallel {\bf x}_1 \parallel_2^2 & \cdots & \parallel {\bf x}_N \parallel_2^2 \\ \vdots & \ddots & \vdots \\ \parallel {\bf x}_1 \parallel_2^2 & \cdots & \parallel {\bf x}_N \parallel_2^2 \end{pmatrix} \\ &=& {\bf H}^T -2{\bf X}^T {\bf X} + {\bf H} \end{eqnarray}$
ただし、
$\begin{eqnarray}{\bf H} = \begin{pmatrix}\parallel {\bf x}_1 \parallel_2^2 & \cdots & \parallel {\bf x}_N \parallel_2^2 \\ \parallel {\bf x}_1 \parallel_2^2 & \cdots & \parallel {\bf x}_N \parallel_2^2 \\ \vdots & \ddots & \vdots \\ \parallel {\bf x}_1 \parallel_2^2 & \cdots & \parallel {\bf x}_N \parallel_2^2 \end{pmatrix} & , & {\bf X} =\begin{pmatrix} {\bf x}_1 {\bf x}_2 \cdots {\bf x}_N \end{pmatrix} \end{eqnarray}$

後は、numpyのnp.diagやnp.tile（MATLABならrepmat）を利用して、
行列 $\bf{H}$ を生成すれば、繰り返し構文なしで
距離の2乗の行列を得ることができます。
サンプルの行列 $\bf{X}$ を引数にとり、
距離の2乗の行列を返す関数のソースコードは以下の通りです。

def GetDistanceMatrix(X):
    H = np.tile(np.diag(np.dot(X.T , X)) , (np.size(X , 1) , 1))
    G = np.dot(X.T , X)
    DistMat = H - 2 * G + H.T
    return DistMat

この関数に3次元1000個のデータ（3行1000列の行列）を入力したところ、
0.05秒程度で計算できるのに対し、forループを回して愚直に計算する関数の場合、
13秒程度かかります。
やはり、行列演算を用いて一括で計算した方が、格段に速いです。

距離行列に基づいて、各データのn近傍のサンプルを導出

np.argsortを用いて、各データにおける距離が小さいサンプル上位n件の
インデックスを求めるだけです。早速ソースコードです。

def GetnNearests(X , n):
    DistMat = GetDistanceMatrix(X)
    NearestsIdx = np.argsort(DistMat , axis=1)
    return NearestsIdx[: , 1:(n + 1)] , DistMat

近傍関係から接続行列（n近傍なら接続）を導出

上記のGetnNearests関数で得られるn近傍のインデックスから、
近傍関係なら1（接続）、そうでないなら0という要素で構成される行列を求めます。
ただし、ここでは対称行列がほしいので、どちらか一方から見て近傍関係に当たるなら、
このペアは接続ということにします。

def GetAdjacencyMatrix(X , n):
    (Idx , DistMat) = GetKNearests(X , n)
    AdjacencyMatrix = np.zeros((X.shape[1] , X.shape[1]) , 'bool');
    AdjacencyMatrix[np.tile(np.arange(Idx.shape[0]) ,
            (Idx.shape[1] , 1)).flatten('F') , Idx.flatten('C')] = True;
    AdjacencyMatrix |= AdjacencyMatrix.T #転置とorをとって対称行列にする
    return AdjacencyMatrix , DistMat

for文を使わない実装にこだわったため、実はここが一番苦労しました。
どうもスマートな方法が思いつかず、np.tileやflattenを駆使して、
1としたい要素のインデックスを無理やり並べて指定しています。
苦労した割にfor文を使用した実装と比べて、サンプル数1000程度では、
さほど速くなるわけではないので、ここは素直にfor文使った方が良いと思います。
もし、もっと良い方法があれば、教えてください。

データ間の距離と接続行列から、Heat Kernel行列を計算

ここまでくれば後は特別なことはしません。素直に実装します。

def GetHeatKernel(X , n , t):
    (AdjMat , DistMat) = GetAdjacencyMatrix(X , n)
    W = np.exp(- (DistMat / t) ) * AdjMat
    return W

np.expにnp.ndarrayのインスタンスを入力すると、
各要素のexpを返してくれます。
また、np.ndarrayクラス同士を*演算子で計算すると、
行列の積でなく、行列の各要素同士を掛け算した行列が返ってきます。

ついでだからLEを実装

Laplacian Eigenmapsです。LEについての解説は論文に譲ります。
Heat Kernel行列さえ求めれば、あとはグラフラプラシアン行列 $\bf{L}$ と
重みの対角行列 $\bf{D}$ 求めて、一般化固有値問題を解くだけです。

def LaplacianEigenmaps(X , n , t , m):
    W = GetHeatKernel(X , n , t)
    D = np.diag(np.sum(W , 0)) #重みの対角行列
    L = D - W #グラフラプラシアン
    (Lambda , V) = sci.linalg.eigh(L , D)
    Y = V[: , 1:(m + 1)].T
    return Y

おなじみの3次元スイスロールデータセット（サンプル数1000個）を2次元に次元削減してみました。

元データ