Upload 199 files

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+causalml/source/docs/_static/img/synthetic_dgp_scatter_plot.png filter=lfs diff=lfs merge=lfs -text
+causalml/source/docs/_static/img/uplift_tree_vis.png filter=lfs diff=lfs merge=lfs -text
+causalml/source/docs/examples/causal_trees_with_synthetic_data_multiple_treatment_groups.ipynb filter=lfs diff=lfs merge=lfs -text
+causalml/source/docs/examples/causal_trees_with_synthetic_data.ipynb filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,18 @@
+FROM python:3.10
+
+RUN useradd -m -u 1000 user && python -m pip install --upgrade pip
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+
+WORKDIR /app
+
+COPY --chown=user ./requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+COPY --chown=user . /app
+ENV MCP_TRANSPORT=http
+ENV MCP_PORT=7860
+
+EXPOSE 7860
+
+CMD ["python", "causalml/mcp_output/start_mcp.py"]

README.md

@@ -1,10 +1,32 @@
 ---
-title: Causalml
-emoji: 📊
-colorFrom: pink
-colorTo: red
+title: Causalml MCP
+emoji: 🤖
+colorFrom: blue
+colorTo: purple
 sdk: docker
+sdk_version: "4.26.0"
+app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Causalml MCP Service
+
+Auto-generated MCP service for causalml.
+
+## Usage
+
+```
+https://None-causalml-mcp.hf.space/mcp
+```
+
+## Connect with Cursor
+
+```json
+{
+  "mcpServers": {
+    "causalml": {
+      "url": "https://None-causalml-mcp.hf.space/mcp"
+    }
+  }
+}
+```

app.py

@@ -0,0 +1,45 @@
+from fastapi import FastAPI
+import os
+import sys
+
+mcp_plugin_path = os.path.join(os.path.dirname(__file__), "causalml", "mcp_output", "mcp_plugin")
+sys.path.insert(0, mcp_plugin_path)
+
+app = FastAPI(
+    title="Causalml MCP Service",
+    description="Auto-generated MCP service for causalml",
+    version="1.0.0"
+)
+
+@app.get("/")
+def root():
+    return {
+        "service": "Causalml MCP Service",
+        "version": "1.0.0",
+        "status": "running",
+        "transport": os.environ.get("MCP_TRANSPORT", "http")
+    }
+
+@app.get("/health")
+def health_check():
+    return {"status": "healthy", "service": "causalml MCP"}
+
+@app.get("/tools")
+def list_tools():
+    try:
+        from mcp_service import create_app
+        mcp_app = create_app()
+        tools = []
+        for tool_name, tool_func in mcp_app.tools.items():
+            tools.append({
+                "name": tool_name,
+                "description": tool_func.__doc__ or "No description available"
+            })
+        return {"tools": tools}
+    except Exception as e:
+        return {"error": f"Failed to load tools: {str(e)}"}
+
+if __name__ == "__main__":
+    import uvicorn
+    port = int(os.environ.get("PORT", 7860))
+    uvicorn.run(app, host="0.0.0.0", port=port)

causalml/mcp_output/README_MCP.md

@@ -0,0 +1,73 @@
+# CausalML: Model Context Protocol (MCP) Service
+
+## Project Introduction
+
+CausalML is a comprehensive Python package designed to provide a suite of uplift modeling and causal inference methods using machine learning algorithms. It is based on recent research and offers a standard interface for estimating the Conditional Average Treatment Effect (CATE) and Individual Treatment Effect (ITE) from experimental or observational data. The package is particularly valuable for real-world applications such as campaign targeting optimization and personalized engagement by estimating the causal impact of interventions on outcomes.
+
+## Installation Method
+
+To install CausalML, ensure you have the following dependencies:
+
+- scikit-learn>=1.6.0
+- xgboost
+- tensorflow>=2.4.0
+- torch
+- scipy>=1.4.1
+- pandas>=0.24.1
+- setuptools
+- Cython
+
+Optional dependencies include:
+
+- pyro-ppl
+- cibuildwheel
+- pytest
+- pytest-cov
+
+You can install CausalML using pip:
+
+```
+pip install causalml
+```
+
+## Quick Start
+
+To quickly get started with CausalML, you can use the following example to estimate treatment effects:
+
+1. Import the necessary modules.
+2. Load your dataset.
+3. Choose a meta-learner or inference method.
+4. Fit the model and predict treatment effects.
+
+Example:
+
+```
+from causalml.inference.meta import BaseTLearner
+from causalml.dataset import make_uplift_classification
+
+X, treatment, y = make_uplift_classification()
+learner = BaseTLearner()
+learner.fit(X, treatment, y)
+ate = learner.estimate_ate(X, treatment, y)
+```
+
+## Available Tools and Endpoints List
+
+- **Meta-Learners**: Includes BaseSLearner, BaseTLearner, BaseXLearner, BaseRLearner, BaseDRLearner, and TMLELearner for various strategies in estimating treatment effects.
+- **Tree-Based Methods**: UpliftTreeClassifier, UpliftRandomForestClassifier, CausalTreeRegressor, and CausalRandomForestRegressor for uplift modeling and causal inference.
+- **Neural Network Methods**: DragonNet and CEVAE for causal inference using TensorFlow and PyTorch.
+- **Instrumental Variable Methods**: DRIVLearner for causal inference.
+- **Metrics and Visualization**: Functions like AUUC, Qini, and plot_lift for evaluating causal inference models.
+- **Optimization Methods**: CounterfactualUnitSelector and CounterfactualValueEstimator for treatment effect estimation and counterfactual analysis.
+
+## Common Issues and Notes
+
+- Ensure all dependencies are correctly installed to avoid import errors.
+- Performance may vary based on the dataset size and complexity of the model chosen.
+- For optimal performance, consider using Cython extensions and leveraging GPU support with TensorFlow or PyTorch.
+
+## Reference Links or Documentation
+
+For more detailed documentation and examples, visit the [CausalML GitHub repository](https://github.com/uber/causalml).
+
+For additional information on methodology and usage, refer to the documentation files within the repository, such as `docs/methodology.rst` and `README.md`.

causalml/mcp_output/analysis.json

@@ -0,0 +1,459 @@
+{
+  "summary": {
+    "repository_url": "https://github.com/uber/causalml",
+    "summary": "Imported via zip fallback, file count: 102",
+    "file_tree": {
+      ".github/.stale.yml": {
+        "size": 683
+      },
+      ".github/ISSUE_TEMPLATE/bug_report.md": {
+        "size": 730
+      },
+      ".github/ISSUE_TEMPLATE/feature_request.md": {
+        "size": 604
+      },
+      ".github/PULL_REQUEST_TEMPLATE.md": {
+        "size": 1616
+      },
+      ".github/workflows/black.yml": {
+        "size": 154
+      },
+      ".github/workflows/python-publish.yml": {
+        "size": 1998
+      },
+      ".github/workflows/python-test.yaml": {
+        "size": 963
+      },
+      ".github/workflows/test-build-from-source.yml": {
+        "size": 2193
+      },
+      ".github/workflows/test-pypi-install.yml": {
+        "size": 1209
+      },
+      ".pre-commit-config.yaml": {
+        "size": 261
+      },
+      ".readthedocs.yml": {
+        "size": 523
+      },
+      "ANTITRUST.md": {
+        "size": 1206
+      },
+      "CHARTER.md": {
+        "size": 4155
+      },
+      "CODE_OF_CONDUCT.md": {
+        "size": 3224
+      },
+      "CONTRIBUTING.md": {
+        "size": 5638
+      },
+      "GOVERNANCE.md": {
+        "size": 3908
+      },
+      "MAINTAINERS.md": {
+        "size": 1310
+      },
+      "README.md": {
+        "size": 9061
+      },
+      "SECURITY.md": {
+        "size": 227
+      },
+      "STEERING_COMMITTEE.md": {
+        "size": 1098
+      },
+      "TRADEMARKS.md": {
+        "size": 4790
+      },
+      "causalml/__init__.py": {
+        "size": 149
+      },
+      "causalml/dataset/__init__.py": {
+        "size": 872
+      },
+      "causalml/dataset/classification.py": {
+        "size": 28659
+      },
+      "causalml/dataset/regression.py": {
+        "size": 8860
+      },
+      "causalml/dataset/semiSynthetic.py": {
+        "size": 36472
+      },
+      "causalml/dataset/synthetic.py": {
+        "size": 24403
+      },
+      "causalml/feature_selection/__init__.py": {
+        "size": 34
+      },
+      "causalml/feature_selection/filters.py": {
+        "size": 27676
+      },
+      "causalml/features.py": {
+        "size": 8366
+      },
+      "causalml/inference/__init__.py": {
+        "size": 0
+      },
+      "causalml/inference/iv/__init__.py": {
+        "size": 117
+      },
+      "causalml/inference/iv/drivlearner.py": {
+        "size": 37605
+      },
+      "causalml/inference/iv/iv_regression.py": {
+        "size": 1420
+      },
+      "causalml/inference/meta/__init__.py": {
+        "size": 474
+      },
+      "causalml/inference/meta/base.py": {
+        "size": 13510
+      },
+      "causalml/inference/meta/drlearner.py": {
+        "size": 24255
+      },
+      "causalml/inference/meta/explainer.py": {
+        "size": 11472
+      },
+      "causalml/inference/meta/rlearner.py": {
+        "size": 29178
+      },
+      "causalml/inference/meta/slearner.py": {
+        "size": 15873
+      },
+      "causalml/inference/meta/tlearner.py": {
+        "size": 15711
+      },
+      "causalml/inference/meta/tmle.py": {
+        "size": 8014
+      },
+      "causalml/inference/meta/utils.py": {
+        "size": 4276
+      },
+      "causalml/inference/meta/xlearner.py": {
+        "size": 26635
+      },
+      "causalml/inference/tf/__init__.py": {
+        "size": 33
+      },
+      "causalml/inference/tf/dragonnet.py": {
+        "size": 10594
+      },
+      "causalml/inference/tf/utils.py": {
+        "size": 6098
+      },
+      "causalml/inference/torch/__init__.py": {
+        "size": 25
+      },
+      "causalml/inference/torch/cevae.py": {
+        "size": 5470
+      },
+      "causalml/inference/tree/__init__.py": {
+        "size": 423
+      },
+      "causalml/inference/tree/_tree/__init__.py": {
+        "size": 302
+      },
+      "causalml/inference/tree/_tree/_classes.py": {
+        "size": 24699
+      },
+      "causalml/inference/tree/causal/__init__.py": {
+        "size": 0
+      },
+      "causalml/inference/tree/causal/_tree.py": {
+        "size": 10355
+      },
+      "causalml/inference/tree/causal/causalforest.py": {
+        "size": 20212
+      },
+      "causalml/inference/tree/causal/causaltree.py": {
+        "size": 17840
+      },
+      "causalml/inference/tree/plot.py": {
+        "size": 24071
+      },
+      "causalml/inference/tree/utils.py": {
+        "size": 11016
+      },
+      "causalml/match.py": {
+        "size": 20210
+      },
+      "causalml/metrics/__init__.py": {
+        "size": 743
+      },
+      "causalml/metrics/classification.py": {
+        "size": 975
+      },
+      "causalml/metrics/const.py": {
+        "size": 12
+      },
+      "causalml/metrics/regression.py": {
+        "size": 3193
+      },
+      "causalml/metrics/sensitivity.py": {
+        "size": 22273
+      },
+      "causalml/metrics/visualize.py": {
+        "size": 36312
+      },
+      "causalml/optimize/__init__.py": {
+        "size": 263
+      },
+      "causalml/optimize/pns.py": {
+        "size": 2765
+      },
+      "causalml/optimize/policylearner.py": {
+        "size": 5973
+      },
+      "causalml/optimize/unit_selection.py": {
+        "size": 9417
+      },
+      "causalml/optimize/utils.py": {
+        "size": 4210
+      },
+      "causalml/optimize/value_optimization.py": {
+        "size": 4094
+      },
+      "causalml/propensity.py": {
+        "size": 7229
+      },
+      "docs/conf.py": {
+        "size": 9212
+      },
+      "docs/environment-py311-rtd.yml": {
+        "size": 600
+      },
+      "docs/environment-py39-rtd.yml": {
+        "size": 605
+      },
+      "docs/issue-859-resolution.md": {
+        "size": 1016
+      },
+      "docs/plans/2026-01-30-scipy-1.16-support.md": {
+        "size": 11647
+      },
+      "docs/requirements.txt": {
+        "size": 114
+      },
+      "pyproject.toml": {
+        "size": 1505
+      },
+      "setup.cfg": {
+        "size": 660
+      },
+      "setup.py": {
+        "size": 1759
+      },
+      "tests/__init__.py": {
+        "size": 0
+      },
+      "tests/conftest.py": {
+        "size": 2622
+      },
+      "tests/const.py": {
+        "size": 392
+      },
+      "tests/test_causal_trees.py": {
+        "size": 9991
+      },
+      "tests/test_cevae.py": {
+        "size": 1490
+      },
+      "tests/test_counterfactual_unit_selection.py": {
+        "size": 2475
+      },
+      "tests/test_datasets.py": {
+        "size": 2379
+      },
+      "tests/test_dragonnet.py": {
+        "size": 661
+      },
+      "tests/test_feature_selection.py": {
+        "size": 2034
+      },
+      "tests/test_features.py": {
+        "size": 1632
+      },
+      "tests/test_ivlearner.py": {
+        "size": 2168
+      },
+      "tests/test_match.py": {
+        "size": 3026
+      },
+      "tests/test_meta_learners.py": {
+        "size": 34264
+      },
+      "tests/test_metrics.py": {
+        "size": 1074
+      },
+      "tests/test_propensity.py": {
+        "size": 1629
+      },
+      "tests/test_sensitivity.py": {
+        "size": 6687
+      },
+      "tests/test_uplift_trees.py": {
+        "size": 11315
+      },
+      "tests/test_utils.py": {
+        "size": 692
+      },
+      "tests/test_value_optimization.py": {
+        "size": 2735
+      },
+      "tests/test_visualize.py": {
+        "size": 1768
+      },
+      "tox.ini": {
+        "size": 307
+      }
+    },
+    "processed_by": "zip_fallback",
+    "success": true
+  },
+  "structure": {
+    "packages": [
+      "source.causalml",
+      "source.causalml.dataset",
+      "source.causalml.feature_selection",
+      "source.causalml.inference",
+      "source.causalml.metrics",
+      "source.causalml.optimize",
+      "source.tests"
+    ]
+  },
+  "dependencies": {
+    "has_environment_yml": false,
+    "has_requirements_txt": false,
+    "pyproject": true,
+    "setup_cfg": true,
+    "setup_py": true
+  },
+  "entry_points": {
+    "imports": [],
+    "cli": [],
+    "modules": []
+  },
+  "llm_analysis": {
+    "core_modules": [
+      {
+        "package": "source.causalml.inference.meta",
+        "module": "meta",
+        "functions": [
+          "fit",
+          "predict",
+          "estimate_ate"
+        ],
+        "classes": [
+          "BaseSLearner",
+          "BaseTLearner",
+          "BaseXLearner",
+          "BaseRLearner",
+          "BaseDRLearner",
+          "TMLELearner"
+        ],
+        "description": "Meta-learners for causal inference, providing various strategies for estimating treatment effects."
+      },
+      {
+        "package": "source.causalml.inference.tree",
+        "module": "tree",
+        "functions": [],
+        "classes": [
+          "UpliftTreeClassifier",
+          "UpliftRandomForestClassifier",
+          "CausalTreeRegressor",
+          "CausalRandomForestRegressor"
+        ],
+        "description": "Tree-based methods for uplift modeling and causal inference, implemented with Cython for performance."
+      },
+      {
+        "package": "source.causalml.inference.nn",
+        "module": "nn",
+        "functions": [],
+        "classes": [
+          "DragonNet",
+          "CEVAE"
+        ],
+        "description": "Neural network methods for causal inference, leveraging TensorFlow and PyTorch."
+      },
+      {
+        "package": "source.causalml.inference.iv",
+        "module": "iv",
+        "functions": [],
+        "classes": [
+          "DRIVLearner"
+        ],
+        "description": "Instrumental variable methods for causal inference."
+      },
+      {
+        "package": "source.causalml.metrics",
+        "module": "metrics",
+        "functions": [
+          "AUUC",
+          "Qini",
+          "plot_lift"
+        ],
+        "classes": [],
+        "description": "Metrics and visualization tools for evaluating causal inference models."
+      },
+      {
+        "package": "source.causalml.optimize",
+        "module": "optimize",
+        "functions": [],
+        "classes": [
+          "CounterfactualUnitSelector",
+          "CounterfactualValueEstimator"
+        ],
+        "description": "Optimization methods for treatment effect estimation and counterfactual analysis."
+      }
+    ],
+    "cli_commands": [],
+    "import_strategy": {
+      "primary": "import",
+      "fallback": "blackbox",
+      "confidence": 0.9
+    },
+    "dependencies": {
+      "required": [
+        "scikit-learn>=1.6.0",
+        "xgboost",
+        "tensorflow>=2.4.0",
+        "torch",
+        "scipy>=1.4.1",
+        "pandas>=0.24.1",
+        "setuptools",
+        "Cython"
+      ],
+      "optional": [
+        "pyro-ppl",
+        "cibuildwheel",
+        "pytest",
+        "pytest-cov"
+      ]
+    },
+    "risk_assessment": {
+      "import_feasibility": 0.8,
+      "intrusiveness_risk": "medium",
+      "complexity": "complex"
+    }
+  },
+  "deepwiki_analysis": {
+    "repo_url": "https://github.com/uber/causalml",
+    "repo_name": "causalml",
+    "content": "uber/causalml\nInstallation and Setup\nPackage Structure\nDevelopment and Contributing\nCore Concepts and Methodology\nCausal Inference and Uplift Modeling\nTreatment Effect Estimation\nPropensity Scores\nInference Methods\nMeta-Learners\nS-Learner and T-Learner\nDR-Learner and DRIV-Learner\nTMLE Learner\nTree-Based Methods\nUplift Trees\nCausal Trees\nNeural Network Methods\nInstrumental Variables\nEvaluation and Interpretation\nMetrics and Visualization\nFeature Importance and Explainability\nSensitivity Analysis\nData Handling and Optimization\nSynthetic Data Generation\nFeature Selection\nMatching Methods\nTreatment Optimization\nExamples and Use Cases\nBasic Usage Examples\nAdvanced Applications\nDocumentation and Configuration\ncausalml/inference/meta/__init__.py\ndocs/about.rst\ndocs/methodology.rst\ndocs/refs.bib\npyproject.toml\nCausalML is a comprehensive Python package that provides a suite of uplift modeling and causal inference methods using machine learning algorithms based on recent research. It offers a standard interface for estimating the Conditional Average Treatment Effect (CATE) and Individual Treatment Effect (ITE) from experimental or observational data. The package estimates the causal impact of interventionTon outcomeYfor users with observed featuresX, without requiring strong assumptions on the model form.\nCausalML is particularly valuable for real-world applications including:\nCampaign targeting optimization: Identifying customers who will have favorable responses to advertising campaigns by estimating KPI effects from ad exposure at the individual level\nPersonalized engagement: Optimizing customer interactions across multiple treatment options (product choices, messaging channels) using heterogeneous treatment effect estimation\nThe library integrates seamlessly with the Python scientific computing ecosystem, building on established frameworks like scikit-learn, XGBoost, TensorFlow, and PyTorch while providing specialized causal inference capabilities not available in general-purpose ML libraries.\nSources:README.md18-28pyproject.toml1-5docs/about.rst4-7\nPackage Architecture\nCausalML is organized into a modular architecture with clear separation of concerns across different aspects of causal inference and uplift modeling.\nSystem Architecture Overview\nBuild SystemExternal DependenciesCausalML Library v0.15.5InfrastructureData & EvaluationCore Inference Enginecausalml.inference.metaS, T, X, R, DR, TMLE Learnerscausalml.inference.treeUplift Trees, Causal Treescausalml.inference.nnDragonNet, CEVAEcausalml.inference.iv2SLS, DRIV Learnercausalml.datasetmake_uplift_classificationcausalml.metricsAUUC, Qini, plot_liftcausalml.feature_selectionFilter methods, LR testcausalml.matchNearestNeighborMatchcausalml.propensityElasticNetPropensityModelcausalml.optimizeCounterfactualUnitSelectorscikit-learn>=1.6.0xgboosttensorflow>=2.4.0torch + pyro-pplscipy>=1.4.1pandas>=0.24.1setuptools + Cythoncibuildwheelpytest + pytest-cov\nBuild System\nExternal Dependencies\nCausalML Library v0.15.5\nInfrastructure\nData & Evaluation\nCore Inference Engine\ncausalml.inference.metaS, T, X, R, DR, TMLE Learners\ncausalml.inference.treeUplift Trees, Causal Trees\ncausalml.inference.nnDragonNet, CEVAE\ncausalml.inference.iv2SLS, DRIV Learner\ncausalml.datasetmake_uplift_classification\ncausalml.metricsAUUC, Qini, plot_lift\ncausalml.feature_selectionFilter methods, LR test\ncausalml.matchNearestNeighborMatch\ncausalml.propensityElasticNetPropensityModel\ncausalml.optimizeCounterfactualUnitSelector\nscikit-learn>=1.6.0\ntensorflow>=2.4.0\ntorch + pyro-ppl\nscipy>=1.4.1\npandas>=0.24.1\nsetuptools + Cython\ncibuildwheel\npytest + pytest-cov\nModule Dependency Analysis: The inference module forms the core engine with heavy integration to the Python ML ecosystem. Tree-based methods utilize Cython extensions for performance, while neural network methods have optional TensorFlow/PyTorch dependencies. The evaluation and infrastructure modules provide supporting functionality for complete causal inference workflows.\nSources:pyproject.toml24-76docs/methodology.rst10-36\nCore Module Structure\ncausalml.inference.meta\nBaseSLearner\nBaseTLearner\nBaseXLearner\nBaseRLearner\nBaseDRLearner\nTMLELearner\ncausalml.inference.tree\nUpliftTreeClassifier\nUpliftRandomForestClassifier\nCausalTreeRegressor\ncausalml.inference.nn\ncausalml.inference.iv\ncausalml.metrics\ncausalml.optimize\nCounterfactualUnitSelector\nCounterfactualValueEstimator\ncausalml.dataset\nmake_uplift_classification\nmake_uplift_regression\ncausalml.feature_selection\nFilterSelect\nLRSelectorRegressor\ncausalml.propensity\nElasticNetPropensityModel\nGradientBoostedPropensityModel\ncausalml.match\nNearestNeighborMatch\nMatchOptimizer\nSources:causalml/inference/meta/__init__.py1-13docs/methodology.rst10-36\nCore Modules and Their Purpose\nfeature_selection\nSources:causalml/__init__.py1-10docs/methodology.rst10-36\nInference Methods\nTheinferencemodule contains the core estimation algorithms for causal effects, organized into several submodules based on methodology.\nInference Methods Architecture\nThecausalml.inferencemodule contains the core estimation algorithms, organized into specialized submodules based on methodological approach.\ncausalml.inference\nMeta-Learner Class Hierarchy\nBaseSLearner\"Single model with treatment indicator\"+model: any ML model+fit(X, treatment, y)+predict(X, treatment)+estimate_ate(X, treatment, y)BaseTLearner\"Separate models per treatment\"+model_c: control model+model_t: treatment model+fit(X, treatment, y)+predict(X, treatment)+estimate_ate(X, treatment, y)BaseXLearner\"Four-stage cross-learning\"+model_c: control outcome model+model_t: treatment outcome model+model_tau_c: control effect model+model_tau_t: treatment effect model+propensity_model: ElasticNetPropensityModel+fit(X, treatment, y)+predict(X, treatment)BaseRLearner\"Residual-based approach\"+model_mu: outcome model+model_tau: treatment effect model+model_p: propensity model+cv: cross-validation folds+fit(X, treatment, y)+predict(X, treatment)BaseDRLearner\"Doubly robust estimation\"+model_mu_c: control outcome model+model_mu_t: treatment outcome model+model_tau: effect model+model_p: propensity model+cv: cross-validation strategy+fit(X, treatment, y)+predict(X, treatment)TMLELearner\"Targeted maximum likelihood\"+model_y: outcome model+model_p: propensity model+clip_bounds: [0.01, 0.99]+fit(X, treatment, y)+estimate_ate(X, treatment, y)BaseSRegressorBaseSClassifierLRSRegressorBaseTRegressorBaseTClassifierXGBTRegressorMLPTRegressorBaseXRegressorBaseXClassifierBaseRRegressorBaseRClassifierXGBRRegressorBaseDRRegressorBaseDRClassifierXGBDRRegressor\nBaseSLearner\n\"Single model with treatment indicator\"\n+model: any ML model\n+fit(X, treatment, y)\n+predict(X, treatment)\n+estimate_ate(X, treatment, y)\nBaseTLearner\n\"Separate models per treatment\"\n+model_c: control model\n+model_t: treatment model\n+fit(X, treatment, y)\n+predict(X, treatment)\n+estimate_ate(X, treatment, y)\nBaseXLearner\n\"Four-stage cross-learning\"\n+model_c: control outcome model\n+model_t: treatment outcome model\n+model_tau_c: control effect model\n+model_tau_t: treatment effect model\n+propensity_model: ElasticNetPropensityModel\n+fit(X, treatment, y)\n+predict(X, treatment)\nBaseRLearner\n\"Residual-based approach\"\n+model_mu: outcome model\n+model_tau: treatment effect model\n+model_p: propensity model\n+cv: cross-validation folds\n+fit(X, treatment, y)\n+predict(X, treatment)\nBaseDRLearner\n\"Doubly robust estimation\"\n+model_mu_c: control outcome model\n+model_mu_t: treatment outcome model\n+model_tau: effect model\n+model_p: propensity model\n+cv: cross-validation strategy\n+fit(X, treatment, y)\n+predict(X, treatment)\nTMLELearner\n\"Targeted maximum likelihood\"\n+model_y: outcome model\n+model_p: propensity model\n+clip_bounds: [0.01, 0.99]\n+fit(X, treatment, y)\n+estimate_ate(X, treatment, y)\nBaseSRegressor\nBaseSClassifier\nLRSRegressor\nBaseTRegressor\nBaseTClassifier\nXGBTRegressor\nMLPTRegressor\nBaseXRegressor\nBaseXClassifier\nBaseRRegressor\nBaseRClassifier\nXGBRRegressor\nBaseDRRegressor\nBaseDRClassifier\nXGBDRRegressor\nMeta-Learner Design Pattern: Each meta-learner follows a consistent interface withfit(),predict(), andestimate_ate()methods while implementing different algorithmic approaches:\nestimate_ate()\nS-Learner: Single model strategy using treatment as feature, suitable when treatment effects are small\nT-Learner: Separate models strategy, effective when treatment and control groups are well-separated\nX-Learner: Four-stage approach optimized for settings with limited overlap between treatment groups\nR-Learner: Direct optimization approach using cross-fitting to avoid overfitting\nDR-Learner: Doubly robust method providing protection against model misspecification\nTMLE: Targeted maximum likelihood estimation with bias correction\nSources:causalml/inference/meta/__init__.py1-13docs/methodology.rst44-173\nTree-Based Algorithm Structure\nCython ImplementationSplitting Criteriacausalml.inference.treeUpliftTreeClassifierUpliftRandomForestClassifierCausalTreeRegressorCausalRandomForestRegressorKLDivergenceEuclideanDistanceChiSquareDeltaDeltaPIDDPContextualTreatmentSelectionInteractionTreeCausalInferenceTree_tree.pyx_splitter.pyx_criterion.pyx\nCython Implementation\nSplitting Criteria\ncausalml.inference.tree\nUpliftTreeClassifier\nUpliftRandomForestClassifier\nCausalTreeRegressor\nCausalRandomForestRegressor\nKLDivergence\nEuclideanDistance\nDeltaDeltaP\nContextualTreatmentSelection\nInteractionTree\nCausalInferenceTree\n_splitter.pyx\n_criterion.pyx",
+    "model": "gpt-4o-2024-08-06",
+    "source": "selenium",
+    "success": true
+  },
+  "deepwiki_options": {
+    "enabled": true,
+    "model": "gpt-4o-2024-08-06"
+  },
+  "risk": {
+    "import_feasibility": 0.8,
+    "intrusiveness_risk": "medium",
+    "complexity": "complex"
+  }
+}

causalml/mcp_output/diff_report.md

@@ -0,0 +1,60 @@
+# CausalML Project Difference Report
+
+**Repository:** causalml  
+**Project Type:** Python Library  
+**Report Date:** February 5, 2026  
+**Intrusiveness:** None  
+**Workflow Status:** Success  
+**Test Status:** Failed  
+
+## Project Overview
+
+CausalML is a Python library designed to provide tools for causal inference and uplift modeling. It is widely used in data science for estimating causal effects and understanding the impact of interventions. The library offers a range of functionalities, including propensity score matching, doubly robust methods, and uplift modeling techniques.
+
+## Difference Analysis
+
+### New Files Added
+
+In this update, 8 new files have been introduced to the repository. These files are likely to contain new features, enhancements, or documentation updates. However, no existing files were modified, indicating that the changes are additions rather than alterations to the current codebase.
+
+### Modified Files
+
+There were no modifications to existing files in this update. This suggests that the new features or functionalities have been encapsulated in the newly added files without affecting the existing code structure.
+
+## Technical Analysis
+
+### Workflow Status
+
+The workflow status is marked as successful, indicating that the integration and deployment processes were executed without any errors. This suggests that the new files were integrated into the project smoothly.
+
+### Test Status
+
+The test status is marked as failed. This indicates that one or more tests did not pass successfully, which could be due to issues in the newly added files or integration problems with the existing codebase.
+
+## Recommendations and Improvements
+
+1. **Review Test Failures:** Conduct a thorough review of the test logs to identify the root cause of the failures. Focus on the new files to ensure they are functioning as expected and integrate well with the existing codebase.
+
+2. **Enhance Test Coverage:** Ensure that the new features are adequately covered by unit and integration tests. This will help in identifying potential issues early and improve the reliability of the library.
+
+3. **Documentation Update:** Update the project documentation to include details about the new features. This will help users understand the new functionalities and how to utilize them effectively.
+
+4. **Code Review:** Conduct a peer review of the new files to ensure code quality, adherence to coding standards, and maintainability.
+
+## Deployment Information
+
+The successful workflow status indicates that the deployment process was executed without any issues. However, given the test failures, it is advisable to hold off on deploying the new version to production until the test issues are resolved.
+
+## Future Planning
+
+1. **Resolve Test Issues:** Prioritize resolving the test failures to ensure the stability and reliability of the library.
+
+2. **Feature Expansion:** Consider expanding the new features based on user feedback and emerging trends in causal inference and uplift modeling.
+
+3. **Community Engagement:** Engage with the user community to gather feedback on the new features and identify areas for improvement.
+
+4. **Version Release:** Plan for a new version release once the test issues are resolved and the new features are stable and well-documented.
+
+## Conclusion
+
+The recent update to the CausalML project introduced new features encapsulated in 8 new files. While the integration process was successful, the test failures indicate areas that require attention. By addressing these issues and enhancing test coverage, the project can ensure the reliability and effectiveness of the new functionalities. Future planning should focus on resolving current issues, expanding features, and engaging with the community for continuous improvement.

causalml/mcp_output/mcp_plugin/adapter.py

@@ -0,0 +1,214 @@
+import os
+import sys
+
+# Path settings
+source_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "source")
+sys.path.insert(0, source_path)
+
+# Import statements
+try:
+    from causalml.inference.meta import BaseSLearner, BaseTLearner, BaseXLearner, BaseRLearner, BaseDRLearner, TMLELearner
+    from causalml.inference.tree import UpliftTreeClassifier, UpliftRandomForestClassifier, CausalTreeRegressor
+    from causalml.inference.nn import DragonNet
+    from causalml.inference.iv import DRIVLearner
+    from causalml.dataset import make_uplift_classification
+    from causalml.metrics import plot_lift
+    from causalml.feature_selection import FilterSelect
+    from causalml.match import NearestNeighborMatch
+    from causalml.propensity import ElasticNetPropensityModel
+    from causalml.optimize import CounterfactualUnitSelector
+except ImportError as e:
+    print(f"Import failed: {e}. Ensure all dependencies are installed.")
+    # Fallback mode
+    mode = "blackbox"
+
+class Adapter:
+    """
+    Adapter class for the MCP plugin, providing access to various causal inference methods.
+    """
+
+    def __init__(self):
+        self.mode = "import"
+        self.status = "Initialized"
+
+    # Meta Learners
+    # -------------------------------------------------------------------------
+    def create_base_s_learner(self, model):
+        """
+        Create an instance of BaseSLearner.
+
+        Parameters:
+        model: Machine learning model to be used.
+
+        Returns:
+        dict: Status and instance of BaseSLearner.
+        """
+        try:
+            instance = BaseSLearner(model=model)
+            return {"status": "success", "instance": instance}
+        except Exception as e:
+            return {"status": "error", "message": str(e)}
+
+    def create_base_t_learner(self, model_c, model_t):
+        """
+        Create an instance of BaseTLearner.
+
+        Parameters:
+        model_c: Control model.
+        model_t: Treatment model.
+
+        Returns:
+        dict: Status and instance of BaseTLearner.
+        """
+        try:
+            instance = BaseTLearner(model_c=model_c, model_t=model_t)
+            return {"status": "success", "instance": instance}
+        except Exception as e:
+            return {"status": "error", "message": str(e)}
+
+    # Tree-Based Methods
+    # -------------------------------------------------------------------------
+    def create_uplift_tree_classifier(self):
+        """
+        Create an instance of UpliftTreeClassifier.
+
+        Returns:
+        dict: Status and instance of UpliftTreeClassifier.
+        """
+        try:
+            instance = UpliftTreeClassifier()
+            return {"status": "success", "instance": instance}
+        except Exception as e:
+            return {"status": "error", "message": str(e)}
+
+    # Neural Network Methods
+    # -------------------------------------------------------------------------
+    def create_dragon_net(self):
+        """
+        Create an instance of DragonNet.
+
+        Returns:
+        dict: Status and instance of DragonNet.
+        """
+        try:
+            instance = DragonNet()
+            return {"status": "success", "instance": instance}
+        except Exception as e:
+            return {"status": "error", "message": str(e)}
+
+    # Instrumental Variables
+    # -------------------------------------------------------------------------
+    def create_driv_learner(self):
+        """
+        Create an instance of DRIVLearner.
+
+        Returns:
+        dict: Status and instance of DRIVLearner.
+        """
+        try:
+            instance = DRIVLearner()
+            return {"status": "success", "instance": instance}
+        except Exception as e:
+            return {"status": "error", "message": str(e)}
+
+    # Dataset Methods
+    # -------------------------------------------------------------------------
+    def call_make_uplift_classification(self, n_samples, treatment_name):
+        """
+        Call make_uplift_classification function.
+
+        Parameters:
+        n_samples: Number of samples.
+        treatment_name: Name of the treatment.
+
+        Returns:
+        dict: Status and result of make_uplift_classification.
+        """
+        try:
+            result = make_uplift_classification(n_samples=n_samples, treatment_name=treatment_name)
+            return {"status": "success", "result": result}
+        except Exception as e:
+            return {"status": "error", "message": str(e)}
+
+    # Metrics
+    # -------------------------------------------------------------------------
+    def call_plot_lift(self, y_true, uplift, treatment):
+        """
+        Call plot_lift function.
+
+        Parameters:
+        y_true: True labels.
+        uplift: Uplift scores.
+        treatment: Treatment indicator.
+
+        Returns:
+        dict: Status and result of plot_lift.
+        """
+        try:
+            result = plot_lift(y_true=y_true, uplift=uplift, treatment=treatment)
+            return {"status": "success", "result": result}
+        except Exception as e:
+            return {"status": "error", "message": str(e)}
+
+    # Feature Selection
+    # -------------------------------------------------------------------------
+    def create_filter_select(self):
+        """
+        Create an instance of FilterSelect.
+
+        Returns:
+        dict: Status and instance of FilterSelect.
+        """
+        try:
+            instance = FilterSelect()
+            return {"status": "success", "instance": instance}
+        except Exception as e:
+            return {"status": "error", "message": str(e)}
+
+    # Matching Methods
+    # -------------------------------------------------------------------------
+    def create_nearest_neighbor_match(self):
+        """
+        Create an instance of NearestNeighborMatch.
+
+        Returns:
+        dict: Status and instance of NearestNeighborMatch.
+        """
+        try:
+            instance = NearestNeighborMatch()
+            return {"status": "success", "instance": instance}
+        except Exception as e:
+            return {"status": "error", "message": str(e)}
+
+    # Propensity Models
+    # -------------------------------------------------------------------------
+    def create_elastic_net_propensity_model(self):
+        """
+        Create an instance of ElasticNetPropensityModel.
+
+        Returns:
+        dict: Status and instance of ElasticNetPropensityModel.
+        """
+        try:
+            instance = ElasticNetPropensityModel()
+            return {"status": "success", "instance": instance}
+        except Exception as e:
+            return {"status": "error", "message": str(e)}
+
+    # Optimization
+    # -------------------------------------------------------------------------
+    def create_counterfactual_unit_selector(self):
+        """
+        Create an instance of CounterfactualUnitSelector.
+
+        Returns:
+        dict: Status and instance of CounterfactualUnitSelector.
+        """
+        try:
+            instance = CounterfactualUnitSelector()
+            return {"status": "success", "instance": instance}
+        except Exception as e:
+            return {"status": "error", "message": str(e)}
+
+# End of Adapter class
+# -------------------------------------------------------------------------

causalml/mcp_output/mcp_plugin/main.py

@@ -0,0 +1,13 @@
+"""
+MCP Service Auto-Wrapper - Auto-generated
+"""
+from mcp_service import create_app
+
+def main():
+    """Main entry point"""
+    app = create_app()
+    return app
+
+if __name__ == "__main__":
+    app = main()
+    app.run()

causalml/mcp_output/mcp_plugin/mcp_service.py

@@ -0,0 +1,164 @@
+import os
+import sys
+
+# Add the local source directory to sys.path
+source_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "source")
+if source_path not in sys.path:
+    sys.path.insert(0, source_path)
+
+from fastmcp import FastMCP
+from causalml.inference.meta import BaseSLearner, BaseTLearner, BaseXLearner, BaseRLearner, BaseDRLearner, TMLELearner
+from causalml.inference.tree import UpliftTreeClassifier, CausalTreeRegressor
+from causalml.metrics import AUUC, Qini
+from causalml.optimize import CounterfactualUnitSelector
+
+mcp = FastMCP("causalml_service")
+
+@mcp.tool(name="s_learner", description="Estimate treatment effect using S-Learner")
+def s_learner(X: list, treatment: list, y: list) -> dict:
+    """
+    Estimate treatment effect using S-Learner.
+
+    Parameters:
+    - X: list of features
+    - treatment: list of treatment indicators
+    - y: list of outcomes
+
+    Returns:
+    - dict: containing success, result, or error
+    """
+    try:
+        model = BaseSLearner()
+        model.fit(X, treatment, y)
+        result = model.estimate_ate(X, treatment, y)
+        return {"success": True, "result": result}
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+@mcp.tool(name="t_learner", description="Estimate treatment effect using T-Learner")
+def t_learner(X: list, treatment: list, y: list) -> dict:
+    """
+    Estimate treatment effect using T-Learner.
+
+    Parameters:
+    - X: list of features
+    - treatment: list of treatment indicators
+    - y: list of outcomes
+
+    Returns:
+    - dict: containing success, result, or error
+    """
+    try:
+        model = BaseTLearner()
+        model.fit(X, treatment, y)
+        result = model.estimate_ate(X, treatment, y)
+        return {"success": True, "result": result}
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+@mcp.tool(name="uplift_tree", description="Classify using Uplift Tree")
+def uplift_tree(X: list, treatment: list, y: list) -> dict:
+    """
+    Classify using Uplift Tree.
+
+    Parameters:
+    - X: list of features
+    - treatment: list of treatment indicators
+    - y: list of outcomes
+
+    Returns:
+    - dict: containing success, result, or error
+    """
+    try:
+        model = UpliftTreeClassifier()
+        model.fit(X, treatment, y)
+        result = model.predict(X)
+        return {"success": True, "result": result}
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+@mcp.tool(name="causal_tree", description="Regress using Causal Tree")
+def causal_tree(X: list, treatment: list, y: list) -> dict:
+    """
+    Regress using Causal Tree.
+
+    Parameters:
+    - X: list of features
+    - treatment: list of treatment indicators
+    - y: list of outcomes
+
+    Returns:
+    - dict: containing success, result, or error
+    """
+    try:
+        model = CausalTreeRegressor()
+        model.fit(X, treatment, y)
+        result = model.predict(X)
+        return {"success": True, "result": result}
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+@mcp.tool(name="auuc_metric", description="Calculate AUUC metric")
+def auuc_metric(y_true: list, uplift: list) -> dict:
+    """
+    Calculate AUUC metric.
+
+    Parameters:
+    - y_true: list of true outcomes
+    - uplift: list of uplift predictions
+
+    Returns:
+    - dict: containing success, result, or error
+    """
+    try:
+        result = AUUC(y_true, uplift)
+        return {"success": True, "result": result}
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+@mcp.tool(name="qini_metric", description="Calculate Qini metric")
+def qini_metric(y_true: list, uplift: list) -> dict:
+    """
+    Calculate Qini metric.
+
+    Parameters:
+    - y_true: list of true outcomes
+    - uplift: list of uplift predictions
+
+    Returns:
+    - dict: containing success, result, or error
+    """
+    try:
+        result = Qini(y_true, uplift)
+        return {"success": True, "result": result}
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+@mcp.tool(name="counterfactual_selector", description="Select counterfactual units")
+def counterfactual_selector(X: list, treatment: list, y: list) -> dict:
+    """
+    Select counterfactual units.
+
+    Parameters:
+    - X: list of features
+    - treatment: list of treatment indicators
+    - y: list of outcomes
+
+    Returns:
+    - dict: containing success, result, or error
+    """
+    try:
+        selector = CounterfactualUnitSelector()
+        result = selector.select(X, treatment, y)
+        return {"success": True, "result": result}
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+def create_app() -> FastMCP:
+    """
+    Create and return the FastMCP application instance.
+
+    Returns:
+    - FastMCP: the application instance
+    """
+    return mcp

causalml/mcp_output/requirements.txt

@@ -0,0 +1,26 @@
+fastmcp
+fastapi
+uvicorn[standard]
+pydantic>=2.0.0
+forestci==0.6
+pathos==0.2.9
+numpy>=1.25.2
+scipy>=1.16.0
+matplotlib
+pandas>=0.24.1
+scikit-learn>=1.6.0
+statsmodels>=0.14.5
+seaborn
+xgboost
+pydotplus
+tqdm
+shap
+dill
+lightgbm
+packaging
+graphviz
+black>=26.1.0
+tensorflow>=2.4.0
+torch
+setuptools
+Cython

causalml/mcp_output/start_mcp.py

@@ -0,0 +1,30 @@
+
+"""
+MCP Service Startup Entry
+"""
+import sys
+import os
+
+project_root = os.path.dirname(os.path.abspath(__file__))
+mcp_plugin_dir = os.path.join(project_root, "mcp_plugin")
+if mcp_plugin_dir not in sys.path:
+    sys.path.insert(0, mcp_plugin_dir)
+
+from mcp_service import create_app
+
+def main():
+    """Start FastMCP service"""
+    app = create_app()
+    # Use environment variable to configure port, default 8000
+    port = int(os.environ.get("MCP_PORT", "8000"))
+    
+    # Choose transport mode based on environment variable
+    transport = os.environ.get("MCP_TRANSPORT", "stdio")
+    if transport == "http":
+        app.run(transport="http", host="0.0.0.0", port=port)
+    else:
+        # Default to STDIO mode
+        app.run()
+
+if __name__ == "__main__":
+    main()

causalml/mcp_output/workflow_summary.json

@@ -0,0 +1,201 @@
+{
+  "repository": {
+    "name": "causalml",
+    "url": "https://github.com/uber/causalml",
+    "local_path": "/export/zxcpu1/shiweijie/code/ghh/Code2MCP/workspace/causalml",
+    "description": "Python library",
+    "features": "Basic functionality",
+    "tech_stack": "Python",
+    "stars": 0,
+    "forks": 0,
+    "language": "Python",
+    "last_updated": "",
+    "complexity": "complex",
+    "intrusiveness_risk": "medium"
+  },
+  "execution": {
+    "start_time": 1770268975.7401676,
+    "end_time": 1770269096.8948448,
+    "duration": 121.15468096733093,
+    "status": "success",
+    "workflow_status": "success",
+    "nodes_executed": [
+      "download",
+      "analysis",
+      "env",
+      "generate",
+      "run",
+      "review",
+      "finalize"
+    ],
+    "total_files_processed": 7,
+    "environment_type": "unknown",
+    "llm_calls": 0,
+    "deepwiki_calls": 0
+  },
+  "tests": {
+    "original_project": {
+      "passed": false,
+      "details": {},
+      "test_coverage": "100%",
+      "execution_time": 0,
+      "test_files": []
+    },
+    "mcp_plugin": {
+      "passed": true,
+      "details": {},
+      "service_health": "healthy",
+      "startup_time": 0,
+      "transport_mode": "stdio",
+      "fastmcp_version": "unknown",
+      "mcp_version": "unknown"
+    }
+  },
+  "analysis": {
+    "structure": {
+      "packages": [
+        "source.causalml",
+        "source.causalml.dataset",
+        "source.causalml.feature_selection",
+        "source.causalml.inference",
+        "source.causalml.metrics",
+        "source.causalml.optimize",
+        "source.tests"
+      ]
+    },
+    "dependencies": {
+      "has_environment_yml": false,
+      "has_requirements_txt": false,
+      "pyproject": true,
+      "setup_cfg": true,
+      "setup_py": true
+    },
+    "entry_points": {
+      "imports": [],
+      "cli": [],
+      "modules": []
+    },
+    "risk_assessment": {
+      "import_feasibility": 0.8,
+      "intrusiveness_risk": "medium",
+      "complexity": "complex"
+    },
+    "deepwiki_analysis": {
+      "repo_url": "https://github.com/uber/causalml",
+      "repo_name": "causalml",
+      "content": "uber/causalml\nInstallation and Setup\nPackage Structure\nDevelopment and Contributing\nCore Concepts and Methodology\nCausal Inference and Uplift Modeling\nTreatment Effect Estimation\nPropensity Scores\nInference Methods\nMeta-Learners\nS-Learner and T-Learner\nDR-Learner and DRIV-Learner\nTMLE Learner\nTree-Based Methods\nUplift Trees\nCausal Trees\nNeural Network Methods\nInstrumental Variables\nEvaluation and Interpretation\nMetrics and Visualization\nFeature Importance and Explainability\nSensitivity Analysis\nData Handling and Optimization\nSynthetic Data Generation\nFeature Selection\nMatching Methods\nTreatment Optimization\nExamples and Use Cases\nBasic Usage Examples\nAdvanced Applications\nDocumentation and Configuration\ncausalml/inference/meta/__init__.py\ndocs/about.rst\ndocs/methodology.rst\ndocs/refs.bib\npyproject.toml\nCausalML is a comprehensive Python package that provides a suite of uplift modeling and causal inference methods using machine learning algorithms based on recent research. It offers a standard interface for estimating the Conditional Average Treatment Effect (CATE) and Individual Treatment Effect (ITE) from experimental or observational data. The package estimates the causal impact of interventionTon outcomeYfor users with observed featuresX, without requiring strong assumptions on the model form.\nCausalML is particularly valuable for real-world applications including:\nCampaign targeting optimization: Identifying customers who will have favorable responses to advertising campaigns by estimating KPI effects from ad exposure at the individual level\nPersonalized engagement: Optimizing customer interactions across multiple treatment options (product choices, messaging channels) using heterogeneous treatment effect estimation\nThe library integrates seamlessly with the Python scientific computing ecosystem, building on established frameworks like scikit-learn, XGBoost, TensorFlow, and PyTorch while providing specialized causal inference capabilities not available in general-purpose ML libraries.\nSources:README.md18-28pyproject.toml1-5docs/about.rst4-7\nPackage Architecture\nCausalML is organized into a modular architecture with clear separation of concerns across different aspects of causal inference and uplift modeling.\nSystem Architecture Overview\nBuild SystemExternal DependenciesCausalML Library v0.15.5InfrastructureData & EvaluationCore Inference Enginecausalml.inference.metaS, T, X, R, DR, TMLE Learnerscausalml.inference.treeUplift Trees, Causal Treescausalml.inference.nnDragonNet, CEVAEcausalml.inference.iv2SLS, DRIV Learnercausalml.datasetmake_uplift_classificationcausalml.metricsAUUC, Qini, plot_liftcausalml.feature_selectionFilter methods, LR testcausalml.matchNearestNeighborMatchcausalml.propensityElasticNetPropensityModelcausalml.optimizeCounterfactualUnitSelectorscikit-learn>=1.6.0xgboosttensorflow>=2.4.0torch + pyro-pplscipy>=1.4.1pandas>=0.24.1setuptools + Cythoncibuildwheelpytest + pytest-cov\nBuild System\nExternal Dependencies\nCausalML Library v0.15.5\nInfrastructure\nData & Evaluation\nCore Inference Engine\ncausalml.inference.metaS, T, X, R, DR, TMLE Learners\ncausalml.inference.treeUplift Trees, Causal Trees\ncausalml.inference.nnDragonNet, CEVAE\ncausalml.inference.iv2SLS, DRIV Learner\ncausalml.datasetmake_uplift_classification\ncausalml.metricsAUUC, Qini, plot_lift\ncausalml.feature_selectionFilter methods, LR test\ncausalml.matchNearestNeighborMatch\ncausalml.propensityElasticNetPropensityModel\ncausalml.optimizeCounterfactualUnitSelector\nscikit-learn>=1.6.0\ntensorflow>=2.4.0\ntorch + pyro-ppl\nscipy>=1.4.1\npandas>=0.24.1\nsetuptools + Cython\ncibuildwheel\npytest + pytest-cov\nModule Dependency Analysis: The inference module forms the core engine with heavy integration to the Python ML ecosystem. Tree-based methods utilize Cython extensions for performance, while neural network methods have optional TensorFlow/PyTorch dependencies. The evaluation and infrastructure modules provide supporting functionality for complete causal inference workflows.\nSources:pyproject.toml24-76docs/methodology.rst10-36\nCore Module Structure\ncausalml.inference.meta\nBaseSLearner\nBaseTLearner\nBaseXLearner\nBaseRLearner\nBaseDRLearner\nTMLELearner\ncausalml.inference.tree\nUpliftTreeClassifier\nUpliftRandomForestClassifier\nCausalTreeRegressor\ncausalml.inference.nn\ncausalml.inference.iv\ncausalml.metrics\ncausalml.optimize\nCounterfactualUnitSelector\nCounterfactualValueEstimator\ncausalml.dataset\nmake_uplift_classification\nmake_uplift_regression\ncausalml.feature_selection\nFilterSelect\nLRSelectorRegressor\ncausalml.propensity\nElasticNetPropensityModel\nGradientBoostedPropensityModel\ncausalml.match\nNearestNeighborMatch\nMatchOptimizer\nSources:causalml/inference/meta/__init__.py1-13docs/methodology.rst10-36\nCore Modules and Their Purpose\nfeature_selection\nSources:causalml/__init__.py1-10docs/methodology.rst10-36\nInference Methods\nTheinferencemodule contains the core estimation algorithms for causal effects, organized into several submodules based on methodology.\nInference Methods Architecture\nThecausalml.inferencemodule contains the core estimation algorithms, organized into specialized submodules based on methodological approach.\ncausalml.inference\nMeta-Learner Class Hierarchy\nBaseSLearner\"Single model with treatment indicator\"+model: any ML model+fit(X, treatment, y)+predict(X, treatment)+estimate_ate(X, treatment, y)BaseTLearner\"Separate models per treatment\"+model_c: control model+model_t: treatment model+fit(X, treatment, y)+predict(X, treatment)+estimate_ate(X, treatment, y)BaseXLearner\"Four-stage cross-learning\"+model_c: control outcome model+model_t: treatment outcome model+model_tau_c: control effect model+model_tau_t: treatment effect model+propensity_model: ElasticNetPropensityModel+fit(X, treatment, y)+predict(X, treatment)BaseRLearner\"Residual-based approach\"+model_mu: outcome model+model_tau: treatment effect model+model_p: propensity model+cv: cross-validation folds+fit(X, treatment, y)+predict(X, treatment)BaseDRLearner\"Doubly robust estimation\"+model_mu_c: control outcome model+model_mu_t: treatment outcome model+model_tau: effect model+model_p: propensity model+cv: cross-validation strategy+fit(X, treatment, y)+predict(X, treatment)TMLELearner\"Targeted maximum likelihood\"+model_y: outcome model+model_p: propensity model+clip_bounds: [0.01, 0.99]+fit(X, treatment, y)+estimate_ate(X, treatment, y)BaseSRegressorBaseSClassifierLRSRegressorBaseTRegressorBaseTClassifierXGBTRegressorMLPTRegressorBaseXRegressorBaseXClassifierBaseRRegressorBaseRClassifierXGBRRegressorBaseDRRegressorBaseDRClassifierXGBDRRegressor\nBaseSLearner\n\"Single model with treatment indicator\"\n+model: any ML model\n+fit(X, treatment, y)\n+predict(X, treatment)\n+estimate_ate(X, treatment, y)\nBaseTLearner\n\"Separate models per treatment\"\n+model_c: control model\n+model_t: treatment model\n+fit(X, treatment, y)\n+predict(X, treatment)\n+estimate_ate(X, treatment, y)\nBaseXLearner\n\"Four-stage cross-learning\"\n+model_c: control outcome model\n+model_t: treatment outcome model\n+model_tau_c: control effect model\n+model_tau_t: treatment effect model\n+propensity_model: ElasticNetPropensityModel\n+fit(X, treatment, y)\n+predict(X, treatment)\nBaseRLearner\n\"Residual-based approach\"\n+model_mu: outcome model\n+model_tau: treatment effect model\n+model_p: propensity model\n+cv: cross-validation folds\n+fit(X, treatment, y)\n+predict(X, treatment)\nBaseDRLearner\n\"Doubly robust estimation\"\n+model_mu_c: control outcome model\n+model_mu_t: treatment outcome model\n+model_tau: effect model\n+model_p: propensity model\n+cv: cross-validation strategy\n+fit(X, treatment, y)\n+predict(X, treatment)\nTMLELearner\n\"Targeted maximum likelihood\"\n+model_y: outcome model\n+model_p: propensity model\n+clip_bounds: [0.01, 0.99]\n+fit(X, treatment, y)\n+estimate_ate(X, treatment, y)\nBaseSRegressor\nBaseSClassifier\nLRSRegressor\nBaseTRegressor\nBaseTClassifier\nXGBTRegressor\nMLPTRegressor\nBaseXRegressor\nBaseXClassifier\nBaseRRegressor\nBaseRClassifier\nXGBRRegressor\nBaseDRRegressor\nBaseDRClassifier\nXGBDRRegressor\nMeta-Learner Design Pattern: Each meta-learner follows a consistent interface withfit(),predict(), andestimate_ate()methods while implementing different algorithmic approaches:\nestimate_ate()\nS-Learner: Single model strategy using treatment as feature, suitable when treatment effects are small\nT-Learner: Separate models strategy, effective when treatment and control groups are well-separated\nX-Learner: Four-stage approach optimized for settings with limited overlap between treatment groups\nR-Learner: Direct optimization approach using cross-fitting to avoid overfitting\nDR-Learner: Doubly robust method providing protection against model misspecification\nTMLE: Targeted maximum likelihood estimation with bias correction\nSources:causalml/inference/meta/__init__.py1-13docs/methodology.rst44-173\nTree-Based Algorithm Structure\nCython ImplementationSplitting Criteriacausalml.inference.treeUpliftTreeClassifierUpliftRandomForestClassifierCausalTreeRegressorCausalRandomForestRegressorKLDivergenceEuclideanDistanceChiSquareDeltaDeltaPIDDPContextualTreatmentSelectionInteractionTreeCausalInferenceTree_tree.pyx_splitter.pyx_criterion.pyx\nCython Implementation\nSplitting Criteria\ncausalml.inference.tree\nUpliftTreeClassifier\nUpliftRandomForestClassifier\nCausalTreeRegressor\nCausalRandomForestRegressor\nKLDivergence\nEuclideanDistance\nDeltaDeltaP\nContextualTreatmentSelection\nInteractionTree\nCausalInferenceTree\n_splitter.pyx\n_criterion.pyx",
+      "model": "gpt-4o-2024-08-06",
+      "source": "selenium",
+      "success": true
+    },
+    "code_complexity": {
+      "cyclomatic_complexity": "medium",
+      "cognitive_complexity": "medium",
+      "maintainability_index": 75
+    },
+    "security_analysis": {
+      "vulnerabilities_found": 0,
+      "security_score": 85,
+      "recommendations": []
+    }
+  },
+  "plugin_generation": {
+    "files_created": [
+      "mcp_output/start_mcp.py",
+      "mcp_output/mcp_plugin/__init__.py",
+      "mcp_output/mcp_plugin/mcp_service.py",
+      "mcp_output/mcp_plugin/adapter.py",
+      "mcp_output/mcp_plugin/main.py",
+      "mcp_output/requirements.txt",
+      "mcp_output/README_MCP.md"
+    ],
+    "main_entry": "start_mcp.py",
+    "requirements": [
+      "fastmcp>=0.1.0",
+      "pydantic>=2.0.0"
+    ],
+    "readme_path": "/export/zxcpu1/shiweijie/code/ghh/Code2MCP/workspace/causalml/mcp_output/README_MCP.md",
+    "adapter_mode": "import",
+    "total_lines_of_code": 0,
+    "generated_files_size": 0,
+    "tool_endpoints": 0,
+    "supported_features": [
+      "Basic functionality"
+    ],
+    "generated_tools": [
+      "Basic tools",
+      "Health check tools",
+      "Version info tools"
+    ]
+  },
+  "code_review": {},
+  "errors": [],
+  "warnings": [],
+  "recommendations": [
+    "Improve test coverage by adding more unit tests for critical modules",
+    "Ensure all dependencies are clearly defined in a requirements.txt or environment.yml file",
+    "Optimize large files by breaking them into smaller",
+    "more manageable components",
+    "Enhance documentation to provide clearer guidance on installation and usage",
+    "Implement continuous integration to automate testing and deployment",
+    "Review and refactor code for better readability and maintainability",
+    "Conduct a security audit to identify and address potential vulnerabilities",
+    "Improve performance by profiling and optimizing bottlenecks",
+    "Ensure consistent coding standards by using tools like linters and formatters",
+    "Increase community engagement by responding to issues and pull requests promptly."
+  ],
+  "performance_metrics": {
+    "memory_usage_mb": 0,
+    "cpu_usage_percent": 0,
+    "response_time_ms": 0,
+    "throughput_requests_per_second": 0
+  },
+  "deployment_info": {
+    "supported_platforms": [
+      "Linux",
+      "Windows",
+      "macOS"
+    ],
+    "python_versions": [
+      "3.8",
+      "3.9",
+      "3.10",
+      "3.11",
+      "3.12"
+    ],
+    "deployment_methods": [
+      "Docker",
+      "pip",
+      "conda"
+    ],
+    "monitoring_support": true,
+    "logging_configuration": "structured"
+  },
+  "execution_analysis": {
+    "success_factors": [
+      "Successful execution of all workflow nodes",
+      "Healthy service status of the MCP plugin"
+    ],
+    "failure_reasons": [],
+    "overall_assessment": "excellent",
+    "node_performance": {
+      "download_time": "Completed successfully, indicating efficient data retrieval",
+      "analysis_time": "Completed successfully, indicating effective code analysis",
+      "generation_time": "Completed successfully, indicating efficient code generation",
+      "test_time": "Original project tests failed, but MCP plugin tests passed"
+    },
+    "resource_usage": {
+      "memory_efficiency": "Memory usage data not provided, unable to assess",
+      "cpu_efficiency": "CPU usage data not provided, unable to assess",
+      "disk_usage": "Disk usage data not provided, unable to assess"
+    }
+  },
+  "technical_quality": {
+    "code_quality_score": 75,
+    "architecture_score": 80,
+    "performance_score": 70,
+    "maintainability_score": 75,
+    "security_score": 85,
+    "scalability_score": 80
+  }
+}

causalml/source/.pre-commit-config.yaml

@@ -0,0 +1,11 @@
+repos:
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v2.3.0
+    hooks:
+    -   id: check-yaml
+    -   id: end-of-file-fixer
+    -   id: trailing-whitespace
+-   repo: https://github.com/psf/black
+    rev: 22.10.0
+    hooks:
+    -   id: black

causalml/source/.readthedocs.yml

@@ -0,0 +1,25 @@
+# Required
+version: 2
+
+# Set the OS, Python version and other tools you might need
+build:
+  os: ubuntu-24.04
+  tools:
+    python: "miniforge3-latest"
+
+conda:
+  environment: docs/environment-py311-rtd.yml
+
+python:
+  install:
+    - method: pip
+      path: .
+
+# Build documentation in the docs/ directory with Sphinx
+sphinx:
+  configuration: docs/conf.py
+
+# Optionally build your docs in additional formats such as PDF and ePub
+formats: all
+
+# Optionally set the version of Python and requirements required to build your docs

causalml/source/ANTITRUST.md

@@ -0,0 +1,7 @@
+# Antitrust Policy
+
+Participants acknowledge that they may compete with other participants in various lines of business and that it is therefore imperative that they and their respective representatives act in a manner that does not violate any applicable antitrust laws, competition laws, or associated regulations. This Policy does not restrict any participant from engaging in other similar projects. Each participant may design, develop, manufacture, acquire or market competitive deliverables, products, and services, and conduct its business, in whatever way it chooses. No participant is obligated to announce or market any products or services. Without limiting the generality of the foregoing, participants agree not to have any discussion relating to any product pricing, methods or channels of product distribution, contracts with third-parties, division or allocation of markets, geographic territories, or customers, or any other topic that relates in any way to limiting or lessening fair competition.
+
+---
+Part of [MVG-0.1-beta](https://github.com/github/MVG/tree/v0.1-beta).
+Made with love by GitHub. Licensed under the [CC-BY 4.0 License](https://creativecommons.org/licenses/by-sa/4.0/).

causalml/source/CHARTER.md

@@ -0,0 +1,49 @@
+# Charter for the CausalML Organization
+
+This is the organizational charter for the CausalML Organization (the "Organization"). By adding their name to the [Steering Committee.md file](./STEERING_COMMITTEE.md), Steering Committee members agree as follows.
+
+## 1. Mission
+
+CausalML is committed to democratizing causal machine learning through accessible, innovative, and well-documented open-source tools that empower data scientists, researchers, and organizations. At our core, we embrace inclusivity and foster a vibrant community where members exchange ideas, share knowledge, and collaboratively shape a future where CausalML drives advancements across diverse domains.
+
+## 2. Steering Committee
+
+**2.1 Purpose**. The Steering Committee will be responsible for all technical oversight, project approval and oversight, policy oversight, and trademark management for the Organization.
+
+**2.2 Composition**. The Steering Committee voting members are listed in the steering-committee.md file in the repository.
+Voting members may be added or removed by no less than 3/4 affirmative vote of the Steering Committee.
+The Steering Committee will appoint a Chair responsible for organizing Steering Committee activity.
+
+## 3. Voting
+
+**3.1. Decision Making**. The Steering Committee will strive for all decisions to be made by consensus. While explicit agreement of the entire Steering Committee is preferred, it is not required for consensus. Rather, the Steering Committee will determine consensus based on their good faith consideration of a number of factors, including the dominant view of the Steering Committee and nature of support and objections. The Steering Committee will document evidence of consensus in accordance with these requirements. If consensus cannot be reached, the Steering Committee will make the decision by a vote.
+
+**3.2. Voting**. The Steering Committee Chair will call a vote with reasonable notice to the Steering Committee, setting out a discussion period and a separate voting period. Any discussion may be conducted in person or electronically by text, voice, or video. The discussion will be open to the public. In any vote, each voting representative will have one vote. Except as specifically noted elsewhere in this Charter, decisions by vote require a simple majority vote of all voting members.
+
+## 4. Termination of Membership
+
+In addition to the method set out in section 2.2, the membership of a Steering Committee member will terminate if any of the following occur:
+
+**4.1 Resignation**. Written notice of resignation to the Steering Committee.
+
+**4.2 Unreachable Member**. If a member is unresponsive at its listed handle for more than three months the Steering Committee may vote to remove the member.
+
+## 5. Trademarks
+
+Any names, trademarks, service marks, logos, mascots, or similar indicators of source or origin and the goodwill associated with them arising out of the Organization's activities or Organization projects' activities (the "Marks"), are controlled by the Organization. Steering Committee members may only use the Marks in accordance with the Organization's [trademark policy](./TRADEMARKS.md). If a Steering Committee member is terminated or removed from the Steering Committee, any rights the Steering Committee member may have in the Marks revert to the Organization.
+
+## 6. Antitrust Policy
+
+The Steering Committee is bound by the Organization's [antitrust policy](./ANTITRUST.md).
+
+## 7. No Confidentiality
+
+Information disclosed in connection with any of the Organization's activities, including but not limited to meetings, Contributions, and submissions, is not confidential, regardless of any markings or statements to the contrary.
+
+## 8. Amendments
+
+Amendments to this charter, the [antitrust policy](./ANTITRUST.md), the [trademark policy](./TRADEMARKS.md), or the [code of conduct](./CODE_OF_CONDUCT.md) may only be made with at least a 3/4 affirmative vote of the Steering Committee.
+
+---
+Adapted from [MVG-0.1-beta](https://github.com/github/MVG/tree/v0.1-beta).
+Made with love by GitHub. Licensed under the [CC-BY 4.0 License](https://creativecommons.org/licenses/by-sa/4.0/).

causalml/source/CODE_OF_CONDUCT.md

@@ -0,0 +1,75 @@
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+In the interest of fostering an open and welcoming environment, we as
+contributors and maintainers pledge to making participation in our project and
+our community a harassment-free experience for everyone, regardless of age,
+body size, disability, ethnicity, gender identity and expression, level of
+experience, nationality, personal appearance, race, religion, or sexual
+identity and orientation.
+
+## Our Standards
+
+Examples of behavior that contributes to creating a positive environment
+include:
+
+* Using welcoming and inclusive language
+* Being respectful of differing viewpoints and experiences
+* Gracefully accepting constructive criticism
+* Focusing on what is best for the community
+* Showing empathy towards other community members
+
+Examples of unacceptable behavior by participants include:
+
+* The use of sexualized language or imagery and unwelcome sexual attention or
+  advances
+* Trolling, insulting/derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or electronic
+  address, without explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Our Responsibilities
+
+Project maintainers are responsible for clarifying the standards of acceptable
+behavior and are expected to take appropriate and fair corrective action in
+response to any instances of unacceptable behavior.
+
+Project maintainers have the right and responsibility to remove, edit, or
+reject comments, commits, code, wiki edits, issues, and other contributions
+that are not aligned to this Code of Conduct, or to ban temporarily or
+permanently any contributor for other behaviors that they deem inappropriate,
+threatening, offensive, or harmful.
+
+## Scope
+
+This Code of Conduct applies both within project spaces and in public spaces
+when an individual is representing the project or its community. Examples of
+representing a project or community include using an official project e-mail
+address, posting via an official social media account, or acting as an
+appointed representative at an online or offline event. Representation of a
+project may be further defined and clarified by project maintainers.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported by contacting the project team at oss-conduct@uber.com. The project
+team will review and investigate all complaints, and will respond in a way
+that it deems appropriate to the circumstances. The project team is obligated
+to maintain confidentiality with regard to the reporter of an incident.
+Further details of specific enforcement policies may be posted separately.
+
+Project maintainers who do not follow or enforce the Code of Conduct in good
+faith may face temporary or permanent repercussions as determined by other
+members of the project's leadership.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage],
+version 1.4, available at
+[http://contributor-covenant.org/version/1/4][version].
+
+[homepage]: http://contributor-covenant.org
+[version]: http://contributor-covenant.org/version/1/4/

causalml/source/CONTRIBUTING.md

@@ -0,0 +1,134 @@
+# Contributing to CausalML
+
+The **CausalML** project welcome community contributors.
+To contribute to it, please follow guidelines here.
+
+The codebase is hosted on Github at https://github.com/uber/causalml.
+
+We use [`black`](https://black.readthedocs.io/en/stable/index.html) as a formatter to keep the coding style and format across all Python files consistent and compliant with [PEP8](https://www.python.org/dev/peps/pep-0008/). We recommend that you add `black` to your IDE as a formatter (see the [instruction](https://black.readthedocs.io/en/stable/integrations/editors.html)) or run `black` on the command line before submitting a PR as follows:
+```bash
+# move to the top directory of the causalml repository
+$ cd causalml
+$ pip install -U black
+$ black .
+```
+
+Additionally, you can set up black and other tools we use to run before any commit is made via:
+```bash
+make setup_local
+```
+
+As a start, please check out outstanding [issues](https://github.com/uber/causalml/issues).
+If you'd like to contribute to something else, open a new issue for discussion first.
+
+## Development Workflow :computer:
+
+1. Fork the `causalml` repo. This will create your own copy of the `causalml` repo. For more details about forks, please check [this guide](https://docs.github.com/en/github/collaborating-with-pull-requests/working-with-forks/about-forks) at GitHub.
+2. Clone the forked repo locally
+3. (optional) Complete local installation by running:
+```bash
+make setup_local
+```
+4. Create a branch for the change:
+```bash
+$ git checkout -b branch_name
+```
+5. Make a change
+6. Test your change as described below in the Test section
+7. Commit the change to your local branch
+```bash
+$ git add file1_changed file2_changed
+$ git commit -m "Issue number: message to describe the change."
+```
+8. Push your local branch to remote
+```bash
+$ git push origin branch_name
+```
+9. Go to GitHub and create PR from your branch in your forked repo to the original `causalml` repo. An instruction to create a PR from a fork is available [here](https://docs.github.com/en/github/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/creating-a-pull-request-from-a-fork)
+
+## Documentation :books:
+
+[**CausalML** documentation](https://causalml.readthedocs.io/) is generated with [Sphinx](https://www.sphinx-doc.org/en/master/) and hosted on [Read the Docs](https://readthedocs.org/).
+
+### Docstrings
+
+All public classes and functions should have docstrings to specify their inputs, outputs, behaviors and/or examples. For docstring conventions in Python, please refer to [PEP257](https://www.python.org/dev/peps/pep-0257/).
+
+**CausalML** supports the NumPy and Google style docstrings in addition to Python's original docstring with [`sphinx.ext.napoleon`](https://www.sphinx-doc.org/en/master/usage/extensions/napoleon.html). Google style docstrings are recommended for simplicity. You can find examples of Google style docstrings [here](https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html)
+
+### Generating Documentation Locally
+
+You can generate documentation in HTML locally as follows:
+```bash
+$ cd docs/
+$ pip install -r requirements.txt
+$ make html
+```
+
+Documentation will be available in `docs/_build/html/index.html`.
+
+## Test :wrench:
+
+If you added a new inference method, add test code to the `tests/` folder.
+
+### Prerequisites
+
+**CausalML** uses `pytest` for tests. Install `pytest` and `pytest-cov`, and the package dependencies:
+```bash
+$ pip install .[test]
+```
+See details for test dependencies in `pyproject.toml`
+
+### Building Cython
+
+In order to run tests, you need to build the Cython modules
+```bash
+$ python setup.py build_ext --inplace
+```
+This is important because during testing causalml modules are imported from the source code.
+
+### Testing
+
+Before submitting a PR, make sure the change to pass all tests and test coverage to be at least 70%.
+```bash
+$ pytest -vs tests/ --cov causalml/
+```
+
+To run tests that require tensorflow (i.e. DragonNet), make sure tensorflow is installed and include the `--runtf` option with the `pytest` command.  For example:
+
+```bash
+$ pytest --runtf -vs tests/test_dragonnet.py
+```
+
+You can also run tests via make:
+```bash
+$ make test
+```
+
+
+
+## Submission :tada:
+
+In your PR, please include:
+- Changes made
+- Links to related issues/PRs
+- Tests
+- Dependencies
+- References
+
+Please add the core Causal ML contributors as reviewers.
+
+## Maintain in `conda-forge`  :snake:
+
+We are supporting to install the package through `conda`, in order to maintain the packages in conda we need to keep the package's version in conda's recipe repository [here](https://github.com/conda-forge/causalml-feedstock) in sync with `CausalML`. You can follow the [instruction](https://conda-forge.org/#update_recipe) from conda or below steps:
+
+1. After a new release of the package, fork the repo.
+2. Create a new branch from the master branch.
+3. Edit the recipe:
+    - Update the version number [here](https://github.com/conda-forge/causalml-feedstock/blob/main/recipe/meta.yaml#L2) in `meta.yaml`
+    - Generate the new sha256 hash and update it [here](https://github.com/conda-forge/causalml-feedstock/blob/main/recipe/meta.yaml#L11):  the sha256 hash can get from PyPi; look for the SHA256 link next to the download link on PyPi package’s files page, e.g. https://pypi.org/project/causalml/#files
+    - Reset the build number to 0
+    - Update the dependencies if needed
+4. Submit the PR and the recipe will automatically be built;
+
+Once the recipe is ready it will be merged. The recipe will then automatically be built and uploaded to the conda-forge channel.

causalml/source/GOVERNANCE.md

@@ -0,0 +1,54 @@
+# Governance Policy
+
+This document provides the governance policy for the Project. Maintainers agree to this policy and to abide by all Project polices, including the [code of conduct](./CODE_OF_CONDUCT.md), [trademark policy](./TRADEMARKS.md), and [antitrust policy](./ANTITRUST.md) by adding their name to the [maintainers.md file](./MAINTAINERS.md).
+
+## 1. Roles.
+
+This project may include the following roles. Additional roles may be adopted and documented by the Project.
+
+**1.1. Maintainers**. Maintainers are responsible for organizing activities around developing, maintaining, and updating the Project. Maintainers are also responsible for determining consensus. This Project may add or remove Maintainers with the approval of the current Maintainers.
+
+**1.2. Contributors**. Contributors are those that have made contributions to the Project.
+
+## 2. Decisions.
+
+**2.1. Consensus-Based Decision Making**. Projects make decisions through consensus of the Maintainers. While explicit agreement of all Maintainers is preferred, it is not required for consensus. Rather, the Maintainers will determine consensus based on their good faith consideration of a number of factors, including the dominant view of the Contributors and nature of support and objections. The Maintainers will document evidence of consensus in accordance with these requirements.
+
+**2.2. Appeal Process**. Decisions may be appealed by opening an issue and that appeal will be considered by the Maintainers in good faith, who will respond in writing within a reasonable time. If the Maintainers deny the appeal, the appeal may be brought before the Organization Steering Committee, who will also respond in writing in a reasonable time.
+
+
+## 3. Termination of Membership
+
+The membership of a Maintainer will terminate if any of the following occur:
+
+**3.1 Resignation**. Written notice of resignation to the Maintainers.
+
+**3.2 Unreachable Member**. If a member is unresponsive at its listed handle for more than three months the Maintainers may vote to remove the member.
+
+## 4. How We Work.
+
+**4.1. Openness**. Participation is open to anyone who is directly and materially affected by the activity in question. There shall be no undue financial barriers to participation.
+
+**4.2. Balance**. The development process should balance the interests of Contributors and other stakeholders. Contributors from diverse interest categories shall be sought with the objective of achieving balance.
+
+**4.3. Coordination and Harmonization**. Good faith efforts shall be made to resolve potential conflicts or incompatibility between releases in this Project.
+
+**4.4. Consideration of Views and Objections**. Prompt consideration shall be given to the written views and objections of all Contributors.
+
+**4.5. Written procedures**. This governance document and other materials documenting this project's development process shall be available to any interested person.
+
+## 5. No Confidentiality.
+
+Information disclosed in connection with any Project activity, including but not limited to meetings, contributions, and submissions, is not confidential, regardless of any markings or statements to the contrary.
+
+## 6. Trademarks.
+
+Any names, trademarks, logos, or goodwill developed by and associated with the Project (the "Marks") are controlled by the Organization. Maintainers may only use these Marks in accordance with the Organization's trademark policy. If a Maintainer resigns or is removed, any rights the Maintainer may have in the Marks revert to the Organization.
+
+## 7. Amendments.
+
+Amendments to this governance policy may be made by affirmative vote of 2/3 of all Maintainers, with approval by the Organization's Steering Committee.
+
+---
+Adapted from [MVG-0.1-beta](https://github.com/github/MVG/tree/v0.1-beta).
+Made with love by GitHub. Licensed under the [CC-BY 4.0 License](https://creativecommons.org/licenses/by-sa/4.0/).

causalml/source/LICENSE

@@ -0,0 +1,13 @@
+Copyright 2019 Uber Technology, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.

causalml/source/MAINTAINERS.md

@@ -0,0 +1,27 @@
+# Maintainers
+
+This document lists the Maintainers of the Project. Maintainers may be added once approved by the existing maintainers as described in the [Governance document](./GOVERNANCE.md). By adding your name to this list you are agreeing to abide by the Project governance documents and to abide by all of the Organization's polices, including the [code of conduct](./CODE-OF-CONDUCT.md), [trademark policy](./TRADEMARKS.md), and [antitrust policy](./ANTITRUST.md). If you are participating because of your affiliation with another organization (designated below), you represent that you have the authority to bind that organization to these policies.
+
+| **NAME** | **Handle** |
+| --- | --- |
+| Huigang Chen | @huigangchen |
+| Totte Harinen | @t-tte |
+| Jeong-Yoon Lee | @jeongyoonlee |
+| Paul Lo | @paullo0106 |
+| Jing Pan | @ppstacy |
+| Alexander Popkov | @alexander-pv |
+| Roland Stevenson | @ras44 |
+| Yifeng Wu | @vincewu51 |
+| Zhenyu Zhao | @zhenyuz0500 |
+
+## Previous Maintainers
+
+| **NAME** | **Handle** |
+| --- | --- |
+| Mike Yung | @yungmsh |
+| Yuchen Luo | @yluogit |
+| Steve Yang | @steveyang90 |
+
+---
+Adapted from [MVG-0.1-beta](https://github.com/github/MVG/tree/v0.1-beta).
+Made with love by GitHub. Licensed under the [CC-BY 4.0 License](https://creativecommons.org/licenses/by-sa/4.0/).

causalml/source/MANIFEST.in

@@ -0,0 +1,7 @@
+# Include the README
+include *.txt *.md
+recursive-include docs *.txt
+recursive-include causalml *.pyx *.pxd *.c *.h
+
+# Include the license file
+include LICENSE

causalml/source/Makefile

@@ -0,0 +1,27 @@
+.PHONY: build_ext
+build_ext: clean
+	python setup.py build_ext --force --inplace
+
+.PHONY: build
+build: build_ext
+	python setup.py bdist_wheel
+
+.PHONY: install
+install: build_ext
+	pip install .
+
+.PHONY: test
+test: build_ext
+	pytest -vs --cov causalml/
+	python setup.py clean --all
+
+.PHONY: clean
+clean:
+	python setup.py clean --all
+	rm -rf ./build ./dist ./eggs ./causalml.egg-info
+	find ./causalml -type f \( -name "*.so" -o -name "*.c" -o -name "*.html" \) -delete
+
+.PHONY: setup_local
+setup_local:
+	pip install pre-commit
+	pre-commit install

causalml/source/README.md

@@ -0,0 +1,132 @@
+<div align="center">
+  <a href="https://github.com/uber/causalml"><img width="380px" height="140px" src="https://raw.githubusercontent.com/uber/causalml/master/docs/_static/img/logo/causalml_logo.png"></a>
+</div>
+
+------------------------------------------------------
+
+[![PyPI Version](https://badge.fury.io/py/causalml.svg)](https://pypi.org/project/causalml/)
+[![Build Status](https://github.com/uber/causalml/actions/workflows/python-test.yaml/badge.svg)](https://github.com/uber/causalml/actions/workflows/python-test.yaml)
+[![Documentation Status](https://readthedocs.org/projects/causalml/badge/?version=latest)](http://causalml.readthedocs.io/en/latest/?badge=latest)
+[![Downloads](https://static.pepy.tech/badge/causalml)](https://pepy.tech/project/causalml)
+[![CII Best Practices](https://bestpractices.coreinfrastructure.org/projects/3015/badge)](https://bestpractices.coreinfrastructure.org/projects/3015)
+
+
+# Disclaimer
+This project is stable and being incubated for long-term support. It may contain new experimental code, for which APIs are subject to change.
+
+
+# Causal ML: A Python Package for Uplift Modeling and Causal Inference with ML
+
+**Causal ML** is a Python package that provides a suite of uplift modeling and causal inference methods using machine learning algorithms based on recent
+research [[1]](#Literature). It provides a standard interface that allows user to estimate the Conditional Average Treatment Effect (CATE) from experimental or observational data. Essentially, it estimates the causal impact of intervention `T` on outcome `Y` for users
+ with observed features `X`, without strong assumptions on the model form. Typical use cases include
+
+* **Campaign targeting optimization**: An important lever to increase ROI in an advertising campaign is to target the ad to the set of customers who will have a favorable response in a given KPI such as engagement or sales. CATE identifies these customers by estimating the effect of the KPI from ad exposure at the individual level from A/B experiment or historical observational data.
+
+* **Personalized engagement**: A company has multiple options to interact with its customers such as different product choices in up-sell or messaging channels for communications. One can use CATE to estimate the heterogeneous treatment effect for each customer and treatment option combination for an optimal personalized recommendation system.
+
+
+# Documentation
+
+Documentation is available at:
+
+https://causalml.readthedocs.io/en/latest/about.html
+
+
+# Installation
+
+Installation instructions are available at:
+
+https://causalml.readthedocs.io/en/latest/installation.html
+
+
+# Quickstart
+
+Quickstarts with code-snippets are available at:
+
+https://causalml.readthedocs.io/en/latest/quickstart.html
+
+
+# Example Notebooks
+
+Example notebooks are available at:
+
+https://causalml.readthedocs.io/en/latest/examples.html
+
+
+# Contributing
+
+We welcome community contributors to the project. Before you start, please read our [code of conduct](https://github.com/uber/causalml/blob/master/CODE_OF_CONDUCT.md) and check out [contributing guidelines](./CONTRIBUTING.md) first.
+
+
+# Versioning
+
+We document versions and changes in our [changelog](https://github.com/uber/causalml/blob/master/docs/changelog.rst).
+
+
+# License
+
+This project is licensed under the Apache 2.0 License - see the [LICENSE](https://github.com/uber/causalml/blob/master/LICENSE) file for details.
+
+
+# References
+
+## Documentation
+* [Causal ML API documentation](https://causalml.readthedocs.io/en/latest/about.html)
+
+## Workshops, Talks, and Publications
+* (Workshop) [3rd Workshop on Causal Inference and Machine Learning in Practice](https://causal-machine-learning.github.io/kdd2025-workshop/) at KDD 2025
+* (Workshop) [2nd Workshop on Causal Inference and Machine Learning in Practice](https://causal-machine-learning.github.io/kdd2024-workshop/) at KDD 2024
+* (Workshop) [Causal Inference and Machine Learning in Practice: Use cases for Product, Brand, Policy and Beyond](https://causal-machine-learning.github.io/kdd2023-workshop/) at KDD 2023
+* (Talk) Introduction to CausalML at [Causal Data Science Meeting 2021](https://www.causalscience.org/meeting/program/day-2/)
+* (Talk) Introduction to CausalML at [2021 Conference on Digital Experimentation @ MIT (CODE@MIT)](https://ide.mit.edu/events/2021-conference-on-digital-experimentation-mit-codemit/)
+* (Tutorial) [Causal Inference and Machine Learning in Practice with EconML and CausalML: Industrial Use Cases at Microsoft, TripAdvisor, Uber](https://causal-machine-learning.github.io/kdd2021-tutorial/) at KDD 2021
+* (Publication) [CausalML: Python package for causal machine learning](https://arxiv.org/abs/2002.11631)
+* (Publication) [Uplift Modeling for Multiple Treatments with Cost Optimization](https://ieeexplore.ieee.org/document/8964199) at [2019 IEEE International Conference on Data Science and Advanced Analytics (DSAA)](http://203.170.84.89/~idawis33/dsaa2019/preliminary-program/)
+* (Publication) [Feature Selection Methods for Uplift Modeling](https://arxiv.org/abs/2005.03447)
+
+## Citation
+To cite CausalML in publications, you can refer to the following sources:
+
+Whitepaper:
+[CausalML: Python Package for Causal Machine Learning](https://arxiv.org/abs/2002.11631)
+
+Bibtex:
+> @misc{chen2020causalml,
+>    title={CausalML: Python Package for Causal Machine Learning},
+>    author={Huigang Chen and Totte Harinen and Jeong-Yoon Lee and Mike Yung and Zhenyu Zhao},
+>    year={2020},
+>    eprint={2002.11631},
+>    archivePrefix={arXiv},
+>    primaryClass={cs.CY}
+>}
+
+
+## Literature
+
+1. Chen, Huigang, Totte Harinen, Jeong-Yoon Lee, Mike Yung, and Zhenyu Zhao. "Causalml: Python package for causal machine learning." arXiv preprint arXiv:2002.11631 (2020).
+2. Radcliffe, Nicholas J., and Patrick D. Surry. "Real-world uplift modelling with significance-based uplift trees." White Paper TR-2011-1, Stochastic Solutions (2011): 1-33.
+3. Zhao, Yan, Xiao Fang, and David Simchi-Levi. "Uplift modeling with multiple treatments and general response types." Proceedings of the 2017 SIAM International Conference on Data Mining. Society for Industrial and Applied Mathematics, 2017.
+4. Hansotia, Behram, and Brad Rukstales. "Incremental value modeling." Journal of Interactive Marketing 16.3 (2002): 35-46.
+5. Jannik Rößler, Richard Guse, and Detlef Schoder. "The Best of Two Worlds: Using Recent Advances from Uplift Modeling and Heterogeneous Treatment Effects to Optimize Targeting Policies". International Conference on Information Systems (2022)
+6. Su, Xiaogang, et al. "Subgroup analysis via recursive partitioning." Journal of Machine Learning Research 10.2 (2009).
+7. Su, Xiaogang, et al. "Facilitating score and causal inference trees for large observational studies." Journal of Machine Learning Research 13 (2012): 2955.
+8. Athey, Susan, and Guido Imbens. "Recursive partitioning for heterogeneous causal effects." Proceedings of the National Academy of Sciences 113.27 (2016): 7353-7360.
+9. Künzel, Sören R., et al. "Metalearners for estimating heterogeneous treatment effects using machine learning." Proceedings of the national academy of sciences 116.10 (2019): 4156-4165.
+10. Nie, Xinkun, and Stefan Wager. "Quasi-oracle estimation of heterogeneous treatment effects." arXiv preprint arXiv:1712.04912 (2017).
+11. Bang, Heejung, and James M. Robins. "Doubly robust estimation in missing data and causal inference models." Biometrics 61.4 (2005): 962-973.
+12. Van Der Laan, Mark J., and Daniel Rubin. "Targeted maximum likelihood learning." The international journal of biostatistics 2.1 (2006).
+13. Kennedy, Edward H. "Optimal doubly robust estimation of heterogeneous causal effects." arXiv preprint arXiv:2004.14497 (2020).
+14. Louizos, Christos, et al. "Causal effect inference with deep latent-variable models." arXiv preprint arXiv:1705.08821 (2017).
+15. Shi, Claudia, David M. Blei, and Victor Veitch. "Adapting neural networks for the estimation of treatment effects." 33rd Conference on Neural Information Processing Systems (NeurIPS 2019), 2019.
+16. Zhao, Zhenyu, Yumin Zhang, Totte Harinen, and Mike Yung. "Feature Selection Methods for Uplift Modeling." arXiv preprint arXiv:2005.03447 (2020).
+17. Zhao, Zhenyu, and Totte Harinen. "Uplift modeling for multiple treatments with cost optimization." In 2019 IEEE International Conference on Data Science and Advanced Analytics (DSAA), pp. 422-431. IEEE, 2019.
+
+
+## Related projects
+
+* [uplift](https://cran.r-project.org/web/packages/uplift/index.html): uplift models in R
+* [grf](https://cran.r-project.org/web/packages/grf/index.html): generalized random forests that include heterogeneous treatment effect estimation in R
+* [rlearner](https://github.com/xnie/rlearner): A R package that implements R-Learner
+* [DoWhy](https://github.com/Microsoft/dowhy):  Causal inference in Python based on Judea Pearl's do-calculus
+* [EconML](https://github.com/microsoft/EconML): A Python package that implements heterogeneous treatment effect estimators from econometrics and machine learning methods

causalml/source/SECURITY.md

@@ -0,0 +1,11 @@
+# Security Policy
+
+## Supported Versions
+
+| Version | Supported          |
+| ------- | ------------------ |
+| all     | :white_check_mark: |
+
+## Reporting a Vulnerability
+
+Please report any vulnerabilities to causalml@uber.com

causalml/source/STEERING_COMMITTEE.md

@@ -0,0 +1,14 @@
+# Steering Committee
+
+This document lists the members of the Organization's Steering Committee. Voting members may be added once approved by the Steering Committee as described in the [charter](./CHARTER.md). By adding your name to this list you are agreeing to abide by all Organization polices, including the [charter](./CHARTER.md), the [code of conduct](./CODE_OF_CONDUCT.md), the [trademark policy](./TRADEMARKS.md), and the [antitrust policy](./ANTITRUST.md). If you are serving on the Steering Committee because of your affiliation with another organization (designated below), you represent that you have authority to bind that organization to these policies.
+
+| **NAME** | **Handle** | **Affiliated Organization** |
+| --- | --- | --- |
+| Huigang Chen | @huigangchen | Meta |
+| Totte Harinen | @t-tte | AirBnB |
+| Jeong-Yoon Lee | @jeongyoonlee | Uber |
+| Zhenyu Zhao | @zhenyuz0500 | Tencent |
+
+---
+Adapted from [MVG-0.1-beta](https://github.com/github/MVG/tree/v0.1-beta).
+Made with love by GitHub. Licensed under the [CC-BY 4.0 License](https://creativecommons.org/licenses/by-sa/4.0/).

causalml/source/TRADEMARKS.md

@@ -0,0 +1,44 @@
+## Introduction
+
+This is the Organization's policy for the use of our trademarks. While our work is available under free and open source software licenses, those licenses do not include a license to use our trademarks.
+
+This policy describes how you may use our trademarks.  Our goal is to strike a balance between: 1) our need to ensure that our trademarks remain reliable indicators of the quality software we release; and 2) our community members' desire to be full participants in our Organization.
+
+## Our Trademarks
+
+This policy covers the name of the Organization and each of the Organization's projects, as well as any associated names, trademarks, service marks, logos, mascots, or similar indicators of source or origin (our "Marks").
+
+## In General
+
+Whenever you use our Marks, you must always do so in a way that does not mislead anyone about exactly who is the source of the software. For example, you cannot say you are distributing the "Mark" software when you're distributing a modified version of it because people will believe they are getting the same software that they can get directly from us when they aren't. You also cannot use our Marks on your website in a way that suggests that your website is an official Organization website or that we endorse your website. But, if true, you can say you like the "Mark" software, that you participate in the "Mark" community, that you are providing an unmodified version of the "Mark" software, or that you wrote a book describing how to use the "Mark" software.
+
+This fundamental requirement, that it is always clear to people what they are getting and from whom, is reflected throughout this policy. It should also serve as your guide if you are not sure about how you are using the Marks.
+
+In addition:
+* You may not use or register, in whole or in part, the Marks as part of your own trademark, service mark, domain name, company name, trade name, product name or service name.
+* Trademark law does not allow your use of names or trademarks that are too similar to ours. You therefore may not use an obvious variation of any of our Marks or any phonetic equivalent, foreign language equivalent, takeoff, or abbreviation for a similar or compatible product or service.
+* You agree that any goodwill generated by your use of the Marks and participation in our community inures solely to our collective benefit.
+
+## Distribution of unmodified source code or unmodified executable code we have compiled
+
+When you redistribute an unmodified copy of our software, you are not changing the quality or nature of it. Therefore, you may retain the Marks we have placed on the software to identify your redistribution. This kind of use only applies if you are redistributing an official distribution from this Project that has not been changed in any way.
+
+## Distribution of executable code that you have compiled, or modified code
+
+You may use any word marks, but not any Organization logos, to truthfully describe the origin of the software that you are providing, that is, that the code you are distributing is a modification of our software. You may say, for example, that "this software is derived from the source code for 'Mark' software."
+
+Of course, you can place your own trademarks or logos on versions of the software to which you have made substantive modifications, because by modifying the software, you have become the origin of that exact version. In that case, you should not use our Marks.
+
+However, you may use our Marks for the distribution of code (source or executable) on the condition that any executable is built from the official Project source code and that any modifications are limited to switching on or off features already included in the software, translations into other languages, and incorporating minor bug-fix patches. Use of our Marks on any further modification is not permitted.
+
+## Statements about your software's relation to our software
+
+You may use the word Marks, but not the Organization's logos, to truthfully describe the relationship between your software and ours. Our Mark should be used after a verb or preposition that describes the relationship between your software and ours. So you may say, for example, "Bob's software for the 'Mark' platform" but may not say "Bob's 'Mark' software." Some other examples that may work for you are:
+
+* [Your software] uses "Mark" software
+* [Your software] is powered by "Mark" software
+* [Your software] runs on "Mark" software
+* [Your software] for use with "Mark" software
+* [Your software] for Mark software
+
+These guidelines are based on the [Model Trademark Guidelines](http://www.modeltrademarkguidelines.org), used under a [Creative Commons Attribution 3.0 Unported license](https://creativecommons.org/licenses/by/3.0/deed.en_US)

causalml/source/__init__.py

@@ -0,0 +1,4 @@
+# -*- coding: utf-8 -*-
+"""
+causalml Project Package Initialization File
+"""

causalml/source/causalml/__init__.py

@@ -0,0 +1,10 @@
+__all__ = [
+    "dataset",
+    "features",
+    "feature_selection",
+    "inference",
+    "match",
+    "metrics",
+    "optimize",
+    "propensity",
+]

causalml/source/causalml/dataset/__init__.py

@@ -0,0 +1,16 @@
+from .regression import synthetic_data
+from .regression import simulate_nuisance_and_easy_treatment
+from .regression import simulate_randomized_trial
+from .regression import simulate_easy_propensity_difficult_baseline
+from .regression import simulate_unrelated_treatment_control
+from .regression import simulate_hidden_confounder
+from .classification import make_uplift_classification
+from .classification import make_uplift_classification_logistic
+
+from .synthetic import get_synthetic_preds, get_synthetic_preds_holdout
+from .synthetic import get_synthetic_summary, get_synthetic_summary_holdout
+from .synthetic import scatter_plot_summary, scatter_plot_summary_holdout
+from .synthetic import bar_plot_summary, bar_plot_summary_holdout
+from .synthetic import distr_plot_single_sim
+from .synthetic import scatter_plot_single_sim
+from .synthetic import get_synthetic_auuc

causalml/source/causalml/dataset/classification.py

@@ -0,0 +1,692 @@
+import random
+import numpy as np
+import pandas as pd
+from sklearn.datasets import make_classification
+from scipy.interpolate import UnivariateSpline
+from scipy.optimize import fsolve
+from scipy.special import expit, logit
+
+
+# ------ Define a list of functions for feature transformation
+def _f_linear(x):
+    """
+    Linear transformation (actually identical transformation)
+    """
+    return np.array(x)
+
+
+def _f_quadratic(x):
+    """
+    Quadratic transformation
+    """
+    return np.array(x) * np.array(x)
+
+
+def _f_cubic(x):
+    """
+    Quadratic transformation
+    """
+    return np.array(x) * np.array(x) * np.array(x)
+
+
+def _f_relu(x):
+    """
+    Relu transformation
+    """
+    x = np.array(x)
+    return np.maximum(x, 0)
+
+
+def _f_sin(x):
+    """
+    Sine transformation
+    """
+    return np.sin(np.array(x) * np.pi)
+
+
+def _f_cos(x):
+    """
+    Cosine transformation
+    """
+    return np.cos(np.array(x) * np.pi)
+
+
+# ------ Generating non-linear splines as feature transformation functions
+def _generate_splines(
+    n_functions=10,
+    n_initial_points=10,
+    s=0.01,
+    x_min=-3,
+    x_max=3,
+    y_min=0,
+    y_max=1,
+    random_seed=2019,
+):
+    """
+    Generate a list of spline functions for feature
+    transformation.
+
+    Parameters
+    ----------
+    n_functions : int, optional
+        Number of spline functions to be created.
+    n_initial_points: int, optional
+        Number of initial random points to be placed on a 2D plot to fit a spline.
+    s:  float or None, optional
+        Positive smoothing factor used to choose the number of knots (arg in scipy.interpolate.UnivariateSpline).
+    x_min: int or float, optional
+        The minimum value of the X range.
+    x_max:  int or float, optional
+        The maximum value of the X range.
+    y_min: int or float, optional
+        The minimum value of the Y range.
+    y_max: int or float, optional
+        The maxium value of the Y range.
+    random_seed: int, optional
+        Random seed.
+
+    Returns
+    -------
+    spls: list
+        List of spline functions.
+    """
+    np.random.seed(random_seed)
+    spls = []
+    for i in range(n_functions):
+        x = np.linspace(x_min, x_max, n_initial_points)
+        y = np.random.uniform(y_min, y_max, n_initial_points)
+        spl = UnivariateSpline(x, y, s=s)
+        spls.append(spl)
+    return spls
+
+
+def _standardize(x):
+    """
+    Standardize a vector to be mean 0 and std 1.
+    """
+    return (np.array(x) - np.mean(x)) / np.std(x)
+
+
+def _fixed_transformation(fs, x, f_index=0):
+    """
+    Transform and standardize a vector by a transformation function.
+    If the given index is within the function list f_index < len(fs), then use fs[f_index] as the transformation
+    function. Otherwise, randomly choose a function from the function list.
+
+    Parameters
+    ----------
+    fs : list
+        A collection of functions for transformation.
+    x : list
+        Feature values to be transformed.
+    f_index : int, optional
+        The function index to be used to select a transformation function.
+    """
+    try:
+        y = fs[f_index](x)
+    except IndexError:
+        y = fs[np.asscalar(np.random.choice(len(fs), 1))](x)
+    y = _standardize(y)
+    return y
+
+
+def _random_transformation(fs, x):
+    """
+    Transform and standardize a vector by a function randomly chosen from
+    the function collection.
+
+    Parameters
+    ----------
+    fs : list
+        A collection of functions (splines) for transformation.
+    x : list
+        Feature values to be transformed.
+    """
+    fi = np.random.choice(range(len(fs)), 1)
+    y = fs[fi[0]](x)
+    y = _standardize(y)
+    return y
+
+
+def _softmax(z, p, xb):
+    """
+    Softmax function. This function is used to reversely solve the constant root value in the linear part to make the
+    softmax function output mean to be a given value.
+
+    Parameters
+    ----------
+    z : float
+        Constant value in the linear part.
+    p : float
+        The target output mean value.
+    xb : list
+        An array, with each element as the sum of product of coefficient and feature value
+    """
+    sm_arr = expit(z + np.array(xb))
+    res = p - np.mean(sm_arr)
+    return res
+
+
+# ------ Data generation function (V2) using logistic regression as underlying model
+def make_uplift_classification_logistic(
+    n_samples=10000,
+    treatment_name=["control", "treatment1", "treatment2", "treatment3"],
+    y_name="conversion",
+    n_classification_features=10,
+    n_classification_informative=5,
+    n_classification_redundant=0,
+    n_classification_repeated=0,
+    n_uplift_dict={"treatment1": 2, "treatment2": 2, "treatment3": 3},
+    n_mix_informative_uplift_dict={"treatment1": 1, "treatment2": 1, "treatment3": 0},
+    delta_uplift_dict={"treatment1": 0.02, "treatment2": 0.05, "treatment3": -0.05},
+    positive_class_proportion=0.1,
+    random_seed=20200101,
+    feature_association_list=["linear", "quadratic", "cubic", "relu", "sin", "cos"],
+    random_select_association=True,
+    error_std=0.05,
+):
+    """Generate a synthetic dataset for classification uplift modeling problem.
+
+    Parameters
+    ----------
+    n_samples : int, optional (default=1000)
+        The number of samples to be generated for each treatment group.
+    treatment_name: list, optional (default = ['control','treatment1','treatment2','treatment3'])
+        The list of treatment names. The first element must be 'control' as control group, and the rest are treated as
+        treatment groups.
+    y_name: string, optional (default = 'conversion')
+        The name of the outcome variable to be used as a column in the output dataframe.
+    n_classification_features: int, optional (default = 10)
+        Total number of features for base classification
+    n_classification_informative: int, optional (default = 5)
+        Total number of informative features for base classification
+    n_classification_redundant: int, optional (default = 0)
+        Total number of redundant features for base classification
+    n_classification_repeated: int, optional (default = 0)
+        Total number of repeated features for base classification
+    n_uplift_dict: dictionary, optional (default: {'treatment1': 2, 'treatment2': 2, 'treatment3': 3})
+        Number of features for generating heterogeneous treatment effects for corresponding treatment group.
+        Dictionary of {treatment_key: number_of_features_for_uplift}.
+    n_mix_informative_uplift_dict: dictionary, optional (default: {'treatment1': 1, 'treatment2': 1, 'treatment3': 1})
+        Number of mix features for each treatment. The mix feature is defined as a linear combination
+        of a randomly selected informative classification feature and a randomly selected uplift feature.
+        The mixture is made by a weighted sum (p*feature1 + (1-p)*feature2), where the weight p is drawn from a uniform
+        distribution between 0 and 1.
+    delta_uplift_dict: dictionary, optional (default: {'treatment1': .02, 'treatment2': .05, 'treatment3': -.05})
+        Treatment effect (delta), can be positive or negative.
+        Dictionary of {treatment_key: delta}.
+    positive_class_proportion: float, optional (default = 0.1)
+        The proportion of positive label (1) in the control group, or the mean of outcome variable for control group.
+    random_seed : int, optional (default = 20200101)
+        The random seed to be used in the data generation process.
+    feature_association_list : list, optional (default = ['linear','quadratic','cubic','relu','sin','cos'])
+        List of uplift feature association patterns to the treatment effect. For example, if the feature pattern is
+        'quadratic', then the treatment effect will increase or decrease quadratically with the feature.
+        The values in the list must be one of ('linear','quadratic','cubic','relu','sin','cos'). However, the same
+        value can appear multiple times in the list.
+    random_select_association : boolean, optional (default = True)
+        How the feature patterns are selected from the feature_association_list to be applied in the data generation
+        process. If random_select_association = True, then for every uplift feature, a random feature association
+        pattern is selected from the list. If random_select_association = False, then the feature association pattern
+        is selected from the list in turns to be applied to each feature one by one.
+    error_std : float, optional (default = 0.05)
+        Standard deviation to be used in the error term of the logistic regression. The error is drawn from a normal
+        distribution with mean 0 and standard deviation specified in this argument.
+
+    Returns
+    -------
+    df1 : DataFrame
+        A data frame containing the treatment label, features, and outcome variable.
+    x_name : list
+        The list of feature names generated.
+    """
+
+    # Set means for each experiment group
+    mean_dict = {}
+    mean_dict[treatment_name[0]] = positive_class_proportion
+    for treatment_key_i in treatment_name[1:]:
+        mean_dict[treatment_key_i] = positive_class_proportion
+        if treatment_key_i in delta_uplift_dict:
+            mean_dict[treatment_key_i] += delta_uplift_dict[treatment_key_i]
+
+    # create data frame
+    df1 = pd.DataFrame()
+    n = n_samples
+
+    # set seed
+    np.random.seed(seed=random_seed)
+
+    # define feature association function list ------------------------------------------------#
+    feature_association_pattern_dict = {
+        "linear": _f_linear,
+        "quadratic": _f_quadratic,
+        "cubic": _f_cubic,
+        "relu": _f_relu,
+        "sin": _f_sin,
+        "cos": _f_cos,
+    }
+    f_list = []
+    for fi in feature_association_list:
+        f_list.append(feature_association_pattern_dict[fi])
+
+    # generate treatment key ------------------------------------------------#
+    treatment_list = []
+    for ti in treatment_name:
+        treatment_list += [ti] * n
+    treatment_list = np.random.permutation(treatment_list)
+    df1["treatment_group_key"] = treatment_list
+
+    # feature name list
+    x_name = []
+
+    x_informative_name = []
+    x_informative_transformed = []
+
+    # generate informative features -----------------------------------------#
+    for xi in range(n_classification_informative):
+        # observed feature
+        x = np.random.normal(0, 1, df1.shape[0])
+        x_name_i = "x" + str(len(x_name) + 1) + "_informative"
+        x_name.append(x_name_i)
+        x_informative_name.append(x_name_i)
+        df1[x_name_i] = x
+        # transformed feature that takes effect in the model
+        x_name_i = x_name_i + "_transformed"
+        df1[x_name_i] = _fixed_transformation(f_list, x, xi)
+        x_informative_transformed.append(x_name_i)
+
+    # generate redundant features (linear) ----------------------------------#
+    # linearly combine informative ones
+    for xi in range(n_classification_redundant):
+        nx = (
+            np.random.choice(n_classification_informative, size=1, replace=False)[0] + 1
+        )
+        bx = np.random.normal(0, 1, size=nx)
+        fx = np.random.choice(
+            n_classification_informative, size=nx, replace=False, p=None
+        )
+        x_name_i = "x" + str(len(x_name) + 1) + "_redundant_linear"
+        for xxi in range(nx):
+            x_name_i += "_x" + str(fx[xxi] + 1)
+        x_name.append(x_name_i)
+        x = np.zeros(df1.shape[0])
+        for xxi in range(nx):
+            x += bx[xxi] * df1[x_name[fx[xxi]]]
+        x = _standardize(x)
+        df1[x_name_i] = x
+
+    # generate repeated features --------------------------------------------#
+    # randomly select from informative ones
+    for xi in range(n_classification_repeated):
+        # [N] sklearn.datasets.make_classification may also draw repeated
+        # features from redundant ones
+        fx = np.random.choice(
+            n_classification_informative, size=1, replace=False, p=None
+        )
+        x_name_i = "x" + str(len(x_name) + 1) + "_repeated" + "_x" + str(fx[0] + 1)
+        x_name.append(x_name_i)
+        df1[x_name_i] = df1[x_name[fx[0]]]
+
+    # generate irrelevant features ------------------------------------------#
+    for xi in range(
+        n_classification_features
+        - n_classification_informative
+        - n_classification_redundant
+        - n_classification_repeated
+    ):
+        x_name_i = "x" + str(len(x_name) + 1) + "_irrelevant"
+        x_name.append(x_name_i)
+        df1[x_name_i] = np.random.normal(0, 1, df1.shape[0])
+
+    # Generate uplift features ------------------------------------------------#
+    x_name_uplift_transformed_dict = dict()
+    for treatment_key_i in treatment_name:
+        treatment_index = df1.index[
+            df1["treatment_group_key"] == treatment_key_i
+        ].tolist()
+        if treatment_key_i in n_uplift_dict and n_uplift_dict[treatment_key_i] > 0:
+            x_name_uplift_transformed = []
+            x_name_uplift = []
+            for xi in range(n_uplift_dict[treatment_key_i]):
+                # observed feature
+                x = np.random.normal(0, 1, df1.shape[0])
+                x_name_i = "x" + str(len(x_name) + 1) + "_uplift"
+                x_name.append(x_name_i)
+                x_name_uplift.append(x_name_i)
+                df1[x_name_i] = x
+                # transformed feature that takes effect in the model
+                x_name_i = x_name_i + "_transformed"
+                if random_select_association:
+                    df1[x_name_i] = _fixed_transformation(
+                        f_list, x, random.randint(0, len(f_list) - 1)
+                    )
+                else:
+                    df1[x_name_i] = _fixed_transformation(f_list, x, xi % len(f_list))
+                x_name_uplift_transformed.append(x_name_i)
+            x_name_uplift_transformed_dict[treatment_key_i] = x_name_uplift_transformed
+
+    # generate mixed informative and uplift features
+    for treatment_key_i in treatment_name:
+        if (
+            treatment_key_i in n_mix_informative_uplift_dict
+            and n_mix_informative_uplift_dict[treatment_key_i] > 0
+        ):
+            for xi in range(n_mix_informative_uplift_dict[treatment_key_i]):
+                x_name_i = "x" + str(len(x_name) + 1) + "_mix"
+                x_name.append(x_name_i)
+                p_weight = np.random.uniform(0, 1)
+                df1[x_name_i] = (
+                    p_weight * df1[np.random.choice(x_informative_name)]
+                    + (1 - p_weight) * df1[np.random.choice(x_name_uplift)]
+                )
+
+    # generate conversion probability ------------------------------------------------#
+    # baseline conversion
+    coef_classify = []
+    for ci in range(n_classification_informative):
+        rcoef = [0]
+        while np.abs(rcoef) < 0.1:
+            rcoef = np.random.randn(1) * np.sqrt(1.0 / n_classification_informative)
+        coef_classify.append(rcoef[0])
+    x_classify = df1[x_informative_transformed].values
+    p1 = positive_class_proportion
+    a10 = logit(p1)
+    err = np.random.normal(0, error_std, df1.shape[0])
+    xb_array = (x_classify * coef_classify).sum(axis=1) + err
+    # solve for the constant value so that the output metric mean equal to the function input positive_class_proportion
+    a1 = fsolve(_softmax, a10, args=(p1, xb_array))[0]
+    df1["conversion_prob_linear"] = a1 + xb_array
+    df1["control_conversion_prob_linear"] = df1["conversion_prob_linear"].values
+
+    # uplift conversion
+    for treatment_key_i in treatment_name:
+        if (
+            treatment_key_i in delta_uplift_dict
+            and np.abs(delta_uplift_dict[treatment_key_i]) > 0.0
+        ):
+            treatment_index = df1.index[
+                df1["treatment_group_key"] == treatment_key_i
+            ].tolist()
+            # coefficient
+            coef_uplift = []
+            for ci in range(n_uplift_dict[treatment_key_i]):
+                coef_uplift.append(0.5)
+            x_uplift = df1.loc[
+                :, x_name_uplift_transformed_dict[treatment_key_i]
+            ].values
+            p2 = mean_dict[treatment_key_i]
+            a20 = np.log(p2 / (1.0 - p2)) - a1
+            xb_array = df1["conversion_prob_linear"].values + (
+                x_uplift * coef_uplift
+            ).sum(axis=1)
+            xb_array_treatment = xb_array[treatment_index]
+            a2 = fsolve(_softmax, a20, args=(p2, xb_array_treatment))[0]
+            df1["%s_conversion_prob_linear" % (treatment_key_i)] = a2 + xb_array
+            df1.loc[treatment_index, "conversion_prob_linear"] = df1.loc[
+                treatment_index, "%s_conversion_prob_linear" % (treatment_key_i)
+            ].values
+        else:
+            df1["%s_conversion_prob_linear" % (treatment_key_i)] = df1[
+                "conversion_prob_linear"
+            ].values
+
+    # generate conversion probability and true treatment effect ---------------------------------#
+    df1["conversion_prob"] = 1 / (1 + np.exp(-df1["conversion_prob_linear"].values))
+    df1["control_conversion_prob"] = 1 / (
+        1 + np.exp(-df1["control_conversion_prob_linear"].values)
+    )
+    for treatment_key_i in treatment_name:
+        df1["%s_conversion_prob" % (treatment_key_i)] = 1 / (
+            1 + np.exp(-df1["%s_conversion_prob_linear" % (treatment_key_i)].values)
+        )
+        df1["%s_true_effect" % (treatment_key_i)] = (
+            df1["%s_conversion_prob" % (treatment_key_i)].values
+            - df1["control_conversion_prob"].values
+        )
+
+    # generate Y ------------------------------------------------------------#
+    df1["conversion_prob"] = np.clip(df1["conversion_prob"].values, 0, 1)
+    df1[y_name] = np.random.binomial(1, df1["conversion_prob"].values)
+
+    return df1, x_name
+
+
+def make_uplift_classification(
+    n_samples=1000,
+    treatment_name=["control", "treatment1", "treatment2", "treatment3"],
+    y_name="conversion",
+    n_classification_features=10,
+    n_classification_informative=5,
+    n_classification_redundant=0,
+    n_classification_repeated=0,
+    n_uplift_increase_dict={"treatment1": 2, "treatment2": 2, "treatment3": 2},
+    n_uplift_decrease_dict={"treatment1": 0, "treatment2": 0, "treatment3": 0},
+    delta_uplift_increase_dict={
+        "treatment1": 0.02,
+        "treatment2": 0.05,
+        "treatment3": 0.1,
+    },
+    delta_uplift_decrease_dict={
+        "treatment1": 0.0,
+        "treatment2": 0.0,
+        "treatment3": 0.0,
+    },
+    n_uplift_increase_mix_informative_dict={
+        "treatment1": 1,
+        "treatment2": 1,
+        "treatment3": 1,
+    },
+    n_uplift_decrease_mix_informative_dict={
+        "treatment1": 0,
+        "treatment2": 0,
+        "treatment3": 0,
+    },
+    positive_class_proportion=0.5,
+    random_seed=20190101,
+):
+    """Generate a synthetic dataset for classification uplift modeling problem.
+
+    Parameters
+    ----------
+    n_samples : int, optional (default=1000)
+        The number of samples to be generated for each treatment group.
+    treatment_name: list, optional (default = ['control','treatment1','treatment2','treatment3'])
+        The list of treatment names.
+    y_name: string, optional (default = 'conversion')
+        The name of the outcome variable to be used as a column in the output dataframe.
+    n_classification_features: int, optional (default = 10)
+        Total number of features for base classification
+    n_classification_informative: int, optional (default = 5)
+        Total number of informative features for base classification
+    n_classification_redundant: int, optional (default = 0)
+        Total number of redundant features for base classification
+    n_classification_repeated: int, optional (default = 0)
+        Total number of repeated features for base classification
+    n_uplift_increase_dict: dictionary, optional (default: {'treatment1': 2, 'treatment2': 2, 'treatment3': 2})
+        Number of features for generating positive treatment effects for corresponding treatment group.
+        Dictionary of {treatment_key: number_of_features_for_increase_uplift}.
+    n_uplift_decrease_dict: dictionary, optional (default: {'treatment1': 0, 'treatment2': 0, 'treatment3': 0})
+        Number of features for generating negative treatment effects for corresponding treatment group.
+        Dictionary of {treatment_key: number_of_features_for_increase_uplift}.
+    delta_uplift_increase_dict: dictionary, optional (default: {'treatment1': .02, 'treatment2': .05, 'treatment3': .1})
+        Positive treatment effect created by the positive uplift features on the base classification label.
+        Dictionary of {treatment_key: increase_delta}.
+    delta_uplift_decrease_dict: dictionary, optional (default: {'treatment1': 0., 'treatment2': 0., 'treatment3': 0.})
+        Negative treatment effect created by the negative uplift features on the base classification label.
+        Dictionary of {treatment_key: increase_delta}.
+    n_uplift_increase_mix_informative_dict: dictionary, optional
+        Number of positive mix features for each treatment. The positive mix feature is defined as a linear combination
+        of a randomly selected informative classification feature and a randomly selected positive uplift feature.
+        The linear combination is made by two coefficients sampled from a uniform distribution between -1 and 1.
+        default: {'treatment1': 1, 'treatment2': 1, 'treatment3': 1}
+    n_uplift_decrease_mix_informative_dict: dictionary, optional
+        Number of negative mix features for each treatment. The negative mix feature is defined as a linear combination
+        of a randomly selected informative classification feature and a randomly selected negative uplift feature. The
+        linear combination is made by two coefficients sampled from a uniform distribution between -1 and 1.
+        default: {'treatment1': 0, 'treatment2': 0, 'treatment3': 0}
+    positive_class_proportion: float, optional (default = 0.5)
+        The proportion of positive label (1) in the control group.
+    random_seed : int, optional (default = 20190101)
+        The random seed to be used in the data generation process.
+
+    Returns
+    -------
+    df_res : DataFrame
+        A data frame containing the treatment label, features, and outcome variable.
+    x_name : list
+        The list of feature names generated.
+
+    Notes
+    -----
+    The algorithm for generating the base classification dataset is adapted from the make_classification method in the
+    sklearn package, that uses the algorithm in Guyon [1] designed to generate the "Madelon" dataset.
+
+    References
+    ----------
+    .. [1] I. Guyon, "Design of experiments for the NIPS 2003 variable
+            selection benchmark", 2003.
+    """
+    # set seed
+    np.random.seed(seed=random_seed)
+
+    # create data frame
+    df_res = pd.DataFrame()
+
+    # generate treatment key
+    n_all = n_samples * len(treatment_name)
+    treatment_list = []
+    for ti in treatment_name:
+        treatment_list += [ti] * n_samples
+    treatment_list = np.random.permutation(treatment_list)
+    df_res["treatment_group_key"] = treatment_list
+
+    # generate features and labels
+    X1, Y1 = make_classification(
+        n_samples=n_all,
+        n_features=n_classification_features,
+        n_informative=n_classification_informative,
+        n_redundant=n_classification_redundant,
+        n_repeated=n_classification_repeated,
+        n_clusters_per_class=1,
+        weights=[1 - positive_class_proportion, positive_class_proportion],
+    )
+
+    x_name = []
+    x_informative_name = []
+    for xi in range(n_classification_informative):
+        x_name_i = "x" + str(len(x_name) + 1) + "_informative"
+        x_name.append(x_name_i)
+        x_informative_name.append(x_name_i)
+        df_res[x_name_i] = X1[:, xi]
+    for xi in range(n_classification_redundant):
+        x_name_i = "x" + str(len(x_name) + 1) + "_redundant"
+        x_name.append(x_name_i)
+        df_res[x_name_i] = X1[:, n_classification_informative + xi]
+    for xi in range(n_classification_repeated):
+        x_name_i = "x" + str(len(x_name) + 1) + "_repeated"
+        x_name.append(x_name_i)
+        df_res[x_name_i] = X1[
+            :, n_classification_informative + n_classification_redundant + xi
+        ]
+
+    for xi in range(
+        n_classification_features
+        - n_classification_informative
+        - n_classification_redundant
+        - n_classification_repeated
+    ):
+        x_name_i = "x" + str(len(x_name) + 1) + "_irrelevant"
+        x_name.append(x_name_i)
+        df_res[x_name_i] = np.random.normal(0, 1, n_all)
+
+    # default treatment effects
+    Y = Y1.copy()
+    Y_increase = np.zeros_like(Y1)
+    Y_decrease = np.zeros_like(Y1)
+
+    # generate uplift (positive)
+    for treatment_key_i in treatment_name:
+        treatment_index = df_res.index[
+            df_res["treatment_group_key"] == treatment_key_i
+        ].tolist()
+        if (
+            treatment_key_i in n_uplift_increase_dict
+            and n_uplift_increase_dict[treatment_key_i] > 0
+        ):
+            x_uplift_increase_name = []
+            adjust_class_proportion = (delta_uplift_increase_dict[treatment_key_i]) / (
+                1 - positive_class_proportion
+            )
+            X_increase, Y_increase = make_classification(
+                n_samples=n_all,
+                n_features=n_uplift_increase_dict[treatment_key_i],
+                n_informative=n_uplift_increase_dict[treatment_key_i],
+                n_redundant=0,
+                n_clusters_per_class=1,
+                weights=[1 - adjust_class_proportion, adjust_class_proportion],
+            )
+            for xi in range(n_uplift_increase_dict[treatment_key_i]):
+                x_name_i = "x" + str(len(x_name) + 1) + "_uplift_increase"
+                x_name.append(x_name_i)
+                x_uplift_increase_name.append(x_name_i)
+                df_res[x_name_i] = X_increase[:, xi]
+            Y[treatment_index] = Y[treatment_index] + Y_increase[treatment_index]
+            if n_uplift_increase_mix_informative_dict[treatment_key_i] > 0:
+                for xi in range(
+                    n_uplift_increase_mix_informative_dict[treatment_key_i]
+                ):
+                    x_name_i = "x" + str(len(x_name) + 1) + "_increase_mix"
+                    x_name.append(x_name_i)
+                    df_res[x_name_i] = (
+                        np.random.uniform(-1, 1)
+                        * df_res[np.random.choice(x_informative_name)]
+                        + np.random.uniform(-1, 1)
+                        * df_res[np.random.choice(x_uplift_increase_name)]
+                    )
+
+    # generate uplift (negative)
+    for treatment_key_i in treatment_name:
+        treatment_index = df_res.index[
+            df_res["treatment_group_key"] == treatment_key_i
+        ].tolist()
+        if (
+            treatment_key_i in n_uplift_decrease_dict
+            and n_uplift_decrease_dict[treatment_key_i] > 0
+        ):
+            x_uplift_decrease_name = []
+            adjust_class_proportion = (delta_uplift_decrease_dict[treatment_key_i]) / (
+                1 - positive_class_proportion
+            )
+            X_decrease, Y_decrease = make_classification(
+                n_samples=n_all,
+                n_features=n_uplift_decrease_dict[treatment_key_i],
+                n_informative=n_uplift_decrease_dict[treatment_key_i],
+                n_redundant=0,
+                n_clusters_per_class=1,
+                weights=[1 - adjust_class_proportion, adjust_class_proportion],
+            )
+            for xi in range(n_uplift_decrease_dict[treatment_key_i]):
+                x_name_i = "x" + str(len(x_name) + 1) + "_uplift_decrease"
+                x_name.append(x_name_i)
+                x_uplift_decrease_name.append(x_name_i)
+                df_res[x_name_i] = X_decrease[:, xi]
+            Y[treatment_index] = Y[treatment_index] - Y_decrease[treatment_index]
+            if n_uplift_decrease_mix_informative_dict[treatment_key_i] > 0:
+                for xi in range(
+                    n_uplift_decrease_mix_informative_dict[treatment_key_i]
+                ):
+                    x_name_i = "x" + str(len(x_name) + 1) + "_decrease_mix"
+                    x_name.append(x_name_i)
+                    df_res[x_name_i] = (
+                        np.random.uniform(-1, 1)
+                        * df_res[np.random.choice(x_informative_name)]
+                        + np.random.uniform(-1, 1)
+                        * df_res[np.random.choice(x_uplift_decrease_name)]
+                    )
+
+    # truncate Y
+    Y = np.clip(Y, 0, 1)
+
+    df_res[y_name] = Y
+    df_res["treatment_effect"] = Y - Y1
+    return df_res, x_name

causalml/source/causalml/dataset/regression.py

@@ -0,0 +1,209 @@
+import logging
+
+import numpy as np
+from scipy.special import expit, logit
+
+logger = logging.getLogger("causalml")
+
+
+def synthetic_data(mode=1, n=1000, p=5, sigma=1.0, adj=0.0):
+    """ Synthetic data in Nie X. and Wager S. (2018) 'Quasi-Oracle Estimation of Heterogeneous Treatment Effects'
+    Args:
+        mode (int, optional): mode of the simulation: \
+            1 for difficult nuisance components and an easy treatment effect. \
+            2 for a randomized trial. \
+            3 for an easy propensity and a difficult baseline. \
+            4 for unrelated treatment and control groups. \
+            5 for a hidden confounder biasing treatment.
+        n (int, optional): number of observations
+        p (int optional): number of covariates (>=5)
+        sigma (float): standard deviation of the error term
+        adj (float): adjustment term for the distribution of propensity, e. Higher values shift the distribution to 0.
+                     It does not apply to mode == 2 or 3.
+    Returns:
+        (tuple): Synthetically generated samples with the following outputs:
+            - y ((n,)-array): outcome variable.
+            - X ((n,p)-ndarray): independent variables.
+            - w ((n,)-array): treatment flag with value 0 or 1.
+            - tau ((n,)-array): individual treatment effect.
+            - b ((n,)-array): expected outcome.
+            - e ((n,)-array): propensity of receiving treatment.
+    """
+
+    catalog = {
+        1: simulate_nuisance_and_easy_treatment,
+        2: simulate_randomized_trial,
+        3: simulate_easy_propensity_difficult_baseline,
+        4: simulate_unrelated_treatment_control,
+        5: simulate_hidden_confounder,
+    }
+
+    assert mode in catalog, "Invalid mode {}. Should be one of {}".format(
+        mode, set(catalog)
+    )
+    return catalog[mode](n, p, sigma, adj)
+
+
+def simulate_nuisance_and_easy_treatment(n=1000, p=5, sigma=1.0, adj=0.0):
+    """Synthetic data with a difficult nuisance components and an easy treatment effect
+        From Setup A in Nie X. and Wager S. (2018) 'Quasi-Oracle Estimation of Heterogeneous Treatment Effects'
+    Args:
+        n (int, optional): number of observations
+        p (int optional): number of covariates (>=5)
+        sigma (float): standard deviation of the error term
+        adj (float): adjustment term for the distribution of propensity, e. Higher values shift the distribution to 0.
+    Returns:
+        (tuple): Synthetically generated samples with the following outputs:
+            - y ((n,)-array): outcome variable.
+            - X ((n,p)-ndarray): independent variables.
+            - w ((n,)-array): treatment flag with value 0 or 1.
+            - tau ((n,)-array): individual treatment effect.
+            - b ((n,)-array): expected outcome.
+            - e ((n,)-array): propensity of receiving treatment.
+    """
+
+    X = np.random.uniform(size=n * p).reshape((n, -1))
+    b = (
+        np.sin(np.pi * X[:, 0] * X[:, 1])
+        + 2 * (X[:, 2] - 0.5) ** 2
+        + X[:, 3]
+        + 0.5 * X[:, 4]
+    )
+    eta = 0.1
+    e = np.maximum(
+        np.repeat(eta, n),
+        np.minimum(np.sin(np.pi * X[:, 0] * X[:, 1]), np.repeat(1 - eta, n)),
+    )
+    e = expit(logit(e) - adj)
+    tau = (X[:, 0] + X[:, 1]) / 2
+
+    w = np.random.binomial(1, e, size=n)
+    y = b + (w - 0.5) * tau + sigma * np.random.normal(size=n)
+
+    return y, X, w, tau, b, e
+
+
+def simulate_randomized_trial(n=1000, p=5, sigma=1.0, adj=0.0):
+    """Synthetic data of a randomized trial
+        From Setup B in Nie X. and Wager S. (2018) 'Quasi-Oracle Estimation of Heterogeneous Treatment Effects'
+    Args:
+        n (int, optional): number of observations
+        p (int optional): number of covariates (>=5)
+        sigma (float): standard deviation of the error term
+        adj (float): no effect. added for consistency
+    Returns:
+        (tuple): Synthetically generated samples with the following outputs:
+            - y ((n,)-array): outcome variable.
+            - X ((n,p)-ndarray): independent variables.
+            - w ((n,)-array): treatment flag with value 0 or 1.
+            - tau ((n,)-array): individual treatment effect.
+            - b ((n,)-array): expected outcome.
+            - e ((n,)-array): propensity of receiving treatment.
+    """
+
+    X = np.random.normal(size=n * p).reshape((n, -1))
+    b = np.maximum.reduce([np.repeat(0.0, n), X[:, 0] + X[:, 1], X[:, 2]]) + np.maximum(
+        np.repeat(0.0, n), X[:, 3] + X[:, 4]
+    )
+    e = np.repeat(0.5, n)
+    tau = X[:, 0] + np.log1p(np.exp(X[:, 1]))
+
+    w = np.random.binomial(1, e, size=n)
+    y = b + (w - 0.5) * tau + sigma * np.random.normal(size=n)
+
+    return y, X, w, tau, b, e
+
+
+def simulate_easy_propensity_difficult_baseline(n=1000, p=5, sigma=1.0, adj=0.0):
+    """Synthetic data with easy propensity and a difficult baseline
+        From Setup C in Nie X. and Wager S. (2018) 'Quasi-Oracle Estimation of Heterogeneous Treatment Effects'
+    Args:
+        n (int, optional): number of observations
+        p (int optional): number of covariates (>=3)
+        sigma (float): standard deviation of the error term
+        adj (float): no effect. added for consistency
+    Returns:
+        (tuple): Synthetically generated samples with the following outputs:
+            - y ((n,)-array): outcome variable.
+            - X ((n,p)-ndarray): independent variables.
+            - w ((n,)-array): treatment flag with value 0 or 1.
+            - tau ((n,)-array): individual treatment effect.
+            - b ((n,)-array): expected outcome.
+            - e ((n,)-array): propensity of receiving treatment.
+    """
+
+    X = np.random.normal(size=n * p).reshape((n, -1))
+    b = 2 * np.log1p(np.exp(X[:, 0] + X[:, 1] + X[:, 2]))
+    e = 1 / (1 + np.exp(X[:, 1] + X[:, 2]))
+    tau = np.repeat(1.0, n)
+
+    w = np.random.binomial(1, e, size=n)
+    y = b + (w - 0.5) * tau + sigma * np.random.normal(size=n)
+
+    return y, X, w, tau, b, e
+
+
+def simulate_unrelated_treatment_control(n=1000, p=5, sigma=1.0, adj=0.0):
+    """Synthetic data with unrelated treatment and control groups.
+        From Setup D in Nie X. and Wager S. (2018) 'Quasi-Oracle Estimation of Heterogeneous Treatment Effects'
+    Args:
+        n (int, optional): number of observations
+        p (int optional): number of covariates (>=3)
+        sigma (float): standard deviation of the error term
+        adj (float): adjustment term for the distribution of propensity, e. Higher values shift the distribution to 0.
+    Returns:
+        (tuple): Synthetically generated samples with the following outputs:
+            - y ((n,)-array): outcome variable.
+            - X ((n,p)-ndarray): independent variables.
+            - w ((n,)-array): treatment flag with value 0 or 1.
+            - tau ((n,)-array): individual treatment effect.
+            - b ((n,)-array): expected outcome.
+            - e ((n,)-array): propensity of receiving treatment.
+    """
+
+    X = np.random.normal(size=n * p).reshape((n, -1))
+    b = (
+        np.maximum(np.repeat(0.0, n), X[:, 0] + X[:, 1] + X[:, 2])
+        + np.maximum(np.repeat(0.0, n), X[:, 3] + X[:, 4])
+    ) / 2
+    e = 1 / (1 + np.exp(-X[:, 0]) + np.exp(-X[:, 1]))
+    e = expit(logit(e) - adj)
+    tau = np.maximum(np.repeat(0.0, n), X[:, 0] + X[:, 1] + X[:, 2]) - np.maximum(
+        np.repeat(0.0, n), X[:, 3] + X[:, 4]
+    )
+
+    w = np.random.binomial(1, e, size=n)
+    y = b + (w - 0.5) * tau + sigma * np.random.normal(size=n)
+
+    return y, X, w, tau, b, e
+
+
+def simulate_hidden_confounder(n=10000, p=5, sigma=1.0, adj=0.0):
+    """Synthetic dataset with a hidden confounder biasing treatment.
+        From Louizos et al. (2018) "Causal Effect Inference with Deep Latent-Variable Models"
+    Args:
+        n (int, optional): number of observations
+        p (int optional): number of covariates (>=3)
+        sigma (float): standard deviation of the error term
+        adj (float): no effect. added for consistency
+    Returns:
+        (tuple): Synthetically generated samples with the following outputs:
+            - y ((n,)-array): outcome variable.
+            - X ((n,p)-ndarray): independent variables.
+            - w ((n,)-array): treatment flag with value 0 or 1.
+            - tau ((n,)-array): individual treatment effect.
+            - b ((n,)-array): expected outcome.
+            - e ((n,)-array): propensity of receiving treatment.
+    """
+    z = np.random.binomial(1, 0.5, size=n).astype(np.double)
+    X = np.random.normal(z, 5 * z + 3 * (1 - z), size=(p, n)).T
+    e = 0.75 * z + 0.25 * (1 - z)
+    w = np.random.binomial(1, e)
+    b = expit(3 * (z + 2 * (2 * w - 2)))
+    y = np.random.binomial(1, b)
+
+    # Compute true ite tau for evaluation (via Monte Carlo approximation).
+    t0_t1 = np.array([[0.0], [1.0]])
+    y_t0, y_t1 = expit(3 * (z + 2 * (2 * t0_t1 - 2)))
+    tau = y_t1 - y_t0
+    return y, X, w, tau, b, e

causalml/source/causalml/dataset/semiSynthetic.py

@@ -0,0 +1,1056 @@
+# Synthetic Validation Dataset Generator according to the paper: "Synth-Validation: Selecting the Best Causal Inference Method for a Given Dataset"
+# https://arxiv.org/pdf/1711.00083
+
+import numpy as np
+import pandas as pd
+from scipy.optimize import minimize
+from sklearn.tree import DecisionTreeRegressor
+from sklearn.ensemble import RandomForestRegressor
+from typing import Callable, List, Optional, Union
+from numpy.typing import ArrayLike
+import multiprocessing as mp
+from functools import partial
+from sklearn.linear_model import LinearRegression
+from causalml.inference.meta import BaseXRegressor, BaseTRegressor
+from scipy.special import expit
+
+
+class SemiSynthDataGenerator:
+    def __init__(
+        self,
+        Q: int = 5,
+        gamma: float = 2.0,
+        train_frac: float = 0.8,
+        val_frac: float = 0.1,
+        B: int = 5,
+        maxdepths: List[int] = [1, 2, 3],
+        lambdas: List[float] = np.logspace(-5, 1, num=5).tolist(),
+        M: int = 30,
+        early_stopping_rounds: int = 3,
+        verbose: bool = False,
+        **kwargs,
+    ):
+        self.Q = Q
+        self.gamma = gamma
+        self.train_frac = train_frac
+        self.val_frac = val_frac
+        self.B = B
+        self.maxdepths = maxdepths
+        self.lambdas = lambdas
+        self.M = M
+        self.early_stopping_rounds = early_stopping_rounds
+        self.verbose = verbose
+        self.kwargs = kwargs
+
+    def fit(
+        self,
+        X: pd.DataFrame,
+        w: pd.Series,
+        y: pd.Series,
+        initial_taus: Optional[List[float]] = None,
+    ):
+        self.X = X
+        self.y = y
+        self.w = w
+        np.random.seed(42)
+        if initial_taus is None:
+            # raw_tau
+            raw_tau = y[w == 1].mean() - y[w == 0].mean()
+            # lm_tau
+            X_lm = pd.concat([w, X], axis=1)
+            lm = LinearRegression().fit(X_lm, y)
+            lm_tau = lm.coef_[0]
+            # x_learner_tau
+            x_learner = BaseXRegressor(DecisionTreeRegressor())
+            x_learner_tau = x_learner.estimate_ate(X=X, treatment=w, y=y)[0]
+            # t_learner_tau
+            t_learner = BaseTRegressor(RandomForestRegressor())
+            t_learner_tau = t_learner.estimate_ate(X=X, treatment=w, y=y)[0]
+            initial_taus = [
+                float(raw_tau),
+                float(lm_tau),
+                float(x_learner_tau),
+                float(t_learner_tau),
+            ]
+        else:
+            initial_taus = [float(t) for t in initial_taus]
+        initial_taus_arr = np.array(initial_taus, dtype=float)
+        initial_taus_range = initial_taus_arr.max() - initial_taus_arr.min()
+        initial_taus_median = np.median(initial_taus_arr)
+        taus = np.linspace(
+            initial_taus_median - self.gamma * initial_taus_range,
+            initial_taus_median + self.gamma * initial_taus_range,
+            self.Q,
+        )
+        self.taus = taus
+        self.dgps = []
+        for real_tau in taus:
+            self.dgps.append(
+                miu_cv(
+                    y=np.asarray(self.y),
+                    w=np.asarray(self.w),
+                    X=self.X,
+                    real_tau=real_tau,
+                    train_frac=self.train_frac,
+                    val_frac=self.val_frac,
+                    B=self.B,
+                    max_depths=self.maxdepths,
+                    lambdas=self.lambdas,
+                    M=self.M,
+                    early_stopping_rounds=self.early_stopping_rounds,
+                    verbose=self.verbose,
+                    **self.kwargs,
+                )
+            )
+
+    def generate(self, K: int = 10, n=None) -> List[List[pd.DataFrame]]:
+        if n is None:
+            n = len(self.X)
+        if all((self.y == 0) | (self.y == 1)):
+            binary_y = True
+        else:
+            binary_y = False
+        ctrl_idx = np.where(self.w == 0)[0]
+        trt_idx = np.where(self.w == 1)[0]
+        ctrl_n = int(n * (len(ctrl_idx) / len(self.X)))
+        trt_n = int(n * (len(trt_idx) / len(self.X)))
+        ans = []
+        for q in range(len(self.dgps)):
+            datasets = []
+            dgp_q = self.dgps[q]["final_model"]
+            data_tau = self.X.copy()
+            y0 = dgp_q[0](data_tau)
+            y1 = dgp_q[1](data_tau)
+            if binary_y:
+                y0 = logistic(y0)
+                y1 = logistic(y1)
+            data_tau["w"] = self.w
+            data_tau["tau_i"] = y1 - y0
+            data_tau["y_w"] = np.where(self.w == 1, y1, y0)
+            resid = self.y - data_tau["y_w"]
+            for k in range(K):
+                rng = np.random.default_rng(seed=k)
+                ctrl_idx_qk = rng.choice(ctrl_idx, size=ctrl_n, replace=True)
+                trt_idx_qk = rng.choice(trt_idx, size=trt_n, replace=True)
+                idx = np.concatenate([ctrl_idx_qk, trt_idx_qk])
+                data_qk = data_tau.iloc[idx].copy()
+                if not binary_y:
+                    data_qk["y"] = data_qk["y_w"] + rng.choice(
+                        resid, size=len(data_qk), replace=True
+                    )  # aka observed y
+                else:
+                    data_qk["y"] = data_qk["y_w"].apply(lambda x: rng.binomial(1, x))
+                data_qk = data_qk[["y", "w", "tau_i"] + list(self.X)]
+                datasets.append(data_qk)
+            ans.append(datasets)
+        return ans
+
+
+def continuous_objective(x, Q, a, d):
+    """
+    Compute the continuous objective function for quadratic optimization.
+
+    Parameters:
+    -----------
+    x : np.ndarray
+        The variable vector to optimize over.
+    Q : np.ndarray
+        The quadratic coefficient matrix.
+    a : np.ndarray
+        The linear coefficient vector.
+    d : float
+        The constant term.
+
+    Returns:
+    --------
+    float
+        The value of the objective function: x^T Q x + a^T x + d
+    """
+    return np.dot(x, Q @ x) + np.dot(a, x) + d
+
+
+def deviance(y, pred):
+    """
+    Compute the binomial deviance loss function.
+
+    Parameters:
+    -----------
+    y : np.ndarray
+        True binary outcomes (0 or 1).
+    pred : np.ndarray
+        Predicted logits.
+
+    Returns:
+    --------
+    float
+        The binomial deviance loss: -2 * mean(y * pred - log(1 + exp(pred)))
+    """
+    return -2.0 * np.mean((y * pred) - np.logaddexp(0.0, pred))
+
+
+def logit(x):
+    """
+    Compute the logit (log-odds) transformation.
+
+    Parameters:
+    -----------
+    x : np.ndarray
+        Input values between 0 and 1.
+
+    Returns:
+    --------
+    np.ndarray
+        Logit-transformed values: log(x / (1 - x))
+    """
+    return np.log(x / (1 - x))
+
+
+def logistic(x):
+    """
+    Compute the logistic (sigmoid) transformation.
+
+    Parameters:
+    -----------
+    x : np.ndarray
+        Input values (can be any real number).
+
+    Returns:
+    --------
+    np.ndarray
+        Logistic-transformed values: 1 / (1 + exp(-x))
+    """
+    return 1 / (1 + np.exp(-x))
+
+
+def binary_objective(x, w, y):
+    """
+    Compute the binary objective function for treatment effect estimation.
+
+    Parameters:
+    -----------
+    x : np.ndarray
+        Parameter vector [x0, x1] where x0 is for control group, x1 for treatment group.
+    w : np.ndarray
+        Treatment assignment vector (0 for control, 1 for treatment).
+    y : np.ndarray
+        Binary outcome vector.
+
+    Returns:
+    --------
+    float
+        The binary deviance loss for the given parameters.
+    """
+    pred = np.where(w == 0, x[0], x[1])
+    return deviance(y, pred)
+
+
+def negative_gradient(y, pred):
+    """
+    Compute the negative gradient for binary outcomes.
+
+    Parameters:
+    -----------
+    y : np.ndarray
+        True binary outcomes (0 or 1).
+    pred : np.ndarray
+        Predicted logits.
+
+    Returns:
+    --------
+    np.ndarray
+        The negative gradient: y - losgistic_sigmoid(pred)
+    """
+    return y - expit(pred.ravel())
+
+
+def miu_m(
+    y: ArrayLike,
+    w: ArrayLike,
+    X: Union[pd.DataFrame, ArrayLike],
+    real_tau: Optional[float] = None,
+    miu_m_minus_1: Optional[List[Callable]] = None,
+    val_y: Optional[ArrayLike] = None,
+    val_w: Optional[ArrayLike] = None,
+    val_X: Optional[Union[pd.DataFrame, ArrayLike]] = None,
+    max_depth: Union[int, float] = 3,
+    lambda_: float = 0.0,
+    **tree_args,
+) -> List[Callable]:
+    """
+    Build the m-th iteration of the MIU (Model-based Imputation with Uncertainty) ensemble.
+
+    This function implements a single iteration of the MIU algorithm, which builds
+    treatment-specific models while maintaining a constraint on the treatment effect.
+
+    Parameters:
+    -----------
+    y : ArrayLike
+        Outcome array. Can be continuous or binary (0/1). Will be converted to np.ndarray.
+    w : ArrayLike
+        Treatment assignment array (0 for control, 1 for treatment). Will be converted to np.ndarray.
+    X : Union[pd.DataFrame, ArrayLike]
+        Covariate matrix for training the models. Will be converted to pd.DataFrame.
+    real_tau : Optional[float], default=None
+        The true treatment effect to constrain the model. Required for m=1.
+    miu_m_minus_1 : Optional[List[Callable]], default=None
+        List of two functions [miu_0, miu_1] from the previous iteration.
+        If None, this is the first iteration (m=1).
+    val_y : Optional[ArrayLike], default=None
+        Validation outcome array. Used for constraint calculation if provided.
+    val_w : Optional[ArrayLike], default=None
+        Validation treatment assignment array.
+    val_X : Optional[Union[pd.DataFrame, ArrayLike]], default=None
+        Validation covariate matrix. Used for constraint calculation if provided.
+    max_depth : Union[int, float], default=3
+        Maximum depth of the decision trees used in this iteration.
+    lambda_ : float, default=0.0
+        L2 regularization parameter for the leaf values.
+    **tree_args
+        Additional arguments passed to DecisionTreeRegressor.
+
+    Returns:
+    --------
+    List[Callable]
+        List containing two functions [miu_0_m, miu_1_m]:
+        - miu_0_m: Function that predicts outcomes for control group (w=0)
+        - miu_1_m: Function that predicts outcomes for treatment group (w=1)
+
+    Notes:
+    ------
+    - For m=1, the function fits simple constant models with treatment effect constraint
+    - For m>1, the function fits regression trees to residuals from previous iteration
+    - The treatment effect constraint ensures honest estimation of treatment effects
+    - Binary outcomes use logistic regression, continuous outcomes use linear regression
+    """
+    # Convert inputs to appropriate types
+    y = np.asarray(y)
+    w = np.asarray(w)
+
+    if not isinstance(X, pd.DataFrame):
+        X = pd.DataFrame(X)
+
+    if val_y is not None:
+        val_y = np.asarray(val_y)
+    if val_w is not None:
+        val_w = np.asarray(val_w)
+    if val_X is not None and not isinstance(val_X, pd.DataFrame):
+        val_X = pd.DataFrame(val_X)
+
+    if all((y == 0) | (y == 1)):
+        binary_y = True
+    else:
+        binary_y = False
+    if miu_m_minus_1 is None:
+        # m == 1
+        if real_tau is None:
+            raise ValueError("For m=1 (first call to miu_m) real_tau must be supplied")
+        x0 = np.zeros(2)
+        if not binary_y:
+            n0 = (w == 0).sum()
+            n1 = (w == 1).sum()
+            Q = np.array([[n0, 0], [0, n1]])
+            a = np.array(
+                [
+                    -2 * y[w == 0].sum(),
+                    -2 * y[w == 1].sum(),
+                ]
+            )
+            d = (y**2).sum()
+
+            constraints = {"type": "eq", "fun": lambda x: x[1] - x[0] - real_tau}
+            res = minimize(
+                fun=continuous_objective,
+                x0=x0,
+                args=(Q, a, d),
+                constraints=constraints,
+                method="SLSQP",
+            )
+        else:
+            constraints = {
+                "type": "eq",
+                "fun": lambda x: logistic(x[1]) - logistic(x[0]) - real_tau,
+            }
+            res = minimize(
+                fun=binary_objective,
+                x0=x0,
+                args=(w, y),
+                constraints=constraints,
+                method="SLSQP",
+            )
+
+        res01, res11 = res.x[0], res.x[1]
+
+        def miu_01(x):
+            return np.repeat(res01, len(x))
+
+        def miu_11(x):
+            return np.repeat(res11, len(x))
+
+        return [miu_01, miu_11]
+    else:
+        # m > 1
+        miu_0_m_minus_1, miu_1_m_minus_1 = miu_m_minus_1
+        # Predict y_hat using previous miu functions
+        y_1 = miu_1_m_minus_1(X)
+        y_0 = miu_0_m_minus_1(X)
+        y_hat = np.where(w == 1, y_1, y_0)
+        if not binary_y:
+            resid = y - y_hat
+        else:
+            resid = negative_gradient(y, y_hat)
+        treat = w == 1
+        # Fit regression trees to residuals
+        b_0m = DecisionTreeRegressor(max_depth=max_depth, **tree_args, random_state=42)
+        b_0m.fit(X.loc[~treat], resid[~treat])
+        b_1m = DecisionTreeRegressor(max_depth=max_depth, **tree_args, random_state=42)
+        b_1m.fit(X.loc[treat], resid[treat])
+        # Predict leaf node for each sample
+        R0 = b_0m.apply(X)
+        R1 = b_1m.apply(X)
+        if not binary_y:
+            # Group sizes and residuals
+            resid0 = (
+                pd.Series(resid)
+                .groupby(R0)
+                .agg(["count", "sum"])
+                .reset_index()
+                .rename(columns={"index": "leaf_node"})
+            )
+            resid1 = (
+                pd.Series(resid)
+                .groupby(R1)
+                .agg(["count", "sum"])
+                .reset_index()
+                .rename(columns={"index": "leaf_node"})
+            )
+            num_params = len(resid0) + len(resid1)
+            Q = np.diag(
+                np.concatenate([resid0["count"].to_numpy(), resid1["count"].to_numpy()])
+                * lambda_
+            )
+            a = -2 * np.concatenate(
+                [resid0["sum"].to_numpy(), resid1["sum"].to_numpy()]
+            )
+            d = (resid**2).sum()
+            if val_X is not None and val_y is not None and val_w is not None:
+                # Optionally add validation data
+                X_full = pd.concat([X, val_X], ignore_index=True, axis=0)
+                R0 = b_0m.apply(X_full)
+                R1 = b_1m.apply(X_full)
+                # Making the constraint apply over the entire dataset - this is still honest
+                resid0 = (
+                    pd.Series(R0)
+                    .groupby(R0)
+                    .agg(["count"])
+                    .reset_index()
+                    .rename(columns={"index": "leaf_node"})
+                )
+                resid1 = (
+                    pd.Series(R1)
+                    .groupby(R1)
+                    .agg(["count"])
+                    .reset_index()
+                    .rename(columns={"index": "leaf_node"})
+                )
+
+            constraints = {
+                "type": "eq",
+                "fun": lambda x: np.dot(resid1["count"].to_numpy(), x[len(resid0) :])
+                - np.dot(resid0["count"].to_numpy(), x[: len(resid0)]),
+            }
+            x0 = np.zeros(num_params)
+            res = minimize(
+                continuous_objective,
+                x0,
+                args=(Q, a, d),
+                constraints=constraints,
+                method="SLSQP",
+            )
+        else:
+            resid0 = pd.DataFrame({"leaf_node": np.unique(R0)})
+            resid1 = pd.DataFrame({"leaf_node": np.unique(R1)})
+            x0 = np.zeros(len(resid0) + len(resid1))
+
+            def binary_objective_m(x, y, w, R0, R1, lambda_):
+                loss = []
+                r0 = np.unique(R0)
+                r1 = np.unique(R1)
+                ctrl_nodes = len(r0)
+                for i in range(len(x)):
+                    if i < ctrl_nodes - 1:
+                        idx = (R0 == r0[i]) & (w == 0)
+                    else:
+                        idx = (R1 == r1[i - ctrl_nodes]) & (w == 1)
+                    pred = np.full(sum(idx), x[i])
+                    loss.append(deviance(y[idx], pred) + sum(idx) * lambda_ * x[i] ** 2)
+                return np.array(loss).sum()
+
+            if val_X is not None and val_y is not None and val_w is not None:
+                # Optionally add validation data
+                X_full = pd.concat([X, val_X], ignore_index=True, axis=0)
+                R0_constraint = b_0m.apply(X_full)
+                R1_constraint = b_1m.apply(X_full)
+                prev0 = miu_0_m_minus_1(X_full)
+                prev1 = miu_1_m_minus_1(X_full)
+                # Making the constraint apply over the entire dataset - this is still honest
+            else:
+                R0_constraint = R0
+                R1_constraint = R1
+                prev0 = y_0
+                prev1 = y_1
+
+            real_tau = (logistic(prev1) - logistic(prev0)).mean()
+
+            def con_m(x, R0_constraint, R1_constraint, prev0, prev1):
+                group_sum = []
+                r0 = np.unique(R0_constraint)
+                r1 = np.unique(R1_constraint)
+                ctrl_nodes = len(r0)
+                for i in range(len(x)):
+                    if i < ctrl_nodes:
+                        idx = R0_constraint == r0[i]
+                        group_sum.append(
+                            (logistic(prev0[idx] + x[i])).sum()
+                            / len(R0_constraint)
+                            * -1
+                        )
+                    else:
+                        idx = R1_constraint == r1[i - ctrl_nodes]
+                        group_sum.append(
+                            (logistic(prev1[idx] + x[i])).sum() / len(R0_constraint)
+                        )
+                return np.array(group_sum).sum() - real_tau
+
+            constraints = {
+                "type": "eq",
+                "fun": lambda x: con_m(x, R0_constraint, R1_constraint, prev0, prev1),
+            }
+
+            res = minimize(
+                fun=binary_objective_m,
+                x0=x0,
+                args=(y, w, R0, R1, lambda_),
+                constraints=constraints,
+                method="SLSQP",
+            )
+
+        # Assign fitted values to leaves
+        resid0["leaf_value"] = res.x[: len(resid0)]
+        resid1["leaf_value"] = res.x[len(resid0) :]
+
+        def miu_0m(x):
+            prev = miu_0_m_minus_1(x)
+            leaves = pd.DataFrame({"leaf_node": b_0m.apply(x)})
+            return (
+                prev
+                + leaves.merge(resid0, on="leaf_node", how="left")[
+                    "leaf_value"
+                ].to_numpy()
+            )
+
+        def miu_1m(x):
+            prev = miu_1_m_minus_1(x)
+            leaves = pd.DataFrame({"leaf_node": b_1m.apply(x)})
+            return (
+                prev
+                + leaves.merge(resid1, on="leaf_node", how="left")[
+                    "leaf_value"
+                ].to_numpy()
+            )
+
+        return [miu_0m, miu_1m]
+
+
+def miu(
+    y: ArrayLike,
+    w: ArrayLike,
+    X: Union[pd.DataFrame, ArrayLike],
+    real_tau: float,
+    val_y: Optional[ArrayLike] = None,
+    val_w: Optional[ArrayLike] = None,
+    val_X: Optional[Union[pd.DataFrame, ArrayLike]] = None,
+    max_depth: Union[int, float] = 3,
+    lambda_: float = 0.0,
+    M: int = 10,
+    early_stopping_rounds: Union[int, float] = float("inf"),
+    verbose: bool = False,
+    **tree_args,
+) -> dict:
+    """
+    Train an ensemble of M MIU models and return the best one.
+
+    This function implements the complete MIU (Model-based Imputation with Uncertainty)
+    algorithm, which builds an ensemble of treatment-specific models while maintaining
+    constraints on the treatment effect for honest estimation.
+
+    Parameters:
+    -----------
+    y : ArrayLike
+        Outcome array. Can be continuous or binary (0/1). Will be converted to np.ndarray.
+    w : ArrayLike
+        Treatment assignment array (0 for control, 1 for treatment). Will be converted to np.ndarray.
+    X : Union[pd.DataFrame, ArrayLike]
+        Covariate matrix for training the models. Will be converted to pd.DataFrame.
+    real_tau : float
+        The true treatment effect to constrain the model. This is used to ensure
+        honest estimation of treatment effects.
+    val_y : Optional[ArrayLike], default=None
+        Validation outcome array. Used for early stopping and model selection.
+    val_w : Optional[ArrayLike], default=None
+        Validation treatment assignment array.
+    val_X : Optional[Union[pd.DataFrame, ArrayLike]], default=None
+        Validation covariate matrix. Used for early stopping and model selection.
+    max_depth : Union[int, float], default=3
+        Maximum depth of the decision trees used in each iteration.
+    lambda_ : float, default=0.0
+        L2 regularization parameter for the leaf values in each iteration.
+    M : int, default=10
+        Maximum number of ensemble iterations to perform.
+    early_stopping_rounds : Union[int, float], default=float('inf')
+        Number of rounds without improvement before stopping early.
+        If val_X is None, this must be float('inf').
+    verbose : bool, default=False
+        Whether to print progress information during training.
+    **tree_args
+        Additional arguments passed to DecisionTreeRegressor in each iteration.
+
+    Returns:
+    --------
+    dict
+        Dictionary containing:
+        - 'best_model': List[Callable] - The best ensemble model [miu_0, miu_1]
+        - 'loss': np.ndarray - Array of validation losses for each iteration
+        - 'best_model_m': int - The iteration number of the best model
+
+    Notes:
+    ------
+    - The algorithm builds an ensemble by iteratively fitting models to residuals
+    - Each iteration maintains the treatment effect constraint using real_tau
+    - Early stopping is based on validation loss if validation data is provided
+    - The best model is selected based on validation loss or training loss
+    - Binary outcomes use logistic regression, continuous outcomes use linear regression
+    """
+    if val_X is None and not np.isinf(early_stopping_rounds):
+        raise ValueError("If val_X is None then early_stopping_rounds must be Inf")
+
+    # Convert inputs to appropriate types
+    y = np.asarray(y)
+    w = np.asarray(w)
+
+    if not isinstance(X, pd.DataFrame):
+        X = pd.DataFrame(X)
+
+    if val_y is not None:
+        val_y = np.asarray(val_y)
+    if val_w is not None:
+        val_w = np.asarray(val_w)
+    if val_X is not None and not isinstance(val_X, pd.DataFrame):
+        val_X = pd.DataFrame(val_X)
+
+    if all((y == 0) | (y == 1)):
+        binary_y = True
+    else:
+        binary_y = False
+
+    loss = np.full(M, np.nan)
+    best_model_ind = 0
+    best_model = None
+
+    for i in range(M):
+        if i == 0:
+            ans = miu_m(
+                y=y,
+                w=w,
+                X=X,
+                real_tau=real_tau,
+                val_X=val_X,
+                val_y=val_y,
+                val_w=val_w,
+                max_depth=max_depth,
+                lambda_=lambda_,
+                **tree_args,
+            )
+            best_model = ans
+        else:
+            ans = miu_m(
+                y=y,
+                w=w,
+                X=X,
+                miu_m_minus_1=ans,
+                val_X=val_X,
+                val_y=val_y,
+                val_w=val_w,
+                max_depth=max_depth,
+                lambda_=lambda_,
+                **tree_args,
+            )
+
+        if val_X is None:
+            # Use training data for loss calculation
+            pred = np.where(w == 1, ans[1](X), ans[0](X))
+            if not binary_y:
+                loss[i] = np.mean((y - pred) ** 2)
+            else:
+                loss[i] = deviance(y, pred)
+        else:
+            # Use validation data for loss calculation
+            pred = np.where(val_w == 1, ans[1](val_X), ans[0](val_X))
+            if not binary_y:
+                loss[i] = np.mean((val_y - pred) ** 2)
+            else:
+                loss[i] = deviance(val_y, pred)
+
+        if np.nanargmin(loss) != best_model_ind:
+            best_model_ind = np.nanargmin(loss)
+            best_model = ans
+        elif i - np.nanargmin(loss) > early_stopping_rounds:
+            if verbose:
+                print(
+                    f"Best tree: {best_model_ind + 1}, best tree loss: {loss[best_model_ind]}"
+                )
+            return {
+                "best_model": best_model,
+                "loss": loss,
+                "best_model_m": best_model_ind + 1,
+            }
+
+    if verbose:
+        print(
+            f"Best tree: {best_model_ind + 1}, best tree loss: {loss[best_model_ind]}"
+        )
+    return {"best_model": best_model, "loss": loss, "best_model_m": best_model_ind + 1}
+
+
+def miu_cv(
+    y: ArrayLike,
+    w: ArrayLike,
+    X: Union[pd.DataFrame, ArrayLike],
+    real_tau: float,
+    train_frac: float = 0.8,
+    val_frac: float = 0.1,
+    B: int = 5,
+    max_depths: List[int] = [1, 3, 5],
+    lambdas: List[float] = np.logspace(
+        -5, 1, num=5
+    ).tolist(),  # range of lambdas is like in glmnet
+    M: int = 30,
+    early_stopping_rounds: Union[int, float] = float("inf"),
+    verbose: bool = False,
+    n_jobs: int = -1,
+    **tree_args,
+) -> dict:
+    """
+    Perform cross-validation to find optimal hyperparameters for the MIU model.
+
+    This function performs bootstrap-based cross-validation to tune the hyperparameters
+    of the MIU algorithm, including max_depth and lambda regularization parameter.
+
+    Parameters:
+    -----------
+    y : ArrayLike
+        Outcome array. Can be continuous or binary (0/1). Will be converted to np.ndarray.
+    w : ArrayLike
+        Treatment assignment array (0 for control, 1 for treatment). Will be converted to np.ndarray.
+    X : Union[pd.DataFrame, ArrayLike]
+        Covariate matrix for training the models. Will be converted to pd.DataFrame.
+    real_tau : float
+        The true treatment effect to constrain the model.
+    train_frac : float, default=0.8
+        Fraction of data to use for training in each bootstrap iteration.
+    val_frac : float, default=0.1
+        Fraction of training data to use for validation. If 0, no validation is performed.
+    B : int, default=5
+        Number of bootstrap iterations for cross-validation.
+    max_depths : List[int], default=[1, 3, 5]
+        List of maximum tree depths to try during hyperparameter tuning.
+    lambdas : List[float], default=np.logspace(-5, 1, 5).tolist()
+        List of L2 regularization parameters to try during hyperparameter tuning.
+        Range is similar to glmnet: from 1e-5 to 10.
+    M : int, default=30
+        Maximum number of ensemble iterations for each model.
+    early_stopping_rounds : Union[int, float], default=float('inf')
+        Number of rounds without improvement before stopping early.
+        If val_frac is 0, this must be float('inf').
+    verbose : bool, default=False
+        Whether to print progress information during cross-validation.
+    n_jobs : int, default=-1
+        Number of jobs to run in parallel. -1 means using all processors - 1.
+    **tree_args
+        Additional arguments passed to DecisionTreeRegressor.
+
+    Returns:
+    --------
+    dict
+        Dictionary containing:
+        - 'final_model': List[Callable] - The best ensemble model trained on full data
+        - 'params_loss': pd.DataFrame - Cross-validation results for all parameter combinations
+
+    Notes:
+    ------
+    - Uses stratified bootstrap sampling to maintain treatment group proportions
+    - Performs parallel processing across parameter combinations and bootstrap iterations
+    - Selects best parameters based on mean test loss across bootstrap iterations
+    - Final model is trained on the full dataset using the best parameters
+    - The params_loss DataFrame contains loss, r_sq, and m for each parameter combination
+    """
+    # Convert inputs to appropriate types
+    y = np.asarray(y)
+    w = np.asarray(w)
+
+    if not isinstance(X, pd.DataFrame):
+        X = pd.DataFrame(X)
+    # Create parameter grid
+    param_combinations = []
+    for max_depth in max_depths:
+        for lambda_ in lambdas:
+            param_combinations.append({"max_depth": max_depth, "lambda_": lambda_})
+
+    params_loss = pd.DataFrame(param_combinations)
+    params_loss["loss"] = np.nan
+    params_loss["r_sq"] = np.nan
+    params_loss["m"] = np.nan
+    params_loss = params_loss.merge(pd.DataFrame({"b": range(B)}), how="cross")
+
+    # Set number of jobs
+    if n_jobs == -1:
+        n_jobs = mp.cpu_count() - 1  # don't freeze the computer
+
+    # Run rows in parallel
+    if n_jobs > 1:
+        with mp.Pool(processes=n_jobs) as pool:
+            params_loss = pool.map(
+                partial(
+                    miu_row,
+                    y=y,
+                    w=w,
+                    X=X,
+                    real_tau=real_tau,
+                    train_frac=train_frac,
+                    val_frac=val_frac,
+                    M=M,
+                    early_stopping_rounds=early_stopping_rounds,
+                    verbose=False,
+                    **tree_args,
+                ),
+                [row for _, row in params_loss.iterrows()],
+            )
+    else:
+        params_loss = [
+            miu_row(
+                row,
+                y=y,
+                w=w,
+                X=X,
+                real_tau=real_tau,
+                train_frac=train_frac,
+                val_frac=val_frac,
+                M=M,
+                early_stopping_rounds=early_stopping_rounds,
+                verbose=False,
+                **tree_args,
+            )
+            for _, row in params_loss.iterrows()
+        ]
+
+    # Aggregate results back into params_loss DataFrame
+    params_loss = pd.concat(params_loss, axis=0)
+    params_loss = (
+        params_loss.groupby(["max_depth", "lambda_"])
+        .agg({"loss": "mean", "r_sq": "mean", "m": "mean"})
+        .reset_index()
+        .assign(m=lambda x: x["m"].astype(int))
+    )
+
+    # Find best parameters
+    best_idx = np.argmin(params_loss["loss"])
+    params_loss["best_params"] = params_loss.index == best_idx
+    best_params = params_loss.iloc[best_idx]
+
+    if verbose:
+        print(
+            f"Best params: max_depth - {best_params['max_depth']}, "
+            f"lambda - {best_params['lambda_']}, m - {best_params['m']}"
+        )
+
+    # Train final model with best parameters
+    final_model = miu(
+        y=y,
+        w=w,
+        X=X,
+        real_tau=real_tau,
+        val_y=None,
+        val_w=None,
+        val_X=None,
+        max_depth=int(best_params["max_depth"]),
+        lambda_=best_params["lambda_"],
+        M=int(best_params["m"]),
+        early_stopping_rounds=float("inf"),
+        verbose=False,
+        **tree_args,
+    )
+
+    return {
+        "final_model": final_model["best_model"],
+        "params_loss": params_loss,
+    }
+
+
+def miu_row(
+    row: pd.Series,
+    y: ArrayLike,
+    w: ArrayLike,
+    X: Union[pd.DataFrame, ArrayLike],
+    real_tau: float,
+    train_frac: float = 0.8,
+    val_frac: float = 0.1,
+    M: int = 30,
+    early_stopping_rounds: Union[int, float] = float("inf"),
+    verbose: bool = False,
+    **tree_args,
+) -> pd.Series:
+    """
+    Train a single MIU model for a specific parameter combination and bootstrap iteration.
+
+    This function is designed to be used in parallel processing for cross-validation.
+    It trains a MIU model with specific hyperparameters on a bootstrap sample and
+    evaluates it on the out-of-bag test set.
+
+    Parameters:
+    -----------
+    row : pd.Series
+        A pandas Series containing the parameter combination to evaluate.
+        Must contain 'max_depth' and 'lambda_' keys.
+    y : ArrayLike
+        Outcome array. Can be continuous or binary (0/1). Will be converted to np.ndarray.
+    w : ArrayLike
+        Treatment assignment array (0 for control, 1 for treatment). Will be converted to np.ndarray.
+    X : Union[pd.DataFrame, ArrayLike]
+        Covariate matrix for training the models. Will be converted to pd.DataFrame.
+    real_tau : float
+        The true treatment effect to constrain the model.
+    train_frac : float, default=0.8
+        Fraction of data to use for training.
+    val_frac : float, default=0.1
+        Fraction of training data to use for validation. If 0, no validation is performed.
+    M : int, default=30
+        Maximum number of ensemble iterations for the model.
+    early_stopping_rounds : Union[int, float], default=float('inf')
+        Number of rounds without improvement before stopping early.
+        If val_frac is 0, this must be float('inf').
+    verbose : bool, default=False
+        Whether to print progress information during training.
+    **tree_args
+        Additional arguments passed to DecisionTreeRegressor.
+
+    Returns:
+    --------
+    pd.Series
+        A pandas Series containing the original parameters plus:
+        - 'm': int - The number of iterations in the best model
+        - 'loss': float - The test loss (MSE for continuous, deviance for binary)
+        - 'r_sq': float - The R-squared value on the test set
+
+    Notes:
+    ------
+    - Performs stratified bootstrap sampling to maintain treatment group proportions
+    - Uses the parameters from 'row' to train the MIU model
+    - Evaluates the model on the out-of-bag test set
+    - Returns results as a pandas Series for easy aggregation
+    - Designed for parallel processing in cross-validation
+    """
+    # Convert inputs to appropriate types
+    y = np.asarray(y)
+    w = np.asarray(w)
+
+    if not isinstance(X, pd.DataFrame):
+        X = pd.DataFrame(X)
+
+    if all((y == 0) | (y == 1)):
+        binary_y = True
+    else:
+        binary_y = False
+
+    # Stratified split data into train/test based on treatment w
+    n_samples = len(y)
+    w_0_indices = np.where(w == 0)[0]
+    w_1_indices = np.where(w == 1)[0]
+
+    # Calculate split sizes for each treatment group
+    train_size_0 = int(train_frac * len(w_0_indices))
+    train_size_1 = int(train_frac * len(w_1_indices))
+
+    # Randomly select train indices for each treatment group
+    train_indices_0 = np.random.choice(w_0_indices, size=train_size_0, replace=False)
+    train_indices_1 = np.random.choice(w_1_indices, size=train_size_1, replace=False)
+    train_indices = np.concatenate([train_indices_0, train_indices_1])
+
+    # Remaining indices go to test
+    test_indices = np.setdiff1d(np.arange(n_samples), train_indices)
+
+    if val_frac > 0:
+        # Further stratified split train into train/validation
+        val_size_0 = int(val_frac * len(train_indices_0))
+        val_size_1 = int(val_frac * len(train_indices_1))
+
+        val_indices_0 = np.random.choice(
+            train_indices_0, size=val_size_0, replace=False
+        )
+        val_indices_1 = np.random.choice(
+            train_indices_1, size=val_size_1, replace=False
+        )
+        val_indices = np.concatenate([val_indices_0, val_indices_1])
+
+        # Remove validation indices from train
+        train_indices = np.setdiff1d(train_indices, val_indices)
+
+        val_X = X.iloc[val_indices]
+        val_y = y[val_indices]
+        val_w = w[val_indices]
+    else:
+        val_X = None
+        val_y = None
+        val_w = None
+
+    train_X = X.iloc[train_indices]
+    train_y = y[train_indices]
+    train_w = w[train_indices]
+    test_X = X.iloc[test_indices]
+    test_y = y[test_indices]
+    test_w = w[test_indices]
+    miu_row = miu(
+        y=train_y,
+        w=train_w,
+        X=train_X,
+        real_tau=real_tau,
+        val_y=val_y,
+        val_w=val_w,
+        val_X=val_X,
+        max_depth=int(row["max_depth"]),
+        lambda_=float(row["lambda_"]),
+        M=M,
+        early_stopping_rounds=early_stopping_rounds,
+        verbose=verbose,
+        **tree_args,
+    )
+
+    # Make predictions on test set
+    pred = np.where(
+        test_w == 1, miu_row["best_model"][1](test_X), miu_row["best_model"][0](test_X)
+    )
+
+    # Create result dict
+    row["m"] = miu_row["best_model_m"]
+
+    if not binary_y:
+        row["loss"] = np.mean((test_y - pred) ** 2)
+        row["r_sq"] = 1 - row["loss"] / np.mean((test_y - np.mean(test_y)) ** 2)
+    else:
+        row["loss"] = deviance(test_y, pred)
+        baseline_pred = np.where(
+            test_w == 1,
+            logit(np.mean(test_y[test_w == 1])),
+            logit(np.mean(test_y[test_w == 0])),
+        )
+        row["r_sq"] = 1 - row["loss"] / deviance(test_y, baseline_pred)
+
+    return row.to_frame().T

causalml/source/causalml/dataset/synthetic.py

@@ -0,0 +1,655 @@
+from matplotlib import pyplot as plt
+import numpy as np
+import pandas as pd
+from sklearn.metrics import mean_squared_error as mse
+from sklearn.metrics import auc
+from sklearn.model_selection import train_test_split
+from sklearn.linear_model import LinearRegression
+from xgboost import XGBRegressor
+from scipy.stats import entropy
+import warnings
+
+from causalml.inference.meta import (
+    BaseXRegressor,
+    BaseRRegressor,
+    BaseSRegressor,
+    BaseTRegressor,
+)
+from causalml.inference.tree.causal.causaltree import CausalTreeRegressor
+from causalml.propensity import ElasticNetPropensityModel
+from causalml.metrics import plot_gain, get_cumgain
+
+plt.style.use("fivethirtyeight")
+warnings.filterwarnings("ignore")
+
+KEY_GENERATED_DATA = "generated_data"
+KEY_ACTUAL = "Actuals"
+
+RANDOM_SEED = 42
+
+
+def get_synthetic_preds(synthetic_data_func, n=1000, estimators={}):
+    """Generate predictions for synthetic data using specified function (single simulation)
+
+    Args:
+        synthetic_data_func (function): synthetic data generation function
+        n (int, optional): number of samples
+        estimators (dict of object): dict of names and objects of treatment effect estimators
+
+    Returns:
+        (dict): dict of the actual and estimates of treatment effects
+    """
+    y, X, w, tau, b, e = synthetic_data_func(n=n)
+
+    preds_dict = {}
+    preds_dict[KEY_ACTUAL] = tau
+    preds_dict[KEY_GENERATED_DATA] = {
+        "y": y,
+        "X": X,
+        "w": w,
+        "tau": tau,
+        "b": b,
+        "e": e,
+    }
+
+    # Predict p_hat because e would not be directly observed in real-life
+    p_model = ElasticNetPropensityModel()
+    p_hat = p_model.fit_predict(X, w)
+
+    if estimators:
+        for name, learner in estimators.items():
+            try:
+                preds_dict[name] = learner.fit_predict(
+                    X=X, treatment=w, y=y, p=p_hat
+                ).flatten()
+            except TypeError:
+                preds_dict[name] = learner.fit_predict(X=X, treatment=w, y=y).flatten()
+    else:
+        for base_learner, label_l in zip(
+            [BaseSRegressor, BaseTRegressor, BaseXRegressor, BaseRRegressor],
+            ["S", "T", "X", "R"],
+        ):
+            for model, label_m in zip([LinearRegression, XGBRegressor], ["LR", "XGB"]):
+                learner = base_learner(model())
+                model_name = "{} Learner ({})".format(label_l, label_m)
+                try:
+                    preds_dict[model_name] = learner.fit_predict(
+                        X=X, treatment=w, y=y, p=p_hat
+                    ).flatten()
+                except TypeError:
+                    preds_dict[model_name] = learner.fit_predict(
+                        X=X, treatment=w, y=y
+                    ).flatten()
+
+        learner = CausalTreeRegressor(random_state=RANDOM_SEED)
+        preds_dict["Causal Tree"] = learner.fit_predict(X=X, treatment=w, y=y).flatten()
+
+    return preds_dict
+
+
+def get_synthetic_summary(synthetic_data_func, n=1000, k=1, estimators={}):
+    """Generate a summary for predictions on synthetic data using specified function
+
+    Args:
+        synthetic_data_func (function): synthetic data generation function
+        n (int, optional): number of samples per simulation
+        k (int, optional): number of simulations
+    """
+    summaries = []
+
+    for i in range(k):
+        synthetic_preds = get_synthetic_preds(
+            synthetic_data_func, n=n, estimators=estimators
+        )
+        actuals = synthetic_preds[KEY_ACTUAL]
+        synthetic_summary = pd.DataFrame(
+            {
+                label: [preds.mean(), mse(preds, actuals)]
+                for label, preds in synthetic_preds.items()
+                if label != KEY_GENERATED_DATA
+            },
+            index=["ATE", "MSE"],
+        ).T
+
+        synthetic_summary["Abs % Error of ATE"] = np.abs(
+            (synthetic_summary["ATE"] / synthetic_summary.loc[KEY_ACTUAL, "ATE"]) - 1
+        )
+
+        for label in synthetic_summary.index:
+            stacked_values = np.hstack((synthetic_preds[label], actuals))
+            stacked_low = np.percentile(stacked_values, 0.1)
+            stacked_high = np.percentile(stacked_values, 99.9)
+            bins = np.linspace(stacked_low, stacked_high, 100)
+
+            distr = np.histogram(synthetic_preds[label], bins=bins)[0]
+            distr = np.clip(distr / distr.sum(), 0.001, 0.999)
+            true_distr = np.histogram(actuals, bins=bins)[0]
+            true_distr = np.clip(true_distr / true_distr.sum(), 0.001, 0.999)
+
+            kl = entropy(distr, true_distr)
+            synthetic_summary.loc[label, "KL Divergence"] = kl
+
+        summaries.append(synthetic_summary)
+
+    summary = sum(summaries) / k
+    return summary[["Abs % Error of ATE", "MSE", "KL Divergence"]]
+
+
+def scatter_plot_summary(synthetic_summary, k, drop_learners=[], drop_cols=[]):
+    """Generates a scatter plot comparing learner performance. Each learner's performance is plotted as a point in the
+    (Abs % Error of ATE, MSE) space.
+
+    Args:
+        synthetic_summary (pd.DataFrame): summary generated by get_synthetic_summary()
+        k (int): number of simulations (used only for plot title text)
+        drop_learners (list, optional): list of learners (str) to omit when plotting
+        drop_cols (list, optional): list of metrics (str) to omit when plotting
+    """
+    plot_data = synthetic_summary.drop(drop_learners).drop(drop_cols, axis=1)
+
+    fig, ax = plt.subplots()
+    fig.set_size_inches(12, 8)
+    xs = plot_data["Abs % Error of ATE"]
+    ys = plot_data["MSE"]
+
+    ax.scatter(xs, ys)
+
+    ylim = ax.get_ylim()
+    xlim = ax.get_xlim()
+
+    for i, txt in enumerate(plot_data.index):
+        ax.annotate(
+            txt,
+            (
+                xs[i] - np.random.binomial(1, 0.5) * xlim[1] * 0.04,
+                ys[i] - ylim[1] * 0.03,
+            ),
+        )
+
+    ax.set_xlabel("Abs % Error of ATE")
+    ax.set_ylabel("MSE")
+    ax.set_title("Learner Performance (averaged over k={} simulations)".format(k))
+
+
+def bar_plot_summary(
+    synthetic_summary,
+    k,
+    drop_learners=[],
+    drop_cols=[],
+    sort_cols=["MSE", "Abs % Error of ATE"],
+):
+    """Generates a bar plot comparing learner performance.
+
+    Args:
+        synthetic_summary (pd.DataFrame): summary generated by get_synthetic_summary()
+        k (int): number of simulations (used only for plot title text)
+        drop_learners (list, optional): list of learners (str) to omit when plotting
+        drop_cols (list, optional): list of metrics (str) to omit when plotting
+        sort_cols (list, optional): list of metrics (str) to sort on when plotting
+    """
+    plot_data = synthetic_summary.sort_values(sort_cols, ascending=True)
+    plot_data = plot_data.drop(drop_learners + [KEY_ACTUAL]).drop(drop_cols, axis=1)
+
+    plot_data.plot(kind="bar", figsize=(12, 8))
+    plt.xticks(rotation=30)
+    plt.title("Learner Performance (averaged over k={} simulations)".format(k))
+
+
+def distr_plot_single_sim(
+    synthetic_preds,
+    kind="kde",
+    drop_learners=[],
+    bins=50,
+    histtype="step",
+    alpha=1,
+    linewidth=1,
+    bw_method=1,
+):
+    """Plots the distribution of each learner's predictions (for a single simulation).
+    Kernel Density Estimation (kde) and actual histogram plots supported.
+
+    Args:
+        synthetic_preds (dict): dictionary of predictions generated by get_synthetic_preds()
+        kind (str, optional): 'kde' or 'hist'
+        drop_learners (list, optional): list of learners (str) to omit when plotting
+        bins (int, optional): number of bins to plot if kind set to 'hist'
+        histtype (str, optional): histogram type if kind set to 'hist'
+        alpha (float, optional): alpha (transparency) for plotting
+        linewidth (int, optional): line width for plotting
+        bw_method (float, optional): parameter for kde
+    """
+    preds_for_plot = synthetic_preds.copy()
+
+    # deleted generated data and assign actual value
+    del preds_for_plot[KEY_GENERATED_DATA]
+    global_lower = np.percentile(np.hstack(list(preds_for_plot.values())), 1)
+    global_upper = np.percentile(np.hstack(list(preds_for_plot.values())), 99)
+    learners = list(preds_for_plot.keys())
+    learners = [learner for learner in learners if learner not in drop_learners]
+
+    # Plotting
+    plt.figure(figsize=(12, 8))
+    colors = [
+        "black",
+        "red",
+        "blue",
+        "green",
+        "cyan",
+        "brown",
+        "grey",
+        "pink",
+        "orange",
+        "yellow",
+    ]
+    for i, (k, v) in enumerate(preds_for_plot.items()):
+        if k in learners:
+            if kind == "kde":
+                v = pd.Series(v.flatten())
+                v = v[v.between(global_lower, global_upper)]
+                v.plot(
+                    kind="kde",
+                    bw_method=bw_method,
+                    label=k,
+                    linewidth=linewidth,
+                    color=colors[i],
+                )
+            elif kind == "hist":
+                plt.hist(
+                    v,
+                    bins=np.linspace(global_lower, global_upper, bins),
+                    label=k,
+                    histtype=histtype,
+                    alpha=alpha,
+                    linewidth=linewidth,
+                    color=colors[i],
+                )
+            else:
+                pass
+
+    plt.xlim(global_lower, global_upper)
+    plt.legend(loc="center left", bbox_to_anchor=(1, 0.5))
+    plt.title("Distribution from a Single Simulation")
+
+
+def scatter_plot_single_sim(synthetic_preds):
+    """Creates a grid of scatter plots comparing each learner's predictions with the truth (for a single simulation).
+
+    Args:
+        synthetic_preds (dict): dictionary of predictions generated by get_synthetic_preds() or
+            get_synthetic_preds_holdout()
+    """
+    preds_for_plot = synthetic_preds.copy()
+
+    # deleted generated data and get actual column name
+    del preds_for_plot[KEY_GENERATED_DATA]
+    n_row = int(np.ceil(len(preds_for_plot.keys()) / 3))
+
+    fig, axes = plt.subplots(n_row, 3, figsize=(5 * n_row, 15))
+    axes = np.ravel(axes)
+
+    for i, (label, preds) in enumerate(preds_for_plot.items()):
+        axes[i].scatter(preds_for_plot[KEY_ACTUAL], preds, s=2, label="Predictions")
+        axes[i].set_title(label, size=12)
+        axes[i].set_xlabel("Actual", size=10)
+        axes[i].set_ylabel("Prediction", size=10)
+        xlim = axes[i].get_xlim()
+        ylim = axes[i].get_xlim()
+        axes[i].plot(
+            [xlim[0], xlim[1]],
+            [ylim[0], ylim[1]],
+            label="Perfect Model",
+            linewidth=1,
+            color="grey",
+        )
+        axes[i].legend(loc=2, prop={"size": 10})
+
+
+def get_synthetic_preds_holdout(
+    synthetic_data_func, n=1000, valid_size=0.2, estimators={}
+):
+    """Generate predictions for synthetic data using specified function (single simulation) for train and holdout
+
+    Args:
+        synthetic_data_func (function): synthetic data generation function
+        n (int, optional): number of samples
+        valid_size(float,optional): validaiton/hold out data size
+        estimators (dict of object): dict of names and objects of treatment effect estimators
+
+    Returns:
+        (tuple): synthetic training and validation data dictionaries:
+
+          - preds_dict_train (dict): synthetic training data dictionary
+          - preds_dict_valid (dict): synthetic validation data dictionary
+    """
+    y, X, w, tau, b, e = synthetic_data_func(n=n)
+
+    (
+        X_train,
+        X_val,
+        y_train,
+        y_val,
+        w_train,
+        w_val,
+        tau_train,
+        tau_val,
+        b_train,
+        b_val,
+        e_train,
+        e_val,
+    ) = train_test_split(
+        X, y, w, tau, b, e, test_size=valid_size, random_state=RANDOM_SEED, shuffle=True
+    )
+
+    preds_dict_train = {}
+    preds_dict_valid = {}
+
+    preds_dict_train[KEY_ACTUAL] = tau_train
+    preds_dict_valid[KEY_ACTUAL] = tau_val
+
+    preds_dict_train["generated_data"] = {
+        "y": y_train,
+        "X": X_train,
+        "w": w_train,
+        "tau": tau_train,
+        "b": b_train,
+        "e": e_train,
+    }
+    preds_dict_valid["generated_data"] = {
+        "y": y_val,
+        "X": X_val,
+        "w": w_val,
+        "tau": tau_val,
+        "b": b_val,
+        "e": e_val,
+    }
+
+    # Predict p_hat because e would not be directly observed in real-life
+    p_model = ElasticNetPropensityModel()
+    p_hat_train = p_model.fit_predict(X_train, w_train)
+    p_hat_val = p_model.fit_predict(X_val, w_val)
+
+    for base_learner, label_l in zip(
+        [BaseSRegressor, BaseTRegressor, BaseXRegressor, BaseRRegressor],
+        ["S", "T", "X", "R"],
+    ):
+        for model, label_m in zip([LinearRegression, XGBRegressor], ["LR", "XGB"]):
+            # RLearner will need to fit on the p_hat
+            if label_l != "R":
+                learner = base_learner(model())
+                # fit the model on training data only
+                learner.fit(X=X_train, treatment=w_train, y=y_train)
+                try:
+                    preds_dict_train["{} Learner ({})".format(label_l, label_m)] = (
+                        learner.predict(X=X_train, p=p_hat_train).flatten()
+                    )
+                    preds_dict_valid["{} Learner ({})".format(label_l, label_m)] = (
+                        learner.predict(X=X_val, p=p_hat_val).flatten()
+                    )
+                except TypeError:
+                    preds_dict_train["{} Learner ({})".format(label_l, label_m)] = (
+                        learner.predict(
+                            X=X_train, treatment=w_train, y=y_train
+                        ).flatten()
+                    )
+                    preds_dict_valid["{} Learner ({})".format(label_l, label_m)] = (
+                        learner.predict(X=X_val, treatment=w_val, y=y_val).flatten()
+                    )
+            else:
+                learner = base_learner(model())
+                learner.fit(X=X_train, p=p_hat_train, treatment=w_train, y=y_train)
+                preds_dict_train["{} Learner ({})".format(label_l, label_m)] = (
+                    learner.predict(X=X_train).flatten()
+                )
+                preds_dict_valid["{} Learner ({})".format(label_l, label_m)] = (
+                    learner.predict(X=X_val).flatten()
+                )
+
+    return preds_dict_train, preds_dict_valid
+
+
+def get_synthetic_summary_holdout(synthetic_data_func, n=1000, valid_size=0.2, k=1):
+    """Generate a summary for predictions on synthetic data for train and holdout using specified function
+
+    Args:
+        synthetic_data_func (function): synthetic data generation function
+        n (int, optional): number of samples per simulation
+        valid_size(float,optional): validation/hold out data size
+        k (int, optional): number of simulations
+
+
+    Returns:
+        (tuple): summary evaluation metrics of predictions for train and validation:
+
+          - summary_train (pandas.DataFrame): training data evaluation summary
+          - summary_train (pandas.DataFrame): validation data evaluation summary
+    """
+
+    summaries_train = []
+    summaries_validation = []
+
+    for i in range(k):
+        preds_dict_train, preds_dict_valid = get_synthetic_preds_holdout(
+            synthetic_data_func, n=n, valid_size=valid_size
+        )
+        actuals_train = preds_dict_train[KEY_ACTUAL]
+        actuals_validation = preds_dict_valid[KEY_ACTUAL]
+
+        synthetic_summary_train = pd.DataFrame(
+            {
+                label: [preds.mean(), mse(preds, actuals_train)]
+                for label, preds in preds_dict_train.items()
+                if KEY_GENERATED_DATA not in label.lower()
+            },
+            index=["ATE", "MSE"],
+        ).T
+        synthetic_summary_train["Abs % Error of ATE"] = np.abs(
+            (
+                synthetic_summary_train["ATE"]
+                / synthetic_summary_train.loc[KEY_ACTUAL, "ATE"]
+            )
+            - 1
+        )
+
+        synthetic_summary_validation = pd.DataFrame(
+            {
+                label: [preds.mean(), mse(preds, actuals_validation)]
+                for label, preds in preds_dict_valid.items()
+                if KEY_GENERATED_DATA not in label.lower()
+            },
+            index=["ATE", "MSE"],
+        ).T
+        synthetic_summary_validation["Abs % Error of ATE"] = np.abs(
+            (
+                synthetic_summary_validation["ATE"]
+                / synthetic_summary_validation.loc[KEY_ACTUAL, "ATE"]
+            )
+            - 1
+        )
+
+        # calculate kl divergence for training
+        for label in synthetic_summary_train.index:
+            stacked_values = np.hstack((preds_dict_train[label], actuals_train))
+            stacked_low = np.percentile(stacked_values, 0.1)
+            stacked_high = np.percentile(stacked_values, 99.9)
+            bins = np.linspace(stacked_low, stacked_high, 100)
+
+            distr = np.histogram(preds_dict_train[label], bins=bins)[0]
+            distr = np.clip(distr / distr.sum(), 0.001, 0.999)
+            true_distr = np.histogram(actuals_train, bins=bins)[0]
+            true_distr = np.clip(true_distr / true_distr.sum(), 0.001, 0.999)
+
+            kl = entropy(distr, true_distr)
+            synthetic_summary_train.loc[label, "KL Divergence"] = kl
+
+        # calculate kl divergence for validation
+        for label in synthetic_summary_validation.index:
+            stacked_values = np.hstack((preds_dict_valid[label], actuals_validation))
+            stacked_low = np.percentile(stacked_values, 0.1)
+            stacked_high = np.percentile(stacked_values, 99.9)
+            bins = np.linspace(stacked_low, stacked_high, 100)
+
+            distr = np.histogram(preds_dict_valid[label], bins=bins)[0]
+            distr = np.clip(distr / distr.sum(), 0.001, 0.999)
+            true_distr = np.histogram(actuals_validation, bins=bins)[0]
+            true_distr = np.clip(true_distr / true_distr.sum(), 0.001, 0.999)
+
+            kl = entropy(distr, true_distr)
+            synthetic_summary_validation.loc[label, "KL Divergence"] = kl
+
+        summaries_train.append(synthetic_summary_train)
+        summaries_validation.append(synthetic_summary_validation)
+
+    summary_train = sum(summaries_train) / k
+    summary_validation = sum(summaries_validation) / k
+    return (
+        summary_train[["Abs % Error of ATE", "MSE", "KL Divergence"]],
+        summary_validation[["Abs % Error of ATE", "MSE", "KL Divergence"]],
+    )
+
+
+def scatter_plot_summary_holdout(
+    train_summary,
+    validation_summary,
+    k,
+    label=["Train", "Validation"],
+    drop_learners=[],
+    drop_cols=[],
+):
+    """Generates a scatter plot comparing learner performance by training and validation.
+
+    Args:
+        train_summary (pd.DataFrame): summary for training synthetic data generated by get_synthetic_summary_holdout()
+        validation_summary (pd.DataFrame): summary for validation synthetic data generated by
+            get_synthetic_summary_holdout()
+        label (string, optional): legend label for plot
+        k (int): number of simulations (used only for plot title text)
+        drop_learners (list, optional): list of learners (str) to omit when plotting
+        drop_cols (list, optional): list of metrics (str) to omit when plotting
+    """
+    train_summary = train_summary.drop(drop_learners).drop(drop_cols, axis=1)
+    validation_summary = validation_summary.drop(drop_learners).drop(drop_cols, axis=1)
+
+    plot_data = pd.concat([train_summary, validation_summary])
+    plot_data["label"] = [i.replace("Train", "") for i in plot_data.index]
+    plot_data["label"] = [i.replace("Validation", "") for i in plot_data.label]
+
+    fig, ax = plt.subplots()
+    fig.set_size_inches(12, 8)
+    xs = plot_data["Abs % Error of ATE"]
+    ys = plot_data["MSE"]
+    group = np.array(
+        [label[0]] * train_summary.shape[0] + [label[1]] * validation_summary.shape[0]
+    )
+    cdict = {label[0]: "red", label[1]: "blue"}
+
+    for g in np.unique(group):
+        ix = np.where(group == g)[0].tolist()
+        ax.scatter(xs[ix], ys[ix], c=cdict[g], label=g, s=100)
+
+    for i, txt in enumerate(plot_data.label[:10]):
+        ax.annotate(txt, (xs[i] + 0.005, ys[i]))
+
+    ax.set_xlabel("Abs % Error of ATE")
+    ax.set_ylabel("MSE")
+    ax.set_title("Learner Performance (averaged over k={} simulations)".format(k))
+    ax.legend(loc="center left", bbox_to_anchor=(1.1, 0.5))
+    plt.show()
+
+
+def bar_plot_summary_holdout(
+    train_summary, validation_summary, k, drop_learners=[], drop_cols=[]
+):
+    """Generates a bar plot comparing learner performance by training and validation
+
+    Args:
+        train_summary (pd.DataFrame): summary for training synthetic data generated by get_synthetic_summary_holdout()
+        validation_summary (pd.DataFrame): summary for validation synthetic data generated by
+            get_synthetic_summary_holdout()
+        k (int): number of simulations (used only for plot title text)
+        drop_learners (list, optional): list of learners (str) to omit when plotting
+        drop_cols (list, optional): list of metrics (str) to omit when plotting
+    """
+    train_summary = train_summary.drop([KEY_ACTUAL])
+    train_summary["Learner"] = train_summary.index
+
+    validation_summary = validation_summary.drop([KEY_ACTUAL])
+    validation_summary["Learner"] = validation_summary.index
+
+    for metric in ["Abs % Error of ATE", "MSE", "KL Divergence"]:
+        plot_data_sub = pd.DataFrame(train_summary.Learner).reset_index(drop=True)
+        plot_data_sub["train"] = train_summary[metric].values
+        plot_data_sub["validation"] = validation_summary[metric].values
+        plot_data_sub = plot_data_sub.set_index("Learner")
+        plot_data_sub = plot_data_sub.drop(drop_learners).drop(drop_cols, axis=1)
+        plot_data_sub = plot_data_sub.sort_values("train", ascending=True)
+
+        plot_data_sub.plot(kind="bar", color=["red", "blue"], figsize=(12, 8))
+        plt.xticks(rotation=30)
+        plt.title(
+            "Learner Performance of {} (averaged over k={} simulations)".format(
+                metric, k
+            )
+        )
+
+
+def get_synthetic_auuc(
+    synthetic_preds,
+    drop_learners=[],
+    outcome_col="y",
+    treatment_col="w",
+    treatment_effect_col="tau",
+    plot=True,
+):
+    """Get auuc values for cumulative gains of model estimates in quantiles.
+
+    For details, reference get_cumgain() and plot_gain()
+    Args:
+        synthetic_preds (dict): dictionary of predictions generated by get_synthetic_preds()
+        or get_synthetic_preds_holdout()
+        outcome_col (str, optional): the column name for the actual outcome
+        treatment_col (str, optional): the column name for the treatment indicator (0 or 1)
+        treatment_effect_col (str, optional): the column name for the true treatment effect
+        plot (boolean,optional): plot the cumulative gain chart or not
+
+    Returns:
+        (pandas.DataFrame): auuc values by learner for cumulative gains of model estimates
+    """
+    synthetic_preds_df = synthetic_preds.copy()
+    generated_data = synthetic_preds_df.pop(KEY_GENERATED_DATA)
+    synthetic_preds_df = pd.DataFrame(synthetic_preds_df)
+    synthetic_preds_df = synthetic_preds_df.drop(drop_learners, axis=1)
+
+    synthetic_preds_df["y"] = generated_data[outcome_col]
+    synthetic_preds_df["w"] = generated_data[treatment_col]
+    if treatment_effect_col in generated_data.keys():
+        synthetic_preds_df["tau"] = generated_data[treatment_effect_col]
+
+    assert (
+        (outcome_col in synthetic_preds_df.columns)
+        and (treatment_col in synthetic_preds_df.columns)
+        or treatment_effect_col in synthetic_preds_df.columns
+    )
+
+    cumlift = get_cumgain(
+        synthetic_preds_df,
+        outcome_col="y",
+        treatment_col="w",
+        treatment_effect_col="tau",
+    )
+    auuc_df = pd.DataFrame(cumlift.columns)
+    auuc_df.columns = ["Learner"]
+    auuc_df["cum_gain_auuc"] = [
+        auc(cumlift.index.values / 100, cumlift[learner].values)
+        for learner in cumlift.columns
+    ]
+    auuc_df = auuc_df.sort_values("cum_gain_auuc", ascending=False)
+
+    if plot:
+        plot_gain(
+            synthetic_preds_df,
+            outcome_col=outcome_col,
+            treatment_col=treatment_col,
+            treatment_effect_col=treatment_effect_col,
+        )
+
+    return auuc_df

causalml/source/causalml/feature_selection/__init__.py

@@ -0,0 +1 @@
+from .filters import FilterSelect

causalml/source/causalml/feature_selection/filters.py

@@ -0,0 +1,663 @@
+"""
+Filter feature selection methods for uplift modeling
+
+- Currently only for classification problem: the outcome variable of uplift model is binary.
+"""
+
+import numpy as np
+import pandas as pd
+import statsmodels.api as sm
+from scipy import stats
+from sklearn.impute import SimpleImputer
+
+
+class FilterSelect:
+    """A class for feature importance methods."""
+
+    def __init__(self):
+        return
+
+    @staticmethod
+    def _filter_F_one_feature(data, treatment_indicator, feature_name, y_name, order=1):
+        """
+        Conduct F-test of the interaction between treatment and one feature.
+
+        Args:
+            data (pd.Dataframe): DataFrame containing outcome, features, and experiment group
+            treatment_indicator (string): the column name for binary indicator of treatment (1) or control (0)
+            feature_name (string): feature name, as one column in the data DataFrame
+            y_name (string): name of the outcome variable
+            order (int): the order of feature to be evaluated with the treatment effect, order takes 3 values: 1,2,3.
+                order = 1 corresponds to linear importance of the feature, order=2 corresponds to quadratic and linear
+                importance of the feature,
+            order= 3 will calculate feature importance up to cubic forms.
+
+        Returns:
+            F_test_result : pd.DataFrame
+                a data frame containing the feature importance statistics
+        """
+        Y = data[y_name]
+        X = data[[treatment_indicator, feature_name]]
+        X = sm.add_constant(X)
+        X["{}-{}".format(treatment_indicator, feature_name)] = X[
+            [treatment_indicator, feature_name]
+        ].product(axis=1)
+
+        if order not in [1, 2, 3]:
+            raise Exception("ValueError: order argument only takes value 1,2,3.")
+
+        if order == 1:
+            pass
+        elif order == 2:
+            x_tmp_name = "{}_o{}".format(feature_name, order)
+            X[x_tmp_name] = X[[feature_name]] ** order
+            X["{}-{}".format(treatment_indicator, x_tmp_name)] = X[
+                [treatment_indicator, x_tmp_name]
+            ].product(axis=1)
+        elif order == 3:
+            x_tmp_name = "{}_o{}".format(feature_name, 2)
+            X[x_tmp_name] = X[[feature_name]] ** 2
+            X["{}-{}".format(treatment_indicator, x_tmp_name)] = X[
+                [treatment_indicator, x_tmp_name]
+            ].product(axis=1)
+
+            x_tmp_name = "{}_o{}".format(feature_name, order)
+            X[x_tmp_name] = X[[feature_name]] ** order
+            X["{}-{}".format(treatment_indicator, x_tmp_name)] = X[
+                [treatment_indicator, x_tmp_name]
+            ].product(axis=1)
+
+        model = sm.OLS(Y, X)
+        result = model.fit()
+
+        if order == 1:
+            F_test = result.f_test(np.array([0, 0, 0, 1]))
+        elif order == 2:
+            F_test = result.f_test(np.array([[0, 0, 0, 1, 0, 0], [0, 0, 0, 0, 0, 1]]))
+        elif order == 3:
+            F_test = result.f_test(
+                np.array(
+                    [
+                        [0, 0, 0, 1, 0, 0, 0, 0],
+                        [0, 0, 0, 0, 0, 1, 0, 0],
+                        [0, 0, 0, 0, 0, 0, 0, 1],
+                    ]
+                )
+            )
+
+        F_test_result = pd.DataFrame(
+            {
+                "feature": feature_name,  # for the interaction, not the main effect
+                "method": "F{} Filter".format(order),
+                "score": float(F_test.fvalue),
+                "p_value": F_test.pvalue,
+                "misc": "df_num: {}, df_denom: {}, order:{}".format(
+                    F_test.df_num, F_test.df_denom, order
+                ),
+            },
+            index=[0],
+        ).reset_index(drop=True)
+
+        return F_test_result
+
+    def filter_F(self, data, treatment_indicator, features, y_name, order=1):
+        """
+        Rank features based on the F-statistics of the interaction.
+
+        Args:
+            data (pd.Dataframe): DataFrame containing outcome, features, and experiment group
+            treatment_indicator (string): the column name for binary indicator of treatment (1) or control (0)
+            features (list of string): list of feature names, that are columns in the data DataFrame
+            y_name (string): name of the outcome variable
+            order (int): the order of feature to be evaluated with the treatment effect, order takes 3 values: 1,2,3.
+                order = 1 corresponds to linear importance of the feature, order=2 corresponds to quadratic and linear
+                importance of the feature,
+            order= 3 will calculate feature importance up to cubic forms.
+
+        Returns:
+            all_result : pd.DataFrame
+                a data frame containing the feature importance statistics
+        """
+        if order not in [1, 2, 3]:
+            raise Exception("ValueError: order argument only takes value 1,2,3.")
+
+        all_result = pd.DataFrame()
+        for x_name_i in features:
+            one_result = self._filter_F_one_feature(
+                data=data,
+                treatment_indicator=treatment_indicator,
+                feature_name=x_name_i,
+                y_name=y_name,
+                order=order,
+            )
+            all_result = pd.concat([all_result, one_result])
+
+        all_result = all_result.sort_values(by="score", ascending=False)
+        all_result["rank"] = all_result["score"].rank(ascending=False)
+
+        return all_result
+
+    @staticmethod
+    def _filter_LR_one_feature(
+        data, treatment_indicator, feature_name, y_name, order=1, disp=True
+    ):
+        """
+        Conduct LR (Likelihood Ratio) test of the interaction between treatment and one feature.
+
+        Args:
+            data (pd.Dataframe): DataFrame containing outcome, features, and experiment group
+            treatment_indicator (string): the column name for binary indicator of treatment (1) or control (0)
+            feature_name (string): feature name, as one column in the data DataFrame
+            y_name (string): name of the outcome variable
+            order (int): the order of feature to be evaluated with the treatment effect, order takes 3 values: 1,2,3.
+                order = 1 corresponds to linear importance of the feature, order=2 corresponds to quadratic and linear
+                importance of the feature,
+            order= 3 will calculate feature importance up to cubic forms.
+
+        Returns:
+            LR_test_result : pd.DataFrame
+                a data frame containing the feature importance statistics
+        """
+        Y = data[y_name]
+
+        # Restricted model
+        x_name_r = ["const", treatment_indicator, feature_name]
+        x_name_f = x_name_r.copy()
+        X = data[[treatment_indicator, feature_name]]
+        X = sm.add_constant(X)
+
+        X["{}-{}".format(treatment_indicator, feature_name)] = X[
+            [treatment_indicator, feature_name]
+        ].product(axis=1)
+        x_name_f.append("{}-{}".format(treatment_indicator, feature_name))
+
+        if order == 2:
+            x_tmp_name = "{}_o{}".format(feature_name, order)
+            X[x_tmp_name] = X[[feature_name]] ** order
+            X["{}-{}".format(treatment_indicator, x_tmp_name)] = X[
+                [treatment_indicator, x_tmp_name]
+            ].product(axis=1)
+            x_name_r.append(x_tmp_name)
+            x_name_f += [x_tmp_name, "{}-{}".format(treatment_indicator, x_tmp_name)]
+        elif order == 3:
+            x_tmp_name = "{}_o{}".format(feature_name, 2)
+            X[x_tmp_name] = X[[feature_name]] ** 2
+            X["{}-{}".format(treatment_indicator, x_tmp_name)] = X[
+                [treatment_indicator, x_tmp_name]
+            ].product(axis=1)
+            x_name_r.append(x_tmp_name)
+            x_name_f += [x_tmp_name, "{}-{}".format(treatment_indicator, x_tmp_name)]
+            x_tmp_name = "{}_o{}".format(feature_name, order)
+            X[x_tmp_name] = X[[feature_name]] ** order
+            X["{}-{}".format(treatment_indicator, x_tmp_name)] = X[
+                [treatment_indicator, x_tmp_name]
+            ].product(axis=1)
+            x_name_r.append(x_tmp_name)
+            x_name_f += [x_tmp_name, "{}-{}".format(treatment_indicator, x_tmp_name)]
+
+        # Full model (with interaction)
+        model_r = sm.Logit(Y, X[x_name_r])
+        result_r = model_r.fit(disp=disp)
+
+        model_f = sm.Logit(Y, X[x_name_f])
+        result_f = model_f.fit(disp=disp)
+
+        LR_stat = -2 * (result_r.llf - result_f.llf)
+        LR_df = len(result_f.params) - len(result_r.params)
+        LR_pvalue = 1 - stats.chi2.cdf(LR_stat, df=LR_df)
+
+        LR_test_result = pd.DataFrame(
+            {
+                "feature": feature_name,  # for the interaction, not the main effect
+                "method": "LR{} Filter".format(order),
+                "score": LR_stat,
+                "p_value": LR_pvalue,
+                "misc": "df: {}, order: {}".format(LR_df, order),
+            },
+            index=[0],
+        ).reset_index(drop=True)
+
+        return LR_test_result
+
+    def filter_LR(
+        self, data, treatment_indicator, features, y_name, order=1, disp=True
+    ):
+        """
+        Rank features based on the LRT-statistics of the interaction.
+
+        Args:
+            data (pd.Dataframe): DataFrame containing outcome, features, and experiment group
+            treatment_indicator (string): the column name for binary indicator of treatment (1) or control (0)
+            feature_name (string): feature name, as one column in the data DataFrame
+            y_name (string): name of the outcome variable
+            order (int): the order of feature to be evaluated with the treatment effect, order takes 3 values: 1,2,3.
+                order = 1 corresponds to linear importance of the feature, order=2 corresponds to quadratic and linear
+                importance of the feature,
+            order= 3 will calculate feature importance up to cubic forms.
+
+        Returns:
+            all_result : pd.DataFrame
+                a data frame containing the feature importance statistics
+        """
+        if order not in [1, 2, 3]:
+            raise Exception("ValueError: order argument only takes value 1,2,3.")
+
+        all_result = pd.DataFrame()
+        for x_name_i in features:
+            one_result = self._filter_LR_one_feature(
+                data=data,
+                treatment_indicator=treatment_indicator,
+                feature_name=x_name_i,
+                y_name=y_name,
+                order=order,
+                disp=disp,
+            )
+            all_result = pd.concat([all_result, one_result])
+
+        all_result = all_result.sort_values(by="score", ascending=False)
+        all_result["rank"] = all_result["score"].rank(ascending=False)
+
+        return all_result
+
+    # Get node summary - a function
+    @staticmethod
+    def _GetNodeSummary(
+        data,
+        experiment_group_column="treatment_group_key",
+        y_name="conversion",
+        smooth=True,
+    ):
+        """
+        To count the conversions and get the probabilities by treatment groups. This function comes from the uplift
+        tree algorithm, that is used for tree node split evaluation.
+
+        Parameters
+        ----------
+        data : DataFrame
+            The DataFrame that contains all the data (in the current "node").
+        experiment_group_column : str
+            Treatment indicator column name.
+        y_name : str
+            Label indicator column name.
+        smooth : bool
+            Smooth label count by adding 1 in case certain labels do not occur
+            naturally with a treatment. Prevents zero divisions.
+
+        Returns
+        -------
+        results : dict
+            Counts of conversions by treatment groups, of the form:
+            {'control': {0: 10, 1: 8}, 'treatment1': {0: 5, 1: 15}}
+        nodeSummary: dict
+            Probability of conversion and group size by treatment groups, of
+            the form:
+            {'control': [0.490, 500], 'treatment1': [0.584, 500]}
+        """
+
+        # Note: results and nodeSummary are both dict with treatment_group_key
+        # as the key.  So we can compute the treatment effect and/or
+        # divergence easily.
+
+        # Counts of conversions by treatment group
+        results_series = data.groupby([experiment_group_column, y_name]).size()
+
+        treatment_group_keys = results_series.index.levels[0].tolist()
+        y_name_keys = results_series.index.levels[1].tolist()
+
+        results = {}
+        for ti in treatment_group_keys:
+            results.update({ti: {}})
+            for ci in y_name_keys:
+                if smooth:
+                    results[ti].update(
+                        {
+                            ci: (
+                                results_series[ti, ci]
+                                if results_series.index.isin([(ti, ci)]).any()
+                                else 1
+                            )
+                        }
+                    )
+                else:
+                    results[ti].update({ci: results_series[ti, ci]})
+
+        # Probability of conversion and group size by treatment group
+        nodeSummary = {}
+        for treatment_group_key in results:
+            n_1 = results[treatment_group_key].get(1, 0)
+            n_total = results[treatment_group_key].get(1, 0) + results[
+                treatment_group_key
+            ].get(0, 0)
+            y_mean = 1.0 * n_1 / n_total
+            nodeSummary[treatment_group_key] = [y_mean, n_total]
+
+        return results, nodeSummary
+
+    # Divergence-related functions, from upliftpy
+    @staticmethod
+    def _kl_divergence(pk, qk):
+        """
+        Calculate KL Divergence for binary classification.
+
+        Args:
+            pk (float): Probability of class 1 in treatment group
+            qk (float): Probability of class 1 in control group
+        """
+        if qk < 0.1**6:
+            qk = 0.1**6
+        elif qk > 1 - 0.1**6:
+            qk = 1 - 0.1**6
+        S = pk * np.log(pk / qk) + (1 - pk) * np.log((1 - pk) / (1 - qk))
+        return S
+
+    def _evaluate_KL(self, nodeSummary, control_group="control"):
+        """
+        Calculate the multi-treatment unconditional D (one node)
+        with KL Divergence as split Evaluation function.
+
+        Args:
+            nodeSummary (dict): a dictionary containing the statistics for a tree node sample
+            control_group (string, optional, default='control'): the name for control group
+
+        Notes
+        -----
+        The function works for more than one non-control treatment groups.
+        """
+        if control_group not in nodeSummary:
+            return 0
+        pc = nodeSummary[control_group][0]
+        d_res = 0
+        for treatment_group in nodeSummary:
+            if treatment_group != control_group:
+                d_res += self._kl_divergence(nodeSummary[treatment_group][0], pc)
+        return d_res
+
+    @staticmethod
+    def _evaluate_ED(nodeSummary, control_group="control"):
+        """
+        Calculate the multi-treatment unconditional D (one node)
+        with Euclidean Distance as split Evaluation function.
+
+        Args:
+            nodeSummary (dict): a dictionary containing the statistics for a tree node sample
+            control_group (string, optional, default='control'): the name for control group
+        """
+        if control_group not in nodeSummary:
+            return 0
+        pc = nodeSummary[control_group][0]
+        d_res = 0
+        for treatment_group in nodeSummary:
+            if treatment_group != control_group:
+                d_res += 2 * (nodeSummary[treatment_group][0] - pc) ** 2
+        return d_res
+
+    @staticmethod
+    def _evaluate_Chi(nodeSummary, control_group="control"):
+        """
+        Calculate the multi-treatment unconditional D (one node)
+        with Chi-Square as split Evaluation function.
+
+        Args:
+            nodeSummary (dict): a dictionary containing the statistics for a tree node sample
+            control_group (string, optional, default='control'): the name for control group
+        """
+        if control_group not in nodeSummary:
+            return 0
+        pc = nodeSummary[control_group][0]
+        d_res = 0
+        for treatment_group in nodeSummary:
+            if treatment_group != control_group:
+                d_res += (nodeSummary[treatment_group][0] - pc) ** 2 / max(
+                    0.1**6, pc
+                ) + (nodeSummary[treatment_group][0] - pc) ** 2 / max(0.1**6, 1 - pc)
+        return d_res
+
+    def _filter_D_one_feature(
+        self,
+        data,
+        feature_name,
+        y_name,
+        n_bins=10,
+        method="KL",
+        control_group="control",
+        experiment_group_column="treatment_group_key",
+        null_impute=None,
+    ):
+        """
+        Calculate the chosen divergence measure for one feature.
+
+        Args:
+            data (pd.Dataframe): DataFrame containing outcome, features, and experiment group
+            treatment_indicator (string): the column name for binary indicator of treatment (1) or control (0)
+            feature_name (string): feature name, as one column in the data DataFrame
+            y_name (string): name of the outcome variable
+            method (string, optional, default = 'KL'): taking one of the following values {'F', 'LR', 'KL', 'ED', 'Chi'}
+                The feature selection method to be used to rank the features.
+                'F' for F-test
+                'LR' for likelihood ratio test
+                'KL', 'ED', 'Chi' for bin-based uplift filter methods, KL divergence, Euclidean distance,
+                Chi-Square respectively
+            experiment_group_column (string, optional, default = 'treatment_group_key'): the experiment column name in
+                the DataFrame, which contains the treatment and control assignment label
+            control_group (string, optional, default = 'control'): name for control group, value in the experiment
+                group column
+            n_bins (int, optional, default = 10): number of bins to be used for bin-based uplift filter methods
+            null_impute (str, optional, default=None): impute np.nan present in the data taking on of the following
+                strategy values {'mean', 'median', 'most_frequent', None}. If Value is None and null is present then
+                exception will be raised
+
+        Returns:
+            D_result : pd.DataFrame
+                a data frame containing the feature importance statistics
+        """
+        # [TODO] Application to categorical features
+
+        if method == "KL":
+            evaluationFunction = self._evaluate_KL
+        elif method == "ED":
+            evaluationFunction = self._evaluate_ED
+        elif method == "Chi":
+            evaluationFunction = self._evaluate_Chi
+
+        totalSize = len(data.index)
+
+        # impute null if enabled
+        if null_impute is not None:
+            data[feature_name] = SimpleImputer(
+                missing_values=np.nan, strategy=null_impute
+            ).fit_transform(data[feature_name].values.reshape(-1, 1))
+        elif data[feature_name].isna().any():
+            raise Exception(
+                "Null value(s) present in column '{}'. Please impute the null value or use null_impute parameter "
+                "provided.".format(feature_name)
+            )
+
+        # drop duplicate edges in pq.cut result to avoid issues
+        x_bin = pd.qcut(
+            data[feature_name].values, n_bins, labels=False, duplicates="drop"
+        )
+
+        d_children = 0
+
+        for i_bin in range(np.nanmax(x_bin).astype(int) + 1):  # range(n_bins):
+            nodeSummary = self._GetNodeSummary(
+                data=data.loc[x_bin == i_bin],
+                experiment_group_column=experiment_group_column,
+                y_name=y_name,
+            )[1]
+            nodeScore = evaluationFunction(nodeSummary, control_group=control_group)
+            nodeSize = sum([x[1] for x in list(nodeSummary.values())])
+            d_children += nodeScore * nodeSize / totalSize
+
+        parentNodeSummary = self._GetNodeSummary(
+            data=data, experiment_group_column=experiment_group_column, y_name=y_name
+        )[1]
+        d_parent = evaluationFunction(parentNodeSummary, control_group=control_group)
+
+        d_res = d_children - d_parent
+
+        D_result = pd.DataFrame(
+            {
+                "feature": feature_name,
+                "method": method,
+                "score": d_res,
+                "p_value": None,
+                "misc": "number_of_bins: {}".format(
+                    min(n_bins, np.nanmax(x_bin).astype(int) + 1)
+                ),  # format(n_bins),
+            },
+            index=[0],
+        ).reset_index(drop=True)
+
+        return D_result
+
+    def filter_D(
+        self,
+        data,
+        features,
+        y_name,
+        n_bins=10,
+        method="KL",
+        control_group="control",
+        experiment_group_column="treatment_group_key",
+        null_impute=None,
+    ):
+        """
+        Rank features based on the chosen divergence measure.
+
+        Args:
+            data (pd.Dataframe): DataFrame containing outcome, features, and experiment group
+            treatment_indicator (string): the column name for binary indicator of treatment (1) or control (0)
+            features (list of string): list of feature names, that are columns in the data DataFrame
+            y_name (string): name of the outcome variable
+            method (string, optional, default = 'KL'): taking one of the following values {'F', 'LR', 'KL', 'ED', 'Chi'}
+                The feature selection method to be used to rank the features.
+                'F' for F-test
+                'LR' for likelihood ratio test
+                'KL', 'ED', 'Chi' for bin-based uplift filter methods, KL divergence, Euclidean distance, Chi-Square
+                respectively
+            experiment_group_column (string, optional, default = 'treatment_group_key'): the experiment column name in
+                the DataFrame, which contains the treatment and control assignment label
+            control_group (string, optional, default = 'control'): name for control group, value in the experiment
+                group column
+            n_bins (int, optional, default = 10): number of bins to be used for bin-based uplift filter methods
+            null_impute (str, optional, default=None): impute np.nan present in the data taking on of the followin
+                strategy values {'mean', 'median', 'most_frequent', None}. If Value is None and null is present then
+                exception will be raised
+
+        Returns:
+            all_result : pd.DataFrame
+                a data frame containing the feature importance statistics
+        """
+
+        all_result = pd.DataFrame()
+
+        for x_name_i in features:
+            one_result = self._filter_D_one_feature(
+                data=data,
+                feature_name=x_name_i,
+                y_name=y_name,
+                n_bins=n_bins,
+                method=method,
+                control_group=control_group,
+                experiment_group_column=experiment_group_column,
+                null_impute=null_impute,
+            )
+            all_result = pd.concat([all_result, one_result])
+
+        all_result = all_result.sort_values(by="score", ascending=False)
+        all_result["rank"] = all_result["score"].rank(ascending=False)
+
+        return all_result
+
+    def get_importance(
+        self,
+        data,
+        features,
+        y_name,
+        method,
+        experiment_group_column="treatment_group_key",
+        control_group="control",
+        treatment_group="treatment",
+        n_bins=5,
+        null_impute=None,
+        order=1,
+        disp=False,
+    ):
+        """
+        Rank features based on the chosen statistic of the interaction.
+
+        Args:
+            data (pd.Dataframe): DataFrame containing outcome, features, and experiment group
+            features (list of string): list of feature names, that are columns in the data DataFrame
+            y_name (string): name of the outcome variable
+            method (string, optional, default = 'KL'): taking one of the following values {'F', 'LR', 'KL', 'ED', 'Chi'}
+                The feature selection method to be used to rank the features.
+                'F' for F-test
+                'LR' for likelihood ratio test
+                'KL', 'ED', 'Chi' for bin-based uplift filter methods, KL divergence, Euclidean distance, Chi-Square
+                respectively
+            experiment_group_column (string): the experiment column name in the DataFrame, which contains the treatment
+                and control assignment label
+            control_group (string): name for control group, value in the experiment group column
+            treatment_group (string): name for treatment group, value in the experiment group column
+            n_bins (int, optional): number of bins to be used for bin-based uplift filter methods
+            null_impute (str, optional, default=None): impute np.nan present in the data taking on of the following
+                strategy values {'mean', 'median', 'most_frequent', None}. If value is None and null is present then
+                exception will be raised
+            order (int): the order of feature to be evaluated with the treatment effect for F filter and LR filter,
+                order takes 3 values: 1,2,3. order = 1 corresponds to linear importance of the feature, order=2
+                corresponds to quadratic and linear importance of the feature,
+            order= 3 will calculate feature importance up to cubic forms.
+            disp (bool): Set to True to print convergence messages for Logistic regression convergence in LR method.
+
+        Returns:
+            all_result : pd.DataFrame
+                a data frame with following columns: ['method', 'feature', 'rank', 'score', 'p_value', 'misc']
+        """
+
+        if method == "F":
+            data = data[
+                data[experiment_group_column].isin([control_group, treatment_group])
+            ]
+            data["treatment_indicator"] = 0
+            data.loc[
+                data[experiment_group_column] == treatment_group, "treatment_indicator"
+            ] = 1
+            all_result = self.filter_F(
+                data=data,
+                treatment_indicator="treatment_indicator",
+                features=features,
+                y_name=y_name,
+                order=order,
+            )
+        elif method == "LR":
+            data = data[
+                data[experiment_group_column].isin([control_group, treatment_group])
+            ]
+            data["treatment_indicator"] = 0
+            data.loc[
+                data[experiment_group_column] == treatment_group, "treatment_indicator"
+            ] = 1
+            all_result = self.filter_LR(
+                data=data,
+                disp=disp,
+                treatment_indicator="treatment_indicator",
+                features=features,
+                y_name=y_name,
+                order=order,
+            )
+        else:
+            all_result = self.filter_D(
+                data=data,
+                method=method,
+                features=features,
+                y_name=y_name,
+                n_bins=n_bins,
+                control_group=control_group,
+                experiment_group_column=experiment_group_column,
+                null_impute=null_impute,
+            )
+
+        all_result["method"] = method + " filter"
+        return all_result[["method", "feature", "rank", "score", "p_value", "misc"]]

causalml/source/causalml/features.py

@@ -0,0 +1,267 @@
+import logging
+import numpy as np
+import pandas as pd
+from scipy import sparse
+from sklearn import base
+
+logger = logging.getLogger("causalml")
+
+
+NAN_INT = -98765  # A random integer to impute missing values with
+
+
+class LabelEncoder(base.BaseEstimator):
+    """Label Encoder that groups infrequent values into one label.
+
+    Code from https://github.com/jeongyoonlee/Kaggler/blob/master/kaggler/preprocessing/data.py
+
+    Attributes:
+        min_obs (int): minimum number of observation to assign a label.
+        label_encoders (list of dict): label encoders for columns
+        label_maxes (list of int): maximum of labels for columns
+    """
+
+    def __init__(self, min_obs=10):
+        """Initialize the LabelEncoder class object.
+
+        Args:
+            min_obs (int): minimum number of observation to assign a label.
+        """
+
+        self.min_obs = min_obs
+
+    def __repr__(self):
+        return ("LabelEncoder(min_obs={})").format(self.min_obs)
+
+    def _get_label_encoder_and_max(self, x):
+        """Return a mapping from values and its maximum of a column to integer labels.
+
+        Args:
+            x (pandas.Series): a categorical column to encode.
+
+        Returns:
+            label_encoder (dict): mapping from values of features to integers
+            max_label (int): maximum label
+        """
+
+        # NaN cannot be used as a key for dict. So replace it with a random integer.
+        label_count = x.fillna(NAN_INT).value_counts()
+        n_uniq = label_count.shape[0]
+
+        label_count = label_count[label_count >= self.min_obs]
+        n_uniq_new = label_count.shape[0]
+
+        # If every label appears more than min_obs, new label starts from 0.
+        # Otherwise, new label starts from 1 and 0 is used for all old labels
+        # that appear less than min_obs.
+        offset = 0 if n_uniq == n_uniq_new else 1
+
+        label_encoder = pd.Series(
+            np.arange(n_uniq_new) + offset, index=label_count.index
+        )
+        max_label = label_encoder.max()
+        label_encoder = label_encoder.to_dict()
+
+        return label_encoder, max_label
+
+    def _transform_col(self, x, i):
+        """Encode one categorical column into labels.
+
+        Args:
+            x (pandas.Series): a categorical column to encode
+            i (int): column index
+
+        Returns:
+            x (pandas.Series): a column with labels.
+        """
+        return x.fillna(NAN_INT).map(self.label_encoders[i]).fillna(0)
+
+    def fit(self, X, y=None):
+        self.label_encoders = [None] * X.shape[1]
+        self.label_maxes = [None] * X.shape[1]
+
+        for i, col in enumerate(X.columns):
+            (
+                self.label_encoders[i],
+                self.label_maxes[i],
+            ) = self._get_label_encoder_and_max(X[col])
+
+        return self
+
+    def transform(self, X):
+        """Encode categorical columns into label encoded columns
+
+        Args:
+            X (pandas.DataFrame): categorical columns to encode
+
+        Returns:
+            X (pandas.DataFrame): label encoded columns
+        """
+        X = X.copy()
+        for i, col in enumerate(X.columns):
+            X[col] = self._transform_col(X[col], i).astype(float)
+
+        return X
+
+    def fit_transform(self, X, y=None):
+        """Encode categorical columns into label encoded columns
+
+        Args:
+            X (pandas.DataFrame): categorical columns to encode
+
+        Returns:
+            X (pandas.DataFrame): label encoded columns
+        """
+        X = X.copy()
+        self.label_encoders = [None] * X.shape[1]
+        self.label_maxes = [None] * X.shape[1]
+
+        for i, col in enumerate(X.columns):
+            (
+                self.label_encoders[i],
+                self.label_maxes[i],
+            ) = self._get_label_encoder_and_max(X[col])
+
+            X[col] = (
+                X[col]
+                .fillna(NAN_INT)
+                .map(self.label_encoders[i])
+                .fillna(0)
+                .astype(float)
+            )
+
+        return X
+
+
+class OneHotEncoder(base.BaseEstimator):
+    """One-Hot-Encoder that groups infrequent values into one dummy variable.
+
+    Code from https://github.com/jeongyoonlee/Kaggler/blob/master/kaggler/preprocessing/data.py
+
+    Attributes:
+        min_obs (int): minimum number of observation to create a dummy variable
+        label_encoders (list of (dict, int)): label encoders and their maximums
+                                              for columns
+    """
+
+    def __init__(self, min_obs=10):
+        """Initialize the OneHotEncoder class object.
+
+        Args:
+            min_obs (int): minimum number of observation to create a dummy variable
+        """
+
+        self.min_obs = min_obs
+        self.label_encoder = LabelEncoder(min_obs)
+
+    def __repr__(self):
+        return ("OneHotEncoder(min_obs={})").format(self.min_obs)
+
+    def _transform_col(self, x, i):
+        """Encode one categorical column into sparse matrix with one-hot-encoding.
+
+        Args:
+            x (pandas.Series): a categorical column to encode
+            i (int): column index
+
+        Returns:
+            X (scipy.sparse.coo_matrix): sparse matrix encoding a categorical
+                                         variable into dummy variables
+        """
+
+        labels = self.label_encoder._transform_col(x, i)
+        label_max = self.label_encoder.label_maxes[i]
+
+        # build row and column index for non-zero values of a sparse matrix
+        index = np.array(range(len(labels)))
+        i = index[labels > 0]
+        j = labels[labels > 0] - 1  # column index starts from 0
+
+        if len(i) > 0:
+            return sparse.coo_matrix(
+                (np.ones_like(i), (i, j)), shape=(x.shape[0], label_max)
+            )
+        else:
+            # if there is no non-zero value, return no matrix
+            return None
+
+    def fit(self, X, y=None):
+        self.label_encoder.fit(X)
+
+        return self
+
+    def transform(self, X):
+        """Encode categorical columns into sparse matrix with one-hot-encoding.
+
+        Args:
+            X (pandas.DataFrame): categorical columns to encode
+
+        Returns:
+            X_new (scipy.sparse.coo_matrix): sparse matrix encoding categorical
+                                             variables into dummy variables
+        """
+
+        X_new = None
+        for i, col in enumerate(X.columns):
+            X_col = self._transform_col(X[col], i)
+            if X_col is not None:
+                if X_new is None:
+                    X_new = X_col
+                else:
+                    X_new = sparse.hstack((X_new, X_col))
+
+            logger.debug(
+                "{} --> {} features".format(col, self.label_encoder.label_maxes[i])
+            )
+
+        assert (
+            X_new is not None
+        ), "no column was transformed, please check your dataframe input"
+        return X_new
+
+    def fit_transform(self, X, y=None):
+        """Encode categorical columns into sparse matrix with one-hot-encoding.
+
+        Args:
+            X (pandas.DataFrame): categorical columns to encode
+
+        Returns:
+            sparse matrix encoding categorical variables into dummy variables
+        """
+
+        self.label_encoder.fit(X)
+
+        return self.transform(X)
+
+
+def load_data(data, features, transformations={}):
+    """Load data and set the feature matrix and label vector.
+
+    Args:
+        data (pandas.DataFrame): total input data
+        features (list of str): column names to be used in the inference model
+        transformation (dict of (str, func)): transformations to be applied to features
+
+    Returns:
+        X (numpy.matrix): a feature matrix
+    """
+
+    df = data[features].copy()
+
+    bool_cols = [col for col in df.columns if df[col].dtype == bool]
+    df.loc[:, bool_cols] = df[bool_cols].astype(int)
+
+    for col, transformation in transformations.items():
+        logger.info("Applying {} to {}".format(transformation.__name__, col))
+        df[col] = df[col].apply(transformation)
+
+    cat_cols = [col for col in features if not pd.api.types.is_numeric_dtype(df[col])]
+    num_cols = [col for col in features if col not in cat_cols]
+
+    logger.info("Applying one-hot-encoding to {}".format(cat_cols))
+    ohe = OneHotEncoder(min_obs=df.shape[0] * 0.001)
+    X_cat = ohe.fit_transform(df[cat_cols]).todense()
+
+    X = np.hstack([df[num_cols].values, X_cat])
+
+    return X

causalml/source/causalml/inference/iv/__init__.py

@@ -0,0 +1,2 @@
+from .iv_regression import IVRegressor
+from .drivlearner import BaseDRIVLearner, BaseDRIVRegressor, XGBDRIVRegressor

causalml/source/causalml/inference/iv/drivlearner.py

@@ -0,0 +1,881 @@
+import logging
+from copy import deepcopy
+
+import numpy as np
+import pandas as pd
+from causalml.inference.meta.explainer import Explainer
+from causalml.inference.meta.utils import (
+    check_treatment_vector,
+    check_p_conditions,
+    convert_pd_to_np,
+)
+from causalml.metrics import regression_metrics
+from causalml.propensity import compute_propensity_score
+from scipy.stats import norm
+from sklearn.model_selection import KFold
+from tqdm import tqdm
+from xgboost import XGBRegressor
+
+logger = logging.getLogger("causalml")
+
+
+class BaseDRIVLearner:
+    """A parent class for DRIV-learner regressor classes.
+
+    A DRIV-learner estimates endogenous treatment effects for compliers with machine learning models.
+
+    Details of DR-learner are available at `Kennedy (2020) <https://arxiv.org/abs/2004.14497>`_.
+    The DR moment condition for LATE comes from
+    `Chernozhukov et al (2018) <https://academic.oup.com/ectj/article/21/1/C1/5056401>`_.
+    """
+
+    def __init__(
+        self,
+        learner=None,
+        control_outcome_learner=None,
+        treatment_outcome_learner=None,
+        treatment_effect_learner=None,
+        ate_alpha=0.05,
+        control_name=0,
+    ):
+        """Initialize a DR-learner.
+
+        Args:
+            learner (optional): a model to estimate outcomes and treatment effects in both the control and treatment
+                groups
+            control_outcome_learner (optional): a model to estimate outcomes in the control group
+            treatment_outcome_learner (optional): a model to estimate outcomes in the treatment group
+            treatment_effect_learner (optional): a model to estimate treatment effects in the treatment group. It needs
+                to take `sample_weight` as an input argument in `fit()`.
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+        """
+        assert (learner is not None) or (
+            (control_outcome_learner is not None)
+            and (treatment_outcome_learner is not None)
+            and (treatment_effect_learner is not None)
+        )
+
+        if control_outcome_learner is None:
+            self.model_mu_c = deepcopy(learner)
+        else:
+            self.model_mu_c = control_outcome_learner
+
+        if treatment_outcome_learner is None:
+            self.model_mu_t = deepcopy(learner)
+        else:
+            self.model_mu_t = treatment_outcome_learner
+
+        if treatment_effect_learner is None:
+            self.model_tau = deepcopy(learner)
+        else:
+            self.model_tau = treatment_effect_learner
+
+        self.ate_alpha = ate_alpha
+        self.control_name = control_name
+
+        self.propensity_1 = None
+        self.propensity_0 = None
+        self.propensity_assign = None
+
+    def __repr__(self):
+        return (
+            "{}(control_outcome_learner={},\n"
+            "\ttreatment_outcome_learner={},\n"
+            "\ttreatment_effect_learner={})".format(
+                self.__class__.__name__,
+                self.model_mu_c.__repr__(),
+                self.model_mu_t.__repr__(),
+                self.model_tau.__repr__(),
+            )
+        )
+
+    def fit(
+        self, X, assignment, treatment, y, p=None, pZ=None, seed=None, calibrate=True
+    ):
+        """Fit the inference model.
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            assignment (np.array or pd.Series): a (0,1)-valued assignment vector. The assignment is the
+                instrumental variable that does not depend on unknown confounders. The assignment status
+                influences treatment in a monotonic way, i.e. one can only be more likely to take the
+                treatment if assigned.
+            treatment (np.array or pd.Series): a treatment vector
+            y (np.array or pd.Series): an outcome vector
+            p (2-tuple of np.ndarray or pd.Series or dict, optional): The first (second) element corresponds to
+                unassigned (assigned) units. Each is an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of float
+                (0,1). If None will run ElasticNetPropensityModel() to generate the propensity scores.
+            pZ (np.array or pd.Series, optional): an array of assignment probability of float (0,1); if None
+                will run ElasticNetPropensityModel() to generate the assignment probability score.
+            seed (int): random seed for cross-fitting
+        """
+        X, treatment, assignment, y = convert_pd_to_np(X, treatment, assignment, y)
+        check_treatment_vector(treatment, self.control_name)
+        self.t_groups = np.unique(treatment[treatment != self.control_name])
+        self.t_groups.sort()
+        self._classes = {group: i for i, group in enumerate(self.t_groups)}
+
+        # The estimator splits the data into 3 partitions for cross-fit on the propensity score estimation,
+        # the outcome regression, and the treatment regression on the doubly robust estimates. The use of
+        # the partitions is rotated so we do not lose on the sample size.  We do not cross-fit the assignment
+        # score estimation as the assignment process is usually simple.
+        cv = KFold(n_splits=3, shuffle=True, random_state=seed)
+        split_indices = [index for _, index in cv.split(y)]
+
+        self.models_mu_c = {
+            group: [
+                deepcopy(self.model_mu_c),
+                deepcopy(self.model_mu_c),
+                deepcopy(self.model_mu_c),
+            ]
+            for group in self.t_groups
+        }
+        self.models_mu_t = {
+            group: [
+                deepcopy(self.model_mu_t),
+                deepcopy(self.model_mu_t),
+                deepcopy(self.model_mu_t),
+            ]
+            for group in self.t_groups
+        }
+        self.models_tau = {
+            group: [
+                deepcopy(self.model_tau),
+                deepcopy(self.model_tau),
+                deepcopy(self.model_tau),
+            ]
+            for group in self.t_groups
+        }
+
+        if p is None:
+            self.propensity_1 = {
+                group: np.zeros(y.shape[0]) for group in self.t_groups
+            }  # propensity scores for those assigned
+            self.propensity_0 = {
+                group: np.zeros(y.shape[0]) for group in self.t_groups
+            }  # propensity scores for those not assigned
+        if pZ is None:
+            self.propensity_assign, _ = compute_propensity_score(
+                X=X,
+                treatment=assignment,
+                X_pred=X,
+                treatment_pred=assignment,
+                calibrate_p=calibrate,
+            )
+        else:
+            self.propensity_assign = pZ
+
+        for ifold in range(3):
+            treatment_idx = split_indices[ifold]
+            outcome_idx = split_indices[(ifold + 1) % 3]
+            tau_idx = split_indices[(ifold + 2) % 3]
+
+            treatment_treat, treatment_out, treatment_tau = (
+                treatment[treatment_idx],
+                treatment[outcome_idx],
+                treatment[tau_idx],
+            )
+            assignment_treat, assignment_out, assignment_tau = (
+                assignment[treatment_idx],
+                assignment[outcome_idx],
+                assignment[tau_idx],
+            )
+            y_out, y_tau = y[outcome_idx], y[tau_idx]
+            X_treat, X_out, X_tau = X[treatment_idx], X[outcome_idx], X[tau_idx]
+            pZ_tau = self.propensity_assign[tau_idx]
+
+            if p is None:
+                logger.info("Generating propensity score")
+                cur_p_1 = dict()
+                cur_p_0 = dict()
+
+                for group in self.t_groups:
+                    mask = (treatment_treat == group) | (
+                        treatment_treat == self.control_name
+                    )
+                    mask_1, mask_0 = (
+                        mask & (assignment_treat == 1),
+                        mask & (assignment_treat == 0),
+                    )
+                    cur_p_1[group], _ = compute_propensity_score(
+                        X=X_treat[mask_1],
+                        treatment=(treatment_treat[mask_1] == group).astype(int),
+                        X_pred=X_tau,
+                        treatment_pred=(treatment_tau == group).astype(int),
+                    )
+                    if (treatment_treat[mask_0] == group).sum() == 0:
+                        cur_p_0[group] = np.zeros(X_tau.shape[0])
+                    else:
+                        cur_p_0[group], _ = compute_propensity_score(
+                            X=X_treat[mask_0],
+                            treatment=(treatment_treat[mask_0] == group).astype(int),
+                            X_pred=X_tau,
+                            treatment_pred=(treatment_tau == group).astype(int),
+                        )
+                    self.propensity_1[group][tau_idx] = cur_p_1[group]
+                    self.propensity_0[group][tau_idx] = cur_p_0[group]
+            else:
+                cur_p_1 = dict()
+                cur_p_0 = dict()
+                if isinstance(p[0], (np.ndarray, pd.Series)):
+                    cur_p_0 = {self.t_groups[0]: convert_pd_to_np(p[0][tau_idx])}
+                else:
+                    cur_p_0 = {g: prop[tau_idx] for g, prop in p[0].items()}
+                check_p_conditions(cur_p_0, self.t_groups)
+
+                if isinstance(p[1], (np.ndarray, pd.Series)):
+                    cur_p_1 = {self.t_groups[0]: convert_pd_to_np(p[1][tau_idx])}
+                else:
+                    cur_p_1 = {g: prop[tau_idx] for g, prop in p[1].items()}
+                check_p_conditions(cur_p_1, self.t_groups)
+
+            logger.info("Generate outcome regressions")
+            for group in self.t_groups:
+                mask = (treatment_out == group) | (treatment_out == self.control_name)
+                mask_1, mask_0 = (
+                    mask & (assignment_out == 1),
+                    mask & (assignment_out == 0),
+                )
+                self.models_mu_c[group][ifold].fit(X_out[mask_0], y_out[mask_0])
+                self.models_mu_t[group][ifold].fit(X_out[mask_1], y_out[mask_1])
+
+            logger.info("Fit pseudo outcomes from the DR formula")
+
+            for group in self.t_groups:
+                mask = (treatment_tau == group) | (treatment_tau == self.control_name)
+                treatment_filt = treatment_tau[mask]
+                X_filt = X_tau[mask]
+                y_filt = y_tau[mask]
+                w_filt = (treatment_filt == group).astype(int)
+                p_1_filt = cur_p_1[group][mask]
+                p_0_filt = cur_p_0[group][mask]
+                z_filt = assignment_tau[mask]
+                pZ_filt = pZ_tau[mask]
+                mu_t = self.models_mu_t[group][ifold].predict(X_filt)
+                mu_c = self.models_mu_c[group][ifold].predict(X_filt)
+                dr = (
+                    z_filt * (y_filt - mu_t) / pZ_filt
+                    - (1 - z_filt) * (y_filt - mu_c) / (1 - pZ_filt)
+                    + mu_t
+                    - mu_c
+                )
+                weight = (
+                    z_filt * (w_filt - p_1_filt) / pZ_filt
+                    - (1 - z_filt) * (w_filt - p_0_filt) / (1 - pZ_filt)
+                    + p_1_filt
+                    - p_0_filt
+                )
+                dr /= weight
+                self.models_tau[group][ifold].fit(X_filt, dr, sample_weight=weight**2)
+
+    def predict(self, X, treatment=None, y=None, return_components=False, verbose=True):
+        """Predict treatment effects.
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series, optional): a treatment vector
+            y (np.array or pd.Series, optional): an outcome vector
+            verbose (bool, optional): whether to output progress logs
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects for compliers, i.e. those individuals
+                who take the treatment only if they are assigned.
+        """
+        X, treatment, y = convert_pd_to_np(X, treatment, y)
+
+        te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+        yhat_cs = {}
+        yhat_ts = {}
+
+        for i, group in enumerate(self.t_groups):
+            models_tau = self.models_tau[group]
+            _te = np.r_[[model.predict(X) for model in models_tau]].mean(axis=0)
+            te[:, i] = np.ravel(_te)
+            yhat_cs[group] = np.r_[
+                [model.predict(X) for model in self.models_mu_c[group]]
+            ].mean(axis=0)
+            yhat_ts[group] = np.r_[
+                [model.predict(X) for model in self.models_mu_t[group]]
+            ].mean(axis=0)
+
+            if (y is not None) and (treatment is not None) and verbose:
+                mask = (treatment == group) | (treatment == self.control_name)
+                treatment_filt = treatment[mask]
+                y_filt = y[mask]
+                w = (treatment_filt == group).astype(int)
+
+                yhat = np.zeros_like(y_filt, dtype=float)
+                yhat[w == 0] = yhat_cs[group][mask][w == 0]
+                yhat[w == 1] = yhat_ts[group][mask][w == 1]
+
+                logger.info("Error metrics for group {}".format(group))
+                regression_metrics(y_filt, yhat, w)
+
+        if not return_components:
+            return te
+        else:
+            return te, yhat_cs, yhat_ts
+
+    def fit_predict(
+        self,
+        X,
+        assignment,
+        treatment,
+        y,
+        p=None,
+        pZ=None,
+        return_ci=False,
+        n_bootstraps=1000,
+        bootstrap_size=10000,
+        return_components=False,
+        verbose=True,
+        seed=None,
+        calibrate=True,
+    ):
+        """Fit the treatment effect and outcome models of the R learner and predict treatment effects.
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            assignment (np.array or pd.Series): a (0,1)-valued assignment vector. The assignment is the
+                instrumental variable that does not depend on unknown confounders. The assignment status
+                influences treatment in a monotonic way, i.e. one can only be more likely to take the
+                treatment if assigned.
+            treatment (np.array or pd.Series): a treatment vector
+            y (np.array or pd.Series): an outcome vector
+            p (2-tuple of np.ndarray or pd.Series or dict, optional): The first (second) element corresponds to
+                unassigned (assigned) units. Each is an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of float
+                (0,1). If None will run ElasticNetPropensityModel() to generate the propensity scores.
+            pZ (np.array or pd.Series, optional): an array of assignment probability of float (0,1); if None
+                will run ElasticNetPropensityModel() to generate the assignment probability score.
+            return_ci (bool): whether to return confidence intervals
+            n_bootstraps (int): number of bootstrap iterations
+            bootstrap_size (int): number of samples per bootstrap
+            return_components (bool, optional): whether to return outcome for treatment and control seperately
+            verbose (str): whether to output progress logs
+            seed (int): random seed for cross-fitting
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects for compliers, , i.e. those individuals
+                who take the treatment only if they are assigned. Output dim: [n_samples, n_treatment]
+                If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment],
+                UB [n_samples, n_treatment]
+        """
+        X, assignment, treatment, y = convert_pd_to_np(X, assignment, treatment, y)
+        self.fit(X, assignment, treatment, y, p, seed, calibrate)
+
+        if p is None:
+            p = (self.propensity_0, self.propensity_1)
+        else:
+            check_p_conditions(p[0], self.t_groups)
+            check_p_conditions(p[1], self.t_groups)
+
+        if isinstance(p[0], (np.ndarray, pd.Series)):
+            treatment_name = self.t_groups[0]
+            p = (
+                {treatment_name: convert_pd_to_np(p[0])},
+                {treatment_name: convert_pd_to_np(p[1])},
+            )
+        elif isinstance(p[0], dict):
+            p = (
+                {
+                    treatment_name: convert_pd_to_np(_p)
+                    for treatment_name, _p in p[0].items()
+                },
+                {
+                    treatment_name: convert_pd_to_np(_p)
+                    for treatment_name, _p in p[1].items()
+                },
+            )
+
+        if pZ is None:
+            pZ = self.propensity_assign
+
+        te = self.predict(
+            X, treatment=treatment, y=y, return_components=return_components
+        )
+
+        if not return_ci:
+            return te
+        else:
+            t_groups_global = self.t_groups
+            _classes_global = self._classes
+            models_mu_c_global = deepcopy(self.models_mu_c)
+            models_mu_t_global = deepcopy(self.models_mu_t)
+            models_tau_global = deepcopy(self.models_tau)
+            te_bootstraps = np.zeros(
+                shape=(X.shape[0], self.t_groups.shape[0], n_bootstraps)
+            )
+
+            logger.info("Bootstrap Confidence Intervals")
+            for i in tqdm(range(n_bootstraps)):
+                te_b = self.bootstrap(
+                    X, assignment, treatment, y, p, pZ, size=bootstrap_size, seed=seed
+                )
+                te_bootstraps[:, :, i] = te_b
+
+            te_lower = np.percentile(te_bootstraps, (self.ate_alpha / 2) * 100, axis=2)
+            te_upper = np.percentile(
+                te_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=2
+            )
+
+            # set member variables back to global (currently last bootstrapped outcome)
+            self.t_groups = t_groups_global
+            self._classes = _classes_global
+            self.models_mu_c = deepcopy(models_mu_c_global)
+            self.models_mu_t = deepcopy(models_mu_t_global)
+            self.models_tau = deepcopy(models_tau_global)
+
+            return (te, te_lower, te_upper)
+
+    def estimate_ate(
+        self,
+        X,
+        assignment,
+        treatment,
+        y,
+        p=None,
+        pZ=None,
+        bootstrap_ci=False,
+        n_bootstraps=1000,
+        bootstrap_size=10000,
+        seed=None,
+        calibrate=True,
+    ):
+        """Estimate the Average Treatment Effect (ATE) for compliers.
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            assignment (np.array or pd.Series): an assignment vector. The assignment is the
+                instrumental variable that does not depend on unknown confounders. The assignment status
+                influences treatment in a monotonic way, i.e. one can only be more likely to take the
+                treatment if assigned.
+            treatment (np.array or pd.Series): a treatment vector
+            y (np.array or pd.Series): an outcome vector
+            p (2-tuple of np.ndarray or pd.Series or dict, optional): The first (second) element corresponds to
+                unassigned (assigned) units. Each is an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of float
+                (0,1). If None will run ElasticNetPropensityModel() to generate the propensity scores.
+            pZ (np.array or pd.Series, optional): an array of assignment probability of float (0,1); if None
+                will run ElasticNetPropensityModel() to generate the assignment probability score.
+            bootstrap_ci (bool): whether run bootstrap for confidence intervals
+            n_bootstraps (int): number of bootstrap iterations
+            bootstrap_size (int): number of samples per bootstrap
+            seed (int): random seed for cross-fitting
+        Returns:
+            The mean and confidence interval (LB, UB) of the ATE estimate.
+        """
+        te, yhat_cs, yhat_ts = self.fit_predict(
+            X,
+            assignment,
+            treatment,
+            y,
+            p,
+            return_components=True,
+            seed=seed,
+            calibrate=calibrate,
+        )
+        X, assignment, treatment, y = convert_pd_to_np(X, assignment, treatment, y)
+
+        if p is None:
+            p = (self.propensity_0, self.propensity_1)
+        else:
+            check_p_conditions(p[0], self.t_groups)
+            check_p_conditions(p[1], self.t_groups)
+
+        if isinstance(p[0], (np.ndarray, pd.Series)):
+            treatment_name = self.t_groups[0]
+            p = (
+                {treatment_name: convert_pd_to_np(p[0])},
+                {treatment_name: convert_pd_to_np(p[1])},
+            )
+        elif isinstance(p[0], dict):
+            p = (
+                {
+                    treatment_name: convert_pd_to_np(_p)
+                    for treatment_name, _p in p[0].items()
+                },
+                {
+                    treatment_name: convert_pd_to_np(_p)
+                    for treatment_name, _p in p[1].items()
+                },
+            )
+
+        ate = np.zeros(self.t_groups.shape[0])
+        ate_lb = np.zeros(self.t_groups.shape[0])
+        ate_ub = np.zeros(self.t_groups.shape[0])
+
+        for i, group in enumerate(self.t_groups):
+            _ate = te[:, i].mean()
+
+            mask = (treatment == group) | (treatment == self.control_name)
+            mask_1, mask_0 = mask & (assignment == 1), mask & (assignment == 0)
+            Gamma = (treatment[mask_1] == group).mean() - (
+                treatment[mask_0] == group
+            ).mean()
+
+            y_filt_1, y_filt_0 = y[mask_1], y[mask_0]
+            yhat_0 = yhat_cs[group][mask_0]
+            yhat_1 = yhat_ts[group][mask_1]
+            treatment_filt_1, treatment_filt_0 = treatment[mask_1], treatment[mask_0]
+            prob_treatment_1, prob_treatment_0 = (
+                p[1][group][mask_1],
+                p[0][group][mask_0],
+            )
+            w = (assignment[mask]).mean()
+
+            part_1 = (
+                (y_filt_1 - yhat_1).var()
+                + _ate**2 * (treatment_filt_1 - prob_treatment_1).var()
+                - 2
+                * _ate
+                * (y_filt_1 * treatment_filt_1 - yhat_1 * prob_treatment_1).mean()
+            )
+            part_0 = (
+                (y_filt_0 - yhat_0).var()
+                + _ate**2 * (treatment_filt_0 - prob_treatment_0).var()
+                - 2
+                * _ate
+                * (y_filt_0 * treatment_filt_0 - yhat_0 * prob_treatment_0).mean()
+            )
+            part_2 = np.mean(
+                (
+                    yhat_ts[group][mask]
+                    - yhat_cs[group][mask]
+                    - _ate * (p[1][group][mask] - p[0][group][mask])
+                )
+                ** 2
+            )
+
+            # SE formula is based on the lower bound formula (9) from Frölich, Markus. 2006.
+            # "Nonparametric IV estimation of local average treatment effects wth covariates."
+            # Journal of Econometrics.
+            se = np.sqrt((part_1 / w + part_0 / (1 - w)) + part_2) / Gamma
+
+            _ate_lb = _ate - se * norm.ppf(1 - self.ate_alpha / 2)
+            _ate_ub = _ate + se * norm.ppf(1 - self.ate_alpha / 2)
+
+            ate[i] = _ate
+            ate_lb[i] = _ate_lb
+            ate_ub[i] = _ate_ub
+
+        if not bootstrap_ci:
+            return ate, ate_lb, ate_ub
+        else:
+            t_groups_global = self.t_groups
+            _classes_global = self._classes
+            models_mu_c_global = deepcopy(self.models_mu_c)
+            models_mu_t_global = deepcopy(self.models_mu_t)
+            models_tau_global = deepcopy(self.models_tau)
+
+            logger.info("Bootstrap Confidence Intervals for ATE")
+            ate_bootstraps = np.zeros(shape=(self.t_groups.shape[0], n_bootstraps))
+
+            for n in tqdm(range(n_bootstraps)):
+                cate_b = self.bootstrap(
+                    X, assignment, treatment, y, p, pZ, size=bootstrap_size, seed=seed
+                )
+                ate_bootstraps[:, n] = cate_b.mean()
+
+            ate_lower = np.percentile(
+                ate_bootstraps, (self.ate_alpha / 2) * 100, axis=1
+            )
+            ate_upper = np.percentile(
+                ate_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=1
+            )
+
+            # set member variables back to global (currently last bootstrapped outcome)
+            self.t_groups = t_groups_global
+            self._classes = _classes_global
+            self.models_mu_c = deepcopy(models_mu_c_global)
+            self.models_mu_t = deepcopy(models_mu_t_global)
+            self.models_tau = deepcopy(models_tau_global)
+            return ate, ate_lower, ate_upper
+
+    def bootstrap(self, X, assignment, treatment, y, p, pZ, size=10000, seed=None):
+        """Runs a single bootstrap. Fits on bootstrapped sample, then predicts on whole population."""
+        idxs = np.random.choice(np.arange(0, X.shape[0]), size=size)
+        X_b = X[idxs]
+
+        if isinstance(p[0], (np.ndarray, pd.Series)):
+            p0_b = {self.t_groups[0]: convert_pd_to_np(p[0][idxs])}
+        else:
+            p0_b = {g: prop[idxs] for g, prop in p[0].items()}
+        if isinstance(p[1], (np.ndarray, pd.Series)):
+            p1_b = {self.t_groups[0]: convert_pd_to_np(p[1][idxs])}
+        else:
+            p1_b = {g: prop[idxs] for g, prop in p[1].items()}
+
+        pZ_b = pZ[idxs]
+        assignment_b = assignment[idxs]
+        treatment_b = treatment[idxs]
+        y_b = y[idxs]
+        self.fit(
+            X=X_b,
+            assignment=assignment_b,
+            treatment=treatment_b,
+            y=y_b,
+            p=(p0_b, p1_b),
+            pZ=pZ_b,
+            seed=seed,
+        )
+        te_b = self.predict(X=X)
+        return te_b
+
+    def get_importance(
+        self,
+        X=None,
+        tau=None,
+        model_tau_feature=None,
+        features=None,
+        method="auto",
+        normalize=True,
+        test_size=0.3,
+        random_state=None,
+    ):
+        """
+        Builds a model (using X to predict estimated/actual tau), and then calculates feature importances
+        based on a specified method.
+
+        Currently supported methods are:
+            - auto (calculates importance based on estimator's default implementation of feature importance;
+                    estimator must be tree-based)
+                    Note: if none provided, it uses lightgbm's LGBMRegressor as estimator, and "gain" as
+                    importance type
+            - permutation (calculates importance based on mean decrease in accuracy when a feature column is permuted;
+                           estimator can be any form)
+        Hint: for permutation, downsample data for better performance especially if X.shape[1] is large
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            tau (np.array): a treatment effect vector (estimated/actual)
+            model_tau_feature (sklearn/lightgbm/xgboost model object): an unfitted model object
+            features (np.array): list/array of feature names. If None, an enumerated list will be used
+            method (str): auto, permutation
+            normalize (bool): normalize by sum of importances if method=auto (defaults to True)
+            test_size (float/int): if float, represents the proportion of the dataset to include in the test split.
+                                   If int, represents the absolute number of test samples (used for estimating
+                                   permutation importance)
+            random_state (int/RandomState instance/None): random state used in permutation importance estimation
+        """
+        explainer = Explainer(
+            method=method,
+            control_name=self.control_name,
+            X=X,
+            tau=tau,
+            model_tau=model_tau_feature,
+            features=features,
+            classes=self._classes,
+            normalize=normalize,
+            test_size=test_size,
+            random_state=random_state,
+        )
+        return explainer.get_importance()
+
+    def get_shap_values(self, X=None, model_tau_feature=None, tau=None, features=None):
+        """
+        Builds a model (using X to predict estimated/actual tau), and then calculates shapley values.
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            tau (np.array): a treatment effect vector (estimated/actual)
+            model_tau_feature (sklearn/lightgbm/xgboost model object): an unfitted model object
+            features (optional, np.array): list/array of feature names. If None, an enumerated list will be used.
+        """
+        explainer = Explainer(
+            method="shapley",
+            control_name=self.control_name,
+            X=X,
+            tau=tau,
+            model_tau=model_tau_feature,
+            features=features,
+            classes=self._classes,
+        )
+        return explainer.get_shap_values()
+
+    def plot_importance(
+        self,
+        X=None,
+        tau=None,
+        model_tau_feature=None,
+        features=None,
+        method="auto",
+        normalize=True,
+        test_size=0.3,
+        random_state=None,
+    ):
+        """
+        Builds a model (using X to predict estimated/actual tau), and then plots feature importances
+        based on a specified method.
+
+        Currently supported methods are:
+            - auto (calculates importance based on estimator's default implementation of feature importance;
+                    estimator must be tree-based)
+                    Note: if none provided, it uses lightgbm's LGBMRegressor as estimator, and "gain" as
+                    importance type
+            - permutation (calculates importance based on mean decrease in accuracy when a feature column is permuted;
+                           estimator can be any form)
+        Hint: for permutation, downsample data for better performance especially if X.shape[1] is large
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            tau (np.array): a treatment effect vector (estimated/actual)
+            model_tau_feature (sklearn/lightgbm/xgboost model object): an unfitted model object
+            features (optional, np.array): list/array of feature names. If None, an enumerated list will be used
+            method (str): auto, permutation
+            normalize (bool): normalize by sum of importances if method=auto (defaults to True)
+            test_size (float/int): if float, represents the proportion of the dataset to include in the test split.
+                                   If int, represents the absolute number of test samples (used for estimating
+                                   permutation importance)
+            random_state (int/RandomState instance/None): random state used in permutation importance estimation
+        """
+        explainer = Explainer(
+            method=method,
+            control_name=self.control_name,
+            X=X,
+            tau=tau,
+            model_tau=model_tau_feature,
+            features=features,
+            classes=self._classes,
+            normalize=normalize,
+            test_size=test_size,
+            random_state=random_state,
+        )
+        explainer.plot_importance()
+
+    def plot_shap_values(
+        self,
+        X=None,
+        tau=None,
+        model_tau_feature=None,
+        features=None,
+        shap_dict=None,
+        **kwargs,
+    ):
+        """
+        Plots distribution of shapley values.
+
+        If shapley values have been pre-computed, pass it through the shap_dict parameter.
+        If shap_dict is not provided, this builds a new model (using X to predict estimated/actual tau),
+        and then calculates shapley values.
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix. Required if shap_dict is None.
+            tau (np.array): a treatment effect vector (estimated/actual)
+            model_tau_feature (sklearn/lightgbm/xgboost model object): an unfitted model object
+            features (optional, np.array): list/array of feature names. If None, an enumerated list will be used.
+            shap_dict (optional, dict): a dict of shapley value matrices. If None, shap_dict will be computed.
+        """
+        override_checks = False if shap_dict is None else True
+        explainer = Explainer(
+            method="shapley",
+            control_name=self.control_name,
+            X=X,
+            tau=tau,
+            model_tau=model_tau_feature,
+            features=features,
+            override_checks=override_checks,
+            classes=self._classes,
+        )
+        explainer.plot_shap_values(shap_dict=shap_dict)
+
+    def plot_shap_dependence(
+        self,
+        treatment_group,
+        feature_idx,
+        X,
+        tau,
+        model_tau_feature=None,
+        features=None,
+        shap_dict=None,
+        interaction_idx="auto",
+        **kwargs,
+    ):
+        """
+        Plots dependency of shapley values for a specified feature, colored by an interaction feature.
+
+        If shapley values have been pre-computed, pass it through the shap_dict parameter.
+        If shap_dict is not provided, this builds a new model (using X to predict estimated/actual tau),
+        and then calculates shapley values.
+
+        This plots the value of the feature on the x-axis and the SHAP value of the same feature
+        on the y-axis. This shows how the model depends on the given feature, and is like a
+        richer extension of the classical partial dependence plots. Vertical dispersion of the
+        data points represents interaction effects.
+
+        Args:
+            treatment_group (str or int): name of treatment group to create dependency plot on
+            feature_idx (str or int): feature index / name to create dependency plot on
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            tau (np.array): a treatment effect vector (estimated/actual)
+            model_tau_feature (sklearn/lightgbm/xgboost model object): an unfitted model object
+            features (optional, np.array): list/array of feature names. If None, an enumerated list will be used.
+            shap_dict (optional, dict): a dict of shapley value matrices. If None, shap_dict will be computed.
+            interaction_idx (optional, str or int): feature index / name used in coloring scheme as interaction feature.
+                If "auto" then shap.common.approximate_interactions is used to pick what seems to be the
+                strongest interaction (note that to find to true strongest interaction you need to compute
+                the SHAP interaction values).
+        """
+        override_checks = False if shap_dict is None else True
+        explainer = Explainer(
+            method="shapley",
+            control_name=self.control_name,
+            X=X,
+            tau=tau,
+            model_tau=model_tau_feature,
+            features=features,
+            override_checks=override_checks,
+            classes=self._classes,
+        )
+        explainer.plot_shap_dependence(
+            treatment_group=treatment_group,
+            feature_idx=feature_idx,
+            shap_dict=shap_dict,
+            interaction_idx=interaction_idx,
+            **kwargs,
+        )
+
+
+class BaseDRIVRegressor(BaseDRIVLearner):
+    """
+    A parent class for DRIV-learner regressor classes.
+    """
+
+    def __init__(
+        self,
+        learner=None,
+        control_outcome_learner=None,
+        treatment_outcome_learner=None,
+        treatment_effect_learner=None,
+        ate_alpha=0.05,
+        control_name=0,
+    ):
+        """Initialize a DRIV-learner regressor.
+
+        Args:
+            learner (optional): a model to estimate outcomes and treatment effects in both the control and treatment
+                groups
+            control_outcome_learner (optional): a model to estimate outcomes in the control group
+            treatment_outcome_learner (optional): a model to estimate outcomes in the treatment group
+            treatment_effect_learner (optional): a model to estimate treatment effects in the treatment group. It needs
+                to take `sample_weight` as an input argument in `fit()`.
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+        """
+        super().__init__(
+            learner=learner,
+            control_outcome_learner=control_outcome_learner,
+            treatment_outcome_learner=treatment_outcome_learner,
+            treatment_effect_learner=treatment_effect_learner,
+            ate_alpha=ate_alpha,
+            control_name=control_name,
+        )
+
+
+class XGBDRIVRegressor(BaseDRIVRegressor):
+    def __init__(self, ate_alpha=0.05, control_name=0, *args, **kwargs):
+        """Initialize a DRIV-learner with two XGBoost models."""
+        super().__init__(
+            learner=XGBRegressor(*args, **kwargs),
+            ate_alpha=ate_alpha,
+            control_name=control_name,
+        )

causalml/source/causalml/inference/iv/iv_regression.py

@@ -0,0 +1,48 @@
+import numpy as np
+
+from causalml.inference.meta.utils import convert_pd_to_np
+import statsmodels.api as sm
+from statsmodels.sandbox.regression.gmm import IV2SLS
+
+
+class IVRegressor:
+    """A wrapper class that uses IV2SLS from statsmodel
+
+    A linear 2SLS model that estimates the average treatment effect with endogenous treatment variable.
+    """
+
+    def __init__(self):
+        """
+        Initializes the class.
+        """
+
+        self.method = "2SLS"
+
+    def fit(self, X, treatment, y, w):
+        """Fits the 2SLS model.
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series): a treatment vector
+            y (np.array or pd.Series): an outcome vector
+            w (np.array or pd.Series): an instrument vector
+        """
+
+        X, treatment, y, w = convert_pd_to_np(X, treatment, y, w)
+
+        exog = sm.add_constant(np.c_[X, treatment])
+        endog = y
+        instrument = sm.add_constant(np.c_[X, w])
+
+        self.iv_model = IV2SLS(endog=endog, exog=exog, instrument=instrument)
+        self.iv_fit = self.iv_model.fit()
+
+    def predict(self):
+        """Returns the average treatment effect and its estimated standard error
+
+        Returns:
+            (float): average treatment effect
+            (float): standard error of the estimation
+        """
+
+        return self.iv_fit.params[-1], self.iv_fit.bse[-1]

causalml/source/causalml/inference/meta/__init__.py

@@ -0,0 +1,12 @@
+from .slearner import LRSRegressor, BaseSLearner, BaseSRegressor, BaseSClassifier
+from .tlearner import (
+    XGBTRegressor,
+    MLPTRegressor,
+    BaseTLearner,
+    BaseTRegressor,
+    BaseTClassifier,
+)
+from .xlearner import BaseXLearner, BaseXRegressor, BaseXClassifier
+from .rlearner import BaseRLearner, BaseRRegressor, BaseRClassifier, XGBRRegressor
+from .tmle import TMLELearner
+from .drlearner import BaseDRLearner, BaseDRRegressor, BaseDRClassifier, XGBDRRegressor

causalml/source/causalml/inference/meta/base.py

@@ -0,0 +1,337 @@
+from abc import ABCMeta, abstractmethod
+import logging
+import numpy as np
+import pandas as pd
+
+from causalml.inference.meta.explainer import Explainer
+from causalml.inference.meta.utils import check_p_conditions, convert_pd_to_np
+from causalml.propensity import compute_propensity_score
+
+logger = logging.getLogger("causalml")
+
+
+class BaseLearner(metaclass=ABCMeta):
+    @classmethod
+    @abstractmethod
+    def fit(self, X, treatment, y, p=None):
+        pass
+
+    @classmethod
+    @abstractmethod
+    def predict(
+        self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
+    ):
+        pass
+
+    def fit_predict(
+        self,
+        X,
+        treatment,
+        y,
+        p=None,
+        return_ci=False,
+        n_bootstraps=1000,
+        bootstrap_size=10000,
+        return_components=False,
+        verbose=True,
+    ):
+        self.fit(X, treatment, y, p)
+        return self.predict(X, treatment, y, p, return_components, verbose)
+
+    @classmethod
+    @abstractmethod
+    def estimate_ate(
+        self,
+        X,
+        treatment,
+        y,
+        p=None,
+        bootstrap_ci=False,
+        n_bootstraps=1000,
+        bootstrap_size=10000,
+    ):
+        pass
+
+    def bootstrap(self, X, treatment, y, p=None, size=10000):
+        """Runs a single bootstrap. Fits on bootstrapped sample, then predicts on whole population."""
+        idxs = np.random.choice(np.arange(0, X.shape[0]), size=size)
+        X_b = X[idxs]
+
+        if p is not None:
+            p_b = {group: _p[idxs] for group, _p in p.items()}
+        else:
+            p_b = None
+
+        treatment_b = treatment[idxs]
+        y_b = y[idxs]
+        self.fit(X=X_b, treatment=treatment_b, y=y_b, p=p_b)
+        return self.predict(X=X, p=p)
+
+    @staticmethod
+    def _format_p(p, t_groups):
+        """Format propensity scores into a dictionary of {treatment group: propensity scores}.
+
+        Args:
+            p (np.ndarray, pd.Series, or dict): propensity scores
+            t_groups (list): treatment group names.
+
+        Returns:
+            dict of {treatment group: propensity scores}
+        """
+        check_p_conditions(p, t_groups)
+
+        if isinstance(p, (np.ndarray, pd.Series)):
+            treatment_name = t_groups[0]
+            p = {treatment_name: convert_pd_to_np(p)}
+        elif isinstance(p, dict):
+            p = {
+                treatment_name: convert_pd_to_np(_p) for treatment_name, _p in p.items()
+            }
+
+        return p
+
+    def _set_propensity_models(self, X, treatment, y):
+        """Set self.propensity and self.propensity_models.
+
+        It trains propensity models for all treatment groups, save them in self.propensity_models, and
+        save propensity scores in self.propensity in dictionaries with treatment groups as keys.
+
+        It will use self.model_p if available to train propensity models. Otherwise, it will use a default
+        PropensityModel (i.e. ElasticNetPropensityModel).
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series): a treatment vector
+            y (np.array or pd.Series): an outcome vector
+        """
+        logger.info("Generating propensity score")
+        p = dict()
+        p_model = dict()
+        for group in self.t_groups:
+            mask = (treatment == group) | (treatment == self.control_name)
+            treatment_filt = treatment[mask]
+            X_filt = X[mask]
+            w_filt = (treatment_filt == group).astype(int)
+            w = (treatment == group).astype(int)
+            propensity_model = self.model_p if hasattr(self, "model_p") else None
+            p[group], p_model[group] = compute_propensity_score(
+                X=X_filt,
+                treatment=w_filt,
+                p_model=propensity_model,
+                X_pred=X,
+                treatment_pred=w,
+            )
+        self.propensity_model = p_model
+        self.propensity = p
+
+    def get_importance(
+        self,
+        X=None,
+        tau=None,
+        model_tau_feature=None,
+        features=None,
+        method="auto",
+        normalize=True,
+        test_size=0.3,
+        random_state=None,
+    ):
+        """
+        Builds a model (using X to predict estimated/actual tau), and then calculates feature importances
+        based on a specified method.
+
+        Currently supported methods are:
+            - auto (calculates importance based on estimator's default implementation of feature importance;
+                    estimator must be tree-based)
+                    Note: if none provided, it uses lightgbm's LGBMRegressor as estimator, and "gain" as
+                    importance type
+            - permutation (calculates importance based on mean decrease in accuracy when a feature column is permuted;
+                           estimator can be any form)
+        Hint: for permutation, downsample data for better performance especially if X.shape[1] is large
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            tau (np.array): a treatment effect vector (estimated/actual)
+            model_tau_feature (sklearn/lightgbm/xgboost model object): an unfitted model object
+            features (np.array): list/array of feature names. If None, an enumerated list will be used
+            method (str): auto, permutation
+            normalize (bool): normalize by sum of importances if method=auto (defaults to True)
+            test_size (float/int): if float, represents the proportion of the dataset to include in the test split.
+                                   If int, represents the absolute number of test samples (used for estimating
+                                   permutation importance)
+            random_state (int/RandomState instance/None): random state used in permutation importance estimation
+        """
+        explainer = Explainer(
+            method=method,
+            control_name=self.control_name,
+            X=X,
+            tau=tau,
+            model_tau=model_tau_feature,
+            features=features,
+            classes=self._classes,
+            normalize=normalize,
+            test_size=test_size,
+            random_state=random_state,
+        )
+        return explainer.get_importance()
+
+    def get_shap_values(self, X=None, model_tau_feature=None, tau=None, features=None):
+        """
+        Builds a model (using X to predict estimated/actual tau), and then calculates shapley values.
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            tau (np.array): a treatment effect vector (estimated/actual)
+            model_tau_feature (sklearn/lightgbm/xgboost model object): an unfitted model object
+            features (optional, np.array): list/array of feature names. If None, an enumerated list will be used.
+        """
+        explainer = Explainer(
+            method="shapley",
+            control_name=self.control_name,
+            X=X,
+            tau=tau,
+            model_tau=model_tau_feature,
+            features=features,
+            classes=self._classes,
+        )
+        return explainer.get_shap_values()
+
+    def plot_importance(
+        self,
+        X=None,
+        tau=None,
+        model_tau_feature=None,
+        features=None,
+        method="auto",
+        normalize=True,
+        test_size=0.3,
+        random_state=None,
+    ):
+        """
+        Builds a model (using X to predict estimated/actual tau), and then plots feature importances
+        based on a specified method.
+
+        Currently supported methods are:
+            - auto (calculates importance based on estimator's default implementation of feature importance;
+                    estimator must be tree-based)
+                    Note: if none provided, it uses lightgbm's LGBMRegressor as estimator, and "gain" as
+                    importance type
+            - permutation (calculates importance based on mean decrease in accuracy when a feature column is permuted;
+                           estimator can be any form)
+        Hint: for permutation, downsample data for better performance especially if X.shape[1] is large
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            tau (np.array): a treatment effect vector (estimated/actual)
+            model_tau_feature (sklearn/lightgbm/xgboost model object): an unfitted model object
+            features (optional, np.array): list/array of feature names. If None, an enumerated list will be used
+            method (str): auto, permutation
+            normalize (bool): normalize by sum of importances if method=auto (defaults to True)
+            test_size (float/int): if float, represents the proportion of the dataset to include in the test split.
+                                   If int, represents the absolute number of test samples (used for estimating
+                                   permutation importance)
+            random_state (int/RandomState instance/None): random state used in permutation importance estimation
+        """
+        explainer = Explainer(
+            method=method,
+            control_name=self.control_name,
+            X=X,
+            tau=tau,
+            model_tau=model_tau_feature,
+            features=features,
+            classes=self._classes,
+            normalize=normalize,
+            test_size=test_size,
+            random_state=random_state,
+        )
+        explainer.plot_importance()
+
+    def plot_shap_values(
+        self,
+        X=None,
+        tau=None,
+        model_tau_feature=None,
+        features=None,
+        shap_dict=None,
+        **kwargs,
+    ):
+        """
+        Plots distribution of shapley values.
+
+        If shapley values have been pre-computed, pass it through the shap_dict parameter.
+        If shap_dict is not provided, this builds a new model (using X to predict estimated/actual tau),
+        and then calculates shapley values.
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix. Required if shap_dict is None.
+            tau (np.array): a treatment effect vector (estimated/actual)
+            model_tau_feature (sklearn/lightgbm/xgboost model object): an unfitted model object
+            features (optional, np.array): list/array of feature names. If None, an enumerated list will be used.
+            shap_dict (optional, dict): a dict of shapley value matrices. If None, shap_dict will be computed.
+        """
+        override_checks = shap_dict is not None
+        explainer = Explainer(
+            method="shapley",
+            control_name=self.control_name,
+            X=X,
+            tau=tau,
+            model_tau=model_tau_feature,
+            features=features,
+            override_checks=override_checks,
+            classes=self._classes,
+        )
+        explainer.plot_shap_values(shap_dict=shap_dict, **kwargs)
+
+    def plot_shap_dependence(
+        self,
+        treatment_group,
+        feature_idx,
+        X,
+        tau,
+        model_tau_feature=None,
+        features=None,
+        shap_dict=None,
+        interaction_idx="auto",
+        **kwargs,
+    ):
+        """
+        Plots dependency of shapley values for a specified feature, colored by an interaction feature.
+
+        If shapley values have been pre-computed, pass it through the shap_dict parameter.
+        If shap_dict is not provided, this builds a new model (using X to predict estimated/actual tau),
+        and then calculates shapley values.
+
+        This plots the value of the feature on the x-axis and the SHAP value of the same feature
+        on the y-axis. This shows how the model depends on the given feature, and is like a
+        richer extension of the classical partial dependence plots. Vertical dispersion of the
+        data points represents interaction effects.
+
+        Args:
+            treatment_group (str or int): name of treatment group to create dependency plot on
+            feature_idx (str or int): feature index / name to create dependency plot on
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            tau (np.array): a treatment effect vector (estimated/actual)
+            model_tau_feature (sklearn/lightgbm/xgboost model object): an unfitted model object
+            features (optional, np.array): list/array of feature names. If None, an enumerated list will be used.
+            shap_dict (optional, dict): a dict of shapley value matrices. If None, shap_dict will be computed.
+            interaction_idx (optional, str or int): feature index / name used in coloring scheme as interaction feature.
+                If "auto" then shap.common.approximate_interactions is used to pick what seems to be the
+                strongest interaction (note that to find to true strongest interaction you need to compute
+                the SHAP interaction values).
+        """
+        override_checks = False if shap_dict is None else True
+        explainer = Explainer(
+            method="shapley",
+            control_name=self.control_name,
+            X=X,
+            tau=tau,
+            model_tau=model_tau_feature,
+            features=features,
+            override_checks=override_checks,
+            classes=self._classes,
+        )
+        explainer.plot_shap_dependence(
+            treatment_group=treatment_group,
+            feature_idx=feature_idx,
+            shap_dict=shap_dict,
+            interaction_idx=interaction_idx,
+            **kwargs,
+        )

causalml/source/causalml/inference/meta/drlearner.py

@@ -0,0 +1,592 @@
+from copy import deepcopy
+import logging
+import numpy as np
+import pandas as pd
+from scipy.stats import norm
+from sklearn.model_selection import KFold
+from tqdm import tqdm
+from xgboost import XGBRegressor
+
+from causalml.inference.meta.base import BaseLearner
+from causalml.inference.meta.utils import (
+    check_treatment_vector,
+    check_p_conditions,
+    convert_pd_to_np,
+)
+from causalml.metrics import regression_metrics, classification_metrics
+from causalml.propensity import compute_propensity_score
+
+logger = logging.getLogger("causalml")
+
+
+class BaseDRLearner(BaseLearner):
+    """A parent class for DR-learner regressor classes.
+
+    A DR-learner estimates treatment effects with machine learning models.
+
+    Details of DR-learner are available at `Kennedy (2020) <https://arxiv.org/abs/2004.14497>`_.
+    """
+
+    def __init__(
+        self,
+        learner=None,
+        control_outcome_learner=None,
+        treatment_outcome_learner=None,
+        treatment_effect_learner=None,
+        ate_alpha=0.05,
+        control_name=0,
+    ):
+        """Initialize a DR-learner.
+
+        Args:
+            learner (optional): a model to estimate outcomes and treatment effects in both the control and treatment
+                groups
+            control_outcome_learner (optional): a model to estimate outcomes in the control group
+            treatment_outcome_learner (optional): a model to estimate outcomes in the treatment group
+            treatment_effect_learner (optional): a model to estimate treatment effects in the treatment group
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+        """
+        assert (learner is not None) or (
+            (control_outcome_learner is not None)
+            and (treatment_outcome_learner is not None)
+            and (treatment_effect_learner is not None)
+        )
+
+        if control_outcome_learner is None:
+            self.model_mu_c = deepcopy(learner)
+        else:
+            self.model_mu_c = control_outcome_learner
+
+        if treatment_outcome_learner is None:
+            self.model_mu_t = deepcopy(learner)
+        else:
+            self.model_mu_t = treatment_outcome_learner
+
+        if treatment_effect_learner is None:
+            self.model_tau = deepcopy(learner)
+        else:
+            self.model_tau = treatment_effect_learner
+
+        self.ate_alpha = ate_alpha
+        self.control_name = control_name
+
+        self.propensity = None
+
+    def __repr__(self):
+        return (
+            "{}(control_outcome_learner={},\n"
+            "\ttreatment_outcome_learner={},\n"
+            "\ttreatment_effect_learner={})".format(
+                self.__class__.__name__,
+                self.model_mu_c.__repr__(),
+                self.model_mu_t.__repr__(),
+                self.model_tau.__repr__(),
+            )
+        )
+
+    def fit(self, X, treatment, y, p=None, seed=None):
+        """Fit the inference model.
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series): a treatment vector
+            y (np.array or pd.Series): an outcome vector
+            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
+            seed (int): random seed for cross-fitting
+        """
+        X, treatment, y = convert_pd_to_np(X, treatment, y)
+        check_treatment_vector(treatment, self.control_name)
+        self.t_groups = np.unique(treatment[treatment != self.control_name])
+        self.t_groups.sort()
+        self._classes = {group: i for i, group in enumerate(self.t_groups)}
+
+        # The estimator splits the data into 3 partitions for cross-fit on the propensity score estimation,
+        # the outcome regression, and the treatment regression on the doubly robust estimates. The use of
+        # the partitions is rotated so we do not lose on the sample size.
+        cv = KFold(n_splits=3, shuffle=True, random_state=seed)
+        split_indices = [index for _, index in cv.split(y)]
+
+        self.models_mu_c = [
+            deepcopy(self.model_mu_c),
+            deepcopy(self.model_mu_c),
+            deepcopy(self.model_mu_c),
+        ]
+        self.models_mu_t = {
+            group: [
+                deepcopy(self.model_mu_t),
+                deepcopy(self.model_mu_t),
+                deepcopy(self.model_mu_t),
+            ]
+            for group in self.t_groups
+        }
+        self.models_tau = {
+            group: [
+                deepcopy(self.model_tau),
+                deepcopy(self.model_tau),
+                deepcopy(self.model_tau),
+            ]
+            for group in self.t_groups
+        }
+        if p is None:
+            self.propensity = {group: np.zeros(y.shape[0]) for group in self.t_groups}
+
+        for ifold in range(3):
+            treatment_idx = split_indices[ifold]
+            outcome_idx = split_indices[(ifold + 1) % 3]
+            tau_idx = split_indices[(ifold + 2) % 3]
+
+            treatment_treat, treatment_out, treatment_tau = (
+                treatment[treatment_idx],
+                treatment[outcome_idx],
+                treatment[tau_idx],
+            )
+            y_out, y_tau = y[outcome_idx], y[tau_idx]
+            X_treat, X_out, X_tau = X[treatment_idx], X[outcome_idx], X[tau_idx]
+
+            if p is None:
+                logger.info("Generating propensity score")
+                cur_p = dict()
+
+                for group in self.t_groups:
+                    mask = (treatment_treat == group) | (
+                        treatment_treat == self.control_name
+                    )
+                    treatment_filt = treatment_treat[mask]
+                    X_filt = X_treat[mask]
+                    w_filt = (treatment_filt == group).astype(int)
+                    w = (treatment_tau == group).astype(int)
+                    cur_p[group], _ = compute_propensity_score(
+                        X=X_filt, treatment=w_filt, X_pred=X_tau, treatment_pred=w
+                    )
+                    self.propensity[group][tau_idx] = cur_p[group]
+            else:
+                cur_p = dict()
+                if isinstance(p, (np.ndarray, pd.Series)):
+                    cur_p = {self.t_groups[0]: convert_pd_to_np(p[tau_idx])}
+                else:
+                    cur_p = {g: prop[tau_idx] for g, prop in p.items()}
+                check_p_conditions(cur_p, self.t_groups)
+
+            logger.info("Generate outcome regressions")
+            self.models_mu_c[ifold].fit(
+                X_out[treatment_out == self.control_name],
+                y_out[treatment_out == self.control_name],
+            )
+            for group in self.t_groups:
+                self.models_mu_t[group][ifold].fit(
+                    X_out[treatment_out == group], y_out[treatment_out == group]
+                )
+
+            logger.info("Fit pseudo outcomes from the DR formula")
+
+            for group in self.t_groups:
+                mask = (treatment_tau == group) | (treatment_tau == self.control_name)
+                treatment_filt = treatment_tau[mask]
+                X_filt = X_tau[mask]
+                y_filt = y_tau[mask]
+                w_filt = (treatment_filt == group).astype(int)
+                p_filt = cur_p[group][mask]
+                mu_t = self.models_mu_t[group][ifold].predict(X_filt)
+                mu_c = self.models_mu_c[ifold].predict(X_filt)
+                dr = (
+                    (w_filt - p_filt)
+                    / p_filt
+                    / (1 - p_filt)
+                    * (y_filt - mu_t * w_filt - mu_c * (1 - w_filt))
+                    + mu_t
+                    - mu_c
+                )
+                self.models_tau[group][ifold].fit(X_filt, dr)
+
+    def predict(
+        self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
+    ):
+        """Predict treatment effects.
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series, optional): a treatment vector
+            y (np.array or pd.Series, optional): an outcome vector
+            verbose (bool, optional): whether to output progress logs
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects.
+        """
+        X, treatment, y = convert_pd_to_np(X, treatment, y)
+
+        te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+        yhat_cs = {}
+        yhat_ts = {}
+
+        for i, group in enumerate(self.t_groups):
+            models_tau = self.models_tau[group]
+            _te = np.r_[[model.predict(X) for model in models_tau]].mean(axis=0)
+            te[:, i] = np.ravel(_te)
+            yhat_cs[group] = np.r_[
+                [model.predict(X) for model in self.models_mu_c]
+            ].mean(axis=0)
+            yhat_ts[group] = np.r_[
+                [model.predict(X) for model in self.models_mu_t[group]]
+            ].mean(axis=0)
+
+            if (y is not None) and (treatment is not None) and verbose:
+                mask = (treatment == group) | (treatment == self.control_name)
+                treatment_filt = treatment[mask]
+                y_filt = y[mask]
+                w = (treatment_filt == group).astype(int)
+
+                yhat = np.zeros_like(y_filt, dtype=float)
+                yhat[w == 0] = yhat_cs[group][mask][w == 0]
+                yhat[w == 1] = yhat_ts[group][mask][w == 1]
+
+                logger.info("Error metrics for group {}".format(group))
+                regression_metrics(y_filt, yhat, w)
+
+        if not return_components:
+            return te
+        else:
+            return te, yhat_cs, yhat_ts
+
+    def fit_predict(
+        self,
+        X,
+        treatment,
+        y,
+        p=None,
+        return_ci=False,
+        n_bootstraps=1000,
+        bootstrap_size=10000,
+        return_components=False,
+        verbose=True,
+        seed=None,
+    ):
+        """Fit the treatment effect and outcome models of the R learner and predict treatment effects.
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series): a treatment vector
+            y (np.array or pd.Series): an outcome vector
+            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
+            return_ci (bool): whether to return confidence intervals
+            n_bootstraps (int): number of bootstrap iterations
+            bootstrap_size (int): number of samples per bootstrap
+            return_components (bool, optional): whether to return outcome for treatment and control seperately
+            verbose (str): whether to output progress logs
+            seed (int): random seed for cross-fitting
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects. Output dim: [n_samples, n_treatment]
+                If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment],
+                UB [n_samples, n_treatment]
+        """
+        X, treatment, y = convert_pd_to_np(X, treatment, y)
+        self.fit(X, treatment, y, p, seed)
+
+        if p is None:
+            p = self.propensity
+
+        check_p_conditions(p, self.t_groups)
+        if isinstance(p, (np.ndarray, pd.Series)):
+            treatment_name = self.t_groups[0]
+            p = {treatment_name: convert_pd_to_np(p)}
+        elif isinstance(p, dict):
+            p = {
+                treatment_name: convert_pd_to_np(_p) for treatment_name, _p in p.items()
+            }
+
+        te = self.predict(
+            X, treatment=treatment, y=y, return_components=return_components
+        )
+
+        if not return_ci:
+            return te
+        else:
+            t_groups_global = self.t_groups
+            _classes_global = self._classes
+            models_mu_c_global = deepcopy(self.models_mu_c)
+            models_mu_t_global = deepcopy(self.models_mu_t)
+            models_tau_global = deepcopy(self.models_tau)
+            te_bootstraps = np.zeros(
+                shape=(X.shape[0], self.t_groups.shape[0], n_bootstraps)
+            )
+
+            logger.info("Bootstrap Confidence Intervals")
+            for i in tqdm(range(n_bootstraps)):
+                te_b = self.bootstrap(X, treatment, y, p, size=bootstrap_size)
+                te_bootstraps[:, :, i] = te_b
+
+            te_lower = np.percentile(te_bootstraps, (self.ate_alpha / 2) * 100, axis=2)
+            te_upper = np.percentile(
+                te_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=2
+            )
+
+            # set member variables back to global (currently last bootstrapped outcome)
+            self.t_groups = t_groups_global
+            self._classes = _classes_global
+            self.models_mu_c = deepcopy(models_mu_c_global)
+            self.models_mu_t = deepcopy(models_mu_t_global)
+            self.models_tau = deepcopy(models_tau_global)
+
+            return (te, te_lower, te_upper)
+
+    def estimate_ate(
+        self,
+        X,
+        treatment,
+        y,
+        p=None,
+        bootstrap_ci=False,
+        n_bootstraps=1000,
+        bootstrap_size=10000,
+        seed=None,
+        pretrain=False,
+    ):
+        """Estimate the Average Treatment Effect (ATE).
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series): a treatment vector
+            y (np.array or pd.Series): an outcome vector
+            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
+            bootstrap_ci (bool): whether run bootstrap for confidence intervals
+            n_bootstraps (int): number of bootstrap iterations
+            bootstrap_size (int): number of samples per bootstrap
+            seed (int): random seed for cross-fitting
+            pretrain (bool): whether a model has been fit, default False.
+        Returns:
+            The mean and confidence interval (LB, UB) of the ATE estimate.
+        """
+        if pretrain:
+            te, yhat_cs, yhat_ts = self.predict(
+                X, treatment, y, p, return_components=True
+            )
+        else:
+            te, yhat_cs, yhat_ts = self.fit_predict(
+                X, treatment, y, p, return_components=True, seed=seed
+            )
+        X, treatment, y = convert_pd_to_np(X, treatment, y)
+
+        if p is None:
+            p = self.propensity
+        else:
+            check_p_conditions(p, self.t_groups)
+        if isinstance(p, (np.ndarray, pd.Series)):
+            treatment_name = self.t_groups[0]
+            p = {treatment_name: convert_pd_to_np(p)}
+        elif isinstance(p, dict):
+            p = {
+                treatment_name: convert_pd_to_np(_p) for treatment_name, _p in p.items()
+            }
+
+        ate = np.zeros(self.t_groups.shape[0])
+        ate_lb = np.zeros(self.t_groups.shape[0])
+        ate_ub = np.zeros(self.t_groups.shape[0])
+
+        for i, group in enumerate(self.t_groups):
+            _ate = te[:, i].mean()
+
+            mask = (treatment == group) | (treatment == self.control_name)
+            treatment_filt = treatment[mask]
+            w = (treatment_filt == group).astype(int)
+            prob_treatment = float(sum(w)) / w.shape[0]
+
+            yhat_c = yhat_cs[group][mask]
+            yhat_t = yhat_ts[group][mask]
+            y_filt = y[mask]
+
+            # SE formula is based on the lower bound formula (7) from Imbens, Guido W., and Jeffrey M. Wooldridge. 2009.
+            # "Recent Developments in the Econometrics of Program Evaluation." Journal of Economic Literature
+            se = np.sqrt(
+                (
+                    (y_filt[w == 0] - yhat_c[w == 0]).var() / (1 - prob_treatment)
+                    + (y_filt[w == 1] - yhat_t[w == 1]).var() / prob_treatment
+                    + (yhat_t - yhat_c).var()
+                )
+                / y_filt.shape[0]
+            )
+
+            _ate_lb = _ate - se * norm.ppf(1 - self.ate_alpha / 2)
+            _ate_ub = _ate + se * norm.ppf(1 - self.ate_alpha / 2)
+
+            ate[i] = _ate
+            ate_lb[i] = _ate_lb
+            ate_ub[i] = _ate_ub
+
+        if not bootstrap_ci:
+            return ate, ate_lb, ate_ub
+        else:
+            t_groups_global = self.t_groups
+            _classes_global = self._classes
+            models_mu_c_global = deepcopy(self.models_mu_c)
+            models_mu_t_global = deepcopy(self.models_mu_t)
+            models_tau_global = deepcopy(self.models_tau)
+
+            logger.info("Bootstrap Confidence Intervals for ATE")
+            ate_bootstraps = np.zeros(shape=(self.t_groups.shape[0], n_bootstraps))
+
+            for n in tqdm(range(n_bootstraps)):
+                cate_b = self.bootstrap(
+                    X, treatment, y, p, size=bootstrap_size, seed=seed
+                )
+                ate_bootstraps[:, n] = cate_b.mean(axis=0)
+
+            ate_lower = np.percentile(
+                ate_bootstraps, (self.ate_alpha / 2) * 100, axis=1
+            )
+            ate_upper = np.percentile(
+                ate_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=1
+            )
+
+            # set member variables back to global (currently last bootstrapped outcome)
+            self.t_groups = t_groups_global
+            self._classes = _classes_global
+            self.models_mu_c = deepcopy(models_mu_c_global)
+            self.models_mu_t = deepcopy(models_mu_t_global)
+            self.models_tau = deepcopy(models_tau_global)
+            return ate, ate_lower, ate_upper
+
+
+class BaseDRRegressor(BaseDRLearner):
+    """
+    A parent class for DR-learner regressor classes.
+    """
+
+    def __init__(
+        self,
+        learner=None,
+        control_outcome_learner=None,
+        treatment_outcome_learner=None,
+        treatment_effect_learner=None,
+        ate_alpha=0.05,
+        control_name=0,
+    ):
+        """Initialize an DR-learner regressor.
+
+        Args:
+            learner (optional): a model to estimate outcomes and treatment effects in both the control and treatment
+                groups
+            control_outcome_learner (optional): a model to estimate outcomes in the control group
+            treatment_outcome_learner (optional): a model to estimate outcomes in the treatment group
+            treatment_effect_learner (optional): a model to estimate treatment effects in the treatment group
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+        """
+        super().__init__(
+            learner=learner,
+            control_outcome_learner=control_outcome_learner,
+            treatment_outcome_learner=treatment_outcome_learner,
+            treatment_effect_learner=treatment_effect_learner,
+            ate_alpha=ate_alpha,
+            control_name=control_name,
+        )
+
+
+class BaseDRClassifier(BaseDRLearner):
+    """
+    A parent class for DR-learner classifier classes.
+    """
+
+    def __init__(
+        self,
+        learner=None,
+        control_outcome_learner=None,
+        treatment_outcome_learner=None,
+        treatment_effect_learner=None,
+        ate_alpha=0.05,
+        control_name=0,
+    ):
+        """Initialize a DR-learner classifier.
+
+        Args:
+            learner (optional): a model to estimate outcomes and treatment effects in both the control and treatment
+                groups. Should have a predict_proba() method for outcome models.
+            control_outcome_learner (optional): a model to estimate outcomes in the control group.
+                Should have a predict_proba() method.
+            treatment_outcome_learner (optional): a model to estimate outcomes in the treatment group.
+                Should have a predict_proba() method.
+            treatment_effect_learner (optional): a model to estimate treatment effects in the treatment group.
+                Should be a regressor.
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+        """
+        super().__init__(
+            learner=learner,
+            control_outcome_learner=control_outcome_learner,
+            treatment_outcome_learner=treatment_outcome_learner,
+            treatment_effect_learner=treatment_effect_learner,
+            ate_alpha=ate_alpha,
+            control_name=control_name,
+        )
+
+    def predict(
+        self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
+    ):
+        """Predict treatment effects.
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series, optional): a treatment vector. Used for computing
+                classification metrics when y is also provided.
+            y (np.array or pd.Series, optional): an outcome vector. Used for computing
+                classification metrics when treatment is also provided.
+            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1). Currently not used in prediction but kept for API consistency.
+            return_components (bool, optional): whether to return outcome probabilities for treatment and control
+                groups separately. Defaults to False.
+            verbose (bool, optional): whether to output progress logs. Defaults to True.
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects.
+            If return_components is True, also returns:
+                - dict: Predicted probabilities for the control group (yhat_cs).
+                - dict: Predicted probabilities for the treatment group (yhat_ts).
+        """
+        X, treatment, y = convert_pd_to_np(X, treatment, y)
+
+        te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+        yhat_cs = {}
+        yhat_ts = {}
+
+        for i, group in enumerate(self.t_groups):
+            models_tau = self.models_tau[group]
+            _te = np.r_[[model.predict(X) for model in models_tau]].mean(axis=0)
+            te[:, i] = np.ravel(_te)
+            yhat_cs[group] = np.r_[
+                [model.predict_proba(X)[:, 1] for model in self.models_mu_c]
+            ].mean(axis=0)
+            yhat_ts[group] = np.r_[
+                [model.predict_proba(X)[:, 1] for model in self.models_mu_t[group]]
+            ].mean(axis=0)
+
+            if (y is not None) and (treatment is not None) and verbose:
+                mask = (treatment == group) | (treatment == self.control_name)
+                treatment_filt = treatment[mask]
+                y_filt = y[mask]
+                w = (treatment_filt == group).astype(int)
+
+                yhat = np.zeros_like(y_filt, dtype=float)
+                yhat[w == 0] = yhat_cs[group][mask][w == 0]
+                yhat[w == 1] = yhat_ts[group][mask][w == 1]
+
+                logger.info("Error metrics for group {}".format(group))
+                classification_metrics(y_filt, yhat, w)
+
+        if not return_components:
+            return te
+        else:
+            return te, yhat_cs, yhat_ts
+
+
+class XGBDRRegressor(BaseDRRegressor):
+    def __init__(self, ate_alpha=0.05, control_name=0, *args, **kwargs):
+        """Initialize a DR-learner with two XGBoost models."""
+        super().__init__(
+            learner=XGBRegressor(*args, **kwargs),
+            ate_alpha=ate_alpha,
+            control_name=control_name,
+        )

causalml/source/causalml/inference/meta/explainer.py

@@ -0,0 +1,278 @@
+import pandas as pd
+import shap
+import matplotlib.pyplot as plt
+from lightgbm import LGBMRegressor
+from sklearn.inspection import permutation_importance
+from sklearn.model_selection import train_test_split
+from copy import deepcopy
+
+from causalml.inference.meta.utils import convert_pd_to_np
+
+VALID_METHODS = ("auto", "permutation", "shapley")
+
+
+class Explainer:
+    def __init__(
+        self,
+        method,
+        control_name,
+        X,
+        tau,
+        classes,
+        model_tau=None,
+        features=None,
+        normalize=True,
+        test_size=0.3,
+        random_state=None,
+        override_checks=False,
+        r_learners=None,
+    ):
+        """
+        The Explainer class handles all feature explanation/interpretation functions, including plotting
+        feature importances, shapley value distributions, and shapley value dependency plots.
+
+        Currently supported methods are:
+            - auto (calculates importance based on estimator's default implementation of feature importance;
+                    estimator must be tree-based)
+                    Note: if none provided, it uses lightgbm's LGBMRegressor as estimator, and "gain" as
+                    importance type
+            - permutation (calculates importance based on mean decrease in accuracy when a feature column is permuted;
+                           estimator can be any form)
+            - shapley (calculates shapley values; estimator must be tree-based)
+        Hint: for permutation, downsample data for better performance especially if X.shape[1] is large
+
+        Args:
+            method (str): auto, permutation, shapley
+            control_name (str/int/float): name of control group
+            X (np.matrix): a feature matrix
+            tau (np.array): a treatment effect vector (estimated/actual)
+            classes (dict): a mapping of treatment names to indices (used for indexing tau array)
+            model_tau (sklearn/lightgbm/xgboost model object): a model object
+            features (np.array): list/array of feature names. If None, an enumerated list will be used.
+            normalize (bool): normalize by sum of importances if method=auto (defaults to True)
+            test_size (float/int): if float, represents the proportion of the dataset to include in the test split.
+                                   If int, represents the absolute number of test samples (used for estimating
+                                   permutation importance)
+            random_state (int/RandomState instance/None): random state used in permutation importance estimation
+            override_checks (bool): overrides self.check_conditions (e.g. if importance/shapley values are pre-computed)
+            r_learners (dict): a mapping of treatment group to fitted R Learners
+        """
+        self.method = method
+        self.control_name = control_name
+        self.X = convert_pd_to_np(X)
+        self.tau = convert_pd_to_np(tau)
+        if self.tau is not None and self.tau.ndim == 1:
+            self.tau = self.tau.reshape(-1, 1)
+        self.classes = classes
+        self.model_tau = (
+            LGBMRegressor(importance_type="gain") if model_tau is None else model_tau
+        )
+        self.features = features
+        self.normalize = normalize
+        self.test_size = test_size
+        self.random_state = random_state
+        self.override_checks = override_checks
+        self.r_learners = r_learners
+
+        if not self.override_checks:
+            self.check_conditions()
+            self.create_feature_names()
+            self.build_new_tau_models()
+
+    def check_conditions(self):
+        """
+        Checks for multiple conditions:
+            - method is valid
+            - X, tau, and classes are specified
+            - model_tau has feature_importances_ attribute after fitting
+        """
+        assert self.method in VALID_METHODS, "Current supported methods: {}".format(
+            ", ".join(VALID_METHODS)
+        )
+
+        assert all(
+            obj is not None for obj in (self.X, self.tau, self.classes)
+        ), "X, tau, and classes must be provided."
+
+        model_test = deepcopy(self.model_tau)
+        model_test.fit(
+            [[0], [1]], [0, 1]
+        )  # Fit w/ dummy data to check for feature_importances_ below
+        assert hasattr(
+            model_test, "feature_importances_"
+        ), "model_tau must have the feature_importances_ method (after fitting)"
+
+    def create_feature_names(self):
+        """
+        Creates feature names (simple enumerated list) if not provided in __init__.
+        """
+        if self.features is None:
+            num_features = self.X.shape[1]
+            self.features = ["Feature_{:03d}".format(i) for i in range(num_features)]
+
+    def build_new_tau_models(self):
+        """
+        Builds tau models (using X to predict estimated/actual tau) for each treatment group.
+        """
+        if self.method in ("permutation"):
+            self.X_train, self.X_test, self.tau_train, self.tau_test = train_test_split(
+                self.X,
+                self.tau,
+                test_size=self.test_size,
+                random_state=self.random_state,
+            )
+        else:
+            self.X_train, self.tau_train = self.X, self.tau
+
+        if self.r_learners is not None:
+            self.models_tau = deepcopy(self.r_learners)
+        else:
+            self.models_tau = {
+                group: deepcopy(self.model_tau) for group in self.classes
+            }
+            for group, idx in self.classes.items():
+                self.models_tau[group].fit(self.X_train, self.tau_train[:, idx])
+
+    def get_importance(self):
+        """
+        Calculates feature importances for each treatment group, based on specified method in __init__.
+        """
+        importance_catalog = {
+            "auto": self.default_importance,
+            "permutation": self.perm_importance,
+        }
+        importance_dict = importance_catalog[self.method]()
+
+        importance_dict = {
+            group: pd.Series(array, index=self.features).sort_values(ascending=False)
+            for group, array in importance_dict.items()
+        }
+        return importance_dict
+
+    def default_importance(self):
+        """
+        Calculates feature importances for each treatment group, based on the model_tau's default implementation.
+        """
+        importance_dict = {}
+        if self.r_learners is not None:
+            self.models_tau = deepcopy(self.r_learners)
+        for group, idx in self.classes.items():
+            importance_dict[group] = self.models_tau[group].feature_importances_
+            if self.normalize:
+                importance_dict[group] = (
+                    importance_dict[group] / importance_dict[group].sum()
+                )
+
+        return importance_dict
+
+    def perm_importance(self):
+        """
+        Calculates feature importances for each treatment group, based on the permutation method.
+        """
+        importance_dict = {}
+        if self.r_learners is not None:
+            self.models_tau = deepcopy(self.r_learners)
+            self.X_test, self.tau_test = self.X, self.tau
+        for group, idx in self.classes.items():
+            perm_estimator = self.models_tau[group]
+            importance_dict[group] = permutation_importance(
+                estimator=perm_estimator,
+                X=self.X_test,
+                y=self.tau_test[:, idx],
+                random_state=self.random_state,
+            ).importances_mean
+
+        return importance_dict
+
+    def get_shap_values(self):
+        """
+        Calculates shapley values for each treatment group.
+        """
+        shap_dict = {}
+        for group, mod in self.models_tau.items():
+            explainer = shap.TreeExplainer(mod)
+            if self.r_learners is not None:
+                explainer.model.original_model.params["objective"] = (
+                    None  # hacky way of running shap without error
+                )
+            shap_values = explainer.shap_values(self.X)
+            shap_dict[group] = shap_values
+
+        return shap_dict
+
+    def plot_importance(self, importance_dict=None, title_prefix="", figsize=(12, 8)):
+        """
+        Calculates and plots feature importances for each treatment group, based on specified method in __init__.
+        Skips the calculation part if importance_dict is given.
+        Args:
+            importance_dict (optional, dict): a dict of feature importance matrics. If None, importance_dict will be
+                computed.
+            title_prefix (optional, str): a prefix to the title of the plot.
+            figsize (optional, tuple): the size of the figure.
+        """
+        if importance_dict is None:
+            importance_dict = self.get_importance()
+        for group, series in importance_dict.items():
+            plt.figure()
+            series.sort_values().plot(kind="barh", figsize=figsize)
+            title = group
+            if title_prefix != "":
+                title = "{} - {}".format(title_prefix, title)
+            plt.title(title)
+
+    def plot_shap_values(self, shap_dict=None, **kwargs):
+        """
+        Calculates and plots the distribution of shapley values of each feature, for each treatment group.
+        Skips the calculation part if shap_dict is given.
+
+        Args:
+            shap_dict (optional, dict): a dict of shapley value matrics. If None, shap_dict will be computed.
+        """
+        if shap_dict is None:
+            shap_dict = self.get_shap_values()
+
+        for group, values in shap_dict.items():
+            plt.title(group)
+            shap.summary_plot(
+                values, features=self.X, feature_names=self.features, **kwargs
+            )
+
+    def plot_shap_dependence(
+        self,
+        treatment_group,
+        feature_idx,
+        shap_dict=None,
+        interaction_idx="auto",
+        **kwargs,
+    ):
+        """
+         Plots dependency of shapley values for a specified feature, colored by an interaction feature.
+         Skips the calculation part if shap_dict is given.
+
+         This plots the value of the feature on the x-axis and the SHAP value of the same feature
+         on the y-axis. This shows how the model depends on the given feature, and is like a
+         richer extension of the classical partial dependence plots. Vertical dispersion of the
+         data points represents interaction effects.
+
+        Args:
+             treatment_group (str or int): name of treatment group to create dependency plot on
+             feature_idx (str or int): feature index/name to create dependency plot on
+             shap_dict (optional, dict): a dict of shapley value matrices. If None, shap_dict will be computed.
+             interaction_idx (optional, str or int): feature index/name used in coloring scheme as interaction feature.
+                 If "auto" then shap.common.approximate_interactions is used to pick what seems to be the
+                 strongest interaction (note that to find to true strongest interaction you need to compute
+                 the SHAP interaction values).
+        """
+        if shap_dict is None:
+            shap_dict = self.get_shap_values()
+
+        shap_values = shap_dict[treatment_group]
+
+        shap.dependence_plot(
+            feature_idx,
+            shap_values,
+            self.X,
+            interaction_index=interaction_idx,
+            feature_names=self.features,
+            **kwargs,
+        )

causalml/source/causalml/inference/meta/rlearner.py

@@ -0,0 +1,695 @@
+from copy import deepcopy
+import logging
+import numpy as np
+from tqdm import tqdm
+from scipy.stats import norm
+from sklearn.model_selection import cross_val_predict, KFold, train_test_split
+from xgboost import XGBRegressor
+
+from causalml.inference.meta.base import BaseLearner
+from causalml.inference.meta.utils import (
+    check_treatment_vector,
+    get_xgboost_objective_metric,
+    convert_pd_to_np,
+    get_weighted_variance,
+)
+from causalml.propensity import ElasticNetPropensityModel
+
+logger = logging.getLogger("causalml")
+
+
+class BaseRLearner(BaseLearner):
+    """A parent class for R-learner classes.
+
+    An R-learner estimates treatment effects with two machine learning models and the propensity score.
+
+    Details of R-learner are available at `Nie and Wager (2019) <https://arxiv.org/abs/1712.04912>`_.
+    """
+
+    def __init__(
+        self,
+        learner=None,
+        outcome_learner=None,
+        effect_learner=None,
+        propensity_learner=ElasticNetPropensityModel(),
+        ate_alpha=0.05,
+        control_name=0,
+        n_fold=5,
+        random_state=None,
+        cv_n_jobs=-1,
+    ):
+        """Initialize an R-learner.
+
+        Args:
+            learner (optional): a model to estimate outcomes and treatment effects
+            outcome_learner (optional): a model to estimate outcomes
+            effect_learner (optional): a model to estimate treatment effects. It needs to take `sample_weight` as an
+                input argument for `fit()`
+            propensity_learner (optional): a model to estimate propensity scores. `ElasticNetPropensityModel()` will
+                be used by default.
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+            n_fold (int, optional): the number of cross validation folds for outcome_learner
+            random_state (int or RandomState, optional): a seed (int) or random number generator (RandomState)
+            cv_n_jobs (int, optional): number of parallel jobs to run for cross_val_predict. -1 means using all
+                processors
+        """
+        assert (learner is not None) or (
+            (outcome_learner is not None) and (effect_learner is not None)
+        )
+        assert propensity_learner is not None
+
+        self.model_mu = (
+            outcome_learner if outcome_learner is not None else deepcopy(learner)
+        )
+        self.model_tau = (
+            effect_learner if effect_learner is not None else deepcopy(learner)
+        )
+        self.model_p = propensity_learner
+
+        self.ate_alpha = ate_alpha
+        self.control_name = control_name
+
+        self.random_state = random_state
+        self.cv = KFold(n_splits=n_fold, shuffle=True, random_state=random_state)
+        self.cv_n_jobs = cv_n_jobs
+
+        self.propensity = None
+        self.propensity_model = None
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}\n"
+            f"\toutcome_learner={self.model_mu.__repr__()}\n"
+            f"\teffect_learner={self.model_tau.__repr__()}\n"
+            f"\tpropensity_learner={self.model_p.__repr__()}"
+        )
+
+    def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
+        """Fit the treatment effect and outcome models of the R learner.
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series): a treatment vector
+            y (np.array or pd.Series): an outcome vector
+            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
+            sample_weight (np.array or pd.Series, optional): an array of sample weights indicating the
+                weight of each observation for `effect_learner`. If None, it assumes equal weight.
+            verbose (bool, optional): whether to output progress logs
+        """
+        X, treatment, y = convert_pd_to_np(X, treatment, y)
+        check_treatment_vector(treatment, self.control_name)
+        if sample_weight is not None:
+            assert len(sample_weight) == len(
+                y
+            ), "Data length must be equal for sample_weight and the input data"
+            sample_weight = convert_pd_to_np(sample_weight)
+        self.t_groups = np.unique(treatment[treatment != self.control_name])
+        self.t_groups.sort()
+
+        if p is None:
+            self._set_propensity_models(X=X, treatment=treatment, y=y)
+            p = self.propensity
+        else:
+            p = self._format_p(p, self.t_groups)
+
+        self._classes = {group: i for i, group in enumerate(self.t_groups)}
+        self.models_tau = {group: deepcopy(self.model_tau) for group in self.t_groups}
+        self.vars_c = {}
+        self.vars_t = {}
+
+        if verbose:
+            logger.info("generating out-of-fold CV outcome estimates")
+        yhat = cross_val_predict(self.model_mu, X, y, cv=self.cv, n_jobs=self.cv_n_jobs)
+
+        for group in self.t_groups:
+            mask = (treatment == group) | (treatment == self.control_name)
+            treatment_filt = treatment[mask]
+            X_filt = X[mask]
+            y_filt = y[mask]
+            yhat_filt = yhat[mask]
+            p_filt = p[group][mask]
+            w = (treatment_filt == group).astype(int)
+
+            weight = (w - p_filt) ** 2
+            diff_c = y_filt[w == 0] - yhat_filt[w == 0]
+            diff_t = y_filt[w == 1] - yhat_filt[w == 1]
+            if sample_weight is not None:
+                sample_weight_filt = sample_weight[mask]
+                sample_weight_filt_c = sample_weight_filt[w == 0]
+                sample_weight_filt_t = sample_weight_filt[w == 1]
+                self.vars_c[group] = get_weighted_variance(diff_c, sample_weight_filt_c)
+                self.vars_t[group] = get_weighted_variance(diff_t, sample_weight_filt_t)
+                weight *= sample_weight_filt  # update weight
+            else:
+                self.vars_c[group] = diff_c.var()
+                self.vars_t[group] = diff_t.var()
+
+            if verbose:
+                logger.info(
+                    "training the treatment effect model for {} with R-loss".format(
+                        group
+                    )
+                )
+            self.models_tau[group].fit(
+                X_filt, (y_filt - yhat_filt) / (w - p_filt), sample_weight=weight
+            )
+
+    def predict(self, X, p=None):
+        """Predict treatment effects.
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects.
+        """
+        X = convert_pd_to_np(X)
+        te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+        for i, group in enumerate(self.t_groups):
+            dhat = self.models_tau[group].predict(X)
+            te[:, i] = dhat
+
+        return te
+
+    def fit_predict(
+        self,
+        X,
+        treatment,
+        y,
+        p=None,
+        sample_weight=None,
+        return_ci=False,
+        n_bootstraps=1000,
+        bootstrap_size=10000,
+        verbose=True,
+    ):
+        """Fit the treatment effect and outcome models of the R learner and predict treatment effects.
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series): a treatment vector
+            y (np.array or pd.Series): an outcome vector
+            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
+            sample_weight (np.array or pd.Series, optional): an array of sample weights indicating the
+                weight of each observation for `effect_learner`. If None, it assumes equal weight.
+            return_ci (bool): whether to return confidence intervals
+            n_bootstraps (int): number of bootstrap iterations
+            bootstrap_size (int): number of samples per bootstrap
+            verbose (bool): whether to output progress logs
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects. Output dim: [n_samples, n_treatment].
+                If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment],
+                UB [n_samples, n_treatment]
+        """
+        X, treatment, y = convert_pd_to_np(X, treatment, y)
+        self.fit(X, treatment, y, p, sample_weight, verbose=verbose)
+        te = self.predict(X)
+
+        if not return_ci:
+            return te
+        else:
+            t_groups_global = self.t_groups
+            _classes_global = self._classes
+            model_mu_global = deepcopy(self.model_mu)
+            models_tau_global = deepcopy(self.models_tau)
+            te_bootstraps = np.zeros(
+                shape=(X.shape[0], self.t_groups.shape[0], n_bootstraps)
+            )
+
+            logger.info("Bootstrap Confidence Intervals")
+            for i in tqdm(range(n_bootstraps)):
+                if p is None:
+                    p = self.propensity
+                else:
+                    p = self._format_p(p, self.t_groups)
+                te_b = self.bootstrap(X, treatment, y, p, size=bootstrap_size)
+                te_bootstraps[:, :, i] = te_b
+
+            te_lower = np.percentile(te_bootstraps, (self.ate_alpha / 2) * 100, axis=2)
+            te_upper = np.percentile(
+                te_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=2
+            )
+
+            # set member variables back to global (currently last bootstrapped outcome)
+            self.t_groups = t_groups_global
+            self._classes = _classes_global
+            self.model_mu = deepcopy(model_mu_global)
+            self.models_tau = deepcopy(models_tau_global)
+
+            return (te, te_lower, te_upper)
+
+    def estimate_ate(
+        self,
+        X,
+        treatment=None,
+        y=None,
+        p=None,
+        sample_weight=None,
+        bootstrap_ci=False,
+        n_bootstraps=1000,
+        bootstrap_size=10000,
+        pretrain=False,
+    ):
+        """Estimate the Average Treatment Effect (ATE).
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series): only needed when pretrain=False, a treatment vector
+            y (np.array or pd.Series):only needed when pretrain=False, an outcome vector
+            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
+            sample_weight (np.array or pd.Series, optional): an array of sample weights indicating the
+                weight of each observation for `effect_learner`. If None, it assumes equal weight.
+            bootstrap_ci (bool): whether run bootstrap for confidence intervals
+            n_bootstraps (int): number of bootstrap iterations
+            bootstrap_size (int): number of samples per bootstrap
+            pretrain (bool): whether a model has been fit, default False.
+        Returns:
+            The mean and confidence interval (LB, UB) of the ATE estimate.
+        """
+        X, treatment, y = convert_pd_to_np(X, treatment, y)
+        if pretrain:
+            te = self.predict(X, p)
+        else:
+            if not len(treatment) or not len(y):
+                raise ValueError("treatmeng and y must be provided when pretrain=False")
+            te = self.fit_predict(X, treatment, y, p, sample_weight, return_ci=False)
+
+        ate = np.zeros(self.t_groups.shape[0])
+        ate_lb = np.zeros(self.t_groups.shape[0])
+        ate_ub = np.zeros(self.t_groups.shape[0])
+
+        for i, group in enumerate(self.t_groups):
+            w = (treatment == group).astype(int)
+            prob_treatment = float(sum(w)) / X.shape[0]
+            _ate = te[:, i].mean()
+
+            se = (
+                np.sqrt(
+                    (self.vars_t[group] / prob_treatment)
+                    + (self.vars_c[group] / (1 - prob_treatment))
+                    + te[:, i].var()
+                )
+                / X.shape[0]
+            )
+
+            _ate_lb = _ate - se * norm.ppf(1 - self.ate_alpha / 2)
+            _ate_ub = _ate + se * norm.ppf(1 - self.ate_alpha / 2)
+
+            ate[i] = _ate
+            ate_lb[i] = _ate_lb
+            ate_ub[i] = _ate_ub
+
+        if not bootstrap_ci:
+            return ate, ate_lb, ate_ub
+        else:
+            t_groups_global = self.t_groups
+            _classes_global = self._classes
+            model_mu_global = deepcopy(self.model_mu)
+            models_tau_global = deepcopy(self.models_tau)
+
+            logger.info("Bootstrap Confidence Intervals for ATE")
+            ate_bootstraps = np.zeros(shape=(self.t_groups.shape[0], n_bootstraps))
+
+            for n in tqdm(range(n_bootstraps)):
+                if p is None:
+                    p = self.propensity
+                else:
+                    p = self._format_p(p, self.t_groups)
+                cate_b = self.bootstrap(X, treatment, y, p, size=bootstrap_size)
+                ate_bootstraps[:, n] = cate_b.mean(axis=0)
+
+            ate_lower = np.percentile(
+                ate_bootstraps, (self.ate_alpha / 2) * 100, axis=1
+            )
+            ate_upper = np.percentile(
+                ate_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=1
+            )
+
+            # set member variables back to global (currently last bootstrapped outcome)
+            self.t_groups = t_groups_global
+            self._classes = _classes_global
+            self.model_mu = deepcopy(model_mu_global)
+            self.models_tau = deepcopy(models_tau_global)
+            return ate, ate_lower, ate_upper
+
+
+class BaseRRegressor(BaseRLearner):
+    """
+    A parent class for R-learner regressor classes.
+    """
+
+    def __init__(
+        self,
+        learner=None,
+        outcome_learner=None,
+        effect_learner=None,
+        propensity_learner=ElasticNetPropensityModel(),
+        ate_alpha=0.05,
+        control_name=0,
+        n_fold=5,
+        random_state=None,
+    ):
+        """Initialize an R-learner regressor.
+
+        Args:
+            learner (optional): a model to estimate outcomes and treatment effects
+            outcome_learner (optional): a model to estimate outcomes
+            effect_learner (optional): a model to estimate treatment effects. It needs to take `sample_weight` as an
+                input argument for `fit()`
+            propensity_learner (optional): a model to estimate propensity scores. `ElasticNetPropensityModel()` will
+                be used by default.
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+            n_fold (int, optional): the number of cross validation folds for outcome_learner
+            random_state (int or RandomState, optional): a seed (int) or random number generator (RandomState)
+        """
+        super().__init__(
+            learner=learner,
+            outcome_learner=outcome_learner,
+            effect_learner=effect_learner,
+            propensity_learner=propensity_learner,
+            ate_alpha=ate_alpha,
+            control_name=control_name,
+            n_fold=n_fold,
+            random_state=random_state,
+        )
+
+
+class BaseRClassifier(BaseRLearner):
+    """
+    A parent class for R-learner classifier classes.
+    """
+
+    def __init__(
+        self,
+        outcome_learner=None,
+        effect_learner=None,
+        propensity_learner=ElasticNetPropensityModel(),
+        ate_alpha=0.05,
+        control_name=0,
+        n_fold=5,
+        random_state=None,
+    ):
+        """Initialize an R-learner classifier.
+
+        Args:
+            outcome_learner: a model to estimate outcomes. Should be a classifier.
+            effect_learner: a model to estimate treatment effects. It needs to take `sample_weight` as an
+                input argument for `fit()`. Should be a regressor.
+            propensity_learner (optional): a model to estimate propensity scores. `ElasticNetPropensityModel()` will
+                be used by default.
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+            n_fold (int, optional): the number of cross validation folds for outcome_learner
+            random_state (int or RandomState, optional): a seed (int) or random number generator (RandomState)
+        """
+        super().__init__(
+            learner=None,
+            outcome_learner=outcome_learner,
+            effect_learner=effect_learner,
+            propensity_learner=propensity_learner,
+            ate_alpha=ate_alpha,
+            control_name=control_name,
+            n_fold=n_fold,
+            random_state=random_state,
+        )
+
+        if (outcome_learner is None) and (effect_learner is None):
+            raise ValueError(
+                "Either the outcome learner or the effect learner must be specified."
+            )
+
+    def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
+        """Fit the treatment effect and outcome models of the R learner.
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series): a treatment vector
+            y (np.array or pd.Series): an outcome vector
+            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
+            sample_weight (np.array or pd.Series, optional): an array of sample weights indicating the
+                weight of each observation for `effect_learner`. If None, it assumes equal weight.
+            verbose (bool, optional): whether to output progress logs
+        """
+        X, treatment, y = convert_pd_to_np(X, treatment, y)
+        check_treatment_vector(treatment, self.control_name)
+        if sample_weight is not None:
+            assert len(sample_weight) == len(
+                y
+            ), "Data length must be equal for sample_weight and the input data"
+            sample_weight = convert_pd_to_np(sample_weight)
+        self.t_groups = np.unique(treatment[treatment != self.control_name])
+        self.t_groups.sort()
+
+        if p is None:
+            self._set_propensity_models(X=X, treatment=treatment, y=y)
+            p = self.propensity
+        else:
+            p = self._format_p(p, self.t_groups)
+
+        self._classes = {group: i for i, group in enumerate(self.t_groups)}
+        self.models_tau = {group: deepcopy(self.model_tau) for group in self.t_groups}
+        self.vars_c = {}
+        self.vars_t = {}
+
+        if verbose:
+            logger.info("generating out-of-fold CV outcome estimates")
+        yhat = cross_val_predict(
+            self.model_mu, X, y, cv=self.cv, method="predict_proba", n_jobs=-1
+        )[:, 1]
+
+        for group in self.t_groups:
+            mask = (treatment == group) | (treatment == self.control_name)
+            treatment_filt = treatment[mask]
+            X_filt = X[mask]
+            y_filt = y[mask]
+            yhat_filt = yhat[mask]
+            p_filt = p[group][mask]
+            w = (treatment_filt == group).astype(int)
+
+            weight = (w - p_filt) ** 2
+            diff_c = y_filt[w == 0] - yhat_filt[w == 0]
+            diff_t = y_filt[w == 1] - yhat_filt[w == 1]
+            if sample_weight is not None:
+                sample_weight_filt = sample_weight[mask]
+                sample_weight_filt_c = sample_weight_filt[w == 0]
+                sample_weight_filt_t = sample_weight_filt[w == 1]
+                self.vars_c[group] = get_weighted_variance(diff_c, sample_weight_filt_c)
+                self.vars_t[group] = get_weighted_variance(diff_t, sample_weight_filt_t)
+                weight *= sample_weight_filt  # update weight
+            else:
+                self.vars_c[group] = diff_c.var()
+                self.vars_t[group] = diff_t.var()
+
+            if verbose:
+                logger.info(
+                    "training the treatment effect model for {} with R-loss".format(
+                        group
+                    )
+                )
+            self.models_tau[group].fit(
+                X_filt, (y_filt - yhat_filt) / (w - p_filt), sample_weight=weight
+            )
+
+    def predict(self, X, p=None):
+        """Predict treatment effects.
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects.
+        """
+        X = convert_pd_to_np(X)
+        te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+        for i, group in enumerate(self.t_groups):
+            dhat = self.models_tau[group].predict(X)
+            te[:, i] = dhat
+
+        return te
+
+
+class XGBRRegressor(BaseRRegressor):
+    def __init__(
+        self,
+        early_stopping=True,
+        test_size=0.3,
+        early_stopping_rounds=30,
+        effect_learner_objective="reg:squarederror",
+        effect_learner_n_estimators=500,
+        random_state=42,
+        *args,
+        **kwargs,
+    ):
+        """Initialize an R-learner regressor with XGBoost model using pairwise ranking objective.
+
+        Args:
+            early_stopping: whether or not to use early stopping when fitting effect learner
+            test_size (float, optional): the proportion of the dataset to use as validation set when early stopping is
+                                         enabled
+            early_stopping_rounds (int, optional): validation metric needs to improve at least once in every
+                                                   early_stopping_rounds round(s) to continue training
+            effect_learner_objective (str, optional): the learning objective for the effect learner
+                                                      (default = 'reg:squarederror')
+            effect_learner_n_estimators (int, optional): number of trees to fit for the effect learner (default = 500)
+        """
+
+        assert isinstance(random_state, int), "random_state should be int."
+
+        objective, metric = get_xgboost_objective_metric(effect_learner_objective)
+        self.effect_learner_objective = objective
+        self.effect_learner_eval_metric = metric
+        self.effect_learner_n_estimators = effect_learner_n_estimators
+        self.early_stopping = early_stopping
+        if self.early_stopping:
+            self.test_size = test_size
+            self.early_stopping_rounds = early_stopping_rounds
+
+            effect_learner = XGBRegressor(
+                objective=self.effect_learner_objective,
+                n_estimators=self.effect_learner_n_estimators,
+                eval_metric=self.effect_learner_eval_metric,
+                early_stopping_rounds=self.early_stopping_rounds,
+                random_state=random_state,
+                *args,
+                **kwargs,
+            )
+        else:
+            effect_learner = XGBRegressor(
+                objective=self.effect_learner_objective,
+                n_estimators=self.effect_learner_n_estimators,
+                eval_metric=self.effect_learner_eval_metric,
+                random_state=random_state,
+                *args,
+                **kwargs,
+            )
+
+        super().__init__(
+            outcome_learner=XGBRegressor(random_state=random_state, *args, **kwargs),
+            effect_learner=effect_learner,
+        )
+
+    def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
+        """Fit the treatment effect and outcome models of the R learner.
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            y (np.array or pd.Series): an outcome vector
+            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
+            sample_weight (np.array or pd.Series, optional): an array of sample weights indicating the
+                weight of each observation for `effect_learner`. If None, it assumes equal weight.
+            verbose (bool, optional): whether to output progress logs
+        """
+        X, treatment, y = convert_pd_to_np(X, treatment, y)
+        check_treatment_vector(treatment, self.control_name)
+        # initialize equal sample weight if it's not provided, for simplicity purpose
+        sample_weight = (
+            convert_pd_to_np(sample_weight)
+            if sample_weight is not None
+            else convert_pd_to_np(np.ones(len(y)))
+        )
+        assert len(sample_weight) == len(
+            y
+        ), "Data length must be equal for sample_weight and the input data"
+        self.t_groups = np.unique(treatment[treatment != self.control_name])
+        self.t_groups.sort()
+
+        if p is None:
+            self._set_propensity_models(X=X, treatment=treatment, y=y)
+            p = self.propensity
+        else:
+            p = self._format_p(p, self.t_groups)
+
+        self._classes = {group: i for i, group in enumerate(self.t_groups)}
+        self.models_tau = {group: deepcopy(self.model_tau) for group in self.t_groups}
+        self.vars_c = {}
+        self.vars_t = {}
+
+        if verbose:
+            logger.info("generating out-of-fold CV outcome estimates")
+        yhat = cross_val_predict(self.model_mu, X, y, cv=self.cv, n_jobs=-1)
+
+        for group in self.t_groups:
+            treatment_mask = (treatment == group) | (treatment == self.control_name)
+            treatment_filt = treatment[treatment_mask]
+            w = (treatment_filt == group).astype(int)
+
+            X_filt = X[treatment_mask]
+            y_filt = y[treatment_mask]
+            yhat_filt = yhat[treatment_mask]
+            p_filt = p[group][treatment_mask]
+            sample_weight_filt = sample_weight[treatment_mask]
+
+            if verbose:
+                logger.info(
+                    "training the treatment effect model for {} with R-loss".format(
+                        group
+                    )
+                )
+
+            if self.early_stopping:
+                (
+                    X_train_filt,
+                    X_test_filt,
+                    y_train_filt,
+                    y_test_filt,
+                    yhat_train_filt,
+                    yhat_test_filt,
+                    w_train,
+                    w_test,
+                    p_train_filt,
+                    p_test_filt,
+                    sample_weight_train_filt,
+                    sample_weight_test_filt,
+                ) = train_test_split(
+                    X_filt,
+                    y_filt,
+                    yhat_filt,
+                    w,
+                    p_filt,
+                    sample_weight_filt,
+                    test_size=self.test_size,
+                    random_state=self.random_state,
+                )
+
+                self.models_tau[group].fit(
+                    X=X_train_filt,
+                    y=(y_train_filt - yhat_train_filt) / (w_train - p_train_filt),
+                    sample_weight=sample_weight_train_filt
+                    * ((w_train - p_train_filt) ** 2),
+                    eval_set=[
+                        (
+                            X_test_filt,
+                            (y_test_filt - yhat_test_filt) / (w_test - p_test_filt),
+                        )
+                    ],
+                    sample_weight_eval_set=[
+                        sample_weight_test_filt * ((w_test - p_test_filt) ** 2)
+                    ],
+                    verbose=verbose,
+                )
+
+            else:
+                self.models_tau[group].fit(
+                    X_filt,
+                    (y_filt - yhat_filt) / (w - p_filt),
+                    sample_weight=sample_weight_filt * ((w - p_filt) ** 2),
+                )
+
+            diff_c = y_filt[w == 0] - yhat_filt[w == 0]
+            diff_t = y_filt[w == 1] - yhat_filt[w == 1]
+            sample_weight_filt_c = sample_weight_filt[w == 0]
+            sample_weight_filt_t = sample_weight_filt[w == 1]
+            self.vars_c[group] = get_weighted_variance(diff_c, sample_weight_filt_c)
+            self.vars_t[group] = get_weighted_variance(diff_t, sample_weight_filt_t)

causalml/source/causalml/inference/meta/slearner.py

@@ -0,0 +1,411 @@
+import logging
+import numpy as np
+from tqdm import tqdm
+from scipy.stats import norm
+from sklearn.dummy import DummyRegressor
+import statsmodels.api as sm
+from copy import deepcopy
+
+from causalml.inference.meta.base import BaseLearner
+from causalml.inference.meta.utils import check_treatment_vector, convert_pd_to_np
+from causalml.metrics import regression_metrics, classification_metrics
+
+logger = logging.getLogger("causalml")
+
+
+class StatsmodelsOLS:
+    """A sklearn style wrapper class for statsmodels' OLS."""
+
+    def __init__(self, cov_type="HC1", alpha=0.05):
+        """Initialize a statsmodels' OLS wrapper class object.
+        Args:
+            cov_type (str, optional): covariance estimator type.
+            alpha (float, optional): the confidence level alpha.
+        """
+        self.cov_type = cov_type
+        self.alpha = alpha
+
+    def fit(self, X, y):
+        """Fit OLS.
+        Args:
+            X (np.matrix): a feature matrix
+            y (np.array): a label vector
+        """
+        # Append ones. The first column is for the treatment indicator.
+        X = sm.add_constant(X, prepend=False, has_constant="add")
+        self.model = sm.OLS(y, X).fit(cov_type=self.cov_type)
+        self.coefficients = self.model.params
+        self.conf_ints = self.model.conf_int(alpha=self.alpha)
+
+    def predict(self, X):
+        # Append ones. The first column is for the treatment indicator.
+        X = sm.add_constant(X, prepend=False, has_constant="add")
+        return self.model.predict(X)
+
+
+class BaseSLearner(BaseLearner):
+    """A parent class for S-learner classes.
+    An S-learner estimates treatment effects with one machine learning model.
+    Details of S-learner are available at `Kunzel et al. (2018) <https://arxiv.org/abs/1706.03461>`_.
+    """
+
+    def __init__(self, learner=None, ate_alpha=0.05, control_name=0):
+        """Initialize an S-learner.
+        Args:
+            learner (optional): a model to estimate the treatment effect
+            control_name (str or int, optional): name of control group
+        """
+        if learner is not None:
+            self.model = learner
+        else:
+            self.model = DummyRegressor()
+        self.ate_alpha = ate_alpha
+        self.control_name = control_name
+
+    def __repr__(self):
+        return "{}(model={})".format(self.__class__.__name__, self.model.__repr__())
+
+    def fit(self, X, treatment, y, p=None):
+        """Fit the inference model
+        Args:
+            X (np.matrix, np.array, or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series): a treatment vector
+            y (np.array or pd.Series): an outcome vector
+        """
+        X, treatment, y = convert_pd_to_np(X, treatment, y)
+        check_treatment_vector(treatment, self.control_name)
+        self.t_groups = np.unique(treatment[treatment != self.control_name])
+        self.t_groups.sort()
+        self._classes = {group: i for i, group in enumerate(self.t_groups)}
+        self.models = {group: deepcopy(self.model) for group in self.t_groups}
+
+        for group in self.t_groups:
+            mask = (treatment == group) | (treatment == self.control_name)
+            treatment_filt = treatment[mask]
+            X_filt = X[mask]
+            y_filt = y[mask]
+
+            w = (treatment_filt == group).astype(int)
+            X_new = np.hstack((w.reshape((-1, 1)), X_filt))
+            self.models[group].fit(X_new, y_filt)
+
+    def predict(
+        self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
+    ):
+        """Predict treatment effects.
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series, optional): a treatment vector
+            y (np.array or pd.Series, optional): an outcome vector
+            return_components (bool, optional): whether to return outcome for treatment and control seperately
+            verbose (bool, optional): whether to output progress logs
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects.
+        """
+        X, treatment, y = convert_pd_to_np(X, treatment, y)
+        yhat_cs = {}
+        yhat_ts = {}
+
+        for group in self.t_groups:
+            model = self.models[group]
+
+            # set the treatment column to zero (the control group)
+            X_new = np.hstack((np.zeros((X.shape[0], 1)), X))
+            yhat_cs[group] = model.predict(X_new)
+
+            # set the treatment column to one (the treatment group)
+            X_new[:, 0] = 1
+            yhat_ts[group] = model.predict(X_new)
+
+            if (y is not None) and (treatment is not None) and verbose:
+                mask = (treatment == group) | (treatment == self.control_name)
+                treatment_filt = treatment[mask]
+                w = (treatment_filt == group).astype(int)
+                y_filt = y[mask]
+
+                yhat = np.zeros_like(y_filt, dtype=float)
+                yhat[w == 0] = yhat_cs[group][mask][w == 0]
+                yhat[w == 1] = yhat_ts[group][mask][w == 1]
+
+                logger.info("Error metrics for group {}".format(group))
+                regression_metrics(y_filt, yhat, w)
+
+        te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+        for i, group in enumerate(self.t_groups):
+            te[:, i] = yhat_ts[group] - yhat_cs[group]
+
+        if not return_components:
+            return te
+        else:
+            return te, yhat_cs, yhat_ts
+
+    def fit_predict(
+        self,
+        X,
+        treatment,
+        y,
+        p=None,
+        return_ci=False,
+        n_bootstraps=1000,
+        bootstrap_size=10000,
+        return_components=False,
+        verbose=True,
+    ):
+        """Fit the inference model of the S learner and predict treatment effects.
+        Args:
+            X (np.matrix, np.array, or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series): a treatment vector
+            y (np.array or pd.Series): an outcome vector
+            return_ci (bool, optional): whether to return confidence intervals
+            n_bootstraps (int, optional): number of bootstrap iterations
+            bootstrap_size (int, optional): number of samples per bootstrap
+            return_components (bool, optional): whether to return outcome for treatment and control seperately
+            verbose (bool, optional): whether to output progress logs
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects. Output dim: [n_samples, n_treatment].
+                If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment],
+                UB [n_samples, n_treatment]
+        """
+        self.fit(X, treatment, y)
+        te = self.predict(X, treatment, y, return_components=return_components)
+
+        if not return_ci:
+            return te
+        else:
+            t_groups_global = self.t_groups
+            _classes_global = self._classes
+            models_global = deepcopy(self.models)
+            te_bootstraps = np.zeros(
+                shape=(X.shape[0], self.t_groups.shape[0], n_bootstraps)
+            )
+
+            logger.info("Bootstrap Confidence Intervals")
+            for i in tqdm(range(n_bootstraps)):
+                te_b = self.bootstrap(X, treatment, y, size=bootstrap_size)
+                te_bootstraps[:, :, i] = te_b
+
+            te_lower = np.percentile(te_bootstraps, (self.ate_alpha / 2) * 100, axis=2)
+            te_upper = np.percentile(
+                te_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=2
+            )
+
+            # set member variables back to global (currently last bootstrapped outcome)
+            self.t_groups = t_groups_global
+            self._classes = _classes_global
+            self.models = deepcopy(models_global)
+
+            return (te, te_lower, te_upper)
+
+    def estimate_ate(
+        self,
+        X,
+        treatment,
+        y,
+        p=None,
+        return_ci=False,
+        bootstrap_ci=False,
+        n_bootstraps=1000,
+        bootstrap_size=10000,
+        pretrain=False,
+    ):
+        """Estimate the Average Treatment Effect (ATE).
+
+        Args:
+            X (np.matrix, np.array, or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series): a treatment vector
+            y (np.array or pd.Series): an outcome vector
+            return_ci (bool, optional): whether to return confidence intervals
+            bootstrap_ci (bool): whether to return confidence intervals
+            n_bootstraps (int): number of bootstrap iterations
+            bootstrap_size (int): number of samples per bootstrap
+            pretrain (bool): whether a model has been fit, default False.
+        Returns:
+            The mean and confidence interval (LB, UB) of the ATE estimate.
+        """
+
+        X, treatment, y = convert_pd_to_np(X, treatment, y)
+        if pretrain:
+            te, yhat_cs, yhat_ts = self.predict(X, treatment, y, return_components=True)
+        else:
+            te, yhat_cs, yhat_ts = self.fit_predict(
+                X, treatment, y, return_components=True
+            )
+
+        ate = np.zeros(self.t_groups.shape[0])
+        ate_lb = np.zeros(self.t_groups.shape[0])
+        ate_ub = np.zeros(self.t_groups.shape[0])
+
+        for i, group in enumerate(self.t_groups):
+            _ate = te[:, i].mean()
+
+            mask = (treatment == group) | (treatment == self.control_name)
+            treatment_filt = treatment[mask]
+            y_filt = y[mask]
+            w = (treatment_filt == group).astype(int)
+            prob_treatment = float(sum(w)) / w.shape[0]
+
+            yhat_c = yhat_cs[group][mask]
+            yhat_t = yhat_ts[group][mask]
+
+            se = np.sqrt(
+                (
+                    (y_filt[w == 0] - yhat_c[w == 0]).var() / (1 - prob_treatment)
+                    + (y_filt[w == 1] - yhat_t[w == 1]).var() / prob_treatment
+                    + (yhat_t - yhat_c).var()
+                )
+                / y_filt.shape[0]
+            )
+
+            _ate_lb = _ate - se * norm.ppf(1 - self.ate_alpha / 2)
+            _ate_ub = _ate + se * norm.ppf(1 - self.ate_alpha / 2)
+
+            ate[i] = _ate
+            ate_lb[i] = _ate_lb
+            ate_ub[i] = _ate_ub
+
+        if not return_ci:
+            return ate
+        elif return_ci and not bootstrap_ci:
+            return ate, ate_lb, ate_ub
+        else:
+            t_groups_global = self.t_groups
+            _classes_global = self._classes
+            models_global = deepcopy(self.models)
+
+            logger.info("Bootstrap Confidence Intervals for ATE")
+            ate_bootstraps = np.zeros(shape=(self.t_groups.shape[0], n_bootstraps))
+
+            for n in tqdm(range(n_bootstraps)):
+                ate_b = self.bootstrap(X, treatment, y, size=bootstrap_size)
+                ate_bootstraps[:, n] = ate_b.mean(axis=0)
+
+            ate_lower = np.percentile(
+                ate_bootstraps, (self.ate_alpha / 2) * 100, axis=1
+            )
+            ate_upper = np.percentile(
+                ate_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=1
+            )
+
+            # set member variables back to global (currently last bootstrapped outcome)
+            self.t_groups = t_groups_global
+            self._classes = _classes_global
+            self.models = deepcopy(models_global)
+
+            return ate, ate_lower, ate_upper
+
+
+class BaseSRegressor(BaseSLearner):
+    """
+    A parent class for S-learner regressor classes.
+    """
+
+    def __init__(self, learner=None, ate_alpha=0.05, control_name=0):
+        """Initialize an S-learner regressor.
+        Args:
+            learner (optional): a model to estimate the treatment effect
+            control_name (str or int, optional): name of control group
+        """
+        super().__init__(
+            learner=learner, ate_alpha=ate_alpha, control_name=control_name
+        )
+
+
+class BaseSClassifier(BaseSLearner):
+    """
+    A parent class for S-learner classifier classes.
+    """
+
+    def __init__(self, learner=None, ate_alpha=0.05, control_name=0):
+        """Initialize an S-learner classifier.
+        Args:
+            learner (optional): a model to estimate the treatment effect.
+                Should have a predict_proba() method.
+            control_name (str or int, optional): name of control group
+        """
+        super().__init__(
+            learner=learner, ate_alpha=ate_alpha, control_name=control_name
+        )
+
+    def predict(
+        self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
+    ):
+        """Predict treatment effects.
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series, optional): a treatment vector
+            y (np.array or pd.Series, optional): an outcome vector
+            return_components (bool, optional): whether to return outcome for treatment and control seperately
+            verbose (bool, optional): whether to output progress logs
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects.
+        """
+        X, treatment, y = convert_pd_to_np(X, treatment, y)
+        yhat_cs = {}
+        yhat_ts = {}
+
+        for group in self.t_groups:
+            model = self.models[group]
+
+            # set the treatment column to zero (the control group)
+            X_new = np.hstack((np.zeros((X.shape[0], 1)), X))
+            yhat_cs[group] = model.predict_proba(X_new)[:, 1]
+
+            # set the treatment column to one (the treatment group)
+            X_new[:, 0] = 1
+            yhat_ts[group] = model.predict_proba(X_new)[:, 1]
+
+            if y is not None and (treatment is not None) and verbose:
+                mask = (treatment == group) | (treatment == self.control_name)
+                treatment_filt = treatment[mask]
+                w = (treatment_filt == group).astype(int)
+                y_filt = y[mask]
+
+                yhat = np.zeros_like(y_filt, dtype=float)
+                yhat[w == 0] = yhat_cs[group][mask][w == 0]
+                yhat[w == 1] = yhat_ts[group][mask][w == 1]
+
+                logger.info("Error metrics for group {}".format(group))
+                classification_metrics(y_filt, yhat, w)
+
+        te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+        for i, group in enumerate(self.t_groups):
+            te[:, i] = yhat_ts[group] - yhat_cs[group]
+
+        if not return_components:
+            return te
+        else:
+            return te, yhat_cs, yhat_ts
+
+
+class LRSRegressor(BaseSRegressor):
+    def __init__(self, ate_alpha=0.05, control_name=0):
+        """Initialize an S-learner with a linear regression model.
+        Args:
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+        """
+        super().__init__(StatsmodelsOLS(alpha=ate_alpha), ate_alpha, control_name)
+
+    def estimate_ate(self, X, treatment, y, p=None, pretrain=False):
+        """Estimate the Average Treatment Effect (ATE).
+        Args:
+            X (np.matrix, np.array, or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series): a treatment vector
+            y (np.array or pd.Series): an outcome vector
+        Returns:
+            The mean and confidence interval (LB, UB) of the ATE estimate.
+        """
+        X, treatment, y = convert_pd_to_np(X, treatment, y)
+        if not pretrain:
+            self.fit(X, treatment, y)
+
+        ate = np.zeros(self.t_groups.shape[0])
+        ate_lb = np.zeros(self.t_groups.shape[0])
+        ate_ub = np.zeros(self.t_groups.shape[0])
+
+        for i, group in enumerate(self.t_groups):
+            ate[i] = self.models[group].coefficients[0]
+            ate_lb[i] = self.models[group].conf_ints[0, 0]
+            ate_ub[i] = self.models[group].conf_ints[0, 1]
+
+        return ate, ate_lb, ate_ub

causalml/source/causalml/inference/meta/tlearner.py

@@ -0,0 +1,423 @@
+from copy import deepcopy
+import logging
+import numpy as np
+from packaging import version
+from scipy.stats import norm
+import sklearn
+from sklearn.exceptions import ConvergenceWarning
+from sklearn.neural_network import MLPRegressor
+
+if version.parse(sklearn.__version__) >= version.parse("0.22.0"):
+    from sklearn.utils._testing import ignore_warnings
+else:
+    from sklearn.utils.testing import ignore_warnings
+from tqdm import tqdm
+from xgboost import XGBRegressor
+
+from causalml.inference.meta.base import BaseLearner
+from causalml.inference.meta.utils import check_treatment_vector, convert_pd_to_np
+from causalml.metrics import regression_metrics, classification_metrics
+
+logger = logging.getLogger("causalml")
+
+
+class BaseTLearner(BaseLearner):
+    """A parent class for T-learner regressor classes.
+
+    A T-learner estimates treatment effects with two machine learning models.
+
+    Details of T-learner are available at `Kunzel et al. (2018) <https://arxiv.org/abs/1706.03461>`_.
+    """
+
+    def __init__(
+        self,
+        learner=None,
+        control_learner=None,
+        treatment_learner=None,
+        ate_alpha=0.05,
+        control_name=0,
+    ):
+        """Initialize a T-learner.
+
+        Args:
+            learner (model): a model to estimate control and treatment outcomes.
+            control_learner (model, optional): a model to estimate control outcomes
+            treatment_learner (model, optional): a model to estimate treatment outcomes
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+        """
+        assert (learner is not None) or (
+            (control_learner is not None) and (treatment_learner is not None)
+        )
+
+        if control_learner is None:
+            self.model_c = deepcopy(learner)
+        else:
+            self.model_c = control_learner
+
+        if treatment_learner is None:
+            self.model_t = deepcopy(learner)
+        else:
+            self.model_t = treatment_learner
+
+        self.ate_alpha = ate_alpha
+        self.control_name = control_name
+
+    def __repr__(self):
+        return "{}(model_c={}, model_t={})".format(
+            self.__class__.__name__, self.model_c.__repr__(), self.model_t.__repr__()
+        )
+
+    @ignore_warnings(category=ConvergenceWarning)
+    def fit(self, X, treatment, y, p=None):
+        """Fit the inference model
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series): a treatment vector
+            y (np.array or pd.Series): an outcome vector
+        """
+        X, treatment, y = convert_pd_to_np(X, treatment, y)
+        check_treatment_vector(treatment, self.control_name)
+        self.t_groups = np.unique(treatment[treatment != self.control_name])
+        self.t_groups.sort()
+        self._classes = {group: i for i, group in enumerate(self.t_groups)}
+        self.models_c = {group: deepcopy(self.model_c) for group in self.t_groups}
+        self.models_t = {group: deepcopy(self.model_t) for group in self.t_groups}
+
+        for group in self.t_groups:
+            mask = (treatment == group) | (treatment == self.control_name)
+            treatment_filt = treatment[mask]
+            X_filt = X[mask]
+            y_filt = y[mask]
+            w = (treatment_filt == group).astype(int)
+
+            self.models_c[group].fit(X_filt[w == 0], y_filt[w == 0])
+            self.models_t[group].fit(X_filt[w == 1], y_filt[w == 1])
+
+    def predict(
+        self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
+    ):
+        """Predict treatment effects.
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series, optional): a treatment vector
+            y (np.array or pd.Series, optional): an outcome vector
+            return_components (bool, optional): whether to return outcome for treatment and control seperately
+            verbose (bool, optional): whether to output progress logs
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects.
+        """
+        X, treatment, y = convert_pd_to_np(X, treatment, y)
+        yhat_cs = {}
+        yhat_ts = {}
+
+        for group in self.t_groups:
+            model_c = self.models_c[group]
+            model_t = self.models_t[group]
+            yhat_cs[group] = model_c.predict(X)
+            yhat_ts[group] = model_t.predict(X)
+
+            if (y is not None) and (treatment is not None) and verbose:
+                mask = (treatment == group) | (treatment == self.control_name)
+                treatment_filt = treatment[mask]
+                y_filt = y[mask]
+                w = (treatment_filt == group).astype(int)
+
+                yhat = np.zeros_like(y_filt, dtype=float)
+                yhat[w == 0] = yhat_cs[group][mask][w == 0]
+                yhat[w == 1] = yhat_ts[group][mask][w == 1]
+
+                logger.info("Error metrics for group {}".format(group))
+                regression_metrics(y_filt, yhat, w)
+
+        te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+        for i, group in enumerate(self.t_groups):
+            te[:, i] = yhat_ts[group] - yhat_cs[group]
+
+        if not return_components:
+            return te
+        else:
+            return te, yhat_cs, yhat_ts
+
+    def fit_predict(
+        self,
+        X,
+        treatment,
+        y,
+        p=None,
+        return_ci=False,
+        n_bootstraps=1000,
+        bootstrap_size=10000,
+        return_components=False,
+        verbose=True,
+    ):
+        """Fit the inference model of the T learner and predict treatment effects.
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series): a treatment vector
+            y (np.array or pd.Series): an outcome vector
+            return_ci (bool): whether to return confidence intervals
+            n_bootstraps (int): number of bootstrap iterations
+            bootstrap_size (int): number of samples per bootstrap
+            return_components (bool, optional): whether to return outcome for treatment and control seperately
+            verbose (str): whether to output progress logs
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects. Output dim: [n_samples, n_treatment].
+                If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment],
+                UB [n_samples, n_treatment]
+        """
+        X, treatment, y = convert_pd_to_np(X, treatment, y)
+        self.fit(X, treatment, y)
+        te = self.predict(X, treatment, y, return_components=return_components)
+
+        if not return_ci:
+            return te
+        else:
+            t_groups_global = self.t_groups
+            _classes_global = self._classes
+            models_c_global = deepcopy(self.models_c)
+            models_t_global = deepcopy(self.models_t)
+            te_bootstraps = np.zeros(
+                shape=(X.shape[0], self.t_groups.shape[0], n_bootstraps)
+            )
+
+            logger.info("Bootstrap Confidence Intervals")
+            for i in tqdm(range(n_bootstraps)):
+                te_b = self.bootstrap(X, treatment, y, size=bootstrap_size)
+                te_bootstraps[:, :, i] = te_b
+
+            te_lower = np.percentile(te_bootstraps, (self.ate_alpha / 2) * 100, axis=2)
+            te_upper = np.percentile(
+                te_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=2
+            )
+
+            # set member variables back to global (currently last bootstrapped outcome)
+            self.t_groups = t_groups_global
+            self._classes = _classes_global
+            self.models_c = deepcopy(models_c_global)
+            self.models_t = deepcopy(models_t_global)
+
+            return (te, te_lower, te_upper)
+
+    def estimate_ate(
+        self,
+        X,
+        treatment,
+        y,
+        p=None,
+        bootstrap_ci=False,
+        n_bootstraps=1000,
+        bootstrap_size=10000,
+        pretrain=False,
+    ):
+        """Estimate the Average Treatment Effect (ATE).
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series): a treatment vector
+            y (np.array or pd.Series): an outcome vector
+            bootstrap_ci (bool): whether to return confidence intervals
+            n_bootstraps (int): number of bootstrap iterations
+            bootstrap_size (int): number of samples per bootstrap
+        Returns:
+            The mean and confidence interval (LB, UB) of the ATE estimate.
+            pretrain (bool): whether a model has been fit, default False.
+        """
+        X, treatment, y = convert_pd_to_np(X, treatment, y)
+        if pretrain:
+            te, yhat_cs, yhat_ts = self.predict(X, treatment, y, return_components=True)
+        else:
+            te, yhat_cs, yhat_ts = self.fit_predict(
+                X, treatment, y, return_components=True
+            )
+
+        ate = np.zeros(self.t_groups.shape[0])
+        ate_lb = np.zeros(self.t_groups.shape[0])
+        ate_ub = np.zeros(self.t_groups.shape[0])
+
+        for i, group in enumerate(self.t_groups):
+            _ate = te[:, i].mean()
+
+            mask = (treatment == group) | (treatment == self.control_name)
+            treatment_filt = treatment[mask]
+            y_filt = y[mask]
+            w = (treatment_filt == group).astype(int)
+            prob_treatment = float(sum(w)) / w.shape[0]
+
+            yhat_c = yhat_cs[group][mask]
+            yhat_t = yhat_ts[group][mask]
+
+            se = np.sqrt(
+                (
+                    (y_filt[w == 0] - yhat_c[w == 0]).var() / (1 - prob_treatment)
+                    + (y_filt[w == 1] - yhat_t[w == 1]).var() / prob_treatment
+                    + (yhat_t - yhat_c).var()
+                )
+                / y_filt.shape[0]
+            )
+
+            _ate_lb = _ate - se * norm.ppf(1 - self.ate_alpha / 2)
+            _ate_ub = _ate + se * norm.ppf(1 - self.ate_alpha / 2)
+
+            ate[i] = _ate
+            ate_lb[i] = _ate_lb
+            ate_ub[i] = _ate_ub
+
+        if not bootstrap_ci:
+            return ate, ate_lb, ate_ub
+        else:
+            t_groups_global = self.t_groups
+            _classes_global = self._classes
+            models_c_global = deepcopy(self.models_c)
+            models_t_global = deepcopy(self.models_t)
+
+            logger.info("Bootstrap Confidence Intervals for ATE")
+            ate_bootstraps = np.zeros(shape=(self.t_groups.shape[0], n_bootstraps))
+
+            for n in tqdm(range(n_bootstraps)):
+                ate_b = self.bootstrap(X, treatment, y, size=bootstrap_size)
+                ate_bootstraps[:, n] = ate_b.mean(axis=0)
+
+            ate_lower = np.percentile(
+                ate_bootstraps, (self.ate_alpha / 2) * 100, axis=1
+            )
+            ate_upper = np.percentile(
+                ate_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=1
+            )
+
+            # set member variables back to global (currently last bootstrapped outcome)
+            self.t_groups = t_groups_global
+            self._classes = _classes_global
+            self.models_c = deepcopy(models_c_global)
+            self.models_t = deepcopy(models_t_global)
+
+            return ate, ate_lower, ate_upper
+
+
+class BaseTRegressor(BaseTLearner):
+    """
+    A parent class for T-learner regressor classes.
+    """
+
+    def __init__(
+        self,
+        learner=None,
+        control_learner=None,
+        treatment_learner=None,
+        ate_alpha=0.05,
+        control_name=0,
+    ):
+        """Initialize a T-learner regressor.
+
+        Args:
+            learner (model): a model to estimate control and treatment outcomes.
+            control_learner (model, optional): a model to estimate control outcomes
+            treatment_learner (model, optional): a model to estimate treatment outcomes
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+        """
+        super().__init__(
+            learner=learner,
+            control_learner=control_learner,
+            treatment_learner=treatment_learner,
+            ate_alpha=ate_alpha,
+            control_name=control_name,
+        )
+
+
+class BaseTClassifier(BaseTLearner):
+    """
+    A parent class for T-learner classifier classes.
+    """
+
+    def __init__(
+        self,
+        learner=None,
+        control_learner=None,
+        treatment_learner=None,
+        ate_alpha=0.05,
+        control_name=0,
+    ):
+        """Initialize a T-learner classifier.
+
+        Args:
+            learner (model): a model to estimate control and treatment outcomes.
+            control_learner (model, optional): a model to estimate control outcomes
+            treatment_learner (model, optional): a model to estimate treatment outcomes
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+        """
+        super().__init__(
+            learner=learner,
+            control_learner=control_learner,
+            treatment_learner=treatment_learner,
+            ate_alpha=ate_alpha,
+            control_name=control_name,
+        )
+
+    def predict(
+        self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
+    ):
+        """Predict treatment effects.
+
+        Args:
+            X (np.matrix or np.array or pd.Dataframe): a feature matrix
+            treatment (np.array or pd.Series, optional): a treatment vector
+            y (np.array or pd.Series, optional): an outcome vector
+            verbose (bool, optional): whether to output progress logs
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects.
+        """
+        yhat_cs = {}
+        yhat_ts = {}
+
+        for group in self.t_groups:
+            model_c = self.models_c[group]
+            model_t = self.models_t[group]
+            yhat_cs[group] = model_c.predict_proba(X)[:, 1]
+            yhat_ts[group] = model_t.predict_proba(X)[:, 1]
+
+            if (y is not None) and (treatment is not None) and verbose:
+                mask = (treatment == group) | (treatment == self.control_name)
+                treatment_filt = treatment[mask]
+                y_filt = y[mask]
+                w = (treatment_filt == group).astype(int)
+
+                yhat = np.zeros_like(y_filt, dtype=float)
+                yhat[w == 0] = yhat_cs[group][mask][w == 0]
+                yhat[w == 1] = yhat_ts[group][mask][w == 1]
+
+                logger.info("Error metrics for group {}".format(group))
+                classification_metrics(y_filt, yhat, w)
+
+        te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+        for i, group in enumerate(self.t_groups):
+            te[:, i] = yhat_ts[group] - yhat_cs[group]
+
+        if not return_components:
+            return te
+        else:
+            return te, yhat_cs, yhat_ts
+
+
+class XGBTRegressor(BaseTRegressor):
+    def __init__(self, ate_alpha=0.05, control_name=0, *args, **kwargs):
+        """Initialize a T-learner with two XGBoost models."""
+        super().__init__(
+            learner=XGBRegressor(*args, **kwargs),
+            ate_alpha=ate_alpha,
+            control_name=control_name,
+        )
+
+
+class MLPTRegressor(BaseTRegressor):
+    def __init__(self, ate_alpha=0.05, control_name=0, *args, **kwargs):
+        """Initialize a T-learner with two MLP models."""
+        super().__init__(
+            learner=MLPRegressor(*args, **kwargs),
+            ate_alpha=ate_alpha,
+            control_name=control_name,
+        )