group-by

create new column with found value from column after group by


I have data as follow:

{'index': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 49, 50, 51, 52, 53, 54, 55, 56, 57],
 'columns': ['Subject', 'Visit', 'Date'],
 'data': [['A', 'Screening', Timestamp('2023-11-15 00:00:00')],
  ['A', 'Week 0', Timestamp('2023-11-29 00:00:00')],
  ['A', 'Week 2', Timestamp('2023-12-12 00:00:00')],
  ['A', 'Week 4', Timestamp('2023-12-27 00:00:00')],
  ['A', 'Week 8', Timestamp('2024-01-22 00:00:00')],
  ['A', 'Week 12', Timestamp('2024-02-21 00:00:00')],
  ['A', 'Week 16', Timestamp('2024-03-17 00:00:00')],
  ['A', 'Week 20', Timestamp('2024-04-17 00:00:00')],
  ['A', 'Week 28', Timestamp('2024-06-06 00:00:00')],
  ['A', 'Week 36', Timestamp('2024-08-08 00:00:00')],
  ['B', 'Screening', Timestamp('2024-02-19 00:00:00')],
  ['B', 'Week 0', Timestamp('2024-03-10 00:00:00')],
  ['B', 'Week 2', Timestamp('2024-03-24 00:00:00')],
  ['B', 'Week 4', Timestamp('2024-04-07 00:00:00')],
  ['B', 'Week 8', Timestamp('2024-05-05 00:00:00')],
  ['B', 'Week 12', Timestamp('2024-06-02 00:00:00')],
  ['B', 'Week 16', Timestamp('2024-06-27 00:00:00')],
  ['B', 'Week 20', Timestamp('2024-07-28 00:00:00')],
  ['B', 'Week 28', Timestamp('2024-09-04 00:00:00')]],
  'index_names': [None],
  'column_names': [None]}

I want to create new column in df named "ScreeningDate" which would contain Screening date for given subject.

Can you please advice?


Solution

  • First, you need to convert the JSON-like / dictionary structure to dataframe.

    json_data = {
        'index': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 49, 50, 51, 52, 53, 54, 55, 56, 57],
        'columns': ['Subject', 'Visit', 'Date'],
        'data': [
            ['A', 'Screening', '2023-11-15'],
            ['A', 'Week 0', '2023-11-29'],
            ['A', 'Week 2', '2023-12-12'],
            ['A', 'Week 4', '2023-12-27'],
            ['A', 'Week 8', '2024-01-22'],
            ['A', 'Week 12', '2024-02-21'],
            ['A', 'Week 16', '2024-03-17'],
            ['A', 'Week 20', '2024-04-17'],
            ['A', 'Week 28', '2024-06-06'],
            ['A', 'Week 36', '2024-08-08'],
            ['B', 'Screening', '2024-02-19'],
            ['B', 'Week 0', '2024-03-10'],
            ['B', 'Week 2', '2024-03-24'],
            ['B', 'Week 4', '2024-04-07'],
            ['B', 'Week 8', '2024-05-05'],
            ['B', 'Week 12', '2024-06-02'],
            ['B', 'Week 16', '2024-06-27'],
            ['B', 'Week 20', '2024-07-28'],
            ['B', 'Week 28', '2024-09-04']
        ]
    }
    

    Converting JSON-like structure to DataFrame

    df = pd.DataFrame(json_data['data'], columns=json_data['columns'])
    
    df['Date'] = pd.to_datetime(df_from_json['Date'])
    
    df
    

    Df:

     Subject      Visit       Date
    0        A  Screening 2023-11-15
    1        A     Week 0 2023-11-29
    2        A     Week 2 2023-12-12
    3        A     Week 4 2023-12-27
    4        A     Week 8 2024-01-22
    5        A    Week 12 2024-02-21
    6        A    Week 16 2024-03-17
    7        A    Week 20 2024-04-17
    8        A    Week 28 2024-06-06
    9        A    Week 36 2024-08-08
    10       B  Screening 2024-02-19
    11       B     Week 0 2024-03-10
    12       B     Week 2 2024-03-24
    13       B     Week 4 2024-04-07
    14       B     Week 8 2024-05-05
    15       B    Week 12 2024-06-02
    16       B    Week 16 2024-06-27
    17       B    Week 20 2024-07-28
    18       B    Week 28 2024-09-04
    

    Group the data by Subject column. Within the groups locate Date that corresponds to the Visit labelled as Screening. Then apply screening date to all rows using transform function, this allows returning a column of the same length as the dataframe.

    Then store the resulting screening date values for each subject in a new column named ScreeningDate.

    data = {
        'Subject': ['A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 
                    'B', 'B', 'B', 'B', 'B', 'B', 'B', 'B', 'B'],
        'Visit': ['Screening', 'Week 0', 'Week 2', 'Week 4', 'Week 8', 
                  'Week 12', 'Week 16', 'Week 20', 'Week 28', 'Week 36', 
                  'Screening', 'Week 0', 'Week 2', 'Week 4', 'Week 8', 
                  'Week 12', 'Week 16', 'Week 20', 'Week 28'],
        'Date': [
            pd.Timestamp('2023-11-15'), pd.Timestamp('2023-11-29'), pd.Timestamp('2023-12-12'), 
            pd.Timestamp('2023-12-27'), pd.Timestamp('2024-01-22'), pd.Timestamp('2024-02-21'), 
            pd.Timestamp('2024-03-17'), pd.Timestamp('2024-04-17'), pd.Timestamp('2024-06-06'), 
            pd.Timestamp('2024-08-08'), pd.Timestamp('2024-02-19'), pd.Timestamp('2024-03-10'), 
            pd.Timestamp('2024-03-24'), pd.Timestamp('2024-04-07'), pd.Timestamp('2024-05-05'), 
            pd.Timestamp('2024-06-02'), pd.Timestamp('2024-06-27'), pd.Timestamp('2024-07-28'), 
            pd.Timestamp('2024-09-04')
        ]
    }
    
    df = pd.DataFrame(data)
    
    df['ScreeningDate'] = df.groupby('Subject')['Date'].transform(lambda x: x.loc[df['Visit'] == 'Screening'].values[0])
    
    df
    

    Df after changes:

     Subject      Visit       Date ScreeningDate
    0        A  Screening 2023-11-15    2023-11-15
    1        A     Week 0 2023-11-29    2023-11-15
    2        A     Week 2 2023-12-12    2023-11-15
    3        A     Week 4 2023-12-27    2023-11-15
    4        A     Week 8 2024-01-22    2023-11-15
    5        A    Week 12 2024-02-21    2023-11-15
    6        A    Week 16 2024-03-17    2023-11-15
    7        A    Week 20 2024-04-17    2023-11-15
    8        A    Week 28 2024-06-06    2023-11-15
    9        A    Week 36 2024-08-08    2023-11-15
    10       B  Screening 2024-02-19    2024-02-19
    11       B     Week 0 2024-03-10    2024-02-19
    12       B     Week 2 2024-03-24    2024-02-19
    13       B     Week 4 2024-04-07    2024-02-19
    14       B     Week 8 2024-05-05    2024-02-19
    15       B    Week 12 2024-06-02    2024-02-19
    16       B    Week 16 2024-06-27    2024-02-19
    17       B    Week 20 2024-07-28    2024-02-19
    18       B    Week 28 2024-09-04    2024-02-19