n 维数组的窗口函数的推广

Question

我有 a 个 n 维数组，我想对其应用 windowing 函数。简而言之，我需要为每个维度构造一个 window 函数并将其乘以 a 数组。例如，我首先为第一个维度构造 window 函数，将其堆叠到剩余维度并将其逐点乘以数组 a。我依次对所有数组维度执行此操作。

我已经能够通过在 if a.ndim == 1: ... elif a.ndim == 2: ... 等条件结构中考虑数组的维度来做到这一点。这是一个具有非通用版本的 MCVE（示例是 1D 和 3D 数组）：

import numpy as np
import scipy.signal as signal

def window_ndim(a, wfunction):
    """
    Performs an in-place windowing on N-dimensional data.
    This is done to mitigate boundary effects in the FFT.
    :param a: Input data to be windowed, modified in place.
    :param wfunction: 1D window generation function. Example: scipy.signal.hamming
    :return: windowed a
    """
    if a.ndim == 1:
        return a * wfunction(len(a))
    elif a.ndim == 2:
        window0 = wfunction(a.shape[0])
        window1 = wfunction(a.shape[1])
        window0 = np.stack([window0] * a.shape[1], axis=1)
        window1 = np.stack([window1] * a.shape[0], axis=0)
        a *= window0*window1
        return a
    elif a.ndim == 3:
        window0 = wfunction(a.shape[0])
        window1 = wfunction(a.shape[1])
        window2 = wfunction(a.shape[2])
        window0 = np.stack([window0] * a.shape[1], axis=1)
        window0 = np.stack([window0] * a.shape[2], axis=2)
        window1 = np.stack([window1] * a.shape[0], axis=0)
        window1 = np.stack([window1] * a.shape[2], axis=2)
        window2 = np.stack([window2] * a.shape[0], axis=0)
        window2 = np.stack([window2] * a.shape[1], axis=1)
        a *= window0*window1*window2
        return a
    else: raise ValueError('Wrong dimensions')

np.random.seed(0)
np.set_printoptions(precision=2)
a = np.random.rand(2,3,4)
# [[[0.55 0.72 0.6  0.54]
#   [0.42 0.65 0.44 0.89]
#   [0.96 0.38 0.79 0.53]]

#  [[0.57 0.93 0.07 0.09]
#   [0.02 0.83 0.78 0.87]
#   [0.98 0.8  0.46 0.78]]]
a_windowed = window_ndim(a, signal.hamming)
# [[[2.81e-04 3.52e-03 2.97e-03 2.79e-04]
#   [2.71e-03 3.98e-02 2.70e-02 5.71e-03]
#   [4.93e-04 1.89e-03 3.90e-03 2.71e-04]]

#  [[2.91e-04 4.56e-03 3.50e-04 4.46e-05]
#   [1.29e-04 5.13e-02 4.79e-02 5.57e-03]
#   [5.01e-04 3.94e-03 2.27e-03 4.00e-04]]]

a = np.random.rand(10) # [0.12 0.64 0.14 0.94 0.52 0.41 0.26 0.77 0.46 0.57]
a_windowed = window_ndim(a, signal.hamming) # [0.01 0.12 0.07 0.73 0.51 0.4  0.2  0.36 0.09 0.05]

我的目标是泛化这个条件结构，这样我就不需要检查数组的维度大小写了。像 for axis, axis_size in enumerate(a.shape):... 这样的东西会更优雅并且可以解释 n 维数组，而不仅仅是 1、2 或 3 维。我的尝试涉及 itertools.cycle 和 itertools.islice 组成的

axis_idxs = np.arange(len(a.shape))
the_cycle = cycle(axis_idxs)
for axis, axis_size in enumerate(a.shape):
    axis_cycle = islice(the_cycle, axis, None)
    next_axis = next(axis_cycle)
    window = wfunction(axis_size)
    window = np.stack([window]*a.shape[next_axis], axis=next_axis)
    ...
    a *= window
return a

但从来没有走得太远，因为 a.ndim == 3 很难从第二个轴构造 window 函数，因为我首先需要先为第一个轴堆叠，然后再为最后一个轴堆叠，这与其他 window 函数（第一个和最后一个轴），我通过循环 axis_cycle.

顺序堆叠在后续轴上

Answer 1

找到了一种将其推广到 n 维数组的方法，方法是始终在第一维开始堆叠 window 数组并跳过自身轴上的堆叠操作。

def window_ndim(a, wfunction):
    for axis, axis_size in enumerate(a.shape):
        window = wfunction(axis_size)
        for i in range(len(a.shape)):
            if i == axis:
                continue
            else:
                window = np.stack([window] * a.shape[i], axis=i)
        a *= window
    return a

这 returns 与问题 MCVE 测试用例中所示的 1D、2D 或 3D 数组的结果相同。

n 维数组的窗口函数的推广

Generalization of a windowing function for n-dimensional array

python

multidimensional-array

windowing

python-3.x