[Numpy-discussion] A Cython apply_along_axis function

[Keith Goodman]

It's hard to write Cython code that can handle all dtypes and
arbitrary number of dimensions. The former is typically dealt with
using templates, but what do people do about the latter?

I'm trying to take baby steps towards writing an apply_along_axis
function that takes as input a cython function, numpy array, and axis.
I'm using the following numpy ticket as a guide but I'm really just
copying and pasting without understanding:

http://projects.scipy.org/numpy/attachment/ticket/1213/_selectmodule.pyx

Can anyone spot why I get a segfault on the call to nanmean_1d in
apply_along_axis?

import numpy as np
cimport numpy as np
import cython

cdef double NAN = <double> np.nan
ctypedef np.float64_t (*func_t)(void *buf, np.npy_intp size, np.npy_intp s)

def apply_along_axis(np.ndarray[np.float64_t, ndim=1] a, int axis):

    cdef func_t nanmean_1d
    cdef np.npy_intp stride, itemsize
    cdef int ndim = a.ndim
    cdef np.float64_t out

    itemsize = <np.npy_intp> a.itemsize

    if ndim == 1:
        stride = a.strides[0] // itemsize # convert stride bytes --> items
        out = nanmean_1d(a.data, a.shape[0], stride)
    else:
        raise ValueError("Not yet coded")

    return out

cdef np.float64_t nanmean_1d(void *buf, np.npy_intp n, np.npy_intp s):
    "nanmean of buffer."
    cdef np.float64_t *a = <np.float64_t *> buf #
    cdef np.npy_intp i, count = 0
    cdef np.float64_t asum, ai
    if s == 1:
        for i in range(n):
            ai = a[i]
            if ai == ai:
                asum += ai
                count += 1
    else:
        for i in range(n):
            ai = a[i*s]
            if ai == ai:
                asum += ai
                count += 1
    if count > 0:
        return asum / count
    else:
        return NAN