{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Latent MNIST data experiments\n",
    "This file contains experiments for the latent MNIST data set. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn, optim\n",
    "import torch.nn.functional as F\n",
    "import torchvision\n",
    "import torch.distributions as dist\n",
    "from scipy.stats import norm\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from GANmodel import *\n",
    "from lineardatatools import p_y_zt_from_true_dist, linear_binary_ty_pydot\n",
    "from matplotlib.lines import Line2D\n",
    "from imagedata import *\n",
    "from imagedatatools import *\n",
    "from plot_functions import *\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "netG = Generator(3)\n",
    "netD = Discriminator()\n",
    "netG.load_state_dict(torch.load(\"./GANmodels/3dGenerator\"))\n",
    "netD.load_state_dict(torch.load(\"./GANmodels/3dDiscriminator\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 7628 characters]",
      "text/plain": [
       "<Figure size 400x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#Example image from the data generating distribution\n",
    "z = torch.randn(1,3)[:,:,None,None]\n",
    "im =((1+netG(z))/2).squeeze()\n",
    "fig, ax = plt.subplots(1,1, figsize=(4,4))\n",
    "plt.imshow(dist.Bernoulli(im).sample())\n",
    "plt.grid(False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate data for loss scaling tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "main_folder = \"imagedata\"\n",
    "sub_folder_data = \"samplesize100000\"\n",
    "\n",
    "data_labels = [\"data100000\"]\n",
    "\n",
    "generator = netG\n",
    "c_x = np.array([1])\n",
    "s_x = np.array([1])\n",
    "t_a = np.array([1])\n",
    "t_b = np.array([0.5])\n",
    "y_a0 = np.array([2])\n",
    "y_b0 = np.array([-0.5])\n",
    "y_a1 = np.array([2])\n",
    "y_b1 = np.array([0.5])\n",
    "param_times = 1\n",
    "gen_zdim=3\n",
    "dataparameters = [[100000, gen_zdim, generator, t_a, t_b, y_a0, y_b0, y_a1, y_b1, c_x, s_x]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step  0\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "({'data100000': {0:              z0        z1        z2        x0    t         y\n",
       "   0     -0.033588 -0.056642 -0.032322  2.496197  1.0  1.092581\n",
       "   1     -0.153361 -1.836679  0.670796 -0.384910  0.0 -0.873604\n",
       "   2      0.331998 -0.223278 -1.932001  0.888722  1.0  1.218269\n",
       "   3     -0.103416 -0.003162 -1.015926  1.221512  1.0 -0.985543\n",
       "   4     -0.790482  3.427965 -0.962994 -1.103791  1.0 -0.549818\n",
       "   ...         ...       ...       ...       ...  ...       ...\n",
       "   99995  1.129985  0.919266 -1.183493  1.170293  1.0  3.258621\n",
       "   99996 -0.208778  0.300943  1.144392  0.508172  1.0  0.002202\n",
       "   99997  0.219733 -2.292035 -1.646471  0.912890  1.0  0.688406\n",
       "   99998  0.201335 -0.870611  1.028542 -0.494065  1.0  1.106175\n",
       "   99999 -1.785558  1.503549 -0.767202 -1.003225  0.0  1.389041\n",
       "   \n",
       "   [100000 rows x 6 columns]}},\n",
       " {'data100000': {0: <imagedata.ImageDataset at 0x107c75a90>}})"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "create_dfs_datasets(generate_image_data, dataparameters, param_times, repeat=True, main_folder=main_folder,\n",
    "                    sub_folder=sub_folder_data, labels=data_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load the generated data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "main_folder = \"imagedata\"\n",
    "sub_folder_data = \"samplesize100000\"\n",
    "param_times = 1\n",
    "data_labels = [\"data100000\"]\n",
    "\n",
    "dfs, datasets = load_dfs_datasets(main_folder, sub_folder_data, param_times)\n",
    "data = dfs['data100000'][0]\n",
    "imageDataset = datasets['data100000'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# remove information from the linear-gaussian proxy X\n",
    "imageDataset.X = torch.ones(size=imageDataset.X.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = data['z0']\n",
    "t = data['t']\n",
    "y = data['y']\n",
    "y_t0 = data.loc[data['t']==0, 'y']\n",
    "y_t1 = data.loc[data['t']==1, 'y']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 39572 characters]",
      "text/plain": [
       "<Figure size 1600x400 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# show distribution of the dataset\n",
    "fig, ax = plt.subplots(1,3, figsize=(16,4),gridspec_kw={\"width_ratios\":[1,3,1]})\n",
    "\n",
    "\n",
    "ax[0].hist(z, bins=50, density=True)\n",
    "ax[0].set_title(\"$z$ distribution\", fontsize=18)\n",
    "ax[0].set_xlabel('$z$', fontsize=16)\n",
    "ax[0].set_ylabel(\"probability\", fontsize=16)\n",
    "ax[0].tick_params(which='major', axis='both', labelsize=16)\n",
    "\n",
    "ax[1].hist(y, bins=30, density=True)\n",
    "ax[1].set_title(\"$y$ distribution\", fontsize=18)\n",
    "ax[1].set_xlabel('$y$', fontsize=16)\n",
    "ax[1].set_ylabel(\"probability\", fontsize=16)\n",
    "ax[1].tick_params(which='major', axis='both', labelsize=16)\n",
    "\n",
    "ax[2].hist(t, bins=30, density=True)\n",
    "ax[2].set_title(\"$t$ distribution\", fontsize=18)\n",
    "ax[2].set_xlabel('$t$', fontsize=16)\n",
    "ax[2].set_ylabel(\"probability\", fontsize=16)\n",
    "ax[2].tick_params(which='major', axis='both', labelsize=16)\n",
    "\n",
    "for i in range(3):\n",
    "    ax[i].spines['right'].set_visible(False)\n",
    "    ax[i].spines['top'].set_visible(False)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "there are 40095 samles t=0\n",
      "there are 59905 samles t=1\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 27884 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# show distribution of p(y|t)\n",
    "print(f\"there are {len(y_t0)} samles t=0\")\n",
    "print(f\"there are {len(y_t1)} samles t=1\")\n",
    "plt.hist(y_t0, bins=30, alpha=0.5, label='p(y|t=0)', density=True)\n",
    "plt.hist(y_t1, bins=30, alpha=0.5, label='p(y|t=1)', density=True)\n",
    "plt.legend(loc='upper right')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Devide the dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Split to two Gaussian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set the distribution for source domain and the target domain\n",
    "dist_z_source = [-1, 1]\n",
    "dist_z_target = [1, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 32296 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "\n",
    "#mu, std = norm.fit(sorted_data['z']) \n",
    "x = np.arange(-3, 3, 0.1)\n",
    "p = norm.pdf(x, loc=0, scale=1)\n",
    "\n",
    "plt.hist(data['z0'], bins=30, density=True, label='P(z)')\n",
    "plt.plot(x, p, label='N(0,1)')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# target z distribution\n",
    "def p_source(x):\n",
    "    return norm.pdf(x, loc=dist_z_source[0], scale=dist_z_source[1])\n",
    "\n",
    "def p_target(x):\n",
    "    return norm.pdf(x, loc=dist_z_target[0], scale=dist_z_target[1])\n",
    "    \n",
    "# z distribution in the dataset\n",
    "def q(x):\n",
    "    return norm.pdf(x, loc=0, scale=1)\n",
    "\n",
    "def rejection_sampling(sample, p, size):\n",
    "    \n",
    "    x = np.arange(-3, 3, 0.1)\n",
    "    k = max(p(x)/q(x))\n",
    "    c = np.random.uniform(0, 1, size=size)\n",
    "    mask = p(sample) / (k * q(sample)) > c\n",
    "    return mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "# rejection sampling\n",
    "sample = np.array(imageDataset.z[:,0].squeeze())\n",
    "mask_source = rejection_sampling(sample, p_source, len(data))\n",
    "source_data = imageDataset.apply_mask(mask_source)\n",
    "\n",
    "mask_target = rejection_sampling(sample, p_target, len(data))\n",
    "target_data = imageDataset.apply_mask(mask_target)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "there are 8084 samles in source domain\n",
      "there are 8879 samles in target domain\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 57700 characters]",
      "text/plain": [
       "<Figure size 1600x400 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<Figure size 640x480 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "print(f\"there are {len(source_data)} samles in source domain\")\n",
    "print(f\"there are {len(target_data)} samles in target domain\")\n",
    "\n",
    "#plt.legend(loc='upper right')\n",
    "#plt.show()\n",
    "\n",
    "plot_souce_target(source_data, target_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cut it by half"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {},
   "outputs": [],
   "source": [
    "sorted_ImageDataset = imageDataset.sort_by_z()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {},
   "outputs": [],
   "source": [
    "dist_z_source = [\"-\"]\n",
    "dist_z_target = [\"+\"]\n",
    "\n",
    "length = sorted_ImageDataset.length\n",
    "\n",
    "if dist_z_source[0] == \"-\":\n",
    "    source_data = ImageDataset(sorted_ImageDataset.images[0: int(length/2)], sorted_ImageDataset.X[0: int(length/2)], sorted_ImageDataset.t[0: int(length/2)], sorted_ImageDataset.y[0: int(length/2)], sorted_ImageDataset.z[0: int(length/2)])\n",
    "\n",
    "\n",
    "    target_data = ImageDataset(sorted_ImageDataset.images[int(length/2):], sorted_ImageDataset.X[int(length/2):], sorted_ImageDataset.t[int(length/2):], sorted_ImageDataset.y[int(length/2):], sorted_ImageDataset.z[int(length/2):])\n",
    "else:\n",
    "    target_data = ImageDataset(sorted_ImageDataset.images[0: int(length/2)], sorted_ImageDataset.X[0: int(length/2)], sorted_ImageDataset.t[0: int(length/2)], sorted_ImageDataset.y[0: int(length/2)], sorted_ImageDataset.z[0: int(length/2)])\n",
    "\n",
    "\n",
    "    source_data = ImageDataset(sorted_ImageDataset.images[int(length/2):], sorted_ImageDataset.X[int(length/2):], sorted_ImageDataset.t[int(length/2):], sorted_ImageDataset.y[int(length/2):], sorted_ImageDataset.z[int(length/2):])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "there are 50000 samles in source domain\n",
      "there are 50000 samles in target domain\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 57476 characters]",
      "text/plain": [
       "<Figure size 1600x400 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "print(f\"there are {len(source_data)} samles in source domain\")\n",
    "print(f\"there are {len(target_data)} samles in target domain\")\n",
    "\n",
    "#plt.legend(loc='upper right')\n",
    "#plt.show()\n",
    "\n",
    "plot_souce_target(source_data, target_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Devide it by digits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train on the source domain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.utils\n",
    "from models import trainZTtoYmodel, trainXTtoYmodel, trainZtoXmodel, GenericDataset, RMSE_loss\n",
    "\n",
    "number_trained_sample = 1000\n",
    "source_data = source_data.sample(number_trained_sample)\n",
    "#target_data = target_data.sample(number_trained_sample)\n",
    "device = 'mps'\n",
    "source_dataframe = source_data.dataset_to_dataframe()\n",
    "target_dataframe = target_data.dataset_to_dataframe()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pretrain a decoder for $p(w \\vert z)$ and $p(y\\vert z, x)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# can be ignored as we have pretrained decoder for N(-1,1), N(1,1), \n",
    "# negative trucated Gaussian and postive trucated Gaussian in \n",
    "# model/imagedata\n",
    "plot_curves=True\n",
    "device = 'mps'\n",
    "\n",
    "z_dim = 3\n",
    "x_dim = 1\n",
    "param_times = 2\n",
    "\n",
    "p_y_zt_nn=True\n",
    "p_y_zt_nn_layers=3\n",
    "p_y_zt_nn_width=30\n",
    "p_t_z_nn=False\n",
    "p_t_z_nn_layers=3\n",
    "p_t_z_nn_width=30\n",
    "p_x_z_nn=True\n",
    "p_x_z_nn_layers=3\n",
    "p_x_z_nn_width=30\n",
    "\n",
    "num_epochs = 500\n",
    "lr_start = 0.001\n",
    "lr_end = 0.0005\n",
    "separate_ty = True\n",
    "BATCH_SIZE = 100\n",
    "\n",
    "model = ImageCEVAE(x_dim, z_dim, device=device, p_y_zt_nn=p_y_zt_nn, p_y_zt_nn_layers=p_y_zt_nn_layers,\n",
    "            p_y_zt_nn_width=p_y_zt_nn_width, p_t_z_nn=p_t_z_nn, p_t_z_nn_layers=p_t_z_nn_layers, p_t_z_nn_width=p_t_z_nn_width,\n",
    "            p_x_z_nn=p_x_z_nn, p_x_z_nn_layers=p_x_z_nn_layers, p_x_z_nn_width=p_x_z_nn_width,separate_ty=separate_ty)\n",
    "source_dataloader = DataLoader(source_data, batch_size=BATCH_SIZE)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0:\n",
      "Image: 5262.101379394531, x: 19.528533577919006, t: 0, y: 19.46196734905243\n",
      "\n",
      "Epoch 50:\n",
      "Image: 1415.2007141113281, x: 18.98470163345337, t: 0, y: 15.275440335273743\n",
      "\n",
      "Epoch 100:\n",
      "Image: 932.5657196044922, x: 18.671436309814453, t: 0, y: 14.745296359062195\n",
      "\n",
      "Epoch 150:\n",
      "Image: 706.3411712646484, x: 18.249229431152344, t: 0, y: 14.660892605781555\n",
      "\n",
      "Epoch 200:\n",
      "Image: 620.1294021606445, x: 17.707452178001404, t: 0, y: 14.610298037528992\n",
      "\n",
      "Epoch 250:\n",
      "Image: 565.1984710693359, x: 17.03713846206665, t: 0, y: 14.576022028923035\n",
      "\n",
      "Epoch 300:\n",
      "Image: 526.5389671325684, x: 16.225585341453552, t: 0, y: 14.54378616809845\n",
      "\n",
      "Epoch 350:\n",
      "Image: 490.6045608520508, x: 15.259230971336365, t: 0, y: 14.509616255760193\n",
      "\n",
      "Epoch 400:\n",
      "Image: 466.0627861022949, x: 14.127265334129333, t: 0, y: 14.471090197563171\n",
      "\n",
      "Epoch 450:\n",
      "Image: 422.1161117553711, x: 12.826404213905334, t: 0, y: 14.429858326911926\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 80224 characters]",
      "text/plain": [
       "<Figure size 1200x800 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 16908 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "decoder = train_decoder(device, model, True, source_dataloader, num_epochs, lr_start, lr_end)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "if p_t_z_nn:\n",
    "    t_net = 'withTnet'\n",
    "else:\n",
    "    t_net = 'withoutTnet'\n",
    "torch.save(decoder.state_dict(), \"./model/{}/{}_NoX_decoder_{}epoches_{}samples_{}\".format(main_folder,dist_z_source[0],num_epochs,number_trained_sample, t_net))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load a pre-trained decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading from ./model/imagedata/-1_X_decoder_500epoches_1000samples_withoutTnet\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "if p_t_z_nn:\n",
    "    t_net = 'withTnet'\n",
    "else:\n",
    "    t_net = 'withoutTnet'\n",
    "transDe_CEVAE = ImageCEVAE(x_dim, z_dim, device=device, p_y_zt_nn=p_y_zt_nn, p_y_zt_nn_layers=p_y_zt_nn_layers,\n",
    "            p_y_zt_nn_width=p_y_zt_nn_width, p_t_z_nn=p_t_z_nn, p_t_z_nn_layers=p_t_z_nn_layers, p_t_z_nn_width=p_t_z_nn_width,\n",
    "            p_x_z_nn=p_x_z_nn, p_x_z_nn_layers=p_x_z_nn_layers, p_x_z_nn_width=p_x_z_nn_width,separate_ty=separate_ty)\n",
    "print(\"Loading from \" + \"./model/{}/{}_X_decoder_{}epoches_{}samples_{}\".format(main_folder,dist_z_source[0],num_epochs,number_trained_sample, t_net))\n",
    "\n",
    "transDe_CEVAE.decoder.load_state_dict(torch.load(\"./model/{}/{}_X_decoder_{}epoches_{}samples_{}\".format(main_folder,dist_z_source[0],num_epochs,number_trained_sample, t_net)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1.9330, 0.0408, 1.6731]], device='mps:0')\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 7604 characters]",
      "text/plain": [
       "<Figure size 400x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "\n",
    "# test the decoder\n",
    "#Example image from the target data distribution\n",
    "z = torch.normal(1,1,size=(1,3)).to(device)\n",
    "print(z)\n",
    "image,_,_,_,_,_,_ = transDe_CEVAE.decoder(z, t=1)\n",
    "out = image.squeeze()\n",
    "fig, ax = plt.subplots(1,1, figsize=(4,4))\n",
    "plt.imshow(dist.Bernoulli(logits=out).sample().detach().cpu())\n",
    "plt.grid(False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train a VAE using source domain data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n",
      "Epoch 0:\n",
      "Image: 4716.683685302734, x: 19.49955379962921, t: 0, y: 19.264492511749268\n",
      "Epoch loss at epoch 0: 4765.4300537109375\n",
      "Image: 4716.683685302734, x: 19.49955379962921, t: 0, y: 19.264492511749268\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 115964 characters]",
      "text/plain": [
       "<Figure size 1500x500 with 10 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 20:\n",
      "Image: 1756.170913696289, x: 17.452929735183716, t: 0, y: 15.31073796749115\n",
      "Epoch 40:\n",
      "Image: 1566.501693725586, x: 15.838437557220459, t: 0, y: 15.053199529647827\n",
      "Epoch loss at epoch 50: 1564.4369049072266\n",
      "Image: 1452.3118896484375, x: 15.064437627792358, t: 0, y: 14.998063921928406\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 87016 characters]",
      "text/plain": [
       "<Figure size 1500x500 with 10 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 60:\n",
      "Image: 1407.2005996704102, x: 14.30426299571991, t: 0, y: 14.981924176216125\n",
      "Epoch 80:\n",
      "Image: 1292.316261291504, x: 12.807417869567871, t: 0, y: 14.887378692626953\n",
      "Epoch 100:\n",
      "Image: 1249.5983581542969, x: 11.336641907691956, t: 0, y: 14.82059121131897\n",
      "Epoch loss at epoch 100: 1371.9647216796875\n",
      "Image: 1249.5983581542969, x: 11.336641907691956, t: 0, y: 14.82059121131897\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 93784 characters]",
      "text/plain": [
       "<Figure size 1500x500 with 10 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 120:\n",
      "Image: 1213.1976623535156, x: 9.92555570602417, t: 0, y: 14.783013105392456\n",
      "Epoch 140:\n",
      "Image: 1121.6950073242188, x: 8.545061826705933, t: 0, y: 14.67937695980072\n",
      "Epoch loss at epoch 150: 1301.880958557129\n",
      "Image: 1184.684341430664, x: 7.841121792793274, t: 0, y: 14.692275881767273\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 84116 characters]",
      "text/plain": [
       "<Figure size 1500x500 with 10 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 160:\n",
      "Image: 1089.1090774536133, x: 7.151559233665466, t: 0, y: 14.589766025543213\n",
      "Epoch 180:\n",
      "Image: 1169.0448608398438, x: 5.872901260852814, t: 0, y: 14.519596576690674\n",
      "Epoch 200:\n",
      "Image: 1033.6638793945312, x: 4.552726536989212, t: 0, y: 14.45666754245758\n",
      "Epoch loss at epoch 200: 1150.0789337158203\n",
      "Image: 1033.6638793945312, x: 4.552726536989212, t: 0, y: 14.45666754245758\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 93088 characters]",
      "text/plain": [
       "<Figure size 1500x500 with 10 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 220:\n",
      "Image: 1013.4090118408203, x: 3.2800978124141693, t: 0, y: 14.401973247528076\n",
      "Epoch 240:\n",
      "Image: 971.6020584106445, x: 2.0619224905967712, t: 0, y: 14.26071572303772\n",
      "Epoch loss at epoch 250: 1094.665641784668\n",
      "Image: 978.1935577392578, x: 1.4423059523105621, t: 0, y: 14.280243515968323\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 81644 characters]",
      "text/plain": [
       "<Figure size 1500x500 with 10 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 260:\n",
      "Image: 900.7694244384766, x: 0.8566272333264351, t: 0, y: 14.181411027908325\n",
      "Epoch 280:\n",
      "Image: 905.3857421875, x: -0.2789800576865673, t: 0, y: 14.126026153564453\n",
      "Epoch 300:\n",
      "Image: 925.5479354858398, x: -1.4799694269895554, t: 0, y: 14.214563488960266\n",
      "Epoch loss at epoch 300: 1037.771957397461\n",
      "Image: 925.5479354858398, x: -1.4799694269895554, t: 0, y: 14.214563488960266\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 84064 characters]",
      "text/plain": [
       "<Figure size 1500x500 with 10 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 320:\n",
      "Image: 823.5486221313477, x: -2.5692599713802338, t: 0, y: 13.997759461402893\n",
      "Epoch 340:\n",
      "Image: 1071.7869415283203, x: -3.652268886566162, t: 0, y: 14.13176965713501\n",
      "Epoch loss at epoch 350: 933.1306762695312\n",
      "Image: 820.9145812988281, x: -4.225954532623291, t: 0, y: 14.000043869018555\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 74840 characters]",
      "text/plain": [
       "<Figure size 1500x500 with 10 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 360:\n",
      "Image: 779.077033996582, x: -4.730424582958221, t: 0, y: 13.891936898231506\n",
      "Epoch 380:\n",
      "Image: 839.3865127563477, x: -5.752850115299225, t: 0, y: 13.921246647834778\n",
      "Epoch 400:\n",
      "Image: 756.6401443481445, x: -6.755788326263428, t: 0, y: 13.818090796470642\n",
      "Epoch loss at epoch 400: 869.028434753418\n",
      "Image: 756.6401443481445, x: -6.755788326263428, t: 0, y: 13.818090796470642\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 75408 characters]",
      "text/plain": [
       "<Figure size 1500x500 with 10 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 420:\n",
      "Image: 738.0509300231934, x: -7.69605940580368, t: 0, y: 13.831464529037476\n",
      "Epoch 440:\n",
      "Image: 736.6894454956055, x: -8.761474788188934, t: 0, y: 13.78146481513977\n",
      "Epoch loss at epoch 450: 796.8942489624023\n",
      "Image: 686.0911178588867, x: -9.205161571502686, t: 0, y: 13.770880222320557\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 79444 characters]",
      "text/plain": [
       "<Figure size 1500x500 with 10 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 460:\n",
      "Image: 672.105842590332, x: -9.631701231002808, t: 0, y: 13.726381659507751\n",
      "Epoch 480:\n",
      "Image: 650.6860656738281, x: -10.469807982444763, t: 0, y: 13.654858708381653\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 106664 characters]",
      "text/plain": [
       "<Figure size 1200x800 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 16908 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# can be ignored as we have pretrained VAE for N(-1,1), N(1,1), \n",
    "# negative trucated Gaussian and postive trucated Gaussian in \n",
    "# model/imagedata\n",
    "source_CEVAE,_ = train_model(device, plot_curves, print_logs=True, train_loader=source_dataloader, num_epochs=num_epochs, lr_start=0.001, lr_end=lr_end, x_dim=x_dim, z_dim=z_dim, p_y_zt_nn=p_y_zt_nn, p_y_zt_nn_layers=p_y_zt_nn_layers, p_y_zt_nn_width=p_y_zt_nn_width, p_t_z_nn=p_t_z_nn, p_t_z_nn_layers=p_t_z_nn_layers, p_t_z_nn_width=p_t_z_nn_width, p_x_z_nn=p_x_z_nn, p_x_z_nn_layers=p_x_z_nn_layers, p_x_z_nn_width=p_x_z_nn_width, loss_scaling=1, separate_ty=separate_ty)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "if p_t_z_nn:\n",
    "    t_net = 'withTnet'\n",
    "else:\n",
    "    t_net = 'withoutTnet'\n",
    "torch.save(source_CEVAE.state_dict(), \"./model/{}/{}_NoX_sourceCEVAE_{}epochs_{}samples_{}\".format(main_folder,dist_z_source[0],num_epochs,number_trained_sample, t_net))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load a pre-trained CEVAE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading from ./model/imagedata/-1_X_sourceCEVAE_500epoches_1000samples_withoutTnet\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "source_CEVAE = ImageCEVAE(x_dim, z_dim, device=device, p_y_zt_nn=p_y_zt_nn, p_y_zt_nn_layers=p_y_zt_nn_layers,\n",
    "            p_y_zt_nn_width=p_y_zt_nn_width, p_t_z_nn=p_t_z_nn, p_t_z_nn_layers=p_t_z_nn_layers, p_t_z_nn_width=p_t_z_nn_width,\n",
    "            p_x_z_nn=p_x_z_nn, p_x_z_nn_layers=p_x_z_nn_layers, p_x_z_nn_width=p_x_z_nn_width,separate_ty=separate_ty)\n",
    "print(\"Loading from \" + \"./model/{}/{}_X_sourceCEVAE_{}epoches_{}samples_{}\".format(main_folder,dist_z_source[0],num_epochs,number_trained_sample, t_net))\n",
    "\n",
    "source_CEVAE.load_state_dict(torch.load(\"./model/{}/{}_X_sourceCEVAE_{}epochs_{}samples_{}\".format(main_folder,dist_z_source[0],num_epochs,number_trained_sample, t_net)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-0.2362,  0.5533,  1.1840]], device='mps:0')\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 7704 characters]",
      "text/plain": [
       "<Figure size 400x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "# test the decoder\n",
    "#Example image from the target data distribution\n",
    "z = torch.normal(0,1,size=(1,3)).to(device)\n",
    "print(z)\n",
    "image,_,_,_,_,_,_ = source_CEVAE.decoder(z, t=1)\n",
    "out = image.squeeze()\n",
    "fig, ax = plt.subplots(1,1, figsize=(4,4))\n",
    "plt.imshow(dist.Bernoulli(logits=out).sample().detach().cpu())\n",
    "plt.grid(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 45468 characters]",
      "text/plain": [
       "<Figure size 1600x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_py_zt_image(source_dataframe, target_dataframe, device, source_CEVAE.decoder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test on the target domain dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test done for vae frequency  1000\n",
      "Test done for vae frequency  200\n",
      "Test done for vae frequency  100\n",
      "Test done for vae frequency  50\n",
      "Test done for vae frequency  25\n",
      "Test done for vae frequency  15\n",
      "Test done for vae frequency  10\n",
      "Test done for vae frequency  8\n",
      "Test done for vae frequency  6\n",
      "Test done for vae frequency  5\n"
     ]
    }
   ],
   "source": [
    "\n",
    "from env import eva_causal_image, eva_random, eva_negative\n",
    "from copy import deepcopy\n",
    "import random\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "seeds = [1,2,3,4,5]\n",
    "\n",
    "# eliminate the randomness of data sampling\n",
    "seed = 1\n",
    "np.random.seed(seed)\n",
    "random.seed(seed)\n",
    "N = 1000\n",
    "#test_data = data.sample(N)        # target domain N(0, 1)\n",
    "test_data = target_data.sample(N)\n",
    "test_data.shuffle() #shuffle the data in case of order of the dataset\n",
    "# Hyperparamer setting\n",
    "beta = 0.1\n",
    "vae_freq = [1000, 200, 100, 50, 25, 15, 10, 8, 6, 5]\n",
    "lr = 0.001\n",
    "batch_size = 32\n",
    "y_separate_heads = True\n",
    "update_decoder = True\n",
    "train = True\n",
    "\n",
    "total_regret_vs_vae_freq = {\"Causal\":[], \"VAE\":[], \"VAE(prior)\": [], \"Random\": [] }\n",
    "\n",
    "for j in range(len(vae_freq)):\n",
    "        reward_causal_transDe_episodes = []     \n",
    "        reward_causal_NoDe_episodes = []\n",
    "        reward_causal_sourceVAE_episodes = []\n",
    "        reward_negative_episodes = []\n",
    "        reward_random_episodes = []\n",
    "\n",
    "        regret_causal_transDe_episodes = []\n",
    "        regret_causal_NoDe_episodes = []\n",
    "        regret_causal_sourceVAE_episodes = []\n",
    "        regret_negative_episodes = []\n",
    "        regret_random_episodes = []\n",
    "\n",
    "        for i in range(len(seeds)):\n",
    "\n",
    "                seed = seeds[i]\n",
    "                np.random.seed(seed)\n",
    "                random.seed(seed)\n",
    "                torch.manual_seed(seed)\n",
    "                torch.cuda.manual_seed_all(seed)\n",
    "\n",
    "                torch.backends.cudnn.deterministic = True\n",
    "                torch.backends.cudnn.benchmark = False\n",
    "\n",
    "                target_vae = ImageCEVAE(x_dim, z_dim, device=device, p_y_zt_nn=p_y_zt_nn, p_y_zt_nn_layers=p_y_zt_nn_layers,\n",
    "            p_y_zt_nn_width=p_y_zt_nn_width, p_t_z_nn=p_t_z_nn, p_t_z_nn_layers=p_t_z_nn_layers, p_t_z_nn_width=p_t_z_nn_width,\n",
    "            p_x_z_nn=p_x_z_nn, p_x_z_nn_layers=p_x_z_nn_layers, p_x_z_nn_width=p_x_z_nn_width,separate_ty=separate_ty)\n",
    "                transDe_vae = ImageCEVAE(x_dim, z_dim, device=device, p_y_zt_nn=p_y_zt_nn,              p_y_zt_nn_layers=p_y_zt_nn_layers,\n",
    "            p_y_zt_nn_width=p_y_zt_nn_width, p_t_z_nn=p_t_z_nn, p_t_z_nn_layers=p_t_z_nn_layers, p_t_z_nn_width=p_t_z_nn_width,\n",
    "            p_x_z_nn=p_x_z_nn, p_x_z_nn_layers=p_x_z_nn_layers, p_x_z_nn_width=p_x_z_nn_width,separate_ty=separate_ty)\n",
    "                transDe_vae.decoder = deepcopy(transDe_CEVAE.decoder)\n",
    "                #transDe_vae.encoder = deepcopy(source_CEVAE.encoder)\n",
    "                #transDe_vae.decoder = deepcopy(source_CEVAE.decoder)\n",
    "                source_vae = deepcopy(source_CEVAE)\n",
    "\n",
    "                reward_causal_transDe, regret_causal_transDe, transVAE_z, true_z = eva_causal_image(device, test_data, transDe_vae, beta=beta, vae_freq=vae_freq[j], lr=lr, batch_size=batch_size, train = train, update_decoder=update_decoder)\n",
    "                reward_causal_NoDe, regret_causal_NoDe, randomVAE_z, _ = eva_causal_image(device, test_data, target_vae, beta=beta, vae_freq=vae_freq[j], lr=lr, batch_size=batch_size, train=train)\n",
    "                reward_causal_sourceVAE, regret_causal_sourceVAE, sourceVAE_z, true_z = eva_causal_image(device, test_data, source_vae, beta=beta, vae_freq=vae_freq[j], lr=lr, batch_size=batch_size, train=train)\n",
    "                reward_random, regret_random = eva_random(test_data, N) \n",
    "\n",
    "                reward_causal_transDe_episodes.append(reward_causal_transDe)\n",
    "                reward_causal_NoDe_episodes.append(reward_causal_NoDe)\n",
    "                reward_causal_sourceVAE_episodes.append(reward_causal_sourceVAE)\n",
    "                reward_random_episodes.append(reward_random)\n",
    "\n",
    "                regret_causal_transDe_episodes.append(regret_causal_transDe)\n",
    "                regret_causal_NoDe_episodes.append(regret_causal_NoDe)\n",
    "                regret_causal_sourceVAE_episodes.append(regret_causal_sourceVAE)\n",
    "                regret_random_episodes.append(regret_random) \n",
    "        print(\"Test done for vae frequency \", vae_freq[j])\n",
    "        total_regret_vs_vae_freq[\"Causal\"].append(np.mean(np.array([regret_causal_transDe_episodes[i][-1] for i in range(len(seeds))])))  \n",
    "        total_regret_vs_vae_freq[\"VAE\"].append(np.mean(np.array([regret_causal_NoDe_episodes[i][-1] for i in range(len(seeds))])))\n",
    "        total_regret_vs_vae_freq[\"VAE(prior)\"].append(np.mean(np.array([regret_causal_sourceVAE_episodes[i][-1] for i in range(len(seeds))])))     \n",
    "        total_regret_vs_vae_freq[\"Random\"].append(np.mean(np.array([regret_random_episodes[i][-1] for i in range(len(seeds))])))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total_regret_-_+.csv\n"
     ]
    }
   ],
   "source": [
    "# store the results in csv file\n",
    "# data structure:\n",
    "# tranDe,randomVAE,sourceVAE,random, gradient step\n",
    "# x1,x2,x3,x4,x5\n",
    "# x1,x2,x3,x4,x5\n",
    "# x1,x2,x3,x4,x6\n",
    "import csv\n",
    "\n",
    "num_update = [int(1000/freq) for freq in vae_freq]\n",
    "name = \"MNIST_total_regret_\"+str(dist_z_source[0])+\"_\"+str(dist_z_target[0])+\".csv\"\n",
    "print(name)\n",
    "# Open a CSV file for writing\n",
    "\n",
    "keys = list(total_regret_vs_vae_freq.keys())\n",
    "keys.append('gradient_steps')\n",
    "\n",
    "# Open a CSV file for writing\n",
    "with open(\"results/\"+name, 'w', newline='') as file:\n",
    "    writer = csv.writer(file)\n",
    "\n",
    "    # Write the header (keys of the dictionary)\n",
    "    writer.writerow(keys)\n",
    "\n",
    "    # Write the rows (zip the ndarray values to create rows)\n",
    "    rows = zip(*total_regret_vs_vae_freq.values(), num_update)\n",
    "    for row in rows:\n",
    "        writer.writerow(row)\n",
    "\n",
    "keys =['posterior_z', 'true_z']\n",
    "\n",
    "\n",
    "with open('results/'+name, 'w', newline='') as file:\n",
    "    writer = csv.writer(file)\n",
    "\n",
    "    writer.writerow(keys)\n",
    "\n",
    "    rows = zip(transVAE_z, true_z)\n",
    "    for row in rows:\n",
    "        writer.writerow(row)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 45844 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# ploting the regret\n",
    "from plot_functions import plot_regret, plot_total_regret_vs_freq\n",
    "\n",
    "#plot_regret(regret_causal_transDe_episodes, regret_causal_NoDe_episodes, regret_causal_sourceVAE_episodes, regret_negative_episodes, regret_random_episodes, dist_z_source=dist_z_source, dist_z_target=dist_z_target)\n",
    "color_list = ['red', 'blue', 'green', 'black']\n",
    "num_update = [int(1000/freq) for freq in vae_freq]\n",
    "plot_total_regret_vs_freq(total_regret_vs_vae_freq, num_update, color_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 32884 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "z = np.random.randn(1000)\n",
    "z = [transVAE_z, true_z]\n",
    "labels = ['posterior', 'ground truth']\n",
    "name = \"MNIST_dist_z_\"+str(dist_z_source[0])+\"_\"+str(dist_z_target[0])+\".tex\"\n",
    "plot_z_comparasion(z, labels, name, z_index=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
