def test_explicit_runner_scaler(self):
p = subsystem.SubSystemFromSequence("NALA ALA CALA")
b = builder.SystemBuilder(explicit_solvent=True)
sys = b.build_system([p])
sys.temperature_scaler = temperature.ConstantTemperatureScaler(300.0 * u.kelvin)
rest2_scaler = temperature.GeometricTemperatureScaler(
0, 1, 300.0 * u.kelvin, 350.0 * u.kelvin
)
opt = options.RunOptions(solvation="explicit")
opt.rest2_scaler = temperature.REST2Scaler(300.0 * u.kelvin, rest2_scaler)
opt.minimize_steps = 100
opt.timesteps = 2
opt.use_rest2 = True
runner = openmm_runner.OpenMMRunner(sys, opt, platform="Reference")
runner.prepare_for_timestep(sys.get_state_template(), 0.0, 1)
pos = sys._coordinates.copy()
vel = np.zeros_like(pos)
alpha = 0.0
energy = 0.0
box_vectors = sys._box_vectors
s = state.SystemState(pos, vel, alpha, energy, box_vectors)
s = runner.minimize_then_run(s)
s = runner.run(s)
assert s
def gen_state(index):
pos = index * np.ones((N_ATOMS, 3))
vel = index * np.ones((N_ATOMS, 3))
alpha = 0.0
energy = 0.0
box_vectors = np.zeros(3)
return state.SystemState(pos, vel, alpha, energy, box_vectors)
def test_implicit_runner_amap(self):
p = subsystem.SubSystemFromSequence("NALA ALA CALA")
b = builder.SystemBuilder()
sys = b.build_system([p])
sys.temperature_scaler = temperature.ConstantTemperatureScaler(300.0 * u.kelvin)
opt = options.RunOptions()
opt.timesteps = 20
opt.use_amap = True
opt.amap_beta_bias = 10
runner = openmm_runner.OpenMMRunner(sys, opt, platform="Reference")
runner.prepare_for_timestep(sys.get_state_template(), 0.0, 1)
pos = sys._coordinates.copy()
vel = np.zeros_like(pos)
alpha = 0.0
energy = 0.0
box_vectors = np.zeros(3)
s = state.SystemState(pos, vel, alpha, energy, box_vectors)
s = runner.minimize_then_run(s)
s = runner.run(s)
assert s
def setUp(self):
self.setUpTempDir()
# setup data store
self.N_ATOMS = 500
self.N_REPLICAS = 16
self.N_DISCRETE = 10
self.N_CONTINUOUS = 5
self.N_MAPPINGS = 10
self.state_template = state.SystemState(
np.zeros((self.N_ATOMS, 3)),
np.zeros((self.N_ATOMS, 3)),
0.0,
0.0,
np.zeros(3),
param_sampling.ParameterState(
np.zeros(self.N_DISCRETE, dtype=np.int32),
np.zeros(self.N_CONTINUOUS, dtype=np.float64),
),
np.arange(self.N_MAPPINGS),
)
# dummy pdb writer; can't use a mock because they can't be pickled
pdb_writer = object()
self.store = vault.DataStore(self.state_template,
self.N_REPLICAS,
pdb_writer,
block_size=10)
self.store.initialize(mode="w")
def load_states(self, stage: int) -> Sequence[interfaces.IState]:
"""
Load states from disk
Args:
stage: stage to load
Returns:
list of states
"""
self._handle_load_stage(stage)
positions = self.load_positions(stage)
velocities = self.load_velocities(stage)
box_vectors = self.load_box_vectors(stage)
alphas = self.load_alphas(stage)
energies = self.load_energies(stage)
discrete_parameters = self.load_discrete_parameters(stage)
continuous_parameters = self.load_continuous_parameters(stage)
mappings = self.load_mappings(stage)
states = []
for i in range(self._n_replicas):
s = state.SystemState(
positions[i],
velocities[i],
alphas[i],
energies[i],
box_vectors[i],
param_sampling.ParameterState(discrete_parameters[i],
continuous_parameters[i]),
mappings[i],
)
states.append(s)
return states
def setUp(self):
self.setUpTempDir()
self.N_ATOMS = 500
self.N_REPLICAS = 16
self.N_DISCRETE = 10
self.N_CONTINUOUS = 5
self.N_MAPPINGS = 10
self.state_template = state.SystemState(
np.zeros((self.N_ATOMS, 3)),
np.zeros((self.N_ATOMS, 3)),
0.0,
0.0,
np.zeros(3),
param_sampling.ParameterState(
np.zeros(self.N_DISCRETE, dtype=np.int32),
np.zeros(self.N_CONTINUOUS, dtype=np.float64),
),
np.arange(self.N_MAPPINGS),
)
# setup objects to save to disk
c = comm.MPICommunicator(self.N_ATOMS, self.N_REPLICAS)
l = ladder.NearestNeighborLadder(n_trials=100)
policy = adaptor.AdaptationPolicy(1.0, 50, 100)
a = adaptor.EqualAcceptanceAdaptor(n_replicas=self.N_REPLICAS,
adaptation_policy=policy)
# make some states
def gen_state(index, n_atoms):
pos = index * np.ones((n_atoms, 3))
vel = index * np.ones((n_atoms, 3))
energy = index
lam = index / 100.0
discrete = np.zeros(self.N_DISCRETE, dtype=np.int32)
continuous = np.zeros(self.N_CONTINUOUS, dtype=np.float64)
params = param_sampling.ParameterState(discrete, continuous)
mappings = np.arange(self.N_MAPPINGS)
return state.SystemState(pos, vel, lam, energy, np.zeros(3),
params, mappings)
states = [gen_state(i, self.N_ATOMS) for i in range(self.N_REPLICAS)]
runner = leader.LeaderReplicaExchangeRunner(self.N_REPLICAS,
max_steps=100,
ladder=l,
adaptor=a)
# dummy pdb writer; can't use a mock because they can't be pickled
pdb_writer = object()
self.store = vault.DataStore(self.state_template, self.N_REPLICAS,
pdb_writer)
self.store.initialize(mode="w")
# save some stuff
self.store.save_data_store()
self.store.save_communicator(c)
self.store.save_remd_runner(runner)
self.store.save_states(states, stage=0)
def gen_state(index, n_atoms):
pos = index * np.ones((n_atoms, 3))
vel = index * np.ones((n_atoms, 3))
energy = index
lam = index / 100.0
discrete = np.zeros(self.N_DISCRETE, dtype=np.int32)
continuous = np.zeros(self.N_CONTINUOUS, dtype=np.float64)
params = param_sampling.ParameterState(discrete, continuous)
return state.SystemState(pos, vel, lam, energy, np.zeros(3), params)
def gen_state(index):
pos = index * np.ones((N_ATOMS, 3))
vel = index * np.ones((N_ATOMS, 3))
alpha = 0.0
energy = 0.0
box_vectors = np.zeros(3)
discrete = np.zeros(N_DISCRETE, dtype=np.int32)
continuous = np.zeros(N_CONTINUOUS, dtype=np.float64)
params = param_sampling.ParameterState(discrete, continuous)
return state.SystemState(pos, vel, alpha, energy, box_vectors, params)
def get_state_template(self):
pos = self._coordinates.copy()
vel = np.zeros_like(pos)
alpha = 0.0
energy = 0.0
box_vectors = self._box_vectors
if box_vectors is None:
box_vectors = np.array([0.0, 0.0, 0.0])
params = self.param_sampler.get_initial_state()
return state.SystemState(pos, vel, alpha, energy, box_vectors, params)
def setUp(self):
self.N_ATOMS = 500
self.N_REPLICAS = 4
self.N_DISCRETE = 10
self.N_CONTINUOUS = 5
self.state_template = state.SystemState(
np.zeros((self.N_ATOMS, 3)),
np.zeros((self.N_ATOMS, 3)),
0.0,
0.0,
np.zeros(3),
param_sampling.ParameterState(
np.zeros(self.N_DISCRETE, dtype=np.int32),
np.zeros(self.N_CONTINUOUS, dtype=np.float64),
),
)
def load_states(self, stage):
"""
Load states from disk
:param stage: integer stage to load
:return: list of SystemState objects
"""
self._handle_load_stage(stage)
positions = self.load_positions(stage)
velocities = self.load_velocities(stage)
box_vectors = self.load_box_vectors(stage)
alphas = self.load_alphas(stage)
energies = self.load_energies(stage)
states = []
for i in range(self._n_replicas):
s = state.SystemState(positions[i], velocities[i], alphas[i],
energies[i], box_vectors[i])
states.append(s)
return states
def gen_state(index):
pos = index * np.ones((N_ATOMS, 3))
vel = index * np.ones((N_ATOMS, 3))
alpha = 0
energy = 0
return state.SystemState(pos, vel, alpha, energy)
def test_lambda_must_be_between_zero_and_one(self):
"should raise RuntimeError if lambda is outside of [0,1]"
bad_lam = -2
with self.assertRaises(RuntimeError):
state.SystemState(self.coords, self.vels, bad_lam, self.energy,
self.box_vectors)
def test_should_raise_if_vels_not_match_coords(self):
"should raise runtime error if vels is not the same shape as coords"
bad_vels = np.zeros((42, 3))
with self.assertRaises(RuntimeError):
state.SystemState(self.coords, bad_vels, self.lam, self.energy,
self.box_vectors)
def test_should_raise_with_coords_second_dim_not_3(self):
"should raise RuntimeError if coords is not (n, 3)"
bad_pos = np.zeros((75, 4))
with self.assertRaises(RuntimeError):
state.SystemState(bad_pos, self.vels, self.lam, self.energy,
self.box_vectors)
def test_should_raise_with_coords_not_2d(self):
"should raise RuntimeError if coords is not 2d"
bad_pos = np.zeros((75, 83, 52))
with self.assertRaises(RuntimeError):
state.SystemState(bad_pos, self.vels, self.lam, self.energy,
self.box_vectors)
