project/01_1_train_VAE.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f0e1208b-9b36-4294-b3fd-910e05a82f2e",
   "metadata": {},
   "source": [
    "# Variational Autoencoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18b5d571-2956-4112-b22c-43d6c2146b06",
   "metadata": {
    "lines_to_next_cell": 2,
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "\n",
    "import logging\n",
    "from functools import partial\n",
    "\n",
    "import pandas as pd\n",
    "import sklearn\n",
    "import torch\n",
    "from fastai import learner\n",
    "from fastai.basics import *\n",
    "from fastai.callback.all import *\n",
    "from fastai.callback.all import EarlyStoppingCallback\n",
    "from fastai.learner import Learner\n",
    "from fastai.torch_basics import *\n",
    "from IPython.display import display\n",
    "from sklearn.impute import SimpleImputer\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from torch.nn import Sigmoid\n",
    "\n",
    "import pimmslearn\n",
    "import pimmslearn.model\n",
    "import pimmslearn.models as models\n",
    "import pimmslearn.nb\n",
    "from pimmslearn.analyzers import analyzers\n",
    "from pimmslearn.io import datasplits\n",
    "# overwriting Recorder callback with custom plot_loss\n",
    "from pimmslearn.models import ae, plot_loss\n",
    "\n",
    "learner.Recorder.plot_loss = plot_loss\n",
    "\n",
    "\n",
    "logger = pimmslearn.logging.setup_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": "code",
   "execution_count": null,
   "id": "5dbc8d89",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# catch passed parameters\n",
    "args = None\n",
    "args = dict(globals()).keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43f062bb-8d2c-4afa-8a81-9e2a328050fe",
   "metadata": {},
   "source": [
    "Papermill script parameters:"
   ]
  },
  {
   "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'  # file format of create splits, default pickle (pkl)\n",
    "# Machine parsed metadata from rawfile workflow\n",
    "fn_rawfile_metadata: str = 'data/dev_datasets/HeLa_6070/files_selected_metadata_N50.csv'\n",
    "# training\n",
    "epochs_max: int = 50  # Maximum number of epochs\n",
    "batch_size: int = 64  # Batch size for training (and evaluation)\n",
    "cuda: bool = True  # Whether to use a GPU for training\n",
    "# model\n",
    "# Dimensionality of encoding dimension (latent space of model)\n",
    "latent_dim: int = 25\n",
    "# A underscore separated string of layers, '256_128' for the encoder, reverse will be use for decoder\n",
    "hidden_layers: str = '256_128'\n",
    "# force_train:bool = True # Force training when saved model could be used. Per default re-train model\n",
    "patience: int = 50  # Patience for early stopping\n",
    "sample_idx_position: int = 0  # position of index which is sample ID\n",
    "model: str = 'VAE'  # model name\n",
    "model_key: str = 'VAE'  # potentially alternative key for model (grid search)\n",
    "save_pred_real_na: bool = True  # Save all predictions for missing values\n",
    "# metadata -> defaults for metadata extracted from machine data\n",
    "meta_date_col: str = None  # date column in meta data\n",
    "meta_cat_col: str = None  # category column in meta data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1da8e04b-2eb2-4ecb-b1d1-2de871cd0a56",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "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": "8083658b",
   "metadata": {
    "lines_to_next_cell": 2,
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "args = pimmslearn.nb.args_from_dict(args)\n",
    "\n",
    "if isinstance(args.hidden_layers, str):\n",
    "    args.overwrite_entry(\"hidden_layers\", [int(x)\n",
    "                         for x in args.hidden_layers.split('_')])\n",
    "else:\n",
    "    raise ValueError(\n",
    "        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": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "TEMPLATE_MODEL_PARAMS = 'model_params_{}.json'"
   ]
  },
  {
   "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.sample(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6045414b",
   "metadata": {},
   "source": [
    "Infer index names from long format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44958473",
   "metadata": {
    "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": [
    "load meta data for splits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5b945aa-9b4e-4487-8b09-dca289e64d9d",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "if args.fn_rawfile_metadata:\n",
    "    df_meta = pd.read_csv(args.fn_rawfile_metadata, index_col=0)\n",
    "    display(df_meta.loc[data.train_X.index.levels[0]])\n",
    "else:\n",
    "    df_meta = None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7ef773a-5588-4b11-9bd7-840b628eeaff",
   "metadata": {},
   "source": [
    "## Initialize Comparison\n",
    "\n",
    "- replicates idea for truely missing values: Define truth as by using n=3 replicates to impute\n",
    "  each sample\n",
    "- real test data:\n",
    "    - Not used for predictions or early stopping.\n",
    "    - [x] add some additional NAs based on distribution of data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98f675b6-e619-45b6-8f04-b75237d212a7",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "freq_feat = pimmslearn.io.datasplits.load_freq(args.data)\n",
    "freq_feat.head()  # training data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e02afb4e-ee91-4753-a118-7fa20fe621c9",
   "metadata": {},
   "source": [
    "### Produce some addional simulated samples"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02d5763b-00fe-44ce-9dfa-b6e506045762",
   "metadata": {},
   "source": [
    "The validation simulated NA is used to by all models to evaluate training performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19eebaff-0e1e-4e44-ae40-12d2f0e75c74",
   "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": "3797a539-84d9-430a-8d16-7cc0eebfe9f5",
   "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": "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": {
    "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": "5cc46767-49c9-487e-bf2d-19aacc9114c3",
   "metadata": {},
   "source": [
    "### Add interpolation performance"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21681578-dfc5-4f3c-ac7d-a0758ed5417b",
   "metadata": {},
   "source": [
    "### Fill Validation data with potentially missing features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7952fe13",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "data.train_X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0a0bcd9-22af-4dd9-af56-b041931ee918",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "data.val_y  # potentially has less features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f0826f9",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "data.val_y = pd.DataFrame(pd.NA, index=data.train_X.index,\n",
    "                          columns=data.train_X.columns).fillna(data.val_y)\n",
    "data.val_y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6de7a68e-34a3-463f-a1b3-7c76a8855ca0",
   "metadata": {},
   "source": [
    "## Variational Autoencoder"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f5c05d4-d0bb-476f-b955-ac5822c38f86",
   "metadata": {},
   "source": [
    "### Analysis: DataLoaders, Model, transform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0d0d02f",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "default_pipeline = sklearn.pipeline.Pipeline(\n",
    "    [\n",
    "        ('normalize', StandardScaler()),\n",
    "        ('impute', SimpleImputer(add_indicator=False))\n",
    "    ])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae2f194f-a879-4190-89b9-5b370a827212",
   "metadata": {},
   "source": [
    "### Analysis: DataLoaders, Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43d49b4a-00ec-4874-8839-28a3cbc0e3b3",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "\n",
    "analysis = ae.AutoEncoderAnalysis(  # datasplits=data,\n",
    "    train_df=data.train_X,\n",
    "    val_df=data.val_y,\n",
    "    model=models.vae.VAE,\n",
    "    model_kwargs=dict(n_features=data.train_X.shape[-1],\n",
    "                      n_neurons=args.hidden_layers,\n",
    "                      # last_encoder_activation=None,\n",
    "                      last_decoder_activation=None,\n",
    "                      dim_latent=args.latent_dim),\n",
    "    transform=default_pipeline,\n",
    "    decode=['normalize'],\n",
    "    bs=args.batch_size)\n",
    "args.n_params = analysis.n_params_ae\n",
    "if args.cuda:\n",
    "    analysis.model = analysis.model.cuda()\n",
    "analysis.model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef5410d9-8b64-45bc-96d3-4b04145043c0",
   "metadata": {},
   "source": [
    "### Training\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cca0e4a4",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "results = []\n",
    "loss_fct = partial(models.vae.loss_fct, results=results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9366b06a",
   "metadata": {
    "lines_to_next_cell": 0,
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "analysis.learn = Learner(dls=analysis.dls,\n",
    "                         model=analysis.model,\n",
    "                         loss_func=loss_fct,\n",
    "                         cbs=[ae.ModelAdapterVAE(),\n",
    "                              EarlyStoppingCallback(patience=args.patience)\n",
    "                              ])\n",
    "\n",
    "analysis.learn.show_training_loop()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2231b67e",
   "metadata": {},
   "source": [
    "Adding a `EarlyStoppingCallback` results in an error.  Potential fix in\n",
    "[PR3509](https://github.com/fastai/fastai/pull/3509) is not yet in\n",
    "current version. Try again later"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4b0aec2",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# learn.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "563a1e0a",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "suggested_lr = analysis.learn.lr_find()\n",
    "analysis.params['suggested_inital_lr'] = suggested_lr.valley\n",
    "suggested_lr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "468565f5",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "results.clear()  # reset results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0f4eaee-f308-4081-aa16-0579fe9df7fe",
   "metadata": {},
   "source": [
    "dump model config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec77d9e5-f619-4355-ab37-2bd44029236d",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# needs class as argument, not instance, but serialization needs instance\n",
    "analysis.params['last_decoder_activation'] = Sigmoid()\n",
    "\n",
    "pimmslearn.io.dump_json(\n",
    "    pimmslearn.io.parse_dict(\n",
    "        analysis.params, types=[\n",
    "            (torch.nn.modules.module.Module, lambda m: str(m))\n",
    "        ]),\n",
    "    args.out_models / TEMPLATE_MODEL_PARAMS.format(args.model_key))\n",
    "\n",
    "# restore original value\n",
    "analysis.params['last_decoder_activation'] = Sigmoid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43d18ab6",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# papermill_description=train\n",
    "analysis.learn.fit_one_cycle(args.epochs_max, lr_max=suggested_lr.valley)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f862ba1-7589-4076-9dae-b7e4868f4b8f",
   "metadata": {},
   "source": [
    "Save number of actually trained epochs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "079f9743-213c-422c-ba61-919c276fd710",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "args.epoch_trained = analysis.learn.epoch + 1\n",
    "args.epoch_trained"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f41e0969",
   "metadata": {},
   "source": [
    "#### Loss normalized by total number of measurements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a4f91f9-c1a2-40c5-99d4-c289fb89cff8",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "N_train_notna = data.train_X.notna().sum().sum()\n",
    "N_val_notna = data.val_y.notna().sum().sum()\n",
    "fig = models.plot_training_losses(analysis.learn, args.model_key,\n",
    "                                  folder=args.out_figures,\n",
    "                                  norm_factors=[N_train_notna, N_val_notna])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9db46826-6ab0-47b5-942c-0edc08a15e35",
   "metadata": {},
   "source": [
    "### Predictions\n",
    "create predictions and select validation data predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e73fb4dd-e73a-48df-82b7-b378fd3ee266",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "analysis.model.eval()\n",
    "pred, target = res = ae.get_preds_from_df(df=data.train_X, learn=analysis.learn,\n",
    "                                          position_pred_tuple=0,\n",
    "                                          transformer=analysis.transform)\n",
    "pred = pred.stack()\n",
    "pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9fc0e36",
   "metadata": {
    "lines_to_next_cell": 2,
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "val_pred_simulated_na['VAE'] = pred  # 'model_key' ?\n",
    "val_pred_simulated_na"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4d1fd73-9eb7-4f25-ad81-c42c6a840e77",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "test_pred_simulated_na['VAE'] = pred  # model_key?\n",
    "test_pred_simulated_na"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "733d4cbe",
   "metadata": {},
   "source": [
    "save missing values predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c43e401f",
   "metadata": {
    "lines_to_next_cell": 2,
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "if args.save_pred_real_na:\n",
    "    pred_real_na = ae.get_missing_values(df_train_wide=data.train_X,\n",
    "                                         val_idx=val_pred_simulated_na.index,\n",
    "                                         test_idx=test_pred_simulated_na.index,\n",
    "                                         pred=pred)\n",
    "    display(pred_real_na)\n",
    "    pred_real_na.to_csv(args.out_preds / f\"pred_real_na_{args.model_key}.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6c3dd2a-a30c-4e72-b5e8-7de2a7469eb2",
   "metadata": {},
   "source": [
    "### Plots\n",
    "\n",
    "- validation data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7999c89e-65fe-4c00-8e20-cb8ab88d1603",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "analysis.model = analysis.model.cpu()\n",
    "# underlying data is train_X for both\n",
    "# assert analysis.dls.valid.data.equals(analysis.dls.train.data)\n",
    "# Reconstruct DataLoader for case that during training singleton batches were dropped\n",
    "_dl = torch.utils.data.DataLoader(\n",
    "    pimmslearn.io.datasets.DatasetWithTarget(\n",
    "        analysis.dls.valid.data),\n",
    "    batch_size=args.batch_size,\n",
    "    shuffle=False)\n",
    "df_latent = pimmslearn.model.get_latent_space(analysis.model.get_mu_and_logvar,\n",
    "                                              dl=_dl,\n",
    "                                              dl_index=analysis.dls.valid.data.index)\n",
    "df_latent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fdd8f86-639e-4e0e-bb89-466f3ba0ef7b",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "ana_latent = analyzers.LatentAnalysis(df_latent,\n",
    "                                      df_meta,\n",
    "                                      args.model_key,\n",
    "                                      folder=args.out_figures)\n",
    "if args.meta_date_col and df_meta is not None:\n",
    "    figures[f'latent_{args.model_key}_by_date'], ax = ana_latent.plot_by_date(\n",
    "        args.meta_date_col)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "763a5633-a9dd-4785-a0c0-91f588346c22",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "if args.meta_cat_col and df_meta is not None:\n",
    "    figures[f'latent_{args.model_key}_by_{\"_\".join(args.meta_cat_col.split())}'], ax = ana_latent.plot_by_category(\n",
    "        args.meta_cat_col)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3fdd5cb-4038-489f-b4d8-54ec6ea913b5",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "feat_freq_val = val_pred_simulated_na['observed'].groupby(level=-1).count()\n",
    "feat_freq_val.name = 'freq_val'\n",
    "ax = feat_freq_val.plot.box()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "446adbaf-81db-4ac5-b041-064744143602",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "feat_freq_val.value_counts().sort_index().head()  # require more than one feat?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d408dbfa",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "errors_val = val_pred_simulated_na.drop('observed', axis=1).sub(\n",
    "    val_pred_simulated_na['observed'], axis=0)\n",
    "errors_val = errors_val.abs().groupby(level=-1).mean()\n",
    "errors_val = errors_val.join(freq_feat).sort_values(by='freq', ascending=True)\n",
    "\n",
    "\n",
    "errors_val_smoothed = errors_val.copy()  # .loc[feat_freq_val > 1]\n",
    "errors_val_smoothed[errors_val.columns[:-1]] = errors_val[errors_val.columns[:-1]\n",
    "                                                          ].rolling(window=200, min_periods=1).mean()\n",
    "ax = errors_val_smoothed.plot(x='freq', figsize=(15, 10))\n",
    "# errors_val_smoothed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e505353-f19e-4961-9279-f1f0f1e4be09",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "errors_val = val_pred_simulated_na.drop('observed', axis=1).sub(\n",
    "    val_pred_simulated_na['observed'], axis=0)\n",
    "errors_val.abs().groupby(level=-1).agg(['mean', 'count'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95022f04-0e0d-47bf-8267-6135a936328f",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "errors_val"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29f16d62-f1c8-4447-bdd3-2d3161b6346f",
   "metadata": {},
   "source": [
    "## Comparisons\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": "markdown",
   "id": "861b9b3c-11d0-4cda-98c9-fadea16e47c4",
   "metadata": {},
   "source": [
    "### Validation data\n",
    "\n",
    "- all measured (identified, observed) peptides in validation data"
   ]
  },
  {
   "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(no_na_key='NA interpolated', with_na_key='NA not interpolated')\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": [
     "hide-input"
    ]
   },
   "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": {
    "lines_to_next_cell": 2
   },
   "source": [
    "### Test Datasplit\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "571ac8d4-bb5d-45db-bba8-59817e476304",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "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": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "pimmslearn.io.dump_json(d_metrics.metrics, args.out_metrics /\n",
    "                        f'metrics_{args.model_key}.json')\n",
    "d_metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d99deb9-9aad-4ba9-b79d-e4b3c6c7f023",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "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": "2b49894b-876e-428b-bbba-18229db8b984",
   "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": {},
   "source": [
    "## Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06e433ec",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "figures  # switch to fnames?"
   ]
  },
  {
   "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"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "cf83e9cb890c7f96eb0ae04f39a82254555f56a1a0ed2f03b23a8b40fe6cd31c"
  },
  "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
 },
 "nbformat": 4,
 "nbformat_minor": 5
}