🐛[BUG]: power_spectrum in metrics.general puts the zero frequency half a cell off along an odd width

Version

main at ff5d19d08123de47ca446caed1d70a225d540184

On which installation method(s) does this occur?

Source

Describe the issue

physicsnemo.metrics.general.power_spectrum.power_spectrum builds the radial wavenumber grid after fftshift with

k = torch.hypot(xx - h // 2, yy - w / 2).to(torch.float32)

After fftshift the zero frequency of an axis of length n sits at index n // 2 for even and odd n alike. The height uses h // 2, the width uses w / 2. For an even width the two agree. For an odd width w / 2 is half a cell away from the true center, so every wavenumber along the width is off by 0.5 and the whole grid is skewed in one direction.

Three things follow for any field with an odd width.

1. The same cosine along the height and along the width gives different spectra. For a 33 by 33 grid and wavenumber 5 the height version puts all power into the bin around k = 5 and the width version splits it between two neighbouring bins, 6.19 and 5.04 instead of 10.08 and 0.
2. The spectrum of a field differs from the spectrum of its transpose, up to about 12 percent per bin on a 33 by 33 random field. For a 32 by 33 field even the returned bin centers differ from the 33 by 32 case.
3. A single cosine along the width fills two bins instead of one, so a peak reads lower and wider than it is.

The existing test test_power_spectrum only uses a 32 by 32 grid, where w / 2 and w // 2 coincide, so it cannot see this.

The fix is to build the grid in float32 and use w // 2 for the width. A one character change to // alone is not enough because torch.hypot has no integer kernel, hence the float32 arange.

Minimum reproducible example

import math
import torch
from physicsnemo.metrics.general.power_spectrum import power_spectrum

size, m = 33, 5
coords = torch.arange(size, dtype=torch.float32)
wave = torch.cos(2 * math.pi * m * coords / size)
along_h = wave.view(size, 1).repeat(1, size)[None, None]
along_w = wave.view(1, size).repeat(size, 1)[None, None]

_, p_h = power_spectrum(along_h)
_, p_w = power_spectrum(along_w)
print(p_h[0, 0])
print(p_w[0, 0])

x = torch.randn(1, 1, 33, 33)
_, p = power_spectrum(x)
_, p_t = power_spectrum(x.transpose(-2, -1))
print(torch.allclose(p, p_t, rtol=1e-4, atol=1e-6))

Relevant log output

tensor([0., 0., 0., 10.0833, 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
tensor([0., 0., 6.1875, 5.0417, 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
False

Tiny values of order 1e-13 are printed as 0 above. A pytest run of the attached regression test against main is at https://github.com/Nicholas022400701/oss-ci/actions/runs/35357438260 with 4 failed and 2 passed, the two passing cases are the even sized 32 by 32 grids.

Environment details

Ubuntu GitHub Actions runner, Python 3.12.14, torch 2.14.0+cpu, numpy 2.5.2, physicsnemo installed from source with pip install -e .

I have a fix with a regression test ready and will open a pull request that links this issue.

Start in physicsnemo.metrics.general.power_spectrum.power_spectrum and read the existing test_power_spectrum test. Run the current tests with odd-sized height and width cases, then add regression coverage for those cases and confirm the returned spectra and bin centers are consistent. The issue author reports having a fix and regression test ready, so check for overlapping work before starting.