evidence/code/armadillo.py

"""Paper-scale Armadillo checks for Claims 1, 3, and 4."""

from __future__ import annotations

import gzip
import hashlib
import json
import math
import os
import csv
import time
import urllib.request
from pathlib import Path

import numpy as np
from scipy.sparse.linalg import LinearOperator, eigsh
from scipy.spatial.distance import cdist

from resolution import ARTIFACT_DIR as CLAIM_2_ARTIFACT_DIR
from resolution import ROOT


ARTIFACT_ROOT = ROOT / ".openresearch" / "artifacts"
USER_AGENT = "OpenResearch-Reproduction-Audit/1.0 (contact: research-local)"


def _download(url: str, expected_sha256: str) -> tuple[bytes, dict]:
    request = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
    with urllib.request.urlopen(request, timeout=120) as response:
        compressed = response.read()
    observed = hashlib.sha256(compressed).hexdigest()
    if observed != expected_sha256:
        raise RuntimeError(
            f"Armadillo hash mismatch: expected {expected_sha256}, got {observed}"
        )
    return compressed, {
        "url": url,
        "retrieval_user_agent": USER_AGENT,
        "sha256": observed,
        "compressed_bytes": len(compressed),
    }


def _parse_armadillo(compressed: bytes) -> tuple[np.ndarray, np.ndarray, dict]:
    raw = gzip.decompress(compressed)
    marker = b"end_header\n"
    header_end = raw.index(marker) + len(marker)
    header = raw[:header_end].decode("ascii")
    if "format binary_big_endian 1.0" not in header:
        raise RuntimeError("unexpected PLY format")

    vertex_count = int(
        next(
            line.split()[2]
            for line in header.splitlines()
            if line.startswith("element vertex ")
        )
    )
    face_count = int(
        next(
            line.split()[2]
            for line in header.splitlines()
            if line.startswith("element face ")
        )
    )
    payload = memoryview(raw)[header_end:]
    vertex_bytes = 12 * vertex_count
    vertices = (
        np.frombuffer(payload[:vertex_bytes], dtype=">f4")
        .astype(np.float64)
        .reshape(vertex_count, 3)
    )
    face_dtype = np.dtype(
        [("intensity", "u1"), ("count", "u1"), ("indices", ">i4", (3,))]
    )
    faces_raw = np.frombuffer(
        payload[vertex_bytes:], dtype=face_dtype, count=face_count
    )
    if faces_raw.size != face_count or not np.all(faces_raw["count"] == 3):
        raise RuntimeError("expected an all-triangle Armadillo mesh")
    faces = faces_raw["indices"].astype(np.int64)
    if faces.min() < 0 or faces.max() >= vertex_count:
        raise RuntimeError("invalid face index")
    return vertices, faces, {
        "ply_format": "binary_big_endian_1.0",
        "vertex_count": vertex_count,
        "face_count": face_count,
        "all_triangles": True,
    }


def _normalize_to_unit_ball(vertices: np.ndarray) -> tuple[np.ndarray, dict]:
    center = 0.5 * (vertices.min(axis=0) + vertices.max(axis=0))
    centered = vertices - center
    radius = float(np.linalg.norm(centered, axis=1).max())
    normalized = centered / radius
    return normalized, {
        "normalization": "bounding-box center then divide by maximum radius",
        "center": center.tolist(),
        "radius_before_scaling": radius,
        "max_radius_after_scaling": float(
            np.linalg.norm(normalized, axis=1).max()
        ),
    }


def _sample_surface(
    vertices: np.ndarray, faces: np.ndarray, count: int, seed: int
) -> tuple[np.ndarray, dict]:
    triangles = vertices[faces]
    cross = np.cross(
        triangles[:, 1] - triangles[:, 0],
        triangles[:, 2] - triangles[:, 0],
    )
    areas = 0.5 * np.linalg.norm(cross, axis=1)
    positive = areas > 0.0
    if not np.all(positive):
        areas = np.where(positive, areas, 0.0)
    probabilities = areas / areas.sum()
    rng = np.random.default_rng(seed)
    selected = rng.choice(faces.shape[0], size=count, p=probabilities)
    chosen = triangles[selected]
    u = rng.random(count)
    v = rng.random(count)
    sqrt_u = np.sqrt(u)
    barycentric = np.column_stack(
        [1.0 - sqrt_u, sqrt_u * (1.0 - v), sqrt_u * v]
    )
    points = np.einsum("ni,nij->nj", barycentric, chosen)
    return points, {
        "sampling": "triangle-area-weighted iid barycentric surface sampling",
        "seed": seed,
        "count": count,
        "positive_area_faces": int(np.count_nonzero(positive)),
        "surface_area_normalized_units": float(areas.sum()),
    }


def _kernel(points: np.ndarray, name: str, sigma: float) -> tuple[np.ndarray, np.ndarray]:
    squared = cdist(points, points, metric="sqeuclidean")
    if name == "gaussian":
        log_kernel = -squared / (2.0 * sigma**2)
    elif name == "exponential":
        log_kernel = -np.sqrt(squared) / sigma
    else:
        raise ValueError(name)
    return np.exp(log_kernel), log_kernel


def _sinkhorn(
    kernel: np.ndarray, weights: np.ndarray, tolerance: float = 1e-12
) -> tuple[np.ndarray, list[float], int, float]:
    scaling = np.ones(kernel.shape[0], dtype=np.float64)
    curve: list[float] = []
    threshold_iteration = -1
    residual_max = math.inf
    for iteration in range(1, 301):
        row_values = scaling * (kernel @ (weights * scaling))
        mean_error = float(np.sum(weights * np.abs(row_values - 1.0)))
        curve.append(mean_error)
        if threshold_iteration < 0 and mean_error < 1e-3:
            threshold_iteration = iteration
        residual_max = float(np.max(np.abs(row_values - 1.0)))
        if residual_max < tolerance:
            return scaling, curve, threshold_iteration, residual_max
        scaling = np.sqrt(
            scaling / np.maximum(kernel @ (weights * scaling), 1e-300)
        )
    raise RuntimeError(f"Sinkhorn did not converge: residual={residual_max}")


def _self_adjoint_relative_error(
    kernel: np.ndarray,
    weights: np.ndarray,
    left_scale: np.ndarray,
    right_scale: np.ndarray,
) -> float:
    maximum = 0.0
    scale = 0.0
    block_size = 500
    for start in range(0, kernel.shape[0], block_size):
        stop = min(start + block_size, kernel.shape[0])
        block = (
            weights[start:stop, None]
            * left_scale[start:stop, None]
            * kernel[start:stop]
            * right_scale[None, :]
            * weights[None, :]
        )
        reverse = (
            weights[start:stop, None]
            * right_scale[start:stop, None]
            * kernel[start:stop]
            * left_scale[None, :]
            * weights[None, :]
        )
        maximum = max(maximum, float(np.max(np.abs(block - reverse))))
        scale = max(scale, float(np.max(np.abs(block))))
    return maximum / max(scale, 1e-300)


def _top_spectrum(
    kernel: np.ndarray, weights: np.ndarray, scaling: np.ndarray
) -> list[float]:
    diagonal = np.sqrt(weights) * scaling

    def matvec(vector: np.ndarray) -> np.ndarray:
        return diagonal * (kernel @ (diagonal * vector))

    operator = LinearOperator(
        kernel.shape, matvec=matvec, rmatvec=matvec, dtype=np.float64
    )
    values = eigsh(
        operator,
        k=3,
        which="LA",
        return_eigenvectors=False,
        tol=2e-10,
        maxiter=500,
    )
    return np.sort(values).tolist()


def _landmark_minimum(
    kernel: np.ndarray,
    weights: np.ndarray,
    scaling: np.ndarray,
    count: int,
) -> float:
    indices = np.linspace(
        0, kernel.shape[0] - 1, count, dtype=np.int64
    )
    diagonal = np.sqrt(weights[indices]) * scaling[indices]
    subset = diagonal[:, None] * kernel[np.ix_(indices, indices)] * diagonal[None, :]
    return float(np.linalg.eigvalsh(subset)[0])


def _normalization_metrics(
    kernel: np.ndarray,
    weights: np.ndarray,
    scaling: np.ndarray,
    kind: str,
) -> dict:
    degree = kernel @ weights
    if kind == "sinkhorn":
        left = scaling
        right = scaling
        mass = scaling * (kernel @ (weights * scaling))
    elif kind == "row":
        left = 1.0 / degree
        right = np.ones_like(degree)
        mass = left * (kernel @ weights)
    elif kind == "symmetric":
        left = 1.0 / np.sqrt(degree)
        right = left
        mass = left * (kernel @ (weights * right))
    else:
        raise ValueError(kind)
    return {
        "normalization": kind,
        "mass_max_error": float(np.max(np.abs(mass - 1.0))),
        "self_adjoint_relative_error": _self_adjoint_relative_error(
            kernel, weights, left, right
        ),
    }


def run_armadillo(config: dict) -> dict:
    start = time.perf_counter()
    spec = config["armadillo"]
    compressed, source = _download(spec["url"], spec["sha256"])
    vertices, faces, mesh = _parse_armadillo(compressed)
    vertices, normalization = _normalize_to_unit_ball(vertices)
    points, sampling = _sample_surface(
        vertices, faces, int(spec["surface_sample_count"]), int(spec["seed"])
    )
    weights = np.full(points.shape[0], 1.0 / points.shape[0])
    kernel_records = []

    for kernel_name in spec["kernels"]:
        kernel, log_kernel = _kernel(points, kernel_name, float(spec["sigma"]))
        scaling, curve, threshold_iteration, final_residual = _sinkhorn(
            kernel, weights
        )
        sinkhorn_metrics = _normalization_metrics(
            kernel, weights, scaling, "sinkhorn"
        )
        row_metrics = _normalization_metrics(
            kernel, weights, scaling, "row"
        )
        symmetric_metrics = _normalization_metrics(
            kernel, weights, scaling, "symmetric"
        )
        top_eigenvalues = _top_spectrum(kernel, weights, scaling)
        landmark_minimum = _landmark_minimum(
            kernel,
            weights,
            scaling,
            int(spec["landmark_spectrum_count"]),
        )
        min_log_entry = float(
            log_kernel.min()
            + 2.0 * np.log(scaling).min()
            + np.log(weights).min()
        )
        kernel_records.append(
            {
                "kernel": kernel_name,
                "sigma": float(spec["sigma"]),
                "n": points.shape[0],
                "iterations_to_mean_error_below_1e-3": threshold_iteration,
                "mean_error_curve_first_12": curve[:12],
                "final_max_mass_residual": final_residual,
                "scaling_min": float(scaling.min()),
                "scaling_max": float(scaling.max()),
                "minimum_log_operator_entry_lower_bound": min_log_entry,
                "float64_zero_kernel_entries": int(
                    np.count_nonzero(kernel == 0.0)
                ),
                "top_eigenvalues": top_eigenvalues,
                "landmark_spectrum_count": int(
                    spec["landmark_spectrum_count"]
                ),
                "landmark_minimum_eigenvalue": landmark_minimum,
                "normalizations": [
                    sinkhorn_metrics,
                    row_metrics,
                    symmetric_metrics,
                ],
            }
        )

    return {
        "schema_version": 1,
        "claims": [1, 3, 4],
        "data_source": source,
        "mesh": mesh,
        "normalization": normalization,
        "sampling": sampling,
        "bandwidth": float(spec["sigma"]),
        "kernels": kernel_records,
        "runtime_seconds": time.perf_counter() - start,
    }


def write_armadillo_artifacts(result: dict) -> Path:
    for claim_id in (1, 3, 4):
        directory = ARTIFACT_ROOT / f"claim_{claim_id}"
        directory.mkdir(parents=True, exist_ok=True)
        (directory / "raw_results.json").write_text(
            json.dumps(result, indent=2) + "\n", encoding="utf-8"
        )
        with (directory / "raw_results.csv").open(
            "w", newline="", encoding="utf-8"
        ) as handle:
            writer = csv.DictWriter(
                handle,
                fieldnames=[
                    "kernel",
                    "n",
                    "sigma",
                    "iterations_to_1e-3",
                    "final_mass_residual",
                    "top_eigenvalue",
                    "landmark_minimum_eigenvalue",
                    "sinkhorn_self_adjoint_error",
                    "row_self_adjoint_error",
                    "symmetric_mass_error",
                ],
            )
            writer.writeheader()
            for record in result["kernels"]:
                normalizations = {
                    item["normalization"]: item
                    for item in record["normalizations"]
                }
                writer.writerow(
                    {
                        "kernel": record["kernel"],
                        "n": record["n"],
                        "sigma": record["sigma"],
                        "iterations_to_1e-3": record[
                            "iterations_to_mean_error_below_1e-3"
                        ],
                        "final_mass_residual": record[
                            "final_max_mass_residual"
                        ],
                        "top_eigenvalue": record["top_eigenvalues"][-1],
                        "landmark_minimum_eigenvalue": record[
                            "landmark_minimum_eigenvalue"
                        ],
                        "sinkhorn_self_adjoint_error": normalizations[
                            "sinkhorn"
                        ]["self_adjoint_relative_error"],
                        "row_self_adjoint_error": normalizations["row"][
                            "self_adjoint_relative_error"
                        ],
                        "symmetric_mass_error": normalizations["symmetric"][
                            "mass_max_error"
                        ],
                    }
                )
    return ARTIFACT_ROOT / "claim_1" / "raw_results.json"


def print_armadillo_summary(result: dict) -> None:
    print(
        "ARMADILLO_SOURCE="
        + json.dumps(result["data_source"], sort_keys=True)
    )
    print("ARMADILLO_MESH=" + json.dumps(result["mesh"], sort_keys=True))
    print(
        "ARMADILLO_SAMPLING="
        + json.dumps(result["sampling"], sort_keys=True)
    )
    for record in result["kernels"]:
        print("ARMADILLO_KERNEL_RESULT=" + json.dumps(record, sort_keys=True))
    print(f"ARMADILLO_RUNTIME_SECONDS={result['runtime_seconds']:.6f}")