scripts/benchmark_dememwm_multiview_selection.py

"""Benchmark DeMemWM dynamic multiview memory selectors on synthetic poses."""

from __future__ import annotations

import argparse
import importlib.util
import statistics
import sys
import time
from pathlib import Path

import numpy as np
import torch


REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

SELECTORS = ("fov_greedy", "pose_plucker_fps")


def _load_memory_selection_module():
    module_path = REPO_ROOT / "datasets" / "video" / "memory_selection.py"
    spec = importlib.util.spec_from_file_location("dememwm_memory_selection", module_path)
    if spec is None or spec.loader is None:
        raise ImportError(f"could not load memory selection module from {module_path}")
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


memory_selection = _load_memory_selection_module()


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--num-frames", type=int, required=True)
    parser.add_argument("--target-start", type=int, required=True)
    parser.add_argument("--target-len", type=int, required=True)
    parser.add_argument("--num-iters", type=int, required=True)
    parser.add_argument("--pose-preselect-topk", type=int, required=True)
    parser.add_argument("--candidate-chunk-size", type=int, required=True)
    parser.add_argument("--selectors", nargs="+", default=list(SELECTORS))
    parser.add_argument("--write-report", type=Path, default=None)
    return parser.parse_args()


def _synthetic_poses(num_frames: int) -> np.ndarray:
    frame = np.arange(num_frames, dtype=np.float32)
    poses = np.zeros((num_frames, 5), dtype=np.float32)
    poses[:, 0] = 0.03 * frame + 24.0 * np.sin(frame * 0.031)
    poses[:, 1] = 4.0 * np.cos(frame * 0.019)
    poses[:, 2] = 0.015 * frame + 24.0 * np.cos(frame * 0.027)
    poses[:, 3] = 18.0 * np.sin(frame * 0.017)
    poses[:, 4] = np.remainder(2.7 * frame + 30.0 * np.sin(frame * 0.011) + 180.0, 360.0) - 180.0
    return poses


def _target_positions(target_start: int, target_len: int, num_frames: int) -> np.ndarray:
    stop = target_start + target_len
    if target_start < 0 or target_len <= 0 or stop > num_frames:
        raise ValueError(
            f"target window [{target_start}, {stop}) must be non-empty and inside num_frames={num_frames}"
        )
    return np.arange(target_start, stop, dtype=np.int64)


def _selection_cfg(selector: str, args: argparse.Namespace) -> dict:
    return {
        "enabled": True,
        "causal": True,
        "max_anchor_frames": 0,
        "max_dynamic_frames": args.target_len,
        "max_revisit_frames": 0,
        "pose_similarity_threshold": 0.0,
        "training_use_plucker": True,
        "training_plucker_weight": 1.0,
        "fov_overlap_threshold": 0.6,
        "min_total_selected_coverage": 0.1,
        "local_context_exclusion_frames": 8,
        "plucker_moment_radius": 30.0,
        "anchor_diverse_selection": True,
        "pose_preselect_topk": args.pose_preselect_topk,
        "candidate_chunk_size": args.candidate_chunk_size,
        "dynamic": {
            "selection_policy": "multiview",
            "multiview_selector": selector,
        },
    }


def _percentile(values: list[float], fraction: float) -> float:
    if not values:
        return 0.0
    ordered = sorted(values)
    index = min(len(ordered) - 1, int(np.ceil(fraction * len(ordered))) - 1)
    return ordered[index]


def _base_candidates(poses: np.ndarray, target_positions: np.ndarray, cfg: dict) -> np.ndarray:
    return memory_selection._memory_candidate_frames(
        len(poses),
        target_positions,
        cfg,
        "training",
        min_candidate_frame=0,
    )


def _fov_candidate_count(poses: np.ndarray, target_positions: np.ndarray, cfg: dict) -> int:
    candidates = _base_candidates(poses, target_positions, cfg)
    poses_t = torch.as_tensor(poses, dtype=torch.float32)
    preselected = memory_selection._pose_preselect(candidates, poses_t, target_positions, cfg)
    return int(len(preselected))


def _pose_plucker_candidate_count(poses: np.ndarray, target_positions: np.ndarray, cfg: dict) -> int:
    candidates = _base_candidates(poses, target_positions, cfg)
    ranked_ids, _, _ = memory_selection._rank_pose_plucker_candidates(poses, candidates, target_positions, cfg)
    topk = memory_selection.cfg_get(cfg, "pose_preselect_topk", 64)
    if topk is not None and int(topk) > 0:
        return int(min(int(topk), ranked_ids.numel()))
    return int(ranked_ids.numel())


def _run_once(selector: str, poses: np.ndarray, target_positions: np.ndarray, cfg: dict, count: int) -> np.ndarray:
    if selector == "fov_greedy":
        candidates = _base_candidates(poses, target_positions, cfg)
        pool = memory_selection._build_fov_candidate_pool(
            poses,
            candidates,
            target_positions,
            cfg,
            use_plucker=True,
        )
        return memory_selection._select_dynamic_multiview(
            poses,
            target_positions,
            cfg,
            count,
            split="training",
            fov_pool=pool,
        )
    if selector == "pose_plucker_fps":
        return memory_selection._select_dynamic_multiview(
            poses,
            target_positions,
            cfg,
            count,
            split="training",
        )
    raise ValueError(f"unknown selector {selector!r}")


def _benchmark_selector(selector: str, poses: np.ndarray, target_positions: np.ndarray, args: argparse.Namespace) -> dict:
    cfg = _selection_cfg(selector, args)
    count = int(args.target_len)
    if selector == "fov_greedy":
        candidate_count = _fov_candidate_count(poses, target_positions, cfg)
        fov_pool_reuse = True
    else:
        candidate_count = _pose_plucker_candidate_count(poses, target_positions, cfg)
        fov_pool_reuse = False

    selected = _run_once(selector, poses, target_positions, cfg, count)
    timings_ms = []
    for _ in range(args.num_iters):
        start = time.perf_counter()
        selected = _run_once(selector, poses, target_positions, cfg, count)
        timings_ms.append((time.perf_counter() - start) * 1000.0)

    return {
        "selector": selector,
        "mean_ms": statistics.fmean(timings_ms),
        "median_ms": statistics.median(timings_ms),
        "p90_ms": _percentile(timings_ms, 0.90),
        "selected_count": int(len(selected)),
        "candidate_count_after_pose_preselection": candidate_count,
        "fov_pool_reuse": fov_pool_reuse,
        "device": str(torch.device("cpu")),
    }


def _format_results(results: list[dict]) -> str:
    lines = [
        "selector mean_ms median_ms p90_ms selected_count candidate_count_after_pose_preselection fov_pool_reuse device"
    ]
    for row in results:
        lines.append(
            "{selector} {mean_ms:.3f} {median_ms:.3f} {p90_ms:.3f} {selected_count} "
            "{candidate_count_after_pose_preselection} {fov_pool_reuse} {device}".format(**row)
        )
    return "\n".join(lines)


def _write_report(path: Path, args: argparse.Namespace, results: list[dict]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    lines = [
        "# DeMemWM Multiview Selection Speed Report",
        "",
        "This benchmark used deterministic synthetic poses only. It is not a substitute for a real dataset sampling benchmark.",
        "",
        "```text",
        "python " + " ".join(sys.argv),
        "```",
        "",
        "| selector | mean ms | median ms | p90 ms | selected frames | pose-preselected candidates | FOV pool reuse | device |",
        "| --- | ---: | ---: | ---: | ---: | ---: | --- | --- |",
    ]
    for row in results:
        lines.append(
            "| {selector} | {mean_ms:.3f} | {median_ms:.3f} | {p90_ms:.3f} | {selected_count} | "
            "{candidate_count_after_pose_preselection} | {fov_pool_reuse} | {device} |".format(**row)
        )
    lines.extend(
        [
            "",
            f"Synthetic frames: {args.num_frames}",
            f"Target window: [{args.target_start}, {args.target_start + args.target_len})",
            f"Iterations: {args.num_iters}",
            f"pose_preselect_topk: {args.pose_preselect_topk}",
            f"candidate_chunk_size: {args.candidate_chunk_size}",
        ]
    )
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def main() -> int:
    args = _parse_args()
    unknown = [selector for selector in args.selectors if selector not in SELECTORS]
    if unknown:
        valid = ", ".join(SELECTORS)
        print(f"unknown selector(s): {', '.join(unknown)}; valid selectors: {valid}", file=sys.stderr)
        return 2

    poses = _synthetic_poses(args.num_frames)
    target_positions = _target_positions(args.target_start, args.target_len, args.num_frames)
    results = [_benchmark_selector(selector, poses, target_positions, args) for selector in args.selectors]
    print(_format_results(results))

    if args.write_report is not None:
        _write_report(args.write_report, args, results)
        print(f"wrote report: {args.write_report}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())