如何加速含iloc的大循环？列车电池充放电周期电量计算优化

我有一个列车电池相关的DataFrame，包含Timestamp（每10秒采样一次）、Current（安培）和Voltage列。已知规则：

• 电流<0时，电池为列车供电
• 电流>10时，处于充电模式
• 0≤电流≤10时，电池维持100%电量

电池周期定义为：从放电开始（为列车供电）到电量回到100%时结束。我需要计算每个周期内电池供给列车的电量（公式：Current*Hour）。

当前实现步骤：

# 创建存储电量的空列
df['Energy_T2'] = "";

# 只保留供电或维持模式的数据（排除充电模式）
df_fullp = df_full[df_full.HMI_IBatt_T2 < 10];
# 注：HMI_IBatt_T2 是电流列

下一步需要遍历电流列，识别周期起始点，将每个时间间隔的供电量存入df['Energy_T2']，再按周期求和。

数据准备代码：

c = df_fullp.columns.get_loc('HMI_IBatt_T2');
d = df_fullp.columns.get_loc('Energy_T2');
e = df_fullp.columns.get_loc('Timestamp');
Vehiclegp = df_fullp.groupby(['Vehicle']);

由于样本对应不同车辆的电池，通过groupby按车辆分组。

我的循环实现如下，但针对26900725行数据运行极慢，无法在合理时间内完成：

df_fullp.iloc[0, df_fullp.columns.get_loc('Energy_T2')] = 0; 

for v in df_fullp.Vehicle.unique():
    s1 = (Vehiclegp.get_group(v));
    for i in range(2, df_fullp.shape[0]):
        if s1.iloc[i,c] >= 0: 
            if s1.iloc[i+1,c] >= 0: 
                df_fullp.iloc[i,d] = 0; # 维持模式，电量为0
            else: # 计算供电量：时间差*电流(A*h)
                df_fullp.iloc[i,d] = (s1.iloc[i,e]-s1.iloc[i-1,e])*s1.iloc[i,c]/3600;  
        else:
            df_fullp.iloc[i,d] = (s1.iloc[i,e]-s1.iloc[i-1,e])*s1.iloc[i,c]/3600;

我刚接触Python编程，对可用库了解有限，请问如何加速这段代码？

针对2600万行的大数据量，Python嵌套循环的效率完全无法满足需求，必须改用Pandas的矢量化操作和分组批量处理，以下是具体优化步骤：

1. 用矢量化计算替代逐行循环 #

Pandas的矢量化操作基于底层C实现，速度比Python循环快100~1000倍。先处理时间差和单条记录的供电量：

import pandas as pd

# 确保Timestamp是datetime类型（如果还不是的话）
df_fullp['Timestamp'] = pd.to_datetime(df_fullp['Timestamp'])

# 计算每条记录与前一条的时间间隔（转换为小时）
df_fullp['hour_diff'] = df_fullp['Timestamp'].diff().dt.total_seconds() / 3600

# 初始化Energy_T2为0
df_fullp['Energy_T2'] = 0.0

# 仅在电流<0时计算供电量，其余情况保持0
supply_mask = df_fullp['HMI_IBatt_T2'] < 0
df_fullp.loc[supply_mask, 'Energy_T2'] = df_fullp.loc[supply_mask, 'HMI_IBatt_T2'] * df_fullp.loc[supply_mask, 'hour_diff']

注：因为采样间隔固定为10秒，也可以直接用10/3600代替hour_diff，进一步提升速度：

df_fullp.loc[supply_mask, 'Energy_T2'] = df_fullp.loc[supply_mask, 'HMI_IBatt_T2'] * (10/3600)

2. 按车辆分组识别电池周期并求和 #

电池周期的核心是从放电开始（电流从≥0变为<0）到回到维持模式并稳定，我们可以通过标记周期起始点来生成周期ID，再按ID求和：

def process_single_vehicle(group):
    # 确保车辆数据按时间排序（关键，否则周期识别会出错）
    group = group.sort_values('Timestamp').reset_index(drop=True)
    
    # 标记周期起始点：前一条记录电流≥0，当前记录电流<0
    group['cycle_start'] = (group['HMI_IBatt_T2'].shift(1) >= 0) & (group['HMI_IBatt_T2'] < 0)
    
    # 生成周期ID：累加起始标记，每次起始对应一个新周期
    group['cycle_id'] = group['cycle_start'].cumsum()
    
    # 过滤掉没有放电的无效周期（cycle_id=0的行）
    group = group[group['cycle_id'] != 0]
    
    # 按周期ID计算总供电量，并合并回原数据
    cycle_total = group.groupby('cycle_id')['Energy_T2'].sum().reset_index(name='Total_Cycle_Energy')
    return group.merge(cycle_total, on='cycle_id', how='left')

# 按Vehicle分组处理所有车辆
df_processed = df_fullp.groupby('Vehicle').apply(process_single_vehicle)

3. 额外优化建议 #

• 数据类型优化：将HMI_IBatt_T2转换为float32或int32（如果精度允许），减少内存占用和计算开销：

  df_fullp['HMI_IBatt_T2'] = df_fullp['HMI_IBatt_T2'].astype('float32')
  

• 避免重复列定位：不要用columns.get_loc()反复获取列索引，直接用列名访问更高效
• 内存管理：如果内存不足，可以用df_fullp = df_fullp.drop(columns=['hour_diff'])删除临时列，或者使用Dask库处理超大数据集

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