chemprep.fingerprints

 1from rdkit import Chem
 2from rdkit.Chem import AllChem, MACCSkeys
 3from rdkit.Chem.rdmolops import RDKFingerprint
 4from rdkit.Chem.AtomPairs import Pairs, Torsions
 5from rdkit.Chem.Fingerprints import FingerprintMols
 6
 7def get_available_fingerprints():
 8    """
 9    Returns a list of available fingerprint types.
10    """
11    return [
12        "MACCS",
13        "RDKit",
14        "Topological/Daylight",
15        "AtomPair",
16        "TopologicalTorsion",
17        "Morgan",
18        "ECFP2",
19        "ECFP4",
20        "ECFP6",
21        "FCFP2",
22        "FCFP4",
23        "FCFP6",
24    ]
25
26def generate_fingerprint(smiles: str, fingerprint_type: str = "Morgan"):
27    """
28    Generates a molecular fingerprint for a given SMILES string.
29
30    Args:
31        smiles: The SMILES string of the molecule.
32        fingerprint_type: The type of fingerprint to generate.
33
34    Returns:
35        The generated fingerprint.
36    """
37    mol = Chem.MolFromSmiles(smiles)
38    if mol is None:
39        raise ValueError("Invalid SMILES string")
40
41    if fingerprint_type == "MACCS":
42        return MACCSkeys.GenMACCSKeys(mol)
43    elif fingerprint_type == "RDKit":
44        return RDKFingerprint(mol)
45    elif fingerprint_type == "Topological/Daylight":
46        return FingerprintMols.FingerprintMol(mol)
47    elif fingerprint_type == "AtomPair":
48        return Pairs.GetAtomPairFingerprint(mol)
49    elif fingerprint_type == "TopologicalTorsion":
50        return Torsions.GetTopologicalTorsionFingerprint(mol)
51    elif fingerprint_type == "Morgan":
52        return AllChem.GetMorganFingerprint(mol, 2)
53    elif fingerprint_type == "ECFP2":
54        return AllChem.GetMorganFingerprint(mol, 1)
55    elif fingerprint_type == "ECFP4":
56        return AllChem.GetMorganFingerprint(mol, 2)
57    elif fingerprint_type == "ECFP6":
58        return AllChem.GetMorganFingerprint(mol, 3)
59    elif fingerprint_type == "FCFP2":
60        return AllChem.GetMorganFingerprint(mol, 1, useFeatures=True)
61    elif fingerprint_type == "FCFP4":
62        return AllChem.GetMorganFingerprint(mol, 2, useFeatures=True)
63    elif fingerprint_type == "FCFP6":
64        return AllChem.GetMorganFingerprint(mol, 3, useFeatures=True)
65    else:
66        raise ValueError(f"Unsupported fingerprint type: {fingerprint_type}")