speed up /generalize some network method

Author: AntSimi

examples/06_grid_manipulation/pet_lavd.py

@@ -44,6 +44,7 @@ def update_axes(ax, mappable=None):
             orientation="horizontal",
         )
         cb.set_label("Vorticity integration along trajectory at initial position")
+        return cb
 
 
 kw_vorticity = dict(vmin=0, vmax=2e-5, cmap="viridis")
@@ -90,7 +91,8 @@ def save(self, *args, **kwargs):
 # Display vorticity field
 fig, ax, _ = start_ax()
 mappable = g.display(ax, abs(g.grid("vort")), **kw_vorticity)
-update_axes(ax, mappable)
+cb = update_axes(ax, mappable)
+cb.set_label("Vorticity")
 
 # %%
 # Particles

notebooks/python_module/02_eddy_identification/pet_eddy_detection.ipynb

@@ -253,7 +253,7 @@
       },
       "outputs": [],
       "source": [
-        "ax = start_axes(\"Speed Radius (km)\")\nkwargs = dict(vmin=10, vmax=50, s=80, ref=-10, cmap=\"magma_r\", factor=0.001)\na.scatter(ax, \"radius_s\", **kwargs)\nm = c.scatter(\n    ax, \"radius_s\", **kwargs\n)\nupdate_axes(ax, m)"
+        "ax = start_axes(\"Speed Radius (km)\")\nkwargs = dict(vmin=10, vmax=50, s=80, ref=-10, cmap=\"magma_r\", factor=0.001)\na.scatter(ax, \"radius_s\", **kwargs)\nm = c.scatter(ax, \"radius_s\", **kwargs)\nupdate_axes(ax, m)"
       ]
     },
     {
@@ -271,7 +271,7 @@
       },
       "outputs": [],
       "source": [
-        "ax = start_axes(\"Effective Radius (km)\")\nkwargs = dict(vmin=10, vmax=80, cmap=\"magma_r\", factor=0.001, lut=14, ref=-10)\na.filled(ax, \"effective_radius\", **kwargs)\nm = c.filled(\n    ax, \"radius_e\", **kwargs\n)\nupdate_axes(ax, m)"
+        "ax = start_axes(\"Effective Radius (km)\")\nkwargs = dict(vmin=10, vmax=80, cmap=\"magma_r\", factor=0.001, lut=14, ref=-10)\na.filled(ax, \"effective_radius\", **kwargs)\nm = c.filled(ax, \"radius_e\", **kwargs)\nupdate_axes(ax, m)"
       ]
     }
   ],

notebooks/python_module/06_grid_manipulation/pet_lavd.ipynb

@@ -37,7 +37,7 @@
       },
       "outputs": [],
       "source": [
-        "def start_ax(title=\"\", dpi=90):\n    fig = plt.figure(figsize=(16, 9), dpi=dpi)\n    ax = fig.add_axes([0, 0, 1, 1], projection=\"full_axes\")\n    ax.set_xlim(0, 32), ax.set_ylim(28, 46)\n    ax.set_title(title)\n    return fig, ax, ax.text(3, 32, \"\", fontsize=20)\n\n\ndef update_axes(ax, mappable=None):\n    ax.grid()\n    if mappable:\n        cb = plt.colorbar(\n            mappable,\n            cax=ax.figure.add_axes([0.05, 0.1, 0.9, 0.01]),\n            orientation=\"horizontal\",\n        )\n        cb.set_label(\"Vorticity integration along trajectory at initial position\")\n\n\nkw_vorticity = dict(vmin=0, vmax=2e-5, cmap=\"viridis\")"
+        "def start_ax(title=\"\", dpi=90):\n    fig = plt.figure(figsize=(16, 9), dpi=dpi)\n    ax = fig.add_axes([0, 0, 1, 1], projection=\"full_axes\")\n    ax.set_xlim(0, 32), ax.set_ylim(28, 46)\n    ax.set_title(title)\n    return fig, ax, ax.text(3, 32, \"\", fontsize=20)\n\n\ndef update_axes(ax, mappable=None):\n    ax.grid()\n    if mappable:\n        cb = plt.colorbar(\n            mappable,\n            cax=ax.figure.add_axes([0.05, 0.1, 0.9, 0.01]),\n            orientation=\"horizontal\",\n        )\n        cb.set_label(\"Vorticity integration along trajectory at initial position\")\n        return cb\n\n\nkw_vorticity = dict(vmin=0, vmax=2e-5, cmap=\"viridis\")"
       ]
     },
     {
@@ -84,7 +84,7 @@
       },
       "outputs": [],
       "source": [
-        "fig, ax, _ = start_ax()\nmappable = g.display(ax, abs(g.grid(\"vort\")), **kw_vorticity)\nupdate_axes(ax, mappable)"
+        "fig, ax, _ = start_ax()\nmappable = g.display(ax, abs(g.grid(\"vort\")), **kw_vorticity)\ncb = update_axes(ax, mappable)\ncb.set_label(\"Vorticity\")"
       ]
     },
     {

src/py_eddy_tracker/observations/network.py

@@ -6,7 +6,7 @@
 from glob import glob
 
 from numba import njit
-from numpy import arange, array, bincount, empty, ones, uint32, unique, zeros
+from numpy import arange, array, bincount, empty, in1d, ones, uint32, unique, zeros
 
 from ..generic import build_index, wrap_longitude
 from ..poly import bbox_intersection, vertice_overlap
@@ -119,13 +119,12 @@ def longer_than(self, nb_day_min=-1, nb_day_max=-1):
         if nb_day_max < 0:
             nb_day_max = 1000000000000
         mask = zeros(self.shape, dtype="bool")
-        for i, b0, b1 in self.iter_on(self.segment_track_array):
+        t = self.time
+        for i, b0, b1 in self.iter_on(self.track):
             nb = i.stop - i.start
             if nb == 0:
                 continue
-            t = self.time[i]
-            dt = t.max() - t.min()
-            if nb_day_min <= dt <= nb_day_max:
+            if nb_day_min <= ptp(t[i]) <= nb_day_max:
                 mask[i] = True
         return self.extract_with_mask(mask)
 
@@ -164,21 +163,26 @@ def obs_relative_order(self, i_obs):
         self.only_one_network()
         return self.segment_relative_order(self.segment[i_obs])
 
-    def connexions(self):
-        self.only_one_network()
+    def connexions(self, multi_network=False):
+        if multi_network:
+            segment = self.segment_track_array
+        else:
+            self.only_one_network()
+            segment = self.segment
         segments_connexion = dict()
 
         def add_seg(father, child):
             if father not in segments_connexion:
                 segments_connexion[father] = list()
             segments_connexion[father].append(child)
 
-        for i, seg, _ in self.iter_on("segment"):
+        previous_obs, next_obs = self.previous_obs, self.next_obs
+        for i, seg, _ in self.iter_on(segment):
             if i.start == i.stop:
                 continue
-            i_p, i_n = self.previous_obs[i.start], self.next_obs[i.stop - 1]
+            i_p, i_n = previous_obs[i.start], next_obs[i.stop - 1]
             # segment of interaction
-            p_seg, n_seg = self.segment[i_p], self.segment[i_n]
+            p_seg, n_seg = segment[i_p], segment[i_n]
             # Where segment are called
             if i_p != -1:
                 add_seg(p_seg, seg)
@@ -395,6 +399,26 @@ def map_segment(self, method, y, same=True, **kw):
             out = array(out)
         return out
 
+    def map_network(self, method, y, same=True, **kw):
+        if same:
+            out = empty(y.shape, **kw)
+        else:
+            out = list()
+        for i, b0, b1 in self.iter_on(self.track):
+            res = method(y[i])
+            if same:
+                out[i] = res
+            else:
+                if isinstance(i, slice):
+                    if i.start == i.stop:
+                        continue
+                elif len(i) == 0:
+                    continue
+                out.append(res)
+        if not same:
+            out = array(out)
+        return out
+
     def scatter_timeline(
         self,
         ax,
@@ -410,7 +434,7 @@ def scatter_timeline(
         Must be call on only one network
         """
         self.only_one_network()
-        y = (self.segment if yfield is None else self[yfield]) * yfactor
+        y = (self.segment if yfield is None else self.parse_varname(yfield)) * yfactor
         if method == "all":
             pass
         else:
@@ -536,23 +560,25 @@ def segment_track_array(self):
     def birth_event(self):
         # FIXME how to manage group 0
         indices = list()
+        previous_obs = self.previous_obs
         for i, _, _ in self.iter_on(self.segment_track_array):
             nb = i.stop - i.start
             if nb == 0:
                 continue
-            i_p = self.previous_obs[i.start]
+            i_p = previous_obs[i.start]
             if i_p == -1:
                 indices.append(i.start)
         return self.extract_event(list(set(indices)))
 
     def death_event(self):
         # FIXME how to manage group 0
         indices = list()
+        next_obs = self.next_obs
         for i, _, _ in self.iter_on(self.segment_track_array):
             nb = i.stop - i.start
             if nb == 0:
                 continue
-            i_n = self.next_obs[i.stop - 1]
+            i_n = next_obs[i.stop - 1]
             if i_n == -1:
                 indices.append(i.stop - 1)
         return self.extract_event(list(set(indices)))
@@ -567,16 +593,16 @@ def merging_event(self, triplet=False):
         if triplet:
             idx_m0_stop = list()
             idx_m0 = list()
-
+        next_obs, previous_obs = self.next_obs, self.previous_obs
         for i, _, _ in self.iter_on(self.segment_track_array):
             nb = i.stop - i.start
             if nb == 0:
                 continue
-            i_n = self.next_obs[i.stop - 1]
+            i_n = next_obs[i.stop - 1]
             if i_n != -1:
                 if triplet:
                     idx_m0_stop.append(i.stop - 1)
-                    idx_m0.append(self.previous_obs[i_n])
+                    idx_m0.append(previous_obs[i_n])
                 idx_m1.append(i_n)
 
         if triplet:
@@ -598,15 +624,16 @@ def spliting_event(self, triplet=False):
         if triplet:
             idx_s1_start = list()
             idx_s1 = list()
+        next_obs, previous_obs = self.next_obs, self.previous_obs
         for i, _, _ in self.iter_on(self.segment_track_array):
             nb = i.stop - i.start
             if nb == 0:
                 continue
-            i_p = self.previous_obs[i.start]
+            i_p = previous_obs[i.start]
             if i_p != -1:
                 if triplet:
                     idx_s1_start.append(i.start)
-                    idx_s1.append(self.next_obs[i_p])
+                    idx_s1.append(next_obs[i_p])
                 idx_s0.append(i_p)
         if triplet:
             return (
@@ -700,32 +727,38 @@ def plot(self, ax, ref=None, color_cycle=None, **kwargs):
             j += 1
         return mappables
 
-    def remove_dead_end(self, nobs=3, recursive=0, mask=None):
+    def remove_dead_end(self, nobs=3, ndays=0, recursive=0, mask=None):
         """
         .. warning::
             It will remove short segment which splits than merges with same segment
         """
-        self.only_one_network()
         segments_keep = list()
-        connexions = self.connexions()
-        for i, b0, _ in self.iter_on("segment"):
-            nb = i.stop - i.start
+        connexions = self.connexions(multi_network=True)
+        t = self.time
+        for i, b0, _ in self.iter_on(self.segment_track_array):
             if mask and mask[i].any():
                 segments_keep.append(b0)
                 continue
-            if nb < nobs and len(connexions.get(b0, tuple())) < 2:
+            nb = i.stop - i.start
+            dt = t[i.stop - 1] - t[i.start]
+            if (nb < nobs or dt < ndays) and len(connexions.get(b0, tuple())) < 2:
                 continue
             segments_keep.append(b0)
         if recursive > 0:
-            return self.extract_segment(segments_keep).remove_dead_end(
-                nobs, recursive - 1
+            return self.extract_segment(segments_keep, absolute=True).remove_dead_end(
+                nobs, ndays, recursive - 1
             )
-        return self.extract_segment(segments_keep)
+        return self.extract_segment(segments_keep, absolute=True)
 
-    def extract_segment(self, segments):
+    def extract_segment(self, segments, absolute=False):
         mask = ones(self.shape, dtype="bool")
-        for i, b0, b1 in self.iter_on("segment"):
-            if b0 not in segments:
+        segments = array(segments)
+        values = self.segment_track_array if absolute else "segment"
+        keep = ones(values.max() + 1, dtype="bool")
+        v = unique(values)
+        keep[v] = in1d(v, segments)
+        for i, b0, b1 in self.iter_on(values):
+            if not keep[b0]:
                 mask[i] = False
         return self.extract_with_mask(mask)
 
@@ -929,3 +962,8 @@ def new_numbering(segs):
             s0 = segs[i]
             j += 1
         segs[i] = j
+
+
+@njit(cache=True)
+def ptp(values):
+    return values.max() - values.min()