bst_correlation.py

import argparse
import os
import time

import pandas as pd
import numpy as np
import scipy.signal
import matplotlib.pyplot as plt

from aeri_tools.io.dmv import housekeeping

def main(path):

    data = housekeeping.get_all_housekeeping(path)
    data = data[np.isfinite(data['BBsupportStructureTemp'])]
    data = data[np.isfinite(data['airNearBBsTemp'])]

    try:
        #subtract the rolling average
        for x in ('BBsupportStructureTemp', 'airNearBBsTemp'):
            tmp = scipy.signal.savgol_filter(data[x].values[:], 201, 3)
            data[x] = data[x] - tmp

        #correlate the data
        correlation = np.correlate(data['BBsupportStructureTemp'].values[:],
                                    data['airNearBBsTemp'].values[:], 'same')

        print('max = ', np.amax(correlation), ' : min = ',
                np.amin(correlation))

        fig, ax = plt.subplots(1, figsize=(15,10), sharex=True)

        '''
        for x, val in ('dmv', 'hk'):
            plt.sca(ax[x])
            plt.plot(list(data.index), data[val].values[:])
        '''

        plt.plot(data.index, correlation)

        name = '/Users/adiebold/Documents/corr_pngs/' + path[-12:-6] + '.png'
        # plt.savefig(name)
        plt.show()
        # plt.clf()
    #probably the scipy.signal.savgol_filter failing
    except:
        print('FAIL')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('filepath')
    args = parser.parse_args()
    print(args.filepath)

    #amount of files to skip
    skip_num = 0
    curr_num = 0
    print('skip_num = ', skip_num, '\n')
    if os.path.isdir(args.filepath):
        for filename_1 in os.listdir(args.filepath):
            filename_1 = args.filepath + '/' + filename_1
            filename_1 = filename_1.replace('//', '/')
            if os.path.isdir(filename_1):
                for filename_2 in os.listdir(filename_1):
                    filename_2 = filename_1 + '/' + filename_2
                    filename_2 = filename_2.replace('//', '/')
                    if (os.path.isfile(filename_2)
                    and filename_2.endswith('B1.CXS')):
                        curr_num += 1
                        if curr_num >= skip_num:
                            print(curr_num, ': ', filename_2)
                            main(filename_2)
                        else:
                            print(curr_num, ': ', filename_2, ' -- SKIPPED')
            elif os.path.isfile(filename_1) and filename_1.endswith('B1.CXS'):
                curr_num += 1
                if curr_num >= skip_num:
                    print(curr_num, ': ', filename_1)
                    main(filename_1)
                else:
                    print(curr_num, ': ', filename_1, ' -- SKIPPED')
    elif os.path.isfile(args.filepath):
        if args.filepath.endswith('B1.CXS'):
            print(args.filepath)
            main(args.filepath)