Get close network in another atlas

Author: AntSimi

CHANGELOG.rst

@@ -31,7 +31,9 @@ Added
 - Save EddyAnim in mp4
 - Add method to get eddy contour which enclosed obs defined with (x,y) coordinates
 - Add **EddyNetworkSubSetter** to subset network which need special tool and operation after subset
-- Add functions to find relatives segments
+- Network:
+    - Add method to find relatives segments
+    - Add method to get cloase network in an other atlas
 
 [3.3.0] - 2020-12-03
 --------------------

src/py_eddy_tracker/observations/network.py

@@ -10,6 +10,7 @@
     arange,
     array,
     bincount,
+    concatenate,
     empty,
     in1d,
     ones,
@@ -298,6 +299,30 @@ def relatives(self, obs, order=2, direct=True, only_past=False, only_future=Fals
 
         return self.extract_with_mask(mask)
 
+    def close_network(self, other, nb_obs_min=10, **kwargs):
+        """
+        Get close network from another atlas.
+
+        :param self other: Atlas to compare
+        :param int nb_obs_min: Minimal number of overlap for one trajectory
+        :param dict kwargs: keyword arguments for match function
+        :return: return other atlas reduce to common track with self
+
+        .. warning::
+            It could be a costly operation for huge dataset
+        """
+        p0, p1 = self.period
+        indexs = list()
+        for i_self, i_other, t0, t1 in self.align_on(other, bins=range(p0, p1 + 2)):
+            i, j, s = self.index(i_self).match(other.index(i_other), **kwargs)
+            indexs.append(other.re_reference_index(j, i_other))
+        indexs = concatenate(indexs)
+        tr, nb = unique(other.track[indexs], return_counts=True)
+        m = zeros(other.track.shape, dtype=bool)
+        for i in tr[nb >= nb_obs_min]:
+            m[other.network_slice(i)] = True
+        return other.extract_with_mask(m)
+
     def numbering_segment(self):
         """
         New numbering of segment
@@ -359,7 +384,9 @@ def display_timeline(
         :param str,array field: yaxis values, if None, segments are used
         :param str method: if None, mean values are used
         :param float factor: to multiply field
-        :param str colors_mode: color of lines. "roll" means looping through colors, "y" means color adapt the y values (for matching color plots)
+        :param str colors_mode:
+            color of lines. "roll" means looping through colors,
+            "y" means color adapt the y values (for matching color plots)
         :return: plot mappable
         """
         self.only_one_network()

src/py_eddy_tracker/observations/observation.py

@@ -200,8 +200,8 @@ def __eq__(self, other):
             return False
         return array_equal(self.obs, other.obs)
 
-    ### colors methods
     def get_color(self, i):
+        """Return colors like a cyclic list"""
         return self.COLORS[i % self.NB_COLORS]
 
     @property
@@ -1076,6 +1076,21 @@ def match(self, other, method="overlap", intern=False, cmin=0, **kwargs):
         m = c >= cmin  # ajout >= pour garder la cmin dans la sélection
         return i[m], j[m], c[m]
 
+    @staticmethod
+    def re_reference_index(index, ref):
+        """
+        Shift index with ref
+
+        :param array,int index: local index to re ref
+        :param slice,array ref:
+            reference could be a slice in this case we juste add start to index
+            or could be indexs and in this case we need to translate
+        """
+        if isinstance(ref, slice):
+            return index + ref.start
+        else:
+            return ref[index]
+
     @classmethod
     def cost_function_common_area(cls, xy_in, xy_out, distance, intern=False):
         """How does it work on x bound ?

src/py_eddy_tracker/observations/tracking.py

@@ -516,13 +516,6 @@ def extract_with_mask(
                 new.track = id_translate[new.track]
         return new
 
-    @staticmethod
-    def re_reference_index(index, ref):
-        if isinstance(ref, slice):
-            return index + ref.start
-        else:
-            return ref[index]
-
     def shape_polygon(self, intern=False):
         """
         Get the polygon enclosing each trajectory.