Fixed bug in obs data reading (ecmwf#698)

* Restored old prediction had functionally. Other adjustments/reverts, in particular in attention.

* Ruff'ed

* Fixed bug in obs data reading so that data violated window

* Fix

* Update data_reader_obs.py

* Restoring to develop

* Fix

* Ruffed

Author: clessig

src/weathergen/datasets/data_reader_base.py

@@ -254,11 +254,9 @@ def check_reader_data(rdata: ReaderData, dtr: DTRange) -> None:
         f"{rdata.datetimes.shape[0]}"
     )
 
-    assert np.logical_and(
-        rdata.datetimes >= dtr.start,
-        # rdata.datetimes < dtr.end  # TODO: enforce monotonicty also for obs
-        rdata.datetimes <= dtr.end,
-    ).all(), f"datetimes for data points violate window {dtr}."
+    assert np.logical_and(rdata.datetimes >= dtr.start, rdata.datetimes < dtr.end).all(), (
+        f"datetimes for data points violate window {dtr}."
+    )
 
 
 class DataReaderBase(metaclass=ABCMeta):

src/weathergen/datasets/data_reader_obs.py

@@ -225,7 +225,7 @@ def _get(self, idx: int, channels_idx: list[int]) -> ReaderData:
                 num_data_fields=len(channels_idx), num_geo_fields=len(self.geoinfo_idx)
             )
 
-        start_row = self.indices_start[idx]
+        start_row = self.indices_start[idx - 1]
         end_row = self.indices_end[idx]
 
         coords = self.data.oindex[start_row:end_row, self.coords_idx]
@@ -238,11 +238,17 @@ def _get(self, idx: int, channels_idx: list[int]) -> ReaderData:
         data = self.data.oindex[start_row:end_row, channels_idx]
         datetimes = self.dt[start_row:end_row][:, 0]
 
+        # indices_start, indices_end above work with [t_start, t_end] and violate
+        # our convention [t_start, t_end) where endpoint is excluded
+        # compute mask to enforce it
+        t_win = self.time_window_handler.window(idx)
+        t_mask = np.logical_and(datetimes >= t_win.start, datetimes < t_win.end)
+
         rdata = ReaderData(
-            coords=coords,
-            geoinfos=geoinfos,
-            data=data,
-            datetimes=datetimes,
+            coords=coords[t_mask],
+            geoinfos=geoinfos[t_mask],
+            data=data[t_mask],
+            datetimes=datetimes[t_mask],
         )
 
         dtr = self.time_window_handler.window(idx)