import numpy as np import pandas as pd from lifelines import CoxPHFitter DEFAULT_YEARS = [3, 5, 10] def build_os_time_days( df: pd.DataFrame, months_col: str = "Months_from_transplant_to_last_followup", death_date_col: str = "Date_of_death", hsct_date_col: str = "HSCT_date", event_col: str = "Event", ) -> pd.Series: """ Build OS_time_days with priority: 1) Months_from_transplant_to_last_followup (converted to days) 2) If missing, and Event == 1, use Date_of_death - HSCT_date Returns ------- pd.Series Survival duration in days. """ out = pd.Series(np.nan, index=df.index, dtype="float64") # First preference: follow-up months already present if months_col in df.columns: months = pd.to_numeric(df[months_col], errors="coerce") valid_months = months.notna() & (months > 0) out.loc[valid_months] = months.loc[valid_months] * 30.44 # Fallback: calculate from death date - HSCT date for death events required_cols = [death_date_col, hsct_date_col, event_col] if all(c in df.columns for c in required_cols): death_dt = pd.to_datetime(df[death_date_col], errors="coerce") hsct_dt = pd.to_datetime(df[hsct_date_col], errors="coerce") event = pd.to_numeric(df[event_col], errors="coerce") mask = out.isna() & death_dt.notna() & hsct_dt.notna() & (event == 1) out.loc[mask] = (death_dt.loc[mask] - hsct_dt.loc[mask]).dt.days.astype(float) # Clean invalid durations out.loc[out <= 0] = np.nan return out def prepare_cox_df( df_surv: pd.DataFrame, preds_surv: np.ndarray, covariates: list[str], duration_col: str = "OS_time_days", event_col: str = "Event_clean", risk_col_name: str = "Predicted_GVHD_Risk", ) -> tuple[pd.DataFrame, list[str], list[str]]: """ Prepare Cox dataframe for OS modeling. Parameters ---------- df_surv : pd.DataFrame Survival-ready dataframe, usually preprocessed dataframe. preds_surv : np.ndarray Predicted risk scores aligned with df_surv rows. covariates : list[str] Covariates to include in Cox model. duration_col : str Name of duration column to use/create. event_col : str Name of cleaned event column to use/create. risk_col_name : str Explicit risk predictor name, e.g.: - Predicted_acute_gvhd_risk - Predicted_chronic_gvhd_risk Returns ------- df_cox : pd.DataFrame Numeric-only Cox dataset including duration and event columns. design_cols : list[str] Final model covariates after one-hot encoding. cat_cols : list[str] Original categorical covariates that were one-hot encoded. """ df = df_surv.copy() # Build OS_time_days if not already present or if missing if duration_col not in df.columns: df[duration_col] = build_os_time_days(df) else: existing_duration = pd.to_numeric(df[duration_col], errors="coerce") rebuilt_duration = build_os_time_days(df) df[duration_col] = existing_duration fill_mask = df[duration_col].isna() df.loc[fill_mask, duration_col] = rebuilt_duration.loc[fill_mask] # Build Event_clean if not already present if event_col not in df.columns: if "Event" not in df.columns: raise ValueError( f"Neither '{event_col}' nor raw 'Event' column found in dataframe." ) df[event_col] = pd.to_numeric(df["Event"], errors="coerce") else: df[event_col] = pd.to_numeric(df[event_col], errors="coerce") # Add predicted GVHD risk as Cox covariate preds_surv = np.asarray(preds_surv).astype(float).ravel() if len(preds_surv) != len(df): raise ValueError( f"Length mismatch: preds_surv has {len(preds_surv)} rows but df_surv has {len(df)} rows." ) df[risk_col_name] = preds_surv # Keep only required columns keep_covariates = [c for c in covariates if c in df.columns] cols = [duration_col, event_col, risk_col_name] + keep_covariates df = df[cols].copy() # Coerce survival columns df[duration_col] = pd.to_numeric(df[duration_col], errors="coerce") df[event_col] = pd.to_numeric(df[event_col], errors="coerce") # Drop invalid rows df = df.dropna(subset=[duration_col, event_col]) df = df[df[duration_col] > 0] df[event_col] = df[event_col].astype(int).clip(0, 1) # Identify categorical covariates from pandas.api.types import is_object_dtype, is_string_dtype, is_categorical_dtype cat_cols = [ c for c in df.columns if c not in [duration_col, event_col] and (is_object_dtype(df[c]) or is_string_dtype(df[c]) or is_categorical_dtype(df[c])) ] # One-hot encode categoricals if cat_cols: df = pd.get_dummies(df, columns=cat_cols, drop_first=True) # Ensure numeric design matrix for c in df.columns: if c in [duration_col, event_col]: continue df[c] = pd.to_numeric(df[c], errors="coerce").fillna(0) design_cols = [c for c in df.columns if c not in [duration_col, event_col]] return df, design_cols, cat_cols def fit_cox( df_cox: pd.DataFrame, duration_col: str = "OS_time_days", event_col: str = "Event_clean", ) -> CoxPHFitter: cph = CoxPHFitter(penalizer=0.1) cph.fit(df_cox, duration_col=duration_col, event_col=event_col) return cph def make_patient_design_row( patient_row: pd.DataFrame, design_cols: list[str], cat_cols_original: list[str], ): x = patient_row.copy() if cat_cols_original: available_cat_cols = [c for c in cat_cols_original if c in x.columns] if available_cat_cols: x = pd.get_dummies(x, columns=available_cat_cols, drop_first=True) unseen = sorted(set(x.columns) - set(design_cols)) x = x.reindex(columns=design_cols, fill_value=0) for c in x.columns: x[c] = pd.to_numeric(x[c], errors="coerce").fillna(0) return x, unseen def predict_patient_survival(cph, patient_design, years=None): if years is None: years = DEFAULT_YEARS surv_fn = cph.predict_survival_function(patient_design) # survival function is a DataFrame: index=time, column=patient s_series = surv_fn.iloc[:, 0] times = s_series.index.values landmarks = {} for y in years: t = 365 * int(y) le_mask = times <= t if le_mask.any(): s = float(s_series.loc[times[le_mask].max()]) else: s = float(s_series.iloc[0]) landmarks[int(y)] = s return surv_fn, landmarks