add indexer

Author: AntSimi

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))

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]]

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],

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

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=}"