added plotting of track/nontrack waves

enstools/feature/identification/african_easterly_waves/plotting.py

@@ -126,6 +126,103 @@ def plot_track(track, fn):
 
     return figure_name
 
+from collections import defaultdict
+def plot_differences(set_graph, tracks, cv=None): # TODO CV too?
+    # 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)
+
+    for track in tracks:
+        track_nodes = [e.parent for e in track.graph.edges]
+
+        for track_node in track_nodes:
+            try:
+                # track node in list, set to True
+                idx = set_nodes.index(track_node)
+                is_in_set_nodes[idx] = True
+            except ValueError:
+                # not in it
+                continue
+
+    # lists which WTs are part of tracks, and which are not part of tracks
+    wts_in_tracks_list = [set_nodes[i] for i, b in enumerate(is_in_set_nodes) if b]
+    wts_not_in_tracks_list = [set_nodes[i] for i, b in enumerate(is_in_set_nodes) if not b]
+
+    # make these lists to dicts with date as key
+    wts_in_tracks = defaultdict(list)
+    for wt_in_track in wts_in_tracks_list:
+        wts_in_tracks[wt_in_track.time].append(wt_in_track)
+    wts_not_in_tracks = defaultdict(list)
+    for wt_not_in_track in wts_not_in_tracks_list:
+        wts_not_in_tracks[wt_not_in_track.time].append(wt_not_in_track)
+
+    dates = set()
+    dates.update(wts_in_tracks.keys())
+    dates.update(wts_not_in_tracks.keys())
+    dates_list = list(dates)
+    dates_list.sort()
+
+    for date in dates_list:
+        fig_name = cfg.plot_dir + "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)
+
+
+
+def plot_ts_part_of_track(wts_part_of_tracks, wt_not_part_of_tracks, fig_name, cv=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 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')
+
+    # 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='green') # 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())
+
+    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

@@ -6,9 +6,10 @@ from datetime import timedelta
 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_track, plot_track_in_ts, plot_timesteps_from_desc, plot_tracks_from_desc
+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 sys
+from enstools.feature.util.data_utils import get_subset_by_description
 
 pipeline = FeaturePipeline(african_easterly_waves_pb2, processing_mode='2d')
 
@@ -50,6 +51,11 @@ for trackable_set in od.sets:
     for track_id, track in enumerate(tracks):
         plot_track(track, "track" + str(track_id))
 
+    ds = pipeline.get_data()
+    ds_set = get_subset_by_description(ds, trackable_set, '2d')
+
+    plot_differences(g, tracks, cv=ds_set.cv)
+
 
 
 # out_netcdf_path = data_path + '_streamers.nc'