sweetberry: calculate statistics for sweetberry readings

This CL provides the tool to calculate statistics for sweetberry
readings and present them in a clear & easy to read format. It
also provides the flag to store raw data and statistics
summary, should the need arise.
There are also some code cleanup for powerlog.py.

BRANCH=None
BUG=b:35578707
TEST=./powerlog.py -b xxx.board -c xxx.scenario --print_stats \
--save_stats --save_raw_data
python -m unittest -v stats_manager_unittest

Change-Id: I4aa732756fe6512f37acfcb59b11d950101887d7
Signed-off-by: Mengqi Guo <mqg@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/667241
Reviewed-by: Nick Sanders <nsanders@chromium.org>

4 files changed, 278 insertions, 27 deletions

diff --git a/extra/usb_power/__init__.py b/extra/usb_power/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
--- /dev/null
+++ b/extra/usb_power/__init__.py
diff --git a/extra/usb_power/powerlog.py b/extra/usb_power/powerlog.py
index 5daafaa07f..5f39d7dc1b 100755
--- a/extra/usb_power/powerlog.py
+++ b/extra/usb_power/powerlog.py
@@ -6,9 +6,12 @@
 """Program to fetch power logging data from a sweetberry device
    or other usb device that exports a USB power logging interface.
 """
+from __future__ import print_function
 import argparse
 import array
+import datetime
 import json
+import os
 import struct
 import sys
 import time
@@ -17,14 +20,16 @@ from pprint import pprint
 
 import usb
 
+from stats_manager import StatsManager
+
 # This can be overridden by -v.
 debug = False
 def debuglog(msg):
   if debug:
-    print msg
+    print(msg)
 
 def logoutput(msg):
-  print msg
+  print(msg)
   sys.stdout.flush()
 
 
@@ -309,11 +314,6 @@ class Spower(object):
     else:
       debuglog("Command START: FAIL")
 
-    title = "ts:%dus" % actual_us
-    for i in range(0, len(self._inas)):
-      name = self._inas[i]['name']
-      title += ", %s uW" % name
-
     return actual_us
 
   def add_ina_name(self, name):
@@ -343,7 +343,7 @@ class Spower(object):
     raise Exception("Power", "Failed to find INA %s" % name)
 
   def set_time(self, timestamp_us):
-    """Set sweetberry tie to match host time.
+    """Set sweetberry time to match host time.
 
     Args:
       timestamp_us: host timestmap in us.
@@ -403,7 +403,7 @@ class Spower(object):
       cmd = struct.pack("<H", 0x0004)
       bytesread = self.wr_command(cmd, read_count=expected_bytes)
     except usb.core.USBError as e:
-      print "READ LINE FAILED %s" % e
+      print("READ LINE FAILED %s" % e)
       return None
 
     if len(bytesread) == 1:
@@ -441,13 +441,13 @@ class Spower(object):
     """
     status, size = struct.unpack("<BB", data[0:2])
     if len(data) != self.report_size(size):
-      print "READ LINE FAILED st:%d size:%d expected:%d len:%d" % (
-          status, size, self.report_size(size), len(data))
+      print("READ LINE FAILED st:%d size:%d expected:%d len:%d" % (
+          status, size, self.report_size(size), len(data)))
     else:
       pass
 
     timestamp = struct.unpack("<Q", data[2:10])[0]
-    debuglog("READ LINE: st:%d size:%d time:%d" % (status, size, timestamp))
+    debuglog("READ LINE: st:%d size:%d time:%dus" % (status, size, timestamp))
     ftimestamp = float(timestamp) / 1000000.
 
     record = {"ts": ftimestamp, "status": status, "berry":self._board}
@@ -457,7 +457,7 @@ class Spower(object):
       raw_w = struct.unpack("<H", data[idx:idx+2])[0]
       uw = raw_w * self._inas[i]['uWscale']
       name = self._inas[i]['name']
-      debuglog("READ %d %s: %fs: %fmW" % (i, name, ftimestamp, uw))
+      debuglog("READ %d %s: %fs: %fuW" % (i, name, ftimestamp, uw))
       record[self._inas[i]['name']] = uw
 
     return record
@@ -484,11 +484,13 @@ class powerlog(object):
     obj = powerlog()
 
   Instance Variables:
-    _pwr[]: Spower objects for individual sweetberries
+    _data: records sweetberries readings and calculates statistics.
+    _pwr[]: Spower objects for individual sweetberries.
   """
 
-  def __init__(self, brdfile, cfgfile, serial_a=None,
-               serial_b=None, sync_date=False, use_ms=False):
+  def __init__(self, brdfile, cfgfile, serial_a=None, serial_b=None,
+               sync_date=False, use_ms=False, print_stats=False,
+               save_stats=False, save_raw_data=False):
     """
     Args:
       brdfile: string name of json file containing board layout.
@@ -498,8 +500,12 @@ class powerlog(object):
       sync_date: report timestamps synced with host datetime.
       use_ms: report timestamps in ms rather than us.
     """
+    self._data = StatsManager()
     self._pwr = {}
     self._use_ms = use_ms
+    self._print_stats = print_stats
+    self._save_stats = save_stats
+    self._save_raw_data = save_raw_data
 
     if not serial_a and not serial_b:
       self._pwr['A'] = Spower('A')
@@ -518,7 +524,7 @@ class powerlog(object):
 
     # Allocate the rails to the appropriate boards.
     used_boards = []
-    for name in names:
+    for name in self._names:
       success = False
       for key in self._pwr.keys():
         if self._pwr[key].add_ina_name(name):
@@ -594,7 +600,8 @@ class powerlog(object):
               aggregate_record[rkey] = record[rkey]
             aggregate_record["boards"].add(record["berry"])
           else:
-            print "break %s, %s" % (record["berry"], aggregate_record["boards"])
+            print("break %s, %s" % (record["berry"],
+                                    aggregate_record["boards"]))
             break
 
           if aggregate_record["boards"] == set(self._pwr.keys()):
@@ -602,6 +609,7 @@ class powerlog(object):
             for name in self._names:
               if name in aggregate_record:
                 csv += ", %.2f" % aggregate_record[name]
+                self._data.AddValue(name, aggregate_record[name])
               else:
                 csv += ", "
             csv += ", %d" % aggregate_record["status"]
@@ -611,10 +619,22 @@ class powerlog(object):
             for r in range(0, len(self._pwr)):
               pending_records.pop(0)
 
+    except KeyboardInterrupt:
+      print('\nCTRL+C caught.')
 
     finally:
       for key in self._pwr:
         self._pwr[key].stop()
+      self._data.CalculateStats()
+      if self._print_stats:
+        self._data.PrintSummary()
+      save_dir = datetime.datetime.now().strftime('Sweetberry%Y%m%d%H%M%S')
+      save_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)),
+                                              save_dir)
+      if self._save_stats:
+        self._data.SaveSummary(save_dir)
+      if self._save_raw_data:
+        self._data.SaveRawData(save_dir)
 
 
 def main():
@@ -624,15 +644,13 @@ def main():
   parser.add_argument('-b', '--board', type=str,
       help="Board configuration file, eg. my.board", default="")
   parser.add_argument('-c', '--config', type=str,
-      help="Rail config to monitor, eg my.config", default="")
+      help="Rail config to monitor, eg my.scenario", default="")
   parser.add_argument('-A', '--serial', type=str,
       help="Serial number of sweetberry A", default="")
   parser.add_argument('-B', '--serial_b', type=str,
       help="Serial number of sweetberry B", default="")
   parser.add_argument('-t', '--integration_us', type=int,
       help="Target integration time for samples", default=100000)
-  parser.add_argument('-n', '--samples', type=int,
-      help="Samples to capture, or none to sample forever.", default=0)
   parser.add_argument('-s', '--seconds', type=float,
       help="Seconds to run capture. Overrides -n", default=0.)
   parser.add_argument('--date', default=False,
@@ -641,6 +659,15 @@ def main():
       help="Print timestamp as milliseconds", action="store_true")
   parser.add_argument('--slow', default=False,
       help="Intentionally overflow", action="store_true")
+  parser.add_argument('--print_stats', default=False,
+      help="Print statistics for sweetberry readings at the end",
+      action="store_true")
+  parser.add_argument('--save_stats', default=False,
+      help="Save statistics for sweetberry readings",
+      action="store_true")
+  parser.add_argument('--save_raw_data', default=False,
+      help="Save raw data for sweetberry readings",
+      action="store_true")
   parser.add_argument('-v', '--verbose', default=False,
       help="Very chatty printout", action="store_true")
 
@@ -658,12 +685,14 @@ def main():
 
   brdfile = args.board
   cfgfile = args.config
-  samples = args.samples
   seconds = args.seconds
   serial_a = args.serial
   serial_b = args.serial_b
   sync_date = args.date
   use_ms = args.ms
+  print_stats = args.print_stats
+  save_stats = args.save_stats
+  save_raw_data = args.save_raw_data
 
   boards = []
 
@@ -671,13 +700,11 @@ def main():
   if args.slow:
     sync_speed = 1.2
 
-  forever = True
-  if samples > 0 or seconds > 0.:
-    forever = False
-
   # Set up logging interface.
   powerlogger = powerlog(brdfile, cfgfile, serial_a=serial_a,
-      serial_b=serial_b, sync_date=sync_date, use_ms=use_ms)
+      serial_b=serial_b, sync_date=sync_date, use_ms=use_ms,
+      print_stats=print_stats, save_stats=save_stats,
+      save_raw_data=save_raw_data)
 
   # Start logging.
   powerlogger.start(integration_us_request, seconds, sync_speed=sync_speed)
diff --git a/extra/usb_power/stats_manager.py b/extra/usb_power/stats_manager.py
new file mode 100644
index 0000000000..02f984f097
--- /dev/null
+++ b/extra/usb_power/stats_manager.py
@@ -0,0 +1,137 @@
+# Copyright 2017 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Calculates statistics for lists of data and pretty print them."""
+
+from __future__ import print_function
+import collections
+import numpy
+import os
+
+STATS_PREFIX = '@@'
+KEY_PREFIX = '__'
+# This prefix is used for keys that should not be shown in the summary tab, such
+# as timeline keys.
+NOSHOW_PREFIX = '!!'
+
+class StatsManager(object):
+  """Calculates statistics for several lists of data(float)."""
+
+  def __init__(self):
+    """Initialize infrastructure for data and their statistics."""
+    self._data = collections.defaultdict(list)
+    self._summary = {}
+
+  def AddValue(self, domain, value):
+    """Add one value for a domain.
+
+    Args:
+      domain: the domain name for the value.
+      value: one time reading for domain, expect type float.
+    """
+    if isinstance(value, int):
+      value = float(value)
+    if isinstance(value, float):
+      self._data[domain].append(value)
+      return
+    print('Warning: value %s for domain %s is not a number, thus ignored.' %
+          (value, domain))
+
+  def CalculateStats(self):
+    """Calculate stats for all domain-data pairs.
+
+    First erases all previous stats, then calculate stats for all data.
+    """
+    self._summary = {}
+    for domain, data in self._data.iteritems():
+      data_np = numpy.array(data)
+      self._summary[domain] = {
+          'mean'     : data_np.mean(),
+          'min'      : data_np.min(),
+          'max'      : data_np.max(),
+          'stddev'   : data_np.std(),
+          'count'    : data_np.size,
+      }
+
+  def _SummaryToString(self, prefix=STATS_PREFIX):
+    """Format summary into a string, ready for pretty print.
+
+    Args:
+      prefix: start every row in summary string with prefix, for easier reading.
+    """
+    headers = ('NAME', 'COUNT', 'MEAN', 'STDDEV', 'MAX', 'MIN')
+    table = [headers]
+    for domain in sorted(self._summary.keys()):
+      if domain.startswith(NOSHOW_PREFIX):
+        continue
+      stats = self._summary[domain]
+      row = [domain.lstrip(KEY_PREFIX)]
+      row.append(str(stats['count']))
+      for entry in headers[2:]:
+        row.append('%.2f' % stats[entry.lower()])
+      table.append(row)
+
+    max_col_width = []
+    for col_idx in range(len(table[0])):
+      col_item_widths = [len(row[col_idx]) for row in table]
+      max_col_width.append(max(col_item_widths))
+
+    formatted_table = []
+    for row in table:
+      formatted_row = prefix + ' '
+      for i in range(len(row)):
+        formatted_row += row[i].rjust(max_col_width[i] + 2)
+      formatted_table.append(formatted_row)
+    return '\n'.join(formatted_table)
+
+  def PrintSummary(self, prefix=STATS_PREFIX):
+    """Print the formatted summary.
+
+    Args:
+      prefix: start every row in summary string with prefix, for easier reading.
+    """
+    summary_str = self._SummaryToString(prefix=prefix)
+    print(summary_str)
+
+  def GetSummary(self):
+    """Getter for summary."""
+    return self._summary
+
+  def SaveSummary(self, directory, fname='summary.txt', prefix=STATS_PREFIX):
+    """Save summary to file.
+
+    Args:
+      directory: directory to save the summary in.
+      fname: filename to save summary under.
+      prefix: start every row in summary string with prefix, for easier reading.
+    """
+    summary_str = self._SummaryToString(prefix=prefix) + '\n'
+
+    if not os.path.exists(directory):
+      os.makedirs(directory)
+    fname = os.path.join(directory, fname)
+    with open(fname, 'w') as f:
+      f.write(summary_str)
+
+  def GetRawData(self):
+    """Getter for all raw_data."""
+    return self._data
+
+  def SaveRawData(self, directory, dirname='raw_data'):
+    """Save raw data to file.
+
+    Args:
+      directory: directory to create the raw data folder in.
+      dirname: folder in which raw data live.
+    """
+    if not os.path.exists(directory):
+      os.makedirs(directory)
+    dirname = os.path.join(directory, dirname)
+    if not os.path.exists(dirname):
+      os.makedirs(dirname)
+    for domain, data in self._data.iteritems():
+      fname = domain + '.txt'
+      fname = os.path.join(dirname, fname)
+      with open(fname, 'w') as f:
+        f.write('\n'.join('%.2f' % value for value in data) + '\n')
diff --git a/extra/usb_power/stats_manager_unittest.py b/extra/usb_power/stats_manager_unittest.py
new file mode 100644
index 0000000000..9b86b15ad4
--- /dev/null
+++ b/extra/usb_power/stats_manager_unittest.py
@@ -0,0 +1,87 @@
+# Copyright 2017 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Unit tests for StatsManager."""
+
+from __future__ import print_function
+import os
+import shutil
+import tempfile
+import unittest
+
+from stats_manager import StatsManager
+
+class TestStatsManager(unittest.TestCase):
+  """Test to verify StatsManager methods work as expected.
+
+  StatsManager should collect raw data, calculate their statistics, and save
+  them in expected format.
+  """
+
+  def setUp(self):
+    """Set up data and create a temporary directory to save data and stats."""
+    self.tempdir = tempfile.mkdtemp()
+    self.data = StatsManager()
+    self.data.AddValue('A', 99999.5)
+    self.data.AddValue('A', 100000.5)
+    self.data.AddValue('A', 'ERROR')
+    self.data.AddValue('B', 1.5)
+    self.data.AddValue('B', 2.5)
+    self.data.AddValue('B', 3.5)
+    self.data.CalculateStats()
+
+  def tearDown(self):
+    """Delete the temporary directory and its content."""
+    shutil.rmtree(self.tempdir)
+
+  def test_GetRawData(self):
+    raw_data = self.data.GetRawData()
+    self.assertListEqual([99999.5, 100000.5], raw_data['A'])
+    self.assertListEqual([1.5, 2.5, 3.5], raw_data['B'])
+
+  def test_GetSummary(self):
+    summary = self.data.GetSummary()
+    self.assertEqual(2, summary['A']['count'])
+    self.assertAlmostEqual(100000.5, summary['A']['max'])
+    self.assertAlmostEqual(99999.5, summary['A']['min'])
+    self.assertAlmostEqual(0.5, summary['A']['stddev'])
+    self.assertAlmostEqual(100000.0, summary['A']['mean'])
+    self.assertEqual(3, summary['B']['count'])
+    self.assertAlmostEqual(3.5, summary['B']['max'])
+    self.assertAlmostEqual(1.5, summary['B']['min'])
+    self.assertAlmostEqual(0.81649658092773, summary['B']['stddev'])
+    self.assertAlmostEqual(2.5, summary['B']['mean'])
+
+  def test_SaveRawData(self):
+    dirname = 'unittest_raw_data'
+    self.data.SaveRawData(self.tempdir, dirname)
+    dirname = os.path.join(self.tempdir, dirname)
+    fileA = os.path.join(dirname, 'A.txt')
+    fileB = os.path.join(dirname, 'B.txt')
+    with open(fileA, 'r') as fA:
+      self.assertEqual('99999.50', fA.readline().strip())
+      self.assertEqual('100000.50', fA.readline().strip())
+    with open(fileB, 'r') as fB:
+      self.assertEqual('1.50', fB.readline().strip())
+      self.assertEqual('2.50', fB.readline().strip())
+      self.assertEqual('3.50', fB.readline().strip())
+
+  def test_SaveSummary(self):
+    fname = 'unittest_summary.txt'
+    self.data.SaveSummary(self.tempdir, fname)
+    fname = os.path.join(self.tempdir, fname)
+    with open(fname, 'r') as f:
+      self.assertEqual(
+          '@@   NAME  COUNT       MEAN  STDDEV        MAX       MIN\n',
+          f.readline())
+      self.assertEqual(
+          '@@      A      2  100000.00    0.50  100000.50  99999.50\n',
+          f.readline())
+      self.assertEqual(
+          '@@      B      3       2.50    0.82       3.50      1.50\n',
+          f.readline())
+
+
+if __name__ == '__main__':
+  unittest.main()