如何在 Python 中向量化这个 for 循环？

Question

我的数据有 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]

Answer 1

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。

有用的答案

虽然矢量化是不可能的，但这并不意味着不能实现显着的加速。 一种简单但非常有效的方法是使用numba-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与原始代码相比，可以实现> 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

Answer 2

我不认为您可以矢量化此操作，因为您有需要运行先前计算的相关行。 话虽如此，您的功能仍有很大的优化空间。 让我们首先用一些随机点检查您的原始实现。

import numpy as np
import pandas as pd
import time

size = 10000000
data = np.random.randint(-2, 10, size=size)
data = data.reshape([size//4, 4])

time_start = time.time()
df = pd.DataFrame(data=data, columns=["A", "B", "C", "D"])
df['E']=np.nan
df.at[0,'E'] = df['D'][0]
for i in range(len(df)-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]

print(f"Operation on pd df took {time.time() - time_start} seconds.")

输出：

Operation on pd df took 84.00791883468628 seconds.

由于对 DataFrame 的操作通常很慢，我们可以改为对底层的 numpy 数组进行操作。

time_start = time.time()
df = pd.DataFrame(data=data, columns=["A", "B", "C", "D"])
c_vals = df["C"].values
d_vals = df["D"].values

e_vals = [d_vals[0]]
last_e = e_vals[0]
for i in range(len(df)-1):
    if last_e < 0:
        last_e = d_vals[i+1]
    else:
        last_e = c_vals[i+1]
    e_vals.append(last_e)

df['E'] = e_vals

print(f"Operation on np array took {time.time() - time_start} seconds.")

输出：

Operation on np array took 2.2387869358062744 seconds.

现在我们可以争辩说 Python 中的 for 循环很慢，我们可以使用可以处理 numpy 数组的 JIT 编译器，例如 numba。

import numba

time_start = time.time()
df = pd.DataFrame(data=data, columns=["A", "B", "C", "D"])
c_vals = df["C"].values
d_vals = df["D"].values

@numba.jit(nopython=True)
def numba_calc(c_vals, d_vals):
    e_vals = [d_vals[0]]
    last_e = e_vals[0]
    for i in range(len(c_vals)-1):
        if last_e < 0:
            last_e = d_vals[i+1]
        else:
            last_e = c_vals[i+1]
        e_vals.append(last_e)

    return e_vals

df["E"] = numba_calc(c_vals, d_vals)

print(f"Operation on np array with numba took {time.time() - time_start} seconds.")

输出：

Operation on np array with numba took 1.2623450756072998 seconds.

因此，特别是对于使用 numba 的更大的 DataFrames 会付出代价，而在原始 numpy 数组上操作通常会提供很好的运行时改进。