如何向量化Scipy实现的函数？优化分组线性回归统计流程

Nice call looking to vectorize this grouped regression workflow—ditching explicit loops not only cleans up your code but can also give a solid performance boost, especially with large datasets. Here's a step-by-step, fully vectorized approach using NumPy that replaces your loop-based calls to fx_stats.

Core Idea: Aggregate Group-Level Stats First #

Instead of looping through each key and running regression individually, we can compute all necessary group-level statistics in a vectorized way using np.bincount. Linear regression metrics like slope, Pearson r, and R² can all be derived from basic sums and counts per group, which np.bincount handles efficiently.

Step-by-Step Implementation #

Let’s assume your data is stored in NumPy arrays (adjust if you’re using pandas—we can map this logic easily):

import numpy as np
from scipy.stats import linregress  # Only for comparison, not needed in vectorized code

# Sample data
keys = np.array(['Apple', 'Apple', 'Banana', 'Banana', 'Banana', 'Cherry'])
days = np.array([1, 2, 3, 1, 2, 3])
calculated_weight = np.array([10, 20, 30, 15, 30, 45])

1. Assign Group Indices

First, map each unique key to an integer index—this lets us use np.bincount to aggregate by group:

unique_keys, group_indices = np.unique(keys, return_inverse=True)

2. Compute Group-Level Aggregates

Calculate all the sums and counts we need for regression using vectorized operations:

# Number of samples per group
group_counts = np.bincount(group_indices)

# Sum of x (days), y (weight), x*y, x², y² per group
sum_x = np.bincount(group_indices, weights=days)
sum_y = np.bincount(group_indices, weights=calculated_weight)
sum_xy = np.bincount(group_indices, weights=days * calculated_weight)
sum_x_squared = np.bincount(group_indices, weights=days ** 2)
sum_y_squared = np.bincount(group_indices, weights=calculated_weight ** 2)

3. Vectorize Regression Metric Calculations

Use the aggregated stats to compute slope, Pearson r, R², and intercept in one go:

# Calculate numerator and denominator for slope and Pearson r
numerator = group_counts * sum_xy - sum_x * sum_y
denominator_slope = group_counts * sum_x_squared - sum_x ** 2
denominator_r = np.sqrt(denominator_slope * (group_counts * sum_y_squared - sum_y ** 2))

# Handle edge cases (e.g., groups with <2 samples where denominator is 0)
valid_mask = denominator_slope != 0

# Compute metrics
slope = np.where(valid_mask, numerator / denominator_slope, np.nan)
pearson_r = np.where(valid_mask & (denominator_r != 0), numerator / denominator_r, np.nan)
r_squared = np.where(valid_mask, pearson_r ** 2, np.nan)
intercept = np.where(valid_mask, (sum_y - slope * sum_x) / group_counts, np.nan)

4. Organize Results

Package the output into a structured format that’s easy to use—either a NumPy structured array or a dictionary:

# Option 1: Structured NumPy array (great for further numerical operations)
results_array = np.rec.fromarrays(
    [unique_keys, slope, intercept, pearson_r, r_squared, group_counts],
    names='key,slope,intercept,pearson_r,r_squared,sample_count'
)

# Option 2: Dictionary (intuitive for lookup)
results_dict = {
    key: {
        'slope': slope[i],
        'intercept': intercept[i],
        'pearson_r': pearson_r[i],
        'r_squared': r_squared[i],
        'sample_count': group_counts[i]
    }
    for i, key in enumerate(unique_keys)
}

Why This Works Better Than Loops #

• Performance: np.bincount is implemented in C, so it’s orders of magnitude faster than Python loops for large datasets.
• Code Clarity: No nested loops or repeated function calls—all logic is consolidated in a few vectorized steps.
• Maintainability: Easy to extend with additional metrics (like standard errors) by adding more aggregated stats.

Verification #

If you want to confirm the results match your original loop-based approach, compare with a quick loop:

# Original loop approach for comparison
loop_results = {}
for key in unique_keys:
    mask = keys == key
    x = days[mask]
    y = calculated_weight[mask]
    res = linregress(x, y)
    loop_results[key] = {
        'slope': res.slope,
        'pearson_r': res.rvalue,
        'r_squared': res.rvalue ** 2
    }

# Check equality (account for floating point precision)
assert np.allclose(results_dict['Apple']['slope'], loop_results['Apple']['slope'])
assert np.allclose(results_dict['Banana']['pearson_r'], loop_results['Banana']['pearson_r'])

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