Compute window std dev, skew and kurtosis together

src/jabs/feature_extraction/feature_base_class.py

@@ -50,6 +50,14 @@ def _linear_detrend_numpy(segment: np.ndarray) -> np.ndarray:
     return segment - trend
 
 
+# window operations computed together by `Feature._shared_moment_features`
+_SHARED_MOMENT_OPERATIONS: dict[str, typing.Callable] = {
+    "std_dev": window_stats.window_std_dev,
+    "skew": window_stats.window_skew,
+    "kurtosis": window_stats.window_kurtosis,
+}
+
+
 class Feature(abc.ABC):
     """Abstract Base Class to define a common interface for classes that implement one or more related features"""
 
@@ -235,18 +243,54 @@ def _window_standard(self, window_size: int, per_frame_values: dict) -> dict:
         Returns:
             dict of window and signal processing feature values
         """
+        shared = self._shared_moment_features(window_size, per_frame_values)
         values = {}
         for op in self._window_operations:
-            values[op] = self._compute_window_feature(
-                per_frame_values,
-                window_size,
-                self._window_operations[op],
-            )
+            if op in shared:
+                values[op] = shared[op]
+            else:
+                values[op] = self._compute_window_feature(
+                    per_frame_values,
+                    window_size,
+                    self._window_operations[op],
+                )
         # Also include signal features
         signal_features = self._window_signal(window_size, per_frame_values)
         values.update(signal_features)
         return values
 
+    def _shared_moment_features(self, window_size: int, per_frame_values: dict) -> dict:
+        """Compute the standard deviation, skew, and kurtosis window features together.
+
+        These three share the window mean, centered values, and variance, so computing
+        them in one pass is faster than running each window operation separately. Only
+        operations still using the default implementations are computed here, so a
+        subclass that replaces one of them gets its own implementation.
+
+        Args:
+            window_size: number of frames (in each direction) to include in the window
+            per_frame_values: dict of per-frame values for this identity
+
+        Returns:
+            dict mapping each of "std_dev", "skew", and "kurtosis" that uses its default
+            implementation to a dict of per-frame feature name to window values
+        """
+        shared_ops = [
+            op
+            for op, default in _SHARED_MOMENT_OPERATIONS.items()
+            if self._window_operations.get(op) is default
+        ]
+        if not shared_ops:
+            return {}
+
+        results: dict[str, dict] = {op: {} for op in _SHARED_MOMENT_OPERATIONS}
+        for key, val in per_frame_values.items():
+            std_dev, skew, kurtosis = window_stats.window_std_skew_kurtosis(val, window_size)
+            results["std_dev"][key] = std_dev
+            results["skew"][key] = skew
+            results["kurtosis"][key] = kurtosis
+        return {op: results[op] for op in shared_ops}
+
     def _window_signal(self, window_size: int, per_frame_values: dict) -> dict:
         """The standard method for computing signal processing window features.
 

src/jabs/feature_extraction/window_operations/window_stats.py

@@ -95,6 +95,70 @@ def _center_windows(values: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
     return counts, values - mean[:, None]
 
 
+def _central_moment_terms(
+    values: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """Compute the terms shared by the window standard deviation, skew, and kurtosis.
+
+    Args:
+        values: 2d array with time on the first axis and the window on the second
+
+    Returns:
+        tuple of (number of non-nan values in each window, window values minus their
+        window mean, the square of those centered values, and the biased variance of
+        each window). nan values stay nan.
+    """
+    counts, centered = _center_windows(values)
+    squared = centered * centered
+    variance = np.nansum(squared, axis=1) / counts
+    return counts, centered, squared, variance
+
+
+def _skew_from_terms(
+    counts: np.ndarray, centered: np.ndarray, squared: np.ndarray, variance: np.ndarray
+) -> np.ndarray:
+    """Biased skew of each window from the terms of `_central_moment_terms`."""
+    return np.nansum(squared * centered, axis=1) / counts / variance**1.5
+
+
+def _kurtosis_from_terms(
+    counts: np.ndarray, squared: np.ndarray, variance: np.ndarray
+) -> np.ndarray:
+    """Pearson kurtosis of each window from the terms of `_central_moment_terms`."""
+    return np.nansum(squared * squared, axis=1) / counts / (variance * variance)
+
+
+def window_std_skew_kurtosis(
+    values: np.ndarray, window: int
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Calculates the masked standard deviation, skew, and kurtosis of each window at once.
+
+    Equal to calling `window_std_dev`, `window_skew`, and `window_kurtosis` separately,
+    but the window mean, centered values, and variance they share are computed only
+    once.
+
+    Args:
+        values: 1d np.ndarray of values
+        window: number of frames to include on each side of the center frame. The
+            window used for each output value spans ``2 * window + 1`` frames.
+
+    Returns:
+        tuple of sliding window standard deviation, skew, and kurtosis values
+    """
+    window_values = pad_sliding_window(values, window, pad_const=np.nan)
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=RuntimeWarning)
+        counts, centered, squared, variance = _central_moment_terms(window_values)
+        # np.nanstd only skips values that were nan in the input; a nan produced by
+        # inf - inf still makes the result nan. Sum the same way so windows containing
+        # infinities give the same result as window_std_dev.
+        input_nan = np.isnan(window_values)
+        std_dev = np.sqrt(np.sum(np.where(input_nan, 0.0, squared), axis=1) / counts)
+        skew = _skew_from_terms(counts, centered, squared, variance)
+        kurtosis = _kurtosis_from_terms(counts, squared, variance)
+    return std_dev, skew, kurtosis
+
+
 def np_kurtosis(values: np.ndarray) -> np.ndarray:
     """Calculates kurtosis in a rolling window faster than scipy
 
@@ -119,10 +183,8 @@ def np_kurtosis(values: np.ndarray) -> np.ndarray:
         numpy emits (does not raise) a RuntimeWarning when an entire window is nan.
         Callers such as `window_kurtosis` suppress that warning.
     """
-    counts, centered = _center_windows(values)
-    squared = centered * centered
-    variance = np.nansum(squared, axis=1) / counts
-    return np.nansum(squared * squared, axis=1) / counts / (variance * variance)
+    counts, _, squared, variance = _central_moment_terms(values)
+    return _kurtosis_from_terms(counts, squared, variance)
 
 
 def window_kurtosis(values: np.ndarray, window: int) -> np.ndarray:
@@ -160,10 +222,8 @@ def np_skew(values: np.ndarray) -> np.ndarray:
         numpy emits (does not raise) a RuntimeWarning when an entire window is nan.
         Callers such as `window_skew` suppress that warning.
     """
-    counts, centered = _center_windows(values)
-    squared = centered * centered
-    variance = np.nansum(squared, axis=1) / counts
-    return np.nansum(squared * centered, axis=1) / counts / variance**1.5
+    counts, centered, squared, variance = _central_moment_terms(values)
+    return _skew_from_terms(counts, centered, squared, variance)
 
 
 def window_skew(values: np.ndarray, window: int) -> np.ndarray:

tests/feature_extraction/test_feature_base_class.py

@@ -1,5 +1,6 @@
 """Tests for module-level helpers in feature_base_class."""
 
+import typing
 import warnings
 from types import SimpleNamespace
 
@@ -9,6 +10,7 @@
 from scipy.signal import detrend
 
 from jabs.feature_extraction.feature_base_class import Feature, _linear_detrend_numpy
+from jabs.feature_extraction.window_operations import window_stats
 
 
 class TestLinearDetrendNumpy:
@@ -278,3 +280,61 @@ def test_circular_window_features_edge_cases() -> None:
     result = Feature._compute_window_features_circular(feature, {"a": values}, mask, 1)
     # window around the last frame holds 170 and -170, whose circular mean is +-180
     assert abs(result["mean"]["a"][5]) == pytest.approx(180.0)
+
+
+class _StandardStubFeature(Feature):
+    """Minimal non-circular Feature for exercising `Feature.window` without a pose file."""
+
+    _name = "standard_stub"
+
+    def per_frame(self, identity: int) -> dict[str, np.ndarray]:
+        """Not used: the tests pass per-frame values to `window` directly."""
+        return {}
+
+
+def _standard_window_inputs() -> tuple[SimpleNamespace, dict[str, np.ndarray]]:
+    """Fake pose data and two per-frame columns with gaps."""
+    rng = np.random.default_rng(seed=5)
+    values = {"a": rng.normal(10.0, 3.0, 80), "b": rng.normal(-5.0, 1.0, 80)}
+    values["a"][20:26] = np.nan
+    return SimpleNamespace(fps=30), values
+
+
+def test_window_std_skew_kurtosis_computed_together_match_separate_ops() -> None:
+    """The shared std/skew/kurtosis path gives the window ops' own results, in op order."""
+    poses, values = _standard_window_inputs()
+
+    with np.errstate(all="ignore"), warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        result = _StandardStubFeature(poses, 1.0).window(0, 3, values)
+        expected = {
+            op: {key: fn(val, 3) for key, val in values.items()}
+            for op, fn in Feature._window_operations.items()
+        }
+
+    assert list(result)[: len(expected)] == list(expected)
+    for op, columns in expected.items():
+        for key, want in columns.items():
+            np.testing.assert_array_equal(result[op][key], want, err_msg=f"{op} {key}")
+
+
+def test_overridden_moment_operation_is_used() -> None:
+    """A subclass that replaces one of the shared ops gets its own implementation."""
+
+    def constant_skew(values: np.ndarray, window: int) -> np.ndarray:
+        return np.full(values.shape, 42.0)
+
+    class _CustomSkewFeature(_StandardStubFeature):
+        _window_operations: typing.ClassVar[dict[str, typing.Callable]] = {
+            **Feature._window_operations,
+            "skew": constant_skew,
+        }
+
+    poses, values = _standard_window_inputs()
+    with np.errstate(all="ignore"), warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        result = _CustomSkewFeature(poses, 1.0).window(0, 3, values)
+        expected_kurtosis = window_stats.window_kurtosis(values["b"], 3)
+
+    np.testing.assert_array_equal(result["skew"]["a"], np.full(80, 42.0))
+    np.testing.assert_array_equal(result["kurtosis"]["b"], expected_kurtosis)

tests/feature_extraction/window_operations/test_window_stats.py

@@ -165,3 +165,33 @@ def test_window_ops_match_per_window_reference(
         expected = _sliding_window_oracle(values, window, oracle_op)
 
     np.testing.assert_allclose(result, expected, rtol=1e-9, atol=1e-9, equal_nan=True)
+
+
+def _infinity_test_series() -> dict[str, np.ndarray]:
+    """Series with infinities, which produce nan via inf - inf when windows are centered."""
+    rng = np.random.default_rng(seed=11)
+    mixed = rng.normal(0.0, 20.0, 60)
+    mixed[[5, 17, 30]] = np.inf
+    mixed[[22, 41]] = -np.inf
+    with_nan = mixed.copy()
+    with_nan[rng.random(60) < 0.2] = np.nan
+    return {"all_inf": np.full(60, np.inf), "mixed_inf": mixed, "inf_and_nan": with_nan}
+
+
+@pytest.mark.parametrize("window", [1, 3, 10], ids=["w1", "w3", "w10"])
+@pytest.mark.parametrize("case", list(_window_test_series()) + list(_infinity_test_series()))
+def test_window_std_skew_kurtosis_matches_separate_ops(case: str, window: int) -> None:
+    """The combined op is bit-identical to running the three window ops separately."""
+    values = {**_window_test_series(), **_infinity_test_series()}[case]
+
+    with np.errstate(all="ignore"), warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        std_dev, skew, kurtosis = window_stats.window_std_skew_kurtosis(values, window)
+        expected = (
+            window_stats.window_std_dev(values, window),
+            window_stats.window_skew(values, window),
+            window_stats.window_kurtosis(values, window),
+        )
+
+    for actual, want in zip((std_dev, skew, kurtosis), expected, strict=True):
+        np.testing.assert_array_equal(actual, want)