Use of mergesort to save time when array is alreay ordered

Author: AntSimi

examples/16_network/pet_follow_particle.py

@@ -17,6 +17,7 @@
 from py_eddy_tracker.data import get_path
 from py_eddy_tracker.dataset.grid import GridCollection
 from py_eddy_tracker.observations.network import NetworkObservations
+from py_eddy_tracker.poly import group_obs
 
 start_logger().setLevel("ERROR")
 
@@ -203,8 +204,12 @@ def get_matrix(i_start, i_end, translate_start, translate_end, i_target, pct):
 # ^^^^^^^^^^^^^^^^^^
 step = 1 / 60.0
 
-x, y = meshgrid(arange(20, 36, step), arange(30, 46, step))
+x, y = meshgrid(arange(24, 36, step), arange(31, 36, step))
 x0, y0 = x.reshape(-1), y.reshape(-1)
+# Pre-order to speed up
+_, i = group_obs(x0, y0, 1, 360)
+x0, y0 = x0[i], y0[i]
+
 t_start, t_end = n.period
 dt = 14
 

notebooks/python_module/16_network/pet_follow_particle.ipynb

@@ -26,7 +26,7 @@
       },
       "outputs": [],
       "source": [
-        "import re\n\nfrom matplotlib import colors\nfrom matplotlib import pyplot as plt\nfrom matplotlib.animation import FuncAnimation\nfrom numba import njit\nfrom numba import types as nb_types\nfrom numpy import arange, meshgrid, ones, unique, where, zeros\n\nfrom py_eddy_tracker import start_logger\nfrom py_eddy_tracker.appli.gui import Anim\nfrom py_eddy_tracker.data import get_path\nfrom py_eddy_tracker.dataset.grid import GridCollection\nfrom py_eddy_tracker.observations.network import NetworkObservations\n\nstart_logger().setLevel(\"ERROR\")"
+        "import re\n\nfrom matplotlib import colors\nfrom matplotlib import pyplot as plt\nfrom matplotlib.animation import FuncAnimation\nfrom numba import njit\nfrom numba import types as nb_types\nfrom numpy import arange, meshgrid, ones, unique, where, zeros\n\nfrom py_eddy_tracker import start_logger\nfrom py_eddy_tracker.appli.gui import Anim\nfrom py_eddy_tracker.data import get_path\nfrom py_eddy_tracker.dataset.grid import GridCollection\nfrom py_eddy_tracker.observations.network import NetworkObservations\nfrom py_eddy_tracker.poly import group_obs\n\nstart_logger().setLevel(\"ERROR\")"
       ]
     },
     {
@@ -138,7 +138,7 @@
       },
       "outputs": [],
       "source": [
-        "step = 1 / 60.0\n\nx, y = meshgrid(arange(20, 36, step), arange(30, 46, step))\nx0, y0 = x.reshape(-1), y.reshape(-1)\nt_start, t_end = n.period\ndt = 14\n\nshape = (n.obs.size, 2)\n# Forward run\ni_target_f, pct_target_f = -ones(shape, dtype=\"i4\"), zeros(shape, dtype=\"i1\")\nfor t in range(t_start, t_end - dt):\n    particle_candidate(x0, y0, c, n, t, i_target_f, pct_target_f, delta_t=dt)\n\n# Backward run\ni_target_b, pct_target_b = -ones(shape, dtype=\"i4\"), zeros(shape, dtype=\"i1\")\nfor t in range(t_start + dt, t_end):\n    particle_candidate(x0, y0, c, n, t, i_target_b, pct_target_b, delta_t=-dt)"
+        "step = 1 / 60.0\n\nx, y = meshgrid(arange(24, 36, step), arange(31, 36, step))\nx0, y0 = x.reshape(-1), y.reshape(-1)\n# Pre-order to speed up\n_, i = group_obs(x0, y0, 1, 360)\nx0, y0 = x0[i], y0[i]\n\nt_start, t_end = n.period\ndt = 14\n\nshape = (n.obs.size, 2)\n# Forward run\ni_target_f, pct_target_f = -ones(shape, dtype=\"i4\"), zeros(shape, dtype=\"i1\")\nfor t in range(t_start, t_end - dt):\n    particle_candidate(x0, y0, c, n, t, i_target_f, pct_target_f, delta_t=dt)\n\n# Backward run\ni_target_b, pct_target_b = -ones(shape, dtype=\"i4\"), zeros(shape, dtype=\"i1\")\nfor t in range(t_start + dt, t_end):\n    particle_candidate(x0, y0, c, n, t, i_target_b, pct_target_b, delta_t=-dt)"
       ]
     },
     {

src/py_eddy_tracker/poly.py

@@ -799,7 +799,7 @@ def group_obs(x, y, step, nb_x):
     i = empty(nb, dtype=numba_types.uint32)
     for k in range(nb):
         i[k] = box_index(x[k], y[k], step, nb_x)
-    return i, i.argsort()
+    return i, i.argsort(kind="mergesort")
 
 
 @njit(cache=True)
@@ -824,7 +824,9 @@ def poly_indexs(x_p, y_p, x_c, y_c):
     :param array x_c: longitude of contours
     :param array y_c: latitude of contours
     """
-    i, i_order = group_obs(x_p, y_p, 1, 360)
+    nb_x = 360
+    step = 1.0
+    i, i_order = group_obs(x_p, y_p, step, nb_x)
     nb_p = x_p.shape[0]
     nb_c = x_c.shape[0]
     indexs = -ones(nb_p, dtype=numba_types.int32)
@@ -837,17 +839,17 @@ def poly_indexs(x_p, y_p, x_c, y_c):
         x_c_min, y_c_min = x_.min(), y_.min()
         x_c_max, y_c_max = x_.max(), y_.max()
         v = create_vertice(x_, y_)
-        i0, j0 = box_indexes(x_c_min, y_c_min, 1)
-        i1, j1 = box_indexes(x_c_max, y_c_max, 1)
+        i0, j0 = box_indexes(x_c_min, y_c_min, step)
+        i1, j1 = box_indexes(x_c_max, y_c_max, step)
         # i0 could be greater than i1, (x_c is always continious) so you could have a contour over bound
         if i0 > i1:
-            i1 += 360
+            i1 += nb_x
         for i_x in range(i0, i1 + 1):
             # we force i_x in 0 360 range
-            i_x %= 360
+            i_x %= nb_x
             for i_y in range(j0, j1 + 1):
                 # Get box indices
-                i_box = i_x + 360 * i_y - i_first
+                i_box = i_x + nb_x * i_y - i_first
                 # Indice must be in table range
                 if i_box < 0 or i_box >= nb_bloc:
                     continue
@@ -860,6 +862,7 @@ def poly_indexs(x_p, y_p, x_c, y_c):
                         continue
                     if y < y_c_min:
                         continue
+                    # Normalize longitude at +-180° around x_c_min
                     x = (x_p[i_p] - x_c_min + 180) % 360 + x_c_min - 180
                     if x > x_c_max:
                         continue