-
-
Notifications
You must be signed in to change notification settings - Fork 198
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Sync example_ConsRiskyContribModel notebook with the .py script (#1336)
- Loading branch information
Showing
1 changed file
with
364 additions
and
0 deletions.
There are no files selected for viewing
364 changes: 364 additions & 0 deletions
364
examples/ConsumptionSaving/example_ConsRiskyContribModel.ipynb
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 | ||
} |