Calculate the duration of a state with a pandas Dataframe

Question

I try to calculate how often a state is entered and how long it lasts. For example I have the three possible states 1,2 and 3, which state is active is logged in a pandas Dataframe:

test = pd.DataFrame([2,2,2,1,1,1,2,2,2,3,2,2,1,1], index=pd.date_range('00:00', freq='1h', periods=14))

For example the state 1 is entered two times (at index 3 and 12), the first time it lasts three hours, the second time two hours (so on average 2.5). State 2 is entered 3 times, on average for 2.66 hours.

I know that I can mask data I'm not interested in, for example to analyize for state 1:

state1 = test.mask(test!=1)

but from there on I can't find a way to go on.

Answer 1

I hope the comments give enough explanation - the key point is you can use a custom rolling window function and then cumsum to group the rows into "clumps" of the same state.

# set things up
freq = "1h"
df = pd.DataFrame(
    [2,2,2,1,1,1,2,2,2,3,2,2,1,1],
    index=pd.date_range('00:00', freq=freq, periods=14)
)

# add a column saying if a row belongs to the same state as the one before it
df["is_first"] = pd.rolling_apply(df, 2, lambda x: x[0] != x[1]).fillna(1)

# the cumulative sum - each "clump" gets its own integer id
df["value_group"] = df["is_first"].cumsum()

# get the rows corresponding to states beginning
start = df.groupby("value_group", as_index=False).nth(0)
# get the rows corresponding to states ending
end = df.groupby("value_group", as_index=False).nth(-1)

# put the timestamp indexes of the "first" and "last" state measurements into
# their own data frame
start_end = pd.DataFrame(
    {
        "start": start.index,
        # add freq to get when the state ended
        "end": end.index + pd.Timedelta(freq),
        "value": start[0]
    }
)
# convert timedeltas to seconds (float)
start_end["duration"] = (
    (start_end["end"] - start_end["start"]).apply(float) / 1e9
)
# get average state length and counts
agg = start_end.groupby("value").agg(["mean", "count"])["duration"]
agg["mean"] = agg["mean"] / (60 * 60)

And the output:

           mean  count
value                 
1      2.500000      2
2      2.666667      3
3      1.000000      1