Commit ba2e0a4c authored by Yutong Zhao's avatar Yutong Zhao
Browse files

yapf

parent f5d04605

deepchem/models/tensorgraph/layers.py

+5 −8
Original line number Diff line number Diff line
@@ -3267,11 +3267,7 @@ class ANIFeat(Layer): 
                                        coordinates)



    out_tensor = tf.concat(

        [

            tf.to_float(tf.expand_dims(atom_numbers, 2)),

            radial_sym,

            angular_sym

        ],

        [tf.to_float(tf.expand_dims(atom_numbers, 2)), radial_sym, angular_sym],

        axis=2)



    if set_tensors:
@@ -3298,7 +3294,9 @@ class ANIFeat(Layer): 
    tensor2 = tf.stack([coordinates] * max_atoms, axis=2)



    # Calculate pairwise distance

    d = tf.sqrt(tf.nn.relu(tf.reduce_sum(tf.squared_difference(tensor1,tensor2), axis=3)))

    d = tf.sqrt(

        tf.nn.relu(

            tf.reduce_sum(tf.squared_difference(tensor1, tensor2), axis=3)))

    # Masking for valid atom index

    d = d * flags

    return d
@@ -3358,7 +3356,6 @@ class ANIFeat(Layer): 
    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)

deepchem/models/tensorgraph/models/symmetry_function_regression.py

+22 −23
Original line number Diff line number Diff line
@@ -109,6 +109,7 @@ class BPSymmetryFunctionRegression(TensorGraph): 
        feed_dict[self.atom_feats] = np.array(X_b[:, :, 1:], dtype=float)

        yield feed_dict





class ANIRegression(TensorGraph):



  def __init__(self,
@@ -133,7 +134,6 @@ class ANIRegression(TensorGraph): 
    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,
@@ -218,7 +218,8 @@ class ANIRegression(TensorGraph): 
    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))

    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):
@@ -301,8 +302,7 @@ class ANIRegression(TensorGraph): 
    self.atom_feats = Feature(shape=(None, self.max_atoms, 4))



    previous_layer = ANIFeat(

      in_layers=self.atom_feats,

      max_atoms=self.max_atoms)

        in_layers=self.atom_feats, max_atoms=self.max_atoms)



    self.featurized = previous_layer
@@ -378,11 +378,11 @@ class ANIRegression(TensorGraph): 
      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_)

      save_dict["_kwargs"] = np.array(

          [json.dumps(self._kwargs)], dtype=np.string_)



      np.savez(path, **save_dict)





  @classmethod

  def load_numpy(cls, model_dir):

    """
@@ -404,7 +404,6 @@ class ANIRegression(TensorGraph): 
    obj = cls(**kwargs)

    obj.build()





    all_ops = []



    g = obj._get_tf("Graph")

deepchem/models/tensorgraph/models/test_symmetry_functions.py

+43 −30
Original line number Diff line number Diff line
@@ -8,17 +8,14 @@ import numpy as np 
from deepchem.models import ANIRegression

import deepchem as dc





class TestANIRegression(unittest.TestCase):



  def setUp(self):



    max_atoms = 3



    X = np.array([

      [1, 5.0, 3.2, 1.1],

      [6, 1.0, 3.4, -1.1],

      [1, 2.3, 3.4, 2.2]

    ])

    X = np.array([[1, 5.0, 3.2, 1.1], [6, 1.0, 3.4, -1.1], [1, 2.3, 3.4, 2.2]])



    X = X.reshape((1, X.shape[0], X.shape[1]))
@@ -39,7 +36,8 @@ class TestANIRegression(unittest.TestCase): 
        "learning_rate": 0.001,

        "use_queue": False,

        "mode": "regression",

      "model_dir": self.model_dir}

        "model_dir": self.model_dir

    }



    model = ANIRegression(**self.kwargs)
@@ -49,13 +47,18 @@ class TestANIRegression(unittest.TestCase): 


    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,

        -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])
@@ -72,7 +75,8 @@ class TestANIRegression(unittest.TestCase): 
      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)

      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)
@@ -81,7 +85,8 @@ class TestANIRegression(unittest.TestCase): 


    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])

    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
@@ -95,7 +100,8 @@ class TestANIRegression(unittest.TestCase): 
    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])

    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]
@@ -115,9 +121,15 @@ class TestANIRegression(unittest.TestCase): 
    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,

        -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])
@@ -135,5 +147,6 @@ class TestANIRegression(unittest.TestCase): 


    assert expected == predicted





if __name__ == '__main__':

  unittest.main()

examples/roitberg/app.py

+4 −4
Original line number Diff line number Diff line
@@ -5,6 +5,7 @@ import flask 


webapp = Flask(__name__)





@webapp.route('/potential', methods=["POST"])

def potential():

  content = request.get_json(force=True)
@@ -16,6 +17,7 @@ def potential(): 
  result = webapp.model.pred_one(x0, a0)

  return flask.jsonify({'y': result.tolist()[0]}), 200





@webapp.route('/gradient', methods=["POST"])

def index():

  content = request.get_json(force=True)
@@ -31,6 +33,7 @@ def index(): 


  return flask.jsonify({'grad': res.tolist()}), 200





@webapp.route('/minimize', methods=["POST"])

def minimize():

  content = request.get_json(force=True)
@@ -44,13 +47,10 @@ def minimize(): 
    constraints = content['constraints']

    print('setting constraints')





  num_atoms = X.shape[0]

  x0 = X[:, 1:]

  a0 = X[:, :1]







  res = webapp.model.minimize_structure(x0, a0, constraints)

  res = res.reshape((num_atoms, 3))

  y = webapp.model.pred_one(res, a0).tolist()[0]

examples/roitberg/pyanitools.py

+9 −15
Original line number Diff line number Diff line
@@ -127,6 +127,7 @@ class anidataloader(object): 
  def cleanup(self):

    self.store.close()





if __name__ == "__main__":

  base_dir = os.environ["ROITBERG_ANI"]
@@ -134,19 +135,12 @@ if __name__ == "__main__": 
  # Start with a smaller dataset before continuing. Use all of them

  # for production

  hdf5files = [

      'ani_gdb_s01.h5',

      'ani_gdb_s02.h5',

      'ani_gdb_s03.h5',

      'ani_gdb_s04.h5',

      'ani_gdb_s05.h5',

      'ani_gdb_s06.h5',

      'ani_gdb_s07.h5',

      'ani_gdb_s08.h5'

      'ani_gdb_s01.h5', 'ani_gdb_s02.h5', 'ani_gdb_s03.h5', 'ani_gdb_s04.h5',

      'ani_gdb_s05.h5', 'ani_gdb_s06.h5', 'ani_gdb_s07.h5', 'ani_gdb_s08.h5'

  ]



  hdf5files = [os.path.join(base_dir, f) for f in hdf5files]





  for hdf5file in hdf5files:

    print("processing", hdf5file)

    adl = anidataloader(hdf5file)

CRA Git | Maintained and supported by SUSTech CRA and CCSE