ocf-data-sampler 0.0.36__py3-none-any.whl → 0.0.37__py3-none-any.whl

ocf_data_sampler/numpy_batch/collate.py

@@ -0,0 +1,79 @@
+from ocf_data_sampler.numpy_batch import NWPBatchKey
+
+import numpy as np
+import logging
+from typing import Union
+
+logger = logging.getLogger(__name__)
+
+
+def stack_np_examples_into_batch(dict_list):
+    """
+    Stacks Numpy examples into a batch
+
+    See also: `unstack_np_batch_into_examples()` for opposite
+
+    Args:
+        dict_list: A list of dict-like Numpy examples to stack
+
+    Returns:
+        The stacked NumpyBatch object
+    """
+    batch = {}
+
+    batch_keys = list(dict_list[0].keys())
+
+    for batch_key in batch_keys:
+        # NWP is nested so treat separately
+        if batch_key == "nwp":
+            nwp_batch: dict[str, NWPBatchKey] = {}
+
+            # Unpack source keys
+            nwp_sources = list(dict_list[0]["nwp"].keys())
+
+            for nwp_source in nwp_sources:
+                # Keys can be different for different NWPs
+                nwp_batch_keys = list(dict_list[0]["nwp"][nwp_source].keys())
+
+                nwp_source_batch = {}
+                for nwp_batch_key in nwp_batch_keys:
+                    nwp_source_batch[nwp_batch_key] = stack_data_list(
+                        [d["nwp"][nwp_source][nwp_batch_key] for d in dict_list],
+                        nwp_batch_key,
+                    )
+
+                nwp_batch[nwp_source] = nwp_source_batch
+
+            batch["nwp"] = nwp_batch
+
+        else:
+            batch[batch_key] = stack_data_list(
+                [d[batch_key] for d in dict_list],
+                batch_key,
+            )
+
+    return batch
+
+
+def _key_is_constant(batch_key):
+    is_constant = batch_key.endswith("t0_idx") or batch_key == NWPBatchKey.channel_names
+    return is_constant
+
+
+def stack_data_list(
+    data_list: list,
+    batch_key: Union[str, NWPBatchKey],
+):
+    """How to combine data entries for each key
+    """
+    if _key_is_constant(batch_key):
+        # These are always the same for all examples.
+        return data_list[0]
+    try:
+        return np.stack(data_list)
+    except Exception as e:
+        logger.debug(f"Could not stack the following shapes together, ({batch_key})")
+        shapes = [example.shape for example in data_list]
+        logger.debug(shapes)
+        logger.error(e)
+        raise e

{ocf_data_sampler-0.0.36.dist-info → ocf_data_sampler-0.0.37.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ocf_data_sampler
-Version: 0.0.36
+Version: 0.0.37
 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
@@ -56,7 +56,7 @@ Requires-Dist: mkdocs-material>=8.0; extra == "docs"
 # ocf-data-sampler
  
 <!-- ALL-CONTRIBUTORS-BADGE:START - Do not remove or modify this section -->
-[![All Contributors](https://img.shields.io/badge/all_contributors-6-orange.svg?style=flat-square)](#contributors-)
+[![All Contributors](https://img.shields.io/badge/all_contributors-7-orange.svg?style=flat-square)](#contributors-)
 <!-- ALL-CONTRIBUTORS-BADGE:END -->
 
 [![tags badge](https://img.shields.io/github/v/tag/openclimatefix/ocf-data-sampler?include_prereleases&sort=semver&color=FFAC5F)](https://github.com/openclimatefix/ocf-data-sampler/tags)
@@ -129,6 +129,7 @@ Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/d
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/peterdudfield"><img src="https://avatars.githubusercontent.com/u/34686298?v=4?s=100" width="100px;" alt="Peter Dudfield"/><br /><sub><b>Peter Dudfield</b></sub></a><br /><a href="https://github.com/openclimatefix/ocf-data-sampler/commits?author=peterdudfield" title="Code">💻</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/VikramsDataScience"><img src="https://avatars.githubusercontent.com/u/45002417?v=4?s=100" width="100px;" alt="Vikram Pande"/><br /><sub><b>Vikram Pande</b></sub></a><br /><a href="https://github.com/openclimatefix/ocf-data-sampler/commits?author=VikramsDataScience" title="Code">💻</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/SophiaLi20"><img src="https://avatars.githubusercontent.com/u/163532536?v=4?s=100" width="100px;" alt="Unnati Bhardwaj"/><br /><sub><b>Unnati Bhardwaj</b></sub></a><br /><a href="https://github.com/openclimatefix/ocf-data-sampler/commits?author=SophiaLi20" title="Documentation">📖</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/alirashidAR"><img src="https://avatars.githubusercontent.com/u/110668489?v=4?s=100" width="100px;" alt="Ali Rashid"/><br /><sub><b>Ali Rashid</b></sub></a><br /><a href="https://github.com/openclimatefix/ocf-data-sampler/commits?author=alirashidAR" title="Code">💻</a></td>
     </tr>
   </tbody>
 </table>

{ocf_data_sampler-0.0.36.dist-info → ocf_data_sampler-0.0.37.dist-info}/RECORD

@@ -19,6 +19,7 @@ ocf_data_sampler/load/nwp/providers/ecmwf.py,sha256=2iR1Iy542lo51rC6XFLV-3pbUE68
 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=8MgRF29rK9bKP4b4iHakaoGwBKUcjWZ-VFKjCcq53QA,336
+ocf_data_sampler/numpy_batch/collate.py,sha256=KyWdDi8AXD5YiokXXiqr2_X1SC1me1GrhnQMelg0Qx8,2202
 ocf_data_sampler/numpy_batch/gsp.py,sha256=QjQ25JmtufvdiSsxUkBTPhxouYGWPnnWze8pXr_aBno,960
 ocf_data_sampler/numpy_batch/nwp.py,sha256=dAehfRo5DL2Yb20ifHHl5cU1QOrm3ZOpQmN39fSUOw8,1255
 ocf_data_sampler/numpy_batch/satellite.py,sha256=3NoE_ElzMHwO60apqJeFAwI6J7eIxD0OWTyAVl-uJi8,903
@@ -41,12 +42,13 @@ ocf_data_sampler/torch_datasets/site.py,sha256=NYuhgm9ti9SRt1dcb_WrFYYo14NgVdOsa
 ocf_data_sampler/torch_datasets/valid_time_periods.py,sha256=Qo65qUHtle_bW5tLTYr7empHTRv-lpjvfx_6GNJj3Xg,4371
 scripts/refactor_site.py,sha256=asZ27hQ4IyXgCCUaFJqcz1ObBNcV2W3ywqHBpSXA_fc,1728
 tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tests/conftest.py,sha256=N-_XgXpWeTRhkwP_NVh2mBORt2LKkM4mbkm-O62RN5I,7363
+tests/conftest.py,sha256=68hH-HPdHPLvLrtYJU8bjfkdGKbhPfNveLKvUs6_Lr0,7970
 tests/config/test_config.py,sha256=eaye_F7-el4tTP4n2vRME8qlV0b2jaKUX4HhgOUpa7E,5203
 tests/load/test_load_gsp.py,sha256=aT_nqaSXmUTcdHzuTT7AmXJr3R31k4OEN-Fv3eLxlQE,424
 tests/load/test_load_nwp.py,sha256=3qyyDkB1q9t3tyAwogfotNrxqUOpXXimco1CImoEWGg,753
 tests/load/test_load_satellite.py,sha256=STX5AqqmOAgUgE9R1xyq_sM3P1b8NKdGjO-hDhayfxM,524
 tests/load/test_load_sites.py,sha256=T9lSEnGPI8FQISudVYHHNTHeplNS62Vrx48jaZ6J_Jo,364
+tests/numpy_batch/test_collate.py,sha256=U0u5LLpkImr7R50xDuTKNo2Of8sUT5pEs9F2ZMW3jEU,836
 tests/numpy_batch/test_gsp.py,sha256=VANXV32K8aLX4dCdhCUnDorJmyNN-Bjc7Wc1N-RzWEk,548
 tests/numpy_batch/test_nwp.py,sha256=Fnj7cR-VR2Z0kMu8SrgnIayjxWnPWrYFjWSjMmnrh4Y,1445
 tests/numpy_batch/test_satellite.py,sha256=8a4ZwMLpsOmYKmwI1oW_su_hwkCNYMEJAEfa0dbsx1k,1179
@@ -57,10 +59,10 @@ tests/select/test_find_contiguous_time_periods.py,sha256=kOga_V7er5We7ewMARXaKdM
 tests/select/test_location.py,sha256=_WZk2FPYeJ-nIfCJS6Sp_yaVEEo7m31DmMFoZzgyCts,2712
 tests/select/test_select_spatial_slice.py,sha256=7EX9b6g-pMdACQx3yefjs5do2s-Rho2UmKevV4oglsU,5147
 tests/select/test_select_time_slice.py,sha256=K1EJR5TwZa9dJf_YTEHxGtvs398iy1xS2lr1BgJZkoo,9603
-tests/torch_datasets/test_pvnet_uk_regional.py,sha256=ZNyrisyhM1vw4q2qcHDuvX2uRi-v3U8Y8lOYx7cd8yM,2635
+tests/torch_datasets/test_pvnet_uk_regional.py,sha256=eqy0nQOWoHnqltlJlGmRlgIiIzPEwOC6o5A6GARryKA,2118
 tests/torch_datasets/test_site.py,sha256=YuVjWTI14_kmEOx23XE5J_RZ8UalCKD2xRv6mqYizB8,2872
-ocf_data_sampler-0.0.36.dist-info/LICENSE,sha256=F-Q3UFCR-BECSocV55BFDpn4YKxve9PKrm-lTt6o_Tg,1073
-ocf_data_sampler-0.0.36.dist-info/METADATA,sha256=t2N8IyEMdDI8Zf9YXModYFV6zmYPddy1y2wg9nEw_Cc,9933
-ocf_data_sampler-0.0.36.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-ocf_data_sampler-0.0.36.dist-info/top_level.txt,sha256=Faob6N6cFdPc5eUpCTYcXgCaNhi4XLLteUL5W5ayYmg,31
-ocf_data_sampler-0.0.36.dist-info/RECORD,,
+ocf_data_sampler-0.0.37.dist-info/LICENSE,sha256=F-Q3UFCR-BECSocV55BFDpn4YKxve9PKrm-lTt6o_Tg,1073
+ocf_data_sampler-0.0.37.dist-info/METADATA,sha256=tKixIA37U0AA76QsYmCIfLzpzE2aSGRmquSx69jX4aY,10290
+ocf_data_sampler-0.0.37.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+ocf_data_sampler-0.0.37.dist-info/top_level.txt,sha256=Faob6N6cFdPc5eUpCTYcXgCaNhi4XLLteUL5W5ayYmg,31
+ocf_data_sampler-0.0.37.dist-info/RECORD,,

tests/conftest.py

@@ -7,6 +7,7 @@ import xarray as xr
 import tempfile
 
 from ocf_data_sampler.config.model import Site
+from ocf_data_sampler.config import load_yaml_configuration, save_yaml_configuration
 
 _top_test_directory = os.path.dirname(os.path.realpath(__file__))
 
@@ -269,3 +270,18 @@ def uk_gsp_zarr_path(ds_uk_gsp):
         ds_uk_gsp.to_zarr(filename)
         yield filename
 
+
+@pytest.fixture()
+def pvnet_config_filename(
+    tmp_path, config_filename, nwp_ukv_zarr_path, uk_gsp_zarr_path, sat_zarr_path
+):
+
+    # adjust config to point to the zarr file
+    config = load_yaml_configuration(config_filename)
+    config.input_data.nwp["ukv"].zarr_path = nwp_ukv_zarr_path
+    config.input_data.satellite.zarr_path = sat_zarr_path
+    config.input_data.gsp.zarr_path = uk_gsp_zarr_path
+
+    filename = f"{tmp_path}/configuration.yaml"
+    save_yaml_configuration(config, filename)
+    return filename

tests/numpy_batch/test_collate.py

@@ -0,0 +1,26 @@
+from ocf_data_sampler.numpy_batch import GSPBatchKey, SatelliteBatchKey
+from ocf_data_sampler.numpy_batch.collate import stack_np_examples_into_batch
+from ocf_data_sampler.torch_datasets import PVNetUKRegionalDataset
+
+
+def test_pvnet(pvnet_config_filename):
+
+    # Create dataset object
+    dataset = PVNetUKRegionalDataset(pvnet_config_filename)
+
+    assert len(dataset.locations) == 317
+    assert len(dataset.valid_t0_times) == 39
+    assert len(dataset) == 317 * 39
+
+    # Generate 2 samples
+    sample1 = dataset[0]
+    sample2 = dataset[1]
+
+    batch = stack_np_examples_into_batch([sample1, sample2])
+
+    assert isinstance(batch, dict)
+    assert "nwp" in batch
+    assert isinstance(batch["nwp"], dict)
+    assert "ukv" in batch["nwp"]
+    assert GSPBatchKey.gsp in batch
+    assert SatelliteBatchKey.satellite_actual in batch

tests/torch_datasets/test_pvnet_uk_regional.py

@@ -6,19 +6,6 @@ from ocf_data_sampler.config import load_yaml_configuration, save_yaml_configura
 from ocf_data_sampler.numpy_batch import NWPBatchKey, GSPBatchKey, SatelliteBatchKey
 
 
-@pytest.fixture()
-def pvnet_config_filename(tmp_path, config_filename, nwp_ukv_zarr_path, uk_gsp_zarr_path, sat_zarr_path):
-
-    # adjust config to point to the zarr file
-    config = load_yaml_configuration(config_filename)
-    config.input_data.nwp['ukv'].zarr_path = nwp_ukv_zarr_path
-    config.input_data.satellite.zarr_path = sat_zarr_path
-    config.input_data.gsp.zarr_path = uk_gsp_zarr_path
-
-    filename = f"{tmp_path}/configuration.yaml"
-    save_yaml_configuration(config, filename)
-    return filename
-
 
 def test_pvnet(pvnet_config_filename):