ocf-data-sampler 0.2.7__py3-none-any.whl → 0.2.9__py3-none-any.whl

ocf_data_sampler/select/__init__.py

@@ -4,5 +4,3 @@ from .find_contiguous_time_periods import (
     intersection_of_multiple_dataframes_of_periods,
 )
 from .location import Location
-from .spatial_slice_for_dataset import slice_datasets_by_space
-from .time_slice_for_dataset import slice_datasets_by_time

ocf_data_sampler/select/dropout.py

@@ -33,7 +33,7 @@ def draw_dropout_time(
         raise ValueError("dropout_frac must be between 0 and 1 inclusive")
 
     if (len(dropout_timedeltas) == 0) or (np.random.uniform() >= dropout_frac):
-        dropout_time = t0
+        dropout_time = None
     else:
         dropout_time = t0 + np.random.choice(dropout_timedeltas)
 
@@ -42,7 +42,7 @@ def draw_dropout_time(
 
 def apply_dropout_time(
     ds: xr.DataArray,
-    dropout_time: pd.Timestamp,
+    dropout_time: pd.Timestamp | None,
 ) -> xr.DataArray:
     """Apply dropout time to the data.
 
@@ -50,5 +50,8 @@ def apply_dropout_time(
         ds: Xarray DataArray with 'time_utc' coordinate
         dropout_time: Time after which data is set to NaN
     """
-    # This replaces the times after the dropout with NaNs
-    return ds.where(ds.time_utc <= dropout_time)
+    if dropout_time is None:
+        return ds
+    else:
+        # This replaces the times after the dropout with NaNs
+        return ds.where(ds.time_utc <= dropout_time)

ocf_data_sampler/torch_datasets/datasets/pvnet_uk.py

@@ -17,16 +17,17 @@ from ocf_data_sampler.numpy_sample import (
     make_sun_position_numpy_sample,
 )
 from ocf_data_sampler.numpy_sample.collate import stack_np_samples_into_batch
+from ocf_data_sampler.numpy_sample.common_types import NumpyBatch, NumpySample
 from ocf_data_sampler.numpy_sample.gsp import GSPSampleKey
 from ocf_data_sampler.numpy_sample.nwp import NWPSampleKey
-from ocf_data_sampler.select import (
-    Location,
-    fill_time_periods,
+from ocf_data_sampler.select import Location, fill_time_periods
+from ocf_data_sampler.select.geospatial import osgb_to_lon_lat
+from ocf_data_sampler.torch_datasets.utils import (
+    channel_dict_to_dataarray,
+    find_valid_time_periods,
     slice_datasets_by_space,
     slice_datasets_by_time,
 )
-from ocf_data_sampler.select.geospatial import osgb_to_lon_lat
-from ocf_data_sampler.torch_datasets.utils import channel_dict_to_dataarray, find_valid_time_periods
 from ocf_data_sampler.torch_datasets.utils.merge_and_fill_utils import (
     fill_nans_in_arrays,
     merge_dicts,
@@ -36,99 +37,6 @@ from ocf_data_sampler.utils import minutes
 xr.set_options(keep_attrs=True)
 
 
-def process_and_combine_datasets(
-    dataset_dict: dict,
-    config: Configuration,
-    t0: pd.Timestamp,
-    location: Location,
-) -> dict:
-    """Normalise and convert data to numpy arrays."""
-    numpy_modalities = []
-
-    if "nwp" in dataset_dict:
-        nwp_numpy_modalities = {}
-
-        for nwp_key, da_nwp in dataset_dict["nwp"].items():
-
-            # Standardise and convert to NumpyBatch
-
-            da_channel_means = channel_dict_to_dataarray(
-                config.input_data.nwp[nwp_key].channel_means,
-            )
-            da_channel_stds = channel_dict_to_dataarray(
-                config.input_data.nwp[nwp_key].channel_stds,
-            )
-
-            da_nwp = (da_nwp - da_channel_means) / da_channel_stds
-
-            nwp_numpy_modalities[nwp_key] = convert_nwp_to_numpy_sample(da_nwp)
-
-        # Combine the NWPs into NumpyBatch
-        numpy_modalities.append({NWPSampleKey.nwp: nwp_numpy_modalities})
-
-    if "sat" in dataset_dict:
-        da_sat = dataset_dict["sat"]
-
-        # Standardise and convert to NumpyBatch
-        da_channel_means = channel_dict_to_dataarray(config.input_data.satellite.channel_means)
-        da_channel_stds = channel_dict_to_dataarray(config.input_data.satellite.channel_stds)
-
-        da_sat = (da_sat - da_channel_means) / da_channel_stds
-
-        numpy_modalities.append(convert_satellite_to_numpy_sample(da_sat))
-
-    if "gsp" in dataset_dict:
-        gsp_config = config.input_data.gsp
-        da_gsp = dataset_dict["gsp"]
-        da_gsp = da_gsp / da_gsp.effective_capacity_mwp
-
-        # Convert to NumpyBatch
-        numpy_modalities.append(
-            convert_gsp_to_numpy_sample(
-                da_gsp,
-                t0_idx=-gsp_config.interval_start_minutes / gsp_config.time_resolution_minutes,
-            ),
-        )
-
-    # Add GSP location data
-    numpy_modalities.append(
-        {
-            GSPSampleKey.gsp_id: location.id,
-            GSPSampleKey.x_osgb: location.x,
-            GSPSampleKey.y_osgb: location.y,
-        },
-    )
-
-    # Only add solar position if explicitly configured
-    has_solar_config = (
-        hasattr(config.input_data, "solar_position") and
-        config.input_data.solar_position is not None
-    )
-
-    if has_solar_config:
-        solar_config = config.input_data.solar_position
-
-        # Create datetime range for solar position calculation
-        datetimes = pd.date_range(
-            t0 + minutes(solar_config.interval_start_minutes),
-            t0 + minutes(solar_config.interval_end_minutes),
-            freq=minutes(solar_config.time_resolution_minutes),
-        )
-
-        # Convert OSGB coordinates to lon/lat
-        lon, lat = osgb_to_lon_lat(location.x, location.y)
-
-        # Calculate solar positions and add to modalities
-        solar_positions = make_sun_position_numpy_sample(datetimes, lon, lat)
-        numpy_modalities.append(solar_positions)
-
-    # Combine all the modalities and fill NaNs
-    combined_sample = merge_dicts(numpy_modalities)
-    combined_sample = fill_nans_in_arrays(combined_sample)
-
-    return combined_sample
-
-
 def compute(xarray_dict: dict) -> dict:
     """Eagerly load a nested dictionary of xarray DataArrays."""
     for k, v in xarray_dict.items():
@@ -139,25 +47,12 @@ def compute(xarray_dict: dict) -> dict:
     return xarray_dict
 
 
-def find_valid_t0_times(datasets_dict: dict, config: Configuration) -> pd.DatetimeIndex:
-    """Find the t0 times where all of the requested input data is available.
+def get_gsp_locations(gsp_ids: list[int] | None = None) -> list[Location]:
+    """Get list of locations of all GSPs.
 
     Args:
-        datasets_dict: A dictionary of input datasets
-        config: Configuration file
+        gsp_ids: List of GSP IDs to include. Defaults to all
     """
-    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.gsp.time_resolution_minutes),
-    )
-    return valid_t0_times
-
-
-def get_gsp_locations(gsp_ids: list[int] | None = None) -> list[Location]:
-    """Get list of locations of all GSPs."""
     if gsp_ids is None:
         gsp_ids = list(range(1, 318))
 
@@ -181,8 +76,8 @@ def get_gsp_locations(gsp_ids: list[int] | None = None) -> list[Location]:
     return locations
 
 
-class PVNetUKRegionalDataset(Dataset):
-    """A torch Dataset for creating PVNet UK regional samples."""
+class AbstractPVNetUKDataset(Dataset):
+    """Abstract class for PVNet UK datasets."""
 
     def __init__(
         self,
@@ -191,7 +86,7 @@ class PVNetUKRegionalDataset(Dataset):
         end_time: str | None = None,
         gsp_ids: list[int] | None = None,
     ) -> None:
-        """A torch Dataset for creating PVNet UK GSP samples.
+        """A torch Dataset for creating PVNet UK samples.
 
         Args:
             config_filename: Path to the configuration file
@@ -199,13 +94,11 @@ class PVNetUKRegionalDataset(Dataset):
             end_time: Limit the init-times to be before this
             gsp_ids: List of GSP IDs to create samples for. Defaults to all
         """
-        # config = load_yaml_configuration(config_filename)
-        config: Configuration = load_yaml_configuration(config_filename)
-
+        config = load_yaml_configuration(config_filename)
         datasets_dict = get_dataset_dict(config.input_data)
 
         # Get t0 times where all input data is available
-        valid_t0_times = find_valid_t0_times(datasets_dict, config)
+        valid_t0_times = self.find_valid_t0_times(datasets_dict, config)
 
         # Filter t0 times to given range
         if start_time is not None:
@@ -215,35 +108,167 @@ class PVNetUKRegionalDataset(Dataset):
             valid_t0_times = valid_t0_times[valid_t0_times <= pd.Timestamp(end_time)]
 
         # Construct list of locations to sample from
-        locations = get_gsp_locations(gsp_ids)
+        self.locations = get_gsp_locations(gsp_ids)
+        self.valid_t0_times = valid_t0_times
+
+        # Assign config and input data to self
+        self.config = config
+        self.datasets_dict = datasets_dict
+
+
+    @staticmethod
+    def process_and_combine_datasets(
+        dataset_dict: dict,
+        config: Configuration,
+        t0: pd.Timestamp,
+        location: Location,
+    ) -> NumpySample:
+        """Normalise and convert data to numpy arrays.
+
+        Args:
+            dataset_dict: Dictionary of xarray datasets
+            config: Configuration object
+            t0: init-time for sample
+            location: location of the sample
+        """
+        numpy_modalities = []
+
+        if "nwp" in dataset_dict:
+            nwp_numpy_modalities = {}
+
+            for nwp_key, da_nwp in dataset_dict["nwp"].items():
+
+                # Standardise and convert to NumpyBatch
+
+                da_channel_means = channel_dict_to_dataarray(
+                    config.input_data.nwp[nwp_key].channel_means,
+                )
+                da_channel_stds = channel_dict_to_dataarray(
+                    config.input_data.nwp[nwp_key].channel_stds,
+                )
+
+                da_nwp = (da_nwp - da_channel_means) / da_channel_stds
+
+                nwp_numpy_modalities[nwp_key] = convert_nwp_to_numpy_sample(da_nwp)
+
+            # Combine the NWPs into NumpyBatch
+            numpy_modalities.append({NWPSampleKey.nwp: nwp_numpy_modalities})
+
+        if "sat" in dataset_dict:
+            da_sat = dataset_dict["sat"]
+
+            # Standardise and convert to NumpyBatch
+            da_channel_means = channel_dict_to_dataarray(config.input_data.satellite.channel_means)
+            da_channel_stds = channel_dict_to_dataarray(config.input_data.satellite.channel_stds)
+
+            da_sat = (da_sat - da_channel_means) / da_channel_stds
+
+            numpy_modalities.append(convert_satellite_to_numpy_sample(da_sat))
+
+        if "gsp" in dataset_dict:
+            gsp_config = config.input_data.gsp
+            da_gsp = dataset_dict["gsp"]
+            da_gsp = da_gsp / da_gsp.effective_capacity_mwp
+
+            # Convert to NumpyBatch
+            numpy_modalities.append(
+                convert_gsp_to_numpy_sample(
+                    da_gsp,
+                    t0_idx=-gsp_config.interval_start_minutes / gsp_config.time_resolution_minutes,
+                ),
+            )
+
+        # Add GSP location data
+        numpy_modalities.append(
+            {
+                GSPSampleKey.gsp_id: location.id,
+                GSPSampleKey.x_osgb: location.x,
+                GSPSampleKey.y_osgb: location.y,
+            },
+        )
+
+        # Only add solar position if explicitly configured
+        has_solar_config = (
+            hasattr(config.input_data, "solar_position") and
+            config.input_data.solar_position is not None
+        )
+
+        if has_solar_config:
+            solar_config = config.input_data.solar_position
+
+            # Create datetime range for solar position calculation
+            datetimes = pd.date_range(
+                t0 + minutes(solar_config.interval_start_minutes),
+                t0 + minutes(solar_config.interval_end_minutes),
+                freq=minutes(solar_config.time_resolution_minutes),
+            )
+
+            # Convert OSGB coordinates to lon/lat
+            lon, lat = osgb_to_lon_lat(location.x, location.y)
+
+            # Calculate solar positions and add to modalities
+            numpy_modalities.append(make_sun_position_numpy_sample(datetimes, lon, lat))
+
+        # Combine all the modalities and fill NaNs
+        combined_sample = merge_dicts(numpy_modalities)
+        combined_sample = fill_nans_in_arrays(combined_sample)
+
+        return combined_sample
+
+    @staticmethod
+    def find_valid_t0_times(datasets_dict: dict, config: Configuration) -> pd.DatetimeIndex:
+        """Find the t0 times where all of the requested input data is available.
+
+        Args:
+            datasets_dict: A dictionary of input datasets
+            config: Configuration 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.gsp.time_resolution_minutes),
+        )
+        return valid_t0_times
+
+
+
+class PVNetUKRegionalDataset(AbstractPVNetUKDataset):
+    """A torch Dataset for creating PVNet UK regional samples."""
+
+    @override
+    def __init__(
+        self,
+        config_filename: str,
+        start_time: str | None = None,
+        end_time: str | None = None,
+        gsp_ids: list[int] | None = None,
+    ) -> None:
+
+        super().__init__(config_filename, start_time, end_time, gsp_ids)
 
         # Construct a lookup for locations - useful for users to construct sample by GSP ID
-        location_lookup = {loc.id: loc for loc in locations}
+        location_lookup = {loc.id: loc for loc in self.locations}
 
         # Construct indices for sampling
         t_index, loc_index = np.meshgrid(
-            np.arange(len(valid_t0_times)),
-            np.arange(len(locations)),
+            np.arange(len(self.valid_t0_times)),
+            np.arange(len(self.locations)),
         )
 
         # Make array of all possible (t0, location) coordinates. Each row is a single coordinate
         index_pairs = np.stack((t_index.ravel(), loc_index.ravel())).T
 
         # Assign coords and indices to self
-        self.valid_t0_times = valid_t0_times
-        self.locations = locations
         self.location_lookup = location_lookup
         self.index_pairs = index_pairs
 
-        # Assign config and input data to self
-        self.datasets_dict = datasets_dict
-        self.config = config
-
     @override
     def __len__(self) -> int:
         return len(self.index_pairs)
 
-    def _get_sample(self, t0: pd.Timestamp, location: Location) -> dict:
+    def _get_sample(self, t0: pd.Timestamp, location: Location) -> NumpySample:
         """Generate the PVNet sample for given coordinates.
 
         Args:
@@ -254,21 +279,18 @@ class PVNetUKRegionalDataset(Dataset):
         sample_dict = slice_datasets_by_time(sample_dict, t0, self.config)
         sample_dict = compute(sample_dict)
 
-        sample = process_and_combine_datasets(sample_dict, self.config, t0, location)
-
-        return sample
+        return self.process_and_combine_datasets(sample_dict, self.config, t0, location)
 
     @override
-    def __getitem__(self, idx: int) -> dict:
+    def __getitem__(self, idx: int) -> NumpySample:
         # Get the coordinates of the sample
         t_index, loc_index = self.index_pairs[idx]
         location = self.locations[loc_index]
         t0 = self.valid_t0_times[t_index]
 
-        # Generate the sample
         return self._get_sample(t0, location)
 
-    def get_sample(self, t0: pd.Timestamp, gsp_id: int) -> dict:
+    def get_sample(self, t0: pd.Timestamp, gsp_id: int) -> NumpySample:
         """Generate a sample for the given coordinates.
 
         Useful for users to generate specific samples.
@@ -288,56 +310,14 @@ class PVNetUKRegionalDataset(Dataset):
         return self._get_sample(t0, location)
 
 
-class PVNetUKConcurrentDataset(Dataset):
+class PVNetUKConcurrentDataset(AbstractPVNetUKDataset):
     """A torch Dataset for creating concurrent PVNet UK regional samples."""
 
-    def __init__(
-        self,
-        config_filename: str,
-        start_time: str | None = None,
-        end_time: str | None = None,
-        gsp_ids: list[int] | None = None,
-    ) -> None:
-        """A torch Dataset for creating concurrent samples of PVNet UK regional data.
-
-        Each concurrent sample includes the data from all GSPs for a single t0 time
-
-        Args:
-            config_filename: Path to the configuration file
-            start_time: Limit the init-times to be after this
-            end_time: Limit the init-times to be before this
-            gsp_ids: List of all GSP IDs included in each sample. Defaults to all
-        """
-        config = load_yaml_configuration(config_filename)
-
-        datasets_dict = get_dataset_dict(config.input_data)
-
-        # Get t0 times where all input data is available
-        valid_t0_times = find_valid_t0_times(datasets_dict, config)
-
-        # Filter t0 times to given range
-        if start_time is not None:
-            valid_t0_times = valid_t0_times[valid_t0_times >= pd.Timestamp(start_time)]
-
-        if end_time is not None:
-            valid_t0_times = valid_t0_times[valid_t0_times <= pd.Timestamp(end_time)]
-
-        # Construct list of locations to sample from
-        locations = get_gsp_locations(gsp_ids)
-
-        # Assign coords and indices to self
-        self.valid_t0_times = valid_t0_times
-        self.locations = locations
-
-        # Assign config and input data to self
-        self.datasets_dict = datasets_dict
-        self.config = config
-
     @override
     def __len__(self) -> int:
         return len(self.valid_t0_times)
 
-    def _get_sample(self, t0: pd.Timestamp) -> dict:
+    def _get_sample(self, t0: pd.Timestamp) -> NumpyBatch:
         """Generate a concurrent PVNet sample for given init-time.
 
         Args:
@@ -352,7 +332,7 @@ class PVNetUKConcurrentDataset(Dataset):
         # Prepare sample for each GSP
         for location in self.locations:
             gsp_sample_dict = slice_datasets_by_space(sample_dict, location, self.config)
-            gsp_numpy_sample = process_and_combine_datasets(
+            gsp_numpy_sample = self.process_and_combine_datasets(
                 gsp_sample_dict,
                 self.config,
                 t0,
@@ -364,10 +344,10 @@ class PVNetUKConcurrentDataset(Dataset):
         return stack_np_samples_into_batch(gsp_samples)
 
     @override
-    def __getitem__(self, idx: int) -> dict:
+    def __getitem__(self, idx: int) -> NumpyBatch:
         return self._get_sample(self.valid_t0_times[idx])
 
-    def get_sample(self, t0: pd.Timestamp) -> dict:
+    def get_sample(self, t0: pd.Timestamp) -> NumpyBatch:
         """Generate a sample for the given init-time.
 
         Useful for users to generate specific samples.

ocf_data_sampler/torch_datasets/datasets/site.py

@@ -1,14 +1,12 @@
 """Torch dataset for sites."""
 
-import logging
-
 import numpy as np
 import pandas as pd
 import xarray as xr
 from torch.utils.data import Dataset
 from typing_extensions import override
 
-from ocf_data_sampler.config import Configuration, load_yaml_configuration
+from ocf_data_sampler.config import load_yaml_configuration
 from ocf_data_sampler.load.load_dataset import get_dataset_dict
 from ocf_data_sampler.numpy_sample import (
     NWPSampleKey,
@@ -18,15 +16,19 @@ from ocf_data_sampler.numpy_sample import (
     make_datetime_numpy_dict,
     make_sun_position_numpy_sample,
 )
+from ocf_data_sampler.numpy_sample.common_types import NumpySample
 from ocf_data_sampler.select import (
     Location,
     fill_time_periods,
     find_contiguous_t0_periods,
     intersection_of_multiple_dataframes_of_periods,
+)
+from ocf_data_sampler.torch_datasets.utils import (
+    channel_dict_to_dataarray,
+    find_valid_time_periods,
     slice_datasets_by_space,
     slice_datasets_by_time,
 )
-from ocf_data_sampler.torch_datasets.utils import channel_dict_to_dataarray, find_valid_time_periods
 from ocf_data_sampler.torch_datasets.utils.merge_and_fill_utils import (
     fill_nans_in_arrays,
     merge_dicts,
@@ -52,7 +54,7 @@ class SitesDataset(Dataset):
             start_time: Limit the init-times to be after this
             end_time: Limit the init-times to be before this
         """
-        config: Configuration = load_yaml_configuration(config_filename)
+        config = load_yaml_configuration(config_filename)
         datasets_dict = get_dataset_dict(config.input_data)
 
         # Assign config and input data to self
@@ -61,6 +63,7 @@ class SitesDataset(Dataset):
 
         # get all locations
         self.locations = self.get_locations(datasets_dict["site"])
+        self.location_lookup = {loc.id: loc for loc in self.locations}
 
         # Get t0 times where all input data is available
         valid_t0_and_site_ids = self.find_valid_t0_and_site_ids(datasets_dict)
@@ -89,7 +92,7 @@ class SitesDataset(Dataset):
         t0, site_id = self.valid_t0_and_site_ids.iloc[idx]
 
         # get location from site id
-        location = self.get_location_from_site_id(site_id)
+        location = self.location_lookup[site_id]
 
         # Generate the sample
         return self._get_sample(t0, location)
@@ -105,8 +108,7 @@ class SitesDataset(Dataset):
         sample_dict = slice_datasets_by_time(sample_dict, t0, self.config)
 
         sample = self.process_and_combine_site_sample_dict(sample_dict, t0)
-        sample = sample.compute()
-        return sample
+        return sample.compute()
 
     def get_sample(self, t0: pd.Timestamp, site_id: int) -> dict:
         """Generate a sample for a given site id and t0.
@@ -117,22 +119,10 @@ class SitesDataset(Dataset):
             t0: init-time for sample
             site_id: site id as int
         """
-        location = self.get_location_from_site_id(site_id)
+        location = self.location_lookup[site_id]
 
         return self._get_sample(t0, location)
 
-    def get_location_from_site_id(self, site_id: int) -> Location:
-        """Get location from system id."""
-        locations = [loc for loc in self.locations if loc.id == site_id]
-        if len(locations) == 0:
-            raise ValueError(f"Location not found for site_id {site_id}")
-
-        if len(locations) > 1:
-            logging.warning(
-                f"Multiple locations found for site_id {site_id}, but will take the first",
-            )
-
-        return locations[0]
 
     def find_valid_t0_and_site_ids(
         self,
@@ -148,24 +138,21 @@ class SitesDataset(Dataset):
             datasets_dict: A dictionary of input datasets
             config: Configuration file
         """
-        # 1. Get valid time period for nwp and satellite
+        # Get valid time period for nwp and satellite
         datasets_without_site = {k: v for k, v in datasets_dict.items() if k != "site"}
         valid_time_periods = find_valid_time_periods(datasets_without_site, self.config)
 
-        # 2. Now lets loop over each location in system id and find the valid periods
-        # Should we have a different option if there are not nans
+        # Loop over each location in system id and obtain valid periods
         sites = datasets_dict["site"]
         site_ids = sites.site_id.values
         site_config = self.config.input_data.site
         valid_t0_and_site_ids = []
         for site_id in site_ids:
             site = sites.sel(site_id=site_id)
-
-            # drop any nan values
-            # not sure this is right?
+            # Drop NaN values
             site = site.dropna(dim="time_utc")
 
-            # Get the valid time periods for this location
+            # Obtain valid time periods for this location
             time_periods = find_contiguous_t0_periods(
                 pd.DatetimeIndex(site["time_utc"]),
                 time_resolution=minutes(site_config.time_resolution_minutes),
@@ -176,7 +163,7 @@ class SitesDataset(Dataset):
                 [valid_time_periods, time_periods],
             )
 
-            # Fill out the contiguous time periods to get the t0 times
+            # Fill out contiguous time periods to get t0 times
             valid_t0_times_per_site = fill_time_periods(
                 valid_time_periods_per_site,
                 freq=minutes(site_config.time_resolution_minutes),
@@ -188,12 +175,15 @@ class SitesDataset(Dataset):
 
         valid_t0_and_site_ids = pd.concat(valid_t0_and_site_ids)
         valid_t0_and_site_ids.index.name = "t0"
-        valid_t0_and_site_ids.reset_index(inplace=True)
+        return valid_t0_and_site_ids.reset_index()
 
-        return valid_t0_and_site_ids
 
     def get_locations(self, site_xr: xr.Dataset) -> list[Location]:
-        """Get list of locations of all sites."""
+        """Get list of locations of all sites.
+
+        Args:
+            site_xr: xarray Dataset of site data
+        """
         locations = []
         for site_id in site_xr.site_id.values:
             site = site_xr.sel(site_id=site_id)
@@ -220,7 +210,6 @@ class SitesDataset(Dataset):
 
         Returns:
             xr.Dataset: A merged Dataset with nans filled in.
-
         """
         data_arrays = []
 
@@ -228,7 +217,6 @@ class SitesDataset(Dataset):
             for nwp_key, da_nwp in dataset_dict["nwp"].items():
                 provider = self.config.input_data.nwp[nwp_key].provider
 
-                # Standardise
                 da_channel_means = channel_dict_to_dataarray(
                     self.config.input_data.nwp[nwp_key].channel_means,
                 )
@@ -237,7 +225,6 @@ class SitesDataset(Dataset):
                 )
 
                 da_nwp = (da_nwp - da_channel_means) / da_channel_stds
-
                 data_arrays.append((f"nwp-{provider}", da_nwp))
 
         if "sat" in dataset_dict:
@@ -251,11 +238,9 @@ class SitesDataset(Dataset):
             )
 
             da_sat = (da_sat - da_channel_means) / da_channel_stds
-
             data_arrays.append(("satellite", da_sat))
 
         if "site" in dataset_dict:
-            # site_config = config.input_data.site
             da_sites = dataset_dict["site"]
             da_sites = da_sites / da_sites.capacity_kwp
             data_arrays.append(("site", da_sites))
@@ -372,12 +357,16 @@ class SitesDataset(Dataset):
 
 
 def convert_netcdf_to_numpy_sample(ds: xr.Dataset) -> dict:
-    """Convert a netcdf dataset to a numpy sample."""
+    """Convert a netcdf dataset to a numpy sample.
+
+    Args:
+        ds: xarray Dataset
+    """
     # convert the single dataset to a dict of arrays
     sample_dict = convert_from_dataset_to_dict_datasets(ds)
 
     if "satellite" in sample_dict:
-        # rename satellite to satellite actual # TODO this could be improves
+        # rename satellite to sat # TODO this could be improved
         sample_dict["sat"] = sample_dict.pop("satellite")
 
     # process and combine the datasets
@@ -408,43 +397,52 @@ def convert_from_dataset_to_dict_datasets(combined_dataset: xr.Dataset) -> dict[
         The uncombined datasets as a dict of xr.Datasets
     """
     # Split into datasets by splitting by the prefix added in combine_to_netcdf
-    datasets = {}
+    datasets: dict[str, xr.DataArray] = {}
+
     # Go through each data variable and split it into a dataset
     for key, dataset in combined_dataset.items():
-        # If 'key_' doesn't exist in a dim or coordinate, remove it
-        dataset_dims = list(dataset.coords)
-        for dim in dataset_dims:
+        # If 'key__' doesn't exist in a dim or coordinate, remove it
+        for dim in list(dataset.coords):
             if f"{key}__" not in dim:
-                dataset: xr.Dataset = dataset.drop(dim)
+                dataset = dataset.drop_vars(dim)
         dataset = dataset.rename(
             {dim: dim.split(f"{key}__")[1] for dim in dataset.dims if dim not in dataset.coords},
         )
-        dataset: xr.Dataset = dataset.rename(
+        dataset = dataset.rename(
             {coord: coord.split(f"{key}__")[1] for coord in dataset.coords},
         )
         # Split the dataset by the prefix
         datasets[key] = dataset
 
     # Unflatten any NWP data
-    datasets = nest_nwp_source_dict(datasets, sep="-")
-    return datasets
+    return nest_nwp_source_dict(datasets, sep="-")
+
 
+def nest_nwp_source_dict(
+    dataset_dict: dict[xr.Dataset],
+    sep: str = "-",
+) -> dict[str, xr.Dataset | dict[xr.Dataset]]:
+    """Re-nest a dictionary where the NWP values are nested under keys 'nwp-<key>'.
 
-def nest_nwp_source_dict(d: dict, sep: str = "/") -> dict:
-    """Re-nest a dictionary where the NWP values are nested under keys 'nwp/<key>'."""
+    Args:
+        dataset_dict: Dictionary of datasets
+        sep: Separator to use to nest NWP keys
+    """
     nwp_prefix = f"nwp{sep}"
-    new_dict = {k: v for k, v in d.items() if not k.startswith(nwp_prefix)}
-    nwp_keys = [k for k in d if k.startswith(nwp_prefix)]
+    new_dict = {k: v for k, v in dataset_dict.items() if not k.startswith(nwp_prefix)}
+    nwp_keys = [k for k in dataset_dict if k.startswith(nwp_prefix)]
     if len(nwp_keys) > 0:
-        nwp_subdict = {k.removeprefix(nwp_prefix): d[k] for k in nwp_keys}
+        nwp_subdict = {k.removeprefix(nwp_prefix): dataset_dict[k] for k in nwp_keys}
         new_dict["nwp"] = nwp_subdict
     return new_dict
 
 
-def convert_to_numpy_and_combine(
-    dataset_dict: dict,
-) -> dict:
-    """Convert input data in a dict to numpy arrays."""
+def convert_to_numpy_and_combine(dataset_dict: dict[xr.Dataset]) -> NumpySample:
+    """Convert input data in a dict to numpy arrays.
+
+    Args:
+        dataset_dict: Dictionary of xarray Datasets
+    """
     numpy_modalities = []
 
     if "nwp" in dataset_dict:
@@ -474,9 +472,7 @@ def convert_to_numpy_and_combine(
 
     # Combine all the modalities and fill NaNs
     combined_sample = merge_dicts(numpy_modalities)
-    combined_sample = fill_nans_in_arrays(combined_sample)
-
-    return combined_sample
+    return fill_nans_in_arrays(combined_sample)
 
 
 def coarsen_data(xr_data: xr.Dataset, coarsen_to_deg: float = 0.1) -> xr.Dataset:

ocf_data_sampler/torch_datasets/sample/__init__.py

@@ -0,0 +1,3 @@
+from .base import SampleBase
+from .uk_regional import UKRegionalSample
+from .site import SiteSample

ocf_data_sampler/{sample → torch_datasets/sample}/site.py

@@ -4,9 +4,10 @@ import xarray as xr
 from typing_extensions import override
 
 from ocf_data_sampler.numpy_sample.common_types import NumpySample
-from ocf_data_sampler.sample.base import SampleBase
 from ocf_data_sampler.torch_datasets.datasets.site import convert_netcdf_to_numpy_sample
 
+from .base import SampleBase
+
 
 class SiteSample(SampleBase):
     """Handles PVNet site specific netCDF operations."""

ocf_data_sampler/{sample → torch_datasets/sample}/uk_regional.py

@@ -9,7 +9,8 @@ from ocf_data_sampler.numpy_sample import (
     SatelliteSampleKey,
 )
 from ocf_data_sampler.numpy_sample.common_types import NumpySample
-from ocf_data_sampler.sample.base import SampleBase
+
+from .base import SampleBase
 
 
 class UKRegionalSample(SampleBase):

ocf_data_sampler/torch_datasets/utils/__init__.py

@@ -1,3 +1,5 @@
 from .channel_dict_to_dataarray import channel_dict_to_dataarray
 from .merge_and_fill_utils import fill_nans_in_arrays, merge_dicts
-from .valid_time_periods import find_valid_time_periods
+from .valid_time_periods import find_valid_time_periods
+from .spatial_slice_for_dataset import slice_datasets_by_space
+from .time_slice_for_dataset import slice_datasets_by_time

ocf_data_sampler/torch_datasets/utils/channel_dict_to_dataarray.py

@@ -1,10 +1,17 @@
-"""Converts a dictionary of channel values to a DataArray."""
+"""Utility function for converting channel dictionaries to xarray DataArrays."""
 
 import xarray as xr
 
 
 def channel_dict_to_dataarray(channel_dict: dict[str, float]) -> xr.DataArray:
-    """Converts a dictionary of channel values to a DataArray."""
+    """Converts a dictionary of channel values to a DataArray.
+
+    Args:
+        channel_dict: Dictionary mapping channel names (str) to their values (float).
+
+    Returns:
+        xr.DataArray: A 1D DataArray with channels as coordinates.
+    """
     return xr.DataArray(
         list(channel_dict.values()),
         coords={"channel": list(channel_dict.keys())},

{ocf_data_sampler-0.2.7.dist-info → ocf_data_sampler-0.2.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: ocf-data-sampler
-Version: 0.2.7
+Version: 0.2.9
 Author: James Fulton, Peter Dudfield
 Author-email: Open Climate Fix team <info@openclimatefix.org>
 License: MIT License
@@ -35,6 +35,7 @@ Requires-Dist: numpy
 Requires-Dist: pandas
 Requires-Dist: xarray
 Requires-Dist: zarr==2.18.3
+Requires-Dist: numcodecs<0.16
 Requires-Dist: dask
 Requires-Dist: matplotlib
 Requires-Dist: ocf_blosc2

{ocf_data_sampler-0.2.7.dist-info → ocf_data_sampler-0.2.9.dist-info}/RECORD

@@ -29,30 +29,30 @@ ocf_data_sampler/numpy_sample/nwp.py,sha256=X9T5XZLVucXX8QAUhdeTnomNBPrsfvsO8I4S
 ocf_data_sampler/numpy_sample/satellite.py,sha256=RaYzYIcB1AmDrKeiqSpn4QVfBH-QMe26F1P5t1az2Jg,1111
 ocf_data_sampler/numpy_sample/site.py,sha256=zfYBjK3CJrIaKH1QdKXU7gwOxTqONt527y3nJ9TRnwc,1325
 ocf_data_sampler/numpy_sample/sun_position.py,sha256=5tt-zNm6aRuZMsxZPaAxyg7HeikswfZCeHWXTHuO2K0,1555
-ocf_data_sampler/sample/__init__.py,sha256=zdS73NTnxFX_j8uh9tT-IXiURB6635wbneM1koWYV1o,169
-ocf_data_sampler/sample/base.py,sha256=cQ1oIyhdmlotejZK8B3Cw6MNvpdnBPD8G_o2h7Ye4Vc,2206
-ocf_data_sampler/sample/site.py,sha256=BhQPygeLUncXJGN_Yd2CL050kN6ktZlobaJw0O0RagI,1290
-ocf_data_sampler/sample/uk_regional.py,sha256=VOby07RnZYvzszExwqoZRVwZ1EbCclRpXq1e9CL16CE,2463
-ocf_data_sampler/select/__init__.py,sha256=E4AJulEbO2K-o0UlG1fgaEteuf_1ZFjHTvrotXSb4YU,332
-ocf_data_sampler/select/dropout.py,sha256=_rzXl8_4VHTY_JMjbaoWopaFCJmLdaBpqfYF4vr24tk,1638
+ocf_data_sampler/select/__init__.py,sha256=mK7Wu_-j9IXGTYrOuDf5yDDuU5a306b0iGKTAooNg_s,210
+ocf_data_sampler/select/dropout.py,sha256=WVOCweTGfIjufAlnfmYiPofz6X38TxQgzkLwtiB3TrU,1712
 ocf_data_sampler/select/fill_time_periods.py,sha256=TlGxp1xiAqnhdWfLy0pv3FuZc00dtimjWdLzr4JoTGA,865
 ocf_data_sampler/select/find_contiguous_time_periods.py,sha256=cEXrQDzk8pXknxB0q3v5DakosagHMoLDAj302B8Xpw0,11537
 ocf_data_sampler/select/geospatial.py,sha256=CDExkl36eZOKmdJPzUr_K0Wn3axHqv5nYo-EkSiINcc,5032
 ocf_data_sampler/select/location.py,sha256=AZvGR8y62opiW7zACGXjoOtBEWRfSLOZIA73O5Deu0c,1037
 ocf_data_sampler/select/select_spatial_slice.py,sha256=qY2Ll00EPA80oBtzwMoR5nk0UIpoWZF9oXl22YwWr0Q,12341
 ocf_data_sampler/select/select_time_slice.py,sha256=HeHbwZ0CP03x0-LaJtpbSdtpLufwVTR73p6wH6O_PS8,5513
-ocf_data_sampler/select/spatial_slice_for_dataset.py,sha256=Hvz0wHSWMYYamf2oHNiGlzJcM4cAH6pL_7ZEvIBL2dE,1882
-ocf_data_sampler/select/time_slice_for_dataset.py,sha256=1DN6VsWWdLvkpJxodZtBRDUgC4vJE2td_RP5J3ZqPNw,4268
 ocf_data_sampler/torch_datasets/datasets/__init__.py,sha256=jfJSFcR0eO1AqeH7S3KnGjsBqVZT5w3oyi784PUR6Q0,146
-ocf_data_sampler/torch_datasets/datasets/pvnet_uk.py,sha256=ZkXm0IQEIzZUi8O-qJJz2PbJr9T4ZvutL424yRQUJhc,12878
-ocf_data_sampler/torch_datasets/datasets/site.py,sha256=j29cWPIcksRbge014MxR0_OgJqoskdki6KqvtoHtxpY,18023
-ocf_data_sampler/torch_datasets/utils/__init__.py,sha256=7Yt4anQVU9y27nj4Wx1tRLqbAQLbzW0ED71UL65LvxA,187
-ocf_data_sampler/torch_datasets/utils/channel_dict_to_dataarray.py,sha256=MGylKhXxXLQC2fYv-8L_GVoYhov3LcEwC0Q21xItDSk,353
+ocf_data_sampler/torch_datasets/datasets/pvnet_uk.py,sha256=tx5Sg64eknhU6VIcONiAaG2PurN6Y8Te6rE3AaWg8t4,12338
+ocf_data_sampler/torch_datasets/datasets/site.py,sha256=nRUlhXQQGVrTuBmE1QnwXAUsPTXz0dsezlQjwK71jIQ,17641
+ocf_data_sampler/torch_datasets/sample/__init__.py,sha256=GL84vdZl_SjHDGVyh9Uekx2XhPYuZ0dnO3l6f6KXnHI,100
+ocf_data_sampler/torch_datasets/sample/base.py,sha256=cQ1oIyhdmlotejZK8B3Cw6MNvpdnBPD8G_o2h7Ye4Vc,2206
+ocf_data_sampler/torch_datasets/sample/site.py,sha256=ZUEgn50g-GmqujOEtezNILF7wjokF80sDAA4OOldcRI,1268
+ocf_data_sampler/torch_datasets/sample/uk_regional.py,sha256=zpCeUw3eljOnoJTSUYW2R4kiWrY6hbuXjK8igJrXgPg,2441
+ocf_data_sampler/torch_datasets/utils/__init__.py,sha256=N7i_hHtWUDiJqsiJoDx4T_QuiYOuvIyulPrn6xEA4TY,309
+ocf_data_sampler/torch_datasets/utils/channel_dict_to_dataarray.py,sha256=un2IiyoAmTDIymdeMiPU899_86iCDMD-oIifjHlNyqw,555
 ocf_data_sampler/torch_datasets/utils/merge_and_fill_utils.py,sha256=we7BTxRH7B7jKayDT7YfNyfI3zZClz2Bk-HXKQIokgU,956
+ocf_data_sampler/torch_datasets/utils/spatial_slice_for_dataset.py,sha256=Hvz0wHSWMYYamf2oHNiGlzJcM4cAH6pL_7ZEvIBL2dE,1882
+ocf_data_sampler/torch_datasets/utils/time_slice_for_dataset.py,sha256=1DN6VsWWdLvkpJxodZtBRDUgC4vJE2td_RP5J3ZqPNw,4268
 ocf_data_sampler/torch_datasets/utils/valid_time_periods.py,sha256=xcy75cVxl0WrglnX5YUAFjXXlO2GwEBHWyqo8TDuiOA,4714
 scripts/refactor_site.py,sha256=skzvsPP0Cn9yTKndzkilyNcGz4DZ88ctvCJ0XrBdc2A,3135
 utils/compute_icon_mean_stddev.py,sha256=a1oWMRMnny39rV-dvu8rcx85sb4bXzPFrR1gkUr4Jpg,2296
-ocf_data_sampler-0.2.7.dist-info/METADATA,sha256=XjzoJNcebQ_9UTsC--_HX9gG7mrBQBKmih3wDGOWtQQ,11594
-ocf_data_sampler-0.2.7.dist-info/WHEEL,sha256=52BFRY2Up02UkjOa29eZOS2VxUrpPORXg1pkohGGUS8,91
-ocf_data_sampler-0.2.7.dist-info/top_level.txt,sha256=LEFU4Uk-PEo72QGLAfnVZIUEm37Q8mKuMeg_Xk-p33g,31
-ocf_data_sampler-0.2.7.dist-info/RECORD,,
+ocf_data_sampler-0.2.9.dist-info/METADATA,sha256=IOkMU53l3Mgeh9dCIyXvcFMZtsJTPj7vhYT_Sb7hWZU,11624
+ocf_data_sampler-0.2.9.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+ocf_data_sampler-0.2.9.dist-info/top_level.txt,sha256=LEFU4Uk-PEo72QGLAfnVZIUEm37Q8mKuMeg_Xk-p33g,31
+ocf_data_sampler-0.2.9.dist-info/RECORD,,

{ocf_data_sampler-0.2.7.dist-info → ocf_data_sampler-0.2.9.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (76.0.0)
+Generator: setuptools (78.1.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

ocf_data_sampler/sample/__init__.py

@@ -1,3 +0,0 @@
-from ocf_data_sampler.sample.base import SampleBase
-from ocf_data_sampler.sample.uk_regional import UKRegionalSample
-from ocf_data_sampler.sample.site import SiteSample