Skip to content

Commit

Permalink
Sync example_ConsRiskyContribModel notebook with the .py script (#1336)
Browse files Browse the repository at this point in the history
  • Loading branch information
MridulS authored Aug 5, 2023
1 parent e6a37c5 commit 7c019cd
Showing 1 changed file with 364 additions and 0 deletions.
364 changes: 364 additions & 0 deletions examples/ConsumptionSaving/example_ConsRiskyContribModel.ipynb
Original file line number Diff line number Diff line change
@@ -0,0 +1,364 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "b67cf3c6",
"metadata": {},
"outputs": [],
"source": [
"\"\"\"\n",
"Example implementations of HARK.ConsumptionSaving.ConsPortfolioModel\n",
"\"\"\"\n",
"from HARK.ConsumptionSaving.ConsRiskyContribModel import (\n",
" RiskyContribConsumerType,\n",
" init_risky_contrib,\n",
")\n",
"from time import time\n",
"import numpy as np"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e0529382",
"metadata": {
"title": "Define a plotting function"
},
"outputs": [],
"source": [
"\n",
"\n",
"def plot_slices_3d(\n",
" functions, bot_x, top_x, y_slices, N=300, y_name=None, titles=None, ax_labs=None\n",
"):\n",
"\n",
" import matplotlib.pyplot as plt\n",
"\n",
" if type(functions) == list:\n",
" function_list = functions\n",
" else:\n",
" function_list = [functions]\n",
"\n",
" nfunc = len(function_list)\n",
"\n",
" # Initialize figure and axes\n",
" fig = plt.figure(figsize=plt.figaspect(1.0 / nfunc))\n",
"\n",
" # Create x grid\n",
" x = np.linspace(bot_x, top_x, N, endpoint=True)\n",
"\n",
" for k in range(nfunc):\n",
" ax = fig.add_subplot(1, nfunc, k + 1)\n",
"\n",
" for y in y_slices:\n",
"\n",
" if y_name is None:\n",
" lab = \"\"\n",
" else:\n",
" lab = y_name + \"=\" + str(y)\n",
"\n",
" z = function_list[k](x, np.ones_like(x) * y)\n",
" ax.plot(x, z, label=lab)\n",
"\n",
" if ax_labs is not None:\n",
" ax.set_xlabel(ax_labs[0])\n",
" ax.set_ylabel(ax_labs[1])\n",
"\n",
" # ax.imshow(Z, extent=[bottom[0],top[0],bottom[1],top[1]], origin='lower')\n",
" # ax.colorbar();\n",
" if titles is not None:\n",
" ax.set_title(titles[k])\n",
"\n",
" ax.set_xlim([bot_x, top_x])\n",
"\n",
" if y_name is not None:\n",
" ax.legend()\n",
"\n",
" plt.show()\n",
"\n",
"\n",
"def plot_slices_4d(\n",
" functions,\n",
" bot_x,\n",
" top_x,\n",
" y_slices,\n",
" w_slices,\n",
" N=300,\n",
" slice_names=None,\n",
" titles=None,\n",
" ax_labs=None,\n",
"):\n",
"\n",
" import matplotlib.pyplot as plt\n",
"\n",
" if type(functions) == list:\n",
" function_list = functions\n",
" else:\n",
" function_list = [functions]\n",
"\n",
" nfunc = len(function_list)\n",
" nws = len(w_slices)\n",
"\n",
" # Initialize figure and axes\n",
" fig = plt.figure(figsize=plt.figaspect(1.0 / nfunc))\n",
"\n",
" # Create x grid\n",
" x = np.linspace(bot_x, top_x, N, endpoint=True)\n",
"\n",
" for j in range(nws):\n",
" w = w_slices[j]\n",
"\n",
" for k in range(nfunc):\n",
" ax = fig.add_subplot(nws, nfunc, j * nfunc + k + 1)\n",
"\n",
" for y in y_slices:\n",
"\n",
" if slice_names is None:\n",
" lab = \"\"\n",
" else:\n",
" lab = (\n",
" slice_names[0]\n",
" + \"=\"\n",
" + str(y)\n",
" + \",\"\n",
" + slice_names[1]\n",
" + \"=\"\n",
" + str(w)\n",
" )\n",
"\n",
" z = function_list[k](x, np.ones_like(x) * y, np.ones_like(x) * w)\n",
" ax.plot(x, z, label=lab)\n",
"\n",
" if ax_labs is not None:\n",
" ax.set_xlabel(ax_labs[0])\n",
" ax.set_ylabel(ax_labs[1])\n",
"\n",
" # ax.imshow(Z, extent=[bottom[0],top[0],bottom[1],top[1]], origin='lower')\n",
" # ax.colorbar();\n",
" if titles is not None:\n",
" ax.set_title(titles[k])\n",
"\n",
" ax.set_xlim([bot_x, top_x])\n",
"\n",
" if slice_names is not None:\n",
" ax.legend()\n",
"\n",
" plt.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "fdee9c15",
"metadata": {},
"outputs": [],
"source": [
"# Solve an infinite horizon version\n",
"\n",
"# Get initial parameters\n",
"par_infinite = init_risky_contrib.copy()\n",
"# And make the problem infinite horizon\n",
"par_infinite[\"cycles\"] = 0\n",
"# and sticky\n",
"par_infinite[\"AdjustPrb\"] = 1.0\n",
"# and with a withdrawal tax\n",
"par_infinite[\"tau\"] = 0.1\n",
"\n",
"par_infinite[\"DiscreteShareBool\"] = False\n",
"par_infinite[\"vFuncBool\"] = False\n",
"\n",
"# Create agent and solve it.\n",
"inf_agent = RiskyContribConsumerType(tolerance=1e-3, **par_infinite)\n",
"print(\"Now solving infinite horizon version\")\n",
"t0 = time()\n",
"inf_agent.solve(verbose=True)\n",
"t1 = time()\n",
"print(\"Converged!\")\n",
"print(\"Solving took \" + str(t1 - t0) + \" seconds.\")\n",
"\n",
"# Plot policy functions\n",
"periods = [0]\n",
"n_slices = [0, 2, 6]\n",
"mMax = 20\n",
"\n",
"dfracFunc = [inf_agent.solution[t].stage_sols[\"Reb\"].dfracFunc_Adj for t in periods]\n",
"ShareFunc = [inf_agent.solution[t].stage_sols[\"Sha\"].ShareFunc_Adj for t in periods]\n",
"cFuncFxd = [inf_agent.solution[t].stage_sols[\"Cns\"].cFunc for t in periods]\n",
"\n",
"# Rebalancing\n",
"plot_slices_3d(\n",
" dfracFunc,\n",
" 0,\n",
" mMax,\n",
" y_slices=n_slices,\n",
" y_name=\"n\",\n",
" titles=[\"t = \" + str(t) for t in periods],\n",
" ax_labs=[\"m\", \"d\"],\n",
")\n",
"# Share\n",
"plot_slices_3d(\n",
" ShareFunc,\n",
" 0,\n",
" mMax,\n",
" y_slices=n_slices,\n",
" y_name=\"n\",\n",
" titles=[\"t = \" + str(t) for t in periods],\n",
" ax_labs=[\"m\", \"S\"],\n",
")\n",
"\n",
"# Consumption\n",
"shares = [0.0, 0.9]\n",
"plot_slices_4d(\n",
" cFuncFxd,\n",
" 0,\n",
" mMax,\n",
" y_slices=n_slices,\n",
" w_slices=shares,\n",
" slice_names=[\"n_til\", \"s\"],\n",
" titles=[\"t = \" + str(t) for t in periods],\n",
" ax_labs=[\"m_til\", \"c\"],\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b8fac948",
"metadata": {},
"outputs": [],
"source": [
"# Solve a short, finite horizon version\n",
"par_finite = init_risky_contrib.copy()\n",
"\n",
"# Four period model\n",
"par_finite[\"PermGroFac\"] = [2.0, 1.0, 0.1, 1.0]\n",
"par_finite[\"PermShkStd\"] = [0.1, 0.1, 0.0, 0.0]\n",
"par_finite[\"TranShkStd\"] = [0.2, 0.2, 0.0, 0.0]\n",
"par_finite[\"AdjustPrb\"] = [0.5, 0.5, 1.0, 1.0]\n",
"par_finite[\"tau\"] = [0.1, 0.1, 0.0, 0.0]\n",
"par_finite[\"LivPrb\"] = [1.0, 1.0, 1.0, 1.0]\n",
"par_finite[\"T_cycle\"] = 4\n",
"par_finite[\"T_retire\"] = 0\n",
"par_finite[\"T_age\"] = 4\n",
"\n",
"# Adjust discounting and returns distribution so that they make sense in a\n",
"# 4-period model\n",
"par_finite[\"DiscFac\"] = 0.95 ** 15\n",
"par_finite[\"Rfree\"] = 1.03 ** 15\n",
"par_finite[\"RiskyAvg\"] = 1.08 ** 15 # Average return of the risky asset\n",
"par_finite[\"RiskyStd\"] = 0.20 * np.sqrt(15) # Standard deviation of (log) risky returns\n",
"\n",
"\n",
"# Create and solve\n",
"contrib_agent = RiskyContribConsumerType(**par_finite)\n",
"print(\"Now solving\")\n",
"t0 = time()\n",
"contrib_agent.solve()\n",
"t1 = time()\n",
"print(\"Solving took \" + str(t1 - t0) + \" seconds.\")\n",
"\n",
"# Plot Policy functions\n",
"periods = [0, 2, 3]\n",
"\n",
"dfracFunc = [contrib_agent.solution[t].stage_sols[\"Reb\"].dfracFunc_Adj for t in periods]\n",
"ShareFunc = [contrib_agent.solution[t].stage_sols[\"Sha\"].ShareFunc_Adj for t in periods]\n",
"cFuncFxd = [contrib_agent.solution[t].stage_sols[\"Cns\"].cFunc for t in periods]\n",
"\n",
"# Rebalancing\n",
"plot_slices_3d(\n",
" dfracFunc,\n",
" 0,\n",
" mMax,\n",
" y_slices=n_slices,\n",
" y_name=\"n\",\n",
" titles=[\"t = \" + str(t) for t in periods],\n",
" ax_labs=[\"m\", \"d\"],\n",
")\n",
"# Share\n",
"plot_slices_3d(\n",
" ShareFunc,\n",
" 0,\n",
" mMax,\n",
" y_slices=n_slices,\n",
" y_name=\"n\",\n",
" titles=[\"t = \" + str(t) for t in periods],\n",
" ax_labs=[\"m\", \"S\"],\n",
")\n",
"# Consumption\n",
"plot_slices_4d(\n",
" cFuncFxd,\n",
" 0,\n",
" mMax,\n",
" y_slices=n_slices,\n",
" w_slices=shares,\n",
" slice_names=[\"n_til\", \"s\"],\n",
" titles=[\"t = \" + str(t) for t in periods],\n",
" ax_labs=[\"m_til\", \"c\"],\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d860387e",
"metadata": {
"title": "Simulate the finite horizon consumer"
},
"outputs": [],
"source": [
"contrib_agent.track_vars = [\n",
" \"pLvl\",\n",
" \"t_age\",\n",
" \"Adjust\",\n",
" \"mNrm\",\n",
" \"nNrm\",\n",
" \"mNrmTilde\",\n",
" \"nNrmTilde\",\n",
" \"aNrm\",\n",
" \"cNrm\",\n",
" \"Share\",\n",
" \"dfrac\",\n",
"]\n",
"contrib_agent.T_sim = 4\n",
"contrib_agent.AgentCount = 10\n",
"contrib_agent.initialize_sim()\n",
"contrib_agent.simulate()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b99aa25b",
"metadata": {
"title": "Format simulation results"
},
"outputs": [],
"source": [
"\n",
"import pandas as pd\n",
"\n",
"df = contrib_agent.history\n",
"\n",
"# Add an id to the simulation results\n",
"agent_id = np.arange(contrib_agent.AgentCount)\n",
"df[\"id\"] = np.tile(agent_id, (contrib_agent.T_sim, 1))\n",
"\n",
"# Flatten variables\n",
"df = {k: v.flatten(order=\"F\") for k, v in df.items()}\n",
"\n",
"# Make dataframe\n",
"df = pd.DataFrame(df)"
]
}
],
"metadata": {
"jupytext": {
"cell_metadata_filter": "title,-all",
"main_language": "python",
"notebook_metadata_filter": "-all"
}
},
"nbformat": 4,
"nbformat_minor": 5
}

0 comments on commit 7c019cd

Please sign in to comment.