Merge pull request #284 from openclimatefix/jacob/downsampling

Add Downsampling and new Mean/Stddev for India

Author: jacobbieker

ocf_datapipes/training/common.py

@@ -1215,3 +1215,11 @@ def create_t0_and_loc_datapipes(
     location_pipe, t0_datapipe = t0_loc_datapipe.unzip(sequence_length=2)
 
     return location_pipe, t0_datapipe
+
+
+def potentially_coarsen(xr_data: xr.Dataset):
+    """Coarsen the data, if it is separated by 0.05 degrees each"""
+    if "latitude" in xr_data.coords and "longitude" in xr_data.coords:
+        if xr_data.latitude.values[1] - xr_data.latitude.values[0] == 0.05:
+            xr_data = xr_data.coarsen(latitude=2, longitude=2)
+    return xr_data

ocf_datapipes/training/pvnet_site.py

@@ -16,6 +16,7 @@
     fill_nans_in_arrays,
     fill_nans_in_pv,
     normalize_gsp,
+    potentially_coarsen,
     slice_datapipes_by_time,
 )
 from ocf_datapipes.utils.consts import (
@@ -231,9 +232,16 @@ def construct_sliced_data_pipeline(
                 roi_height_pixels=conf_nwp[nwp_key].nwp_image_size_pixels_height,
                 roi_width_pixels=conf_nwp[nwp_key].nwp_image_size_pixels_width,
             )
+            # Coarsen the data, if it is separated by 0.05 degrees each
+            nwp_datapipe = nwp_datapipe.map(potentially_coarsen)
+            # Somewhat hacky way for India specifically, need different mean/std for ECMWF data
+            if conf_nwp[nwp_key].nwp_provider in ["ecmwf"]:
+                normalize_provider = "ecmwf_india"
+            else:
+                normalize_provider = conf_nwp[nwp_key].nwp_provider
             nwp_datapipes_dict[nwp_key] = nwp_datapipe.normalize(
-                mean=NWP_MEANS[conf_nwp[nwp_key].nwp_provider],
-                std=NWP_STDS[conf_nwp[nwp_key].nwp_provider],
+                mean=NWP_MEANS[normalize_provider],
+                std=NWP_STDS[normalize_provider],
             )
 
     if "sat" in datapipes_dict:

ocf_datapipes/training/windnet.py

@@ -17,6 +17,7 @@
     fill_nans_in_pv,
     normalize_gsp,
     normalize_wind,
+    potentially_coarsen,
     slice_datapipes_by_time,
 )
 from ocf_datapipes.utils.consts import (
@@ -236,9 +237,15 @@ def construct_sliced_data_pipeline(
                 roi_height_pixels=conf_nwp[nwp_key].nwp_image_size_pixels_height,
                 roi_width_pixels=conf_nwp[nwp_key].nwp_image_size_pixels_width,
             )
+            nwp_datapipe = nwp_datapipe.map(potentially_coarsen)
+            # Somewhat hacky way for India specifically, need different mean/std for ECMWF data
+            if conf_nwp[nwp_key].nwp_provider in ["ecmwf"]:
+                normalize_provider = "ecmwf_india"
+            else:
+                normalize_provider = conf_nwp[nwp_key].nwp_provider
             nwp_datapipes_dict[nwp_key] = nwp_datapipe.normalize(
-                mean=NWP_MEANS[conf_nwp[nwp_key].nwp_provider],
-                std=NWP_STDS[conf_nwp[nwp_key].nwp_provider],
+                mean=NWP_MEANS[normalize_provider],
+                std=NWP_STDS[normalize_provider],
             )
 
             if upsample_nwp:

ocf_datapipes/utils/consts.py

@@ -35,7 +35,7 @@ def __getitem__(self, key):
 
 # --------------------------- NWP ------------------------------------
 
-NWP_PROVIDERS = ["ukv", "gfs", "icon-eu", "icon-global", "ecmwf"]
+NWP_PROVIDERS = ["ukv", "gfs", "icon-eu", "icon-global", "ecmwf", "ecmwf_india"]
 
 # ------ UKV
 # Means and std computed WITH version_7 and higher, MetOffice values
@@ -189,19 +189,67 @@ def __getitem__(self, key):
 ECMWF_STD = _to_data_array(ECMWF_STD)
 ECMWF_MEAN = _to_data_array(ECMWF_MEAN)
 
+# These were calculated from 200 random init times of ECMWF data from 2020-2024
+INDIA_ECMWF_MEAN = {
+    "dlwrf": 56931292.0,
+    "dswrf": 31114434.0,
+    "duvrs": 3363442.75,
+    "hcc": 0.2640938460826874,
+    "lcc": 0.09548956900835037,
+    "mcc": 0.11118797957897186,
+    "prate": 3.028022729267832e-05,
+    "sde": 0.0015021926956251264,
+    "sr": 31403020.0,
+    "t2m": 298.24462890625,
+    "tcc": 0.33420485258102417,
+    "u10": 0.7755107283592224,
+    "u100": 1.0605329275131226,
+    "u200": 1.2298915386199951,
+    "v10": 0.02332865633070469,
+    "v100": -0.07577426731586456,
+    "v200": -0.1255049854516983,
+}
+
+INDIA_ECMWF_STD = {
+    "dlwrf": 34551808.0,
+    "dswrf": 21211150.0,
+    "duvrs": 2300205.5,
+    "hcc": 0.3942722678184509,
+    "lcc": 0.22802403569221497,
+    "mcc": 0.2254289835691452,
+    "prate": 0.00023954990319907665,
+    "sde": 0.09043189883232117,
+    "sr": 23481620.0,
+    "t2m": 7.574307918548584,
+    "tcc": 0.4046371579170227,
+    "u10": 2.7440550327301025,
+    "u100": 4.084362506866455,
+    "u200": 4.770451068878174,
+    "v10": 2.401158571243286,
+    "v100": 3.5278923511505127,
+    "v200": 3.974159002304077,
+}
+
+
+INDIA_ECMWF_VARIABLE_NAMES = tuple(INDIA_ECMWF_MEAN.keys())
+INDIA_ECMWF_STD = _to_data_array(INDIA_ECMWF_STD)
+INDIA_ECMWF_MEAN = _to_data_array(INDIA_ECMWF_MEAN)
+
 # ------ ALL NWPS
 # These dictionaries are for convenience
 NWP_VARIABLE_NAMES = NWPStatDict(
     ukv=UKV_VARIABLE_NAMES,
     gfs=GFS_VARIABLE_NAMES,
     ecmwf=ECMWF_VARIABLE_NAMES,
+    ecmwf_india=INDIA_ECMWF_VARIABLE_NAMES,
 )
 NWP_STDS = NWPStatDict(
     ukv=UKV_STD,
     gfs=GFS_STD,
     ecmwf=ECMWF_STD,
+    ecmwf_india=INDIA_ECMWF_STD,
 )
-NWP_MEANS = NWPStatDict(ukv=UKV_MEAN, gfs=GFS_MEAN, ecmwf=ECMWF_MEAN)
+NWP_MEANS = NWPStatDict(ukv=UKV_MEAN, gfs=GFS_MEAN, ecmwf=ECMWF_MEAN, ecmwf_india=INDIA_ECMWF_MEAN)
 
 # --------------------------- SATELLITE ------------------------------