ml-jku/tox21_xgboost_classifier

1import json2import numpy as np3import pandas as pd4 5from datasets import load_dataset6 7from rdkit import Chem, DataStructs8from rdkit.Chem import Descriptors, rdFingerprintGenerator, MACCSkeys9from rdkit.Chem.rdchem import Mol10 11from .utils import USED_200_DESCR, TOX_SMARTS_PATH, Standardizer12 13 14def create_cleaned_mol_objects(smiles: list[str]) -> tuple[list[Mol], np.ndarray]:15    """This function creates cleaned RDKit mol objects from a list of SMILES.16    Taken from https://huggingface.co/spaces/ml-jku/mhnfs/blob/main/src/data_preprocessing/create_descriptors.py17    Modification by Antonia Ebner:18        - skip uncleanable molecules19        - return clean molecule mask20 21    Args:22        smiles (list[str]): list of SMILES23 24    Returns:25        list[Mol]: list of cleaned molecules26        np.ndarray[bool]: mask that contains False at index `i`, if molecule in `smiles` at27            index `i` could not be cleaned and was removed.28    """29    sm = Standardizer(canon_taut=True)30 31    clean_mol_mask = list()32    mols = list()33    for i, smile in enumerate(smiles):34        mol = Chem.MolFromSmiles(smile)35        standardized_mol, _ = sm.standardize_mol(mol)36        is_cleaned = standardized_mol is not None37        clean_mol_mask.append(is_cleaned)38        if not is_cleaned:39            continue40        can_mol = Chem.MolFromSmiles(Chem.MolToSmiles(standardized_mol))41        mols.append(can_mol)42 43    return mols, np.array(clean_mol_mask)44 45 46def create_ecfp_fps(mols: list[Mol], radius=3, fpsize=2048, **kwargs) -> np.ndarray:47    """This function ECFP fingerprints for a list of molecules.48    Inspired by from https://huggingface.co/spaces/ml-jku/mhnfs/blob/main/src/data_preprocessing/create_descriptors.py49 50    Args:51        mols (list[Mol]): list of molecules52 53    Returns:54        np.ndarray: ECFP fingerprints of molecules55    """56    ecfps = list()57 58    for mol in mols:59        gen = rdFingerprintGenerator.GetMorganGenerator(60            countSimulation=True, fpSize=fpsize, radius=radius61        )62        fp_sparse_vec = gen.GetCountFingerprint(mol)63 64        fp = np.zeros((0,), np.int8)65        DataStructs.ConvertToNumpyArray(fp_sparse_vec, fp)66 67        ecfps.append(fp)68 69    return np.array(ecfps)70 71 72def create_maccs_keys(mols: list[Mol]) -> np.ndarray:73    """This function creates MACCS keys for a list of molecules.74 75    Args:76        mols (list[Mol]): list of molecules77 78    Returns:79        np.ndarray: MACCS keys of molecules80    """81    maccs = [MACCSkeys.GenMACCSKeys(x) for x in mols]82    return np.array(maccs)83 84 85def get_tox_patterns(filepath: str):86    """This retrieves the tox features defined in filepath.87    Args:88        filepath (str): A list of tox features89    """90    # load patterns91    with open(filepath) as f:92        smarts_list = [s[1] for s in json.load(f)]93 94    # Code does not work for this case95    assert len([s for s in smarts_list if ("AND" in s) and ("OR" in s)]) == 096 97    # Chem.MolFromSmarts takes a long time so it pays of to parse all the smarts first98    # and then use them for all molecules. This gives a huge speedup over existing code.99    # a list of patterns, whether to negate the match result and how to join them to obtain one boolean value100    all_patterns = []101    for smarts in smarts_list:102        patterns = []  # list of smarts-patterns103        # value for each of the patterns above. Negates the values of the above later.104        negations = []105 106        if " AND " in smarts:107            smarts = smarts.split(" AND ")108            merge_any = False  # If an ' AND ' is found all 'subsmarts' have to match109        else:110            # If there is an ' OR ' present it's enough is any of the 'subsmarts' match.111            # This also accumulates smarts where neither ' OR ' nor ' AND ' occur112            smarts = smarts.split(" OR ")113            merge_any = True114 115        # for all subsmarts check if they are preceded by 'NOT '116        for s in smarts:117            neg = s.startswith("NOT ")118            if neg:119                s = s[4:]120            patterns.append(Chem.MolFromSmarts(s))121            negations.append(neg)122 123        all_patterns.append((patterns, negations, merge_any))124    return all_patterns125 126 127def create_tox_features(mols: list[Mol], patterns: list) -> np.ndarray:128    """Matches the tox patterns against a molecule. Returns a boolean array"""129    tox_data = []130    for mol in mols:131        mol_features = []132        for patts, negations, merge_any in patterns:133            matches = [mol.HasSubstructMatch(p) for p in patts]134            matches = [m != n for m, n in zip(matches, negations)]135            if merge_any:136                pres = any(matches)137            else:138                pres = all(matches)139            mol_features.append(pres)140 141        tox_data.append(np.array(mol_features))142 143    return np.array(tox_data)144 145 146def create_rdkit_descriptors(mols: list[Mol]) -> np.ndarray:147    """This function creates RDKit descriptors for a list of molecules.148    Taken from https://huggingface.co/spaces/ml-jku/mhnfs/blob/main/src/data_preprocessing/create_descriptors.py149 150    Args:151        mols (list[Mol]): list of molecules152 153    Returns:154        np.ndarray: RDKit descriptors of molecules155    """156    rdkit_descriptors = list()157 158    for mol in mols:159        descrs = []160        for _, descr_calc_fn in Descriptors._descList:161            descrs.append(descr_calc_fn(mol))162 163        descrs = np.array(descrs)164        descrs = descrs[USED_200_DESCR]165        rdkit_descriptors.append(descrs)166 167    return np.array(rdkit_descriptors)168 169 170def create_quantiles(raw_features: np.ndarray, ecdfs: list) -> np.ndarray:171    """Create quantile values for given features using the columns172    Taken from https://huggingface.co/spaces/ml-jku/mhnfs/blob/main/src/data_preprocessing/create_descriptors.py173 174    Args:175        raw_features (np.ndarray): values to put into quantiles176        ecdfs (list): ECDFs to use177 178    Returns:179        np.ndarray: computed quantiles180    """181    quantiles = np.zeros_like(raw_features)182 183    for column in range(raw_features.shape[1]):184        raw_values = raw_features[:, column].reshape(-1)185        ecdf = ecdfs[column]186        q = ecdf(raw_values)187        quantiles[:, column] = q188 189    return quantiles190 191 192def fill(features, mask, value=np.nan):193    n_mols = len(mask)194    n_features = features.shape[1]195 196    data = np.zeros(shape=(n_mols, n_features))197    data.fill(value)198    data[~mask] = features199    return data200 201 202def create_descriptors(203    smiles,204    descriptors,205    **ecfp_kwargs,206):207    """Generate molecular descriptors for multiple SMILES strings.208    Inspired by https://huggingface.co/spaces/ml-jku/mhnfs/blob/main/src/data_preprocessing/create_descriptors.py209 210    Each SMILES is processed and sanitized using RDKit.211    SMILES that cannot be sanitized are encoded with NaNs, and a corresponding boolean mask212    is returned to indicate which inputs were successfully processed.213 214    Args:215        smiles (list[str]): List of SMILES strings for which to generate descriptors.216        descriptors (list[str]): List of descriptor types to compute.217            Supported values include:218            ['ecfps', 'tox', 'maccs', 'rdkit_descrs'].219 220    Returns:221        tuple[dict[str, np.ndarray], np.ndarray]:222            - A dictionary mapping descriptor names to their computed arrays.223            - A boolean mask of shape (len(smiles),) indicating which SMILES224            were successfully sanitized and processed.225    """226    # Create cleanded rdkit mol objects227    mols, clean_mol_mask = create_cleaned_mol_objects(smiles)228    print(f"Cleaned molecules, {(~clean_mol_mask).sum()} could not be sanitized")229 230    # Create fingerprints and descriptors231    if "ecfps" in descriptors:232        ecfps = create_ecfp_fps(mols, **ecfp_kwargs)233        ecfps = fill(ecfps, ~clean_mol_mask)234        print("Created ECFP fingerprints")235 236    if "tox" in descriptors:237        tox_patterns = get_tox_patterns(TOX_SMARTS_PATH)238        tox = create_tox_features(mols, tox_patterns)239        tox = fill(tox, ~clean_mol_mask)240        print("Created Tox features")241 242    if "maccs" in descriptors:243        maccs = create_maccs_keys(mols)244        maccs = fill(maccs, ~clean_mol_mask)245        print("Created MACCS keys")246 247    if "rdkit_descrs" in descriptors:248        rdkit_descrs = create_rdkit_descriptors(mols)249        rdkit_descrs = fill(rdkit_descrs, ~clean_mol_mask)250        print("Created RDKit descriptors")251 252    # concatenate features253    features = {}254    for descr in descriptors:255        features[descr] = vars()[descr]256 257    return features, clean_mol_mask258 259 260def get_tox21_split(token, cvfold=None):261    """Retrieve Tox21 splits from HuggingFace with respect to given cvfold."""262    ds = load_dataset("ml-jku/tox21", token=token)263 264    train_df = ds["train"].to_pandas()265    val_df = ds["validation"].to_pandas()266 267    if cvfold is None:268        return {"train": train_df, "validation": val_df}269 270    combined_df = pd.concat([train_df, val_df], ignore_index=True)271    cvfold = float(cvfold)272 273    # create new splits274    cvfold = float(cvfold)275    train_df = combined_df[combined_df.CVfold != cvfold]276    val_df = combined_df[combined_df.CVfold == cvfold]277 278    # exclude train mols that occur in the validation split279    val_inchikeys = set(val_df["inchikey"])280    train_df = train_df[~train_df["inchikey"].isin(val_inchikeys)]281 282    return {283        "train": train_df.reset_index(drop=True),284        "validation": val_df.reset_index(drop=True),285    }286