From 57e04516bc98fe89988032ce1a1ae07c0b7031ac Mon Sep 17 00:00:00 2001
From: Sukhil Patel <sukhil@Sukhils-MacBook-Pro.local>
Date: Mon, 6 Jan 2025 15:32:30 +0000
Subject: [PATCH] Update normalisation constants

---
 ocf_datapipes/training/pvnet_site.py |  8 +++++--
 ocf_datapipes/utils/consts.py        | 36 +++++++++++++++++++++++-----
 2 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/ocf_datapipes/training/pvnet_site.py b/ocf_datapipes/training/pvnet_site.py
index a0b3e808..75635e51 100644
--- a/ocf_datapipes/training/pvnet_site.py
+++ b/ocf_datapipes/training/pvnet_site.py
@@ -188,6 +188,7 @@ def construct_sliced_data_pipeline(
     location_pipe: IterDataPipe,
     t0_datapipe: IterDataPipe,
     production: bool = False,
+    new_normalisation_constants: bool = False
 ) -> dict:
     """Constructs data pipeline for the input data config file.
 
@@ -198,6 +199,7 @@ def construct_sliced_data_pipeline(
         location_pipe: Datapipe yielding locations.
         t0_datapipe: Datapipe yielding times.
         production: Whether constucting pipeline for production inference.
+        new_normalisation_constants: whether new normalisation constants are used.
     """
 
     datapipes_dict = _get_datapipes_dict(
@@ -245,8 +247,10 @@ def construct_sliced_data_pipeline(
             # 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"
-            elif conf_nwp[nwp_key].nwp_provider in ["gfs"]:
-                normalize_provider = "gfs_india"
+            elif new_normalisation_constants and conf_nwp[nwp_key].nwp_provider in ["mo_global"]:
+                normalize_provider = "mo_global_new_india"
+            elif new_normalisation_constants and conf_nwp[nwp_key].nwp_provider in ["gfs"]:
+                    normalize_provider = "gfs_india"
             else:
                 normalize_provider = conf_nwp[nwp_key].nwp_provider
 
diff --git a/ocf_datapipes/utils/consts.py b/ocf_datapipes/utils/consts.py
index f4f02cd9..955e048c 100644
--- a/ocf_datapipes/utils/consts.py
+++ b/ocf_datapipes/utils/consts.py
@@ -48,6 +48,7 @@ def __getitem__(self, key):
     "merra2",
     "merra2_uk",
     "mo_global",
+    "mo_global_new_india"
 ]
 
 # ------ UKV
@@ -133,9 +134,29 @@ def __getitem__(self, key):
 UKV_STD = _to_data_array(UKV_STD)
 UKV_MEAN = _to_data_array(UKV_MEAN)
 
-# --- MO Global
+# --- MO Global (partial initial constants)
 
 MO_GLOBAL_INDIA_MEAN = {
+    "temperature_sl": 298.2,
+    "wind_u_component_10m": 0.5732,
+    "wind_v_component_10m": -0.2831,
+}
+
+MO_GLOBAL_INDIA_STD = {
+    "temperature_sl": 8.473,
+    "wind_u_component_10m": 2.599,
+    "wind_v_component_10m": 2.016,
+}
+
+
+MO_GLOBAL_VARIABLE_NAMES = tuple(MO_GLOBAL_INDIA_MEAN.keys())
+MO_GLOBAL_INDIA_STD = _to_data_array(MO_GLOBAL_INDIA_STD)
+MO_GLOBAL_INDIA_MEAN = _to_data_array(MO_GLOBAL_INDIA_MEAN)
+
+
+# --- MO Global New
+
+MO_GLOBAL_INDIA_NEW_MEAN = {
     "temperature_sl": 295.34392488,
     "wind_u_component_10m": 0.83223102,
     "wind_v_component_10m": 0.0802083,
@@ -148,7 +169,7 @@ def __getitem__(self, key):
     "visibility_sl": 23181.81547681,
 }
 
-MO_GLOBAL_INDIA_STD = {
+MO_GLOBAL_INDIA_NEW_STD = {
     "temperature_sl": 12.26983825,
     "wind_u_component_10m": 3.45169835,
     "wind_v_component_10m": 2.9825603,
@@ -162,9 +183,9 @@ def __getitem__(self, key):
 }
 
 
-MO_GLOBAL_VARIABLE_NAMES = tuple(MO_GLOBAL_INDIA_MEAN.keys())
-MO_GLOBAL_INDIA_STD = _to_data_array(MO_GLOBAL_INDIA_STD)
-MO_GLOBAL_INDIA_MEAN = _to_data_array(MO_GLOBAL_INDIA_MEAN)
+MO_GLOBAL_NEW_VARIABLE_NAMES = tuple(MO_GLOBAL_INDIA_NEW_MEAN.keys())
+MO_GLOBAL_INDIA_NEW_STD = _to_data_array(MO_GLOBAL_INDIA_NEW_STD)
+MO_GLOBAL_INDIA_NEW_MEAN = _to_data_array(MO_GLOBAL_INDIA_NEW_MEAN)
 
 
 # ------ GFS
@@ -213,7 +234,7 @@ def __getitem__(self, key):
 GFS_MEAN = _to_data_array(GFS_MEAN)
 
 
-# ------ GFS
+# ------ GFS India
 GFS_INDIA_STD_DICT = {
     "t": 14.93798,
     "prate": 5.965701e-05,
@@ -434,6 +455,7 @@ def __getitem__(self, key):
     merra2=MERRA2_VARIABLE_NAMES,
     merra2_uk=UK_MERRA2_VARIABLE_NAMES,
     mo_global=MO_GLOBAL_VARIABLE_NAMES,
+    mo_global_new_india=MO_GLOBAL_NEW_VARIABLE_NAMES
 )
 NWP_STDS = NWPStatDict(
     ukv=UKV_STD,
@@ -445,6 +467,7 @@ def __getitem__(self, key):
     merra2=MERRA2_STD,
     merra2_uk=UK_MERRA2_STD,
     mo_global=MO_GLOBAL_INDIA_STD,
+    mo_global_new_india=MO_GLOBAL_INDIA_NEW_STD,
 )
 NWP_MEANS = NWPStatDict(
     ukv=UKV_MEAN,
@@ -456,6 +479,7 @@ def __getitem__(self, key):
     merra2=MERRA2_MEAN,
     merra2_uk=UK_MERRA2_MEAN,
     mo_global=MO_GLOBAL_INDIA_MEAN,
+    mo_global_new_india=MO_GLOBAL_INDIA_NEW_MEAN
 )
 
 # --------------------------- SATELLITE ------------------------------