neural_lam.models.forecasters.autoregressive
============================================

.. py:module:: neural_lam.models.forecasters.autoregressive

.. autoapi-nested-parse::

   Forecaster that uses an auto-regressive strategy to unroll a forecast.



Classes
-------

.. autoapisummary::

   neural_lam.models.forecasters.autoregressive.ARForecaster


Module Contents
---------------

.. py:class:: ARForecaster(predictor: neural_lam.models.step_predictors.base.StepPredictor, datastore: neural_lam.datastore.BaseDatastore)

   Bases: :py:obj:`neural_lam.models.forecasters.base.Forecaster`


   Subclass of Forecaster that uses an auto-regressive strategy to
   unroll a forecast. Makes use of a StepPredictor at each AR step.

   Initialize the ARForecaster.

   :param predictor: The predictor to use for each step.
   :type predictor: StepPredictor
   :param datastore: The datastore providing grid metadata and boundary masks.
   :type datastore: BaseDatastore


   .. py:method:: forward(init_states: torch.Tensor, forcing_features: torch.Tensor, boundary_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]

      Unroll the autoregressive model: at each step ``i`` call
      ``self.predictor`` to produce the next state, then overwrite boundary
      nodes with the true value from ``boundary_states[:, i]``.

      :param init_states: Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial
                          states ``[X_{t-1}, X_t]`` used to start the rollout from. Dims:
                          ``B`` is batch size, ``2`` initial time steps (``X_{t-1}, X_t``),
                          ``num_grid_nodes`` is the number of spatial nodes, and
                          ``num_state_vars`` is the number of state variables.
      :type init_states: torch.Tensor
      :param forcing_features: Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``.
                               Forcing features for each predicted step; ``pred_steps`` defines
                               the rollout length. Dims: ``B`` is batch size, ``pred_steps`` is
                               the number of predicted steps, ``num_grid_nodes`` is the
                               number of spatial nodes, and ``num_forcing_vars`` is the number
                               of forcing variables (already concatenated past/current/future
                               windows).
      :type forcing_features: torch.Tensor
      :param boundary_states: Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``.
                              True state values used ONLY to overwrite boundary nodes at
                              each AR step. The interior prediction at step ``i`` must not
                              depend on ``boundary_states[:, i]`` in any other way. Dims:
                              ``B`` is batch size, ``pred_steps`` is the number of
                              predicted steps, ``num_grid_nodes`` is the number of spatial
                              nodes, and
                              ``num_state_vars`` is the state feature dimension.
      :type boundary_states: torch.Tensor

      :returns: * **prediction** (*torch.Tensor*) -- Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. Stacked
                  per-step forecasts (with boundary overwritten by the true
                  value). Dims: same as ``boundary_states``.
                * **pred_std** (*torch.Tensor or None*) -- Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)`` when the
                  wrapped predictor outputs an std, otherwise ``None`` (in which
                  case ``ForecasterModule`` substitutes the constant
                  per-variable std). Dims: same as ``prediction``.



   .. py:attribute:: boundary_mask
      :type:  torch.Tensor


   .. py:attribute:: interior_mask
      :type:  torch.Tensor


   .. py:attribute:: predictor


   .. py:property:: predicts_std
      :type: bool


      Whether the forecaster predicts standard deviation.

      :returns: ``True`` if the forecaster predicts standard deviation,
                ``False`` otherwise.
      :rtype: bool