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

ocf_data_sampler/load/gsp.py

@@ -29,5 +29,3 @@ def open_gsp(zarr_path: str | Path) -> xr.DataArray:
     )
 
     return ds.generation_mw
-
-

ocf_data_sampler/load/nwp/providers/ecmwf.py

@@ -2,7 +2,13 @@
 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:
     """
@@ -25,7 +31,7 @@ def open_ifs(zarr_path: Path | str | list[Path] | list[str]) -> xr.DataArray:
         }
     )
 
-    # Check the timestmps are unique and increasing
+    # Check the timestamps are unique and increasing
     check_time_unique_increasing(ds.init_time_utc)
 
     # Make sure the spatial coords are in increasing order
@@ -34,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:
@@ -31,7 +35,7 @@ def open_ukv(zarr_path: Path | str | list[Path] | list[str]) -> xr.DataArray:
         }
     )
 
-    # Check the timestmps are unique and increasing
+    # Check the timestamps are unique and increasing
     check_time_unique_increasing(ds.init_time_utc)
 
     # Make sure the spatial coords are in increasing order
@@ -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

@@ -3,12 +3,15 @@
 import subprocess
 from pathlib import Path
 
-import pandas as pd
 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.DataArray:
+def _get_single_sat_data(zarr_path: Path | str) -> xr.Dataset:
     """Helper function to open a zarr from either local or GCP path.
 
     The local or GCP path may contain wildcard matching (*)
@@ -89,7 +92,7 @@ def open_sat_data(zarr_path: Path | str | list[Path] | list[str]) -> xr.DataArra
         }
     )
 
-    # Check the timestmps are unique and increasing
+    # Check the timestamps are unique and increasing
     check_time_unique_increasing(ds.time_utc)
 
     # Make sure the spatial coords are in increasing order
@@ -98,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

@@ -7,6 +7,7 @@ def check_time_unique_increasing(datetimes) -> None:
     assert time.is_unique
     assert time.is_monotonic_increasing
 
+
 def make_spatial_coords_increasing(ds: xr.Dataset, x_coord: str, y_coord: str) -> xr.Dataset:
     """Make sure the spatial coordinates are in increasing order
     
@@ -26,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/select/dropout.py

@@ -27,7 +27,7 @@ def draw_dropout_time(
 
 
 def apply_dropout_time(
-        ds: xr.Dataset,
+        ds: xr.DataArray,
         dropout_time: pd.Timestamp | None,
     ):
 

ocf_data_sampler/select/select_time_slice.py

@@ -37,7 +37,7 @@ def _sel_fillinterp(
 
 
 def select_time_slice(
-    ds: xr.Dataset | xr.DataArray,
+    ds: xr.DataArray,
     t0: pd.Timestamp,
     sample_period_duration: pd.Timedelta,
     history_duration: pd.Timedelta | None = None,
@@ -75,7 +75,7 @@ def select_time_slice(
 
 
 def select_time_slice_nwp(
-    ds: xr.Dataset | xr.DataArray,
+    ds: xr.DataArray,
     t0: pd.Timestamp,
     sample_period_duration: pd.Timedelta,
     history_duration: pd.Timedelta,

{ocf_data_sampler-0.0.14.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.14
+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.14.dist-info → ocf_data_sampler-0.0.16.dist-info}/RECORD

@@ -1,14 +1,14 @@
 ocf_data_sampler/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 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=ewuypEah4dfj_gIcYwZoLHiA2ZjUj7KPeyBYKwT1qdo,761
-ocf_data_sampler/load/satellite.py,sha256=RcF0HmpV2PKedOdqcTc6dDk4qdQZAdTLYwmMuNpSTZQ,2906
-ocf_data_sampler/load/utils.py,sha256=tkhuhL3YzJucAtaCH572OxBkYvcEDpSed83yg02O8jg,966
+ocf_data_sampler/load/gsp.py,sha256=Gcr1JVUOPKhFRDCSHtfPDjxx0BtyyEhXrZvGEKLPJ5I,759
+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=1AFfdhReLjVZIfoA6N1fnKuBS-c8s_6W660A4k5KfA8,1095
-ocf_data_sampler/load/nwp/providers/ukv.py,sha256=IyZ9F9W9exfm08z1ndKsDN6BG2iir5bYf6RI4pOlGb4,1124
+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
@@ -16,17 +16,17 @@ ocf_data_sampler/numpy_batch/nwp.py,sha256=Rv0yfDj902Z2oCwdlRjOs3Kh-F5Fgxjjylh99
 ocf_data_sampler/numpy_batch/satellite.py,sha256=e6eoNmiiHtzZbDVtBolFzDuE3qwhHN6bL9H86emAUsk,732
 ocf_data_sampler/numpy_batch/sun_position.py,sha256=UW6-WtjrKdCkcguolHUDSLhYFfarknQzzjlCX8YdEOM,1700
 ocf_data_sampler/select/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-ocf_data_sampler/select/dropout.py,sha256=JYbjG5e8d48te7xj4I9pTWk43d6ksjGeyKFLSTuAOlY,1062
+ocf_data_sampler/select/dropout.py,sha256=3aDqlL3U9kzzIkIHV44h5_ZbfOUSg1iChHKingXk7-s,1064
 ocf_data_sampler/select/fill_time_periods.py,sha256=iTtMjIPFYG5xtUYYedAFBLjTWWUa7t7WQ0-yksWf0-E,440
 ocf_data_sampler/select/find_contiguous_time_periods.py,sha256=6ioB8LeFpFNBMgKDxrgG3zqzNjkBF_jlV9yye2ZYT2E,11925
 ocf_data_sampler/select/select_spatial_slice.py,sha256=7BSzOFPMSBWpBWXSajWTfI8luUVsSgh4zN-rkr-AuUs,11470
-ocf_data_sampler/select/select_time_slice.py,sha256=XuksC9N03c5rV9OeWtxjGuoGyeJJGy4JMJe3w7m6oaw,6654
+ocf_data_sampler/select/select_time_slice.py,sha256=41cch1fQr59fZgv7UHsNGc3OvoynrixT3bmr3_1d7cU,6628
 ocf_data_sampler/torch_datasets/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 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.14.dist-info/LICENSE,sha256=F-Q3UFCR-BECSocV55BFDpn4YKxve9PKrm-lTt6o_Tg,1073
-ocf_data_sampler-0.0.14.dist-info/METADATA,sha256=3aN9lKWnmbNdjsF-J69AKAPwvc2WwUDkb0Nnyorr92c,588
-ocf_data_sampler-0.0.14.dist-info/WHEEL,sha256=UvcQYKBHoFqaQd6LKyqHw9fxEolWLQnlzP0h_LgJAfI,91
-ocf_data_sampler-0.0.14.dist-info/top_level.txt,sha256=KaQn5qzkJGJP6hKWqsVAc9t0cMLjVvSTk8-kTrW79SA,23
-ocf_data_sampler-0.0.14.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,,