test: simple unit tests for moving average

automated_test.py

@@ -6,6 +6,7 @@
 
 import kimimaro.intake
 import kimimaro.skeletontricks
+from kimimaro.utility import moving_average
 
 def test_empty_image():
   labels = np.zeros( (256, 256, 256), dtype=bool)  
@@ -469,7 +470,37 @@ def test_cross_sectional_area():
   assert np.all(skel.cross_sectional_area == 9)
 
 
+def test_moving_average():
 
+  data = np.array([])
+  assert np.all(moving_average(data, 1) == data)
+  assert np.all(moving_average(data, 2) == data)
+
+  data = np.array([1,1,1,1,1,1,1,1,1,1,1])
+  assert np.all(moving_average(data, 1) == data)
+
+  data = np.array([1,1,1,1,1,1,1,1,1,1,1,1])
+  assert np.all(moving_average(data, 1) == data)
+
+  data = np.array([1,1,1,1,1,10,1,1,1,1,1])
+  assert np.all(moving_average(data, 1) == data)
+
+  data = np.array([1,1,1,1,1,1,1,1,1,1,1])
+  assert np.all(moving_average(data, 2) == data)
+
+  data = np.array([0,1,1,1,1,1,1,1,1,1,0])
+  ans = np.array([
+    0,0.5,1,1,1,1,1,1,1,1,0.5
+  ])
+  assert np.all(moving_average(data, 2) == ans)
+
+  data = np.array([0,1,1,1,1,1,1,1,1,1,0])
+  ans = np.array([
+    1/3,1/3,2/3,1,1,1,1,1,1,1,2/3
+  ])
+  res = moving_average(data, 3)
+  assert np.all(res == ans)
+  assert len(ans) == len(data)
 
 
 

kimimaro/utility.py

@@ -194,13 +194,17 @@ def moving_average(a:np.ndarray, n:int) -> np.ndarray:
     raise ValueError(f"Window size ({n}), must be >= 1.")
   elif n == 1:
     return a
-  mirror = (n - 1) / 2
-  extra = 0
-  if mirror != int(mirror):
-    extra = 1
-  mirror = int(mirror)
-  a = np.pad(a, [[mirror, mirror+extra],[0,0]], mode="symmetric")
+
+  if len(a) == 0:
+    return a
+
+  if a.ndim == 2:
+    a = np.pad(a, [[n, n],[0,0]], mode="symmetric")
+  else:
+    a = np.pad(a, [n, n], mode="symmetric")
+
   ret = np.cumsum(a, dtype=float, axis=0)
-  ret[n:] = ret[n:] - ret[:-n]
-  return ret[n - 1:] / n
+  ret = (ret[n:] - ret[:-n])[:-n]
+  ret /= float(n)
+  return ret