protac_splitter/graphs/e3_clustering.py

from typing import List, Optional, Tuple, Any, Dict
import functools

import pandas as pd
import numpy as np
from tqdm import tqdm
from sklearn.cluster import AgglomerativeClustering, KMeans
from scipy.stats import skew
from sklearn.metrics import silhouette_score, davies_bouldin_score, calinski_harabasz_score

from rdkit import Chem, DataStructs
from rdkit.Chem import rdFingerprintGenerator

from protac_splitter.graphs.utils import get_fp, numpy_to_rdkit_fp
from protac_splitter.chemoinformatics import remove_dummy_atoms


def get_umap_clusters_fp(fp_list: List[str], n_clusters: int = 7) -> np.ndarray:
    """
    Cluster a list of SMILES strings using the umap clustering algorithm.
    From Scaffold Splits Overestimate Virtual Screening Performance
    https://arxiv.org/abs/2406.00873

    Args:
        fp_list (List[str]): List of SMILES strings.
        n_clusters (int): The number of clusters to use for clustering.

    Returns:
        np.ndarray: Array of cluster labels corresponding to each SMILES string in the input list.
    """
    ac = AgglomerativeClustering(n_clusters=n_clusters)
    ac.fit_predict(np.stack(fp_list))
    return ac.labels_

def get_kmeans_clusters_fp(fp_list: List[str], n_clusters: int = 10, return_centroids: bool = False) -> np.ndarray:
    """
    Cluster a list of SMILES strings using the KMeans clustering algorithm.

    Args:
        fp_list (List[str]): List of SMILES strings.
        n_clusters (int): The number of clusters to use for clustering.
        return_centroids (bool): If True, return the cluster centroids as well.

    Returns:
        np.ndarray: Array of cluster labels corresponding to each SMILES string in the input list.
    """
    km = KMeans(n_clusters=n_clusters, n_init='auto', random_state=42, max_iter=1000)
    if return_centroids:
        km.fit(np.stack(fp_list))
        return km.labels_, km.cluster_centers_
    return km.fit_predict(np.stack(fp_list))

def evaluate_clusters(X: np.array, clusters: np.ndarray) -> Dict[str, float]:
    """ Compute clustering metrics and assess cluster size distribution.
    
    Args:
        X (np.array): The input data used for clustering.
        clusters (np.ndarray): The cluster labels for each data point in X.
        
    Returns:
        Dict[str, float]: A dictionary containing various clustering metrics:
            - silhouette: Silhouette score of the clustering.
            - davies_bouldin: Davies-Bouldin index of the clustering.
            - calinski_harabasz: Calinski-Harabasz index of the clustering.
            - avg_cluster_size: Average size of clusters.
            - avg_cluster_data_ratio: Ratio of average cluster size to total data size.
            - std_cluster_size: Standard deviation of cluster sizes.
            - min_cluster_size: Minimum size of clusters.
            - median_cluster_size: Median size of clusters.
            - max_cluster_size: Maximum size of clusters.
            - cluster_size_skewness: Skewness of cluster sizes indicating imbalance.
            - num_clusters: Number of unique clusters found.
    """
    
    unique_clusters = list(set(clusters))
    
    if len(unique_clusters) < 2:  # Avoid single-cluster issues
        return {
            "silhouette": -1,
            "davies_bouldin": float("inf"),
            "calinski_harabasz": -1,
            "avg_cluster_size": len(X),
            "avg_cluster_data_ratio": 1,
            "std_cluster_size": 0,
            "min_cluster_size": len(X),
            "median_cluster_size": len(X),
            "max_cluster_size": len(X),
            "cluster_size_skewness": 0,
            "num_clusters": 1,
        }

    # Compute standard clustering metrics
    silhouette = silhouette_score(X, clusters)
    davies_bouldin = davies_bouldin_score(X, clusters)
    calinski_harabasz = calinski_harabasz_score(X, clusters)

    # Compute cluster size statistics
    cluster_sizes = [len(np.where(clusters == i)[0]) for i in np.unique(clusters)]
    avg_cluster_size = np.mean(cluster_sizes)
    avg_cluster_data_ratio = avg_cluster_size / len(X)
    std_cluster_size = np.std(cluster_sizes)
    median_cluster_size = np.median(cluster_sizes)
    min_cluster_size = np.min(cluster_sizes)
    max_cluster_size = np.max(cluster_sizes)
    cluster_size_skewness = skew(cluster_sizes, nan_policy="omit")  # Indicates imbalance in cluster sizes

    return {
        "silhouette": silhouette,
        "davies_bouldin": davies_bouldin,
        "calinski_harabasz": calinski_harabasz,
        "avg_cluster_size": avg_cluster_size,
        "avg_cluster_data_ratio": avg_cluster_data_ratio,
        "std_cluster_size": std_cluster_size,
        "min_cluster_size": min_cluster_size,
        "median_cluster_size": median_cluster_size,
        "max_cluster_size": max_cluster_size,
        "cluster_size_skewness": cluster_size_skewness,
        "num_clusters": len(unique_clusters),
    }

def get_representative_e3s(
    train_df: pd.DataFrame,
    fp_generator: Optional[Any] = None,
    n_clusters_candidates: List[int] = [10, 25, 50, 100, 150],
    e3_column: str = 'E3 Binder SMILES with direction',
) -> Tuple[List[str], List[Any], int, pd.DataFrame]:
    """
    Get representative E3 ligands from a DataFrame of training data by clustering their fingerprints.
    This function computes Morgan fingerprints for unique E3 ligands, clusters them using KMeans and UMAP,
    evaluates the clusters using silhouette, Davies-Bouldin, and Calinski-Harabasz scores, and identifies
    the optimal number of clusters based on these metrics.
    It returns the representative E3 ligands, their fingerprints, the best number of clusters, and a DataFrame
    containing the clustering metrics.
    
    Parameters:
        train_df (pd.DataFrame): DataFrame containing training data with E3 ligands.
        fp_generator (Optional[Any]): RDKit fingerprint generator. If None, a default Morgan fingerprint generator with 1024 bits and radius 6 is used.
        n_clusters_candidates (List[int]): List of candidate numbers of clusters to evaluate.
        e3_column (str): The column name in the DataFrame that contains the E3 ligand SMILES strings.
        
    Returns:
        Tuple[List[str], List[Any], int, pd.DataFrame]: A tuple containing:
            - List of representative E3 ligand SMILES strings.
            - List of RDKit fingerprints corresponding to the representative E3 ligands.
            - The best number of clusters determined from the clustering metrics.
            - DataFrame containing clustering metrics for each candidate number of clusters.
    """
    if e3_column not in train_df.columns:
        raise ValueError(f"Column '{e3_column}' not found in the DataFrame.")

    if fp_generator is None:
        fp_generator = rdFingerprintGenerator.GetMorganGenerator(
            radius=16,
            fpSize=1024,
            useBondTypes=True,
            includeChirality=True,
        )

    fp_dict = {}
    for smi in tqdm(train_df[e3_column].unique()):
        fp = get_fp(remove_dummy_atoms(smi), fp_generator)
        if fp is not None:
            fp_dict[smi] = fp

    fp_list = list(fp_dict.values())
    fp2smiles = {fp.tobytes(): smi for smi, fp in fp_dict.items() if fp is not None}

    centroids_dict = {}
    clusters_dict = {}
    metrics_df = []
    for n_clusters in tqdm(n_clusters_candidates, desc="Clustering and evaluating"):
        clusters, centroids = get_kmeans_clusters_fp(fp_list, n_clusters=n_clusters, return_centroids=True)
        metrics = evaluate_clusters(fp_list, clusters)
        clusters_dict[f'kmeans_n{n_clusters}'] = clusters.copy()
        centroids_dict[n_clusters] = centroids.copy()

        metrics['num_clusters'] = n_clusters
        metrics['cluster_algorithm'] = 'kmeans'
        metrics_df.append(metrics.copy())
        
        clusters = get_umap_clusters_fp(fp_list, n_clusters=n_clusters)
        metrics = evaluate_clusters(fp_list, clusters)
        clusters_dict[f'umap_n{n_clusters}'] = clusters.copy()

        metrics['num_clusters'] = n_clusters
        metrics['cluster_algorithm'] = 'umap'
        metrics_df.append(metrics.copy())

    metrics_df = pd.DataFrame(metrics_df)

    # Get the sweet spot for the number of clusters
    # Flip davies_bouldin so that all metrics are to be maximized
    metrics_df['-davies_bouldin'] = -metrics_df['davies_bouldin']

    # Normalize all three metrics (by group if you want per algorithm)
    metrics = ['silhouette', '-davies_bouldin', 'calinski_harabasz']
    df_norm = metrics_df.copy()
    df_norm[metrics] = df_norm.groupby('cluster_algorithm')[metrics].transform(
        lambda x: (x - x.min()) / (x.max() - x.min())
    )

    # Measure divergence: standard deviation of normalized metrics per row
    df_norm['metric_divergence'] = df_norm[metrics].std(axis=1)

    # Pick the point with lowest divergence, possibly applying constraints (e.g. not too many clusters)
    sweet_spots = df_norm.loc[df_norm.groupby('cluster_algorithm')['metric_divergence'].idxmin()]

    best_n_clusters = sweet_spots[['num_clusters']]['num_clusters'].unique()[0]

    # Get the centroids of the clusters
    centroids = centroids_dict[best_n_clusters]

    # Get the cluster labels for the centroids
    clusters = np.array(clusters_dict[f'kmeans_n{n_clusters}'])
    representative_e3s = []
    representative_e3s_fp = []
    for label, centroid in enumerate(centroids):
        # Isolate the FP with the same label as the centroid
        fp_cluster = np.array(fp_list)[clusters == label]
        # Get the closest FP for the centroid, use euclidean distance
        distances = np.linalg.norm(fp_cluster - centroid, axis=1)
        closest_fp = np.argmin(distances)
        # To get the SMILES from the FP, use the fp2smiles dictionary
        closest_smiles = fp2smiles[fp_cluster[closest_fp].tobytes()]
        # Append the closest SMILES to the representative_e3s list
        representative_e3s.append(closest_smiles)
        representative_e3s_fp.append(fp_cluster[closest_fp])

    # Convert the representative E3s to RDKit fingerprints
    representative_e3s_fp = [numpy_to_rdkit_fp(fp) for fp in representative_e3s_fp]
    
    return representative_e3s, representative_e3s_fp, best_n_clusters, metrics_df


DEFAULT_REPRESENTATIVE_E3S = [
    'Cc1ncsc1-c1ccc(CNC(=O)[C@@H]2C[C@@H](O)CN2C(=O)CN[*:2])cc1',
    'O=C1CCC(N2Cc3c(N=[*:2])cccc3C2=O)C(=O)N1',
    'CC(=O)NC(C(=O)N1CC(O)CC1C(=O)[*:2])C(C)(C)C',
    'CN[C@@H](C)C(=O)N[C@H](C(=O)N1C[C@@H](Oc2ccccc2[*:2])C[C@H]1C(=O)N[C@@H]1CCCc2ccccc21)C1CCCCC1',
    'Cc1ncsc1-c1ccc(CNC(=O)C2CC(O)CN2C(=O)C(NC(=O)CCO[*:2])C(C)(C)C)cc1',
    'O=C1CCC(N2Cc3ccc([*:2])cc3C2=O)C(=O)N1',
    'COc1ccc(C2=N[C@@H](c3ccc(Cl)cc3)[C@@H](c3ccc(Cl)cc3)N2C(=O)N2CCN(CC(=O)[*:2])C(=O)C2)c(OC(C)C)c1',
    'CC(NC(=O)C1CC(O)CN1C(=O)C(N[*:2])C(C)(C)C)c1ccc(C2CC2)cc1',
    'CCOc1cc(C(C)(C)C)ccc1C1=NC(c2ccc(Cl)cc2)C(c2ccc(Cl)cc2)N1C(=O)N1CCN(CCCC[*:2])CC1',
    'CNC(C)C(=O)NC(C(=O)N1CCCC1c1cncc(C(=O)c2cccc([*:2])c2)c1)C1CCCCC1',
    'CN[C@@H](C)C(=O)N[C@H](C(=O)N1CCC[C@H]1c1nc(C(=O)c2ccc([*:2])cc2)cs1)C1CCCCC1',
    'O=C1CCC(N2C(=O)c3cccc(OC[*:2])c3C2=O)C(=O)N1',
    'CCOc1cc(C(C)(C)C)ccc1C1=NC(c2ccc(Cl)cc2)C(c2ccc(Cl)cc2)N1C(=O)N1CCN([*:2])CC1',
    'Cc1ncsc1-c1ccc(CNC(=O)[C@H]2C[C@H](O)CN2C(=O)C(N[*:2])C(C)(C)C)cc1',
    'Cc1ncsc1-c1ccc([C@H](C)NC(=O)[C@@H]2C[C@@H](O)CN2C(=O)[C@@H](N[*:2])C(C)(C)C)cc1',
    'CN[C@@H](C)C(=O)N[C@H](C(=O)N1CCC[C@H]1c1cncc(C(=O)c2cccc([*:2])c2)c1)C1CCCCC1',
    'Cc1ncsc1-c1ccc(CNC(=O)[C@@H]2C[C@@H](O)CN2C(=O)[C@@H](N[*:2])C(C)(C)C)c(OC2CCNCC2)c1',
    'CNC(C)C(=O)NC(C(=O)N1CC(Oc2ccc([*:2])cc2)CC1C(=O)NC1CCCc2ccccc21)C1CCCCC1',
    'C[C@H](NC(=O)[C@@H]1C[C@@H](O)CN1C(=O)[C@@H](N[*:2])C(C)(C)C)c1ccc(C(C)(C)C)cc1',
    'CNC(C)C(=O)NC(C(=O)N1CCCC1c1nc(C(=O)c2ccc([*:2])cc2)cs1)C1CCCCC1',
    'CC(=O)NC(C(=O)N1CC(O)CC1C(=O)NCc1ccc(-c2scnc2C)cc1[*:2])C(C)(C)C',
    'Cc1ncsc1-c1ccc(CNC(=O)[C@@H]2C[C@@H](O)CN2C(=O)[C@@H](NC(=O)C2(F)CC2)C(C)(C)C)c([*:2])c1',
    'CCOc1cc(C(C)(C)C)ccc1C1=NC(C)(c2ccc(Cl)cc2)C(C)(c2ccc(Cl)cc2)N1C(=O)N1CCN(CC(=O)[*:2])CC1',
    'COc1ccc(C(=O)[*:2])cc1N1CCC(=O)NC1=O',
    'CN[C@@H](C)C(=O)N[C@H](C(=O)N[C@H]1C[C@H]2CC[C@@H]1N(CCc1ccc([*:2])cc1)C2)C1CCCCC1',
    'CNC(C)C(=O)NC(C(=O)N1CC(N[*:2])CC1C(=O)NC1CCCc2ccccc21)C1CCCCC1',
    'CN[C@@H](C)C(=O)N[C@@H](CCCCN[*:2])C(=O)N1CCC[C@H]1C(=O)Nc1snnc1-c1ccccc1',
    'CNC(C)C(=O)NC(C(=O)NC1CC2CCC1N(CCc1cccc([*:2])c1)C2)C1CCCCC1',
    'O=C1CCC(N2C(=O)c3ccc(N[*:2])cc3C2=O)C(=O)N1',
    'CNC(C)C(=O)NC(C(=O)N1CC(NC(=O)CC[*:2])CC1C(=O)Nc1c(F)cccc1F)C(C)(C)C',
    'Cc1ncsc1-c1ccc(CNC(=O)[C@@H]2C[C@@H](O)CN2C(=O)[C@H](N[*:2])C(C)(C)C)cc1',
    'Cc1nc[nH]c1-c1ccc(CNC(=O)C2CC(O)CN2C(=O)C(N[*:2])C(C)(C)C)cc1',
    'Cc1ncsc1-c1ccc(C(C)NC(=O)C2CC(O)CN2C(=O)C(N[*:2])C(C)(C)C)cc1',
    'Cc1ncsc1-c1ccc(CNC(=O)[C@@H]2C[C@@H](O)CN2C(=O)[C@@H](N[*:2])C(C)(C)C)cc1',
    'O=C1CCC(c2cccc([*:2])c2)C(=O)N1',
    'CC(=O)N[C@H](C(=O)N1C[C@@H](O)C[C@@H]1C(=O)N[C@@H](CC(=O)N1CCC([*:2])CC1)c1ccccc1)C(C)C',
    'O=C(CCl)[*:2]',
    'CC[C@@H](NC(=O)[C@@H]1C[C@H](N[*:2])CN1C(=O)[C@@H](NC(=O)[C@H](C)NC)C(C)(C)C)c1ccccc1',
    'CN[C@H](C)C(=O)N[C@@H]1CCO[C@@H]2CC(C)(C)[C@H](C(=O)N[C@@H]3CCCc4cc([*:2])ccc43)N2C1=O',
    'CN[C@@H](C)C(=O)N[C@H](C(=O)N1CCC[C@H]1c1nc(C(=O)c2ccc(F)cc2)cs1)C1CCN(C[*:2])CC1',
    'Cc1ncsc1-c1ccc(CNC(=O)C2CC(O)CN2C(=O)C(N[*:2])C(C)(C)C)cc1',
    'CNC(C)C(=O)NC(CCCCN[*:2])C(=O)N1CCCC1C(=O)Nc1snnc1-c1ccccc1',
    'O=C1CCC(N2C(=O)c3cccc([*:2])c3C2=O)C(=O)O1',
    'COc1ccc(C2=N[C@@H](c3ccc(Cl)cc3)[C@@H](c3ccc(Cl)cc3)N2C(=O)N2CCN(CC(=O)[*:2])C(=O)C2)cc1OC(C)C',
    'Cc1ncsc1-c1ccc(CNC(=O)C2CC(O)CN2C(=O)C(N[*:2])C(C)(C)C)c(OC2CCNCC2)c1',
    'CNC(C)C(=O)NC(C(=O)N1CCCC1c1cncc(-n2ccc3c(C(=O)[*:2])cccc32)c1)C(C)C',
    'CCN1CCN(Cc2ccc(NC(=O)c3cccc(-c4ccc5nc(N[*:2])sc5n4)c3)cc2C(F)(F)F)CC1',
    'CN[C@@H](C)C(=O)N[C@H](C(=O)N1C[C@@H](NC(=O)CC[*:2])C[C@H]1C(=O)Nc1c(F)cccc1F)C(C)(C)C',
    'CNC(C)C(=O)NC(C(=O)N1CCCC1C(=O)NC(C(=O)[*:2])C(c1ccccc1)c1ccccc1)C1CCCCC1',
    'CC(=O)NCC(C(=O)N1CC(O)CC1C(=O)NC(CC(=O)N1CCC(N2CCC([*:2])CC2)CC1)c1ccccc1)C(C)C',
]


@functools.lru_cache(maxsize=1, typed=False)
def get_representative_e3s_fp(
    e3_list: Optional[List[str]] = None,
    fp_generator: Optional[Any] = None,
    verbose: int = 0,
) -> List[DataStructs.ExplicitBitVect]:
    """
    Generate Morgan fingerprints for a list of E3 ligands. If no list is provided,
    it uses a default list of representative E3 ligands.
    
    Parameters:
        e3_list (Optional[List[str]]): List of SMILES strings for E3 ligands. If None, uses a default list.
        fp_generator (Optional[Any]): RDKit fingerprint generator. If None, a default Morgan fingerprint generator is used.
        
    Returns:
        List[DataStructs.ExplicitBitVect]: List of RDKit Morgan fingerprints for the E3 ligands.
    """
    representative_e3s_fp = []
    if verbose > 0:
        iterable = tqdm(e3_list or DEFAULT_REPRESENTATIVE_E3S, desc="Generating fingerprints for E3 ligands")
    else:
        iterable = e3_list or DEFAULT_REPRESENTATIVE_E3S
    for smi in iterable:
        # Get the Morgan fingerprint for the SMILES string
        fp = get_fp(remove_dummy_atoms(smi), fp_generator, return_np=False)
        if fp is not None:
            representative_e3s_fp.append(fp)
        else:
            print(f"Warning: Invalid SMILES string '{smi}' encountered, skipping.")
    if not representative_e3s_fp:
        raise ValueError("No valid E3 ligands found in the provided list.")
    return representative_e3s_fp