Incorrect window correction applied in _window_corrrection_factor() when scaling="spectrum"

[colinbeyers]

When using xrft.power_spectrum() with scaling="spectrum" and window_correction=True, the window correction appears to be incorrectly applied. This does not happen when using scaling="density", where the window correction is correctly applied. The issue seems to be coming from _window_corrrection_factor(), where when scaling="spectrum", the window correction is calculated as windows.mean() ** 2 instead of being correctly calculated as (windows**2).mean(), which is done when scaling="density".

The code below reproduces this discrepancy with synthetic data, and shows that when using xrft.power_spectrum() with scaling="spectrum" and window_correction=False and applying the correct correction manually, the variance calculated from the spectrum matches the variance of the original signal within about <1%.

import numpy as np
import xarray as xr
import xrft

# -----------------------------
# Parameters and signal setup
# -----------------------------
Nx, Ny = 256, 256
np.random.seed(42)
u = np.random.randn(Ny, Nx)
da = xr.DataArray(u, dims=['y', 'x'])

# Demean and detrend
da_preproc = da - da.mean()

# -----------------------------
# 1. Real-space variance
# -----------------------------
var_real = da_preproc.var().item()

# -----------------------------
# 2. PSD (density), with window correction [WORKS]
# -----------------------------
psd = xrft.power_spectrum(
    da_preproc,
    dim=["y", "x"],
    window="hann",
    detrend="linear",
    scaling="density",
    window_correction=True,
)
# Estimate grid spacing from frequency coords
dx = float(psd["freq_x"].spacing)
dy = float(psd["freq_y"].spacing)
var_psd = psd.sum().item() * dx * dy  # Parseval

# -----------------------------
# 3. PS (spectrum), with window correction [SUSPECTED BUG]
# -----------------------------
ps = xrft.power_spectrum(
    da_preproc,
    dim=["y", "x"],
    window="hann",
    detrend="linear",
    scaling="spectrum",
    window_correction=True,
)
var_ps_wrong = ps.sum().item()

# -----------------------------
# 4. PS (spectrum), with NO window correction [UNDERESTIMATES (as expected)]
# -----------------------------
ps_nocorr = xrft.power_spectrum(
    da_preproc,
    dim=["y", "x"],
    window="hann",
    detrend="linear",
    scaling="spectrum",
    window_correction=False,
)
var_ps_uncorrected = ps_nocorr.sum().item()

# -----------------------------
# 5. PS (spectrum), manually corrected [FIX]
# -----------------------------
# Compute manual Hann correction factor for 2D
win_x = np.hanning(Nx)
win_y = np.hanning(Ny)
corr_factor = 1 / (np.mean(win_x**2) * np.mean(win_y**2))

ps_manual_corr = ps_nocorr * corr_factor
var_ps_manual = ps_manual_corr.sum().item()

# -----------------------------
# Print comparison
# -----------------------------
print(f"{'Source':<30} {'Variance':>15}")
print("-" * 50)
print(f"{'Real-space variance':<30} {var_real:15.5e}")
print(f"{'XRFT PSD (density) + correction':<30} {var_psd:15.5e}")
print(f"{'XRFT PS (spectrum) + correction':<30} {var_ps_wrong:15.5e}")
print(f"{'XRFT PS (spectrum) uncorrected':<30} {var_ps_uncorrected:15.5e}")
print(f"{'Manual-corrected PS':<30} {var_ps_manual:15.5e}")