BiomedSciAI / causallib

A Python package for modular causal inference analysis and model evaluations
Apache License 2.0
728 stars 97 forks source link

Model selection within weight-based survival models #47

Closed ehudkr closed 1 year ago

ehudkr commented 1 year ago

Allows to incorporate model-selection model (e.g. GridSearchCV) within weight-based survival models. Namely, it is now possible to:

from causallib.datasets import load_nhefs_survival
from causallib.estimation import IPW
from causallib.model_selection import GridSearchCV
from causallib.survival import WeightedSurvival
from sklearn.linear_model import LogisticRegression

data = load_nhefs_survival()

ipw = IPW(LogisticRegression(penalty="none", solver="saga", max_iter=10000))
ipw_grid_model = GridSearchCV(
    ipw,
    param_grid=dict(clip_min=[0.2, 0.3]),
    scoring="weighted_roc_auc_error",
)
survival_model = WeightedSurvival(
    weight_model=ipw_grid_model,
)
survival_model.fit(data.X, data.a, data.t, data.y)
outcomes = survival_model.estimate_population_outcome(
    data.X, data.a, data.t, data.y, 
)