kw func

enstools/feature/identification/african_easterly_waves/run_identify.py

@@ -10,7 +10,7 @@ from os.path import expanduser, join
 from enstools.feature.util.graph import DataGraph
 from enstools.feature.identification.african_easterly_waves.plotting import plot_differences, plot_track, plot_track_in_ts, plot_timesteps_from_desc, plot_tracks_from_desc
 import enstools.feature.identification.african_easterly_waves.configuration as cfg
-import os, sys
+import os, sys, glob
 from enstools.feature.util.data_utils import get_subset_by_description
 
 pipeline = FeaturePipeline(african_easterly_waves_pb2, processing_mode='2d')
@@ -24,24 +24,50 @@ pipeline = FeaturePipeline(african_easterly_waves_pb2, processing_mode='2d')
 #     if len(sys.argv) > 2:
 #         proc_month_of_year = int(sys.argv[2])
 
-if len(sys.argv) == 3 and sys.argv[1] == '-kw' and sys.argv[2] == 'ecmwf_fc':
-    # kitweather:
+if len(sys.argv) == 3 and sys.argv[1] == '-kw' and sys.argv[2] == 'ana':
+    # kitweather: make plots from available cached ecmwf analysis
+    data_fc_root = cfg.aew_kitweather_ecmwf_dir
+    in_file = data_fc_root + "*/*000h_tropicalvars.nc"
+    print("Executing for: " + in_file)
+
+elif len(sys.argv) == 3 and sys.argv[1] == '-kw' and sys.argv[2] == 'ecmwf_fc':
+    # kitweather: use last 7 days of analysis and the ecmwf forecast
+
     # get latest subdirectory time as what we choose
     data_fc_root = cfg.aew_kitweather_ecmwf_dir
-    all_subdirs = [d[0] for d in os.walk(data_fc_root)]
+    all_subdirs = glob.glob(data_fc_root + "*/")
     
+    print("Collecting forecast files...")
     latest_subdir = max(all_subdirs, key=os.path.getmtime)
+    forecast_files_glob = latest_subdir + "ecmwf-hres_latlon_1.0deg_*"
+    fc_file_list = glob.glob(forecast_files_glob)
 
-    # if len([name for name in latest_subdir if os.path.isfile(name)]) < 41: # should have 41 files
-    #     print("Less than 41 files in newest directory " + latest_subdir)
-    #     exit()
-    in_file = latest_subdir + "/*"
+    if len(fc_file_list) < 41:
+        print("Expected 41 files in " + forecast_files_glob + ", got " + str(len(fc_file_list)))
+        print("Exit.")
+        exit(1)
+    print("Found all 41 forecast files at " + forecast_files_glob + ".")
 
-elif len(sys.argv) == 3 and sys.argv[1] == '-kw' and sys.argv[2] == 'ana':
-    # kitweather: make plots from available cached ecmwf analysis
+    # get last 7 days of analysis: 000h from previous runs
+    
+    print("Collecting analysis files...")
     data_fc_root = cfg.aew_kitweather_ecmwf_dir
-    in_file = data_fc_root + "*/*000h_tropicalvars.nc"
-    print("Executing for: " + in_file)
+    all_subdirs_by_time = sorted(all_subdirs, key=os.path.getmtime)
+    last_7d_ana_subdirs = all_subdirs_by_time[-28:] # last 28 timesteps = last 7 days
+    last_7d_ana_glob = [sd + "ecmwf-hres_latlon_1.0deg_*_000h_tropicalvars.nc" for sd in last_7d_ana_subdirs]
+    ana_file_list = []
+    for ana_ts_glob in last_7d_ana_glob:
+        cur_g = glob.glob(ana_ts_glob)
+        if len(cur_g) != 1:
+            print("Found " + str(len(cur_g)) + " files at " + ana_ts_glob + ", expected 1. Exit.")
+            exit(1)
+        ana_file_list.extend(cur_g)
+
+    print("Found " + str(len(ana_file_list)) + " analysis files.")
+    in_file = list(set(ana_file_list + fc_file_list)) # current 000h twice.
+    for f in in_file:
+        print(f)
+    print("Done collecting files.")
 
 else:
     in_file = cfg.in_files
@@ -55,6 +81,7 @@ pipeline.set_identification_strategy(i_strat)
 pipeline.set_tracking_strategy(t_strat)
 
 pipeline.set_data_path(in_file)
+
 # execute pipeline
 pipeline.execute()