使用 Cython 在 numpy 网格上评估 C 函数

Question

Simple wrapping of C code with cython 中的示例很好地描述了如何在从 numpy 传递的数组上评估用 C 编写的函数并将结果返回到 numpy 数组中。

如何做同样的事情但返回一个二维数组？ IE。我想在由两个 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 个元素的一维数组而不是二维矩阵。

更新 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;
}