py: added plotting functionality for ex2

src/plotters.py

@@ -50,3 +50,95 @@ def plot_result_ex1_24(result, mode_probabilities):
     ax[1].set_xlim([0, N])
     plt.tight_layout()
     return fig
+
+def plot_result_ex2_2(result, trajs):
+    fig = plt.figure(constrained_layout=True, figsize=(20, 10))
+    gs = fig.add_gridspec(2, 2)
+    ax = [fig.add_subplot(gs[0,0]),
+         fig.add_subplot(gs[:, 1]),
+         fig.add_subplot(gs[1, 0])]
+    Yt = np.hstack(result['Y'])
+    yx, yy = radar_to_pos(Yt)
+
+    for key, T in trajs.items():
+        ax[1].plot(T[0, :], T[1, :], color='k', lw=2)
+
+    for track in result['tracks']:
+        x = np.vstack(track['x'])
+        t = np.hstack(track['t']).flatten()
+        assoc = np.hstack(track['associations']).flatten()
+        if track in result['confirmed_tracks']:
+            ls = '-'
+            l = ax[0].plot(t, track['identity']*np.ones(t.shape), ls=ls, markersize=3)[0]
+            ax[0].plot(assoc, track['identity']*np.ones(assoc.shape), ls='', color=l.get_color(), marker='x', markersize=6)
+            ax[1].plot(x[:, 0], x[:, 1], ls=ls, color=l.get_color(), lw=3)
+        else:
+            ls = '--'
+            ax[1].plot(x[:, 0], x[:, 1], ls=ls, lw=2)
+
+    ax[0].set_ylabel('Track identity')
+    ax[0].set_title('Confirmed tracks over time')
+    ax[0].set_xlabel('Time index, k')
+    ax[1].plot(yx, yy, '.', color='k')
+    ax[1].set_xlabel(r'$p_x$')
+    ax[1].set_ylabel(r'$p_y$')
+    ax[1].set_title('Measurements and measurement predictions + tracks')
+
+    # Plot the RMSE for the matched trajectories
+    for track_id, key in result['matches'].items():
+        T = trajs[key]
+        track = [track for track in result['confirmed_tracks'] if track['identity'] == track_id][0]
+        x = np.vstack(track['x']).T
+        t = np.hstack(track['t']).flatten()
+        tf = t[-1]
+        gtf = T.shape[1]
+        # Find starting point
+        n = np.argmin(np.sum((T-x[:2, [0]])**2, axis=0))
+        # Find end point
+        if T.shape[1] > x.shape[1]:
+            N = x.shape[1]
+        else:
+            N = T.shape[1]
+        xrmse = np.sum((T[:, n:N]-x[:2, :N-n])**2, axis=0)/N
+        terr = t[0]+np.abs(gtf-tf)
+        if N-n > 0:
+            T = np.linspace(0, 1, N-n)
+            ax[2].plot(T, xrmse, label='{} -- Time error: {} timesteps'.format(key, terr))
+
+    ax[2].set_xlabel('Trajectory index')
+    ax[2].set_ylabel('Positional RMSE')
+    ax[2].legend()
+    return fig
+
+def plot_result_ex2_24(result):
+    fig = plt.figure(constrained_layout=True, figsize=(20, 16))
+    gs = fig.add_gridspec(4, 2)
+    ax = [fig.add_subplot(gs[0, :]),
+         fig.add_subplot(gs[1:, :])]
+    Yt = np.hstack(result['Y'])
+    yx, yy = radar_to_pos(Yt)
+
+    for track in result['tracks']:
+        x = np.vstack(track['x'])
+        t = np.hstack(track['t']).flatten()
+        assoc = np.hstack(track['associations']).flatten()
+
+        if track in result['confirmed_tracks']:
+            ls = '-'
+            l = ax[0].plot(t, track['identity']*np.ones(t.shape), ls=ls, markersize=3)[0]
+            ax[0].plot(assoc, track['identity']*np.ones(assoc.shape), ls='', color=l.get_color(), marker='x', markersize=6)
+            ax[1].plot(x[:, 0], x[:, 1], ls=ls, color=l.get_color(), lw=3)
+        else:
+            ls = '--'
+            ax[1].plot(x[:, 0], x[:, 1], ls=ls, lw=2)
+
+    ax[0].set_ylabel('Track identity')
+    ax[0].set_title('Confirmed tracks over time')
+    ax[0].set_xlabel('Time index, k')
+    ax[1].plot(yx, yy, '.', color='k')
+    ax[1].set_xlabel(r'$p_x$')
+    ax[1].set_ylabel(r'$p_y$')
+    ax[1].set_title('Measurements and measurement predictions + tracks')
+    ax[1].set_xlim([-2000, 2000])
+    ax[1].set_ylim([-21000, -17000])
+    return fig