ocf-data-sampler 0.4.0__tar.gz → 0.5.0__tar.gz

{ocf_data_sampler-0.4.0 → ocf_data_sampler-0.5.0}/PKG-INFO

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

{ocf_data_sampler-0.4.0 → ocf_data_sampler-0.5.0}/ocf_data_sampler/numpy_sample/__init__.py

@@ -5,5 +5,4 @@ from .gsp import convert_gsp_to_numpy_sample, GSPSampleKey
 from .nwp import convert_nwp_to_numpy_sample, NWPSampleKey
 from .satellite import convert_satellite_to_numpy_sample, SatelliteSampleKey
 from .sun_position import make_sun_position_numpy_sample
-from .site import convert_site_to_numpy_sample
-
+from .site import convert_site_to_numpy_sample, SiteSampleKey

ocf_data_sampler-0.5.0/ocf_data_sampler/torch_datasets/datasets/__init__.py

@@ -0,0 +1,2 @@
+from .pvnet_uk import PVNetUKRegionalDataset, PVNetUKConcurrentDataset
+from .site import SitesDataset

{ocf_data_sampler-0.4.0 → ocf_data_sampler-0.5.0}/ocf_data_sampler/torch_datasets/datasets/site.py

@@ -58,6 +58,96 @@ def get_locations(site_xr: xr.Dataset) -> list[Location]:
 
     return locations
 
+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,
+            ),
+        )
+
+        # add datetime features
+        datetimes = pd.DatetimeIndex(da_sites.time_utc.values)
+        datetime_features = encode_datetimes(datetimes=datetimes)
+
+        numpy_modalities.append(datetime_features)
+
+    # Only add solar position if explicitly configured
+    if config.input_data.solar_position is not None:
+        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
+
 
 class SitesDataset(Dataset):
     """A torch Dataset for creating PVNet Site samples."""
@@ -181,8 +271,9 @@ class SitesDataset(Dataset):
         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()
+        sample_dict = compute(sample_dict)
+
+        return process_and_combine_datasets(sample_dict, self.config, t0)
 
     def get_sample(self, t0: pd.Timestamp, site_id: int) -> dict:
         """Generate a sample for a given site id and t0.
@@ -197,159 +288,6 @@ class SitesDataset(Dataset):
 
         return self._get_sample(t0, location)
 
-    def process_and_combine_site_sample_dict(
-        self,
-        dataset_dict: dict,
-        t0: pd.Timestamp,
-    ) -> xr.Dataset:
-        """Normalize and combine data into a single xr Dataset.
-
-        Args:
-            dataset_dict: dict containing sliced xr DataArrays
-            t0: The initial timestamp of the sample
-
-        Returns:
-            xr.Dataset: A merged Dataset with nans filled in.
-        """
-        data_arrays = []
-
-        if "nwp" in dataset_dict:
-            for nwp_key, da_nwp in dataset_dict["nwp"].items():
-                provider = self.config.input_data.nwp[nwp_key].provider
-
-                da_channel_means = channel_dict_to_dataarray(
-                    self.config.input_data.nwp[nwp_key].channel_means,
-                )
-                da_channel_stds = channel_dict_to_dataarray(
-                    self.config.input_data.nwp[nwp_key].channel_stds,
-                )
-
-                da_nwp = (da_nwp - da_channel_means) / da_channel_stds
-                data_arrays.append((f"nwp-{provider}", da_nwp))
-
-        if "sat" in dataset_dict:
-            da_sat = dataset_dict["sat"]
-
-            da_channel_means = channel_dict_to_dataarray(
-                self.config.input_data.satellite.channel_means,
-            )
-            da_channel_stds = channel_dict_to_dataarray(
-                self.config.input_data.satellite.channel_stds,
-            )
-
-            da_sat = (da_sat - da_channel_means) / da_channel_stds
-            data_arrays.append(("satellite", da_sat))
-
-        if "site" in dataset_dict:
-            da_sites = dataset_dict["site"]
-            da_sites = da_sites / da_sites.capacity_kwp
-            data_arrays.append(("site", da_sites))
-
-        combined_sample_dataset = self.merge_data_arrays(data_arrays)
-
-        # add datetime features
-        datetimes = pd.DatetimeIndex(combined_sample_dataset.site__time_utc.values)
-        datetime_features = encode_datetimes(datetimes=datetimes)
-        combined_sample_dataset = combined_sample_dataset.assign_coords(
-            {k: ("site__time_utc", v) for k, v in datetime_features.items()},
-        )
-
-        # 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
-        )
-
-        if has_solar_config:
-            solar_config = self.config.input_data.solar_position
-
-            # Datetime range - solar config params
-            solar_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 sun position features
-            sun_position_features = make_sun_position_numpy_sample(
-                datetimes=solar_datetimes,
-                lon=combined_sample_dataset.site__longitude.values,
-                lat=combined_sample_dataset.site__latitude.values,
-            )
-
-            # Use existing dimension for solar positions
-            # TODO decouple this as a separate data varaible
-            solar_dim_name = "site__time_utc"
-
-            # Assign solar position values
-            for key, values in sun_position_features.items():
-                combined_sample_dataset = combined_sample_dataset.assign_coords(
-                    {key: (solar_dim_name, values)},
-                )
-
-        # TODO include t0_index in xr dataset?
-
-        # Fill any nan values
-        return combined_sample_dataset.fillna(0.0)
-
-    def merge_data_arrays(
-        self,
-        normalised_data_arrays: list[tuple[str, xr.DataArray]],
-    ) -> xr.Dataset:
-        """Combine a list of DataArrays into a single Dataset with unique naming conventions.
-
-        Args:
-            normalised_data_arrays: List of tuples where each tuple contains:
-                - A string (key name).
-                - An xarray.DataArray.
-
-        Returns:
-            xr.Dataset: A merged Dataset with uniquely named variables, coordinates, and dimensions.
-        """
-        datasets = []
-
-        for key, data_array in normalised_data_arrays:
-            # Ensure all attributes are strings for consistency
-            data_array = data_array.assign_attrs(
-                {attr_key: str(attr_value) for attr_key, attr_value in data_array.attrs.items()},
-            )
-
-            # Convert DataArray to Dataset with the variable name as the key
-            dataset = data_array.to_dataset(name=key)
-
-            # Prepend key name to all dimension and coordinate names for uniqueness
-            dataset = dataset.rename(
-                {dim: f"{key}__{dim}" for dim in dataset.dims if dim not in dataset.coords},
-            )
-            dataset = dataset.rename(
-                {coord: f"{key}__{coord}" for coord in dataset.coords},
-            )
-
-            # Handle concatenation dimension if applicable
-            concat_dim = (
-                f"{key}__target_time_utc"
-                if f"{key}__target_time_utc" in dataset.coords
-                else f"{key}__time_utc"
-            )
-
-            if f"{key}__init_time_utc" in dataset.coords:
-                init_coord = f"{key}__init_time_utc"
-                if dataset[init_coord].ndim == 0:  # Check if scalar
-                    expanded_init_times = [dataset[init_coord].values] * len(dataset[concat_dim])
-                    dataset = dataset.assign_coords({init_coord: (concat_dim, expanded_init_times)})
-
-            datasets.append(dataset)
-
-        # Ensure all datasets are valid xarray.Dataset objects
-        for ds in datasets:
-            if not isinstance(ds, xr.Dataset):
-                raise ValueError(f"Object is not an xr.Dataset: {type(ds)}")
-
-        # Merge all prepared datasets
-        combined_dataset = xr.merge(datasets)
-
-        return combined_dataset
-
 
 class SitesDatasetConcurrent(Dataset):
     """A torch Dataset for creating PVNet Site batches with samples for all sites."""
@@ -394,89 +332,6 @@ class SitesDatasetConcurrent(Dataset):
         # 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,
@@ -547,7 +402,7 @@ class SitesDatasetConcurrent(Dataset):
 
         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_numpy_sample = process_and_combine_datasets(
                 site_sample_dict,
                 self.config,
                 t0,
@@ -557,122 +412,6 @@ class SitesDatasetConcurrent(Dataset):
         return stack_np_samples_into_batch(site_samples)
 
 
-# ----- functions to load presaved samples ------
-
-
-def convert_netcdf_to_numpy_sample(ds: xr.Dataset) -> dict:
-    """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 sat # TODO this could be improved
-        sample_dict["sat"] = sample_dict.pop("satellite")
-
-    # process and combine the datasets
-    sample = convert_to_numpy_and_combine(dataset_dict=sample_dict)
-
-    # Add solar coord and datetime features
-    keys = ["solar_azimuth", "solar_elevation", "date_sin", "date_cos", "time_sin", "time_cos"]
-    for key in keys:
-        if key in ds.coords:
-            sample[key] = ds.coords[key].values
-
-    # TODO think about normalization:
-    # * maybe its done not in sample creation, maybe its done afterwards,
-    #   to allow it to be flexible
-
-    return sample
-
-
-def convert_from_dataset_to_dict_datasets(combined_dataset: xr.Dataset) -> dict[str, xr.DataArray]:
-    """Convert a combined sample dataset to a dict of datasets for each input.
-
-    Args:
-        combined_dataset: The combined NetCDF dataset
-
-    Returns:
-        The uncombined datasets as a dict of xr.Datasets
-    """
-    # Split into datasets by splitting by the prefix added in combine_to_netcdf
-    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
-        for dim in list(dataset.coords):
-            if f"{key}__" not in 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 = 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
-    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>'.
-
-    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 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): 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[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:
-        nwp_numpy_modalities = {}
-        for nwp_key, da_nwp in dataset_dict["nwp"].items():
-            # Convert to NumpySample
-            nwp_numpy_modalities[nwp_key] = convert_nwp_to_numpy_sample(da_nwp)
-
-        # Combine the NWPs into NumpySample
-        numpy_modalities.append({NWPSampleKey.nwp: nwp_numpy_modalities})
-
-    if "sat" in dataset_dict:
-        # Satellite is already in the range [0-1] so no need to standardise
-        da_sat = dataset_dict["sat"]
-
-        # Convert to NumpySample
-        numpy_modalities.append(convert_satellite_to_numpy_sample(da_sat))
-
-    if "site" in dataset_dict:
-        da_sites = dataset_dict["site"]
-
-        numpy_modalities.append(convert_site_to_numpy_sample(da_sites))
-
-    # Combine all the modalities and fill NaNs
-    combined_sample = merge_dicts(numpy_modalities)
-    return fill_nans_in_arrays(combined_sample)
-
-
 def coarsen_data(xr_data: xr.Dataset, coarsen_to_deg: float = 0.1) -> xr.Dataset:
     """Coarsen the data to a specified resolution in degrees.
 

ocf_data_sampler-0.5.0/ocf_data_sampler/torch_datasets/sample/site.py

@@ -0,0 +1,48 @@
+"""PVNet Site sample implementation for netCDF data handling and conversion."""
+
+import torch
+from typing_extensions import override
+
+from ocf_data_sampler.numpy_sample.common_types import NumpySample
+
+from .base import SampleBase
+
+
+# TODO this is now similar to the UKRegionalSample
+# We should consider just having one Sample class for all datasets
+class SiteSample(SampleBase):
+    """Handles SiteSample specific operations."""
+
+    def __init__(self, data: NumpySample) -> None:
+        """Initializes the SiteSample object with the given NumpySample."""
+        self._data = data
+
+    @override
+    def to_numpy(self) -> NumpySample:
+        return self._data
+
+    @override
+    def save(self, path: str) -> None:
+        """Saves sample to the specified path in pickle format."""
+        # Saves to pickle format
+        torch.save(self._data, path)
+
+    @classmethod
+    @override
+    def load(cls, path: str) -> "SiteSample":
+        """Loads sample from the specified path.
+
+        Args:
+            path: Path to the saved sample file.
+
+        Returns:
+            A SiteSample instance with the loaded data.
+        """
+        # Loads from .pt format
+        # TODO: We should move away from using torch.load(..., weights_only=False)
+        return cls(torch.load(path, weights_only=False))
+
+    @override
+    def plot(self) -> None:
+        # TODO - placeholder for now
+        raise NotImplementedError("Plotting not yet implemented for SiteSample")

{ocf_data_sampler-0.4.0 → ocf_data_sampler-0.5.0}/ocf_data_sampler.egg-info/PKG-INFO

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

ocf_data_sampler-0.4.0/ocf_data_sampler/torch_datasets/datasets/__init__.py

@@ -1,6 +0,0 @@
-from .pvnet_uk import PVNetUKRegionalDataset, PVNetUKConcurrentDataset
-
-from .site import (
-    convert_netcdf_to_numpy_sample,
-    SitesDataset
-)

ocf_data_sampler-0.4.0/ocf_data_sampler/torch_datasets/sample/site.py

@@ -1,39 +0,0 @@
-"""PVNet Site sample implementation for netCDF data handling and conversion."""
-
-import xarray as xr
-from typing_extensions import override
-
-from ocf_data_sampler.numpy_sample.common_types import NumpySample
-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."""
-
-    def __init__(self, data: xr.Dataset) -> None:
-        """Initializes the SiteSample object with the given xarray Dataset."""
-        if not isinstance(data, xr.Dataset):
-            raise TypeError(f"Data must be xarray Dataset - Found type {type(data)}")
-        self._data = data
-
-    @override
-    def to_numpy(self) -> NumpySample:
-        return convert_netcdf_to_numpy_sample(self._data)
-
-    @override
-    def save(self, path: str) -> None:
-        # Saves as NetCDF
-        self._data.to_netcdf(path, mode="w", engine="h5netcdf")
-
-    @classmethod
-    @override
-    def load(cls, path: str) -> "SiteSample":
-        # Loads from NetCDF
-        return cls(xr.open_dataset(path, decode_timedelta=False))
-
-    @override
-    def plot(self) -> None:
-        # TODO - placeholder for now
-        raise NotImplementedError("Plotting not yet implemented for SiteSample")