From 141e941c28249cb95b63deb094de534a98f27da4 Mon Sep 17 00:00:00 2001 From: "A. Delepoulle" <36040805+AntSimi@users.noreply.github.com> Date: Thu, 26 Jun 2025 09:23:49 +0000 Subject: [PATCH] add indexer --- src/py_eddy_tracker/dataset/grid.py | 25 ++++++++-- src/py_eddy_tracker/observations/groups.py | 19 ++++++++ src/py_eddy_tracker/poly.py | 8 +--- tests/test_generic.py | 48 ++++++++++++++----- tests/test_poly.py | 56 +++++++++++++++++++--- 5 files changed, 128 insertions(+), 28 deletions(-) diff --git a/src/py_eddy_tracker/dataset/grid.py b/src/py_eddy_tracker/dataset/grid.py index f15503b2..05b4b8bd 100644 --- a/src/py_eddy_tracker/dataset/grid.py +++ b/src/py_eddy_tracker/dataset/grid.py @@ -326,6 +326,9 @@ def clean(self): self.contours = None self.vars = dict() + def parse_varname(self, name): + return self.grid(name) if isinstance(name, str) else name + @property def is_centered(self): """Give True if pixel is described with its center's position or @@ -526,6 +529,8 @@ def grid(self, varname, indexs=None): .. minigallery:: py_eddy_tracker.GridDataset.grid """ + # Ensure to have acces to coordinates + self.populate() if indexs is None: indexs = dict() if varname not in self.vars: @@ -1914,6 +1919,18 @@ def init_speed_coef(self, uname="u", vname="v"): u, v = self.grid(uname), self.grid(vname) self._speed_ev = sqrt(u * u + v * v) + def text(self, ax, name, label="{v:.2f}", **kwargs): + """ + :param matplotlib.axes.Axes ax: matplotlib axes used to draw + :param str,array name: variable to display, could be an array + :param str label: string to format value + :param dict kwargs: look at :py:meth:`matplotlib.axes.Axes.pcolormesh` + """ + data = self.parse_varname(name) + for ix, x in enumerate(self.x_c): + for iy, y in enumerate(self.y_c): + ax.text(x, y, label.format(v=data[ix,iy]), **kwargs) + def display(self, ax, name, factor=1, ref=None, **kwargs): """ :param matplotlib.axes.Axes ax: matplotlib axes used to draw @@ -1926,7 +1943,7 @@ def display(self, ax, name, factor=1, ref=None, **kwargs): """ if "cmap" not in kwargs: kwargs["cmap"] = "coolwarm" - data = self.grid(name) if isinstance(name, str) else name + data = self.parse_varname(name) if ref is None: x = self.x_bounds else: @@ -1946,7 +1963,7 @@ def contour(self, ax, name, factor=1, ref=None, **kwargs): .. minigallery:: py_eddy_tracker.RegularGridDataset.contour """ - data = self.grid(name) if isinstance(name, str) else name + data = self.parse_varname(name) if ref is None: x = self.x_c else: @@ -2020,8 +2037,8 @@ def uv_for_advection( self.add_uv(h_name) self.vars.pop(h_name, None) - u = (self.grid(u_name) if isinstance(u_name, str) else u_name).copy() - v = (self.grid(v_name) if isinstance(v_name, str) else v_name).copy() + u = self.parse_varname(u_name).copy() + v = self.parse_varname(v_name).copy() # N seconds / 1 degrees in m coef = time_step * 180 / pi / self.EARTH_RADIUS * factor u *= coef / cos(radians(self.y_c)) diff --git a/src/py_eddy_tracker/observations/groups.py b/src/py_eddy_tracker/observations/groups.py index 81929e1e..58692caa 100644 --- a/src/py_eddy_tracker/observations/groups.py +++ b/src/py_eddy_tracker/observations/groups.py @@ -268,6 +268,12 @@ def get_matrix(i_start, i_end, translate_start, translate_end, i_target, pct): class GroupEddiesObservations(EddiesObservations, ABC): + __slots__ = ('indexers',) + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.indexers = dict() + @abstractmethod def fix_next_previous_obs(self): pass @@ -474,3 +480,16 @@ def merge_particle_result( i_local_targets[m] = i_end[i_local_targets[m]] i_targets[i_origin] = i_local_targets percents[i_origin] = local_percents + + def daily_time_indexer(self, t, dt=.5): + """Get all index cover by time window + + :param float t: time centered of window + :param float dt: delta +/- around t + """ + if 'time' not in self.indexers: + x0, x1 = self.period + self.indexers['time'] = x0, *window_index(self.time, np.arange(x0, x1), 0.5) + x0, i_ordered, first_index, last_index = self.indexers['time'] + j0, j1 = max(0, int(np.ceil(t - x0 - dt))), min(int(np.trunc(t - x0 + dt)), last_index.size - 1) + return i_ordered[first_index[j0]:last_index[j1]] diff --git a/src/py_eddy_tracker/poly.py b/src/py_eddy_tracker/poly.py index b5849610..ef200f23 100644 --- a/src/py_eddy_tracker/poly.py +++ b/src/py_eddy_tracker/poly.py @@ -717,12 +717,8 @@ def get_pixel_in_regular(vertices, x_c, y_c, x_start, x_stop, y_start, y_stop): :param int y_stop: north index of contour """ if x_stop < x_start: - x_ref = vertices[0, 0] - x_array = ( - (concatenate((x_c[x_start:], x_c[:x_stop])) - x_ref + 180) % 360 - + x_ref - - 180 - ) + x_ref = vertices[0, 0] - 180 + x_array = (concatenate((x_c[x_start:], x_c[:x_stop])) - x_ref) % 360 + x_ref return winding_number_grid_in_poly( x_array, y_c[y_start:y_stop], diff --git a/tests/test_generic.py b/tests/test_generic.py index ee2d7881..8d522c75 100644 --- a/tests/test_generic.py +++ b/tests/test_generic.py @@ -1,39 +1,39 @@ -from numpy import arange, array, nan, ones, zeros +import numpy as np -from py_eddy_tracker.generic import cumsum_by_track, simplify, wrap_longitude +from py_eddy_tracker.generic import cumsum_by_track, simplify, wrap_longitude, bbox_indice_regular def test_simplify(): - x = arange(10, dtype="f4") - y = zeros(10, dtype="f4") + x = np.arange(10, dtype="f4") + y = np.zeros(10, dtype="f4") # Will jump one value on two x_, y_ = simplify(x, y, precision=1) assert x_.shape[0] == 5 x_, y_ = simplify(x, y, precision=0.99) assert x_.shape[0] == 10 # check nan management - x[4] = nan + x[4] = np.nan x_, y_ = simplify(x, y, precision=1) assert x_.shape[0] == 6 - x[3] = nan + x[3] = np.nan x_, y_ = simplify(x, y, precision=1) assert x_.shape[0] == 6 - x[:4] = nan + x[:4] = np.nan x_, y_ = simplify(x, y, precision=1) assert x_.shape[0] == 3 - x[:] = nan + x[:] = np.nan x_, y_ = simplify(x, y, precision=1) assert x_.shape[0] == 0 def test_cumsum_by_track(): - a = ones(10, dtype="i4") * 2 - track = array([1, 1, 2, 2, 2, 2, 44, 44, 44, 48]) + a = np.ones(10, dtype="i4") * 2 + track = np.array([1, 1, 2, 2, 2, 2, 44, 44, 44, 48]) assert (cumsum_by_track(a, track) == [2, 4, 2, 4, 6, 8, 2, 4, 6, 2]).all() def test_wrapping(): - y = x = arange(-5, 5, dtype="f4") + y = x = np.arange(-5, 5, dtype="f4") x_, _ = wrap_longitude(x, y, ref=-10) assert (x_ == x).all() x_, _ = wrap_longitude(x, y, ref=1) @@ -49,3 +49,29 @@ def test_wrapping(): # x %= 360 # x_, _ = wrap_longitude(x, y, ref=-10, cut=True) # assert x.size == x_.size + + +def test_bbox_fit_contour(): + v= np.array([ + [200.6, 10.5], + [203.4, 10.4], + [204.4, 9.6], + [198.4, 9.3], + [200.6, 10.5], + ]) + xc0, yc0 = v.min(axis=0) + xc1, yc1 = v.max(axis=0) + for step in [1, .5, .2, .133333, .125,.1,.07, .025]: + x_c = np.arange(0, 360, step) + y_c = np.arange(-80, 80, step) + xstep, ystep = x_c[1] - x_c[0], y_c[1] - y_c[0] + x0, y0 = x_c - xstep / 2.0, y_c - ystep / 2.0 + nb_x = x_c.shape[0] + (x_start, x_stop), (y_start, y_stop) = bbox_indice_regular(v, x0, y0, xstep, ystep, 0, True, nb_x) + i_x = np.where((x_c >= xc0) * (x_c <= xc1))[0] + i_y = np.where((y_c >= yc0) * (y_c <= yc1))[0] + (x_start_raw, x_stop_raw), (y_start_raw, y_stop_raw) = (i_x.min(), i_x.max()), (i_y.min(), i_y.max()) + assert x_start == x_start_raw + assert x_stop == x_stop_raw + 1 + assert y_start == y_start_raw + assert y_stop == y_stop_raw + 1 \ No newline at end of file diff --git a/tests/test_poly.py b/tests/test_poly.py index a780f64d..0c102d93 100644 --- a/tests/test_poly.py +++ b/tests/test_poly.py @@ -1,4 +1,4 @@ -from numpy import array, pi, roll +import numpy as np from pytest import approx from py_eddy_tracker.poly import ( @@ -7,11 +7,13 @@ get_convex_hull, poly_area_vertice, visvalingam, + get_pixel_in_regular, ) +from py_eddy_tracker.generic import bbox_indice_regular # Vertices for next test -V = array(((2, 2, 3, 3, 2), (-10, -9, -9, -10, -10))) -V_concave = array(((2, 2, 2.5, 3, 3, 2), (-10, -9, -9.5, -9, -10, -10))) +V = np.array(((2, 2, 3, 3, 2), (-10, -9, -9, -10, -10))) +V_concave = np.array(((2, 2, 2.5, 3, 3, 2), (-10, -9, -9.5, -9, -10, -10))) def test_poly_area(): @@ -23,7 +25,7 @@ def test_fit_circle(): assert x0 == approx(2.5, rel=1e-10) assert y0 == approx(-9.5, rel=1e-10) assert r == approx(2**0.5 / 2, rel=1e-10) - assert err == approx((1 - 2 / pi) * 100, rel=1e-10) + assert err == approx((1 - 2 / np.pi) * 100, rel=1e-10) def test_convex(): @@ -38,8 +40,8 @@ def test_convex_hull(): def test_visvalingam(): - x = array([1, 2, 3, 4, 5, 6.75, 6, 1]) - y = array([-0.5, -1.5, -1, -1.75, -1, -1, -0.5, -0.5]) + x = np.array([1, 2, 3, 4, 5, 6.75, 6, 1]) + y = np.array([-0.5, -1.5, -1, -1.75, -1, -1, -0.5, -0.5]) x_target = [1, 2, 3, 4, 6, 1] y_target = [-0.5, -1.5, -1, -1.75, -0.5, -0.5] x_, y_ = visvalingam(x, y, 6) @@ -48,6 +50,46 @@ def test_visvalingam(): x_, y_ = visvalingam(x[:-1], y[:-1], 6) assert (x_target == x_).all() assert (y_target == y_).all() - x_, y_ = visvalingam(roll(x, 2), roll(y, 2), 6) + x_, y_ = visvalingam(np.roll(x, 2), np.roll(y, 2), 6) assert (x_target[:-1] == x_[1:]).all() assert (y_target[:-1] == y_[1:]).all() + + +def test_pixel_in_contour(): + xmin, xmax, ymin, ymax = -1, 31, -2, 11 + v = np.array([ + [xmin, ymax], + [xmax, ymax], + [xmax, ymin], + [xmin, ymin], + [xmin, ymax], + ]) + xstep = ystep = 7.5 + # Global grid + for x0_ in [-10, 0]: + x0, y0 = (x0_, x0_ + 360), (-10, 20) + x_c, y_c = np.arange(*x0, xstep), np.arange(*y0, ystep) + (x_start, x_stop), (y_start, y_stop) = bbox_indice_regular(v, x0, y0, xstep, ystep, 1, True, int(360/5)) + i, j = get_pixel_in_regular(v, x_c, y_c, x_start, x_stop, y_start, y_stop) + assert (x_c[i] < xmax).all(), f"{x0=}" + assert (x_c[i] > xmin).all(), f"{x0=}" + + # Regional grid + # contour over two bounds + x0, y0 = (20, 370), (-10, 20) + x_c, y_c = np.arange(*x0, xstep), np.arange(*y0, ystep) + (x_start, x_stop), (y_start, y_stop) = bbox_indice_regular(v, x0, y0, xstep, ystep, 1, False, int(360/5)) + i, j = get_pixel_in_regular(v, x_c, y_c, x_start, x_stop, y_start, y_stop) + assert (x_c[i] < xmax).all(), f"{x0=}" + assert (x_c[i] > xmin).all(), f"{x0=}" + + # first case grid fully cover contour, and not in second case + for x0_ in [-2, -.5]: + x0, y0 = (x0_, 100), (-10, 20) + x_c, y_c = np.arange(*x0, xstep), np.arange(*y0, ystep) + (x_start, x_stop), (y_start, y_stop) = bbox_indice_regular(v, x0, y0, xstep, ystep, 1, False, int(360/5)) + i, j = get_pixel_in_regular(v, x_c, y_c, x_start, x_stop, y_start, y_stop) + assert (x_c[i] < xmax).all(), f"{x0=}" + assert (x_c[i] > xmin).all(), f"{x0=}" + assert (y_c[j] < ymax).all(), f"{y0=}" + assert (y_c[j] > ymin).all(), f"{y0=}" \ No newline at end of file