def test_brownian_simulation_saved_potential_shape():
# Test shape of brownian (overdamped langevin) simulation with both
# forces and potential saved
# Generate simulation
my_sim = Simulation(model,
initial_coordinates,
length=sim_length,
save_interval=save_interval,
save_potential=True)
traj = my_sim.simulate()
# Here, we verify that the trajectory shape corresponds with the
# choices of simulation length and saving frequency above
# We also verify that the forces and the potential are saved
assert traj.shape == (frames, sim_length // save_interval, beads, dims)
assert my_sim.simulated_forces is None
assert my_sim.simulated_potential.shape == (frames,
sim_length // save_interval,
beads, 1)
assert my_sim.kinetic_energies is None
def test_langevin_simulation_saved_forces_shape():
# Test shape of langevin simulation with only forces saved
# Generate simulation
my_sim = Simulation(model,
initial_coordinates,
length=sim_length,
save_interval=save_interval,
save_forces=True,
friction=friction,
masses=masses)
traj = my_sim.simulate()
# Here, we verify that the trajectory shape corresponds with the
# choices of simulation length and saving frequency above
# We also verify that the forces, but not the potential, is saved
assert traj.shape == (frames, sim_length // save_interval, beads, dims)
assert my_sim.simulated_forces.shape == (frames,
sim_length // save_interval,
beads, dims)
assert my_sim.simulated_potential is None
assert my_sim.kinetic_energies.shape == (frames,
sim_length // save_interval)
def test_cgnet_simulation():
# Tests a simulation from a CGnet built with the GeometryFeature
# for the shapes of its coordinate, force, and potential outputs
# First, we set up a bond harmonic prior and a GeometryFeature layer
bonds_idx = geom_stats.return_indices('Bonds')
bonds_interactions, _ = geom_stats.get_prior_statistics(features='Bonds',
as_list=True)
harmonic_potential = HarmonicLayer(bonds_idx, bonds_interactions)
feature_layer = GeometryFeature(feature_tuples='all_backbone',
n_beads=beads)
num_feats = feature_layer(coords).size()[1]
# Next, we create a 4 layer hidden architecture with a random width
# and with a scalar output
rand = np.random.randint(1, 10)
arch = (LinearLayer(num_feats, rand, bias=True, activation=nn.Tanh()) +
LinearLayer(rand, rand, bias=True, activation=nn.Tanh()) +
LinearLayer(rand, rand, bias=True, activation=nn.Tanh()) +
LinearLayer(rand, rand, bias=True, activation=nn.Tanh()) +
LinearLayer(rand, 1, bias=True, activation=None))
# Next, we instance a CGnet model using the above objects
# with force matching as a loss criterion
model = CGnet(arch,
ForceLoss(),
feature=feature_layer,
priors=[harmonic_potential])
model.eval()
# Here, we produce mock target protein force data
forces = torch.randn((frames, beads, 3), requires_grad=False)
# Here, we create an optimizer for traning the model,
# and we train it for one epoch
optimizer = torch.optim.Adam(model.parameters(), lr=0.05, weight_decay=0)
optimizer.zero_grad()
energy, pred_forces = model.forward(coords)
loss = model.criterion(pred_forces, forces)
loss.backward()
optimizer.step()
# Here, we define random simulation frame lengths
# as well as randomly choosing to save every 2 or 4 frames
length = np.random.choice([2, 4]) * 2
save = np.random.choice([2, 4])
# Here we instance a simulation class and produce a CG trajectory
my_sim = Simulation(model,
coords,
beta=geom_stats.beta,
length=length,
save_interval=save,
save_forces=True,
save_potential=True)
traj = my_sim.simulate()
# We test to see if the trajectory is the proper shape based on the above
# choices for simulation length and frame saving
assert traj.shape == (frames, length // save, beads, dims)
assert my_sim.simulated_forces.shape == (frames, length // save, beads,
dims)
assert my_sim.simulated_potential.shape == (frames, length // save, 1)
def test_single_model_simulation_vs_multimodelsimulation():
# Tests to make sure that Simulation and MultiModelSimulation return
# the same simulation results (coordinates, forces, potential energy,
# and kinetic energies) if a single model is used in both cases.
# First, we generate a random integer seed
seed = np.random.randint(0, 1e6)
# Next, we set up a model
save_interval = 1
dt = 0.001 * np.random.randint(1, 11)
friction = 10 * np.random.randint(1, 11)
k = np.random.randint(1, 6)
n_particles = np.random.randint(1, 101)
n_sims = np.random.randint(1, 11)
initial_coordinates = torch.randn((n_sims, n_particles, 3))
masses = n_particles * [np.random.randint(low=1, high=5)]
sim_length = np.random.randint(2, 11)
model = HarmonicPotential(k=k,
T=300,
n_particles=n_particles,
dt=dt,
friction=friction,
n_sims=n_sims,
sim_length=sim_length,
save_interval=save_interval)
# Next, we simulate both models. We wrap both of the simulations in
# a temporary directory as to not generate permanent simulation files
with tempfile.TemporaryDirectory() as tmp:
sim = Simulation(model,
initial_coordinates,
embeddings=None,
length=sim_length,
save_interval=save_interval,
masses=masses,
dt=dt,
save_forces=True,
save_potential=True,
friction=friction,
random_seed=seed,
filename=tmp + '/test')
multi_sim = MultiModelSimulation([model],
initial_coordinates,
embeddings=None,
length=sim_length,
save_interval=save_interval,
masses=masses,
dt=dt,
save_forces=True,
save_potential=True,
friction=friction,
random_seed=seed,
filename=tmp + '/test_copy')
trajectory = sim.simulate()
trajectory_from_multi = multi_sim.simulate()
# Here, we test the equality of the two simulation results
assert trajectory.shape == trajectory_from_multi.shape
assert sim.simulated_forces.shape == multi_sim.simulated_forces.shape
assert sim.simulated_potential.shape == multi_sim.simulated_potential.shape
assert sim.kinetic_energies.shape == multi_sim.kinetic_energies.shape
np.testing.assert_array_equal(trajectory, trajectory_from_multi)
np.testing.assert_array_equal(sim.simulated_potential,
multi_sim.simulated_potential)
np.testing.assert_array_equal(sim.simulated_forces,
multi_sim.simulated_forces)
np.testing.assert_array_equal(sim.kinetic_energies,
multi_sim.kinetic_energies)
def test_saving_all_quantities():
# Tests, using a temporary directory, the saving of coordinates,
# forces, potential, and kinetic energies from a Langevin simulation
# (i) That the number of numpy files saved is correct
# (ii) That the saved numpy files have the proper shapes
# (iii) That the contatenation of the saved numpy files are equal to the
# trajectory output from the simulation
n_sims = np.random.randint(1, high=5)
sim_length = np.random.choice([24, 36])
npy_interval = np.random.choice([6, 12])
save_interval = np.random.choice([2, 3])
n_expected_files = sim_length / npy_interval
model = HarmonicPotential(k=1,
T=300,
n_particles=10,
dt=0.001,
friction=10,
n_sims=n_sims,
sim_length=sim_length,
save_interval=save_interval)
initial_coordinates = torch.zeros((model.n_sims, model.n_particles, 3))
with tempfile.TemporaryDirectory() as tmp:
my_sim = Simulation(model,
initial_coordinates,
embeddings=None,
beta=model.beta,
length=model.sim_length,
friction=model.friction,
dt=model.dt,
save_forces=True,
save_potential=True,
masses=model.masses,
save_interval=model.save_interval,
export_interval=npy_interval,
filename=tmp + '/test')
traj = my_sim.simulate()
assert traj.shape[1] == sim_length / save_interval
file_list = os.listdir(tmp)
assert len(
file_list) == n_expected_files * 4 # coords, forces, pot, ke
coords_file_list = sorted(
[file for file in file_list if 'coords' in file])
force_file_list = sorted(
[file for file in file_list if 'forces' in file])
potential_file_list = sorted(
[file for file in file_list if 'potential' in file])
ke_file_list = sorted(
[file for file in file_list if 'kineticenergy' in file])
file_list_list = [
coords_file_list, force_file_list, potential_file_list,
ke_file_list
]
expected_chunk_length = npy_interval / save_interval
# needed for (iii)
running_coords = None
running_forces = None
running_potential = None
running_ke = None
running_list = [
running_coords, running_forces, running_potential, running_ke
]
obs_list = [
my_sim.simulated_coords, my_sim.simulated_forces,
my_sim.simulated_potential, my_sim.kinetic_energies
]
for j, obs_file_list in enumerate(file_list_list):
for i in range(len(obs_file_list)):
temp_traj = np.load(tmp + '/' + obs_file_list[i])
# Test (ii)
if j < 3:
np.testing.assert_array_equal(
temp_traj.shape,
[n_sims, expected_chunk_length, model.n_particles, 3])
else:
np.testing.assert_array_equal(
temp_traj.shape, [n_sims, expected_chunk_length])
if running_list[j] is None:
running_list[j] = temp_traj
else:
running_list[j] = np.concatenate(
[running_list[j], temp_traj], axis=1)
# Test (iii)
np.testing.assert_array_equal(obs_list[j], running_list[j])
def test_saving_numpy_coordinates():
# Tests, using a temporary directory, the saving of *coordinates*
# from a Brownian (overdamped Langevin) simulation
# (i) That the number of numpy files saved is correct
# (ii) That the saved numpy files have the proper shapes
# (iii) That the contatenation of the saved numpy files are equal to the
# trajectory output from the simulation
n_sims = np.random.randint(1, high=5)
sim_length = np.random.choice([24, 36])
npy_interval = np.random.choice([6, 12])
save_interval = np.random.choice([2, 3])
n_expected_files = sim_length / npy_interval
model = HarmonicPotential(k=1,
T=300,
n_particles=10,
dt=0.001,
friction=None,
n_sims=n_sims,
sim_length=sim_length,
save_interval=save_interval)
initial_coordinates = torch.zeros((model.n_sims, model.n_particles, 3))
with tempfile.TemporaryDirectory() as tmp:
my_sim = Simulation(model,
initial_coordinates,
embeddings=None,
beta=model.beta,
length=model.sim_length,
friction=model.friction,
dt=model.dt,
save_forces=False,
save_potential=False,
save_interval=model.save_interval,
export_interval=npy_interval,
filename=tmp + '/test')
traj = my_sim.simulate()
assert traj.shape[1] == sim_length / save_interval
file_list = os.listdir(tmp)
assert len(file_list) == n_expected_files
expected_chunk_length = npy_interval / save_interval
running_traj = None # needed for (iii)
for i in range(len(file_list)):
temp_traj = np.load(tmp + '/' + file_list[i])
# Test (ii)
np.testing.assert_array_equal(
temp_traj.shape,
[n_sims, expected_chunk_length, model.n_particles, 3])
if running_traj is None:
running_traj = temp_traj
else:
running_traj = np.concatenate([running_traj, temp_traj],
axis=1)
# Test (iii)
np.testing.assert_array_equal(traj, running_traj)