Commit 1e810c94 authored by Bharath Ramsundar's avatar Bharath Ramsundar Committed by Nathan Frey
Browse files

First cut at atomic conv fix

parent 8f688953

deepchem/models/atomic_conv.py

+17 −10
Original line number Diff line number Diff line 
__author__ = "Joseph Gomes"

__copyright__ = "Copyright 2017, Stanford University"

__license__ = "MIT"



import sys



from deepchem.models import KerasModel

from deepchem.models.layers import AtomicConvolution

from deepchem.models.losses import L2Loss

from tensorflow.keras.layers import Input, Layer

from tensorflow.keras.layers import Input, Layer, Dense, Flatten, Concatenate



import numpy as np

import tensorflow as tf
@@ -224,19 +220,29 @@ class AtomicConvModel(KerasModel): 
    self._frag1_conv = AtomicConvolution(

        atom_types=self.atom_types, radial_params=rp,

        boxsize=None)([frag1_X, frag1_nbrs, frag1_nbrs_z])

    flattened1 = Flatten()(self._frag1_conv)



    self._frag2_conv = AtomicConvolution(

        atom_types=self.atom_types, radial_params=rp,

        boxsize=None)([frag2_X, frag2_nbrs, frag2_nbrs_z])

    flattened2 = Flatten()(self._frag2_conv)



    self._complex_conv = AtomicConvolution(

        atom_types=self.atom_types, radial_params=rp,

        boxsize=None)([complex_X, complex_nbrs, complex_nbrs_z])

    flattened3 = Flatten()(self._complex_conv)



    concat = Concatenate()([flattened1, flattened2, flattened3])



    score = AtomicConvScore(self.atom_types, layer_sizes)([

        self._frag1_conv, self._frag2_conv, self._complex_conv, frag1_z,

        frag2_z, complex_z

    ])

    #score = AtomicConvScore(self.atom_types, layer_sizes)([

    #    self._frag1_conv, self._frag2_conv, self._complex_conv, frag1_z,

    #    frag2_z, complex_z

    #])

    layer = Dense(100)(concat)

    output = Dense(1)(layer)

    print("output")

    print(output)

    #loss = dc.models.losses.L2Loss()



    model = tf.keras.Model(

        inputs=[
@@ -244,7 +250,8 @@ class AtomicConvModel(KerasModel): 
            frag2_nbrs_z, frag2_z, complex_X, complex_nbrs, complex_nbrs_z,

            complex_z

        ],

        outputs=score)

        #outputs=score)

        outputs=output)

    super(AtomicConvModel, self).__init__(

        model, L2Loss(), batch_size=batch_size, **kwargs)

deepchem/models/tests/test_atomic_conv.py

+80 −74
Original line number Diff line number Diff line
@@ -16,10 +16,8 @@ from deepchem.data import NumpyDataset 
from deepchem.feat import ComplexNeighborListFragmentAtomicCoordinates





class TestAtomicConv(unittest.TestCase):



@pytest.mark.slow

  def test_atomic_conv(self):

def test_atomic_conv():

  """A simple test that initializes and fits an AtomicConvModel."""

  # For simplicity, let's assume both molecules have same number of

  # atoms.
@@ -63,77 +61,85 @@ class TestAtomicConv(unittest.TestCase): 
  features = np.asarray(features)

  labels = np.random.rand(batch_size)

  train = NumpyDataset(features, labels)

    atomic_convnet.fit(train, nb_epoch=300)

  #atomic_convnet.fit(train, nb_epoch=300)

  atomic_convnet.fit(train, nb_epoch=100)

  print("labels")

  print(labels)

  print("atomic_convnet.predict(train)")

  print(atomic_convnet.predict(train))

  assert np.allclose(labels, atomic_convnet.predict(train), atol=0.01)



  @pytest.mark.slow

  def test_atomic_conv_variable(self):

    """A simple test that initializes and fits an AtomicConvModel on variable input size."""

    # For simplicity, let's assume both molecules have same number of

    # atoms.

    frag1_num_atoms = 1000

    frag2_num_atoms = 1200

    complex_num_atoms = frag1_num_atoms + frag2_num_atoms

    batch_size = 1

    atomic_convnet = atomic_conv.AtomicConvModel(

        batch_size=batch_size,

        frag1_num_atoms=frag1_num_atoms,

        frag2_num_atoms=frag2_num_atoms,

        complex_num_atoms=complex_num_atoms)



    # Creates a set of dummy features that contain the coordinate and

    # neighbor-list features required by the AtomicConvModel.

    features = []

    frag1_coords = np.random.rand(frag1_num_atoms, 3)

    frag1_nbr_list = {i: [] for i in range(frag1_num_atoms)}

    frag1_z = np.random.randint(10, size=(frag1_num_atoms))

    frag2_coords = np.random.rand(frag2_num_atoms, 3)

    frag2_nbr_list = {i: [] for i in range(frag2_num_atoms)}

    frag2_z = np.random.randint(10, size=(frag2_num_atoms))

    system_coords = np.random.rand(complex_num_atoms, 3)

    system_nbr_list = {i: [] for i in range(complex_num_atoms)}

    system_z = np.random.randint(10, size=(complex_num_atoms))



    features.append(

        (frag1_coords, frag1_nbr_list, frag1_z, frag2_coords, frag2_nbr_list,

         frag2_z, system_coords, system_nbr_list, system_z))

    features = np.asarray(features)

    labels = np.zeros(batch_size)

    train = NumpyDataset(features, labels)

    atomic_convnet.fit(train, nb_epoch=1)



  @pytest.mark.slow

  def test_atomic_conv_with_feat(self):

    """A simple test for running an atomic convolution on featurized data."""

    dir_path = os.path.dirname(os.path.realpath(__file__))

    ligand_file = os.path.join(dir_path,

                               "../../feat/tests/data/3zso_ligand_hyd.pdb")

    protein_file = os.path.join(dir_path,

                                "../../feat/tests/data/3zso_protein.pdb")

    # Pulled from PDB files. For larger datasets with more PDBs, would use

    # max num atoms instead of exact.

    frag1_num_atoms = 44  # for ligand atoms

    frag2_num_atoms = 2336  # for protein atoms

    complex_num_atoms = 2380  # in total

    max_num_neighbors = 4

    # Cutoff in angstroms

    neighbor_cutoff = 4

    complex_featurizer = ComplexNeighborListFragmentAtomicCoordinates(

        frag1_num_atoms, frag2_num_atoms, complex_num_atoms, max_num_neighbors,

        neighbor_cutoff)

    # arbitrary label

    labels = np.array([0])

    features, _ = complex_featurizer.featurize([ligand_file], [protein_file])

    dataset = deepchem.data.DiskDataset.from_numpy(features, labels)



    batch_size = 1

    print("Constructing Atomic Conv model")

    atomic_convnet = atomic_conv.AtomicConvModel(

        batch_size=batch_size,

        frag1_num_atoms=frag1_num_atoms,

        frag2_num_atoms=frag2_num_atoms,

        complex_num_atoms=complex_num_atoms)



    print("About to call fit")

    # Run a fitting operation

    atomic_convnet.fit(dataset)
 
#class TestAtomicConv(unittest.TestCase):

#

#  @pytest.mark.slow

#  def test_atomic_conv_variable(self):

#    """A simple test that initializes and fits an AtomicConvModel on variable input size."""

#    # For simplicity, let's assume both molecules have same number of

#    # atoms.

#    frag1_num_atoms = 1000

#    frag2_num_atoms = 1200

#    complex_num_atoms = frag1_num_atoms + frag2_num_atoms

#    batch_size = 1

#    atomic_convnet = atomic_conv.AtomicConvModel(

#        batch_size=batch_size,

#        frag1_num_atoms=frag1_num_atoms,

#        frag2_num_atoms=frag2_num_atoms,

#        complex_num_atoms=complex_num_atoms)

#

#    # Creates a set of dummy features that contain the coordinate and

#    # neighbor-list features required by the AtomicConvModel.

#    features = []

#    frag1_coords = np.random.rand(frag1_num_atoms, 3)

#    frag1_nbr_list = {i: [] for i in range(frag1_num_atoms)}

#    frag1_z = np.random.randint(10, size=(frag1_num_atoms))

#    frag2_coords = np.random.rand(frag2_num_atoms, 3)

#    frag2_nbr_list = {i: [] for i in range(frag2_num_atoms)}

#    frag2_z = np.random.randint(10, size=(frag2_num_atoms))

#    system_coords = np.random.rand(complex_num_atoms, 3)

#    system_nbr_list = {i: [] for i in range(complex_num_atoms)}

#    system_z = np.random.randint(10, size=(complex_num_atoms))

#

#    features.append(

#        (frag1_coords, frag1_nbr_list, frag1_z, frag2_coords, frag2_nbr_list,

#         frag2_z, system_coords, system_nbr_list, system_z))

#    features = np.asarray(features)

#    labels = np.zeros(batch_size)

#    train = NumpyDataset(features, labels)

#    atomic_convnet.fit(train, nb_epoch=1)

#

#  @pytest.mark.slow

#  def test_atomic_conv_with_feat(self):

#    """A simple test for running an atomic convolution on featurized data."""

#    dir_path = os.path.dirname(os.path.realpath(__file__))

#    ligand_file = os.path.join(dir_path,

#                               "../../feat/tests/data/3zso_ligand_hyd.pdb")

#    protein_file = os.path.join(dir_path,

#                                "../../feat/tests/data/3zso_protein.pdb")

#    # Pulled from PDB files. For larger datasets with more PDBs, would use

#    # max num atoms instead of exact.

#    frag1_num_atoms = 44  # for ligand atoms

#    frag2_num_atoms = 2336  # for protein atoms

#    complex_num_atoms = 2380  # in total

#    max_num_neighbors = 4

#    # Cutoff in angstroms

#    neighbor_cutoff = 4

#    complex_featurizer = ComplexNeighborListFragmentAtomicCoordinates(

#        frag1_num_atoms, frag2_num_atoms, complex_num_atoms, max_num_neighbors,

#        neighbor_cutoff)

#    # arbitrary label

#    labels = np.array([0])

#    features, _ = complex_featurizer.featurize([ligand_file], [protein_file])

#    dataset = deepchem.data.DiskDataset.from_numpy(features, labels)

#

#    batch_size = 1

#    print("Constructing Atomic Conv model")

#    atomic_convnet = atomic_conv.AtomicConvModel(

#        batch_size=batch_size,

#        frag1_num_atoms=frag1_num_atoms,

#        frag2_num_atoms=frag2_num_atoms,

#        complex_num_atoms=complex_num_atoms)

#

#    print("About to call fit")

#    # Run a fitting operation

#    atomic_convnet.fit(dataset)

CRA Git | Maintained and supported by SUSTech CRA and CCSE