VibrationAnalysis.py

#!/usr/bin/python
# (c) João Paulo Barraca <jpbarraca@gmail.com>
#
# TauLabs Vibration Analysis script to process the output of the VibrationAnalysis module.
#

def analyse_vibration(vibrationAnalysis,  bins=32, sample_rate=0.005):
    data = {'x': [], 'y': [], 'z': []}

    e = {'x': [], 'y': [], 'z': []}

    stats = {'x': {}, 'y': {}, 'z': {}}

    power_vals = 1+int(max(max(vibrationAnalysis['x']), max(vibrationAnalysis['y']), max(vibrationAnalysis['x']))[0])

    stats['x'] = {'data': [0] * bins, 'dist': [0] * power_vals, 'max': [0] * bins, 'count': 0, 'peak_freq': 0}
    stats['y'] = {'data': [0] * bins, 'dist': [0] * power_vals, 'max': [0] * bins, 'count': 0, 'peak_freq': 0}
    stats['z'] = {'data': [0] * bins, 'dist': [0] * power_vals, 'max': [0] * bins, 'count': 0, 'peak_freq': 0}

    for s in vibrationAnalysis:
        curr_bin = s['inst_id']
        if len(e['x']) > curr_bin:
            if len(e['x']) == bins:
                data['x'].append(e['x'])
                data['y'].append(e['y'])
                data['z'].append(e['z'])
                for axis in ['x', 'y', 'z']:
                    if sum(e[axis]) > 0:
                        for i, v in enumerate(e[axis]):
                            stats[axis]['data'][i] += v
                            if v > 0:
                                stats[axis]['dist'][int(v)] += 1

                            # Not really max in order to discard outliers
                            if max(stats[axis]['max'][i], v) == v:
                                stats[axis]['max'][i] = (stats[axis]['max'][i] + v) / 2.0

                        stats[axis]['count'] += 1

            e = {'x': [], 'y': [], 'z': []}

        e['x'].append(s['x'][0])
        e['y'].append(s['y'][0])
        e['z'].append(s['z'][0])

    if len(e['x']) == bins:
        data['x'].append(e['x'])
        data['y'].append(e['y'])
        data['z'].append(e['z'])

    #Calculate Average Vibration for each bin
    stats['x']['avg'] = [x/stats['x']['count'] for x in stats['x']['data']]
    stats['y']['avg'] = [y/stats['y']['count'] for y in stats['y']['data']]
    stats['z']['avg'] = [z/stats['z']['count'] for z in stats['z']['data']]

    #Calculate Frequency of Peak Average Values (Dominant Vibration Frequency)
    for axis in ['x', 'y', 'z']:
        m = -1
        stats[axis]['peak_freq'] = None
        for i,v in enumerate(stats[axis]['avg']):
            if v > m:
                m = v
                stats[axis]['peak_freq'] = i/(sample_rate * bins*2.0)

    print("##### Stats ##### ")
    print("Peak Vibration Intensity:\n\tx=%6.2f\n\ty=%6.2f\n\tz=%6.2f\n" % (max(stats['x']['max']),max(stats['y']['max']),max(stats['z']['max'])))
    print("Dominant Vibration Frequency:\n\tx=%6.2fHz (%4d RPM)\n\ty=%6.2fHz (%4d RPM)\n\tz=%6.2fHz (%4d RPM)\n" %
        (stats['x']['peak_freq'],stats['x']['peak_freq']/0.0166666666667,
         stats['y']['peak_freq'],stats['y']['peak_freq']/0.0166666666667,
         stats['z']['peak_freq'],stats['z']['peak_freq']/0.0166666666667,))


    print("Average Peak Vibration Intensity:\n\tx=%6.2f\n\ty=%6.2f\n\tz=%6.2f\n" %
          (max(stats['x']['data'])/stats['x']['count'],
           max(stats['y']['data'])/stats['y']['count'],
           max(stats['z']['data'])/stats['z']['count']))
    print("Cumulative Average Vibration Intensity:\n\tx=%6.2f\n\ty=%6.2f\n\tz=%6.2f\n" %
          (sum(stats['x']['data'])/stats['x']['count']/bins,
           sum(stats['y']['data'])/stats['y']['count']/bins,
           sum(stats['z']['data'])/stats['z']['count']/bins))


    return data, stats


def plot_waterfall(data, axis, bins, sample_rate):
    import matplotlib.pyplot as plt
    import numpy as np

    matrix = np.array(data[axis])
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(1, 1, 1)
    ax.set_aspect('equal')
    plt.imshow(matrix, interpolation='catrom', aspect='auto', cmap=plt.cm.gnuplot)
    plt.colorbar()
    ax.grid(True)
    ax.xaxis.set_major_locator(plt.MaxNLocator(bins+1))
    ax.set_xticklabels(np.linspace(0, (bins-1)/(sample_rate*bins*2), num=bins, dtype=int))
    plt.title('Vibration Intensity in %s direction' % axis)
    ax.xaxis.set_label_text('Frequency (Hz)')
    ax.yaxis.set_label_text('Sample Number')
    plt.tight_layout()
    fig.savefig('vibration_waterfall-%s.png' % axis, format='png')

    plt.show()


def plot_vibration_frequency(stats, bins, sample_rate):
    import matplotlib.pyplot as plt
    import numpy as np

    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(1, 1, 1)

    bins = len(stats['x']['data'])
    x = np.linspace(0, (bins)/(sample_rate*float(2*bins)), num=bins, dtype=int)

    plt.plot(x, stats['x']['max'], label="X Peak", ls='dashed', lw=1, color='red')
    plt.plot(x, stats['y']['max'], label="Y Peak", ls='dashed', lw=1, color='blue')
    plt.plot(x, stats['z']['max'], label="Z Peak", ls='dashed', lw=1, color='green')

    plt.plot(x, stats['x']['avg'], label="X Average", lw=2, color='red')
    plt.plot(x, stats['y']['avg'], label="Y Average", lw=2, color='blue')
    plt.plot(x, stats['z']['avg'], label="Z Average", lw=2, color='green')
    ax.xaxis.set_major_locator(plt.MaxNLocator(bins))
    ax.grid(True)
    plt.legend(prop={'size': 8})
    plt.title('Vibration Frequency Power')
    ax.xaxis.set_label_text('Frequency (Hz)')
    ax.yaxis.set_label_text('Average Intensity')
    plt.tight_layout()
    fig.savefig('vibration_frequency.png', format='png')
    plt.show()


def plot_vibration_density(stats, bins, sample_rate):
    import matplotlib.pyplot as plt

    bins = len(stats['x'])

    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(1, 1, 1)

    plt.plot(stats['x']['dist'], label="X", lw=2)
    plt.plot(stats['y']['dist'], label="Y", lw=2)
    plt.plot(stats['z']['dist'], label="Z", lw=2)
    ax.xaxis.set_major_locator(plt.MaxNLocator(bins))
    ax.xaxis.set_label_text('Intensity')
    ax.yaxis.set_label_text('Count')
    ax.grid(True)
    plt.title('Vibration Intensity Distribution')
    plt.legend()
    plt.tight_layout()
    fig.savefig('vibration_density.png', format='png')
    plt.show()

if __name__ == "__main__":
    import taulabs

    uavo_list = taulabs.telemetry.get_telemetry_by_args()

    vibrationAnalysis = uavo_list.as_numpy_array(taulabs.uavo.UAVO_VibrationAnalysisOutput)
    vibrationAnalysisSettings = uavo_list.as_numpy_array(taulabs.uavo.UAVO_VibrationAnalysisSettings)

    sample_rate = 0.001
    bins = 32
    try:
        sample_rate = vibrationAnalysisSettings['SampleRate'][0][0] * 0.001
        bins = max(vibrationAnalysis['inst_id'])+1
    except Exception, e:
        pass
    print vibrationAnalysisSettings
    print("SampleRate: %f Bins: %d Samples: %d" % (sample_rate, bins, len(vibrationAnalysis['time'])))

    data, stats = analyse_vibration(vibrationAnalysis, bins, sample_rate)
    plot_vibration_density(stats, bins, sample_rate)
    plot_vibration_frequency(stats, bins, sample_rate)
    plot_waterfall(data, 'x', bins, sample_rate)
    plot_waterfall(data, 'y', bins, sample_rate)
    plot_waterfall(data, 'z', bins, sample_rate)