[Numpy-discussion] A Cython apply_along_axis function

[Keith Goodman]

On Wed, Dec 1, 2010 at 6:07 PM, Keith Goodman <kwgoodman@gmail.com> wrote:
> On Wed, Dec 1, 2010 at 5:53 PM, David <david@silveregg.co.jp> wrote:
>
>> On 12/02/2010 04:47 AM, Keith Goodman wrote:
>>> 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?
>>
>> The only way that I know to do that systematically is iterator. There is
>> a relatively simple example in scipy/signal (lfilter.c.src).
>>
>> I wonder if it would be possible to add better support for numpy
>> iterators in cython...
>
> Thanks for the tip. I'm starting to think that for now I should just
> template both dtype and ndim.

I ended up templating both dtype and axis. For the axis templating I
used two functions: looper and loop_cdef.

LOOPER

    Make a 3d loop template:

    >>> loop = '''
    .... for iINDEX0 in range(nINDEX0):
    ....         for iINDEX1 in range(nINDEX1):
    ....             amin = MAXDTYPE
    ....         for iINDEX2 in range(nINDEX2):
    ....                 ai = a[INDEXALL]
    ....             if ai <= amin:
    ....                 amin = ai
    ....         y[INDEXPOP] = amin
    .... '''

    Import the looper function:

    >>> from bottleneck.src.template.template import looper

    Make a loop over axis=0:

    >>> print looper(loop, ndim=3, axis=0)
    for i1 in range(n1):
        for i2 in range(n2):
            amin = MAXDTYPE
            for i0 in range(n0):
                ai = a[i0, i1, i2]
                if ai <= amin:
                    amin = ai
            y[i1, i2] = amin

    Make a loop over axis=1:

    >>> print looper(loop, ndim=3, axis=1)
    for i0 in range(n0):
        for i2 in range(n2):
            amin = MAXDTYPE
            for i1 in range(n1):
                ai = a[i0, i1, i2]
                if ai <= amin:
                    amin = ai
            y[i0, i2] = amin

LOOP_CDEF

    Define parameters:

    >>> ndim = 3
    >>> dtype = 'float64'
    >>> axis = 1
    >>> is_reducing_function = True

    Import loop_cdef:

    >>> from bottleneck.src.template.template import loop_cdef

    Make loop initialization code:

    >>> print loop_cdef(ndim, dtype, axis, is_reducing_function)
        cdef Py_ssize_t i0, i1, i2
        cdef int n0 = a.shape[0]
        cdef int n1 = a.shape[1]
        cdef int n2 = a.shape[2]
        cdef np.npy_intp *dims = [n0, n2]
        cdef np.ndarray[np.float64_t, ndim=2] y = PyArray_EMPTY(2, dims,
                                                  NPY_float64, 0)