Python中具有inf值的滚动平均值

import numpy as np

def rolling_average_with_inf(arr, window_size):
    # 将inf值替换为NaN
    arr = np.where(arr == np.inf, np.nan, arr)
    
    # 使用numpy的rolling函数计算滚动平均值
    rolling_avg = np.nanmean(np.pad(arr, (window_size-1, 0), mode='constant'), axis=0)
    
    return rolling_avg

# 示例数据
data = np.array([1, 2, 3, np.inf, 5, 6, np.inf, 8, 9])

# 计算滚动平均值
window_size = 3
result = rolling_average_with_inf(data, window_size)
print(result)