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)