ocf-data-sampler 0.1.10__py3-none-any.whl → 0.1.16__py3-none-any.whl

ocf_data_sampler/select/time_slice_for_dataset.py

@@ -1,25 +1,26 @@
-""" Slice datasets by time"""
+"""Slice datasets by time."""
+
 import pandas as pd
 import xarray as xr
 
 from ocf_data_sampler.config import Configuration
-from ocf_data_sampler.select.dropout import draw_dropout_time, apply_dropout_time
-from ocf_data_sampler.select.select_time_slice import select_time_slice_nwp, select_time_slice
+from ocf_data_sampler.select.dropout import apply_dropout_time, draw_dropout_time
+from ocf_data_sampler.select.select_time_slice import select_time_slice, select_time_slice_nwp
 from ocf_data_sampler.utils import minutes
 
+
 def slice_datasets_by_time(
     datasets_dict: dict,
     t0: pd.Timestamp,
     config: Configuration,
 ) -> dict:
-    """Slice the dictionary of input data sources around a given t0 time
+    """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.
     """
-
     sliced_datasets_dict = {}
 
     if "nwp" in datasets_dict:
@@ -31,7 +32,7 @@ def slice_datasets_by_time(
             sliced_datasets_dict["nwp"][nwp_key] = select_time_slice_nwp(
                 da_nwp,
                 t0,
-                sample_period_duration=minutes(nwp_config.time_resolution_minutes),
+                time_resolution=minutes(nwp_config.time_resolution_minutes),
                 interval_start=minutes(nwp_config.interval_start_minutes),
                 interval_end=minutes(nwp_config.interval_end_minutes),
                 dropout_timedeltas=minutes(nwp_config.dropout_timedeltas_minutes),
@@ -45,7 +46,7 @@ def slice_datasets_by_time(
         sliced_datasets_dict["sat"] = select_time_slice(
             datasets_dict["sat"],
             t0,
-            sample_period_duration=minutes(sat_config.time_resolution_minutes),
+            time_resolution=minutes(sat_config.time_resolution_minutes),
             interval_start=minutes(sat_config.interval_start_minutes),
             interval_end=minutes(sat_config.interval_end_minutes),
         )
@@ -65,11 +66,11 @@ def slice_datasets_by_time(
 
     if "gsp" in datasets_dict:
         gsp_config = config.input_data.gsp
-        
+
         da_gsp_past = select_time_slice(
             datasets_dict["gsp"],
             t0,
-            sample_period_duration=minutes(gsp_config.time_resolution_minutes),
+            time_resolution=minutes(gsp_config.time_resolution_minutes),
             interval_start=minutes(gsp_config.interval_start_minutes),
             interval_end=minutes(0),
         )
@@ -82,18 +83,18 @@ def slice_datasets_by_time(
         )
 
         da_gsp_past = apply_dropout_time(
-            da_gsp_past, 
-            gsp_dropout_time
+            da_gsp_past,
+            gsp_dropout_time,
         )
-        
+
         da_gsp_future = select_time_slice(
             datasets_dict["gsp"],
             t0,
-            sample_period_duration=minutes(gsp_config.time_resolution_minutes),
+            time_resolution=minutes(gsp_config.time_resolution_minutes),
             interval_start=minutes(gsp_config.time_resolution_minutes),
             interval_end=minutes(gsp_config.interval_end_minutes),
         )
-        
+
         sliced_datasets_dict["gsp"] = xr.concat([da_gsp_past, da_gsp_future], dim="time_utc")
 
     if "site" in datasets_dict:
@@ -102,7 +103,7 @@ def slice_datasets_by_time(
         sliced_datasets_dict["site"] = select_time_slice(
             datasets_dict["site"],
             t0,
-            sample_period_duration=minutes(site_config.time_resolution_minutes),
+            time_resolution=minutes(site_config.time_resolution_minutes),
             interval_start=minutes(site_config.interval_start_minutes),
             interval_end=minutes(site_config.interval_end_minutes),
         )
@@ -120,4 +121,4 @@ def slice_datasets_by_time(
             site_dropout_time,
         )
 
-    return sliced_datasets_dict
+    return sliced_datasets_dict

ocf_data_sampler/torch_datasets/datasets/pvnet_uk.py

@@ -1,41 +1,42 @@
-"""Torch dataset for UK PVNet"""
+"""Torch dataset for UK PVNet."""
 
-import pkg_resources
+from importlib.resources import files
 
 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.load.load_dataset import get_dataset_dict
-from ocf_data_sampler.select import (
-    fill_time_periods, 
-    Location, 
-    slice_datasets_by_space, 
-    slice_datasets_by_time,
-)
-from ocf_data_sampler.utils import minutes
 from ocf_data_sampler.constants import NWP_MEANS, NWP_STDS, RSS_MEAN, RSS_STD
+from ocf_data_sampler.load.load_dataset import get_dataset_dict
 from ocf_data_sampler.numpy_sample import (
+    convert_gsp_to_numpy_sample,
     convert_nwp_to_numpy_sample,
     convert_satellite_to_numpy_sample,
-    convert_gsp_to_numpy_sample,
     make_sun_position_numpy_sample,
 )
+from ocf_data_sampler.numpy_sample.collate import stack_np_samples_into_batch
 from ocf_data_sampler.numpy_sample.gsp import GSPSampleKey
 from ocf_data_sampler.numpy_sample.nwp import NWPSampleKey
-from ocf_data_sampler.numpy_sample.collate import stack_np_samples_into_batch
+from ocf_data_sampler.select import (
+    Location,
+    fill_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.valid_time_periods import find_valid_time_periods
 from ocf_data_sampler.torch_datasets.utils.merge_and_fill_utils import (
-    merge_dicts,
     fill_nans_in_arrays,
+    merge_dicts,
 )
+from ocf_data_sampler.torch_datasets.utils.valid_time_periods import find_valid_time_periods
 from ocf_data_sampler.torch_datasets.utils.validate_channels import (
     validate_nwp_channels,
     validate_satellite_channels,
 )
-
+from ocf_data_sampler.utils import minutes
 
 xr.set_options(keep_attrs=True)
 
@@ -45,14 +46,12 @@ def process_and_combine_datasets(
     config: Configuration,
     t0: pd.Timestamp,
     location: Location,
-    target_key: str = 'gsp'
 ) -> dict:
-
-    """Normalise and convert data to numpy arrays"""
+    """Normalise and convert data to numpy arrays."""
     numpy_modalities = []
 
     if "nwp" in dataset_dict:
-        nwp_numpy_modalities = dict()
+        nwp_numpy_modalities = {}
 
         for nwp_key, da_nwp in dataset_dict["nwp"].items():
             provider = config.input_data.nwp[nwp_key].provider
@@ -71,41 +70,50 @@ def process_and_combine_datasets(
         da_sat = (da_sat - RSS_MEAN) / RSS_STD
         numpy_modalities.append(convert_satellite_to_numpy_sample(da_sat))
 
-    gsp_config = config.input_data.gsp
-
     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
-            )
+                da_gsp,
+                t0_idx=-gsp_config.interval_start_minutes / gsp_config.time_resolution_minutes,
+            ),
         )
 
-    if target_key == 'gsp':
-        # Make sun coords NumpySample
+    # 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(gsp_config.interval_start_minutes),
-            t0+minutes(gsp_config.interval_end_minutes),
-            freq=minutes(gsp_config.time_resolution_minutes),
+            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)
 
-        numpy_modalities.append(
-            {
-                GSPSampleKey.gsp_id: location.id,
-                GSPSampleKey.x_osgb: location.x,
-                GSPSampleKey.y_osgb: location.y,
-            }
-        )
-
-    numpy_modalities.append(
-        make_sun_position_numpy_sample(datetimes, lon, lat, key_prefix=target_key)
-    )
+        # 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)
@@ -115,7 +123,7 @@ def process_and_combine_datasets(
 
 
 def compute(xarray_dict: dict) -> dict:
-    """Eagerly load a nested dictionary of xarray DataArrays"""
+    """Eagerly load a nested dictionary of xarray DataArrays."""
     for k, v in xarray_dict.items():
         if isinstance(v, dict):
             xarray_dict[k] = compute(v)
@@ -125,59 +133,58 @@ def compute(xarray_dict: dict) -> 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
+    """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)
+        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"""
-    
+    """Get list of locations of all GSPs."""
     if gsp_ids is None:
-        gsp_ids = [i for i in range(1, 318)]
-    
+        gsp_ids = list(range(1, 318))
+
     locations = []
 
     # Load UK GSP locations
     df_gsp_loc = pd.read_csv(
-        pkg_resources.resource_filename(__name__, "../../data/uk_gsp_locations.csv"),
+        files("ocf_data_sampler.data").joinpath("uk_gsp_locations.csv"),
         index_col="gsp_id",
     )
 
     for gsp_id in gsp_ids:
         locations.append(
             Location(
-                coordinate_system = "osgb",
+                coordinate_system="osgb",
                 x=df_gsp_loc.loc[gsp_id].x_osgb,
                 y=df_gsp_loc.loc[gsp_id].y_osgb,
                 id=gsp_id,
-            )
+            ),
         )
     return locations
 
 
 class PVNetUKRegionalDataset(Dataset):
+    """A torch Dataset for creating PVNet UK regional samples."""
+
     def __init__(
-        self, 
-        config_filename: str, 
+        self,
+        config_filename: str,
         start_time: str | None = None,
         end_time: str | None = None,
         gsp_ids: list[int] | None = None,
-    ):
-        """A torch Dataset for creating PVNet UK GSP samples
+    ) -> None:
+        """A torch Dataset for creating PVNet UK GSP samples.
 
         Args:
             config_filename: Path to the configuration file
@@ -185,31 +192,30 @@ 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)
         validate_nwp_channels(config)
         validate_satellite_channels(config)
 
         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)]
-            
+            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)]
+            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)
 
         # Construct a lookup for locations - useful for users to construct sample by GSP ID
         location_lookup = {loc.id: loc for loc in locations}
-        
-        # Construct indices for sampling
+
+        # Construct indices for sampling
         t_index, loc_index = np.meshgrid(
             np.arange(len(valid_t0_times)),
             np.arange(len(locations)),
@@ -217,7 +223,7 @@ class PVNetUKRegionalDataset(Dataset):
 
         # 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
@@ -227,15 +233,14 @@ class PVNetUKRegionalDataset(Dataset):
         # Assign config and input data to self
         self.datasets_dict = datasets_dict
         self.config = config
-        
-        
-    def __len__(self):
+
+    @override
+    def __len__(self) -> int:
         return len(self.index_pairs)
-    
-    
+
     def _get_sample(self, t0: pd.Timestamp, location: Location) -> dict:
-        """Generate the PVNet sample for given coordinates
-        
+        """Generate the PVNet sample for given coordinates.
+
         Args:
             t0: init-time for sample
             location: location for sample
@@ -245,49 +250,51 @@ class PVNetUKRegionalDataset(Dataset):
         sample_dict = compute(sample_dict)
 
         sample = process_and_combine_datasets(sample_dict, self.config, t0, location)
-        
+
         return sample
-    
-        
-    def __getitem__(self, idx):
-        
+
+    @override
+    def __getitem__(self, idx: int) -> dict:
         # 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:
-        """Generate a sample for the given coordinates. 
-        
+        """Generate a sample for the given coordinates.
+
         Useful for users to generate specific samples.
-        
+
         Args:
             t0: init-time for sample
             gsp_id: GSP ID
         """
         # Check the user has asked for a sample which we have the data for
-        assert t0 in self.valid_t0_times
-        assert gsp_id in self.location_lookup
+        if t0 not in self.valid_t0_times:
+            raise ValueError(f"Input init time '{t0!s}' not in valid times")
+        if gsp_id not in self.location_lookup:
+            raise ValueError(f"Input GSP '{gsp_id}' not known")
 
         location = self.location_lookup[gsp_id]
-        
+
         return self._get_sample(t0, location)
-    
-    
+
+
 class PVNetUKConcurrentDataset(Dataset):
+    """A torch Dataset for creating concurrent PVNet UK regional samples."""
+
     def __init__(
-        self, 
-        config_filename: str, 
+        self,
+        config_filename: str,
         start_time: str | None = None,
         end_time: str | None = None,
         gsp_ids: list[int] | None = None,
-    ):
-        """A torch Dataset for creating concurrent samples of PVNet UK regional data
-        
+    ) -> 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:
@@ -296,7 +303,6 @@ class PVNetUKConcurrentDataset(Dataset):
             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)
 
         # Validate channels for NWP and satellite data
@@ -304,20 +310,20 @@ class PVNetUKConcurrentDataset(Dataset):
         validate_satellite_channels(config)
 
         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)]
-            
+            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)]
+            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
@@ -325,48 +331,50 @@ class PVNetUKConcurrentDataset(Dataset):
         # Assign config and input data to self
         self.datasets_dict = datasets_dict
         self.config = config
-        
-        
-    def __len__(self):
+
+    @override
+    def __len__(self) -> int:
         return len(self.valid_t0_times)
-    
-    
+
     def _get_sample(self, t0: pd.Timestamp) -> dict:
-        """Generate a concurrent PVNet sample for given init-time
-        
+        """Generate a concurrent PVNet sample for given init-time.
+
         Args:
             t0: init-time for sample
         """
         # 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 = compute(sample_dict)
-        
+
         gsp_samples = []
-        
+
         # 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_sample_dict, self.config, t0, location
+                gsp_sample_dict,
+                self.config,
+                t0,
+                location,
             )
             gsp_samples.append(gsp_numpy_sample)
-        
+
         # Stack GSP samples
         return stack_np_samples_into_batch(gsp_samples)
-    
-        
-    def __getitem__(self, idx):
+
+    @override
+    def __getitem__(self, idx: int) -> dict:
         return self._get_sample(self.valid_t0_times[idx])
-    
 
     def get_sample(self, t0: pd.Timestamp) -> dict:
-        """Generate a sample for the given init-time. 
-        
+        """Generate a sample for the given init-time.
+
         Useful for users to generate specific samples.
-        
+
         Args:
             t0: init-time for sample
         """
         # Check data is availablle for init-time t0
-        assert t0 in self.valid_t0_times        
+        if t0 not in self.valid_t0_times:
+            raise ValueError(f"Input init time '{t0!s}' not in valid times")
         return self._get_sample(t0)