ocf-data-sampler 0.5.7__py3-none-any.whl → 0.5.8__py3-none-any.whl

ocf_data_sampler/torch_datasets/datasets/pvnet_uk.py

@@ -30,7 +30,7 @@ from ocf_data_sampler.torch_datasets.utils.merge_and_fill_utils import (
     fill_nans_in_arrays,
     merge_dicts,
 )
-from ocf_data_sampler.utils import compute, minutes
+from ocf_data_sampler.utils import minutes, tensorstore_compute
 
 xr.set_options(keep_attrs=True)
 
@@ -254,7 +254,7 @@ class PVNetUKRegionalDataset(AbstractPVNetUKDataset):
         """
         sample_dict = slice_datasets_by_space(self.datasets_dict, location, self.config)
         sample_dict = slice_datasets_by_time(sample_dict, t0, self.config)
-        sample_dict = compute(sample_dict)
+        sample_dict = tensorstore_compute(sample_dict)
 
         return self.process_and_combine_datasets(sample_dict, t0, location)
 
@@ -313,7 +313,7 @@ class PVNetUKConcurrentDataset(AbstractPVNetUKDataset):
         """
         # 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)
+        sample_dict = tensorstore_compute(sample_dict)
 
         gsp_samples = []
 

ocf_data_sampler/torch_datasets/datasets/site.py

@@ -34,7 +34,7 @@ from ocf_data_sampler.torch_datasets.utils.merge_and_fill_utils import (
     fill_nans_in_arrays,
     merge_dicts,
 )
-from ocf_data_sampler.utils import compute, minutes
+from ocf_data_sampler.utils import minutes, tensorstore_compute
 
 xr.set_options(keep_attrs=True)
 
@@ -272,7 +272,7 @@ 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_dict = compute(sample_dict)
+        sample_dict = tensorstore_compute(sample_dict)
 
         return process_and_combine_datasets(
             sample_dict,
@@ -408,7 +408,7 @@ class SitesDatasetConcurrent(Dataset):
         """
         # 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)
+        sample_dict = tensorstore_compute(sample_dict)
 
         site_samples = []
 

ocf_data_sampler/utils.py

@@ -1,6 +1,7 @@
 """Miscellaneous helper functions."""
 
 import pandas as pd
+from xarray_tensorstore import read
 
 
 def minutes(minutes: int | list[float]) -> pd.Timedelta | pd.TimedeltaIndex:
@@ -11,11 +12,26 @@ def minutes(minutes: int | list[float]) -> pd.Timedelta | pd.TimedeltaIndex:
     """
     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")
+            xarray_dict[k] = v.compute()
     return xarray_dict
+
+
+def tensorstore_compute(xarray_dict: dict) -> dict:
+    """Eagerly read and load a nested dictionary of xarray-tensorstore DataArrays."""
+    # Kick off the tensorstore async reading
+    for k, v in xarray_dict.items():
+        if isinstance(v, dict):
+            xarray_dict[k] = tensorstore_compute(v)
+        else:
+            xarray_dict[k] = read(v)
+
+    # Running the compute function will wait until all arrays have been read
+    return compute(xarray_dict)
+

{ocf_data_sampler-0.5.7.dist-info → ocf_data_sampler-0.5.8.dist-info}/METADATA

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

{ocf_data_sampler-0.5.7.dist-info → ocf_data_sampler-0.5.8.dist-info}/RECORD

@@ -1,5 +1,5 @@
 ocf_data_sampler/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-ocf_data_sampler/utils.py,sha256=2NEl70ySdTpr0pbLRk4LGklvXe1Nv1hun9XKcDw7-44,610
+ocf_data_sampler/utils.py,sha256=0Wlx7SNOJE5ZNs_F3m-fkKsF58haF-IxDxcQD2hKN34,1088
 ocf_data_sampler/config/__init__.py,sha256=O29mbH0XG2gIY1g3BaveGCnpBO2SFqdu-qzJ7a6evl0,223
 ocf_data_sampler/config/load.py,sha256=LL-7wemI8o4KPkx35j-wQ3HjsMvDgqXr7G46IcASfnU,632
 ocf_data_sampler/config/model.py,sha256=Jss8UDJAaQIBDr9megX2pERoT0ocFmwLNFC8pCWN6VA,12386
@@ -40,8 +40,8 @@ ocf_data_sampler/select/location.py,sha256=AZvGR8y62opiW7zACGXjoOtBEWRfSLOZIA73O
 ocf_data_sampler/select/select_spatial_slice.py,sha256=Hd4jGRUfIZRoWCirOQZeoLpaUnStB6KyFSTPX69wZLw,8790
 ocf_data_sampler/select/select_time_slice.py,sha256=HeHbwZ0CP03x0-LaJtpbSdtpLufwVTR73p6wH6O_PS8,5513
 ocf_data_sampler/torch_datasets/datasets/__init__.py,sha256=o0SsEXXZ6k9iL__5_RN1Sf60lw_eqK91P3UFEHAD2k0,102
-ocf_data_sampler/torch_datasets/datasets/pvnet_uk.py,sha256=876oLukvb1nLtZQ8HBN3PWfN7urKH2xa45tVar7XrbM,12010
-ocf_data_sampler/torch_datasets/datasets/site.py,sha256=nn6N8daGxllYwCCiFKbCJANTl84NrDRl-nbNGcfXc3U,15429
+ocf_data_sampler/torch_datasets/datasets/pvnet_uk.py,sha256=qbyvTOZZNcGioeH-DDoJmSf_KLRidiuBQRnrvZXD6ts,12046
+ocf_data_sampler/torch_datasets/datasets/site.py,sha256=_FUV_KDe5k7acAmjE9Z2kYgxCFJZrLjziaZssIi1ipg,15465
 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
@@ -56,7 +56,7 @@ ocf_data_sampler/torch_datasets/utils/validation_utils.py,sha256=YqmT-lExWlI8_ul
 scripts/download_gsp_location_data.py,sha256=rRDXMoqX-RYY4jPdxhdlxJGhWdl6r245F5UARgKV6P4,3121
 scripts/refactor_site.py,sha256=skzvsPP0Cn9yTKndzkilyNcGz4DZ88ctvCJ0XrBdc2A,3135
 utils/compute_icon_mean_stddev.py,sha256=a1oWMRMnny39rV-dvu8rcx85sb4bXzPFrR1gkUr4Jpg,2296
-ocf_data_sampler-0.5.7.dist-info/METADATA,sha256=Nu2RLYiLYyU6nkLu8g__Q8EPFIgYMLu5cZLcLXAckXs,12816
-ocf_data_sampler-0.5.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ocf_data_sampler-0.5.7.dist-info/top_level.txt,sha256=LEFU4Uk-PEo72QGLAfnVZIUEm37Q8mKuMeg_Xk-p33g,31
-ocf_data_sampler-0.5.7.dist-info/RECORD,,
+ocf_data_sampler-0.5.8.dist-info/METADATA,sha256=qWknNvSrAAMj-4FydO9gdFkO2z7nZ3Cc2GZPxzjq75Q,12816
+ocf_data_sampler-0.5.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ocf_data_sampler-0.5.8.dist-info/top_level.txt,sha256=LEFU4Uk-PEo72QGLAfnVZIUEm37Q8mKuMeg_Xk-p33g,31
+ocf_data_sampler-0.5.8.dist-info/RECORD,,