feat: add probabilistic forecaster interface (decoupled from Graph-EFM)#678
Closed
Sir-Sloth-The-Lazy wants to merge 24 commits into
Closed
feat: add probabilistic forecaster interface (decoupled from Graph-EFM)#678Sir-Sloth-The-Lazy wants to merge 24 commits into
Sir-Sloth-The-Lazy wants to merge 24 commits into
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>
Introduce the latent-variable forecasting interface as three layers with a clean boundary: the predictor produces distributions only, the forecaster owns the rollout and the KL between those distributions, and the module assembles the ELBO. - Add LatentStepPredictor (step_predictors/probabilistic.py): abstract step_distributions -> (prior, posterior, pred_mean, pred_std); forward is the prior-sampling special case. Likelihood/KL/ELBO live outside it. - Add ProbabilisticARForecaster: training_rollout (posterior-conditioned, reduces per-step KL) and sample_trajectories (prior ensemble). - Add ProbabilisticForecasterModule: assembles the beta-weighted negative ELBO in training_step, substituting the constant per_var_std fallback. - Add crps_ens metric for ensemble CRPS. - Add interface tests driven by a graph-free dummy LatentStepPredictor. Graph-EFM is not yet wired onto this interface; that lands in a follow-up.
Sir-Sloth-The-Lazy
force-pushed
the
feat/probabilistic-forecaster-module
branch
from
77e68cf to
2cf6150
Compare
Contributor
Author
|
2cf6150 This is the new addition |
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
Introduces the latent-variable forecasting interface as three layers with a single clean boundary the predictor produces distributions only, the forecaster owns the rollout + the KL between those distributions, and the module assembles the ELBO. This is the interface-only step; Graph-EFM is intentionally not wired onto it yet (that lands in a follow-up), so the contract stays focused.
Files belonging to this PR:
neural_lam/models/step_predictors/probabilistic.py:LatentStepPredictor: abstractstep_distributions(...) -> (prior_dist, posterior_dist, pred_mean, pred_std);forwardis the prior-sampling special case. The predictor computes no likelihood, no KL, noper_var_std.targetis optional and used for one thing only: conditioning the variational posterior.neural_lam/models/forecasters/probabilistic.py:ProbabilisticARForecaster:training_rollout(posterior-conditioned rollout, reduces the per-step KL between the model's own prior/posterior) andsample_trajectories(prior ensemble for CRPS / ensemble eval).neural_lam/models/probabilistic_module.py:ProbabilisticForecasterModule: assembles the β-weighted negative ELBO intraining_step, substituting the constantper_var_stdfallback (owned by the baseForecasterModule) when the predictor emits no std.neural_lam/metrics.py:crps_ensmetric (ensemble CRPS).tests/test_probabilistic_interface.py: interface tests driven by a graph-free dummyLatentStepPredictor, verifying the contract independently of any concrete predictor.neural_lam/models/__init__.py: export the new classes.Motivation: an earlier review flagged that the step predictor computing likelihood / KL /
per_var_stdfelt wrong, and that the methods were hard to judge without seeing how the final optimized loss is derived. This PR answers that by defining the loss-computation interface explicitly: likelihood andper_var_stdmove to the module, the full ELBO derivation is visible in one place (training_step), andtargetreaches the predictor only to form the posterior.Open design question for review: KL is currently reduced in the forecaster (between the model's own two distributions), with the module owning only the
kl_betaweighting and ELBO assembly. The alternative is to surface the raw per-step distributions up to the module and reduce KL there, making the module the single sink for the whole loss. Feedback welcome on which boundary you prefer.Issue Link
Part of the Graph-EFM port (#62) stack (follows #648). No standalone issue.
Type of change
Checklist before requesting a review
Author checklist after completed review