posts/Comparing_velocity_clamp_experiments/spectrum.py

   1 #!/usr/bin/env python2
   2
   3 "Plot vibration data with higher low-frequency resolution"
   4
   5 import argparse as _argparse
   6
   7 import FFT_tools as _FFT_tools
   8 import h5py as _h5py
   9 import numpy as _numpy
  10 import matplotlib.pyplot as _pyplot
  11 import calibcant.vibration_analyze as _vibration_analyze
  12
  13
  14 def plot_spectrum(path, group='/', frequency=1, title=None, filename=None):
  15     with _h5py.File(path) as f:
  16         data = f[group][...]
  17     print('loaded {} points from {}:{}'.format(data.shape, path, group))
  18     dt = 1./frequency
  19     time_axis = _numpy.arange(len(data)) * dt
  20     vib_data = data - data.mean()  # give the FFT routine an easier time
  21     freq_axis,power = _FFT_tools.unitary_power_spectrum(
  22         vib_data, freq=frequency)
  23
  24     figure = _pyplot.figure()
  25
  26     time_axes = figure.add_subplot(2, 1, 1)
  27     time_axes.autoscale(tight=True)
  28     if title:
  29         time_axes.set_title(title)
  30     time_axes.set_xlabel('time (seconds)')
  31     time_axes.set_ylabel('deflection (bits)')
  32     time_axes.plot(time_axis, data, 'b.')
  33
  34     freq_axes = figure.add_subplot(2, 1, 2)
  35     freq_axes.autoscale(tight=True)
  36     freq_axes.set_xscale('log')
  37     freq_axes.set_yscale('log')
  38     freq_axes.set_xlabel('frequency (Hz)')
  39     freq_axes.set_ylabel('power (bits$^2$/Hz)')
  40     freq_axes.plot(freq_axis, power, 'b.')
  41     if filename:
  42         figure.savefig(filename)
  43
  44 def run(*args, **kwargs):
  45     parser = _argparse.ArgumentParser(description=__doc__)
  46     parser.add_argument(
  47         '-F', '--frequency', type=float,
  48         help='frequency at which the data was aquired (in Hz)')
  49     parser.add_argument(
  50         '-t', '--title', help='plot title')
  51     parser.add_argument(
  52         '-f', '--filename',
  53         help='filename for saving the generated figure')
  54     parser.add_argument(
  55         'path', help='path to HDF5 file')
  56     parser.add_argument(
  57         'group', default='/', nargs='?', help='data group in the HDF5 file')
  58
  59     args = parser.parse_args(*args, **kwargs)
  60
  61     plot_spectrum(
  62         path=args.path, group=args.group, frequency=args.frequency,
  63         title=args.title, filename=args.filename)
  64     return args.filename
  65
  66
  67 if __name__ == '__main__':
  68     filename = run()
  69     if not filename:
  70         _pyplot.show()