groupby中的条件前填充

Question

I have a data frame for patients and their visits to the clinic. Patients may take a drug at some visits, and only the initial dose is recorded, or when the dose is changed. If the dose doesn't change at the next visit, what's recorded is "drug ongoing? Yes. Dose changed? No". What I need to get is the exact dose for each visit.

I tried forward fill with groupby (groupby patient_id ), but I'm stuck at how to insert the condition that only fill missing when drug is ongoing and dose is not changed.

df = pd.DataFrame({'patient_id': ['a', 'a', 'a', 'b', 'b', 'b', 'c', 'c', 'c'], \
              'visit_number':[1, 2, 3, 2, 3, 4, 10, 11, 12], \
             'drug_ongoing':[np.nan, 1, 1, np.nan, 0, 1, 1, 1, 0], \
             'drug_dose_changed':[0, 0, 0, 0, np.nan,0, 0, 1, np.nan], \
             'dose':[40, np.nan, np.nan, 60, np.nan, 70, 80, np.nan, np.nan]})

I tried:

df['dose_filled'] = df.groupby('patient_id')['dose'].ffill()

But in this way, all the missing is filled.

The desired new column 'dose_filled' is [40, 40, 40, 60, np.nan, 70, 80, np.nan, np.nan]

Answer 1

In your case , filter before ffill

s=df.loc[(df['drug_ongoing'].eq(1)&df['drug_dose_changed'].eq(0))|df.visit_number.eq(df.groupby('patient_id').visit_number.transform('first'))].groupby('patient_id').dose.ffill()
df.dose.fillna(s,inplace=True)
df
Out[38]: 
  patient_id  visit_number  drug_ongoing  drug_dose_changed  dose
0          a             1           NaN                0.0  40.0
1          a             2           1.0                0.0  40.0
2          a             3           1.0                0.0  40.0
3          b             2           NaN                0.0  60.0
4          b             3           0.0                NaN   NaN
5          b             4           1.0                0.0  70.0
6          c            10           1.0                0.0  80.0
7          c            11           1.0                1.0   NaN
8          c            12           0.0                NaN   NaN

Answer 2

I think you need:

np.where(~df.drug_dose_changed.astype(bool),df.dose.ffill(),df.dose)

Output:

array([40., 40., 40., 60., nan, 70., 80., nan, nan])