How to access updated previous row data when iterating through a dataframe (using itertuples/iterrows) without using df.at/df.loc?

user_19337821 0
Viewed 31

I am currently working with a dataframe where calculations must be done sequentially (i.e the calculation for each row requires the calculation from the previous row). As such vectorization is simply not possible (as df.shift is not viable) and therefore I must iterate.

I have written a function which does these calculations successfully using a for loop and df.itertuples(). However, to access the newly calculated value from the previous row I am using df.at[index-1, "Column Name"]. And although this works, due to the vast number of columns on which I must conduct calculations, accessing so many values by index has bottlenecked my code.

Here is a simplified example of what my code looks like:

(these are arbitrary calculations but I hope they get my thought process across)

df=pd.DataFrame({"Column 0":[1,2,3,4,5,6,7,8,9,10]})
df["Column1"]=np.nan
df["Column2"]=np.nan
df["Column3"]=np.nan
for row in df.itertuples():
  index=row.Index
  if index==0:
    df.at[index, "Column1"]=0
    df.at[index, "Column2"]=1
    df.at[index, "Column3"]=10
  else:   
    df.at[index, "Column1"]=df.at[index-1, "Column2"]+df.at[index-1, "Column3"]
    df.at[index, "Column2"]=df.at[index, "Column1"]+df.at[index-1, "Column3"]
    df.at[index, "Column3"]=df.at[index, "Column2"]*row.Column0 + df.at[index-1, "Column1"]

So my question; is it is possible to avoid having to access these newly calculated values by index and index-1?

pandas vectorization dataframe python loops
1 Answers

Example how to speed up the code with numba:

from numba import jit
from timeit import timeit


def create_df():
    df = pd.DataFrame({"Column 0": np.arange(1, 1000)})   # <--- here I use bigger dataframe to show up more the speed difference

    df["Column1"] = np.nan
    df["Column2"] = np.nan
    df["Column3"] = np.nan

    return df


def fn1(df):                            # <--- standard pandas function
    for index, row in df.iterrows():
        if index == 0:
            df.at[index, "Column1"] = 0
            df.at[index, "Column2"] = 1
            df.at[index, "Column3"] = 10
        else:
            df.at[index, "Column1"] = (
                df.at[index - 1, "Column2"] + df.at[index - 1, "Column3"]
            )
            df.at[index, "Column2"] = (
                df.at[index, "Column1"] + df.at[index - 1, "Column3"]
            )
            df.at[index, "Column3"] = (
                df.at[index, "Column2"] * row["Column 0"]
                + df.at[index - 1, "Column1"]
            )
    return df


@jit(nopython=True)
def fn2(vals):                        # <---- function that is Jit-ed with Numba
    for row in range(len(vals)):
        if row == 0:
            vals[row][1] = 0
            vals[row][2] = 1
            vals[row][3] = 10
        else:
            vals[row][1] = vals[row - 1][2] + vals[row - 1][3]
            vals[row][2] = vals[row][1] + vals[row - 1][3]
            vals[row][3] = vals[row][2] * vals[row][0] + vals[row - 1][1]
    return vals


d1 = fn1(create_df())

d2 = create_df()
d2[:] = fn2(d2.values)                  # <--- compute values with numpy and assign it back to Dataframe

assert d1.eq(d2).all().all() == True    # <---- make sure that output from standard function is the same as output from Numba function

t1 = timeit("fn1(df)", "df=create_df()", number=100, globals=globals())
t2 = timeit(
    "df[:] = fn2(df.values)", "df=create_df()", number=100, globals=globals()
)
print(t1)
print(t2)

Prints on my machine (Python 3.9.7/Pandas 1.4.2/numba 0.55.1):

7.869270268129185
0.06915179989300668

The Numba function is ~100 times faster.

user_100359850
Related