ocf-data-sampler 0.5.18__tar.gz → 0.5.20__tar.gz

{ocf_data_sampler-0.5.18 → ocf_data_sampler-0.5.20}/PKG-INFO

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

{ocf_data_sampler-0.5.18 → ocf_data_sampler-0.5.20}/ocf_data_sampler/load/gsp.py

@@ -32,7 +32,7 @@ def open_gsp(
     boundaries_version: str = "20220314",
     public: bool = False,
 ) -> xr.DataArray:
-    """Open the GSP data and validates its data types.
+    """Open and eagerly load the GSP data and validates its data types.
 
     Args:
         zarr_path: Path to the GSP zarr data
@@ -93,4 +93,6 @@ def open_gsp(
             dtype = gsp_da.coords[coord].dtype
             raise TypeError(f"{coord} should be {expected_dtype.__name__}, not {dtype}")
 
-    return gsp_da
+    # Below we load the data eagerly into memory - this makes the dataset faster to sample from, but
+    # at the cost of a little extra memory usage
+    return gsp_da.compute()

ocf_data_sampler-0.5.20/ocf_data_sampler/torch_datasets/datasets/picklecache.py

@@ -0,0 +1,33 @@
+"""A mixin to handle pickling and caching a dataset's state."""
+
+import os
+import pickle
+
+
+class PickleCacheMixin:
+    """A mixin for classes that need to cache their state using pickle."""
+    def __init__(self, *args: list, **kwargs: dict) -> None:
+        """Initialize the pickle path and call the parent constructor."""
+        self._pickle_path = None
+        super().__init__(*args, **kwargs)  # cooperative multiple inheritance
+
+    def presave_pickle(self, pickle_path: str) -> None:
+        """Save the full object state to a pickle file and store the pickle path."""
+        self._pickle_path = pickle_path
+        with open(pickle_path, "wb") as f:
+            pickle.dump(self.__dict__, f)
+
+    def __getstate__(self) -> dict:
+        """If presaved, only pickle reference. Otherwise pickle everything."""
+        if self._pickle_path:
+            return {"_pickle_path": self._pickle_path}
+        else:
+            return self.__dict__
+
+    def __setstate__(self, state: dict) -> None:
+        """Restore object from pickle, reloading from presaved file if possible."""
+        self.__dict__.update(state)
+        if self._pickle_path and os.path.exists(self._pickle_path):
+            with open(self._pickle_path, "rb") as f:
+                saved_state = pickle.load(f) # noqa: S301
+                self.__dict__.update(saved_state)

{ocf_data_sampler-0.5.18 → ocf_data_sampler-0.5.20}/ocf_data_sampler/torch_datasets/datasets/pvnet_uk.py

@@ -19,6 +19,7 @@ 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.torch_datasets.datasets.picklecache import PickleCacheMixin
 from ocf_data_sampler.torch_datasets.utils import (
     add_alterate_coordinate_projections,
     config_normalization_values_to_dicts,
@@ -67,7 +68,7 @@ def get_gsp_locations(
     return locations
 
 
-class AbstractPVNetUKDataset(Dataset):
+class AbstractPVNetUKDataset(PickleCacheMixin, Dataset):
     """Abstract class for PVNet UK datasets."""
 
     def __init__(
@@ -85,6 +86,8 @@ class AbstractPVNetUKDataset(Dataset):
             end_time: Limit the init-times to be before this
             gsp_ids: List of GSP IDs to create samples for. Defaults to all
         """
+        super().__init__()
+
         config = load_yaml_configuration(config_filename)
         datasets_dict = get_dataset_dict(config.input_data, gsp_ids=gsp_ids)
 
@@ -225,7 +228,6 @@ class AbstractPVNetUKDataset(Dataset):
         return valid_t0_times
 
 
-
 class PVNetUKRegionalDataset(AbstractPVNetUKDataset):
     """A torch Dataset for creating PVNet UK regional samples."""
 

{ocf_data_sampler-0.5.18 → ocf_data_sampler-0.5.20}/ocf_data_sampler/torch_datasets/datasets/site.py

@@ -24,6 +24,7 @@ from ocf_data_sampler.select import (
     find_contiguous_t0_periods,
     intersection_of_multiple_dataframes_of_periods,
 )
+from ocf_data_sampler.torch_datasets.datasets.picklecache import PickleCacheMixin
 from ocf_data_sampler.torch_datasets.utils import (
     add_alterate_coordinate_projections,
     config_normalization_values_to_dicts,
@@ -144,7 +145,7 @@ def process_and_combine_datasets(
     return combined_sample
 
 
-class SitesDataset(Dataset):
+class SitesDataset(PickleCacheMixin, Dataset):
     """A torch Dataset for creating PVNet Site samples."""
 
     def __init__(
@@ -160,6 +161,8 @@ class SitesDataset(Dataset):
             start_time: Limit the init-times to be after this
             end_time: Limit the init-times to be before this
         """
+        super().__init__()
+
         config = load_yaml_configuration(config_filename)
         datasets_dict = get_dataset_dict(config.input_data)
 
@@ -301,7 +304,7 @@ class SitesDataset(Dataset):
         return self._get_sample(t0, location)
 
 
-class SitesDatasetConcurrent(Dataset):
+class SitesDatasetConcurrent(PickleCacheMixin, Dataset):
     """A torch Dataset for creating PVNet Site batches with samples for all sites."""
 
     def __init__(
@@ -317,6 +320,8 @@ class SitesDatasetConcurrent(Dataset):
             start_time: Limit the init-times to be after this
             end_time: Limit the init-times to be before this
         """
+        super().__init__()
+
         config = load_yaml_configuration(config_filename)
         datasets_dict = get_dataset_dict(config.input_data)
 

{ocf_data_sampler-0.5.18 → ocf_data_sampler-0.5.20}/ocf_data_sampler.egg-info/PKG-INFO

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

{ocf_data_sampler-0.5.18 → ocf_data_sampler-0.5.20}/ocf_data_sampler.egg-info/SOURCES.txt

@@ -49,6 +49,7 @@ ocf_data_sampler/select/location.py
 ocf_data_sampler/select/select_spatial_slice.py
 ocf_data_sampler/select/select_time_slice.py
 ocf_data_sampler/torch_datasets/datasets/__init__.py
+ocf_data_sampler/torch_datasets/datasets/picklecache.py
 ocf_data_sampler/torch_datasets/datasets/pvnet_uk.py
 ocf_data_sampler/torch_datasets/datasets/site.py
 ocf_data_sampler/torch_datasets/sample/__init__.py