化工参数去除率计算代码优化：提升运行速度需求

问题背景 #

需要计算污水处理厂氧化工艺后化学/生物参数的去除率，现有代码功能正常但运行缓慢：处理15×80的DataFrame时，笔记本需10秒，PC需4秒，数据量增大后效率骤降。

计算规则：

• 单步去除率：1 - n(i)/n(i-1)（i为同组内的顺序位置）
• 总去除率：1 - n(i)/n(0)（n(0)为同组内第一个测量值）
• 按测量点ID分组执行计算，结果存入原DataFrame的S_removal（单步）和T_removal（总去除率）列。

原示例代码及实现函数如下：

示例代码：

import pandas as pd
import numpy as np

data = {"ID": ["X1_P0001", "X2_P0001", "X3_P0001", "X1_P0002", "X2_P0002", "X3_P0002", "X4_P0002","X5_P0002", "X1_P0003", "X2_P0003", "X3_P0003"],
    "Measurement": [100, 80, 60, 120,90,70,50,25, 85,65,35]}
df = pd.DataFrame(data)
df["S_removal"]= np.nan
df["T_removal"]= np.nan

原实现函数：

def removal_TEST(Rem1, Measure, Rem2):
    lst = [i.split("_")[1] for i in df["ID"]] #提取相关ID信息
    y = np.unique(lst) #存储唯一ID值用于循环
    for ID in y: 

        id_list = []                           
        for i in range(0, len(df["ID"])):
            if ID in df["ID"][i]:
                id_list.append(i)
            else:                        #仅将相关ID存入新列表
                id_list.append(np.nan)

        indexlist = pd.Series(id_list)           
        first_index = indexlist.first_valid_index() #获取ID列表的首尾索引
        last_index = indexlist.last_valid_index()  
        
        col_indizes = []

        for i in range(first_index, last_index+1):
            col_indizes.append(i)
       
        for i in col_indizes:
            if i == 0:                            
                continue      #i=0无前置元素，跳过计算
                   
            else: 
                Rem1[i]= 1-(Measure[i]/Measure[i-1])
                
        Rem1[first_index]= np.nan      #每个ID的首个条目设为NaN

        for i in range(first_index, last_index+1):
            col_indizes.append(i)

        for i in range(len(Rem2)):
            for i in col_indizes:
                Rem2[i]= 1-(Measure[i]/Measure[first_index])

        Rem2[first_index]= np.nan

优化思路与实现 #

原代码大量嵌套循环且直接操作DataFrame元素，是效率低下的核心原因。利用Pandas的分组(groupby)和向量化函数可完全替代循环，大幅提升速度。

步骤1：提取分组标识 #

先从ID列中提取分组键（如P0001、P0002），避免重复计算：

df['group'] = df['ID'].str.split('_').str[1]

步骤2：计算单步去除率（S_removal） #

使用groupby结合shift函数获取同组内前一行的测量值，直接按公式计算：

# 同组内当前值除以前一行值，再计算去除率
df['S_removal'] = 1 - df['Measurement'] / df.groupby('group')['Measurement'].shift(1)
# 每组第一个值设为NaN（无前置数据）
df.loc[df.groupby('group').head(1).index, 'S_removal'] = np.nan

步骤3：计算总去除率（T_removal） #

使用groupby结合transform('first')获取每组的第一个测量值，批量计算：

# 获取每组的初始测量值
group_first = df.groupby('group')['Measurement'].transform('first')
# 计算总去除率
df['T_removal'] = 1 - df['Measurement'] / group_first
# 每组第一个值设为NaN
df.loc[df.groupby('group').head(1).index, 'T_removal'] = np.nan

完整优化代码 #

import pandas as pd
import numpy as np

data = {"ID": ["X1_P0001", "X2_P0001", "X3_P0001", "X1_P0002", "X2_P0002", "X3_P0002", "X4_P0002","X5_P0002", "X1_P0003", "X2_P0003", "X3_P0003"],
    "Measurement": [100, 80, 60, 120,90,70,50,25, 85,65,35]}
df = pd.DataFrame(data)

# 提取分组键
df['group'] = df['ID'].str.split('_').str[1]

# 计算单步去除率
df['S_removal'] = 1 - df['Measurement'] / df.groupby('group')['Measurement'].shift(1)
df.loc[df.groupby('group').head(1).index, 'S_removal'] = np.nan

# 计算总去除率
group_first = df.groupby('group')['Measurement'].transform('first')
df['T_removal'] = 1 - df['Measurement'] / group_first
df.loc[df.groupby('group').head(1).index, 'T_removal'] = np.nan

# 可选：删除临时分组列
df.drop('group', axis=1, inplace=True)

print(df)

优化效果说明 #

• 完全抛弃嵌套循环，采用Pandas原生向量化操作，处理速度提升几十到上百倍（测试同规模数据耗时可降至毫秒级）。
• 代码更简洁易读，减少出错概率，维护成本更低。
• 支持任意规模的DataFrame扩展，不会因数据量增大出现明显性能瓶颈。

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