{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matlab.engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('..')\n",
    "import lipMIP as lm\n",
    "import neural_nets.data_loaders as dl\n",
    "import neural_nets.train as nn_train\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import math\n",
    "from relu_nets import ReLUNet\n",
    "import lipMIP as lm\n",
    "import torch\n",
    "from hyperbox import Hyperbox\n",
    "from pre_activation_bounds import PreactivationBounds\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from other_methods.lip_sdp import LipSDPCustom\n",
    "from other_methods.fast_lip import FastLip\n",
    "from other_methods.naive_methods import NaiveUB, RandomLB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 13404 characters]\n",
      "text/plain": [
       "<Figure size 576x576 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# First make a simple randomK parameter dataset\n",
    "#data_params = dl.RandomKParameters(10, 25, radius=0.02)\n",
    "data_params = dl.EricParameters(25, 0.05)\n",
    "dataset = dl.Random2DBinaryDataset(data_params, batch_size=32, random_seed=420)\n",
    "dataset.plot_2d()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# And then make a simple network to classify it \n",
    "simple_net = ReLUNet(layer_sizes=(2, 8, 16, 16, 16, 2), bias=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# And train this thing:\n",
    "trainset, valset = dataset.split_train_val(1.0)\n",
    "train_params = nn_train.TrainParameters(trainset, trainset, 1000, test_after_epoch=200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nn_train.training_loop(simple_net, train_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = simple_net.display_decision_bounds([0, 1], [0,1], 100)\n",
    "dataset.plot_2d(ax=ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### PREACTIVATION CHECK \n",
    "small_hbox = Hyperbox.build_linf_ball(np.array([0.5, 0.5]), 0.2)\n",
    "x = torch.Tensor([0.5, 0.5])\n",
    "bounds = PreactivationBounds.naive_ia_from_hyperbox(simple_net, small_hbox)\n",
    "def check_preacts_within_bounds(x):\n",
    "    preacts = simple_net(x, return_preacts=True)\n",
    "    for i, preact in enumerate(preacts):\n",
    "        preact = preact.data.numpy().reshape(-1)\n",
    "        ith_bound = bounds.get_ith_layer_bounds(i)\n",
    "        print(i, len(ith_bound[0]), sum(ith_bound[0] <= preact), sum(ith_bound[1] >= preact))\n",
    "        \n",
    "check_preacts_within_bounds(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To compute the Local Lipschitz Constant of a DNN, we need:\n",
    "# 1) A region to compute over \n",
    "# 2) A neural network to compute\n",
    "# 3) Function parameters {LP_norm, preactivation_techniques, etc}\n",
    "\n",
    "# LipParameters holds items (1, 3)\n",
    "# LipProblem holds items (1, 2, 3)\n",
    "# LipMIPResult holds all info from a completed lipschitz object\n",
    "\n",
    "\n",
    "# EvaluationParameters holds parameters for how to evaluate lipschitz\n",
    "# so let's do this later\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Goal: pick a random point and evaluate the 0.1 hyperbox lipschitz \n",
    "hbox = Hyperbox.build_linf_ball(np.array([0.5, 0.5]), 0.5)\n",
    "c_vector = torch.tensor([1.0, -1.0])\n",
    "\n",
    "params = lm.LipParameters(hbox, c_vector, verbose=False)\n",
    "problem = lm.LipProblem(simple_net, params)\n",
    "result = problem.compute_max_lipschitz()\n",
    "result\n",
    "print(result.value, result.best_x)\n",
    "print(simple_net.get_grad_at_point([0.3741, 0.5873], c_vector).norm(p=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# and do the other methods \n",
    "nub = NaiveUB(simple_net, c_vector, 'linf')\n",
    "nub.compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rlb = RandomLB(simple_net, c_vector, hbox, 'linf')\n",
    "rlb.compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "simple_net.random_max_grad(hbox, c_vector, 1000, pnorm=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "simple_net.get_grad_at_point(np.array([0.40050755, 0.6]), c_vector)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result.squire.get_grad_at_point(np.array([0.5005, 0.4886]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Simple idea: get lipschitz constants as we train\n",
    "new_net = simple_net = ReLUNet(layer_sizes=(2, 20, 20, 20, 2), bias=True)\n",
    "train_params = nn_train.TrainParameters(trainset, trainset, 50, test_after_epoch=50)\n",
    "done_epochs = 0 \n",
    "lip_results = []\n",
    "lip_params = lm.LipParameters(hbox, c_vector, verbose=False)\n",
    "while done_epochs < 1000:\n",
    "    nn_train.training_loop(new_net, train_params, epoch_start_no=done_epochs)\n",
    "    problem = lm.LipProblem(simple_net, lip_params)\n",
    "    lip_results.append(problem.compute_max_lipschitz())\n",
    "    done_epochs += train_params.num_epochs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[_.value for _ in lip_results]"
   ]
  },
  {
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
