Add tool to :

- Get azimuth track
- Count event by track
- Use array to compute stat

Author: AntSimi

src/py_eddy_tracker/observations/observation.py

@@ -1532,20 +1532,26 @@ def filled(
         c.norm = Normalize(vmin=vmin, vmax=vmax)
         return c
 
-    def bins_stat(self, name, bins=None):
+    def bins_stat(self, xname, bins=None, yname=None, method=None):
         """
-        :param str name: var which will be use
+        :param str,array xname: var which will be use
         :param array, None bins: bins to perform statistics,if None method will get min and max of variable
+        :param None,str yname: var which will be use to apply method
+        :param None,str method: If None method count in each bin, could also do mean, std
         :return: x array and y array
         :rtype: array,array
 
         .. minigallery:: py_eddy_tracker.EddiesObservations.bins_stat
         """
-        v = self[name]
+        v = self[xname] if isinstance(xname, str) else xname
         if bins is None:
             bins = arange(v.min(), v.max() + 2)
         y, x = hist_numba(v, bins=bins)
         x = (x[1:] + x[:-1]) / 2
+        if method == "mean":
+            y_v = self[yname]
+            y_, _ = histogram(v, bins=bins, weights=y_v)
+            y = y_ / y
         return x, y
 
     def display(
@@ -1679,12 +1685,13 @@ def grid_count(self, bins, intern=False, center=False):
             grid.mask = grid == 0
         return regular_grid
 
-    def grid_stat(self, bins, varname):
+    def grid_stat(self, bins, varname, data=None):
         """
         Compute mean of eddies in each bin
 
         :param (numpy.array,numpy.array) bins: bins to compute count
         :param str varname: name of variable to compute mean
+        :param array data: Array used to compute stat if defined
         :return: return grid of mean
         :rtype: py_eddy_tracker.dataset.grid.RegularGridDataset
 
@@ -1693,8 +1700,14 @@ def grid_stat(self, bins, varname):
         x_bins, y_bins = arange(*bins[0]), arange(*bins[1])
         x0 = bins[0][0]
         x, y = (self.longitude - x0) % 360 + x0, self.latitude
-        sum_obs = histogram2d(x, y, (x_bins, y_bins), weights=self[varname])[0]
-        nb_obs = histogram2d(x, y, (x_bins, y_bins))[0]
+        data = self[varname] if data is None else data
+        if hasattr(data, 'mask'):
+            m = ~data.mask
+            sum_obs = histogram2d(x[m], y[m], (x_bins, y_bins), weights=data[m])[0]
+            nb_obs = histogram2d(x[m], y[m], (x_bins, y_bins))[0]
+        else:
+            sum_obs = histogram2d(x, y, (x_bins, y_bins), weights=data)[0]
+            nb_obs = histogram2d(x, y, (x_bins, y_bins))[0]
         from ..dataset.grid import RegularGridDataset
 
         regular_grid = RegularGridDataset.with_array(

src/py_eddy_tracker/observations/tracking.py

@@ -15,6 +15,12 @@
     array,
     median,
     histogram,
+    sin,
+    cos,
+    arctan2,
+    radians,
+    degrees,
+    tan,
 )
 from datetime import datetime, timedelta
 from numba import njit
@@ -236,18 +242,30 @@ def extract_with_period(self, period, **kwargs):
             mask *= self.time <= (dataset_period[1] + p_max)
         return self.extract_with_mask(mask, **kwargs)
 
-    def get_azimuth(self):
+    def get_azimuth(self, equatorward=False):
         """
         Return azimuth for each tracks.
 
         Azimuth is compute with first and last observation
 
+        :param bool equatorward: If True,  Poleward are positive and equatorward negative
         :rtype: array
         """
-        i0 = self.index_from_track
-        i1 = i0 - 1 + self.nb_obs_by_track
+        i0, nb = self.index_from_track, self.nb_obs_by_track
+        i0 = i0[nb != 0]
+        i1 = i0 - 1 + nb[nb != 0]
         lat0, lon0 = self.latitude[i0], self.longitude[i0]
         lat1, lon1 = self.latitude[i1], self.longitude[i1]
+        lat0, lon0 = radians(lat0), radians(lon0)
+        lat1, lon1 = radians(lat1), radians(lon1)
+        dlon = lon1 - lon0
+        x = cos(lat0) * sin(lat1) - sin(lat0) * cos(lat1) * cos(dlon)
+        y = sin(dlon) * cos(lat1)
+        azimuth = degrees(arctan2(y, x)) + 90
+        if equatorward:
+            south = lat0 < 0
+            azimuth[south] *= -1
+        return azimuth
 
     def get_mask_from_id(self, tracks):
         mask = zeros(self.tracks.shape, dtype=bool_)
@@ -269,6 +287,19 @@ def compute_index(self):
             compute_index(self.tracks, self.__first_index_of_track, self.__obs_by_track)
             logger.debug("... OK")
 
+    def count_by_track(self, mask):
+        """
+        Count by track
+
+        :param array[bool] mask: Mask of boolean count +1 if true
+        :return: Return count by track
+        :rtype: array
+        """
+        s = self.tracks.max() + 1
+        obs_by_track = zeros(s, "i4")
+        count_by_track(self.tracks, mask, obs_by_track)
+        return obs_by_track
+
     @property
     def index_from_track(self):
         self.compute_index()
@@ -599,6 +630,13 @@ def compute_index(tracks, index, number):
         previous_track = track
 
 
+@njit(cache=True)
+def count_by_track(tracks, mask, number):
+    for track, test in zip(tracks, mask):
+        if test:
+            number[track] += 1
+
+
 @njit(cache=True)
 def compute_mask_from_id(tracks, first_index, number_of_obs, mask):
     for track in tracks: