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

Question

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

下面的段落是图片，因为 StackOverflow 不支持 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()

Answer 1

没有理由在这里使用循环。而且你真的不应该为这些东西使用 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()