chore: Add test for checking physical limits and zeroes in NWP data #…

ocf_datapipes/load/nwp/nwp.py

@@ -5,8 +5,9 @@
 from typing import Union
 
 import dask
+import dask.array
+import numpy as np
 import xarray as xr
-from constants import NWP_LIMITS
 from ocf_blosc2 import Blosc2  # noqa: F401
 from torch.utils.data import IterDataPipe, functional_datapipe
 
@@ -17,6 +18,8 @@
 from ocf_datapipes.load.nwp.providers.merra2 import open_merra2
 from ocf_datapipes.load.nwp.providers.ukv import open_ukv
 
+from .constants import NWP_LIMITS
+
 logger = logging.getLogger(__name__)
 
 
@@ -30,6 +33,7 @@ def __init__(
         provider: str = "ukv",
         check_for_zeros: bool = False,
         check_physical_limits: bool = False,
+        check_for_nans: bool = False,
     ):
         """
         Opens NWP Zarr and yields it
@@ -39,10 +43,12 @@ def __init__(
             provider: NWP provider
             check_for_zeros: Check for zeros in the NWP data
             check_physical_limits: Check the physical limits of nwp data (e.g. -100<temperature<100)
+            check_for_nans: Check for NaNs in the NWP data
         """
         self.zarr_path = zarr_path
         self.check_for_zeros = check_for_zeros
         self.check_physical_limits = check_physical_limits
+        self.check_for_nans = check_for_nans
         self.limits = NWP_LIMITS
 
         logger.info(f"Using {provider.lower()}")
@@ -71,6 +77,8 @@ def __iter__(self) -> Union[xr.DataArray, xr.Dataset]:  # type: ignore
             self.check_if_zeros(nwp)
         if self.check_physical_limits:
             self.check_if_physical_limits(nwp)
+        if self.check_for_nans:
+            self.check_if_nans(nwp)
         while True:
             yield nwp
 
@@ -124,3 +132,20 @@ def check_if_physical_limits(self, nwp: Union[xr.DataArray, xr.Dataset]):
                         raise ValueError(
                             f"NWP data {var_name} is outside physical limits: ({lower},{upper})"
                         )
+    def check_if_nans(self, nwp: Union[xr.DataArray, xr.Dataset]):
+        """Checks if the NWP data contains NaNs"""
+        if isinstance(nwp, xr.DataArray):
+            if dask.is_dask_collection(nwp.data):
+                if dask.array.isnan(nwp.data).any().compute():
+                    raise ValueError("NWP data contains NaNs")
+            else:
+                if np.isnan(nwp.data).any():
+                    raise ValueError("NWP DataArray contains NaNs")
+        elif isinstance(nwp, xr.Dataset):
+            for var in nwp.data_vars:
+                if dask.is_dask_collection(nwp[var].data):
+                    if dask.array.isnan(nwp[var].data).any().compute():
+                        raise ValueError(f"NWP Dataset variable{var} contains NaNs")
+                else:
+                    if np.isnan(nwp[var].data).any():
+                        raise ValueError(f"NWP Dataset variable{var} contains NaNs")

tests/batch/test_merge_numpy_examples_to_batch.py

@@ -40,49 +40,13 @@ def _single_batch_sample(fill_value):
     return sample
 
 
-def _single_batch_sample_nan(fill_value):
-    """This function allows us to create batches with different filled values"""
-
-    sample: NumpyBatch = {}
-    sample[BatchKey.satellite_actual] = np.full(
-        (12, 10, 24, 24), fill_value, dtype=np.float32
-    )  # shape: (time, channel, x, y)
-    sample[BatchKey.gsp_id] = np.full((1,), fill_value)  # shape: (1,)
-    sample[BatchKey.gsp_t0_idx] = 4  # scalar and constant across all samples
-
-    sample_nwp_ukv: NWPNumpyBatch = {}
-    sample_nwp_ukv[NWPBatchKey.nwp] = np.full(
-        (8, 2, 24, 24), fill_value, dtype=np.float32
-    )  # shape: (time, variable, x, y)
-    sample_nwp_ukv[NWPBatchKey.nwp][0, 0, 0, 0] = np.nan
-
-    sample_nwp_ukv[NWPBatchKey.nwp_channel_names] = ["a", "b"]  # shape: (variable,)
-
-    sample_nwp_ecmwf: NWPNumpyBatch = {}
-    sample_nwp_ecmwf[NWPBatchKey.nwp] = np.full(
-        (8, 4, 12, 12), fill_value
-    )  # shape: (time, variable, x, y)
-
-    sample[BatchKey.nwp] = {
-        "ukv": sample_nwp_ukv,
-        "ecmwf": sample_nwp_ecmwf,
-    }
-    # print(sample[BatchKey.nwp]["ukv"])
-    return sample
-
 
 @pytest.fixture
 def numpy_sample_datapipe():
     dp = IterableWrapper([_single_batch_sample(i) for i in range(8)])
     return dp
 
 
-@pytest.fixture
-def numpy_nan_sample_datapipe():
-    dp = IterableWrapper([_single_batch_sample_nan(i) for i in range(8)])
-    return dp
-
-
 def test_merge_numpy_batch(numpy_sample_datapipe):
     dp = MergeNumpyBatchIterDataPipe(numpy_sample_datapipe.batch(4))
     dp_iter = iter(dp)
@@ -99,13 +63,6 @@ def test_merge_numpy_batch(numpy_sample_datapipe):
         assert nwp_batch[NWPBatchKey.nwp_channel_names] == ["a", "b"]
 
 
-def test_merge_numpy_batch_for_nans(numpy_nan_sample_datapipe):
-    with pytest.raises(
-        ValueError
-    ):  # checks for Error raised if NWP/BatchKey DataArray contains Nans
-        dp = MergeNumpyBatchIterDataPipe(numpy_nan_sample_datapipe.batch(4))
-        dp_iter = iter(dp)
-        metadata = next(dp_iter)
 
 
 def test_merge_numpy_examples_to_batch(numpy_sample_datapipe):

tests/load/nwp/test_load_nwp.py

@@ -111,7 +111,7 @@ def test_check_for_zeros():
     # to generate data with zeros and limits:
     original_store_path = "tests/data/nwp_data/test.zarr"
     original_store = zarr.open(original_store_path, mode="r")
-    new_store_path = "tests/data/nwp_data/test_with_zeros_n_limits.zarr"
+    new_store_path = "tests/data/nwp_data/test_with_zeros_n_limits_n_nans.zarr"
     # Optionally, clear the destination store if it already exists
     shutil.rmtree(new_store_path, ignore_errors=True)
     with zarr.open(new_store_path, mode="w") as new_store:
@@ -120,9 +120,11 @@ def test_check_for_zeros():
 
         new_store["UKV"][0, 0, 0, 0] = 0
         new_store["UKV"][0, 0, 0, 1] = np.random.uniform(190, 360, size=(548,))
+        new_store["UKV"][0, 0, 0, 2] = np.nan
+
     shutil.copy(
         "tests/data/nwp_data/test.zarr/.zmetadata",
-        "tests/data/nwp_data/test_with_zeros_n_limits.zarr/.zmetadata",
+        "tests/data/nwp_data/test_with_zeros_n_limits_n_nans.zarr/.zmetadata",
     )
 
     # positive test case
@@ -142,7 +144,7 @@ def test_check_for_zeros():
 def test_check_physical_limits():
     # positive test case
     nwp_datapipe1 = OpenNWP(
-        zarr_path="tests/data/nwp_data/test_with_zeros_n_limits.zarr", check_physical_limits=True
+        zarr_path="tests/data/nwp_data/test_with_zeros_n_limits_n_nans.zarr", check_physical_limits=True
     )
     with pytest.raises(
         ValueError
@@ -154,6 +156,19 @@ def test_check_physical_limits():
     metadata = next(iter(nwp_datapipe2))
     assert metadata is not None
 
+def test_check_if_nans():
+    # positive test case
+    nwp_datapipe1 = OpenNWP(
+        zarr_path="tests/data/nwp_data/test_with_zeros_n_limits_n_nans.zarr", check_for_nans=True
+    )
+    with pytest.raises(ValueError):  # checks for Error raised if NWP DataArray contains nans
+        metadata = next(iter(nwp_datapipe1))
+
+    # negative test case
+    nwp_datapipe2 = OpenNWP(zarr_path="tests/data/nwp_data/test.zarr", check_for_nans=True)
+    metadata = next(iter(nwp_datapipe2))
+    assert metadata is not None
+
     shutil.rmtree(
-        "tests/data/nwp_data/test_with_zeros_n_limits.zarr"
-    )  # removes the zarr file created for testing
+        "tests/data/nwp_data/test_with_zeros_n_limits_n_nans.zarr"
+    )  # removes the zarr file created for testing