return sorting argument

Author: AntSimi

src/py_eddy_tracker/observations/network.py

@@ -106,12 +106,14 @@ def fix_next_previous_obs(next_obs, previous_obs, flag_virtual):
 
 class NetworkObservations(GroupEddiesObservations):
 
-    __slots__ = ("_index_network",)
-
+    __slots__ = ("_index_network", "_index_segment_track", "_segment_track_array")
     NOGROUP = 0
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        self.reset_index()
+    
+    def reset_index(self):
         self._index_network = None
         self._index_segment_track = None
         self._segment_track_array = None
@@ -251,9 +253,8 @@ def elements(self):
 
     def astype(self, cls):
         new = cls.new_like(self, self.shape)
-        print()
-        for k in new.obs.dtype.names:
-            if k in self.obs.dtype.names:
+        for k in new.fields:
+            if k in self.fields:
                 new[k][:] = self[k][:]
         new.sign_type = self.sign_type
         return new
@@ -371,23 +372,27 @@ def correct_close_events(self, nb_days_max=20):
 
         self.segment[:] = segment_copy
         self.previous_obs[:] = previous_obs
-
-        self.sort()
+        return self.sort()
 
     def sort(self, order=("track", "segment", "time")):
         """
         Sort observations
 
         :param tuple order: order or sorting. Given to :func:`numpy.argsort`
         """
-        index_order = self.obs.argsort(order=order)
-        for field in self.elements:
+        index_order = self.obs.argsort(order=order, kind="mergesort")
+        self.reset_index()
+        for field in self.fields:
             self[field][:] = self[field][index_order]
 
-        translate = -ones(index_order.max() + 2, dtype="i4")
-        translate[index_order] = arange(index_order.shape[0])
+        nb_obs = len(self)
+        # we add 1 for -1 index return index -1
+        translate = -ones(nb_obs + 1, dtype="i4")
+        translate[index_order] = arange(nb_obs)
+        # next & previous must be re-indexed
         self.next_obs[:] = translate[self.next_obs]
         self.previous_obs[:] = translate[self.previous_obs]
+        return index_order, translate
 
     def obs_relative_order(self, i_obs):
         self.only_one_network()
@@ -654,16 +659,16 @@ def normalize_longitude(self):
         lon0 = (self.lon[i_start] - 180).repeat(i_stop - i_start)
         logger.debug("Normalize longitude")
         self.lon[:] = (self.lon - lon0) % 360 + lon0
-        if "lon_max" in self.obs.dtype.names:
+        if "lon_max" in self.fields:
             logger.debug("Normalize longitude_max")
             self.lon_max[:] = (self.lon_max - self.lon + 180) % 360 + self.lon - 180
         if not self.raw_data:
-            if "contour_lon_e" in self.obs.dtype.names:
+            if "contour_lon_e" in self.fields:
                 logger.debug("Normalize effective contour longitude")
                 self.contour_lon_e[:] = (
                     (self.contour_lon_e.T - self.lon + 180) % 360 + self.lon - 180
                 ).T
-            if "contour_lon_s" in self.obs.dtype.names:
+            if "contour_lon_s" in self.fields:
                 logger.debug("Normalize speed contour longitude")
                 self.contour_lon_s[:] = (
                     (self.contour_lon_s.T - self.lon + 180) % 360 + self.lon - 180
@@ -1071,7 +1076,7 @@ def extract_event(self, indices):
             raw_data=self.raw_data,
         )
 
-        for k in new.obs.dtype.names:
+        for k in new.fields:
             new[k][:] = self[k][indices]
         new.sign_type = self.sign_type
         return new
@@ -1194,27 +1199,11 @@ def dissociate_network(self):
         """
         Dissociate networks with no known interaction (splitting/merging)
         """
-
         tags = self.tag_segment(multi_network=True)
         if self.track[0] == 0:
             tags -= 1
-
         self.track[:] = tags[self.segment_track_array]
-
-        i_sort = self.obs.argsort(order=("track", "segment", "time"), kind="mergesort")
-        # Sort directly obs, with hope to save memory
-        self.obs.sort(order=("track", "segment", "time"), kind="mergesort")
-        self._index_network = None
-
-        # n & p must be re-indexed
-        n, p = self.next_obs, self.previous_obs
-        # we add 2 for -1 index return index -1
-        nb_obs = len(self)
-        translate = -ones(nb_obs + 1, dtype="i4")
-        translate[:-1][i_sort] = arange(nb_obs)
-        self.next_obs[:] = translate[n]
-        self.previous_obs[:] = translate[p]
-        return translate
+        return self.sort()
 
     def network_segment(self, id_network, id_segment):
         return self.extract_with_mask(self.segment_slice(id_network, id_segment))

src/py_eddy_tracker/observations/observation.py

@@ -306,12 +306,16 @@ def box_display(value):
         """Return values evenly spaced with few numbers"""
         return "".join([f"{v_:10.2f}" for v_ in value])
 
+    @property
+    def fields(self):
+        return list(self.obs.dtype.names)
+
     def field_table(self):
         """
         Produce description table of the fields available in this object
         """
         rows = [("Name (Unit)", "Long name", "Scale factor", "Offset")]
-        names = list(self.obs.dtype.names)
+        names = self.fields
         names.sort()
         for field in names:
             infos = VAR_DESCR[field]
@@ -414,7 +418,7 @@ def remove_fields(self, *fields):
         """
         nb_obs = self.obs.shape[0]
         fields = set(fields)
-        only_variables = set(self.obs.dtype.names) - fields
+        only_variables = set(self.fields) - fields
         track_extra_variables = set(self.track_extra_variables) - fields
         array_variables = set(self.array_variables) - fields
         new = self.__class__(
@@ -426,7 +430,7 @@ def remove_fields(self, *fields):
             raw_data=self.raw_data,
         )
         new.sign_type = self.sign_type
-        for name in new.obs.dtype.names:
+        for name in new.fields:
             logger.debug("Copy of field %s ...", name)
             new.obs[name] = self.obs[name]
         return new
@@ -444,12 +448,12 @@ def add_fields(self, fields=list(), array_fields=list()):
             track_array_variables=self.track_array_variables,
             array_variables=list(concatenate((self.array_variables, array_fields))),
             only_variables=list(
-                concatenate((self.obs.dtype.names, fields, array_fields))
+                concatenate((self.fields, fields, array_fields))
             ),
             raw_data=self.raw_data,
         )
         new.sign_type = self.sign_type
-        for name in self.obs.dtype.names:
+        for name in self.fields:
             logger.debug("Copy of field %s ...", name)
             new.obs[name] = self.obs[name]
         return new
@@ -467,8 +471,8 @@ def circle_contour(self, only_virtual=False, factor=1):
         """
         angle = radians(linspace(0, 360, self.track_array_variables))
         x_norm, y_norm = cos(angle), sin(angle)
-        radius_s = "contour_lon_s" in self.obs.dtype.names
-        radius_e = "contour_lon_e" in self.obs.dtype.names
+        radius_s = "contour_lon_s" in self.fields
+        radius_e = "contour_lon_e" in self.fields
         for i, obs in enumerate(self):
             if only_virtual and not obs["virtual"]:
                 continue
@@ -684,7 +688,7 @@ def distance(self, other):
 
     def __copy__(self):
         eddies = self.new_like(self, len(self))
-        for k in self.obs.dtype.names:
+        for k in self.fields:
             eddies[k][:] = self[k][:]
         eddies.sign_type = self.sign_type
         return eddies

src/py_eddy_tracker/observations/tracking.py

@@ -183,16 +183,16 @@ def normalize_longitude(self):
         lon0 = (self.lon[self.index_from_track] - 180).repeat(self.nb_obs_by_track)
         logger.debug("Normalize longitude")
         self.lon[:] = (self.lon - lon0) % 360 + lon0
-        if "lon_max" in self.obs.dtype.names:
+        if "lon_max" in self.fields:
             logger.debug("Normalize longitude_max")
             self.lon_max[:] = (self.lon_max - self.lon + 180) % 360 + self.lon - 180
         if not self.raw_data:
-            if "contour_lon_e" in self.obs.dtype.names:
+            if "contour_lon_e" in self.fields:
                 logger.debug("Normalize effective contour longitude")
                 self.contour_lon_e[:] = (
                     (self.contour_lon_e.T - self.lon + 180) % 360 + self.lon - 180
                 ).T
-            if "contour_lon_s" in self.obs.dtype.names:
+            if "contour_lon_s" in self.fields:
                 logger.debug("Normalize speed contour longitude")
                 self.contour_lon_s[:] = (
                     (self.contour_lon_s.T - self.lon + 180) % 360 + self.lon - 180

src/py_eddy_tracker/tracking.py

@@ -658,7 +658,7 @@ def merge(self, until=-1, raw_data=True):
         # Set type of eddy with first file
         eddies.sign_type = self.current_obs.sign_type
         # Fields to copy
-        fields = self.current_obs.obs.dtype.names
+        fields = self.current_obs.fields
 
         # To know if the track start
         first_obs_save_in_tracks = zeros(self.i_current_by_tracks.shape, dtype=bool_)
@@ -707,7 +707,7 @@ def merge(self, until=-1, raw_data=True):
             # Index in the current file
             index_current = self[i]["out"]
 
-            if "cost_association" in eddies.obs.dtype.names:
+            if "cost_association" in eddies.fields:
                 eddies["cost_association"][index_final - 1] = self[i]["cost_value"]
             # Copy all variable
             for field in fields: