diff --git a/src/ezmsg/nwb/clockmodel.py b/src/ezmsg/nwb/clockmodel.py index 8d71c57..97c2fd1 100644 --- a/src/ezmsg/nwb/clockmodel.py +++ b/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"], diff --git a/tests/test_clockmodel.py b/tests/test_clockmodel.py index 5349d9b..27f67ce 100644 --- a/tests/test_clockmodel.py +++ b/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 ---