Pandas: How to compute a conditional rolling/accumulative maximum within a group

user_12427902 0
Viewed 199

I would like to achieve the following results in the column condrolmax (based on column close) (conditional rolling/accumulative max) without using a stupidly slow for loop.

Index    close    bool       condrolmax
0        1        True       1
1        3        True       3
2        2        True       3
3        5        True       5
4        3        False      5
5        3        True       3 --> rolling/accumulative maximum reset (False cond above)
6        4        True       4
7        5        False      4
8        7        False      4
9        5        True       5 --> rolling/accumulative maximum reset (False cond above)
10       7        False      5
11       8        False      5
12       6        True       6 --> rolling/accumulative maximum reset (False cond above)
13       8        True       8
14       5        False      8
15       5        True       5 --> rolling/accumulative maximum reset (False cond above)
16       7        True       7
17       15       True       15
18       16       True       16

The code to create this dataframe:

# initialise data of lists.
data = {'close':[1,3,2,5,3,3,4,5,7,5,7,8,6,8,5,5,7,15,16],
        'bool':[True, True, True, True, False, True, True, False, False, True, False,
                False, True, True, False, True, True, True, True],
        'condrolmax': [1,3,3,5,5,3,4,4,4,5,5,5,6,8,8,5,7,15,16]}
 
# Create DataFrame
df = pd.DataFrame(data)

I am sure it is possible to vectorize that (one liner). Any suggestions ?

Thanks again !

pandas conditional-statements python
3 Answers

First make groups using your condition (bool changing from False to True) and cumsum, then apply your rolling after a groupby:

group = (df['bool']&(~df['bool']).shift()).cumsum()
df.groupby(group)['close'].rolling(2, min_periods=1).max()

output:

0     0      1.0
      1      3.0
      2      3.0
      3      5.0
      4      5.0
1     5      3.0
      6      4.0
      7      5.0
      8      7.0
2     9      5.0
      10     7.0
      11     8.0
3     12     6.0
      13     8.0
      14     8.0
4     15     5.0
      16     7.0
      17    15.0
      18    16.0
Name: close, dtype: float64

To insert back as a column:

df['condrolmax'] = df.groupby(group)['close'].rolling(2, min_periods=1).max().droplevel(0)

output:

    close   bool  condrolmax
0       1   True         1.0
1       3   True         3.0
2       2   True         3.0
3       5   True         5.0
4       3  False         5.0
5       3   True         3.0
6       4   True         4.0
7       5  False         5.0
8       7  False         7.0
9       5   True         5.0
10      7  False         7.0
11      8  False         8.0
12      6   True         6.0
13      8   True         8.0
14      5  False         8.0
15      5   True         5.0
16      7   True         7.0
17     15   True        15.0
18     16   True        16.0

NB. if you want the boundary to be included in the rolling, use min_periods=1 in rolling

user_163434640

You can set group and then use cummax(), as follows:

# Set group: New group if current row `bool` is True and last row `bool` is False
g = (df['bool'] & (~df['bool']).shift()).cumsum()   

# Get cumulative max of column `close` within the group 
df['condrolmax'] = df.groupby(g)['close'].cummax()

Result:

print(df)

    close   bool  condrolmax
0       1   True           1
1       3   True           3
2       2   True           3
3       5   True           5
4       3  False           5
5       3   True           3
6       4   True           4
7       5  False           5
8       7  False           7
9       5   True           5
10      7  False           7
11      8  False           8
12      6   True           6
13      8   True           8
14      5  False           8
15      5   True           5
16      7   True           7
17     15   True          15
18     16   True          16
user_150706970

I'm not sure how we can use linear algebra and vectorizing to make this function faster, but using list comprehension, we write a faster algorithm. First, define the function as:

def faster_condrolmax(df):
    df['cond_index'] = [df.index[i] if df['bool'][i]==False else 0 for i in 
    df.index]
    df['cond_comp_index'] = [np.max(df.cond_index[0:i]) for i in df.index]
    df['cond_comp_index'] = df['cond_comp_index'].fillna(0).astype(int)
    df['condrolmax'] = np.zeros(len(df.close))
    df['condrolmax'] = [np.max(df.close[df.cond_comp_index[i]:i]) if 
               df.cond_comp_index[i]<i else df.close[i] for 
               i in range(len(df.close))]
    return df

Then, you can use:

!pip install line_profiler
%load_ext line_profiler

to add and load the line profiler and see how long each line of the code takes with this:

%lprun -f faster_condrolmax faster_condrolmax(df)

which will result as: Each line profiling results

of just see how long the whole function takes:

%timeit faster_condrolmax(df)

which will result as: Total algorithm profiling result

If you use the SeaBean's function, you can get better results half the speed it takes for my proposed functions. However, the speed estimated for SeaBean's doesn't seem robust, and to estimate his functions, you should run it on a larger dataset and then decide. That's all because %timeit reports like this: SeaBean's function profiling result

user_150940660
Related