enh: vectorise spectral integrals and speed up spectrum construction

src/waveresponse/_core.py

@@ -1230,7 +1230,10 @@ def _dir_spectrum(self, degrees=None):
 
         f, d, vv = self.grid(freq_hz=False, degrees=degrees)
 
-        s = np.array([trapezoid(vv_f, f) for vv_f in vv.T])
+        # Integrate all directions in one call. The contiguous copy keeps
+        # the summation order identical to integrating one direction at a
+        # time, so the result is bit-for-bit the same as a per-row loop.
+        s = trapezoid(np.ascontiguousarray(vv.T), f, axis=-1)
 
         return d, s
 
@@ -1344,8 +1347,12 @@ def tz(self):
                 "Mean zero-crossing period is only defined for positive real-valued spectra."
             )
 
-        m0 = self.moment(0, freq_hz=True)
-        m2 = self.moment(2, freq_hz=True)
+        # Integrate out the directional domain once for both moments.
+        # ``f**0 * s`` is exactly ``s``, so this equals
+        # ``moment(0) / moment(2)`` bit for bit.
+        f, s = self._freq_spectrum(freq_hz=True)
+        m0 = trapezoid(s, f)
+        m2 = trapezoid((f**2) * s, f)
 
         return np.sqrt(m0 / m2)
 
@@ -1493,7 +1500,10 @@ def _freq_spectrum(self, freq_hz=None):
         d = self._full_range_dir(d)
         vv = self.interpolate(f, d, freq_hz=freq_hz, degrees=False)
 
-        s = np.array([trapezoid(vv_d, d) for vv_d in vv])
+        # Integrate all frequencies in one call. The contiguous copy keeps
+        # the summation order identical to integrating one frequency at a
+        # time, so the result is bit-for-bit the same as a per-row loop.
+        s = trapezoid(np.ascontiguousarray(vv), d, axis=-1)
 
         return f, s
 
@@ -1570,10 +1580,14 @@ def from_spectrum1d(
         else:
             period = 2.0 * np.pi
 
-        for (idx_f, idx_d), val_i in np.ndenumerate(vals):
+        # The direction offset does not depend on frequency; reduce it
+        # once per direction instead of once per grid point.
+        dirs_rel = [_robust_modulus(d_i - dirp, period) for d_i in dirs]
+
+        for idx_f in range(len(vals)):
             f_i = freq[idx_f]
-            d_i = _robust_modulus(dirs[idx_d] - dirp, period)
-            vals[idx_f, idx_d] = spread_fun(f_i, d_i) * val_i
+            for idx_d, d_i in enumerate(dirs_rel):
+                vals[idx_f, idx_d] = spread_fun(f_i, d_i) * vals[idx_f, idx_d]
 
         return cls(
             freq,

src/waveresponse/_utils.py

@@ -8,7 +8,17 @@ def _robust_modulus(x, periodicity):
     Similar to ``x % periodicity``, but ensures that it is robust w.r.t. floating
     point numbers.
     """
-    x = np.asarray_chkfinite(x % periodicity).copy()
+    x = np.asarray(x % periodicity)
+
+    if x.ndim == 0:
+        # Scalar fast path: same arithmetic without the array machinery.
+        if not np.isfinite(x):
+            raise ValueError("array must not contain infs or NaNs")
+        if x == periodicity:
+            x = np.asarray(np.nextafter(x, -1, dtype=x.dtype))
+        return x
+
+    x = np.asarray_chkfinite(x).copy()
     return np.nextafter(x, -1, where=(x == periodicity), out=x)
 
 

tests/test_core.py

@@ -5,7 +5,7 @@
 import numpy as np
 import pandas as pd
 import pytest
-from scipy.integrate import quad
+from scipy.integrate import quad, trapezoid
 from scipy.interpolate import RegularGridInterpolator as RGI
 
 import waveresponse as wr
@@ -3514,6 +3514,31 @@ def spread_fun(f, d):
 
         np.testing.assert_array_almost_equal(spectrum._vals, vals_expect)
 
+    def test_from_spectrum1d_nonuniform_dirs_wrapping_dirp(self):
+        freq = np.array([0.2, 0.5, 1.0, 1.7])
+        dirs = np.array([0.0, 30.0, 45.0, 120.0, 200.0, 330.0])
+        spectrum1d = np.array([1.0, 2.5, 3.0, 0.5])
+        dirp = 300.0
+
+        def spread_fun(f, d):
+            return (1.0 + f) * np.cos(np.radians(d) / 2.0) ** 2
+
+        spectrum = DirectionalSpectrum.from_spectrum1d(
+            freq, dirs, spectrum1d, spread_fun, dirp, freq_hz=True, degrees=True
+        )
+
+        vals_expect = np.array(
+            [
+                [spread_fun(f_i, (d_i - dirp) % 360.0) * s_i for d_i in dirs]
+                for f_i, s_i in zip(freq, spectrum1d)
+            ]
+        )
+        spectrum_expect = DirectionalSpectrum(
+            freq, dirs, vals_expect, freq_hz=True, degrees=True
+        )
+
+        np.testing.assert_array_almost_equal(spectrum._vals, spectrum_expect._vals)
+
     def test_from_spectrum1d_and_integrate_back_rad(self):
         freq = np.linspace(0.0, 1.0, 50)
         dirs = np.linspace(0.0, 2.0 * np.pi, endpoint=False)
@@ -3751,6 +3776,32 @@ def test_spectrum1d_axis0_rad(self):
         np.testing.assert_array_almost_equal(dir_out, dir_expect)
         np.testing.assert_array_almost_equal(spectrum1d_out, spectrum1d_expect)
 
+    def test_spectrum1d_axis1_equals_rowwise_integration(self):
+        freq = np.linspace(0.1, 3.0, 61)
+        dirs = np.radians([5.0, 20.0, 60.0, 90.0, 150.0, 200.0, 250.0, 300.0, 340.0])
+        vals = np.random.default_rng(0).random((len(freq), len(dirs)))
+        spectrum = DirectionalSpectrum(freq, dirs, vals)
+
+        _, spectrum1d_out = spectrum.spectrum1d(axis=1)
+
+        d = DirectionalSpectrum._full_range_dir(dirs)
+        _, _, vv = spectrum.reshape(freq, d).grid()
+        spectrum1d_expect = np.array([trapezoid(vv_i, d) for vv_i in vv])
+
+        np.testing.assert_array_equal(spectrum1d_out, spectrum1d_expect)
+
+    def test_spectrum1d_axis0_equals_rowwise_integration(self):
+        freq = np.linspace(0.1, 3.0, 61)
+        dirs = np.radians([5.0, 20.0, 60.0, 90.0, 150.0, 200.0, 250.0, 300.0, 340.0])
+        vals = np.random.default_rng(0).random((len(freq), len(dirs)))
+        spectrum = DirectionalSpectrum(freq, dirs, vals)
+
+        _, spectrum1d_out = spectrum.spectrum1d(axis=0)
+
+        spectrum1d_expect = np.array([trapezoid(vv_i, freq) for vv_i in vals.T])
+
+        np.testing.assert_array_equal(spectrum1d_out, spectrum1d_expect)
+
     def test_moment_m0_hz(self):
         f0 = 0.0
         f1 = 2.0
@@ -3963,6 +4014,17 @@ def test_tz(self):
 
         assert tz_out == pytest.approx(tz_expect, rel=0.1)
 
+    def test_tz_equals_moment_ratio(self):
+        freq = np.linspace(0.05, 2.0, 40)
+        dirs = np.arange(5, 360, 10)
+        vals = np.random.default_rng(1).random((len(freq), len(dirs)))
+        spectrum = DirectionalSpectrum(freq, dirs, vals, freq_hz=True, degrees=True)
+
+        m0 = spectrum.moment(0, freq_hz=True)
+        m2 = spectrum.moment(2, freq_hz=True)
+
+        assert spectrum.tz == np.sqrt(m0 / m2)
+
     def test_tz_raises_complex(self):
         f0 = 0.0
         f1 = 2.0
@@ -4681,6 +4743,19 @@ def test_spectrum1d_axis0_rad(self):
         np.testing.assert_array_almost_equal(dir_out, dir_expect)
         np.testing.assert_array_almost_equal(spectrum1d_out, spectrum1d_expect)
 
+    def test_spectrum1d_axis0_equals_rowwise_integration(self):
+        freq = np.linspace(0.1, 3.0, 61)
+        dirs = np.radians([5.0, 20.0, 60.0, 90.0, 150.0, 200.0, 250.0, 300.0, 340.0])
+        vals = np.random.default_rng(0).random((len(freq), len(dirs)))
+        spectrum = DirectionalBinSpectrum(freq, dirs, vals)
+
+        _, spectrum1d_out = spectrum.spectrum1d(axis=0)
+
+        _, _, vv = spectrum.grid()
+        spectrum1d_expect = np.array([trapezoid(vv_i, freq) for vv_i in vv.T])
+
+        np.testing.assert_array_equal(spectrum1d_out, spectrum1d_expect)
+
     def test_moment_m0_hz(self):
         f0 = 0.0
         f1 = 2.0
@@ -4897,6 +4972,17 @@ def test_tz(self):
 
         assert tz_out == pytest.approx(tz_expect, rel=0.1)
 
+    def test_tz_equals_moment_ratio(self):
+        freq = np.linspace(0.05, 2.0, 40)
+        dirs = np.arange(5, 360, 10)
+        vals = np.random.default_rng(1).random((len(freq), len(dirs)))
+        spectrum = DirectionalBinSpectrum(freq, dirs, vals, freq_hz=True, degrees=True)
+
+        m0 = spectrum.moment(0, freq_hz=True)
+        m2 = spectrum.moment(2, freq_hz=True)
+
+        assert spectrum.tz == np.sqrt(m0 / m2)
+
     def test_tz_raises_complex(self):
         f0 = 0.0
         f1 = 2.0

tests/test_utils.py

@@ -44,6 +44,27 @@ def test_array_deg(self):
 
         np.testing.assert_almost_equal(x_mod, expect)
 
+    def test_scalar_returns_0d_array(self):
+        x_mod = _robust_modulus(2.5, 2.0)
+
+        assert isinstance(x_mod, np.ndarray)
+        assert x_mod.ndim == 0
+
+    def test_0d_array(self):
+        x_mod = _robust_modulus(np.array(4.0), 2.0)
+
+        assert x_mod.ndim == 0
+        assert x_mod == 0.0
+
+    @pytest.mark.parametrize("x", [np.inf, -np.inf, np.nan])
+    def test_scalar_nonfinite_raises(self, x):
+        with pytest.raises(ValueError):
+            _robust_modulus(x, 2.0)
+
+    def test_array_nonfinite_raises(self):
+        with np.errstate(invalid="ignore"), pytest.raises(ValueError):
+            _robust_modulus(np.array([0.5, np.inf]), 2.0)
+
 
 class Test_complex_to_polar:
     def test_deg(self):