Merge pull request #59 from ecmwf/162-datashader-does-not-handle-well-nans-in-ocean-plots

fix(training, plots) Exclude nans from error colorbars

Author: anaprietonem

training/src/anemoi/training/diagnostics/plots.py

@@ -584,6 +584,7 @@ def error_plot_in_degrees(array1: np.ndarray, array2: np.ndarray) -> np.ndarray:
         # For 'errors', only persistence and increments need identical colorbar-limits
         combined_error = np.concatenate(((pred - input_), (truth - input_)))
         norm = Normalize(vmin=np.nanmin(combined_data), vmax=np.nanmax(combined_data))
+        norm_error = TwoSlopeNorm(vmin=np.nanmin(combined_error), vcenter=0.0, vmax=np.nanmax(combined_error))
         single_plot(fig, ax[1], lon, lat, truth, norm=norm, title=f"{vname} target", datashader=datashader)
         single_plot(fig, ax[2], lon, lat, pred, norm=norm, title=f"{vname} pred", datashader=datashader)
         single_plot(
@@ -688,7 +689,7 @@ def error_plot_in_degrees(array1: np.ndarray, array2: np.ndarray) -> np.ndarray:
                 lat,
                 pred - input_,
                 cmap="bwr",
-                norm=TwoSlopeNorm(vmin=combined_error.min(), vcenter=0.0, vmax=combined_error.max()),
+                norm=norm_error,
                 title=f"{vname} increment [pred - input]",
                 datashader=datashader,
             )
@@ -699,7 +700,7 @@ def error_plot_in_degrees(array1: np.ndarray, array2: np.ndarray) -> np.ndarray:
                 lat,
                 truth - input_,
                 cmap="bwr",
-                norm=TwoSlopeNorm(vmin=combined_error.min(), vcenter=0.0, vmax=combined_error.max()),
+                norm=norm_error,
                 title=f"{vname} persist err",
                 datashader=datashader,
             )