Correct cumsum_by_track and add test

Author: AntSimi

src/py_eddy_tracker/generic.py

@@ -169,10 +169,9 @@ def cumsum_by_track(field, track):
         tr = track[i]
         if tr != tr_previous:
             d_cum = 0
-        cumsum_array[i] = d_cum
         d_cum += field[i]
+        cumsum_array[i] = d_cum
         tr_previous = tr
-    cumsum_array[i + 1] = d_cum
     return cumsum_array
 
 

tests/test_generic.py

@@ -1,6 +1,6 @@
-from numpy import arange, nan, zeros
+from numpy import arange, array, nan, ones, zeros
 
-from py_eddy_tracker.generic import simplify
+from py_eddy_tracker.generic import cumsum_by_track, simplify
 
 
 def test_simplify():
@@ -24,3 +24,9 @@ def test_simplify():
     x[:] = 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])
+    assert (cumsum_by_track(a, track) == [2, 4, 2, 4, 6, 8, 2, 4, 6, 2]).all()