refactoring rdkit grid featurizer

"""

Topological fingerprints for macromolecular structures.

"""

import numpy as np

import logging

from deepchem.utils.hash_utils import hash_ecfp

from deepchem.feat import ComplexFeaturizer

from deepchem.utils.hash_utils import vectorize

from deepchem.utils.rdkit_util import compute_all_ecfp

from deepchem.utils.rdkit_util import compute_pairwise_distances

logger = logging.getLogger(__name__)

def featurize_binding_pocket_ecfp(protein_xyz,

                                  protein,

                                  ligand_xyz,

                                  ligand,

                                  pairwise_distances=None,

                                  cutoff=4.5,

                                  ecfp_degree=2):

  """Computes ECFP dicts for ligand and binding pocket of the protein.

  Parameters

  ----------

  protein_xyz: np.ndarray

    Of shape (N_protein_atoms, 3)

  protein: rdkit.rdchem.Mol

    Contains more metadata.

  ligand_xyz: np.ndarray

    Of shape (N_ligand_atoms, 3)

  ligand: rdkit.rdchem.Mol

    Contains more metadata

  pairwise_distances: np.ndarray

    Array of pairwise protein-ligand distances (Angstroms)

  cutoff: float

    Cutoff distance for contact consideration

  ecfp_degree: int

    ECFP radius

  """

  if pairwise_distances is None:

    pairwise_distances = compute_pairwise_distances(protein_xyz, ligand_xyz)

  contacts = np.nonzero((pairwise_distances < cutoff))

  protein_atoms = set([int(c) for c in contacts[0].tolist()])

  protein_ecfp_dict = compute_all_ecfp(

      protein, indices=protein_atoms, degree=ecfp_degree)

  ligand_ecfp_dict = compute_all_ecfp(ligand, degree=ecfp_degree)

  return (protein_ecfp_dict, ligand_ecfp_dict)

class ContactCircularFingerprint(ComplexFeaturizer):

  """Compute (Morgan) fingerprints near contact points of macromolecular complexes.

  Given a macromolecular complex made up of multiple constituent molecules, first compute the contact points where atoms from different molecules come close to one another. For atoms within "contact regions," compute radial "ECFP" fragments which are sub-molecules centered at atoms in the contact region.

  """

  def __init__(self, 

               cutoff=4.5,

               radius=2,

               size=8):

    """

    Parameters

    ----------

    cutoff: float (default 4.5)

      Distance cutoff in angstroms for molecules in complex.

    radius : int, optional (default 2)

        Fingerprint radius.

    size : int, optional (default 8)

      Length of generated bit vector.

    """

    self.cutoff = cutoff

    self.radius = radius

    self.size = size

  def _featurize_complex(self, mol, protein):

    """

    Compute featurization for a single mol/protein complex

    TODO(rbharath): This is very not ergonomic. I'd much prefer

    returning an vector instead of a list of two vectors. In

    addition, there's a question of efficiency.

    RdkitGridFeaturizer caches rotated versions etc internally.

    To make things work out of box, we are accepting that

    kludgey input. This needs to be cleaned up before full

    merge.

    Parameters

    ----------

    mol: object

      Representation of the molecule

    protein: object

      Representation of the protein

    """

    (lig_xyz, lig_rdk), (prot_xyz, prot_rdk) = mol, protein

    distances = compute_pairwise_distances(prot_xyz, lig_xyz)

    return [

        vectorize(

            hash_ecfp, feature_dict=ecfp_dict, size=self.size)

        for ecfp_dict in featurize_binding_pocket_ecfp(

            prot_xyz,

            prot_rdk,

            lig_xyz,

            lig_rdk,

            distances,

            cutoff=self.cutoff,

            ecfp_degree=self.radius)

    ]

"""

SPLIF Fingeprints for protein-ligand complexes.

"""

import logging

import numpy as np

from deepchem.utils.hash_utils import hash_ecfp_pair

from deepchem.utils.rdkit_util import compute_all_ecfp

from deepchem.feat import ComplexFeaturizer

from deepchem.utils.hash_utils import vectorize

from deepchem.utils.rdkit_util import compute_pairwise_distances

logger = logging.getLogger(__name__)

SPLIF_CONTACT_BINS = [(0, 2.0), (2.0, 3.0), (3.0, 4.5)]

def compute_splif_features_in_range(protein,

                                    ligand,

                                    pairwise_distances,

                                    contact_bin,

                                    ecfp_degree=2):

  """Computes SPLIF features for protein atoms close to ligand atoms.

  Finds all protein atoms that are > contact_bin[0] and <

  contact_bin[1] away from ligand atoms. Then, finds the ECFP

  fingerprints for the contacting atoms. Returns a dictionary

  mapping (protein_index_i, ligand_index_j) --> (protein_ecfp_i,

  ligand_ecfp_j)

  """

  contacts = np.nonzero((pairwise_distances > contact_bin[0]) &

                        (pairwise_distances < contact_bin[1]))

  protein_atoms = set([int(c) for c in contacts[0].tolist()])

  contacts = zip(contacts[0], contacts[1])

  protein_ecfp_dict = compute_all_ecfp(

      protein, indices=protein_atoms, degree=ecfp_degree)

  ligand_ecfp_dict = compute_all_ecfp(ligand, degree=ecfp_degree)

  splif_dict = {

      contact: (protein_ecfp_dict[contact[0]], ligand_ecfp_dict[contact[1]])

      for contact in contacts

  }

  return (splif_dict)

def featurize_splif(protein_xyz, protein, ligand_xyz, ligand,

contact_bins,

                    pairwise_distances, ecfp_degree):

  """Computes SPLIF featurization of protein-ligand binding pocket.

  For each contact range (i.e. 1 A to 2 A, 2 A to 3 A, etc.)

  compute a dictionary mapping (protein_index_i, ligand_index_j)

  tuples --> (protein_ecfp_i, ligand_ecfp_j) tuples. Return a

  list of such splif dictionaries.

  """

  splif_dicts = []

  for i, contact_bin in enumerate(contact_bins):

    splif_dicts.append(

        compute_splif_features_in_range(protein, ligand, pairwise_distances,

                                        contact_bin, ecfp_degree))

  return (splif_dicts)

class SplifFingerprint(ComplexFeaturizer):

  """Computes SPLIF Fingerprints for a macromolecular complex.

  SPLIF fingerprints are based on a technique introduced in the

  following paper. 

  Da, C., and D. Kireev. "Structural protein–ligand interaction

  fingerprints (SPLIF) for structure-based virtual screening:

  method and benchmark study." Journal of chemical information

  and modeling 54.9 (2014): 2555-2561.

  SPLIF fingerprints are a subclass of `ComplexFeaturizer`. It

  requires 3D coordinates for a protein-ligand complex. For

  each ligand atom, it identifies close protein atoms. These

  atom pairs are expanded to 2D circular fragments and a

  fingerprint for the union is turned on in the bit vector.

  """

  def __init__(self, 

               contact_bins=None,

               radius=2,

               size=8):

    """

    Parameters

    ----------

    contact_bins: list[tuple] 

      List of contact bins. If not specified is set to default

      `[(0, 2.0), (2.0, 3.0), (3.0, 4.5)]`.

    radius : int, optional (default 2)

        Fingerprint radius used for circular fingerprints.

    size: int, optional (default 8)

      Length of generated bit vector.

    """

    if contact_bins is None:

      self.contact_bins = SPLIF_CONTACT_BINS

    else:

      self.contact_bins = contact_bins

    self.size = size

    self.radius = radius

  def _featurize_complex(self, mol, protein):

    """

    Compute featurization for a single mol/protein complex

    TODO(rbharath): This is very not ergonomic. I'd much prefer

    returning an vector instead of a list of two vectors. In

    addition, there's a question of efficiency.

    RdkitGridFeaturizer caches rotated versions etc internally.

    To make things work out of box, we are accepting that

    kludgey input. This needs to be cleaned up before full

    merge.

    Parameters

    ----------

    mol: object

      Representation of the molecule

    protein: object

      Representation of the protein

    """

    (lig_xyz, lig_rdk), (prot_xyz, prot_rdk) = mol, protein

    distances = compute_pairwise_distances(prot_xyz, lig_xyz)

    return [

        vectorize(hash_ecfp_pair, feature_dict=splif_dict,

            size=self.size) for splif_dict in featurize_splif(

                prot_xyz, prot_rdk, lig_xyz, lig_rdk, self.contact_bins, distances, self.radius)

    ]

np.random.seed(123)

from deepchem.utils import rdkit_util

from deepchem.utils import hash_utils

from deepchem.utils import voxel_utils

from deepchem.feat import rdkit_grid_featurizer as rgf

def random_string(length, chars=None):

  import string

  if chars is None:

    chars = list(string.ascii_letters + string.ascii_letters + '()[]+-.=#@/\\')

  return ''.join(np.random.choice(chars, length))

from deepchem.feat.splif_fingerprint import compute_splif_features_in_range

class TestHelperFunctions(unittest.TestCase):

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

    self.protein_file = os.path.join(current_dir, 'data',

                                     '3ws9_protein_fixer_rdkit.pdb')

    self.ligand_file = os.path.join(current_dir, 'data', '3ws9_ligand.sdf')

  def test_hash_ecfp(self):

    for power in (2, 16, 64):

      for _ in range(10):

        string = random_string(10)

        string_hash = rgf.hash_ecfp(string, power)

        self.assertIsInstance(string_hash, int)

        self.assertLess(string_hash, 2**power)

        self.assertGreaterEqual(string_hash, 0)

  def test_hash_ecfp_pair(self):

    for power in (2, 16, 64):

      for _ in range(10):

        string1 = random_string(10)

        string2 = random_string(10)

        pair_hash = rgf.hash_ecfp_pair((string1, string2), power)

        self.assertIsInstance(pair_hash, int)

        self.assertLess(pair_hash, 2**power)

        self.assertGreaterEqual(pair_hash, 0)

    self.ligand_file = os.path.join(current_dir, '3ws9_ligand.sdf')

  def test_convert_atom_to_voxel(self):

    # 20 points with coords between -5 and 5, centered at 0

    distance = rdkit_util.compute_pairwise_distances(prot_xyz, lig_xyz)

    for bins in ((0, 2), (2, 3)):

      splif_dict = rgf.compute_splif_features_in_range(

      splif_dict = compute_splif_features_in_range(

          prot_rdk,

          lig_rdk,

          distance,

        pairwise_distances=distance,

        ecfp_degree=2)

    expected_dicts = [

        rgf.compute_splif_features_in_range(

        compute_splif_features_in_range(

            prot_rdk, lig_rdk, distance, c_bin, ecfp_degree=2) for c_bin in bins

    ]

    self.assertIsInstance(dicts, list)

        flatten=True,

        sanitize=True)

    prot_tensor = rgf_featurizer._voxelize(

    prot_tensor = voxel_utils.voxelize(

        rgf.convert_atom_to_voxel,

        rgf_featurizer.voxels_per_edge,

        rgf_featurizer.box_width,

        rgf_featurizer.voxel_width,

        rgf.hash_ecfp,

        prot_xyz,

        feature_dict=prot_ecfp_dict,

        channel_power=f_power)

        nb_channel=2**f_power)

    self.assertEqual(prot_tensor.shape, tuple([box_w] * 3 + [2**f_power]))

    all_features = prot_tensor.sum()

    # protein is too big for the box, some features should be missing

    self.assertGreater(all_features, 0)

    self.assertLess(all_features, prot_rdk.GetNumAtoms())

    lig_tensor = rgf_featurizer._voxelize(

    lig_tensor = voxel_utils.voxelize(

        rgf.convert_atom_to_voxel,

        rgf_featurizer.voxels_per_edge,

        rgf_featurizer.box_width,

        rgf_featurizer.voxel_width,

        rgf.hash_ecfp,

        lig_xyz,

        feature_dict=lig_ecfp_dict,

        channel_power=f_power)

        nb_channel=2**f_power)

    self.assertEqual(lig_tensor.shape, tuple([box_w] * 3 + [2**f_power]))

    all_features = lig_tensor.sum()

    # whole ligand should fit in the box

logger = logging.getLogger(__name__)

def hash_ecfp(ecfp, size):

  """

  Returns an int < size representing given ECFP fragment.

  ecfp_hash = int(digest, 16) % (size)

  return (ecfp_hash)

def hash_ecfp_pair(ecfp_pair, size):

  """Returns an int < size representing that ECFP pair.

  ecfp_hash = int(digest, 16) % (size)

  return (ecfp_hash)

def vectorize(hash_function, feature_dict=None, size=1024):

def vectorize(hash_function,

              feature_dict=None,

              size=1024):

  """Helper function to vectorize a spatial description from a hash.

  Hash functions are used to perform spatial featurizations in

  feature_vector = np.zeros(size)

  if feature_dict is not None:

    on_channels = [

        hash_function(feature, size) for key, feature in feature_dict.items()

        hash_function(feature, size)

        for key, feature in feature_dict.items()

    ]

    feature_vector[on_channels] += 1