Cython 函数接收可变大小矩阵输入

Question

我正在尝试把一个原生的Python函数的一部分转换成Cython，以提高计算速度。我想专门为那个耗时的循环部分写一个Cython函数（正如ipython lprun友好地告诉我的）。不过，这个函数需要处理大小不一的矩阵……我不知道怎么把这个问题轻松地转到静态类型的Cython上。

for index1 in range(0,num_products):
    for index2 in range(0,num_products):
        cond_prob = (data[index1] * data[index2]).sum() / max(col_sums[index1], col_sums[index2])
        prox[index1][index2] = cond_prob

问题在于，num_products每年都在变化，所以矩阵（数据）的大小也是不固定的。

在这种情况下，最好的策略是什么呢？

1. 我应该写两个C函数吗？一个用memalloc创建一个特定尺寸的矩阵，然后再写一个函数来对这个创建的矩阵进行循环处理？
2. 有没有什么高级的Cython/Numpy技巧可以帮助解决这个问题？我能否写一个C函数，接收一个大小不一的Numpy数组，并传递它的大小？

Answer 1

你有没有试过在numpy中去掉for循环呢？

对于你方程的第一部分，你可以试试这样做：

(data[ np.newaxis,:] * data[:,np.newaxis]).sum(2) 

如果内存不够用的话，你还可以使用np.einsum()这个函数。至于第二部分，如果你还没试过的话，可能也可以想办法用numpy表达式来实现（这会稍微复杂一点）。

Answer 2

Cython代码是（有策略地）静态类型的，但这并不意味着数组的大小必须是固定的。在纯C语言中，传递一个多维数组给函数可能会有点麻烦，但在Cython中，你应该可以做到类似下面的操作：

注意，我从你的后续问题中取了函数和变量的名字。

import numpy as np
cimport numpy as np
cimport cython

@cython.boundscheck(False)
@cython.cdivision(True)
def cooccurance_probability_cy(double[:,:] X):
    cdef int P, i, j, k
    P = X.shape[0]
    cdef double item
    cdef double [:] CS = np.sum(X, axis=1)
    cdef double [:,:] D = np.empty((P, P), dtype=np.float)

    for i in range(P):
        for j in range(P):
            item = 0
            for k in range(P):
                item += X[i,k] * X[j,k]
            D[i,j] = item / max(CS[i], CS[j])
    return D

另一方面，如果你使用正确的函数和一些广播功能，单纯使用Numpy对于这个问题也应该相当快。实际上，由于计算的复杂性主要由矩阵乘法决定，我发现下面的代码比上面的Cython代码快得多（np.inner使用了高度优化的BLAS例程）：

def new(X):
    CS = np.sum(X, axis=1, keepdims=True)
    D = np.inner(X,X) / np.maximum(CS, CS.T)
    return D