예제 #1
0
def objective(**kwargs):
    groups = kwargs["groups"]
    hmc_integrators.guess_force_groups(system,
                                       others=groups[0],
                                       nonbonded=groups[1],
                                       fft=groups[2])
    integrator = lb_loader.kw_to_int(**kwargs)
    simulation = lb_loader.build(testsystem,
                                 integrator,
                                 temperature,
                                 precision=precision,
                                 platform_name=platform_name,
                                 state=state)
    integrator.reset_time()
    integrator.step(n_steps)
    new_state = simulation.context.getState(getPositions=True,
                                            getParameters=True)
    print("eff_ns_per_day=%f, eff_dt=%f" %
          (integrator.effective_ns_per_day,
           integrator.effective_timestep / u.femtoseconds))
    results = dict(
        n_steps=n_steps,
        intname=integrator.__class__.__name__,
        effective_ns_per_day=integrator.effective_ns_per_day,
        acceptance_rate=integrator.acceptance_rate,
        effective_timestep=integrator.effective_timestep / u.femtoseconds,
        time_per_step=integrator.time_per_step,
    )
    print(results)
    return results, new_state
예제 #2
0
def objective(**kwargs):
    groups = kwargs["groups"]
    hmc_integrators.guess_force_groups(system, others=groups[0], nonbonded=groups[1], fft=groups[2])
    integrator = lb_loader.kw_to_int(**kwargs)
    simulation = lb_loader.build(testsystem, integrator, temperature, precision=precision, platform_name=platform_name, state=state)
    integrator.reset_time()
    integrator.step(n_steps)
    new_state = simulation.context.getState(getPositions=True, getParameters=True)
    print("eff_ns_per_day=%f, eff_dt=%f" % (integrator.effective_ns_per_day, integrator.effective_timestep / u.femtoseconds))
    results = dict(
    n_steps=n_steps,
    intname=integrator.__class__.__name__,
    effective_ns_per_day=integrator.effective_ns_per_day,
    acceptance_rate=integrator.acceptance_rate,
    effective_timestep=integrator.effective_timestep / u.femtoseconds,
    time_per_step=integrator.time_per_step,
    )
    print(results)
    return results, new_state
예제 #3
0
def set_respa(system, scheme):
    """INPLACE on system!"""
    if len(scheme.values()) == 1:
        hmc_integrators.guess_force_groups(system, nonbonded=0, fft=0, others=0)
        groups = [(0, 1)]
    if len(scheme.values()) == 2:
        hmc_integrators.guess_force_groups(system, nonbonded=0, fft=1, others=0)
        groups = [(0, scheme["others"]), (1, scheme["nonbonded"])]
    if len(scheme.values()) == 3:
        hmc_integrators.guess_force_groups(system, nonbonded=1, fft=2, others=0)
        groups = [(0, scheme["others"]), (1, scheme["nonbonded"]), (2, scheme["fft"])]
    return groups
예제 #4
0
testsystem = testsystems.DHFRExplicit(rigid_water=False, constraints=None)

system, topology, positions = testsystem.system, testsystem.topology, testsystem.positions

platform = mm.Platform.getPlatformByName('CUDA')
properties = {'CudaPrecision': "single"}

temperature = 1 * u.kelvin
timestep = 1.0 * u.femtoseconds
steps = 5000

#hmc_integrators.guess_force_groups(system, nonbonded=0, others=1, fft=0)
#groups = [(0, 1), (1, 4)]

hmc_integrators.guess_force_groups(system, nonbonded=0, others=0, fft=0)
groups = [(0, 1)]


integrator = mm.MTSIntegrator(timestep, groups)
simulation = app.Simulation(topology, system, integrator, platform=platform, platformProperties=properties)

simulation.context.setPositions(positions)
simulation.context.setVelocitiesToTemperature(temperature)

integrator.step(1)

t0 = time.time()
integrator.step(steps)
dt = time.time() - t0
outer_per_day = steps / dt * 60 * 60 * 24
예제 #5
0
def load(sysname):
    cutoff = 0.9 * u.nanometers
    temperature = 300. * u.kelvin
    langevin_timestep = 0.5 * u.femtoseconds
    timestep = 2 * u.femtoseconds
    equil_steps = 40000
    groups = [(0, 1)]
    steps_per_hmc = 25

    if sysname == "chargedljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj(
            charge=0.15 * u.elementary_charge)

    if sysname == "chargedswitchedljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj(
            charge=0.15 * u.elementary_charge,
            switch_width=0.34 * u.nanometers)

    if sysname == "chargedlongljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj(
            cutoff=1.333 * u.nanometers, charge=0.15 * u.elementary_charge)

    if sysname == "chargedswitchedlongljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj(
            cutoff=1.333 * u.nanometers,
            charge=0.15 * u.elementary_charge,
            switch_width=0.34 * u.nanometers)

    if sysname == "chargedswitchedaccuratelongljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj(
            cutoff=1.333 * u.nanometers,
            charge=0.15 * u.elementary_charge,
            switch_width=0.34 * u.nanometers,
            ewaldErrorTolerance=5E-5)

    if sysname == "chargedswitchedaccurateljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj(
            charge=0.15 * u.elementary_charge,
            switch_width=0.34 * u.nanometers,
            ewaldErrorTolerance=5E-5)

    if sysname == "ljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj()

    if sysname == "longljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj(
            cutoff=1.333 * u.nanometers)

    if sysname == "shortljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj(
            cutoff=0.90 * u.nanometers)

    if sysname == "shiftedljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj(
            shift=True)

    if sysname == "switchedljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj(
            switch_width=0.34 * u.nanometers)

    if sysname == "switchedshortljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj(
            cutoff=0.90 * u.nanometers, switch_width=0.34 * u.nanometers)

    if sysname == "cluster":
        testsystem = testsystems.LennardJonesCluster(nx=8, ny=8, nz=8)
        system, positions = testsystem.system, testsystem.positions
        hmc_integrators.guess_force_groups(system, nonbonded=0)
        groups = [(0, 1)]
        temperature = 25. * u.kelvin
        timestep = 40 * u.femtoseconds

    if sysname == "shortbigcluster":
        testsystem = testsystems.LennardJonesCluster(nx=20,
                                                     ny=20,
                                                     nz=20,
                                                     cutoff=0.75 *
                                                     u.nanometers)
        system, positions = testsystem.system, testsystem.positions
        hmc_integrators.guess_force_groups(system, nonbonded=0)
        groups = [(0, 1)]
        temperature = 25. * u.kelvin
        timestep = 10 * u.femtoseconds

    if sysname == "switchedshortbigcluster":
        testsystem = testsystems.LennardJonesCluster(
            nx=20,
            ny=20,
            nz=20,
            cutoff=0.75 * u.nanometers,
            switch_width=0.1 * u.nanometers)
        system, positions = testsystem.system, testsystem.positions
        hmc_integrators.guess_force_groups(system, nonbonded=0)
        groups = [(0, 1)]
        temperature = 25. * u.kelvin
        timestep = 10 * u.femtoseconds

    if sysname == "bigcluster":
        testsystem = testsystems.LennardJonesCluster(nx=20,
                                                     ny=20,
                                                     nz=20,
                                                     cutoff=1.25 *
                                                     u.nanometers)
        system, positions = testsystem.system, testsystem.positions
        hmc_integrators.guess_force_groups(system, nonbonded=0)
        groups = [(0, 1)]
        temperature = 25. * u.kelvin
        timestep = 10 * u.femtoseconds

    if sysname == "shortcluster":
        testsystem = testsystems.LennardJonesCluster(nx=8,
                                                     ny=8,
                                                     nz=8,
                                                     cutoff=0.75 *
                                                     u.nanometers)
        system, positions = testsystem.system, testsystem.positions
        hmc_integrators.guess_force_groups(system, nonbonded=0)
        groups = [(0, 1)]
        temperature = 25. * u.kelvin
        timestep = 40 * u.femtoseconds

    if sysname == "shortwater":
        testsystem = testsystems.WaterBox(
            box_edge=3.18 * u.nanometers,
            cutoff=0.9 * u.nanometers,
            switch_width=None)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions

    if sysname == "shortswitchedwater":
        testsystem = testsystems.WaterBox(
            box_edge=3.18 * u.nanometers,
            cutoff=0.9 * u.nanometers,
            switch_width=3.0 * u.angstroms)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions

    if sysname == "water":
        testsystem = testsystems.WaterBox(
            box_edge=3.18 * u.nanometers,
            cutoff=1.1 * u.nanometers,
            switch_width=None)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions

    if sysname == "switchedwater":
        testsystem = testsystems.WaterBox(
            box_edge=3.18 * u.nanometers,
            cutoff=1.1 * u.nanometers,
            switch_width=3.0 * u.angstroms)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions

    if sysname == "switchedaccuratewater":
        testsystem = testsystems.WaterBox(
            box_edge=3.18 * u.nanometers,
            cutoff=1.1 * u.nanometers,
            switch_width=3.0 * u.angstroms,
            ewaldErrorTolerance=5E-5)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions
        hmc_integrators.guess_force_groups(
            system, nonbonded=0, fft=1, others=0)  # We may want to try reduce

    if sysname == "switchedaccurateflexiblewater":
        testsystem = testsystems.WaterBox(
            box_edge=3.18 * u.nanometers,
            cutoff=1.1 * u.nanometers,
            switch_width=3.0 * u.angstroms,
            ewaldErrorTolerance=5E-5,
            constrained=False)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions
        # Using these groups for hyperopt-optimal RESPA integrators
        #hmc_integrators.guess_force_groups(system, nonbonded=0, others=1, fft=0)
        hmc_integrators.guess_force_groups(system,
                                           nonbonded=1,
                                           others=1,
                                           fft=0)
        equil_steps = 100000
        langevin_timestep = 0.25 * u.femtoseconds

    if sysname == "switchedaccuratebigflexiblewater":
        testsystem = testsystems.WaterBox(
            box_edge=10.0 * u.nanometers,
            cutoff=1.1 * u.nanometers,
            switch_width=3.0 * u.angstroms,
            ewaldErrorTolerance=5E-5,
            constrained=False)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions
        # Using these groups for hyperopt-optimal RESPA integrators
        hmc_integrators.guess_force_groups(system,
                                           nonbonded=1,
                                           others=0,
                                           fft=1)
        equil_steps = 100000
        langevin_timestep = 0.1 * u.femtoseconds

    if sysname == "switchedaccuratestiffflexiblewater":

        testsystem = testsystems.WaterBox(
            box_edge=3.18 * u.nanometers,
            cutoff=1.1 * u.nanometers,
            switch_width=3.0 * u.angstroms,
            ewaldErrorTolerance=5E-5,
            constrained=False)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions
        # Using these groups for hyperopt-optimal RESPA integrators
        #hmc_integrators.guess_force_groups(system, nonbonded=0, others=1, fft=0)
        hmc_integrators.guess_force_groups(system,
                                           nonbonded=1,
                                           others=1,
                                           fft=0)
        equil_steps = 100000
        langevin_timestep = 0.25 * u.femtoseconds

        FACTOR = 10.0
        # Make bonded terms stiffer to highlight RESPA advantage
        for force in system.getForces():
            if type(force) == mm.HarmonicBondForce:
                for i in range(force.getNumBonds()):
                    (a0, a1, length, strength) = force.getBondParameters(i)
                    force.setBondParameters(i, a0, a1, length,
                                            strength * FACTOR)

            elif type(force) == mm.HarmonicAngleForce:
                for i in range(force.getNumAngles()):
                    (a0, a1, a2, length,
                     strength) = force.getAngleParameters(i)
                    force.setAngleParameters(i, a0, a1, a2, length,
                                             strength * FACTOR)

    if sysname == "switchedaccuratenptwater":
        testsystem = testsystems.WaterBox(
            box_edge=3.18 * u.nanometers,
            cutoff=1.1 * u.nanometers,
            switch_width=3.0 * u.angstroms,
            ewaldErrorTolerance=5E-5)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions
        system.addForce(
            mm.MonteCarloBarostat(1.0 * u.atmospheres, temperature, 1))

    if sysname == "longswitchedwater":
        testsystem = testsystems.WaterBox(
            box_edge=3.18 * u.nanometers,
            cutoff=1.5 * u.nanometers,
            switch_width=3.0 * u.angstroms)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions

    if sysname == "rfwater":
        testsystem = testsystems.WaterBox(
            box_edge=3.18 * u.nanometers,
            cutoff=1.1 * u.nanometers,
            nonbondedMethod=app.CutoffPeriodic,
            switch_width=None)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions

    if sysname == "switchedrfwater":
        testsystem = testsystems.WaterBox(
            box_edge=3.18 * u.nanometers,
            cutoff=1.1 * u.nanometers,
            nonbondedMethod=app.CutoffPeriodic,
            switch_width=3.0 * u.angstroms)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions

    if sysname == "longswitchedrfwater":
        testsystem = testsystems.WaterBox(
            box_edge=3.18 * u.nanometers,
            cutoff=1.5 * u.nanometers,
            nonbondedMethod=app.CutoffPeriodic,
            switch_width=3.0 * u.angstroms)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions

    if sysname == "density":
        system, positions = load_lb(hydrogenMass=3.0 * u.amu)
        hmc_integrators.guess_force_groups(system,
                                           nonbonded=1,
                                           fft=1,
                                           others=0)
        groups = [(0, 4), (1, 1)]
        timestep = 2.0 * u.femtoseconds

    if sysname == "dhfr":
        testsystem = testsystems.DHFRExplicit(nonbondedCutoff=1.1 *
                                              u.nanometers,
                                              nonbondedMethod=app.PME,
                                              switch_width=2.0 * u.angstroms,
                                              ewaldErrorTolerance=5E-5)
        system, positions = testsystem.system, testsystem.positions
        hmc_integrators.guess_force_groups(system,
                                           nonbonded=1,
                                           fft=2,
                                           others=0)
        groups = [(0, 4), (1, 2), (2, 1)]
        timestep = 0.75 * u.femtoseconds
        equil_steps = 10000

    if sysname == "src":
        testsystem = testsystems.SrcExplicit(nonbondedCutoff=1.1 *
                                             u.nanometers,
                                             nonbondedMethod=app.PME,
                                             switch_width=2.0 * u.angstroms,
                                             ewaldErrorTolerance=5E-5)
        system, positions = testsystem.system, testsystem.positions
        hmc_integrators.guess_force_groups(system,
                                           nonbonded=1,
                                           fft=1,
                                           others=0)
        groups = [(0, 2), (1, 1)]
        timestep = 0.25 * u.femtoseconds
        equil_steps = 1000

    if sysname == "ho":
        K = 90.0 * u.kilocalories_per_mole / u.angstroms**2
        mass = 39.948 * u.amu
        timestep = np.sqrt(mass / K) * 0.4
        testsystem = testsystems.HarmonicOscillatorArray()
        system, positions = testsystem.system, testsystem.positions
        groups = [(0, 1)]

    if sysname == "customho":
        timestep = 1000.0 * u.femtoseconds
        testsystem = testsystems.CustomExternalForcesTestSystem()
        system, positions = testsystem.system, testsystem.positions
        groups = [(0, 1)]

    if sysname == "customsplitho":
        timestep = 1000.0 * u.femtoseconds
        energy_expressions = ("0.75 * (x^2 + y^2 + z^2)",
                              "0.25 * (x^2 + y^2 + z^2)")
        testsystem = testsystems.CustomExternalForcesTestSystem(
            energy_expressions=energy_expressions)
        system, positions = testsystem.system, testsystem.positions
        groups = [(0, 2), (1, 1)]

    if sysname == "alanine":
        testsystem = testsystems.AlanineDipeptideImplicit()
        system, positions = testsystem.system, testsystem.positions
        groups = [(0, 2), (1, 1)]
        timestep = 2.0 * u.femtoseconds
        hmc_integrators.guess_force_groups(system, nonbonded=1, others=0)
        #remove_cmm(system)  # Unrolled shouldn't need this
        equil_steps = 10000

    if sysname == "alanineexplicit":
        testsystem = testsystems.AlanineDipeptideExplicit(
            cutoff=1.1 * u.nanometers,
            switch_width=2 * u.angstrom,
            ewaldErrorTolerance=5E-5)
        system, positions = testsystem.system, testsystem.positions
        #groups = [(0, 2), (1, 1), (2, 1)]
        #groups = [(0, 2), (1, 1)]
        timestep = 1.0 * u.femtoseconds
        #hmc_integrators.guess_force_groups(system, nonbonded=1, others=0, fft=2)
        #hmc_integrators.guess_force_groups(system, nonbonded=0, others=0, fft=1)
        groups = [(0, 1), (1, 2)]
        hmc_integrators.guess_force_groups(system,
                                           nonbonded=0,
                                           others=1,
                                           fft=0)

        #remove_cmm(system)  # Unrolled doesn't need this
        equil_steps = 4000

    # guess force groups

    if "ljbox" in sysname:
        timestep = 25 * u.femtoseconds
        temperature = 25. * u.kelvin
        system, positions = testsystem.system, testsystem.positions
        hmc_integrators.guess_force_groups(system, nonbonded=0, fft=1)
        groups = [(0, 2), (1, 1)]
        equil_steps = 10000
        steps_per_hmc = 15

    if sysname == "amoeba":
        testsystem = testsystems.AMOEBAIonBox()
        system, positions = testsystem.system, testsystem.positions
        hmc_integrators.guess_force_groups(system,
                                           nonbonded=0,
                                           fft=1,
                                           others=0)

    return system, positions, groups, temperature, timestep, langevin_timestep, testsystem, equil_steps, steps_per_hmc
import lb_loader
import simtk.openmm.app as app
import numpy as np
import pandas as pd
import simtk.openmm as mm
from simtk import unit as u
from openmmtools import hmc_integrators, testsystems
pd.set_option('display.width', 1000)

n_steps = 3000
temperature = 300. * u.kelvin

hydrogenMass = 3.0 * u.amu

system, positions = lb_loader.load_amoeba()
hmc_integrators.guess_force_groups(system, multipole=2)

#system, positions = lb_loader.load_lb(hydrogenMass=hydrogenMass)
#hmc_integrators.guess_force_groups(system, nonbonded=1, fft=1, others=0)

integrator = mm.LangevinIntegrator(temperature, 1.0 / u.picoseconds, 0.25 * u.femtoseconds)
context = mm.Context(system, integrator)
context.setPositions(positions)
context.setVelocitiesToTemperature(temperature)

integrator.step(1000)
positions = context.getState(getPositions=True).getPositions()
state = context.getState(getEnergy=True)
energy = state.getPotentialEnergy() + state.getKineticEnergy()
energy, state.getPotentialEnergy(), state.getKineticEnergy()
예제 #7
0
import lb_loader
import numpy as np
import pandas as pd
import simtk.openmm as mm
from simtk import unit as u
from openmmtools import hmc_integrators, testsystems

collision_rate = 1.0 / u.picoseconds
n_steps = 1500
temperature = 300. * u.kelvin

system, positions = lb_loader.load_lb()
hmc_integrators.guess_force_groups(system)

positions = lb_loader.pre_equil(system, positions, temperature)

def test_hmc(timestep, steps_per_hmc):
    timestep = timestep * u.femtoseconds
    integrator = hmc_integrators.GHMC2(temperature, steps_per_hmc, timestep)
    context = mm.Context(system, integrator)
    context.setPositions(positions)
    context.setVelocitiesToTemperature(temperature)
    integrator.step(n_steps)
    return integrator.acceptance_rate

timestep_list = np.linspace(1.5, 2.25, 3)
steps_per_hmc_list = np.array([5, 10, 25])
#steps_per_hmc_list = np.array([10, 5, 25])
data = []
for i, timestep in enumerate(timestep_list):
    for j, steps_per_hmc in enumerate(steps_per_hmc_list):
positions, boxes, state = lb_loader.equilibrate(testsystem, temperature, langevin_timestep, steps=equil_steps, minimize=True, use_hmc=False, precision=precision, platform_name=platform_name)

int_args = dict(
temperature = temperature,
timestep = 1.0,
steps_per_hmc = 1,
collision_rate = 1.0 / u.picoseconds,
extra_chances = 0,
group0_iterations = 2,
groups=(0,1,1),
)

collision_rate = int_args["collision_rate"]

groups = int_args["groups"]
hmc_integrators.guess_force_groups(system, others=groups[0], nonbonded=groups[1], fft=groups[2])

steps = 1000

integrator = lb_loader.kw_to_int(**int_args)
#integrator = integrators.HMCIntegrator(temperature=temperature, nsteps=steps_per_hmc, timestep=timestep)
#integrator = hmc_integrators.GHMCIntegrator(temperature=temperature, steps_per_hmc=steps_per_hmc, timestep=timestep, collision_rate=collision_rate)

#temperature = int_args["temperature"]
#steps_per_hmc = int_args["steps_per_hmc"]
#timestep = int_args["timestep"]
groups = [(0, 1)]
extra_chances = 1
#collision_rate = None
integrator = hmc_integrators.GHMCIntegrator(temperature=temperature, steps_per_hmc=steps_per_hmc, timestep=timestep, collision_rate=collision_rate)
#integrator = hmc_integrators.GHMCRESPAIntegrator(temperature, steps_per_hmc=steps_per_hmc, timestep=timestep, groups=groups)
예제 #9
0
def load(sysname):
    cutoff = 0.9 * u.nanometers
    temperature = 300. * u.kelvin
    langevin_timestep = 0.5 * u.femtoseconds
    timestep = 2 * u.femtoseconds
    equil_steps = 40000
    groups = [(0, 1)]
    steps_per_hmc = 25

    if sysname == "chargedljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj(charge=0.15*u.elementary_charge)

    if sysname == "chargedswitchedljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj(charge=0.15*u.elementary_charge, switch_width=0.34*u.nanometers)

    if sysname == "chargedlongljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj(cutoff=1.333*u.nanometers, charge=0.15*u.elementary_charge)

    if sysname == "chargedswitchedlongljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj(cutoff=1.333*u.nanometers, charge=0.15*u.elementary_charge, switch_width=0.34*u.nanometers)

    if sysname == "chargedswitchedaccuratelongljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj(cutoff=1.333*u.nanometers, charge=0.15*u.elementary_charge, switch_width=0.34*u.nanometers, ewaldErrorTolerance=5E-5)

    if sysname == "chargedswitchedaccurateljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj(charge=0.15*u.elementary_charge, switch_width=0.34*u.nanometers, ewaldErrorTolerance=5E-5)

    if sysname == "ljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj()

    if sysname == "longljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj(cutoff=1.333*u.nanometers)

    if sysname == "shortljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj(cutoff=0.90*u.nanometers)

    if sysname == "shiftedljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj(shift=True)

    if sysname == "switchedljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj(switch_width=0.34*u.nanometers)

    if sysname == "switchedshortljbox":
        testsystem, system, positions, timestep, langevin_timestep = load_lj(cutoff=0.90*u.nanometers, switch_width=0.34*u.nanometers)

    if sysname == "cluster":
        testsystem = testsystems.LennardJonesCluster(nx=8, ny=8, nz=8)
        system, positions = testsystem.system, testsystem.positions
        hmc_integrators.guess_force_groups(system, nonbonded=0)
        groups = [(0, 1)]
        temperature = 25. * u.kelvin
        timestep = 40 * u.femtoseconds

    if sysname == "shortbigcluster":
        testsystem = testsystems.LennardJonesCluster(nx=20, ny=20, nz=20, cutoff=0.75*u.nanometers)
        system, positions = testsystem.system, testsystem.positions
        hmc_integrators.guess_force_groups(system, nonbonded=0)
        groups = [(0, 1)]
        temperature = 25. * u.kelvin
        timestep = 10 * u.femtoseconds

    if sysname == "switchedshortbigcluster":
        testsystem = testsystems.LennardJonesCluster(nx=20, ny=20, nz=20, cutoff=0.75*u.nanometers, switch_width=0.1*u.nanometers)
        system, positions = testsystem.system, testsystem.positions
        hmc_integrators.guess_force_groups(system, nonbonded=0)
        groups = [(0, 1)]
        temperature = 25. * u.kelvin
        timestep = 10 * u.femtoseconds

    if sysname == "bigcluster":
        testsystem = testsystems.LennardJonesCluster(nx=20, ny=20, nz=20, cutoff=1.25*u.nanometers)
        system, positions = testsystem.system, testsystem.positions
        hmc_integrators.guess_force_groups(system, nonbonded=0)
        groups = [(0, 1)]
        temperature = 25. * u.kelvin
        timestep = 10 * u.femtoseconds

    if sysname == "shortcluster":
        testsystem = testsystems.LennardJonesCluster(nx=8, ny=8, nz=8, cutoff=0.75*u.nanometers)
        system, positions = testsystem.system, testsystem.positions
        hmc_integrators.guess_force_groups(system, nonbonded=0)
        groups = [(0, 1)]
        temperature = 25. * u.kelvin
        timestep = 40 * u.femtoseconds

    if sysname == "shortwater":
        testsystem = testsystems.WaterBox(box_edge=3.18 * u.nanometers, cutoff=0.9*u.nanometers, switch_width=None)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions

    if sysname == "shortswitchedwater":
        testsystem = testsystems.WaterBox(box_edge=3.18 * u.nanometers, cutoff=0.9*u.nanometers, switch_width=3.0*u.angstroms)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions

    if sysname == "water":
        testsystem = testsystems.WaterBox(box_edge=3.18 * u.nanometers, cutoff=1.1*u.nanometers, switch_width=None)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions

    if sysname == "switchedwater":
        testsystem = testsystems.WaterBox(box_edge=3.18 * u.nanometers, cutoff=1.1*u.nanometers, switch_width=3.0*u.angstroms)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions

    if sysname == "switchedaccuratewater":
        testsystem = testsystems.WaterBox(box_edge=3.18 * u.nanometers, cutoff=1.1*u.nanometers, switch_width=3.0*u.angstroms, ewaldErrorTolerance=5E-5)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions
        hmc_integrators.guess_force_groups(system, nonbonded=0, fft=1, others=0)  # We may want to try reduce

    if sysname == "switchedaccurateflexiblewater":
        testsystem = testsystems.WaterBox(box_edge=3.18 * u.nanometers, cutoff=1.1*u.nanometers, switch_width=3.0*u.angstroms, ewaldErrorTolerance=5E-5, constrained=False)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions
        # Using these groups for hyperopt-optimal RESPA integrators
        #hmc_integrators.guess_force_groups(system, nonbonded=0, others=1, fft=0)
        hmc_integrators.guess_force_groups(system, nonbonded=1, others=1, fft=0)
        equil_steps = 100000
        langevin_timestep = 0.25 * u.femtoseconds

    if sysname == "switchedaccuratebigflexiblewater":
        testsystem = testsystems.WaterBox(box_edge=10.0 * u.nanometers, cutoff=1.1*u.nanometers, switch_width=3.0*u.angstroms, ewaldErrorTolerance=5E-5, constrained=False)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions
        # Using these groups for hyperopt-optimal RESPA integrators
        hmc_integrators.guess_force_groups(system, nonbonded=1, others=0, fft=1)
        equil_steps = 100000
        langevin_timestep = 0.1 * u.femtoseconds


    if sysname == "switchedaccuratestiffflexiblewater":

        testsystem = testsystems.WaterBox(box_edge=3.18 * u.nanometers, cutoff=1.1*u.nanometers, switch_width=3.0*u.angstroms, ewaldErrorTolerance=5E-5, constrained=False)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions
        # Using these groups for hyperopt-optimal RESPA integrators
        #hmc_integrators.guess_force_groups(system, nonbonded=0, others=1, fft=0)
        hmc_integrators.guess_force_groups(system, nonbonded=1, others=1, fft=0)
        equil_steps = 100000
        langevin_timestep = 0.25 * u.femtoseconds

        FACTOR = 10.0
        # Make bonded terms stiffer to highlight RESPA advantage
        for force in system.getForces():
            if type(force) == mm.HarmonicBondForce:
                for i in range(force.getNumBonds()):
                    (a0, a1, length, strength) = force.getBondParameters(i)
                    force.setBondParameters(i, a0, a1, length, strength * FACTOR)

            elif type(force) == mm.HarmonicAngleForce:
                for i in range(force.getNumAngles()):
                    (a0, a1, a2, length, strength) = force.getAngleParameters(i)
                    force.setAngleParameters(i, a0, a1, a2, length, strength * FACTOR)




    if sysname == "switchedaccuratenptwater":
        testsystem = testsystems.WaterBox(box_edge=3.18 * u.nanometers, cutoff=1.1*u.nanometers, switch_width=3.0*u.angstroms, ewaldErrorTolerance=5E-5)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions
        system.addForce(mm.MonteCarloBarostat(1.0 * u.atmospheres, temperature, 1))

    if sysname == "longswitchedwater":
        testsystem = testsystems.WaterBox(box_edge=3.18 * u.nanometers, cutoff=1.5*u.nanometers, switch_width=3.0*u.angstroms)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions

    if sysname == "rfwater":
        testsystem = testsystems.WaterBox(box_edge=3.18 * u.nanometers, cutoff=1.1*u.nanometers, nonbondedMethod=app.CutoffPeriodic, switch_width=None)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions

    if sysname == "switchedrfwater":
        testsystem = testsystems.WaterBox(box_edge=3.18 * u.nanometers, cutoff=1.1*u.nanometers, nonbondedMethod=app.CutoffPeriodic, switch_width=3.0*u.angstroms)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions

    if sysname == "longswitchedrfwater":
        testsystem = testsystems.WaterBox(box_edge=3.18 * u.nanometers, cutoff=1.5*u.nanometers, nonbondedMethod=app.CutoffPeriodic, switch_width=3.0*u.angstroms)  # Around 1060 molecules of water
        system, positions = testsystem.system, testsystem.positions

    if sysname == "density":
        system, positions = load_lb(hydrogenMass=3.0 * u.amu)
        hmc_integrators.guess_force_groups(system, nonbonded=1, fft=1, others=0)
        groups = [(0, 4), (1, 1)]
        timestep = 2.0 * u.femtoseconds

    if sysname == "dhfr":
        testsystem = testsystems.DHFRExplicit(nonbondedCutoff=1.1*u.nanometers, nonbondedMethod=app.PME, switch_width=2.0*u.angstroms, ewaldErrorTolerance=5E-5)
        system, positions = testsystem.system, testsystem.positions
        hmc_integrators.guess_force_groups(system, nonbonded=1, fft=2, others=0)
        groups = [(0, 4), (1, 2), (2, 1)]
        timestep = 0.75 * u.femtoseconds
        equil_steps = 10000


    if sysname == "src":
        testsystem = testsystems.SrcExplicit(nonbondedCutoff=1.1*u.nanometers, nonbondedMethod=app.PME, switch_width=2.0*u.angstroms, ewaldErrorTolerance=5E-5)
        system, positions = testsystem.system, testsystem.positions
        hmc_integrators.guess_force_groups(system, nonbonded=1, fft=1, others=0)
        groups = [(0, 2), (1, 1)]
        timestep = 0.25 * u.femtoseconds
        equil_steps = 1000

    if sysname == "ho":
        K = 90.0 * u.kilocalories_per_mole / u.angstroms**2
        mass = 39.948 * u.amu
        timestep = np.sqrt(mass / K) * 0.4
        testsystem = testsystems.HarmonicOscillatorArray()
        system, positions = testsystem.system, testsystem.positions
        groups = [(0, 1)]

    if sysname == "customho":
        timestep = 1000.0 * u.femtoseconds
        testsystem = testsystems.CustomExternalForcesTestSystem()
        system, positions = testsystem.system, testsystem.positions
        groups = [(0, 1)]

    if sysname == "customsplitho":
        timestep = 1000.0 * u.femtoseconds
        energy_expressions = ("0.75 * (x^2 + y^2 + z^2)", "0.25 * (x^2 + y^2 + z^2)")
        testsystem = testsystems.CustomExternalForcesTestSystem(energy_expressions=energy_expressions)
        system, positions = testsystem.system, testsystem.positions
        groups = [(0, 2), (1, 1)]

    if sysname == "alanine":
        testsystem = testsystems.AlanineDipeptideImplicit()
        system, positions = testsystem.system, testsystem.positions
        groups = [(0, 2), (1, 1)]
        timestep = 2.0 * u.femtoseconds
        hmc_integrators.guess_force_groups(system, nonbonded=1, others=0)
        #remove_cmm(system)  # Unrolled shouldn't need this
        equil_steps = 10000

    if sysname == "alanineexplicit":
        testsystem = testsystems.AlanineDipeptideExplicit(cutoff=1.1*u.nanometers, switch_width=2*u.angstrom, ewaldErrorTolerance=5E-5)
        system, positions = testsystem.system, testsystem.positions
        #groups = [(0, 2), (1, 1), (2, 1)]
        #groups = [(0, 2), (1, 1)]
        timestep = 1.0 * u.femtoseconds
        #hmc_integrators.guess_force_groups(system, nonbonded=1, others=0, fft=2)
        #hmc_integrators.guess_force_groups(system, nonbonded=0, others=0, fft=1)
        groups = [(0, 1), (1, 2)]
        hmc_integrators.guess_force_groups(system, nonbonded=0, others=1, fft=0)

        #remove_cmm(system)  # Unrolled doesn't need this
        equil_steps = 4000


    # guess force groups

    if "ljbox" in sysname:
        timestep = 25 * u.femtoseconds
        temperature = 25. * u.kelvin
        system, positions = testsystem.system, testsystem.positions
        hmc_integrators.guess_force_groups(system, nonbonded=0, fft=1)
        groups = [(0, 2), (1, 1)]
        equil_steps = 10000
        steps_per_hmc = 15

    if sysname == "amoeba":
        testsystem = testsystems.AMOEBAIonBox()
        system, positions = testsystem.system, testsystem.positions
        hmc_integrators.guess_force_groups(system, nonbonded=0, fft=1, others=0)


    return system, positions, groups, temperature, timestep, langevin_timestep, testsystem, equil_steps, steps_per_hmc
import lb_loader
import simtk.openmm.app as app
import numpy as np
import pandas as pd
import simtk.openmm as mm
from simtk import unit as u
from openmmtools import hmc_integrators, testsystems
pd.set_option('display.width', 1000)

n_steps = 3000
temperature = 300. * u.kelvin

hydrogenMass = 3.0 * u.amu

system, positions = lb_loader.load_amoeba()
hmc_integrators.guess_force_groups(system, multipole=2)

#system, positions = lb_loader.load_lb(hydrogenMass=hydrogenMass)
#hmc_integrators.guess_force_groups(system, nonbonded=1, fft=1, others=0)

integrator = mm.LangevinIntegrator(temperature, 1.0 / u.picoseconds,
                                   0.25 * u.femtoseconds)
context = mm.Context(system, integrator)
context.setPositions(positions)
context.setVelocitiesToTemperature(temperature)

integrator.step(1000)
positions = context.getState(getPositions=True).getPositions()
state = context.getState(getEnergy=True)
energy = state.getPotentialEnergy() + state.getKineticEnergy()
energy, state.getPotentialEnergy(), state.getKineticEnergy()
int_args = dict(
    temperature=temperature,
    timestep=1.0,
    steps_per_hmc=1,
    collision_rate=1.0 / u.picoseconds,
    extra_chances=0,
    group0_iterations=2,
    groups=(0, 1, 1),
)

collision_rate = int_args["collision_rate"]

groups = int_args["groups"]
hmc_integrators.guess_force_groups(system,
                                   others=groups[0],
                                   nonbonded=groups[1],
                                   fft=groups[2])

steps = 1000

integrator = lb_loader.kw_to_int(**int_args)
#integrator = integrators.HMCIntegrator(temperature=temperature, nsteps=steps_per_hmc, timestep=timestep)
#integrator = hmc_integrators.GHMCIntegrator(temperature=temperature, steps_per_hmc=steps_per_hmc, timestep=timestep, collision_rate=collision_rate)

#temperature = int_args["temperature"]
#steps_per_hmc = int_args["steps_per_hmc"]
#timestep = int_args["timestep"]
groups = [(0, 1)]
extra_chances = 1
#collision_rate = None
integrator = hmc_integrators.GHMCIntegrator(temperature=temperature,