From d021915b748ad71b32ea42aad1c67dcf077b4928 Mon Sep 17 00:00:00 2001
From: peterdudfield <peter.dudfield@hotmail.com>
Date: Fri, 29 Nov 2024 15:03:37 +0000
Subject: [PATCH] add regridding

---
 .../data/mo_global/india_coords.nc            | Bin 0 -> 3224 bytes
 india_forecast_app/data/nwp.py                |  57 ++++++++++++++++++
 india_forecast_app/models/pvnet/model.py      |  42 +++++++++----
 india_forecast_app/models/pvnet/utils.py      |  31 +++++++---
 tests/data/test_nwp.py                        |  34 +++++++++++
 5 files changed, 142 insertions(+), 22 deletions(-)
 create mode 100644 india_forecast_app/data/mo_global/india_coords.nc
 create mode 100644 india_forecast_app/data/nwp.py
 create mode 100644 tests/data/test_nwp.py

diff --git a/india_forecast_app/data/mo_global/india_coords.nc b/india_forecast_app/data/mo_global/india_coords.nc
new file mode 100644
index 0000000000000000000000000000000000000000..b97874c4aa754768fff0db989cfd7f161ebad112
GIT binary patch
literal 3224
zcmb7GX>b%}5DkYR1QZBYIAZRDBaqEzvpIINZ$+Yjz|K8|$Yngh0Lr0Uf=XG62%?rC
z2#9hk2wI2~DvIEVfPyz7pi)2y3Zk4|Uw_^H@t;#QH8Z<2{dK>7@AcQ~+S-Ome59%i
z-KW$KpD=R5#JXG5{SkFGtsgt)wts(6&-2y!@6k-%{P&v%@@I2(rTf&OwIl26-Qo2U
zDPmGmvN{{gxS;Rlxk-vV9ah)>u17KG{2%ceZ&dHAqsj66^ff6dw&VPsuiHJ;Ew*2-
z2hVQq{oF^}V?W6A@!#KhY;e)J*eUt_%%so$JiECzc2@pBclQUID*cwR3-aFO32%&=
ze(^v|y+m=4r2ZVye?J=VcWAi9`02r`x1bS!k1N9c-dp|4QY3R6G@;Mlo{1EWhtyDP
z+NsBHM>CF#=Cs;~0k~4)w+QM3EoGgSR87O?XeH~Vr9Y{@oPji1w>71+?^m>z_1i3%
zGw@J#9okABZNE&7-nslsw3B?=>u`P1Uh+zJ8&gE2OMV@mx<5JZU{7?EJUbFRy`M)X
z$+vTxDMzk*sSP?y-Wir&7#WiP)s)`!OURTuGWVTXd->RCWJx_)vna9C$dbCUi)i64
z$d>xDHT)xxBX#E7K5k>y&c(=)dUI0(1t%tbj$Elbcl9qzUMao`xl(_gN^Z)L>yalo
z<Q=P@-Rr6QkuP}U6PkUtB42PRP|-g7^EVU-J_V|KDj*64r@~wHfkMHnh|uY{3`K%l
z(aIkP0q8P{&Y_ELcN)3~j$Mupi_Z$*fUbgPS49@Z7<3g}yHU7jx1pQh+imAFD&hrn
z7o5A-4eGouC4laNcd>5QC=?6s#mlLX6(|<`d#GYnis&JI=%LeVkDk(xo-@<NESQr(
zPw7iZD+2gklt_O{hHC&(B7G`(nE?6=CDN}z0|KBE0@AlYLKldD^e?cC3a}B7J_b%x
zf_I@*`dJE{07|8=rSlYh&%A(A>2K*_1y{fK5tKd$OR^{+g3|9GD;+`UdvLE}`=AU2
zrT=A_+o+Q$6CRX}C3?6h6F!u^Me-YuGT}w2#bBxgA>l`8s0!D49YVsB&<53=yFWxo
z_!9c-;(?q`j-p(6Q$9cggL2_d`BEw*6Xn9A@}Kp>2n(OWHTp$_g;(LZM5Kub3%|nO
zX`-M)cvjIv69E;%w~BkI(YsI~ysOx)8;J_hZ)FCR+ZdI?!^+W*E2OhgDSWJaQ@KHE
z6qUlusw>GZ2BAv$S#=Z9@g7tOPpj6eXz5R)O88p!hk|e1MpO%Lt8e&N9lKF2{H<QB
z*jjJ`)xzWIlO(@nM1;?g2-#^aBEsv)qkXmx3B?c*en-C5#6nbf9_>!yCL=0*k4~k=
z%|}#tAKj&aMpXD;(}{rDiyF~E%^g(8Nz{lQYF^h1qegTQYep~TBPRNY-Ke@d>>9*G
zC$SgFE`}f`dWl^k`<n`gZjjIVhUf?V8Bd6gFp>Qpq9-h2d?C8RyX*%Lec=@29T%O&
z+pzz}MQ`y4;~y8@#p~EV;-bIAb?k=;mwhd9i!Kz23H)ASCJ}iB5>Kcy<Fk;U>6A7W
ziH)q=#9OSN#2)5x;vn-WagzI*xXkz(sY(_fY=Y5Yp*j}9C?bGozzFmFF=|ObePK8}
zPmDUA8^(B^55@z^h6WyjL2Od~V=QAIFgCCrjcu$8V>jbze9JsBe&T)^muRe|!feX(
z#7tK-QFFm8WPdlyiJ;41_GWxdTN@C}o7w-&an#ljOr2gQnDg`w0_IZIx4E8mX1>9A
zo4c5I<~Q7b^90$zFqjue7`3pP=zRyQcG~!0<?DS0tTMfSfYplxz8#jS#@REgU=3Ge
zv3@ozEyufH&0wFh=CWT{OIWwoI@XW%I`i22NXaMXQ&?Z?hQm6}^U*rbbIwlYd1bd{
zUfOv)KkOh0s6Fg!l-;iV1Gd3FZ_{haZ|$+{d-inpAA1h-)LzVbv|nUh*js5F6~g|I
z_a*jMy#KI|@jk;or-cirG4s}Gt&JH@4(|h;0Bx5{IN-VA==Ns78Nxp5jAlP`rm?R$
zvsvHHB1O|NPs3SDc5yG9E!r^PywCIA`GV)Rb5s)=&TndJs5$C3(#8ijO>O&r-vKw9
zgt!K73C|%{%X2*30qo1}P3(_uJ^Pe9mHonfL@hh*XSfS#5bcM%hWCB$YrNlc-_wQ*
z_jA^}`y&~`0JvwhvA}Ccw%h?;E1sub7SA=0U{>?Rt6`t_`m_GM!R&k99qd2eWTMNp
z@E+#<lDB|$?yX{-d#|$2y`8LcZ$InY`<`{~onf8(Nvw0f1?$|`^2mf=%sTfYtaHB~
z>)ap2I`>Di&izTObN?aMxj&zE?yqE>`<qzj{tni;zmIk9A7P#Qr&;Iz-&*JY06E`3
A#{d8T

literal 0
HcmV?d00001

diff --git a/india_forecast_app/data/nwp.py b/india_forecast_app/data/nwp.py
new file mode 100644
index 0000000..de68896
--- /dev/null
+++ b/india_forecast_app/data/nwp.py
@@ -0,0 +1,57 @@
+import xarray as xr
+import os
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def regrid_nwp_data(nwp_zarr: str, target_coords_path: str, nwp_zarr_save: str):
+    """This function loads the  NWP data, then regrids and saves it back out if the data is not
+    on the same grid as expected. The data is resaved in-place.
+
+    method can be 'conservative' or 'bilinear'
+    """
+
+    logger.info(f"Regridding NWP data {nwp_zarr} to expected grid to {target_coords_path}")
+
+    ds_raw = xr.open_zarr(nwp_zarr)
+
+    # These are the coords we are aiming for
+    ds_target_coords = xr.load_dataset(target_coords_path)
+
+    # Check if regridding step needs to be done
+    needs_regridding = not (
+        ds_raw.latitude.equals(ds_target_coords.latitude)
+        and ds_raw.longitude.equals(ds_target_coords.longitude)
+    )
+
+    if not needs_regridding:
+        logger.info(f"No NWP regridding required for {nwp_zarr} - skipping this step")
+        return
+
+    logger.info(f"Regridding NWP {nwp_zarr} to expected grid")
+
+    # Pull the raw data into RAM
+    ds_raw = ds_raw.compute()
+
+    # regrid
+    ds_regridded = ds_raw.interp(
+        latitude=ds_target_coords.latitude, longitude=ds_target_coords.longitude
+    )
+
+    # Re-save - including rechunking
+    os.system(f"rm -rf {nwp_zarr_save}")
+    ds_regridded["variable"] = ds_regridded["variable"].astype(str)
+
+    # Rechunk to these dimensions when saving
+    save_chunk_dict = {
+        "step": 5,
+        "latitude": 100,
+        "longitude": 100,
+        "x": 100,
+        "y": 100,
+    }
+
+    ds_regridded.chunk(
+        {k: save_chunk_dict[k] for k in list(ds_raw.xindexes) if k in save_chunk_dict}
+    ).to_zarr(nwp_zarr_save)
diff --git a/india_forecast_app/models/pvnet/model.py b/india_forecast_app/models/pvnet/model.py
index d6604cc..47444f7 100644
--- a/india_forecast_app/models/pvnet/model.py
+++ b/india_forecast_app/models/pvnet/model.py
@@ -36,6 +36,7 @@
 )
 from .utils import (
     download_satellite_data,
+    NWPProcessAndCacheConfig,
     populate_data_config_sources,
     process_and_cache_nwp,
     save_batch,
@@ -214,26 +215,41 @@ def _prepare_data_sources(self):
         satellite_source_file_path = os.getenv("SATELLITE_ZARR_PATH", None)
 
         # only load nwp that we need
-        nwp_paths = []
-        nwp_source_file_paths = []
+        nwp_configs = []
         nwp_keys = self.config["input_data"]["nwp"].keys()
         if "ecmwf" in nwp_keys:
-            nwp_ecmwf_source_file_path = os.environ["NWP_ECMWF_ZARR_PATH"]
-            nwp_source_file_paths.append(nwp_ecmwf_source_file_path)
-            nwp_paths.append(nwp_ecmwf_path)
+
+            nwp_configs.append(
+                NWPProcessAndCacheConfig(
+                    source_nwp_path=os.environ["NWP_ECMWF_ZARR_PATH"],
+                    dest_nwp_path=nwp_ecmwf_path,
+                    source="ecmwf",
+                )
+            )
+
         if "gfs" in nwp_keys:
-            nwp_gfs_source_file_path = os.environ["NWP_GFS_ZARR_PATH"]
-            nwp_source_file_paths.append(nwp_gfs_source_file_path)
-            nwp_paths.append(nwp_gfs_path)
+
+            nwp_configs.append(
+                NWPProcessAndCacheConfig(
+                    source_nwp_path=os.environ["NWP_GFS_ZARR_PATH"],
+                    dest_nwp_path=nwp_gfs_path,
+                    source="gfs",
+                )
+            )
+
         if "mo_global" in nwp_keys:
-            nwp_mo_global_source_file_path = os.environ["NWP_MO_GLOBAL_ZARR_PATH"]
-            nwp_source_file_paths.append(nwp_mo_global_source_file_path)
-            nwp_paths.append(nwp_mo_global_path)
+            nwp_configs.append(
+                NWPProcessAndCacheConfig(
+                    source_nwp_path=os.environ["NWP_MO_GLOBAL_ZARR_PATH"],
+                    dest_nwp_path=nwp_mo_global_path,
+                    source="mo_global",
+                )
+            )
 
         # Remove local cached zarr if already exists
-        for nwp_source_file_path, nwp_path in zip(nwp_source_file_paths, nwp_paths, strict=False):
+        for nwp_config in nwp_configs:
             # Process/cache remote zarr locally
-            process_and_cache_nwp(nwp_source_file_path, nwp_path)
+            process_and_cache_nwp(nwp_config)
         if use_satellite and "satellite" in self.config["input_data"].keys():
             shutil.rmtree(satellite_path, ignore_errors=True)
             download_satellite_data(satellite_source_file_path)
diff --git a/india_forecast_app/models/pvnet/utils.py b/india_forecast_app/models/pvnet/utils.py
index 32b02b2..f6e023e 100644
--- a/india_forecast_app/models/pvnet/utils.py
+++ b/india_forecast_app/models/pvnet/utils.py
@@ -10,7 +10,9 @@
 import yaml
 from ocf_datapipes.batch import BatchKey
 from ocf_datapipes.utils.consts import ELEVATION_MEAN, ELEVATION_STD
+from pydantic import BaseModel
 
+from india_forecast_app.data.nwp import regrid_nwp_data
 from .consts import (
     nwp_ecmwf_path,
     nwp_gfs_path,
@@ -25,6 +27,12 @@
 log = logging.getLogger(__name__)
 
 
+class NWPProcessAndCacheConfig(BaseModel):
+    source_nwp_path: str
+    dest_nwp_path: str
+    source: str
+
+
 def worker_init_fn(worker_id):
     """
     Clear reference to the loop and thread.
@@ -92,12 +100,13 @@ def populate_data_config_sources(input_path, output_path):
     return config
 
 
-def process_and_cache_nwp(source_nwp_path: str, dest_nwp_path: str):
+def process_and_cache_nwp(nwp_config: NWPProcessAndCacheConfig):
     """Reads zarr file, renames t variable to t2m and saves zarr to new destination"""
 
-    log.info(
-        f"Processing and caching NWP data for {source_nwp_path}, " f"and saving to {dest_nwp_path}"
-    )
+    source_nwp_path = nwp_config.source_nwp_path
+    dest_nwp_path = nwp_config.dest_nwp_path
+
+    log.info(f"Processing and caching NWP data for {source_nwp_path} and saving to {dest_nwp_path}")
 
     if os.path.exists(dest_nwp_path):
         log.info(f"File already exists at {dest_nwp_path}")
@@ -115,10 +124,7 @@ def process_and_cache_nwp(source_nwp_path: str, dest_nwp_path: str):
         if ds[v].dtype == object:
             ds[v].encoding.clear()
 
-    is_gfs = "gfs" in source_nwp_path.lower()
-    is_ecmwf = "ecmwf" in source_nwp_path.lower()
-
-    if is_ecmwf:
+    if nwp_config.source == "ecmwf":
         # Rename t variable to t2m
         variables = list(ds.variable.values)
         new_variables = []
@@ -134,14 +140,21 @@ def process_and_cache_nwp(source_nwp_path: str, dest_nwp_path: str):
         ds.__setitem__("variable", new_variables)
 
     # Hack to resolve some NWP data format differences between providers
-    elif is_gfs:
+    elif nwp_config.source == "gfs":
         data_var = ds[list(ds.data_vars.keys())[0]]
         # # Use .to_dataset() to split the data variable based on 'variable' dim
         ds = data_var.to_dataset(dim="variable")
         ds = ds.rename({"t2m": "t"})
+
     # Save destination path
     ds.to_zarr(dest_nwp_path, mode="a")
 
+    if nwp_config.source == "mo_global":
+        # regrid data
+        regrid_nwp_data(
+            dest_nwp_path, "india_forecast_app/data/mo_global/india_coords.nc", dest_nwp_path
+        )
+
 
 def download_satellite_data(satellite_source_file_path: str) -> None:
     """Download the sat data"""
diff --git a/tests/data/test_nwp.py b/tests/data/test_nwp.py
new file mode 100644
index 0000000..4c9617b
--- /dev/null
+++ b/tests/data/test_nwp.py
@@ -0,0 +1,34 @@
+""" Tests for the nwp regridding module """
+import os
+import tempfile
+
+import xarray as xr
+
+from india_forecast_app.data.nwp import regrid_nwp_data
+
+
+def test_regrid_nwp_data(nwp_mo_global_data):
+    """Test the regridding of the nwp data"""
+
+    # create a temporary dir
+    with tempfile.TemporaryDirectory() as temp_dir:
+
+        # save mo data to zarr
+        nwp_zarr = os.environ["NWP_MO_GLOBAL_ZARR_PATH"]
+
+        # regrid the data
+        nwp_zarr_save = f"{temp_dir}/nwp_regrid.zarr"
+        regrid_nwp_data(
+            nwp_zarr, "india_forecast_app/data/mo_global/india_coords.nc", nwp_zarr_save
+        )
+
+        # open the regridded data
+        nwp_xr = xr.open_zarr(nwp_zarr)
+        nwp_xr_regridded = xr.open_zarr(nwp_zarr_save)
+
+        # check the data is different in latitude and longitude
+        assert not nwp_xr_regridded.latitude.equals(nwp_xr.latitude)
+        assert not nwp_xr_regridded.longitude.equals(nwp_xr.longitude)
+
+        assert len(nwp_xr_regridded.latitude) == 225
+        assert len(nwp_xr_regridded.longitude) == 150