add tracker

Author: AntSimi

src/py_eddy_tracker/featured_tracking/area_tracker.py

@@ -0,0 +1,34 @@
+from ..observations.observation import EddiesObservations as Model
+from ..dataset.grid import RegularGridDataset
+from numpy import ma
+import logging
+
+logger = logging.getLogger("pet")
+
+
+class AreaTracker(Model):
+    def tracking(self, other):
+        shape = (self.shape[0], other.shape[0])
+        i, j, c = self.match(other, intern=False)
+        cost_mat = ma.empty(shape, dtype="f4")
+        cost_mat.mask = ma.ones(shape, dtype="bool")
+        m = c > 0
+        i, j, c = i[m], j[m], c[m]
+        cost_mat[i, j] = 1 - c
+
+        i_self, i_other = self.solve_function(cost_mat)
+        i_self, i_other = self.post_process_link(other, i_self, i_other)
+        logger.debug("%d matched with previous", i_self.shape[0])
+        return i_self, i_other, cost_mat[i_self, i_other]
+
+    def propagate(self, previous_obs, current_obs, obs_to_extend, dead_track, nb_next, model):
+        virtual = super().propagate(previous_obs, current_obs, obs_to_extend, dead_track, nb_next, model)
+        nb_dead = len(previous_obs)
+        nb_virtual_extend = nb_next - nb_dead
+        for key in model.elements:
+            if "contour_" not in key:
+                continue
+            virtual[key][:nb_dead] = current_obs[key]
+            if nb_virtual_extend > 0:
+                virtual[key][nb_dead:] = obs_to_extend[key]
+        return virtual

src/py_eddy_tracker/observations/tracking.py

@@ -27,12 +27,24 @@
 ===========================================================================
 
 """
-from numpy import empty, arange, where, unique, interp, ones, bool_, zeros, array, median
+from numpy import (
+    empty,
+    arange,
+    where,
+    unique,
+    interp,
+    ones,
+    bool_,
+    zeros,
+    array,
+    median,
+)
 from .. import VAR_DESCR_inv
 import logging
 from datetime import datetime, timedelta
 from .observation import EddiesObservations
 from numba import njit
+from ..generic import split_line
 
 logger = logging.getLogger("pet")
 
@@ -224,28 +236,30 @@ def extract_ids(self, tracks):
         return self.__extract_with_mask(mask)
 
     def extract_first_obs_in_box(self, res):
-        data = empty(self.obs.shape, dtype=[('lon', 'f4'), ('lat', 'f4'), ('track', 'i4')])
-        data['lon'] = self.longitude - self.longitude % res
-        data['lat'] = self.latitude - self.latitude % res
-        data['track'] = self.obs["track"]
+        data = empty(
+            self.obs.shape, dtype=[("lon", "f4"), ("lat", "f4"), ("track", "i4")]
+        )
+        data["lon"] = self.longitude - self.longitude % res
+        data["lat"] = self.latitude - self.latitude % res
+        data["track"] = self.obs["track"]
         _, indexs = unique(data, return_index=True)
-        mask = zeros(self.obs.shape, dtype='bool')
+        mask = zeros(self.obs.shape, dtype="bool")
         mask[indexs] = True
         return self.__extract_with_mask(mask)
 
     def extract_in_direction(self, direction, value=0):
         nb_obs = self.nb_obs_by_track
         i_start = self.index_from_track
         i_stop = i_start + nb_obs - 1
-        if direction in ('S', 'N'):
+        if direction in ("S", "N"):
             d_lat = self.latitude[i_stop] - self.latitude[i_start]
-            mask = d_lat < 0 if 'S' == direction else d_lat > 0
+            mask = d_lat < 0 if "S" == direction else d_lat > 0
             mask &= abs(d_lat) > value
         else:
-            lon_start , lon_end = self.longitude[i_start], self.longitude[i_stop]
+            lon_start, lon_end = self.longitude[i_start], self.longitude[i_stop]
             lon_end = (lon_end - (lon_start - 180)) % 360 + lon_start - 180
             d_lon = lon_end - lon_start
-            mask = d_lon < 0 if 'W' == direction else d_lon > 0
+            mask = d_lon < 0 if "W" == direction else d_lon > 0
             mask &= abs(d_lon) > value
         mask = mask.repeat(nb_obs)
         return self.__extract_with_mask(mask)
@@ -280,7 +294,12 @@ def median_filter(self, half_window, xfield, yfield, inplace=True):
             self.obs[yfield] = result
 
     def __extract_with_mask(
-        self, mask, full_path=False, remove_incomplete=False, compress_id=False, reject_virtual=False,
+        self,
+        mask,
+        full_path=False,
+        remove_incomplete=False,
+        compress_id=False,
+        reject_virtual=False,
     ):
         """
         Extract a subset of observations
@@ -302,7 +321,7 @@ def __extract_with_mask(
 
         if full_path:
             if reject_virtual:
-                mask *= ~self.obs['virtual'].astype('bool')
+                mask *= ~self.obs["virtual"].astype("bool")
             tracks = unique(self.tracks[mask])
             mask = self.get_mask_from_id(tracks)
         elif remove_incomplete:
@@ -333,6 +352,14 @@ def __extract_with_mask(
                 new.obs["track"] = id_translate[new.obs["track"]]
         return new
 
+    def plot(self, ax, ref=None, ** kwargs):
+        if "label" in kwargs:
+            kwargs["label"] += " (%s eddies)" % (self.nb_obs_by_track != 0).sum()
+        x, y = split_line(self.longitude, self.latitude, self.tracks)
+        if ref is not None:
+            x = (x - ref) % 360 + ref
+        return ax.plot(x, y, **kwargs)
+
 
 @njit(cache=True)
 def compute_index(tracks, index, number):
@@ -373,7 +400,9 @@ def track_loess_filter(half_window, x, y, track):
         if i != 0:
             i_previous = i - 1
             dx = x[i] - x[i_previous]
-            while dx < half_window and i_previous != 0 and cur_track == track[i_previous]:
+            while (
+                dx < half_window and i_previous != 0 and cur_track == track[i_previous]
+            ):
                 w = (1 - (dx / half_window) ** 3) ** 3
                 y_sum += y[i_previous] * w
                 w_sum += w
@@ -412,7 +441,11 @@ def track_median_filter(half_window, x, y, track):
         cur_track = track[i]
         while x[i] - x[i_previous] > half_window or cur_track != track[i_previous]:
             i_previous += 1
-        while i_next < nb and x[i_next] - x[i] <= half_window and cur_track == track[i_next]:
+        while (
+            i_next < nb
+            and x[i_next] - x[i] <= half_window
+            and cur_track == track[i_next]
+        ):
             i_next += 1
         y_new[i] = median(y[i_previous:i_next])
     return y_new