feat: add latent encoder/decoder infrastructure for Graph-EFM port by Sir-Sloth-The-Lazy · Pull Request #648 · mllam/neural-lam

Sir-Sloth-The-Lazy · 2026-05-27T23:59:50Z

Describe your changes

Adds neural_lam/models/latent/ the encoder and decoder submodules that the probabilistic Graph-EFM model needs. This is infrastructure-only: no model uses these classes yet. They are consumed by the upcoming GraphEFMPredictor (StepPredictor subclass) which will close #62.

New modules in neural_lam/models/latent/:

File	What it adds
`base_encoder.py`	`BaseLatentEncoder` : abstract base; handles isotropic / diagonal Gaussian output
`base_decoder.py`	`BaseGraphLatentDecoder` : abstract base; residual grid MLP + latent embedder + param map
`constant_encoder.py`	`ConstantLatentEncoder` : input-independent prior (used when `learn_prior=False`)
`graph_encoder.py`	`GraphLatentEncoder` : flat graph: grid → mesh via PropagationNet + InteractionNet stack
`graph_decoder.py`	`GraphLatentDecoder` : flat graph: grid + latent → grid via g2m / processor / m2g
`hi_graph_encoder.py`	`HiGraphLatentEncoder` : hierarchical mesh: propagates up to top level, reads out latent dist
`hi_graph_decoder.py`	`HiGraphLatentDecoder` : hierarchical mesh: up + latent fusion + down pass back to grid

Adaptations from prob_model_lam for the current main architecture:

constants.GRID_STATE_DIM (removed) → num_state_vars constructor arg on all decoders
from neural_lam.interaction_net import ... → from neural_lam.gnn_layers import ...
GraphLatentDecoder.processor unified with the other four GNN-seq constructions to use utils.make_gnn_seq, which also handles processor_layers=0 gracefully
HiGraph{Encoder,Decoder} raise ValueError for single-level meshes (where the latent would be silently ignored); points users to the flat variants

Also adds to neural_lam/utils.py:

IdentityModule pass-through nn.Module for multi-arg pyg.nn.Sequential pipelines
make_gnn_seq builds a pyg.nn.Sequential of InteractionNet layers, or IdentityModule when num_gnn_layers=0; lazy-imports gnn_layers to avoid the existing gnn_layers → utils circular dependency

Open question flagged in constant_encoder.py docstring: the static prior returns Normal(mean=1, std=1) faithful to prob_model_lam but the --learn_prior CLI help on that branch describes it as "mean 0". One of the two is wrong; will raise it separately with @joeloskarsson.

Dependencies:

PropagationNet and InteractionNet are already on main (Adds PropagationNet GNN layer and makes it optionally usable in existing deterministic models #507 merged).
feat/add-ensemble-metrics-crps-spread #226 which brings the ensemble metrics on the current main.
feat: Add batched tensor expansion utilities for O(S) ensemble unroll… #489 add ensemble branch fold utility

Issue Link

Partially addresses #62 (prerequisite for the GraphEFMPredictor PR).

Type of change

🐛 Bug fix (non-breaking change that fixes an issue)
✨ New feature (non-breaking change that adds functionality)
💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
📖 Documentation (Addition or improvements to documentation)

Checklist before requesting a review

My branch is up-to-date with the target branch
I have performed a self-review of my code
For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
I have placed in-line comments to clarify the intent of any hard-to-understand passages of my code
I have updated the README to cover introduced code changes
I have added tests that prove my fix is effective or that my feature works
I have given the PR a name that clearly describes the change, written in imperative form
I have requested a reviewer and an assignee

Checklist for reviewers

the code is readable
the code is well tested
the code is documented (including return types and parameters)
the code is easy to maintain

Author checklist after completed review

I have added a line to the CHANGELOG describing this change

Checklist for assignee

PR is up to date with the base branch
the tests pass
the PR is assigned to the next milestone
author has added an entry to the changelog

Sir-Sloth-The-Lazy · 2026-05-28T00:18:35Z

Pinging @joeloskarsson !

joeloskarsson

Had a look over the encoders now, decoders still TODO :)

Sir-Sloth-The-Lazy · 2026-06-05T01:19:17Z

On it !

Sir-Sloth-The-Lazy · 2026-06-05T11:42:00Z

I have pushed the latest changes , please have a look whenever you have time

joeloskarsson

Good work with fixing the earlier things! I now did a complete readthrough of all of this, which resulted in a number of comments, some small, some more substantial. There are also a few that require more discussion, so do expect us to have some back and forth around these.

As you are fixing these things, please leave me some comments very briefly explaining the fix and link to the commit that fixes it (a good reason to fix one thing in each commit :)).

joeloskarsson

Good job with all the nice fixes! I went over some of the previous comments and resolved most of them. Still have some left, but submitting these comments for now just so you can see them.

joeloskarsson

Now I have gone over everything, and resolved or commented on things needing further attention 😄This is coming along super nicely I think, and especially nice that we are finding some shortcomings of both the earlier graph-efm code and in the repo, and able to resolve much of that.

Most things left are smaller fixes. The bigger question is how this will fit with the larger Module/Forecaster framework. This should be the priority to hammer out now. I have also not yet looked much at the methods in GraphEFM that deal with loss computation, simply because I expect these might change a bit once we figure out the details of the probabilistic model training interface.

Sir-Sloth-The-Lazy · 2026-06-18T10:02:52Z

Thank you @joeloskarsson for the review 🤩. I would be happy to take the work of utils refactor. I will be opening another PR soon solving that issue

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.

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.

Sir-Sloth-The-Lazy · 2026-07-11T03:36:38Z

@observingClouds, please spare sometime and review this work as well

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.

joeloskarsson · 2026-07-17T16:46:22Z

I am pretty happy with this now, just

figuring out some details around the prior network instantiation above.
need to get this back up to date with main and look over tests and linting.

@observingClouds do you want to give this a look over as well before we merge it in? Recall that this will not just be the model components that make up the Graph-EFM model as a step-predictor, not any of the mechanisms to actually train it and compute the related losses. That will be a follow up building on #700, instantiating the interface there.

Add JetBrains IDE project directory to the ignore list alongside the existing .vim/.vscode entries.

observingClouds · 2026-07-18T22:00:29Z

I'm good with this once your concerns/comments are addressed

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.

Sir-Sloth-The-Lazy · 2026-07-21T00:15:56Z

Completed the requested changes and resolved the conflict ! @joeloskarsson @observingClouds

joeloskarsson · 2026-07-22T19:31:00Z

Gave this a complete read through before merging. Unfortunately I found something related to the clamping that I had missed earlier :/ (see above). But otherwise this all looks good. So if we can just sort out clamping for this model then I will go ahead and merge.

Move the residual connection (current state + decoder increment) out of the latent decoder and into the predictor, which then applies get_clamped_new_state so output clamping applies to Graph-EFM like it does for the deterministic models. The latent decoders now output the state increment only (dropping the last_state argument).

Use an empty list for the intra-level (m2m) mesh embedding when no intra-level GNNs are configured, instead of passing raw unembedded edge features that are never consumed. Gate the corresponding m2m access in the hierarchical decoder so the empty placeholder is never indexed.

Sir-Sloth-The-Lazy · 2026-07-25T04:05:26Z

@joeloskarsson I have completed the changes !

joeloskarsson

Thanks for sorting out the clamping as well! This all looks good now, so I will go ahead and merge.

Main moved the graph loading and buffer registration out of BaseGraphModel into utils.load_and_register_graph (mllam#648), so the HeteroData wiring moves there with it. It is now applied wherever a module loads a graph through that helper, which includes the new Graph-EFM step predictors, and the grid node count is read from the datastore directly.

joeloskarsson reviewed Jun 1, 2026

View reviewed changes

Comment thread neural_lam/models/forecaster_module.py Outdated

joeloskarsson reviewed Jun 1, 2026

View reviewed changes

Comment thread neural_lam/models/latent/__init__.py

joeloskarsson added this to the v0.7.0 (proposed) milestone Jun 4, 2026

joeloskarsson reviewed Jun 4, 2026

View reviewed changes

joeloskarsson self-assigned this Jun 4, 2026

Sir-Sloth-The-Lazy requested a review from joeloskarsson June 5, 2026 11:41

sadamov mentioned this pull request Jun 6, 2026

Probabilistic metrics & ensemble plumbing - consolidated track for v0.8.0 #649

Open

joeloskarsson modified the milestones: v0.7.0 (proposed), v0.7.0 Jun 8, 2026

joeloskarsson requested changes Jun 11, 2026

View reviewed changes

joeloskarsson requested changes Jun 16, 2026

View reviewed changes

joeloskarsson reviewed Jun 17, 2026

View reviewed changes

Comment thread neural_lam/utils.py Outdated

joeloskarsson requested changes Jun 17, 2026

View reviewed changes

Comment thread neural_lam/utils.py Outdated

Comment thread neural_lam/models/step_predictors/graph/graph_efm.py Outdated

This was referenced Jun 21, 2026

feat: add probabilistic forecaster interface (decoupled from Graph-EFM) Sir-Sloth-The-Lazy/neural-lam#2

Closed

feat: add probabilistic forecaster interface (decoupled from Graph-EFM) #678

Closed

Sir-Sloth-The-Lazy and others added 11 commits June 21, 2026 17:08

Update neural_lam/models/latent/constant_encoder.py

6c23ff3

Co-authored-by: Joel Oskarsson <joel.oskarsson@outlook.com>

style: apply black formatting to hierarchical latent modules

df632b7

docs: convert latent module and make_gnn_seq docstrings to numpy style

0e50d9c

Add proper Parameters/Returns sections following the numpydoc convention, per review.

docs: align ConstantLatentEncoder.compute_dist_params docstring with …

d2d249b

…base class Use the base class summary and expand on it with the constant-specific behavior, per review.

Sir-Sloth-The-Lazy added 2 commits July 11, 2026 08:38

docs: clarify GNN-type flag scope in train_model.py

f5c5bc0

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.

Sir-Sloth-The-Lazy added 2 commits July 15, 2026 02:25

Merge branch 'main' into feat/latent-encoder-decoder-infra

c086298

joeloskarsson reviewed Jul 17, 2026

View reviewed changes

Comment thread neural_lam/utils.py Outdated

joeloskarsson marked this pull request as ready for review July 17, 2026 16:43

Ignore .idea directory in .gitignore

98ab692

Add JetBrains IDE project directory to the ignore list alongside the existing .vim/.vscode entries.

Sir-Sloth-The-Lazy added 2 commits July 21, 2026 03:33

Sir-Sloth-The-Lazy requested a review from joeloskarsson July 20, 2026 23:07

joeloskarsson reviewed Jul 22, 2026

View reviewed changes

Comment thread neural_lam/models/step_predictors/graph/graph_efm.py Outdated

joeloskarsson reviewed Jul 22, 2026

View reviewed changes

Comment thread neural_lam/models/step_predictors/graph/graph_efm.py Outdated

Sir-Sloth-The-Lazy mentioned this pull request Jul 25, 2026

feat: Graph-EFM ensemble forecasting model #712

Draft

21 tasks

Sir-Sloth-The-Lazy added 3 commits July 25, 2026 09:21

docs: add CHANGELOG entry for Graph-EFM latent infra and step predictors

ac36d08

Sir-Sloth-The-Lazy requested a review from joeloskarsson July 25, 2026 04:05

Address PR review: rescale decoder output with difference statistics

d10b078

joeloskarsson approved these changes Aug 2, 2026

View reviewed changes

joeloskarsson merged commit c27ff8b into mllam:main Aug 2, 2026
9 checks passed

This was referenced Aug 7, 2026

Support hierarchical graphs in the HeteroData representation #713

Open

Load flat graphs into pyg.HeteroData and wire into BaseGraphModel #711

Open

Uh oh!

Conversation

Sir-Sloth-The-Lazy commented May 27, 2026 • edited by joeloskarsson Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Describe your changes

Issue Link

Type of change

Checklist before requesting a review

Checklist for reviewers

Author checklist after completed review

Checklist for assignee

Uh oh!

Sir-Sloth-The-Lazy commented May 28, 2026

Uh oh!

Uh oh!

Uh oh!

joeloskarsson left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Sir-Sloth-The-Lazy commented Jun 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Sir-Sloth-The-Lazy commented Jun 5, 2026

Uh oh!

joeloskarsson left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

joeloskarsson left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

joeloskarsson left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Sir-Sloth-The-Lazy commented Jun 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Sir-Sloth-The-Lazy commented Jul 11, 2026

Uh oh!

Uh oh!

joeloskarsson commented Jul 17, 2026

Uh oh!

observingClouds commented Jul 18, 2026

Uh oh!

Sir-Sloth-The-Lazy commented Jul 21, 2026

Uh oh!

Uh oh!

Uh oh!

joeloskarsson commented Jul 22, 2026

Uh oh!

Sir-Sloth-The-Lazy commented Jul 25, 2026

Uh oh!

joeloskarsson left a comment

Choose a reason for hiding this comment

Uh oh!

Sir-Sloth-The-Lazy commented May 27, 2026 •

edited by joeloskarsson

Loading

Sir-Sloth-The-Lazy commented Jun 5, 2026 •

edited

Loading

Sir-Sloth-The-Lazy commented Jun 18, 2026 •

edited

Loading