使用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()