Solution for issue AntSimi#11, wrap longitude if asked

Author: AntSimi

src/py_eddy_tracker/generic.py

@@ -14,6 +14,7 @@
     linspace,
     interp,
     where,
+    isnan,
 )
 from numba import njit, prange
 from numpy.linalg import lstsq
@@ -287,3 +288,50 @@ def split_line(x, y, i):
             new_y[new_j] = nan
             new_j += 1
     return new_x, new_y
+
+
+@njit(cache=True)
+def wrap_longitude(x, y, ref, cut=False):
+    if cut:
+        indexs = list()
+        nb = x.shape[0]
+        new_previous = (x[0] - ref) % 360
+        x_previous = x[0]
+        for i in range(1, nb):
+            x_ = x[i]
+            new_x = (x_ - ref) % 360
+            if not isnan(x_) and not isnan(x_previous):
+                d_new = new_x - new_previous
+                d = x_ - x_previous
+                if d != d_new:
+                    indexs.append(i)
+            x_previous, new_previous = x_, new_x
+
+        nb_indexs = len(indexs)
+        new_size = nb + nb_indexs * 3
+        out_x = empty(new_size, dtype=x.dtype)
+        out_y = empty(new_size, dtype=y.dtype)
+        i_ = 0
+        j = 0
+        for i in range(nb):
+            if j < nb_indexs and i == indexs[j]:
+                j += 1
+                cor = 360 if x[i - 1] > x[i] else -360
+                out_x[i + i_] = (x[i] - ref) % 360 + ref - cor
+                out_y[i + i_] = y[i]
+                out_x[i + i_ + 1] = nan
+                out_y[i + i_ + 1] = nan
+                out_x[i + i_ + 2] = (x[i - 1] - ref) % 360 + ref + cor
+                out_y[i + i_ + 2] = y[i - 1]
+                i_ += 3
+            out_x[i + i_] = (x[i] - ref) % 360 + ref
+            out_y[i + i_] = y[i]
+        return out_x, out_y
+
+    else:
+        nb = x.shape[0]
+        out = empty(nb, dtype=x.dtype)
+        for i in range(nb):
+            out[i] = (x[i] - ref) % 360 + ref
+        return out, y
+

src/py_eddy_tracker/observations/observation.py

@@ -58,7 +58,7 @@
 from tokenize import TokenError
 from matplotlib.path import Path as BasePath
 from .. import VAR_DESCR, VAR_DESCR_inv
-from ..generic import distance_grid, distance, flatten_line_matrix
+from ..generic import distance_grid, distance, flatten_line_matrix, wrap_longitude
 from ..poly import bbox_intersection, common_area
 
 logger = logging.getLogger("pet")
@@ -1242,7 +1242,7 @@ def set_global_attr_netcdf(self, h_nc):
         for key, item in self.global_attr.items():
             h_nc.setncattr(key, item)
 
-    def scatter(self, ax, name, ref=None, factor=1, ** kwargs):
+    def scatter(self, ax, name, ref=None, factor=1, **kwargs):
         x = self.longitude
         if ref is not None:
             x = (x - ref) % 360 + ref
@@ -1260,17 +1260,15 @@ def display(self, ax, ref=None, extern_only=False, intern_only=False, **kwargs):
         kwargs_e = kwargs.copy()
         if not extern_only:
             kwargs_e.pop("label", None)
-        if ref is None:
-            if not extern_only:
-                ax.plot(lon_s, lat_s, **kwargs)
-            if not intern_only:
-                ax.plot(lon_e, lat_e, linestyle="-.", **kwargs_e)
-        else:
-            # FIXME : ref could split eddies
-            if not extern_only:
-                ax.plot((lon_s - ref) % 360 + ref, lat_s, **kwargs)
-            if not intern_only:
-                ax.plot((lon_e - ref) % 360 + ref, lat_e, linestyle="-.", **kwargs_e)
+
+        if not extern_only:
+            if ref is not None:
+                lon_s, lat_s = wrap_longitude(lon_s, lat_s, ref, cut=True)
+            ax.plot(lon_s, lat_s, **kwargs)
+        if not intern_only:
+            if ref is not None:
+                lon_e, lat_e = wrap_longitude(lon_e, lat_e, ref, cut=True)
+            ax.plot(lon_e, lat_e, linestyle="-.", **kwargs_e)
 
     def grid_count(self, bins, intern=False, center=False):
         x_name, y_name = self.intern(intern)

src/py_eddy_tracker/observations/tracking.py

@@ -44,7 +44,7 @@
 from datetime import datetime, timedelta
 from .observation import EddiesObservations
 from numba import njit
-from ..generic import split_line
+from ..generic import split_line, wrap_longitude
 
 logger = logging.getLogger("pet")
 
@@ -365,10 +365,9 @@ def __extract_with_mask(
     def plot(self, ax, ref=None, **kwargs):
         if "label" in kwargs:
             kwargs["label"] += " (%s eddies)" % (self.nb_obs_by_track != 0).sum()
-        x = self.longitude
+        x, y = split_line(self.longitude, self.latitude, self.tracks)
         if ref is not None:
-            x = (x - ref) % 360 + ref
-        x, y = split_line(x, self.latitude, self.tracks)
+            x, y = wrap_longitude(x, y, ref, cut=True)
         return ax.plot(x, y, **kwargs)