Bottlenecks-CLUB/exp_colored_mnist_P1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"7\"\n",
    "import tensorflow as tf\n",
    "import logging\n",
    "tf.get_logger().setLevel(logging.ERROR)\n",
    "import tensorflow.keras.backend as K\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from tensorflow.keras import activations\n",
    "from tensorflow.keras.models import Model,Sequential\n",
    "from tensorflow.keras.layers import Dense,Conv2D,BatchNormalization,Lambda,Activation,Input,Flatten,Conv2DTranspose,concatenate,Reshape,LeakyReLU"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set some options and read data based on selected ones"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#constants\n",
    "IMG_DIM = 28\n",
    "# manual options\n",
    "sensitive = 'class' # 'color' or 'class'\n",
    "prior_dist = 0 # 0: simple gaussian 1: guassian mixture Note: if prior generator model enabled, this option is not used\n",
    "mnist_type = 'unbalanced' # 'balanced' 'unbalanced\n",
    "supervised = True #False:Unsupervised (use input data as utility)\n",
    "data_path = \"./Colored_MNIST\"\n",
    "\n",
    "# automatically set variables based on selected options\n",
    "utility = 'class' if sensitive=='color' else 'color'\n",
    "prior_type = 'Isotropic' if prior_dist == 0 else 'Mixture'\n",
    "data_type = 'biased' if mnist_type=='unbalanced' else 'uniform'\n",
    "\n",
    "exp_info = f\"s_{sensitive}_u_{utility}_{prior_type}_{data_type}_supervised\" if supervised else f\"s_{sensitive}_{prior_type}_{data_type}_unsupervised\"\n",
    "exp_info += \"_P1\"\n",
    "\n",
    "data = np.load(f\"{data_path}/colored_mnist_{mnist_type}.npz\")\n",
    "\n",
    "x_train = data['x_train']\n",
    "x_train = x_train.astype(np.float32) / 255.0\n",
    "s_train = data['s_train'] if(sensitive == 'color') else data['y_train']\n",
    "u_train = data['s_train'] if(sensitive == 'class') else data['y_train']\n",
    "if not supervised: u_train = x_train\n",
    "\n",
    "\n",
    "x_test = data['x_test']\n",
    "x_test = x_test.astype(np.float32) / 255.0\n",
    "s_test = data['s_test'] if(sensitive == 'color') else data['y_test']\n",
    "u_test = data['s_test'] if(sensitive == 'class') else data['y_test']\n",
    "if not supervised: u_test = x_test\n",
    "\n",
    "\n",
    "exp_info"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# one hot S and U (if supervised) in order to ready them for the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def one_hot(labels, num_labels):\n",
    "    num_labels_data = labels.shape[0]\n",
    "    one_hot_encoding = np.zeros((num_labels_data,num_labels))\n",
    "    one_hot_encoding[np.arange(num_labels_data),labels] = 1\n",
    "    one_hot_encoding = np.reshape(one_hot_encoding, [-1, num_labels])\n",
    "    return one_hot_encoding\n",
    "\n",
    "DIM_S = 3 if(sensitive == 'color') else 10\n",
    "if supervised:\n",
    "    DIM_U = 10 if(sensitive == 'color') else 3\n",
    "else:\n",
    "    DIM_U = x_train[0].shape\n",
    "\n",
    "if supervised: u_train = one_hot(u_train, DIM_U).astype(np.float32)\n",
    "s_train = one_hot(s_train, DIM_S).astype(np.float32)\n",
    "if supervised: u_test = one_hot(u_test, DIM_U).astype(np.float32)\n",
    "s_test = one_hot(s_test, DIM_S).astype(np.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class sampling(tf.keras.layers.Layer):\n",
    "    \"\"\"Uses (z_mean, z_log_var) to sample z\"\"\"\n",
    "\n",
    "    def call(self, inputs):\n",
    "        z_mean, z_log_var = inputs\n",
    "        batch = tf.shape(z_mean)[0]\n",
    "        dim = tf.shape(z_mean)[1]\n",
    "        epsilon = tf.keras.backend.random_normal(shape=(batch, dim))\n",
    "        return tf.math.add(z_mean, tf.exp(0.5 * z_log_var) * epsilon, name=\"sample\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_encoder(DIM_Z, input_x):\n",
    "    stride = 2\n",
    "#     input_x = Input( shape = [IMG_DIM,IMG_DIM,3], name=\"x\" )\n",
    "\n",
    "    #first hidden layer\n",
    "    x = Conv2D(64, 5, strides=stride, padding=\"same\", name=\"enc_h1\")(input_x)\n",
    "    x = BatchNormalization(name=\"enc_h1_normalized\")(x)\n",
    "    x = Activation(LeakyReLU(), name=\"enc_h1_activation\")(x)\n",
    "    #second hidden layer\n",
    "    x = Conv2D(128, 5, strides=stride, padding=\"same\", name=\"enc_h2\")(x)\n",
    "    x = BatchNormalization(name=\"enc_h2_normalized\")(x)\n",
    "    x = Activation(LeakyReLU(), name=\"enc_h2_activation\")(x)\n",
    "\n",
    "    shape = K.int_shape(x)\n",
    "    x = Flatten()(x)\n",
    "    x = Dense(DIM_Z*4, name=\"enc_dense_1\")(x)\n",
    "    x = BatchNormalization()(x)\n",
    "    x = Activation('tanh')(x)\n",
    "\n",
    "    z_mean = Dense(DIM_Z, name=\"z_mean\")(x)\n",
    "    z_log_sigma_sq = Dense(DIM_Z, name=\"z_sigma\")(x)\n",
    "#     z = Lambda(sampling, output_shape=DIM_Z, name='sample_z')([z_mean, z_log_sigma_sq])\n",
    "    z = sampling(name=\"sample_z\")([z_mean, z_log_sigma_sq])\n",
    "#     prior_loss = K.mean(-0.5 * K.sum(1 + z_log_sigma_sq - K.square(z_mean) - K.exp(z_log_sigma_sq), axis=-1))\n",
    "    \n",
    "    encoder = Model(input_x, z, name = \"Encoder\")\n",
    "#     encoder.add_loss((alpha+beta) * prior_loss)\n",
    "    \n",
    "    return (encoder, z_mean, z_log_sigma_sq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_decoder(DIM_Z):\n",
    "    stride = 2\n",
    "    input_z = Input(shape = (DIM_Z,), name=\"z_encoder\")\n",
    "    input_s = Input( shape = (DIM_S,), name=\"s_input\")\n",
    "    z_with_s = concatenate([input_z, input_s], name=\"concat_layer\")\n",
    "    x = Dense(7*7*128, name=\"dec_dense_1\")(z_with_s)\n",
    "    x = BatchNormalization(name=\"dec_dense_1_normalized\")(x)\n",
    "    x = Activation(LeakyReLU(), name=\"dec_dense_1_activation\")(x)\n",
    "\n",
    "    #Unflatten\n",
    "    x = Reshape((7,7,128), name=\"unflatten1\")(x)\n",
    "    x = BatchNormalization(name=\"unflatten1_normalized\")(x)\n",
    "    x = Activation(LeakyReLU(), name=\"unflatten1_activation\")(x)\n",
    "\n",
    "    x = Conv2DTranspose(64, 5, strides=stride, padding = \"same\", name=\"dec_h1\")(x)\n",
    "    x = BatchNormalization(name=\"dec_h1_normalized\")(x)\n",
    "    x = Activation(LeakyReLU(), name=\"dec_h1_activation\")(x)\n",
    "\n",
    "    x = Conv2DTranspose(3, 5, strides=stride, padding = \"same\", name=\"dec_h2\")(x)\n",
    "    x_hat = Activation('sigmoid', name=\"x_hat\")(x)\n",
    "\n",
    "    decoder = Model([input_z, input_s], x_hat, name = \"Uncertainty_Decoder\")\n",
    "    return decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if supervised:\n",
    "    def get_utility_model(DIM_Z):\n",
    "        model = Sequential(name=\"Utility_Decoder_supervised\")\n",
    "        model.add(Dense(DIM_Z * 4, input_dim=DIM_Z))\n",
    "        model.add(BatchNormalization())\n",
    "        model.add(LeakyReLU(alpha=0.2))\n",
    "        model.add(Dense(DIM_U, activation='softmax', name=\"u_hat\"))\n",
    "        return model\n",
    "else:\n",
    "    def get_utility_model(DIM_Z):\n",
    "        stride = 2\n",
    "        model = Sequential(name=\"Utility_Decoder_unsupervised\")\n",
    "        model.add(Dense(7*7*128, input_dim=DIM_Z))\n",
    "        model.add(BatchNormalization())\n",
    "        model.add(LeakyReLU(alpha=0.2))\n",
    "        \n",
    "        model.add(Reshape((7,7,128)))\n",
    "        \n",
    "        model.add(Conv2DTranspose(64, 5, strides=stride, padding=\"same\"))\n",
    "        model.add(BatchNormalization())\n",
    "        model.add(LeakyReLU(alpha=0.2))\n",
    "        \n",
    "        model.add(Conv2DTranspose(3, 5, strides=stride, activation='sigmoid',padding=\"same\", name=\"u_hat\"))\n",
    "        return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_z_discriminator(DIM_Z):\n",
    "    model = Sequential(name=\"Latent_Space_Discriminator\")\n",
    "    \n",
    "    model.add(Dense(512, input_dim=DIM_Z))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(LeakyReLU(alpha=0.2))\n",
    "    \n",
    "    model.add(Dense(128))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(LeakyReLU(alpha=0.2))\n",
    "    \n",
    "    model.add(Dense(1, activation=\"sigmoid\"))\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_utility_discriminator():\n",
    "    stride = 2\n",
    "    model = Sequential(name=\"Attribute_Class_Discriminator\")\n",
    "\n",
    "    model.add(Dense(DIM_U * 4, input_dim=DIM_U))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(LeakyReLU(alpha=0.2))\n",
    "    \n",
    "    model.add(Dense(DIM_U * 4))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(LeakyReLU(alpha=0.2))\n",
    "    \n",
    "    model.add(Dense(1, activation=\"sigmoid\"))\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_visible_space_discriminator():\n",
    "    stride = 2\n",
    "    model = Sequential(name=\"Visible_Space_Discriminator\")\n",
    "    \n",
    "    model.add(Input(shape = [IMG_DIM,IMG_DIM,3]))\n",
    "    \n",
    "    model.add(Conv2D(32, 3, strides=stride, padding=\"same\"))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(LeakyReLU(alpha=0.2))\n",
    "    \n",
    "    model.add(Conv2D(64, 3, strides=stride, padding=\"same\"))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(LeakyReLU(alpha=0.2))\n",
    "    \n",
    "    model.add(Conv2D(128, 3, strides=stride, padding=\"same\"))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(LeakyReLU(alpha=0.2))\n",
    "    \n",
    "    model.add(Flatten())\n",
    "    \n",
    "    model.add(Dense(128))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(LeakyReLU(alpha=0.2))\n",
    "    \n",
    "    model.add(Dense(1, activation=\"sigmoid\"))\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_prior_generator(DIM_Z, noise_dim=50):\n",
    "    model = Sequential(name=\"Prior_Distribution_Generator\")\n",
    "    \n",
    "    model.add(Dense(noise_dim, input_dim=noise_dim))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(LeakyReLU(alpha=0.2))\n",
    "    model.add(Dense(noise_dim))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(LeakyReLU(alpha=0.2))\n",
    "    model.add(Dense(DIM_Z, name=\"z_prior_generator\"))\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build and train a model for evaluating S performance\n",
    "The following two block cells are **not** part of the framework, it's an evaluator for sensitivity which is required for the experiments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_s_eval_model():\n",
    "    stride = 2\n",
    "    model = Sequential(name=\"CNN_Model_sensitivity_evaluator\")\n",
    "    model.add(Input(shape = [IMG_DIM,IMG_DIM,3]))\n",
    "\n",
    "    model.add(Conv2D(64, 5, strides=stride, padding=\"same\"))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(Activation(LeakyReLU(alpha=0.2)))\n",
    "\n",
    "    model.add(Conv2D(128, 5, strides=stride, padding=\"same\"))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(Activation(LeakyReLU(alpha=0.2)))\n",
    "\n",
    "    model.add(Flatten())\n",
    "\n",
    "    model.add(Dense(512 * DIM_S))\n",
    "    model.add(BatchNormalization())\n",
    "    model.add(Activation(\"relu\"))\n",
    "\n",
    "    model.add(Dense(DIM_S, activation=\"softmax\", name='s_hat'))\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s_eval_model = get_s_eval_model()\n",
    "s_eval_model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.001), loss='mse')\n",
    "\n",
    "force_train = False\n",
    "\n",
    "if not os.path.exists(f\"saved_models/mnist_eval_model_s_{sensitive}.h5\") or force_train:\n",
    "    print(\"Training sensitive attribute evaluator\")\n",
    "    history = s_eval_model.fit(x_train, s_train, batch_size=1024, epochs=100, shuffle=True, verbose=2)\n",
    "    s_eval_model.save_weights(f\"saved_models/mnist_eval_model_s_{sensitive}.h5\")\n",
    "else:\n",
    "    print(\"Loading sensitive attribute evaluator from file\")\n",
    "    s_eval_model.load_weights(f\"saved_models/mnist_eval_model_s_{sensitive}.h5\")\n",
    "\n",
    "## Evaluating evaluator!\n",
    "s_hat_test = s_eval_model.predict(x_test)\n",
    "print(f\"Evaluator accuracy = {(sum(np.argmax(s_hat_test,axis=-1)==np.argmax(s_test,axis=-1)) / s_test.shape[0]) * 100}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Putting all togther - define CLUB model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z_gen_enabled = False\n",
    "## utility and reconstruction losses\n",
    "if supervised:\n",
    "    def loss_u(u_true, u_pred):\n",
    "        return K.mean(K.sum(K.square(u_true-u_pred), axis=-1))\n",
    "else:\n",
    "    def loss_u(u_true, u_pred):\n",
    "        return K.mean(K.sum(K.square(u_true-u_pred), axis=(1,2,3)))\n",
    "\n",
    "def loss_x(alpha):\n",
    "    def loss(inp_x, out_x):\n",
    "        return alpha * K.mean(K.sum(K.square(inp_x-out_x), axis=(1,2,3)))\n",
    "    return loss\n",
    "\n",
    "bce = tf.keras.losses.BinaryCrossentropy()\n",
    "# Weighted cross-entropy loss\n",
    "def loss_wce(coef):\n",
    "    def loss(y, y_pred):\n",
    "         return coef * bce(y, y_pred)\n",
    "    return loss\n",
    "\n",
    "########## Inputs\n",
    "def get_full_model(DIM_Z, alpha, beta, learning_rate=0.0001, dim_noise = 100):\n",
    "\n",
    "    input_x = Input( shape=[IMG_DIM,IMG_DIM,3], name=\"x\" )\n",
    "    input_s = Input( shape = (DIM_S,), name=\"s\" )\n",
    "    input_z = Input( shape = (DIM_Z,), name=\"z\" )\n",
    "\n",
    "    input_noise = Input( shape = (dim_noise,), name=\"Noise\")\n",
    "\n",
    "    ########## Define AE: Encoder, Utility Decoder and Uncertainty Decoder\n",
    "    encoder,z_mean,z_log_sigma_sq = get_encoder(DIM_Z, input_x)\n",
    "    uncertainty_decoder = get_decoder(DIM_Z)\n",
    "    utility_decoder = get_utility_model(DIM_Z)\n",
    "\n",
    "    z = encoder(input_x)\n",
    "    x_hat = uncertainty_decoder([z, input_s])\n",
    "    u_hat = utility_decoder(z)\n",
    "\n",
    "    autoencoder = Model([input_x, input_s], [x_hat, u_hat], name=\"CLUB_Autoencoder\")\n",
    "    prior_loss = (alpha+beta) * K.mean(-0.5 * K.sum(1 + z_log_sigma_sq - K.square(z_mean) - K.exp(z_log_sigma_sq), axis=-1))\n",
    "    prior_loss = tf.identity(prior_loss, name=\"kl_loss\")\n",
    "    autoencoder.add_loss(prior_loss)\n",
    "    autoencoder.compile(loss=[loss_x(alpha), loss_u], optimizer=tf.keras.optimizers.Adam(lr=learning_rate*5))\n",
    "\n",
    "    ########## Define Latent Space Discriminator\n",
    "    z_discriminator = get_z_discriminator(DIM_Z)\n",
    "    z_discriminator.compile(loss=loss_wce(0.1*(alpha+beta)), optimizer=tf.keras.optimizers.Adam(lr=learning_rate))\n",
    "\n",
    "    ##########\n",
    "    #z_discriminator should train separately\n",
    "    z_discriminator.trainable = False\n",
    "    prior_generator = get_prior_generator(DIM_Z, dim_noise)\n",
    "    prior_gen_zdiscriminator = Model(input_noise, z_discriminator(prior_generator(input_noise)), name=\"CLUB_generator_zdiscriminator\")\n",
    "    prior_gen_zdiscriminator.compile(loss=loss_wce(-0.1*(alpha+beta)), optimizer=tf.keras.optimizers.Adam(lr=learning_rate))\n",
    "\n",
    "    ##########\n",
    "    encoder_zdiscriminator = Model(input_x, z_discriminator(encoder(input_x)), name=\"CLUB_encoder_zdiscriminator\")\n",
    "    encoder_zdiscriminator.compile(loss=loss_wce(-0.1*(alpha+beta)), optimizer=tf.keras.optimizers.Adam(lr=learning_rate))\n",
    "\n",
    "    ##########\n",
    "    u_dircriminator = get_utility_discriminator() if supervised else get_visible_space_discriminator()\n",
    "    u_dircriminator.compile(loss=loss_wce(0.1), optimizer=tf.keras.optimizers.Adam(lr=learning_rate))\n",
    "\n",
    "    ##########\n",
    "    #u_dircriminator should train separately\n",
    "    u_dircriminator.trainable = False\n",
    "\n",
    "    if z_gen_enabled:    \n",
    "        decoder_udiscriminator = Model(input_noise, u_dircriminator(utility_decoder(prior_generator(input_noise))), name=\"CLUB_decoder_discriminator\")\n",
    "    else:\n",
    "        decoder_udiscriminator = Model(input_z, u_dircriminator(utility_decoder(input_z)), name=\"CLUB_decoder_discriminator\")\n",
    "\n",
    "    decoder_udiscriminator.compile(loss=loss_wce(-0.1), optimizer=tf.keras.optimizers.Adam(lr=learning_rate))\n",
    "    \n",
    "    return encoder, uncertainty_decoder, utility_decoder,autoencoder,z_discriminator,prior_generator,prior_gen_zdiscriminator,encoder_zdiscriminator,u_dircriminator,decoder_udiscriminator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for test\n",
    "# encoder, uncertainty_decoder, utility_decoder,autoencoder,z_discriminator,prior_generator,prior_gen_zdiscriminator,encoder_zdiscriminator,u_dircriminator,decoder_udiscriminator= get_full_model(32,0.1,0.1)\n",
    "# tf.keras.utils.plot_model(decoder_udiscriminator, expand_nested=True, show_shapes=True)\n",
    "# tf.keras.utils.plot_model(autoencoder, show_shapes=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pre-training implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pre_train_ae(dim_z, alpha, beta, force_train=False, max_itr=50, batch_size=1024, verbose=2):\n",
    "    info_str = f\"d_{dim_z}_beta_{beta}_alpha_{alpha}_{exp_info}\"\n",
    "    if not os.path.exists(f\"saved_models/mnist_pretrain_{info_str}.h5\") or force_train:\n",
    "        print(f\"Pre-Training with {info_str}\")\n",
    "        autoencoder.fit([x_train, s_train], [x_train, u_train], batch_size=batch_size, epochs=max_itr, shuffle=True, verbose=verbose)\n",
    "        autoencoder.save_weights(f\"saved_models/mnist_pretrain_{info_str}.h5\")\n",
    "    else:\n",
    "        print(f\"Loading model from file with {info_str}\")\n",
    "        autoencoder.load_weights(f\"saved_models/mnist_pretrain_{info_str}.h5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Full model training loop "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## set prior distribution function based on user config (The following cell will be ignored if the prior-generator network enabled)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.sampler import gaussian, gaussian_mixture\n",
    "\n",
    "def sample_prior(latent_dim, batch_size):\n",
    "    if prior_dist == 0:\n",
    "        return gaussian(batch_size, latent_dim)\n",
    "    elif prior_dist == 1:\n",
    "        return gaussian_mixture(batch_size, latent_dim, num_labels=min(10, DIM_U))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## customized training loop for block-wised algorithm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main_train(dim_z, max_itr=2000, batch_size=1024, z_disc_enabled=True,u_disc_enabled=True,verbose=0, dim_noise=100):\n",
    "    ones = np.ones((batch_size, 1))\n",
    "    zeros = np.zeros((batch_size, 1))\n",
    "    start_time = time.time()\n",
    "    for epoch in range(max_itr):\n",
    "        start_time_epoch = time.time()\n",
    "        # Select a random batch of images\n",
    "        idx = np.random.randint(0, x_train.shape[0], batch_size)\n",
    "        x = x_train[idx]\n",
    "        s = s_train[idx]\n",
    "        u = u_train[idx]\n",
    "\n",
    "        # ---------------------\n",
    "        #  1- Train the Encoder, Utility Decoder, Uncertainty Decoder\n",
    "        # ---------------------\n",
    "        ae_loss = autoencoder.train_on_batch([x, s], [x, u])\n",
    "\n",
    "        if z_disc_enabled:\n",
    "        # ---------------------\n",
    "        #  2- Train the Latent Space Discriminator\n",
    "        # ---------------------\n",
    "\n",
    "            if z_gen_enabled:\n",
    "                noise = np.random.uniform(-1.0, 1.0, size=[batch_size, dim_noise])\n",
    "                latent_prior = prior_generator(noise)\n",
    "            else:\n",
    "                latent_prior = sample_prior(dim_z, batch_size)\n",
    "\n",
    "            idx = np.random.randint(0, x_train.shape[0], batch_size)\n",
    "            x = x_train[idx]\n",
    "            latent_enc = encoder(x)\n",
    "\n",
    "            d_loss_prior = z_discriminator.train_on_batch(latent_prior, zeros)\n",
    "            d_loss_enc = z_discriminator.train_on_batch(latent_enc, ones)\n",
    "            dz_loss = np.add(d_loss_prior, d_loss_enc)\n",
    "\n",
    "            # ---------------------\n",
    "            # 3- Train the Encoder and Prior Distribution Generator Adversarially\n",
    "            # ---------------------\n",
    "            \n",
    "            if z_gen_enabled:\n",
    "                noise = np.random.uniform(-1.0, 1.0, size=[batch_size, dim_noise])\n",
    "                prior_loss = prior_gen_zdiscriminator.train_on_batch(noise, zeros)\n",
    "            else:\n",
    "                prior_loss = 0.\n",
    "\n",
    "            idx = np.random.randint(0, x_train.shape[0], batch_size)\n",
    "            x = x_train[idx]\n",
    "            edz_loss = encoder_zdiscriminator.train_on_batch(x, ones)\n",
    "        else:\n",
    "            dz_loss = 0.\n",
    "            prior_loss = 0.\n",
    "            edz_loss = 0.\n",
    "\n",
    "        if u_disc_enabled:\n",
    "        # ---------------------\n",
    "        #  4- Train Visible_Space/Attribute_Class Discriminator \n",
    "        # ---------------------\n",
    "\n",
    "            if z_gen_enabled:\n",
    "                noise = np.random.uniform(-1.0, 1.0, size=[batch_size, dim_noise])\n",
    "                latent = prior_generator(noise)\n",
    "            else:\n",
    "                latent = sample_prior(dim_z, batch_size)\n",
    "            u_dec = utility_decoder(latent)\n",
    "\n",
    "            idx = np.random.randint(0, x_train.shape[0], batch_size)\n",
    "            u = u_train[idx]\n",
    "\n",
    "            d_loss_real = u_dircriminator.train_on_batch(u, ones)\n",
    "            d_loss_fake = u_dircriminator.train_on_batch(u_dec, zeros)\n",
    "            du_loss = np.add(d_loss_real, d_loss_fake)\n",
    "\n",
    "        # ---------------------\n",
    "        #  5- Train the Prior Distribution Generator and Utility Decoder Adversarially\n",
    "        # ---------------------\n",
    "            if z_gen_enabled:\n",
    "                noise = np.random.uniform(-1.0, 1.0, size=[batch_size, dim_noise])\n",
    "                gdu_loss = decoder_udiscriminator.train_on_batch(noise, zeros)\n",
    "            else:    \n",
    "                latent = sample_prior(dimz, batch_size)\n",
    "                gdu_loss = decoder_udiscriminator.train_on_batch(latent, zeros)\n",
    "        else:\n",
    "            gdu_loss = 0\n",
    "            du_loss = 0\n",
    "\n",
    "        # ---------------------\n",
    "        #  Print stats info\n",
    "        # ---------------------\n",
    "        if verbose != 0 and epoch % 20 == 0:\n",
    "            print(f\"{epoch}, mse:{ae_loss[1]:.4f}, u:{ae_loss[2]:.4f}, dz:{dz_loss:.4f}, edz:{edz_loss:.4f}, prior:{prior_loss:.4f}, du:{du_loss:.4f}, gdu:{gdu_loss:.4f}\")\n",
    "            print(f\"One epoch execution time: {(time.time() - start_time_epoch):.5f} seconds\")\n",
    "\n",
    "    total_time = (time.time() - start_time) / 60\n",
    "    print(f\"Total Execution Time: {total_time:.4f} minutes\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The following function used for showing the uncertainty-decoder outputs (reconsterction) with different S values on the list of test data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.pics_tools import plot_image_grid\n",
    "def save_figures(model, dim_z, alpha, beta):\n",
    "    info_str = f\"d_{dim_z}_beta_{beta}_alpha_{alpha}_{exp_info}\"\n",
    "    n_plot_samps = 10\n",
    "    ind_list = [3, 2, 1, 18, 4 ,15, 11, 0, 128, 9]\n",
    "    if sensitive == 'color':    \n",
    "        x_hat_no_color, *_ = model.predict([x_test[ind_list], np.zeros((n_plot_samps,DIM_S)) ])\n",
    "        x_hat_all_color, *_ = model.predict([x_test[ind_list], np.ones((n_plot_samps,DIM_S)) ])\n",
    "        s0 = np.zeros((n_plot_samps,3))\n",
    "        s0[:, 0] = 1\n",
    "        s1 = np.zeros((n_plot_samps,3))\n",
    "        s1[:, 1] = 1\n",
    "        s2 = np.zeros((n_plot_samps,3))\n",
    "        s2[:, 2] = 1\n",
    "        x_hat_0, *_ = model.predict([x_test[ind_list], s0 ])\n",
    "        x_hat_1, *_ = model.predict([x_test[ind_list], s1 ])\n",
    "        x_hat_2, *_ = model.predict([x_test[ind_list], s2 ])\n",
    "\n",
    "        fig = plot_image_grid(np.concatenate([x_test[ind_list], x_hat_no_color, x_hat_all_color, x_hat_0, x_hat_1, x_hat_2], axis=0),\n",
    "          [IMG_DIM, IMG_DIM, 3], (6,n_plot_samps))\n",
    "        fig.set_size_inches(16.2, 10)\n",
    "\n",
    "    else:\n",
    "        x_hat_no_color, *_ = model.predict([x_test[ind_list], np.zeros((n_plot_samps,DIM_S)) ])\n",
    "        x_inp = x_test[ind_list]\n",
    "\n",
    "        for i in range(DIM_S):\n",
    "            s = np.zeros((n_plot_samps,DIM_S))\n",
    "            s[:, i] = 1\n",
    "            x_hat, *_ = model.predict([x_test[ind_list], s ])\n",
    "            x_inp = np.concatenate([x_inp, x_hat], axis=0)\n",
    "\n",
    "        fig = plot_image_grid( x_inp, [IMG_DIM, IMG_DIM, 3], (11, n_plot_samps))\n",
    "        fig.set_size_inches(10.6, 12)\n",
    "\n",
    "    plt.subplots_adjust(wspace=0, hspace=0)    \n",
    "    plt.savefig(f'./saved_figures_mnist/colored_mnist_{info_str}.svg', format='svg')\n",
    "    plt.savefig(f'./saved_figures_mnist/colored_mnist_{info_str}.png', format='png')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Experiment 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.mine import MINE\n",
    "alpha_list = [0.005, 0.2, 0.7, 0.9]\n",
    "beta_list = [0.0001, 0.1, 0.3, 0.5, 0.8, 0.95]\n",
    "DIM_Z = [8, 64]\n",
    "\n",
    "util_acc_tr = np.zeros((len(DIM_Z),len(alpha_list),len(beta_list)))\n",
    "util_acc_ts = np.zeros((len(DIM_Z),len(alpha_list),len(beta_list)))\n",
    "\n",
    "sens_mae_tr = np.zeros((DIM_S,len(DIM_Z),len(alpha_list),len(beta_list)))\n",
    "sens_mae_ts = np.zeros((DIM_S,len(DIM_Z),len(alpha_list),len(beta_list)))\n",
    "sens_acc_tr = np.zeros((DIM_S,len(DIM_Z),len(alpha_list),len(beta_list)))\n",
    "sens_acc_ts = np.zeros((DIM_S,len(DIM_Z),len(alpha_list),len(beta_list)))\n",
    "\n",
    "mi_u_ts = np.zeros((len(DIM_Z),len(alpha_list),len(beta_list)))\n",
    "mi_u_tr = np.zeros((len(DIM_Z),len(alpha_list),len(beta_list)))\n",
    "mi_s_ts = np.zeros((len(DIM_Z),len(alpha_list),len(beta_list)))\n",
    "mi_s_tr = np.zeros((len(DIM_Z),len(alpha_list),len(beta_list)))\n",
    "\n",
    "\n",
    "s_test_state = np.zeros((x_test.shape[0],DIM_S,DIM_S))\n",
    "s_train_state = np.zeros((x_train.shape[0],DIM_S,DIM_S))\n",
    "for i in range(DIM_S):\n",
    "    s_test_state[:, i, i] = 1\n",
    "    s_train_state[:, i, i] = 1\n",
    "\n",
    "for i, dimz in enumerate(DIM_Z):\n",
    "    for j, alpha in enumerate(alpha_list):\n",
    "        for k, beta in enumerate(beta_list):\n",
    "            #Pre training\n",
    "            encoder, uncertainty_decoder, utility_decoder,autoencoder,z_discriminator,prior_generator,prior_gen_zdiscriminator,encoder_zdiscriminator,u_dircriminator,decoder_udiscriminator = get_full_model(dimz, alpha, beta)\n",
    "            pre_train_ae(dimz, alpha, beta, force_train=False, max_itr=100, batch_size=1024, verbose=0)\n",
    "            \n",
    "            print(\"Full model training\")\n",
    "            main_train(dimz, max_itr=100, batch_size=1024, verbose=1)\n",
    "\n",
    "            print(\"Evaluate performance U\")\n",
    "            z_test = encoder.predict(x_test)\n",
    "            z_train = encoder.predict(x_train)\n",
    "            \n",
    "            u_test_hat = utility_decoder.predict(z_test)\n",
    "            u_train_hat = utility_decoder.predict(z_train)\n",
    "            \n",
    "            u_test_hat = utility_decoder.predict(z_test)\n",
    "            u_train_hat = utility_decoder.predict(z_train)\n",
    "            \n",
    "            if supervised:\n",
    "                util_acc_ts[i][j][k] = np.mean(np.argmax(u_test_hat,axis=1)==np.argmax(u_test,axis=1)) * 100\n",
    "                util_acc_tr[i][j][k] = np.mean(np.argmax(u_train_hat,axis=1)==np.argmax(u_train,axis=1)) * 100\n",
    "            else:\n",
    "                util_acc_ts[i][j][k] = np.mean(np.sum(np.square(u_test_hat-u_test), axis=(1,2,3)))\n",
    "                util_acc_tr[i][j][k] = np.mean(np.sum(np.square(u_train_hat-u_train), axis=(1,2,3)))\n",
    "            \n",
    "            save_figures(autoencoder, dimz, alpha, beta)\n",
    "            \n",
    "            print(\"Evaluate performance for S states\")\n",
    "            \n",
    "            for m in range(DIM_S):\n",
    "                x_test_hat = uncertainty_decoder.predict([z_test, s_test_state[:,:,m]])\n",
    "                s_test_hat = s_eval_model.predict(x_test_hat)\n",
    "                sens_mae_ts[m][i][j][k] = np.mean(np.abs(s_test_hat - s_test_state[:,:,m]))\n",
    "                sens_acc_ts[m][i][j][k] = np.mean(np.argmax(s_test_hat,axis=1)==m) * 100\n",
    "                \n",
    "                x_train_hat = uncertainty_decoder.predict([z_train, s_train_state[:,:,m]])\n",
    "                s_train_hat = s_eval_model.predict(x_train_hat)\n",
    "                sens_mae_tr[m][i][j][k] = np.mean(np.abs(s_train_hat - s_train_state[:,:,m]))\n",
    "                sens_acc_tr[m][i][j][k] = np.mean(np.argmax(s_train_hat,axis=1)==m) * 100\n",
    "            \n",
    "            if supervised:\n",
    "                print(\"Evaluate Mutual Information I(U,Z)\")\n",
    "                mine = MINE(x_dim=dimz, y_dim=DIM_U)\n",
    "                _, mi_u_ts[i][j][k] = mine.fit(z_test, u_test, epochs=250, batch_size=1024, verbose=0)\n",
    "                _, mi_u_tr[i][j][k] = mine.fit(z_train, u_train, epochs=250, batch_size=1024, verbose=0)\n",
    "            \n",
    "            print(\"Evaluate Mutual Information I(S,Z)\")\n",
    "            mine = MINE(x_dim=dimz, y_dim=DIM_S)\n",
    "            _, mi_s_ts[i][j][k] = mine.fit(z_test, s_test, epochs=250, batch_size=1024, verbose=0)\n",
    "            _, mi_s_tr[i][j][k] = mine.fit(z_train, s_train, epochs=250, batch_size=1024, verbose=0) \n",
    "            \n",
    "print(\"Finished!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if supervised:\n",
    "    print(np.mean(np.argmax(u_test_hat,axis=1)==np.argmax(u_test,axis=1)) * 100)\n",
    "else:\n",
    "    print(np.mean(np.sum(np.square(u_train_hat-u_train),axis=(1,2,3))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot results\n",
    "## Plot accuracy U (or MSE in unsupervised mode) (for each z we save one fig for all betas on the Test data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams.update({'font.size': 16})\n",
    "plt.rcParams['figure.dpi'] = 300\n",
    "\n",
    "linestyles = ['-', '--', '-.', ':']\n",
    "for i in range(util_acc_ts.shape[0]):\n",
    "    plt.figure(figsize=(15,8))\n",
    "#     if supervised:\n",
    "#         plt.title(r\"Utility attribute's accuracy with $\\beta$ $\\in$ \" + f'{beta_list} {exp_info}_d={DIM_Z[i]}', fontsize=16)\n",
    "#     else:\n",
    "#         plt.title(r\"Utility attribute's MSE with $\\beta$ $\\in$ \" + f'{beta_list} {exp_info}_d={DIM_Z[i]}', fontsize=16)\n",
    "    for j in range(util_acc_ts.shape[1]):\n",
    "        if supervised:\n",
    "            plt.plot(util_acc_ts[i][j], label=r'Utility Att Acc., Test, $\\alpha=$' + f\"{alpha_list[j]}\", linestyle=linestyles[0], linewidth=2)\n",
    "#             plt.plot(util_acc_tr[i][j], label=r'Utility Att Acc., Train, $\\alpha=$' + f\"{alpha_list[j]}\", linestyle=linestyles[1], linewidth=3)\n",
    "        else:\n",
    "            plt.plot(util_acc_ts[i][j], label=r'Utility Att MSE., Test, $\\alpha=$' + f\"{alpha_list[j]}\", linestyle=linestyles[0], linewidth=2)\n",
    "#             plt.plot(util_acc_ts[i][j], label=r'Utility Att MSE., Train, $\\alpha=$' + f\"{alpha_list[j]}\", linestyle=linestyles[1], linewidth=3)\n",
    "\n",
    "\n",
    "    plt.legend(prop={'size': 16})\n",
    "    plt.grid()\n",
    "    plt.xticks(list(range(len(beta_list))), beta_list, fontsize=18, rotation=90)\n",
    "    plt.xlabel(r'$\\beta$', fontsize=24)\n",
    "    if supervised:\n",
    "        plt.ylabel(r'Accuracy of the $\\mathbf{U}$', fontsize=16)\n",
    "        plt.savefig(f'./saved_figures/chart_colored_mnist_acc_d_{DIM_Z[i]}_{exp_info}.svg', format='svg', bbox_inches='tight')\n",
    "        plt.savefig(f'./saved_figures/chart_colored_mnist_acc_d_{DIM_Z[i]}_{exp_info}.png', format='png', bbox_inches='tight')\n",
    "    else:\n",
    "        plt.ylabel(r'Mean Square Error (MSE) of the $\\mathbf{U}$', fontsize=16)\n",
    "        plt.savefig(f'./saved_figures/chart_colored_mnist_mse_u_d_{DIM_Z[i]}_{exp_info}.svg', format='svg', bbox_inches='tight')\n",
    "        plt.savefig(f'./saved_figures/chart_colored_mnist_mse_u_d_{DIM_Z[i]}_{exp_info}.png', format='png', bbox_inches='tight')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot accuracy U (for each z and alpha we save and plot a figure on the Test and Train data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams.update({'font.size': 16})\n",
    "plt.rcParams['figure.dpi'] = 300\n",
    "\n",
    "linestyles = ['-', '--', '-.', ':']\n",
    "for i in range(util_acc_ts.shape[0]):\n",
    "    plt.figure(figsize=(15,8))\n",
    "    for j in range(util_acc_ts.shape[1]):\n",
    "        if supervised:\n",
    "#             plt.title(r\"Utility attribute's accuracy with $\\beta$ $\\in$ \" + f'{beta_list} {exp_info}_d={DIM_Z[i]}_alpha={alpha_list[j]}', fontsize=16)\n",
    "            plt.plot(util_acc_ts[i][j], label=r'Utility Att Acc., Test $\\alpha=$' + f\"{alpha_list[j]}\", linestyle=linestyles[0], linewidth=2)\n",
    "            plt.plot(util_acc_tr[i][j], label=r'Utility Att Acc., Train $\\alpha=$' + f\"{alpha_list[j]}\", linestyle=linestyles[1], linewidth=3)\n",
    "        else:\n",
    "#             plt.title(r\"Utility attribute's MSE with $\\beta$ $\\in$ \" + f'{beta_list} {exp_info}_d={DIM_Z[i]}_alpha={alpha_list[j]}', fontsize=16)\n",
    "            plt.plot(util_acc_ts[i][j], label=r'Utility Att MSE., Test $\\alpha=$' + f\"{alpha_list[j]}\", linestyle=linestyles[0], linewidth=2)\n",
    "            plt.plot(util_acc_tr[i][j], label=r'Utility Att MSE., Train $\\alpha=$' + f\"{alpha_list[j]}\", linestyle=linestyles[1], linewidth=3)\n",
    "\n",
    "\n",
    "        plt.legend(prop={'size': 16})\n",
    "        plt.grid()\n",
    "        plt.xticks(list(range(len(beta_list))), beta_list, fontsize=18, rotation=90)\n",
    "\n",
    "        plt.xlabel(r'$\\beta$', fontsize=24)\n",
    "        if supervised:\n",
    "            plt.ylabel(r'Accuracy of the $\\mathbf{U}$', fontsize=16)\n",
    "            plt.savefig(f'./saved_figures/chart_colored_mnist_acc_d_{DIM_Z[i]}_alpha_{alpha_list[j]}_{exp_info}.svg', format='svg', bbox_inches='tight')\n",
    "            plt.savefig(f'./saved_figures/chart_colored_mnist_acc_d_{DIM_Z[i]}_alpha_{alpha_list[j]}_{exp_info}.png', format='png', bbox_inches='tight')\n",
    "        else:\n",
    "            plt.ylabel(r'Mean Square Error (MSE) of the $\\mathbf{U}$', fontsize=16)\n",
    "            plt.savefig(f'./saved_figures/chart_colored_mnist_mse_u_d_{DIM_Z[i]}_alpha_{alpha_list[j]}_{exp_info}.svg', format='svg', bbox_inches='tight')\n",
    "            plt.savefig(f'./saved_figures/chart_colored_mnist_mse_u_d_{DIM_Z[i]}_alpha_{alpha_list[j]}_{exp_info}.png', format='png', bbox_inches='tight')\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot accuracy S"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams.update({'font.size': 16})\n",
    "plt.rcParams['figure.dpi'] = 300\n",
    "\n",
    "linestyles = ['-', '--', '-.', ':']\n",
    "for i in range(sens_acc_ts.shape[1]):\n",
    "    plt.figure(figsize=(15,8))\n",
    "#     plt.title(r'Sensitive attribute accuracy with $\\beta$ $\\in$ ' + f'{beta_list} {exp_info}_d={DIM_Z[i]}', fontsize=16)\n",
    "    for j in range(sens_mae_ts.shape[2]):\n",
    "        for k in range(DIM_S):\n",
    "            plt.plot(sens_acc_ts[k][i][j], label=r'Sensitive Att Acc., Test, $\\alpha=$' + f'{alpha_list[j]}, s = {k}', linestyle=linestyles[0], linewidth=2)\n",
    "#             plt.plot(sens_acc_tr[k][i][j], label=r'Sensiitive Att Acc., Train, $\\alpha=$' + f'{alpha_list[j]}, s = {k}', linestyle=linestyles[1], linewidth=3)\n",
    "\n",
    "\n",
    "    plt.legend(prop={'size': 16})\n",
    "    plt.grid()\n",
    "    plt.xticks(list(range(len(beta_list))), beta_list, fontsize=18, rotation=90)\n",
    "    plt.xlabel(r'$\\beta$', fontsize=24)\n",
    "    plt.ylabel(r'Accuracy on the $\\mathbf{S}$', fontsize=16)\n",
    "    plt.savefig(f'./saved_figures/chart_colored_mnist_acc_s_d_{DIM_Z[i]}_{exp_info}.svg', format='svg', bbox_inches='tight')\n",
    "    plt.savefig(f'./saved_figures/chart_colored_mnist_acc_s_d_{DIM_Z[i]}_{exp_info}.png', format='png', bbox_inches='tight')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot Mean absolute error (MAE) on all sensitivity states"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams.update({'font.size': 16})\n",
    "plt.rcParams['figure.dpi'] = 300\n",
    "\n",
    "linestyles = ['-', '--', '-.', ':']\n",
    "for i in range(sens_mae_ts.shape[1]):\n",
    "    plt.figure(figsize=(15,8))\n",
    "#     plt.title(r'Sensitive attribute MAE with $\\beta$ $\\in$ ' + f'{beta_list} {exp_info}_d={DIM_Z[i]}', fontsize=16)\n",
    "    for j in range(sens_mae_ts.shape[2]):\n",
    "        for k in range(DIM_S):\n",
    "            plt.plot(sens_mae_ts[k][i][j], label=r'Sensitive Att MAE., Test, $\\alpha=$' + f'{alpha_list[j]}, s = {k}', linestyle=linestyles[0], linewidth=2)\n",
    "\n",
    "\n",
    "    plt.legend(prop={'size': 16})\n",
    "    plt.grid()\n",
    "    plt.xticks(list(range(len(beta_list))), beta_list, fontsize=18, rotation=90)\n",
    "    plt.xlabel(r'$\\beta$', fontsize=24)\n",
    "    plt.ylabel(r'Mean Absolute Error (MAE) of the $\\mathbf{S}$', fontsize=16)\n",
    "    plt.savefig(f'./saved_figures/chart_colored_mnist_mae_d_{DIM_Z[i]}_{exp_info}.svg', format='svg', bbox_inches='tight')\n",
    "    plt.savefig(f'./saved_figures/chart_colored_mnist_mae_d_{DIM_Z[i]}_{exp_info}.png', format='png', bbox_inches='tight')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot mutual information between Z and S"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams.update({'font.size': 16})\n",
    "plt.rcParams['figure.dpi'] = 300\n",
    "plt.figure(figsize=(15,8))\n",
    "# plt.title(r'Mutual Information with $\\beta$ $\\in$ ' + f'{beta_list}_{exp_info}', fontsize=20)\n",
    "linestyles = ['-', '--', '-.', ':']\n",
    "colors = ['b', 'r', 'g', 'c']\n",
    "for i in range(mi_s_ts.shape[0]):\n",
    "    for j in range(mi_s_ts.shape[1]):\n",
    "        plt.plot(mi_s_ts[i][j], label=r'Test, $d_z=$' + f'{DIM_Z[i]}' + r', $\\alpha=$' + f\"{alpha_list[j]}\", linestyle=linestyles[0] ,linewidth=3)\n",
    "#         plt.plot(mi_s_tr[i][j], label=r'Train, $d_z=$' + f'{DIM_Z[i]}' + r', $\\alpha=$' + f\"{alpha_list[j]}\", linestyle=linestyles[1], c=colors[j%len(colors)] ,linewidth=3)\n",
    "        \n",
    "# plt.xscale('linear')\n",
    "plt.xlabel(r'$\\beta$', fontsize=24)\n",
    "plt.xticks(list(range(len(beta_list))), beta_list, rotation=90)\n",
    "\n",
    "plt.ylabel(r'I($\\mathbf{S}$;$\\mathbf{Z}$)', fontsize=24)\n",
    "plt.legend(prop={'size': 16})\n",
    "plt.grid()\n",
    "plt.savefig(f'./saved_figures/chart_colored_mnist_mi_zs_{exp_info}.svg', format='svg', bbox_inches='tight')\n",
    "plt.savefig(f'./saved_figures/chart_colored_mnist_mi_zs_{exp_info}.png', format='png', bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot mutual information between Z and U"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if supervised:\n",
    "    plt.rcParams.update({'font.size': 16})\n",
    "    plt.rcParams['figure.dpi'] = 300\n",
    "    plt.figure(figsize=(15,8))\n",
    "    plt.title(r'Mutual Information with $\\beta$ $\\in$ ' + f'{beta_list}_{exp_info}', fontsize=20)\n",
    "    linestyles = ['-', '--', '-.', ':']\n",
    "    colors = ['b', 'r', 'g', 'c']\n",
    "    for i in range(mi_u_ts.shape[0]):\n",
    "        for j in range(mi_u_ts.shape[1]):\n",
    "            plt.plot(mi_u_ts[i][j], label=r'Test, $d_z=$' + f'{DIM_Z[i]}' + r', $\\alpha=$' + f\"{alpha_list[j]}\", linestyle=linestyles[0] ,linewidth=3)\n",
    "    #         plt.plot(mi_u_tr[i][j], label=r'Train, $d_z=$' + f'{DIM_Z[i]}' + r', $\\alpha=$' + f\"{alpha_list[j]}\", linestyle=linestyles[1], c=colors[j%len(colors)] ,linewidth=3)\n",
    "\n",
    "    # plt.xscale('linear')\n",
    "    plt.xlabel(r'$\\beta$', fontsize=24)\n",
    "    plt.xticks(list(range(len(beta_list))), beta_list, rotation=90)\n",
    "\n",
    "    plt.ylabel(r'I($\\mathbf{U}$;$\\mathbf{Z}$)', fontsize=24)\n",
    "    plt.legend(prop={'size': 16})\n",
    "    plt.grid()\n",
    "    plt.savefig(f'./saved_figures/chart_colored_mnist_mi_zu_{exp_info}.svg', format='svg', bbox_inches='tight')\n",
    "    plt.savefig(f'./saved_figures/chart_colored_mnist_mi_zu_{exp_info}.png', format='png', bbox_inches='tight')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, u_test_hat = autoencoder.predict([x_test, s_test])\n",
    "if supervised:\n",
    "    print(\"Acc = \",np.mean(np.argmax(u_test_hat,axis=1)==np.argmax(u_test,axis=1)) * 100)\n",
    "else:\n",
    "    print(\"MSE = \", np.mean(np.sum(np.square(u_train_hat-u_train), axis=(1,2,3))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Save results values to file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.io as sio\n",
    "sio.savemat(f'./saved_data/acc_colored_mnist_{exp_info}.mat', {'util_acc_ts':util_acc_ts, 'util_acc_tr':util_acc_tr, 'sens_acc_ts':sens_acc_ts, 'sens_acc_tr':sens_acc_tr, 'sens_mae_ts':sens_mae_ts, 'sens_mae_tr':sens_mae_tr})\n",
    "sio.savemat(f'./saved_data/mi_colored_mnist_{exp_info}.mat', {'mi_s_ts':mi_s_ts, 'mi_s_tr':mi_s_tr, 'mi_u_ts':mi_u_ts, 'mi_u_tr':mi_u_tr})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load results from file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.io as sio\n",
    "mat_contents  = sio.loadmat(f'./saved_data/acc_colored_mnist_{exp_info}.mat')\n",
    "util_acc_ts = mat_contents['util_acc_ts']\n",
    "util_acc_tr = mat_contents['util_acc_tr']\n",
    "sens_acc_ts = mat_contents['sens_acc_ts']\n",
    "sens_acc_tr = mat_contents['sens_acc_tr']\n",
    "sens_mae_ts = mat_contents['sens_mae_ts']\n",
    "sens_mae_tr = mat_contents['sens_mae_tr']\n",
    "\n",
    "\n",
    "mat_contents  = sio.loadmat(f'./saved_data/mi_colored_mnist_{exp_info}.mat')\n",
    "mi_s_tr = mat_contents['mi_s_tr']\n",
    "mi_s_ts = mat_contents['mi_s_ts']\n",
    "mi_u_ts = mat_contents['mi_u_ts']\n",
    "mi_u_tr = mat_contents['mi_u_tr']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha_list = [0.005, 0.2, 0.7, 0.9]\n",
    "beta_list = [0.0001, 0.1, 0.3, 0.5, 0.8, 0.95]\n",
    "DIM_Z = [8, 64]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"test\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:tensorflow-24]",
   "language": "python",
   "name": "conda-env-tensorflow-24-py"
  },
  "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}