Add Exponential Moving Average (EMA) of Model Weights via PyTorch Lightning Callback

Problem

neural-lam uses raw optimizer weights for checkpointing and evaluation. There is no Exponential Moving Average (EMA) of model weights — a technique standard in every major neural weather prediction system. EMA maintains a shadow copy of parameters as a running average (θ_ema ← α·θ_ema + (1−α)·θ_current), used for validation/test/inference while training continues on raw weights. This is especially critical for autoregressive rollout, where per-step noise compounds across forecast horizons. Currently:

• configure_optimizers() returns bare AdamW — no weight averaging (ar_model.py:L201)
• Zero pl.Callback subclasses exist in the codebase
• ModelCheckpoint saves raw (noisy) weights (train_model.py:L310)

Why It Matters

• Forecast quality: GraphCast reports 5–12% RMSE reduction with EMA vs. raw weights over 10-day rollouts
• Checkpoint reliability: EMA gives lower-variance val_mean_loss → ModelCheckpoint selects more stable snapshots
• Best-practice gap: neural-lam is the only major open-source neural weather framework without EMA
• Reproducibility: Published benchmarks report EMA-evaluated metrics; without EMA, scores can't be matched

Proposed Implementation

1. New file neural_lam/callbacks.py — EMACallback(pl.Callback):

• on_fit_start: clone model params as initial EMA weights
• on_train_batch_end: update via torch.lerp_ (efficient in-place)
• on_validation_start / on_test_start: swap EMA weights in
• on_validation_end / on_test_end: swap original weights back
• on_save_checkpoint /
• on_load_checkpoint: persist EMA state

2. Modify :

• train_model.py:Add --ema_decay CLI arg (default None = disabled, e.g. 0.999)
• Append EMACallback to trainer callbacks list when enabled

3. Tests (tests/test_ema_callback.py):

• Mathematical correctness of running average after N steps
• Weights swap correctly during validation vs. training
• EMA state survives checkpoint save/load
• No behavior change when --ema_decay is not set

Dependencies: None beyond existing PyTorch + PyTorch Lightning