path: root/worldalign/ricci_control.py

"""Ricci-flow control for the on-manifold assignment gate.

Hypothesis under test: a discrete Ricci flow that smooths each modality's
relational geometry before matching could repair the local ordering that
static relation fields fail. The flow is run independently per modality
with identical hyperparameters; hidden pairs never touch the flow or the
energy and only score orderings, as in the base gate.

Two geometries are gated per condition: heat-kernel channels on the raw
kNN graph (diffusion without flow) and the same construction after
Ollivier-Ricci weight evolution. Differences between them are attributable
to the flow itself rather than to the diffusion representation.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import ot
import torch
from scipy.sparse import csr_matrix
from scipy.sparse.csgraph import shortest_path

from .common import seed_everything, write_json
from .manifold_gate import (
    gate_a_global_ranking,
    gate_b_transpositions,
    load_flickr,
    load_vg,
    standardize_relation,
    steepest_descent,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", choices=["flickr", "vg"], default="flickr")
    parser.add_argument("--manifest", default="artifacts/manifest.json")
    parser.add_argument("--vision", default="artifacts/vision.pt")
    parser.add_argument("--text", default="artifacts/text.pt")
    parser.add_argument("--text-orbits", default="artifacts/text_orbits_qwen0p5b.pt")
    parser.add_argument("--text-mode", choices=["single", "orbit_mean"], default="orbit_mean")
    parser.add_argument("--vg-vision", default="artifacts/vg_5k/vision_features.pt")
    parser.add_argument("--vg-text", default="artifacts/vg_5k/text_features.pt")
    parser.add_argument(
        "--vg-ground-truth", default="artifacts/vg_5k/ground_truth.private.jsonl"
    )
    parser.add_argument("--vg-bundle-channels", action="store_true")
    parser.add_argument("--split", choices=["val", "test"], default="test")
    parser.add_argument("--samples", type=int, default=512)
    parser.add_argument("--subset-seed", type=int, default=0)
    parser.add_argument("--neighbors", type=int, default=16)
    parser.add_argument("--flow-iterations", type=int, default=10)
    parser.add_argument("--flow-step", type=float, default=0.4)
    parser.add_argument("--lazy-alpha", type=float, default=0.5)
    parser.add_argument("--heat-times", default="1.0,4.0")
    parser.add_argument("--random-perms", type=int, default=300)
    parser.add_argument("--derangement-samples", type=int, default=50)
    parser.add_argument("--descent-restarts", type=int, default=2)
    parser.add_argument("--descent-max-steps", type=int, default=200000)
    parser.add_argument("--descent-objective", default="mse", choices=["mse", "m30_total"])
    parser.add_argument("--descent-verify-top", type=int, default=64)
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--seed", type=int, default=20260729)
    parser.add_argument("--output", default="artifacts/manifold_gate/ricci_control.json")
    parser.add_argument("--trajectory-output")
    return parser.parse_args()


def knn_edges(distance: np.ndarray, k: int) -> list[tuple[int, int]]:
    order = distance.argsort(-1)[:, 1 : k + 1]
    edges = {
        (min(i, int(j)), max(i, int(j)))
        for i in range(len(distance))
        for j in order[i]
    }
    return sorted(edges)


def graph_apsp(size: int, weights: dict[tuple[int, int], float]) -> np.ndarray:
    rows, cols, vals = [], [], []
    for (i, j), w in weights.items():
        rows.extend((i, j))
        cols.extend((j, i))
        vals.extend((w, w))
    graph = csr_matrix((vals, (rows, cols)), shape=(size, size))
    return shortest_path(graph, method="D", directed=False)


def ollivier_ricci_apsp(
    distance: np.ndarray, args: argparse.Namespace, iterations: int
) -> np.ndarray:
    """All-pairs geodesics after Ollivier-Ricci weight evolution.

    Lazy uniform neighbor measures, W1 ground costs from current geodesics,
    multiplicative weight update w <- w * (1 - step * kappa), total edge
    mass renormalized each iteration. iterations=0 gives the un-flowed
    graph geometry for the diffusion-only control.
    """
    edges = knn_edges(distance, args.neighbors)
    weights = {edge: max(float(distance[edge]), 1e-9) for edge in edges}
    total = sum(weights.values())
    neighbor_map: dict[int, list[int]] = {}
    for i, j in edges:
        neighbor_map.setdefault(i, []).append(j)
        neighbor_map.setdefault(j, []).append(i)
    apsp = graph_apsp(len(distance), weights)
    for _ in range(iterations):
        updated: dict[tuple[int, int], float] = {}
        for i, j in edges:
            support_i = [i] + neighbor_map[i]
            support_j = [j] + neighbor_map[j]
            mass_i = np.full(len(support_i), (1 - args.lazy_alpha) / len(neighbor_map[i]))
            mass_i[0] = args.lazy_alpha
            mass_j = np.full(len(support_j), (1 - args.lazy_alpha) / len(neighbor_map[j]))
            mass_j[0] = args.lazy_alpha
            ground = apsp[np.ix_(support_i, support_j)]
            if not np.isfinite(ground).all():
                finite_max = apsp[np.isfinite(apsp)].max()
                ground = np.where(np.isfinite(ground), ground, 2.0 * finite_max)
            wasserstein = ot.emd2(mass_i, mass_j, ground)
            geodesic = max(float(apsp[i, j]), 1e-9)
            curvature = 1.0 - wasserstein / geodesic
            updated[(i, j)] = max(1e-9, weights[(i, j)] * (1.0 - args.flow_step * curvature))
        scale = total / sum(updated.values())
        weights = {edge: w * scale for edge, w in updated.items()}
        apsp = graph_apsp(len(distance), weights)
    return apsp


def geometry_channels(
    apsp: np.ndarray, heat_times: tuple[float, ...]
) -> torch.Tensor:
    """Standardized relation channels of a flowed geometry.

    Channel 0 is the negative geodesic field; the rest are heat kernels of
    the normalized Laplacian of a geodesic-scale affinity.
    """
    finite = apsp[np.isfinite(apsp) & (apsp > 0)]
    scale = np.median(finite)
    capped = np.where(np.isfinite(apsp), apsp, finite.max() * 2.0)
    affinity = np.exp(-capped / scale)
    np.fill_diagonal(affinity, 1.0)
    degree = affinity.sum(-1)
    normalized = affinity / np.sqrt(degree[:, None] * degree[None, :])
    values, vectors = np.linalg.eigh((normalized + normalized.T) / 2.0)
    laplacian_eigen = 1.0 - values
    channels = [torch.from_numpy(-capped).double()]
    for t in heat_times:
        heat = (vectors * np.exp(-t * laplacian_eigen)) @ vectors.T
        channels.append(torch.from_numpy(heat).double())
    return torch.stack(
        [standardize_relation(channel)[0] for channel in channels]
    )


def run_gates(
    text_channels: torch.Tensor,
    visual_channels: torch.Tensor,
    args: argparse.Namespace,
    generator: torch.Generator,
) -> dict:
    text_relation = text_channels[0]
    visual_relation = visual_channels[0]
    report = {
        "gate_a": gate_a_global_ranking(
            text_channels, visual_channels, text_relation, visual_relation, args, generator
        ),
        "gate_b": gate_b_transpositions(
            text_channels, visual_channels, text_relation, visual_relation, None
        ),
        "descent_from_true": steepest_descent(
            text_channels,
            visual_channels,
            text_relation,
            visual_relation,
            torch.arange(len(visual_relation)),
            args,
        ),
    }
    report["descent_from_true"].pop("final_permutation", None)
    report["descent_from_random"] = []
    for _ in range(args.descent_restarts):
        start = torch.argsort(torch.rand(len(visual_relation), generator=generator))
        result = steepest_descent(
            text_channels, visual_channels, text_relation, visual_relation, start, args
        )
        result.pop("final_permutation", None)
        result.pop("trajectory", None)
        report["descent_from_random"].append(result)
    return report


def main() -> None:
    args = parse_args()
    seed_everything(args.seed)
    data = load_flickr(args) if args.dataset == "flickr" else load_vg(args)
    heat_times = tuple(float(t) for t in args.heat_times.split(","))
    states = {
        "visual": torch.nn.functional.normalize(
            data["visual_views"].double().mean(1), dim=-1
        ),
        "text": torch.nn.functional.normalize(
            data["text_views"].double().mean(1), dim=-1
        ),
    }
    distances = {
        side: (1.0 - state @ state.T).clamp_min(0.0).numpy()
        for side, state in states.items()
    }
    report: dict = {
        "protocol": (
            "Each modality's kNN geometry evolves independently under "
            "Ollivier-Ricci flow; the identical heat-kernel channels are "
            "gated with and without the flow. Hidden pairs score orderings "
            "only."
        ),
        "meta": {**data["meta"], "flow": vars(args)},
        "conditions": {},
    }
    for label, iterations in (
        ("diffusion_no_flow", 0),
        ("ricci_flow", args.flow_iterations),
    ):
        channels = {}
        for side in ("visual", "text"):
            apsp = ollivier_ricci_apsp(distances[side], args, iterations)
            channels[side] = geometry_channels(apsp, heat_times)
        generator = torch.Generator().manual_seed(args.seed)
        result = run_gates(channels["text"], channels["visual"], args, generator)
        report["conditions"][label] = result
        print(
            json.dumps(
                {
                    label: {
                        "true_z_mse": result["gate_a"]["random"]["mse"]["true_z"],
                        "improving_fraction": result["gate_b"]["improving_fraction"],
                        "descent_keeps": result["descent_from_true"]["final_accuracy"],
                        "counterfeit_found": any(
                            r["final_objective"]
                            < result["gate_a"]["true"]["mse"]
                            and r["final_accuracy"] < 0.5
                            for r in result["descent_from_random"]
                        ),
                    }
                }
            )
        )
    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
    write_json(args.output, report)
    print(f"Wrote {args.output}")


if __name__ == "__main__":
    main()