Changed RDKitMol to RDKitAtom for atom_features, added example docstring,...

from deepchem.feat.base_classes import MolecularFeaturizer

from deepchem.utils.typing import RDKitMol

from deepchem.utils.molecule_feature_utils import one_hot_encode

import numpy as np

from rdkit import Chem

    - Node Features

    - Adjacency Matrix

    - Distance Matrix

  References

  ---------

  .. [1] Lukasz Maziarka et al. "Molecule Attention Transformer`<https://arxiv.org/abs/2002.08264>`"

  Examples

  --------

  >>> import deepchem as dc

  >>> feat = dc.feat.MATFeaturizer()

  >>> out = feat.featurize("CCC")

  """

  def __init__(

      self,

      mol: Chem.rdchem.Mol,

      add_dummy_node: bool = True,

      one_hot_formal_charge: bool = True,

  ):

    """

    Parameters

    ----------

    mol (rdchem.Mol): rdkit.Chem.rdchem.Mol

      RDKit Mol object.

    add_dummy_node: bool, default True

      If True, a dummy node will be added to the molecular graph.

    one_hot_formal_charge: bool, default True

      If True, formal charges on atoms are one-hot encoded.

    References

    ---------

    .. [1] Lukasz Maziarka et al. "Molecule Attention Transformer`<https://arxiv.org/abs/2002.08264>`"

    """

    self.mol = mol

    self.add_dummy_node = add_dummy_node

    self.one_hot_formal_charge = one_hot_formal_charge

  def atom_features(self, atom):

    """

    Deepchem already contains an atom_features function, however we are defining a new one here due to the need to handle features specific to MAT.

    Since we need new features like Mol GetNeighbors and IsInRing, and the number of features required for MAT is a fraction of what the Deepchem atom_features function computes, we can speed up computation by defining a custom function.

    Since we need new features like Atom GetNeighbors and IsInRing, and the number of features required for MAT is a fraction of what the Deepchem atom_features function computes, we can speed up computation by defining a custom function.

    Parameters

    ----------

    mol: rdkit.Chem.rdchem.Mol

      RDKit mol object.

    atom: RDKitAtom

      RDKit Atom object.

    Returns

    ----------

    Parameters

    ----------

    mol: rdkit.Chem.rdchem.Mol

    mol: RDKitMol

      RDKit mol object.

    Returns