docs: standardize docstrings for ARModel init

neural_lam/metrics.py

@@ -86,21 +86,30 @@ def wmse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
 
 def mse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     """
-    (Unweighted) Mean Squared Error
-
-    (...,) is any number of batch dimensions, potentially different
-        but broadcastable
-    pred: (..., N, d_state), prediction
-    target: (..., N, d_state), target
-    pred_std: (..., N, d_state) or (d_state,), predicted std.-dev.
-    mask: (N,), boolean mask describing which grid nodes to use in metric
-    average_grid: boolean, if grid dimension -2 should be reduced (mean over N)
-    sum_vars: boolean, if variable dimension -1 should be reduced (sum
-        over d_state)
-
-    Returns:
-    metric_val: One of (...,), (..., d_state), (..., N), (..., N, d_state),
-    depending on reduction arguments.
+    (Unweighted) Mean Squared Error.
+
+    Parameters
+    ----------
+    pred : torch.Tensor
+        Prediction tensor of shape (..., N, d_state).
+    target : torch.Tensor
+        Target tensor of shape (..., N, d_state).
+    pred_std : torch.Tensor
+        Predicted standard deviation of shape (..., N, d_state) or (d_state,).
+    mask : torch.Tensor, optional
+        Boolean mask of shape (N,) describing which grid nodes to use, 
+        by default None.
+    average_grid : bool, optional
+        Whether the grid dimension -2 should be reduced (mean over N), 
+        by default True.
+    sum_vars : bool, optional
+        Whether the variable dimension -1 should be reduced (sum over d_state), 
+        by default True.
+
+    Returns
+    -------
+    torch.Tensor
+        Metric value(s) depending on reduction arguments.
     """
     # Replace pred_std with constant ones
     return wmse(

neural_lam/models/ar_model.py

@@ -36,6 +36,20 @@ def __init__(
         config: NeuralLAMConfig,
         datastore: BaseDatastore,
     ):
+        """
+        Initialize the Auto-Regressive model.
+
+        Parameters
+        ----------
+        args : argparse.Namespace
+            Command-line arguments containing model hyperparameters and 
+            run configurations.
+        config : NeuralLAMConfig
+            Configuration object for the NeuralLAM project.
+        datastore : BaseDatastore
+            Datastore object used to retrieve weather data and 
+            standardization statistics.
+        """
         super().__init__()
         self.save_hyperparameters(ignore=["datastore"])
         self.args = args
@@ -220,10 +234,22 @@ def expand_to_batch(x, batch_size):
 
     def predict_step(self, prev_state, prev_prev_state, forcing):
         """
-        Step state one step ahead using prediction model, X_{t-1}, X_t -> X_t+1
-        prev_state: (B, num_grid_nodes, feature_dim), X_t prev_prev_state: (B,
-        num_grid_nodes, feature_dim), X_{t-1} forcing: (B, num_grid_nodes,
-        forcing_dim)
+        Step state one step ahead using the prediction model.
+        Computes X_{t+1} given X_t and X_{t-1}.
+
+        Parameters
+        ----------
+        prev_state : torch.Tensor
+            State at time t (X_t), shape (B, num_grid_nodes, feature_dim).
+        prev_prev_state : torch.Tensor
+            State at time t-1 (X_{t-1}), shape (B, num_grid_nodes, feature_dim).
+        forcing : torch.Tensor
+            Forcing features at time t, shape (B, num_grid_nodes, forcing_dim).
+
+        Raises
+        ------
+        NotImplementedError
+            This is an abstract method that must be implemented by subclasses.
         """
         raise NotImplementedError("No prediction step implemented")