ocf-data-sampler 0.0.10__py3-none-any.whl → 0.0.11__py3-none-any.whl

ocf_data_sampler/load/nwp/providers/ecmwf.py

@@ -0,0 +1,37 @@
+"""ECMWF provider loaders"""
+from pathlib import Path
+import xarray as xr
+from ocf_data_sampler.load.nwp.providers.utils import open_zarr_paths
+from ocf_data_sampler.load.utils import check_time_unique_increasing, make_spatial_coords_increasing
+
+def open_ifs(zarr_path: Path | str | list[Path] | list[str]) -> xr.DataArray:
+    """
+    Opens the ECMWF IFS NWP data
+
+    Args:
+        zarr_path: Path to the zarr to open
+
+    Returns:
+        Xarray DataArray of the NWP data
+    """
+    # Open the data
+    ds = open_zarr_paths(zarr_path)
+
+    # Rename
+    ds = ds.rename(
+        {
+            "init_time": "init_time_utc",
+            "variable": "channel",
+        }
+    )
+
+    # Check the timestmps are unique and increasing
+    check_time_unique_increasing(ds.init_time_utc)
+
+    # Make sure the spatial coords are in increasing order
+    ds = make_spatial_coords_increasing(ds, x_coord="longitude", y_coord="latitude")
+
+    ds = ds.transpose("init_time_utc", "step", "channel", "longitude", "latitude")
+    
+    # TODO: should we control the dtype of the DataArray?
+    return ds.ECMWF_UK

ocf_data_sampler/load/nwp/providers/ukv.py

@@ -0,0 +1,45 @@
+"""UKV provider loaders"""
+
+import xarray as xr
+
+from pathlib import Path
+
+from ocf_data_sampler.load.nwp.providers.utils import open_zarr_paths
+from ocf_data_sampler.load.utils import check_time_unique_increasing, make_spatial_coords_increasing
+
+
+def open_ukv(zarr_path: Path | str | list[Path] | list[str]) -> xr.DataArray:
+    """
+    Opens the NWP data
+
+    Args:
+        zarr_path: Path to the zarr to open
+
+    Returns:
+        Xarray DataArray of the NWP data
+    """
+    # Open the data
+    ds = open_zarr_paths(zarr_path)
+
+    # Rename
+    ds = ds.rename(
+        {
+            "init_time": "init_time_utc",
+            "variable": "channel",
+            "x": "x_osgb",
+            "y": "y_osgb",
+        }
+    )
+
+    # Check the timestmps are unique and increasing
+    check_time_unique_increasing(ds.init_time_utc)
+
+    # Make sure the spatial coords are in increasing order
+    ds = make_spatial_coords_increasing(ds, x_coord="x_osgb", y_coord="y_osgb")
+
+    ds = ds.transpose("init_time_utc", "step", "channel", "x_osgb", "y_osgb")
+
+    # TODO: should we control the dtype of the DataArray?
+    return ds.UKV
+
+

ocf_data_sampler/load/nwp/providers/utils.py

@@ -0,0 +1,34 @@
+from pathlib import Path
+import xarray as xr
+
+
+def open_zarr_paths(
+        zarr_path: Path | str | list[Path] | list[str], 
+        time_dim: str = "init_time"
+    ) -> xr.Dataset:
+    """Opens the NWP data
+
+    Args:
+        zarr_path: Path to the zarr(s) to open
+        time_dim: Name of the time dimension
+
+    Returns:
+        The opened Xarray Dataset
+    """
+    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",
+        ).sortby(time_dim)
+    else:
+        ds = xr.open_dataset(
+            zarr_path,
+            engine="zarr",
+            consolidated=True,
+            mode="r",
+            chunks="auto",
+        )
+    return ds

{ocf_data_sampler-0.0.10.dist-info → ocf_data_sampler-0.0.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ocf_data_sampler
-Version: 0.0.10
+Version: 0.0.11
 Summary: Sample from weather data for renewable energy prediction
 Author: James Fulton, Peter Dudfield, and the Open Climate Fix team
 Author-email: info@openclimatefix.org

{ocf_data_sampler-0.0.10.dist-info → ocf_data_sampler-0.0.11.dist-info}/RECORD

@@ -6,6 +6,10 @@ ocf_data_sampler/load/satellite.py,sha256=RcF0HmpV2PKedOdqcTc6dDk4qdQZAdTLYwmMuN
 ocf_data_sampler/load/utils.py,sha256=tkhuhL3YzJucAtaCH572OxBkYvcEDpSed83yg02O8jg,966
 ocf_data_sampler/load/nwp/__init__.py,sha256=SmcrnbygO5xtCKmGR4wtHrj-HI7nOAvnAtfuvRufBGQ,25
 ocf_data_sampler/load/nwp/nwp.py,sha256=O4QnajEZem8BvBgTcYYDBhRhgqPYuJkolHmpMRmrXEA,610
+ocf_data_sampler/load/nwp/providers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ocf_data_sampler/load/nwp/providers/ecmwf.py,sha256=1AFfdhReLjVZIfoA6N1fnKuBS-c8s_6W660A4k5KfA8,1095
+ocf_data_sampler/load/nwp/providers/ukv.py,sha256=IyZ9F9W9exfm08z1ndKsDN6BG2iir5bYf6RI4pOlGb4,1124
+ocf_data_sampler/load/nwp/providers/utils.py,sha256=Sy2exG1wpXLLhMXYdsfR-DZMR3txG1_bBmBdchlc-yA,848
 ocf_data_sampler/numpy_batch/__init__.py,sha256=mrtqwbGik5Zc9MYP5byfCTBm08wMtS2XnTsypC4fPMo,245
 ocf_data_sampler/numpy_batch/gsp.py,sha256=EL0_cJJNyvkQQcOat9vFA61pF4lema3BP_vB4ZS788U,805
 ocf_data_sampler/numpy_batch/nwp.py,sha256=Rv0yfDj902Z2oCwdlRjOs3Kh-F5Fgxjjylh99-lQ9ws,1105
@@ -21,8 +25,8 @@ ocf_data_sampler/torch_datasets/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOC
 ocf_data_sampler/torch_datasets/pvnet_uk_regional.py,sha256=rVKFfoHqSfm4C-eOXiqi5GwBJdMewRMIikvpjEJXi1s,17477
 tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/conftest.py,sha256=OcArgF60paroZQqoP7xExRBF34nEyMuXd7dS7hD6p3w,5393
-ocf_data_sampler-0.0.10.dist-info/LICENSE,sha256=F-Q3UFCR-BECSocV55BFDpn4YKxve9PKrm-lTt6o_Tg,1073
-ocf_data_sampler-0.0.10.dist-info/METADATA,sha256=eXXMZ7ogKXx0j3Krj508liDesXTjUzM2c0NlzOz1P6Q,588
-ocf_data_sampler-0.0.10.dist-info/WHEEL,sha256=HiCZjzuy6Dw0hdX5R3LCFPDmFS4BWl8H-8W39XfmgX4,91
-ocf_data_sampler-0.0.10.dist-info/top_level.txt,sha256=KaQn5qzkJGJP6hKWqsVAc9t0cMLjVvSTk8-kTrW79SA,23
-ocf_data_sampler-0.0.10.dist-info/RECORD,,
+ocf_data_sampler-0.0.11.dist-info/LICENSE,sha256=F-Q3UFCR-BECSocV55BFDpn4YKxve9PKrm-lTt6o_Tg,1073
+ocf_data_sampler-0.0.11.dist-info/METADATA,sha256=UygNspgLmps-YBxZwJd9wWq1ATU9mPtvRZ4b3-Qbr24,588
+ocf_data_sampler-0.0.11.dist-info/WHEEL,sha256=HiCZjzuy6Dw0hdX5R3LCFPDmFS4BWl8H-8W39XfmgX4,91
+ocf_data_sampler-0.0.11.dist-info/top_level.txt,sha256=KaQn5qzkJGJP6hKWqsVAc9t0cMLjVvSTk8-kTrW79SA,23
+ocf_data_sampler-0.0.11.dist-info/RECORD,,