use TimeWindowHandler to convert forecast steps to indices

Author: grassesi

src/weathergen/datasets/data_reader_base.py

@@ -188,6 +188,9 @@ def window(self, idx: TIndex) -> DTRange:
 
         return DTRange(t_start_win, t_end_win)
 
+    def get_n_steps(self, forecast_step: int) -> int:
+        return (int(self.t_window_len) * forecast_step) // int(self.t_window_step)
+
 
 @dataclass
 class ReaderData:

src/weathergen/datasets/multi_stream_data_sampler.py

@@ -243,7 +243,7 @@ def reset(self):
         # value in worker_workset()
         self.rng = np.random.default_rng(self.data_loader_rng_seed)
 
-        fsm = (
+        fsm: int = (
             self.forecast_steps[min(self.epoch, len(self.forecast_steps) - 1)]
             if self.forecast_policy != "random"
             else self.forecast_steps.max()
@@ -255,7 +255,7 @@ def reset(self):
         index_range = self.time_window_handler.get_index_range()
         idx_end = index_range.end
         # native length of datasets, independent of epoch length that has potentially been specified
-        forecast_len = (self.len_hrs * (fsm + 1)) // self.step_hrs
+        forecast_len = self.time_window_handler.get_n_steps(fsm + 1)
         idx_end -= forecast_len + self.forecast_offset
         assert idx_end > 0, "dataset size too small for forecast range"
         self.perms = np.arange(index_range.start, idx_end)
@@ -373,12 +373,10 @@ def __iter__(self):
                         for fstep in range(
                             self.forecast_offset, self.forecast_offset + forecast_dt + 1
                         ):
-                            step_forecast_dt = (
-                                idx + (self.forecast_delta_hrs * fstep) // self.step_hrs
-                            )
-                            time_win2 = self.time_window_handler.window(step_forecast_dt)
+                            forecast_idx = idx + self.time_window_handler.get_n_steps(fstep)
+                            time_win2 = self.time_window_handler.window(forecast_idx)
 
-                            rdata = ds.get_target(step_forecast_dt)
+                            rdata = ds.get_target(forecast_idx)
 
                             sample_is_empty = rdata.is_empty()
                             if sample_is_empty: