Add write speed to data

add_write_speed.py

import argparse

import numpy as np


def add_write_speed_to_file(filename):
    if 'test' in filename:
        write_score_filename = 'relevant_epoch_write_zscore_test.npbin'
    else:
        write_score_filename = 'relevant_epoch_write_zscore_train.npbin'
    write_data = np.fromfile(write_score_filename).reshape((-1, 2))
    other_data = np.fromfile(filename).reshape((-1, 22))
    file_offset = 0
    new_array = np.array([])

    for write_speed, ts in write_data:
        row = other_data[file_offset]
        row_ts = row[0]

        while not (row_ts <= ts <= row_ts + 300):
            # Instead of adding a new entry to new_array, just skip these values and increment
            # the file offset.
            file_offset += 1
            row = other_data[file_offset]
            row_ts = row[0]

        while row_ts <= ts <= row_ts + 300:
            new_array = np.append(new_array, np.append(row, [write_speed]))
            file_offset += 1
            row = other_data[file_offset]
            row_ts = row[0]

        if (file_offset % 1000) % 10 == 0:
                print "Still progressing. Now {0} lines into the file".format(file_offset)

        if file_offset >= len(other_data):
            break

    print "Successfully added all write speeds. Now writing new file..."
    with open('{filename}_with_write_speed.npbin'.format(filename=filename.split('.')[0]),
              'wb') as new_file:
        new_array.tofile(new_file)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--test_data', action='store_true', default=False)
    parser.add_argument('--train_data', action='store_true', default=False)
    args = parser.parse_args()

    if args.test_data:
        print "Adding write speed to test data..."
        add_write_speed_to_file('test_20pc.npbin')

    if args.train_data:
        print "Adding write speed to training data..."
        add_write_speed_to_file('train_20pc.npbin')