时间戳

首页 > 解决方案 > 如何在 Python 中对这个 for 循环进行矢量化？

python - 如何在 Python 中对这个 for 循环进行矢量化？

问题描述

我的数据有 4 列 AD 有整数。我正在添加一个新列 E，它的第一个值与 D 列中的第一个值相同。如果 E 列中的前一个值为负数，E 中的下一个值应该是 D 列中的对应值，否则它在列中取对应值C。

import pandas as pd
import numpy as np
from pandas import Series, DataFrame
data=pd.read_excel('/Users/xxxx/Documents/PY Notebooks/Data/yyyy.xlsx')
data1=data.copy()
data1['E']=np.nan
data1.at[0,'E']=data1['D'][0]
l=len(data1)
for i in range(l-1):
    if data1['E'][i]<0:
        data1.at[i+1,'E']=data1['D'][i+1]
    else:
        data1.at[i+1,'E']=data1['C'][i+1]

标签： pythonpandasfor-loopvectorization

解决方案

TL;DR：转到基准代码并使用方法 1。

简答

没有。矢量化是不可能的。

长答案

定理：对于这个特定的任务，给定行的输出不能使用任何小于该行的部分长度的向后滚动窗口的有限长度来确定。

因此，无法以矢量化方式处理此输出逻辑。（有关在 CPU 中执行矢量化的想法，请参阅此答案）。只能从数据帧的开头计算输出。

证明：考虑一个数据框的目标行df。假设有一个大小为 的向后滚动窗口，因此窗口之前存在n < partial length一个先前的值。df["E"]我们用 表示这个先前的值state。

考虑一个特殊情况：df["C"] == -1并且df["D"] == 1在窗口内。

  • 情况 1 ( state < 0)：此滚动窗口内的输出将是 [1, -1, 1, -1, .....]，使最后一个元素(-1)^(n-1)
  • 情况 2 ( state >= 0)：输出将是 [-1, 1, -1, 1, .....]，使得最后一个元素(-1)^(n)

因此，目标行的输出df["E"]可能依赖于窗口外的状态变量。QED。

有用的答案

虽然矢量化是不可能的，但这并不意味着不能实现显着的加速。一种简单但非常有效的方法是使用 anumba-compiled generator来执行顺序生成。它只需要将您的逻辑重新写入生成器函数并添加两行：

import numba

@numba.njit
def my_generator_func():
    ....

当然，您可能必须先安装 numba。如果这是不可能的，那么使用没有 numba 优化的普通生成器也可以。

基准

基准测试在 i5-8250U (4C8T) 笔记本电脑上执行，配备 16GB RAM，运行 64 位 debian 10。Python 版本为 3.7.9，pandas 为 1.1.3。n = 10^7（1000 万）条记录用于基准测试。

结果：

1. numba-njit: 2.48s
2. plain generator (no numba): 5.13s
3. original: 271.15s

> 100x可以针对原始代码实现效率增益。

代码

from datetime import datetime
import pandas as pd
import numpy as np

n = 10000000  # a large number of rows
df = pd.DataFrame({"C": -np.ones(n), "D": np.ones(n)})
#print(df.head())

# ========== Method 1. generator + numba njit ==========
ti = datetime.now()

import numba

@numba.njit
def gen(plus: np.array, minus: np.array):
    l = len(plus)
    assert len(minus) == l
    # first
    state = minus[0]
    yield state
    # second to last
    for i in range(l-1):
        state = minus[i+1] if state < 0 else plus[i+1]
        yield state

df["E"] = [i for i in gen(df["C"].values, df["D"].values)]

tf = datetime.now()
print(f"1. numba-njit: {(tf-ti).total_seconds():.2f}s")  # 1. numba-njit: 0.47s

# ========== Method 2. Generator without numba ==========
df = pd.DataFrame({"C": -np.ones(n), "D": np.ones(n)})
ti = datetime.now()

def gen_plain(plus: np.array, minus: np.array):
    l = len(plus)
    assert len(minus) == l
    # first
    state = minus[0]
    yield state
    # second to last
    for i in range(l-1):
        state = minus[i+1] if state < 0 else plus[i+1]
        yield state

df["E"] = [i for i in gen_plain(df["C"].values, df["D"].values)]

tf = datetime.now()
print(f"2. plain generator (no numba): {(tf-ti).total_seconds():.2f}s")  #

# ========== Method 3. Direct iteration ==========
df = pd.DataFrame({"C": -np.ones(n), "D": np.ones(n)})
ti = datetime.now()

# code provided by the OP
df['E']=np.nan
df.at[0,'E'] = df['D'][0]
l=len(df)
for i in range(l - 1):
    if df['E'][i] < 0:
        df.at[i+1,'E'] = df['D'][i+1]
    else:
        df.at[i+1,'E'] = df['C'][i+1]

tf = datetime.now()
print(f"3. original: {(tf-ti).total_seconds():.2f}s") # 2. 26.61s

推荐阅读