ocf-data-sampler 0.2.19__tar.gz → 0.2.21__tar.gz

{ocf_data_sampler-0.2.19 → ocf_data_sampler-0.2.21}/PKG-INFO

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

{ocf_data_sampler-0.2.19 → ocf_data_sampler-0.2.21}/ocf_data_sampler/config/model.py

@@ -211,6 +211,7 @@ class NWP(TimeWindowMixin, DropoutMixin, SpatialWindowMixin, NormalisationConsta
         " used to construct an example. If set to None, then the max staleness is set according to"
         " the maximum forecast horizon of the NWP and the requested forecast length.",
     )
+    public: bool = Field(False, description="Whether the NWP data is public or private")
 
     @field_validator("provider")
     def validate_provider(cls, v: str) -> str:

{ocf_data_sampler-0.2.19 → ocf_data_sampler-0.2.21}/ocf_data_sampler/load/load_dataset.py

@@ -38,7 +38,11 @@ def get_dataset_dict(
     if input_config.nwp:
         datasets_dict["nwp"] = {}
         for nwp_source, nwp_config in input_config.nwp.items():
-            da_nwp = open_nwp(nwp_config.zarr_path, provider=nwp_config.provider)
+            da_nwp = open_nwp(
+                zarr_path=nwp_config.zarr_path,
+                provider=nwp_config.provider,
+                public=nwp_config.public,
+            )
 
             da_nwp = da_nwp.sel(channel=list(nwp_config.channels))
 

{ocf_data_sampler-0.2.19 → ocf_data_sampler-0.2.21}/ocf_data_sampler/load/nwp/nwp.py

@@ -9,18 +9,23 @@ from ocf_data_sampler.load.nwp.providers.icon import open_icon_eu
 from ocf_data_sampler.load.nwp.providers.ukv import open_ukv
 
 
-def open_nwp(zarr_path: str | list[str], provider: str) -> xr.DataArray:
+def open_nwp(zarr_path: str | list[str], provider: str, public: bool = False) -> xr.DataArray:
     """Opens NWP zarr.
 
     Args:
         zarr_path: path to the zarr file
         provider: NWP provider
+        public: Whether the data is public or private (only for GFS)
 
     Returns:
         Xarray DataArray of the NWP data
     """
     provider = provider.lower()
 
+    kwargs = {
+        "zarr_path": zarr_path,
+    }
+
     if provider == "ukv":
         _open_nwp = open_ukv
     elif provider == "ecmwf":
@@ -29,9 +34,14 @@ def open_nwp(zarr_path: str | list[str], provider: str) -> xr.DataArray:
         _open_nwp = open_icon_eu
     elif provider == "gfs":
         _open_nwp = open_gfs
+
+        # GFS has a public/private flag
+        if public:
+            kwargs["public"] = True
+
     elif provider == "cloudcasting":
         _open_nwp = open_cloudcasting
     else:
         raise ValueError(f"Unknown provider: {provider}")
 
-    return _open_nwp(zarr_path)
+    return _open_nwp(**kwargs)

{ocf_data_sampler-0.2.19 → ocf_data_sampler-0.2.21}/ocf_data_sampler/load/nwp/providers/cloudcasting.py

@@ -1,7 +1,5 @@
 """Cloudcasting provider loader."""
 
-from pathlib import Path
-
 import xarray as xr
 
 from ocf_data_sampler.load.nwp.providers.utils import open_zarr_paths
@@ -12,14 +10,14 @@ from ocf_data_sampler.load.utils import (
 )
 
 
-def open_cloudcasting(zarr_path: Path | str | list[Path] | list[str]) -> xr.DataArray:
+def open_cloudcasting(zarr_path: str | list[str]) -> xr.DataArray:
     """Opens the satellite predictions from cloudcasting.
 
     Cloudcasting is a OCF forecast product. We forecast future satellite images from recent
     satellite images. More information can be found in the references below.
 
     Args:
-        zarr_path: Path to the zarr to open
+        zarr_path: Path to the zarr(s) to open
 
     Returns:
         Xarray DataArray of the cloudcasting data

{ocf_data_sampler-0.2.19 → ocf_data_sampler-0.2.21}/ocf_data_sampler/load/nwp/providers/ecmwf.py

@@ -14,7 +14,7 @@ def open_ifs(zarr_path: str | list[str]) -> xr.DataArray:
     """Opens the ECMWF IFS NWP data.
 
     Args:
-        zarr_path: Path to the zarr to open
+        zarr_path: Path to the zarr(s) to open
 
     Returns:
         Xarray DataArray of the NWP data

{ocf_data_sampler-0.2.19 → ocf_data_sampler-0.2.21}/ocf_data_sampler/load/nwp/providers/gfs.py

@@ -10,11 +10,12 @@ from ocf_data_sampler.load.utils import check_time_unique_increasing, make_spati
 _log = logging.getLogger(__name__)
 
 
-def open_gfs(zarr_path: str | list[str]) -> xr.DataArray:
+def open_gfs(zarr_path: str | list[str], public: bool = False) -> xr.DataArray:
     """Opens the GFS data.
 
     Args:
-        zarr_path: Path to the zarr to open
+        zarr_path: Path to the zarr(s) to open
+        public: Whether the data is public or private
 
     Returns:
         Xarray DataArray of the NWP data
@@ -22,7 +23,7 @@ def open_gfs(zarr_path: str | list[str]) -> xr.DataArray:
     _log.info("Loading NWP GFS data")
 
     # Open data
-    gfs: xr.Dataset = open_zarr_paths(zarr_path, time_dim="init_time_utc")
+    gfs: xr.Dataset = open_zarr_paths(zarr_path, time_dim="init_time_utc", public=public)
     nwp: xr.DataArray = gfs.to_array()
     nwp = nwp.rename({"variable": "channel"})  # `variable` appears when using `to_array`
 

{ocf_data_sampler-0.2.19 → ocf_data_sampler-0.2.21}/ocf_data_sampler/load/nwp/providers/icon.py

@@ -19,7 +19,7 @@ def remove_isobaric_lelvels_from_coords(nwp: xr.Dataset) -> xr.Dataset:
     return nwp.drop_vars(["isobaricInhPa", *variables_to_drop])
 
 
-def open_icon_eu(zarr_path: str) -> xr.Dataset:
+def open_icon_eu(zarr_path: str | list[str]) -> xr.Dataset:
     """Opens the ICON data.
 
     ICON EU Data is on a regular lat/lon grid
@@ -27,7 +27,7 @@ def open_icon_eu(zarr_path: str) -> xr.Dataset:
     Each of the variables is its own data variable
 
     Args:
-        zarr_path: Path to the zarr to open
+        zarr_path: Path to the zarr(s) to open
 
     Returns:
         Xarray DataArray of the NWP data

{ocf_data_sampler-0.2.19 → ocf_data_sampler-0.2.21}/ocf_data_sampler/load/nwp/providers/ukv.py

@@ -14,7 +14,7 @@ def open_ukv(zarr_path: str | list[str]) -> xr.DataArray:
     """Opens the NWP data.
 
     Args:
-        zarr_path: Path to the zarr to open
+        zarr_path: Path to the zarr(s) to open
 
     Returns:
         Xarray DataArray of the NWP data

{ocf_data_sampler-0.2.19 → ocf_data_sampler-0.2.21}/ocf_data_sampler/load/nwp/providers/utils.py

@@ -3,32 +3,41 @@
 import xarray as xr
 
 
-def open_zarr_paths(zarr_path: str | list[str], time_dim: str = "init_time") -> xr.Dataset:
+def open_zarr_paths(
+    zarr_path: str | list[str], time_dim: str = "init_time", public: bool = False,
+) -> xr.Dataset:
     """Opens the NWP data.
 
     Args:
         zarr_path: Path to the zarr(s) to open
         time_dim: Name of the time dimension
+        public: Whether the data is public or private
 
     Returns:
         The opened Xarray Dataset
     """
+    general_kwargs = {
+        "engine": "zarr",
+        "chunks": "auto",
+        "decode_timedelta": True,
+    }
+
+    if public:
+        # note this only works for s3 zarr paths at the moment
+        general_kwargs["storage_options"] = {"anon": True}
+
     if type(zarr_path) in [list, tuple] or "*" in str(zarr_path):  # Multi-file dataset
         ds = xr.open_mfdataset(
             zarr_path,
-            engine="zarr",
             concat_dim=time_dim,
             combine="nested",
-            chunks="auto",
-            decode_timedelta=True,
+            **general_kwargs,
         ).sortby(time_dim)
     else:
         ds = xr.open_dataset(
             zarr_path,
-            engine="zarr",
             consolidated=True,
             mode="r",
-            chunks="auto",
-            decode_timedelta=True,
+            **general_kwargs,
         )
     return ds

{ocf_data_sampler-0.2.19 → ocf_data_sampler-0.2.21}/ocf_data_sampler/torch_datasets/datasets/pvnet_uk.py

@@ -1,6 +1,5 @@
 """Torch dataset for UK PVNet."""
 
-import numpy as np
 import pandas as pd
 import xarray as xr
 from torch.utils.data import Dataset
@@ -257,22 +256,12 @@ class PVNetUKRegionalDataset(AbstractPVNetUKDataset):
         # Construct a lookup for locations - useful for users to construct sample by GSP ID
         location_lookup = {loc.id: loc for loc in self.locations}
 
-        # Construct indices for sampling
-        t_index, loc_index = np.meshgrid(
-            np.arange(len(self.valid_t0_times)),
-            np.arange(len(self.locations)),
-        )
-
-        # Make array of all possible (t0, location) coordinates. Each row is a single coordinate
-        index_pairs = np.stack((t_index.ravel(), loc_index.ravel())).T
-
         # Assign coords and indices to self
         self.location_lookup = location_lookup
-        self.index_pairs = index_pairs
 
     @override
     def __len__(self) -> int:
-        return len(self.index_pairs)
+        return len(self.locations)*len(self.valid_t0_times)
 
     def _get_sample(self, t0: pd.Timestamp, location: Location) -> NumpySample:
         """Generate the PVNet sample for given coordinates.
@@ -290,7 +279,16 @@ class PVNetUKRegionalDataset(AbstractPVNetUKDataset):
     @override
     def __getitem__(self, idx: int) -> NumpySample:
         # Get the coordinates of the sample
-        t_index, loc_index = self.index_pairs[idx]
+
+        if idx >= len(self):
+            raise ValueError(f"Index {idx} out of range for dataset of length {len(self)}")
+
+        # t_index will be between 0 and len(self.valid_t0_times)-1
+        t_index = idx % len(self.valid_t0_times)
+
+        # For each location, there are len(self.valid_t0_times) possible samples
+        loc_index = idx // len(self.valid_t0_times)
+
         location = self.locations[loc_index]
         t0 = self.valid_t0_times[t_index]
 

{ocf_data_sampler-0.2.19 → ocf_data_sampler-0.2.21}/ocf_data_sampler.egg-info/PKG-INFO

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