motors/pistol_grip/generate_cogging.py

#!/usr/bin/python3

import sys
from matplotlib import pylab

# TODO(austin): Plot flag.

def main(argv):
  if len(argv) < 4:
    print 'Args: input output.cc struct_name'
    return 1
  data_sum = [0.0] * 4096
  data_count = [0] * 4096
  data_list_absolute = []
  data_list_current = []

  with open(argv[1], 'r') as fd:
    for line in fd:
      if line.startswith('reading'):
        split_line = line.split()
        data_absolute = int(split_line[1])
        data_index = int(split_line[3][2:])
        data_current = int(split_line[2]) / 10000.0
        data_sum[data_index] += data_current
        data_count[data_index] += 1
        data_list_absolute.append(data_absolute)
        data_list_current.append(data_current)
  data = [0.0] * 4096
  min_zero = 4096
  max_zero = 0
  for i in range(0, 4096):
    if data_count[i] == 0:
      min_zero = min(i, min_zero)
      max_zero = max(i, min_zero)

  for i in range(0, 4096):
    if data_count[i] != 0:
      data[i] = data_sum[i] / data_count[i]
  if min_zero == 0 and max_zero == 4095:
    for i in range(0, 4096):
      if data_count[i] != 0:
        while i > 0:
          data[i - 1] = data[i]
          i -= 1
        break;

    for i in reversed(range(0, 4096)):
      if data_count[i] != 0:
        while i < 4095:
          data[i + 1] = data[i]
          i += 1
        break;
  else:
    for i in range(0, 4096):
      if data_count[i] == 0:
        if i < (min_zero + max_zero) / 2:
          data[i] = data[min_zero - 1]
        else:
          data[i] = data[max_zero + 1]

  pylab.plot(range(0, 4096), data)
  pylab.figure()
  pylab.plot(data_list_absolute, data_list_current)
  pylab.show()
  with open(argv[2], 'w') as out_fd:
    out_fd.write('extern const float %s[4096];\n' % argv[3])
    out_fd.write('const float %s[4096] = {\n' % argv[3])
    for datapoint in data:
      out_fd.write('    %ff,\n' % datapoint)
    out_fd.write('};')

  return 0

if __name__ == '__main__':
  sys.exit(main(sys.argv))