CODEWITHSUNDEEP
pythonpandasdataframediffmean
MarcoC
MarcoC

Reputation: 119

Calculating mean and diff conditional to date values

I have the following dataframe, where a given assignment work_id is performed by a student s_id in a date work_date with a relative score score. For each student, the dates are sorted in descending order.

df = pd.DataFrame(columns=['work_id', 's_id', 'score','work_date'],
...                   data =[['a3', 'p01', np.nan,'2020-05-01'],
...                          ['a2', 'p01',10,'2020-06-10'],
...                          ['a1','p01', 5, '2020-06-15'],
...                          ['a5','p02', 5, '2019-10-10'],
...                          ['a7','p02', 11, '2020-03-01'],
...                          ['a6','p02', np.nan, '2020-04-01'],
...                          ['a4','p02', 4, '2020-06-20'],
...                          ])

>>> df
  work_id s_id  score   work_date
0      a3  p01    NaN  2020-05-01
1      a2  p01   10.0  2020-06-10
2      a1  p01    5.0  2020-06-15
3      a5  p02    5.0  2019-10-10
4      a7  p02   11.0  2020-03-01
5      a6  p02    NaN  2020-04-01
6      a4  p02    4.0  2020-06-20

I would like to add two columns: mean_score and diff_score. Column mean_score should show the averaged score obtained by each student, where the average is calculated including all the scores obtained in the previous assignments. Column diff_score should contain the difference between the current score and the previous one (which is not a NaN). The final dataframe must thus look like this:

work_id s_id  score   work_date  mean_score  diff_score
0      a3  p01    9.0  2020-05-01         NaN         NaN
1      a2  p01   10.0  2020-06-10    10.00000         NaN
2      a1  p01    5.0  2020-06-15    7.500000        -5.0
3      a5  p02    5.0  2019-10-10    5.000000         NaN
4      a7  p02   11.0  2020-03-01    8.000000         6.0
5      a6  p02    NaN  2020-04-01         NaN         NaN
6      a4  p02    4.0  2020-06-20    6.666667        -7.0

I can get this in a cumbersome way, by defining the following two functions (which take care of the possible presence of NaN entries) and using apply/lambda:

def calculate_mean(workid):
    date = df[df.work_id == workid].work_date.iloc[0]
    sid = df[df.work_id == workid].s_id.iloc[0]
    if df[(df.work_id==workid) & (df.s_id==sid) & (df.work_date == date)].score.notnull().item():
        mean = df[(df.s_id == sid) & (df.work_date <= date)].score.mean()
    else:
        mean = np.nan
    return mean

def calculate_diff(workid):
    date = df[df.work_id == workid].work_date.iloc[0]
    sid = df[df.work_id == workid].s_id.iloc[0]
    try:
        if df[(df.s_id==sid) & (df.work_date == date)].score.notnull().item():
            delta = df[(df.s_id == sid) & (df.work_date <= date) & (df.score.notnull())].score.diff().iloc[-1]
        else:
            delta = np.nan
    except:
        delta = np.nan
    return delta 

df['mean_score'] = df['work_id'].apply(lambda x: calculate_mean(x) )
df['diff_score'] = df['work_id'].apply(lambda x: calculate_diff(x) )

I need a more efficient way (perhaps using groupby) as this one is very slow on large dataframes.

Upvotes: 1

Views: 46

Answers (1)

Chris
Chris

Reputation: 29742

IIUC, use pandas.DataFrame.groupby with expanding.mean and diff:

g = df.groupby("s_id")["score"]
s1 = g.apply(lambda x: x.dropna().expanding().mean())
s2 = g.apply(lambda x: x.dropna().diff())

df["mean_score"] = s1.reset_index(level=0, drop=True)
df["diff_score"] = s2.reset_index(level=0, drop=True)
print(df)

Or make a function:

def mean_and_diff(series):
    s = series.dropna()
    d = {"mean_score": s.expanding().mean(), "diff_score": s.diff()}
    return pd.DataFrame(d)
tmp = df.groupby("s_id")["score"].apply(mean_and_diff).reset_index(level=0, drop=True)
df[["mean_score", "diff_score"]] = tmp[["mean_score", "diff_score"]]

Output:

  work_id s_id  score   work_date  mean_score  diff_score
0      a3  p01    NaN  2020-05-01         NaN         NaN
1      a2  p01   10.0  2020-06-10   10.000000         NaN
2      a1  p01    5.0  2020-06-15    7.500000        -5.0
3      a5  p02    5.0  2019-10-10    5.000000         NaN
4      a7  p02   11.0  2020-03-01    8.000000         6.0
5      a6  p02    NaN  2020-04-01         NaN         NaN
6      a4  p02    4.0  2020-06-20    6.666667        -7.0

Upvotes: 1

Related Questions