add if statement

ocf_datapipes/training/windnet.py

@@ -128,12 +128,15 @@ def __iter__(self):
             for filename in self.filenames:
                 dataset = xr.open_dataset(filename)
                 datasets = uncombine_from_single_dataset(dataset)
-                # print(datasets)
-                datasets["nwp"]["ecmwf"] = potentially_coarsen(datasets["nwp"]["ecmwf"])
-                # Select the specific keys desired
-                datasets["nwp"]["ecmwf"] = datasets["nwp"]["ecmwf"].sel(
-                    channel=["u10", "u100", "u200", "v10", "v100", "v200"]
-                )
+
+                if "ecmwf" in datasets["nwp"]:
+                    datasets["nwp"]["ecmwf"] = potentially_coarsen(datasets["nwp"]["ecmwf"])
+
+                    # Select the specific keys desired
+                    datasets["nwp"]["ecmwf"] = datasets["nwp"]["ecmwf"].sel(
+                        channel=["u10", "u100", "u200", "v10", "v100", "v200"]
+                    )
+
                 # Yield a dictionary of the data, using the keys in self.keys
                 # print(datasets)
                 dataset_dict = {}

tests/batch/test_utils.py

@@ -2,29 +2,24 @@
 import torch
 
 from ocf_datapipes.batch import BatchKey, NumpyBatch, TensorBatch
-from ocf_datapipes.batch import copy_batch_to_device, batch_to_tensor
+from ocf_datapipes.batch.utils import copy_batch_to_device, batch_to_tensor
+import pytest
 
 
-def _create_test_batch() -> NumpyBatch:
+@pytest.fixture()
+def sample_numpy_batch():
     sample: NumpyBatch = {}
     sample[BatchKey.satellite_actual] = np.full((12, 10, 24, 24), 0)
     return sample
 
 
-def test_batch_to_tensor() -> None:
-    batch: NumpyBatch = _create_test_batch()
-    tensor_batch = batch_to_tensor(batch)
+def test_batch_to_tensor(sample_numpy_batch):
+    tensor_batch = batch_to_tensor(sample_numpy_batch)
     assert isinstance(tensor_batch[BatchKey.satellite_actual], torch.Tensor)
 
 
-def test_copy_batch_to_device() -> None:
-    batch = _create_test_batch()
-    tensor_batch = batch_to_tensor(batch)
+def test_copy_batch_to_device(sample_numpy_batch):
+    tensor_batch = batch_to_tensor(sample_numpy_batch)
     device = torch.device("cpu")
     batch_copy: TensorBatch = copy_batch_to_device(tensor_batch, device)
     assert batch_copy[BatchKey.satellite_actual].device == device  # type: ignore
-
-
-if __name__ == "__main__":
-    test_batch_to_tensor()
-    test_copy_batch_to_device()