Fix UnboundLocalError in compute_tke_spectrum when length is given

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

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, :, :, :]
 

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)