Keep reconstructed timestamps monotone across segment joins

src/ezmsg/nwb/clockmodel.py

@@ -86,6 +86,22 @@
 across 51 s on a clean stream that self-fit to 0.08 ms). So the sibling's clock
 is borrowed only when the target is genuinely too corrupted to trust its own."""
 
+SHARED_MODEL_JITTER_FLOOR_PERIODS = 1.0
+"""Absolute jitter floor for the shared model, in sample periods: a target must
+also be jittier than this before a sibling's clock is borrowed.
+
+Serves two purposes. It keeps the ratio test meaningful when the cleanest member
+is *perfect* -- a stream whose stamps were regenerated as ``i / fs`` self-fits to
+~1e-14 s (or exactly zero), and a pure ratio test against that either rescues
+everything or, with a ``> 0`` guard, rescues nothing precisely when the reference
+is most trustworthy. And it stops a large *relative* difference between two
+already-clean streams from triggering a rescue that cannot help: sub-sample
+jitter is not corruption, whatever its ratio to a cleaner sibling.
+
+One period separates the observed populations with room to spare: clean
+recordings self-fit to 0.3-0.5 periods, while a stream whose stamps were written
+at the wrong intra-chunk rate lands at ~37."""
+
 CACHE_DIR = Path("~").expanduser() / ".ezmsg" / "nwb-cache" / "dejitter"
 """Where reconstructed timestamp vectors are cached across file opens, alongside
 remfile's ``nwb-cache``. Multi-pass workloads (e.g. training over many epochs of
@@ -265,6 +281,50 @@ def _smooth_eval(x: np.ndarray, y: np.ndarray, n_knots: int, method: str) -> np.
     return fit_map(x, y, n_knots, method)(x)
 
 
+def _enforce_join_monotonicity(
+    out: np.ndarray, segments: list[tuple[int, int]], period: float, raw: np.ndarray
+) -> np.ndarray:
+    """Shift whole segments forward so no segment boundary runs backwards.
+
+    Each segment is fit independently so a real gap survives, which leaves the
+    *joins* unconstrained: a segment's smoothed first sample is an estimate, and
+    in the shared-model path each segment also carries its own median anchor, so
+    segment ``k+1`` can begin before segment ``k`` ended. Observed on real data
+    at 33 sample periods (1.1 ms) of backward step -- far past rounding, and a
+    contract violation for callers (the gap-splitter and chunk anchoring both
+    read these timestamps as monotone).
+
+    A broken join is re-cut to the jump the *raw* timestamps show across it
+    (floored at one period), and the whole following segment is shifted to match,
+    cumulatively. Using the raw jump is what keeps a real gap intact: shifting by
+    just enough to restore order would collapse the gap to a single period, and
+    the raw stamps are the only remaining evidence of how long the recording
+    actually stopped. Healthy joins are left alone.
+
+    Shifting rather than clamping is what preserves the fit: a segment keeps its
+    internal shape and its duration. ``np.maximum.accumulate`` would instead
+    flatten the overlap into a plateau of duplicate stamps, discarding the
+    segment's own timing.
+
+    Shifts are cumulative, so a late sample can move by the sum of all earlier
+    corrections. That is the honest cost, and it is paid only where the
+    reconstruction had already gone wrong.
+    """
+    if len(segments) < 2:
+        return out
+    shift = 0.0
+    for a, b in segments[1:]:
+        if shift:
+            out[a:b] += shift
+        jump = out[a] - out[a - 1]
+        if jump < period:
+            want = max(float(raw[a] - raw[a - 1]), period)
+            extra = want - jump
+            out[a:b] += extra
+            shift += extra
+    return out
+
+
 # --- Reconstruction -------------------------------------------------------
 
 
@@ -292,12 +352,13 @@ def reconstruct_self(
     if gaps.size == 0:
         return base
     out = base.copy()
-    for a, b in _segments(n, gaps):
+    segs = _segments(n, gaps)
+    for a, b in segs:
         if b - a >= 2:
             out[a:b] = _smooth_eval(np.arange(b - a, dtype=float), ds[a:b], n_knots, method)
         else:
             out[a:b] = ds[a:b]
-    return out
+    return _enforce_join_monotonicity(out, segs, _nominal_period(ds), ds)
 
 
 def reconstruct_shared(
@@ -343,10 +404,14 @@ def through_c(dev_clean: np.ndarray, tgt_seg: np.ndarray | None) -> np.ndarray:
     if gaps.size == 0:
         return through_c(dev_base, tgt)
     out = np.empty(n, dtype=float)
-    for a, b in _segments(n, gaps):
+    segs = _segments(n, gaps)
+    for a, b in segs:
         dev_clean = _smooth_eval(np.arange(b - a, dtype=float), dev[a:b], n_knots, method) if b - a >= 2 else dev[a:b]
         out[a:b] = through_c(dev_clean, None if tgt is None else tgt[a:b])
-    return out
+    # Each segment carries its own median anchor, so joins are doubly
+    # unconstrained here -- this is the path where the 33-period backward step
+    # was observed.
+    return _enforce_join_monotonicity(out, segs, _nominal_period(dev), dev)
 
 
 def _detect_gaps(times: np.ndarray, gap_threshold_s: float | None) -> np.ndarray:
@@ -374,11 +439,13 @@ def reconstruct_group(
     may be ``None`` for a lone stream). The cleanest member -- least jitter about
     its own self-fit -- is dejittered directly and can supply the clock model
     ``C`` for the rest. A member is reconstructed via that shared model only when
-    it is at least :data:`SHARED_MODEL_JITTER_RATIO` times jitterier than the
-    clean source (and carries device timestamps to map through); otherwise it is
+    it carries device timestamps and is jitterier than both
+    :data:`SHARED_MODEL_JITTER_RATIO` times the clean source and
+    :data:`SHARED_MODEL_JITTER_FLOOR_PERIODS` sample periods; otherwise it is
     self-fit, which never imposes a sibling's clock. A single-member group falls
     back to a self-fit. Genuine data gaps are preserved per member
-    (``gap_threshold_s`` forwarded; ``None`` = auto). Returns
+    (``gap_threshold_s`` forwarded; ``None`` = auto), and segment joins are kept
+    monotone (:func:`_enforce_join_monotonicity`). Returns
     ``{key: reconstructed_dataset_times}``.
 
     This is the file-agnostic entry point shared by the read-time slicer and any
@@ -402,12 +469,16 @@ def reconstruct_group(
             continue
         # Borrow the clean sibling's clock only for a genuinely corrupted target;
         # an already-clean stream self-fits to sub-sample accuracy and the shared
-        # model would only inject the sibling's slightly-different rate.
-        rescue = (
-            m["device"] is not None
-            and clean_resid > 0.0
-            and resids[m["key"]] >= SHARED_MODEL_JITTER_RATIO * clean_resid
+        # model would only inject the sibling's slightly-different rate. The
+        # target must clear both bars -- jitterier than the source by the ratio,
+        # AND jittery in absolute terms -- so a perfect source (ratio threshold
+        # 0) still rescues a corrupted sibling, while two clean streams never
+        # trigger each other however far apart their ratio happens to fall.
+        threshold = max(
+            SHARED_MODEL_JITTER_RATIO * clean_resid,
+            SHARED_MODEL_JITTER_FLOOR_PERIODS * _nominal_period(np.asarray(m["dataset"], dtype=float)),
         )
+        rescue = m["device"] is not None and resids[m["key"]] >= threshold
         if rescue:
             out[m["key"]] = reconstruct_shared(
                 target_device=m["device"],

tests/test_clockmodel.py

@@ -5,6 +5,7 @@
 
 from ezmsg.nwb.clockmodel import (
     ClockModel,
+    _enforce_join_monotonicity,
     cache_lookup,
     cache_store,
     find_real_gaps,
@@ -229,6 +230,114 @@ def test_reconstruct_shared_preserves_gap_on_device_clock():
     assert abs(np.diff(recon)[big[0]] - 0.3) < 0.02
 
 
+# --- segment joins ---
+
+
+def test_enforce_join_monotonicity_recuts_a_backward_join():
+    """A join that runs backwards is re-cut to the jump the raw stamps show."""
+    period = GAIN
+    raw = np.array([0.0, 0.001, 0.002, 0.012, 0.013])  # 10 ms real gap at index 3
+    out = np.array([0.0, 0.001, 0.002, 0.0015, 0.0025])  # segment 2 anchored 0.5 ms early
+    fixed = _enforce_join_monotonicity(out.copy(), [(0, 3), (3, 5)], period, raw)
+
+    assert np.all(np.diff(fixed) > 0)
+    # The gap is restored from the raw stamps, not collapsed to one period.
+    assert fixed[3] - fixed[2] == pytest.approx(0.010)
+    # The moved segment keeps its own shape.
+    assert np.allclose(np.diff(fixed[3:]), np.diff(out[3:]))
+
+
+def test_enforce_join_monotonicity_leaves_healthy_joins_alone():
+    period = GAIN
+    raw = np.array([0.0, 0.001, 0.002, 0.012, 0.013])
+    out = np.array([0.0, 0.001, 0.002, 0.012, 0.013])
+    fixed = _enforce_join_monotonicity(out.copy(), [(0, 3), (3, 5)], period, raw)
+    assert np.array_equal(fixed, out)
+
+
+def test_enforce_join_monotonicity_shifts_cumulatively():
+    """Two bad joins: the third segment carries both corrections."""
+    period = GAIN
+    raw = np.array([0.0, 0.001, 0.011, 0.012, 0.022, 0.023])
+    out = np.array([0.0, 0.001, 0.0005, 0.0015, 0.0010, 0.0020])
+    fixed = _enforce_join_monotonicity(out.copy(), [(0, 2), (2, 4), (4, 6)], period, raw)
+
+    assert np.all(np.diff(fixed) > 0)
+    assert fixed[2] - fixed[1] == pytest.approx(0.010)
+    assert fixed[4] - fixed[3] == pytest.approx(0.010)
+
+
+def test_reconstruct_shared_join_stays_monotone_when_anchors_disagree():
+    """Per-segment median anchors must not let a segment start before the last ended.
+
+    Regression from real data: each segment of the shared path re-anchors on its
+    own median, so a target whose latency shifts across a gap produced a 33
+    sample-period backward step in the reconstructed output -- which callers read
+    as monotone.
+    """
+    truth = _smooth_truth()
+    epoch = 1.7e9
+    clean_device = epoch + np.arange(N) / RATE
+    target_device = epoch + np.arange(N) / RATE
+    target_device[3000:] += 0.005  # small real gap: 5 ms
+
+    target_dataset = _jitter(truth, seed=5)
+    target_dataset[3000:] += 0.005
+    target_dataset[3000:] -= 0.010  # latency shift larger than the gap
+
+    recon = reconstruct_shared(target_device, clean_device, truth, target_dataset=target_dataset, gap_threshold_s=0.002)
+    assert np.all(np.diff(recon) >= 0), "reconstruction must not run backwards across segment joins"
+    # The device clock's jump survives the repair rather than collapsing to one
+    # period: 5 ms of missing data plus the sample period that spans it.
+    assert np.diff(recon)[2999] == pytest.approx(0.005 + GAIN, abs=5e-4)
+
+
+def test_reconstruct_self_stays_monotone_across_gaps():
+    truth = _smooth_truth()
+    gapped_truth = truth.copy()
+    gapped_truth[2500:] += 0.5
+    recon = reconstruct_self(_jitter(gapped_truth), gap_threshold_s=0.05)
+    assert np.all(np.diff(recon) >= 0)
+
+
+# --- shared-model gate ---
+
+
+def test_shared_model_rescues_when_the_clean_member_is_perfect():
+    """A flawless reference must not disable the rescue.
+
+    Regenerated timestamps (``i / fs``) self-fit to ~0, so a pure ratio test
+    against them either rescues everything or, guarded with ``> 0``, rescues
+    nothing -- exactly when the reference is most worth borrowing.
+    """
+    epoch = 1.7e9
+    device = epoch + np.arange(N) / RATE
+    perfect = np.arange(N) / RATE
+    members = [
+        {"key": "perfect", "device": device.copy(), "dataset": perfect},
+        {"key": "broken", "device": device.copy(), "dataset": _jitter(perfect, scale_periods=8.0, seed=4)},
+    ]
+    out = reconstruct_group(members)
+    # The broken member is pulled back onto the shared clock, far inside its own
+    # 8-period jitter.
+    assert np.abs(out["broken"] - perfect).max() < 2.0 * GAIN
+
+
+def test_shared_model_not_triggered_by_subsample_jitter_ratio():
+    """A big ratio between two sub-sample-clean streams is not corruption."""
+    truth = _smooth_truth()
+    epoch = 1.7e9
+    device = epoch + np.arange(N) / RATE
+    truth_b = np.arange(N) / (RATE * 1.0001) + 0.003 * np.sin(2 * np.pi * np.arange(N) / N)
+    members = [
+        {"key": "a", "device": device.copy(), "dataset": _jitter(truth, scale_periods=0.002, seed=2)},
+        {"key": "b", "device": device.copy(), "dataset": _jitter(truth_b, scale_periods=0.2, seed=3)},
+    ]  # ratio ~100x, but "b" is still well under one sample period of jitter
+    out = reconstruct_group(members)
+    # "b" self-fits: it tracks its OWN clock, not "a"'s slightly different rate.
+    assert np.abs(out["b"] - truth_b).max() < 5e-4
+
+
 # --- group_clocks ---