From 44d9793ae67f8582ca76723aeb272008d11b3673 Mon Sep 17 00:00:00 2001 From: Gaurav Sharma Date: Thu, 19 Mar 2026 19:37:11 +0530 Subject: [PATCH] docs: fix formatting errors in ARModel docstrings --- neural_lam/models/ar_model.py | 34 ++++++++++++++++++++-------------- 1 file changed, 20 insertions(+), 14 deletions(-) diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py index f1bcb461d..7b499571a 100644 --- a/neural_lam/models/ar_model.py +++ b/neural_lam/models/ar_model.py @@ -278,11 +278,14 @@ def unroll_prediction(self, init_states, forcing_features, true_states): def common_step(self, batch): """ - Predict on single batch batch consists of: init_states: (B, 2, - num_grid_nodes, d_features) target_states: (B, pred_steps, - num_grid_nodes, d_features) forcing_features: (B, pred_steps, - num_grid_nodes, d_forcing), - where index 0 corresponds to index 1 of init_states + Predict on single batch. + + Batch consists of: + init_states: (B, 2, num_grid_nodes, d_features) + target_states: (B, pred_steps, num_grid_nodes, d_features) + forcing_features: (B, pred_steps, num_grid_nodes, d_forcing) + + Note: index 0 corresponds to index 1 of init_states. """ (init_states, target_states, forcing_features, batch_times) = batch @@ -465,12 +468,14 @@ def test_step(self, batch, batch_idx): def plot_examples(self, batch, n_examples, split, prediction=None): """ - Plot the first n_examples forecasts from batch + Plot the first n_examples forecasts from batch. - batch: batch with data to plot corresponding forecasts for n_examples: - number of forecasts to plot prediction: (B, pred_steps, num_grid_nodes, - d_f), existing prediction. - Generate if None. + Args: + batch: batch with data to plot corresponding forecasts for. + n_examples: number of forecasts to plot. + split: the data split to evaluate. + prediction: (B, pred_steps, num_grid_nodes, d_f), existing prediction. + Generate if None. """ if prediction is None: prediction, target, _, _ = self.common_step(batch) @@ -635,11 +640,12 @@ def create_metric_log_dict(self, metric_tensor, prefix, metric_name): def aggregate_and_plot_metrics(self, metrics_dict, prefix): """ - Aggregate and create error map plots for all metrics in metrics_dict + Aggregate and create error map plots for all metrics in metrics_dict. - metrics_dict: dictionary with metric_names and list of tensors - with step-evals. - prefix: string, prefix to use for logging + Args: + metrics_dict: dictionary with metric_names and list of tensors + with step-evals. + prefix: string, prefix to use for logging. """ log_dict = {} for metric_name, metric_val_list in metrics_dict.items():