Converts an OpenMM CustomIntegrator into a readable list of computation steps

pretty_print_customintegrator.py

from openmmtools.integrators import GHMCIntegrator, GradientDescentMinimizationIntegrator

integrators = [GHMCIntegrator(), GradientDescentMinimizationIntegrator()]

step_type_dict = {
    0 : "{target} <- {expr}",
    1: "{target} <- {expr}",
    2: "{target} <- sum({expr})",
    3: "constrain positions",
    4: "constrain velocities",
    5: "allow forces to update the context state",
    6: "if:",
    7: "while:",
    8: "end"
}

for integrator in integrators:
    scheme = integrator.__class__.__name__
    readable_lines = []
    for i in range(integrator.getNumComputations()):
        step_type, target, expr = integrator.getComputationStep(i)
        readable_line = step_type_dict[step_type].format(target=target, expr=expr)
        readable_lines.append(readable_line + "\n")
        print(readable_line)

    # Save pretty-printed results to .txt
    with open("{}.txt".format(scheme), "w") as f:
        f.writelines(readable_lines)