Callback PlotHistogram breaks if only one variable is specified. (ecm…

…wf#165)

* enable plothistogram and plotspectrum for only one variable

CHANGELOG.md

@@ -11,7 +11,8 @@ Keep it human-readable, your future self will thank you!
 ## [Unreleased](https://github.com/ecmwf/anemoi-training/compare/0.3.0...HEAD)
 ### Fixed
 
-Fixed bug in power spectra plotting for the n320 resolution.
+- Fixed bug in power spectra plotting for the n320 resolution.
+- Allow histogram and spectrum plot for one variable [#165](https://github.com/ecmwf/anemoi-training/pull/165)
 
 ### Added
 - Introduce variable to configure (Cosine Annealing) optimizer warm up [#155](https://github.com/ecmwf/anemoi-training/pull/155)

src/anemoi/training/diagnostics/plots.py

@@ -166,6 +166,8 @@ def plot_power_spectrum(
 
     figsize = (n_plots_y * 4, n_plots_x * 3)
     fig, ax = plt.subplots(n_plots_x, n_plots_y, figsize=figsize, layout=LAYOUT)
+    if n_plots_x == 1:
+        ax = [ax]
 
     pc_lat, pc_lon = equirectangular_projection(latlons)
 
@@ -293,6 +295,8 @@ def plot_histogram(
 
     figsize = (n_plots_y * 4, n_plots_x * 3)
     fig, ax = plt.subplots(n_plots_x, n_plots_y, figsize=figsize, layout=LAYOUT)
+    if n_plots_x == 1:
+        ax = [ax]
 
     for plot_idx, (variable_idx, (variable_name, output_only)) in enumerate(parameters.items()):
         yt = y_true[..., variable_idx].squeeze()