{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": true
   },
   "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 pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as pat\n",
    "import scipy.linalg as sl\n",
    "import scipy.stats as ss\n",
    "import scipy.spatial as ssp\n",
    "import numpy as np\n",
    "import math\n",
    "import pickle\n",
    "\n",
    "from copy import deepcopy\n",
    "import itertools\n",
    "import tqdm\n",
    "\n",
    "# additionally requires libcpab (https://github.com/SkafteNicki/libcpab)\n",
    "import nwarp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def augment_cps(cps):\n",
    "    return np.concatenate(([0], cps, [T])).astype(int)\n",
    "\n",
    "def hausdorff(cps1, cps2):\n",
    "    \"\"\"The provided cps must NOT include 0 and T.\"\"\"\n",
    "    return max(\n",
    "        ssp.distance.directed_hausdorff(cps1.reshape(-1, 1), cps2.reshape(-1, 1))[0],\n",
    "        ssp.distance.directed_hausdorff(cps2.reshape(-1, 1), cps1.reshape(-1, 1))[0]\n",
    "    )\n",
    "\n",
    "def compute_frobenius_matrix(cps):\n",
    "    \"\"\"The provided cps MUST include 0 and T.\"\"\"\n",
    "    T = cps[-1]\n",
    "    timestep_to_segment = (\n",
    "        (cps[:-1] <= np.tile(np.linspace(0, T-1, T).reshape(-1,1), len(cps)-1)) &\n",
    "        (np.tile(np.linspace(0, T-1, T).reshape(-1,1), len(cps)-1) < cps[1:])\n",
    "    )\n",
    "    timesteps_in_same_segment = np.matmul(timestep_to_segment, timestep_to_segment.transpose())\n",
    "    frobenius = timestep_to_segment/np.sqrt(np.tile(timestep_to_segment.sum(axis=0).reshape(-1, 1), T).transpose())\n",
    "    frobenius = np.matmul(frobenius, frobenius.transpose())\n",
    "    return frobenius\n",
    "\n",
    "def frobenius(cps1, cps2, m1=None, m2=None):\n",
    "    \"\"\"The provided cps must NOT include 0 and T.\"\"\"\n",
    "    if m1 is None:\n",
    "        m1 = compute_frobenius_matrix(augment_cps(cps1))\n",
    "    if m2 is None:\n",
    "        m2 = compute_frobenius_matrix(augment_cps(cps2))\n",
    "    return sl.norm(m1 - m2, 'fro')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def iidsample_with_cps(cps, dists):\n",
    "    #for d in dists:\n",
    "    #    print(f'{d.mean():5.2f}\\t{d.var():5.2f}')\n",
    "    sample = np.empty(cps[-1])\n",
    "    olddist = None\n",
    "    dist = np.random.choice(dists)\n",
    "    for k in range(len(cps)-1):\n",
    "        while olddist == dist:\n",
    "            dist = np.random.choice(dists)\n",
    "        sample[cps[k]:cps[k+1]] = dist.rvs(cps[k+1]-cps[k])\n",
    "        olddist = dist\n",
    "    return sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 18392 characters]\n",
      "text/plain": [
       "<Figure size 1080x144 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Arlot2019 scenario 1\n",
    "cps = np.array([0, 100, 130, 220, 320, 370, 520, 620, 740, 790, 870, 1000], dtype=int)\n",
    "dists = [\n",
    "    ss.binom(10, 0.2),\n",
    "    ss.nbinom(3, 0.7),\n",
    "    ss.hypergeom(10, 5, 2),\n",
    "    ss.norm(2.5, scale=np.sqrt(0.25)),\n",
    "    ss.gamma(0.5, scale=5),\n",
    "    ss.weibull_min(2, scale=5),\n",
    "    ss.pareto(3, scale=1.5)\n",
    "]\n",
    "\n",
    "T = cps[-1]\n",
    "K = len(cps)-1\n",
    "\n",
    "sample = iidsample_with_cps(cps, dists)\n",
    "\n",
    "plt.figure(figsize=(15,2))\n",
    "plt.scatter(range(cps[-1]), sample)\n",
    "for cp in cps[1:-1]:\n",
    "    plt.axvline(cp, lw=1, color='black')\n",
    "plt.xlim(0, cps[-1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Segmented model estimation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# PIECEWISE GAUSSIAN MODEL (sensitive to changes in mean and variance)\n",
    "# TSP-based parameter warping with constant mode vector\n",
    "\n",
    "Dparamg = 0 # number of global parameters (same in every segment)\n",
    "Dparaml = 2 # number of local parameters (different in every segment)\n",
    "\n",
    "paramwarpTSPB = nwarp.ParameterWarp(\n",
    "    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",
    "    verbose=False\n",
    ")\n",
    "\n",
    "paramwarpVanilla = nwarp.ParameterWarp(\n",
    "    K, Dparamg, Dparaml,\n",
    "    nwarp.VanillaWarp(nwarp.Constant((T-1,)), nonneg_trans='square'),\n",
    "    verbose=False\n",
    ")\n",
    "\n",
    "paramwarpCPAB = nwarp.ParameterWarp(\n",
    "    K, Dparamg, Dparaml,\n",
    "    nwarp.CPABWarp(K+1, nwarp.Constant((K,))),\n",
    "    verbose=False\n",
    ")\n",
    "\n",
    "model = nwarp.NormalDistribution()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training and Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_train = torch.Tensor(iidsample_with_cps(cps, dists)).reshape(1, -1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NormalDistribution()\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "ParameterWarp(\n",
       "  (warp): TSPStepWarp(\n",
       "    (loc_net): Constant()\n",
       "  )\n",
       "  (resample): Resample()\n",
       ")"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "n_restarts    =  10 # number of randomized restarts\n",
    "n_epochs      = 300 # total number of epochs per restart\n",
    "n_epochs_hard = 100 # use hard segmentation for the last X epochs\n",
    "lr            = 1e-1\n",
    "loss_fn       = nn.MSELoss()\n",
    "paramwarp     = paramwarpTSPB\n",
    "show_plots    = True\n",
    "\n",
    "print(model)\n",
    "paramwarp.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "restart  1/10: 100%|██████████| 300/300 [00:01<00:00, 177.31it/s, max=4.4681, cur=1.4688]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 35984 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:01<00:00, 184.33it/s, max=4.5742, cur=1.4290]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 34512 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:01<00:00, 175.82it/s, max=3.4459, cur=1.5621]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 35324 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:01<00:00, 185.99it/s, max=3.8600, cur=1.5091]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 33728 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:01<00:00, 184.00it/s, max=4.9252, cur=1.4925]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 36964 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:01<00:00, 175.79it/s, max=4.1195, cur=1.5470]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 30760 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:01<00:00, 175.41it/s, max=3.9629, cur=1.4408]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 37696 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:01<00:00, 161.43it/s, max=3.8070, cur=1.4422]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 35688 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:01<00:00, 181.69it/s, max=3.6206, cur=1.4335]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 35272 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:01<00:00, 187.42it/s, max=4.5126, cur=1.4492]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 36596 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 = 1.4290\n",
      "mean loss = 1.4774 +-0.0458\n"
     ]
    }
   ],
   "source": [
    "best_loss = np.inf\n",
    "loss_curves = np.empty((n_restarts, n_epochs))\n",
    "for r in range(n_restarts):\n",
    "    # reset everything\n",
    "    optimizer = torch.optim.Adam(paramwarp.parameters(), lr=lr)\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",
    "    _ = paramwarp.apply(nwarp.reset_parameters)\n",
    "\n",
    "    # perform training\n",
    "    paramwarp.train()\n",
    "    for epoch in epoch_counter:\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        if epoch == n_epochs - n_epochs_hard:\n",
    "            optimizer = torch.optim.Adam(paramwarp.parameters(), lr=lr)\n",
    "            resample_kernel = \"integer\"\n",
    "\n",
    "        param_hat_train = paramwarp(x_train, resample_kernel=resample_kernel)\n",
    "        nllloss = -model(x_train, param_hat_train).mean()\n",
    "        loss = nllloss\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",
    "    loss_curves[r,:] = train_losses\n",
    "    if train_losses[-1] < best_loss:\n",
    "        best_paramwarp_state = deepcopy(paramwarp.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",
    "paramwarp.load_state_dict(best_paramwarp_state)\n",
    "print(f\"best loss = {best_loss:.4f}\")\n",
    "print(f\"mean loss = {loss_curves[:,-1].mean():.4f} +-{loss_curves[:,-1].std():.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 35432 characters]\n",
      "text/plain": [
       "<Figure size 1080x144 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "param_hat_train = paramwarp(x_train, resample_kernel=resample_kernel)\n",
    "cps_hat = paramwarp._almat.sum(dim=1).cumsum(dim=1).squeeze()[:-1]\n",
    "\n",
    "plt.figure(figsize=(15,2))\n",
    "plt.xlim(0,T)\n",
    "plt.plot(x_train[0,:,0], alpha=1, zorder=-1)\n",
    "plt.fill_between(x=range(T),\n",
    "                 y1=param_hat_train[0,:,0].detach()-2*torch.sqrt(torch.exp(param_hat_train[0,:,1].detach())),\n",
    "                 y2=param_hat_train[0,:,0].detach()+2*torch.sqrt(torch.exp(param_hat_train[0,:,1].detach())),\n",
    "                color='orange', alpha=0.5)\n",
    "plt.plot(param_hat_train[0,:,0].detach(), color='orange', lw=3)\n",
    "for cp in cps_hat:\n",
    "    plt.axvline(cp, lw=3, color='black', alpha=0.25)\n",
    "for cp in cps:\n",
    "    plt.axvline(cp, lw=1, color='black')\n",
    "plt.title('mean +- 2 sigma')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "d_hdf = 103\n",
      "d_fro = 2.54\n"
     ]
    }
   ],
   "source": [
    "print(f'd_hdf = {hausdorff(cps[1:-1], cps_hat):.0f}')\n",
    "print(f'd_fro = {frobenius(cps[1:-1], cps_hat):.2f}')"
   ]
  },
  {
   "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": 2
}
