notebooks/chapter06_book.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hessian and variance estimation for neural network"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def towdir(s):\n",
    "    return (str('./datasets_book/'+s))\n",
    "\n",
    "import deepglmlib.utils as utils\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import importlib\n",
    "importlib.reload(utils)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "beta = np.array([-0.5,1.0,0.8,0.7,0.4,0.2]).reshape((6,1))\n",
    "beta_true = beta\n",
    "\n",
    "n = 30\n",
    "x = np.random.uniform(0,1,n*5).reshape((n,5))\n",
    "X = np.hstack([np.ones((n,1)), x])\n",
    "e = np.random.randn(n).reshape((n,1))/5\n",
    "y = X @ beta + e\n",
    "\n",
    "x_train, x_test, y_train, y_test = utils.f_splitData(x,y,percentage=0.25)\n",
    "\n",
    "x_train, x_test = utils.f_normalizeData(x_train,x_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_train, p_train = x_train.shape\n",
    "n_test, p_test   = x_test.shape\n",
    "X_train          = np.hstack([np.ones((n_train,1)), x_train])\n",
    "X_test           = np.hstack([np.ones((n_test,1)), x_test])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_train, p_train, n_test, p_test "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Variance estimation of regression parameters without  pytorch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Standard-deviation estimation from **statsmodels**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import statsmodels.api as stm\n",
    "# n_train       = x_train.shape[0]\n",
    "ols           = stm.OLS(y_train, X_train)\n",
    "fit_ols_train = ols.fit()\n",
    "olssumy       = fit_ols_train.summary()\n",
    "\n",
    "print(olssumy.tables[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Recall for some algebra on the variance of the coefficients\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Standard-deviation estimation from **numpy**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f_varthetahat(X,y,printed=False):\n",
    "    n            = X.shape[0]\n",
    "    p            = X.shape[1] #intercept not counting\n",
    "    beta_hat     = np.linalg.inv(X.T @ X) @ X.T @ y\n",
    "    y_hat        = X @ beta_hat\n",
    "    residual     = y - y_hat\n",
    "    sigma2_hat   = np.sum(residual**2) / (n - p)\n",
    "    var_beta_hat = sigma2_hat * np.linalg.inv(X.T @ X)\n",
    "    if printed:\n",
    "        for p_ in range(p):\n",
    "            standard_error = var_beta_hat[p_, p_] ** 0.5\n",
    "            print(f\"SE(beta_hat[{p_}]): {standard_error}\")\n",
    "    return beta_hat.ravel(), var_beta_hat, sigma2_hat\n",
    "\n",
    "beta_hat_train, var_beta_hat_train, sigma2_hat_train = \\\n",
    "   f_varthetahat(X_train,y_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The covariance matrix estimate is equal to:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Cov_betahat = var_beta_hat_train\n",
    "\n",
    "print(\"Cov_betahat=\\n\", np.round(Cov_betahat,4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Variance parameters with pytorch: hessian computation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "\n",
    "dataset_train = TensorDataset( torch.Tensor(x_train), torch.Tensor(y_train) )\n",
    "dataset_test = TensorDataset( torch.Tensor(x_test), torch.Tensor(y_test) )\n",
    "\n",
    "dl_train = DataLoader(dataset_train,shuffle=True,batch_size=10)\n",
    "dl_test = DataLoader(dataset_test,shuffle=True,batch_size=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"cuda.is_available()     = \", torch.cuda.is_available())\n",
    "print(\"cuda.get_device_name(0) = \",torch.cuda.get_device_name(0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" \\\n",
    "  if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import copy\n",
    "\n",
    "px = x_train.shape[1]\n",
    "nbmax_epoqs = 150\n",
    "alpha_t     = 1e-3\n",
    "debug_out   = 10\n",
    "\n",
    "layers_regress = [ nn.Linear(px,1,bias=True) ]\n",
    "\n",
    "model     =  utils.GNLMRegression(\"LinearRegression\",\n",
    "                                  copy.deepcopy(layers_regress))\n",
    "loss      = torch.nn.MSELoss(reduction='sum')\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=alpha_t, momentum=0.0)\n",
    "monitor   = utils.MyMonitorTest(model,loss,dl_train,dl_test,nbmax_epoqs,\n",
    "                                debug_out,device=device)\n",
    "\n",
    "loss_train_s,tmax,monistopc  = \\\n",
    "  utils.f_train_glmr(dl_train,model,optimizer,loss,\n",
    "                     monitor,device=device,printed=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_train_s = loss_train_s\n",
    "t_train = range(len(loss_train_s))\n",
    "\n",
    "loss_test_s = monitor.loss_test_s[monitor.loss_test_s>0]\n",
    "t_test = monitor.step_test_s[monitor.loss_test_s>0].astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "fig, ax = plt.subplots()\n",
    "utils.f_draw_s([ t_train, t_test ],\n",
    "               [ loss_train_s/n_train,loss_test_s/n_test],\n",
    "               [\"b-\", \"r-\"] ,\"t\",[ \"loss train\", \"loss test\"], \" \", ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thus, the vector of regression coefficient is written,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#  utils.extract_weights_lin(model)\n",
    "def extract_weights_lin(model,keybias=\"lin.bias\",keyweight=\"lin.weight\"):\n",
    "    weight_ = bias_ = None\n",
    "    for param_tensor in model.state_dict():\n",
    "        if (param_tensor==keyweight):\n",
    "            weight_= model.state_dict()[param_tensor]\n",
    "        if (param_tensor==keybias):\n",
    "            bias_= model.state_dict()[param_tensor]    \n",
    "    return np.append(bias_,weight_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta_hat_train2, var_beta_hat_train2, sigma2_hat_train2 = \\\n",
    "   f_varthetahat(X_train,y_train)\n",
    "\n",
    "std_beta_hat_train2 = np.sqrt(np.diag(var_beta_hat_train2))\n",
    "\n",
    "np.round(std_beta_hat_train2,3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example with real data for classification"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Recall about the hessian and variance estimation for logistic regression"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hessian and parameter variance for abalone dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "abalone = pd.read_csv(towdir(\"./abalone_prep.csv\"))\n",
    "\n",
    "x       = abalone.drop(columns=\"rings\")\n",
    "y       = abalone[\"rings\"].values - 1\n",
    "y       = (y>np.median(y)).astype(int)\n",
    "\n",
    "x.shape, y.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let sample the train and set samples for the training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_train, x_test, y_train, y_test = utils.f_splitData(x.values,y,percentage=0.333)\n",
    "x_train, x_test                  = utils.f_normalizeData(x_train,x_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Results from the module statsmodels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import statsmodels.api as stm\n",
    "n_train       = x_train.shape[0]\n",
    "X_train      = np.hstack([np.ones((n_train, 1)), x_train])\n",
    "lgt           = stm.Logit(y_train, X_train)\n",
    "fit_lgt_train = lgt.fit(maxiter=300)\n",
    "lgtsumy       = fit_lgt_train.summary()\n",
    "\n",
    "print(lgtsumy.tables[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Results from the module sklearn and variance with numpy "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from sklearn import linear_model\n",
    "logit     = linear_model.LogisticRegression(penalty='none',fit_intercept=False,max_iter=300)\n",
    "resLogit  = logit.fit(X_train, y_train)\n",
    "predProbs = resLogit.predict_proba(X_train)\n",
    "Omega     = np.diagflat(np.product(predProbs, axis=1))\n",
    "cov_theta = np.linalg.inv(np.dot(np.dot(X_train.T, Omega), X_train))\n",
    "\n",
    "betahat_skl = resLogit.coef_.ravel()\n",
    "print(\"betahat_skl: \", np.round(betahat_skl,4))\n",
    "print(\"stdhat_skl: \", np.round(np.sqrt(np.diag(cov_theta)),3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Training with pytorch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "\n",
    "dataset_train     = TensorDataset( torch.Tensor(x_train), \n",
    "                                       torch.Tensor(y_train) )\n",
    "\n",
    "dataset_test     = TensorDataset( torch.Tensor(x_test), \n",
    "                                       torch.Tensor(y_test) )\n",
    "\n",
    "dl_train = DataLoader(dataset_train,shuffle=True,batch_size=100)\n",
    "dl_test = DataLoader(dataset_test,shuffle=True,batch_size=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" \\\n",
    "  if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import copy\n",
    "\n",
    "px = x_train.shape[1]\n",
    "nbmax_epoqs = 800\n",
    "alpha_t     = 1e-2\n",
    "debug_out   = 10\n",
    "\n",
    "layers_regress = [ nn.Linear(px,1,bias=True) ]\n",
    "\n",
    "model     =  utils.GNLMRegression(\"LogisticRegression\",\n",
    "                                  copy.deepcopy(layers_regress))\n",
    "\n",
    "loss      = torch.nn.BCEWithLogitsLoss(reduction='sum')\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=alpha_t, momentum=0.0)\n",
    "monitor   = utils.MyMonitorTest(model,loss,dl_train,dl_test,nbmax_epoqs,debug_out,device=device)\n",
    "\n",
    "loss_train_s,tmax,monistopc  = utils.f_train_glmr(dl_train,model,optimizer,loss,monitor,device=device,\n",
    "                                            transform_yb=utils.transform_yb,\n",
    "                                            transform_yhatb=utils.transform_yhatb,\n",
    "                                            printed=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_train, n_test = x_train.shape[0], x_test.shape[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_weights_lin(model,keybias=\"lin.bias\",keyweight=\"lin.weight\"):\n",
    "    weight_ = bias_ = None\n",
    "    for param_tensor in model.state_dict():\n",
    "        if (param_tensor==keyweight):\n",
    "            weight_= model.state_dict()[param_tensor]\n",
    "        if (param_tensor==keybias):\n",
    "            bias_= model.state_dict()[param_tensor]\n",
    "    return np.append(bias_,weight_)\n",
    "\n",
    "thetahat_torch = extract_weights_lin(model.cpu(),keybias=\"net.0.bias\",keyweight=\"net.0.weight\")\n",
    "print( np.round(thetahat_torch,4) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logLik_torch=0\n",
    "with torch.no_grad():\n",
    "    for b, (Xb,yb) in enumerate(dl_train):\n",
    "        yhatb = model(Xb)\n",
    "        logLik_torch -= loss(yhatb.reshape(yb.shape), yb)\n",
    "\n",
    "float(logLik_torch.detach().numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(np.round(resLogit.coef_.ravel(),4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "phat = np.exp(X_train @ resLogit.coef_.ravel())\n",
    "phat = phat/(1+phat)\n",
    "np.sum(y_train * np.log(phat) + (1-y_train)*np.log(1-phat))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Hessian computation with pytorch from second-order derivative"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch, torchvision\n",
    "from torch.autograd import Variable, grad\n",
    "import torch.distributions as td\n",
    "import math\n",
    "from torch.optim import Adam\n",
    "import scipy.stats\n",
    "\n",
    "theta_train_ = torch.Tensor(thetahat_torch) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p_model = 0\n",
    "for p in model.parameters():\n",
    "    if len(p.shape)>1:\n",
    "        p_model += p.shape[1]\n",
    "    else:\n",
    "        p_model += 1 #scalar (bias=intercept=wk0)\n",
    "p_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "I         = torch.zeros((p_model,p_model))\n",
    "thessian  = torch.autograd.functional.hessian\n",
    "\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=0.0, momentum=0.0)\n",
    "theta_hat = theta_train_\n",
    "\n",
    "t=0\n",
    "for b,(Xb,yb) in enumerate(dl_train):\n",
    "    print(\".\", end = '')\n",
    "    Xb = torch.Tensor(np.hstack([np.ones((len(Xb), 1)), Xb]))\n",
    "    #\n",
    "    def log_lik_b(theta):\n",
    "        p_b = torch.exp(Xb@theta)\n",
    "        p_b = p_b/(1+p_b)\n",
    "        return torch.log(p_b.T)@ yb +  torch.log(1-p_b.T) @ (1-yb)\n",
    "    optimizer.zero_grad()\n",
    "    I_b = -thessian(log_lik_b, theta_hat) / n_train\n",
    "    I = I + I_b#.squeeze()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The hessian could be computed for the loss from the model, directly, this more advanced implementation is not given here and can be found as an extern module for pytorch, from repositories websites. This is mostly wanted for deep neural neworks where hidden layers are trained too and with their weighted added to the list of parameters.\n",
    "\n",
    "The numpy version of the pytorch Tensor object which contains the hessian matrix is extracted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "I = I.detach().numpy()\n",
    "\n",
    "print()\n",
    "print(np.round(I,3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.asarray(lgtsumy.tables[1].data)[:,2][1:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.round(np.sqrt(np.diag(np.linalg.inv(I))/n_train),3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is slightly equal to the standard-deviations from the estimation before with the dedicated python module. The difference comes from that the final solution for the regression coefficients were only nearby as the training is for a nonlinear function with different initial values and different inferential procedures. The first solution is a second-order procedure while the second one with pytorch is a first-order one which as for more carefull settings. This is the price to pay in order to get a more scalable algorithm without requiring the hessian at each step of the training, in order to avoid a non necessary costly numerical burden."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Hessian computation with pytorch from second-order derivative (bis)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let revisit this python code with a more general setting with the model for computation of the predicted target variable. This is required for more advanced model with hidden layers for instance. First, let remember that the weights are structured as follows in pytorch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for p in model.parameters(): print(p.data, end= \" \")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The weights are also available from the dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mn = model.net[0]\n",
    "print(mn.state_dict()['bias'], mn.state_dict()['weight'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There is a more direct access to the weights which allows to update their values too."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(model.net[0].bias.data, model.net[0].weight.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(theta_hat[0], theta_hat[1:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f_get_p_model(model):\n",
    "    p_model = 0\n",
    "    for p in model.parameters():\n",
    "        if len(p.shape)>1:        # matrix          : array 2 dims\n",
    "            p_model += p.shape[0] * p.shape[1]\n",
    "        else:\n",
    "            p_model += p.shape[0] # vector or scalar:  array 1 dim\n",
    "    return p_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f_varianceMatrixFromFullHessian_ForParameters_modelNN(model, loss, \n",
    "                                                          dataloader,n_train,\n",
    "                                                          loss_yy_model = None,\n",
    "                                                          device=None):\n",
    "    if device is not None: model = model.to(device)\n",
    "    model.eval()\n",
    "    \n",
    "    p_model = f_get_p_model(model)\n",
    "    Imodel  = torch.zeros((p_model,p_model)) \n",
    "    if device is not None: Imodel = Imodel.to(device)\n",
    "    \n",
    "    optimizer = torch.optim.SGD(model.parameters(), lr=0.0, momentum=0.0)\n",
    "    optimizer.zero_grad()\n",
    "    \n",
    "    for b,(Xb,yb) in enumerate(dataloader):\n",
    "        print(\".\", end = '')\n",
    "        \n",
    "        if device is not None: Xb=Xb.to(device)\n",
    "            \n",
    "        yhatb = model(Xb)\n",
    "        if device is not None: yhatb = yhatb.to(device) # transformation ?????? <- linear & logit ! \n",
    "        \n",
    "        #loss_b      = loss(yhatb, yb.reshape(yhatb.shape))\n",
    "        # if loss_yy_model is None:\n",
    "        \n",
    "        if device is not None: yb = yb.to(device)\n",
    "        #yb = transform_yb(yb, model.name,yhatb, device)\n",
    "            \n",
    "        loss_b = loss(yhatb, yb.reshape(yhatb.shape))\n",
    "        # else:\n",
    "        # loss_b = loss_yy_model(loss(yhatb, yb.reshape(yhatb.shape)),model)\n",
    "               \n",
    "        grad1rds_list = torch.autograd.grad(loss_b, model.parameters(), \\\n",
    "                                         create_graph=True, \\\n",
    "                                         retain_graph=True, \\\n",
    "                                         allow_unused=True)\n",
    "\n",
    "        grad1rds_vec = torch.cat([g.view(-1) for g in grad1rds_list]) #.squeeze()\n",
    "        \n",
    "        grad2rds_list = []\n",
    "        for grad1rd in  grad1rds_vec:\n",
    "            grad2rds_1row = torch.autograd.grad(grad1rd, model.parameters(), \\\n",
    "                                          create_graph=True, \\\n",
    "                                          retain_graph=True, \\\n",
    "                                          allow_unused=True)\n",
    "\n",
    "            grad2rds_1vect = torch.cat([g.view(-1) for g in grad2rds_1row]) #.squeeze()\n",
    "\n",
    "            grad2rds_list.append( grad2rds_1vect )\n",
    "\n",
    "        for k in range(p_model):\n",
    "            Imodel[k,:] += grad2rds_list[k] / n_train\n",
    "        \n",
    "    return Imodel, p_model\n",
    "\n",
    "###\n",
    "Imodel_unroll, p_model = \\\n",
    "  f_varianceMatrixFromFullHessian_ForParameters_modelNN(model, loss, dl_train, n_train,\n",
    "                                                        device=device)\n",
    "\n",
    "Imodel_unroll = Imodel_unroll.detach().cpu().numpy()\n",
    "Vtheta_hat = np.linalg.inv(Imodel_unroll) / n_train\n",
    "\n",
    "print()\n",
    "print(np.round(np.sqrt(np.diag(Vtheta_hat)),3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Hessian computation with pytorch from first-order derivative"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f_varianceMatrixFromGradients_ForParameters_modelNN(model, loss, dataloader, \n",
    "                                                        n_train, device=None):\n",
    "    if device is not None: model = model.to(device)   \n",
    "    p_model = f_get_p_model(model)\n",
    "    Iapprox = torch.zeros((p_model,p_model))\n",
    "    if device is not None: Iapprox = Iapprox.to(device)\n",
    "    optimizer = torch.optim.SGD(model.parameters(), lr=0.0, momentum=0.0)\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "    for b,(Xb,yb) in enumerate(dataloader):\n",
    "        print(\".\", end = '')\n",
    "        for i in range(Xb.shape[0]):\n",
    "            Xb_i = Xb[i,:]\n",
    "            yb_i = yb[i].ravel()\n",
    "            if device is not None: Xb_i=Xb_i.to(device)\n",
    "            if device is not None: yb_i = yb_i.to(device)\n",
    "            yhatb_i = model(Xb_i).ravel()\n",
    "            loss_b = loss(yhatb_i, yb_i)\n",
    "            optimizer.zero_grad()\n",
    "            loss_b.backward()\n",
    "            gradient_vect = []\n",
    "            with torch.no_grad():\n",
    "                for p in model.parameters():\n",
    "                    gradient_vect.append(p.grad.view(-1))\n",
    "                gradient_vect = torch.cat(gradient_vect)\n",
    "                gradient_vect = gradient_vect.reshape((p_model,1))\n",
    "                Iapprox = Iapprox + gradient_vect @ gradient_vect.T /n_train\n",
    "    return Iapprox, p_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Iapprox, p_model = \\\n",
    "   f_varianceMatrixFromGradients_ForParameters_modelNN(model, loss, dl_train, \n",
    "                                                       n_train, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Iapprox          = Iapprox.detach().cpu().numpy() "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stdI1 = np.sqrt(np.diag(np.linalg.inv(Imodel_unroll))/n_train)\n",
    "stdI2 = np.sqrt(np.diag(np.linalg.inv(Iapprox))/n_train)\n",
    "\n",
    "\n",
    "stdI1 = np.roll(stdI1,1)\n",
    "stdI2 = np.roll(stdI2,1)\n",
    "\n",
    "print(np.round(stdI1,3))\n",
    "print(np.round(stdI2,3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.asarray(lgtsumy.tables[1].data)[:,2][1:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "myenv",
   "language": "python",
   "name": "myenv"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}