Evaluating C function on a numpy grid using Cython

Question

The example in Simple wrapping of C code with cython describes nicely how to evaluate a function written in C on an array passed from numpy and return the result in a numpy array.

How would one go about doing the same thing but returning a 2D array? Ie I'd like to evaluate a C function on a grid defined by two numpy arrays, and return the result as a numpy 2D array.

It would be something like this (using same functions as in the link above). Obviously one can't use double z[] now, but I'm not sure how to pass a 2D numpy array to 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;
}

This is the original .pyx file (see below).

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

Many thanks.

Answer 1

You can use a normal array for a 2D Matrix. You need only give the length of the dimension to the function.

In the C file do something as that: (z is now an array of length 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;
}

In Python you need to do the same, so you can use a 1D Array with N*N elements instead of an 2D Matrix.

Update 3D case

(z is now an array of length 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;
}