多嵌套循环计算产品衰减效率低，求并行/向量化优化方案

计算产品衰减时遇到性能瓶颈：需要处理大量参数组合，基于MultiIndex DataFrame的多层嵌套循环耗时极高，目前仅完成德国一国的计算，还需扩展到其他国家。
尝试用joblib做并行计算，但速度反而变慢；试过向量化优化，也没提升效率，求有效的并行或向量化优化方法。

核心计算代码 #

##############################################
#The actual computation
##############################################
import time 
start_time = time.time()
#Loop over gr, ip and year
for gr in np.arange(-2.0, 2.0, 0.5):  # Loop over -2.0, -1.5, ..., 2.0
    for l_s in np.arange(-10.0, 10.0, 1.0):
        for ip in [70, 80, 90]:
            for year in np.arange(1942.0, 2100.0, 1.0):  # Loop over years from 1942.0 to 2099.0
                subs = sales_df.loc[(gr, l_s, ip, year)]['Germany'].values                
                diagonal_values = subs
                result_array = np.zeros((len(diagonal_values) + len(decay) - 1, len(diagonal_values)))

                for i, val in enumerate(diagonal_values):
                    result_array[i:i+len(decay), i] = val * decay
                
                decay_total = result_array.sum(axis=1).round(0)
                decay_total_insert = decay_total[0:300]
                
                # Inserting the decayed values into the DataFrame
                subsetted_df.loc[(gr, l_s, ip, year), 'Germany'] = decay_total_insert
end_time = time.time()
execution_time = end_time - start_time

测试数据生成代码 #

import numpy as np
import pandas as pd
##############################################
#Creating sample data
##############################################
# Sample input data can be done smarter, but doesn't matter since I have real data.
years = np.arange(1942, 2100)
gstep = np.array([-2, -1.5, -1, -0.5, 0, 0.5, 1, 1.5, 2])
lstep = np.arange(-10, 11)
iprange = np.arange(70, 91, step=10)
countries = ["Germany", "Sweden", "Norway", "Austria", "Belgium", "France"]  # Add 17 more countries here

# Generate all possible combinations of input variables
inputcombs = np.array(np.meshgrid(years, gstep, lstep, iprange)).T.reshape(-1, 4)
input_df = pd.DataFrame(inputcombs, columns=["year", "gr", "lschange", "ip"])

# Generating sample sales data
np.random.seed(42)  # For reproducibility
sdata = np.random.randint(0, 100, size=(input_df.shape[0], len(countries)))
input_df[countries] = sdata

# Set the specified columns as the index
input_df.set_index(["gr", "lschange", "ip", "year"], inplace=True)

# Sorting index by specified columns
input_df.sort_index(level=["gr", "lschange", "ip", "year"], inplace=True)

# Create the additional index "actual_year" ranging from 1943 to 2100
actual_years = np.arange(1943, 2401)
index_levels = input_df.index.levels + [actual_years]
multi_index = pd.MultiIndex.from_product(index_levels, names=input_df.index.names + ["actual_year"])
output_df = pd.DataFrame(index=multi_index, columns=countries)

# Subsetting the DataFrame to get only values that are higher than year in index.
subsetting_condition = (output_df.index.get_level_values('actual_year') >= (output_df.index.get_level_values('year') + 1)) & \
                      (output_df.index.get_level_values('actual_year') <= (output_df.index.get_level_values('year') + 300))
subsetted_df = output_df[subsetting_condition]

# Get the shape of the existing DataFrame
sales_df = subsetted_df.copy()
# Get the shape of the existing DataFrame
rows, cols = sales_df.shape

# Generate an array of random integers between 5 and 1000 with the same shape as the DataFrame
random_integers = np.random.randint(5, 1001, size=(rows, cols))

# Replace the DataFrame values with the random integers
sales_df.values[:] = random_integers

##############################################
#Creating decay vector
##############################################

# Parameters for the log-normal distribution
mean = 0.0  # Mean of the log-normal distribution
std_dev = 1.0  # Standard deviation of the log-normal distribution
size = 86  # Number of elements in the vector

# Generate log-normal values
log_normal_values = np.random.lognormal(mean, std_dev, size)

# Normalize the vector to sum up to 1
decay = log_normal_values / np.sum(log_normal_values)

优化方案 #

1. 用卷积替代嵌套循环（核心优化） #

你的核心计算是对序列做滑动加权求和，这完全可以用numpy.convolve直接实现，替代内层循环和矩阵构建：

# 替换原内层循环的核心逻辑
decay_total = np.convolve(diagonal_values, decay, mode='full').round(0)
decay_total_insert = decay_total[:300]

这一步将时间复杂度从O(n*k)（n为序列长度，k为decay长度）降到O(n+k)，效率提升显著。

2. 批量处理所有国家 #

不要逐个国家单独计算，直接对DataFrame的所有列做向量化操作：

import time
start_time = time.time()

# 遍历参数组合，一次性处理所有国家
for gr in np.arange(-2.0, 2.0, 0.5):
    for l_s in np.arange(-10.0, 10.0, 1.0):
        for ip in [70, 80, 90]:
            for year in np.arange(1942.0, 2100.0, 1.0):
                # 获取当前参数下所有国家的销售数据
                subs = sales_df.loc[(gr, l_s, ip, year)].values
                # 对每一列（国家）执行卷积计算
                decay_totals = np.array([np.convolve(col, decay, mode='full')[:300].round(0) for col in subs.T])
                # 批量写入结果
                subsetted_df.loc[(gr, l_s, ip, year)] = decay_totals.T

end_time = time.time()
print(f"执行时间: {end_time - start_time:.2f}秒")

3. 高效并行：用multiprocessing替代joblib #

joblib在小任务并行时的进程开销过高，改用multiprocessing.Pool对参数组合做并行，同时结合向量化：

from multiprocessing import Pool
import numpy as np
import pandas as pd

def process_param_comb(param):
    gr, l_s, ip, year = param
    subs = sales_df.loc[(gr, l_s, ip, year)].values
    decay_totals = np.array([np.convolve(col, decay, mode='full')[:300].round(0) for col in subs.T])
    return (gr, l_s, ip, year), decay_totals.T

if __name__ == "__main__":
    start_time = time.time()
    # 生成所有参数组合
    params = []
    for gr in np.arange(-2.0, 2.0, 0.5):
        for l_s in np.arange(-10.0, 10.0, 1.0):
            for ip in [70, 80, 90]:
                for year in np.arange(1942.0, 2100.0, 1.0):
                    params.append((gr, l_s, ip, year))
    # 并行处理参数组合
    with Pool() as pool:
        results = pool.map(process_param_comb, params)
    # 批量写入结果
    for (gr, l_s, ip, year), values in results:
        subsetted_df.loc[(gr, l_s, ip, year)] = values
    end_time = time.time()
    print(f"并行执行时间: {end_time - start_time:.2f}秒")

4. 数据结构优化 #

• 确保sales_df和subsetted_df的索引完全对齐，避免loc查找时的隐式转换开销；
• 若参数值为浮点数，可提前转换为整数编码（例如将gr的-2.0映射为0，-1.5映射为1等），加速索引查找；
• 用df.xs方法快速切片MultiIndex，替代多层loc查找。

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