Skip to content

Commit b27f1eb

Browse files
jstacclaudemmcky
authored
Reorganize and rename optimal savings lectures (#735)
* Reorganize and rename optimal savings lectures This commit restructures the cake eating lecture series into a more coherent "Introduction to Optimal Savings" section with clearer naming and terminology. Changes: - Created new "Introduction to Optimal Savings" section in table of contents - Renamed all 6 lectures with "os" prefix for consistency: * cake_eating.md → os.md (Optimal Savings I: Cake Eating) * cake_eating_numerical.md → os_numerical.md (Optimal Savings II: Numerical Cake Eating) * cake_eating_stochastic.md → os_stochastic.md (Optimal Savings III: Stochastic Returns) * cake_eating_time_iter.md → os_time_iter.md (Optimal Savings IV: Time Iteration) * cake_eating_egm.md → os_egm.md (Optimal Savings V: The Endogenous Grid Method) * cake_eating_egm_jax.md → os_egm_jax.md (Optimal Savings VI: EGM with JAX) - Updated all lecture titles to use "Optimal Savings I-VI" naming convention - Replaced "cake eating" terminology with "optimal savings" throughout all lectures - Updated terminology in os_stochastic.md: "cake" → "wealth/harvest" to better reflect stochastic growth - Updated all cross-references across the codebase to use new filenames - Made cross-references robust to future title changes by using {doc}`filename` format 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * Add IFP I lecture and restructure IFP series - Add new ifp_discrete.md as "IFP I: Discretization and VFI" - Implements income fluctuation problem using discretization and value function iteration - Uses Model NamedTuple for clean parameter management - Includes both Python loop and jax.lax.while_loop implementations for comparison - Demonstrates proper JAX benchmarking with block_until_ready() - Shows 3-4x speedup from using jax.lax.while_loop over plain Python - Converts vmap implementation section to exercise format - Rename ifp.md to ifp_egm.md as "IFP II: The Endogenous Grid Method" - Updated title and cross-references - Maintains continuity with new IFP I lecture - Update _toc.yml to reflect new lecture order - Update cross-references in ifp_advanced.md 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * Fix build errors in ifp_discrete.md - Add URL for Epstein-Zin preferences link - Remove reference to non-existent opt_savings_2 lecture - Fix exercise syntax: use {exercise} with {solution-start}/{solution-end} 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * Fix exercise syntax in ifp_egm.md - Change {exercise-start}/{exercise-end} to {exercise} - Keep {solution-start}/{solution-end} syntax for solutions 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * Add IFP II lecture on Optimistic Policy Iteration - New ifp_opi.md lecture implementing OPI for income fluctuation problem - Uses Model NamedTuple structure consistent with ifp_discrete.md - Implements policy operator T_σ and iterate_policy_operator - Provides comprehensive timing comparisons between VFI and OPI - Shows 2-2.5x speedup from OPI over VFI - Includes exercise exploring parameter sensitivity - Update _toc.yml to include ifp_opi between ifp_discrete and ifp_egm - Update ifp_egm.md title from "IFP II" to "IFP III" - Reflects new lecture ordering with OPI as second lecture Technical highlights: - Uses @jax.jit decorators and jax.lax.while_loop for performance - Properly handles JAX traced values with jnp.where instead of Python if - Tests multiple values of m (policy iteration steps) to find optimal - Visualization comparing OPI performance across different m values Tested: Converts to Python and runs successfully with realistic speedups 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * Fix exercise labels in ifp_egm.md - Change ifp_ex1, ifp_ex2, ifp_ex3 to ifp_egm_ex1, ifp_egm_ex2, ifp_egm_ex3 - Prevents label conflicts with ifp_opi.md exercises - Fix exercise-start/exercise-end syntax for ifp_ex3 to use {exercise} 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * Add @jax.jit decorators to VFI and OPI functions - Improves performance by JIT compiling the main iteration routines - Speedup now consistently 2.5-2.6x (vs 2.4-2.5x before) - Absolute times also improved (OPI: ~0.3-0.4s vs ~0.4-0.5s) 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * Extend m_vals range in OPI performance testing - Add m=1, 200, 400 to test range - Now tests m_vals = [1, 5, 10, 25, 50, 100, 200, 400] - Shows performance across wider range of policy iteration steps 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * Add explanation of OPI performance across different m values - Explains why m=1 is slower than VFI (implementation overhead) - Notes optimal performance at m=25-50 (3x speedup) - Explains degradation for large m (200, 400) - Emphasizes the 'sweet spot' concept for choosing m 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * Fix build errors in ifp_egm.md - Change equation label from 'eqvfs' to 'eqvfs_egm' to avoid duplicate with ifp_discrete.md - Add missing closing backticks after solution-end directives 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * fix: nested note admonition in exercise * fix: exercise gates * fix: link to earlier exercise * Minor edits to code and markdown * fix: JAX compatibility and code improvements in IFP and OS lectures Fixed JAX implementation issues and improved code quality across multiple lectures: ## ifp_egm.md - Fixed compute_asset_stationary() argument order (c_vals, ae_vals, ifp) - Fixed jax.vmap() to use in_axes parameter instead of axes - Fixed fori_loop update function signature (t, state) instead of (state, t) - Fixed jax.random.fold_in argument order - Added int32 type casting for JAX compatibility - Improved code comments and documentation - Reorganized simulation section before exercises ## os_numerical.md - Simplified maximize() function by removing unused args parameter - Renamed state_action_value() to B() for clarity - Improved function documentation and code organization - Fixed code examples to use simplified function signatures ## Minor edits to ifp_advanced.md and os.md All lectures now convert to Python via jupytext and run without errors. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * fix: improve parameter naming and function signatures in ifp_egm - Simplify K_numpy parameter names from c_vals_init/ae_vals_init to c_vals/ae_vals - Update solve_model and solve_model_numpy signatures to accept both initial conditions - Fix argument order in compute_asset_stationary calls - Add clear comments for initial conditions setup - Standardize parameter ordering across NumPy and JAX implementations 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * fix: resolve build errors in ifp_egm and os_numerical ifp_egm.md fixes: - Add missing initialization in solve_model_numpy (c_vals, ae_vals = c_vals_init, ae_vals_init) - Fix update function parameter order in simulate_household (t, state instead of state, t) - Fix jax.random.fold_in argument order (key, t instead of t, key) - Add .astype(jnp.int32) to z_next_idx to fix dtype mismatch - Update jax.vmap to use in_axes instead of axes - Add missing arguments to sim_all_households call - Fix compute_asset_stationary argument order in all calls os_numerical.md fixes: - Simplify maximize function signature from (g, upper_bound, args) to (g, upper_bound) - Remove unused args parameter and tuple unpacking All changes tested by converting to Python with jupytext and running successfully. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> --------- Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Matt McKay <mmcky@users.noreply.github.com> Co-authored-by: mmcky <mamckay@gmail.com>
1 parent 96662c4 commit b27f1eb

File tree

13 files changed

+1283
-294
lines changed

13 files changed

+1283
-294
lines changed

lectures/_toc.yml

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -72,16 +72,21 @@ parts:
7272
- file: jv
7373
- file: odu
7474
- file: mccall_q
75+
- caption: Introduction to Optimal Savings
76+
numbered: true
77+
chapters:
78+
- file: os
79+
- file: os_numerical
80+
- file: os_stochastic
81+
- file: os_time_iter
82+
- file: os_egm
83+
- file: os_egm_jax
7584
- caption: Household Problems
7685
numbered: true
7786
chapters:
78-
- file: cake_eating
79-
- file: cake_eating_numerical
80-
- file: cake_eating_stochastic
81-
- file: cake_eating_time_iter
82-
- file: cake_eating_egm
83-
- file: cake_eating_egm_jax
84-
- file: ifp
87+
- file: ifp_discrete
88+
- file: ifp_opi
89+
- file: ifp_egm
8590
- file: ifp_advanced
8691
- caption: LQ Control
8792
numbered: true

lectures/ifp_advanced.md

Lines changed: 9 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ kernelspec:
1717
</div>
1818
```
1919

20-
# {index}`The Income Fluctuation Problem II: Stochastic Returns on Assets <single: The Income Fluctuation Problem II: Stochastic Returns on Assets>`
20+
# {index}`The Income Fluctuation Problem IV: Stochastic Returns on Assets <single: The Income Fluctuation Problem IV: Stochastic Returns on Assets>`
2121

2222
```{contents} Contents
2323
:depth: 2
@@ -34,7 +34,7 @@ tags: [hide-output]
3434

3535
## Overview
3636

37-
In this lecture, we continue our study of the {doc}`income fluctuation problem <ifp>`.
37+
In this lecture, we continue our study of the income fluctuation problem described in {doc}`ifp_egm`.
3838

3939
While the interest rate was previously taken to be fixed, we now allow
4040
returns on assets to be state-dependent.
@@ -112,7 +112,7 @@ where
112112

113113
Let $P$ represent the Markov matrix for the chain $\{Z_t\}_{t \geq 0}$.
114114

115-
Our assumptions on preferences are the same as our {doc}`previous lecture <ifp>` on the income fluctuation problem.
115+
Our assumptions on preferences are the same as in {doc}`ifp_egm`.
116116

117117
As before, $\mathbb E_z \hat X$ means expectation of next period value
118118
$\hat X$ given current value $Z = z$.
@@ -160,8 +160,7 @@ the IID and CRRA environment of {cite}`benhabib2015`.
160160

161161
### Optimality
162162

163-
Let the class of candidate consumption policies $\mathscr C$ be defined
164-
{doc}`as before <ifp>`.
163+
Let the class of candidate consumption policies $\mathscr C$ be defined as in {doc}`ifp_egm`.
165164

166165
In {cite}`ma2020income` it is shown that, under the stated assumptions,
167166

@@ -182,8 +181,7 @@ In the present setting, the Euler equation takes the form
182181
\right\}
183182
```
184183

185-
(Intuition and derivation are similar to our {doc}`earlier lecture <ifp>` on
186-
the income fluctuation problem.)
184+
(Intuition and derivation are similar to {doc}`ifp_egm`.)
187185

188186
We again solve the Euler equation using time iteration, iterating with a
189187
Coleman--Reffett operator $K$ defined to match the Euler equation
@@ -197,8 +195,7 @@ Coleman--Reffett operator $K$ defined to match the Euler equation
197195
### A Time Iteration Operator
198196

199197
Our definition of the candidate class $\sigma \in \mathscr C$ of consumption
200-
policies is the same as in our {doc}`earlier lecture <ifp>` on the income
201-
fluctuation problem.
198+
policies is the same as in {doc}`ifp_egm`.
202199

203200
For fixed $\sigma \in \mathscr C$ and $(a,z) \in \mathbf S$, the value
204201
$K\sigma(a,z)$ of the function $K\sigma$ at $(a,z)$ is defined as the
@@ -251,7 +248,7 @@ convergence (as measured by the distance $\rho$).
251248
### Using an Endogenous Grid
252249

253250
In the study of that model we found that it was possible to further
254-
accelerate time iteration via the {doc}`endogenous grid method <cake_eating_egm>`.
251+
accelerate time iteration via the {doc}`endogenous grid method <os_egm>`.
255252

256253
We will use the same method here.
257254

@@ -578,7 +575,7 @@ In contrast, when $z=1$ (good state), higher expected future income allows the h
578575
Let's try to get some idea of what will happen to assets over the long run
579576
under this consumption policy.
580577

581-
As with our {doc}`earlier lecture <ifp>` on the income fluctuation problem, we
578+
As in {doc}`ifp_egm`, we
582579
begin by producing a 45 degree diagram showing the law of motion for assets
583580

584581
```{code-cell} python3
@@ -911,7 +908,7 @@ The JAX implementation provides several advantages:
911908
```{exercise}
912909
:label: ifpa_ex1
913910
914-
Let's repeat our {ref}`earlier exercise <ifp_ex2>` on the long-run
911+
Let's repeat our {ref}`earlier exercise <ifp_egm_ex2>` on the long-run
915912
cross sectional distribution of assets.
916913
917914
In that exercise, we used a relatively simple income fluctuation model.

0 commit comments

Comments
 (0)