ocf-data-sampler 0.2.36__tar.gz → 0.2.38__tar.gz

{ocf_data_sampler-0.2.36 → ocf_data_sampler-0.2.38}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ocf-data-sampler
-Version: 0.2.36
+Version: 0.2.38
 Author: James Fulton, Peter Dudfield
 Author-email: Open Climate Fix team <info@openclimatefix.org>
 License: MIT License

{ocf_data_sampler-0.2.36 → ocf_data_sampler-0.2.38}/ocf_data_sampler/load/nwp/providers/gfs.py

@@ -24,8 +24,7 @@ def open_gfs(zarr_path: str | list[str], public: bool = False) -> xr.DataArray:
 
     # Open data
     gfs: xr.Dataset = open_zarr_paths(zarr_path, time_dim="init_time_utc", public=public)
-    nwp: xr.DataArray = gfs.to_array()
-    nwp = nwp.rename({"variable": "channel"})  # `variable` appears when using `to_array`
+    nwp: xr.DataArray = gfs.to_array(dim="channel")
 
     del gfs
 

{ocf_data_sampler-0.2.36 → ocf_data_sampler-0.2.38}/ocf_data_sampler/select/find_contiguous_time_periods.py

@@ -242,6 +242,11 @@ def intersection_of_2_dataframes_of_periods(a: pd.DataFrame, b: pd.DataFrame) ->
     if a.empty or b.empty:
         return pd.DataFrame(columns=["start_dt", "end_dt"])
 
+    # Maybe switch these for efficiency in the next section. We will do the native python loop over
+    # the shorter dataframe
+    if len(a) > len(b):
+        a, b = b, a
+
     all_intersecting_periods = []
     for a_period in a.itertuples():
         # Five ways in which two periods may overlap:
@@ -250,12 +255,12 @@ def intersection_of_2_dataframes_of_periods(a: pd.DataFrame, b: pd.DataFrame) ->
         # In all five, `a` must always start before (or equal to) where `b` ends,
         # and `a` must always end after (or equal to) where `b` starts.
 
-        overlapping_periods = b[(a_period.start_dt <= b.end_dt) & (a_period.end_dt >= b.start_dt)]
-
         # There are two ways in which two periods may *not* overlap:
         # a: |---|        or        |---|
         # b:       |---|      |---|
-        # `overlapping` will not include periods which do *not* overlap.
+        # `overlapping_periods` will not include periods which do *not* overlap.
+
+        overlapping_periods = b[(a_period.start_dt <= b.end_dt) & (a_period.end_dt >= b.start_dt)]
 
         # Now find the intersection of each period in `overlapping_periods` with
         # the period from `a` that starts at `a_start_dt` and ends at `a_end_dt`.
@@ -269,5 +274,6 @@ def intersection_of_2_dataframes_of_periods(a: pd.DataFrame, b: pd.DataFrame) ->
 
         all_intersecting_periods.append(intersection)
 
+
     all_intersecting_periods = pd.concat(all_intersecting_periods)
     return all_intersecting_periods.sort_values(by="start_dt").reset_index(drop=True)

{ocf_data_sampler-0.2.36 → ocf_data_sampler-0.2.38}/ocf_data_sampler/torch_datasets/datasets/pvnet_uk.py

@@ -30,21 +30,11 @@ from ocf_data_sampler.torch_datasets.utils.merge_and_fill_utils import (
     fill_nans_in_arrays,
     merge_dicts,
 )
-from ocf_data_sampler.utils import minutes
+from ocf_data_sampler.utils import compute, minutes
 
 xr.set_options(keep_attrs=True)
 
 
-def compute(xarray_dict: dict) -> dict:
-    """Eagerly load a nested dictionary of xarray DataArrays."""
-    for k, v in xarray_dict.items():
-        if isinstance(v, dict):
-            xarray_dict[k] = compute(v)
-        else:
-            xarray_dict[k] = v.compute(scheduler="single-threaded")
-    return xarray_dict
-
-
 def get_gsp_locations(
     gsp_ids: list[int] | None = None,
     version: str = "20220314",

{ocf_data_sampler-0.2.36 → ocf_data_sampler-0.2.38}/ocf_data_sampler/torch_datasets/datasets/site.py

@@ -6,7 +6,7 @@ import xarray as xr
 from torch.utils.data import Dataset
 from typing_extensions import override
 
-from ocf_data_sampler.config import load_yaml_configuration
+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.numpy_sample import (
     NWPSampleKey,
@@ -16,6 +16,7 @@ from ocf_data_sampler.numpy_sample import (
     make_datetime_numpy_dict,
     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 NumpySample
 from ocf_data_sampler.select import (
     Location,
@@ -33,11 +34,31 @@ from ocf_data_sampler.torch_datasets.utils.merge_and_fill_utils import (
     fill_nans_in_arrays,
     merge_dicts,
 )
-from ocf_data_sampler.utils import minutes
+from ocf_data_sampler.utils import compute, minutes
 
 xr.set_options(keep_attrs=True)
 
 
+def get_locations(site_xr: xr.Dataset) -> list[Location]:
+    """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)
+        location = Location(
+            id=site_id,
+            x=site.longitude.values,
+            y=site.latitude.values,
+            coordinate_system="lon_lat",
+        )
+        locations.append(location)
+
+    return locations
+
+
 class SitesDataset(Dataset):
     """A torch Dataset for creating PVNet Site samples."""
 
@@ -62,7 +83,7 @@ class SitesDataset(Dataset):
         self.config = config
 
         # get all locations
-        self.locations = self.get_locations(datasets_dict["site"])
+        self.locations = 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
@@ -82,48 +103,6 @@ class SitesDataset(Dataset):
         # Assign coords and indices to self
         self.valid_t0_and_site_ids = valid_t0_and_site_ids
 
-    @override
-    def __len__(self) -> int:
-        return len(self.valid_t0_and_site_ids)
-
-    @override
-    def __getitem__(self, idx: int) -> dict:
-        # Get the coordinates of the sample
-        t0, site_id = self.valid_t0_and_site_ids.iloc[idx]
-
-        # get location from site id
-        location = self.location_lookup[site_id]
-
-        # Generate the sample
-        return self._get_sample(t0, location)
-
-    def _get_sample(self, t0: pd.Timestamp, location: Location) -> dict:
-        """Generate the PVNet sample for given coordinates.
-
-        Args:
-            t0: init-time for sample
-            location: location for sample
-        """
-        sample_dict = slice_datasets_by_space(self.datasets_dict, location, self.config)
-        sample_dict = slice_datasets_by_time(sample_dict, t0, self.config)
-
-        sample = self.process_and_combine_site_sample_dict(sample_dict, t0)
-        return sample.compute()
-
-    def get_sample(self, t0: pd.Timestamp, site_id: int) -> dict:
-        """Generate a sample for a given site id and t0.
-
-        Useful for users to generate samples by t0 and site id
-
-        Args:
-            t0: init-time for sample
-            site_id: site id as int
-        """
-        location = self.location_lookup[site_id]
-
-        return self._get_sample(t0, location)
-
-
     def find_valid_t0_and_site_ids(
         self,
         datasets_dict: dict,
@@ -177,25 +156,46 @@ class SitesDataset(Dataset):
         valid_t0_and_site_ids.index.name = "t0"
         return valid_t0_and_site_ids.reset_index()
 
+    @override
+    def __len__(self) -> int:
+        return len(self.valid_t0_and_site_ids)
+
+    @override
+    def __getitem__(self, idx: int) -> dict:
+        # Get the coordinates of the sample
+        t0, site_id = self.valid_t0_and_site_ids.iloc[idx]
 
-    def get_locations(self, site_xr: xr.Dataset) -> list[Location]:
-        """Get list of locations of all sites.
+        # get location from site id
+        location = self.location_lookup[site_id]
+
+        # Generate the sample
+        return self._get_sample(t0, location)
+
+    def _get_sample(self, t0: pd.Timestamp, location: Location) -> dict:
+        """Generate the PVNet sample for given coordinates.
 
         Args:
-            site_xr: xarray Dataset of site data
+            t0: init-time for sample
+            location: location for sample
         """
-        locations = []
-        for site_id in site_xr.site_id.values:
-            site = site_xr.sel(site_id=site_id)
-            location = Location(
-                id=site_id,
-                x=site.longitude.values,
-                y=site.latitude.values,
-                coordinate_system="lon_lat",
-            )
-            locations.append(location)
+        sample_dict = slice_datasets_by_space(self.datasets_dict, location, self.config)
+        sample_dict = slice_datasets_by_time(sample_dict, t0, self.config)
+
+        sample = self.process_and_combine_site_sample_dict(sample_dict, t0)
+        return sample.compute()
+
+    def get_sample(self, t0: pd.Timestamp, site_id: int) -> dict:
+        """Generate a sample for a given site id and t0.
+
+        Useful for users to generate samples by t0 and site id
 
-        return locations
+        Args:
+            t0: init-time for sample
+            site_id: site id as int
+        """
+        location = self.location_lookup[site_id]
+
+        return self._get_sample(t0, location)
 
     def process_and_combine_site_sample_dict(
         self,
@@ -256,8 +256,8 @@ class SitesDataset(Dataset):
 
         # Only add solar position if explicitly configured
         has_solar_config = (
-            hasattr(self.config.input_data, "solar_position") and
-            self.config.input_data.solar_position is not None
+            hasattr(self.config.input_data, "solar_position")
+            and self.config.input_data.solar_position is not None
         )
 
         if has_solar_config:
@@ -351,6 +351,216 @@ class SitesDataset(Dataset):
         return combined_dataset
 
 
+class SitesDatasetConcurrent(Dataset):
+    """A torch Dataset for creating PVNet Site batches with samples for all sites."""
+
+    def __init__(
+        self,
+        config_filename: str,
+        start_time: str | None = None,
+        end_time: str | None = None,
+    ) -> None:
+        """A torch Dataset for creating PVNet Site samples.
+
+        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
+        """
+        config = load_yaml_configuration(config_filename)
+        datasets_dict = get_dataset_dict(config.input_data)
+
+        # Assign config and input data to self
+        self.datasets_dict = datasets_dict
+        self.config = config
+
+        # get all locations
+        self.locations = get_locations(datasets_dict["site"])
+
+        # Get t0 times where all input data is available
+        valid_t0s = self.find_valid_t0s(datasets_dict)
+
+        # Filter t0 times to given range
+        if start_time is not None:
+            valid_t0s = valid_t0s[
+                valid_t0s >= pd.Timestamp(start_time)
+            ]
+
+        if end_time is not None:
+            valid_t0s = valid_t0s[
+                valid_t0s <= pd.Timestamp(end_time)
+            ]
+
+        # Assign coords and indices to self
+        self.valid_t0s = valid_t0s
+
+    @staticmethod
+    def process_and_combine_datasets(
+        dataset_dict: dict,
+        config: Configuration,
+        t0: pd.Timestamp,
+    ) -> NumpySample:
+        """Normalise and convert data to numpy arrays.
+
+        Args:
+            dataset_dict: Dictionary of xarray datasets
+            config: Configuration object
+            t0: init-time for 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 "site" in dataset_dict:
+            da_sites = dataset_dict["site"]
+            da_sites = da_sites / da_sites.capacity_kwp
+
+            # Convert to NumpyBatch
+            numpy_modalities.append(
+                convert_site_to_numpy_sample(
+                    da_sites,
+                ),
+            )
+
+        # 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),
+            )
+
+            # Calculate solar positions and add to modalities
+            numpy_modalities.append(
+                make_sun_position_numpy_sample(
+                    datetimes, da_sites.longitude.values, da_sites.latitude.values,
+                ),
+            )
+
+        # 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 find_valid_t0s(
+        self,
+        datasets_dict: dict,
+    ) -> pd.DataFrame:
+        """Find the t0 times where all of the requested input data is available.
+
+        The idea is to
+        1. Get valid time period for nwp and satellite
+        2. For the first site location, find valid periods for that location
+        Note there is an assumption that all sites have the same t0 data available
+
+        Args:
+            datasets_dict: A dictionary of input datasets
+        """
+        # 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)
+        sites = datasets_dict["site"]
+
+        # Taking just the first site value, assume t0s the same for all of them
+        site_id = sites.site_id.values[0]
+        site_config = self.config.input_data.site
+        site = sites.sel(site_id=site_id)
+        # Drop NaN values
+        site = site.dropna(dim="time_utc")
+
+        # 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),
+            interval_start=minutes(site_config.interval_start_minutes),
+            interval_end=minutes(site_config.interval_end_minutes),
+        )
+        valid_time_periods_per_site = intersection_of_multiple_dataframes_of_periods(
+            [valid_time_periods, time_periods],
+        )
+
+        # 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),
+        )
+
+        return valid_t0_times_per_site
+
+    @override
+    def __len__(self) -> int:
+        return len(self.valid_t0s)
+
+    @override
+    def __getitem__(self, idx: int) -> dict:
+        # Get the coordinates of the sample
+        t0 = self.valid_t0s[idx]
+
+        return self._get_batch(t0)
+
+    def _get_batch(self, t0: pd.Timestamp) -> dict:
+        """Generate the PVNet batch for given coordinates.
+
+        Args:
+            t0: init-time for sample
+        """
+        # slice by time first as we want to keep all site id info
+        sample_dict = slice_datasets_by_time(self.datasets_dict, t0, self.config)
+        sample_dict = compute(sample_dict)
+
+        site_samples = []
+
+        for location in self.locations:
+            site_sample_dict = slice_datasets_by_space(sample_dict, location, self.config)
+            site_numpy_sample = self.process_and_combine_datasets(
+                site_sample_dict,
+                self.config,
+                t0,
+            )
+            site_samples.append(site_numpy_sample)
+
+        return stack_np_samples_into_batch(site_samples)
+
+
 # ----- functions to load presaved samples ------
 
 

ocf_data_sampler-0.2.38/ocf_data_sampler/utils.py

@@ -0,0 +1,21 @@
+"""Miscellaneous helper functions."""
+
+import pandas as pd
+
+
+def minutes(minutes: int | list[float]) -> pd.Timedelta | pd.TimedeltaIndex:
+    """Timedelta minutes.
+
+    Args:
+        minutes: the number of minutes, single value or list
+    """
+    return pd.to_timedelta(minutes, unit="m")
+
+def compute(xarray_dict: dict) -> dict:
+    """Eagerly load a nested dictionary of xarray DataArrays."""
+    for k, v in xarray_dict.items():
+        if isinstance(v, dict):
+            xarray_dict[k] = compute(v)
+        else:
+            xarray_dict[k] = v.compute(scheduler="single-threaded")
+    return xarray_dict

{ocf_data_sampler-0.2.36 → ocf_data_sampler-0.2.38}/ocf_data_sampler.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ocf-data-sampler
-Version: 0.2.36
+Version: 0.2.38
 Author: James Fulton, Peter Dudfield
 Author-email: Open Climate Fix team <info@openclimatefix.org>
 License: MIT License

ocf_data_sampler-0.2.36/ocf_data_sampler/utils.py

@@ -1,12 +0,0 @@
-"""Miscellaneous helper functions."""
-
-import pandas as pd
-
-
-def minutes(minutes: int | list[float]) -> pd.Timedelta | pd.TimedeltaIndex:
-    """Timedelta minutes.
-
-    Args:
-        minutes: the number of minutes, single value or list
-    """
-    return pd.to_timedelta(minutes, unit="m")