PAPER_RESULTS.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SsOAZVkmTCUT"
   },
   "source": [
    "## **1**. Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "executionInfo": {
     "elapsed": 1644,
     "status": "ok",
     "timestamp": 1619043221554,
     "user": {
      "displayName": "Henrique Silva",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh0sTZzx8mKsKV9-j3HdK09d6njRJoJnyBA8jbau1I=s64",
      "userId": "15429421149160718308"
     },
     "user_tz": 180
    },
    "id": "2zP0oI7j9s6F"
   },
   "outputs": [],
   "source": [
    "ON_COLAB=False\n",
    "if ON_COLAB:\n",
    "    from google.colab import drive\n",
    "    drive.mount('/content/drive',force_remount=True)\n",
    "    BASEDIR='/content/drive/My Drive/MentalHealthShared/'\n",
    "    PYTHONDIR=BASEDIR+'src'\n",
    "    RESULTSDIR=BASEDIR+'results/'\n",
    "    MODELSDIR=BASEDIR+'model/'\n",
    "    DATADIR=BASEDIR+'data/'\n",
    "    import models\n",
    "    import pytorchtools\n",
    "    from utils import reset_seeds, count_parameters, evaluate, train_over_nepochs, createTensorDataset\n",
    "    from training_functions import load_df, compute_bin_weights, save_stats_tensors, load_stats_tensors, get_znorm_params, get_subreddit_range, split_indices, get_subreddit_weights, get_baselines_df, WeightedL1Loss, WeightedMSELoss, grid_search_train\n",
    "else:\n",
    "    import os\n",
    "    BASEDIR = os.getcwd() + \"/\"\n",
    "    dirs = [\"results\",\"model\",\"data\"]\n",
    "    for dirc in dirs:\n",
    "        if dirc not in os.listdir(): \n",
    "            os.makedirs(os.path.join(BASEDIR,dirc))\n",
    "    PYTHONDIR=BASEDIR+'src/'\n",
    "    RESULTSDIR=BASEDIR+'results/'\n",
    "    MODELSDIR=BASEDIR+'model/'\n",
    "    DATADIR=BASEDIR+'data/'\n",
    "    from src import models\n",
    "    from src import pytorchtools\n",
    "    from src.utils import reset_seeds, count_parameters, evaluate, train_over_nepochs, createTensorDataset\n",
    "    from src.training_functions import load_df, compute_bin_weights, save_stats_tensors, load_stats_tensors, get_znorm_params, get_subreddit_range, split_indices, get_subreddit_weights, get_baselines_df, WeightedL1Loss, WeightedMSELoss, grid_search_train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 1643,
     "status": "ok",
     "timestamp": 1619043221555,
     "user": {
      "displayName": "Henrique Silva",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh0sTZzx8mKsKV9-j3HdK09d6njRJoJnyBA8jbau1I=s64",
      "userId": "15429421149160718308"
     },
     "user_tz": 180
    },
    "id": "W4zubEqFtB4l"
   },
   "outputs": [],
   "source": [
    "# DEFINE SUBREDDITS\n",
    "SUBREDDITS  = ['Anxiety','bipolar','depression','SuicideWatch'] \n",
    "\n",
    "subreddit2title = {'depression':'DEP','suicidewatch':'SUI','anxiety':'ANX','bipolar':'BIP'}\n",
    "\n",
    "EXTENSION = '.parquet'\n",
    "STRATIFIED = True  # If true, will apply a stratified K-fold cross-validation to the dataset.\n",
    "TRAIN = True # If true, will train the models\n",
    "\n",
    "USE_GRU = True # If true, will use the GRU RNN model in training\n",
    "USE_XGB = True # If true, will use the XGBoost model in training\n",
    "\n",
    "GRID_SEARCH_XGB = False # If true, will do a grid search for the XGBoost model\n",
    "PLOT_XGB = False # If true, will plot the results found by the XGBoost model(Only available if TRAIN = True)\n",
    "\n",
    "TEST  = False # If true, will test the models\n",
    "\n",
    "ZNORMALIZE = False # If true, will apply a z-normalization to the dataset\n",
    "\n",
    "\n",
    "KEEP_TEXT = True # Only True if using Section 6 for Case Study\n",
    "\n",
    "\n",
    "assert TRAIN ^ TEST # either TRAIN or TEST\n",
    "\n",
    "FILTERED=True # If true will use filtered seq_len for the threads\n",
    "BIN_WIDTH=0.2 # Controls the width of the bins used to calculate the weighted L1 loss\n",
    "MIN_VALUE = -1 # Controls the minimum output value(set to -1)\n",
    "\n",
    "\n",
    "INCLUDE_TARGET = 0\n",
    "\n",
    "\n",
    "MAX_BRANCH_LEN = 16   # not including authors' last comment\n",
    "MAX_THREAD_LEN = 64   # not including authors' last comment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 4594,
     "status": "ok",
     "timestamp": 1619043224518,
     "user": {
      "displayName": "Henrique Silva",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh0sTZzx8mKsKV9-j3HdK09d6njRJoJnyBA8jbau1I=s64",
      "userId": "15429421149160718308"
     },
     "user_tz": 180
    },
    "id": "KfJCklnZc63x",
    "outputId": "37c9fefc-77a4-4ee0-8ccf-29e06730bf53"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "if device == torch.device('cuda'):\n",
    "    print(f\"Device successfully set to cuda\")\n",
    "else:\n",
    "    print(\"WARNING! DEVICE IS NOT SET TO CUDA\")\n",
    "\n",
    "print(torch.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 6913,
     "status": "ok",
     "timestamp": 1619043226841,
     "user": {
      "displayName": "Henrique Silva",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh0sTZzx8mKsKV9-j3HdK09d6njRJoJnyBA8jbau1I=s64",
      "userId": "15429421149160718308"
     },
     "user_tz": 180
    },
    "id": "sKyjZsDKKrPY"
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(PYTHONDIR)\n",
    "import random\n",
    "import pickle\n",
    "import time\n",
    "import pdb\n",
    "import importlib\n",
    "import itertools\n",
    "import pprint\n",
    "import copy\n",
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.model_selection import GridSearchCV,train_test_split\n",
    "from sklearn.metrics import mean_squared_error\n",
    "\n",
    "from tqdm.notebook import tqdm\n",
    "from scipy import stats\n",
    "\n",
    "import torch.optim as optim\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset, TensorDataset, DataLoader, random_split, Sampler, SubsetRandomSampler, Subset\n",
    "\n",
    "from xgboost import XGBRegressor, Booster, DMatrix\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qfiTmnNDV367"
   },
   "source": [
    "# **2**. Creating a RedditDataset or TensorDataset instance\n",
    "\n",
    "Instead of using a custom Dataset class (former ```RedditDataset```), we store the data required by the experiments in a ```TensorDataset``` to speed up the retrieval of the batchs.\n",
    "\n",
    "```set_threads``` is a list of indexes containing the threads to be used as a dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CdCFbGprpA5s"
   },
   "source": [
    "## **2.1** Build the source data\n",
    "\n",
    "* Build the source data in order to predict the EmT based on the sequence of comments in a thread"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 102691,
     "status": "ok",
     "timestamp": 1619020752729,
     "user": {
      "displayName": "Henrique Silva",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh0sTZzx8mKsKV9-j3HdK09d6njRJoJnyBA8jbau1I=s64",
      "userId": "15429421149160718308"
     },
     "user_tz": 180
    },
    "id": "wv1F_eR3El09",
    "outputId": "2f26453a-968a-4d4b-d220-82c0138189fc"
   },
   "outputs": [],
   "source": [
    "suffix='_distilbert_filtered_posts' + EXTENSION\n",
    "df_list = {subreddit:load_df(DATADIR+subreddit+suffix, MAX_THREAD_LEN) \\\n",
    "           for subreddit in SUBREDDITS}\n",
    "\n",
    "post_df = pd.concat((df_list[subreddit] for subreddit in SUBREDDITS), keys=SUBREDDITS)\n",
    "del df_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 102679,
     "status": "ok",
     "timestamp": 1619020752731,
     "user": {
      "displayName": "Henrique Silva",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh0sTZzx8mKsKV9-j3HdK09d6njRJoJnyBA8jbau1I=s64",
      "userId": "15429421149160718308"
     },
     "user_tz": 180
    },
    "id": "MyZzK_aB6YHJ",
    "outputId": "d800a3a6-9fc4-4bde-f570-9f17f726a46d"
   },
   "outputs": [],
   "source": [
    "if not KEEP_TEXT:\n",
    "  post_df = post_df[['created_utc', 'seq_len','score', 'features', 'filtered_seqlen','valid_branches']]\n",
    "\n",
    "# new strategy to construct observations: follow branches of every discussion tree \n",
    "\n",
    "if FILTERED:\n",
    "    # dropna on filtered_seqlen, then replace seq_len by filtered_seqlen\n",
    "    post_df.dropna(subset=['filtered_seqlen'], inplace=True)\n",
    "    post_df.filtered_seqlen = post_df.filtered_seqlen.astype(int)\n",
    "\n",
    "    post_df.drop(columns='seq_len',inplace=True)\n",
    "    post_df.rename(columns={'filtered_seqlen':'seq_len'},inplace=True)\n",
    "else:\n",
    "    post_df.drop(columns='filtered_seqlen',inplace=True)\n",
    "\n",
    "print(f'Fraction of threads that had to be truncated: {(post_df.seq_len>(MAX_THREAD_LEN+1)).mean()}')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 106633,
     "status": "ok",
     "timestamp": 1619020756692,
     "user": {
      "displayName": "Henrique Silva",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh0sTZzx8mKsKV9-j3HdK09d6njRJoJnyBA8jbau1I=s64",
      "userId": "15429421149160718308"
     },
     "user_tz": 180
    },
    "id": "vt5koa7q5QVF",
    "outputId": "ba005eb6-ef00-4588-ae06-c3c46efcc377"
   },
   "outputs": [],
   "source": [
    "if ZNORMALIZE:\n",
    "  prefix = ''\n",
    "else:\n",
    "  prefix = 'unnorm_'\n",
    "\n",
    "if TRAIN:\n",
    "  if ZNORMALIZE:\n",
    "    src_m, src_s = get_znorm_params(post_df)\n",
    "  else:\n",
    "    shape = [1,post_df.iloc[0].features.shape[-1]]\n",
    "    src_m = torch.zeros(shape)\n",
    "    src_s = torch.ones(shape)\n",
    "\n",
    "  save_stats_tensors(src_m,src_s,f'{BASEDIR}data/{prefix}')\n",
    "  \n",
    "if TEST:\n",
    "  src_m, src_s = load_stats_tensors(f'{BASEDIR}data/{prefix}')\n",
    "\n",
    "score_m = float(src_m[0,-2])\n",
    "score_s = float(src_s[0,-2])\n",
    "\n",
    "print(f'Average score in dataset is {score_m}')\n",
    "\n",
    "subreddit2range = get_subreddit_range(post_df)\n",
    "print(subreddit2range[SUBREDDITS[0]])\n",
    "\n",
    "suffix = \"random\"\n",
    "if STRATIFIED:\n",
    "  suffix += '_strat'\n",
    "\n",
    "# get indices\n",
    "if TRAIN:\n",
    "  if len(SUBREDDITS) > 1:\n",
    "      train_inds, valid_inds, test_inds  = split_indices(post_df, STRATIFIED, MIN_VALUE, BIN_WIDTH)\n",
    "  else:\n",
    "    subreddit = SUBREDDITS[0]\n",
    "    with open(f'{DATADIR}{subreddit}_{suffix}_splits.pkl','rb') as infile:\n",
    "      splits = pickle.load(infile)\n",
    "      train_locs =  [(subreddit,loc) for loc in splits[0]]\n",
    "      valid_locs =  [(subreddit,loc) for loc in splits[1]]\n",
    "      test_locs  =  [(subreddit,loc) for loc in splits[2]]\n",
    "\n",
    "    train_inds = post_df.index.get_indexer_for(train_locs)\n",
    "    valid_inds = post_df.index.get_indexer_for(valid_locs)\n",
    "    test_inds  = post_df.index.get_indexer_for(test_locs)\n",
    "\n",
    "\n",
    "if TEST:\n",
    "  train_locs = []\n",
    "  valid_locs = []\n",
    "  test_locs  = []\n",
    "  for subreddit in SUBREDDITS:\n",
    "      with open(f'{DATADIR}{subreddit}_{suffix}_splits.pkl','rb') as infile:\n",
    "        splits = pickle.load(infile)\n",
    "        test_locs += [(subreddit,loc) for loc in splits[2]]\n",
    "\n",
    "  # extract instances\n",
    "  print(len(post_df))\n",
    "  post_df = post_df.loc[test_locs]\n",
    "  print(len(post_df))\n",
    "\n",
    "# compute weights for Weighted L1 Loss\n",
    "subreddit2weights = get_subreddit_weights(post_df, BIN_WIDTH,MIN_VALUE,device)\n",
    "print(subreddit2weights)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 115044,
     "status": "ok",
     "timestamp": 1619020765124,
     "user": {
      "displayName": "Henrique Silva",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh0sTZzx8mKsKV9-j3HdK09d6njRJoJnyBA8jbau1I=s64",
      "userId": "15429421149160718308"
     },
     "user_tz": 180
    },
    "id": "heaO2d6xcAvh",
    "outputId": "b91da058-dba9-46e7-a2d4-34b0e31f1b41"
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "print('Creating src')\n",
    "\n",
    "print('Creating y')\n",
    "y = torch.Tensor(post_df.apply(lambda p: p.score[p.seq_len-1], axis=1).values)\n",
    "\n",
    "print('Creating src_len_series')\n",
    "src_len_series = post_df.seq_len-1\n",
    "max_length=MAX_THREAD_LEN\n",
    "\n",
    "if USE_GRU:\n",
    "    src = nn.utils.rnn.pad_sequence(\n",
    "    [ p.features[:min(MAX_THREAD_LEN,p.seq_len-1),:] for index, p in post_df.iterrows()], batch_first=True)\n",
    "    # src = nn.utils.rnn.pad_sequence(\n",
    "    # [ (p.features[:min(MAX_THREAD_LEN,p.seq_len-1),:]-src_m)/src_s for index, p in post_df.iterrows()], batch_first=True)\n",
    "    print(f\"GRU src tensor size: {src.size()}\")\n",
    "\n",
    "if USE_XGB:      \n",
    "    src_xgb = torch.cat((\n",
    "        torch.cat([torch.mean((p.features[:min(MAX_THREAD_LEN,p.seq_len-1),:]-src_m)/src_s,keepdims=True,dim=0) for index, p in post_df.iterrows()],axis=0),\n",
    "        torch.cat([torch.max((p.features[:min(MAX_THREAD_LEN,p.seq_len-1),:]-src_m)/src_s,keepdims=True,dim=0)[0] for index, p in post_df.iterrows()],axis=0)\n",
    "        ),1)\n",
    "    print(f\"XGB src tensor size: {src_xgb.size()}\")\n",
    "\n",
    "\n",
    "print(f'y tensor size: {y.size()}')\n",
    "\n",
    "\n",
    "if INCLUDE_TARGET:\n",
    "  tgt = nn.utils.rnn.pad_sequence(\n",
    "    [(p.features[b[-2]]-src_m)/src_s for index, p in post_df.iterrows() for b in p.valid_branches],\n",
    "  batch_first=True)\n",
    "else:\n",
    "  tgt = None\n",
    "\n",
    "  # clean up memory\n",
    "  #if (not TEST) and (not KEEP_TEXT):\n",
    "  if (not KEEP_TEXT):\n",
    "    del post_df\n",
    "\n",
    "  print('Creating dataset')\n",
    "  if USE_GRU:\n",
    "    #dataset = createTensorDataset(src, src_len_series, y, tgt=tgt, max_length=max_length) # all threads, lim 63 comments\n",
    "    dataset = createTensorDataset(src, src_len_series, y, max_length=max_length) # all threads, lim 63 comments\n",
    "    del src\n",
    "  if USE_XGB:\n",
    "    dataset_xgb = createTensorDataset(src_xgb, src_len_series, y, max_length=max_length)\n",
    "    del src_xgb\n",
    "  \n",
    "  del src_len_series, tgt, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 115037,
     "status": "ok",
     "timestamp": 1619020765124,
     "user": {
      "displayName": "Henrique Silva",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh0sTZzx8mKsKV9-j3HdK09d6njRJoJnyBA8jbau1I=s64",
      "userId": "15429421149160718308"
     },
     "user_tz": 180
    },
    "id": "6MLMLb6LAidE",
    "outputId": "5c277d0a-7c5b-416a-a289-2e17bec8c37e"
   },
   "outputs": [],
   "source": [
    "if USE_GRU:\n",
    "  batch = dataset[0]\n",
    "  EMBEDDING_DIM = batch[0].shape[-1]\n",
    "  print(f\"embedding dimension for GRU: {EMBEDDING_DIM}\")\n",
    "if USE_XGB:\n",
    "  batch = dataset_xgb[0]\n",
    "  EMBEDDING_DIM_XGB = batch[0].shape[-1]\n",
    "  print(f\"embedding dimension for XGB: {EMBEDDING_DIM_XGB}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "39QaZ0BUpLeO"
   },
   "source": [
    "## **2.2** Create training, validation and test sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 115030,
     "status": "ok",
     "timestamp": 1619020765125,
     "user": {
      "displayName": "Henrique Silva",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh0sTZzx8mKsKV9-j3HdK09d6njRJoJnyBA8jbau1I=s64",
      "userId": "15429421149160718308"
     },
     "user_tz": 180
    },
    "id": "5LGqkeO8916A",
    "outputId": "95127a0d-07be-4883-9af5-69c364265cf7"
   },
   "outputs": [],
   "source": [
    "#%%time\n",
    "\n",
    "print(\"Creating dataloaders...\")\n",
    "if TRAIN:\n",
    "  if USE_XGB:\n",
    "    train_loader_xgb = DataLoader(Subset(dataset_xgb,train_inds), batch_size=len(train_inds), shuffle=True, num_workers=0)\n",
    "    valid_loader_xgb = DataLoader(Subset(dataset_xgb,valid_inds), batch_size=len(valid_inds), shuffle=False, num_workers=0)\n",
    "    test_loader_xgb  = DataLoader(Subset(dataset_xgb, test_inds), batch_size=len(test_inds), shuffle=False, num_workers=1) \n",
    "\n",
    "  if USE_GRU: \n",
    "      train_loader = DataLoader(Subset(dataset,train_inds), batch_size=32, shuffle=True, num_workers=1)\n",
    "      valid_loader = DataLoader(Subset(dataset,valid_inds), batch_size=len(valid_inds), shuffle=False, num_workers=1)\n",
    "      test_loader  = DataLoader(Subset(dataset, test_inds), batch_size=len(test_inds), shuffle=False, num_workers=1)\n",
    "\n",
    "\n",
    "if TEST:\n",
    "  if USE_GRU:\n",
    "    test_loader  = DataLoader(dataset, batch_size=len(dataset), shuffle=False, num_workers=1)\n",
    "  if USE_XGB:\n",
    "    test_loader_xgb  = DataLoader(dataset_xgb, batch_size=len(dataset_xgb), shuffle=False, num_workers=1)\n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9UimhhZOUkm_"
   },
   "source": [
    "# **3**. Training the models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6k5SkAKY8IIO"
   },
   "source": [
    "## **3.1** (Stacked) (Bi-) GRU model\n",
    "\n",
    "Takes one source sequence as input."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3PEaxt-rU5I4"
   },
   "source": [
    "Example of how to run the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 396,
     "status": "ok",
     "timestamp": 1619020801356,
     "user": {
      "displayName": "Henrique Silva",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh0sTZzx8mKsKV9-j3HdK09d6njRJoJnyBA8jbau1I=s64",
      "userId": "15429421149160718308"
     },
     "user_tz": 180
    },
    "id": "f4X21d2JN1j8",
    "outputId": "dd239396-6a4d-48bc-b637-88d952b2d13a"
   },
   "outputs": [],
   "source": [
    "\n",
    "N_EPOCHS=20 # maximum number of epochs to train the model\n",
    "PATIENCE=3 # constant that controls the Early Stopping mechanism of the grid search for the GRU model\n",
    "\n",
    "# make sure it is a list\n",
    "if type(SUBREDDITS) != list:\n",
    "  SUBREDDITS = list(SUBREDDITS)\n",
    "\n",
    "if TRAIN:\n",
    "  if len(SUBREDDITS) == 4:\n",
    "    SRC_DATASET='all'\n",
    "  if len(SUBREDDITS) == 1:\n",
    "    SRC_DATASET=SUBREDDITS[0]\n",
    "\n",
    "if TEST:\n",
    "  SRC_DATASET='all'\n",
    "#   SRC_DATASET='whatisthisthing'  \n",
    "#   SRC_DATASET=SUBREDDITS[0] # use model trained with own dataset\n",
    "\n",
    "PREFIX ='random'\n",
    "if STRATIFIED:\n",
    "  PREFIX += '_strat'\n",
    "\n",
    "results_filename=f'{RESULTSDIR}{PREFIX}_{SRC_DATASET}_b{int(2./BIN_WIDTH):02}_f{FILTERED}_n{N_EPOCHS}_p{PATIENCE}_mGRUv1.pkl'\n",
    "\n",
    "#results_filename=f'{DATADIR}{PREFIX}_{SRC_DATASET}_b{int(2./BIN_WIDTH):02}_f{FILTERED}_n{N_EPOCHS}_p{PATIENCE}_mGRUv1.pkl'\n",
    "print(f'Results file is {results_filename}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000,
     "referenced_widgets": [
      "55d1f3b86d8340569bae6a15b41aec03",
      "39628f0a8a194d3a9ee3444855cb6f74",
      "3316b0f32b55459b935d26e6c48e7ee4",
      "c6c2f2e7d58b4dadb4f74634026258dd",
      "7795ad26f8ea41ffaa3c594368366f19",
      "85a7151619b24a2fab13311fa9aa9d9a",
      "6502b37cbea84f7caa46809835f87df6",
      "41759595873a4053b9989b65b87d9877",
      "c00efd99693d493380ed14311ba9e60c",
      "52f51fa792864e0ea8e8e19ad305cc17",
      "ce7f42e2ebbb47c6a833373ca9cf3bc1",
      "2c652df58ecf444aa8674b7d85bd22a1",
      "e68fdc498d36479fbc8eddc33072d2d1",
      "93481f22f8ef4f5dafb8163c10e8548a",
      "e7678cf453504b2abbdc9248c2321eac",
      "a78c0fe71c5c43149aea5054fcda6acc",
      "e933ea33c0794790bdd61898adca5f36",
      "376b1768105e416fa0288d763c7e24c5",
      "ff25d80f657c420ab67b73a09e7c93aa",
      "28d291d02eae48ce9db267d818f01643",
      "807b5764492646ea9d85d1e1fcfae51e",
      "3cc198966c94474f852cb61a1b8e5b19",
      "1ab015a4bab24295905e7762bc0c8e1a",
      "900d3201900e463e9359b01252a56821",
      "a68f26679c0243159d9233a2c38b6579",
      "8beb51b17bb34c80a8a494d1115a279b",
      "ed9401335ebf482185b42d096dfcd665",
      "723fd87edfa74e4299ae5078735d081e",
      "012c427cbd34484dbccfcc17859486b3",
      "6d298d1b1af4407897041958d2d7cdee",
      "629eb24f8c204a029a4af0b3ff2fe7aa",
      "1981e9b6e2b94fbbb706c0a106d642cd",
      "7fb1a5b471974905946d1f95110600d1",
      "df97ed292e8c4b70aee9b3632ec63a12",
      "67dfec88be88456499b20bfe6e295a9f",
      "7e14b73da03742599230d92cf1ec0ad4",
      "2330c93336ee4093a3e5667247de5551",
      "3ba97b1cee714239adad8e2f2cd9bdf1",
      "11fbf9e4df024e73957206c43984487d",
      "d9bb357dbd5f473983df2f4bad28ccc6"
     ]
    },
    "executionInfo": {
     "elapsed": 1598,
     "status": "ok",
     "timestamp": 1619020807799,
     "user": {
      "displayName": "Henrique Silva",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh0sTZzx8mKsKV9-j3HdK09d6njRJoJnyBA8jbau1I=s64",
      "userId": "15429421149160718308"
     },
     "user_tz": 180
    },
    "id": "xZcTVAmavrWq",
    "outputId": "f9527cdd-454f-4a75-a4e2-cd12e070aaeb"
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "GRID_SEARCH = False\n",
    "USE_PAPER_SETUP = True\n",
    "assert GRID_SEARCH != USE_PAPER_SETUP\n",
    "\n",
    "if TRAIN:\n",
    "\n",
    "    # test_params\n",
    "    if USE_PAPER_SETUP:\n",
    "        hidden_size_list     = [64]\n",
    "        bidirectional_list   = [False]\n",
    "        nlayers_dropout_list = [(2,0.0)]\n",
    "    if GRID_SEARCH:\n",
    "        hidden_size_list     = [4, 16, 64]\n",
    "        bidirectional_list   = [False, True]\n",
    "        nlayers_dropout_list = [(1,0.)] + [(2,prob) for prob in [0., .1, .2, .5]]\n",
    "        \n",
    "    train_criteria=[WeightedL1Loss( subreddit2weights['all'],BIN_WIDTH, MIN_VALUE)]\n",
    "    test_criteria =[WeightedL1Loss( subreddit2weights['all'],BIN_WIDTH, MIN_VALUE),\n",
    "              WeightedMSELoss(subreddit2weights['all'],BIN_WIDTH, MIN_VALUE),\n",
    "              nn.L1Loss(), nn.MSELoss()]\n",
    "\n",
    "    results_df = grid_search_train(train_loader,valid_loader,hidden_size_list, bidirectional_list, nlayers_dropout_list, train_criteria, test_criteria, results_filename,EMBEDDING_DIM=EMBEDDING_DIM, PLOT=False,N_EPOCHS=N_EPOCHS,PATIENCE=PATIENCE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 736
    },
    "executionInfo": {
     "elapsed": 781,
     "status": "ok",
     "timestamp": 1619043965779,
     "user": {
      "displayName": "Henrique Silva",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh0sTZzx8mKsKV9-j3HdK09d6njRJoJnyBA8jbau1I=s64",
      "userId": "15429421149160718308"
     },
     "user_tz": 180
    },
    "id": "djNhmv3gmcAD",
    "outputId": "cc78614f-50e9-45d4-a525-485f4a2a7cd5"
   },
   "outputs": [],
   "source": [
    "# show best results\n",
    "print(f'Loading {results_filename}...')\n",
    "results_df = pd.read_pickle(results_filename)\n",
    "\n",
    "results_df_view = pd.concat((results_df,results_df['params'].apply(pd.Series)), axis=1).drop(\n",
    "    columns=['params','dropout_out','input_size','output_size','uses_two_series_as_input','para','para_best'])\n",
    "\n",
    "results_df_view.sort_values('WeightedL1Loss').head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 297
    },
    "executionInfo": {
     "elapsed": 859,
     "status": "ok",
     "timestamp": 1619004266692,
     "user": {
      "displayName": "Fabricio Murai",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gg5cssoFD3bf2OVCNPF9ZU6hPGXYZtax8bHszJ-Rgc=s64",
      "userId": "04196534601882645786"
     },
     "user_tz": 180
    },
    "id": "vNC-65BYqJF7",
    "outputId": "67c57948-4b6f-47d6-cf8d-f5dd2ef831bb"
   },
   "outputs": [],
   "source": [
    "# plot relationship between losses\n",
    "results_df.plot.scatter('WeightedL1Loss_best', 'WeightedL1Loss')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_GmJY2KaDGAn"
   },
   "source": [
    "### Load the best model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 4191,
     "status": "ok",
     "timestamp": 1619020830217,
     "user": {
      "displayName": "Henrique Silva",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh0sTZzx8mKsKV9-j3HdK09d6njRJoJnyBA8jbau1I=s64",
      "userId": "15429421149160718308"
     },
     "user_tz": 180
    },
    "id": "S2kTb4nxnKR0",
    "outputId": "70bf0675-0261-4439-cd25-07b18b47ce42"
   },
   "outputs": [],
   "source": [
    "# load best model\n",
    "\n",
    "RETRAIN = False\n",
    "USE_LOSS_BEST = False\n",
    "USE_BEST_VALIDATION = True\n",
    "criterion = WeightedL1Loss(subreddit2weights['all'], BIN_WIDTH, MIN_VALUE)\n",
    "#criterion = nn.MSELoss()\n",
    "criterion_name = lambda x: x.__class__.__name__.split('.')[-1]\n",
    "\n",
    "\n",
    "# see if the file with best params is available\n",
    "# results_df = pd.read_pickle(results_filename)\n",
    "\n",
    "if USE_LOSS_BEST:\n",
    "    best_result = results_df.loc[results_df[criterion_name(criterion)+'_best'].argmin()]\n",
    "else:\n",
    "    best_result = results_df.loc[results_df[criterion_name(criterion)].argmin()]\n",
    "\n",
    "# uncomment to load specific model instead\n",
    "# best_result = results_df.iloc[0]\n",
    "if RETRAIN:\n",
    "    reset_seeds()\n",
    "    model = models.GRUSentiment(best_result.params)\n",
    "    _, valid_loss, _ = train_over_nepochs( model, train_loader, valid_loader,\n",
    "                                        criterion=criterion, device=device,\n",
    "                                        patience=3, n_epochs=N_EPOCHS)\n",
    "\n",
    "    if USE_BEST_VALIDATION:\n",
    "        model.load_state_dict(torch.load('checkpoint.pt', map_location=lambda storage, loc: storage))\n",
    "\n",
    "else:\n",
    "    model = models.GRUSentiment(best_result.params)\n",
    "    if USE_BEST_VALIDATION:\n",
    "        model.load_state_dict(best_result.para)\n",
    "    else:\n",
    "        model.load_state_dict(best_result.para_best)\n",
    "\n",
    "model.to(device)\n",
    "\n",
    "test_loss, outputs = evaluate(model, iter(test_loader), criterion=criterion, device=device, return_predictions=True)\n",
    "model_yhat = outputs[0][0].cpu().numpy() # extract data from outputs\n",
    "ilocs = outputs[0][2].cpu().numpy()\n",
    "model_series = pd.Series(model_yhat.ravel(),index=ilocs)\n",
    "# del outputs\n",
    "\n",
    "# print model params\n",
    "print(best_result.params)\n",
    "print(f'Test loss: {test_loss:.3f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SJPMs6SyvAgk"
   },
   "outputs": [],
   "source": [
    "tmp_df = get_baselines_df(test_loader, score_s, score_m)\n",
    "tmp_df['model'] = model_series\n",
    "tmp_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "R6OH-uXnVIyt"
   },
   "source": [
    "## **3.2** Training XGBoost Regressor as a baseline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "if ON_COLAB:\n",
    "    from xgb_utils import plot_model_error, hyperParameterTuning_xgb, getKeysByValue, grid_search_xgb\n",
    "else:    \n",
    "    from src.xgb_utils import plot_model_error, hyperParameterTuning_xgb, getKeysByValue, grid_search_xgb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "S7BaaVLuUDrG"
   },
   "source": [
    "### XGB optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HNG79n01riaF"
   },
   "outputs": [],
   "source": [
    "if USE_XGB:\n",
    "    if TRAIN:\n",
    "        # LOADING DATA\n",
    "        print(\"Loading data...\")\n",
    "        batch_train = next(iter(train_loader_xgb))\n",
    "        batch_val   = next(iter(valid_loader_xgb))\n",
    "        batch_test  = next(iter(test_loader_xgb))\n",
    "\n",
    "        X_train,_,y_train,_ = batch_train\n",
    "        X_val,  _,y_val,  _ = batch_val\n",
    "        X_test, _,y_test, _ = batch_test\n",
    "\n",
    "        # CONVERTING TO NUMPY ARRAYS\n",
    "        X_train = X_train.to('cpu').numpy()\n",
    "        y_train = y_train.to('cpu').numpy()\n",
    "        X_val   = X_val.to('cpu').numpy()\n",
    "        y_val   = y_val.to('cpu').numpy()\n",
    "        X_test  = X_test.to('cpu').numpy()\n",
    "        y_test  = y_test.to('cpu').numpy()   \n",
    "\n",
    "        # DEFINING WEIGHTS\n",
    "        bin_weights = subreddit2weights['all'].cpu().numpy()\n",
    "        eval_set = [(X_train,y_train),(X_test,y_test),(X_val,y_val)]\n",
    "\n",
    "        if GRID_SEARCH_XGB:\n",
    "            param_tuning_xgb_mse = {\n",
    "                'learning_rate': [0.001,0.01, 0.1],\n",
    "                'max_depth': [1, 3, 5],\n",
    "                'min_child_weight': [1, 3, 5],\n",
    "                'subsample': [0.5, 0.7],\n",
    "                'colsample_bytree': [0.5, 0.7],\n",
    "                'n_estimators' : [100, 200, 500],\n",
    "                'objective': ['reg:squarederror']\n",
    "            }\n",
    "\n",
    "            print(\"Starting grid search...\")\n",
    "            best_parameters_mse,_ = grid_search_xgb(param_tuning_xgb_mse,early_stop=3,X_train=X_train,y_train=y_train,X_val=X_val,y_val=y_val,n_jobs=4)\n",
    "            best_parameters_mse['objective'] = 'reg:squarederror'\n",
    "            print(f\"Best parameters for the MSE model found by grid search:{best_parameters_mse}\") \n",
    "            best_para_xgb = open(MODELSDIR + \"best_para_xgb_mse.pkl\", \"wb\")\n",
    "            pickle.dump(best_parameters_mse, best_para_xgb)\n",
    "            best_para_xgb.close()\n",
    "\n",
    "            print(\"Fitting the model with MSE\")\n",
    "            xgb_model_mse = XGBRegressor(**best_parameters_mse)\n",
    "            xgb_model_mse.n_jobs = 4\n",
    "\n",
    "            xgb_model_mse.fit(X_train,y_train,eval_set=eval_set,early_stopping_rounds=3)\n",
    "            print(f\"Best ntree_limit for the MSE model:{xgb_model_mse.best_ntree_limit}\")\n",
    "            f = open(MODELSDIR + \"best_ntree_mse.txt\", \"w\")\n",
    "            f.write(str(xgb_model_mse.best_ntree_limit))\n",
    "            f.close()\n",
    "\n",
    "            xgb_model_mse.save_model(MODELSDIR + \"xgb_mse.model\")    \n",
    "\n",
    "            if PLOT_XGB:\n",
    "                #MSE MODEL\n",
    "                plot_model_error(xgb_model_mse,'rmse')      \n",
    "\n",
    "            #EVAL\n",
    "            print(\"Evaluating model...\")\n",
    "            test_criteria =[WeightedL1Loss( subreddit2weights['all'],BIN_WIDTH,MIN_VALUE),\n",
    "                            nn.L1Loss(), nn.MSELoss()] \n",
    "\n",
    "            print(\"MSE Model test error:\")\n",
    "            for criterion in test_criteria:\n",
    "                test_error = criterion(torch.FloatTensor(xgb_model_mse.predict(X_test)).to(device),torch.FloatTensor(y_test).to(device))\n",
    "                print(criterion.__class__.__name__,test_error)\n",
    "            \n",
    "        else:\n",
    "            print(\"Loading best parameters...\")\n",
    "            try:\n",
    "                best_parameters_mse = open(MODELSDIR + \"best_para_xgb_mse.pkl\", \"rb\")\n",
    "                best_parameters_mse = pickle.load(best_parameters_mse)\n",
    "                print(\"Successfully loaded best parameters for the MSE model\")\n",
    "            except:\n",
    "                print(\"MSE model parameters file not found, using standard ones instead\")\n",
    "                best_parameters_mse = {\n",
    "            'learning_rate': 0.1,\n",
    "            'max_depth':  5,\n",
    "            'min_child_weight': 3,\n",
    "            'subsample': 0.5,\n",
    "            'colsample_bytree': 0.5,\n",
    "            'n_estimators' :  200,\n",
    "            'objective': 'reg:squarederror'}            \n",
    "\n",
    "            print(\"Fitting the model with MSE\")\n",
    "            xgb_model_mse = XGBRegressor(**best_parameters_mse, n_jobs=4)\n",
    "            xgb_model_mse.fit(X_train,y_train,eval_set=eval_set,early_stopping_rounds=3)\n",
    "\n",
    "            print(f\"Best ntree_limit for the model:{xgb_model_mse.best_ntree_limit}\")\n",
    "            f = open(MODELSDIR + \"best_ntree_mse.txt\", \"w\")\n",
    "            f.write(str(xgb_model_mse.best_ntree_limit))\n",
    "            f.close()\n",
    "\n",
    "            xgb_model_mse.save_model(MODELSDIR + \"xgb_mse.model\")     \n",
    "\n",
    "            if PLOT_XGB:\n",
    "                #MSE MODEL\n",
    "                plot_model_error(xgb_model_mse,'rmse')\n",
    "\n",
    "            #EVAL\n",
    "            print(\"Evaluating model...\")\n",
    "            print(\"-------------------------\")\n",
    "            test_criteria =[WeightedL1Loss( subreddit2weights['all'],BIN_WIDTH,MIN_VALUE),\n",
    "                            nn.L1Loss(), nn.MSELoss()] \n",
    "                            \n",
    "            print(\"MSE Model test error:\")\n",
    "            for criterion in test_criteria:\n",
    "                test_error = criterion(torch.FloatTensor(xgb_model_mse.predict(X_test)).to(device),torch.FloatTensor(y_test).to(device))\n",
    "                print(criterion.__class__.__name__,test_error)\n",
    "\n",
    "\n",
    "    if TEST:\n",
    "        # LOADING DATA\n",
    "        batch_test = next(iter(test_loader_xgb))\n",
    "        X_test,_,y_test,_ = batch_test\n",
    "        # CONVERTING TO NUMPY ARRAYS\n",
    "        X_test = X_test.to('cpu').numpy()\n",
    "        y_test = y_test.to('cpu').numpy() \n",
    "        \n",
    "        #LOADING MODELS\n",
    "        print(\"Loading trained xgb model...\")\n",
    "        xgb_model_mse = XGBRegressor()\n",
    "        xgb_model_mse.load_model(MODELSDIR + \"xgb_mse.model\")\n",
    "        \n",
    "        \n",
    "        best_ntree_mse = open(MODELSDIR + \"best_ntree_mse.txt\", \"r\")\n",
    "        best_ntree_mse = int(best_ntree_mse.read())  \n",
    "\n",
    "        print(\"Model and parameters loaded\")\n",
    "        print(\"-------------------------\")\n",
    "\n",
    "        #EVAL\n",
    "        print(\"Testing model...\")\n",
    "        print(\"-------------------------\")\n",
    "        test_criteria =[WeightedL1Loss( subreddit2weights['all'],BIN_WIDTH,MIN_VALUE),\n",
    "                        nn.L1Loss(), nn.MSELoss()] \n",
    "                        \n",
    "        print(\"MSE Model test error:\")\n",
    "        for criterion in test_criteria:\n",
    "            test_error = criterion(torch.FloatTensor(xgb_model_mse.predict(X_test)).to(device),torch.FloatTensor(y_test).to(device))\n",
    "            print(criterion.__class__.__name__,test_error)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZUgz6Zj8WFgW"
   },
   "source": [
    "## **4**. Comparing the results with simple baselines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if ON_COLAB:\n",
    "    from utils import compute_error\n",
    "else:\n",
    "    from src.utils import compute_error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 142
    },
    "executionInfo": {
     "elapsed": 502,
     "status": "ok",
     "timestamp": 1618601090693,
     "user": {
      "displayName": "Henrique Silva",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh0sTZzx8mKsKV9-j3HdK09d6njRJoJnyBA8jbau1I=s64",
      "userId": "15429421149160718308"
     },
     "user_tz": 180
    },
    "id": "jMsAYYKxkcSX",
    "outputId": "de83e1eb-dabe-47f5-a0f9-78558acfbd76"
   },
   "outputs": [],
   "source": [
    "# Computes the loss metrics for all models\n",
    "pred_column_names = ['unchanged', 'mean','last']\n",
    "loss2title = {'L1Loss':'L1 Loss','MSELoss':'MSE Loss','WeightedL1Loss':'Weighted L1 Loss'}\n",
    "subreddit='all'\n",
    "\n",
    "criteria=[WeightedL1Loss(subreddit2weights[subreddit],BIN_WIDTH,MIN_VALUE)]\n",
    "criteria_names = [criterion.__class__.__name__.split('.')[-1] for criterion in criteria]\n",
    "\n",
    "test_loss = dict()\n",
    "for pred_column_name in pred_column_names:\n",
    "    test_loss[pred_column_name] = compute_error(tmp_df, pred_column_name, criteria, device)\n",
    "\n",
    "df = pd.DataFrame.from_dict(test_loss,orient='index')\n",
    "df.columns = df.columns.map(loss2title)\n",
    "df.index = df.index.map(str.upper)\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "e90fuJJsGBpP"
   },
   "source": [
    "### Generate results table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 448,
     "status": "ok",
     "timestamp": 1617887807565,
     "user": {
      "displayName": "Henrique Silva",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh0sTZzx8mKsKV9-j3HdK09d6njRJoJnyBA8jbau1I=s64",
      "userId": "15429421149160718308"
     },
     "user_tz": 180
    },
    "id": "uvGxRB7iEPW9",
    "outputId": "83593dfe-961a-42ad-b691-a6e0eec45adc"
   },
   "outputs": [],
   "source": [
    "if TRAIN:\n",
    "    subreddit2range2 = {}\n",
    "    for subreddit,v in subreddit2range.items():\n",
    "        if type(v) == slice:\n",
    "            inds = [ind for ind in tmp_df.index if v.start <= ind < v.stop]\n",
    "        else:\n",
    "            inds = list(tmp_df.index)\n",
    "        subreddit2range2[subreddit] = inds\n",
    "elif TEST:\n",
    "    test_locs_df = pd.DataFrame(test_locs)\n",
    "    subreddit2range2 = {'all': test_locs_df.index}\n",
    "    for subreddit in SUBREDDITS:\n",
    "      subreddit2range2[subreddit] = test_locs_df[test_locs_df[0] == subreddit].index\n",
    "\n",
    "\n",
    "pred_column_names = ['unchanged', 'mean','last','model']\n",
    "loss2title = {'L1Loss':'L1 Loss','MSELoss':'MSE Loss','WeightedL1Loss':'Weighted L1 Loss'}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# tmp_df['model'] = model_yhat\n",
    "# tmp_df['xgboost'] = xgb_model_l1.predict(X_test)\n",
    "# tmp_df['xgb-mse'] = xgb_model_mse.predict(X_test)\n",
    "\n",
    "\n",
    "df_list = []\n",
    "for subreddit in ['all']+SUBREDDITS:\n",
    "  criteria=[WeightedL1Loss(subreddit2weights[subreddit],BIN_WIDTH,MIN_VALUE)]\n",
    "  #criteria=[nn.L1Loss(),nn.MSELoss()]\n",
    "  criteria_names = [criterion.__class__.__name__.split('.')[-1] for criterion in criteria]\n",
    "\n",
    "  test_loss = dict()\n",
    "  for pred_column_name in pred_column_names:\n",
    "    test_loss[pred_column_name] = compute_error(tmp_df.loc[subreddit2range2[subreddit]], pred_column_name, criteria, device)\n",
    "\n",
    "  df = pd.DataFrame.from_dict(test_loss,orient='index')\n",
    "  df.columns = df.columns.map(loss2title)\n",
    "  df.index = df.index.map(str.upper)\n",
    "  df_list.append(df)\n",
    "\n",
    "test_df = pd.concat(df_list,keys=['ALL']+[subreddit2title[subreddit.lower()] for subreddit in SUBREDDITS],axis=1)\n",
    "test_df = test_df.swaplevel(0,1,axis=1).sort_index(1)\n",
    "n=len(SUBREDDITS)+1\n",
    "test_df = test_df.round(3)\n",
    "test_df = test_df.loc[['UNCHANGED','MEAN','LAST','MODEL'],:]\n",
    "test_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 724,
     "status": "ok",
     "timestamp": 1617651565145,
     "user": {
      "displayName": "Henrique Silva",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh0sTZzx8mKsKV9-j3HdK09d6njRJoJnyBA8jbau1I=s64",
      "userId": "15429421149160718308"
     },
     "user_tz": 180
    },
    "id": "oZTROBs2u-33",
    "outputId": "43a41e86-1867-491d-9257-e9ad240f51f0"
   },
   "outputs": [],
   "source": [
    "formatters = []\n",
    "for column in test_df.columns:\n",
    "  formatters.append(lambda x, column=column: '\\\\textbf{%s}'%( ('%.3f' % float(x)).lstrip('0') ) if float(x) == test_df[column].min() else ('%.3f'% float(x)).lstrip('0') )\n",
    "formatters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 442,
     "status": "ok",
     "timestamp": 1617651565989,
     "user": {
      "displayName": "Henrique Silva",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh0sTZzx8mKsKV9-j3HdK09d6njRJoJnyBA8jbau1I=s64",
      "userId": "15429421149160718308"
     },
     "user_tz": 180
    },
    "id": "CaFLHseZqJPo",
    "outputId": "a3e28912-6f4a-4f7c-edd2-2a3875c71fd2"
   },
   "outputs": [],
   "source": [
    "print(test_df.to_latex(\n",
    "    escape=False,column_format='l|ccccc|ccccc|ccccc',multicolumn_format='c',\n",
    "    label='tab:test_loss', formatters=formatters,\n",
    "    caption='Results w.r.t.\\\\ L1, MSE and proposed Weighted L1 Loss, which gives more weight to more extreme responses. '+\n",
    "    'Model is trained on entire 2017 dataset whereas \\\\textsc{Model-Subreddit} is trained on target subreddit (column). '+\n",
    "    '\\\\textsc{Mean} yields lowest L1 and MSE, but Model outperforms baselines w.r.t.\\\\ Weighted L1 loss.'\n",
    "    ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 227
    },
    "executionInfo": {
     "elapsed": 1251,
     "status": "ok",
     "timestamp": 1616437162358,
     "user": {
      "displayName": "Henrique Silva",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh0sTZzx8mKsKV9-j3HdK09d6njRJoJnyBA8jbau1I=s64",
      "userId": "15429421149160718308"
     },
     "user_tz": 180
    },
    "id": "rckqwFtlu0-y",
    "outputId": "6ed75af2-7613-418d-d0d7-3d3eb455568c"
   },
   "outputs": [],
   "source": [
    "pred_column_names2 = ['final score']+pred_column_names\n",
    "fig, axs = plt.subplots(1,len(pred_column_names2), figsize=(3*len(pred_column_names2),3))\n",
    "for ax, column_name in zip(axs,pred_column_names2):\n",
    "  ax = tmp_df[column_name].hist(ax=ax,bins=np.arange(-1.0,1.0,.1))\n",
    "  ax.set_xlim(-1,1)\n",
    "  ax.set_ylim(0,700)\n",
    "  ax.set_title(column_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EZ0qQnrfIN4V"
   },
   "source": [
    "# **5**. Plotting the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "59ZyHgIm6V3b"
   },
   "outputs": [],
   "source": [
    "# TODO: add doc to function\n",
    "def plot_prediction_density(tmp_df, pred_column_names, subreddit=None, axs=None, choice_inds=None, cmap=plt.cm.gist_earth_r, vmax=None):\n",
    "  if subreddit is not None:\n",
    "    tmp_df = tmp_df.loc[subreddit2range2[subreddit]]\n",
    "\n",
    "  xmin = ymin = -1.\n",
    "  xmax = ymax =  1.\n",
    "\n",
    "  # deltaX = (max(x) - min(x))/10\n",
    "  # deltaY = (max(y) - min(y))/10\n",
    "  # xmin = min(x) - deltaX\n",
    "  # xmax = max(x) + deltaX\n",
    "  # ymin = min(y) - deltaY\n",
    "  # ymax = max(y) + deltaY\n",
    "  # xx, yy = np.mgrid[xmin:xmax:50j, ymin:ymax:50j]\n",
    "\n",
    "  xx, yy = np.mgrid[xmin:xmax:50j, ymin:ymax:50j]\n",
    "  positions = np.vstack([xx.ravel(), yy.ravel()])\n",
    "  reset_seeds()\n",
    "\n",
    "  if choice_inds is None:\n",
    "    choice_inds = np.random.choice(len(tmp_df), 100)\n",
    "  print(tmp_df.iloc[choice_inds]['thread ncom'].max())\n",
    "  # choice_inds = np.random.choice(nz[0], 100)\n",
    "\n",
    "  if axs is None:\n",
    "   fig, axs = plt.subplots(2,2,figsize=(3*2,3*2),sharex=True,sharey=True)\n",
    "   plt.subplots_adjust(wspace=.15,hspace=.15)\n",
    "\n",
    "  for idx, (ax, column_name) in enumerate(zip(axs.ravel(),pred_column_names)):\n",
    "    loss_array = tmp_df[['final score', column_name]].values\n",
    "    kernel = stats.gaussian_kde(loss_array.T)\n",
    "    Z = np.reshape(kernel(positions).T, xx.shape)\n",
    "    # print(subreddit,idx,Z.min(),Z.max())\n",
    "\n",
    "    if vmax is not None:\n",
    "      levels = np.linspace(vmax/8.,vmax,8)\n",
    "      # levels = np.insert(np.linspace(vmax*2/7.,vmax,7),0,0)\n",
    "      vmin=vmax/8.\n",
    "    else:\n",
    "      levels = 8\n",
    "      vmin=None\n",
    "\n",
    "    ax.plot([xmin, xmax], [ymin, ymax], ls='-', c='k',alpha=.5)#,transform=ax.transAxes)\n",
    "\n",
    "    # cfset = ax.contourf(xx, yy, Z, cmap=cmap, levels=levels)\n",
    "    hb = ax.imshow(np.rot90(Z), cmap=cmap, aspect='auto',\n",
    "              extent=[xmin, xmax, ymin, ymax], vmin=0,vmax=vmax)\n",
    "    # cset = ax.contour(xx, yy, Z, colors='k', linewidths=.25, levels=levels)\n",
    "    # plt.colorbar(hb, ax=ax)\n",
    "    # ax.clabel(cset, inline=1, fontsize=8)\n",
    "\n",
    "\n",
    "    ncom = tmp_df['branch ncom'].values if column_name.startswith('branch') else tmp_df['thread ncom'].values\n",
    "    ax.scatter(loss_array[choice_inds,0], loss_array[choice_inds,1], 2*ncom[choice_inds], 'gray', edgecolor='k')\n",
    "\n",
    "    ax.set_xlim([xmin, xmax])\n",
    "    ax.set_ylim([ymin, ymax])\n",
    "    # ax.yaxis.tick_right()\n",
    "    # ax.set_xticks(np.arange(-1.,1.,.5))\n",
    "    ax.set_yticks(np.arange(-1.,1.01,.5))\n",
    "    ax.grid(ls='--', c='gray', alpha=.5)\n",
    "\n",
    "    # ax.axhline((ymin+ymax)/2,color='k',alpha=.5)\n",
    "    # ax.axvline((xmin+xmax)/2,color='k',alpha=.5)\n",
    "    # x = np.array([xmin,xmax])\n",
    "    # ax.plot(x, 1.5*x, ls='--', c='k',alpha=.5)\n",
    "    # ax.plot(x, 0.5*x, ls='--', c='k',alpha=.5)\n",
    "    # ax.set_title(column_name)\n",
    "\n",
    "\n",
    "  # divider = make_axes_locatable(plt.gca())\n",
    "  # cax = divider.append_axes(\"right\", size=\"2%\", pad=0.00)\n",
    "  # cbar = fig.colorbar(hb, cax=cax)\n",
    "\n",
    "  # fig.text(0.5, 0.04, r'EmT($c_n$) (true value)', ha='center')\n",
    "  # fig.text(0.04, 0.5, 'prediction', va='center', rotation='vertical')\n",
    "\n",
    "  # plt.savefig(RESULTSDIR+subreddit+'.pdf')\n",
    "  return hb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 862
    },
    "executionInfo": {
     "elapsed": 4455,
     "status": "ok",
     "timestamp": 1617889040733,
     "user": {
      "displayName": "Henrique Silva",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Gh0sTZzx8mKsKV9-j3HdK09d6njRJoJnyBA8jbau1I=s64",
      "userId": "15429421149160718308"
     },
     "user_tz": 180
    },
    "id": "8aFVmlL-b-cq",
    "outputId": "57bad078-aa88-4785-f580-1615e1630c2e"
   },
   "outputs": [],
   "source": [
    "# plot density\n",
    "vmax_list = [2, 3.2, 2., 2.]\n",
    "vmax_list = [2.4, 2.4, 2.4, 2.4]\n",
    "\n",
    "choice_inds = None\n",
    "\n",
    "fig, axs = plt.subplots(4,4,figsize=(3*4,3*4),sharex=True,sharey=True)\n",
    "plt.subplots_adjust(wspace=.1,hspace=.1)\n",
    "\n",
    "for ix, subreddit in enumerate(SUBREDDITS):\n",
    "  hb = plot_prediction_density(tmp_df, pred_column_names[1:], subreddit, axs[ix,:], choice_inds=choice_inds, vmax=vmax_list[ix])\n",
    "\n",
    "fig.text(0.5 , 0.08, r'EmT($c_n$) (true value)', ha='center', fontsize='x-large')\n",
    "fig.text(0.0, 0.5 , 'prediction', va='center', rotation='vertical', fontsize='x-large')\n",
    "\n",
    "cbaxes = fig.add_axes([0.2, 0.93, .6, 0.01])\n",
    "\n",
    "for ax, col in zip(axs[0], pred_column_names[1:]):\n",
    "  ax.set_title(col.upper(), size='large')\n",
    "\n",
    "for ax, row in zip(axs[:,0], SUBREDDITS):\n",
    "  # ax.yaxis.set_label_position(\"right\")\n",
    "  ax.set_ylabel(subreddit2title[row.lower()], rotation=0, size='large', horizontalalignment='right')\n",
    "\n",
    "cb = plt.colorbar(hb, cax = cbaxes, orientation='horizontal')\n",
    "cb.ax.set_title('density', size='large')\n",
    "plt.savefig('heatmaps.pdf')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DTiiIhnY3ZlH"
   },
   "source": [
    "# **6**. Case Study"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ToQJkqGNiTAQ"
   },
   "outputs": [],
   "source": [
    "# TODO: explain what we are reading\n",
    "REDDIT='SuicideWatch'\n",
    "suffix = '.pkl'\n",
    "with open(DATADIR+REDDIT+'_post2data'+suffix,'rb') as infile:\n",
    "  post2data = pickle.load(infile)\n",
    "\n",
    "with open(DATADIR+REDDIT+'_comment2data'+suffix,'rb') as infile:\n",
    "  comment2data = pickle.load(infile)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 80
    },
    "id": "k7zqUhb4y02d",
    "outputId": "cb5f0cea-bbef-4e02-8d4d-0294e52ce606"
   },
   "outputs": [],
   "source": [
    "# 5855 (original iloc = 58762)\n",
    "z = tmp_df.loc[subreddit2range2['SuicideWatch']]\n",
    "z[(z['final score'] > 0.75) & (z['model'] < -0.75)].sort_values('final score',ascending=True).head(20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 390
    },
    "id": "Nwn1vtGDpRwH",
    "outputId": "c7785207-0144-4886-9c59-a96d4ee21fed"
   },
   "outputs": [],
   "source": [
    "# 374 (original iloc = 4001)\n",
    "z = tmp_df.loc[subreddit2range2['Anxiety']]\n",
    "z[(z['final score'] > 0.75) & (z['model'] > 0.6)].sort_values('thread ncom',ascending=True).head(20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ceTPXvaRh5e5"
   },
   "outputs": [],
   "source": [
    "def print_post_and_comments(index, suffix='', infile_extension='.pkl'):\n",
    "  post = post_df.iloc[index]\n",
    "\n",
    "  print(f'Post id: {post.name[1]}')\n",
    "  print(f'Author (score {(post[\"score\"][0]-0.5)*2}):')\n",
    "\n",
    "  subreddit = post.name[0]\n",
    "  with open(f'{DATADIR}{subreddit}_comment2data{suffix}{infile_extension}','rb') as infile:\n",
    "    comment2data = pickle.load(infile)\n",
    "\n",
    "  comment_ids = post.comments\n",
    "  pprint.pprint(post.text)\n",
    "  for cx in range(max(0,post.seq_len-10),post.seq_len-1):\n",
    "    print('Author' if post.is_post_author[cx+1] else 'Commenter', end='')\n",
    "    print(f' (score {(post[\"score\"][cx+1])}):')\n",
    "    if 'body' in comment2data[comment_ids[cx]]:\n",
    "      pprint.pprint(comment2data[comment_ids[cx]]['body'])\n",
    "    else:\n",
    "      pprint.pprint(comment2data[comment_ids[cx]]['text'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "JToAeaHz0C21",
    "outputId": "af5691f8-e648-438c-d8da-f08d0581a0fe"
   },
   "outputs": [],
   "source": [
    "# TODO: remove for in all cells below, just set ind to each value\n",
    "#for ind in [5855]: [58762]:\n",
    "ind = np.where(post_df.reset_index().level_1 == '712ctw')[0][0]\n",
    "print(post_df.iloc[ind])\n",
    "print(tmp_df.loc[ind])\n",
    "print_post_and_comments(ind)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 111
    },
    "id": "LKxnE2XKp2uv",
    "outputId": "85c8e635-9a41-4dfa-fa53-3218a94fa09f"
   },
   "outputs": [],
   "source": [
    "# 1336 (original iloc = 13443)\n",
    "z = tmp_df.loc[subreddit2range2['bipolar']]\n",
    "z[(z['final score'] < -0.75) & (z['model'] < 0.1) & (z['model'] > 0)].sort_values('thread ncom',ascending=False).head(20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "L8_trjzesdY8"
   },
   "outputs": [],
   "source": [
    "# 1422 (original iloc = 14185)\n",
    "z[(z['final score'] > -0.5) & (z['model'] < -0.5)].sort_values('thread ncom',ascending=False).head(20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 173
    },
    "id": "5TCslO_M8Vtv",
    "outputId": "18c68189-03f8-4b1e-84d8-dbc88c478df4"
   },
   "outputs": [],
   "source": [
    "#  5090 (original iloc = 50549)\n",
    "z = tmp_df.loc[subreddit2range2['SuicideWatch']]\n",
    "z[(z['final score'] > .6) & (z['final score'] < .75)  & (z['model'] > .6) & (z['model'] < .75)].sort_values('thread ncom',ascending=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "pzPzTgplycrV",
    "outputId": "476d96cf-1957-48b0-a856-ba75931ceab5"
   },
   "outputs": [],
   "source": [
    "# for ind in [5090]: [50549]:\n",
    "ind = np.where(post_df.reset_index().level_1 == '5nr93y')[0][0]\n",
    "print(post_df.iloc[ind])\n",
    "print(tmp_df.loc[ind])\n",
    "print_post_and_comments(ind)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "NTbhpa51aoW3",
    "outputId": "4e8757e0-aef2-47b4-fc1d-c549719e279b"
   },
   "outputs": [],
   "source": [
    "#for ind in [13443]:\n",
    "ind = np.where(post_df.reset_index().level_1 == '5x1ll5')[0][0]\n",
    "print(post_df.iloc[ind])\n",
    "print(tmp_df.loc[ind])\n",
    "print_post_and_comments(ind)\n",
    "print('------------------------------------------------------')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Vmj4PJ5IpGPm"
   },
   "source": [
    "# **7**. Can we predict large fluctuations from one author comment to the next?\n",
    "\n",
    "1. Plot the distribution of changes from one comment to the next.\n",
    "2. Take the largest fluctuations and see how well we predict those. Is this large or small? Are the baselines predicting those better?\n",
    "3. For those that we do well, inspect content of the post, \"last checkpoint of author\", following comment and last comment of author."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 282
    },
    "id": "OnaALYi9pD4V",
    "outputId": "76f6d845-cac8-42c8-e00f-b99b0140d8c9"
   },
   "outputs": [],
   "source": [
    "#TODO: alternate the cells with the descriptions 1,2,3 above\n",
    "fluctuations = []\n",
    "\n",
    "\n",
    "for batch in test_loader:\n",
    "  is_author_value = batch[0][0][0,-1] # this has to be retrieved because the feature was z-normalized\n",
    "  batch_size = batch[0].size(0)\n",
    "\n",
    "for ix in range(batch_size):\n",
    "    ncom = batch[1][ix]\n",
    "    for jx in range(ncom-1,-1,-1):\n",
    "      if batch[0][ix][jx,-1] == is_author_value:\n",
    "          fluctuations.append(batch[2][ix]-(batch[0][ix][jx,-2]*score_s+score_m))\n",
    "          break   \n",
    "\n",
    "fluctuations = np.array(fluctuations)\n",
    "# fluctuation_series = pd.Series(fluctuations, index=test_locs)\n",
    "fluctuation_series = pd.Series(fluctuations)\n",
    "fluctuation_series.hist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "6kJuH9vCvw1m",
    "outputId": "1dc4a534-66fc-4822-bb2a-4ad86c2e6049"
   },
   "outputs": [],
   "source": [
    "fluctuation_series = tmp_df['final score']-tmp_df['unchanged']\n",
    "fluctuation_series.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "qxR8FoaLUV35",
    "outputId": "abd2f179-342c-4a8a-ce8a-68cd70b869c7"
   },
   "outputs": [],
   "source": [
    "# positive variation\n",
    "\n",
    "\n",
    "\n",
    "def get_error_stats(series, lower_bound=None,upper_bound=None):\n",
    "  assert (lower_bound is not None) ^ (upper_bound is not None)\n",
    "  if upper_bound is not None:\n",
    "    selected_inds = series[series < upper_bound].index\n",
    "    print(f'[Upper bound: {upper_bound:.3f}] ',sep='')\n",
    "  if lower_bound is not None:\n",
    "    selected_inds = series[series > lower_bound].index\n",
    "    print(f'[Lower bound: {lower_bound:.3f}] ',sep='')\n",
    "\n",
    "  sub_df = tmp_df.loc[selected_inds]\n",
    "\n",
    "  model_error = (sub_df['final score']-sub_df['model']).abs()\n",
    "  results = {'MODEL L1':model_error.mean()}\n",
    "  for baseline in ['mean','last','xgb']:\n",
    "    baseline_error = (sub_df['final score']-sub_df[baseline]).abs()\n",
    "    succ_fraction = (model_error < baseline_error).mean()\n",
    "    print( f\"Model outperforms {baseline} in {succ_fraction*100:.1f}% of all cases\")\n",
    "    print( f\"Model error: {model_error.mean():.3f}\\t Baseline error: {baseline_error.mean():.3f}\")\n",
    "    results[baseline] = {f'M outp. {baseline.upper()} (%)': succ_fraction,\n",
    "                         f'{baseline.upper()} L1': baseline_error.mean() }\n",
    "\n",
    "  print(f\"Cases: {len(selected_inds)}\\n\\n\")\n",
    "  return results\n",
    "\n",
    "n=len(fluctuation_series)\n",
    "upper_bound=fluctuation_series.sort_values().iloc[int(0.05*n)]\n",
    "lower_bound=fluctuation_series.sort_values().iloc[int(0.95*n)]\n",
    "\n",
    "extreme_stats = {}\n",
    "\n",
    "print('XXXX Fluctuation from previous comment: <5-th and >95-th percentile')\n",
    "extreme_stats['EmT shift $>$ 95th perc.'] = get_error_stats(fluctuation_series, lower_bound=lower_bound)\n",
    "extreme_stats['EmT shift $>$ +1.0'] = get_error_stats(fluctuation_series, lower_bound=1.)\n",
    "extreme_stats['Final EmT $>$ +0.8'] = get_error_stats(tmp_df['final score'], lower_bound=.8)\n",
    "\n",
    "extreme_stats['EmT shift $<$  5th perc.'] = get_error_stats(fluctuation_series, upper_bound=upper_bound)\n",
    "extreme_stats['EmT shift $<$ -1.0'] = get_error_stats(fluctuation_series, upper_bound=-1.)\n",
    "extreme_stats['Final EmT $<$ -0.8'] = get_error_stats(tmp_df['final score'], upper_bound=-.8)\n",
    "\n",
    "print('XXX Fluctuation from previous comment: <-0.5 and >0.5')\n",
    "\n",
    "print('XXX Final score: <0.1 and >0.9')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 227
    },
    "id": "iQAmJYNOkJow",
    "outputId": "794b53ac-3c4e-497a-916b-92ce91a75b4a"
   },
   "outputs": [],
   "source": [
    "extreme_df = pd.DataFrame.from_dict(extreme_stats, orient='index')\n",
    "extreme_df = pd.concat((extreme_df,extreme_df['mean'].apply(pd.Series),extreme_df['last'].apply(pd.Series),extreme_df['xgb'].apply(pd.Series)), axis=1).drop(columns=['last','mean','xgb'])\n",
    "#extreme_df[['MEAN L1','MODEL L1','LAST L1','XGB L1']] = extreme_df[['MEAN L1','MODEL L1','LAST L1']].multiply(2)\n",
    "extreme_df = extreme_df.round(3)\n",
    "extreme_df['M outp. MEAN (%)'] = extreme_df['M outp. MEAN (%)']*100\n",
    "extreme_df['M outp. LAST (%)'] = extreme_df['M outp. LAST (%)']*100\n",
    "extreme_df['M outp. XGB (%)'] = extreme_df['M outp. XGB (%)']*100\n",
    "\n",
    "new_columns = [extreme_df.columns[ix] for ix in [0,2,4,6,1,3,5]] \n",
    "extreme_df = extreme_df[new_columns]\n",
    "extreme_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "fFmXodVuxbd6",
    "outputId": "4d2d3515-d885-4dc2-be50-05d9853a42d6"
   },
   "outputs": [],
   "source": [
    "formatters = []\n",
    "for ix, column in enumerate(extreme_df.columns):\n",
    "  formatters.append(lambda x, column_ix=ix: '\\\\textbf{%s}'%( ('%.1f' % float(x)).lstrip('0') ) if (column_ix >=3 and x > 50) else ('%.3f'% float(x)).lstrip('0') )\n",
    "formatters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "axdxkJRsp8NR",
    "outputId": "58f9b969-ea73-418a-8d32-5f65449bf3df"
   },
   "outputs": [],
   "source": [
    "print(extreme_df.to_latex(\n",
    "    column_format='l|rrrr|ccc',  escape=False, label='tab:extreme_values', formatters=formatters,\n",
    "    caption='Performance on extreme values. Central columns show L1 error for \\\\textsc{Mean}, \\\\textsc{Model} and \\\\textsc{Last}. '+\n",
    "    'First (last) column show fraction of cases (selected from test set according to row description) where \\\\textsc{Model} outperforms '+\n",
    "    '\\\\textsc{Mean} (\\\\textsc{Last}). For extremely positive shifts, \\\\textsc{Last} performs best because a large improvement is often '+\n",
    "    'preceded by a very positive comment. For extremely negative shifts, \\\\textsc{Model} performs best.'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9huinO6ZwFzK"
   },
   "outputs": [],
   "source": [
    "for ind in large_change_inds:\n",
    "  print(ind, branch_data[ind])\n",
    "  index, _, bx = branch_data[ind]\n",
    "  b = post_df.loc[index].valid_branches[bx]\n",
    "  print(index, bx, b)\n",
    "  print(post_df.loc[index].features[b,-2])\n",
    "  print(post_df.loc[index].features[b,-1])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wLxXDy4duJaY"
   },
   "outputs": [],
   "source": [
    "for pred_column_name in pred_column_names:\n",
    "  test_loss[pred_column_name] = compute_error(tmp_df.iloc[large_change_inds], pred_column_name, criteria)\n",
    "pd.DataFrame.from_dict(test_loss, orient='index').round(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aqPcSjZ4OP3r"
   },
   "outputs": [],
   "source": [
    "attn_weights = attn_weights.squeeze(1).cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fnvgR6XjPJs5"
   },
   "outputs": [],
   "source": [
    "count_nz = np.count_nonzero( attn_weights,  axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "h1fkn23wCbfH"
   },
   "outputs": [],
   "source": [
    "attn_weights_sorted = -np.sort(-attn_weights, axis=1)\n",
    "imbalance = (attn_weights_sorted[:,0]-attn_weights_sorted[:,1])\n",
    "imbalance_index = np.argsort(-imbalance)\n",
    "imbalance_index[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "s96tEQ1cQb4f"
   },
   "outputs": [],
   "source": [
    "X = np.vstack((pred_error,imbalance,initial_scores,y,yhat,te_diff, count_nz))\n",
    "X = X.T\n",
    "tmp_df = pd.DataFrame(X, columns=['pred_error','imbalance','initial_score','final_score','prediction','te_diff', 'count_nz'], index=post_ids)\n",
    "tmp_df.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Z3N8Vsgk0yON"
   },
   "outputs": [],
   "source": [
    "pred_error = np.absolute(z[:,0]-z[:,1])\n",
    "cost = pred_error-2*imbalance.cpu().numpy()-0.6*np.absolute(z[:,1])-0.1*(attn_weights[:,0,2].cpu().numpy()>0)\n",
    "index_array = np.argsort(cost)\n",
    "index_array[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Bw2bZTtWDwpk"
   },
   "outputs": [],
   "source": [
    "index = index_array[2]\n",
    "print(z[index])\n",
    "attn_weights[index]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OYnM5XNmBk5e"
   },
   "outputs": [],
   "source": [
    "post_id = post_ids[index]\n",
    "post=post_df.loc[post_id]\n",
    "print(post)\n",
    "#pprint.pprint(post['text'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HeHcA6DXK9Xt"
   },
   "outputs": [],
   "source": [
    "comment = comment_df.loc[post_df.loc[post_id,'comments'][7]]\n",
    "print(comment) \n",
    "pprint.pprint(comment['text'])"
   ]
  }
 ],
 "metadata": {
  "accelerator": "TPU",
  "colab": {
   "collapsed_sections": [
    "OKpd9gQqT1Z-"
   ],
   "machine_shape": "hm",
   "name": "CLEAN_CODE.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "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.9.4"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "012c427cbd34484dbccfcc17859486b3": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": "initial"
     }
    },
    "11fbf9e4df024e73957206c43984487d": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "1981e9b6e2b94fbbb706c0a106d642cd": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "1ab015a4bab24295905e7762bc0c8e1a": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "2330c93336ee4093a3e5667247de5551": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": "initial"
     }
    },
    "28d291d02eae48ce9db267d818f01643": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_900d3201900e463e9359b01252a56821",
      "placeholder": "​",
      "style": "IPY_MODEL_1ab015a4bab24295905e7762bc0c8e1a",
      "value": " 2/? [00:00&lt;00:00, 29.51it/s]"
     }
    },
    "2c652df58ecf444aa8674b7d85bd22a1": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_a78c0fe71c5c43149aea5054fcda6acc",
      "placeholder": "​",
      "style": "IPY_MODEL_e7678cf453504b2abbdc9248c2321eac",
      "value": " 3/? [00:00&lt;00:00, 13.57it/s]"
     }
    },
    "3316b0f32b55459b935d26e6c48e7ee4": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "L0: criterion: ",
      "description_tooltip": null,
      "layout": "IPY_MODEL_85a7151619b24a2fab13311fa9aa9d9a",
      "max": 1,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_7795ad26f8ea41ffaa3c594368366f19",
      "value": 1
     }
    },
    "376b1768105e416fa0288d763c7e24c5": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "39628f0a8a194d3a9ee3444855cb6f74": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "3ba97b1cee714239adad8e2f2cd9bdf1": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "3cc198966c94474f852cb61a1b8e5b19": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "41759595873a4053b9989b65b87d9877": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "52f51fa792864e0ea8e8e19ad305cc17": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "55d1f3b86d8340569bae6a15b41aec03": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_3316b0f32b55459b935d26e6c48e7ee4",
       "IPY_MODEL_c6c2f2e7d58b4dadb4f74634026258dd"
      ],
      "layout": "IPY_MODEL_39628f0a8a194d3a9ee3444855cb6f74"
     }
    },
    "629eb24f8c204a029a4af0b3ff2fe7aa": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "6502b37cbea84f7caa46809835f87df6": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "67dfec88be88456499b20bfe6e295a9f": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "L2: bidirectional: ",
      "description_tooltip": null,
      "layout": "IPY_MODEL_3ba97b1cee714239adad8e2f2cd9bdf1",
      "max": 1,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_2330c93336ee4093a3e5667247de5551",
      "value": 1
     }
    },
    "6d298d1b1af4407897041958d2d7cdee": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "723fd87edfa74e4299ae5078735d081e": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_1981e9b6e2b94fbbb706c0a106d642cd",
      "placeholder": "​",
      "style": "IPY_MODEL_629eb24f8c204a029a4af0b3ff2fe7aa",
      "value": " 2/? [00:20&lt;00:00, 10.30s/it]"
     }
    },
    "7795ad26f8ea41ffaa3c594368366f19": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": "initial"
     }
    },
    "7e14b73da03742599230d92cf1ec0ad4": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_d9bb357dbd5f473983df2f4bad28ccc6",
      "placeholder": "​",
      "style": "IPY_MODEL_11fbf9e4df024e73957206c43984487d",
      "value": " 2/? [00:20&lt;00:00, 10.27s/it]"
     }
    },
    "7fb1a5b471974905946d1f95110600d1": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_67dfec88be88456499b20bfe6e295a9f",
       "IPY_MODEL_7e14b73da03742599230d92cf1ec0ad4"
      ],
      "layout": "IPY_MODEL_df97ed292e8c4b70aee9b3632ec63a12"
     }
    },
    "807b5764492646ea9d85d1e1fcfae51e": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": "initial"
     }
    },
    "85a7151619b24a2fab13311fa9aa9d9a": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "8beb51b17bb34c80a8a494d1115a279b": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "900d3201900e463e9359b01252a56821": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "93481f22f8ef4f5dafb8163c10e8548a": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "a68f26679c0243159d9233a2c38b6579": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_ed9401335ebf482185b42d096dfcd665",
       "IPY_MODEL_723fd87edfa74e4299ae5078735d081e"
      ],
      "layout": "IPY_MODEL_8beb51b17bb34c80a8a494d1115a279b"
     }
    },
    "a78c0fe71c5c43149aea5054fcda6acc": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "c00efd99693d493380ed14311ba9e60c": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_ce7f42e2ebbb47c6a833373ca9cf3bc1",
       "IPY_MODEL_2c652df58ecf444aa8674b7d85bd22a1"
      ],
      "layout": "IPY_MODEL_52f51fa792864e0ea8e8e19ad305cc17"
     }
    },
    "c6c2f2e7d58b4dadb4f74634026258dd": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_41759595873a4053b9989b65b87d9877",
      "placeholder": "​",
      "style": "IPY_MODEL_6502b37cbea84f7caa46809835f87df6",
      "value": " 1/? [00:00&lt;00:00,  4.06it/s]"
     }
    },
    "ce7f42e2ebbb47c6a833373ca9cf3bc1": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "L1: hidden_size: ",
      "description_tooltip": null,
      "layout": "IPY_MODEL_93481f22f8ef4f5dafb8163c10e8548a",
      "max": 1,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_e68fdc498d36479fbc8eddc33072d2d1",
      "value": 1
     }
    },
    "d9bb357dbd5f473983df2f4bad28ccc6": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "df97ed292e8c4b70aee9b3632ec63a12": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "e68fdc498d36479fbc8eddc33072d2d1": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": "initial"
     }
    },
    "e7678cf453504b2abbdc9248c2321eac": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "e933ea33c0794790bdd61898adca5f36": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_ff25d80f657c420ab67b73a09e7c93aa",
       "IPY_MODEL_28d291d02eae48ce9db267d818f01643"
      ],
      "layout": "IPY_MODEL_376b1768105e416fa0288d763c7e24c5"
     }
    },
    "ed9401335ebf482185b42d096dfcd665": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "L2: bidirectional: ",
      "description_tooltip": null,
      "layout": "IPY_MODEL_6d298d1b1af4407897041958d2d7cdee",
      "max": 1,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_012c427cbd34484dbccfcc17859486b3",
      "value": 1
     }
    },
    "ff25d80f657c420ab67b73a09e7c93aa": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "L2: bidirectional: ",
      "description_tooltip": null,
      "layout": "IPY_MODEL_3cc198966c94474f852cb61a1b8e5b19",
      "max": 1,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_807b5764492646ea9d85d1e1fcfae51e",
      "value": 1
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}