用numpy(或其他向量化方法)优化这个函数

Optimize this function with numpy (or other vectorization methods)

我正在计算 Python 种群遗传学领域的经典计算。我很清楚有很多算法可以完成这项工作,但出于某种原因我想构建自己的算法。

下面的段落是图片,因为 Whosebug 不支持 MathJax

我想要一个有效的算法来计算那些 Fst。目前我只设法进行循环,没有计算被矢量化 我如何使用 numpy(或其他矢量化方法)进行此计算?


这是我认为应该完成的代码:

def Fst(W, p):
    I = len(p[0])
    K = len(p)
    H_T = 0
    H_S = 0
    for i in xrange(I):
        bar_p_i = 0
        for k in xrange(K):
            bar_p_i += W[k] * p[k][i]
            H_S += W[k] * p[k][i] * p[k][i]
        H_T += bar_p_i*bar_p_i
    H_T = 1 - H_T
    H_S = 1 - H_S
    return (H_T - H_S) / H_T

def main():
    W = [0.2, 0.1, 0.2, 0.5]
    p = [[0.1,0.3,0.6],[0,0,1],[0.4,0.5,0.1],[0,0.1,0.9]]
    F = Fst(W,p)
    print("Fst = " + str(F))
    return

main()

这里没有理由使用循环。而且你真的不应该为这些东西使用 Numba 或 Cython - 像你拥有的线性代数表达式是 Numpy 中矢量化操作背后的全部原因。

由于如果您继续使用 Numpy,此类问题将一次又一次地出现,我建议您掌握 Numpy 中线性代数的基本知识。您可能会发现这本书的章节有帮助:

https://www.safaribooksonline.com/library/view/python-for-data/9781449323592/ch04.html

至于你的具体情况:从你的变量创建 numpy 数组开始:

import numpy as np
W = np.array(W)
p = np.array(p)

现在,您的 \bar p_i^2 由点积定义。这很简单:

bar_p_i = p.T.dot(W)

注意 T,用于转置,因为点积取第一个矩阵的最后一个索引和第二个矩阵的第一个索引索引的元素的总和。转置会反转索引,因此第一个索引成为最后一个索引。

你H_t是由一个总和定义的。这也很简单:

H_T = 1 - bar_p_i.sum()

同样适用于您的 H_S:

H_S = 1 - ((bar_p_i**2).T.dot(W)).sum()