add complete kmeans pubchem dataset

Author: sfluegel

chebai/preprocessing/datasets/pubchem.py

@@ -12,6 +12,7 @@
 import random
 import shutil
 import tempfile
+from scipy import spatial
 
 import pandas as pd
 from sklearn.model_selection import train_test_split
@@ -28,9 +29,11 @@
     ChEBIOver50,
     ChEBIOver100,
     ChEBIOverX,
+    _ChEBIDataExtractor,
 )
 from rdkit import Chem, DataStructs
 from rdkit.Chem import AllChem
+from sklearn.cluster import KMeans
 
 
 class PubChem(XYBaseDataModule):
@@ -215,73 +218,247 @@ def download(self):
 
 class PubChemKMeans(PubChem):
 
-    def __init__(self, *args, n_clusters=1e4, random_size=1e6, **kwargs):
+    def __init__(
+        self,
+        *args,
+        n_clusters=1e4,
+        random_size=1e6,
+        exclude_data_from: _ChEBIDataExtractor = None,
+        validation_size_limit=4000,
+        include_min_n_clusters=100,
+        **kwargs,
+    ):
         """k: number of entries in this dataset,
         n_random_subsets: number of subsets of random data from which to draw
         the most dissimilar molecules,
         random_size_factor: size of random subsets (in total) in relation to k"""
         self.n_clusters = int(n_clusters)
+        self.exclude_data_from = exclude_data_from
+        self.validation_size_limit = validation_size_limit
+        self.include_min_n_clusters = include_min_n_clusters
         super(PubChemKMeans, self).__init__(*args, k=int(random_size), **kwargs)
+        self._fingerprints = None
+        self._cluster_centers = None
+        self._fingerprints_clustered = None
+        self._exclusion_data_clustered = None
+        self._cluster_centers_superclustered = None
 
     @property
     def _name(self):
         return f"PubchemKMeans"
 
-    def download(self):
-        if self._k == PubChem.FULL:
-            super().download()
+    @property
+    def split_label(self):
+        if self._k and self._k != self.FULL:
+            return f"{self.n_clusters}_centers_out_of_{self._k}"
         else:
-            print(f"Loading random dataset (size: {self._k})...")
-            random_dataset = PubChem(k=self._k)
-            random_dataset.download()
+            return f"{self.n_clusters}_centers_out_of_full"
+
+    @property
+    def raw_file_names(self):
+        return ["train.txt", "validation.txt", "test.txt"]
+
+    @property
+    def fingerprints(self):
+        if self._fingerprints is None:
             fingerprints_path = os.path.join(self.raw_dir, "fingerprints.pkl")
             if not os.path.exists(fingerprints_path):
+                print(f"No fingerprints found...")
+                print(f"Loading random dataset (size: {self._k})...")
+                random_dataset = PubChem(k=self._k)
+                random_dataset.download()
                 with open(
                     os.path.join(random_dataset.raw_dir, "smiles.txt"), "r"
                 ) as f_in:
                     random_smiles = [s.split("\t")[1].strip() for s in f_in.readlines()]
                     fpgen = AllChem.GetRDKitFPGenerator()
-                    selected_smiles = []
                     print(f"Converting SMILES to molecules...")
                     mols = [Chem.MolFromSmiles(s) for s in tqdm.tqdm(random_smiles)]
                     print(f"Generating Fingerprints...")
                     fps = [
                         fpgen.GetFingerprint(m) if m is not None else m
                         for m in tqdm.tqdm(mols)
                     ]
-                    similarity = []
                     d = {"smiles": random_smiles, "fps": fps}
-                    df = pd.DataFrame(d, columns=["smiles", "fps"])
-                    df = df.dropna()
-                    df.to_pickle(open(fingerprints_path, "wb"))
+                    fingerprints_df = pd.DataFrame(d, columns=["smiles", "fps"])
+                    fingerprints_df = fingerprints_df.dropna()
+                    fingerprints_df.to_pickle(open(fingerprints_path, "wb"))
+                    self._fingerprints = fingerprints_df
             else:
-                df = pd.read_pickle(open(fingerprints_path, "rb"))
-            fps = np.array([list(vec) for vec in df["fps"].tolist()])
-            print(f"Starting k-means clustering...")
-            start_time = time.perf_counter()
-            kmeans = KMeans(n_clusters=self.n_clusters, random_state=0, n_init="auto")
-            kmeans.fit(fps)
-            print(f"Finished k-means in {time.perf_counter() - start_time:.2f} seconds")
-            df["label"] = kmeans.labels_
-            df.to_pickle(
-                open(
-                    os.path.join(
-                        self.raw_dir, f"fingerprints_labeled_{self.n_clusters}.pkl"
-                    ),
-                    "wb",
-                )
+                self._fingerprints = pd.read_pickle(open(fingerprints_path, "rb"))
+        return self._fingerprints
+
+    def _build_clusters(self):
+        fingerprints_clustered_path = os.path.join(
+            self.raw_dir, "fingerprints_clustered.pkl"
+        )
+        cluster_centers_path = os.path.join(self.raw_dir, f"cluster_centers.pkl")
+        print(f"Starting k-means clustering...")
+        start_time = time.perf_counter()
+        kmeans = KMeans(n_clusters=self.n_clusters, random_state=0, n_init="auto")
+        fps = np.array([list(vec) for vec in self.fingerprints["fps"].tolist()])
+        kmeans.fit(fps)
+        print(f"Finished k-means in {time.perf_counter() - start_time:.2f} seconds")
+        fingerprints_df = self.fingerprints
+        fingerprints_df["label"] = kmeans.labels_
+        fingerprints_df.to_pickle(
+            open(
+                fingerprints_clustered_path,
+                "wb",
             )
-            cluster_df = pd.DataFrame(
-                data={"centers": [center for center in kmeans.cluster_centers_]}
+        )
+        cluster_df = pd.DataFrame(
+            data={"centers": [center for center in kmeans.cluster_centers_]}
+        )
+        cluster_df.to_pickle(
+            open(
+                cluster_centers_path,
+                "wb",
             )
-            cluster_df.to_pickle(
-                open(
-                    os.path.join(
-                        self.raw_dir, f"cluster_centers_{self.n_clusters}.pkl"
-                    ),
-                    "wb",
+        )
+
+        return cluster_df, fingerprints_df
+
+    def _exclude_clusters(self, cluster_centers):
+        exclusion_data_path = os.path.join(self.raw_dir, "exclusion_data_clustered.pkl")
+        cluster_centers_np = np.array(
+            [
+                [cci for cci in cluster_center]
+                for cluster_center in cluster_centers["centers"]
+            ]
+        )
+        if self.exclude_data_from is not None:
+            if not os.path.exists(exclusion_data_path):
+                print(f"Loading data for exclusion of clusters...")
+                raw_chebi = []
+                for filename in self.exclude_data_from.raw_file_names:
+                    raw_chebi.append(
+                        pd.read_pickle(
+                            open(
+                                os.path.join(self.exclude_data_from.raw_dir, filename),
+                                "rb",
+                            )
+                        )
+                    )
+                raw_chebi = pd.concat(raw_chebi)
+                raw_chebi_smiles = np.array(raw_chebi["SMILES"])
+                fpgen = AllChem.GetRDKitFPGenerator()
+                print(f"Converting SMILES to molecules...")
+                mols = [Chem.MolFromSmiles(s) for s in tqdm.tqdm(raw_chebi_smiles)]
+                print(f"Generating Fingerprints...")
+                chebi_fps = [
+                    fpgen.GetFingerprint(m) if m is not None else m
+                    for m in tqdm.tqdm(mols)
+                ]
+                print(f"Finding cluster for each instance from exclusion-data")
+                chebi_fps = np.array([list(fp) for fp in chebi_fps if fp is not None])
+                tree = spatial.KDTree(cluster_centers_np)
+                chebi_clusters = [tree.query(fp)[1] for fp in chebi_fps]
+                chebi_clusters_df = pd.DataFrame(
+                    {"fp": [fp for fp in chebi_fps], "center_id": chebi_clusters},
+                    columns=["fp", "center_id"],
                 )
+                chebi_clusters_df.to_pickle(open(exclusion_data_path, "wb"))
+            else:
+                chebi_clusters_df = pd.read_pickle(open(exclusion_data_path, "rb"))
+            # filter pubchem clusters and remove all that contain data from the exclusion set
+            print(f"Removing clusters with data from exclusion-set")
+            counts = chebi_clusters_df["center_id"].value_counts()
+            cluster_centers["n_chebi_instances"] = counts
+            cluster_centers["n_chebi_instances"].fillna(0, inplace=True)
+            cluster_centers.sort_values(
+                by="n_chebi_instances", ascending=False, inplace=True
             )
+            zero_centers = cluster_centers[cluster_centers["n_chebi_instances"] == 0]
+            if len(zero_centers) > self.include_min_n_clusters:
+                cluster_centers = zero_centers
+            else:
+                cluster_centers = cluster_centers[-self.include_min_n_clusters :]
+        return cluster_centers
+
+    @property
+    def cluster_centers(self):
+        cluster_centers_path = os.path.join(self.raw_dir, f"cluster_centers.pkl")
+        if self._cluster_centers is None:
+            if os.path.exists(cluster_centers_path):
+                self._cluster_centers = pd.read_pickle(open(cluster_centers_path, "rb"))
+            else:
+                self._cluster_centers = self._build_clusters()[0]
+        return self._cluster_centers
+
+    @property
+    def fingerprints_clustered(self):
+        fingerprints_path = os.path.join(self.raw_dir, f"fingerprints_clustered.pkl")
+        if self._fingerprints_clustered is None:
+            if os.path.exists(fingerprints_path):
+                self._fingerprints_clustered = pd.read_pickle(
+                    open(fingerprints_path, "rb")
+                )
+            else:
+                self._fingerprints_clustered = self._build_clusters()[1]
+        return self._fingerprints_clustered
+
+    @property
+    def cluster_centers_superclustered(self):
+        cluster_centers_path = os.path.join(
+            self.raw_dir, f"cluster_centers_superclustered.pkl"
+        )
+        if self._cluster_centers_superclustered is None:
+            if not os.path.exists(cluster_centers_path):
+                clusters_filtered = self._exclude_clusters(self.cluster_centers)
+                print(f"Superclustering PubChem clusters")
+                kmeans = KMeans(n_clusters=3, random_state=0, n_init="auto")
+                clusters_np = np.array(
+                    [[cci for cci in center] for center in clusters_filtered["centers"]]
+                )
+                kmeans.fit(clusters_np)
+                clusters_filtered["label"] = kmeans.labels_
+                clusters_filtered.to_pickle(
+                    open(
+                        os.path.join(
+                            self.raw_dir, "cluster_centers_superclustered.pkl"
+                        ),
+                        "wb",
+                    )
+                )
+                self._cluster_centers_superclustered = clusters_filtered
+            else:
+                self._cluster_centers_superclustered = pd.read_pickle(
+                    open(
+                        os.path.join(
+                            self.raw_dir, f"cluster_centers_superclustered.pkl"
+                        ),
+                        "rb",
+                    )
+                )
+        return self._cluster_centers_superclustered
+
+    def download(self):
+        if self._k == PubChem.FULL:
+            super().download()
+        else:
+            if not all(
+                os.path.exists(os.path.join(self.raw_dir, file))
+                for file in self.raw_file_names
+            ):
+                fingerprints = self.fingerprints_clustered
+                fingerprints["big_cluster_assignment"] = fingerprints["label"].apply(
+                    lambda l: (
+                        -1
+                        if l not in self.cluster_centers_superclustered.index
+                        else self.cluster_centers_superclustered.loc[int(l), "label"]
+                    )
+                )
+                fp_grouped = fingerprints.groupby("big_cluster_assignment")
+                splits = [fp_grouped.get_group(g) for g in fp_grouped.groups if g != -1]
+                splits[0] = splits[0][: self.validation_size_limit]
+                splits.sort(key=lambda x: len(x))
+                for i, name in enumerate(["validation", "test", "train"]):
+                    if not os.path.exists(os.path.join(self.raw_dir, f"{name}.txt")):
+                        open(os.path.join(self.raw_dir, f"{name}.txt"), "x").close()
+                    with open(os.path.join(self.raw_dir, f"{name}.txt"), "w") as f:
+                        for id, row in splits[i].iterrows():
+                            f.writelines(f"{id}\t{row['smiles']}\n")
 
 
 class PubChemDissimilarSMILES(PubChemDissimilar):
@@ -383,7 +560,12 @@ def download(self):
 
 
 if __name__ == "__main__":
-    kmeans_data = PubChemKMeans()
+    kmeans_data = PubChemKMeans(
+        n_clusters=100,
+        random_size=10000,
+        exclude_data_from=ChEBIOver100(chebi_version=231),
+        include_min_n_clusters=10,
+    )
     kmeans_data.download()