幂律拟合向量化实现求助：迁移至PyTorch GPU并消除for循环

我正在将代码库迁移至使用PyTorch张量在GPU上运行，GPU上的for循环（尤其针对小尺寸数据）性能极差。我尝试对以下幂律拟合函数进行向量化（即消除所有for循环），尤其需要针对Ds[]部分的实现提供帮助。完成向量化的NumPy版本后，转换为张量的过程会非常简单。

原始代码 #

import numpy as np

data = np.array([11, 2, 3, 40, 5, 6, 7, 8, 9, 10])
N = len(data)
log_data = np.log(data)
alphas = np.zeros(N-1)
Ds = np.zeros(N-1)

for i, xmin in enumerate(data[:-1]):
    n = float(N - i)
    alpha = 1 + n / (np.sum(log_data[i:]) - n * log_data[i])
    alphas[i] = alpha
    if alpha > 1:
        Ds[i] = np.max(np.abs(
            1 - (data[i:] / xmin) ** (-alpha + 1)  # 理论CDF
            - np.arange(n) / n                     # 实际CDF
        ))

sigmas = (alphas - 1) / np.sqrt(N - np.arange(N - 1))

尝试代码（存在错误） #

import numpy as np

def fit_power_law(data):
    # Assuming 'data' is already a NumPy array
    logdata = np.log(data)
    N = len(data)

    # Compute cumulative sums from the end
    sums = np.cumsum(logdata[::-1])[::-1]

    # Create array of n values
    ns = np.arange(N, 0, -1, dtype=np.float64)

    # Compute alphas
    alphas = 1 + ns / (sums - ns * logdata)

    # Filter alphas to only include those > 1
    valid_indices = alphas > 1
    filtered_alphas = alphas[valid_indices]
    filtered_data = data[valid_indices]

    M = len(filtered_data)

    # Create a matrix of NaNs
    matrix = np.full((M, M), np.nan)

    # Create an index array for broadcasting
    idx = np.arange(M)
    shifted_indices = idx[:, None] + idx

    # Mask to avoid out-of-bound indices
    mask = shifted_indices < M

    # Fill the matrix with the appropriate values
    matrix[mask] = filtered_data[shifted_indices[mask]]

    # Divide each row by the data array
    tt = matrix / filtered_data[:, np.newaxis]

    # Calculate the first term  [1 - (data[i:] / xmin) ** (-alpha + 1)]
    first_term = 1 - tt ** (-filtered_alphas[:, np.newaxis] + 1)

    
    # Reverse the filtered data array
    data_reversed = filtered_data

    # Create a range array and a mask
    range_array = np.arange(int(data.shape[0]))    
    mask = range_array < data_reversed[:, np.newaxis]

    # Create the result array
    result = np.where(mask, range_array, np.nan)


    # Calculate the length of each inner array (excluding NaNs)
    lengths = np.sum(~np.isnan(result), axis=1)

    # Ensure lengths are at least 1 to avoid division by zero
    lengths = np.maximum(lengths, 1)

    # Broadcast lengths to match the shape of result for division
    lengths_broadcast = lengths[:, np.newaxis]

    # np.arange(n) / n #This is wrong
    second_term = np.where(~np.isnan(result), result / lengths_broadcast, np.nan)

    # Calculate the difference
    min_answer = first_term - second_term

    # Return the maximum absolute difference
    ds = np.nanmax(np.abs(min_answer), axis=1)[:-1]


    alphas = alphas[:-1]

    sigma = (alphas - 1) / np.sqrt(N - np.arange(N - 1, dtype=np.float64))

    ds_final = np.zeros(N-1, dtype=np.float64)

    ds_final[valid_indices] = ds

    return ds_final, alphas, sigma

核心问题在于对np.arange(n)/n的向量化逻辑错误，原始代码中每个i对应的该数组是长度为n=N-i的序列：0/n,1/n,...,(n-1)/n。以下是修正后的完整向量化实现：

import numpy as np

def fit_power_law(data):
    data = np.asarray(data)
    N = len(data)
    log_data = np.log(data)
    
    # 向量化计算每个位置i到末尾的log和，替代循环中的np.sum(log_data[i:])
    cum_sum_log = np.cumsum(log_data[::-1])[::-1]
    # 生成每个i对应的n值：N-i，截取前N-1个对应原始循环的data[:-1]
    ns = np.arange(N, 0, -1, dtype=np.float64)[:-1]
    
    # 计算所有alphas值
    alphas = 1 + ns / (cum_sum_log[:-1] - ns * log_data[:-1])
    
    # 初始化Ds数组
    Ds = np.zeros(N-1, dtype=np.float64)
    # 筛选alpha>1的有效索引
    valid_idx = alphas > 1
    
    # 处理有效索引对应的计算
    if np.any(valid_idx):
        valid_alphas = alphas[valid_idx]
        valid_xmins = data[:-1][valid_idx]
        valid_ns = ns[valid_idx].astype(int)
        max_n = valid_ns.max()
        
        # 用滑动窗口生成所有data[i:i+n]的矩阵，避免手动索引偏移
        data_window = np.lib.stride_tricks.sliding_window_view(data, max_n)[:-1][valid_idx]
        # 按每个i对应的实际n长度截断窗口
        data_window = np.array([row[:n] for row, n in zip(data_window, valid_ns)])
        
        # 计算理论CDF
        theoretical_cdf = 1 - (data_window / valid_xmins[:, np.newaxis]) ** (-valid_alphas[:, np.newaxis] + 1)
        
        # 向量化生成实际CDF：每个i对应长度n的序列k/n（k从0到n-1）
        actual_cdf = np.arange(max_n) / valid_ns[:, np.newaxis]
        # 同样按实际n长度截断
        actual_cdf = np.array([row[:n] for row, n in zip(actual_cdf, valid_ns)])
        
        # 计算绝对差的最大值并赋值给Ds
        Ds[valid_idx] = np.max(np.abs(theoretical_cdf - actual_cdf), axis=1)
    
    # 计算sigmas
    sigmas = (alphas - 1) / np.sqrt(N - np.arange(N-1, dtype=np.float64))
    
    return Ds, alphas, sigmas

关键修正说明 #

• 滑动窗口处理data[i:]：使用np.lib.stride_tricks.sliding_window_view生成所有滑动窗口，替代手动构建矩阵的冗余操作，逻辑更清晰且高效。
• 实际CDF正确生成：通过广播np.arange(max_n)和valid_ns直接生成每个i对应的k/n序列，完全匹配原始循环的逻辑。
• 索引对齐优化：确保有效索引对应的alpha、xmin、n和数据窗口严格对应，避免维度不匹配问题。
• 适配PyTorch迁移：代码结构简洁，后续转换为PyTorch张量时，只需将np替换为torch，滑动窗口可改用torch.unfold实现，无需大幅修改。

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