JP-4323: Improve performance of tso median calculation in outlier detection (#10452)

Author: braingram

changes/10452.outlier_detection.rst

@@ -0,0 +1 @@
+Reduce memory usage and improve runtime for rolling median used for tso data. For a MIRI dataset including 2 files each with 40 integrations the memory usage was reduced by a factor of 7 and runtime improved by a factor of 3.

jwst/outlier_detection/tests/test_algorithms.py

@@ -1,15 +1,23 @@
 import numpy as np
+import pytest
 
 from jwst.outlier_detection.tso import moving_median_over_zeroth_axis
 
 
-def test_rolling_median():
+@pytest.mark.parametrize(
+    "w, expected_time_axis",
+    [
+        (3, [1, 1, 2, 3, 4, 5, 6, 7, 8, 8]),
+        (4, [1.5, 1.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, 7.5]),
+    ],
+)
+def test_rolling_median(w, expected_time_axis):
     time_axis = np.arange(10)
-    expected_time_axis = np.array([1, 1, 2, 3, 4, 5, 6, 7, 8, 8])
     spatial_axis = np.ones((5, 5))
     arr = time_axis[:, np.newaxis, np.newaxis] * spatial_axis[np.newaxis, :, :]
 
-    w = 3
     result = moving_median_over_zeroth_axis(arr, w)
-    expected = expected_time_axis[:, np.newaxis, np.newaxis] * spatial_axis[np.newaxis, :, :]
+    expected = (
+        np.array(expected_time_axis)[:, np.newaxis, np.newaxis] * spatial_axis[np.newaxis, :, :]
+    )
     assert np.allclose(result, expected)

jwst/outlier_detection/tso.py

@@ -163,11 +163,14 @@ def moving_median_over_zeroth_axis(x: np.ndarray, w: int) -> np.ndarray:
     """
     Calculate the median of a moving window over the zeroth axis of an N-d array.
 
-    Algorithm works by expanding the array into an additional dimension
-    where the new axis has the same length as the window size. Each entry in that
-    axis is a copy of the original array shifted by 1 with respect to the previous
-    entry, such that the rolling median is simply the median over the new axis.
-    modified from https://stackoverflow.com/a/71154394, see link for more details.
+    Slide a window of size w over the array along axis 0, and for each position,
+    calculate the median of the values inside that window (across axis 0 only).
+    The result at each step is stored in the center position of the window,
+    producing an output array with the same shape as the input.
+
+    Because the window cannot fully overlap the data at the beginning and end,
+    those edge positions are filled with the nearest computed median value to
+    avoid missing data.
 
     Parameters
     ----------
@@ -183,17 +186,15 @@ def moving_median_over_zeroth_axis(x: np.ndarray, w: int) -> np.ndarray:
     """
     if w <= 1:
         raise ValueError("Rolling median window size must be greater than 1.")
-    shifted = np.zeros((x.shape[0] + w - 1, w, *x.shape[1:])) * np.nan
-    for idx in range(w - 1):
-        shifted[idx : -w + idx + 1, idx] = x
-    shifted[idx + 1 :, idx + 1] = x
-    medians: np.ndarray = np.median(shifted, axis=1)
-    for idx in range(w - 1):
-        medians[idx] = np.median(shifted[idx, : idx + 1])
-        medians[-idx - 1] = np.median(shifted[-idx - 1, -idx - 1 :])
-    medians = medians[(w - 1) // 2 : -(w - 1) // 2]
-
+    out = np.full(x.shape, np.nan)
+    hw, odd_window = divmod(w, 2)
+    for start_index in range(x.shape[0] - w + 1):
+        end_index = start_index + w
+        np.median(x[start_index:end_index], axis=0, out=out[start_index + hw])
     # Fill in the edges with the nearest valid value
-    medians[: w // 2] = medians[w // 2]
-    medians[-w // 2 :] = medians[-w // 2]
-    return medians
+    out[:hw] = out[hw]
+    if odd_window:
+        out[-hw:] = out[-hw - 1]
+    else:
+        out[-hw + 1 :] = out[-hw]
+    return out