geometric_toolkit.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HHm1cOIx_2sz"
   },
   "source": [
    "# Organization\n",
    "\n",
    "\n",
    "1.   Training/loading VAE\n",
    "2.   Computing geodesics on data manifold\n",
    "3. Geodesics restricted to straight line\n",
    "4. Derivative direction vs. interpolation direction\n",
    "5. Greedy geodesics vs. straight-line geodesics\n",
    "6. Studying the geodesics in more detail\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "uYvcEm6YBtRH"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import argparse\n",
    "import torch\n",
    "import torch.utils.data\n",
    "from torch import nn, optim\n",
    "from torch.nn import functional as F\n",
    "from torchvision import datasets, transforms\n",
    "from torchvision.utils import save_image\n",
    "import pickle\n",
    "import sys\n",
    "import numpy as np\n",
    "import scipy\n",
    "from scipy import stats\n",
    "import matplotlib.pyplot as plt\n",
    "import math\n",
    "\n",
    " # To use the Google Colab GPU acceleration, go to Edit --> Notebook Settings.\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "RYElsJqABfjU"
   },
   "outputs": [],
   "source": [
    "\"\"\" The following code will store the trained models in your Google drive, so you\n",
    "do not need to retrain the models every single time the notebook restarts.\n",
    "\"\"\"\n",
    "\"\"\"from google.colab import drive\n",
    "drive.mount('/content/drive')\n",
    "notebook_data_path = '/content/drive/MyDrive/geomtoolkit/'\n",
    "if not os.path.exists(notebook_data_path):\n",
    "  os.mkdir(notebook_data_path)\"\"\"\n",
    "notebook_data_path = '/home/gridsan/hanlaw/' # change per-person"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "L8IqjEcR_RMK"
   },
   "source": [
    "# Training/loading the VAE\n",
    "\n",
    "The following code either trains the VAE, or loads a pre-trained VAE from a pickle file. The latter is preferable if the VAE has not been yet trained."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "aaRcZ5fgCcrz"
   },
   "outputs": [],
   "source": [
    "#@title Code: VAE class definition\n",
    "\"\"\"\n",
    "The following code is a slightly modified version of the pytorch library's\n",
    "example directory for representing a simple VAE.\n",
    "\"\"\"\n",
    "class VAE(nn.Module):\n",
    "    def __init__(self, latent_d=20):\n",
    "        super(VAE, self).__init__()\n",
    "\n",
    "        self.latent_d = latent_d\n",
    "        self.fc1 = nn.Linear(784, 400)\n",
    "        self.fc21 = nn.Linear(400, latent_d)\n",
    "        self.fc22 = nn.Linear(400, latent_d)\n",
    "        self.fc3 = nn.Linear(latent_d, 400)\n",
    "        self.fc4 = nn.Linear(400, 784)\n",
    "\n",
    "    def encode(self, x):\n",
    "        h1 = F.relu(self.fc1(x))\n",
    "        return self.fc21(h1), self.fc22(h1)\n",
    "\n",
    "    def reparameterize(self, mu, logvar):\n",
    "        std = torch.exp(0.5*logvar)\n",
    "        eps = torch.randn_like(std)\n",
    "        return mu + eps*std\n",
    "\n",
    "    def decode(self, z):\n",
    "        h3 = F.relu(self.fc3(z))\n",
    "        return torch.sigmoid(self.fc4(h3))\n",
    "\n",
    "    def forward(self, x):\n",
    "        mu, logvar = self.encode(x.view(-1, 784))\n",
    "        z = self.reparameterize(mu, logvar)\n",
    "        return self.decode(z), mu, logvar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "Tg9D01vDJuFq"
   },
   "outputs": [],
   "source": [
    "def get_trained_vae(dataset_name,latent_d):\n",
    "  model_path = notebook_data_path + 'trained_models/'\n",
    "  if not os.path.exists(model_path):\n",
    "    os.mkdir(model_path)\n",
    "  vae_filename = model_path + dataset_name + '_vae_d' + str(latent_d) + '.pkl'\n",
    "  if not os.path.exists(vae_filename):\n",
    "    print('VAE does not exist. Training it now.')\n",
    "    train_vae(dataset_name,latent_d,vae_filename)\n",
    "  model = pickle.load(open(vae_filename, 'rb'))\n",
    "  return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "5zGwB-XOLrhT"
   },
   "outputs": [],
   "source": [
    "#@title Code: VAE training\n",
    "\"\"\"\n",
    "The following code is a slightly modified version of the pytorch library's\n",
    "example directory for training a simple VAE.\n",
    "\"\"\"\n",
    "def train_vae(datasetname,latent_d,filename):\n",
    "    sys.argv = ['-f']\n",
    "\n",
    "    parser = argparse.ArgumentParser(description='VAE MNIST Example')\n",
    "    parser.add_argument('--batch-size', type=int, default=128, metavar='N',\n",
    "                        help='input batch size for training (default: 128)')\n",
    "    parser.add_argument('--epochs', type=int, default=10, metavar='N',\n",
    "                        help='number of epochs to train (default: 10)')\n",
    "    parser.add_argument('--no-cuda', action='store_true', default=False,\n",
    "                        help='disables CUDA training')\n",
    "    parser.add_argument('--seed', type=int, default=1, metavar='S',\n",
    "                        help='random seed (default: 1)')\n",
    "    parser.add_argument('--log-interval', type=int, default=10, metavar='N',\n",
    "                        help='how many batches to wait before logging training status')\n",
    "    parser.add_argument('--vis-interval',type=int, default=10, metavar='N',\n",
    "                        help='how many batches to wait before dumping visualization')\n",
    "    args = parser.parse_args()\n",
    "    args.cuda = not args.no_cuda and torch.cuda.is_available()\n",
    "\n",
    "    # torch.manual_seed(args.seed)\n",
    "\n",
    "    kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}\n",
    "\n",
    "\n",
    "    dataset_name = datasetname\n",
    "    if dataset_name == 'mnist':\n",
    "        traindataset = datasets.MNIST(notebook_data_path + 'data/', train=True, download=True,\n",
    "                      transform=transforms.ToTensor())\n",
    "\n",
    "        testdataset = datasets.MNIST(notebook_data_path + 'data/', train=False, transform=transforms.ToTensor())\n",
    "\n",
    "    elif dataset_name in ['angle', 'circleangle']:\n",
    "        transformlist = transforms.Compose([transforms.Grayscale(num_output_channels=1),\n",
    "                                        transforms.ToTensor()])\n",
    "        traindataset = datasets.ImageFolder('./data/' + dataset_name + '/', transform=transformlist)\n",
    "        testdataset = datasets.ImageFolder('./data/' + dataset_name + '/', transform=transformlist)\n",
    "\n",
    "    elif dataset_name == 'untrained':\n",
    "        args.epochs = 0\n",
    "        # dummy code\n",
    "        traindataset = datasets.MNIST(notebook_data_path + 'data/', train=True, download=True,\n",
    "              transform=transforms.ToTensor())\n",
    "\n",
    "        testdataset = datasets.MNIST(notebook_data_path + 'data/', train=False, transform=transforms.ToTensor())\n",
    "\n",
    "\n",
    "    train_loader = torch.utils.data.DataLoader(traindataset,\n",
    "        batch_size=args.batch_size, shuffle=True, **kwargs)\n",
    "    test_loader = torch.utils.data.DataLoader(testdataset,\n",
    "        batch_size=args.batch_size, shuffle=True, **kwargs)\n",
    "\n",
    "    # Reconstruction + KL divergence losses summed over all elements and batch\n",
    "    def vae_loss_function(recon_x, x, mu, logvar):\n",
    "        BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction='sum')\n",
    "\n",
    "        # see Appendix B from VAE paper:\n",
    "        # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014\n",
    "        # https://arxiv.org/abs/1312.6114\n",
    "        # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)\n",
    "        KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())\n",
    "\n",
    "        return BCE + KLD\n",
    "\n",
    "    def vae_train(model, optimizer, epoch):\n",
    "        model.train()\n",
    "        train_loss = 0\n",
    "        for batch_idx, (data, _) in enumerate(train_loader):\n",
    "            data = data.to(device)\n",
    "            optimizer.zero_grad()\n",
    "            recon_batch, mu, logvar = model(data)\n",
    "            loss = vae_loss_function(recon_batch, data, mu, logvar)\n",
    "            loss.backward()\n",
    "            train_loss += loss.item()\n",
    "            optimizer.step()\n",
    "            if batch_idx % args.log_interval == 0:\n",
    "                print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
    "                    epoch, batch_idx * len(data), len(train_loader.dataset),\n",
    "                    100. * batch_idx / len(train_loader),\n",
    "                    loss.item() / len(data)))\n",
    "            # if batch_idx % args.vis_interval == 0:\n",
    "                # vis_weights(model)\n",
    "\n",
    "        print('====> Epoch: {} Average loss: {:.4f}'.format(epoch, train_loss / len(train_loader.dataset)))\n",
    "\n",
    "    model = VAE(latent_d = latent_d).to(device)\n",
    "    optimizer = optim.Adam(model.parameters(), lr=1e-3)\n",
    "    for epoch in range(1, args.epochs + 1):\n",
    "        vae_train(model, optimizer, epoch)\n",
    "    pickle.dump(model, open(filename, 'wb'))\n",
    "    print('trained',latent_d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "KLfpTUGRErMU"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VAE does not exist. Training it now.\n",
      "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to /home/gridsan/hanlaw/data/MNIST/raw/train-images-idx3-ubyte.gz\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "613704d5278b489c84ce7a08fb92fb24",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value=''), FloatProgress(value=1.0, bar_style='info', layout=Layout(width='20px'), max=1.0…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "URLError",
     "evalue": "<urlopen error [Errno 101] Network is unreachable>",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mOSError\u001b[0m                                   Traceback (most recent call last)",
      "\u001b[0;32m/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.7/urllib/request.py\u001b[0m in \u001b[0;36mdo_open\u001b[0;34m(self, http_class, req, **http_conn_args)\u001b[0m\n\u001b[1;32m   1316\u001b[0m                 h.request(req.get_method(), req.selector, req.data, headers,\n\u001b[0;32m-> 1317\u001b[0;31m                           encode_chunked=req.has_header('Transfer-encoding'))\n\u001b[0m\u001b[1;32m   1318\u001b[0m             \u001b[0;32mexcept\u001b[0m \u001b[0mOSError\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0merr\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;31m# timeout error\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.7/http/client.py\u001b[0m in \u001b[0;36mrequest\u001b[0;34m(self, method, url, body, headers, encode_chunked)\u001b[0m\n\u001b[1;32m   1228\u001b[0m         \u001b[0;34m\"\"\"Send a complete request to the server.\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1229\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_send_request\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmethod\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0murl\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbody\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mheaders\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mencode_chunked\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1230\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.7/http/client.py\u001b[0m in \u001b[0;36m_send_request\u001b[0;34m(self, method, url, body, headers, encode_chunked)\u001b[0m\n\u001b[1;32m   1274\u001b[0m             \u001b[0mbody\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_encode\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbody\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'body'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1275\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mendheaders\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbody\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mencode_chunked\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mencode_chunked\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1276\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.7/http/client.py\u001b[0m in \u001b[0;36mendheaders\u001b[0;34m(self, message_body, encode_chunked)\u001b[0m\n\u001b[1;32m   1223\u001b[0m             \u001b[0;32mraise\u001b[0m \u001b[0mCannotSendHeader\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1224\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_send_output\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmessage_body\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mencode_chunked\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mencode_chunked\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1225\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.7/http/client.py\u001b[0m in \u001b[0;36m_send_output\u001b[0;34m(self, message_body, encode_chunked)\u001b[0m\n\u001b[1;32m   1015\u001b[0m         \u001b[0;32mdel\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_buffer\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1016\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmsg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1017\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.7/http/client.py\u001b[0m in \u001b[0;36msend\u001b[0;34m(self, data)\u001b[0m\n\u001b[1;32m    955\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mauto_open\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 956\u001b[0;31m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconnect\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    957\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.7/http/client.py\u001b[0m in \u001b[0;36mconnect\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    927\u001b[0m         self.sock = self._create_connection(\n\u001b[0;32m--> 928\u001b[0;31m             (self.host,self.port), self.timeout, self.source_address)\n\u001b[0m\u001b[1;32m    929\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msock\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msetsockopt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msocket\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mIPPROTO_TCP\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msocket\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTCP_NODELAY\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.7/socket.py\u001b[0m in \u001b[0;36mcreate_connection\u001b[0;34m(address, timeout, source_address)\u001b[0m\n\u001b[1;32m    726\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0merr\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 727\u001b[0;31m         \u001b[0;32mraise\u001b[0m \u001b[0merr\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    728\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.7/socket.py\u001b[0m in \u001b[0;36mcreate_connection\u001b[0;34m(address, timeout, source_address)\u001b[0m\n\u001b[1;32m    715\u001b[0m                 \u001b[0msock\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbind\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msource_address\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 716\u001b[0;31m             \u001b[0msock\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconnect\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msa\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    717\u001b[0m             \u001b[0;31m# Break explicitly a reference cycle\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mOSError\u001b[0m: [Errno 101] Network is unreachable",
      "\nDuring handling of the above exception, another exception occurred:\n",
      "\u001b[0;31mURLError\u001b[0m                                  Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-12-38df74e1edf9>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0mDrive\u001b[0m \u001b[0mto\u001b[0m \u001b[0mallow\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mre\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0muse\u001b[0m \u001b[0mlater\u001b[0m \u001b[0mon\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \"\"\"\n\u001b[0;32m----> 6\u001b[0;31m \u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_trained_vae\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'mnist'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m5\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-10-c5ddea160dd4>\u001b[0m in \u001b[0;36mget_trained_vae\u001b[0;34m(dataset_name, latent_d)\u001b[0m\n\u001b[1;32m      6\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexists\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mvae_filename\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'VAE does not exist. Training it now.'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 8\u001b[0;31m     \u001b[0mtrain_vae\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdataset_name\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mlatent_d\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mvae_filename\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      9\u001b[0m   \u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpickle\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mload\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mvae_filename\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'rb'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     10\u001b[0m   \u001b[0;32mreturn\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-11-2b9395c5a64d>\u001b[0m in \u001b[0;36mtrain_vae\u001b[0;34m(datasetname, latent_d, filename)\u001b[0m\n\u001b[1;32m     31\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mdataset_name\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'mnist'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     32\u001b[0m         traindataset = datasets.MNIST(notebook_data_path + 'data/', train=True, download=True,\n\u001b[0;32m---> 33\u001b[0;31m                       transform=transforms.ToTensor())\n\u001b[0m\u001b[1;32m     34\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     35\u001b[0m         \u001b[0mtestdataset\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdatasets\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mMNIST\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnotebook_data_path\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m'data/'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrain\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtransform\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtransforms\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mToTensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.7/site-packages/torchvision/datasets/mnist.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, root, train, transform, target_transform, download)\u001b[0m\n\u001b[1;32m     68\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     69\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mdownload\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 70\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdownload\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     71\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     72\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_check_exists\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.7/site-packages/torchvision/datasets/mnist.py\u001b[0m in \u001b[0;36mdownload\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    135\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0murl\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmd5\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mresources\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    136\u001b[0m             \u001b[0mfilename\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0murl\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrpartition\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'/'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 137\u001b[0;31m             \u001b[0mdownload_and_extract_archive\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0murl\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdownload_root\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mraw_folder\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfilename\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfilename\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmd5\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mmd5\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    138\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    139\u001b[0m         \u001b[0;31m# process and save as torch files\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.7/site-packages/torchvision/datasets/utils.py\u001b[0m in \u001b[0;36mdownload_and_extract_archive\u001b[0;34m(url, download_root, extract_root, filename, md5, remove_finished)\u001b[0m\n\u001b[1;32m    233\u001b[0m         \u001b[0mfilename\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbasename\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0murl\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    234\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 235\u001b[0;31m     \u001b[0mdownload_url\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0murl\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdownload_root\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfilename\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmd5\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    236\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    237\u001b[0m     \u001b[0marchive\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjoin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdownload_root\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfilename\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.7/site-packages/torchvision/datasets/utils.py\u001b[0m in \u001b[0;36mdownload_url\u001b[0;34m(url, root, filename, md5)\u001b[0m\n\u001b[1;32m     81\u001b[0m                 )\n\u001b[1;32m     82\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 83\u001b[0;31m                 \u001b[0;32mraise\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     84\u001b[0m         \u001b[0;31m# check integrity of downloaded file\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     85\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mcheck_integrity\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfpath\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmd5\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.7/site-packages/torchvision/datasets/utils.py\u001b[0m in \u001b[0;36mdownload_url\u001b[0;34m(url, root, filename, md5)\u001b[0m\n\u001b[1;32m     69\u001b[0m             urllib.request.urlretrieve(\n\u001b[1;32m     70\u001b[0m                 \u001b[0murl\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfpath\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 71\u001b[0;31m                 \u001b[0mreporthook\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mgen_bar_updater\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     72\u001b[0m             )\n\u001b[1;32m     73\u001b[0m         \u001b[0;32mexcept\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0murllib\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0merror\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mURLError\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mIOError\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.7/urllib/request.py\u001b[0m in \u001b[0;36murlretrieve\u001b[0;34m(url, filename, reporthook, data)\u001b[0m\n\u001b[1;32m    245\u001b[0m     \u001b[0murl_type\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpath\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msplittype\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0murl\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    246\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 247\u001b[0;31m     \u001b[0;32mwith\u001b[0m \u001b[0mcontextlib\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclosing\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0murlopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0murl\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mfp\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    248\u001b[0m         \u001b[0mheaders\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minfo\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    249\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.7/urllib/request.py\u001b[0m in \u001b[0;36murlopen\u001b[0;34m(url, data, timeout, cafile, capath, cadefault, context)\u001b[0m\n\u001b[1;32m    220\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    221\u001b[0m         \u001b[0mopener\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_opener\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 222\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mopener\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0murl\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtimeout\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    223\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    224\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0minstall_opener\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mopener\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.7/urllib/request.py\u001b[0m in \u001b[0;36mopen\u001b[0;34m(self, fullurl, data, timeout)\u001b[0m\n\u001b[1;32m    523\u001b[0m             \u001b[0mreq\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmeth\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mreq\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    524\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 525\u001b[0;31m         \u001b[0mresponse\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_open\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mreq\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    526\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    527\u001b[0m         \u001b[0;31m# post-process response\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.7/urllib/request.py\u001b[0m in \u001b[0;36m_open\u001b[0;34m(self, req, data)\u001b[0m\n\u001b[1;32m    541\u001b[0m         \u001b[0mprotocol\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mreq\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtype\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    542\u001b[0m         result = self._call_chain(self.handle_open, protocol, protocol +\n\u001b[0;32m--> 543\u001b[0;31m                                   '_open', req)\n\u001b[0m\u001b[1;32m    544\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    545\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.7/urllib/request.py\u001b[0m in \u001b[0;36m_call_chain\u001b[0;34m(self, chain, kind, meth_name, *args)\u001b[0m\n\u001b[1;32m    501\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhandler\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mhandlers\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    502\u001b[0m             \u001b[0mfunc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgetattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhandler\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmeth_name\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 503\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    504\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mresult\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    505\u001b[0m                 \u001b[0;32mreturn\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.7/urllib/request.py\u001b[0m in \u001b[0;36mhttp_open\u001b[0;34m(self, req)\u001b[0m\n\u001b[1;32m   1343\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1344\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mhttp_open\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreq\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1345\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdo_open\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhttp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclient\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mHTTPConnection\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreq\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1346\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1347\u001b[0m     \u001b[0mhttp_request\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mAbstractHTTPHandler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdo_request_\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/state/partition1/llgrid/pkg/anaconda/anaconda3-2020b/lib/python3.7/urllib/request.py\u001b[0m in \u001b[0;36mdo_open\u001b[0;34m(self, http_class, req, **http_conn_args)\u001b[0m\n\u001b[1;32m   1317\u001b[0m                           encode_chunked=req.has_header('Transfer-encoding'))\n\u001b[1;32m   1318\u001b[0m             \u001b[0;32mexcept\u001b[0m \u001b[0mOSError\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0merr\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;31m# timeout error\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1319\u001b[0;31m                 \u001b[0;32mraise\u001b[0m \u001b[0mURLError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0merr\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1320\u001b[0m             \u001b[0mr\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mh\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgetresponse\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1321\u001b[0m         \u001b[0;32mexcept\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mURLError\u001b[0m: <urlopen error [Errno 101] Network is unreachable>"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "Train VAE on MNIST dataset with latent dimension 5.\n",
    "The first time that you do this, the model will save the model to your Google\n",
    "Drive to allow for re-use later on.\n",
    "\"\"\"\n",
    "model = get_trained_vae('mnist', 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ANWBDXhL_WDA"
   },
   "source": [
    "# Computing geodesics on manifolds\n",
    "\n",
    "Implementation of the \"classic\" approach where you start with straight-line interpolation in latent space and perturb it to obtain shorter a path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "chXS4CY6YqDS"
   },
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "def minimize_path_energy(model, z0, z1, numpts=10,numtrainiter=1000,learning_rate=0.001,verbose=False):\n",
    "  \"\"\"\n",
    "  Initialization: a linear interpolating path between z0 and z1 with numpts number of points.\n",
    "  Output: a path with numpts between z0 and z1 trained to minimize energy with respect to the model.\n",
    "  \"\"\"\n",
    "\n",
    "  # Initialize the interpolation between z0 and z1 with a linear interpolation.\n",
    "  t = torch.linspace(1 / (numpts + 1), numpts / (numpts + 1), numpts).to(device)\n",
    "  interp_points = torch.outer(1 - t, torch.flatten(z0)) + torch.outer(t, torch.flatten(z1))\n",
    "  interp_points.requires_grad = True\n",
    "\n",
    "  # For convenience, precompute x0 and x1.\n",
    "  x0 = model.decode(z0).detach()\n",
    "  x1 = model.decode(z1).detach()\n",
    "\n",
    "  optimizer = optim.Adam([interp_points], lr=learning_rate)\n",
    "  for i in range(numtrainiter):\n",
    "    imgs = model.decode(interp_points)\n",
    "\n",
    "    # compute the energy of the path\n",
    "    loss = torch.sum((torch.cat((x0, imgs)) - torch.cat((imgs, x1)))**2)\n",
    "\n",
    "    # # compute the length of the path\n",
    "    # loss = torch.sqrt(torch.sum((torch.cat((x0, imgs)) - torch.cat((imgs, x1)))**2))\n",
    "\n",
    "    if verbose:\n",
    "      with torch.no_grad():\n",
    "        if i % 1000 == 0:\n",
    "          print('loss',i,loss.to('cpu'))\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "  imgs = model.decode(interp_points)\n",
    "\n",
    "  # compute the energy of the path\n",
    "  loss = torch.sum((torch.cat((x0, imgs)) - torch.cat((imgs, x1)))**2)\n",
    "\n",
    "  # # compute the length of the path\n",
    "  # loss = torch.sqrt(torch.sum((torch.cat((x0, imgs)) - torch.cat((imgs, x1)))**2))\n",
    "\n",
    "  return torch.cat((z0, interp_points.detach(), z1)), loss.detach()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Fzh830szO8tC"
   },
   "outputs": [],
   "source": [
    "def get_curve_length(model, zseq):\n",
    "  length = 0\n",
    "  for i in range(zseq.shape[0]-1):\n",
    "    length = length + torch.sqrt(torch.sum((model.decode(zseq[i,:]) - model.decode(zseq[i+1,:])) ** 2))\n",
    "  return length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pZq3o_25rRtW"
   },
   "outputs": [],
   "source": [
    "def disp_vec_img(a,filename=None):\n",
    "  with torch.no_grad():\n",
    "    a = np.asarray(a.to('cpu'))\n",
    "    pixels = a.reshape((28, 28))\n",
    "    plt.imshow(pixels, cmap='gray')\n",
    "    if filename is None:\n",
    "      plt.show()\n",
    "    else:\n",
    "      plt.savefig(filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0x7RYA_G3i0G"
   },
   "outputs": [],
   "source": [
    "import imageio\n",
    "from base64 import b64encode\n",
    "from IPython.display import HTML\n",
    "\n",
    "def display_imageseq_video(zseq, video_file='test.mp4'):\n",
    "  with torch.no_grad():\n",
    "    imageseq = model.decode(zseq)\n",
    "    imageseq = (imageseq * 255).byte()\n",
    "  imageseq = np.asarray(imageseq.to('cpu'))\n",
    "  imageseq = [np.pad(imageseq[i,:].reshape((28,28)),((2,2),(2,2))) for i in range(imageseq.shape[0])]\n",
    "  imageio.mimwrite('test.mp4', imageseq, ffmpeg_params=['-sws_flags', 'neighbor', '-vf', 'scale=320:320'], fps=200); \n",
    "  mp4 = open('test.mp4','rb').read()\n",
    "  data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
    "  htmlobj = HTML(\"\"\"\n",
    "  <video width=320 controls>\n",
    "        <source src=\"%s\" type=\"video/mp4\">\n",
    "  </video>\n",
    "  \"\"\" % data_url)\n",
    "  return htmlobj"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5EChSI0o6KWK"
   },
   "outputs": [],
   "source": [
    "z0 = torch.randn([1, model.latent_d]).to(device)\n",
    "z1 = torch.randn([1, model.latent_d]).to(device)\n",
    "print('We will be interpolating between the following two random digits:')\n",
    "disp_vec_img(model.decode(z0))\n",
    "disp_vec_img(model.decode(z1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7Ga_9seq6FLy"
   },
   "outputs": [],
   "source": [
    "numpts = 1000\n",
    "zlinearseq, _ = minimize_path_energy(model, z0, z1, numpts=numpts,numtrainiter=0)\n",
    "zseq, _ = minimize_path_energy(model, z0, z1, numpts=numpts,numtrainiter=3001,learning_rate=0.0005,verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "J4Yk_p80PSG9"
   },
   "outputs": [],
   "source": [
    "get_curve_length(model, zlinearseq)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8D6t42dCPXRt"
   },
   "outputs": [],
   "source": [
    "get_curve_length(model, zseq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Wf8RerMUnF4b"
   },
   "outputs": [],
   "source": [
    "def component_orth_to(vec, subspace):\n",
    "  vec2 = torch.clone(vec)\n",
    "  sub2 = torch.clone(subspace)\n",
    "  # gram-schmidt orthogonalization of sub2\n",
    "  for i in range(0,sub2.shape[0]):\n",
    "    for j in range(i):\n",
    "      sub2[i,:] = sub2[i,:] - sub2[j,:] * torch.dot(sub2[j,:], sub2[i,:])\n",
    "    sub2[i,:] = sub2[i,:] / torch.norm(sub2[i,:])\n",
    "  # print(sub2)\n",
    "  \n",
    "  for i in range(vec.shape[0]):\n",
    "    for j in range(0,sub2.shape[0]):\n",
    "      vec2[i,:] = vec2[i,:] - sub2[j,:] * torch.dot(sub2[j,:],vec2[i,:])\n",
    "\n",
    "  return vec2\n",
    "\n",
    "\n",
    "orth_to_1d = component_orth_to(zseq - z0, z1 - z0)\n",
    "orth_to_2d = component_orth_to(zseq - z0, torch.cat((z0,z1)))\n",
    "# print(torch.norm(ans[0,:]))\n",
    "# print(torch.norm(ans[1,:]))\n",
    "# print(torch.dot(ans[0,:],ans[1,:]))\n",
    "print(torch.norm(zseq,dim=1))\n",
    "print(torch.norm(orth_to_1d,dim=1))\n",
    "print(torch.norm(orth_to_2d,dim=1))\n",
    "plt.plot(np.asarray((torch.norm(orth_to_1d,dim=1) / torch.norm(orth_to_2d,dim=1)).detach().cpu()))\n",
    "\n",
    "# The takeaway of this code block when run on the untrained network\n",
    "# is that projection of zseq to (z0,z1) plane is roughly\n",
    "# the same as projection to z0 -> z1 line. This indicates that any deviations\n",
    "# from the straight line interpolation are not due to \"curvature of the manifold\",\n",
    "# but rather to \"noise\" effects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nUTaOJKzmn4a"
   },
   "outputs": [],
   "source": [
    "zseq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fc6YNa-WmtJK"
   },
   "outputs": [],
   "source": [
    "display_imageseq_video(zlinearseq,video_file='test1.mp4')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sLQJDojW38JV"
   },
   "outputs": [],
   "source": [
    "display_imageseq_video(zseq,'test2.mp4')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aPxilmUDGG8E"
   },
   "source": [
    "# Geodesics restricted to the straight line\n",
    "\n",
    "In this section, we optimize over geodesics restricted to the straight line. Obviously, the only improvement over the original geodesic is due to the spacing of the points on the straight line improving.\n",
    "\n",
    "This can be computed in two ways:\n",
    "\n",
    "(1) We can take the arc length of the straight line (computed using Euler integration), and derive the optimal energy of an interpolation from it.\n",
    "\n",
    "(2) We can optimize with Adam.\n",
    "\n",
    "Naturally, I will implement and compare both approaches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "W1rtHNVLGrO8"
   },
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "def minimize_straight_line_energy(model, z0, z1, numpts=10,numtrainiter=1000,learning_rate=0.001,verbose=False):\n",
    "  \"\"\"\n",
    "  Initialization: a linear interpolating path between z0 and z1 with numpts number of points.\n",
    "  Output: a path with numpts between z0 and z1 trained to minimize energy with respect to the model.\n",
    "  \"\"\"\n",
    "\n",
    "  # Initialize the interpolation between z0 and z1 with a linear interpolation.\n",
    "  t = torch.linspace(1 / (numpts + 1), numpts / (numpts + 1), numpts).to(device)\n",
    "  t.requires_grad = True\n",
    "\n",
    "  # For convenience, precompute x0 and x1.\n",
    "  x0 = model.decode(z0).detach()\n",
    "  x1 = model.decode(z1).detach()\n",
    "\n",
    "  optimizer = optim.Adam([t], lr=learning_rate)\n",
    "  for i in range(numtrainiter):\n",
    "    interp_points = torch.outer(1 - t, torch.flatten(z0)) + torch.outer(t, torch.flatten(z1))\n",
    "    imgs = model.decode(interp_points)\n",
    "\n",
    "    # compute the energy of the path\n",
    "    loss = torch.sum((torch.cat((x0, imgs)) - torch.cat((imgs, x1)))**2)\n",
    "\n",
    "    if verbose:\n",
    "      with torch.no_grad():\n",
    "        if i % 1000 == 0:\n",
    "          print('loss',i,loss.to('cpu'))\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "  interp_points = torch.outer(1 - t, torch.flatten(z0)) + torch.outer(t, torch.flatten(z1))\n",
    "  imgs = model.decode(interp_points)\n",
    "  loss = torch.sum((torch.cat((x0, imgs)) - torch.cat((imgs, x1)))**2)\n",
    "  return torch.cat((z0, interp_points.detach(), z1)), loss.detach()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "m8aDSFVNghvp"
   },
   "outputs": [],
   "source": [
    "numtrials = 1\n",
    "dist_parameter = 4\n",
    "diff = torch.randn([numtrials, model.latent_d]).to(device) * math.sqrt(2)\n",
    "diff = diff / torch.norm(diff,dim=1).view((numtrials,1)) * dist_parameter\n",
    "z0 = torch.randn([numtrials, model.latent_d]).to(device) * math.sqrt(1/2) + diff / 2\n",
    "z1 = z0 + diff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "D9TBIX_qQMZ2"
   },
   "outputs": [],
   "source": [
    "z0 = torch.randn([1, model.latent_d]).to(device)\n",
    "z1 = torch.randn([1, model.latent_d]).to(device)\n",
    "print('We will be interpolating between the following two random digits:')\n",
    "disp_vec_img(model.decode(z0))\n",
    "disp_vec_img(model.decode(z1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NNxiDkq_HW-G"
   },
   "outputs": [],
   "source": [
    "numpts = 1000\n",
    "znaiveseq, naiveenergy = minimize_path_energy(model, z0, z1, numpts=numpts,numtrainiter=0)\n",
    "zstraightseq, straightenergy = minimize_straight_line_energy(model, z0, z1, numpts=numpts,numtrainiter=10000,learning_rate=0.0005,verbose=True)\n",
    "zseq, pathenergy = minimize_path_energy(model, z0, z1, numpts=numpts,numtrainiter=10000,learning_rate=0.0005,verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6J_Q-2qVQTlW"
   },
   "outputs": [],
   "source": [
    "print(get_curve_length(model,znaiveseq))\n",
    "print(get_curve_length(model,zstraightseq))\n",
    "print(get_curve_length(model,zseq))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Zgj0uCG4RPXY"
   },
   "outputs": [],
   "source": [
    "model = get_trained_vae('untrained',5)\n",
    "\n",
    "naive_lengths = []\n",
    "straight_lengths = []\n",
    "optimized_lengths = []\n",
    "numtrials = 1\n",
    "\n",
    "for trial_num in range(numtrials):\n",
    "  z0 = torch.randn([1, model.latent_d]).to(device)\n",
    "  z1 = torch.randn([1, model.latent_d]).to(device)\n",
    "\n",
    "  # dist_parameter = 100\n",
    "  # diff = torch.randn([1, model.latent_d]).to(device) * math.sqrt(2)\n",
    "  # diff = diff / torch.norm(diff,dim=1).view((1,1)) * dist_parameter\n",
    "  # z0 = torch.randn([1, model.latent_d]).to(device) * math.sqrt(1/2) + diff / 2\n",
    "  # z1 = z0 + diff\n",
    "\n",
    "  # print('We will be interpolating between the following two random digits:')\n",
    "  # disp_vec_img(model.decode(z0))\n",
    "  # disp_vec_img(model.decode(z1))\n",
    "  print('trial',trial_num)\n",
    "  numpts = 20\n",
    "  znaiveseq, naiveenergy = minimize_path_energy(model, z0, z1, numpts=numpts,numtrainiter=0)\n",
    "  # zstraightseq, straightenergy = minimize_straight_line_energy(model, z0, z1, numpts=numpts,numtrainiter=3001,learning_rate=0.0005,verbose=True)\n",
    "  zseq, pathenergy = minimize_path_energy(model, z0, z1, numpts=numpts,numtrainiter=10001,learning_rate=0.0005,verbose=True)\n",
    "\n",
    "  with torch.no_grad():\n",
    "    naive_lengths.append(get_curve_length(model,znaiveseq).cpu())\n",
    "    # straight_lengths.append(get_curve_length(model,zstraightseq).cpu())\n",
    "    optimized_lengths.append(get_curve_length(model,zseq).cpu())\n",
    "\n",
    "naive_lengths = np.asarray(naive_lengths)\n",
    "# straight_lengths = np.asarray(straight_lengths)\n",
    "optimized_lengths = np.asarray(optimized_lengths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "x43FcJJiSTnZ"
   },
   "outputs": [],
   "source": [
    "plt.scatter(np.linspace(0,25,100),np.linspace(0,25,100))\n",
    "plt.scatter(np.asarray(naive_lengths), np.asarray(optimized_lengths))\n",
    "# plt.scatter(np.asarray(naive_lengths), np.asarray(straight_lengths))\n",
    "plt.show()\n",
    "plt.scatter(np.asarray(naive_lengths),np.asarray(optimized_lengths) / np.asarray(naive_lengths))\n",
    "plt.show()\n",
    "# straight_lengths - naive_lengths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mwEe8Ws_RzMk"
   },
   "outputs": [],
   "source": [
    "np.asarray([x.cpu() for x in optimized_lengths])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rRRL2vR_RqvO"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8c_2JZ9XpsfJ"
   },
   "outputs": [],
   "source": [
    "zseq = minimize_path_energy(model, z0, z1, numpts=numpts,numtrainiter=100000,learning_rate=0.0005,verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cTuSpTmUrvJM"
   },
   "source": [
    "Below, we implement the other method, which is a simple analytical calculation based on the arc-length [TODO]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XjOxFlSssBKy"
   },
   "source": [
    "# Directions of derivatives\n",
    "\n",
    "Although the geodesics are not straight lines in the latent space, they are pretty close! Why could this be?\n",
    "\n",
    "#### Hypothesis:\n",
    "For any $z_0,z_1$ in the latent space, we have $\\frac{\\partial (g(z) - g(z_1))^2}{\\partial z}|_{z = z_0} \\approx z_1 - z_0$, where $\\approx$ means that the directions are roughly aligned. This would be the case if, for example, $g$ were a linear model. The hypothesis says that, to first-order approximation, the neural network is a linear model.\n",
    "\n",
    "#### Takeaway:\n",
    "The hypothesis is roughly correct, although a little less so if z0,z1 are far apart or if the dimension is large. See below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9nll22xOr4Z_"
   },
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "def grad_gz0_to_gz1(model, z0, z1):\n",
    "  \"\"\"\n",
    "  Given z0, z1, find the direction of the gradient at z0 that moves g(z0) to g(z1)\n",
    "  \"\"\"\n",
    "  z0.requires_grad = True\n",
    "  x0 = model.decode(z0)\n",
    "  x1 = model.decode(z1).detach()\n",
    "\n",
    "  loss = torch.sum((x0 - x1)**2)\n",
    "  loss.backward()\n",
    "  \n",
    "  return z0.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RQ6AI_9JvpOx"
   },
   "outputs": [],
   "source": [
    "import math\n",
    "for latent_d in [100]:\n",
    "  for dist_parameter in [None, 0.5, 2, 4, 8, 16]:\n",
    "    model = get_trained_vae('untrained',latent_d)\n",
    "\n",
    "    numtrials = 10000\n",
    "\n",
    "    if dist_parameter is None:\n",
    "      z0 = torch.randn([numtrials, model.latent_d]).to(device)\n",
    "      z1 = torch.randn([numtrials, model.latent_d]).to(device)\n",
    "\n",
    "    else:\n",
    "      diff = torch.randn([numtrials, model.latent_d]).to(device) * math.sqrt(2)\n",
    "      diff = diff / torch.norm(diff,dim=1).view((numtrials,1)) * dist_parameter\n",
    "      z0 = torch.randn([numtrials, model.latent_d]).to(device) * math.sqrt(1/2) + diff / 2\n",
    "      z1 = z0 + diff\n",
    "\n",
    "    z0grad = grad_gz0_to_gz1(model, z0, z1)\n",
    "    z0delta = z1 - z0\n",
    "    sample_distance = torch.norm(z0delta,dim=1).detach().cpu()\n",
    "\n",
    "    z0grad = z0grad / torch.norm(z0grad,dim=1).view((numtrials,1))\n",
    "\n",
    "    z0delta = z0delta / torch.norm(z0delta,dim=1).view((numtrials,1))\n",
    "\n",
    "    grad_delta_alignment = torch.sum(z0grad * z0delta, dim=1).detach().cpu()\n",
    "    plt.hist(grad_delta_alignment,bins=20)\n",
    "    plt.title('Dot product of grad with linear interpolation direction. latent_d =' + str(latent_d) )\n",
    "    plt.show()\n",
    "\n",
    "    plt.scatter(sample_distance, grad_delta_alignment,marker='.')\n",
    "    plt.xlabel('Distance of z0 to z1')\n",
    "    plt.ylabel('Alignment of v = z1 - z0 vs. w = direction of gradient')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uTavzcKaTS_y"
   },
   "source": [
    "# Greedy geodesic (start at z0 and follow the gradient so that g(z0) reaches g(z1))\n",
    "\n",
    "The takeaway that I got from my limited experiments here is that straight-line interpolation is better than the path found by the greedy geodesic algorithm.\n",
    "\n",
    "However, I only ran two trials, with d = 5, so this should be taken with a grain of salt.\n",
    "\n",
    "Also, this code is not very well written since I was jumping around, so it will not run if executed in a straight line."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "C5_r0daSTxdN"
   },
   "outputs": [],
   "source": [
    "for latent_d in [5]:\n",
    "  numtrials = 1\n",
    "  numsteps = 10000\n",
    "  step_size = 0.001\n",
    "  z0 = torch.randn([numtrials, model.latent_d]).to(device)\n",
    "  z1 = torch.randn([numtrials, model.latent_d]).to(device)\n",
    "\n",
    "  ptlist = []\n",
    "  currz = torch.clone(z0)\n",
    "\n",
    "  currz.requires_grad = True\n",
    "  x1 = model.decode(z1).detach()\n",
    "  optimizer = optim.Adam([currz], lr=0.0001)\n",
    "  for step in range(numsteps):\n",
    "    ptlist.append(torch.clone(currz))\n",
    "    currx = model.decode(currz)\n",
    "    loss = torch.sum((currx - x1)**2)\n",
    "    print(loss)\n",
    "    if loss < 0.02:\n",
    "      break\n",
    "\n",
    "    loss.backward()\n",
    "    # print(torch.norm(currz.grad,dim=1))\n",
    "    # print(currz.grad)\n",
    "    with torch.no_grad():\n",
    "      graddir = currz.grad / torch.norm(currz.grad, dim=1).view(numtrials, 1)\n",
    "      currz.add_(-graddir, alpha=step_size)\n",
    "\n",
    "      # print(graddir)\n",
    "\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "    # z0grad = z0grad / torch.norm(z0grad,dim=1).view((numtrials,1))\n",
    "    # print(currz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oFfjyTufagQQ"
   },
   "outputs": [],
   "source": [
    "ptlist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3sIRb8VRT22i"
   },
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "def minimize_path_energy(model, z0, z1, init_path=None, numpts=10,numtrainiter=1000,learning_rate=0.001,verbose=False):\n",
    "  \"\"\"\n",
    "  Initialization: a linear interpolating path between z0 and z1 with numpts number of points.\n",
    "  Output: a path with numpts between z0 and z1 trained to minimize energy with respect to the model.\n",
    "  \"\"\"\n",
    "\n",
    "  if init_path is None:\n",
    "    # Initialize the interpolation between z0 and z1 with a linear interpolation.\n",
    "    t = torch.linspace(1 / (numpts + 1), numpts / (numpts + 1), numpts).to(device)\n",
    "    interp_points = torch.outer(1 - t, torch.flatten(z0)) + torch.outer(t, torch.flatten(z1))\n",
    "  else:\n",
    "    interp_points = init_path\n",
    "  interp_points.requires_grad = True\n",
    "\n",
    "  # For convenience, precompute x0 and x1.\n",
    "  x0 = model.decode(z0).detach()\n",
    "  x1 = model.decode(z1).detach()\n",
    "\n",
    "  optimizer = optim.Adam([interp_points], lr=learning_rate)\n",
    "  for i in range(numtrainiter):\n",
    "    imgs = model.decode(interp_points)\n",
    "\n",
    "    # compute the energy of the path\n",
    "    loss = torch.sum((torch.cat((x0, imgs)) - torch.cat((imgs, x1)))**2)\n",
    "\n",
    "    if verbose:\n",
    "      with torch.no_grad():\n",
    "        if i % 1000 == 0:\n",
    "          print('loss',i,loss.to('cpu'))\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "  imgs = model.decode(interp_points)\n",
    "  loss = torch.sum((torch.cat((x0, imgs)) - torch.cat((imgs, x1)))**2)\n",
    "  return torch.cat((z0, interp_points.detach(), z1)), loss.detach()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iG5VSnIWbAE_"
   },
   "outputs": [],
   "source": [
    "path, loss = minimize_path_energy(model,z0,z1,numpts=4852,numtrainiter=10000,learning_rate=0.0001,verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iMK0IwOKdInb"
   },
   "outputs": [],
   "source": [
    "compute_path_length(model, path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eZ0rvzH2cVtb"
   },
   "outputs": [],
   "source": [
    "def compute_path_length(model,path):\n",
    "  l = 0\n",
    "  for i in range(path.shape[0]-1):\n",
    "    xip1 = model.decode(path[i+1,:])\n",
    "    xi = model.decode(path[i,:])\n",
    "    l += torch.norm(xip1 - xi)\n",
    "  return l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HXSzo1NBeKXd"
   },
   "outputs": [],
   "source": [
    "len(ptlist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "75zNFCRtcyZm"
   },
   "outputs": [],
   "source": [
    "compute_path_length(model,torch.cat(ptlist).detach())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "od9KorF_bN7B"
   },
   "outputs": [],
   "source": [
    "minimize_path_energy(model,z0,z1,init_path=torch.cat(ptlist).detach(), numpts=numsteps,numtrainiter=20000,verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xHXngPyYfzTq"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZHeHI11wZc_j"
   },
   "source": [
    "# Deprecated code: Geodesics of networks at initialization using analytic formula.\n",
    "\n",
    "This runs into some errors in the calculation of the derivative of the arccosine, and the analytic formula does not account for the random biases. I found it better to simply use the 'untrained' network instead.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "sNSqEqSpZgS1"
   },
   "outputs": [],
   "source": [
    "#@title Deprecated code: analytic distance calculation\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "def get_pairwise_energy_random_net(z1,z2,L):\n",
    "  norm1 = torch.sum(torch.square(z1))\n",
    "  norm2 = torch.sum(torch.square(z2))\n",
    "  dotprod = torch.sum(z1 * z2)\n",
    "\n",
    "  normprod = torch.sqrt(norm1 * norm2)\n",
    "  latentangle = torch.acos(dotprod / normprod)\n",
    "\n",
    "  currangle = latentangle\n",
    "  for i in range(L):\n",
    "    currangle = torch.acos(((np.pi - currangle) * torch.cos(currangle) + torch.sin(currangle))/np.pi)\n",
    "  cosgenangle = torch.cos(currangle)\n",
    "  # cosgenangle = ((np.pi - latentangle) * torch.cos(latentangle) + torch.sin(latentangle))/np.pi\n",
    "  # cosgenangle = torch.cos(latentangle)\n",
    "\n",
    "  energy = norm1 + norm2 - 2 * normprod * cosgenangle\n",
    "\n",
    "  return energy\n",
    "  # print(energy)\n",
    "  # length = torch.sqrt(energy)\n",
    "  # length = energy ** (1/2)\n",
    "  # print('Returning length')\n",
    "  # # print(length)\n",
    "  # # print(norm1, norm2, dotprod)\n",
    "  # # print(energy)\n",
    "  # # return torch.sqrt(energy)\n",
    "  # return length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "lXyEu4uwXEIW"
   },
   "outputs": [],
   "source": [
    "#@title Deprecated code: minimize path energy random net\n",
    "\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "def minimize_path_energy_random_net(z0, z1, L, numpts=10,numtrainiter=1000,learning_rate=0.001,verbose=False):\n",
    "  \"\"\"\n",
    "  Initialization: a linear interpolating path between z0 and z1 with numpts number of points.\n",
    "  Output: a path with numpts between z0 and z1 trained to minimize energy with respect to the model.\n",
    "  \"\"\"\n",
    "\n",
    "  # Initialize the interpolation between z0 and z1 with a linear interpolation.\n",
    "  t = torch.linspace(1 / (numpts + 1), numpts / (numpts + 1), numpts).to(device)\n",
    "  interp_points = torch.outer(1 - t, torch.flatten(z0)) + torch.outer(t, torch.flatten(z1))\n",
    "  # interp_points = torch.rand_like(z0)\n",
    "  interp_points.requires_grad = True\n",
    "\n",
    "  def totloss(interp_points):\n",
    "    # compute the energy of the path\n",
    "    loss = get_pairwise_energy_random_net(z0,interp_points[0,:],L)\n",
    "    for i in range(numpts-1):\n",
    "      loss += get_pairwise_energy_random_net(interp_points[i,:],interp_points[i+1,:],L)\n",
    "    loss += get_pairwise_energy_random_net(interp_points[numpts-1,:],z1,L)\n",
    "    return loss\n",
    "\n",
    "  optimizer = optim.Adam([interp_points], lr=learning_rate)\n",
    "  for i in range(numtrainiter):\n",
    "    imgs = model.decode(interp_points)\n",
    "\n",
    "    loss = totloss(interp_points)\n",
    "\n",
    "    # print(loss)\n",
    "\n",
    "\n",
    "    if verbose:\n",
    "      with torch.no_grad():\n",
    "        if i % 100 == 0:\n",
    "          print('loss',i,loss.to('cpu'))\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "  loss = totloss(interp_points)\n",
    "  return torch.cat((z0, interp_points.detach(), z1)), loss.detach()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "9qPJ1QU7ZgS2"
   },
   "outputs": [],
   "source": [
    "#@title Deprecated: optimize paths using analytic distance for random net\n",
    "\n",
    "import math\n",
    "for latent_d in [5]:\n",
    "  # for dist_parameter in [None, 0.5, 2, 4, 8, 16]:\n",
    "  for dist_parameter in [None]:\n",
    "    model = get_trained_vae('mnist',latent_d)\n",
    "\n",
    "    numtrials = 1\n",
    "    assert(numtrials == 1)\n",
    "\n",
    "    if dist_parameter is None:\n",
    "      z0 = torch.randn([numtrials, model.latent_d]).to(device)\n",
    "      z1 = torch.randn([numtrials, model.latent_d]).to(device)\n",
    "\n",
    "    else:\n",
    "      diff = torch.randn([numtrials, model.latent_d]).to(device) * math.sqrt(2)\n",
    "      diff = diff / torch.norm(diff,dim=1).view((numtrials,1)) * dist_parameter\n",
    "      z0 = torch.randn([numtrials, model.latent_d]).to(device) * math.sqrt(1/2) + diff / 2\n",
    "      z1 = z0 + diff\n",
    "\n",
    "    print(get_pairwise_energy_random_net(z0,z1,L=2))\n",
    "    torch.autograd.set_detect_anomaly(True)\n",
    "    outputs = minimize_path_energy_random_net(z0,z1,L=1,numpts=10,learning_rate=0.01,verbose=True)\n",
    "    print(outputs)\n",
    "\n",
    "    # z0grad = grad_gz0_to_gz1(model, z0, z1)\n",
    "    # z0delta = z1 - z0\n",
    "    # sample_distance = torch.norm(z0delta,dim=1).detach().cpu()\n",
    "\n",
    "    # z0grad = z0grad / torch.norm(z0grad,dim=1).view((numtrials,1))\n",
    "\n",
    "    # z0delta = z0delta / torch.norm(z0delta,dim=1).view((numtrials,1))\n",
    "\n",
    "    # grad_delta_alignment = torch.sum(z0grad * z0delta, dim=1).detach().cpu()\n",
    "    # plt.hist(grad_delta_alignment,bins=20)\n",
    "    # plt.title('Dot product of grad with linear interpolation direction. latent_d =' + str(latent_d) )\n",
    "    # plt.show()\n",
    "\n",
    "    # plt.scatter(sample_distance, grad_delta_alignment,marker='.')\n",
    "    # plt.xlabel('Distance of z0 to z1')\n",
    "    # plt.ylabel('Alignment of v = z1 - z0 vs. w = direction of gradient')\n",
    "    # plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nBOfwfXLVPEX"
   },
   "outputs": [],
   "source": [
    "# George Stepaniants Code: Computing Jacobian Statistics at Sample Points on Manifold\n",
    "model_example = get_trained_vae('untrained', 20)\n",
    "\n",
    "trials = 1000\n",
    "conds = []\n",
    "lmaxs = []\n",
    "lmins = []\n",
    "for i in range(trials):\n",
    "  print(i)\n",
    "  z = torch.randn([model_example.latent_d]).to(device)\n",
    "  Jg = torch.autograd.functional.jacobian(model_example.decode, z)\n",
    "  G = torch.matmul(Jg.t(), Jg)\n",
    "\n",
    "  cond = torch.linalg.cond(G).item()\n",
    "  lmax = torch.lobpcg(G)[0].item()\n",
    "  lmin = torch.lobpcg(G, largest=False)[0].item()\n",
    "\n",
    "  conds.append(cond)\n",
    "  lmaxs.append(lmax)\n",
    "  lmins.append(lmin)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VAtDTGV0W4zd"
   },
   "outputs": [],
   "source": [
    "plt.figure(1)\n",
    "plt.title('Condition Number of Jg^T * Jg')\n",
    "plt.hist(conds, bins=30, density=True)\n",
    "\n",
    "plt.figure(2)\n",
    "plt.title('Maximum Eigenvalue of Jg^T * Jg')\n",
    "plt.hist(lmaxs, bins=30, density=True)\n",
    "\n",
    "plt.figure(3)\n",
    "plt.title('Minimum Eigenvalue of Jg^T * Jg')\n",
    "plt.hist(lmins, bins=30, density=True)\n",
    "\n",
    "plt.figure(4)\n",
    "plt.title('Minimum vs. Maximum Eigenvalues')\n",
    "plt.scatter(lmins, lmaxs)\n",
    "lims = [np.min([plt.xlim(), plt.ylim()]), np.max([plt.xlim(), plt.ylim()])]\n",
    "plt.plot(lims, lims, 'r-', alpha=0.75, zorder=0)\n",
    "plt.xlabel('Minimum Eigenvalue')\n",
    "plt.ylabel('Maximum Eigenvalue')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "37-HpwbtDd75"
   },
   "outputs": [],
   "source": [
    "print(np.mean(conds))\n",
    "import statistics\n",
    "print(statistics.median(conds))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lafZ0_G4ES-6"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "geometric_toolkit.ipynb",
   "private_outputs": true,
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}