Apply 1D function on one axis of nd-array

user_841562 0
Viewed 819

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!

cython numpy arrays python
3 Answers

To answer your question:

If you really just want to iterate over all but a given axis, you can use:

for s in itertools.product(map(range, arr.shape[:axis]+arr.shape[axis+1:]):
    arr[s[:axis] + (slice(None),) + s[axis:]]

Maybe there's a more elegant way to do it, but this should work.

But, don't iterate:

For your problem, I would just rewrite your function to work on a given axis of an ndarray. I think this should work:

def transfnd(f, x, y, axis, out):
    s = list(f.shape)
    s.insert(axis, 1)
    yx = [y.size, x.size] + [1]*(f.ndim - axis - 1)
    out[...] = np.sum(f.reshape(*s)*np.exp(-1j*x[None,:]*y[:,None]).reshape(*yx), axis+1)

It's really just the generalization of your current implementation, but instead of inserting a new axis in F at the beginning, it inserts it at axis (there might be a better way to do this than with the list(shape) method, but that was all I could do. Finally, you have to add trailing new axes to your yx outer product, to match as many trailing indices you have in F.

I didn't really know how to test this, but the shapes all work out, so please test it and let me know whether it works.

user_17306740
Related