Correction for reduced energy when windowing

Applying window weighting (e.g. Hann) to a signal reduces its energy. I think it’s common to want to correct for this. For example, it appears that such a correction is applied by default in scipy.signal.welch:

import xrft
import numpy as np
import scipy.signal
import xarray as xr
import matplotlib.pyplot as plt 

x = xr.DataArray(np.random.normal(0, 1, size=10000), 
                 dims=['t'], 
                 coords=[range(0,10000)])
n_segment = 100

f_scipy, p_scipy = scipy.signal.welch(x.values,
                                      window='hann',
                                      nperseg=n_segment,
                                      noverlap=0,
                                      return_onesided=False)
                
p_xrft = xrft.xrft.power_spectrum(x.chunk({'t': n_segment}), 
                                  dim=['t'],
                                  detrend='constant',
                                  window=True,
                                  chunks_to_segments=True).mean('t_segment')

# Correction tends to 8/3 for infinitely long Hann window
p_xrft_corrected = (1 / np.mean(np.hanning(n_segment)**2))*p_xrft 

plt.plot(f_scipy, p_scipy, 
         label='scipy psd')
plt.plot(p_xrft.freq_t, p_xrft, 
         color='C2', label='xrft psd')
plt.plot(p_xrft.freq_t, p_xrft_corrected, 
         linestyle='--', color='C1', label='corrected xrft psd')
plt.ylim(0,1.6)
plt.xlabel('freq')
plt.ylabel('psd')
plt.legend()

Is there an appetite (beyond my own) to have an option to apply such a correction when using window=True in xrft?

[roxyboy]

If you pass window='windows[0]' to sps.welch() rather than window='hann', then you’ll get exactly the same values as our corrected results. Perhaps it would be a good idea for xrft to also use scipy.windows? This would make extending to other windows very easy.

I agree with this but we’ll keep it for another PR.

[roxyboy]

Sorry, I forgot to mention that I changed the _apply_window function to also return the windows as below:

def _apply_window(da, dims, window_type="hanning"):
    """Creating windows in dimensions dims."""

    if window_type not in ["hanning"]:
        raise NotImplementedError("Only hanning window is supported for now.")

    numpy_win_func = getattr(np, window_type)

    if da.chunks:

        def dask_win_func(n):
            return dsar.from_delayed(delayed(numpy_win_func, pure=True)(n), (n,), float)

        win_func = dask_win_func
    else:
        win_func = numpy_win_func

    windows = [
        xr.DataArray(win_func(len(da[d])), dims=da[d].dims, coords=da[d].coords)
        for d in dims
    ]

    return windows, da * reduce(operator.mul, windows[::-1])