added possibility to rerun runs, changed colors and crop

enstools/feature/identification/african_easterly_waves/colorpalette_dyamond_prec_rate.txt

+163,205,231
+120,149,223
+ 81, 92,216
+ 63,164, 52
+ 86,209, 76
+255,245, 75
+254,185, 63
+251,136, 49
+251, 22, 33
+184, 18, 26
+135,  7, 18

enstools/feature/identification/african_easterly_waves/plotting.py

@@ -2,7 +2,9 @@ import os.path
 
 from matplotlib import pyplot as plt
 import numpy as np
+from PIL import Image
 import matplotlib
+import matplotlib as mpl
 import matplotlib.patches as patches
 import cartopy.crs as ccrs
 import cartopy.feature as cfeature
@@ -10,6 +12,140 @@ from datetime import datetime
 from enstools.feature.util.data_utils import pb_str_to_datetime
 from pathlib import Path
 import enstools.feature.identification.african_easterly_waves.configuration as cfg
+from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER
+
+def get_kitweather_rain_cm():
+    rgb_colors = []
+    pathtotxtfile = '/home/iconeps/icon_data/additional_data/colorpalettes/'
+    filename_colorpalette = 'colorpalette_dyamond_prec_rate.txt'
+    with open(pathtotxtfile + filename_colorpalette, 'r') as f:
+        lines = f.readlines()
+    for i, line in enumerate(lines):
+        rgb_colors.append([float(line[0:3])/255, float(line[4:7])/255, float(line[8:11])/255, 1])
+    rgb_colors = [[1, 1, 1, 0]] + rgb_colors + [[0.35, 0, 0.4, 1]]
+    cmap = mpl.colors.ListedColormap(rgb_colors[1:-1]) # , name=colorpalette
+    cmap = cmap.with_extremes(bad='white', under=rgb_colors[0], over=rgb_colors[-1])
+
+    levels = [0.1,0.2,0.3,0.5,1,2,3,5,10,20,30,50]
+    norm = mpl.colors.BoundaryNorm(levels, cmap.N)
+
+    return levels, cmap, norm
+
+
+def crop_top_bottom_whitespace(path):
+
+    # pixels from image left where a vertical column is scanned from top and bottom for non-white pixels
+    x_scan_position = 450
+    add_bottom_delta = 20
+
+    im = Image.open(path)
+    image_array_y = np.where(np.asarray(im.convert('L')) < 255, 1, 0)[:, x_scan_position]
+    vmargins = [np.where(image_array_y[2:] == 1)[0][0] + 2 + 1,
+                image_array_y[:-2].shape[0] - np.where(image_array_y[:-2] == 1)[0][-1] + 2]
+    im_cropped = Image.new('RGBA',(im.size[0], im.size[1] - vmargins[0] - vmargins[1] + add_bottom_delta), (0, 0, 0, 0))
+    im_cropped.paste(im.crop((0, vmargins[0], im.size[0], im.size[1] - vmargins[1] + add_bottom_delta)), (0, 0))
+    im.close()
+    im_cropped.save(path, 'png')
+    im_cropped.close()
+
+    return
+
+
+## MAIN PLOTTING FUNC FOR KITWEATHER PLOTS
+def plot_ts_filtered_waves(wts_part_of_tracks, fig_name, ds=None, tp=None):
+    
+    from timeit import default_timer as timer
+    t1 = timer()
+
+    resolution = 1600
+    cbar_space_px = 80
+    subplotparameters = mpl.figure.SubplotParams(left=0, bottom=0, right=1 - cbar_space_px / resolution, top=1,
+                                             wspace=0, hspace=0)
+    fig, ax = plt.subplots(figsize=(resolution / 100, resolution / 100),
+                       dpi=100,
+                       subplotpars=subplotparameters,
+                       subplot_kw=dict(projection = ccrs.PlateCarree()))
+
+    extent = [-100, 35, -10, 35]
+
+
+    levels_rain, rain_cm, norm = get_kitweather_rain_cm()
+    distance_plot_to_cbar = 0.010
+    axins = ax.inset_axes([1 + distance_plot_to_cbar, 0.05, 0.015, 0.93],
+                      transform=ax.transAxes)
+    ticks_list = levels_rain
+    cbar = fig.colorbar(mpl.cm.ScalarMappable(cmap=rain_cm, norm=norm),
+                    cax=axins, extend='both', extendfrac=0.03,
+                    ticks=ticks_list)
+    axins.text(0.5, -0.06, 'mm/hr', transform=axins.transAxes,
+           horizontalalignment='left', verticalalignment='center')
+
+    t2 = timer()
+
+    if ds is not None:
+        ds.plot.streamplot(x='lon', y='lat', u='u', v='v', linewidth=0.6,
+                   arrowsize = 0.5,
+                   density=6,
+                   color='black')
+
+    t3 = timer()
+
+    if tp is not None:
+        # transform to mm
+        tp.plot.contourf(levels=levels_rain, extend='max', subplot_kws={'transform_first': True},
+            cmap=rain_cm, norm=norm, add_colorbar=False)
+
+    t4 = timer()
+
+    # generate plot per pressure level, per time step
+
+    for obj_idx, node in enumerate(wts_part_of_tracks):
+        line_pts = node.object.properties.line_pts
+        line = patches.Path([[p.lon, p.lat] for p in line_pts])
+        patch = patches.PathPatch(line, linewidth=3, facecolor='none', edgecolor='crimson') # cmap(time_weight)
+        ax.add_patch(patch)
+
+    t5 = timer()
+
+    # ax.coastlines()
+    ax.add_feature(cfeature.BORDERS.with_scale('50m'), linewidth=0.3)
+    ax.add_feature(cfeature.COASTLINE.with_scale('50m'), linewidth=0.3)
+    ax.set_extent(extent, crs=ccrs.PlateCarree())
+    ax.add_feature(cfeature.LAND.with_scale('50m'), facecolor=list(np.array([255, 225, 171])/255))
+
+    ax.get_xaxis().set_ticklabels([])
+    ax.get_yaxis().set_ticklabels([])
+
+    gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True, linewidth=0.5, color='gray', alpha=0.5, linestyle='--')
+    gl.top_labels = False
+    gl.right_labels = False
+    gl.xformatter = LONGITUDE_FORMATTER
+    gl.yformatter = LATITUDE_FORMATTER
+
+    ax.set_title("")
+    fig.tight_layout()
+
+    plt.savefig(fig_name, format='png', backend='agg')
+
+    plt.figure().clear()
+    plt.close()
+    plt.cla()
+    plt.clf()
+
+    crop_top_bottom_whitespace(fig_name)
+
+    t6 = timer()
+    """
+    print("Init: " + str(t2 - t1))
+    print("Streamplot: " + str(t3 - t2))
+    print("Rain: " + str(t4 - t3))
+    print("Wavetroughs: " + str(t5 - t4))
+    print("Finalize: " + str(t6 - t5))
+    """
+    print("Saved to " +  fig_name)
+    return
+
+
 
 # plots the wave state (all wavetroughs given specific timestep in a set) ts: pb2.Timestep
 def plot_wavetroughs(ts, fig_name, cv=None):
@@ -174,7 +310,7 @@ def plot_differences(set_graph, tracks, ds=None, tp=None, plot_prefix=None):
         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)
+            plot_ts_filtered_waves(wts_in_tracks[date], fig_name, ds=ds_ss, tp=tp_ss.tp) # wts_not_in_tracks[date], 
         except KeyError:
             print("No rain data for " + str(date))
             
@@ -183,64 +319,6 @@ def plot_differences(set_graph, tracks, ds=None, tp=None, plot_prefix=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()))
-
-    x_ticks = [-100, -95, -85, -75, -65, -55, -45, -35, -25, -15, -5, 5, 15, 25, 35]
-    y_ticks = [0, 10, 20, 30]
-    extent = [-100, -45, -10, 35]
-
-    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
-
-    # colors per time step
-    # min_time = wave_thr_list[0].time.astype('float64')
-    # max_time = wave_thr_list[-1].time.astype('float64')
-    # cmap = matplotlib.cm.get_cmap('rainbow')
-    # color_wgts = np.linspace(0.0, 1.0, len(wave_thr_list))
-    # colors = ['red', 'yellow', 'green', 'blue', 'purple']
-
-    for obj_idx, node in enumerate(wts_part_of_tracks):
-        line_pts = node.object.properties.line_pts
-        line = patches.Path([[p.lon, p.lat] for p in line_pts])
-        patch = patches.PathPatch(line, linewidth=2, facecolor='none', edgecolor='orange') # cmap(time_weight)
-        ax.add_patch(patch)
-
-    # for obj_idx, node in enumerate(wt_not_part_of_tracks):
-    #     line_pts = node.object.properties.line_pts
-    #     line = patches.Path([[p.lon, p.lat] for p in line_pts])
-    #     patch = patches.PathPatch(line, linewidth=2, facecolor='none', edgecolor='red') # cmap(time_weight)
-    #     ax.add_patch(patch)
-
-    ax.coastlines()
-    ax.add_feature(cfeature.BORDERS.with_scale('50m'))
-    ax.set_extent(extent, crs=ccrs.PlateCarree())
-    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')
-
-    plt.figure().clear()
-    plt.close()
-    plt.cla()
-    plt.clf()
-
-    return
-
 
 def plot_track_from_graph(track_desc, fig_name_prefix, cv=None):
 

enstools/feature/identification/african_easterly_waves/run_identify.py

@@ -15,6 +15,7 @@ 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
+from pprint import pprint
 
 
 pipeline = FeaturePipeline(african_easterly_waves_pb2, processing_mode='2d')
@@ -34,35 +35,48 @@ if len(sys.argv) == 3 and sys.argv[1] == '-kw' and sys.argv[2] == 'ana':
     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':
+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
+    data_fc_root = cfg.aew_kitweather_ecmwf_dir
 
     # get latest subdirectory time as what we choose
-    data_fc_root = cfg.aew_kitweather_ecmwf_dir
     all_subdirs = glob.glob(data_fc_root + "*/")
-    
-    print(data_fc_root + "*/")
 
-    print("Collecting forecast files...")
-    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)
+    # if as additional argument run_YYYYMMDDHH, use this, otherwise find newest.
+    if len(sys.argv) == 4:
+        latest_subdir = data_fc_root + sys.argv[3] + '/'
+        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]
+        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. Found " + str(len(fc_file_list)) + " files in forecast.")
+            print("Exit.")
+            exit(1)
+    else:
+    
+        print(data_fc_root + "*/")
 
-    if len(fc_file_list) < 41:
-        print("Expected 41 files in " + forecast_files_glob + ", got " + str(len(fc_file_list)))
-        print("Trying previous timestep...")
-        latest_subdir = sorted_subdirs[-2]
+        print("Collecting forecast files...")
+        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("Missing files as well. Found " + str(len(fc_file_list)) + " files.")
-            print("Exit.")
-            exit(1)
+            print("Expected 41 files in " + forecast_files_glob + ", got " + str(len(fc_file_list)))
+            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
@@ -70,8 +84,12 @@ elif len(sys.argv) == 3 and sys.argv[1] == '-kw' and sys.argv[2] == 'ecmwf_fc':
     print("Collecting analysis files...")
     data_fc_root = cfg.aew_kitweather_ecmwf_dir
     all_subdirs_by_time = sorted(all_subdirs)
-    last_7d_ana_subdirs = all_subdirs_by_time[-28:-1] # last 28 timesteps = last 7 days
+
+    # get analysis files starting at latest_subdirs 7 days in the past
+    latest_subdir_rel_index = all_subdirs_by_time.index(latest_subdir) - len(all_subdirs_by_time)
+    last_7d_ana_subdirs = all_subdirs_by_time[latest_subdir_rel_index-28:latest_subdir_rel_index] # 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)
@@ -82,6 +100,7 @@ elif len(sys.argv) == 3 and sys.argv[1] == '-kw' and sys.argv[2] == 'ecmwf_fc':
 
     print("Found " + str(len(ana_file_list)) + " analysis files.")
     in_file = sorted(list(set(ana_file_list + fc_file_list))) # current 000h twice.
+    pprint(in_file)
 
     print("Collecting rain data...")
     # just get all tp 0.4deg 6h
@@ -158,10 +177,10 @@ for trackable_set in od.sets:
     ds = pipeline.get_data()
     ds_set = get_subset_by_description(ds, trackable_set, '2d')
 
-    if len(sys.argv) == 3 and sys.argv[1] == '-kw' and sys.argv[2] == 'ecmwf_fc':
+    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, 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':
+    elif len(sys.argv) >= 3 and sys.argv[1] == '-kw' and sys.argv[2] == 'ana':
         plot_differences(g, tracks, ds=ds_set, tp=rain_ds, plot_prefix=cfg.plot_dir + "ana/")
     else:
         plot_differences(g, tracks, ds=ds_set)
@@ -177,11 +196,21 @@ if sys.argv[1] == '-kw':
             print("Removing directory " + str(os.path.join(cfg.plot_dir, sd)))
             shutil.rmtree(os.path.join(cfg.plot_dir, sd))
 
+
     # All done. Update text file containing time of latest finished run.
     yyyymmddhh = time_dir[4:]
 
     with open(cfg.latest_run_info_file, 'w+') as info_file:
         info_file.write(yyyymmddhh)
+        
+    # All done. scp data over to webserver.
+    # TODO when port 22 free test.
+    # TODO improve plotting performance
+    path_webserver = '/home/iconeps/Data3/plots/ecmwf/aew_prediction_maps/'
+    print('scp -r ' + cfg.plot_dir + time_dir + '/ ' + 'iconeps@imk-tss-web.imk-tro.kit.edu:' + path_webserver)
+    os.system('scp -r ' + cfg.plot_dir + time_dir + '/ ' + 'iconeps@imk-tss-web.imk-tro.kit.edu:' + path_webserver)
+    os.system('scp ' + cfg.latest_run_info_file + ' iconeps@imk-tss-web.imk-tro.kit.edu:' + path_webserver)
+
     exit()
 
 # out_netcdf_path = data_path + '_streamers.nc'