feat: Graph-EFM ensemble forecasting model - #712
Draft
Sir-Sloth-The-Lazy wants to merge 56 commits into
Draft
Conversation
Adds neural_lam/models/latent/ with the encoder and decoder submodules needed by the probabilistic GraphEFM model (issue mllam#62). Ported from the prob_model_lam branch with adaptations for the current main architecture: - constants.GRID_STATE_DIM replaced by a num_state_vars constructor arg - interaction_net imports updated to neural_lam.gnn_layers - GraphLatentDecoder.processor unified with the other four GNN-seq constructions to use utils.make_gnn_seq (handles processor_layers=0) - HiGraph{Encoder,Decoder} guard against single-level meshes where the latent variable would be silently ignored - ConstantLatentEncoder docstring documents the N(1,1) vs N(0,1) discrepancy with the prob_model_lam CLI help (open question upstream) Also adds to neural_lam/utils.py: - IdentityModule: pass-through nn.Module for multi-arg sequential GNNs - make_gnn_seq: builds a pyg.nn.Sequential of InteractionNets, or an IdentityModule when num_gnn_layers=0; lazy-imports gnn_layers to avoid the existing gnn_layers -> utils circular dependency 17 tests in tests/test_latent_modules.py cover output shapes, distribution properties, backprop to every parameter, 2- and 3-level hierarchical graphs, intra_level_layers=0, and the single-level guard.
Co-authored-by: Joel Oskarsson <joel.oskarsson@outlook.com>
Make GNN types configurable and tidy up the latent modules per PR review: - make_gnn_seq: accept a gnn_type arg (resolved via get_gnn_class) so it is not limited to InteractionNet, and make it strict (raise on num_gnn_layers < 1) instead of silently returning an IdentityModule; callers now own the no-op (identity) case explicitly. - graph/hi encoders and decoders: expose g2m/m2g/mesh_up/mesh_down gnn_type parameters wired to get_gnn_class, with defaults matching prob_model_lam. - graph encoder/decoder: rename processor_layers -> m2m_layers (and the self.processor attribute -> self.m2m_gnns); "processor" was misleading in an encoder/decoder context. - ConstantLatentEncoder: return zeros instead of ones so the static prior is mean 0 (fixes the prob_model_lam mean-1 bug; matches its own CLI help). - tests: update for the renamed arg and strict make_gnn_seq, add coverage for the flat zero-m2m identity path, and assert the constant prior is N(0, 1).
Port prob_model_lam's GraphEFM single-step half onto the StepPredictor interface, reusing the latent encoder/decoder infra. The predictor owns its conditional prior, variational encoder, and latent decoder, plus the per-step ELBO pieces (compute_step_loss) and sampling helpers; rollout, ELBO assembly, ensemble logic, and logging stay outside it. - forward is source's predict_step (prior rsample -> decode -> sampled next state); no rescaling/clamping - loss_fn and interior_mask are threaded parameters, not predictor state; compute_step_loss takes compute_kl (kl_term=None when off) - per_var_std mirrors ForecasterModule's formula, hence the config arg - one class for flat + hierarchical meshes, resolved from self.hierarchical - not registered in MODELS yet (needs config / no mesh_aggr); config-aware assembly deferred to the ensemble-forecaster PR Adds tests/test_graph_efm_predictor.py covering forward shapes, output_std, compute_step_loss + KL toggle, differentiability, member stochasticity, sample_obs_noise, and the per_var_std formula (flat + hierarchical).
…s for the rest Per review discussion: the architecturally constrained edge sets in the hierarchical latent modules get fixed GNN types instead of parameters: - HiGraphLatentEncoder mesh-up: PropagationNet (must push grid info up into the latent readout) - HiGraphLatentDecoder mesh-up: InteractionNet (PropagationNet residual would bypass Z at the top level, leaving it unused at initialization) - HiGraphLatentDecoder mesh-down: PropagationNet (must push Z down the hierarchy to reach the grid output) All remaining choices (g2m/m2g) stay configurable and default to InteractionNet for consistency with the rest of the codebase. GraphEFM now accepts g2m_gnn_type/m2g_gnn_type and passes them through to the prior, encoder and decoder, ready for wiring to the existing argparse flags.
Upstream main added an interrogate pre-commit hook requiring 100% docstring coverage, which failed on this branch's CI after merging. - Remove the branch's pre-reorganization duplicates (forecaster.py, ar_forecaster.py, step_predictor.py, forecaster_module.py); main carries the same code under models/forecasters/, models/ step_predictors/ and models/module.py, and all imports already go through the new layout. - Add the missing module and __init__ docstrings (numpy style) in the latent modules, GraphEFM and utils.IdentityModule.forward.
Remove references to the original prob_model_lam implementation and other work meta-information from docstrings and comments, per review. Docstrings now describe what each class/function does; usage context is left to call sites.
Add proper Parameters/Returns sections following the numpydoc convention, per review.
…dentityModule When m2m_layers / intra_level_layers is 0, the latent modules now set the corresponding GNN attribute to None and skip the update in the forward pass, instead of routing representations through a no-op IdentityModule. This makes it clear from the forward code that no processing happens in that case. IdentityModule is removed from utils. The hierarchical up/down loops index levels explicitly to accommodate the conditional; outputs are unchanged (verified bit-identical against the previous implementation).
…base class Use the base class summary and expand on it with the constant-specific behavior, per review.
Co-authored-by: Joel Oskarsson <joel.oskarsson@outlook.com>
Describe the role of each representation in the message passing (sender/receiver, where the latent enters, purpose of the residual grid rep) in BaseGraphLatentDecoder and the inheriting decoders, per review.
Replace the single GraphEFM class that resolved flat vs hierarchical at construction with an explicit class per graph type, per review: - BaseGraphEFM: graph-type independent setup (graph loading, grid and grid-mesh edge embedders, per-variable std) and all shared behavior (forward, compute_step_loss, estimate_likelihood, sampling helpers). Validates the loaded graph against the subclass's requires_hierarchical and exposes an embedd_mesh hook used by embedd_all. - GraphEFM: hierarchical mesh graphs; builds per-level mesh embedders and HiGraphLatentEncoder/HiGraphLatentDecoder modules. - GraphEFMMS: flat (e.g. multi-scale) mesh graphs; builds the flat mesh embedders and GraphLatentEncoder/GraphLatentDecoder modules. learn_prior remains a constructor flag on both subclasses. Tests select the class per graph type and cover the graph-type mismatch error.
Rename the layer-count parameters to say what the layers are, matching the latent module parameter names: prior/encoder/decoder_intra_level_ layers on GraphEFM (hierarchical) and prior/encoder/decoder_m2m_layers on GraphEFMMS (flat), per review.
Make the GraphEFM and GraphEFMMS __init__ docstrings self-contained with the full parameter list, instead of pointing at the base class for the shared ones, per review.
Sampling uncorrelated Gaussian observation noise per grid node is not useful in practice, so the option is removed entirely rather than left to tempt users, per review. forward now returns the decoder mean directly (the prediction is stochastic only through the latent sample) and the trivial sample_next_state helper is dropped.
BaseGraphModel and BaseGraphEFM duplicated their graph loading, buffer registration and grid-input-dim computation. Factor these into two utils helpers used by both: - utils.load_and_register_graph(module, datastore, graph_name): loads the graph and registers its tensors/BufferLists on the module, returning whether it is hierarchical. - utils.grid_input_dim(datastore, grid_static_dim, num_past_forcing_ steps, num_future_forcing_steps): the total grid input dimensionality. This keeps the two model families' grid-feature setup in one place (e.g. for a future boundary-forcing input) without coupling their differing forward passes or submodule sets via inheritance.
Rename embedd_all -> embedd_grid_and_graph (embeds the grid for states up to t-1 plus the full graph) and embedd_current -> embedd_grid_with_target (embeds the grid including the target state, for the encoder), per review.
In forward, prev_state is already X_t, so pass it directly to the decoder instead of aliasing it to last_state, per review.
Co-authored-by: Joel Oskarsson <joel.oskarsson@outlook.com>
Move the constant-prior branch (identical across subclasses) into BaseGraphEFM.build_prior and delegate the graph-specific learnable prior to build_learnable_prior, which GraphEFM and GraphEFMMS implement.
Spell out the multi-scale variant's name instead of the EFMMS abbreviation, keeping the GraphEFM prefix shared with the hierarchical variant.
…_dim The static feature count is already available from the datastore via get_num_data_vars(category="static"), so drop the redundant grid_static_dim argument and query it inside the function.
get_num_data_vars("static") can report a nonzero count even when the
datastore provides no static dataarray, in which case the grid static
buffer is empty. Mirror the buffer construction by treating a None static
dataarray as zero static features, fixing the no-static-features case.
Add abstract Forecaster.compute_training_loss returning a finished (loss, loss_components) pair, so each forecaster owns its complete training objective. ForecasterModule.training_step now only injects the configured scoring rule, interior mask and per_var_std, and logs the result. The deterministic ARForecaster loss is unchanged in value. Add the abstract ProbabilisticForecaster (sample_ensemble capability), ProbabilisticARForecaster (sequential sampled rollouts, trains on the configured score of the ensemble mean) and a minimal ProbabilisticForecasterModule whose validation samples an ensemble and logs the RMSE of the ensemble mean. Interface design from mllam#685.
score_metric/per_var_std were injected into compute_training_loss by ForecasterModule and also used directly for val/test loss reporting, duplicating config the forecaster already needs for its own objective. ARForecaster/ProbabilisticARForecaster now own self.loss and self.per_var_std (computed from an optional config ctor arg), and ForecasterModule reads them off self.forecaster instead. Also trims the CHANGELOG entry for mllam#685 down to one sentence per review feedback.
A Forecaster built without config now silently has per_var_std=None when its predictor doesn't output its own std. Previously per_var_std was always computed by ForecasterModule itself, so this gap didn't exist; now that construction is split across two calls, catch it at ForecasterModule init instead of crashing at the first val/test step.
Each member's predicted std is its own, not a spread computed across the ensemble, so ensemble_std was a misleading name. Document on ProbabilisticForecaster that a per-member std makes the predictive distribution a mixture of Gaussians, and note in ProbabilisticARForecaster.compute_training_loss that averaging the per-member stds is a simplification of the true mixture variance (which also includes the spread between member means).
Co-authored-by: Joel Oskarsson <joel.oskarsson@outlook.com>
Scoring the ensemble mean with a pointwise metric only rewards the mean being right, giving the model no incentive to keep a calibrated spread, and risks training it to collapse to a point estimate. Redeclare compute_training_loss as abstract on ProbabilisticARForecaster instead of providing that as a default (it would otherwise silently fall back to ARForecaster's single-rollout objective via MRO, not even the ensemble mean). Concrete subclasses must define their own objective. Tests that only need an instantiable forecaster now use a local ConcreteProbabilisticARForecaster example (ensemble-mean scoring, moved out of the library code); a new test locks in that the base class itself cannot be instantiated.
Drop ProbabilisticARForecaster's ensemble_size constructor arg and the implicit num_members=None -> self.ensemble_size fallback in sample_ensemble; num_members is now always required. Baking a default member count into the forecaster's state was unnecessary now that compute_training_loss is abstract too (nothing in the shared base class path used it) and just adds an implicit default callers could silently rely on instead of deciding explicitly. The num_members < 1 validation moves from __init__ to sample_ensemble accordingly. ProbabilisticForecasterModule.eval_ensemble_size follows suit: it no longer defaults to None with a forecaster fallback, it's required. Test-only ConcreteProbabilisticARForecaster (used wherever a concrete probabilistic forecaster is needed for testing) gains its own train_num_members for the training objective, since deciding how many members to sample during training is now the concrete subclass's call.
Mirrors validation_step: samples eval_ensemble_size members and scores the ensemble mean, same as validation. Factored the shared sampling + scoring + logging into _ensemble_step(batch, phase) rather than duplicating the block, since validation_step and test_step differ only in their log-key prefix and which metrics dict collects the result. Overrides on_test_epoch_end (rather than inheriting ForecasterModule's) since this module's test_step doesn't populate spatial_loss_maps or plot examples - the inherited version would crash on torch.cat of an empty list.
Rename the ensemble-mean diagnostic keys from *_loss_unroll/*_mean_loss to *_ens_rmse_unroll/*_mean_ens_rmse so they aren't conflated with the training loss, per review feedback.
Remove explanatory comments around the per_var_std overrides; the assignments are clear on their own.
…crete deterministic/probabilistic modules Introduce BaseForecasterModule (abstract) under models/forecasters/ holding shared plumbing (training_step, common_step, batch standardization, checkpoint compatibility, plotting/aggregation helpers), with validation_step, test_step and on_test_epoch_end left abstract since they differ meaningfully between evaluation modes. Rename ForecasterModule to DeterministicForecasterModule and move it, alongside ProbabilisticForecasterModule, into forecasters/ as siblings implementing the shared contract, rather than one subclassing the other.
Note that g2m/m2g flags apply to Graph-EFM too, while mesh_up/mesh_down flags only affect Hi-LAM since Graph-EFM hard-codes those GNN types.
Move prior construction fully into the base constructor instead of each subclass calling self.build_prior() itself: BaseGraphEFM.__init__ now takes latent_dim/learn_prior/prior_dist/prior_layers/g2m_gnn_type, validates the loaded graph against a new expects_hierarchical class attribute, derives num_mesh_nodes generically, and builds prior_model directly. GraphEFM and GraphEFMMultiScale forward these to super() and drop their duplicated graph-type check, num_mesh_nodes assignment and build_prior call, reusing self.latent_dim instead.
Removes estimate_likelihood/compute_step_loss and the per_var_std buffer they existed to feed, plus the now-unused config constructor arg (its only use was building per_var_std). Per the objective now living on the Forecaster (mllam#700), the predictor stays a pure network construct: encoder/prior/decoder plus the forward sampling path.
Add JetBrains IDE project directory to the ignore list alongside the existing .vim/.vscode entries.
…aster DeterministicForecasterModule.validation_step/test_step read self.forecaster.loss and self.forecaster.per_var_std directly, so the Module still knew how to compute a loss from a prediction. Add an abstract Forecaster.score (implemented on ARForecaster) that resolves the pred_std fallback and applies a scoring rule internally; the Module now only calls forecaster.score(...) and never touches loss/per_var_std itself.
…ster BaseForecasterModule.__init__ raised if forecaster.per_var_std was None and the forecaster didn't predict its own std -- a check on the Forecaster's configuration living in the wrong class. Move it into ARForecaster via a new _resolve_pred_std helper, shared by score() and compute_training_loss(): construction with config=None now always succeeds (a valid state for forecasters only ever used for inference), and the ValueError instead fires from the Forecaster itself, only once scoring is attempted without any std to use.
…dules/ BaseForecasterModule, DeterministicForecasterModule and ProbabilisticForecasterModule (Lightning wrappers around a Forecaster) were mixed in with the Forecaster classes themselves under neural_lam/models/forecasters/. Move them into their own neural_lam/models/modules/ package (base.py, deterministic.py, probabilistic.py) so the two concerns: what a forecaster is vs how it's trained/evaluated by Lightning live in separate directories. Pure move: internal imports and neural_lam/models/__init__.py updated accordingly, no behavioural change.
Resolves the conflict on neural_lam/models/module.py: main's PR mllam#675 added --train_steps_to_log and deduplicated the val/test/train prediction+loss+logging pattern into _compute_prediction_and_loss / _log_step_loss, on top of the pre-mllam#700 module.py where the Module still owned self.loss/self.per_var_std directly. This branch had since deleted that file (split into modules/base.py, deterministic.py, probabilistic.py, with loss ownership moved onto the Forecaster). Ports mllam#675's dedup/logging helpers (_warn_skipped_steps, _log_step_loss, the train_steps_to_log hparam) into modules/base.py, and rewires DeterministicForecasterModule's validation_step/test_step to use them via forecaster.score() instead of the removed self.loss/self.per_var_std. _compute_prediction_and_loss is dropped rather than ported: it assumes a single per-step-decomposable loss, which conflicts with compute_training_loss's contract that a forecaster's training objective may not decompose per step (e.g. an ELBO). training_step therefore still only logs the aggregate train_loss; train_steps_to_log is accepted for CLI/checkpoint compatibility but does not yet produce a train_loss_unroll{i} breakdown -- documented inline as a deliberate scope boundary, not an oversight. Also fixes two call sites (deterministic.py, probabilistic.py) that referenced the pre-rename _warn_skipped_val_steps, and absorbs the unrelated utils.py -> utils/ package split from mllam#682 (auto-merged cleanly, this branch never touched utils.py).
Resolves the modify/delete conflict on neural_lam/utils.py: main split it into the utils/ package (mllam#682), while this branch had added load_and_register_graph, compute_grid_input_dim and make_gnn_seq on top of the monolithic file. Ports the three into utils/graph.py (the first two) and utils/networks.py (make_gnn_seq), re-exported from utils/__init__.py. Also fixes a latent bug uncovered while porting: load_and_register_graph never gained the mesh_node_features_scaling parameter that main's mllam#323 added to load_graph, so every caller was either broken (graph_efm.py, missing arg -> TypeError) or working around it by loading the graph a second time by hand (step_predictors/graph/base.py, leaving dead/duplicate code and a double self.hierarchical assignment). Both callers now compute grid_xy_max_span and pass it through load_and_register_graph once.
…asses Three changes to BaseGraphEFM per review: - Inline build_prior's body into __init__ (it was only ever called once, from there) and remove the method. - Remove the expects_hierarchical class attribute and the generic hierarchical-vs-flat branch it drove in the base class; replace with an abstract latent_spatial_dim property that each subclass implements from its own knowledge of its graph shape (top mesh level for GraphEFM, all mesh nodes for GraphEFMMultiScale), used to size the constant prior. - Move the graph-type ValueError back out of the base class into the subclasses via a new check_graph_type hook, so the "hierarchical or flat" concern lives only in GraphEFM/GraphEFMMultiScale, not the base class. check_graph_type is a method the base class calls (not code left for the subclass constructor to run after super().__init__() returns): since the base class now builds the prior -- calling build_learnable_prior, which assumes edge_index tensors of the shape the subclass expects -- the check has to happen before that runs, not after. Doing it the latter way was tried first and left test_graph_type_mismatch_raises[GraphEFMMultiScale- hierarchical] hitting a RuntimeError deep in gnn_layers.py's InteractionNet (BufferList indexed as if it were a single edge_index tensor) instead of the intended ValueError, because build_learnable_prior ran during super().__init__(), before the subclass's post-super() check ever got a chance to fire.
…Datastore utils/graph.py's TYPE_CHECKING-guarded BaseDatastore import predates the utils/ package split (mllam#682); this checks whether it's still needed now that it's merged. It still was, but not for the reason the guard's comment implied. The cycle isn't the monolithic-vs-package structure -- it's datastore/mdp.py importing log_on_rank_zero from the utils *package* (`from ..utils import log_on_rank_zero`), which only resolves once utils/__init__.py has fully run. utils/__init__.py imports .graph before .logging, so an eager BaseDatastore import in graph.py (which pulls in the datastore package, which imports mdp.py) hits log_on_rank_zero before it's bound. Fix: import log_on_rank_zero from the utils.logging *submodule* directly in mdp.py, not the package. That resolves independently of utils/__init__.py's progress, breaking the cycle without relying on import order (unlike reordering utils/__init__.py, which also works but is fragile and silently reintroducible). The TYPE_CHECKING guard is no longer needed.
…nfig/CLI) Add GraphEFMForecaster, which trains the GraphEFM/GraphEFMMultiScale latent step predictors via their ELBO through the predictors' step_distributions interface, with the static graph embedding hoisted out of the rollout. Add a probabilistic config section (ProbabilisticConfig) and config-aware train_model wiring registering graph_efm (hierarchical) and graph_efm_ms (flat). Move the decoder residual into the predictor so the predicted mean is clamped to the configured output range like the deterministic models, and use an empty m2m placeholder when no intra-level layers are configured.
Sir-Sloth-The-Lazy
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Describe your changes
Instantiates the full Graph-EFM ensemble forecasting model end to end by combining two in-flight PRs and adding the missing forecaster:
GraphEFM(hierarchical) /GraphEFMMultiScale(flat) step predictors.ProbabilisticForecaster,ProbabilisticARForecaster,ProbabilisticForecasterModule).GraphEFMForecaster, aProbabilisticARForecasterthat trains the Graph-EFM step predictors via their ELBO (reconstruction likelihood +kl_beta-weighted KL), driving each step through the predictors'step_distributionsinterface with the static graph embedding hoisted out of the rollout. Plus aprobabilisticconfig section (ProbabilisticConfig) holding the Graph-EFM hyperparameters, and config-aware assembly intrain_model.pyregisteringgraph_efm(hierarchical) andgraph_efm_ms(flat).Both variants train end to end and are evaluated as ensembles (
ProbabilisticForecasterModule).Draft / WIP : not for merge as-is. This branch is stacked on top of #648 and #700, so the diff below includes their commits; it is opened for visibility of the full model coming together. The intent is to land it as reviewable pieces: the predictor-side review fixes (clamping the predicted mean, dropping the raw-features m2m placeholder) are already on #648, and the forecaster/config/CLI will follow as its own PR on top of #700.
Dependencies: requires #648 and #700 to be merged first.
Issue Link
Part of #685 and #62 (a general probabilistic forecasting interface). Depends on #648 and #700.
Type of change
Checklist before requesting a review
pullwith--rebaseoption if possible).Checklist for reviewers
Each PR comes with its own improvements and flaws. The reviewer should check the following:
Author checklist after completed review
reflecting type of change (add section where missing):
Checklist for assignee