scripts: Small improvements for input shaper and accelerometer scripts

Signed-off-by: Dmitry Butyugin <dmbutyugin@google.com>

1 files changed, 21 insertions, 5 deletions

diff --git a/scripts/graph_accelerometer.py b/scripts/graph_accelerometer.py
index 0e975bf5..f6d7cf89 100755
--- a/scripts/graph_accelerometer.py
+++ b/scripts/graph_accelerometer.py
@@ -12,10 +12,20 @@ sys.path.append(os.path.join(os.path.dirname(os.path.realpath(__file__)),
                              '..', 'klippy', 'extras'))
 from shaper_calibrate import ShaperCalibrate
 
-MAX_TITLE_LENGTH=80
+MAX_TITLE_LENGTH=65
 
-def parse_log(logname):
-    return np.loadtxt(logname, comments='#', delimiter=',')
+def parse_log(logname, opts):
+    with open(logname) as f:
+        for header in f:
+            if not header.startswith('#'):
+                break
+        if not header.startswith('freq,psd_x,psd_y,psd_z,psd_xyz'):
+            # Raw accelerometer data
+            return np.loadtxt(logname, comments='#', delimiter=',')
+    # Power spectral density data or shaper calibration data
+    opts.error("File %s does not contain raw accelerometer data and therefore "
+               "is not supported by graph_accelerometer.py script. Please use "
+               "calibrate_shaper.py script to process it instead." % (logname,))
 
 ######################################################################
 # Raw accelerometer graphing
@@ -185,7 +195,7 @@ def setup_matplotlib(output):
 
 def main():
     # Parse command-line arguments
-    usage = "%prog [options] <logs>"
+    usage = "%prog [options] <raw logs>"
     opts = optparse.OptionParser(usage)
     opts.add_option("-o", "--output", type="string", dest="output",
                     default=None, help="filename of output graph")
@@ -205,7 +215,7 @@ def main():
         opts.error("Incorrect number of arguments")
 
     # Parse data
-    datas = [parse_log(fn) for fn in args]
+    datas = [parse_log(fn, opts) for fn in args]
 
     setup_matplotlib(options.output)
 
@@ -215,6 +225,8 @@ def main():
         if options.compare:
             opts.error("comparison mode is not supported with csv output")
         if options.specgram:
+            if len(args) > 1:
+                opts.error("Only 1 input is supported in specgram mode")
             pdata, bins, t = calc_specgram(datas[0], options.axis)
             write_specgram(pdata, bins, t, options.output)
         else:
@@ -223,8 +235,12 @@ def main():
 
     # Draw graph
     if options.raw:
+        if len(args) > 1:
+            opts.error("Only 1 input is supported in raw mode")
         fig = plot_accel(datas[0], args[0])
     elif options.specgram:
+        if len(args) > 1:
+            opts.error("Only 1 input is supported in specgram mode")
         fig = plot_specgram(datas[0], args[0], options.max_freq, options.axis)
     elif options.compare:
         fig = plot_compare_frequency(datas, args, options.max_freq)