ocf-data-sampler 0.5.6__tar.gz → 0.5.8__tar.gz

{ocf_data_sampler-0.5.6 → ocf_data_sampler-0.5.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ocf-data-sampler
-Version: 0.5.6
+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.6 → ocf_data_sampler-0.5.8}/ocf_data_sampler/load/open_xarray_tensorstore.py

@@ -148,7 +148,7 @@ def open_zarrs(
         dim=concat_dim,
         data_vars="minimal",
         compat="equals",
-        combine_attrs="no_conflicts",
+        combine_attrs="drop_conflicts",
     )
 
     if mask_and_scale:

{ocf_data_sampler-0.5.6 → ocf_data_sampler-0.5.8}/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-0.5.6 → ocf_data_sampler-0.5.8}/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-0.5.8/ocf_data_sampler/utils.py

@@ -0,0 +1,37 @@
+"""Miscellaneous helper functions."""
+
+import pandas as pd
+from xarray_tensorstore import read
+
+
+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()
+    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.6 → ocf_data_sampler-0.5.8}/ocf_data_sampler.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ocf-data-sampler
-Version: 0.5.6
+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.6/ocf_data_sampler/utils.py

@@ -1,21 +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")
-
-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