Skip to content

Conversation

@jstac
Copy link
Contributor

@jstac jstac commented Nov 2, 2025

Summary

This PR refines the Lake Model lecture with several improvements to code quality, performance, and pedagogical flow:

Code improvements:

  • Remove redundant @jax.jit decorators from helper functions (compute_matrices, stock_update, rate_update, markov_update)
    • These functions are only called from within jitted functions, making the decorators unnecessary
    • Removing them allows JAX's compiler to better optimize the entire computation graph
    • Prevents compilation boundaries that can inhibit operation fusion and other optimizations
  • Refactor aggregate dynamics plotting code to use a for loop instead of repetitive statements
  • Remove hardcoded color specifications to allow matplotlib's default color cycling

Content organization:

  • Move rate definitions ($e_t$, $u_t$) from the introduction to the "Laws of motion for rates" section where they are more relevant
  • Relocate Exercise 1 to appear immediately before "Dynamics of an individual worker" section for improved narrative flow
  • Simplify Exercise 1 to focus on demonstrating JAX's vmap functionality, removing less pedagogically interesting parameter comparisons

Test plan

  • Converted lecture to Python script using jupytext
  • Executed all code successfully - all outputs match expected results
  • Verified all functionality is preserved
  • Code is cleaner and better organized

🤖 Generated with Claude Code

…ation

This commit improves the Lake Model lecture with several refinements:

**Code improvements:**
- Remove redundant `@jax.jit` decorators from `compute_matrices`, `stock_update`, `rate_update`, and `markov_update` (these functions are only called from within jitted functions, so the decorators are unnecessary and can inhibit compiler optimization)
- Refactor aggregate dynamics plot to use a for loop instead of repetitive code
- Remove hardcoded colors ('r') from plots to use matplotlib's default color cycle

**Content organization:**
- Move rate definitions ($e_t$, $u_t$) to "Laws of motion for rates" section where they logically belong
- Relocate Exercise 1 to appear immediately before "Dynamics of an individual worker" section for better flow
- Simplify Exercise 1 to focus on the pedagogically interesting `vmap` usage, removing less interesting parameter comparison parts

These changes improve code clarity, performance, and pedagogical flow without changing functionality.

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

Co-Authored-By: Claude <noreply@anthropic.com>
@github-actions
Copy link

github-actions bot commented Nov 2, 2025

📖 Netlify Preview Ready!

Preview URL: https://pr-664--sunny-cactus-210e3e.netlify.app (27bb29b)

📚 Changed Lecture Pages: lake_model

jstac and others added 2 commits November 2, 2025 17:40
Major improvements:
- Added create_lake_model() function to generate model instances with precomputed matrices A, R, and g
- Replaced A_hat notation with R throughout (code and LaTeX) for cleaner notation
- Updated LakeModel NamedTuple to store computed matrices A and R
- Modified all functions to unpack model using tuple unpacking for efficiency
- Added type hints to stock_update(), rate_update(), and create_lake_model()
- Simplified generate_path() function by removing unused time parameter
- Updated rate_steady_state() to use Perron-Frobenius theorem (argmax instead of searching for eigenvalue near 1)
- Converted all LakeModel() instantiations to use create_lake_model()
- Updated markov simulation to use dedicated simulate_markov() function

Benefits:
- Matrices computed once at model creation instead of repeatedly
- Cleaner mathematical notation using R instead of \hat{A}
- More efficient code with direct tuple unpacking
- Better type safety with added annotations
- More mathematically rigorous using Perron-Frobenius theorem

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

Co-Authored-By: Claude <noreply@anthropic.com>
@github-actions
Copy link

github-actions bot commented Nov 2, 2025

📖 Netlify Preview Ready!

Preview URL: https://pr-664--sunny-cactus-210e3e.netlify.app (8a84067)

📚 Changed Lecture Pages: lake_model

@jstac jstac merged commit 1c98e34 into main Nov 2, 2025
1 check passed
@jstac jstac deleted the lake-misc-edits branch November 2, 2025 09:36
@jstac
Copy link
Contributor Author

jstac commented Nov 2, 2025

Additional Refactoring: Improved Model Structure and Notation

I've made additional improvements to lake_model.md:

Key Changes:

  1. Added create_lake_model() function: Generates model instances with precomputed matrices A, R, and g stored within the model

    • Matrices are computed once at creation instead of repeatedly
    • More efficient as functions now just unpack and use precomputed values
  2. Replaced A_hat notation with R throughout (both code and LaTeX):

    • Cleaner, more concise notation
    • Updated all equations: x_{t+1} = R x_t instead of x_{t+1} = \hat{A} x_t
  3. Improved function efficiency:

    • All functions now unpack model using tuple unpacking: λ, α, b, d, A, R, g = model
    • Removed compute_matrices() function (no longer needed)
    • Added type hints to stock_update(), rate_update(), and create_lake_model()
  4. Enhanced steady state computation:

    • Updated rate_steady_state() to use Perron-Frobenius theorem explicitly
    • Uses argmax to find largest eigenvalue instead of searching for eigenvalue near 1
    • More mathematically rigorous with improved docstring explaining the theorem
  5. Simplified generate_path() function:

    • Removed unused time parameter for cleaner interface
    • Updated Markov simulation with dedicated simulate_markov() helper

All code has been tested and verified to work correctly. ✅

@github-actions
Copy link

github-actions bot commented Nov 2, 2025

📖 Netlify Preview Ready!

Preview URL: https://pr-664--sunny-cactus-210e3e.netlify.app (8cce1a7)

📚 Changed Lecture Pages: lake_model

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants