path: root/worldalign/field_geometry.py

"""The space and the operators, which we never varied.

Everything measured so far fixed one geometry without saying so. Scene states
are compared by Euclidean inner product, the relation field holds those raw
values, and two fields are compared by Pearson correlation -- three
commitments to flatness made by default rather than by argument. The
measured deficit is that the two modalities disagree about which scenes are
similar, and "similar" is exactly what those three choices define.

The mismatch that costs the most is likely the cheapest to remove. Two
independently trained encoders have no reason to put their similarities on
the same scale: if vision similarity is any monotone but nonlinear function
of text similarity, the two fields describe identical relational structure
and Pearson correlation still reports disagreement. A rank transform is
invariant to every monotone reparametrisation and costs one sort.

Density is the second. A scene sitting in a crowded region of one modality
and a sparse region of the other has systematically shifted similarities to
everything, which local scaling and the diffusion kernel both remove -- the
latter by replacing similarity with how probability spreads, a quantity
defined by the graph rather than by whatever units the encoder happened to
emit.

Each transform is applied to each modality independently, so nothing here
uses the pairing or crosses the gauge.
"""

from __future__ import annotations

import argparse
import json

import numpy as np
import torch
from scipy.optimize import linear_sum_assignment
from scipy.stats import rankdata

from .common import write_json


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--fields", default="artifacts/vg_5k/np_n256.pt")
    parser.add_argument("--label", default="natural-best")
    parser.add_argument("--repeats", type=int, default=5)
    parser.add_argument("--output", default="artifacts/vg_5k/field_geometry.json")
    return parser.parse_args()


def offdiagonal_mask(size: int) -> np.ndarray:
    return ~np.eye(size, dtype=bool)


def standardise(matrix: np.ndarray) -> np.ndarray:
    mask = offdiagonal_mask(len(matrix))
    values = matrix[mask]
    out = (matrix - values.mean()) / values.std()
    np.fill_diagonal(out, 0.0)
    return out


# --- geometries, each applied within one modality -------------------------

def identity(field: np.ndarray) -> np.ndarray:
    return standardise(field)


def global_rank(field: np.ndarray) -> np.ndarray:
    """Copula transform: keep the ordering of similarities, discard the scale."""
    size = len(field)
    mask = offdiagonal_mask(size)
    out = np.zeros_like(field)
    out[mask] = rankdata(field[mask]) / mask.sum()
    return standardise(out)


def row_rank(field: np.ndarray) -> np.ndarray:
    """Rank within each scene's own neighbourhood, removing per-scene scale."""
    size = len(field)
    out = np.zeros_like(field)
    for row in range(size):
        others = np.delete(np.arange(size), row)
        out[row, others] = rankdata(field[row, others]) / len(others)
    out = (out + out.T) / 2.0
    return standardise(out)


def local_scaling(field: np.ndarray, neighbours: int = 10) -> np.ndarray:
    """Self-tuning: divide by each scene's own neighbourhood scale (CSLS-like)."""
    size = len(field)
    work = field.copy()
    np.fill_diagonal(work, -np.inf)
    order = np.sort(work, axis=1)[:, ::-1]
    scale = order[:, :neighbours].mean(1)
    out = field - 0.5 * (scale[:, None] + scale[None, :])
    np.fill_diagonal(out, 0.0)
    return standardise(out)


def diffusion(field: np.ndarray, steps: int = 3) -> np.ndarray:
    """Replace similarity by how probability spreads on the similarity graph."""
    size = len(field)
    affinity = field - field.min()
    np.fill_diagonal(affinity, 0.0)
    degree = affinity.sum(1).clip(1e-9)
    walk = affinity / degree[:, None]
    powered = np.linalg.matrix_power(walk, steps)
    out = powered @ powered.T
    np.fill_diagonal(out, 0.0)
    return standardise(out)


def heat_kernel(field: np.ndarray, time: float = 1.0) -> np.ndarray:
    """Heat flow on the normalised Laplacian: a metric the graph defines itself."""
    size = len(field)
    affinity = field - field.min()
    np.fill_diagonal(affinity, 0.0)
    degree = affinity.sum(1).clip(1e-9)
    normalised = affinity / np.sqrt(np.outer(degree, degree))
    values, vectors = np.linalg.eigh((normalised + normalised.T) / 2.0)
    out = (vectors * np.exp(time * (values - values.max()))) @ vectors.T
    np.fill_diagonal(out, 0.0)
    return standardise(out)


def hyperbolic(field: np.ndarray, dimension: int = 16) -> np.ndarray:
    """Embed each field in hyperbolic space and re-derive distances there.

    If scene similarity is hierarchical -- broad categories containing finer
    ones -- a flat inner product cannot hold it without distortion, while
    negative curvature can. The field is embedded by eigenmap and read back
    as Lorentzian distance.
    """
    size = len(field)
    values, vectors = np.linalg.eigh((field + field.T) / 2.0)
    order = np.argsort(values)[::-1][:dimension]
    coordinates = vectors[:, order] * np.sqrt(np.abs(values[order]))
    scale = np.abs(coordinates).max().clip(1e-9)
    spatial = coordinates / scale * 0.9
    time = np.sqrt(1.0 + (spatial ** 2).sum(1))
    product = np.outer(time, time) - spatial @ spatial.T
    out = -np.arccosh(np.clip(product, 1.0, None))
    np.fill_diagonal(out, 0.0)
    return standardise(out)


GEOMETRIES = {
    "euclidean (current)": identity,
    "global rank (copula)": global_rank,
    "row rank": row_rank,
    "local scaling": local_scaling,
    "diffusion t=3": diffusion,
    "heat kernel": heat_kernel,
    "hyperbolic d=16": hyperbolic,
}


# --- comparison operators -------------------------------------------------

def anchor_bound(first: np.ndarray, second: np.ndarray, repeats: int,
                 spearman: bool = False) -> float:
    size = len(first)
    scores = []
    for repeat in range(repeats):
        generator = np.random.default_rng(repeat)
        order = generator.permutation(size)
        half = size // 2
        anchors, probe = order[:half], order[half:]

        def profile(field):
            block = field[np.ix_(probe, anchors)]
            if spearman:
                block = np.apply_along_axis(rankdata, 1, block)
            block = block - block.mean(1, keepdims=True)
            return block / block.std(1, keepdims=True).clip(1e-9)

        similarity = profile(first) @ profile(second).T / half
        _, columns = linear_sum_assignment(-similarity)
        scores.append(float((columns == np.arange(len(probe))).mean()))
    return float(np.mean(scores))


def main() -> None:
    args = parse_args()
    state = torch.load(args.fields, map_location="cpu", weights_only=False)
    visual = state["visual_field"].double().numpy()
    text = state["text_field"].double().numpy()
    mask = offdiagonal_mask(len(visual))

    rows = []
    print("%-24s %8s %10s %12s" % ("geometry", "rho", "bound", "bound+rank-op"))
    for name, transform in GEOMETRIES.items():
        try:
            first, second = transform(visual), transform(text)
        except Exception as error:  # a geometry that cannot be formed is a result
            print("%-24s  failed: %s" % (name, error))
            continue
        rho = float(np.corrcoef(first[mask], second[mask])[0, 1])
        plain = anchor_bound(first, second, args.repeats)
        ranked = anchor_bound(first, second, args.repeats, spearman=True)
        rows.append({"geometry": name, "correlation": rho,
                     "anchor_bound": plain, "anchor_bound_rank_operator": ranked})
        print("%-24s %8.3f %10.4f %12.4f" % (name, rho, plain, ranked))

    best = max(rows, key=lambda row: max(row["anchor_bound"],
                                         row["anchor_bound_rank_operator"]))
    write_json(args.output, {
        "protocol": (
            "Each geometry is applied to each modality independently, so no "
            "transform crosses the gauge or uses the pairing. The bound is the "
            "anchor upper bound; the last column re-computes it with a rank "
            "comparison operator instead of a linear one."
        ),
        "label": args.label,
        "rows": rows,
        "best": best,
    })
    print(json.dumps(best))


if __name__ == "__main__":
    main()