leak_proof_cmap_00_define_splits.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright © The University of Edinburgh, 2024.\n",
    "\n",
    "Development has been supported by GSK.\n",
    "\n",
    "# Setup\n",
    "Load CMAP data from Phenonaut objects and define directories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "import pandas as pd\n",
    "import leakproofcmap\n",
    "working_dir=Path(\"working_dir\")\n",
    "packaged_dataset_dir=Path(\"/local_scratch/data/phenonaut_datasets/cmap/\")\n",
    "pickle_dir=(working_dir/Path(\"pickles\")).resolve()\n",
    "split_data=(working_dir/Path(\"split_data\")).resolve()\n",
    "if not split_data.exists():\n",
    "    split_data.mkdir(parents=True)\n",
    "\n",
    "phe=leakproofcmap.get_cmap_phenonaut_object(phenonaut_packaged_dataset_dir=packaged_dataset_dir,working_dir=working_dir,pickle_dir=pickle_dir)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Write out MOA and pert_iname counts\n",
    "\n",
    "This is purely for accounting and seeing how many of each we have\n",
    "\n",
    "Writes 'moa_counts.tsv' and 'pert_iname_counts.tsv' to working_dir/split_data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp_df=phe.df.moa.value_counts().to_frame(\"count\")\n",
    "tmp_df.index.name = 'moa'\n",
    "tmp_df.to_csv(split_data/\"moa_counts.tsv\", sep=\"\\t\")\n",
    "del tmp_df\n",
    "tmp_df=phe.df.pert_iname.value_counts().to_frame(\"counts\")\n",
    "tmp_df.index.name = 'pert_iname'\n",
    "tmp_df.to_csv(split_data/\"pert_iname_counts.tsv\", sep=\"\\t\")\n",
    "del tmp_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Display the number of cell lines assessed against each MOA\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "moas_in_cell_line_counts = {moa:len(phe.df.query(\"moa == @moa\")['cell_id'].unique()) for moa in tqdm(phe.df.moa.unique())}\n",
    "print({k:v for k,v in sorted(moas_in_cell_line_counts.items(), key=lambda x: -x[1])})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate cell line splits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Write out splits based on cell line response to median moa perturbations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import phenonaut\n",
    "from phenonaut.transforms.imputers import KNNImputer\n",
    "\n",
    "cell_line_to_moa_median_responses_df_file=split_data/\"cell_line_to_moa_median_imputed.tsv\"\n",
    "\n",
    "if not cell_line_to_moa_median_responses_df_file.exists():\n",
    "    print(\"Doesnt exists\")\n",
    "    cell_line_to_moa_df_with_nans_file=split_data/\"cell_line_to_moa_median_with_nans.tsv\"\n",
    "    if not cell_line_to_moa_df_with_nans_file.exists():\n",
    "        counts_df=phe.df.groupby([\"cell_id\", \"moa\"])['DDR1'].count().reset_index().rename(columns={'DDR1':'count'}).pivot(index='cell_id', columns=['moa'], values='count')\n",
    "        counts_df.fillna(0).astype(int).to_csv(split_data/\"cell_line_to_moa_count.tsv\", sep=\"\\t\")\n",
    "\n",
    "        cell_line_to_moa_df=pd.DataFrame(columns=[f\"{moa}_{f}\" for moa in counts_df.columns for f in phe.ds.features], index=counts_df.index)\n",
    "        for cell_line in counts_df.index:\n",
    "            groups=phe.df.query(\"cell_id==@cell_line\").groupby('moa')\n",
    "            for moa, g_df in groups:\n",
    "                cell_line_to_moa_df.loc[cell_line, [f\"{moa}_{f}\" for f in phe.ds.features]]=g_df[phe.ds.features].median().values\n",
    "        cell_line_to_moa_df.to_csv(cell_line_to_moa_df_with_nans_file, sep=\"\\t\")\n",
    "    else:\n",
    "        print(\"Found file, reading\")\n",
    "        cell_line_to_moa_df=pd.read_csv(cell_line_to_moa_df_with_nans_file, sep=\"\\t\", index_col=[0])\n",
    "    counts_df=phe.df.groupby([\"cell_id\", \"moa\"])['DDR1'].count().reset_index().rename(columns={'DDR1':'count'}).pivot(index='cell_id', columns=['moa'], values='count')\n",
    "    phe_cell_line_to_moa_response=phenonaut.Phenonaut(phenonaut.data.Dataset(\"cell_line_to_moa_response\", input_file_path_or_df=cell_line_to_moa_df, features=[f\"{m}_{f}\" for f in phe.ds.features for m in counts_df.columns]))\n",
    "    imputer=KNNImputer()\n",
    "    print(\"Running imputation\")\n",
    "    imputer(phe_cell_line_to_moa_response.ds)\n",
    "    tmp_df=phe_cell_line_to_moa_response.df[phe_cell_line_to_moa_response.ds.features]\n",
    "    tmp_df.to_csv(cell_line_to_moa_median_responses_df_file, sep=\"\\t\")\n",
    "    del tmp_df\n",
    "    del phe_cell_line_to_moa_response\n",
    "\n",
    "cell_line_to_moa_median_responses_df=pd.read_csv(cell_line_to_moa_median_responses_df_file, sep=\"\\t\", index_col=[0])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define function to perform diverse splits\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numpy.random import default_rng\n",
    "from scipy.spatial.distance import pdist, squareform\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "def distribute_data_into_splits(response_df:pd.DataFrame, n_splits:int=5, random_state:int=7)->pd.DataFrame:\n",
    "    np_rng=np.random.default_rng(random_state)\n",
    "    splits=[[] for _ in range(n_splits)]\n",
    "    distances=pd.DataFrame(squareform(pdist(response_df, metric='euclidean')), index=response_df.index, columns=response_df.index)\n",
    "    distances.values[np.eye(len(distances), dtype=bool)]=np.nan\n",
    "    # Add most diverse to splits\n",
    "    for split, md in zip(splits, np_rng.permutation(distances.sum(axis=0).sort_values(ascending=False).index[:n_splits])):\n",
    "        split.append(md)\n",
    "        distances.loc[:,md]=np.nan\n",
    "\n",
    "    # Add the first most similar sample to the split (dealing with closest to split[:, 0])\n",
    "    for split in splits:\n",
    "        min_dist_index=np.nanargmin(distances.loc[split[0]].values)\n",
    "        split.append(distances.index.tolist()[min_dist_index])\n",
    "        distances.loc[:, distances.index.tolist()[min_dist_index]]=np.nan\n",
    "\n",
    "    cur_split_idx=0\n",
    "    # As above, add most similar to any of the others in the split\n",
    "    while len([item for split in splits for item in split])<len(response_df):\n",
    "        split_dists=distances.loc[splits[cur_split_idx],:]\n",
    "        min_dist_idx=np.unravel_index(np.nanargmin(distances.loc[splits[cur_split_idx],:]),split_dists.shape)\n",
    "        splits[cur_split_idx].append(split_dists.columns.tolist()[min_dist_idx[1]])\n",
    "        distances.loc[:,split_dists.columns.tolist()[min_dist_idx[1]]]=np.nan\n",
    "        cur_split_idx+=1\n",
    "        if cur_split_idx==len(splits):\n",
    "            cur_split_idx=0\n",
    "    return pd.DataFrame(splits, index=[f'split_{n}' for n in range(1, n_splits+1)])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perform splitting and write to file\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.decomposition import PCA\n",
    "from sklearn.pipeline import Pipeline\n",
    "pipeline = Pipeline([('scaling', StandardScaler()), ('pca', PCA(n_components=None))])\n",
    "\n",
    "cell_line_to_moa_median_responses_df_stdscale_pca=pd.DataFrame(pipeline.fit_transform(cell_line_to_moa_median_responses_df), index=cell_line_to_moa_median_responses_df.index)\n",
    "\n",
    "splits_by_cell_line_based_on_moa_df=distribute_data_into_splits(cell_line_to_moa_median_responses_df_stdscale_pca, 5)\n",
    "splits_by_cell_line_based_on_moa_df.to_csv(split_data/\"cmap_splits_by_cell_line_based_on_moa.tsv\", sep=\"\\t\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate MOA splits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Gather info on cell line to (MOA), define cell line splits by MOA\n",
    "\n",
    "import phenonaut\n",
    "from phenonaut.transforms.imputers import KNNImputer, RFImputer\n",
    "\n",
    "moa_to_cell_line_median_responses_df_file=split_data/\"moa_to_cell_line_median_imputed.tsv\"\n",
    "if not moa_to_cell_line_median_responses_df_file.exists():\n",
    "    moa_to_cell_line_df_with_nans_file=split_data/\"moa_to_cell_line_median_with_nans.tsv\"\n",
    "    if not moa_to_cell_line_df_with_nans_file.exists():\n",
    "        print(f\"Generating {moa_to_cell_line_median_responses_df_file}\")\n",
    "        counts_df=phe.df.groupby([\"cell_id\", \"moa\"])['DDR1'].count().reset_index().rename(columns={'DDR1':'count'}).pivot(index='moa', columns=['cell_id'], values='count')\n",
    "        counts_df.fillna(0).astype(int).to_csv(split_data/\"moa_to_cell_line_count.tsv\", sep=\"\\t\")\n",
    "\n",
    "        moa_to_cell_line_df=pd.DataFrame(columns=[f\"{cl}_{f}\" for cl in counts_df.columns for f in phe.ds.features], index=counts_df.index)\n",
    "        for moa in counts_df.index:\n",
    "            groups=phe.df.query(\"moa==@moa\").groupby('cell_id')\n",
    "            for cell_line, g_df in groups:\n",
    "                moa_to_cell_line_df.loc[moa, [f\"{cell_line}_{f}\" for f in phe.ds.features]]=g_df[phe.ds.features].median().values\n",
    "            moa_to_cell_line_df.dropna(axis=1, how='all')\n",
    "        moa_to_cell_line_df.to_csv(moa_to_cell_line_df_with_nans_file, sep=\"\\t\")\n",
    "    else:\n",
    "        print(\"Found file, reading\")\n",
    "        moa_to_cell_line_df=pd.read_csv(moa_to_cell_line_df_with_nans_file, sep=\"\\t\", index_col=[0])\n",
    "    phe_cell_line_to_moa_response=phenonaut.Phenonaut(phenonaut.data.Dataset(\"moa_to_cell_line_response\", input_file_path_or_df=moa_to_cell_line_df, features=[f\"{cl}_{f}\" for cl in counts_df.columns for f in phe.ds.features]))\n",
    "    imputer=KNNImputer()\n",
    "    print(\"Running imputation\")\n",
    "    imputer(phe_cell_line_to_moa_response.ds)\n",
    "    tmp_df=phe_cell_line_to_moa_response.df[phe_cell_line_to_moa_response.ds.features]\n",
    "    tmp_df.to_csv(moa_to_cell_line_median_responses_df_file, sep=\"\\t\")\n",
    "moa_to_cell_line_median_responses_df=pd.read_csv(moa_to_cell_line_median_responses_df_file, sep=\"\\t\", index_col=[0])\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perform splitting and write to file\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.decomposition import PCA\n",
    "from sklearn.pipeline import Pipeline\n",
    "print(moa_to_cell_line_median_responses_df.shape)\n",
    "moa_to_cell_line_median_responses_df_no_dmso=moa_to_cell_line_median_responses_df.drop(index='control vehicle')\n",
    "print(moa_to_cell_line_median_responses_df_no_dmso.shape)\n",
    "pipeline = Pipeline([('scaling', StandardScaler()), ('pca', PCA(n_components=None))])\n",
    "\n",
    "moa_to_cell_line_median_responses_df_stdscale_pca=pd.DataFrame(pipeline.fit_transform(moa_to_cell_line_median_responses_df_no_dmso), index=moa_to_cell_line_median_responses_df_no_dmso.index)\n",
    "\n",
    "splits_by_moa_based_on_cell_line_df=distribute_data_into_splits(moa_to_cell_line_median_responses_df_stdscale_pca, 5)\n",
    "splits_by_moa_based_on_cell_line_df.to_csv(split_data/\"cmap_splits_by_moa_based_on_cell_line.tsv\", sep=\"\\t\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Write split data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "import pandas as pd\n",
    "import leakproofcmap\n",
    "\n",
    "working_dir=Path(\"working_dir\")\n",
    "packaged_dataset_dir=Path(\"/local_scratch/data/phenonaut_datasets/cmap/\")\n",
    "pickle_dir=(working_dir/Path(\"pickles\")).resolve()\n",
    "split_data=(working_dir/Path(\"split_data\")).resolve()\n",
    "if not split_data.exists():\n",
    "    split_data.mkdir(parents=True)\n",
    "\n",
    "phe=leakproofcmap.get_cmap_phenonaut_object(phenonaut_packaged_dataset_dir=packaged_dataset_dir,working_dir=working_dir,pickle_dir=pickle_dir)\n",
    "\n",
    "split_data=working_dir/\"split_data\"\n",
    "if not split_data.exists():\n",
    "    split_data.mkdir(parents=True)\n",
    "for cell_id_split_i in range(1, 5 + 1):\n",
    "    for moa_split_i in range(1, 5 + 1):\n",
    "        name=f\"cellidsplit{cell_id_split_i}_moasplit{moa_split_i}\"\n",
    "        split = leakproofcmap.CMAPSplit.from_splits_tsv(\n",
    "            cell_id_split_file=split_data / \"cmap_splits_by_cell_line_based_on_moa.tsv\",\n",
    "            moa_split_file=split_data / \"cmap_splits_by_moa_based_on_cell_line.tsv\",\n",
    "            cell_line_split_number=cell_id_split_i,\n",
    "            moa_split_number=moa_split_i,df=phe.df, features=phe.ds.features,\n",
    "            name=name,\n",
    "            np_rng_seed=42,\n",
    "            is_scalar_and_pca_target=True if (cell_id_split_i==1 and moa_split_i==1) else False,\n",
    "        )\n",
    "        split.save(split_data/f\"cmap_split_{name}.json\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Write split cpd and moa counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "ds_counts_df=pd.DataFrame(columns=[\"n_records\",\"unique_moas\",\"unique_compounds\"], index=pd.MultiIndex.from_product([range(1,6), range(1,6), range(1,6), [\"fold_type\"]], names=['cli', 'moai', 'fold_i', 'fold_type'])).sort_index()\n",
    "\n",
    "for cell_line_i in range(1,6):\n",
    "    for moa_i in range(1,6):\n",
    "        for fold_i in range(1,6):\n",
    "            split_object=leakproofcmap.CMAPSplit.load(split_data/Path(f\"cmap_split_cellidsplit{cell_line_i}_moasplit{moa_i}.json\"))\n",
    "            df, _ = split_object.get_df(phe.df, phe.ds.features, False, False, \"train\", fold_number=fold_i)\n",
    "            ds_counts_df.loc[(cell_line_i,moa_i, fold_i, \"train\"),\"n_records\"]=df.shape[0]\n",
    "            ds_counts_df=ds_counts_df.sort_index()\n",
    "\n",
    "            ds_counts_df.loc[(cell_line_i,moa_i, fold_i, \"train\"),\"unique_moas\"]=len(df.moa.unique())\n",
    "            ds_counts_df=ds_counts_df.sort_index()\n",
    "            ds_counts_df.loc[(cell_line_i,moa_i, fold_i, \"train\"),\"unique_compounds\"]=len(df.pert_iname.unique())\n",
    "            ds_counts_df=ds_counts_df.sort_index()\n",
    "            df, _ = split_object.get_df(phe.df, phe.ds.features, False, False, \"val\", fold_number=fold_i)\n",
    "            ds_counts_df.loc[(cell_line_i,moa_i, fold_i, \"val\"),\"n_records\"]=df.shape[0]\n",
    "            ds_counts_df=ds_counts_df.sort_index()\n",
    "            ds_counts_df.loc[(cell_line_i,moa_i, fold_i, \"val\"),\"unique_moas\"]=len(df.moa.unique())\n",
    "            ds_counts_df=ds_counts_df.sort_index()\n",
    "            ds_counts_df.loc[(cell_line_i,moa_i, fold_i, \"val\"),\"unique_compounds\"]=len(df.pert_iname.unique())\n",
    "        df, _ = split_object.get_df(phe.df, phe.ds.features, False, False, \"test\")\n",
    "        ds_counts_df.loc[(cell_line_i,moa_i, \"NA\", \"test\"),\"n_records\"]=df.shape[0]\n",
    "        ds_counts_df=ds_counts_df.sort_index()\n",
    "        ds_counts_df.loc[(cell_line_i,moa_i, \"NA\", \"test\"),\"unique_moas\"]=len(df.moa.unique())\n",
    "        ds_counts_df=ds_counts_df.sort_index()\n",
    "        ds_counts_df.loc[(cell_line_i,moa_i, \"NA\", \"test\"),\"unique_compounds\"]=len(df.pert_iname.unique())\n",
    "        ds_counts_df=ds_counts_df.sort_index()\n",
    "ds_counts_df=ds_counts_df.dropna(how='any', axis=0)\n",
    "print(ds_counts_df)\n",
    "ds_counts_df.to_csv(split_data/\"split_info.tsv\", sep=\"\\t\")\n",
    "ds_counts_df.to_csv(split_data/\"split_info.csv\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}