Iluise/fix empty io 819 plotting (ecmwf#826)

* Fix to io problems.

* Fix issues in input

* fix plotting

* ruffed

---------

Co-authored-by: Christian Lessig <christian.lessig@ecmwf.int>

Author: iluise

packages/common/src/weathergen/common/io.py

@@ -436,9 +436,11 @@ def _offset_key(self, key: ItemKey):
 
     def _extract_coordinates(self, stream_idx, offset_key, datapoints) -> DataCoordinates:
         _coords = self.targets_coords[offset_key.forecast_step][stream_idx][datapoints].numpy()
+
         # ensure _coords has size (?,2)
         if len(_coords) == 0:
             _coords = np.zeros((0, 2), dtype=np.float32)
+
         coords = _coords[..., :2]  # first two columns are lat,lon
         geoinfo = _coords[..., 2:]  # the rest is geoinfo => potentially empty
         if geoinfo.size > 0:  # TODO: set geoinfo to be empty for now

packages/evaluate/src/weathergen/evaluate/plotter.py

@@ -212,7 +212,9 @@ def create_histograms_per_sample(
 
             for (valid_time, targ_t), (_, prd_t) in groups:
                 if valid_time is not None:
-                    _logger.debug(f"Plotting map for {var} at valid_time {valid_time}")
+                    _logger.debug(
+                        f"Plotting histogram for {var} at valid_time {valid_time}"
+                    )
 
                 name = self.plot_histogram(targ_t, prd_t, hist_output_dir, var, tag=tag)
                 plot_names.append(name)
@@ -341,7 +343,6 @@ def create_maps_per_sample(
                 self.stream
             )
 
-
         # Basic map output directory for this stream
         map_output_dir = self.get_map_output_dir(tag)
 

packages/evaluate/src/weathergen/evaluate/utils.py

@@ -116,6 +116,8 @@ def get_data(
         else:
             points_per_sample = None
 
+        fsteps_final = fsteps
+
         for fstep in fsteps:
             _logger.info(f"RUN {run_id} - {stream}: Processing fstep {fstep}...")
             da_tars_fs, da_preds_fs = [], []
@@ -135,6 +137,12 @@ def get_data(
                     pred = bbox.apply_mask(pred)
 
                 npoints = len(target.ipoint)
+                if npoints == 0:
+                    _logger.info(
+                        f"Skipping {stream} sample {sample} forecast step: {fstep}. Dataset is empty."
+                    )
+                    fsteps_final.remove(fstep)
+                    continue
 
                 da_tars_fs.append(target.squeeze())
                 da_preds_fs.append(pred.squeeze())
@@ -143,31 +151,37 @@ def get_data(
             _logger.debug(
                 f"Concatenating targets and predictions for stream {stream}, forecast_step {fstep}..."
             )
-            da_tars_fs = xr.concat(da_tars_fs, dim="ipoint")
-            da_preds_fs = xr.concat(da_preds_fs, dim="ipoint")
 
-            if set(channels) != set(all_channels):
-                _logger.debug(
-                    f"Restricting targets and predictions to channels {channels} for stream {stream}..."
-                )
-                available_channels = da_tars_fs.channel.values
-                existing_channels = [ch for ch in channels if ch in available_channels]
-                if len(existing_channels) < len(channels):
-                    _logger.warning(
-                        f"The following channels were not found: {list(set(channels) - set(existing_channels))}. Skipping them."
+            if da_tars_fs:
+                da_tars_fs = xr.concat(da_tars_fs, dim="ipoint")
+                da_preds_fs = xr.concat(da_preds_fs, dim="ipoint")
+
+                if set(channels) != set(all_channels):
+                    _logger.debug(
+                        f"Restricting targets and predictions to channels {channels} for stream {stream}..."
                     )
+                    available_channels = da_tars_fs.channel.values
+                    existing_channels = [
+                        ch for ch in channels if ch in available_channels
+                    ]
+                    if len(existing_channels) < len(channels):
+                        _logger.warning(
+                            f"The following channels were not found: {list(set(channels) - set(existing_channels))}. Skipping them."
+                        )
 
-                da_tars_fs = da_tars_fs.sel(channel=existing_channels)
-                da_preds_fs = da_preds_fs.sel(channel=existing_channels)
+                    da_tars_fs = da_tars_fs.sel(channel=existing_channels)
+                    da_preds_fs = da_preds_fs.sel(channel=existing_channels)
 
             da_tars.append(da_tars_fs)
             da_preds.append(da_preds_fs)
             if return_counts:
                 points_per_sample.loc[{"forecast_step": fstep}] = np.array(pps)
 
         # Safer than a list
-        da_tars = {fstep: da for fstep, da in zip(fsteps, da_tars, strict=False)}
-        da_preds = {fstep: da for fstep, da in zip(fsteps, da_preds, strict=False)}
+        da_tars = {fstep: da for fstep, da in zip(fsteps_final, da_tars, strict=False)}
+        da_preds = {
+            fstep: da for fstep, da in zip(fsteps_final, da_preds, strict=False)
+        }
 
         return WeatherGeneratorOutput(
             target=da_tars, prediction=da_preds, points_per_sample=points_per_sample
@@ -396,6 +410,10 @@ def plot_data(
     da_tars = model_output.target
     da_preds = model_output.prediction
 
+    if not da_tars:
+        _logger.info(f"Skipping Plot Data for {stream}. Targets are empty.")
+        return
+
     maps_config = common_ranges(da_tars, da_preds, plot_chs, maps_config)
 
     plot_names = []