程序设计与计算思维¶

Computer Programming and Computational Thinking

第 22 讲：二维的平流与扩散¶

2025—2026学年度春季学期

清华大学 韩文弢

问题引入¶

上一讲研究了一维空间中的平流和扩散，将温度 $T$ 视为时间 $t$ 和单一空间坐标 $x$ 的函数 $T(t, x)$，并引入了偏微分方程（PDE）的概念。

但现实世界是多维的。海洋和大气中的温度、盐度、湿度等物理量不仅随时间变化，还在两个甚至三个空间维度上分布。例如，要模拟洋流如何将赤道的热量输送到极地，就必须考虑经度和纬度两个方向上的温度变化。

从一维扩展到二维，核心的计算思想并没有本质改变——仍然是离散化和时间步进。但需要引入两个重要的新工具：

  • 模板（stencil）：一种统一描述邻域操作的框架，将微分算子的离散化表示为一个小矩阵
  • 幽灵单元（ghost cells）：一种优雅处理边界条件的技巧，使代码无需对边界做特殊判断

本讲将首先介绍这两个工具，然后回顾上一讲的一维平流-扩散模型，最后将所有内容推广到二维，模拟真实的海洋洋流对热量的传输。

准备工作¶

In [1]:
import numpy as np
import matplotlib.pyplot as plt
from ipywidgets import interact, FloatSlider, IntSlider

import time
from numba import njit

plt.rcParams['font.family'] = ['Noto Sans', 'Noto Sans CJK SC']

模板计算¶

在上一讲中，一维平流和扩散的数值方法都涉及对相邻网格点的操作：

  • 平流（上游差分）：T_new[i] = T[i] + dt * U * (T[i-1] - T[i]) / dx
  • 扩散：T_new[i] = T[i] + dt * D * (T[i+1] - 2*T[i] + T[i-1]) / dx**2

这些操作的共同特点是：每个网格点的新值取决于它自身和相邻点的旧值。在二维情况下，"相邻"的含义扩展为上下左右四个方向，形成一个 3×3 的邻域。

模板 (stencil) 就是一种将这种邻域操作系统化的方法：定义一个小的权重矩阵（模板），将其与每个点周围的邻域逐元素相乘再求和，就得到了该点处微分算子的近似值。

下面从最基础的二维数组操作开始，逐步构建模板计算的工具箱。

In [2]:
np.random.seed(42)
data = np.random.randint(1, 10, size=(6, 7))
print('原始数组 data：')
print(data)
print(f'形状: {data.shape}')

原始数组 data：
[[7 4 8 5 7 3 7]
 [8 5 4 8 8 3 6]
 [5 2 8 6 2 5 1]
 [6 9 1 3 7 4 9]
 [3 5 3 7 5 9 7]
 [2 4 9 2 9 5 2]]
形状: (6, 7)

扩展数组与幽灵单元¶

当模板窗口滑动到数组边缘时，朴素的做法需要进行特殊处理（回忆一维情况）。

可以通过在原始数组周围添加一圈幽灵单元来实现相同效果。幽灵单元提供了边界值，使得模板操作无需对边界做特殊处理。

下面将原始数据嵌入一个更大的数组中，四周各留一行/列作为幽灵单元：

In [3]:
rows, cols = data.shape
A = np.zeros((rows + 2, cols + 2), dtype=data.dtype)

A[1:-1, 1:-1] = data

print('扩展数组 A（四周为幽灵单元，初始值为0）：')
print(A)
print(f'形状: {A.shape}')
print(f'\nA[1,1] = {A[1,1]}  (对应 data[0,0] = {data[0,0]})')
print(f'A[0,0] = {A[0,0]}  (幽灵单元)')

扩展数组 A（四周为幽灵单元，初始值为0）：
[[0 0 0 0 0 0 0 0 0]
 [0 7 4 8 5 7 3 7 0]
 [0 8 5 4 8 8 3 6 0]
 [0 5 2 8 6 2 5 1 0]
 [0 6 9 1 3 7 4 9 0]
 [0 3 5 3 7 5 9 7 0]
 [0 2 4 9 2 9 5 2 0]
 [0 0 0 0 0 0 0 0 0]]
形状: (8, 9)

A[1,1] = 7  (对应 data[0,0] = 7)
A[0,0] = 0  (幽灵单元)

邻域：3×3 窗口¶

模板操作需要一个邻域——以当前点为中心的 3×3 窗口。对于数组中的每个位置 (i, j)，邻域为 A[i-1:i+2, j-1:j+2]。

下面展示几个具有代表性的点的 3×3 邻域（注意边缘处幽灵单元的值为 0）：

In [4]:
for i, j, name in [
    (1, 1, '左上角'),
    (1, 3, '上边界一点'),
    (2, 3, '内部一点'),
    (2, 7, '右边界一点'),
]:
    neighborhood = A[i-1:i+2, j-1:j+2]
    print(f'({i}, {j})（{name}）的邻域 =')
    print(neighborhood)
    print()

(1, 1)（左上角）的邻域 =
[[0 0 0]
 [0 7 4]
 [0 8 5]]

(1, 3)（上边界一点）的邻域 =
[[0 0 0]
 [4 8 5]
 [5 4 8]]

(2, 3)（内部一点）的邻域 =
[[4 8 5]
 [5 4 8]
 [2 8 6]]

(2, 7)（右边界一点）的邻域 =
[[3 7 0]
 [3 6 0]
 [5 1 0]]

模板¶

定义一个 3×3 的模板矩阵。这是离散拉普拉斯算子的一种形式：

$$\text{stencil} = \begin{pmatrix} 0 & -1 & 0 \\ -1 & 4 & -1 \\ 0 & -1 & 0 \end{pmatrix}$$

将模板与邻域逐元素相乘后求和，就得到了拉普拉斯算子在该点的近似值。

模板操作等价于计算：

$$\nabla^2 T_{i,j} \approx \frac{T_{i+1,j} + T_{i-1,j} + T_{i,j+1} + T_{i,j-1} - 4T_{i,j}}{(\Delta x)^2}$$

In [5]:
stencil = np.array([[ 0, -1,  0],
                    [-1,  4, -1],
                    [ 0, -1,  0]])

result = np.zeros_like(data)
for i in range(rows):
    for j in range(cols):
        neighborhood = A[i:i+3, j:j+3]
        result[i, j] = np.sum(neighborhood * stencil)

print('模板作用结果（与原始数据同大小）：')
print(result)
print(f'形状: {result.shape}  (与 data 的形状 {data.shape} 相同)')

模板作用结果（与原始数据同大小）：
[[ 16  -4  19  -3  12  -5  19]
 [ 15   2 -13   9  12 -10  13]
 [  4 -19  19   3 -18  10 -16]
 [  7  22 -19  -9  14 -14  24]
 [ -1   1 -10  15 -12  15   8]
 [  1   0  27 -17  24   0  -4]]
形状: (6, 7)  (与 data 的形状 (6, 7) 相同)

边界条件¶

注意结果与原始数据大小相同，模板在边缘处也"正常工作"了——因为幽灵单元默认值为 0。这对应了零边界条件（Dirichlet）：假设边界外的值为 0。

但还有其他常见的边界条件：

边界条件 含义 幽灵单元设置 应用场景
零值（Dirichlet） 边界外值为 0 幽灵单元 = 0 固定温度的边界（如恒温槽壁面）
周期性（Periodic） 数组左右/上下环绕 幽灵单元 = 对侧的值 环形域（如全球经度方向）
零导数（Neumann） 边界处导数为 0 幽灵单元 = 相邻内部值 无通量边界（如海洋-陆地交界）
In [6]:
B_periodic = np.zeros((rows + 2, cols + 2), dtype=int)
B_periodic[1:-1, 1:-1] = data

# 周期性边界条件
B_periodic[0, :] = B_periodic[rows, :]  # 上幽灵行 = 底部数据行
B_periodic[-1, :] = B_periodic[1, :]    # 下幽灵行 = 顶部数据行
B_periodic[:, 0] = B_periodic[:, cols]  # 左幽灵列 = 右侧数据列
B_periodic[:, -1] = B_periodic[:, 1]    # 右幽灵列 = 左侧数据列

print('周期性边界条件下的扩展数组：')
print(B_periodic)
print(f'\n验证周期性：')
print(f'  B_periodic[0, 1] = {B_periodic[0, 1]} == data[-1, 0] = {data[-1, 0]}  (底部环绕到顶部)')
print(f'  B_periodic[1, 0] = {B_periodic[1, 0]} == data[0, -1] = {data[0, -1]}  (右侧环绕到左侧)')

周期性边界条件下的扩展数组：
[[2 2 4 9 2 9 5 2 2]
 [7 7 4 8 5 7 3 7 7]
 [6 8 5 4 8 8 3 6 8]
 [1 5 2 8 6 2 5 1 5]
 [9 6 9 1 3 7 4 9 6]
 [7 3 5 3 7 5 9 7 3]
 [2 2 4 9 2 9 5 2 2]
 [7 7 4 8 5 7 3 7 7]]

验证周期性：
  B_periodic[0, 1] = 2 == data[-1, 0] = 2  (底部环绕到顶部)
  B_periodic[1, 0] = 7 == data[0, -1] = 7  (右侧环绕到左侧)
In [7]:
result_periodic = np.zeros_like(data)
for i in range(rows):
    for j in range(cols):
        neighborhood = B_periodic[i:i+3, j:j+3]
        result_periodic[i, j] = np.sum(neighborhood * stencil)

print('周期性边界条件下的模板结果：')
print(result_periodic)

周期性边界条件下的模板结果：
[[  7  -8  10  -5   3 -10  10]
 [  9   2 -13   9  12 -10   5]
 [  3 -19  19   3 -18  10 -21]
 [ -2  22 -19  -9  14 -14  18]
 [ -8   1 -10  15 -12  15   5]
 [ -8  -4  19 -22  17  -3 -13]]
In [8]:
B_neumann = np.zeros((rows + 2, cols + 2), dtype=int)
B_neumann[1:-1, 1:-1] = data

B_neumann[0, :] = B_neumann[1, :]
B_neumann[-1, :] = B_neumann[-2, :]
B_neumann[:, 0] = B_neumann[:, 1]
B_neumann[:, -1] = B_neumann[:, -2]

result_neumann = np.zeros_like(data)
for i in range(rows):
    for j in range(cols):
        neighborhood = B_neumann[i:i+3, j:j+3]
        result_neumann[i, j] = np.sum(neighborhood * stencil)

print('零导数边界条件下的模板结果：')
print(result_neumann)

零导数边界条件下的模板结果：
[[  2  -8  11  -8   5  -8   5]
 [  7   2 -13   9  12 -10   7]
 [ -1 -19  19   3 -18  10 -17]
 [  1  22 -19  -9  14 -14  15]
 [ -4   1 -10  15 -12  15   1]
 [ -3  -4  18 -19  15  -5  -8]]
In [9]:
fig, axes = plt.subplots(2, 2, figsize=(12, 10))

ax = axes[0, 0]
im = ax.imshow(data, cmap='RdBu_r', vmax=10, vmin=-30)
ax.set_title('原始数据')
plt.colorbar(im, ax=ax)

ax = axes[0, 1]
im = ax.imshow(result, cmap='RdBu_r', vmax=10, vmin=-30)
ax.set_title('零值边界（Dirichlet）')
plt.colorbar(im, ax=ax)

ax = axes[1, 0]
im = ax.imshow(result_periodic, cmap='RdBu_r', vmax=10, vmin=-30)
ax.set_title('周期性边界（Periodic）')
plt.colorbar(im, ax=ax)

ax = axes[1, 1]
im = ax.imshow(result_neumann, cmap='RdBu_r', vmax=10, vmin=-30)
ax.set_title('零导数边界（Neumann）')
plt.colorbar(im, ax=ax)

plt.tight_layout()
plt.show()
No description has been provided for this image

一维平流-扩散问题回顾¶

在将平流-扩散方程推广到二维之前，先回顾上一讲的核心结果。

上一讲考虑了定义在区间 $[0, L_x]$ 上的温度场 $T(t, x)$，满足以下 PDE：

$$\frac{\partial T}{\partial t} = -U \frac{\partial T}{\partial x} + \kappa \frac{\partial^2 T}{\partial x^2}$$

其中 $U$ 是均匀流速（平流），$\kappa$ 是热扩散率（扩散）。

数值方法的关键步骤：

  1. 空间离散化：将区间等分为 $N_x$ 个单元，在每个单元中心存储温度值，得到向量 $\mathbf{T}^n = (T^n_1, \ldots, T^n_{N_x})$
  2. 空间导数的近似：
    • 一阶导数（平流项）：中心差分 $\frac{\partial T}{\partial x} \approx \frac{T_{i+1} - T_{i-1}}{2\delta x}$
    • 二阶导数（扩散项）：$\frac{\partial^2 T}{\partial x^2} \approx \frac{T_{i+1} - 2T_i + T_{i-1}}{\delta x^2}$
  3. 时间步进：前向欧拉法 $T^{n+1}_i = T^n_i + \delta t \cdot (\text{平流} + \text{扩散})$
  4. 边界条件：周期性边界条件（$T_0 = T_{N_x}$，$T_{N_x+1} = T_1$）

在二维情况下，温度场变为 $T(t, x, y)$，需要在两个方向上分别近似空间导数。每个方向的操作与一维完全相同——这正是模板方法的优势所在：只需定义好二维模板，就可以统一处理任意方向的导数。

下面将看到，二维扩散项就是将 $x$ 方向和 $y$ 方向的二阶导数模板叠加，得到的就是前面介绍的离散拉普拉斯算子模板。

二维平流-扩散方程¶

二维平流-扩散方程描述温度场 $T = T(x, y, t)$ 在速度场 $\vec{u}(x,y) = (u, v)$ 中的演变：

$$\frac{\partial T(x,y,t)}{\partial t} = -u(x,y) \frac{\partial T}{\partial x} - v(x,y) \frac{\partial T}{\partial y} + \kappa \left( \frac{\partial^{2} T}{\partial x^{2}} + \frac{\partial^{2} T}{\partial y^{2}} \right),$$

其中 $\kappa$ 是热扩散率。$x$ 为经度方向（西→东为正），$y$ 为纬度方向（南→北为正）。

矢量形式¶

利用梯度 $\nabla$ 和拉普拉斯 $\nabla^2$ 算子，方程可简写为：

$$\frac{\partial T}{\partial t} = -\vec{u} \cdot \nabla T + \kappa \nabla^{2} T$$

由于海水近似不可压缩（$\nabla \cdot \vec{u} = 0$），利用乘积法则可以改写为通量形式：

$$\frac{\partial T}{\partial t} = -\nabla \cdot (\vec{u}T) + \kappa \nabla^{2} T$$

二维平流-扩散的模板¶

离散化一阶偏导数¶

$x$ 方向的一阶偏导数使用中心差分：

$$\frac{\partial T(x_i, y_j, t_n)}{\partial x} \approx \frac{T_{i+1,\,j}^{n} - T_{i-1,\,j}^{n}}{2 \Delta x}$$

对应的 $x$-梯度模板和 $y$-梯度模板如下：

In [10]:
fig, axes = plt.subplots(2, 2, figsize=(10, 8))

xgrad = np.array([[0, 0, 0], [-1, 0, 1], [0, 0, 0]]) / 2.0
ygrad = np.array([[0, -1, 0], [0, 0, 0], [0, 1, 0]]) / 2.0
xdiff = np.array([[0, 0, 0], [1, -2, 1], [0, 0, 0]])
ydiff = np.array([[0, 1, 0], [0, -2, 0], [0, 1, 0]])

for ax, kernel, title in [
    (axes[0, 0], xgrad, 'x-梯度模板 ∂T/∂x'),
    (axes[0, 1], ygrad, 'y-梯度模板 ∂T/∂y'),
    (axes[1, 0], xdiff, 'x-二阶导数模板 ∂²T/∂x²'),
    (axes[1, 1], ydiff, 'y-二阶导数模板 ∂²T/∂y²'),
]:
    vmax = max(abs(kernel.min()), abs(kernel.max()))
    im = ax.imshow(kernel, cmap='RdBu_r', vmin=-vmax, vmax=vmax)
    ax.set_title(title)
    ax.set_xticks([])
    ax.set_yticks([])
    for ii in range(3):
        for jj in range(3):
            ax.text(jj, ii, f'{kernel[ii, jj]:.1f}', ha='center', va='center',
                    color='black' if abs(kernel[ii, jj]) < vmax * 0.5 else 'white')
    plt.colorbar(im, ax=ax, fraction=0.046)

plt.tight_layout()
plt.show()
No description has been provided for this image

离散化扩散项¶

扩散项的离散化是两个方向二阶导数的叠加——将上面两个二阶导数模板相加，就得到了前面介绍的离散拉普拉斯算子：

$$\kappa \left( \frac{\partial^{2} T}{\partial x^{2}} + \frac{\partial^{2} T}{\partial y^{2}} \right) \approx \kappa \left( \frac{T_{i+1,j} - 2T_{i,j} + T_{i-1,j}}{(\Delta x)^2} + \frac{T_{i,j+1} - 2T_{i,j} + T_{i,j-1}}{(\Delta y)^2} \right)$$

无通量边界条件¶

在 $x$ 和 $y$ 边界上施加无通量条件（即上面介绍的 Neumann 边界条件）：$u\frac{\partial T}{\partial x} = \kappa \frac{\partial T}{\partial x} = 0$。

通过幽灵单元实现：$T_{1,j} = T_{2,j}$，$T_{N_x,j} = T_{N_x-1,j}$（$y$ 方向类似）。这正是前面零导数边界条件的做法。

In [11]:
class Grid:
    """网格对象"""
    
    def __init__(self, N, L):
        """
        N: 一个维度上的网格数量
        L: 一个维度的长度
        """
        self.N = N
        self.L = L
        self.dx = L / N
        self.dy = L / N
        
        self.x = np.arange(-self.dx/2, L + self.dx, self.dx)
        self.y = np.arange(-L - self.dy/2, L + self.dy, self.dy)
        self.Nx = len(self.x)
        self.Ny = len(self.y)
        
        self.X, self.Y = np.meshgrid(self.x, self.y)


class OceanModel:
    """海洋模式"""
    
    def __init__(self, grid, kappa=1e4, u=None, v=None):
        self.G = grid
        self.kappa = kappa
        if u is None:
            self.u = np.zeros((grid.Ny, grid.Nx))
        else:
            self.u = u
        if v is None:
            self.v = np.zeros((grid.Ny, grid.Nx))
        else:
            self.v = v


def update_ghost_cells(T):
    T[0, :] = T[1, :]
    T[-1, :] = T[-2, :]
    T[:, 0] = T[:, 1]
    T[:, -1] = T[:, -2]


def advect(T, u, v, dx, dy):
    Ny, Nx = T.shape
    tendency = np.zeros((Ny, Nx))
    tendency[1:-1, 1:-1] = -(
        u[1:-1, 1:-1] * (T[1:-1, 2:] - T[1:-1, :-2]) / (2 * dx) +
        v[1:-1, 1:-1] * (T[2:, 1:-1] - T[:-2, 1:-1]) / (2 * dy)
    )
    return tendency


def diffuse(T, kappa, dx, dy):
    Ny, Nx = T.shape
    tendency = np.zeros((Ny, Nx))
    tendency[1:-1, 1:-1] = kappa * (
        (T[1:-1, 2:] - 2*T[1:-1, 1:-1] + T[1:-1, :-2]) / dx**2 +
        (T[2:, 1:-1] - 2*T[1:-1, 1:-1] + T[:-2, 1:-1]) / dy**2
    )
    return tendency


def timestep(T, model):
    update_ghost_cells(T)
    adv = advect(T, model.u, model.v, model.G.dx, model.G.dy)
    diff = diffuse(T, model.kappa, model.G.dx, model.G.dy)
    T[1:-1, 1:-1] += model.dt * (adv[1:-1, 1:-1] + diff[1:-1, 1:-1])
    return T

双环流海洋速度场¶

速度场采用经典的风驱动环流解析解。该解通过边界层渐近方法求解四阶 PDE 得到：

$$- \epsilon_{M} \hat{\nabla}^{4} \hat{\Psi} + \frac{\partial \hat{\Psi}}{\partial \hat{x}} = \nabla \times \hat{\tau} \mathbf{z}$$

其中 $\epsilon_M = \nu / (\beta L^3)$ 是无量纲参数。

In [12]:
def double_gyre(G, beta=2e-11, tau0=0.1, rho0=1e3, nu=1e5, H=1000.):
    epsM = nu / (beta * G.L**3)
    eps = epsM**(1/3)
    
    x = G.X / G.L
    y = G.Y / G.L
    
    psi = np.pi * np.sin(np.pi * y) * (
        1 - x - np.exp(-x/(2*eps)) * (
            np.cos(np.sqrt(3)*x/(2*eps)) +
            (1/np.sqrt(3)) * np.sin(np.sqrt(3)*x/(2*eps))
        ) + eps * np.exp((x - 1)/eps)
    )
    
    dpsi_dy, dpsi_dx = np.gradient(psi, G.dy, G.dx)
    u = dpsi_dy
    v = -dpsi_dx
    
    scale = (tau0/rho0) / (beta * G.L * H)
    u = scale * u
    v = scale * v
    
    u[0, :] = 0; v[0, :] = 0
    u[-1, :] = 0; v[-1, :] = 0
    u[:, 0] = 0; v[:, 0] = 0
    u[:, -1] = 0; v[:, -1] = 0
    
    return u, v
In [13]:
N_grid = 30
L_domain = 6e6
G = Grid(N_grid, L_domain)
u_dg, v_dg = double_gyre(G)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

speed = np.sqrt(u_dg**2 + v_dg**2)
im1 = ax1.pcolormesh(G.X * 1e-3, G.Y * 1e-3, speed, cmap='viridis', shading='auto')
plt.colorbar(im1, ax=ax1, label='流速 [m/s]')
ax1.set_xlabel('经度距离 [km]')
ax1.set_ylabel('纬度距离 [km]')
ax1.set_title('双环流速度场大小')

skip = 3
ax2.quiver(G.X[::skip, ::skip] * 1e-3, G.Y[::skip, ::skip] * 1e-3,
           u_dg[::skip, ::skip], v_dg[::skip, ::skip],
           speed[::skip, ::skip], cmap='viridis', alpha=0.8)
ax2.set_xlabel('经度距离 [km]')
ax2.set_ylabel('纬度距离 [km]')
ax2.set_title('双环流速度矢量场')

plt.tight_layout()
plt.show()
No description has been provided for this image

初始条件¶

纬度线性温度¶

模拟从赤道到极地的线性温度分布：

$$T_0(x, y) = 0.5 \left(1 - \frac{y}{L}\right)$$

热点初始条件¶

在中纬度放置一个温度为 1°C 的热点区域。

In [14]:
def linear_temperature(G):
    return 0.5 * (1 - G.Y / G.L)


def init_box(G, value=1.0, nx=2, ny=2, xspan=False, yspan=False):
    T = np.zeros((G.Ny, G.Nx))
    cy, cx = G.Ny // 2, G.Nx // 2
    T[cy-ny:cy+ny, cx-nx:cx+nx] = value
    if xspan:
        T[cy-ny:cy+ny, :] = value
    if yspan:
        T[:, cx-nx:cx+nx] = value
    return T


fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(16, 4))

G_test = Grid(30, 6e6)
im1 = ax1.pcolormesh(G_test.X*1e-3, G_test.Y*1e-3, linear_temperature(G_test),
                      cmap='RdBu_r', vmin=-1, vmax=1, shading='auto')
ax1.set_title('线性温度')
ax1.set_ylabel('纬度 [km]')
plt.colorbar(im1, ax=ax1, label='T [°C]')

im2 = ax2.pcolormesh(G_test.X*1e-3, G_test.Y*1e-3, init_box(G_test),
                      cmap='RdBu_r', vmin=-1, vmax=1, shading='auto')
ax2.set_title('中心热点')
plt.colorbar(im2, ax=ax2, label='T [°C]')

im3 = ax3.pcolormesh(G_test.X*1e-3, G_test.Y*1e-3, init_box(G_test, xspan=True),
                      cmap='RdBu_r', vmin=-1, vmax=1, shading='auto')
ax3.set_title('纬向热点')
plt.colorbar(im3, ax=ax3, label='T [°C]')

for ax in [ax1, ax2, ax3]:
    ax.set_xlabel('经度 [km]')

plt.tight_layout()
plt.show()
No description has been provided for this image

模拟洋流的热量传输¶

使用双环流速度场运行二维平流-扩散模拟。调整参数观察不同的物理效果：

  • U 倍数：放大/缩小洋流速度
  • 扩散率 $\kappa$：控制热量扩散的速率
In [15]:
def run_simulation(G, kappa, U_mult, IC, dt=12*3600, n_steps=500):
    u, v = double_gyre(G)
    model = OceanModel(G, kappa=kappa, u=u * U_mult, v=v * U_mult)
    model.dt = dt
    
    T = IC.copy()
    snapshots = [T.copy()]
    times = [0]
    
    cfl = np.max(np.sqrt(model.u**2 + model.v**2)) * dt / G.dx
    
    for step in range(n_steps):
        T = timestep(T, model)
        if (step + 1) % (n_steps // 10) == 0:
            snapshots.append(T.copy())
            times.append((step + 1) * dt / (3600 * 24))
    
    return T, snapshots, times, cfl


G_sim = Grid(30, 6e6)
IC_sim = init_box(G_sim, xspan=True)

T_final, snapshots, times, cfl = run_simulation(
    G_sim, kappa=1e4, U_mult=1.0, IC=IC_sim, n_steps=500
)

print(f'CFL = {cfl:.3f}')
print(f'模拟天数: {times[-1]:.0f} 天')

n_panels = min(len(snapshots), 6)
indices = np.linspace(0, len(snapshots)-1, n_panels, dtype=int)

fig, axes = plt.subplots(2, 3, figsize=(16, 10))
for idx, ax in zip(indices, axes.flat):
    im = ax.pcolormesh(G_sim.X*1e-3, G_sim.Y*1e-3, snapshots[idx],
                        cmap='RdBu_r', vmin=-0.5, vmax=1.0, shading='auto')
    ax.set_title(f't = {times[idx]:.0f} 天')
    ax.set_xlabel('经度 [km]')
    ax.set_ylabel('纬度 [km]')
    plt.colorbar(im, ax=ax, label='T [°C]')

plt.tight_layout()
plt.show()

CFL = 0.000
模拟天数: 250 天
No description has been provided for this image
In [16]:
def plot_simulation_interactive(U_mult=1.0, kappa=1e4, n_steps=300):
    G = Grid(30, 6e6)
    IC = init_box(G, xspan=True)
    T_final, snaps, day_list, cfl = run_simulation(
        G, kappa=kappa, U_mult=U_mult, IC=IC, n_steps=n_steps
    )
    
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))
    
    im1 = ax1.pcolormesh(G.X*1e-3, G.Y*1e-3, snaps[-1],
                          cmap='RdBu_r', vmin=-0.5, vmax=1.0, shading='auto')
    ax1.set_xlabel('经度 [km]')
    ax1.set_ylabel('纬度 [km]')
    ax1.set_title(f'温度场 (t = {day_list[-1]:.0f} 天)')
    plt.colorbar(im1, ax=ax1, label='T [°C]')
    
    anomaly = snaps[-1] - IC
    im2 = ax2.pcolormesh(G.X*1e-3, G.Y*1e-3, anomaly,
                          cmap='RdBu_r', vmin=-1.0, vmax=1.0, shading='auto')
    ax2.set_xlabel('经度 [km]')
    ax2.set_ylabel('纬度 [km]')
    ax2.set_title('温度异常')
    plt.colorbar(im2, ax=ax2, label='ΔT [°C]')
    
    plt.suptitle(f'U 倍数 = {U_mult:.1f}, κ = {kappa:.0f} m²/s, CFL = {cfl:.2f}')
    plt.tight_layout()
    plt.show()

interact(plot_simulation_interactive,
         U_mult=FloatSlider(min=0.0, max=8.0, step=0.5, value=1.0,
                            description='U 倍数'),
         kappa=FloatSlider(min=0, max=1e5, step=1e3, value=1e4,
                           description='κ [m²/s]'),
         n_steps=IntSlider(min=100, max=1000, step=100, value=300,
                          description='时间步数'));

interactive(children=(FloatSlider(value=1.0, description='U 倍数', max=8.0, step=0.5), FloatSlider(value=10000.0…
In [17]:
G_check = Grid(30, 6e6)
IC_check = init_box(G_check, xspan=True)

T_end, snaps_check, days_check, cfl_check = run_simulation(
    G_check, kappa=1e4, U_mult=1.0, IC=IC_check, n_steps=500
)

heat_capacity = 51.0
dx_dy = G_check.dx * G_check.dy

heat_content = [heat_capacity * np.sum(s) * dx_dy * 1e-15 for s in snaps_check]
mean_temps = [np.mean(s) for s in snaps_check]

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 4))

ax1.plot(days_check, heat_content, lw=2)
ax1.set_xlabel('时间 [天]')
ax1.set_ylabel('总热量 [Peta-Joules]')
ax1.set_title(f'总热量守恒 (变化 < {abs(heat_content[-1]-heat_content[0])/heat_content[0]*100:.3f}%)')

ax2.plot(days_check, mean_temps, lw=2, color='orange')
ax2.set_xlabel('时间 [天]')
ax2.set_ylabel('平均温度 [°C]')
ax2.set_title(f'平均温度 (初始={mean_temps[0]:.4f}, 最终={mean_temps[-1]:.4f})')

plt.tight_layout()
plt.show()
No description has been provided for this image

计算性能¶

模拟的计算量与网格点数的平方成正比（二维网格）。下面测试不同网格分辨率下每个时间步的耗时。

同时使用 Numba 的 @njit 装饰器对核心计算函数进行即时编译（JIT）加速，与纯 Python/NumPy 实现进行对比。

In [18]:
@njit
def update_ghost_cells_jit(T):
    T[0, :] = T[1, :]
    T[-1, :] = T[-2, :]
    T[:, 0] = T[:, 1]
    T[:, -1] = T[:, -2]


@njit
def advect_jit(T, u, v, dx, dy):
    Ny, Nx = T.shape
    tendency = np.zeros((Ny, Nx))
    for i in range(1, Ny - 1):
        for j in range(1, Nx - 1):
            tendency[i, j] = -(
                u[i, j] * (T[i, j+1] - T[i, j-1]) / (2 * dx) +
                v[i, j] * (T[i+1, j] - T[i-1, j]) / (2 * dy)
            )
    return tendency


@njit
def diffuse_jit(T, kappa, dx, dy):
    Ny, Nx = T.shape
    tendency = np.zeros((Ny, Nx))
    for i in range(1, Ny - 1):
        for j in range(1, Nx - 1):
            tendency[i, j] = kappa * (
                (T[i, j+1] - 2*T[i, j] + T[i, j-1]) / dx**2 +
                (T[i+1, j] - 2*T[i, j] + T[i-1, j]) / dy**2
            )
    return tendency


# 预热 JIT 编译
_G_warmup = Grid(8, 6e6)
_T_warmup = init_box(_G_warmup)
_u_warmup, _v_warmup = double_gyre(_G_warmup)
update_ghost_cells_jit(_T_warmup.copy())
_ = advect_jit(_T_warmup.copy(), _u_warmup, _v_warmup, _G_warmup.dx, _G_warmup.dy)
_ = diffuse_jit(_T_warmup.copy(), 1e4, _G_warmup.dx, _G_warmup.dy)
print('Numba JIT 编译完成')


Nvec = range(1, 20)
tvec = []
tvec_jit = []

for Npower in Nvec:
    G_scale = Grid(8 * Npower, 6e6)
    IC_scale = init_box(G_scale)
    u_s, v_s = double_gyre(G_scale)
    model = OceanModel(G_scale, kappa=1e4, u=u_s, v=v_s)
    model.dt = 6 * 3600

    T_test = IC_scale.copy()
    T_test_jit = IC_scale.copy()

    # 纯 NumPy
    t0 = time.time()
    update_ghost_cells(T_test)
    adv = advect(T_test, model.u, model.v, model.G.dx, model.G.dy)
    diff = diffuse(T_test, model.kappa, model.G.dx, model.G.dy)
    elapsed = time.time() - t0
    tvec.append(elapsed)

    # Numba JIT
    t0 = time.time()
    update_ghost_cells_jit(T_test_jit)
    adv_jit = advect_jit(T_test_jit, model.u, model.v, model.G.dx, model.G.dy)
    diff_jit = diffuse_jit(T_test_jit, model.kappa, model.G.dx, model.G.dy)
    elapsed_jit = time.time() - t0
    tvec_jit.append(elapsed_jit)

# 绘图
fig, ax = plt.subplots(figsize=(8, 5))
N_vals = [8 * n for n in Nvec]

ax.plot(N_vals, tvec, 'o-', lw=2, label='NumPy（纯 Python）')
ax.plot(N_vals, tvec_jit, 's-', lw=2, color='green', label='Numba JIT')

ax.set_xlabel('x 方向网格点数')
ax.set_ylabel('每步耗时 [秒]')
ax.set_title('计算时间 vs 网格分辨率：NumPy vs Numba')
ax.grid(True, alpha=0.3)
ax.legend()

if len(tvec) > 2:
    coeffs = np.polyfit(np.log(N_vals[3:]), np.log(tvec[3:]), 1)
    ax.plot(N_vals, np.exp(coeffs[1]) * np.array(N_vals)**coeffs[0],
            '--', color='blue', alpha=0.5, label=f'NumPy 拟合: O(N^{coeffs[0]:.1f})')

if len(tvec_jit) > 2:
    coeffs_jit = np.polyfit(np.log(N_vals[3:]), np.log(tvec_jit[3:]), 1)
    ax.plot(N_vals, np.exp(coeffs_jit[1]) * np.array(N_vals)**coeffs_jit[0],
            '--', color='green', alpha=0.5, label=f'Numba 拟合: O(N^{coeffs_jit[0]:.1f})')

ax.legend()
plt.tight_layout()
plt.show()

# 打印最大加速比
if len(tvec) > 0 and len(tvec_jit) > 0:
    max_speedup = max(t / j for t, j in zip(tvec, tvec_jit) if j > 0)
    max_idx = max(range(len(tvec)), key=lambda k: tvec[k] / tvec_jit[k] if tvec_jit[k] > 0 else 0)
    print(f'最大加速比: {max_speedup:.1f}x (N = {N_vals[max_idx]})')

Numba JIT 编译完成
No description has been provided for this image 
最大加速比: 5.6x (N = 136)

现实世界中的气候模型¶

模板操作是气候建模的核心计算工具。本节从基础的模板概念出发，逐步构建了完整的二维海洋热量传输模拟。

在前面的一维纬度能量平衡模型中，扩散项

$$D \frac{\partial}{\partial x}\left[(1-x^2)\frac{\partial T}{\partial x}\right]$$

正是通过模板来离散化的。在更复杂的三维气候模型中，同样的原理应用于全球网格：

  • 大气环流模型（GCM）：在经度、纬度、高度三维网格上求解流体力学方程
  • 海洋环流模型：在海量网格点上计算温度、盐度、流速的演变
  • 每个网格点的更新都需要模板操作——将邻域信息与微分算子的离散近似相结合

边界条件的选择同样关键：

  • 极地通常使用周期性或对称边界条件
  • 海洋-陆地交界处使用特殊的通量条件
  • 大气顶部和地表使用辐射边界条件
In [19]:
%%HTML

<video width="960" controls>
    <source src="https://pacman.cs.tsinghua.edu.cn/~hanwentao/cpct/22/sea-surface-temperatures-2017.mp4" type="video/mp4">
</video>

本讲小结¶

本讲从一维 PDE 推广到二维 PDE，引入了模板计算的核心工具。从计算思维的角度：

1. 抽象：模板作为统一的算子表示

上一讲中，平流和扩散的数值方法分别用不同的代码实现。本讲引入模板这一抽象：任何线性微分算子的离散化都可以表示为一个小矩阵（3×3），与邻域逐元素相乘再求和。同一个框架统一处理了梯度、拉普拉斯等不同算子，这是计算思维中"抽象"的典型体现。

2. 维度的扩展

从一维到二维，温度场从向量 $T_i$ 变为矩阵 $T_{i,j}$，空间导数的近似从一维切片变为二维切片。代码结构几乎不变——advect() 和 diffuse() 函数中的 NumPy 切片语法从一维推广到二维，核心逻辑相同。

3. 幽灵单元：边界条件的优雅处理

边界条件是数值计算中的经典难题。本讲引入的幽灵单元技巧，通过在数据周围添加一圈额外单元来统一处理边界：

  • 零值边界（Dirichlet）：幽灵单元 = 0
  • 周期性边界（Periodic）：幽灵单元 = 对侧值
  • 零导数边界（Neumann）：幽灵单元 = 相邻值

这使得核心计算代码无需任何边界判断，体现了"分离关注点"的设计思想。

4. 建模：从工具到物理

将模板、幽灵单元和时间步进组合在一起，成功模拟了二维海洋洋流对热量的传输。速度场来自风驱动环流的解析解（双环流），温度场的演变同时受平流和扩散的影响——这与第 19 讲的气候模型和第 21 讲的一维模型形成了完整的知识链条。