ANI-1 elements (1e545237) · Commits · 钟慕尧 / deepchem

deepchem/models/tensorgraph/models/symmetry_function_regression.py

+9 −4

Original line number	Diff line number	Diff line
		@@ -18,7 +18,7 @@ from deepchem.models.tensorgraph.symmetry_functions import DistanceMatrix, \

		class BPSymmetryFunctionRegression(TensorGraph):

		def __init__(self, n_tasks, max_atoms, n_hidden=10, n_embedding=10, **kwargs):
		def __init__(self, n_tasks, max_atoms, n_hidden=40, n_embedding=10, **kwargs):
		"""
		Parameters
		----------
		@@ -52,7 +52,8 @@ class BPSymmetryFunctionRegression(TensorGraph):
		angular_symmetry = AngularSymmetry(
		self.max_atoms,
		in_layers=[distance_cutoff, distance_matrix, self.atom_coordinates])
		atom_embedding = DTNNEmbedding(n_embedding=self.n_embedding, in_layers=[self.atom_numbers])
		atom_embedding = DTNNEmbedding(
		n_embedding=self.n_embedding, in_layers=[self.atom_numbers])

		feature_merge = BPFeatureMerge(
		self.max_atoms,
		@@ -64,11 +65,15 @@ class BPSymmetryFunctionRegression(TensorGraph):
		out_channels=self.n_hidden,
		activation_fn=tf.nn.tanh,
		in_layers=[feature_merge])

		Hidden2 = Dense(
		out_channels=self.n_hidden,
		activation_fn=tf.nn.tanh,
		in_layers=[Hidden])
		costs = []
		self.labels_fd = []
		for task in range(self.n_tasks):
		regression = Dense(out_channels=1, activation_fn=None, in_layers=[Hidden])
		regression = Dense(
		out_channels=1, activation_fn=None, in_layers=[Hidden2])
		output = BPGather(self.max_atoms, in_layers=[regression, self.atom_flags])
		self.add_output(output)

deepchem/models/tensorgraph/symmetry_functions.py

+210 −21

Original line number	Diff line number	Diff line
		@@ -17,6 +17,7 @@ from deepchem.nn import model_ops

		from deepchem.models.tensorgraph.layers import Layer
		from deepchem.models.tensorgraph.layers import convert_to_layers
		from deepchem.metrics import to_one_hot


		class DistanceMatrix(Layer):
		@@ -86,15 +87,36 @@ class DistanceCutoff(Layer):
		class RadialSymmetry(Layer):
		""" Radial Symmetry Function """

		def __init__(self, max_atoms, **kwargs):
		def __init__(self,
		max_atoms,
		Rs_init=None,
		ita_init=None,
		atomic_number_differentiated=False,
		atom_numbers=[1, 6, 7, 8],
		**kwargs):
		self.max_atoms = max_atoms
		self.atomic_number_differentiated = atomic_number_differentiated
		self.atom_number_cases = atom_numbers
		if Rs_init is None:
		self.Rs_init = np.array([0.5, 1.17, 1.83, 2.5, 3.17, 3.83, 4.5, 5.17])
		else:
		self.Rs_init = np.array(Rs_init)
		if ita_init is None:
		self.ita_init = np.array([4.])
		else:
		self.ita_init = np.array(ita_init)

		super(RadialSymmetry, self).__init__(**kwargs)

		def build(self):
		""" Parameters for the Gaussian """
		self.Rs = tf.Variable(tf.constant(0.))
		#self.ita = tf.exp(tf.Variable(tf.constant(0.)))
		self.ita = 1.
		len_Rs = len(self.Rs_init)
		len_ita = len(self.ita_init)
		self.length = len_Rs * len_ita
		Rs_init, ita_init = np.meshgrid(self.Rs_init, self.ita_init)
		self.Rs = tf.constant(Rs_init.flatten(), dtype=tf.float32)
		self.ita = tf.constant(ita_init.flatten(), dtype=tf.float32)
		self.atom_number_embedding = tf.eye(max(self.atom_number_cases))

		def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):
		""" Generate Radial Symmetry Function """
		@@ -105,24 +127,65 @@ class RadialSymmetry(Layer):
		self.build()
		d_cutoff = in_layers[0].out_tensor
		d = in_layers[1].out_tensor
		out_tensor = tf.exp(-self.ita * tf.square(d - self.Rs)) * d_cutoff
		if self.atomic_number_differentiated:
		atom_numbers = in_layers[2].out_tensor
		atom_number_embedded = tf.nn.embedding_lookup(self.atom_number_embedding,
		atom_numbers)
		d_cutoff = tf.stack([d_cutoff] * self.length, axis=3)
		d = tf.stack([d] * self.length, axis=3)
		Rs = tf.reshape(self.Rs, (1, 1, 1, -1))
		ita = tf.reshape(self.ita, (1, 1, 1, -1))
		out_tensor = tf.exp(-ita * tf.square(d - Rs)) * d_cutoff
		if self.atomic_number_differentiated:
		out_tensors = []
		for atom_type in self.atom_number_cases:
		selected_atoms = tf.expand_dims(
		tf.expand_dims(atom_number_embedded[:, :, atom_type], axis=1),
		axis=3)
		out_tensors.append(tf.reduce_sum(out_tensor * selected_atoms, axis=2))
		self.out_tensor = tf.concat(out_tensors, axis=2)
		else:
		self.out_tensor = tf.reduce_sum(out_tensor, axis=2)


		class AngularSymmetry(Layer):
		""" Angular Symmetry Function """

		def __init__(self, max_atoms, **kwargs):
		def __init__(self,
		max_atoms,
		lambd_init=None,
		ita_init=None,
		zeta_init=None,
		**kwargs):
		self.max_atoms = max_atoms
		if lambd_init is None:
		self.lambd_init = np.array([1., -1.])
		else:
		self.lambd_init = np.array(lambd_init)

		if ita_init is None:
		self.ita_init = np.array([4.])
		else:
		self.ita_init = np.array(ita_init)

		if zeta_init is None:
		self.zeta_init = np.array([2., 4., 8.])
		else:
		self.zeta_init = np.array(zeta_init)

		super(AngularSymmetry, self).__init__(**kwargs)

		def build(self):
		#self.lambd = tf.Variable(tf.constant(1.))
		self.lambd = 1.
		#self.ita = tf.exp(tf.Variable(tf.constant(0.)))
		self.ita = 1.
		#self.zeta = tf.Variable(tf.constant(0.8))
		self.zeta = 0.8
		len_lambd = len(self.lambd_init)
		len_ita = len(self.ita_init)
		len_zeta = len(self.zeta_init)
		self.length = len_lambd * len_ita * len_zeta

		lambd_init, ita_init, zeta_init = np.meshgrid(self.lambd_init,
		self.ita_init, self.zeta_init)
		self.lambd = tf.constant(lambd_init.flatten(), dtype=tf.float32)
		self.ita = tf.constant(ita_init.flatten(), dtype=tf.float32)
		self.zeta = tf.constant(zeta_init.flatten(), dtype=tf.float32)

		def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):
		""" Generate Angular Symmetry Function """
		@@ -150,11 +213,141 @@ class AngularSymmetry(Layer):

		theta = tf.div(theta, R_ij * R_ik + 1e-5)

		out_tensor = tf.pow(1+self.lambdtf.cos(theta), self.zeta) \
		tf.exp(-self.ita(tf.square(R_ij)+tf.square(R_ik)+tf.square(R_jk))) \
		R_ij = tf.stack([R_ij] * self.length, axis=4)
		R_ik = tf.stack([R_ik] * self.length, axis=4)
		R_jk = tf.stack([R_jk] * self.length, axis=4)
		f_R_ij = tf.stack([f_R_ij] * self.length, axis=4)
		f_R_ik = tf.stack([f_R_ik] * self.length, axis=4)
		f_R_jk = tf.stack([f_R_jk] * self.length, axis=4)

		theta = tf.stack([theta] * self.length, axis=4)
		lambd = tf.reshape(self.lambd, (1, 1, 1, 1, -1))
		zeta = tf.reshape(self.zeta, (1, 1, 1, 1, -1))
		ita = tf.reshape(self.ita, (1, 1, 1, 1, -1))

		out_tensor = tf.pow(1+lambdtf.cos(theta), zeta) \
		tf.exp(-ita(tf.square(R_ij)+tf.square(R_ik)+tf.square(R_jk))) \
		f_R_ij * f_R_ik * f_R_jk
		self.out_tensor = tf.reduce_sum(out_tensor, axis=[2, 3]) * \
		tf.pow(tf.constant(2.), 1-self.zeta)
		tf.pow(tf.constant(2.), 1-tf.reshape(self.zeta, (1,1,-1)))


		class AngularSymmetryMod(Layer):
		""" Angular Symmetry Function """

		def __init__(self,
		max_atoms,
		lambd_init=None,
		ita_init=None,
		zeta_init=None,
		Rs_init=None,
		thetas_init=None,
		atomic_number_differentiated=False,
		atom_numbers=[1, 6, 7, 8],
		**kwargs):
		self.max_atoms = max_atoms
		self.atomic_number_differentiated = atomic_number_differentiated
		self.atom_number_cases = atom_numbers
		if lambd_init is None:
		self.lambd_init = np.array([1., -1.])
		else:
		self.lambd_init = np.array(lambd_init)

		if ita_init is None:
		self.ita_init = np.array([4.])
		else:
		self.ita_init = np.array(ita_init)

		if zeta_init is None:
		self.zeta_init = np.array([2., 4., 8.])
		else:
		self.zeta_init = np.array(zeta_init)

		if Rs_init is None:
		self.Rs_init = np.array([0.5, 1.17, 1.83, 2.5, 3.17, 3.83, 4.5, 5.17])
		else:
		self.Rs_init = np.array(Rs_init)

		if thetas_init is None:
		self.thetas_init = np.array([0., 1.57, 3.14, 4.71])
		else:
		self.thetas_init = np.array(thetas_init)
		super(AngularSymmetryMod, self).__init__(**kwargs)

		def build(self):
		len_lambd = len(self.lambd_init)
		len_ita = len(self.ita_init)
		len_zeta = len(self.zeta_init)
		len_Rs = len(self.Rs_init)
		len_thetas = len(self.thetas_init)
		self.length = len_lambd * len_ita * len_zeta * len_Rs * len_thetas

		lambd_init, ita_init, zeta_init, Rs_init, thetas_init = \
		np.meshgrid(self.lambd_init, self.ita_init, self.zeta_init, self.Rs_init, self.thetas_init)
		self.lambd = tf.constant(lambd_init.flatten(), dtype=tf.float32)
		self.ita = tf.constant(ita_init.flatten(), dtype=tf.float32)
		self.zeta = tf.constant(zeta_init.flatten(), dtype=tf.float32)
		self.Rs = tf.constant(Rs_init.flatten(), dtype=tf.float32)
		self.thetas = tf.constant(thetas_init.flatten(), dtype=tf.float32)
		self.atom_number_embedding = tf.eye(max(self.atom_number_cases))

		def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):
		""" Generate Angular Symmetry Function """
		if in_layers is None:
		in_layers = self.in_layers
		in_layers = convert_to_layers(in_layers)

		self.build()
		max_atoms = self.max_atoms
		d_cutoff = in_layers[0].out_tensor
		d = in_layers[1].out_tensor
		atom_coordinates = in_layers[2].out_tensor
		if self.atomic_number_differentiated:
		atom_numbers = in_layers[3].out_tensor
		atom_number_embedded = tf.nn.embedding_lookup(self.atom_number_embedding,
		atom_numbers)

		vector_distances = tf.tile(tf.expand_dims(atom_coordinates, axis=2), (1,1,max_atoms,1)) - \
		tf.tile(tf.expand_dims(atom_coordinates, axis=1), (1,max_atoms,1,1))
		R_ij = tf.tile(tf.expand_dims(d, axis=3), (1, 1, 1, max_atoms))
		R_ik = tf.tile(tf.expand_dims(d, axis=2), (1, 1, max_atoms, 1))
		f_R_ij = tf.tile(tf.expand_dims(d_cutoff, axis=3), (1, 1, 1, max_atoms))
		f_R_ik = tf.tile(tf.expand_dims(d_cutoff, axis=2), (1, 1, max_atoms, 1))

		# Define angle theta = R_ij(Vector) dot R_ik(Vector)/R_ij(distance)/R_ik(distance)
		theta = tf.reduce_sum(tf.tile(tf.expand_dims(vector_distances, axis=3), (1,1,1,max_atoms,1)) * \
		tf.tile(tf.expand_dims(vector_distances, axis=2), (1,1,max_atoms,1,1)), axis=4)

		theta = tf.div(theta, R_ij * R_ik + 1e-5)

		R_ij = tf.stack([R_ij] * self.length, axis=4)
		R_ik = tf.stack([R_ik] * self.length, axis=4)
		f_R_ij = tf.stack([f_R_ij] * self.length, axis=4)
		f_R_ik = tf.stack([f_R_ik] * self.length, axis=4)

		theta = tf.stack([theta] * self.length, axis=4)
		lambd = tf.reshape(self.lambd, (1, 1, 1, 1, -1))
		zeta = tf.reshape(self.zeta, (1, 1, 1, 1, -1))
		ita = tf.reshape(self.ita, (1, 1, 1, 1, -1))
		Rs = tf.reshape(self.Rs, (1, 1, 1, 1, -1))
		thetas = tf.reshape(self.thetas, (1, 1, 1, 1, -1))

		out_tensor = tf.pow(1+lambdtf.cos(theta - thetas), zeta) \
		tf.exp(-itatf.square((R_ij+R_ik)/2-Rs)) \
		f_R_ij * f_R_ik * tf.pow(tf.constant(2.), 1 - zeta)
		if self.atomic_number_differentiated:
		out_tensors = []
		for atom_type_j in self.atom_number_cases:
		for atom_type_k in self.atom_number_cases:
		selected_atoms = tf.stack([atom_number_embedded[:, :, atom_type_j]] * max_atoms, axis=2) * \
		tf.stack([atom_number_embedded[:, :, atom_type_k]] * max_atoms, axis=1)
		selected_atoms = tf.expand_dims(
		tf.expand_dims(selected_atoms, axis=1), axis=4)
		out_tensors.append(
		tf.reduce_sum(out_tensor * selected_atoms, axis=[2, 3]))
		self.out_tensor = tf.concat(out_tensors, axis=2)
		else:
		self.out_tensor = tf.reduce_sum(out_tensor, axis=[2, 3])


		class BPFeatureMerge(Layer):
		@@ -175,11 +368,7 @@ class BPFeatureMerge(Layer):
		atom_flags = in_layers[3].out_tensor

		out_tensor = tf.concat(
		[
		atom_embedding, tf.expand_dims(radial_symmetry, 2),
		tf.expand_dims(angular_symmetry, 2)
		],
		axis=2)
		[atom_embedding, radial_symmetry, angular_symmetry], axis=2)
		self.out_tensor = out_tensor * atom_flags[:, :, 0:1]

examples/qm7/qm7_BPSymmetryFunction.py

+3 −3

Original line number	Diff line number	Diff line
		@@ -24,7 +24,7 @@ metric = [

		# Batch size of models
		max_atoms = 23
		n_hidden = 10
		n_hidden = 40
		n_embedding = 0
		batch_size = 16

		@@ -34,12 +34,12 @@ model = dc.models.BPSymmetryFunctionRegression(
		n_hidden=n_hidden,
		n_embedding=n_embedding,
		batch_size=batch_size,
		learning_rate=0.0001,
		learning_rate=0.001,
		use_queue=False,
		mode="regression")

		# Fit trained model
		model.fit(train_dataset, nb_epoch=200, checkpoint_interval=100)
		model.fit(train_dataset, nb_epoch=20, checkpoint_interval=1000)

		print("Evaluating model")
		train_scores = model.evaluate(train_dataset, metric, transformers)

Admin message