recalibration_friedman1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "d0dedc23",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"/Users/fm03-mb03/Repositories/proper_calibration_errors\""
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# code taken from Wiedman et al 2021\n",
    "# https://github.com/devmotion/Calibration_ICLR2021\n",
    "\n",
    "using Base.Filesystem\n",
    "\n",
    "using Arrow\n",
    "using CairoMakie\n",
    "using CalibrationErrors\n",
    "using CalibrationErrorsDistributions\n",
    "using CalibrationTests\n",
    "using CSV\n",
    "using DataFrames\n",
    "using Distributions\n",
    "using Flux\n",
    "using Random\n",
    "using ProgressLogging\n",
    "using Query\n",
    "using Optim\n",
    "\n",
    "using Logging: with_logger\n",
    "using TerminalLoggers: TerminalLogger\n",
    "using ColorSchemes: Dark2_8\n",
    "\n",
    "pwd()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "97255c76",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "sample_data (generic function with 1 method)"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "friedman1_var(x) = 10 * sinpi(x[1] * x[2]) + 20 * (x[3] - 1//2)^2 + 10 * x[4] + 5 * x[5]\n",
    "\n",
    "function sample_data(n::Int)\n",
    "    ## sample inputs\n",
    "    xs = rand(10, n)\n",
    "\n",
    "    ## sample targets\n",
    "    ys = map(eachcol(xs)) do x\n",
    "        # sigma = a + b * x[6]\n",
    "        # epsilon ~ N(0, sigma^2)\n",
    "        return friedman1_var(x) + randn() * (0.5 + 1*x[6])\n",
    "    end\n",
    "\n",
    "    return xs, ys\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "bad90b11",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_models = 5\n",
    "\n",
    "Random.seed!(100)\n",
    "train_data = sample_data(100);\n",
    "\n",
    "# validation data\n",
    "\n",
    "Random.seed!(200)\n",
    "val_data = sample_data(100);\n",
    "\n",
    "# For the evaluation of the models we use another data set of 100 samples that is\n",
    "# sampled according to the same law.\n",
    "\n",
    "Random.seed!(300)\n",
    "test_data = sample_data(100);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "e09f3d91",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "mean_var (generic function with 1 method)"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# constant for numerical stability of division\n",
    "eps = 1e-10\n",
    "\n",
    "function glorot_uniform(nout::Int, nin::Int)\n",
    "    return (rand(nout, nin) .- 0.5) .* sqrt(24 / (nout + nin))\n",
    "end\n",
    "\n",
    "function nn_model()\n",
    "    ## initial parameters\n",
    "    f = Chain(\n",
    "        Dense(10 => 200, relu; init=glorot_uniform),\n",
    "        Dense(200 => 50, relu; init=glorot_uniform),\n",
    "        Dense(50 => 2; init=glorot_uniform),\n",
    "    )\n",
    "    # due to a lack of julia expertise, the variance output is in the log space\n",
    "    # and will be transformed in each loss later\n",
    "    return f\n",
    "end\n",
    "\n",
    "function pmcc(ps, ys)\n",
    "    vars = exp.(ps[2,:]) .+ eps\n",
    "    scores = (ps[1,:] .- ys).^2 ./ vars .+ log.(vars)\n",
    "    return mean(scores)\n",
    "end\n",
    "\n",
    "function skce_biased(ps, ys)\n",
    "    kernel = WassersteinExponentialKernel() ⊗ SqExponentialKernel()\n",
    "    estimator = BiasedSKCE(kernel)\n",
    "    n = size(ps, 2)\n",
    "    predictions = [Normal(ps[1, i], sqrt(exp(ps[2, i]) + eps)) for i in 1:n]\n",
    "    return calibrationerror(estimator, vec(predictions), ys)\n",
    "end\n",
    "\n",
    "function skce_unbiased(ps, ys)\n",
    "    kernel = WassersteinExponentialKernel() ⊗ SqExponentialKernel()\n",
    "    estimator = UnbiasedSKCE(kernel)\n",
    "    n = size(ps, 2)\n",
    "    predictions = [Normal(ps[1, i], sqrt(exp(ps[2, i]) + eps)) for i in 1:n]\n",
    "    return calibrationerror(estimator, vec(predictions), ys)\n",
    "end\n",
    "\n",
    "function mse(ps, ys)\n",
    "    scores = (ps[1,:] .- ys).^2\n",
    "    return mean(scores)\n",
    "end\n",
    "\n",
    "function mean_var(ps)\n",
    "    vars = exp.(ps[2,:]) .+ eps\n",
    "    return mean(vars)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "d809226e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "recal (generic function with 1 method)"
      ]
     },
     "execution_count": 77,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function recal(f, val_xs, val_ys)\n",
    "    preds = f(val_xs)\n",
    "    \n",
    "    function helper(w)\n",
    "        cals = transpose(hcat(preds[1, :], log.(exp.(preds[2, :]) .* w[1] .+ w[2])))\n",
    "        return pmcc(cals, val_ys)\n",
    "    end\n",
    "\n",
    "    lower = [eps, eps]\n",
    "    upper = [Inf, Inf]\n",
    "    res = optimize(helper, lower, upper, [1.0, 2*eps], Fminbox(LBFGS()))\n",
    "    w = Optim.minimizer(res)\n",
    "\n",
    "    return x -> transpose(hcat(x[1, :], log.(exp.(preds[2, :]) .* w[1] .+ w[2])))\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "3ed4196d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "train (generic function with 1 method)"
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# ## Training\n",
    "#\n",
    "# We use a maximum likelihood approach and train the parameters $\\theta$ of the model\n",
    "# for 1000 iterations by minimizing the DSS on the training data set\n",
    "# using ADAM.\n",
    "#\n",
    "# We train 5 models and compute the predicted distributions on the training and test data sets\n",
    "# in each iteration step.\n",
    "#\n",
    "# The initial values of the weight matrices of the neural networks are sampled from the\n",
    "# [uniform Glorot initialization](http://proceedings.mlr.press/v9/glorot10a/glorot10a.pdf)\n",
    "# and the offset vectors are initialized with zeros. The model parameters are learnt by\n",
    "# iteratively minimizing the DSS on the training data set.\n",
    "# The parameters of the neural networks are trained by gradient descent with the\n",
    "# [Adam optimization algorithm](https://arxiv.org/pdf/1412.6980.pdf) (default\n",
    "# settings in [Flux.jl](https://github.com/FluxML/Flux.jl)).\n",
    "\n",
    "function train(id, (train_xs, train_ys), (val_xs, val_ys), (test_xs, _))\n",
    "    ## check if file exists\n",
    "    filename = joinpath(\"data\", \"friedman1_var\", \"predictions_id=$(id).arrow\")\n",
    "    isfile(filename) && return nothing\n",
    "\n",
    "    ## compute the predictions of the initial neural network\n",
    "    f = nn_model()\n",
    "    train_preds = f(train_xs)\n",
    "    # get recalibration model\n",
    "    g = recal(f, val_xs, val_ys)\n",
    "    test_rc_preds = g(f(test_xs))\n",
    "    test_preds = f(test_xs)\n",
    "\n",
    "    ## save the initial model and its predictions\n",
    "    niters = 1000\n",
    "    train_predss = Vector{typeof(train_preds)}(undef, niters + 1)\n",
    "    test_rc_predss = Vector{typeof(test_rc_preds)}(undef, niters + 1)\n",
    "    test_predss = Vector{typeof(test_preds)}(undef, niters + 1)\n",
    "    train_predss[1] = train_preds\n",
    "    test_rc_predss[1] = test_rc_preds\n",
    "    test_predss[1] = test_preds\n",
    "\n",
    "    ## train with ADAM\n",
    "    params = Flux.Params(Flux.params(f))\n",
    "    opt = ADAM()\n",
    "    @progress name = \"training (id = $id)\" for i in 2:(niters + 1)\n",
    "        ## compute gradients\n",
    "        gradients = gradient(params) do\n",
    "            return pmcc(f(train_xs), train_ys)\n",
    "        end\n",
    "\n",
    "        ## update the parameters\n",
    "        Flux.Optimise.update!(opt, params, gradients)\n",
    "\n",
    "        ## save the model and its predictions\n",
    "        train_predss[i] = f(train_xs)\n",
    "        # get recalibration model\n",
    "        g = recal(f, val_xs, val_ys)\n",
    "        test_rc_predss[i] = g(f(test_xs))\n",
    "        test_predss[i] = f(test_xs)\n",
    "    end\n",
    "\n",
    "    ## save the predictions\n",
    "    mkpath(dirname(filename))\n",
    "    preds = (train_preds=train_predss, test_rc_preds=test_rc_predss, test_preds=test_predss)\n",
    "    Arrow.write(filename, preds)\n",
    "\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "f0b8323e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "┌ Info: training NN model: run 1\n",
      "└ @ Main In[79]:3\n",
      "┌ Info: training NN model: run 2\n",
      "└ @ Main In[79]:3\n",
      "┌ Info: training NN model: run 3\n",
      "└ @ Main In[79]:3\n",
      "┌ Info: training NN model: run 4\n",
      "└ @ Main In[79]:3\n",
      "┌ Info: training NN model: run 5\n",
      "└ @ Main In[79]:3\n"
     ]
    }
   ],
   "source": [
    "Random.seed!(100)\n",
    "for (id, seed) in enumerate(rand(UInt, n_models))\n",
    "    @info \"training NN model: run $id\"\n",
    "    Random.seed!(seed)\n",
    "    train(id, train_data, val_data, test_data)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "2c8a451a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "evaluate_stats (generic function with 1 method)"
      ]
     },
     "execution_count": 80,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# ## Evaluations\n",
    "#\n",
    "# SKCE (biased & unbiased), DSS, MSE, avg predicted variance\n",
    "\n",
    "function evaluate_models(dataset, id, ys)\n",
    "    ## output file\n",
    "    out = joinpath(\"data\", \"friedman1_var\", \"statistics_id=$(id)_dataset=$(dataset).csv\")\n",
    "    isfile(out) && return nothing\n",
    "\n",
    "    ## load data\n",
    "    filename = joinpath(\"data\", \"friedman1_var\", \"predictions_id=$(id).arrow\")\n",
    "    isfile(filename) || error(\"predictions for run \", id, \" not found\")\n",
    "    tbl = Arrow.Table(filename)\n",
    "    predss = getproperty(tbl, Symbol(dataset, :_preds))\n",
    "    predictionss = map(predss) do preds\n",
    "        return map(preds) do pred\n",
    "            return pred\n",
    "        end\n",
    "    end\n",
    "\n",
    "    return evaluate_stats(out, predictionss, ys)\n",
    "end\n",
    "\n",
    "function evaluate_stats(file, predictionss, ys)\n",
    "    mkpath(dirname(file))\n",
    "    open(file, \"w\") do f\n",
    "        ## print headers\n",
    "        println(f, \"iteration,statistic,estimate\")\n",
    "\n",
    "        @progress name = \"iterations\" for (i, predictions) in enumerate(predictionss)\n",
    "            preds = reshape(predictions, 2, trunc(Int, length(predictions)/2))   \n",
    "            ## mean squared error\n",
    "            mse_v = mse(preds, ys)\n",
    "            println(f, i - 1, \",MSE,\", mse_v)            \n",
    "            \n",
    "            ## mean-variance score\n",
    "            pmcc_v = pmcc(preds, ys)\n",
    "            println(f, i - 1, \",PMCC,\", pmcc_v)\n",
    "\n",
    "            ## unbiased estimator of SKCE\n",
    "            skce = skce_unbiased(preds, ys)\n",
    "            println(f, i - 1, \",SKCE (unbiased),\", skce)\n",
    "\n",
    "            ## biased estimator of SKCE\n",
    "            skce = skce_biased(preds, ys)\n",
    "            println(f, i - 1, \",SKCE (biased),\", skce)\n",
    "\n",
    "            ## mean predicted var\n",
    "            var = mean_var(preds)\n",
    "            println(f, i - 1, \",Avg Var,\", var)\n",
    "        end\n",
    "    end\n",
    "\n",
    "    return nothing\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "3faa9232",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "┌ Info: evaluating training statistics: run 1\n",
      "└ @ Main In[81]:5\n",
      "┌ Info: evaluating test statistics: run 1\n",
      "└ @ Main In[81]:10\n",
      "┌ Info: evaluating test rc statistics: run 1\n",
      "└ @ Main In[81]:15\n",
      "┌ Info: evaluating training statistics: run 2\n",
      "└ @ Main In[81]:5\n",
      "┌ Info: evaluating test statistics: run 2\n",
      "└ @ Main In[81]:10\n",
      "┌ Info: evaluating test rc statistics: run 2\n",
      "└ @ Main In[81]:15\n",
      "┌ Info: evaluating training statistics: run 3\n",
      "└ @ Main In[81]:5\n",
      "┌ Info: evaluating test statistics: run 3\n",
      "└ @ Main In[81]:10\n",
      "┌ Info: evaluating test rc statistics: run 3\n",
      "└ @ Main In[81]:15\n",
      "┌ Info: evaluating training statistics: run 4\n",
      "└ @ Main In[81]:5\n",
      "┌ Info: evaluating test statistics: run 4\n",
      "└ @ Main In[81]:10\n",
      "┌ Info: evaluating test rc statistics: run 4\n",
      "└ @ Main In[81]:15\n",
      "┌ Info: evaluating training statistics: run 5\n",
      "└ @ Main In[81]:5\n",
      "┌ Info: evaluating test statistics: run 5\n",
      "└ @ Main In[81]:10\n",
      "┌ Info: evaluating test rc statistics: run 5\n",
      "└ @ Main In[81]:15\n"
     ]
    }
   ],
   "source": [
    "Random.seed!(300)\n",
    "for (id, seed) in enumerate(rand(UInt, n_models))\n",
    "    \n",
    "    ## evaluate models on training data set\n",
    "    @info \"evaluating training statistics: run $id\"\n",
    "    Random.seed!(seed)\n",
    "    evaluate_models(\"train\", id, train_data[2])\n",
    "\n",
    "    ## evaluate models on test data set\n",
    "    @info \"evaluating test statistics: run $id\"\n",
    "    Random.seed!(seed)\n",
    "    evaluate_models(\"test\", id, test_data[2])\n",
    "\n",
    "    ## evaluate models on recalibrated test data set\n",
    "    @info \"evaluating test rc statistics: run $id\"\n",
    "    Random.seed!(seed)\n",
    "    evaluate_models(\"test_rc\", id, test_data[2])\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a988251",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.7.2",
   "language": "julia",
   "name": "julia-1.7"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}