Updates to IFP (#682)

* misc

* Update ifp.md: Convert from Numba to JAX with optimized EGM implementation

Converted the Income Fluctuation Problem lecture from Numba to JAX implementation with significant improvements:

**Key Changes:**
- Replaced NamedTuple syntax errors (brackets to proper syntax)
- Added missing imports: `jax`, `from typing import NamedTuple`
- Fixed `create_ifp()` function: corrected assertion to use local variables instead of `self`
- Implemented efficient vectorized K operator using JAX vmap (~4,400 solves/second)
- Added comprehensive step-by-step comments explaining the Endogenous Grid Method algorithm
- Fixed all variable naming issues (a_grid → asset_grid, σ_array → σ, model → ifp)
- Corrected initial guess: σ_init = R * asset_grid[:, None] + y(z_grid)
- Updated all test code and examples to use correct function names and variables

**Performance:**
- Optimized K operator eliminates all Python for loops
- Vectorized expected marginal utility computation: u_prime_vals @ Π[j, :]
- Used jax.vmap for efficient parallelization over income states
- Result: ~0.23 ms per solve with proper block_until_ready()

**Documentation:**
- Added detailed 5-step breakdown of EGM algorithm in K operator
- Included shape annotations for all intermediate arrays
- Explained economic interpretation of each computational step

All code tested and verified to satisfy budget constraints (0 ≤ c ≤ R*a + y).

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

* Fix line length and exercise code in ifp.md

- Ensured all Python code lines are ≤80 characters
- Fixed all exercises to use create_ifp() instead of IFP()
- Fixed compute_asset_series to accept σ_init parameter
- Fixed simulation to use correct budget constraint: a_{t+1} = R*a_t + y - c
- Fixed all variable references in exercises (a_grid → asset_grid, etc.)
- Tested by converting to .py and running successfully

All code now runs correctly from md → py conversion.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

* Update ifp.md: Optimize simulation and fix parameter stability

Key improvements to the Income Fluctuation Problem lecture:

**Simulation optimization:**
- Replaced sequential single-household simulation with parallel multi-household approach
- Simulates 50,000 households for 500 periods using JAX's vmap for efficiency
- Leverages ergodicity: cross-sectional distribution approximates stationary distribution
- Uses jax.lax.scan with pre-split random keys for 2x performance vs fori_loop
- Changed variable naming from 'carry' to 'state' for clarity

**Parameter fixes:**
- Increased β from 0.96 to 0.98 for non-degenerate stationary distribution
- Increased asset grid max from 16 to 20, then to 40 to prevent grid boundary issues
- Reduced good shock from 0.25 to 0.2 for stable asset accumulation
- Restricted interest rate ranges to ensure R*β < 1 stability condition
- Added random initial assets to avoid zero-asset absorbing state

**Code quality:**
- Standardized all code cells to use 'ipython' language
- Fixed plot axes in Exercise 3 (interest rate on x-axis, capital on y-axis)
- Added debug output for mean assets calculation
- Removed old inefficient simulation approach

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

* Fix equation references: replace eqeul0 with ee00-ee01

The references to non-existent equation label 'eqeul0' were causing
build warnings. Updated to reference the correct Euler equation
labels 'ee00' and 'ee01'.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

* misc

* Update IFP lecture: Change timing to a' = R(a - c) + Y'

This commit updates the Income Fluctuation Problem (IFP) lecture to use
a different timing convention for the budget constraint:
- Changed from: a' + c ≤ Ra + Y to a' = R(a - c) + Y'
- Updated timing description to reflect that income is realized next period
- Revised EGM implementation to use savings grid approach (s_i)
- Updated all code including K operator, simulations, and initial guesses
- Fixed boundary condition: consume everything (c = a) when constrained
- Added ifp.py generated from ifp.md using jupytext

All numerical examples and exercises have been verified to run correctly.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

* Refine IFP lecture: Improve mathematical clarity and plotting

Key improvements:
- Added clarification that Euler equation (eqeul1) holds only for interior solutions
- Specified savings grid is strictly increasing
- Fixed EGM boundary condition: anchor interpolation at (0,0) instead of c=a
- Renamed asset_grid to savings_grid throughout for conceptual accuracy
- Updated default parameters to match main branch (β=0.96, grid_max=16, y≈(0,2))
- Created get_endogenous_grid() helper function
- Updated all plots to use endogenous grid (a = c + s) rather than exogenous savings grid
- Removed intermediate ifp.py file (can be regenerated from ifp.md using jupytext)

Mathematical corrections ensure the lecture accurately represents the EGM algorithm
where the Euler equation is solved on the savings grid and policies are plotted on
the resulting endogenous asset grid.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

* Fix IFP lecture: Complete NumPy and JAX implementations with verification

Major fixes and improvements:

1. NumPy Implementation (K_numpy and solve_model_numpy):
   - Fixed type hints (np.ndarray instead of jnp.ndarray)
   - Fixed loop iteration bug (range(n_z) instead of n_z)
   - Fixed undefined variable references (savings_grid[i], z_grid[k])
   - Added missing γ parameter to u_prime() and u_prime_inv() calls
   - Fixed boundary condition assignment (new_c_vals instead of c_vals)
   - Fixed broadcasting in endogenous grid calculation
   - Fixed return values (c_vals, ae_vals instead of undefined σ)

2. JAX Implementation (K and solve_model):
   - Fully implemented K operator using jax.vmap for vectorization
   - Fixed solve_model body function variable names
   - Added JAX 64-bit precision configuration for numerical accuracy
   - Verified results match NumPy to machine precision (~10^-15)

3. Code Throughout Lecture:
   - Fixed all solve_model unpacking (returns tuple, not single value)
   - Fixed asset law of motion plot (interpolate on endogenous grid)
   - Fixed cake eating sanity check (removed non-existent function)
   - Fixed compute_asset_stationary to use correct interpolation grid
   - Fixed all exercise solutions (1, 2, 3)
   - Improved variable naming consistency (σ_init → c_vals_init, z → k, lb → label)

4. Added Verification Section:
   - Compares NumPy and JAX implementations
   - Shows numerical differences at machine precision
   - Validates correctness of JAX implementation

All code tested end-to-end and verified working correctly.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

* misc

* Refine IFP lecture: Improve clarity, notation, and code structure

Major improvements:

- **Documentation enhancements**: Renamed to "The Household Problem", introduced
  "income fluctuation problem" terminology, added comprehensive note explaining
  timing convention and how it differs from traditional formulation

- **Mathematical precision**: Changed "measurability" to "adaptedness", used :=
  for definitions, improved EGM algorithm description with explicit definitions

- **Code refactoring**: Renamed `savings_grid` → `s` for brevity, simplified
  vmap by eliminating lambda wrapper using `in_axes=(None, 0)`, improved variable
  naming (e.g., `mu` for marginal utility)

- **Structural improvements**: Split utility setup into separate cell, added
  "Dynamics" subsection, reorganized JAX verification, improved exercise specs

- **Terminology**: Changed "Coleman-Reffett operator" → "Euler equation operator"
  in docstrings for accuracy

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

* misc

---------

Co-authored-by: Claude <[email protected]>

Author: jstac