speed up/optimize pandas code I am making a data-driven bond screening and I have as input a big dataset of 1526 columns and 2412 rows. For 10 columns it takes 2 minutes processing time at the moment, which is too much. The following function takes 90% of the time:

Input of the function is df: a pandas series, where the index is a time series and the first column has floats, like this:

https://imgur.com/a/3pQSQZC

def future_returns(df):
    grid_columns = np.arange(len(df))
    grid = pd.DataFrame(index=df.index, columns=grid_columns)

    # fill grid with copies of df, shifted 1 element forward for each column
    for no, idx in enumerate(grid.columns):
        grid.loc[:, idx] = df.shift(-no)

    # calculate future returns from every point in the index
    future_returns = grid.divide(grid.iloc[:, 0], axis=0) - 1

    return future_returns
code

Your code itself is clear, with and has only a few improvements

df

your parameter df expects actually a Series, and not a DataFrame, so I would rename this.

composition

now you first make an empty DataFrame and then change the values. More clear would be to generate it directly with the correct data:

def future_returns_2(data):
    grid = pd.DataFrame(
        index=data.index, 
        columns=np.arange(len(data)), 
        data=[data.shift(-i).values for i in range(len(data))],
    )
    return grid.divide(data, axis=0) - 1


Conveniently, this is also about faster

numpy

If you really want it a lot faster, you should stay in the numpy space for as long as possible, and only generate the DataFrame at the last possible time.

You can use numpy.roll

arr = data.values
result = np.array(
    [np.roll(arr, -i) for i in range(len(arr))],
    copy=False,
) / arr - 1 


Since numpy.roll doesn't make the lower triangle of the result NaN, You should add this yourself:

mask = np.rot90(np.tri(l,), k=-1)
mask[np.where(1 - mask)] = np.nan
mask



array([[ 1.,  1.,  1.,  1.,  1.],
       [ 1.,  1.,  1.,  1., nan],
       [ 1.,  1.,  1., nan, nan],
       [ 1.,  1., nan, nan, nan],
       [ 1., nan, nan, nan, nan]])



Now you can deduct this mask instead of 1

def future_returns_numpy(data):
    arr = data.values
    l = len(arr)

    mask = np.rot90(np.tri(l), k=-1)
    mask[np.where(1 - mask)] = np.nan

    result = np.array(
        [np.roll(arr, -i) for i in range(l)], 
        copy=False,
    ) / arr - mask

    return pd.DataFrame(data = result.T, index = data.index)


I find this code less clear than the pandas algorithm, but if speed is important, I would use this.

timings

For this dummy data

size=1000
np.random.seed(0)
data = pd.Series(
    np.random.random(size), 
    index= pd.date_range(start='20190101', freq='1d', periods = size),
)



OP: 10.4 s ± 528 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
future_returns_2: 722 ms ± 29.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
future_returns_numpy: 79 ms ± 7.62 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)