project/04_1_train_pimms_models.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "eae0a078",
   "metadata": {},
   "source": [
    "# PIMMS Tutorial: Scikit-learn style transformers\n",
    "\n",
    "1. Load data into pandas dataframe\n",
    "2. Fit model on training data, potentially specify validation data\n",
    "3. Impute only missing values with predictions from model\n",
    "\n",
    "Autoencoders need wide training data, i.e. a sample with all its features' intensities, whereas\n",
    "Collaborative Filtering needs long training data, i.e. sample identifier a feature identifier and the intensity.\n",
    "Both data formats can be transformed into each other, but models using long data format do not need to\n",
    "take care of missing values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0650846",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from importlib import metadata\n",
    "\n",
    "IN_COLAB = 'COLAB_GPU' in os.environ\n",
    "if IN_COLAB:\n",
    "    try:\n",
    "        _v = metadata.version('pimms-learn')\n",
    "        print(f\"Running in colab and pimms-learn ({_v}) is installed.\")\n",
    "    except metadata.PackageNotFoundError:\n",
    "        print(\"Install PIMMS...\")\n",
    "        # !pip install git+https://github.com/RasmussenLab/pimms.git\n",
    "        !pip install pimms-learn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36b5a27d",
   "metadata": {},
   "source": [
    "If on colab, please restart the environment and run everything from here on.\n",
    "\n",
    "Specify example data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65b5bdaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "IN_COLAB = 'COLAB_GPU' in os.environ\n",
    "\n",
    "fn_intensities = 'data/dev_datasets/HeLa_6070/protein_groups_wide_N50.csv'\n",
    "if IN_COLAB:\n",
    "    fn_intensities = ('https://raw.githubusercontent.com/RasmussenLab/pimms/main/'\n",
    "                      'project/data/dev_datasets/HeLa_6070/protein_groups_wide_N50.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d628ce9",
   "metadata": {},
   "source": [
    "Load package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c289d17",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "import random\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from IPython.display import display\n",
    "\n",
    "import pimmslearn.filter\n",
    "import pimmslearn.plotting.data\n",
    "import pimmslearn.sampling\n",
    "from pimmslearn.plotting.defaults import color_model_mapping\n",
    "from pimmslearn.sklearn.ae_transformer import AETransformer\n",
    "from pimmslearn.sklearn.cf_transformer import CollaborativeFilteringTransformer\n",
    "\n",
    "pimmslearn.plotting.make_large_descriptors(8)\n",
    "\n",
    "\n",
    "logger = logger = pimmslearn.logging.setup_nb_logger()\n",
    "logging.getLogger('fontTools').setLevel(logging.ERROR)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59282a30",
   "metadata": {},
   "source": [
    "## Parameters\n",
    "Can be set by papermill on the command line or manually in the (colab) notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab9b68d6",
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "fn_intensities: str = fn_intensities  # path or url to the data file in csv format\n",
    "index_name: str = 'Sample ID'  # name of the index column\n",
    "column_name: str = 'protein group'  # name of the column index\n",
    "select_features: bool = True  # Whether to select features based on prevalence\n",
    "feat_prevalence: float = 0.2  # minimum prevalence of a feature to be included\n",
    "sample_completeness: float = 0.3  # minimum completeness of a sample to be included\n",
    "sample_splits: bool = True  # Whether to sample validation and test data\n",
    "frac_non_train: float = 0.1  # fraction of non training data (validation and test split)\n",
    "frac_mnar: float = 0.0  # fraction of missing not at random data, rest: missing completely at random\n",
    "random_state: int = 42  # random state for reproducibility"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "231b6650",
   "metadata": {},
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a3edbdd",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "df = pd.read_csv(fn_intensities, index_col=0)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "727b3ace",
   "metadata": {},
   "source": [
    "We will need the data in long format for Collaborative Filtering.\n",
    "Naming both the row and column index assures\n",
    "that the data can be transformed very easily into long format:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fde25e9",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "df.index.name = index_name  # already set\n",
    "df.columns.name = column_name  # not set due to csv disk file format\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6c788f0",
   "metadata": {},
   "source": [
    "### Data transformation: log2 transformation\n",
    "Transform the data using the logarithm, here using base 2:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "554d4fa7",
   "metadata": {
    "lines_to_next_cell": 2,
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "df = np.log2(df + 1)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbce73d1",
   "metadata": {},
   "source": [
    "### two plots inspecting data availability\n",
    "\n",
    "1. proportion of missing values per feature median (N = protein groups)\n",
    "2. CDF of available intensities per protein group"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "536793bb",
   "metadata": {
    "lines_to_next_cell": 2,
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "ax = pimmslearn.plotting.data.plot_feat_median_over_prop_missing(\n",
    "    data=df, type='boxplot')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "313bc55e",
   "metadata": {
    "lines_to_next_cell": 2,
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "df.notna().sum().sort_values().plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf54a3af",
   "metadata": {},
   "source": [
    "### Data selection\n",
    "define a minimum feature and sample frequency for a feature to be included"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28ad27de",
   "metadata": {
    "lines_to_next_cell": 2,
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "if select_features:\n",
    "    # potentially this can take a few iterations to stabilize.\n",
    "    df = pimmslearn.filter.select_features(df, feat_prevalence=feat_prevalence)\n",
    "    df = pimmslearn.filter.select_features(df=df, feat_prevalence=sample_completeness, axis=1)\n",
    "df.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "693b1ee5",
   "metadata": {},
   "source": [
    "Transform to long-data format:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "646ea5bb",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "df = df.stack().to_frame('intensity')\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54d4b781",
   "metadata": {},
   "source": [
    "## Optionally: Sample data\n",
    "- models can be trained without subsetting the data\n",
    "- allows evaluation of the models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2586a1c",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "if sample_splits:\n",
    "    splits, thresholds, fake_na_mcar, fake_na_mnar = pimmslearn.sampling.sample_mnar_mcar(\n",
    "        df_long=df,\n",
    "        frac_non_train=frac_non_train,\n",
    "        frac_mnar=frac_mnar,\n",
    "        random_state=random_state,\n",
    "    )\n",
    "    splits = pimmslearn.sampling.check_split_integrity(splits)\n",
    "else:\n",
    "    splits = pimmslearn.sampling.DataSplits(is_wide_format=False)\n",
    "    splits.train_X = df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7792ce6e",
   "metadata": {},
   "source": [
    "The resulting DataFrame with one column has an `MulitIndex` with the sample and feature identifier."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2390859",
   "metadata": {},
   "source": [
    "## Collaborative Filtering\n",
    "\n",
    "Inspect annotations of the scikit-learn like Transformer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "567854c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # CollaborativeFilteringTransformer?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a76ba4ce",
   "metadata": {},
   "source": [
    "Let's set up collaborative filtering without a validation or test set, using\n",
    "all the data there is."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b547a77",
   "metadata": {},
   "outputs": [],
   "source": [
    "cf_model = CollaborativeFilteringTransformer(\n",
    "    target_column='intensity',\n",
    "    sample_column='Sample ID',\n",
    "    item_column='protein group',\n",
    "    out_folder='runs/scikit_interface')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f86364d4",
   "metadata": {},
   "source": [
    "We use `fit` and `transform` to train the model and impute the missing values.\n",
    "> Scikit learns interface requires a `X` and `y`. `y` is the validation data in our context.\n",
    "> We might have to change the interface to allow usage within pipelines (-> `y` is not needed).\n",
    "> This will probably mean setting up a validation set within the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb5ac432",
   "metadata": {},
   "outputs": [],
   "source": [
    "cf_model.fit(splits.train_X,\n",
    "             splits.val_y,\n",
    "             cuda=False,\n",
    "             epochs_max=20,\n",
    "             )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7db2f4c4",
   "metadata": {},
   "source": [
    "Impute missing values usin `transform` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3dac537",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_imputed = cf_model.transform(df).unstack()\n",
    "assert df_imputed.isna().sum().sum() == 0\n",
    "df_imputed.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b94ba21b",
   "metadata": {},
   "source": [
    "Let's plot the distribution of the imputed values vs the ones used for training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99ff7ecd",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "df_imputed = df_imputed.stack()  # long-format\n",
    "observed = df_imputed.loc[df.index]\n",
    "imputed = df_imputed.loc[df_imputed.index.difference(df.index)]\n",
    "df_imputed = df_imputed.unstack()  # back to wide-format\n",
    "# some checks\n",
    "assert len(df) == len(observed)\n",
    "assert df_imputed.shape[0] * df_imputed.shape[1] == len(imputed) + len(observed)\n",
    "\n",
    "fig, axes = plt.subplots(2, figsize=(8, 4))\n",
    "\n",
    "min_max = pimmslearn.plotting.data.get_min_max_iterable(\n",
    "    [observed, imputed])\n",
    "label_template = '{method} (N={n:,d})'\n",
    "ax, _ = pimmslearn.plotting.data.plot_histogram_intensities(\n",
    "    observed,\n",
    "    ax=axes[0],\n",
    "    min_max=min_max,\n",
    "    label=label_template.format(method='measured',\n",
    "                                n=len(observed),\n",
    "                                ),\n",
    "    color='grey',\n",
    "    alpha=1)\n",
    "_ = ax.legend()\n",
    "ax, _ = pimmslearn.plotting.data.plot_histogram_intensities(\n",
    "    imputed,\n",
    "    ax=axes[1],\n",
    "    min_max=min_max,\n",
    "    label=label_template.format(method='CF imputed',\n",
    "                                n=len(imputed),\n",
    "                                ),\n",
    "    color=color_model_mapping['CF'],\n",
    "    alpha=1)\n",
    "_ = ax.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6d6552c",
   "metadata": {},
   "source": [
    "## AutoEncoder architectures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7184c2e",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# Use wide format of data\n",
    "splits.to_wide_format()\n",
    "splits.train_X"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae52c6fe",
   "metadata": {},
   "source": [
    "Validation data for early stopping (if specified)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bbd0017",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "splits.val_y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6da6c4e2",
   "metadata": {},
   "source": [
    "Training and validation need the same shape:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99e690f9",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "if splits.val_y is not None:\n",
    "    splits.val_y = pd.DataFrame(pd.NA, index=splits.train_X.index,\n",
    "                                columns=splits.train_X.columns).fillna(splits.val_y)\n",
    "\n",
    "    print(splits.train_X.shape, splits.val_y.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45b9c22c",
   "metadata": {},
   "source": [
    "Select either `DAE` or `VAE` model by chance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "873f3668",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_selected = random.choice(['DAE', 'VAE'])\n",
    "print(\"Selected model by chance:\", model_selected)\n",
    "model = AETransformer(\n",
    "    model=model_selected,\n",
    "    hidden_layers=[512,],\n",
    "    latent_dim=50,\n",
    "    out_folder='runs/scikit_interface',\n",
    "    batch_size=10,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aea5dc44",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.fit(splits.train_X, splits.val_y,\n",
    "          epochs_max=50,\n",
    "          cuda=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f89fb41f",
   "metadata": {},
   "source": [
    "Impute missing values using `transform` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62cd0721",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_imputed = model.transform(splits.train_X).stack()\n",
    "df_imputed"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17398941",
   "metadata": {},
   "source": [
    "Evaluate the model using the validation data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a664072",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "if splits.val_y is not None:\n",
    "    pred_val = splits.val_y.stack().to_frame('observed')\n",
    "    pred_val[model_selected] = df_imputed\n",
    "    val_metrics = pimmslearn.models.calculte_metrics(pred_val, 'observed')\n",
    "    display(val_metrics)\n",
    "\n",
    "    fig, ax = plt.subplots(figsize=(8, 2))\n",
    "\n",
    "    ax, errors_binned = pimmslearn.plotting.errors.plot_errors_by_median(\n",
    "        pred=pred_val,\n",
    "        target_col='observed',\n",
    "        feat_medians=splits.train_X.median(),\n",
    "        ax=ax,\n",
    "        metric_name='MAE',\n",
    "        palette=color_model_mapping)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e215dba2",
   "metadata": {},
   "source": [
    "replace predicted values with validation data values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f89799e8",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "splits.to_long_format()\n",
    "df_imputed = df_imputed.replace(splits.val_y).replace(splits.test_y)\n",
    "df_imputed"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "506de661",
   "metadata": {},
   "source": [
    "Plot the distribution of the imputed values vs the observed data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "235fdb66",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "observed = df_imputed.loc[df.index].squeeze()\n",
    "imputed = df_imputed.loc[df_imputed.index.difference(df.index)].squeeze()\n",
    "\n",
    "fig, axes = plt.subplots(2, figsize=(8, 4))\n",
    "\n",
    "min_max = pimmslearn.plotting.data.get_min_max_iterable([observed, imputed])\n",
    "\n",
    "label_template = '{method} (N={n:,d})'\n",
    "ax, _ = pimmslearn.plotting.data.plot_histogram_intensities(\n",
    "    observed,\n",
    "    ax=axes[0],\n",
    "    min_max=min_max,\n",
    "    label=label_template.format(method='measured',\n",
    "                                n=len(observed),\n",
    "                                ),\n",
    "    color='grey',\n",
    "    alpha=1)\n",
    "_ = ax.legend()\n",
    "ax, _ = pimmslearn.plotting.data.plot_histogram_intensities(\n",
    "    imputed,\n",
    "    ax=axes[1],\n",
    "    min_max=min_max,\n",
    "    label=label_template.format(method=f'{model_selected} imputed',\n",
    "                                n=len(imputed),\n",
    "                                ),\n",
    "    color=color_model_mapping[model_selected],\n",
    "    alpha=1)\n",
    "_ = ax.legend()"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  },
  "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.8.17"
  },
  "mystnb": {
   "execution_raise_on_error": true,
   "execution_timeout": 120
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}