Python：多分布计算的Kullback Leibler散度的快速高效实现

Question

我遇到了一个问题：我有一个包含大约2万组离散分布（直方图）的矩阵，我需要计算这些分布之间的KL散度（KLD）。最简单的方法就是用两个循环，逐一计算每一对分布的KLD。这种方法很耗时间，真的很耗时间。我在想有没有什么基于矩阵或数组的计算方法可以解决这个问题。肯定不是我第一个遇到这个问题的人。可惜我对Python还不太熟悉。任何帮助都会非常感谢。谢谢！A.

Answer 1

你可以试试 numba

import numpy as np
from numba import njit

@njit
def kld(p, q):
    """
    Epsilon is used here to avoid infinity result
    """
    epsilon = 1e-8
    
    p = p+epsilon
    q = q+epsilon

    divergence = np.sum(p*np.log(p/q))
    return divergence

Answer 2

这个应该可以正常运行。

def kullback_leibler_divergence(X):

    """
    Finds the pairwise Kullback-Leibler divergence
    matrix between all rows in X.

    Parameters
    ----------
    X : array_like, shape (n_samples, n_features)
        Array of probability data. Each row must sum to 1.

    Returns
    -------
    D : ndarray, shape (n_samples, n_samples)
        The Kullback-Leibler divergence matrix. A pairwise matrix D such that D_{i, j}
        is the divergence between the ith and jth vectors of the given matrix X.

    Notes
    -----
    Based on code from Gordon J. Berman et al.
    (https://github.com/gordonberman/MotionMapper)

    References:
    -----------
    Berman, G. J., Choi, D. M., Bialek, W., & Shaevitz, J. W. (2014). 
    Mapping the stereotyped behaviour of freely moving fruit flies. 
    Journal of The Royal Society Interface, 11(99), 20140672.
    """

    X_log = np.log(X)
    X_log[np.isinf(X_log) | np.isnan(X_log)] = 0

    entropies = -np.sum(X * X_log, axis=1)

    D = np.matmul(-X, X_log.T)
    D = D - entropies
    D = D / np.log(2)
    D *= (1 - np.eye(D.shape[0]))

    return D

Answer 3

itertools是一个强大的工具，可以用来处理不同对象之间的各种排列组合等。它的一个功能是通过product函数实现双重“for”循环：

from itertools import product
KL = [kl(hist_mat[:,i],hist_mat[:,j]) for i,j in product( range(0,K), range(0,K) )]

K表示要成对比较的直方图数量，
hist_mat包含了每个直方图，存放在它的每一列中。
kl是你选择的Kullback-Leibler散度的实现方法。

唯一的缺点是，对于大数据集，它的运行速度还是比较慢。

祝好，
A.

Answer 4

我发现了一个关于如何使用numpy计算文本文件之间的Kullback-Leibler（KL）距离的帖子，里面提到SciPy已经实现了KL距离的计算，代码如下：

scipy.stats.entropy(pk, qk=None, base=None)

你可以在这个链接找到相关的文档：http://docs.scipy.org/doc/scipy-dev/reference/generated/scipy.stats.entropy.html；使用这个方法可以让计算变得更快。

另外，还有一个用numpy实现的例子，可以参考这个链接：https://gist.github.com/larsmans/3104581：

import numpy as np
 
def kl(p, q):
    """Kullback-Leibler divergence D(P || Q) for discrete distributions
    
    Parameters
    ----------
    p, q : array-like, dtype=float, shape=n
    Discrete probability distributions.
    """
    p = np.asarray(p, dtype=np.float)
    q = np.asarray(q, dtype=np.float)
    
    return np.sum(np.where(p != 0, p * np.log(p / q), 0))

需要注意的是，把数据转换成numpy数组和从numpy数组转换回去的过程比较慢。我不太确定是保持20,000个一维的numpy数组会更好（这样可能会占用很多内存），还是保持一个二维的numpy数组，然后对其中的切片进行操作。