👌 update for review comments

from deepchem.utils.typing import RDKitAtom, RDKitBond, RDKitMol

from deepchem.feat.graph_data import GraphData

from deepchem.feat.base_classes import MolecularFeaturizer

from deepchem.utils.graph_conv_utils import get_atom_type_one_hot, \

from deepchem.utils.molecule_feature_utils import get_atom_type_one_hot, \

  construct_hydrogen_bonding_info, get_atom_hydrogen_bonding_one_hot, \

  get_atom_is_in_aromatic_one_hot, get_atom_hybridization_one_hot, \

  get_atom_total_num_Hs_one_hot, get_atom_chirality_one_hot, get_atom_formal_charge, \

  - Chirality: A one-hot vector of the chirality, "R" or "S".

  - Formal charge: Integer electronic charge.

  - Partial charge: Calculated partial charge.

  - Ring sizes: A one-hot vector of the number of rings (3-8) that include this atom.

  - Ring sizes: A one-hot vector of the size (3-8) of rings that include this atom.

  - Hybridization: A one-hot vector of "sp", "sp2", "sp3".

  - Hydrogen bonding: A one-hot vector of whether this atom is a hydrogen bond donor or acceptor.

  - Aromatic: A one-hot vector of whether the atom belongs to an aromatic ring.

  - Stereo: A one-hot vector of the stereo configuration of a bond.

  If you want to know more details about features, please check the paper [1]_ and

  utilities in deepchem.utils.graph_conv_utils.py.

  utilities in deepchem.utils.molecule_feature_utils.py.

  Examples

  --------

  This class requires RDKit to be installed.

  """

  def __init__(self, add_self_loop: bool = False):

  def __init__(self, add_self_edges: bool = False):

    """

    Parameters

    ----------

    add_self_loop: bool, default False

    add_self_edges: bool, default False

      Whether to add self-connected edges or not. If you want to use DGL,

      you sometimes need to add explict self-connected edges.

    """

    self.add_self_loop = add_self_loop

    self.add_self_edges = add_self_edges

  def _featurize(self, mol: RDKitMol) -> GraphData:

    """Calculate molecule graph features from RDKit mol object.

    )

    # construct edge (bond) information

    src, dist, bond_features = [], [], []

    src, dest, bond_features = [], [], []

    for bond in mol.GetBonds():

      # add edge list considering a directed graph

      start, end = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()

      src += [start, end]

      dist += [end, start]

      dest += [end, start]

      bond_features += 2 * [_construct_bond_feature(bond)]

    if self.add_self_loop:

      num_atoms = mol.GetNumAtoms()

      src += [i for i in range(num_atoms)]

      dist += [i for i in range(num_atoms)]

      dest += [i for i in range(num_atoms)]

      # add dummy edge features

      bond_fea_length = len(bond_features[0])

      bond_features += num_atoms * [[0 for _ in range(bond_fea_length)]]

    return GraphData(

        node_features=atom_features,

        edge_index=np.array([src, dist], dtype=np.int),

        edge_index=np.array([src, dest], dtype=np.int),

        edge_features=np.array(bond_features, dtype=np.float))

  def test_featurizer_with_self_loop(self):

    smiles = ["C1=CC=CN=C1", "O=C(NCc1cc(OC)c(O)cc1)CCCC/C=C/C(C)C"]

    featurizer = MolGraphConvFeaturizer(add_self_loop=True)

    featurizer = MolGraphConvFeaturizer(add_self_edges=True)

    graph_feat = featurizer.featurize(smiles)

    assert len(graph_feat) == 2

      hidden_node_dim: int = 64,

      in_edge_dim: int = 41,

      num_conv: int = 3,

      predicator_hidden_feats: int = 128,

      predictor_hidden_feats: int = 128,

      n_tasks: int = 1,

      mode: str = 'regression',

      n_classes: int = 2,

      based on default setting of CGCNNFeaturizer.

    num_conv: int, default 3

      The number of convolutional layers.

    predicator_hidden_feats: int, default 128

    predictor_hidden_feats: int, default 128

      The size for hidden representations in the output MLP predictor.

    n_tasks: int, default 1

      The number of the output size.

            batch_norm=True) for _ in range(num_conv)

    ])

    self.pooling = dgl.mean_nodes

    self.fc = nn.Linear(hidden_node_dim, predicator_hidden_feats)

    self.fc = nn.Linear(hidden_node_dim, predictor_hidden_feats)

    if self.mode == 'regression':

      self.out = nn.Linear(predicator_hidden_feats, n_tasks)

      self.out = nn.Linear(predictor_hidden_feats, n_tasks)

    else:

      self.out = nn.Linear(predicator_hidden_feats, n_tasks * n_classes)

      self.out = nn.Linear(predictor_hidden_feats, n_tasks * n_classes)

  def forward(self, dgl_graph):

    """Predict labels

               hidden_node_dim: int = 64,

               in_edge_dim: int = 41,

               num_conv: int = 3,

               predicator_hidden_feats: int = 128,

               predictor_hidden_feats: int = 128,

               n_tasks: int = 1,

               mode: str = 'regression',

               n_classes: int = 2,

      based on default setting of CGCNNFeaturizer.

    num_conv: int, default 3

      The number of convolutional layers.

    predicator_hidden_feats: int, default 128

    predictor_hidden_feats: int, default 128

      The size for hidden representations in the output MLP predictor.

    n_tasks: int, default 1

      The number of the output size.

      This class accepts all the keyword arguments from TorchModel.

    """

    model = CGCNN(in_node_dim, hidden_node_dim, in_edge_dim, num_conv,

                  predicator_hidden_feats, n_tasks, mode, n_classes)

                  predictor_hidden_feats, n_tasks, mode, n_classes)

    if mode == "regression":

      loss: Loss = L2Loss()

      output_types = ['prediction']

      in_node_dim: int = 38,

      hidden_node_dim: int = 64,

      heads: int = 4,

      dropout_rate: float = 0.0,

      dropout: float = 0.0,

      num_conv: int = 3,

      predicator_hidden_feats: int = 32,

      predictor_hidden_feats: int = 32,

      n_tasks: int = 1,

  ):

    """

      The length of the hidden node feature vectors.

    heads: int, default 4

      The number of multi-head-attentions.

    dropout_rate: float, default 0.0

    dropout: float, default 0.0

      The dropout probability for each convolutional layer.

    num_conv: int, default 3

      The number of convolutional layers.

    predicator_hidden_feats: int, default 32

    predictor_hidden_feats: int, default 32

      The size for hidden representations in the output MLP predictor, default to 32.

    n_tasks: int, default 1

      The number of the output size, default to 1.

            out_channels=hidden_node_dim,

            heads=heads,

            concat=False,

            dropout=dropout_rate) for _ in range(num_conv)

            dropout=dropout) for _ in range(num_conv)

    ])

    self.pooling = global_mean_pool

    self.fc = nn.Linear(hidden_node_dim, predicator_hidden_feats)

  >> import deepchem as dc

  >> featurizer = dc.feat.MolGraphConvFeaturizer()

  >> dataset_config = {"reload": False, "featurizer": featurizer, "transformers": []}

  >> tasks, datasets, transformers = dc.molnet.load_tox21(**dataset_config)

  >> tasks, datasets, transformers = dc.molnet.load_tox21(reload=False, featurizer=featurizer, transformers=[])

  >> train, valid, test = datasets

  >> model = dc.models.GATModel(loss=dc.models.losses.SoftmaxCrossEntropy(), batch_size=32, learning_rate=0.001)

  >> model.fit(train, nb_epoch=50)

               in_node_dim: int = 38,

               hidden_node_dim: int = 64,

               heads: int = 4,

               dropout_rate: float = 0.0,

               dropout: float = 0.0,

               num_conv: int = 3,

               predicator_hidden_feats: int = 32,

               n_tasks: int = 1,

      The length of the hidden node feature vectors.

    heads: int, default 4

      The number of multi-head-attentions.

    dropout_rate: float, default 0.0

    dropout: float, default 0.0

      The dropout probability for each convolutional layer.

    num_conv: int, default 3

      The number of convolutional layers.

        in_node_dim,

        hidden_node_dim,

        heads,

        dropout_rate,

        dropout,

        num_conv,

        predicator_hidden_feats,

        n_tasks,