Skip to content

Commit 08f937c

Browse files
committed
Speed up pixel counting
1 parent 427e4ca commit 08f937c

File tree

2 files changed

+86
-16
lines changed

2 files changed

+86
-16
lines changed

src/py_eddy_tracker/dataset/grid.py

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1198,10 +1198,23 @@ def bbox_indice(self, vertices):
11981198
)
11991199

12001200
def get_pixels_in(self, contour):
1201-
(x_start, x_stop), (y_start, y_stop) = contour.bbox_slice
1202-
return get_pixel_in_regular(
1203-
contour.vertices, self.x_c, self.y_c, x_start, x_stop, y_start, y_stop
1204-
)
1201+
"""
1202+
Get indices of pixels in contour.
1203+
1204+
:param vertice,Path contour: Contour which enclosed some pixels
1205+
:return: Indices of grid in contour
1206+
:rtype: array[int],array[int]
1207+
"""
1208+
if isinstance(contour, BasePath):
1209+
(x_start, x_stop), (y_start, y_stop) = contour.bbox_slice
1210+
return get_pixel_in_regular(
1211+
contour.vertices, self.x_c, self.y_c, x_start, x_stop, y_start, y_stop
1212+
)
1213+
else:
1214+
(x_start, x_stop), (y_start, y_stop) = self.bbox_indice(contour)
1215+
return get_pixel_in_regular(
1216+
contour, self.x_c, self.y_c, x_start, x_stop, y_start, y_stop
1217+
)
12051218

12061219
def normalize_x_indice(self, indices):
12071220
return indices % self.x_size

src/py_eddy_tracker/observations/observation.py

Lines changed: 69 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1912,24 +1912,28 @@ def grid_count(self, bins, intern=False, center=False, filter=slice(None)):
19121912
),
19131913
centered=True,
19141914
)
1915-
debug_active = logger.getEffectiveLevel() == logging.DEBUG
19161915
if center:
19171916
x, y = (self.longitude[filter] - x0) % 360 + x0, self.latitude[filter]
19181917
grid[:] = histogram2d(x, y, (x_bins, y_bins))[0]
19191918
grid.mask = grid.data == 0
19201919
else:
19211920
x_ref = ((self.longitude[filter] - x0) % 360 + x0 - 180).reshape(-1, 1)
1922-
nb = x_ref.shape[0]
1923-
for i_, (x, y_) in enumerate(
1924-
zip(self[x_name][filter], self[y_name][filter])
1925-
):
1926-
x_ = (x - x_ref[i_]) % 360 + x_ref[i_]
1927-
if debug_active and i_ % 10000 == 0:
1928-
print(f"{i_}/{nb}", end="\r")
1929-
i, j = BasePath(create_vertice(x_, y_)).pixels_in(regular_grid)
1930-
grid_count_(grid, i, j)
1931-
if debug_active:
1932-
print()
1921+
x_contour, y_contour = self[x_name][filter], self[y_name][filter]
1922+
grid_count_pixel_in(
1923+
grid,
1924+
x_contour,
1925+
y_contour,
1926+
x_ref,
1927+
regular_grid.x_bounds,
1928+
regular_grid.y_bounds,
1929+
regular_grid.xstep,
1930+
regular_grid.ystep,
1931+
regular_grid.N,
1932+
regular_grid.is_circular(),
1933+
regular_grid.x_size,
1934+
regular_grid.x_c,
1935+
regular_grid.y_c,
1936+
)
19331937
grid.mask = grid == 0
19341938
return regular_grid
19351939

@@ -2083,6 +2087,59 @@ def grid_count_(grid, i, j):
20832087
grid[i_, j_] += 1
20842088

20852089

2090+
@njit(cache=True)
2091+
def grid_count_pixel_in(
2092+
grid,
2093+
x,
2094+
y,
2095+
x_ref,
2096+
x_bounds,
2097+
y_bounds,
2098+
xstep,
2099+
ystep,
2100+
N,
2101+
is_circular,
2102+
x_size,
2103+
x_c,
2104+
y_c,
2105+
):
2106+
"""
2107+
Count how many time a pixel is used.
2108+
2109+
:param array grid:
2110+
:param array x: x for all contour
2111+
:param array y: y for all contour
2112+
:param array x_ref: x reference for wrapping
2113+
:param array x_bounds: grid longitude
2114+
:param array y_bounds: grid latitude
2115+
:param float xstep: step between two longitude
2116+
:param float ystep: step between two latitude
2117+
:param int N: shift of index to enlarge window
2118+
:param bool is_circular: To know if grid is wrappable
2119+
:param int x_size: Number of longitude
2120+
:param array x_c: longitude coordinate of grid
2121+
:param array y_c: latitude coordinate of grid
2122+
"""
2123+
nb = x_ref.shape[0]
2124+
for i_ in range(nb):
2125+
x_, y_, x_ref_ = x[i_], y[i_], x_ref[i_]
2126+
x_ = (x_ - x_ref_) % 360 + x_ref_
2127+
v = create_vertice(x_, y_)
2128+
(x_start, x_stop), (y_start, y_stop) = bbox_indice_regular(
2129+
v,
2130+
x_bounds,
2131+
y_bounds,
2132+
xstep,
2133+
ystep,
2134+
N,
2135+
is_circular,
2136+
x_size,
2137+
)
2138+
2139+
i, j = get_pixel_in_regular(v, x_c, y_c, x_start, x_stop, y_start, y_stop)
2140+
grid_count_(grid, i, j)
2141+
2142+
20862143
@njit(cache=True)
20872144
def insidepoly(x_p, y_p, x_c, y_c):
20882145
"""

0 commit comments

Comments
 (0)