Generalizing to Probabilistic Forecasting Models
The existing Neural-lam probabilistic forecasting model, GraphEFM, currently exists in an isolated branch (prob_model_lam) with a monolithic ARModel class encapsulating both the training, rollout, and single-step prediction functionality in a single file. Meanwhile, PR #208 refactors the deterministic models into a clean three-layer hierarchy: ForecasterModule (Lightning training harness), Forecaster (rollout strategy), and StepPredictor (single-step prediction). The GraphEFM model cannot be merged into main until it is ported to this new hierarchy. This project ports the GraphEFM model into the new hierarchy by refactoring it into three new classes that mirror the existing separation of concerns: GraphEFMPredictor (extends StepPredictor : owns the encoder, prior, decoder, latent variable logic), EnsembleForecaster (extends Forecaster : owns the AR rollout with separate prior and encoder paths), and EnsembleForecasterModule (extends ForecasterModule : owns the ELBO training, ensemble metrics, and visualization). Key design decisions include using explicit forward() and forward_encoder_path() instead of gating on self.training (handled automatically by PyTorch Lightning), making the predictor self-contained with its own per_var_std buffer, and replacing all hardcoded constants with the datastore API introduced in PR #208. Deliverables: - All components of the encoder/decoder moved, GraphEFM predictor that fulfills the contract of StepPredictor, EnsembleForecaster with both paths tested, and end-to-end graph efM training of --model graph_efm on MEPS, along with ensemble metrics. - KL/CRPS warmup schedules, ensemble parameters defined via configuration, Zarr-based ensemble prediction saving, additional ensemble visualizations, spread maps, and spaghetti plot, checkpoint backward compatibility, and architecture documentation.
Project details
Technologies
Not listed in the archive