enstools/feature/identification/african_easterly_waves/configuration.py

@@ -8,14 +8,16 @@ import numpy as np
 # latS = -35
 # lonW = -100
 # lonE = 45
-data_lat = (-35, 35)
+
+data_lat = (3, 35)
 data_lon = (-100, 45)
 
-aew_clim_dir = '/home/ws/he7273/phd_all/data/aew/clim/cv_clim_era5.nc' # 'C:\\Users\\Christoph\\phd\\data\\enstools-feature\\cv_clim_era5.nc' # '/home/christoph/phd/data/aew/clim/cv_clim_era5.nc' # '/home/christoph/phd/data/framework_example_ds/aew/' # '/project/meteo/w2w/C3/fischer/belanger/aew_clim/' #
-in_files = '/home/ws/he7273/phd_all/data/coll_oper/jja2021/jja2021.nc' # 'C:\\Users\\Christoph\\phd\\data\\enstools-feature\\2008_sum_uv.nc' # '/home/christoph/phd/data/framework_example_ds/aew/cv_aug_08.nc'
+aew_clim_dir = '/project/meteo/w2w/C3/fischer/belanger/aew_clim/cv_clim_era5.nc' # '/lsdf/MOD/Gruppe_Transregio/Gruppe_Knippertz/kitweather/data/era5/cv_clim_era5.nc' 'C:\\Users\\Christoph\\phd\\data\\enstools-feature\\cv_clim_era5.nc' # '/home/christoph/phd/data/aew/clim/cv_clim_era5.nc' # '/home/christoph/phd/data/framework_example_ds/aew/' # '/project/meteo/w2w/C3/fischer/belanger/aew_clim/' #
+in_files = '/project/meteo/w2w/C3/fischer/data/jja2022.nc' # '/home/ws/he7273/phd_all/data/coll_oper/jja2021/jja2021.nc'  # 'C:\\Users\\Christoph\\phd\\data\\enstools-feature\\2008_sum_uv.nc' # '/home/christoph/phd/data/framework_example_ds/aew/cv_aug_08.nc'
+out_dir = '/project/meteo/w2w/C3/fischer/belanger/out/' # join('/home/ws/he7273/phd_all/data/aew/out/') # '/project/meteo/w2w/C3/fischer/belanger/out/'
+out_json_path = out_dir + 'jja2022_fc_test.json'
 
-out_dir = join('/home/ws/he7273/phd_all/data/aew/out/') # '/project/meteo/w2w/C3/fischer/belanger/out/'
-plot_dir = join('/home/ws/he7273/phd_all/data/aew/plots/') # '/project/meteo/w2w/C3/fischer/belanger/plots/'
+plot_dir = '/project/meteo/w2w/C3/fischer/belanger/plots/' # join('/home/ws/he7273/phd_all/data/aew/plots/') # '/project/meteo/w2w/C3/fischer/belanger/plots/'
 
 timedelta_ana = np.timedelta64(7, 'D')
 
@@ -23,14 +25,23 @@ timedelta_ana = np.timedelta64(7, 'D')
 # num_files_ecmwf = 41 # files in finished forecast data
 # file_prefix_ecmwf = 'ecmwf-hres_latlon_1.0deg_'
 # file_prefix_ecmwf_rain = 'ecmwf-hres_latlon_0.4deg_'
-plot_dir_ecmwf = '/home/iconeps/plots/aew/ecmwf-hres/' # '/project/meteo/w2w/C3/fischer/belanger/plots/'
 
 # aew_kitweather_icon_dir = '/lsdf/MOD/Gruppe_Transregio/Gruppe_Knippertz/kitweather/data/icon-global-det/forecasts/'
 # num_files_icon = 31 # files in finished forecast data
 # file_prefix_icon = 'icon-global-det_latlon_1.0deg_'
 # file_prefix_icon_rain = 'icon-global-det_latlon_0.25deg_'
+plot_dir_ecmwf = '/home/iconeps/plots/aew/ecmwf-hres/' # '/project/meteo/w2w/C3/fischer/belanger/plots/'
+plot_dir_ecmwf_prefix = 'aew_ecmwf-hres_'
+plot_dir_ecmwf_ens = '/home/iconeps/plots/aew/ecmwf-ens/' # '/project/meteo/w2w/C3/fischer/belanger/plots/'
+plot_dir_ecmwf_ens_prefix = 'aew_ecmwf-ens_'
 plot_dir_icon = '/home/iconeps/plots/aew/icon-global-det/' # '/project/meteo/w2w/C3/fischer/belanger/plots/'
+plot_dir_icon_prefix = 'aew_icon-global-det_'
 plot_dir_gfs = '/home/iconeps/plots/aew/gfs/'
+plot_dir_gfs_prefix = 'aew_gfs_'
+plot_dir_gfs_ens = '/home/iconeps/plots/aew/gefs/' # '/project/meteo/w2w/C3/fischer/belanger/plots/'
+plot_dir_gfs_ens_prefix = 'aew_gefs_'
+
+ens_rain_threshold = 1.0 # display prob of exceeding X mm hr-1
 
 latest_run_dir = '/home/iconeps/plots/aew/'
 latest_run_info_file_name = 'latest_aew_run.txt'
@@ -52,18 +63,18 @@ def get_clim_file():
 
 
 # wave area to be extracted: at least one point of trough needs to be in this range
-wave_filter_lat = (0, 30)
+wave_filter_lat = (-5, 30) # TODO see bottom method -> more adaptive on lons.
 wave_filter_lon = (-110, 55)
 
 levels = [700]  # 700 hPa
-u_dim = 'u700_rea'
-v_dim = 'v700_rea'
+u_dim = 'xxu700_fc'
+v_dim = 'v700_fc'
 
 # time of interest, if None all
 # june-oct is AEW season
 
-start_date = None # '2008-08-01T00:00' #  # '2008-08-01T00:00'
-end_date = None # '2008-08-15T00:00' # None # '2008-08-03T00:00'
+start_date = None # '2022-08-01T00:00' # None # '2008-08-01T00:00' #  # '2008-08-01T00:00'
+end_date = None # '2022-08-15T00:00'  # None # '2008-08-15T00:00' # None # '2008-08-03T00:00'
 
 # Algorithm parameters
 # max u wind (m/s) (0 = only keep west-propagating). Belanger: 2.5; Berry: 0.0
@@ -96,3 +107,17 @@ avg_speed_min_deg_per_h = avg_speed_min_m_per_s * 3.6 * 0.00914
 # at 10°N we have in longitude direction 0.00914 degrees/km (360/(40,000*cos(10deg)))
 speed_range_m_per_s = [3.0, 12.0] # [5,10], but be more gentle with polygons.
 speed_deg_per_h = [-m_per_s * 3.6 * 0.00914 for m_per_s in speed_range_m_per_s]  # negative -> westward [-5, -10]
+
+def is_point_in_area(lon, lat):
+
+    if lon < -65:
+        if 5 < lat < 28:
+            return True
+    elif lon < -50:
+        if 3 < lat < 27:
+            return True
+    else:
+        if 0 < lat < 26:
+            return True
+    
+    return False

enstools/feature/identification/african_easterly_waves/identification.py

@@ -59,8 +59,8 @@ class AEWIdentification(IdentificationStrategy):
         lat_range = self.config.data_lat
         clim_file = self.config.get_clim_file()
 
-        u_name = self.config.u_dim if self.config.u_dim is not None else get_u_var(dataset)
-        v_name = self.config.v_dim if self.config.v_dim is not None else get_v_var(dataset)
+        u_name = self.config.u_dim if hasattr(self.config, 'u_dim') and (self.config.u_dim is not None) else get_u_var(dataset)
+        v_name = self.config.v_dim if hasattr(self.config, 'v_dim') and (self.config.v_dim is not None) else get_v_var(dataset)
 
         if u_name is None or v_name is None:
             print("Could not locate u and v fields in dataset. Needed to compute advection terms.")
@@ -129,7 +129,7 @@ class AEWIdentification(IdentificationStrategy):
             dataset[u_name] = dataset[u_name].expand_dims(level=self.config.levels)
             dataset[v_name] = dataset[v_name].expand_dims(level=self.config.levels)
             level_str = 'level'
-
+        
         # rename cv_clim dimensions to be same as in data.
         cv_clim = cv_clim.rename({'lat': lat_str, 'lon': lon_str})
         if 'plev' in cv_clim.dims and 'plev' != level_str:
@@ -220,6 +220,7 @@ class AEWIdentification(IdentificationStrategy):
         return dataset
 
     def identify(self, data_chunk: xr.Dataset, **kwargs):
+
         objs = []
         trough_mask_cur = data_chunk.trough_mask
         if np.isnan(trough_mask_cur).all():
@@ -246,11 +247,12 @@ class AEWIdentification(IdentificationStrategy):
 
         id_ = 1
         for path in paths:
+
             # get new object, set id
             o = self.get_new_object()
             o.id = id_
             # populate it
-            populate_object(o.properties, path)
+            populate_object(o.properties, path, self.config)
 
             # add to objects if keep
             if not self.keep_wavetrough(o.properties):
@@ -397,6 +399,7 @@ class AEWIdentification(IdentificationStrategy):
         Called for each wavetrough, check if kept based on filtering heuristics:
         - WT requires any point of wavetrough in config.wave_filter range
         - WT requires minimum length threshold
+        - WT lat center not in 5..25
 
         Parameters
         ----------
@@ -420,5 +423,9 @@ class AEWIdentification(IdentificationStrategy):
 
         if properties.length_deg <= self.config.degree_len_thr:  # too small
             return False
+        
+        mid_lat = properties.bb.max.lat - properties.bb.min.lat
+        if mid_lat < 5.0 or mid_lat > 25.0:
+            return False
 
         return True

enstools/feature/identification/african_easterly_waves/plotting.py

@@ -14,10 +14,17 @@ 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():
+def get_kitweather_rain_cm(ens_mode):
     rgb_colors = []
-    pathtotxtfile = '/home/iconeps/icon_data/additional_data/colorpalettes/'
+    
+    pathtotxtfile = '/project/meteo/w2w/C3/fischer/belanger/enstools-feature/enstools/feature/identification/african_easterly_waves/' # '/home/iconeps/icon_data/additional_data/colorpalettes/'
     filename_colorpalette = 'colorpalette_dyamond_prec_rate.txt'
+    
+    if ens_mode:
+        filename_colorpalette = 'colormap_WhiteBeigeGreenBlue_16.txt'
+    else:
+        filename_colorpalette = 'colorpalette_dyamond_prec_rate.txt'
+
     with open(pathtotxtfile + filename_colorpalette, 'r') as f:
         lines = f.readlines()
     for i, line in enumerate(lines):
@@ -26,7 +33,10 @@ def get_kitweather_rain_cm():
     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]
+    if ens_mode:
+        levels = [0.01,0.02,0.05,0.1,0.12,0.15,0.2,0.25,0.3,0.4,0.5,0.75]
+    else:
+        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
@@ -51,8 +61,8 @@ def crop_top_bottom_whitespace(path):
     return
 
 
-## MAIN PLOTTING FUNC FOR KITWEATHER PLOTS
-def plot_ts_filtered_waves(wts_part_of_tracks, fig_name, ds=None, tp=None):
+## MAIN PLOTTING FUNC FOR KITWEATHER PLOTS - DETERMINISTIC MODE
+def plot_ts_filtered_waves(wts_part_of_tracks, fig_name, ds=None, tp=None, ens_mode=False):
     
     from timeit import default_timer as timer
     t1 = timer()
@@ -68,8 +78,7 @@ def plot_ts_filtered_waves(wts_part_of_tracks, fig_name, ds=None, tp=None):
 
     extent = [-100, 35, -10, 35]
 
-
-    levels_rain, rain_cm, norm = get_kitweather_rain_cm()
+    levels_rain, rain_cm, norm = get_kitweather_rain_cm(ens_mode)
     distance_plot_to_cbar = 0.010
     axins = ax.inset_axes([1 + distance_plot_to_cbar, 0.05, 0.015, 0.93],
                       transform=ax.transAxes)
@@ -77,12 +86,15 @@ def plot_ts_filtered_waves(wts_part_of_tracks, fig_name, ds=None, tp=None):
     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,
+    unit_text = '>1 mm/hr\nprobability' if ens_mode else 'mm/hr'
+    y_off = -0.075 if ens_mode else -0.06
+    axins.text(0.25, y_off, unit_text, transform=axins.transAxes,
            horizontalalignment='left', verticalalignment='center')
 
     t2 = timer()
 
-    if ds is not None:
+    if ds is not None and not ens_mode: # no uv for ens
+
         # print("Before dec")
         # streamplot_func = _add_transform_first_to_streamplot(ds.plot.streamplot)
         # print("After dec")
@@ -105,7 +117,10 @@ def plot_ts_filtered_waves(wts_part_of_tracks, fig_name, ds=None, tp=None):
     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)
+        if ens_mode: # ensemble way thinner
+            patch = patches.PathPatch(line, linewidth=1, facecolor='none', edgecolor='crimson')
+        else:
+            patch = patches.PathPatch(line, linewidth=3, facecolor='none', edgecolor='crimson') # cmap(time_weight)
         ax.add_patch(patch)
 
     t5 = timer()
@@ -152,6 +167,98 @@ def plot_ts_filtered_waves(wts_part_of_tracks, fig_name, ds=None, tp=None):
     return
 
 
+"""
+## MAIN PLOTTING FUNC FOR KITWEATHER PLOTS - ENSEMBLE MODE
+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:
+
+        # print("Before dec")
+        # streamplot_func = _add_transform_first_to_streamplot(ds.plot.streamplot)
+        # print("After dec")
+        ds.plot.streamplot(x='lon', y='lat', u='u', v='v', linewidth=0.6,
+                   arrowsize = 0.5,
+                   density=6,
+                   color='black') # , transform_first=True not working, or is already implemented. still slow.
+        
+    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("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):
@@ -268,8 +375,11 @@ def plot_track(track, fn):
 
     return figure_name
 
+
 from collections import defaultdict
 def plot_differences(set_graph, tracks, ds=None, tp=None, plot_prefix=None):
+    print("plot_differences() deprecated")
+    exit() # OLD FUNC.
     # plot the differences of the total graph and the tracks
     # so check which WTs are part of a track and which have been dropped.
 
@@ -314,10 +424,13 @@ def plot_differences(set_graph, tracks, ds=None, tp=None, plot_prefix=None):
 
     for date in dates_list:
         fig_name = plot_prefix + date[0:4] + date[5:7] + date[8:10] + "T" + date[11:13] + ".png"
-        ds_ss = ds.sel(time=date)
+        try:
+            ds_ss = ds.sel(time=date)
+        except (KeyError, AttributeError) as e:
+            print("No ds data for " + str(date))
+            ds_ss = None
         try:
             tp_ss = tp.sel(time=date).tp
-
         except (KeyError, AttributeError) as e:
             print("No rain data for " + str(date))
             tp_ss = None
@@ -329,6 +442,48 @@ def plot_differences(set_graph, tracks, ds=None, tp=None, plot_prefix=None):
             # plot_ts_part_of_track(wts_in_tracks[date], wts_not_in_tracks[date], fig_name, ds=ds_ss, tp=tp_ss)
 
 
+import pandas as pd
+def plot_kw(tracks, ds, tp=None, plot_prefix=None, ens_mode=False):
+
+    dates_dt = pd.to_datetime(ds.time.values)
+
+    wavetroughs_in_tracks = dict() # by time
+    for date in dates_dt:
+        wavetroughs_in_tracks[date] = []
+
+    # put nodes of all tracks in buckets by time
+    for track in tracks:
+        track_nodes = [e.parent for e in track.edges]
+        for track_node in track_nodes:
+            
+            track_node_time = pb_str_to_datetime(track_node.time)
+            wavetroughs_in_tracks[track_node_time].append(track_node)
+
+    if plot_prefix is None:
+        plot_prefix = cfg.plot_dir
+        # create subdirs if needed
+    else:
+        plot_dir = '/'.join(plot_prefix.split('/')[:-1]) + '/'
+    os.makedirs(plot_dir, exist_ok=True)
+
+    if ens_mode and tp is not None:
+        tp_thr = (tp > cfg.ens_rain_threshold).astype(dtype=float)
+        tp_prob = tp_thr.mean(dim="member") # TODO get_member_dim?
+        tp = tp_prob
+
+    # call plotting for each date
+    for date in dates_dt:
+        fig_name = plot_prefix + date.strftime("%Y%m%dT%H") + ".png"
+
+        ds_ss = ds.sel(time=date).squeeze()
+        try:
+            tp_ss = tp.sel(time=date)
+
+        except (KeyError, AttributeError) as e:
+            print("No rain data for " + str(date))
+            tp_ss = None
+        plot_ts_filtered_waves(wavetroughs_in_tracks[date], fig_name, ds=ds_ss, tp=tp_ss, ens_mode=ens_mode)
+
 
 def plot_track_from_graph(track_desc, fig_name_prefix, cv=None):
 

enstools/feature/identification/african_easterly_waves/processing.py

@@ -211,14 +211,24 @@ def compute_cv(dataset, u_str, v_str, cv_str):
     return dataset
 
 
-def populate_object(obj_props, path):
+def populate_object(obj_props, path, cfg):
     # obj_props.num_nodes = len(path)
     # fill the properties defined in the .proto file.
 
+    # first, remove vertices out of area
+    filtered_vertices = []
+    for v_idx, v in enumerate(path.vertices):
+        lon = v[0]
+        lat = v[1]
+
+        # dont use vertex not fulfilling area restriction
+        if cfg.is_point_in_area(lon, lat):
+            filtered_vertices.append(v)
+
     # vertices of path
     min_lat, max_lat, min_lon, max_lon = 90.0, -90.0, 180.0, -180.0
     dist_deg = 0.0
-    for v_idx, v in enumerate(path.vertices):
+    for v_idx, v in enumerate(filtered_vertices):
         line_pt = obj_props.line_pts.add()
         line_pt.lon = v[0]
         line_pt.lat = v[1]
@@ -234,7 +244,7 @@ def populate_object(obj_props, path):
 
         if v_idx > 0:
             dist_deg = dist_deg + (
-                        ((path.vertices[v_idx - 1][0] - v[0]) ** 2 + (path.vertices[v_idx - 1][1] - v[1]) ** 2) ** 0.5)
+                        ((filtered_vertices[v_idx - 1][0] - v[0]) ** 2 + (filtered_vertices[v_idx - 1][1] - v[1]) ** 2) ** 0.5)
 
     # bounding box
     obj_props.bb.min.lat = min_lat

enstools/feature/identification/african_easterly_waves/run_identify.py

@@ -8,7 +8,7 @@ from datetime import timedelta, datetime
 from enstools.feature.identification._proto_gen import african_easterly_waves_pb2
 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
+from enstools.feature.identification.african_easterly_waves.plotting import plot_kw, 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, shutil
 from enstools.feature.util.data_utils import get_subset_by_description
@@ -32,10 +32,29 @@ if len(sys.argv) >= 3 and sys.argv[1] == '-kw':
 
     if sys.argv[2] == 'ecmwf_fc':
         ds, run = get_ecmwf_forecast(run=run, include_analysis_delta=cfg.timedelta_ana, add_prec_rate=True)
+        plot_dir = cfg.plot_dir_ecmwf
+        plot_file_prefix = cfg.plot_dir_ecmwf_prefix
+        ens_mode = False
     elif sys.argv[2] == 'icon_fc':
         ds, run = get_icon_forecast(run=run, include_analysis_delta=cfg.timedelta_ana, add_prec_rate=True)
+        plot_dir = cfg.plot_dir_icon
+        plot_file_prefix = cfg.plot_dir_icon_prefix
+        ens_mode = False
     elif sys.argv[2] == 'gfs_fc':
         ds, run = get_gfs_forecast(run=run, include_analysis_delta=cfg.timedelta_ana, add_prec_rate=True)
+        plot_dir = cfg.plot_dir_gfs
+        plot_file_prefix = cfg.plot_dir_gfs_prefix
+        ens_mode = False
+    elif sys.argv[2] == 'ecmwf_ens':
+        ds, run = get_ecmwf_ensemble(run=run, add_prec_rate=True)
+        plot_dir = cfg.plot_dir_ecmwf_ens
+        plot_file_prefix = cfg.plot_dir_ecmwf_ens_prefix
+        ens_mode = True
+    elif sys.argv[2] == 'gfs_ens':
+        ds, run = get_gfs_ensemble(run=run, add_prec_rate=True)
+        plot_dir = cfg.plot_dir_gfs_ens
+        plot_file_prefix = cfg.plot_dir_gfs_ens_prefix
+        ens_mode = True
     else:
         print("Unknown command line parameter " + sys.argv[2] + ", expected ecmwf_fc or icon_fc.")
         exit(1)
@@ -62,7 +81,8 @@ pipeline.execute()
 
 od = pipeline.get_object_desc()
 
-for trackable_set in od.sets:
+all_tracks = []
+for set_id, trackable_set in enumerate(od.sets):
 
     # generate graph out of tracked data
     g = DataGraph(trackable_set, t_strat)
@@ -76,13 +96,27 @@ for trackable_set in od.sets:
     # children of a node: track.get_children(track.graph.edges[0].parent)
 
     # plot tracks
+
     for track_id, track in enumerate(tracks):
-        plot_track(track, "track" + str(track_id))
+        plot_track(track, "track" + "{:03d}".format(set_id) + "_" + "{:03d}".format(track_id))
 
     ds = pipeline.get_data()
     ds_set = get_subset_by_description(ds, trackable_set, '2d')
     # plot differences (passed filtering / did not pass)
-    # plot_differences(g, tracks, cv=ds_set.cv)
+    # plot_differences(g, tracks) TODO need to fix if we want this
+
+    # for track_id, track in enumerate(tracks):
+    #     plot_track(track, "track" + str(track_id))
+
+    all_tracks.extend(tracks)
+# ds_set = get_subset_by_description(ds, trackable_set, '2d')
+    
+ds = pipeline.get_data()
+
+# plot kitweather mode
+if len(sys.argv) > 1 and sys.argv[1] == '-kw':
+    plot_kw(all_tracks, ds=ds, tp=rain_ds.tp, plot_prefix=(plot_dir + run + '/' + plot_file_prefix), ens_mode=ens_mode) # TODO need g?
+
 
 
 # no out data besides plots on kitweather
@@ -97,7 +131,7 @@ if len(sys.argv) > 1 and sys.argv[1] == '-kw':
 
 
     # All done. Update text file containing time of latest finished run.
-    if os.path.exists(cfg.plot_dir_ecmwf + run) and os.path.exists(cfg.plot_dir_icon + run) and os.path.exists(cfg.plot_dir_gfs + run):
+    if os.path.exists(cfg.plot_dir_ecmwf + run) and os.path.exists(cfg.plot_dir_icon + run) and os.path.exists(cfg.plot_dir_gfs + run) and not ens_mode:
         print("Run " + run + " done for ECMWF, ICON, GFS. Update latest run info file.") 
         yyyymmddhh = run[4:]
 
@@ -120,8 +154,8 @@ if len(sys.argv) > 1 and sys.argv[1] == '-kw':
 # out_netcdf_path = data_path + '_streamers.nc'
 """
 if len(sys.argv) == 1:
-    out_json_path = out_dir + 'aew_desc.json'
-    out_dataset_path = out_dir + '05_wt.nc'
+    out_json_path = out_dir + 'aew_desc_w21.json'
+    out_dataset_path = out_dir + '05_wt_w21.nc'
 elif len(sys.argv) == 2:
     out_json_path = out_dir + 'aew_desc_' + str(proc_summer_of_year) + '.json'
     out_dataset_path = out_dir + '05_wt_' + str(proc_summer_of_year) + '.nc'
@@ -132,4 +166,5 @@ else:
     
 """
 
-pipeline.save_result(description_type='json', description_path=out_json_path) # , dataset_path=out_dataset_path) # dataset_path=out_dataset_path,
+# no out dataset here.
+pipeline.save_result(description_type='json', description_path=cfg.out_json_path) #, dataset_path=out_dataset_path) # dataset_path=out_dataset_path,