From e2601b4beaadd394d9d928c45f47b0a949b1cf12 Mon Sep 17 00:00:00 2001 From: Mario Santa Cruz Date: Fri, 13 Sep 2024 09:46:06 +0000 Subject: [PATCH] fix: edge case 1 set of nodes with 1 node attribute --- src/anemoi/graphs/plotting/displots.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/anemoi/graphs/plotting/displots.py b/src/anemoi/graphs/plotting/displots.py index 22316f5..5fe37ee 100644 --- a/src/anemoi/graphs/plotting/displots.py +++ b/src/anemoi/graphs/plotting/displots.py @@ -5,6 +5,7 @@ from typing import Union import matplotlib.pyplot as plt +import numpy as np import torch from torch_geometric.data import HeteroData from torch_geometric.data.storage import EdgeStorage @@ -83,7 +84,9 @@ def plot_distribution_attributes( # Define the layout _, axs = plt.subplots(num_items, dim_attrs, figsize=(10 * num_items, 10)) - if axs.ndim == 1: + if num_items == dim_attrs == 1: + axs = np.array([[axs]]) + elif axs.ndim == 1: axs = axs.reshape(num_items, dim_attrs) for i, (item_name, item_store) in enumerate(graph_items):