如何用稀疏矩阵实现非均匀域移动平均？

非均匀域移动平均现有实现 #

以下代码可以正常运行：

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

def convolve(s, f):
    """计算序列S与通用函数F的卷积
    本质是基于S.index的权重F对S做移动平均"""
    index_v = s.index.values
    weight_mx = f(index_v - index_v[:, np.newaxis])
    weighted_sum = np.sum(s.values[:, np.newaxis] * weight_mx, axis=0)
    normalization = np.sum(weight_mx, axis=0)
    return pd.Series(weighted_sum/normalization, index=s.index)

size = 1000
df = pd.DataFrame({"x":np.random.normal(size=size)},
                  index=np.random.exponential(size=size)).sort_index()
def f(x):
    return np.exp(-x*x*30)
df["avg"] = convolve(df.x, f)
plt.scatter(df.index, df.avg, s=1, label="平均值")
plt.scatter(df.index, df.x, s=1, label="随机值")
plt.title("随机数据的移动平均")
plt.legend()

但这个实现有严重的内存问题：它会生成一个**O(n²)**的权重矩阵（原问题中误写为O(size³)，实际是二维数组），当数据量达到14454条时，会触发内存错误：

MemoryError: Unable to allocate 22.0 TiB for an array with shape (14454, 14454, 14454) and data type float64

已知f函数仅在一个很窄的取值范围内返回非零值。

优化方案：基于有序索引的稀疏计算 #

因为输入序列的索引已经排序（原代码中调用了sort_index()），我们可以为每个索引点定位出权重函数f有非零输出的有效区间，只计算该区间内的权重，完全避免生成全量n×n矩阵。

优化后的convolve_sparse函数

def convolve_sparse(s, f, tol=1e-12):
    """稀疏化非均匀域移动平均计算
    参数:
        s: 输入pandas Series，索引需已排序
        f: 权重函数，仅在窄范围内返回非零值
        tol: 权重阈值，小于该值的权重视为0，跳过计算
    """
    index_v = s.index.values
    values = s.values
    n = len(index_v)
    weighted_sum = np.zeros(n)
    normalization = np.zeros(n)
    
    # 二分查找权重函数的有效半宽（找到最大x使得f(x) >= tol）
    def get_half_width(f, tol):
        low, high = 0.0, 1.0
        while f(high) >= tol:
            high *= 2
        while high - low > 1e-8:
            mid = (low + high) / 2
            if f(mid) >= tol:
                low = mid
            else:
                high = mid
        return low
    
    half_width = get_half_width(f, tol)
    
    for i in range(n):
        curr_idx = index_v[i]
        # 定位有效区间的左右边界
        left = np.searchsorted(index_v, curr_idx - half_width, side='left')
        right = np.searchsorted(index_v, curr_idx + half_width, side='right')
        # 计算区间内的索引差值与权重
        idx_diffs = index_v[left:right] - curr_idx
        weights = f(idx_diffs)
        # 过滤极小权重（可选，进一步降低计算量）
        valid_mask = weights >= tol
        valid_weights = weights[valid_mask]
        valid_vals = values[left:right][valid_mask]
        # 累加加权和与归一化项
        weighted_sum[i] = np.sum(valid_vals * valid_weights)
        normalization[i] = np.sum(valid_weights)
    
    # 处理极端情况（归一化项为0）
    normalization[normalization == 0] = 1e-12
    return pd.Series(weighted_sum / normalization, index=s.index)

验证优化效果

# 原参数测试
size = 1000
df = pd.DataFrame({"x":np.random.normal(size=size)},
                  index=np.random.exponential(size=size)).sort_index()
def f(x):
    return np.exp(-x*x*30)
df["avg_sparse"] = convolve_sparse(df.x, f)

# 对比原实现结果（误差可忽略）
print(f"最大误差：{np.max(np.abs(df['avg'] - df['avg_sparse'])):.10f}")

# 绘制对比图
plt.scatter(df.index, df.avg, s=1, label="原实现")
plt.scatter(df.index, df.avg_sparse, s=1, label="稀疏优化")
plt.title("原实现与稀疏优化结果对比")
plt.legend()
plt.show()

优化核心逻辑 #

1. 区间定位：利用np.searchsorted在有序索引中快速找到每个点的有效权重区间，单点位定位时间复杂度为O(log n)。
2. 内存优化：仅存储当前点有效区间内的权重，内存占用从O(n²)降至O(k)（k为单条数据的平均有效区间长度，远小于n）。
3. 通用性：通过二分查找自动适配不同的权重函数f，只要f具有窄带非零特性即可。

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