-lazy cube management

-event statistics

Author: AntSimi

src/py_eddy_tracker/dataset/grid.py

@@ -304,14 +304,26 @@ def __init__(
                 "We assume pixel position of grid is centered for %s", filename
             )
         if not unset:
-            self.load_general_features()
-            self.load()
+            self.populate()
 
     def populate(self):
         if self.dimensions is None:
             self.load_general_features()
             self.load()
 
+    def clean(self):
+        self.dimensions = None
+        self.variables_description = None
+        self.global_attrs = None
+        self.x_c = None
+        self.y_c = None
+        self.x_bounds = None
+        self.y_bounds = None
+        self.x_dim = None
+        self.y_dim = None
+        self.contours = None
+        self.vars = dict()
+
     @property
     def is_centered(self):
         """Give True if pixel is described with its center's position or
@@ -429,7 +441,7 @@ def c_to_bounds(c):
     def setup_coordinates(self):
         x_name, y_name = self.coordinates
         if self.is_centered:
-            logger.info("Grid center")
+            # logger.info("Grid center")
             self.x_c = self.vars[x_name].astype("float64")
             self.y_c = self.vars[y_name].astype("float64")
 
@@ -1968,14 +1980,21 @@ def interp(self, grid_name, lons, lats, method="bilinear"):
             self.x_c, self.y_c, g, m, lons, lats, nearest=method == "nearest"
         )
 
-    def uv_for_advection(self, u_name, v_name, time_step=600, backward=False, factor=1):
+    def uv_for_advection(self, u_name=None, v_name=None, time_step=600, h_name=None, backward=False, factor=1):
         """
         Get U,V to be used in degrees with precomputed time step
 
-        :param str,array u_name: U field to advect obs
-        :param str,array v_name: V field to advect obs
+        :param None,str,array u_name: U field to advect obs, if h_name is None
+        :param None,str,array v_name: V field to advect obs, if h_name is None
+        :param None,str,array h_name: H field to compute UV to advect obs, if u_name and v_name are None
         :param int time_step: Number of second for each advection
         """
+        if h_name is not None:
+            u_name, v_name = 'u', 'v'
+            if u_name not in self.vars:
+                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()
         # N seconds / 1 degrees in m
@@ -2318,6 +2337,14 @@ def from_netcdf_list(
             new.datasets.append((_t, d))
         return new
 
+    @property
+    def are_loaded(self):
+        return ~array([d.dimensions is None for _, d in self.datasets])
+
+    def __repr__(self):
+        nb_dataset = len(self.datasets)
+        return f"{self.are_loaded.sum()}/{nb_dataset} datasets loaded"
+
     def shift_files(self, t, filename, heigth=None, **rgd_kwargs):
         """Add next file to the list and remove the oldest"""
 
@@ -2440,36 +2467,42 @@ def filament(
             t += dt
             yield t, f_x, f_y
 
+    def reset_grids(self, N=None):
+        if N is not None:
+            m = self.are_loaded
+            if m.sum() > N:
+                for i in where(m)[0]:
+                    self.datasets[i][1].clean()
+
     def advect(
         self,
         x,
         y,
-        u_name,
-        v_name,
         t_init,
         mask_particule=None,
         nb_step=10,
         time_step=600,
         rk4=True,
+        reset_grid=None,
         **kw,
     ):
         """
         At each call it will update position in place with u & v field
 
         :param array x: Longitude of obs to move
         :param array y: Latitude of obs to move
-        :param str,array u_name: U field to advect obs
-        :param str,array v_name: V field to advect obs
         :param float t_init: time to start advection
         :param array,None mask_particule: advect only i mask is True
         :param int nb_step: Number of iteration before to release data
         :param int time_step: Number of second for each advection
         :param bool rk4: Use rk4 algorithm instead of finite difference
+        :param int reset_grid: Delete all loaded data in cube if there are more than N grid loaded
 
         :return: t,x,y position
 
         .. minigallery:: py_eddy_tracker.GridCollection.advect
         """
+        self.reset_grids(reset_grid)
         backward = kw.get("backward", False)
         if backward:
             generator = self.get_previous_time_step(t_init)
@@ -2480,9 +2513,9 @@ def advect(
             dt = nb_step * time_step
             t_step = time_step
         t0, d0 = generator.__next__()
-        u0, v0, m0 = d0.uv_for_advection(u_name, v_name, time_step, **kw)
+        u0, v0, m0 = d0.uv_for_advection(time_step=time_step, **kw)
         t1, d1 = generator.__next__()
-        u1, v1, m1 = d1.uv_for_advection(u_name, v_name, time_step, **kw)
+        u1, v1, m1 = d1.uv_for_advection(time_step=time_step, **kw)
         t0 = t0 * 86400
         t1 = t1 * 86400
         t = t_init * 86400
@@ -2497,7 +2530,7 @@ def advect(
                 t0, u0, v0, m0 = t1, u1, v1, m1
                 t1, d1 = generator.__next__()
                 t1 = t1 * 86400
-                u1, v1, m1 = d1.uv_for_advection(u_name, v_name, time_step, **kw)
+                u1, v1, m1 = d1.uv_for_advection(time_step=time_step, **kw)
             w = 1 - (arange(t, t + dt, t_step) - t0) / (t1 - t0)
             half_w = t_step / 2.0 / (t1 - t0)
             advect_(

src/py_eddy_tracker/observations/groups.py

@@ -114,7 +114,7 @@ def particle_candidate_step(
     # Advect particles
     kw = dict(nb_step=day_fraction, time_step=86400 / day_fraction)
     p = c.advect(x, y, t_init=t_start, **kwargs, **kw)
-    for _ in range(dt):
+    for _ in range(abs(dt)):
         _, x, y = p.__next__()
     m = ~(isnan(x) + isnan(y))
     i_end = full(x.shape, -1, dtype="i4")
@@ -352,7 +352,7 @@ def keep_tracks_by_date(self, date, nb_days):
         return self.extract_with_mask(mask)
 
     def particle_candidate_atlas(
-        self, cube, space_step, dt, start_intern=False, end_intern=False, **kwargs
+        self, cube, space_step, dt, start_intern=False, end_intern=False, callback_coherence=None, finalize_coherence=None, **kwargs
     ):
         """Select particles within eddies, advect them, return target observation and associated percentages
 
@@ -361,7 +361,9 @@ def particle_candidate_atlas(
         :param int dt: duration of advection
         :param bool start_intern: Use intern or extern contour at injection, defaults to False
         :param bool end_intern: Use intern or extern contour at end of advection, defaults to False
-        :params dict kwargs: dict of params given to advection
+        :param dict kwargs: dict of params given to advection
+        :param func callback_coherence: if None we will use cls.fill_coherence
+        :param func finalize_coherence: to apply on results of callback_coherence
         :return (np.array,np.array): return target index and percent associate
         """
         t_start, t_end = int(self.period[0]), int(self.period[1])
@@ -374,23 +376,62 @@ def particle_candidate_atlas(
         i_target, pct = full(shape, -1, dtype="i4"), zeros(shape, dtype="i1")
         # Backward or forward
         times = arange(t_start, t_end - dt) if dt > 0 else arange(t_start + dt, t_end)
+
+        if callback_coherence is None:
+            callback_coherence = self.fill_coherence
+        indexs = dict()
+        results = list()
+        kw_coherence = dict(space_step=space_step, dt=dt, c=cube)
+        kw_coherence.update(kwargs)
         for t in times:
+            logger.info("Coherence for time step : %s in [%s:%s]", t, times[0], times[-1])
             # Get index for origin
             i = t - t_start
             indexs0 = i_sort[i_start[i] : i_end[i]]
             # Get index for end
             i = t + dt - t_start
             indexs1 = i_sort[i_start[i] : i_end[i]]
-            # Get contour data
-            contours0 = [self[label][indexs0] for label in self.intern(start_intern)]
-            contours1 = [self[label][indexs1] for label in self.intern(end_intern)]
-            # Get local result
-            i_target_, pct_ = particle_candidate_step(
-                t, contours0, contours1, space_step, dt, cube, **kwargs
-            )
-            # Merge result
-            m = i_target_ != -1
-            i_target_[m] = indexs1[i_target_[m]]
-            i_target[indexs0] = i_target_
-            pct[indexs0] = pct_
+            if indexs0.size == 0 or indexs1.size == 0:
+                continue
+
+            results.append(callback_coherence(self, i_target, pct, indexs0, indexs1, start_intern, end_intern, t_start=t, **kw_coherence))
+            indexs[results[-1]] = indexs0, indexs1
+
+        if finalize_coherence is not None:
+            finalize_coherence(results, indexs, i_target, pct)
         return i_target, pct
+
+    @classmethod
+    def fill_coherence(cls, network, i_targets, percents, i_origin, i_end, start_intern, end_intern, **kwargs):
+        """_summary_
+
+        :param array i_targets: global target
+        :param array percents:
+        :param array i_origin: indices of origins
+        :param array i_end: indices of ends
+        :param bool start_intern: Use intern or extern contour at injection
+        :param bool end_intern: Use intern or extern contour at end of advection
+        """
+        # Get contour data
+        contours_start = [network[label][i_origin] for label in cls.intern(start_intern)]
+        contours_end = [network[label][i_end] for label in cls.intern(end_intern)]
+        # Compute local coherence
+        i_local_targets, local_percents = particle_candidate_step(contours_start=contours_start, contours_end=contours_end,**kwargs)
+        # Store
+        cls.merge_particle_result(i_targets, percents, i_local_targets, local_percents, i_origin, i_end)
+    
+    @staticmethod
+    def merge_particle_result(i_targets, percents, i_local_targets, local_percents, i_origin, i_end):
+        """Copy local result in merged result with global indexation
+
+        :param array i_targets: global target
+        :param array percents: 
+        :param array i_local_targets: local index target
+        :param array local_percents: 
+        :param array i_origin: indices of origins
+        :param array i_end: indices of ends
+        """
+        m = i_local_targets != -1
+        i_local_targets[m] = i_end[i_local_targets[m]]
+        i_targets[i_origin] = i_local_targets
+        percents[i_origin] = local_percents

src/py_eddy_tracker/observations/network.py

@@ -7,7 +7,7 @@
 import time
 
 import netCDF4
-from numba import njit
+from numba import njit, types as nb_types
 from numba.typed import List
 from numpy import (
     arange,
@@ -23,6 +23,8 @@
     unique,
     where,
     zeros,
+    percentile,
+    nan
 )
 import zarr
 
@@ -1743,6 +1745,7 @@ def segment_coherence_forward(
         step_mesh=1.0 / 50,
         contour_start="speed",
         contour_end="speed",
+        **kwargs,
     ):
 
         """
@@ -1801,6 +1804,7 @@ def date2file(julian_day):
                 n_days=n_days,
                 contour_start=contour_start,
                 contour_end=contour_end,
+                **kwargs
             )
             logger.info(
                 (
@@ -1996,7 +2000,69 @@ def build_dataset(self, group, raw_data=True):
             print()
         eddies.track[new_i] = group
         return eddies
+    
+@njit(cache=True)
+def get_percentile_on_following_obs(i, indexs, percents, follow_obs, t, segment, i_target, window, q=50, nb_min=1):
+    """Get stat on a part of segment close of an event
+
+    :param int i: index to follow
+    :param array indexs: indexs from coherence
+    :param array percents: percent from coherence
+    :param array[int] follow_obs: give index for the following observation
+    :param array t: time for each observation
+    :param array segment: segment for each observation
+    :param int i_target: index of target
+    :param int window: time window of search
+    :param int q: Percentile from 0 to 100, defaults to 50
+    :param int nb_min: Number minimal of observation to provide statistics, defaults to 1
+    :return float : return statistic
+    """
+    last_t, segment_follow = t[i], segment[i]
+    segment_target = segment[i_target]
+    percent_target = empty(window, dtype=percents.dtype)
+    j = 0
+    while abs(last_t - t[i]) < window and i != -1 and segment_follow == segment[i]:
+        # Iter on primary & secondary
+        for index, percent in zip(indexs[i], percents[i]):
+            if index != -1 and segment[index] == segment_target:
+                percent_target[j] = percent
+                j += 1
+        i = follow_obs[i]
+    if j < nb_min:
+        return nan
+    return percentile(percent_target[:j], q)
 
+@njit(cache=True)
+def get_percentile_around_event(i, i1, i2, ind, pct, follow_obs, t, segment, window=10, follow_parent=False, q=50, nb_min=1):
+    """Get stat around event
+
+    :param array[int] i: Indexs of target
+    :param array[int] i1: Indexs of primary origin
+    :param array[int] i2: Indexs of secondary origin
+    :param array ind: indexs from coherence
+    :param array pct: percent from coherence
+    :param array[int] follow_obs: give index for the following observation
+    :param array t: time for each observation
+    :param array segment: segment for each observation
+    :param int window: time window of search, defaults to 10
+    :param bool follow_parent: Follow parent instead of child, defaults to False
+    :param int q: Percentile from 0 to 100, defaults to 50
+    :param int nb_min: Number minimal of observation to provide statistics, defaults to 1
+    :return (array,array) : statistic for each event
+    """
+    stat1 = empty(i.size, dtype=nb_types.float32)
+    stat2 = empty(i.size, dtype=nb_types.float32)
+    # iter on event
+    for j, (i_, i1_, i2_) in enumerate(zip(i, i1, i2)):
+        if follow_parent:
+            # We follow parent
+            stat1[j] = get_percentile_on_following_obs(i_, ind, pct, follow_obs, t, segment, i1_, window, q, nb_min)
+            stat2[j] = get_percentile_on_following_obs(i_, ind, pct, follow_obs, t, segment, i2_, window, q, nb_min)
+        else:
+            # We follow child
+            stat1[j] = get_percentile_on_following_obs(i1_, ind, pct, follow_obs, t, segment, i_, window, q, nb_min)
+            stat2[j] = get_percentile_on_following_obs(i2_, ind, pct, follow_obs, t, segment, i_, window, q, nb_min)
+    return stat1, stat2
 
 @njit(cache=True)
 def get_next_index(gr):

src/py_eddy_tracker/observations/observation.py

@@ -1656,6 +1656,31 @@ def create_variable(
         except ValueError:
             logger.warning("Data is empty")
 
+    @staticmethod
+    def get_filters_zarr(name):
+        """Get filters to store in zarr for known variable
+
+        :param str name: private variable name
+        :return list: filters list
+        """
+        content = VAR_DESCR.get(name)
+        filters = list()
+        store_dtype = content["output_type"]
+        scale_factor, add_offset = content.get("scale_factor", None), content.get("add_offset", None)
+        if scale_factor is not None or add_offset is not None:
+            if add_offset is None:
+                add_offset = 0
+            filters.append(
+                zarr.FixedScaleOffset(
+                    offset=add_offset,
+                    scale=1 / scale_factor,
+                    dtype=content["nc_type"],
+                    astype=store_dtype,
+                )
+            )
+        filters.extend(content.get("filters", []))
+        return filters
+
     def create_variable_zarr(
         self,
         handler_zarr,