Skip to content

Commit f7f820d

Browse files
committed
Move out argsort from numba to speed up
1 parent a39435d commit f7f820d

File tree

1 file changed

+6
-2
lines changed

1 file changed

+6
-2
lines changed

src/py_eddy_tracker/generic.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -615,7 +615,6 @@ def build_circle(x0, y0, r):
615615
return x_norm * r + x0, y_norm * r + y0
616616

617617

618-
@njit(cache=True)
619618
def window_index(x, x0, half_window=1):
620619
"""
621620
Give for a fixed half_window each start and end index for each x0, in
@@ -626,7 +625,12 @@ def window_index(x, x0, half_window=1):
626625
:param float half_window: half window
627626
"""
628627
# Sort array, bounds will be sort also
629-
i_ordered = x.argsort()
628+
i_ordered = x.argsort(kind="mergesort")
629+
return window_index_(x, i_ordered, x0, half_window)
630+
631+
632+
@njit(cache=True)
633+
def window_index_(x, i_ordered, x0, half_window=1):
630634
nb_x, nb_pt = x.size, x0.size
631635
first_index = empty(nb_pt, dtype=i_ordered.dtype)
632636
last_index = empty(nb_pt, dtype=i_ordered.dtype)

0 commit comments

Comments
 (0)