将Matlab的function_AutoCorr自相关函数转换为等效Python函数

功能说明 #

该函数生成的矩阵对角线元素为自相关值，非对角线元素为交叉相关值，输入输出和原Matlab函数完全一致：

• 输入参数1 tau：滞后值，例如tau=1
• 输入参数2 y：输入数据集，例如10行9列的矩阵
• 输出：维度为N×N的相关矩阵（N为输入数据集的列数）

原Matlab实现代码 #

function CORR=function_AutoCorr(tau,y)

% 该函数生成矩阵，对角线元素为自相关值，非对角线元素为交叉相关值
% y 为数据集，例如10×9的矩阵
% tau 为滞后值，例如tau=1

Size=size(y);
N=Size(1,2); % 列数
T=Size(1,1); % 行数（序列长度）
for i=1:N
  for j=1:N
      temp1=0;
    for t=1:T-tau
    G=0.5*((y(t+tau,i)*y(t,j))+(y(t+tau,j)*y(t,i)));
    temp1=temp1+G;
    end
    CORR(i,j)=temp1/(T-tau);
  end
end
end

等效Python实现代码 #

基于numpy实现，逻辑和原Matlab代码完全对齐，输出结果完全一致：

import numpy as np

def function_AutoCorr(tau, y):
    """
    生成对角线为自相关值、非对角线为交叉相关值的矩阵
    参数:
        tau: 滞后值
        y: 输入数据集，格式为(行数T, 列数N)的numpy数组
    返回:
        CORR: 维度为(N, N)的相关矩阵
    """
    T, N = y.shape  # T为序列长度（行数），N为变量数（列数）
    CORR = np.zeros((N, N))
    
    for i in range(N):
        for j in range(N):
            temp1 = 0.0
            # 遍历所有有效时间步，完全对齐Matlab的循环逻辑
            for t in range(T - tau):
                G = 0.5 * ((y[t+tau, i] * y[t, j]) + (y[t+tau, j] * y[t, i]))
                temp1 += G
            CORR[i, j] = temp1 / (T - tau)
    return CORR

# 调用示例
if __name__ == "__main__":
    # 生成示例数据：10行9列的随机矩阵
    y = np.random.rand(10, 9)
    tau = 1
    corr_matrix = function_AutoCorr(tau, y)
    print(corr_matrix)

优化版（向量化实现，大矩阵计算效率更高，结果完全一致） #

def function_AutoCorr_vectorized(tau, y):
    T, N = y.shape
    CORR = np.zeros((N, N))
    y_lag = y[tau:, :]  # 滞后tau步的序列
    y_ori = y[:-tau, :] # 原始序列
    for i in range(N):
        for j in range(N):
            CORR[i,j] = 0.5 * (np.mean(y_lag[:,i] * y_ori[:,j]) + np.mean(y_lag[:,j] * y_ori[:,i]))
    return CORR

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