克里金插值计算速度优化求助：大矩阵下lstsq运算过慢

问题分析 #

你的代码核心瓶颈有两个：

1. Python层级的双重循环：外层10000次、内层6000次的循环完全在Python解释器中执行，效率极低；
2. 重复调用np.linalg.lstsq：每次求解都要重新对常量矩阵matrx做分解，这是完全冗余的计算——克里金法中的协方差/变异函数矩阵只需要分解一次即可复用。

下面是针对性的优化方案，从向量化、预计算、编译加速三个维度入手：

优化方案1：向量化计算 + 预矩阵分解（最推荐） #

利用numpy的广播机制消除内层循环，同时预计算matrx的分解结果，避免重复计算：

步骤1：预计算矩阵分解

克里金的matrx（通常是协方差矩阵或变异函数矩阵）是正定矩阵，我们可以用Cholesky分解替代lstsq，分解一次后后续求解线性方程组的速度会提升几个数量级：

import numpy as np

# 预计算Cholesky分解（克里金场景下matrx基本都是正定矩阵）
L = np.linalg.cholesky(matrx)
L_T = L.T  # 预存转置矩阵，避免重复计算

步骤2：向量化计算所有距离和半方差

用numpy广播一次性计算所有gdata与data的点对距离，彻底消除内层循环：

# 计算所有gdata点到data点的距离，shape: (10000, 6000)
distances = np.sqrt(
    (gdata[:, None, 0] - data[:, 0])**2 + 
    (gdata[:, None, 1] - data[:, 1])**2
)

# 向量化计算半方差，shape: (10000, 6000)
semivars = c1 + c2 * (1 - np.exp(-(distances / c3)**c4))

# 给每个半方差向量添加末尾的1（对应克里金的拉格朗日项），shape: (10000, 6001)
semivariance_matrix = np.hstack([semivars, np.ones((len(gdata), 1))])

步骤3：批量求解线性方程组

利用预分解的矩阵快速求解所有10000个线性方程组：

predictions = []
for semivar_vec in semivariance_matrix:
    # 用Cholesky分解两步求解：先解Ly = b，再解L^T x = y
    y = np.linalg.solve(L, semivar_vec)
    x = np.linalg.solve(L_T, y)
    # 取前6000个权重（去掉最后一个拉格朗日乘数）
    weights = x[:-1]
    predictions.append(np.sum(data[:, 2] * weights))

# 转成numpy数组方便后续处理
predictions = np.array(predictions)

如果你的matrx不是正定矩阵，可以用QR分解替代Cholesky：

Q, R = np.linalg.qr(matrx)
# 求解时替换为：x = np.linalg.solve(R, Q.T @ semivar_vec)

优化方案2：用Numba加速循环（适合小幅修改代码） #

如果不想重构代码为向量化风格，可以用Numba的JIT编译加速Python循环，直接将原循环编译为机器码：

import numba as nb
import numpy as np

# 预计算matrx的Cholesky分解
L = np.linalg.cholesky(matrx)
L_T = L.T

# 用Numba JIT装饰，开启多线程并行
@nb.njit(parallel=True)
def kriging_prediction(gdata, data, c1, c2, c3, c4, L, L_T):
    predictions = np.zeros(len(gdata))
    data_vals = data[:, 2]
    data_x = data[:, 0]
    data_y = data[:, 1]
    
    # 并行遍历所有gdata点
    for i in nb.prange(len(gdata)):
        gx, gy = gdata[i]
        semivariance = np.zeros(len(data) + 1)
        
        # 内层循环计算半方差（Numba会编译为机器码）
        for j in range(len(data)):
            dx = gx - data_x[j]
            dy = gy - data_y[j]
            distance = np.sqrt(dx*dx + dy*dy)
            semivariance[j] = c1 + c2 * (1 - np.exp(-(distance/c3)**c4))
        
        semivariance[-1] = 1.0
        
        # 用预分解矩阵快速求解
        y = np.linalg.solve(L, semivariance)
        x = np.linalg.solve(L_T, y)
        weights = x[:-1]
        
        predictions[i] = np.sum(data_vals * weights)
    
    return predictions

# 调用函数得到结果
predictions = kriging_prediction(gdata, data, c1, c2, c3, c4, L, L_T)

@nb.njit(parallel=True)会自动利用多CPU核心并行处理外层循环，速度比纯Python循环快几十倍甚至上百倍。

额外优化建议 #

1. 内存优化：如果10000×6000的距离矩阵占用内存过大（约480MB，float64类型），可以分块处理gdata，比如每次处理1000个点，避免内存溢出；
2. 去重预处理：如果data中有重复坐标点，提前去重可以减少后续计算量；
3. 选择精度合适的数据类型：如果不需要高精度，将float64转为float32可以减少内存占用并提升计算速度。

内容的提问来源于stack exchange，提问作者user3185171