Merge pull request #27 from Arcascope/grid-every-window

senpy 4.2.0: transform every window with at least min_samples samples

Author: ericcanton

README.md

@@ -40,7 +40,8 @@ window grid, and they all put the same samples in the same windows:
 
 * Window `k` spans `[origin + k*hop, origin + k*hop + window_s)`, with `hop = window_s - overlap_s`.
   Starts are `k*hop`, never an accumulated sum. Samples before the origin belong to no window.
-* The grid runs while a window ends no later than one median sample period past the last sample.
+* The grid runs through the last window holding at least `min_samples` samples, so the trailing
+  windows the recording stops partway through are transformed like any other window with data.
 * `times` are window centres, `k*hop + window_s/2`, measured from the origin.
 
 `senpy.window_grid(timestamps, window_s, overlap_s, origin_s=..., min_samples=...)` returns that
@@ -246,11 +247,10 @@ band at 30 s windows costs about 150 000 multiply-accumulates per second of stre
   subwindow straddles a window edge; one that did could not be shared by the windows either side.
 * **Sample rate.** Supplied rather than measured: it sets the magnitude scale and the grid size.
   `compute_nustft` takes the median spacing over the whole recording, which a stream cannot see.
-* **The trailing window.** `push` reports only windows the stream has passed the end of, which is 
-  all a live stream can honestly say. `compute_nustft` knows where the recording stops and also 
-  emits a final window ending within one sample period of the last timestamp; that one comes out 
-  of `flush()`. `compute_nustft_streaming` applies this rule for you and is the function to compare
-  the two paths with.
+* **The trailing windows.** `push` reports only windows the stream has passed the end of, which is
+  all a live stream can honestly say. The windows the data stops partway through come out of
+  `flush()`, as `compute_nustft` reports them. `compute_nustft_streaming` does both and is the
+  function to compare the two paths with.
 * **The Nyquist bin** (present only when `fmax` is unset) is the true $+N/2$ coefficient.
   `compute_nustft` reports its conjugate there, an artifact of reading that bin out of the aliased
   FINUFFT mode. Magnitudes are identical.

senpy/_grid.py

@@ -4,9 +4,9 @@
 on the timestamps' own clock, in seconds. Its start is computed by
 multiplication, never by accumulating ``hop``, so every backend -- the C++
 batch transform, the streaming transform, the JAX transform and the JAX packer
--- puts the same sample in the same window. The grid runs while a window ends
-no later than one median sample period past the last timestamp, the bound the
-batch transforms have always used.
+-- puts the same sample in the same window. The grid runs through the last
+window holding at least ``min_samples`` samples: every window with enough data
+is transformed, at the end of a recording as anywhere else.
 
 Timestamps are measured from the origin once, by :func:`relative_seconds`,
 and every backend is handed those relative times. That is what makes the
@@ -137,22 +137,38 @@ def median_spacing(t_s: NDArray[np.float64]) -> float:
     return float(positive[positive.size // 2])
 
 
-def window_count(last_relative_s: float, dt_median_s: float, window_s: float, hop_s: float) -> int:
-    """How many windows fit: those with ``k * hop + window <= last + dt_median``.
+def last_window_start(last_relative_s: float, hop_s: float) -> int:
+    """Index of the last window starting at or before the last sample; -1 if none.
 
-    ``last_relative_s`` is the last timestamp minus the origin. The estimate is
-    corrected against the exact comparison, so the count agrees with a loop
-    that tests each window in turn.
+    Windows starting later hold no samples. The estimate is corrected against
+    the exact comparison, so it agrees with a loop that tests each start.
     """
-    limit = last_relative_s + dt_median_s
-    if not math.isfinite(limit) or window_s > limit:
-        return 0
-    n = int(math.floor((limit - window_s) / hop_s)) + 1
-    while n > 0 and (n - 1) * hop_s + window_s > limit:
-        n -= 1
-    while n * hop_s + window_s <= limit:
-        n += 1
-    return n
+    if not math.isfinite(last_relative_s) or last_relative_s < 0.0:
+        return -1
+    k = int(math.floor(last_relative_s / hop_s))
+    while k >= 0 and k * hop_s > last_relative_s:
+        k -= 1
+    while (k + 1) * hop_s <= last_relative_s:
+        k += 1
+    return k
+
+
+def window_count(
+    t_relative_s: NDArray[np.float64], window_s: float, hop_s: float, min_samples: int
+) -> int:
+    """How many windows the grid has: through the last one with ``min_samples`` samples.
+
+    ``t_relative_s`` are sorted timestamps minus the origin. Windows before
+    that last one stay on the grid whatever they hold; with
+    ``empty_windows="keep"`` those without enough samples are NaN rows.
+    """
+    t = np.asarray(t_relative_s, dtype=np.float64)
+    starts = np.arange(last_window_start(float(t[-1]), hop_s) + 1, dtype=np.int64) * hop_s
+    counts = np.searchsorted(t, starts + window_s, side="left") - np.searchsorted(
+        t, starts, side="left"
+    )
+    enough = np.flatnonzero(counts >= min_samples)
+    return int(enough[-1]) + 1 if enough.size else 0
 
 
 @dataclass(frozen=True)
@@ -164,7 +180,7 @@ class WindowGrid:
             seconds. A result's ``times`` are measured from it.
         window_s: Window duration in seconds.
         hop_s: Window hop in seconds (``window_s - overlap_s``).
-        dt_median_s: Upper median sample spacing, which sets the grid's end.
+        dt_median_s: Upper median sample spacing, which sets the transform size.
         first: Index of each window's first sample, shape ``(n_windows,)``.
         stop: One past each window's last sample, shape ``(n_windows,)``.
         min_samples: Fewest samples a window needs to be transformed.
@@ -263,7 +279,7 @@ def grid_from_relative(
     min_samples: int,
 ) -> WindowGrid:
     """Build the grid from timestamps already expressed relative to the origin."""
-    n = window_count(float(t_relative_s[-1]), dt_median_s, window_s, hop_s)
+    n = window_count(t_relative_s, window_s, hop_s, min_samples)
     starts = np.arange(n, dtype=np.int64) * hop_s
     first = np.searchsorted(t_relative_s, starts, side="left").astype(np.int64)
     stop = np.searchsorted(t_relative_s, starts + window_s, side="left").astype(np.int64)

senpy/_version.py

@@ -1 +1 @@
-__version__ = "4.1.0"
+__version__ = "4.2.0"

senpy/api.py

@@ -1075,11 +1075,9 @@ def compute_nustft_streaming(
     windows: List[StreamingWindow] = []
     for start in range(0, len(t), chunk):
         windows.extend(transform.push(t[start : start + chunk], values[start : start + chunk]))
-    # push() only reports windows the stream has passed the end of. compute_nustft, which sees
-    # where the recording stops, also emits a final window that ends within one sample period
-    # of the last timestamp; take that one out of the flush and drop anything past the grid.
+    # push() only reports windows the stream has passed the end of; the trailing windows the
+    # data stops partway through come out of the flush.
     windows.extend(transform.flush())
-    windows = [w for w in windows if w.index < grid.n_windows]
 
     if grid.n_windows == 0 or (empty_windows == "drop" and not windows):
         raise ValueError("compute_nustft_streaming requires enough data for at least one window")

senpy/jax_backend.py

@@ -34,13 +34,13 @@
 from ._grid import (
     DEFAULT_MIN_SAMPLES,
     OriginSpec,
+    last_window_start,
     median_spacing,
     relative_seconds,
     resolve_empty_windows,
     resolve_origin,
     timestamp_scale,
     validate_min_samples,
-    window_count,
 )
 
 
@@ -794,10 +794,9 @@ def _compute_nustft(
     n_pos_freqs = nfft_padded // 2 + 1
 
     # The same grid senpy._grid.window_grid builds on the host: window k starts
-    # at k * hop after the origin.
+    # at k * hop after the origin, through the last window with min_samples.
     t_end = float(host_t[-1]) if host_t is not None else float(jax.device_get(t[-1]))
-    n_windows = window_count(t_end, dt_median, window_s, hop_s)
-    starts = np.arange(n_windows, dtype=np.int64) * hop_s
+    starts = np.arange(last_window_start(t_end, hop_s) + 1, dtype=np.int64) * hop_s
 
     complex_dtype = jnp.result_type(s, 1j)
     frequencies = jnp.arange(n_pos_freqs, dtype=t.dtype) / window_s
@@ -817,7 +816,7 @@ def empty_result() -> JaxNUSTFTResult:
             origin_s=origin,
         )
 
-    if n_windows == 0:
+    if starts.size == 0:
         return empty_result()
 
     if host_t is not None:
@@ -829,6 +828,13 @@ def empty_result() -> JaxNUSTFTResult:
         end_indices = np.asarray(
             jax.device_get(jnp.searchsorted(t, starts_device + window_s, side="left"))
         )
+    enough = np.flatnonzero(end_indices - start_indices >= min_samples)
+    n_windows = int(enough[-1]) + 1 if enough.size else 0
+    if n_windows == 0:
+        return empty_result()
+    starts = starts[:n_windows]
+    start_indices = start_indices[:n_windows]
+    end_indices = end_indices[:n_windows]
     grid_counts = (end_indices - start_indices).astype(np.int64)
     mode_indices = jnp.concatenate(
         (jnp.arange(nfft_padded // 2, nfft_padded), jnp.array([0]))

senpy/tests/test_jax_api.py

@@ -64,9 +64,10 @@ def test_compute_nustft_returns_device_arrays_with_known_tone():
     )
 
     assert result.axis_order == "time_frequency"
-    assert result.coefficients.shape == (3, 33)
+    # 8 s of data, 4 s windows every 2 s: the one starting at 6 s holds the last 2 s.
+    assert result.coefficients.shape == (4, 33)
     assert result.frequencies.shape == (33,)
-    assert result.times.shape == (3,)
+    assert result.times.shape == (4,)
     peak = result.frequencies[jnp.argmax(jnp.mean(result.power, axis=0))]
     assert float(peak) == pytest.approx(3.0)
 
@@ -82,7 +83,7 @@ def test_compute_nufft_spectrogram_and_welch_stay_jax_backed():
         timestamps, signal, window_s=4.0, overlap_s=2.0
     )
 
-    assert spectrogram.Sxx.shape == (3, 129)
+    assert spectrogram.Sxx.shape == (4, 129)
     assert frequencies.shape == welch.shape == (129,)
 
 

senpy/tests/test_spectral_api.py

@@ -276,4 +276,4 @@ def test_removed_nufft_alias_and_uniform_compatibility_alias_warns():
 
 
 def test_version_is_consolidated():
-    assert senpy.__version__ == sp.__version__ == "4.1.0"
+    assert senpy.__version__ == sp.__version__ == "4.2.0"

senpy/tests/test_streaming_nustft.py

@@ -57,14 +57,10 @@ def batch_equivalent(t, signal, chunk=137, **kwargs):
     """Streams the whole array and returns the windows ``compute_nustft`` would report.
 
     ``push`` only reports windows the stream has passed the end of, which is all a live stream
-    can honestly say. ``compute_nustft`` knows where the recording stops, so it also emits a
-    final window ending within one sample period of the last timestamp; that one comes out of
-    the flush.
+    can honestly say. The windows the data stops partway through come out of the flush.
     """
     transform, windows = stream_through(t, signal, chunk=chunk, **kwargs)
-    limit = float(t[-1]) + 1.0 / FS
-    windows = windows + [w for w in transform.flush() if w.start + WINDOW_S <= limit]
-    return transform, windows
+    return transform, windows + transform.flush()
 
 
 def relative_deviation(actual, expected):
@@ -144,15 +140,17 @@ def test_dropouts_leave_holes_rather_than_shifting_windows():
     assert windows[indices.index(3)].sample_count < 3000
 
 
-def test_flush_reports_the_trailing_window_compute_nustft_stops_short_of():
+def test_flush_reports_the_trailing_windows_push_cannot():
     t, signal = strap_stream(seconds=125, gaps=())
     transform, windows = stream_through(t, signal)
     reference = senpy.compute_nustft(t, signal, window_s=WINDOW_S, overlap_s=0.0)
 
-    assert len(windows) == len(reference.times)
+    # Windows 0-3 are complete; window 4 (120-150 s) holds the last 5 s.
+    assert [w.index for w in windows] == [0, 1, 2, 3]
     tail = transform.flush()
     assert [w.index for w in tail] == [4]
     assert tail[0].sample_count == 500
+    assert len(reference.times) == 5
 
 
 def test_convenience_wrapper_reproduces_compute_nustft():

senpy/tests/test_window_grid.py

@@ -56,14 +56,41 @@ def test_grid_starts_are_whole_hops_after_the_origin():
         assert inside.max() < grid.starts[k] + WINDOW_S
 
 
-def test_grid_ends_where_the_batch_transforms_always_have():
-    # Evenly spaced samples: the last window may end at most one sample past the data.
+def test_grid_runs_through_the_last_window_with_enough_samples():
+    # Evenly spaced samples every 0.5 s up to 59.5 s: the window starting at 58 s
+    # holds 4 samples (58.0-59.5), the one at 60 s none.
     t = np.arange(0.0, 60.0, 0.5)
     grid = senpy.window_grid(t, WINDOW_S, OVERLAP_S)
-    last_end = grid.starts[-1] + WINDOW_S
 
-    assert last_end <= t[-1] + 0.5
-    assert last_end + HOP_S > t[-1] + 0.5
+    assert grid.starts[-1] == 58.0
+    assert grid.sample_count[-1] == 4 and grid.valid[-1]
+    # The windows the data stops partway through are on the grid and valid.
+    assert grid.valid[grid.starts + WINDOW_S > t[-1]].all()
+    assert senpy.window_grid(t, WINDOW_S, OVERLAP_S, min_samples=5).starts[-1] == 56.0
+
+
+def test_trailing_sparse_windows_end_the_grid_but_interior_ones_stay():
+    t = np.concatenate([np.arange(0.0, 20.0, 0.5), [40.0, 59.0]])
+    grid = senpy.window_grid(t, WINDOW_S, OVERLAP_S)
+
+    # The last window with 4 samples starts at 18 s; windows up to it that hold
+    # fewer stay on the grid as invalid, the sparse ones after it are not on it.
+    assert grid.starts[-1] == 18.0
+    assert grid.valid.all()
+    kept = senpy.window_grid(np.concatenate([t, np.arange(80.0, 82.0, 0.5)]), WINDOW_S, OVERLAP_S)
+    assert kept.starts[-1] == 80.0
+    assert not kept.valid[(kept.starts > 18.0) & (kept.starts < 72.0)].any()
+
+
+def test_the_end_of_the_grid_does_not_hinge_on_rounding():
+    # np.arange drift leaves the last sample ~2e-10 s short of 239.98 s. The
+    # window 230-240 s is full either way, and the grid runs on to the window
+    # starting at 238 s, which holds the last 2 s.
+    t = np.arange(7.0, 240.0, 1 / 50)
+    grid = senpy.window_grid(t, WINDOW_S, OVERLAP_S, origin_s=0.0)
+
+    assert grid.starts[-1] == 238.0
+    assert grid.sample_count[grid.starts == 230.0][0] >= 499
 
 
 def test_dropout_windows_are_counted_not_lost():
@@ -347,13 +374,12 @@ def test_jax_returns_zero_rows_when_there_is_nothing_to_report(kwargs):
     assert result.window_index.size == result.sample_count.size == result.valid.size == 0
 
 
-def test_keep_mode_with_every_window_too_sparse_is_all_nan_not_an_error():
+def test_keep_mode_with_every_window_too_sparse_has_nothing_to_report():
     t, x = recording()
-    result = nustft(t, x, empty_windows="keep", min_samples=100_000)
 
-    assert result.valid.size == senpy.window_grid(t, WINDOW_S, OVERLAP_S).n_windows
-    assert not result.valid.any()
-    assert np.isnan(result.coefficients).all()
+    assert senpy.window_grid(t, WINDOW_S, OVERLAP_S, min_samples=100_000).n_windows == 0
+    with pytest.raises(ValueError, match="at least one window"):
+        nustft(t, x, empty_windows="keep", min_samples=100_000)
 
 
 # ── the 5.0 default change ──────────────────────────────────────────

src/sensor_processing_native.cpp

@@ -317,8 +317,9 @@ class SensorProcessor
     // Window k spans [origin + k*hop, origin + k*hop + secperseg). origin defaults (NaN) to the
     // first timestamp; samples before it belong to no window. Starts are k*hop, never an
     // accumulated sum, so senpy/_grid.py, the streaming transform and the JAX backend place every
-    // sample in the same window. Windows with fewer than min_samples samples are not transformed;
-    // grid_sample_counts still reports them, so callers can rebuild the full grid.
+    // sample in the same window. The grid runs through the last window with at least min_samples
+    // samples. Windows before it with fewer are not transformed; grid_sample_counts still reports
+    // them, so callers can rebuild the full grid.
     static NUSTFTResult computeNUSTFT(
         const std::vector<double> &timestamps,
         const std::vector<double> &signal,
@@ -406,7 +407,9 @@ class SensorProcessor
         std::vector<double> t_rel(n);
         for (size_t i = 0; i < n; ++i)
             t_rel[i] = timestamps[i] - origin;
-        const double limit = t_rel[n - 1] + dt_median;
+        // Windows starting after the last sample hold none; the grid then ends at the last
+        // window with min_samples samples (trimmed below).
+        const double t_last = t_rel[n - 1];
 
         // Use two-pointer indices to avoid rescanning timestamps for every window
         size_t start_idx = 0;
@@ -415,7 +418,7 @@ class SensorProcessor
         {
             const double win_start = static_cast<double>(w) * hop_dur;
             const double win_end = win_start + win_dur;
-            if (!(win_end <= limit))
+            if (!(win_start <= t_last))
                 break;
 
             // Advance start_idx to the first sample >= win_start
@@ -533,6 +536,9 @@ class SensorProcessor
             window_indices.push_back(w);
         }
 
+        // The grid ends at the last window with enough data to transform.
+        grid_sample_counts.resize(window_indices.empty() ? 0 : static_cast<size_t>(window_indices.back()) + 1);
+
         // Build result
         NUSTFTResult result;
         result.freqs = freqs;
@@ -1668,8 +1674,6 @@ class StreamingNUSTFT
     }
 
     // Emits every window the stream has passed the end of, oldest first, and forgets it.
-    // The readiness test matches computeNUSTFT's loop bound, which admits a final window whose
-    // end is within one sample period of the last timestamp.
     std::vector<StreamingNUSTFTWindow> drain()
     {
         std::vector<StreamingNUSTFTWindow> ready;
@@ -1680,8 +1684,8 @@ class StreamingNUSTFT
         // input, no sample belonging to it can arrive after that. Reporting a window even one
         // sample period early would strand the samples that a packet whose clock drifted
         // backwards still owes it -- and computeNUSTFT, which sees the whole recording, would
-        // have counted them. The trailing window computeNUSTFT admits when it is within one
-        // sample period of complete comes out of flush() instead.
+        // have counted them. The trailing windows the data stops partway through come out of
+        // flush() instead.
         const double reach = newest_ - origin_ - win_dur_;
         if (reach < 0.0)
             return ready;
@@ -1703,8 +1707,8 @@ class StreamingNUSTFT
         return ready;
     }
 
-    // Emits everything still open, however partial, and resets to empty. Unlike drain(), this
-    // reports the trailing window computeNUSTFT would have stopped short of.
+    // Emits everything still open, however partial, and resets to empty: the trailing windows
+    // the data stops partway through, which computeNUSTFT also reports.
     std::vector<StreamingNUSTFTWindow> flush()
     {
         std::vector<StreamingNUSTFTWindow> ready;