diff --git a/features/esm_SAM.py b/features/esm_SAM.py index 0c5ba21..f59282c 100644 --- a/features/esm_SAM.py +++ b/features/esm_SAM.py @@ -36,8 +36,17 @@ def extract_stressful_events(df_esm: pd.DataFrame) -> pd.DataFrame: df_esm_event_threat_challenge_mean_wide = calculate_threat_challenge_means( df_esm_sam_clean ) + # 2. + df_esm_event_stress = detect_stressful_event(df_esm_sam_clean) - return df_esm_event_threat_challenge_mean_wide + # Join all the features related to the events. + df_esm_events = df_esm_event_threat_challenge_mean_wide.join( + df_esm_event_stress[ + GROUP_QUESTIONNAIRES_BY + ["event_present", "event_stressfulness"] + ].set_index(GROUP_QUESTIONNAIRES_BY) + ) + + return df_esm_events def calculate_threat_challenge_means(df_esm_sam_clean: pd.DataFrame) -> pd.DataFrame: @@ -78,3 +87,14 @@ def calculate_threat_challenge_means(df_esm_sam_clean: pd.DataFrame) -> pd.DataF values=["esm_numeric_mean"], ) return df_esm_event_threat_challenge_mean_wide + + +def detect_stressful_event(df_esm_sam_clean: pd.DataFrame) -> pd.DataFrame: + df_esm_event_stress = df_esm_sam_clean[ + df_esm_sam_clean["questionnaire_id"] == QUESTIONNAIRE_ID_SAM.get("event_stress") + ] + df_esm_event_stress = df_esm_event_stress.assign( + event_present=lambda x: x.esm_user_answer_numeric > 0, + event_stressfulness=lambda x: x.esm_user_answer_numeric, + ) + return df_esm_event_stress