nyc_restaurants_8.py

# Create next grade date
df["NEXT_GRADE_DATE"] = df.GRADE_DATE.shift()

# Remove next grade date for most recent grades
df.loc[df.drop_duplicates("KEY").index,"NEXT_GRADE_DATE"] = pd.NaT

df["TIME_AT_GRADE"] = df.NEXT_GRADE_DATE - df.GRADE_DATE

f, ax = plt.subplots(figsize=(10,5))

for grade, color in colors.items():
    days_at_grade = df[(~df.TIME_AT_GRADE.isna()) & (df.GRADE == grade)].TIME_AT_GRADE.dt.days
    if days_at_grade.size > 0:
        ax.hist(days_at_grade, color=color, label=grade, alpha=0.5, bins=100, weights=np.zeros_like(days_at_grade) + 1. / days_at_grade.size)
ax.legend()
f.tight_layout()