Comment out resolution check and fix precip plot

MeteoSwiss · Nov 21, 2024 · 0e6c00a · 0e6c00a
1 parent 4b564e9
commit 0e6c00a
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 2 deletions.
diff --git a/src/anemoi/training/data/datamodule.py b/src/anemoi/training/data/datamodule.py
@@ -189,7 +189,7 @@ def _get_dataset(
             shuffle=shuffle,
             label=label,
         )
-        self._check_resolution(data.resolution)
+        #self._check_resolution(data.resolution)
         return data
 
     def _get_dataloader(self, ds: NativeGridDataset, stage: str) -> DataLoader:

diff --git a/src/anemoi/training/diagnostics/plots.py b/src/anemoi/training/diagnostics/plots.py
@@ -535,7 +535,41 @@ def error_plot_in_degrees(array1: np.ndarray, array2: np.ndarray) -> np.ndarray:
         )
 
     if sum(input_) != 0:
-        if vname == "mwd":
+        if vname in precip_and_related_fields:
+            # Create a custom colormap for precipitation
+            nws_precip_colors = cmap_precip
+            precip_colormap = ListedColormap(nws_precip_colors)
+
+            # Defining the actual precipitation accumulation levels in mm
+            cummulation_lvls = clevels
+            norm = BoundaryNorm(cummulation_lvls, len(cummulation_lvls) + 1)
+
+            # converting to mm from m
+            input_ *= 1000.0
+            truth *= 1000.0
+            pred *= 1000.0
+            scatter_plot(fig, ax[0], lon=lon, lat=lat, data=input_, cmap=precip_colormap, norm=norm, title=f"{vname} input")
+            scatter_plot(
+                fig,
+                ax[4],
+                lon=lon,
+                lat=lat,
+                data=pred - input_,
+                cmap="bwr",
+                norm=TwoSlopeNorm(vcenter=0.0),
+                title=f"{vname} increment [pred - input]",
+            )
+            scatter_plot(
+                fig,
+                ax[5],
+                lon=lon,
+                lat=lat,
+                data=truth - input_,
+                cmap="bwr",
+                norm=TwoSlopeNorm(vcenter=0.0),
+                title=f"{vname} persist err",
+            )
+        elif vname == "mwd":
             scatter_plot(fig, ax[0], lon=lon, lat=lat, data=input_, cmap=cyclic_colormap, title=f"{vname} input")
             err_plot = error_plot_in_degrees(pred, input_)
             scatter_plot(