You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Copy file name to clipboardExpand all lines: lectures/newtons_method.md
+40-20Lines changed: 40 additions & 20 deletions
Display the source diff
Display the rich diff
Original file line number
Diff line number
Diff line change
@@ -26,9 +26,13 @@ Alternatively, if you have your own GPU, you can follow the [instructions](https
26
26
27
27
## Overview
28
28
29
-
Continuing from the [Newton's Method lecture](https://python.quantecon.org/newton_method.html), we are going to solve the multidimensional problem with `JAX`.
29
+
In this lecture we highlight some of the capabilities of JAX, including JIT
30
+
compilation and automatic differentiation.
30
31
31
-
More information about JAX can be found [here](https://python-programming.quantecon.org/jax_intro.html).
32
+
The application is computing equilibria via Newton's method, which we discussed
33
+
in [a more elementary QuantEcon lecture](https://python.quantecon.org/newton_method.html)
34
+
35
+
Here our focus is on how to apply JAX to this problem.
32
36
33
37
We use the following imports in this lecture
34
38
@@ -38,11 +42,21 @@ import jax.numpy as jnp
38
42
from scipy.optimize import root
39
43
```
40
44
41
-
## The Two Goods Market Equilibrium
45
+
## The Equilibrium Problem
46
+
47
+
In this section we describe the market equilibrium problem we will solve with
48
+
JAX.
49
+
50
+
We begin with a two good case,
51
+
which is borrowed from [an earlier lecture](https://python.quantecon.org/newton_method.html).
42
52
43
-
Let's have a quick recap of this problem -- a more detailed explanation and derivation can be found at [A Two Goods Market Equilibrium](https://python.quantecon.org/newton_method.html#two-goods-market).
53
+
Then we shift to higher dimensions.
44
54
45
-
Assume we have a market for two complementary goods where demand depends on the price of both components.
55
+
56
+
### The Two Goods Market Equilibrium
57
+
58
+
Assume we have a market for two complementary goods where demand depends on the
59
+
price of both components.
46
60
47
61
We label them good 0 and good 1, with price vector $p = (p_0, p_1)$.
48
62
@@ -90,23 +104,24 @@ $$
90
104
for this particular question.
91
105
92
106
93
-
### The Multivariable Market Equilibrium
107
+
### A High-Dimensional Version
94
108
95
-
We can now easily get the multivariable version of the problem above.
109
+
Let's now shift to a linear algebra formulation, which alllows us to handle
110
+
arbitrarily many goods.
96
111
97
112
The supply function remains unchanged,
98
113
99
114
$$
100
-
q^s (p) =b \sqrt{p}
115
+
q^s (p) =b \sqrt{p}
101
116
$$
102
117
103
-
The demand function is,
118
+
The demand function becomes
104
119
105
120
$$
106
-
q^d (p) = \text{exp}(- A \cdot p) + c
121
+
q^d (p) = \text{exp}(- A \cdot p) + c
107
122
$$
108
123
109
-
Our new excess demand function is,
124
+
Our new excess demand function is
110
125
111
126
$$
112
127
e(p) = \text{exp}(- A \cdot p) + c - b \sqrt{p}
@@ -120,9 +135,17 @@ def e(p, A, b, c):
120
135
```
121
136
122
137
123
-
## Using Newton's Method
124
138
125
-
Now let's use the multivariate version of Newton's method to compute the equilibrium price
139
+
## Computation
140
+
141
+
In this section we describe and then implement the solution method.
142
+
143
+
144
+
### Newton's Method
145
+
146
+
We use a multivariate version of Newton's method to compute the equilibrium price.
147
+
148
+
The rule for updating a guess $p_n$ of the price vector is
Here $J_e(p_n)$ is the Jacobian of $e$ evaluated at $p_n$.
133
156
134
-
The iteration starts from some initial guess of the price vector $p_0$.
135
-
136
-
Here, instead of coding Jacobian by hand, We use the `jax.jacobian()` function to auto-differentiate and calculate the Jacobian.
157
+
Iteration starts from initial guess $p_0$.
137
158
138
-
With only slight modification, we can generalize [our previous attempt](https://python.quantecon.org/newton_method.html#first-newton-attempt) to multi-dimensional problems
159
+
Instead of coding the Jacobian by hand, we use `jax.jacobian()`.
0 commit comments