def test_smear_gaussian_trainable(): dist = torch.tensor([[[0.0, 1.0, 1.5, 0.25], [0.5, 1.5, 3.0, 1.0]]]) # smear using 5 Gaussian functions with 0.75 spacing smear = GaussianSmearing(start=1., stop=4., n_gaussians=5, trainable=True) # absolute value of centered distances expt = torch.tensor([[[[1, 1.75, 2.5, 3.25, 4.], [0, 0.75, 1.5, 2.25, 3.], [0.5, 0.25, 1., 1.75, 2.5], [0.75, 1.5, 2.25, 3., 3.75]], [[0.5, 1.25, 2., 2.75, 3.5], [0.5, 0.25, 1., 1.75, 2.5], [2., 1.25, 0.5, 0.25, 1.], [0, 0.75, 1.5, 2.25, 3.]]]]) expt = torch.exp((-0.5 / 0.75**2) * expt**2) assert torch.allclose(expt, smear(dist), atol=0.0, rtol=1.0e-7) params = list(smear.parameters()) assert len(params) == 2 assert len(params[0]) == 5 assert len(params[1]) == 5 # centered = True smear = GaussianSmearing(start=1., stop=4., n_gaussians=5, trainable=True, centered=True) expt = -0.5 / torch.tensor([1, 1.75, 2.5, 3.25, 4])**2 expt = torch.exp(expt * dist[:, :, :, None]**2) assert torch.allclose(expt, smear(dist), atol=0.0, rtol=1.0e-7) params = list(smear.parameters()) assert len(params) == 2 assert len(params[0]) == 5 assert len(params[1]) == 5
def test_smear_gaussian(): dist = torch.tensor([[[0.0, 1.0, 1.5], [0.5, 1.5, 3.0]]]) # smear using 4 Gaussian functions with 1. spacing smear = GaussianSmearing(start=1., stop=4., n_gaussians=4) # absolute value of centered distances expt = torch.tensor([[[[1, 2, 3, 4], [0, 1, 2, 3], [0.5, 0.5, 1.5, 2.5]], [[.5, 1.5, 2.5, 3.5], [.5, .5, 1.5, 2.5], [2, 1, 0, 1]]]]) expt = torch.exp(-0.5 * expt**2) assert torch.allclose(expt, smear(dist), atol=0.0, rtol=1.0e-7) assert list(smear.parameters()) == [] # centered = True smear = GaussianSmearing(start=1., stop=4., n_gaussians=4, centered=True) expt = torch.exp((-0.5 / torch.tensor([1, 2, 3, 4.])**2) * dist[:, :, :, None]**2) assert torch.allclose(expt, smear(dist), atol=0.0, rtol=1.0e-7) assert list(smear.parameters()) == []
def test_smear_gaussian_one_distance(): # case of one distance dist = torch.tensor([[[1.0]]]) # trainable = False smear = GaussianSmearing(n_gaussians=6, centered=False, trainable=False) expt = torch.exp(-0.5 * torch.tensor([[[1., 0., 1., 4., 9., 16.]]])) assert torch.allclose(expt, smear(dist), atol=0.0, rtol=1.0e-7) assert list(smear.parameters()) == [] # trainable = True smear = GaussianSmearing(n_gaussians=6, centered=False, trainable=True) assert torch.allclose(expt, smear(dist), atol=0.0, rtol=1.0e-7) params = list(smear.parameters()) assert len(params) == 2 assert len(params[0]) == 6 assert len(params[1]) == 6 # centered = True smear = GaussianSmearing(n_gaussians=6, centered=True) expt = -0.5 / torch.tensor([0., 1, 2, 3, 4, 5])**2 expt = torch.exp(expt * dist[:, :, :, None]**2) assert torch.allclose(expt, smear(dist), atol=0.0, rtol=1.0e-7) assert list(smear.parameters()) == []
def test_gaussian_smearing(n_spatial_basis, distances): model = GaussianSmearing(n_gaussians=n_spatial_basis) out_shape = [*list(distances.shape), n_spatial_basis] inputs = [distances] assert_equal_shape(model, inputs, out_shape)
def __init__(self, n_atom_basis=128, n_filters=128, n_interactions=3, cutoff=5.0, n_gaussians=25, normalize_filter=False, coupled_interactions=False, return_intermediate=False, max_z=100, cutoff_network=HardCutoff, trainable_gaussians=False, distance_expansion=None, charged_systems=False, use_noise=False, noise_mean=0, noise_std=1, chargeEmbedding=True, ownFeatures=False, nFeatures=8, finalFeature=None, finalFeatureStart=7, finalFeatureStop=8): super(SchNet, self).__init__() self.finalFeature = finalFeature self.finalFeatureStart = finalFeatureStart self.finalFeatureStop = finalFeatureStop self.chargeEmbedding = chargeEmbedding self.ownFeatures = ownFeatures self.n_atom_basis = n_atom_basis # make a lookup table to store embeddings for each element (up to atomic # number max_z) each of which is a vector of size n_atom_basis if chargeEmbedding and not ownFeatures: self.embedding = nn.Embedding(max_z, n_atom_basis, padding_idx=0) elif chargeEmbedding and ownFeatures: if nFeatures is None: raise NotImplementedError self.embedding = nn.Embedding(max_z, int(n_atom_basis / 2), padding_idx=0) self.denseEmbedding = Dense(nFeatures, int(n_atom_basis / 2)) elif ownFeatures and not chargeEmbedding: if nFeatures is None: raise NotImplementedError self.denseEmbedding = Dense(nFeatures, n_atom_basis) else: raise NotImplementedError # layer for computing interatomic distances self.distances = AtomDistances() # layer for expanding interatomic distances in a basis if distance_expansion is None: self.distance_expansion = GaussianSmearing( 0.0, cutoff, n_gaussians, trainable=trainable_gaussians) else: self.distance_expansion = distance_expansion # block for computing interaction if isinstance(n_filters, list): self.interactions = nn.ModuleList([ SchNetInteraction( n_atom_basis=n_atom_basis, n_spatial_basis=n_gaussians, n_filters=n_filters[i], cutoff_network=cutoff_network, cutoff=cutoff, normalize_filter=normalize_filter, ) for i in range(n_interactions) ]) elif coupled_interactions: # use the same SchNetInteraction instance (hence the same weights) self.interactions = nn.ModuleList([ SchNetInteraction( n_atom_basis=n_atom_basis, n_spatial_basis=n_gaussians, n_filters=n_filters, cutoff_network=cutoff_network, cutoff=cutoff, normalize_filter=normalize_filter, ) ] * n_interactions) else: # use one SchNetInteraction instance for each interaction self.interactions = nn.ModuleList([ SchNetInteraction( n_atom_basis=n_atom_basis, n_spatial_basis=n_gaussians, n_filters=n_filters, cutoff_network=cutoff_network, cutoff=cutoff, normalize_filter=normalize_filter, ) for _ in range(n_interactions) ]) # set attributes self.use_noise = use_noise self.noise_mean = noise_mean self.noise_std = noise_std self.return_intermediate = return_intermediate self.charged_systems = charged_systems if charged_systems: self.charge = nn.Parameter(torch.Tensor(1, n_atom_basis)) self.charge.data.normal_(0, 1.0 / n_atom_basis**0.5)
def __init__( self, n_atom_basis=128, n_filters=128, n_interactions=3, cutoff=5.0, n_gaussians=25, normalize_filter=False, coupled_interactions=False, return_intermediate=False, max_z=100, cutoff_network=CosineCutoff, trainable_gaussians=False, distance_expansion=None, charged_systems=False, ): super(SchNet, self).__init__() self.n_atom_basis = n_atom_basis # make a lookup table to store embeddings for each element (up to atomic # number max_z) each of which is a vector of size n_atom_basis self.embedding = nn.Embedding(max_z, n_atom_basis, padding_idx=0) # layer for computing interatomic distances self.distances = AtomDistances() # layer for expanding interatomic distances in a basis if distance_expansion is None: self.distance_expansion = GaussianSmearing( 0.0, cutoff, n_gaussians, trainable=trainable_gaussians) else: self.distance_expansion = distance_expansion # block for computing interaction if coupled_interactions: # use the same SchNetInteraction instance (hence the same weights) self.interactions = nn.ModuleList([ SchNetInteraction( n_atom_basis=n_atom_basis, n_spatial_basis=n_gaussians, n_filters=n_filters, cutoff_network=cutoff_network, cutoff=cutoff, normalize_filter=normalize_filter, ) ] * n_interactions) else: # use one SchNetInteraction instance for each interaction self.interactions = nn.ModuleList([ SchNetInteraction( n_atom_basis=n_atom_basis, n_spatial_basis=n_gaussians, n_filters=n_filters, cutoff_network=cutoff_network, cutoff=cutoff, normalize_filter=normalize_filter, ) for _ in range(n_interactions) ]) # set attributes self.return_intermediate = return_intermediate self.charged_systems = charged_systems if charged_systems: self.charge = nn.Parameter(torch.Tensor(1, n_atom_basis)) self.charge.data.normal_(0, 1.0 / n_atom_basis**0.5)