ocf-data-sampler 0.5.19__py3-none-any.whl → 0.5.20__py3-none-any.whl

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/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/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.19.dist-info → ocf_data_sampler-0.5.20.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ocf-data-sampler
-Version: 0.5.19
+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.19.dist-info → ocf_data_sampler-0.5.20.dist-info}/RECORD

@@ -41,8 +41,9 @@ ocf_data_sampler/select/location.py,sha256=Qp0di-Pgq8WLjN9IBcTVTaRM3lckhr4ZVzaDR
 ocf_data_sampler/select/select_spatial_slice.py,sha256=NB6NtZBc_Mb5zPCItzBIEa_Nroj2kEsjUIsa_kdWoj0,7081
 ocf_data_sampler/select/select_time_slice.py,sha256=cpkdovJMvcjxSGfq9G0OJK5aDAeCXg7exWYrJnR4N2w,4116
 ocf_data_sampler/torch_datasets/datasets/__init__.py,sha256=o0SsEXXZ6k9iL__5_RN1Sf60lw_eqK91P3UFEHAD2k0,102
-ocf_data_sampler/torch_datasets/datasets/pvnet_uk.py,sha256=wUsIZ0Fhq5bbE8v02C0UPcFWIhWI7kfSka9UrWP0_m4,12240
-ocf_data_sampler/torch_datasets/datasets/site.py,sha256=OXrYSRrWUdQbEjsEPPJjam10zJKU6S3r5kA07RbpzFU,15680
+ocf_data_sampler/torch_datasets/datasets/picklecache.py,sha256=b8T5lgKfiPXLwuVQuFpCQBlU-HNBrA-Z-eSwYICKvsQ,1350
+ocf_data_sampler/torch_datasets/datasets/pvnet_uk.py,sha256=n0NYbhrMR_LtGz4EiryljeQG-tI4fWfJLifKclMThcU,12367
+ocf_data_sampler/torch_datasets/datasets/site.py,sha256=fZyGwUmcFMYuEJd05fZc96mLEya8ehYnVRVRssX52qM,15854
 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=40NwNTqjL1WVhPdwe02zDHHfDLG2u_bvCfRCtGAtFc0,1466
@@ -58,7 +59,7 @@ ocf_data_sampler/torch_datasets/utils/valid_time_periods.py,sha256=xcy75cVxl0Wrg
 ocf_data_sampler/torch_datasets/utils/validation_utils.py,sha256=YqmT-lExWlI8_ul3l0EP73Ik002fStr_bhsZh9mQqEU,4735
 scripts/download_gsp_location_data.py,sha256=rRDXMoqX-RYY4jPdxhdlxJGhWdl6r245F5UARgKV6P4,3121
 scripts/refactor_site.py,sha256=skzvsPP0Cn9yTKndzkilyNcGz4DZ88ctvCJ0XrBdc2A,3135
-ocf_data_sampler-0.5.19.dist-info/METADATA,sha256=CyMcURh5QvAzzxVvu_DzbK4krMLtGtSlVjr4KURpFiw,12817
-ocf_data_sampler-0.5.19.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ocf_data_sampler-0.5.19.dist-info/top_level.txt,sha256=deUxqmsONNAGZDNbsntbXH7BRA1MqWaUeAJrCo6q_xA,25
-ocf_data_sampler-0.5.19.dist-info/RECORD,,
+ocf_data_sampler-0.5.20.dist-info/METADATA,sha256=AKZ_mCJGD1NgN_ZG9YF9pa0CZ9fEZOwWamZPmM8jBXc,12817
+ocf_data_sampler-0.5.20.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ocf_data_sampler-0.5.20.dist-info/top_level.txt,sha256=deUxqmsONNAGZDNbsntbXH7BRA1MqWaUeAJrCo6q_xA,25
+ocf_data_sampler-0.5.20.dist-info/RECORD,,