Apply 1D function on one axis of nd-array

Question

What I want:

I want to apply a 1D function to an arbitrarily shaped ndarray, such that it modifies a certain axis. Similar to the axis argument in numpy.fft.fft.

Take the following example:

import numpy as np


def transf1d(f, x, y, out):
    """Transform `f(x)` to `g(y)`.

    This function is actually a C-function that is far more complicated
    and should not be modified. It only takes 1D arrays as parameters.    

    """
    out[...] = (f[None,:]*np.exp(-1j*x[None,:]*y[:,None])).sum(-1)


def transf_all(F, x, y, axis=-1, out=None):
    """General N-D transform.

    Perform `transf1d` along the given `axis`.

    Given the following:
      F.shape == (2, 3, 100, 4, 5)
      x.shape == (100,)
      y.shape == (50,)
      axis == 2

    Then the output shape would be:
      out.shape == (2, 3, 50, 4, 5)

    This function should wrap `transf1d` such that it works on arbitrarily
    shaped (compatible) arrays `F`, and `out`.

    """
    if out is None:
        shape = list(np.shape(F))
        shape[axis] = np.size(y)

    for f, o in magic_iterator(F, out):
        # Given above shapes:
        #   f.shape == (100,)
        #   o.shape == (50,)
        transf1d(f, x, y, o)

    return out

The function transf1d takes a 1D ndarray f, and two more 1D arrays x, and y. It performs a fourier transform of f(x) from the x-axis to the y-axis. The result is stored in the out argument.

Now I want to wrap this in a more general function transf_all, that can take ndarrays of arbitrary shape along with an axis argument, that specifies along which axis to transform.

Notes:

• My code is actually written in Cython. Ideally, the magic_iterator would be fast in Cython.
• The function transf1d actually is a C-function that returns its output in the out argument. Hence, I couldn't get it to work with numpy.apply_along_axis.
• Because transf1d is actually a pretty complicated C-function I cannot rewrite it to work on arbitrary arrays. I need to wrap it in a Cython function that deals with the additional dimensions.
• Note, that the arrays x, and y can differ in their lengths.

My question:

How can I do this? How can I iterate over arbitrary dimensions of an ndarray such that at each iteration I will get a 1D array containing the specified axis?

I had a look at nditer, but I'm not sure if that is actually the right tool for this job.

Cheers!

Answer 1

import numpy as np


def transf1d(f, x, y, out):
    """Transform `f(x)` to `g(y)`.

    This function is actually a C-function that is far more complicated
    and should not be modified. It only takes 1D arrays as parameters.

    """
    out[...] = (f[None,:]*np.exp(-1j*x[None,:]*y[:,None])).sum(-1)


def transf_all(F, x, y, axis=-1, out=None):
    """General N-D transform.

    Perform `transf1d` along the given `axis`.

    Given the following:
      F.shape == (2, 3, 100, 4, 5)
      x.shape == (100,)
      y.shape == (50,)
      axis == 2

    Then the output shape would be:
      out.shape == (2, 3, 50, 4, 5)

    This function should wrap `transf1d` such that it works on arbitrarily
    shaped (compatible) arrays `F`, and `out`.

    """

    def wrapper(f):
        """
        wrap transf1d for apply_along_axis compatibility
        that is, having a signature of F.shape[axis] -> out.shape[axis]
        """
        out = np.empty_like(y)
        transf1d(f, x, y, out)
        return out
    return np.apply_along_axis(wrapper, axis, F)

I believe this should do what you want, although I havnt tested it. Note that the looping happening inside apply_along_axis has python-level performance though, so this only vectorizes the operation in terms of style, not in terms of performance. However, that is quite probably of no concern, assuming the decision to resort to external C code for the inner loop is justified by it being a nontrivial operation in the first place.