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() # 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 __init__( self, positions: np.ndarray, velocities: np.ndarray, alpha: float, energy: float, box_vector: np.ndarray, parameters: Optional[param_sampling.ParameterState] = None, ) -> None: """ Initialize a SystemState Params: positions: coordinates of structure, shape(n_atoms, 3) velocities: velocities for structure, shape(n_atoms, 3) alpha: alpha value, within ``[0, 1]`` energy: total potential energy, including restraints, in kJ/mol box_vector: the box vectors, shape(3, 3) in nm """ self.positions = positions self.velocities = velocities self.box_vector = box_vector self.n_atoms = positions.shape[0] self.alpha = alpha self.energy = energy if parameters is None: self.parameters = param_sampling.ParameterState( discrete=np.array([], dtype=np.int32), continuous=np.array([], dtype=np.float64), ) else: self.parameters = parameters self._validate()
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 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), ), )