Update imports for python 3 compatibility

benchmark/evaluation/__init__.py

@@ -1,7 +1,7 @@
-from analysis import estimate_nonequilibrium_free_energy, compute_free_energy
-from entropy import estimate_entropy, estimate_marginal_entropies
+from .analysis import estimate_nonequilibrium_free_energy, compute_free_energy
+from .entropy import estimate_entropy, estimate_marginal_entropies
 __all__ = ["estimate_nonequilibrium_free_energy",
            "compute_free_energy",
            "estimate_entropy",
            "estimate_marginal_entropies"
-           ]
+           ]

benchmark/evaluation/analysis.py

@@ -1,6 +1,6 @@
 import numpy as np
 from pickle import load
-from entropy import estimate_marginal_entropies, estimate_entropy
+from .entropy import estimate_marginal_entropies, estimate_entropy
 import matplotlib
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt
@@ -58,7 +58,7 @@ def compute_free_energy(xv, potential, kinetic_energy, beta):
 if __name__ == "__main__":
     print("reading and re-analyzing results")
     name = "alanine_unconstrained_null_results.pkl"
-    with open(name, "r") as f: results = load(f)
+    with open(name, "rb") as f: results = load(f)
     for scheme in results.keys():
         print(scheme)
         W_shads_F, W_shads_R, DeltaF_neq, sq_uncertainty, W_midpoints = results[scheme]

benchmark/experiments/baoab_vs_vvvr.py

@@ -6,24 +6,26 @@
 import pickle
 import numpy as np
 from simtk import unit
-from benchmark.plotting import plot_scheme_comparison
+#from baoab_vs_aboba_analysis import plot_results
 from benchmark.evaluation.analysis import estimate_nonequilibrium_free_energy
+from benchmark.plotting import plot, savefig
 
 if __name__ == "__main__":
     n_protocol_samples, protocol_length = 1000, 50
-    system_name = "alanine_unconstrained"
-    equilibrium_simulator = benchmark.testsystems.alanine_unconstrained
-    target_filename = os.path.join(DATA_PATH, "scheme_comparison_{}.pkl".format(system_name))
+    system_name = "alanine_constrained"
+    equilibrium_simulator = benchmark.testsystems.alanine_constrained
+    target_filename = os.path.join(DATA_PATH, "baoab_vs_aboba_{}.pkl".format(system_name))
 
-    schemes = {"BAO": "V R O", "BABO": "V R V O",
-               "ABO": "R V O", "ABAO": "R V R O"}
-    timesteps = np.linspace(0.1, 1.5, 5)
+    schemes = {"BAOAB": "V R O O R V", "VVVR": "O V R R V O"}
+    timesteps = np.linspace(0.1, 3.0, 10)
     noneq_simulators = {}
     for name, scheme in schemes.items():
         for timestep in timesteps:
             noneq_simulators[(name, timestep)] = NonequilibriumSimulator(equilibrium_simulator,
                                                          LangevinSplittingIntegrator(
-                                                             splitting=scheme, timestep=timestep * unit.femtosecond))
+                                                             splitting=scheme,
+                                                             timestep=timestep * unit.femtosecond,
+                                                             collision_rate=1.0/unit.picoseconds))
     # need to catch "Exception: Particle coordinate is nan"
 
     results = {}
@@ -40,7 +42,7 @@
             DeltaF_neq, squared_uncertainty = estimate_nonequilibrium_free_energy(W_shads_F, W_shads_R)
             print("\t{:.3f} +/- {:.3f}".format(DeltaF_neq, np.sqrt(squared_uncertainty)))
 
-    with open(target_filename, "w") as f:
+    with open(target_filename, "wb") as f:
         pickle.dump(results, f)
 
-    plot_scheme_comparison(target_filename, system_name)
+    # TODO: Plot stuff

benchmark/experiments/potential_energy_distributions.py

@@ -3,7 +3,7 @@
 
 import matplotlib
 
-from code.testsystems import system_params
+from benchmark.testsystems import system_params
 
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt
@@ -13,7 +13,7 @@
 from simtk import unit
 import numpy as np
 from tqdm import tqdm
-from code.utils import stderr
+from benchmark.utilities import stderr
 histstyle = {"alpha":0.3,
              "histtype":"stepfilled",
              "bins":50
@@ -98,4 +98,4 @@ def get_samples(simulation):
 
     from pickle import dump
 
-    with open("constrained_energies.pkl", "w") as f: dump(results, f)
+    with open("constrained_energies.pkl", "wb") as f: dump(results, f)

benchmark/experiments/quartic_kl_validation.py

@@ -58,7 +58,7 @@
             DeltaF_neq, squared_uncertainty = estimate_nonequilibrium_free_energy(W_shads_F, W_shads_R)
             print("\t{:.5f} +/- {:.5f}".format(DeltaF_neq, np.sqrt(squared_uncertainty)))
 
-    with open(target_filename, "w") as f:
+    with open(target_filename, "wb") as f:
         pickle.dump(results, f)
 
-    plot_scheme_comparison(target_filename, system_name)
+    plot_scheme_comparison(target_filename, system_name)

benchmark/experiments/quartic_protocol_length_dependence.py

@@ -48,7 +48,7 @@
         errors.append(np.sqrt(squared_uncertainty))
         print("\t{:.5f} +/- {:.5f}".format(DeltaF_neq, np.sqrt(squared_uncertainty)))
 
-    with open(data_filename, "w") as f:
+    with open(data_filename, "wb") as f:
         pickle.dump(results, f)
 
     plt.figure()

benchmark/experiments/replicate_paper_plot.py

@@ -5,12 +5,12 @@
 
 from benchmark.experiments.benchmark import get_equilibrium_samples, null_midpoint_operator
 from benchmark.integrators.integrators import LangevinSplittingIntegrator
-from code.testsystems import system_params
+from benchmark.testsystems import system_params
 
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt
 
-from code.timestep_search import sweep_over_timesteps
+from benchmark.experiments.timestep_search import sweep_over_timesteps
 
 if __name__ == "__main__":
     unconstrained_timesteps_to_try = [0.5, 0.75, 1, 1.25, 1.5]

benchmark/experiments/self_averaging.py

@@ -51,5 +51,5 @@ def oscillator_factory(coupling_strength=0.0):
         results[coupling_strength] = W_shads
 
 
-    with open(target_filename, "w") as f:
+    with open(target_filename, "wb") as f:
         pickle.dump(results, f)

benchmark/experiments/benchmark.py → benchmark/experiments/tools.py

@@ -2,16 +2,18 @@
 from simtk import unit
 from simtk.openmm import app
 
-from code.integrators import LangevinSplittingIntegrator
-from code.testsystems import system_params
+from benchmark.integrators import LangevinSplittingIntegrator
+from benchmark.testsystems import system_params
 
 W_unit = unit.kilojoule_per_mole
 from benchmark.evaluation.analysis import estimate_nonequilibrium_free_energy
 from pickle import dump
 from openmmtools.integrators import GHMCIntegrator
 from tqdm import tqdm
 
-from code.utils import plot, get_total_energy, get_summary_string, measure_shadow_work
+from benchmark.plotting import plot
+from benchmark.utilities import get_total_energy, get_summary_string
+from benchmark.tests.measure_shadow_work import measure_shadow_work
 
 
 def randomization_midpoint_operator(simulation, temperature):
@@ -93,7 +95,7 @@ def collect_and_save_results(schemes, simulation_factory, equilibrium_samples,
                            M=M, midpoint_operator=midpoint_operator, temperature=temperature)
         print("\n".join((scheme, get_summary_string(results[scheme], linebreaks=True))))
 
-    with open("{}_results.pkl".format(name), "w") as f: dump(results, f)
+    with open("{}_results.pkl".format(name), "wb") as f: dump(results, f)
 
     return results
 

benchmark/experiments/waterbox_test.py

@@ -12,11 +12,11 @@
 from simtk.openmm import app
 from tqdm import tqdm
 
-from benchmark.experiments.benchmark import estimate_nonequilibrium_free_energy
-from benchmark.experiments.benchmark import get_equilibrium_samples, \
+from benchmark.experiments.tools import estimate_nonequilibrium_free_energy
+from benchmark.experiments.tools import get_equilibrium_samples, \
     randomization_midpoint_operator, apply_protocol
 from benchmark.integrators.integrators import LangevinSplittingIntegrator
-from code.testsystems import system_params
+from benchmark.testsystems import system_params
 
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt
@@ -133,4 +133,4 @@ def plot_deltaF_neq(schemes, Ms, midpoint_operator):
     plot_deltaF_neq(schemes, Ms,
                     lambda simulation: randomization_midpoint_operator(simulation, temperature))
     plt.savefig("../figures/{}_conf_DeltaFs_over_Ms.jpg".format(name))
-    plt.close()
+    plt.close()

benchmark/integrators/__init__.py

@@ -1,9 +1,9 @@
-from ghmc import CustomizableGHMC
-from langevin import LangevinSplittingIntegrator
-from numba_integrators import baoab_factory, vvvr_factory, metropolis_hastings_factory, aboba_factory
-from mts_utilities import condense_splitting, generate_sequential_BAOAB_string, generate_all_BAOAB_permutation_strings
+from .ghmc import CustomizableGHMC
+from .langevin import LangevinSplittingIntegrator
+from .numba_integrators import baoab_factory, vvvr_factory, metropolis_hastings_factory, aboba_factory
+from .mts_utilities import condense_splitting, generate_sequential_BAOAB_string, generate_all_BAOAB_permutation_strings, generate_solvent_solute_splitting_string
 
 __all__ = ["CustomizableGHMC", "LangevinSplittingIntegrator",
            "baoab_factory", "vvvr_factory", "aboba_factory", "metropolis_hastings_factory",
            "condense_splitting", "generate_sequential_BAOAB_string", "generate_all_BAOAB_permutation_strings"
-           ]
+           ]

benchmark/integrators/langevin.py

@@ -122,6 +122,9 @@ def __init__(self,
 
         # Define substep functions
         def R_step():
+            self.addConstrainPositions()  # TODO: Constrain initial step only?
+            self.addConstrainVelocities() # TODO: Constrain initial step only?
+
             if measure_shadow_work:
                 self.addComputeGlobal("old_pe", "energy")
                 self.addComputeSum("old_ke", kinetic_energy)
@@ -147,6 +150,9 @@ def V_step(fg):
                 Force group to use in this substep.
                 "" means all forces, "0" means force-group 0, etc.
             """
+            self.addConstrainPositions()  # TODO: Constrain initial step only?
+            self.addConstrainVelocities() # TODO: Constrain initial step only?
+
             if measure_shadow_work:
                 self.addComputeSum("old_ke", kinetic_energy)
 
@@ -163,6 +169,9 @@ def V_step(fg):
                 self.addComputeGlobal("shadow_work", "shadow_work + (new_ke - old_ke)")
 
         def O_step():
+            self.addConstrainPositions()  # TODO: Constrain initial step only?
+            self.addConstrainVelocities() # TODO: Constrain initial step only?
+
             if measure_heat:
                 self.addComputeSum("old_ke", kinetic_energy)
 

benchmark/plotting/__init__.py

@@ -1,3 +1,3 @@
-from plotting_utilities import generate_figure_filename, plot_scheme_comparison
+from .plotting_utilities import savefig, plot, generate_figure_filename, plot_scheme_comparison
 
-__all__ = ["generate_figure_filename", "plot_scheme_comparison"]
+__all__ = ["generate_figure_filename", "plot_scheme_comparison"]

benchmark/plotting/plotting_utilities.py

@@ -6,6 +6,7 @@
 import pickle
 import os
 from benchmark import FIGURE_PATH
+from benchmark.utilities import unpack_trajs
 figure_directory = FIGURE_PATH
 figure_format = ".jpg"
 
@@ -129,7 +130,7 @@ def plot_scheme_comparison(target_filename, name):
     and one for the
         i.e. results[marginal][(name, timestep)] = W_shads_F, W_shads_R
     """
-    with open(target_filename, "r") as f:
+    with open(target_filename, "rb") as f:
         results = pickle.load(f)
 
     def get_schemes(result_dict):
@@ -172,4 +173,4 @@ def plot_curves(results):
         plt.savefig(generate_figure_filename("scheme_comparison_{}.jpg".format(name)), dpi=300)
         plt.close()
 
-    plot_curves(results)
+    plot_curves(results)

benchmark/serialization/serialize_integrators.py

@@ -25,7 +25,7 @@
     integrator = LangevinSplittingIntegrator(scheme)
 
     # Export integrator to XML
-    with open(os.path.join(DATA_PATH, "serialized_integrators/{}.xml".format(scheme), "w")) as f:
+    with open(os.path.join(DATA_PATH, "serialized_integrators/{}.xml".format(scheme), "wb")) as f:
         f.writelines(XmlSerializer.serialize(integrator))
 
     # Also pretty-print (using `step_type_dict`)
@@ -36,5 +36,5 @@
     print(readable_lines)
 
     # Save pretty-printed results to .txt
-    with open(os.path.join(DATA_PATH, "serialized_integrators/{}.txt".format(scheme), "w")) as f:
+    with open(os.path.join(DATA_PATH, "serialized_integrators/{}.txt".format(scheme), "wb")) as f:
         f.writelines(readable_lines)

benchmark/serialization/serialize_testsystems.py

@@ -8,5 +8,5 @@
         top, sys, pos = system_params[name]["loader"](constrained)
         c = "constrained"
         if not constrained: c = "unconstrained"
-        with open("serialized_testsystems/{}_{}.xml".format(name, c), "w") as f:
+        with open("serialized_testsystems/{}_{}.xml".format(name, c), "wb") as f:
             f.writelines(XmlSerializer.serialize(sys))

tests/measure_shadow_work.py → benchmark/tests/measure_shadow_work.py

@@ -89,4 +89,4 @@ def measure_shadow_work(simulation, n_steps):
     elif ("W_shad" in global_variable_names):
         return measure_shadow_work_via_W_shad(simulation, n_steps)
     else:
-        raise (RuntimeError("Simulation doesn't support shadow work computation"))
+        raise (RuntimeError("Simulation doesn't support shadow work computation"))

tests/substep_test.py → benchmark/tests/substep_test.py

@@ -1,5 +1,5 @@
-from tests import compare_substep_energy_changes, simulation_factory
-from code.utils import generate_solvent_solute_splitting_string
+from .tests import compare_substep_energy_changes, simulation_factory
+from benchmark.integrators import generate_solvent_solute_splitting_string
 
 for constrained in [True, False]:
     if constrained: print("\n\nWith constraints\n")
@@ -12,4 +12,4 @@
                    ]:
         simulation = simulation_factory(scheme)
         print("Scheme: {}".format(scheme))
-        compare_substep_energy_changes(simulation, n_steps=10)
+        compare_substep_energy_changes(simulation, n_steps=10)

tests/tests.py → benchmark/tests/tests.py

@@ -1,11 +1,13 @@
 import numpy as np
-from benchmark.integrators.integrators import LangevinSplittingIntegrator
+from benchmark.integrators import LangevinSplittingIntegrator
 from openmmtools.testsystems import AlanineDipeptideVacuum
 from simtk import unit
 from simtk.openmm import app
 
 W_unit = unit.kilojoule_per_mole
-from code.utils import configure_platform, get_total_energy, strip_unit, generate_solvent_solute_splitting_string
+from benchmark.testsystems.configuration import configure_platform
+from benchmark.utilities import get_total_energy, strip_unit
+from benchmark.integrators import generate_solvent_solute_splitting_string
 from benchmark.testsystems.testsystems import load_alanine
 
 
@@ -124,4 +126,4 @@ def check_W_shad_consistency(energy_changes):
     simulation.context.setPositions(testsystem.positions)
     simulation.context.setVelocitiesToTemperature(temperature)
     simulation.step(100)
-    check_W_shad_consistency(record_energy_changes(simulation, W_shad_name="shadow_work"))
+    check_W_shad_consistency(record_energy_changes(simulation, W_shad_name="shadow_work"))

benchmark/testsystems/__init__.py

@@ -1,9 +1,11 @@
-from alanine_dipeptide import alanine_constrained, alanine_unconstrained, solvated_alanine_unconstrained
-from waterbox import flexible_waterbox, waterbox_constrained
-from coupled_power_oscillators import coupled_power_oscillators
-from low_dimensional_systems import quartic, NumbaNonequilibriumSimulator
-from bookkeepers import EquilibriumSimulator, NonequilibriumSimulator
+from .alanine_dipeptide import alanine_constrained, alanine_unconstrained
+#from .alanine_dipeptide import solvated_alanine_constrained, solvated_alanine_unconstrained # TODO: Uncomment when this is working
+from .waterbox import flexible_waterbox, waterbox_constrained
+from .coupled_power_oscillators import coupled_power_oscillators
+from .low_dimensional_systems import quartic, NumbaNonequilibriumSimulator
+from .bookkeepers import EquilibriumSimulator, NonequilibriumSimulator
+from .testsystems import system_params
 
 __all__ = ["alanine_constrained", "alanine_unconstrained", "solvated_alanine_unconstrained",
            "flexible_waterbox", "waterbox_constrained", "coupled_power_oscillators",
-           "EquilibriumSimulator", "NonequilibriumSimulator", "quartic", "NumbaNonequilibriumSimulator"]
+           "EquilibriumSimulator", "NonequilibriumSimulator", "quartic", "NumbaNonequilibriumSimulator"]