src/data.py

"""
This files includes data processing tools.
"""

import os
import argparse
import json
from typing import Iterable, Literal

import numpy as np
import pandas as pd

from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.preprocessing import StandardScaler
from sklearn.feature_selection import VarianceThreshold
from statsmodels.distributions.empirical_distribution import ECDF
from datasets import load_dataset

import torch

from rdkit import Chem, DataStructs
from rdkit.Chem import Descriptors, rdFingerprintGenerator, MACCSkeys
from rdkit.Chem.rdchem import Mol

from src.utils import (
    TASKS,
    HF_TOKEN,
    USED_200_DESCR,
    Standardizer,
    load_pickle,
    write_pickle,
    KNOWN_DESCR,
)

class SquashScaler(TransformerMixin, BaseEstimator):
    """
    Scaler that performs sequential standardization, nonlinearity (tanh), and
    re-standardization. Inspired by DeepTox (Mayr et al., 2016)
    """

    def __init__(self):
        self.scaler1 = StandardScaler()
        self.scaler2 = StandardScaler()

    def fit(self, X):
        _X = X.copy()
        _X = self.scaler1.fit_transform(_X)
        _X = np.tanh(_X)
        _X = self.scaler2.fit(_X)
        self.is_fitted_ = True
        return self

    def transform(self, X):
        _X = X.copy()
        _X = self.scaler1.transform(_X)
        _X = np.tanh(_X)
        return self.scaler2.transform(_X)


def create_cleaned_mol_objects(smiles: list[str]) -> tuple[list[Mol], np.ndarray]:
    """This function creates cleaned RDKit mol objects from a list of SMILES.

    Args:
        smiles (list[str]): list of SMILES

    Returns:
        list[Mol]: list of cleaned molecules
        np.ndarray[bool]: mask that contains False at index `i`, if molecule in `smiles` at
            index `i` could not be cleaned and was removed.
    """
    sm = Standardizer(canon_taut=True)

    clean_mol_mask = list()
    mols = list()
    for i, smile in enumerate(smiles):
        mol = Chem.MolFromSmiles(smile)
        standardized_mol, _ = sm.standardize_mol(mol)
        is_cleaned = standardized_mol is not None
        clean_mol_mask.append(is_cleaned)
        if not is_cleaned:
            continue
        can_mol = Chem.MolFromSmiles(Chem.MolToSmiles(standardized_mol))
        mols.append(can_mol)

    return mols, np.array(clean_mol_mask)


def create_ecfp_fps(mols: list[Mol], radius=None, fpsize=None) -> np.ndarray:
    """This function ECFP fingerprints for a list of molecules.

    Args:
        mols (list[Mol]): list of molecules

    Returns:
        np.ndarray: ECFP fingerprints of molecules
    """
    ecfps = list()

    kwargs = {}
    if not fpsize is None:
        kwargs["fpSize"] = fpsize
    if not radius is None:
        kwargs["radius"] = radius
    for mol in mols:
        gen = rdFingerprintGenerator.GetMorganGenerator(countSimulation=True, **kwargs)
        fp_sparse_vec = gen.GetCountFingerprint(mol)

        fp = np.zeros((0,), np.int8)
        DataStructs.ConvertToNumpyArray(fp_sparse_vec, fp)

        ecfps.append(fp)

    return np.array(ecfps)


def create_maccs_keys(mols: list[Mol]) -> np.ndarray:
    maccs = [MACCSkeys.GenMACCSKeys(x) for x in mols]
    return np.array(maccs)


def get_tox_patterns(filepath: str):
    """This calculates tox features defined in tox_smarts.json.
    Args:
        mols: A list of Mol
        n_jobs: If >1 multiprocessing is used
    """
    # load patterns
    with open(filepath) as f:
        smarts_list = [s[1] for s in json.load(f)]

    # Code does not work for this case
    assert len([s for s in smarts_list if ("AND" in s) and ("OR" in s)]) == 0

    # Chem.MolFromSmarts takes a long time so it pays of to parse all the smarts first
    # and then use them for all molecules. This gives a huge speedup over existing code.
    # a list of patterns, whether to negate the match result and how to join them to obtain one boolean value
    all_patterns = []
    for smarts in smarts_list:
        patterns = []  # list of smarts-patterns
        # value for each of the patterns above. Negates the values of the above later.
        negations = []

        if " AND " in smarts:
            smarts = smarts.split(" AND ")
            merge_any = False  # If an ' AND ' is found all 'subsmarts' have to match
        else:
            # If there is an ' OR ' present it's enough is any of the 'subsmarts' match.
            # This also accumulates smarts where neither ' OR ' nor ' AND ' occur
            smarts = smarts.split(" OR ")
            merge_any = True

        # for all subsmarts check if they are preceded by 'NOT '
        for s in smarts:
            neg = s.startswith("NOT ")
            if neg:
                s = s[4:]
            patterns.append(Chem.MolFromSmarts(s))
            negations.append(neg)

        all_patterns.append((patterns, negations, merge_any))
    return all_patterns


def create_tox_features(mols: list[Mol], patterns: list) -> np.ndarray:
    """Matches the tox patterns against a molecule. Returns a boolean array"""
    tox_data = []
    for mol in mols:
        mol_features = []
        for patts, negations, merge_any in patterns:
            matches = [mol.HasSubstructMatch(p) for p in patts]
            matches = [m != n for m, n in zip(matches, negations)]
            if merge_any:
                pres = any(matches)
            else:
                pres = all(matches)
            mol_features.append(pres)

        tox_data.append(np.array(mol_features))

    return np.array(tox_data)


def create_rdkit_descriptors(mols: list[Mol]) -> np.ndarray:
    """This function creates RDKit descriptors for a list of molecules.

    Args:
        mols (list[Mol]): list of molecules

    Returns:
        np.ndarray: RDKit descriptors of molecules
    """
    rdkit_descriptors = list()

    for mol in mols:
        descrs = []
        for _, descr_calc_fn in Descriptors._descList:
            descrs.append(descr_calc_fn(mol))

        descrs = np.array(descrs)
        descrs = descrs[USED_200_DESCR]
        rdkit_descriptors.append(descrs)

    return np.array(rdkit_descriptors)


def create_quantiles(raw_features: np.ndarray, ecdfs: list) -> np.ndarray:
    """Create quantile values for given features using the columns

    Args:
        raw_features (np.ndarray): values to put into quantiles
        ecdfs (list): ECDFs to use

    Returns:
        np.ndarray: computed quantiles
    """
    quantiles = np.zeros_like(raw_features)

    for column in range(raw_features.shape[1]):
        raw_values = raw_features[:, column].reshape(-1)
        ecdf = ecdfs[column]
        q = ecdf(raw_values)
        quantiles[:, column] = q

    return quantiles


def fill(features, mask, value=np.nan):
    n_mols = len(mask)
    n_features = features.shape[1]

    data = np.zeros(shape=(n_mols, n_features))
    data.fill(value)
    data[~mask] = features
    return data


def get_descriptor_dataset(
    data_path: str,
    descriptors: Iterable[str] | Literal["all"],
    scaler=None,
    save_scaler_path: str = "data/scaler.pkl",
    verbose=True,
    normalize: str = "standard",
):
    if descriptors == "all":
        descriptors = KNOWN_DESCR

    assert isinstance(descriptors, Iterable), "Passed descriptors are not iterable!"
    assert all(
        [descr in KNOWN_DESCR for descr in descriptors]
    ), f"Passed descriptors contains unknown descriptor types. Allowed descriptors: {KNOWN_DESCR}"
    print(f"Load: {data_path}")
    datafile = np.load(data_path)

    if not isinstance(datafile, np.ndarray):
        # concatenate all descriptors and normalize

        if "features" in datafile:
            data = datafile["features"]
            print("Features are already concatenated")
        else:
            data = np.concatenate([datafile[descr] for descr in descriptors], axis=1)
            print(f"Concatenated features with order: {descriptors}")
        labels = datafile["labels"]

    else:
        print("NPY file passed, cannot select specific descriptors")
        data, labels = datafile[:, :-12], datafile[:, -12:]

    if normalize != "none":
        data, scaler = normalize_features(
            data,
            scaler=scaler,
            save_scaler_path=save_scaler_path,
            verbose=verbose,
            normalization=normalize,
        )

    # filter out unsanitized molecules
    mask = ~np.isnan(data).any(axis=1)
    data = data[mask]
    labels = labels[mask]

    assert data.shape[0] == labels.shape[0], (
        f"Mismatch between data and labels: "
        f"data has {data.shape[0]} samples, but labels has {labels.shape[0]} samples."
    )

    return (data, labels, scaler)


def get_torch_descriptor_dataset(
    data_path: str,
    descriptors: list[str],
    scaler=None,
    save_scaler_path: str = "data/scaler.pkl",
    nan_to_num: int = -100,
    verbose=True,
    normalize: str = "standard",
) -> torch.utils.data.TensorDataset:
    data, labels, scaler = get_descriptor_dataset(
        data_path,
        descriptors,
        scaler,
        save_scaler_path,
        verbose=verbose,
        normalize=normalize,
    )

    labels = np.nan_to_num(labels, nan=nan_to_num)

    dataset = torch.utils.data.TensorDataset(
        torch.FloatTensor(data), torch.LongTensor(labels)
    )
    return dataset, scaler



def get_tox21_split(token="", cvfold=None):
    """Retrieve Tox21 splits from HuggingFace with respect to given cvfold."""
    ds = load_dataset("ml-jku/tox21", token=token)

    train_df = ds["train"].to_pandas()
    val_df = ds["validation"].to_pandas()

    if cvfold is None:
        return {"train": train_df, "validation": val_df}

    combined_df = pd.concat([train_df, val_df], ignore_index=True)
    cvfold = float(cvfold)

    # create new splits
    cvfold = float(cvfold)
    train_df = combined_df[combined_df.CVfold != cvfold]
    val_df = combined_df[combined_df.CVfold == cvfold]

    # exclude train mols that occur in the validation split
    val_inchikeys = set(val_df["inchikey"])
    train_df = train_df[~train_df["inchikey"].isin(val_inchikeys)]

    return {
        "train": train_df.reset_index(drop=True),
        "validation": val_df.reset_index(drop=True),
    }


def normalize_features(
    raw_features,
    scaler=None,
    save_scaler_path: str = "",
    verbose=True,
    normalization: str = "standard",
):
    if scaler is None:
        if normalization == "standard":
            scaler = StandardScaler()
        elif normalization == "squash":
            scaler = SquashScaler()
        scaler.fit(raw_features)
        if verbose:
            print("Fitted the StandardScaler")
        if save_scaler_path:
            write_pickle(save_scaler_path, scaler)
            if verbose:
                print(f"Saved the StandardScaler under {save_scaler_path}")

    # Normalize feature vectors
    normalized_features = scaler.transform(raw_features)
    if verbose:
        print("Normalized molecule features")
    return normalized_features, scaler