Loading deepchem/models/tensorgraph/layers.py +179 −0 Original line number Diff line number Diff line Loading @@ -3248,6 +3248,185 @@ class BetaShare(Layer): self.out_tensor, self.betas = tensor class ANIFeat(Layer): """Performs transform from 3D coordinates to ANI symmetry functions """ def __init__(self, in_layers, max_atoms=23, radial_cutoff=4.6, angular_cutoff=3.1, radial_length=32, angular_length=8, atom_cases=[1, 6, 7, 8, 16], atomic_number_differentiated=True, coordinates_in_bohr=True, **kwargs): """ Only X can be transformed """ self.max_atoms = max_atoms self.radial_cutoff = radial_cutoff self.angular_cutoff = angular_cutoff self.radial_length = radial_length self.angular_length = angular_length self.atom_cases = atom_cases self.atomic_number_differentiated = atomic_number_differentiated self.coordinates_in_bohr = coordinates_in_bohr super(ANIFeat, self).__init__(in_layers, **kwargs) def create_tensor(self, in_layers=None, set_tensors=True, **kwargs): """ In layers should be of shape dtype tf.float32, (None, self.max_atoms, 4) """ inputs = self._get_input_tensors(in_layers)[0] atom_numbers = tf.cast(inputs[:, :, 0], tf.int32) flags = tf.sign(atom_numbers) flags = tf.to_float(tf.expand_dims(flags, 1) * tf.expand_dims(flags, 2)) coordinates = inputs[:, :, 1:] if self.coordinates_in_bohr: coordinates = coordinates * 0.52917721092 d = self.distance_matrix(coordinates, flags) d_radial_cutoff = self.distance_cutoff(d, self.radial_cutoff, flags) d_angular_cutoff = self.distance_cutoff(d, self.angular_cutoff, flags) radial_sym = self.radial_symmetry(d_radial_cutoff, d, atom_numbers) angular_sym = self.angular_symmetry(d_angular_cutoff, d, atom_numbers, coordinates) out_tensor = tf.concat( [tf.to_float(tf.expand_dims(atom_numbers, 2)), radial_sym, angular_sym], axis=2) if set_tensors: self.out_tensor = out_tensor return out_tensor def distance_matrix(self, coordinates, flags): """ Generate distance matrix """ # (TODO YTZ:) faster, less memory intensive way # r = tf.reduce_sum(tf.square(coordinates), 2) # r = tf.expand_dims(r, -1) # inner = 2*tf.matmul(coordinates, tf.transpose(coordinates, perm=[0,2,1])) # # inner = 2*tf.matmul(coordinates, coordinates, transpose_b=True) # d = r - inner + tf.transpose(r, perm=[0,2,1]) # d = tf.nn.relu(d) # fix numerical instabilities about diagonal # d = tf.sqrt(d) # does this have negative elements? may be unstable for diagonals max_atoms = self.max_atoms tensor1 = tf.stack([coordinates] * max_atoms, axis=1) tensor2 = tf.stack([coordinates] * max_atoms, axis=2) # Calculate pairwise distance d = tf.sqrt( tf.reduce_sum(tf.squared_difference(tensor1, tensor2), axis=3) + 1e-7) d = d * flags return d def distance_cutoff(self, d, cutoff, flags): """ Generate distance matrix with trainable cutoff """ # Cutoff with threshold Rc d_flag = flags * tf.sign(cutoff - d) d_flag = tf.nn.relu(d_flag) d_flag = d_flag * tf.expand_dims((1 - tf.eye(self.max_atoms)), 0) d = 0.5 * (tf.cos(np.pi * d / cutoff) + 1) return d * d_flag # return d def radial_symmetry(self, d_cutoff, d, atom_numbers): """ Radial Symmetry Function """ embedding = tf.eye(np.max(self.atom_cases) + 1) atom_numbers_embedded = tf.nn.embedding_lookup(embedding, atom_numbers) Rs = np.linspace(0., self.radial_cutoff, self.radial_length) ita = np.ones_like(Rs) * 3 / (Rs[1] - Rs[0])**2 Rs = tf.to_float(np.reshape(Rs, (1, 1, 1, -1))) ita = tf.to_float(np.reshape(ita, (1, 1, 1, -1))) length = ita.get_shape().as_list()[-1] d_cutoff = tf.stack([d_cutoff] * length, axis=3) d = tf.stack([d] * length, axis=3) out = tf.exp(-ita * tf.square(d - Rs)) * d_cutoff if self.atomic_number_differentiated: out_tensors = [] for atom_type in self.atom_cases: selected_atoms = tf.expand_dims( tf.expand_dims(atom_numbers_embedded[:, :, atom_type], axis=1), axis=3) out_tensors.append(tf.reduce_sum(out * selected_atoms, axis=2)) return tf.concat(out_tensors, axis=2) else: return tf.reduce_sum(out, axis=2) def angular_symmetry(self, d_cutoff, d, atom_numbers, coordinates): """ Angular Symmetry Function """ max_atoms = self.max_atoms embedding = tf.eye(np.max(self.atom_cases) + 1) atom_numbers_embedded = tf.nn.embedding_lookup(embedding, atom_numbers) Rs = np.linspace(0., self.angular_cutoff, self.angular_length) ita = 3 / (Rs[1] - Rs[0])**2 thetas = np.linspace(0., np.pi, self.angular_length) zeta = float(self.angular_length**2) ita, zeta, Rs, thetas = np.meshgrid(ita, zeta, Rs, thetas) zeta = tf.to_float(np.reshape(zeta, (1, 1, 1, 1, -1))) ita = tf.to_float(np.reshape(ita, (1, 1, 1, 1, -1))) Rs = tf.to_float(np.reshape(Rs, (1, 1, 1, 1, -1))) thetas = tf.to_float(np.reshape(thetas, (1, 1, 1, 1, -1))) length = zeta.get_shape().as_list()[-1] # tf.stack issues again... vector_distances = tf.stack([coordinates] * max_atoms, 1) - tf.stack( [coordinates] * max_atoms, 2) R_ij = tf.stack([d] * max_atoms, axis=3) R_ik = tf.stack([d] * max_atoms, axis=2) f_R_ij = tf.stack([d_cutoff] * max_atoms, axis=3) f_R_ik = tf.stack([d_cutoff] * max_atoms, axis=2) # Define angle theta = arccos(R_ij(Vector) dot R_ik(Vector)/R_ij(distance)/R_ik(distance)) vector_mul = tf.reduce_sum(tf.stack([vector_distances]*max_atoms, axis=3) * \ tf.stack([vector_distances]*max_atoms, axis=2), axis=4) vector_mul = vector_mul * tf.sign(f_R_ij) * tf.sign(f_R_ik) theta = tf.acos(tf.div(vector_mul, R_ij * R_ik + 1e-5)) R_ij = tf.stack([R_ij] * length, axis=4) R_ik = tf.stack([R_ik] * length, axis=4) f_R_ij = tf.stack([f_R_ij] * length, axis=4) f_R_ik = tf.stack([f_R_ik] * length, axis=4) theta = tf.stack([theta] * length, axis=4) out_tensor = tf.pow((1. + tf.cos(theta - thetas))/2., zeta) * \ tf.exp(-ita * tf.square((R_ij + R_ik)/2. - Rs)) * f_R_ij * f_R_ik * 2 if self.atomic_number_differentiated: out_tensors = [] for id_j, atom_type_j in enumerate(self.atom_cases): for atom_type_k in self.atom_cases[id_j:]: selected_atoms = tf.stack([atom_numbers_embedded[:, :, atom_type_j]] * max_atoms, axis=2) * \ tf.stack([atom_numbers_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))) return tf.concat(out_tensors, axis=2) else: return tf.reduce_sum(out_tensor, axis=(2, 3)) def get_num_feats(self): n_feat = self.outputs.get_shape().as_list()[-1] return n_feat class PassThroughLayer(Layer): """ Layer which takes a tensor from in_tensor[0].out_tensors at an index Loading deepchem/models/tensorgraph/models/symmetry_function_regression.py +238 −6 Original line number Diff line number Diff line Loading @@ -4,11 +4,18 @@ Created on Thu Jul 6 20:31:47 2017 @author: zqwu @contributors: ytz """ import os import numpy as np import json import scipy.optimize import tensorflow as tf from deepchem.models.tensorgraph.layers import Dense, Concat, WeightedError, Stack import deepchem as dc from deepchem.models.tensorgraph.layers import Dense, Concat, WeightedError, Stack, Layer, ANIFeat from deepchem.models.tensorgraph.layers import L2Loss, Label, Weights, Feature from deepchem.models.tensorgraph.tensor_graph import TensorGraph from deepchem.models.tensorgraph.graph_layers import DTNNEmbedding Loading Loading @@ -39,7 +46,9 @@ class BPSymmetryFunctionRegression(TensorGraph): self.max_atoms = max_atoms self.n_feat = n_feat self.layer_structures = layer_structures super(BPSymmetryFunctionRegression, self).__init__(**kwargs) self.build_graph() def build_graph(self): Loading Loading @@ -106,7 +115,6 @@ class ANIRegression(TensorGraph): def __init__(self, n_tasks, max_atoms, n_feat, layer_structures=[128, 64], atom_number_cases=[1, 6, 7, 8, 16], **kwargs): Loading @@ -122,17 +130,181 @@ class ANIRegression(TensorGraph): """ self.n_tasks = n_tasks self.max_atoms = max_atoms self.n_feat = n_feat self.layer_structures = layer_structures self.atom_number_cases = atom_number_cases super(ANIRegression, self).__init__(**kwargs) # (ytz): this is really dirty but needed for restoring models self._kwargs = { "n_tasks": n_tasks, "max_atoms": max_atoms, "layer_structures": layer_structures, "atom_number_cases": atom_number_cases } self._kwargs.update(kwargs) self.build_graph() self.grad = None def save(self): self.grad = None # recompute grad on restore super(ANIRegression, self).save() def build_grad(self): self.grad = tf.gradients(self.outputs, self.atom_feats) def compute_grad(self, dataset, upper_lim=1): """ Computes a batched gradients given an input dataset. Parameters ---------- dataset: dc.Dataset dataset-like object whose X values will be used to compute gradients from upper_lim: int subset of dataset used. Returns ------- np.array Gradients of the input of shape (max_atoms, 4). Note that it is up to the end user to slice this matrix into the correct shape, since it's very likely the derivatives with respect to the atomic numbers are zero. """ with self._get_tf("Graph").as_default(): if not self.built: self.build() if not self.grad: self.build_grad() feed_dict = dict() X = dataset.X flags = np.sign(np.array(X[:upper_lim, :, 0])) feed_dict[self.atom_flags] = np.stack([flags]*self.max_atoms, axis=2)*\ np.stack([flags]*self.max_atoms, axis=1) feed_dict[self.atom_numbers] = np.array(X[:upper_lim, :, 0], dtype=int) feed_dict[self.atom_feats] = np.array(X[:upper_lim, :, :], dtype=float) return self.session.run([self.grad], feed_dict=feed_dict) def pred_one(self, X, atomic_nums, constraints=None): """ Makes an energy prediction for a set of atomic coordinates. Parameters ---------- X: np.array numpy array of shape (a, 3) where a <= max_atoms and dtype is float-like atomic_nums: np.array numpy array of shape (a,) where a is the same as that of X. constraints: unused This parameter is mainly for compatibility purposes for scipy optimize Returns ------- float Predicted energy. Note that the meaning of the returned value is dependent on the training y-values both in semantics (relative vs absolute) and units (kcal/mol vs Hartrees) """ num_atoms = atomic_nums.shape[0] X = X.reshape((num_atoms, 3)) A = atomic_nums.reshape((atomic_nums.shape[0], 1)) Z = np.zeros((self.max_atoms, 4)) Z[:X.shape[0], 1:X.shape[1] + 1] = X Z[:A.shape[0], :A.shape[1]] = A X = Z dd = dc.data.NumpyDataset( np.array(X).reshape((1, self.max_atoms, 4)), np.array(0), np.array(1)) return self.predict(dd)[0] def grad_one(self, X, atomic_nums, constraints=None): """ Computes gradients for that of a single structure. Parameters ---------- X: np.array numpy array of shape (a, 3) where a <= max_atoms and dtype is float-like atomic_nums: np.array numpy array of shape (a,) where a is the same as that of X. constraints: np.array numpy array of indices of X used for constraining a subset of the atoms of the molecule. Returns ------- np.array derivatives of the same shape and type as input parameter X. """ num_atoms = atomic_nums.shape[0] X = X.reshape((num_atoms, 3)) A = atomic_nums.reshape((atomic_nums.shape[0], 1)) Z = np.zeros((self.max_atoms, 4)) Z[:X.shape[0], 1:X.shape[1] + 1] = X Z[:A.shape[0], :A.shape[1]] = A X = Z inp = np.array(X).reshape((1, self.max_atoms, 4)) dd = dc.data.NumpyDataset(inp, np.array([1]), np.array([1])) res = self.compute_grad(dd)[0][0][0] res = res[:num_atoms, 1:] if constraints is not None: for idx in constraints: res[idx][0] = 0 res[idx][1] = 0 res[idx][2] = 0 return res.reshape((num_atoms * 3,)) def minimize_structure(self, X, atomic_nums, constraints=None): """ Minimizes a structure, as defined by a set of coordinates and their atomic numbers. Parameters ---------- X: np.array numpy array of shape (a, 3) where a <= max_atoms and dtype is float-like atomic_nums: np.array numpy array of shape (a,) where a is the same as that of X. Returns ------- np.array minimized coordinates of the same shape and type as input parameter X. """ num_atoms = atomic_nums.shape[0] res = scipy.optimize.minimize( self.pred_one, X, args=(atomic_nums, constraints), jac=self.grad_one, method="BFGS", tol=1e-6, options={'disp': True}) return res.x.reshape((num_atoms, 3)) def build_graph(self): self.atom_numbers = Feature(shape=(None, self.max_atoms), dtype=tf.int32) self.atom_flags = Feature(shape=(None, self.max_atoms, self.max_atoms)) self.atom_feats = Feature(shape=(None, self.max_atoms, self.n_feat)) previous_layer = self.atom_feats self.atom_feats = Feature(shape=(None, self.max_atoms, 4)) previous_layer = ANIFeat( in_layers=self.atom_feats, max_atoms=self.max_atoms) self.featurized = previous_layer Hiddens = [] for n_hidden in self.layer_structures: Loading Loading @@ -188,5 +360,65 @@ class ANIRegression(TensorGraph): feed_dict[self.atom_flags] = np.stack([flags]*self.max_atoms, axis=2)*\ np.stack([flags]*self.max_atoms, axis=1) feed_dict[self.atom_numbers] = np.array(X_b[:, :, 0], dtype=int) feed_dict[self.atom_feats] = np.array(X_b[:, :, 1:], dtype=float) feed_dict[self.atom_feats] = np.array(X_b[:, :, :], dtype=float) yield feed_dict def save_numpy(self): """ Save to a portable numpy file. Note that this relies on the names to be consistent across different versions. The file is saved as save_pickle.npz under the model_dir. """ path = os.path.join(self.model_dir, "save_pickle.npz") with self._get_tf("Graph").as_default(): all_vars = tf.trainable_variables() all_vals = self.session.run(all_vars) save_dict = {} for idx, _ in enumerate(all_vars): save_dict[all_vars[idx].name] = all_vals[idx] save_dict["_kwargs"] = np.array( [json.dumps(self._kwargs)], dtype=np.string_) np.savez(path, **save_dict) @classmethod def load_numpy(cls, model_dir): """ Load from a portable numpy file. Parameters ---------- model_dir: str Location of the model directory. """ path = os.path.join(model_dir, "save_pickle.npz") npo = np.load(path) json_blob = npo["_kwargs"][0].decode('UTF-8') kwargs = json.loads(json_blob) obj = cls(**kwargs) obj.build() all_ops = [] g = obj._get_tf("Graph") with g.as_default(): all_vars = tf.trainable_variables() for k in npo.keys(): if k == "_kwargs": continue val = npo[k] tensor = g.get_tensor_by_name(k) all_ops.append(tf.assign(tensor, val)) obj.session.run(all_ops) return obj deepchem/models/tensorgraph/models/test_symmetry_functions.py 0 → 100644 +159 −0 Original line number Diff line number Diff line import os import unittest import tempfile import scipy import numpy as np from deepchem.models import ANIRegression import deepchem as dc class TestANIRegression(unittest.TestCase): def setUp(self): max_atoms = 4 X = np.array([[ [1, 5.0, 3.2, 1.1], [6, 1.0, 3.4, -1.1], [1, 2.3, 3.4, 2.2], [0, 0, 0, 0], ], [ [8, 2.0, -1.4, -1.1], [7, 6.3, 2.4, 3.2], [0, 0, 0, 0], [0, 0, 0, 0], ]]) y = np.array([2.0, 1.1]) layer_structures = [128, 128, 64] atom_number_cases = [1, 6, 7, 8] self.model_dir = tempfile.mkdtemp() self.kwargs = { "n_tasks": 1, "max_atoms": max_atoms, "layer_structures": layer_structures, "atom_number_cases": atom_number_cases, "batch_size": 2, "learning_rate": 0.001, "use_queue": False, "mode": "regression", "model_dir": self.model_dir } model = ANIRegression(**self.kwargs) train_dataset = dc.data.NumpyDataset(X, y, n_tasks=1) model.fit(train_dataset, nb_epoch=2, checkpoint_interval=100) self.model = model def test_gradients(self): new_x = np.array([ -2.0, 1.2, 2.1, 1.3, -6.4, 3.1, -2.5, 2.4, 5.6, ]) new_atomic_nums = np.array([1, 1, 6]) delta = 1e-2 # use central difference since forward difference has a pretty high # approximation error grad_approx = [] for idx in range(new_x.shape[0]): d_new_x_plus = np.array(new_x) d_new_x_plus[idx] += delta d_new_x_minus = np.array(new_x) d_new_x_minus[idx] -= delta dydx = (self.model.pred_one(d_new_x_plus, new_atomic_nums) - self.model.pred_one(d_new_x_minus, new_atomic_nums)) / (2 * delta) grad_approx.append(dydx[0]) grad_approx = np.array(grad_approx) grad_exact = self.model.grad_one(new_x, new_atomic_nums) np.testing.assert_array_almost_equal(grad_approx, grad_exact, decimal=3) grad_exact_constrained = self.model.grad_one( new_x, new_atomic_nums, constraints=[0, 2]) assert grad_exact_constrained[0] == 0 assert grad_exact_constrained[1] == 0 assert grad_exact_constrained[2] == 0 assert grad_exact_constrained[3] == grad_exact[3] assert grad_exact_constrained[4] == grad_exact[4] assert grad_exact_constrained[5] == grad_exact[5] assert grad_exact_constrained[6] == 0 assert grad_exact_constrained[7] == 0 assert grad_exact_constrained[8] == 0 min_coords = self.model.minimize_structure( new_x, new_atomic_nums, constraints=[0, 2]) assert min_coords[0][0] == new_x[0] assert min_coords[0][1] == new_x[1] assert min_coords[0][2] == new_x[2] # assert min_coords[1][0] != new_x[3] # assert min_coords[1][1] != new_x[4] # assert min_coords[1][2] != new_x[5] assert min_coords[2][0] == new_x[6] assert min_coords[2][1] == new_x[7] assert min_coords[2][2] == new_x[8] def test_numpy_save_load(self): self.model.save_numpy() restored_model = ANIRegression.load_numpy(self.model_dir) new_x = np.array([ -2.0, 1.2, 2.1, 1.3, -6.4, 3.1, -2.5, 2.4, 5.6, ]) new_atomic_nums = np.array([1, 1, 6]) expected = self.model.pred_one(new_x, new_atomic_nums) predicted = restored_model.pred_one(new_x, new_atomic_nums) assert self.model.n_tasks == restored_model.n_tasks assert self.model.max_atoms == restored_model.max_atoms assert self.model.layer_structures == restored_model.layer_structures assert self.model.atom_number_cases == restored_model.atom_number_cases assert self.model.batch_size == restored_model.batch_size assert self.model.learning_rate == restored_model.learning_rate assert self.model.use_queue == restored_model.use_queue assert expected == predicted if __name__ == '__main__': unittest.main() deepchem/models/tests/test_overfit.py +0 −19 Original line number Diff line number Diff line Loading @@ -1005,23 +1005,14 @@ class TestOverfit(test_util.TensorFlowTestCase): transformers = [ dc.trans.NormalizationTransformer(transform_y=True, dataset=dataset), dc.trans.ANITransformer( max_atoms=13, atom_cases=[1, 6, 7, 8], radial_cutoff=8., angular_cutoff=5., radial_length=8, angular_length=4) ] for transformer in transformers: dataset = transformer.transform(dataset) n_feat = transformers[-1].get_num_feats() - 1 model = dc.models.ANIRegression( n_tasks, 13, n_feat, atom_number_cases=[1, 6, 7, 8], batch_size=batch_size, learning_rate=0.001, Loading Loading @@ -1054,24 +1045,14 @@ class TestOverfit(test_util.TensorFlowTestCase): transformers = [ dc.trans.NormalizationTransformer(transform_y=True, dataset=dataset), dc.trans.ANITransformer( max_atoms=13, atom_cases=[1, 6, 7, 8], atomic_number_differentiated=False, radial_cutoff=8., angular_cutoff=5., radial_length=8, angular_length=4) ] for transformer in transformers: dataset = transformer.transform(dataset) n_feat = transformers[-1].get_num_feats() - 1 model = dc.models.ANIRegression( n_tasks, 13, n_feat, atom_number_cases=[1, 6, 7, 8], batch_size=batch_size, learning_rate=0.001, Loading deepchem/splits/splitters.py +2 −1 Original line number Diff line number Diff line Loading @@ -195,7 +195,7 @@ class Splitter(object): class RandomGroupSplitter(Splitter): def __init__(self, groups): def __init__(self, groups, *args, **kwargs): """ A splitter class that splits on groupings. An example use case is when there are multiple conformations of the same molecule that share the same topology. Loading @@ -221,6 +221,7 @@ class RandomGroupSplitter(Splitter): """ self.groups = groups super(RandomGroupSplitter, self).__init__(*args, **kwargs) def split(self, dataset, Loading Loading
deepchem/models/tensorgraph/layers.py +179 −0 Original line number Diff line number Diff line Loading @@ -3248,6 +3248,185 @@ class BetaShare(Layer): self.out_tensor, self.betas = tensor class ANIFeat(Layer): """Performs transform from 3D coordinates to ANI symmetry functions """ def __init__(self, in_layers, max_atoms=23, radial_cutoff=4.6, angular_cutoff=3.1, radial_length=32, angular_length=8, atom_cases=[1, 6, 7, 8, 16], atomic_number_differentiated=True, coordinates_in_bohr=True, **kwargs): """ Only X can be transformed """ self.max_atoms = max_atoms self.radial_cutoff = radial_cutoff self.angular_cutoff = angular_cutoff self.radial_length = radial_length self.angular_length = angular_length self.atom_cases = atom_cases self.atomic_number_differentiated = atomic_number_differentiated self.coordinates_in_bohr = coordinates_in_bohr super(ANIFeat, self).__init__(in_layers, **kwargs) def create_tensor(self, in_layers=None, set_tensors=True, **kwargs): """ In layers should be of shape dtype tf.float32, (None, self.max_atoms, 4) """ inputs = self._get_input_tensors(in_layers)[0] atom_numbers = tf.cast(inputs[:, :, 0], tf.int32) flags = tf.sign(atom_numbers) flags = tf.to_float(tf.expand_dims(flags, 1) * tf.expand_dims(flags, 2)) coordinates = inputs[:, :, 1:] if self.coordinates_in_bohr: coordinates = coordinates * 0.52917721092 d = self.distance_matrix(coordinates, flags) d_radial_cutoff = self.distance_cutoff(d, self.radial_cutoff, flags) d_angular_cutoff = self.distance_cutoff(d, self.angular_cutoff, flags) radial_sym = self.radial_symmetry(d_radial_cutoff, d, atom_numbers) angular_sym = self.angular_symmetry(d_angular_cutoff, d, atom_numbers, coordinates) out_tensor = tf.concat( [tf.to_float(tf.expand_dims(atom_numbers, 2)), radial_sym, angular_sym], axis=2) if set_tensors: self.out_tensor = out_tensor return out_tensor def distance_matrix(self, coordinates, flags): """ Generate distance matrix """ # (TODO YTZ:) faster, less memory intensive way # r = tf.reduce_sum(tf.square(coordinates), 2) # r = tf.expand_dims(r, -1) # inner = 2*tf.matmul(coordinates, tf.transpose(coordinates, perm=[0,2,1])) # # inner = 2*tf.matmul(coordinates, coordinates, transpose_b=True) # d = r - inner + tf.transpose(r, perm=[0,2,1]) # d = tf.nn.relu(d) # fix numerical instabilities about diagonal # d = tf.sqrt(d) # does this have negative elements? may be unstable for diagonals max_atoms = self.max_atoms tensor1 = tf.stack([coordinates] * max_atoms, axis=1) tensor2 = tf.stack([coordinates] * max_atoms, axis=2) # Calculate pairwise distance d = tf.sqrt( tf.reduce_sum(tf.squared_difference(tensor1, tensor2), axis=3) + 1e-7) d = d * flags return d def distance_cutoff(self, d, cutoff, flags): """ Generate distance matrix with trainable cutoff """ # Cutoff with threshold Rc d_flag = flags * tf.sign(cutoff - d) d_flag = tf.nn.relu(d_flag) d_flag = d_flag * tf.expand_dims((1 - tf.eye(self.max_atoms)), 0) d = 0.5 * (tf.cos(np.pi * d / cutoff) + 1) return d * d_flag # return d def radial_symmetry(self, d_cutoff, d, atom_numbers): """ Radial Symmetry Function """ embedding = tf.eye(np.max(self.atom_cases) + 1) atom_numbers_embedded = tf.nn.embedding_lookup(embedding, atom_numbers) Rs = np.linspace(0., self.radial_cutoff, self.radial_length) ita = np.ones_like(Rs) * 3 / (Rs[1] - Rs[0])**2 Rs = tf.to_float(np.reshape(Rs, (1, 1, 1, -1))) ita = tf.to_float(np.reshape(ita, (1, 1, 1, -1))) length = ita.get_shape().as_list()[-1] d_cutoff = tf.stack([d_cutoff] * length, axis=3) d = tf.stack([d] * length, axis=3) out = tf.exp(-ita * tf.square(d - Rs)) * d_cutoff if self.atomic_number_differentiated: out_tensors = [] for atom_type in self.atom_cases: selected_atoms = tf.expand_dims( tf.expand_dims(atom_numbers_embedded[:, :, atom_type], axis=1), axis=3) out_tensors.append(tf.reduce_sum(out * selected_atoms, axis=2)) return tf.concat(out_tensors, axis=2) else: return tf.reduce_sum(out, axis=2) def angular_symmetry(self, d_cutoff, d, atom_numbers, coordinates): """ Angular Symmetry Function """ max_atoms = self.max_atoms embedding = tf.eye(np.max(self.atom_cases) + 1) atom_numbers_embedded = tf.nn.embedding_lookup(embedding, atom_numbers) Rs = np.linspace(0., self.angular_cutoff, self.angular_length) ita = 3 / (Rs[1] - Rs[0])**2 thetas = np.linspace(0., np.pi, self.angular_length) zeta = float(self.angular_length**2) ita, zeta, Rs, thetas = np.meshgrid(ita, zeta, Rs, thetas) zeta = tf.to_float(np.reshape(zeta, (1, 1, 1, 1, -1))) ita = tf.to_float(np.reshape(ita, (1, 1, 1, 1, -1))) Rs = tf.to_float(np.reshape(Rs, (1, 1, 1, 1, -1))) thetas = tf.to_float(np.reshape(thetas, (1, 1, 1, 1, -1))) length = zeta.get_shape().as_list()[-1] # tf.stack issues again... vector_distances = tf.stack([coordinates] * max_atoms, 1) - tf.stack( [coordinates] * max_atoms, 2) R_ij = tf.stack([d] * max_atoms, axis=3) R_ik = tf.stack([d] * max_atoms, axis=2) f_R_ij = tf.stack([d_cutoff] * max_atoms, axis=3) f_R_ik = tf.stack([d_cutoff] * max_atoms, axis=2) # Define angle theta = arccos(R_ij(Vector) dot R_ik(Vector)/R_ij(distance)/R_ik(distance)) vector_mul = tf.reduce_sum(tf.stack([vector_distances]*max_atoms, axis=3) * \ tf.stack([vector_distances]*max_atoms, axis=2), axis=4) vector_mul = vector_mul * tf.sign(f_R_ij) * tf.sign(f_R_ik) theta = tf.acos(tf.div(vector_mul, R_ij * R_ik + 1e-5)) R_ij = tf.stack([R_ij] * length, axis=4) R_ik = tf.stack([R_ik] * length, axis=4) f_R_ij = tf.stack([f_R_ij] * length, axis=4) f_R_ik = tf.stack([f_R_ik] * length, axis=4) theta = tf.stack([theta] * length, axis=4) out_tensor = tf.pow((1. + tf.cos(theta - thetas))/2., zeta) * \ tf.exp(-ita * tf.square((R_ij + R_ik)/2. - Rs)) * f_R_ij * f_R_ik * 2 if self.atomic_number_differentiated: out_tensors = [] for id_j, atom_type_j in enumerate(self.atom_cases): for atom_type_k in self.atom_cases[id_j:]: selected_atoms = tf.stack([atom_numbers_embedded[:, :, atom_type_j]] * max_atoms, axis=2) * \ tf.stack([atom_numbers_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))) return tf.concat(out_tensors, axis=2) else: return tf.reduce_sum(out_tensor, axis=(2, 3)) def get_num_feats(self): n_feat = self.outputs.get_shape().as_list()[-1] return n_feat class PassThroughLayer(Layer): """ Layer which takes a tensor from in_tensor[0].out_tensors at an index Loading
deepchem/models/tensorgraph/models/symmetry_function_regression.py +238 −6 Original line number Diff line number Diff line Loading @@ -4,11 +4,18 @@ Created on Thu Jul 6 20:31:47 2017 @author: zqwu @contributors: ytz """ import os import numpy as np import json import scipy.optimize import tensorflow as tf from deepchem.models.tensorgraph.layers import Dense, Concat, WeightedError, Stack import deepchem as dc from deepchem.models.tensorgraph.layers import Dense, Concat, WeightedError, Stack, Layer, ANIFeat from deepchem.models.tensorgraph.layers import L2Loss, Label, Weights, Feature from deepchem.models.tensorgraph.tensor_graph import TensorGraph from deepchem.models.tensorgraph.graph_layers import DTNNEmbedding Loading Loading @@ -39,7 +46,9 @@ class BPSymmetryFunctionRegression(TensorGraph): self.max_atoms = max_atoms self.n_feat = n_feat self.layer_structures = layer_structures super(BPSymmetryFunctionRegression, self).__init__(**kwargs) self.build_graph() def build_graph(self): Loading Loading @@ -106,7 +115,6 @@ class ANIRegression(TensorGraph): def __init__(self, n_tasks, max_atoms, n_feat, layer_structures=[128, 64], atom_number_cases=[1, 6, 7, 8, 16], **kwargs): Loading @@ -122,17 +130,181 @@ class ANIRegression(TensorGraph): """ self.n_tasks = n_tasks self.max_atoms = max_atoms self.n_feat = n_feat self.layer_structures = layer_structures self.atom_number_cases = atom_number_cases super(ANIRegression, self).__init__(**kwargs) # (ytz): this is really dirty but needed for restoring models self._kwargs = { "n_tasks": n_tasks, "max_atoms": max_atoms, "layer_structures": layer_structures, "atom_number_cases": atom_number_cases } self._kwargs.update(kwargs) self.build_graph() self.grad = None def save(self): self.grad = None # recompute grad on restore super(ANIRegression, self).save() def build_grad(self): self.grad = tf.gradients(self.outputs, self.atom_feats) def compute_grad(self, dataset, upper_lim=1): """ Computes a batched gradients given an input dataset. Parameters ---------- dataset: dc.Dataset dataset-like object whose X values will be used to compute gradients from upper_lim: int subset of dataset used. Returns ------- np.array Gradients of the input of shape (max_atoms, 4). Note that it is up to the end user to slice this matrix into the correct shape, since it's very likely the derivatives with respect to the atomic numbers are zero. """ with self._get_tf("Graph").as_default(): if not self.built: self.build() if not self.grad: self.build_grad() feed_dict = dict() X = dataset.X flags = np.sign(np.array(X[:upper_lim, :, 0])) feed_dict[self.atom_flags] = np.stack([flags]*self.max_atoms, axis=2)*\ np.stack([flags]*self.max_atoms, axis=1) feed_dict[self.atom_numbers] = np.array(X[:upper_lim, :, 0], dtype=int) feed_dict[self.atom_feats] = np.array(X[:upper_lim, :, :], dtype=float) return self.session.run([self.grad], feed_dict=feed_dict) def pred_one(self, X, atomic_nums, constraints=None): """ Makes an energy prediction for a set of atomic coordinates. Parameters ---------- X: np.array numpy array of shape (a, 3) where a <= max_atoms and dtype is float-like atomic_nums: np.array numpy array of shape (a,) where a is the same as that of X. constraints: unused This parameter is mainly for compatibility purposes for scipy optimize Returns ------- float Predicted energy. Note that the meaning of the returned value is dependent on the training y-values both in semantics (relative vs absolute) and units (kcal/mol vs Hartrees) """ num_atoms = atomic_nums.shape[0] X = X.reshape((num_atoms, 3)) A = atomic_nums.reshape((atomic_nums.shape[0], 1)) Z = np.zeros((self.max_atoms, 4)) Z[:X.shape[0], 1:X.shape[1] + 1] = X Z[:A.shape[0], :A.shape[1]] = A X = Z dd = dc.data.NumpyDataset( np.array(X).reshape((1, self.max_atoms, 4)), np.array(0), np.array(1)) return self.predict(dd)[0] def grad_one(self, X, atomic_nums, constraints=None): """ Computes gradients for that of a single structure. Parameters ---------- X: np.array numpy array of shape (a, 3) where a <= max_atoms and dtype is float-like atomic_nums: np.array numpy array of shape (a,) where a is the same as that of X. constraints: np.array numpy array of indices of X used for constraining a subset of the atoms of the molecule. Returns ------- np.array derivatives of the same shape and type as input parameter X. """ num_atoms = atomic_nums.shape[0] X = X.reshape((num_atoms, 3)) A = atomic_nums.reshape((atomic_nums.shape[0], 1)) Z = np.zeros((self.max_atoms, 4)) Z[:X.shape[0], 1:X.shape[1] + 1] = X Z[:A.shape[0], :A.shape[1]] = A X = Z inp = np.array(X).reshape((1, self.max_atoms, 4)) dd = dc.data.NumpyDataset(inp, np.array([1]), np.array([1])) res = self.compute_grad(dd)[0][0][0] res = res[:num_atoms, 1:] if constraints is not None: for idx in constraints: res[idx][0] = 0 res[idx][1] = 0 res[idx][2] = 0 return res.reshape((num_atoms * 3,)) def minimize_structure(self, X, atomic_nums, constraints=None): """ Minimizes a structure, as defined by a set of coordinates and their atomic numbers. Parameters ---------- X: np.array numpy array of shape (a, 3) where a <= max_atoms and dtype is float-like atomic_nums: np.array numpy array of shape (a,) where a is the same as that of X. Returns ------- np.array minimized coordinates of the same shape and type as input parameter X. """ num_atoms = atomic_nums.shape[0] res = scipy.optimize.minimize( self.pred_one, X, args=(atomic_nums, constraints), jac=self.grad_one, method="BFGS", tol=1e-6, options={'disp': True}) return res.x.reshape((num_atoms, 3)) def build_graph(self): self.atom_numbers = Feature(shape=(None, self.max_atoms), dtype=tf.int32) self.atom_flags = Feature(shape=(None, self.max_atoms, self.max_atoms)) self.atom_feats = Feature(shape=(None, self.max_atoms, self.n_feat)) previous_layer = self.atom_feats self.atom_feats = Feature(shape=(None, self.max_atoms, 4)) previous_layer = ANIFeat( in_layers=self.atom_feats, max_atoms=self.max_atoms) self.featurized = previous_layer Hiddens = [] for n_hidden in self.layer_structures: Loading Loading @@ -188,5 +360,65 @@ class ANIRegression(TensorGraph): feed_dict[self.atom_flags] = np.stack([flags]*self.max_atoms, axis=2)*\ np.stack([flags]*self.max_atoms, axis=1) feed_dict[self.atom_numbers] = np.array(X_b[:, :, 0], dtype=int) feed_dict[self.atom_feats] = np.array(X_b[:, :, 1:], dtype=float) feed_dict[self.atom_feats] = np.array(X_b[:, :, :], dtype=float) yield feed_dict def save_numpy(self): """ Save to a portable numpy file. Note that this relies on the names to be consistent across different versions. The file is saved as save_pickle.npz under the model_dir. """ path = os.path.join(self.model_dir, "save_pickle.npz") with self._get_tf("Graph").as_default(): all_vars = tf.trainable_variables() all_vals = self.session.run(all_vars) save_dict = {} for idx, _ in enumerate(all_vars): save_dict[all_vars[idx].name] = all_vals[idx] save_dict["_kwargs"] = np.array( [json.dumps(self._kwargs)], dtype=np.string_) np.savez(path, **save_dict) @classmethod def load_numpy(cls, model_dir): """ Load from a portable numpy file. Parameters ---------- model_dir: str Location of the model directory. """ path = os.path.join(model_dir, "save_pickle.npz") npo = np.load(path) json_blob = npo["_kwargs"][0].decode('UTF-8') kwargs = json.loads(json_blob) obj = cls(**kwargs) obj.build() all_ops = [] g = obj._get_tf("Graph") with g.as_default(): all_vars = tf.trainable_variables() for k in npo.keys(): if k == "_kwargs": continue val = npo[k] tensor = g.get_tensor_by_name(k) all_ops.append(tf.assign(tensor, val)) obj.session.run(all_ops) return obj
deepchem/models/tensorgraph/models/test_symmetry_functions.py 0 → 100644 +159 −0 Original line number Diff line number Diff line import os import unittest import tempfile import scipy import numpy as np from deepchem.models import ANIRegression import deepchem as dc class TestANIRegression(unittest.TestCase): def setUp(self): max_atoms = 4 X = np.array([[ [1, 5.0, 3.2, 1.1], [6, 1.0, 3.4, -1.1], [1, 2.3, 3.4, 2.2], [0, 0, 0, 0], ], [ [8, 2.0, -1.4, -1.1], [7, 6.3, 2.4, 3.2], [0, 0, 0, 0], [0, 0, 0, 0], ]]) y = np.array([2.0, 1.1]) layer_structures = [128, 128, 64] atom_number_cases = [1, 6, 7, 8] self.model_dir = tempfile.mkdtemp() self.kwargs = { "n_tasks": 1, "max_atoms": max_atoms, "layer_structures": layer_structures, "atom_number_cases": atom_number_cases, "batch_size": 2, "learning_rate": 0.001, "use_queue": False, "mode": "regression", "model_dir": self.model_dir } model = ANIRegression(**self.kwargs) train_dataset = dc.data.NumpyDataset(X, y, n_tasks=1) model.fit(train_dataset, nb_epoch=2, checkpoint_interval=100) self.model = model def test_gradients(self): new_x = np.array([ -2.0, 1.2, 2.1, 1.3, -6.4, 3.1, -2.5, 2.4, 5.6, ]) new_atomic_nums = np.array([1, 1, 6]) delta = 1e-2 # use central difference since forward difference has a pretty high # approximation error grad_approx = [] for idx in range(new_x.shape[0]): d_new_x_plus = np.array(new_x) d_new_x_plus[idx] += delta d_new_x_minus = np.array(new_x) d_new_x_minus[idx] -= delta dydx = (self.model.pred_one(d_new_x_plus, new_atomic_nums) - self.model.pred_one(d_new_x_minus, new_atomic_nums)) / (2 * delta) grad_approx.append(dydx[0]) grad_approx = np.array(grad_approx) grad_exact = self.model.grad_one(new_x, new_atomic_nums) np.testing.assert_array_almost_equal(grad_approx, grad_exact, decimal=3) grad_exact_constrained = self.model.grad_one( new_x, new_atomic_nums, constraints=[0, 2]) assert grad_exact_constrained[0] == 0 assert grad_exact_constrained[1] == 0 assert grad_exact_constrained[2] == 0 assert grad_exact_constrained[3] == grad_exact[3] assert grad_exact_constrained[4] == grad_exact[4] assert grad_exact_constrained[5] == grad_exact[5] assert grad_exact_constrained[6] == 0 assert grad_exact_constrained[7] == 0 assert grad_exact_constrained[8] == 0 min_coords = self.model.minimize_structure( new_x, new_atomic_nums, constraints=[0, 2]) assert min_coords[0][0] == new_x[0] assert min_coords[0][1] == new_x[1] assert min_coords[0][2] == new_x[2] # assert min_coords[1][0] != new_x[3] # assert min_coords[1][1] != new_x[4] # assert min_coords[1][2] != new_x[5] assert min_coords[2][0] == new_x[6] assert min_coords[2][1] == new_x[7] assert min_coords[2][2] == new_x[8] def test_numpy_save_load(self): self.model.save_numpy() restored_model = ANIRegression.load_numpy(self.model_dir) new_x = np.array([ -2.0, 1.2, 2.1, 1.3, -6.4, 3.1, -2.5, 2.4, 5.6, ]) new_atomic_nums = np.array([1, 1, 6]) expected = self.model.pred_one(new_x, new_atomic_nums) predicted = restored_model.pred_one(new_x, new_atomic_nums) assert self.model.n_tasks == restored_model.n_tasks assert self.model.max_atoms == restored_model.max_atoms assert self.model.layer_structures == restored_model.layer_structures assert self.model.atom_number_cases == restored_model.atom_number_cases assert self.model.batch_size == restored_model.batch_size assert self.model.learning_rate == restored_model.learning_rate assert self.model.use_queue == restored_model.use_queue assert expected == predicted if __name__ == '__main__': unittest.main()
deepchem/models/tests/test_overfit.py +0 −19 Original line number Diff line number Diff line Loading @@ -1005,23 +1005,14 @@ class TestOverfit(test_util.TensorFlowTestCase): transformers = [ dc.trans.NormalizationTransformer(transform_y=True, dataset=dataset), dc.trans.ANITransformer( max_atoms=13, atom_cases=[1, 6, 7, 8], radial_cutoff=8., angular_cutoff=5., radial_length=8, angular_length=4) ] for transformer in transformers: dataset = transformer.transform(dataset) n_feat = transformers[-1].get_num_feats() - 1 model = dc.models.ANIRegression( n_tasks, 13, n_feat, atom_number_cases=[1, 6, 7, 8], batch_size=batch_size, learning_rate=0.001, Loading Loading @@ -1054,24 +1045,14 @@ class TestOverfit(test_util.TensorFlowTestCase): transformers = [ dc.trans.NormalizationTransformer(transform_y=True, dataset=dataset), dc.trans.ANITransformer( max_atoms=13, atom_cases=[1, 6, 7, 8], atomic_number_differentiated=False, radial_cutoff=8., angular_cutoff=5., radial_length=8, angular_length=4) ] for transformer in transformers: dataset = transformer.transform(dataset) n_feat = transformers[-1].get_num_feats() - 1 model = dc.models.ANIRegression( n_tasks, 13, n_feat, atom_number_cases=[1, 6, 7, 8], batch_size=batch_size, learning_rate=0.001, Loading
deepchem/splits/splitters.py +2 −1 Original line number Diff line number Diff line Loading @@ -195,7 +195,7 @@ class Splitter(object): class RandomGroupSplitter(Splitter): def __init__(self, groups): def __init__(self, groups, *args, **kwargs): """ A splitter class that splits on groupings. An example use case is when there are multiple conformations of the same molecule that share the same topology. Loading @@ -221,6 +221,7 @@ class RandomGroupSplitter(Splitter): """ self.groups = groups super(RandomGroupSplitter, self).__init__(*args, **kwargs) def split(self, dataset, Loading
