diff --git a/src/ocf_data_sampler/config/__init__.py b/src/ocf_data_sampler/config/__init__.py index 78ead638..b4040a48 100644 --- a/src/ocf_data_sampler/config/__init__.py +++ b/src/ocf_data_sampler/config/__init__.py @@ -1,5 +1,5 @@ """Configuration model.""" from ocf_data_sampler.config.load import load_yaml_configuration -from ocf_data_sampler.config.model import Configuration, InputData +from ocf_data_sampler.config.model import PVNetDataConfig from ocf_data_sampler.config.save import save_yaml_configuration diff --git a/src/ocf_data_sampler/config/load.py b/src/ocf_data_sampler/config/load.py index 150262e4..8fa08f39 100644 --- a/src/ocf_data_sampler/config/load.py +++ b/src/ocf_data_sampler/config/load.py @@ -3,10 +3,10 @@ import fsspec from pyaml_env import parse_config -from ocf_data_sampler.config.model import Configuration +from ocf_data_sampler.config.model import PVNetDataConfig -def load_yaml_configuration(filename: str) -> Configuration: +def load_yaml_configuration(filename: str) -> PVNetDataConfig: """Load a yaml file which has a configuration in it. Args: @@ -19,4 +19,4 @@ def load_yaml_configuration(filename: str) -> Configuration: with fsspec.open(filename, mode="r") as stream: configuration = parse_config(data=stream) - return Configuration(**configuration) + return PVNetDataConfig(**configuration) diff --git a/src/ocf_data_sampler/config/model.py b/src/ocf_data_sampler/config/model.py index f8c674ee..ab2e37f4 100644 --- a/src/ocf_data_sampler/config/model.py +++ b/src/ocf_data_sampler/config/model.py @@ -1,8 +1,4 @@ -"""Configuration model for the dataset. - -Absolute or relative zarr filepath(s). -Prefix with a protocol like s3:// to read from alternative filesystems. -""" +"""Configuration model for the PVNet dataset.""" from collections.abc import Iterator from typing import Literal @@ -20,16 +16,6 @@ class Base(BaseModel): model_config = ConfigDict(extra="forbid") -class General(Base): - """General pydantic model.""" - - name: str = Field("example", description="The name of this configuration file") - description: str = Field( - "example configuration", - description="Description of this configuration file", - ) - - class TimeWindowMixin(Base): """Mixin class, to add interval start, end and resolution minutes.""" @@ -72,7 +58,16 @@ def validate_intervals(self) -> "TimeWindowMixin": return self -class DropoutMixin(Base): +class FillValueMixin(Base): + """Mixin class, to add a value used for filling missing data.""" + + dropout_fill_value: float = Field( + default=0.0, + description="The value used to fill in dropped out data or any missing values." + ) + + +class DropoutMixin(FillValueMixin): """Mixin class, to add dropout minutes.""" dropout_timedeltas_minutes: list[int] = Field( @@ -87,11 +82,6 @@ class DropoutMixin(Base): "floats (probability that dropout of the corresponding timedelta is applied)", ) - dropout_fill_value: float = Field( - default=0.0, - description="The value used to fill in dropped out data or any missing values." - ) - @field_validator("dropout_timedeltas_minutes") def dropout_timedeltas_minutes_negative(cls, v: list[int]) -> list[int]: """Validate 'dropout_timedeltas_minutes'.""" @@ -235,7 +225,7 @@ class NWP(TimeWindowMixin, DropoutMixin, SpatialWindowMixin, NormalisationConsta provider: str = Field(..., description="The provider of the NWP data") - accum_channels: list[str] = Field([], description="the nwp channels which need to be diffed") + accum_channels: list[str] = Field([], description="The NWP channels which need to be diffed") max_staleness_minutes: int | None = Field( None, @@ -319,7 +309,48 @@ def items(self) -> Iterator[tuple[str, NWP]]: return self.root.items() -class Generation(TimeWindowMixin, DropoutMixin): +class GenerationWindow(Base): + """Interval bounds for a generation window. + + Unlike `TimeWindowMixin`, no temporal resolution here - it belongs to the shared + generation data source (`Generation.time_resolution_minutes`), not to an individual window. + """ + + interval_start_minutes: int = Field( + ..., + description="Data interval starts at `t0 + interval_start_minutes`", + ) + + interval_end_minutes: int = Field( + ..., + description="Data interval ends at `t0 + interval_end_minutes`", + ) + + @model_validator(mode="after") + def validate_interval_order(self) -> "GenerationWindow": + """Validator for time interval fields.""" + start = self.interval_start_minutes + end = self.interval_end_minutes + if start > end: + raise ValueError( + f"interval_start_minutes ({start}) must be <= interval_end_minutes ({end})", + ) + return self + + +class GenerationInputWindow(GenerationWindow, DropoutMixin): + """A generation window with dropout configuration. + + Dropout is configurable here since only the input window (not the prediction target) + should ever be randomly masked out. + """ + + +class GenerationTargetWindow(GenerationWindow, FillValueMixin): + """Generation target window configuration model.""" + + +class Generation(Base): """Generation configuration model.""" zarr_path: str = Field( @@ -328,11 +359,77 @@ class Generation(TimeWindowMixin, DropoutMixin): "to read from alternative filesystems.", ) + time_resolution_minutes: int = Field( + ..., + gt=0, + description="The temporal resolution of the generation data in minutes", + ) + + input: GenerationInputWindow | None = None + target: GenerationTargetWindow | None = None + + @model_validator(mode="after") + def validate_windows(self) -> "Generation": + """Validate the input/target windows are set and divisible by the shared resolution.""" + if self.input is None and self.target is None: + raise ValueError( + "At least one of `generation.input` or `generation.target` must be configured", + ) + + for name, window in (("input", self.input), ("target", self.target)): + if window is None: + continue + for bound_name, bound in ( + ("interval_start_minutes", window.interval_start_minutes), + ("interval_end_minutes", window.interval_end_minutes), + ): + if bound % self.time_resolution_minutes != 0: + raise ValueError( + f"generation.{name}.{bound_name} ({bound}) must be divisible by " + f"generation.time_resolution_minutes ({self.time_resolution_minutes})", + ) + return self + + +class SamplingGrid(Base): + """Configuration for the (location, time) grid that t0 times are sampled from.""" + + locations_csv_path: str = Field( + ..., + description="Absolute or relative CSV filepath to the locations metadata (location IDs " + "and their coordinates) - see `ocf_data_sampler.load.locations.open_locations`.", + ) + + exclude_location_ids: list[int] = Field( + default=[], + description="Location IDs from the locations metadata to drop from the sampling grid. " + "Every ID listed must be present in the locations data.", + ) + + t0_resolution_minutes: int = Field( + ..., + gt=0, + description="The cadence t0 candidates are enumerated at, needed to compute valid t0 " + "times regardless of which other input sources are configured.", + ) + + @field_validator("exclude_location_ids") + def validate_exclude_location_ids_unique(cls, v: list[int]) -> list[int]: + """Validate 'exclude_location_ids'.""" + duplicates = {i for i in v if v.count(i) > 1} + if duplicates: + raise ValueError(f"exclude_location_ids contains duplicates: {sorted(duplicates)}") + return v + class SolarPosition(TimeWindowMixin): """Solar position configuration model.""" +class DatetimeEncoding(TimeWindowMixin): + """Datetime encoding configuration model.""" + + _embedding_type = list[tuple[str, Literal["cyclic", "linear"]]] class T0Embedding(Base): """Configuration for the t0 time embedding.""" @@ -374,18 +471,13 @@ def validate_embeddings(cls, embeddings: _embedding_type) -> _embedding_type: return embeddings -class InputData(Base): - """Input data model.""" +class PVNetDataConfig(Base): + """Configuration model for the PVNet dataset.""" + sampling_grid: SamplingGrid satellite: Satellite | None = None nwp: MultiNWP | None = None generation: Generation | None = None solar_position: SolarPosition | None = None + datetime_encoding: DatetimeEncoding | None = None t0_embedding: T0Embedding | None = None - - -class Configuration(Base): - """Configuration model for the dataset.""" - - general: General = General() - input_data: InputData = InputData() diff --git a/src/ocf_data_sampler/config/save.py b/src/ocf_data_sampler/config/save.py index 3bb15869..57bd96dc 100644 --- a/src/ocf_data_sampler/config/save.py +++ b/src/ocf_data_sampler/config/save.py @@ -10,14 +10,14 @@ import fsspec import yaml -from ocf_data_sampler.config.model import Configuration +from ocf_data_sampler.config.model import PVNetDataConfig -def save_yaml_configuration(configuration: Configuration, filename: str) -> None: +def save_yaml_configuration(configuration: PVNetDataConfig, filename: str) -> None: """Save a configuration object to a YAML file. Args: - configuration: Configuration object containing the settings to save + configuration: PVNetDataConfig object containing the settings to save filename: Destination path for the YAML file. Can be a local path or cloud storage URL (e.g., 'gs://', 's3://'). For local paths, absolute paths are recommended. diff --git a/src/ocf_data_sampler/datasets/pvnet/dataset.py b/src/ocf_data_sampler/datasets/pvnet/dataset.py index c1f3dfed..94f59a8e 100644 --- a/src/ocf_data_sampler/datasets/pvnet/dataset.py +++ b/src/ocf_data_sampler/datasets/pvnet/dataset.py @@ -12,15 +12,13 @@ from ocf_data_sampler.common.lightarray import LightDataArray from ocf_data_sampler.common.time_utils import date_range, get_posix_timestamp, minutes from ocf_data_sampler.config.load import load_yaml_configuration -from ocf_data_sampler.config.model import Configuration +from ocf_data_sampler.config.model import PVNetDataConfig from ocf_data_sampler.datasets.cache import PickleCacheMixin from ocf_data_sampler.datasets.pvnet.loading import get_dataset_dict from ocf_data_sampler.datasets.pvnet.materialise import load_data_dict from ocf_data_sampler.datasets.pvnet.preprocess import ( - apply_dropout_to_datasets, config_normalization_values_to_dicts, - diff_nwp_data, - fill_nans_in_dataset_dicts, + preprocess_dataset_dict, ) from ocf_data_sampler.datasets.pvnet.sample import ( convert_to_numpy_sample, @@ -35,6 +33,7 @@ from ocf_data_sampler.datasets.pvnet.types import NumpySample, SourceDict, TensorBatch from ocf_data_sampler.datasets.pvnet.valid_t0s import find_valid_time_periods from ocf_data_sampler.features.time_encodings import encode_datetimes +from ocf_data_sampler.load import open_locations from ocf_data_sampler.select import ( fill_time_periods, find_contiguous_t0_periods, @@ -46,27 +45,36 @@ -def get_locations(generation_data: xr.DataArray) -> list[Location]: - """Get list of locations of all locations. +def get_locations(csv_path: str, exclude_ids: list[int] | None = None) -> list[Location]: + """Load the locations metadata and build the list of all locations. Args: - generation_data: xarray dataarray of generation data with location info + csv_path: Path to the locations CSV data + exclude_ids: Location IDs to drop from the returned locations """ - locations = [] - location_ids = generation_data["location_id"].values - - for location_id in location_ids: - gen_data = generation_data.sel(location_id=location_id) - locations.append( - Location( - x=gen_data["longitude"].values, - y=gen_data["latitude"].values, - coord_system="lon_lat", - id=int(location_id), - ), - ) + locations_data = open_locations(csv_path) - return locations + if exclude_ids: + missing_ids = np.setdiff1d(exclude_ids, locations_data["location_id"].values) + if len(missing_ids) > 0: + raise ValueError( + f"Cannot exclude location IDs which are not in the locations data: {missing_ids}", + ) + + locations_data = locations_data[~locations_data["location_id"].isin(exclude_ids)] + + if len(locations_data) == 0: + raise ValueError("All location IDs in the locations data have been excluded") + + return [ + Location( + x=row.longitude, + y=row.latitude, + coord_system="lon_lat", + id=int(row.location_id), + ) + for row in locations_data.itertuples() + ] def xarray_to_lightarray_dict( @@ -161,6 +169,69 @@ def add_alternate_coordinate_projections( return locations +def build_numpy_sample( + dataset_dict: SourceDict, + t0: np.datetime64, + location: Location, + config: PVNetDataConfig, + include_extra_metadata: bool = False, +) -> NumpySample: + """Convert data to numpy arrays and add auxiliary features. + + Note: the data in `dataset_dict` is expected to already be preprocessed - see + `preprocess_dataset_dict`. + + Args: + dataset_dict: Dictionary of xarray datasets + t0: init-time for sample + location: location of the sample + config: PVNetDataConfig object + include_extra_metadata: Whether to add additional non-essential metadata to the sample + """ + # Convert all xarray modalities to a single NumpySample + sample = convert_to_numpy_sample(dataset_dict, include_extra_metadata) + + sample["location_id"] = location.id + lon, lat = location.in_coord_system("lon_lat") + + if include_extra_metadata: + sample["location_longitude"] = lon + sample["location_latitude"] = lat + + # Add t0 embedding if configured + if config.t0_embedding is not None: + sample.update( + make_t0_encoding_numpy_sample(t0, config.t0_embedding.embeddings), + ) + + # Add datetime encodings if configured + if config.datetime_encoding is not None: + dt_config = config.datetime_encoding + + datetimes = date_range( + t0 + minutes(dt_config.interval_start_minutes), + t0 + minutes(dt_config.interval_end_minutes), + freq=minutes(dt_config.time_resolution_minutes), + ) + sample.update(encode_datetimes(datetimes=datetimes)) + + # Add solar position if configured + if config.solar_position is not None: + solar_config = config.solar_position + + # Create datetime range for solar position calculation + datetimes = date_range( + t0 + minutes(solar_config.interval_start_minutes), + t0 + minutes(solar_config.interval_end_minutes), + freq=minutes(solar_config.time_resolution_minutes), + ) + + sample.update(make_sun_position_numpy_sample(datetimes, lon=lon, lat=lat)) + + sample["t0"] = get_posix_timestamp(t0) + + return sample + class AbstractPVNetDataset(PickleCacheMixin, Dataset): """Abstract class for PVNet datasets.""" @@ -191,10 +262,28 @@ def __init__( config = load_yaml_configuration(config_filename) - datasets_dict = get_dataset_dict(config.input_data) + locations = get_locations( + config.sampling_grid.locations_csv_path, + config.sampling_grid.exclude_location_ids, + ) + + datasets_dict = get_dataset_dict(config) + + if "generation" in datasets_dict: + location_ids = [loc.id for loc in locations] + missing = np.setdiff1d(location_ids, datasets_dict["generation"]["location_id"].values) + if len(missing) > 0: + raise ValueError(f"Generation data is missing for location IDs: {missing}") - # Check if generation data has nans - self.complete_generation = not datasets_dict["generation"].isnull().any() + # Slice the generation data to only include the specified locations. This allows us to + # quality check the generation data for nans and find valid t0 times for each location. + datasets_dict["generation"] = datasets_dict["generation"].sel(location_id=location_ids) + + # Check if generation data has nans. If generation isn't configured at all, there's no + # per-location data availability to consider, so a single global t0 grid still applies. + self.complete_generation = ( + "generation" not in datasets_dict or not datasets_dict["generation"].isnull().any() + ) if self.complete_generation: valid_t0_times = self.find_valid_t0_times(datasets_dict, config) @@ -210,7 +299,9 @@ def __init__( "Generation data has nans so t0s are handled separately for each location_id.", ) # If non-identical times per location, find valid t0s per location id - valid_t0_and_location_ids = self.find_valid_t0_and_location_ids(datasets_dict, config) + valid_t0_and_location_ids = self.find_valid_t0_and_location_ids( + datasets_dict, locations, config, + ) # Filter t0 times to given range if time_periods is not None: @@ -219,9 +310,6 @@ def __init__( self.valid_t0_and_location_ids = valid_t0_and_location_ids - # Construct list of locations to sample from - locations = get_locations(generation_data=datasets_dict["generation"]) - self.locations = add_alternate_coordinate_projections(locations, datasets_dict) self.config = config @@ -232,12 +320,6 @@ def __init__( else: self.datasets_dict = xarray_to_lightarray_dict(datasets_dict) - # Assign t0 idx value - self.t0_idx = ( - -config.input_data.generation.interval_start_minutes - // config.input_data.generation.time_resolution_minutes - ) - # Extract the normalisation values from the config for faster access mean_dict, std_dict, clip_min_dict, clip_max_dict = ( config_normalization_values_to_dicts(config) @@ -262,113 +344,31 @@ def _sanitise_index(self, idx: int) -> int: return index - def process_and_combine_datasets( - self, - dataset_dict: SourceDict, - t0: np.datetime64, - location: Location, - ) -> NumpySample: - """Normalise and convert data to numpy arrays. - - Args: - dataset_dict: Dictionary of xarray datasets - t0: init-time for sample - location: location of the sample - """ - # Normalise NWP - if "nwp" in dataset_dict: - for nwp_key, da_nwp in dataset_dict["nwp"].items(): - channel_means = self.mean_dict["nwp"][nwp_key] - channel_stds = self.std_dict["nwp"][nwp_key] - channel_mins = self.clip_min_dict["nwp"][nwp_key] - channel_maxs = self.clip_max_dict["nwp"][nwp_key] - dataset_dict["nwp"][nwp_key].data = ( - (da_nwp.data.clip(channel_mins, channel_maxs) - channel_means) - / channel_stds - ) - - # Normalise satellite - if "sat" in dataset_dict: - channel_means = self.mean_dict["sat"] - channel_stds = self.std_dict["sat"] - channel_mins = self.clip_min_dict["sat"] - channel_maxs = self.clip_max_dict["sat"] - dataset_dict["sat"].data = ( - (dataset_dict["sat"].data.clip(channel_mins, channel_maxs) - channel_means) - / channel_stds - ) - - # Fill NaNs - dataset_dict = fill_nans_in_dataset_dicts(dataset_dict, config=self.config) - - # Convert all xarray modalities to a single NumpySample - sample = convert_to_numpy_sample(dataset_dict, self.t0_idx, self.include_extra_metadata) - - # Add location metadata not present on the DataArray - if "generation" in dataset_dict: - sample["location_id"] = location.id - - # Add datetime encodings over the full generation time range - generation_config = self.config.input_data.generation - datetimes = date_range( - t0 + minutes(generation_config.interval_start_minutes), - t0 + minutes(generation_config.interval_end_minutes), - freq=minutes(generation_config.time_resolution_minutes), - ) - sample.update(encode_datetimes(datetimes=datetimes)) - - # Add t0 embedding if configured - if self.config.input_data.t0_embedding is not None: - sample.update( - make_t0_encoding_numpy_sample(t0, self.config.input_data.t0_embedding.embeddings), - ) - - # Add solar position if configured - if self.config.input_data.solar_position is not None: - solar_config = self.config.input_data.solar_position - - # Create datetime range for solar position calculation - datetimes = date_range( - t0 + minutes(solar_config.interval_start_minutes), - t0 + minutes(solar_config.interval_end_minutes), - freq=minutes(solar_config.time_resolution_minutes), - ) - sample.update( - make_sun_position_numpy_sample( - datetimes, - dataset_dict["generation"]["longitude"].values, - dataset_dict["generation"]["latitude"].values, - ), - ) - - sample["t0"] = get_posix_timestamp(t0) - - return sample - @staticmethod def find_valid_t0_times( datasets_dict: SourceDict, - config: Configuration, + config: PVNetDataConfig, ) -> NDArray[np.datetime64]: """Find the t0 times where all of the requested input data is available. Args: datasets_dict: A dictionary of input datasets - config: Configuration file + config: PVNetDataConfig file """ valid_time_periods = find_valid_time_periods(datasets_dict, config) # Fill out the contiguous time periods to get the t0 times valid_t0_times = fill_time_periods( valid_time_periods, - freq=minutes(config.input_data.generation.time_resolution_minutes), + freq=minutes(config.sampling_grid.t0_resolution_minutes), ) return valid_t0_times @staticmethod def find_valid_t0_and_location_ids( datasets_dict: SourceDict, - config: Configuration, + locations: list[Location], + config: PVNetDataConfig, ) -> pd.DataFrame: """Find the t0 times where all of the requested input data is available for each location. @@ -378,41 +378,48 @@ def find_valid_t0_and_location_ids( Args: datasets_dict: A dictionary of input datasets - config: Configuration file + 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) # Loop over each location in system id and obtain valid periods - generations = datasets_dict["generation"] - location_ids = generations.location_id.values - generation_config = config.input_data.generation + generation_windows = [ + w for w in (config.generation.input, config.generation.target) if w is not None + ] valid_t0_and_location_ids = [] - for location_id in location_ids: - generation = generations.sel(location_id=location_id) - # Drop NaN values - generation = generation.dropna(dim="time_utc") - - # Obtain valid time periods for this location - time_periods = find_contiguous_t0_periods( - generation["time_utc"].values, - time_resolution=minutes(generation_config.time_resolution_minutes), - interval_start=minutes(generation_config.interval_start_minutes), - interval_end=minutes(generation_config.interval_end_minutes), + for location in locations: + # Drop NaN values for location + generation = ( + datasets_dict["generation"] + .sel(location_id=location.id) + .dropna(dim="time_utc") ) + + # Obtain valid time periods for this location, for each configured window + time_periods_per_window = [ + find_contiguous_t0_periods( + generation["time_utc"].values, + time_resolution=minutes(config.generation.time_resolution_minutes), + interval_start=minutes(window_config.interval_start_minutes), + interval_end=minutes(window_config.interval_end_minutes), + ) + for window_config in generation_windows + ] valid_time_periods_per_location = intersect_time_periods( - [valid_time_periods, time_periods], + [valid_time_periods, *time_periods_per_window], ) # Fill out contiguous time periods to get t0 times valid_t0_times_per_location = fill_time_periods( valid_time_periods_per_location, - freq=minutes(generation_config.time_resolution_minutes), + 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_per_location["location_id"] = location.id valid_t0_and_location_ids.append(valid_t0_per_location) valid_t0_and_location_ids = pd.concat(valid_t0_and_location_ids) @@ -451,10 +458,13 @@ def _get_sample(self, t0: np.datetime64, location: Location) -> NumpySample: sample_dict = slice_datasets_by_space(self.datasets_dict, location, self.config) sample_dict = slice_datasets_by_time(sample_dict, t0, self.config) sample_dict = load_data_dict(sample_dict) - # Apply dropout to the data sources in-place - apply_dropout_to_datasets(sample_dict, t0, self.config) - sample_dict = diff_nwp_data(sample_dict, self.config) - return self.process_and_combine_datasets(sample_dict, t0, location) + sample_dict = preprocess_dataset_dict( + sample_dict, t0, self.config, + self.mean_dict, self.std_dict, self.clip_min_dict, self.clip_max_dict, + ) + return build_numpy_sample( + sample_dict, t0, location, self.config, self.include_extra_metadata, + ) @override def __getitem__(self, idx: int) -> NumpySample: @@ -549,19 +559,19 @@ def _get_sample(self, t0: np.datetime64) -> TensorBatch: # Slice by time then load to avoid loading the data multiple times from disk sample_dict = slice_datasets_by_time(self.datasets_dict, t0, self.config) sample_dict = load_data_dict(sample_dict) - # Apply dropout to the data sources in-place - apply_dropout_to_datasets(sample_dict, t0, self.config) - sample_dict = diff_nwp_data(sample_dict, self.config) + # Preprocessing is location-independent, so do it once before slicing per-location below + sample_dict = preprocess_dataset_dict( + sample_dict, t0, self.config, + self.mean_dict, self.std_dict, self.clip_min_dict, self.clip_max_dict, + ) samples = [] # Prepare sample for each location for location in self.locations: sliced_sample_dict = slice_datasets_by_space(sample_dict, location, self.config) - numpy_sample = self.process_and_combine_datasets( - sliced_sample_dict, - t0, - location, + numpy_sample = build_numpy_sample( + sliced_sample_dict, t0, location, self.config, self.include_extra_metadata, ) samples.append(numpy_sample) diff --git a/src/ocf_data_sampler/datasets/pvnet/loading.py b/src/ocf_data_sampler/datasets/pvnet/loading.py index ef3d359e..63e4e5f8 100644 --- a/src/ocf_data_sampler/datasets/pvnet/loading.py +++ b/src/ocf_data_sampler/datasets/pvnet/loading.py @@ -1,41 +1,31 @@ """Loads all data sources.""" -import logging - import xarray as xr -from ocf_data_sampler.config import InputData +from ocf_data_sampler.config import PVNetDataConfig from ocf_data_sampler.datasets.pvnet.types import SourceDict from ocf_data_sampler.load import open_generation, open_nwp, open_sat_data -logger = logging.getLogger(__name__) +def get_dataset_dict(config: PVNetDataConfig) -> SourceDict[xr.DataArray]: + """Construct dictionary of all of the per-sample input data sources. -def get_dataset_dict(input_config: InputData) -> SourceDict[xr.DataArray]: - """Construct dictionary of all of the input data sources. + Locations metadata is deliberately excluded - it isn't a per-sample source, so the caller + loads it separately. Args: - input_config: InputData configuration object + config: PVNetDataConfig configuration object """ datasets_dict = {} - # Load generation data unless the path is None - if input_config.generation and input_config.generation.zarr_path: - da_generation = open_generation(zarr_path=input_config.generation.zarr_path) - - # Remove location_id 0 if more than one location present - if len(da_generation["location_id"]) > 1 and 0 in da_generation["location_id"].values: - da_generation = da_generation.drop_sel(location_id=0) - logger.info( - "Id 0 has been filtered out, this is only used for summation models.", - ) - - datasets_dict["generation"] = da_generation + # Load generation data if in config + if config.generation is not None: + datasets_dict["generation"] = open_generation(zarr_path=config.generation.zarr_path) # Load NWP data if in config - if input_config.nwp: + if config.nwp: datasets_dict["nwp"] = {} - for nwp_source, nwp_config in input_config.nwp.items(): + for nwp_source, nwp_config in config.nwp.items(): da_nwp = open_nwp(zarr_path=nwp_config.zarr_path, provider=nwp_config.provider) da_nwp = da_nwp.sel(channel=list(nwp_config.channels)) @@ -43,12 +33,11 @@ def get_dataset_dict(input_config: InputData) -> SourceDict[xr.DataArray]: datasets_dict["nwp"][nwp_source] = da_nwp # Load satellite data if in config - if input_config.satellite: - sat_config = input_config.satellite + if config.satellite: - da_sat = open_sat_data(sat_config.zarr_path) + da_sat = open_sat_data(config.satellite.zarr_path) - da_sat = da_sat.sel(channel=list(sat_config.channels)) + da_sat = da_sat.sel(channel=list(config.satellite.channels)) datasets_dict["sat"] = da_sat diff --git a/src/ocf_data_sampler/datasets/pvnet/preprocess.py b/src/ocf_data_sampler/datasets/pvnet/preprocess.py index e7cf44c6..144b374f 100644 --- a/src/ocf_data_sampler/datasets/pvnet/preprocess.py +++ b/src/ocf_data_sampler/datasets/pvnet/preprocess.py @@ -4,14 +4,14 @@ from ocf_data_sampler.common.time_utils import minutes from ocf_data_sampler.common.types import TArray -from ocf_data_sampler.config.model import Configuration +from ocf_data_sampler.config.model import PVNetDataConfig from ocf_data_sampler.datasets.pvnet.types import SourceDict from ocf_data_sampler.features.diff_channels import diff_channels -from ocf_data_sampler.select.dropout import apply_history_dropout +from ocf_data_sampler.select.dropout import apply_dropout def config_normalization_values_to_dicts( - config: Configuration, + config: PVNetDataConfig, ) -> tuple[dict[str, np.ndarray | dict[str, np.ndarray]]]: """Construct numpy arrays of mean, std, and clip values from the config normalisation constants. @@ -29,15 +29,14 @@ def config_normalization_values_to_dicts( clip_min_dict = {} clip_max_dict = {} - if config.input_data.nwp is not None: + if config.nwp is not None: means_dict["nwp"] = {} stds_dict["nwp"] = {} clip_min_dict["nwp"] = {} clip_max_dict["nwp"] = {} - for nwp_key in config.input_data.nwp: - nwp_config = config.input_data.nwp[nwp_key] + for nwp_key, nwp_config in config.nwp.items(): means_list = [] stds_list = [] @@ -61,16 +60,15 @@ def config_normalization_values_to_dicts( clip_min_dict["nwp"][nwp_key] = np.array(clip_min_list)[None, :, None, None] clip_max_dict["nwp"][nwp_key] = np.array(clip_max_list)[None, :, None, None] - if config.input_data.satellite is not None: - sat_config = config.input_data.satellite + if config.satellite is not None: means_list = [] stds_list = [] clip_min_list = [] clip_max_list = [] - for channel in list(sat_config.channels): - norm_conf = sat_config.normalisation_constants[channel] + for channel in list(config.satellite.channels): + norm_conf = config.satellite.normalisation_constants[channel] means_list.append(norm_conf.mean) stds_list.append(norm_conf.std) clip_min_list.append(-np.inf if norm_conf.clip_min is None else norm_conf.clip_min) @@ -85,16 +83,93 @@ def config_normalization_values_to_dicts( return means_dict, stds_dict, clip_min_dict, clip_max_dict -def diff_nwp_data(dataset_dict: SourceDict, config: Configuration) -> SourceDict: +def normalise_dataset_dicts( + dataset_dict: SourceDict, + mean_dict: dict[str, np.ndarray | dict[str, np.ndarray]], + std_dict: dict[str, np.ndarray | dict[str, np.ndarray]], + clip_min_dict: dict[str, np.ndarray | dict[str, np.ndarray]], + clip_max_dict: dict[str, np.ndarray | dict[str, np.ndarray]], +) -> SourceDict: + """Normalise the NWP, satellite, and generation data in-place. + + Args: + dataset_dict: Dictionary of xarray datasets + mean_dict: Means, as constructed by `config_normalization_values_to_dicts` + std_dict: Standard deviations, as constructed by `config_normalization_values_to_dicts` + clip_min_dict: Clip minimums, as constructed by `config_normalization_values_to_dicts` + clip_max_dict: Clip maximums, as constructed by `config_normalization_values_to_dicts` + """ + if "nwp" in dataset_dict: + for nwp_key, da_nwp in dataset_dict["nwp"].items(): + channel_means = mean_dict["nwp"][nwp_key] + channel_stds = std_dict["nwp"][nwp_key] + channel_mins = clip_min_dict["nwp"][nwp_key] + channel_maxs = clip_max_dict["nwp"][nwp_key] + dataset_dict["nwp"][nwp_key].data = ( + (da_nwp.data.clip(channel_mins, channel_maxs) - channel_means) + / channel_stds + ) + + if "sat" in dataset_dict: + channel_means = mean_dict["sat"] + channel_stds = std_dict["sat"] + channel_mins = clip_min_dict["sat"] + channel_maxs = clip_max_dict["sat"] + dataset_dict["sat"].data = ( + (dataset_dict["sat"].data.clip(channel_mins, channel_maxs) - channel_means) + / channel_stds + ) + + for key in ("generation_input", "generation_target"): + if key in dataset_dict: + dataset_dict[key] = normalise_generation_by_capacity(dataset_dict[key]) + + return dataset_dict + + +def normalise_generation_by_capacity(da: TArray) -> TArray: + """Rescale `generation_mw` to a capacity factor, leaving `capacity_mwp` unchanged. + + Zero capacity means no plant, so the capacity factor is taken as 0 rather than the undefined + 0/0. Emitting NaN instead would route it through the dropout fill, which signals missing data + rather than a known-zero output. + + Args: + da: Generation DataArray-like with a `gen_param` dimension + """ + gen_params = list(da["gen_param"].values) + gen_idx = gen_params.index("generation_mw") + cap_idx = gen_params.index("capacity_mwp") + + generation_values = da.isel(gen_param=gen_idx).values + capacity_values = da.isel(gen_param=cap_idx).values + + normalised = np.divide( + generation_values, + capacity_values, + out=np.zeros_like(generation_values, dtype=float), + where=capacity_values != 0, + ) + + new_data = da.data.copy() + index = [slice(None)] * new_data.ndim + index[da.dims.index("gen_param")] = gen_idx + new_data[tuple(index)] = normalised + da.data = new_data + + return da + + +def diff_nwp_data(dataset_dict: SourceDict, config: PVNetDataConfig) -> SourceDict: """Take the in-place diff of some channels of the NWP data. Args: dataset_dict: Dictionary of xarray datasets - config: Configuration object + config: PVNetDataConfig object """ if "nwp" in dataset_dict: for nwp_key, da_nwp in dataset_dict["nwp"].items(): - accum_channels = config.input_data.nwp[nwp_key].accum_channels + accum_channels = config.nwp[nwp_key].accum_channels if len(accum_channels)>0: # diff_channels() is an in-place operation and modifies the input dataset_dict["nwp"][nwp_key] = diff_channels(da_nwp, accum_channels) @@ -104,61 +179,62 @@ def diff_nwp_data(dataset_dict: SourceDict, config: Configuration) -> SourceDict def apply_dropout_to_datasets( datasets_dict: SourceDict, t0: np.datetime64, - config: Configuration, + config: PVNetDataConfig, ) -> None: - """Apply dropout in-placeto the dictionary of input data sources around a given t0 time. + """Apply dropout in-place to the dictionary of input data sources around a given t0 time. Args: datasets_dict: Dictionary of the input data sources t0: The init-time - config: Configuration object. + config: PVNetDataConfig object. Returns: None. The input datasets_dict is modified in place. """ if "sat" in datasets_dict: - sat_config = config.input_data.satellite - - apply_history_dropout( + apply_dropout( datasets_dict["sat"], t0, - dropout_timedeltas=minutes(sat_config.dropout_timedeltas_minutes), - dropout_frac=sat_config.dropout_fraction, + dropout_timedeltas=minutes(config.satellite.dropout_timedeltas_minutes), + dropout_frac=config.satellite.dropout_fraction, ) - if "generation" in datasets_dict: - generation_config = config.input_data.generation + # generation_target is never dropped out since it's the prediction target + if "generation_input" in datasets_dict: - # Dropout on the past generation, but not the future generation - apply_history_dropout( - datasets_dict["generation"], + # Note: capacity_mwp is dropped out along with generation_mw + apply_dropout( + datasets_dict["generation_input"], t0, - dropout_timedeltas=minutes(generation_config.dropout_timedeltas_minutes), - dropout_frac=generation_config.dropout_fraction, + dropout_timedeltas=minutes(config.generation.input.dropout_timedeltas_minutes), + dropout_frac=config.generation.input.dropout_fraction, ) return -def fill_nans_in_dataset_dicts(datasets_dict: SourceDict, config: Configuration) -> SourceDict: +def fill_nans_in_dataset_dicts(datasets_dict: SourceDict, config: PVNetDataConfig) -> SourceDict: """Fills all NaN values in the dataarrays in-place. Args: datasets_dict: Dictionary of the input data sources - config: Configuration object. + config: PVNetDataConfig object. """ - conf_in = config.input_data - if "generation" in datasets_dict: - datasets_dict["generation"] = fill_nans( - datasets_dict["generation"], - conf_in.generation.dropout_fill_value, - ) + if config.generation is not None: + for key, window_config in ( + ("generation_input", config.generation.input), + ("generation_target", config.generation.target), + ): + if key in datasets_dict: + datasets_dict[key] = fill_nans( + datasets_dict[key], window_config.dropout_fill_value, + ) if "sat" in datasets_dict: - datasets_dict["sat"] = fill_nans(datasets_dict["sat"], conf_in.satellite.dropout_fill_value) + datasets_dict["sat"] = fill_nans(datasets_dict["sat"], config.satellite.dropout_fill_value) if "nwp" in datasets_dict: - for nwp_key, nwp_config in config.input_data.nwp.items(): + for nwp_key, nwp_config in config.nwp.items(): datasets_dict["nwp"][nwp_key] = fill_nans( datasets_dict["nwp"][nwp_key], nwp_config.dropout_fill_value, @@ -172,3 +248,42 @@ def fill_nans(da: TArray, fill_value: float) -> TArray: if np.isnan(da.data).any(): da.data = np.nan_to_num(da.data, copy=True, nan=fill_value) return da + + +def preprocess_dataset_dict( + dataset_dict: SourceDict, + t0: np.datetime64, + config: PVNetDataConfig, + mean_dict: dict[str, np.ndarray | dict[str, np.ndarray]], + std_dict: dict[str, np.ndarray | dict[str, np.ndarray]], + clip_min_dict: dict[str, np.ndarray | dict[str, np.ndarray]], + clip_max_dict: dict[str, np.ndarray | dict[str, np.ndarray]], +) -> SourceDict: + """Diff, normalise, dropout, and fill NaNs in the dictionary of input data sources. + + These steps are always applied in the order listed, since some steps depend on the output of + previous steps. For example, + - NWP channel differencing must be done before normalisation, since the diffed channels have + different statistics + - Dropout must be applied after normalisation, since the fill value is specified in normalised + units + - NaN filling must be done after dropout, since dropout introduces NaNs in the data + + Note: `dataset_dict` is expected to already be loaded - see `load_data_dict`. + + Args: + dataset_dict: Dictionary of xarray datasets + t0: The init-time + config: PVNetDataConfig object + mean_dict: Means, as constructed by `config_normalization_values_to_dicts` + std_dict: Standard deviations, as constructed by `config_normalization_values_to_dicts` + clip_min_dict: Clip minimums, as constructed by `config_normalization_values_to_dicts` + clip_max_dict: Clip maximums, as constructed by `config_normalization_values_to_dicts` + """ + dataset_dict = diff_nwp_data(dataset_dict, config) + dataset_dict = normalise_dataset_dicts( + dataset_dict, mean_dict, std_dict, clip_min_dict, clip_max_dict, + ) + apply_dropout_to_datasets(dataset_dict, t0, config) + dataset_dict = fill_nans_in_dataset_dicts(dataset_dict, config=config) + return dataset_dict diff --git a/src/ocf_data_sampler/datasets/pvnet/sample.py b/src/ocf_data_sampler/datasets/pvnet/sample.py index 9c6b33d5..3297fcf1 100644 --- a/src/ocf_data_sampler/datasets/pvnet/sample.py +++ b/src/ocf_data_sampler/datasets/pvnet/sample.py @@ -10,7 +10,6 @@ def convert_to_numpy_sample( datasets_dict: SourceDict, - t0_idx: int, include_extra_metadata: bool = False, ) -> NumpySample: """Convert a dictionary of xarray objects to a NumpySample. @@ -18,10 +17,10 @@ def convert_to_numpy_sample( Args: datasets_dict: Dictionary of xarray DataArrays, with same structure as used inside PVNetDataset classes. Expected keys are any of following: - - "generation": DataArray of generation data + - "generation_input": DataArray of generation data used as model input + - "generation_target": DataArray of generation data used as the prediction target - "sat": DataArray of satellite data - "nwp": dict of DataArrays by provider name (e.g. {"ukv": da, "ecmwf": da}) - t0_idx: Index of t0 within generation include_extra_metadata: Whether to add additional non-essential metadata to the batch Returns: @@ -29,36 +28,25 @@ def convert_to_numpy_sample( """ numpy_sample: NumpySample = {} - if "generation" in datasets_dict: - da = datasets_dict["generation"] + for key in ("generation_input", "generation_target"): + if key not in datasets_dict: + continue - # Get the position index of the generation and capacities + da = datasets_dict[key] + + # generation_mw is expected to already been normalised by capacity so should be in the + # range [0, 1]. capacity_mwp is still expected to be in MW gen_idx = list(da["gen_param"].values).index("generation_mw") cap_idx = list(da["gen_param"].values).index("capacity_mwp") - generation_values = da.isel(gen_param=gen_idx).values - capacity_value = da.isel(gen_param=cap_idx).values[0] - - if capacity_value!=0: - generation_values = generation_values/capacity_value - numpy_sample.update( { - "generation": generation_values, - "capacity_mwp": capacity_value, - "generation_t0_idx": int(t0_idx), - "generation_time_utc": da["time_utc"].values.astype(float), + key: da.isel(gen_param=gen_idx).values, + f"{key}_capacity_mwp": da.isel(gen_param=cap_idx).values, + f"{key}_time_utc": da["time_utc"].values.astype(float), }, ) - if include_extra_metadata: - numpy_sample.update( - { - "location_longitude": float(da["longitude"].values), - "location_latitude": float(da["latitude"].values), - }, - ) - if "sat" in datasets_dict: da = datasets_dict["sat"] numpy_sample.update({"satellite": da.values}) diff --git a/src/ocf_data_sampler/datasets/pvnet/slicing.py b/src/ocf_data_sampler/datasets/pvnet/slicing.py index 0bf33029..e2a8497d 100644 --- a/src/ocf_data_sampler/datasets/pvnet/slicing.py +++ b/src/ocf_data_sampler/datasets/pvnet/slicing.py @@ -4,7 +4,7 @@ from ocf_data_sampler.common.indexing import get_indices_in_sorted_unique from ocf_data_sampler.common.time_utils import minutes -from ocf_data_sampler.config.model import Configuration +from ocf_data_sampler.config.model import PVNetDataConfig from ocf_data_sampler.datasets.pvnet.types import SourceDict from ocf_data_sampler.select.spatial_slice import ( select_spatial_slice_pixels, @@ -17,29 +17,24 @@ def slice_datasets_by_space( datasets_dict: SourceDict, location: Location, - config: Configuration, + config: PVNetDataConfig, ) -> SourceDict: """Slice the dictionary of input data sources around a given location. Args: datasets_dict: Dictionary of the input data sources location: The location to sample around - config: Configuration object. + config: PVNetDataConfig object. Returns: A dictionary of the sliced input data sources. """ - if not set(datasets_dict.keys()).issubset({"nwp", "sat", "generation"}): - raise ValueError( - "'datasets_dict' should only contain keys 'nwp', 'sat', 'generation'", - ) - sliced_datasets_dict = {} if "nwp" in datasets_dict: sliced_datasets_dict["nwp"] = {} - for nwp_key, nwp_config in config.input_data.nwp.items(): + for nwp_key, nwp_config in config.nwp.items(): sliced_datasets_dict["nwp"][nwp_key] = select_spatial_slice_pixels( datasets_dict["nwp"][nwp_key], location, @@ -48,21 +43,24 @@ def slice_datasets_by_space( ) if "sat" in datasets_dict: - sat_config = config.input_data.satellite - sliced_datasets_dict["sat"] = select_spatial_slice_pixels( datasets_dict["sat"], location, - height_pixels=sat_config.image_size_pixels_height, - width_pixels=sat_config.image_size_pixels_width, + height_pixels=config.satellite.image_size_pixels_height, + width_pixels=config.satellite.image_size_pixels_width, ) - if "generation" in datasets_dict: + # Depending on whether this is called before or after time-slicing, the generation data is + # under a single "generation" key (raw, pre-split) or "generation_input"/"generation_target" + # (post-split) - slice whichever of these are present by location. + for key in ("generation", "generation_input", "generation_target"): + if key not in datasets_dict: + continue - location_ids = datasets_dict["generation"]["location_id"].values + location_ids = datasets_dict[key]["location_id"].values loc_index = get_indices_in_sorted_unique(location_ids, location.id) - sliced_datasets_dict["generation"] = datasets_dict["generation"].isel(location_id=loc_index) + sliced_datasets_dict[key] = datasets_dict[key].isel(location_id=loc_index) return sliced_datasets_dict @@ -70,14 +68,14 @@ def slice_datasets_by_space( def reduce_spatial_extent_of_datasets( datasets_dict: SourceDict, locations: list[Location], - config: Configuration, + config: PVNetDataConfig, ) -> SourceDict: """Reduce the spatial extent of the datasets to only cover the locations. Args: datasets_dict: Dictionary of the input data sources locations: List of locations to reduce to - config: Configuration object + config: PVNetDataConfig object Returns: A dictionary of the reduced input data sources. @@ -87,7 +85,7 @@ def reduce_spatial_extent_of_datasets( if "nwp" in datasets_dict: sliced_datasets_dict["nwp"] = {} - for nwp_key, nwp_config in config.input_data.nwp.items(): + for nwp_key, nwp_config in config.nwp.items(): sliced_datasets_dict["nwp"][nwp_key] = select_spatial_slice_pixels_multiple( datasets_dict["nwp"][nwp_key], locations, @@ -97,13 +95,12 @@ def reduce_spatial_extent_of_datasets( if "sat" in datasets_dict: - sat_config = config.input_data.satellite sliced_datasets_dict["sat"] = select_spatial_slice_pixels_multiple( datasets_dict["sat"], locations, - height_pixels=sat_config.image_size_pixels_height, - width_pixels=sat_config.image_size_pixels_width, + height_pixels=config.satellite.image_size_pixels_height, + width_pixels=config.satellite.image_size_pixels_width, ) if "generation" in datasets_dict: @@ -115,14 +112,14 @@ def reduce_spatial_extent_of_datasets( def slice_datasets_by_time( datasets_dict: SourceDict, t0: np.datetime64, - config: Configuration, + config: PVNetDataConfig, ) -> SourceDict: """Slice the dictionary of input data sources around a given t0 time. Args: datasets_dict: Dictionary of the input data sources t0: The init-time - config: Configuration object. + config: PVNetDataConfig object. Returns: A dictionary of the sliced input data sources. @@ -133,7 +130,7 @@ def slice_datasets_by_time( sliced_datasets_dict["nwp"] = {} for nwp_key, da_nwp in datasets_dict["nwp"].items(): - nwp_config = config.input_data.nwp[nwp_key] + nwp_config = config.nwp[nwp_key] # Add a buffer if we need to diff some of the channels in time if len(nwp_config.accum_channels)>0: @@ -155,25 +152,29 @@ def slice_datasets_by_time( ) if "sat" in datasets_dict: - sat_config = config.input_data.satellite sliced_datasets_dict["sat"] = select_time_slice( datasets_dict["sat"], t0, - time_resolution=minutes(sat_config.time_resolution_minutes), - interval_start=minutes(sat_config.interval_start_minutes), - interval_end=minutes(sat_config.interval_end_minutes), + time_resolution=minutes(config.satellite.time_resolution_minutes), + interval_start=minutes(config.satellite.interval_start_minutes), + interval_end=minutes(config.satellite.interval_end_minutes), ) if "generation" in datasets_dict: - generation_config = config.input_data.generation - - sliced_datasets_dict["generation"] = select_time_slice( - datasets_dict["generation"], - t0, - time_resolution=minutes(generation_config.time_resolution_minutes), - interval_start=minutes(generation_config.interval_start_minutes), - interval_end=minutes(generation_config.interval_end_minutes), - ) + for key, window_config in ( + ("generation_input", config.generation.input), + ("generation_target", config.generation.target), + ): + if window_config is None: + continue + + sliced_datasets_dict[key] = select_time_slice( + datasets_dict["generation"], + t0, + time_resolution=minutes(config.generation.time_resolution_minutes), + interval_start=minutes(window_config.interval_start_minutes), + interval_end=minutes(window_config.interval_end_minutes), + ) return sliced_datasets_dict diff --git a/src/ocf_data_sampler/datasets/pvnet/valid_t0s.py b/src/ocf_data_sampler/datasets/pvnet/valid_t0s.py index 8514fe12..7921bc97 100644 --- a/src/ocf_data_sampler/datasets/pvnet/valid_t0s.py +++ b/src/ocf_data_sampler/datasets/pvnet/valid_t0s.py @@ -4,7 +4,7 @@ import pandas as pd from ocf_data_sampler.common.time_utils import minutes -from ocf_data_sampler.config.model import Configuration +from ocf_data_sampler.config.model import PVNetDataConfig from ocf_data_sampler.datasets.pvnet.types import SourceDict from ocf_data_sampler.select.time_periods import ( find_contiguous_t0_periods, @@ -15,20 +15,20 @@ def find_valid_time_periods( datasets_dict: SourceDict, - config: Configuration, + config: PVNetDataConfig, ) -> pd.DataFrame: """Find the t0 times where all of the requested input data is available. Args: datasets_dict: A dictionary of input datasets - config: Configuration file + config: PVNetDataConfig file Returns: A DataFrame containing the valid t0 time periods. """ contiguous_time_periods: list[pd.DataFrame] = [] if "nwp" in datasets_dict: - for nwp_key, nwp_config in config.input_data.nwp.items(): + for nwp_key, nwp_config in config.nwp.items(): da = datasets_dict["nwp"][nwp_key] # Extract the max extents of the forecast steps @@ -68,13 +68,11 @@ def find_valid_time_periods( contiguous_time_periods.append(time_periods) if "sat" in datasets_dict: - sat_config = config.input_data.satellite - time_periods = find_contiguous_t0_periods( datasets_dict["sat"]["time_utc"].values, - time_resolution=minutes(sat_config.time_resolution_minutes), - interval_start=minutes(sat_config.interval_start_minutes), - interval_end=minutes(sat_config.interval_end_minutes), + time_resolution=minutes(config.satellite.time_resolution_minutes), + interval_start=minutes(config.satellite.interval_start_minutes), + interval_end=minutes(config.satellite.interval_end_minutes), ) contiguous_time_periods.append(time_periods) @@ -83,24 +81,30 @@ def find_valid_time_periods( raise ValueError("No valid t0 periods found for satellite data") if "generation" in datasets_dict: - generation_config = config.input_data.generation - - time_periods = find_contiguous_t0_periods( - datasets_dict["generation"]["time_utc"].values, - time_resolution=minutes(generation_config.time_resolution_minutes), - interval_start=minutes(generation_config.interval_start_minutes), - interval_end=minutes(generation_config.interval_end_minutes), - ) + for window_name, window_config in ( + ("input", config.generation.input), + ("target", config.generation.target), + ): + if window_config is None: + continue + + time_periods = find_contiguous_t0_periods( + datasets_dict["generation"]["time_utc"].values, + time_resolution=minutes(config.generation.time_resolution_minutes), + interval_start=minutes(window_config.interval_start_minutes), + interval_end=minutes(window_config.interval_end_minutes), + ) - if len(time_periods) == 0: - raise ValueError("No valid t0 periods found for generation data") + if len(time_periods) == 0: + raise ValueError( + f"No valid t0 periods found for {window_name} generation data", + ) - contiguous_time_periods.append(time_periods) + contiguous_time_periods.append(time_periods) # Find joint overlapping contiguous time periods valid_time_periods = intersect_time_periods(contiguous_time_periods) - # check there are some valid time periods if len(valid_time_periods) == 0: raise ValueError(f"No valid time periods found, {contiguous_time_periods=}") diff --git a/src/ocf_data_sampler/load/__init__.py b/src/ocf_data_sampler/load/__init__.py index f2fe9000..17bdcf17 100644 --- a/src/ocf_data_sampler/load/__init__.py +++ b/src/ocf_data_sampler/load/__init__.py @@ -1,3 +1,4 @@ from ocf_data_sampler.load.generation import open_generation +from ocf_data_sampler.load.locations import open_locations from ocf_data_sampler.load.nwp import open_nwp from ocf_data_sampler.load.satellite import open_sat_data diff --git a/src/ocf_data_sampler/load/generation.py b/src/ocf_data_sampler/load/generation.py index b55d2f52..0d7691dd 100644 --- a/src/ocf_data_sampler/load/generation.py +++ b/src/ocf_data_sampler/load/generation.py @@ -9,9 +9,6 @@ Coordinates: time_utc (time_utc): The datetimes associated with each generation and capacity value location_id (location_id): The integer IDs of the locations - longitude (location_id): The longitudes of the locations - latitude (location_id): The latitudes of the locations - """ import numpy as np @@ -43,8 +40,6 @@ def open_generation(zarr_path: str) -> xr.DataArray: coord_dtypes = { "time_utc": np.datetime64, "location_id": np.integer, - "longitude": np.number, - "latitude": np.number, } validate_coords( ds, diff --git a/src/ocf_data_sampler/load/locations.py b/src/ocf_data_sampler/load/locations.py new file mode 100644 index 00000000..03112e92 --- /dev/null +++ b/src/ocf_data_sampler/load/locations.py @@ -0,0 +1,51 @@ +"""Functions for loading locations metadata. + +Locations data schema: a CSV file with the following columns: + + location_id: The integer IDs of the locations + longitude: The longitudes of the locations + latitude: The latitudes of the locations + +Rows must be in increasing `location_id` order, and no value may be left blank (i.e. no NaNs). + +A CSV is used rather than zarr since this catalogue is small and benefits from being human +readable and hand editable. Additional columns may be included to make the file easier to work with, +but are dropped on load. +""" + +import pandas as pd + +from ocf_data_sampler.common.indexing import assert_values_unique_increasing + + +def open_locations(csv_path: str) -> pd.DataFrame: + """Open the locations metadata and validate its columns and data types. + + Args: + csv_path: Path to the locations CSV data + + Returns: + pd.DataFrame: The locations metadata, in increasing location ID order + """ + column_dtypes = { + "location_id": "int64", + "longitude": "float64", + "latitude": "float64", + } + df = pd.read_csv(csv_path, dtype=column_dtypes) + + if missing_columns := set(column_dtypes) - set(df.columns): + raise ValueError( + f"Locations data should have columns {list(column_dtypes)}, but the following " + f"were missing: {missing_columns}", + ) + + # Extra columns are allowed in the file, but are dropped on load to avoid bloat + df = df[list(column_dtypes)] + + if df.isna().any().any(): + raise ValueError("Locations data must not contain missing values") + + assert_values_unique_increasing(df["location_id"].values, "location_id") + + return df diff --git a/src/ocf_data_sampler/select/dropout.py b/src/ocf_data_sampler/select/dropout.py index 8972ce8f..5465eb3f 100644 --- a/src/ocf_data_sampler/select/dropout.py +++ b/src/ocf_data_sampler/select/dropout.py @@ -5,15 +5,16 @@ from ocf_data_sampler.common.types import TArray -def apply_history_dropout( +def apply_dropout( da: TArray, t0: np.datetime64, dropout_timedeltas: list[np.timedelta64], dropout_frac: float | list[float], ) -> TArray: - """Apply in-place random dropout to the historical part of some sequence data. + """Apply in-place random dropout to some sequence data. - Dropped out data is replaced with NaNs. + A timedelta relative to t0 is randomly sampled, and all data after that point in time is + dropped out (replaced with NaNs). This helper requires a NumPy-backed DataArray-like object. It mutates the backing array in place, so it should be called after any lazy data has been materialised. @@ -61,9 +62,11 @@ def apply_history_dropout( return da else: times = da["time_utc"].values - keep = (times <= t0 + timedelta_choice) | (times > t0) + keep = times <= t0 + timedelta_choice axis = da.dims.index("time_utc") mask = np.expand_dims(keep, tuple(i for i in range(da.data.ndim) if i != axis)) + mask = np.broadcast_to(mask, da.data.shape) + da.data = np.where(mask, da.data, np.nan) return da diff --git a/tests/config/test_config.py b/tests/config/test_config.py index bb60dad2..4b7f1719 100644 --- a/tests/config/test_config.py +++ b/tests/config/test_config.py @@ -2,34 +2,36 @@ from pydantic import ValidationError from ocf_data_sampler.config.load import load_yaml_configuration -from ocf_data_sampler.config.model import Configuration +from ocf_data_sampler.config.model import PVNetDataConfig + +_MINIMAL_SAMPLING_GRID = {"locations_csv_path": "locations.csv", "t0_resolution_minutes": 30} def _load_config_and_provider(config_path): config = load_yaml_configuration(config_path) - provider = next(iter(config.input_data.nwp.root.keys())) + provider = next(iter(config.nwp.root.keys())) return config, provider def _validate_configuration(config): """Recreate config instance from dict to trigger validation.""" - return Configuration(**config.model_dump()) + return PVNetDataConfig(**config.model_dump()) def test_default_configuration(): - """Test default pydantic class""" - _ = Configuration() + """Test default pydantic class - sampling_grid is the only required field""" + _ = PVNetDataConfig(sampling_grid=_MINIMAL_SAMPLING_GRID) def test_extra_field_error(): """ Check an extra parameters in config causes error """ - configuration = Configuration() + configuration = PVNetDataConfig(sampling_grid=_MINIMAL_SAMPLING_GRID) configuration_dict = configuration.model_dump() configuration_dict["extra_field"] = "extra_value" with pytest.raises(ValidationError, match="Extra inputs are not permitted"): - Configuration(**configuration_dict) + PVNetDataConfig(**configuration_dict) def test_incorrect_interval_start_minutes(config_filename): @@ -37,7 +39,7 @@ def test_incorrect_interval_start_minutes(config_filename): Check a history length not divisible by time resolution causes error """ configuration, provider = _load_config_and_provider(config_filename) - configuration.input_data.nwp[provider].interval_start_minutes = -1111 + configuration.nwp[provider].interval_start_minutes = -1111 with pytest.raises( ValueError, match=r"interval_start_minutes \(-1111\) " @@ -51,7 +53,7 @@ def test_incorrect_interval_end_minutes(config_filename): Check a forecast length not divisible by time resolution causes error """ configuration, provider = _load_config_and_provider(config_filename) - configuration.input_data.nwp[provider].interval_end_minutes = 1111 + configuration.nwp[provider].interval_end_minutes = 1111 with pytest.raises( ValueError, match=r"interval_end_minutes \(1111\) " @@ -65,7 +67,7 @@ def test_incorrect_nwp_provider(config_filename): Check an unexpected nwp provider causes error """ configuration, provider = _load_config_and_provider(config_filename) - configuration.input_data.nwp[provider].provider = "unexpected_provider" + configuration.nwp[provider].provider = "unexpected_provider" with pytest.raises(ValidationError, match="Unknown NWP provider"): _validate_configuration(configuration) @@ -73,11 +75,11 @@ def test_incorrect_nwp_provider(config_filename): def test_nwp_provider_is_canonicalized(config_filename): """NWP provider names are stored using their canonical lowercase spelling.""" configuration, provider = _load_config_and_provider(config_filename) - configuration.input_data.nwp[provider].provider = "UKV" + configuration.nwp[provider].provider = "UKV" validated = _validate_configuration(configuration) - assert validated.input_data.nwp[provider].provider == "ukv" + assert validated.nwp[provider].provider == "ukv" def test_incorrect_dropout(config_filename): @@ -87,12 +89,12 @@ def test_incorrect_dropout(config_filename): configuration, provider = _load_config_and_provider(config_filename) # Check that a positive number is not allowed - configuration.input_data.nwp[provider].dropout_timedeltas_minutes = [120] + configuration.nwp[provider].dropout_timedeltas_minutes = [120] with pytest.raises(Exception, match="Dropout timedeltas must be negative"): _validate_configuration(configuration) # Check that zero is allowed - configuration.input_data.nwp[provider].dropout_timedeltas_minutes = [0] + configuration.nwp[provider].dropout_timedeltas_minutes = [0] _validate_configuration(configuration) @@ -102,23 +104,23 @@ def test_incorrect_dropout_fraction(config_filename): """ configuration, provider = _load_config_and_provider(config_filename) - configuration.input_data.nwp[provider].dropout_fraction = 1.1 + configuration.nwp[provider].dropout_fraction = 1.1 with pytest.raises(ValidationError, match=r"Dropout fractions must be in range *"): _validate_configuration(configuration) - configuration.input_data.nwp[provider].dropout_fraction = -0.1 + configuration.nwp[provider].dropout_fraction = -0.1 with pytest.raises(ValidationError, match=r"Dropout fractions must be in range *"): _validate_configuration(configuration) - configuration.input_data.nwp[provider].dropout_fraction = [1.0, 0.1] + configuration.nwp[provider].dropout_fraction = [1.0, 0.1] with pytest.raises(ValidationError, match=r"The sum of dropout fractions must be in range *"): _validate_configuration(configuration) - configuration.input_data.nwp[provider].dropout_fraction = [-0.1, 1.1] + configuration.nwp[provider].dropout_fraction = [-0.1, 1.1] with pytest.raises(ValidationError, match=r"All dropout fractions must be in range *"): _validate_configuration(configuration) - configuration.input_data.nwp[provider].dropout_fraction = [] + configuration.nwp[provider].dropout_fraction = [] with pytest.raises(ValidationError, match="List cannot be empty"): _validate_configuration(configuration) @@ -127,8 +129,8 @@ def test_dropout_fraction_list_length_matches_timedeltas(config_filename): """List dropout fractions must align with dropout timedeltas one-to-one.""" configuration, provider = _load_config_and_provider(config_filename) - configuration.input_data.nwp[provider].dropout_timedeltas_minutes = [-60, -120] - configuration.input_data.nwp[provider].dropout_fraction = [0.5] + configuration.nwp[provider].dropout_timedeltas_minutes = [-60, -120] + configuration.nwp[provider].dropout_fraction = [0.5] with pytest.raises( ValidationError, @@ -142,16 +144,16 @@ def test_inconsistent_dropout_use(config_filename): Check dropout fraction outside of range causes error """ configuration = load_yaml_configuration(config_filename) - configuration.input_data.satellite.dropout_fraction = 1.0 - configuration.input_data.satellite.dropout_timedeltas_minutes = [] + configuration.satellite.dropout_fraction = 1.0 + configuration.satellite.dropout_timedeltas_minutes = [] with pytest.raises( ValueError, match="To dropout fraction > 0 requires a list of dropout timedeltas", ): _validate_configuration(configuration) - configuration.input_data.satellite.dropout_fraction = 0.0 - configuration.input_data.satellite.dropout_timedeltas_minutes = [-120, -60] + configuration.satellite.dropout_fraction = 0.0 + configuration.satellite.dropout_timedeltas_minutes = [-120, -60] with pytest.raises( ValueError, match="To use dropout timedeltas dropout fraction should be > 0", @@ -165,15 +167,41 @@ def test_accum_channels_validation(config_filename): # Test invalid channel scenario invalid_config = config.model_copy(deep=True) - invalid_nwp = invalid_config.input_data.nwp.root[nwp_name] + invalid_nwp = invalid_config.nwp.root[nwp_name] invalid_nwp.accum_channels = ["invalid_channel"] # Verify exact error message expected_error = ( - rf"input_data.nwp.{nwp_name}\n" + rf"nwp.{nwp_name}\n" fr" Value error, NWP provider '{nwp_name}': all values in 'accum_channels' " r"should be present in 'channels'\. " r"Extra values found: {'invalid_channel'}.*" ) with pytest.raises(ValidationError, match=expected_error): _validate_configuration(invalid_config) + + +def test_generation_interval_divisibility_raises(config_filename): + """generation.input/target must be divisible by generation.time_resolution_minutes.""" + configuration = load_yaml_configuration(config_filename) + configuration.generation.input.interval_start_minutes = -45 + + with pytest.raises( + ValueError, + match=r"generation\.input\.interval_start_minutes \(-45\) must be divisible by " + r"generation\.time_resolution_minutes \(30\)", + ): + _validate_configuration(configuration) + + +def test_generation_requires_input_or_target(config_filename): + """At least one of generation.input or generation.target must be configured.""" + configuration = load_yaml_configuration(config_filename) + configuration.generation.input = None + configuration.generation.target = None + + with pytest.raises( + ValueError, + match=r"At least one of `generation\.input` or `generation\.target` must be configured", + ): + _validate_configuration(configuration) diff --git a/tests/config/test_load.py b/tests/config/test_load.py index 36f971a3..e2d1b1e9 100644 --- a/tests/config/test_load.py +++ b/tests/config/test_load.py @@ -1,6 +1,6 @@ from ocf_data_sampler.config.load import load_yaml_configuration -from ocf_data_sampler.config.model import Configuration +from ocf_data_sampler.config.model import PVNetDataConfig def test_load_yaml_configuration(config_filename): - assert isinstance(load_yaml_configuration(config_filename), Configuration) + assert isinstance(load_yaml_configuration(config_filename), PVNetDataConfig) diff --git a/tests/conftest.py b/tests/conftest.py index 1f93d88e..685e242a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -14,6 +14,12 @@ NWP_FREQ = pd.Timedelta("3h") RANDOM_SEED = 42 +# The LOCATION_IDS catalog mirrors the GSPs: ID 0 is the national aggregate and IDs 1-317 are the +# regional GSPs +LOCATION_IDS = tuple(range(318)) +# The SITE_LOCATION_IDS catalog has no national aggregate since they mirror different sites +SITE_LOCATION_IDS = tuple(range(1, 11)) + UK_SAT_AREA = """msg_seviri_rss_3km: description: MSG SEVIRI Rapid Scanning Service area definition with 3 km resolution projection: @@ -62,6 +68,13 @@ def create_xr_dataset(coords, data, name, attrs=None): return da.to_dataset(name=name) +def save_csv(df, path, filename): + """Save dataframe to csv""" + csv_path = path / filename + df.to_csv(csv_path, index=False) + return str(csv_path) + + def save_zarr(ds, path, filename, chunks=None): """Save dataset to zarr""" if chunks: @@ -192,54 +205,66 @@ def nwp_cloudcasting_zarr_path(session_tmp_path, session_rng): yield save_zarr(ds, session_tmp_path, "cloudcasting.zarr", chunks) -@pytest.fixture(scope="session") -def ds_generation(session_rng): - times = pd.date_range("2023-01-01 00:00", "2023-01-02 00:00", freq="30min") - location_ids = np.arange(318) - # Rough UK bounding box +def _locations_dataframe(session_rng, location_ids): + """Build a locations catalog over the given IDs, with random points in a rough UK bbox.""" lat_min, lat_max = 49.9, 58.7 lon_min, lon_max = -8.6, 1.8 - # Generate random uniform points - longitudes = session_rng.uniform(lon_min, lon_max, len(location_ids)).astype("float64") - latitudes = session_rng.uniform(lat_min, lat_max, len(location_ids)).astype("float64") + return pd.DataFrame( + { + "location_id": location_ids, + "longitude": session_rng.uniform(lon_min, lon_max, len(location_ids)), + "latitude": session_rng.uniform(lat_min, lat_max, len(location_ids)), + }, + ) + + +@pytest.fixture(scope="session") +def df_locations(session_rng): + """The locations catalog - the source of truth for which locations are samplable.""" + return _locations_dataframe(session_rng, LOCATION_IDS) + + +@pytest.fixture(scope="session") +def df_site_locations(session_rng): + """The locations catalog for the site-level fixtures.""" + return _locations_dataframe(session_rng, SITE_LOCATION_IDS) + - capacity = np.ones((len(times), len(location_ids))) +@pytest.fixture(scope="session") +def ds_generation(session_rng, df_locations): + """Generation for every catalogued location with no missing generation data""" + + times = pd.date_range("2023-01-01 00:00", "2023-01-02 00:00", freq="30min") + location_ids = df_locations["location_id"].values + shape = (len(times), len(location_ids)) - generation = session_rng.uniform(0, 200, (len(times), len(location_ids))).astype(np.float32) + capacity = 200 + capacities = np.full(shape, fill_value=capacity, dtype="float32") + generations = session_rng.uniform(0, capacity, shape).astype("float32") - # Build Dataset return xr.Dataset( data_vars={ - "capacity_mwp": (("time_utc", "location_id"), capacity), - "generation_mw": (("time_utc", "location_id"), generation), + "capacity_mwp": (("time_utc", "location_id"), capacities), + "generation_mw": (("time_utc", "location_id"), generations), }, coords={ "time_utc": times, "location_id": location_ids, - "longitude": ("location_id", longitudes), - "latitude": ("location_id", latitudes), }, ) # location data (non overlapping time periods) and starting with id 1 @pytest.fixture(scope="session") -def ds_site_generation(session_rng): +def ds_site_generation(session_rng, df_site_locations): # Define a global time range (covers all possible site periods) global_times = pd.date_range("2023-01-01 00:00", "2023-01-02 00:00", freq="30min") n_times = len(global_times) - location_ids = np.arange(1, 11) + location_ids = df_site_locations["location_id"].values n_sites = len(location_ids) - # Rough UK bounding box - lat_min, lat_max = 49.9, 58.7 - lon_min, lon_max = -8.6, 1.8 - - longitudes = session_rng.uniform(lon_min, lon_max, n_sites).astype("float64") - latitudes = session_rng.uniform(lat_min, lat_max, n_sites).astype("float64") - # Initialize with NaNs capacity = np.full((n_times, n_sites), np.nan, dtype="float32") generation = np.full((n_times, n_sites), np.nan, dtype="float32") @@ -258,7 +283,6 @@ def ds_site_generation(session_rng): "float32", ) - # Build Dataset return xr.Dataset( data_vars={ "capacity_mwp": (("time_utc", "location_id"), capacity), @@ -267,8 +291,6 @@ def ds_site_generation(session_rng): coords={ "time_utc": global_times, "location_id": location_ids, - "longitude": ("location_id", longitudes), - "latitude": ("location_id", latitudes), }, ) @@ -283,18 +305,30 @@ def site_generation_zarr_path(session_tmp_path, ds_site_generation): yield save_zarr(ds_site_generation, session_tmp_path, "site_generation.zarr") +@pytest.fixture(scope="session") +def locations_csv_path(session_tmp_path, df_locations): + yield save_csv(df_locations, session_tmp_path, "locations.csv") + + +@pytest.fixture(scope="session") +def site_locations_csv_path(session_tmp_path, df_site_locations): + yield save_csv(df_site_locations, session_tmp_path, "site_locations.csv") + + @pytest.fixture() def pvnet_config_filename( tmp_path, config_filename, nwp_ukv_zarr_path, generation_zarr_path, + locations_csv_path, sat_zarr_path, ): config = load_yaml_configuration(config_filename) - config.input_data.nwp["ukv"].zarr_path = nwp_ukv_zarr_path - config.input_data.satellite.zarr_path = sat_zarr_path - config.input_data.generation.zarr_path = generation_zarr_path + config.nwp["ukv"].zarr_path = nwp_ukv_zarr_path + config.satellite.zarr_path = sat_zarr_path + config.generation.zarr_path = generation_zarr_path + config.sampling_grid.locations_csv_path = locations_csv_path path = tmp_path / "configuration.yaml" save_yaml_configuration(config, str(path)) @@ -307,12 +341,16 @@ def pvnet_site_config_filename( config_filename, nwp_ukv_zarr_path, site_generation_zarr_path, + site_locations_csv_path, sat_zarr_path, ): config = load_yaml_configuration(config_filename) - config.input_data.nwp["ukv"].zarr_path = nwp_ukv_zarr_path - config.input_data.satellite.zarr_path = sat_zarr_path - config.input_data.generation.zarr_path = site_generation_zarr_path + config.nwp["ukv"].zarr_path = nwp_ukv_zarr_path + config.satellite.zarr_path = sat_zarr_path + config.generation.zarr_path = site_generation_zarr_path + config.sampling_grid.locations_csv_path = site_locations_csv_path + # The site catalog has no national aggregate, so nothing to exclude + config.sampling_grid.exclude_location_ids = [] path = session_tmp_path / "configuration.yaml" save_yaml_configuration(config, str(path)) diff --git a/tests/datasets/pvnet/test_dataset.py b/tests/datasets/pvnet/test_dataset.py index 96d8ccce..53694def 100644 --- a/tests/datasets/pvnet/test_dataset.py +++ b/tests/datasets/pvnet/test_dataset.py @@ -11,8 +11,10 @@ from ocf_data_sampler.datasets.pvnet.dataset import ( PVNetConcurrentDataset, PVNetDataset, + get_locations, get_time_periods_mask, ) +from tests.conftest import LOCATION_IDS, SITE_LOCATION_IDS def _pvnet_dataset_sample_check(sample, config, batch_dim = None): @@ -24,21 +26,23 @@ def _pvnet_dataset_sample_check(sample, config, batch_dim = None): assert isinstance(sample, dict) # Specific keys should always be present - required_keys = ["nwp_ukv", "satellite", "generation", "t0", "t0_embedding"] + required_keys = [ + "nwp_ukv", "satellite", "generation_input", "generation_target", "t0", "t0_embedding", + ] for key in required_keys: assert key in sample solar_keys = ["solar_azimuth", "solar_elevation"] - if config.input_data.solar_position is not None: + if config.solar_position is not None: # Test solar position keys are present when configured for key in solar_keys: assert key in sample, f"Solar position key {key} should be present in sample" # Get expected time steps from config expected_time_steps = ( - config.input_data.solar_position.interval_end_minutes - - config.input_data.solar_position.interval_start_minutes - ) // config.input_data.solar_position.time_resolution_minutes + 1 + config.solar_position.interval_end_minutes + - config.solar_position.interval_start_minutes + ) // config.solar_position.time_resolution_minutes + 1 # Test solar angle shapes based on config assert sample["solar_azimuth"].shape == (*batch_dim, expected_time_steps) @@ -48,16 +52,35 @@ def _pvnet_dataset_sample_check(sample, config, batch_dim = None): for key in solar_keys: assert key not in sample, f"Solar position key {key} should not be present" + datetime_encoding_keys = ["date_sin", "date_cos", "time_sin", "time_cos"] + if config.datetime_encoding is not None: + # Test datetime encoding keys are present when configured + for key in datetime_encoding_keys: + assert key in sample, f"Datetime encoding key {key} should be present in sample" + + # Get expected time steps from config + expected_time_steps = ( + config.datetime_encoding.interval_end_minutes + - config.datetime_encoding.interval_start_minutes + ) // config.datetime_encoding.time_resolution_minutes + 1 + + # Test datetime encoding shapes based on config + for key in datetime_encoding_keys: + assert sample[key].shape == (*batch_dim, expected_time_steps) + else: + # Assert that datetime encoding keys are not present + for key in datetime_encoding_keys: + assert key not in sample, f"Datetime encoding key {key} should not be present" + # Check the shape of the data is correct # 30 minutes of 5 minute data (inclusive), one channel, 2x2 pixels assert sample["satellite"].shape == (*batch_dim, 7, 1, 2, 2) # 3 hours of 60 minute data (inclusive), one channel, 2x2 pixels assert sample["nwp_ukv"].shape == (*batch_dim, 4, 1, 2, 2) - # 3 hours of 30 minute data (inclusive) - assert sample["generation"].shape == (*batch_dim, 7) - # Datetime encoding keys same shape as the generation - for datetime_key in ["date_sin", "date_cos", "time_sin", "time_cos"]: - assert sample[datetime_key].shape == (*batch_dim, 7) + # generation_input: 1 hour of 30 minute data (inclusive) = 3 steps + # generation_target: 2 hours of 30 minute data (inclusive) = 5 steps + assert sample["generation_input"].shape == (*batch_dim, 3) + assert sample["generation_target"].shape == (*batch_dim, 5) # The config uses 3 periods each of which generates a sin and cos embedding assert sample["t0_embedding"].shape == (*batch_dim, 6) @@ -93,6 +116,11 @@ def test_get_time_periods_mask(): assert np.array_equal(mask, expected_mask), f"Expected {expected_mask} but got {mask}" +def _expected_num_locations(dataset, catalog_ids): + """The catalogued locations which survive the config's exclusion list.""" + return len(catalog_ids) - len(dataset.config.sampling_grid.exclude_location_ids) + + def test_pvnet_dataset(pvnet_config_filename): dataset = PVNetDataset( pvnet_config_filename, @@ -103,7 +131,7 @@ def test_pvnet_dataset(pvnet_config_filename): ) expected_t0s = 6 # 2 time periods each with 3 t0s (inclusive) at 30 minute intervals - num_locs = 317 # Quantity of regional GSPs + num_locs = _expected_num_locations(dataset, LOCATION_IDS) assert len(dataset.locations) == num_locs assert len(dataset.valid_t0_times) == expected_t0s @@ -114,6 +142,25 @@ def test_pvnet_dataset(pvnet_config_filename): _pvnet_dataset_sample_check(sample, dataset.config) +def test_get_locations_exclude_ids(locations_csv_path): + excluded_ids = [LOCATION_IDS[0], LOCATION_IDS[5], LOCATION_IDS[-1]] + locations = get_locations(locations_csv_path, exclude_ids=excluded_ids) + + assert len(locations) == len(LOCATION_IDS) - len(excluded_ids) + assert not set(excluded_ids) & {loc.id for loc in locations} + + +def test_get_locations_exclude_unknown_id(locations_csv_path): + unknown_id = max(LOCATION_IDS) + 1 + with pytest.raises(ValueError, match="not in the locations data"): + get_locations(locations_csv_path, exclude_ids=[unknown_id]) + + +def test_get_locations_exclude_all_ids(locations_csv_path): + with pytest.raises(ValueError, match=r"All location IDs .* have been excluded"): + get_locations(locations_csv_path, exclude_ids=list(LOCATION_IDS)) + + def test_pvnet_dataset_sites(pvnet_site_config_filename): dataset = PVNetDataset( pvnet_site_config_filename, @@ -124,7 +171,7 @@ def test_pvnet_dataset_sites(pvnet_site_config_filename): ) expected_t0s = 6 # 2 time periods each with 3 t0s (inclusive) at 30 minute intervals - num_locs = 10 + 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 @@ -155,14 +202,14 @@ def check_samples_equal(sample0, sample1): def test_pvnet_concurrent_dataset(pvnet_config_filename): # Create dataset object using limited set of GSPs dataset = PVNetConcurrentDataset(pvnet_config_filename) - num_gsps = 317 - assert len(dataset.locations) == num_gsps # Quantity of regional GSPs + num_locations = _expected_num_locations(dataset, LOCATION_IDS) + assert len(dataset.locations) == num_locations # NB. I have not checked the value (39 below) is in fact correct assert len(dataset.valid_t0_times) == 39 assert len(dataset) == 39 sample = dataset[0] - _pvnet_dataset_sample_check(sample, dataset.config, (num_gsps,)) + _pvnet_dataset_sample_check(sample, dataset.config, (num_locations,)) def test_pvnet_dataset_getitem_bounds(pvnet_config_filename): @@ -184,11 +231,11 @@ def test_solar_position_decoupling(tmp_path, pvnet_config_filename): """Test that solar position calculations are properly decoupled from data sources.""" config = load_yaml_configuration(pvnet_config_filename) config_without_solar = config.model_copy(deep=True) - config_without_solar.input_data.solar_position = None + config_without_solar.solar_position = None # Create version with explicit solar position configuration config_with_solar = config.model_copy(deep=True) - config_with_solar.input_data.solar_position = SolarPosition( + config_with_solar.solar_position = SolarPosition( time_resolution_minutes=30, interval_start_minutes=0, interval_end_minutes=180, @@ -217,6 +264,29 @@ def test_solar_position_decoupling(tmp_path, pvnet_config_filename): assert key in sample_with_solar, f"Solar key {key} should be in sample" +def test_pvnet_dataset_without_generation(tmp_path, pvnet_config_filename): + """Test that a dataset can be built with locations/NWP/satellite but no generation at all.""" + config = load_yaml_configuration(pvnet_config_filename) + config.generation = None + + config_path = tmp_path / "config_without_generation.yaml" + save_yaml_configuration(config, config_path) + + dataset = PVNetDataset(config_path) + + # With no generation data, there's nothing to be incomplete about + assert dataset.complete_generation + + # All locations from the locations catalog are available - none to filter out + assert len(dataset.locations) == _expected_num_locations(dataset, LOCATION_IDS) + + sample = dataset[0] + assert "generation_input" not in sample + assert "generation_target" not in sample + assert "nwp_ukv" in sample + assert "satellite" in sample + + def test_pvnet_dataset_raw_sample_iteration(pvnet_config_filename): """Tests iterating raw samples (dict of tensors) from PVNetDataset""" dataset = PVNetDataset(pvnet_config_filename) @@ -239,9 +309,14 @@ def test_pvnet_dataset_raw_sample_iteration(pvnet_config_filename): required_keys = [ "nwp_ukv", "satellite", - "generation", + "generation_input", + "generation_target", "solar_azimuth", "solar_elevation", + "date_sin", + "date_cos", + "time_sin", + "time_cos", "location_id", ] for key in required_keys: @@ -249,7 +324,8 @@ def test_pvnet_dataset_raw_sample_iteration(pvnet_config_filename): # Type assertions assert isinstance(raw_sample["satellite"], torch.Tensor) - assert isinstance(raw_sample["generation"], torch.Tensor) + assert isinstance(raw_sample["generation_input"], torch.Tensor) + assert isinstance(raw_sample["generation_target"], torch.Tensor) assert isinstance(raw_sample["solar_azimuth"], torch.Tensor) assert isinstance(raw_sample["solar_elevation"], torch.Tensor) assert isinstance(raw_sample["nwp_ukv"], torch.Tensor) @@ -257,16 +333,25 @@ def test_pvnet_dataset_raw_sample_iteration(pvnet_config_filename): # Shape assertions assert raw_sample["satellite"].shape == (7, 1, 2, 2) assert raw_sample["nwp_ukv"].shape == (4, 1, 2, 2) - assert raw_sample["generation"].shape == (7,) + assert raw_sample["generation_input"].shape == (3,) + assert raw_sample["generation_target"].shape == (5,) # Solar position shapes - no batch dimension - solar_config = dataset.config.input_data.solar_position + solar_config = dataset.config.solar_position expected_time_steps = ( solar_config.interval_end_minutes - solar_config.interval_start_minutes ) // solar_config.time_resolution_minutes + 1 assert raw_sample["solar_azimuth"].shape == (expected_time_steps,) assert raw_sample["solar_elevation"].shape == (expected_time_steps,) + # Datetime encoding shapes - no batch dimension + dt_config = dataset.config.datetime_encoding + expected_dt_time_steps = ( + dt_config.interval_end_minutes - dt_config.interval_start_minutes + ) // dt_config.time_resolution_minutes + 1 + for key in ("date_sin", "date_cos", "time_sin", "time_cos"): + assert raw_sample[key].shape == (expected_dt_time_steps,) + assert isinstance(raw_sample["location_id"], int | np.integer) diff --git a/tests/datasets/pvnet/test_preprocess.py b/tests/datasets/pvnet/test_preprocess.py index 63dd9e9c..245bbc7c 100644 --- a/tests/datasets/pvnet/test_preprocess.py +++ b/tests/datasets/pvnet/test_preprocess.py @@ -1,80 +1,131 @@ import numpy as np +import pandas as pd import xarray as xr from ocf_data_sampler.config import load_yaml_configuration from ocf_data_sampler.datasets.pvnet.preprocess import ( apply_dropout_to_datasets, fill_nans_in_dataset_dicts, + normalise_dataset_dicts, ) +def _generation_da(generation_mw: np.ndarray, capacity_mwp: np.ndarray) -> xr.DataArray: + """Build a generation DataArray from 2D `(time_utc, location_id)` MW and capacity arrays.""" + n_times, n_locations = generation_mw.shape + return xr.DataArray( + np.stack([generation_mw, capacity_mwp], axis=-1), + coords={ + "time_utc": pd.date_range("2023-01-01", periods=n_times, freq="30min"), + "location_id": np.arange(1, n_locations + 1), + "gen_param": ["generation_mw", "capacity_mwp"], + }, + dims=("time_utc", "location_id", "gen_param"), + ) + + +def _assert_dropout_applied(da: xr.DataArray, cutoff_time: np.datetime64): + """Assert that all values after the cutoff time are NaN, and all values before are not NaN. + """ + assert not np.any(np.isnan(da.sel(time_utc=slice(None, cutoff_time)))) + assert np.all(np.isnan(da.sel(time_utc=slice(cutoff_time + np.timedelta64(1, "s"), None)))) + + def test_fill_nans_in_dataset_dicts(config_filename): """Test the fill_nans_in_arrays function from configuration""" configuration = load_yaml_configuration(config_filename) # Set custom satellite and nwp values, generation is left as default 0.0 - configuration.input_data.satellite.dropout_fill_value = -1.0 - configuration.input_data.nwp["ukv"].dropout_fill_value = -2.0 + configuration.satellite.dropout_fill_value = -1.0 + configuration.nwp["ukv"].dropout_fill_value = -2.0 gen = np.array([1.0, np.nan, 3.0, np.nan]) sat = np.array([1.0, np.nan, 3.0, np.nan]) ukv = np.array([np.nan, 3.0, np.nan]) datasets_dict = { - "generation": xr.DataArray(gen), + "generation_input": xr.DataArray(gen), + "generation_target": xr.DataArray(gen.copy()), "sat": xr.DataArray(sat), "nwp": {"ukv": xr.DataArray(ukv)}, } datasets_dict = fill_nans_in_dataset_dicts(datasets_dict, config=configuration) - assert np.array_equal(datasets_dict["generation"].values, np.array([1.0, 0.0, 3.0, 0.0])) + expected_gen = np.array([1.0, 0.0, 3.0, 0.0]) + assert np.array_equal(datasets_dict["generation_input"].values, expected_gen) + assert np.array_equal(datasets_dict["generation_target"].values, expected_gen) assert np.array_equal(datasets_dict["sat"].values, np.array([1.0, -1.0, 3.0, -1.0])) assert np.array_equal(datasets_dict["nwp"]["ukv"].values, np.array([-2.0, 3.0, -2.0])) +def test_normalise_dataset_dicts_generation(): + """Generation is normalised to a capacity factor (generation_mw / capacity_mwp).""" + capacity_mwp = np.array([[100.0, 0.0], [100.0, 40.0]]) + generation = _generation_da(np.array([[50.0, 0.0], [100.0, 20.0]]), capacity_mwp) + + datasets_dict = { + "generation_input": generation, + "generation_target": generation.copy(deep=True), + } + + datasets_dict = normalise_dataset_dicts(datasets_dict, {}, {}, {}, {}) + + for key in ("generation_input", "generation_target"): + result = datasets_dict[key] + + # generation_mw is rescaled element-wise by capacity_mwp + normalised = result.sel(gen_param="generation_mw").values + assert normalised[0, 0] == 0.5 + assert normalised[1, 0] == 1.0 + assert normalised[1, 1] == 0.5 + + # Zero capacity normalises to 0 rather than raw MW or a division error + assert normalised[0, 1] == 0.0 + + # capacity_mwp itself is left unchanged + assert np.array_equal(result.sel(gen_param="capacity_mwp").values, capacity_mwp) + + def test_apply_dropout_to_datasets(pvnet_config_filename): config = load_yaml_configuration(pvnet_config_filename) - # Set dropout - config.input_data.generation.dropout_timedeltas_minutes = [-30] - config.input_data.generation.dropout_fraction = 1.0 - config.input_data.satellite.dropout_timedeltas_minutes = [] - config.input_data.satellite.dropout_fraction = 0 - - t0 = np.datetime64("2023-01-01 12:00") - times = np.array( - [ - "2023-01-01T11:00", - "2023-01-01T11:30", - "2023-01-01T12:00", - "2023-01-01T12:30", - ], - dtype="datetime64[m]", - ) - generation = xr.DataArray( - np.arange(4 * 2, dtype=float).reshape(4, 2), - coords={"time_utc": times, "location_id": [1, 2]}, - dims=("time_utc", "location_id"), - ) + config.generation.input.dropout_timedeltas_minutes = [-30] + config.generation.input.dropout_fraction = 1.0 + + config.satellite.dropout_timedeltas_minutes = [-60] + config.satellite.dropout_fraction = 1.0 + + generation_input = _generation_da(generation_mw=np.ones((4, 2)), capacity_mwp=np.ones((4, 2))) + generation_target = generation_input.copy(deep=True) + + t0 = np.datetime64("2023-01-01 00:00") + + times = pd.date_range(end=t0, periods=4, freq="30min").values + sat = xr.DataArray( np.arange(4, dtype=float), coords={"time_utc": times}, dims=("time_utc",), ) - datasets_dict = {"generation": generation, "sat": sat} + datasets_dict = { + "generation_input": generation_input, + "generation_target": generation_target, + "sat": sat + } apply_dropout_to_datasets(datasets_dict, t0, config) - ds_gen = datasets_dict["generation"] + # Generation dropout with a -30 minute cutoff should blank everything from t0 + _assert_dropout_applied(datasets_dict["generation_input"], t0 - np.timedelta64(30, "m")) + + # No dropout is applied to generation.target, so it should have no NaNs + assert not np.any(np.isnan(datasets_dict["generation_target"])) + + # Satellite dropout with a -60 minute cutoff should blank everything from t0 - 60 minutes + _assert_dropout_applied(datasets_dict["sat"], t0 - np.timedelta64(60, "m")) - # Generation dropout with a -30 minute history should blank only the t0 timestep. - assert not np.any(np.isnan(ds_gen.sel(time_utc=slice(None, "2023-01-01T11:30")))) - assert np.all(np.isnan(ds_gen.sel(time_utc=t0))) - assert not np.any(np.isnan(ds_gen.sel(time_utc=slice("2023-01-01T12:30", None)))) - # Satellite dropout is disabled, so the helper should leave it untouched. - xr.testing.assert_equal(datasets_dict["sat"], sat) diff --git a/tests/datasets/pvnet/test_sample.py b/tests/datasets/pvnet/test_sample.py index d876f515..33d5dd93 100644 --- a/tests/datasets/pvnet/test_sample.py +++ b/tests/datasets/pvnet/test_sample.py @@ -19,37 +19,38 @@ def test_make_sun_position_numpy_sample(): def test_convert_generation_to_numpy_sample(generation_zarr_path): + """convert_to_numpy_sample just extracts generation_mw/capacity_mwp as-is - normalising + generation_mw to a capacity factor is normalise_dataset_dicts's job (see test_preprocess.py). + """ da = open_generation(generation_zarr_path).isel(time_utc=slice(0, 10)).sel(location_id=1) - t0_idx = 0 - numpy_sample = convert_to_numpy_sample({"generation": da}, t0_idx=t0_idx) + numpy_sample = convert_to_numpy_sample({"generation_input": da, "generation_target": da}) - # Assert structure - assert isinstance(numpy_sample, dict) - assert "generation" in numpy_sample - assert "capacity_mwp" in numpy_sample - assert "generation_time_utc" in numpy_sample + generation_mw = da.sel(gen_param="generation_mw").values + capacity = da.sel(gen_param="capacity_mwp").values + + for key in ("generation_input", "generation_target"): + # Assert structure + assert isinstance(numpy_sample, dict) + assert key in numpy_sample + assert f"{key}_capacity_mwp" in numpy_sample + assert f"{key}_time_utc" in numpy_sample - # Assert content and capacity values - assert np.array_equal(numpy_sample["generation"], da.sel(gen_param="generation_mw").values) - assert isinstance(numpy_sample["generation_time_utc"], np.ndarray) - assert numpy_sample["generation_time_utc"].dtype == float - assert numpy_sample["capacity_mwp"] == da.sel(gen_param="capacity_mwp").isel(time_utc=0).values + # Assert content is passed through unchanged + assert np.array_equal(numpy_sample[key], generation_mw) + assert np.array_equal(numpy_sample[f"{key}_capacity_mwp"], capacity) + assert isinstance(numpy_sample[f"{key}_time_utc"], np.ndarray) + assert numpy_sample[f"{key}_time_utc"].dtype == float def test_convert_nwp_to_numpy_sample(ds_nwp_ukv_time_sliced): - t0_idx = 0 - numpy_sample = convert_to_numpy_sample( - {"nwp": {"ukv": ds_nwp_ukv_time_sliced}}, - t0_idx=t0_idx, - ) + numpy_sample = convert_to_numpy_sample({"nwp": {"ukv": ds_nwp_ukv_time_sliced}}) assert isinstance(numpy_sample, dict) assert (numpy_sample["nwp_ukv"] == ds_nwp_ukv_time_sliced.values).all() def test_convert_satellite_to_numpy_sample(da_sat_like): - t0_idx = 0 - numpy_sample = convert_to_numpy_sample({"sat": da_sat_like}, t0_idx=t0_idx) + numpy_sample = convert_to_numpy_sample({"sat": da_sat_like}) assert isinstance(numpy_sample, dict) assert (numpy_sample["satellite"] == da_sat_like.values).all() diff --git a/tests/fixtures/configs/pvnet_test_config.yaml b/tests/fixtures/configs/pvnet_test_config.yaml index 38664772..34a398e3 100644 --- a/tests/fixtures/configs/pvnet_test_config.yaml +++ b/tests/fixtures/configs/pvnet_test_config.yaml @@ -1,58 +1,66 @@ -general: - description: Test config for PVNet - name: pvnet_test +sampling_grid: + locations_csv_path: set_in_temp_file + t0_resolution_minutes: 30 + exclude_location_ids: [0] -input_data: - - generation: - zarr_path: set_in_temp_file +generation: + zarr_path: set_in_temp_file + time_resolution_minutes: 30 + input: interval_start_minutes: -60 - interval_end_minutes: 120 - time_resolution_minutes: 30 + interval_end_minutes: 0 dropout_timedeltas_minutes: [] dropout_fraction: 0 + target: + interval_start_minutes: 0 + interval_end_minutes: 120 - nwp: - ukv: - provider: ukv - zarr_path: set_in_temp_file - interval_start_minutes: -60 - interval_end_minutes: 120 - time_resolution_minutes: 60 - channels: - - t # 2-metre temperature - image_size_pixels_height: 2 - image_size_pixels_width: 2 - dropout_timedeltas_minutes: [-180] - dropout_fraction: 1.0 - max_staleness_minutes: null - normalisation_constants: - t: - mean: 283.64913206 - std: 4.38818501 - clip_min: 270 - clip_max: 310 - - satellite: +nwp: + ukv: + provider: ukv zarr_path: set_in_temp_file - interval_start_minutes: -30 - interval_end_minutes: 0 - time_resolution_minutes: 5 + interval_start_minutes: -60 + interval_end_minutes: 120 + time_resolution_minutes: 60 channels: - - IR_016 + - t # 2-metre temperature image_size_pixels_height: 2 image_size_pixels_width: 2 - dropout_timedeltas_minutes: [] - dropout_fraction: 0 + dropout_timedeltas_minutes: [-180] + dropout_fraction: 1.0 + max_staleness_minutes: null normalisation_constants: - IR_016: - mean: 0.17594202 - std: 0.21462157 + t: + mean: 283.64913206 + std: 4.38818501 + clip_min: 270 + clip_max: 310 + +satellite: + zarr_path: set_in_temp_file + interval_start_minutes: -30 + interval_end_minutes: 0 + time_resolution_minutes: 5 + channels: + - IR_016 + image_size_pixels_height: 2 + image_size_pixels_width: 2 + dropout_timedeltas_minutes: [] + dropout_fraction: 0 + normalisation_constants: + IR_016: + mean: 0.17594202 + std: 0.21462157 + +solar_position: + interval_start_minutes: -15 + interval_end_minutes: 15 + time_resolution_minutes: 5 - solar_position: - interval_start_minutes: -15 - interval_end_minutes: 15 - time_resolution_minutes: 5 +datetime_encoding: + interval_start_minutes: -15 + interval_end_minutes: 15 + time_resolution_minutes: 5 - t0_embedding: - embeddings: [["1h", "cyclic"], ["24h", "cyclic"], ["1y", "cyclic"]] +t0_embedding: + embeddings: [["1h", "cyclic"], ["24h", "cyclic"], ["1y", "cyclic"]] diff --git a/tests/load/test_load_generation.py b/tests/load/test_load_generation.py index c425a2f1..83d6439e 100644 --- a/tests/load/test_load_generation.py +++ b/tests/load/test_load_generation.py @@ -6,6 +6,7 @@ import xarray as xr from ocf_data_sampler.load.generation import open_generation +from tests.conftest import LOCATION_IDS def test_open_generation(generation_zarr_path): @@ -14,8 +15,8 @@ def test_open_generation(generation_zarr_path): assert isinstance(da, xr.DataArray) assert da.dims == ("time_utc", "location_id", "gen_param") - assert {"longitude", "latitude"}.issubset(da.coords) - assert da.shape == (49, 318, 2) + # 24 hours of 30 minute data (inclusive), every catalogued location, capacity + generation + assert da.shape == (49, len(LOCATION_IDS), 2) assert len(np.unique(da.coords["location_id"])) == da.shape[1] diff --git a/tests/load/test_load_locations.py b/tests/load/test_load_locations.py new file mode 100644 index 00000000..197e0bb5 --- /dev/null +++ b/tests/load/test_load_locations.py @@ -0,0 +1,92 @@ +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest + +from ocf_data_sampler.load.locations import open_locations + + +def _write_csv(tmp_path: Path, df: pd.DataFrame) -> str: + csv_path = tmp_path / "locations.csv" + df.to_csv(csv_path, index=False) + return str(csv_path) + + +def test_open_locations(locations_csv_path): + """Test the locations data loader with valid data.""" + df = open_locations(locations_csv_path) + + assert isinstance(df, pd.DataFrame) + assert list(df.columns) == ["location_id", "longitude", "latitude"] + assert df["location_id"].dtype == "int64" + assert df["longitude"].dtype == "float64" + assert df["latitude"].dtype == "float64" + assert df["location_id"].is_unique + + +def test_open_locations_unsorted_ids(tmp_path: Path): + """The file must be kept in increasing location ID order.""" + csv_path = _write_csv( + tmp_path, + pd.DataFrame( + { + "location_id": [7, 2, 5], + "longitude": [1.0, 0.0, 0.5], + "latitude": [51.0, 50.0, 50.5], + }, + ), + ) + + with pytest.raises(ValueError, match="location_id must be strictly increasing"): + open_locations(csv_path) + + +def test_open_locations_missing_value(tmp_path: Path): + """A row with a blank coordinate is rejected rather than loaded as NaN.""" + csv_path = _write_csv( + tmp_path, + pd.DataFrame({"location_id": [1], "longitude": [0.0], "latitude": [np.nan]}), + ) + + with pytest.raises(ValueError, match="must not contain missing values"): + open_locations(csv_path) + + +def test_open_locations_missing_column(tmp_path: Path): + """Test that open_locations raises a ValueError when a required column is missing.""" + csv_path = _write_csv( + tmp_path, + pd.DataFrame({"location_id": [1, 2], "longitude": [0.0, 1.0]}), + ) + + with pytest.raises(ValueError, match="Locations data should have columns"): + open_locations(csv_path) + + +def test_open_locations_duplicate_ids(tmp_path: Path): + """Test that open_locations raises a ValueError on repeated location IDs.""" + csv_path = _write_csv( + tmp_path, + pd.DataFrame( + { + "location_id": [1, 1], + "longitude": [0.0, 1.0], + "latitude": [50.0, 51.0], + }, + ), + ) + + with pytest.raises(ValueError, match="location_id must be strictly increasing"): + open_locations(csv_path) + + +def test_open_locations_non_integer_id(tmp_path: Path): + """Test that open_locations rejects location IDs which are not integers.""" + csv_path = _write_csv( + tmp_path, + pd.DataFrame({"location_id": [1.5], "longitude": [0.0], "latitude": [50.0]}), + ) + + with pytest.raises(ValueError, match="int64"): + open_locations(csv_path) diff --git a/tests/select/test_dropout.py b/tests/select/test_dropout.py index 4652f31d..5220d226 100644 --- a/tests/select/test_dropout.py +++ b/tests/select/test_dropout.py @@ -3,17 +3,17 @@ import xarray as xr from ocf_data_sampler.common.time_utils import minutes -from ocf_data_sampler.select.dropout import apply_history_dropout +from ocf_data_sampler.select.dropout import apply_dropout -def test_apply_history_dropout_multiple_timedeltas(da_sample): +def test_apply_dropout_multiple_timedeltas(da_sample): # Dropout edits the input in-place, so make a copy to avoid affecting other tests da_sample = da_sample.copy(deep=True) t0 = da_sample["time_utc"].values[-1] - da_sample_dropout = apply_history_dropout( + da_sample_dropout = apply_dropout( da_sample, t0, dropout_timedeltas=minutes([-30, -45]), @@ -30,14 +30,14 @@ def test_apply_history_dropout_multiple_timedeltas(da_sample): ) -def test_apply_history_dropout_none(da_sample): +def test_apply_dropout_none(da_sample): # Dropout edits the input in-place, so make a copy to avoid affecting other tests da_sample = da_sample.copy(deep=True) t0 = da_sample["time_utc"].values[-1] - da_sample_dropout = apply_history_dropout( + da_sample_dropout = apply_dropout( da_sample, t0, dropout_timedeltas=[minutes(-30)], @@ -45,7 +45,7 @@ def test_apply_history_dropout_none(da_sample): ) xr.testing.assert_equal(da_sample_dropout, da_sample) - da_sample_dropout = apply_history_dropout( + da_sample_dropout = apply_dropout( da_sample, t0, dropout_timedeltas=[], @@ -54,14 +54,14 @@ def test_apply_history_dropout_none(da_sample): xr.testing.assert_equal(da_sample_dropout, da_sample) -def test_apply_history_dropout_list(da_sample): +def test_apply_dropout_list(da_sample): # Dropout edits the input in-place, so make a copy to avoid affecting other tests da_sample = da_sample.copy(deep=True) t0 = da_sample["time_utc"].values[-1] - da_sample_dropout = apply_history_dropout( + da_sample_dropout = apply_dropout( da_sample, t0, dropout_timedeltas=minutes([-30, -45]), @@ -79,7 +79,7 @@ def test_apply_history_dropout_list(da_sample): @pytest.mark.parametrize("t0_str", ["12:30", "13:00", "13:30"]) -def test_apply_history_dropout(da_sample, t0_str): +def test_apply_dropout(da_sample, t0_str): # Dropout edits the input in-place, so make a copy to avoid affecting other tests da_sample = da_sample.copy(deep=True) @@ -87,12 +87,13 @@ def test_apply_history_dropout(da_sample, t0_str): t0_time = pd.Timestamp(f"2024-01-01 {t0_str}") dropout_time = t0_time + minutes(-30) - da_dropout = apply_history_dropout( + da_dropout = apply_dropout( da_sample, t0_time, dropout_timedeltas=[minutes(-30)], dropout_frac=1.0, ) + # Everything after the dropout cut-off is masked, including any data beyond t0. assert da_dropout.sel(time_utc=slice(None, dropout_time)).notnull().all() - assert da_dropout.sel(time_utc=slice(dropout_time + minutes(5), t0_time)).isnull().all() + assert da_dropout.sel(time_utc=slice(dropout_time + minutes(5), None)).isnull().all()