Feature Request: Add a Graph Neural Operator Predictor

[Milan933-coder]

Summary

Add a gno_predictor model to Neural-LAM as a new StepPredictor for
operator-style weather prediction on node and graph data.

Duplicate search

I searched existing mllam/neural-lam issues for GNO,
Graph Neural Operator, and neural operator, and found no duplicates.

Motivation

Neural-LAM already represents weather states as node tensors and supports
graph-based predictors, so a Graph Neural Operator is a natural fit.

Unlike a CNN predictor, a GNO can work directly on node/graph data without
requiring a regular (x, y) grid reshape.

This could provide an operator-learning baseline for limited-area weather
prediction and make it easier to compare message-passing GNNs with
mesh/discretization-aware neural operator layers.

Research background

• Li et al., "Neural Operator: Graph Kernel Network for Partial Differential
  Equations" formulates operator learning with graph-kernel integral layers
  computed by message passing.
  https://arxiv.org/abs/2003.03485
• Li et al., "Multipole Graph Neural Operator for Parametric Partial
  Differential Equations" adds multi-level graph structure to capture
  long-range interactions with linear complexity.
  https://arxiv.org/abs/2006.09535
• Kovachki et al., "Neural Operator: Learning Maps Between Function Spaces"
  frames neural operators as maps between function spaces and includes graph
  neural operators as one parameterization.
  https://arxiv.org/abs/2108.08481

Proposed implementation

Add a new predictor package:

neural_lam/models/step_predictors/gno/
  __init__.py
  gno_layers.py
  gno_predictor.py

Register it as:

"gno_predictor": GNOPredictor

The predictor should follow the existing StepPredictor interface:

prev_state, prev_prev_state, forcing, static features
-> graph neural operator layers
-> standardized one-step state delta
-> diff-stat rescaling
-> get_clamped_new_state
-> ARForecaster rollout

Suggested scope

• Add GNOLayer with coordinate/edge-conditioned kernel messages.
• Add GNOPredictor(StepPredictor).
• Reuse existing graph edge features or datastore coordinates.
• Support output_std with positive softplus standard deviations.
• Keep boundary handling inside ARForecaster.
• Register gno_predictor in MODELS.
• Add CLI construction arguments for hidden size, layers, radius/graph
  choice, and aggregation.
• Add tests for forward shape, output_std, no-static-feature behavior,
  invalid inputs, and autoregressive rollout.

Minimal usage example

from neural_lam.models import ARForecaster, MODELS

predictor = MODELS["gno_predictor"](
    datastore=datastore,
    hidden_dim=64,
    processor_layers=4,
    num_past_forcing_steps=1,
    num_future_forcing_steps=1,
    output_std=False,
    output_clamping_lower=config.training.output_clamping.lower,
    output_clamping_upper=config.training.output_clamping.upper,
)
forecaster = ARForecaster(predictor, datastore)

Validation checklist

• Compare one-step output shapes against GraphLAM and CNNPredictor.
• Run focused unit tests for GNOLayer and GNOPredictor.
• Run a short training smoke test on DummyDatastore.
• Verify checkpoint save/load reconstruction.
• Add a CHANGELOG.md entry under Added.

Notes

This should be a separate feature from the current CNN predictor work.

[Milan933-coder]

@claude @joeloskarsson have a look on this issue.

[joeloskarsson]

Is there some specific capability that the existing graph-based models lack, that you would want from a GNO? To me they seem very similar to existing message passing GNNs.

While we can for sure add additional models, preferably the motivation should be that there is some real interest in a capability offered by these, that existing models lack, rather than just adding more models because we can. I think it would be good to outline how this model would differ from existing GNNs in the codebase before starting work here.

[Milan933-coder]

As far as i know Neural operators can map in the function space ratehr than parameter space .
But I have to Some research on it.
If it is good Then We can work upon it .

The original LAM Was based upon physics based simulation and solving Partial differnetial Equation That's Neural Operators Are introduced .

I just recently Got to know about it that's why i asked ?

I will do some lit review and i will let you know

[sadamov]

One use-case could be when we allow various input sources living on heterogeneous grids and timesteps. In theory neural operators should allow to generate outputs at arbitrary locations and temporal resolutions, inter-/extrapolating from the training data. Such a model would be very flexibly about its inputs and outputs. But other neural operators like FourCastNet did actually not generalize that well to unseen resolutions, in practice (paper, paper). I think this generalization aspect is very interesting to investigate.