|
| 1 | +"""The vertical axis of a native 3-D file goes through the same preprocessing as xy.""" |
| 2 | + |
| 3 | +from __future__ import annotations |
| 4 | + |
| 5 | +from pathlib import Path |
| 6 | + |
| 7 | +import numpy as np |
| 8 | + |
| 9 | +from track2data.core.models import ( |
| 10 | + GapFillCfg, |
| 11 | + IdSwitchCfg, |
| 12 | + JumpCfg, |
| 13 | + PreprocessConfig, |
| 14 | + Session, |
| 15 | + SmoothCfg, |
| 16 | + VideoInfo, |
| 17 | +) |
| 18 | +from track2data.preprocess.pipeline import run |
| 19 | + |
| 20 | + |
| 21 | +def _session(xy: np.ndarray, z: np.ndarray, intervals=None) -> Session: |
| 22 | + return Session( |
| 23 | + velocity_threshold_px_frame=20.0, |
| 24 | + session_id="s", |
| 25 | + folder=Path("/tmp/s"), |
| 26 | + reader="test", |
| 27 | + video=VideoInfo(fps=30.0, n_frames=len(xy), width_px=0, height_px=0), |
| 28 | + n_animals=1, |
| 29 | + trajectory_variant="with_gaps", |
| 30 | + has_stable_identities=True, |
| 31 | + raw_xy=xy, |
| 32 | + raw_z=z.reshape(len(z), 1), |
| 33 | + has_pixel_frame=False, |
| 34 | + coordinate_unit="cm", |
| 35 | + tracking_intervals=intervals, |
| 36 | + ) |
| 37 | + |
| 38 | + |
| 39 | +def _cfg(*, gap=False, jump=False, smooth=None, across=False, max_gap=30) -> PreprocessConfig: |
| 40 | + return PreprocessConfig( |
| 41 | + gap_fill=GapFillCfg( |
| 42 | + enabled=gap, max_gap_frames=max_gap, across_tracking_intervals=across |
| 43 | + ), |
| 44 | + jump=JumpCfg( |
| 45 | + enabled=jump, method="idtracker_velocity_threshold", replacement="linear_interp" |
| 46 | + ), |
| 47 | + identity_switch=IdSwitchCfg(enabled=False), |
| 48 | + smoothing=smooth or SmoothCfg(enabled=False), |
| 49 | + ) |
| 50 | + |
| 51 | + |
| 52 | +def _line(n: int = 60): |
| 53 | + xy = np.zeros((n, 1, 2)) |
| 54 | + xy[:, 0, 0] = np.arange(n, dtype=float) |
| 55 | + z = np.arange(n, dtype=float) * 0.5 |
| 56 | + return xy, z |
| 57 | + |
| 58 | + |
| 59 | +def _path(z: np.ndarray) -> float: |
| 60 | + return float(np.nansum(np.abs(np.diff(z)))) |
| 61 | + |
| 62 | + |
| 63 | +def test_z_is_smoothed_with_the_xy_settings() -> None: |
| 64 | + rng = np.random.default_rng(0) |
| 65 | + xy, _ = _line(300) |
| 66 | + z = np.full(300, 5.0) + rng.normal(0.0, 0.2, 300) |
| 67 | + cfg = _cfg(smooth=SmoothCfg(enabled=True, method="savgol", window=25, polyorder=2)) |
| 68 | + |
| 69 | + psess = run(_session(xy, z), cfg) |
| 70 | + |
| 71 | + assert psess.z is not None |
| 72 | + # a flat true track: smoothing must remove most of the noise path, not leave it raw |
| 73 | + assert _path(psess.z[:, 0]) < 0.25 * _path(z) |
| 74 | + |
| 75 | + |
| 76 | +def test_z_is_unchanged_when_smoothing_and_gap_fill_are_off() -> None: |
| 77 | + xy, z = _line() |
| 78 | + z = z + np.random.default_rng(1).normal(0, 0.1, len(z)) |
| 79 | + psess = run(_session(xy, z), _cfg()) |
| 80 | + assert psess.z is not None |
| 81 | + np.testing.assert_array_equal(psess.z[:, 0], z) |
| 82 | + |
| 83 | + |
| 84 | +def test_short_z_gap_is_filled_like_xy() -> None: |
| 85 | + xy, z = _line() |
| 86 | + xy[10:15] = np.nan |
| 87 | + z[10:15] = np.nan |
| 88 | + psess = run(_session(xy, z), _cfg(gap=True)) |
| 89 | + assert psess.z is not None |
| 90 | + np.testing.assert_allclose(psess.z[10:15, 0], np.arange(10, 15) * 0.5) |
| 91 | + |
| 92 | + |
| 93 | +def test_z_gap_longer_than_max_stays_missing() -> None: |
| 94 | + xy, z = _line() |
| 95 | + xy[10:20] = np.nan |
| 96 | + z[10:20] = np.nan |
| 97 | + psess = run(_session(xy, z), _cfg(gap=True, max_gap=5)) |
| 98 | + assert psess.z is not None |
| 99 | + assert np.isnan(psess.z[10:20, 0]).all() |
| 100 | + |
| 101 | + |
| 102 | +def test_a_z_gap_xy_does_not_have_is_not_invented_beyond_the_anchors() -> None: |
| 103 | + """The leading edge has no left anchor, so it stays missing exactly as xy would.""" |
| 104 | + xy, z = _line() |
| 105 | + z[:3] = np.nan |
| 106 | + psess = run(_session(xy, z), _cfg(gap=True)) |
| 107 | + assert psess.z is not None |
| 108 | + assert np.isnan(psess.z[:3, 0]).all() |
| 109 | + |
| 110 | + |
| 111 | +def test_jump_replaced_z_is_refilled_not_kept_raw() -> None: |
| 112 | + xy, z = _line() |
| 113 | + xy[30, 0, 0] += 500.0 # triangulation glitch in the horizontal plane |
| 114 | + z[30] += 40.0 # and in the vertical one |
| 115 | + cfg = _cfg(gap=True, jump=True) |
| 116 | + |
| 117 | + psess = run(_session(xy, z), cfg) |
| 118 | + |
| 119 | + assert psess.jump_replaced is not None and psess.jump_replaced[30, 0] |
| 120 | + assert psess.z is not None |
| 121 | + np.testing.assert_allclose(psess.z[30, 0], 0.5 * (z[29] + z[31])) |
| 122 | + |
| 123 | + |
| 124 | +def test_jump_replaced_z_stays_missing_when_gap_fill_is_off() -> None: |
| 125 | + xy, z = _line() |
| 126 | + xy[30, 0, 0] += 500.0 |
| 127 | + psess = run(_session(xy, z), _cfg(jump=True)) |
| 128 | + assert psess.z is not None |
| 129 | + assert np.isnan(psess.z[30, 0]) |
| 130 | + |
| 131 | + |
| 132 | +def test_z_is_bridged_across_a_tracking_interval_gap_like_xy() -> None: |
| 133 | + n = 20 |
| 134 | + xy = np.zeros((n, 1, 2)) |
| 135 | + xy[:, 0, 0] = np.arange(n, dtype=float) |
| 136 | + z = np.arange(n, dtype=float) + 100.0 |
| 137 | + session = _session(xy, z, intervals=[(0, 10), (12, 22)]) |
| 138 | + |
| 139 | + psess = run(session, _cfg(gap=True, across=True), bridge_allowed=True) |
| 140 | + |
| 141 | + assert psess.z is not None and psess.tracked_mask is not None |
| 142 | + assert psess.separator_mask is not None |
| 143 | + bridged = ~psess.tracked_mask & ~psess.separator_mask |
| 144 | + assert bridged.any() |
| 145 | + # anchors: z=109 at video frame 9, z=110 at frame 12; video frames 10 and 11 lie on the line |
| 146 | + np.testing.assert_allclose(psess.z[bridged, 0], [109 + 1 / 3, 109 + 2 / 3]) |
| 147 | + |
| 148 | + |
| 149 | +def test_separator_rows_stay_missing_in_z() -> None: |
| 150 | + xy, z = _line(22) |
| 151 | + session = _session(xy, z, intervals=[(0, 10), (1000, 1012)]) |
| 152 | + psess = run(session, _cfg(gap=True, smooth=SmoothCfg(enabled=True))) |
| 153 | + assert psess.separator_mask is not None and psess.separator_mask.any() |
| 154 | + assert psess.z is not None |
| 155 | + assert np.isnan(psess.z[psess.separator_mask]).all() |
| 156 | + |
| 157 | + |
| 158 | +def _noisy_helix(n: int = 600, sigma: float = 0.2): |
| 159 | + """p7 of the 3-D review: a helix with 2 mm (0.2 cm) noise on every axis, 30 fps.""" |
| 160 | + rng = np.random.default_rng(3) |
| 161 | + t = np.arange(n) / 30.0 |
| 162 | + xy = np.zeros((n, 1, 2)) |
| 163 | + xy[:, 0, 0] = 10 * np.cos(2 * np.pi * 0.2 * t) + rng.normal(0, sigma, n) |
| 164 | + xy[:, 0, 1] = 10 * np.sin(2 * np.pi * 0.2 * t) + rng.normal(0, sigma, n) |
| 165 | + true_z = 5 + 2 * np.sin(2 * np.pi * 0.1 * t) |
| 166 | + return xy, true_z + rng.normal(0, sigma, n), true_z |
| 167 | + |
| 168 | + |
| 169 | +def test_vertical_path_falls_with_the_smoothing_window_toward_the_truth() -> None: |
| 170 | + """Before: the vertical path was 129 cm at every window against a true 16 cm.""" |
| 171 | + xy, z, true_z = _noisy_helix() |
| 172 | + paths = [] |
| 173 | + for window in (5, 15, 41): |
| 174 | + cfg = _cfg(smooth=SmoothCfg(enabled=True, method="savgol", window=window, polyorder=2)) |
| 175 | + psess = run(_session(xy, z), cfg) |
| 176 | + assert psess.z is not None |
| 177 | + paths.append(_path(psess.z[:, 0])) |
| 178 | + |
| 179 | + assert paths[0] > paths[1] > paths[2] |
| 180 | + assert paths[2] < 0.25 * _path(z) |
| 181 | + assert paths[2] < 1.5 * _path(true_z) |
| 182 | + |
| 183 | + |
| 184 | +def test_a_sensitivity_grid_point_changes_z() -> None: |
| 185 | + from track2data.core.models import PreprocessConfig as Cfg |
| 186 | + from track2data.sensitivity import GridPoint, _config_for |
| 187 | + |
| 188 | + xy, z, _ = _noisy_helix() |
| 189 | + session = _session(xy, z) |
| 190 | + narrow = run(session, _config_for(Cfg(), GridPoint(smoothing_window=5, max_gap_frames=10))) |
| 191 | + wide = run(session, _config_for(Cfg(), GridPoint(smoothing_window=41, max_gap_frames=10))) |
| 192 | + assert narrow.z is not None and wide.z is not None |
| 193 | + assert _path(wide.z[:, 0]) < _path(narrow.z[:, 0]) |
0 commit comments