data_gen_microarray.py

import json
import os
import warnings
from ast import literal_eval
from collections import Counter
from pathlib import Path

import click
import joblib
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

from signature_sampling.hyperparameter_factory import SAMPLING_FACTORY
from signature_sampling.utils import fpkm, fpkm_normalised_df

warnings.filterwarnings("ignore")


@click.command()
@click.option(
    "--cv_splits_path",
    required=True,
    type=click.Path(path_type=Path, exists=True),
    help="Path to json file containing cross-validation splits wrt indices.",
)
@click.option(
    "--ref_df_path",
    required=True,
    type=click.Path(path_type=Path, exists=True),
    help="Path to dataset to be split and augmented in csv format.",
)
@click.option(
    "--ref_labels_path",
    required=True,
    type=click.Path(path_type=Path, exists=True),
    help="Path to labels associated with the dataset in csv format.",
)
@click.option(
    "--probemap_path",
    required=True,
    type=click.Path(path_type=Path, exists=True),
    help="Path to probemap file that maps genes to ensembl ids and specifies gene lengths.",
)
@click.option(
    "--signature_path",
    required=True,
    type=click.Path(path_type=Path, exists=True),
    help="Path to file containing info about signature genes and their subtypes.",
)
@click.option(
    "--sampling_method",
    required=True,
    type=click.Choice(sorted(SAMPLING_FACTORY.keys())),
    help="Sampling method to use for augmentation.",
)
@click.option("--class_size", type=str, help="Total desired size for each class.")
@click.option("--target", type=str, help="Target variable name.")
@click.option(
    "--validation_size",
    type=float,
    help="Size of validation set when splitting train set for hyperparam tuning.",
)
@click.option("--seed", type=int, help="Seed for train-test split function.")
@click.option(
    "--save_dir",
    type=click.Path(path_type=Path),
    help="Directory in which to save the augmented data.",
)
def main(
    cv_splits_path: Path,
    ref_df_path: Path,
    ref_labels_path: Path,
    probemap_path: Path,
    signature_path: Path,
    sampling_method: str,
    class_size: int,
    target: str,
    validation_size: float,
    seed: int,
    save_dir: Path,
):

    if class_size == "max":
        class_size = None
    else:
        class_size = literal_eval(class_size)
    probemap = pd.read_csv(probemap_path, sep="\t")
    probemap.index = [i.split(".")[0] for i in probemap.id]
    probemap = probemap.drop(columns="id")
    probemap["length"] = probemap.chromEnd - probemap.chromStart
    # Need to fix gene name
    probemap.loc["ENSG00000140105", "gene"] = "WARS1"
    probemap.set_index("gene", inplace=True)

    with open(cv_splits_path, "r") as f:
        cv_splits = json.load(f)

    ref_df = pd.read_csv(ref_df_path, index_col=0)
    ref_labels = pd.read_csv(ref_labels_path, index_col=0)
    assert all(ref_df.index == ref_labels.index)

    # remove ensembl version ID if present
    if any("ENS" in x for x in ref_df.columns):
        ref_df.columns = [i.split(".")[0] for i in ref_df.columns]
        gene_type = "ensemblid"
    else:
        gene_type = "SYMBOL"

    # TODO: add primary tumour check

    if os.path.exists(signature_path):
        genes = pd.read_csv(signature_path)
    else:
        genes = dict()

    signature_genes = dict(
        genes.groupby("subtype").apply(lambda x: x[gene_type].values)
    )

    splits = len(cv_splits.keys())

    for i in range(splits):
        train_idx = cv_splits[str(i)]["train"]
        train_df = ref_df.loc[train_idx, :]
        train_labels = pd.DataFrame(ref_labels.loc[train_idx, :], columns=[target])
        train_df.index.name = None
        train_labels.index.name = None
        # load and process test
        test_idx = cv_splits[str(i)]["test"]
        test_df = ref_df.loc[test_idx, :]
        test_labels = pd.DataFrame(ref_labels.loc[test_idx, :], columns=[target])
        test_df.index.name = None
        test_labels.index.name = None
        target_count = dict(Counter(train_labels[target]))
        # sampler_obj = Sampler()
        if sampling_method != "unaugmented":
            save_path = save_dir / f"{i+1}"
            if len(set(target_count.values())) == 1 and class_size is None:
                augmented_df = train_df
                augmented_labels = train_labels
            else:
                sampler = SAMPLING_FACTORY[sampling_method](sampling_method, class_size)
                sampler.init_target_signatures(signature_genes)
                sampled_df, sampled_labels = sampler.sample(
                    train_df,
                    train_labels,
                    target,
                    **{
                        "gamma_poisson_r": 5,
                        "poisson_r": 5,
                        "overlapping_genes": ["MCM2", "MKI67"],
                    },
                )
                augmented_df = pd.concat([train_df, sampled_df])
                augmented_labels = pd.concat([train_labels, sampled_labels])
                shuffle_idx = np.random.permutation(range(len(augmented_df)))

                augmented_df = augmented_df.iloc[shuffle_idx]
                augmented_labels = augmented_labels.iloc[shuffle_idx]

        else:
            augmented_df = train_df
            augmented_labels = train_labels
            save_path = save_dir.parent / f"{i+1}"

        os.makedirs(save_path, exist_ok=True)
        augmented_df.to_csv(os.path.join(save_path, "train.csv"))
        augmented_labels.to_csv(os.path.join(save_path, "train_labels.csv"))

        Xtrain, Xval, ytrain, yval = train_test_split(
            augmented_df,
            augmented_labels,
            test_size=validation_size,
            random_state=seed,
            stratify=augmented_labels,
        )

        scaler = StandardScaler()
        train_augmented_df_stdz = pd.DataFrame(
            scaler.fit_transform(Xtrain), columns=Xtrain.columns, index=Xtrain.index
        )
        valid_augmented_df_stdz = pd.DataFrame(
            scaler.transform(Xval), columns=Xval.columns, index=Xval.index
        )
        joblib.dump(scaler, save_path / "scaler.pkl")

        train_augmented_df_stdz.to_csv(save_path / "train_logrma_stdz.csv")
        valid_augmented_df_stdz.to_csv(save_path / "valid_logrma_stdz.csv")
        ytrain.to_csv(os.path.join(save_path, "train_labels_logrma.csv"))
        yval.to_csv(os.path.join(save_path, "valid_labels_logrma.csv"))

        # standardise test set based on training set
        assert all(train_augmented_df_stdz.columns == test_df.columns)
        test_stdz = pd.DataFrame(
            scaler.transform(test_df),
            columns=test_df.columns,
            index=test_df.index,
        )

        test_stdz.to_csv(os.path.join(save_path, "test_logrma_stdz.csv"))
        test_labels.to_csv(os.path.join(save_path, "test_labels_stdz.csv"))


if __name__ == "__main__":
    main()