ocf-data-sampler 0.2.19__py3-none-any.whl → 0.2.21__py3-none-any.whl

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/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/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/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/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/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/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/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/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/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.dist-info → ocf_data_sampler-0.2.21.dist-info}/METADATA

@@ -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.dist-info → ocf_data_sampler-0.2.21.dist-info}/RECORD

@@ -2,25 +2,25 @@ ocf_data_sampler/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,
 ocf_data_sampler/utils.py,sha256=DjuneGGisl08ENvPZV_lrcX4b2NCKJC1ZpXgIpxuQi4,290
 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=SyjtlSK6gzQHWUfgX3VNKYLODyiKuD0Mu4hlm9GoHeg,10427
+ocf_data_sampler/config/model.py,sha256=iqffLs_VDqw9jOTLWchVFK4c6FWxHAzCngSfkjLUyCY,10516
 ocf_data_sampler/config/save.py,sha256=m8SPw5rXjkMm1rByjh3pK5StdBi4e8ysnn3jQopdRaI,1064
 ocf_data_sampler/data/uk_gsp_locations_20220314.csv,sha256=RSh7DRh55E3n8lVAaWXGTaXXHevZZtI58td4d4DhGos,10415772
 ocf_data_sampler/data/uk_gsp_locations_20250109.csv,sha256=XZISFatnbpO9j8LwaxNKFzQSjs6hcHFsV8a9uDDpy2E,9055334
 ocf_data_sampler/load/__init__.py,sha256=-vQP9g0UOWdVbjEGyVX_ipa7R1btmiETIKAf6aw4d78,201
 ocf_data_sampler/load/gsp.py,sha256=UfPxwHw2Dw2xYSO5Al28oTamgnEM_n_4bYXsqGwY5Tc,1884
-ocf_data_sampler/load/load_dataset.py,sha256=sIi0nkijR_-1fRfW5JcXNTR0ccGbpkhxb7JX_zjJ-W4,1956
+ocf_data_sampler/load/load_dataset.py,sha256=wSXPUQKgGRM6HC-yBXQ2IcDBQDckOSllmbGnhqikFMQ,2055
 ocf_data_sampler/load/satellite.py,sha256=E7Ln7Y60Qr1RTV-_R71YoxXQM-Ca7Y1faIo3oKB2eFk,2292
 ocf_data_sampler/load/site.py,sha256=zOzlWk6pYZBB5daqG8URGksmDXWKrkutUvN8uALAIh8,1468
 ocf_data_sampler/load/utils.py,sha256=sZ0-zzconcLkVQwAkCYrqKDo98Hrh5ChdiQJv5Bh91g,2040
 ocf_data_sampler/load/nwp/__init__.py,sha256=SmcrnbygO5xtCKmGR4wtHrj-HI7nOAvnAtfuvRufBGQ,25
-ocf_data_sampler/load/nwp/nwp.py,sha256=OUK6thsKfEcqGUj4WoxDiyqGcVwwoujePdmQRis8H8I,1076
+ocf_data_sampler/load/nwp/nwp.py,sha256=PNNYYREEGQT4sxGilNzuthKKOmVJdQL8R2r8bvzyEr0,1317
 ocf_data_sampler/load/nwp/providers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ocf_data_sampler/load/nwp/providers/cloudcasting.py,sha256=mGsEqo5nVGQhfIK1KC-ywLsaFKU18-UFKqQVtB8gpnw,1564
-ocf_data_sampler/load/nwp/providers/ecmwf.py,sha256=5AzktPJgertCx6oij6aePRosPuZHGFznMxTgtkk_mgc,994
-ocf_data_sampler/load/nwp/providers/gfs.py,sha256=gTXfBmjMK8NnbTz_PyPvznfcy0eNHFGGa7eQlDarErc,994
-ocf_data_sampler/load/nwp/providers/icon.py,sha256=q2ZJIKYptXFvoHzb9y3iNKSJZ0PNtnH54h2gwYbgIJ8,1569
-ocf_data_sampler/load/nwp/providers/ukv.py,sha256=-0v8JCLH8ypz8GMXZ6Rrx-I0LoHuHO8sXFupbC1RpM0,1013
-ocf_data_sampler/load/nwp/providers/utils.py,sha256=Rau6j6it2D4YLavoey-bdOlixgn5gDQioToAmRPAvtA,905
+ocf_data_sampler/load/nwp/providers/cloudcasting.py,sha256=fozXpB3a2rNqQgnpRDC7xunxffh7Wwmc0kkCiYmDVJ4,1521
+ocf_data_sampler/load/nwp/providers/ecmwf.py,sha256=an-gXsZwkPQvRXeza1U_4MNU5yEnVm0_8tn03rxTudI,997
+ocf_data_sampler/load/nwp/providers/gfs.py,sha256=glBbo2kXtcTjQv_VNqA32lsdCCGB114Ovm-cibRWxTA,1088
+ocf_data_sampler/load/nwp/providers/icon.py,sha256=6MkOfUk5dmv0XJZLrKMy1e8xipj2fHCTkYXuff7MgUY,1584
+ocf_data_sampler/load/nwp/providers/ukv.py,sha256=Ka1KFZcJYPwr5vuxo-xWGVQC0pudheqGBonUnbyJCMg,1016
+ocf_data_sampler/load/nwp/providers/utils.py,sha256=NrzE3JAtoc6oEywJHxPUdi_I4UJgJ_l5dxLZ4DLKvcg,1124
 ocf_data_sampler/numpy_sample/__init__.py,sha256=nY5C6CcuxiWZ_jrXRzWtN7WyKXhJImSiVTIG6Rz4B_4,401
 ocf_data_sampler/numpy_sample/collate.py,sha256=hoxIc5SoHoIs3Nx37aRZzWChpswjy9lHUgaKgHIoo80,2039
 ocf_data_sampler/numpy_sample/common_types.py,sha256=9CjYHkUTx0ObduWh43fhsybZCTXvexql7qC2ptMDoek,377
@@ -39,7 +39,7 @@ ocf_data_sampler/select/location.py,sha256=AZvGR8y62opiW7zACGXjoOtBEWRfSLOZIA73O
 ocf_data_sampler/select/select_spatial_slice.py,sha256=liAqIa-Amj58pOqx5r16i99HURj9oQ41j7gnPgRDQP4,8201
 ocf_data_sampler/select/select_time_slice.py,sha256=HeHbwZ0CP03x0-LaJtpbSdtpLufwVTR73p6wH6O_PS8,5513
 ocf_data_sampler/torch_datasets/datasets/__init__.py,sha256=jfJSFcR0eO1AqeH7S3KnGjsBqVZT5w3oyi784PUR6Q0,146
-ocf_data_sampler/torch_datasets/datasets/pvnet_uk.py,sha256=ZV2FoMPxFU2aPTWipj9HhJhGfrEg9MYOJRNR8aFcmvs,12613
+ocf_data_sampler/torch_datasets/datasets/pvnet_uk.py,sha256=cd4IyzYu8rMFgLHRXqYpnOIAZe4Yl21YdLmDQw45F7o,12545
 ocf_data_sampler/torch_datasets/datasets/site.py,sha256=nRUlhXQQGVrTuBmE1QnwXAUsPTXz0dsezlQjwK71jIQ,17641
 ocf_data_sampler/torch_datasets/sample/__init__.py,sha256=GL84vdZl_SjHDGVyh9Uekx2XhPYuZ0dnO3l6f6KXnHI,100
 ocf_data_sampler/torch_datasets/sample/base.py,sha256=cQ1oIyhdmlotejZK8B3Cw6MNvpdnBPD8G_o2h7Ye4Vc,2206
@@ -55,7 +55,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.2.19.dist-info/METADATA,sha256=PODEt_bufdiQ_WJnSKXcZ39N_8Gylbpj9CN2_Yl5PUw,11581
-ocf_data_sampler-0.2.19.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
-ocf_data_sampler-0.2.19.dist-info/top_level.txt,sha256=LEFU4Uk-PEo72QGLAfnVZIUEm37Q8mKuMeg_Xk-p33g,31
-ocf_data_sampler-0.2.19.dist-info/RECORD,,
+ocf_data_sampler-0.2.21.dist-info/METADATA,sha256=PKoq-iYwiK9VGZLHgAQCQJm-U97WMkb7HToAsrPNzlw,11581
+ocf_data_sampler-0.2.21.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+ocf_data_sampler-0.2.21.dist-info/top_level.txt,sha256=LEFU4Uk-PEo72QGLAfnVZIUEm37Q8mKuMeg_Xk-p33g,31
+ocf_data_sampler-0.2.21.dist-info/RECORD,,

{ocf_data_sampler-0.2.19.dist-info → ocf_data_sampler-0.2.21.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.3.1)
+Generator: setuptools (80.7.1)
 Root-Is-Purelib: true
 Tag: py3-none-any