使用Cython在numpy網格上評估C函數

Question

用cython簡單包裝C代碼的示例很好地描述了如何在從numpy傳遞的數組上評估用C編寫的函數並將結果返回到numpy數組中。

如何做同樣的事情但返回2D數組呢？ 即我想在由兩個Numpy數組定義的網格上評估C函數，並將結果作為Numpy 2D數組返回。

就像這樣（使用與上面鏈接相同的功能）。 顯然，現在不能使用double z []，但是我不確定如何將2D numpy數組傳遞給C。

/*  fc.cpp    */
int fc( int N, const double a[], const double b[], double z[] )
    {
    for( int i = 0;  i < N;  i ++ ){
        for( int j = 0;  j < N;  j ++ ){
            z[i][j] = somefunction(a[i],b[j]);
    }
    return N;
}

這是原始的.pyx文件（請參見下文）。

import numpy as np
cimport numpy as np
cdef extern from "fc.h": 
    int fc( int N, double* a, double* b, double* z )  # z = a + b

def fpy( N,
    np.ndarray[np.double_t,ndim=1] A,
    np.ndarray[np.double_t,ndim=1] B,
    np.ndarray[np.double_t,ndim=1] Z ):
    """ wrap np arrays to fc( a.data ... ) """
       assert N <= len(A) == len(B) == len(Z)
       fcret = fc( N, <double*> A.data, <double*> B.data, <double*> Z.data )

    return fcret

非常感謝。

Answer 1

您可以將普通陣列用於2D矩陣。 您只需要為函數指定尺寸的長度即可。

在C文件中執行以下操作：（z現在是長度為N * N的數組）

int fc( int N, const double a[], const double b[], double z[] )
{
    for( int i = 0;  i < N;  i++ ){
        for( int j = 0;  j < N;  j ++ ){
            z[(i*N)+j] = somefunction(a[i],b[j]);
    }
    return N;
}

在Python中，您需要執行相同的操作，因此可以使用具有N * N個元素的1D數組而不是2D矩陣。

更新3D外殼

（z現在是長度為N * N * N的數組）

int fc( int N, const double a[], const double b[],const double c[], double z[] )
{
    for( int i = 0;  i < N;  i++ ){
        for( int j = 0;  j < N;  j ++ ){
           for( int k = 0;  k < N;  k ++ ){
            z[((i*N)+j)*N+k] = somefunction(a[i],b[j],c[k]);
    }
    return N;
}