Skip to content

Instantly share code, notes, and snippets.

@sadatnfs

sadatnfs/tfp_linear_model_testing_Fixed.ipynb

Created September 12, 2018 17:48
Show Gist options
  • Save sadatnfs/15c596272e480676842a4d8d7964a4d0 to your computer and use it in GitHub Desktop.
Save sadatnfs/15c596272e480676842a4d8d7964a4d0 to your computer and use it in GitHub Desktop.
Download ZIP
Fitting model in TFP with fixed effect coefficient
Display the source blob
Display the rendered blob
Raw
tfp_linear_model_testing_Fixed.ipynb
{
"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": [
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>iso3</th>\n",
" <th>year</th>\n",
" <th>ln_gdppc</th>\n",
" <th>ln_TFR</th>\n",
" <th>ln_pop</th>\n",
" <th>intercept</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>0</th>\n",
" <td>0</td>\n",
" <td>1960</td>\n",
" <td>7.626012</td>\n",
" <td>2.008214</td>\n",
" <td>16.012330</td>\n",
" <td>1.0</td>\n",
" </tr>\n",
" <tr>\n",
" <th>1</th>\n",
" <td>0</td>\n",
" <td>1961</td>\n",
" <td>7.613911</td>\n",
" <td>2.008214</td>\n",
" <td>16.031095</td>\n",
" <td>1.0</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2</th>\n",
" <td>0</td>\n",
" <td>1962</td>\n",
" <td>7.610066</td>\n",
" <td>2.008214</td>\n",
" <td>16.050445</td>\n",
" <td>1.0</td>\n",
" </tr>\n",
" <tr>\n",
" <th>3</th>\n",
" <td>0</td>\n",
" <td>1963</td>\n",
" <td>7.607401</td>\n",
" <td>2.008214</td>\n",
" <td>16.070370</td>\n",
" <td>1.0</td>\n",
" </tr>\n",
" <tr>\n",
" <th>4</th>\n",
" <td>0</td>\n",
" <td>1964</td>\n",
" <td>7.603515</td>\n",
" <td>2.008214</td>\n",
" <td>16.090864</td>\n",
" <td>1.0</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>"
],
"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": [
"<table class=\"simpletable\">\n",
"<tr>\n",
" <td>Model:</td> <td>MixedLM</td> <td>Dependent Variable:</td> <td>ln_gdppc</td> \n",
"</tr>\n",
"<tr>\n",
" <td>No. Observations:</td> <td>11700</td> <td>Method:</td> <td>REML</td> \n",
"</tr>\n",
"<tr>\n",
" <td>No. Groups:</td> <td>195</td> <td>Scale:</td> <td>0.1200</td> \n",
"</tr>\n",
"<tr>\n",
" <td>Min. group size:</td> <td>60</td> <td>Likelihood:</td> <td>-4798.2007</td>\n",
"</tr>\n",
"<tr>\n",
" <td>Max. group size:</td> <td>60</td> <td>Converged:</td> <td>Yes</td> \n",
"</tr>\n",
"<tr>\n",
" <td>Mean group size:</td> <td>60.0</td> <td></td> <td></td> \n",
"</tr>\n",
"</table>\n",
"<table class=\"simpletable\">\n",
"<tr>\n",
" <td></td> <th>Coef.</th> <th>Std.Err.</th> <th>z</th> <th>P>|z|</th> <th>[0.025</th> <th>0.975]</th>\n",
"</tr>\n",
"<tr>\n",
" <th>Intercept</th> <td>9.929</td> <td>0.070</td> <td>142.555</td> <td>0.000</td> <td>9.793</td> <td>10.066</td>\n",
"</tr>\n",
"<tr>\n",
" <th>ln_TFR</th> <td>-0.876</td> <td>0.011</td> <td>-82.073</td> <td>0.000</td> <td>-0.897</td> <td>-0.855</td>\n",
"</tr>\n",
"<tr>\n",
" <th>Group Var</th> <td>0.909</td> <td>0.270</td> <td></td> <td></td> <td></td> <td></td> \n",
"</tr>\n",
"</table>"
],
"text/plain": [
"<class 'statsmodels.iolib.summary2.Summary'>\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 fixed TFRs"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"## Keeping TFR fixed\n",
"def linear_mixed_effects_model(features):\n",
" \n",
" # Set up fixed effects and other parameters.\n",
" intercept = tf.get_variable(\"intercept\", []) # alpha in eq\n",
" \n",
" ## Constraint is not needed, but can be done as such\n",
" # beta in eq \n",
" effect_ln_TFR = tf.get_variable(\"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"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"# Unnormalized target density as a function of states of REs.\n",
"def target_log_prob_fn(effect_iso3): \n",
" return log_joint( \n",
" features=features_train,\n",
" effect_iso3=effect_iso3, \n",
" ln_gdppc=labels_train)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Set up the graph, and the backend of the E-M algorithm (fixed TFR)"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"tf.reset_default_graph()\n",
"model_template = tf.make_template(\"model\", linear_mixed_effects_model)\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",
"\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,\n",
" step_size=0.02,\n",
" num_leapfrog_steps=2)\n",
"\n",
"## RE state to update\n",
"current_state = [effect_iso3]\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",
"\n",
"# Set up M-step (gradient descent), using Adam Optimizer\n",
"with tf.control_dependencies([expectation_update]):\n",
" loss = -target_log_prob_fn(effect_iso3) \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": "markdown",
"metadata": {},
"source": [
"### Fit the first run (fixed TFR)"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {
"scrolled": false
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Warm-Up Iteration: 0 Acceptance Rate: 1.000\n",
"Warm-Up Iteration: 249 Acceptance Rate: 0.880\n",
"Iteration: 0 Loss: 114910.953,ln_TFR:-0.557, intercept:-1.116\n",
"Iteration: 20 Loss: 28770.760,ln_TFR:-0.044, intercept:-0.431\n",
"Iteration: 40 Loss: 16941.445,ln_TFR:-0.141, intercept:-0.050\n",
"Iteration: 60 Loss: 11933.211,ln_TFR:-0.419, intercept:0.242\n",
"Iteration: 80 Loss: 9291.289,ln_TFR:-0.648, intercept:0.470\n",
"CPU times: user 7.85 s, sys: 333 ms, total: 8.18 s\n",
"Wall time: 5.87 s\n"
]
}
],
"source": [
"%%time\n",
"\n",
"num_warmup_iters = 250\n",
"num_iters = 100\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",
" ln_tfr_wghts = tf.get_variable(name='effect_ln_TFR', dtype=np.float32)\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",
" ln_tfr_wghts, \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 \n",
" effect_int_samples[t] = int_val\n",
" num_accepted+= is_accepted_val\n",
" loss_history[t] = loss_val\n",
" if t % 20 == 0:\n",
" print(\"Iteration: {} Loss: {:.3f},ln_TFR:{:.3f}, intercept:{:.3f}\".\n",
" format(t, loss_val, ln_tfr_val, 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
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment