pop-synth-vae.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import json\n",
    "import pickle\n",
    "import pandas as pd\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import random as rn\n",
    "import copy\n",
    "from pprint import pprint\n",
    "import tensorflow as tf\n",
    "import keras"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reproducible results \n",
    "# https://keras.io/getting-started/faq/#how-can-i-obtain-reproducible-results-using-keras-during-development\n",
    "\n",
    "# The below is necessary in Python 3.2.3 onwards to\n",
    "# have reproducible behavior for certain hash-based operations.\n",
    "# See these references for further details:\n",
    "# https://docs.python.org/3.4/using/cmdline.html#envvar-PYTHONHASHSEED\n",
    "# https://github.com/keras-team/keras/issues/2280#issuecomment-306959926\n",
    "\n",
    "os.environ['PYTHONHASHSEED'] = '0'\n",
    "\n",
    "# The below is necessary for starting Numpy generated random numbers\n",
    "# in a well-defined initial state.\n",
    "\n",
    "np.random.seed(42)\n",
    "\n",
    "# The below is necessary for starting core Python generated random numbers\n",
    "# in a well-defined state.\n",
    "\n",
    "rn.seed(12345)\n",
    "\n",
    "# Force TensorFlow to use single thread.\n",
    "# Multiple threads are a potential source of\n",
    "# non-reproducible results.\n",
    "# For further details, see: https://stackoverflow.com/questions/42022950/which-seeds-have-to-be-set-where-to-realize-100-reproducibility-of-training-res\n",
    "\n",
    "#session_conf = tf.ConfigProto(intra_op_parallelism_threads=1, inter_op_parallelism_threads=1)\n",
    "\n",
    "# The below tf.set_random_seed() will make random number generation\n",
    "# in the TensorFlow backend have a well-defined initial state.\n",
    "# For further details, see: https://www.tensorflow.org/api_docs/python/tf/set_random_seed\n",
    "\n",
    "tf.set_random_seed(1234)\n",
    "\n",
    "#sess = tf.Session(graph=tf.get_default_graph(), config=session_conf)\n",
    "#set_session(sess)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# prevent tensorflow from allocating the entire GPU memory at once\n",
    "from keras.backend.tensorflow_backend import set_session\n",
    "config = tf.ConfigProto()\n",
    "config.gpu_options.allow_growth = True\n",
    "config.gpu_options.visible_device_list = \"0\"\n",
    "set_session(tf.Session(config=config))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "input data format (csv):\n",
    "\n",
    "\"attribute_1_name\", \"attribute_2_name\", ..., \"attribute_n_name\"\n",
    "agent_1_attribute_1_val, agent_1_attribute_1_val, ..., agent_1_attribute_1_val\n",
    "...\n",
    "agent_N_attribute_1_val, agent_N_attribute_2_val, ..., agent_N_attribute_n_val\n",
    "\n",
    "\"\"\" \n",
    "\n",
    "data_dir = 'data/dir'\n",
    "data_filename = \"data.csv\"\n",
    "data_file = os.path.join(data_dir, data_filename)\n",
    "delimiter = ','\n",
    "df_all = pd.read_csv(data_file, delimiter=delimiter)\n",
    "print(df_all.info())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# data statistics\n",
    "df_all.describe().append(df_all.isnull().sum().rename('isnull'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# names of categorical attributes\n",
    "categorical = [\n",
    "    \"categorical_attribute_name_1\",\n",
    "    #...\n",
    "    \"categorical_attribute_name_N_c\"\n",
    "]\n",
    "\n",
    "# names of numerical attributes\n",
    "numerical = [\n",
    "    \"mumerical_attribute_name_1\",\n",
    "    #...\n",
    "    \"mumerical_attribute_name_N_n\"\n",
    "]\n",
    "\n",
    "# names of real-valued numerical attributes (a subset of numerical)\n",
    "numerical_float = [\n",
    "    \"mumerical_float_attribute_name_i\",\n",
    "    #...\n",
    "]\n",
    "\n",
    "# names of integer-valued numerical attributes (a subset of numerical)\n",
    "numerical_int = [\n",
    "    \"mumerical_int_attribute_name_i\",\n",
    "    #...\n",
    "]\n",
    "\n",
    "columns_all = categorical + numerical\n",
    "\n",
    "# attributes to model\n",
    "columns_to_use = columns_all\n",
    "\n",
    "# removing duplicates just in case\n",
    "columns_to_use = list(set(columns_to_use))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# bins used for the comparison of the joint distribution / discretization of the numerical variables\n",
    "min_max_bins = { \n",
    "    'scheme_1': {\n",
    "        # name: [minimum_value, maximum_value, number of bins]\n",
    "        'mumerical_float_attribute_name_i': [0, 99, 100], # 100 bins: [0,1), [1,2), ..., [99,>99)        \n",
    "    },\n",
    "    # ...\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# caching routine\n",
    "def get_file_suffix(train_test_ratio, columns_to_use, data_dir, data_filename_name):\n",
    "    index_file = data_filename_name + '_index' + '.txt'\n",
    "    index_file = os.path.join(data_dir, index_file)\n",
    "    is_index_file = os.path.isfile(index_file) \n",
    "    if not is_index_file:\n",
    "        if not os.path.exists(data_dir): \n",
    "            os.makedirs(data_dir)\n",
    "        with open(index_file, 'w') as f:\n",
    "            f.write('')\n",
    "    with open(index_file) as f:    \n",
    "        index = f.readlines()\n",
    "        index = [x.strip() for x in index]\n",
    "    key = str(train_test_ratio) + '---'\n",
    "    key += '---'.join(sorted(columns_to_use))\n",
    "    if key in index:\n",
    "        file_suffix = str(index.index(key))\n",
    "    else:\n",
    "        file_suffix = str(len(index))\n",
    "        index.append(key)\n",
    "        with open(index_file, 'w') as f:\n",
    "            f.write('\\n'.join(index))\n",
    "    return file_suffix\n",
    "\n",
    "\n",
    "# get record ids for the train/test split\n",
    "def get_split_ids(df, split_data_dir, train_test_ratio, file_suffix):\n",
    "    ids = df.index.tolist()\n",
    "    num_samples = len(year_ids)\n",
    "    num_samples_train = int(num_samples * train_test_ratio)\n",
    "    train_ids = year_ids[:num_samples_train]\n",
    "    test_ids = year_ids[num_samples_train:]  \n",
    "    #print(len(train_ids), len(test_ids))\n",
    "    return train_ids, test_ids\n",
    "\n",
    "\n",
    "# data cleaning and one-hot encoding of categorical variables\n",
    "def prepare_data(df, columns_to_use, categorical, numerical):\n",
    "    numerical = [x for x in numerical if x in columns_to_use]\n",
    "    categorical = [x for x in categorical if x in columns_to_use]\n",
    "    df = df[columns_to_use].copy()\n",
    "    if numerical:\n",
    "        df[numerical] = df[numerical].apply(pd.to_numeric, errors='coerce')\n",
    "        df.fillna(df.mean()[numerical], inplace=True)\n",
    "    if categorical:\n",
    "        names = list(df[categorical].select_dtypes(include=['float64', 'int64']).columns)\n",
    "        df[names] = df[names].fillna(-1).astype(int).astype(str)\n",
    "        # one-hot encoding of categorical variables\n",
    "        df = pd.get_dummies(data=df, columns=cols)\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert numerical variables to categorical\n",
    "def add_num_cat_res(df):\n",
    "    col_cats = []\n",
    "    for col_name in min_max_bins['scheme_1']:\n",
    "        if col_name not in df: continue\n",
    "        b0, b1, bn = min_max_bins['scheme_1'][col_name]\n",
    "        bin_width = (b1 - b0) / bn\n",
    "        bins = [b0 + i * bin_width for i in range(bn + 1)]\n",
    "        col_cat_name = col_name + '--cat-res'\n",
    "        df[col_cat_name] = df[col_name]\n",
    "        df[col_cat_name] = df[col_cat_name].apply(lambda x: x if x < b1 else b1 - 10e-10)\n",
    "        df[col_cat_name] = df[col_cat_name].apply(lambda x: x if x >= b0 else b0)\n",
    "        df[col_cat_name] = df[col_cat_name].apply(lambda x: b0 + bin_width * int((x - b0) / (b1 - b0) * bn))\n",
    "        df[col_cat_name] = df[col_cat_name].astype(int)\n",
    "        categorical += col_cat_name\n",
    "    df = df.drop(list(min_max_bins['scheme_1'].keys()), axis=1)\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "reshuffle = False # reshuffle previous random test/train split\n",
    "convert_numerical_to_categorical = False\n",
    "#\n",
    "data_filename_name, data_filename_ext = data_filename.split('.')\n",
    "shuffled_data_dir = os.path.join(data_dir, 'shuffled_one_hot')\n",
    "if not os.path.exists(shuffled_data_dir): \n",
    "    os.makedirs(shuffled_data_dir)\n",
    "shuffled_data_file = os.path.join(shuffled_data_dir, data_filename)   \n",
    "is_shuffled_data_file = os.path.isfile(shuffled_data_file)\n",
    "if reshuffle or not is_shuffled_data_file:\n",
    "    print(\"preparing df...\")\n",
    "    if convert_numerical_to_categorical:\n",
    "        df = add_num_cat_res(df)\n",
    "    df = prepare_data(df_all, columns_all, categorical, numerical)\n",
    "    df = df.sample(frac=1).reset_index(drop=True) # reshuffle dataframe records\n",
    "    df.to_pickle(shuffled_data_file)\n",
    "else:\n",
    "    df = pd.read_pickle(shuffled_data_file)\n",
    "    print(\"df loaded from file\")\n",
    "#\n",
    "print(shuffled_data_file)\n",
    "print(df.info())\n",
    "print(df.columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_test_ratio = 0.2 # 20% train / 80% test\n",
    "#\n",
    "split_data_dir = 'split'\n",
    "file_suffix = get_file_suffix(train_test_ratio, columns_to_use, split_data_dir, data_filename_name)\n",
    "train_ids, test_ids = get_split_ids(df, split_data_dir, train_test_ratio, file_suffix)\n",
    "# drop unsused variables in the dataframe\n",
    "numerical = [x for x in numerical if x in columns_to_use]\n",
    "categorical = [x for x in categorical if x in columns_to_use]\n",
    "categorical_one_hot = []\n",
    "for cat in categorical:\n",
    "    for col in df:\n",
    "        if cat + '_' in col:\n",
    "            categorical_one_hot.append(col)\n",
    "df = df[numerical + categorical_one_hot]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# column indices for numerical and categorical variables\n",
    "# numerical variables: always the first numerical_col_n columns\n",
    "# categorical variables: cat_groups [first_index, last_index+1]\n",
    "numerical_col, categorical_col, cat_groups = [], [], []\n",
    "col_name_prev = ''\n",
    "for i, col_name in enumerate(df.columns):\n",
    "    col_name = col_name.split(\"_\")[0]\n",
    "    if col_name in numerical:\n",
    "        numerical_col.append(i)\n",
    "    elif col_name in categorical:\n",
    "        categorical_col.append(i)\n",
    "        if col_name != col_name_prev:\n",
    "            cat_groups.append(i)\n",
    "    else:\n",
    "        raise Exception(\"unknown column type: \" + col_name)\n",
    "    col_name_prev = col_name\n",
    "numerical_col_n = len(numerical_col)\n",
    "categorical_col_n = len(categorical_col)\n",
    "cat_groups.append(i + 1)\n",
    "cat_groups_n = len(cat_groups) - 1\n",
    "#\n",
    "print(numerical_col_n, numerical_col)\n",
    "print(categorical_col_n, categorical_col)\n",
    "print(cat_groups_n, cat_groups)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "data = df.as_matrix()\n",
    "data = data.astype(np.float32, copy=False)\n",
    "features_n = len(data[0])\n",
    "#\n",
    "data_train = data[train_ids]\n",
    "data_test = data[test_ids]\n",
    "#\n",
    "print(features_n)\n",
    "print(data_train.shape)\n",
    "print(data_test.shape)\n",
    "#print(data[:2, 0:29])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('\\n--- CATEGORICAL ---')\n",
    "col_names = df.columns.tolist()\n",
    "for g_i in range(cat_groups_n):\n",
    "    g_i_beg = cat_groups[g_i]\n",
    "    g_i_end = cat_groups[g_i + 1]\n",
    "    group_i_names = col_names[g_i_beg: g_i_end]\n",
    "    col_name = '_'.join(col_names[g_i_beg].split('_')[:-1])\n",
    "    print(80*'-')\n",
    "    print(col_name, len(group_i_names))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Validation and caching helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# caching and temp directories\n",
    "model_cache_dir = 'model_cache'\n",
    "temp_dir = 'temp'\n",
    "if not os.path.exists(temp_dir): \n",
    "    os.makedirs(temp_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.ticker as ticker\n",
    "from matplotlib import rcParams\n",
    "rcParams.update({'figure.autolayout': True})\n",
    "\n",
    "    \n",
    "# check and plot marginals \n",
    "def check_marginals(data, bins=None, title=None):\n",
    "    fig, ax = plt.subplots(figsize=(3, 3), dpi=200, facecolor='w', edgecolor='k')\n",
    "    colors = ['cornflowerblue', 'seagreen', 'tomato', 'darkorchid']\n",
    "    labels = ['true', 'pred', 'bla1', 'bla2']\n",
    "    data_len = len(data)\n",
    "    colors = colors[:data_len]\n",
    "    labels = labels[:data_len]\n",
    "    weights = [np.ones_like(d)/float(len(d)) for d in data]\n",
    "    if bins is not None:\n",
    "        counts, bins, patches = ax.hist(data, histtype='bar', color=colors, \n",
    "                                        label=labels, normed=False, bins=bins, weights=weights, rwidth=None)#5\n",
    "    else:\n",
    "        counts, bins, patches = ax.hist(data, histtype='bar', color=colors, \n",
    "                                        label=labels, normed=False, weights=weights, rwidth=None)#20\n",
    "    ax.set_yscale('log')\n",
    "    ax.xaxis.set_major_formatter(ticker.FormatStrFormatter('%0.2f'))\n",
    "    plt.tick_params(\n",
    "        axis='x',          # changes apply to the x-axis\n",
    "        which='both',      # both major and minor ticks are affected\n",
    "        bottom=False,      # ticks along the bottom edge are off\n",
    "        top=False,         # ticks along the top edge are off\n",
    "        labelbottom=False) # labels along the bottom edge are off\n",
    "    if title is not None:\n",
    "        plt.title(title)\n",
    "    plt.ylim((1e-5, 1))\n",
    "    plt.tight_layout()\n",
    "    fig.savefig(os.path.join(temp_dir, '{}.pdf'.format(title)), dpi=fig.dpi)\n",
    "    plt.show()\n",
    "    corr = np.corrcoef(counts[0], counts[1])[0, 1]\n",
    "    rmse = ((counts[0] - counts[1]) ** 2).mean() ** .5\n",
    "    print(\"corr =\", corr)\n",
    "    print(\"rmse =\", rmse)\n",
    "    \n",
    "    \n",
    "# compare two vectors: Y_test (true) and Y_pred (predicted)\n",
    "def compute_stat(Y_test, Y_pred, do_plot, plot_log):\n",
    "    Y_test, Y_pred = np.array(Y_test), np.array(Y_pred)\n",
    "    corr_mat = np.corrcoef(Y_test, Y_pred)\n",
    "    corr = corr_mat[0, 1]\n",
    "    if np.isnan(corr): corr = 0.0\n",
    "    # MAE\n",
    "    mae = np.absolute(Y_test - Y_pred).mean()\n",
    "    # RMSE\n",
    "    rmse = np.linalg.norm(Y_test - Y_pred) / np.sqrt(len(Y_test))\n",
    "    # SRMSE\n",
    "    ybar = Y_test.mean()\n",
    "    srmse = rmse / ybar\n",
    "    # r-square\n",
    "    u = np.sum((Y_pred - Y_test)**2)\n",
    "    v = np.sum((Y_test - ybar)**2)\n",
    "    r2 = 1.0 - u / v\n",
    "    stat = {'mae': mae, 'rmse': rmse, 'r2': r2, 'srmse': srmse, 'corr': corr}\n",
    "    if do_plot:\n",
    "        fig = plt.figure(figsize=(3, 3), dpi=200, facecolor='w', edgecolor='k')\n",
    "        #plot\n",
    "        print('corr = %f' % (corr))\n",
    "        print('MAE = %f' % (mae))\n",
    "        print('RMSE = %f' % (rmse))\n",
    "        print('SRMSE = %f' % (srmse))\n",
    "        print('r2 = %f' % (r2))\n",
    "        min_Y = min([min(Y_test),min(Y_pred)])\n",
    "        max_Y = max([max(Y_test),max(Y_pred)])\n",
    "        w = max_Y - min_Y\n",
    "        max_Y += w * 0.05\n",
    "        text = ['SMRSE = {:.3f}'.format(stat['srmse']),\n",
    "                'Corr = {:.3f}'.format(stat['corr']),\n",
    "                '$R^2$ = {:.3f}'.format(stat['r2'])]\n",
    "        text = '\\n'.join(text)\n",
    "        plt.text(w * 0.08, w * 0.8, text)\n",
    "        plt.plot(Y_test, Y_pred, '.', alpha=0.5, ms=10, color='seagreen', markeredgewidth=0)\n",
    "        plt.plot([min_Y, max_Y], [min_Y, max_Y], ls='--', color='gray', linewidth=1.0)\n",
    "        plt.axis([min_Y, max_Y, min_Y, max_Y])\n",
    "        plt.xlabel('true')\n",
    "        plt.ylabel('predicted')\n",
    "        if plot_log:\n",
    "            eps = 1e-6\n",
    "            plt.axis([max(min_Y, eps), max_Y, max(min_Y, eps), max_Y])\n",
    "            plt.yscale('log')\n",
    "            plt.xscale('log')\n",
    "        #fig.savefig(os.path.join(temp_dir, '{}.pdf'.format(ylabel)), dpi=fig.dpi)\n",
    "        plt.show()\n",
    "    return stat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The functions below are for the comparison of joint distributions \n",
    "# and calculation of conditionals for the Gibbs sampler\n",
    "# ! Global variables, like \"min_max_bins\", are used !\n",
    "\n",
    "# get bin id for a numerical value\n",
    "def get_bin_i(min_val, max_val, val, bins_n):\n",
    "    #print(val, max_val, min_val)\n",
    "    if val >= max_val: val = max_val - 1e-10\n",
    "    if val < min_val: val = min_val\n",
    "    return int((val - min_val) / (max_val - min_val) * bins_n)\n",
    "\n",
    "\n",
    "# calculate frequency of bins\n",
    "def calc_data_freq(data, min_max_num_bin, var_is_ignore, var_g_is_ignore):\n",
    "    data_freq = {}\n",
    "    #print(var_is_ignore, var_g_is_ignore)\n",
    "    for vals in data:\n",
    "        bin_key = get_bin_key(vals, min_max_num_bin, var_is_ignore, var_g_is_ignore)\n",
    "        #print(bin_key)\n",
    "        #raise Exception(\"test\")\n",
    "        if bin_key not in data_freq:\n",
    "            data_freq[bin_key] = 0\n",
    "        data_freq[bin_key] += 1\n",
    "    for k in data_freq.keys():\n",
    "        data_freq[k] /= len(data)\n",
    "    return data_freq\n",
    "\n",
    "\n",
    "# frequency of bins + caching\n",
    "def get_data_freq(data, min_max_num_bin, var_is_ignore, var_g_is_ignore, cache_file=None):\n",
    "    print(cache_file)\n",
    "    if cache_file is not None:\n",
    "        is_cache_file = os.path.isfile(cache_file)\n",
    "        if is_cache_file:\n",
    "            with open(cache_file, 'rb') as f:\n",
    "                data_freq = pickle.load(f)\n",
    "            print(\"data_freq loaded from file:\", cache_file)\n",
    "        else:\n",
    "            print('calc_data_freq started')\n",
    "            start = time.time()\n",
    "            data_freq = calc_data_freq(data, min_max_num_bin, var_is_ignore, var_g_is_ignore)\n",
    "            duration = time.time() - start\n",
    "            print(\"Took {} seconds to calc_data_freq\".format(duration))\n",
    "            with open(cache_file, 'wb') as f:\n",
    "                pickle.dump(data_freq, f)\n",
    "            print(\"data_freq saved to file:\", cache_file)\n",
    "    else:\n",
    "        data_freq = calc_data_freq(data, min_max_num_bin, var_is_ignore, var_g_is_ignore)\n",
    "    return data_freq\n",
    "\n",
    "\n",
    "# calculate min_max_bins for an arbitrary discretization scheme\n",
    "def get_min_max_num_bin(data_test, data_pred, num_bin):\n",
    "    if num_bin in min_max_bins:\n",
    "        min_max_num_bin = min_max_bins[num_bin]\n",
    "    else:\n",
    "        col_names_num = df.columns.tolist()[:numerical_col_n]\n",
    "        # calculate min max for nmerical columns\n",
    "        min_max = {col_names_num[i]: [0, 0, num_bin] for i in range(numerical_col_n)}\n",
    "        for i in range(numerical_col_n):\n",
    "            col_test = data_test[:, i]\n",
    "            col_pred = data_pred[:, i]\n",
    "            min_max[col_names_num[i]][0] = min(np.min(col_test), np.min(col_pred))\n",
    "            min_max[col_names_num[i]][1] = max(np.max(col_test), np.max(col_pred))\n",
    "        #print(min_max)\n",
    "        eps = 1e-10\n",
    "        for i in range(numerical_col_n):\n",
    "            min_max[col_names_num[i]][0] -= eps\n",
    "            min_max[col_names_num[i]][1] += eps\n",
    "        #\n",
    "        min_max_num_bin = min_max\n",
    "    return min_max_num_bin\n",
    "\n",
    "\n",
    "# comare distributions\n",
    "def compare_joints_bins(data_test, data_pred, num_bin, \n",
    "                        data_test_freq=None, min_max_num_bin=None,\n",
    "                        var_is_ignore=None, var_g_is_ignore=None,\n",
    "                        cache_file_test=None, cache_file_pred=None):\n",
    "    if min_max_num_bin is None:\n",
    "        min_max_num_bin = get_min_max_num_bin(data_test, data_pred, num_bin)\n",
    "    #\n",
    "    if data_test_freq is None:\n",
    "        data_test_freq = get_data_freq(data_test, min_max_num_bin, \n",
    "                                        var_is_ignore, var_g_is_ignore,\n",
    "                                        cache_file_test) \n",
    "    data_pred_freq = get_data_freq(data_pred, min_max_num_bin,\n",
    "                                    var_is_ignore, var_g_is_ignore,\n",
    "                                    cache_file_pred) \n",
    "    #print(\"data_test_freq:\", data_test_freq)\n",
    "    #print(\"data_pred_freq:\", data_pred_freq)\n",
    "    #\n",
    "    print(80 * '-')\n",
    "    print(\"numerical var bins_n =\", num_bin)\n",
    "    data_test_freq_list, data_pred_freq_list = [], []\n",
    "    for k in data_test_freq.keys():\n",
    "        data_test_freq_list.append(data_test_freq[k])\n",
    "        if k in data_pred_freq:\n",
    "            data_pred_freq_list.append(data_pred_freq[k])\n",
    "        else:\n",
    "            data_pred_freq_list.append(0.0)\n",
    "    for k in data_pred_freq.keys():\n",
    "        if k not in data_test_freq:\n",
    "            data_pred_freq_list.append(data_pred_freq[k])\n",
    "            data_test_freq_list.append(0.0)\n",
    "    #\n",
    "    data_space_size = len(data_test_freq_list)\n",
    "    stat = compute_stat(data_test_freq_list, data_pred_freq_list, True, False)\n",
    "    print(\"data_space_size =\", data_space_size)\n",
    "    print(stat)\n",
    "    print(80 * '-')\n",
    "    return stat\n",
    "\n",
    "\n",
    "# basically the same as get_min_max_num_bin but the min/max values are used for one set of samples\n",
    "def calc_min_max_num_bin(data, num_bin):\n",
    "    col_names_num = df.columns.tolist()[:numerical_col_n]\n",
    "    # calculate min max for nmerical columns\n",
    "    min_max = {col_names_num[i]: [0, 0, num_bin] for i in range(numerical_col_n)}\n",
    "    for i in range(numerical_col_n):\n",
    "        vals = data[:, i]\n",
    "        min_max[col_names_num[i]][0] = np.min(vals)\n",
    "        min_max[col_names_num[i]][1] = np.max(vals)\n",
    "    #print(min_max)\n",
    "    eps = 1e-10\n",
    "    for i in range(numerical_col_n):\n",
    "        min_max[col_names_num[i]][0] -= eps\n",
    "        min_max[col_names_num[i]][1] += eps\n",
    "    return min_max\n",
    "\n",
    "\n",
    "# key of a bin in the frequecy table (aka joint distribution), like \"var_1=val_1,var_2=_val_2,...\"\n",
    "def get_bin_key(vals, min_max_num_bin, var_is_ignore, var_g_is_ignore):\n",
    "    bin_key = []\n",
    "    col_names_num = df.columns.tolist()[:numerical_col_n]\n",
    "    for i, val in enumerate(vals[: numerical_col_n]):\n",
    "        if var_is_ignore is not None and i in var_is_ignore: continue\n",
    "        bin_i = get_bin_i(min_max_num_bin[col_names_num[i]][0], \n",
    "                          min_max_num_bin[col_names_num[i]][1], \n",
    "                          val, \n",
    "                          min_max_num_bin[col_names_num[i]][2])\n",
    "        bin_key.append('num' + str(i) + \":bin\" + str(bin_i))\n",
    "    for g_i in range(cat_groups_n): \n",
    "        if var_g_is_ignore is not None and g_i in var_g_is_ignore: continue\n",
    "        g_i_beg = cat_groups[g_i]\n",
    "        g_i_end = cat_groups[g_i + 1]\n",
    "        val_cat = vals[g_i_beg: g_i_end]\n",
    "        bin_i = np.argmax(val_cat)\n",
    "        bin_key.append('cat' + str(g_i) + \":bin\" + str(bin_i))\n",
    "    bin_key = \"--\".join(bin_key)\n",
    "    return bin_key\n",
    "\n",
    "\n",
    "# conditionals of the numerical variables for the Gibbs sampler\n",
    "def calc_full_conditionals_num(args):\n",
    "    var_i, full_cond, data, min_max_num_bin = args\n",
    "    print(\"var_i={}\".format(var_i))\n",
    "    var_name = 'num' + str(var_i)\n",
    "    full_cond[var_name] = {}\n",
    "    for vals in data:\n",
    "        bin_key = get_bin_key(vals, min_max_num_bin, [var_i], None)\n",
    "        if bin_key not in full_cond[var_name]:\n",
    "            full_cond[var_name][bin_key] = []\n",
    "        full_cond[var_name][bin_key].append(vals[var_i])\n",
    "\n",
    "\n",
    "# conditionals of the categorical variables for the Gibbs sampler\n",
    "def calc_full_conditionals_cat(args):\n",
    "    var_g_i, full_cond, data, min_max_num_bin, cat_groups = args\n",
    "    print(\"var_g_i={}\".format(var_g_i))\n",
    "    var_name = 'cat' + str(var_g_i)\n",
    "    full_cond[var_name] = {}\n",
    "    for vals in data:\n",
    "        bin_key = get_bin_key(vals, min_max_num_bin, None, [var_g_i])\n",
    "        if bin_key not in full_cond[var_name]:\n",
    "            full_cond[var_name][bin_key] = []\n",
    "        var_g_i_beg = cat_groups[var_g_i]\n",
    "        var_g_i_end = cat_groups[var_g_i + 1]\n",
    "        full_cond[var_name][bin_key].append(vals[var_g_i_beg: var_g_i_end])\n",
    "        \n",
    "\n",
    "# conditionals for the Gibbs sampler\n",
    "def calc_full_conditionals(data, num_bin):\n",
    "    print(\"caclulating conditional for num_bin={}\".format(num_bin))\n",
    "    # full_cond = {var_name: {bin_key: [vals_of_var_name]}}\n",
    "    full_cond = {}\n",
    "    #\n",
    "    if num_bin in min_max_bins:\n",
    "        min_max_num_bin = min_max_bins[num_bin]\n",
    "    else:\n",
    "        min_max_num_bin = calc_min_max_num_bin(data, num_bin)\n",
    "    for var_i in range(numerical_col_n):\n",
    "        args = (var_i, full_cond, data, min_max_num_bin)\n",
    "        calc_full_conditionals_num(args)\n",
    "    for var_g_i in range(cat_groups_n):\n",
    "        args = (var_g_i, full_cond, data, min_max_num_bin, cat_groups)\n",
    "        calc_full_conditionals_cat(args)\n",
    "    data_space_size = 0\n",
    "    for var_name in full_cond:\n",
    "        data_space_size += len(full_cond[var_name])\n",
    "    print(\"data_space_size =\", data_space_size)\n",
    "    #\n",
    "    return full_cond"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compare marginal distributions of the samples\n",
    "# argument data = [samples_vae, samples_test, ...]\n",
    "\n",
    "def check_marginals_numerical(data, num_bin):\n",
    "    print('\\n--- NUMERICAL ---')\n",
    "    col_names_num = df.columns.tolist()[:numerical_col_n]\n",
    "    for col_ind, col_name in enumerate(col_names_num):\n",
    "        print(80*'-')\n",
    "        print(col_name)\n",
    "        d_cols = []\n",
    "        for d in data:\n",
    "            d_col = d[:, col_ind]\n",
    "            if col_name in numerical_int:\n",
    "                d_col = np.around(d_col)\n",
    "                print(\"//int//\")\n",
    "            d_col = np.clip(d_col, \n",
    "                            min_max_bins[num_bin][col_name][0], \n",
    "                            min_max_bins[num_bin][col_name][1])\n",
    "            d_cols.append(d_col)\n",
    "            print(d_col[:10])\n",
    "        title = col_name.split('--')[0]\n",
    "        check_marginals(d_cols, bins=min_max_bins[num_bin][col_name][2], title=title)\n",
    "        print(80*'-')\n",
    "        \n",
    "    \n",
    "def check_marginals_categorical(data):\n",
    "    print('\\n--- CATEGORICAL ---')\n",
    "    col_names = df.columns.tolist()\n",
    "    for g_i in range(cat_groups_n):\n",
    "        g_i_beg = cat_groups[g_i]\n",
    "        g_i_end = cat_groups[g_i + 1]\n",
    "        group_i_names = col_names[g_i_beg: g_i_end]\n",
    "        col_name = '_'.join(col_names[g_i_beg].split('_')[:-1])\n",
    "        print(80*'-')\n",
    "        print(col_name, len(group_i_names))\n",
    "        d_cols = []\n",
    "        for d in data:\n",
    "            d_col = d[:, g_i_beg:g_i_end]\n",
    "            d_col = np.argmax(d_col, axis=1)\n",
    "            d_cols.append(d_col)\n",
    "            print(d_col[:10])\n",
    "        title = col_name.split('--')[0]\n",
    "        check_marginals(d_cols, title=title)\n",
    "        print(80*'-')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# finding nearest samples for the diversity test\n",
    "\n",
    "from sklearn.neighbors import BallTree\n",
    "\n",
    "\n",
    "def get_k_nearest_samples_dist(data_test, data_pred, k, N, var_is_ignore, var_g_is_ignore):\n",
    "    cols_to_use = []\n",
    "    for col_i in range(numerical_col_n + categorical_col_n):\n",
    "        add_col = True\n",
    "        if var_is_ignore is not None and col_i in var_is_ignore: \n",
    "            add_col = False\n",
    "        if var_g_is_ignore is not None:\n",
    "            for g_i in var_g_is_ignore:\n",
    "                if cat_groups[g_i] <= col_i < cat_groups[g_i + 1]:\n",
    "                    add_col = False\n",
    "        if add_col:\n",
    "            cols_to_use.append(col_i)\n",
    "    #\n",
    "    tree = BallTree(data_test[:, cols_to_use], leaf_size=k)\n",
    "    dist_all = []\n",
    "    data_ind = np.arange(len(data_pred))\n",
    "    np.random.shuffle(data_ind)\n",
    "    data_ind = data_ind[:N]\n",
    "    for i in data_ind:\n",
    "        sample_pred = data_pred[i, cols_to_use]\n",
    "        dist, ind = tree.query([sample_pred], k=k)\n",
    "        dist[0] /= np.sqrt(len(cols_to_use))\n",
    "        dist_all += dist[0].tolist()\n",
    "    return np.array(dist_all), data_ind"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate data for the plots of models' performance (used in the plot_res function)\n",
    "def get_perform_plot_data(stat, bin_n_comparisons, samples_keys, sort=False):\n",
    "    errors, diversity_1, diversity_2, xticks = {}, {}, {}, {}\n",
    "    x = range(len(samples_keys))\n",
    "    for bin_n_c in bin_n_comparisons:\n",
    "        print(80 * \"-\")\n",
    "        print(\"bin_n_comparison =\", bin_n_c)\n",
    "        errors[bin_n_c] = [stat[bin_n_c][k]['joint_pdf'][error_measure] \n",
    "                           for k in samples_keys]\n",
    "        diversity_1[bin_n_c] = [stat[bin_n_c][k]['nearest_samples_dist'][diversity_measure_1]\n",
    "                                for k in samples_keys]\n",
    "        diversity_2[bin_n_c] = [stat[bin_n_c][k]['nearest_samples_dist'][diversity_measure_2]\n",
    "                                for k in samples_keys]\n",
    "        #\n",
    "        if sort:\n",
    "            errors_sorted = list(enumerate(errors[bin_n_c]))\n",
    "            errors_sorted = sorted(errors_sorted, key=lambda x: x[1], reverse=True)\n",
    "            ind = [x[0] for x in errors_sorted]\n",
    "            errors[bin_n_c] = [errors[bin_n_c][i] for i in ind]\n",
    "            diversity_1[bin_n_c] = [diversity_1[bin_n_c][i] for i in ind]\n",
    "            diversity_2[bin_n_c] = [diversity_2[bin_n_c][i] for i in ind]\n",
    "            xticks[bin_n_c] = [samples_keys[i] for i in ind]\n",
    "        else:\n",
    "            xticks[bin_n_c] = samples_keys\n",
    "        #\n",
    "        plot_res(x, errors[bin_n_c], xticks=xticks[bin_n_c],\n",
    "                 **error_plot_params[error_measure])\n",
    "        plot_res(x, diversity_1[bin_n_c], xticks=xticks[bin_n_c], \n",
    "                 **diversity_plot_params[diversity_measure_1])\n",
    "        plot_res(x, diversity_2[bin_n_c], xticks=xticks[bin_n_c], \n",
    "                 **diversity_plot_params[diversity_measure_2])\n",
    "        print(80 * \"-\")\n",
    "    return errors, diversity_1, diversity_2, xticks\n",
    "\n",
    "\n",
    "# plot performance of the models\n",
    "def plot_res(x, y, ylim=None, color=None, ylabel=None, xticks=None, \n",
    "             h_lines=None, h_labels=None, title=None):\n",
    "    figsize = (18, 9)\n",
    "    #linestyles = ['--', '-.', ':', '-']\n",
    "    linestyles = [\n",
    "        (0, ()),\n",
    "        (0, (1, 10)),\n",
    "        (0, (1, 1)),\n",
    "        (0, (5, 5)),\n",
    "        (0, (3, 5, 1, 5)),\n",
    "        (0, (3, 5, 1, 5, 1, 5)),\n",
    "    ]\n",
    "    \"\"\"\n",
    "    linestyles = OrderedDict(\n",
    "    [('solid',               (0, ())),\n",
    "     ('loosely dotted',      (0, (1, 10))),\n",
    "     ('dotted',              (0, (1, 5))),\n",
    "     ('densely dotted',      (0, (1, 1))),\n",
    "\n",
    "     ('loosely dashed',      (0, (5, 10))),\n",
    "     ('dashed',              (0, (5, 5))),\n",
    "     ('densely dashed',      (0, (5, 1))),\n",
    "\n",
    "     ('loosely dashdotted',  (0, (3, 10, 1, 10))),\n",
    "     ('dashdotted',          (0, (3, 5, 1, 5))),\n",
    "     ('densely dashdotted',  (0, (3, 1, 1, 1))),\n",
    "\n",
    "     ('loosely dashdotdotted', (0, (3, 10, 1, 10, 1, 10))),\n",
    "     ('dashdotdotted',         (0, (3, 5, 1, 5, 1, 5))),\n",
    "     ('densely dashdotdotted', (0, (3, 1, 1, 1, 1, 1)))])\n",
    "    \"\"\"\n",
    "    fig, ax = plt.subplots(figsize=figsize)\n",
    "    if xticks is not None:\n",
    "        plt.xticks(x, xticks, rotation='vertical')\n",
    "    ax.set_ylabel(ylabel, color=color)\n",
    "    ax.plot(x, y, color=color)\n",
    "    ax.plot(x, y, 'o', color=color)\n",
    "    #\n",
    "    text = []\n",
    "    text.append('model min = {}'.format(min(y)))\n",
    "    #\n",
    "    if h_lines is not None:\n",
    "        for i, line_y in enumerate(h_lines):\n",
    "            if ylim[0] <= line_y <= ylim[1]:\n",
    "                ax.axhline(line_y, ls=linestyles[i], color=color)\n",
    "                if h_labels is not None:\n",
    "                    ax.text(0.0, line_y + (ylim[1] - ylim[0]) / 250, h_labels[i], color=color)\n",
    "            text.append('{} = {:.3f}'.format(h_labels[i], line_y))\n",
    "    #\n",
    "    text = '\\n'.join(text)\n",
    "    ax.text(0.0, ylim[1], text, color=color)\n",
    "    print(text)\n",
    "    #\n",
    "    ax.tick_params(axis='y', labelcolor=color)\n",
    "    ax.set_ylim(ylim)\n",
    "    ax.set_title(title)\n",
    "    #\n",
    "    #fig.savefig('{}.pdf'.format(ylabel), dpi=fig.dpi)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate performance of the models + caching\n",
    "def get_stat(samples_true, samples_pred, samples_train, bin_n_comparisons, \n",
    "             samples_true_cache_dir, samples_pred_cache_dir,\n",
    "             var_is_ignore=None, var_g_is_ignore=None):\n",
    "    tree_k = 1\n",
    "    stat = {}\n",
    "    N = len(bin_n_comparisons) * len(samples_pred)\n",
    "    i = 0\n",
    "    if not os.path.exists(samples_true_cache_dir): \n",
    "        os.makedirs(samples_true_cache_dir)\n",
    "    cache_filename = 'freq_cache_file'\n",
    "    cache_filename += '--' + 'var_is_ignore=' + str(var_is_ignore)\n",
    "    cache_filename += '--' + 'var_g_is_ignore=' + str(var_g_is_ignore)\n",
    "    if len(cache_filename) > 255:\n",
    "        cache_filename = cache_filename.replace(' ', '')\n",
    "    #\n",
    "    for bin_n_c in bin_n_comparisons:\n",
    "        #\n",
    "        data_test_freq_cache_file = os.path.join(samples_true_cache_dir, cache_filename \n",
    "                                                 + '--' + 'bin_n_c=' + str(bin_n_c) \n",
    "                                                 + '--' + file_suffix + '.pickle')\n",
    "        print(data_test_freq_cache_file)\n",
    "        #\n",
    "        min_max_num_bin = get_min_max_num_bin(data_test, None, bin_n_c)\n",
    "        data_test_freq = get_data_freq(data_test, min_max_num_bin, \n",
    "                                       var_is_ignore, var_g_is_ignore,\n",
    "                                       data_test_freq_cache_file)\n",
    "        stat[bin_n_c] = {}\n",
    "        for k in samples_pred:\n",
    "            print(80 * '+')\n",
    "            print(\"bin_n_c =\", bin_n_c)\n",
    "            print(\"samples:\", k)\n",
    "            i += 1\n",
    "            print(\"Calculating {} out of {}\".format(i, N))\n",
    "            #\n",
    "            if np.isnan(samples_pred[k]).any():\n",
    "                print(\"!!! nan in samples_pred !!!\")\n",
    "                samples_pred[k] = np.nan_to_num(samples_pred[k])\n",
    "            #\n",
    "            if '.' in k:\n",
    "                samples_pred_cache_dir_ = os.path.join(samples_pred_cache_dir, \n",
    "                                                       \",\".join(k.split('.')[:-1]))\n",
    "            else:\n",
    "                samples_pred_cache_dir_ = os.path.join(samples_pred_cache_dir, k)\n",
    "            if not os.path.exists(samples_pred_cache_dir_): \n",
    "                os.makedirs(samples_pred_cache_dir_)\n",
    "            #print(k, samples_pred_cache_dir_)\n",
    "            data_pred_freq_cache_file = os.path.join(samples_pred_cache_dir_, cache_filename \n",
    "                                                     + '--' + 'bin_n_c=' + str(bin_n_c) \n",
    "                                                     + '--' + file_suffix + '.pickle')\n",
    "            #\n",
    "            stat[bin_n_c][k] = {}\n",
    "            start = time.time()\n",
    "            stat[bin_n_c][k]['joint_pdf'] = compare_joints_bins(samples_true, samples_pred[k], bin_n_c\n",
    "                                                   , data_test_freq, min_max_num_bin,\n",
    "                                                    var_is_ignore, var_g_is_ignore,\n",
    "                                                    data_test_freq_cache_file, data_pred_freq_cache_file\n",
    "                                                  )\n",
    "            duration = time.time() - start\n",
    "            print(\"Took {} seconds to compare_joints_bins\".format(duration))\n",
    "            nearest_samples_dist_filename = 'nearest_samples_dist--'\\\n",
    "                                                     + 'samples_train--'\\\n",
    "                                                     + cache_filename\\\n",
    "                                                     + '--' + file_suffix + '.pickle'\n",
    "            nearest_samples_ind_filename = 'nearest_samples_ind--'\\\n",
    "                                                     + 'samples_train--'\\\n",
    "                                                     + cache_filename\\\n",
    "                                                     + '--' + file_suffix + '.pickle'\n",
    "            if len(nearest_samples_dist_filename) > 255:\n",
    "                nearest_samples_dist_filename = nearest_samples_dist_filename.replace(' ', '')\n",
    "            if len(nearest_samples_ind_filename) > 255:\n",
    "                nearest_samples_ind_filename = nearest_samples_ind_filename.replace(' ', '')\n",
    "            nearest_samples_dist_file = os.path.join(samples_pred_cache_dir_, \n",
    "                                                     nearest_samples_dist_filename)\n",
    "            nearest_samples_ind_file = os.path.join(samples_pred_cache_dir_, \n",
    "                                                     nearest_samples_ind_filename)\n",
    "            is_nearest_samples_dist_file = os.path.isfile(nearest_samples_dist_file)\n",
    "            is_nearest_samples_ind_file = os.path.isfile(nearest_samples_ind_file)\n",
    "            if is_nearest_samples_dist_file and is_nearest_samples_ind_file:\n",
    "                with open(nearest_samples_dist_file, 'rb') as f:\n",
    "                    nearest_samples_dist = pickle.load(f)\n",
    "                print(\"nearest_samples_dist loaded from file:\", nearest_samples_dist_file)\n",
    "                with open(nearest_samples_ind_file, 'rb') as f:\n",
    "                    nearest_samples_ind = pickle.load(f)\n",
    "                print(\"nearest_samples_ind loaded from file:\", nearest_samples_ind_file)\n",
    "            else:\n",
    "                start = time.time()\n",
    "                nearest_samples_dist, nearest_samples_ind = get_k_nearest_samples_dist(samples_train, \n",
    "                        samples_pred[k], tree_k, len(samples_train), var_is_ignore, var_g_is_ignore)\n",
    "                duration = time.time() - start\n",
    "                print(\"Took {} seconds to get_k_nearest_samples_dist\".format(duration))\n",
    "                with open(nearest_samples_dist_file, 'wb') as f:\n",
    "                    pickle.dump(nearest_samples_dist, f)\n",
    "                print(\"nearest_samples_dist saved to file:\", nearest_samples_dist_file)\n",
    "                with open(nearest_samples_ind_file, 'wb') as f:\n",
    "                    pickle.dump(nearest_samples_ind, f)\n",
    "                print(\"nearest_samples_ind saved to file:\", nearest_samples_ind_file)\n",
    "            \n",
    "            # calc errors / diversity\n",
    "            stat[bin_n_c][k]['nearest_samples_dist'] = {}\n",
    "            stat[bin_n_c][k]['nearest_samples_dist']['mean'] = nearest_samples_dist.mean()\n",
    "            stat[bin_n_c][k]['nearest_samples_dist']['stddev'] = nearest_samples_dist.std()\n",
    "            stat[bin_n_c][k]['nearest_samples_dist']['frac0'] = \\\n",
    "                np.count_nonzero(nearest_samples_dist == 0) / len(nearest_samples_dist)\n",
    "            #\n",
    "            stat[bin_n_c][k]['nearest_samples_dist']['dist'] = nearest_samples_dist\n",
    "            stat[bin_n_c][k]['nearest_samples_dist']['ind'] = nearest_samples_ind\n",
    "            #\n",
    "            print('nearest_samples_dist =', stat[bin_n_c][k]['nearest_samples_dist'])\n",
    "            print(80 * '+')\n",
    "    return stat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate attribute column ids to ignore during comparison of the models\n",
    "# e.g. ignore all attributes except for the Basic\n",
    "def get_columns_to_ignore(selected_columns):\n",
    "    col_names = df.columns.tolist()\n",
    "    #print(len(col_names), col_names)\n",
    "    #print(cat_groups_n, cat_groups)\n",
    "    var_is_ignore, var_g_is_ignore = [], []\n",
    "    for col_i, col_name in enumerate(col_names):\n",
    "        ignore_column = True\n",
    "        for col_name_selected in selected_columns:\n",
    "            if col_name_selected == col_name.split(\"_\")[0]:\n",
    "                ignore_column = False\n",
    "                break\n",
    "        if ignore_column:\n",
    "            if col_i in numerical_col:\n",
    "                var_is_ignore.append(col_i)\n",
    "            else:\n",
    "                var_g_i = -1\n",
    "                for g_i, upper_limit in enumerate(cat_groups):\n",
    "                    #print(col_i, upper_limit)\n",
    "                    if col_i < upper_limit:\n",
    "                        var_g_i = g_i - 1\n",
    "                        break\n",
    "                var_g_is_ignore.append(var_g_i)\n",
    "    var_g_is_ignore = list(set(var_g_is_ignore))\n",
    "    return var_is_ignore, var_g_is_ignore"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Performance measure parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compare numercal distributions using 1) scheme_1 in min_max_bins, 2) 10 uniform bins\n",
    "bin_n_comparisons = ['scheme_1', 10]\n",
    "error_measure = 'srmse'\n",
    "diversity_measure_1 = 'mean'\n",
    "diversity_measure_2 = 'stddev'\n",
    "samples_true_cache_dir = os.path.join(model_cache_dir, 'data_test')\n",
    "samples_train_cache_dir = os.path.join(model_cache_dir, 'data_train')\n",
    "#\n",
    "error_plot_params = {\n",
    "    'srmse': {\n",
    "        'ylabel': 'srmse',\n",
    "        'ylim': (0, 5),\n",
    "        #'color': 'tab:red',\n",
    "        'color': 'k',\n",
    "        'title': 'real vs synthetic samples: joint pdf comparison'\n",
    "    },\n",
    "    'r2': {\n",
    "        'ylabel': 'r2',\n",
    "        'ylim': (-5, 1.1),\n",
    "        #'color': 'tab:red',\n",
    "        'color': 'k',\n",
    "        'title': 'real vs synthetic samples: joint pdf comparison'\n",
    "    },\n",
    "    'corr': {\n",
    "        'ylabel': 'corr',\n",
    "        'ylim': (0, 1.1),\n",
    "        #'color': 'tab:red',\n",
    "        'color': 'k',\n",
    "        'title': 'real vs synthetic samples: joint pdf comparison'\n",
    "    }\n",
    "}\n",
    "diversity_plot_params = {\n",
    "    'mean': {\n",
    "        'ylabel': 'mean',\n",
    "        'ylim': (0, 0.15),\n",
    "        #'color': 'tab:blue',\n",
    "        'color': 'k',\n",
    "        'title': 'distance to the closest sample in the training data'\n",
    "    },\n",
    "    'stddev': {\n",
    "        'ylabel': 'stddev',\n",
    "        'ylim': (0, 0.2),\n",
    "        #'color': 'tab:green',\n",
    "        'color': 'k',\n",
    "        'title': 'distance to the closest sample in the training data'\n",
    "    },\n",
    "    'frac0': {\n",
    "        'ylabel': 'fraction of 0',\n",
    "        'ylim': (0, 1.1),\n",
    "        #'color': 'tab:blue',\n",
    "        'color': 'k',\n",
    "        'title': 'distance to the closest sample in the training data'\n",
    "    }\n",
    "}\n",
    "# compare subset of features for high-dimensional cases. The rest will be ignored\n",
    "selected_columns = [\n",
    "    \"attribute_1_name\",\n",
    "    # ...\n",
    "]\n",
    "#\n",
    "selected_columns = list(set(selected_columns))\n",
    "var_is_ignore, var_g_is_ignore = get_columns_to_ignore(selected_columns)\n",
    "if not var_is_ignore:\n",
    "    var_is_ignore = None\n",
    "else:\n",
    "    print(len(var_is_ignore), var_is_ignore)\n",
    "if not var_g_is_ignore:\n",
    "    var_g_is_ignore = None\n",
    "else:\n",
    "    print(len(var_g_is_ignore), var_g_is_ignore)\n",
    "print(selected_columns)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train vs test: performance of the micro sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "samples_train_keys =[\"data_train\"]\n",
    "samples_pred_cache_dir = os.path.join(model_cache_dir, 'data_train')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# performance for all atributes\n",
    "stat_train = get_stat(data_test, {samples_train_keys[0]: data_train},\n",
    "                      data_test, bin_n_comparisons,\n",
    "                      samples_true_cache_dir, samples_pred_cache_dir,\n",
    "                     )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# performance for all atributes\n",
    "(errors_train, \n",
    " diversity_1_train, \n",
    " diversity_2_train,\n",
    " xticks_train) = get_perform_plot_data(stat_train, bin_n_comparisons, samples_train_keys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# performance for the selected atributes\n",
    "stat_train_selected = get_stat(data_test, {samples_train_keys[0]: data_train}, \n",
    "                               data_test, bin_n_comparisons, \n",
    "                               samples_true_cache_dir, samples_pred_cache_dir,\n",
    "                               var_is_ignore, var_g_is_ignore)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# performance for the selected atributes\n",
    "(errors_train_selected, \n",
    " diversity_1_train_selected, \n",
    " diversity_2_train_selected,\n",
    " xticks_train_selected) = get_perform_plot_data(stat_train_selected, \n",
    "                                            bin_n_comparisons, \n",
    "                                            samples_train_keys)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# VAE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.stats import norm\n",
    "from keras.layers import Input, Dense, Lambda, Layer, Reshape, BatchNormalization, Flatten, Concatenate\n",
    "from keras.models import Model\n",
    "from keras import backend as K\n",
    "from keras import metrics, losses\n",
    "from keras import optimizers\n",
    "from sklearn.preprocessing import StandardScaler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# updatable plot for the training\n",
    "\n",
    "from IPython.display import clear_output\n",
    "\n",
    "class PlotLosses(keras.callbacks.Callback):\n",
    "    def on_train_begin(self, logs={}):\n",
    "        self.i = 0\n",
    "        self.x = []\n",
    "        self.losses = []\n",
    "        self.val_losses = []        \n",
    "        self.fig = plt.figure()        \n",
    "        self.logs = []\n",
    "\n",
    "    def on_epoch_end(self, epoch, logs={}):\n",
    "        m1, m2 = 'loss', 'val_loss'\n",
    "        self.logs.append(logs)\n",
    "        self.x.append(self.i)\n",
    "        self.losses.append(logs.get(m1))\n",
    "        self.val_losses.append(logs.get(m2))\n",
    "        self.i += 1\n",
    "        clear_output(wait=True)\n",
    "        plt.plot(self.x, self.losses, label=m1)\n",
    "        plt.plot(self.x, self.val_losses, label=m2)\n",
    "        plt.legend()\n",
    "        plt.show();\n",
    "        \n",
    "plot_learning = PlotLosses()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparameters grid\n",
    "activation_grid = ['tanh'] # ANN activation\n",
    "cat_loss_weight_grid = [1.0] # weight of the categorical loss term\n",
    "beta_grid = [0.01, 0.05, 0.1, 0.5, 1.0, 10.0, 100.0] # weight of the KL loss term\n",
    "latent_dim_grid = [5, 10, 25] # latent dimensionality\n",
    "hidden_layers_shape_grid = [\n",
    "    [100, 50, 25], [100, 50], [100], [50, 25], [50], [25]\n",
    "] # number of neurons in hidden layers [neurons_number_hidden_layer_1, neurons_number_hidden_layer_2, ...]\n",
    "if numerical_col_n == 0 or categorical_col_n == 0:\n",
    "    cat_loss_weight_grid = [1.0]\n",
    "# number of epochs and mini-batch size\n",
    "epochs = 100\n",
    "batch_size = 64\n",
    "# number of samples from VAE\n",
    "n_samples = 100000\n",
    "# train / development set (to tune hyperparameters) ratio\n",
    "train_dev_ratio = 0.8\n",
    "samples_dev_cache_dir = os.path.join(model_cache_dir, 'data_dev_{}-{}'.format(train_dev_ratio, file_suffix))\n",
    "#\n",
    "pprint(cat_loss_weight_grid)\n",
    "pprint(beta_grid)\n",
    "pprint(latent_dim_grid)\n",
    "pprint(hidden_layers_shape_grid)\n",
    "pprint(activation_grid)\n",
    "print(len(cat_loss_weight_grid) * len(beta_grid) * len(latent_dim_grid) \n",
    "      * len(hidden_layers_shape_grid) * len(activation_grid))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train VAE\n",
    "# ! a lot of global variables, like data_train, numerical_col_n, etc, are used !\n",
    "# ...like probably everywhere in this code :(\n",
    "def vae_train(cat_loss_weight, beta, activation, latent_dim, hidden_layers_shape):\n",
    "    # split data\n",
    "    train_dev_ind = int(len(data_train) * train_dev_ratio)\n",
    "    data_train_01 = copy.deepcopy(data_train[:train_dev_ind])\n",
    "    data_test_01 = copy.deepcopy(data_train[train_dev_ind:])\n",
    "    # scale numerical variables to N(0, 1)\n",
    "    if numerical_col_n > 0:\n",
    "        scaler = StandardScaler()\n",
    "        scaler = scaler.fit(data_train[:, :numerical_col_n])\n",
    "        data_train_01[:, :numerical_col_n] = scaler.transform(\n",
    "            data_train_01[:, :numerical_col_n])\n",
    "        data_test_01[:, :numerical_col_n] = scaler.transform(\n",
    "            data_test_01[:, :numerical_col_n])\n",
    "    else:\n",
    "        scaler = None\n",
    "    #\n",
    "    original_dim = features_n\n",
    "    # specify optimizer\n",
    "    optimizer = optimizers.RMSprop(lr=0.001, rho=0.9, epsilon=None, decay=0.0)\n",
    "    # encoder\n",
    "    x = Input(shape=(original_dim,))\n",
    "    h = x\n",
    "    for hl_shape in hidden_layers_shape:\n",
    "        h = Dense(hl_shape, activation=activation)(h)\n",
    "    z_mean = Dense(latent_dim)(h)\n",
    "    z_log_var = Dense(latent_dim)(h)\n",
    "    # reparameterization\n",
    "    def sampling(args):\n",
    "        z_mean, z_log_var = args\n",
    "        epsilon = K.random_normal(shape=(K.shape(z_mean)[0], latent_dim), mean=0.0, stddev=1.0)\n",
    "        return z_mean + K.exp(z_log_var / 2) * epsilon\n",
    "    z = Lambda(sampling)([z_mean, z_log_var])\n",
    "    # decoder\n",
    "    decoder_h = [Dense(hl_shape, activation=activation) for hl_shape in reversed(hidden_layers_shape)]\n",
    "    if numerical_col_n > 0:\n",
    "        x_decoder_mean_num = Dense(numerical_col_n)\n",
    "    if categorical_col_n > 0:\n",
    "        x_decoder_mean_cat = [Dense(cat_groups[g_i + 1] - cat_groups[g_i], activation='softmax') \n",
    "                              for g_i in range(cat_groups_n)]\n",
    "    h_decoded = decoder_h[0](z)\n",
    "    for decoder_h_ in decoder_h[1:]:\n",
    "        h_decoded = decoder_h_(h_decoded)\n",
    "    if numerical_col_n > 0:\n",
    "        x_decoded_mean_num = x_decoder_mean_num(h_decoded)\n",
    "    if categorical_col_n > 0:\n",
    "        x_decoded_mean_cat = [x_decoder_mean_cat_(h_decoded)\n",
    "                              for x_decoder_mean_cat_ in x_decoder_mean_cat]\n",
    "    if numerical_col_n > 0 and categorical_col_n > 0:\n",
    "        x_decoded_mean = Concatenate()([x_decoded_mean_num] + x_decoded_mean_cat)\n",
    "    elif numerical_col_n > 0:\n",
    "        x_decoded_mean = x_decoded_mean_num\n",
    "    elif categorical_col_n > 0:\n",
    "        x_decoded_mean = Concatenate()(x_decoded_mean_cat)\n",
    "    else:\n",
    "        raise Exception(\"NO FEATURES\")\n",
    "    # custom losses\n",
    "    # total loss (used as loss in model.compile below)\n",
    "    def vae_loss(y_true, y_pred):\n",
    "        # numerical loss\n",
    "        recon_num = 0\n",
    "        if numerical_col_n > 0:\n",
    "            recon_num = K.sum(K.square(y_pred[:, :numerical_col_n] - y_true[:, :numerical_col_n]), axis=-1)\n",
    "        # categorical loss\n",
    "        recon_cat = 0\n",
    "        for g_i in range(cat_groups_n):    \n",
    "            g_i_beg = cat_groups[g_i]\n",
    "            g_i_end = cat_groups[g_i + 1]\n",
    "            recon_cat += metrics.categorical_crossentropy(y_true[:, g_i_beg:g_i_end], y_pred[:, g_i_beg:g_i_end])\n",
    "        # KL loss\n",
    "        kl = - 0.5 * K.sum(1. + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)\n",
    "        # total loss\n",
    "        loss = K.mean(recon_num + recon_cat * cat_loss_weight + beta * kl)\n",
    "        return loss\n",
    "    # KL loss (used as metrics in model.compile below)\n",
    "    def KL_loss(y_true, y_pred):\n",
    "        kl = - 0.5 * K.sum(1. + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)\n",
    "        return K.mean(beta * kl)\n",
    "    # numerical loss (used as metrics in model.compile below)\n",
    "    def recon_loss_num(y_true, y_pred):\n",
    "        recon_num = 0\n",
    "        if numerical_col_n > 0:\n",
    "            recon_num = K.sum(K.square(y_pred[:, :numerical_col_n] - y_true[:, :numerical_col_n]), axis=-1)\n",
    "        return K.mean(recon_num)\n",
    "    # categorical loss (used as metrics in model.compile below)\n",
    "    def recon_loss_cat(y_true, y_pred):\n",
    "        recon_cat = 0\n",
    "        for g_i in range(cat_groups_n):    \n",
    "            g_i_beg = cat_groups[g_i]\n",
    "            g_i_end = cat_groups[g_i + 1]\n",
    "            recon_cat += metrics.categorical_crossentropy(y_true[:, g_i_beg:g_i_end], y_pred[:, g_i_beg:g_i_end])\n",
    "        return K.mean(recon_cat * cat_loss_weight)\n",
    "    # build a VAE model (encoder+decoder)\n",
    "    model = Model(x, x_decoded_mean)\n",
    "    # Run training\n",
    "    if numerical_col_n > 0 and categorical_col_n > 0:\n",
    "        model.compile(optimizer=optimizer, loss=vae_loss, metrics=[KL_loss, recon_loss_num, recon_loss_cat])\n",
    "    elif numerical_col_n > 0:\n",
    "        model.compile(optimizer=optimizer, loss=vae_loss, metrics=[KL_loss, recon_loss_num])\n",
    "    elif categorical_col_n > 0:\n",
    "        model.compile(optimizer=optimizer, loss=vae_loss, metrics=[KL_loss, recon_loss_cat])\n",
    "    else:\n",
    "        raise Exception(\"NO FEATURES\")\n",
    "    model.fit(data_train_01, data_train_01,\n",
    "              shuffle=True,\n",
    "              batch_size=batch_size,\n",
    "              epochs=epochs,\n",
    "              validation_data=(data_test_01, data_test_01),\n",
    "              callbacks=[plot_learning],\n",
    "             )\n",
    "    # build a model to project inputs on the latent space (encoder only)\n",
    "    encoder = Model(x, z_mean)\n",
    "    # display a 2D plot to check the latent space embedding of the training data\n",
    "    x_train_encoded = encoder.predict(data_train_01, batch_size=batch_size)\n",
    "    # dummy labels are used\n",
    "    labels = [sum(np.count_nonzero(a) for a in x) for x in data_train_01]\n",
    "    plt.figure(figsize=(6, 6))\n",
    "    plt.scatter(x_train_encoded[:, 0], x_train_encoded[:, 1], c=labels, alpha=0.05)\n",
    "    plt.colorbar()\n",
    "    plt.show()\n",
    "    # define a sampling model (decoder only)\n",
    "    z_sample = Input(shape=(latent_dim,))\n",
    "    h_decoded_sample = decoder_h[0](z_sample)\n",
    "    for decoder_h_ in decoder_h[1:]:\n",
    "        h_decoded_sample = decoder_h_(h_decoded_sample)\n",
    "    if numerical_col_n > 0:\n",
    "        x_decoded_sample_num = x_decoder_mean_num(h_decoded_sample)\n",
    "    if categorical_col_n > 0:\n",
    "        x_decoded_sample_cat = [x_decoder_mean_cat_(h_decoded_sample)\n",
    "                                for x_decoder_mean_cat_ in x_decoder_mean_cat]\n",
    "    if numerical_col_n > 0 and categorical_col_n > 0:\n",
    "        x_decoded_sample = Concatenate()([x_decoded_sample_num] + x_decoded_sample_cat)\n",
    "    elif numerical_col_n > 0:\n",
    "        x_decoded_sample = x_decoded_sample_num\n",
    "    elif categorical_col_n > 0:\n",
    "        x_decoded_sample = Concatenate()(x_decoded_sample_cat)\n",
    "    else:\n",
    "        raise Exception(\"NO FEATURES\")\n",
    "    # biuld the sampling model\n",
    "    model_sample = Model(z_sample, x_decoded_sample)\n",
    "    # the full VAE model is not saved\n",
    "    return model_sample, scaler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# get samples from VAE\n",
    "def vae_get_samples(model_sample, n_samples, scaler):\n",
    "    z_sample = np.random.normal(0., 1.0, size=(n_samples, latent_dim))\n",
    "    samples = model_sample.predict(z_sample)\n",
    "    # scale back\n",
    "    if numerical_col_n > 0:\n",
    "        samples[:, :numerical_col_n] = scaler.inverse_transform(samples[:, :numerical_col_n])\n",
    "    # back from categorical softmax to one-hot\n",
    "    for g_i in range(cat_groups_n):\n",
    "        g_i_beg = cat_groups[g_i]\n",
    "        g_i_end = cat_groups[g_i + 1]\n",
    "        data_pred_col = samples[:, g_i_beg:g_i_end]\n",
    "        data_pred_col = np.argmax(data_pred_col, axis=1)\n",
    "        for row_i, row in enumerate(samples):\n",
    "            for col_i in range(g_i_beg, g_i_end):\n",
    "                if col_i - g_i_beg != data_pred_col[row_i]:\n",
    "                    samples[row_i, col_i] = 0\n",
    "                else:\n",
    "                    samples[row_i, col_i] = 1\n",
    "    # deal with the integer variables and clip using min/max values\n",
    "    if numerical_col_n > 0:\n",
    "        min_max_scheme = 'scheme_1'\n",
    "        col_names_num = df.columns.tolist()[:numerical_col_n]\n",
    "        for col_ind, col_name in enumerate(col_names_num):\n",
    "            if col_name in numerical_int:\n",
    "                samples[:, col_ind] = np.around(samples[:, col_ind])\n",
    "            # \n",
    "            samples[:, col_ind] = np.clip(samples[:, col_ind], \n",
    "                                    min_max_bins[min_max_scheme][col_name][0], \n",
    "                                    min_max_bins[min_max_scheme][col_name][1])\n",
    "    return samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# training VAE for the hyperparameters grid + caching\n",
    "train_models = True # if False then just load all pretrained models\n",
    "vae_model_cache_dirname = 'vae--epochs={}--batch_size={}'.format(epochs, batch_size) # cache folder\n",
    "#\n",
    "vae_model_cache_dir = os.path.join(model_cache_dir, vae_model_cache_dirname)\n",
    "if not os.path.exists(vae_model_cache_dir): \n",
    "    os.makedirs(vae_model_cache_dir)\n",
    "if train_models: # get combinations of hyperparameters\n",
    "    model_cache_files_vae, model_parameters_vae = [], []\n",
    "    for cat_loss_weight in cat_loss_weight_grid:\n",
    "        for beta in beta_grid:\n",
    "            for activation in activation_grid:\n",
    "                for latent_dim in latent_dim_grid:\n",
    "                    for hidden_layers_shape in hidden_layers_shape_grid:\n",
    "                        model_cache_filename = 'model'\n",
    "                        if beta != 1.0:\n",
    "                            model_cache_filename += '--' + 'beta=' + str(beta).replace('.', ',') \n",
    "                        if cat_loss_weight != 1.0:\n",
    "                            model_cache_filename += '--' + 'clw=' + str(cat_loss_weight).replace('.', ',') \n",
    "                        model_cache_filename += '--' + 'activation=' + str(activation) \n",
    "                        model_cache_filename += '--' + 'latent_dim=' + str(latent_dim)\n",
    "                        model_cache_filename += '--' + 'hidden_layers_shape=' + str(hidden_layers_shape) \n",
    "                        model_cache_filename += '--' + file_suffix + '.h5'\n",
    "                        model_cache_files_vae.append(model_cache_filename)\n",
    "                        model_parameters_vae.append((cat_loss_weight, beta, \n",
    "                                                     activation, latent_dim, hidden_layers_shape))\n",
    "else: # just load all pretrained models\n",
    "    import glob, os\n",
    "    cwd = os.getcwd()\n",
    "    filename_patten = \"model*--\" + file_suffix + \".h5\"\n",
    "    model_cache_files_vae = list(glob.glob(os.path.join(cwd, vae_model_cache_dir, filename_patten)))\n",
    "    model_cache_files_vae = [os.path.split(x)[1] for x in model_cache_files_vae]\n",
    "# number of models\n",
    "n_calc = len(model_cache_files_vae)\n",
    "print(\"len(model_cache_files_vae) =\", len(model_cache_files_vae))\n",
    "# iterate through the models\n",
    "samples_vae = {}\n",
    "for i, model_cache_filename in enumerate(model_cache_files_vae):\n",
    "    # get a model\n",
    "    scaler_cache_filename = 'scaler_' + model_cache_filename\n",
    "    model_cache_file = os.path.join(vae_model_cache_dir, model_cache_filename)\n",
    "    scaler_cache_file = os.path.join(vae_model_cache_dir, scaler_cache_filename)\n",
    "    print(\"model_cache_filename:\", model_cache_filename)\n",
    "    if train_models:\n",
    "        is_model_cache_file = os.path.isfile(model_cache_file)            \n",
    "        # It is not recommended to use pickle or cPickle to save a Keras model\n",
    "        if is_model_cache_file:\n",
    "            model_sample = keras.models.load_model(model_cache_file)\n",
    "            with open(scaler_cache_file, 'rb') as f:\n",
    "                scaler = pickle.load(f)\n",
    "            print('model loaded from file: ' + model_cache_file)\n",
    "        else:\n",
    "            print('training model')\n",
    "            start = time.time()\n",
    "            model_sample, scaler = vae_train(*model_parameters_vae[i])\n",
    "            duration = time.time() - start\n",
    "            print(model_cache_filename)\n",
    "            print(\"Took {} seconds to train vae\".format(duration))\n",
    "            print(\"Trained model {} out of {}\".format(i + 1, n_calc))\n",
    "            model_sample.save(model_cache_file)\n",
    "            with open(scaler_cache_file, 'wb') as f:\n",
    "                pickle.dump(scaler, f)\n",
    "            print('model saved to file: ' + model_cache_file)\n",
    "    # generate samples\n",
    "    samples_cache_filename = 'samples'\n",
    "    samples_cache_filename += '--' + model_cache_filename\n",
    "    samples_cache_filename += '--n_samples=' + str(n_samples)\n",
    "    samples_cache_filename += '--' + file_suffix + '.pickle'\n",
    "    samples_cache_file = os.path.join(vae_model_cache_dir, samples_cache_filename)\n",
    "    print(\"samples_cache_filename:\", samples_cache_filename)\n",
    "    is_samples_cache_file = os.path.isfile(samples_cache_file)\n",
    "    if is_samples_cache_file:\n",
    "        with open(samples_cache_file, 'rb') as f:\n",
    "            samples_vae[model_cache_filename] = pickle.load(f)\n",
    "        print('samples loaded from file')\n",
    "    else:\n",
    "        print('generating samples')\n",
    "        samples_vae[model_cache_filename] = vae_get_samples(model_sample, n_samples, scaler)\n",
    "        with open(samples_cache_file, 'wb') as f:\n",
    "            pickle.dump(samples_vae[model_cache_filename], f)\n",
    "        print('samples saved to file')\n",
    "    print(80 * '-')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# performance for all atributes (development set)\n",
    "stat_vae_dev = get_stat(data_train[train_dev_ind:], samples_vae, data_train, bin_n_comparisons,\n",
    "                        samples_dev_cache_dir, vae_model_cache_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# performance for all atributes (development set)\n",
    "(errors_vae_dev, \n",
    " diversity_1_vae_dev, \n",
    " diversity_2_vae_dev,\n",
    " xticks_vae_dev) = get_perform_plot_data(stat_vae_dev, bin_n_comparisons, model_cache_files_vae, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# performance for the selected atributes (development set)\n",
    "stat_vae_dev_selected = get_stat(data_train[train_dev_ind:], samples_vae, data_train, bin_n_comparisons,\n",
    "                                 samples_dev_cache_dir, vae_model_cache_dir,\n",
    "                                 var_is_ignore, var_g_is_ignore)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# performance for the selected atributes (development set)\n",
    "(errors_vae_dev_selected, \n",
    " diversity_1_vae_dev_selected, \n",
    " diversity_2_vae_dev_selected,\n",
    " xticks_vae_dev_selected) = get_perform_plot_data(stat_vae_dev_selected, bin_n_comparisons, \n",
    "                                                  model_cache_files_vae, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# performance for all atributes (test set)\n",
    "stat_vae = get_stat(data_test, samples_vae, data_train, bin_n_comparisons,\n",
    "                    samples_true_cache_dir, vae_model_cache_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# performance for all atributes (test set)\n",
    "(errors_vae, \n",
    " diversity_1_vae, \n",
    " diversity_2_vae,\n",
    " xticks_vae) = get_perform_plot_data(stat_vae, bin_n_comparisons, model_cache_files_vae, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# performance for the selected atributes (test set)\n",
    "stat_vae_selected = get_stat(data_test, samples_vae, data_train, bin_n_comparisons,\n",
    "                             samples_true_cache_dir, vae_model_cache_dir,\n",
    "                             var_is_ignore, var_g_is_ignore)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# performance for the selected atributes (test set)\n",
    "(errors_vae_selected, \n",
    " diversity_1_vae_selected, \n",
    " diversity_2_vae_selected,\n",
    " xticks_vae_selected) = get_perform_plot_data(stat_vae_selected, bin_n_comparisons, \n",
    "                                              model_cache_files_vae, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot marginals\n",
    "for k in model_cache_files_vae:\n",
    "    print(80 * '-')\n",
    "    print(k)\n",
    "    check_marginals_numerical([data_train, samples_vae[k]], bin_n_comparisons)\n",
    "    check_marginals_categorical([data_train, samples_vae[k]])   \n",
    "    print(80 * '-')\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MCMC (Gibbs) Full Conditional Baseline: p(x\\_i | x\\_-i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get samples from a single starting point\n",
    "def full_cond_get_samples(full_cond, num_bin, n_samples):\n",
    "    print('full cond bin_n={}: generating {} samples'.format(num_bin, n_samples))\n",
    "    #\n",
    "    samples_cond = []\n",
    "    #\n",
    "    if num_bin in min_max_bins:\n",
    "        min_max_num_bin = min_max_bins[num_bin]\n",
    "    else:\n",
    "        min_max_num_bin = calc_min_max_num_bin(data_train, num_bin)\n",
    "    #\n",
    "    idx = np.random.randint(len(data_train))\n",
    "    sample = data_train[idx]\n",
    "    for i in range(burn_out + n_samples * thinning):\n",
    "        if i % 1000 == 0: print(i)\n",
    "        for var_i in range(numerical_col_n):\n",
    "            var_name = 'num' + str(var_i)\n",
    "            bin_key = get_bin_key(sample, min_max_num_bin, [var_i], None)\n",
    "            if bin_key in full_cond[var_name]:\n",
    "                val_i = np.random.randint(len(full_cond[var_name][bin_key]))\n",
    "                sample[var_i] = full_cond[var_name][bin_key][val_i]\n",
    "            else:\n",
    "                print(\"empty bin\")\n",
    "        for var_g_i in range(cat_groups_n):\n",
    "            var_name = 'cat' + str(var_g_i)\n",
    "            var_g_i_beg = cat_groups[var_g_i]\n",
    "            var_g_i_end = cat_groups[var_g_i + 1]\n",
    "            bin_key = get_bin_key(sample, min_max_num_bin, None, [var_g_i])\n",
    "            if bin_key in full_cond[var_name]:\n",
    "                val_i = np.random.randint(len(full_cond[var_name][bin_key]))\n",
    "                sample[var_g_i_beg: var_g_i_end] = full_cond[var_name][bin_key][val_i]\n",
    "            else:\n",
    "                print(\"empty bin\")\n",
    "        if i >= burn_out and (i - burn_out) % thinning == 0:\n",
    "            samples_cond.append(copy.deepcopy(sample))\n",
    "    samples_cond = np.array(samples_cond)\n",
    "    return samples_cond\n",
    "\n",
    "\n",
    "# get samples from several starting points\n",
    "def full_cond_get_samples_over_x0(full_cond, num_bin, n_samples):\n",
    "    print('full cond bin_n={}: generating {} samples'.format(num_bin, n_samples))\n",
    "    #\n",
    "    samples_cond = []\n",
    "    print_time_per_sample = True\n",
    "    #\n",
    "    if num_bin in min_max_bins:\n",
    "        min_max_num_bin = min_max_bins[num_bin]\n",
    "    else:\n",
    "        min_max_num_bin = calc_min_max_num_bin(data_train, num_bin)\n",
    "    #\n",
    "    for n_sample in range(n_samples):\n",
    "        if print_time_per_sample:\n",
    "            start = time.time()\n",
    "        if n_sample % 100 == 0: \n",
    "            print('n_sample =', n_sample)\n",
    "        idx = np.random.randint(len(data_train))\n",
    "        sample = data_train[idx]\n",
    "        for i in range(burn_out):\n",
    "            for var_i in range(numerical_col_n):\n",
    "                var_name = 'num' + str(var_i)\n",
    "                bin_key = get_bin_key(sample, min_max_num_bin, [var_i], None)\n",
    "                if bin_key in full_cond[var_name]:\n",
    "                    val_i = np.random.randint(len(full_cond[var_name][bin_key]))\n",
    "                    sample[var_i] = full_cond[var_name][bin_key][val_i]\n",
    "                else:\n",
    "                    #print(\"n_sample={}, step i={}: empty bin (num)\".format(n_sample, i))\n",
    "                    pass\n",
    "            for var_g_i in range(cat_groups_n):\n",
    "                var_name = 'cat' + str(var_g_i)\n",
    "                var_g_i_beg = cat_groups[var_g_i]\n",
    "                var_g_i_end = cat_groups[var_g_i + 1]\n",
    "                bin_key = get_bin_key(sample, min_max_num_bin, None, [var_g_i])\n",
    "                if bin_key in full_cond[var_name]:\n",
    "                    val_i = np.random.randint(len(full_cond[var_name][bin_key]))\n",
    "                    sample[var_g_i_beg: var_g_i_end] = full_cond[var_name][bin_key][val_i]\n",
    "                else:\n",
    "                    #print(\"n_sample={}, step i={}: empty bin (cat)\".format(n_sample, i))\n",
    "                    pass\n",
    "        samples_cond.append(copy.deepcopy(sample))\n",
    "        if print_time_per_sample:\n",
    "            duration = time.time() - start\n",
    "            print(\"Took {} seconds to get one sample\".format(duration))\n",
    "            print_time_per_sample = False\n",
    "    samples_cond = np.array(samples_cond)\n",
    "    return samples_cond"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# parameters of the Gibbs sampler\n",
    "calc_over_x0 = False\n",
    "if calc_over_x0:\n",
    "    n_samples = 1000 # it is very computationally expencive\n",
    "    burn_out = 10000\n",
    "    thinning = -1\n",
    "else:\n",
    "    n_samples = 100000\n",
    "    burn_out = 20000\n",
    "    thinning = 20    \n",
    "full_cond_model_cache_dir = os.path.join(model_cache_dir, 'full_cond')\n",
    "if not os.path.exists(full_cond_model_cache_dir): \n",
    "    os.makedirs(full_cond_model_cache_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# calculate sampling probabilits for numerical variables using 1) scheme_1 in min_max_bins, 2) 10 uniform bins\n",
    "bin_ns_cond = ['scheme_1', 10]\n",
    "#\n",
    "samples_cond = {}\n",
    "for bin_n in bin_ns_cond:\n",
    "    # get model\n",
    "    model_cache_file = 'model_bin_n=' + str(bin_n) + '--' + file_suffix + '.pickle'\n",
    "    model_cache_file = os.path.join(full_cond_model_cache_dir, model_cache_file)\n",
    "    print(\"full cond cache file:\", model_cache_file)\n",
    "    is_model_cache_file = os.path.isfile(model_cache_file)\n",
    "    if is_model_cache_file:\n",
    "        with open(model_cache_file, 'rb') as f:\n",
    "            full_cond = pickle.load(f)\n",
    "        print(\"full cond loaded from file\")\n",
    "    else:\n",
    "        print(\"calculating full cond\")\n",
    "        start = time.time()\n",
    "        full_cond = calc_full_conditionals(data_train, bin_n)\n",
    "        duration = time.time() - start\n",
    "        print(\"Took {} seconds to calculate full cond\".format(duration))\n",
    "        with open(model_cache_file, 'wb') as f:\n",
    "            pickle.dump(full_cond, f)\n",
    "    # generate samples\n",
    "    samples_cache_file = 'samples_bin_n=' + str(bin_n)\n",
    "    samples_cache_file += '--n_samples=' + str(n_samples)\n",
    "    samples_cache_file += '--burn_out=' + str(burn_out)\n",
    "    if thinning >= 0:\n",
    "        samples_cache_file += '--thinning=' + str(thinning)\n",
    "    else:\n",
    "        samples_cache_file += '--over_x0'\n",
    "    samples_cache_file += '--' + file_suffix + '.pickle'\n",
    "    samples_cache_file = os.path.join(full_cond_model_cache_dir, samples_cache_file)\n",
    "    print(\"samples_cache_file:\", samples_cache_file)\n",
    "    is_samples_cache_file = os.path.isfile(samples_cache_file)\n",
    "    if is_samples_cache_file:\n",
    "        with open(samples_cache_file, 'rb') as f:\n",
    "            samples_cond[bin_n] = pickle.load(f)\n",
    "        print(\"samples loaded from file\")\n",
    "    else:\n",
    "        print('generating samples')\n",
    "        start = time.time()\n",
    "        if thinning >= 0:\n",
    "            samples_cond[bin_n] = full_cond_get_samples(full_cond, bin_n, n_samples)\n",
    "        else:\n",
    "            samples_cond[bin_n] = full_cond_get_samples_over_x0(full_cond, bin_n, n_samples)\n",
    "        duration = time.time() - start\n",
    "        print(\"Took {} seconds to generate samples\".format(duration))\n",
    "        with open(samples_cache_file, 'wb') as f:\n",
    "            pickle.dump(samples_cond[bin_n], f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# performance for all atributes\n",
    "stat_cond = get_stat(data_test, samples_cond, data_train, bin_n_comparisons,\n",
    "                     samples_true_cache_dir, full_cond_model_cache_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# performance for all atributes\n",
    "(errors_cond, \n",
    " diversity_1_cond, \n",
    " diversity_2_cond,\n",
    " xticks_cond) = get_perform_plot_data(stat_cond, bin_n_comparisons, bin_ns_cond)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# performance for the selected atributes\n",
    "stat_cond_selected = get_stat(data_test, samples_cond, data_train, bin_n_comparisons,\n",
    "                              samples_true_cache_dir, full_cond_model_cache_dir,\n",
    "                              var_is_ignore, var_g_is_ignore)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# performance for the selected atributes\n",
    "(errors_cond_selected, \n",
    " diversity_1_cond_selected, \n",
    " diversity_2_cond_selected,\n",
    " xticks_cond_selected) = get_perform_plot_data(stat_cond_selected, \n",
    "                                               bin_n_comparisons, bin_ns_cond)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot marginals\n",
    "bin_n_comparison = 'scheme_1'\n",
    "for bin_n in bin_ns_cond:\n",
    "    print(80 * '-')\n",
    "    print(bin_n)\n",
    "    check_marginals_numerical([data_train, samples_cond[bin_n]], bin_n_comparison)\n",
    "    check_marginals_categorical([data_train, samples_cond[bin_n]])\n",
    "    print(80 * '-')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MCMC Marginal baseline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_samples_marg(data, n_samples):\n",
    "    samples_marg = []\n",
    "    for i in range(n_samples):\n",
    "        if i % 1000 == 0: print(i)\n",
    "        sample = []\n",
    "        for i in range(numerical_col_n):\n",
    "            val = np.random.choice(data[:, i], 1)[0]\n",
    "            #print(val)\n",
    "            #raise Exception(\"stop\")\n",
    "            sample.append(val)\n",
    "        for g_i in range(cat_groups_n):\n",
    "            g_i_beg = cat_groups[g_i]\n",
    "            g_i_end = cat_groups[g_i + 1]\n",
    "            idx = np.random.randint(len(data))\n",
    "            val = data[idx, g_i_beg: g_i_end]\n",
    "            #print(sample, val)\n",
    "            sample.extend(val)\n",
    "        samples_marg.append(sample)\n",
    "    samples_marg = np.array(samples_marg)\n",
    "    return samples_marg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "n_samples = 100000\n",
    "samples_marg_keys = ['marg']\n",
    "marg_cache_dir = os.path.join(model_cache_dir, 'marginal')\n",
    "if not os.path.exists(marg_cache_dir): \n",
    "    os.makedirs(marg_cache_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "samples_marg = {}\n",
    "# generate samples\n",
    "samples_cache_file = 'samples_marg'\n",
    "samples_cache_file += '--n_samples=' + str(n_samples)\n",
    "samples_cache_file += '--' + file_suffix + '.pickle'\n",
    "samples_cache_file = os.path.join(marg_cache_dir, samples_cache_file)\n",
    "print(\"samples_cache_file:\", samples_cache_file)\n",
    "is_samples_cache_file = os.path.isfile(samples_cache_file)\n",
    "if is_samples_cache_file:\n",
    "    with open(samples_cache_file, 'rb') as f:\n",
    "        samples_marg[samples_marg_keys[0]] = pickle.load(f)\n",
    "    print(\"samples loaded from file\")\n",
    "else:\n",
    "    print('generating samples')\n",
    "    start = time.time()\n",
    "    samples_marg[samples_marg_keys[0]] = get_samples_marg(data_train, n_samples)\n",
    "    duration = time.time() - start\n",
    "    print(\"Took {} seconds to generate samples\".format(duration))\n",
    "    with open(samples_cache_file, 'wb') as f:\n",
    "        pickle.dump(samples_marg[samples_marg_keys[0]], f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# performance for all atributes\n",
    "stat_marg = get_stat(data_test, samples_marg, data_train, bin_n_comparisons,\n",
    "                     samples_true_cache_dir, marg_cache_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# performance for all atributes\n",
    "(errors_marg, \n",
    " diversity_1_marg, \n",
    " diversity_2_marg,\n",
    " xticks_marg) = get_perform_plot_data(stat_marg, bin_n_comparisons, samples_marg_keys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# performance for the selected atributes\n",
    "stat_marg_selected = get_stat(data_test, samples_marg, data_train, bin_n_comparisons,\n",
    "                              samples_true_cache_dir, marg_cache_dir,\n",
    "                              var_is_ignore, var_g_is_ignore)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# performance for the selected atributes\n",
    "(errors_marg_selected, \n",
    " diversity_1_marg_selected, \n",
    " diversity_2_marg_selected,\n",
    " xticks_marg_selected) = get_perform_plot_data(stat_marg_selected, \n",
    "                                               bin_n_comparisons, samples_marg_keys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot marginals\n",
    "for bin_n_c in bin_n_comparisons:\n",
    "    print(80 * '-')\n",
    "    print(bin_n_c)\n",
    "    check_marginals_numerical([data_train, samples_marg[samples_marg_keys[0]]], bin_n_c)\n",
    "    check_marginals_categorical([data_train, samples_marg[samples_marg_keys[0]]])\n",
    "    print(80 * '-')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# comparison for all attributes\n",
    "h_labels = bin_ns_cond + samples_train_keys + samples_marg_keys\n",
    "x_vae = range(len(model_cache_files_vae))\n",
    "for bin_n_c in bin_n_comparisons:\n",
    "    print(80 * \"-\")\n",
    "    print(\"bin_n_comparison =\", bin_n_c)\n",
    "    plot_res(x_vae, errors_vae[bin_n_c], xticks=xticks_vae[bin_n_c],\n",
    "             h_lines=errors_cond[bin_n_c] + errors_train[bin_n_c] \n",
    "                     + errors_marg[bin_n_c]\n",
    "             , \n",
    "             h_labels=h_labels,\n",
    "             **error_plot_params[error_measure])\n",
    "    plot_res(x_vae, diversity_1_vae[bin_n_c], xticks=xticks_vae[bin_n_c],\n",
    "             h_lines=diversity_1_cond[bin_n_c] + diversity_1_train[bin_n_c] \n",
    "                     + diversity_1_marg[bin_n_c]\n",
    "             , \n",
    "             h_labels=h_labels,\n",
    "             **diversity_plot_params[diversity_measure_1])\n",
    "    plot_res(x_vae, diversity_2_vae[bin_n_c], xticks=xticks_vae[bin_n_c],\n",
    "             h_lines=diversity_2_cond[bin_n_c] + diversity_2_train[bin_n_c] \n",
    "                     + diversity_2_marg[bin_n_c]\n",
    "             , \n",
    "             h_labels=h_labels,\n",
    "             **diversity_plot_params[diversity_measure_2])\n",
    "    print(80 * \"-\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# comparison for the selected attributes\n",
    "h_labels = bin_ns_cond + samples_train_keys + samples_marg_keys\n",
    "x_vae = range(len(model_cache_files_vae))\n",
    "for bin_n_c in bin_n_comparisons:\n",
    "    print(80 * \"-\")\n",
    "    print(\"bin_n_comparison =\", bin_n_c)\n",
    "    plot_res(x_vae, errors_vae_selected[bin_n_c], xticks=xticks_vae_selected[bin_n_c],\n",
    "             h_lines=errors_cond_selected[bin_n_c] + errors_train_selected[bin_n_c] \n",
    "                     + errors_marg_selected[bin_n_c]\n",
    "             , \n",
    "             h_labels=h_labels,\n",
    "             **error_plot_params[error_measure])\n",
    "    plot_res(x_vae, diversity_1_vae_selected[bin_n_c], xticks=xticks_vae_selected[bin_n_c],\n",
    "             h_lines=diversity_1_cond_selected[bin_n_c] + diversity_1_train_selected[bin_n_c] \n",
    "                     + diversity_1_marg_selected[bin_n_c]\n",
    "             , \n",
    "             h_labels=h_labels,\n",
    "             **diversity_plot_params[diversity_measure_1])\n",
    "    plot_res(x_vae, diversity_2_vae_selected[bin_n_c], xticks=xticks_vae_selected[bin_n_c],\n",
    "             h_lines=diversity_2_cond_selected[bin_n_c] + diversity_2_train_selected[bin_n_c] \n",
    "                     + diversity_2_marg_selected[bin_n_c]\n",
    "             , \n",
    "             h_labels=h_labels,\n",
    "             **diversity_plot_params[diversity_measure_2])\n",
    "    print(80 * \"-\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# find the best performing model\n",
    "def get_best_model(stat):\n",
    "    names = {}\n",
    "    lb = ['srmse', 'rmse', 'mae']\n",
    "    for bin_n_c in stat:\n",
    "        best_k = {}\n",
    "        for k in stat[bin_n_c]:\n",
    "            for name in stat[bin_n_c][k]['joint_pdf']:\n",
    "                if name not in best_k:\n",
    "                    best_k[name] = (k, stat[bin_n_c][k]['joint_pdf'][name])\n",
    "                else:\n",
    "                    if name in lb and stat[bin_n_c][k]['joint_pdf'][name] < best_k[name][1]:\n",
    "                        best_k[name] = (k, stat[bin_n_c][k]['joint_pdf'][name])\n",
    "                    elif name not in lb and stat[bin_n_c][k]['joint_pdf'][name] > best_k[name][1]:\n",
    "                        best_k[name] = (k, stat[bin_n_c][k]['joint_pdf'][name])            \n",
    "        print(80 * '-')\n",
    "        for name in best_k:\n",
    "            names[name] = best_k[name][0]\n",
    "            print(\"best\", name)\n",
    "            print(best_k[name][0])\n",
    "            pprint(stat[bin_n_c][best_k[name][0]]['joint_pdf'])\n",
    "            print('.....')\n",
    "        print(80 * '-')\n",
    "    return names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print('VAE')\n",
    "get_best_model(stat_vae_dev_selected)\n",
    "get_best_model(stat_vae_selected)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print('cond')\n",
    "get_best_model(stat_cond_selected)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print('marg')\n",
    "get_best_model(stat_marg_selected)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print('data_train')\n",
    "get_best_model(stat_train_selected)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py3",
   "language": "python",
   "name": "py3"
  },
  "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}