基于Pandas/NumPy的多ID对滚动窗口相关系数高效计算问询

Hey there, let's tackle this performance bottleneck with rolling correlations between your two DataFrames. Brute-forcing loops over 15k ids, 300 id_2s, and 300 dates is definitely going to crawl—so let's shift to vectorized operations and optimized matrix calculations to cut down runtime drastically.

Problem Recap #

You have two DataFrames:

1. df1 with columns id, date, value (15k unique ids)
2. df2 with columns id_2, date, value (300 unique id_2s)

Requirement: For every id-id_2 pair, compute the rolling 6-month Pearson correlation between their value time series at each date (correlation uses the 6 months ending at the target date). Output a long-format DataFrame with id_pair, date, corr.

Why Your Current Loop Is Slow #

A naive loop over all pairs and dates has a time complexity of O(MNT) (15k * 300 * 300 = 1.35e9 operations). Python loops are not optimized for this scale—we need to leverage vectorized operations from Pandas/NumPy to batch compute correlations.

Optimized Solution #

Here's a step-by-step approach that cuts runtime by orders of magnitude:

1. Reshape Data to Wide Format

First, convert both DataFrames to a wide format where each row is a date, and each column is an id/id_2 time series. This lets us use matrix operations to compute correlations in bulk.

import pandas as pd
import numpy as np

# Sample data (replace with your actual DataFrames)
df1 = pd.DataFrame(
    [["f130","200701",0.016196], ["f130","200702",-0.027798], ...],
    columns=["id", "date", "value"]
)
df2 = pd.DataFrame(
    [["d_1","200701",0.026316], ["d_1","200702",-0.004487], ...],
    columns=["id_2", "date", "value"]
)

# Convert to wide format: rows = dates, columns = ids/id_2s
df1_wide = df1.pivot(index="date", columns="id", values="value").sort_index()
df2_wide = df2.pivot(index="date", columns="id_2", values="value").sort_index()

2. Rolling Standardization

Pearson correlation can be computed using standardized time series (mean=0, std=1). For each rolling window, we'll normalize both datasets first:

window_size = 6

# Compute rolling mean and standard deviation for df1
df1_roll_mean = df1_wide.rolling(window=window_size).mean()
df1_roll_std = df1_wide.rolling(window=window_size).std(ddof=1)  # Sample std
df1_std = (df1_wide - df1_roll_mean) / df1_roll_std

# Repeat for df2
df2_roll_mean = df2_wide.rolling(window=window_size).mean()
df2_roll_std = df2_wide.rolling(window=window_size).std(ddof=1)
df2_std = (df2_wide - df2_roll_mean) / df2_roll_std

3. Batch Compute Rolling Correlations

For each valid date (starting from the 6th month), compute the correlation matrix between all id and id_2 pairs using matrix multiplication. This replaces M*N individual calculations with a single matrix operation per date:

# Get dates where rolling window is valid
valid_dates = df1_wide.index[window_size-1:]
results = []

for date in valid_dates:
    # Extract the 6-month window ending at current date
    window_df1 = df1_std.loc[:date].tail(window_size)
    window_df2 = df2_std.loc[:date].tail(window_size)
    
    # Compute correlation matrix: (M x window) @ (window x N) = M x N
    # Divide by (window_size - 1) for sample correlation
    corr_matrix = (window_df1.T @ window_df2) / (window_size - 1)
    
    # Reshape matrix to long format
    corr_long = corr_matrix.stack().reset_index()
    corr_long.columns = ["id", "id_2", "corr"]
    corr_long["date"] = date
    corr_long["id_pair"] = corr_long["id_2"].str.replace("_", "") + "_" + corr_long["id"]
    
    results.append(corr_long[["id_pair", "date", "corr"]])

# Combine all results into final output
final_df = pd.concat(results, ignore_index=True)
final_df["corr"] = final_df["corr"].round(9)  # Match sample output precision

4. Optional: Parallelize for Extra Speed

Since each date's calculation is independent, we can parallelize the loop using joblib to utilize multiple CPU cores:

from joblib import Parallel, delayed

def compute_corr_for_date(date):
    window_df1 = df1_std.loc[:date].tail(window_size)
    window_df2 = df2_std.loc[:date].tail(window_size)
    corr_matrix = (window_df1.T @ window_df2) / (window_size - 1)
    corr_long = corr_matrix.stack().reset_index()
    corr_long.columns = ["id", "id_2", "corr"]
    corr_long["date"] = date
    corr_long["id_pair"] = corr_long["id_2"].str.replace("_", "") + "_" + corr_long["id"]
    return corr_long[["id_pair", "date", "corr"]]

# Run in parallel (n_jobs=-1 uses all available cores)
results = Parallel(n_jobs=-1)(
    delayed(compute_corr_for_date)(date) for date in valid_dates
)

final_df = pd.concat(results, ignore_index=True)

Key Performance Improvements #

• Matrix Multiplication: Uses optimized NumPy/Pandas C-level operations instead of Python loops, which are 100-1000x faster for large datasets.
• Reduced Loop Count: Instead of looping over 4.5M pairs, we loop over just 300 dates.
• Parallelization: Splits date calculations across CPU cores, cutting runtime further on multi-core machines.

Notes #

• Missing Values: If your data has gaps, use df1_wide.asfreq("MS") (for monthly data) to fill missing dates, then handle NaNs with fillna() or dropna() before computing rolling stats.
• Window Validation: Ensure dates are sorted and continuous to avoid incorrect window selections.

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