scripts/random-muscle-activation.py: Implement signal smoothing

1 files changed, 75 insertions, 12 deletions

diff --git a/scripts/random-muscle-activation.py b/scripts/random-muscle-activation.py
index e168769..2af82b5 100755
--- a/scripts/random-muscle-activation.py
+++ b/scripts/random-muscle-activation.py
@@ -13,6 +13,7 @@ import os
 import sys
 import getopt
 import numpy as np
+import matplotlib.pyplot as plt
 import plotcsv
 
 DEFAULT_N = 1000
@@ -23,24 +24,30 @@ DEFAULT_YMAX = 50.0
 DEFAULT_YDMIN = -0.5
 DEFAULT_YDMAX = 0.5
 
+WINDOWS = [ 'flat', 'hanning', 'hamming', 'bartlett', 'blackman' ]
+WINDOW_LEN = 11
+
 def usage():
     print("""usage: {} [OPTION...] FILE...
 
 Generate random muscle activation data in CSV.
 
 options:
-    -f    force overwrite of existing file
-    -n N  set number of data points to N (default: {})
-    -p    plot data after generating it
-    -i N  set minimum initial y value to N (default: {})
-    -I N  set maximum initial y value to N (default: {})
-    -x N  set minimum x value to N (default: {})
-    -X N  set maximum x value to N (default: {})
-    -y N  set minimum y value to N (default: {})
-    -Y N  set maximum y value to N (default: {})
-    -d N  set minimum y change value to N (default: {})
-    -D N  set maximum y change value to N (default: {})""".format(
+    -f       force overwrite of existing file
+    -n N     set number of data points to N (default: {})
+    -p       plot data after generating it
+    -s WIN   smoothen the data using WIN window function
+             WIN should be one of {}
+    -i N     set minimum initial y value to N (default: {})
+    -I N     set maximum initial y value to N (default: {})
+    -x N     set minimum x value to N (default: {})
+    -X N     set maximum x value to N (default: {})
+    -y N     set minimum y value to N (default: {})
+    -Y N     set maximum y value to N (default: {})
+    -d N     set minimum y change value to N (default: {})
+    -D N     set maximum y change value to N (default: {})""".format(
         os.path.basename(sys.argv[0]), DEFAULT_N,
+        ", ".join(WINDOWS),
         DEFAULT_YMIN, DEFAULT_YMAX,
         DEFAULT_XMIN, DEFAULT_XMAX,
         DEFAULT_YMIN, DEFAULT_YMAX,
@@ -63,9 +70,44 @@ def generate_activation_data(N, xmin, xmax, ymin, ymax, ydmin, ydmax, imin, imax
 
     return act
 
+# based on scipy.org/Cookbook/SignalSmooth
+def smoothen(x, window='hanning', window_len=11):
+    """
+    smoothen signal using a window with requested size
+
+    input:
+      x:          the input signal
+      window:     type of window from 'flat', 'hanning', 'hamming', 'bartlett',
+                  'blackman'; the 'flat' method will apply a moving average
+      window_len: dimension of smoothing window, must be an odd integer
+
+    output:
+      smoothened signal
+    """
+
+    if x.ndim != 1:
+        raise ValueError, "only 1-dimensional arrays can be smoothened"
+    if x.size < window_len:
+        raise ValueError, "input vector needs to be bigger than window size"
+    if type(window_len) != int or window_len % 2 != 1:
+        raise ValueError, "window length must be an odd integer"
+
+    if window_len < 3:
+        return x
+    if not window in WINDOWS:
+        raise ValueError, "window type must be one of " + ", ".join([ "'" + w + "'" for w in WINDOWS ])
+
+    s = np.r_[x[window_len-1:0:-1],x,x[-1:-window_len:-1]]
+    if window == 'flat':
+        w = np.ones(window_len, 'd')
+    else:
+        w = eval('np.' + window + '(window_len)')
+
+    return np.convolve(w / w.sum(), s, mode='valid')
+
 def main():
     try:
-        opts, args = getopt.getopt(sys.argv[1:], "fn:pi:I:x:X:y:Y:d:D:h")
+        opts, args = getopt.getopt(sys.argv[1:], "fn:ps:i:I:x:X:y:Y:d:D:h")
     except getopt.GetoptError as err:
         print(str(err))
         usage()
@@ -73,6 +115,7 @@ def main():
 
     do_overwrite = False
     do_plot = False
+    window = None
     N = DEFAULT_N
     imin = DEFAULT_YMIN
     imax = DEFAULT_YMAX
@@ -94,6 +137,8 @@ def main():
                 sys.exit(-1)
         elif o == '-p':
             do_plot = True
+        elif o == '-s':
+            window = a
         elif o == '-i':
             try:
                 imin = float(a)
@@ -198,6 +243,24 @@ def main():
     for a in args:
         with open(a, 'w') as f:
             act = generate_activation_data(N, xmin, xmax, ymin, ymax, ydmin, ydmax, imin, imax)
+            if window:
+                sact = smoothen(act[:,1], window, WINDOW_LEN)
+                sact = sact[(WINDOW_LEN - 1) / 2:-(WINDOW_LEN - 1) / 2]
+                plt.plot(act[:,1])
+                plt.plot(sact)
+                plt.legend(['original', 'smoothened'])
+                plt.show()
+                act[:,1] = sact
+            # XXX: temp
+            #    plt.plot(act[:,1])
+            #    l = ['original']
+            #    for w in WINDOWS:
+            #        sact = smoothen(act[:,1], w, WINDOW_LEN)
+            #        plt.plot(sact[(WINDOW_LEN-1)/2:-(WINDOW_LEN-1)/2])
+            #        l.append(w)
+            #    plt.legend(l)
+            #    plt.show()
+
             f.write('x,y\n')
             for x, y in act:
                 f.write('{},{}\n'.format(x, y))