58. 最优增长 II：使用Numba加速代码#
除了Anaconda中已有的库外，本讲座还需要以下库：
!pip install quantecon
58.1. 概述#
在上一讲中，我们研究了一个具有代表性个体的随机最优增长模型。
我们使用动态规划方法对该模型进行了求解。
在编写代码时，我们注重清晰性和灵活性。
尽管这些特性十分重要，但在实际应用中，灵活性与运行速度之间往往存在权衡。
其原因在于，当代码的灵活性降低时，我们能够更容易地利用模型的结构特征。
（这一点在算法与数学问题中普遍成立：越具体的问题往往具有更强的结构特征，而经过适当设计，这些结构可被有效利用，从而获得更优的结果。）
因此，在本讲中，我们将牺牲一定的灵活性以换取更高的运行速度，采用即时(JIT)编译来加速代码的执行。
下面让我们从导入相关库开始：
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']
import numpy as np
from numba import jit, jit
from quantecon.optimize.scalar_maximization import brent_max
函数brent_max同样被设计用于嵌入JIT编译代码中。
这些函数可作为SciPy中类似函数的替代方案（遗憾的是，SciPy中的相关函数目前尚不支持JIT）。
58.2. 模型#
本节所使用的模型与我们在前一讲关于最优增长的讲授中所讨论的模型相同。
我们从对数型效用函数开始：
并继续作如下假设：
生产函数为 \(f(k) = k^{\alpha}\)；
随机冲击项 \(\xi\) 的分布为 \(\phi\)，其中 \(\xi := \exp(\mu + s \zeta)\) 且 \(\zeta\) 为标准正态分布。
我们将再次使用价值函数迭代(VFI)来求解这个模型。
具体来说，算法保持不变，唯一的区别在于具体的实现方式。
和之前一样，我们会对比本次计算所得结果与真实解。
def v_star(y, α, β, μ):
"""
真实价值函数
"""
c1 = np.log(1 - α * β) / (1 - β)
c2 = (μ + α * np.log(α * β)) / (1 - α)
c3 = 1 / (1 - β)
c4 = 1 / (1 - α * β)
return c1 + c2 * (c3 - c4) + c4 * np.log(y)
def σ_star(y, α, β):
"""
真实最优策略
"""
return (1 - α * β) * y
58.3. 计算#
我们将再次把最优增长模型的基本要素封装在一个类中。
然而，与此前不同的是，我们将使用Numba的 @jitclass装饰器来对该类进行JIT编译。
由于我们计划使用Numba来编译该类，因此需要明确指定数据类型。
在代码中，你将看到一个名为opt_growth_data的列表，该列表定义在类的上方，用于说明这些类型。
与上一讲不同的是，这里我们将生产和效用函数的具体形式直接写入类中，而非保持一般性形式。
也就是说，我们在此牺牲了一定的灵活性，以换取更高的运行速度。
from numba import float64
from numba.experimental import jitclass
opt_growth_data = [
('α', float64), # 生产参数
('β', float64), # 折现因子
('μ', float64), # 冲击的均值参数
('s', float64), # 冲击的尺度参数
('grid', float64[:]), # 网格（数组）
('shocks', float64[:]) # 冲击样本（数组）
]
@jitclass(opt_growth_data)
class OptimalGrowthModel:
def __init__(self,
α=0.4,
β=0.96,
μ=0,
s=0.1,
grid_max=4,
grid_size=120,
shock_size=250,
seed=1234):
self.α, self.β, self.μ, self.s = α, β, μ, s
# 设置网格
self.grid = np.linspace(1e-5, grid_max, grid_size)
# 存储冲击（设置随机种子以确保结果可重复）
np.random.seed(seed)
self.shocks = np.exp(μ + s * np.random.randn(shock_size))
def f(self, k):
"生产函数"
return k**self.α
def u(self, c):
"效用函数"
return np.log(c)
def f_prime(self, k):
"生产函数的一阶导数"
return self.α * (k**(self.α - 1))
def u_prime(self, c):
"效用函数的一阶导数"
return 1/c
def u_prime_inv(self, c):
"效用函数一阶导数的反函数"
return 1/c
Cell In[4], line 38
return k**self.α
^
IndentationError: unexpected indent
该类还包含若干方法，例如u_prime，虽然在当前讲义中尚未使用，但将在后续课程中发挥作用。
58.3.1. 贝尔曼算子#
我们将使用JIT编译来加速贝尔曼算子的计算。
首先，定义一个函数，用于计算在给定状态y下，某一特定消费选择c所对应的价值。该函数基于贝尔曼方程(57.9)。
@jit
def state_action_value(c, y, v_array, og):
"""
贝尔曼方程右侧。
* c是消费
* y是收入
* og是OptimalGrowthModel的一个实例
* v_array表示网格上的值函数猜测
"""
u, f, β, shocks = og.u, og.f, og.β, og.shocks
v = lambda x: np.interp(x, og.grid, v_array)
return u(c) + β * np.mean(v(f(y - c) * shocks))
现在我们可以实现贝尔曼算子，它用于最大化贝尔曼方程的右侧：
@jit
def T(v, og):
"""
贝尔曼算子。
* og 是 OptimalGrowthModel 的一个实例
* v 是一个数组，表示价值函数的猜测值
"""
v_new = np.empty_like(v)
v_greedy = np.empty_like(v)
for i in range(len(og.grid)):
y = og.grid[i]
# 在状态 y 下最大化贝尔曼方程的右侧
result = brent_max(state_action_value, 1e-10, y, args=(y, v, og))
v_greedy[i], v_new[i] = result[0], result[1]
return v_greedy, v_new
我们使用solve_model函数进行迭代直到收敛。
def solve_model(og,
tol=1e-4,
max_iter=1000,
verbose=True,
print_skip=25):
"""
通过迭代贝尔曼算子求解
"""
# 设置迭代循环
v = og.u(og.grid) # 初始条件
i = 0
error = tol + 1
while i < max_iter and error > tol:
v_greedy, v_new = T(v, og)
error = np.max(np.abs(v - v_new))
i += 1
if verbose and i % print_skip == 0:
print(f"第 {i} 次迭代的误差为 {error}。")
v = v_new
if error > tol:
print("未能收敛！")
elif verbose:
print(f"\n在 {i} 次迭代后收敛。")
return v_greedy, v_new
让我们用默认参数计算近似解。
首先创建一个实例：
og = OptimalGrowthModel()
现在我们调用solve_model，使用%%time魔法指令来记录运行时间。
%%time
v_greedy, v_solution = solve_model(og)
你会发现，这比我们的原始实现要快得多。
下面，生成近似策略与真实策略的对比图：
fig, ax = plt.subplots()
ax.plot(og.grid, v_greedy, lw=2,
alpha=0.8, label='近似策略函数')
ax.plot(og.grid, σ_star(og.grid, og.α, og.β), 'k--',
lw=2, alpha=0.8, label='真实策略函数')
ax.legend()
plt.show()
与之前一样，拟合效果非常好 — 这是意料之中的，因为我们没有改变算法。
两种策略之间的最大绝对偏差是
np.max(np.abs(v_greedy - σ_star(og.grid, og.α, og.β)))
58.4. 练习#
练习 58.1
在默认参数设定下，从给定的初始条件 \(v(y) = u(y)\) 开始，对贝尔曼算子进行 20 次迭代，并记录整个迭代过程所耗费的时间。
解答 练习 58.1
设置初始条件：
v = og.u(og.grid)
计时：
%%time
for i in range(20):
v_greedy, v_new = T(v, og)
v = v_new
与非编译版本的价值函数迭代的用时相比，JIT编译的代码通常快一个数量级。
练习 58.2
将最优增长模型修改为采用CRRA效用函数的设定：
设定γ = 1.5为默认值，并保持其他模型设定不变。
（注意，jitclass目前不支持类继承，因此你需要复制原有的类，并相应修改相关的参数与方法。）
计算最优策略的估计值，并绘制其图像。将所得图像与第一讲最优增长模型中对应练习的图表进行比较。
同时，对比两种实现的运行时间。
解答 练习 58.2
这是CRRA版本的OptimalGrowthModel：
from numba import float64
from numba.experimental import jitclass
opt_growth_data = [
('α', float64), # 生产参数
('β', float64), # 折现因子
('μ', float64), # 冲击的均值参数
('γ', float64), # 偏好参数
('s', float64), # 冲击的尺度参数
('grid', float64[:]), # 网格（数组）
('shocks', float64[:]) # 冲击样本（数组）
]
@jitclass(opt_growth_data)
class OptimalGrowthModel_CRRA:
def __init__(self,
α=0.4,
β=0.96,
μ=0,
s=0.1,
γ=1.5,
grid_max=4,
grid_size=120,
shock_size=250,
seed=1234):
self.α, self.β, self.γ, self.μ, self.s = α, β, γ, μ, s
# 设置网格
self.grid = np.linspace(1e-5, grid_max, grid_size)
# 存储冲击（设置随机种子以确保结果可重复）
np.random.seed(seed)
self.shocks = np.exp(μ + s * np.random.randn(shock_size))
def f(self, k):
"生产函数"
return k**self.α
def u(self, c):
"效用函数"
return c**(1 - self.γ) / (1 - self.γ)
def f_prime(self, k):
"生产函数的一阶导数"
return self.α * (k**(self.α - 1))
def u_prime(self, c):
"效用函数的一阶导数"
return c**(-self.γ)
def u_prime_inv(self, c):
"效用函数一阶导数的反函数"
return c**(-1 / self.γ)
创建一个实例：
og_crra = OptimalGrowthModel_CRRA()
调用solve_model，使用%%time魔术命令来记录运行时间。
%%time
v_greedy, v_solution = solve_model(og_crra)
以下是得到的策略图：
fig, ax = plt.subplots()
ax.plot(og.grid, v_greedy, lw=2,
alpha=0.6, label='近似价值函数')
ax.legend(loc='lower right')
plt.show()
这与我们在练习中使用非jit代码得到的答案相符，但执行时间快了一个数量级。
练习 58.3
在本练习中，我们回到最初的对数型效用函数设定。
当给定最优消费政策 \(\sigma\) 后，收入的动态演化如下：
下图展示了该序列在三种不同贴现因子（因而对应三种不同政策）下的模拟结果，每个序列包含 100 个样本点。
在每个序列中，初始条件都是 \(y_0 = 0.1\)。
贴现因子分别为discount_factors = (0.8, 0.9, 0.98)。
我们还通过设置s = 0.05稍微降低了冲击的幅度。
除此之外，参数和原始设定与前面讨论的对数线性模型相同。
注意，更有耐心的个体通常拥有更高的财富。
请在保持随机性结构的前提下，复现该图像。
解答 练习 58.3
参考答案：
def simulate_og(σ_func, og, y0=0.1, ts_length=100):
'''
根据消费策略σ计算时间序列。
'''
y = np.empty(ts_length)
ξ = np.random.randn(ts_length-1)
y[0] = y0
for t in range(ts_length-1):
y[t+1] = (y[t] - σ_func(y[t]))**og.α * np.exp(og.μ + og.s * ξ[t])
return y
fig, ax = plt.subplots()
for β in (0.8, 0.9, 0.98):
og = OptimalGrowthModel(β=β, s=0.05)
v_greedy, v_solution = solve_model(og, verbose=False)
# 定义最优策略函数
σ_func = lambda x: np.interp(x, og.grid, v_greedy)
y = simulate_og(σ_func, og)
ax.plot(y, lw=2, alpha=0.6, label=rf'$\beta = {β}$')
ax.legend(loc='lower right')
plt.show()