{
 "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 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 find_activation, ActivationFunctionPos, ActivationFunctionPosTraining\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"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6562a621-c537-40fa-b7ae-64698187d5b1",
   "metadata": {},
   "source": [
    "# Compute a *positive* interpolation of the activation function $\\phi_{\\theta}$\n",
    "\n",
    "The densities should have been first computed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e8ed4866-e940-4bea-9ebf-240a89909996",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the densities\n",
    "dct_densities = {}\n",
    "for theta in lst_theta:\n",
    "    dct_load = pickle.load(open(path_objects + 'DensityFunc_theta-{:.2f}.pkl'.format(theta), 'rb'))\n",
    "    dct_densities[theta] = dct_load['density']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ab6faed6-dd48-46e0-aa7b-b2c0a6072589",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_density_no_torch(density):\n",
    "    def density_no_torch(x):\n",
    "        with torch.no_grad():\n",
    "            return density(x)\n",
    "        \n",
    "    return density_no_torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ce851a9f-107f-44c4-ae5b-06b383f37de1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Interpolation parameters\n",
    "xmin = -10\n",
    "xmax = 10\n",
    "xnum = 800\n",
    "inputs = torch.linspace(xmin, xmax, xnum)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "beba314c-1c0b-42bd-b9ac-c26a5514db57",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize storage dicts\n",
    "dct_act_interpolation_x = {}\n",
    "dct_act_interpolation_y = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "aeafb822-f4dd-4a27-a9b5-52779447de77",
   "metadata": {},
   "outputs": [],
   "source": [
    "for theta in lst_theta:\n",
    "    # Build the density function without autograd (improves the computation speed)\n",
    "    density = build_density_no_torch(dct_densities[theta])\n",
    "    \n",
    "    # Compute an interpolation of the related activation function\n",
    "    act_interpolation = torch.empty_like(inputs)\n",
    "    for i, x in enumerate(inputs):\n",
    "        y = stats.norm.sf(x)\n",
    "        act_interpolation[i] = optimize.brentq(lambda t: integrate.quad(density, t, np.inf)[0] - y, 0, 100)\n",
    "    \n",
    "    act_interpolation[0] = 0\n",
    "    \n",
    "    # Store the result\n",
    "    dct_act_interpolation_x[theta] = inputs\n",
    "    dct_act_interpolation_y[theta] = act_interpolation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8f270517-3464-47f7-9bfc-d15b69653ee8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Store the interpolations\n",
    "for theta in lst_theta:\n",
    "    act_inter_x = dct_act_interpolation_x[theta]\n",
    "    act_inter_y = dct_act_interpolation_y[theta]\n",
    "    pickle.dump({'act_interpolation_x': act_inter_x, 'act_interpolation_y': act_inter_y}, \\\n",
    "        file = open(path_objects + 'ActInterpolation_theta-{:.2f}_pos.pkl'.format(theta), 'wb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3deba90f-dd36-4ea9-a86b-c319913ff5a5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 50676 characters]",
      "text/plain": [
       "<Figure size 1000x670 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot the interpolations\n",
    "for theta in lst_theta:\n",
    "    act_inter_x = dct_act_interpolation_x[theta]\n",
    "    act_inter_y = dct_act_interpolation_y[theta]\n",
    "    plt.plot(act_inter_x, act_inter_y, 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.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "2229e26d-0677-40fa-8c6b-7075b2812621",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 63124 characters]",
      "text/plain": [
       "<Figure size 1000x670 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot the interpolations\n",
    "for theta in lst_theta:\n",
    "    act_inter_x = dct_act_interpolation_x[theta]\n",
    "    act_inter_y = dct_act_interpolation_y[theta]\n",
    "    plt.plot(act_inter_x, act_inter_y, 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(-2.5, 2.5)\n",
    "plt.ylim(-2.5, 2.5)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "295acd38-2a6c-4c3f-9248-28fe4bee6102",
   "metadata": {},
   "source": [
    "# Find an approximation of $\\phi_{\\theta}$ that is computationally efficient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c56452b3-a348-4458-b70c-f8c8d2da25c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the interpolations\n",
    "dct_act_interpolation_x = {}\n",
    "dct_act_interpolation_y = {}\n",
    "for theta in lst_theta:\n",
    "    dct_load = pickle.load(open(path_objects + 'ActInterpolation_theta-{:.2f}_pos.pkl'.format(theta), 'rb'))\n",
    "    dct_act_interpolation_x[theta] = dct_load['act_interpolation_x']\n",
    "    dct_act_interpolation_y[theta] = dct_load['act_interpolation_y']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e7131180-6174-4479-a6b0-a0391e02982b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the initial parameters of the activation function\n",
    "a = 1\n",
    "b = 1\n",
    "c = 1\n",
    "d = .01\n",
    "\n",
    "f = 1\n",
    "g = .01\n",
    "h = 0\n",
    "\n",
    "# Build the inputs\n",
    "\"\"\"\n",
    "In the context of Gaussian pre-activations, we must ensure that our proposition of activation function\n",
    "fits well the interpolation where the Normal(0, 1) is concentrated. This is why we build 'inputs_gauss'.\n",
    "\n",
    "However, the set 'inputs_gauss' is highly concentrated around 0, and it would not impose any constraint,\n",
    "for instance, on the interval [4, 10]. So, we artificially add some points to 'inputs_gauss' to build\n",
    "the new set 'inputs', which covers the interval [-xmax, xmax] (where xmax can be set to 10).\n",
    "\"\"\"\n",
    "xnum_gauss = 1000\n",
    "inputs_gauss = torch.tensor(np.sqrt(2) * special.erfinv(np.linspace(-1 + 1/xnum_gauss, 1 - 1/xnum_gauss, \n",
    "                                                                    2*xnum_gauss - 1)))\n",
    "xmax = 10\n",
    "xnum_compl = 50\n",
    "inputs = torch.concat((torch.linspace(-xmax, inputs_gauss[0], xnum_compl),\n",
    "                       inputs_gauss,\n",
    "                       torch.linspace(inputs_gauss[-1], xmax, xnum_compl)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "63abbadd-acc6-4286-878d-408225ee8855",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optimization parameters\n",
    "lr = .01\n",
    "epochs = 10000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "d05d7847-c476-48c2-9ad4-38df4857f85e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize storage dicts\n",
    "dct_activations = {}\n",
    "dct_errors = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "944950d2-18e0-4c91-9ca9-7925962cd293",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the activation function for each theta\n",
    "for theta in lst_theta:\n",
    "    # Get the interpolation\n",
    "    act_inter_x = dct_act_interpolation_x[theta]\n",
    "    act_inter_y = dct_act_interpolation_y[theta]\n",
    "    \n",
    "    # Initialize the activation function to train\n",
    "    activation = ActivationFunctionPosTraining(theta, a, b, c, d, f, g, h)\n",
    "    \n",
    "    # Perform the training\n",
    "    activation, error = find_activation(theta, activation, act_inter_x, act_inter_y, inputs,\n",
    "            lr = lr, epochs = epochs)\n",
    "    \n",
    "    dct_activations[theta] = activation\n",
    "    dct_errors[theta] = error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "42a71a50-af43-42d9-8201-ea270498562f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Store the activation functions\n",
    "for theta in lst_theta:\n",
    "    activation = dct_activations[theta]\n",
    "    error = dct_errors[theta]\n",
    "    torch.save({'activation': ActivationFunctionPos(activation), 'error': error}, \\\n",
    "               path_objects + 'ActFunctionPos_theta-{:.2f}.pkl'.format(theta))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "bb6582d5-a21b-4f7a-b443-dfe37f02654b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 50580 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[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.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d496306-4253-4306-8fb0-319b7f259b6b",
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
