9 various tidy-ups and loud failing

README.md

@@ -22,7 +22,7 @@ It includes and adapts the `ephemeris()` function from the [pvlib-python](https:
 
 ## Documentation
 
-**ocf-data-sampler** doesn't have external documentation _yet_; you can read a bit about how our torch datasets work in the README [here](ocf_data_sampler/torch_datasets/README.md).
+**ocf-data-sampler** doesn't have external documentation _yet_; you can read a bit about how our torch datasets work in the README [here](src/ocf_data_sampler/datasets/pvnet/README.md).
 
 ## FAQ
 

src/ocf_data_sampler/config/model.py

@@ -1,6 +1,6 @@
 """Configuration model for the PVNet dataset."""
 
-from collections.abc import Iterator
+from collections.abc import ItemsView, Iterator, KeysView
 from typing import Literal
 
 from pydantic import BaseModel, ConfigDict, Field, RootModel, field_validator, model_validator
@@ -63,7 +63,8 @@ class FillValueMixin(Base):
 
     dropout_fill_value: float = Field(
         default=0.0,
-        description="The value used to fill in dropped out data or any missing values."
+        description="The value used to fill in dropped out data or any missing values. This is "
+        "applied after normalisation, so it is in normalised units."
     )
 
 
@@ -280,13 +281,12 @@ def check_all_channel_have_normalisation_constants(self) -> "NWP":
 
 
 class MultiNWP(RootModel):
-    """Configuration for multiple NWPs."""
+    """Configuration for multiple NWPs.
 
-    root: dict[str, NWP]
+    The NWP sources are accessed dict-style - `config.nwp["ukv"]`, not `config.nwp.ukv`.
+    """
 
-    @override
-    def __getattr__(self, item: str) -> NWP:
-        return self.root[item]
+    root: dict[str, NWP]
 
     @override
     def __getitem__(self, item: str) -> NWP:
@@ -300,11 +300,11 @@ def __len__(self) -> int:
     def __iter__(self) -> Iterator:
         return iter(self.root)
 
-    def keys(self) -> Iterator[str]:
+    def keys(self) -> KeysView[str]:
         """Returns dictionary-like keys."""
         return self.root.keys()
 
-    def items(self) -> Iterator[tuple[str, NWP]]:
+    def items(self) -> ItemsView[str, NWP]:
         """Returns dictionary-like items."""
         return self.root.items()
 

src/ocf_data_sampler/datasets/cache.py

@@ -6,13 +6,22 @@
 class PickleCacheMixin:
     """A mixin for classes that need to cache their state using pickle."""
 
-    def __init__(self, *args: list, **kwargs: dict) -> None:
+    def __init__(self, *args: object, **kwargs: object) -> None:
         """Initialize the pickle path and call the parent constructor."""
         self._pickle_path = None
         super().__init__(*args, **kwargs)  # cooperative multiple inheritance
 
     def presave_pickle(self, pickle_path: str) -> None:
-        """Save the full object state to a pickle file and store the pickle path."""
+        """Save the full object state to a pickle file and store the pickle path.
+
+        The object will be pickled by reference after calling this function, so the
+        file must be readable wherever it is unpickled.
+
+        The saved state is a snapshot - call this again after mutating the object.
+
+        Args:
+            pickle_path: Where to write the object state to.
+        """
         self._pickle_path = pickle_path
         with open(pickle_path, "wb") as f:
             pickle.dump(self.__dict__, f)
@@ -22,12 +31,23 @@ def __getstate__(self) -> dict:
         if self._pickle_path:
             return {"_pickle_path": self._pickle_path}
         else:
-            return self.__dict__
+            # Copied so that the pickled state can't be mutated through the live object
+            return dict(self.__dict__)
 
     def __setstate__(self, state: dict) -> None:
         """Restore object from pickle, reloading from presaved file if possible."""
         self.__dict__.update(state)
-        if self._pickle_path and os.path.exists(self._pickle_path):
-            with open(self._pickle_path, "rb") as f:
-                saved_state = pickle.load(f)  # noqa: S301
-                self.__dict__.update(saved_state)
+
+        if not self._pickle_path:
+            return
+
+        if not os.path.exists(self._pickle_path):
+            raise FileNotFoundError(
+                f"Presaved state file not found: {self._pickle_path}. This object was pickled by "
+                "reference to that path - see `presave_pickle`.",
+            )
+
+        with open(self._pickle_path, "rb") as f:
+            saved_state = pickle.load(f)  # noqa: S301
+
+        self.__dict__.update(saved_state)

src/ocf_data_sampler/datasets/pvnet/README.md

@@ -53,3 +53,13 @@ graph TD
     E0[Extra processing and add features like sun/time encodings]
     E0 --> F[Sample]
 ```
+
+### Randomness
+
+Dropout time-slice choices draw from **global** NumPy state (`np.random`). Under a PyTorch
+`DataLoader` this is seeded per worker automatically, so a run is reproducible from a single
+global seed — with Lightning, use `seed_everything(seed, workers=True)`. 
+
+If you use the dataset outside a `DataLoader`, you should seed `np.random` yourself. Note that 
+a given index (i.e. t0 and location pair) does always choose the same dropout. The dropout is
+chosen base don the global random state at the time of sampling.

src/ocf_data_sampler/datasets/pvnet/dataset.py

@@ -98,8 +98,10 @@ def get_time_periods_mask(
 ) -> np.ndarray:
     """Get a boolean mask showing which times fall within any of the specified time periods.
 
+    A `None` bound means the period is unbounded in that direction.
+
     Args:
-        times: DatetimeIndex of times to filter
+        times: Array of times to filter
         time_periods: List of tuples specifying the start and end times for each period
     """
     if len(time_periods)==0:
@@ -109,10 +111,15 @@ def get_time_periods_mask(
 
     for start_time, end_time in time_periods:
 
-        start_time = times[0] if start_time is None else np.datetime64(start_time)
-        end_time = times[-1] if end_time is None else np.datetime64(end_time)
+        this_period_mask = np.full(len(times), True)
+
+        # Inclusive of start_time, exclusive of end_time
+        if start_time is not None:
+            this_period_mask &= times >= np.datetime64(start_time)
+        if end_time is not None:
+            this_period_mask &= times < np.datetime64(end_time)
 
-        mask |= (times >= start_time) & (times <= end_time)
+        mask |= this_period_mask
 
     return mask
 
@@ -291,6 +298,8 @@ def __init__(
             # Filter t0 times to given range
             if time_periods is not None:
                 mask = get_time_periods_mask(valid_t0_times, time_periods)
+                if not mask.any():
+                    raise ValueError(f"`time_periods` {time_periods} excluded all valid t0 times")
                 valid_t0_times = valid_t0_times[mask]
 
             self.valid_t0_times = valid_t0_times
@@ -305,8 +314,10 @@ def __init__(
 
             # Filter t0 times to given range
             if time_periods is not None:
-                mask = get_time_periods_mask(valid_t0_and_location_ids["t0"], time_periods)
-                valid_t0_and_location_ids = valid_t0_and_location_ids[mask]
+                mask = get_time_periods_mask(valid_t0_and_location_ids["t0"].values, time_periods)
+                if not mask.any():
+                    raise ValueError(f"`time_periods` {time_periods} excluded all valid t0 times")
+                valid_t0_and_location_ids = valid_t0_and_location_ids.iloc[mask]
 
             self.valid_t0_and_location_ids = valid_t0_and_location_ids
 
@@ -362,6 +373,13 @@ def find_valid_t0_times(
             valid_time_periods,
             freq=minutes(config.sampling_grid.t0_resolution_minutes),
         )
+
+        if len(valid_t0_times) == 0:
+            raise ValueError(
+                f"The inputs overlap, but no period is long enough to contain a t0 on the "
+                f"{config.sampling_grid.t0_resolution_minutes}-minute grid:\n{valid_time_periods}",
+            )
+
         return valid_t0_times
 
     @staticmethod
@@ -381,15 +399,19 @@ def find_valid_t0_and_location_ids(
             locations: The locations to find valid t0 times for
             config: PVNetDataConfig file
         """
-        # Get valid time period for nwp and satellite
-        datasets_without_generation = {k: v for k, v in datasets_dict.items() if k != "generation"}
-        valid_time_periods = find_valid_time_periods(datasets_without_generation, config)
+        # Get valid time periods for inputs other than generation
+        non_gen_time_periods = find_valid_time_periods(
+            datasets_dict={k: v for k, v in datasets_dict.items() if k != "generation"},
+            config=config,
+        )
 
-        # Loop over each location in system id and obtain valid periods
+        # There are separate input and target generation slices
         generation_windows = [
             w for w in (config.generation.input, config.generation.target) if w is not None
         ]
-        valid_t0_and_location_ids = []
+
+        # Loop over each location in system id and obtain valid periods
+        valid_t0_and_location_ids: list[pd.DataFrame] = []
         for location in locations:
             # Drop NaN values for location
             generation = (
@@ -398,8 +420,8 @@ def find_valid_t0_and_location_ids(
                 .dropna(dim="time_utc")
             )
 
-            # Obtain valid time periods for this location, for each configured window
-            time_periods_per_window = [
+            # Obtain valid time periods for this location for both input and target generation
+            gen_time_periods = [
                 find_contiguous_t0_periods(
                     generation["time_utc"].values,
                     time_resolution=minutes(config.generation.time_resolution_minutes),
@@ -408,22 +430,29 @@ def find_valid_t0_and_location_ids(
                 )
                 for window_config in generation_windows
             ]
-            valid_time_periods_per_location = intersect_time_periods(
-                [valid_time_periods, *time_periods_per_window],
+            valid_time_periods = intersect_time_periods(
+                [non_gen_time_periods, *gen_time_periods],
             )
 
             # Fill out contiguous time periods to get t0 times
-            valid_t0_times_per_location = fill_time_periods(
-                valid_time_periods_per_location,
+            valid_t0_times = fill_time_periods(
+                valid_time_periods,
                 freq=minutes(config.sampling_grid.t0_resolution_minutes),
             )
 
-            valid_t0_per_location = pd.DataFrame(index=valid_t0_times_per_location)
-            valid_t0_per_location["location_id"] = location.id
-            valid_t0_and_location_ids.append(valid_t0_per_location)
+            if len(valid_t0_times) == 0:
+                logger.warning(f"No valid t0 times found for location {location.id}")
+
+            valid_t0_and_location_ids.append(
+                pd.DataFrame({"t0": valid_t0_times, "location_id": location.id})
+            )
+
+        all_valid_t0_and_location_ids = pd.concat(valid_t0_and_location_ids, ignore_index=True)
+
+        if len(all_valid_t0_and_location_ids) == 0:
+            raise ValueError("No valid t0 times found for any location")
 
-        valid_t0_and_location_ids = pd.concat(valid_t0_and_location_ids)
-        return valid_t0_and_location_ids.reset_index(names="t0")
+        return all_valid_t0_and_location_ids
 
 
 class PVNetDataset(AbstractPVNetDataset):

src/ocf_data_sampler/datasets/pvnet/valid_t0s.py

@@ -62,9 +62,6 @@ def find_valid_time_periods(
                 max_staleness=max_staleness,
             )
 
-            if len(time_periods) == 0:
-                raise ValueError(f"No valid t0 periods found for {nwp_key} NWP data")
-
             contiguous_time_periods.append(time_periods)
 
     if "sat" in datasets_dict:
@@ -77,14 +74,8 @@ def find_valid_time_periods(
 
         contiguous_time_periods.append(time_periods)
 
-        if len(time_periods) == 0:
-            raise ValueError("No valid t0 periods found for satellite data")
-
     if "generation" in datasets_dict:
-        for window_name, window_config in (
-            ("input", config.generation.input),
-            ("target", config.generation.target),
-        ):
+        for window_config in (config.generation.input, config.generation.target):
             if window_config is None:
                 continue
 
@@ -95,11 +86,6 @@ def find_valid_time_periods(
                 interval_end=minutes(window_config.interval_end_minutes),
             )
 
-            if len(time_periods) == 0:
-                raise ValueError(
-                    f"No valid t0 periods found for {window_name} generation data",
-                )
-
             contiguous_time_periods.append(time_periods)
 
     # Find joint overlapping contiguous time periods

src/ocf_data_sampler/features/time_encodings.py

@@ -27,14 +27,16 @@ def encode_datetimes(datetimes: NDArray[np.datetime64]) -> dict[str, NDArray[np.
     day_fraction = get_day_fraction(datetimes)
     day_of_year = get_day_of_year(datetimes)
 
+    days_in_year = 365 + get_is_leap_year(datetimes).astype(int)
+
     time_in_radians = (2 * np.pi) * day_fraction
-    date_in_radians = (2 * np.pi) * (day_of_year / 365)
+    date_in_radians = (2 * np.pi) * ((day_of_year - 1) / days_in_year)
 
     return {
-        "date_sin": np.sin(date_in_radians),
-        "date_cos": np.cos(date_in_radians),
-        "time_sin": np.sin(time_in_radians),
-        "time_cos": np.cos(time_in_radians),
+        "date_sin": np.sin(date_in_radians).astype(np.float32),
+        "date_cos": np.cos(date_in_radians).astype(np.float32),
+        "time_sin": np.sin(time_in_radians).astype(np.float32),
+        "time_cos": np.cos(time_in_radians).astype(np.float32),
     }
 
 

src/ocf_data_sampler/load/nwp.py

@@ -5,6 +5,7 @@
 
 import numpy as np
 import xarray as xr
+from numpy.typing import NDArray
 
 from ocf_data_sampler.common.indexing import assert_values_unique_increasing
 from ocf_data_sampler.common.xr_tensorstore import ZarrSource, open_zarr_paths
@@ -58,6 +59,7 @@ def open_nwp(zarr_path: ZarrSource, provider: str) -> xr.DataArray:
 
     assert_values_unique_increasing(ds["init_time_utc"].values, "init_time_utc")
     assert_values_unique_increasing(ds["step"].values, "step")
+    _assert_steps_uniformly_spaced(ds["step"].values, provider)
 
     da = extract_single_data_array(ds)
     da = da.transpose("init_time_utc", "step", "channel", x_coord, y_coord)
@@ -68,6 +70,22 @@ def open_nwp(zarr_path: ZarrSource, provider: str) -> xr.DataArray:
     return da
 
 
+def _assert_steps_uniformly_spaced(steps: NDArray[np.timedelta64], provider: str) -> None:
+    """Assert that the forecast steps are evenly spaced.
+
+    Args:
+        steps: The forecast steps, already checked to be unique and increasing.
+        provider: The NWP provider name, used in the error message.
+    """
+    spacings = np.unique(np.diff(steps))
+
+    if len(spacings) > 1:
+        raise ValueError(
+            f"NWP steps for provider {provider!r} must be evenly spaced, but the steps are "
+            f"spaced by {spacings}. Mixed step spacing is not supported.",
+        )
+
+
 def _rename(ds: xr.Dataset, name_mapping: Mapping[str, str]) -> xr.Dataset:
     """Renames coordinates in the dataset based on the provided mapping.
 

src/ocf_data_sampler/select/spatial_slice.py

@@ -166,6 +166,9 @@ def select_spatial_slice_pixels_multiple(
     Returns:
         The selected DataArray-like slice.
     """
+    if len(locations) == 0:
+        raise ValueError("`locations` is empty - there is no region to cover")
+
     target_coords, x_dim, y_dim = find_coord_system(da)
 
     x_values = da[x_dim].values

src/ocf_data_sampler/select/time_periods.py

@@ -150,6 +150,14 @@ def find_contiguous_t0_periods_nwp(
     else:
         init_end_timedelta = min(last_forecast_step - interval_end, max_staleness)
 
+    # Each init-time is unusable before init_start_timedelta and after init_end_timedelta. If they
+    # are the wrong way round then no t0 can ever use it, and every period below will be empty
+    if init_end_timedelta < init_start_timedelta:
+        raise ValueError(
+            f"Init-times are usable from {init_start_timedelta} after init-time, but only until "
+            f"{init_end_timedelta}, so no t0s are available.",
+        )
+
     # Store contiguous periods
     contiguous_periods: list[list[np.datetime64]] = []