QuantEcon
diff --git a/‎lectures/os_egm.md‎
Lines changed: 29 additions & 33 deletions b/‎lectures/os_egm.md‎
Lines changed: 29 additions & 33 deletions
diff --git a/‎lectures/os_egm_jax.md‎
Lines changed: 3 additions & 8 deletions b/‎lectures/os_egm_jax.md‎
Lines changed: 3 additions & 8 deletions
diff --git a/‎lectures/os_stochastic.md‎
Lines changed: 15 additions & 22 deletions b/‎lectures/os_stochastic.md‎
Lines changed: 15 additions & 22 deletions
diff --git a/‎lectures/os_time_iter.md‎
Lines changed: 18 additions & 16 deletions b/‎lectures/os_time_iter.md‎
Lines changed: 18 additions & 16 deletions
@@ -60,16 +60,8 @@ First we remind ourselves of the theory and then we turn to numerical methods.
 
 We work with the model set out in {doc}`os_time_iter`, following the same terminology and notation.
 
-The Euler equation is
-
-```{math}
-:label: egm_euler
-
-(u'\circ \sigma^*)(x)
-= \beta \int (u'\circ \sigma^*)(f(x - \sigma^*(x)) z) f'(x - \sigma^*(x)) z \phi(dz)
-```
-
-As we saw, the Coleman-Reffett operator is a nonlinear operator $K$ engineered so that $\sigma^*$ is a fixed point of $K$.
+As we saw, the Coleman-Reffett operator is a nonlinear operator $K$ engineered so that the optimal policy 
+$\sigma^*$ is a fixed point of $K$.
 
 It takes as its argument a continuous strictly increasing consumption policy $\sigma \in \Sigma$.
 
@@ -85,9 +77,7 @@ u'(c)
 ### Exogenous Grid
 
 As discussed in {doc}`os_time_iter`, to implement the method on a
-computer, we need numerical approximation.
-
-In particular, we represent a policy function by a set of values on a finite grid.
+computer, we represent a policy function by a set of values on a finite grid.
 
 The function itself is reconstructed from this representation when necessary,
 using interpolation or some other method.
@@ -191,19 +181,21 @@ class Model(NamedTuple):
     u_prime_inv: Callable # inverse of u_prime
 
 
-def create_model(u: Callable,
-                 f: Callable,
-                 β: float = 0.96,
-                 μ: float = 0.0,
-                 ν: float = 0.1,
-                 grid_max: float = 4.0,
-                 grid_size: int = 120,
-                 shock_size: int = 250,
-                 seed: int = 1234,
-                 α: float = 0.4,
-                 u_prime: Callable = None,
-                 f_prime: Callable = None,
-                 u_prime_inv: Callable = None) -> Model:
+def create_model(
+        u: Callable,
+        f: Callable,
+        β: float = 0.96,
+        μ: float = 0.0,
+        ν: float = 0.1,
+        grid_max: float = 4.0,
+        grid_size: int = 120,
+        shock_size: int = 250,
+        seed: int = 1234,
+        α: float = 0.4,
+        u_prime: Callable = None,
+        f_prime: Callable = None,
+        u_prime_inv: Callable = None
+    ) -> Model:
     """
     Creates an instance of the optimal savings model.
     """
@@ -214,7 +206,9 @@ def create_model(u: Callable,
     np.random.seed(seed)
     shocks = np.exp(μ + ν * np.random.randn(shock_size))
 
-    return Model(u, f, β, μ, ν, s_grid, shocks, α, u_prime, f_prime, u_prime_inv)
+    return Model(
+        u, f, β, μ, ν, s_grid, shocks, α, u_prime, f_prime, u_prime_inv
+    )
 ```
 
 ### The Operator
@@ -282,12 +276,14 @@ s_grid = model.s_grid
 Here's our solver routine:
 
 ```{code-cell} python3
-def solve_model_time_iter(model: Model,
-                          c_init: np.ndarray,
-                          x_init: np.ndarray,
-                          tol: float = 1e-5,
-                          max_iter: int = 1000,
-                          verbose: bool = True):
+def solve_model_time_iter(
+        model: Model,             # Model details
+        c_init: np.ndarray,       # initial guess of consumption on EG
+        x_init: np.ndarray,       # initial guess of endogenous grid
+        tol: float = 1e-5,        # Error tolerance
+        max_iter: int = 1000,     # Max number of iterations of K
+        verbose: bool = True      # If true print output
+    ):
     """
     Solve the model using time iteration with EGM.
     """
 
@@ -39,7 +39,7 @@ We'll also use JAX's `vmap` function to fully vectorize the Coleman-Reffett oper
 
 Let's start with some standard imports:
 
-```{code-cell} ipython
+```{code-cell} python3
 import matplotlib.pyplot as plt
 import jax
 import jax.numpy as jnp
@@ -113,7 +113,7 @@ def create_model(
     key = jax.random.PRNGKey(seed)
     shocks = jnp.exp(μ + s * jax.random.normal(key, shape=(shock_size,)))
 
-    return Model(β=β, μ=μ, s=s, s_grid=s_grid, shocks=shocks, α=α)
+    return Model(β, μ, s, s_grid, shocks, α)
 ```
 
 
@@ -141,12 +141,7 @@ def K(
     The Coleman-Reffett operator using EGM
 
     """
-
-    # Simplify names
-    β, α = model.β, model.α
-    s_grid, shocks = model.s_grid, model.shocks
-
-    # Linear interpolation of policy using endogenous grid
+    β, μ, s, s_grid, shocks, α = model
     σ = lambda x_val: jnp.interp(x_val, x_in, c_in)
 
     # Define function to compute consumption at a single grid point
 
@@ -62,7 +62,7 @@ More information on this savings problem can be found in
 
 Let's start with some imports:
 
-```{code-cell} ipython
+```{code-cell} python3
 import matplotlib.pyplot as plt
 import numpy as np
 from scipy.interpolate import interp1d
@@ -270,7 +270,7 @@ The term $\int v(f(x - c) z) \phi(dz)$ can be understood as the expected next pe
 * the state is $x$
 * consumption is set to $c$
 
-As shown in [DP1](https://dp.quantecon.org/), Theorem 10.1.11 and a range of other texts,
+As shown in [DP1](https://dp.quantecon.org/) and a range of other texts,
 the value function $v^*$ satisfies the Bellman equation.
 
 In other words, {eq}`fpb30` holds when $v=v^*$.
@@ -313,10 +313,9 @@ function.
 
 In our setting, we have the following key result
 
-* A feasible consumption policy is optimal if and only if it is $v^*$-greedy.
-
-The intuition is similar to the intuition for the Bellman equation, which was
-provided after {eq}`fpb30`.
+```{prf:theorem}
+A feasible consumption policy is optimal if and only if it is $v^*$-greedy.
+```
 
 See, for example, Theorem 10.1.11 of [EDTC](https://johnstachurski.net/edtc.html).
 
@@ -359,7 +358,7 @@ v(x)
 = Tv(x)
 = \max_{0 \leq c \leq x}
 \left\{
-    u(c) + \beta \int v^*(f(x - c) z) \phi(dz)
+    u(c) + \beta \int v(f(x - c) z) \phi(dz)
 \right\}
 $$
 
@@ -515,20 +514,18 @@ def create_model(
 We set up the right-hand side of the Bellman equation
 
 $$
-    B(x, c, v) := u(c) + \beta \int v^*(f(x - c) z) \phi(dz)
+    B(x, c, v) := u(c) + \beta \int v(f(x - c) z) \phi(dz)
 $$
 
 
 ```{code-cell} python3
 def B(
-        x: float,
-        c: float,
-        v_array: np.ndarray,
-        model: Model
-    ) -> float:
-    """
-    Right hand side of the Bellman equation.
-    """
+        x: float,              # State
+        c: float,              # Action
+        v_array: np.ndarray,   # Array representing a guess of the value fn
+        model: Model           # An instance of Model containing parameters
+    ):
+
     u, f, β, μ, ν, x_grid, shocks = model
     v = interp1d(x_grid, v_array)
 
@@ -569,7 +566,7 @@ def T(v: np.ndarray, model: Model) -> tuple[np.ndarray, np.ndarray]:
 
     for i in range(len(x_grid)):
         x = x_grid[i]
-        c_star, v_max = maximize(lambda c: B(x, c, v, model), x)
+        _, v_max = maximize(lambda c: B(x, c, v, model), x)
         v_new[i] = v_max
 
     return v_new
@@ -731,12 +728,8 @@ def solve_model(
         verbose: bool = True,   # print iteration info
         print_skip: int = 25    # iterations between prints
     ):
-    """
-    Solve model by iterating with the Bellman operator.
-
-    """
+    " Solve by value function iteration. "
 
-    # Set up loop
     v = model.u(model.x_grid)  # Initial condition
     i = 0
     error = tol + 1
 
@@ -84,22 +84,23 @@ Recall the Bellman equation
 ```{math}
 :label: cpi_fpb30
 
-v^*(x) = \max_{0 \leq c \leq x}
+v(x) = \max_{0 \leq c \leq x}
     \left\{
-        u(c) + \beta \int v^*(f(x - c) z) \phi(dz)
+        u(c) + \beta \int v(f(x - c) z) \phi(dz)
     \right\}
 \quad \text{for all} \quad
 x \in \mathbb R_+
 ```
 
-Let the optimal consumption policy be denoted by $\sigma^*$.
+Let $v^*$ be the value function and let $\sigma^*$ be the optimal consumption policy.
 
-We know that $\sigma^*$ is a $v^*$-greedy policy so that $\sigma^*(x)$ is the maximizer in {eq}`cpi_fpb30`.
+We know that $\sigma^*$ is a $v^*$-greedy policy.
 
 The conditions above imply that
 
 * $\sigma^*$ is the unique optimal policy for the optimal savings problem
-* the optimal policy is continuous, strictly increasing and also **interior**, in the sense that $0 < \sigma^*(x) < x$ for all strictly positive $x$, and
+* the optimal policy is continuous, strictly increasing and also **interior**,
+  in the sense that $0 < \sigma^*(x) < x$ for all strictly positive $x$, and
 * the value function is strictly concave and continuously differentiable, with
 
 ```{math}
@@ -108,7 +109,8 @@ The conditions above imply that
 (v^*)'(x) = u' (\sigma^*(x) ) := (u' \circ \sigma^*)(x)
 ```
 
-The last result is called the **envelope condition** due to its relationship with the [envelope theorem](https://en.wikipedia.org/wiki/Envelope_theorem).
+The last result is called the **envelope condition** due to its relationship
+with the [envelope theorem](https://en.wikipedia.org/wiki/Envelope_theorem).
 
 To see why {eq}`cpi_env` holds, write the Bellman equation in the equivalent
 form
@@ -278,7 +280,7 @@ As in {doc}`os_stochastic`, we assume that
 This allows us to compare our results to the analytical solutions we obtained in
 that lecture:
 
-```{code-cell} python3
+```{code-cell} ipython
 def v_star(x, α, β, μ):
     """
     True value function
@@ -303,7 +305,7 @@ For this we need access to the functions $u'$ and $f, f'$.
 
 We use the same `Model` structure from {doc}`os_stochastic`.
 
-```{code-cell} python3
+```{code-cell} ipython
 class Model(NamedTuple):
     u: Callable        # utility function
     f: Callable        # production function
@@ -379,7 +381,7 @@ state $x$ and $σ$, the current guess of the policy.
 
 Here's the operator $K$, that implements the root-finding step.
 
-```{code-cell} ipython3
+```{code-cell} ipython
 def K(σ: np.ndarray, model: Model) -> np.ndarray:
     """
     The Coleman-Reffett operator
@@ -402,7 +404,7 @@ def K(σ: np.ndarray, model: Model) -> np.ndarray:
 
 Let's generate an instance and plot some iterates of $K$, starting from $σ(x) = x$.
 
-```{code-cell} python3
+```{code-cell} ipython
 # Define utility and production functions with derivatives
 α = 0.4
 u = lambda c: np.log(c)
@@ -441,7 +443,7 @@ Here is a function called `solve_model_time_iter` that takes an instance of
 using time iteration.
 
 
-```{code-cell} python3
+```{code-cell} ipython
 def solve_model_time_iter(
         model: Model,
         σ_init: np.ndarray,
@@ -473,7 +475,7 @@ def solve_model_time_iter(
 
 Let's call it:
 
-```{code-cell} python3
+```{code-cell} ipython
 # Unpack
 grid = model.grid
 
@@ -483,7 +485,7 @@ grid = model.grid
 
 Here is a plot of the resulting policy, compared with the true policy:
 
-```{code-cell} python3
+```{code-cell} ipython
 # Unpack
 grid, α, β = model.grid, model.α, model.β
 
@@ -503,7 +505,7 @@ Again, the fit is excellent.
 
 The maximal absolute deviation between the two policies is
 
-```{code-cell} python3
+```{code-cell} ipython
 # Unpack
 grid, α, β = model.grid, model.α, model.β
 
@@ -540,7 +542,7 @@ Compute and plot the optimal policy.
 
 We define the CRRA utility function and its derivative.
 
-```{code-cell} python3
+```{code-cell} ipython
 γ = 1.5
 
 def u_crra(c):
@@ -556,7 +558,7 @@ model_crra = create_model(u=u_crra, f=f, α=α,
 
 Now we solve and plot the policy:
 
-```{code-cell} python3
+```{code-cell} ipython
 %%time
 # Unpack
 grid = model_crra.grid