neural_lam.models.step_predictors.graph.base ============================================ .. py:module:: neural_lam.models.step_predictors.graph.base .. autoapi-nested-parse:: Base class for graph-based step predictors. Classes ------- .. autoapisummary:: neural_lam.models.step_predictors.graph.base.BaseGraphModel Module Contents --------------- .. py:class:: BaseGraphModel(datastore: neural_lam.datastore.BaseDatastore, graph_name: str = 'multiscale', hidden_dim: int = 64, hidden_layers: int = 1, processor_layers: int = 4, mesh_aggr: str = 'sum', num_past_forcing_steps: int = 1, num_future_forcing_steps: int = 1, output_std: bool = False, output_clamping_lower: dict[str, float] | None = None, output_clamping_upper: dict[str, float] | None = None, g2m_gnn_type: str = 'InteractionNet', m2g_gnn_type: str = 'InteractionNet') Bases: :py:obj:`neural_lam.models.step_predictors.base.StepPredictor` Base (abstract) class for graph-based models building on the encode-process-decode idea. Initialize the BaseGraphModel. :param datastore: Datastore supplying data and information about dataset and forecast region. :type datastore: BaseDatastore :param graph_name: The name of the graph to load. :type graph_name: str, default "multiscale" :param hidden_dim: The dimension of the hidden representations. :type hidden_dim: int, default 64 :param hidden_layers: The number of hidden layers in the MLPs. :type hidden_layers: int, default 1 :param processor_layers: The number of processor layers in the GNN. :type processor_layers: int, default 4 :param mesh_aggr: The aggregation method for mesh nodes. :type mesh_aggr: str, default "sum" :param num_past_forcing_steps: The number of past forcing steps to include. :type num_past_forcing_steps: int, default 1 :param num_future_forcing_steps: The number of future forcing steps to include. :type num_future_forcing_steps: int, default 1 :param output_std: Whether to output a predicted standard deviation. :type output_std: bool, default False :param output_clamping_lower: Lower clamping limits for state variables. :type output_clamping_lower: dict, optional :param output_clamping_upper: Upper clamping limits for state variables. :type output_clamping_upper: dict, optional .. py:method:: embedd_mesh_nodes() -> torch.Tensor :abstractmethod: Embed static mesh node features. :returns: Shape ``(num_mesh_nodes, hidden_dim)``. Embedded mesh node representations. Dims: ``num_mesh_nodes`` is the number of mesh nodes and ``hidden_dim`` is the hidden dimension. :rtype: torch.Tensor .. py:method:: forward(prev_state: torch.Tensor, prev_prev_state: torch.Tensor, forcing: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None] Advance the state by one step using the encode-process-decode graph pipeline: embed grid + edge features, map grid -> mesh via the g2m GNN, run the (subclass-defined) processor on the mesh, decode mesh -> grid via the m2g GNN, project to a state delta, rescale with difference statistics, and add to ``prev_state`` (with optional clamping). Returns ``(X_{t+1}, optional std)``. :param prev_state: Shape ``(B, num_grid_nodes, num_state_vars)``. The current state ``X_t``. Dims: ``B`` is batch size, ``num_grid_nodes`` is the number of spatial grid nodes, and ``num_state_vars`` is the number of state variables. :type prev_state: torch.Tensor :param prev_prev_state: Shape ``(B, num_grid_nodes, num_state_vars)``. The previous state ``X_{t-1}``, used as additional conditioning. Dims: same as ``prev_state``. :type prev_prev_state: torch.Tensor :param forcing: Shape ``(B, num_grid_nodes, num_forcing_vars)``. External forcings for this step (already concatenated past/current/future windows). Dims: ``B`` is batch size, ``num_grid_nodes`` is the number of spatial grid nodes, and ``num_forcing_vars`` is the forcing feature dimension. :type forcing: torch.Tensor :returns: * **new_state** (*torch.Tensor*) -- Shape ``(B, num_grid_nodes, num_state_vars)``. The predicted next state ``X_{t+1}`` after delta-add and clamping. Dims: same as ``prev_state``. * **pred_std** (*torch.Tensor or None*) -- Shape ``(B, num_grid_nodes, num_state_vars)`` when ``output_std`` is True, otherwise ``None``. Per-feature predicted standard deviation (raw softplus output, not rescaled by diff statistics). Dims: same as ``prev_state``. .. py:method:: get_num_mesh() -> tuple[int, int] :abstractmethod: Compute number of mesh nodes from loaded features, and number of mesh nodes that should be ignored in encoding/decoding. :returns: * **num_mesh_nodes** (*int*) -- The number of mesh nodes. * **num_ignore_mesh_nodes** (*int*) -- The number of mesh nodes to ignore. .. py:method:: process_step(mesh_rep: torch.Tensor) -> torch.Tensor :abstractmethod: Process the mesh representation in the encode-process-decode framework, running one or more message-passing steps. :param mesh_rep: Shape ``(B, num_mesh_nodes, hidden_dim)``. Current mesh node representations. Dims: ``B`` is batch size, ``num_mesh_nodes`` is the number of mesh nodes, and ``hidden_dim`` is the hidden dimension. :type mesh_rep: torch.Tensor :returns: Shape ``(B, num_mesh_nodes, hidden_dim)``. Updated mesh node representations. Dims: same as ``mesh_rep``. :rtype: torch.Tensor .. py:attribute:: diff_mean :type: torch.Tensor .. py:attribute:: diff_std :type: torch.Tensor .. py:attribute:: encoding_grid_mlp .. py:attribute:: g2m_edge_index :type: torch.Tensor .. py:attribute:: g2m_embedder .. py:attribute:: g2m_features :type: torch.Tensor .. py:attribute:: g2m_gnn .. py:attribute:: g2m_gnn_type :value: 'InteractionNet' .. py:attribute:: grid_embedder .. py:attribute:: grid_input_dim .. py:attribute:: hidden_dim :value: 64 .. py:attribute:: hidden_layers :value: 1 .. py:attribute:: hierarchical .. py:attribute:: m2g_edge_index :type: torch.Tensor .. py:attribute:: m2g_embedder .. py:attribute:: m2g_features :type: torch.Tensor .. py:attribute:: m2g_gnn .. py:attribute:: m2g_gnn_type :value: 'InteractionNet' .. py:attribute:: m2m_edge_index :type: torch.Tensor | list[torch.Tensor] .. py:attribute:: m2m_features :type: torch.Tensor | list[torch.Tensor] .. py:attribute:: mesh_aggr :value: 'sum' .. py:attribute:: mesh_static_features :type: torch.Tensor | list[torch.Tensor] .. py:attribute:: mlp_blueprint_end .. py:attribute:: output_map .. py:attribute:: processor_layers :value: 4