scripts/analyze_feature_cache.py

#!/usr/bin/env python3
"""Analyze timestep and cross-chunk feature reuse in Self-Forcing.

The script runs the released four-step causal checkpoint without changing its
outputs. Forward hooks capture sampled DiT hidden states, residual updates, the
final velocity, and a low-dimensional dense projection used for optical-flow
alignment. It then produces:

* within-chunk and cross-chunk feature-pair metrics;
* leave-one-prompt-out channel-wise probes for conditional chunk information;
* shuffled, wrong-step, distant-chunk, zero, and noise controls;
* motion-stratified raw/global/dense-flow alignment measurements.

Each prompt is saved independently, so interrupted generation can be resumed.
"""

from __future__ import annotations

import argparse
import csv
import json
import math
import os
import random
import sys
import time
from collections import defaultdict
from pathlib import Path
from typing import Any, Iterable


def _preparse_gpu() -> str:
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument("--gpu", default="2")
    args, _ = parser.parse_known_args()
    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
    return str(args.gpu)


PHYSICAL_GPU = _preparse_gpu()

import cv2
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F
from omegaconf import OmegaConf

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

from pipeline import CausalInferencePipeline
from utils.misc import set_seed


EXPECTED_LATENT_HEIGHT = 60
EXPECTED_LATENT_WIDTH = 104
FRAME_TOKEN_HEIGHT = 30
FRAME_TOKEN_WIDTH = 52
FRAME_SEQ_LENGTH = FRAME_TOKEN_HEIGHT * FRAME_TOKEN_WIDTH


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Self-Forcing cross-chunk feature-cache analysis"
    )
    parser.add_argument("--gpu", default=PHYSICAL_GPU)
    parser.add_argument(
        "--config_path", type=Path, default=Path("configs/self_forcing_dmd.yaml")
    )
    parser.add_argument(
        "--checkpoint_path",
        type=Path,
        default=Path("checkpoints/self_forcing_dmd.pt"),
    )
    parser.add_argument(
        "--prompt_path",
        type=Path,
        default=Path("prompts/MovieGenVideoBench_extended.txt"),
    )
    parser.add_argument("--output_dir", type=Path, required=True)
    parser.add_argument("--num_prompts", type=int, default=3)
    parser.add_argument("--num_frames", type=int, default=21)
    parser.add_argument("--seed", type=int, default=20260728)
    parser.add_argument(
        "--same_seed",
        action="store_true",
        help="Reset every prompt to --seed for paired cross-model evaluation.",
    )
    parser.add_argument(
        "--layers", type=int, nargs="+", default=[0, 9, 19, 29]
    )
    parser.add_argument("--max_tokens", type=int, default=256)
    parser.add_argument("--projection_dim", type=int, default=64)
    parser.add_argument("--ridge", type=float, default=1e-4)
    parser.add_argument("--use_ema", action="store_true", default=True)
    parser.add_argument("--no_ema", action="store_false", dest="use_ema")
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument(
        "--cosine_only",
        action="store_true",
        help="Generate pair metrics/heatmap only; skip probes and motion analysis.",
    )
    parser.add_argument(
        "--analysis_only",
        action="store_true",
        help="Skip model loading and analyze existing per-prompt snapshots.",
    )
    parser.add_argument(
        "--save_preview",
        action="store_true",
        help="Save a compact MP4 preview when torchvision video IO is available.",
    )
    args = parser.parse_args()
    if args.num_frames % 3 != 0:
        parser.error("--num_frames must be divisible by the configured 3-frame chunk")
    if args.num_prompts < 3 and not args.cosine_only:
        parser.error("--num_prompts must be at least 3 for held-out/shuffle controls")
    return args


def resolve_path(path: Path) -> Path:
    return path if path.is_absolute() else REPO_ROOT / path


def write_csv(path: Path, rows: list[dict[str, Any]]) -> None:
    if not rows:
        return
    path.parent.mkdir(parents=True, exist_ok=True)
    fields: list[str] = []
    for row in rows:
        for key in row:
            if key not in fields:
                fields.append(key)
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields, extrasaction="ignore")
        writer.writeheader()
        writer.writerows(rows)


def read_prompts(path: Path, count: int) -> list[str]:
    with path.open("r", encoding="utf-8") as handle:
        prompts = [line.strip() for line in handle if line.strip()]
    if len(prompts) < count:
        raise ValueError(f"Requested {count} prompts, found {len(prompts)} in {path}")
    return prompts[:count]


def regular_grid_indices(
    frames: int, height: int, width: int, max_tokens: int, device: torch.device
) -> tuple[torch.Tensor, torch.Tensor]:
    total = frames * height * width
    if max_tokens >= total:
        coords = torch.cartesian_prod(
            torch.arange(frames, device=device),
            torch.arange(height, device=device),
            torch.arange(width, device=device),
        )
    else:
        per_frame = max(1, max_tokens // frames)
        h_count = max(
            1, min(height, int(round(math.sqrt(per_frame * height / width))))
        )
        w_count = max(1, min(width, per_frame // h_count))
        while frames * h_count * w_count > max_tokens and w_count > 1:
            w_count -= 1
        while frames * h_count * w_count > max_tokens and h_count > 1:
            h_count -= 1
        hs = (
            torch.linspace(0, height - 1, h_count, device=device)
            .round()
            .long()
            .unique()
        )
        ws = (
            torch.linspace(0, width - 1, w_count, device=device)
            .round()
            .long()
            .unique()
        )
        coords = torch.cartesian_prod(
            torch.arange(frames, device=device), hs, ws
        )
    flat = (
        coords[:, 0] * height * width + coords[:, 1] * width + coords[:, 2]
    )
    return flat.long(), coords.long()


class FeatureRecorder:
    def __init__(
        self,
        model: torch.nn.Module,
        layers: list[int],
        denoising_timesteps: Iterable[float],
        num_frame_per_block: int,
        max_tokens: int,
        projection_dim: int,
    ) -> None:
        self.model = model
        self.layers = sorted(set(int(value) for value in layers))
        self.timesteps = [float(value) for value in denoising_timesteps]
        self.num_frame_per_block = int(num_frame_per_block)
        self.max_tokens = int(max_tokens)
        self.projection_dim = int(projection_dim)
        self.projection_layer = self.layers[-1]
        self.records: dict[str, dict[str, torch.Tensor]] = defaultdict(dict)
        self.projected: dict[str, torch.Tensor] = {}
        self.projected_by_layer: dict[str, torch.Tensor] = {}
        self.sample_coords: dict[str, torch.Tensor] = {}
        self.current: dict[str, Any] = {"active": False}
        self.handles: list[Any] = []
        self._projection_cache: dict[tuple[int, str], torch.Tensor] = {}
        self._register()

    def _register(self) -> None:
        self.handles.append(
            self.model.register_forward_pre_hook(self._model_pre_hook, with_kwargs=True)
        )
        self.handles.append(self.model.register_forward_hook(self._model_output_hook))
        for layer in self.layers:
            if layer < 0 or layer >= len(self.model.blocks):
                raise ValueError(
                    f"Layer {layer} outside model block range 0..{len(self.model.blocks)-1}"
                )
            self.handles.append(
                self.model.blocks[layer].register_forward_hook(
                    self._make_block_hook(layer)
                )
            )

    def close(self) -> None:
        for handle in self.handles:
            handle.remove()
        self.handles.clear()

    def reset(self) -> None:
        self.records = defaultdict(dict)
        self.projected = {}
        self.projected_by_layer = {}
        self.sample_coords = {}
        self.current = {"active": False}

    def _model_pre_hook(
        self, _module: torch.nn.Module, _args: tuple[Any, ...], kwargs: dict[str, Any]
    ) -> None:
        timestep = kwargs.get("t")
        current_start = int(kwargs.get("current_start", 0) or 0)
        if not isinstance(timestep, torch.Tensor) or timestep.numel() == 0:
            self.current = {"active": False}
            return
        value = float(timestep.detach().float().reshape(-1)[0].item())
        distances = [abs(value - expected) for expected in self.timesteps]
        step = int(np.argmin(distances))
        if distances[step] > 0.5:
            self.current = {"active": False, "timestep": value}
            return
        frames = int(timestep.shape[-1]) if timestep.ndim > 1 else 1
        if frames != self.num_frame_per_block:
            self.current = {"active": False, "timestep": value}
            return
        start_frame = current_start // FRAME_SEQ_LENGTH
        chunk = start_frame // self.num_frame_per_block
        self.current = {
            "active": True,
            "chunk": int(chunk),
            "step": step,
            "timestep": value,
            "frames": frames,
        }

    def _key(self) -> str:
        return f"{self.current['chunk']}:{self.current['step']}"

    def _hidden_indices(
        self, tokens: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        frames = int(self.current["frames"])
        if tokens.shape[1] != frames * FRAME_TOKEN_HEIGHT * FRAME_TOKEN_WIDTH:
            raise ValueError(
                f"Unexpected hidden token count {tokens.shape[1]} for {frames} frames"
            )
        return regular_grid_indices(
            frames,
            FRAME_TOKEN_HEIGHT,
            FRAME_TOKEN_WIDTH,
            self.max_tokens,
            tokens.device,
        )

    def _projection(self, dim: int, device: torch.device) -> torch.Tensor:
        cache_key = (dim, str(device))
        if cache_key not in self._projection_cache:
            generator = torch.Generator(device="cpu").manual_seed(20260728 + dim)
            signs = torch.randint(
                0,
                2,
                (dim, self.projection_dim),
                generator=generator,
                dtype=torch.int8,
            )
            projection = (
                signs.float().mul_(2).sub_(1).div_(math.sqrt(self.projection_dim))
            )
            self._projection_cache[cache_key] = projection.to(device)
        return self._projection_cache[cache_key]

    def _make_block_hook(self, layer: int):
        def hook(
            _module: torch.nn.Module,
            inputs: tuple[torch.Tensor, ...],
            output: torch.Tensor,
        ) -> None:
            if not self.current.get("active", False):
                return
            if not inputs or not isinstance(output, torch.Tensor):
                return
            key = self._key()
            hidden_input = inputs[0]
            indices, coords = self._hidden_indices(output)
            hidden = output[0].index_select(0, indices)
            delta = (output - hidden_input)[0].index_select(0, indices)
            self.records[f"block_{layer}_hidden"][key] = (
                hidden.detach().to(dtype=torch.float16, device="cpu")
            )
            self.records[f"block_{layer}_delta"][key] = (
                delta.detach().to(dtype=torch.float16, device="cpu")
            )
            self.sample_coords["hidden"] = coords.detach().cpu()

            if layer in self.layers:
                projection = self._projection(output.shape[-1], output.device)
                dense = torch.matmul(output[0].float(), projection)
                frames = int(self.current["frames"])
                dense = dense.reshape(
                    frames,
                    FRAME_TOKEN_HEIGHT,
                    FRAME_TOKEN_WIDTH,
                    self.projection_dim,
                )
                dense_cpu = dense.detach().to(dtype=torch.float16, device="cpu")
                self.projected_by_layer[f"{layer}:{key}"] = dense_cpu
                # Preserve the legacy key layout for existing final-layer analyses.
                if layer == self.projection_layer:
                    self.projected[key] = dense_cpu

        return hook

    def _model_output_hook(
        self,
        _module: torch.nn.Module,
        _inputs: tuple[Any, ...],
        output: torch.Tensor,
    ) -> None:
        if not self.current.get("active", False):
            return
        if not isinstance(output, torch.Tensor) or output.ndim != 5:
            return
        batch, channels, frames, height, width = output.shape
        if batch != 1:
            raise ValueError(f"Analysis expects batch size 1, got {batch}")
        tokens = output.permute(0, 2, 3, 4, 1).reshape(
            batch, frames * height * width, channels
        )
        indices, coords = regular_grid_indices(
            frames, height, width, self.max_tokens, output.device
        )
        sampled = tokens[0].index_select(0, indices)
        self.records["dit_output"][self._key()] = sampled.detach().to(
            dtype=torch.float16, device="cpu"
        )
        self.sample_coords["dit_output"] = coords.detach().cpu()

    def state_dict(self) -> dict[str, Any]:
        return {
            "timesteps": self.timesteps,
            "layers": self.layers,
            "projection_layer": self.projection_layer,
            "projection_dim": self.projection_dim,
            "records": {stage: dict(values) for stage, values in self.records.items()},
            "projected": dict(self.projected),
            "projected_by_layer": dict(self.projected_by_layer),
            "sample_coords": dict(self.sample_coords),
        }


def build_pipeline(args: argparse.Namespace) -> tuple[CausalInferencePipeline, Any]:
    config = OmegaConf.load(resolve_path(args.config_path))
    default_config = OmegaConf.load(REPO_ROOT / "configs/default_config.yaml")
    config = OmegaConf.merge(default_config, config)
    device = torch.device("cuda")
    pipeline = CausalInferencePipeline(config, device=device)
    checkpoint = torch.load(
        resolve_path(args.checkpoint_path), map_location="cpu", weights_only=False
    )
    state_key = "generator_ema" if args.use_ema else "generator"
    pipeline.generator.load_state_dict(checkpoint[state_key])
    del checkpoint
    pipeline = pipeline.to(dtype=torch.bfloat16)
    pipeline.text_encoder.to(device=device)
    pipeline.generator.to(device=device)
    pipeline.vae.to(device=device)
    pipeline.eval()
    return pipeline, config


def downsample_anchors(video: torch.Tensor, latent_frames: int) -> np.ndarray:
    value = video[0].detach().float().cpu()
    frame_count = value.shape[0]
    if frame_count == latent_frames:
        indices = np.arange(latent_frames)
    else:
        indices = np.linspace(0, frame_count - 1, latent_frames).round().astype(int)
    anchors = value[indices].permute(0, 2, 3, 1).clamp(0, 1).numpy()
    result = []
    for frame in anchors:
        frame_u8 = np.uint8(np.round(frame * 255.0))
        result.append(cv2.resize(frame_u8, (416, 240), interpolation=cv2.INTER_AREA))
    return np.stack(result)


def maybe_save_preview(path: Path, anchors: np.ndarray) -> None:
    try:
        from torchvision.io import write_video

        repeated = np.repeat(anchors, 4, axis=0)
        tensor = torch.from_numpy(repeated)
        write_video(str(path), tensor, fps=16)
    except Exception as error:
        print(f"[preview] skipped: {error}", flush=True)


@torch.inference_mode()
def generate_snapshots(args: argparse.Namespace) -> list[Path]:
    output_dir = args.output_dir
    runs_dir = output_dir / "runs"
    runs_dir.mkdir(parents=True, exist_ok=True)
    prompts = read_prompts(resolve_path(args.prompt_path), args.num_prompts)
    expected_paths = [runs_dir / f"prompt_{index:02d}.pt" for index in range(len(prompts))]
    missing = [
        path
        for path in expected_paths
        if args.overwrite or not path.exists()
    ]
    if not missing:
        print("[generation] all prompt snapshots already exist", flush=True)
        return expected_paths

    pipeline, config = build_pipeline(args)
    denoising_timesteps = [
        float(value) for value in pipeline.denoising_step_list.detach().cpu().tolist()
    ]
    recorder = FeatureRecorder(
        model=pipeline.generator.model,
        layers=args.layers,
        denoising_timesteps=denoising_timesteps,
        num_frame_per_block=pipeline.num_frame_per_block,
        max_tokens=args.max_tokens,
        projection_dim=args.projection_dim,
    )
    metadata = {
        "physical_gpu": args.gpu,
        "config_path": str(resolve_path(args.config_path)),
        "checkpoint_path": str(resolve_path(args.checkpoint_path)),
        "use_ema": args.use_ema,
        "num_prompts": args.num_prompts,
        "num_frames": args.num_frames,
        "num_frame_per_block": pipeline.num_frame_per_block,
        "denoising_timesteps": denoising_timesteps,
        "layers": args.layers,
        "max_tokens": args.max_tokens,
        "projection_dim": args.projection_dim,
        "seed": args.seed,
        "dtype": "bfloat16",
    }
    (output_dir / "experiment_config.json").write_text(
        json.dumps(metadata, indent=2, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )

    try:
        for index, (prompt, path) in enumerate(zip(prompts, expected_paths)):
            if path.exists() and not args.overwrite:
                print(f"[generation] skip existing {path.name}", flush=True)
                continue
            recorder.reset()
            run_seed = args.seed if args.same_seed else args.seed + index
            set_seed(run_seed)
            noise = torch.randn(
                1,
                args.num_frames,
                16,
                EXPECTED_LATENT_HEIGHT,
                EXPECTED_LATENT_WIDTH,
                device="cuda",
                dtype=torch.bfloat16,
            )
            torch.cuda.reset_peak_memory_stats()
            torch.cuda.synchronize()
            start = time.perf_counter()
            print(
                f"[generation] prompt {index + 1}/{len(prompts)} seed={run_seed}",
                flush=True,
            )
            video, latents = pipeline.inference(
                noise=noise,
                text_prompts=[prompt],
                return_latents=True,
                initial_latent=None,
                low_memory=False,
            )
            torch.cuda.synchronize()
            elapsed = time.perf_counter() - start
            peak_gib = torch.cuda.max_memory_allocated() / (1024**3)
            anchors = downsample_anchors(video, args.num_frames)
            state = {
                "run_index": index,
                "prompt": prompt,
                "seed": run_seed,
                "elapsed_s": elapsed,
                "peak_gpu_gib": peak_gib,
                "num_frames": args.num_frames,
                "num_frame_per_block": pipeline.num_frame_per_block,
                "latents": latents[0].detach().to(
                    dtype=torch.float16, device="cpu"
                ),
                **recorder.state_dict(),
            }
            torch.save(state, path)
            np.savez_compressed(path.with_suffix(".anchors.npz"), frames=anchors)
            if args.save_preview:
                maybe_save_preview(path.with_suffix(".mp4"), anchors)
            print(
                f"[generation] saved {path.name}: {elapsed:.1f}s, peak={peak_gib:.1f} GiB",
                flush=True,
            )
            del video, latents, noise, state
            pipeline.vae.model.clear_cache()
            torch.cuda.empty_cache()
    finally:
        recorder.close()
    return expected_paths


def load_runs(paths: list[Path]) -> list[dict[str, Any]]:
    runs = []
    for path in paths:
        if not path.exists():
            raise FileNotFoundError(path)
        run = torch.load(path, map_location="cpu", weights_only=False)
        anchor_path = path.with_suffix(".anchors.npz")
        if not anchor_path.exists():
            raise FileNotFoundError(anchor_path)
        run["anchors"] = np.load(anchor_path, allow_pickle=False)["frames"]
        run["path"] = str(path)
        runs.append(run)
    return runs


def feature(run: dict[str, Any], stage: str, chunk: int, step: int) -> torch.Tensor:
    return run["records"][stage][f"{chunk}:{step}"].float()


def projected_feature(
    run: dict[str, Any], chunk: int, step: int
) -> torch.Tensor:
    return run["projected"][f"{chunk}:{step}"].float()


def available_chunks(run: dict[str, Any], stage: str) -> list[int]:
    return sorted(
        {int(key.split(":")[0]) for key in run["records"][stage].keys()}
    )


def available_steps(run: dict[str, Any], stage: str) -> list[int]:
    return sorted(
        {int(key.split(":")[1]) for key in run["records"][stage].keys()}
    )


def pair_metrics(
    reference: torch.Tensor,
    target: torch.Tensor,
    compute_cka: bool = True,
) -> dict[str, float]:
    if reference.shape != target.shape:
        raise ValueError(f"Pair shape mismatch: {reference.shape} vs {target.shape}")
    eps = 1e-8
    x = reference.float()
    y = target.float()
    xf = x.reshape(-1)
    yf = y.reshape(-1)
    diff = yf - xf
    cosine = F.cosine_similarity(xf[None], yf[None], dim=1, eps=eps)[0]
    xc_flat = xf - xf.mean()
    yc_flat = yf - yf.mean()
    centered_cosine = torch.dot(xc_flat, yc_flat) / (
        torch.linalg.vector_norm(xc_flat)
        * torch.linalg.vector_norm(yc_flat)
        + eps
    )
    token_cosine = F.cosine_similarity(x, y, dim=1, eps=eps)
    rel_l2 = diff.square().mean().sqrt() / (xf.square().mean().sqrt() + eps)
    nmse = diff.square().mean() / (yc_flat.square().mean() + eps)

    if compute_cka:
        xc = x - x.mean(dim=0, keepdim=True)
        yc = y - y.mean(dim=0, keepdim=True)
        gram_x = xc @ xc.T
        gram_y = yc @ yc.T
        cka = (gram_x * gram_y).sum() / (
            (gram_x.square().sum() * gram_y.square().sum()).sqrt() + eps
        )
    else:
        cka = torch.tensor(float("nan"))
    quantiles = torch.quantile(
        token_cosine,
        torch.tensor([0.1, 0.5, 0.9], dtype=token_cosine.dtype),
    )
    return {
        "cosine": float(cosine),
        "centered_cosine": float(centered_cosine),
        "linear_cka": float(cka),
        "rel_l2": float(rel_l2),
        "nmse": float(nmse),
        "token_cosine_mean": float(token_cosine.mean()),
        "token_cosine_p10": float(quantiles[0]),
        "token_cosine_p50": float(quantiles[1]),
        "token_cosine_p90": float(quantiles[2]),
        "reference_rms": float(xf.square().mean().sqrt()),
        "target_rms": float(yf.square().mean().sqrt()),
    }


def collect_pair_rows(
    runs: list[dict[str, Any]], cosine_only: bool = False
) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    stages = sorted(runs[0]["records"])
    if cosine_only:
        stages = [
            stage
            for stage in stages
            if stage.startswith("block_") and stage.endswith("_hidden")
        ]
    for run_index, run in enumerate(runs):
        shuffled_run = runs[(run_index + 1) % len(runs)]
        for stage in stages:
            chunks = available_chunks(run, stage)
            steps = available_steps(run, stage)

            def append(
                comparison: str,
                ref_run: dict[str, Any],
                ref_chunk: int,
                ref_step: int,
                target_chunk: int,
                target_step: int,
            ) -> None:
                metrics = pair_metrics(
                    feature(ref_run, stage, ref_chunk, ref_step),
                    feature(run, stage, target_chunk, target_step),
                    compute_cka=not cosine_only,
                )
                rows.append(
                    {
                        "run": run_index,
                        "comparison": comparison,
                        "stage": stage,
                        "reference_chunk": ref_chunk,
                        "target_chunk": target_chunk,
                        "reference_step": ref_step,
                        "target_step": target_step,
                        **metrics,
                    }
                )

            for chunk in chunks:
                for step in steps[1:]:
                    append(
                        "within_adjacent",
                        run,
                        chunk,
                        step - 1,
                        chunk,
                        step,
                    )
                if chunk < 1:
                    continue
                for step in steps:
                    append(
                        "cross_same",
                        run,
                        chunk - 1,
                        step,
                        chunk,
                        step,
                    )
                    if cosine_only:
                        continue
                    append(
                        "cross_video_shuffle",
                        shuffled_run,
                        chunk - 1,
                        step,
                        chunk,
                        step,
                    )
                    if step > 0:
                        append(
                            "cross_wrong_step",
                            run,
                            chunk - 1,
                            step - 1,
                            chunk,
                            step,
                        )
                    if chunk > 1:
                        append(
                            "cross_distant",
                            run,
                            chunk - 2,
                            step,
                            chunk,
                            step,
                        )
    return rows


def predictor_columns(
    name: str,
    within: torch.Tensor,
    cross: torch.Tensor,
    distant: torch.Tensor,
    wrong: torch.Tensor,
    batch: torch.Tensor,
    noise_seed: int,
) -> list[torch.Tensor]:
    ones = torch.ones_like(within)
    shifted = cross.roll(shifts=max(1, cross.shape[0] // 2), dims=0)
    generator = torch.Generator(device="cpu").manual_seed(noise_seed)
    noise = torch.randn(
        cross.shape, generator=generator, dtype=cross.dtype
    )
    noise = noise * cross.std(dim=0, keepdim=True).clamp_min(1e-6)
    noise = noise + cross.mean(dim=0, keepdim=True)
    mapping = {
        "within_affine": [within, ones],
        "within_quadratic": [within, within.square(), ones],
        "cross_affine": [cross, ones],
        "fusion_same": [within, cross, ones],
        "fusion_distant": [within, distant, ones],
        "fusion_token_shift": [within, shifted, ones],
        "fusion_wrong_step": [within, wrong, ones],
        "fusion_batch_shuffle": [within, batch, ones],
        "fusion_zero": [within, torch.zeros_like(cross), ones],
        "fusion_noise": [within, noise, ones],
    }
    return mapping[name]


PROBE_NAMES = [
    "within_affine",
    "within_quadratic",
    "cross_affine",
    "fusion_same",
    "fusion_distant",
    "fusion_token_shift",
    "fusion_wrong_step",
    "fusion_batch_shuffle",
    "fusion_zero",
    "fusion_noise",
]


def gather_probe_data(
    runs: list[dict[str, Any]],
    run_indices: list[int],
    stage: str,
    step: int,
) -> dict[str, torch.Tensor]:
    buckets: dict[str, list[torch.Tensor]] = defaultdict(list)
    for run_index in run_indices:
        run = runs[run_index]
        other = runs[(run_index + 1) % len(runs)]
        chunks = available_chunks(run, stage)
        for chunk in chunks:
            # Use c >= 2 for every probe so correct, distant, and all other
            # controls are evaluated on exactly the same target tokens.
            if chunk < 2:
                continue
            buckets["target"].append(feature(run, stage, chunk, step))
            buckets["within"].append(feature(run, stage, chunk, step - 1))
            buckets["cross"].append(feature(run, stage, chunk - 1, step))
            buckets["distant"].append(feature(run, stage, chunk - 2, step))
            buckets["wrong"].append(feature(run, stage, chunk - 1, step - 1))
            buckets["batch"].append(feature(other, stage, chunk - 1, step))
    if not buckets:
        raise ValueError(f"No probe data for stage={stage}, step={step}")
    return {key: torch.cat(values, dim=0).float() for key, values in buckets.items()}


def fit_channelwise_probe(
    columns: list[torch.Tensor],
    target: torch.Tensor,
    ridge: float,
) -> torch.Tensor:
    design = torch.stack(columns, dim=-1).double()
    y = target.double()
    gram = torch.einsum("ndp,ndq->dpq", design, design)
    rhs = torch.einsum("ndp,nd->dp", design, y)
    feature_count = gram.shape[-1]
    diagonal_scale = (
        gram.diagonal(dim1=-2, dim2=-1).mean(dim=-1).clamp_min(1e-8)
    )
    regularizer = (
        torch.eye(feature_count, dtype=gram.dtype)[None]
        * (ridge * diagonal_scale)[:, None, None]
    )
    regularizer[:, -1, -1] = 0.0
    try:
        weights = torch.linalg.solve(gram + regularizer, rhs.unsqueeze(-1)).squeeze(-1)
    except torch.linalg.LinAlgError:
        weights = (
            torch.linalg.pinv(gram + regularizer) @ rhs.unsqueeze(-1)
        ).squeeze(-1)
    return weights.float()


def apply_channelwise_probe(
    columns: list[torch.Tensor], weights: torch.Tensor
) -> torch.Tensor:
    design = torch.stack(columns, dim=-1).float()
    return torch.einsum("ndp,dp->nd", design, weights)


def prediction_metrics(prediction: torch.Tensor, target: torch.Tensor) -> dict[str, float]:
    eps = 1e-8
    pred = prediction.float()
    y = target.float()
    error = pred - y
    mse = error.square().mean()
    variance = (y - y.mean()).square().mean()
    nrmse = mse.sqrt() / (variance.sqrt() + eps)
    r2 = 1.0 - mse / (variance + eps)
    cosine = F.cosine_similarity(
        pred.reshape(1, -1), y.reshape(1, -1), dim=1, eps=eps
    )[0]
    return {
        "mse": float(mse),
        "nrmse": float(nrmse),
        "r2": float(r2),
        "cosine": float(cosine),
    }


def run_conditional_probes(
    runs: list[dict[str, Any]], ridge: float
) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    stages = sorted(runs[0]["records"])
    steps = available_steps(runs[0], stages[0])
    for stage in stages:
        for step in steps[1:]:
            for held_out in range(len(runs)):
                train_indices = [index for index in range(len(runs)) if index != held_out]
                train = gather_probe_data(runs, train_indices, stage, step)
                test = gather_probe_data(runs, [held_out], stage, step)
                for probe_name in PROBE_NAMES:
                    train_columns = predictor_columns(
                        probe_name,
                        train["within"],
                        train["cross"],
                        train["distant"],
                        train["wrong"],
                        train["batch"],
                        noise_seed=1000 + held_out * 100 + step,
                    )
                    test_columns = predictor_columns(
                        probe_name,
                        test["within"],
                        test["cross"],
                        test["distant"],
                        test["wrong"],
                        test["batch"],
                        noise_seed=2000 + held_out * 100 + step,
                    )
                    weights = fit_channelwise_probe(
                        train_columns, train["target"], ridge=ridge
                    )
                    prediction = apply_channelwise_probe(test_columns, weights)
                    rows.append(
                        {
                            "held_out_run": held_out,
                            "stage": stage,
                            "step": step,
                            "probe": probe_name,
                            "train_tokens": int(train["target"].shape[0]),
                            "test_tokens": int(test["target"].shape[0]),
                            **prediction_metrics(prediction, test["target"]),
                        }
                    )
    baseline_lookup = {
        (row["held_out_run"], row["stage"], row["step"], row["probe"]): row
        for row in rows
        if row["probe"] in {"within_affine", "within_quadratic"}
    }
    for row in rows:
        key = (row["held_out_run"], row["stage"], row["step"])
        for baseline_name in ("within_affine", "within_quadratic"):
            baseline = baseline_lookup[(*key, baseline_name)]
            row[f"mse_reduction_vs_{baseline_name}"] = (
                baseline["mse"] - row["mse"]
            ) / max(baseline["mse"], 1e-12)
            row[f"r2_gain_vs_{baseline_name}"] = row["r2"] - baseline["r2"]
    return rows


def gray(frame: np.ndarray) -> np.ndarray:
    return cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)


def farneback(source: np.ndarray, target: np.ndarray) -> np.ndarray:
    return cv2.calcOpticalFlowFarneback(
        gray(source),
        gray(target),
        None,
        pyr_scale=0.5,
        levels=4,
        winsize=21,
        iterations=5,
        poly_n=7,
        poly_sigma=1.5,
        flags=0,
    )


def resize_flow(flow: np.ndarray, height: int, width: int) -> torch.Tensor:
    source_height, source_width = flow.shape[:2]
    resized = cv2.resize(flow, (width, height), interpolation=cv2.INTER_AREA)
    resized[..., 0] *= width / source_width
    resized[..., 1] *= height / source_height
    return torch.from_numpy(resized).permute(2, 0, 1).float()


def warp_feature(
    source: torch.Tensor, target_to_source_flow: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    height, width, _ = source.shape
    yy, xx = torch.meshgrid(
        torch.arange(height, dtype=torch.float32),
        torch.arange(width, dtype=torch.float32),
        indexing="ij",
    )
    sample_x = xx + target_to_source_flow[0]
    sample_y = yy + target_to_source_flow[1]
    grid = torch.stack(
        [
            2.0 * sample_x / max(width - 1, 1) - 1.0,
            2.0 * sample_y / max(height - 1, 1) - 1.0,
        ],
        dim=-1,
    )[None]
    value = source.permute(2, 0, 1)[None].float()
    warped = F.grid_sample(
        value, grid, mode="bilinear", padding_mode="zeros", align_corners=True
    )[0].permute(1, 2, 0)
    mask = (
        (sample_x >= 0)
        & (sample_x <= width - 1)
        & (sample_y >= 0)
        & (sample_y <= height - 1)
    )
    return warped, mask


def masked_cosine(
    left: torch.Tensor, right: torch.Tensor, mask: torch.Tensor | None = None
) -> float:
    value = F.cosine_similarity(left.float(), right.float(), dim=-1, eps=1e-8)
    if mask is not None:
        if not bool(mask.any()):
            return float("nan")
        value = value[mask]
    return float(value.mean())


def collect_motion_rows(runs: list[dict[str, Any]]) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for run_index, run in enumerate(runs):
        anchors = run["anchors"]
        chunk_size = int(run["num_frame_per_block"])
        chunk_count = int(run["num_frames"]) // chunk_size
        steps = sorted(
            {int(key.split(":")[1]) for key in run["projected"].keys()}
        )
        for chunk in range(1, chunk_count):
            previous_anchor_index = chunk * chunk_size - 1
            previous_frame = anchors[previous_anchor_index]
            boundary_flow = farneback(
                previous_frame, anchors[chunk * chunk_size]
            )
            median_flow = np.median(
                boundary_flow.reshape(-1, 2), axis=0
            )
            camera_motion = float(np.linalg.norm(median_flow))
            residual = boundary_flow - median_flow[None, None]
            object_motion = float(
                np.linalg.norm(residual, axis=-1).mean()
            )
            total_motion = float(
                np.linalg.norm(boundary_flow, axis=-1).mean()
            )

            for step in steps:
                source_map = projected_feature(run, chunk - 1, step)[-1]
                target_map = projected_feature(run, chunk, step)
                same_slot_source = projected_feature(run, chunk - 1, step)
                per_slot: list[dict[str, float]] = []
                for slot in range(chunk_size):
                    target = target_map[slot]
                    same_slot = same_slot_source[slot]
                    raw_boundary = source_map
                    target_frame = anchors[chunk * chunk_size + slot]
                    backward = farneback(target_frame, previous_frame)
                    feature_flow = resize_flow(
                        backward, FRAME_TOKEN_HEIGHT, FRAME_TOKEN_WIDTH
                    )
                    global_flow = torch.zeros_like(feature_flow)
                    global_flow[0].fill_(float(np.median(feature_flow[0].numpy())))
                    global_flow[1].fill_(float(np.median(feature_flow[1].numpy())))
                    global_aligned, global_mask = warp_feature(
                        source_map, global_flow
                    )
                    flow_aligned, flow_mask = warp_feature(source_map, feature_flow)
                    per_slot.append(
                        {
                            "same_slot_cosine": masked_cosine(target, same_slot),
                            "boundary_raw_cosine": masked_cosine(
                                target, raw_boundary
                            ),
                            "global_aligned_cosine": masked_cosine(
                                target, global_aligned, global_mask
                            ),
                            "flow_aligned_cosine": masked_cosine(
                                target, flow_aligned, flow_mask
                            ),
                            "valid_flow_ratio": float(flow_mask.float().mean()),
                        }
                    )
                rows.append(
                    {
                        "run": run_index,
                        "chunk": chunk,
                        "step": step,
                        "total_motion": total_motion,
                        "camera_motion": camera_motion,
                        "object_motion": object_motion,
                        **{
                            key: float(np.nanmean([item[key] for item in per_slot]))
                            for key in per_slot[0]
                        },
                    }
                )
    motion_values = np.asarray([row["total_motion"] for row in rows])
    if len(motion_values) >= 3:
        low, high = np.quantile(motion_values, [1 / 3, 2 / 3])
        for row in rows:
            value = row["total_motion"]
            row["motion_bin"] = "low" if value <= low else "high" if value > high else "medium"
    return rows


def group_mean(
    rows: list[dict[str, Any]], keys: list[str], metrics: list[str]
) -> list[dict[str, Any]]:
    groups: dict[tuple[Any, ...], list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        groups[tuple(row[key] for key in keys)].append(row)
    result = []
    for group, values in sorted(groups.items(), key=lambda item: tuple(map(str, item[0]))):
        output = {key: value for key, value in zip(keys, group)}
        output["count"] = len(values)
        for metric in metrics:
            finite = [
                float(row[metric])
                for row in values
                if metric in row and math.isfinite(float(row[metric]))
            ]
            output[metric] = float(np.mean(finite)) if finite else float("nan")
        result.append(output)
    return result


def paired_probe_reduction(
    rows: list[dict[str, Any]],
    stage: str,
    reference_probe: str,
    candidate_probe: str = "fusion_same",
) -> dict[str, Any]:
    lookup = {
        (int(row["held_out_run"]), int(row["step"]), row["probe"]): row
        for row in rows
        if row["stage"] == stage
        and row["probe"] in {reference_probe, candidate_probe}
    }
    pairs = sorted(
        {
            (held_out, step)
            for held_out, step, probe in lookup
            if probe == candidate_probe
            and (held_out, step, reference_probe) in lookup
        }
    )
    reductions = []
    for held_out, step in pairs:
        reference = float(lookup[(held_out, step, reference_probe)]["mse"])
        candidate = float(lookup[(held_out, step, candidate_probe)]["mse"])
        reductions.append((reference - candidate) / max(reference, 1e-12))
    return {
        "reference": reference_probe,
        "paired_count": len(reductions),
        "mean_mse_reduction": (
            float(np.mean(reductions)) if reductions else float("nan")
        ),
        "median_mse_reduction": (
            float(np.median(reductions)) if reductions else float("nan")
        ),
        "wins": int(sum(value > 0 for value in reductions)),
    }


def plot_similarity(pair_rows: list[dict[str, Any]], output_dir: Path) -> None:
    stages = [
        stage
        for stage in sorted({row["stage"] for row in pair_rows})
        if stage.endswith("_hidden")
    ]
    comparisons = [
        "within_adjacent",
        "cross_same",
        "cross_wrong_step",
        "cross_distant",
        "cross_video_shuffle",
    ]
    steps = sorted({int(row["target_step"]) for row in pair_rows})
    fig, axes = plt.subplots(
        len(stages), 2, figsize=(12, max(3.2, 2.8 * len(stages))), squeeze=False
    )
    for stage_index, stage in enumerate(stages):
        for metric_index, metric in enumerate(["cosine", "nmse"]):
            matrix = np.full((len(comparisons), len(steps)), np.nan)
            for row_index, comparison in enumerate(comparisons):
                for col_index, step in enumerate(steps):
                    values = [
                        float(row[metric])
                        for row in pair_rows
                        if row["stage"] == stage
                        and row["comparison"] == comparison
                        and int(row["target_step"]) == step
                    ]
                    if values:
                        matrix[row_index, col_index] = np.mean(values)
            ax = axes[stage_index, metric_index]
            image = ax.imshow(
                matrix,
                aspect="auto",
                cmap="viridis_r" if metric == "nmse" else "viridis",
            )
            ax.set_title(f"{stage}: {metric}")
            ax.set_xticks(range(len(steps)), labels=steps)
            ax.set_yticks(range(len(comparisons)), labels=comparisons)
            ax.set_xlabel("target denoising step")
            fig.colorbar(image, ax=ax, fraction=0.03)
    fig.tight_layout()
    fig.savefig(output_dir / "feature_redundancy_heatmap.png", dpi=180)
    plt.close(fig)


def plot_probe_gain(probe_rows: list[dict[str, Any]], output_dir: Path) -> None:
    aggregate = group_mean(
        probe_rows,
        ["stage", "step", "probe"],
        ["mse", "nrmse", "r2", "mse_reduction_vs_within_quadratic"],
    )
    stages = [
        stage
        for stage in sorted({row["stage"] for row in aggregate})
        if stage.endswith("_hidden")
    ]
    probes = [
        "fusion_same",
        "fusion_distant",
        "fusion_token_shift",
        "fusion_wrong_step",
        "fusion_batch_shuffle",
    ]
    steps = sorted({int(row["step"]) for row in aggregate})
    fig, axes = plt.subplots(
        len(stages), 1, figsize=(9, max(3.4, 3.0 * len(stages))), squeeze=False
    )
    for stage_index, stage in enumerate(stages):
        ax = axes[stage_index, 0]
        for probe in probes:
            values = []
            for step in steps:
                match = [
                    row
                    for row in aggregate
                    if row["stage"] == stage
                    and int(row["step"]) == step
                    and row["probe"] == probe
                ]
                values.append(
                    match[0]["mse_reduction_vs_within_quadratic"]
                    if match
                    else np.nan
                )
            ax.plot(steps, values, marker="o", label=probe)
        ax.axhline(0, color="black", linewidth=0.8)
        ax.set_title(stage)
        ax.set_ylabel("MSE reduction vs within quadratic")
        ax.set_xlabel("target denoising step")
        ax.legend(fontsize=8)
        ax.grid(alpha=0.25)
    fig.tight_layout()
    fig.savefig(output_dir / "conditional_chunk_gain.png", dpi=180)
    plt.close(fig)


def plot_motion(motion_rows: list[dict[str, Any]], output_dir: Path) -> None:
    fig, axes = plt.subplots(1, 2, figsize=(12, 4.5))
    axes[0].scatter(
        [row["total_motion"] for row in motion_rows],
        [row["boundary_raw_cosine"] for row in motion_rows],
        s=18,
        alpha=0.65,
        label="boundary raw",
    )
    axes[0].scatter(
        [row["total_motion"] for row in motion_rows],
        [row["flow_aligned_cosine"] for row in motion_rows],
        s=18,
        alpha=0.65,
        label="dense-flow aligned",
    )
    axes[0].set_xlabel("optical-flow magnitude")
    axes[0].set_ylabel("feature cosine")
    axes[0].legend()
    axes[0].grid(alpha=0.25)

    aggregate = group_mean(
        motion_rows,
        ["motion_bin"],
        [
            "same_slot_cosine",
            "boundary_raw_cosine",
            "global_aligned_cosine",
            "flow_aligned_cosine",
        ],
    )
    bins = ["low", "medium", "high"]
    metrics = [
        "same_slot_cosine",
        "boundary_raw_cosine",
        "global_aligned_cosine",
        "flow_aligned_cosine",
    ]
    width = 0.18
    x = np.arange(len(bins))
    for metric_index, metric in enumerate(metrics):
        values = []
        for name in bins:
            match = [row for row in aggregate if row["motion_bin"] == name]
            values.append(match[0][metric] if match else np.nan)
        axes[1].bar(
            x + (metric_index - 1.5) * width,
            values,
            width=width,
            label=metric.replace("_cosine", ""),
        )
    axes[1].set_xticks(x, bins)
    axes[1].set_ylabel("mean feature cosine")
    axes[1].set_xlabel("motion bin")
    axes[1].legend(fontsize=8)
    axes[1].grid(axis="y", alpha=0.25)
    fig.tight_layout()
    fig.savefig(output_dir / "motion_alignment_analysis.png", dpi=180)
    plt.close(fig)


def markdown_table(
    rows: list[dict[str, Any]], columns: list[str], digits: int = 4
) -> str:
    lines = [
        "| " + " | ".join(columns) + " |",
        "|" + "|".join(["---"] * len(columns)) + "|",
    ]
    for row in rows:
        values = []
        for column in columns:
            value = row.get(column, "")
            if isinstance(value, float):
                values.append(f"{value:.{digits}f}")
            else:
                values.append(str(value))
        lines.append("| " + " | ".join(values) + " |")
    return "\n".join(lines)


def build_report(
    runs: list[dict[str, Any]],
    pair_rows: list[dict[str, Any]],
    probe_rows: list[dict[str, Any]],
    motion_rows: list[dict[str, Any]],
    output_dir: Path,
) -> None:
    final_stage = f"block_{runs[0]['projection_layer']}_hidden"
    pair_summary = group_mean(
        [
            row
            for row in pair_rows
            if row["stage"] == final_stage
        ],
        ["comparison"],
        ["cosine", "linear_cka", "rel_l2", "nmse", "token_cosine_p10"],
    )
    probe_summary = group_mean(
        [
            row
            for row in probe_rows
            if row["stage"] == final_stage
            and row["probe"] in set(PROBE_NAMES)
        ],
        ["probe"],
        [
            "nrmse",
            "r2",
            "mse_reduction_vs_within_affine",
            "mse_reduction_vs_within_quadratic",
        ],
    )
    motion_summary = group_mean(
        motion_rows,
        ["motion_bin"],
        [
            "total_motion",
            "same_slot_cosine",
            "boundary_raw_cosine",
            "global_aligned_cosine",
            "flow_aligned_cosine",
        ],
    )
    control_order = [
        "within_affine",
        "within_quadratic",
        "fusion_distant",
        "fusion_wrong_step",
        "fusion_token_shift",
        "fusion_batch_shuffle",
        "fusion_zero",
        "fusion_noise",
    ]
    control_summary = [
        paired_probe_reduction(
            probe_rows,
            stage=final_stage,
            reference_probe=control,
        )
        for control in control_order
    ]
    control_lookup = {row["reference"]: row for row in control_summary}
    affine = control_lookup["within_affine"]
    shuffled = control_lookup["fusion_batch_shuffle"]
    conditional_statement = (
        f"在主层 `{final_stage}` 上，加入正确的前一 chunk 同 timestep 特征，"
        f"相对 within-only affine probe 平均降低 held-out MSE "
        f"{100*affine['mean_mse_reduction']:.2f}%，并在 "
        f"{affine['wins']}/{affine['paired_count']} 个 "
        f"prompt–timestep 配对中取得改善。相对参数量一致的跨视频 "
        f"shuffle 对照，MSE 平均降低 "
        f"{100*shuffled['mean_mse_reduction']:.2f}%。"
    )

    pair_lookup = {row["comparison"]: row for row in pair_summary}
    same_pair = pair_lookup.get("cross_same")
    random_pair = pair_lookup.get("cross_video_shuffle")
    pair_statement = ""
    if same_pair is not None and random_pair is not None:
        pair_statement = (
            f"正确相邻 chunk 的平均 cosine/CKA 为 "
            f"{same_pair['cosine']:.4f}/{same_pair['linear_cka']:.4f}，"
            f"跨视频 shuffle 为 "
            f"{random_pair['cosine']:.4f}/{random_pair['linear_cka']:.4f}。"
        )

    total_motion = np.asarray(
        [float(row["total_motion"]) for row in motion_rows], dtype=np.float64
    )
    raw_cosine = np.asarray(
        [float(row["boundary_raw_cosine"]) for row in motion_rows],
        dtype=np.float64,
    )
    flow_cosine = np.asarray(
        [float(row["flow_aligned_cosine"]) for row in motion_rows],
        dtype=np.float64,
    )
    motion_correlation = (
        float(np.corrcoef(total_motion, raw_cosine)[0, 1])
        if len(motion_rows) > 1
        else float("nan")
    )
    flow_gain = float(np.mean(flow_cosine - raw_cosine))
    flow_wins = int(np.sum(flow_cosine > raw_cosine))
    motion_statement = (
        f"运动强度与未对齐跨 chunk cosine 的 Pearson 相关系数为 "
        f"{motion_correlation:.3f}；dense-flow 对齐平均恢复 "
        f"{flow_gain:.4f} cosine，并在 {flow_wins}/{len(motion_rows)} "
        f"个 chunk–timestep 样本上改善。"
    )

    report = f"""# Self-Forcing Feature Cache 分析结果

## 实验配置

- Prompts：{len(runs)}
- 每条视频 latent frames：{runs[0]['num_frames']}
- 每个 AR chunk latent frames：{runs[0]['num_frame_per_block']}
- Denoising timesteps：{runs[0]['timesteps']}
- Hook layers：{runs[0]['layers']}
- 主分析 stage：`{final_stage}`

## 核心观察

{conditional_statement}

{pair_statement}

{motion_statement}

以下结果是 3 条 prompt 的 pilot 分析。条件 probe 按 prompt 留一测试，
统计单位是 held-out prompt 与 timestep；不能将 token 数量解释为独立视频样本数，
也不能据此声称数据集级统计显著性。

## 主层特征配对

{markdown_table(pair_summary, ['comparison', 'count', 'cosine', 'linear_cka', 'rel_l2', 'nmse', 'token_cosine_p10'])}

## 条件 Probe

Probe 仅使用 chunk index `c ≥ 2` 的目标 chunk，使 `correct`、`c-2 distant`
及其他控制组在完全相同的 token 上比较。

{markdown_table(probe_summary, ['probe', 'count', 'nrmse', 'r2', 'mse_reduction_vs_within_affine', 'mse_reduction_vs_within_quadratic'])}

### `fusion_same` 的成对控制实验

正值表示正确前一 chunk 同 timestep 输入的 MSE 更低。

{markdown_table(control_summary, ['reference', 'paired_count', 'mean_mse_reduction', 'median_mse_reduction', 'wins'])}

## 运动与对齐

{markdown_table(motion_summary, ['motion_bin', 'count', 'total_motion', 'same_slot_cosine', 'boundary_raw_cosine', 'global_aligned_cosine', 'flow_aligned_cosine'])}

这里的 `global_aligned` 是由光流中位数估计的全局平移对齐，
`flow_aligned` 是 dense optical-flow oracle；本轮未实现 homography。

## 本轮结论边界

- 已完成：四步 DMD 模型的 hidden/residual-delta 特征采集、相似度与
  nMSE/CKA、linear/ridge 条件 probe、负对照、运动分桶和光流对齐。
- 尚未完成：接入用户现有的小型非线性预测网络、真实 cache 替换干预、
  最终视频质量与端到端加速评估、规模化多视频置信区间。
- 因而当前结果支持“前一 chunk 提供额外且具有空间对应性的条件信息”，
  但还不能单独证明最终生成质量或真实加速收益。

## 输出文件

- `feature_pair_metrics.csv`
- `feature_pair_summary.csv`
- `conditional_probe_folds.csv`
- `conditional_probe_summary.csv`
- `motion_alignment_metrics.csv`
- `motion_alignment_summary.csv`
- `feature_redundancy_heatmap.png`
- `conditional_chunk_gain.png`
- `motion_alignment_analysis.png`
"""
    (output_dir / "REPORT.md").write_text(report, encoding="utf-8")


def analyze(args: argparse.Namespace, paths: list[Path]) -> None:
    print("[analysis] loading snapshots", flush=True)
    runs = load_runs(paths)
    pair_rows = collect_pair_rows(runs, cosine_only=args.cosine_only)
    if args.cosine_only:
        pair_summary = group_mean(
            pair_rows,
            ["comparison", "stage", "target_step"],
            [
                "cosine",
                "centered_cosine",
                "linear_cka",
                "rel_l2",
                "nmse",
                "token_cosine_mean",
                "token_cosine_p10",
                "token_cosine_p50",
                "token_cosine_p90",
            ],
        )
        write_csv(args.output_dir / "feature_pair_metrics.csv", pair_rows)
        write_csv(args.output_dir / "feature_pair_summary.csv", pair_summary)
        plot_similarity(pair_rows, args.output_dir)
        summary = {
            "runs": len(runs),
            "pair_rows": len(pair_rows),
            "cosine_only": True,
        }
        (args.output_dir / "summary.json").write_text(
            json.dumps(summary, indent=2) + "\n", encoding="utf-8"
        )
        print(f"[analysis] cosine-only complete: {args.output_dir}", flush=True)
        return
    probe_rows = run_conditional_probes(runs, ridge=args.ridge)
    motion_rows = collect_motion_rows(runs)

    pair_summary = group_mean(
        pair_rows,
        ["comparison", "stage", "target_step"],
        [
            "cosine",
            "centered_cosine",
            "linear_cka",
            "rel_l2",
            "nmse",
            "token_cosine_mean",
            "token_cosine_p10",
            "token_cosine_p50",
            "token_cosine_p90",
        ],
    )
    probe_summary = group_mean(
        probe_rows,
        ["stage", "step", "probe"],
        [
            "mse",
            "nrmse",
            "r2",
            "cosine",
            "mse_reduction_vs_within_affine",
            "mse_reduction_vs_within_quadratic",
            "r2_gain_vs_within_affine",
            "r2_gain_vs_within_quadratic",
        ],
    )
    motion_summary = group_mean(
        motion_rows,
        ["motion_bin", "step"],
        [
            "total_motion",
            "camera_motion",
            "object_motion",
            "same_slot_cosine",
            "boundary_raw_cosine",
            "global_aligned_cosine",
            "flow_aligned_cosine",
            "valid_flow_ratio",
        ],
    )

    write_csv(args.output_dir / "feature_pair_metrics.csv", pair_rows)
    write_csv(args.output_dir / "feature_pair_summary.csv", pair_summary)
    write_csv(args.output_dir / "conditional_probe_folds.csv", probe_rows)
    write_csv(args.output_dir / "conditional_probe_summary.csv", probe_summary)
    write_csv(args.output_dir / "motion_alignment_metrics.csv", motion_rows)
    write_csv(args.output_dir / "motion_alignment_summary.csv", motion_summary)
    plot_similarity(pair_rows, args.output_dir)
    plot_probe_gain(probe_rows, args.output_dir)
    plot_motion(motion_rows, args.output_dir)
    build_report(runs, pair_rows, probe_rows, motion_rows, args.output_dir)
    summary = {
        "runs": len(runs),
        "pair_rows": len(pair_rows),
        "probe_rows": len(probe_rows),
        "motion_rows": len(motion_rows),
        "report": str(args.output_dir / "REPORT.md"),
    }
    (args.output_dir / "summary.json").write_text(
        json.dumps(summary, indent=2) + "\n", encoding="utf-8"
    )
    print(f"[analysis] complete: {args.output_dir / 'REPORT.md'}", flush=True)


def main() -> None:
    args = parse_args()
    args.config_path = resolve_path(args.config_path)
    args.checkpoint_path = resolve_path(args.checkpoint_path)
    args.prompt_path = resolve_path(args.prompt_path)
    args.output_dir = resolve_path(args.output_dir)
    args.output_dir.mkdir(parents=True, exist_ok=True)
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.set_grad_enabled(False)

    run_paths = [
        args.output_dir / "runs" / f"prompt_{index:02d}.pt"
        for index in range(args.num_prompts)
    ]
    if not args.analysis_only:
        run_paths = generate_snapshots(args)
    analyze(args, run_paths)


if __name__ == "__main__":
    main()