notebooks/SurrogateModelOptim.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MSM.jl and SurrogateModelOptim.jl\n",
    "\n",
    "This notebook shows how one can estimate a model using [MSM.jl](https://github.com/JulienPascal/MSM.jl) and [SurrogateModelOptim.jl](https://github.com/MrUrq/SurrogateModelOptim.jl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because the surrogate model can be **trained in parallel**, let's first increase the number of workers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "using ClusterManagers\n",
    "using Distributed\n",
    "OnCluster = false #set to false to run locally\n",
    "addWorkers = false #set to false to run serially\n",
    "println(\"OnCluster = $(OnCluster)\")\n",
    "\n",
    "# Current number of workers\n",
    "#--------------------------\n",
    "currentWorkers = nworkers()\n",
    "println(\"Initial number of workers = $(currentWorkers)\")\n",
    "\n",
    "# Increase the number of workers available\n",
    "#-----------------------------------------\n",
    "maxNumberWorkers = 3\n",
    "if addWorkers == true\n",
    "\tif OnCluster == true\n",
    "\t  addprocs(SlurmManager(maxNumberWorkers))\n",
    "\telse\n",
    "\t  addprocs(maxNumberWorkers)\n",
    "\tend\n",
    "end\n",
    "\n",
    "\n",
    "# Sanity checks\n",
    "#-------------\n",
    "hosts = []\n",
    "pids = []\n",
    "for i in workers()\n",
    "\thost, pid = fetch(@spawnat i (gethostname(), getpid()))\n",
    "\tprintln(\"Hello I am worker $(i), my host is $(host)\")\n",
    "\tpush!(hosts, host)\n",
    "\tpush!(pids, pid)\n",
    "end\n",
    "\n",
    "currentWorkers = nworkers()\n",
    "println(\"Number of workers = $(currentWorkers)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's construct the MSM objective function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "using Plots\n",
    "using ParallelDataTransfer\n",
    "\n",
    "@everywhere using MSM\n",
    "@everywhere using DataStructures\n",
    "@everywhere using OrderedCollections\n",
    "@everywhere using Distributions\n",
    "@everywhere using Random\n",
    "@everywhere using DataStructures\n",
    "@everywhere using Statistics\n",
    "@everywhere using LinearAlgebra\n",
    "\n",
    "Random.seed!(1234)  #for replicability reasons\n",
    "T = 100000          #number of periods\n",
    "P = 2               #number of dependent variables\n",
    "beta0 = rand(P)     #choose true coefficients by drawing from a uniform distribution on [0,1]\n",
    "alpha0 = rand(1)[]  #intercept\n",
    "theta0 = 0.0        #coefficient to create serial correlation in the error terms\n",
    "println(\"True intercept = $(alpha0)\")\n",
    "println(\"True coefficient beta0 = $(beta0)\")\n",
    "println(\"Serial correlation coefficient theta0 = $(theta0)\")\n",
    "\n",
    "# Generation of error terms\n",
    "# row = individual dimension\n",
    "# column = time dimension \n",
    "U = zeros(T)\n",
    "d = Normal()\n",
    "U[1] = rand(d, 1)[] #first error term\n",
    "# loop over time periods\n",
    "for t = 2:T\n",
    "    U[t] = rand(d, 1)[] + theta0*U[t-1]\n",
    "end\n",
    "\n",
    "# Let's simulate the dependent variables x_t\n",
    "x = zeros(T, P)\n",
    "\n",
    "d = Uniform(0, 5)\n",
    "for p = 1:P  \n",
    "    x[:,p] = rand(d, T)\n",
    "end\n",
    "\n",
    "# Let's calculate the resulting y_t\n",
    "y = zeros(T)\n",
    "\n",
    "for t=1:T\n",
    "    y[t] = alpha0 + x[t,1]*beta0[1] + x[t,2]*beta0[2] + U[t]\n",
    "end\n",
    "\n",
    "# Send simulated data to workers\n",
    "sendto(workers(), y=y)\n",
    "sendto(workers(), x=x)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize data\n",
    "p1 = scatter(x[1:100,1], y[1:100], xlabel = \"x1\", ylabel = \"y\", legend=:none, smooth=true)\n",
    "p2 = scatter(x[1:100,2], y[1:100], xlabel = \"x2\", ylabel = \"y\", legend=:none, smooth=true)\n",
    "plot(p1, p2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define locally\n",
    "optionsSMM = MSMOptions(maxFuncEvals=100, globalOptimizer = :dxnes, localOptimizer = :NelderMead)\n",
    "myProblem = MSMProblem(options = optionsSMM);\n",
    "\n",
    "# Send to workers\n",
    "sendto(workers(), optionsSMM=optionsSMM)\n",
    "sendto(workers(), myProblem=myProblem)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Priors\n",
    "# Define locally\n",
    "dictPriors = OrderedDict{String,Array{Float64,1}}()\n",
    "dictPriors[\"alpha\"] = [0.5, 0.001, 1.0]\n",
    "dictPriors[\"beta1\"] = [0.5, 0.001, 1.0]\n",
    "dictPriors[\"beta2\"] = [0.5, 0.001, 1.0]\n",
    "# Send to workers\n",
    "sendto(workers(), dictPriors=dictPriors)\n",
    "\n",
    "# Empirical moments\n",
    "# Define locally\n",
    "dictEmpiricalMoments = OrderedDict{String,Array{Float64,1}}()\n",
    "dictEmpiricalMoments[\"mean\"] = [mean(y)] #informative on the intercept\n",
    "dictEmpiricalMoments[\"mean^2\"] = [mean(y.^2)] #informative on the intercept\n",
    "dictEmpiricalMoments[\"mean^3\"] = [mean(y.^3)] #informative on the intercept\n",
    "dictEmpiricalMoments[\"mean_x1y\"] = [mean(x[:,1] .* y)] #informative on betas\n",
    "dictEmpiricalMoments[\"mean_x2y\"] = [mean(x[:,2] .* y)] #informative on betas\n",
    "dictEmpiricalMoments[\"mean_x1y^2\"] = [mean((x[:,1] .* y).^2)] #informative on betas\n",
    "dictEmpiricalMoments[\"mean_x2y^2\"] = [mean((x[:,2] .* y).^2)] #informative on betas\n",
    "# Send to workers\n",
    "sendto(workers(), dictEmpiricalMoments=dictEmpiricalMoments)\n",
    "\n",
    "W = Matrix(1.0 .* I(length(dictEmpiricalMoments)))#initialization\n",
    "#Special case: diagonal matrix\n",
    "#(you may choose something else)\n",
    "for (indexMoment, k) in enumerate(keys(dictEmpiricalMoments))\n",
    "    W[indexMoment,indexMoment] = 1.0/(dictEmpiricalMoments[k][1])^2\n",
    "end\n",
    "sendto(workers(), W=W)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@everywhere set_priors!(myProblem, dictPriors)\n",
    "@everywhere set_empirical_moments!(myProblem, dictEmpiricalMoments)\n",
    "@everywhere set_weight_matrix!(myProblem, W)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# x[1] corresponds to the intercept, x[2] corresponds to beta1, x[3] corresponds to beta2\n",
    "@everywhere function functionLinearModel(x; uniform_draws::Array{Float64,1}, simX::Array{Float64,2}, nbDraws::Int64 = length(uniform_draws), burnInPerc::Int64 = 0)\n",
    "    T = nbDraws\n",
    "    P = 2       #number of dependent variables\n",
    "\n",
    "    alpha = x[1]\n",
    "    beta = x[2:end]\n",
    "    theta = 0.0     #coefficient to create serial correlation in the error terms\n",
    "\n",
    "    # Creation of error terms\n",
    "    # row = individual dimension\n",
    "    # column = time dimension\n",
    "    U = zeros(T)\n",
    "    d = Normal()\n",
    "    # Inverse cdf (i.e. quantile)\n",
    "    gaussian_draws = quantile.(d, uniform_draws)\n",
    "    U[1] = gaussian_draws[1] #first error term\n",
    "\n",
    "    # loop over time periods\n",
    "    for t = 2:T\n",
    "        U[t] = gaussian_draws[t] + theta*U[t-1]\n",
    "    end\n",
    "\n",
    "    # Let's calculate the resulting y_t\n",
    "    y = zeros(T)\n",
    "\n",
    "    for t=1:T\n",
    "        y[t] = alpha + simX[t,1]*beta[1] + simX[t,2]*beta[2] + U[t]\n",
    "    end\n",
    "\n",
    "    # Get rid of the burn-in phase:\n",
    "    #------------------------------\n",
    "    startT = max(1, Int(nbDraws * (burnInPerc / 100)))\n",
    "\n",
    "    # Moments:\n",
    "    #---------\n",
    "    output = OrderedDict{String,Float64}()\n",
    "    output[\"mean\"] = mean(y[startT:nbDraws])\n",
    "    output[\"mean^2\"] = mean(y[startT:nbDraws].^2)\n",
    "    output[\"mean_x1y\"] = mean(simX[startT:nbDraws,1] .* y[startT:nbDraws])\n",
    "    output[\"mean_x2y\"] = mean(simX[startT:nbDraws,2] .* y[startT:nbDraws])\n",
    "    output[\"mean_x1y^2\"] = mean((simX[startT:nbDraws,1] .* y[startT:nbDraws]).^2)\n",
    "    output[\"mean_x2y^2\"] = mean((simX[startT:nbDraws,2] .* y[startT:nbDraws]).^2)\n",
    "\n",
    "    return output\n",
    "end\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's freeze the randomness during the minimization\n",
    "d_Uni = Uniform(0,1)\n",
    "nbDraws = T #number of draws in the simulated data\n",
    "uniform_draws = rand(d_Uni, nbDraws)\n",
    "simX = zeros(length(uniform_draws), 2)\n",
    "d = Uniform(0, 5)\n",
    "for p = 1:2\n",
    "  simX[:,p] = rand(d, length(uniform_draws))\n",
    "end\n",
    "\n",
    "# Send to workers\n",
    "sendto(workers(), simX=simX)\n",
    "sendto(workers(), uniform_draws=uniform_draws)\n",
    "\n",
    "# Construct the objective function everywhere\n",
    "@everywhere set_simulate_empirical_moments!(myProblem, x -> functionLinearModel(x, uniform_draws = uniform_draws, simX = simX))\n",
    "@everywhere construct_objective_function!(myProblem)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Test of the objective function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Safety check: value on the master node == values on slave nodes?\n",
    "using Test\n",
    "val_local = myProblem.objective_function([dictPriors[k][1] for k in keys(dictPriors)]); #local execution\n",
    "val_workers = [];\n",
    "for w in workers() #Execution on workers\n",
    "    push!(val_workers, @fetchfrom w myProblem.objective_function([dictPriors[k][1] for k in keys(dictPriors)]))\n",
    "end\n",
    "for (wIndex, w) in enumerate(workers())\n",
    "    @test abs(val_local - val_workers[wIndex]) < 10e-10\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "using SurrogateModelOptim\n",
    "\n",
    "search_range=generate_bbSearchRange(myProblem)\n",
    "\n",
    "@time result = smoptimize(x -> myProblem.objective_function(x), search_range;\n",
    "                    options=SurrogateModelOptim.Options(\n",
    "                    iterations=50,\n",
    "                    num_interpolants=1*nworkers(), \n",
    "                    num_start_samples=5,\n",
    "                    create_final_surrogate=true,));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Best candidate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "minimizer = best_candidate(result)\n",
    "min_function = best_fitness(result)\n",
    "\n",
    "println(\"Estimated value for alpha = $(minimizer[1])\")\n",
    "println(\"True value for alpha = $(alpha0[1]) \\n\")\n",
    "\n",
    "println(\"Estimated value for beta1 = $(minimizer[2])\")\n",
    "println(\"True value for beta1 = $(beta0[1]) \\n\")\n",
    "\n",
    "println(\"Estimated value for beta2 = $(minimizer[3])\")\n",
    "println(\"True value for beta2 = $(beta0[2]) \\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting the surrogate model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "lower_bound = create_lower_bound(myProblem)\n",
    "upper_bound = create_upper_bound(myProblem)\n",
    "\n",
    "using Plots\n",
    "using Statistics\n",
    "\n",
    "gr()\n",
    "\n",
    "p1 = plot(collect(lower_bound[2]:0.1:upper_bound[2]), collect(lower_bound[3]:0.1:upper_bound[3]), (x, y) -> median(result.sm_interpolant([alpha0; x; y])), linetype=:surface)\n",
    "scatter!([minimizer[2]], [minimizer[3]], [min_function], label=\"min\")\n",
    "\n",
    "p2 = contour(collect(lower_bound[2]:0.1:upper_bound[2]), collect(lower_bound[3]:0.1:upper_bound[3]), (x, y) -> median(result.sm_interpolant([alpha0; x; y])))\n",
    "scatter!([minimizer[2]], [minimizer[3]], label=\"min\")\n",
    "\n",
    "plot(p1, p2, title=\"Surrogate Model\") \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p1 = plot(collect(lower_bound[1]:0.1:upper_bound[1]), collect(lower_bound[2]:0.1:upper_bound[2]), (x, y) -> median(result.sm_interpolant([x; y; beta0[2]])), linetype=:surface)\n",
    "scatter!([minimizer[1]], [minimizer[2]], [min_function], label=\"min\")\n",
    "\n",
    "p2 = contour(collect(lower_bound[1]:0.1:upper_bound[1]), collect(lower_bound[2]:0.1:upper_bound[2]), (x, y) -> median(result.sm_interpolant([x; y; beta0[2]])))\n",
    "scatter!([minimizer[1]], [minimizer[2]], label=\"min\")\n",
    "\n",
    "plot(p1, p2, title=\"Surrogate Model\") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "versioninfo()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.5.2",
   "language": "julia",
   "name": "julia-1.5"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}