project/01_1_train_CF.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f0e1208b-9b36-4294-b3fd-910e05a82f2e",
   "metadata": {},
   "source": [
    "# Collaborative Filtering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18b5d571-2956-4112-b22c-43d6c2146b06",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "from pprint import pprint\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "# overwriting Recorder callback with custom plot_loss\n",
    "from fastai import learner\n",
    "from fastai.collab import *\n",
    "from fastai.collab import (EarlyStoppingCallback, EmbeddingDotBias, Learner,\n",
    "                           MSELossFlat, default_device)\n",
    "from fastai.tabular.all import *\n",
    "\n",
    "import pimmslearn\n",
    "import pimmslearn.model\n",
    "import pimmslearn.models as models\n",
    "import pimmslearn.nb\n",
    "from pimmslearn.io import datasplits\n",
    "from pimmslearn.logging import setup_logger\n",
    "from pimmslearn.models import RecorderDump, plot_loss\n",
    "\n",
    "learner.Recorder.plot_loss = plot_loss\n",
    "# import fastai.callback.hook # Learner.summary\n",
    "\n",
    "\n",
    "logger = setup_logger(logger=logging.getLogger('pimmslearn'))\n",
    "logger.info(\n",
    "    \"Experiment 03 - Analysis of latent spaces and performance comparisions\")\n",
    "\n",
    "figures = {}  # collection of ax or figures"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43f062bb-8d2c-4afa-8a81-9e2a328050fe",
   "metadata": {},
   "source": [
    "Papermill script parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85c7d6f9",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# catch passed parameters\n",
    "args = None\n",
    "args = dict(globals()).keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "297f14bc-3c37-43fa-8217-f790f0593d78",
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# files and folders\n",
    "# Datasplit folder with data for experiment\n",
    "folder_experiment: str = 'runs/example'\n",
    "folder_data: str = ''  # specify data directory if needed\n",
    "file_format: str = 'csv'  # change default to pickled files\n",
    "# training\n",
    "epochs_max: int = 20  # Maximum number of epochs\n",
    "# early_stopping:bool = True # Wheather to use early stopping or not\n",
    "patience: int = 1  # Patience for early stopping\n",
    "batch_size: int = 32_768  # Batch size for training (and evaluation)\n",
    "cuda: bool = True  # Use the GPU for training?\n",
    "# model\n",
    "# Dimensionality of encoding dimension (latent space of model)\n",
    "latent_dim: int = 10\n",
    "sample_idx_position: int = 0  # position of index which is sample ID\n",
    "model: str = 'CF'  # model name\n",
    "model_key: str = 'CF'  # potentially alternative key for model (grid search)\n",
    "save_pred_real_na: bool = True  # Save all predictions for missing values"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1da8e04b-2eb2-4ecb-b1d1-2de871cd0a56",
   "metadata": {},
   "source": [
    "Some argument transformations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0746e70f-0259-48d5-90ef-25fe4b59f9ac",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "args = pimmslearn.nb.get_params(args, globals=globals())\n",
    "args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "100bbf80",
   "metadata": {
    "lines_to_next_cell": 2,
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "args = pimmslearn.nb.args_from_dict(args)\n",
    "\n",
    "# # Currently not needed -> DotProduct used, not a FNN\n",
    "# if isinstance(args.hidden_layers, str):\n",
    "#     args.overwrite_entry(\"hidden_layers\", [int(x) for x in args.hidden_layers.split('_')])\n",
    "# else:\n",
    "#     raise ValueError(f\"hidden_layers is of unknown type {type(args.hidden_layers)}\")\n",
    "args"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d59f6d8-9cb1-461a-8d62-2ab4458cab60",
   "metadata": {},
   "source": [
    "Some naming conventions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a19fe098-a029-4f71-b7fb-e652a9c16ac7",
   "metadata": {
    "lines_to_next_cell": 2,
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "TEMPLATE_MODEL_PARAMS = 'model_params_{}.json'\n",
    "\n",
    "if not args.cuda:\n",
    "    default_device(use=False)  # set to cpu"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "976cc005-0c5f-4e88-9656-c549e613ca68",
   "metadata": {},
   "source": [
    "## Load data in long format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d9cc7bd-6b6f-40b9-8db7-c8228e4b03e3",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "data = datasplits.DataSplits.from_folder(\n",
    "    args.data, file_format=args.file_format)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ca12fc2-bf34-42ac-99b2-a78ff9fe7722",
   "metadata": {},
   "source": [
    "data is loaded in long format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02bb6bf5-0eb1-4c73-9723-414b14eaf7c8",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "data.train_X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3311709",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# ! add check that specified data is available\n",
    "# silent error in fastai if e.g. target column is not available"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6045414b",
   "metadata": {},
   "source": [
    "Infer index names from long format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44958473",
   "metadata": {
    "lines_to_next_cell": 2,
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "index_columns = list(data.train_X.index.names)\n",
    "sample_id = index_columns.pop(args.sample_idx_position)\n",
    "if len(index_columns) == 1:\n",
    "    index_column = index_columns.pop()\n",
    "    index_columns = None\n",
    "    logger.info(f\"{sample_id = }, single feature: {index_column = }\")\n",
    "else:\n",
    "    logger.info(f\"{sample_id = }, multiple features: {index_columns = }\")\n",
    "\n",
    "if not index_columns:\n",
    "    index_columns = [sample_id, index_column]\n",
    "else:\n",
    "    raise NotImplementedError(\n",
    "        \"More than one feature: Needs to be implemented. see above logging output.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a94ae06d-64dd-4f4f-abde-8485a8c8458e",
   "metadata": {},
   "source": [
    "### Use some simulated missing for evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "760dbbf8",
   "metadata": {},
   "source": [
    "The validation simulated NA is used to by all models to evaluate training performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5b945aa-9b4e-4487-8b09-dca289e64d9d",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "val_pred_simulated_na = data.val_y.to_frame(name='observed')\n",
    "val_pred_simulated_na"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98558b10",
   "metadata": {
    "lines_to_next_cell": 2,
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "test_pred_simulated_na = data.test_y.to_frame(name='observed')\n",
    "test_pred_simulated_na.describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d00ea8e9-7f48-4f72-b013-7010666aa1a2",
   "metadata": {},
   "source": [
    "## Collaborative Filtering\n",
    "\n",
    "- save custom collab batch size (increase AE batch size by a factor), could be setup separately.\n",
    "- the test data is used to evaluate the performance after training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ee54305-266a-479a-b677-f151ddde250a",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# larger mini-batches speed up training\n",
    "ana_collab = models.collab.CollabAnalysis(\n",
    "    datasplits=data,\n",
    "    sample_column=sample_id,\n",
    "    item_column=index_column,  # not generic\n",
    "    target_column='intensity',\n",
    "    model_kwargs=dict(n_factors=args.latent_dim,\n",
    "                      y_range=(int(data.train_X.min()),\n",
    "                               int(data.train_X.max()) + 1)\n",
    "                      ),\n",
    "    batch_size=args.batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12ffa243-151e-4220-a1d5-247f8aba3429",
   "metadata": {
    "lines_to_next_cell": 2,
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "print(\"Args:\")\n",
    "pprint(ana_collab.model_kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a02e061-6789-4f3d-8031-a40879c496c8",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "ana_collab.model = EmbeddingDotBias.from_classes(\n",
    "    classes=ana_collab.dls.classes,\n",
    "    **ana_collab.model_kwargs)\n",
    "\n",
    "args.n_params = models.calc_net_weight_count(ana_collab.model)\n",
    "ana_collab.params['n_parameters'] = args.n_params\n",
    "ana_collab.learn = Learner(dls=ana_collab.dls, model=ana_collab.model, loss_func=MSELossFlat(),\n",
    "                           cbs=EarlyStoppingCallback(patience=args.patience),\n",
    "                           model_dir=args.out_models)\n",
    "if args.cuda:\n",
    "    ana_collab.learn.model = ana_collab.learn.model.cuda()\n",
    "else:\n",
    "    # try to set explicitly cpu in case not cuda\n",
    "    # MPS logic might not work properly in fastai yet https://github.com/fastai/fastai/pull/3858\n",
    "    ana_collab.learn.model = ana_collab.learn.model.cpu()\n",
    "\n",
    "# learn.summary() # see comment at DAE"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f82a7346-0b44-44a4-b995-d655e05656f8",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8317c9e1-d128-4ab4-8d60-775cb85ef535",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# papermill_description=train_collab\n",
    "suggested_lr = ana_collab.learn.lr_find()\n",
    "print(f\"{suggested_lr.valley = :.5f}\")\n",
    "ana_collab.learn.fit_one_cycle(args.epochs_max, lr_max=suggested_lr.valley)\n",
    "args.epoch_trained = ana_collab.learn.epoch + 1\n",
    "# ana_collab.learn.fit_one_cycle(args.epochs_max, lr_max=1e-3)\n",
    "ana_collab.model_kwargs['suggested_inital_lr'] = suggested_lr.valley\n",
    "ana_collab.learn.save('collab_model')\n",
    "fig, ax = plt.subplots(figsize=(15, 8))\n",
    "ax.set_title('CF loss: Reconstruction loss')\n",
    "ana_collab.learn.recorder.plot_loss(skip_start=5, ax=ax)\n",
    "recorder_dump = RecorderDump(\n",
    "    recorder=ana_collab.learn.recorder, name='CF')\n",
    "recorder_dump.save(args.out_figures)\n",
    "del recorder_dump\n",
    "pimmslearn.savefig(fig, name='collab_training',\n",
    "                   folder=args.out_figures)\n",
    "ana_collab.model_kwargs['batch_size'] = ana_collab.batch_size\n",
    "pimmslearn.io.dump_json(ana_collab.model_kwargs, args.out_models /\n",
    "                        TEMPLATE_MODEL_PARAMS.format('CF'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e979b7f0-a673-4d3d-9d53-6ac02618eaed",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ee5b225-d50a-4189-9995-ad99c4d47a45",
   "metadata": {},
   "source": [
    "Compare simulated_na data predictions to original values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb76e6c5-e135-41c4-95e8-a56c3764c731",
   "metadata": {
    "lines_to_next_cell": 2,
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# this could be done using the validation data laoder now\n",
    "ana_collab.test_dl = ana_collab.dls.test_dl(\n",
    "    data.val_y.reset_index())  # test_dl is here validation data\n",
    "val_pred_simulated_na['CF'], _ = ana_collab.learn.get_preds(\n",
    "    dl=ana_collab.test_dl)\n",
    "val_pred_simulated_na"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd22f63f-be3f-4f25-ad60-bafd6b028bd7",
   "metadata": {},
   "source": [
    "select test data predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7f0c597-d3c7-42d0-a6ef-3bc4c13121b8",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "ana_collab.test_dl = ana_collab.dls.test_dl(data.test_y.reset_index())\n",
    "test_pred_simulated_na['CF'], _ = ana_collab.learn.get_preds(dl=ana_collab.test_dl)\n",
    "test_pred_simulated_na"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cd76df6",
   "metadata": {
    "lines_to_next_cell": 2,
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "if args.save_pred_real_na:\n",
    "    pred_real_na = models.collab.get_missing_values(\n",
    "        df_train_long=data.train_X,\n",
    "        val_idx=data.val_y.index,\n",
    "        test_idx=data.test_y.index,\n",
    "        analysis_collab=ana_collab)\n",
    "    pred_real_na.to_csv(args.out_preds / f\"pred_real_na_{args.model_key}.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49eced7a-6cd7-414e-b974-4bd1dbe3a787",
   "metadata": {},
   "source": [
    "## Data in wide format\n",
    "\n",
    "- Autoencoder need data in wide format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cff8caf4-ccc9-4a36-a992-2cc596abe51a",
   "metadata": {
    "lines_to_next_cell": 2,
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "data.to_wide_format()\n",
    "args.M = data.train_X.shape[-1]\n",
    "data.train_X.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "861b9b3c-11d0-4cda-98c9-fadea16e47c4",
   "metadata": {},
   "source": [
    "### Validation data\n",
    "\n",
    "- all measured (identified, observed) peptides in validation data\n",
    "\n",
    "> Does not make to much sense to compare collab and AEs,\n",
    "> as the setup differs of training and validation data differs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d825e38e-f3d6-4bca-b621-150267e7b7bc",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# papermill_description=metrics\n",
    "d_metrics = models.Metrics()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55b9b887-a644-4086-8399-e27b533bd22a",
   "metadata": {},
   "source": [
    "The simulated NA for the validation step are real test data (not used for training nor early stopping)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "855a7a6f-93fd-4612-9d8d-96541a2441be",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "added_metrics = d_metrics.add_metrics(val_pred_simulated_na, 'valid_simulated_na')\n",
    "added_metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80fda7f1-3024-46bc-a3c1-17f509016bd4",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Test Datasplit\n",
    "\n",
    "Simulated NAs : Artificially created NAs. Some data was sampled and set\n",
    "explicitly to misssing before it was fed to the model for\n",
    "reconstruction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "571ac8d4-bb5d-45db-bba8-59817e476304",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "added_metrics = d_metrics.add_metrics(test_pred_simulated_na, 'test_simulated_na')\n",
    "added_metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b388de9d-0387-486a-9481-d673e0fec88b",
   "metadata": {},
   "source": [
    "Save all metrics as json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87910434-7d07-4e8e-8380-c92fc515bd16",
   "metadata": {
    "lines_to_next_cell": 2,
    "tags": []
   },
   "outputs": [],
   "source": [
    "pimmslearn.io.dump_json(d_metrics.metrics, args.out_metrics /\n",
    "                        f'metrics_{args.model_key}.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d99deb9-9aad-4ba9-b79d-e4b3c6c7f023",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "metrics_df = models.get_df_from_nested_dict(\n",
    "    d_metrics.metrics, column_levels=['model', 'metric_name']).T\n",
    "metrics_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d8d6b3e-1316-4795-bc3d-4bdf270b890e",
   "metadata": {},
   "source": [
    "## Save predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "782636ac-c979-4f8b-9fc0-66fd0c7a3a8b",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# save simulated missing values for both splits\n",
    "val_pred_simulated_na.to_csv(args.out_preds / f\"pred_val_{args.model_key}.csv\")\n",
    "test_pred_simulated_na.to_csv(args.out_preds / f\"pred_test_{args.model_key}.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42f2c184-e9ef-42be-ae87-7cbf52a0c0b6",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f13cb38-abf0-4b56-9399-3d11d32f7fbc",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "args.dump(fname=args.out_models / f\"model_config_{args.model_key}.yaml\")\n",
    "args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "408b261a",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,py:percent"
  },
  "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.13"
  },
  "toc-autonumbering": true,
  "vscode": {
   "interpreter": {
    "hash": "cf83e9cb890c7f96eb0ae04f39a82254555f56a1a0ed2f03b23a8b40fe6cd31c"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}