{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "66a26b46-f935-46c6-a91f-063fea50ffa1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import pickle\n",
    "import itertools\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "import scipy\n",
    "from scipy import special, integrate, optimize, stats\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "import gaussian_preact\n",
    "from gaussian_preact import ActivationFunction, ActivationFunctionPos, sampler_normal, build_sampler_weibull\n",
    "from gaussian_preact.post_process import build_color, build_label, build_linestyle\n",
    "\n",
    "torch.set_default_dtype(torch.float64)\n",
    "\n",
    "plt.rcParams['text.usetex'] = 'True'\n",
    "plt.rcParams['text.latex.preamble'] = r'\\usepackage{lmodern}'\n",
    "\n",
    "matplotlib.rcParams['mathtext.fontset'] = 'stix'\n",
    "matplotlib.rcParams['font.family'] = 'STIXGeneral'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "adc2e810-f4e3-4782-ab56-d83dd49520d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "matplotlib.rcParams['figure.figsize'] = [10, 6.7]\n",
    "path_objects = './objects/'\n",
    "\n",
    "#lst_theta = [2.05, 2.1, 2.2, 2.5, 3., 4., 5., 6., 7., 8., 9., 10.]    # extensive\n",
    "lst_theta = [2.05, 2.50, 3.00, 4.00, 5.00, 7.00, 10.00]    # as in the paper\n",
    "#lst_theta = [2.11]    # for testing\n",
    "lst_act_types = ['odd', 'pos']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83ccbb52-f4d1-4abe-98fd-b6cb51e39f4f",
   "metadata": {},
   "source": [
    "# Preliminary loads and plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3d2ddd00-5246-4478-b9fb-ba73dca91051",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the activation functions\n",
    "dct_activations = {}\n",
    "for theta in lst_theta:\n",
    "    dct = torch.load(path_objects + 'ActFunction_theta-{:.2f}.pkl'.format(theta), weights_only = False)\n",
    "    dct_activations[('odd', theta)] = dct['activation']\n",
    "    \n",
    "    dct = torch.load(path_objects + 'ActFunctionPos_theta-{:.2f}.pkl'.format(theta), weights_only = False)\n",
    "    dct_activations[('pos', theta)] = dct['activation']\n",
    "dct_activations[\"relu\"] = torch.relu\n",
    "dct_activations[\"tanh\"] = torch.tanh\n",
    "\n",
    "# Create initialization distributions\n",
    "dct_samplers = {}\n",
    "for theta in lst_theta:\n",
    "    dct_samplers[theta] = build_sampler_weibull(theta)\n",
    "dct_samplers[\"relu\"] = sampler_normal\n",
    "dct_samplers[\"tanh\"] = sampler_normal\n",
    "\n",
    "# Create sigma_w\n",
    "dct_sigma_w = {\"relu\": 1.4142135623730951, \"tanh\": 1.2083045973594573}\n",
    "dct_sigma_b = {\"relu\": 0., \"tanh\": 0.11401754250991379}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bb6582d5-a21b-4f7a-b443-dfe37f02654b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 52296 characters]",
      "text/plain": [
       "<Figure size 1000x670 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "matplotlib.rcParams['figure.figsize'] = [10, 6.7]\n",
    "\n",
    "tens_x = torch.linspace(-10, 10, 400)\n",
    "\n",
    "# Plot the activation functions\n",
    "for theta in lst_theta:\n",
    "    activation = dct_activations['pos', theta]\n",
    "    plt.plot(tens_x, activation.forward_ng(tens_x), label = r'$\\theta = {:.2f}$'.format(theta),\n",
    "            c = build_color({'act_function': 'weibull', 'act_theta': theta}))\n",
    "\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "plt.xlim(-10, 10)\n",
    "plt.ylim(-10, 10)\n",
    "plt.title('Positive activation functions')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f7571b53-bc69-4f04-9d12-e702abf643e8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 79612 characters]",
      "text/plain": [
       "<Figure size 1000x670 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "matplotlib.rcParams['figure.figsize'] = [10, 6.7]\n",
    "\n",
    "tens_x = torch.linspace(-10, 10, 400)\n",
    "\n",
    "# Plot the activation functions\n",
    "for theta in lst_theta:\n",
    "    activation = dct_activations['odd', theta]\n",
    "    plt.plot(tens_x, activation.forward_ng(tens_x), label = r'$\\theta = {:.2f}$'.format(theta),\n",
    "            c = build_color({'act_function': 'weibull', 'act_theta': theta}))\n",
    "\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "plt.xlim(-10, 10)\n",
    "plt.ylim(-10, 10)\n",
    "plt.title('Odd activation functions')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a39bf9d1-bbef-4d07-99ab-8927d35ccf28",
   "metadata": {},
   "source": [
    "# Test gaussianity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b6230d57-63b6-4d4f-9524-bd75616f7329",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 33984 characters]",
      "text/plain": [
       "<Figure size 1000x670 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Rapid test of gaussianity in a simple case: \n",
    "# Z = W * phi(X), where X is Gaussian\n",
    "\n",
    "matplotlib.rcParams['figure.figsize'] = [10, 6.7]\n",
    "\n",
    "n_samples = 100000\n",
    "theta = 10\n",
    "\n",
    "activation = dct_activations['odd', theta]\n",
    "sampler_weibull = dct_samplers[theta]\n",
    "\n",
    "z_in = torch.randn(n_samples)\n",
    "x = activation(z_in)\n",
    "w = sampler_weibull([n_samples])\n",
    "z_out = x * w\n",
    "\n",
    "_ = stats.probplot(z_out, dist = 'norm', plot = plt)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d48b413-6026-4308-bd44-b2f3683cf719",
   "metadata": {},
   "source": [
    "## KS-statistic for a 2-layer neural network\n",
    "\n",
    "*Minimal model with dependence between pre-activations.*\n",
    "\n",
    "TODO: documentation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "081768c4-8907-49c0-8f4e-48907fce5c5d",
   "metadata": {},
   "source": [
    "### Computation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "4e03470d-151b-400d-a2c5-179eb0fdd314",
   "metadata": {},
   "outputs": [],
   "source": [
    "def propagate(x, out_dim, sampler, sigma_w = 1., sigma_b = 0.):\n",
    "    w = sampler([x.size(0), out_dim, x.size(1)], device = x.device, dtype = x.dtype)\n",
    "    w.mul_(sigma_w)\n",
    "    ret = torch.bmm(w, x.unsqueeze(2)).squeeze() / np.sqrt(x.size(1))\n",
    "    if ret.dim() == 1:\n",
    "        ret = ret.unsqueeze(1)\n",
    "    if sigma_b != 0.:\n",
    "        ret += sigma_b * sampler([x.size(0), out_dim], device = x.device, dtype = x.dtype)\n",
    "    \n",
    "    return ret\n",
    "\n",
    "def full_propagation(x, lst_nl, sampler, activation, sigma_w = 1., sigma_b = 0.):\n",
    "    for i, nl in enumerate(lst_nl):\n",
    "        if i == 0:\n",
    "            sampler_ = sampler_normal\n",
    "        else:\n",
    "            sampler_ = sampler\n",
    "        z = propagate(x, nl, sampler_, sigma_w, sigma_b)\n",
    "        x = activation(z)\n",
    "    \n",
    "    return z[:,0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "68409271-3d5a-4d50-8861-ab3e63f01c50",
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma_w = 1.\n",
    "n_samples = 200000\n",
    "\n",
    "lst_nb_neur = list(range(1, 11))\n",
    "in_dim = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "82f75a8e-485a-4a18-89b2-cd1f87668463",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_lst_nl(idx_ks, n):\n",
    "    if idx_ks == 0:\n",
    "        return [1, n, 1]\n",
    "    elif idx_ks == 1:\n",
    "        return [n, 2, 1]\n",
    "    elif idx_ks == 2:\n",
    "        return [n, n, 1]\n",
    "    else:\n",
    "        raise IndexError(f'Index \"idx_ks\" should be in [0, 2], got {idx_ks}.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b32b99d4-54ea-46a0-becd-fb62f7726acc",
   "metadata": {},
   "outputs": [],
   "source": [
    "lst_ks = []\n",
    "\n",
    "for idx_ks in range(3):\n",
    "    dct_ks = {}\n",
    "    for act_type, theta in itertools.product(lst_act_types, lst_theta):\n",
    "        dct_ks[(act_type, theta)] = []\n",
    "        for n in lst_nb_neur:\n",
    "            lst_nl = build_lst_nl(idx_ks, n)\n",
    "    \n",
    "            x_in = torch.ones(1, in_dim).expand(n_samples, in_dim)\n",
    "            z = full_propagation(x_in, lst_nl, dct_samplers[theta], dct_activations[act_type, theta])\n",
    "            dct_ks[act_type, theta].append(stats.kstest(z, stats.norm.cdf)[0])\n",
    "\n",
    "    for act_function in [\"relu\", \"tanh\"]:\n",
    "        dct_ks[act_function] = []\n",
    "        for n in lst_nb_neur:\n",
    "            lst_nl = build_lst_nl(idx_ks, n)\n",
    "    \n",
    "            x_in = torch.ones(1, in_dim).expand(n_samples, in_dim)\n",
    "            z = full_propagation(x_in, lst_nl, dct_samplers[act_function], dct_activations[act_function],\n",
    "                                sigma_w = dct_sigma_w[act_function], sigma_b = dct_sigma_b[act_function])\n",
    "            dct_ks[act_function].append(stats.kstest(z, stats.norm.cdf)[0])\n",
    "    \n",
    "    lst_ks.append(dct_ks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "9e745bb3-c928-40c2-98fc-fd22201d1103",
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx_ks in range(3):\n",
    "    dct_ks = lst_ks[idx_ks]\n",
    "    for act_function in [\"relu\", \"tanh\"]:\n",
    "        dct_ks[act_function] = []\n",
    "        for n in lst_nb_neur:\n",
    "            lst_nl = build_lst_nl(idx_ks, n)\n",
    "    \n",
    "            x_in = torch.ones(1, in_dim).expand(n_samples, in_dim)\n",
    "            z = full_propagation(x_in, lst_nl, dct_samplers[act_function], dct_activations[act_function],\n",
    "                                sigma_w = dct_sigma_w[act_function], sigma_b = dct_sigma_b[act_function])\n",
    "            dct_ks[act_function].append(stats.kstest(z, stats.norm.cdf)[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "bd6fd0df-aa43-40bd-84f7-f775c45314a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(lst_ks, 'Data_Init_KS_2layerNN_stat.pkl')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10593395-d387-406c-9ec1-dfeef623e6b2",
   "metadata": {},
   "source": [
    "### Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "275c0a99-2ba2-4b0b-8dad-fe844be90387",
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma_w = 1.\n",
    "n_samples = 200000\n",
    "\n",
    "lst_nb_neur = list(range(1, 11))\n",
    "in_dim = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "bc1656a6-7aed-4108-930c-3c8ae6ef56eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "lst_ks = torch.load('Data_Init_KS_2layerNN_stat.pkl', weights_only = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "f00a796e-5fe6-4fd0-af08-46529a6a5c2a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 331952 characters]",
      "text/plain": [
       "<Figure size 1700x940 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Prepare the plots\n",
    "matplotlib.rcParams['figure.figsize'] = [17, 9.4]\n",
    "matplotlib.rc('font', size = 15, weight = 'normal')\n",
    "titlesize = 25\n",
    "\n",
    "y_min, y_max = 1e-3, 4e-1\n",
    "\n",
    "figure, axis = plt.subplots(2, 3)\n",
    "\n",
    "# Set up some variables\n",
    "def build_title_ks(idx_ks, act_type):\n",
    "    if act_type == \"pos\":\n",
    "        name_act = \"$\\\\phi_{\\\\theta}^{\\\\mathrm{p}}$\"\n",
    "    else:\n",
    "        name_act = \"$\\\\phi_{\\\\theta}^{\\\\mathrm{o}}$\"\n",
    "        \n",
    "    if idx_ks == 0:\n",
    "        return f'$n_0 = 1$, varying $n_1$, {name_act}'\n",
    "    elif idx_ks == 1:\n",
    "        return f'$n_1 = 2$, varying $n_0$, {name_act}'\n",
    "    elif idx_ks == 2:\n",
    "        return f'varying $n_0$, $n_1$, {name_act}'\n",
    "    else:\n",
    "        raise IndexError(f'Index \"idx_ks\" should be in [0, 2], got {idx_ks}.')\n",
    "\n",
    "def build_label_x(idx_ks):\n",
    "    if idx_ks == 0:\n",
    "        return '$n_1$'\n",
    "    elif idx_ks == 1:\n",
    "        return '$n_0$'\n",
    "    elif idx_ks == 2:\n",
    "        return '$n_0 = n_1$'\n",
    "    else:\n",
    "        raise IndexError(f'Index \"idx_ks\" should be in [0, 2], got {idx_ks}.')\n",
    "\n",
    "for idx_ks, dct_ks in enumerate(lst_ks):\n",
    "    for idx, act_type in enumerate(['odd', 'pos']):\n",
    "        for theta in lst_theta:\n",
    "            axis[idx, idx_ks].plot(list(range(1, 11)), dct_ks[act_type, theta], label = r'$\\theta = {:.2f}$'.format(theta),\n",
    "                    c = build_color({'act_function': 'weibull', 'act_theta': theta}))\n",
    "            \n",
    "        axis[idx, idx_ks].set_title(build_title_ks(idx_ks, act_type), fontsize = titlesize)\n",
    "        \n",
    "        axis[idx, idx_ks].set_yscale('log')\n",
    "        axis[idx, idx_ks].set_ylim(y_min, y_max)\n",
    "        \n",
    "        axis[idx, idx_ks].set_xlabel(build_label_x(idx_ks))\n",
    "        axis[idx, idx_ks].set_xticks(lst_nb_neur)\n",
    "\n",
    "        for act_function in [\"relu\", \"tanh\"]:\n",
    "            axis[idx, idx_ks].plot(list(range(1, 11)), dct_ks[act_function], \n",
    "                                   label = build_label({\"act_function\": act_function, \"act_theta\": 0.}),\n",
    "                                   c = build_color({'act_function': act_function, 'act_theta': 0.}))\n",
    "\n",
    "for i in range(2):\n",
    "    for j in range(3):\n",
    "        axis[i, j].grid(axis = 'y')\n",
    "\n",
    "axis[0,0].legend(ncol = 2)\n",
    "\n",
    "# Global title\n",
    "plt.suptitle('KS-statistic for various layer sizes in a 2-layer NN', fontsize = titlesize)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig('Plot_Init_KS_2layerNN.pdf')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6da3d0f-b99e-4050-9949-e3ca22eca018",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
