Make Pandas groupby apply behave like transform? - pandas

Question

Is there a way to apply a function that return iterables of the same lengths as the groups, and "transform" the results as a new column?
I have some time series data of serveral groups that are interweaved.
For example:
import pandas as pd
import numpy as np
import seaborn as sns
iris_data = sns.load_dataset('iris')
df = iris_data.head(5).copy()
df = pd.concat([df, iris_data.iloc[50:55]])
df = pd.concat([df, iris_data.iloc[100:105]])
df = df.sample(frac=1, random_state=0).reset_index(drop=True)
df = df.drop(columns=["sepal_width", "petal_length", "petal_width"])
df.index = pd.date_range("2020/12/16", periods=15)
df
sepal_length species
2020-12-16 4.9 setosa
2020-12-17 6.4 versicolor
2020-12-18 5.5 versicolor
2020-12-19 6.5 versicolor
2020-12-20 6.5 virginica
2020-12-21 5.0 setosa
2020-12-22 4.7 setosa
2020-12-23 6.3 virginica
2020-12-24 6.3 virginica
2020-12-25 6.9 versicolor
2020-12-26 5.8 virginica
2020-12-27 4.6 setosa
2020-12-28 5.1 setosa
2020-12-29 7.0 versicolor
2020-12-30 7.1 virginica
I want to apply some time-dependent logic to each group (in this case, each species), and append the results as a new column.
For example, if I want to calculate an "is new-high" indicator in each series: (again THIS IS NOT THE POINT OF THE QUESTION, just an example of time-dependent calculations)
def is_new_high(group):
"""
This function simulates some order-dependent calculations as an illustration.
It labels True if the current sepal_length is longer than the previous one in the same species.
This probably can be achieved without groupby and apply, but that is not the point here.
"""
group_res = [False] * len(group)
running_max = 0
for i, (row_idx, row) in enumerate(group.iterrows()):
if row.sepal_length > running_max:
running_max = row.sepal_length
group_res[i] = True
return(group_res)
which when applied to the df
grouped_results = df.groupby(df.species).apply(is_new_high)
grouped_results
yields a Series that indicates whether an observation is a new-high in that species:
species
setosa [True, True, False, False, True]
versicolor [True, False, True, True, True]
virginica [True, False, False, False, True]
dtype: object
With a little bit of (inefficient) Pandas trickery (see below), this can be merged back into the original df:
sepal_length species new_high
2020-12-16 4.9 setosa True
2020-12-17 6.4 versicolor True
2020-12-18 5.5 versicolor False
2020-12-19 6.5 versicolor True
2020-12-20 6.5 virginica True
2020-12-21 5.0 setosa True
2020-12-22 4.7 setosa False
2020-12-23 6.3 virginica False
2020-12-24 6.3 virginica False
2020-12-25 6.9 versicolor True
2020-12-26 5.8 virginica False
2020-12-27 4.6 setosa False
2020-12-28 5.1 setosa True
2020-12-29 7.0 versicolor True
2020-12-30 7.1 virginica True
Since the function applied is time- or order-dependent, I cannot think of a way to turn it into a Pandas "transform" function that acts on columns and not on a group (i.e., DataFrame). So it seems I'm stuck with using apply that returns a vector for each group.
My question is the following: Is there a way to apply a function that return iterables of the same lengths as the groups, and "transform" the results as a new column, to achieve the same result as in my example?
Here's the inefficient and ad-hoc way I did it:
def is_new_high(group):
group_res = [(idx, False) for idx in group.index]
temp = 0
for i, (row_idx, row) in enumerate(group.iterrows()):
if row.sepal_length > temp:
temp = row.sepal_length
group_res[i] = (row_idx, True)
return(group_res)
grouped_results_exploded = grouped_results.explode()
res = pd.DataFrame(grouped_results_exploded.tolist(), index=grouped_results_exploded.index,)
res.columns=['timestamp', 'new_high']
res = res.reset_index().set_index("timestamp")
df.join(res, rsuffix="_groupby")