auto_math_kernel_density_and_coherence_1

Snippet from Math-Kernel-Density-and-Coherence.md

import numpy as np
from scipy.stats import gaussian_kde
import matplotlib.pyplot as plt

class CoherenceEstimator:
    """
    Density-based coherence estimation in whitened space.
    
    Estimates local density to assess query alignment with trace manifold.
    """
    
    def __init__(self, bandwidth: float = 0.1, threshold_coherent: float = -1.0,
                 threshold_incoherent: float = -3.0):
        """
        Initialize coherence estimator.
        
        Args:
            bandwidth: Smoothing parameter h (uses Silverman if None).
            threshold_coherent: Above this value indicates alignment.
            threshold_incoherent: Below this value indicates misalignment.
        """
        self.h = bandwidth
        self.threshold_high = threshold_coherent
        self.threshold_low = threshold_incoherent
        self.kde = None
        self.M_inv_sqrt = None
    
    def _whiten(self, X: np.ndarray, M: np.ndarray) -> np.ndarray:
        """
        Apply isotropic transformation via metric inverse square root.
        
        Args:
            X: Points (n x d).
            M: SPD matrix (d x d).
        
        Returns:
            Z: Whitened coordinates (n x d).
        """
        eigenvalues, U = np.linalg.eigh(M)
        s_inv = np.where(eigenvalues > 1e-10, 1/np.sqrt(eigenvalues), 0)
        self.M_inv_sqrt = U @ np.diag(s_inv) @ U.T
        return X @ self.M_inv_sqrt.T
    
    def fit(self, traces: np.ndarray, M: np.ndarray) -> 'CoherenceEstimator':
        """
        Fit KDE on whitened traces.
        
        Args:
            traces: Training points (n x d).
            M: Metric matrix (d x d).
        
        Returns:
            Self for chaining.
        """
        Z = self._whiten(traces, M)
        self.kde = gaussian_kde(Z.T, bw_method=self.h)
        return self
    
    def score(self, query: np.ndarray) -> tuple[float, str]:
        """
        Compute log density and coherence label.
        
        Args:
            query: Query point (d,) or (1 x d).
        
        Returns:
            Tuple of (log probability, label).
        """
        if self.kde is None:
            raise ValueError("Must call fit() before scoring.")
        
        z_query = query.reshape(1, -1) @ self.M_inv_sqrt.T
        log_p = np.log(self.kde(z_query.T)[0] + 1e-10)
        
        if log_p > self.threshold_high:
            label = "high"
        elif log_p > self.threshold_low:
            label = "moderate"
        else:
            label = "low"
        
        return log_p, label
    
    def visualize_density(self, traces: np.ndarray, query: np.ndarray,
                          M: np.ndarray, grid_size: int = 50):
        """
        Generate 2D density heatmap (requires d=2).
        
        Args:
            traces: Training points (n x 2).
            query: Query point (2-vector).
            M: Metric matrix (2 x 2).
            grid_size: Grid resolution for visualization.
        """
        if traces.shape[1] != 2:
            raise ValueError("Visualization requires 2D data.")
        
        Z_traces = self._whiten(traces, M)
        z_query = self._whiten(query.reshape(1, -1), M)
        
        x_min, x_max = Z_traces[:, 0].min() - 1, Z_traces[:, 0].max() + 1
        y_min, y_max = Z_traces[:, 1].min() - 1, Z_traces[:, 1].max() + 1
        xx, yy = np.meshgrid(np.linspace(x_min, x_max, grid_size),
                             np.linspace(y_min, y_max, grid_size))
        
        grid_points = np.c_[xx.ravel(), yy.ravel()]
        densities = np.exp(self.kde.logpdf(grid_points.T))
        densities = densities.reshape(xx.shape)
        
        plt.figure(figsize=(10, 8))
        plt.contourf(xx, yy, densities, levels=20, cmap='viridis', alpha=0.8)
        plt.colorbar(label='Density p-hat(z)')
        plt.scatter(Z_traces[:, 0], Z_traces[:, 1], c='white',
                    s=30, edgecolors='black', label='Traces')
        plt.scatter(z_query[:, 0], z_query[:, 1], c='red',
                    s=200, marker='*', edgecolors='black', label='Query')
        
        log_p_query = self.kde.logpdf(z_query.T)[0]
        plt.contour(xx, yy, np.log(densities + 1e-10),
                    levels=[self.threshold_low, self.threshold_high, log_p_query],
                    colors=['red', 'orange', 'yellow'], linewidths=2)
        
        plt.xlabel('Z_1')
        plt.ylabel('Z_2')
        plt.title(f'Density Map (h={self.h:.3f})')
        plt.legend()
        plt.grid(alpha=0.3)
        plt.show()

# Example Usage
np.random.seed(42)
traces = np.random.multivariate_normal([0, 0], [[1, 0.3], [0.3, 1]], 100)
M = np.eye(2)
q_coherent = np.array([0.1, 0.1])
q_incoherent = np.array([3.5, 3.5])

estimator = CoherenceEstimator(bandwidth=0.2)
estimator.fit(traces, M)

log_p_coh, label_coh = estimator.score(q_coherent)
log_p_inc, label_inc = estimator.score(q_incoherent)

print("Coherence Analysis")
print("=" * 50)
print(f"Coherent query: log p-hat = {log_p_coh:.3f} ({label_coh})")
print(f"Incoherent query: log p-hat = {log_p_inc:.3f} ({label_inc})")
print("\nInterpretation:")
print(f"  > -1.0: High (on-manifold)")
print(f"  -1 to -3: Moderate (edge)")
print(f"  < -3.0: Low (off-manifold)")

# Uncomment to visualize:
# estimator.visualize_density(traces, q_coherent, M)