在NumPy中向量化带充放电递归约束的电池SOC更新

你的电池SOC模拟属于状态依赖的序列递推问题，每一步的充放电决策和SOC更新都强依赖前一步的结果，纯numpy无法实现完全无循环的向量化（因为向量化操作要求元素计算独立或有固定并行规律）。但可以通过以下两种高效方案实现加速，同时保证逻辑完全对齐原代码：

方案1：Numba JIT编译（最推荐，简单高效） #

Numba可以将Python循环编译为机器码，速度比纯Python循环快几十至上百倍，且完全保留原逻辑，不需要修改算法。

代码实现 #

import numpy as np
import numba

@numba.njit()  # 启用No-Python模式JIT编译，首次运行编译后速度极快
def simulate_battery(E_max, P_B_max, SOC_min, surplus, remaining, initial_SOC):
    total_timesteps = len(surplus)
    SOC = np.zeros(total_timesteps + 1)
    SOC[0] = initial_SOC
    charge = np.zeros(total_timesteps)
    discharge = np.zeros(total_timesteps)
    
    for t in range(total_timesteps):
        current_SOC = SOC[t]
        charge[t] = min(surplus[t], P_B_max, (E_max - current_SOC))
        discharge[t] = min(remaining[t], P_B_max, (current_SOC - SOC_min))
        SOC[t+1] = np.clip(current_SOC + charge[t] - discharge[t], SOC_min, E_max)
    
    return SOC, charge, discharge

# 参数初始化
E_max = 100.0
P_B_max = 10.0
SOC_min = 20.0
surplus = np.array([5, 3, 0, 7, 8, 0, 0, 5, 4, 3])
remaining = np.array([0, 2, 5, 1, 0, 4, 6, 0, 2, 5])
initial_SOC = 0.5 * E_max

# 执行模拟
SOC, charge, discharge = simulate_battery(E_max, P_B_max, SOC_min, surplus, remaining, initial_SOC)

优势 #

• 代码几乎和原逻辑完全一致，无需重构
• 编译后速度接近C语言，处理大规模时间步时优势明显
• 支持多线程并行（可添加parallel=True参数进一步优化）

方案2：numpy状态封装递推（技术探索，实用性有限） #

由于numpy的accumulate仅支持单状态或简单多状态的累积，对于本场景的多变量递推，仍需显式循环，但可将状态更新逻辑抽象为独立函数，提升代码可读性：

代码实现 #

import numpy as np

def update_state(prev_state, current_vals):
    """自定义状态更新函数：输入前一步状态，输出当前步状态"""
    prev_SOC, _, _ = prev_state
    s, r = current_vals
    charge = min(s, P_B_max, E_max - prev_SOC)
    discharge = min(r, P_B_max, prev_SOC - SOC_min)
    new_SOC = np.clip(prev_SOC + charge - discharge, SOC_min, E_max)
    return (new_SOC, charge, discharge)

# 参数初始化
E_max = 100.0
P_B_max = 10.0
SOC_min = 20.0
surplus = np.array([5, 3, 0, 7, 8, 0, 0, 5, 4, 3])
remaining = np.array([0, 2, 5, 1, 0, 4, 6, 0, 2, 5])
total_timesteps = len(surplus)

# 准备输入序列：将surplus和remaining配对
input_seq = np.stack([surplus, remaining], axis=1)
# 初始化状态数组：存储每一步的(SOC, charge, discharge)
state_history = np.zeros((total_timesteps + 1, 3))
state_history[0] = (0.5 * E_max, 0.0, 0.0)

# 递推更新状态
for i in range(total_timesteps):
    state_history[i+1] = update_state(state_history[i], input_seq[i])

# 拆分结果
SOC = state_history[:, 0]
charge = state_history[1:, 1]
discharge = state_history[1:, 2]

说明 #

这种方式并未真正移除循环，只是将逻辑解耦，代码可读性略有提升，但速度和纯Python循环无明显差异，仅适合小规模场景测试。

为什么完全无循环的向量化不可行？ #

本场景的核心矛盾是：每一步的充放电量计算动态依赖前一步的SOC结果，而SOC又受之前所有步的充放电决策影响。numpy的向量化操作要求所有元素的计算是独立的，或依赖固定位置的预计算值，无法表达这种动态递推的依赖关系，因此无法实现完全无循环的纯numpy向量化。

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