From ee869e251dc25d9136ca85ed584a5e84a86f206b Mon Sep 17 00:00:00 2001
From: kp992 <kpl992@outlook.com>
Date: Mon, 18 Aug 2025 20:09:08 -0700
Subject: [PATCH 1/6] update code

---
 lectures/newton_method.md | 194 +++++++++++++++++++-------------------
 1 file changed, 95 insertions(+), 99 deletions(-)

diff --git a/lectures/newton_method.md b/lectures/newton_method.md
index 4f2788cf1..a8730f632 100644
--- a/lectures/newton_method.md
+++ b/lectures/newton_method.md
@@ -4,7 +4,7 @@ jupytext:
     extension: .md
     format_name: myst
     format_version: 0.13
-    jupytext_version: 1.14.4
+    jupytext_version: 1.16.7
 kernelspec:
   display_name: Python 3 (ipykernel)
   language: python
@@ -85,16 +85,15 @@ We use the following imports in this lecture
 
 ```{code-cell} ipython3
 import matplotlib.pyplot as plt
-from collections import namedtuple
+from typing import NamedTuple
 from scipy.optimize import root
 from autograd import jacobian
+
 # Thinly-wrapped numpy to enable automatic differentiation
 import autograd.numpy as np
-
-plt.rcParams["figure.figsize"] = (10, 5.7)
 ```
 
-## Fixed Point Computation Using Newton's Method
+## Fixed point computation using Newton's method
 
 In this section we solve the fixed point of the law of motion for capital in
 the setting of the [Solow growth
@@ -104,7 +103,7 @@ We will inspect the fixed point visually, solve it by successive
 approximation, and then apply Newton's method to achieve faster convergence.
 
 (solow)=
-### The Solow Model
+### The Solow model
 
 In the Solow growth model, assuming Cobb-Douglas production technology and
 zero population growth, the law of motion for capital is
@@ -136,17 +135,21 @@ $$ k^* = \left(\frac{s A}{δ}\right)^{1/(1 - α)}  $$
 
 ### Implementation
 
-Let's store our parameters in [`namedtuple`](https://docs.python.org/3/library/collections.html#collections.namedtuple) to help us keep our code clean and concise.
+Let's store our parameters in [`NamedTuple`](https://typing.python.org/en/latest/spec/namedtuples.html) to help us keep our code clean and concise.
 
 ```{code-cell} ipython3
-SolowParameters = namedtuple("SolowParameters", ('A', 's', 'α', 'δ'))
+class SolowParameters(NamedTuple):
+    A: float
+    s: float
+    α: float
+    δ: float
 ```
 
-This function creates a suitable `namedtuple` with default parameter values.
+This function creates a suitable `SolowParameters` with default parameter values.
 
 ```{code-cell} ipython3
 def create_solow_params(A=2.0, s=0.3, α=0.3, δ=0.4):
-    "Creates a Solow model parameterization with default values."
+    """Creates a Solow model parameterization with default values."""
     return SolowParameters(A=A, s=s, α=α, δ=δ)
 ```
 
@@ -156,35 +159,38 @@ The next two functions implement the law of motion [](motion_law) and store the
 def g(k, params):
     A, s, α, δ = params
     return A * s * k**α + (1 - δ) * k
-    
+
+
 def exact_fixed_point(params):
     A, s, α, δ = params
-    return ((s * A) / δ)**(1/(1 - α))
+    return ((s * A) / δ) ** (1 / (1 - α))
 ```
 
 Here is a function to provide a 45 degree plot of the dynamics.
 
 ```{code-cell} ipython3
 def plot_45(params, ax, fontsize=14):
-    
+
     k_min, k_max = 0.0, 3.0
     k_grid = np.linspace(k_min, k_max, 1200)
 
     # Plot the functions
     lb = r"$g(k) = sAk^{\alpha} + (1 - \delta)k$"
-    ax.plot(k_grid, g(k_grid, params),  lw=2, alpha=0.6, label=lb)
+    ax.plot(k_grid, g(k_grid, params), lw=2, alpha=0.6, label=lb)
     ax.plot(k_grid, k_grid, "k--", lw=1, alpha=0.7, label="45")
 
     # Show and annotate the fixed point
     kstar = exact_fixed_point(params)
     fps = (kstar,)
     ax.plot(fps, fps, "go", ms=10, alpha=0.6)
-    ax.annotate(r"$k^* = (sA / \delta)^{\frac{1}{1-\alpha}}$", 
-             xy=(kstar, kstar),
-             xycoords="data",
-             xytext=(20, -20),
-             textcoords="offset points",
-             fontsize=fontsize)
+    ax.annotate(
+        r"$k^* = (sA / \delta)^{\frac{1}{1-\alpha}}$",
+        xy=(kstar, kstar),
+        xycoords="data",
+        xytext=(20, -20),
+        textcoords="offset points",
+        fontsize=fontsize,
+    )
 
     ax.legend(loc="upper left", frameon=False, fontsize=fontsize)
 
@@ -214,7 +220,7 @@ plt.show()
 We see that $k^*$ is indeed the unique positive fixed point.
 
 
-#### Successive Approximation
+#### Successive approximation
 
 First let's compute the fixed point using successive approximation.
 
@@ -225,7 +231,7 @@ Here's a time series from a particular choice of $k_0$.
 
 ```{code-cell} ipython3
 def compute_iterates(k_0, f, params, n=25):
-    "Compute time series of length n generated by arbitrary function f."
+    """Compute time series of length n generated by arbitrary function f."""
     k = k_0
     k_iterates = []
     for t in range(n):
@@ -241,8 +247,8 @@ k_series = compute_iterates(k_0, g, params)
 k_star = exact_fixed_point(params)
 
 fig, ax = plt.subplots()
-ax.plot(k_series, 'o')
-ax.plot([k_star] * len(k_series), 'k--')
+ax.plot(k_series, "o")
+ax.plot([k_star] * len(k_series), "k--")
 ax.set_ylim(0, 3)
 plt.show()
 ```
@@ -263,7 +269,7 @@ This is close to the true value.
 k_star
 ```
 
-#### Newton's Method 
+#### Newton's method 
 
 In general, when applying Newton's fixed point method to some function $g$, 
 we start with a guess $x_0$ of the fixed
@@ -310,7 +316,7 @@ Let's define this:
 ```{code-cell} ipython3
 def Dg(k, params):
     A, s, α, δ = params
-    return α * A * s * k**(α-1) + (1 - δ)
+    return α * A * s * k ** (α - 1) + (1 - δ)
 ```
 
 Here's a function $q$ representing [](newtons_method).
@@ -323,11 +329,13 @@ def q(k, params):
 Now let's plot some trajectories.
 
 ```{code-cell} ipython3
-def plot_trajectories(params, 
-                      k0_a=0.8,  # first initial condition
-                      k0_b=3.1,  # second initial condition
-                      n=20,      # length of time series
-                      fs=14):    # fontsize
+def plot_trajectories(
+    params,
+    k0_a=0.8,  # first initial condition
+    k0_b=3.1,  # second initial condition
+    n=20,  # length of time series
+    fs=14,
+):  # fontsize
 
     fig, axes = plt.subplots(2, 1, figsize=(10, 6))
     ax1, ax2 = axes
@@ -351,7 +359,7 @@ def plot_trajectories(params,
         ax.set_yticks((k_star,))
         ax.set_yticklabels(("$k^*$",), fontsize=fs)
         ax.set_xticks(np.linspace(0, 19, 20))
-        
+
     plt.show()
 ```
 
@@ -363,7 +371,7 @@ plot_trajectories(params)
 We can see that Newton's method converges faster than successive approximation.
 
 
-## Root-Finding in One Dimension
+## Root-Finding in one dimension
 
 In the previous section we computed fixed points.
 
@@ -375,7 +383,7 @@ the problem of finding fixed points.
 
 
 
-### Newton's Method for Zeros
+### Newton's method for zeros
 
 Let's suppose we want to find an $x$ such that $f(x)=0$ for some smooth
 function $f$ mapping real numbers to real numbers.
@@ -421,12 +429,12 @@ def newton(f, Df, x_0, tol=1e-7, max_iter=100_000):
     n = 0
     while error > tol:
         n += 1
-        if(n > max_iter):
-            raise Exception('Max iteration reached without convergence')
+        if n > max_iter:
+            raise Exception("Max iteration reached without convergence")
         y = q(x)
         error = np.abs(x - y)
         x = y
-        print(f'iteration {n}, error = {error:.5f}')
+        print(f"iteration {n}, error = {error:.5f}")
     return x
 ```
 
@@ -438,7 +446,7 @@ automatic differentiation or GPU acceleration, it will be helpful to know how
 to implement Newton's method ourselves.)
 
 
-### Application to Finding Fixed Points
+### Application to finding fixed points
 
 Now consider again the Solow fixed-point calculation, where we solve for $k$
 satisfying $g(k) = k$.
@@ -451,9 +459,9 @@ Let's apply this idea to the Solow problem
 
 ```{code-cell} ipython3
 params = create_solow_params()
-k_star_approx_newton = newton(f=lambda x: g(x, params) - x,
-                              Df=lambda x: Dg(x, params) - 1,
-                              x_0=0.8)
+k_star_approx_newton = newton(
+    f=lambda x: g(x, params) - x, Df=lambda x: Dg(x, params) - 1, x_0=0.8
+)
 ```
 
 ```{code-cell} ipython3
@@ -464,7 +472,7 @@ The result confirms the descent we saw in the graphs above: a very accurate resu
 
 
 
-## Multivariate Newton’s Method
+## Multivariate Newton’s method
 
 In this section, we introduce a two-good problem, present a
 visualization of the problem, and solve for the equilibrium of the two-good market
@@ -477,7 +485,7 @@ We will see a significant performance gain when using Netwon's method.
 
 
 (two_goods_market)=
-### A Two Goods Market Equilibrium
+### A Two Goods market equilibrium
 
 Let's start by computing the market equilibrium of a two-good problem.
 
@@ -531,7 +539,7 @@ $$
 
 for this particular question.
 
-#### A Graphical Exploration
+#### A graphical exploration
 
 Since our problem is only two-dimensional, we can use graphical analysis to visualize and help understand the problem.
 
@@ -549,7 +557,7 @@ The function below calculates the excess demand for given parameters
 
 ```{code-cell} ipython3
 def e(p, A, b, c):
-    return np.exp(- A @ p) + c - b * np.sqrt(p)
+    return np.exp(-A @ p) + c - b * np.sqrt(p)
 ```
 
 Our default parameter values will be
@@ -573,10 +581,7 @@ A = \begin{bmatrix}
 $$
 
 ```{code-cell} ipython3
-A = np.array([
-    [0.5, 0.4],
-    [0.8, 0.2]
-])
+A = np.array([[0.5, 0.4], [0.8, 0.2]])
 b = np.ones(2)
 c = np.ones(2)
 ```
@@ -586,8 +591,10 @@ At a price level of $p = (1, 0.5)$, the excess demand is
 ```{code-cell} ipython3
 ex_demand = e((1.0, 0.5), A, b, c)
 
-print(f'The excess demand for good 0 is {ex_demand[0]:.3f} \n'
-      f'The excess demand for good 1 is {ex_demand[1]:.3f}')
+print(
+    f"The excess demand for good 0 is {ex_demand[0]:.3f} \n"
+    f"The excess demand for good 1 is {ex_demand[1]:.3f}"
+)
 ```
 
 Next we plot the two functions $e_0$ and $e_1$ on a grid of $(p_0, p_1)$ values, using contour surfaces and lines.
@@ -612,7 +619,7 @@ def plot_excess_demand(ax, good=0, grid_size=100, grid_max=4, surface=True):
     ctr1 = ax.contour(p_grid, p_grid, z.T, levels=[0.0])
     ax.set_xlabel("$p_0$")
     ax.set_ylabel("$p_1$")
-    ax.set_title(f'Excess Demand for Good {good}')
+    ax.set_title(f"Excess Demand for Good {good}")
     plt.clabel(ctr1, inline=1, fontsize=13)
 ```
 
@@ -648,7 +655,7 @@ plt.show()
 It seems there is an equilibrium close to $p = (1.6, 1.5)$.
 
 
-#### Using a Multidimensional Root Finder
+#### Using a multidimensional root finder
 
 To solve for $p^*$ more precisely, we use a zero-finding algorithm from `scipy.optimize`.
 
@@ -662,7 +669,7 @@ This uses the [modified Powell method](https://docs.scipy.org/doc/scipy/referenc
 
 ```{code-cell} ipython3
 %%time
-solution = root(lambda p: e(p, A, b, c), init_p, method='hybr')
+solution = root(lambda p: e(p, A, b, c), init_p, method="hybr")
 ```
 
 Here's the resulting value:
@@ -675,13 +682,14 @@ p
 This looks close to our guess from observing the figure. We can plug it back into $e$ to test that $e(p) \approx 0$:
 
 ```{code-cell} ipython3
-np.max(np.abs(e(p, A, b, c)))
+e_p = np.max(np.abs(e(p, A, b, c)))
+e_p.item()
 ```
 
 This is indeed a very small error.
 
 
-#### Adding Gradient Information
+#### Adding gradient information
 
 In many cases, for zero-finding algorithms applied to smooth functions, supplying the [Jacobian](https://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant) of the function leads to better convergence properties.
 
@@ -700,31 +708,30 @@ def jacobian_e(p, A, b, c):
     p_0, p_1 = p
     a_00, a_01 = A[0, :]
     a_10, a_11 = A[1, :]
-    j_00 = -a_00 * np.exp(-a_00 * p_0) - (b[0]/2) * p_0**(-1/2)
+    j_00 = -a_00 * np.exp(-a_00 * p_0) - (b[0] / 2) * p_0 ** (-1 / 2)
     j_01 = -a_01 * np.exp(-a_01 * p_1)
     j_10 = -a_10 * np.exp(-a_10 * p_0)
-    j_11 = -a_11 * np.exp(-a_11 * p_1) - (b[1]/2) * p_1**(-1/2)
-    J = [[j_00, j_01],
-         [j_10, j_11]]
+    j_11 = -a_11 * np.exp(-a_11 * p_1) - (b[1] / 2) * p_1 ** (-1 / 2)
+    J = [[j_00, j_01], [j_10, j_11]]
     return np.array(J)
 ```
 
 ```{code-cell} ipython3
 %%time
-solution = root(lambda p: e(p, A, b, c),
-                init_p, 
-                jac=lambda p: jacobian_e(p, A, b, c), 
-                method='hybr')
+solution = root(
+    lambda p: e(p, A, b, c), init_p, jac=lambda p: jacobian_e(p, A, b, c), method="hybr"
+)
 ```
 
 Now the solution is even more accurate (although, in this low-dimensional problem, the difference is quite small):
 
 ```{code-cell} ipython3
 p = solution.x
-np.max(np.abs(e(p, A, b, c)))
+e_p = np.max(np.abs(e(p, A, b, c)))
+e_p.item()
 ```
 
-#### Using Newton's Method
+#### Using Newton's method
 
 Now let's use Newton's method to compute the equilibrium price using the multivariate version of Newton's method
 
@@ -751,22 +758,22 @@ def newton(f, x_0, tol=1e-5, max_iter=10):
     error = tol + 1
     n = 0
     while error > tol:
-        n+=1
-        if(n > max_iter):
-            raise Exception('Max iteration reached without convergence')
+        n += 1
+        if n > max_iter:
+            raise Exception("Max iteration reached without convergence")
         y = q(x)
-        if(any(np.isnan(y))):
-            raise Exception('Solution not found with NaN generated')
+        if any(np.isnan(y)):
+            raise Exception("Solution not found with NaN generated")
         error = np.linalg.norm(x - y)
         x = y
-        print(f'iteration {n}, error = {error:.5f}')
-    print('\n' + f'Result = {x} \n')
+        print(f"iteration {n}, error = {error:.5f}")
+    print("\n" + f"Result = {x} \n")
     return x
 ```
 
 ```{code-cell} ipython3
 def e(p, A, b, c):
-    return np.exp(- np.dot(A, p)) + c - b * np.sqrt(p)
+    return np.exp(-np.dot(A, p)) + c - b * np.sqrt(p)
 ```
 
 We find the algorithm terminates in 4 steps
@@ -777,7 +784,8 @@ p = newton(lambda p: e(p, A, b, c), init_p)
 ```
 
 ```{code-cell} ipython3
-np.max(np.abs(e(p, A, b, c)))
+e_p = np.max(np.abs(e(p, A, b, c)))
+e_p.item()
 ```
 
 The result is very accurate. 
@@ -785,7 +793,7 @@ The result is very accurate.
 With the larger overhead, the speed is not better than the optimized `scipy` function.
 
 
-### A High-Dimensional Problem
+### A high-dimensional problem
 
 Our next step is to investigate a large market with 3,000 goods.
 
@@ -821,7 +829,8 @@ p = newton(lambda p: e(p, A, b, c), init_p)
 ```
 
 ```{code-cell} ipython3
-np.max(np.abs(e(p, A, b, c)))
+e_p = np.max(np.abs(e(p, A, b, c)))
+e_p.item()
 ```
 
 With the same tolerance, we compare the runtime and accuracy of Newton's method to SciPy's `root` function
@@ -837,10 +846,10 @@ solution = root(lambda p: e(p, A, b, c),
 
 ```{code-cell} ipython3
 p = solution.x
-np.max(np.abs(e(p, A, b, c)))
+e_p = np.max(np.abs(e(p, A, b, c)))
+e_p.item()
 ```
 
-
 ## Exercises
 
 ```{exercise-start}
@@ -914,24 +923,20 @@ The result should converge to the [analytical solution](solved_k).
 Let's first define the parameters for this problem
 
 ```{code-cell} ipython3
-A = np.array([[2.0, 3.0, 3.0],
-              [2.0, 4.0, 2.0],
-              [1.0, 5.0, 1.0]])
+A = np.array([[2.0, 3.0, 3.0], [2.0, 4.0, 2.0], [1.0, 5.0, 1.0]])
 
 s = 0.2
 α = 0.5
 δ = 0.8
 
-initLs = [np.ones(3),
-          np.array([3.0, 5.0, 5.0]),
-          np.repeat(50.0, 3)]
+initLs = [np.ones(3), np.array([3.0, 5.0, 5.0]), np.repeat(50.0, 3)]
 ```
 
 Then define the multivariate version of the formula for the [law of motion of captial](motion_law)
 
 ```{code-cell} ipython3
 def multivariate_solow(k, A=A, s=s, α=α, δ=δ):
-    return (s * np.dot(A, k**α) + (1 - δ) * k)
+    return s * np.dot(A, k**α) + (1 - δ) * k
 ```
 
 Let's run through each starting value and see the output
@@ -978,7 +983,7 @@ init = np.repeat(1.0, 3)
 
 The result is very close to the ground truth but still slightly different.
 
-```{code-cell} python3
+```{code-cell} ipython3
 %time k = newton(lambda k: multivariate_solow(k, A=A, s=s, α=α, δ=δ) - k, \
                  init,\
                  tol=1e-7)
@@ -1021,7 +1026,6 @@ $$
 For this exercise, use the following extreme price vectors as initial values:
 
 $$
-
 \begin{aligned}
     p1_{0} &= (5, 5, 5) \\
     p2_{0} &= (1, 1, 1) \\
@@ -1041,26 +1045,18 @@ Set the tolerance to $0.0$ for more accurate output.
 Define parameters and initial values
 
 ```{code-cell} ipython3
-A = np.array([
-    [0.2, 0.1, 0.7],
-    [0.3, 0.2, 0.5],
-    [0.1, 0.8, 0.1]
-])
+A = np.array([[0.2, 0.1, 0.7], [0.3, 0.2, 0.5], [0.1, 0.8, 0.1]])
 
 b = np.array([1.0, 1.0, 1.0])
 c = np.array([1.0, 1.0, 1.0])
 
-initLs = [np.repeat(5.0, 3),
-          np.ones(3),
-          np.array([4.5, 0.1, 4.0])] 
+initLs = [np.repeat(5.0, 3), np.ones(3), np.array([4.5, 0.1, 4.0])]
 ```
 
 Let’s run through each initial guess and check the output
 
 ```{code-cell} ipython3
----
-tags: [raises-exception]
----
+:tags: [raises-exception]
 
 attempt = 1
 for init in initLs:

From edb44f571d13096311ab2eea3c0fc96bc14a9204 Mon Sep 17 00:00:00 2001
From: kp992 <kpl992@outlook.com>
Date: Mon, 18 Aug 2025 20:20:33 -0700
Subject: [PATCH 2/6] fix typos

---
 lectures/newton_method.md | 64 +++++++++++++++++++--------------------
 1 file changed, 32 insertions(+), 32 deletions(-)

diff --git a/lectures/newton_method.md b/lectures/newton_method.md
index a8730f632..700790415 100644
--- a/lectures/newton_method.md
+++ b/lectures/newton_method.md
@@ -37,7 +37,7 @@ on a `GPU` is [available here](https://jax.quantecon.org/newtons_method.html)
 
 Many economic problems involve finding [fixed
 points](https://en.wikipedia.org/wiki/Fixed_point_(mathematics)) or
-[zeros](https://en.wikipedia.org/wiki/Zero_of_a_function) (sometimes called
+[zeros](https://en.wikipedia.org/wiki/Zero_of_a_function) (also called
 "roots") of functions.
 
 For example, in a simple supply and demand model, an equilibrium price is one
@@ -55,7 +55,7 @@ Newton's method does not always work but, in situations where it does,
 convergence is often fast when compared to other methods.
 
 The lecture will apply Newton's method in one-dimensional and
-multi-dimensional settings to solve fixed-point and zero-finding problems. 
+multidimensional settings to solve fixed-point and zero-finding problems. 
 
 * When finding the fixed point of a function $f$, Newton's method updates
   an existing guess of the fixed point by solving for the fixed point of a
@@ -69,10 +69,10 @@ To build intuition, we first consider an easy, one-dimensional fixed point
 problem where we know the solution and solve it using both successive
 approximation and Newton's method.
 
-Then we apply Newton's method to multi-dimensional settings to solve
+Then we apply Newton's method to multidimensional settings to solve
 market for equilibria with multiple goods.
 
-At the end of the lecture we leverage the power of automatic
+At the end of the lecture, we leverage the power of automatic
 differentiation in [`autograd`](https://github.com/HIPS/autograd) to solve a very high-dimensional equilibrium problem
 
 ```{code-cell} ipython3
@@ -117,21 +117,21 @@ zero population growth, the law of motion for capital is
 Here
 
 - $k_t$ is capital stock per worker,
-- $A, \alpha>0$ are production parameters, $\alpha<1$
+- $A, \alpha>0$ are production parameters, $\alpha < 1$
 - $s>0$ is a savings rate, and
 - $\delta \in(0,1)$ is a rate of depreciation
 
 In this example, we wish to calculate the unique strictly positive fixed point
 of $g$, the law of motion for capital.
 
-In other words, we seek a $k^* > 0$ such that $g(k^*)=k^*$.
+In other words, we seek a $k^* > 0$ such that $g(k^*) = k^*$.
 
 * such a $k^*$ is called a [steady state](https://en.wikipedia.org/wiki/Steady_state),
   since $k_t = k^*$ implies $k_{t+1} = k^*$.
 
-Using pencil and paper to solve $g(k)=k$, you will be able to confirm that
+Using pencil and paper to solve $g(k) = k$, you will be able to confirm that
 
-$$ k^* = \left(\frac{s A}{δ}\right)^{1/(1 - α)}  $$
+$$ k^* = \left(\frac{s A}{\delta}\right)^{1/(1 - \alpha)}  $$
 
 ### Implementation
 
@@ -283,13 +283,13 @@ the function
 ```{math}
 :label: motivation
 
-\hat g(x) \approx g(x_0)+g'(x_0)(x-x_0)
+\hat g(x) \approx g(x_0) + g'(x_0)(x - x_0)
 ```
 
 We solve for the fixed point of $\hat g$ by calculating the $x_1$ that solves
 
 $$
-x_1=\frac{g(x_0)-g'(x_0) x_0}{1-g'(x_0)}
+x_1 = \frac{g(x_0) - g'(x_0) x_0}{1 - g'(x_0)}
 $$
 
 Generalising the process above, Newton's fixed point method iterates on 
@@ -307,7 +307,7 @@ To implement Newton's method we observe that the derivative of the law of motion
 ```{math}
 :label: newton_method2
 
-g'(k) = \alpha s A k^{\alpha-1} + (1-\delta)
+g'(k) = \alpha s A k^{\alpha - 1} + (1 - \delta)
 
 ```
 
@@ -385,7 +385,7 @@ the problem of finding fixed points.
 
 ### Newton's method for zeros
 
-Let's suppose we want to find an $x$ such that $f(x)=0$ for some smooth
+Let's suppose we want to find an $x$ such that $f(x) = 0$ for some smooth
 function $f$ mapping real numbers to real numbers.
 
 Suppose we have a guess $x_0$ and we want to update it to a new point $x_1$.
@@ -393,7 +393,7 @@ Suppose we have a guess $x_0$ and we want to update it to a new point $x_1$.
 As a first step, we take the first-order approximation of $f$ around $x_0$:
 
 $$
-\hat f(x) \approx f\left(x_0\right)+f^{\prime}\left(x_0\right)\left(x-x_0\right)
+\hat f(x) \approx f\left(x_0\right) + f^{\prime}\left(x_0\right)\left(x - x_0\right)
 $$
 
 Now we solve for the zero of $\hat f$.  
@@ -451,7 +451,7 @@ to implement Newton's method ourselves.)
 Now consider again the Solow fixed-point calculation, where we solve for $k$
 satisfying $g(k) = k$.
 
-We can convert to this to a zero-finding problem by setting $f(x) := g(x)-x$.
+We can convert to this to a zero-finding problem by setting $f(x) := g(x) - x$.
 
 Any zero of $f$ is clearly a fixed point of $g$.
 
@@ -468,11 +468,11 @@ k_star_approx_newton = newton(
 k_star_approx_newton
 ```
 
-The result confirms the descent we saw in the graphs above: a very accurate result is reached with only 5 iterations.
+The result confirms the convergence we saw in the graphs above: a very accurate result is reached with only 5 iterations.
 
 
 
-## Multivariate Newton’s method
+## Multivariate Newton's method
 
 In this section, we introduce a two-good problem, present a
 visualization of the problem, and solve for the equilibrium of the two-good market
@@ -481,11 +481,11 @@ using both a zero finder in `SciPy` and Newton's method.
 We then expand the idea to a larger market with 5,000 goods and compare the
 performance of the two methods again.
 
-We will see a significant performance gain when using Netwon's method.
+We will see a significant performance gain when using Newton's method.
 
 
 (two_goods_market)=
-### A Two Goods market equilibrium
+### A two-goods market equilibrium
 
 Let's start by computing the market equilibrium of a two-good problem.
 
@@ -639,9 +639,9 @@ plot_excess_demand(ax, good=1)
 plt.show()
 ```
 
-We see the black contour line of zero, which tells us when $e_i(p)=0$.
+We see the black contour line of zero, which tells us when $e_i(p) = 0$.
 
-For a price vector $p$ such that $e_i(p)=0$ we know that good $i$ is in equilibrium (demand equals supply).
+For a price vector $p$ such that $e_i(p) = 0$ we know that good $i$ is in equilibrium (demand equals supply).
 
 If these two contour lines cross at some price vector $p^*$, then $p^*$ is an equilibrium price vector.
 
@@ -693,7 +693,7 @@ This is indeed a very small error.
 
 In many cases, for zero-finding algorithms applied to smooth functions, supplying the [Jacobian](https://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant) of the function leads to better convergence properties.
 
-Here we manually calculate the elements of the Jacobian
+Here, we manually calculate the elements of the Jacobian
 
 $$
 J(p) = 
@@ -747,9 +747,9 @@ This is a multivariate version of [](oneD-newton)
 
 The iteration starts from some initial guess of the price vector $p_0$. 
 
-Here, instead of coding Jacobian by hand, We use the `jacobian()` function in the `autograd` library to auto-differentiate and calculate the Jacobian.
+Here, instead of coding Jacobian by hand, we use the `jacobian()` function in the `autograd` library to auto-differentiate and calculate the Jacobian.
 
-With only slight modification, we can generalize [our previous attempt](first_newton_attempt) to multi-dimensional problems
+With only slight modification, we can generalize [our previous attempt](first_newton_attempt) to multidimensional problems
 
 ```{code-cell} ipython3
 def newton(f, x_0, tol=1e-5, max_iter=10):
@@ -865,7 +865,7 @@ A = \begin{bmatrix}
             1 & 5 & 1 \\
         \end{bmatrix},
             \quad
-s = 0.2, \quad α = 0.5, \quad δ = 0.8
+s = 0.2, \quad \alpha = 0.5, \quad \delta = 0.8
 $$
 
 As before the law of motion is
@@ -875,7 +875,7 @@ As before the law of motion is
     g(k) := sAk^\alpha + (1-\delta) k
 ```
 
-However $k_t$ is now a $3 \times 1$ vector.
+However, $k_t$ is now a $3 \times 1$ vector.
 
 Solve for the fixed point using Newton's method with the following initial values:
 
@@ -890,7 +890,7 @@ $$
 ````{hint} 
 :class: dropdown
 
-- The computation of the fixed point is equivalent to computing $k^*$ such that $f(k^*) - k^* = 0$.
+- The computation of the fixed point is equivalent to computing $k^*$ such that $g(k^*) - k^* = 0$.
 
 - If you are unsure about your solution, you can start with the solved example:
 
@@ -902,7 +902,7 @@ A = \begin{bmatrix}
         \end{bmatrix}
 ```
 
-with $s = 0.3$, $α = 0.3$, and $δ = 0.4$ and starting value: 
+with $s = 0.3$, $\alpha = 0.3$, and $\delta = 0.4$ and starting value: 
 
 
 ```{math}
@@ -932,7 +932,7 @@ s = 0.2
 initLs = [np.ones(3), np.array([3.0, 5.0, 5.0]), np.repeat(50.0, 3)]
 ```
 
-Then define the multivariate version of the formula for the [law of motion of captial](motion_law)
+Then define the multivariate version of the formula for the [law of motion of capital](motion_law)
 
 ```{code-cell} ipython3
 def multivariate_solow(k, A=A, s=s, α=α, δ=δ):
@@ -955,7 +955,7 @@ We find that the results are invariant to the starting values given the well-def
 
 But the number of iterations it takes to converge is dependent on the starting values.
 
-Let substitute the output back to the formulate to check our last result
+Let's substitute the output back into the formula to check our last result
 
 ```{code-cell} ipython3
 multivariate_solow(k) - k
@@ -1033,7 +1033,7 @@ $$
 \end{aligned}
 $$
 
-Set the tolerance to $0.0$ for more accurate output.
+Set the tolerance to $1e-15$ for more accurate output.
 
 ```{exercise-end}
 ```
@@ -1053,7 +1053,7 @@ c = np.array([1.0, 1.0, 1.0])
 initLs = [np.repeat(5.0, 3), np.ones(3), np.array([4.5, 0.1, 4.0])]
 ```
 
-Let’s run through each initial guess and check the output
+Let's run through each initial guess and check the output
 
 ```{code-cell} ipython3
 :tags: [raises-exception]
@@ -1069,7 +1069,7 @@ for init in initLs:
     attempt += 1
 ```
 
-We can find that Newton's method may fail for some starting values.
+We can see that Newton's method may fail for some starting values.
 
 Sometimes it may take a few initial guesses to achieve convergence.
 

From c63f2946941686f1aa20fe24841d6928f1244e87 Mon Sep 17 00:00:00 2001
From: kp992 <kpl992@outlook.com>
Date: Thu, 21 Aug 2025 18:36:17 -0700
Subject: [PATCH 3/6] update lecture to use JAX

---
 lectures/newton_method.md | 159 ++++++++++++++++++++------------------
 1 file changed, 84 insertions(+), 75 deletions(-)

diff --git a/lectures/newton_method.md b/lectures/newton_method.md
index 700790415..0965ae919 100644
--- a/lectures/newton_method.md
+++ b/lectures/newton_method.md
@@ -73,13 +73,8 @@ Then we apply Newton's method to multidimensional settings to solve
 market for equilibria with multiple goods.
 
 At the end of the lecture, we leverage the power of automatic
-differentiation in [`autograd`](https://github.com/HIPS/autograd) to solve a very high-dimensional equilibrium problem
+differentiation in [`jax`](https://docs.jax.dev/en/latest/_autosummary/jax.grad.html) to solve a very high-dimensional equilibrium problem
 
-```{code-cell} ipython3
-:tags: [hide-output]
-
-!pip install autograd
-```
 
 We use the following imports in this lecture
 
@@ -87,10 +82,11 @@ We use the following imports in this lecture
 import matplotlib.pyplot as plt
 from typing import NamedTuple
 from scipy.optimize import root
-from autograd import jacobian
+import jax.numpy as jnp
+import jax
 
-# Thinly-wrapped numpy to enable automatic differentiation
-import autograd.numpy as np
+# Enable 64-bit precision
+jax.config.update("jax_enable_x64", True)
 ```
 
 ## Fixed point computation using Newton's method
@@ -172,7 +168,7 @@ Here is a function to provide a 45 degree plot of the dynamics.
 def plot_45(params, ax, fontsize=14):
 
     k_min, k_max = 0.0, 3.0
-    k_grid = np.linspace(k_min, k_max, 1200)
+    k_grid = jnp.linspace(k_min, k_max, 1200)
 
     # Plot the functions
     lb = r"$g(k) = sAk^{\alpha} + (1 - \delta)k$"
@@ -353,12 +349,12 @@ def plot_trajectories(
     ax2.plot(ks4, "-o", label="newton steps")
 
     for ax in axes:
-        ax.plot(k_star * np.ones(n), "k--")
+        ax.plot(k_star * jnp.ones(n), "k--")
         ax.legend(fontsize=fs, frameon=False)
         ax.set_ylim(0.6, 3.2)
         ax.set_yticks((k_star,))
         ax.set_yticklabels(("$k^*$",), fontsize=fs)
-        ax.set_xticks(np.linspace(0, 19, 20))
+        ax.set_xticks(jnp.linspace(0, 19, 20))
 
     plt.show()
 ```
@@ -418,10 +414,12 @@ The following code implements the iteration [](oneD-newton)
 (first_newton_attempt)=
 
 ```{code-cell} ipython3
-def newton(f, Df, x_0, tol=1e-7, max_iter=100_000):
+def newton(f, x_0, tol=1e-7, max_iter=100_000):
     x = x_0
+    Df = jax.grad(f)
 
     # Implement the zero-finding formula
+    @jax.jit
     def q(x):
         return x - f(x) / Df(x)
 
@@ -432,10 +430,10 @@ def newton(f, Df, x_0, tol=1e-7, max_iter=100_000):
         if n > max_iter:
             raise Exception("Max iteration reached without convergence")
         y = q(x)
-        error = np.abs(x - y)
+        error = jnp.abs(x - y)
         x = y
         print(f"iteration {n}, error = {error:.5f}")
-    return x
+    return x.item()
 ```
 
 Numerous libraries implement Newton's method in one dimension, including
@@ -459,9 +457,7 @@ Let's apply this idea to the Solow problem
 
 ```{code-cell} ipython3
 params = create_solow_params()
-k_star_approx_newton = newton(
-    f=lambda x: g(x, params) - x, Df=lambda x: Dg(x, params) - 1, x_0=0.8
-)
+k_star_approx_newton = newton(f=lambda x: g(x, params) - x, x_0=0.8)
 ```
 
 ```{code-cell} ipython3
@@ -556,8 +552,9 @@ $$
 The function below calculates the excess demand for given parameters
 
 ```{code-cell} ipython3
+@jax.jit
 def e(p, A, b, c):
-    return np.exp(-A @ p) + c - b * np.sqrt(p)
+    return jnp.exp(-A @ p) + c - b * jnp.sqrt(p)
 ```
 
 Our default parameter values will be
@@ -581,15 +578,16 @@ A = \begin{bmatrix}
 $$
 
 ```{code-cell} ipython3
-A = np.array([[0.5, 0.4], [0.8, 0.2]])
-b = np.ones(2)
-c = np.ones(2)
+A = jnp.array([[0.5, 0.4], [0.8, 0.2]])
+b = jnp.ones(2)
+c = jnp.ones(2)
 ```
 
 At a price level of $p = (1, 0.5)$, the excess demand is
 
 ```{code-cell} ipython3
-ex_demand = e((1.0, 0.5), A, b, c)
+p = jnp.array([1, 0.5])
+ex_demand = e(p, A, b, c)
 
 print(
     f"The excess demand for good 0 is {ex_demand[0]:.3f} \n"
@@ -597,20 +595,30 @@ print(
 )
 ```
 
+To increase the efficiency of computation, we will use the power of vectorization using [`jax.vmap`](https://docs.jax.dev/en/latest/_autosummary/jax.vmap.html). This is much faster than the python loops.
+
+```{code-cell} ipython3
+# Create vectorization on the first axis of p.
+e_vectorized_p_1 = jax.vmap(e, in_axes=(0, None, None, None))
+# Create vectorization on the second axis of p.
+e_vectorized = jax.vmap(e_vectorized_p_1, in_axes=(0, None, None, None))
+```
+
 Next we plot the two functions $e_0$ and $e_1$ on a grid of $(p_0, p_1)$ values, using contour surfaces and lines.
 
 We will use the following function to build the contour plots
 
 ```{code-cell} ipython3
 def plot_excess_demand(ax, good=0, grid_size=100, grid_max=4, surface=True):
+    p_grid = jnp.linspace(0, grid_max, grid_size)
+    # Create meshgrid for all combinations of p_1 and p_2
+    P1, P2 = jnp.meshgrid(p_grid, p_grid, indexing="ij")
+    # Stack to create array of shape (grid_size, grid_size, 2)
+    P = jnp.stack([P1, P2], axis=-1)
 
-    # Create a 100x100 grid
-    p_grid = np.linspace(0, grid_max, grid_size)
-    z = np.empty((100, 100))
-
-    for i, p_1 in enumerate(p_grid):
-        for j, p_2 in enumerate(p_grid):
-            z[i, j] = e((p_1, p_2), A, b, c)[good]
+    # Compute all values at once using vectorized function
+    z_full = e_vectorized(P, A, b, c)
+    z = z_full[:, :, good]
 
     if surface:
         cs1 = ax.contourf(p_grid, p_grid, z.T, alpha=0.5)
@@ -662,7 +670,7 @@ To solve for $p^*$ more precisely, we use a zero-finding algorithm from `scipy.o
 We supply $p = (1, 1)$ as our initial guess.
 
 ```{code-cell} ipython3
-init_p = np.ones(2)
+init_p = jnp.ones(2)
 ```
 
 This uses the [modified Powell method](https://docs.scipy.org/doc/scipy/reference/optimize.root-hybr.html#optimize-root-hybr) to find the zero
@@ -682,7 +690,7 @@ p
 This looks close to our guess from observing the figure. We can plug it back into $e$ to test that $e(p) \approx 0$:
 
 ```{code-cell} ipython3
-e_p = np.max(np.abs(e(p, A, b, c)))
+e_p = jnp.max(jnp.abs(e(p, A, b, c)))
 e_p.item()
 ```
 
@@ -708,12 +716,12 @@ def jacobian_e(p, A, b, c):
     p_0, p_1 = p
     a_00, a_01 = A[0, :]
     a_10, a_11 = A[1, :]
-    j_00 = -a_00 * np.exp(-a_00 * p_0) - (b[0] / 2) * p_0 ** (-1 / 2)
-    j_01 = -a_01 * np.exp(-a_01 * p_1)
-    j_10 = -a_10 * np.exp(-a_10 * p_0)
-    j_11 = -a_11 * np.exp(-a_11 * p_1) - (b[1] / 2) * p_1 ** (-1 / 2)
+    j_00 = -a_00 * jnp.exp(-a_00 * p_0) - (b[0] / 2) * p_0 ** (-1 / 2)
+    j_01 = -a_01 * jnp.exp(-a_01 * p_1)
+    j_10 = -a_10 * jnp.exp(-a_10 * p_0)
+    j_11 = -a_11 * jnp.exp(-a_11 * p_1) - (b[1] / 2) * p_1 ** (-1 / 2)
     J = [[j_00, j_01], [j_10, j_11]]
-    return np.array(J)
+    return jnp.array(J)
 ```
 
 ```{code-cell} ipython3
@@ -727,7 +735,7 @@ Now the solution is even more accurate (although, in this low-dimensional proble
 
 ```{code-cell} ipython3
 p = solution.x
-e_p = np.max(np.abs(e(p, A, b, c)))
+e_p = jnp.max(jnp.abs(e(p, A, b, c)))
 e_p.item()
 ```
 
@@ -747,14 +755,19 @@ This is a multivariate version of [](oneD-newton)
 
 The iteration starts from some initial guess of the price vector $p_0$. 
 
-Here, instead of coding Jacobian by hand, we use the `jacobian()` function in the `autograd` library to auto-differentiate and calculate the Jacobian.
+Here, instead of coding Jacobian by hand, we use the `jacobian()` function in the `jax` library to auto-differentiate and calculate the Jacobian.
 
 With only slight modification, we can generalize [our previous attempt](first_newton_attempt) to multidimensional problems
 
 ```{code-cell} ipython3
 def newton(f, x_0, tol=1e-5, max_iter=10):
     x = x_0
-    q = lambda x: x - np.linalg.solve(jacobian(f)(x), f(x))
+    f_jac = jax.jacobian(f)
+
+    @jax.jit
+    def q(x):
+        return x - jnp.linalg.solve(f_jac(x), f(x))
+
     error = tol + 1
     n = 0
     while error > tol:
@@ -762,20 +775,15 @@ def newton(f, x_0, tol=1e-5, max_iter=10):
         if n > max_iter:
             raise Exception("Max iteration reached without convergence")
         y = q(x)
-        if any(np.isnan(y)):
+        if any(jnp.isnan(y)):
             raise Exception("Solution not found with NaN generated")
-        error = np.linalg.norm(x - y)
+        error = jnp.linalg.norm(x - y)
         x = y
         print(f"iteration {n}, error = {error:.5f}")
     print("\n" + f"Result = {x} \n")
     return x
 ```
 
-```{code-cell} ipython3
-def e(p, A, b, c):
-    return np.exp(-np.dot(A, p)) + c - b * np.sqrt(p)
-```
-
 We find the algorithm terminates in 4 steps
 
 ```{code-cell} ipython3
@@ -784,7 +792,7 @@ p = newton(lambda p: e(p, A, b, c), init_p)
 ```
 
 ```{code-cell} ipython3
-e_p = np.max(np.abs(e(p, A, b, c)))
+e_p = jnp.max(jnp.abs(e(p, A, b, c)))
 e_p.item()
 ```
 
@@ -797,30 +805,29 @@ With the larger overhead, the speed is not better than the optimized `scipy` fun
 
 Our next step is to investigate a large market with 3,000 goods.
 
-A JAX version of this section using GPU accelerated linear algebra and
-automatic differentiation is available [here](https://jax.quantecon.org/newtons_method.html#application)
 
 The excess demand function is essentially the same, but now the matrix $A$ is $3000 \times 3000$ and the parameter vectors $b$ and $c$ are $3000 \times 1$.
 
 ```{code-cell} ipython3
 dim = 3000
-np.random.seed(123)
 
-# Create a random matrix A and normalize the rows to sum to one
-A = np.random.rand(dim, dim)
-A = np.asarray(A)
-s = np.sum(A, axis=0)
+# Create JAX random key
+key = jax.random.PRNGKey(123)
+
+# Create a random matrix A and normalize the columns to sum to one
+A = jax.random.uniform(key, (dim, dim))
+s = jnp.sum(A, axis=0)
 A = A / s
 
 # Set up b and c
-b = np.ones(dim)
-c = np.ones(dim)
+b = jnp.ones(dim)
+c = jnp.ones(dim)
 ```
 
 Here's our initial condition
 
 ```{code-cell} ipython3
-init_p = np.ones(dim)
+init_p = jnp.ones(dim)
 ```
 
 ```{code-cell} ipython3
@@ -829,7 +836,7 @@ p = newton(lambda p: e(p, A, b, c), init_p)
 ```
 
 ```{code-cell} ipython3
-e_p = np.max(np.abs(e(p, A, b, c)))
+e_p = jnp.max(jnp.abs(e(p, A, b, c)))
 e_p.item()
 ```
 
@@ -837,16 +844,18 @@ With the same tolerance, we compare the runtime and accuracy of Newton's method
 
 ```{code-cell} ipython3
 %%time
-solution = root(lambda p: e(p, A, b, c),
-                init_p, 
-                jac=lambda p: jacobian(e)(p, A, b, c), 
-                method='hybr',
-                tol=1e-5)
+solution = root(
+    lambda p: e(p, A, b, c),
+    init_p,
+    jac=lambda p: jax.jacobian(e)(p, A, b, c),
+    method="hybr",
+    tol=1e-5,
+)
 ```
 
 ```{code-cell} ipython3
 p = solution.x
-e_p = np.max(np.abs(e(p, A, b, c)))
+e_p = jnp.max(jnp.abs(e(p, A, b, c)))
 e_p.item()
 ```
 
@@ -923,20 +932,21 @@ The result should converge to the [analytical solution](solved_k).
 Let's first define the parameters for this problem
 
 ```{code-cell} ipython3
-A = np.array([[2.0, 3.0, 3.0], [2.0, 4.0, 2.0], [1.0, 5.0, 1.0]])
+A = jnp.array([[2.0, 3.0, 3.0], [2.0, 4.0, 2.0], [1.0, 5.0, 1.0]])
 
 s = 0.2
 α = 0.5
 δ = 0.8
 
-initLs = [np.ones(3), np.array([3.0, 5.0, 5.0]), np.repeat(50.0, 3)]
+initLs = [jnp.ones(3), jnp.array([3.0, 5.0, 5.0]), jnp.repeat(50.0, 3)]
 ```
 
 Then define the multivariate version of the formula for the [law of motion of capital](motion_law)
 
 ```{code-cell} ipython3
+@jax.jit
 def multivariate_solow(k, A=A, s=s, α=α, δ=δ):
-    return s * np.dot(A, k**α) + (1 - δ) * k
+    return s * jnp.dot(A, k**α) + (1 - δ) * k
 ```
 
 Let's run through each starting value and see the output
@@ -966,7 +976,7 @@ Note the error is very small.
 We can also test our results on the known solution
 
 ```{code-cell} ipython3
-A = np.array([[2.0, 0.0, 0.0],
+A = jnp.array([[2.0, 0.0, 0.0],
                [0.0, 2.0, 0.0],
                [0.0, 0.0, 2.0]])
 
@@ -974,7 +984,7 @@ s = 0.3
 α = 0.3
 δ = 0.4
 
-init = np.repeat(1.0, 3)
+init = jnp.repeat(1.0, 3)
 
 
 %time k = newton(lambda k: multivariate_solow(k, A=A, s=s, α=α, δ=δ) - k, \
@@ -1045,12 +1055,11 @@ Set the tolerance to $1e-15$ for more accurate output.
 Define parameters and initial values
 
 ```{code-cell} ipython3
-A = np.array([[0.2, 0.1, 0.7], [0.3, 0.2, 0.5], [0.1, 0.8, 0.1]])
-
-b = np.array([1.0, 1.0, 1.0])
-c = np.array([1.0, 1.0, 1.0])
+A = jnp.array([[0.2, 0.1, 0.7], [0.3, 0.2, 0.5], [0.1, 0.8, 0.1]])
+b = jnp.array([1.0, 1.0, 1.0])
+c = jnp.array([1.0, 1.0, 1.0])
 
-initLs = [np.repeat(5.0, 3), np.ones(3), np.array([4.5, 0.1, 4.0])]
+initLs = [jnp.repeat(5.0, 3), jnp.ones(3), jnp.array([4.5, 0.1, 4.0])]
 ```
 
 Let's run through each initial guess and check the output

From d15bf8be877a4a29c72d09a243a900d28b46af6a Mon Sep 17 00:00:00 2001
From: kp992 <kpl992@outlook.com>
Date: Fri, 22 Aug 2025 21:26:51 -0700
Subject: [PATCH 4/6] apply suggestions

---
 lectures/newton_method.md | 26 +++++++++++++++++---------
 1 file changed, 17 insertions(+), 9 deletions(-)

diff --git a/lectures/newton_method.md b/lectures/newton_method.md
index 0965ae919..118c3fc11 100644
--- a/lectures/newton_method.md
+++ b/lectures/newton_method.md
@@ -127,7 +127,9 @@ In other words, we seek a $k^* > 0$ such that $g(k^*) = k^*$.
 
 Using pencil and paper to solve $g(k) = k$, you will be able to confirm that
 
-$$ k^* = \left(\frac{s A}{\delta}\right)^{1/(1 - \alpha)}  $$
+$$
+k^* = \left(\frac{s A}{\delta}\right)^{1/(1 - \alpha)}
+$$
 
 ### Implementation
 
@@ -329,9 +331,9 @@ def plot_trajectories(
     params,
     k0_a=0.8,  # first initial condition
     k0_b=3.1,  # second initial condition
-    n=20,  # length of time series
-    fs=14,
-):  # fontsize
+    n=20,      # length of time series
+    fs=14,     # fontsize
+):
 
     fig, axes = plt.subplots(2, 1, figsize=(10, 6))
     ax1, ax2 = axes
@@ -457,14 +459,14 @@ Let's apply this idea to the Solow problem
 
 ```{code-cell} ipython3
 params = create_solow_params()
-k_star_approx_newton = newton(f=lambda x: g(x, params) - x, x_0=0.8)
+k_star_approx_newton = newton(f = lambda x: g(x, params) - x, x_0=0.8)
 ```
 
 ```{code-cell} ipython3
 k_star_approx_newton
 ```
 
-The result confirms the convergence we saw in the graphs above: a very accurate result is reached with only 5 iterations.
+The result confirms convergence we saw in the graphs above: a very accurate result is reached with only 5 iterations.
 
 
 
@@ -600,6 +602,7 @@ To increase the efficiency of computation, we will use the power of vectorizatio
 ```{code-cell} ipython3
 # Create vectorization on the first axis of p.
 e_vectorized_p_1 = jax.vmap(e, in_axes=(0, None, None, None))
+
 # Create vectorization on the second axis of p.
 e_vectorized = jax.vmap(e_vectorized_p_1, in_axes=(0, None, None, None))
 ```
@@ -611,8 +614,10 @@ We will use the following function to build the contour plots
 ```{code-cell} ipython3
 def plot_excess_demand(ax, good=0, grid_size=100, grid_max=4, surface=True):
     p_grid = jnp.linspace(0, grid_max, grid_size)
+
     # Create meshgrid for all combinations of p_1 and p_2
     P1, P2 = jnp.meshgrid(p_grid, p_grid, indexing="ij")
+
     # Stack to create array of shape (grid_size, grid_size, 2)
     P = jnp.stack([P1, P2], axis=-1)
 
@@ -627,7 +632,7 @@ def plot_excess_demand(ax, good=0, grid_size=100, grid_max=4, surface=True):
     ctr1 = ax.contour(p_grid, p_grid, z.T, levels=[0.0])
     ax.set_xlabel("$p_0$")
     ax.set_ylabel("$p_1$")
-    ax.set_title(f"Excess Demand for Good {good}")
+    ax.set_title(f"Excess demand for good {good}")
     plt.clabel(ctr1, inline=1, fontsize=13)
 ```
 
@@ -727,7 +732,10 @@ def jacobian_e(p, A, b, c):
 ```{code-cell} ipython3
 %%time
 solution = root(
-    lambda p: e(p, A, b, c), init_p, jac=lambda p: jacobian_e(p, A, b, c), method="hybr"
+    lambda p: e(p, A, b, c),
+    init_p,
+    jac = lambda p: jacobian_e(p, A, b, c),
+    method="hybr"
 )
 ```
 
@@ -847,7 +855,7 @@ With the same tolerance, we compare the runtime and accuracy of Newton's method
 solution = root(
     lambda p: e(p, A, b, c),
     init_p,
-    jac=lambda p: jax.jacobian(e)(p, A, b, c),
+    jac = lambda p: jax.jacobian(e)(p, A, b, c),
     method="hybr",
     tol=1e-5,
 )

From bed18732fb4af027432604f3fc99550e9a3dc951 Mon Sep 17 00:00:00 2001
From: kp992 <kpl992@outlook.com>
Date: Mon, 8 Sep 2025 21:59:50 -0700
Subject: [PATCH 5/6] update code line length

---
 lectures/newton_method.md | 33 +++++++++++++++++----------------
 1 file changed, 17 insertions(+), 16 deletions(-)

diff --git a/lectures/newton_method.md b/lectures/newton_method.md
index 0575e828d..827bd93b8 100644
--- a/lectures/newton_method.md
+++ b/lectures/newton_method.md
@@ -83,7 +83,9 @@ from typing import NamedTuple
 from scipy.optimize import root
 import jax.numpy as jnp
 import jax
-jax.config.update("jax_enable_x64", True) # Enable 64-bit precision
+
+# Enable 64-bit precision
+jax.config.update("jax_enable_x64", True)
 ```
 
 ## Fixed point computation using Newton's method
@@ -732,7 +734,7 @@ solution = root(
     lambda p: e(p, A, b, c),
     init_p,
     jac = lambda p: jacobian_e(p, A, b, c),
-    method="hybr"
+    method="hybr",
 )
 ```
 
@@ -980,27 +982,29 @@ Note the error is very small.
 We can also test our results on the known solution
 
 ```{code-cell} ipython3
-A = jnp.array([[2.0, 0.0, 0.0],
-               [0.0, 2.0, 0.0],
-               [0.0, 0.0, 2.0]])
+A = jnp.array([[2.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 2.0]])
 
 s = 0.3
 α = 0.3
 δ = 0.4
 
 init = jnp.repeat(1.0, 3)
+```
 
+```{code-cell} ipython3
+%%time
 
-%time k = newton(lambda k: multivariate_solow(k, A=A, s=s, α=α, δ=δ) - k, \
-                 init)
+k = newton(lambda k: multivariate_solow(k, A=A, s=s, α=α, δ=δ) - k, init)
 ```
 
 The result is very close to the ground truth but still slightly different.
 
 ```{code-cell} ipython3
-%time k = newton(lambda k: multivariate_solow(k, A=A, s=s, α=α, δ=δ) - k, \
-                 init,\
-                 tol=1e-7)
+%%time
+
+k = newton(
+    lambda k: multivariate_solow(k, A=A, s=s, α=α, δ=δ) - k, init, tol=1e-7
+)
 ```
 
 We can see it steps towards a more accurate solution.
@@ -1073,12 +1077,9 @@ Let's run through each initial guess and check the output
 
 attempt = 1
 for init in initLs:
-    print(f'Attempt {attempt}: Starting value is {init} \n')
-    %time p = newton(lambda p: e(p, A, b, c), \
-                init, \
-                tol=1e-15, \
-                max_iter=15)
-    print('-'*64)
+    print(f"Attempt {attempt}: Starting value is {init} \n")
+    %time p = newton(lambda p: e(p, A, b, c), init, tol=1e-15, max_iter=15)
+    print("-" * 64)
     attempt += 1
 ```
 

From e3115df03e9e066b14fefede3a46d1f58e777bed Mon Sep 17 00:00:00 2001
From: Humphrey Yang <u6474961@anu.edu.au>
Date: Wed, 10 Sep 2025 13:45:49 +1000
Subject: [PATCH 6/6] minor updates

---
 lectures/newton_method.md | 22 +++++++++-------------
 1 file changed, 9 insertions(+), 13 deletions(-)

diff --git a/lectures/newton_method.md b/lectures/newton_method.md
index 827bd93b8..eb170aa63 100644
--- a/lectures/newton_method.md
+++ b/lectures/newton_method.md
@@ -28,10 +28,6 @@ kernelspec:
 :depth: 2
 ```
 
-```{seealso}
-A version of this lecture using [JAX](https://github.com/jax-ml/jax) is {doc}`available here <jax:newtons_method>`
-```
-
 ## Overview
 
 Many economic problems involve finding [fixed
@@ -69,10 +65,10 @@ problem where we know the solution and solve it using both successive
 approximation and Newton's method.
 
 Then we apply Newton's method to multidimensional settings to solve
-market for equilibria with multiple goods.
+for market equilibria with multiple goods.
 
 At the end of the lecture, we leverage the power of automatic
-differentiation in [`jax`](https://docs.jax.dev/en/latest/_autosummary/jax.grad.html) to solve a very high-dimensional equilibrium problem
+differentiation in [`jax`](https://docs.jax.dev/en/latest/_autosummary/jax.grad.html) to solve a very high-dimensional equilibrium problem.
 
 
 We use the following imports in this lecture
@@ -112,7 +108,7 @@ zero population growth, the law of motion for capital is
 Here
 
 - $k_t$ is capital stock per worker,
-- $A, \alpha>0$ are production parameters, $\alpha < 1$
+- $A, \alpha>0$ are production parameters with $\alpha < 1$
 - $s>0$ is a savings rate, and
 - $\delta \in(0,1)$ is a rate of depreciation
 
@@ -121,7 +117,7 @@ of $g$, the law of motion for capital.
 
 In other words, we seek a $k^* > 0$ such that $g(k^*) = k^*$.
 
-* such a $k^*$ is called a [steady state](https://en.wikipedia.org/wiki/Steady_state),
+* Such a $k^*$ is called a [steady state](https://en.wikipedia.org/wiki/Steady_state),
   since $k_t = k^*$ implies $k_{t+1} = k^*$.
 
 Using pencil and paper to solve $g(k) = k$, you will be able to confirm that
@@ -228,7 +224,7 @@ Here's a time series from a particular choice of $k_0$.
 
 ```{code-cell} ipython3
 def compute_iterates(k_0, f, params, n=25):
-    """Compute time series of length n generated by arbitrary function f."""
+    """Compute time series of length n generated by function f."""
     k = k_0
     k_iterates = []
     for t in range(n):
@@ -489,13 +485,13 @@ Let's start by computing the market equilibrium of a two-good problem.
 We consider a market for two related products, good 0 and good 1, with
 price vector $p = (p_0, p_1)$
 
-Supply of good $i$ at price $p$,
+Supply of good $i$ at price $p$ is
 
 $$ 
 q^s_i (p) = b_i \sqrt{p_i} 
 $$
 
-Demand of good $i$ at price $p$ is,
+Demand of good $i$ at price $p$ is
 
 $$ 
 q^d_i (p) = \exp(-(a_{i0} p_0 + a_{i1} p_1)) + c_i
@@ -505,7 +501,7 @@ Here $c_i$, $b_i$ and $a_{ij}$ are parameters.
 
 For example, the two goods might be computer components that are typically used together, in which case they are complements. Hence demand depends on the price of both components.
 
-The excess demand function is,
+The excess demand function is
 
 $$
 e_i(p) = q^d_i(p) - q^s_i(p), \quad i = 0, 1
@@ -818,7 +814,7 @@ The excess demand function is essentially the same, but now the matrix $A$ is $3
 dim = 3000
 
 # Create JAX random key
-key = jax.random.PRNGKey(123)
+key = jax.random.PRNGKey(0)
 
 # Create a random matrix A and normalize the columns to sum to one
 A = jax.random.uniform(key, (dim, dim))