ocf-data-sampler 0.0.48__py3-none-any.whl → 0.0.49__py3-none-any.whl

ocf_data_sampler/torch_datasets/datasets/site.py

@@ -241,29 +241,30 @@ class SitesDataset(Dataset):
 
         # add datetime features
         datetimes = pd.DatetimeIndex(combined_sample_dataset.site__time_utc.values)
-        datetime_features = make_datetime_numpy_dict(datetimes=datetimes, key_prefix="site")
-        datetime_features_xr = xr.Dataset(datetime_features, coords={"site__time_utc": datetimes})
-        combined_sample_dataset = xr.merge([combined_sample_dataset, datetime_features_xr])
+        datetime_features = make_datetime_numpy_dict(datetimes=datetimes, key_prefix="site_")
+        combined_sample_dataset = combined_sample_dataset.assign_coords(
+            {k: ("site__time_utc", v) for k, v in datetime_features.items()}
+        )
 
         # add sun features
         sun_position_features = make_sun_position_numpy_sample(
             datetimes=datetimes,
             lon=combined_sample_dataset.site__longitude.values,
             lat=combined_sample_dataset.site__latitude.values,
-            key_prefix="site",
+            key_prefix="site_",
         )
-        sun_position_features_xr = xr.Dataset(
-            sun_position_features, coords={"site__time_utc": datetimes}
+        combined_sample_dataset = combined_sample_dataset.assign_coords(
+            {k: ("site__time_utc", v) for k, v in sun_position_features.items()}
         )
-        combined_sample_dataset = xr.merge([combined_sample_dataset, sun_position_features_xr])
 
         # TODO include t0_index in xr dataset? 
 
         # Fill any nan values
         return combined_sample_dataset.fillna(0.0)
 
-
-    def merge_data_arrays(self, normalised_data_arrays: list[Tuple[str, xr.DataArray]]) -> xr.Dataset:
+    def merge_data_arrays(
+        self, normalised_data_arrays: list[Tuple[str, xr.DataArray]]
+    ) -> xr.Dataset:
         """
         Combine a list of DataArrays into a single Dataset with unique naming conventions.
 

{ocf_data_sampler-0.0.48.dist-info → ocf_data_sampler-0.0.49.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: ocf_data_sampler
-Version: 0.0.48
+Version: 0.0.49
 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.48.dist-info → ocf_data_sampler-0.0.49.dist-info}/RECORD

@@ -38,7 +38,7 @@ ocf_data_sampler/select/spatial_slice_for_dataset.py,sha256=3tRrMBXr7s4CnClbVSIq
 ocf_data_sampler/select/time_slice_for_dataset.py,sha256=BFjNwWAzhcb1hpqx7UPi5RF9WWt15owbZp1WB-uGA6Q,4305
 ocf_data_sampler/torch_datasets/datasets/__init__.py,sha256=nJUa2KzVa84ZoM0PT2AbDz26ennmAYc7M7WJVfypPMs,85
 ocf_data_sampler/torch_datasets/datasets/pvnet_uk_regional.py,sha256=xxeX4Js9LQpydehi3BS7k9psqkYGzgJuM17uTYux40M,8742
-ocf_data_sampler/torch_datasets/datasets/site.py,sha256=75M0oDstOLyLZBySVIS6fLJSbEjfxcWBlgGP_ewui7s,15334
+ocf_data_sampler/torch_datasets/datasets/site.py,sha256=v7plMF_WJPkfwnJAUFf_8gXAy8SXE5Og_fgZMEm4c20,15257
 ocf_data_sampler/torch_datasets/utils/merge_and_fill_utils.py,sha256=hIbekql64eXsNDFIoEc--GWxwdVWrh2qKegdOi70Bow,874
 ocf_data_sampler/torch_datasets/utils/valid_time_periods.py,sha256=Qo65qUHtle_bW5tLTYr7empHTRv-lpjvfx_6GNJj3Xg,4371
 scripts/refactor_site.py,sha256=asZ27hQ4IyXgCCUaFJqcz1ObBNcV2W3ywqHBpSXA_fc,1728
@@ -65,9 +65,9 @@ tests/select/test_select_time_slice.py,sha256=K1EJR5TwZa9dJf_YTEHxGtvs398iy1xS2l
 tests/torch_datasets/conftest.py,sha256=eRCzHE7cxS4AoskExkCGFDBeqItktAYNAdkfpMoFCeE,629
 tests/torch_datasets/test_merge_and_fill_utils.py,sha256=ueA0A7gZaWEgNdsU8p3CnKuvSnlleTUjEhSw2HUUROM,1229
 tests/torch_datasets/test_pvnet_uk_regional.py,sha256=FCiFueeFqrsXe7gWguSjBz5ZeUrvyhGbGw81gaVvkHM,5087
-tests/torch_datasets/test_site.py,sha256=0tnjgx6z4VlzjoF_V2p3Y2t2Z1d0o_07Vwb-FH_c3tU,4640
-ocf_data_sampler-0.0.48.dist-info/LICENSE,sha256=F-Q3UFCR-BECSocV55BFDpn4YKxve9PKrm-lTt6o_Tg,1073
-ocf_data_sampler-0.0.48.dist-info/METADATA,sha256=Wl37iFa0xFXGRsHNEHLn7WYeXbWvbg8SGgLXcksExQA,11788
-ocf_data_sampler-0.0.48.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-ocf_data_sampler-0.0.48.dist-info/top_level.txt,sha256=Faob6N6cFdPc5eUpCTYcXgCaNhi4XLLteUL5W5ayYmg,31
-ocf_data_sampler-0.0.48.dist-info/RECORD,,
+tests/torch_datasets/test_site.py,sha256=0gT_7k086BBnxqbvOayiUeI-vzJsYXlx3KvACC0c6lk,6114
+ocf_data_sampler-0.0.49.dist-info/LICENSE,sha256=F-Q3UFCR-BECSocV55BFDpn4YKxve9PKrm-lTt6o_Tg,1073
+ocf_data_sampler-0.0.49.dist-info/METADATA,sha256=GuLd3IDZ7qU9W9wwV84AQ5tN8rlouhF4ZpDThHsVUKo,11788
+ocf_data_sampler-0.0.49.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+ocf_data_sampler-0.0.49.dist-info/top_level.txt,sha256=Faob6N6cFdPc5eUpCTYcXgCaNhi4XLLteUL5W5ayYmg,31
+ocf_data_sampler-0.0.49.dist-info/RECORD,,

tests/torch_datasets/test_site.py

@@ -3,6 +3,8 @@ import numpy as np
 from ocf_data_sampler.torch_datasets.datasets.site import SitesDataset, convert_from_dataset_to_dict_datasets
 from xarray import Dataset, DataArray
 
+from torch.utils.data import DataLoader
+
 
 def test_site(site_config_filename):
 
@@ -18,17 +20,45 @@ def test_site(site_config_filename):
     assert isinstance(sample, Dataset)
 
     # Expected dimensions and data variables
-    expected_dims = {'satellite__x_geostationary', 'site__time_utc', 'nwp-ukv__target_time_utc',
-                     'nwp-ukv__x_osgb', 'satellite__channel', 'satellite__y_geostationary',
-                     'satellite__time_utc', 'nwp-ukv__channel', 'nwp-ukv__y_osgb', 'site_solar_azimuth',
-                     'site_solar_elevation', 'site_date_cos', 'site_time_cos', 'site_time_sin', 'site_date_sin'}
+    expected_dims = {
+        "satellite__x_geostationary",
+        "site__time_utc",
+        "nwp-ukv__target_time_utc",
+        "nwp-ukv__x_osgb",
+        "satellite__channel",
+        "satellite__y_geostationary",
+        "satellite__time_utc",
+        "nwp-ukv__channel",
+        "nwp-ukv__y_osgb",
+    }
+
+    expected_coords_subset = {
+        "site__solar_azimuth",
+        "site__solar_elevation",
+        "site__date_cos",
+        "site__time_cos",
+        "site__time_sin",
+        "site__date_sin",
+    }
 
     expected_data_vars = {"nwp-ukv", "satellite", "site"}
 
+    import xarray as xr
+
+    sample.to_netcdf("sample.nc")
+    sample = xr.open_dataset("sample.nc")
+
     # Check dimensions
-    assert set(sample.dims) == expected_dims, f"Missing or extra dimensions: {set(sample.dims) ^ expected_dims}"
+    assert (
+        set(sample.dims) == expected_dims
+    ), f"Missing or extra dimensions: {set(sample.dims) ^ expected_dims}"
     # Check data variables
-    assert set(sample.data_vars) == expected_data_vars, f"Missing or extra data variables: {set(sample.data_vars) ^ expected_data_vars}"
+    assert (
+        set(sample.data_vars) == expected_data_vars
+    ), f"Missing or extra data variables: {set(sample.data_vars) ^ expected_data_vars}"
+
+    for coords in expected_coords_subset:
+        assert coords in sample.coords
 
     # check the shape of the data is correct
     # 30 minutes of 5 minute data (inclusive), one channel, 2x2 pixels
@@ -38,6 +68,7 @@ def test_site(site_config_filename):
     # 1.5 hours of 30 minute data (inclusive)
     assert sample["site"].values.shape == (4,)
 
+
 def test_site_time_filter_start(site_config_filename):
 
     # Create dataset object
@@ -74,11 +105,51 @@ def test_convert_from_dataset_to_dict_datasets(site_config_filename):
 
     assert isinstance(sample, dict)
 
-    print(sample.keys())
-
     for key in ["nwp", "satellite", "site"]:
         assert key in sample
 
+
+def test_site_dataset_with_dataloader(site_config_filename):
+    # Create dataset object
+    dataset = SitesDataset(site_config_filename)
+
+    expected_coods = {
+        "site__solar_azimuth",
+        "site__solar_elevation",
+        "site__date_cos",
+        "site__time_cos",
+        "site__time_sin",
+        "site__date_sin",
+    }
+
+    sample = dataset[0]
+    for key in expected_coods:
+        assert key in sample
+
+    dataloader_kwargs = dict(
+        shuffle=False,
+        batch_size=None,
+        sampler=None,
+        batch_sampler=None,
+        num_workers=1,
+        collate_fn=None,
+        pin_memory=False,  # Only using CPU to prepare samples so pinning is not beneficial
+        drop_last=False,
+        timeout=0,
+        worker_init_fn=None,
+        prefetch_factor=1,
+        persistent_workers=False,  # Not needed since we only enter the dataloader loop once
+    )
+
+    dataloader = DataLoader(dataset, collate_fn=None, batch_size=None)
+
+    for i, sample in zip(range(1), dataloader):
+
+        # check that expected_dims is in the sample
+        for key in expected_coods:
+            assert key in sample
+
+
 def test_process_and_combine_site_sample_dict(site_config_filename):
     # Load config
     # config = load_yaml_configuration(pvnet_config_filename)