def test_langevin_simulation_safety():
# Test whether a brownian (overdamped langevin) simulation indeed will
# refuse to overwrite existing data unless overwrite is set to true
# Generate simulation
sim = Simulation(model,
initial_coordinates,
length=sim_length,
save_interval=save_interval,
friction=friction,
masses=masses)
# Check that no simulation is stored
assert not sim._simulated
traj = sim.simulate()
assert sim.kinetic_energies is not None
# Check that a simulation is stored now
assert sim._simulated
# Check that it can't be overwritten by default
np.testing.assert_raises(RuntimeError, sim.simulate)
# Check that it can be overwritten with overwrite=True; i.e. that
# this command raises no error
traj2 = sim.simulate(overwrite=True)
assert sim._simulated
def test_harmonic_potential_zero_friction():
# Test that zero friction returns a traj of zeroes and kinetic energy
# of zeroes. Zero friction means vscale will be zero, which means
# that velocities will only ever be updated by zero. We therefore
# expect zero friction to leave our velocities completely unchanged.
# Particularly, in the case where the starting velocities are zero,
# we expect them to remain zero, and thus the positions do not
# change either.
# set up model, internal coords, and sim using class attirbutes
model = HarmonicPotential(k=1,
T=300,
n_particles=1000,
dt=0.001,
friction=0,
n_sims=1,
sim_length=500)
initial_coordinates = torch.zeros((model.n_sims, model.n_particles, 3))
my_sim = Simulation(model,
initial_coordinates,
embeddings=None,
beta=model.beta,
length=model.sim_length,
friction=model.friction,
dt=model.dt,
masses=model.masses,
save_interval=model.save_interval)
traj = my_sim.simulate()
np.testing.assert_array_equal(traj, np.zeros(traj.shape))
np.testing.assert_array_equal(my_sim.kinetic_energies,
np.zeros(my_sim.kinetic_energies.shape))
def test_cgschnet_simulation_shapes():
# Test simulation with embeddings and make sure the shapes of
# the simulated coordinates, forces, and potential are correct
schnet_feature, embedding_property, feature_size = _get_random_schnet_feature(
calculate_geometry=True)
layer_list = [schnet_feature]
feature_combiner = FeatureCombiner(layer_list)
# Next, we make aa CGnet with a random hidden architecture
arch = _get_random_architecture(feature_size)
model = CGnet(arch, ForceLoss(), feature=feature_combiner)
model.eval()
sim_length = np.random.randint(10, 20)
sim = Simulation(model, coords_torch, embedding_property, length=sim_length,
save_interval=1, beta=1., save_forces=True,
save_potential=True)
traj = sim.simulate()
np.testing.assert_array_equal(sim.simulated_coords.shape,
[n_frames, sim_length, n_beads, 3])
np.testing.assert_array_equal(sim.simulated_forces.shape,
[n_frames, sim_length, n_beads, 3])
np.testing.assert_array_equal(sim.simulated_potential.shape,
[n_frames, sim_length, 1])
def test_harmonic_potential_several_temperatures():
# Tests several harmonic potential simulations for correct temperature.
# The standard deviation in measured temperature across the simulation
# is expected to increase as the temperature increases. Heursitically I
# observed it doesn't tend to exceed a standard deviation of 30 for
# simulation lengths of 500 and max temperatures of 900.
temp_parameter = [100, 300, 500, 700, 900]
mean_temp_measured = []
std_temp_measured = []
for temp in temp_parameter:
# set up model, internal coords, and sim using class attirbutes
model = HarmonicPotential(k=1,
T=temp,
n_particles=1000,
dt=0.001,
friction=100,
n_sims=1,
sim_length=500)
initial_coordinates = torch.zeros((model.n_sims, model.n_particles, 3))
my_sim = Simulation(model,
initial_coordinates,
embeddings=None,
beta=model.beta,
length=model.sim_length,
friction=model.friction,
dt=model.dt,
masses=model.masses,
save_interval=model.save_interval)
traj = my_sim.simulate()
n_dofs = 3 * model.n_particles
sim_temps = my_sim.kinetic_energies * 2 / n_dofs / model.kB
sim_temps = sim_temps[:, 20:]
# store mean
sim_temp_mean = np.mean(sim_temps, axis=1)[0] # only one simulation
mean_temp_measured.append(sim_temp_mean)
# store stdev
sim_temp_std = np.std(sim_temps, axis=1)[0] # only one simulation
std_temp_measured.append(sim_temp_std)
# Test that the means are all about the right temperature
np.testing.assert_allclose(temp_parameter, mean_temp_measured, rtol=1)
# Test that the stdevs are all less than 25 (heuristic)
np.testing.assert_array_less(std_temp_measured,
np.repeat(30, len(temp_parameter)))
# Test that the stdevs go up as the temperature goes up
np.testing.assert_array_equal(std_temp_measured, sorted(std_temp_measured))
def test_log_file_basics():
# Tests whether the log file exists, is named correctly, and has the
# correct number of lines
n_sims = np.random.randint(1, high=5)
sim_length = np.random.choice([24, 36])
log_interval = np.random.choice([6, 12])
save_interval = np.random.choice([2, 3])
n_expected_logs = sim_length / log_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,
log_interval=log_interval,
log_type='write',
filename=tmp + '/test')
traj = my_sim.simulate()
assert traj.shape[1] == sim_length / save_interval
file_list = os.listdir(tmp)
# Check that one file exists in the temp directory
assert len(file_list) == 1
# Check that it has the proper name
assert file_list[0] == 'test_log.txt'
# Gather its lines
with open(tmp + '/' + file_list[0]) as f:
line_list = f.readlines()
# We expect the log file to contain the expected number of logs, plus two
# extra lines: one at the start and one at the end.
assert len(line_list) == n_expected_logs + 2
def test_harmonic_potential_shape_and_temperature():
# Tests a single harmonic potential simulation for shape and temperature
# - Tests shapes of trajectory and kinetic energies
# - Tests that average temperature is about 300
# - Tests that the standard deviation is less than 10; this is just
# heuristic based on trying out the data in a notebook when
# writing the test
# set up model, internal coords, and sim using class attirbutes
model = HarmonicPotential(k=1,
T=300,
n_particles=1000,
dt=0.001,
friction=100,
n_sims=1,
sim_length=500)
initial_coordinates = torch.zeros((model.n_sims, model.n_particles, 3))
my_sim = Simulation(model,
initial_coordinates,
embeddings=None,
beta=model.beta,
length=model.sim_length,
friction=model.friction,
dt=model.dt,
masses=model.masses,
save_interval=model.save_interval)
traj = my_sim.simulate()
# check shape of trajectory and kinetic energies
assert traj.shape == (model.n_sims, model.sim_length, model.n_particles, 3)
assert my_sim.kinetic_energies.shape == (1, model.sim_length)
# Calculate temperatures, removing the first 20 time points (this is
# about what it takes to get to a constant temperature)
n_dofs = 3 * model.n_particles
temperatures = my_sim.kinetic_energies * 2 / n_dofs / model.kB
temperatures = temperatures[:, 20:]
mean_temps = np.mean(temperatures, axis=1)
# Test that the means are all about the right temperature
np.testing.assert_allclose(np.mean(temperatures, axis=1),
np.repeat(model.T, model.n_sims),
rtol=1)
# Test that the stdevs are all less than 25 (heuristic)
np.testing.assert_array_less(np.std(temperatures, axis=1),
np.repeat(10, model.n_sims))
def test_schnet_simulation_safety():
# This test confirms that a simulation cannot be instantiated if
# schnet features but no embeddings are provided
with assert_raises(RuntimeError):
Simulation(schnet_model,
initial_coordinates,
length=sim_length,
save_interval=save_interval)
def test_brownian_simulation_shape():
# Test shape of Brownian (overdamped Langevin) simulation without
# saving the forces or the potential
# Generate simulation
my_sim = Simulation(model,
initial_coordinates,
length=sim_length,
save_interval=save_interval)
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 potential and the forces are not saved
assert traj.shape == (frames, sim_length // save_interval, beads, dims)
assert my_sim.simulated_forces is None
assert my_sim.simulated_potential is None
assert my_sim.kinetic_energies is None
def test_brownian_simulation_seeding():
# Test determinism of Brownian (overdamped langevin) simulation with
# random seed. If the same seed is used for two separate simulations,
# the results (trajectory, forces, potential) should be identical
seed = np.random.randint(1000) # Save random seed for simulations
# Generate simulation number one
sim1 = Simulation(model,
initial_coordinates,
length=sim_length,
save_interval=save_interval,
save_forces=True,
save_potential=True,
random_seed=seed)
traj1 = sim1.simulate()
# Generate simulation umber two
sim2 = Simulation(model,
initial_coordinates,
length=sim_length,
save_interval=save_interval,
save_forces=True,
save_potential=True,
random_seed=seed)
traj2 = sim2.simulate()
# Verify that all components of each simulation are equal.
np.testing.assert_array_equal(traj1, traj2)
np.testing.assert_array_equal(sim1.simulated_forces, sim2.simulated_forces)
np.testing.assert_array_equal(sim1.simulated_potential,
sim2.simulated_potential)
assert sim1.kinetic_energies is None
assert sim2.kinetic_energies is None
def test_langevin_simulation_saved_potential_shape():
# Test shape of 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,
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 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.shape == (frames,
sim_length // save_interval)
def test_langevin_simulation_seeding():
# Test determinism of Langevin simulation with random seed. If the
# same seed is used for two separate simulations, the results
# (trajectory, forces, potential) should be identical
seed = np.random.randint(1000) # Save random seed for simulations
# Generate simulation number one
sim1 = Simulation(model,
initial_coordinates,
length=sim_length,
save_interval=save_interval,
save_forces=True,
save_potential=True,
random_seed=seed,
friction=friction,
masses=masses)
traj1 = sim1.simulate()
# Generate simulation umber two
sim2 = Simulation(model,
initial_coordinates,
length=sim_length,
save_interval=save_interval,
save_forces=True,
save_potential=True,
random_seed=seed,
friction=friction,
masses=masses)
traj2 = sim2.simulate()
# Verify that all components of each simulation are equal, and not
# because they're None
assert traj1 is not None
assert traj2 is not None
assert sim1.simulated_forces is not None
assert sim2.simulated_forces is not None
assert sim1.simulated_potential is not None
assert sim2.simulated_potential is not None
assert sim1.kinetic_energies is not None
assert sim2.kinetic_energies is not None
np.testing.assert_array_equal(traj1, traj2)
np.testing.assert_array_equal(sim1.simulated_forces, sim2.simulated_forces)
np.testing.assert_array_equal(sim1.simulated_potential,
sim2.simulated_potential)
np.testing.assert_array_equal(sim1.kinetic_energies, sim2.kinetic_energies)
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)