ocf-data-sampler 0.5.7__py3-none-any.whl → 0.5.9__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.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ocf-data-sampler
-Version: 0.5.7
+Version: 0.5.9
 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.9.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
@@ -55,8 +55,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
-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.9.dist-info/METADATA,sha256=LUgQmrakbDwIEfeP_3IojePDYDdvm15iUtftl5o8Rps,12816
+ocf_data_sampler-0.5.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ocf_data_sampler-0.5.9.dist-info/top_level.txt,sha256=deUxqmsONNAGZDNbsntbXH7BRA1MqWaUeAJrCo6q_xA,25
+ocf_data_sampler-0.5.9.dist-info/RECORD,,

{ocf_data_sampler-0.5.7.dist-info → ocf_data_sampler-0.5.9.dist-info}/top_level.txt

@@ -1,3 +1,2 @@
 ocf_data_sampler
 scripts
-utils

utils/compute_icon_mean_stddev.py

@@ -1,72 +0,0 @@
-"""Script to compute normalisation constants from NWP data."""
-
-import argparse
-import glob
-import logging
-
-import numpy as np
-import xarray as xr
-
-from ocf_data_sampler.load.nwp.providers.icon import open_icon_eu
-
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# Add argument parser
-parser = argparse.ArgumentParser(description="Compute normalization constants from NWP data")
-parser.add_argument("--data-path", type=str, required=True,
-                    help='Path pattern to zarr files (e.g., "/path/to/data/*.zarr.zip")')
-parser.add_argument("--n-samples", type=int, default=2000,
-                    help="Number of random samples to use (default: 2000)")
-
-args = parser.parse_args()
-
-zarr_files = glob.glob(args.data_path)
-n_samples = args.n_samples
-
-ds = open_icon_eu(zarr_files)
-
-n_init_times = ds.sizes["init_time_utc"]
-n_lats = ds.sizes["latitude"]
-n_longs = ds.sizes["longitude"]
-n_steps = ds.sizes["step"]
-
-random_init_times = np.random.choice(n_init_times, size=n_samples, replace=True)
-random_lats = np.random.choice(n_lats, size=n_samples, replace=True)
-random_longs = np.random.choice(n_longs, size=n_samples, replace=True)
-random_steps = np.random.choice(n_steps, size=n_samples, replace=True)
-
-samples = []
-for i in range(n_samples):
-    sample = ds.isel(init_time_utc=random_init_times[i],
-                    latitude=random_lats[i],
-                    longitude=random_longs[i],
-                    step=random_steps[i])
-    samples.append(sample)
-
-samples_stack = xr.concat(samples, dim="samples")
-
-
-available_channels = samples_stack.channel.values.tolist()
-logger.info("Available channels: %s", available_channels)
-
-ICON_EU_MEAN = {}
-ICON_EU_STD = {}
-
-for var in available_channels:
-    if var not in available_channels:
-        logger.warning("Variable '%s' not found in the channel coordinate; skipping.", var)
-        continue
-    var_data = samples_stack.sel(channel=var)
-    var_mean = float(var_data.mean().compute())
-    var_std = float(var_data.std().compute())
-
-    ICON_EU_MEAN[var] = var_mean
-    ICON_EU_STD[var] = var_std
-
-    logger.info("Processed %s: mean=%.4f, std=%.4f", var, var_mean, var_std)
-
-logger.info("\nMean values:\n%s", ICON_EU_MEAN)
-logger.info("\nStandard deviations:\n%s", ICON_EU_STD)
-