diff --git a/plan/dev-plan-ml-pipeline.md b/plan/dev-plan-ml-pipeline.md index dc6a36b..5b6bcf3 100644 --- a/plan/dev-plan-ml-pipeline.md +++ b/plan/dev-plan-ml-pipeline.md @@ -555,9 +555,9 @@ class LOAFWeatherEntity(WeatherEntity): ### Milestone 3: Training Pipeline - [x] Create PyTorch Dataset class (`loaf/data/loaders/dataset.py`) -- [ ] Implement training loop (`loaf/training/trainer.py`) -- [ ] Implement evaluation metrics (`loaf/training/evaluate.py`) -- [ ] Write `scripts/train.py` CLI +- [x] Implement training loop (`loaf/training/trainer.py`) +- [x] Implement evaluation metrics (`loaf/training/evaluate.py`) +- [x] Write `scripts/train.py` CLI - [ ] **Verify:** Train on 1 month data, loss decreases, metrics improve ### Milestone 4: Inference & API @@ -704,4 +704,6 @@ Key files to study/adapt from https://github.com/Earth-Intelligence-Lab/Localize **Next coding task:** ~~Implement data loaders (`loaf/data/loaders/`) for ERA5, HRRR, and IEM~~ ✅ DONE -**Next coding task:** Implement training loop and evaluation metrics (Milestone 3) +**Next coding task:** ~~Implement training loop and evaluation metrics (Milestone 3)~~ ✅ DONE + +**Next coding task:** Verify training with real data, then implement inference pipeline (Milestone 4) diff --git a/pyproject.toml b/pyproject.toml index 95ab5ca..d6a2a72 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,31 +25,27 @@ dependencies = [ # ML Framework "torch>=2.0.0", "torch_geometric>=2.4.0", - # Data Access "herbie-data>=2024.0.0", "cdsapi>=0.6.0", "xarray>=2024.0.0", "netCDF4>=1.6.0", "cfgrib>=0.9.0", - # Data Processing "pandas>=2.0.0", "numpy>=1.24.0", "scipy>=1.11.0", "scikit-learn>=1.3.0", - # Graph Construction "networkx>=3.0", - # API & Integration "flask>=3.0.0", "requests>=2.31.0", "pyyaml>=6.0", - # Utilities "tqdm>=4.65.0", "python-dotenv>=1.0.0", + "dask>=2026.1.2", ] [project.optional-dependencies] diff --git a/software/QUICKSTART.md b/software/QUICKSTART.md new file mode 100644 index 0000000..cb017e7 --- /dev/null +++ b/software/QUICKSTART.md @@ -0,0 +1,141 @@ +# LOAF Quickstart Guide + +## Prerequisites + +1. **uv** - Fast Python package manager: `curl -LsSf https://astral.sh/uv/install.sh | sh` +2. **ERA5 API Access** (optional but recommended): Register at https://cds.climate.copernicus.eu/ + +## Setup + +```bash +cd /home/keenan/code/LOAF +uv sync +cd software # All commands below assume you're in the software/ directory +``` + +## Step 1: Download Data + +### Option A: Quick Start (IEM + HRRR only, no registration) + +```bash + +# Download IEM station observations (no auth required) +uv run python -m loaf.data.download.iem \ + --start-date 2024-10-01 \ + --end-date 2024-10-31 \ + --output data/iem/iem_2024_10.parquet + +# Download HRRR forecasts (no auth required) +uv run python -m loaf.data.download.hrrr \ + --start-date 2024-10-01 \ + --end-date 2024-10-31 \ + --output-dir data/hrrr +``` + +### Option B: Full Setup (includes ERA5) + +First, configure ERA5 API credentials in `~/.cdsapirc`: +``` +url: https://cds.climate.copernicus.eu/api +key: YOUR_UID:YOUR_API_KEY +``` + +Then download: +```bash +# Download ERA5 reanalysis +uv run python -m loaf.data.download.era5 \ + --year 2024 \ + --month 10 \ + --output data/era5/era5_2024_10.nc + +# Download IEM + HRRR as above +``` + +## Step 2: Train Model + +```bash +# Train GNN model (default) +uv run python scripts/train.py \ + --config config/seattle.yaml \ + --data-dir data \ + --year 2024 \ + --epochs 50 \ + --checkpoint checkpoints/ + +# Or train ViT model +uv run python scripts/train.py \ + --config config/seattle.yaml \ + --model vit \ + --epochs 50 +``` + +### Training Options + +| Flag | Description | Default | +|------|-------------|---------| +| `--config` | Config YAML file | `config/seattle.yaml` | +| `--model` | Model type: `gnn` or `vit` | `gnn` | +| `--epochs` | Number of epochs | 100 | +| `--batch-size` | Batch size | 64 | +| `--lr` | Learning rate | 1e-4 | +| `--checkpoint` | Directory to save models | None | +| `--resume` | Resume from checkpoint | None | +| `--device` | Device: `auto`, `cuda`, `cpu` | `auto` | + +## Step 3: Verify Installation + +```bash +# Test model forward pass +uv run python scripts/test_model_forward.py + +# Test training loop with synthetic data +uv run python scripts/test_training.py +``` + +## Data Directory Structure + +After downloading, your data directory should look like: +``` +software/data/ +├── hrrr/ +│ ├── hrrr_20241001.nc +│ ├── hrrr_20241002.nc +│ └── ... +├── iem/ +│ └── iem_2024_10.parquet +└── era5/ + └── era5_2024_10.nc +``` + +## CLI Commands + +```bash +uv run loaf-download-hrrr --help # Download HRRR data +uv run loaf-download-era5 --help # Download ERA5 data +uv run loaf-download-iem --help # Download IEM station data +``` + +## Minimal Example (1 day of data for testing) + +```bash +# Download just 1 day of data for testing +uv run python -m loaf.data.download.iem \ + --start-date 2024-10-15 \ + --end-date 2024-10-15 \ + --output data/iem/test.parquet + +uv run python -m loaf.data.download.hrrr \ + --start-date 2024-10-15 \ + --end-date 2024-10-15 \ + --output-dir data/hrrr + +# Quick training test (3 epochs) +uv run python scripts/train.py --epochs 3 --batch-size 8 +``` + +## Troubleshooting + +- **HRRR download fails**: Data may not be available for very recent dates. Try dates 2+ days ago. +- **ERA5 slow**: ERA5 downloads are queued on the CDS server. Initial requests may take 30+ minutes. +- **Out of memory**: Reduce `--batch-size` (try 16 or 8) +- **torch_geometric not found**: Run `uv sync` to install all dependencies diff --git a/software/loaf/data/loaders/dataset.py b/software/loaf/data/loaders/dataset.py index 0d9cbb4..14e9586 100644 --- a/software/loaf/data/loaders/dataset.py +++ b/software/loaf/data/loaders/dataset.py @@ -79,11 +79,26 @@ def __init__( self._setup_coord_normalizers() def _generate_timeline(self, year: int) -> pd.DatetimeIndex: - """Generate hourly timeline for a year (timezone-naive for xarray compat).""" + """Generate hourly timeline for a year (timezone-naive for xarray compat). + + If a grid loader is available, restricts timeline to times that exist + in the grid data to avoid empty data fetches. + """ start = datetime(year, 1, 1) end = datetime(year + 1, 1, 1) times = list(rrule.rrule(rrule.HOURLY, dtstart=start, until=end))[:-1] - return pd.DatetimeIndex(times) # tz-naive for xarray compatibility + timeline = pd.DatetimeIndex(times) # tz-naive for xarray compatibility + + # Restrict timeline to times available in grid data + if self.grid_loader is not None: + grid_times = pd.DatetimeIndex(self.grid_loader.data.time.values) + timeline = timeline.intersection(grid_times) + + # Also restrict to times available in station data + station_times = pd.DatetimeIndex(self.station_loader.data.time.values) + timeline = timeline.intersection(station_times) + + return timeline def _compute_statistics(self) -> None: """Compute normalization statistics.""" diff --git a/software/loaf/data/loaders/hrrr.py b/software/loaf/data/loaders/hrrr.py index 7a30535..3e742b6 100644 --- a/software/loaf/data/loaders/hrrr.py +++ b/software/loaf/data/loaders/hrrr.py @@ -184,6 +184,8 @@ def get_sample( for var in variables: if var in subset.data_vars: values = subset[var].values.astype(np.float32) + # Replace NaN values with 0 + values = np.nan_to_num(values, nan=0.0) # Reshape to (n_nodes, n_times) if values.ndim == 3: # (time, y, x) -> (n_nodes, time) diff --git a/software/loaf/data/loaders/iem.py b/software/loaf/data/loaders/iem.py index 33f7750..5d8896c 100644 --- a/software/loaf/data/loaders/iem.py +++ b/software/loaf/data/loaders/iem.py @@ -245,6 +245,8 @@ def get_sample( for var in variables: if var in subset.data_vars: values = subset[var].values.astype(np.float32) + # Replace any remaining NaN values with 0 + values = np.nan_to_num(values, nan=0.0) result[var] = torch.from_numpy(values) # Also include "is_real" mask if available diff --git a/software/loaf/training/__init__.py b/software/loaf/training/__init__.py index c998f6b..42ec969 100644 --- a/software/loaf/training/__init__.py +++ b/software/loaf/training/__init__.py @@ -1 +1,39 @@ -"""Training pipeline modules.""" +"""Training pipeline modules. + +This module provides the training infrastructure for LOAF models: +- Trainer: Main training loop with checkpointing and early stopping +- TrainingConfig: Configuration dataclass for training parameters +- Loss functions: MSE, WindSpeed, MAE, Combined +- Metrics: Evaluation metrics for weather forecasting +""" + +from loaf.training.evaluate import ( + CombinedLoss, + MAELoss, + Metrics, + MSELoss, + WindSpeedLoss, + get_loss_function, + wind_speed_error, +) +from loaf.training.trainer import ( + Trainer, + TrainingConfig, + TrainingState, +) + +__all__ = [ + # Trainer + "Trainer", + "TrainingConfig", + "TrainingState", + # Loss functions + "MSELoss", + "MAELoss", + "WindSpeedLoss", + "CombinedLoss", + "get_loss_function", + # Metrics + "Metrics", + "wind_speed_error", +] diff --git a/software/loaf/training/evaluate.py b/software/loaf/training/evaluate.py new file mode 100644 index 0000000..3f73a0e --- /dev/null +++ b/software/loaf/training/evaluate.py @@ -0,0 +1,355 @@ +"""Evaluation metrics and validation for weather forecasting. + +Adapted from LocalizedWeather: EvaluateModel.py and Utils/LossFunctions.py +Authors: Qidong Yang & Jonathan Giezendanner (original) +""" + +import torch +from torch import nn + + +def wind_speed_error(output: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + """Compute wind speed error from u/v components. + + This computes the magnitude of the error vector, which is the standard + metric for wind forecast evaluation. + + Args: + output: Predictions with u, v in first two channels (..., 2+) + target: Targets with u, v in first two channels (..., 2+) + + Returns: + Wind speed error for each sample. + """ + u_error = output[..., 0] - target[..., 0] + v_error = output[..., 1] - target[..., 1] + + # Add epsilon for numerical stability + error = torch.sqrt(u_error**2 + v_error**2 + torch.finfo(torch.float32).eps) + return error + + +class WindSpeedLoss(nn.Module): + """Loss function based on wind speed error (magnitude of error vector).""" + + def __init__(self, reduction: str = "mean"): + super().__init__() + self.reduction = reduction + + def forward( + self, + output: torch.Tensor, + target: torch.Tensor, + mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Compute wind speed loss. + + Args: + output: Predictions (..., 2) with u, v components. + target: Targets (..., 2) with u, v components. + mask: Optional boolean mask for valid samples. + + Returns: + Scalar loss value. + """ + error = wind_speed_error(output, target) + + if mask is not None: + error = error[mask] + + if self.reduction == "mean": + return error.mean() + elif self.reduction == "sum": + return error.sum() + else: + return error + + +class MSELoss(nn.Module): + """Standard MSE loss with optional masking.""" + + def __init__(self, reduction: str = "mean"): + super().__init__() + self.reduction = reduction + self.mse = nn.MSELoss(reduction="none") + + def forward( + self, + output: torch.Tensor, + target: torch.Tensor, + mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Compute MSE loss. + + Args: + output: Predictions. + target: Targets. + mask: Optional boolean mask for valid samples. + + Returns: + Scalar loss value. + """ + error = self.mse(output, target) + + if mask is not None: + error = error[mask] + + if self.reduction == "mean": + return error.mean() + elif self.reduction == "sum": + return error.sum() + else: + return error + + +class MAELoss(nn.Module): + """Mean Absolute Error loss with optional masking.""" + + def __init__(self, reduction: str = "mean"): + super().__init__() + self.reduction = reduction + + def forward( + self, + output: torch.Tensor, + target: torch.Tensor, + mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Compute MAE loss. + + Args: + output: Predictions. + target: Targets. + mask: Optional boolean mask for valid samples. + + Returns: + Scalar loss value. + """ + error = torch.abs(output - target) + + if mask is not None: + error = error[mask] + + if self.reduction == "mean": + return error.mean() + elif self.reduction == "sum": + return error.sum() + else: + return error + + +class CombinedLoss(nn.Module): + """Combined loss for wind (u, v) and other variables. + + Uses wind speed error for u/v components and RMSE for other variables. + This matches the LocalizedWeather paper methodology. + """ + + def __init__( + self, + output_vars: list[str], + reduction: str = "mean", + ): + """Initialize combined loss. + + Args: + output_vars: List of output variable names (e.g., ["u", "v", "temp"]). + reduction: Reduction method ("mean" or "sum"). + """ + super().__init__() + self.output_vars = output_vars + self.reduction = reduction + + # Find indices for u, v, and other variables + self.u_idx = output_vars.index("u") if "u" in output_vars else None + self.v_idx = output_vars.index("v") if "v" in output_vars else None + self.other_indices = [ + i for i, v in enumerate(output_vars) if v not in ["u", "v"] + ] + + def forward( + self, + output: torch.Tensor, + target: torch.Tensor, + mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Compute combined loss. + + Args: + output: Predictions (..., n_vars). + target: Targets (..., n_vars). + mask: Optional boolean mask for valid samples. + + Returns: + Scalar loss value. + """ + errors = [] + + # Wind speed error for u, v components + if self.u_idx is not None and self.v_idx is not None: + u_error = output[..., self.u_idx] - target[..., self.u_idx] + v_error = output[..., self.v_idx] - target[..., self.v_idx] + wind_error = torch.sqrt( + u_error**2 + v_error**2 + torch.finfo(torch.float32).eps + ) + if mask is not None: + wind_error = wind_error[mask[..., 0] if mask.dim() > 1 else mask] + errors.append(wind_error) + + # RMSE for other variables + for idx in self.other_indices: + var_error = torch.sqrt( + (output[..., idx] - target[..., idx]) ** 2 + + torch.finfo(torch.float32).eps + ) + if mask is not None: + var_error = var_error[mask[..., idx] if mask.dim() > 1 else mask] + errors.append(var_error) + + if not errors: + return torch.tensor(0.0, device=output.device) + + # Concatenate all errors + all_errors = torch.cat([e.flatten() for e in errors]) + + if self.reduction == "mean": + return all_errors.mean() + elif self.reduction == "sum": + return all_errors.sum() + else: + return all_errors + + +class Metrics: + """Compute and track evaluation metrics.""" + + def __init__(self, output_vars: list[str]): + """Initialize metrics tracker. + + Args: + output_vars: List of output variable names. + """ + self.output_vars = output_vars + self.reset() + + def reset(self) -> None: + """Reset accumulated metrics.""" + self.mse_sum = {var: 0.0 for var in self.output_vars} + self.mae_sum = {var: 0.0 for var in self.output_vars} + self.count = {var: 0 for var in self.output_vars} + + # Wind-specific metrics + if "u" in self.output_vars and "v" in self.output_vars: + self.wind_speed_error_sum = 0.0 + self.wind_count = 0 + + def update( + self, + output: torch.Tensor, + target: torch.Tensor, + mask: torch.Tensor | None = None, + ) -> None: + """Update metrics with a batch of predictions. + + Args: + output: Predictions (..., n_vars). + target: Targets (..., n_vars). + mask: Optional boolean mask for valid samples (..., n_vars). + """ + output = output.detach() + target = target.detach() + + for i, var in enumerate(self.output_vars): + pred = output[..., i] + tgt = target[..., i] + + if mask is not None: + var_mask = mask[..., i] if mask.dim() > 1 else mask + pred = pred[var_mask] + tgt = tgt[var_mask] + + n = pred.numel() + if n > 0: + self.mse_sum[var] += ((pred - tgt) ** 2).sum().item() + self.mae_sum[var] += torch.abs(pred - tgt).sum().item() + self.count[var] += n + + # Wind speed error + if "u" in self.output_vars and "v" in self.output_vars: + u_idx = self.output_vars.index("u") + v_idx = self.output_vars.index("v") + + u_pred = output[..., u_idx] + v_pred = output[..., v_idx] + u_tgt = target[..., u_idx] + v_tgt = target[..., v_idx] + + if mask is not None: + u_mask = mask[..., u_idx] if mask.dim() > 1 else mask + u_pred = u_pred[u_mask] + v_pred = v_pred[u_mask] + u_tgt = u_tgt[u_mask] + v_tgt = v_tgt[u_mask] + + n = u_pred.numel() + if n > 0: + wind_err = torch.sqrt( + (u_pred - u_tgt) ** 2 + + (v_pred - v_tgt) ** 2 + + torch.finfo(torch.float32).eps + ) + self.wind_speed_error_sum += wind_err.sum().item() + self.wind_count += n + + def compute(self) -> dict[str, float]: + """Compute final metrics. + + Returns: + Dictionary of metric names to values. + """ + results = {} + + # Per-variable metrics + for var in self.output_vars: + if self.count[var] > 0: + results[f"mse_{var}"] = self.mse_sum[var] / self.count[var] + results[f"rmse_{var}"] = (self.mse_sum[var] / self.count[var]) ** 0.5 + results[f"mae_{var}"] = self.mae_sum[var] / self.count[var] + + # Overall metrics + total_mse = sum(self.mse_sum.values()) + total_count = sum(self.count.values()) + if total_count > 0: + results["mse"] = total_mse / total_count + results["rmse"] = (total_mse / total_count) ** 0.5 + + # Wind speed error + if hasattr(self, "wind_count") and self.wind_count > 0: + results["wind_speed_error"] = self.wind_speed_error_sum / self.wind_count + + return results + + +def get_loss_function( + loss_type: str, + output_vars: list[str], +) -> nn.Module: + """Get loss function by name. + + Args: + loss_type: Loss function type ("mse", "wind_speed", "combined"). + output_vars: List of output variable names. + + Returns: + Loss function module. + """ + if loss_type == "mse": + return MSELoss() + elif loss_type == "wind_speed": + return WindSpeedLoss() + elif loss_type == "combined": + return CombinedLoss(output_vars) + elif loss_type == "mae": + return MAELoss() + else: + raise ValueError(f"Unknown loss type: {loss_type}") diff --git a/software/loaf/training/trainer.py b/software/loaf/training/trainer.py new file mode 100644 index 0000000..b7a2571 --- /dev/null +++ b/software/loaf/training/trainer.py @@ -0,0 +1,645 @@ +"""Training loop for weather forecasting models. + +Adapted from LocalizedWeather: Main.py and EvaluateModel.py +Authors: Qidong Yang & Jonathan Giezendanner (original) +""" + +import logging +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import torch +from torch import nn +from torch.optim import AdamW +from torch.optim.lr_scheduler import ReduceLROnPlateau +from torch.utils.data import DataLoader +from tqdm import tqdm + +from loaf.training.evaluate import Metrics, get_loss_function + +logger = logging.getLogger(__name__) + + +@dataclass +class TrainingConfig: + """Configuration for training.""" + + # Data + back_hrs: int = 24 # Number of historical hours for input + + # Optimization + learning_rate: float = 1e-4 + weight_decay: float = 1e-4 + epochs: int = 100 + batch_size: int = 64 + + # Early stopping + patience: int = 10 + min_delta: float = 1e-6 + + # Gradient clipping + max_grad_norm: float = 1.0 + + # Learning rate scheduling + lr_scheduler: str = "reduce_on_plateau" + lr_factor: float = 0.5 + lr_patience: int = 5 + + # Checkpointing + checkpoint_dir: Path | None = None + save_best_only: bool = True + save_every: int = 10 + + # Logging + log_interval: int = 10 + + # Loss function + loss_type: str = "mse" + + # Output variables + output_vars: list[str] = field(default_factory=lambda: ["u", "v"]) + + # Device + device: str = "auto" + + # Random seed + seed: int = 42 + + +@dataclass +class TrainingState: + """Training state for checkpointing and resumption.""" + + epoch: int = 0 + best_val_loss: float = float("inf") + epochs_without_improvement: int = 0 + train_losses: list[float] = field(default_factory=list) + val_losses: list[float] = field(default_factory=list) + + +class Trainer: + """Trainer for weather forecasting models. + + Supports both GNN (MPNN) and Transformer (ViT) models. + """ + + def __init__( + self, + model: nn.Module, + config: TrainingConfig | None = None, + model_type: str = "gnn", + ): + """Initialize trainer. + + Args: + model: PyTorch model to train. + config: Training configuration. + model_type: Type of model ("gnn" or "vit"). + """ + self.config = config or TrainingConfig() + self.model_type = model_type + + # Set device + if self.config.device == "auto": + if torch.cuda.is_available(): + self.device = torch.device("cuda") + elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + self.device = torch.device("mps") + else: + self.device = torch.device("cpu") + else: + self.device = torch.device(self.config.device) + + logger.info(f"Using device: {self.device}") + + # Move model to device + self.model = model.to(self.device) + + # Multi-GPU support + if torch.cuda.device_count() > 1: + logger.info(f"Using {torch.cuda.device_count()} GPUs") + self.model = nn.DataParallel(self.model) + + # Initialize optimizer + self.optimizer = AdamW( + self.model.parameters(), + lr=self.config.learning_rate, + weight_decay=self.config.weight_decay, + ) + + # Initialize scheduler + if self.config.lr_scheduler == "reduce_on_plateau": + self.scheduler = ReduceLROnPlateau( + self.optimizer, + mode="min", + factor=self.config.lr_factor, + patience=self.config.lr_patience, + ) + else: + self.scheduler = None + + # Initialize loss function + self.loss_fn = get_loss_function( + self.config.loss_type, + self.config.output_vars, + ) + + # Initialize metrics + self.metrics = Metrics(self.config.output_vars) + + # Training state + self.state = TrainingState() + + # Set random seed + torch.manual_seed(self.config.seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(self.config.seed) + + def train( + self, + train_loader: DataLoader, + val_loader: DataLoader | None = None, + ) -> dict[str, Any]: + """Run training loop. + + Args: + train_loader: Training data loader. + val_loader: Validation data loader (optional). + + Returns: + Dictionary with training history and final metrics. + """ + logger.info("Starting training...") + logger.info(f" Epochs: {self.config.epochs}") + logger.info(f" Batch size: {self.config.batch_size}") + logger.info(f" Learning rate: {self.config.learning_rate}") + logger.info(f" Device: {self.device}") + + history = { + "train_loss": [], + "val_loss": [], + "train_metrics": [], + "val_metrics": [], + } + + for epoch in range(self.state.epoch, self.config.epochs): + self.state.epoch = epoch + + # Training epoch + train_loss, train_metrics = self._train_epoch(train_loader, epoch) + history["train_loss"].append(train_loss) + history["train_metrics"].append(train_metrics) + self.state.train_losses.append(train_loss) + + # Validation epoch + if val_loader is not None: + val_loss, val_metrics = self._validate_epoch(val_loader, epoch) + history["val_loss"].append(val_loss) + history["val_metrics"].append(val_metrics) + self.state.val_losses.append(val_loss) + + # Learning rate scheduling + if self.scheduler is not None: + self.scheduler.step(val_loss) + + # Early stopping check + if val_loss < self.state.best_val_loss - self.config.min_delta: + self.state.best_val_loss = val_loss + self.state.epochs_without_improvement = 0 + + # Save best model + if self.config.checkpoint_dir and self.config.save_best_only: + self._save_checkpoint("best.pt") + else: + self.state.epochs_without_improvement += 1 + if self.state.epochs_without_improvement >= self.config.patience: + logger.info( + f"Early stopping at epoch {epoch + 1} " + f"(no improvement for {self.config.patience} epochs)" + ) + break + + # Periodic checkpointing + if ( + self.config.checkpoint_dir + and not self.config.save_best_only + and (epoch + 1) % self.config.save_every == 0 + ): + self._save_checkpoint(f"epoch_{epoch + 1}.pt") + + # Log epoch summary + self._log_epoch_summary(epoch, train_loss, train_metrics, val_loader is not None) + + # Save final checkpoint + if self.config.checkpoint_dir: + self._save_checkpoint("final.pt") + + logger.info("Training complete!") + logger.info(f"Best validation loss: {self.state.best_val_loss:.6f}") + + return history + + def _train_epoch( + self, + loader: DataLoader, + epoch: int, + ) -> tuple[float, dict[str, float]]: + """Run a single training epoch. + + Args: + loader: Training data loader. + epoch: Current epoch number. + + Returns: + Tuple of (loss, metrics dictionary). + """ + self.model.train() + self.metrics.reset() + + total_loss = 0.0 + n_batches = 0 + + pbar = tqdm(loader, desc=f"Epoch {epoch + 1} [Train]", leave=False) + + for batch_idx, batch in enumerate(pbar): + loss, output, target = self._train_step(batch) + + total_loss += loss.item() + n_batches += 1 + + # Update metrics + self.metrics.update(output, target) + + # Update progress bar + pbar.set_postfix({"loss": f"{loss.item():.4f}"}) + + # Log periodically + if (batch_idx + 1) % self.config.log_interval == 0: + avg_loss = total_loss / n_batches + logger.debug( + f"Epoch {epoch + 1}, Batch {batch_idx + 1}/{len(loader)}, " + f"Loss: {avg_loss:.4f}" + ) + + avg_loss = total_loss / n_batches + metrics = self.metrics.compute() + + return avg_loss, metrics + + def _train_step( + self, + batch: dict[str, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Run a single training step. + + Args: + batch: Batch of data from dataloader. + + Returns: + Tuple of (loss, predictions, targets). + """ + self.optimizer.zero_grad() + + # Forward pass + output, target = self._forward(batch) + + # Compute loss + loss = self.loss_fn(output, target) + + # Backward pass + loss.backward() + + # Gradient clipping + if self.config.max_grad_norm > 0: + if isinstance(self.model, nn.DataParallel): + torch.nn.utils.clip_grad_norm_( + self.model.module.parameters(), + self.config.max_grad_norm, + ) + else: + torch.nn.utils.clip_grad_norm_( + self.model.parameters(), + self.config.max_grad_norm, + ) + + # Optimizer step + self.optimizer.step() + + return loss, output.detach(), target.detach() + + def _validate_epoch( + self, + loader: DataLoader, + epoch: int, + ) -> tuple[float, dict[str, float]]: + """Run a single validation epoch. + + Args: + loader: Validation data loader. + epoch: Current epoch number. + + Returns: + Tuple of (loss, metrics dictionary). + """ + self.model.eval() + self.metrics.reset() + + total_loss = 0.0 + n_batches = 0 + + pbar = tqdm(loader, desc=f"Epoch {epoch + 1} [Val]", leave=False) + + with torch.no_grad(): + for batch in pbar: + output, target = self._forward(batch) + loss = self.loss_fn(output, target) + + total_loss += loss.item() + n_batches += 1 + + # Update metrics + self.metrics.update(output, target) + + pbar.set_postfix({"loss": f"{loss.item():.4f}"}) + + avg_loss = total_loss / n_batches + metrics = self.metrics.compute() + + return avg_loss, metrics + + def _forward( + self, + batch: dict[str, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor]: + """Run forward pass through the model. + + Args: + batch: Batch of data from dataloader. + + Returns: + Tuple of (predictions, targets). + """ + # Extract data from batch based on model type + if self.model_type == "gnn": + return self._forward_gnn(batch) + elif self.model_type == "vit": + return self._forward_vit(batch) + else: + raise ValueError(f"Unknown model type: {self.model_type}") + + def _forward_gnn( + self, + batch: dict[str, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor]: + """Forward pass for GNN (MPNN) model. + + Args: + batch: Batch of data from dataloader. + + Returns: + Tuple of (predictions, targets). + """ + # Get station data + # Station observations: (batch, time, n_stations, n_vars) -> need to reshape + # The dataset returns: {var: (batch, n_stations, time)} for each var + + # Stack input variables + input_vars = [] + for var in ["u", "v", "temp", "dewpoint"]: + if var in batch: + input_vars.append(batch[var].to(self.device)) + + if input_vars: + # Stack along last dimension: (batch, n_stations, time, n_vars) + # Dataset returns (n_stations, n_times) per var, after batching: (batch, n_stations, n_times) + # After stacking: (batch, n_stations, n_times, n_vars) + # MPNN expects: (batch, n_stations, time, n_vars) - already correct! + madis_x = torch.stack(input_vars, dim=-1) + # Only use first back_hrs timesteps for input (rest is for targets) + madis_x = madis_x[:, :, : self.config.back_hrs, :] + else: + raise ValueError("No station input variables found in batch") + + # Get station positions + madis_lon = batch["station_lon"].to(self.device) + madis_lat = batch["station_lat"].to(self.device) + + # Add extra dimension for positions if needed + if madis_lon.dim() == 2: + madis_lon = madis_lon.unsqueeze(-1) + madis_lat = madis_lat.unsqueeze(-1) + + # Get station graph edges + edge_index = batch["k_edge_index"].to(self.device) + + # Get external grid data if available + ext_vars = [] + for var in ["ext_u", "ext_v", "ext_temp", "ext_dewpoint"]: + if var in batch: + ext_vars.append(batch[var].to(self.device)) + + if ext_vars: + # Grid data also returns (n_nodes, n_times) per var + # After batching and stacking: (batch, n_nodes, n_times, n_vars) - already correct + ext_x = torch.stack(ext_vars, dim=-1) + # Only use first back_hrs timesteps for input + ext_x = ext_x[:, :, : self.config.back_hrs, :] + + ext_lon = batch["grid_lon"].to(self.device) + ext_lat = batch["grid_lat"].to(self.device) + + if ext_lon.dim() == 2: + ext_lon = ext_lon.unsqueeze(-1) + ext_lat = ext_lat.unsqueeze(-1) + + edge_index_e2m = batch["ex2m_edge_index"].to(self.device) + else: + ext_x = None + ext_lon = None + ext_lat = None + edge_index_e2m = None + + # Forward pass + output = self.model( + madis_x, + madis_lon, + madis_lat, + edge_index, + ext_lon, + ext_lat, + ext_x, + edge_index_e2m, + ) + # output: (batch, n_stations, n_out_vars) + + # Get target (last timestep of input vars that are in output_vars) + # batch[var] shape: (batch, n_stations, n_times) + target_vars = [] + for var in self.config.output_vars: + if var in batch: + # Get last timestep: [:, :, -1] -> (batch, n_stations) + target_vars.append(batch[var][:, :, -1].to(self.device)) + + target = torch.stack(target_vars, dim=-1) + # target: (batch, n_stations, n_out_vars) + + return output, target + + def _forward_vit( + self, + batch: dict[str, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor]: + """Forward pass for ViT model. + + Args: + batch: Batch of data from dataloader. + + Returns: + Tuple of (predictions, targets). + """ + # Stack input variables + input_vars = [] + for var in ["u", "v", "temp", "dewpoint"]: + if var in batch: + input_vars.append(batch[var].to(self.device)) + + if input_vars: + # Stack along last dimension: (batch, time, n_stations, n_vars) + madis_x = torch.stack(input_vars, dim=-1) + # ViT expects: (batch, n_stations, time, n_vars) + if madis_x.dim() == 4: + madis_x = madis_x.permute(0, 2, 1, 3).contiguous() + else: + raise ValueError("No station input variables found in batch") + + # Get ERA5 data if available + era5_vars = [] + for var in ["ext_u", "ext_v", "ext_temp", "ext_dewpoint"]: + if var in batch: + era5_vars.append(batch[var].to(self.device)) + + if era5_vars: + era5_x = torch.stack(era5_vars, dim=-1) + if era5_x.dim() == 4: + era5_x = era5_x.permute(0, 2, 1, 3).contiguous() + else: + era5_x = None + + # Forward pass + output, _ = self.model(madis_x, era5_x, return_attn=False) + # output: (batch, n_stations, n_out_vars) + + # Get target (last timestep) + target_vars = [] + for var in self.config.output_vars: + if var in batch: + target_vars.append(batch[var][:, -1, :].to(self.device)) + + target = torch.stack(target_vars, dim=-1) + + return output, target + + def _log_epoch_summary( + self, + epoch: int, + train_loss: float, + train_metrics: dict[str, float], + has_val: bool, + ) -> None: + """Log epoch summary.""" + msg = f"Epoch {epoch + 1}/{self.config.epochs}" + msg += f" | Train Loss: {train_loss:.4f}" + + if "wind_speed_error" in train_metrics: + msg += f" | Wind Err: {train_metrics['wind_speed_error']:.4f}" + + if has_val and self.state.val_losses: + msg += f" | Val Loss: {self.state.val_losses[-1]:.4f}" + + msg += f" | LR: {self.optimizer.param_groups[0]['lr']:.2e}" + + logger.info(msg) + + def _save_checkpoint(self, filename: str) -> None: + """Save model checkpoint. + + Args: + filename: Checkpoint filename. + """ + if self.config.checkpoint_dir is None: + return + + checkpoint_dir = Path(self.config.checkpoint_dir) + checkpoint_dir.mkdir(parents=True, exist_ok=True) + + checkpoint = { + "epoch": self.state.epoch, + "model_state_dict": ( + self.model.module.state_dict() + if isinstance(self.model, nn.DataParallel) + else self.model.state_dict() + ), + "optimizer_state_dict": self.optimizer.state_dict(), + "scheduler_state_dict": ( + self.scheduler.state_dict() if self.scheduler else None + ), + "state": { + "epoch": self.state.epoch, + "best_val_loss": self.state.best_val_loss, + "epochs_without_improvement": self.state.epochs_without_improvement, + "train_losses": self.state.train_losses, + "val_losses": self.state.val_losses, + }, + "config": self.config, + } + + path = checkpoint_dir / filename + torch.save(checkpoint, path) + logger.info(f"Saved checkpoint: {path}") + + def load_checkpoint(self, path: str | Path) -> None: + """Load model checkpoint. + + Args: + path: Path to checkpoint file. + """ + path = Path(path) + checkpoint = torch.load(path, map_location=self.device) + + # Load model state + if isinstance(self.model, nn.DataParallel): + self.model.module.load_state_dict(checkpoint["model_state_dict"]) + else: + self.model.load_state_dict(checkpoint["model_state_dict"]) + + # Load optimizer state + self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) + + # Load scheduler state + if self.scheduler and checkpoint["scheduler_state_dict"]: + self.scheduler.load_state_dict(checkpoint["scheduler_state_dict"]) + + # Load training state + state_dict = checkpoint["state"] + self.state.epoch = state_dict["epoch"] + self.state.best_val_loss = state_dict["best_val_loss"] + self.state.epochs_without_improvement = state_dict["epochs_without_improvement"] + self.state.train_losses = state_dict["train_losses"] + self.state.val_losses = state_dict["val_losses"] + + logger.info(f"Loaded checkpoint from epoch {self.state.epoch}") + + def evaluate( + self, + loader: DataLoader, + ) -> tuple[float, dict[str, float]]: + """Evaluate model on a dataset. + + Args: + loader: Data loader for evaluation. + + Returns: + Tuple of (loss, metrics dictionary). + """ + return self._validate_epoch(loader, 0) diff --git a/software/scripts/test_training.py b/software/scripts/test_training.py new file mode 100644 index 0000000..dffc8ab --- /dev/null +++ b/software/scripts/test_training.py @@ -0,0 +1,306 @@ +#!/usr/bin/env python3 +"""Test script to verify training loop works with synthetic data. + +This script tests: +1. Loss functions +2. Metrics computation +3. Trainer with synthetic batches +""" + +import sys +from pathlib import Path + +import torch +from torch.utils.data import DataLoader, TensorDataset + +# Add software directory to path for imports +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from loaf.model import MPNN, VisionTransformer +from loaf.training import ( + Trainer, + TrainingConfig, + Metrics, + MSELoss, + WindSpeedLoss, + CombinedLoss, + wind_speed_error, +) + + +def test_wind_speed_error(): + """Test wind speed error calculation.""" + print("=" * 60) + print("Testing wind_speed_error") + print("=" * 60) + + # Perfect prediction (epsilon is added for numerical stability, so not exactly 0) + output = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + target = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + error = wind_speed_error(output, target) + # Error should be very small (just epsilon) + assert error.max() < 1e-3 + print(" Perfect prediction: PASSED") + + # Known error - should be approximately 1.0 (sqrt(1^2 + 0^2 + eps)) + output = torch.tensor([[1.0, 0.0]]) + target = torch.tensor([[0.0, 0.0]]) + error = wind_speed_error(output, target) + assert torch.abs(error - 1.0).max() < 1e-3 + print(" Unit error: PASSED") + + print("\nwind_speed_error: PASSED") + + +def test_loss_functions(): + """Test loss function implementations.""" + print("\n" + "=" * 60) + print("Testing Loss Functions") + print("=" * 60) + + batch_size = 4 + n_stations = 10 + n_vars = 2 + + output = torch.randn(batch_size, n_stations, n_vars) + target = torch.randn(batch_size, n_stations, n_vars) + + # MSE Loss + mse_loss = MSELoss() + loss = mse_loss(output, target) + assert loss.dim() == 0 # Scalar + assert loss.item() >= 0 + print(f" MSELoss: {loss.item():.4f} (shape: {loss.shape})") + + # Wind Speed Loss + ws_loss = WindSpeedLoss() + loss = ws_loss(output, target) + assert loss.dim() == 0 + assert loss.item() >= 0 + print(f" WindSpeedLoss: {loss.item():.4f}") + + # Combined Loss + combined_loss = CombinedLoss(["u", "v"]) + loss = combined_loss(output, target) + assert loss.dim() == 0 + assert loss.item() >= 0 + print(f" CombinedLoss: {loss.item():.4f}") + + # Test with mask + mask = torch.rand(batch_size, n_stations) > 0.3 + loss = mse_loss(output, target, mask) + print(f" MSELoss with mask: {loss.item():.4f}") + + print("\nLoss functions: PASSED") + + +def test_metrics(): + """Test metrics computation.""" + print("\n" + "=" * 60) + print("Testing Metrics") + print("=" * 60) + + metrics = Metrics(["u", "v"]) + + # Generate some batches + for _ in range(5): + output = torch.randn(4, 10, 2) + target = torch.randn(4, 10, 2) + metrics.update(output, target) + + results = metrics.compute() + + print(" Computed metrics:") + for name, value in results.items(): + print(f" {name}: {value:.4f}") + + assert "mse_u" in results + assert "mse_v" in results + assert "rmse_u" in results + assert "rmse_v" in results + assert "mae_u" in results + assert "mae_v" in results + assert "wind_speed_error" in results + + # Reset should clear all + metrics.reset() + results = metrics.compute() + assert len([v for v in results.values() if v > 0]) == 0 + + print("\nMetrics: PASSED") + + +def create_synthetic_dataloader(n_samples, n_stations, n_hours, n_vars, batch_size): + """Create a dataloader with synthetic weather data.""" + # Create synthetic tensors matching expected dataset format + # For simplicity, we create a TensorDataset that yields dicts + + class SyntheticDataset(torch.utils.data.Dataset): + def __init__(self, n_samples, n_stations, n_hours, n_vars): + self.n_samples = n_samples + self.n_stations = n_stations + self.n_hours = n_hours + self.n_vars = n_vars + + # Pre-generate some data + self.station_lon = torch.rand(n_stations) * 3 - 124 # -124 to -121 + self.station_lat = torch.rand(n_stations) * 2.5 + 46.5 # 46.5 to 49 + + # Create edge index (simple ring graph) + edges = [] + for i in range(n_stations): + edges.append([i, (i + 1) % n_stations]) + edges.append([(i + 1) % n_stations, i]) + self.edge_index = torch.tensor(edges, dtype=torch.long).t() + + def __len__(self): + return self.n_samples + + def __getitem__(self, idx): + return { + "u": torch.randn(self.n_hours, self.n_stations), + "v": torch.randn(self.n_hours, self.n_stations), + "temp": torch.randn(self.n_hours, self.n_stations), + "dewpoint": torch.randn(self.n_hours, self.n_stations), + "station_lon": self.station_lon, + "station_lat": self.station_lat, + "k_edge_index": self.edge_index, + } + + dataset = SyntheticDataset(n_samples, n_stations, n_hours, n_vars) + return DataLoader(dataset, batch_size=batch_size, shuffle=True) + + +def test_trainer_gnn(): + """Test trainer with GNN model.""" + print("\n" + "=" * 60) + print("Testing Trainer with MPNN") + print("=" * 60) + + # Model parameters + n_stations = 10 + n_hours = 24 + n_vars = 4 + hidden_dim = 32 + n_out = 2 + + # Create model + model = MPNN( + n_passing=2, + lead_hrs=6, + n_node_features_m=n_hours * n_vars, + n_node_features_e=n_hours * n_vars, + n_out_features=n_out, + hidden_dim=hidden_dim, + ) + print(f" Model parameters: {sum(p.numel() for p in model.parameters()):,}") + + # Create synthetic dataloaders + train_loader = create_synthetic_dataloader( + n_samples=32, n_stations=n_stations, n_hours=n_hours, n_vars=n_vars, batch_size=8 + ) + val_loader = create_synthetic_dataloader( + n_samples=16, n_stations=n_stations, n_hours=n_hours, n_vars=n_vars, batch_size=8 + ) + + # Create trainer + config = TrainingConfig( + epochs=3, + learning_rate=1e-3, + batch_size=8, + patience=5, + device="cpu", + output_vars=["u", "v"], + ) + + trainer = Trainer(model, config, model_type="gnn") + + # Train for a few epochs + print(" Training for 3 epochs...") + history = trainer.train(train_loader, val_loader) + + assert len(history["train_loss"]) == 3 + assert len(history["val_loss"]) == 3 + print(f" Final train loss: {history['train_loss'][-1]:.4f}") + print(f" Final val loss: {history['val_loss'][-1]:.4f}") + + print("\nTrainer with MPNN: PASSED") + + +def test_trainer_vit(): + """Test trainer with ViT model.""" + print("\n" + "=" * 60) + print("Testing Trainer with VisionTransformer") + print("=" * 60) + + # Model parameters + n_stations = 10 + n_hours = 24 + n_vars = 4 + n_out = 2 + + # Create model + model = VisionTransformer( + n_stations=n_stations, + madis_len=n_hours, + madis_n_vars_i=n_vars, + madis_n_vars_o=n_out, + dim=32, + attn_dim=16, + mlp_dim=64, + num_heads=2, + num_layers=2, + ) + print(f" Model parameters: {sum(p.numel() for p in model.parameters()):,}") + + # Create synthetic dataloaders + train_loader = create_synthetic_dataloader( + n_samples=32, n_stations=n_stations, n_hours=n_hours, n_vars=n_vars, batch_size=8 + ) + val_loader = create_synthetic_dataloader( + n_samples=16, n_stations=n_stations, n_hours=n_hours, n_vars=n_vars, batch_size=8 + ) + + # Create trainer + config = TrainingConfig( + epochs=3, + learning_rate=1e-3, + batch_size=8, + patience=5, + device="cpu", + output_vars=["u", "v"], + ) + + trainer = Trainer(model, config, model_type="vit") + + # Train for a few epochs + print(" Training for 3 epochs...") + history = trainer.train(train_loader, val_loader) + + assert len(history["train_loss"]) == 3 + assert len(history["val_loss"]) == 3 + print(f" Final train loss: {history['train_loss'][-1]:.4f}") + print(f" Final val loss: {history['val_loss'][-1]:.4f}") + + print("\nTrainer with VisionTransformer: PASSED") + + +def main(): + """Run all tests.""" + print("\n" + "#" * 60) + print("# LOAF Training Pipeline Tests") + print("#" * 60) + + test_wind_speed_error() + test_loss_functions() + test_metrics() + test_trainer_gnn() + test_trainer_vit() + + print("\n" + "=" * 60) + print("ALL TRAINING TESTS PASSED!") + print("=" * 60 + "\n") + + +if __name__ == "__main__": + main() diff --git a/software/scripts/train.py b/software/scripts/train.py new file mode 100644 index 0000000..17636fa --- /dev/null +++ b/software/scripts/train.py @@ -0,0 +1,355 @@ +#!/usr/bin/env python3 +"""Training script for LOAF weather forecasting models. + +Usage: + python scripts/train.py --config config/seattle.yaml + python scripts/train.py --config config/seattle.yaml --model vit --epochs 50 + python scripts/train.py --config config/seattle.yaml --checkpoint checkpoints/ +""" + +import argparse +import logging +import sys +from pathlib import Path + +import torch + +# Add software directory to path for imports +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from loaf.config import get_config +from loaf.data.loaders.dataset import WeatherDataset, create_dataloaders +from loaf.data.loaders.iem import IEMLoader +from loaf.data.loaders.hrrr import HRRRLoader +from loaf.data.loaders.stations import StationMetadata +from loaf.model import MPNN, VisionTransformer +from loaf.training.trainer import Trainer, TrainingConfig + +# Configure logging +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", +) +logger = logging.getLogger(__name__) + + +def parse_args(): + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="Train LOAF weather forecasting model", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + + # Configuration + parser.add_argument( + "--config", + type=str, + default="config/seattle.yaml", + help="Path to configuration YAML file", + ) + + # Data + parser.add_argument( + "--data-dir", + type=str, + default="data", + help="Path to data directory", + ) + parser.add_argument( + "--year", + type=int, + default=2024, + help="Year of data to use for training", + ) + + # Model + parser.add_argument( + "--model", + type=str, + choices=["gnn", "vit"], + default="gnn", + help="Model type to train", + ) + + # Training overrides + parser.add_argument( + "--epochs", + type=int, + default=None, + help="Number of training epochs (overrides config)", + ) + parser.add_argument( + "--batch-size", + type=int, + default=None, + help="Batch size (overrides config)", + ) + parser.add_argument( + "--lr", + type=float, + default=None, + help="Learning rate (overrides config)", + ) + parser.add_argument( + "--patience", + type=int, + default=None, + help="Early stopping patience (overrides config)", + ) + + # Checkpointing + parser.add_argument( + "--checkpoint", + type=str, + default=None, + help="Checkpoint directory for saving models", + ) + parser.add_argument( + "--resume", + type=str, + default=None, + help="Path to checkpoint to resume from", + ) + + # Device + parser.add_argument( + "--device", + type=str, + default="auto", + help="Device to use (auto, cuda, cpu, mps)", + ) + + # Misc + parser.add_argument( + "--seed", + type=int, + default=42, + help="Random seed for reproducibility", + ) + parser.add_argument( + "--debug", + action="store_true", + help="Enable debug logging", + ) + + return parser.parse_args() + + +def create_model( + model_type: str, + config, + n_stations: int, + back_hrs: int, + n_vars_in: int = 4, + n_vars_out: int = 2, + n_grid: int | None = None, +): + """Create model based on type and configuration. + + Args: + model_type: "gnn" or "vit". + config: Configuration object. + n_stations: Number of weather stations. + back_hrs: Number of historical hours. + n_vars_in: Number of input variables. + n_vars_out: Number of output variables. + n_grid: Number of grid points (for GNN external layers). + + Returns: + PyTorch model. + """ + hidden_dim = getattr(config.model, "hidden_dim", 128) + num_gnn_layers = getattr(config.model, "num_gnn_layers", 2) + num_transformer_layers = getattr(config.model, "num_transformer_layers", 5) + num_heads = getattr(config.model, "num_heads", 3) + + # Calculate feature dimensions + n_node_features_m = back_hrs * n_vars_in + n_node_features_e = back_hrs * n_vars_in if n_grid else 0 + + if model_type == "gnn": + model = MPNN( + n_passing=num_gnn_layers, + lead_hrs=getattr(config.data, "lead_hrs", 48), + n_node_features_m=n_node_features_m, + n_node_features_e=n_node_features_e, + n_out_features=n_vars_out, + hidden_dim=hidden_dim, + ) + elif model_type == "vit": + model = VisionTransformer( + n_stations=n_stations, + madis_len=back_hrs, + madis_n_vars_i=n_vars_in, + madis_n_vars_o=n_vars_out, + dim=hidden_dim, + attn_dim=hidden_dim // 2, + mlp_dim=hidden_dim, + num_heads=num_heads, + num_layers=num_transformer_layers, + ) + else: + raise ValueError(f"Unknown model type: {model_type}") + + # Log model info + n_params = sum(p.numel() for p in model.parameters() if p.requires_grad) + logger.info(f"Created {model_type.upper()} model with {n_params:,} parameters") + + return model + + +def main(): + """Main training function.""" + args = parse_args() + + # Set debug logging + if args.debug: + logging.getLogger().setLevel(logging.DEBUG) + + logger.info("=" * 60) + logger.info("LOAF Weather Forecasting Model Training") + logger.info("=" * 60) + + # Load configuration + config_path = Path(args.config) + if not config_path.exists(): + # Try relative to software directory + config_path = Path(__file__).parent.parent / args.config + if not config_path.exists(): + logger.error(f"Configuration file not found: {args.config}") + sys.exit(1) + + config = get_config(config_path) + logger.info(f"Loaded configuration from: {config_path}") + + # Get data directory + data_dir = Path(args.data_dir) + if not data_dir.is_absolute(): + data_dir = Path(__file__).parent.parent / data_dir + + # Get training parameters (CLI overrides config) + epochs = args.epochs or getattr(config.training, "epochs", 100) + batch_size = args.batch_size or getattr(config.training, "batch_size", 64) + learning_rate = args.lr or getattr(config.training, "learning_rate", 1e-4) + patience = args.patience or getattr(config.training, "patience", 10) + val_split = getattr(config.training, "val_split", 0.15) + + # Get data parameters + back_hrs = getattr(config.data, "back_hrs", 24) + lead_hrs = getattr(config.data, "lead_hrs", 48) + + # Get region bounds + lat_bounds = (config.region.lat_min, config.region.lat_max) + lon_bounds = (config.region.lon_min, config.region.lon_max) + + logger.info(f"Region: {config.region.name}") + logger.info(f" Lat: {lat_bounds[0]:.2f} to {lat_bounds[1]:.2f}") + logger.info(f" Lon: {lon_bounds[0]:.2f} to {lon_bounds[1]:.2f}") + logger.info(f"Training parameters:") + logger.info(f" Epochs: {epochs}") + logger.info(f" Batch size: {batch_size}") + logger.info(f" Learning rate: {learning_rate}") + logger.info(f" Early stopping patience: {patience}") + logger.info(f" Historical window: {back_hrs} hours") + logger.info(f" Forecast horizon: {lead_hrs} hours") + + # Create dataloaders + logger.info("Loading data...") + try: + train_loader, val_loader = create_dataloaders( + data_dir=data_dir, + year=args.year, + back_hrs=back_hrs, + lead_hours=lead_hrs, + batch_size=batch_size, + val_split=val_split, + lat_bounds=lat_bounds, + lon_bounds=lon_bounds, + use_hrrr=True, + use_era5=False, + ) + logger.info(f"Training samples: {len(train_loader.dataset)}") + logger.info(f"Validation samples: {len(val_loader.dataset)}") + except FileNotFoundError as e: + logger.error(f"Data not found: {e}") + logger.error( + "Please ensure you have downloaded the required data using the download scripts." + ) + sys.exit(1) + + # Get dataset info from the first batch + sample_batch = next(iter(train_loader)) + n_stations = sample_batch["station_lon"].shape[1] + n_grid = sample_batch.get("grid_lon", torch.zeros(1, 0)).shape[1] + logger.info(f"Number of stations: {n_stations}") + if n_grid > 0: + logger.info(f"Number of grid points: {n_grid}") + + # Create model + model = create_model( + model_type=args.model, + config=config, + n_stations=n_stations, + back_hrs=back_hrs, + n_vars_in=4, # u, v, temp, dewpoint + n_vars_out=2, # u, v + n_grid=n_grid if n_grid > 0 else None, + ) + + # Create trainer config + checkpoint_dir = Path(args.checkpoint) if args.checkpoint else None + if checkpoint_dir: + checkpoint_dir.mkdir(parents=True, exist_ok=True) + + training_config = TrainingConfig( + back_hrs=back_hrs, + learning_rate=learning_rate, + weight_decay=getattr(config.training, "weight_decay", 1e-4), + epochs=epochs, + batch_size=batch_size, + patience=patience, + max_grad_norm=getattr(config.training, "max_grad_norm", 1.0), + checkpoint_dir=checkpoint_dir, + save_best_only=True, + loss_type="mse", + output_vars=["u", "v"], + device=args.device, + seed=args.seed, + ) + + # Create trainer + trainer = Trainer( + model=model, + config=training_config, + model_type=args.model, + ) + + # Resume from checkpoint if specified + if args.resume: + logger.info(f"Resuming from checkpoint: {args.resume}") + trainer.load_checkpoint(args.resume) + + # Train + logger.info("Starting training...") + history = trainer.train(train_loader, val_loader) + + # Final evaluation + logger.info("=" * 60) + logger.info("Training Complete!") + logger.info("=" * 60) + logger.info(f"Best validation loss: {trainer.state.best_val_loss:.6f}") + + if history["val_metrics"]: + final_metrics = history["val_metrics"][-1] + for name, value in final_metrics.items(): + logger.info(f" {name}: {value:.4f}") + + if checkpoint_dir: + logger.info(f"Checkpoints saved to: {checkpoint_dir}") + + return history + + +if __name__ == "__main__": + main() diff --git a/software/train_log.md b/software/train_log.md new file mode 100644 index 0000000..979482f --- /dev/null +++ b/software/train_log.md @@ -0,0 +1,16 @@ +train log + +2026-02-01 01:29:25 [INFO] __main__: Validation samples: 100 +2026-02-01 01:29:26 [INFO] __main__: Number of stations: 14 +2026-02-01 01:29:26 [INFO] __main__: Number of grid points: 10878 +2026-02-01 01:29:26 [INFO] __main__: Created GNN model with 525,186 parameters +2026-02-01 01:29:26 [INFO] loaf.training.trainer: Using device: cuda +2026-02-01 01:29:26 [INFO] __main__: Starting training... +2026-02-01 01:29:26 [INFO] loaf.training.trainer: Starting training... +2026-02-01 01:29:26 [INFO] loaf.training.trainer: Epochs: 100 +2026-02-01 01:29:26 [INFO] loaf.training.trainer: Batch size: 8 +2026-02-01 01:29:26 [INFO] loaf.training.trainer: Learning rate: 0.0001 +2026-02-01 01:29:26 [INFO] loaf.training.trainer: Device: cuda +2026-02-01 01:29:29 [INFO] loaf.training.trainer: Epoch 1/100 | Train Loss: nan | Wind Err: nan | Val Loss: nan | LR: 1.00e-04 +2026-02-01 01:29:32 [INFO] loaf.training.trainer: Epoch 2/100 | Train Loss: nan | Wind Err: nan | Val Loss: nan | LR: 1.00e-04 +2026-02-01 01:29:35 [INFO] loaf.training.trainer: Epoch 3/100 | Train Loss: nan | Wind Err: nan | Val Loss: nan | LR: 1.00e-04 \ No newline at end of file diff --git a/uv.lock b/uv.lock index 936e804..8f35fe3 100644 --- a/uv.lock +++ b/uv.lock @@ -546,6 +546,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/98/78/01c019cdb5d6498122777c1a43056ebb3ebfeef2076d9d026bfe15583b2b/click-8.3.1-py3-none-any.whl", hash = "sha256:981153a64e25f12d547d3426c367a4857371575ee7ad18df2a6183ab0545b2a6", size = 108274, upload-time = "2025-11-15T20:45:41.139Z" }, ] +[[package]] +name = "cloudpickle" +version = "3.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" }, +] + [[package]] name = "colorama" version = "0.4.6" @@ -684,6 +693,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0b/02/4dbe7568a42e46582248942f54dc64ad094769532adbe21e525e4edf7bc4/cuda_pathfinder-1.3.3-py3-none-any.whl", hash = "sha256:9984b664e404f7c134954a771be8775dfd6180ea1e1aef4a5a37d4be05d9bbb1", size = 27154, upload-time = "2025-12-04T22:35:08.996Z" }, ] +[[package]] +name = "dask" +version = "2026.1.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "cloudpickle" }, + { name = "fsspec" }, + { name = "importlib-metadata", marker = "python_full_version < '3.12'" }, + { name = "packaging" }, + { name = "partd" }, + { name = "pyyaml" }, + { name = "toolz" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/bd/52/b0f9172b22778def907db1ff173249e4eb41f054b46a9c83b1528aaf811f/dask-2026.1.2.tar.gz", hash = "sha256:1136683de2750d98ea792670f7434e6c1cfce90cab2cc2f2495a9e60fd25a4fc", size = 10997838, upload-time = "2026-01-30T21:04:20.54Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e5/23/d39ccc4ed76222db31530b0a7d38876fdb7673e23f838e8d8f0ed4651a4f/dask-2026.1.2-py3-none-any.whl", hash = "sha256:46a0cf3b8d87f78a3d2e6b145aea4418a6d6d606fe6a16c79bd8ca2bb862bc91", size = 1482084, upload-time = "2026-01-30T21:04:18.363Z" }, +] + [[package]] name = "eccodes" version = "2.45.0" @@ -1085,6 +1113,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0e/61/66938bbb5fc52dbdf84594873d5b51fb1f7c7794e9c0f5bd885f30bc507b/idna-3.11-py3-none-any.whl", hash = "sha256:771a87f49d9defaf64091e6e6fe9c18d4833f140bd19464795bc32d966ca37ea", size = 71008, upload-time = "2025-10-12T14:55:18.883Z" }, ] +[[package]] +name = "importlib-metadata" +version = "8.7.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "zipp", marker = "python_full_version < '3.12'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f3/49/3b30cad09e7771a4982d9975a8cbf64f00d4a1ececb53297f1d9a7be1b10/importlib_metadata-8.7.1.tar.gz", hash = "sha256:49fef1ae6440c182052f407c8d34a68f72efc36db9ca90dc0113398f2fdde8bb", size = 57107, upload-time = "2025-12-21T10:00:19.278Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fa/5e/f8e9a1d23b9c20a551a8a02ea3637b4642e22c2626e3a13a9a29cdea99eb/importlib_metadata-8.7.1-py3-none-any.whl", hash = "sha256:5a1f80bf1daa489495071efbb095d75a634cf28a8bc299581244063b53176151", size = 27865, upload-time = "2025-12-21T10:00:18.329Z" }, +] + [[package]] name = "iniconfig" version = "2.3.0" @@ -1204,6 +1244,7 @@ source = { editable = "." } dependencies = [ { name = "cdsapi" }, { name = "cfgrib" }, + { name = "dask" }, { name = "flask" }, { name = "herbie-data", version = "2025.12.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "herbie-data", version = "2026.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, @@ -1248,6 +1289,7 @@ requires-dist = [ { name = "black", marker = "extra == 'dev'", specifier = ">=23.0.0" }, { name = "cdsapi", specifier = ">=0.6.0" }, { name = "cfgrib", specifier = ">=0.9.0" }, + { name = "dask", specifier = ">=2026.1.2" }, { name = "flask", specifier = ">=3.0.0" }, { name = "herbie-data", specifier = ">=2024.0.0" }, { name = "mkdocs-material", marker = "extra == 'docs'", specifier = ">=9.7.1" }, @@ -1273,6 +1315,15 @@ requires-dist = [ ] provides-extras = ["dev", "docs"] +[[package]] +name = "locket" +version = "1.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2f/83/97b29fe05cb6ae28d2dbd30b81e2e402a3eed5f460c26e9eaa5895ceacf5/locket-1.0.0.tar.gz", hash = "sha256:5c0d4c052a8bbbf750e056a8e65ccd309086f4f0f18a2eac306a8dfa4112a632", size = 4350, upload-time = "2022-04-20T22:04:44.312Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/db/bc/83e112abc66cd466c6b83f99118035867cecd41802f8d044638aa78a106e/locket-1.0.0-py2.py3-none-any.whl", hash = "sha256:b6c819a722f7b6bd955b80781788e4a66a55628b858d347536b7e81325a3a5e3", size = 4398, upload-time = "2022-04-20T22:04:42.23Z" }, +] + [[package]] name = "markdown" version = "3.10" @@ -2289,6 +2340,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e6/3f/a80ac00acbc6b35166b42850e98a4f466e2c0d9c64054161ba9620f95680/pandas-3.0.0-cp314-cp314t-win_arm64.whl", hash = "sha256:1c39eab3ad38f2d7a249095f0a3d8f8c22cc0f847e98ccf5bbe732b272e2d9fa", size = 9441003, upload-time = "2026-01-21T15:52:02.281Z" }, ] +[[package]] +name = "partd" +version = "1.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "locket" }, + { name = "toolz" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b2/3a/3f06f34820a31257ddcabdfafc2672c5816be79c7e353b02c1f318daa7d4/partd-1.4.2.tar.gz", hash = "sha256:d022c33afbdc8405c226621b015e8067888173d85f7f5ecebb3cafed9a20f02c", size = 21029, upload-time = "2024-05-06T19:51:41.945Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/e7/40fb618334dcdf7c5a316c0e7343c5cd82d3d866edc100d98e29bc945ecd/partd-1.4.2-py3-none-any.whl", hash = "sha256:978e4ac767ec4ba5b86c6eaa52e5a2a3bc748a2ca839e8cc798f1cc6ce6efb0f", size = 18905, upload-time = "2024-05-06T19:51:39.271Z" }, +] + [[package]] name = "pathspec" version = "1.0.3" @@ -3222,6 +3286,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/23/d1/136eb2cb77520a31e1f64cbae9d33ec6df0d78bdf4160398e86eec8a8754/tomli-2.4.0-py3-none-any.whl", hash = "sha256:1f776e7d669ebceb01dee46484485f43a4048746235e683bcdffacdf1fb4785a", size = 14477, upload-time = "2026-01-11T11:22:37.446Z" }, ] +[[package]] +name = "toolz" +version = "1.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/11/d6/114b492226588d6ff54579d95847662fc69196bdeec318eb45393b24c192/toolz-1.1.0.tar.gz", hash = "sha256:27a5c770d068c110d9ed9323f24f1543e83b2f300a687b7891c1a6d56b697b5b", size = 52613, upload-time = "2025-10-17T04:03:21.661Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/12/5911ae3eeec47800503a238d971e51722ccea5feb8569b735184d5fcdbc0/toolz-1.1.0-py3-none-any.whl", hash = "sha256:15ccc861ac51c53696de0a5d6d4607f99c210739caf987b5d2054f3efed429d8", size = 58093, upload-time = "2025-10-17T04:03:20.435Z" }, +] + [[package]] name = "torch" version = "2.10.0" @@ -3691,3 +3764,12 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/48/b7/503c98092fb3b344a179579f55814b613c1fbb1c23b3ec14a7b008a66a6e/yarl-1.22.0-cp314-cp314t-win_arm64.whl", hash = "sha256:9f6d73c1436b934e3f01df1e1b21ff765cd1d28c77dfb9ace207f746d4610ee1", size = 85171, upload-time = "2025-10-06T14:12:16.935Z" }, { url = "https://files.pythonhosted.org/packages/73/ae/b48f95715333080afb75a4504487cbe142cae1268afc482d06692d605ae6/yarl-1.22.0-py3-none-any.whl", hash = "sha256:1380560bdba02b6b6c90de54133c81c9f2a453dee9912fe58c1dcced1edb7cff", size = 46814, upload-time = "2025-10-06T14:12:53.872Z" }, ] + +[[package]] +name = "zipp" +version = "3.23.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e3/02/0f2892c661036d50ede074e376733dca2ae7c6eb617489437771209d4180/zipp-3.23.0.tar.gz", hash = "sha256:a07157588a12518c9d4034df3fbbee09c814741a33ff63c05fa29d26a2404166", size = 25547, upload-time = "2025-06-08T17:06:39.4Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2e/54/647ade08bf0db230bfea292f893923872fd20be6ac6f53b2b936ba839d75/zipp-3.23.0-py3-none-any.whl", hash = "sha256:071652d6115ed432f5ce1d34c336c0adfd6a884660d1e9712a256d3d3bd4b14e", size = 10276, upload-time = "2025-06-08T17:06:38.034Z" }, +]