MetaLearners for Conditional Average Treatment Effect (CATE) estimation

The library focuses on providing

• Methodologically sound cross-fitting
• Convenient access to and reuse of base models
• Consistent APIs across Metalearners
• Support for more than binary treatment variants
• Integrations with pandas, shap, lime, optuna and soon onnx

df = ...

from metalearners import RLearner
from lightgbm import LGBMClassifier, LGBMRegressor

rlearner = RLearner(
    nuisance_model_factory=LGBMRegressor,
    propensity_model_factory=LGBMClassifier,
    treatment_model_factory=LGBMRegressor,
    is_classification=False,
    n_variants=2,
)

features = ["age", "weight", "height"]
rlearner.fit(df[features], df["treatment"], df["outcomes"])
cate_estimates = rlearner.predict(df[features], is_oos=False)

Please refer to our docs for many more in-depth and reproducible examples.

metalearners can either be installed via PyPI with

$ pip install metalearners

or via conda-forge with

$ conda install metalearners -c conda-forge

Development instructions can be found here.