radon.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Radon Multi-level Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we use the radon data and the multi-level model for modeling it. \n",
    "Next, we train a decision-maker (neural network) to show that it is capable of correcting errors due to posterior approximation.\n",
    "The model is implemented using the publicly available [Stan code](http://mc-stan.org/users/documentation/case-studies/radon.html), using the Minnesota subset of the data and also otherwise conforming to their details. \n",
    "The inference is carried out using automatic differentiation variational inference, and we split the data randomly into equally sized training and test set."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Proposed Experiment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since variational approximations are trained with iterative algorithms, we can construct a controlled experiment by terminating the inference algorithm early and varying the termination point. In general, approximations trained for shorter time are further away from the final approximation and the true posterior. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pystan\n",
    "import pickle\n",
    "import numpy as np\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using CPU\n"
     ]
    }
   ],
   "source": [
    "if torch.cuda.is_available():\n",
    "    torch.set_default_tensor_type('torch.cuda.FloatTensor')\n",
    "    env = torch.cuda\n",
    "    device = torch.device('cuda')\n",
    "    print(\"Using GPU\")\n",
    "else:\n",
    "    torch.set_default_tensor_type('torch.FloatTensor')\n",
    "    env = torch\n",
    "    device = torch.device('cpu')\n",
    "    print(\"Using CPU\")\n",
    "    \n",
    "    \n",
    "def totorch(array):\n",
    "    return torch.tensor(array, dtype=env.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from aux import tonumpy, parse_script_args, dict2str, print2\n",
    "import losses\n",
    "import regression_data\n",
    "import radon"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "parsing: <-f>\n"
     ]
    }
   ],
   "source": [
    "args = parse_script_args() # arguments can be passed in the format of NAME1=FLOATVAL1,NAME2=[STRVAL2],NAME3=INTVAL3,..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# optimization general parmeters\n",
    "SEED = args.get(\"SEED\", 5)\n",
    "VI_NITER = int(args.get(\"VI_NITER\", 1e6))\n",
    "HMC_NITER = int(args.get(\"HMC_NITER\", 1e6))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# regularization\n",
    "REGULARIZATION = args.get(\"REGULARIZATION\", \"boot1\").lower() # const/boot1\n",
    "LAMBDA = args.get(\"LAMBDA\", 1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# selected loss: tilted/squared/exptilted/expsquared\n",
    "LOSS = args.get(\"LOSS\", \"tilted\")\n",
    "TILTED_Q = args.get(\"TILTED_Q\", 0.5) # relevant only for tilted and exptilted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CONFIGURATION SUMMARY: HMC_NITER=0 LAMBDA=1.0 LOSS=tilted REGULARIZATION=boot1 SEED=5 TILTED_Q=0.5 VI_NITER=10000\n"
     ]
    }
   ],
   "source": [
    "print(\"CONFIGURATION SUMMARY: %s\" % dict2str(globals()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x7fb5ba09b5d0>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Losses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[LossFactory] Configuration: TILTED_Q=0.5 LINEX_C=None\n",
      "> <tilted> loss: tilted_loss_fixedq with (analytical/Bayes estimator) h: tilted_optimal_h_fixedq\n"
     ]
    }
   ],
   "source": [
    "loss, optimal_h_bayes_estimator = losses.LossFactory(**globals()).create(LOSS)\n",
    "print2(\"> <%s> loss: %s with (analytical/Bayes estimator) h: %s\" % \n",
    "        (LOSS, loss.__name__, optimal_h_bayes_estimator.__name__))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "empirical_risk = lambda preds, y: loss(preds, totorch(y)).mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_radon, n_county, county, u, floor_measure = regression_data.load_radon()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAIN_MASK = np.array([i%2==0 for i in range(len(log_radon))], dtype=bool)\n",
    "TEST_MASK = ~TRAIN_MASK\n",
    "\n",
    "COUNTY_TRAIN = county[TRAIN_MASK]\n",
    "U_TRAIN = u[TRAIN_MASK]\n",
    "FLOOR_MEASURE_TRAIN = floor_measure[TRAIN_MASK]\n",
    "LOG_RADON_TRAIN = log_radon[TRAIN_MASK]\n",
    "\n",
    "COUNTY_TEST = county[TEST_MASK]\n",
    "U_TEST = u[TEST_MASK]\n",
    "FLOOR_MEASURE_TEST = floor_measure[TEST_MASK]\n",
    "LOG_RADON_TEST = log_radon[TEST_MASK]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(460,) (459,)\n"
     ]
    }
   ],
   "source": [
    "Y_TRAIN, Y_TEST = LOG_RADON_TRAIN, LOG_RADON_TEST\n",
    "print(Y_TRAIN.shape, Y_TEST.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:pystan:COMPILING THE C++ CODE FOR MODEL anon_model_09c24493b150e61b5f7babd2ed515f09 NOW.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading model from radon_model.pkl\n",
      "Failed. Recomputing and saving model to radon_model.pkl\n"
     ]
    }
   ],
   "source": [
    "path = \"radon_model.pkl\"\n",
    "try:\n",
    "    print(\"Loading model from %s\" % path)\n",
    "    sm = pickle.load(open(path, \"rb\"))\n",
    "except:\n",
    "    print(\"Failed. Recomputing and saving model to %s\" % path)\n",
    "    sm = pystan.StanModel(model_code=radon.hierarchical_intercept)\n",
    "    pickle.dump(sm, open(path, \"wb\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "hierarchical_intercept_data = radon.create_data(LOG_RADON_TRAIN, n_county, \n",
    "                                                COUNTY_TRAIN, U_TRAIN, FLOOR_MEASURE_TRAIN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### HMC evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "if HMC_NITER<=0:\n",
    "    hmc_train_risk = 0.\n",
    "    hmc_test_risk = 0.\n",
    "else:\n",
    "    hmc_fit = sm.sampling(data=hierarchical_intercept_data, iter=HMC_NITER, n_jobs=1)\n",
    "    ys_train_hmc = totorch(radon.sample_predictive_y0_hmc(hmc_fit, FLOOR_MEASURE_TRAIN, U_TRAIN, COUNTY_TRAIN))\n",
    "    ys_test_hmc = totorch(radon.sample_predictive_y0_hmc(hmc_fit, FLOOR_MEASURE_TEST, U_TEST, COUNTY_TEST))\n",
    "    hstar_train_hmc = optimal_h_bayes_estimator(ys_train_hmc)\n",
    "    hstar_test_hmc = optimal_h_bayes_estimator(ys_test_hmc)\n",
    "    hmc_train_risk = empirical_risk(hstar_train_hmc, Y_TRAIN).item()\n",
    "    hmc_test_risk = empirical_risk(hstar_test_hmc, Y_TEST).item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HMC optimal risk: train:0.0000 test:0.0000\n"
     ]
    }
   ],
   "source": [
    "print(\"HMC optimal risk: train:%.4f test:%.4f\" % (hmc_train_risk, hmc_test_risk))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### VI Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:pystan:Automatic Differentiation Variational Inference (ADVI) is an EXPERIMENTAL ALGORITHM.\n",
      "WARNING:pystan:ADVI samples may be found on the filesystem in the file `/tmp/tmpaq2uvqef/output.csv`\n"
     ]
    }
   ],
   "source": [
    "#vi_fit = sm.vb(data=hierarchical_intercept_data, iter=1000000, tol_rel_obj=0.0001)\n",
    "vi_fit = sm.vb(data=hierarchical_intercept_data, iter=VI_NITER, tol_rel_obj=0.0001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "ys_train = totorch(radon.sample_predictive_y0_vi(vi_fit, FLOOR_MEASURE_TRAIN, U_TRAIN, COUNTY_TRAIN))\n",
    "ys_test = totorch(radon.sample_predictive_y0_vi(vi_fit, FLOOR_MEASURE_TEST, U_TEST, COUNTY_TEST))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "hstar_train = optimal_h_bayes_estimator(ys_train)\n",
    "hstar_test = optimal_h_bayes_estimator(ys_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VI optimal risk: train:0.2686 test:0.2884\n"
     ]
    }
   ],
   "source": [
    "vi_train_risk = empirical_risk(hstar_train, Y_TRAIN).item()\n",
    "vi_test_risk = empirical_risk(hstar_test, Y_TEST).item()\n",
    "print(\"VI optimal risk: train:%.4f test:%.4f\" % (vi_train_risk, vi_test_risk))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Regularization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[regularization] Using regularization: boot1 with lambda=1.0\n",
      "[regularization] Bootstrap-based regularization\n",
      "[get_bootstrap_decisions] 0/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:pystan:Automatic Differentiation Variational Inference (ADVI) is an EXPERIMENTAL ALGORITHM.\n",
      "WARNING:pystan:ADVI samples may be found on the filesystem in the file `/tmp/tmpj6_fj7q6/output.csv`\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[get_bootstrap_decisions] 1/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:pystan:Automatic Differentiation Variational Inference (ADVI) is an EXPERIMENTAL ALGORITHM.\n",
      "WARNING:pystan:ADVI samples may be found on the filesystem in the file `/tmp/tmpm11nz9gc/output.csv`\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[get_bootstrap_decisions] 2/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:pystan:Automatic Differentiation Variational Inference (ADVI) is an EXPERIMENTAL ALGORITHM.\n",
      "WARNING:pystan:ADVI samples may be found on the filesystem in the file `/tmp/tmplvn1y50n/output.csv`\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[get_bootstrap_decisions] 3/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:pystan:Automatic Differentiation Variational Inference (ADVI) is an EXPERIMENTAL ALGORITHM.\n",
      "WARNING:pystan:ADVI samples may be found on the filesystem in the file `/tmp/tmp4hywr9qj/output.csv`\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[get_bootstrap_decisions] 4/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:pystan:Automatic Differentiation Variational Inference (ADVI) is an EXPERIMENTAL ALGORITHM.\n",
      "WARNING:pystan:ADVI samples may be found on the filesystem in the file `/tmp/tmpfvfo1uy2/output.csv`\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[regularization] Bootstrap-based regularization: overwritting regularization mean\n"
     ]
    }
   ],
   "source": [
    "import regression_regularization as regularization\n",
    "\n",
    "HSTAR_TRAIN = hstar_train # default regularization mean\n",
    "\n",
    "print(\"[regularization] Using regularization: %s with lambda=%s\" % (REGULARIZATION, LAMBDA))\n",
    "if REGULARIZATION.startswith(\"boot\"):     \n",
    "    print(\"[regularization] Bootstrap-based regularization\")    \n",
    "    stan_data_generator = radon.yield_data_bootstrap(LOG_RADON_TRAIN, n_county, COUNTY_TRAIN, \n",
    "                                                     U_TRAIN, FLOOR_MEASURE_TRAIN)\n",
    "    sample_predictive_y0_train = lambda vi_fit: radon.sample_predictive_y0_vi(vi_fit, \n",
    "                                                    FLOOR_MEASURE_TRAIN, U_TRAIN, COUNTY_TRAIN)\n",
    "    sample_predictive_y0_test = lambda vi_fit: radon.sample_predictive_y0_vi(vi_fit, \n",
    "                                                    FLOOR_MEASURE_TEST, U_TEST, COUNTY_TEST)\n",
    "    optimal_h_bayes_estimator_np = lambda ys: tonumpy(optimal_h_bayes_estimator(totorch(ys)))\n",
    "\n",
    "    hstar1_train, _ = regularization.get_bootstrap_decisions(sm, stan_data_generator, \n",
    "                                        sample_predictive_y0_train, sample_predictive_y0_test, \n",
    "                                        optimal_h_bayes_estimator_np, num_repetitions=5, iter=VI_NITER)        \n",
    "    LAMBDA_TRAIN = LAMBDA * 1.0/(2. * hstar1_train.std(0)**2)\n",
    "        \n",
    "    if REGULARIZATION.startswith(\"boot1\"):\n",
    "        print(\"[regularization] Bootstrap-based regularization: overwritting regularization mean\")\n",
    "        HSTAR_TRAIN = totorch( hstar1_train.mean(0) )\n",
    "        \n",
    "elif REGULARIZATION.startswith(\"const\"): \n",
    "    LAMBDA_TRAIN = regularization.get_regularization_constant(ys_train, LAMBDA)\n",
    "elif REGULARIZATION.startswith(\"std\"): \n",
    "    LAMBDA_TRAIN = regularization.get_regularization_std(ys_train, LAMBDA)\n",
    "elif REGULARIZATION.startswith(\"qdiff\"): \n",
    "    LAMBDA_TRAIN = regularization.get_regularization_qdiff(ys_train, LAMBDA)\n",
    "else:    \n",
    "    raise Exception(\"[regularization] Unknown regularization method name: %s\" % method)    \n",
    "\n",
    "LAMBDA_TRAIN = totorch(LAMBDA_TRAIN)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Quantile optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def training_empirical_risk(q):    \n",
    "    q = max(0., min(1., q[0]))\n",
    "    h = losses.tilted_optimal_h(ys_train, q) # obtaing a quantile in a very convoluted way :)\n",
    "    risk = empirical_risk(h, Y_TRAIN).item() \n",
    "    regularizer = ( LAMBDA_TRAIN * (h-HSTAR_TRAIN)**2 ).mean()\n",
    "    obj = risk + regularizer\n",
    "    print(\"evaluating training risk @ q=%.2f => risk=%.3f => obj=%.3f\" % (q, risk, obj))\n",
    "    return obj                      "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.579\n",
      "evaluating training risk @ q=0.53 => risk=0.268 => obj=0.648\n",
      "evaluating training risk @ q=0.47 => risk=0.271 => obj=0.733\n",
      "evaluating training risk @ q=0.51 => risk=0.268 => obj=0.597\n",
      "evaluating training risk @ q=0.49 => risk=0.270 => obj=0.645\n",
      "evaluating training risk @ q=0.51 => risk=0.269 => obj=0.579\n",
      "evaluating training risk @ q=0.49 => risk=0.269 => obj=0.594\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.579\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.579\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.594\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.579\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.579\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.594\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.579\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.579\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.594\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.579\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.579\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.594\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.579\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.579\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.594\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.579\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.579\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.594\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.579\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.579\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.594\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.579\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.579\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.594\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.579\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.579\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.594\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.579\n",
      "evaluating training risk @ q=0.50 => risk=0.269 => obj=0.579\n",
      "Optimization terminated successfully.\n",
      "         Current function value: 0.579088\n",
      "         Iterations: 13\n",
      "         Function evaluations: 36\n"
     ]
    }
   ],
   "source": [
    "from scipy.optimize import minimize\n",
    "res = minimize(training_empirical_risk, [TILTED_Q], \n",
    "               method='Nelder-Mead', options={'xtol': 1e-5, 'disp': True, \"maxiter\": 100})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Quantile optimal risk: train:0.269 test:0.2884\n"
     ]
    }
   ],
   "source": [
    "q = max(0., min(1., res[\"x\"][0]))\n",
    "q_train_risk = empirical_risk(losses.tilted_optimal_h(ys_train, q), Y_TRAIN).item() \n",
    "q_test_risk = empirical_risk(losses.tilted_optimal_h(ys_test, q), Y_TEST).item() \n",
    "print(\"Quantile optimal risk: train:%.3f test:%.4f\" % (q_train_risk, q_test_risk))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Decision maker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x7fb5ba09b5d0>"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(SEED)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare features \n",
    "Quantiles = np.array([np.percentile(ys_train, int(q*100), axis=0) for q in np.arange(0., 1.01, 0.05)])\n",
    "X_train = Quantiles\n",
    "Quantiles = np.array([np.percentile(ys_test, int(q*100), axis=0) for q in np.arange(0., 1.01, 0.05)])\n",
    "X_test = Quantiles\n",
    "\n",
    "X_train, X_test = totorch(X_train.T), totorch(X_test.T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sequential(\n",
      "  (0): Linear(in_features=21, out_features=20, bias=True)\n",
      "  (1): ReLU()\n",
      "  (2): Linear(in_features=20, out_features=10, bias=True)\n",
      "  (3): ReLU()\n",
      "  (4): Linear(in_features=10, out_features=1, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "decision_maker = nn.Sequential(\n",
    "  nn.Linear(X_train.shape[1], 20),\n",
    "  nn.ReLU(),\n",
    "  nn.Linear(20, 10),\n",
    "  nn.ReLU(),\n",
    "  nn.Linear(10, 1)\n",
    ")\n",
    "\n",
    "print(decision_maker)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.02s] 0. iteration: training batch risk: 0.754 regularizer: 113.243 train risk: 0.743 test risk: 0.712\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/anonymized/anaconda3/lib/python3.6/site-packages/ipykernel_launcher.py:14: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[3.40s] 1000. iteration: training batch risk: 0.267 regularizer: 0.233 train risk: 0.267 test risk: 0.289\n",
      "[6.13s] 2000. iteration: training batch risk: 0.267 regularizer: 0.223 train risk: 0.267 test risk: 0.289\n",
      "[8.71s] 3000. iteration: training batch risk: 0.267 regularizer: 0.218 train risk: 0.267 test risk: 0.289\n",
      "[11.15s] 4000. iteration: training batch risk: 0.266 regularizer: 0.211 train risk: 0.266 test risk: 0.289\n",
      "[13.76s] 5000. iteration: training batch risk: 0.266 regularizer: 0.196 train risk: 0.266 test risk: 0.289\n",
      "[15.88s] 6000. iteration: training batch risk: 0.266 regularizer: 0.194 train risk: 0.266 test risk: 0.290\n",
      "[18.58s] 7000. iteration: training batch risk: 0.266 regularizer: 0.192 train risk: 0.266 test risk: 0.289\n",
      "[20.62s] 8000. iteration: training batch risk: 0.266 regularizer: 0.192 train risk: 0.266 test risk: 0.290\n",
      "[23.89s] 9000. iteration: training batch risk: 0.266 regularizer: 0.192 train risk: 0.266 test risk: 0.290\n",
      "[27.86s] 10000. iteration: training batch risk: 0.266 regularizer: 0.191 train risk: 0.266 test risk: 0.290\n",
      "[30.53s] 11000. iteration: training batch risk: 0.266 regularizer: 0.189 train risk: 0.266 test risk: 0.289\n",
      "[33.04s] 12000. iteration: training batch risk: 0.266 regularizer: 0.198 train risk: 0.265 test risk: 0.290\n",
      "[35.66s] 13000. iteration: training batch risk: 0.266 regularizer: 0.192 train risk: 0.265 test risk: 0.290\n",
      "[37.43s] 14000. iteration: training batch risk: 0.266 regularizer: 0.197 train risk: 0.265 test risk: 0.290\n",
      "[39.46s] 15000. iteration: training batch risk: 0.266 regularizer: 0.192 train risk: 0.265 test risk: 0.290\n",
      "[42.03s] 16000. iteration: training batch risk: 0.266 regularizer: 0.191 train risk: 0.265 test risk: 0.290\n",
      "[44.39s] 17000. iteration: training batch risk: 0.266 regularizer: 0.189 train risk: 0.265 test risk: 0.290\n",
      "[47.55s] 18000. iteration: training batch risk: 0.266 regularizer: 0.186 train risk: 0.266 test risk: 0.289\n",
      "[50.03s] 19000. iteration: training batch risk: 0.266 regularizer: 0.165 train risk: 0.266 test risk: 0.290\n"
     ]
    }
   ],
   "source": [
    "optimizer = torch.optim.Adam(decision_maker.parameters(), lr=0.01)\n",
    "start = time.time()\n",
    "for i in range(20000):\n",
    "        h = decision_maker(X_train).view(-1)\n",
    "        train_loss = empirical_risk(h, Y_TRAIN)  \n",
    "        regularizer = ( totorch(LAMBDA_TRAIN) * (h-HSTAR_TRAIN)**2 ).mean()\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        (train_loss+regularizer).backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        if i%1000==0:\n",
    "            print(\"[%.2fs] %i. iteration: training batch risk: %.3f regularizer: %.3f train risk: %.3f test risk: %.3f\" % \n",
    "              (time.time()-start, i, train_loss.item(), regularizer.item(),\n",
    "               empirical_risk(decision_maker(X_train).view(-1), Y_TRAIN).item(),\n",
    "               empirical_risk(decision_maker(X_test).view(-1), Y_TEST).item()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DM optimal risk: train:0.266 test:0.2904\n"
     ]
    }
   ],
   "source": [
    "dm_train_risk = empirical_risk(decision_maker(X_train).view(-1), Y_TRAIN).item()\n",
    "dm_test_risk = empirical_risk(decision_maker(X_test).view(-1), Y_TEST).item()\n",
    "print(\"DM optimal risk: train:%.3f test:%.4f\" % (dm_train_risk, dm_test_risk))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Store evaluation results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = [[SEED, LOSS, TILTED_Q, REGULARIZATION, LAMBDA, VI_NITER, HMC_NITER,\n",
    "            vi_train_risk, vi_test_risk, \n",
    "            q_train_risk, q_test_risk, \n",
    "            dm_train_risk, dm_test_risk,\n",
    "            hmc_train_risk, hmc_test_risk]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving to radon_5_tilted_0.5_boot1_1.0_10000.csv\n"
     ]
    }
   ],
   "source": [
    "path = \"radon_%i_%s_%s_%s_%s_%s.csv\" % \\\n",
    "            (SEED, LOSS, TILTED_Q, REGULARIZATION, LAMBDA, VI_NITER)\n",
    "print(\"Saving to %s\" % path)\n",
    "pd.DataFrame(results).to_csv(path, header=False, index=False)"
   ]
  },
  {
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}