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),
),
)