A script to modify OpenFE transformation JSON files

shorten_runs.py

import argparse
import pathlib
import logging
import json

# this takes about 9 minutes to run on my Mac
SHORT_CONFIG = {
    'integrator_settings.n_steps.magnitude': 10,
    'simulation_settings.checkpoint_interval.magnitude': 10,
    'simulation_settings.equilibration_length.magnitude': 40,
    'simulation_settings.equilibration_length.unit': "femtosecond",
    'simulation_settings.production_length.magnitude': 1.0,
    'simulation_settings.production_length.unit': "picosecond",
}

def replace_config_item(config, loc, val):
    path = loc.split('.')
    key = path.pop(0)
    if len(path) == 0:
        config[key] = val
    else:
        next_config = config[key]
        next_loc = ".".join(path)
        replace_config_item(next_config, next_loc, val)


def replace_config(config, short_config):
    for loc, val in short_config.items():
        replace_config_item(config, loc, val)
    return config


if __name__ == "__main__":
    import sys
    logging.basicConfig(level=logging.INFO)
    with open(sys.argv[1]) as f:
        config = json.loads(f.read())

    settings = config['protocol']['settings']
    new_settings = replace_config(settings, SHORT_CONFIG)
    config['protocol']['settings'] = new_settings
    print(json.dumps(config))