{ "cells": [ { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "%%capture \n", "\n", "%matplotlib inline\n", "\n", "import IPython\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", "import pandas as pd\n", "import requests\n", "import tensorflow as tf\n", "import tensorflow_probability as tfp\n", "import warnings\n", "\n", "from tensorflow_probability import edward2 as ed\n", "tfd = tfp.distributions\n", "\n", "from keras.constraints import non_neg\n", "\n", "import statsmodels.api as sm\n", "import statsmodels.formula.api as smf\n", " \n", "plt.style.use('ggplot')" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
iso3yearln_gdppcln_TFRln_popintercept
0019607.6260122.00821416.0123301.0
1019617.6139112.00821416.0310951.0
2019627.6100662.00821416.0504451.0
3019637.6074012.00821416.0703701.0
4019647.6035152.00821416.0908641.0
\n", "
" ], "text/plain": [ " iso3 year ln_gdppc ln_TFR ln_pop intercept\n", "0 0 1960 7.626012 2.008214 16.012330 1.0\n", "1 0 1961 7.613911 2.008214 16.031095 1.0\n", "2 0 1962 7.610066 2.008214 16.050445 1.0\n", "3 0 1963 7.607401 2.008214 16.070370 1.0\n", "4 0 1964 7.603515 2.008214 16.090864 1.0" ] }, "execution_count": 2, "metadata": {}, "output_type": "execute_result" } ], "source": [ "## Get GDP data for testing\n", "data = pd.read_csv('/home/j/Project/IRH/Forecasting/gdp/data/RT_2018_GDP_use.csv')\n", "\n", "## Prep data\n", "data = data[['iso3', 'year', 'ln_gdppc', 'ln_TFR', 'ln_pop']]\n", "data['intercept'] = 1.\n", "data = data.dropna()\n", "\n", "# Remap categories to start from 0 and end at max(category).\n", "data['iso3'] = data['iso3'].astype('category').cat.codes\n", "\n", "## Number of REs\n", "n_res = max(data.iso3) + 1\n", "\n", "data.head()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Model:\n", "$$ ln(GDPpc) = \\alpha + \\alpha_i + \\beta \\ln(TFR) + \\epsilon_{i,t} $$ \n", "$$ \\epsilon_{i,t} \\sim \\mathcal{N}(0, \\sigma^2) $$ \n", "$$ \\alpha_i \\sim \\mathcal{N}(0, \\sigma_a^2) $$\n" ] }, { "cell_type": "code", "execution_count": 3, "metadata": { "scrolled": true }, "outputs": [ { "data": { "text/html": [ "\n", "\n", " \n", "\n", "\n", " \n", "\n", "\n", " \n", "\n", "\n", " \n", "\n", "\n", " \n", "\n", "\n", " \n", "\n", "
Model: MixedLM Dependent Variable: ln_gdppc
No. Observations: 11700 Method: REML
No. Groups: 195 Scale: 0.1200
Min. group size: 60 Likelihood: -4798.2007
Max. group size: 60 Converged: Yes
Mean group size: 60.0
\n", "\n", "\n", " \n", "\n", "\n", " \n", "\n", "\n", " \n", "\n", "\n", " \n", "\n", "
Coef. Std.Err. z P>|z| [0.025 0.975]
Intercept 9.929 0.070 142.555 0.000 9.793 10.066
ln_TFR -0.876 0.011 -82.073 0.000 -0.897 -0.855
Group Var 0.909 0.270
" ], "text/plain": [ "\n", "\"\"\"\n", " Mixed Linear Model Regression Results\n", "========================================================\n", "Model: MixedLM Dependent Variable: ln_gdppc \n", "No. Observations: 11700 Method: REML \n", "No. Groups: 195 Scale: 0.1200 \n", "Min. group size: 60 Likelihood: -4798.2007\n", "Max. group size: 60 Converged: Yes \n", "Mean group size: 60.0 \n", "--------------------------------------------------------\n", " Coef. Std.Err. z P>|z| [0.025 0.975]\n", "--------------------------------------------------------\n", "Intercept 9.929 0.070 142.555 0.000 9.793 10.066\n", "ln_TFR -0.876 0.011 -82.073 0.000 -0.897 -0.855\n", "Group Var 0.909 0.270 \n", "========================================================\n", "\n", "\"\"\"" ] }, "execution_count": 3, "metadata": {}, "output_type": "execute_result" } ], "source": [ "## Fit in StatsModels first\n", "md = smf.mixedlm(\"ln_gdppc ~ 1 + ln_TFR\", data, groups=data[\"iso3\"])\n", "sm_fit = md.fit() \n", "sm_fit.summary()" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [], "source": [ "## Get the data for 'features' and 'labels' (xvars and yvar)\n", "get_value = lambda dataframe, key, dtype: dataframe[key].values.astype(dtype)\n", "\n", "## We get the country codes as integers and then attach on the fixed effects\n", "features_train = {\n", " k: get_value(data, key=k, dtype=np.int32)\n", " for k in ['iso3']}\n", "features_train.update({\n", " k: get_value(data, key=k, dtype=np.float64)\n", " for k in ['intercept', 'ln_TFR']})\n", "\n", "## Get yvar\n", "labels_train = get_value(data, key='ln_gdppc', dtype=np.float64)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Define our mixed effects model, returning the prediction for TFR coefficient converted to a 'random' variable" ] }, { "cell_type": "code", "execution_count": 19, "metadata": {}, "outputs": [], "source": [ "## Using Random TFR (with a meh prior)\n", "def linear_mixed_effects_model_randomTFR(features):\n", " \n", " # Set up fixed effects and other parameters.\n", " intercept = tf.get_variable(\"intercept\", []) # alpha in eq\n", " effect_ln_TFR = ed.Normal(loc = tf.zeros(1), scale = tf.ones(1), name = \"effect_ln_TFR\") \n", " \n", " ## The two variances (exp to force positive estimate)\n", " stddev_iso3 = tf.exp(tf.get_variable(\"stddev_iso3\", [])) \n", " model_stddev = tf.exp(tf.get_variable(\"model_stddev\", []))\n", " \n", " ## Set up random effects.\n", " effect_iso3 = ed.MultivariateNormalDiag( \n", " loc=tf.zeros(n_res),\n", " scale_identity_multiplier=stddev_iso3,\n", " name=\"effect_iso3\")\n", " \n", " ## Likelihood\n", " ln_gdppc = ed.Normal(\n", " loc=((intercept * features['intercept']) +\n", " (effect_ln_TFR * features[\"ln_TFR\"]) + \n", " tf.gather(effect_iso3, features[\"iso3\"])),\n", " scale=model_stddev, name=\"ln_gdppc\")\n", " return ln_gdppc\n" ] }, { "cell_type": "code", "execution_count": 20, "metadata": {}, "outputs": [], "source": [ "# Unnormalized target density as a function of states of REs.\n", "def target_log_prob_fn_randomTFR(effect_iso3, effect_ln_TFR): \n", " return log_joint(\n", " features=features_train,\n", " effect_iso3=effect_iso3, \n", " effect_ln_TFR=effect_ln_TFR,\n", " ln_gdppc=labels_train)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### EM and optimization" ] }, { "cell_type": "code", "execution_count": 21, "metadata": {}, "outputs": [], "source": [ "tf.reset_default_graph()\n", "model_template = tf.make_template(\"model\", linear_mixed_effects_model_randomTFR)\n", "\n", "## Make joint LL\n", "log_joint = ed.make_log_joint_fn(model_template)\n", "\n", "## Set up E-step (MCMC) for RE vars\n", "effect_iso3 = tf.get_variable(\"effect_iso3\", [n_res], trainable=False)\n", "effect_ln_TFR = tf.get_variable(\"effect_ln_TFR\", [1], trainable=False)\n", "\n", "## Global step variable holder\n", "global_step = tf.Variable(0, trainable=False, name=\"global_step\")\n", "\n", "## Set up MCMC object\n", "hmc = tfp.mcmc.HamiltonianMonteCarlo(\n", " target_log_prob_fn=target_log_prob_fn_randomTFR,\n", " step_size=0.02,\n", " num_leapfrog_steps=2)\n", "\n", "## RE state to update\n", "current_state = [effect_iso3, effect_ln_TFR]\n", "\n", "## Update state\n", "with warnings.catch_warnings():\n", " warnings.simplefilter(\"ignore\")\n", " next_state, kernel_results = hmc.one_step(current_state=current_state,\n", " previous_kernel_results=hmc.bootstrap_results(current_state))\n", "\n", "## Update Expectation\n", "expectation_update = tf.group(effect_iso3.assign(next_state[0]),\\\n", " effect_ln_TFR.assign(next_state[1]))\n", "\n", "\n", "# Set up M-step (gradient descent), using Adam Optimizer\n", "with tf.control_dependencies([expectation_update]):\n", " loss = -target_log_prob_fn_randomTFR(effect_iso3, effect_ln_TFR) \n", " learning_rate = tf.train.exponential_decay(learning_rate=0.01, \n", " global_step=global_step,\n", " decay_steps=20, \n", " decay_rate=0.90, staircase=True)\n", " optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)\n", " minimization_update = optimizer.minimize(loss)" ] }, { "cell_type": "code", "execution_count": 22, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Warm-Up Iteration: 0 Acceptance Rate: 1.000\n", "Warm-Up Iteration: 250 Acceptance Rate: 0.466\n", "Warm-Up Iteration: 499 Acceptance Rate: 0.414\n", "Iteration: 0 Loss: 23264.594,ln_TFR:4.226, intercept:0.601\n", "Iteration: 10 Loss: 21832.441,ln_TFR:3.764, intercept:1.088\n", "Iteration: 20 Loss: 20732.273,ln_TFR:3.336, intercept:1.546\n", "Iteration: 30 Loss: 20502.531,ln_TFR:3.293, intercept:1.797\n", "Iteration: 40 Loss: 20492.750,ln_TFR:3.293, intercept:1.786\n", "Iteration: 50 Loss: 20489.379,ln_TFR:3.293, intercept:1.787\n", "Iteration: 60 Loss: 20488.174,ln_TFR:3.293, intercept:1.787\n", "Iteration: 70 Loss: 20487.762,ln_TFR:3.293, intercept:1.787\n", "Iteration: 80 Loss: 20487.631,ln_TFR:3.293, intercept:1.787\n", "Iteration: 90 Loss: 20487.594,ln_TFR:3.293, intercept:1.787\n", "Iteration: 100 Loss: 20487.580,ln_TFR:3.293, intercept:1.787\n", "Iteration: 110 Loss: 20487.578,ln_TFR:3.293, intercept:1.787\n", "Iteration: 120 Loss: 20487.580,ln_TFR:3.293, intercept:1.787\n", "Iteration: 130 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 140 Loss: 20487.578,ln_TFR:3.293, intercept:1.787\n", "Iteration: 150 Loss: 20487.578,ln_TFR:3.293, intercept:1.787\n", "Iteration: 160 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 170 Loss: 20487.578,ln_TFR:3.293, intercept:1.787\n", "Iteration: 180 Loss: 20487.578,ln_TFR:3.293, intercept:1.787\n", "Iteration: 190 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 200 Loss: 20487.578,ln_TFR:3.293, intercept:1.787\n", "Iteration: 210 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 220 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 230 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 240 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 250 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 260 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 270 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 280 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 290 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 300 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 310 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 320 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 330 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 340 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 350 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 360 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 370 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 380 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 390 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 400 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 410 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 420 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 430 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 440 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 450 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 460 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 470 Loss: 20487.578,ln_TFR:3.293, intercept:1.787\n", "Iteration: 480 Loss: 20487.576,ln_TFR:3.293, intercept:1.787\n", "Iteration: 490 Loss: 20487.578,ln_TFR:3.293, intercept:1.787\n", "CPU times: user 7min 33s, sys: 20min 30s, total: 28min 3s\n", "Wall time: 32.5 s\n" ] } ], "source": [ "%%time\n", "\n", "num_warmup_iters = 500\n", "num_iters = 500\n", "num_accepted = 0\n", "\n", "## Record samples of the fitted params and the loss\n", "effect_iso3_samples = np.zeros([num_iters, n_res])\n", "effect_ln_TFR_samples = np.zeros([num_iters])\n", "effect_int_samples = np.zeros([num_iters]) \n", "loss_history = np.zeros([num_iters])\n", "\n", "## Push fixed effects into scope so that we can record them \n", "## (by default, only the REs get recorded since they're in the state)\n", "with tf.variable_scope('model', reuse=True):\n", " int_wghts = tf.get_variable(name='intercept', dtype=np.float32) \n", "\n", "## Start our sesh\n", "with tf.Session() as sess:\n", " sess.run(tf.global_variables_initializer())\n", "\n", " # Run warm-up stage.\n", " for t in range(num_warmup_iters):\n", " _, is_accepted_val = sess.run(\n", " [expectation_update, kernel_results.is_accepted])\n", " num_accepted += is_accepted_val\n", " if t % 250 == 0 or t == num_warmup_iters - 1:\n", " print(\"Warm-Up Iteration: {:>3} Acceptance Rate: {:.3f}\".format(\n", " t, num_accepted / (t + 1)))\n", "\n", " num_accepted = 0 # reset acceptance rate counter\n", "\n", " # Run iterations (no convergence criteria)\n", " for t in range(num_iters):\n", " for _ in range(5): # run 5 MCMC iterations before every joint EM update\n", " _ = sess.run(expectation_update)\n", " [\n", " _,\n", " _,\n", " effect_iso3_val,\n", " ln_tfr_val, \n", " int_val,\n", " is_accepted_val,\n", " loss_val,\n", " ] = sess.run([\n", " expectation_update,\n", " minimization_update,\n", " effect_iso3, \n", " effect_ln_TFR,\n", " int_wghts,\n", " kernel_results.is_accepted,\n", " loss,\n", " ])\n", " effect_iso3_samples[t, :] = effect_iso3_val \n", " effect_ln_TFR_samples[t] = ln_tfr_val[0] \n", " effect_int_samples[t] = int_val\n", " num_accepted+= is_accepted_val\n", " loss_history[t] = loss_val\n", " if (t % 10 == 0):\n", " print(\"Iteration: {} Loss: {:.3f},ln_TFR:{:.3f}, intercept:{:.3f}\".\n", " format(t, loss_val, ln_tfr_val[0], int_val))" ] } ], "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.5" } }, "nbformat": 4, "nbformat_minor": 2 }