Refactor cake_eating_egm lecture: update K operator signature and improve clarity

[jstac]

Summary

This PR refactors the cake_eating_egm lecture to improve code clarity and fix several issues:

• Updated K operator signature: Changed K to take (c_in, x_in) parameters for interpolation and return (c_out, x_out), making the endogenous grid method more explicit
• Renamed grid to s_grid: Changed all references from grid to s_grid to clearly indicate this is the exogenous savings grid
• Changed shock parameter: Renamed shock scale parameter from s to ν (nu) to avoid confusion with savings-related variables
• Updated calling code: Modified solve_model_time_iter and all calling code to work with the new K signature
• Fixed grammar issues: Corrected verb forms and added missing punctuation
• Added performance note: Included a note about Python loops and reference to the JAX lecture for vectorized solution

Test plan

• Converted to Python using jupytext (no warnings)
• Ran the script successfully
• Verified convergence in 14 iterations with error < 1e-05
• All code executes without errors

🤖 Generated with Claude Code

[jstac]

Detailed Changes

1. K Operator Signature Update

Before:

def K(σ_array: np.ndarray, model: Model) -> np.ndarray:
    # Determine endogenous grid
    x = grid + σ_array
    # ... computation ...
    return c

After:

def K(
        c_in: np.ndarray,   # Consumption values on the endogenous grid
        x_in: np.ndarray,   # Current endogenous grid
        model: Model        # Model specification
    ):
    # ... computation ...
    return c_out, x_out

This makes it explicit that K takes consumption and endogenous grid values as inputs and returns updated values for both.

2. Grid Renaming (grid → s_grid)

All references to grid have been renamed to s_grid to emphasize that this is the exogenous savings grid, matching the mathematical notation where we fix grid ${s_i}$ for savings.

Changes:

• Model.grid → Model.s_grid with comment "exogenous savings grid"
• create_model: Comment updated to "Set up exogenous savings grid"
• All code using grid now uses s_grid

3. Shock Scale Parameter (s → ν)

Changed shock scale parameter from s to ν (Greek nu) to avoid confusion with savings.

Changes:

• Model.s → Model.ν
• create_model(s=0.1) → create_model(ν=0.1)
• shocks = np.exp(μ + s * ...) → shocks = np.exp(μ + ν * ...)

4. Updated Calling Code

solve_model_time_iter:

• Now takes (c_init, x_init) instead of just σ_init
• Calls c_new, x_new = K(c, x, model) and unpacks both outputs
• Returns (c, x) tuple

Initialization:

c_init = np.copy(s_grid)
x_init = s_grid + c_init
c, x = solve_model_time_iter(model, c_init, x_init)

Plotting:

• Removed x = grid + σ calculation (now returned from solver)
• Changed all σ references to c

5. Grammar and Style Fixes

• Line 92-93: "fixed/calculated" → "fix/calculate" (correct infinitive after "was to")
• Line 150: Added missing period
• Line 139: Removed extra space in "is determined"
• Line 221: Expanded "EG" → "endogenous grid"
• Line 44: Changed {code-cell} ipython → {code-cell} python3 for consistency

6. Performance Note

Added a note after the K function implementation:

Note the lack of any root-finding algorithm.

​```{note}
The routine is still not particularly fast because we are using pure Python loops.

But in the next lecture ({doc}`cake_eating_egm_jax`) we will use a fully vectorized and efficient solution.
​```

---

All changes have been tested by converting to Python with jupytext and running successfully.

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-730--sunny-cactus-210e3e.netlify.app (080c21e)

📚 Changed Lecture Pages: cake_eating_egm

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-730--sunny-cactus-210e3e.netlify.app (a96ac9d)

📚 Changed Lecture Pages: cake_eating_egm