MoleculeRankAI/predict_activity.py at main · 21buddyboyal25-web/MoleculeRankAI · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
import pandas as pd
import numpy as np
import joblib

from rdkit import Chem
from rdkit.Chem import AllChem


# -----------------------------------------
# Load molecules
# -----------------------------------------

admet = pd.read_csv(
    "data/admet_predictions.csv"
)


# -----------------------------------------
# Load trained RF model
# -----------------------------------------

model = joblib.load(
    "models/rf_model.pkl"
)


# -----------------------------------------
# Generate fingerprints
# -----------------------------------------

def fingerprint(smiles):

    mol = Chem.MolFromSmiles(smiles)

    if mol is None:
        return None

    fp = AllChem.GetMorganFingerprintAsBitVect(
        mol,
        radius=2,
        nBits=2048
    )

    return np.array(fp)


X = []

valid_smiles = []

for smiles in admet["canonical_smiles"]:

    fp = fingerprint(smiles)

    if fp is not None:

        X.append(fp)
        valid_smiles.append(smiles)


X = np.array(X)


# -----------------------------------------
# Predict activity
# -----------------------------------------

predictions = model.predict(X)


results = pd.DataFrame(
    {
        "canonical_smiles": valid_smiles,
        "predicted_activity": predictions
    }
)


# -----------------------------------------
# Save predictions
# -----------------------------------------

results.to_csv(
    "data/activity_predictions.csv",
    index=False
)


print("Activity prediction complete")
print("Molecules predicted:", len(results))

print(results.head())