File Freezer for Dr. Es's Chemistry Research

freeze_files.py

#! /usr/bin/env python

import os
import sys


def generate_permutations(value_1, value_2, increment, iterations):
    permutations = set()
    for i in range(iterations):
        value_1 += increment * i
        value_2 += increment * i
        new_permutations = [
            (value_1, value_2),
            (value_1-increment, value_2),
            (value_1+increment, value_2),
            (value_1, value_2-increment),
            (value_1, value_2+increment),
            (value_1+increment, value_2+increment),
            (value_1-increment, value_2-increment),
            (value_1-increment, value_2+increment),
            (value_1+increment, value_2-increment),
        ]
        permutations |= set(new_permutations)
    return permutations


def freeze_file(file_name, increment, iterations):
    with open(file_name, 'r') as f:
        original = f.readlines()
    for i, lines in enumerate(original):
        if lines[0] == 'B':
            bond_index = i
            break
    bond_line_1 = original.pop(bond_index).split()
    bond_line_2 = original.pop(bond_index).split()
    value_1 = float(bond_line_1[-2])
    value_2 = float(bond_line_2[-2])
    permutations = generate_permutations(
        value_1, value_2, increment, iterations)
    for i, p in enumerate(permutations):
        bond_line_1[-2] = p[0]
        bond_line_2[-2] = p[1]
        original.insert(
            bond_index, ' '.join(str(s) for s in bond_line_2) + '\n')
        original.insert(
            bond_index, ' '.join(str(s) for s in bond_line_1) + '\n')
        dot = file_name.find('.')
        new_file_name = (file_name[0:dot]+'_'+str(i+1)+file_name[dot:])
        with open(new_file_name, 'w') as f:
            f.writelines(original)
        original.pop(bond_index)
        original.pop(bond_index)


def main(argv):
    if len(argv) < 4:
        print('usage: freeze_files increment iterations file_1 [file_2 ...]')
        return
    increment = float(argv[1])
    iterations = int(argv[2])
    files = argv[3:]
    for f in files:
        try:
            freeze_file(f, increment, iterations)
        except IOError as e:
            print(e)
            continue


if __name__ == '__main__':
    main(sys.argv)