I have a n-dimensional array I want to apply a windowing function to. In short, I need to construct a window function for each dimension and multiply it to the a array. For example, I first construct the window function for the first dimension, stack it for the remaining dimensions and multiply it point-wise to the array a. I sequentially do this for all the array dimensions.
I have been able to do so by accounting for the dimensions of the array in a conditional structure such as if a.ndim == 1: ... elif a.ndim == 2: ... and so on. Here is a MCVE with the non-generalized version that does this (examples are 1D and 3D arrays):
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]
My goal is to generalize this conditional structure so I do not need to check the dimensions case of the array. Something like for axis, axis_size in enumerate(a.shape):... would be more elegant and account for a n-dimensional array, instead of just 1, 2 or 3 dimensions. My attempt has involved something with itertools.cycle and itertools.islice consisting of
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
But never got very far since for a.ndim == 3 is difficult to construct the window function from the second axis as I first need to stack for the first axis first and then last axis, contrary to the other window functions (first and last axis) where I stack sequentially on the following axis by cycling through axis_cycle.
Found a way to generalize it for a n-dimensional array by always start stacking the window array on the first dimension and skipping the stacking operation on the own axis.
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
This returns the same result for a 1D, 2D or 3D array as showed in the question MCVE test cases.