From 3f16e360a5d6b65bca946596cd619d0c385e5932 Mon Sep 17 00:00:00 2001 From: Udaya Tejas Date: Sat, 31 Oct 2026 15:45:04 -0400 Subject: [PATCH] Fix UnboundLocalError in compute_tke_spectrum when length is given compute_tke_spectrum only assigned the domain lengths lx, ly, lz when length was None, so passing an explicit length raised UnboundLocalError at the first fftfreq call. Use the given length (defaulting to 2*pi) for all three axes. Signed-off-by: Udaya Tejas --- CHANGELOG.md | 3 +++ physicsnemo/metrics/cae/cfd.py | 3 ++- test/metrics/test_metrics_cfd.py | 15 +++++++++++++++ 3 files changed, 20 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index dc4e708b8d..2b9fa1bd67 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -388,6 +388,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- `physicsnemo.metrics.cae.cfd.compute_tke_spectrum` no longer raises + `UnboundLocalError` when `length` is passed; the domain length now sets the + grid spacing used for the wavenumbers. - The GLOBE DrivAerML example trains under `torch.compile` again on PyTorch 2.13 and 2.14. It and the AirfRANS example now set `torch._dynamo.config.specialize_float`, which avoids a Dynamo crash in GLOBE's diff --git a/physicsnemo/metrics/cae/cfd.py b/physicsnemo/metrics/cae/cfd.py index cee219d1a2..54ba6f58ad 100644 --- a/physicsnemo/metrics/cae/cfd.py +++ b/physicsnemo/metrics/cae/cfd.py @@ -246,7 +246,8 @@ def compute_tke_spectrum( """ if length is None: - lx, ly, lz = 2 * np.pi, 2 * np.pi, 2 * np.pi + length = 2 * np.pi + lx, ly, lz = length, length, length u, v, w = field[0, :, :, :], field[1, :, :, :], field[2, :, :, :] diff --git a/test/metrics/test_metrics_cfd.py b/test/metrics/test_metrics_cfd.py index 101b095e67..6daa6373c0 100644 --- a/test/metrics/test_metrics_cfd.py +++ b/test/metrics/test_metrics_cfd.py @@ -114,3 +114,18 @@ def test_compute_tke_spectrum(): assert E.shape == (nx * ny * nz,) assert wv.shape == (n + 1,) assert tke.shape == (n + 1,) + + +def test_compute_tke_spectrum_length(): + n = 8 + field = np.random.rand(3, n, n, n) + + mk_default, E_default, _, tke_default = compute_tke_spectrum(field) + mk_2pi, E_2pi, _, tke_2pi = compute_tke_spectrum(field, length=2 * np.pi) + np.testing.assert_allclose(mk_2pi, mk_default) + np.testing.assert_allclose(E_2pi, E_default) + np.testing.assert_allclose(tke_2pi, tke_default) + + # Wavenumbers scale inversely with the domain length + mk_unit, _, _, _ = compute_tke_spectrum(field, length=1.0) + np.testing.assert_allclose(mk_unit, mk_default * 2 * np.pi)