"""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}")