ocf-data-sampler 0.2.18__tar.gz → 0.2.20__tar.gz

{ocf_data_sampler-0.2.18 → ocf_data_sampler-0.2.20}/PKG-INFO

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

{ocf_data_sampler-0.2.18 → ocf_data_sampler-0.2.20}/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.18 → ocf_data_sampler-0.2.20}/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.18 → ocf_data_sampler-0.2.20}/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.18 → ocf_data_sampler-0.2.20}/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.18 → ocf_data_sampler-0.2.20}/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.18 → ocf_data_sampler-0.2.20}/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,12 +23,12 @@ 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`
 
     del gfs
 
-    nwp = nwp.rename({"variable": "channel","init_time": "init_time_utc"})
     check_time_unique_increasing(nwp.init_time_utc)
     nwp = make_spatial_coords_increasing(nwp, x_coord="longitude", y_coord="latitude")
 

{ocf_data_sampler-0.2.18 → ocf_data_sampler-0.2.20}/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.18 → ocf_data_sampler-0.2.20}/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.18 → ocf_data_sampler-0.2.20}/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.18 → ocf_data_sampler-0.2.20}/ocf_data_sampler.egg-info/PKG-INFO

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