text: #!/usr/bin/env python
  
  '''
  extract ISBH and ISBD messages from AP_Logging files and produce FFT plots
  '''
  from __future__ import print_function
  
  import numpy
  import os
  import pylab
  import sys
  import time
  import copy
  
  from argparse import ArgumentParser
  import scipy.signal as signal
  
  parser = ArgumentParser(description=__doc__)
  parser.add_argument("--condition", default=None, help="select packets by condition")
  parser.add_argument("logs", metavar="LOG", nargs="+")
  parser.add_argument("--scale", dest='fft_scale', default='db', action='store', help="scale to use for displaying frequency information: 'db' or 'linear'")
  parser.add_argument("--window", dest='fft_window', default='hanning', action='store', help="windowing function to use for processing the data: 'hanning', 'blackman' or 'None'")
  parser.add_argument("--overlap", dest='fft_overlap', default=False, action='store_true', help="whether or not to use window overlap when analysing data")
  parser.add_argument("--output", dest='fft_output', default='psd', action='store', help="whether to output frequency spectrum information as power or linear spectral density: 'psd' or 'lsd'")
  parser.add_argument("--notch-params", default=False, action='store_true', help="whether to output estimated harmonic notch parameters")
  parser.add_argument("--notch-peak", dest='fft_peak', default=0, action='store', help="peak to select when setting notch parameters")
  parser.add_argument("--axis", dest='axis', default='XYZ', action='store', help="Select which axis to plot default: XYZ")
  
  args = parser.parse_args()
  
  from pymavlink import mavutil
  
  def mavfft_fttd(logfile, multi_log):
      '''display fft for raw ACC data in logfile'''
  
      '''object to store data about a single FFT plot'''
      class PlotData(object):
          def __init__(self, ffth):
              self.seqno = -1
              self.fftnum = ffth.N
              self.sensor_type = ffth.type
              self.instance = ffth.instance
              self.sample_rate_hz = ffth.smp_rate
              self.multiplier = ffth.mul
              self.data = {}
              self.data["X"] = []
              self.data["Y"] = []
              self.data["Z"] = []
              self.holes = False
              self.freq = None
  
          def add_fftd(self, fftd):
              if fftd.N != self.fftnum:
                  print("Skipping ISBD with wrong fftnum (%u vs %u)\n" % (fftd.N, self.fftnum))
                  return
              if self.holes:
                  print("Skipping ISBD(%u) for ISBH(%u) with holes in it" % (fftd.seqno, self.fftnum))
                  return
              if fftd.seqno != self.seqno+1:
                  print("ISBH(%u) has holes in it" % fftd.N)
                  self.holes = True
                  return
              self.seqno += 1
              self.data["X"].extend(fftd.x)
              self.data["Y"].extend(fftd.y)
              self.data["Z"].extend(fftd.z)
  
          def overlap_windows(self, plotdata):
              newplotdata = copy.deepcopy(self)
              if plotdata.tag() != self.tag():
                  print("Invalid FFT data tag (%s vs %s) for window overlap" % (plotdata.tag(), self.tag()))
                  return self
              if plotdata.fftnum <= self.fftnum:
                  print("Invalid FFT sequence (%u vs %u) for window overlap" % (plotdata.fftnum, self.fftnum))
                  return self
              newplotdata.data["X"] = numpy.split(numpy.asarray(self.data["X"]), 2)[1].tolist() + numpy.split(numpy.asarray(plotdata.data["X"]), 2)[0].tolist()
              newplotdata.data["Y"] = numpy.split(numpy.asarray(self.data["Y"]), 2)[1].tolist() + numpy.split(numpy.asarray(plotdata.data["Y"]), 2)[0].tolist()
              newplotdata.data["Z"] = numpy.split(numpy.asarray(self.data["Z"]), 2)[1].tolist() + numpy.split(numpy.asarray(plotdata.data["Z"]), 2)[0].tolist()
              return newplotdata
  
          def prefix(self):
              if self.sensor_type == 0:
                  return "Accel"
              elif self.sensor_type == 1:
                  return "Gyro"
              else:
                  return "?Unknown Sensor Type?"
  
          def tag(self):
              return str(self)
  
          def __str__(self):
              return "%s[%u]" % (self.prefix(), self.instance)
  
      print("Processing log %s" % logfile)
      mlog = mavutil.mavlink_connection(logfile)
  
      # see https://holometer.fnal.gov/GH_FFT.pdf for a description of the techniques used here
      things_to_plot = []
      plotdata = None
      prev_plotdata = {}
      msgcount = 0
      hntch_mode = None
      hntch_option = None
      batch_mode = None
      start_time = time.time()
      thr_total = 0.
      thr_count = 0
  
      while True:
          m = mlog.recv_match(condition=args.condition)
          if m is None:
              break
          msgcount += 1
          if msgcount % 1000 == 0:
              sys.stderr.write(".")
          msg_type = m.get_type()
          if msg_type == "ISBH":
              if plotdata is not None:
                  sensor = plotdata.tag()
                  # close off previous data collection
                  # in order to get 50% window overlap we need half the old data + half the new data and then the new data itself
                  if args.fft_overlap and sensor in prev_plotdata:
                      things_to_plot.append(prev_plotdata[sensor].overlap_windows(plotdata))
                  things_to_plot.append(plotdata)
                  prev_plotdata[sensor] = plotdata
              plotdata = PlotData(m)
              continue
  
          if msg_type == "ISBD":
              if plotdata is None:
                  sys.stderr.write("?(fftnum=%u)" % m.N)
                  continue
              plotdata.add_fftd(m)
          
          if msg_type == "PARM":
              if m.Name == "INS_HNTCH_MODE":
                  hntch_mode = m.Value
              elif m.Name == "INS_HNTCH_OPTS":
                  hntch_option = m.Value
              elif m.Name == "INS_LOG_BAT_OPT":
                  batch_mode = m.Value
  
          # get an average read of the throttle value assuming a stable hover above 1m
          if args.notch_params and msg_type == "CTUN":
              if m.Alt > 1:
                  thr_total += m.ThO
                  thr_count = thr_count + 1
  
      print("", file=sys.stderr)
      time_delta = time.time() - start_time
      print("%us messages  %u messages/second  %u kB/second" % (msgcount, msgcount/time_delta, os.stat(filename).st_size/time_delta))
      print("Extracted %u fft data sets" % len(things_to_plot), file=sys.stderr)
      if args.notch_params:
          thr_ref = thr_total / thr_count
          print("Throttle average %f" % thr_ref)
  
      sum_fft = {}
      freqmap = {}
      sample_rates = {}
      counts = {}
      window = {}
      S2 = {}
      hntch_mode_names = { 0:"No", 1:"Throttle", 2:"RPM", 3:"ESC", 4:"FFT"}
      hntch_option_names = { 0:"Single", 1:"Double", 2:"Dynamic", 4:"Loop-Rate"}
      batch_mode_names = { 0:"Pre-filter", 1:"Sensor-rate", 2:"Post-filter" }
      fft_peak = int(args.fft_peak)
  
      first_freq = None
      for thing_to_plot in things_to_plot:
  
          fft_len = len(thing_to_plot.data["X"])
  
          if thing_to_plot.tag() not in sum_fft:
              sum_fft[thing_to_plot.tag()] = {
                  "X": numpy.zeros(fft_len//2+1),
                  "Y": numpy.zeros(fft_len//2+1),
                  "Z": numpy.zeros(fft_len//2+1),
              }
              counts[thing_to_plot.tag()] = 0
  
          if thing_to_plot.tag() not in window:
              if args.fft_window == 'hanning':
                  # create hanning window
                  window[thing_to_plot.tag()] = numpy.hanning(fft_len)
              elif args.fft_window == 'blackman':
                  # create blackman window
                  window[thing_to_plot.tag()] = numpy.blackman(fft_len)
              else:
                  window[thing_to_plot.tag()] = numpy.arange(fft_len)
                  window[thing_to_plot.tag()].fill(1)
              # calculate NEBW constant
              S2[thing_to_plot.tag()] = numpy.inner(window[thing_to_plot.tag()], window[thing_to_plot.tag()])
  
          for axis in [ "X","Y","Z" ]:
              # only plot the selected axis
              if axis not in args.axis:
                  continue
              # normalize data and convert to dps in order to produce more meaningful magnitudes
              if thing_to_plot.sensor_type == 1:
                  d = numpy.array(numpy.degrees(thing_to_plot.data[axis])) / float(thing_to_plot.multiplier)
              else:
                  d = numpy.array(thing_to_plot.data[axis]) / float(thing_to_plot.multiplier)
              if len(d) == 0:
                  print("No data?!?!?!")
                  continue
              if len(window[thing_to_plot.tag()]) != fft_len:
                  print("Skipping corrupted frame of length %d" % fft_len)
                  continue
  
              # apply window to the input
              d *= window[thing_to_plot.tag()]
              # perform RFFT
              d_fft = numpy.fft.rfft(d)
              # convert to squared complex magnitude
              d_fft = numpy.square(abs(d_fft))
              # remove DC component
              d_fft[0] = 0
              d_fft[-1] = 0
              # accumulate the sums
              sum_fft[thing_to_plot.tag()][axis] += d_fft
              freq = numpy.fft.rfftfreq(len(d), 1.0/thing_to_plot.sample_rate_hz)
              freqmap[thing_to_plot.tag()] = freq
  
          sample_rates[thing_to_plot.tag()] = thing_to_plot.sample_rate_hz
          counts[thing_to_plot.tag()] += 1
  
      numpy.seterr(divide = 'ignore')
      for sensor in sum_fft:
          print("Sensor: %s" % str(sensor))
          fig = pylab.figure(str(sensor))
          for axis in [ "X","Y","Z" ]:
              # only plot the selected axis
              if axis not in args.axis:
                  continue
  
              # normalize output to averaged PSD
              psd = 2 * (sum_fft[sensor][axis] / counts[sensor]) / (sample_rates[sensor] * S2[sensor])
  
              # calculate peaks from linear accel data
              # the accel data is less noisy than the gyro data
              if sensor == 'Accel[0]' and axis == "X" and args.notch_params:
                  linear_psd = numpy.sqrt(psd)
                  peaks, _ = signal.find_peaks(psd, prominence=0.1)
                  peak_freqs = freqmap[sensor][peaks]
                  print("Peaks: %s" % str(peak_freqs))
                  print("INS_HNTCH_REF = %.4f" % thr_ref)
                  print("INS_HNTCH_FREQ = %.1f" % float(peak_freqs[fft_peak]))
                  print("INS_HNTCH_BW = %.1f" % (float(peak_freqs[fft_peak])/2.))
  
              # linerize of requested
              if args.fft_output == 'lsd':
                  psd = numpy.sqrt(psd)
              # convert to db if requested
              if args.fft_scale == 'db':
                  psd = 10 * numpy.log10 (psd)
              # add file name to axis if doing multi-file plot
              if multi_log:
                  # remove extension and path
                  (log_name, _, _) = logfile.rpartition('.')
                  (_, _, log_name) = log_name.rpartition('/')
                  legend_label = '%s (%s)' % (axis, log_name)
              else:
                  legend_label = axis
              pylab.plot(freqmap[sensor], psd, label=legend_label)
          pylab.legend(loc='upper right')
          pylab.xlabel('Hz')
          scale_label=''
          psd_label='PSD'
          if args.fft_scale =='db':
              scale_label="dB "
          if args.fft_output == 'lsd':
              if str(sensor).startswith("Gyro"):
                  pylab.ylabel('LSD $' + scale_label + 'd/s/\sqrt{Hz}$')
              else:
                  pylab.ylabel('LSD $' + scale_label + 'm/s^2/\sqrt{Hz}$')
          else:
              if str(sensor).startswith("Gyro"):
                  pylab.ylabel('PSD $' + scale_label + 'd^2/s^2/Hz$')
              else:
                  pylab.ylabel('PSD $' + scale_label + 'm^2/s^4/Hz$')
  
          if hntch_mode is not None and hntch_option is not None and batch_mode is not None:
              option_label = ""
              for hopt in hntch_option_names:
                  if hopt & int(hntch_option) != 0:
                      if len(option_label) > 0:
                          option_label += "+"
                      option_label += hntch_option_names[hopt]
              textstr = '\n'.join((
                  r'%s tracking' % (hntch_mode_names[hntch_mode], ),
                  r'%s notch' % (option_label, ),
                  r'%s sampling' % (batch_mode_names[batch_mode], )))
  
              props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)
  
              pylab.text(0.5, 0.95, textstr, fontsize=12,
                  verticalalignment='top', bbox=props, transform=pylab.gca().transAxes)
  
  # auto set option for adding log names to legend label
  multi_log = False
  if len(args.logs) > 1:
      multi_log = True
  
  for filename in args.logs:
      mavfft_fttd(filename, multi_log)
  
  pylab.show()
  