update animation example

Author: ludwigVonKoopa

examples/16_network/pet_segmentation_anim.py

@@ -37,12 +37,14 @@ def save(self, *args, **kwargs):
 # %%
 # Overlaod of class to pick up
 TRACKS = list()
+INDICES = list()
 
 
 class MyTrack(TrackEddiesObservations):
     @staticmethod
     def get_next_obs(i_current, ids, x, y, time_s, time_e, time_ref, window, **kwargs):
         TRACKS.append(ids["track"].copy())
+        INDICES.append(i_current)
         return TrackEddiesObservations.get_next_obs(
             i_current, ids, x, y, time_s, time_e, time_ref, window, **kwargs
         )
@@ -70,9 +72,16 @@ def get_next_obs(i_current, ids, x, y, time_s, time_e, time_ref, window, **kwarg
 def update(i_frame):
     tr = TRACKS[i_frame]
     mappable_tracks.set_array(tr)
-    s = 80 * ones(tr.shape)
+    s = 40 * ones(tr.shape)
     s[tr == 0] = 4
     mappable_tracks.set_sizes(s)
+
+    indices_frames = INDICES[i_frame]
+    mappable_CONTOUR.set_data(
+        e.contour_lon_e[indices_frames],
+        e.contour_lat_e[indices_frames],
+    )
+    mappable_CONTOUR.set_color(cmap.colors[tr[indices_frames] % len(cmap.colors)])
     return (mappable_tracks,)
 
 
@@ -85,6 +94,9 @@ def update(i_frame):
 mappable_tracks = ax.scatter(
     e.lon, e.lat, c=TRACKS[0], cmap=cmap, vmin=0, vmax=vmax, s=20
 )
+mappable_CONTOUR = ax.plot(
+    e.contour_lon_e[INDICES[0]], e.contour_lat_e[INDICES[0]], color=cmap.colors[0]
+)[0]
 ani = VideoAnimation(
     fig, update, frames=range(1, len(TRACKS), 4), interval=125, blit=True
 )