Skip to content

Commit 70b7718

Browse files
committed
Replace numpy function by numba one
1 parent 0c6df51 commit 70b7718

File tree

2 files changed

+26
-8
lines changed

2 files changed

+26
-8
lines changed

src/py_eddy_tracker/dataset/grid.py

Lines changed: 25 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -711,16 +711,14 @@ def eddy_identification(
711711
i_x_in, i_y_in = contour.pixels_in(self)
712712

713713
# Check if pixels in contour are masked
714-
if data.mask[i_x_in, i_y_in].any():
714+
if has_masked_value(data.mask, i_x_in, i_y_in):
715715
continue
716716

717717
# Test to know cyclone or anticyclone
718-
if anticyclonic_search:
719-
if (data[i_x_in, i_y_in] < cvalues).any():
720-
continue
721-
else:
722-
if (data[i_x_in, i_y_in] > cvalues).any():
723-
continue
718+
if has_value(
719+
data, i_x_in, i_y_in, cvalues, below=anticyclonic_search
720+
):
721+
continue
724722

725723
# FIXME : Maybe limit max must be replace with a maximum of surface
726724
if (
@@ -1857,3 +1855,23 @@ def _nearest_grd_indice(x, y, x0, y0, xstep, ystep):
18571855
numba_types.int32(round(((x - x0[0]) % 360.0) / xstep)),
18581856
numba_types.int32(round((y - y0[0]) / ystep)),
18591857
)
1858+
1859+
1860+
@njit(cache=True)
1861+
def has_masked_value(grid, i_x, i_y):
1862+
for i, j in zip(i_x, i_y):
1863+
if grid[i, j]:
1864+
return True
1865+
return False
1866+
1867+
1868+
@njit(cache=True)
1869+
def has_value(grid, i_x, i_y, value, below=False):
1870+
for i, j in zip(i_x, i_y):
1871+
if below:
1872+
if grid[i, j] < value:
1873+
return True
1874+
else:
1875+
if grid[i, j] > value:
1876+
return True
1877+
return False

tests/test_id.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,5 +10,5 @@
1010
def test_id():
1111
g.add_uv("adt")
1212
a, c = g.eddy_identification("adt", "u", "v", datetime(2019, 2, 23))
13-
assert len(a) == 35
13+
assert len(a) == 36
1414
assert len(c) == 36

0 commit comments

Comments
 (0)