{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cf52d36d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import pickle\n",
    "from pathlib import Path\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_density, find_activation, ParameterizedFunction, Integrand, ActivationFunction, ActivationFunctionTraining\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": "86c9123d",
   "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": "f380b930",
   "metadata": {},
   "source": [
    "# Compute an approximation $g_{\\Lambda}$ of the density $f_{|Y|}$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a94189fb",
   "metadata": {},
   "source": [
    "In this section, we compute an approximation $g_{\\Lambda}$ of the desired density function $f_{|Y|}$. The approximation we compute is a function with parameters, which should be optimized. \n",
    "\n",
    "The key step is the computation of the integral: $F_{\\Lambda}(z) = \\int_0^{\\infty} S_W(z/x) g_{\\Lambda}(x) \\mathrm{d}x$, in such a way that $F_{\\Lambda}(z)$ can be automatically differentiated by PyTorch according to the parameters of $g_{\\Lambda}$. See the documentation of `gaussian_preact.integration` for more information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e3c2d3b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the initial parameters of the density\n",
    "theta = None\n",
    "lambd = 1.\n",
    "alpha = 3.\n",
    "lambd_1 = 1.\n",
    "lambd_2 = 1.\n",
    "gamma = 1.\n",
    "\n",
    "# Build the dataset\n",
    "zmin = 0\n",
    "zmax = 5\n",
    "znum = 200\n",
    "inputs = torch.linspace(zmin, zmax, znum)\n",
    "targets = gaussian_preact.surv_Gaussian_abs(inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1cc70655",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optimization parameters\n",
    "lr = .01\n",
    "epochs = 600"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1cbdc016",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize storage dicts\n",
    "dct_densities = {}\n",
    "dct_losses = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "647624e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the densities for each theta\n",
    "for theta in lst_theta:\n",
    "    # Initialize objects depending on theta\n",
    "    theta_conj = theta_conj = 1/(1/2 - 1/theta)\n",
    "    surv_kernel_Weib = gaussian_preact.build_surv_kernel_Weib(theta, lambd)\n",
    "    density = ParameterizedFunction(theta, theta_conj, alpha, lambd_1, lambd_2, gamma)\n",
    "    integrand = Integrand(density, surv_kernel_Weib)\n",
    "    \n",
    "    # Train and get the final density\n",
    "    dct_densities[theta], dct_losses[theta] = find_density(theta, surv_kernel_Weib, density, integrand,\n",
    "            inputs, targets, lr = lr, epochs = epochs, scipy_update_period = 50, theta_conj_phases = 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "fc7db9a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Store the densities\n",
    "for theta in lst_theta:\n",
    "    density = dct_densities[theta]\n",
    "    loss = dct_losses[theta]\n",
    "    pickle.dump({'density': density, 'loss': loss}, \\\n",
    "        file = open(path_objects + 'DensityFunc_theta-{:.2f}.pkl'.format(theta), 'wb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "c84507e2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 136444 characters]",
      "text/plain": [
       "<Figure size 1000x670 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot the densities\n",
    "tens_x = torch.linspace(0, 3, 200)\n",
    "with torch.no_grad():\n",
    "    for theta, density in dct_densities.items():\n",
    "        tens_y = density(tens_x)\n",
    "        sym_tens_x = torch.concat((-tens_x.flip(0), tens_x))\n",
    "        sym_tens_y = torch.concat((tens_y.flip(0), tens_y))\n",
    "        \n",
    "        plt.plot(sym_tens_x, sym_tens_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(-3, 3)\n",
    "plt.ylim(-.1, 1.1)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9782f6f",
   "metadata": {},
   "source": [
    "# Compute an interpolation of the activation function $\\phi_{\\theta}$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7b64b32",
   "metadata": {},
   "source": [
    "Now that we have computed an estimation of $f_{|Y|}$, we are able to compute numerically $\\phi_{\\theta}$ at any point.\n",
    "\n",
    "In this section, we compute it at `xnum = 400` points between $0$ and `xmax = 10`. As a result, we can build an interpolation of $\\phi_{\\theta}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4ccfe4ee",
   "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": 10,
   "id": "b116167a",
   "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": 11,
   "id": "058915b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Interpolation parameters\n",
    "xmax = 10\n",
    "xnum = 400\n",
    "inputs = torch.linspace(0, xmax, xnum)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "db556d46",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize storage dicts\n",
    "dct_act_interpolation_x = {}\n",
    "dct_act_interpolation_y = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "e0896a2d",
   "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 = 2 * 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": 14,
   "id": "349ae057",
   "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}.pkl'.format(theta), 'wb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "0d2e2df9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 76828 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",
    "    sym_tens_x = torch.concat((-act_inter_x.flip(0), act_inter_x))\n",
    "    sym_tens_y = torch.concat((-act_inter_y.flip(0), act_inter_y))\n",
    "    plt.plot(sym_tens_x, sym_tens_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": 30,
   "id": "d40ce9a9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 80196 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",
    "    sym_tens_x = torch.concat((-act_inter_x.flip(0), act_inter_x))\n",
    "    sym_tens_y = torch.concat((-act_inter_y.flip(0), act_inter_y))\n",
    "    plt.plot(sym_tens_x, sym_tens_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(-3, 3)\n",
    "plt.ylim(-3, 3)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6aac1288",
   "metadata": {},
   "source": [
    "# Find an approximation of $\\phi_{\\theta}$ that is computationally efficient"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7430232b",
   "metadata": {},
   "source": [
    "The interpolation we have build is not fully satisfying if we want to use it in pytorch modules. For instance, computing it at a specific point $x$ needs to determine in each interval of the interpolation $x$ lies. This kind of operation slows down the computation, while making it more difficult to parallelize.\n",
    "\n",
    "So, we have made the choice to approximate it by a parameterized function $\\tilde{\\phi}_{\\theta}$. The parameters of this function are optimized in such a way that $\\tilde{\\phi}_{\\theta}$ is the closest possible to the interpolation computed in the preceding section.\n",
    "\n",
    "The family of functions has been chosen based on the shape of the interpolations, and the theoretic asymptotic growth rate, which is roughly $x^{1 - 2/\\theta}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "285ff18f",
   "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}.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": 18,
   "id": "ce3c7551",
   "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",
    "\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": 19,
   "id": "513f4c52",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optimization parameters\n",
    "lr = .01\n",
    "epochs = 10000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "25d03cf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize storage dicts\n",
    "dct_activations = {}\n",
    "dct_errors = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "90c229bb",
   "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 = ActivationFunctionTraining(theta, a, b, c, d, f, g)\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": 22,
   "id": "30b09980",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Store the activation functions\n",
    "for theta in lst_theta:\n",
    "    activation = ActivationFunction(dct_activations[theta])\n",
    "    error = dct_errors[theta]\n",
    "    torch.save({'activation': activation, 'error': error}, \\\n",
    "        path_objects + 'ActFunction_theta-{:.2f}.pkl'.format(theta))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "4ee5ab3b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 76924 characters]",
      "text/plain": [
       "<Figure size 1000x670 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "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": 32,
   "id": "352f5808",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 81880 characters]",
      "text/plain": [
       "<Figure size 1000x670 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "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(-3, 3)\n",
    "plt.ylim(-3, 3)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "3b0aca9c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_417791/3204192819.py:4: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  dct_load = torch.load(path_objects + 'ActFunction_theta-{:.2f}.pkl'.format(theta))\n"
     ]
    }
   ],
   "source": [
    "# Load the activation functions\n",
    "dct_activations = {}\n",
    "for theta in lst_theta:\n",
    "    dct_load = torch.load(path_objects + 'ActFunction_theta-{:.2f}.pkl'.format(theta))\n",
    "    dct_activations[theta] = dct_load['activation']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4f18a23-d475-45d5-aa78-54063169e3f8",
   "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
}
