plot_blochwalk_robust.py

#!/usr/bin/env python3
"""
Robust plotting script with data sanity checks.
"""

import numpy as np
import torch
import matplotlib.pyplot as plt
from pathlib import Path

def plot_bloch_disk(beliefs_true, beliefs_pred, save_path, rmse):
    """Plot Bloch disk belief geometry."""
    # Flatten to (N*T, 3)
    beliefs_true_flat = beliefs_true.reshape(-1, 3)
    beliefs_pred_flat = beliefs_pred.reshape(-1, 3)
    
    # Extract x and z components (y should be ~0)
    b_x_true = beliefs_true_flat[:, 0]
    b_z_true = beliefs_true_flat[:, 2]
    
    b_x_pred = beliefs_pred_flat[:, 0]
    b_z_pred = beliefs_pred_flat[:, 2]
    
    # Create figure with 3 subplots
    fig, axes = plt.subplots(1, 3, figsize=(18, 6))
    
    # Subsample for plotting
    n_plot = min(5000, len(b_x_true))
    idx = np.random.RandomState(42).choice(len(b_x_true), n_plot, replace=False)
    
    # Plot 1: Ground Truth
    ax = axes[0]
    scatter = ax.scatter(b_x_true[idx], b_z_true[idx], c=b_z_true[idx], 
                        cmap='viridis', s=1, alpha=0.5)
    
    # Draw unit circle boundary
    theta = np.linspace(0, 2*np.pi, 100)
    ax.plot(np.cos(theta), np.sin(theta), 'r-', linewidth=2, alpha=0.3, label='Bloch sphere boundary')
    
    ax.set_xlabel('$b_x$', fontsize=14)
    ax.set_ylabel('$b_z$', fontsize=14)
    ax.set_title('Ground Truth\n(Bloch Walk Quantum Process)', fontsize=14, fontweight='bold')
    ax.set_aspect('equal')
    ax.grid(True, alpha=0.3)
    ax.set_xlim(-1.1, 1.1)
    ax.set_ylim(-1.1, 1.1)
    ax.legend()
    plt.colorbar(scatter, ax=ax, label='$b_z$ value')
    
    # Plot 2: Transformer Learned
    ax = axes[1]
    scatter = ax.scatter(b_x_pred[idx], b_z_pred[idx], c=b_z_true[idx], 
                        cmap='viridis', s=1, alpha=0.5)
    ax.plot(np.cos(theta), np.sin(theta), 'r-', linewidth=2, alpha=0.3, label='Bloch sphere boundary')
    
    ax.set_xlabel('$b_x$', fontsize=14)
    ax.set_ylabel('$b_z$', fontsize=14)
    ax.set_title(f'Transformer Learned\n(RMSE={rmse:.4f})', fontsize=14, fontweight='bold')
    ax.set_aspect('equal')
    ax.grid(True, alpha=0.3)
    # Use adaptive limits based on predicted data
    x_lim = max(abs(b_x_pred).max(), 1.1) + 0.1
    z_lim = max(abs(b_z_pred).max(), 1.1) + 0.1
    ax.set_xlim(-x_lim, x_lim)
    ax.set_ylim(-z_lim, z_lim)
    ax.legend()
    plt.colorbar(scatter, ax=ax, label='Ground truth $b_z$')
    
    # Plot 3: Error visualization
    ax = axes[2]
    errors = np.sqrt((b_x_true[idx] - b_x_pred[idx])**2 + (b_z_true[idx] - b_z_pred[idx])**2)
    scatter = ax.scatter(b_x_true[idx], b_z_true[idx], c=errors, 
                        cmap='Reds', s=1, alpha=0.5, vmin=0, vmax=np.percentile(errors, 95))
    ax.plot(np.cos(theta), np.sin(theta), 'k-', linewidth=2, alpha=0.3)
    
    ax.set_xlabel('$b_x$ (ground truth)', fontsize=14)
    ax.set_ylabel('$b_z$ (ground truth)', fontsize=14)
    ax.set_title('Prediction Error\n(Euclidean Distance)', fontsize=14, fontweight='bold')
    ax.set_aspect('equal')
    ax.grid(True, alpha=0.3)
    ax.set_xlim(-1.1, 1.1)
    ax.set_ylim(-1.1, 1.1)
    plt.colorbar(scatter, ax=ax, label='Error')
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    print(f"✓ Saved Bloch disk plot to {save_path}")
    plt.close()


def main():
    """Load data and create plots."""
    print("=" * 70)
    print("Bloch Walk Visualization (Robust)")
    print("=" * 70)
    print()
    
    # Load analysis batch
    analysis_path = Path("runs/blochwalk_test/analysis_batch.pt")
    print(f"Loading {analysis_path}...")
    data = torch.load(analysis_path)
    residuals = data["residuals"].numpy()  # (N, T, D)
    beliefs = data["beliefs"].numpy()       # (N, T, 3)
    
    print(f"  Residuals shape: {residuals.shape}")
    print(f"  Beliefs shape: {beliefs.shape}")
    print()
    
    # Check for NaN/Inf
    print("Data sanity checks:")
    print(f"  Residuals NaN: {np.isnan(residuals).sum()}")
    print(f"  Residuals Inf: {np.isinf(residuals).sum()}")
    print(f"  Beliefs NaN: {np.isnan(beliefs).sum()}")
    print(f"  Beliefs Inf: {np.isinf(beliefs).sum()}")
    print()
    
    if np.isnan(residuals).any() or np.isinf(residuals).any():
        print("⚠ Warning: Residuals contain NaN/Inf. Filtering...")
        # Find valid samples
        valid_mask = ~(np.isnan(residuals).any(axis=(1, 2)) | np.isinf(residuals).any(axis=(1, 2)))
        residuals = residuals[valid_mask]
        beliefs = beliefs[valid_mask]
        print(f"  Kept {residuals.shape[0]} valid samples")
        print()
    
    # Flatten
    N, T, D = residuals.shape
    X = residuals.reshape(-1, D)  # (N*T, D)
    Y = beliefs.reshape(-1, 3)     # (N*T, 3)
    
    print(f"Flattened: X={X.shape}, Y={Y.shape}")
    print(f"  X range: [{X.min():.4f}, {X.max():.4f}]")
    print(f"  Y range: [{Y.min():.4f}, {Y.max():.4f}]")
    print()
    
    # Normalize X for numerical stability
    X_mean = X.mean(axis=0, keepdims=True)
    X_std = X.std(axis=0, keepdims=True) + 1e-8
    X_norm = (X - X_mean) / X_std
    
    # Simple linear regression with normal equations
    print("Running linear regression (normal equations)...")
    # Add bias
    X_aug = np.column_stack([np.ones(len(X_norm)), X_norm])
    
    # Solve: (X^T X) W = X^T Y
    # Add regularization for stability
    XtX = X_aug.T @ X_aug + 0.01 * np.eye(X_aug.shape[1])
    XtY = X_aug.T @ Y
    
    try:
        W = np.linalg.solve(XtX, XtY)
        Y_pred = X_aug @ W
        
        # Compute RMSE
        rmse = np.sqrt(((Y - Y_pred)**2).mean())
        print(f"  RMSE: {rmse:.6f}")
        print()
        
        # Reshape back
        beliefs_pred = Y_pred.reshape(N, T, 3)
        
        # Plot
        print("Generating plots...")
        plot_path = Path("runs/blochwalk_test/bloch_disk_comparison.png")
        plot_bloch_disk(beliefs, beliefs_pred, plot_path, rmse)
        print()
        
        # Additional stats
        print("Statistics:")
        print(f"  Ground truth b_x range: [{Y[:, 0].min():.4f}, {Y[:, 0].max():.4f}]")
        print(f"  Ground truth b_z range: [{Y[:, 2].min():.4f}, {Y[:, 2].max():.4f}]")
        print(f"  Predicted b_x range: [{Y_pred[:, 0].min():.4f}, {Y_pred[:, 0].max():.4f}]")
        print(f"  Predicted b_z range: [{Y_pred[:, 2].min():.4f}, {Y_pred[:, 2].max():.4f}]")
        print()
        
        print("=" * 70)
        print("✅ COMPLETE!")
        print("=" * 70)
        
    except np.linalg.LinAlgError as e:
        print(f"❌ Linear algebra error: {e}")
        print("The regression failed. This might be due to ill-conditioned data.")


if __name__ == "__main__":
    main()