Merge pull request #312 from openclimatefix/NWP-patch

Fix hardcode NWP time resolution and merra2 loader bug

Author: AUdaltsova

ocf_datapipes/load/README.md

@@ -4,3 +4,49 @@ This module contains the code for loading the data off disk.
 This data is always opened into Xarray `DataSet` or `DataArray` objects for further processing.
 The opened data should only have dimensions renamed to the common format, and minimal processing done.
 Any real processing should be performed in the `transform` module.
+
+## NWP
+
+The NWP data is loaded into an ```IterDataPipe``` in `ocf_datapipes/load/nwp/nwp.py` using a provider from `ocf_datapipes/load/nwp/providers`. Providers open the data file and transform the data into a `DataAray` of a standardised shape the datapipes use; normally this means having the following 5 dimensions: ```init_time_utc, step, channel, latitude, longitude```.
+
+Example of loaded ECMWF data:
+
+```
+<xarray.DataArray 'ECMWF_BLAH' (init_time_utc: 1, step: 2, channel: 2,
+                                latitude: 221, longitude: 221)> Size: 781kB
+dask.array<transpose, shape=(1, 2, 2, 221, 221), dtype=float32, chunksize=(1, 2, 2, 221, 221), chunktype=numpy.ndarray>
+Coordinates:
+  * init_time_utc  (init_time_utc) datetime64[ns] 8B 2023-09-25T12:00:00
+  * latitude       (latitude) float64 2kB 31.0 30.95 30.9 ... 20.1 20.05 20.0
+  * longitude      (longitude) float64 2kB 68.0 68.05 68.1 ... 78.9 78.95 79.0
+  * step           (step) timedelta64[ns] 16B 00:00:00 01:00:00
+  * channel        (channel) <U5 40B 'dlwrf' 'dswrf'
+Attributes:
+    Conventions:             CF-1.7
+    GRIB_centre:             ecmf
+    GRIB_centreDescription:  European Centre for Medium-Range Weather Forecasts
+    GRIB_subCentre:          0
+    institution:             European Centre for Medium-Range Weather Forecasts
+```
+
+There are exceptions, e.g. ICON Global uses an isohedral grid, so it is differently organised and does not have ```latitude``` and ```longitude``` dimensions.
+
+### Adding an NWP provider
+
+1. Add a [provider].py file to `ocf_datapipes/load/nwp/providers` that uses `open_zarr_paths` from `ocf_datapipes.load.nwp.providers.utils` to load the file(s) and returns your data in the right shape, where the dimensions contain:
+   - `init_time_utc`: when the data was initialised
+   - `step`: distance from datapoint to its init_time
+   - `channel`: list of variables
+   - `latitude`: latitude
+   - `longitude`: longitude
+
+Add sanity checks to ensure time is unique and monotonic
+
+2. Add your provider as an option to the `IterDataPipe` in `ocf_datapipes/load/nwp/nwp.py`
+
+3. Add test data to `ocf_datapipes/tests/data` and create a test in `ocf_datapipes/tests/load/nwp/test_load_nwp.py`
+Current tests include:
+    - checking the loaded data is not None
+    - checking all expected dimensions are present
+    - (for some data) checking returns a xr.DataArray
+4. Calculate the mean and std of your data and add to `ocf_datapipes/utils/consts.py`

ocf_datapipes/load/nwp/providers/ecmwf.py

@@ -22,7 +22,7 @@ def open_ifs(zarr_path) -> xr.DataArray:
         raise Exception("Too many TLDVs")
     else:
         dataVar = dataVars[0]
-    ifs: xr.Dataset = nwp[dataVar]
+    ifs: xr.DataArray = nwp[dataVar]
     del nwp
     ifs = ifs.transpose("init_time", "step", "variable", "latitude", "longitude")
     ifs = ifs.rename(

ocf_datapipes/load/nwp/providers/excarta.py

@@ -26,7 +26,7 @@ def preprocess_excarta(ds: xr.Dataset) -> xr.Dataset:
     return ds
 
 
-def open_excarta(zarr_path) -> xr.Dataset:
+def open_excarta(zarr_path) -> xr.DataArray:
     """
     Opens the Excarta hindcast data
 

ocf_datapipes/load/nwp/providers/merra2.py

@@ -26,9 +26,11 @@ def open_merra2(zarr_path) -> xr.DataArray:
         {"channel": list(nwp.keys())}
     )
     nwp = nwp.transpose("init_time_utc", "step", "channel", "latitude", "longitude")
+    aodana: xr.DataArray = nwp["AODANA"]
+    del nwp
 
     # Sanity checks.
-    time = pd.DatetimeIndex(nwp.step + nwp.init_time_utc.values)
+    time = pd.DatetimeIndex(aodana.step + aodana.init_time_utc.values)
     assert time.is_unique
     assert time.is_monotonic_increasing
-    return nwp
+    return aodana

ocf_datapipes/training/common.py

@@ -160,7 +160,7 @@ def open_and_return_datapipes(
                 )
                 .filter_channels(nwp_conf.nwp_channels)
                 .add_t0_idx_and_sample_period_duration(
-                    sample_period_duration=timedelta(hours=1),
+                    sample_period_duration=minutes(nwp_conf.time_resolution_minutes),
                     history_duration=minutes(nwp_conf.history_minutes),
                 )
             )
@@ -731,7 +731,7 @@ def slice_datapipes_by_time(
 
             datapipes_dict["nwp"][nwp_key] = dp.select_time_slice_nwp(
                 t0_datapipe=get_t0_datapipe(f"nwp/{nwp_key}"),
-                sample_period_duration=minutes(60),
+                sample_period_duration=minutes(conf_in.nwp[nwp_key].time_resolution_minutes),
                 history_duration=minutes(conf_in.nwp[nwp_key].history_minutes),
                 forecast_duration=minutes(conf_in.nwp[nwp_key].forecast_minutes),
                 dropout_timedeltas=dropout_timedeltas,

tests/load/nwp/test_load_nwp.py

@@ -1,4 +1,5 @@
 import pandas as pd
+from xarray import DataArray
 
 from ocf_datapipes.load import OpenNWP
 
@@ -49,6 +50,7 @@ def test_load_ecmwf():
     )
     metadata = next(iter(nwp_datapipe))
     assert metadata is not None
+    assert type(next(enumerate(metadata))[1]) == DataArray
     dim_keys = set(["channel", "init_time_utc", "latitude", "longitude", "step"])
     if bool(dim_keys - set(metadata.dims)):
         raise ValueError(
@@ -63,6 +65,7 @@ def test_load_merra2():
     )
     metadata = next(iter(nwp_datapipe))
     assert metadata is not None
+    assert type(next(enumerate(metadata))[1]) == DataArray
     dim_keys = set(["channel", "init_time_utc", "latitude", "longitude", "step"])
     if bool(dim_keys - set(metadata.dims)):
         raise ValueError(