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_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)
dims = 3
coords = np.random.randn(frames, beads, dims).astype('float32')
forces = np.random.randn(frames, beads, dims).astype('float32')
dataset = MoleculeDataset(coords, forces)
# Here we construct a single hidden layer architecture with random
# widths and a terminal contraction to a scalar output
arch = (LinearLayer(dims, dims, activation=nn.Tanh()) +
LinearLayer(dims, 1, activation=None))
# Here we construct a CGnet model using the above architecture
# as well as variables to be used in CG simulation tests
model = CGnet(arch, ForceLoss()).float()
model.eval()
sim_length = np.random.choice([2, 4]) * 2 # Number of frames to simulate
# Frequency with which to save simulation
save_interval = np.random.choice([2, 4])
# frames (choice of 2 or 4)
# Grab intitial coordinates as a simulation starting configuration
# from the moleular dataset
initial_coordinates = dataset[:][0].reshape(-1, beads, dims)
# Langevin simulation parameters
masses = np.ones(beads)
friction = np.random.randint(10, 20)
# SchNet model
feature_size = np.random.randint(5, 10) # random feature size