Package: catenets
CATENets - Conditional Average Treatment Effect Estimation Using Neural NetworksCode Author: Alicia Curth (amc253@cam.ac.uk)This repo contains Jax-based, sklearn-style implementations of Neural Network-based Conditional
Average Treatment Effect (CATE) Estimators, which were used in the AISTATS21 paper'Nonparametric Estimation of Heterogeneous Treatment Effects: From Theory to Learning
Algorithms'(Curth & vd Schaar, 2021a) as well as the follow up
NeurIPS21 paper"On Inductive Biases for Heterogeneous Treatment Effect Estimation"(Curth & vd
Schaar, 2021b) and the NeurIPS21 Datasets & Benchmarks track paper"Really Doing Great at Estimating CATE? A Critical Look at ML Benchmarking Practices in Treatment Effect Estimation"(Curth et al, 2021).We implement the SNet-class we introduce in Curth & vd Schaar (2021a), as well as FlexTENet and
OffsetNet as discussed in Curth & vd Schaar (2021b), and re-implement a number of
NN-based algorithms from existing literature (Shalit et al (2017), Shi et al (2019), Hassanpour
& Greiner (2020)). We also provide Neural Network (NN)-based instantiations of a number of so-called
meta-learners for CATE estimation, including two-step pseudo-outcome regression estimators (the
DR-learner (Kennedy, 2020) and single-robust propensity-weighted (PW) and regression-adjusted (RA) learners), Nie & Wager (2017)'s R-learner and Kuenzel et al (2019)'s X-learner. The jax implementations incatenets.models.jaxwere used in all papers listed; additionally, pytorch versions of some models (catenets.models.torch) were contributed byBogdan Cebere.InterfaceThe repo contains a packagecatenets, which contains all general code used for modeling and evaluation, and a folderexperiments, in which the code for replicating experimental results is contained. All implemented learning algorithms incatenets(SNet, FlexTENet, OffsetNet, TNet, SNet1 (TARNet), SNet2 (DragonNet), SNet3, DRNet, RANet, PWNet, RNet, XNet) come with a sklearn-style wrapper,  implementing a.fit(X, y, w)and a.predict(X)method, where predict returns CATE by default. All hyperparameters are documented in detail in the respective files incatenets.modelsfolder.Example usage:fromcatenets.models.jaximportTNet,SNetfromcatenets.experiment_utils.simulation_utilsimportsimulate_treatment_setup# simulate some data (here: unconfounded, 10 prognostic variables and 5 predictive variables)X,y,w,p,cate=simulate_treatment_setup(n=2000,n_o=10,n_t=5,n_c=0)# estimate CATE using TNett=TNet()t.fit(X,y,w)cate_pred_t=t.predict(X)# without potential outcomescate_pred_t,po0_pred_t,po1_pred_t=t.predict(X,return_po=True)# predict potential outcomes too# estimate CATE using SNets=SNet(penalty_orthogonal=0.01)s.fit(X,y,w)cate_pred_s=s.predict(X)All experiments in Curth & vd Schaar (2021a) can be replicated using this repository; the necessary
code is inexperiments.experiments_AISTATS21. To do so from shell, clone the repo, create a new
virtual environment and runpipinstall-rrequirements.txt#install requirementspythonrun_experiments_AISTATS.pyOptions:
--experiment# defaults to 'simulation', 'ihdp' will run ihdp experiments--setting# different simulation settings in synthetic experiments (can be 1-5)--models# defaults to None which will train all models considered in paper,# can be string of model name (e.g 'TNet'), 'plug' for all plugin models,# 'pseudo' for all pseudo-outcome regression models--file_name# base file name to write to, defaults to 'results'--n_repeats# number of experiments to run for each configuration, defaults to 10 (should be set to 100 for IHDP)Similarly, the experiments in Curth & vd Schaar (2021b) can be replicated using the code inexperiments.experiments_inductivebias_NeurIPS21(or from shell usingpython run_experiments_inductive_bias_NeurIPS.py) and the experiments in Curth et al (2021) can be replicated using the code inexperiments.experiments_benchmarks_NeurIPS21(the catenets experiments can also be run from shell usingpython run_experiments_benchmarks_NeurIPS).The code can also be installed as a python package (catenets). From a local copy of the repo, runpython setup.py install.Note: jax is currently only supported on macOS and linux, but can be run from windows using WSL (the windows subsystem for linux).CitingIf you use this software please cite the corresponding paper(s):@inproceedings{curth2021nonparametric,
  title={Nonparametric Estimation of Heterogeneous Treatment Effects: From Theory to Learning Algorithms},
  author={Curth, Alicia and van der Schaar, Mihaela},
    year={2021},
  booktitle={Proceedings of the 24th International Conference on Artificial
  Intelligence and Statistics (AISTATS)},
  organization={PMLR}
}

@article{curth2021inductive,
  title={On Inductive Biases for Heterogeneous Treatment Effect Estimation},
  author={Curth, Alicia and van der Schaar, Mihaela},
  booktitle={Proceedings of the Thirty-Fifth Conference on Neural Information Processing Systems},
  year={2021}
}


@article{curth2021really,
  title={Really Doing Great at Estimating CATE? A Critical Look at ML Benchmarking Practices in Treatment Effect Estimation},
  author={Curth, Alicia and Svensson, David and Weatherall, James and van der Schaar, Mihaela},
  booktitle={Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks},
  year={2021}
}