spectrum.py: add a simple script for plotting power spectra

Sometimes you want more frequency detail than the standard
calibcant-plot.py scripts will give you.  This script plots a single
(not chunked and averaged) power spectrum, which increases frequency
resolution at the expense of statistical significance.  Example usage
would be something like

  $ spectrum.py -f detail.png -t 'Vibration 15' -F 50000 \
    calibcant/2012-10-01T16-51-15-calibcant-data.h5 /vibration/15/raw/data

diff --git a/posts/Comparing_velocity_clamp_experiments/spectrum.py b/posts/Comparing_velocity_clamp_experiments/spectrum.py
new file mode 100755 (executable)
index 0000000..c9c6d47
--- /dev/null
+++ b/posts/Comparing_velocity_clamp_experiments/spectrum.py
@@ -0,0 +1,70 @@
+#!/usr/bin/env python2
+
+"Plot vibration data with higher low-frequency resolution"
+
+import argparse as _argparse
+
+import FFT_tools as _FFT_tools
+import h5py as _h5py
+import numpy as _numpy
+import matplotlib.pyplot as _pyplot
+import calibcant.vibration_analyze as _vibration_analyze
+
+
+def plot_spectrum(path, group='/', frequency=1, title=None, filename=None):
+    with _h5py.File(path) as f:
+        data = f[group][...]
+    print('loaded {} points from {}:{}'.format(data.shape, path, group))
+    dt = 1./frequency
+    time_axis = _numpy.arange(len(data)) * dt
+    vib_data = data - data.mean()  # give the FFT routine an easier time
+    freq_axis,power = _FFT_tools.unitary_power_spectrum(
+        vib_data, freq=frequency)
+
+    figure = _pyplot.figure()
+
+    time_axes = figure.add_subplot(2, 1, 1)
+    time_axes.autoscale(tight=True)
+    if title:
+        time_axes.set_title(title)
+    time_axes.set_xlabel('time (seconds)')
+    time_axes.set_ylabel('deflection (bits)')
+    time_axes.plot(time_axis, data, 'b.')
+
+    freq_axes = figure.add_subplot(2, 1, 2)
+    freq_axes.autoscale(tight=True)
+    freq_axes.set_xscale('log')
+    freq_axes.set_yscale('log')    
+    freq_axes.set_xlabel('frequency (Hz)')
+    freq_axes.set_ylabel('power (bits$^2$/Hz)')
+    freq_axes.plot(freq_axis, power, 'b.')
+    if filename:
+        figure.savefig(filename)
+
+def run(*args, **kwargs):
+    parser = _argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        '-F', '--frequency', type=float,
+        help='frequency at which the data was aquired (in Hz)')
+    parser.add_argument(
+        '-t', '--title', help='plot title')
+    parser.add_argument(
+        '-f', '--filename',
+        help='filename for saving the generated figure')
+    parser.add_argument(
+        'path', help='path to HDF5 file')
+    parser.add_argument(
+        'group', default='/', nargs='?', help='data group in the HDF5 file')
+
+    args = parser.parse_args(*args, **kwargs)
+
+    plot_spectrum(
+        path=args.path, group=args.group, frequency=args.frequency,
+        title=args.title, filename=args.filename)
+    return args.filename
+
+
+if __name__ == '__main__':
+    filename = run()
+    if not filename:
+        _pyplot.show()