diff --git a/README.md b/README.md index b25405de..5e0f0936 100644 --- a/README.md +++ b/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 diff --git a/src/ocf_data_sampler/config/model.py b/src/ocf_data_sampler/config/model.py index ab2e37f4..a7afd1e4 100644 --- a/src/ocf_data_sampler/config/model.py +++ b/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() diff --git a/src/ocf_data_sampler/datasets/cache.py b/src/ocf_data_sampler/datasets/cache.py index d6b6c2b2..6d06ec69 100644 --- a/src/ocf_data_sampler/datasets/cache.py +++ b/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) diff --git a/src/ocf_data_sampler/datasets/pvnet/README.md b/src/ocf_data_sampler/datasets/pvnet/README.md index 0626c52c..31342e1c 100644 --- a/src/ocf_data_sampler/datasets/pvnet/README.md +++ b/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. diff --git a/src/ocf_data_sampler/datasets/pvnet/dataset.py b/src/ocf_data_sampler/datasets/pvnet/dataset.py index 94f59a8e..d6a9cf62 100644 --- a/src/ocf_data_sampler/datasets/pvnet/dataset.py +++ b/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): diff --git a/src/ocf_data_sampler/datasets/pvnet/valid_t0s.py b/src/ocf_data_sampler/datasets/pvnet/valid_t0s.py index 7921bc97..60009890 100644 --- a/src/ocf_data_sampler/datasets/pvnet/valid_t0s.py +++ b/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 diff --git a/src/ocf_data_sampler/features/time_encodings.py b/src/ocf_data_sampler/features/time_encodings.py index 4becd969..035b14b7 100644 --- a/src/ocf_data_sampler/features/time_encodings.py +++ b/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), } diff --git a/src/ocf_data_sampler/load/nwp.py b/src/ocf_data_sampler/load/nwp.py index d65e5b07..be61a21c 100755 --- a/src/ocf_data_sampler/load/nwp.py +++ b/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. diff --git a/src/ocf_data_sampler/select/spatial_slice.py b/src/ocf_data_sampler/select/spatial_slice.py index 953da01f..5bf1bdf7 100644 --- a/src/ocf_data_sampler/select/spatial_slice.py +++ b/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 diff --git a/src/ocf_data_sampler/select/time_periods.py b/src/ocf_data_sampler/select/time_periods.py index 1aef8e63..08e2ba77 100644 --- a/src/ocf_data_sampler/select/time_periods.py +++ b/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]] = [] diff --git a/tests/datasets/pvnet/test_dataset.py b/tests/datasets/pvnet/test_dataset.py index 53694def..a95834cc 100644 --- a/tests/datasets/pvnet/test_dataset.py +++ b/tests/datasets/pvnet/test_dataset.py @@ -1,3 +1,4 @@ +import os import pickle import numpy as np @@ -87,6 +88,7 @@ def _pvnet_dataset_sample_check(sample, config, batch_dim = None): def test_get_time_periods_mask(): + # The periods are half-open - inclusive of the start time, exclusive of the end time times = pd.to_datetime([ "2023-01-01 05:00", "2023-01-01 06:00", @@ -96,7 +98,7 @@ def test_get_time_periods_mask(): "2023-01-01 12:00", "2023-01-01 12:30", "2023-01-01 13:00", - ]) + ]).values mask = get_time_periods_mask( times, @@ -105,16 +107,27 @@ def test_get_time_periods_mask(): ("2023-01-01 12:00", "2023-01-01 13:00"), ], ) - expected_mask = np.array([False, True, True, True, False, True, True, True]) + expected_mask = np.array([False, True, True, False, False, True, True, False]) assert np.array_equal(mask, expected_mask), f"Expected {expected_mask} but got {mask}" mask = get_time_periods_mask( times, time_periods=[(None, "2023-01-01 07:00")], ) - expected_mask = np.array([True, True, True, True, False, False, False, False]) + expected_mask = np.array([True, True, True, False, False, False, False, False]) assert np.array_equal(mask, expected_mask), f"Expected {expected_mask} but got {mask}" + mask = get_time_periods_mask( + times, + time_periods=[("2023-01-01 12:30", None)], + ) + expected_mask = np.array([False, False, False, False, False, False, True, True]) + assert np.array_equal(mask, expected_mask), f"Expected {expected_mask} but got {mask}" + + # Unbounded in both directions - no time is filtered out, including the last one + mask = get_time_periods_mask(times, time_periods=[(None, None)]) + assert mask.all(), f"Expected all times to be kept but got {mask}" + def _expected_num_locations(dataset, catalog_ids): """The catalogued locations which survive the config's exclusion list.""" @@ -130,8 +143,9 @@ def test_pvnet_dataset(pvnet_config_filename): ], ) - expected_t0s = 6 # 2 time periods each with 3 t0s (inclusive) at 30 minute intervals + expected_t0s = 4 # 2 half-open time periods each with 2 t0s at 30 minute intervals num_locs = _expected_num_locations(dataset, LOCATION_IDS) + assert len(dataset.locations) == num_locs assert len(dataset.valid_t0_times) == expected_t0s @@ -170,8 +184,9 @@ def test_pvnet_dataset_sites(pvnet_site_config_filename): ], ) - expected_t0s = 6 # 2 time periods each with 3 t0s (inclusive) at 30 minute intervals + expected_t0s = 4 # 2 half-open time periods each with 2 t0s at 30 minute intervals num_locs = _expected_num_locations(dataset, SITE_LOCATION_IDS) + assert len(dataset.locations) == num_locs # Should be less than num_locs * expected_t0s as not all locations have data for all t0s # in the time periods @@ -181,6 +196,14 @@ def test_pvnet_dataset_sites(pvnet_site_config_filename): _pvnet_dataset_sample_check(sample, dataset.config) +def test_pvnet_dataset_sites_unbounded_time_periods(pvnet_site_config_filename): + """Unbounded time periods on the per-location t0 path (NaN-bearing generation).""" + dataset = PVNetDataset(pvnet_site_config_filename, time_periods=[(None, None)]) + assert not dataset.complete_generation + # An unbounded period should filter out nothing + assert len(dataset) == len(PVNetDataset(pvnet_site_config_filename)) + + def test_pvnet_dataset_noxarray_mode(pvnet_config_filename): dataset = PVNetDataset(pvnet_config_filename, use_xarray=True) sample = dataset[0] @@ -373,6 +396,19 @@ def test_pvnet_dataset_pickle(tmp_path, pvnet_config_filename): _ = pickle.loads(pickle_bytes) # noqa: S301 +def test_pvnet_dataset_pickle_missing_presaved_file(tmp_path, pvnet_config_filename): + """Unpickling must fail loudly if the presaved state file has gone missing.""" + pickle_path = f"{tmp_path}.pkl" + dataset = PVNetDataset(pvnet_config_filename) + dataset.presave_pickle(pickle_path) + pickle_bytes = pickle.dumps(dataset) + + os.remove(pickle_path) + + with pytest.raises(FileNotFoundError, match="Presaved state file not found"): + _ = pickle.loads(pickle_bytes) # noqa: S301 + + def test_pvnet_dataset_get_sample(pvnet_config_filename): dataset = PVNetDataset( pvnet_config_filename, diff --git a/tests/features/test_time_encodings.py b/tests/features/test_time_encodings.py index fcf13147..7125cc27 100644 --- a/tests/features/test_time_encodings.py +++ b/tests/features/test_time_encodings.py @@ -19,6 +19,17 @@ def test_encode_datetimes(): # Values should be between -1 and 1 for key in ("date_sin", "date_cos", "time_sin", "time_cos"): assert np.all(np.abs(features[key]) <= 1) + assert features[key].dtype == np.float32 + + # The date encoding must agree with encode_t0 and must not alias across the year boundary + for datetime in (np.datetime64("2023-01-01"), np.datetime64("2024-12-31")): + date_sin = encode_datetimes(np.array([datetime]))["date_sin"][0] + assert date_sin == encode_t0(datetime, [("1y", "cyclic")])[0] + + assert ( + encode_datetimes(np.array([np.datetime64("2023-01-01")]))["date_sin"][0] + != encode_datetimes(np.array([np.datetime64("2024-12-31")]))["date_sin"][0] + ) def test_encode_t0(): diff --git a/tests/load/test_load_nwp.py b/tests/load/test_load_nwp.py index b6a8edb9..e838ce17 100755 --- a/tests/load/test_load_nwp.py +++ b/tests/load/test_load_nwp.py @@ -158,6 +158,27 @@ def test_load_ukv_bad_dtype_step(tmp_path): open_nwp(zarr_path=str(zarr_path), provider="ukv") +def test_load_ukv_rejects_uneven_steps(tmp_path): + """Test steps must be evenly spaced - the valid-t0 calculation assumes it.""" + zarr_path = tmp_path / "uneven_ukv_steps.zarr" + array = DataArray( + np.random.rand(1, 3, 1, 1, 1).astype(np.float32), + dims=("init_time_utc", "step", "channel", "x_osgb", "y_osgb"), + coords={ + "init_time_utc": [np.datetime64("2023-01-01")], + # Hourly, then 3-hourly + "step": np.array([1, 2, 5], dtype="timedelta64[h]"), + "channel": ["t"], + "x_osgb": np.array([0], dtype=np.float32), + "y_osgb": np.array([50], dtype=np.float32), + }, + ) + array.to_zarr(zarr_path) + + with pytest.raises(ValueError, match="must be evenly spaced"): + open_nwp(zarr_path=str(zarr_path), provider="ukv") + + def test_load_ecmwf_bad_dtype_longitude(tmp_path): """Test validation fails for ECMWF with a non-numeric longitude dtype.""" zarr_path = tmp_path / "bad_ecmwf_longitude.zarr" diff --git a/tests/select/test_spatial_slice.py b/tests/select/test_spatial_slice.py index 1470b9ce..b4d6cb06 100644 --- a/tests/select/test_spatial_slice.py +++ b/tests/select/test_spatial_slice.py @@ -134,6 +134,17 @@ def test_select_spatial_slice_pixels_out_of_bounds(da): assert "Location(id=456, coord_systems=['osgb'], coordinates={'osgb': (90.1, 90.1)})" in msg +def test_select_spatial_slice_pixels_multiple_empty_locations(da): + """Test that an empty locations list raises rather than returning an empty slice.""" + with pytest.raises(ValueError, match="`locations` is empty"): + select_spatial_slice_pixels_multiple( + da, + locations=[], + width_pixels=3, + height_pixels=3, + ) + + def test_select_spatial_slice_pixels_multiple_out_of_bounds(da): """Test error includes all location context for multi-location spatial slice requests.""" with pytest.raises(ValueError) as excinfo: diff --git a/tests/select/test_time_periods.py b/tests/select/test_time_periods.py index e15a035b..dc575847 100644 --- a/tests/select/test_time_periods.py +++ b/tests/select/test_time_periods.py @@ -1,5 +1,6 @@ import numpy as np import pandas as pd +import pytest from ocf_data_sampler.select.time_periods import ( fill_time_periods, @@ -179,6 +180,32 @@ def test_find_contiguous_t0_periods_nwp(): assert time_periods.equals(expected) +def test_find_contiguous_t0_periods_nwp_forecast_too_short(): + """Test that a forecast too short to serve any t0 raises rather than emitting empty periods.""" + init_times = pd.date_range("2023-01-01 00:00", "2023-01-02 00:00", freq="6h", unit="ns").values + + # The forecast only reaches 2 hours ahead, but each sample needs 3 hours of future data + with pytest.raises(ValueError, match="no t0s are available"): + find_contiguous_t0_periods_nwp( + init_times=init_times, + interval_start=np.timedelta64(0, "h"), + interval_end=np.timedelta64(3, "h"), + first_forecast_step=np.timedelta64(0, "h"), + last_forecast_step=np.timedelta64(2, "h"), + ) + + # The same failure via a max_staleness shorter than the wait imposed by first_forecast_step + with pytest.raises(ValueError, match="no t0s are available"): + find_contiguous_t0_periods_nwp( + init_times=init_times, + interval_start=np.timedelta64(0, "h"), + interval_end=np.timedelta64(3, "h"), + first_forecast_step=np.timedelta64(6, "h"), + last_forecast_step=np.timedelta64(36, "h"), + max_staleness=np.timedelta64(3, "h"), + ) + + def test_intersect_time_periods_with_2_inputs(): def assert_expected_result_with_reverse(a, b, expected_result): """Assert the calculated intersection is as expected with and without a and b switched"""