Skip to content

Quantco/metalearners

Repository files navigation

metalearners

CI Documentation Status Conda-forge PypiVersion codecov.io

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

Example

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["outcomes"], df["treatment"])
cate_estimates = rlearner.predict(df[features], is_oos=False)

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

Installation

metalearners can either be installed via PyPI with

$ pip install metalearners

or via conda-forge with

$ conda install metalearners -c conda-forge

Development

Development instructions can be found here.