在nd-array的一个轴上应用1D功能

Question

我想要的是:

我想将1D函数应用于任意形状的ndarray,以便它修改某个轴.类似于中的axis论点numpy.fft.fft.

请看以下示例:

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 ndarray f和另外两个1D阵列x,和y.它执行f(x)从x-axis到y-axis 的傅里叶变换.结果存储在out参数中.

现在我想将它包装在一个更通用的函数中transf_all,它可以采用任意形状的ndarray和一个axis参数,指定要转换的轴.

笔记:

• 我的代码实际上是用Cython编写的.理想情况下,magic_iterator在Cython中会很快.
• 该函数transf1d实际上是一个C函数,它在out参数中返回其输出.因此,我无法使用它numpy.apply_along_axis.
• 因为transf1d实际上是一个非常复杂的C函数,我无法重写它来处理任意数组.我需要将它包装在一个处理其他维度的Cython函数中.
• 注意,数组x,并且y它们的长度可以不同.

我的问题:

我怎样才能做到这一点？如何迭代ndarray的任意维度,以便在每次迭代时我将获得包含指定的一维数组axis？

我看了一下nditer,但我不确定这是否真的是这项工作的正确工具.

干杯!

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)

我相信这应该可以满足你的要求，尽管我还没有测试过。请注意，apply_along_axis 内部发生的循环具有 python 级别的性能，因此这仅在样式方面对操作进行矢量化，而不是在性能方面。然而，这很可能并不重要，假设内部循环采用外部 C 代码的决定是合理的，因为它首先是一个不平凡的操作。