From dc900f3a35fe841442353aef93da017303b91801 Mon Sep 17 00:00:00 2001
From: dominicWC <135868541+DominicWC@users.noreply.github.com>
Date: Fri, 6 Sep 2024 09:48:56 -0500
Subject: [PATCH 01/26] Fixes a few issues in the latex models

---
 HARK/ConsumptionSaving/ConsBequestModel.py  | 14 +++++---------
 HARK/ConsumptionSaving/ConsIndShockModel.py |  2 +-
 2 files changed, 6 insertions(+), 10 deletions(-)

diff --git a/HARK/ConsumptionSaving/ConsBequestModel.py b/HARK/ConsumptionSaving/ConsBequestModel.py
index 4a265341a..f0bbeb801 100644
--- a/HARK/ConsumptionSaving/ConsBequestModel.py
+++ b/HARK/ConsumptionSaving/ConsBequestModel.py
@@ -273,12 +273,8 @@ def make_warmglow_portfolio_solution_terminal(
 
 class BequestWarmGlowConsumerType(IndShockConsumerType):
     r"""
-    A consumer type with idiosyncratic shocks to permanent and transitory income.
-    Their problem is defined by a sequence of income distributions, survival probabilities
-    (:math:`1-\mathsf{D}`), and permanent income growth rates (:math:`\Gamma`), as well
-    as time invariant values for risk aversion (:math:`\rho`), discount factor (:math:`\beta`),
-    the interest rate (:math:`\mathsf{R}`), the grid of end-of-period assets, and an artificial
-    borrowing constraint (:math:`\underline{a}`).
+    A consumer type with based on IndShockConsumerType, with an additional bequest motive.
+    They gain utility for any wealth they leave when they die, according to a Stone-Geary utility.
 
     .. math::
         \newcommand{\CRRA}{\rho}
@@ -295,7 +291,7 @@ class BequestWarmGlowConsumerType(IndShockConsumerType):
         (\psi_{t+1},\theta_{t+1}) &\sim F_{t+1}, \\
         \mathbb{E}[\psi]=\mathbb{E}[\theta] &= 1, \\
         u(c) &= \frac{c^{1-\CRRA}}{1-\CRRA} \\
-        u_{Beq} (a) = \textbf{BeqFac} \frac{(a+\textbf{BeqShift})^{1-\CRRA_{Beq}}}{1-\CRRA_{Beq}}
+        u_{Beq} (a) &= \textbf{BeqFac} \frac{(a+\textbf{BeqShift})^{1-\CRRA_{Beq}}}{1-\CRRA_{Beq}} \\
         \end{align*}
 
 
@@ -1277,7 +1273,7 @@ def calc_EndOfPrd_v(S, a, z):
 
 class BequestWarmGlowPortfolioType(PortfolioConsumerType):
     r"""
-    A consumer type with based on IndShockConsumerType, with an additional bequest motive.
+    A consumer type with based on PortfolioConsumerType, with an additional bequest motive.
     They gain utility for any wealth they leave when they die, according to a Stone-Geary utility.
 
     .. math::
@@ -1300,7 +1296,7 @@ class BequestWarmGlowPortfolioType(PortfolioConsumerType):
         (\psi_{t+1},\theta_{t+1},\phi_{t+1},p_t) &\sim F_{t+1}, \\
         \mathbb{E}[\psi]=\mathbb{E}[\theta] &= 1. \\
         u(c) &= \frac{c^{1-\CRRA}}{1-\CRRA} \\
-        u_{Beq} (a) = \textbf{BeqFac} \frac{(a+\textbf{BeqShift})^{1-\CRRA_{Beq}}}{1-\CRRA_{Beq}}
+        u_{Beq} (a) &= \textbf{BeqFac} \frac{(a+\textbf{BeqShift})^{1-\CRRA_{Beq}}}{1-\CRRA_{Beq}} \\
         \end{align*}
 
 
diff --git a/HARK/ConsumptionSaving/ConsIndShockModel.py b/HARK/ConsumptionSaving/ConsIndShockModel.py
index f3146c4c1..098c51ed5 100644
--- a/HARK/ConsumptionSaving/ConsIndShockModel.py
+++ b/HARK/ConsumptionSaving/ConsIndShockModel.py
@@ -2798,7 +2798,7 @@ class KinkedRconsumerType(IndShockConsumerType):
         \end{cases}\\
         \Rfree_{boro} &> \Rfree_{save}, \\
         (\psi_{t+1},\theta_{t+1}) &\sim F_{t+1}, \\
-        \mathbb{E}[\psi]=\mathbb{E}[\theta] &= 1.
+        \mathbb{E}[\psi]=\mathbb{E}[\theta] &= 1.\\
         u(c) &= \frac{c^{1-\CRRA}}{1-\CRRA} \\
         \end{align*}
 

From 46e4173c708e6af17a602014bcd44267cb54bb42 Mon Sep 17 00:00:00 2001
From: "Matthew N. White" <mnwhite@gmail.com>
Date: Tue, 17 Sep 2024 14:15:22 -0400
Subject: [PATCH 02/26] WealthPortfolioConsumerType uses constructors

Updated this AgentType subclass to use constructors, have standalone default dictionary. Still has no tests nor documentation.
---
 .../ConsWealthPortfolioModel.py               | 157 ++++++++++++++++--
 1 file changed, 147 insertions(+), 10 deletions(-)

diff --git a/HARK/ConsumptionSaving/ConsWealthPortfolioModel.py b/HARK/ConsumptionSaving/ConsWealthPortfolioModel.py
index a42686c5d..f8937d40f 100644
--- a/HARK/ConsumptionSaving/ConsWealthPortfolioModel.py
+++ b/HARK/ConsumptionSaving/ConsWealthPortfolioModel.py
@@ -3,8 +3,8 @@
 import numpy as np
 from HARK.ConsumptionSaving.ConsPortfolioModel import (
     PortfolioConsumerType,
-    init_portfolio,
     PortfolioSolution,
+    make_portfolio_solution_terminal,
 )
 from HARK.distribution import expected
 from HARK.interpolation import (
@@ -14,11 +14,22 @@
     MargValueFuncCRRA,
     ValueFuncCRRA,
 )
+from HARK.Calibration.Assets.AssetProcesses import (
+    make_lognormal_RiskyDstn,
+    combine_IncShkDstn_and_RiskyDstn,
+    calc_ShareLimit_for_CRRA,
+)
+from HARK.Calibration.Income.IncomeProcesses import (
+    construct_lognormal_income_process_unemployment,
+    get_PermShkDstn_from_IncShkDstn,
+    get_TranShkDstn_from_IncShkDstn,
+)
+from HARK.ConsumptionSaving.ConsRiskyAssetModel import (
+    make_simple_ShareGrid,
+)
 from HARK.rewards import UtilityFuncCRRA
-from HARK.utilities import NullFunc
+from HARK.utilities import NullFunc, make_assets_grid
 
-
-import numpy as np
 from HARK.interpolation import LinearInterp
 
 
@@ -100,6 +111,129 @@ def __call__(self, omega):
         return np.nan_to_num(chi)
 
 
+# Trivial constructor function
+def make_ChiFromOmega_function(CRRA, WealthShare, ChiFromOmega_N, ChiFromOmega_bound):
+    return ChiFromOmegaFunction(
+        CRRA, WealthShare, N=ChiFromOmega_N, z_bound=ChiFromOmega_bound
+    )
+
+
+###############################################################################
+
+# Make a dictionary of constructors for the wealth-in-utility portfolio choice consumer type
+WealthPortfolioConsumerType_constructors_default = {
+    "IncShkDstn": construct_lognormal_income_process_unemployment,
+    "PermShkDstn": get_PermShkDstn_from_IncShkDstn,
+    "TranShkDstn": get_TranShkDstn_from_IncShkDstn,
+    "aXtraGrid": make_assets_grid,
+    "RiskyDstn": make_lognormal_RiskyDstn,
+    "ShockDstn": combine_IncShkDstn_and_RiskyDstn,
+    "ShareLimit": calc_ShareLimit_for_CRRA,
+    "ShareGrid": make_simple_ShareGrid,
+    "ChiFunc": make_ChiFromOmega_function,
+    "solution_terminal": make_portfolio_solution_terminal,
+}
+
+# Default parameters to make IncShkDstn using construct_lognormal_income_process_unemployment
+WealthPortfolioConsumerType_IncShkDstn_default = {
+    "PermShkStd": [0.1],  # Standard deviation of log permanent income shocks
+    "PermShkCount": 7,  # Number of points in discrete approximation to permanent income shocks
+    "TranShkStd": [0.1],  # Standard deviation of log transitory income shocks
+    "TranShkCount": 7,  # Number of points in discrete approximation to transitory income shocks
+    "UnempPrb": 0.05,  # Probability of unemployment while working
+    "IncUnemp": 0.3,  # Unemployment benefits replacement rate while working
+    "T_retire": 0,  # Period of retirement (0 --> no retirement)
+    "UnempPrbRet": 0.005,  # Probability of "unemployment" while retired
+    "IncUnempRet": 0.0,  # "Unemployment" benefits when retired
+}
+
+# Default parameters to make aXtraGrid using make_assets_grid
+WealthPortfolioConsumerType_aXtraGrid_default = {
+    "aXtraMin": 0.001,  # Minimum end-of-period "assets above minimum" value
+    "aXtraMax": 100,  # Maximum end-of-period "assets above minimum" value
+    "aXtraNestFac": 1,  # Exponential nesting factor for aXtraGrid
+    "aXtraCount": 200,  # Number of points in the grid of "assets above minimum"
+    "aXtraExtra": None,  # Additional other values to add in grid (optional)
+}
+
+# Default parameters to make RiskyDstn with make_lognormal_RiskyDstn (and uniform ShareGrid)
+WealthPortfolioConsumerType_RiskyDstn_default = {
+    "RiskyAvg": 1.08,  # Mean return factor of risky asset
+    "RiskyStd": 0.18362634887,  # Stdev of log returns on risky asset
+    "RiskyCount": 5,  # Number of integration nodes to use in approximation of risky returns
+}
+
+WealthPortfolioConsumerType_ShareGrid_default = {
+    "ShareCount": 25  # Number of discrete points in the risky share approximation
+}
+
+# Default parameters to make ChiFunc with make_ChiFromOmega_function
+WealthPortfolioConsumerType_ChiFunc_default = {
+    "ChiFromOmega_N": 501,  # Number of gridpoints in chi-from-omega function
+    "ChiFromOmega_bound": 15,  # Highest gridpoint to use for it
+}
+
+# Make a dictionary to specify a risky asset consumer type
+WealthPortfolioConsumerType_solving_default = {
+    # BASIC HARK PARAMETERS REQUIRED TO SOLVE THE MODEL
+    "cycles": 1,  # Finite, non-cyclic model
+    "T_cycle": 1,  # Number of periods in the cycle for this agent type
+    "constructors": WealthPortfolioConsumerType_constructors_default,  # See dictionary above
+    # PRIMITIVE RAW PARAMETERS REQUIRED TO SOLVE THE MODEL
+    "CRRA": 5.0,  # Coefficient of relative risk aversion
+    "Rfree": 1.03,  # Return factor on risk free asset
+    "DiscFac": 0.90,  # Intertemporal discount factor
+    "LivPrb": [0.98],  # Survival probability after each period
+    "PermGroFac": [1.01],  # Permanent income growth factor
+    "BoroCnstArt": 0.0,  # Artificial borrowing constraint
+    "WealthShare": 0.5,  # Share of wealth in Cobb-Douglas aggregator in utility function
+    "WealthShift": 0.1,  # Shifter for wealth in utility function
+    "DiscreteShareBool": False,  # Whether risky asset share is restricted to discrete values
+    "vFuncBool": False,  # Whether to calculate the value function during solution
+    "CubicBool": False,  # Whether to use cubic spline interpolation when True
+    # (Uses linear spline interpolation for cFunc when False)
+    "AdjustPrb": 1.0,  # Probability that the agent can update their risky portfolio share each period
+    "sim_common_Rrisky": True,  # Whether risky returns have a shared/common value across agents
+}
+WealthPortfolioConsumerType_simulation_default = {
+    # PARAMETERS REQUIRED TO SIMULATE THE MODEL
+    "AgentCount": 10000,  # Number of agents of this type
+    "T_age": None,  # Age after which simulated agents are automatically killed
+    "aNrmInitMean": 0.0,  # Mean of log initial assets
+    "aNrmInitStd": 1.0,  # Standard deviation of log initial assets
+    "pLvlInitMean": 0.0,  # Mean of log initial permanent income
+    "pLvlInitStd": 0.0,  # Standard deviation of log initial permanent income
+    "PermGroFacAgg": 1.0,  # Aggregate permanent income growth factor
+    # (The portion of PermGroFac attributable to aggregate productivity growth)
+    "NewbornTransShk": False,  # Whether Newborns have transitory shock
+    # ADDITIONAL OPTIONAL PARAMETERS
+    "PerfMITShk": False,  # Do Perfect Foresight MIT Shock
+    # (Forces Newborns to follow solution path of the agent they replaced if True)
+    "neutral_measure": False,  # Whether to use permanent income neutral measure (see Harmenberg 2021)
+}
+
+# Assemble the default dictionary
+WealthPortfolioConsumerType_default = {}
+WealthPortfolioConsumerType_default.update(WealthPortfolioConsumerType_solving_default)
+WealthPortfolioConsumerType_default.update(
+    WealthPortfolioConsumerType_simulation_default
+)
+WealthPortfolioConsumerType_default.update(
+    WealthPortfolioConsumerType_aXtraGrid_default
+)
+WealthPortfolioConsumerType_default.update(
+    WealthPortfolioConsumerType_ShareGrid_default
+)
+WealthPortfolioConsumerType_default.update(
+    WealthPortfolioConsumerType_IncShkDstn_default
+)
+WealthPortfolioConsumerType_default.update(
+    WealthPortfolioConsumerType_RiskyDstn_default
+)
+WealthPortfolioConsumerType_default.update(WealthPortfolioConsumerType_ChiFunc_default)
+init_wealth_portfolio = WealthPortfolioConsumerType_default
+
+
 class WealthPortfolioConsumerType(PortfolioConsumerType):
     time_inv_ = deepcopy(PortfolioConsumerType.time_inv_)
     time_inv_ = time_inv_ + ["WealthShare", "WealthShift", "ChiFunc"]
@@ -114,10 +248,18 @@ def __init__(self, **kwds):
 
         self.solve_one_period = solve_one_period_WealthPortfolio
 
+    def update(self):
+        super().update()
+        self.update_ChiFunc()
+
+    def update_ChiFunc(self):
         if self.WealthShare == 0.0:
             self.ChiFunc = None
         else:
-            self.ChiFunc = ChiFromOmegaFunction(self.CRRA, self.WealthShare)
+            self.construct("ChiFunc")
+
+
+###############################################################################
 
 
 def utility(c, a, CRRA, share=0.0, intercept=0.0):
@@ -502,8 +644,3 @@ def solve_one_period_WealthPortfolio(
         vFuncAdj=vFuncNow,
     )
     return solution_now
-
-
-init_wealth_portfolio = init_portfolio.copy()
-init_wealth_portfolio["WealthShare"] = 0.5
-init_wealth_portfolio["WealthShift"] = 0.1

From bc8cd6551019dfab278fd7552b869e97a92359ba Mon Sep 17 00:00:00 2001
From: akshay_shanker <andy.mathur@icloud.com>
Date: Wed, 16 Oct 2024 17:12:36 +1100
Subject: [PATCH 03/26] test

---
 examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
index c9de4c1a8..aaf5a44f0 100644
--- a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
+++ b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
@@ -111,7 +111,7 @@
    "source": [
     "## \n",
     "\n",
-    "Given these steady state choices, we will need to find $K, Z$ to clear the firm first order conditions.\n"
+    "We first find steady state capital (`K') and productivity  (`Z') implied by the firm's first order condition and the above steady state values of labor (`L_ss'), the real interest rate (`r_ss') and output (`Y_ss'). \n"
    ]
   },
   {

From 16e9d0e730b2305a94c007f62ad4ce35002cf9aa Mon Sep 17 00:00:00 2001
From: akshay_shanker <andy.mathur@icloud.com>
Date: Wed, 16 Oct 2024 17:30:29 +1100
Subject: [PATCH 04/26] edits.

---
 .../KS-HARK-presentation.ipynb                | 68 +++++++++++++------
 1 file changed, 49 insertions(+), 19 deletions(-)

diff --git a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
index aaf5a44f0..c22f5793f 100644
--- a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
+++ b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
@@ -60,6 +60,26 @@
     "from HARK.ConsumptionSaving.ConsNewKeynesianModel import NewKeynesianConsumerType"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "28d24bc5",
+   "metadata": {},
+   "source": [
+    "# Calibration and setup\n",
+    "\n",
+    "## The firm"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "df58ea5a",
+   "metadata": {},
+   "source": [
+    "\n",
+    "\n",
+    "Start by collecting the calibration for the economy in a dictionary."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 3,
@@ -75,13 +95,12 @@
    "outputs": [],
    "source": [
     "calibration = {\n",
-    "    \"eis\": 1,\n",
-    "    \"delta\": 0.025,\n",
-    "    \"alpha\": 0.11,\n",
-    "    \"L\": 1.0,\n",
-    "    \"K\": 1.0,\n",
-    "    \"Y\": 1.0,\n",
-    "    \"r\": 0.01,\n",
+    "    \"eis\": 1,       # Elasticity of intertemporal substitution \n",
+    "    \"delta\": 0.025, # Depreciation rate\n",
+    "    \"alpha\": 0.11,  # Capital share of income\n",
+    "    \"L_ss\": 1.0,    # Steady state labor\n",
+    "    \"Y_ss\": 1.0,    # Steady state output\n",
+    "    \"r_ss\": 0.01,   # Steady state real interest rate\n",
     "}"
    ]
   },
@@ -98,20 +117,16 @@
     }
    },
    "outputs": [],
-   "source": [
-    "L_ss = 1.0  # Steady state labor\n",
-    "r_ss = 0.01  # steady state interest rate\n",
-    "Y_ss = 1.0  # steady state output"
-   ]
+   "source": []
   },
   {
    "cell_type": "markdown",
    "id": "1bac99bf",
    "metadata": {},
    "source": [
-    "## \n",
+    "### Find the steady state capital and productivity \n",
     "\n",
-    "We first find steady state capital (`K') and productivity  (`Z') implied by the firm's first order condition and the above steady state values of labor (`L_ss'), the real interest rate (`r_ss') and output (`Y_ss'). \n"
+    "Find steady state capital ('K_ss') and productivity  ('Z') implied by the firm's first order condition and the above steady state values of labor ('L_ss'), the real interest rate ('r_ss') and output (`Y_ss'). "
    ]
   },
   {
@@ -130,12 +145,16 @@
    "source": [
     "from scipy.optimize import root\n",
     "\n",
-    "\n",
     "def your_funcs(X):\n",
+    "    \n",
     "    L = calibration[\"L\"]\n",
     "    alpha = calibration[\"alpha\"]\n",
     "    delta = calibration[\"delta\"]\n",
     "\n",
+    "    L_ss = calibration['L_ss']  \n",
+    "    r_ss = calibration['r_ss']  \n",
+    "    Y_ss = calibration['Y_ss'] \n",
+    "\n",
     "    K, Z = X\n",
     "    # all RHS have to be 0\n",
     "    f = [\n",
@@ -148,7 +167,10 @@
     "\n",
     "sol = root(your_funcs, [1.0, 1.0])  # find roots\n",
     "\n",
-    "K_ss, Z_ss = sol.x"
+    "K_ss, Z_ss = sol.x\n",
+    "\n",
+    "calibration[\"K_ss\"] = K_ss\n",
+    "calibration[\"Z_ss\"] = Z_ss"
    ]
   },
   {
@@ -191,7 +213,7 @@
    "id": "922dce5d",
    "metadata": {},
    "source": [
-    " Let's double check the roots we find produce our chosen steady state values for $ r, Y , L$."
+    " Let's double check the roots we find produce our chosen steady state values for $ r$, $Y$ and L$."
    ]
   },
   {
@@ -249,6 +271,14 @@
     "print(r_ss, w_ss, Y_ss)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "7ca88d16",
+   "metadata": {},
+   "source": [
+    "## HARK agent"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 9,
@@ -309,7 +339,7 @@
    "id": "3be9593e",
    "metadata": {},
    "source": [
-    "# Create HARK agent"
+    "### Create HARK agent"
    ]
   },
   {
@@ -334,7 +364,7 @@
    "id": "fa266888",
    "metadata": {},
    "source": [
-    "# Find Steady state discount factor clear asset market\n",
+    "### Find Steady state discount factor clear asset market\n",
     "\n",
     "We will estimate the discount factor to ensure that asset supply equals the steady state capital we found earlier. \n",
     "\n"

From fa4a3f6a169b19df93d6461b8e2c5f82770446f7 Mon Sep 17 00:00:00 2001
From: akshay_shanker <andy.mathur@icloud.com>
Date: Wed, 16 Oct 2024 18:08:47 +1100
Subject: [PATCH 05/26] edit

---
 .../KS-HARK-presentation.ipynb                | 113 +++++++++++-------
 1 file changed, 72 insertions(+), 41 deletions(-)

diff --git a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
index c22f5793f..a275be3ab 100644
--- a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
+++ b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
@@ -126,7 +126,7 @@
    "source": [
     "### Find the steady state capital and productivity \n",
     "\n",
-    "Find steady state capital ('K_ss') and productivity  ('Z') implied by the firm's first order condition and the above steady state values of labor ('L_ss'), the real interest rate ('r_ss') and output (`Y_ss'). "
+    "Find steady state capital (`K_ss`) and productivity  (`Z`) implied by the firm's first order condition and the above steady state values of labor (`L_ss`), the real interest rate (`r_ss`) and output (`Y_ss`). "
    ]
   },
   {
@@ -147,19 +147,17 @@
     "\n",
     "def your_funcs(X):\n",
     "    \n",
-    "    L = calibration[\"L\"]\n",
+    "    L_ss = calibration[\"L_ss\"]\n",
     "    alpha = calibration[\"alpha\"]\n",
     "    delta = calibration[\"delta\"]\n",
-    "\n",
-    "    L_ss = calibration['L_ss']  \n",
     "    r_ss = calibration['r_ss']  \n",
     "    Y_ss = calibration['Y_ss'] \n",
     "\n",
     "    K, Z = X\n",
     "    # all RHS have to be 0\n",
     "    f = [\n",
-    "        alpha * Z * (K / L) ** (alpha - 1) - delta - r_ss,  # r = MPK\n",
-    "        Z * K**alpha * L ** (1 - alpha) - Y_ss,  # Y = Z*F(K,L)\n",
+    "        alpha * Z * (K / L_ss) ** (alpha - 1) - delta - r_ss,  # r = MPK\n",
+    "        Z * K**alpha * L_ss ** (1 - alpha) - Y_ss,  # Y = Z*F(K,L)\n",
     "    ]\n",
     "\n",
     "    return f\n",
@@ -213,7 +211,7 @@
    "id": "922dce5d",
    "metadata": {},
    "source": [
-    " Let's double check the roots we find produce our chosen steady state values for $ r$, $Y$ and L$."
+    " Let's double check the roots we find produce our chosen steady state values. "
    ]
   },
   {
@@ -243,7 +241,9 @@
     "    return r, w, Y\n",
     "\n",
     "\n",
-    "r_ss, w_ss, Y_ss = firm(sol.x[0], sol.x[1])"
+    "r_ss, w_ss, Y_ss = firm(sol.x[0], sol.x[1])\n",
+    "\n",
+    "calibration[\"w_ss\"] = w_ss"
    ]
   },
   {
@@ -268,7 +268,7 @@
     }
    ],
    "source": [
-    "print(r_ss, w_ss, Y_ss)"
+    "print(f\"Steady state capital: {K_ss}\", f\"Steady state productivity: {Z_ss}\", f\"Steady state output: {Y_ss}\", f\"Steady state real interest rate: {r_ss}\", f\"Steady state wage: {w_ss}\")"
    ]
   },
   {
@@ -276,7 +276,11 @@
    "id": "7ca88d16",
    "metadata": {},
    "source": [
-    "## HARK agent"
+    "## HARK agent\n",
+    "\n",
+    "We now create an instance of an HARK agent. First specify the agent's parameters in a dictionary. \n",
+    "\n",
+    "For now we start the dictionary with a discount factor of 0.98, but will update the below so the steady state is a general equilibrium."
    ]
   },
   {
@@ -300,38 +304,38 @@
     "    \"DiscFac\": 0.98,  # Intertemporal discount factor\n",
     "    \"LivPrb\": [0.99375],  # Survival probability\n",
     "    \"PermGroFac\": [1.00],  # Permanent income growth factor\n",
+    "\n",
     "    # Parameters that specify the income distribution over the lifecycle\n",
-    "    # Standard deviation of log permanent shocks to income\n",
-    "    \"PermShkStd\": [0.06],\n",
-    "    \"PermShkCount\": 5,  # Number of points in discrete approximation to permanent income shocks\n",
-    "    # Standard deviation of log transitory shocks to income\n",
-    "    \"TranShkStd\": [0.2],\n",
-    "    \"TranShkCount\": 5,\n",
-    "    # HANK params\n",
-    "    \"tax_rate\": [\n",
-    "        0,\n",
-    "    ],  # set to 0.0 because we are going to assume that labor here is actually after tax income\n",
-    "    \"labor\": [L_ss],\n",
-    "    \"wage\": [w_ss],\n",
-    "    # Number of points in discrete approximation to transitory income shocks\n",
+    "    \"PermShkStd\": [0.06],  # Std dev of log permanent shocks to income\n",
+    "    \"PermShkCount\": 5,  # Number of points in discrete approx to perm income shocks\n",
+    "    \"TranShkStd\": [0.2],  # Std dev of log transitory shocks to income\n",
+    "    \"TranShkCount\": 5,  # Number of points in discrete approx to transitory shocks\n",
+    "\n",
+    "    # HANK parameters\n",
+    "    \"tax_rate\": [0],  # Assume labor is after-tax income\n",
+    "    \"labor\": calibration[\"L_ss\"],\n",
+    "    \"wage\": calibration[\"w_ss\"],\n",
+    "\n",
+    "    # Unemployment parameters\n",
     "    \"UnempPrb\": 0.0,  # Probability of unemployment while working\n",
     "    \"IncUnemp\": 0.0,  # Unemployment benefits replacement rate\n",
-    "    \"UnempPrbRet\": 0.0000,  # Probability of \"unemployment\" while retired\n",
+    "    \"UnempPrbRet\": 0.0,  # Probability of unemployment while retired\n",
     "    \"IncUnempRet\": 0.0,  # \"Unemployment\" benefits when retired\n",
     "    \"T_retire\": 0.0,  # Period of retirement (0 --> no retirement)\n",
+    "\n",
     "    # Parameters for constructing the \"assets above minimum\" grid\n",
     "    \"aXtraMin\": 0.0001,  # Minimum end-of-period \"assets above minimum\" value\n",
     "    \"aXtraMax\": 2000,  # Maximum end-of-period \"assets above minimum\" value\n",
     "    \"aXtraCount\": 200,  # Number of points in the base grid of \"assets above minimum\"\n",
-    "    # Exponential nesting factor when constructing \"assets above minimum\" grid\n",
-    "    \"aXtraNestFac\": 3,\n",
+    "    \"aXtraNestFac\": 3,  # Exponential nesting factor for \"assets above minimum\" grid\n",
     "    \"aXtraExtra\": None,  # Additional values to add to aXtraGrid\n",
-    "    # Transition Matrix simulation parameters\n",
-    "    \"mCount\": 200,\n",
-    "    \"mMax\": 2000,\n",
-    "    \"mMin\": 0.0001,\n",
-    "    \"mFac\": 3,\n",
-    "}"
+    "\n",
+    "    # Transition matrix simulation parameters\n",
+    "    \"mCount\": 200,  # Number of discrete points for post-shock market resources\n",
+    "    \"mMax\": 2000,  # Maximum post-shock market resources\n",
+    "    \"mMin\": 0.0001,  # Minimum post-shock market resources\n",
+    "    \"mFac\": 3,  # Exponential nesting factor for market resources grid\n",
+    "}\n"
    ]
   },
   {
@@ -339,7 +343,7 @@
    "id": "3be9593e",
    "metadata": {},
    "source": [
-    "### Create HARK agent"
+    "The `NewKeynesianConsumerType` class contains within itself the Python code that constructs the solution for equilibrium agents using the SSJ toolkit -- including computing the steady state policy functions and Jacobian matrices in response to a policy shock. "
    ]
   },
   {
@@ -356,7 +360,7 @@
    },
    "outputs": [],
    "source": [
-    "Agent = NewKeynesianConsumerType(**HANK_Dict)"
+    "tempAgent = NewKeynesianConsumerType(**HANK_Dict)"
    ]
   },
   {
@@ -364,10 +368,26 @@
    "id": "fa266888",
    "metadata": {},
    "source": [
-    "### Find Steady state discount factor clear asset market\n",
-    "\n",
-    "We will estimate the discount factor to ensure that asset supply equals the steady state capital we found earlier. \n",
-    "\n"
+    "Given a discount factor, the agent will supply a steady state capital stock `A_ss`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ebe1792e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "A_ss = tempAgent.compute_steady_state()[0]\n",
+    "print(\"Steady state agent asset supply is: {A_ss}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7636f41f",
+   "metadata": {},
+   "source": [
+    "We can use root finding to find the discount factor that ensures the agent supplies the capital stock `K_ss` which satisfies the firm equilibrium condition"
    ]
   },
   {
@@ -395,7 +415,10 @@
     "def A_ss_func(beta):\n",
     "    HANK_Dict[\"DiscFac\"] = beta\n",
     "\n",
+    "    # For a given value of beta, we need to re-create the Agent\n",
     "    Agent_func = NewKeynesianConsumerType(**HANK_Dict, verbose=False)\n",
+    "\n",
+    "    # And then solve for the steady state supply \n",
     "    A_ss = Agent_func.compute_steady_state()[0]\n",
     "\n",
     "    return A_ss\n",
@@ -413,6 +436,14 @@
     "print(\"Time taken to solve for steady state\", time.time() - start)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "01be6a8a",
+   "metadata": {},
+   "source": [
+    "Now we create the steady state agent using the general equilibrium discount factor."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 12,
@@ -456,9 +487,9 @@
     }
    ],
    "source": [
-    "# to make sure goods and asset markets clear\n",
-    "print(\"goods_clearing\", Y_ss - C_ss - calibration[\"delta\"] * K_ss)\n",
-    "print(\"asset_clearing\", A_ss - K_ss)"
+    "# To make sure goods and asset markets clear\n",
+    "print(\"goods_clearing\", calibration['Y_ss'] - C_ss - calibration[\"delta\"] * calibration['K_ss'])\n",
+    "print(\"asset_clearing\", A_ss - calibration['K_ss'])"
    ]
   },
   {

From f06bf950849b951c72b4ab49739454fb6f1be38f Mon Sep 17 00:00:00 2001
From: akshay_shanker <andy.mathur@icloud.com>
Date: Wed, 16 Oct 2024 19:10:17 +1100
Subject: [PATCH 06/26] explanations.

---
 .../KS-HARK-presentation.ipynb                | 163 ++++++++++--------
 1 file changed, 89 insertions(+), 74 deletions(-)

diff --git a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
index a275be3ab..96c407d89 100644
--- a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
+++ b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
@@ -104,21 +104,6 @@
     "}"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "id": "2a97701f",
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-07-11T15:18:16.176291Z",
-     "iopub.status.busy": "2024-07-11T15:18:16.176051Z",
-     "iopub.status.idle": "2024-07-11T15:18:16.178499Z",
-     "shell.execute_reply": "2024-07-11T15:18:16.178018Z"
-    }
-   },
-   "outputs": [],
-   "source": []
-  },
   {
    "cell_type": "markdown",
    "id": "1bac99bf",
@@ -231,13 +216,13 @@
     "def firm(\n",
     "    K,\n",
     "    Z,\n",
-    "    L=calibration[\"L\"],\n",
+    "    L_ss =calibration[\"L_ss\"],\n",
     "    alpha=calibration[\"alpha\"],\n",
     "    delta=calibration[\"delta\"],\n",
     "):\n",
-    "    r = alpha * Z * (K / L) ** (alpha - 1) - delta\n",
-    "    w = (1 - alpha) * Z * (K / L) ** alpha\n",
-    "    Y = Z * K**alpha * L ** (1 - alpha)\n",
+    "    r = alpha * Z * (K / L_ss) ** (alpha - 1) - delta\n",
+    "    w = (1 - alpha) * Z * (K / L_ss) ** alpha\n",
+    "    Y = Z * K**alpha * L_ss ** (1 - alpha)\n",
     "    return r, w, Y\n",
     "\n",
     "\n",
@@ -268,7 +253,13 @@
     }
    ],
    "source": [
-    "print(f\"Steady state capital: {K_ss}\", f\"Steady state productivity: {Z_ss}\", f\"Steady state output: {Y_ss}\", f\"Steady state real interest rate: {r_ss}\", f\"Steady state wage: {w_ss}\")"
+    "print(\"--------------------------------------+--------\")\n",
+    "print(f\"Steady state capital                  | {K_ss}\")\n",
+    "print(f\"Steady state productivity             | {Z_ss}\")\n",
+    "print(f\"Steady state output                   | {Y_ss}\")\n",
+    "print(f\"Steady state real interest rate       | {r_ss}\")\n",
+    "print(f\"Steady state wage                     | {w_ss}\")\n",
+    "print(\"--------------------------------------+--------\")"
    ]
   },
   {
@@ -278,9 +269,9 @@
    "source": [
     "## HARK agent\n",
     "\n",
-    "We now create an instance of an HARK agent. First specify the agent's parameters in a dictionary. \n",
+    "We now create an instance of an HARK HANK agent.\n",
     "\n",
-    "For now we start the dictionary with a discount factor of 0.98, but will update the below so the steady state is a general equilibrium."
+    "First specify the agent's parameters in a dictionary. For now we start the dictionary with a discount factor of 0.98, but will update the below so the steady state is a general equilibrium."
    ]
   },
   {
@@ -297,45 +288,52 @@
    },
    "outputs": [],
    "source": [
+    "L_ss = calibration[\"L_ss\"]\n",
+    "w_ss = calibration[\"w_ss\"]\n",
+    "r_ss = calibration[\"r_ss\"]\n",
+    "\n",
     "HANK_Dict = {\n",
-    "    # Parameters shared with the perfect foresight model\n",
+    "    \n",
+    "    # Individual agent 'preferences' (shared with perfect foresight model)\n",
     "    \"CRRA\": calibration[\"eis\"],  # Coefficient of relative risk aversion\n",
-    "    \"Rfree\": (1 + r_ss),  # Interest factor on assets\n",
     "    \"DiscFac\": 0.98,  # Intertemporal discount factor\n",
     "    \"LivPrb\": [0.99375],  # Survival probability\n",
+    "    \n",
+    "    # Individual lifcycle income process parameters \n",
     "    \"PermGroFac\": [1.00],  # Permanent income growth factor\n",
+    "    \"PermShkStd\": [0.06], # Standard deviation of log permanent shocks to income\n",
+    "    \"PermShkCount\": 5, # Number of points in discrete approximation to permanent income shocks\n",
+    "    \"TranShkStd\": [0.2],# Standard deviation of log transitory shocks to income\n",
+    "    \"TranShkCount\": 5,# Number of points in discrete approximation to transitory income shocks \n",
     "\n",
-    "    # Parameters that specify the income distribution over the lifecycle\n",
-    "    \"PermShkStd\": [0.06],  # Std dev of log permanent shocks to income\n",
-    "    \"PermShkCount\": 5,  # Number of points in discrete approx to perm income shocks\n",
-    "    \"TranShkStd\": [0.2],  # Std dev of log transitory shocks to income\n",
-    "    \"TranShkCount\": 5,  # Number of points in discrete approx to transitory shocks\n",
-    "\n",
-    "    # HANK parameters\n",
-    "    \"tax_rate\": [0],  # Assume labor is after-tax income\n",
-    "    \"labor\": calibration[\"L_ss\"],\n",
-    "    \"wage\": calibration[\"w_ss\"],\n",
-    "\n",
-    "    # Unemployment parameters\n",
+    "    # Parameters related to unemployment and retirement \n",
     "    \"UnempPrb\": 0.0,  # Probability of unemployment while working\n",
     "    \"IncUnemp\": 0.0,  # Unemployment benefits replacement rate\n",
-    "    \"UnempPrbRet\": 0.0,  # Probability of unemployment while retired\n",
+    "    \"UnempPrbRet\": 0.0000,  # Probability of \"unemployment\" while retired\n",
     "    \"IncUnempRet\": 0.0,  # \"Unemployment\" benefits when retired\n",
     "    \"T_retire\": 0.0,  # Period of retirement (0 --> no retirement)\n",
-    "\n",
-    "    # Parameters for constructing the \"assets above minimum\" grid\n",
+    "    \n",
+    "    # General equilibrium parameters \n",
+    "    \"Rfree\": (1 + r_ss),  # Interest factor for assets faced by agents \n",
+    "    \"wage\": [w_ss],   # Wage rate faced by agents \n",
+    "    \"tax_rate\": [0], # set to 0.0 because we are going to assume that labor here is actually after tax income\n",
+    "    \"labor\": [L_ss], # Aggregate (mean) labor supply\n",
+    "    \n",
+    "    # Parameters for constructing \"assets above minimum\" grid\n",
     "    \"aXtraMin\": 0.0001,  # Minimum end-of-period \"assets above minimum\" value\n",
     "    \"aXtraMax\": 2000,  # Maximum end-of-period \"assets above minimum\" value\n",
     "    \"aXtraCount\": 200,  # Number of points in the base grid of \"assets above minimum\"\n",
-    "    \"aXtraNestFac\": 3,  # Exponential nesting factor for \"assets above minimum\" grid\n",
+    "    \n",
+    "    # Exponential nesting factor when constructing \"assets above minimum\" grid\n",
+    "    \"aXtraNestFac\": 3,\n",
     "    \"aXtraExtra\": None,  # Additional values to add to aXtraGrid\n",
-    "\n",
+    "    \n",
     "    # Transition matrix simulation parameters\n",
-    "    \"mCount\": 200,  # Number of discrete points for post-shock market resources\n",
-    "    \"mMax\": 2000,  # Maximum post-shock market resources\n",
-    "    \"mMin\": 0.0001,  # Minimum post-shock market resources\n",
-    "    \"mFac\": 3,  # Exponential nesting factor for market resources grid\n",
-    "}\n"
+    "    \"mCount\": 200,\n",
+    "    \"mMax\": 2000,\n",
+    "    \"mMin\": 0.0001,\n",
+    "    \"mFac\": 3,\n",
+    "}"
    ]
   },
   {
@@ -343,7 +341,7 @@
    "id": "3be9593e",
    "metadata": {},
    "source": [
-    "The `NewKeynesianConsumerType` class contains within itself the Python code that constructs the solution for equilibrium agents using the SSJ toolkit -- including computing the steady state policy functions and Jacobian matrices in response to a policy shock. "
+    "The `NewKeynesianConsumerType` class contains within itself the Python code that constructs the solution for equilibrium agents using the SSJ toolkit -- including computing the steady state policy functions and Jacobian matrices in response to a policy shock. Let create a temporary instance of the agent."
    ]
   },
   {
@@ -379,15 +377,7 @@
    "outputs": [],
    "source": [
     "A_ss = tempAgent.compute_steady_state()[0]\n",
-    "print(\"Steady state agent asset supply is: {A_ss}\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "7636f41f",
-   "metadata": {},
-   "source": [
-    "We can use root finding to find the discount factor that ensures the agent supplies the capital stock `K_ss` which satisfies the firm equilibrium condition"
+    "print(f\"Steady state agent asset supply for beta = 0.98 is: {A_ss}\")"
    ]
   },
   {
@@ -423,7 +413,6 @@
     "\n",
     "    return A_ss\n",
     "\n",
-    "\n",
     "def ss_dif(beta):\n",
     "    return A_ss_func(beta) - Asset_target\n",
     "\n",
@@ -460,8 +449,8 @@
    "source": [
     "# Create a new agent\n",
     "HANK_Dict[\"DiscFac\"] = DiscFac\n",
-    "Agent_GE = NewKeynesianConsumerType(**HANK_Dict, verbose=False)\n",
-    "A_ss, C_ss = Agent_GE.compute_steady_state()"
+    "steadyHANK = NewKeynesianConsumerType(**HANK_Dict, verbose=False)\n",
+    "A_ss, C_ss = steadyHANK.compute_steady_state()"
    ]
   },
   {
@@ -488,16 +477,32 @@
    ],
    "source": [
     "# To make sure goods and asset markets clear\n",
-    "print(\"goods_clearing\", calibration['Y_ss'] - C_ss - calibration[\"delta\"] * calibration['K_ss'])\n",
-    "print(\"asset_clearing\", A_ss - calibration['K_ss'])"
+    "print(\"Final goods clearing:\", calibration['Y_ss'] - C_ss - calibration[\"delta\"] * calibration['K_ss'])\n",
+    "print(\"Asset clearing:\", A_ss - calibration['K_ss'])"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "86e05d73",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
   {
    "cell_type": "markdown",
    "id": "a8167383",
    "metadata": {},
    "source": [
-    "# Computing Heterogenous Agent Jacobians"
+    "## Computing Jacobians using SSJ\n",
+    "\n",
+    "With the steady state agent in hand, we can now compute the Jacobians of the steady state HANK agent's policy functions with respect to the aggregate state variables!\n",
+    "\n",
+    "The `calc_jacobian` method of our HANK agent computes the Jacobians of the steady state aggregate consumption (`CJACW`) and assets (`AJACW`) with respect to a pertubation of the specified variable.\n",
+    "\n",
+    "Recall `CJACW[s,t]` is the time $t$ response to a shock at time $s$. \n",
+    "\n",
+    "Here is an example where we shock the wage rate. "
    ]
   },
   {
@@ -524,9 +529,9 @@
    "source": [
     "start = time.time()\n",
     "\n",
-    "CJACW, AJACW = Agent_GE.calc_jacobian(\"wage\", 300)  # Wage jacobians\n",
+    "CJACW, AJACW = steadyHANK.calc_jacobian(\"wage\", 300)  # Wage jacobians\n",
     "\n",
-    "print(\"Time taken to compute jacobians\", time.time() - start)"
+    "print(f\"Time taken to compute wage Jacobians: {time.time() - start:.3f} seconds\")"
    ]
   },
   {
@@ -554,15 +559,16 @@
     }
    ],
    "source": [
-    "plt.plot(CJACW.T[0])\n",
-    "plt.plot(CJACW.T[20])\n",
-    "plt.plot(CJACW.T[50])\n",
-    "plt.plot(CJACW.T[100])\n",
+    "plt.plot(CJACW.T[0], label = \"s =0\")\n",
+    "plt.plot(CJACW.T[20], label = 's=20')\n",
+    "plt.plot(CJACW.T[50], label = 's=50')\n",
+    "plt.plot(CJACW.T[100], label = 's=100')\n",
     "plt.xlim(-2, 300)\n",
     "plt.plot(np.arange(300), np.zeros(300), color=\"k\")\n",
     "plt.title(\"Consumption Jacobian Wage\")\n",
     "plt.xlabel(\"quarters\")\n",
     "plt.ylabel(\"C response\")\n",
+    "plt.legend()\n",
     "plt.show()"
    ]
   },
@@ -590,9 +596,9 @@
    "source": [
     "start = time.time()\n",
     "\n",
-    "CJACR, AJACR = Agent_GE.calc_jacobian(\"Rfree\", 300)  # Rfree jacobians\n",
+    "CJACR, AJACR = steadyHANK.calc_jacobian(\"Rfree\", 300)  # Rfree jacobians\n",
     "\n",
-    "print(\"Time taken to compute jacobians\", time.time() - start)"
+    "print(f\"Time taken to compute return factor Jacobians: {time.time() - start:.3f} seconds\")"
    ]
   },
   {
@@ -620,9 +626,10 @@
     }
    ],
    "source": [
-    "plt.plot(CJACR.T[0])\n",
-    "plt.plot(CJACR.T[30])\n",
-    "plt.plot(CJACR.T[50])\n",
+    "plt.plot(CJACR.T[0], label = \"s =0\")\n",
+    "plt.plot(CJACR.T[20], label = 's=20')\n",
+    "plt.plot(CJACR.T[50], label = 's=50')\n",
+    "plt.plot(CJACR.T[100], label = 's=100')\n",
     "plt.plot(np.arange(300), np.zeros(300), color=\"k\")\n",
     "plt.title(\"Consumption Jacobian interest rate\")\n",
     "plt.xlabel(\"quarters\")\n",
@@ -630,6 +637,14 @@
     "plt.show()"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "c5d0b70e",
+   "metadata": {},
+   "source": [
+    "Let's store the Jacobians in a dictionary for later use."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 18,
@@ -682,7 +697,7 @@
    "id": "d3669fab",
    "metadata": {},
    "source": [
-    "## Other Blocks of the Model"
+    "## Impulse resopnse functions and simulations "
    ]
   },
   {

From 1df99feaab13ed4128dbd92908d5f7acaf901361 Mon Sep 17 00:00:00 2001
From: akshay_shanker <andy.mathur@icloud.com>
Date: Wed, 16 Oct 2024 19:30:53 +1100
Subject: [PATCH 07/26] misc.

---
 .../KS-HARK-presentation.ipynb                | 38 ++++++++-----------
 1 file changed, 15 insertions(+), 23 deletions(-)

diff --git a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
index 96c407d89..810053194 100644
--- a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
+++ b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
@@ -253,13 +253,13 @@
     }
    ],
    "source": [
-    "print(\"--------------------------------------+--------\")\n",
-    "print(f\"Steady state capital                  | {K_ss}\")\n",
-    "print(f\"Steady state productivity             | {Z_ss}\")\n",
-    "print(f\"Steady state output                   | {Y_ss}\")\n",
-    "print(f\"Steady state real interest rate       | {r_ss}\")\n",
-    "print(f\"Steady state wage                     | {w_ss}\")\n",
-    "print(\"--------------------------------------+--------\")"
+    "print(\"--------------------------------------+------------\")\n",
+    "print(f\"Steady state capital                  | {K_ss:.4f}\")\n",
+    "print(f\"Steady state productivity             | {Z_ss:.4f}\")\n",
+    "print(f\"Steady state output                   | {Y_ss:.4f}\")\n",
+    "print(f\"Steady state real interest rate       | {r_ss:.4f}\")\n",
+    "print(f\"Steady state wage                     | {w_ss:.4f}\")\n",
+    "print(\"--------------------------------------+------------\")"
    ]
   },
   {
@@ -271,7 +271,7 @@
     "\n",
     "We now create an instance of an HARK HANK agent.\n",
     "\n",
-    "First specify the agent's parameters in a dictionary. For now we start the dictionary with a discount factor of 0.98, but will update the below so the steady state is a general equilibrium."
+    "First specify the agent's parameters in a dictionary. For now we use a discount factor of 0.98, but will update the below so the steady state is a general equilibrium."
    ]
   },
   {
@@ -377,7 +377,7 @@
    "outputs": [],
    "source": [
     "A_ss = tempAgent.compute_steady_state()[0]\n",
-    "print(f\"Steady state agent asset supply for beta = 0.98 is: {A_ss}\")"
+    "print(f\"Steady state agent asset supply for beta = 0.98 is: {A_ss:.3f}\")"
    ]
   },
   {
@@ -422,7 +422,7 @@
     "\n",
     "DiscFac = optimize.brentq(ss_dif, 0.8, 0.9999)\n",
     "\n",
-    "print(\"Time taken to solve for steady state\", time.time() - start)"
+    "print(f\"Time taken to solve for steady state {time.time() - start:.3f} secs.\")"
    ]
   },
   {
@@ -697,7 +697,7 @@
    "id": "d3669fab",
    "metadata": {},
    "source": [
-    "## Impulse resopnse functions and simulations "
+    "## Impulse response functions and simulations "
    ]
   },
   {
@@ -751,7 +751,7 @@
    "id": "f3d569b5",
    "metadata": {},
    "source": [
-    "# Solving for Impulse Responses"
+    "### Solving for Impulse Responses"
    ]
   },
   {
@@ -768,8 +768,8 @@
    },
    "outputs": [],
    "source": [
-    "T = 300  # <-- the length of the IRF\n",
-    "rho_Z = 0.8  # persistence of IRF shock\n",
+    "T = 300  # Length of the IRF\n",
+    "rho_Z = 0.8  # Persistence of IRF shock\n",
     "dZ = 0.001 * Z_ss * rho_Z ** np.arange(T)\n",
     "shocks = {\"Z\": dZ}\n",
     "\n",
@@ -853,7 +853,7 @@
    "id": "a0130fc3",
    "metadata": {},
    "source": [
-    "# Simulating the model"
+    "### Simulating the model"
    ]
   },
   {
@@ -920,14 +920,6 @@
     "data_simul = simulate(list(impulses.values()), outputs, T_sim)\n",
     "plot_timeseries(data_simul, (1, 3), figsize=(12, 4))"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "27732edd",
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {

From 38c0a031f7ab0b97a99962a5a7b6b1352cf30ca3 Mon Sep 17 00:00:00 2001
From: akshay_shanker <andy.mathur@icloud.com>
Date: Wed, 16 Oct 2024 19:35:49 +1100
Subject: [PATCH 08/26] Update KS-HARK-presentation.ipynb

---
 examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
index 810053194..b6c14c2c2 100644
--- a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
+++ b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
@@ -67,7 +67,7 @@
    "source": [
     "# Calibration and setup\n",
     "\n",
-    "## The firm"
+    "## Firm setup"
    ]
   },
   {

From aed0be4482cc938ef5971ee126216779a6700efb Mon Sep 17 00:00:00 2001
From: akshay_shanker <andy.mathur@icloud.com>
Date: Wed, 16 Oct 2024 19:36:12 +1100
Subject: [PATCH 09/26] Update KS-HARK-presentation.ipynb

---
 examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
index b6c14c2c2..21db78e85 100644
--- a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
+++ b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
@@ -77,7 +77,7 @@
    "source": [
     "\n",
     "\n",
-    "Start by collecting the calibration for the economy in a dictionary."
+    "Collect the calibration for the economy in a dictionary."
    ]
   },
   {

From d6d0b84495fd583e9f90076b344a0d027c8e71be Mon Sep 17 00:00:00 2001
From: akshay_shanker <andy.mathur@icloud.com>
Date: Wed, 16 Oct 2024 19:36:39 +1100
Subject: [PATCH 10/26] Update KS-HARK-presentation.ipynb

---
 examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
index 21db78e85..93e605597 100644
--- a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
+++ b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
@@ -77,7 +77,7 @@
    "source": [
     "\n",
     "\n",
-    "Collect the calibration for the economy in a dictionary."
+    "Collect calibration for the production economy in a dictionary."
    ]
   },
   {

From 076f870930c85f22da02d347679e75932cec62e8 Mon Sep 17 00:00:00 2001
From: akshay_shanker <andy.mathur@icloud.com>
Date: Wed, 16 Oct 2024 19:53:08 +1100
Subject: [PATCH 11/26] misc.

---
 examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
index 93e605597..f7129bc0f 100644
--- a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
+++ b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
@@ -269,9 +269,11 @@
    "source": [
     "## HARK agent\n",
     "\n",
-    "We now create an instance of an HARK HANK agent.\n",
+    "We now create an instance of an agent.\n",
     "\n",
-    "First specify the agent's parameters in a dictionary. For now we use a discount factor of 0.98, but will update the below so the steady state is a general equilibrium."
+    "HARK represents solving agent problems where aggregate variables can effect an agent's income process as instances of `HANKAgentType`, a subclass of `AgentType`. \n",
+    "\n",
+    "First specify the agent's parameters in a dictionary. For now we use a discount factor of 0.98, but we will update `DiscFac` below so the steady state is a general equilibrium."
    ]
   },
   {

From a36ba6abd97b25a326b0b98a1532e2ac65050be9 Mon Sep 17 00:00:00 2001
From: akshay_shanker <andy.mathur@icloud.com>
Date: Wed, 16 Oct 2024 19:54:05 +1100
Subject: [PATCH 12/26] misc.

---
 examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
index f7129bc0f..e8a1cdc80 100644
--- a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
+++ b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
@@ -273,6 +273,8 @@
     "\n",
     "HARK represents solving agent problems where aggregate variables can effect an agent's income process as instances of `HANKAgentType`, a subclass of `AgentType`. \n",
     "\n",
+    "**_NOTE:_**  An `HANKAgentType` is still a microeconomic agent solving an income fluctuation problem. The only difference from `IndShockConsumerType` is that additional aggregate labour income variables are passed to the agent's income process. This allows the agent's problem to be defined within a general equilirium framework such as (but not restricted to) HANK. \n",
+    "\n",
     "First specify the agent's parameters in a dictionary. For now we use a discount factor of 0.98, but we will update `DiscFac` below so the steady state is a general equilibrium."
    ]
   },
@@ -343,7 +345,9 @@
    "id": "3be9593e",
    "metadata": {},
    "source": [
-    "The `NewKeynesianConsumerType` class contains within itself the Python code that constructs the solution for equilibrium agents using the SSJ toolkit -- including computing the steady state policy functions and Jacobian matrices in response to a policy shock. Let create a temporary instance of the agent."
+    "The `NewKeynesianConsumerType` class contains within itself the Python code that constructs the solution for equilibrium agents using the SSJ toolkit, including computing the steady state policy functions and Jacobian matrices in response to a policy shock. \n",
+    "\n",
+    "Let's create a temporary instance of an HANK agent."
    ]
   },
   {

From ec71a9c0aaa9bd2eb841a31b6fa83480f8f313ff Mon Sep 17 00:00:00 2001
From: akshay_shanker <andy.mathur@icloud.com>
Date: Wed, 16 Oct 2024 20:19:51 +1100
Subject: [PATCH 13/26] figure formatting.

---
 .../KS-HARK-presentation.ipynb                | 50 +++++++++----------
 1 file changed, 25 insertions(+), 25 deletions(-)

diff --git a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
index e8a1cdc80..62f58df2a 100644
--- a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
+++ b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
@@ -67,6 +67,8 @@
    "source": [
     "# Calibration and setup\n",
     "\n",
+    "This notebook uses the HARK toolkit to solve the [Krusell Smith (1998)](https://www.journals.uchicago.edu/doi/abs/10.1086/250034) model using the [Sequence Space Jacobian](https://onlinelibrary.wiley.com/doi/abs/10.3982/ECTA17434) method and [SSJ toolkit](https://github.com/shade-econ/sequence-jacobian). \n",
+    "\n",
     "## Firm setup"
    ]
   },
@@ -75,8 +77,6 @@
    "id": "df58ea5a",
    "metadata": {},
    "source": [
-    "\n",
-    "\n",
     "Collect calibration for the production economy in a dictionary."
    ]
   },
@@ -191,12 +191,18 @@
     "print(sol)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "e09205ba",
+   "metadata": {},
+   "source": []
+  },
   {
    "cell_type": "markdown",
    "id": "922dce5d",
    "metadata": {},
    "source": [
-    " Let's double check the roots we find produce our chosen steady state values. "
+    "Double check the roots we find produce our chosen steady state values. "
    ]
   },
   {
@@ -269,13 +275,13 @@
    "source": [
     "## HARK agent\n",
     "\n",
-    "We now create an instance of an agent.\n",
+    "HARK represents solving agent problems where aggregate variables can affect an agent's income process as instances of `NewKeynesianConsumerType`, a subclass of `AgentType` (see [A Gentle Introduction to HARK](https://docs.econ-ark.org/examples/Gentle-Intro/Gentle-Intro-To-HARK.html)).\n",
     "\n",
-    "HARK represents solving agent problems where aggregate variables can effect an agent's income process as instances of `HANKAgentType`, a subclass of `AgentType`. \n",
+    "**_NOTE:_**  An `HANKAgentType` is still a microeconomic agent solving an income fluctuation problem. The only difference from [`IndShockConsumerType`](https://docs.econ-ark.org/examples/HowWeSolveIndShockConsumerType/HowWeSolveIndShockConsumerType.html) is that additional aggregate labour income variables are passed to the agent's income process. This allows the agent's problem to be defined within a general equilibrium framework such as (but not restricted to) HANK. \n",
     "\n",
-    "**_NOTE:_**  An `HANKAgentType` is still a microeconomic agent solving an income fluctuation problem. The only difference from `IndShockConsumerType` is that additional aggregate labour income variables are passed to the agent's income process. This allows the agent's problem to be defined within a general equilirium framework such as (but not restricted to) HANK. \n",
+    "**_NOTE:_** Researchers can also create their own agent types by subclassing `AgentType`. See [HARK's documentation](https://docs.econ-ark.org/examples/HowWeSolveIndShockConsumerType/HowWeSolveIndShockConsumerType.html) for more information.\n",
     "\n",
-    "First specify the agent's parameters in a dictionary. For now we use a discount factor of 0.98, but we will update `DiscFac` below so the steady state is a general equilibrium."
+    "We now create an instance of an agent; first specify the agent's parameters in a dictionary. To start, we use a discount factor of 0.98."
    ]
   },
   {
@@ -345,7 +351,7 @@
    "id": "3be9593e",
    "metadata": {},
    "source": [
-    "The `NewKeynesianConsumerType` class contains within itself the Python code that constructs the solution for equilibrium agents using the SSJ toolkit, including computing the steady state policy functions and Jacobian matrices in response to a policy shock. \n",
+    "The `NewKeynesianConsumerType` class contains within itself the Python code that constructs the solution for equilibrium agents using the [SSJ toolkit] (https://github.com/shade-econ/sequence-jacobian), including computing the steady state policy functions and Jacobian matrices in response to a policy shock. \n",
     "\n",
     "Let's create a temporary instance of an HANK agent."
    ]
@@ -436,7 +442,7 @@
    "id": "01be6a8a",
    "metadata": {},
    "source": [
-    "Now we create the steady state agent using the general equilibrium discount factor."
+    "We can now create the steady state agent using the general equilibrium discount factor."
    ]
   },
   {
@@ -487,14 +493,6 @@
     "print(\"Asset clearing:\", A_ss - calibration['K_ss'])"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "86e05d73",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
   {
    "cell_type": "markdown",
    "id": "a8167383",
@@ -502,7 +500,7 @@
    "source": [
     "## Computing Jacobians using SSJ\n",
     "\n",
-    "With the steady state agent in hand, we can now compute the Jacobians of the steady state HANK agent's policy functions with respect to the aggregate state variables!\n",
+    "With the steady state agent in hand, we can compute the Jacobians of the steady state HANK agent's policy functions with respect to the aggregate state variables!\n",
     "\n",
     "The `calc_jacobian` method of our HANK agent computes the Jacobians of the steady state aggregate consumption (`CJACW`) and assets (`AJACW`) with respect to a pertubation of the specified variable.\n",
     "\n",
@@ -783,7 +781,6 @@
     "unknowns = [\"K\"]\n",
     "targets = [\"asset_mkt\"]\n",
     "\n",
-    "\n",
     "irfs_Z = ks.solve_impulse_linear(SteadyState_Dict, unknowns, targets, shocks)"
    ]
   },
@@ -805,7 +802,7 @@
     "    irfs_list,\n",
     "    variables,\n",
     "    labels=[\" \"],\n",
-    "    ylabel=r\"Percentage points (dev. from ss)\",\n",
+    "    ylabel=r\"Percentage point (dev. from ss)\",\n",
     "    T_plot=50,\n",
     "    figsize=(18, 6),\n",
     "):\n",
@@ -816,13 +813,17 @@
     "    for i in range(n_var):\n",
     "        # plot all irfs\n",
     "        for j, irf in enumerate(irfs_list):\n",
-    "            ax[i].plot(irf[variables[i]][:50], label=labels[j])\n",
+    "            ax[i].plot(irf[variables[i]][:50]*1e2, label=labels[j])\n",
     "        ax[i].set_title(variables[i])\n",
-    "        ax[i].set_xlabel(r\"$t$\")\n",
+    "        ax[i].set_xlabel(r\"Quarters\")\n",
     "        if i == 0:\n",
     "            ax[i].set_ylabel(ylabel)\n",
-    "        ax[i].legend()\n",
-    "    plt.show()"
+    "        #ax[i].legend()\n",
+    "        #ax[i].x\n",
+    "    \n",
+    "    plt.tight_layout()\n",
+    "    plt.show()\n",
+    "    #plt.savefig(\"irf.png\")"
    ]
   },
   {
@@ -911,7 +912,6 @@
     "rhos = {\"Z\": 0.8}\n",
     "impulses = {}\n",
     "\n",
-    "\n",
     "for i in inputs:\n",
     "    own_shock = {i: sigmas[i] * rhos[i] ** np.arange(T)}\n",
     "    impulses[i] = ks.solve_impulse_linear(\n",

From 0c200b76d08e85c6d7d3ba1593d655680d792dfc Mon Sep 17 00:00:00 2001
From: akshay_shanker <andy.mathur@icloud.com>
Date: Wed, 16 Oct 2024 20:20:49 +1100
Subject: [PATCH 14/26] links

---
 examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
index 62f58df2a..ec0808c6e 100644
--- a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
+++ b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
@@ -502,7 +502,7 @@
     "\n",
     "With the steady state agent in hand, we can compute the Jacobians of the steady state HANK agent's policy functions with respect to the aggregate state variables!\n",
     "\n",
-    "The `calc_jacobian` method of our HANK agent computes the Jacobians of the steady state aggregate consumption (`CJACW`) and assets (`AJACW`) with respect to a pertubation of the specified variable.\n",
+    "The []`calc_jacobian` method of our HANK agent computes the Jacobians of the steady state aggregate consumption (`CJACW`) and assets (`AJACW`) with respect to a pertubation of the specified variable.\n",
     "\n",
     "Recall `CJACW[s,t]` is the time $t$ response to a shock at time $s$. \n",
     "\n",

From 135208cc1eb187faa440f519287a91f3accd20e2 Mon Sep 17 00:00:00 2001
From: akshay_shanker <andy.mathur@icloud.com>
Date: Wed, 16 Oct 2024 20:24:49 +1100
Subject: [PATCH 15/26] minor.

---
 examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
index ec0808c6e..18f39dec4 100644
--- a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
+++ b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
@@ -277,7 +277,7 @@
     "\n",
     "HARK represents solving agent problems where aggregate variables can affect an agent's income process as instances of `NewKeynesianConsumerType`, a subclass of `AgentType` (see [A Gentle Introduction to HARK](https://docs.econ-ark.org/examples/Gentle-Intro/Gentle-Intro-To-HARK.html)).\n",
     "\n",
-    "**_NOTE:_**  An `HANKAgentType` is still a microeconomic agent solving an income fluctuation problem. The only difference from [`IndShockConsumerType`](https://docs.econ-ark.org/examples/HowWeSolveIndShockConsumerType/HowWeSolveIndShockConsumerType.html) is that additional aggregate labour income variables are passed to the agent's income process. This allows the agent's problem to be defined within a general equilibrium framework such as (but not restricted to) HANK. \n",
+    "**_NOTE:_**  An `NewKeynesianConsumerType` is still a microeconomic agent solving an income fluctuation problem. The only difference from [`IndShockConsumerType`](https://docs.econ-ark.org/examples/HowWeSolveIndShockConsumerType/HowWeSolveIndShockConsumerType.html) is that additional aggregate labour income variables are passed to the agent's income process. This allows the agent's problem to be defined within a general equilibrium framework such as (but not restricted to) HANK. \n",
     "\n",
     "**_NOTE:_** Researchers can also create their own agent types by subclassing `AgentType`. See [HARK's documentation](https://docs.econ-ark.org/examples/HowWeSolveIndShockConsumerType/HowWeSolveIndShockConsumerType.html) for more information.\n",
     "\n",
@@ -351,7 +351,7 @@
    "id": "3be9593e",
    "metadata": {},
    "source": [
-    "The `NewKeynesianConsumerType` class contains within itself the Python code that constructs the solution for equilibrium agents using the [SSJ toolkit] (https://github.com/shade-econ/sequence-jacobian), including computing the steady state policy functions and Jacobian matrices in response to a policy shock. \n",
+    "The `NewKeynesianConsumerType` class contains within itself the Python code that constructs the solution for equilibrium agents using the [SSJ toolkit](https://github.com/shade-econ/sequence-jacobian), including computing the steady state policy functions and Jacobian matrices in response to a policy shock. \n",
     "\n",
     "Let's create a temporary instance of an HANK agent."
    ]
@@ -502,7 +502,7 @@
     "\n",
     "With the steady state agent in hand, we can compute the Jacobians of the steady state HANK agent's policy functions with respect to the aggregate state variables!\n",
     "\n",
-    "The []`calc_jacobian` method of our HANK agent computes the Jacobians of the steady state aggregate consumption (`CJACW`) and assets (`AJACW`) with respect to a pertubation of the specified variable.\n",
+    "The `calc_jacobian` method of our HANK agent computes the Jacobians of the steady state aggregate consumption (`CJACW`) and assets (`AJACW`) with respect to a pertubation of the specified variable.\n",
     "\n",
     "Recall `CJACW[s,t]` is the time $t$ response to a shock at time $s$. \n",
     "\n",

From d5424a88ea9c0ce835c85823778e69ec573e7787 Mon Sep 17 00:00:00 2001
From: akshay_shanker <andy.mathur@icloud.com>
Date: Wed, 16 Oct 2024 20:28:57 +1100
Subject: [PATCH 16/26] headings.

---
 .../KS-HARK-presentation.ipynb                | 32 +++++++++----------
 1 file changed, 15 insertions(+), 17 deletions(-)

diff --git a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
index 18f39dec4..547156f77 100644
--- a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
+++ b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
@@ -69,14 +69,8 @@
     "\n",
     "This notebook uses the HARK toolkit to solve the [Krusell Smith (1998)](https://www.journals.uchicago.edu/doi/abs/10.1086/250034) model using the [Sequence Space Jacobian](https://onlinelibrary.wiley.com/doi/abs/10.3982/ECTA17434) method and [SSJ toolkit](https://github.com/shade-econ/sequence-jacobian). \n",
     "\n",
-    "## Firm setup"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "df58ea5a",
-   "metadata": {},
-   "source": [
+    "## Firm setup\n",
+    "\n",
     "Collect calibration for the production economy in a dictionary."
    ]
   },
@@ -109,7 +103,7 @@
    "id": "1bac99bf",
    "metadata": {},
    "source": [
-    "### Find the steady state capital and productivity \n",
+    "### Steady State Capital \n",
     "\n",
     "Find steady state capital (`K_ss`) and productivity  (`Z`) implied by the firm's first order condition and the above steady state values of labor (`L_ss`), the real interest rate (`r_ss`) and output (`Y_ss`). "
    ]
@@ -191,12 +185,6 @@
     "print(sol)"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "id": "e09205ba",
-   "metadata": {},
-   "source": []
-  },
   {
    "cell_type": "markdown",
    "id": "922dce5d",
@@ -392,6 +380,16 @@
     "print(f\"Steady state agent asset supply for beta = 0.98 is: {A_ss:.3f}\")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "911805ab",
+   "metadata": {},
+   "source": [
+    "### Steady State Assets \n",
+    "\n",
+    "We can find `discFac` which clears the asset market given steady state capital demand by the firm of `K_ss`. "
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 11,
@@ -498,7 +496,7 @@
    "id": "a8167383",
    "metadata": {},
    "source": [
-    "## Computing Jacobians using SSJ\n",
+    "# Computing Jacobians using SSJ\n",
     "\n",
     "With the steady state agent in hand, we can compute the Jacobians of the steady state HANK agent's policy functions with respect to the aggregate state variables!\n",
     "\n",
@@ -701,7 +699,7 @@
    "id": "d3669fab",
    "metadata": {},
    "source": [
-    "## Impulse response functions and simulations "
+    "## Impulse Response Functions and Simulations "
    ]
   },
   {

From 54b1ef361ae02906a3a4c2c5d4d4083fb2e8ab12 Mon Sep 17 00:00:00 2001
From: akshay_shanker <andy.mathur@icloud.com>
Date: Wed, 16 Oct 2024 20:33:22 +1100
Subject: [PATCH 17/26] minor.

---
 examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
index 547156f77..9d0e33f2e 100644
--- a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
+++ b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
@@ -265,7 +265,7 @@
     "\n",
     "HARK represents solving agent problems where aggregate variables can affect an agent's income process as instances of `NewKeynesianConsumerType`, a subclass of `AgentType` (see [A Gentle Introduction to HARK](https://docs.econ-ark.org/examples/Gentle-Intro/Gentle-Intro-To-HARK.html)).\n",
     "\n",
-    "**_NOTE:_**  An `NewKeynesianConsumerType` is still a microeconomic agent solving an income fluctuation problem. The only difference from [`IndShockConsumerType`](https://docs.econ-ark.org/examples/HowWeSolveIndShockConsumerType/HowWeSolveIndShockConsumerType.html) is that additional aggregate labour income variables are passed to the agent's income process. This allows the agent's problem to be defined within a general equilibrium framework such as (but not restricted to) HANK. \n",
+    "**_NOTE:_**  `NewKeynesianConsumerType` is still a microeconomic agent solving an income fluctuation problem. The only difference from [`IndShockConsumerType`](https://docs.econ-ark.org/examples/HowWeSolveIndShockConsumerType/HowWeSolveIndShockConsumerType.html) is that additional aggregate labour income variables are passed to the agent's income process. This allows the agent's problem to be defined within a general equilibrium framework such as (but not restricted to) HANK. \n",
     "\n",
     "**_NOTE:_** Researchers can also create their own agent types by subclassing `AgentType`. See [HARK's documentation](https://docs.econ-ark.org/examples/HowWeSolveIndShockConsumerType/HowWeSolveIndShockConsumerType.html) for more information.\n",
     "\n",

From 96962b7dfb77dc13c1710d7a18932f6dae4d9003 Mon Sep 17 00:00:00 2001
From: akshay_shanker <andy.mathur@icloud.com>
Date: Wed, 16 Oct 2024 22:02:40 +1100
Subject: [PATCH 18/26] Update KS-HARK-presentation.ipynb

---
 examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
index 9d0e33f2e..9f14534c1 100644
--- a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
+++ b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
@@ -67,7 +67,7 @@
    "source": [
     "# Calibration and setup\n",
     "\n",
-    "This notebook uses the HARK toolkit to solve the [Krusell Smith (1998)](https://www.journals.uchicago.edu/doi/abs/10.1086/250034) model using the [Sequence Space Jacobian](https://onlinelibrary.wiley.com/doi/abs/10.3982/ECTA17434) method and [SSJ toolkit](https://github.com/shade-econ/sequence-jacobian). \n",
+    "This notebook uses the HARK toolkit to solve [Krusell and Smith (1998)](https://www.journals.uchicago.edu/doi/abs/10.1086/250034), applying the [Sequence Space Jacobian](https://onlinelibrary.wiley.com/doi/abs/10.3982/ECTA17434) and [SSJ toolkit](https://github.com/shade-econ/sequence-jacobian). \n",
     "\n",
     "## Firm setup\n",
     "\n",

From c8e13ff0bf5fe944c8106e4c1b620fe7db1a556d Mon Sep 17 00:00:00 2001
From: akshay_shanker <andy.mathur@icloud.com>
Date: Wed, 16 Oct 2024 22:07:21 +1100
Subject: [PATCH 19/26] Update KS-HARK-presentation.ipynb

---
 examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
index 9f14534c1..48bd6299b 100644
--- a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
+++ b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
@@ -339,7 +339,7 @@
    "id": "3be9593e",
    "metadata": {},
    "source": [
-    "The `NewKeynesianConsumerType` class contains within itself the Python code that constructs the solution for equilibrium agents using the [SSJ toolkit](https://github.com/shade-econ/sequence-jacobian), including computing the steady state policy functions and Jacobian matrices in response to a policy shock. \n",
+    "Instances of the `NewKeynesianConsumerType` class contain functions that solve policy functions of optimizing agents and compute Jacobian matrices in response to a policy shocks using the [SSJ toolkit](https://github.com/shade-econ/sequence-jacobian). \n",
     "\n",
     "Let's create a temporary instance of an HANK agent."
    ]

From afb628e6e10ca111b12f9b163236d0ec0fe352a5 Mon Sep 17 00:00:00 2001
From: akshay_shanker <andy.mathur@icloud.com>
Date: Wed, 16 Oct 2024 22:08:04 +1100
Subject: [PATCH 20/26] Update KS-HARK-presentation.ipynb

---
 examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
index 48bd6299b..c1749c406 100644
--- a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
+++ b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
@@ -341,7 +341,7 @@
    "source": [
     "Instances of the `NewKeynesianConsumerType` class contain functions that solve policy functions of optimizing agents and compute Jacobian matrices in response to a policy shocks using the [SSJ toolkit](https://github.com/shade-econ/sequence-jacobian). \n",
     "\n",
-    "Let's create a temporary instance of an HANK agent."
+    "Let's create a temporary instance of a `NewKeynesianConsumerType` agent."
    ]
   },
   {

From 88fe518af9a825ba222a54e6fbe3dbf4d1c6dec7 Mon Sep 17 00:00:00 2001
From: akshay_shanker <andy.mathur@icloud.com>
Date: Wed, 16 Oct 2024 22:34:51 +1100
Subject: [PATCH 21/26] Minor polish.

---
 .../KS-HARK-presentation.ipynb                | 46 +++++++++++--------
 1 file changed, 27 insertions(+), 19 deletions(-)

diff --git a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
index c1749c406..6172f8374 100644
--- a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
+++ b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
@@ -105,7 +105,7 @@
    "source": [
     "### Steady State Capital \n",
     "\n",
-    "Find steady state capital (`K_ss`) and productivity  (`Z`) implied by the firm's first order condition and the above steady state values of labor (`L_ss`), the real interest rate (`r_ss`) and output (`Y_ss`). "
+    "Find steady state capital (`K_ss`) and productivity  (`Z`) implied by the firm's first order condition, aggregate resource constraint and the above steady state values of labor (`L_ss`), the real interest rate (`r_ss`) and output (`Y_ss`). "
    ]
   },
   {
@@ -133,9 +133,10 @@
     "    Y_ss = calibration['Y_ss'] \n",
     "\n",
     "    K, Z = X\n",
-    "    # all RHS have to be 0\n",
+    "\n",
+    "    # Firms first order condition and aggregate resource constraint\n",
     "    f = [\n",
-    "        alpha * Z * (K / L_ss) ** (alpha - 1) - delta - r_ss,  # r = MPK\n",
+    "        alpha * Z * (K / L_ss) ** (alpha - 1) - delta - r_ss,  # r_ss = MPK\n",
     "        Z * K**alpha * L_ss ** (1 - alpha) - Y_ss,  # Y = Z*F(K,L)\n",
     "    ]\n",
     "\n",
@@ -249,7 +250,6 @@
    "source": [
     "print(\"--------------------------------------+------------\")\n",
     "print(f\"Steady state capital                  | {K_ss:.4f}\")\n",
-    "print(f\"Steady state productivity             | {Z_ss:.4f}\")\n",
     "print(f\"Steady state output                   | {Y_ss:.4f}\")\n",
     "print(f\"Steady state real interest rate       | {r_ss:.4f}\")\n",
     "print(f\"Steady state wage                     | {w_ss:.4f}\")\n",
@@ -263,13 +263,15 @@
    "source": [
     "## HARK agent\n",
     "\n",
-    "HARK represents solving agent problems where aggregate variables can affect an agent's income process as instances of `NewKeynesianConsumerType`, a subclass of `AgentType` (see [A Gentle Introduction to HARK](https://docs.econ-ark.org/examples/Gentle-Intro/Gentle-Intro-To-HARK.html)).\n",
+    "HARK represents agent problems where aggregate variables can affect an individual income process as instances of `NewKeynesianConsumerType`, a subclass of `AgentType` (see [A Gentle Introduction to HARK](https://docs.econ-ark.org/examples/Gentle-Intro/Gentle-Intro-To-HARK.html)).\n",
+    "\n",
+    "Instances of the `NewKeynesianConsumerType` class contain functions that solve policy functions of optimizing agents and compute Jacobian matrices in response to a policy shocks using the [SSJ toolkit](https://github.com/shade-econ/sequence-jacobian). \n",
     "\n",
     "**_NOTE:_**  `NewKeynesianConsumerType` is still a microeconomic agent solving an income fluctuation problem. The only difference from [`IndShockConsumerType`](https://docs.econ-ark.org/examples/HowWeSolveIndShockConsumerType/HowWeSolveIndShockConsumerType.html) is that additional aggregate labour income variables are passed to the agent's income process. This allows the agent's problem to be defined within a general equilibrium framework such as (but not restricted to) HANK. \n",
     "\n",
     "**_NOTE:_** Researchers can also create their own agent types by subclassing `AgentType`. See [HARK's documentation](https://docs.econ-ark.org/examples/HowWeSolveIndShockConsumerType/HowWeSolveIndShockConsumerType.html) for more information.\n",
     "\n",
-    "We now create an instance of an agent; first specify the agent's parameters in a dictionary. To start, we use a discount factor of 0.98."
+    "To create an instance of `NewKeynesianConsumerType`, first specify parameters in a dictionary. To start, we use a discount factor of 0.98."
    ]
   },
   {
@@ -311,7 +313,7 @@
     "    \"IncUnempRet\": 0.0,  # \"Unemployment\" benefits when retired\n",
     "    \"T_retire\": 0.0,  # Period of retirement (0 --> no retirement)\n",
     "    \n",
-    "    # General equilibrium parameters \n",
+    "    # Aggregates affecting the agent's decision \n",
     "    \"Rfree\": (1 + r_ss),  # Interest factor for assets faced by agents \n",
     "    \"wage\": [w_ss],   # Wage rate faced by agents \n",
     "    \"tax_rate\": [0], # set to 0.0 because we are going to assume that labor here is actually after tax income\n",
@@ -339,8 +341,6 @@
    "id": "3be9593e",
    "metadata": {},
    "source": [
-    "Instances of the `NewKeynesianConsumerType` class contain functions that solve policy functions of optimizing agents and compute Jacobian matrices in response to a policy shocks using the [SSJ toolkit](https://github.com/shade-econ/sequence-jacobian). \n",
-    "\n",
     "Let's create a temporary instance of a `NewKeynesianConsumerType` agent."
    ]
   },
@@ -387,7 +387,7 @@
    "source": [
     "### Steady State Assets \n",
     "\n",
-    "We can find `discFac` which clears the asset market given steady state capital demand by the firm of `K_ss`. "
+    "Since we are interested in computing a steady state equilibrium, we find `discFac` such that `A_ss`  clears the asset market given that `K_ss` is the steady state firm capital demand. "
    ]
   },
   {
@@ -498,9 +498,9 @@
    "source": [
     "# Computing Jacobians using SSJ\n",
     "\n",
-    "With the steady state agent in hand, we can compute the Jacobians of the steady state HANK agent's policy functions with respect to the aggregate state variables!\n",
+    "With the steady state agent in hand, we can compute the Jacobians of the steady state HANK agent's policy with respect to the aggregate state variables.\n",
     "\n",
-    "The `calc_jacobian` method of our HANK agent computes the Jacobians of the steady state aggregate consumption (`CJACW`) and assets (`AJACW`) with respect to a pertubation of the specified variable.\n",
+    "The `calc_jacobian` method of our `NewKeynesianConsumerType` agent computes the Jacobians of the steady state aggregate consumption (`CJACW`) and assets (`AJACW`) with respect to a pertubation of a specified variable.\n",
     "\n",
     "Recall `CJACW[s,t]` is the time $t$ response to a shock at time $s$. \n",
     "\n",
@@ -561,15 +561,15 @@
     }
    ],
    "source": [
-    "plt.plot(CJACW.T[0], label = \"s =0\")\n",
+    "plt.plot(CJACW.T[0], label = \"s=0\")\n",
     "plt.plot(CJACW.T[20], label = 's=20')\n",
     "plt.plot(CJACW.T[50], label = 's=50')\n",
     "plt.plot(CJACW.T[100], label = 's=100')\n",
     "plt.xlim(-2, 300)\n",
     "plt.plot(np.arange(300), np.zeros(300), color=\"k\")\n",
-    "plt.title(\"Consumption Jacobian Wage\")\n",
-    "plt.xlabel(\"quarters\")\n",
-    "plt.ylabel(\"C response\")\n",
+    "plt.title(\"Consumption Jacobian (Wage) \")\n",
+    "plt.xlabel(\"Quarters\")\n",
+    "plt.ylabel(\"Consumption response\")\n",
     "plt.legend()\n",
     "plt.show()"
    ]
@@ -603,6 +603,14 @@
     "print(f\"Time taken to compute return factor Jacobians: {time.time() - start:.3f} seconds\")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "311a3a60",
+   "metadata": {},
+   "source": [
+    "Here is an example where we shock the return factor. "
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 17,
@@ -633,9 +641,9 @@
     "plt.plot(CJACR.T[50], label = 's=50')\n",
     "plt.plot(CJACR.T[100], label = 's=100')\n",
     "plt.plot(np.arange(300), np.zeros(300), color=\"k\")\n",
-    "plt.title(\"Consumption Jacobian interest rate\")\n",
-    "plt.xlabel(\"quarters\")\n",
-    "plt.ylabel(\"C response\")\n",
+    "plt.title(\"Consumption Jacobian (Return Factor)\")\n",
+    "plt.xlabel(\"Quarters\")\n",
+    "plt.ylabel(\"Consumption response\")\n",
     "plt.show()"
    ]
   },

From 1ad90e9711fe4df08ee1a2425f7bc1fc8e6fe5c5 Mon Sep 17 00:00:00 2001
From: alanlujan91 <alanlujan91@gmail.com>
Date: Wed, 16 Oct 2024 08:19:05 -0400
Subject: [PATCH 22/26] Update KS-HARK-presentation.ipynb

---
 .../KS-HARK-presentation.ipynb                | 89 ++++++++++---------
 1 file changed, 45 insertions(+), 44 deletions(-)

diff --git a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
index 6172f8374..5d24aa7c8 100644
--- a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
+++ b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
@@ -89,12 +89,12 @@
    "outputs": [],
    "source": [
     "calibration = {\n",
-    "    \"eis\": 1,       # Elasticity of intertemporal substitution \n",
-    "    \"delta\": 0.025, # Depreciation rate\n",
+    "    \"eis\": 1,  # Elasticity of intertemporal substitution\n",
+    "    \"delta\": 0.025,  # Depreciation rate\n",
     "    \"alpha\": 0.11,  # Capital share of income\n",
-    "    \"L_ss\": 1.0,    # Steady state labor\n",
-    "    \"Y_ss\": 1.0,    # Steady state output\n",
-    "    \"r_ss\": 0.01,   # Steady state real interest rate\n",
+    "    \"L_ss\": 1.0,  # Steady state labor\n",
+    "    \"Y_ss\": 1.0,  # Steady state output\n",
+    "    \"r_ss\": 0.01,  # Steady state real interest rate\n",
     "}"
    ]
   },
@@ -124,13 +124,13 @@
    "source": [
     "from scipy.optimize import root\n",
     "\n",
+    "\n",
     "def your_funcs(X):\n",
-    "    \n",
     "    L_ss = calibration[\"L_ss\"]\n",
     "    alpha = calibration[\"alpha\"]\n",
     "    delta = calibration[\"delta\"]\n",
-    "    r_ss = calibration['r_ss']  \n",
-    "    Y_ss = calibration['Y_ss'] \n",
+    "    r_ss = calibration[\"r_ss\"]\n",
+    "    Y_ss = calibration[\"Y_ss\"]\n",
     "\n",
     "    K, Z = X\n",
     "\n",
@@ -211,7 +211,7 @@
     "def firm(\n",
     "    K,\n",
     "    Z,\n",
-    "    L_ss =calibration[\"L_ss\"],\n",
+    "    L_ss=calibration[\"L_ss\"],\n",
     "    alpha=calibration[\"alpha\"],\n",
     "    delta=calibration[\"delta\"],\n",
     "):\n",
@@ -293,41 +293,36 @@
     "r_ss = calibration[\"r_ss\"]\n",
     "\n",
     "HANK_Dict = {\n",
-    "    \n",
     "    # Individual agent 'preferences' (shared with perfect foresight model)\n",
     "    \"CRRA\": calibration[\"eis\"],  # Coefficient of relative risk aversion\n",
     "    \"DiscFac\": 0.98,  # Intertemporal discount factor\n",
     "    \"LivPrb\": [0.99375],  # Survival probability\n",
-    "    \n",
-    "    # Individual lifcycle income process parameters \n",
+    "    # Individual lifcycle income process parameters\n",
     "    \"PermGroFac\": [1.00],  # Permanent income growth factor\n",
-    "    \"PermShkStd\": [0.06], # Standard deviation of log permanent shocks to income\n",
-    "    \"PermShkCount\": 5, # Number of points in discrete approximation to permanent income shocks\n",
-    "    \"TranShkStd\": [0.2],# Standard deviation of log transitory shocks to income\n",
-    "    \"TranShkCount\": 5,# Number of points in discrete approximation to transitory income shocks \n",
-    "\n",
-    "    # Parameters related to unemployment and retirement \n",
+    "    \"PermShkStd\": [0.06],  # Standard deviation of log permanent shocks to income\n",
+    "    \"PermShkCount\": 5,  # Number of points in discrete approximation to permanent income shocks\n",
+    "    \"TranShkStd\": [0.2],  # Standard deviation of log transitory shocks to income\n",
+    "    \"TranShkCount\": 5,  # Number of points in discrete approximation to transitory income shocks\n",
+    "    # Parameters related to unemployment and retirement\n",
     "    \"UnempPrb\": 0.0,  # Probability of unemployment while working\n",
     "    \"IncUnemp\": 0.0,  # Unemployment benefits replacement rate\n",
     "    \"UnempPrbRet\": 0.0000,  # Probability of \"unemployment\" while retired\n",
     "    \"IncUnempRet\": 0.0,  # \"Unemployment\" benefits when retired\n",
     "    \"T_retire\": 0.0,  # Period of retirement (0 --> no retirement)\n",
-    "    \n",
-    "    # Aggregates affecting the agent's decision \n",
-    "    \"Rfree\": (1 + r_ss),  # Interest factor for assets faced by agents \n",
-    "    \"wage\": [w_ss],   # Wage rate faced by agents \n",
-    "    \"tax_rate\": [0], # set to 0.0 because we are going to assume that labor here is actually after tax income\n",
-    "    \"labor\": [L_ss], # Aggregate (mean) labor supply\n",
-    "    \n",
+    "    # Aggregates affecting the agent's decision\n",
+    "    \"Rfree\": (1 + r_ss),  # Interest factor for assets faced by agents\n",
+    "    \"wage\": [w_ss],  # Wage rate faced by agents\n",
+    "    \"tax_rate\": [\n",
+    "        0\n",
+    "    ],  # set to 0.0 because we are going to assume that labor here is actually after tax income\n",
+    "    \"labor\": [L_ss],  # Aggregate (mean) labor supply\n",
     "    # Parameters for constructing \"assets above minimum\" grid\n",
     "    \"aXtraMin\": 0.0001,  # Minimum end-of-period \"assets above minimum\" value\n",
     "    \"aXtraMax\": 2000,  # Maximum end-of-period \"assets above minimum\" value\n",
     "    \"aXtraCount\": 200,  # Number of points in the base grid of \"assets above minimum\"\n",
-    "    \n",
     "    # Exponential nesting factor when constructing \"assets above minimum\" grid\n",
     "    \"aXtraNestFac\": 3,\n",
     "    \"aXtraExtra\": None,  # Additional values to add to aXtraGrid\n",
-    "    \n",
     "    # Transition matrix simulation parameters\n",
     "    \"mCount\": 200,\n",
     "    \"mMax\": 2000,\n",
@@ -418,11 +413,12 @@
     "    # For a given value of beta, we need to re-create the Agent\n",
     "    Agent_func = NewKeynesianConsumerType(**HANK_Dict, verbose=False)\n",
     "\n",
-    "    # And then solve for the steady state supply \n",
+    "    # And then solve for the steady state supply\n",
     "    A_ss = Agent_func.compute_steady_state()[0]\n",
     "\n",
     "    return A_ss\n",
     "\n",
+    "\n",
     "def ss_dif(beta):\n",
     "    return A_ss_func(beta) - Asset_target\n",
     "\n",
@@ -487,8 +483,11 @@
    ],
    "source": [
     "# To make sure goods and asset markets clear\n",
-    "print(\"Final goods clearing:\", calibration['Y_ss'] - C_ss - calibration[\"delta\"] * calibration['K_ss'])\n",
-    "print(\"Asset clearing:\", A_ss - calibration['K_ss'])"
+    "print(\n",
+    "    \"Final goods clearing:\",\n",
+    "    calibration[\"Y_ss\"] - C_ss - calibration[\"delta\"] * calibration[\"K_ss\"],\n",
+    ")\n",
+    "print(\"Asset clearing:\", A_ss - calibration[\"K_ss\"])"
    ]
   },
   {
@@ -561,10 +560,10 @@
     }
    ],
    "source": [
-    "plt.plot(CJACW.T[0], label = \"s=0\")\n",
-    "plt.plot(CJACW.T[20], label = 's=20')\n",
-    "plt.plot(CJACW.T[50], label = 's=50')\n",
-    "plt.plot(CJACW.T[100], label = 's=100')\n",
+    "plt.plot(CJACW.T[0], label=\"s=0\")\n",
+    "plt.plot(CJACW.T[20], label=\"s=20\")\n",
+    "plt.plot(CJACW.T[50], label=\"s=50\")\n",
+    "plt.plot(CJACW.T[100], label=\"s=100\")\n",
     "plt.xlim(-2, 300)\n",
     "plt.plot(np.arange(300), np.zeros(300), color=\"k\")\n",
     "plt.title(\"Consumption Jacobian (Wage) \")\n",
@@ -600,7 +599,9 @@
     "\n",
     "CJACR, AJACR = steadyHANK.calc_jacobian(\"Rfree\", 300)  # Rfree jacobians\n",
     "\n",
-    "print(f\"Time taken to compute return factor Jacobians: {time.time() - start:.3f} seconds\")"
+    "print(\n",
+    "    f\"Time taken to compute return factor Jacobians: {time.time() - start:.3f} seconds\"\n",
+    ")"
    ]
   },
   {
@@ -636,10 +637,10 @@
     }
    ],
    "source": [
-    "plt.plot(CJACR.T[0], label = \"s =0\")\n",
-    "plt.plot(CJACR.T[20], label = 's=20')\n",
-    "plt.plot(CJACR.T[50], label = 's=50')\n",
-    "plt.plot(CJACR.T[100], label = 's=100')\n",
+    "plt.plot(CJACR.T[0], label=\"s =0\")\n",
+    "plt.plot(CJACR.T[20], label=\"s=20\")\n",
+    "plt.plot(CJACR.T[50], label=\"s=50\")\n",
+    "plt.plot(CJACR.T[100], label=\"s=100\")\n",
     "plt.plot(np.arange(300), np.zeros(300), color=\"k\")\n",
     "plt.title(\"Consumption Jacobian (Return Factor)\")\n",
     "plt.xlabel(\"Quarters\")\n",
@@ -819,17 +820,17 @@
     "    for i in range(n_var):\n",
     "        # plot all irfs\n",
     "        for j, irf in enumerate(irfs_list):\n",
-    "            ax[i].plot(irf[variables[i]][:50]*1e2, label=labels[j])\n",
+    "            ax[i].plot(irf[variables[i]][:50] * 1e2, label=labels[j])\n",
     "        ax[i].set_title(variables[i])\n",
     "        ax[i].set_xlabel(r\"Quarters\")\n",
     "        if i == 0:\n",
     "            ax[i].set_ylabel(ylabel)\n",
-    "        #ax[i].legend()\n",
-    "        #ax[i].x\n",
-    "    \n",
+    "        # ax[i].legend()\n",
+    "        # ax[i].x\n",
+    "\n",
     "    plt.tight_layout()\n",
     "    plt.show()\n",
-    "    #plt.savefig(\"irf.png\")"
+    "    # plt.savefig(\"irf.png\")"
    ]
   },
   {

From f22a3ca2a086f3fa793a8b36ece989bf5d5fcbe2 Mon Sep 17 00:00:00 2001
From: Mridul Seth <mail@mriduls.com>
Date: Wed, 16 Oct 2024 16:55:12 +0200
Subject: [PATCH 23/26] MAINT: Clean up binder folder, (#1502)

* MAINT: Clean up binder folder, requirements are controlled by pyproject.toml

* Delete binder/postBuild

* Delete binder/postBuild.bat
---
 binder/postBuild        | 10 ----------
 binder/postBuild.bat    |  3 ---
 binder/requirements.txt |  7 -------
 3 files changed, 20 deletions(-)
 delete mode 100755 binder/postBuild
 delete mode 100644 binder/postBuild.bat

diff --git a/binder/postBuild b/binder/postBuild
deleted file mode 100755
index 4e4b22864..000000000
--- a/binder/postBuild
+++ /dev/null
@@ -1,10 +0,0 @@
-#!/bin/sh
-jupyter contrib nbextension install --user
-jupyter nbextension enable codefolding/main
-python -m cite2c.install
-result=$(python <<EOF
-from notebook.services.config.manager import ConfigManager
-cm = ConfigManager()
-cm.update('cite2c', {'zotero':{"user_id": "5043554","username": "econ-ark","access_token": "XZpH9NsoAZmDMmjLKiy8xMXX"}})
-EOF
-)
diff --git a/binder/postBuild.bat b/binder/postBuild.bat
deleted file mode 100644
index b95f9ed64..000000000
--- a/binder/postBuild.bat
+++ /dev/null
@@ -1,3 +0,0 @@
-jupyter contrib nbextension install --user
-jupyter nbextension enable codefolding/main
-python -m cite2c.install
diff --git a/binder/requirements.txt b/binder/requirements.txt
index 78add7d76..5b74430ab 100644
--- a/binder/requirements.txt
+++ b/binder/requirements.txt
@@ -1,8 +1 @@
-cite2c
 git+https://github.com/econ-ark/hark@master
-ipywidgets
-jupyter_contrib_nbextensions
-matplotlib
-numpy
-scipy
-seaborn

From 73db66b66195d8f74886bb1fba325ecbbbbf521b Mon Sep 17 00:00:00 2001
From: Mridul Seth <mail@mriduls.com>
Date: Wed, 16 Oct 2024 17:10:00 +0200
Subject: [PATCH 24/26] Add sequence_jacobian to requirements for binder 
 (#1503)

* add sequence_jacobian to requirements for binder

* Remove manual pip install
---
 binder/requirements.txt                       |  1 +
 .../KS-HARK-presentation.ipynb                | 25 -------------------
 2 files changed, 1 insertion(+), 25 deletions(-)

diff --git a/binder/requirements.txt b/binder/requirements.txt
index 5b74430ab..ed0b97cfd 100644
--- a/binder/requirements.txt
+++ b/binder/requirements.txt
@@ -1 +1,2 @@
 git+https://github.com/econ-ark/hark@master
+sequence_jacobian # required for KS-HARK-presentation example
diff --git a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
index 5d24aa7c8..8da2bb30d 100644
--- a/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
+++ b/examples/ConsNewKeynesianModel/KS-HARK-presentation.ipynb
@@ -10,31 +10,6 @@
     "By William Du"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "id": "46e98a1e",
-   "metadata": {
-    "execution": {
-     "iopub.execute_input": "2024-07-11T15:18:12.929098Z",
-     "iopub.status.busy": "2024-07-11T15:18:12.928854Z",
-     "iopub.status.idle": "2024-07-11T15:18:14.164584Z",
-     "shell.execute_reply": "2024-07-11T15:18:14.163965Z"
-    }
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Note: you may need to restart the kernel to use updated packages.\n"
-     ]
-    }
-   ],
-   "source": [
-    "%pip install sequence_jacobian -q"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": 2,

From 3a9cfa76415d79fd3328e0c39870cd4e0d22c2ef Mon Sep 17 00:00:00 2001
From: sidd3888 <siddarthv2011@gmail.com>
Date: Sun, 27 Oct 2024 15:07:12 +0530
Subject: [PATCH 25/26] tail points and divide by zero fix

---
 HARK/distribution.py | 47 +++++++++++++++++++++++++++++---------------
 1 file changed, 31 insertions(+), 16 deletions(-)

diff --git a/HARK/distribution.py b/HARK/distribution.py
index 274765b2d..eb3f7ab38 100644
--- a/HARK/distribution.py
+++ b/HARK/distribution.py
@@ -963,11 +963,11 @@ def _approx_equiprobable(
             )
 
         if np.array_equal(self.Sigma, np.diag(np.diag(self.Sigma))):
-            ind_atoms = np.empty((self.M, N))
+            ind_atoms = np.empty((self.M, N + 2 * tail_N))
 
             for i in range(self.M):
                 if self.Sigma[i, i] == 0.0:
-                    x_atoms = np.repeat(np.exp(self.mu[i]), N)
+                    x_atoms = np.repeat(np.exp(self.mu[i]), N + 2 * tail_N)
                     ind_atoms[i] = x_atoms
                 else:
                     x_atoms = (
@@ -983,7 +983,22 @@ def _approx_equiprobable(
             atoms = np.stack(
                 [ar.flatten() for ar in list(np.meshgrid(*atoms_list))], axis=1
             ).T
-            pmv = np.repeat(1 / (N**self.M), N**self.M)
+
+            interiors = np.empty([self.M, (N + 2 * tail_N) ** (self.M)])
+
+            inners = np.zeros(N + 2 * tail_N)
+
+            if tail_N > 0:
+                inners[:tail_N] = [(tail_N - i) for i in range(tail_N)]
+                inners[-tail_N:] = [(i + 1) for i in range(tail_N)]
+
+            for i in range(self.M):
+                inners_i = [inners for _ in range((N + 2 * tail_N) ** i)]
+
+                interiors[i] = np.repeat(
+                    [*inners_i], (N + 2 * tail_N) ** (self.M - (i + 1))
+                )
+
         else:
             if tail_bound is not None:
                 if type(tail_bound) is float:
@@ -1024,10 +1039,10 @@ def eval(params, z):
                 excl = []
 
                 for j in range(len(z)):
-                    if z[j, 0] != z[j, 1]:
-                        inds.append(j)
-                    else:
+                    if z[j, 0] == z[j, 1]:
                         excl.append(j)
+                    elif params[j] != 0.0:
+                        inds.append(j)
 
                 dim = len(inds)
 
@@ -1111,21 +1126,21 @@ def eval(params, z):
 
                     atoms[i] = xi_atoms
 
-            max_locs = np.argmax(np.abs(interiors), axis=0)
+        max_locs = np.argmax(np.abs(interiors), axis=0)
 
-            max_inds = np.stack([max_locs, np.arange(len(max_locs))], axis=1)
+        max_inds = np.stack([max_locs, np.arange(len(max_locs))], axis=1)
 
-            prob_locs = interiors[max_inds[:, 0], max_inds[:, 1]]
+        prob_locs = interiors[max_inds[:, 0], max_inds[:, 1]]
 
-            def prob_assign(x):
-                if x == 0:
-                    return 1 / (N**self.M)
-                else:
-                    return 0.0
+        def prob_assign(x):
+            if x == 0:
+                return 1 / (N**self.M)
+            else:
+                return 0.0
 
-            prob_vec = np.vectorize(prob_assign)
+        prob_vec = np.vectorize(prob_assign)
 
-            pmv = prob_vec(prob_locs)
+        pmv = prob_vec(prob_locs)
 
         limit = {
             "dist": self,

From 894d53e9c7e425d3fad0828903bb551e42f4026e Mon Sep 17 00:00:00 2001
From: Emmanuel Ferdman <emmanuelferdman@gmail.com>
Date: Thu, 31 Oct 2024 22:03:14 +0200
Subject: [PATCH 26/26] update transition matrices notebook reference (#1508)

Signed-off-by: Emmanuel Ferdman <emmanuelferdman@gmail.com>
---
 examples/Journeys/Journey-Policymaker.ipynb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/Journeys/Journey-Policymaker.ipynb b/examples/Journeys/Journey-Policymaker.ipynb
index 8a8fe51c9..d3ce854c2 100644
--- a/examples/Journeys/Journey-Policymaker.ipynb
+++ b/examples/Journeys/Journey-Policymaker.ipynb
@@ -325,7 +325,7 @@
    "metadata": {},
    "source": [
     "### 2.3 Simulation\n",
-    "After we solved the model backwards, we can simulate agents forward using Monte Carlo or [Transition Matrices](https://github.com/econ-ark/HARK/tree/master/examples/ConsIndShockModel/IndShockConsumerType_Transition_Matrix_Example.ipynb). These results can be used for in-model regressions or plotting distributions of assets or consumption."
+    "After we solved the model backwards, we can simulate agents forward using Monte Carlo or [Transition Matrices](https://github.com/econ-ark/HARK/tree/master/examples/ConsNewKeynesianModel/Transition_Matrix_Example.ipynb). These results can be used for in-model regressions or plotting distributions of assets or consumption."
    ]
   },
   {