evaluate.py

#!/usr/bin/env python3
"""
Evaluation script for all localization methods.

This script is to evaluate different localization methods (DDPM, GCN, MLP, MLE, UNet).

Usage Examples:
    python evaluate.py --method ddpm
    python evaluate.py --method gcn --epochs 1000
    python evaluate.py --method mle --noise-level 0.2
    python evaluate.py --method ddpm --prob-loc --num-samples 100
"""

import os
import sys
import argparse
import numpy as np
import torch
import torch.optim as optim

# Add parent directory to path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from ddpm_loc.methods import DDPMLoc
from ddpm_loc.methods.baselines import GCNLoc, MLPLoc, MLELoc, UNetLoc
from ddpm_loc.models import DenoiseNet, GCN, MLP, UNet
from ddpm_loc.utils.data_utils import generate_agent_data, generate_anchor_positions_by_layout, generate_noise, split_train_test_data
from ddpm_loc.utils.visualization import plot_scenario_setup, plot_prob_loc_progress


def parse_arguments():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(description='Evaluate localization methods')
    
    # Method selection
    parser.add_argument('--method', type=str, required=True, 
                        choices=['ddpm', 'gcn', 'mlp', 'mle', 'unet'], 
                        help='Localization method to evaluate')
    
    # Data parameters
    parser.add_argument('--num-samples', type=int, default=2000,
                        help='Number of agent samples (default: 2000)')
    parser.add_argument('--num-anchors', type=int, default=50,
                        help='Number of anchor points (default: 50)')
    parser.add_argument('--noise-level', type=float, default=0.1,
                        help='Measurement noise level \\sigma (default: 0.1)')
    parser.add_argument('--agent-noise-level', type=float, default=0.0,
                        help='Agent position generation noise level (default: 0.0)')
    parser.add_argument('--noise-type', type=str, default='rayleigh',
                        choices=['normal', 'rayleigh', 'mixture'],
                        help='Type of noise (default: rayleigh)')
    parser.add_argument('--range-limit', type=float, default=1.0,
                        help='Range limit for data generation (default: 1.0)')
    parser.add_argument('--layout', type=str, default='spiral',
                        choices=['spiral', 'grid', 'random', 'circle', 'cluster'],
                        help='Agent data layout type (default: spiral)')
    parser.add_argument('--anchor-layout', type=str, default='random',
                        choices=['random', 'grid', 'circle', 'uniform', 'boundary'],
                        help='Anchor layout type (default: random)')
    parser.add_argument('--test-size', type=float, default=0.5,
                        help='Test agents set fraction (default: 0.5)')
    parser.add_argument('--random-seed', type=int, default=42,
                        help='Random seed for reproducibility (default: 42)')
    
    # Training parameters
    parser.add_argument('--epochs', type=int, default=2000,
                        help='Number of training epochs (default: 2000)')
    parser.add_argument('--learning-rate', type=float, default=1e-3,
                        help='Learning rate (default: 1e-3)')
    
    # Method-specific parameters
    # DDPM parameters
    parser.add_argument('--ddpm-num-steps', type=int, default=1000,
                        help='DDPM number of diffusion steps (default: 1000)')
    parser.add_argument('--ddpm-beta-start', type=float, default=1e-5,
                        help='DDPM beta start value (default: 1e-5)')
    parser.add_argument('--ddpm-beta-end', type=float, default=1e-3,
                        help='DDPM beta end value (default: 1e-3)')
    
    # Run probabilistic localization parameters
    parser.add_argument('--prob-loc', action='store_true',
                        help='Run probabilistic localization (DDPM only)')
    parser.add_argument('--num-sampling-times', type=int, default=100,
                        help='Number of sampling times for probabilistic localization (default: 100)')
    
    # GCN parameters
    parser.add_argument('--gcn-hidden-dim', type=int, default=2000,
                        help='GCN hidden dimension (default: 2000)')
    parser.add_argument('--gcn-threshold', type=float, default=10.0,
                        help='GCN threshold (default: 10.0)')
    
    # MLP parameters
    parser.add_argument('--mlp-hidden-dims', type=int, nargs='+', 
                        default=[128, 64, 32, 32, 64, 128],
                        help='MLP hidden dimensions (default: [128, 64, 32, 32, 64, 128])')
    
    # UNet parameters
    parser.add_argument('--unet-embed-dim', type=int, default=128,
                        help='UNet embedding dimension (default: 128)')
    
    # Output parameters
    parser.add_argument('--plot', action='store_true',
                        help='Plot results')
    parser.add_argument('--save-plots', action='store_true',
                        help='Save plots to files')
    parser.add_argument('--plot-dir', type=str, default='plots',
                        help='Directory to save plots (default: plots)')
    
    return parser.parse_args()


def setup_directories(args):
    """Create necessary directories."""
    if args.save_plots:
        os.makedirs(args.plot_dir, exist_ok=True)
        # Create method-specific directory with agent and anchor layout
        method_dir = os.path.join(args.plot_dir, args.method, args.layout, args.anchor_layout)
        os.makedirs(method_dir, exist_ok=True)
        # Also create method directory for backward compatibility
        os.makedirs(os.path.join(args.plot_dir, args.method), exist_ok=True)


def generate_data(args):
    """Generate experimental data."""
    print(f"📊 Generating experimental data (layout: {args.layout})...")
    
    # Set random seeds
    np.random.seed(args.random_seed)
    torch.manual_seed(args.random_seed)
    
    # Generate agent and anchor locations
    agent_loc, _ = generate_agent_data(
        layout=args.layout,
        num_samples=args.num_samples,
        noise_level=args.agent_noise_level,
        range_limit=args.range_limit,
        seed=args.random_seed
    )
    anchor_loc = generate_anchor_positions_by_layout(
        layout=args.anchor_layout,
        num_anchors=args.num_anchors,
        range_limit=args.range_limit,
        seed=args.random_seed
    )
    
    print(f"🔊 Noise type: {args.noise_type}, Noise level: {args.noise_level}")
    
    # Generate noise
    noise = generate_noise(
        len(anchor_loc), len(agent_loc), args.noise_level, 
        args.noise_type,
        seed=args.random_seed
    )
    # Ensure noise is a numpy array
    noise = np.asarray(noise)
    
    # Split noise into anchor-anchor and agent-anchor components
    anchor_2_anchor_noise = noise[:len(anchor_loc), :]
    agent_2_anchor_noise = noise[len(anchor_loc):len(anchor_loc) + len(agent_loc), :]
    
    # Split data
    train_agent_loc, test_agent_loc, train_agent_2_anchor_noise, test_agent_2_anchor_noise = split_train_test_data(
        agent_loc, agent_2_anchor_noise, 
        test_size=args.test_size,
        random_state=args.random_seed
    )
    
    # Convert to numpy arrays with consistent dtype
    train_agent_loc = np.asarray(train_agent_loc)
    test_agent_loc = np.asarray(test_agent_loc)
    anchor_loc = np.asarray(anchor_loc)
    
    # Compute distances and measurements
    anchor_2_anchor_distance = np.linalg.norm(
        anchor_loc[:, np.newaxis] - anchor_loc[np.newaxis, :], axis=2
    )
    train_agent_2_anchor_distance = np.linalg.norm(
        train_agent_loc[:, np.newaxis] - anchor_loc[np.newaxis, :], axis=2
    )
    test_agent_2_anchor_distance = np.linalg.norm(
        test_agent_loc[:, np.newaxis] - anchor_loc[np.newaxis, :], axis=2
    )
    
    # Add noise to distances
    anchor_2_anchor_measurements = anchor_2_anchor_distance + anchor_2_anchor_noise
    train_measurements = np.vstack([
        anchor_2_anchor_measurements,
        train_agent_2_anchor_distance + train_agent_2_anchor_noise
    ])
    test_measurements = np.vstack([
        anchor_2_anchor_measurements,
        test_agent_2_anchor_distance + test_agent_2_anchor_noise
    ])
    
    # Define indices
    idx_anchor = list(range(len(anchor_loc)))
    idx_agent_in_train = list(range(len(anchor_loc), len(anchor_loc) + len(train_agent_loc)))
    idx_agent_in_test = list(range(len(anchor_loc), len(anchor_loc) + len(test_agent_loc)))
    
    return {
        'anchor_loc': anchor_loc,
        'train_agent_loc': train_agent_loc,
        'test_agent_loc': test_agent_loc,
        'train_measurements': train_measurements,
        'test_measurements': test_measurements,
        'idx_anchor': idx_anchor,
        'idx_agent_in_train': idx_agent_in_train,
        'idx_agent_in_test': idx_agent_in_test
    }


def evaluate_ddpm(args, data):
    """Evaluate DDPM method."""
    print("🤖 Initializing DDPM model...")
    model = DenoiseNet(cond_dim=len(data['anchor_loc']))
    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
    
    ddpm_method = DDPMLoc(
        model, optimizer,
        num_steps=args.ddpm_num_steps,
        beta_start=args.ddpm_beta_start,
        beta_end=args.ddpm_beta_end
    )
    
    # Train the model
    print("⌛️ Training DDPM model...")
    save_path = os.path.join(args.plot_dir, 'ddpm', args.layout, args.anchor_layout, 'training_curve.png') if args.save_plots else None
    losses, grad_norms = ddpm_method.train(
        data['anchor_loc'], data['train_measurements'], data['idx_anchor'],
        epochs=args.epochs, plot=args.plot, save_path=save_path
    )
    
    print("✅ Training completed")
    
    if args.prob_loc:
        # Probabilistic localization
        print("🧪 Testing DDPM model (probabilistic)...")
        all_predictions = []
        for i in range(args.num_sampling_times):
            pred_loc_np, _ = ddpm_method.test(
                data['anchor_loc'], data['test_agent_loc'], data['test_measurements'],
                data['idx_anchor'], data['idx_agent_in_test'], 
                sample_plot=False, test_plot=False
            )
            all_predictions.append(pred_loc_np)
        
        all_predictions = np.array(all_predictions)
        sample_means = np.mean(all_predictions[:, len(data['idx_anchor']):, :], axis=0)
        
        # Plot probabilistic results
        if args.plot:
            save_path = os.path.join(args.plot_dir, 'ddpm', args.layout, args.anchor_layout, 'prob_loc_results.png') if args.save_plots else None
            plot_prob_loc_progress(
                all_predictions, sample_means, data['anchor_loc'], data['test_agent_loc'],
                data['idx_anchor'], data['idx_agent_in_test'], save_path=save_path
            )
        
        rmse = np.sqrt(np.mean(np.sum((sample_means - data['test_agent_loc'])**2, axis=1)))
        return sample_means, rmse
    else:
        # Standard localization
        print("🧪 Testing DDPM model...")
        test_save_path = os.path.join(args.plot_dir, 'ddpm', args.layout, args.anchor_layout, 'test_results.png') if args.save_plots else None
        sample_save_path = os.path.join(args.plot_dir, 'ddpm', args.layout, args.anchor_layout, 'sampling') if args.save_plots else None
        
        pred_loc_np, rmse = ddpm_method.test(
            data['anchor_loc'], data['test_agent_loc'], data['test_measurements'],
            data['idx_anchor'], data['idx_agent_in_test'], 
            sample_plot=args.plot, sample_save_path=sample_save_path,
            test_plot=args.plot, test_save_path=test_save_path
        )
        
        return pred_loc_np, rmse


def evaluate_gcn(args, data):
    """Evaluate GCN method."""
    print("🤖 Initializing GCN model...")
    model = GCN(nfeat=len(data['anchor_loc']), nhid1=args.gcn_hidden_dim, nout=2)
    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
    
    gcn_method = GCNLoc(
        model, optimizer,
        threshold=args.gcn_threshold
    )
    
    # Train the model
    print("⌛️ Training GCN model...")
    save_path = os.path.join(args.plot_dir, 'gcn', args.layout, args.anchor_layout, 'training_curve.png') if args.save_plots else None
    losses, grad_norms = gcn_method.train(
        data['anchor_loc'], data['train_agent_loc'], data['train_measurements'], 
        data['idx_anchor'], data['idx_agent_in_train'],
        epochs=args.epochs, plot=args.plot, save_path=save_path
    )
    
    print("✅ Training completed")
    
    # Test the model
    print("🧪 Testing GCN model...")
    save_path = os.path.join(args.plot_dir, 'gcn', args.layout, args.anchor_layout, 'test_results.png') if args.save_plots else None
    pred_loc_np, rmse = gcn_method.test(
        data['anchor_loc'], data['test_agent_loc'], data['test_measurements'],
        data['idx_anchor'], data['idx_agent_in_test'], 
        plot=args.plot, save_path=save_path
    )
    
    return pred_loc_np, rmse


def evaluate_mlp(args, data):
    """Evaluate MLP method."""
    print("🤖 Initializing MLP model...")
    model = MLP(input_dim=len(data['anchor_loc']), hidden_dims=args.mlp_hidden_dims, output_dim=2)
    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
    
    mlp_method = MLPLoc(model, optimizer)
    
    # Train the model
    print("⌛️ Training MLP model...")
    save_path = os.path.join(args.plot_dir, 'mlp', args.layout, args.anchor_layout, 'training_curve.png') if args.save_plots else None
    losses, grad_norms = mlp_method.train(
        data['anchor_loc'], data['train_measurements'], data['idx_anchor'], 
        epochs=args.epochs, plot=args.plot, save_path=save_path
    )
    
    print("✅ Training completed")
    
    # Test the model
    print("🧪 Testing MLP model...")
    save_path = os.path.join(args.plot_dir, 'mlp', args.layout, args.anchor_layout, 'test_results.png') if args.save_plots else None
    pred_loc_np, rmse = mlp_method.test(
        data['anchor_loc'], data['test_agent_loc'], data['test_measurements'],
        data['idx_anchor'], data['idx_agent_in_test'], 
        plot=args.plot, save_path=save_path
    )
    
    return pred_loc_np, rmse


def evaluate_mle(args, data):
    """Evaluate MLE method."""
    print("🤖 Initializing MLE method...")
    mle_method = MLELoc()
    
    # Test the model (MLE doesn't require training)
    print("🧪 Testing MLE method...")
    save_path = os.path.join(args.plot_dir, 'mle', args.layout, args.anchor_layout, 'test_results.png') if args.save_plots else None
    
    # MLE uses only the test agent measurements
    test_agent_measurements = data['test_measurements'][len(data['anchor_loc']):len(data['anchor_loc']) + len(data['test_agent_loc']), :]
    
    pred_loc_np, rmse = mle_method.test(
        data['anchor_loc'], data['test_agent_loc'], test_agent_measurements,
        data['idx_anchor'], data['idx_agent_in_test'], 
        plot=args.plot, save_path=save_path
    )
    
    return pred_loc_np, rmse


def evaluate_unet(args, data):
    """Evaluate UNet method."""
    print("🤖 Initializing UNet model...")
    model = UNet(input_dim=len(data['anchor_loc']), embed_dim=args.unet_embed_dim, output_dim=2)
    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
    
    unet_method = UNetLoc(model, optimizer)
    
    # Train the model
    print("⌛️ Training UNet model...")
    save_path = os.path.join(args.plot_dir, 'unet', args.layout, args.anchor_layout, 'training_curve.png') if args.save_plots else None
    losses, grad_norms = unet_method.train(
        data['anchor_loc'], data['train_measurements'], data['idx_anchor'], 
        epochs=args.epochs, plot=args.plot, save_path=save_path
    )
    
    print("✅ Training completed")
    
    # Test the model
    print("🧪 Testing UNet model...")
    save_path = os.path.join(args.plot_dir, 'unet', args.layout, args.anchor_layout, 'test_results.png') if args.save_plots else None
    pred_loc_np, rmse = unet_method.test(
        data['anchor_loc'], data['test_agent_loc'], data['test_measurements'],
        data['idx_anchor'], data['idx_agent_in_test'], 
        plot=args.plot, save_path=save_path
    )
    
    return pred_loc_np, rmse


def main():
    """Main evaluation function."""
    args = parse_arguments()
    
    print(f"🔬 {args.method.upper()} Localization Evaluation")
    print("=" * 40)
    
    # Setup directories
    setup_directories(args)
    
    # Generate data
    data = generate_data(args)
    
    # Plot scenario setup
    if args.plot:
        save_path = os.path.join(args.plot_dir, f'agents_and_anchors_{args.layout}_{args.anchor_layout}.png') if args.save_plots else None
        plot_scenario_setup(data['anchor_loc'], data['test_agent_loc'], save_path=save_path)
    
    # Evaluate the selected method
    if args.method == 'ddpm':
        pred_loc_np, rmse = evaluate_ddpm(args, data)
    elif args.method == 'gcn':
        pred_loc_np, rmse = evaluate_gcn(args, data)
    elif args.method == 'mlp':
        pred_loc_np, rmse = evaluate_mlp(args, data)
    elif args.method == 'mle':
        pred_loc_np, rmse = evaluate_mle(args, data)
    elif args.method == 'unet':
        pred_loc_np, rmse = evaluate_unet(args, data)
    else:
        raise ValueError(f"Unknown method: {args.method}")
    
    # Print results
    print("=" * 40)
    print("📊 Results Summary:")
    print(f"Method: {args.method.upper()}")
    print(f"RMSE: {rmse:.4f}")
    if args.method == 'ddpm' and args.prob_loc:
        print("Mode: Probabilistic Localization")
    print("🎉 Evaluation completed successfully!")


if __name__ == "__main__":
    main()