neural_lam.models.step_predictors.base ====================================== .. py:module:: neural_lam.models.step_predictors.base .. autoapi-nested-parse:: Base class for step predictors. Classes ------- .. autoapisummary:: neural_lam.models.step_predictors.base.StepPredictor Module Contents --------------- .. py:class:: StepPredictor(datastore: neural_lam.datastore.BaseDatastore, output_std: bool = False, output_clamping_lower: dict[str, float] | None = None, output_clamping_upper: dict[str, float] | None = None) Bases: :py:obj:`torch.nn.Module`, :py:obj:`abc.ABC` Abstract base class for step predictors mapping from the two previous time steps plus forcing into a prediction of the next state. Initialize the StepPredictor. :param datastore: The datastore providing grid metadata and data access. :type datastore: BaseDatastore :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:: expand_to_batch(x: torch.Tensor, batch_size: int) -> torch.Tensor Expand a shared node-feature tensor into a batch of copies. :param x: Shape ``(N, d)``. Tensor to expand. Dims: ``N`` is the number of nodes and ``d`` is the feature dimension. :type x: torch.Tensor :param batch_size: Target batch size ``B``. :type batch_size: int :returns: Shape ``(B, N, d)``. Batch-expanded view of ``x``. :rtype: torch.Tensor .. py:method:: forward(prev_state: torch.Tensor, prev_prev_state: torch.Tensor, forcing: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None] :abstractmethod: Advance the state by one step: ``(X_t, X_{t-1}, forcing_t) -> X_{t+1}``. :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 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 nodes, and ``num_forcing_vars`` is the forcing feature dimension. :type forcing: torch.Tensor :returns: * **pred_state** (*torch.Tensor*) -- Shape ``(B, num_grid_nodes, num_state_vars)``. The predicted next state ``X_{t+1}``. 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. Dims: same as ``prev_state``. .. py:method:: get_clamped_new_state(state_delta: torch.Tensor, prev_state: torch.Tensor) -> torch.Tensor Clamp the next-state prediction to its valid feature range. Instead of the plain residual update ``X_{t+1} = X_t + delta``, the clamped form is ``X_{t+1} = f(f^{-1}(X_t) + delta)`` so the model learns to output increments in the domain of the inverse clamping function. :param state_delta: Shape ``(B, num_grid_nodes, num_state_vars)``. Raw predicted state increment (network output, already rescaled). Dims: ``B`` is batch size, ``num_grid_nodes`` is the number of spatial nodes, and ``num_state_vars`` is the number of state variables. :type state_delta: torch.Tensor :param prev_state: Shape ``(B, num_grid_nodes, num_state_vars)``. Current state ``X_t`` used as the base for the clamped update. Dims: same as ``state_delta``. :type prev_state: torch.Tensor :returns: Shape ``(B, num_grid_nodes, num_state_vars)``. Clamped next state. Dims: same as ``state_delta``. :rtype: torch.Tensor .. py:method:: prepare_clamping_params(datastore: neural_lam.datastore.BaseDatastore) -> None Prepare parameters for clamping predicted values to valid range. Reads the per-feature lower/upper limits from ``self._output_clamping_lower`` and ``self._output_clamping_upper`` (set in ``__init__``) and registers the buffers and clamping functions used by ``get_clamped_new_state``. :param datastore: The datastore providing variable names. :type datastore: BaseDatastore .. py:attribute:: clamp_lower_idx :type: torch.Tensor .. py:attribute:: clamp_lower_upper_idx :type: torch.Tensor .. py:attribute:: clamp_upper_idx :type: torch.Tensor .. py:attribute:: grid_static_features :type: torch.Tensor .. py:attribute:: output_std :value: False .. py:property:: predicts_std :type: bool Whether this predictor outputs a predicted standard deviation. :returns: ``True`` if the predictor predicts standard deviation, ``False`` otherwise. :rtype: bool .. py:attribute:: sigmoid_lower_lims :type: torch.Tensor .. py:attribute:: sigmoid_upper_lims :type: torch.Tensor .. py:attribute:: softplus_lower_lims :type: torch.Tensor .. py:attribute:: softplus_upper_lims :type: torch.Tensor .. py:attribute:: state_mean :type: torch.Tensor .. py:attribute:: state_std :type: torch.Tensor .. py:property:: trainable :type: bool Whether the predictor has learnable parameters and can be trained. :returns: ``True`` if the predictor has trainable parameters, ``False`` otherwise. :rtype: bool