Fix censor_flg encoding in mimic_2() (#227)

* fixed mimic_2 dataset loader so that censor_flg is encoded correctly

* handle the case if censor_flg passed in the columns_obs_only argument

---------

Co-authored-by: Eljas Roellin <65244425+eroell@users.noreply.github.com>

Author: sueoglu

src/ehrdata/dt/datasets.py

@@ -917,6 +917,15 @@ def mimic_2(
         columns_obs_only=columns_obs_only,
     )
 
+    # In the raw dataset, the variable censor_flg is encoded inversely (0=death, 1=censored)
+    # We flip it here so it follows the standard convention (0=censored, 1=event happened)
+
+    censor_col = "censor_flg"
+    if censor_col in edata.var.index:
+        edata[:, [censor_col]].X = np.where(edata[:, [censor_col]].X == 0, 1, 0)
+    elif censor_col in edata.obs.columns:
+        edata.obs[censor_col] = np.where(edata.obs[censor_col] == 0, 1, 0)
+
     return edata