How to group entries by consecutive dates in Python/Pandas

Question

I have a pandas series, called hot_days that looks like the following:

0     1980-06-04
1     1981-08-05
2     1982-06-04
3     1982-06-05
4     1982-07-08
         ...    
294   2019-07-25
295   2019-08-24
296   2019-08-25
297   2019-08-26
298   2019-08-27

It is a list of dates where the temperature in a given location is above a threshold. I want to detect and record when a heatwave occurs, which is when the temperature is over this threshold for three or more days. I want to end up with a dataframe containing the date the heatwave started, and its length. By applying:

new_series = (hot_days == hot_days.shift(2)+pd.Timedelta("2 days")) * (hot_days.groupby((hot_days == hot_days.shift(2)+pd.Timedelta("2 days")).cumsum()).cumcount()+1)

I get the series:

1      0
2      0
3      0
4      0
      ..
294    1
295    0
296    0
297    1
298    1

Which has a 1 for dates during a heatwave, and 0 for dates that are not in a heatwave, which I believe is a step in the right direction. However, since I'm new to pandas, I'm not quite sure how I can achieve my goal. I know I can use loops, however I understand this is 'un-pythonic' as loops are slow in Python, so I'd rather find a more elegant solution (although the dataset is small enough that loops will work in a reasonable amount of time).

Answer 1

Let's call s the initial Series.

Identify the heat wave days:

waves = s.eq(s.shift(1)+pd.DateOffset(days=1)) & s.eq(s.shift(2)+pd.DateOffset(days=2))

Create a DataFrame with wave and wave groups:

df = pd.concat({'date': s,
                'wave': waves,
                'group': waves.diff(1).ne(0).cumsum()
                }, axis=1)

List the waves and their duration:

pd.DataFrame({gid: pd.Series({'start': g.iloc[0]['date'],
                              'end': g.iloc[-1]['date'],
                              'duration': len(g)})
              for gid, g in df[df['wave']].groupby('group')
              }).T

output:

       start        end duration
2 2019-08-26 2019-08-27        2
         

NB. I have slightly different results due ton incomplete dataset

edit: here is how the waves.diff(1).ne(0).cumsum() works:

    bool   diff  diff_int  diff_not  diff_not_int  diff_not_cumsum
0   True    NaN       NaN      True             1                1
1  False   True      -1.0      True             1                2
2  False  False       0.0     False             0                2
3   True   True       1.0      True             1                3
4   True  False       0.0     False             0                3

Answer 2

We could use shifting to get an added count so you can later count up in a loop.

s = pd.Series([0,1,1,0,0,0,1,1,1,1,1,0,0,1,1,0,0,0,0,1,1,0,0,1,0,0])
s * (s.groupby((y != s.shift()).cumsum()).cumcount() + 1)

To get:

0     0
1     1
2     2
3     0
4     0
5     0
6     1
7     2
8     3
9     4
10    5
11    0
12    0
13    1
14    2
15    0
16    0
17    0
18    0
19    1
20    2
21    0
22    0
23    1
24    0
25    0

Or we can loop through a group by in order to get separate lists.

df = pd.DataFrame({"a":s})
for i, g in df.groupby([(df.a != df.a.shift()).cumsum()]):
    print (i,end="")
    print (g)
    print (g.a.tolist())
    print("--")

To get:

1   a
0  0
[0]
--
2   a
1  1
2  1
[1, 1]
--
3   a
3  0
4  0
5  0
[0, 0, 0]
--
4    a
6   1
7   1
8   1
9   1
10  1
[1, 1, 1, 1, 1]
--
5    a
11  0
12  0
[0, 0]
--
6    a
13  1
14  1
[1, 1]
--
7    a
15  0
16  0
17  0
18  0
[0, 0, 0, 0]
--
8    a
19  1
20  1
[1, 1]
--
9    a
21  0
22  0
[0, 0]
--
10    a
23  1
[1]
--
11    a
24  0
25  0
[0, 0]
--