用Python和NumPy手把手实现Sinkhorn-Knopp算法：从成本矩阵到最优分配的可视化-育师

用Python和NumPy手把手实现Sinkhorn-Knopp算法：从成本矩阵到最优分配的可视化

在数据科学和机器学习领域，最优传输问题是一个经典且实用的课题。想象一下这样的场景：你手头有8个仓库的货物需要分配到5个不同的商店，每个仓库的库存量不同，每个商店的需求量也不同，而且运输成本因仓库和商店的位置而异。如何找到一个最优的分配方案，使得总运输成本最低？这正是Sinkhorn-Knopp算法大显身手的地方。

与传统的最优化方法不同，Sinkhorn-Knopp算法通过引入熵正则化，将硬约束问题转化为更易求解的软约束问题。本文将带你用Python和NumPy一步步实现这个算法，并通过可视化直观地理解分配结果。我们不会深入复杂的数学推导，而是聚焦于"怎么做"和"怎么看"，让你在动手实践中掌握算法的精髓。

1. 环境准备与数据初始化

在开始编码之前，我们需要准备好Python环境和必要的库。建议使用Python 3.6+版本，并安装以下库：

pip install numpy pandas matplotlib

接下来，我们定义问题的输入数据。最优传输问题需要三个核心输入：

成本矩阵M：表示从每个供应点到每个需求点的运输成本
供应分布r：表示每个供应点的资源总量
需求分布c：表示每个需求点的需求量

让我们创建一个具体的例子：

import numpy as np # 供应分布 (8个仓库的库存) r = np.array([3, 3, 3, 4, 2, 2, 2, 1]) # 需求分布 (5个商店的需求) c = np.array([4, 2, 6, 4, 4]) # 成本矩阵 (8x5，表示从仓库到商店的运输成本) M = np.array([ [2, 2, 1, 0, 0], [0, -2, -2, -2, -2], [1, 2, 2, 2, -1], [2, 1, 0, 1, -1], [0.5, 2, 2, 1, 0], [0, 1, 1, 1, -1], [-2, 2, 2, 1, 1], [2, 1, 2, 1, -1] ], dtype=float)

注意：在实际应用中，成本矩阵的值通常为正数。本例中包含负数是为了展示算法对各类情况的适应性。

2. Sinkhorn-Knopp算法核心实现

Sinkhorn-Knopp算法的核心思想是通过交替的行列归一化，找到一个满足边际约束的传输矩阵P。下面是算法的具体实现步骤：

2.1 算法初始化

首先，我们需要定义几个关键参数：

lam：熵正则化系数，控制正则化的强度
epsilon：收敛阈值，用于判断算法是否收敛

def compute_optimal_transport(M, r, c, lam, epsilon=1e-8): """ 计算最优传输矩阵和Sinkhorn距离 参数: M : 成本矩阵 (n x m) r : 供应分布 (n,) c : 需求分布 (m,) lam : 熵正则化系数 epsilon : 收敛阈值 返回: P : 最优传输矩阵 (n x m) dist : Sinkhorn距离 """ n, m = M.shape # 初始化传输矩阵P P = np.exp(-lam * M) P /= P.sum() # 初始归一化 u = np.zeros(n) # 用于存储行和 # 开始迭代 while np.max(np.abs(u - P.sum(1))) > epsilon: u = P.sum(1) # 计算当前行和 P *= (r / u).reshape((-1, 1)) # 行归一化 v = P.sum(0) # 计算当前列和 P *= (c / v).reshape((1, -1)) # 列归一化 return P, np.sum(P * M) # 返回传输矩阵和总成本

2.2 NumPy广播机制的应用

在实现中，我们巧妙地利用了NumPy的广播机制来简化行列归一化的操作。广播机制允许不同形状的数组进行算术运算，这是NumPy的强大特性之一。

例如，在行归一化这一步：

P *= (r / u).reshape((-1, 1))

这里(r / u)得到一个形状为(n,)的数组，通过reshape((-1, 1))将其转换为(n,1)的列向量。当这个列向量与(n,m)的矩阵P相乘时，NumPy会自动将列向量广播到m列，实现每行的独立缩放。

3. 运行算法与结果分析

现在我们可以运行算法并查看结果了：

lam = 10 # 正则化系数 P, d = compute_optimal_transport(M, r, c, lam=lam) print("最优传输矩阵P:") print(np.round(P, 3)) print("\nSinkhorn距离:", d)

运行结果可能如下：

最优传输矩阵P: [[0. 0. 0. 0. 0. ] [0. 0. 0. 0. 0. ] [0. 0. 0. 0. 0. ] [0. 0. 0. 0. 0. ] [0. 0. 0. 0. 0. ] [0. 0. 0. 0. 0. ] [0. 0. 0. 0. 0. ] [0. 0. 0. 0. 0. ]] Sinkhorn距离: 0.0

提示：实际运行结果会根据输入参数不同而变化。这个示例输出是占位符，真实结果需要运行代码获得。

4. 结果可视化

理解一个矩阵的最好方式就是将其可视化。我们将使用堆叠条形图来展示分配结果：

import pandas as pd import matplotlib.pyplot as plt # 将传输矩阵转换为DataFrame便于可视化 partition = pd.DataFrame(P, index=[f"仓库{i+1}" for i in range(len(r))], columns=[f"商店{j+1}" for j in range(len(c))]) # 绘制堆叠条形图 plt.figure(figsize=(10, 6)) ax = partition.plot(kind='bar', stacked=True) ax.set_ylabel('分配比例') ax.set_title(f'最优分配方案 (λ={lam})') plt.legend(title='目标商店', bbox_to_anchor=(1.05, 1), loc='upper left') plt.tight_layout() plt.show()

可视化结果可以清晰地展示：

每个仓库的资源如何分配到各个商店
哪些仓库-商店对之间的分配比例较高
整体分配是否符合供应和需求的约束

5. 参数调优与算法分析

5.1 正则化系数λ的影响

λ值控制着正则化的强度，对结果有重要影响：

λ值较大：强调熵正则化，分配更加均匀
λ值较小：更注重最小化运输成本，分配更加集中

我们可以通过实验观察λ的影响：

lambdas = [1, 10, 100] results = {} for lam in lambdas: P, d = compute_optimal_transport(M, r, c, lam=lam) results[f"λ={lam}"] = P.flatten() # 比较不同λ下的分配分布 pd.DataFrame(results).plot(kind='kde', title='不同λ值下的分配分布') plt.xlabel('分配比例') plt.show()

5.2 收敛性分析

了解算法的收敛速度对实际应用很重要。我们可以记录每次迭代的变化：

def compute_optimal_transport_with_trace(M, r, c, lam, epsilon=1e-8): n, m = M.shape P = np.exp(-lam * M) P /= P.sum() u = np.zeros(n) trace = [] while True: u = P.sum(1) trace.append(np.max(np.abs(u - r))) if trace[-1] <= epsilon: break P *= (r / u).reshape((-1, 1)) P *= (c / P.sum(0)).reshape((1, -1)) return P, np.sum(P * M), trace _, _, trace = compute_optimal_transport_with_trace(M, r, c, lam=10) plt.plot(trace) plt.title('收敛过程') plt.xlabel('迭代次数') plt.ylabel('最大行和误差') plt.yscale('log') plt.show()

6. 实际应用案例

让我们考虑一个更实际的例子：广告投放优化。假设有：

5个广告位（供应）
3个广告主（需求）
成本矩阵表示不同广告位对不同广告主的效果（点击率）

# 广告位展示量 ad_slots = np.array([10, 15, 20, 15, 10]) # 广告主预算 advertisers = np.array([25, 20, 25]) # 效果矩阵 (越高表示效果越好) effectiveness = np.array([ [0.8, 0.5, 0.3], [0.6, 0.7, 0.4], [0.5, 0.6, 0.9], [0.7, 0.4, 0.5], [0.4, 0.8, 0.6] ]) # 将效果转换为成本 (效果越低，成本越高) M_ads = 1 - effectiveness # 计算最优分配 P_ads, _ = compute_optimal_transport(M_ads, ad_slots, advertisers, lam=20) # 可视化 pd.DataFrame(P_ads, index=[f"广告位{i+1}" for i in range(5)], columns=[f"广告主{j+1}" for j in range(3)]).plot(kind='bar', stacked=True) plt.ylabel('展示量分配') plt.title('广告位最优分配') plt.show()

这个例子展示了Sinkhorn-Knopp算法在资源分配问题中的实际价值。通过调整成本矩阵的定义，我们可以将算法应用于各种场景，如：