added better plots and adaptive directory handling for kitweather

enstools/feature/identification/african_easterly_waves/configuration.py

@@ -10,15 +10,15 @@ from datetime import timedelta
 data_lat = (-35, 35)
 data_lon = (-100, 45)
 
-aew_clim_dir = '/home/iconeps/filter_scripts/aew_script/cv_clim_era5.nc' # '/home/christoph/phd/data/framework_example_ds/aew/' # '/project/meteo/w2w/C3/fischer/belanger/aew_clim/' # 
-aew_kitweather_ecmwf_dir = '/home/iconeps/ecmwf_data/model_data/ecmwf-hres/forecasts/'
+aew_clim_dir = '/lsdf/MOD/Gruppe_Transregio/Gruppe_Knippertz/kitweather/data/era5/cv_clim_era5.nc' # TODO # '/home/christoph/phd/data/framework_example_ds/aew/' # '/project/meteo/w2w/C3/fischer/belanger/aew_clim/' # 
+aew_kitweather_ecmwf_dir = '/lsdf/MOD/Gruppe_Transregio/Gruppe_Knippertz/kitweather/data/ecmwf-hres/forecasts/'
 
 in_files = '/home/christoph/phd/data/aew/cv/2008cv.nc' # '/home/christoph/phd/data/framework_example_ds/aew/cv_aug_08.nc'
 
 # out_dir = join(expanduser("~") + '/phd/data/aew/out/') # '/project/meteo/w2w/C3/fischer/belanger/out/'
 plot_dir = join('/home/iconeps/plots/aew/') # '/project/meteo/w2w/C3/fischer/belanger/plots/'
 
-cv_data_dir = '/project/meteo/w2w/C2/athul/data_athul/AEW_cv_data/' # reference where ALL cv data is (for clim calc.)
+cv_data_dir = '/project/meteo/w2w/C2/athul/data_athul/AEW_cv_data/' # reference where ALL cv data is (only for clim calc.)
 cv_data_ex = join(cv_data_dir, '*.nc')
 
 

enstools/feature/identification/african_easterly_waves/plotting.py

@@ -127,14 +127,10 @@ def plot_track(track, fn):
     return figure_name
 
 from collections import defaultdict
-def plot_differences(set_graph, tracks, cv=None, plot_prefix=None):
+def plot_differences(set_graph, tracks, ds=None, tp=None, plot_prefix=None):
     # plot the differences of the total graph and the tracks
     # so check which WTs are part of a track and which have been dropped.
 
-    # TODO
-    #   join tracks list
-    #   set_graph elements not in tracks
-    #   for eaach timestep plot two mengen
     set_nodes = [e.parent for e in set_graph.graph.edges]
     is_in_set_nodes = [False] * len(set_nodes)
 
@@ -174,13 +170,20 @@ def plot_differences(set_graph, tracks, cv=None, plot_prefix=None):
     os.makedirs(plot_prefix, exist_ok=True)
 
     for date in dates_list:
-        fig_name = plot_prefix + "part_of_wave_" + date.replace(':', '_') + ".png"
-        cv_ss = cv.sel(time=date)
-        plot_ts_part_of_track(wts_in_tracks[date], wts_not_in_tracks[date], fig_name, cv_ss)
+        fig_name = plot_prefix + "aew_" + date[0:4] + date[5:7] + date[8:10] + "T" + date[11:13] + ".png"
+        ds_ss = ds.sel(time=date)
+        try:
+            tp_ss = tp.sel(time=date)
+            plot_ts_part_of_track(wts_in_tracks[date], wts_not_in_tracks[date], fig_name, ds=ds_ss, tp=tp_ss.tp)
+        except KeyError:
+            print("No rain data for " + str(date))
+            
+            tp_ss = None
+            # plot_ts_part_of_track(wts_in_tracks[date], wts_not_in_tracks[date], fig_name, ds=ds_ss, tp=tp_ss)
 
 
 
-def plot_ts_part_of_track(wts_part_of_tracks, wt_not_part_of_tracks, fig_name, cv=None):
+def plot_ts_part_of_track(wts_part_of_tracks, wt_not_part_of_tracks, fig_name, ds=None, tp=None):
 
     fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(11, 4), subplot_kw=dict(projection=ccrs.PlateCarree()))
 
@@ -188,9 +191,16 @@ def plot_ts_part_of_track(wts_part_of_tracks, wt_not_part_of_tracks, fig_name, c
     y_ticks = [0, 10, 20, 30]
     extent = [-100, -45, -10, 35]
 
-    if cv is not None:
-        levelfc = np.asarray([0, 0.5, 1, 2, 3]) * 1e-5
-        cv.plot.contourf(levels=levelfc, vmin=0, extend='max', cmap='Blues')
+    if ds is not None:
+        ds.plot.streamplot(x='lon', y='lat', u='u', v='v', linewidth=0.3,
+                   arrowsize = 0.4,
+                   density=6,
+                   color='grey')
+
+    if tp is not None:
+        # transform to mm
+        levels_rain = np.asarray([0.1, 0.5, 1, 2, 5, 10])
+        tp.plot.contourf(levels=levels_rain, vmin=0, extend='max', cmap='Blues')
 
     # generate plot per pressure level, per time step
 
@@ -219,6 +229,8 @@ def plot_ts_part_of_track(wts_part_of_tracks, wt_not_part_of_tracks, fig_name, c
     yt1 = ax.set_yticks(y_ticks, crs=ccrs.PlateCarree())
     xt1 = ax.set_xticks(x_ticks, crs=ccrs.PlateCarree())
 
+    ax.set_title("")
+
     print("Save to " +  fig_name)
     plt.savefig(fig_name, format='png')
 

enstools/feature/identification/african_easterly_waves/run_identify.py

@@ -10,8 +10,12 @@ 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, glob
+import os, sys, glob, shutil
 from enstools.feature.util.data_utils import get_subset_by_description
+import xarray as xr
+xr.set_options(keep_attrs=True)
+import numpy as np
+
 
 pipeline = FeaturePipeline(african_easterly_waves_pb2, processing_mode='2d')
 
@@ -37,23 +41,36 @@ elif len(sys.argv) == 3 and sys.argv[1] == '-kw' and sys.argv[2] == 'ecmwf_fc':
     data_fc_root = cfg.aew_kitweather_ecmwf_dir
     all_subdirs = glob.glob(data_fc_root + "*/")
     
+    print(data_fc_root + "*/")
+
     print("Collecting forecast files...")
-    latest_subdir = max(all_subdirs, key=os.path.getmtime)
+    sorted_subdirs = sorted(all_subdirs)
+    latest_subdir = sorted_subdirs[-1]
     forecast_files_glob = latest_subdir + "ecmwf-hres_latlon_1.0deg_*"
     fc_file_list = glob.glob(forecast_files_glob)
 
+    fc_rain_file_list = [fc_file.replace("1.0", "0.4").replace("tropicalvars", "tp") for fc_file in fc_file_list if "000h" not in fc_file]
+
     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("Trying previous timestep...")
+        latest_subdir = sorted_subdirs[-2]
+        forecast_files_glob = latest_subdir + "ecmwf-hres_latlon_1.0deg_*"
+        fc_file_list = glob.glob(forecast_files_glob)
+        fc_rain_file_list = [fc_file.replace("1.0", "0.4").replace("tropicalvars", "tp") for fc_file in fc_file_list if "000h" not in fc_file]
+
+        if len(fc_file_list) < 41:
+            print("Missing files as well. Found " + str(len(fc_file_list)) + " files.")
+            print("Exit.")
+            exit(1)
     print("Found all 41 forecast files at " + forecast_files_glob + ".")
 
     # get last 7 days of analysis: 000h from previous runs
     
     print("Collecting analysis files...")
     data_fc_root = cfg.aew_kitweather_ecmwf_dir
-    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
+    all_subdirs_by_time = sorted(all_subdirs)
+    last_7d_ana_subdirs = all_subdirs_by_time[-28:-1] # 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:
@@ -64,10 +81,40 @@ elif len(sys.argv) == 3 and sys.argv[1] == '-kw' and sys.argv[2] == 'ecmwf_fc':
         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)
+    in_file = sorted(list(set(ana_file_list + fc_file_list))) # current 000h twice.
+
+    print("Collecting rain data...")
+    # just get all tp 0.4deg 6h
+    ana_rain_files = glob.glob(data_fc_root + "*/ecmwf-hres_latlon_0.4deg_*_006h_tp.nc")
+
+    # load forecast files separately: need to compute deltas from tp.
+    rain_fc_ds = xr.open_mfdataset(fc_rain_file_list)
+    rain_fc_tp = rain_fc_ds.tp
+    # rain_fc_tp_diff = rain_fc_tp.differentiate(coord="time", datetime_unit="6h")
+
+    times = rain_fc_tp.time.values
+    for t_idx, time in reversed(list(enumerate(times))):
+        if t_idx > 0:
+            rain_fc_tp.loc[dict(time=time)] = rain_fc_tp.isel(time=t_idx) - rain_fc_tp.isel(time=(t_idx-1))
+
+    rain_fc_ds['tp'] = rain_fc_tp
+
+    ana_rain_files = sorted(ana_rain_files)
+
+    rain_ana_ds = xr.open_mfdataset(ana_rain_files)
+    latest_ana_dt = rain_ana_ds.time[-1]
+    earliest_ana_dt = latest_ana_dt - np.timedelta64(7, 'D')
+    rain_ana_ds = rain_ana_ds.sel(time=(slice(earliest_ana_dt, latest_ana_dt + np.timedelta64(1, 'D')))) # TODO could remove last?
+
+    all_rain_ds = xr.merge([rain_ana_ds, rain_fc_ds])
+
+    # change total precipitation to hourly precip
+    all_rain_ds.tp.attrs['units'] = 'mm hr-1'
+    all_rain_ds.tp.attrs['long_name'] = 'Precipitation rate'
+    all_rain_ds['tp'] = all_rain_ds.tp / 6.0 * 1000.0 # from 6hrly (downloaded) to hourly rate and m to mm
+
     print("Done collecting files.")
+    all_rain_ds.to_netcdf("/home/iconeps/test.nc")
 
 else:
     in_file = cfg.in_files
@@ -107,16 +154,20 @@ for trackable_set in od.sets:
 
     if len(sys.argv) == 3 and sys.argv[1] == '-kw' and sys.argv[2] == 'ecmwf_fc':
         time_dir = os.path.basename(os.path.normpath(latest_subdir))
-        plot_differences(g, tracks, cv=ds_set.cv, plot_prefix=cfg.plot_dir + time_dir + "/")
+        plot_differences(g, tracks, ds=ds_set, tp=all_rain_ds, plot_prefix=cfg.plot_dir + time_dir + "/")
     elif len(sys.argv) == 3 and sys.argv[1] == '-kw' and sys.argv[2] == 'ana':
-        plot_differences(g, tracks, cv=ds_set.cv, plot_prefix=cfg.plot_dir + "ana/")
+        plot_differences(g, tracks, ds=ds_set, tp=rain_ds, plot_prefix=cfg.plot_dir + "ana/")
     else:
-        plot_differences(g, tracks, cv=ds_set.cv)
+        plot_differences(g, tracks, ds=ds_set)
 
 # no out data besides plots on kitweather
 if sys.argv[1] == '-kw':
-    exit()
-
+    # delete old plots
+    subdirs = [dI for dI in os.listdir(cfg.plot_dir) if os.path.isdir(os.path.join(cfg.plot_dir,dI))]
+    for sd in subdirs: # for each subdir in plot dir
+        if not sd == time_dir: # if not current dir -> del it
+            print("Removing directory " + str(os.path.join(cfg.plot_dir, sd)))
+            shutil.rmtree(os.path.join(cfg.plot_dir, sd))
 
 # out_netcdf_path = data_path + '_streamers.nc'