Add training scripts

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)

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]

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

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."""

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)

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

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",
+]