Merge pull request #3 from UCLCheminformatics/simple_functions

Added Simple functions

Author: OliverBScott

scaffoldgraph/core/graph.py

@@ -8,13 +8,15 @@
 from collections import Counter
 
 import tqdm
+import networkx as nx
+
 from loguru import logger
-from networkx import DiGraph
 from rdkit import RDLogger
 from rdkit.Chem import rdMolHash, MolToSmiles, rdmolops
 from rdkit.Chem.rdMolDescriptors import CalcNumRings
 
 from scaffoldgraph.io import *
+from scaffoldgraph.utils import canonize_smiles
 from .fragment import get_murcko_scaffold, get_annotated_murcko_scaffold
 from .scaffold import Scaffold
 
@@ -28,7 +30,7 @@ def init_molecule_name(mol):
         mol.SetProp('_Name', n)
 
 
-class ScaffoldGraph(DiGraph, ABC):
+class ScaffoldGraph(nx.DiGraph, ABC):
     """Abstract base class for ScaffoldGraphs"""
 
     def __init__(self, graph=None, fragmenter=None):
@@ -92,29 +94,125 @@ def _recursive_constructor(self, child):
     @property
     def num_scaffold_nodes(self):
         """Return the number of scaffolds in the graph"""
-        return len(list(self.get_scaffold_nodes()))
+        count = 0
+        for _ in self.get_scaffold_nodes():
+            count += 1
+        return count
 
     @property
     def num_molecule_nodes(self):
         """Return the number of molecules in the graph"""
-        return len(list(self.get_molecule_nodes()))
+        count = 0
+        for _ in self.get_molecule_nodes():
+            count += 1
+        return count
 
     def get_scaffold_nodes(self, data=False):
+        """Return a generator of all scaffold nodes in the graph"""
         if data is True:
             return ((n, self.nodes[n]) for n, d in self.nodes(data='type') if d == 'scaffold')
         else:
             return (n for n, d in self.nodes(data='type') if d == 'scaffold')
 
     def get_molecule_nodes(self, data=False):
+        """Return a generator of all molecule nodes in the graph"""
         if data is True:
             return ((n, self.nodes[n]) for n, d in self.nodes(data='type') if d == 'molecule')
         else:
             return (n for n, d in self.nodes(data='type') if d == 'molecule')
 
     def get_hierarchy_sizes(self):
+        """Return a collections.Counter object indicating the number of scaffolds
+        within each hierarchy level"""
         hierarchy = (d['hierarchy'] for _, d in self.get_scaffold_nodes(data=True))
         return Counter(hierarchy)
 
+    def max_hierarchy(self):
+        """Return the largest hierarchy level"""
+        return max(self.get_hierarchy_sizes())
+
+    def min_hierarchy(self):
+        """Return the smallest hierarchy level"""
+        return min(self.get_hierarchy_sizes())
+
+    def get_scaffolds_in_hierarchy(self, hierarchy):
+        """Return a generator of all scaffolds within a specified hierarchy"""
+        for s, d in self.get_scaffold_nodes(data=True):
+            if d['hierarchy'] == int(hierarchy):
+                yield s
+
+    def scaffold_in_graph(self, scaffold_smiles):
+        """Returns True if specified scaffold SMILES is in the scaffold graph
+
+        Parameters
+        ----------
+        scaffold_smiles : (str) SMILES of query scaffold.
+        """
+        result = scaffold_smiles in self
+        if result is not True:
+            scaffold_smiles = canonize_smiles(scaffold_smiles, failsafe=True)
+            result = scaffold_smiles in self
+        return result
+
+    def molecule_in_graph(self, molecule_id):
+        """Returns True if specified molecule ID is in the scaffold graph
+
+        Parameters
+        ----------
+        molecule_id: (str) ID of query molecule.
+        """
+        return str(molecule_id) in self
+
+    def get_molecules_for_scaffold(self, scaffold_smiles):
+        """Return a list of molecule IDs which are represented by a scaffold in the graph.
+
+        Note: This is determined by traversing the graph. In the case of a scaffold tree
+        the results represent the rules used to prioritize the scaffolds.
+
+        Parameters
+        ----------
+        scaffold_smiles : (str) SMILES of query scaffold.
+        """
+        molecules = []
+        if scaffold_smiles not in self:
+            scaffold_smiles = canonize_smiles(scaffold_smiles, failsafe=True)
+            if scaffold_smiles not in self:
+                return molecules
+        for succ in nx.bfs_tree(self, scaffold_smiles, reverse=False):
+            if self.nodes[succ]['type'] == 'molecule':
+                molecules.append(succ)
+        return molecules
+
+    def get_scaffolds_for_molecule(self, molecule_id):
+        """Return a list of scaffold SMILES connected to a query molecule ID
+
+        Parameters
+        ----------
+        molecule_id: (str) ID of query molecule.
+        """
+        scaffolds = []
+        if molecule_id not in self:
+            return scaffolds
+        for succ in nx.bfs_tree(self, molecule_id, reverse=True):
+            if self.nodes[succ]['type'] == 'scaffold':
+                scaffolds.append(succ)
+        return scaffolds
+
+    def separate_disconnected_components(self, sort=False):
+        """Separate disconnected components into distinct ScaffoldGraph objects.
+
+        Parameters
+        ----------
+        sort: if True sort components in descending order according
+            to the number of nodes in the subgraph.
+        """
+        components = []
+        for c in nx.weakly_connected_components(self):
+            components.append(self.subgraph(c).copy())
+        if sort:
+            return sorted(components, key=len, reverse=True)
+        return components
+
     def add_molecule_node(self, molecule, **attr):
         name = molecule.GetProp('_Name')
         default_attr = dict(type='molecule', smiles=MolToSmiles(molecule))

scaffoldgraph/utils/__init__.py

@@ -2,8 +2,10 @@
 scaffoldgraph.utils
 """
 
+from .misc import canonize_smiles
 from .aggregate import aggregate
 
 __all__ = [
+    'canonize_smiles',
     'aggregate'
 ]

scaffoldgraph/utils/misc.py

@@ -0,0 +1,15 @@
+"""
+scaffoldgraph.utils.misc
+
+Defines miscellaneous functions used within scaffoldgraph
+"""
+
+from rdkit import Chem
+
+
+def canonize_smiles(smiles, failsafe=True):
+    """Canonize a SMILES string (with failsafe)"""
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None and failsafe:
+        return smiles
+    return Chem.MolToSmiles(mol)

tests/core/test_graph.py

@@ -17,5 +17,5 @@ def test_init_molecule_name():
 
 
 def test_graph_subclass():
-    assert issubclass(ScaffoldGraph, DiGraph)
+    assert issubclass(ScaffoldGraph, nx.DiGraph)
     assert issubclass(ScaffoldGraph, ABC)

tests/test_network.py

@@ -15,6 +15,12 @@ def test_network(sdf_file):
     return network
 
 
+@pytest.fixture(name='network')
+def long_test_network():
+    network = sg.ScaffoldNetwork.from_smiles_file('tests/test_smiles.smi')
+    return network
+
+
 def test_network_from_sdf(sdf_file):
     network = sg.ScaffoldNetwork.from_sdf(sdf_file)
     assert network.num_scaffold_nodes == 8
@@ -33,5 +39,47 @@ def test_hiers(sdf_file):
     assert network.num_molecule_nodes == 2
 
 
+def test_hierarchy_functions(network):
+    hierarchy_sizes = network.get_hierarchy_sizes()
+    assert hierarchy_sizes[1] == 7
+    assert hierarchy_sizes[2] == 10
+    assert hierarchy_sizes[3] == 7
+    assert hierarchy_sizes[4] == 1
+    assert network.max_hierarchy() == 4
+    assert network.min_hierarchy() == 1
+    s_in_h2 = {
+        'C1=Cn2cnnc2CN=C1', 'O=C1CN=C(c2ccccn2)C=CN1', 'O=C1CN=Cc2ccccc2N1',
+        'C1=CC(c2ccccc2)=[NH+]CC=N1', 'C1=Nc2ccccc2C=[NH+]C1',
+        'O=C1CC(=O)N(c2ccccc2)C=CN1', 'O=C1CC(=O)Nc2ccccc2N1',
+        'O=C1CN=C(c2ccccc2)C=CN1', 'O=C1C[NH+]=Cc2ccccc2N1',
+        'O=C1C[NH+]=C(c2ccccc2)C=CN1'
+    }
+    assert s_in_h2 == set(network.get_scaffolds_in_hierarchy(2))
+
+
+def test_simple_functions(network):
+    assert network.scaffold_in_graph('C1=Cn2cnnc2CN=C1') is True
+    # Below is the non-canonical SMILES of the above
+    assert network.scaffold_in_graph('C1=C-n2:c:n:n:c:2-C-N=C-1') is True
+    assert network.scaffold_in_graph('c1ccccc1CCNc2ccccc2') is False
+    assert network.molecule_in_graph('Adinazolam') is True
+    assert network.molecule_in_graph('Citalopram') is False
+
+
+def test_traversal(network):
+    s_for_adinazolam = {
+        'c1ccc(C2=NCc3nncn3-c3ccccc32)cc1', 'C1=NCc2nncn2-c2ccccc21',
+        'C1=Cn2cnnc2CN=C1c1ccccc1', 'C1=Cn2cnnc2CN=C1', 'c1nnc[nH]1'
+    }
+    assert set(network.get_scaffolds_for_molecule('Adinazolam')) == s_for_adinazolam
+    m_for_scaffold = {'Adinazolam', 'Alprazolam'}
+    assert set(network.get_molecules_for_scaffold('c1nnc[nH]1')) == m_for_scaffold
+
+
+def test_separate_disconnected(network):
+    assert len(network.separate_disconnected_components()) == 2
+    assert type(network.separate_disconnected_components()[0]) == type(network)
+
+
 def test_repr(test_net):
     assert repr(test_net) == '<ScaffoldNetwork at {}>'.format(hex(id(test_net)))

tests/test_smiles.smi

@@ -0,0 +1,10 @@
+CN(C)Cc1n-2c(nn1)CN=C(c1ccccc1)c1cc(Cl)ccc12 Adinazolam
+Cc1n-2c(nn1)CN=C(c1ccccc1)c1cc(Cl)ccc12 Alprazolam
+Brc1cc2c(cc1)NC(=O)CN=C2c1ncccc1 Bromazepam
+CNC1=Nc2c(cc(Cl)cc2)C(c2ccccc2)=[N+]([O-])C1 Chlordiazepoxide
+CN1c2ccc(Cl)cc2N(c2ccccc2)C(=O)CC1=O Clobazam
+[O-][N+](=O)c1cc2c(cc1)NC(=O)CN=C2c1c(Cl)cccc1 Clonazepam
+OC(=O)C1N=C(c2ccccc2)c2cc(Cl)ccc2NC1=O Clorazepate
+Clc1cc2c(cc1)NC(=O)CN=C2c1c(Cl)cccc1 Delorazepam
+[O-][N+]1=C(c2ccccc2)c2cc(Cl)ccc2NC(=O)C1 Demoxepam
+Clc1cc2c(cc1)NC(=O)CC(=O)N2c1ccccc1 Desmethylclobazam