scripts/parameter_sweep.py

"""Parameter sweep script for KdV PINN validation runs.

This script performs overnight parameter sweeps over different numbers of solitons,
kappas, and initial positions to generate a comprehensive dataset of validation runs.
"""
import os
import numpy as np
import itertools
from kdv_pinn.validation import validate_run

def run_parameter_sweep(base_dir='sweep_results'):
    """Run parameter sweep over multiple soliton configurations.

    Sweeps over:
    - Number of solitons: n = 1, 2, 3, 5, 7
    - Different kappa values for each n
    - Different x0 positions for each n

    Args:
        base_dir: Base directory for all sweep results
    """
    os.makedirs(base_dir, exist_ok=True)

    # Explicit sweep configurations with interesting collision dynamics
    # Focus on: fast chasers, multiple crossings, cascade collisions, speed variety
    sweep_configs = {
        1: {
            'kappas': [[1.0], [1.5], [2.0], [2.5]],
            'x0s': [[0], [5], [-5], [10]]
        },
        2: {
            'kappas': [
                # Fast chaser scenarios - dramatic catch-ups
                [2.5, 0.8],  # Very fast catching very slow
                [2.0, 1.0],  # Classic fast chaser
                [3.0, 1.5],  # Super fast catching medium
                # Similar speeds - minimal interaction
                [1.5, 1.4],  # Almost matched speeds, should barely interact
                [2.0, 1.9],  # Very close speeds
                # Reverse - slow in front (no collision expected)
                [0.9, 2.0],  # Slow leads, fast behind - shouldn't collide
                # Widely separated
                [3.5, 0.6],  # Extreme speed difference
                [2.2, 1.1],  # 2:1 speed ratio
            ],
            'x0s': [
                [-10, 5],    # Fast far behind, has to chase
                [-8, 0],     # Fast behind slow
                [-15, 3],    # Fast very far behind
                [0, 0.5],    # Nearly overlapping, similar speeds
                [0, 1],      # Slightly separated, similar speeds
                [5, -10],    # Slow in front, fast way behind (no collision)
                [-12, 2],    # Extreme chase scenario
                [-10, 0],    # Fast chaser from behind
            ]
        },
        3: {
            'kappas': [
                # Cascade: fast catches middle, then catches slow
                [3.0, 1.5, 0.8],   # Clear cascade progression
                [2.5, 1.3, 0.7],   # Another cascade
                # One fast, two slow and close
                [3.0, 1.0, 0.9],   # Fast guy catches tight pair
                # Fast guy in middle (by speed, not position)
                [2.5, 0.9, 1.6],   # Slow in middle by speed
                [2.0, 1.0, 1.8],   # Middle is slowest
                # Fast guy way in back catches everyone
                [3.5, 1.2, 1.0],   # One super fast chaser
                # Cluster of similar speeds
                [1.5, 1.4, 1.45],  # Very tight cluster - minimal dynamics
                # Wide spread
                [3.0, 1.5, 0.6],   # Evenly spread speeds
            ],
            'x0s': [
                [-15, -2, 5],      # Fast way back, catches middle then slow
                [-12, -3, 4],      # Cascade scenario
                [-15, 2, 3],       # Fast catches tight pair
                [-10, 3, -4],      # Slow in middle position
                [-8, 5, -6],       # Similar mixed ordering
                [-20, 2, 0],       # Fast very far behind everyone
                [0, 0.5, 1],       # Tight cluster in space
                [-15, -5, 5],      # Evenly spaced
            ]
        },
        5: {
            'kappas': [
                # One super fast guy chasing pack
                [4.0, 1.2, 1.0, 0.8, 0.6],   # Single fast chaser
                [3.5, 1.4, 1.1, 0.9, 0.7],   # Fast chaser with graduated pack
                # Two fast chasers
                [3.0, 2.5, 0.9, 0.8, 0.7],   # Two fast, three slow clustered
                [2.8, 2.2, 1.0, 0.9, 0.8],   # Two fast, tight slow pack
                # Mixed speeds (not monotonic)
                [2.5, 1.0, 2.0, 0.8, 1.5],   # Interleaved speeds
                [3.0, 0.7, 2.0, 1.0, 1.5],   # More speed mixing
            ],
            'x0s': [
                [-25, 3, 5, 7, 9],           # Fast way behind pack
                [-20, 2, 4, 6, 8],           # Fast chaser
                [-15, -12, 4, 6, 8],         # Two fast chasers
                [-12, -10, 3, 5, 7],         # Two fast back, tight pack ahead
                [-10, 5, -8, 8, -2],         # Interleaved positions
                [-15, 10, -5, 3, 0],         # Mixed positions
            ]
        },
        7: {
            'kappas': [
                # Multiple fast chasers
                [4.0, 3.0, 1.0, 0.9, 0.8, 0.7, 0.6],  # Two fast, pack of slow
                [3.5, 2.5, 2.0, 0.9, 0.8, 0.75, 0.7], # Three fast, pack of slow
                # Interleaved
                [3.0, 1.0, 2.5, 0.9, 2.0, 0.8, 1.5],  # Completely mixed speeds
            ],
            'x0s': [
                [-30, -20, 5, 7, 9, 11, 13],          # Two fast way back
                [-25, -15, -10, 6, 8, 10, 12],        # Three fast back
                [-15, 8, -12, 10, -8, 12, -2],        # Interleaved positions
            ]
        }
    }

    # Track all runs
    total_runs = sum(len(cfg['kappas']) for cfg in sweep_configs.values())
    run_count = 0
    successful_runs = []
    failed_runs = []

    print(f"Starting parameter sweep with {total_runs} total configurations")
    print("=" * 80)

    # Run sweep for each soliton number
    for n_solitons in [1, 2, 3, 5, 7]:
        config = sweep_configs[n_solitons]
        kappa_list = config['kappas']
        x0_list = config['x0s']

        # Run all combinations for this soliton number
        for idx, (kappas, x0s) in enumerate(zip(kappa_list, x0_list)):
            run_count += 1

            # Create descriptive name for this run
            kappa_str = '_'.join([f'{k:.2f}' for k in kappas])
            x0_str = '_'.join([f'{x:.1f}' for x in x0s])
            run_name = f'n{n_solitons}_run{idx:02d}_k{kappa_str}_x{x0_str}'
            output_dir = os.path.join(base_dir, run_name)

            print(f"\n[{run_count}/{total_runs}] Running: {run_name}")
            print(f"  Kappas: {kappas}")
            print(f"  x0s: {x0s}")

            try:
                # Import here to set config per run
                from kdv_pinn.configuration import kdv_config

                # Configure this run
                kdv_config.kappas = kappas
                kdv_config.x0s = x0s
                kdv_config.num_epochs = 5000
                kdv_config.num_pretrain_epochs = 2500
                kdv_config.num_samp_bulk = 96
                kdv_config.num_samp_eval = 128
                kdv_config.plot_interval = 50
                kdv_config.MLP = [2, 128, 128, 128, 1]
                kdv_config.lr = 1e-3
                kdv_config.T = 1
                kdv_config.L = 10
                kdv_config.Lmax = 15
                kdv_config.lambda_BC = 1
                kdv_config.lambda_kdv = 1
                kdv_config.vmin = 0
                kdv_config.vmax = 3

                # Mark config as pre-configured so validate_run doesn't override
                kdv_config._configured = True

                # Run validation
                metrics = validate_run(output_dir=output_dir)

                # Clear the flag for next run
                delattr(kdv_config, '_configured')

                successful_runs.append({
                    'name': run_name,
                    'n_solitons': n_solitons,
                    'kappas': kappas,
                    'x0s': x0s,
                    'metrics': metrics
                })
                print(f"  ✓ SUCCESS")

            except Exception as e:
                print(f"  ✗ FAILED: {str(e)}")
                failed_runs.append({
                    'name': run_name,
                    'n_solitons': n_solitons,
                    'kappas': kappas,
                    'x0s': x0s,
                    'error': str(e)
                })

    # Write summary
    print("\n" + "=" * 80)
    print("SWEEP COMPLETE")
    print("=" * 80)
    print(f"Successful runs: {len(successful_runs)}/{total_runs}")
    print(f"Failed runs: {len(failed_runs)}/{total_runs}")

    summary_file = os.path.join(base_dir, 'sweep_summary.txt')
    with open(summary_file, 'w') as f:
        f.write("PARAMETER SWEEP SUMMARY\n")
        f.write("=" * 80 + "\n\n")
        f.write(f"Total runs: {total_runs}\n")
        f.write(f"Successful: {len(successful_runs)}\n")
        f.write(f"Failed: {len(failed_runs)}\n\n")

        f.write("SUCCESSFUL RUNS:\n")
        f.write("-" * 80 + "\n")
        for run in successful_runs:
            f.write(f"\n{run['name']}\n")
            f.write(f"  N solitons: {run['n_solitons']}\n")
            f.write(f"  Kappas: {run['kappas']}\n")
            f.write(f"  x0s: {run['x0s']}\n")

            # Write key metrics
            if run['metrics']:
                for key in ['u_mae', 'u_rmse']:
                    if key in run['metrics']:
                        f.write(f"  {key}: {run['metrics'][key]:.6e}\n")

        if failed_runs:
            f.write("\n\nFAILED RUNS:\n")
            f.write("-" * 80 + "\n")
            for run in failed_runs:
                f.write(f"\n{run['name']}\n")
                f.write(f"  N solitons: {run['n_solitons']}\n")
                f.write(f"  Kappas: {run['kappas']}\n")
                f.write(f"  x0s: {run['x0s']}\n")
                f.write(f"  Error: {run['error']}\n")

    print(f"\nSummary written to: {summary_file}")
    print(f"All results saved to: {base_dir}/")

    return successful_runs, failed_runs


if __name__ == "__main__":
    import torch

    # Set random seed for reproducibility
    torch.manual_seed(42)
    np.random.seed(42)

    print("KdV PINN Parameter Sweep")
    print("This will run overnight. Results will be saved to sweep_results/")
    print("\nStarting sweep...\n")

    successful, failed = run_parameter_sweep()

    print("\nSweep complete! Check sweep_results/ for all plots and metrics.")