使用numpy（或其他矢量化方法）优化此函数

Question

我正在用Python计算一个群体遗传学领域的经典计算。我很清楚，有很多算法，做这项工作，但我想建立自己的一些原因。你知道吗

下面的段落是一张图片，因为StackOverflow不支持MathJax

我想有一个有效的算法来计算这些Fst。目前我只做For循环，没有矢量化计算如何使用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\u i^2由点积定义。这很简单：

bar_p_i = p.T.dot(W)

注意T，对于转置，因为点积取第一个矩阵的最后一个索引和第二个矩阵的第一个索引所索引的元素之和。转置反转索引，使第一个索引成为最后一个索引。你知道吗

你可以用一个和来定义它。这也很简单：

H_T = 1 - bar_p_i.sum()

同样地，对于您的H\ S：

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