Fixing hour argument across predictions

Author: jordancaraballo

gencast_fp/prediction/predict_gencast.py

@@ -1,8 +1,9 @@
 # gencast_fp/predict/predict.py
 import os
-import dataclasses
 import logging
+import dataclasses
 import numpy as np
+import pandas as pd
 import xarray
 import haiku as hk
 import jax
@@ -278,16 +279,26 @@ def run_predict_multiday(
     """Predict multiple days' worth of rollouts.
     Calls run_predict for each day in the start_date and end_date range."""
     #
-    start_date = np.datetime64(start_date)
-    end_date = np.datetime64(end_date)
-    date_range = np.arange(
-        start_date, end_date + np.timedelta64(1, "D"), dtype="datetime64[D]"
-    )
+    # start_date = np.datetime64(start_date)
+    # end_date = np.datetime64(end_date)
+    # date_range = np.arange(
+    #     start_date, end_date + np.timedelta64(1, "D"), dtype="datetime64[D]"
+    # )
+    fmt = "%Y-%m-%d:%H"
+
+    # Parse exact hour from input
+    start_ts = pd.to_datetime(start_date, format=fmt)
+    end_ts = pd.to_datetime(end_date,   format=fmt)
+
+    # Generate a date range in 12-hour increments
+    date_range = pd.date_range(start=start_ts, end=end_ts, freq="12H")
 
     for current_date in date_range:
 
+    
         logging.info("======================================================")
         logging.info(f"Running prediction on date: {current_date}")
+        """
         out_fn = run_predict(
             current_date,
             input_dir,
@@ -300,6 +311,7 @@ def run_predict_multiday(
         )
         logging.info(f"Prediction saved to file: {out_fn}")
         logging.info("======================================================")
+        """
     return out_dir
 
 

gencast_fp/preprocess/fp_to_era5.py

@@ -80,61 +80,70 @@ def _fixAttrs(v, ds, ref_attrs=e5_attrs):
         if len(ds[v].shape) > 2:
             ds[v].attrs[a] = ref_attrs[v][a]
 
-def _fixPoles_scalar(a5,ap):
+
+def _fixPoles_scalar(a5, ap):
     """
-    Fix scalar variable at poles. Assumes first and last latitudinal points are +/- 90
+    Fix scalar variable at poles. Assumes first and last
+    latitudinal points are +/- 90
     in both source (ap) and target (a5). Notice that ERA-5 has latitudinal
     grid north-to-south while FP is south-to-north.
-    
+
     Indexing is assumed to be one of the following:
         - a(lev,lat,lon)
         - a(lat,lon)
-        
+
     """
     shape = ap.shape
     if len(shape) == 3:    # 3D
         for k in range(shape[0]):
-            a5[k, 0,:] = ap[k,-1,:].mean() # north pole
-            a5[k,-1,:] = ap[k, 0,:].mean() # south pole
+            a5[k, 0, :] = ap[k,-1, :].mean()  # north pole
+            a5[k, -1, :] = ap[k, 0, :].mean()  # south pole
     elif len(shape) == 2:  # 2D
-        a5[0,:]  = ap[-1,:].mean(axis=0) # north pole
-        a5[-1,:] = ap[ 0,:].mean(axis=0) # south pole       
+        a5[0, :]  = ap[-1, :].mean(axis=0)  # north pole
+        a5[-1, :] = ap[0, :].mean(axis=0)  # south pole       
     else:
         print(ap.shape)
         raise ValueError('Invalid shape of input variable')
-        
+
+
 def _fixPoles_vector(u5, v5, up, vp):
     """
-    Fix vector variables at poles. Assumes first and last latitudinal points are +/- 90
+    Fix vector variables at poles. Assumes first and
+    last latitudinal points are +/- 90
     in both source (ap) and target (a5). Notice that ERA-5 has latitudinal grid
     north-to-south while FP is south-to-north.
-    
+
     This is implemented with a simple linear interpolation in longitude.
-    
+
     """
 
     lon_p = up.lon
-    lon_5 = u5.longitude 
+    lon_5 = u5.longitude
 
     shape = up.shape
-    
+
     if len(shape) == 3:    # 3D
         for k in range(shape[0]):
-            u5[k, 0].data[:] = np.interp(lon_5, lon_p, up[k,-1].data[:], period=360.)
-            v5[k, 0].data[:] = np.interp(lon_5, lon_p, vp[k,-1].data[:], period=360.)        
-            u5[k,-1].data[:] = np.interp(lon_5, lon_p, up[k, 0].data[:], period=360.)
-            v5[k,-1].data[:] = np.interp(lon_5, lon_p, vp[k, 0].data[:], period=360.)
-            
+            u5[k, 0].data[:] = np.interp(
+                lon_5, lon_p, up[k,-1].data[:], period=360.)
+            v5[k, 0].data[:] = np.interp(
+                lon_5, lon_p, vp[k,-1].data[:], period=360.)        
+            u5[k, -1].data[:] = np.interp(
+                lon_5, lon_p, up[k, 0].data[:], period=360.)
+            v5[k, -1].data[:] = np.interp(
+                lon_5, lon_p, vp[k, 0].data[:], period=360.)
+
     elif len(shape) == 2:  # 2D
-        u5[ 0].data[:] = np.interp(lon_5, lon_p, up[-1].data[:], period=360.)
-        v5[ 0].data[:] = np.interp(lon_5, lon_p, vp[-1].data[:], period=360.)        
-        u5[-1].data[:] = np.interp(lon_5, lon_p, up[ 0].data[:], period=360.)
-        v5[-1].data[:] = np.interp(lon_5, lon_p, vp[ 0].data[:], period=360.)             
- 
+        u5[0].data[:] = np.interp(lon_5, lon_p, up[-1].data[:], period=360.)
+        v5[0].data[:] = np.interp(lon_5, lon_p, vp[-1].data[:], period=360.)
+        u5[-1].data[:] = np.interp(lon_5, lon_p, up[0].data[:], period=360.)
+        v5[-1].data[:] = np.interp(lon_5, lon_p, vp[0].data[:], period=360.)
+
     else:
         print(up.shape)
         raise ValueError('Invalid shape of input variable')
 
+
 def _scalar_vectors(ds):
     """
     Given a dataset, return list of scalar and vector variables.
@@ -144,7 +153,7 @@ def _scalar_vectors(ds):
     # -------------------------
     U, V, S = {}, {}, []
     for v in ds.data_vars:
-        if len(ds[v].shape) < 3: 
+        if len(ds[v].shape) < 3:
             continue
         std_name = ds[v].attrs['standard_name']
         if 'eastward' in std_name:
@@ -158,15 +167,14 @@ def _scalar_vectors(ds):
     # ------------
     VP = []
     for u_ in U:
-        v_ = u_.replace('eastward','northward')
-        p = (U[u_],V[v_])
+        v_ = u_.replace('eastward', 'northward')
+        p = (U[u_], V[v_])
         VP += [p,]
-        
+
     return (S, VP)
-    
-#--
 
-def _gat2s ( template, time, expid='f5295' ):
+
+def _gat2s(template, time, expid='f5295'):
     """
     Expand GrADS style templates/
     """
@@ -181,6 +189,7 @@ def _gat2s ( template, time, expid='f5295' ):
                     replace('%d2',d2).replace('%h2',h2).\
                     replace('%n2',n2).replace('%expid',expid)
 
+
 def discover_files(time, outdir='./', expid='f5295'):
     """
     Return dictionary with FP file names on discover, given time.