Skip to content

docs: fix formatting errors in ARModel docstrings #442

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
alphaleporus wants to merge 1 commit into from
Closed

docs: fix formatting errors in ARModel docstrings #442

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 20 additions & 14 deletions neural_lam/models/ar_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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():
Expand Down