在ndarray的一个轴上应用1D函数

Question

我想要的：

我想对任意形状的ndarray应用一个1D函数，这样它就可以修改某个轴。类似于numpy.fft.fft中的axis参数。在

以下面的例子为例：

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

函数transf1d接受一个1D数组f，以及另外两个1D数组x，和{}。它执行f(x)从x-轴到y-轴的傅立叶变换。结果存储在out参数中。在

现在我想用一个更通用的函数transf_all来包装这个函数，它可以接受任意形状的ndarrays以及一个axis参数，该参数指定沿哪个轴进行变换。在

注意事项：

• 我的代码实际上是用Cython编写的。理想情况下，magic_iterator在Cython中会很快。在
• 函数transf1d实际上是一个C函数，它在out参数中返回其输出。因此，我无法让它与numpy.apply_along_axis一起工作。在
• 因为transf1d实际上是一个非常复杂的C函数，我不能重写它来处理任意数组。我需要把它包装在一个Cython函数中，这个函数处理额外的维度。在
• 注意，数组x和{}的长度可能不同。在

我的问题：

我该怎么做？如何在nArray的任意维上迭代，以便在每次迭代时都能得到一个包含指定的axis的一维数组？在

我看了一下nditer，但我不确定这是否是这项工作的正确工具。在

干杯！在

Answer 1

我发现了一种在Cython中使用numpyc-API（下面的代码）迭代除了一个轴以外的所有轴的方法。然而，它并不漂亮。是否值得这样做取决于内部函数和数据的大小。在

如果有人知道一个更优雅的方式来做这件事，请告诉我。在

我比较了Eelco的解决方案，它们在大参数下运行的速度相当。对于较小的参数，C-API解决方案更快：

In [5]: y=linspace(-1,1,100);

In [6]: %timeit transf.apply_along(f, x, y, axis=1)
1 loops, best of 3: 5.28 s per loop

In [7]: %timeit transf.transfnd(f, x, y, axis=1)
1 loops, best of 3: 5.16 s per loop

如您所见，对于这个输入，两个函数的速度大致相同。在

^{pr2}$

但是，对于较小的输入数组，C-API方法更快。在

代码

#cython: boundscheck=False
#cython: wraparound=False
#cython: cdivision=True

import numpy as np
cimport numpy as np
np.import_array()

cdef extern from "complex.h":
    double complex cexp(double complex z) nogil

cdef void transf1d(double complex[:] f,
                   double[:] x,
                   double[:] y,
                   double complex[:] out,
                   int Nx,
                   int Ny) nogil:
    cdef int i, j

    for i in xrange(Ny):
        out[i] = 0
        for j in xrange(Nx):
            out[i] = out[i] + f[j]*cexp(-1j*x[j]*y[i])


def transfnd(F, x, y, axis=-1, out=None):
    # Make sure everything is a numpy array.
    F = np.asanyarray(F, dtype=complex)
    x = np.asanyarray(x, dtype=float)
    y = np.asanyarray(y, dtype=float)

    # Calculate absolute axis.
    cdef int ax = axis
    if ax < 0:
        ax = np.ndim(F) + ax

    # Calculate lengths of the axes `x`, and `y`.
    cdef int Nx = np.size(x), Ny = np.size(y)

    # Output array.
    if out is None:
        shape = list(np.shape(F))
        shape[axis] = Ny
        out = np.empty(shape, dtype=complex)
    else:
        out = np.asanyarray(out, dtype=complex)

    # Error check.
    assert np.shape(F)[axis] == Nx, \
            'Array length mismatch between `F`, and `x`!'
    assert np.shape(out)[axis] == Ny, \
            'Array length mismatch between `out`, and `y`!'
    f_shape = list(np.shape(F))
    o_shape = list(np.shape(out))
    f_shape[axis] = 0
    o_shape[axis] = 0
    assert f_shape == o_shape, 'Array shape mismatch between `F`, and `out`!'

    # Construct iterator over all but one axis.
    cdef np.flatiter itf = np.PyArray_IterAllButAxis(F, &ax)
    cdef np.flatiter ito = np.PyArray_IterAllButAxis(out, &ax)
    cdef int f_stride = F.strides[axis]
    cdef int o_stride = out.strides[axis]

    # Memoryview to access one slice per iteration.
    cdef double complex[:] fdat
    cdef double complex[:] odat
    cdef double[:] xdat = x
    cdef double[:] ydat = y

    while np.PyArray_ITER_NOTDONE(itf):
        # View the current `x`, and `y` axes.
        fdat = <double complex[:Nx]> np.PyArray_ITER_DATA(itf)
        fdat.strides[0] = f_stride
        odat = <double complex[:Ny]> np.PyArray_ITER_DATA(ito)
        odat.strides[0] = o_stride

        # Perform the 1D-transformation on one slice.
        transf1d(fdat, xdat, ydat, odat, Nx, Ny)

        # Go to next step.
        np.PyArray_ITER_NEXT(itf)
        np.PyArray_ITER_NEXT(ito)

    return out


# For comparison
def apply_along(F, x, y, axis=-1):
    # Make sure everything is a numpy array.
    F = np.asanyarray(F, dtype=complex)
    x = np.asanyarray(x, dtype=float)
    y = np.asanyarray(y, dtype=float)

    # Calculate absolute axis.
    cdef int ax = axis
    if ax < 0:
        ax = np.ndim(F) + ax

    # Calculate lengths of the axes `x`, and `y`.
    cdef int Nx = np.size(x), Ny = np.size(y)

    # Error check.
    assert np.shape(F)[axis] == Nx, \
            'Array length mismatch between `F`, and `x`!'

    def wrapper(f):
        out = np.empty(Ny, complex)
        transf1d(f, x, y, out, Nx, Ny)
        return out

    return np.apply_along_axis(wrapper, axis, F)

使用以下内容生成setup.py

from distutils.core import setup
from Cython.Build import cythonize
import numpy as np

setup(
    name = 'transf',
    ext_modules = cythonize('transf.pyx'),
    include_dirs = [np.get_include()],
)

Answer 2

回答您的问题：

如果您真的只想迭代给定轴以外的所有轴，可以使用：

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

也许有更优雅的方式来做，但这应该行得通。在

但是，不要重复：

对于您的问题，我将重写您的函数，使其在ndarray的给定轴上工作。我认为这应该行得通：

^{pr2}$

它实际上只是您当前实现的泛化，但是它没有在开始时在F中插入一个新的轴，而是在axis处插入它（可能有比使用list(shape)方法更好的方法，但我只能这样做。最后，您必须在yx外积中添加新的尾随轴，以匹配F中的任意多个尾随索引。在

我真的不知道如何测试这个，但形状都可以，所以请测试一下，让我知道它是否有效。在

Answer 3

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)

我相信这应该是你想要的，尽管我还没有测试过。请注意，apply_内沿_轴发生的循环具有python级别的性能，因此这只是在风格方面而不是在性能方面对操作进行矢量化。然而，这很可能并不重要，假设在内部循环中使用外部C代码的决定是正确的，因为它首先是一个非常重要的操作。在