Files
rdkit/Code/GraphMol/Fingerprints/Wrap/testGenerators.py

421 lines
18 KiB
Python

import unittest
import numpy as np
from rdkit import Chem, DataStructs
from rdkit.Chem import rdFingerprintGenerator
class TestCase(unittest.TestCase):
def setUp(self):
pass
def testAtomPairGenerator(self):
m = Chem.MolFromSmiles('CCC')
g = rdFingerprintGenerator.GetAtomPairGenerator()
fp = g.GetSparseCountFingerprint(m)
nz = fp.GetNonzeroElements()
self.assertEqual(len(nz), 2)
fp = g.GetCountFingerprint(m)
nz = fp.GetNonzeroElements()
self.assertEqual(len(nz), 2)
fp = g.GetSparseFingerprint(m)
nzc = fp.GetNumOnBits()
self.assertEqual(nzc, 3)
fp = g.GetFingerprint(m)
nzc = fp.GetNumOnBits()
self.assertEqual(nzc, 3)
g = rdFingerprintGenerator.GetAtomPairGenerator(
atomInvariantsGenerator=rdFingerprintGenerator.GetAtomPairAtomInvGen())
fp = g.GetSparseCountFingerprint(m)
nz = fp.GetNonzeroElements()
self.assertEqual(len(nz), 2)
g = rdFingerprintGenerator.GetAtomPairGenerator(minDistance=2)
fp = g.GetSparseCountFingerprint(m)
nz = fp.GetNonzeroElements()
self.assertEqual(len(nz), 1)
g = rdFingerprintGenerator.GetAtomPairGenerator(maxDistance=1)
fp = g.GetSparseCountFingerprint(m)
nz = fp.GetNonzeroElements()
self.assertEqual(len(nz), 1)
g = rdFingerprintGenerator.GetAtomPairGenerator(countSimulation=False)
fp = g.GetSparseFingerprint(m)
nzc = fp.GetNumOnBits()
self.assertEqual(nzc, 2)
invGen = rdFingerprintGenerator.GetAtomPairAtomInvGen(includeChirality=False)
invGenChirality = rdFingerprintGenerator.GetAtomPairAtomInvGen(includeChirality=True)
g = rdFingerprintGenerator.GetAtomPairGenerator(includeChirality=False,
atomInvariantsGenerator=invGen)
gChirality = rdFingerprintGenerator.GetAtomPairGenerator(
includeChirality=True, atomInvariantsGenerator=invGenChirality)
fp = g.GetSparseCountFingerprint(m)
nz = fp.GetNonzeroElements()
fpChirality = gChirality.GetSparseCountFingerprint(m)
nzChirality = fpChirality.GetNonzeroElements()
self.assertNotEqual(nz.keys(), nzChirality.keys())
def testMorganGenerator(self):
m = Chem.MolFromSmiles('CCCC(=O)O')
g = rdFingerprintGenerator.GetMorganGenerator(3)
fp = g.GetSparseCountFingerprint(m)
nz = fp.GetNonzeroElements()
self.assertEqual(len(nz), 14)
invgen = rdFingerprintGenerator.GetMorganAtomInvGen(True)
g = rdFingerprintGenerator.GetMorganGenerator(radius=3, atomInvariantsGenerator=invgen)
fp = g.GetSparseCountFingerprint(m)
nz = fp.GetNonzeroElements()
self.assertEqual(len(nz), 14)
invgen = rdFingerprintGenerator.GetMorganFeatureAtomInvGen()
g = rdFingerprintGenerator.GetMorganGenerator(radius=3, atomInvariantsGenerator=invgen)
fp = g.GetSparseCountFingerprint(m)
nz = fp.GetNonzeroElements()
self.assertEqual(len(nz), 13)
ms = [Chem.MolFromSmiles(x, sanitize=False) for x in ('C1=CC=CN=N1', 'C1C=CC=NN=1')]
for m in ms:
m.UpdatePropertyCache()
Chem.GetSymmSSSR(m)
g = rdFingerprintGenerator.GetMorganGenerator(radius=2, useBondTypes=True)
self.assertNotEqual(g.GetSparseCountFingerprint(ms[0]), g.GetSparseCountFingerprint(ms[1]))
g = rdFingerprintGenerator.GetMorganGenerator(radius=2, useBondTypes=False)
self.assertEqual(g.GetSparseCountFingerprint(ms[0]), g.GetSparseCountFingerprint(ms[1]))
binvgen = rdFingerprintGenerator.GetMorganBondInvGen(useBondTypes=False)
g2 = rdFingerprintGenerator.GetMorganGenerator(radius=2, bondInvariantsGenerator=binvgen)
self.assertEqual(g.GetSparseCountFingerprint(ms[0]), g2.GetSparseCountFingerprint(ms[0]))
self.assertEqual(g.GetSparseCountFingerprint(ms[1]), g2.GetSparseCountFingerprint(ms[1]))
def testRDKitFPGenerator(self):
m = Chem.MolFromSmiles('CCCCC')
g = rdFingerprintGenerator.GetRDKitFPGenerator()
fp = g.GetSparseCountFingerprint(m)
nz = fp.GetNonzeroElements()
self.assertEqual(len(nz), 8)
def testTopologicalTorsionGenerator(self):
m = Chem.MolFromSmiles('CCCCC')
g = rdFingerprintGenerator.GetTopologicalTorsionGenerator()
fp = g.GetSparseCountFingerprint(m)
nz = fp.GetNonzeroElements()
self.assertEqual(len(nz), 1)
def testBulk(self):
m1 = Chem.MolFromSmiles('CCC')
m2 = Chem.MolFromSmiles('OCCCCC')
m3 = Chem.MolFromSmiles('CCCCC')
g = rdFingerprintGenerator.GetAtomPairGenerator()
results = rdFingerprintGenerator.GetSparseCountFPs([m1, m2, m3],
rdFingerprintGenerator.AtomPairFP)
self.assertEqual(results[0], g.GetSparseCountFingerprint(m1))
self.assertEqual(results[1], g.GetSparseCountFingerprint(m2))
self.assertEqual(results[2], g.GetSparseCountFingerprint(m3))
self.assertEqual(len(results), 3)
g = rdFingerprintGenerator.GetMorganGenerator(2)
results = rdFingerprintGenerator.GetSparseCountFPs([m1, m2, m3],
rdFingerprintGenerator.MorganFP)
self.assertEqual(results[0], g.GetSparseCountFingerprint(m1))
self.assertEqual(results[1], g.GetSparseCountFingerprint(m2))
self.assertEqual(results[2], g.GetSparseCountFingerprint(m3))
self.assertEqual(len(results), 3)
g = rdFingerprintGenerator.GetRDKitFPGenerator()
results = rdFingerprintGenerator.GetSparseCountFPs([m1, m2, m3], rdFingerprintGenerator.RDKitFP)
self.assertEqual(results[0], g.GetSparseCountFingerprint(m1))
self.assertEqual(results[1], g.GetSparseCountFingerprint(m2))
self.assertEqual(results[2], g.GetSparseCountFingerprint(m3))
self.assertEqual(len(results), 3)
g = rdFingerprintGenerator.GetTopologicalTorsionGenerator()
results = rdFingerprintGenerator.GetSparseCountFPs([m1, m2, m3],
rdFingerprintGenerator.TopologicalTorsionFP)
self.assertEqual(results[0], g.GetSparseCountFingerprint(m1))
self.assertEqual(results[1], g.GetSparseCountFingerprint(m2))
self.assertEqual(results[2], g.GetSparseCountFingerprint(m3))
self.assertEqual(len(results), 3)
g = rdFingerprintGenerator.GetAtomPairGenerator()
results = rdFingerprintGenerator.GetSparseFPs([m1, m2, m3], rdFingerprintGenerator.AtomPairFP)
self.assertEqual(results[0], g.GetSparseFingerprint(m1))
self.assertEqual(results[1], g.GetSparseFingerprint(m2))
self.assertEqual(results[2], g.GetSparseFingerprint(m3))
self.assertEqual(len(results), 3)
g = rdFingerprintGenerator.GetMorganGenerator(2)
results = rdFingerprintGenerator.GetSparseFPs([m1, m2, m3], rdFingerprintGenerator.MorganFP)
self.assertEqual(results[0], g.GetSparseFingerprint(m1))
self.assertEqual(results[1], g.GetSparseFingerprint(m2))
self.assertEqual(results[2], g.GetSparseFingerprint(m3))
self.assertEqual(len(results), 3)
g = rdFingerprintGenerator.GetRDKitFPGenerator()
results = rdFingerprintGenerator.GetSparseFPs([m1, m2, m3], rdFingerprintGenerator.RDKitFP)
self.assertEqual(results[0], g.GetSparseFingerprint(m1))
self.assertEqual(results[1], g.GetSparseFingerprint(m2))
self.assertEqual(results[2], g.GetSparseFingerprint(m3))
self.assertEqual(len(results), 3)
g = rdFingerprintGenerator.GetTopologicalTorsionGenerator()
results = rdFingerprintGenerator.GetSparseFPs([m1, m2, m3],
rdFingerprintGenerator.TopologicalTorsionFP)
self.assertEqual(results[0], g.GetSparseFingerprint(m1))
self.assertEqual(results[1], g.GetSparseFingerprint(m2))
self.assertEqual(results[2], g.GetSparseFingerprint(m3))
self.assertEqual(len(results), 3)
g = rdFingerprintGenerator.GetAtomPairGenerator()
results = rdFingerprintGenerator.GetCountFPs([m1, m2, m3], rdFingerprintGenerator.AtomPairFP)
self.assertEqual(results[0], g.GetCountFingerprint(m1))
self.assertEqual(results[1], g.GetCountFingerprint(m2))
self.assertEqual(results[2], g.GetCountFingerprint(m3))
self.assertEqual(len(results), 3)
g = rdFingerprintGenerator.GetMorganGenerator(2)
results = rdFingerprintGenerator.GetCountFPs([m1, m2, m3], rdFingerprintGenerator.MorganFP)
self.assertEqual(results[0], g.GetCountFingerprint(m1))
self.assertEqual(results[1], g.GetCountFingerprint(m2))
self.assertEqual(results[2], g.GetCountFingerprint(m3))
self.assertEqual(len(results), 3)
g = rdFingerprintGenerator.GetRDKitFPGenerator()
results = rdFingerprintGenerator.GetCountFPs([m1, m2, m3], rdFingerprintGenerator.RDKitFP)
self.assertEqual(results[0], g.GetCountFingerprint(m1))
self.assertEqual(results[1], g.GetCountFingerprint(m2))
self.assertEqual(results[2], g.GetCountFingerprint(m3))
self.assertEqual(len(results), 3)
g = rdFingerprintGenerator.GetTopologicalTorsionGenerator()
results = rdFingerprintGenerator.GetCountFPs([m1, m2, m3],
rdFingerprintGenerator.TopologicalTorsionFP)
self.assertEqual(results[0], g.GetCountFingerprint(m1))
self.assertEqual(results[1], g.GetCountFingerprint(m2))
self.assertEqual(results[2], g.GetCountFingerprint(m3))
self.assertEqual(len(results), 3)
def testNumBitsPerFeature(self):
m1 = Chem.MolFromSmiles('CCCO')
g = rdFingerprintGenerator.GetRDKitFPGenerator(minPath=1, maxPath=2)
fp = g.GetFingerprint(m1)
self.assertEqual(fp.GetNumOnBits(), 8)
g = rdFingerprintGenerator.GetRDKitFPGenerator(minPath=1, maxPath=2, numBitsPerFeature=1)
fp = g.GetFingerprint(m1)
self.assertEqual(fp.GetNumOnBits(), 4)
def testAdditionalOutput(self):
m1 = Chem.MolFromSmiles('CCO')
g = rdFingerprintGenerator.GetAtomPairGenerator()
ao = rdFingerprintGenerator.AdditionalOutput()
ao.AllocateAtomCounts()
fp = g.GetFingerprint(m1, additionalOutput=ao)
self.assertEqual(ao.GetAtomCounts(), (2, 2, 2))
self.assertIsNone(ao.GetAtomToBits())
self.assertIsNone(ao.GetBitInfoMap())
self.assertIsNone(ao.GetBitPaths())
ao = rdFingerprintGenerator.AdditionalOutput()
ao.AllocateAtomToBits()
fp = g.GetFingerprint(m1, additionalOutput=ao)
self.assertIsNone(ao.GetAtomCounts())
self.assertEqual(ao.GetAtomToBits(), ((1404, 1916), (1404, 1596), (1596, 1916)))
self.assertIsNone(ao.GetBitInfoMap())
self.assertIsNone(ao.GetBitPaths())
ao = rdFingerprintGenerator.AdditionalOutput()
ao.AllocateBitInfoMap()
fp = g.GetFingerprint(m1, additionalOutput=ao)
self.assertIsNone(ao.GetAtomCounts())
self.assertIsNone(ao.GetAtomToBits())
self.assertEqual(ao.GetBitInfoMap(), {1404: ((0, 1), ), 1596: ((1, 2), ), 1916: ((0, 2), )})
self.assertIsNone(ao.GetBitPaths())
def testCountBounds(self):
m = Chem.MolFromSmiles('COc1ccc(CCNC(=O)c2ccccc2C(=O)NCCc2ccc(OC)cc2)cc1')
fp1 = rdFingerprintGenerator.GetRDKitFPGenerator(fpSize=2048,
countSimulation=True).GetFingerprint(m)
fp2 = rdFingerprintGenerator.GetRDKitFPGenerator(fpSize=2048, countSimulation=True,
countBounds=(1, 8, 16, 32)).GetFingerprint(m)
self.assertNotEqual(fp1.GetNumOnBits(), fp2.GetNumOnBits())
fp1 = rdFingerprintGenerator.GetTopologicalTorsionGenerator(
fpSize=2048, countSimulation=True).GetFingerprint(m)
fp2 = rdFingerprintGenerator.GetTopologicalTorsionGenerator(fpSize=2048, countSimulation=True,
countBounds=(1, 8, 16,
32)).GetFingerprint(m)
self.assertNotEqual(fp1.GetNumOnBits(), fp2.GetNumOnBits())
fp1 = rdFingerprintGenerator.GetMorganGenerator(fpSize=2048,
countSimulation=True).GetFingerprint(m)
fp2 = rdFingerprintGenerator.GetMorganGenerator(fpSize=2048, countSimulation=True,
countBounds=(1, 8, 16, 32)).GetFingerprint(m)
self.assertNotEqual(fp1.GetNumOnBits(), fp2.GetNumOnBits())
fp1 = rdFingerprintGenerator.GetAtomPairGenerator(fpSize=2048,
countSimulation=True).GetFingerprint(m)
fp2 = rdFingerprintGenerator.GetAtomPairGenerator(fpSize=2048, countSimulation=True,
countBounds=(1, 8, 16, 32)).GetFingerprint(m)
self.assertNotEqual(fp1.GetNumOnBits(), fp2.GetNumOnBits())
def testNumpyFingerprints(self):
m = Chem.MolFromSmiles('COc1ccc(CCNC(=O)c2ccccc2C(=O)NCCc2ccc(OC)cc2)cc1')
for fn in (rdFingerprintGenerator.GetRDKitFPGenerator,
rdFingerprintGenerator.GetMorganGenerator,
rdFingerprintGenerator.GetAtomPairGenerator,
rdFingerprintGenerator.GetTopologicalTorsionGenerator):
gen = fn(fpSize=2048)
bv = gen.GetFingerprint(m)
oarr = np.zeros((bv.GetNumBits(), ), 'u1')
DataStructs.ConvertToNumpyArray(bv, oarr)
arr = gen.GetFingerprintAsNumPy(m)
np.testing.assert_array_equal(oarr, arr)
fp = gen.GetCountFingerprint(m)
oarr = np.zeros((fp.GetLength(), ), 'u4')
DataStructs.ConvertToNumpyArray(fp, oarr)
arr = gen.GetCountFingerprintAsNumPy(m)
np.testing.assert_array_equal(oarr, arr)
def testMorganRedundantEnvironments(self):
m = Chem.MolFromSmiles('CC(=O)O')
g = rdFingerprintGenerator.GetMorganGenerator(2)
fp = g.GetSparseCountFingerprint(m)
self.assertEqual(fp.GetTotalVal(), 8)
g = rdFingerprintGenerator.GetMorganGenerator(2, includeRedundantEnvironments=True)
fp = g.GetSparseCountFingerprint(m)
self.assertEqual(fp.GetTotalVal(), 12)
def testFingerprintOptions(self):
m = Chem.MolFromSmiles('CC(=O)O')
g = rdFingerprintGenerator.GetMorganGenerator(2)
fp = g.GetSparseCountFingerprint(m)
self.assertEqual(fp.GetTotalVal(), 8)
g.GetOptions().numBitsPerFeature = 2
fp = g.GetSparseCountFingerprint(m)
self.assertEqual(fp.GetTotalVal(), 16)
g.GetOptions().includeRedundantEnvironments = True
fp = g.GetSparseCountFingerprint(m)
self.assertEqual(fp.GetTotalVal(), 24)
m = Chem.MolFromSmiles('CC(=O)C')
g = rdFingerprintGenerator.GetAtomPairGenerator()
fp = g.GetFingerprint(m)
self.assertEqual(fp.GetNumOnBits(), 6)
g.GetOptions().countSimulation = False
fp = g.GetFingerprint(m)
self.assertEqual(fp.GetNumOnBits(), 4)
g.GetOptions().maxDistance = 1
fp = g.GetFingerprint(m)
self.assertEqual(fp.GetNumOnBits(), 2)
m = Chem.MolFromSmiles('OC(C)(C)C')
g = rdFingerprintGenerator.GetAtomPairGenerator()
fp = g.GetFingerprint(m)
self.assertEqual(fp.GetNumOnBits(), 7)
g.GetOptions().SetCountBounds((1, 2, 3, 4))
fp = g.GetFingerprint(m)
self.assertEqual(fp.GetNumOnBits(), 10)
def testTopologicalTorsionShortestPaths(self):
m = Chem.MolFromSmiles('CC1CCC1')
g = rdFingerprintGenerator.GetTopologicalTorsionGenerator()
fp = g.GetSparseCountFingerprint(m)
nz = fp.GetNonzeroElements()
self.assertEqual(len(nz), 3)
g.GetOptions().onlyShortestPaths = True
fp = g.GetSparseCountFingerprint(m)
nz = fp.GetNonzeroElements()
self.assertEqual(len(nz), 1)
def testMorganGeneratorMultiMol(self):
smis = [
'CC1CCC1',
'CCC1CCC1',
'CCCC1CCC1',
'CC1CC(O)C1',
'CC1CC(OC)C1',
]
for i in range(4):
smis = smis + smis
ms = [Chem.MolFromSmiles(smi) for smi in smis]
g = rdFingerprintGenerator.GetMorganGenerator()
ofps = tuple([g.GetFingerprint(m) for m in ms])
tfps = g.GetFingerprints(ms, numThreads=1)
for ofp, tfp in zip(ofps, tfps):
self.assertEqual(ofp, tfp)
tfps = g.GetFingerprints(ms, numThreads=4)
for ofp, tfp in zip(ofps, tfps):
self.assertEqual(ofp, tfp)
ofps = tuple([g.GetCountFingerprint(m) for m in ms])
tfps = g.GetCountFingerprints(ms, numThreads=1)
for ofp, tfp in zip(ofps, tfps):
self.assertEqual(ofp, tfp)
tfps = g.GetCountFingerprints(ms, numThreads=4)
for ofp, tfp in zip(ofps, tfps):
self.assertEqual(ofp, tfp)
ofps = tuple([g.GetSparseFingerprint(m) for m in ms])
tfps = g.GetSparseFingerprints(ms, numThreads=1)
for ofp, tfp in zip(ofps, tfps):
self.assertEqual(ofp, tfp)
tfps = g.GetSparseFingerprints(ms, numThreads=4)
for ofp, tfp in zip(ofps, tfps):
self.assertEqual(ofp, tfp)
ofps = tuple([g.GetSparseCountFingerprint(m) for m in ms])
tfps = g.GetSparseCountFingerprints(ms, numThreads=1)
for ofp, tfp in zip(ofps, tfps):
self.assertEqual(ofp, tfp)
tfps = g.GetSparseCountFingerprints(ms, numThreads=4)
for ofp, tfp in zip(ofps, tfps):
self.assertEqual(ofp, tfp)
def testFingerprintGeneratorOptionsLifetime(self):
# this should not result in a seg fault
import inspect
inspect.getmembers(rdFingerprintGenerator.GetRDKitFPGenerator().GetOptions())
def testGetAtomsPerBit(self):
fpg = rdFingerprintGenerator.GetTopologicalTorsionGenerator()
ao = rdFingerprintGenerator.AdditionalOutput()
ao.AllocateAtomsPerBit()
m = Chem.MolFromSmiles('c1ccccn1')
fp = fpg.GetFingerprint(m,additionalOutput=ao)
apb = ao.GetAtomsPerBit()
self.assertEqual(fp.GetNumOnBits(),len(apb))
self.assertEqual(apb[140],((2, 1, 0, 5), (2, 3, 4, 5)))
ao = rdFingerprintGenerator.AdditionalOutput()
_ = fpg.GetFingerprint(m,additionalOutput=ao)
self.assertIsNone(ao.GetAtomsPerBit())
def testJSONSerialization(self):
m = Chem.MolFromSmiles('CCC')
g1 = rdFingerprintGenerator.GetAtomPairGenerator()
fp1 = g1.GetFingerprint(m)
jsonStr = g1.ToJSON()
g2 = rdFingerprintGenerator.FingerprintGeneratorFromJSON(jsonStr)
fp2 = g2.GetFingerprint(m)
self.assertEqual(fp1, fp2)
if __name__ == '__main__':
unittest.main()