ocf-data-sampler 0.0.15__py3-none-any.whl → 0.0.16__py3-none-any.whl

ocf_data_sampler/load/nwp/providers/ecmwf.py

@@ -2,7 +2,12 @@
 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
+from ocf_data_sampler.load.utils import (
+    check_time_unique_increasing,
+    make_spatial_coords_increasing,
+    get_xr_data_array_from_xr_dataset
+)
+
 
 
 def open_ifs(zarr_path: Path | str | list[Path] | list[str]) -> xr.DataArray:
@@ -35,4 +40,4 @@ def open_ifs(zarr_path: Path | str | list[Path] | list[str]) -> xr.DataArray:
     ds = ds.transpose("init_time_utc", "step", "channel", "longitude", "latitude")
     
     # TODO: should we control the dtype of the DataArray?
-    return ds.ECMWF_UK
+    return get_xr_data_array_from_xr_dataset(ds)

ocf_data_sampler/load/nwp/providers/ukv.py

@@ -5,7 +5,11 @@ 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
+from ocf_data_sampler.load.utils import (
+    check_time_unique_increasing,
+    make_spatial_coords_increasing,
+    get_xr_data_array_from_xr_dataset
+)
 
 
 def open_ukv(zarr_path: Path | str | list[Path] | list[str]) -> xr.DataArray:
@@ -40,6 +44,6 @@ def open_ukv(zarr_path: Path | str | list[Path] | list[str]) -> xr.DataArray:
     ds = ds.transpose("init_time_utc", "step", "channel", "x_osgb", "y_osgb")
 
     # TODO: should we control the dtype of the DataArray?
-    return ds.UKV
+    return get_xr_data_array_from_xr_dataset(ds)
 
 

ocf_data_sampler/load/satellite.py

@@ -4,7 +4,11 @@ import subprocess
 from pathlib import Path
 
 import xarray as xr
-from ocf_data_sampler.load.utils import check_time_unique_increasing, make_spatial_coords_increasing
+from ocf_data_sampler.load.utils import (
+    check_time_unique_increasing,
+    make_spatial_coords_increasing,
+    get_xr_data_array_from_xr_dataset
+)
 
 
 def _get_single_sat_data(zarr_path: Path | str) -> xr.Dataset:
@@ -97,4 +101,4 @@ def open_sat_data(zarr_path: Path | str | list[Path] | list[str]) -> xr.DataArra
     ds = ds.transpose("time_utc", "channel", "x_geostationary", "y_geostationary")
 
     # TODO: should we control the dtype of the DataArray?
-    return ds.data
+    return get_xr_data_array_from_xr_dataset(ds)

ocf_data_sampler/load/utils.py

@@ -27,4 +27,17 @@ def make_spatial_coords_increasing(ds: xr.Dataset, x_coord: str, y_coord: str) -
     assert (ds[x_coord].diff(dim=x_coord) > 0).all()
     assert (ds[y_coord].diff(dim=y_coord) > 0).all()
 
-    return ds
+    return ds
+
+
+def get_xr_data_array_from_xr_dataset(ds: xr.Dataset) -> xr.DataArray:
+    """Return underlying xr.DataArray from passed xr.Dataset. 
+    Checks only one variable is present and returns it as an xr.DataArray.
+
+    Args:
+        ds: xr.Dataset to extract xr.DataArray from
+    """
+
+    datavars = list(ds.var())
+    assert len(datavars) == 1, "Cannot open as xr.DataArray: dataset contains multiple variables"
+    return ds[datavars[0]]

{ocf_data_sampler-0.0.15.dist-info → ocf_data_sampler-0.0.16.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ocf_data_sampler
-Version: 0.0.15
+Version: 0.0.16
 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.15.dist-info → ocf_data_sampler-0.0.16.dist-info}/RECORD

@@ -2,13 +2,13 @@ ocf_data_sampler/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,
 ocf_data_sampler/data/uk_gsp_locations.csv,sha256=RSh7DRh55E3n8lVAaWXGTaXXHevZZtI58td4d4DhGos,10415772
 ocf_data_sampler/load/__init__.py,sha256=MjgfxilTzyz1RYFoBEeAXmE9hyjknLvdmlHPmlAoiQY,44
 ocf_data_sampler/load/gsp.py,sha256=Gcr1JVUOPKhFRDCSHtfPDjxx0BtyyEhXrZvGEKLPJ5I,759
-ocf_data_sampler/load/satellite.py,sha256=SkcL5ZSiAs5iOU71Ds12GQd8re1N5vEg8IDKuoUFoOU,2885
-ocf_data_sampler/load/utils.py,sha256=kQM9BHwKG--soehlnJOe8VqEI9WEoky-L4aep8ANNOM,967
+ocf_data_sampler/load/satellite.py,sha256=3KlA1fx4SwxdzM-jC1WRaONXO0D6m0WxORnEnwUnZrA,2967
+ocf_data_sampler/load/utils.py,sha256=EQGvVWlGMoSOdbDYuMfVAa0v6wmAOPmHIAemdrTB5v4,1406
 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=hLfq23DJI_NmPdtjELU4uuw2zoDeIUHVY6HYpxo776o,1097
-ocf_data_sampler/load/nwp/providers/ukv.py,sha256=y9DiQ4hK618MnYmskXZcC-ImYcWM7y5vA9gFeuseEUY,1125
+ocf_data_sampler/load/nwp/providers/ecmwf.py,sha256=vW-p3vCyQ-CofKo555-gE7VDi5hlpjtjTLfHqWF0HEE,1175
+ocf_data_sampler/load/nwp/providers/ukv.py,sha256=79Bm7q-K_GJPYMy62SUIZbRWRF4-tIaB1dYPEgLD9vo,1207
 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
@@ -25,8 +25,8 @@ ocf_data_sampler/torch_datasets/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOC
 ocf_data_sampler/torch_datasets/pvnet_uk_regional.py,sha256=haljV4FAZI4-Qf-65nq-JIJOIQNhR6YRncjTBWMYkY4,18502
 tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/conftest.py,sha256=OcArgF60paroZQqoP7xExRBF34nEyMuXd7dS7hD6p3w,5393
-ocf_data_sampler-0.0.15.dist-info/LICENSE,sha256=F-Q3UFCR-BECSocV55BFDpn4YKxve9PKrm-lTt6o_Tg,1073
-ocf_data_sampler-0.0.15.dist-info/METADATA,sha256=LaNhiRYMOobKgxuQvARTv5hzlswCaINSJtxchuT-lgY,588
-ocf_data_sampler-0.0.15.dist-info/WHEEL,sha256=UvcQYKBHoFqaQd6LKyqHw9fxEolWLQnlzP0h_LgJAfI,91
-ocf_data_sampler-0.0.15.dist-info/top_level.txt,sha256=KaQn5qzkJGJP6hKWqsVAc9t0cMLjVvSTk8-kTrW79SA,23
-ocf_data_sampler-0.0.15.dist-info/RECORD,,
+ocf_data_sampler-0.0.16.dist-info/LICENSE,sha256=F-Q3UFCR-BECSocV55BFDpn4YKxve9PKrm-lTt6o_Tg,1073
+ocf_data_sampler-0.0.16.dist-info/METADATA,sha256=lV0PmJljbd63TmlSErEOXN8zSn5031q4WHItd5mJGs4,588
+ocf_data_sampler-0.0.16.dist-info/WHEEL,sha256=UvcQYKBHoFqaQd6LKyqHw9fxEolWLQnlzP0h_LgJAfI,91
+ocf_data_sampler-0.0.16.dist-info/top_level.txt,sha256=KaQn5qzkJGJP6hKWqsVAc9t0cMLjVvSTk8-kTrW79SA,23
+ocf_data_sampler-0.0.16.dist-info/RECORD,,