"""actionrope/geometry.py 的纯 CPU 单测：帧/cell 映射、known/new 掩码几何、日志采样、转场连续性。 /opt/dlami/nvme/zhiyangdeng/ActionRoPE/.venv/bin/python -m pytest tests/test_geometry.py """ from __future__ import annotations import json import os import sys import numpy as np import pytest sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from actionrope import geometry as G # noqa: E402 # ------------------------------------------------------------ cell ↔ 帧 ---- def test_cell_frames_partition(): """21 个 cell 恰好不重不漏地覆盖 81 帧，cell 0 只含帧 0。""" assert G.cell_frames(0) == [0] assert G.cell_frames(1) == [1, 2, 3, 4] assert G.cell_frames(20) == [77, 78, 79, 80] seen = [f for k in range(G.CELLS) for f in G.cell_frames(k)] assert seen == list(range(G.FRAMES)) with pytest.raises(ValueError): G.cell_frames(21) def test_frames_to_cells_mean(): off = np.zeros((G.FRAMES, 2)) off[:, 0] = np.arange(G.FRAMES) * 2.0 # dx 每帧 +2 off[:, 1] = -np.arange(G.FRAMES) * 0.5 # dy 每帧 −0.5 cells = G.frames_to_cells(off) assert cells.shape == (G.CELLS, 2) assert np.allclose(cells[0], 0.0) # cell 1 = 帧 1..4，均值帧号 2.5 assert np.allclose(cells[1], [5.0, -1.25]) assert np.allclose(cells[20], [np.mean([77, 78, 79, 80]) * 2.0, -np.mean([77, 78, 79, 80]) * 0.5]) with pytest.raises(ValueError): G.frames_to_cells(np.zeros((80, 2))) # ------------------------------------------------------------ known / new ---- def test_known_mask_zero_offset_all_known(): off = np.zeros((3, 2)) assert G.known_mask_tokens(off).all() assert G.known_mask_latent(off).all() assert G.known_mask_tokens(off).shape == (3, G.TOK_H, G.TOK_W) assert G.known_mask_latent(off).shape == (3, G.LAT_H, G.LAT_W) def test_known_mask_dx_plus_64_right_two_columns_new(): """玩家右走 64 px ⇒ 背景左移 ⇒ 屏幕最右 2 列 token 是首帧没见过的。""" m = G.known_mask_tokens(np.array([[64.0, 0.0]]))[0] assert not m[:, -2:].any() assert m[:, :-2].all() # latent 分辨率：64 px = 4 个 latent 像素 ml = G.known_mask_latent(np.array([[64.0, 0.0]]))[0] assert not ml[:, -4:].any() assert ml[:, :-4].all() def test_known_mask_dy_minus_32_top_row_new(): """玩家上走 32 px（dy<0）⇒ 背景下移 ⇒ 顶部 1 行 token 是新的。""" m = G.known_mask_tokens(np.array([[0.0, -32.0]]))[0] assert not m[0].any() assert m[1:].all() ml = G.known_mask_latent(np.array([[0.0, -32.0]]))[0] assert not ml[:2].any() assert ml[2:].all() def test_known_mask_fractional_and_boundary(): """半格偏移落在闭区间边界上仍算 known；再多一点就是 new。""" m = G.known_mask_tokens(np.array([[16.0, 0.0]]))[0] # 0.5 token assert m.all() m = G.known_mask_tokens(np.array([[16.5, 0.0]]))[0] assert not m[:, -1].any() and m[:, :-1].all() # 左上方向：dx=−40 ⇒ −1.25 token ⇒ 只有最左 1 列 new（j=0 → −1.25 < −0.5；j=1 → −0.25 ≥ −0.5 仍 known） m = G.known_mask_tokens(np.array([[-40.0, 0.0]]))[0] assert not m[:, 0].any() and m[:, 1:].all() m = G.known_mask_tokens(np.array([[-48.0, 0.0]]))[0] # −1.5 ⇒ j=1 恰在边界 −0.5，known assert not m[:, 0].any() and m[:, 1:].all() m = G.known_mask_tokens(np.array([[-49.0, 0.0]]))[0] assert not m[:, :2].any() and m[:, 2:].all() def test_known_mask_per_frame_independent(): off = np.array([[0.0, 0.0], [64.0, 0.0], [0.0, -32.0], [1000.0, 0.0]]) m = G.known_mask_tokens(off) assert m[0].all() assert not m[1][:, -2:].any() assert not m[2][0].any() assert not m[3].any() # 位移超过整幅 ⇒ 全 new # ------------------------------------------------------------ 转场连续性 ---- def test_transition_offsets_continuity(): rng = np.random.default_rng(0) a = np.cumsum(rng.normal(size=(G.FRAMES, 2)), 0); a -= a[0] b = np.cumsum(rng.normal(size=(G.FRAMES, 2)), 0); b -= b[0] n_pre = 17 out = G.transition_frame_offsets(a, b, n_pre) assert out.shape == (G.FRAMES, 2) assert np.allclose(out[:n_pre], a[:n_pre]) # 接缝处的增量来自 B：off[n_pre] − off[n_pre−1] = B[n_pre] − B[n_pre−1] assert np.allclose(out[n_pre] - out[n_pre - 1], b[n_pre] - b[n_pre - 1]) assert np.allclose(np.diff(out[n_pre:], axis=0), np.diff(b[n_pre:], axis=0)) assert np.allclose(out[0], 0.0) # 全程与 A 或 B 同源时退化为原序列 assert np.allclose(G.transition_frame_offsets(a, a, n_pre), a) with pytest.raises(ValueError): G.transition_frame_offsets(a, b, 0) # ------------------------------------------------------------ 日志采样 ---- def _synthetic_tracks(tmp_path, offset_segments, M_px1280, n=400, dt=0.05, gap=None): """造一个最小 session：直线匀速走 x，可选在中间挖一段日志断档。""" vt = np.arange(n) * dt x = 2.0 * vt # 2 世界单位/s y = np.zeros(n); z = np.zeros(n) if gap is not None: keep = (vt < gap[0]) | (vt > gap[1]) vt, x, y, z = vt[keep], x[keep], y[keep], z[keep] # 故意加重复时间戳，load_session_tracks 必须去掉 vt = np.concatenate([vt, vt[:3]]); x = np.concatenate([x, x[:3]]) y = np.concatenate([y, y[:3]]); z = np.concatenate([z, z[:3]]) d = tmp_path / "sess" d.mkdir() np.savez(d / "state.npz", vt=vt, x=x.astype(np.float32), y=y.astype(np.float32), z=z.astype(np.float32), qf=np.arange(vt.size), rg=np.zeros(vt.size, np.int32)) (d / "offsets.json").write_text(json.dumps(dict(segments=offset_segments, M_px1280=M_px1280))) return G.load_session_tracks(str(tmp_path), "sess") def test_load_session_tracks_dedupe_and_scale(tmp_path): M = [[-40.0, 0.0], [0.0, 18.0], [0.0, 27.0]] tr = _synthetic_tracks(tmp_path, [dict(vt_lo=0.0, vt_hi=1e12, offset=0.1)], M) assert tr["vt"].size == 400 and np.all(np.diff(tr["vt"]) > 0) assert np.allclose(tr["M_px832"][:, 0], np.array(M)[:, 0] * 832 / 1280) assert np.allclose(tr["M_px832"][:, 1], np.array(M)[:, 1] * 480 / 720) assert np.allclose(tr["vid_t"], tr["vt"] + 0.1) u, v = G.project_px832(tr, 1.0, 0.0, 0.0) assert np.isclose(u, -40.0 * 832 / 1280) and np.isclose(v, 0.0) def test_offset_piecewise_constant(tmp_path): M = [[-40.0, 0.0], [0.0, 18.0], [0.0, 27.0]] segs = [dict(vt_lo=0.0, vt_hi=10.0, offset=0.1), dict(vt_lo=10.0, vt_hi=1e12, offset=-4.0)] tr = _synthetic_tracks(tmp_path, segs, M) assert G.offset_at(tr, 5.0) == 0.1 assert G.offset_at(tr, 10.0) == -4.0 # 段以 vt_lo 为界，右侧闭 assert G.offset_at(tr, 15.0) == -4.0 def test_sample_frame_offsets_linear_walk(tmp_path): """匀速沿世界 +x 走 2/s：M[0,0]=−40 ⇒ 1280 画幅下 −80 px/s ⇒ 832 画幅 −52 px/s。""" M = [[-40.0, 0.0], [0.0, 18.0], [0.0, 27.0]] tr = _synthetic_tracks(tmp_path, [dict(vt_lo=0.0, vt_hi=1e12, offset=0.1)], M) off, valid, reason = G.sample_frame_offsets(tr, video_t0=2.1, return_valid=True) assert valid and reason == "" assert off.shape == (G.FRAMES, 2) assert np.allclose(off[0], 0.0) # 最近邻采样：20 Hz 日志对 16 fps 的帧最多差 0.025 s ⇒ 1.3 px expect = -52.0 * np.arange(G.FRAMES) / G.FPS assert np.abs(off[:, 0] - expect).max() < 1.5 assert np.allclose(off[:, 1], 0.0) assert np.isfinite(G.sample_frame_offsets(tr, 2.1)).all() def test_sample_frame_offsets_gap_invalid(tmp_path): M = [[-40.0, 0.0], [0.0, 18.0], [0.0, 27.0]] tr = _synthetic_tracks(tmp_path, [dict(vt_lo=0.0, vt_hi=1e12, offset=0.0)], M, gap=(5.0, 5.6)) off, valid, reason = G.sample_frame_offsets(tr, video_t0=3.0, return_valid=True) assert not valid and reason.startswith("log_gap") assert np.isnan(G.sample_frame_offsets(tr, 3.0)).all() # 阈值放宽到 0.5 s 就有效（最近邻误差最大 0.3 s） _, valid2, _ = G.sample_frame_offsets(tr, 3.0, return_valid=True, max_sample_dt=0.5) assert valid2 # 完全落在日志范围之外 ⇒ 断档 _, valid3, reason3 = G.sample_frame_offsets(tr, video_t0=100.0, return_valid=True) assert not valid3 and reason3.startswith("log_gap") def test_sample_frame_offsets_teleport_invalid(tmp_path): M = [[-40.0, 0.0], [0.0, 18.0], [0.0, 27.0]] tr = _synthetic_tracks(tmp_path, [dict(vt_lo=0.0, vt_hi=1e12, offset=0.0)], M) tr["x"] = tr["x"].copy(); tr["x"][tr["vt"] > 4.0] += 50.0 # 传送 50 世界单位 = 1300 px _, valid, reason = G.sample_frame_offsets(tr, video_t0=3.0, return_valid=True) assert not valid and reason.startswith("teleport") def test_sample_frame_offsets_log_jump_invalid(tmp_path): """offset 台阶：vt=10 处 offset 从 0 跳到 −4 ⇒ vid_t 局部非单调，跨过台阶的 clip 最近邻会跳到 4 s 之后的日志样本（时间差仍 < 0.15 s），必须按 log_jump 判 invalid。""" M = [[-40.0, 0.0], [0.0, 18.0], [0.0, 27.0]] segs = [dict(vt_lo=0.0, vt_hi=10.0, offset=0.0), dict(vt_lo=10.0, vt_hi=1e12, offset=-4.0)] tr = _synthetic_tracks(tmp_path, segs, M) assert np.diff(tr["vid_t"]).min() < 0 _, valid, reason = G.sample_frame_offsets(tr, video_t0=4.0, return_valid=True) assert not valid and reason.startswith("log_jump") # 台阶之前 / 之后整段落在单调区间里的 clip 不受影响 _, valid_a, _ = G.sample_frame_offsets(tr, video_t0=0.5, return_valid=True) _, valid_b, _ = G.sample_frame_offsets(tr, video_t0=8.0, return_valid=True) assert valid_a and valid_b def test_action_dirs_unit(): n = np.hypot(G.ACTION_DIRS[:, 0], G.ACTION_DIRS[:, 1]) assert np.isclose(n[0], 0.0) and np.allclose(n[1:], 1.0) assert np.allclose(G.ACTION_DIRS[7], [1.0, 0.0]) # moving right assert np.allclose(G.ACTION_DIRS[1], [0.0, -1.0]) # moving up（屏幕 v 向下为正） def test_shaped_loss_weight(): """dist 形状：known=1，紧贴足迹边界 ≈1，远处到 new_weight；σ 调制只作用于额外部分。""" import torch from actionrope.arope import loss_weight_map off = torch.tensor([[[0.0, 0.0], [96.0, 0.0]]]) # 帧 1 右移 96 px = 6 latent 格 ⇒ 最右 6 列 new flat = loss_weight_map(off, 30, 52, new_weight=3.0)[0, 0] assert flat[1, :, :46].eq(1).all() and flat[1, :, 46:].eq(3).all() and flat[0].eq(1).all() w = loss_weight_map(off, 30, 52, new_weight=3.0, ramp_latent=4.0)[0, 0] assert w[0].eq(1).all() and w[1, :, :46].eq(1).all() col = w[1, 0, 46:].tolist() # 距边界 0.5,1.5,…,5.5 格 ⇒ 1+2·min(d/4,1) assert abs(col[0] - 1.25) < 1e-5 and abs(col[1] - 1.75) < 1e-5 and abs(col[3] - 2.75) < 1e-5 and col[4] == 3.0 and col[5] == 3.0 ws = loss_weight_map(off, 30, 52, new_weight=3.0, ramp_latent=4.0, sigma=torch.tensor([1.0]))[0, 0] assert abs(ws[1, 0, 51] - 4.0) < 1e-5 # 1 + 2·1.5 w0 = loss_weight_map(off, 30, 52, new_weight=3.0, ramp_latent=4.0, sigma=torch.tensor([0.0]))[0, 0] assert abs(w0[1, 0, 51] - 2.0) < 1e-5 # 1 + 2·0.5 assert ws[1, :, :46].eq(1).all()