如何用Dask快速为滚动窗口值拟合二次多项式？

1. 修复自定义函数的报错问题 #

你的window_curve_fit报错是因为返回列表而非numpy数组，且map_blocks参数传递方式有误。先修改函数：

import numpy as np
from numpy.lib.stride_tricks import sliding_window_view

def window_curve_fit(arr, deg=2, size=5, x_step=0.01):
    # 生成中心对齐的窗口x值（比如size=5时为[-0.02, -0.01, 0, 0.01, 0.02]）
    x = np.linspace(-(size//2)*x_step, (size//2)*x_step, size)
    # 预计算拟合所需的X矩阵，避免重复计算
    X = np.vstack([x**2, x, np.ones(size)]).T
    # 对每列生成滑动窗口，形状为(rows - size +1, cols, size)
    windows = sliding_window_view(arr, window_shape=size, axis=0)
    # 用伪逆矩阵批量计算所有窗口的二次项系数a
    X_inv = np.linalg.pinv(X)
    a = windows @ X_inv[:, 0]
    return 2 * a  # 返回曲率值，形状为(rows - size +1, cols)

调用map_blocks时，需指定输出块形状，确保Dask能正确拼接结果：

import dask.array as da

n_window = 5
x_step = 0.01

# 计算输出数组的形状和块结构
output_shape = (scene2.shape[0] - n_window + 1, scene2.shape[1])
output_chunks = (tuple(c - n_window +1 for c in scene2.chunks[0]), scene2.chunks[1])

rslt = da.map_blocks(
    window_curve_fit,
    scene2,
    deg=2,
    size=n_window,
    x_step=x_step,
    dtype=np.float32,
    chunks=output_chunks
).compute()

2. 提速核心：向量化替代循环+预计算 #

之前方法慢的核心原因是逐个窗口循环拟合，改用以下方式可大幅提速：

• 预计算X的伪逆矩阵，避免每个窗口重复执行最小二乘法计算
• 用numpy广播/矩阵乘法批量处理所有窗口，完全利用CPU向量化指令
• 弃用Polynomial.fit（面向对象类开销远高于numpy底层矩阵运算）

3. 优化Dask并行效率 #

• 保持现有块大小(500, 3000)：列数单块存储，避免跨块处理列数据（滚动窗口沿行方向，行分块不影响）
• 根据CPU核心数调整Dask并行资源：

  from dask.distributed import Client
  client = Client(n_workers=4, threads_per_worker=2)  # 按需调整参数
  

• 确保每个块计算量足够大，降低Dask调度开销占比

4. 简洁替代方案：Dask rolling+apply #

若偏好Dask原生滚动窗口API，可使用以下代码：

def fit_quadratic(window):
    x = np.linspace(-(window.shape[0]//2)*0.01, (window.shape[0]//2)*0.01, window.shape[0])
    X = np.vstack([x**2, x, np.ones(window.shape[0])]).T
    a = np.linalg.lstsq(X, window, rcond=None)[0][0]
    return 2*a

# 沿行方向做中心滚动窗口拟合
rslt = scene2.rolling(window=n_window, axis=0, center=True).apply(fit_quadratic, dtype=np.float32).compute()

注：该方案代码更简洁，但并行效率略低于map_blocks+向量化实现。

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