neural_lam.models.step_predictors.graph.graph_efm ================================================= .. py:module:: neural_lam.models.step_predictors.graph.graph_efm .. autoapi-nested-parse:: Graph-based Ensemble Forecasting Model (Graph-EFM) single-step predictors, for hierarchical (GraphEFM) and flat (GraphEFMMultiScale) mesh graphs. Classes ------- .. autoapisummary:: neural_lam.models.step_predictors.graph.graph_efm.BaseGraphEFM neural_lam.models.step_predictors.graph.graph_efm.GraphEFM neural_lam.models.step_predictors.graph.graph_efm.GraphEFMMultiScale Module Contents --------------- .. py:class:: BaseGraphEFM(datastore: neural_lam.datastore.BaseDatastore, graph_name: str, hidden_dim: int = 64, hidden_layers: int = 1, latent_dim: Optional[int] = None, learn_prior: bool = True, prior_dist: str = 'isotropic', prior_layers: int = 2, g2m_gnn_type: str = 'InteractionNet', num_past_forcing_steps: int = 1, num_future_forcing_steps: int = 1, output_std: bool = False, output_clamping_lower: Optional[Dict[str, float]] = None, output_clamping_upper: Optional[Dict[str, float]] = None) Bases: :py:obj:`neural_lam.models.step_predictors.base.StepPredictor` Base class for Graph-based Ensemble Forecasting Model single-step predictors. A latent-variable step predictor consisting of a conditional prior, a variational encoder and a latent decoder, each of which carries its own grid-to-mesh, on-mesh and mesh-to-grid GNNs. The encode-process-decode backbone of ``BaseGraphModel`` therefore does not apply -- this extends ``StepPredictor`` directly. ``forward`` samples a single step from the prior. The encoder (variational posterior) and prior are exposed as ``self.encoder``/``self.prior_model`` for a forecaster to condition on the target and assemble a training objective from; the predictor itself does not compute any loss. This base class sets up everything that is independent of the mesh graph type: it loads the graph, calls the subclass's :meth:`check_graph_type` to verify the loaded graph matches what it requires, then builds the prior, delegating to the subclass's :meth:`build_learnable_prior` when learned and to :attr:`latent_spatial_dim` for the constant prior's node count. Concrete subclasses build the mesh embedders and the encoder/decoder latent modules, and implement :meth:`embedd_mesh`, :meth:`check_graph_type`, :meth:`build_learnable_prior` and :attr:`latent_spatial_dim`. See :class:`GraphEFM` (hierarchical graph) and :class:`GraphEFMMultiScale` (flat graph). Set up the graph-type independent parts of the predictor. Loads the graph, builds the grid embedders, the grid-mesh edge embedders and the prior. Building the mesh embedders and the encoder/decoder latent modules is left to the subclass constructor. :param datastore: Datastore providing static features, standardization statistics and variable counts. :type datastore: BaseDatastore :param graph_name: Name of the graph directory (under ``/graph``) to load. Must be of the graph type required by the concrete subclass. :type graph_name: str :param hidden_dim: Dimensionality of internal node and edge representations. :type hidden_dim: int :param hidden_layers: Number of hidden layers in internal MLPs. :type hidden_layers: int :param latent_dim: Dimensionality of the latent variable at each latent-carrying mesh node (top level for hierarchical graphs, all mesh nodes for flat graphs); defaults to ``hidden_dim`` when None. The resolved value is stored as ``self.latent_dim`` for the subclass to reuse when building its encoder/decoder. :type latent_dim: int, optional :param learn_prior: If True, the prior is the graph-type specific learnable encoder built by :meth:`build_learnable_prior`, conditioned on the previous state; if False, a constant ``Normal(0, 1)`` prior is used. :type learn_prior: bool :param prior_dist: Output distribution of the prior: ``"isotropic"`` or ``"diagonal"``. :type prior_dist: str :param prior_layers: Number of on-mesh GNN layers in the learnable prior. :type prior_layers: int :param g2m_gnn_type: GNN type for the grid-to-mesh step of the prior (key in ``gnn_layers.GNN_TYPES``). :type g2m_gnn_type: str :param num_past_forcing_steps: Number of past forcing steps included in the input window. :type num_past_forcing_steps: int :param num_future_forcing_steps: Number of future forcing steps included in the input window. :type num_future_forcing_steps: int :param output_std: If True, the decoder outputs a per-variable std alongside the mean; if False, ``forward`` returns ``None`` for the std. :type output_std: bool :param output_clamping_lower: Lower clamping limits per output variable. :type output_clamping_lower: dict of str to float, optional :param output_clamping_upper: Upper clamping limits per output variable. :type output_clamping_upper: dict of str to float, optional .. py:method:: build_learnable_prior(latent_dim: int, hidden_dim: int, hidden_layers: int, g2m_gnn_type: str, prior_dist: str, prior_layers: int) -> torch.nn.Module :abstractmethod: Build the graph-type specific learnable prior encoder. Implemented by the concrete subclass, which knows the mesh graph type and therefore the appropriate latent encoder class. :param latent_dim: Dimensionality of the latent variable at each mesh node. :type latent_dim: int :param hidden_dim: Dimensionality of internal node and edge representations. :type hidden_dim: int :param hidden_layers: Number of hidden layers in internal MLPs. :type hidden_layers: int :param g2m_gnn_type: GNN type for the grid-to-mesh step of the prior. :type g2m_gnn_type: str :param prior_dist: Output distribution of the prior: ``"isotropic"`` or ``"diagonal"``. :type prior_dist: str :param prior_layers: Number of on-mesh GNN layers in the prior. :type prior_layers: int :returns: The learnable prior latent encoder. :rtype: torch.nn.Module .. py:method:: check_graph_type(graph_name: str) -> None :abstractmethod: Verify the loaded graph (``self.hierarchical``) is of the type this predictor requires. Implemented by the concrete subclass, which is the only place that knows whether it requires a hierarchical or flat mesh graph. Called by this base class right after loading the graph, before anything that assumes a specific graph shape. :param graph_name: Name of the graph directory that was loaded, for the error message. :type graph_name: str :raises ValueError: If ``self.hierarchical`` does not match what this predictor requires. .. py:method:: embedd_grid_and_graph(prev_state: torch.Tensor, prev_prev_state: torch.Tensor, forcing: torch.Tensor) -> tuple[torch.Tensor, dict[str, Any]] Embed the grid (states up to t-1) and the full graph. :param prev_state: Shape ``(B, num_grid_nodes, d_state)``. ``X_t``. :type prev_state: torch.Tensor :param prev_prev_state: Shape ``(B, num_grid_nodes, d_state)``. ``X_{t-1}``. :type prev_prev_state: torch.Tensor :param forcing: Shape ``(B, num_grid_nodes, d_forcing)``. :type forcing: torch.Tensor :returns: * **grid_emb** (*torch.Tensor*) -- Shape ``(B, num_grid_nodes, d_h)``. Grid embedding. * **graph_emb** (*dict*) -- Edge/mesh embeddings, each entry of shape ``(B, *, d_h)``. .. py:method:: embedd_grid_with_target(prev_state: torch.Tensor, prev_prev_state: torch.Tensor, forcing: torch.Tensor, current_state: torch.Tensor) -> torch.Tensor Embed the grid representation including the current (target) state. Used as input to the encoder, which is conditioned also on the target. :param prev_state: Shape ``(B, num_grid_nodes, d_state)``. ``X_t``. :type prev_state: torch.Tensor :param prev_prev_state: Shape ``(B, num_grid_nodes, d_state)``. ``X_{t-1}``. :type prev_prev_state: torch.Tensor :param forcing: Shape ``(B, num_grid_nodes, d_forcing)``. :type forcing: torch.Tensor :param current_state: Shape ``(B, num_grid_nodes, d_state)``. ``X_{t+1}`` (target). :type current_state: torch.Tensor :returns: Shape ``(B, num_grid_nodes, d_h)``. Grid embedding. :rtype: torch.Tensor .. py:method:: embedd_mesh(batch_size: int) -> dict[str, Any] :abstractmethod: Embed static mesh node and intra-mesh edge features. :param batch_size: Batch size to expand the embeddings to. :type batch_size: int :returns: Mesh-related entries of the graph embedding (``mesh``, ``m2m`` and, for hierarchical graphs, ``mesh_up`` and ``mesh_down``). Entries are tensors of shape ``(B, *, d_h)`` for flat graphs and per-level lists of such tensors for hierarchical graphs. :rtype: dict .. py:method:: forward(prev_state: torch.Tensor, prev_prev_state: torch.Tensor, forcing: torch.Tensor) -> tuple[torch.Tensor, Optional[torch.Tensor]] Sample one time step prediction: embed features, sample the latent from the prior, decode, and return the predicted next state. The prediction is stochastic only through the latent sample; no observation noise is added. :param prev_state: Shape ``(B, num_grid_nodes, d_state)``. ``X_t``. :type prev_state: torch.Tensor :param prev_prev_state: Shape ``(B, num_grid_nodes, d_state)``. ``X_{t-1}``. :type prev_prev_state: torch.Tensor :param forcing: Shape ``(B, num_grid_nodes, d_forcing)``. :type forcing: torch.Tensor :returns: * **new_state** (*torch.Tensor*) -- Shape ``(B, num_grid_nodes, d_state)``. Predicted ``X_{t+1}`` (the decoder mean, given the sampled latent). * **pred_std** (*torch.Tensor or None*) -- Shape ``(B, num_grid_nodes, d_state)`` when ``output_std`` is True, otherwise None. .. py:attribute:: g2m_edge_index :type: torch.Tensor .. py:attribute:: g2m_embedder .. py:attribute:: grid_current_embedder .. py:attribute:: grid_dim .. py:attribute:: grid_prev_embedder .. py:attribute:: hierarchical .. py:attribute:: latent_dim .. py:property:: latent_spatial_dim :type: int :abstractmethod: Number of mesh nodes the latent variable lives on. Implemented by the concrete subclass, which knows the mesh graph type: the top mesh level for hierarchical graphs, or every mesh node for flat graphs. :returns: Number of latent-carrying mesh nodes. :rtype: int .. py:attribute:: m2g_edge_index :type: torch.Tensor .. py:attribute:: m2g_embedder .. py:attribute:: mlp_blueprint_end .. py:attribute:: num_state_vars .. py:class:: GraphEFM(datastore: neural_lam.datastore.BaseDatastore, graph_name: str = 'hierarchical', hidden_dim: int = 64, hidden_layers: int = 1, latent_dim: Optional[int] = None, prior_intra_level_layers: int = 2, encoder_intra_level_layers: int = 2, decoder_intra_level_layers: int = 4, learn_prior: bool = True, prior_dist: str = 'isotropic', num_past_forcing_steps: int = 1, num_future_forcing_steps: int = 1, g2m_gnn_type: str = 'InteractionNet', m2g_gnn_type: str = 'InteractionNet', output_std: bool = False, output_clamping_lower: Optional[Dict[str, float]] = None, output_clamping_upper: Optional[Dict[str, float]] = None) Bases: :py:obj:`BaseGraphEFM` Graph-based Ensemble Forecasting Model on a hierarchical mesh graph. The latent variable lives on the top level of the mesh hierarchy. The prior and variational encoder are ``HiGraphLatentEncoder``s and the decoder is a ``HiGraphLatentDecoder``. Build the mesh embedders and the hierarchical encoder/decoder latent modules. The prior is built by the base class. :param datastore: Datastore providing static features, standardization statistics and variable counts. :type datastore: BaseDatastore :param graph_name: Name of the graph directory (under ``/graph``) to load. Must be a hierarchical graph. :type graph_name: str :param hidden_dim: Dimensionality of internal node and edge representations. :type hidden_dim: int :param hidden_layers: Number of hidden layers in internal MLPs. :type hidden_layers: int :param latent_dim: Dimensionality of the latent variable at each top-level mesh node; defaults to ``hidden_dim`` when None. Forwarded to the base class, which resolves the default and stores it as ``self.latent_dim``. :type latent_dim: int, optional :param prior_intra_level_layers: Number of intra-level GNN layers in the (learned) prior. Forwarded to the base class as ``prior_layers``. :type prior_intra_level_layers: int :param encoder_intra_level_layers: Number of intra-level GNN layers in the variational encoder. :type encoder_intra_level_layers: int :param decoder_intra_level_layers: Number of intra-level GNN layers in the latent decoder. :type decoder_intra_level_layers: int :param learn_prior: If True, the prior is a hierarchical graph encoder conditioned on the previous state; if False, a constant ``Normal(0, 1)`` prior is used. Forwarded to the base class. :type learn_prior: bool :param prior_dist: Output distribution of the prior: ``"isotropic"`` or ``"diagonal"``. Forwarded to the base class. :type prior_dist: str :param num_past_forcing_steps: Number of past forcing steps included in the input window. :type num_past_forcing_steps: int :param num_future_forcing_steps: Number of future forcing steps included in the input window. :type num_future_forcing_steps: int :param g2m_gnn_type: GNN type for the grid-to-mesh steps of the prior, encoder and decoder (key in ``gnn_layers.GNN_TYPES``). :type g2m_gnn_type: str :param m2g_gnn_type: GNN type for the mesh-to-grid step of the decoder (key in ``gnn_layers.GNN_TYPES``). :type m2g_gnn_type: str :param output_std: If True, the decoder outputs a per-variable std alongside the mean; if False, ``forward`` returns ``None`` for the std. :type output_std: bool :param output_clamping_lower: Lower clamping limits per output variable. :type output_clamping_lower: dict of str to float, optional :param output_clamping_upper: Upper clamping limits per output variable. :type output_clamping_upper: dict of str to float, optional .. py:method:: build_learnable_prior(latent_dim: int, hidden_dim: int, hidden_layers: int, g2m_gnn_type: str, prior_dist: str, prior_layers: int) -> neural_lam.models.latent.HiGraphLatentEncoder Build the hierarchical learnable prior encoder. :param latent_dim: Dimensionality of the latent variable at each top-level mesh node. :type latent_dim: int :param hidden_dim: Dimensionality of internal node and edge representations. :type hidden_dim: int :param hidden_layers: Number of hidden layers in internal MLPs. :type hidden_layers: int :param g2m_gnn_type: GNN type for the grid-to-mesh step of the prior. :type g2m_gnn_type: str :param prior_dist: Output distribution of the prior: ``"isotropic"`` or ``"diagonal"``. :type prior_dist: str :param prior_layers: Number of intra-level GNN layers in the prior. :type prior_layers: int :returns: The learnable prior latent encoder. :rtype: HiGraphLatentEncoder .. py:method:: check_graph_type(graph_name: str) -> None Verify the loaded graph is hierarchical. :param graph_name: Name of the graph directory that was loaded, for the error message. :type graph_name: str :raises ValueError: If the loaded graph is flat. .. py:method:: embedd_mesh(batch_size: int) -> dict[str, Any] Embed static mesh node and intra-mesh edge features per level. :param batch_size: Batch size to expand the embeddings to. :type batch_size: int :returns: Entries ``mesh``, ``m2m``, ``mesh_up`` and ``mesh_down``, each a list with one ``(B, *, d_h)`` tensor per mesh level (or inter-level connection). :rtype: dict .. py:attribute:: decoder .. py:attribute:: embedd_m2m .. py:attribute:: encoder .. py:property:: latent_spatial_dim :type: int Number of mesh nodes on the top mesh level, where the latent variable lives. :returns: Number of top-level mesh nodes. :rtype: int .. py:attribute:: m2m_edge_index :type: neural_lam.utils.BufferList .. py:attribute:: m2m_features :type: list[torch.Tensor] .. py:attribute:: mesh_down_edge_index :type: neural_lam.utils.BufferList .. py:attribute:: mesh_down_embedders .. py:attribute:: mesh_down_features :type: list[torch.Tensor] .. py:attribute:: mesh_embedders .. py:attribute:: mesh_static_features :type: list[torch.Tensor] .. py:attribute:: mesh_up_edge_index :type: neural_lam.utils.BufferList .. py:attribute:: mesh_up_embedders .. py:attribute:: mesh_up_features :type: list[torch.Tensor] .. py:class:: GraphEFMMultiScale(datastore: neural_lam.datastore.BaseDatastore, graph_name: str = 'multiscale', hidden_dim: int = 64, hidden_layers: int = 1, latent_dim: Optional[int] = None, prior_m2m_layers: int = 2, encoder_m2m_layers: int = 2, decoder_m2m_layers: int = 4, learn_prior: bool = True, prior_dist: str = 'isotropic', num_past_forcing_steps: int = 1, num_future_forcing_steps: int = 1, g2m_gnn_type: str = 'InteractionNet', m2g_gnn_type: str = 'InteractionNet', output_std: bool = False, output_clamping_lower: Optional[Dict[str, float]] = None, output_clamping_upper: Optional[Dict[str, float]] = None) Bases: :py:obj:`BaseGraphEFM` Graph-based Ensemble Forecasting Model on a flat mesh graph (e.g. a multi-scale graph). The latent variable lives on the (single) level of the mesh graph. The prior and variational encoder are ``GraphLatentEncoder``s and the decoder is a ``GraphLatentDecoder``. Build the mesh embedders and the flat-graph encoder/decoder latent modules. The prior is built by the base class. :param datastore: Datastore providing static features, standardization statistics and variable counts. :type datastore: BaseDatastore :param graph_name: Name of the graph directory (under ``/graph``) to load. Must be a flat graph. :type graph_name: str :param hidden_dim: Dimensionality of internal node and edge representations. :type hidden_dim: int :param hidden_layers: Number of hidden layers in internal MLPs. :type hidden_layers: int :param latent_dim: Dimensionality of the latent variable at each mesh node; defaults to ``hidden_dim`` when None. Forwarded to the base class, which resolves the default and stores it as ``self.latent_dim``. :type latent_dim: int, optional :param prior_m2m_layers: Number of on-mesh (m2m) GNN layers in the (learned) prior. Forwarded to the base class as ``prior_layers``. :type prior_m2m_layers: int :param encoder_m2m_layers: Number of on-mesh (m2m) GNN layers in the variational encoder. :type encoder_m2m_layers: int :param decoder_m2m_layers: Number of on-mesh (m2m) GNN layers in the latent decoder. :type decoder_m2m_layers: int :param learn_prior: If True, the prior is a graph encoder conditioned on the previous state; if False, a constant ``Normal(0, 1)`` prior is used. Forwarded to the base class. :type learn_prior: bool :param prior_dist: Output distribution of the prior: ``"isotropic"`` or ``"diagonal"``. Forwarded to the base class. :type prior_dist: str :param num_past_forcing_steps: Number of past forcing steps included in the input window. :type num_past_forcing_steps: int :param num_future_forcing_steps: Number of future forcing steps included in the input window. :type num_future_forcing_steps: int :param g2m_gnn_type: GNN type for the grid-to-mesh steps of the prior, encoder and decoder (key in ``gnn_layers.GNN_TYPES``). :type g2m_gnn_type: str :param m2g_gnn_type: GNN type for the mesh-to-grid step of the decoder (key in ``gnn_layers.GNN_TYPES``). :type m2g_gnn_type: str :param output_std: If True, the decoder outputs a per-variable std alongside the mean; if False, ``forward`` returns ``None`` for the std. :type output_std: bool :param output_clamping_lower: Lower clamping limits per output variable. :type output_clamping_lower: dict of str to float, optional :param output_clamping_upper: Upper clamping limits per output variable. :type output_clamping_upper: dict of str to float, optional .. py:method:: build_learnable_prior(latent_dim: int, hidden_dim: int, hidden_layers: int, g2m_gnn_type: str, prior_dist: str, prior_layers: int) -> neural_lam.models.latent.GraphLatentEncoder Build the flat-graph learnable prior encoder. :param latent_dim: Dimensionality of the latent variable at each mesh node. :type latent_dim: int :param hidden_dim: Dimensionality of internal node and edge representations. :type hidden_dim: int :param hidden_layers: Number of hidden layers in internal MLPs. :type hidden_layers: int :param g2m_gnn_type: GNN type for the grid-to-mesh step of the prior. :type g2m_gnn_type: str :param prior_dist: Output distribution of the prior: ``"isotropic"`` or ``"diagonal"``. :type prior_dist: str :param prior_layers: Number of on-mesh (m2m) GNN layers in the prior. :type prior_layers: int :returns: The learnable prior latent encoder. :rtype: GraphLatentEncoder .. py:method:: check_graph_type(graph_name: str) -> None Verify the loaded graph is flat. :param graph_name: Name of the graph directory that was loaded, for the error message. :type graph_name: str :raises ValueError: If the loaded graph is hierarchical. .. py:method:: embedd_mesh(batch_size: int) -> dict[str, Any] Embed static mesh node and intra-mesh edge features. :param batch_size: Batch size to expand the embeddings to. :type batch_size: int :returns: Entries ``mesh``: ``(B, num_mesh_nodes, d_h)`` and ``m2m``: ``(B, M_m2m, d_h)``. :rtype: dict .. py:attribute:: decoder .. py:attribute:: encoder .. py:property:: latent_spatial_dim :type: int Number of mesh nodes, where the latent variable lives. :returns: Number of mesh nodes. :rtype: int .. py:attribute:: m2m_edge_index :type: torch.Tensor .. py:attribute:: m2m_embedder .. py:attribute:: m2m_features :type: torch.Tensor .. py:attribute:: mesh_embedder .. py:attribute:: mesh_static_features :type: torch.Tensor