{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Figure size 432x288 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<Figure size 432x288 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<Figure size 432x288 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.linalg as sl\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import arff\n",
    "from copy import deepcopy\n",
    "import tqdm\n",
    "\n",
    "# additionally requires libcpab (https://github.com/SkafteNicki/libcpab)\n",
    "import nwarp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(52848, 33)\n"
     ]
    }
   ],
   "source": [
    "# Please download the USP Data Stream repository from: https://sites.google.com/view/uspdsrepository\n",
    "path_to_usp = '../../../../../../Documents/data/usp-stream-data'\n",
    "\n",
    "dataset = 'INSECTS-abrupt_balanced'\n",
    "#dataset = 'INSECTS-incremental_balanced'\n",
    "#dataset = 'INSECTS-gradual_balanced'\n",
    "#dataset = 'INSECTS-incremental-abrupt_balanced'\n",
    "#dataset = 'INSECTS-incremental-reoccurring_balanced'\n",
    "\n",
    "data_arff = arff.load(f'{path_to_usp}/{dataset}_norm.arff')\n",
    "data = pd.DataFrame(list(data_arff))\n",
    "data_x = data[range(33)]\n",
    "data_y1hot = pd.get_dummies(data[33])\n",
    "print(data_x.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_train = torch.as_tensor(data_x.values.astype(np.float32)).unsqueeze(0)\n",
    "y_train1hot = torch.as_tensor(data_y1hot.values.astype(np.float32))\n",
    "y_train = y_train1hot.argmax(-1).squeeze()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sequential(\n",
      "  (0): Linear(in_features=33, out_features=8, bias=True)\n",
      "  (1): ReLU()\n",
      ")\n",
      "ParameterWarp(\n",
      "  (warp): TSPStepWarp(\n",
      "    (loc_net): Constant()\n",
      "  )\n",
      "  (resample): Resample()\n",
      ")\n",
      "GeneralizedLinearModel(\n",
      "  (invlink_func): Identity()\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "N, T, Din = x_train.shape\n",
    "_, Dout = y_train1hot.shape\n",
    "\n",
    "K = 6\n",
    "Dhidden = 8 # hidden dimension\n",
    "\n",
    "# softmax regression with cross entropy loss (softmax performed in loss function)\n",
    "glm = nwarp.GeneralizedLinearModel(Dout, invlink_func=nn.Identity())\n",
    "\n",
    "# TSP-based parameter warping with constant mode vector\n",
    "Dparamg = 0 # number of global parameters (same in every segment)\n",
    "Dparaml = Dhidden*Dout # number of local parameters (different in every segment)\n",
    "paramwarp = nwarp.ParameterWarp(K, Dparamg, Dparaml,\n",
    "               nwarp.TSPStepWarp(nwarp.Constant((K-1,)),\n",
    "                                 width=0.125, power=16.,\n",
    "                                 min_step=0.0001, max_step=0.9999))\n",
    "\n",
    "# feature transformation for the covariates\n",
    "# map shape (T, Din) to shape (T, Dhidden)\n",
    "covariates = nn.Sequential(\n",
    "    nn.Linear(Din, Dhidden),\n",
    "    nn.ReLU(),\n",
    ")\n",
    "\n",
    "print(covariates)\n",
    "print(paramwarp)\n",
    "print(glm)\n",
    "#glm(covariates(x_train), paramwarp(x_train)[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_restarts    =  10 # number of randomized restarts\n",
    "n_epochs      = 300 # total number of epochs\n",
    "n_epochs_hard = 100 # use hard segmentation for the last X epochs\n",
    "show_plots    = True\n",
    "loss_fn       = nn.CrossEntropyLoss(reduction='mean')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "restart  1/10: 100%|██████████| 300/300 [00:33<00:00,  9.02it/s, max=1.7896, cur=0.6174]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 42152 characters]\n",
      "text/plain": [
       "<Figure size 1080x144 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "restart  2/10: 100%|██████████| 300/300 [00:33<00:00,  8.97it/s, max=1.7920, cur=0.6425]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 37644 characters]\n",
      "text/plain": [
       "<Figure size 1080x144 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "restart  3/10: 100%|██████████| 300/300 [00:32<00:00,  9.16it/s, max=1.7923, cur=0.6617]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 37940 characters]\n",
      "text/plain": [
       "<Figure size 1080x144 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "restart  4/10: 100%|██████████| 300/300 [00:33<00:00,  8.90it/s, max=1.7947, cur=0.6918]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 45992 characters]\n",
      "text/plain": [
       "<Figure size 1080x144 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "restart  5/10: 100%|██████████| 300/300 [00:34<00:00,  8.73it/s, max=1.7928, cur=0.7052]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 41548 characters]\n",
      "text/plain": [
       "<Figure size 1080x144 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "restart  6/10: 100%|██████████| 300/300 [00:31<00:00,  9.54it/s, max=1.7911, cur=0.6514]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 43936 characters]\n",
      "text/plain": [
       "<Figure size 1080x144 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "restart  7/10: 100%|██████████| 300/300 [00:30<00:00,  9.72it/s, max=1.7902, cur=0.6142]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 40556 characters]\n",
      "text/plain": [
       "<Figure size 1080x144 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "restart  8/10: 100%|██████████| 300/300 [00:32<00:00,  9.21it/s, max=1.7905, cur=0.6437]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 37668 characters]\n",
      "text/plain": [
       "<Figure size 1080x144 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "restart  9/10: 100%|██████████| 300/300 [00:29<00:00, 10.04it/s, max=1.7916, cur=0.6741]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 40556 characters]\n",
      "text/plain": [
       "<Figure size 1080x144 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "restart 10/10: 100%|██████████| 300/300 [00:33<00:00,  8.97it/s, max=1.7891, cur=0.6498]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 37396 characters]\n",
      "text/plain": [
       "<Figure size 1080x144 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "best loss = 0.6142\n"
     ]
    }
   ],
   "source": [
    "best_loss = np.inf\n",
    "for r in range(n_restarts):\n",
    "    # reset everything\n",
    "    optimizer = torch.optim.Adam([\n",
    "                        {'params': paramwarp.parameters(), 'lr': 1e-1},\n",
    "                        {'params': covariates.parameters(), 'lr': 1e-1}\n",
    "                ], weight_decay=0.0)\n",
    "    param_norm = []\n",
    "    grad_norm = []\n",
    "    train_losses = []\n",
    "    resample_kernel = 'linear'\n",
    "    epoch_counter = tqdm.tqdm(range(n_epochs), desc=f'restart {(r+1):2d}/{n_restarts:2d}')\n",
    "\n",
    "    # initialize parameters\n",
    "    _ = covariates.apply(nwarp.reset_parameters)\n",
    "    _ = paramwarp.apply(nwarp.reset_parameters)\n",
    "    nn.init.uniform_(paramwarp.warp.loc_net.const, -1., 0.) # segmentation\n",
    "\n",
    "    # perform training\n",
    "    paramwarp.train()\n",
    "    covariates.train()    \n",
    "    for epoch in epoch_counter:\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        if epoch == n_epochs - n_epochs_hard:\n",
    "            resample_kernel = 'integer'\n",
    "\n",
    "        param_hat_train = paramwarp(x_train, # input is ignored, but must have shape (N, T, Din)\n",
    "                              resample_kernel=resample_kernel)[0]\n",
    "        y_hat_train = glm(covariates(x_train), param_hat_train)\n",
    "\n",
    "        loss = loss_fn(y_hat_train.squeeze(), y_train)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        train_losses.append(loss.item())\n",
    "        param_norm.append([sl.norm(p.detach()) for p in paramwarp.parameters() if len(p)>0])\n",
    "        grad_norm.append([sl.norm(p.grad.detach()) for p in paramwarp.parameters() if len(p)>0])\n",
    "        epoch_counter.set_postfix({'max': f'{max(train_losses):.4f}', 'cur': f'{loss.item():.4f}'})\n",
    "\n",
    "    if train_losses[-1] < best_loss:\n",
    "        best_paramwarp_state = deepcopy(paramwarp.state_dict())\n",
    "        best_covariates_state = deepcopy(covariates.state_dict())\n",
    "        best_loss = train_losses[-1]\n",
    "\n",
    "    if show_plots:\n",
    "        plt.figure(figsize=(15,2))\n",
    "        plt.subplot(131)\n",
    "        plt.title('loss')\n",
    "        plt.ylim(np.min(train_losses), np.percentile(train_losses, 95))\n",
    "        plt.plot(train_losses)\n",
    "\n",
    "        plt.subplot(132)\n",
    "        plt.title('parameter norm')\n",
    "        lines = plt.plot(np.array(param_norm)/np.array(param_norm).max(axis=0))\n",
    "        plt.legend(lines, [' x '.join([str(d) for d in p.size()]) for p in paramwarp.parameters() if len(p)>0])\n",
    "\n",
    "        plt.subplot(133)\n",
    "        plt.title('gradient norm')\n",
    "        normalized_grad_norm = np.array(grad_norm)/np.array(grad_norm).max(axis=0)\n",
    "        lines = plt.plot(normalized_grad_norm)\n",
    "        plt.legend(lines, [' x '.join([str(d) for d in p.size()]) for p in paramwarp.parameters() if len(p)>0])\n",
    "        plt.ylim(np.min(normalized_grad_norm), np.percentile(normalized_grad_norm, 95))\n",
    "\n",
    "        plt.show()\n",
    "\n",
    "paramwarp.eval()\n",
    "covariates.eval()\n",
    "paramwarp.load_state_dict(best_paramwarp_state)\n",
    "covariates.load_state_dict(best_covariates_state)\n",
    "print(f'best loss = {best_loss:.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "param_hat_train, almat_hat_train, gamma_hat_train = paramwarp(x_train, # input is ignored, but must have shape (N, T, Din)\n",
    "                                                              resample_kernel=resample_kernel)\n",
    "y_hat_train = glm(covariates(x_train), param_hat_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INSECTS-abrupt_balanced\n",
      " 6 0.74 \n",
      "cps = 12748 18068 28104 38266 46879 \n"
     ]
    }
   ],
   "source": [
    "print(dataset)\n",
    "print(f'{K:2.0f}',\n",
    "      f'{torch.sum(y_hat_train.argmax(-1) == y_train).item()/T:.2f}',\n",
    "      end=' ')\n",
    "print()\n",
    "print('cps =', end=' ')\n",
    "for cp in almat_hat_train.sum(dim=1).cumsum(dim=1).squeeze()[:-1]:\n",
    "    print(f'{cp.item():.0f}', end=' ')\n",
    "print()"
   ]
  },
  {
   "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
