7. 使用牛顿法求解经济模型#

GPU

本讲座是在配有GPU的机器上构建的——不过没有GPU也可以运行。

Google Colab 提供带GPU的免费套餐，使用方法如下：

  1. 点击右上角的”播放”图标

  2. 选择 Colab

  3. 将运行时环境设置为包含GPU

除了 Anaconda 中已有的库之外，本讲座还需要以下库：

!pip install jax

Hide code cell output

 
Requirement already satisfied: jax in /home/runner/miniconda3/envs/quantecon/lib/python3.13/site-packages (0.11.1)
Requirement already satisfied: jaxlib<=0.11.1,>=0.11.1 in /home/runner/miniconda3/envs/quantecon/lib/python3.13/site-packages (from jax) (0.11.1)
Requirement already satisfied: ml_dtypes>=0.5.0 in /home/runner/miniconda3/envs/quantecon/lib/python3.13/site-packages (from jax) (0.6.0)
Requirement already satisfied: numpy>=2.1 in /home/runner/miniconda3/envs/quantecon/lib/python3.13/site-packages (from jax) (2.4.6)
Requirement already satisfied: opt_einsum in /home/runner/miniconda3/envs/quantecon/lib/python3.13/site-packages (from jax) (3.4.0)
Requirement already satisfied: scipy>=1.15 in /home/runner/miniconda3/envs/quantecon/lib/python3.13/site-packages (from jax) (1.18.0)

7.1. 概述#

许多经济问题涉及寻找不动点或零点（有时也称为”根”）。

例如，在简单的供需模型中，均衡价格是使超额需求为零的价格。

换句话说，均衡是超额需求函数的零点。

有各种算法可用于求解不动点和零点。

在本讲中，我们将学习一种重要的基于梯度的技术，称为牛顿法。

牛顿法并非总是有效，但在适用的情况下，其收敛速度通常比其他方法更快。

本讲将在一维和多维环境中应用牛顿法来解决不动点和零点计算问题。

  • 在寻找函数 \(f\) 的不动点时，牛顿法通过求解对函数 \(f\) 的线性近似的不动点，来更新已有的不动点猜测值。

  • 在寻找函数 \(f\) 的零点时，牛顿法通过求解对函数 \(f\) 的线性近似的零点，来更新已有的猜测值。

为了建立直观认识，我们首先考虑一个简单的一维不动点问题，其中我们已知解，并使用连续近似和牛顿法来求解。

然后我们将牛顿法应用到多维环境中，求解多种商品的市场均衡问题。

在本讲的最后，我们将利用 jax 中自动微分的能力来求解一个非常高维的均衡问题。

我们在本讲中使用以下导入语句

import matplotlib.pyplot as plt
import matplotlib as mpl
FONTPATH = "fonts/SourceHanSerifSC-SemiBold.otf"
mpl.font_manager.fontManager.addfont(FONTPATH)
plt.rcParams['font.family'] = ['Source Han Serif SC']
from typing import NamedTuple
from scipy.optimize import root
import jax.numpy as jnp
import jax

# 启用64位精度
jax.config.update("jax_enable_x64", True)

7.2. 用牛顿法计算不动点#

在本节中，我们将在索洛增长模型的框架下求解资本运动规律的不动点。

我们将通过可视化方式检查不动点，用连续逼近法求解，然后应用牛顿法来实现更快的收敛。

7.2.1. 索洛模型#

在索洛增长模型中，假设采用柯布-道格拉斯生产技术且人口零增长，资本的运动规律为

(7.1)#\[ k_{t+1} = g(k_t) \quad \text{where} \quad g(k) := sAk^\alpha + (1-\delta) k\]

其中

  • \(k_t\) 是人均资本存量

  • \(A, \alpha>0\) 是生产参数，\(\alpha<1\)

  • \(s>0\) 是储蓄率

  • \(\delta \in(0,1)\) 是折旧率

在这个例子中，我们希望计算资本运动规律\(g\)的唯一严格正不动点。

换句话说，我们要寻找一个 \(k^* > 0\) 使得 \(g(k^*)=k^*\)。

  • 这样的 \(k^*\) 被称为稳态， 因为当 \(k_t = k^*\) 时，\(k_{t+1} = k^*\)。

用纸笔解方程 \(g(k)=k\)，你可以验证

\[ k^* = \left(\frac{s A}{\delta}\right)^{1/(1 - \alpha)} \]

7.2.2. 实现#

让我们使用 NamedTuple 来存储我们的参数，这有助于保持代码的整洁和简洁。

class SolowParameters(NamedTuple):
    A: float
    s: float
    α: float
    δ: float

此函数创建一个带有默认参数值的适当的 SolowParameters。

def create_solow_params(A=2.0, s=0.3, α=0.3, δ=0.4):
    """Creates a Solow model parameterization with default values."""
    return SolowParameters(A=A, s=s, α=α, δ=δ)

接下来的两个函数实现运动定律(7.1)并存储真实的不动点\(k^*\)。

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 - α))

这是一个用于绘制45度动态图的函数。

def plot_45(params, ax, fontsize=14):

    k_min, k_max = 0.0, 3.0
    k_grid = jnp.linspace(k_min, k_max, 1200)

    # 绘制函数
    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, k_grid, "k--", lw=1, alpha=0.7, label="45")

    # 显示并标注固定点
    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.legend(loc="upper left", frameon=False, fontsize=fontsize)

    ax.set_yticks((0, 1, 2, 3))
    ax.set_yticklabels((0.0, 1.0, 2.0, 3.0), fontsize=fontsize)
    ax.set_ylim(0, 3)
    ax.set_xlabel("$k_t$", fontsize=fontsize)
    ax.set_ylabel("$k_{t+1}$", fontsize=fontsize)

让我们看看两个参数化的45度图。

params = create_solow_params()
fig, ax = plt.subplots(figsize=(8, 8))
plot_45(params, ax)
plt.show()

findfont: Failed to find font weight normal, now using 600.
findfont: Failed to find font weight normal, now using 600.
findfont: Failed to find font weight normal, now using 600.
_images/9f917a4ad0a888127c6cdae46033f9fab00e44a1f11eceea76b5fef6c2078e66.png 
params = create_solow_params(α=0.05, δ=0.5)
fig, ax = plt.subplots(figsize=(8, 8))
plot_45(params, ax)
plt.show()
_images/5d4dcdab96f5d6227b1f82d66b3c593b28fb6fe41443ae46a909b2a1f41b6fb8.png

我们看到 \(k^*\) 确实是唯一的正固定点。

7.2.2.1. 连续近似法#

首先让我们用连续近似法来计算固定点。

在这种情况下，连续近似法意味着从某个初始状态 \(k_0\) 开始，使用运动规律反复更新资本。

这里是以特定的 \(k_0\) 为初始值得到的时间序列。

def compute_iterates(k_0, f, params, n=25):
    """Compute time series of length n generated by function f."""
    k = k_0
    k_iterates = []
    for t in range(n):
        k_iterates.append(k)
        k = f(k, params)
    return k_iterates

params = create_solow_params()
k_0 = 0.25
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.set_ylim(0, 3)
plt.show()
_images/da54eb6c62e4912aefbbeb97d1f348abd45152b05fe76e848ae49b35b4b372af.png

让我们看看长时间序列的输出。

k_series = compute_iterates(k_0, g, params, n=10_000)
k_star_approx = k_series[-1]
k_star_approx

1.7846741842265788

这接近真实值。

k_star

1.7846741842265788

7.2.2.2. 牛顿法#

一般来说，当对某个函数\(g\)应用牛顿不动点法时，我们从一个不动点的猜测值\(x_0\)开始，然后通过求解\(x_0\)处切线的不动点来更新。

首先，我们回顾一下\(g\)在\(x_0\)处的一阶近似（即\(g\)在\(x_0\)处的一阶泰勒近似）是以下函数：

(7.2)#\[\hat g(x) \approx g(x_0) + g'(x_0)(x - x_0)\]

我们通过计算满足以下等式的\(x_1\)来求解\(\hat g\)的不动点：

\[ x_1 = \frac{g(x_0) - g'(x_0) x_0}{1 - g'(x_0)} \]

推广上述过程，牛顿不动点法的迭代公式为：

(7.3)#\[x_{t+1} = \frac{g(x_t) - g'(x_t) x_t}{ 1 - g'(x_t) }, \quad x_0 \text{ 给定}\]

要实现牛顿法，我们观察到资本运动定律(7.1)的导数为：

(7.4)#\[g'(k) = \alpha s A k^{\alpha - 1} + (1 - \delta)\]

让我们定义这个函数

def Dg(k, params):
    A, s, α, δ = params
    return α * A * s * k ** (α - 1) + (1 - δ)

下面的函数 \(q\) 表示 (7.3)。

def q(k, params):
    return (g(k, params) - Dg(k, params) * k) / (1 - Dg(k, params))

现在让我们绘制一些轨迹。

def plot_trajectories(
    params,
    k0_a=0.8,  # 第一个初始条件
    k0_b=3.1,  # 第二个初始条件
    n=20,      # 时间序列长度
    fs=14,     # 字体大小
):

    fig, axes = plt.subplots(2, 1, figsize=(10, 6))
    ax1, ax2 = axes

    ks1 = compute_iterates(k0_a, g, params, n)
    ax1.plot(ks1, "-o", label="连续近似")

    ks2 = compute_iterates(k0_b, g, params, n)
    ax2.plot(ks2, "-o", label="连续近似")

    ks3 = compute_iterates(k0_a, q, params, n)
    ax1.plot(ks3, "-o", label="牛顿法")

    ks4 = compute_iterates(k0_b, q, params, n)
    ax2.plot(ks4, "-o", label="牛顿法")

    for ax in axes:
        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(jnp.linspace(0, 19, 20))

    plt.show()

params = create_solow_params()
plot_trajectories(params)
_images/74e5310838e733aae9095c1cd896203636c44566e9e20edcdf8dd6f514f7f697.png

我们可以看到牛顿法比连续逼近法收敛得更快。

7.3. 一维求根#

在上一节中我们计算了不动点。

事实上，牛顿法更常与寻找函数零点的问题相关联。

让我们讨论这个”求根”问题，然后说明它与寻找不动点的问题是如何联系的。

7.3.1. 牛顿法求零点#

假设我们想要找到一个 \(x\) 使得对某个光滑函数 \(f\) (从实数映射到实数)有 \(f(x)=0\)。

假设我们有一个猜测值 \(x_0\) 并且想要将其更新为新的点 \(x_1\)。

作为第一步，我们取 \(f\) 在 \(x_0\) 处的一阶近似：

\[ \hat f(x) \approx f\left(x_0\right)+f^{\prime}\left(x_0\right)\left(x-x_0\right) \]

现在我们求解 \(\hat f\) 的零点。

具体来说，我们令 \(\hat{f}(x_1) = 0\) 并求解 \(x_1\)，得到

\[ x_1 = x_0 - \frac{ f(x_0) }{ f'(x_0) }, \quad x_0 \text{ 已知} \]

对于一维零点查找问题，牛顿法的迭代公式可以概括为

(7.5)#\[x_{t+1} = x_t - \frac{ f(x_t) }{ f'(x_t) }, \quad x_0 \text{ 给定}\]

以下代码实现了迭代公式 (7.5)

def newton(f, x_0, tol=1e-7, max_iter=100_000):
    x = x_0
    Df = jax.grad(f)

    # 实现零点查找公式
    @jax.jit
    def q(x):
        return x - f(x) / Df(x)

    error = tol + 1
    n = 0
    while error > tol:
        n += 1
        if n > max_iter:
            raise Exception('达到最大迭代次数但未收敛')
        y = q(x)
        error = jnp.abs(x - y)
        x = y
        print(f'迭代 {n}, 误差 = {error:.5f}')
    return x.item()

许多库都实现了一维牛顿法，包括SciPy，所以这里的代码仅作说明用途。

（话虽如此，当我们想要使用自动微分或GPU加速等技术来应用牛顿法时，了解如何自己实现牛顿法会很有帮助。）

7.3.2. 在寻找不动点中的应用#

现在再次考虑索洛不动点计算，我们要求解满足\(g(k) = k\)的\(k\)值。

我们可以通过设定\(f(x) := g(x)-x\)将其转换为零点寻找问题。

显然，\(f\)的任何零点都是\(g\)的不动点。

让我们将这个想法应用到索洛问题中

params = create_solow_params()
k_star_approx_newton = newton(f = lambda x: g(x, params) - x, x_0=0.8)

迭代 1, 误差 = 1.27209
迭代 2, 误差 = 0.28180
迭代 3, 误差 = 0.00561
迭代 4, 误差 = 0.00000
迭代 5, 误差 = 0.00000

k_star_approx_newton

1.7846741842265788

结果证实了我们在上面图表中看到的收敛情况：仅需5次迭代就达到了非常精确的结果。

7.4. 多元牛顿法#

在本节中，我们将介绍一个双商品问题，可视化问题，并使用SciPy中的零点查找器和牛顿法来求解这个双商品市场的均衡。

然后，我们将这个概念扩展到一个包含5,000种商品的更大市场，并再次比较这两种方法的性能。

我们将看到使用牛顿法时能获得显著的性能提升。

7.4.1. 双商品市场均衡#

让我们从计算双商品问题的市场均衡开始。

我们考虑一个包含两种相关产品的市场，商品0和商品1，价格向量为\(p = (p_0, p_1)\)

在价格\(p\)下，商品\(i\)的供给为，

\[ q^s_i (p) = b_i \sqrt{p_i} \]

在价格\(p\)下，商品\(i\)的需求为，

\[ q^d_i (p) = \exp(-(a_{i0} p_0 + a_{i1} p_1)) + c_i \]

这里的\(c_i\)、\(b_i\)和\(a_{ij}\)都是参数。

例如，这两种商品可能是通常一起使用的计算机组件，在这种情况下它们是互补品。因此需求取决于两种组件的价格。

超额需求函数为，

\[ e_i(p) = q^d_i(p) - q^s_i(p), \quad i = 0, 1 \]

均衡价格向量\(p^*\)满足\(e_i(p^*) = 0\)。

我们设定

\[\begin{split} A = \begin{bmatrix} a_{00} & a_{01} \\ a_{10} & a_{11} \end{bmatrix}, \qquad b = \begin{bmatrix} b_0 \\ b_1 \end{bmatrix} \qquad \text{和} \qquad c = \begin{bmatrix} c_0 \\ c_1 \end{bmatrix} \end{split}\]

用于这个特定问题。

7.4.1.1. 图形化探索#

由于我们的问题只是二维的，我们可以使用图形分析来可视化并帮助理解这个问题。

我们的第一步是定义超额需求函数

\[\begin{split} e(p) = \begin{bmatrix} e_0(p) \\ e_1(p) \end{bmatrix} \end{split}\]

下面的函数计算给定参数的超额需求

@jax.jit
def e(p, A, b, c):
    return jnp.exp(-A @ p) + c - b * jnp.sqrt(p)

我们的默认参数值将是

\[\begin{split} A = \begin{bmatrix} 0.5 & 0.4 \\ 0.8 & 0.2 \end{bmatrix}, \qquad b = \begin{bmatrix} 1 \\ 1 \end{bmatrix} \qquad \text{和} \qquad c = \begin{bmatrix} 1 \\ 1 \end{bmatrix} \end{split}\]
A = jnp.array([[0.5, 0.4], [0.8, 0.2]])
b = jnp.ones(2)
c = jnp.ones(2)

在价格水平 \(p = (1, 0.5)\) 时，超额需求为

p = jnp.array([1, 0.5])
ex_demand = e(p, A, b, c)

print(
    f"商品0的超额需求为 {ex_demand[0]:.3f} \n"
    f"商品1的超额需求为 {ex_demand[1]:.3f}"
)

商品0的超额需求为 0.497 
商品1的超额需求为 0.699

为了提高计算效率，我们将使用 jax.vmap 提供的向量化能力。这比使用python循环快得多。

# 在p的第一个轴上创建向量化。
e_vectorized_p_1 = jax.vmap(e, in_axes=(0, None, None, None))

# 在p的第二个轴上创建向量化。
e_vectorized = jax.vmap(e_vectorized_p_1, in_axes=(0, None, None, None))

接下来我们在\((p_0, p_1)\)值的网格上绘制两个函数\(e_0\)和\(e_1\)的等高线图和曲面图。

我们将使用以下函数来构建等高线图

def plot_excess_demand(ax, good=0, grid_size=100, grid_max=4, surface=True):
    p_grid = jnp.linspace(0, grid_max, grid_size)

    # 为p_1和p_2的所有组合创建网格
    P1, P2 = jnp.meshgrid(p_grid, p_grid, indexing="ij")

    # 堆叠成形状为(grid_size, grid_size, 2)的数组
    P = jnp.stack([P1, P2], axis=-1)

    # 使用向量化函数一次性计算所有值
    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)
        plt.colorbar(cs1, ax=ax, format="%.6f")

    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'超额需求函数 {good}')
    plt.clabel(ctr1, inline=1, fontsize=13)

这是 \(e_0\) 的图

fig, ax = plt.subplots()
plot_excess_demand(ax, good=0)
plt.show()

findfont: Failed to find font weight normal, now using 600.
findfont: Failed to find font weight normal, now using 600.
findfont: Failed to find font weight normal, now using 600.
_images/9b67c46277d6a6aa987dba5ef17b8fdd968aebb6461cc55c83fc18e13741533f.png

这是 \(e_1\) 的图

fig, ax = plt.subplots()
plot_excess_demand(ax, good=1)
plt.show()
_images/93677a47d6c06cb28534be5b4789fcdb3234623d953bd5317015395a4169d1f5.png

我们看到黑色的零等高线，它告诉我们何时\(e_i(p)=0\)。

对于使得\(e_i(p)=0\)的价格向量\(p\)，我们知道商品\(i\)处于均衡状态（需求等于供给）。

如果这两条等高线在某个价格向量\(p^*\)处相交，那么\(p^*\)就是一个均衡价格向量。

fig, ax = plt.subplots(figsize=(10, 5.7))
for good in (0, 1):
    plot_excess_demand(ax, good=good, surface=False)
plt.show()
_images/2444b843196c3662e6f603cd1a446f41ce16f54f9abc150090b6a96ad3ca7174.png

看起来在 \(p = (1.6, 1.5)\) 附近存在一个均衡点。

7.4.1.2. 使用多维根查找器#

为了更精确地求解 \(p^*\)，我们使用 scipy.optimize 中的零点查找算法。

我们以 \(p = (1, 1)\) 作为初始猜测值。

init_p = jnp.ones(2)

这个算法使用改进的Powell方法来寻找零点

%%time
solution = root(lambda p: e(p, A, b, c), init_p, method="hybr")

CPU times: user 1.57 ms, sys: 42 μs, total: 1.61 ms
Wall time: 1.39 ms

这是得到的值

p = solution.x
p

array([1.57080182, 1.46928838])

这个结果看起来和我们从图中观察到的猜测很接近。我们可以把它代回到 \(e\) 中验证 \(e(p) \approx 0\)

e_p = jnp.max(jnp.abs(e(p, A, b, c)))
e_p.item()

2.0383694732117874e-13

这确实是一个很小的误差。

7.4.1.3. 添加梯度信息#

在许多情况下，对于应用于光滑函数的零点查找算法，提供函数的雅可比矩阵可以带来更好的收敛性质。

这里我们手动计算雅可比矩阵的元素

\[\begin{split} J(p) = \begin{bmatrix} \frac{\partial e_0}{\partial p_0}(p) & \frac{\partial e_0}{\partial p_1}(p) \\ \frac{\partial e_1}{\partial p_0}(p) & \frac{\partial e_1}{\partial p_1}(p) \end{bmatrix} \end{split}\]
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 * 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 jnp.array(J)

%%time
solution = root(
    lambda p: e(p, A, b, c),
    init_p,
    jac = lambda p: jacobian_e(p, A, b, c),
    method="hybr",
)

CPU times: user 214 ms, sys: 2 μs, total: 214 ms
Wall time: 205 ms

现在的解更加精确了（尽管在这个低维问题中，差异非常小）：

p = solution.x
e_p = jnp.max(jnp.abs(e(p, A, b, c)))
e_p.item()

1.1102230246251565e-15

7.4.1.4. 使用牛顿法#

现在让我们使用牛顿法来计算均衡价格，采用多变量版本的牛顿法

(7.6)#\[p_{n+1} = p_n - J_e(p_n)^{-1} e(p_n)\]

这是(7.5)的多变量版本

（这里的\(J_e(p_n)\)是在\(p_n\)处计算的\(e\)的雅可比矩阵。）

迭代从价格向量\(p_0\)的某个初始猜测开始。

在这里，我们不手动编写雅可比矩阵，而是使用jax库中的jacobian()函数来自动求导并计算雅可比矩阵。

只需稍作修改，我们就可以将我们之前的尝试推广到多维问题

def newton(f, x_0, tol=1e-5, max_iter=10):
    x = x_0
    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:
        n += 1
        if n > max_iter:
            raise Exception("Max iteration reached without convergence")
        y = q(x)
        if any(jnp.isnan(y)):
            raise Exception("Solution not found with NaN generated")
        error = jnp.linalg.norm(x - y)
        x = y
        print(f'迭代 {n}, 误差 = {error:.5f}')
    print("\n" + f"Result = {x} \n")
    return x

我们发现算法在4步内终止

%%time
p = newton(lambda p: e(p, A, b, c), init_p)

迭代 1, 误差 = 0.62515
迭代 2, 误差 = 0.11152
迭代 3, 误差 = 0.00258
迭代 4, 误差 = 0.00000

Result = [1.57080182 1.46928838] 

CPU times: user 374 ms, sys: 1.73 ms, total: 376 ms
Wall time: 241 ms

e_p = jnp.max(jnp.abs(e(p, A, b, c)))
e_p.item()

1.461053500406706e-13

结果非常准确。

在开销较大的情况下，速度并不比经过优化的scipy函数更快。

7.4.2. 高维问题#

我们的下一步是研究一个有3,000种商品的大型市场。

超额需求函数基本相同，但现在矩阵 \(A\) 是 \(3000 \times 3000\) 的，参数向量 \(b\) 和 \(c\) 是 \(3000 \times 1\) 的。

dim = 3000

# 创建JAX随机密钥
key = jax.random.key(0)

# 创建随机矩阵A并将列归一化使其和为1
A = jax.random.uniform(key, (dim, dim))
s = jnp.sum(A, axis=0)
A = A / s

# 设置b和c
b = jnp.ones(dim)
c = jnp.ones(dim)

这是我们的初始条件

init_p = jnp.ones(dim)

%%time
p = newton(lambda p: e(p, A, b, c), init_p)

迭代 1, 误差 = 23.22262
迭代 2, 误差 = 3.94537
迭代 3, 误差 = 0.08500
迭代 4, 误差 = 0.00004
迭代 5, 误差 = 0.00000

Result = [1.50723773 1.51041603 1.50134795 ... 1.49941629 1.49033692 1.49666807] 

CPU times: user 17.3 s, sys: 616 ms, total: 17.9 s
Wall time: 5.58 s

e_p = jnp.max(jnp.abs(e(p, A, b, c)))
e_p.item()

6.661338147750939e-16

在相同的容差条件下，我们比较牛顿法与SciPy的root函数的运行时间和精确度

%%time
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,
)

CPU times: user 28.3 s, sys: 277 ms, total: 28.6 s
Wall time: 24.4 s

p = solution.x
e_p = jnp.max(jnp.abs(e(p, A, b, c)))
e_p.item()

9.209231102147442e-07

7.5. 练习#

练习 7.1

考虑索洛固定点问题的三维扩展，其中

\[\begin{split} A = \begin{bmatrix} 2 & 3 & 3 \\ 2 & 4 & 2 \\ 1 & 5 & 1 \\ \end{bmatrix}, \quad s = 0.2, \quad \alpha = 0.5, \quad \delta = 0.8 \end{split}\]

和之前一样，运动方程为

\[ k_{t+1} = g(k_t) \quad \text{where} \quad g(k) := sAk^\alpha + (1-\delta) k\]

但现在 \(k_t\) 是一个 \(3 \times 1\) 向量。

使用牛顿法求解固定点，初始值如下：

\[\begin{split} \begin{aligned} k1_{0} &= (1, 1, 1) \\ k2_{0} &= (3, 5, 5) \\ k3_{0} &= (50, 50, 50) \end{aligned} \end{split}\]

提示

  • 固定点的计算等价于计算满足 \(g(k^*) - k^* = 0\) 的 \(k^*\)。

  • 如果你对你的解决方案不确定，可以从已解决的示例开始：

\[\begin{split}A = \begin{bmatrix} 2 & 0 & 0 \\ 0 & 2 & 0 \\ 0 & 0 & 2 \\ \end{bmatrix}\end{split}\]

其中 \(s = 0.3\)、\(\alpha = 0.3\) 和 \(\delta = 0.4\)，初始值为：

\[k_0 = (1, 1, 1)\]

结果应该收敛到解析解。

练习 7.2

在这个练习中，让我们尝试不同的初始值，看看牛顿法对不同起始点的反应如何。

让我们定义一个具有以下默认值的三商品问题：

\[\begin{split} A = \begin{bmatrix} 0.2 & 0.1 & 0.7 \\ 0.3 & 0.2 & 0.5 \\ 0.1 & 0.8 & 0.1 \\ \end{bmatrix}, \qquad b = \begin{bmatrix} 1 \\ 1 \\ 1 \end{bmatrix} \qquad \text{和} \qquad c = \begin{bmatrix} 1 \\ 1 \\ 1 \end{bmatrix} \end{split}\]

对于这个练习，使用以下极端价格向量作为初始值：

\[\begin{split} \begin{aligned} p1_{0} &= (5, 5, 5) \\ p2_{0} &= (1, 1, 1) \\ p3_{0} &= (4.5, 0.1, 4) \end{aligned} \end{split}\]

将容差设置为\(1e-15\)以获得更精确的输出。