def test_schnet_feature_geometry():
# Tests SchnetFeature's calls to the Geometry class for
# distance calculations
# First, we instance a SchnetFeature that can call to Geometry
rbf_layer = GaussianRBF()
schnet_feature = SchnetFeature(feature_size=n_feats,
embedding_layer=None,
rbf_layer=rbf_layer,
n_interaction_blocks=2,
calculate_geometry=True,
n_beads=beads)
# Next we instance a geom_stats that only calculates distances
# and compare distance pair tuples
geom_stats = GeometryStatistics(coords,
backbone_inds='all',
get_all_distances=True,
get_backbone_angles=False,
get_backbone_dihedrals=False,
get_redundant_distance_mapping=True)
# Here, we make sure that schnet_feature's calls to Geometry
# can replicate those of a GeometryStatistics instance
schnet_distance_pairs, _ = schnet_feature.geometry.get_distance_indices(
beads, [], [])
schnet_red_dist_map = schnet_feature.geometry.get_redundant_distance_mapping(
schnet_distance_pairs)
assert schnet_distance_pairs == geom_stats._distance_pairs
np.testing.assert_equal(schnet_red_dist_map,
geom_stats.redundant_distance_mapping)
def test_bead_energy_masking():
# Tests to make sure that masked energies and forces are properly zeroed
# by the bead mask used with variable sized input
# We create a simple random embedding layer and some
# mock, padded embeddings that originally have varying length
num_feats = np.random.randint(10, 50)
n_embeddings = np.random.randint(beads, 2 * beads)
embedding_layer = CGBeadEmbedding(n_embeddings=n_embeddings,
embedding_dim=num_feats)
variable_beads = np.random.randint(3, beads,
size=frames) # random protein sizes
variable_embeddings = [
np.random.randint(1, high=beads, size=bead) for bead in variable_beads
]
padded_embedding_list = []
for embedding in variable_embeddings:
pads_needed = beads - embedding.shape[0]
padded_embeddings = np.hstack((embedding, np.zeros(pads_needed)))
padded_embedding_list.append(padded_embeddings)
embedding_property = torch.tensor(padded_embedding_list).long()
# we create a simple 2 layer random width terminal network
rand = np.random.randint(1, 10)
arch = (LinearLayer(num_feats, rand, bias=True, activation=nn.Tanh()) +
LinearLayer(rand, 1, bias=True, activation=nn.Tanh()))
# Next we create a basic SchnetFeature
rbf_layer = GaussianRBF()
feature = SchnetFeature(num_feats,
embedding_layer=embedding_layer,
rbf_layer=rbf_layer,
n_interaction_blocks=np.random.randint(2, 5),
calculate_geometry=True,
n_beads=beads)
# Next, we instance a CGSchNet model using the above objects
# with force matching as a loss criterion. We forward the coords
# and the embedding property through as well
model = CGnet(arch, ForceLoss(), feature=feature)
energy, force = model.forward(coords,
embedding_property=embedding_property)
# the force components for masked beads should all be zero if the padding
# due to variable length input is masked properly
# We check each frame of the above output individually:
for i in range(frames):
masked_forces = force[i][variable_beads[i]:]
zero_forces = np.zeros((beads - variable_beads[i], 3))
np.testing.assert_array_equal(masked_forces.detach().numpy(),
zero_forces)
def test_schnet_feature():
# Tests proper forwarding through SchNet wrapper class
# Create random embedding properties
embedding_property = torch.randint(low=1,
high=n_embeddings,
size=(frames, beads))
# Initialize the embedding and SchnetFeature class
embedding_layer = CGBeadEmbedding(n_embeddings=n_embeddings,
embedding_dim=n_feats)
rbf_layer = GaussianRBF()
schnet_feature = SchnetFeature(feature_size=n_feats,
embedding_layer=embedding_layer,
rbf_layer=rbf_layer,
n_interaction_blocks=2,
calculate_geometry=True,
n_beads=beads,
neighbor_cutoff=neighbor_cutoff)
# Next, we check that forwarding cartesian coordinates through SchnetFeature
# produces the correct output feature sizes
schnet_features = schnet_feature(torch.from_numpy(coords),
embedding_property)
assert schnet_features.size() == (frames, beads, n_feats)
# To test the internal logic of SchnetFeature, a full forward pass is
# computed using the same internal components as SchnetFeature.
# First, the embedding feature and radial basis function expansion are
# computed.
features = embedding_layer.forward(embedding_property)
rbf_expansion = schnet_feature.rbf_layer(distances=distances)
# Then, the embedding features are used as an input to the first
# interaction block and as a residual connection, respectively. Depending
# on the amount of interaction blocks, the resulting features are then used
# again as an input for the next interaction block and as a residual
# connection, respectively.
for interaction_block in schnet_feature.interaction_blocks:
interaction_features = interaction_block(features=features,
rbf_expansion=rbf_expansion,
neighbor_list=test_nbh,
neighbor_mask=test_nbh_mask)
features = features + interaction_features
# The manually computed features and the feature from the forward pass
# should be the same.
np.testing.assert_allclose(schnet_features.detach(), features.detach())
def test_schnet_activation_default():
# We check to see if the default activation, ShiftedSoftplus,
# is correctly placed in the SchnetFeature
interaction_blocks = np.random.randint(1, high=5)
rbf_layer = GaussianRBF()
schnet_feature = SchnetFeature(feature_size=n_feats,
embedding_layer=None,
rbf_layer=rbf_layer,
n_interaction_blocks=interaction_blocks,
calculate_geometry=True,
n_beads=beads)
# check all atom-wise layers and the filter generator networks
# in both cases, the second index of the nn.Sequential objects
# that hold the LinearLayers
for interaction_block in schnet_feature.interaction_blocks:
assert isinstance(interaction_block.output_dense[1], ShiftedSoftplus)
assert isinstance(interaction_block.cfconv.filter_generator[1],
ShiftedSoftplus)
def _get_random_schnet_feature(calculate_geometry=False):
""" Helper function for producing SchnetFeature instances with
random initializaitons
Parameters
----------
calculate_geometry : bool (default=False)
specifies whether or not the returned SchnetFeature should
calculate pairwise distances
Returns
-------
SchnetFeature : SchnetFeature
Instance of SchnetFeature with random initialization variables
embedding_property: torch.Tensor
Random embedding property for using in forwarding through
FeatureCombiner or SchnetFeature in tests
feature_size : int
Random feature size for schnet interaction blocks
"""
feature_size = np.random.randint(5, high=10) # random feature size
n_embeddings = np.random.randint(3, high=5) # random embedding number
embedding_dim = feature_size # embedding property size
n_interaction_blocks = np.random.randint(
1, 3) # random number of interactions
neighbor_cutoff = np.random.uniform(0, 1) # random neighbor cutoff
# random embedding property
embedding_property = torch.randint(low=1, high=n_embeddings,
size=(n_frames, n_beads))
embedding_layer = CGBeadEmbedding(n_embeddings=n_embeddings,
embedding_dim=embedding_dim)
rbf_layer = GaussianRBF()
schnet_feature = SchnetFeature(feature_size=feature_size,
embedding_layer=embedding_layer,
rbf_layer=rbf_layer,
n_interaction_blocks=n_interaction_blocks,
calculate_geometry=calculate_geometry,
n_beads=n_beads,
neighbor_cutoff=neighbor_cutoff)
return schnet_feature, embedding_property, feature_size
def test_schnet_activations():
# Tests whether setting the activation function from the SchnetFeature
# level correctly sets the activation for the InteractionBlocks and
# ContinuousFilterConvolutions
# Here, we instance some common activation functions and shuffle them
alt_activations = [nn.Tanh(), nn.ReLU(), nn.ELU(), nn.Sigmoid()]
alt_activation_classes = [nn.Tanh, nn.ReLU, nn.ELU, nn.Sigmoid]
alt_lists = list(zip(alt_activations, alt_activation_classes))
np.random.shuffle(alt_lists)
alt_activations, alt_activation_classes = zip(*alt_lists)
# Here we instance an random number of interaction blocks
interaction_blocks = np.random.randint(1,
high=5,
size=len(alt_activations))
rbf_layer = GaussianRBF()
# Here, we loop through all the activations and make sure that
# they appear where they should in the model
for activation, activation_class, iblock in zip(alt_activations,
alt_activation_classes,
interaction_blocks):
schnet_feature = SchnetFeature(feature_size=n_feats,
embedding_layer=None,
rbf_layer=rbf_layer,
activation=activation,
n_interaction_blocks=iblock,
calculate_geometry=True,
n_beads=beads)
# check all atom-wise layers and the filter generator networks
# in both cases, the second index of the nn.Sequential objects
# that hold the LinearLayers
for interaction_block in schnet_feature.interaction_blocks:
assert isinstance(interaction_block.output_dense[1],
activation_class)
assert isinstance(interaction_block.cfconv.filter_generator[1],
activation_class)
def test_shared_weights():
# Tests the weight sharing functionality of the interaction block
feature_size = np.random.randint(4, 8)
rbf_layer = GaussianRBF()
# Initialize two Schnet networks
# With and without weight sharing, respectively.
schnet_feature_no_shared_weights = SchnetFeature(feature_size=feature_size,
embedding_layer=None,
rbf_layer=rbf_layer,
n_interaction_blocks=2,
n_beads=beads,
share_weights=False)
schnet_feature_shared_weights = SchnetFeature(feature_size=feature_size,
embedding_layer=None,
rbf_layer=rbf_layer,
n_interaction_blocks=2,
n_beads=beads,
share_weights=True)
# Loop over all parameters in both interaction blocks
# If the weights are shared, the parameters in both interaction blocks
# should be equal, respectively.
for param1, param2 in zip(
schnet_feature_shared_weights.interaction_blocks[0].parameters(),
schnet_feature_shared_weights.interaction_blocks[1].parameters()):
assert np.array_equal(param1.detach().numpy(), param2.detach().numpy())
# If the weights are not shared, the parameters should be different.
for param1, param2 in zip(
schnet_feature_no_shared_weights.interaction_blocks[0].parameters(
),
schnet_feature_no_shared_weights.interaction_blocks[1].parameters(
)):
assert not np.array_equal(param1.detach().numpy(),
param2.detach().numpy())
# 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
embedding_dim = 10 # embedding property size
n_interaction_blocks = np.random.randint(1, 3) # random number of interactions
neighbor_cutoff = np.random.uniform(0, 1) # random neighbor cutoff
embedding_layer = CGBeadEmbedding(n_embeddings=embedding_dim,
embedding_dim=feature_size)
rbf_layer = GaussianRBF()
# Here we use the above variables to create the SchnetFeature
schnet_feature = SchnetFeature(feature_size=feature_size,
embedding_layer=embedding_layer,
rbf_layer=rbf_layer,
n_interaction_blocks=n_interaction_blocks,
calculate_geometry=True,
n_beads=beads,
neighbor_cutoff=neighbor_cutoff)
schnet_model = CGnet(arch, ForceLoss(), feature=schnet_feature).eval()
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# The following tests probe basic shapes/functionalities of the simulation #
# class and are repeated for Brownian (i.e., overdamped Langevin) and. #
# Langevin simulations. Checks are routinely made to make sure that. #