ocf-data-sampler 0.5.7__tar.gz → 0.5.9__tar.gz

{ocf_data_sampler-0.5.7 → ocf_data_sampler-0.5.9}/PKG-INFO

@@ -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 → ocf_data_sampler-0.5.9}/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.7 → ocf_data_sampler-0.5.9}/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.9/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.7 → ocf_data_sampler-0.5.9}/ocf_data_sampler.egg-info/PKG-INFO

@@ -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 → ocf_data_sampler-0.5.9}/ocf_data_sampler.egg-info/SOURCES.txt

@@ -62,5 +62,4 @@ ocf_data_sampler/torch_datasets/utils/time_slice_for_dataset.py
 ocf_data_sampler/torch_datasets/utils/valid_time_periods.py
 ocf_data_sampler/torch_datasets/utils/validation_utils.py
 scripts/download_gsp_location_data.py
-scripts/refactor_site.py
-utils/compute_icon_mean_stddev.py
+scripts/refactor_site.py

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

@@ -1,4 +1,3 @@
 dist
 ocf_data_sampler
 scripts
-utils

ocf_data_sampler-0.5.7/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

ocf_data_sampler-0.5.7/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)
-