evidence/code/spectra.py

"""Cross-modality Armadillo spectral reproduction for Claims 5 and 6."""

from __future__ import annotations

import json
import time
import csv
from pathlib import Path

import numpy as np
from scipy.sparse import coo_matrix, diags
from scipy.sparse.linalg import LinearOperator, eigsh
from scipy.spatial import cKDTree
from scipy.spatial.distance import cdist
from scipy.special import logsumexp

from armadillo import (
    ARTIFACT_ROOT,
    _download,
    _kernel,
    _normalize_to_unit_ball,
    _parse_armadillo,
    _sample_surface,
    _sinkhorn,
)


def _diffusion_eigendecomposition(
    kernel: np.ndarray,
    weights: np.ndarray,
    count: int,
) -> tuple[np.ndarray, np.ndarray, int]:
    scaling, _, threshold_iteration, _ = _sinkhorn(kernel, weights)
    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, vectors = eigsh(
        operator,
        k=count,
        which="LA",
        tol=2e-9,
        maxiter=1_000,
    )
    order = np.argsort(values)[::-1]
    functions = vectors[:, order] / np.sqrt(weights)[:, None]
    return values[order], functions, threshold_iteration


def _diffusion_eigenvalues(
    kernel: np.ndarray,
    weights: np.ndarray,
    count: int,
) -> tuple[np.ndarray, int]:
    values, _, threshold_iteration = _diffusion_eigendecomposition(
        kernel, weights, count
    )
    return values, threshold_iteration


def _kmeans_plus_plus(
    points: np.ndarray, components: int, seed: int
) -> tuple[np.ndarray, np.ndarray]:
    rng = np.random.default_rng(seed)
    centers = np.empty((components, points.shape[1]), dtype=np.float64)
    centers[0] = points[rng.integers(points.shape[0])]
    minimum = np.sum((points - centers[0]) ** 2, axis=1)
    for index in range(1, components):
        probabilities = minimum / minimum.sum()
        centers[index] = points[rng.choice(points.shape[0], p=probabilities)]
        candidate = np.sum((points - centers[index]) ** 2, axis=1)
        minimum = np.minimum(minimum, candidate)
    labels = np.zeros(points.shape[0], dtype=np.int64)
    for _ in range(6):
        distances = cdist(points, centers, metric="sqeuclidean")
        labels = np.argmin(distances, axis=1)
        for index in range(components):
            selected = labels == index
            if np.any(selected):
                centers[index] = points[selected].mean(axis=0)
            else:
                centers[index] = points[np.argmax(distances.min(axis=1))]
    return centers, labels


def _fit_gmm(
    points: np.ndarray,
    components: int,
    iterations: int,
    seed: int,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, list[float]]:
    means, labels = _kmeans_plus_plus(points, components, seed)
    dimension = points.shape[1]
    covariances = np.empty((components, dimension, dimension))
    weights = np.empty(components)
    global_covariance = np.cov(points.T) + 1e-5 * np.eye(dimension)
    for index in range(components):
        selected_points = points[labels == index]
        weights[index] = max(selected_points.shape[0], 1)
        if selected_points.shape[0] >= 4:
            covariances[index] = (
                np.cov(selected_points.T) + 1e-5 * np.eye(dimension)
            )
        else:
            covariances[index] = 0.05 * global_covariance
    weights /= weights.sum()
    lower_bounds: list[float] = []

    for _ in range(iterations):
        log_probabilities = np.empty((points.shape[0], components))
        for index in range(components):
            inverse = np.linalg.inv(covariances[index])
            sign, logdet = np.linalg.slogdet(covariances[index])
            if sign <= 0:
                raise RuntimeError("non-positive GMM covariance")
            difference = points - means[index]
            quadratic = np.einsum(
                "ni,ij,nj->n", difference, inverse, difference
            )
            log_probabilities[:, index] = (
                np.log(max(weights[index], 1e-300))
                - 0.5
                * (
                    quadratic
                    + logdet
                    + dimension * np.log(2.0 * np.pi)
                )
            )
        normalizer = logsumexp(log_probabilities, axis=1)
        lower_bounds.append(float(normalizer.mean()))
        if (
            len(lower_bounds) >= 3
            and abs(lower_bounds[-1] - lower_bounds[-2]) < 1e-3
        ):
            break
        responsibilities = np.exp(
            log_probabilities - normalizer[:, None]
        )
        effective = responsibilities.sum(axis=0) + 1e-12
        weights = effective / effective.sum()
        means = (responsibilities.T @ points) / effective[:, None]
        for index in range(components):
            difference = points - means[index]
            covariances[index] = (
                np.einsum(
                    "n,ni,nj->ij",
                    responsibilities[:, index],
                    difference,
                    difference,
                )
                / effective[index]
                + 1e-6 * np.eye(dimension)
            )
    return weights, means, covariances, lower_bounds


def _gmm_kernel(
    means: np.ndarray, covariances: np.ndarray, sigma: float
) -> np.ndarray:
    count, dimension = means.shape
    kernel = np.empty((count, count), dtype=np.float64)
    isotropic = sigma**2 * np.eye(dimension)
    for index in range(count):
        combined = isotropic[None, :, :] + covariances[index] + covariances
        inverses = np.linalg.inv(combined)
        differences = means[index] - means
        quadratic = np.einsum(
            "ni,nij,nj->n", differences, inverses, differences
        )
        kernel[index] = np.exp(-0.5 * quadratic)
    return 0.5 * (kernel + kernel.T)


def _surface_voxels(
    vertices: np.ndarray,
    faces: np.ndarray,
    edge: float,
    sample_count: int,
    seed: int,
) -> tuple[np.ndarray, np.ndarray, dict]:
    samples, _ = _sample_surface(
        vertices, faces, sample_count, seed + 101
    )
    indices = np.floor((samples + 1.0) / edge).astype(np.int64)
    indices = np.unique(indices, axis=0)
    centers = -1.0 + edge * (indices.astype(np.float64) + 0.5)
    squared = cdist(centers, centers, metric="sqeuclidean")
    density = np.exp(-squared / (2.0 * (3.0 * edge) ** 2)).sum(axis=1)
    weights = 1.0 / density
    weights /= weights.sum()
    return centers, weights, {
        "edge": edge,
        "nonempty_voxels": centers.shape[0],
        "rasterization_surface_samples": sample_count,
        "mass_formula": "Eq.46 inverse Gaussian KDE with std=3 voxels",
    }


def _clustered_cotan_spectrum(
    vertices: np.ndarray,
    faces: np.ndarray,
    cluster_edge: float,
    count: int,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, dict]:
    keys = np.floor((vertices + 1.0) / cluster_edge).astype(np.int64)
    _, inverse = np.unique(keys, axis=0, return_inverse=True)
    cluster_count = int(inverse.max()) + 1
    clustered = np.zeros((cluster_count, 3), dtype=np.float64)
    cluster_sizes = np.bincount(inverse, minlength=cluster_count)
    np.add.at(clustered, inverse, vertices)
    clustered /= cluster_sizes[:, None]
    mapped_faces = inverse[faces]
    nondegenerate = (
        (mapped_faces[:, 0] != mapped_faces[:, 1])
        & (mapped_faces[:, 1] != mapped_faces[:, 2])
        & (mapped_faces[:, 2] != mapped_faces[:, 0])
    )
    mapped_faces = mapped_faces[nondegenerate]
    canonical = np.sort(mapped_faces, axis=1)
    _, unique_indices = np.unique(canonical, axis=0, return_index=True)
    mapped_faces = mapped_faces[np.sort(unique_indices)]

    triangles = clustered[mapped_faces]
    twice_area = np.linalg.norm(
        np.cross(
            triangles[:, 1] - triangles[:, 0],
            triangles[:, 2] - triangles[:, 0],
        ),
        axis=1,
    )
    valid = twice_area > 1e-14
    triangles = triangles[valid]
    mapped_faces = mapped_faces[valid]
    twice_area = twice_area[valid]
    mass = np.zeros(cluster_count, dtype=np.float64)
    for local in range(3):
        np.add.at(mass, mapped_faces[:, local], twice_area / 6.0)

    edge_rows = []
    edge_cols = []
    edge_values = []
    for first, second, opposite in ((0, 1, 2), (1, 2, 0), (2, 0, 1)):
        u = triangles[:, first] - triangles[:, opposite]
        v = triangles[:, second] - triangles[:, opposite]
        cotangent = np.einsum("ni,ni->n", u, v) / twice_area
        weight = 0.5 * cotangent
        left = mapped_faces[:, first]
        right = mapped_faces[:, second]
        edge_rows.extend([left, right])
        edge_cols.extend([right, left])
        edge_values.extend([-weight, -weight])
    rows = np.concatenate(edge_rows)
    cols = np.concatenate(edge_cols)
    values = np.concatenate(edge_values)
    stiffness = coo_matrix(
        (values, (rows, cols)), shape=(cluster_count, cluster_count)
    ).tocsr()
    stiffness = stiffness + diags(-np.asarray(stiffness.sum(axis=1)).ravel())
    active = mass > 0.0
    stiffness = stiffness[active][:, active]
    mass = mass[active]
    eigenvalues, eigenvectors = eigsh(
        stiffness,
        k=count,
        M=diags(mass),
        sigma=1e-8,
        which="LM",
        tol=2e-7,
        maxiter=2_000,
    )
    order = np.argsort(eigenvalues)
    eigenvalues = np.maximum(eigenvalues[order], 0.0)
    eigenvectors = eigenvectors[:, order]
    return eigenvalues, eigenvectors, clustered[active], {
        "method": "cotangent Laplacian after deterministic voxel vertex clustering",
        "cluster_edge": cluster_edge,
        "original_vertices": vertices.shape[0],
        "clustered_vertices": cluster_count,
        "active_vertices": int(np.count_nonzero(active)),
        "clustered_faces": mapped_faces.shape[0],
    }


def _laplacian_estimate(
    diffusion_values: np.ndarray, denominator: float
) -> np.ndarray:
    clipped = np.clip(diffusion_values, 1e-300, 1.0)
    return -2.0 * np.log(clipped) / denominator


def _comparison(reference: np.ndarray, candidate: np.ndarray) -> dict:
    indices = np.arange(1, min(reference.size, candidate.size))
    low = indices[indices <= 14]
    scale = float(
        np.dot(reference[low], candidate[low])
        / max(np.dot(candidate[low], candidate[low]), 1e-300)
    )
    aligned = scale * candidate
    relative = np.abs(aligned - reference) / np.maximum(reference, 1e-12)
    correlation = float(np.corrcoef(reference[low], candidate[low])[0, 1])
    divergence_candidates = indices[
        (indices >= 10) & (relative[indices] > 0.25)
    ]
    divergence_index = (
        int(divergence_candidates[0] + 1)
        if divergence_candidates.size
        else None
    )
    baseline = relative[low]
    baseline_median = float(np.median(baseline))
    baseline_mad = float(
        np.median(np.abs(baseline - baseline_median))
    )
    sustained_threshold = max(
        0.10, baseline_median + 2.0 * baseline_mad
    )
    sustained_divergence = next(
        (
            start + 1
            for start in range(15, relative.size - 4)
            if np.count_nonzero(
                relative[start : start + 5] > sustained_threshold
            )
            >= 3
        ),
        None,
    )
    return {
        "least_squares_scale_indices_2_to_15": scale,
        "pearson_indices_2_to_15": correlation,
        "median_relative_error_indices_2_to_15": float(
            np.median(relative[low])
        ),
        "first_index_after_10_relative_error_above_25pct": divergence_index,
        "low_mode_error_median": baseline_median,
        "low_mode_error_mad": baseline_mad,
        "sustained_divergence_threshold": sustained_threshold,
        "sustained_divergence_rule": (
            "first mode >=16 with at least 3 of 5 aligned errors above "
            "max(0.10, low-mode median + 2*MAD)"
        ),
        "sustained_divergence_index": sustained_divergence,
        "aligned_relative_errors": relative.tolist(),
    }


def run_spectral_analysis(config: dict) -> dict:
    start = time.perf_counter()
    armadillo_spec = config["armadillo"]
    spectral_spec = config["spectral_analysis"]
    count = int(spectral_spec["eigenvalue_count"])
    sigma = float(armadillo_spec["sigma"])
    compressed, source = _download(
        armadillo_spec["url"], armadillo_spec["sha256"]
    )
    vertices, faces, mesh = _parse_armadillo(compressed)
    vertices, normalization = _normalize_to_unit_ball(vertices)
    points, sampling = _sample_surface(
        vertices,
        faces,
        int(armadillo_spec["surface_sample_count"]),
        int(spectral_spec["seed"]),
    )

    point_weights = np.full(points.shape[0], 1.0 / points.shape[0])
    point_kernel, _ = _kernel(points, "gaussian", sigma)
    point_diffusion, point_functions, point_iterations = (
        _diffusion_eigendecomposition(
            point_kernel, point_weights, count
        )
    )
    point_laplacian = _laplacian_estimate(point_diffusion, sigma**2)

    gmm_weights, gmm_means, gmm_covariances, lower_bounds = _fit_gmm(
        points,
        int(spectral_spec["gmm_components"]),
        int(spectral_spec["gmm_em_iterations"]),
        int(spectral_spec["seed"]),
    )
    gmm_kernel = _gmm_kernel(gmm_means, gmm_covariances, sigma)
    gmm_diffusion, gmm_functions, gmm_iterations = (
        _diffusion_eigendecomposition(
            gmm_kernel, gmm_weights, count
        )
    )
    average_trace = float(
        np.sum(gmm_weights * np.trace(gmm_covariances, axis1=1, axis2=2))
    )
    gmm_denominator = sigma**2 + average_trace
    gmm_laplacian = _laplacian_estimate(
        gmm_diffusion, gmm_denominator
    )

    voxel_points, voxel_weights, voxel_metadata = _surface_voxels(
        vertices,
        faces,
        float(spectral_spec["voxel_edge"]),
        int(spectral_spec["voxel_surface_samples"]),
        int(spectral_spec["seed"]),
    )
    voxel_kernel, _ = _kernel(voxel_points, "gaussian", sigma)
    voxel_diffusion, voxel_functions, voxel_iterations = (
        _diffusion_eigendecomposition(
            voxel_kernel, voxel_weights, count
        )
    )
    voxel_laplacian = _laplacian_estimate(voxel_diffusion, sigma**2)

    (
        cotan_laplacian,
        cotan_functions,
        cotan_points,
        cotan_metadata,
    ) = _clustered_cotan_spectrum(
        vertices,
        faces,
        float(spectral_spec["cotan_cluster_edge"]),
        count,
    )
    from volume_spectra import (
        _canonical_correlations,
        run_volume_spectral_analysis,
    )

    cotan_tree = cKDTree(cotan_points)
    surface_reference_points = cotan_functions[
        cotan_tree.query(points, workers=1)[1]
    ]
    surface_reference_gmm = cotan_functions[
        cotan_tree.query(gmm_means, workers=1)[1]
    ]
    surface_reference_voxels = cotan_functions[
        cotan_tree.query(voxel_points, workers=1)[1]
    ]
    surface_eigenspaces = {
        "point_5000": {
            "modes_2_to_10": _canonical_correlations(
                surface_reference_points,
                point_functions,
                point_weights,
                1,
                10,
            ),
            "modes_8_to_12": _canonical_correlations(
                surface_reference_points,
                point_functions,
                point_weights,
                7,
                12,
            ),
        },
        "gmm_500": {
            "modes_2_to_10": _canonical_correlations(
                surface_reference_gmm,
                gmm_functions,
                gmm_weights,
                1,
                10,
            ),
            "modes_8_to_12": _canonical_correlations(
                surface_reference_gmm,
                gmm_functions,
                gmm_weights,
                7,
                12,
            ),
        },
        "surface_voxels": {
            "modes_2_to_10": _canonical_correlations(
                surface_reference_voxels,
                voxel_functions,
                voxel_weights,
                1,
                10,
            ),
            "modes_8_to_12": _canonical_correlations(
                surface_reference_voxels,
                voxel_functions,
                voxel_weights,
                7,
                12,
            ),
        },
    }
    modalities = {
        "point_5000": {
            "count": points.shape[0],
            "sinkhorn_iterations_to_1e-3": point_iterations,
            "diffusion_eigenvalues": point_diffusion.tolist(),
            "estimated_laplacian_eigenvalues": point_laplacian.tolist(),
            "conversion": "Eq.47",
            "reference_interpolation": "nearest clustered cotan vertex",
            "eigenspaces": surface_eigenspaces["point_5000"],
        },
        "gmm_500": {
            "count": gmm_means.shape[0],
            "kernel": "Eq.6 covariance-aware Gaussian overlap",
            "covariance_matrices": int(gmm_covariances.shape[0]),
            "covariance_dimension": int(gmm_covariances.shape[1]),
            "minimum_covariance_eigenvalue": float(
                np.linalg.eigvalsh(gmm_covariances).min()
            ),
            "sinkhorn_iterations_to_1e-3": gmm_iterations,
            "diffusion_eigenvalues": gmm_diffusion.tolist(),
            "estimated_laplacian_eigenvalues": gmm_laplacian.tolist(),
            "conversion": "Eq.48 with d=2",
            "average_covariance_trace": average_trace,
            "effective_denominator": gmm_denominator,
            "em_lower_bounds": lower_bounds,
            "reference_interpolation": "nearest clustered cotan vertex",
            "eigenspaces": surface_eigenspaces["gmm_500"],
        },
        "surface_voxels": {
            **voxel_metadata,
            "sinkhorn_iterations_to_1e-3": voxel_iterations,
            "diffusion_eigenvalues": voxel_diffusion.tolist(),
            "estimated_laplacian_eigenvalues": voxel_laplacian.tolist(),
            "conversion": "Eq.47",
            "reference_interpolation": "nearest clustered cotan vertex",
            "eigenspaces": surface_eigenspaces["surface_voxels"],
        },
        "cotan_reference": {
            **cotan_metadata,
            "laplacian_eigenvalues": cotan_laplacian.tolist(),
        },
    }
    comparisons = {
        "point_5000": _comparison(cotan_laplacian, point_laplacian),
        "gmm_500": _comparison(cotan_laplacian, gmm_laplacian),
        "surface_voxels": _comparison(cotan_laplacian, voxel_laplacian),
    }
    volume = None
    if spectral_spec.get("volume_enabled", False):
        volume = run_volume_spectral_analysis(
            vertices,
            faces,
            spectral_spec,
            sigma,
            _fit_gmm,
            _gmm_kernel,
            _laplacian_estimate,
            _comparison,
            _sinkhorn,
        )
        modalities.update(volume["modalities"])
    return {
        "schema_version": 2,
        "claims": [5, 6],
        "data_source": source,
        "mesh": mesh,
        "normalization": normalization,
        "sampling": sampling,
        "sigma": sigma,
        "eigenvalue_count": count,
        "modalities": modalities,
        "comparisons_to_cotan": comparisons,
        "comparisons_to_fem": (
            {} if volume is None else volume["comparisons_to_fem"]
        ),
        "volume_seed_sweep": (
            [] if volume is None else volume["seed_sweep"]
        ),
        "tetrahedralization": (
            None if volume is None else volume["tetrahedralization"]
        ),
        "divergence_index_window": spectral_spec.get(
            "divergence_index_window"
        ),
        "sampling_resolution_boundary": (
            None
            if volume is None
            else volume["sampling_resolution_boundary"]
        ),
        "volume_runtime_seconds": (
            0.0 if volume is None else volume["runtime_seconds"]
        ),
        "runtime_seconds": time.perf_counter() - start,
    }


def write_spectral_artifacts(result: dict) -> Path:
    for claim_id in (5, 6):
        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=["modality", "index", "laplacian_eigenvalue"],
            )
            writer.writeheader()
            for modality, record in result["modalities"].items():
                values = record.get(
                    "estimated_laplacian_eigenvalues",
                    record.get("laplacian_eigenvalues"),
                )
                for index, value in enumerate(values, start=1):
                    writer.writerow(
                        {
                            "modality": modality,
                            "index": index,
                            "laplacian_eigenvalue": f"{value:.17g}",
                        }
                    )
    return ARTIFACT_ROOT / "claim_6" / "raw_results.json"


def print_spectral_summary(result: dict) -> None:
    print(
        "SPECTRAL_MODALITY_METADATA="
        + json.dumps(
            {
                key: {
                    field: value
                    for field, value in record.items()
                    if "eigenvalues" not in field
                    and field not in {"eigenspaces", "em_lower_bounds"}
                }
                for key, record in result["modalities"].items()
            },
            sort_keys=True,
        )
    )
    print(
        "SPECTRAL_COMPARISONS="
        + json.dumps(result["comparisons_to_cotan"], sort_keys=True)
    )
    print(
        "VOLUME_COMPARISONS_TO_FEM="
        + json.dumps(result["comparisons_to_fem"], sort_keys=True)
    )
    print(
        "EIGENSPACE_SUMMARY="
        + json.dumps(
            {
                modality: {
                    window: {
                        "median": diagnostics[
                            "median_canonical_correlation"
                        ],
                        "minimum": diagnostics[
                            "minimum_canonical_correlation"
                        ],
                    }
                    for window, diagnostics in record.get(
                        "eigenspaces", {}
                    ).items()
                }
                for modality, record in result["modalities"].items()
                if "eigenspaces" in record
            },
            sort_keys=True,
        )
    )
    print(
        "VOLUME_SEED_SWEEP="
        + json.dumps(
            [
                {
                    "seed": record["seed"],
                    "comparisons_to_fem": record["comparisons_to_fem"],
                    "eigenspace_summary": record["eigenspace_summary"],
                }
                for record in result["volume_seed_sweep"]
            ],
            sort_keys=True,
        )
    )
    for key, record in result["modalities"].items():
        eigenvalues = record.get(
            "estimated_laplacian_eigenvalues",
            record.get("laplacian_eigenvalues"),
        )
        print(f"SPECTRUM_{key}=" + json.dumps(eigenvalues))
    print(f"SPECTRAL_RUNTIME_SECONDS={result['runtime_seconds']:.6f}")