"""Main entry point for KdV PINN training and validation.""" import argparse import os import torch from types import SimpleNamespace from kdv_pinn.models import KdV_pinn from kdv_pinn.train import train_pinn, pretrain from kdv_pinn.scattering import ScatteringData, SchrodingerSolver from kdv_pinn.configuration import kdv_config, config_to_dict, load_config from kdv_pinn.device import get_device from kdv_pinn.validation import ( validate_run, validate_from_checkpoint, validate_analytical_only ) def main(): parser = argparse.ArgumentParser( description='Train and validate PINN for KdV equation', formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: # Run full validation with animation python run_training.py --mode validate --config config.json --animate # Generate outputs from existing model (no re-training) python run_training.py --mode validate --load-model output/pinn_model.pt # Run analytical-only mode (no training) python run_training.py --mode analytical --config config.json --animate # Simple training run python run_training.py --mode train --epochs 5000 """ ) # Mode selection parser.add_argument('--mode', type=str, default='train', choices=['train', 'validate', 'analytical'], help='Operation mode: train (simple training), validate (full validation), analytical (analytical solution only)') # Configuration parser.add_argument('--config', type=str, nargs='+', help='Path(s) to JSON configuration file(s). Can specify multiple files.') parser.add_argument('--output-dir', type=str, help='Output directory for plots and results (overrides config file)') # Training parameters parser.add_argument('--epochs', type=int, default=None, help='Number of training epochs (overrides config)') parser.add_argument('--pretrain-epochs', type=int, default=None, help='Number of pretraining epochs (overrides config)') # Model loading parser.add_argument('--load-model', type=str, help='Path to saved model checkpoint (.pt file) to load and generate outputs without re-training') # Animation parser.add_argument('--animate', action='store_true', help='Generate animation of potential and eigenfunctions') args = parser.parse_args() # Handle validation mode with model loading (no training) if args.load_model: print(f"Loading model from {args.load_model} for validation...") validate_from_checkpoint( args.load_model, output_dir=args.output_dir, generate_anim=args.animate ) return # Handle analytical-only mode if args.mode == 'analytical': # Process config files if args.config: config_files = args.config print(f"Running analytical validation on {len(config_files)} configuration(s)") for i, config_path in enumerate(config_files, 1): print(f"\n{'='*80}") print(f"Configuration {i}/{len(config_files)}: {config_path}") print(f"{'='*80}") config = load_config(config_path) # Override epochs to 0 for analytical mode config.num_epochs = 0 config.num_pretrain_epochs = 0 # Determine output directory if args.output_dir: if len(config_files) == 1: output_dir = args.output_dir else: config_name = getattr(config, 'name', f'run_{i}') output_dir = os.path.join(args.output_dir, config_name) else: output_dir = None try: validate_analytical_only( config=config, output_dir=output_dir, generate_anim=args.animate ) print(f"\n✓ Configuration {i}/{len(config_files)} completed successfully") except Exception as e: print(f"\n✗ Configuration {i}/{len(config_files)} failed with error:") print(f" {type(e).__name__}: {e}") import traceback traceback.print_exc() continue print(f"\n{'='*80}") print(f"All configurations complete: {len(config_files)} total") print(f"{'='*80}") else: # Use default config validate_analytical_only( config=None, output_dir=args.output_dir, generate_anim=args.animate ) return # Handle full validation mode if args.mode == 'validate': if args.config: config_files = args.config print(f"Running validation on {len(config_files)} configuration(s)") for i, config_path in enumerate(config_files, 1): print(f"\n{'='*80}") print(f"Configuration {i}/{len(config_files)}: {config_path}") print(f"{'='*80}") config = load_config(config_path) # Override epochs if specified if args.epochs is not None: config.num_epochs = args.epochs if args.pretrain_epochs is not None: config.num_pretrain_epochs = args.pretrain_epochs # Determine output directory if args.output_dir: if len(config_files) == 1: output_dir = args.output_dir else: config_name = getattr(config, 'name', f'run_{i}') output_dir = os.path.join(args.output_dir, config_name) else: output_dir = None try: validate_run( config=config, output_dir=output_dir, generate_anim=args.animate ) print(f"\n✓ Configuration {i}/{len(config_files)} completed successfully") except Exception as e: print(f"\n✗ Configuration {i}/{len(config_files)} failed with error:") print(f" {type(e).__name__}: {e}") import traceback traceback.print_exc() continue print(f"\n{'='*80}") print(f"All configurations complete: {len(config_files)} total") print(f"{'='*80}") else: # Use default config config = SimpleNamespace(**vars(kdv_config)) if args.epochs is not None: config.num_epochs = args.epochs if args.pretrain_epochs is not None: config.num_pretrain_epochs = args.pretrain_epochs validate_run( config=config, output_dir=args.output_dir, generate_anim=args.animate ) return # Handle simple training mode (legacy behavior) device = get_device() print(f"Using {device} device.") training_config = kdv_config if args.epochs is not None: training_config.num_epochs = args.epochs if args.pretrain_epochs is not None: training_config.num_pretrain_epochs = args.pretrain_epochs torch.manual_seed(training_config.seed) model = KdV_pinn(training_config).to(device) result = train_pinn(model, training_config, device) print("Training complete!") output_path = args.output_dir if args.output_dir else 'pinn_model.pt' if os.path.isdir(output_path): output_path = os.path.join(output_path, 'pinn_model.pt') torch.save({ 'model_state_dict': result['model'].state_dict(), 'optimizer_state_dict': result['optimizer'].state_dict(), 'config': config_to_dict(training_config), 'metrics': result['metrics'] }, output_path) print(f"Model saved to {output_path}") return result if __name__ == '__main__': main()