Skip to content

Quantco/metalearners

This branch is up to date with main.

Folders and files

NameName
Last commit message
Last commit date

Latest commit

a4fe514 · Feb 15, 2025
Aug 27, 2024
Aug 15, 2024
Jul 23, 2024
Dec 3, 2024
Feb 15, 2025
Feb 15, 2025
Feb 5, 2024
Feb 5, 2024
Dec 25, 2024
Jun 6, 2024
Feb 5, 2024
Jun 25, 2024
Jan 29, 2025
Jun 18, 2024
Oct 6, 2024
Aug 27, 2024
Feb 15, 2025
Feb 15, 2025
Jul 4, 2024

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.