diff --git a/.gitattributes b/.gitattributes
index a6344aac8c09253b3b630fb776ae94478aa0275b..e7ce3571fbfe93c015b0f13cbf872b872f54023f 100644
--- a/.gitattributes
+++ b/.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
diff --git a/Dockerfile b/Dockerfile
new file mode 100644
index 0000000000000000000000000000000000000000..74f3c1279df888a11f2570416a174e4e0f7e0148
--- /dev/null
+++ b/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"]
diff --git a/README.md b/README.md
index 2917ed2c3298382a167df01d31e297b3258c1df9..4605be707c4a08246651ed8846d8be6ff3414f1c 100644
--- a/README.md
+++ b/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"
+ }
+ }
+}
+```
diff --git a/app.py b/app.py
new file mode 100644
index 0000000000000000000000000000000000000000..4a2f73be8c1de6b4d63fca4f2835d165638cf1b8
--- /dev/null
+++ b/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)
diff --git a/causalml/mcp_output/README_MCP.md b/causalml/mcp_output/README_MCP.md
new file mode 100644
index 0000000000000000000000000000000000000000..9c0b9c6dda7e497427ce7c988aed97de4df2e8b4
--- /dev/null
+++ b/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`.
\ No newline at end of file
diff --git a/causalml/mcp_output/analysis.json b/causalml/mcp_output/analysis.json
new file mode 100644
index 0000000000000000000000000000000000000000..5921d30a9d7031b43dbeff59e74aa81c68f1a0d8
--- /dev/null
+++ b/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"
+ }
+}
\ No newline at end of file
diff --git a/causalml/mcp_output/diff_report.md b/causalml/mcp_output/diff_report.md
new file mode 100644
index 0000000000000000000000000000000000000000..4cf48c2109b4c672f6ebe5f0951ce658a8b9807c
--- /dev/null
+++ b/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.
\ No newline at end of file
diff --git a/causalml/mcp_output/mcp_plugin/__init__.py b/causalml/mcp_output/mcp_plugin/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/causalml/mcp_output/mcp_plugin/adapter.py b/causalml/mcp_output/mcp_plugin/adapter.py
new file mode 100644
index 0000000000000000000000000000000000000000..70a9ba49e7cff0d807076de7bd53cb121e559ce9
--- /dev/null
+++ b/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
+# -------------------------------------------------------------------------
\ No newline at end of file
diff --git a/causalml/mcp_output/mcp_plugin/main.py b/causalml/mcp_output/mcp_plugin/main.py
new file mode 100644
index 0000000000000000000000000000000000000000..fca6ec384e22f703b287550e94cc00baaaa4c4a7
--- /dev/null
+++ b/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()
\ No newline at end of file
diff --git a/causalml/mcp_output/mcp_plugin/mcp_service.py b/causalml/mcp_output/mcp_plugin/mcp_service.py
new file mode 100644
index 0000000000000000000000000000000000000000..48385ea64525d6c73a9fb49de1a1164275ca7559
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/causalml/mcp_output/requirements.txt b/causalml/mcp_output/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..b5f4dc25e406b94168459bc13be5be0a7c511696
--- /dev/null
+++ b/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
diff --git a/causalml/mcp_output/start_mcp.py b/causalml/mcp_output/start_mcp.py
new file mode 100644
index 0000000000000000000000000000000000000000..fc7fcbd9646ad53f089fc94af8129043a703325a
--- /dev/null
+++ b/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()
diff --git a/causalml/mcp_output/workflow_summary.json b/causalml/mcp_output/workflow_summary.json
new file mode 100644
index 0000000000000000000000000000000000000000..482aea0eb5bcd28fc6e294f3d3e9ffc122268e0c
--- /dev/null
+++ b/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
+ }
+}
\ No newline at end of file
diff --git a/causalml/source/.pre-commit-config.yaml b/causalml/source/.pre-commit-config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..48ee0308d3ccab3f18b36291d2b8d1183a4a563e
--- /dev/null
+++ b/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
diff --git a/causalml/source/.readthedocs.yml b/causalml/source/.readthedocs.yml
new file mode 100644
index 0000000000000000000000000000000000000000..8da3cc7ec2f01351ef762bcd73bb4a3ab9bd9be4
--- /dev/null
+++ b/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
diff --git a/causalml/source/ANTITRUST.md b/causalml/source/ANTITRUST.md
new file mode 100644
index 0000000000000000000000000000000000000000..b44893eec5f0f2a0a03695ffaaa82399db708985
--- /dev/null
+++ b/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/).
diff --git a/causalml/source/CHARTER.md b/causalml/source/CHARTER.md
new file mode 100644
index 0000000000000000000000000000000000000000..f529693ab37f2bd4ce6d3b7a2696c4ea8c7cc164
--- /dev/null
+++ b/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/).
diff --git a/causalml/source/CODE_OF_CONDUCT.md b/causalml/source/CODE_OF_CONDUCT.md
new file mode 100644
index 0000000000000000000000000000000000000000..e327d9aa5cd0a1ab2a9f3d602f2ecefa6e7f72c1
--- /dev/null
+++ b/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/
diff --git a/causalml/source/CONTRIBUTING.md b/causalml/source/CONTRIBUTING.md
new file mode 100644
index 0000000000000000000000000000000000000000..fec4445eff5481ccec7eaf35ff075d64e03ee2ff
--- /dev/null
+++ b/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.
diff --git a/causalml/source/GOVERNANCE.md b/causalml/source/GOVERNANCE.md
new file mode 100644
index 0000000000000000000000000000000000000000..937f22c9e85c87a0253d7510933989a3385157af
--- /dev/null
+++ b/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/).
diff --git a/causalml/source/LICENSE b/causalml/source/LICENSE
new file mode 100644
index 0000000000000000000000000000000000000000..9b0216ed73224bf357530eee275caf57e5c3b1df
--- /dev/null
+++ b/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.
\ No newline at end of file
diff --git a/causalml/source/MAINTAINERS.md b/causalml/source/MAINTAINERS.md
new file mode 100644
index 0000000000000000000000000000000000000000..1aed73ec33fcb33ee42f196422cc644ea4099362
--- /dev/null
+++ b/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/).
diff --git a/causalml/source/MANIFEST.in b/causalml/source/MANIFEST.in
new file mode 100644
index 0000000000000000000000000000000000000000..358d03575ad93ed849149c2313e603148214aeff
--- /dev/null
+++ b/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
diff --git a/causalml/source/Makefile b/causalml/source/Makefile
new file mode 100644
index 0000000000000000000000000000000000000000..69150bbb92e168012774871b74310ea11f1cc3f2
--- /dev/null
+++ b/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
diff --git a/causalml/source/README.md b/causalml/source/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..737750800a3b47624768b5b0eae1857e364621b7
--- /dev/null
+++ b/causalml/source/README.md
@@ -0,0 +1,132 @@
+
+
+
+
+------------------------------------------------------
+
+[](https://pypi.org/project/causalml/)
+[](https://github.com/uber/causalml/actions/workflows/python-test.yaml)
+[](http://causalml.readthedocs.io/en/latest/?badge=latest)
+[](https://pepy.tech/project/causalml)
+[](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
diff --git a/causalml/source/SECURITY.md b/causalml/source/SECURITY.md
new file mode 100644
index 0000000000000000000000000000000000000000..d4f13550698f9de12af3f822b331b722bca88469
--- /dev/null
+++ b/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
diff --git a/causalml/source/STEERING_COMMITTEE.md b/causalml/source/STEERING_COMMITTEE.md
new file mode 100644
index 0000000000000000000000000000000000000000..aa4efd60529b4ae922b5020c6c608d37e56a3a78
--- /dev/null
+++ b/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/).
diff --git a/causalml/source/TRADEMARKS.md b/causalml/source/TRADEMARKS.md
new file mode 100644
index 0000000000000000000000000000000000000000..d58e223eba55aeca3c92f4f8edf7bbe1e0508862
--- /dev/null
+++ b/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)
diff --git a/causalml/source/__init__.py b/causalml/source/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..4835847f8291fd6e38998f9d35e337876fef6cfa
--- /dev/null
+++ b/causalml/source/__init__.py
@@ -0,0 +1,4 @@
+# -*- coding: utf-8 -*-
+"""
+causalml Project Package Initialization File
+"""
diff --git a/causalml/source/causalml/__init__.py b/causalml/source/causalml/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..129cba0004e3946c7ae78faa6497e3c24cc5af2a
--- /dev/null
+++ b/causalml/source/causalml/__init__.py
@@ -0,0 +1,10 @@
+__all__ = [
+ "dataset",
+ "features",
+ "feature_selection",
+ "inference",
+ "match",
+ "metrics",
+ "optimize",
+ "propensity",
+]
diff --git a/causalml/source/causalml/dataset/__init__.py b/causalml/source/causalml/dataset/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..f4678b11d78e69110d5598e1f84c58e51d9c5762
--- /dev/null
+++ b/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
diff --git a/causalml/source/causalml/dataset/classification.py b/causalml/source/causalml/dataset/classification.py
new file mode 100644
index 0000000000000000000000000000000000000000..892cdb6acd7d2ef72646afd2149a01d13c3e0be8
--- /dev/null
+++ b/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
diff --git a/causalml/source/causalml/dataset/regression.py b/causalml/source/causalml/dataset/regression.py
new file mode 100644
index 0000000000000000000000000000000000000000..0beb3ce2d01925981a8dbc14aefde833fcae51d1
--- /dev/null
+++ b/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
diff --git a/causalml/source/causalml/dataset/semiSynthetic.py b/causalml/source/causalml/dataset/semiSynthetic.py
new file mode 100644
index 0000000000000000000000000000000000000000..20a284860d2b77ce23a32aa598f3c05a82cd6414
--- /dev/null
+++ b/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
diff --git a/causalml/source/causalml/dataset/synthetic.py b/causalml/source/causalml/dataset/synthetic.py
new file mode 100644
index 0000000000000000000000000000000000000000..983784fd335dbc78b364ff3277be67470e73166e
--- /dev/null
+++ b/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
diff --git a/causalml/source/causalml/feature_selection/__init__.py b/causalml/source/causalml/feature_selection/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..3c8a623c494dabdd4f3e7572dd7e98107ae0985f
--- /dev/null
+++ b/causalml/source/causalml/feature_selection/__init__.py
@@ -0,0 +1 @@
+from .filters import FilterSelect
diff --git a/causalml/source/causalml/feature_selection/filters.py b/causalml/source/causalml/feature_selection/filters.py
new file mode 100644
index 0000000000000000000000000000000000000000..f92497887a02771e2693bb75d13a5033b8d4443b
--- /dev/null
+++ b/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"]]
diff --git a/causalml/source/causalml/features.py b/causalml/source/causalml/features.py
new file mode 100644
index 0000000000000000000000000000000000000000..f4b1e1c2c9cb089ae0d8b16ad7b5d78d822c2c23
--- /dev/null
+++ b/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
diff --git a/causalml/source/causalml/inference/__init__.py b/causalml/source/causalml/inference/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/causalml/source/causalml/inference/iv/__init__.py b/causalml/source/causalml/inference/iv/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..544379ac2806b9b761334855fd288058a5ea9f0a
--- /dev/null
+++ b/causalml/source/causalml/inference/iv/__init__.py
@@ -0,0 +1,2 @@
+from .iv_regression import IVRegressor
+from .drivlearner import BaseDRIVLearner, BaseDRIVRegressor, XGBDRIVRegressor
diff --git a/causalml/source/causalml/inference/iv/drivlearner.py b/causalml/source/causalml/inference/iv/drivlearner.py
new file mode 100644
index 0000000000000000000000000000000000000000..bdc6d21d95069e732ab2baa9116c135c542f756b
--- /dev/null
+++ b/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) `_.
+ The DR moment condition for LATE comes from
+ `Chernozhukov et al (2018) `_.
+ """
+
+ 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,
+ )
diff --git a/causalml/source/causalml/inference/iv/iv_regression.py b/causalml/source/causalml/inference/iv/iv_regression.py
new file mode 100644
index 0000000000000000000000000000000000000000..612c8b8e30f194c7d096e33ee4c1e299829a964d
--- /dev/null
+++ b/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]
diff --git a/causalml/source/causalml/inference/meta/__init__.py b/causalml/source/causalml/inference/meta/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..2127e5e6dd0a54ce0f000b74d78fa9b5700f302f
--- /dev/null
+++ b/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
diff --git a/causalml/source/causalml/inference/meta/base.py b/causalml/source/causalml/inference/meta/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..3f509ceac8834587b79d97b7388b785a4029f96e
--- /dev/null
+++ b/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,
+ )
diff --git a/causalml/source/causalml/inference/meta/drlearner.py b/causalml/source/causalml/inference/meta/drlearner.py
new file mode 100644
index 0000000000000000000000000000000000000000..300a1aa4482e33d2612df14ce33e6e78c1349600
--- /dev/null
+++ b/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) `_.
+ """
+
+ 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,
+ )
diff --git a/causalml/source/causalml/inference/meta/explainer.py b/causalml/source/causalml/inference/meta/explainer.py
new file mode 100644
index 0000000000000000000000000000000000000000..b92fee8603fb065f1439105764699cfb7cdf13da
--- /dev/null
+++ b/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,
+ )
diff --git a/causalml/source/causalml/inference/meta/rlearner.py b/causalml/source/causalml/inference/meta/rlearner.py
new file mode 100644
index 0000000000000000000000000000000000000000..53563ccec2ee319b393eb0115590388a66b5d2b5
--- /dev/null
+++ b/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) `_.
+ """
+
+ 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)
diff --git a/causalml/source/causalml/inference/meta/slearner.py b/causalml/source/causalml/inference/meta/slearner.py
new file mode 100644
index 0000000000000000000000000000000000000000..796ac11f9e5dad5604b3df06987e36b6033845b5
--- /dev/null
+++ b/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) `_.
+ """
+
+ 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
diff --git a/causalml/source/causalml/inference/meta/tlearner.py b/causalml/source/causalml/inference/meta/tlearner.py
new file mode 100644
index 0000000000000000000000000000000000000000..04ca796f346708ce4ad205a4229ca320426fdb2c
--- /dev/null
+++ b/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) `_.
+ """
+
+ 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,
+ )
diff --git a/causalml/source/causalml/inference/meta/tmle.py b/causalml/source/causalml/inference/meta/tmle.py
new file mode 100644
index 0000000000000000000000000000000000000000..372d0f3d72624ffd147bdaf8910fbb55005cd29e
--- /dev/null
+++ b/causalml/source/causalml/inference/meta/tmle.py
@@ -0,0 +1,221 @@
+import logging
+import numpy as np
+import pandas as pd
+from scipy.optimize import minimize
+from scipy.special import expit, logit
+from scipy.stats import norm
+from sklearn.preprocessing import MinMaxScaler
+
+from causalml.inference.meta.utils import (
+ check_treatment_vector,
+ check_p_conditions,
+ convert_pd_to_np,
+)
+
+logger = logging.getLogger("causalml")
+
+
+def logit_tmle(x, y, a, h0, h1):
+ p = expit(a + x[0] * h0 + x[1] * h1)
+ return np.mean(-np.log(np.power(p, y) * np.power(1 - p, 1 - y)))
+
+
+def logit_tmle_grad(x, y, a, h0, h1):
+ p = expit(a + x[0] * h0 + x[1] * h1)
+ return np.array([-np.mean((y - p) * h0), -np.mean((y - p) * h1)])
+
+
+def logit_tmle_hess(x, y, a, h0, h1):
+ p = expit(a + x[0] * h0 + x[1] * h1)
+ return np.array(
+ [
+ [np.mean(p * (1 - p) * h0 * h0), np.mean(p * (1 - p) * h0 * h1)],
+ [np.mean(p * (1 - p) * h0 * h1), np.mean(p * (1 - p) * h1 * h1)],
+ ]
+ )
+
+
+def simple_tmle(y, w, q0w, q1w, p, alpha=0.0001):
+ """Calculate the ATE and variances with the simplified TMLE method.
+
+ Args:
+ y (numpy.array): an outcome vector
+ w (numpy.array): a treatment vector
+ q0w (numpy.array): an outcome prediction vector given no treatment
+ q1w (numpy.array): an outcome prediction vector given treatment
+ p (numpy.array): a propensity score vector
+ alpha (float, optional): a clipping threshold for predictions
+
+ Returns:
+ (tuple)
+
+ - ate (float): ATE
+ - se (float): The standard error of ATE
+ """
+ scaler = MinMaxScaler()
+ ystar = scaler.fit_transform(y.reshape(-1, 1)).flatten()
+
+ q0 = np.clip(scaler.transform(q0w.reshape(-1, 1)).flatten(), alpha, 1 - alpha)
+ q1 = np.clip(scaler.transform(q1w.reshape(-1, 1)).flatten(), alpha, 1 - alpha)
+ qaw = q0 * (1 - w) + q1 * w
+ intercept = logit(qaw)
+
+ h1 = w / p
+ h0 = (1 - w) / (1 - p)
+ sol = minimize(
+ logit_tmle,
+ np.zeros(2),
+ args=(ystar, intercept, h0, h1),
+ method="Newton-CG",
+ jac=logit_tmle_grad,
+ hess=logit_tmle_hess,
+ )
+
+ qawstar = scaler.inverse_transform(
+ expit(intercept + sol.x[0] * h0 + sol.x[1] * h1).reshape(-1, 1)
+ ).flatten()
+ q0star = scaler.inverse_transform(
+ expit(logit(q0) + sol.x[0] / (1 - p)).reshape(-1, 1)
+ ).flatten()
+ q1star = scaler.inverse_transform(
+ expit(logit(q1) + sol.x[1] / p).reshape(-1, 1)
+ ).flatten()
+
+ ic = (
+ (w / p - (1 - w) / (1 - p)) * (y - qawstar)
+ + q1star
+ - q0star
+ - np.mean(q1star - q0star)
+ )
+
+ return np.mean(q1star - q0star), np.sqrt(np.var(ic) / np.size(y))
+
+
+class TMLELearner:
+ """Targeted maximum likelihood estimation.
+
+ Ref: Gruber, S., & Van Der Laan, M. J. (2009). Targeted maximum likelihood estimation: A gentle introduction.
+ """
+
+ def __init__(
+ self,
+ learner,
+ ate_alpha=0.05,
+ control_name=0,
+ cv=None,
+ ):
+ """Initialize a TMLE learner.
+
+ Args:
+ learner: a model to estimate the outcome
+ ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+ control_name (str or int, optional): the name of the control group
+ cv (sklearn.model_selection._BaseKFold, optional): sklearn CV object
+ """
+ self.model_tau = learner
+ self.ate_alpha = ate_alpha
+ self.control_name = control_name
+ self.cv = cv
+
+ def __repr__(self):
+ return "{}(model={}, cv={})".format(
+ self.__class__.__name__, self.model_tau.__repr__(), self.cv
+ )
+
+ def estimate_ate(self, X, treatment, y, p, segment=None, return_ci=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): 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)
+ segment (np.array, optional): An optional segment vector of int. If given, the ATE and its CI will be
+ estimated for each segment.
+ return_ci (bool, optional): Whether to return confidence intervals
+
+ Returns:
+ (tuple): The ATE and its confidence interval (LB, UB) for each treatment, t and segment, s
+ """
+ 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()
+
+ 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 = []
+ ate_lb = []
+ ate_ub = []
+
+ for _, group in enumerate(self.t_groups):
+ logger.info("Estimating ATE for group {}.".format(group))
+ w_group = (treatment == group).astype(int)
+ p_group = p[group]
+
+ yhat_c = np.zeros_like(y, dtype=float)
+ yhat_t = np.zeros_like(y, dtype=float)
+ if self.cv:
+ for i_fold, (i_trn, i_val) in enumerate(self.cv.split(X, y), 1):
+ logger.info("Training an outcome model for CV #{}".format(i_fold))
+ self.model_tau.fit(
+ np.hstack((X[i_trn], w_group[i_trn].reshape(-1, 1))), y[i_trn]
+ )
+
+ yhat_c[i_val] = self.model_tau.predict(
+ np.hstack((X[i_val], np.zeros((len(i_val), 1))))
+ )
+ yhat_t[i_val] = self.model_tau.predict(
+ np.hstack((X[i_val], np.ones((len(i_val), 1))))
+ )
+
+ else:
+ self.model_tau.fit(np.hstack((X, w_group.reshape(-1, 1))), y)
+
+ yhat_c = self.model_tau.predict(np.hstack((X, np.zeros((len(y), 1)))))
+ yhat_t = self.model_tau.predict(np.hstack((X, np.ones((len(y), 1)))))
+
+ if segment is None:
+ logger.info("Training the TMLE learner.")
+ _ate, se = simple_tmle(y, w_group, yhat_c, yhat_t, p_group)
+ _ate_lb = _ate - se * norm.ppf(1 - self.ate_alpha / 2)
+ _ate_ub = _ate + se * norm.ppf(1 - self.ate_alpha / 2)
+ else:
+ assert (
+ segment.shape[0] == X.shape[0] and segment.ndim == 1
+ ), "Segment must be the 1-d np.array of int."
+ segments = np.unique(segment)
+
+ _ate = []
+ _ate_lb = []
+ _ate_ub = []
+ for s in sorted(segments):
+ logger.info("Training the TMLE learner for segment {}.".format(s))
+ filt = (segment == s) & (yhat_c < np.quantile(yhat_c, q=0.99))
+ _ate_s, se = simple_tmle(
+ y[filt],
+ w_group[filt],
+ yhat_c[filt],
+ yhat_t[filt],
+ p_group[filt],
+ )
+ _ate_lb_s = _ate_s - se * norm.ppf(1 - self.ate_alpha / 2)
+ _ate_ub_s = _ate_s + se * norm.ppf(1 - self.ate_alpha / 2)
+
+ _ate.append(_ate_s)
+ _ate_lb.append(_ate_lb_s)
+ _ate_ub.append(_ate_ub_s)
+
+ ate.append(_ate)
+ ate_lb.append(_ate_lb)
+ ate_ub.append(_ate_ub)
+
+ return np.array(ate), np.array(ate_lb), np.array(ate_ub)
diff --git a/causalml/source/causalml/inference/meta/utils.py b/causalml/source/causalml/inference/meta/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..157eeaf6ed3daea5cfcf1ca6603f5737d8ccc365
--- /dev/null
+++ b/causalml/source/causalml/inference/meta/utils.py
@@ -0,0 +1,136 @@
+import pandas as pd
+import numpy as np
+
+from packaging import version
+from xgboost import __version__ as xgboost_version
+
+
+def convert_pd_to_np(*args):
+ output = [obj.to_numpy() if hasattr(obj, "to_numpy") else obj for obj in args]
+ return output if len(output) > 1 else output[0]
+
+
+def check_treatment_vector(treatment, control_name=None):
+ n_unique_treatments = np.unique(treatment).shape[0]
+ assert n_unique_treatments > 1, "Treatment vector must have at least two levels."
+ if control_name is not None:
+ assert (
+ control_name in treatment
+ ), "Control group level {} not found in treatment vector.".format(control_name)
+
+
+def check_p_conditions(p, t_groups):
+ eps = np.finfo(float).eps
+ assert isinstance(
+ p, (np.ndarray, pd.Series, dict)
+ ), "p must be an np.ndarray, pd.Series, or dict type"
+ if isinstance(p, (np.ndarray, pd.Series)):
+ assert (
+ t_groups.shape[0] == 1
+ ), "If p is passed as an np.ndarray, there must be only 1 unique non-control group in the treatment vector."
+ assert (0 + eps < p).all() and (
+ p < 1 - eps
+ ).all(), "The values of p should lie within the (0, 1) interval."
+
+ if isinstance(p, dict):
+ for t_name in t_groups:
+ assert (0 + eps < p[t_name]).all() and (
+ p[t_name] < 1 - eps
+ ).all(), "The values of p should lie within the (0, 1) interval."
+
+
+def check_explain_conditions(method, models, X=None, treatment=None, y=None):
+ valid_methods = ["gini", "permutation", "shapley"]
+ assert method in valid_methods, "Current supported methods: {}".format(
+ ", ".join(valid_methods)
+ )
+
+ if method in ("gini", "shapley"):
+ conds = [hasattr(mod, "feature_importances_") for mod in models]
+ assert all(
+ conds
+ ), "Both models must have .feature_importances_ attribute if method = {}".format(
+ method
+ )
+
+ if method in ("permutation", "shapley"):
+ assert all(
+ arr is not None for arr in (X, treatment, y)
+ ), "X, treatment, and y must be provided if method = {}".format(method)
+
+
+def clean_xgboost_objective(objective):
+ """
+ Translate objective to be compatible with loaded xgboost version
+
+ Args
+ ----
+
+ objective : string
+ The objective to translate.
+
+ Returns
+ -------
+ The translated objective, or original if no translation was required.
+ """
+ compat_before_v83 = {"reg:squarederror": "reg:linear"}
+ compat_v83_or_later = {"reg:linear": "reg:squarederror"}
+ if version.parse(xgboost_version) < version.parse("0.83"):
+ if objective in compat_before_v83:
+ objective = compat_before_v83[objective]
+ else:
+ if objective in compat_v83_or_later:
+ objective = compat_v83_or_later[objective]
+ return objective
+
+
+def get_xgboost_objective_metric(objective):
+ """
+ Get the xgboost version-compatible objective and evaluation metric from a potentially version-incompatible input.
+
+ Args
+ ----
+
+ objective : string
+ An xgboost objective that may be incompatible with the installed version.
+
+ Returns
+ -------
+ A tuple with the translated objective and evaluation metric.
+ """
+
+ def clean_dict_keys(orig):
+ return {clean_xgboost_objective(k): v for (k, v) in orig.items()}
+
+ metric_mapping = clean_dict_keys(
+ {"rank:pairwise": "auc", "reg:squarederror": "rmse"}
+ )
+
+ objective = clean_xgboost_objective(objective)
+
+ assert (
+ objective in metric_mapping
+ ), "Effect learner objective must be one of: " + ", ".join(metric_mapping)
+ return objective, metric_mapping[objective]
+
+
+def get_weighted_variance(x, sample_weight):
+ """
+ Calculate the variance of array x with sample_weight.
+
+ Args
+ ----
+
+ x : (np.array)
+ A list of number
+
+ sample_weight (np.array or list): an array of sample weights indicating the
+ weight of each observation for `effect_learner`. If None, it assumes equal weight.
+
+ Returns
+ -------
+ The variance of x with sample weight
+ """
+ average = np.average(x, weights=sample_weight)
+ variance = np.average((x - average) ** 2, weights=sample_weight)
+ return variance
diff --git a/causalml/source/causalml/inference/meta/xlearner.py b/causalml/source/causalml/inference/meta/xlearner.py
new file mode 100644
index 0000000000000000000000000000000000000000..88b5dc1da45cb4071641b6f3f70de9a25a7ab0d4
--- /dev/null
+++ b/causalml/source/causalml/inference/meta/xlearner.py
@@ -0,0 +1,642 @@
+from copy import deepcopy
+import logging
+import numpy as np
+from tqdm import tqdm
+from scipy.stats import norm
+
+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 BaseXLearner(BaseLearner):
+ """A parent class for X-learner regressor classes.
+
+ An X-learner estimates treatment effects with four machine learning models.
+
+ Details of X-learner are available at `Kunzel et al. (2018) `_.
+ """
+
+ def __init__(
+ self,
+ learner=None,
+ control_outcome_learner=None,
+ treatment_outcome_learner=None,
+ control_effect_learner=None,
+ treatment_effect_learner=None,
+ ate_alpha=0.05,
+ control_name=0,
+ ):
+ """Initialize a X-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
+ control_effect_learner (optional): a model to estimate treatment effects in the control 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 (control_effect_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 control_effect_learner is None:
+ self.model_tau_c = deepcopy(learner)
+ else:
+ self.model_tau_c = control_effect_learner
+
+ if treatment_effect_learner is None:
+ self.model_tau_t = deepcopy(learner)
+ else:
+ self.model_tau_t = treatment_effect_learner
+
+ self.ate_alpha = ate_alpha
+ self.control_name = control_name
+
+ self.propensity = None
+ self.propensity_model = None
+
+ def __repr__(self):
+ return (
+ "{}(control_outcome_learner={},\n"
+ "\ttreatment_outcome_learner={},\n"
+ "\tcontrol_effect_learner={},\n"
+ "\ttreatment_effect_learner={})".format(
+ self.__class__.__name__,
+ self.model_mu_c.__repr__(),
+ self.model_mu_t.__repr__(),
+ self.model_tau_c.__repr__(),
+ self.model_tau_t.__repr__(),
+ )
+ )
+
+ 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
+ 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.
+ """
+ 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()
+
+ 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_mu_c = {group: deepcopy(self.model_mu_c) for group in self.t_groups}
+ self.models_mu_t = {group: deepcopy(self.model_mu_t) for group in self.t_groups}
+ self.models_tau_c = {
+ group: deepcopy(self.model_tau_c) for group in self.t_groups
+ }
+ self.models_tau_t = {
+ group: deepcopy(self.model_tau_t) for group in self.t_groups
+ }
+ self.vars_c = {}
+ self.vars_t = {}
+
+ 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)
+
+ # Train outcome models
+ self.models_mu_c[group].fit(X_filt[w == 0], y_filt[w == 0])
+ self.models_mu_t[group].fit(X_filt[w == 1], y_filt[w == 1])
+
+ # Calculate variances and treatment effects
+ var_c = (
+ y_filt[w == 0] - self.models_mu_c[group].predict(X_filt[w == 0])
+ ).var()
+ self.vars_c[group] = var_c
+ var_t = (
+ y_filt[w == 1] - self.models_mu_t[group].predict(X_filt[w == 1])
+ ).var()
+ self.vars_t[group] = var_t
+
+ # Train treatment models
+ d_c = self.models_mu_t[group].predict(X_filt[w == 0]) - y_filt[w == 0]
+ d_t = y_filt[w == 1] - self.models_mu_c[group].predict(X_filt[w == 1])
+ self.models_tau_c[group].fit(X_filt[w == 0], d_c)
+ self.models_tau_t[group].fit(X_filt[w == 1], d_t)
+
+ 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
+ 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_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)
+
+ if p is None:
+ logger.info("Generating propensity score")
+ p = dict()
+ for group in self.t_groups:
+ p_model = self.propensity_model[group]
+ p[group] = p_model.predict(X)
+ else:
+ p = self._format_p(p, self.t_groups)
+
+ te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+ dhat_cs = {}
+ dhat_ts = {}
+
+ for i, group in enumerate(self.t_groups):
+ model_tau_c = self.models_tau_c[group]
+ model_tau_t = self.models_tau_t[group]
+ dhat_cs[group] = model_tau_c.predict(X)
+ dhat_ts[group] = model_tau_t.predict(X)
+
+ _te = (p[group] * dhat_cs[group] + (1 - p[group]) * dhat_ts[group]).reshape(
+ -1, 1
+ )
+ te[:, i] = np.ravel(_te)
+
+ if (y is not None) and (treatment is not None) and verbose:
+ 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)
+
+ yhat = np.zeros_like(y_filt, dtype=float)
+ yhat[w == 0] = self.models_mu_c[group].predict(X_filt[w == 0])
+ yhat[w == 1] = self.models_mu_t[group].predict(X_filt[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, dhat_cs, dhat_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 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
+ 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)
+
+ if p is None:
+ p = self.propensity
+ else:
+ p = self._format_p(p, self.t_groups)
+
+ te = self.predict(
+ X, treatment=treatment, y=y, p=p, 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_c_global = deepcopy(self.models_tau_c)
+ models_tau_t_global = deepcopy(self.models_tau_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, 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_c = deepcopy(models_tau_c_global)
+ self.models_tau_t = deepcopy(models_tau_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
+ 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
+ pretrain (bool): whether a model has been fit, default False.
+ Returns:
+ The mean and confidence interval (LB, UB) of the ATE estimate.
+ """
+ if pretrain:
+ if p is None:
+ # when p is null, use pretrain propensity score
+ if not self.propensity:
+ raise ValueError("no propensity score, please call fit() first")
+ te, dhat_cs, dhat_ts = self.predict(
+ X, treatment, y, p=self.propensity, return_components=True
+ )
+ else:
+ p = self._format_p(p, self.t_groups)
+ te, dhat_cs, dhat_ts = self.predict(
+ X, treatment, y, p=p, return_components=True
+ )
+ else:
+ te, dhat_cs, dhat_ts = self.fit_predict(
+ X, treatment, y, p, return_components=True
+ )
+ X, treatment, y = convert_pd_to_np(X, treatment, y)
+
+ if p is None:
+ p = self.propensity
+ else:
+ p = self._format_p(p, self.t_groups)
+
+ 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]
+
+ dhat_c = dhat_cs[group][mask]
+ dhat_t = dhat_ts[group][mask]
+ p_filt = p[group][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(
+ (
+ self.vars_t[group] / prob_treatment
+ + self.vars_c[group] / (1 - prob_treatment)
+ + (p_filt * dhat_c + (1 - p_filt) * dhat_t).var()
+ )
+ / w.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_c_global = deepcopy(self.models_tau_c)
+ models_tau_t_global = deepcopy(self.models_tau_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)):
+ 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.models_mu_c = deepcopy(models_mu_c_global)
+ self.models_mu_t = deepcopy(models_mu_t_global)
+ self.models_tau_c = deepcopy(models_tau_c_global)
+ self.models_tau_t = deepcopy(models_tau_t_global)
+ return ate, ate_lower, ate_upper
+
+
+class BaseXRegressor(BaseXLearner):
+ """
+ A parent class for X-learner regressor classes.
+ """
+
+ def __init__(
+ self,
+ learner=None,
+ control_outcome_learner=None,
+ treatment_outcome_learner=None,
+ control_effect_learner=None,
+ treatment_effect_learner=None,
+ ate_alpha=0.05,
+ control_name=0,
+ ):
+ """Initialize an X-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
+ control_effect_learner (optional): a model to estimate treatment effects in the control 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,
+ control_effect_learner=control_effect_learner,
+ treatment_effect_learner=treatment_effect_learner,
+ ate_alpha=ate_alpha,
+ control_name=control_name,
+ )
+
+
+class BaseXClassifier(BaseXLearner):
+ """
+ A parent class for X-learner classifier classes.
+ """
+
+ def __init__(
+ self,
+ outcome_learner=None,
+ effect_learner=None,
+ control_outcome_learner=None,
+ treatment_outcome_learner=None,
+ control_effect_learner=None,
+ treatment_effect_learner=None,
+ ate_alpha=0.05,
+ control_name=0,
+ ):
+ """Initialize an X-learner classifier.
+
+ Args:
+ outcome_learner (optional): a model to estimate outcomes in both the control and treatment groups.
+ Should be a classifier.
+ effect_learner (optional): a model to estimate treatment effects in both the control and treatment groups.
+ Should be a regressor.
+ control_outcome_learner (optional): a model to estimate outcomes in the control group.
+ Should be a classifier.
+ treatment_outcome_learner (optional): a model to estimate outcomes in the treatment group.
+ Should be a classifier.
+ control_effect_learner (optional): a model to estimate treatment effects in the control group.
+ Should be a regressor.
+ 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
+ """
+ if outcome_learner is not None:
+ control_outcome_learner = outcome_learner
+ treatment_outcome_learner = outcome_learner
+ if effect_learner is not None:
+ control_effect_learner = effect_learner
+ treatment_effect_learner = effect_learner
+
+ super().__init__(
+ learner=None,
+ control_outcome_learner=control_outcome_learner,
+ treatment_outcome_learner=treatment_outcome_learner,
+ control_effect_learner=control_effect_learner,
+ treatment_effect_learner=treatment_effect_learner,
+ ate_alpha=ate_alpha,
+ control_name=control_name,
+ )
+
+ if (
+ (control_outcome_learner is None) or (treatment_outcome_learner is None)
+ ) and ((control_effect_learner is None) or (treatment_effect_learner is None)):
+ raise ValueError(
+ "Either the outcome learner or the effect learner pair must be specified."
+ )
+
+ 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
+ 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.
+ """
+ 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()
+
+ 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_mu_c = {group: deepcopy(self.model_mu_c) for group in self.t_groups}
+ self.models_mu_t = {group: deepcopy(self.model_mu_t) for group in self.t_groups}
+ self.models_tau_c = {
+ group: deepcopy(self.model_tau_c) for group in self.t_groups
+ }
+ self.models_tau_t = {
+ group: deepcopy(self.model_tau_t) for group in self.t_groups
+ }
+ self.vars_c = {}
+ self.vars_t = {}
+
+ 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)
+
+ # Train outcome models
+ self.models_mu_c[group].fit(X_filt[w == 0], y_filt[w == 0])
+ self.models_mu_t[group].fit(X_filt[w == 1], y_filt[w == 1])
+
+ # Calculate variances and treatment effects
+ var_c = (
+ y_filt[w == 0]
+ - self.models_mu_c[group].predict_proba(X_filt[w == 0])[:, 1]
+ ).var()
+ self.vars_c[group] = var_c
+ var_t = (
+ y_filt[w == 1]
+ - self.models_mu_t[group].predict_proba(X_filt[w == 1])[:, 1]
+ ).var()
+ self.vars_t[group] = var_t
+
+ # Train treatment models
+ d_c = (
+ self.models_mu_t[group].predict_proba(X_filt[w == 0])[:, 1]
+ - y_filt[w == 0]
+ )
+ d_t = (
+ y_filt[w == 1]
+ - self.models_mu_c[group].predict_proba(X_filt[w == 1])[:, 1]
+ )
+ self.models_tau_c[group].fit(X_filt[w == 0], d_c)
+ self.models_tau_t[group].fit(X_filt[w == 1], d_t)
+
+ 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
+ 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_components (bool, optional): whether to return outcome for treatment and control seperately
+ return_p_score (bool, optional): whether to return propensity score
+ 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)
+
+ if p is None:
+ logger.info("Generating propensity score")
+ p = dict()
+ for group in self.t_groups:
+ p_model = self.propensity_model[group]
+ p[group] = p_model.predict(X)
+ else:
+ p = self._format_p(p, self.t_groups)
+
+ te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+ dhat_cs = {}
+ dhat_ts = {}
+
+ for i, group in enumerate(self.t_groups):
+ model_tau_c = self.models_tau_c[group]
+ model_tau_t = self.models_tau_t[group]
+ dhat_cs[group] = model_tau_c.predict(X)
+ dhat_ts[group] = model_tau_t.predict(X)
+
+ _te = (p[group] * dhat_cs[group] + (1 - p[group]) * dhat_ts[group]).reshape(
+ -1, 1
+ )
+ te[:, i] = np.ravel(_te)
+
+ if (y is not None) and (treatment is not None) and verbose:
+ 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)
+
+ yhat = np.zeros_like(y_filt, dtype=float)
+ yhat[w == 0] = self.models_mu_c[group].predict_proba(X_filt[w == 0])[
+ :, 1
+ ]
+ yhat[w == 1] = self.models_mu_t[group].predict_proba(X_filt[w == 1])[
+ :, 1
+ ]
+
+ logger.info("Error metrics for group {}".format(group))
+ classification_metrics(y_filt, yhat, w)
+
+ if not return_components:
+ return te
+ else:
+ return te, dhat_cs, dhat_ts
diff --git a/causalml/source/causalml/inference/tf/__init__.py b/causalml/source/causalml/inference/tf/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..27407c70c28a7909385ae9be2d1564b97aa83b11
--- /dev/null
+++ b/causalml/source/causalml/inference/tf/__init__.py
@@ -0,0 +1 @@
+from .dragonnet import DragonNet
diff --git a/causalml/source/causalml/inference/tf/dragonnet.py b/causalml/source/causalml/inference/tf/dragonnet.py
new file mode 100644
index 0000000000000000000000000000000000000000..f8cd3ed3665ac4224571e9bc646138c20936869e
--- /dev/null
+++ b/causalml/source/causalml/inference/tf/dragonnet.py
@@ -0,0 +1,326 @@
+"""
+This module implements the Dragonnet [1], which adapts the design and training of neural networks to improve
+the quality of treatment effect estimates. The authors propose two adaptations:
+
+- A new architecture, the Dragonnet, that exploits the sufficiency of the propensity score for estimation adjustment.
+- A regularization procedure, targeted regularization, that induces a bias towards models that have non-parametrically
+ optimal asymptotic properties ‘out-of-the-box’. Studies on benchmark datasets for causal inference show these
+ adaptations outperform existing methods. Code is available at github.com/claudiashi57/dragonnet
+
+**References**
+
+[1] C. Shi, D. Blei, V. Veitch (2019).
+ | Adapting Neural Networks for the Estimation of Treatment Effects.
+ | https://arxiv.org/pdf/1906.02120.pdf
+ | https://github.com/claudiashi57/dragonnet
+"""
+
+import numpy as np
+from tensorflow.keras import Input, Model
+from tensorflow.keras.callbacks import EarlyStopping, ReduceLROnPlateau, TerminateOnNaN
+from tensorflow.keras.layers import Dense, Concatenate
+from tensorflow.keras.optimizers import SGD, Adam
+from tensorflow.keras.regularizers import l2
+from tensorflow.keras.models import load_model
+
+from causalml.inference.tf.utils import (
+ dragonnet_loss_binarycross,
+ EpsilonLayer,
+ regression_loss,
+ binary_classification_loss,
+ treatment_accuracy,
+ track_epsilon,
+ make_tarreg_loss,
+)
+from causalml.inference.meta.utils import convert_pd_to_np
+
+
+class DragonNet:
+ def __init__(
+ self,
+ neurons_per_layer=200,
+ targeted_reg=True,
+ ratio=1.0,
+ val_split=0.2,
+ batch_size=64,
+ epochs=100,
+ learning_rate=1e-5,
+ momentum=0.9,
+ reg_l2=0.01,
+ use_adam=True,
+ adam_epochs=30,
+ adam_learning_rate=1e-3,
+ loss_func=dragonnet_loss_binarycross,
+ verbose=True,
+ ):
+ """
+ Initializes a Dragonnet.
+ """
+ self.neurons_per_layer = neurons_per_layer
+ self.targeted_reg = targeted_reg
+ self.ratio = ratio
+ self.val_split = val_split
+ self.batch_size = batch_size
+ self.epochs = epochs
+ self.learning_rate = learning_rate
+ self.momentum = momentum
+ self.use_adam = use_adam
+ self.adam_learning_rate = adam_learning_rate
+ self.adam_epochs = adam_epochs
+ self.reg_l2 = reg_l2
+ self.loss_func = loss_func
+ self.verbose = verbose
+
+ def make_dragonnet(self, input_dim):
+ """
+ Neural net predictive model. The dragon has three heads.
+
+ Args:
+ input_dim (int): number of rows in input
+ Returns:
+ model (keras.models.Model): DragonNet model
+ """
+ inputs = Input(shape=(input_dim,), name="input")
+
+ # representation
+ x = Dense(
+ units=self.neurons_per_layer,
+ activation="elu",
+ kernel_initializer="RandomNormal",
+ )(inputs)
+ x = Dense(
+ units=self.neurons_per_layer,
+ activation="elu",
+ kernel_initializer="RandomNormal",
+ )(x)
+ x = Dense(
+ units=self.neurons_per_layer,
+ activation="elu",
+ kernel_initializer="RandomNormal",
+ )(x)
+
+ t_predictions = Dense(units=1, activation="sigmoid")(x)
+
+ # HYPOTHESIS
+ y0_hidden = Dense(
+ units=int(self.neurons_per_layer / 2),
+ activation="elu",
+ kernel_regularizer=l2(self.reg_l2),
+ )(x)
+ y1_hidden = Dense(
+ units=int(self.neurons_per_layer / 2),
+ activation="elu",
+ kernel_regularizer=l2(self.reg_l2),
+ )(x)
+
+ # second layer
+ y0_hidden = Dense(
+ units=int(self.neurons_per_layer / 2),
+ activation="elu",
+ kernel_regularizer=l2(self.reg_l2),
+ )(y0_hidden)
+ y1_hidden = Dense(
+ units=int(self.neurons_per_layer / 2),
+ activation="elu",
+ kernel_regularizer=l2(self.reg_l2),
+ )(y1_hidden)
+
+ # third
+ y0_predictions = Dense(
+ units=1,
+ activation=None,
+ kernel_regularizer=l2(self.reg_l2),
+ name="y0_predictions",
+ )(y0_hidden)
+ y1_predictions = Dense(
+ units=1,
+ activation=None,
+ kernel_regularizer=l2(self.reg_l2),
+ name="y1_predictions",
+ )(y1_hidden)
+
+ dl = EpsilonLayer()
+ epsilons = dl(t_predictions, name="epsilon")
+ concat_pred = Concatenate(1)(
+ [y0_predictions, y1_predictions, t_predictions, epsilons]
+ )
+ model = Model(inputs=inputs, outputs=concat_pred)
+
+ return model
+
+ def fit(self, X, treatment, y, p=None):
+ """
+ Fits the DragonNet 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)
+
+ y = np.hstack((y.reshape(-1, 1), treatment.reshape(-1, 1)))
+
+ self.dragonnet = self.make_dragonnet(X.shape[1])
+
+ metrics = [
+ regression_loss,
+ binary_classification_loss,
+ treatment_accuracy,
+ track_epsilon,
+ ]
+
+ if self.targeted_reg:
+ loss = make_tarreg_loss(ratio=self.ratio, dragonnet_loss=self.loss_func)
+ else:
+ loss = self.loss_func
+
+ if self.use_adam:
+ self.dragonnet.compile(
+ optimizer=Adam(learning_rate=self.adam_learning_rate),
+ loss=loss,
+ metrics=metrics,
+ )
+
+ adam_callbacks = [
+ TerminateOnNaN(),
+ EarlyStopping(monitor="val_loss", patience=2, min_delta=0.0),
+ ReduceLROnPlateau(
+ monitor="loss",
+ factor=0.5,
+ patience=5,
+ verbose=self.verbose,
+ mode="auto",
+ min_delta=1e-8,
+ cooldown=0,
+ min_lr=0,
+ ),
+ ]
+
+ self.dragonnet.fit(
+ X,
+ y,
+ callbacks=adam_callbacks,
+ validation_split=self.val_split,
+ epochs=self.adam_epochs,
+ batch_size=self.batch_size,
+ verbose=self.verbose,
+ )
+
+ sgd_callbacks = [
+ TerminateOnNaN(),
+ EarlyStopping(monitor="val_loss", patience=40, min_delta=0.0),
+ ReduceLROnPlateau(
+ monitor="loss",
+ factor=0.5,
+ patience=5,
+ verbose=self.verbose,
+ mode="auto",
+ min_delta=0.0,
+ cooldown=0,
+ min_lr=0,
+ ),
+ ]
+
+ self.dragonnet.compile(
+ optimizer=SGD(
+ learning_rate=self.learning_rate, momentum=self.momentum, nesterov=True
+ ),
+ loss=loss,
+ metrics=metrics,
+ )
+ self.dragonnet.fit(
+ X,
+ y,
+ callbacks=sgd_callbacks,
+ validation_split=self.val_split,
+ epochs=self.epochs,
+ batch_size=self.batch_size,
+ verbose=self.verbose,
+ )
+
+ def predict(self, X, treatment=None, y=None, p=None):
+ """
+ Calls predict on fitted DragonNet.
+
+ Args:
+ X (np.matrix or np.array or pd.Dataframe): a feature matrix
+ Returns:
+ (np.array): a 2D array with shape (X.shape[0], 4),
+ where each row takes the form of (outcome do(t=0), outcome do(t=1), propensity, epsilon)
+ """
+ return self.dragonnet.predict(X)
+
+ def predict_propensity(self, X):
+ """
+ Predicts the individual propensity scores.
+
+ Args:
+ X (np.matrix or np.array or pd.Dataframe): a feature matrix
+ Returns:
+ (np.array): propensity score vector
+ """
+ preds = self.predict(X)
+ return preds[:, 2]
+
+ def predict_tau(self, X):
+ """
+ Predicts the individual treatment effect (tau / "ITE").
+
+ Args:
+ X (np.matrix or np.array or pd.Dataframe): a feature matrix
+ Returns:
+ (np.array): treatment effect vector
+ """
+ preds = self.predict(X)
+ return (preds[:, 1] - preds[:, 0]).reshape(-1, 1)
+
+ def fit_predict(self, X, treatment, y, p=None, return_components=False):
+ """
+ Fits the DragonNet model and then predicts.
+
+ 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_components (bool, optional): whether to return
+ Returns:
+ (np.array): predictions based on return_components flag
+ if return_components=False (default), each row is treatment effect
+ if return_components=True, each row is (outcome do(t=0), outcome do(t=1), propensity, epsilon)
+ """
+ self.fit(X, treatment, y)
+ return self.predict_tau(X)
+
+ def save(self, h5_filepath):
+ """
+ Save the dragonnet model as a H5 file.
+
+ Args:
+ h5_filepath (H5 file path): H5 file path
+ """
+ self.dragonnet.save(h5_filepath)
+
+ def load(self, h5_filepath, ratio=1.0, dragonnet_loss=dragonnet_loss_binarycross):
+ """
+ Load the dragonnet model from a H5 file.
+
+ Args:
+ h5_filepath (H5 file path): H5 file path
+ ratio (float): weight assigned to the targeted regularization loss component
+ dragonnet_loss (function): a loss function
+ """
+ self.dragonnet = load_model(
+ h5_filepath,
+ custom_objects={
+ "EpsilonLayer": EpsilonLayer,
+ "dragonnet_loss_binarycross": dragonnet_loss_binarycross,
+ "tarreg_ATE_unbounded_domain_loss": make_tarreg_loss(
+ ratio=ratio, dragonnet_loss=dragonnet_loss
+ ),
+ "regression_loss": regression_loss,
+ "binary_classification_loss": binary_classification_loss,
+ "treatment_accuracy": treatment_accuracy,
+ "track_epsilon": track_epsilon,
+ },
+ )
diff --git a/causalml/source/causalml/inference/tf/utils.py b/causalml/source/causalml/inference/tf/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..dd01de2f4ac975df2f01ce5786da9c875562871d
--- /dev/null
+++ b/causalml/source/causalml/inference/tf/utils.py
@@ -0,0 +1,172 @@
+import tensorflow as tf
+from tensorflow.keras import backend as K
+from tensorflow.keras.layers import Layer
+from tensorflow.keras.metrics import binary_accuracy
+
+
+def binary_classification_loss(concat_true, concat_pred):
+ """
+ Implements a classification (binary cross-entropy) loss function for DragonNet architecture.
+
+ Args:
+ - concat_true (tf.tensor): tensor of true samples, with shape (n_samples, 2)
+ Each row in concat_true is comprised of (y, treatment)
+ - concat_pred (tf.tensor): tensor of predictions, with shape (n_samples, 4)
+ Each row in concat_pred is comprised of (y0, y1, propensity, epsilon)
+ Returns:
+ - (float): binary cross-entropy loss
+ """
+ t_true = concat_true[:, 1]
+ t_pred = concat_pred[:, 2]
+ t_pred = (t_pred + 0.001) / 1.002
+ losst = tf.reduce_sum(K.binary_crossentropy(t_true, t_pred))
+
+ return losst
+
+
+def regression_loss(concat_true, concat_pred):
+ """
+ Implements a regression (squared error) loss function for DragonNet architecture.
+
+ Args:
+ - concat_true (tf.tensor): tensor of true samples, with shape (n_samples, 2)
+ Each row in concat_true is comprised of (y, treatment)
+ - concat_pred (tf.tensor): tensor of predictions, with shape (n_samples, 4)
+ Each row in concat_pred is comprised of (y0, y1, propensity, epsilon)
+ Returns:
+ - (float): aggregated regression loss
+ """
+ y_true = concat_true[:, 0]
+ t_true = concat_true[:, 1]
+
+ y0_pred = concat_pred[:, 0]
+ y1_pred = concat_pred[:, 1]
+
+ loss0 = tf.reduce_sum((1.0 - t_true) * tf.square(y_true - y0_pred))
+ loss1 = tf.reduce_sum(t_true * tf.square(y_true - y1_pred))
+
+ return loss0 + loss1
+
+
+def dragonnet_loss_binarycross(concat_true, concat_pred):
+ """
+ Implements regression + classification loss in one wrapper function.
+
+ Args:
+ - concat_true (tf.tensor): tensor of true samples, with shape (n_samples, 2)
+ Each row in concat_true is comprised of (y, treatment)
+ - concat_pred (tf.tensor): tensor of predictions, with shape (n_samples, 4)
+ Each row in concat_pred is comprised of (y0, y1, propensity, epsilon)
+ Returns:
+ - (float): aggregated regression + classification loss
+ """
+ return regression_loss(concat_true, concat_pred) + binary_classification_loss(
+ concat_true, concat_pred
+ )
+
+
+def treatment_accuracy(concat_true, concat_pred):
+ """
+ Returns keras' binary_accuracy between treatment and prediction of propensity.
+
+ Args:
+ - concat_true (tf.tensor): tensor of true samples, with shape (n_samples, 2)
+ Each row in concat_true is comprised of (y, treatment)
+ - concat_pred (tf.tensor): tensor of predictions, with shape (n_samples, 4)
+ Each row in concat_pred is comprised of (y0, y1, propensity, epsilon)
+ Returns:
+ - (float): binary accuracy
+ """
+ t_true = concat_true[:, 1]
+ t_pred = concat_pred[:, 2]
+ return binary_accuracy(t_true, t_pred)
+
+
+def track_epsilon(concat_true, concat_pred):
+ """
+ Tracks the mean absolute value of epsilon.
+
+ Args:
+ - concat_true (tf.tensor): tensor of true samples, with shape (n_samples, 2)
+ Each row in concat_true is comprised of (y, treatment)
+ - concat_pred (tf.tensor): tensor of predictions, with shape (n_samples, 4)
+ Each row in concat_pred is comprised of (y0, y1, propensity, epsilon)
+ Returns:
+ - (float): mean absolute value of epsilon
+ """
+ epsilons = concat_pred[:, 3]
+ return tf.abs(tf.reduce_mean(epsilons))
+
+
+def make_tarreg_loss(ratio=1.0, dragonnet_loss=dragonnet_loss_binarycross):
+ """
+ Given a specified loss function, returns the same loss function with targeted regularization.
+
+ Args:
+ ratio (float): weight assigned to the targeted regularization loss component
+ dragonnet_loss (function): a loss function
+ Returns:
+ (function): loss function with targeted regularization, weighted by specified ratio
+ """
+
+ def tarreg_ATE_unbounded_domain_loss(concat_true, concat_pred):
+ """
+ Returns the loss function (specified in outer function) with targeted regularization.
+ """
+ vanilla_loss = dragonnet_loss(concat_true, concat_pred)
+
+ y_true = concat_true[:, 0]
+ t_true = concat_true[:, 1]
+
+ y0_pred = concat_pred[:, 0]
+ y1_pred = concat_pred[:, 1]
+ t_pred = concat_pred[:, 2]
+
+ epsilons = concat_pred[:, 3]
+ t_pred = (t_pred + 0.01) / 1.02
+ # t_pred = tf.clip_by_value(t_pred,0.01, 0.99,name='t_pred')
+
+ y_pred = t_true * y1_pred + (1 - t_true) * y0_pred
+
+ h = t_true / t_pred - (1 - t_true) / (1 - t_pred)
+
+ y_pert = y_pred + epsilons * h
+ targeted_regularization = tf.reduce_sum(tf.square(y_true - y_pert))
+
+ # final
+ loss = vanilla_loss + ratio * targeted_regularization
+ return loss
+
+ return tarreg_ATE_unbounded_domain_loss
+
+
+class EpsilonLayer(Layer):
+ """
+ Custom keras layer to allow epsilon to be learned during training process.
+ """
+
+ def __init__(self, **kwargs):
+ """
+ Inherits keras' Layer object.
+ """
+ super(EpsilonLayer, self).__init__(**kwargs)
+
+ def build(self, input_shape):
+ """
+ Creates a trainable weight variable for this layer.
+ """
+ self.epsilon = self.add_weight(
+ name="epsilon", shape=[1, 1], initializer="RandomNormal", trainable=True
+ )
+ super(EpsilonLayer, self).build(input_shape)
+
+ def call(self, inputs, **kwargs):
+ return self.epsilon * tf.ones_like(inputs)[:, 0:1]
+
+ def get_config(self):
+ config = super().get_config()
+ return config
+
+ @classmethod
+ def from_config(cls, config):
+ return cls(**config)
diff --git a/causalml/source/causalml/inference/torch/__init__.py b/causalml/source/causalml/inference/torch/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..b12f42b8b7a85adefef57413af429028b6ae0d1b
--- /dev/null
+++ b/causalml/source/causalml/inference/torch/__init__.py
@@ -0,0 +1 @@
+from .cevae import CEVAE
diff --git a/causalml/source/causalml/inference/torch/cevae.py b/causalml/source/causalml/inference/torch/cevae.py
new file mode 100644
index 0000000000000000000000000000000000000000..3e7b0f79f26a209058d80436fdeed8741784357a
--- /dev/null
+++ b/causalml/source/causalml/inference/torch/cevae.py
@@ -0,0 +1,142 @@
+"""
+This module calls the CEVAE[1] function implemented by pyro team. CEVAE demonstrates a number of innovations including:
+
+- A generative model for causal effect inference with hidden confounders;
+- A model and guide with twin neural nets to allow imbalanced treatment; and
+- A custom training loss that includes both ELBO terms and extra terms needed to train the guide to be able to answer
+counterfactual queries.
+
+Generative model for a causal model with latent confounder z and binary treatment w:
+ z ~ p(z) # latent confounder
+ x ~ p(x|z) # partial noisy observation of z
+ w ~ p(w|z) # treatment, whose application is biased by z
+ y ~ p(y|t,z) # outcome
+Each of these distributions is defined by a neural network. The y distribution is defined by a disjoint pair of neural
+networks defining p(y|t=0,z) and p(y|t=1,z); this allows highly imbalanced treatment.
+
+**References**
+
+[1] C. Louizos, U. Shalit, J. Mooij, D. Sontag, R. Zemel, M. Welling (2017).
+ | Causal Effect Inference with Deep Latent-Variable Models.
+ | http://papers.nips.cc/paper/7223-causal-effect-inference-with-deep-latent-variable-models.pdf
+ | https://github.com/AMLab-Amsterdam/CEVAE
+"""
+
+import logging
+import torch
+from pyro.contrib.cevae import CEVAE as CEVAEModel
+
+from causalml.inference.meta.utils import convert_pd_to_np
+
+pyro_logger = logging.getLogger("pyro")
+pyro_logger.setLevel(logging.DEBUG)
+if pyro_logger.handlers:
+ pyro_logger.handlers[0].setLevel(logging.DEBUG)
+
+
+class CEVAE:
+ def __init__(
+ self,
+ outcome_dist="studentt",
+ latent_dim=20,
+ hidden_dim=200,
+ num_epochs=50,
+ num_layers=3,
+ batch_size=100,
+ learning_rate=1e-3,
+ learning_rate_decay=0.1,
+ num_samples=1000,
+ weight_decay=1e-4,
+ ):
+ """
+ Initializes CEVAE.
+
+ Args:
+ outcome_dist (str): Outcome distribution as one of: "bernoulli" , "exponential", "laplace", "normal",
+ and "studentt"
+ latent_dim (int) : Dimension of the latent variable
+ hidden_dim (int) : Dimension of hidden layers of fully connected networks
+ num_epochs (int): Number of training epochs
+ num_layers (int): Number of hidden layers in fully connected networks
+ batch_size (int): Batch size
+ learning_rate (int): Learning rate
+ learning_rate_decay (float/int): Learning rate decay over all epochs; the per-step decay rate will
+ depend on batch size and number of epochs such that the initial
+ learning rate will be learning_rate and the
+ final learning rate will be learning_rate * learning_rate_decay
+ num_samples (int) : Number of samples to calculate ITE
+ weight_decay (float) : Weight decay
+ """
+ self.outcome_dist = outcome_dist
+ self.latent_dim = latent_dim
+ self.hidden_dim = hidden_dim
+ self.num_epochs = num_epochs
+ self.num_layers = num_layers
+ self.batch_size = batch_size
+ self.learning_rate = learning_rate
+ self.learning_rate_decay = learning_rate_decay
+ self.num_samples = num_samples
+ self.weight_decay = weight_decay
+
+ def fit(self, X, treatment, y, p=None):
+ """
+ Fits CEVAE.
+
+ 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)
+
+ self.cevae = CEVAEModel(
+ outcome_dist=self.outcome_dist,
+ feature_dim=X.shape[-1],
+ latent_dim=self.latent_dim,
+ hidden_dim=self.hidden_dim,
+ num_layers=self.num_layers,
+ )
+
+ self.cevae.fit(
+ x=torch.tensor(X, dtype=torch.float),
+ t=torch.tensor(treatment, dtype=torch.float),
+ y=torch.tensor(y, dtype=torch.float),
+ num_epochs=self.num_epochs,
+ batch_size=self.batch_size,
+ learning_rate=self.learning_rate,
+ learning_rate_decay=self.learning_rate_decay,
+ weight_decay=self.weight_decay,
+ )
+
+ def predict(self, X, treatment=None, y=None, p=None):
+ """
+ Calls predict on fitted DragonNet.
+
+ Args:
+ X (np.matrix or np.array or pd.Dataframe): a feature matrix
+ Returns:
+ (np.ndarray): Predictions of treatment effects.
+ """
+ return (
+ self.cevae.ite(
+ torch.tensor(X, dtype=torch.float),
+ num_samples=self.num_samples,
+ batch_size=self.batch_size,
+ )
+ .cpu()
+ .numpy()
+ )
+
+ def fit_predict(self, X, treatment, y, p=None):
+ """
+ Fits the CEVAE model and then predicts.
+
+ 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
+ Returns:
+ (np.ndarray): Predictions of treatment effects.
+ """
+ self.fit(X, treatment, y)
+ return self.predict(X)
diff --git a/causalml/source/causalml/inference/tree/__init__.py b/causalml/source/causalml/inference/tree/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..ee318b4a53c8033220cba092bb7c026218cdc3bb
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/__init__.py
@@ -0,0 +1,12 @@
+from .causal.causaltree import CausalTreeRegressor
+from .causal.causalforest import CausalRandomForestRegressor
+from .plot import uplift_tree_string, uplift_tree_plot, plot_dist_tree_leaves_values
+from .uplift import DecisionTree, UpliftTreeClassifier, UpliftRandomForestClassifier
+from .utils import (
+ cat_group,
+ cat_transform,
+ cv_fold_index,
+ cat_continuous,
+ kpi_transform,
+ get_tree_leaves_mask,
+)
diff --git a/causalml/source/causalml/inference/tree/_tree/__init__.py b/causalml/source/causalml/inference/tree/_tree/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..da91f0eaffbcfeedb9f44d5f2c6c25cdea6740eb
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/_tree/__init__.py
@@ -0,0 +1,14 @@
+"""
+This part of tree structures definition was initially borrowed from
+https://github.com/scikit-learn/scikit-learn/tree/1.5.2/sklearn/tree
+"""
+
+"""Decision tree based models for classification and regression."""
+
+from ._classes import (
+ BaseDecisionTree,
+)
+
+__all__ = [
+ "BaseDecisionTree",
+]
diff --git a/causalml/source/causalml/inference/tree/_tree/_classes.py b/causalml/source/causalml/inference/tree/_tree/_classes.py
new file mode 100644
index 0000000000000000000000000000000000000000..a9c8c8765bb4cd2f8e01ee8f828f56374f14ccb3
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/_tree/_classes.py
@@ -0,0 +1,685 @@
+"""
+This module gathers tree-based methods, including decision, regression and
+randomized trees. Single and multi-output problems are both handled.
+"""
+
+# Authors: Gilles Louppe
+# Peter Prettenhofer
+# Brian Holt
+# Noel Dawe
+# Satrajit Gosh
+# Joly Arnaud
+# Fares Hedayati
+# Nelson Liu
+#
+# License: BSD 3 clause
+
+import copy
+import numbers
+from abc import ABCMeta, abstractmethod
+from math import ceil
+from numbers import Integral, Real
+
+import numpy as np
+from scipy.sparse import issparse
+
+from sklearn.base import (
+ BaseEstimator,
+ ClassifierMixin,
+ MultiOutputMixin,
+ RegressorMixin,
+ _fit_context,
+ clone,
+ is_classifier,
+)
+from sklearn.utils import Bunch, check_random_state, compute_sample_weight
+from sklearn.utils._param_validation import Hidden, Interval, RealNotInt, StrOptions
+from sklearn.utils.multiclass import check_classification_targets
+from sklearn.utils.validation import (
+ _assert_all_finite_element_wise,
+ _check_sample_weight,
+ assert_all_finite,
+ check_is_fitted,
+ validate_data,
+)
+from . import _criterion, _splitter, _tree
+from ._criterion import Criterion
+from ._splitter import Splitter
+from ._tree import (
+ BestFirstTreeBuilder,
+ DepthFirstTreeBuilder,
+ Tree,
+ _build_pruned_tree_ccp,
+ ccp_pruning_path,
+)
+from ._utils import _any_isnan_axis0
+
+# =============================================================================
+# Types and constants
+# =============================================================================
+
+DTYPE = _tree.DTYPE
+DOUBLE = _tree.DOUBLE
+INT = _tree.INT
+
+CRITERIA_CLF = {
+ "gini": _criterion.Gini,
+ "log_loss": _criterion.Entropy,
+ "entropy": _criterion.Entropy,
+}
+CRITERIA_REG = {
+ "squared_error": _criterion.MSE,
+ "friedman_mse": _criterion.FriedmanMSE,
+ "absolute_error": _criterion.MAE,
+ "poisson": _criterion.Poisson,
+}
+
+DENSE_SPLITTERS = {"best": _splitter.BestSplitter, "random": _splitter.RandomSplitter}
+
+SPARSE_SPLITTERS = {
+ "best": _splitter.BestSparseSplitter,
+ "random": _splitter.RandomSparseSplitter,
+}
+
+# =============================================================================
+# Base decision tree
+# =============================================================================
+
+
+class BaseDecisionTree(MultiOutputMixin, BaseEstimator, metaclass=ABCMeta):
+ """Base class for decision trees.
+
+ Warning: This class should not be used directly.
+ Use derived classes instead.
+ """
+
+ _parameter_constraints: dict = {
+ "splitter": [StrOptions({"best", "random"})],
+ "max_depth": [Interval(Integral, 1, None, closed="left"), None],
+ "min_samples_split": [
+ Interval(Integral, 2, None, closed="left"),
+ Interval(RealNotInt, 0.0, 1.0, closed="right"),
+ ],
+ "min_samples_leaf": [
+ Interval(Integral, 1, None, closed="left"),
+ Interval(RealNotInt, 0.0, 1.0, closed="neither"),
+ ],
+ "min_weight_fraction_leaf": [Interval(Real, 0.0, 0.5, closed="both")],
+ "max_features": [
+ Interval(Integral, 1, None, closed="left"),
+ Interval(RealNotInt, 0.0, 1.0, closed="right"),
+ StrOptions({"sqrt", "log2"}),
+ None,
+ ],
+ "random_state": ["random_state"],
+ "max_leaf_nodes": [Interval(Integral, 2, None, closed="left"), None],
+ "min_impurity_decrease": [Interval(Real, 0.0, None, closed="left")],
+ "ccp_alpha": [Interval(Real, 0.0, None, closed="left")],
+ "monotonic_cst": ["array-like", None],
+ }
+
+ @abstractmethod
+ def __init__(
+ self,
+ *,
+ criterion,
+ splitter,
+ max_depth,
+ min_samples_split,
+ min_samples_leaf,
+ min_weight_fraction_leaf,
+ max_features,
+ max_leaf_nodes,
+ random_state,
+ min_impurity_decrease,
+ class_weight=None,
+ ccp_alpha=0.0,
+ monotonic_cst=None,
+ ):
+ self.criterion = criterion
+ self.splitter = splitter
+ self.max_depth = max_depth
+ self.min_samples_split = min_samples_split
+ self.min_samples_leaf = min_samples_leaf
+ self.min_weight_fraction_leaf = min_weight_fraction_leaf
+ self.max_features = max_features
+ self.max_leaf_nodes = max_leaf_nodes
+ self.random_state = random_state
+ self.min_impurity_decrease = min_impurity_decrease
+ self.class_weight = class_weight
+ self.ccp_alpha = ccp_alpha
+ self.monotonic_cst = monotonic_cst
+
+ def get_depth(self):
+ """Return the depth of the decision tree.
+
+ The depth of a tree is the maximum distance between the root
+ and any leaf.
+
+ Returns
+ -------
+ self.tree_.max_depth : int
+ The maximum depth of the tree.
+ """
+ check_is_fitted(self)
+ return self.tree_.max_depth
+
+ def get_n_leaves(self):
+ """Return the number of leaves of the decision tree.
+
+ Returns
+ -------
+ self.tree_.n_leaves : int
+ Number of leaves.
+ """
+ check_is_fitted(self)
+ return self.tree_.n_leaves
+
+ def _support_missing_values(self, X):
+ return (
+ not issparse(X)
+ and self._get_tags()["allow_nan"]
+ and self.monotonic_cst is None
+ )
+
+ def _compute_missing_values_in_feature_mask(self, X, estimator_name=None):
+ """Return boolean mask denoting if there are missing values for each feature.
+
+ This method also ensures that X is finite.
+
+ Parameter
+ ---------
+ X : array-like of shape (n_samples, n_features), dtype=DOUBLE
+ Input data.
+
+ estimator_name : str or None, default=None
+ Name to use when raising an error. Defaults to the class name.
+
+ Returns
+ -------
+ missing_values_in_feature_mask : ndarray of shape (n_features,), or None
+ Missing value mask. If missing values are not supported or there
+ are no missing values, return None.
+ """
+ estimator_name = estimator_name or self.__class__.__name__
+ common_kwargs = dict(estimator_name=estimator_name, input_name="X")
+
+ if not self._support_missing_values(X):
+ assert_all_finite(X, **common_kwargs)
+ return None
+
+ with np.errstate(over="ignore"):
+ overall_sum = np.sum(X)
+
+ if not np.isfinite(overall_sum):
+ # Raise a ValueError in case of the presence of an infinite element.
+ _assert_all_finite_element_wise(X, xp=np, allow_nan=True, **common_kwargs)
+
+ # If the sum is not nan, then there are no missing values
+ if not np.isnan(overall_sum):
+ return None
+
+ missing_values_in_feature_mask = _any_isnan_axis0(X)
+ return missing_values_in_feature_mask
+
+ def _fit(
+ self,
+ X,
+ y,
+ sample_weight=None,
+ check_input=True,
+ missing_values_in_feature_mask=None,
+ ):
+ random_state = check_random_state(self.random_state)
+
+ if check_input:
+ # Need to validate separately here.
+ # We can't pass multi_output=True because that would allow y to be
+ # csr.
+
+ # _compute_missing_values_in_feature_mask will check for finite values and
+ # compute the missing mask if the tree supports missing values
+ check_X_params = dict(
+ dtype=DTYPE, accept_sparse="csc", ensure_all_finite=False
+ )
+ check_y_params = dict(ensure_2d=False, dtype=None)
+ X, y = validate_data(
+ self, X, y, validate_separately=(check_X_params, check_y_params)
+ )
+
+ missing_values_in_feature_mask = (
+ self._compute_missing_values_in_feature_mask(X)
+ )
+ if issparse(X):
+ X.sort_indices()
+
+ if X.indices.dtype != np.intc or X.indptr.dtype != np.intc:
+ raise ValueError(
+ "No support for np.int64 index based sparse matrices"
+ )
+
+ if self.criterion == "poisson":
+ if np.any(y < 0):
+ raise ValueError(
+ "Some value(s) of y are negative which is"
+ " not allowed for Poisson regression."
+ )
+ if np.sum(y) <= 0:
+ raise ValueError(
+ "Sum of y is not positive which is "
+ "necessary for Poisson regression."
+ )
+
+ # Determine output settings
+ n_samples, self.n_features_in_ = X.shape
+ is_classification = is_classifier(self)
+
+ y = np.atleast_1d(y)
+ expanded_class_weight = None
+
+ if y.ndim == 1:
+ # reshape is necessary to preserve the data contiguity against vs
+ # [:, np.newaxis] that does not.
+ y = np.reshape(y, (-1, 1))
+
+ self.n_outputs_ = y.shape[1]
+
+ if is_classification:
+ check_classification_targets(y)
+ y = np.copy(y)
+
+ self.classes_ = []
+ self.n_classes_ = []
+
+ if self.class_weight is not None:
+ y_original = np.copy(y)
+
+ y_encoded = np.zeros(y.shape, dtype=int)
+ for k in range(self.n_outputs_):
+ classes_k, y_encoded[:, k] = np.unique(y[:, k], return_inverse=True)
+ self.classes_.append(classes_k)
+ self.n_classes_.append(classes_k.shape[0])
+ y = y_encoded
+
+ if self.class_weight is not None:
+ expanded_class_weight = compute_sample_weight(
+ self.class_weight, y_original
+ )
+
+ self.n_classes_ = np.array(self.n_classes_, dtype=np.intp)
+
+ if getattr(y, "dtype", None) != DOUBLE or not y.flags.contiguous:
+ y = np.ascontiguousarray(y, dtype=DOUBLE)
+
+ max_depth = np.iinfo(np.int32).max if self.max_depth is None else self.max_depth
+
+ if isinstance(self.min_samples_leaf, numbers.Integral):
+ min_samples_leaf = self.min_samples_leaf
+ else: # float
+ min_samples_leaf = int(ceil(self.min_samples_leaf * n_samples))
+
+ if isinstance(self.min_samples_split, numbers.Integral):
+ min_samples_split = self.min_samples_split
+ else: # float
+ min_samples_split = int(ceil(self.min_samples_split * n_samples))
+ min_samples_split = max(2, min_samples_split)
+
+ min_samples_split = max(min_samples_split, 2 * min_samples_leaf)
+
+ if isinstance(self.max_features, str):
+ if self.max_features == "sqrt":
+ max_features = max(1, int(np.sqrt(self.n_features_in_)))
+ elif self.max_features == "log2":
+ max_features = max(1, int(np.log2(self.n_features_in_)))
+ elif self.max_features is None:
+ max_features = self.n_features_in_
+ elif isinstance(self.max_features, numbers.Integral):
+ max_features = self.max_features
+ else: # float
+ if self.max_features > 0.0:
+ max_features = max(1, int(self.max_features * self.n_features_in_))
+ else:
+ max_features = 0
+
+ self.max_features_ = max_features
+
+ max_leaf_nodes = -1 if self.max_leaf_nodes is None else self.max_leaf_nodes
+
+ if len(y) != n_samples:
+ raise ValueError(
+ "Number of labels=%d does not match number of samples=%d"
+ % (len(y), n_samples)
+ )
+
+ if sample_weight is not None:
+ sample_weight = _check_sample_weight(sample_weight, X, DOUBLE)
+
+ if expanded_class_weight is not None:
+ if sample_weight is not None:
+ sample_weight = sample_weight * expanded_class_weight
+ else:
+ sample_weight = expanded_class_weight
+
+ # Set min_weight_leaf from min_weight_fraction_leaf
+ if sample_weight is None:
+ min_weight_leaf = self.min_weight_fraction_leaf * n_samples
+ else:
+ min_weight_leaf = self.min_weight_fraction_leaf * np.sum(sample_weight)
+
+ # Build tree
+ criterion = self.criterion
+ if not isinstance(criterion, Criterion):
+ if is_classification:
+ criterion = CRITERIA_CLF[self.criterion](
+ self.n_outputs_, self.n_classes_
+ )
+ else:
+ criterion = CRITERIA_REG[self.criterion](self.n_outputs_, n_samples)
+ else:
+ # Make a deepcopy in case the criterion has mutable attributes that
+ # might be shared and modified concurrently during parallel fitting
+ criterion = copy.deepcopy(criterion)
+
+ SPLITTERS = SPARSE_SPLITTERS if issparse(X) else DENSE_SPLITTERS
+
+ splitter = self.splitter
+ if self.monotonic_cst is None:
+ monotonic_cst = None
+ else:
+ if self.n_outputs_ > 1:
+ raise ValueError(
+ "Monotonicity constraints are not supported with multiple outputs."
+ )
+ # Check to correct monotonicity constraint' specification,
+ # by applying element-wise logical conjunction
+ # Note: we do not cast `np.asarray(self.monotonic_cst, dtype=np.int8)`
+ # straight away here so as to generate error messages for invalid
+ # values using the original values prior to any dtype related conversion.
+ monotonic_cst = np.asarray(self.monotonic_cst)
+ if monotonic_cst.shape[0] != X.shape[1]:
+ raise ValueError(
+ "monotonic_cst has shape {} but the input data "
+ "X has {} features.".format(monotonic_cst.shape[0], X.shape[1])
+ )
+ valid_constraints = np.isin(monotonic_cst, (-1, 0, 1))
+ if not np.all(valid_constraints):
+ unique_constaints_value = np.unique(monotonic_cst)
+ raise ValueError(
+ "monotonic_cst must be None or an array-like of -1, 0 or 1, but"
+ f" got {unique_constaints_value}"
+ )
+ monotonic_cst = np.asarray(monotonic_cst, dtype=np.int8)
+ if is_classifier(self):
+ if self.n_classes_[0] > 2:
+ raise ValueError(
+ "Monotonicity constraints are not supported with multiclass "
+ "classification"
+ )
+ # Binary classification trees are built by constraining probabilities
+ # of the *negative class* in order to make the implementation similar
+ # to regression trees.
+ # Since self.monotonic_cst encodes constraints on probabilities of the
+ # *positive class*, all signs must be flipped.
+ monotonic_cst *= -1
+
+ if not isinstance(self.splitter, Splitter):
+ splitter = SPLITTERS[self.splitter](
+ criterion,
+ self.max_features_,
+ min_samples_leaf,
+ min_weight_leaf,
+ random_state,
+ monotonic_cst,
+ )
+
+ if is_classifier(self):
+ self.tree_ = Tree(self.n_features_in_, self.n_classes_, self.n_outputs_)
+ else:
+ self.tree_ = Tree(
+ self.n_features_in_,
+ # TODO: tree shouldn't need this in this case
+ np.array([1] * self.n_outputs_, dtype=np.intp),
+ self.n_outputs_,
+ )
+
+ # Use BestFirst if max_leaf_nodes given; use DepthFirst otherwise
+ if max_leaf_nodes < 0:
+ builder = DepthFirstTreeBuilder(
+ splitter,
+ min_samples_split,
+ min_samples_leaf,
+ min_weight_leaf,
+ max_depth,
+ self.min_impurity_decrease,
+ )
+ else:
+ builder = BestFirstTreeBuilder(
+ splitter,
+ min_samples_split,
+ min_samples_leaf,
+ min_weight_leaf,
+ max_depth,
+ max_leaf_nodes,
+ self.min_impurity_decrease,
+ )
+
+ builder.build(self.tree_, X, y, sample_weight, missing_values_in_feature_mask)
+
+ if self.n_outputs_ == 1 and is_classifier(self):
+ self.n_classes_ = self.n_classes_[0]
+ self.classes_ = self.classes_[0]
+
+ self._prune_tree()
+
+ return self
+
+ def _validate_X_predict(self, X, check_input):
+ """Validate the training data on predict (probabilities)."""
+ if check_input:
+ if self._support_missing_values(X):
+ ensure_all_finite = "allow-nan"
+ else:
+ ensure_all_finite = True
+ X = validate_data(
+ self,
+ X,
+ dtype=DTYPE,
+ accept_sparse="csr",
+ reset=False,
+ ensure_all_finite=ensure_all_finite,
+ )
+ if issparse(X) and (
+ X.indices.dtype != np.intc or X.indptr.dtype != np.intc
+ ):
+ raise ValueError("No support for np.int64 index based sparse matrices")
+ else:
+ # The number of features is checked regardless of `check_input`
+ self._check_n_features(X, reset=False)
+ return X
+
+ def predict(self, X, check_input=True):
+ """Predict class or regression value for X.
+
+ For a classification model, the predicted class for each sample in X is
+ returned. For a regression model, the predicted value based on X is
+ returned.
+
+ Parameters
+ ----------
+ X : {array-like, sparse matrix} of shape (n_samples, n_features)
+ The input samples. Internally, it will be converted to
+ ``dtype=np.float32`` and if a sparse matrix is provided
+ to a sparse ``csr_matrix``.
+
+ check_input : bool, default=True
+ Allow to bypass several input checking.
+ Don't use this parameter unless you know what you're doing.
+
+ Returns
+ -------
+ y : array-like of shape (n_samples,) or (n_samples, n_outputs)
+ The predicted classes, or the predict values.
+ """
+ check_is_fitted(self)
+ X = self._validate_X_predict(X, check_input)
+ proba = self.tree_.predict(X)
+ n_samples = X.shape[0]
+
+ # Classification
+ if is_classifier(self):
+ if self.n_outputs_ == 1:
+ return self.classes_.take(np.argmax(proba, axis=1), axis=0)
+
+ else:
+ class_type = self.classes_[0].dtype
+ predictions = np.zeros((n_samples, self.n_outputs_), dtype=class_type)
+ for k in range(self.n_outputs_):
+ predictions[:, k] = self.classes_[k].take(
+ np.argmax(proba[:, k], axis=1), axis=0
+ )
+
+ return predictions
+
+ # Regression
+ else:
+ if self.n_outputs_ == 1:
+ return proba[:, 0]
+
+ else:
+ return proba[:, :, 0]
+
+ def apply(self, X, check_input=True):
+ """Return the index of the leaf that each sample is predicted as.
+
+ .. versionadded:: 0.17
+
+ Parameters
+ ----------
+ X : {array-like, sparse matrix} of shape (n_samples, n_features)
+ The input samples. Internally, it will be converted to
+ ``dtype=np.float32`` and if a sparse matrix is provided
+ to a sparse ``csr_matrix``.
+
+ check_input : bool, default=True
+ Allow to bypass several input checking.
+ Don't use this parameter unless you know what you're doing.
+
+ Returns
+ -------
+ X_leaves : array-like of shape (n_samples,)
+ For each datapoint x in X, return the index of the leaf x
+ ends up in. Leaves are numbered within
+ ``[0; self.tree_.node_count)``, possibly with gaps in the
+ numbering.
+ """
+ check_is_fitted(self)
+ X = self._validate_X_predict(X, check_input)
+ return self.tree_.apply(X)
+
+ def decision_path(self, X, check_input=True):
+ """Return the decision path in the tree.
+
+ .. versionadded:: 0.18
+
+ Parameters
+ ----------
+ X : {array-like, sparse matrix} of shape (n_samples, n_features)
+ The input samples. Internally, it will be converted to
+ ``dtype=np.float32`` and if a sparse matrix is provided
+ to a sparse ``csr_matrix``.
+
+ check_input : bool, default=True
+ Allow to bypass several input checking.
+ Don't use this parameter unless you know what you're doing.
+
+ Returns
+ -------
+ indicator : sparse matrix of shape (n_samples, n_nodes)
+ Return a node indicator CSR matrix where non zero elements
+ indicates that the samples goes through the nodes.
+ """
+ X = self._validate_X_predict(X, check_input)
+ return self.tree_.decision_path(X)
+
+ def _prune_tree(self):
+ """Prune tree using Minimal Cost-Complexity Pruning."""
+ check_is_fitted(self)
+
+ if self.ccp_alpha == 0.0:
+ return
+
+ # build pruned tree
+ if is_classifier(self):
+ n_classes = np.atleast_1d(self.n_classes_)
+ pruned_tree = Tree(self.n_features_in_, n_classes, self.n_outputs_)
+ else:
+ pruned_tree = Tree(
+ self.n_features_in_,
+ # TODO: the tree shouldn't need this param
+ np.array([1] * self.n_outputs_, dtype=np.intp),
+ self.n_outputs_,
+ )
+ _build_pruned_tree_ccp(pruned_tree, self.tree_, self.ccp_alpha)
+
+ self.tree_ = pruned_tree
+
+ def cost_complexity_pruning_path(self, X, y, sample_weight=None):
+ """Compute the pruning path during Minimal Cost-Complexity Pruning.
+
+ See :ref:`minimal_cost_complexity_pruning` for details on the pruning
+ process.
+
+ Parameters
+ ----------
+ X : {array-like, sparse matrix} of shape (n_samples, n_features)
+ The training input samples. Internally, it will be converted to
+ ``dtype=np.float32`` and if a sparse matrix is provided
+ to a sparse ``csc_matrix``.
+
+ y : array-like of shape (n_samples,) or (n_samples, n_outputs)
+ The target values (class labels) as integers or strings.
+
+ sample_weight : array-like of shape (n_samples,), default=None
+ Sample weights. If None, then samples are equally weighted. Splits
+ that would create child nodes with net zero or negative weight are
+ ignored while searching for a split in each node. Splits are also
+ ignored if they would result in any single class carrying a
+ negative weight in either child node.
+
+ Returns
+ -------
+ ccp_path : :class:`~sklearn.utils.Bunch`
+ Dictionary-like object, with the following attributes.
+
+ ccp_alphas : ndarray
+ Effective alphas of subtree during pruning.
+
+ impurities : ndarray
+ Sum of the impurities of the subtree leaves for the
+ corresponding alpha value in ``ccp_alphas``.
+ """
+ est = clone(self).set_params(ccp_alpha=0.0)
+ est.fit(X, y, sample_weight=sample_weight)
+ return Bunch(**ccp_pruning_path(est.tree_))
+
+ @property
+ def feature_importances_(self):
+ """Return the feature importances.
+
+ The importance of a feature is computed as the (normalized) total
+ reduction of the criterion brought by that feature.
+ It is also known as the Gini importance.
+
+ Warning: impurity-based feature importances can be misleading for
+ high cardinality features (many unique values). See
+ :func:`sklearn.inspection.permutation_importance` as an alternative.
+
+ Returns
+ -------
+ feature_importances_ : ndarray of shape (n_features,)
+ Normalized total reduction of criteria by feature
+ (Gini importance).
+ """
+ check_is_fitted(self)
+
+ return self.tree_.compute_feature_importances()
diff --git a/causalml/source/causalml/inference/tree/_tree/_criterion.pxd b/causalml/source/causalml/inference/tree/_tree/_criterion.pxd
new file mode 100644
index 0000000000000000000000000000000000000000..53785299f14ac9941a3c794fc4248103124c8344
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/_tree/_criterion.pxd
@@ -0,0 +1,121 @@
+# Authors: Gilles Louppe
+# Peter Prettenhofer
+# Brian Holt
+# Joel Nothman
+# Arnaud Joly
+# Jacob Schreiber
+#
+# License: BSD 3 clause
+
+# cython: cdivision=True
+# cython: boundscheck=False
+# cython: wraparound=False
+# cython: language_level=3
+# cython: linetrace=True
+
+# See _criterion.pyx for implementation details.
+from ._typedefs cimport float64_t, int8_t, int32_t, intp_t
+
+
+cdef class Criterion:
+ # The criterion computes the impurity of a node and the reduction of
+ # impurity of a split on that node. It also computes the output statistics
+ # such as the mean in regression and class probabilities in classification.
+
+ # Internal structures
+ cdef const float64_t[:, ::1] y # Values of y
+ cdef const float64_t[:] sample_weight # Sample weights
+
+ cdef const intp_t[:] sample_indices # Sample indices in X, y
+ cdef intp_t start # samples[start:pos] are the samples in the left node
+ cdef intp_t pos # samples[pos:end] are the samples in the right node
+ cdef intp_t end
+ cdef intp_t n_missing # Number of missing values for the feature being evaluated
+ cdef bint missing_go_to_left # Whether missing values go to the left node
+
+ cdef intp_t n_outputs # Number of outputs
+ cdef intp_t n_samples # Number of samples
+ cdef intp_t n_node_samples # Number of samples in the node (end-start)
+ cdef float64_t weighted_n_samples # Weighted number of samples (in total)
+ cdef float64_t weighted_n_node_samples # Weighted number of samples in the node
+ cdef float64_t weighted_n_left # Weighted number of samples in the left node
+ cdef float64_t weighted_n_right # Weighted number of samples in the right node
+ cdef float64_t weighted_n_missing # Weighted number of samples that are missing
+
+ # The criterion object is maintained such that left and right collected
+ # statistics correspond to samples[start:pos] and samples[pos:end].
+
+ # Methods
+ cdef int init(
+ self,
+ const float64_t[:, ::1] y,
+ const float64_t[:] sample_weight,
+ float64_t weighted_n_samples,
+ const intp_t[:] sample_indices,
+ intp_t start,
+ intp_t end
+ ) except -1 nogil
+ cdef void init_sum_missing(self)
+ cdef void init_missing(self, intp_t n_missing) noexcept nogil
+ cdef int reset(self) except -1 nogil
+ cdef int reverse_reset(self) except -1 nogil
+ cdef int update(self, intp_t new_pos) except -1 nogil
+ cdef float64_t node_impurity(self) noexcept nogil
+ cdef void children_impurity(
+ self,
+ float64_t* impurity_left,
+ float64_t* impurity_right
+ ) noexcept nogil
+ cdef void node_value(
+ self,
+ float64_t* dest
+ ) noexcept nogil
+ cdef void clip_node_value(
+ self,
+ float64_t* dest,
+ float64_t lower_bound,
+ float64_t upper_bound
+ ) noexcept nogil
+ cdef float64_t middle_value(self) noexcept nogil
+ cdef float64_t impurity_improvement(
+ self,
+ float64_t impurity_parent,
+ float64_t impurity_left,
+ float64_t impurity_right
+ ) noexcept nogil
+ cdef float64_t proxy_impurity_improvement(self) noexcept nogil
+ cdef bint check_monotonicity(
+ self,
+ int8_t monotonic_cst,
+ float64_t lower_bound,
+ float64_t upper_bound,
+ ) noexcept nogil
+ cdef inline bint _check_monotonicity(
+ self,
+ int8_t monotonic_cst,
+ float64_t lower_bound,
+ float64_t upper_bound,
+ float64_t sum_left,
+ float64_t sum_right,
+ ) noexcept nogil
+
+cdef class ClassificationCriterion(Criterion):
+ """Abstract criterion for classification."""
+
+ cdef intp_t[::1] n_classes
+ cdef intp_t max_n_classes
+
+ cdef float64_t[:, ::1] sum_total # The sum of the weighted count of each label.
+ cdef float64_t[:, ::1] sum_left # Same as above, but for the left side of the split
+ cdef float64_t[:, ::1] sum_right # Same as above, but for the right side of the split
+ cdef float64_t[:, ::1] sum_missing # Same as above, but for missing values in X
+
+cdef class RegressionCriterion(Criterion):
+ """Abstract regression criterion."""
+
+ cdef float64_t sq_sum_total
+
+ cdef float64_t[::1] sum_total # The sum of w*y.
+ cdef float64_t[::1] sum_left # Same as above, but for the left side of the split
+ cdef float64_t[::1] sum_right # Same as above, but for the right side of the split
+ cdef float64_t[::1] sum_missing # Same as above, but for missing values in X
diff --git a/causalml/source/causalml/inference/tree/_tree/_criterion.pyx b/causalml/source/causalml/inference/tree/_tree/_criterion.pyx
new file mode 100644
index 0000000000000000000000000000000000000000..0f0fd34921a2db510cad43c7f2c178c0b40604c6
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/_tree/_criterion.pyx
@@ -0,0 +1,1714 @@
+# Authors: Gilles Louppe
+# Peter Prettenhofer
+# Brian Holt
+# Noel Dawe
+# Satrajit Gosh
+# Lars Buitinck
+# Arnaud Joly
+# Joel Nothman
+# Fares Hedayati
+# Jacob Schreiber
+# Nelson Liu
+#
+# License: BSD 3 clause
+
+# cython: cdivision=True
+# cython: boundscheck=False
+# cython: wraparound=False
+# cython: language_level=3
+# cython: linetrace=True
+
+from libc.string cimport memcpy
+from libc.string cimport memset
+from libc.math cimport fabs, INFINITY
+
+import numpy as np
+cimport numpy as cnp
+cnp.import_array()
+
+from scipy.special.cython_special cimport xlogy
+
+from ._utils cimport log
+from ._utils cimport WeightedMedianCalculator
+
+# EPSILON is used in the Poisson criterion
+cdef float64_t EPSILON = 10 * np.finfo('double').eps
+
+cdef class Criterion:
+ """Interface for impurity criteria.
+
+ This object stores methods on how to calculate how good a split is using
+ different metrics.
+ """
+ def __getstate__(self):
+ return {}
+
+ def __setstate__(self, d):
+ pass
+
+ cdef int init(
+ self,
+ const float64_t[:, ::1] y,
+ const float64_t[:] sample_weight,
+ float64_t weighted_n_samples,
+ const intp_t[:] sample_indices,
+ intp_t start,
+ intp_t end,
+ ) except -1 nogil:
+ """Placeholder for a method which will initialize the criterion.
+
+ Returns -1 in case of failure to allocate memory (and raise MemoryError)
+ or 0 otherwise.
+
+ Parameters
+ ----------
+ y : ndarray, dtype=float64_t
+ y is a buffer that can store values for n_outputs target variables
+ stored as a Cython memoryview.
+ sample_weight : ndarray, dtype=float64_t
+ The weight of each sample stored as a Cython memoryview.
+ weighted_n_samples : float64_t
+ The total weight of the samples being considered
+ sample_indices : ndarray, dtype=intp_t
+ A mask on the samples. Indices of the samples in X and y we want to use,
+ where sample_indices[start:end] correspond to the samples in this node.
+ start : intp_t
+ The first sample to be used on this node
+ end : intp_t
+ The last sample used on this node
+
+ """
+ pass
+
+ cdef void init_missing(self, intp_t n_missing) noexcept nogil:
+ """Initialize sum_missing if there are missing values.
+
+ This method assumes that caller placed the missing samples in
+ self.sample_indices[-n_missing:]
+
+ Parameters
+ ----------
+ n_missing: intp_t
+ Number of missing values for specific feature.
+ """
+ pass
+
+ cdef int reset(self) except -1 nogil:
+ """Reset the criterion at pos=start.
+
+ This method must be implemented by the subclass.
+ """
+ pass
+
+ cdef int reverse_reset(self) except -1 nogil:
+ """Reset the criterion at pos=end.
+
+ This method must be implemented by the subclass.
+ """
+ pass
+
+ cdef int update(self, intp_t new_pos) except -1 nogil:
+ """Updated statistics by moving sample_indices[pos:new_pos] to the left child.
+
+ This updates the collected statistics by moving sample_indices[pos:new_pos]
+ from the right child to the left child. It must be implemented by
+ the subclass.
+
+ Parameters
+ ----------
+ new_pos : intp_t
+ New starting index position of the sample_indices in the right child
+ """
+ pass
+
+ cdef float64_t node_impurity(self) noexcept nogil:
+ """Placeholder for calculating the impurity of the node.
+
+ Placeholder for a method which will evaluate the impurity of
+ the current node, i.e. the impurity of sample_indices[start:end]. This is the
+ primary function of the criterion class. The smaller the impurity the
+ better.
+ """
+ pass
+
+ cdef void children_impurity(self, float64_t* impurity_left,
+ float64_t* impurity_right) noexcept nogil:
+ """Placeholder for calculating the impurity of children.
+
+ Placeholder for a method which evaluates the impurity in
+ children nodes, i.e. the impurity of sample_indices[start:pos] + the impurity
+ of sample_indices[pos:end].
+
+ Parameters
+ ----------
+ impurity_left : float64_t pointer
+ The memory address where the impurity of the left child should be
+ stored.
+ impurity_right : float64_t pointer
+ The memory address where the impurity of the right child should be
+ stored
+ """
+ pass
+
+ cdef void node_value(self, float64_t* dest) noexcept nogil:
+ """Placeholder for storing the node value.
+
+ Placeholder for a method which will compute the node value
+ of sample_indices[start:end] and save the value into dest.
+
+ Parameters
+ ----------
+ dest : float64_t pointer
+ The memory address where the node value should be stored.
+ """
+ pass
+
+ cdef void clip_node_value(self, float64_t* dest, float64_t lower_bound, float64_t upper_bound) noexcept nogil:
+ pass
+
+ cdef float64_t middle_value(self) noexcept nogil:
+ """Compute the middle value of a split for monotonicity constraints
+
+ This method is implemented in ClassificationCriterion and RegressionCriterion.
+ """
+ pass
+
+ cdef float64_t proxy_impurity_improvement(self) noexcept nogil:
+ """Compute a proxy of the impurity reduction.
+
+ This method is used to speed up the search for the best split.
+ It is a proxy quantity such that the split that maximizes this value
+ also maximizes the impurity improvement. It neglects all constant terms
+ of the impurity decrease for a given split.
+
+ The absolute impurity improvement is only computed by the
+ impurity_improvement method once the best split has been found.
+ """
+ cdef float64_t impurity_left
+ cdef float64_t impurity_right
+ self.children_impurity(&impurity_left, &impurity_right)
+
+ return (- self.weighted_n_right * impurity_right
+ - self.weighted_n_left * impurity_left)
+
+ cdef float64_t impurity_improvement(self, float64_t impurity_parent,
+ float64_t impurity_left,
+ float64_t impurity_right) noexcept nogil:
+ """Compute the improvement in impurity.
+
+ This method computes the improvement in impurity when a split occurs.
+ The weighted impurity improvement equation is the following:
+
+ N_t / N * (impurity - N_t_R / N_t * right_impurity
+ - N_t_L / N_t * left_impurity)
+
+ where N is the total number of samples, N_t is the number of samples
+ at the current node, N_t_L is the number of samples in the left child,
+ and N_t_R is the number of samples in the right child,
+
+ Parameters
+ ----------
+ impurity_parent : float64_t
+ The initial impurity of the parent node before the split
+
+ impurity_left : float64_t
+ The impurity of the left child
+
+ impurity_right : float64_t
+ The impurity of the right child
+
+ Return
+ ------
+ float64_t : improvement in impurity after the split occurs
+ """
+ return ((self.weighted_n_node_samples / self.weighted_n_samples) *
+ (impurity_parent - (self.weighted_n_right /
+ self.weighted_n_node_samples * impurity_right)
+ - (self.weighted_n_left /
+ self.weighted_n_node_samples * impurity_left)))
+
+ cdef bint check_monotonicity(
+ self,
+ cnp.int8_t monotonic_cst,
+ float64_t lower_bound,
+ float64_t upper_bound,
+ ) noexcept nogil:
+ pass
+
+ cdef inline bint _check_monotonicity(
+ self,
+ cnp.int8_t monotonic_cst,
+ float64_t lower_bound,
+ float64_t upper_bound,
+ float64_t value_left,
+ float64_t value_right,
+ ) noexcept nogil:
+ cdef:
+ bint check_lower_bound = (
+ (value_left >= lower_bound) &
+ (value_right >= lower_bound)
+ )
+ bint check_upper_bound = (
+ (value_left <= upper_bound) &
+ (value_right <= upper_bound)
+ )
+ bint check_monotonic_cst = (
+ (value_left - value_right) * monotonic_cst <= 0
+ )
+ return check_lower_bound & check_upper_bound & check_monotonic_cst
+
+ cdef void init_sum_missing(self):
+ """Init sum_missing to hold sums for missing values."""
+
+cdef inline void _move_sums_classification(
+ ClassificationCriterion criterion,
+ float64_t[:, ::1] sum_1,
+ float64_t[:, ::1] sum_2,
+ float64_t* weighted_n_1,
+ float64_t* weighted_n_2,
+ bint put_missing_in_1,
+) noexcept nogil:
+ """Distribute sum_total and sum_missing into sum_1 and sum_2.
+
+ If there are missing values and:
+ - put_missing_in_1 is True, then missing values to go sum_1. Specifically:
+ sum_1 = sum_missing
+ sum_2 = sum_total - sum_missing
+
+ - put_missing_in_1 is False, then missing values go to sum_2. Specifically:
+ sum_1 = 0
+ sum_2 = sum_total
+ """
+ cdef intp_t k, c, n_bytes
+ if criterion.n_missing != 0 and put_missing_in_1:
+ for k in range(criterion.n_outputs):
+ n_bytes = criterion.n_classes[k] * sizeof(float64_t)
+ memcpy(&sum_1[k, 0], &criterion.sum_missing[k, 0], n_bytes)
+
+ for k in range(criterion.n_outputs):
+ for c in range(criterion.n_classes[k]):
+ sum_2[k, c] = criterion.sum_total[k, c] - criterion.sum_missing[k, c]
+
+ weighted_n_1[0] = criterion.weighted_n_missing
+ weighted_n_2[0] = criterion.weighted_n_node_samples - criterion.weighted_n_missing
+ else:
+ # Assigning sum_2 = sum_total for all outputs.
+ for k in range(criterion.n_outputs):
+ n_bytes = criterion.n_classes[k] * sizeof(float64_t)
+ memset(&sum_1[k, 0], 0, n_bytes)
+ memcpy(&sum_2[k, 0], &criterion.sum_total[k, 0], n_bytes)
+
+ weighted_n_1[0] = 0.0
+ weighted_n_2[0] = criterion.weighted_n_node_samples
+
+
+cdef class ClassificationCriterion(Criterion):
+ """Abstract criterion for classification."""
+
+ def __cinit__(self, intp_t n_outputs,
+ cnp.ndarray[intp_t, ndim=1] n_classes):
+ """Initialize attributes for this criterion.
+
+ Parameters
+ ----------
+ n_outputs : intp_t
+ The number of targets, the dimensionality of the prediction
+ n_classes : numpy.ndarray, dtype=intp_t
+ The number of unique classes in each target
+ """
+ self.start = 0
+ self.pos = 0
+ self.end = 0
+ self.missing_go_to_left = 0
+
+ self.n_outputs = n_outputs
+ self.n_samples = 0
+ self.n_node_samples = 0
+ self.weighted_n_node_samples = 0.0
+ self.weighted_n_left = 0.0
+ self.weighted_n_right = 0.0
+ self.weighted_n_missing = 0.0
+
+ self.n_classes = np.empty(n_outputs, dtype=np.intp)
+
+ cdef intp_t k = 0
+ cdef intp_t max_n_classes = 0
+
+ # For each target, set the number of unique classes in that target,
+ # and also compute the maximal stride of all targets
+ for k in range(n_outputs):
+ self.n_classes[k] = n_classes[k]
+
+ if n_classes[k] > max_n_classes:
+ max_n_classes = n_classes[k]
+
+ self.max_n_classes = max_n_classes
+
+ # Count labels for each output
+ self.sum_total = np.zeros((n_outputs, max_n_classes), dtype=np.float64)
+ self.sum_left = np.zeros((n_outputs, max_n_classes), dtype=np.float64)
+ self.sum_right = np.zeros((n_outputs, max_n_classes), dtype=np.float64)
+
+ def __reduce__(self):
+ return (type(self),
+ (self.n_outputs, np.asarray(self.n_classes)), self.__getstate__())
+
+ cdef int init(
+ self,
+ const float64_t[:, ::1] y,
+ const float64_t[:] sample_weight,
+ float64_t weighted_n_samples,
+ const intp_t[:] sample_indices,
+ intp_t start,
+ intp_t end
+ ) except -1 nogil:
+ """Initialize the criterion.
+
+ This initializes the criterion at node sample_indices[start:end] and children
+ sample_indices[start:start] and sample_indices[start:end].
+
+ Returns -1 in case of failure to allocate memory (and raise MemoryError)
+ or 0 otherwise.
+
+ Parameters
+ ----------
+ y : ndarray, dtype=float64_t
+ The target stored as a buffer for memory efficiency.
+ sample_weight : ndarray, dtype=float64_t
+ The weight of each sample stored as a Cython memoryview.
+ weighted_n_samples : float64_t
+ The total weight of all samples
+ sample_indices : ndarray, dtype=intp_t
+ A mask on the samples. Indices of the samples in X and y we want to use,
+ where sample_indices[start:end] correspond to the samples in this node.
+ start : intp_t
+ The first sample to use in the mask
+ end : intp_t
+ The last sample to use in the mask
+ """
+ self.y = y
+ self.sample_weight = sample_weight
+ self.sample_indices = sample_indices
+ self.start = start
+ self.end = end
+ self.n_node_samples = end - start
+ self.weighted_n_samples = weighted_n_samples
+ self.weighted_n_node_samples = 0.0
+
+ cdef intp_t i
+ cdef intp_t p
+ cdef intp_t k
+ cdef intp_t c
+ cdef float64_t w = 1.0
+
+ for k in range(self.n_outputs):
+ memset(&self.sum_total[k, 0], 0, self.n_classes[k] * sizeof(float64_t))
+
+ for p in range(start, end):
+ i = sample_indices[p]
+
+ # w is originally set to be 1.0, meaning that if no sample weights
+ # are given, the default weight of each sample is 1.0.
+ if sample_weight is not None:
+ w = sample_weight[i]
+
+ # Count weighted class frequency for each target
+ for k in range(self.n_outputs):
+ c = self.y[i, k]
+ self.sum_total[k, c] += w
+
+ self.weighted_n_node_samples += w
+
+ # Reset to pos=start
+ self.reset()
+ return 0
+
+ cdef void init_sum_missing(self):
+ """Init sum_missing to hold sums for missing values."""
+ self.sum_missing = np.zeros((self.n_outputs, self.max_n_classes), dtype=np.float64)
+
+ cdef void init_missing(self, intp_t n_missing) noexcept nogil:
+ """Initialize sum_missing if there are missing values.
+
+ This method assumes that caller placed the missing samples in
+ self.sample_indices[-n_missing:]
+ """
+ cdef intp_t i, p, k, c
+ cdef float64_t w = 1.0
+
+ self.n_missing = n_missing
+ if n_missing == 0:
+ return
+
+ memset(&self.sum_missing[0, 0], 0, self.max_n_classes * self.n_outputs * sizeof(float64_t))
+
+ self.weighted_n_missing = 0.0
+
+ # The missing samples are assumed to be in self.sample_indices[-n_missing:]
+ for p in range(self.end - n_missing, self.end):
+ i = self.sample_indices[p]
+ if self.sample_weight is not None:
+ w = self.sample_weight[i]
+
+ for k in range(self.n_outputs):
+ c = self.y[i, k]
+ self.sum_missing[k, c] += w
+
+ self.weighted_n_missing += w
+
+ cdef int reset(self) except -1 nogil:
+ """Reset the criterion at pos=start.
+
+ Returns -1 in case of failure to allocate memory (and raise MemoryError)
+ or 0 otherwise.
+ """
+ self.pos = self.start
+ _move_sums_classification(
+ self,
+ self.sum_left,
+ self.sum_right,
+ &self.weighted_n_left,
+ &self.weighted_n_right,
+ self.missing_go_to_left,
+ )
+ return 0
+
+ cdef int reverse_reset(self) except -1 nogil:
+ """Reset the criterion at pos=end.
+
+ Returns -1 in case of failure to allocate memory (and raise MemoryError)
+ or 0 otherwise.
+ """
+ self.pos = self.end
+ _move_sums_classification(
+ self,
+ self.sum_right,
+ self.sum_left,
+ &self.weighted_n_right,
+ &self.weighted_n_left,
+ not self.missing_go_to_left
+ )
+ return 0
+
+ cdef int update(self, intp_t new_pos) except -1 nogil:
+ """Updated statistics by moving sample_indices[pos:new_pos] to the left child.
+
+ Returns -1 in case of failure to allocate memory (and raise MemoryError)
+ or 0 otherwise.
+
+ Parameters
+ ----------
+ new_pos : intp_t
+ The new ending position for which to move sample_indices from the right
+ child to the left child.
+ """
+ cdef intp_t pos = self.pos
+ # The missing samples are assumed to be in
+ # self.sample_indices[-self.n_missing:] that is
+ # self.sample_indices[end_non_missing:self.end].
+ cdef intp_t end_non_missing = self.end - self.n_missing
+
+ cdef const intp_t[:] sample_indices = self.sample_indices
+ cdef const float64_t[:] sample_weight = self.sample_weight
+
+ cdef intp_t i
+ cdef intp_t p
+ cdef intp_t k
+ cdef intp_t c
+ cdef float64_t w = 1.0
+
+ # Update statistics up to new_pos
+ #
+ # Given that
+ # sum_left[x] + sum_right[x] = sum_total[x]
+ # and that sum_total is known, we are going to update
+ # sum_left from the direction that require the least amount
+ # of computations, i.e. from pos to new_pos or from end to new_po.
+ if (new_pos - pos) <= (end_non_missing - new_pos):
+ for p in range(pos, new_pos):
+ i = sample_indices[p]
+
+ if sample_weight is not None:
+ w = sample_weight[i]
+
+ for k in range(self.n_outputs):
+ self.sum_left[k, self.y[i, k]] += w
+
+ self.weighted_n_left += w
+
+ else:
+ self.reverse_reset()
+
+ for p in range(end_non_missing - 1, new_pos - 1, -1):
+ i = sample_indices[p]
+
+ if sample_weight is not None:
+ w = sample_weight[i]
+
+ for k in range(self.n_outputs):
+ self.sum_left[k, self.y[i, k]] -= w
+
+ self.weighted_n_left -= w
+
+ # Update right part statistics
+ self.weighted_n_right = self.weighted_n_node_samples - self.weighted_n_left
+ for k in range(self.n_outputs):
+ for c in range(self.n_classes[k]):
+ self.sum_right[k, c] = self.sum_total[k, c] - self.sum_left[k, c]
+
+ self.pos = new_pos
+ return 0
+
+ cdef float64_t node_impurity(self) noexcept nogil:
+ pass
+
+ cdef void children_impurity(self, float64_t* impurity_left,
+ float64_t* impurity_right) noexcept nogil:
+ pass
+
+ cdef void node_value(self, float64_t* dest) noexcept nogil:
+ """Compute the node value of sample_indices[start:end] and save it into dest.
+
+ Parameters
+ ----------
+ dest : float64_t pointer
+ The memory address which we will save the node value into.
+ """
+ cdef intp_t k, c
+
+ for k in range(self.n_outputs):
+ for c in range(self.n_classes[k]):
+ dest[c] = self.sum_total[k, c] / self.weighted_n_node_samples
+ dest += self.max_n_classes
+
+ cdef inline void clip_node_value(
+ self, float64_t * dest, float64_t lower_bound, float64_t upper_bound
+ ) noexcept nogil:
+ """Clip the values in dest such that predicted probabilities stay between
+ `lower_bound` and `upper_bound` when monotonic constraints are enforced.
+ Note that monotonicity constraints are only supported for:
+ - single-output trees and
+ - binary classifications.
+ """
+ if dest[0] < lower_bound:
+ dest[0] = lower_bound
+ elif dest[0] > upper_bound:
+ dest[0] = upper_bound
+
+ # Values for binary classification must sum to 1.
+ dest[1] = 1 - dest[0]
+
+ cdef inline float64_t middle_value(self) noexcept nogil:
+ """Compute the middle value of a split for monotonicity constraints as the simple average
+ of the left and right children values.
+
+ Note that monotonicity constraints are only supported for:
+ - single-output trees and
+ - binary classifications.
+ """
+ return (
+ (self.sum_left[0, 0] / (2 * self.weighted_n_left)) +
+ (self.sum_right[0, 0] / (2 * self.weighted_n_right))
+ )
+
+ cdef inline bint check_monotonicity(
+ self,
+ cnp.int8_t monotonic_cst,
+ float64_t lower_bound,
+ float64_t upper_bound,
+ ) noexcept nogil:
+ """Check monotonicity constraint is satisfied at the current classification split"""
+ cdef:
+ float64_t value_left = self.sum_left[0][0] / self.weighted_n_left
+ float64_t value_right = self.sum_right[0][0] / self.weighted_n_right
+
+ return self._check_monotonicity(monotonic_cst, lower_bound, upper_bound, value_left, value_right)
+
+
+cdef class Entropy(ClassificationCriterion):
+ r"""Cross Entropy impurity criterion.
+
+ This handles cases where the target is a classification taking values
+ 0, 1, ... K-2, K-1. If node m represents a region Rm with Nm observations,
+ then let
+
+ count_k = 1 / Nm \sum_{x_i in Rm} I(yi = k)
+
+ be the proportion of class k observations in node m.
+
+ The cross-entropy is then defined as
+
+ cross-entropy = -\sum_{k=0}^{K-1} count_k log(count_k)
+ """
+
+ cdef float64_t node_impurity(self) noexcept nogil:
+ """Evaluate the impurity of the current node.
+
+ Evaluate the cross-entropy criterion as impurity of the current node,
+ i.e. the impurity of sample_indices[start:end]. The smaller the impurity the
+ better.
+ """
+ cdef float64_t entropy = 0.0
+ cdef float64_t count_k
+ cdef intp_t k
+ cdef intp_t c
+
+ for k in range(self.n_outputs):
+ for c in range(self.n_classes[k]):
+ count_k = self.sum_total[k, c]
+ if count_k > 0.0:
+ count_k /= self.weighted_n_node_samples
+ entropy -= count_k * log(count_k)
+
+ return entropy / self.n_outputs
+
+ cdef void children_impurity(self, float64_t* impurity_left,
+ float64_t* impurity_right) noexcept nogil:
+ """Evaluate the impurity in children nodes.
+
+ i.e. the impurity of the left child (sample_indices[start:pos]) and the
+ impurity the right child (sample_indices[pos:end]).
+
+ Parameters
+ ----------
+ impurity_left : float64_t pointer
+ The memory address to save the impurity of the left node
+ impurity_right : float64_t pointer
+ The memory address to save the impurity of the right node
+ """
+ cdef float64_t entropy_left = 0.0
+ cdef float64_t entropy_right = 0.0
+ cdef float64_t count_k
+ cdef intp_t k
+ cdef intp_t c
+
+ for k in range(self.n_outputs):
+ for c in range(self.n_classes[k]):
+ count_k = self.sum_left[k, c]
+ if count_k > 0.0:
+ count_k /= self.weighted_n_left
+ entropy_left -= count_k * log(count_k)
+
+ count_k = self.sum_right[k, c]
+ if count_k > 0.0:
+ count_k /= self.weighted_n_right
+ entropy_right -= count_k * log(count_k)
+
+ impurity_left[0] = entropy_left / self.n_outputs
+ impurity_right[0] = entropy_right / self.n_outputs
+
+
+cdef class Gini(ClassificationCriterion):
+ r"""Gini Index impurity criterion.
+
+ This handles cases where the target is a classification taking values
+ 0, 1, ... K-2, K-1. If node m represents a region Rm with Nm observations,
+ then let
+
+ count_k = 1/ Nm \sum_{x_i in Rm} I(yi = k)
+
+ be the proportion of class k observations in node m.
+
+ The Gini Index is then defined as:
+
+ index = \sum_{k=0}^{K-1} count_k (1 - count_k)
+ = 1 - \sum_{k=0}^{K-1} count_k ** 2
+ """
+
+ cdef float64_t node_impurity(self) noexcept nogil:
+ """Evaluate the impurity of the current node.
+
+ Evaluate the Gini criterion as impurity of the current node,
+ i.e. the impurity of sample_indices[start:end]. The smaller the impurity the
+ better.
+ """
+ cdef float64_t gini = 0.0
+ cdef float64_t sq_count
+ cdef float64_t count_k
+ cdef intp_t k
+ cdef intp_t c
+
+ for k in range(self.n_outputs):
+ sq_count = 0.0
+
+ for c in range(self.n_classes[k]):
+ count_k = self.sum_total[k, c]
+ sq_count += count_k * count_k
+
+ gini += 1.0 - sq_count / (self.weighted_n_node_samples *
+ self.weighted_n_node_samples)
+
+ return gini / self.n_outputs
+
+ cdef void children_impurity(self, float64_t* impurity_left,
+ float64_t* impurity_right) noexcept nogil:
+ """Evaluate the impurity in children nodes.
+
+ i.e. the impurity of the left child (sample_indices[start:pos]) and the
+ impurity the right child (sample_indices[pos:end]) using the Gini index.
+
+ Parameters
+ ----------
+ impurity_left : float64_t pointer
+ The memory address to save the impurity of the left node to
+ impurity_right : float64_t pointer
+ The memory address to save the impurity of the right node to
+ """
+ cdef float64_t gini_left = 0.0
+ cdef float64_t gini_right = 0.0
+ cdef float64_t sq_count_left
+ cdef float64_t sq_count_right
+ cdef float64_t count_k
+ cdef intp_t k
+ cdef intp_t c
+
+ for k in range(self.n_outputs):
+ sq_count_left = 0.0
+ sq_count_right = 0.0
+
+ for c in range(self.n_classes[k]):
+ count_k = self.sum_left[k, c]
+ sq_count_left += count_k * count_k
+
+ count_k = self.sum_right[k, c]
+ sq_count_right += count_k * count_k
+
+ gini_left += 1.0 - sq_count_left / (self.weighted_n_left *
+ self.weighted_n_left)
+
+ gini_right += 1.0 - sq_count_right / (self.weighted_n_right *
+ self.weighted_n_right)
+
+ impurity_left[0] = gini_left / self.n_outputs
+ impurity_right[0] = gini_right / self.n_outputs
+
+
+cdef inline void _move_sums_regression(
+ RegressionCriterion criterion,
+ float64_t[::1] sum_1,
+ float64_t[::1] sum_2,
+ float64_t* weighted_n_1,
+ float64_t* weighted_n_2,
+ bint put_missing_in_1,
+) noexcept nogil:
+ """Distribute sum_total and sum_missing into sum_1 and sum_2.
+
+ If there are missing values and:
+ - put_missing_in_1 is True, then missing values to go sum_1. Specifically:
+ sum_1 = sum_missing
+ sum_2 = sum_total - sum_missing
+
+ - put_missing_in_1 is False, then missing values go to sum_2. Specifically:
+ sum_1 = 0
+ sum_2 = sum_total
+ """
+ cdef:
+ intp_t i
+ intp_t n_bytes = criterion.n_outputs * sizeof(float64_t)
+ bint has_missing = criterion.n_missing != 0
+
+ if has_missing and put_missing_in_1:
+ memcpy(&sum_1[0], &criterion.sum_missing[0], n_bytes)
+ for i in range(criterion.n_outputs):
+ sum_2[i] = criterion.sum_total[i] - criterion.sum_missing[i]
+ weighted_n_1[0] = criterion.weighted_n_missing
+ weighted_n_2[0] = criterion.weighted_n_node_samples - criterion.weighted_n_missing
+ else:
+ memset(&sum_1[0], 0, n_bytes)
+ # Assigning sum_2 = sum_total for all outputs.
+ memcpy(&sum_2[0], &criterion.sum_total[0], n_bytes)
+ weighted_n_1[0] = 0.0
+ weighted_n_2[0] = criterion.weighted_n_node_samples
+
+
+cdef class RegressionCriterion(Criterion):
+ r"""Abstract regression criterion.
+
+ This handles cases where the target is a continuous value, and is
+ evaluated by computing the variance of the target values left and right
+ of the split point. The computation takes linear time with `n_samples`
+ by using ::
+
+ var = \sum_i^n (y_i - y_bar) ** 2
+ = (\sum_i^n y_i ** 2) - n_samples * y_bar ** 2
+ """
+
+ def __cinit__(self, intp_t n_outputs, intp_t n_samples):
+ """Initialize parameters for this criterion.
+
+ Parameters
+ ----------
+ n_outputs : intp_t
+ The number of targets to be predicted
+
+ n_samples : intp_t
+ The total number of samples to fit on
+ """
+ # Default values
+ self.start = 0
+ self.pos = 0
+ self.end = 0
+
+ self.n_outputs = n_outputs
+ self.n_samples = n_samples
+ self.n_node_samples = 0
+ self.weighted_n_node_samples = 0.0
+ self.weighted_n_left = 0.0
+ self.weighted_n_right = 0.0
+ self.weighted_n_missing = 0.0
+
+ self.sq_sum_total = 0.0
+
+ self.sum_total = np.zeros(n_outputs, dtype=np.float64)
+ self.sum_left = np.zeros(n_outputs, dtype=np.float64)
+ self.sum_right = np.zeros(n_outputs, dtype=np.float64)
+
+ def __reduce__(self):
+ return (type(self), (self.n_outputs, self.n_samples), self.__getstate__())
+
+ cdef int init(
+ self,
+ const float64_t[:, ::1] y,
+ const float64_t[:] sample_weight,
+ float64_t weighted_n_samples,
+ const intp_t[:] sample_indices,
+ intp_t start,
+ intp_t end,
+ ) except -1 nogil:
+ """Initialize the criterion.
+
+ This initializes the criterion at node sample_indices[start:end] and children
+ sample_indices[start:start] and sample_indices[start:end].
+ """
+ # Initialize fields
+ self.y = y
+ self.sample_weight = sample_weight
+ self.sample_indices = sample_indices
+ self.start = start
+ self.end = end
+ self.n_node_samples = end - start
+ self.weighted_n_samples = weighted_n_samples
+ self.weighted_n_node_samples = 0.
+
+ cdef intp_t i
+ cdef intp_t p
+ cdef intp_t k
+ cdef float64_t y_ik
+ cdef float64_t w_y_ik
+ cdef float64_t w = 1.0
+ self.sq_sum_total = 0.0
+ memset(&self.sum_total[0], 0, self.n_outputs * sizeof(float64_t))
+
+ for p in range(start, end):
+ i = sample_indices[p]
+
+ if sample_weight is not None:
+ w = sample_weight[i]
+
+ for k in range(self.n_outputs):
+ y_ik = self.y[i, k]
+ w_y_ik = w * y_ik
+ self.sum_total[k] += w_y_ik
+ self.sq_sum_total += w_y_ik * y_ik
+
+ self.weighted_n_node_samples += w
+
+ # Reset to pos=start
+ self.reset()
+ return 0
+
+ cdef void init_sum_missing(self):
+ """Init sum_missing to hold sums for missing values."""
+ self.sum_missing = np.zeros(self.n_outputs, dtype=np.float64)
+
+ cdef void init_missing(self, intp_t n_missing) noexcept nogil:
+ """Initialize sum_missing if there are missing values.
+
+ This method assumes that caller placed the missing samples in
+ self.sample_indices[-n_missing:]
+ """
+ cdef intp_t i, p, k
+ cdef float64_t y_ik
+ cdef float64_t w_y_ik
+ cdef float64_t w = 1.0
+
+ self.n_missing = n_missing
+ if n_missing == 0:
+ return
+
+ memset(&self.sum_missing[0], 0, self.n_outputs * sizeof(float64_t))
+
+ self.weighted_n_missing = 0.0
+
+ # The missing samples are assumed to be in self.sample_indices[-n_missing:]
+ for p in range(self.end - n_missing, self.end):
+ i = self.sample_indices[p]
+ if self.sample_weight is not None:
+ w = self.sample_weight[i]
+
+ for k in range(self.n_outputs):
+ y_ik = self.y[i, k]
+ w_y_ik = w * y_ik
+ self.sum_missing[k] += w_y_ik
+
+ self.weighted_n_missing += w
+
+ cdef int reset(self) except -1 nogil:
+ """Reset the criterion at pos=start."""
+ self.pos = self.start
+ _move_sums_regression(
+ self,
+ self.sum_left,
+ self.sum_right,
+ &self.weighted_n_left,
+ &self.weighted_n_right,
+ self.missing_go_to_left
+ )
+ return 0
+
+ cdef int reverse_reset(self) except -1 nogil:
+ """Reset the criterion at pos=end."""
+ self.pos = self.end
+ _move_sums_regression(
+ self,
+ self.sum_right,
+ self.sum_left,
+ &self.weighted_n_right,
+ &self.weighted_n_left,
+ not self.missing_go_to_left
+ )
+ return 0
+
+ cdef int update(self, intp_t new_pos) except -1 nogil:
+ """Updated statistics by moving sample_indices[pos:new_pos] to the left."""
+ cdef const float64_t[:] sample_weight = self.sample_weight
+ cdef const intp_t[:] sample_indices = self.sample_indices
+
+ cdef intp_t pos = self.pos
+
+ # The missing samples are assumed to be in
+ # self.sample_indices[-self.n_missing:] that is
+ # self.sample_indices[end_non_missing:self.end].
+ cdef intp_t end_non_missing = self.end - self.n_missing
+ cdef intp_t i
+ cdef intp_t p
+ cdef intp_t k
+ cdef float64_t w = 1.0
+
+ # Update statistics up to new_pos
+ #
+ # Given that
+ # sum_left[x] + sum_right[x] = sum_total[x]
+ # and that sum_total is known, we are going to update
+ # sum_left from the direction that require the least amount
+ # of computations, i.e. from pos to new_pos or from end to new_pos.
+ if (new_pos - pos) <= (end_non_missing - new_pos):
+ for p in range(pos, new_pos):
+ i = sample_indices[p]
+
+ if sample_weight is not None:
+ w = sample_weight[i]
+
+ for k in range(self.n_outputs):
+ self.sum_left[k] += w * self.y[i, k]
+
+ self.weighted_n_left += w
+ else:
+ self.reverse_reset()
+
+ for p in range(end_non_missing - 1, new_pos - 1, -1):
+ i = sample_indices[p]
+
+ if sample_weight is not None:
+ w = sample_weight[i]
+
+ for k in range(self.n_outputs):
+ self.sum_left[k] -= w * self.y[i, k]
+
+ self.weighted_n_left -= w
+
+ self.weighted_n_right = (self.weighted_n_node_samples -
+ self.weighted_n_left)
+ for k in range(self.n_outputs):
+ self.sum_right[k] = self.sum_total[k] - self.sum_left[k]
+
+ self.pos = new_pos
+ return 0
+
+ cdef float64_t node_impurity(self) noexcept nogil:
+ pass
+
+ cdef void children_impurity(self, float64_t* impurity_left,
+ float64_t* impurity_right) noexcept nogil:
+ pass
+
+ cdef void node_value(self, float64_t* dest) noexcept nogil:
+ """Compute the node value of sample_indices[start:end] into dest."""
+ cdef intp_t k
+
+ for k in range(self.n_outputs):
+ dest[k] = self.sum_total[k] / self.weighted_n_node_samples
+
+ cdef inline void clip_node_value(self, float64_t* dest, float64_t lower_bound, float64_t upper_bound) noexcept nogil:
+ """Clip the value in dest between lower_bound and upper_bound for monotonic constraints."""
+ if dest[0] < lower_bound:
+ dest[0] = lower_bound
+ elif dest[0] > upper_bound:
+ dest[0] = upper_bound
+
+ cdef float64_t middle_value(self) noexcept nogil:
+ """Compute the middle value of a split for monotonicity constraints as the simple average
+ of the left and right children values.
+
+ Monotonicity constraints are only supported for single-output trees we can safely assume
+ n_outputs == 1.
+ """
+ return (
+ (self.sum_left[0] / (2 * self.weighted_n_left)) +
+ (self.sum_right[0] / (2 * self.weighted_n_right))
+ )
+
+ cdef bint check_monotonicity(
+ self,
+ cnp.int8_t monotonic_cst,
+ float64_t lower_bound,
+ float64_t upper_bound,
+ ) noexcept nogil:
+ """Check monotonicity constraint is satisfied at the current regression split"""
+ cdef:
+ float64_t value_left = self.sum_left[0] / self.weighted_n_left
+ float64_t value_right = self.sum_right[0] / self.weighted_n_right
+
+ return self._check_monotonicity(monotonic_cst, lower_bound, upper_bound, value_left, value_right)
+
+cdef class MSE(RegressionCriterion):
+ """Mean squared error impurity criterion.
+
+ MSE = var_left + var_right
+ """
+
+ cdef float64_t node_impurity(self) noexcept nogil:
+ """Evaluate the impurity of the current node.
+
+ Evaluate the MSE criterion as impurity of the current node,
+ i.e. the impurity of sample_indices[start:end]. The smaller the impurity the
+ better.
+ """
+ cdef float64_t impurity
+ cdef intp_t k
+
+ impurity = self.sq_sum_total / self.weighted_n_node_samples
+ for k in range(self.n_outputs):
+ impurity -= (self.sum_total[k] / self.weighted_n_node_samples)**2.0
+
+ return impurity / self.n_outputs
+
+ cdef float64_t proxy_impurity_improvement(self) noexcept nogil:
+ """Compute a proxy of the impurity reduction.
+
+ This method is used to speed up the search for the best split.
+ It is a proxy quantity such that the split that maximizes this value
+ also maximizes the impurity improvement. It neglects all constant terms
+ of the impurity decrease for a given split.
+
+ The absolute impurity improvement is only computed by the
+ impurity_improvement method once the best split has been found.
+
+ The MSE proxy is derived from
+
+ sum_{i left}(y_i - y_pred_L)^2 + sum_{i right}(y_i - y_pred_R)^2
+ = sum(y_i^2) - n_L * mean_{i left}(y_i)^2 - n_R * mean_{i right}(y_i)^2
+
+ Neglecting constant terms, this gives:
+
+ - 1/n_L * sum_{i left}(y_i)^2 - 1/n_R * sum_{i right}(y_i)^2
+ """
+ cdef intp_t k
+ cdef float64_t proxy_impurity_left = 0.0
+ cdef float64_t proxy_impurity_right = 0.0
+
+ for k in range(self.n_outputs):
+ proxy_impurity_left += self.sum_left[k] * self.sum_left[k]
+ proxy_impurity_right += self.sum_right[k] * self.sum_right[k]
+
+ return (proxy_impurity_left / self.weighted_n_left +
+ proxy_impurity_right / self.weighted_n_right)
+
+ cdef void children_impurity(self, float64_t* impurity_left,
+ float64_t* impurity_right) noexcept nogil:
+ """Evaluate the impurity in children nodes.
+
+ i.e. the impurity of the left child (sample_indices[start:pos]) and the
+ impurity the right child (sample_indices[pos:end]).
+ """
+ cdef const float64_t[:] sample_weight = self.sample_weight
+ cdef const intp_t[:] sample_indices = self.sample_indices
+ cdef intp_t pos = self.pos
+ cdef intp_t start = self.start
+
+ cdef float64_t y_ik
+
+ cdef float64_t sq_sum_left = 0.0
+ cdef float64_t sq_sum_right
+
+ cdef intp_t i
+ cdef intp_t p
+ cdef intp_t k
+ cdef float64_t w = 1.0
+
+ cdef intp_t end_non_missing
+
+ for p in range(start, pos):
+ i = sample_indices[p]
+
+ if sample_weight is not None:
+ w = sample_weight[i]
+
+ for k in range(self.n_outputs):
+ y_ik = self.y[i, k]
+ sq_sum_left += w * y_ik * y_ik
+
+ if self.missing_go_to_left:
+ # add up the impact of these missing values on the left child
+ # statistics.
+ # Note: this only impacts the square sum as the sum
+ # is modified elsewhere.
+ end_non_missing = self.end - self.n_missing
+
+ for p in range(end_non_missing, self.end):
+ i = sample_indices[p]
+ if sample_weight is not None:
+ w = sample_weight[i]
+
+ for k in range(self.n_outputs):
+ y_ik = self.y[i, k]
+ sq_sum_left += w * y_ik * y_ik
+
+ sq_sum_right = self.sq_sum_total - sq_sum_left
+
+ impurity_left[0] = sq_sum_left / self.weighted_n_left
+ impurity_right[0] = sq_sum_right / self.weighted_n_right
+
+ for k in range(self.n_outputs):
+ impurity_left[0] -= (self.sum_left[k] / self.weighted_n_left) ** 2.0
+ impurity_right[0] -= (self.sum_right[k] / self.weighted_n_right) ** 2.0
+
+ impurity_left[0] /= self.n_outputs
+ impurity_right[0] /= self.n_outputs
+
+
+cdef class MAE(RegressionCriterion):
+ r"""Mean absolute error impurity criterion.
+
+ MAE = (1 / n)*(\sum_i |y_i - f_i|), where y_i is the true
+ value and f_i is the predicted value."""
+
+ cdef cnp.ndarray left_child
+ cdef cnp.ndarray right_child
+ cdef void** left_child_ptr
+ cdef void** right_child_ptr
+ cdef float64_t[::1] node_medians
+
+ def __cinit__(self, intp_t n_outputs, intp_t n_samples):
+ """Initialize parameters for this criterion.
+
+ Parameters
+ ----------
+ n_outputs : intp_t
+ The number of targets to be predicted
+
+ n_samples : intp_t
+ The total number of samples to fit on
+ """
+ # Default values
+ self.start = 0
+ self.pos = 0
+ self.end = 0
+
+ self.n_outputs = n_outputs
+ self.n_samples = n_samples
+ self.n_node_samples = 0
+ self.weighted_n_node_samples = 0.0
+ self.weighted_n_left = 0.0
+ self.weighted_n_right = 0.0
+
+ self.node_medians = np.zeros(n_outputs, dtype=np.float64)
+
+ self.left_child = np.empty(n_outputs, dtype='object')
+ self.right_child = np.empty(n_outputs, dtype='object')
+ # initialize WeightedMedianCalculators
+ for k in range(n_outputs):
+ self.left_child[k] = WeightedMedianCalculator(n_samples)
+ self.right_child[k] = WeightedMedianCalculator(n_samples)
+
+ self.left_child_ptr = cnp.PyArray_DATA(self.left_child)
+ self.right_child_ptr = cnp.PyArray_DATA(self.right_child)
+
+ cdef int init(
+ self,
+ const float64_t[:, ::1] y,
+ const float64_t[:] sample_weight,
+ float64_t weighted_n_samples,
+ const intp_t[:] sample_indices,
+ intp_t start,
+ intp_t end,
+ ) except -1 nogil:
+ """Initialize the criterion.
+
+ This initializes the criterion at node sample_indices[start:end] and children
+ sample_indices[start:start] and sample_indices[start:end].
+ """
+ cdef intp_t i, p, k
+ cdef float64_t w = 1.0
+
+ # Initialize fields
+ self.y = y
+ self.sample_weight = sample_weight
+ self.sample_indices = sample_indices
+ self.start = start
+ self.end = end
+ self.n_node_samples = end - start
+ self.weighted_n_samples = weighted_n_samples
+ self.weighted_n_node_samples = 0.
+
+ cdef void** left_child = self.left_child_ptr
+ cdef void** right_child = self.right_child_ptr
+
+ for k in range(self.n_outputs):
+ ( left_child[k]).reset()
+ ( right_child[k]).reset()
+
+ for p in range(start, end):
+ i = sample_indices[p]
+
+ if sample_weight is not None:
+ w = sample_weight[i]
+
+ for k in range(self.n_outputs):
+ # push method ends up calling safe_realloc, hence `except -1`
+ # push all values to the right side,
+ # since pos = start initially anyway
+ ( right_child[k]).push(self.y[i, k], w)
+
+ self.weighted_n_node_samples += w
+ # calculate the node medians
+ for k in range(self.n_outputs):
+ self.node_medians[k] = ( right_child[k]).get_median()
+
+ # Reset to pos=start
+ self.reset()
+ return 0
+
+ cdef void init_missing(self, intp_t n_missing) noexcept nogil:
+ """Raise error if n_missing != 0."""
+ if n_missing == 0:
+ return
+ with gil:
+ raise ValueError("missing values is not supported for MAE.")
+
+ cdef int reset(self) except -1 nogil:
+ """Reset the criterion at pos=start.
+
+ Returns -1 in case of failure to allocate memory (and raise MemoryError)
+ or 0 otherwise.
+ """
+ cdef intp_t i, k
+ cdef float64_t value
+ cdef float64_t weight
+
+ cdef void** left_child = self.left_child_ptr
+ cdef void** right_child = self.right_child_ptr
+
+ self.weighted_n_left = 0.0
+ self.weighted_n_right = self.weighted_n_node_samples
+ self.pos = self.start
+
+ # reset the WeightedMedianCalculators, left should have no
+ # elements and right should have all elements.
+
+ for k in range(self.n_outputs):
+ # if left has no elements, it's already reset
+ for i in range(( left_child[k]).size()):
+ # remove everything from left and put it into right
+ ( left_child[k]).pop(&value,
+ &weight)
+ # push method ends up calling safe_realloc, hence `except -1`
+ ( right_child[k]).push(value,
+ weight)
+ return 0
+
+ cdef int reverse_reset(self) except -1 nogil:
+ """Reset the criterion at pos=end.
+
+ Returns -1 in case of failure to allocate memory (and raise MemoryError)
+ or 0 otherwise.
+ """
+ self.weighted_n_right = 0.0
+ self.weighted_n_left = self.weighted_n_node_samples
+ self.pos = self.end
+
+ cdef float64_t value
+ cdef float64_t weight
+ cdef void** left_child = self.left_child_ptr
+ cdef void** right_child = self.right_child_ptr
+
+ # reverse reset the WeightedMedianCalculators, right should have no
+ # elements and left should have all elements.
+ for k in range(self.n_outputs):
+ # if right has no elements, it's already reset
+ for i in range(( right_child[k]).size()):
+ # remove everything from right and put it into left
+ ( right_child[k]).pop(&value,
+ &weight)
+ # push method ends up calling safe_realloc, hence `except -1`
+ ( left_child[k]).push(value,
+ weight)
+ return 0
+
+ cdef int update(self, intp_t new_pos) except -1 nogil:
+ """Updated statistics by moving sample_indices[pos:new_pos] to the left.
+
+ Returns -1 in case of failure to allocate memory (and raise MemoryError)
+ or 0 otherwise.
+ """
+ cdef const float64_t[:] sample_weight = self.sample_weight
+ cdef const intp_t[:] sample_indices = self.sample_indices
+
+ cdef void** left_child = self.left_child_ptr
+ cdef void** right_child = self.right_child_ptr
+
+ cdef intp_t pos = self.pos
+ cdef intp_t end = self.end
+ cdef intp_t i, p, k
+ cdef float64_t w = 1.0
+
+ # Update statistics up to new_pos
+ #
+ # We are going to update right_child and left_child
+ # from the direction that require the least amount of
+ # computations, i.e. from pos to new_pos or from end to new_pos.
+ if (new_pos - pos) <= (end - new_pos):
+ for p in range(pos, new_pos):
+ i = sample_indices[p]
+
+ if sample_weight is not None:
+ w = sample_weight[i]
+
+ for k in range(self.n_outputs):
+ # remove y_ik and its weight w from right and add to left
+ ( right_child[k]).remove(self.y[i, k], w)
+ # push method ends up calling safe_realloc, hence except -1
+ ( left_child[k]).push(self.y[i, k], w)
+
+ self.weighted_n_left += w
+ else:
+ self.reverse_reset()
+
+ for p in range(end - 1, new_pos - 1, -1):
+ i = sample_indices[p]
+
+ if sample_weight is not None:
+ w = sample_weight[i]
+
+ for k in range(self.n_outputs):
+ # remove y_ik and its weight w from left and add to right
+ ( left_child[k]).remove(self.y[i, k], w)
+ ( right_child[k]).push(self.y[i, k], w)
+
+ self.weighted_n_left -= w
+
+ self.weighted_n_right = (self.weighted_n_node_samples -
+ self.weighted_n_left)
+ self.pos = new_pos
+ return 0
+
+ cdef void node_value(self, float64_t* dest) noexcept nogil:
+ """Computes the node value of sample_indices[start:end] into dest."""
+ cdef intp_t k
+ for k in range(self.n_outputs):
+ dest[k] = self.node_medians[k]
+
+ cdef inline float64_t middle_value(self) noexcept nogil:
+ """Compute the middle value of a split for monotonicity constraints as the simple average
+ of the left and right children values.
+
+ Monotonicity constraints are only supported for single-output trees we can safely assume
+ n_outputs == 1.
+ """
+ return (
+ ( self.left_child_ptr[0]).get_median() +
+ ( self.right_child_ptr[0]).get_median()
+ ) / 2
+
+ cdef inline bint check_monotonicity(
+ self,
+ cnp.int8_t monotonic_cst,
+ float64_t lower_bound,
+ float64_t upper_bound,
+ ) noexcept nogil:
+ """Check monotonicity constraint is satisfied at the current regression split"""
+ cdef:
+ float64_t value_left = ( self.left_child_ptr[0]).get_median()
+ float64_t value_right = ( self.right_child_ptr[0]).get_median()
+
+ return self._check_monotonicity(monotonic_cst, lower_bound, upper_bound, value_left, value_right)
+
+ cdef float64_t node_impurity(self) noexcept nogil:
+ """Evaluate the impurity of the current node.
+
+ Evaluate the MAE criterion as impurity of the current node,
+ i.e. the impurity of sample_indices[start:end]. The smaller the impurity the
+ better.
+ """
+ cdef const float64_t[:] sample_weight = self.sample_weight
+ cdef const intp_t[:] sample_indices = self.sample_indices
+ cdef intp_t i, p, k
+ cdef float64_t w = 1.0
+ cdef float64_t impurity = 0.0
+
+ for k in range(self.n_outputs):
+ for p in range(self.start, self.end):
+ i = sample_indices[p]
+
+ if sample_weight is not None:
+ w = sample_weight[i]
+
+ impurity += fabs(self.y[i, k] - self.node_medians[k]) * w
+
+ return impurity / (self.weighted_n_node_samples * self.n_outputs)
+
+ cdef void children_impurity(self, float64_t* p_impurity_left,
+ float64_t* p_impurity_right) noexcept nogil:
+ """Evaluate the impurity in children nodes.
+
+ i.e. the impurity of the left child (sample_indices[start:pos]) and the
+ impurity the right child (sample_indices[pos:end]).
+ """
+ cdef const float64_t[:] sample_weight = self.sample_weight
+ cdef const intp_t[:] sample_indices = self.sample_indices
+
+ cdef intp_t start = self.start
+ cdef intp_t pos = self.pos
+ cdef intp_t end = self.end
+
+ cdef intp_t i, p, k
+ cdef float64_t median
+ cdef float64_t w = 1.0
+ cdef float64_t impurity_left = 0.0
+ cdef float64_t impurity_right = 0.0
+
+ cdef void** left_child = self.left_child_ptr
+ cdef void** right_child = self.right_child_ptr
+
+ for k in range(self.n_outputs):
+ median = ( left_child[k]).get_median()
+ for p in range(start, pos):
+ i = sample_indices[p]
+
+ if sample_weight is not None:
+ w = sample_weight[i]
+
+ impurity_left += fabs(self.y[i, k] - median) * w
+ p_impurity_left[0] = impurity_left / (self.weighted_n_left *
+ self.n_outputs)
+
+ for k in range(self.n_outputs):
+ median = ( right_child[k]).get_median()
+ for p in range(pos, end):
+ i = sample_indices[p]
+
+ if sample_weight is not None:
+ w = sample_weight[i]
+
+ impurity_right += fabs(self.y[i, k] - median) * w
+ p_impurity_right[0] = impurity_right / (self.weighted_n_right *
+ self.n_outputs)
+
+
+cdef class FriedmanMSE(MSE):
+ """Mean squared error impurity criterion with improvement score by Friedman.
+
+ Uses the formula (35) in Friedman's original Gradient Boosting paper:
+
+ diff = mean_left - mean_right
+ improvement = n_left * n_right * diff^2 / (n_left + n_right)
+ """
+
+ cdef float64_t proxy_impurity_improvement(self) noexcept nogil:
+ """Compute a proxy of the impurity reduction.
+
+ This method is used to speed up the search for the best split.
+ It is a proxy quantity such that the split that maximizes this value
+ also maximizes the impurity improvement. It neglects all constant terms
+ of the impurity decrease for a given split.
+
+ The absolute impurity improvement is only computed by the
+ impurity_improvement method once the best split has been found.
+ """
+ cdef float64_t total_sum_left = 0.0
+ cdef float64_t total_sum_right = 0.0
+
+ cdef intp_t k
+ cdef float64_t diff = 0.0
+
+ for k in range(self.n_outputs):
+ total_sum_left += self.sum_left[k]
+ total_sum_right += self.sum_right[k]
+
+ diff = (self.weighted_n_right * total_sum_left -
+ self.weighted_n_left * total_sum_right)
+
+ return diff * diff / (self.weighted_n_left * self.weighted_n_right)
+
+ cdef float64_t impurity_improvement(self, float64_t impurity_parent, float64_t
+ impurity_left, float64_t impurity_right) noexcept nogil:
+ # Note: none of the arguments are used here
+ cdef float64_t total_sum_left = 0.0
+ cdef float64_t total_sum_right = 0.0
+
+ cdef intp_t k
+ cdef float64_t diff = 0.0
+
+ for k in range(self.n_outputs):
+ total_sum_left += self.sum_left[k]
+ total_sum_right += self.sum_right[k]
+
+ diff = (self.weighted_n_right * total_sum_left -
+ self.weighted_n_left * total_sum_right) / self.n_outputs
+
+ return (diff * diff / (self.weighted_n_left * self.weighted_n_right *
+ self.weighted_n_node_samples))
+
+
+cdef class Poisson(RegressionCriterion):
+ """Half Poisson deviance as impurity criterion.
+
+ Poisson deviance = 2/n * sum(y_true * log(y_true/y_pred) + y_pred - y_true)
+
+ Note that the deviance is >= 0, and since we have `y_pred = mean(y_true)`
+ at the leaves, one always has `sum(y_pred - y_true) = 0`. It remains the
+ implemented impurity (factor 2 is skipped):
+ 1/n * sum(y_true * log(y_true/y_pred)
+ """
+ # FIXME in 1.0:
+ # min_impurity_split with default = 0 forces us to use a non-negative
+ # impurity like the Poisson deviance. Without this restriction, one could
+ # throw away the 'constant' term sum(y_true * log(y_true)) and just use
+ # Poisson loss = - 1/n * sum(y_true * log(y_pred))
+ # = - 1/n * sum(y_true * log(mean(y_true))
+ # = - mean(y_true) * log(mean(y_true))
+ # With this trick (used in proxy_impurity_improvement()), as for MSE,
+ # children_impurity would only need to go over left xor right split, not
+ # both. This could be faster.
+
+ cdef float64_t node_impurity(self) noexcept nogil:
+ """Evaluate the impurity of the current node.
+
+ Evaluate the Poisson criterion as impurity of the current node,
+ i.e. the impurity of sample_indices[start:end]. The smaller the impurity the
+ better.
+ """
+ return self.poisson_loss(self.start, self.end, self.sum_total,
+ self.weighted_n_node_samples)
+
+ cdef float64_t proxy_impurity_improvement(self) noexcept nogil:
+ """Compute a proxy of the impurity reduction.
+
+ This method is used to speed up the search for the best split.
+ It is a proxy quantity such that the split that maximizes this value
+ also maximizes the impurity improvement. It neglects all constant terms
+ of the impurity decrease for a given split.
+
+ The absolute impurity improvement is only computed by the
+ impurity_improvement method once the best split has been found.
+
+ The Poisson proxy is derived from:
+
+ sum_{i left }(y_i * log(y_i / y_pred_L))
+ + sum_{i right}(y_i * log(y_i / y_pred_R))
+ = sum(y_i * log(y_i) - n_L * mean_{i left}(y_i) * log(mean_{i left}(y_i))
+ - n_R * mean_{i right}(y_i) * log(mean_{i right}(y_i))
+
+ Neglecting constant terms, this gives
+
+ - sum{i left }(y_i) * log(mean{i left}(y_i))
+ - sum{i right}(y_i) * log(mean{i right}(y_i))
+ """
+ cdef intp_t k
+ cdef float64_t proxy_impurity_left = 0.0
+ cdef float64_t proxy_impurity_right = 0.0
+ cdef float64_t y_mean_left = 0.
+ cdef float64_t y_mean_right = 0.
+
+ for k in range(self.n_outputs):
+ if (self.sum_left[k] <= EPSILON) or (self.sum_right[k] <= EPSILON):
+ # Poisson loss does not allow non-positive predictions. We
+ # therefore forbid splits that have child nodes with
+ # sum(y_i) <= 0.
+ # Since sum_right = sum_total - sum_left, it can lead to
+ # floating point rounding error and will not give zero. Thus,
+ # we relax the above comparison to sum(y_i) <= EPSILON.
+ return -INFINITY
+ else:
+ y_mean_left = self.sum_left[k] / self.weighted_n_left
+ y_mean_right = self.sum_right[k] / self.weighted_n_right
+ proxy_impurity_left -= self.sum_left[k] * log(y_mean_left)
+ proxy_impurity_right -= self.sum_right[k] * log(y_mean_right)
+
+ return - proxy_impurity_left - proxy_impurity_right
+
+ cdef void children_impurity(self, float64_t* impurity_left,
+ float64_t* impurity_right) noexcept nogil:
+ """Evaluate the impurity in children nodes.
+
+ i.e. the impurity of the left child (sample_indices[start:pos]) and the
+ impurity of the right child (sample_indices[pos:end]) for Poisson.
+ """
+ cdef intp_t start = self.start
+ cdef intp_t pos = self.pos
+ cdef intp_t end = self.end
+
+ impurity_left[0] = self.poisson_loss(start, pos, self.sum_left,
+ self.weighted_n_left)
+
+ impurity_right[0] = self.poisson_loss(pos, end, self.sum_right,
+ self.weighted_n_right)
+
+ cdef inline float64_t poisson_loss(
+ self,
+ intp_t start,
+ intp_t end,
+ const float64_t[::1] y_sum,
+ float64_t weight_sum
+ ) noexcept nogil:
+ """Helper function to compute Poisson loss (~deviance) of a given node.
+ """
+ cdef const float64_t[:, ::1] y = self.y
+ cdef const float64_t[:] sample_weight = self.sample_weight
+ cdef const intp_t[:] sample_indices = self.sample_indices
+
+ cdef float64_t y_mean = 0.
+ cdef float64_t poisson_loss = 0.
+ cdef float64_t w = 1.0
+ cdef intp_t i, k, p
+ cdef intp_t n_outputs = self.n_outputs
+
+ for k in range(n_outputs):
+ if y_sum[k] <= EPSILON:
+ # y_sum could be computed from the subtraction
+ # sum_right = sum_total - sum_left leading to a potential
+ # floating point rounding error.
+ # Thus, we relax the comparison y_sum <= 0 to
+ # y_sum <= EPSILON.
+ return INFINITY
+
+ y_mean = y_sum[k] / weight_sum
+
+ for p in range(start, end):
+ i = sample_indices[p]
+
+ if sample_weight is not None:
+ w = sample_weight[i]
+
+ poisson_loss += w * xlogy(y[i, k], y[i, k] / y_mean)
+ return poisson_loss / (weight_sum * n_outputs)
diff --git a/causalml/source/causalml/inference/tree/_tree/_splitter.pxd b/causalml/source/causalml/inference/tree/_tree/_splitter.pxd
new file mode 100644
index 0000000000000000000000000000000000000000..c02c9a150d8c485bf05ce4e7e974a52487951f29
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/_tree/_splitter.pxd
@@ -0,0 +1,117 @@
+# Authors: Gilles Louppe
+# Peter Prettenhofer
+# Brian Holt
+# Joel Nothman
+# Arnaud Joly
+# Jacob Schreiber
+#
+# License: BSD 3 clause
+
+# distutils: language = c++
+# cython: cdivision=True
+# cython: boundscheck=False
+# cython: wraparound=False
+# cython: language_level=3
+# cython: linetrace=True
+
+# See _splitter.pyx for details.
+from ._criterion cimport Criterion
+from ._tree cimport ParentInfo
+
+from ._typedefs cimport float32_t, float64_t, intp_t, int8_t, int32_t, uint32_t
+
+
+cdef struct SplitRecord:
+ # Data to track sample split
+ intp_t feature # Which feature to split on.
+ intp_t pos # Split samples array at the given position,
+ # # i.e. count of samples below threshold for feature.
+ # # pos is >= end if the node is a leaf.
+ float64_t threshold # Threshold to split at.
+ float64_t improvement # Impurity improvement given parent node.
+ float64_t impurity_left # Impurity of the left split.
+ float64_t impurity_right # Impurity of the right split.
+ float64_t lower_bound # Lower bound on value of both children for monotonicity
+ float64_t upper_bound # Upper bound on value of both children for monotonicity
+ unsigned char missing_go_to_left # Controls if missing values go to the left node.
+ intp_t n_missing # Number of missing values for the feature being split on
+
+cdef class Splitter:
+ # The splitter searches in the input space for a feature and a threshold
+ # to split the samples samples[start:end].
+ #
+ # The impurity computations are delegated to a criterion object.
+
+ # Internal structures
+ cdef public Criterion criterion # Impurity criterion
+ cdef public intp_t max_features # Number of features to test
+ cdef public intp_t min_samples_leaf # Min samples in a leaf
+ cdef public float64_t min_weight_leaf # Minimum weight in a leaf
+
+ cdef object random_state # Random state
+ cdef uint32_t rand_r_state # sklearn_rand_r random number state
+
+ cdef intp_t[::1] samples # Sample indices in X, y
+ cdef intp_t n_samples # X.shape[0]
+ cdef float64_t weighted_n_samples # Weighted number of samples
+ cdef intp_t[::1] features # Feature indices in X
+ cdef intp_t[::1] constant_features # Constant features indices
+ cdef intp_t n_features # X.shape[1]
+ cdef float32_t[::1] feature_values # temp. array holding feature values
+
+ cdef intp_t start # Start position for the current node
+ cdef intp_t end # End position for the current node
+
+ cdef const float64_t[:, ::1] y
+ # Monotonicity constraints for each feature.
+ # The encoding is as follows:
+ # -1: monotonic decrease
+ # 0: no constraint
+ # +1: monotonic increase
+ cdef const int8_t[:] monotonic_cst
+ cdef bint with_monotonic_cst
+ cdef const float64_t[:] sample_weight
+
+ # The samples vector `samples` is maintained by the Splitter object such
+ # that the samples contained in a node are contiguous. With this setting,
+ # `node_split` reorganizes the node samples `samples[start:end]` in two
+ # subsets `samples[start:pos]` and `samples[pos:end]`.
+
+ # The 1-d `features` array of size n_features contains the features
+ # indices and allows fast sampling without replacement of features.
+
+ # The 1-d `constant_features` array of size n_features holds in
+ # `constant_features[:n_constant_features]` the feature ids with
+ # constant values for all the samples that reached a specific node.
+ # The value `n_constant_features` is given by the parent node to its
+ # child nodes. The content of the range `[n_constant_features:]` is left
+ # undefined, but preallocated for performance reasons
+ # This allows optimization with depth-based tree building.
+
+ # Methods
+ cdef int init(
+ self,
+ object X,
+ const float64_t[:, ::1] y,
+ const float64_t[:] sample_weight,
+ const unsigned char[::1] missing_values_in_feature_mask,
+ ) except -1
+
+ cdef int node_reset(
+ self,
+ intp_t start,
+ intp_t end,
+ float64_t* weighted_n_node_samples
+ ) except -1 nogil
+
+ cdef int node_split(
+ self,
+ ParentInfo* parent,
+ SplitRecord* split,
+ ) except -1 nogil
+
+ cdef void node_value(self, float64_t* dest) noexcept nogil
+
+ cdef void clip_node_value(self, float64_t* dest, float64_t lower_bound, float64_t upper_bound) noexcept nogil
+
+ cdef float64_t node_impurity(self) noexcept nogil
diff --git a/causalml/source/causalml/inference/tree/_tree/_splitter.pyx b/causalml/source/causalml/inference/tree/_tree/_splitter.pyx
new file mode 100644
index 0000000000000000000000000000000000000000..2eac9f4d3f439be06aaf10e79036d580513ca49a
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/_tree/_splitter.pyx
@@ -0,0 +1,1622 @@
+# Authors: Gilles Louppe
+# Peter Prettenhofer
+# Brian Holt
+# Noel Dawe
+# Satrajit Gosh
+# Lars Buitinck
+# Arnaud Joly
+# Joel Nothman
+# Fares Hedayati
+# Jacob Schreiber
+#
+# License: BSD 3 clause
+
+# cython: cdivision=True
+# cython: boundscheck=False
+# cython: wraparound=False
+# cython: language_level=3
+# cython: linetrace=True
+
+from cython cimport final
+from libc.math cimport isnan
+from libc.stdlib cimport qsort
+from libc.string cimport memcpy
+
+from ._criterion cimport Criterion
+from ._utils cimport log
+from ._utils cimport rand_int
+from ._utils cimport rand_uniform
+from ._utils cimport RAND_R_MAX
+from ._typedefs cimport int8_t
+
+import numpy as np
+from scipy.sparse import issparse
+
+
+cdef float64_t INFINITY = np.inf
+
+# Mitigate precision differences between 32 bit and 64 bit
+cdef float32_t FEATURE_THRESHOLD = 1e-7
+
+# Constant to switch between algorithm non zero value extract algorithm
+# in SparsePartitioner
+cdef float32_t EXTRACT_NNZ_SWITCH = 0.1
+
+cdef inline void _init_split(SplitRecord* self, intp_t start_pos) noexcept nogil:
+ self.impurity_left = INFINITY
+ self.impurity_right = INFINITY
+ self.pos = start_pos
+ self.feature = 0
+ self.threshold = 0.
+ self.improvement = -INFINITY
+ self.missing_go_to_left = False
+ self.n_missing = 0
+
+cdef class Splitter:
+ """Abstract splitter class.
+
+ Splitters are called by tree builders to find the best splits on both
+ sparse and dense data, one split at a time.
+ """
+
+ def __cinit__(
+ self,
+ Criterion criterion,
+ intp_t max_features,
+ intp_t min_samples_leaf,
+ float64_t min_weight_leaf,
+ object random_state,
+ const int8_t[:] monotonic_cst,
+ ):
+ """
+ Parameters
+ ----------
+ criterion : Criterion
+ The criterion to measure the quality of a split.
+
+ max_features : intp_t
+ The maximal number of randomly selected features which can be
+ considered for a split.
+
+ min_samples_leaf : intp_t
+ The minimal number of samples each leaf can have, where splits
+ which would result in having less samples in a leaf are not
+ considered.
+
+ min_weight_leaf : float64_t
+ The minimal weight each leaf can have, where the weight is the sum
+ of the weights of each sample in it.
+
+ random_state : object
+ The user inputted random state to be used for pseudo-randomness
+
+ monotonic_cst : const int8_t[:]
+ Monotonicity constraints
+
+ """
+
+ self.criterion = criterion
+
+ self.n_samples = 0
+ self.n_features = 0
+
+ self.max_features = max_features
+ self.min_samples_leaf = min_samples_leaf
+ self.min_weight_leaf = min_weight_leaf
+ self.random_state = random_state
+ self.monotonic_cst = monotonic_cst
+ self.with_monotonic_cst = monotonic_cst is not None
+
+ def __getstate__(self):
+ return {}
+
+ def __setstate__(self, d):
+ pass
+
+ def __reduce__(self):
+ return (type(self), (self.criterion,
+ self.max_features,
+ self.min_samples_leaf,
+ self.min_weight_leaf,
+ self.random_state,
+ self.monotonic_cst), self.__getstate__())
+
+ cdef int init(
+ self,
+ object X,
+ const float64_t[:, ::1] y,
+ const float64_t[:] sample_weight,
+ const unsigned char[::1] missing_values_in_feature_mask,
+ ) except -1:
+ """Initialize the splitter.
+
+ Take in the input data X, the target Y, and optional sample weights.
+
+ Returns -1 in case of failure to allocate memory (and raise MemoryError)
+ or 0 otherwise.
+
+ Parameters
+ ----------
+ X : object
+ This contains the inputs. Usually it is a 2d numpy array.
+
+ y : ndarray, dtype=float64_t
+ This is the vector of targets, or true labels, for the samples represented
+ as a Cython memoryview.
+
+ sample_weight : ndarray, dtype=float64_t
+ The weights of the samples, where higher weighted samples are fit
+ closer than lower weight samples. If not provided, all samples
+ are assumed to have uniform weight. This is represented
+ as a Cython memoryview.
+
+ has_missing : bool
+ At least one missing values is in X.
+ """
+
+ self.rand_r_state = self.random_state.randint(0, RAND_R_MAX)
+ cdef intp_t n_samples = X.shape[0]
+
+ # Create a new array which will be used to store nonzero
+ # samples from the feature of interest
+ self.samples = np.empty(n_samples, dtype=np.intp)
+ cdef intp_t[::1] samples = self.samples
+
+ cdef intp_t i, j
+ cdef float64_t weighted_n_samples = 0.0
+ j = 0
+
+ for i in range(n_samples):
+ # Only work with positively weighted samples
+ if sample_weight is None or sample_weight[i] != 0.0:
+ samples[j] = i
+ j += 1
+
+ if sample_weight is not None:
+ weighted_n_samples += sample_weight[i]
+ else:
+ weighted_n_samples += 1.0
+
+ # Number of samples is number of positively weighted samples
+ self.n_samples = j
+ self.weighted_n_samples = weighted_n_samples
+
+ cdef intp_t n_features = X.shape[1]
+ self.features = np.arange(n_features, dtype=np.intp)
+ self.n_features = n_features
+
+ self.feature_values = np.empty(n_samples, dtype=np.float32)
+ self.constant_features = np.empty(n_features, dtype=np.intp)
+
+ self.y = y
+
+ self.sample_weight = sample_weight
+ if missing_values_in_feature_mask is not None:
+ self.criterion.init_sum_missing()
+ return 0
+
+ cdef int node_reset(
+ self,
+ intp_t start,
+ intp_t end,
+ float64_t* weighted_n_node_samples
+ ) except -1 nogil:
+ """Reset splitter on node samples[start:end].
+
+ Returns -1 in case of failure to allocate memory (and raise MemoryError)
+ or 0 otherwise.
+
+ Parameters
+ ----------
+ start : intp_t
+ The index of the first sample to consider
+ end : intp_t
+ The index of the last sample to consider
+ weighted_n_node_samples : ndarray, dtype=float64_t pointer
+ The total weight of those samples
+ """
+
+ self.start = start
+ self.end = end
+
+ self.criterion.init(
+ self.y,
+ self.sample_weight,
+ self.weighted_n_samples,
+ self.samples,
+ start,
+ end
+ )
+
+ weighted_n_node_samples[0] = self.criterion.weighted_n_node_samples
+ return 0
+
+ cdef int node_split(
+ self,
+ ParentInfo* parent_record,
+ SplitRecord* split,
+ ) except -1 nogil:
+
+ """Find the best split on node samples[start:end].
+
+ This is a placeholder method. The majority of computation will be done
+ here.
+
+ It should return -1 upon errors.
+ """
+
+ pass
+
+ cdef void node_value(self, float64_t* dest) noexcept nogil:
+ """Copy the value of node samples[start:end] into dest."""
+
+ self.criterion.node_value(dest)
+
+ cdef inline void clip_node_value(self, float64_t* dest, float64_t lower_bound, float64_t upper_bound) noexcept nogil:
+ """Clip the value in dest between lower_bound and upper_bound for monotonic constraints."""
+
+ self.criterion.clip_node_value(dest, lower_bound, upper_bound)
+
+ cdef float64_t node_impurity(self) noexcept nogil:
+ """Return the impurity of the current node."""
+
+ return self.criterion.node_impurity()
+
+cdef inline void shift_missing_values_to_left_if_required(
+ SplitRecord* best,
+ intp_t[::1] samples,
+ intp_t end,
+) noexcept nogil:
+ """Shift missing value sample indices to the left of the split if required.
+
+ Note: this should always be called at the very end because it will
+ move samples around, thereby affecting the criterion.
+ This affects the computation of the children impurity, which affects
+ the computation of the next node.
+ """
+ cdef intp_t i, p, current_end
+ # The partitioner partitions the data such that the missing values are in
+ # samples[-n_missing:] for the criterion to consume. If the missing values
+ # are going to the right node, then the missing values are already in the
+ # correct position. If the missing values go left, then we move the missing
+ # values to samples[best.pos:best.pos+n_missing] and update `best.pos`.
+ if best.n_missing > 0 and best.missing_go_to_left:
+ for p in range(best.n_missing):
+ i = best.pos + p
+ current_end = end - 1 - p
+ samples[i], samples[current_end] = samples[current_end], samples[i]
+ best.pos += best.n_missing
+
+# Introduce a fused-class to make it possible to share the split implementation
+# between the dense and sparse cases in the node_split_best and node_split_random
+# functions. The alternative would have been to use inheritance-based polymorphism
+# but it would have resulted in a ~10% overall tree fitting performance
+# degradation caused by the overhead frequent virtual method lookups.
+ctypedef fused Partitioner:
+ DensePartitioner
+ SparsePartitioner
+
+cdef inline int node_split_best(
+ Splitter splitter,
+ Partitioner partitioner,
+ Criterion criterion,
+ SplitRecord* split,
+ ParentInfo* parent_record,
+ bint with_monotonic_cst,
+ const int8_t[:] monotonic_cst,
+) except -1 nogil:
+ """Find the best split on node samples[start:end]
+
+ Returns -1 in case of failure to allocate memory (and raise MemoryError)
+ or 0 otherwise.
+ """
+ # Find the best split
+ cdef intp_t start = splitter.start
+ cdef intp_t end = splitter.end
+ cdef intp_t end_non_missing
+ cdef intp_t n_missing = 0
+ cdef bint has_missing = 0
+ cdef intp_t n_searches
+ cdef intp_t n_left, n_right
+ cdef bint missing_go_to_left
+
+ cdef intp_t[::1] samples = splitter.samples
+ cdef intp_t[::1] features = splitter.features
+ cdef intp_t[::1] constant_features = splitter.constant_features
+ cdef intp_t n_features = splitter.n_features
+
+ cdef float32_t[::1] feature_values = splitter.feature_values
+ cdef intp_t max_features = splitter.max_features
+ cdef intp_t min_samples_leaf = splitter.min_samples_leaf
+ cdef float64_t min_weight_leaf = splitter.min_weight_leaf
+ cdef uint32_t* random_state = &splitter.rand_r_state
+
+ cdef SplitRecord best_split, current_split
+ cdef float64_t current_proxy_improvement = -INFINITY
+ cdef float64_t best_proxy_improvement = -INFINITY
+
+ cdef float64_t impurity = parent_record.impurity
+ cdef float64_t lower_bound = parent_record.lower_bound
+ cdef float64_t upper_bound = parent_record.upper_bound
+
+ cdef intp_t f_i = n_features
+ cdef intp_t f_j
+ cdef intp_t p
+ cdef intp_t p_prev
+
+ cdef intp_t n_visited_features = 0
+ # Number of features discovered to be constant during the split search
+ cdef intp_t n_found_constants = 0
+ # Number of features known to be constant and drawn without replacement
+ cdef intp_t n_drawn_constants = 0
+ cdef intp_t n_known_constants = parent_record.n_constant_features
+ # n_total_constants = n_known_constants + n_found_constants
+ cdef intp_t n_total_constants = n_known_constants
+
+ _init_split(&best_split, end)
+
+ partitioner.init_node_split(start, end)
+
+ # Sample up to max_features without replacement using a
+ # Fisher-Yates-based algorithm (using the local variables `f_i` and
+ # `f_j` to compute a permutation of the `features` array).
+ #
+ # Skip the CPU intensive evaluation of the impurity criterion for
+ # features that were already detected as constant (hence not suitable
+ # for good splitting) by ancestor nodes and save the information on
+ # newly discovered constant features to spare computation on descendant
+ # nodes.
+ while (f_i > n_total_constants and # Stop early if remaining features
+ # are constant
+ (n_visited_features < max_features or
+ # At least one drawn features must be non constant
+ n_visited_features <= n_found_constants + n_drawn_constants)):
+
+ n_visited_features += 1
+
+ # Loop invariant: elements of features in
+ # - [:n_drawn_constant[ holds drawn and known constant features;
+ # - [n_drawn_constant:n_known_constant[ holds known constant
+ # features that haven't been drawn yet;
+ # - [n_known_constant:n_total_constant[ holds newly found constant
+ # features;
+ # - [n_total_constant:f_i[ holds features that haven't been drawn
+ # yet and aren't constant apriori.
+ # - [f_i:n_features[ holds features that have been drawn
+ # and aren't constant.
+
+ # Draw a feature at random
+ f_j = rand_int(n_drawn_constants, f_i - n_found_constants,
+ random_state)
+
+ if f_j < n_known_constants:
+ # f_j in the interval [n_drawn_constants, n_known_constants[
+ features[n_drawn_constants], features[f_j] = features[f_j], features[n_drawn_constants]
+
+ n_drawn_constants += 1
+ continue
+
+ # f_j in the interval [n_known_constants, f_i - n_found_constants[
+ f_j += n_found_constants
+ # f_j in the interval [n_total_constants, f_i[
+ current_split.feature = features[f_j]
+ partitioner.sort_samples_and_feature_values(current_split.feature)
+ n_missing = partitioner.n_missing
+ end_non_missing = end - n_missing
+
+ if (
+ # All values for this feature are missing, or
+ end_non_missing == start or
+ # This feature is considered constant (max - min <= FEATURE_THRESHOLD)
+ feature_values[end_non_missing - 1] <= feature_values[start] + FEATURE_THRESHOLD
+ ):
+ # We consider this feature constant in this case.
+ # Since finding a split among constant feature is not valuable,
+ # we do not consider this feature for splitting.
+ features[f_j], features[n_total_constants] = features[n_total_constants], features[f_j]
+
+ n_found_constants += 1
+ n_total_constants += 1
+ continue
+
+ f_i -= 1
+ features[f_i], features[f_j] = features[f_j], features[f_i]
+ has_missing = n_missing != 0
+ criterion.init_missing(n_missing) # initialize even when n_missing == 0
+
+ # Evaluate all splits
+
+ # If there are missing values, then we search twice for the most optimal split.
+ # The first search will have all the missing values going to the right node.
+ # The second search will have all the missing values going to the left node.
+ # If there are no missing values, then we search only once for the most
+ # optimal split.
+ n_searches = 2 if has_missing else 1
+
+ for i in range(n_searches):
+ missing_go_to_left = i == 1
+ criterion.missing_go_to_left = missing_go_to_left
+ criterion.reset()
+
+ p = start
+
+ while p < end_non_missing:
+ partitioner.next_p(&p_prev, &p)
+
+ if p >= end_non_missing:
+ continue
+
+ if missing_go_to_left:
+ n_left = p - start + n_missing
+ n_right = end_non_missing - p
+ else:
+ n_left = p - start
+ n_right = end_non_missing - p + n_missing
+
+ # Reject if min_samples_leaf is not guaranteed
+ if n_left < min_samples_leaf or n_right < min_samples_leaf:
+ continue
+
+ current_split.pos = p
+ criterion.update(current_split.pos)
+
+ # Reject if monotonicity constraints are not satisfied
+ if (
+ with_monotonic_cst and
+ monotonic_cst[current_split.feature] != 0 and
+ not criterion.check_monotonicity(
+ monotonic_cst[current_split.feature],
+ lower_bound,
+ upper_bound,
+ )
+ ):
+ continue
+
+ # Reject if min_weight_leaf is not satisfied
+ if ((criterion.weighted_n_left < min_weight_leaf) or
+ (criterion.weighted_n_right < min_weight_leaf)):
+ continue
+
+ current_proxy_improvement = criterion.proxy_impurity_improvement()
+
+ if current_proxy_improvement > best_proxy_improvement:
+ best_proxy_improvement = current_proxy_improvement
+ # sum of halves is used to avoid infinite value
+ current_split.threshold = (
+ feature_values[p_prev] / 2.0 + feature_values[p] / 2.0
+ )
+
+ if (
+ current_split.threshold == feature_values[p] or
+ current_split.threshold == INFINITY or
+ current_split.threshold == -INFINITY
+ ):
+ current_split.threshold = feature_values[p_prev]
+
+ current_split.n_missing = n_missing
+ if n_missing == 0:
+ current_split.missing_go_to_left = n_left > n_right
+ else:
+ current_split.missing_go_to_left = missing_go_to_left
+
+ best_split = current_split # copy
+
+ # Evaluate when there are missing values and all missing values goes
+ # to the right node and non-missing values goes to the left node.
+ if has_missing:
+ n_left, n_right = end - start - n_missing, n_missing
+ p = end - n_missing
+ missing_go_to_left = 0
+
+ if not (n_left < min_samples_leaf or n_right < min_samples_leaf):
+ criterion.missing_go_to_left = missing_go_to_left
+ criterion.update(p)
+
+ if not ((criterion.weighted_n_left < min_weight_leaf) or
+ (criterion.weighted_n_right < min_weight_leaf)):
+ current_proxy_improvement = criterion.proxy_impurity_improvement()
+
+ if current_proxy_improvement > best_proxy_improvement:
+ best_proxy_improvement = current_proxy_improvement
+ current_split.threshold = INFINITY
+ current_split.missing_go_to_left = missing_go_to_left
+ current_split.n_missing = n_missing
+ current_split.pos = p
+ best_split = current_split
+
+ # Reorganize into samples[start:best_split.pos] + samples[best_split.pos:end]
+ if best_split.pos < end:
+ partitioner.partition_samples_final(
+ best_split.pos,
+ best_split.threshold,
+ best_split.feature,
+ best_split.n_missing
+ )
+ criterion.init_missing(best_split.n_missing)
+ criterion.missing_go_to_left = best_split.missing_go_to_left
+
+ criterion.reset()
+ criterion.update(best_split.pos)
+ criterion.children_impurity(
+ &best_split.impurity_left, &best_split.impurity_right
+ )
+ best_split.improvement = criterion.impurity_improvement(
+ impurity,
+ best_split.impurity_left,
+ best_split.impurity_right
+ )
+
+ shift_missing_values_to_left_if_required(&best_split, samples, end)
+
+ # Respect invariant for constant features: the original order of
+ # element in features[:n_known_constants] must be preserved for sibling
+ # and child nodes
+ memcpy(&features[0], &constant_features[0], sizeof(intp_t) * n_known_constants)
+
+ # Copy newly found constant features
+ memcpy(&constant_features[n_known_constants],
+ &features[n_known_constants],
+ sizeof(intp_t) * n_found_constants)
+
+ # Return values
+ parent_record.n_constant_features = n_total_constants
+ split[0] = best_split
+ return 0
+
+
+# Sort n-element arrays pointed to by feature_values and samples, simultaneously,
+# by the values in feature_values. Algorithm: Introsort (Musser, SP&E, 1997).
+cdef inline void sort(float32_t* feature_values, intp_t* samples, intp_t n) noexcept nogil:
+ if n == 0:
+ return
+ cdef intp_t maxd = 2 * log(n)
+ introsort(feature_values, samples, n, maxd)
+
+
+cdef inline void swap(float32_t* feature_values, intp_t* samples,
+ intp_t i, intp_t j) noexcept nogil:
+ # Helper for sort
+ feature_values[i], feature_values[j] = feature_values[j], feature_values[i]
+ samples[i], samples[j] = samples[j], samples[i]
+
+
+cdef inline float32_t median3(float32_t* feature_values, intp_t n) noexcept nogil:
+ # Median of three pivot selection, after Bentley and McIlroy (1993).
+ # Engineering a sort function. SP&E. Requires 8/3 comparisons on average.
+ cdef float32_t a = feature_values[0], b = feature_values[n / 2], c = feature_values[n - 1]
+ if a < b:
+ if b < c:
+ return b
+ elif a < c:
+ return c
+ else:
+ return a
+ elif b < c:
+ if a < c:
+ return a
+ else:
+ return c
+ else:
+ return b
+
+
+# Introsort with median of 3 pivot selection and 3-way partition function
+# (robust to repeated elements, e.g. lots of zero features).
+cdef void introsort(float32_t* feature_values, intp_t *samples,
+ intp_t n, intp_t maxd) noexcept nogil:
+ cdef float32_t pivot
+ cdef intp_t i, l, r
+
+ while n > 1:
+ if maxd <= 0: # max depth limit exceeded ("gone quadratic")
+ heapsort(feature_values, samples, n)
+ return
+ maxd -= 1
+
+ pivot = median3(feature_values, n)
+
+ # Three-way partition.
+ i = l = 0
+ r = n
+ while i < r:
+ if feature_values[i] < pivot:
+ swap(feature_values, samples, i, l)
+ i += 1
+ l += 1
+ elif feature_values[i] > pivot:
+ r -= 1
+ swap(feature_values, samples, i, r)
+ else:
+ i += 1
+
+ introsort(feature_values, samples, l, maxd)
+ feature_values += r
+ samples += r
+ n -= r
+
+
+cdef inline void sift_down(float32_t* feature_values, intp_t* samples,
+ intp_t start, intp_t end) noexcept nogil:
+ # Restore heap order in feature_values[start:end] by moving the max element to start.
+ cdef intp_t child, maxind, root
+
+ root = start
+ while True:
+ child = root * 2 + 1
+
+ # find max of root, left child, right child
+ maxind = root
+ if child < end and feature_values[maxind] < feature_values[child]:
+ maxind = child
+ if child + 1 < end and feature_values[maxind] < feature_values[child + 1]:
+ maxind = child + 1
+
+ if maxind == root:
+ break
+ else:
+ swap(feature_values, samples, root, maxind)
+ root = maxind
+
+
+cdef void heapsort(float32_t* feature_values, intp_t* samples, intp_t n) noexcept nogil:
+ cdef intp_t start, end
+
+ # heapify
+ start = (n - 2) / 2
+ end = n
+ while True:
+ sift_down(feature_values, samples, start, end)
+ if start == 0:
+ break
+ start -= 1
+
+ # sort by shrinking the heap, putting the max element immediately after it
+ end = n - 1
+ while end > 0:
+ swap(feature_values, samples, 0, end)
+ sift_down(feature_values, samples, 0, end)
+ end = end - 1
+
+cdef inline int node_split_random(
+ Splitter splitter,
+ Partitioner partitioner,
+ Criterion criterion,
+ SplitRecord* split,
+ ParentInfo* parent_record,
+ bint with_monotonic_cst,
+ const int8_t[:] monotonic_cst,
+) except -1 nogil:
+ """Find the best random split on node samples[start:end]
+
+ Returns -1 in case of failure to allocate memory (and raise MemoryError)
+ or 0 otherwise.
+ """
+ # Draw random splits and pick the best
+ cdef intp_t start = splitter.start
+ cdef intp_t end = splitter.end
+
+ cdef intp_t[::1] features = splitter.features
+ cdef intp_t[::1] constant_features = splitter.constant_features
+ cdef intp_t n_features = splitter.n_features
+
+ cdef intp_t max_features = splitter.max_features
+ cdef intp_t min_samples_leaf = splitter.min_samples_leaf
+ cdef float64_t min_weight_leaf = splitter.min_weight_leaf
+ cdef uint32_t* random_state = &splitter.rand_r_state
+
+ cdef SplitRecord best_split, current_split
+ cdef float64_t current_proxy_improvement = - INFINITY
+ cdef float64_t best_proxy_improvement = - INFINITY
+
+ cdef float64_t impurity = parent_record.impurity
+ cdef float64_t lower_bound = parent_record.lower_bound
+ cdef float64_t upper_bound = parent_record.upper_bound
+
+ cdef intp_t f_i = n_features
+ cdef intp_t f_j
+ # Number of features discovered to be constant during the split search
+ cdef intp_t n_found_constants = 0
+ # Number of features known to be constant and drawn without replacement
+ cdef intp_t n_drawn_constants = 0
+ cdef intp_t n_known_constants = parent_record.n_constant_features
+ # n_total_constants = n_known_constants + n_found_constants
+ cdef intp_t n_total_constants = n_known_constants
+ cdef intp_t n_visited_features = 0
+ cdef float32_t min_feature_value
+ cdef float32_t max_feature_value
+
+ _init_split(&best_split, end)
+
+ partitioner.init_node_split(start, end)
+
+ # Sample up to max_features without replacement using a
+ # Fisher-Yates-based algorithm (using the local variables `f_i` and
+ # `f_j` to compute a permutation of the `features` array).
+ #
+ # Skip the CPU intensive evaluation of the impurity criterion for
+ # features that were already detected as constant (hence not suitable
+ # for good splitting) by ancestor nodes and save the information on
+ # newly discovered constant features to spare computation on descendant
+ # nodes.
+ while (f_i > n_total_constants and # Stop early if remaining features
+ # are constant
+ (n_visited_features < max_features or
+ # At least one drawn features must be non constant
+ n_visited_features <= n_found_constants + n_drawn_constants)):
+ n_visited_features += 1
+
+ # Loop invariant: elements of features in
+ # - [:n_drawn_constant[ holds drawn and known constant features;
+ # - [n_drawn_constant:n_known_constant[ holds known constant
+ # features that haven't been drawn yet;
+ # - [n_known_constant:n_total_constant[ holds newly found constant
+ # features;
+ # - [n_total_constant:f_i[ holds features that haven't been drawn
+ # yet and aren't constant apriori.
+ # - [f_i:n_features[ holds features that have been drawn
+ # and aren't constant.
+
+ # Draw a feature at random
+ f_j = rand_int(n_drawn_constants, f_i - n_found_constants,
+ random_state)
+
+ if f_j < n_known_constants:
+ # f_j in the interval [n_drawn_constants, n_known_constants[
+ features[n_drawn_constants], features[f_j] = features[f_j], features[n_drawn_constants]
+ n_drawn_constants += 1
+ continue
+
+ # f_j in the interval [n_known_constants, f_i - n_found_constants[
+ f_j += n_found_constants
+ # f_j in the interval [n_total_constants, f_i[
+
+ current_split.feature = features[f_j]
+
+ # Find min, max
+ partitioner.find_min_max(
+ current_split.feature, &min_feature_value, &max_feature_value
+ )
+
+ if max_feature_value <= min_feature_value + FEATURE_THRESHOLD:
+ features[f_j], features[n_total_constants] = features[n_total_constants], current_split.feature
+
+ n_found_constants += 1
+ n_total_constants += 1
+ continue
+
+ f_i -= 1
+ features[f_i], features[f_j] = features[f_j], features[f_i]
+
+ # Draw a random threshold
+ current_split.threshold = rand_uniform(
+ min_feature_value,
+ max_feature_value,
+ random_state,
+ )
+
+ if current_split.threshold == max_feature_value:
+ current_split.threshold = min_feature_value
+
+ # Partition
+ current_split.pos = partitioner.partition_samples(current_split.threshold)
+
+ # Reject if min_samples_leaf is not guaranteed
+ if (((current_split.pos - start) < min_samples_leaf) or
+ ((end - current_split.pos) < min_samples_leaf)):
+ continue
+
+ # Evaluate split
+ # At this point, the criterion has a view into the samples that was partitioned
+ # by the partitioner. The criterion will use the partition to evaluating the split.
+ criterion.reset()
+ criterion.update(current_split.pos)
+
+ # Reject if min_weight_leaf is not satisfied
+ if ((criterion.weighted_n_left < min_weight_leaf) or
+ (criterion.weighted_n_right < min_weight_leaf)):
+ continue
+
+ # Reject if monotonicity constraints are not satisfied
+ if (
+ with_monotonic_cst and
+ monotonic_cst[current_split.feature] != 0 and
+ not criterion.check_monotonicity(
+ monotonic_cst[current_split.feature],
+ lower_bound,
+ upper_bound,
+ )
+ ):
+ continue
+
+ current_proxy_improvement = criterion.proxy_impurity_improvement()
+
+ if current_proxy_improvement > best_proxy_improvement:
+ best_proxy_improvement = current_proxy_improvement
+ best_split = current_split # copy
+
+ # Reorganize into samples[start:best.pos] + samples[best.pos:end]
+ if best_split.pos < end:
+ if current_split.feature != best_split.feature:
+ # TODO: Pass in best.n_missing when random splitter supports missing values.
+ partitioner.partition_samples_final(
+ best_split.pos, best_split.threshold, best_split.feature, 0
+ )
+
+ criterion.reset()
+ criterion.update(best_split.pos)
+ criterion.children_impurity(
+ &best_split.impurity_left, &best_split.impurity_right
+ )
+ best_split.improvement = criterion.impurity_improvement(
+ impurity, best_split.impurity_left, best_split.impurity_right
+ )
+
+ # Respect invariant for constant features: the original order of
+ # element in features[:n_known_constants] must be preserved for sibling
+ # and child nodes
+ memcpy(&features[0], &constant_features[0], sizeof(intp_t) * n_known_constants)
+
+ # Copy newly found constant features
+ memcpy(&constant_features[n_known_constants],
+ &features[n_known_constants],
+ sizeof(intp_t) * n_found_constants)
+
+ # Return values
+ parent_record.n_constant_features = n_total_constants
+ split[0] = best_split
+ return 0
+
+
+@final
+cdef class DensePartitioner:
+ """Partitioner specialized for dense data.
+
+ Note that this partitioner is agnostic to the splitting strategy (best vs. random).
+ """
+ cdef:
+ const float32_t[:, :] X
+ cdef intp_t[::1] samples
+ cdef float32_t[::1] feature_values
+ cdef intp_t start
+ cdef intp_t end
+ cdef intp_t n_missing
+ cdef const unsigned char[::1] missing_values_in_feature_mask
+
+ def __init__(
+ self,
+ const float32_t[:, :] X,
+ intp_t[::1] samples,
+ float32_t[::1] feature_values,
+ const unsigned char[::1] missing_values_in_feature_mask,
+ ):
+ self.X = X
+ self.samples = samples
+ self.feature_values = feature_values
+ self.missing_values_in_feature_mask = missing_values_in_feature_mask
+
+ cdef inline void init_node_split(self, intp_t start, intp_t end) noexcept nogil:
+ """Initialize splitter at the beginning of node_split."""
+ self.start = start
+ self.end = end
+ self.n_missing = 0
+
+ cdef inline void sort_samples_and_feature_values(
+ self, intp_t current_feature
+ ) noexcept nogil:
+ """Simultaneously sort based on the feature_values.
+
+ Missing values are stored at the end of feature_values.
+ The number of missing values observed in feature_values is stored
+ in self.n_missing.
+ """
+ cdef:
+ intp_t i, current_end
+ float32_t[::1] feature_values = self.feature_values
+ const float32_t[:, :] X = self.X
+ intp_t[::1] samples = self.samples
+ intp_t n_missing = 0
+ const unsigned char[::1] missing_values_in_feature_mask = self.missing_values_in_feature_mask
+
+ # Sort samples along that feature; by
+ # copying the values into an array and
+ # sorting the array in a manner which utilizes the cache more
+ # effectively.
+ if missing_values_in_feature_mask is not None and missing_values_in_feature_mask[current_feature]:
+ i, current_end = self.start, self.end - 1
+ # Missing values are placed at the end and do not participate in the sorting.
+ while i <= current_end:
+ # Finds the right-most value that is not missing so that
+ # it can be swapped with missing values at its left.
+ if isnan(X[samples[current_end], current_feature]):
+ n_missing += 1
+ current_end -= 1
+ continue
+
+ # X[samples[current_end], current_feature] is a non-missing value
+ if isnan(X[samples[i], current_feature]):
+ samples[i], samples[current_end] = samples[current_end], samples[i]
+ n_missing += 1
+ current_end -= 1
+
+ feature_values[i] = X[samples[i], current_feature]
+ i += 1
+ else:
+ # When there are no missing values, we only need to copy the data into
+ # feature_values
+ for i in range(self.start, self.end):
+ feature_values[i] = X[samples[i], current_feature]
+
+ sort(&feature_values[self.start], &samples[self.start], self.end - self.start - n_missing)
+ self.n_missing = n_missing
+
+ cdef inline void find_min_max(
+ self,
+ intp_t current_feature,
+ float32_t* min_feature_value_out,
+ float32_t* max_feature_value_out,
+ ) noexcept nogil:
+ """Find the minimum and maximum value for current_feature."""
+ cdef:
+ intp_t p
+ float32_t current_feature_value
+ const float32_t[:, :] X = self.X
+ intp_t[::1] samples = self.samples
+ float32_t min_feature_value = X[samples[self.start], current_feature]
+ float32_t max_feature_value = min_feature_value
+ float32_t[::1] feature_values = self.feature_values
+
+ feature_values[self.start] = min_feature_value
+
+ for p in range(self.start + 1, self.end):
+ current_feature_value = X[samples[p], current_feature]
+ feature_values[p] = current_feature_value
+
+ if current_feature_value < min_feature_value:
+ min_feature_value = current_feature_value
+ elif current_feature_value > max_feature_value:
+ max_feature_value = current_feature_value
+
+ min_feature_value_out[0] = min_feature_value
+ max_feature_value_out[0] = max_feature_value
+
+ cdef inline void next_p(self, intp_t* p_prev, intp_t* p) noexcept nogil:
+ """Compute the next p_prev and p for iteratiing over feature values.
+
+ The missing values are not included when iterating through the feature values.
+ """
+ cdef:
+ float32_t[::1] feature_values = self.feature_values
+ intp_t end_non_missing = self.end - self.n_missing
+
+ while (
+ p[0] + 1 < end_non_missing and
+ feature_values[p[0] + 1] <= feature_values[p[0]] + FEATURE_THRESHOLD
+ ):
+ p[0] += 1
+
+ p_prev[0] = p[0]
+
+ # By adding 1, we have
+ # (feature_values[p] >= end) or (feature_values[p] > feature_values[p - 1])
+ p[0] += 1
+
+ cdef inline intp_t partition_samples(self, float64_t current_threshold) noexcept nogil:
+ """Partition samples for feature_values at the current_threshold."""
+ cdef:
+ intp_t p = self.start
+ intp_t partition_end = self.end
+ intp_t[::1] samples = self.samples
+ float32_t[::1] feature_values = self.feature_values
+
+ while p < partition_end:
+ if feature_values[p] <= current_threshold:
+ p += 1
+ else:
+ partition_end -= 1
+
+ feature_values[p], feature_values[partition_end] = (
+ feature_values[partition_end], feature_values[p]
+ )
+ samples[p], samples[partition_end] = samples[partition_end], samples[p]
+
+ return partition_end
+
+ cdef inline void partition_samples_final(
+ self,
+ intp_t best_pos,
+ float64_t best_threshold,
+ intp_t best_feature,
+ intp_t best_n_missing,
+ ) noexcept nogil:
+ """Partition samples for X at the best_threshold and best_feature.
+
+ If missing values are present, this method partitions `samples`
+ so that the `best_n_missing` missing values' indices are in the
+ right-most end of `samples`, that is `samples[end_non_missing:end]`.
+ """
+ cdef:
+ # Local invariance: start <= p <= partition_end <= end
+ intp_t start = self.start
+ intp_t p = start
+ intp_t end = self.end - 1
+ intp_t partition_end = end - best_n_missing
+ intp_t[::1] samples = self.samples
+ const float32_t[:, :] X = self.X
+ float32_t current_value
+
+ if best_n_missing != 0:
+ # Move samples with missing values to the end while partitioning the
+ # non-missing samples
+ while p < partition_end:
+ # Keep samples with missing values at the end
+ if isnan(X[samples[end], best_feature]):
+ end -= 1
+ continue
+
+ # Swap sample with missing values with the sample at the end
+ current_value = X[samples[p], best_feature]
+ if isnan(current_value):
+ samples[p], samples[end] = samples[end], samples[p]
+ end -= 1
+
+ # The swapped sample at the end is always a non-missing value, so
+ # we can continue the algorithm without checking for missingness.
+ current_value = X[samples[p], best_feature]
+
+ # Partition the non-missing samples
+ if current_value <= best_threshold:
+ p += 1
+ else:
+ samples[p], samples[partition_end] = samples[partition_end], samples[p]
+ partition_end -= 1
+ else:
+ # Partitioning routine when there are no missing values
+ while p < partition_end:
+ if X[samples[p], best_feature] <= best_threshold:
+ p += 1
+ else:
+ samples[p], samples[partition_end] = samples[partition_end], samples[p]
+ partition_end -= 1
+
+
+@final
+cdef class SparsePartitioner:
+ """Partitioner specialized for sparse CSC data.
+
+ Note that this partitioner is agnostic to the splitting strategy (best vs. random).
+ """
+ cdef intp_t[::1] samples
+ cdef float32_t[::1] feature_values
+ cdef intp_t start
+ cdef intp_t end
+ cdef intp_t n_missing
+ cdef const unsigned char[::1] missing_values_in_feature_mask
+
+ cdef const float32_t[::1] X_data
+ cdef const int32_t[::1] X_indices
+ cdef const int32_t[::1] X_indptr
+
+ cdef intp_t n_total_samples
+
+ cdef intp_t[::1] index_to_samples
+ cdef intp_t[::1] sorted_samples
+
+ cdef intp_t start_positive
+ cdef intp_t end_negative
+ cdef bint is_samples_sorted
+
+ def __init__(
+ self,
+ object X,
+ intp_t[::1] samples,
+ intp_t n_samples,
+ float32_t[::1] feature_values,
+ const unsigned char[::1] missing_values_in_feature_mask,
+ ):
+ if not (issparse(X) and X.format == "csc"):
+ raise ValueError("X should be in csc format")
+
+ self.samples = samples
+ self.feature_values = feature_values
+
+ # Initialize X
+ cdef intp_t n_total_samples = X.shape[0]
+
+ self.X_data = X.data
+ self.X_indices = X.indices
+ self.X_indptr = X.indptr
+ self.n_total_samples = n_total_samples
+
+ # Initialize auxiliary array used to perform split
+ self.index_to_samples = np.full(n_total_samples, fill_value=-1, dtype=np.intp)
+ self.sorted_samples = np.empty(n_samples, dtype=np.intp)
+
+ cdef intp_t p
+ for p in range(n_samples):
+ self.index_to_samples[samples[p]] = p
+
+ self.missing_values_in_feature_mask = missing_values_in_feature_mask
+
+ cdef inline void init_node_split(self, intp_t start, intp_t end) noexcept nogil:
+ """Initialize splitter at the beginning of node_split."""
+ self.start = start
+ self.end = end
+ self.is_samples_sorted = 0
+ self.n_missing = 0
+
+ cdef inline void sort_samples_and_feature_values(
+ self, intp_t current_feature
+ ) noexcept nogil:
+ """Simultaneously sort based on the feature_values."""
+ cdef:
+ float32_t[::1] feature_values = self.feature_values
+ intp_t[::1] index_to_samples = self.index_to_samples
+ intp_t[::1] samples = self.samples
+
+ self.extract_nnz(current_feature)
+ # Sort the positive and negative parts of `feature_values`
+ sort(&feature_values[self.start], &samples[self.start], self.end_negative - self.start)
+ if self.start_positive < self.end:
+ sort(
+ &feature_values[self.start_positive],
+ &samples[self.start_positive],
+ self.end - self.start_positive
+ )
+
+ # Update index_to_samples to take into account the sort
+ for p in range(self.start, self.end_negative):
+ index_to_samples[samples[p]] = p
+ for p in range(self.start_positive, self.end):
+ index_to_samples[samples[p]] = p
+
+ # Add one or two zeros in feature_values, if there is any
+ if self.end_negative < self.start_positive:
+ self.start_positive -= 1
+ feature_values[self.start_positive] = 0.
+
+ if self.end_negative != self.start_positive:
+ feature_values[self.end_negative] = 0.
+ self.end_negative += 1
+
+ # XXX: When sparse supports missing values, this should be set to the
+ # number of missing values for current_feature
+ self.n_missing = 0
+
+ cdef inline void find_min_max(
+ self,
+ intp_t current_feature,
+ float32_t* min_feature_value_out,
+ float32_t* max_feature_value_out,
+ ) noexcept nogil:
+ """Find the minimum and maximum value for current_feature."""
+ cdef:
+ intp_t p
+ float32_t current_feature_value, min_feature_value, max_feature_value
+ float32_t[::1] feature_values = self.feature_values
+
+ self.extract_nnz(current_feature)
+
+ if self.end_negative != self.start_positive:
+ # There is a zero
+ min_feature_value = 0
+ max_feature_value = 0
+ else:
+ min_feature_value = feature_values[self.start]
+ max_feature_value = min_feature_value
+
+ # Find min, max in feature_values[start:end_negative]
+ for p in range(self.start, self.end_negative):
+ current_feature_value = feature_values[p]
+
+ if current_feature_value < min_feature_value:
+ min_feature_value = current_feature_value
+ elif current_feature_value > max_feature_value:
+ max_feature_value = current_feature_value
+
+ # Update min, max given feature_values[start_positive:end]
+ for p in range(self.start_positive, self.end):
+ current_feature_value = feature_values[p]
+
+ if current_feature_value < min_feature_value:
+ min_feature_value = current_feature_value
+ elif current_feature_value > max_feature_value:
+ max_feature_value = current_feature_value
+
+ min_feature_value_out[0] = min_feature_value
+ max_feature_value_out[0] = max_feature_value
+
+ cdef inline void next_p(self, intp_t* p_prev, intp_t* p) noexcept nogil:
+ """Compute the next p_prev and p for iteratiing over feature values."""
+ cdef:
+ intp_t p_next
+ float32_t[::1] feature_values = self.feature_values
+
+ if p[0] + 1 != self.end_negative:
+ p_next = p[0] + 1
+ else:
+ p_next = self.start_positive
+
+ while (p_next < self.end and
+ feature_values[p_next] <= feature_values[p[0]] + FEATURE_THRESHOLD):
+ p[0] = p_next
+ if p[0] + 1 != self.end_negative:
+ p_next = p[0] + 1
+ else:
+ p_next = self.start_positive
+
+ p_prev[0] = p[0]
+ p[0] = p_next
+
+ cdef inline intp_t partition_samples(self, float64_t current_threshold) noexcept nogil:
+ """Partition samples for feature_values at the current_threshold."""
+ return self._partition(current_threshold, self.start_positive)
+
+ cdef inline void partition_samples_final(
+ self,
+ intp_t best_pos,
+ float64_t best_threshold,
+ intp_t best_feature,
+ intp_t n_missing,
+ ) noexcept nogil:
+ """Partition samples for X at the best_threshold and best_feature."""
+ self.extract_nnz(best_feature)
+ self._partition(best_threshold, best_pos)
+
+ cdef inline intp_t _partition(self, float64_t threshold, intp_t zero_pos) noexcept nogil:
+ """Partition samples[start:end] based on threshold."""
+ cdef:
+ intp_t p, partition_end
+ intp_t[::1] index_to_samples = self.index_to_samples
+ float32_t[::1] feature_values = self.feature_values
+ intp_t[::1] samples = self.samples
+
+ if threshold < 0.:
+ p = self.start
+ partition_end = self.end_negative
+ elif threshold > 0.:
+ p = self.start_positive
+ partition_end = self.end
+ else:
+ # Data are already split
+ return zero_pos
+
+ while p < partition_end:
+ if feature_values[p] <= threshold:
+ p += 1
+
+ else:
+ partition_end -= 1
+
+ feature_values[p], feature_values[partition_end] = (
+ feature_values[partition_end], feature_values[p]
+ )
+ sparse_swap(index_to_samples, samples, p, partition_end)
+
+ return partition_end
+
+ cdef inline void extract_nnz(self, intp_t feature) noexcept nogil:
+ """Extract and partition values for a given feature.
+
+ The extracted values are partitioned between negative values
+ feature_values[start:end_negative[0]] and positive values
+ feature_values[start_positive[0]:end].
+ The samples and index_to_samples are modified according to this
+ partition.
+
+ The extraction corresponds to the intersection between the arrays
+ X_indices[indptr_start:indptr_end] and samples[start:end].
+ This is done efficiently using either an index_to_samples based approach
+ or binary search based approach.
+
+ Parameters
+ ----------
+ feature : intp_t,
+ Index of the feature we want to extract non zero value.
+ """
+ cdef intp_t[::1] samples = self.samples
+ cdef float32_t[::1] feature_values = self.feature_values
+ cdef intp_t indptr_start = self.X_indptr[feature],
+ cdef intp_t indptr_end = self.X_indptr[feature + 1]
+ cdef intp_t n_indices = (indptr_end - indptr_start)
+ cdef intp_t n_samples = self.end - self.start
+ cdef intp_t[::1] index_to_samples = self.index_to_samples
+ cdef intp_t[::1] sorted_samples = self.sorted_samples
+ cdef const int32_t[::1] X_indices = self.X_indices
+ cdef const float32_t[::1] X_data = self.X_data
+
+ # Use binary search if n_samples * log(n_indices) <
+ # n_indices and index_to_samples approach otherwise.
+ # O(n_samples * log(n_indices)) is the running time of binary
+ # search and O(n_indices) is the running time of index_to_samples
+ # approach.
+ if ((1 - self.is_samples_sorted) * n_samples * log(n_samples) +
+ n_samples * log(n_indices) < EXTRACT_NNZ_SWITCH * n_indices):
+ extract_nnz_binary_search(X_indices, X_data,
+ indptr_start, indptr_end,
+ samples, self.start, self.end,
+ index_to_samples,
+ feature_values,
+ &self.end_negative, &self.start_positive,
+ sorted_samples, &self.is_samples_sorted)
+
+ # Using an index to samples technique to extract non zero values
+ # index_to_samples is a mapping from X_indices to samples
+ else:
+ extract_nnz_index_to_samples(X_indices, X_data,
+ indptr_start, indptr_end,
+ samples, self.start, self.end,
+ index_to_samples,
+ feature_values,
+ &self.end_negative, &self.start_positive)
+
+
+cdef int compare_intp_t(const void* a, const void* b) noexcept nogil:
+ """Comparison function for sort.
+
+ This must return an `int` as it is used by stdlib's qsort, which expects
+ an `int` return value.
+ """
+ return ((a)[0] - (b)[0])
+
+
+cdef inline void binary_search(const int32_t[::1] sorted_array,
+ int32_t start, int32_t end,
+ intp_t value, intp_t* index,
+ int32_t* new_start) noexcept nogil:
+ """Return the index of value in the sorted array.
+
+ If not found, return -1. new_start is the last pivot + 1
+ """
+ cdef int32_t pivot
+ index[0] = -1
+ while start < end:
+ pivot = start + (end - start) / 2
+
+ if sorted_array[pivot] == value:
+ index[0] = pivot
+ start = pivot + 1
+ break
+
+ if sorted_array[pivot] < value:
+ start = pivot + 1
+ else:
+ end = pivot
+ new_start[0] = start
+
+
+cdef inline void extract_nnz_index_to_samples(const int32_t[::1] X_indices,
+ const float32_t[::1] X_data,
+ int32_t indptr_start,
+ int32_t indptr_end,
+ intp_t[::1] samples,
+ intp_t start,
+ intp_t end,
+ intp_t[::1] index_to_samples,
+ float32_t[::1] feature_values,
+ intp_t* end_negative,
+ intp_t* start_positive) noexcept nogil:
+ """Extract and partition values for a feature using index_to_samples.
+
+ Complexity is O(indptr_end - indptr_start).
+ """
+ cdef int32_t k
+ cdef intp_t index
+ cdef intp_t end_negative_ = start
+ cdef intp_t start_positive_ = end
+
+ for k in range(indptr_start, indptr_end):
+ if start <= index_to_samples[X_indices[k]] < end:
+ if X_data[k] > 0:
+ start_positive_ -= 1
+ feature_values[start_positive_] = X_data[k]
+ index = index_to_samples[X_indices[k]]
+ sparse_swap(index_to_samples, samples, index, start_positive_)
+
+ elif X_data[k] < 0:
+ feature_values[end_negative_] = X_data[k]
+ index = index_to_samples[X_indices[k]]
+ sparse_swap(index_to_samples, samples, index, end_negative_)
+ end_negative_ += 1
+
+ # Returned values
+ end_negative[0] = end_negative_
+ start_positive[0] = start_positive_
+
+
+cdef inline void extract_nnz_binary_search(const int32_t[::1] X_indices,
+ const float32_t[::1] X_data,
+ int32_t indptr_start,
+ int32_t indptr_end,
+ intp_t[::1] samples,
+ intp_t start,
+ intp_t end,
+ intp_t[::1] index_to_samples,
+ float32_t[::1] feature_values,
+ intp_t* end_negative,
+ intp_t* start_positive,
+ intp_t[::1] sorted_samples,
+ bint* is_samples_sorted) noexcept nogil:
+ """Extract and partition values for a given feature using binary search.
+
+ If n_samples = end - start and n_indices = indptr_end - indptr_start,
+ the complexity is
+
+ O((1 - is_samples_sorted[0]) * n_samples * log(n_samples) +
+ n_samples * log(n_indices)).
+ """
+ cdef intp_t n_samples
+
+ if not is_samples_sorted[0]:
+ n_samples = end - start
+ memcpy(&sorted_samples[start], &samples[start],
+ n_samples * sizeof(intp_t))
+ qsort(&sorted_samples[start], n_samples, sizeof(intp_t),
+ compare_intp_t)
+ is_samples_sorted[0] = 1
+
+ while (indptr_start < indptr_end and
+ sorted_samples[start] > X_indices[indptr_start]):
+ indptr_start += 1
+
+ while (indptr_start < indptr_end and
+ sorted_samples[end - 1] < X_indices[indptr_end - 1]):
+ indptr_end -= 1
+
+ cdef intp_t p = start
+ cdef intp_t index
+ cdef intp_t k
+ cdef intp_t end_negative_ = start
+ cdef intp_t start_positive_ = end
+
+ while (p < end and indptr_start < indptr_end):
+ # Find index of sorted_samples[p] in X_indices
+ binary_search(X_indices, indptr_start, indptr_end,
+ sorted_samples[p], &k, &indptr_start)
+
+ if k != -1:
+ # If k != -1, we have found a non zero value
+
+ if X_data[k] > 0:
+ start_positive_ -= 1
+ feature_values[start_positive_] = X_data[k]
+ index = index_to_samples[X_indices[k]]
+ sparse_swap(index_to_samples, samples, index, start_positive_)
+
+ elif X_data[k] < 0:
+ feature_values[end_negative_] = X_data[k]
+ index = index_to_samples[X_indices[k]]
+ sparse_swap(index_to_samples, samples, index, end_negative_)
+ end_negative_ += 1
+ p += 1
+
+ # Returned values
+ end_negative[0] = end_negative_
+ start_positive[0] = start_positive_
+
+
+cdef inline void sparse_swap(intp_t[::1] index_to_samples, intp_t[::1] samples,
+ intp_t pos_1, intp_t pos_2) noexcept nogil:
+ """Swap sample pos_1 and pos_2 preserving sparse invariant."""
+ samples[pos_1], samples[pos_2] = samples[pos_2], samples[pos_1]
+ index_to_samples[samples[pos_1]] = pos_1
+ index_to_samples[samples[pos_2]] = pos_2
+
+
+cdef class BestSplitter(Splitter):
+ """Splitter for finding the best split on dense data."""
+ cdef DensePartitioner partitioner
+ cdef int init(
+ self,
+ object X,
+ const float64_t[:, ::1] y,
+ const float64_t[:] sample_weight,
+ const unsigned char[::1] missing_values_in_feature_mask,
+ ) except -1:
+ Splitter.init(self, X, y, sample_weight, missing_values_in_feature_mask)
+ self.partitioner = DensePartitioner(
+ X, self.samples, self.feature_values, missing_values_in_feature_mask
+ )
+
+ cdef int node_split(
+ self,
+ ParentInfo* parent_record,
+ SplitRecord* split,
+ ) except -1 nogil:
+ return node_split_best(
+ self,
+ self.partitioner,
+ self.criterion,
+ split,
+ parent_record,
+ self.with_monotonic_cst,
+ self.monotonic_cst,
+ )
+
+cdef class BestSparseSplitter(Splitter):
+ """Splitter for finding the best split, using the sparse data."""
+ cdef SparsePartitioner partitioner
+ cdef int init(
+ self,
+ object X,
+ const float64_t[:, ::1] y,
+ const float64_t[:] sample_weight,
+ const unsigned char[::1] missing_values_in_feature_mask,
+ ) except -1:
+ Splitter.init(self, X, y, sample_weight, missing_values_in_feature_mask)
+ self.partitioner = SparsePartitioner(
+ X, self.samples, self.n_samples, self.feature_values, missing_values_in_feature_mask
+ )
+
+ cdef int node_split(
+ self,
+ ParentInfo* parent_record,
+ SplitRecord* split,
+ ) except -1 nogil:
+ return node_split_best(
+ self,
+ self.partitioner,
+ self.criterion,
+ split,
+ parent_record,
+ self.with_monotonic_cst,
+ self.monotonic_cst,
+ )
+
+cdef class RandomSplitter(Splitter):
+ """Splitter for finding the best random split on dense data."""
+ cdef DensePartitioner partitioner
+ cdef int init(
+ self,
+ object X,
+ const float64_t[:, ::1] y,
+ const float64_t[:] sample_weight,
+ const unsigned char[::1] missing_values_in_feature_mask,
+ ) except -1:
+ Splitter.init(self, X, y, sample_weight, missing_values_in_feature_mask)
+ self.partitioner = DensePartitioner(
+ X, self.samples, self.feature_values, missing_values_in_feature_mask
+ )
+
+ cdef int node_split(
+ self,
+ ParentInfo* parent_record,
+ SplitRecord* split,
+ ) except -1 nogil:
+ return node_split_random(
+ self,
+ self.partitioner,
+ self.criterion,
+ split,
+ parent_record,
+ self.with_monotonic_cst,
+ self.monotonic_cst,
+ )
+
+cdef class RandomSparseSplitter(Splitter):
+ """Splitter for finding the best random split, using the sparse data."""
+ cdef SparsePartitioner partitioner
+ cdef int init(
+ self,
+ object X,
+ const float64_t[:, ::1] y,
+ const float64_t[:] sample_weight,
+ const unsigned char[::1] missing_values_in_feature_mask,
+ ) except -1:
+ Splitter.init(self, X, y, sample_weight, missing_values_in_feature_mask)
+ self.partitioner = SparsePartitioner(
+ X, self.samples, self.n_samples, self.feature_values, missing_values_in_feature_mask
+ )
+ cdef int node_split(
+ self,
+ ParentInfo* parent_record,
+ SplitRecord* split,
+ ) except -1 nogil:
+ return node_split_random(
+ self,
+ self.partitioner,
+ self.criterion,
+ split,
+ parent_record,
+ self.with_monotonic_cst,
+ self.monotonic_cst,
+ )
diff --git a/causalml/source/causalml/inference/tree/_tree/_tree.pxd b/causalml/source/causalml/inference/tree/_tree/_tree.pxd
new file mode 100644
index 0000000000000000000000000000000000000000..0a06ff72ed21ccfaddfe7bdb52159a7063d16a44
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/_tree/_tree.pxd
@@ -0,0 +1,161 @@
+# Authors: Gilles Louppe
+# Peter Prettenhofer
+# Brian Holt
+# Joel Nothman
+# Arnaud Joly
+# Jacob Schreiber
+# Nelson Liu
+#
+# License: BSD 3 clause
+
+# distutils: language = c++
+# cython: cdivision=True
+# cython: boundscheck=False
+# cython: wraparound=False
+# cython: language_level=3
+# cython: linetrace=True
+
+# See _tree.pyx for details.
+
+import numpy as np
+cimport numpy as cnp
+
+from ._typedefs cimport float32_t, float64_t, intp_t, int32_t, uint32_t
+
+from ._splitter cimport Splitter
+from ._splitter cimport SplitRecord
+
+cdef struct Node:
+ # Base storage structure for the nodes in a Tree object
+
+ intp_t left_child # id of the left child of the node
+ intp_t right_child # id of the right child of the node
+ intp_t feature # Feature used for splitting the node
+ float64_t threshold # Threshold value at the node
+ float64_t impurity # Impurity of the node (i.e., the value of the criterion)
+ intp_t n_node_samples # Number of samples at the node
+ float64_t weighted_n_node_samples # Weighted number of samples at the node
+ unsigned char missing_go_to_left # Whether features have missing values
+
+cdef void _init_parent_record(ParentInfo* record) noexcept nogil
+
+cdef struct ParentInfo:
+ # Structure to store information about the parent of a node
+ # This is passed to the splitter, to provide information about the previous split
+
+ float64_t lower_bound # the lower bound of the parent's impurity
+ float64_t upper_bound # the upper bound of the parent's impurity
+ float64_t impurity # the impurity of the parent
+ intp_t n_constant_features # the number of constant features found in parent
+
+cdef class Tree:
+ # The Tree object is a binary tree structure constructed by the
+ # TreeBuilder. The tree structure is used for predictions and
+ # feature importances.
+
+ # Input/Output layout
+ cdef public intp_t n_features # Number of features in X
+ cdef intp_t* n_classes # Number of classes in y[:, k]
+ cdef public intp_t n_outputs # Number of outputs in y
+ cdef public intp_t max_n_classes # max(n_classes)
+
+ # Inner structures: values are stored separately from node structure,
+ # since size is determined at runtime.
+ cdef public intp_t max_depth # Max depth of the tree
+ cdef public intp_t node_count # Counter for node IDs
+ cdef public intp_t capacity # Capacity of tree, in terms of nodes
+ cdef Node* nodes # Array of nodes
+ cdef float64_t* value # (capacity, n_outputs, max_n_classes) array of values
+ cdef intp_t value_stride # = n_outputs * max_n_classes
+
+ # Methods
+ cdef intp_t _add_node(self, intp_t parent, bint is_left, bint is_leaf,
+ intp_t feature, float64_t threshold, float64_t impurity,
+ intp_t n_node_samples,
+ float64_t weighted_n_node_samples,
+ unsigned char missing_go_to_left) except -1 nogil
+ cdef int _resize(self, intp_t capacity) except -1 nogil
+ cdef int _resize_c(self, intp_t capacity=*) except -1 nogil
+
+ cdef cnp.ndarray _get_value_ndarray(self)
+ cdef cnp.ndarray _get_node_ndarray(self)
+
+ cpdef cnp.ndarray predict(self, object X)
+
+ cpdef cnp.ndarray apply(self, object X)
+ cdef cnp.ndarray _apply_dense(self, object X)
+ cdef cnp.ndarray _apply_sparse_csr(self, object X)
+
+ cpdef object decision_path(self, object X)
+ cdef object _decision_path_dense(self, object X)
+ cdef object _decision_path_sparse_csr(self, object X)
+
+ cpdef compute_node_depths(self)
+ cpdef compute_feature_importances(self, normalize=*)
+
+
+# =============================================================================
+# Tree builder
+# =============================================================================
+
+cdef class TreeBuilder:
+ # The TreeBuilder recursively builds a Tree object from training samples,
+ # using a Splitter object for splitting internal nodes and assigning
+ # values to leaves.
+ #
+ # This class controls the various stopping criteria and the node splitting
+ # evaluation order, e.g. depth-first or best-first.
+
+ cdef Splitter splitter # Splitting algorithm
+
+ cdef intp_t min_samples_split # Minimum number of samples in an internal node
+ cdef intp_t min_samples_leaf # Minimum number of samples in a leaf
+ cdef float64_t min_weight_leaf # Minimum weight in a leaf
+ cdef intp_t max_depth # Maximal tree depth
+ cdef float64_t min_impurity_decrease # Impurity threshold for early stopping
+
+ cpdef build(
+ self,
+ Tree tree,
+ object X,
+ const float64_t[:, ::1] y,
+ const float64_t[:] sample_weight=*,
+ const unsigned char[::1] missing_values_in_feature_mask=*,
+ )
+
+ cdef _check_input(
+ self,
+ object X,
+ const float64_t[:, ::1] y,
+ const float64_t[:] sample_weight,
+ )
+
+cdef struct FrontierRecord:
+ # Record of information of a Node, the frontier for a split. Those records are
+ # maintained in a heap to access the Node with the best improvement in impurity,
+ # allowing growing trees greedily on this improvement.
+ intp_t node_id
+ intp_t start
+ intp_t end
+ intp_t pos
+ intp_t depth
+ bint is_leaf
+ float64_t impurity
+ float64_t impurity_left
+ float64_t impurity_right
+ float64_t improvement
+ float64_t lower_bound
+ float64_t upper_bound
+ float64_t middle_value
+
+# A record on the stack for depth-first tree growing
+cdef struct StackRecord:
+ intp_t start
+ intp_t end
+ intp_t depth
+ intp_t parent
+ bint is_left
+ float64_t impurity
+ intp_t n_constant_features
+ float64_t lower_bound
+ float64_t upper_bound
diff --git a/causalml/source/causalml/inference/tree/_tree/_tree.pyx b/causalml/source/causalml/inference/tree/_tree/_tree.pyx
new file mode 100644
index 0000000000000000000000000000000000000000..06a231edab872105b228b1221864b2be0157d8f0
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/_tree/_tree.pyx
@@ -0,0 +1,1962 @@
+# Authors: Gilles Louppe
+# Peter Prettenhofer
+# Brian Holt
+# Noel Dawe
+# Satrajit Gosh
+# Lars Buitinck
+# Arnaud Joly
+# Joel Nothman
+# Fares Hedayati
+# Jacob Schreiber
+# Nelson Liu
+#
+# License: BSD 3 clause
+
+# distutils: language = c++
+# cython: cdivision=True
+# cython: boundscheck=False
+# cython: wraparound=False
+# cython: language_level=3
+# cython: linetrace=True
+
+from cpython cimport Py_INCREF, PyObject, PyTypeObject
+
+from libc.stdlib cimport free
+from libc.string cimport memcpy
+from libc.string cimport memset
+from libc.stdint cimport INTPTR_MAX
+from libc.math cimport isnan
+from libcpp.vector cimport vector
+from libcpp.algorithm cimport pop_heap
+from libcpp.algorithm cimport push_heap
+from libcpp cimport bool
+
+import struct
+
+import numpy as np
+cimport numpy as cnp
+cnp.import_array()
+
+from scipy.sparse import issparse
+from scipy.sparse import csr_matrix
+
+from ._utils cimport safe_realloc
+from ._utils cimport sizet_ptr_to_ndarray
+
+cdef extern from "numpy/arrayobject.h":
+ object PyArray_NewFromDescr(PyTypeObject* subtype, cnp.dtype descr,
+ int nd, cnp.npy_intp* dims,
+ cnp.npy_intp* strides,
+ void* data, int flags, object obj)
+ int PyArray_SetBaseObject(cnp.ndarray arr, PyObject* obj)
+
+cdef extern from "" namespace "std" nogil:
+ cdef cppclass stack[T]:
+ ctypedef T value_type
+ stack() except +
+ bint empty()
+ void pop()
+ void push(T&) except + # Raise c++ exception for bad_alloc -> MemoryError
+ T& top()
+
+# =============================================================================
+# Types and constants
+# =============================================================================
+
+from numpy import float32 as DTYPE
+from numpy import float64 as DOUBLE
+from numpy import int32 as INT
+
+cdef float64_t INFINITY = np.inf
+cdef float64_t EPSILON = np.finfo('double').eps
+
+# Some handy constants (BestFirstTreeBuilder)
+cdef bint IS_FIRST = 1
+cdef bint IS_NOT_FIRST = 0
+cdef bint IS_LEFT = 1
+cdef bint IS_NOT_LEFT = 0
+
+TREE_LEAF = -1
+TREE_UNDEFINED = -2
+cdef intp_t _TREE_LEAF = TREE_LEAF
+cdef intp_t _TREE_UNDEFINED = TREE_UNDEFINED
+
+# Build the corresponding numpy dtype for Node.
+# This works by casting `dummy` to an array of Node of length 1, which numpy
+# can construct a `dtype`-object for. See https://stackoverflow.com/q/62448946
+# for a more detailed explanation.
+cdef Node dummy
+NODE_DTYPE = np.asarray((&dummy)).dtype
+
+cdef void _init_parent_record(ParentInfo* record) noexcept nogil:
+ record.n_constant_features = 0
+ record.impurity = INFINITY
+ record.lower_bound = -INFINITY
+ record.upper_bound = INFINITY
+
+# =============================================================================
+# TreeBuilder
+# =============================================================================
+
+cdef class TreeBuilder:
+ """Interface for different tree building strategies."""
+
+ cpdef build(
+ self,
+ Tree tree,
+ object X,
+ const float64_t[:, ::1] y,
+ const float64_t[:] sample_weight=None,
+ const unsigned char[::1] missing_values_in_feature_mask=None,
+ ):
+ """Build a decision tree from the training set (X, y)."""
+ pass
+
+ cdef inline _check_input(
+ self,
+ object X,
+ const float64_t[:, ::1] y,
+ const float64_t[:] sample_weight,
+ ):
+ """Check input dtype, layout and format"""
+ if issparse(X):
+ X = X.tocsc()
+ X.sort_indices()
+
+ if X.data.dtype != DTYPE:
+ X.data = np.ascontiguousarray(X.data, dtype=DTYPE)
+
+ if X.indices.dtype != np.int32 or X.indptr.dtype != np.int32:
+ raise ValueError("No support for np.int64 index based "
+ "sparse matrices")
+
+ elif X.dtype != DTYPE:
+ # since we have to copy we will make it fortran for efficiency
+ X = np.asfortranarray(X, dtype=DTYPE)
+
+ # TODO: This check for y seems to be redundant, as it is also
+ # present in the BaseDecisionTree's fit method, and therefore
+ # can be removed.
+ if y.base.dtype != DOUBLE or not y.base.flags.contiguous:
+ y = np.ascontiguousarray(y, dtype=DOUBLE)
+
+ if (
+ sample_weight is not None and
+ (
+ sample_weight.base.dtype != DOUBLE or
+ not sample_weight.base.flags.contiguous
+ )
+ ):
+ sample_weight = np.asarray(sample_weight, dtype=DOUBLE, order="C")
+
+ return X, y, sample_weight
+
+# Depth first builder ---------------------------------------------------------
+cdef class DepthFirstTreeBuilder(TreeBuilder):
+ """Build a decision tree in depth-first fashion."""
+
+ def __cinit__(self, Splitter splitter, intp_t min_samples_split,
+ intp_t min_samples_leaf, float64_t min_weight_leaf,
+ intp_t max_depth, float64_t min_impurity_decrease,
+ *args, **kwargs):
+ self.splitter = splitter
+ self.min_samples_split = min_samples_split
+ self.min_samples_leaf = min_samples_leaf
+ self.min_weight_leaf = min_weight_leaf
+ self.max_depth = max_depth
+ self.min_impurity_decrease = min_impurity_decrease
+
+ cpdef build(
+ self,
+ Tree tree,
+ object X,
+ const float64_t[:, ::1] y,
+ const float64_t[:] sample_weight=None,
+ const unsigned char[::1] missing_values_in_feature_mask=None,
+ ):
+ """Build a decision tree from the training set (X, y)."""
+
+ # check input
+ X, y, sample_weight = self._check_input(X, y, sample_weight)
+
+ # Initial capacity
+ cdef intp_t init_capacity
+
+ if tree.max_depth <= 10:
+ init_capacity = (2 ** (tree.max_depth + 1)) - 1
+ else:
+ init_capacity = 2047
+
+ tree._resize(init_capacity)
+
+ # Parameters
+ cdef Splitter splitter = self.splitter
+ cdef intp_t max_depth = self.max_depth
+ cdef intp_t min_samples_leaf = self.min_samples_leaf
+ cdef float64_t min_weight_leaf = self.min_weight_leaf
+ cdef intp_t min_samples_split = self.min_samples_split
+ cdef float64_t min_impurity_decrease = self.min_impurity_decrease
+
+ # Recursive partition (without actual recursion)
+ splitter.init(X, y, sample_weight, missing_values_in_feature_mask)
+
+ cdef intp_t start
+ cdef intp_t end
+ cdef intp_t depth
+ cdef intp_t parent
+ cdef bint is_left
+ cdef intp_t n_node_samples = splitter.n_samples
+ cdef float64_t weighted_n_node_samples
+ cdef SplitRecord split
+ cdef intp_t node_id
+
+ cdef float64_t middle_value
+ cdef float64_t left_child_min
+ cdef float64_t left_child_max
+ cdef float64_t right_child_min
+ cdef float64_t right_child_max
+ cdef bint is_leaf
+ cdef bint first = 1
+ cdef intp_t max_depth_seen = -1
+ cdef int rc = 0
+
+ cdef stack[StackRecord] builder_stack
+ cdef StackRecord stack_record
+
+ cdef ParentInfo parent_record
+ _init_parent_record(&parent_record)
+
+ with nogil:
+ # push root node onto stack
+ builder_stack.push({
+ "start": 0,
+ "end": n_node_samples,
+ "depth": 0,
+ "parent": _TREE_UNDEFINED,
+ "is_left": 0,
+ "impurity": INFINITY,
+ "n_constant_features": 0,
+ "lower_bound": -INFINITY,
+ "upper_bound": INFINITY,
+ })
+
+ while not builder_stack.empty():
+ stack_record = builder_stack.top()
+ builder_stack.pop()
+
+ start = stack_record.start
+ end = stack_record.end
+ depth = stack_record.depth
+ parent = stack_record.parent
+ is_left = stack_record.is_left
+ parent_record.impurity = stack_record.impurity
+ parent_record.n_constant_features = stack_record.n_constant_features
+ parent_record.lower_bound = stack_record.lower_bound
+ parent_record.upper_bound = stack_record.upper_bound
+
+ n_node_samples = end - start
+ splitter.node_reset(start, end, &weighted_n_node_samples)
+
+ is_leaf = (depth >= max_depth or
+ n_node_samples < min_samples_split or
+ n_node_samples < 2 * min_samples_leaf or
+ weighted_n_node_samples < 2 * min_weight_leaf)
+
+ if first:
+ parent_record.impurity = splitter.node_impurity()
+ first = 0
+
+ # impurity == 0 with tolerance due to rounding errors
+ is_leaf = is_leaf or parent_record.impurity <= EPSILON
+
+ if not is_leaf:
+ splitter.node_split(
+ &parent_record,
+ &split,
+ )
+ # If EPSILON=0 in the below comparison, float precision
+ # issues stop splitting, producing trees that are
+ # dissimilar to v0.18
+ is_leaf = (is_leaf or split.pos >= end or
+ (split.improvement + EPSILON <
+ min_impurity_decrease))
+
+ node_id = tree._add_node(parent, is_left, is_leaf, split.feature,
+ split.threshold, parent_record.impurity,
+ n_node_samples, weighted_n_node_samples,
+ split.missing_go_to_left)
+
+ if node_id == INTPTR_MAX:
+ rc = -1
+ break
+
+ # Store value for all nodes, to facilitate tree/model
+ # inspection and interpretation
+ splitter.node_value(tree.value + node_id * tree.value_stride)
+ if splitter.with_monotonic_cst:
+ splitter.clip_node_value(tree.value + node_id * tree.value_stride, parent_record.lower_bound, parent_record.upper_bound)
+
+ if not is_leaf:
+ if (
+ not splitter.with_monotonic_cst or
+ splitter.monotonic_cst[split.feature] == 0
+ ):
+ # Split on a feature with no monotonicity constraint
+
+ # Current bounds must always be propagated to both children.
+ # If a monotonic constraint is active, bounds are used in
+ # node value clipping.
+ left_child_min = right_child_min = parent_record.lower_bound
+ left_child_max = right_child_max = parent_record.upper_bound
+ elif splitter.monotonic_cst[split.feature] == 1:
+ # Split on a feature with monotonic increase constraint
+ left_child_min = parent_record.lower_bound
+ right_child_max = parent_record.upper_bound
+
+ # Lower bound for right child and upper bound for left child
+ # are set to the same value.
+ middle_value = splitter.criterion.middle_value()
+ right_child_min = middle_value
+ left_child_max = middle_value
+ else: # i.e. splitter.monotonic_cst[split.feature] == -1
+ # Split on a feature with monotonic decrease constraint
+ right_child_min = parent_record.lower_bound
+ left_child_max = parent_record.upper_bound
+
+ # Lower bound for left child and upper bound for right child
+ # are set to the same value.
+ middle_value = splitter.criterion.middle_value()
+ left_child_min = middle_value
+ right_child_max = middle_value
+
+ # Push right child on stack
+ builder_stack.push({
+ "start": split.pos,
+ "end": end,
+ "depth": depth + 1,
+ "parent": node_id,
+ "is_left": 0,
+ "impurity": split.impurity_right,
+ "n_constant_features": parent_record.n_constant_features,
+ "lower_bound": right_child_min,
+ "upper_bound": right_child_max,
+ })
+
+ # Push left child on stack
+ builder_stack.push({
+ "start": start,
+ "end": split.pos,
+ "depth": depth + 1,
+ "parent": node_id,
+ "is_left": 1,
+ "impurity": split.impurity_left,
+ "n_constant_features": parent_record.n_constant_features,
+ "lower_bound": left_child_min,
+ "upper_bound": left_child_max,
+ })
+
+ if depth > max_depth_seen:
+ max_depth_seen = depth
+
+ if rc >= 0:
+ rc = tree._resize_c(tree.node_count)
+
+ if rc >= 0:
+ tree.max_depth = max_depth_seen
+ if rc == -1:
+ raise MemoryError()
+
+
+# Best first builder ----------------------------------------------------------
+cdef inline bool _compare_records(
+ const FrontierRecord& left,
+ const FrontierRecord& right,
+):
+ return left.improvement < right.improvement
+
+cdef inline void _add_to_frontier(
+ FrontierRecord rec,
+ vector[FrontierRecord]& frontier,
+) noexcept nogil:
+ """Adds record `rec` to the priority queue `frontier`."""
+ frontier.push_back(rec)
+ push_heap(frontier.begin(), frontier.end(), &_compare_records)
+
+
+cdef class BestFirstTreeBuilder(TreeBuilder):
+ """Build a decision tree in best-first fashion.
+
+ The best node to expand is given by the node at the frontier that has the
+ highest impurity improvement.
+ """
+ cdef intp_t max_leaf_nodes
+
+ def __cinit__(self, Splitter splitter, intp_t min_samples_split,
+ intp_t min_samples_leaf, min_weight_leaf,
+ intp_t max_depth, intp_t max_leaf_nodes,
+ float64_t min_impurity_decrease,
+ *args, **kwargs):
+ self.splitter = splitter
+ self.min_samples_split = min_samples_split
+ self.min_samples_leaf = min_samples_leaf
+ self.min_weight_leaf = min_weight_leaf
+ self.max_depth = max_depth
+ self.max_leaf_nodes = max_leaf_nodes
+ self.min_impurity_decrease = min_impurity_decrease
+
+ cpdef build(
+ self,
+ Tree tree,
+ object X,
+ const float64_t[:, ::1] y,
+ const float64_t[:] sample_weight=None,
+ const unsigned char[::1] missing_values_in_feature_mask=None,
+ ):
+ """Build a decision tree from the training set (X, y)."""
+
+ # check input
+ X, y, sample_weight = self._check_input(X, y, sample_weight)
+
+ # Parameters
+ cdef Splitter splitter = self.splitter
+ cdef intp_t max_leaf_nodes = self.max_leaf_nodes
+
+ # Recursive partition (without actual recursion)
+ splitter.init(X, y, sample_weight, missing_values_in_feature_mask)
+
+ cdef vector[FrontierRecord] frontier
+ cdef FrontierRecord record
+ cdef FrontierRecord split_node_left
+ cdef FrontierRecord split_node_right
+ cdef float64_t left_child_min
+ cdef float64_t left_child_max
+ cdef float64_t right_child_min
+ cdef float64_t right_child_max
+
+ cdef intp_t n_node_samples = splitter.n_samples
+ cdef intp_t max_split_nodes = max_leaf_nodes - 1
+ cdef bint is_leaf
+ cdef intp_t max_depth_seen = -1
+ cdef int rc = 0
+ cdef Node* node
+
+ cdef ParentInfo parent_record
+ _init_parent_record(&parent_record)
+
+ # Initial capacity
+ cdef intp_t init_capacity = max_split_nodes + max_leaf_nodes
+ tree._resize(init_capacity)
+
+ with nogil:
+ # add root to frontier
+ rc = self._add_split_node(
+ splitter=splitter,
+ tree=tree,
+ start=0,
+ end=n_node_samples,
+ is_first=IS_FIRST,
+ is_left=IS_LEFT,
+ parent=NULL,
+ depth=0,
+ parent_record=&parent_record,
+ res=&split_node_left,
+ )
+ if rc >= 0:
+ _add_to_frontier(split_node_left, frontier)
+
+ while not frontier.empty():
+ pop_heap(frontier.begin(), frontier.end(), &_compare_records)
+ record = frontier.back()
+ frontier.pop_back()
+
+ node = &tree.nodes[record.node_id]
+ is_leaf = (record.is_leaf or max_split_nodes <= 0)
+
+ if is_leaf:
+ # Node is not expandable; set node as leaf
+ node.left_child = _TREE_LEAF
+ node.right_child = _TREE_LEAF
+ node.feature = _TREE_UNDEFINED
+ node.threshold = _TREE_UNDEFINED
+
+ else:
+ # Node is expandable
+
+ if (
+ not splitter.with_monotonic_cst or
+ splitter.monotonic_cst[node.feature] == 0
+ ):
+ # Split on a feature with no monotonicity constraint
+
+ # Current bounds must always be propagated to both children.
+ # If a monotonic constraint is active, bounds are used in
+ # node value clipping.
+ left_child_min = right_child_min = record.lower_bound
+ left_child_max = right_child_max = record.upper_bound
+ elif splitter.monotonic_cst[node.feature] == 1:
+ # Split on a feature with monotonic increase constraint
+ left_child_min = record.lower_bound
+ right_child_max = record.upper_bound
+
+ # Lower bound for right child and upper bound for left child
+ # are set to the same value.
+ right_child_min = record.middle_value
+ left_child_max = record.middle_value
+ else: # i.e. splitter.monotonic_cst[split.feature] == -1
+ # Split on a feature with monotonic decrease constraint
+ right_child_min = record.lower_bound
+ left_child_max = record.upper_bound
+
+ # Lower bound for left child and upper bound for right child
+ # are set to the same value.
+ left_child_min = record.middle_value
+ right_child_max = record.middle_value
+
+ # Decrement number of split nodes available
+ max_split_nodes -= 1
+
+ # Compute left split node
+ parent_record.lower_bound = left_child_min
+ parent_record.upper_bound = left_child_max
+ parent_record.impurity = record.impurity_left
+ rc = self._add_split_node(
+ splitter=splitter,
+ tree=tree,
+ start=record.start,
+ end=record.pos,
+ is_first=IS_NOT_FIRST,
+ is_left=IS_LEFT,
+ parent=node,
+ depth=record.depth + 1,
+ parent_record=&parent_record,
+ res=&split_node_left,
+ )
+ if rc == -1:
+ break
+
+ # tree.nodes may have changed
+ node = &tree.nodes[record.node_id]
+
+ # Compute right split node
+ parent_record.lower_bound = right_child_min
+ parent_record.upper_bound = right_child_max
+ parent_record.impurity = record.impurity_right
+ rc = self._add_split_node(
+ splitter=splitter,
+ tree=tree,
+ start=record.pos,
+ end=record.end,
+ is_first=IS_NOT_FIRST,
+ is_left=IS_NOT_LEFT,
+ parent=node,
+ depth=record.depth + 1,
+ parent_record=&parent_record,
+ res=&split_node_right,
+ )
+ if rc == -1:
+ break
+
+ # Add nodes to queue
+ _add_to_frontier(split_node_left, frontier)
+ _add_to_frontier(split_node_right, frontier)
+
+ if record.depth > max_depth_seen:
+ max_depth_seen = record.depth
+
+ if rc >= 0:
+ rc = tree._resize_c(tree.node_count)
+
+ if rc >= 0:
+ tree.max_depth = max_depth_seen
+
+ if rc == -1:
+ raise MemoryError()
+
+ cdef inline int _add_split_node(
+ self,
+ Splitter splitter,
+ Tree tree,
+ intp_t start,
+ intp_t end,
+ bint is_first,
+ bint is_left,
+ Node* parent,
+ intp_t depth,
+ ParentInfo* parent_record,
+ FrontierRecord* res
+ ) except -1 nogil:
+ """Adds node w/ partition ``[start, end)`` to the frontier. """
+ cdef SplitRecord split
+ cdef intp_t node_id
+ cdef intp_t n_node_samples
+ cdef float64_t min_impurity_decrease = self.min_impurity_decrease
+ cdef float64_t weighted_n_node_samples
+ cdef bint is_leaf
+
+ splitter.node_reset(start, end, &weighted_n_node_samples)
+
+ # reset n_constant_features for this specific split before beginning split search
+ parent_record.n_constant_features = 0
+
+ if is_first:
+ parent_record.impurity = splitter.node_impurity()
+
+ n_node_samples = end - start
+ is_leaf = (depth >= self.max_depth or
+ n_node_samples < self.min_samples_split or
+ n_node_samples < 2 * self.min_samples_leaf or
+ weighted_n_node_samples < 2 * self.min_weight_leaf or
+ parent_record.impurity <= EPSILON # impurity == 0 with tolerance
+ )
+
+ if not is_leaf:
+ splitter.node_split(
+ parent_record,
+ &split,
+ )
+ # If EPSILON=0 in the below comparison, float precision issues stop
+ # splitting early, producing trees that are dissimilar to v0.18
+ is_leaf = (is_leaf or split.pos >= end or
+ split.improvement + EPSILON < min_impurity_decrease)
+
+ node_id = tree._add_node(parent - tree.nodes
+ if parent != NULL
+ else _TREE_UNDEFINED,
+ is_left, is_leaf,
+ split.feature, split.threshold, parent_record.impurity,
+ n_node_samples, weighted_n_node_samples,
+ split.missing_go_to_left)
+ if node_id == INTPTR_MAX:
+ return -1
+
+ # compute values also for split nodes (might become leafs later).
+ splitter.node_value(tree.value + node_id * tree.value_stride)
+ if splitter.with_monotonic_cst:
+ splitter.clip_node_value(tree.value + node_id * tree.value_stride, parent_record.lower_bound, parent_record.upper_bound)
+
+ res.node_id = node_id
+ res.start = start
+ res.end = end
+ res.depth = depth
+ res.impurity = parent_record.impurity
+ res.lower_bound = parent_record.lower_bound
+ res.upper_bound = parent_record.upper_bound
+ res.middle_value = splitter.criterion.middle_value()
+
+ if not is_leaf:
+ # is split node
+ res.pos = split.pos
+ res.is_leaf = 0
+ res.improvement = split.improvement
+ res.impurity_left = split.impurity_left
+ res.impurity_right = split.impurity_right
+
+ else:
+ # is leaf => 0 improvement
+ res.pos = end
+ res.is_leaf = 1
+ res.improvement = 0.0
+ res.impurity_left = parent_record.impurity
+ res.impurity_right = parent_record.impurity
+
+ return 0
+
+
+# =============================================================================
+# Tree
+# =============================================================================
+
+cdef class Tree:
+ """Array-based representation of a binary decision tree.
+
+ The binary tree is represented as a number of parallel arrays. The i-th
+ element of each array holds information about the node `i`. Node 0 is the
+ tree's root. You can find a detailed description of all arrays in
+ `_tree.pxd`. NOTE: Some of the arrays only apply to either leaves or split
+ nodes, resp. In this case the values of nodes of the other type are
+ arbitrary!
+
+ Attributes
+ ----------
+ node_count : intp_t
+ The number of nodes (internal nodes + leaves) in the tree.
+
+ capacity : intp_t
+ The current capacity (i.e., size) of the arrays, which is at least as
+ great as `node_count`.
+
+ max_depth : intp_t
+ The depth of the tree, i.e. the maximum depth of its leaves.
+
+ children_left : array of intp_t, shape [node_count]
+ children_left[i] holds the node id of the left child of node i.
+ For leaves, children_left[i] == TREE_LEAF. Otherwise,
+ children_left[i] > i. This child handles the case where
+ X[:, feature[i]] <= threshold[i].
+
+ children_right : array of intp_t, shape [node_count]
+ children_right[i] holds the node id of the right child of node i.
+ For leaves, children_right[i] == TREE_LEAF. Otherwise,
+ children_right[i] > i. This child handles the case where
+ X[:, feature[i]] > threshold[i].
+
+ n_leaves : intp_t
+ Number of leaves in the tree.
+
+ feature : array of intp_t, shape [node_count]
+ feature[i] holds the feature to split on, for the internal node i.
+
+ threshold : array of float64_t, shape [node_count]
+ threshold[i] holds the threshold for the internal node i.
+
+ value : array of float64_t, shape [node_count, n_outputs, max_n_classes]
+ Contains the constant prediction value of each node.
+
+ impurity : array of float64_t, shape [node_count]
+ impurity[i] holds the impurity (i.e., the value of the splitting
+ criterion) at node i.
+
+ n_node_samples : array of intp_t, shape [node_count]
+ n_node_samples[i] holds the number of training samples reaching node i.
+
+ weighted_n_node_samples : array of float64_t, shape [node_count]
+ weighted_n_node_samples[i] holds the weighted number of training samples
+ reaching node i.
+
+ missing_go_to_left : array of bool, shape [node_count]
+ missing_go_to_left[i] holds a bool indicating whether or not there were
+ missing values at node i.
+ """
+ # Wrap for outside world.
+ # WARNING: these reference the current `nodes` and `value` buffers, which
+ # must not be freed by a subsequent memory allocation.
+ # (i.e. through `_resize` or `__setstate__`)
+ @property
+ def n_classes(self):
+ return sizet_ptr_to_ndarray(self.n_classes, self.n_outputs)
+
+ @property
+ def children_left(self):
+ return self._get_node_ndarray()['left_child'][:self.node_count]
+
+ @property
+ def children_right(self):
+ return self._get_node_ndarray()['right_child'][:self.node_count]
+
+ @property
+ def n_leaves(self):
+ return np.sum(np.logical_and(
+ self.children_left == -1,
+ self.children_right == -1))
+
+ @property
+ def feature(self):
+ return self._get_node_ndarray()['feature'][:self.node_count]
+
+ @property
+ def threshold(self):
+ return self._get_node_ndarray()['threshold'][:self.node_count]
+
+ @property
+ def impurity(self):
+ return self._get_node_ndarray()['impurity'][:self.node_count]
+
+ @property
+ def n_node_samples(self):
+ return self._get_node_ndarray()['n_node_samples'][:self.node_count]
+
+ @property
+ def weighted_n_node_samples(self):
+ return self._get_node_ndarray()['weighted_n_node_samples'][:self.node_count]
+
+ @property
+ def missing_go_to_left(self):
+ return self._get_node_ndarray()['missing_go_to_left'][:self.node_count]
+
+ @property
+ def value(self):
+ return self._get_value_ndarray()[:self.node_count]
+
+ # TODO: Convert n_classes to cython.integral memory view once
+ # https://github.com/cython/cython/issues/5243 is fixed
+ def __cinit__(self, intp_t n_features, cnp.ndarray n_classes, intp_t n_outputs):
+ """Constructor."""
+ cdef intp_t dummy = 0
+ intp_t_dtype = np.array(dummy).dtype
+
+ n_classes = _check_n_classes(n_classes, intp_t_dtype)
+
+ # Input/Output layout
+ self.n_features = n_features
+ self.n_outputs = n_outputs
+ self.n_classes = NULL
+ safe_realloc(&self.n_classes, n_outputs)
+
+ self.max_n_classes = np.max(n_classes)
+ self.value_stride = n_outputs * self.max_n_classes
+
+ cdef intp_t k
+ for k in range(n_outputs):
+ self.n_classes[k] = n_classes[k]
+
+ # Inner structures
+ self.max_depth = 0
+ self.node_count = 0
+ self.capacity = 0
+ self.value = NULL
+ self.nodes = NULL
+
+ def __dealloc__(self):
+ """Destructor."""
+ # Free all inner structures
+ free(self.n_classes)
+ free(self.value)
+ free(self.nodes)
+
+ def __reduce__(self):
+ """Reduce re-implementation, for pickling."""
+ return (Tree, (self.n_features,
+ sizet_ptr_to_ndarray(self.n_classes, self.n_outputs),
+ self.n_outputs), self.__getstate__())
+
+ def __getstate__(self):
+ """Getstate re-implementation, for pickling."""
+ d = {}
+ # capacity is inferred during the __setstate__ using nodes
+ d["max_depth"] = self.max_depth
+ d["node_count"] = self.node_count
+ d["nodes"] = self._get_node_ndarray()
+ d["values"] = self._get_value_ndarray()
+ return d
+
+ def __setstate__(self, d):
+ """Setstate re-implementation, for unpickling."""
+ self.max_depth = d["max_depth"]
+ self.node_count = d["node_count"]
+
+ if 'nodes' not in d:
+ raise ValueError('You have loaded Tree version which '
+ 'cannot be imported')
+
+ node_ndarray = d['nodes']
+ value_ndarray = d['values']
+
+ value_shape = (node_ndarray.shape[0], self.n_outputs,
+ self.max_n_classes)
+
+ node_ndarray = _check_node_ndarray(node_ndarray, expected_dtype=NODE_DTYPE)
+ value_ndarray = _check_value_ndarray(
+ value_ndarray,
+ expected_dtype=np.dtype(np.float64),
+ expected_shape=value_shape
+ )
+
+ self.capacity = node_ndarray.shape[0]
+ if self._resize_c(self.capacity) != 0:
+ raise MemoryError("resizing tree to %d" % self.capacity)
+
+ memcpy(self.nodes, cnp.PyArray_DATA(node_ndarray),
+ self.capacity * sizeof(Node))
+ memcpy(self.value, cnp.PyArray_DATA(value_ndarray),
+ self.capacity * self.value_stride * sizeof(float64_t))
+
+ cdef int _resize(self, intp_t capacity) except -1 nogil:
+ """Resize all inner arrays to `capacity`, if `capacity` == -1, then
+ double the size of the inner arrays.
+
+ Returns -1 in case of failure to allocate memory (and raise MemoryError)
+ or 0 otherwise.
+ """
+ if self._resize_c(capacity) != 0:
+ # Acquire gil only if we need to raise
+ with gil:
+ raise MemoryError()
+
+ cdef int _resize_c(self, intp_t capacity=INTPTR_MAX) except -1 nogil:
+ """Guts of _resize
+
+ Returns -1 in case of failure to allocate memory (and raise MemoryError)
+ or 0 otherwise.
+ """
+ if capacity == self.capacity and self.nodes != NULL:
+ return 0
+
+ if capacity == INTPTR_MAX:
+ if self.capacity == 0:
+ capacity = 3 # default initial value
+ else:
+ capacity = 2 * self.capacity
+
+ safe_realloc(&self.nodes, capacity)
+ safe_realloc(&self.value, capacity * self.value_stride)
+
+ if capacity > self.capacity:
+ # value memory is initialised to 0 to enable classifier argmax
+ memset((self.value + self.capacity * self.value_stride), 0,
+ (capacity - self.capacity) * self.value_stride *
+ sizeof(float64_t))
+ # node memory is initialised to 0 to ensure deterministic pickle (padding in Node struct)
+ memset((self.nodes + self.capacity), 0, (capacity - self.capacity) * sizeof(Node))
+
+ # if capacity smaller than node_count, adjust the counter
+ if capacity < self.node_count:
+ self.node_count = capacity
+
+ self.capacity = capacity
+ return 0
+
+ cdef intp_t _add_node(self, intp_t parent, bint is_left, bint is_leaf,
+ intp_t feature, float64_t threshold, float64_t impurity,
+ intp_t n_node_samples,
+ float64_t weighted_n_node_samples,
+ unsigned char missing_go_to_left) except -1 nogil:
+ """Add a node to the tree.
+
+ The new node registers itself as the child of its parent.
+
+ Returns (intp_t)(-1) on error.
+ """
+ cdef intp_t node_id = self.node_count
+
+ if node_id >= self.capacity:
+ if self._resize_c() != 0:
+ return INTPTR_MAX
+
+ cdef Node* node = &self.nodes[node_id]
+ node.impurity = impurity
+ node.n_node_samples = n_node_samples
+ node.weighted_n_node_samples = weighted_n_node_samples
+
+ if parent != _TREE_UNDEFINED:
+ if is_left:
+ self.nodes[parent].left_child = node_id
+ else:
+ self.nodes[parent].right_child = node_id
+
+ if is_leaf:
+ node.left_child = _TREE_LEAF
+ node.right_child = _TREE_LEAF
+ node.feature = _TREE_UNDEFINED
+ node.threshold = _TREE_UNDEFINED
+
+ else:
+ # left_child and right_child will be set later
+ node.feature = feature
+ node.threshold = threshold
+ node.missing_go_to_left = missing_go_to_left
+
+ self.node_count += 1
+
+ return node_id
+
+ cpdef cnp.ndarray predict(self, object X):
+ """Predict target for X."""
+ out = self._get_value_ndarray().take(self.apply(X), axis=0,
+ mode='clip')
+ if self.n_outputs == 1:
+ out = out.reshape(X.shape[0], self.max_n_classes)
+ return out
+
+ cpdef cnp.ndarray apply(self, object X):
+ """Finds the terminal region (=leaf node) for each sample in X."""
+ if issparse(X):
+ return self._apply_sparse_csr(X)
+ else:
+ return self._apply_dense(X)
+
+ cdef inline cnp.ndarray _apply_dense(self, object X):
+ """Finds the terminal region (=leaf node) for each sample in X."""
+
+ # Check input
+ if not isinstance(X, np.ndarray):
+ raise ValueError("X should be in np.ndarray format, got %s"
+ % type(X))
+
+ if X.dtype != DTYPE:
+ raise ValueError("X.dtype should be np.float32, got %s" % X.dtype)
+
+ # Extract input
+ cdef const float32_t[:, :] X_ndarray = X
+ cdef intp_t n_samples = X.shape[0]
+ cdef float32_t X_i_node_feature
+
+ # Initialize output
+ cdef intp_t[:] out = np.zeros(n_samples, dtype=np.intp)
+
+ # Initialize auxiliary data-structure
+ cdef Node* node = NULL
+ cdef intp_t i = 0
+
+ with nogil:
+ for i in range(n_samples):
+ node = self.nodes
+ # While node not a leaf
+ while node.left_child != _TREE_LEAF:
+ X_i_node_feature = X_ndarray[i, node.feature]
+ # ... and node.right_child != _TREE_LEAF:
+ if isnan(X_i_node_feature):
+ if node.missing_go_to_left:
+ node = &self.nodes[node.left_child]
+ else:
+ node = &self.nodes[node.right_child]
+ elif X_i_node_feature <= node.threshold:
+ node = &self.nodes[node.left_child]
+ else:
+ node = &self.nodes[node.right_child]
+
+ out[i] = (node - self.nodes) # node offset
+
+ return np.asarray(out)
+
+ cdef inline cnp.ndarray _apply_sparse_csr(self, object X):
+ """Finds the terminal region (=leaf node) for each sample in sparse X.
+ """
+ # Check input
+ if not (issparse(X) and X.format == 'csr'):
+ raise ValueError("X should be in csr_matrix format, got %s"
+ % type(X))
+
+ if X.dtype != DTYPE:
+ raise ValueError("X.dtype should be np.float32, got %s" % X.dtype)
+
+ # Extract input
+ cdef const float32_t[:] X_data = X.data
+ cdef const int32_t[:] X_indices = X.indices
+ cdef const int32_t[:] X_indptr = X.indptr
+
+ cdef intp_t n_samples = X.shape[0]
+ cdef intp_t n_features = X.shape[1]
+
+ # Initialize output
+ cdef intp_t[:] out = np.zeros(n_samples, dtype=np.intp)
+
+ # Initialize auxiliary data-structure
+ cdef float32_t feature_value = 0.
+ cdef Node* node = NULL
+ cdef float32_t* X_sample = NULL
+ cdef intp_t i = 0
+ cdef int32_t k = 0
+
+ # feature_to_sample as a data structure records the last seen sample
+ # for each feature; functionally, it is an efficient way to identify
+ # which features are nonzero in the present sample.
+ cdef intp_t* feature_to_sample = NULL
+
+ safe_realloc(&X_sample, n_features)
+ safe_realloc(&feature_to_sample, n_features)
+
+ with nogil:
+ memset(feature_to_sample, -1, n_features * sizeof(intp_t))
+
+ for i in range(n_samples):
+ node = self.nodes
+
+ for k in range(X_indptr[i], X_indptr[i + 1]):
+ feature_to_sample[X_indices[k]] = i
+ X_sample[X_indices[k]] = X_data[k]
+
+ # While node not a leaf
+ while node.left_child != _TREE_LEAF:
+ # ... and node.right_child != _TREE_LEAF:
+ if feature_to_sample[node.feature] == i:
+ feature_value = X_sample[node.feature]
+
+ else:
+ feature_value = 0.
+
+ if feature_value <= node.threshold:
+ node = &self.nodes[node.left_child]
+ else:
+ node = &self.nodes[node.right_child]
+
+ out[i] = (node - self.nodes) # node offset
+
+ # Free auxiliary arrays
+ free(X_sample)
+ free(feature_to_sample)
+
+ return np.asarray(out)
+
+ cpdef object decision_path(self, object X):
+ """Finds the decision path (=node) for each sample in X."""
+ if issparse(X):
+ return self._decision_path_sparse_csr(X)
+ else:
+ return self._decision_path_dense(X)
+
+ cdef inline object _decision_path_dense(self, object X):
+ """Finds the decision path (=node) for each sample in X."""
+
+ # Check input
+ if not isinstance(X, np.ndarray):
+ raise ValueError("X should be in np.ndarray format, got %s"
+ % type(X))
+
+ if X.dtype != DTYPE:
+ raise ValueError("X.dtype should be np.float32, got %s" % X.dtype)
+
+ # Extract input
+ cdef const float32_t[:, :] X_ndarray = X
+ cdef intp_t n_samples = X.shape[0]
+
+ # Initialize output
+ cdef intp_t[:] indptr = np.zeros(n_samples + 1, dtype=np.intp)
+ cdef intp_t[:] indices = np.zeros(
+ n_samples * (1 + self.max_depth), dtype=np.intp
+ )
+
+ # Initialize auxiliary data-structure
+ cdef Node* node = NULL
+ cdef intp_t i = 0
+
+ with nogil:
+ for i in range(n_samples):
+ node = self.nodes
+ indptr[i + 1] = indptr[i]
+
+ # Add all external nodes
+ while node.left_child != _TREE_LEAF:
+ # ... and node.right_child != _TREE_LEAF:
+ indices[indptr[i + 1]] = (node - self.nodes)
+ indptr[i + 1] += 1
+
+ if X_ndarray[i, node.feature] <= node.threshold:
+ node = &self.nodes[node.left_child]
+ else:
+ node = &self.nodes[node.right_child]
+
+ # Add the leave node
+ indices[indptr[i + 1]] = (node - self.nodes)
+ indptr[i + 1] += 1
+
+ indices = indices[:indptr[n_samples]]
+ cdef intp_t[:] data = np.ones(shape=len(indices), dtype=np.intp)
+ out = csr_matrix((data, indices, indptr),
+ shape=(n_samples, self.node_count))
+
+ return out
+
+ cdef inline object _decision_path_sparse_csr(self, object X):
+ """Finds the decision path (=node) for each sample in X."""
+
+ # Check input
+ if not (issparse(X) and X.format == "csr"):
+ raise ValueError("X should be in csr_matrix format, got %s"
+ % type(X))
+
+ if X.dtype != DTYPE:
+ raise ValueError("X.dtype should be np.float32, got %s" % X.dtype)
+
+ # Extract input
+ cdef const float32_t[:] X_data = X.data
+ cdef const int32_t[:] X_indices = X.indices
+ cdef const int32_t[:] X_indptr = X.indptr
+
+ cdef intp_t n_samples = X.shape[0]
+ cdef intp_t n_features = X.shape[1]
+
+ # Initialize output
+ cdef intp_t[:] indptr = np.zeros(n_samples + 1, dtype=np.intp)
+ cdef intp_t[:] indices = np.zeros(
+ n_samples * (1 + self.max_depth), dtype=np.intp
+ )
+
+ # Initialize auxiliary data-structure
+ cdef float32_t feature_value = 0.
+ cdef Node* node = NULL
+ cdef float32_t* X_sample = NULL
+ cdef intp_t i = 0
+ cdef int32_t k = 0
+
+ # feature_to_sample as a data structure records the last seen sample
+ # for each feature; functionally, it is an efficient way to identify
+ # which features are nonzero in the present sample.
+ cdef intp_t* feature_to_sample = NULL
+
+ safe_realloc(&X_sample, n_features)
+ safe_realloc(&feature_to_sample, n_features)
+
+ with nogil:
+ memset(feature_to_sample, -1, n_features * sizeof(intp_t))
+
+ for i in range(n_samples):
+ node = self.nodes
+ indptr[i + 1] = indptr[i]
+
+ for k in range(X_indptr[i], X_indptr[i + 1]):
+ feature_to_sample[X_indices[k]] = i
+ X_sample[X_indices[k]] = X_data[k]
+
+ # While node not a leaf
+ while node.left_child != _TREE_LEAF:
+ # ... and node.right_child != _TREE_LEAF:
+
+ indices[indptr[i + 1]] = (node - self.nodes)
+ indptr[i + 1] += 1
+
+ if feature_to_sample[node.feature] == i:
+ feature_value = X_sample[node.feature]
+
+ else:
+ feature_value = 0.
+
+ if feature_value <= node.threshold:
+ node = &self.nodes[node.left_child]
+ else:
+ node = &self.nodes[node.right_child]
+
+ # Add the leave node
+ indices[indptr[i + 1]] = (node - self.nodes)
+ indptr[i + 1] += 1
+
+ # Free auxiliary arrays
+ free(X_sample)
+ free(feature_to_sample)
+
+ indices = indices[:indptr[n_samples]]
+ cdef intp_t[:] data = np.ones(shape=len(indices), dtype=np.intp)
+ out = csr_matrix((data, indices, indptr),
+ shape=(n_samples, self.node_count))
+
+ return out
+
+ cpdef compute_node_depths(self):
+ """Compute the depth of each node in a tree.
+
+ .. versionadded:: 1.3
+
+ Returns
+ -------
+ depths : ndarray of shape (self.node_count,), dtype=np.int64
+ The depth of each node in the tree.
+ """
+ cdef:
+ cnp.int64_t[::1] depths = np.empty(self.node_count, dtype=np.int64)
+ cnp.npy_intp[:] children_left = self.children_left
+ cnp.npy_intp[:] children_right = self.children_right
+ cnp.npy_intp node_id
+ cnp.npy_intp node_count = self.node_count
+ cnp.int64_t depth
+
+ depths[0] = 1 # init root node
+ for node_id in range(node_count):
+ if children_left[node_id] != _TREE_LEAF:
+ depth = depths[node_id] + 1
+ depths[children_left[node_id]] = depth
+ depths[children_right[node_id]] = depth
+
+ return depths.base
+
+ cpdef compute_feature_importances(self, normalize=True):
+ """Computes the importance of each feature (aka variable)."""
+ cdef Node* left
+ cdef Node* right
+ cdef Node* nodes = self.nodes
+ cdef Node* node = nodes
+ cdef Node* end_node = node + self.node_count
+
+ cdef float64_t normalizer = 0.
+
+ cdef cnp.float64_t[:] importances = np.zeros(self.n_features)
+
+ with nogil:
+ while node != end_node:
+ if node.left_child != _TREE_LEAF:
+ # ... and node.right_child != _TREE_LEAF:
+ left = &nodes[node.left_child]
+ right = &nodes[node.right_child]
+
+ importances[node.feature] += (
+ node.weighted_n_node_samples * node.impurity -
+ left.weighted_n_node_samples * left.impurity -
+ right.weighted_n_node_samples * right.impurity)
+ node += 1
+
+ for i in range(self.n_features):
+ importances[i] /= nodes[0].weighted_n_node_samples
+
+ if normalize:
+ normalizer = np.sum(importances)
+
+ if normalizer > 0.0:
+ # Avoid dividing by zero (e.g., when root is pure)
+ for i in range(self.n_features):
+ importances[i] /= normalizer
+
+ return np.asarray(importances)
+
+ cdef cnp.ndarray _get_value_ndarray(self):
+ """Wraps value as a 3-d NumPy array.
+
+ The array keeps a reference to this Tree, which manages the underlying
+ memory.
+ """
+ cdef cnp.npy_intp shape[3]
+ shape[0] = self.node_count
+ shape[1] = self.n_outputs
+ shape[2] = self.max_n_classes
+ cdef cnp.ndarray arr
+ arr = cnp.PyArray_SimpleNewFromData(3, shape, cnp.NPY_DOUBLE, self.value)
+ Py_INCREF(self)
+ if PyArray_SetBaseObject(arr, self) < 0:
+ raise ValueError("Can't initialize array.")
+ return arr
+
+ cdef cnp.ndarray _get_node_ndarray(self):
+ """Wraps nodes as a NumPy struct array.
+
+ The array keeps a reference to this Tree, which manages the underlying
+ memory. Individual fields are publicly accessible as properties of the
+ Tree.
+ """
+ cdef cnp.npy_intp shape[1]
+ shape[0] = self.node_count
+ cdef cnp.npy_intp strides[1]
+ strides[0] = sizeof(Node)
+ cdef cnp.ndarray arr
+ Py_INCREF(NODE_DTYPE)
+ arr = PyArray_NewFromDescr( cnp.ndarray,
+ NODE_DTYPE, 1, shape,
+ strides, self.nodes,
+ cnp.NPY_ARRAY_DEFAULT, None)
+ Py_INCREF(self)
+ if PyArray_SetBaseObject(arr, self) < 0:
+ raise ValueError("Can't initialize array.")
+ return arr
+
+ def compute_partial_dependence(self, float32_t[:, ::1] X,
+ const intp_t[::1] target_features,
+ float64_t[::1] out):
+ """Partial dependence of the response on the ``target_feature`` set.
+
+ For each sample in ``X`` a tree traversal is performed.
+ Each traversal starts from the root with weight 1.0.
+
+ At each non-leaf node that splits on a target feature, either
+ the left child or the right child is visited based on the feature
+ value of the current sample, and the weight is not modified.
+ At each non-leaf node that splits on a complementary feature,
+ both children are visited and the weight is multiplied by the fraction
+ of training samples which went to each child.
+
+ At each leaf, the value of the node is multiplied by the current
+ weight (weights sum to 1 for all visited terminal nodes).
+
+ Parameters
+ ----------
+ X : view on 2d ndarray, shape (n_samples, n_target_features)
+ The grid points on which the partial dependence should be
+ evaluated.
+ target_features : view on 1d ndarray, shape (n_target_features)
+ The set of target features for which the partial dependence
+ should be evaluated.
+ out : view on 1d ndarray, shape (n_samples)
+ The value of the partial dependence function on each grid
+ point.
+ """
+ cdef:
+ float64_t[::1] weight_stack = np.zeros(self.node_count,
+ dtype=np.float64)
+ intp_t[::1] node_idx_stack = np.zeros(self.node_count,
+ dtype=np.intp)
+ intp_t sample_idx
+ intp_t feature_idx
+ intp_t stack_size
+ float64_t left_sample_frac
+ float64_t current_weight
+ float64_t total_weight # used for sanity check only
+ Node *current_node # use a pointer to avoid copying attributes
+ intp_t current_node_idx
+ bint is_target_feature
+ intp_t _TREE_LEAF = TREE_LEAF # to avoid python interactions
+
+ for sample_idx in range(X.shape[0]):
+ # init stacks for current sample
+ stack_size = 1
+ node_idx_stack[0] = 0 # root node
+ weight_stack[0] = 1 # all the samples are in the root node
+ total_weight = 0
+
+ while stack_size > 0:
+ # pop the stack
+ stack_size -= 1
+ current_node_idx = node_idx_stack[stack_size]
+ current_node = &self.nodes[current_node_idx]
+
+ if current_node.left_child == _TREE_LEAF:
+ # leaf node
+ out[sample_idx] += (weight_stack[stack_size] *
+ self.value[current_node_idx])
+ total_weight += weight_stack[stack_size]
+ else:
+ # non-leaf node
+
+ # determine if the split feature is a target feature
+ is_target_feature = False
+ for feature_idx in range(target_features.shape[0]):
+ if target_features[feature_idx] == current_node.feature:
+ is_target_feature = True
+ break
+
+ if is_target_feature:
+ # In this case, we push left or right child on stack
+ if X[sample_idx, feature_idx] <= current_node.threshold:
+ node_idx_stack[stack_size] = current_node.left_child
+ else:
+ node_idx_stack[stack_size] = current_node.right_child
+ stack_size += 1
+ else:
+ # In this case, we push both children onto the stack,
+ # and give a weight proportional to the number of
+ # samples going through each branch.
+
+ # push left child
+ node_idx_stack[stack_size] = current_node.left_child
+ left_sample_frac = (
+ self.nodes[current_node.left_child].weighted_n_node_samples /
+ current_node.weighted_n_node_samples)
+ current_weight = weight_stack[stack_size]
+ weight_stack[stack_size] = current_weight * left_sample_frac
+ stack_size += 1
+
+ # push right child
+ node_idx_stack[stack_size] = current_node.right_child
+ weight_stack[stack_size] = (
+ current_weight * (1 - left_sample_frac))
+ stack_size += 1
+
+ # Sanity check. Should never happen.
+ if not (0.999 < total_weight < 1.001):
+ raise ValueError("Total weight should be 1.0 but was %.9f" %
+ total_weight)
+
+
+def _check_n_classes(n_classes, expected_dtype):
+ if n_classes.ndim != 1:
+ raise ValueError(
+ f"Wrong dimensions for n_classes from the pickle: "
+ f"expected 1, got {n_classes.ndim}"
+ )
+
+ if n_classes.dtype == expected_dtype:
+ return n_classes
+
+ # Handles both different endianness and different bitness
+ if n_classes.dtype.kind == "i" and n_classes.dtype.itemsize in [4, 8]:
+ return n_classes.astype(expected_dtype, casting="same_kind")
+
+ raise ValueError(
+ "n_classes from the pickle has an incompatible dtype:\n"
+ f"- expected: {expected_dtype}\n"
+ f"- got: {n_classes.dtype}"
+ )
+
+
+def _check_value_ndarray(value_ndarray, expected_dtype, expected_shape):
+ if value_ndarray.shape != expected_shape:
+ raise ValueError(
+ "Wrong shape for value array from the pickle: "
+ f"expected {expected_shape}, got {value_ndarray.shape}"
+ )
+
+ if not value_ndarray.flags.c_contiguous:
+ raise ValueError(
+ "value array from the pickle should be a C-contiguous array"
+ )
+
+ if value_ndarray.dtype == expected_dtype:
+ return value_ndarray
+
+ # Handles different endianness
+ if value_ndarray.dtype.str.endswith('f8'):
+ return value_ndarray.astype(expected_dtype, casting='equiv')
+
+ raise ValueError(
+ "value array from the pickle has an incompatible dtype:\n"
+ f"- expected: {expected_dtype}\n"
+ f"- got: {value_ndarray.dtype}"
+ )
+
+
+def _dtype_to_dict(dtype):
+ return {name: dt.str for name, (dt, *rest) in dtype.fields.items()}
+
+
+def _dtype_dict_with_modified_bitness(dtype_dict):
+ # field names in Node struct with intp_t types (see sklearn/tree/_tree.pxd)
+ indexing_field_names = ["left_child", "right_child", "feature", "n_node_samples"]
+
+ expected_dtype_size = str(struct.calcsize("P"))
+ allowed_dtype_size = "8" if expected_dtype_size == "4" else "4"
+
+ allowed_dtype_dict = dtype_dict.copy()
+ for name in indexing_field_names:
+ allowed_dtype_dict[name] = allowed_dtype_dict[name].replace(
+ expected_dtype_size, allowed_dtype_size
+ )
+
+ return allowed_dtype_dict
+
+
+def _all_compatible_dtype_dicts(dtype):
+ # The Cython code for decision trees uses platform-specific intp_t
+ # typed indexing fields that correspond to either i4 or i8 dtypes for
+ # the matching fields in the numpy array depending on the bitness of
+ # the platform (32 bit or 64 bit respectively).
+ #
+ # We need to cast the indexing fields of the NODE_DTYPE-dtyped array at
+ # pickle load time to enable cross-bitness deployment scenarios. We
+ # typically want to make it possible to run the expensive fit method of
+ # a tree estimator on a 64 bit server platform, pickle the estimator
+ # for deployment and run the predict method of a low power 32 bit edge
+ # platform.
+ #
+ # A similar thing happens for endianness, the machine where the pickle was
+ # saved can have a different endianness than the machine where the pickle
+ # is loaded
+
+ dtype_dict = _dtype_to_dict(dtype)
+ dtype_dict_with_modified_bitness = _dtype_dict_with_modified_bitness(dtype_dict)
+ dtype_dict_with_modified_endianness = _dtype_to_dict(dtype.newbyteorder())
+ dtype_dict_with_modified_bitness_and_endianness = _dtype_dict_with_modified_bitness(
+ dtype_dict_with_modified_endianness
+ )
+
+ return [
+ dtype_dict,
+ dtype_dict_with_modified_bitness,
+ dtype_dict_with_modified_endianness,
+ dtype_dict_with_modified_bitness_and_endianness,
+ ]
+
+
+def _check_node_ndarray(node_ndarray, expected_dtype):
+ if node_ndarray.ndim != 1:
+ raise ValueError(
+ "Wrong dimensions for node array from the pickle: "
+ f"expected 1, got {node_ndarray.ndim}"
+ )
+
+ if not node_ndarray.flags.c_contiguous:
+ raise ValueError(
+ "node array from the pickle should be a C-contiguous array"
+ )
+
+ node_ndarray_dtype = node_ndarray.dtype
+ if node_ndarray_dtype == expected_dtype:
+ return node_ndarray
+
+ node_ndarray_dtype_dict = _dtype_to_dict(node_ndarray_dtype)
+ all_compatible_dtype_dicts = _all_compatible_dtype_dicts(expected_dtype)
+
+ if node_ndarray_dtype_dict not in all_compatible_dtype_dicts:
+ raise ValueError(
+ "node array from the pickle has an incompatible dtype:\n"
+ f"- expected: {expected_dtype}\n"
+ f"- got : {node_ndarray_dtype}"
+ )
+
+ return node_ndarray.astype(expected_dtype, casting="same_kind")
+
+
+# =============================================================================
+# Build Pruned Tree
+# =============================================================================
+
+
+cdef class _CCPPruneController:
+ """Base class used by build_pruned_tree_ccp and ccp_pruning_path
+ to control pruning.
+ """
+ cdef bint stop_pruning(self, float64_t effective_alpha) noexcept nogil:
+ """Return 1 to stop pruning and 0 to continue pruning"""
+ return 0
+
+ cdef void save_metrics(self, float64_t effective_alpha,
+ float64_t subtree_impurities) noexcept nogil:
+ """Save metrics when pruning"""
+ pass
+
+ cdef void after_pruning(self, unsigned char[:] in_subtree) noexcept nogil:
+ """Called after pruning"""
+ pass
+
+
+cdef class _AlphaPruner(_CCPPruneController):
+ """Use alpha to control when to stop pruning."""
+ cdef float64_t ccp_alpha
+ cdef intp_t capacity
+
+ def __cinit__(self, float64_t ccp_alpha):
+ self.ccp_alpha = ccp_alpha
+ self.capacity = 0
+
+ cdef bint stop_pruning(self, float64_t effective_alpha) noexcept nogil:
+ # The subtree on the previous iteration has the greatest ccp_alpha
+ # less than or equal to self.ccp_alpha
+ return self.ccp_alpha < effective_alpha
+
+ cdef void after_pruning(self, unsigned char[:] in_subtree) noexcept nogil:
+ """Updates the number of leaves in subtree"""
+ for i in range(in_subtree.shape[0]):
+ if in_subtree[i]:
+ self.capacity += 1
+
+
+cdef class _PathFinder(_CCPPruneController):
+ """Record metrics used to return the cost complexity path."""
+ cdef float64_t[:] ccp_alphas
+ cdef float64_t[:] impurities
+ cdef uint32_t count
+
+ def __cinit__(self, intp_t node_count):
+ self.ccp_alphas = np.zeros(shape=(node_count), dtype=np.float64)
+ self.impurities = np.zeros(shape=(node_count), dtype=np.float64)
+ self.count = 0
+
+ cdef void save_metrics(self,
+ float64_t effective_alpha,
+ float64_t subtree_impurities) noexcept nogil:
+ self.ccp_alphas[self.count] = effective_alpha
+ self.impurities[self.count] = subtree_impurities
+ self.count += 1
+
+
+cdef struct CostComplexityPruningRecord:
+ intp_t node_idx
+ intp_t parent
+
+cdef _cost_complexity_prune(unsigned char[:] leaves_in_subtree, # OUT
+ Tree orig_tree,
+ _CCPPruneController controller):
+ """Perform cost complexity pruning.
+
+ This function takes an already grown tree, `orig_tree` and outputs a
+ boolean mask `leaves_in_subtree` which are the leaves in the pruned tree.
+ During the pruning process, the controller is passed the effective alpha and
+ the subtree impurities. Furthermore, the controller signals when to stop
+ pruning.
+
+ Parameters
+ ----------
+ leaves_in_subtree : unsigned char[:]
+ Output for leaves of subtree
+ orig_tree : Tree
+ Original tree
+ ccp_controller : _CCPPruneController
+ Cost complexity controller
+ """
+
+ cdef:
+ intp_t i
+ intp_t n_nodes = orig_tree.node_count
+ # prior probability using weighted samples
+ float64_t[:] weighted_n_node_samples = orig_tree.weighted_n_node_samples
+ float64_t total_sum_weights = weighted_n_node_samples[0]
+ float64_t[:] impurity = orig_tree.impurity
+ # weighted impurity of each node
+ float64_t[:] r_node = np.empty(shape=n_nodes, dtype=np.float64)
+
+ intp_t[:] child_l = orig_tree.children_left
+ intp_t[:] child_r = orig_tree.children_right
+ intp_t[:] parent = np.zeros(shape=n_nodes, dtype=np.intp)
+
+ stack[CostComplexityPruningRecord] ccp_stack
+ CostComplexityPruningRecord stack_record
+ intp_t node_idx
+ stack[intp_t] node_indices_stack
+
+ intp_t[:] n_leaves = np.zeros(shape=n_nodes, dtype=np.intp)
+ float64_t[:] r_branch = np.zeros(shape=n_nodes, dtype=np.float64)
+ float64_t current_r
+ intp_t leaf_idx
+ intp_t parent_idx
+
+ # candidate nodes that can be pruned
+ unsigned char[:] candidate_nodes = np.zeros(shape=n_nodes,
+ dtype=np.uint8)
+ # nodes in subtree
+ unsigned char[:] in_subtree = np.ones(shape=n_nodes, dtype=np.uint8)
+ intp_t pruned_branch_node_idx
+ float64_t subtree_alpha
+ float64_t effective_alpha
+ intp_t n_pruned_leaves
+ float64_t r_diff
+ float64_t max_float64 = np.finfo(np.float64).max
+
+ # find parent node ids and leaves
+ with nogil:
+
+ for i in range(r_node.shape[0]):
+ r_node[i] = (
+ weighted_n_node_samples[i] * impurity[i] / total_sum_weights)
+
+ # Push the root node
+ ccp_stack.push({"node_idx": 0, "parent": _TREE_UNDEFINED})
+
+ while not ccp_stack.empty():
+ stack_record = ccp_stack.top()
+ ccp_stack.pop()
+
+ node_idx = stack_record.node_idx
+ parent[node_idx] = stack_record.parent
+
+ if child_l[node_idx] == _TREE_LEAF:
+ # ... and child_r[node_idx] == _TREE_LEAF:
+ leaves_in_subtree[node_idx] = 1
+ else:
+ ccp_stack.push({"node_idx": child_l[node_idx], "parent": node_idx})
+ ccp_stack.push({"node_idx": child_r[node_idx], "parent": node_idx})
+
+ # computes number of leaves in all branches and the overall impurity of
+ # the branch. The overall impurity is the sum of r_node in its leaves.
+ for leaf_idx in range(leaves_in_subtree.shape[0]):
+ if not leaves_in_subtree[leaf_idx]:
+ continue
+ r_branch[leaf_idx] = r_node[leaf_idx]
+
+ # bubble up values to ancestor nodes
+ current_r = r_node[leaf_idx]
+ while leaf_idx != 0:
+ parent_idx = parent[leaf_idx]
+ r_branch[parent_idx] += current_r
+ n_leaves[parent_idx] += 1
+ leaf_idx = parent_idx
+
+ for i in range(leaves_in_subtree.shape[0]):
+ candidate_nodes[i] = not leaves_in_subtree[i]
+
+ # save metrics before pruning
+ controller.save_metrics(0.0, r_branch[0])
+
+ # while root node is not a leaf
+ while candidate_nodes[0]:
+
+ # computes ccp_alpha for subtrees and finds the minimal alpha
+ effective_alpha = max_float64
+ for i in range(n_nodes):
+ if not candidate_nodes[i]:
+ continue
+ subtree_alpha = (r_node[i] - r_branch[i]) / (n_leaves[i] - 1)
+ if subtree_alpha < effective_alpha:
+ effective_alpha = subtree_alpha
+ pruned_branch_node_idx = i
+
+ if controller.stop_pruning(effective_alpha):
+ break
+
+ node_indices_stack.push(pruned_branch_node_idx)
+
+ # descendants of branch are not in subtree
+ while not node_indices_stack.empty():
+ node_idx = node_indices_stack.top()
+ node_indices_stack.pop()
+
+ if not in_subtree[node_idx]:
+ continue # branch has already been marked for pruning
+ candidate_nodes[node_idx] = 0
+ leaves_in_subtree[node_idx] = 0
+ in_subtree[node_idx] = 0
+
+ if child_l[node_idx] != _TREE_LEAF:
+ # ... and child_r[node_idx] != _TREE_LEAF:
+ node_indices_stack.push(child_l[node_idx])
+ node_indices_stack.push(child_r[node_idx])
+ leaves_in_subtree[pruned_branch_node_idx] = 1
+ in_subtree[pruned_branch_node_idx] = 1
+
+ # updates number of leaves
+ n_pruned_leaves = n_leaves[pruned_branch_node_idx] - 1
+ n_leaves[pruned_branch_node_idx] = 0
+
+ # computes the increase in r_branch to bubble up
+ r_diff = r_node[pruned_branch_node_idx] - r_branch[pruned_branch_node_idx]
+ r_branch[pruned_branch_node_idx] = r_node[pruned_branch_node_idx]
+
+ # bubble up values to ancestors
+ node_idx = parent[pruned_branch_node_idx]
+ while node_idx != _TREE_UNDEFINED:
+ n_leaves[node_idx] -= n_pruned_leaves
+ r_branch[node_idx] += r_diff
+ node_idx = parent[node_idx]
+
+ controller.save_metrics(effective_alpha, r_branch[0])
+
+ controller.after_pruning(in_subtree)
+
+
+def _build_pruned_tree_ccp(
+ Tree tree, # OUT
+ Tree orig_tree,
+ float64_t ccp_alpha
+):
+ """Build a pruned tree from the original tree using cost complexity
+ pruning.
+
+ The values and nodes from the original tree are copied into the pruned
+ tree.
+
+ Parameters
+ ----------
+ tree : Tree
+ Location to place the pruned tree
+ orig_tree : Tree
+ Original tree
+ ccp_alpha : positive float64_t
+ Complexity parameter. The subtree with the largest cost complexity
+ that is smaller than ``ccp_alpha`` will be chosen. By default,
+ no pruning is performed.
+ """
+
+ cdef:
+ intp_t n_nodes = orig_tree.node_count
+ unsigned char[:] leaves_in_subtree = np.zeros(
+ shape=n_nodes, dtype=np.uint8)
+
+ pruning_controller = _AlphaPruner(ccp_alpha=ccp_alpha)
+
+ _cost_complexity_prune(leaves_in_subtree, orig_tree, pruning_controller)
+
+ _build_pruned_tree(tree, orig_tree, leaves_in_subtree,
+ pruning_controller.capacity)
+
+
+def ccp_pruning_path(Tree orig_tree):
+ """Computes the cost complexity pruning path.
+
+ Parameters
+ ----------
+ tree : Tree
+ Original tree.
+
+ Returns
+ -------
+ path_info : dict
+ Information about pruning path with attributes:
+
+ ccp_alphas : ndarray
+ Effective alphas of subtree during pruning.
+
+ impurities : ndarray
+ Sum of the impurities of the subtree leaves for the
+ corresponding alpha value in ``ccp_alphas``.
+ """
+ cdef:
+ unsigned char[:] leaves_in_subtree = np.zeros(
+ shape=orig_tree.node_count, dtype=np.uint8)
+
+ path_finder = _PathFinder(orig_tree.node_count)
+
+ _cost_complexity_prune(leaves_in_subtree, orig_tree, path_finder)
+
+ cdef:
+ uint32_t total_items = path_finder.count
+ float64_t[:] ccp_alphas = np.empty(shape=total_items, dtype=np.float64)
+ float64_t[:] impurities = np.empty(shape=total_items, dtype=np.float64)
+ uint32_t count = 0
+
+ while count < total_items:
+ ccp_alphas[count] = path_finder.ccp_alphas[count]
+ impurities[count] = path_finder.impurities[count]
+ count += 1
+
+ return {
+ 'ccp_alphas': np.asarray(ccp_alphas),
+ 'impurities': np.asarray(impurities),
+ }
+
+
+cdef struct BuildPrunedRecord:
+ intp_t start
+ intp_t depth
+ intp_t parent
+ bint is_left
+
+cdef _build_pruned_tree(
+ Tree tree, # OUT
+ Tree orig_tree,
+ const unsigned char[:] leaves_in_subtree,
+ intp_t capacity
+):
+ """Build a pruned tree.
+
+ Build a pruned tree from the original tree by transforming the nodes in
+ ``leaves_in_subtree`` into leaves.
+
+ Parameters
+ ----------
+ tree : Tree
+ Location to place the pruned tree
+ orig_tree : Tree
+ Original tree
+ leaves_in_subtree : unsigned char memoryview, shape=(node_count, )
+ Boolean mask for leaves to include in subtree
+ capacity : intp_t
+ Number of nodes to initially allocate in pruned tree
+ """
+ tree._resize(capacity)
+
+ cdef:
+ intp_t orig_node_id
+ intp_t new_node_id
+ intp_t depth
+ intp_t parent
+ bint is_left
+ bint is_leaf
+
+ # value_stride for original tree and new tree are the same
+ intp_t value_stride = orig_tree.value_stride
+ intp_t max_depth_seen = -1
+ int rc = 0
+ Node* node
+ float64_t* orig_value_ptr
+ float64_t* new_value_ptr
+
+ stack[BuildPrunedRecord] prune_stack
+ BuildPrunedRecord stack_record
+
+ with nogil:
+ # push root node onto stack
+ prune_stack.push({"start": 0, "depth": 0, "parent": _TREE_UNDEFINED, "is_left": 0})
+
+ while not prune_stack.empty():
+ stack_record = prune_stack.top()
+ prune_stack.pop()
+
+ orig_node_id = stack_record.start
+ depth = stack_record.depth
+ parent = stack_record.parent
+ is_left = stack_record.is_left
+
+ is_leaf = leaves_in_subtree[orig_node_id]
+ node = &orig_tree.nodes[orig_node_id]
+
+ new_node_id = tree._add_node(
+ parent, is_left, is_leaf, node.feature, node.threshold,
+ node.impurity, node.n_node_samples,
+ node.weighted_n_node_samples, node.missing_go_to_left)
+
+ if new_node_id == INTPTR_MAX:
+ rc = -1
+ break
+
+ # copy value from original tree to new tree
+ orig_value_ptr = orig_tree.value + value_stride * orig_node_id
+ new_value_ptr = tree.value + value_stride * new_node_id
+ memcpy(new_value_ptr, orig_value_ptr, sizeof(float64_t) * value_stride)
+
+ if not is_leaf:
+ # Push right child on stack
+ prune_stack.push({"start": node.right_child, "depth": depth + 1,
+ "parent": new_node_id, "is_left": 0})
+ # push left child on stack
+ prune_stack.push({"start": node.left_child, "depth": depth + 1,
+ "parent": new_node_id, "is_left": 1})
+
+ if depth > max_depth_seen:
+ max_depth_seen = depth
+
+ if rc >= 0:
+ tree.max_depth = max_depth_seen
+ if rc == -1:
+ raise MemoryError("pruning tree")
diff --git a/causalml/source/causalml/inference/tree/_tree/_typedefs.pxd b/causalml/source/causalml/inference/tree/_tree/_typedefs.pxd
new file mode 100644
index 0000000000000000000000000000000000000000..f77227466158055ba2a2d4c8db51f235400c7801
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/_tree/_typedefs.pxd
@@ -0,0 +1,41 @@
+# Commonly used types
+# These are redefinitions of the ones defined by numpy in
+# https://github.com/numpy/numpy/blob/main/numpy/__init__.pxd.
+# It will eventually avoid having to always include the numpy headers even when we
+# would only use it for the types.
+#
+# When used to declare variables that will receive values from numpy arrays, it
+# should match the dtype of the array. For example, to declare a variable that will
+# receive values from a numpy array of dtype np.float64, the type float64_t must be
+# used.
+#
+# TODO: Stop defining custom types locally or globally like DTYPE_t and friends and
+# use these consistently throughout the codebase.
+# NOTE: Extend this list as needed when converting more cython extensions.
+ctypedef unsigned char uint8_t
+ctypedef unsigned int uint32_t
+ctypedef unsigned long long uint64_t
+# Note: In NumPy 2, indexing always happens with npy_intp which is an alias for
+# the Py_ssize_t type, see PEP 353.
+#
+# Note that on most platforms Py_ssize_t is equivalent to C99's intptr_t,
+# but they can differ on architecture with segmented memory (none
+# supported by scikit-learn at the time of writing).
+#
+# intp_t/np.intp should be used to index arrays in a platform dependent way.
+# Storing arrays with platform dependent dtypes as attribute on picklable
+# objects is not recommended as it requires special care when loading and
+# using such datastructures on a host with different bitness. Instead one
+# should rather use fixed width integer types such as int32 or uint32 when we know
+# that the number of elements to index is not larger to 2 or 4 billions.
+ctypedef Py_ssize_t intp_t
+ctypedef float float32_t
+ctypedef double float64_t
+# Sparse matrices indices and indices' pointers arrays must use int32_t over
+# intp_t because intp_t is platform dependent.
+# When large sparse matrices are supported, indexing must use int64_t.
+# See https://github.com/scikit-learn/scikit-learn/issues/23653 which tracks the
+# ongoing work to support large sparse matrices.
+ctypedef signed char int8_t
+ctypedef signed int int32_t
+ctypedef signed long long int64_t
diff --git a/causalml/source/causalml/inference/tree/_tree/_typedefs.pyx b/causalml/source/causalml/inference/tree/_tree/_typedefs.pyx
new file mode 100644
index 0000000000000000000000000000000000000000..2d8eaab49e1b7d1b209760a8744cb52c051c35fc
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/_tree/_typedefs.pyx
@@ -0,0 +1,23 @@
+# _typedefs is a declaration only module
+#
+# The functions implemented here are for testing purpose only.
+
+
+import numpy as np
+
+
+ctypedef fused testing_type_t:
+ float32_t
+ float64_t
+ int8_t
+ int32_t
+ int64_t
+ intp_t
+ uint8_t
+ uint32_t
+ uint64_t
+
+
+def testing_make_array_from_typed_val(testing_type_t val):
+ cdef testing_type_t[:] val_view = &val
+ return np.asarray(val_view)
diff --git a/causalml/source/causalml/inference/tree/_tree/_utils.pxd b/causalml/source/causalml/inference/tree/_tree/_utils.pxd
new file mode 100644
index 0000000000000000000000000000000000000000..4a2280327b2179089eecb0f24100a64c754461dc
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/_tree/_utils.pxd
@@ -0,0 +1,109 @@
+# Authors: Gilles Louppe
+# Peter Prettenhofer
+# Arnaud Joly
+# Jacob Schreiber
+# Nelson Liu
+#
+# License: BSD 3 clause
+
+# distutils: language = c++
+# cython: cdivision=True
+# cython: boundscheck=False
+# cython: wraparound=False
+# cython: language_level=3
+# cython: linetrace=True
+
+# See _utils.pyx for details.
+
+cimport numpy as cnp
+from ._tree cimport Node
+from ._typedefs cimport float32_t, float64_t, intp_t, int32_t, uint32_t
+
+cdef enum:
+ # Max value for our rand_r replacement (near the bottom).
+ # We don't use RAND_MAX because it's different across platforms and
+ # particularly tiny on Windows/MSVC.
+ # It corresponds to the maximum representable value for
+ # 32-bit signed integers (i.e. 2^31 - 1).
+ RAND_R_MAX = 2147483647
+
+
+# safe_realloc(&p, n) resizes the allocation of p to n * sizeof(*p) bytes or
+# raises a MemoryError. It never calls free, since that's __dealloc__'s job.
+# cdef float32_t *p = NULL
+# safe_realloc(&p, n)
+# is equivalent to p = malloc(n * sizeof(*p)) with error checking.
+ctypedef fused realloc_ptr:
+ # Add pointer types here as needed.
+ (float32_t*)
+ (intp_t*)
+ (unsigned char*)
+ (WeightedPQueueRecord*)
+ (float64_t*)
+ (float64_t**)
+ (Node*)
+ (Node**)
+
+cdef int safe_realloc(realloc_ptr* p, intp_t nelems) except -1 nogil
+
+
+cdef cnp.ndarray sizet_ptr_to_ndarray(intp_t* data, intp_t size)
+
+
+cdef intp_t rand_int(intp_t low, intp_t high,
+ uint32_t* random_state) noexcept nogil
+
+
+cdef float64_t rand_uniform(float64_t low, float64_t high,
+ uint32_t* random_state) noexcept nogil
+
+
+cdef float64_t log(float64_t x) noexcept nogil
+
+# =============================================================================
+# WeightedPQueue data structure
+# =============================================================================
+
+# A record stored in the WeightedPQueue
+cdef struct WeightedPQueueRecord:
+ float64_t data
+ float64_t weight
+
+cdef class WeightedPQueue:
+ cdef intp_t capacity
+ cdef intp_t array_ptr
+ cdef WeightedPQueueRecord* array_
+
+ cdef bint is_empty(self) noexcept nogil
+ cdef int reset(self) except -1 nogil
+ cdef intp_t size(self) noexcept nogil
+ cdef int push(self, float64_t data, float64_t weight) except -1 nogil
+ cdef int remove(self, float64_t data, float64_t weight) noexcept nogil
+ cdef int pop(self, float64_t* data, float64_t* weight) noexcept nogil
+ cdef int peek(self, float64_t* data, float64_t* weight) noexcept nogil
+ cdef float64_t get_weight_from_index(self, intp_t index) noexcept nogil
+ cdef float64_t get_value_from_index(self, intp_t index) noexcept nogil
+
+
+# =============================================================================
+# WeightedMedianCalculator data structure
+# =============================================================================
+
+cdef class WeightedMedianCalculator:
+ cdef intp_t initial_capacity
+ cdef WeightedPQueue samples
+ cdef float64_t total_weight
+ cdef intp_t k
+ cdef float64_t sum_w_0_k # represents sum(weights[0:k]) = w[0] + w[1] + ... + w[k-1]
+ cdef intp_t size(self) noexcept nogil
+ cdef int push(self, float64_t data, float64_t weight) except -1 nogil
+ cdef int reset(self) except -1 nogil
+ cdef int update_median_parameters_post_push(
+ self, float64_t data, float64_t weight,
+ float64_t original_median) noexcept nogil
+ cdef int remove(self, float64_t data, float64_t weight) noexcept nogil
+ cdef int pop(self, float64_t* data, float64_t* weight) noexcept nogil
+ cdef int update_median_parameters_post_remove(
+ self, float64_t data, float64_t weight,
+ float64_t original_median) noexcept nogil
+ cdef float64_t get_median(self) noexcept nogil
diff --git a/causalml/source/causalml/inference/tree/_tree/_utils.pyx b/causalml/source/causalml/inference/tree/_tree/_utils.pyx
new file mode 100644
index 0000000000000000000000000000000000000000..28676589f44849dba51cd3855d0a9c3a85ce78b5
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/_tree/_utils.pyx
@@ -0,0 +1,492 @@
+# Authors: Gilles Louppe
+# Peter Prettenhofer
+# Arnaud Joly
+# Jacob Schreiber
+# Nelson Liu
+#
+#
+# License: BSD 3 clause
+
+from libc.stdlib cimport free
+from libc.stdlib cimport realloc
+from libc.math cimport log as ln
+from libc.math cimport isnan
+
+import numpy as np
+cimport numpy as cnp
+cnp.import_array()
+
+# Random number generation utilities
+# Copied from sklearn.utils._random to avoid DEFAULT_SEED signature mismatch
+# Original authors: The scikit-learn developers
+# License: BSD-3-Clause
+# Copied from sklearn 1.6+ _random.pxd to avoid signature mismatch issues
+
+from ._typedefs cimport uint32_t
+
+cdef const uint32_t DEFAULT_SEED = 1
+
+# rand_r replacement using a 32bit XorShift generator
+# See http://www.jstatsoft.org/v08/i14/paper for details
+cdef inline uint32_t our_rand_r(uint32_t* seed) nogil:
+ """Generate a pseudo-random np.uint32 from a np.uint32 seed"""
+ # seed shouldn't ever be 0.
+ if (seed[0] == 0):
+ seed[0] = DEFAULT_SEED
+
+ seed[0] ^= (seed[0] << 13)
+ seed[0] ^= (seed[0] >> 17)
+ seed[0] ^= (seed[0] << 5)
+
+ # Use the modulo to ensure we don't return values greater than
+ # the maximum representable value for signed 32bit integers.
+ return seed[0] % ((RAND_R_MAX) + 1)
+
+# =============================================================================
+# Helper functions
+# =============================================================================
+
+cdef int safe_realloc(realloc_ptr* p, intp_t nelems) except -1 nogil:
+ # sizeof(realloc_ptr[0]) would be more like idiomatic C, but causes Cython
+ # 0.20.1 to crash.
+ cdef intp_t nbytes = nelems * sizeof(p[0][0])
+ if nbytes / sizeof(p[0][0]) != nelems:
+ # Overflow in the multiplication
+ raise MemoryError(f"could not allocate ({nelems} * {sizeof(p[0][0])}) bytes")
+
+ cdef realloc_ptr tmp = realloc(p[0], nbytes)
+ if tmp == NULL:
+ raise MemoryError(f"could not allocate {nbytes} bytes")
+
+ p[0] = tmp
+ return 0
+
+
+"""TODO: fix Cython compile error
+def _realloc_test():
+ # Helper for tests. Tries to allocate (-1) / 2 * sizeof(intp_t)
+ # bytes, which will always overflow.
+ cdef intp_t* p = NULL
+ safe_realloc(&p, (-1) / 2)
+ if p != NULL:
+ free(p)
+ assert False
+"""
+
+
+cdef inline cnp.ndarray sizet_ptr_to_ndarray(intp_t* data, intp_t size):
+ """Return copied data as 1D numpy array of intp's."""
+ cdef cnp.npy_intp shape[1]
+ shape[0] = size
+ return cnp.PyArray_SimpleNewFromData(1, shape, cnp.NPY_INTP, data).copy()
+
+
+cdef inline intp_t rand_int(intp_t low, intp_t high,
+ uint32_t* random_state) noexcept nogil:
+ """Generate a random integer in [low; end)."""
+ return low + our_rand_r(random_state) % (high - low)
+
+
+cdef inline float64_t rand_uniform(float64_t low, float64_t high,
+ uint32_t* random_state) noexcept nogil:
+ """Generate a random float64_t in [low; high)."""
+ return ((high - low) * our_rand_r(random_state) /
+ RAND_R_MAX) + low
+
+
+cdef inline float64_t log(float64_t x) noexcept nogil:
+ return ln(x) / ln(2.0)
+
+# =============================================================================
+# WeightedPQueue data structure
+# =============================================================================
+
+cdef class WeightedPQueue:
+ """A priority queue class, always sorted in increasing order.
+
+ Attributes
+ ----------
+ capacity : intp_t
+ The capacity of the priority queue.
+
+ array_ptr : intp_t
+ The water mark of the priority queue; the priority queue grows from
+ left to right in the array ``array_``. ``array_ptr`` is always
+ less than ``capacity``.
+
+ array_ : WeightedPQueueRecord*
+ The array of priority queue records. The minimum element is on the
+ left at index 0, and the maximum element is on the right at index
+ ``array_ptr-1``.
+ """
+
+ def __cinit__(self, intp_t capacity):
+ self.capacity = capacity
+ self.array_ptr = 0
+ safe_realloc(&self.array_, capacity)
+
+ def __dealloc__(self):
+ free(self.array_)
+
+ cdef int reset(self) except -1 nogil:
+ """Reset the WeightedPQueue to its state at construction
+
+ Return -1 in case of failure to allocate memory (and raise MemoryError)
+ or 0 otherwise.
+ """
+ self.array_ptr = 0
+ # Since safe_realloc can raise MemoryError, use `except -1`
+ safe_realloc(&self.array_, self.capacity)
+ return 0
+
+ cdef bint is_empty(self) noexcept nogil:
+ return self.array_ptr <= 0
+
+ cdef intp_t size(self) noexcept nogil:
+ return self.array_ptr
+
+ cdef int push(self, float64_t data, float64_t weight) except -1 nogil:
+ """Push record on the array.
+
+ Return -1 in case of failure to allocate memory (and raise MemoryError)
+ or 0 otherwise.
+ """
+ cdef intp_t array_ptr = self.array_ptr
+ cdef WeightedPQueueRecord* array = NULL
+ cdef intp_t i
+
+ # Resize if capacity not sufficient
+ if array_ptr >= self.capacity:
+ self.capacity *= 2
+ # Since safe_realloc can raise MemoryError, use `except -1`
+ safe_realloc(&self.array_, self.capacity)
+
+ # Put element as last element of array
+ array = self.array_
+ array[array_ptr].data = data
+ array[array_ptr].weight = weight
+
+ # bubble last element up according until it is sorted
+ # in ascending order
+ i = array_ptr
+ while(i != 0 and array[i].data < array[i-1].data):
+ array[i], array[i-1] = array[i-1], array[i]
+ i -= 1
+
+ # Increase element count
+ self.array_ptr = array_ptr + 1
+ return 0
+
+ cdef int remove(self, float64_t data, float64_t weight) noexcept nogil:
+ """Remove a specific value/weight record from the array.
+ Returns 0 if successful, -1 if record not found."""
+ cdef intp_t array_ptr = self.array_ptr
+ cdef WeightedPQueueRecord* array = self.array_
+ cdef intp_t idx_to_remove = -1
+ cdef intp_t i
+
+ if array_ptr <= 0:
+ return -1
+
+ # find element to remove
+ for i in range(array_ptr):
+ if array[i].data == data and array[i].weight == weight:
+ idx_to_remove = i
+ break
+
+ if idx_to_remove == -1:
+ return -1
+
+ # shift the elements after the removed element
+ # to the left.
+ for i in range(idx_to_remove, array_ptr-1):
+ array[i] = array[i+1]
+
+ self.array_ptr = array_ptr - 1
+ return 0
+
+ cdef int pop(self, float64_t* data, float64_t* weight) noexcept nogil:
+ """Remove the top (minimum) element from array.
+ Returns 0 if successful, -1 if nothing to remove."""
+ cdef intp_t array_ptr = self.array_ptr
+ cdef WeightedPQueueRecord* array = self.array_
+ cdef intp_t i
+
+ if array_ptr <= 0:
+ return -1
+
+ data[0] = array[0].data
+ weight[0] = array[0].weight
+
+ # shift the elements after the removed element
+ # to the left.
+ for i in range(0, array_ptr-1):
+ array[i] = array[i+1]
+
+ self.array_ptr = array_ptr - 1
+ return 0
+
+ cdef int peek(self, float64_t* data, float64_t* weight) noexcept nogil:
+ """Write the top element from array to a pointer.
+ Returns 0 if successful, -1 if nothing to write."""
+ cdef WeightedPQueueRecord* array = self.array_
+ if self.array_ptr <= 0:
+ return -1
+ # Take first value
+ data[0] = array[0].data
+ weight[0] = array[0].weight
+ return 0
+
+ cdef float64_t get_weight_from_index(self, intp_t index) noexcept nogil:
+ """Given an index between [0,self.current_capacity], access
+ the appropriate heap and return the requested weight"""
+ cdef WeightedPQueueRecord* array = self.array_
+
+ # get weight at index
+ return array[index].weight
+
+ cdef float64_t get_value_from_index(self, intp_t index) noexcept nogil:
+ """Given an index between [0,self.current_capacity], access
+ the appropriate heap and return the requested value"""
+ cdef WeightedPQueueRecord* array = self.array_
+
+ # get value at index
+ return array[index].data
+
+# =============================================================================
+# WeightedMedianCalculator data structure
+# =============================================================================
+
+cdef class WeightedMedianCalculator:
+ """A class to handle calculation of the weighted median from streams of
+ data. To do so, it maintains a parameter ``k`` such that the sum of the
+ weights in the range [0,k) is greater than or equal to half of the total
+ weight. By minimizing the value of ``k`` that fulfills this constraint,
+ calculating the median is done by either taking the value of the sample
+ at index ``k-1`` of ``samples`` (samples[k-1].data) or the average of
+ the samples at index ``k-1`` and ``k`` of ``samples``
+ ((samples[k-1] + samples[k]) / 2).
+
+ Attributes
+ ----------
+ initial_capacity : intp_t
+ The initial capacity of the WeightedMedianCalculator.
+
+ samples : WeightedPQueue
+ Holds the samples (consisting of values and their weights) used in the
+ weighted median calculation.
+
+ total_weight : float64_t
+ The sum of the weights of items in ``samples``. Represents the total
+ weight of all samples used in the median calculation.
+
+ k : intp_t
+ Index used to calculate the median.
+
+ sum_w_0_k : float64_t
+ The sum of the weights from samples[0:k]. Used in the weighted
+ median calculation; minimizing the value of ``k`` such that
+ ``sum_w_0_k`` >= ``total_weight / 2`` provides a mechanism for
+ calculating the median in constant time.
+
+ """
+
+ def __cinit__(self, intp_t initial_capacity):
+ self.initial_capacity = initial_capacity
+ self.samples = WeightedPQueue(initial_capacity)
+ self.total_weight = 0
+ self.k = 0
+ self.sum_w_0_k = 0
+
+ cdef intp_t size(self) noexcept nogil:
+ """Return the number of samples in the
+ WeightedMedianCalculator"""
+ return self.samples.size()
+
+ cdef int reset(self) except -1 nogil:
+ """Reset the WeightedMedianCalculator to its state at construction
+
+ Return -1 in case of failure to allocate memory (and raise MemoryError)
+ or 0 otherwise.
+ """
+ # samples.reset (WeightedPQueue.reset) uses safe_realloc, hence
+ # except -1
+ self.samples.reset()
+ self.total_weight = 0
+ self.k = 0
+ self.sum_w_0_k = 0
+ return 0
+
+ cdef int push(self, float64_t data, float64_t weight) except -1 nogil:
+ """Push a value and its associated weight to the WeightedMedianCalculator
+
+ Return -1 in case of failure to allocate memory (and raise MemoryError)
+ or 0 otherwise.
+ """
+ cdef int return_value
+ cdef float64_t original_median = 0.0
+
+ if self.size() != 0:
+ original_median = self.get_median()
+ # samples.push (WeightedPQueue.push) uses safe_realloc, hence except -1
+ return_value = self.samples.push(data, weight)
+ self.update_median_parameters_post_push(data, weight,
+ original_median)
+ return return_value
+
+ cdef int update_median_parameters_post_push(
+ self, float64_t data, float64_t weight,
+ float64_t original_median) noexcept nogil:
+ """Update the parameters used in the median calculation,
+ namely `k` and `sum_w_0_k` after an insertion"""
+
+ # trivial case of one element.
+ if self.size() == 1:
+ self.k = 1
+ self.total_weight = weight
+ self.sum_w_0_k = self.total_weight
+ return 0
+
+ # get the original weighted median
+ self.total_weight += weight
+
+ if data < original_median:
+ # inserting below the median, so increment k and
+ # then update self.sum_w_0_k accordingly by adding
+ # the weight that was added.
+ self.k += 1
+ # update sum_w_0_k by adding the weight added
+ self.sum_w_0_k += weight
+
+ # minimize k such that sum(W[0:k]) >= total_weight / 2
+ # minimum value of k is 1
+ while(self.k > 1 and ((self.sum_w_0_k -
+ self.samples.get_weight_from_index(self.k-1))
+ >= self.total_weight / 2.0)):
+ self.k -= 1
+ self.sum_w_0_k -= self.samples.get_weight_from_index(self.k)
+ return 0
+
+ if data >= original_median:
+ # inserting above or at the median
+ # minimize k such that sum(W[0:k]) >= total_weight / 2
+ while(self.k < self.samples.size() and
+ (self.sum_w_0_k < self.total_weight / 2.0)):
+ self.k += 1
+ self.sum_w_0_k += self.samples.get_weight_from_index(self.k-1)
+ return 0
+
+ cdef int remove(self, float64_t data, float64_t weight) noexcept nogil:
+ """Remove a value from the MedianHeap, removing it
+ from consideration in the median calculation
+ """
+ cdef int return_value
+ cdef float64_t original_median = 0.0
+
+ if self.size() != 0:
+ original_median = self.get_median()
+
+ return_value = self.samples.remove(data, weight)
+ self.update_median_parameters_post_remove(data, weight,
+ original_median)
+ return return_value
+
+ cdef int pop(self, float64_t* data, float64_t* weight) noexcept nogil:
+ """Pop a value from the MedianHeap, starting from the
+ left and moving to the right.
+ """
+ cdef int return_value
+ cdef float64_t original_median = 0.0
+
+ if self.size() != 0:
+ original_median = self.get_median()
+
+ # no elements to pop
+ if self.samples.size() == 0:
+ return -1
+
+ return_value = self.samples.pop(data, weight)
+ self.update_median_parameters_post_remove(data[0],
+ weight[0],
+ original_median)
+ return return_value
+
+ cdef int update_median_parameters_post_remove(
+ self, float64_t data, float64_t weight,
+ float64_t original_median) noexcept nogil:
+ """Update the parameters used in the median calculation,
+ namely `k` and `sum_w_0_k` after a removal"""
+ # reset parameters because it there are no elements
+ if self.samples.size() == 0:
+ self.k = 0
+ self.total_weight = 0
+ self.sum_w_0_k = 0
+ return 0
+
+ # trivial case of one element.
+ if self.samples.size() == 1:
+ self.k = 1
+ self.total_weight -= weight
+ self.sum_w_0_k = self.total_weight
+ return 0
+
+ # get the current weighted median
+ self.total_weight -= weight
+
+ if data < original_median:
+ # removing below the median, so decrement k and
+ # then update self.sum_w_0_k accordingly by subtracting
+ # the removed weight
+
+ self.k -= 1
+ # update sum_w_0_k by removing the weight at index k
+ self.sum_w_0_k -= weight
+
+ # minimize k such that sum(W[0:k]) >= total_weight / 2
+ # by incrementing k and updating sum_w_0_k accordingly
+ # until the condition is met.
+ while(self.k < self.samples.size() and
+ (self.sum_w_0_k < self.total_weight / 2.0)):
+ self.k += 1
+ self.sum_w_0_k += self.samples.get_weight_from_index(self.k-1)
+ return 0
+
+ if data >= original_median:
+ # removing above the median
+ # minimize k such that sum(W[0:k]) >= total_weight / 2
+ while(self.k > 1 and ((self.sum_w_0_k -
+ self.samples.get_weight_from_index(self.k-1))
+ >= self.total_weight / 2.0)):
+ self.k -= 1
+ self.sum_w_0_k -= self.samples.get_weight_from_index(self.k)
+ return 0
+
+ cdef float64_t get_median(self) noexcept nogil:
+ """Write the median to a pointer, taking into account
+ sample weights."""
+ if self.sum_w_0_k == (self.total_weight / 2.0):
+ # split median
+ return (self.samples.get_value_from_index(self.k) +
+ self.samples.get_value_from_index(self.k-1)) / 2.0
+ if self.sum_w_0_k > (self.total_weight / 2.0):
+ # whole median
+ return self.samples.get_value_from_index(self.k-1)
+
+
+def _any_isnan_axis0(const float32_t[:, :] X):
+ """Same as np.any(np.isnan(X), axis=0)"""
+ cdef:
+ intp_t i, j
+ intp_t n_samples = X.shape[0]
+ intp_t n_features = X.shape[1]
+ unsigned char[::1] isnan_out = np.zeros(X.shape[1], dtype=np.bool_)
+
+ with nogil:
+ for i in range(n_samples):
+ for j in range(n_features):
+ if isnan_out[j]:
+ continue
+ if isnan(X[i, j]):
+ isnan_out[j] = True
+ break
+ return np.asarray(isnan_out)
diff --git a/causalml/source/causalml/inference/tree/causal/__init__.py b/causalml/source/causalml/inference/tree/causal/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/causalml/source/causalml/inference/tree/causal/_builder.pxd b/causalml/source/causalml/inference/tree/causal/_builder.pxd
new file mode 100644
index 0000000000000000000000000000000000000000..bdd9b4675ee222fdf592d81a19063a36f2c77ea6
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/causal/_builder.pxd
@@ -0,0 +1,11 @@
+# cython: cdivision=True
+# cython: boundscheck=False
+# cython: wraparound=False
+# cython: language_level=3
+# cython: linetrace=True
+
+from .._tree._tree cimport Node, Tree, TreeBuilder
+from .._tree._splitter cimport Splitter, SplitRecord
+from .._tree._typedefs cimport intp_t, int32_t, int64_t, float32_t, float64_t
+from .._tree._tree cimport FrontierRecord, StackRecord
+from .._tree._tree cimport ParentInfo, _init_parent_record
diff --git a/causalml/source/causalml/inference/tree/causal/_builder.pyx b/causalml/source/causalml/inference/tree/causal/_builder.pyx
new file mode 100644
index 0000000000000000000000000000000000000000..f056fa9e1bfda6489c7fc13c5e7cdf41b24c3c61
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/causal/_builder.pyx
@@ -0,0 +1,572 @@
+# distutils: language = c++
+# cython: cdivision=True
+# cython: boundscheck=False
+# cython: wraparound=False
+# cython: language_level=3
+# cython: linetrace=True
+
+
+from libc.stdint cimport INTPTR_MAX
+from libcpp cimport bool
+from libcpp.stack cimport stack
+from libcpp.vector cimport vector
+from libcpp.algorithm cimport pop_heap
+from libcpp.algorithm cimport push_heap
+
+from ._criterion cimport CausalRegressionCriterion
+
+import numpy as np
+cimport numpy as np
+np.import_array()
+
+
+cdef float64_t INFINITY = np.inf
+cdef float64_t EPSILON = np.finfo('double').eps
+
+cdef int IS_FIRST = 1
+cdef int IS_NOT_FIRST = 0
+cdef int IS_LEFT = 1
+cdef int IS_NOT_LEFT = 0
+
+TREE_LEAF = -1
+TREE_UNDEFINED = -2
+cdef intp_t _TREE_LEAF = TREE_LEAF
+cdef intp_t _TREE_UNDEFINED = TREE_UNDEFINED
+
+
+cdef class DepthFirstCausalTreeBuilder(TreeBuilder):
+ """Build a decision tree in depth-first fashion.
+ DepthFirstTreeBuilder modified for causal trees
+ Source: https://github.com/scikit-learn/scikit-learn/blob/main/sklearn/tree/_tree.pyx
+ """
+
+ cdef intp_t min_group_samples
+
+ def __cinit__(self, Splitter splitter, intp_t min_samples_split,
+ intp_t min_samples_leaf, float64_t min_weight_leaf,
+ intp_t max_depth, float64_t min_impurity_decrease,
+ intp_t min_group_samples):
+ self.splitter = splitter
+ self.min_samples_split = min_samples_split
+ self.min_samples_leaf = min_samples_leaf
+ self.min_weight_leaf = min_weight_leaf
+ self.max_depth = max_depth
+ self.min_impurity_decrease = min_impurity_decrease
+ self.min_group_samples = min_group_samples
+
+ cpdef build(self, Tree tree, object X,
+ const float64_t[:, ::1] y,
+ const float64_t[:] sample_weight=None,
+ const unsigned char[::1] missing_values_in_feature_mask=None,
+ ):
+ """Build a decision tree from the training set (X, y)."""
+
+ # check input
+ X, y, sample_weight = self._check_input(X, y, sample_weight)
+
+ # Initial capacity
+ cdef intp_t init_capacity
+
+ if tree.max_depth <= 10:
+ init_capacity = (2 ** (tree.max_depth + 1)) - 1
+ else:
+ init_capacity = 2047
+
+ tree._resize(init_capacity)
+
+ # Parameters
+ cdef Splitter splitter = self.splitter
+ cdef intp_t max_depth = self.max_depth
+ cdef intp_t min_samples_leaf = self.min_samples_leaf
+ cdef float64_t min_weight_leaf = self.min_weight_leaf
+ cdef intp_t min_samples_split = self.min_samples_split
+ cdef float64_t min_impurity_decrease = self.min_impurity_decrease
+ cdef intp_t min_group_samples = self.min_group_samples
+
+ # Recursive partition (without actual recursion)
+ splitter.init(X, y, sample_weight, missing_values_in_feature_mask)
+
+ cdef intp_t start
+ cdef intp_t end
+ cdef intp_t depth
+ cdef intp_t parent
+ cdef bint is_left
+ cdef intp_t n_node_samples = splitter.n_samples
+ cdef float64_t weighted_n_samples = splitter.weighted_n_samples
+ cdef float64_t weighted_n_node_samples
+ cdef SplitRecord split
+ cdef intp_t node_id
+
+ # Groups statistic
+ cdef int64_t tr_count_mean
+ cdef int32_t ct_count
+ cdef int32_t groups_count
+ cdef int32_t min_size
+
+ cdef float64_t middle_value
+ cdef float64_t left_child_min
+ cdef float64_t left_child_max
+ cdef float64_t right_child_min
+ cdef float64_t right_child_max
+ cdef intp_t n_constant_features
+ cdef bint is_leaf
+ cdef bint first = 1
+ cdef intp_t max_depth_seen = -1
+ cdef int rc = 0
+
+ cdef stack[StackRecord] builder_stack
+ cdef StackRecord stack_record
+
+ cdef ParentInfo parent_record
+ _init_parent_record(&parent_record)
+
+ with nogil:
+ # push root node onto stack
+ builder_stack.push({
+ "start": 0,
+ "end": n_node_samples,
+ "depth": 0,
+ "parent": _TREE_UNDEFINED,
+ "is_left": 0,
+ "impurity": INFINITY,
+ "n_constant_features": 0,
+ "lower_bound": -INFINITY,
+ "upper_bound": INFINITY,
+ })
+
+ while not builder_stack.empty():
+ stack_record = builder_stack.top()
+ builder_stack.pop()
+
+ start = stack_record.start
+ end = stack_record.end
+ depth = stack_record.depth
+ parent = stack_record.parent
+ is_left = stack_record.is_left
+ parent_record.impurity = stack_record.impurity
+ parent_record.n_constant_features = stack_record.n_constant_features
+ parent_record.lower_bound = stack_record.lower_bound
+ parent_record.upper_bound = stack_record.upper_bound
+
+ n_node_samples = end - start
+ splitter.node_reset(start, end, &weighted_n_node_samples)
+
+ ( splitter.criterion).get_group_stats(&groups_count, &tr_count_mean, &ct_count, &min_size)
+
+ is_leaf = (depth >= max_depth or
+ n_node_samples < min_samples_split or
+ n_node_samples < 2 * min_samples_leaf or
+ tr_count_mean < min_samples_split // groups_count or
+ ct_count < min_samples_split // groups_count or
+ tr_count_mean < min_samples_leaf or
+ ct_count < min_samples_leaf or
+ min_size < min_group_samples or
+ weighted_n_node_samples < 2 * min_weight_leaf)
+
+ if first:
+ parent_record.impurity = splitter.node_impurity()
+ first = 0
+
+ if not is_leaf:
+ splitter.node_split(&parent_record, &split,)
+
+ is_leaf = (is_leaf or split.pos >= end or
+ (split.improvement + EPSILON < min_impurity_decrease))
+
+ node_id = tree._add_node(parent, is_left, is_leaf, split.feature,
+ split.threshold, parent_record.impurity,
+ n_node_samples, weighted_n_node_samples,
+ split.missing_go_to_left)
+
+ if node_id == INTPTR_MAX:
+ rc = -1
+ break
+
+ # Store value for all nodes, to facilitate tree/model
+ # inspection and interpretation
+ splitter.node_value(tree.value + node_id * tree.value_stride)
+ if splitter.with_monotonic_cst:
+ splitter.clip_node_value(tree.value + node_id * tree.value_stride, parent_record.lower_bound, parent_record.upper_bound)
+
+ if not is_leaf:
+ if (
+ not splitter.with_monotonic_cst or
+ splitter.monotonic_cst[split.feature] == 0
+ ):
+ # Split on a feature with no monotonicity constraint
+
+ # Current bounds must always be propagated to both children.
+ # If a monotonic constraint is active, bounds are used in
+ # node value clipping.
+ left_child_min = right_child_min = parent_record.lower_bound
+ left_child_max = right_child_max = parent_record.upper_bound
+ elif splitter.monotonic_cst[split.feature] == 1:
+ # Split on a feature with monotonic increase constraint
+ left_child_min = parent_record.lower_bound
+ right_child_max = parent_record.upper_bound
+
+ # Lower bound for right child and upper bound for left child
+ # are set to the same value.
+ middle_value = splitter.criterion.middle_value()
+ right_child_min = middle_value
+ left_child_max = middle_value
+ else: # i.e. splitter.monotonic_cst[split.feature] == -1
+ # Split on a feature with monotonic decrease constraint
+ right_child_min = parent_record.lower_bound
+ left_child_max = parent_record.upper_bound
+
+ # Lower bound for left child and upper bound for right child
+ # are set to the same value.
+ middle_value = splitter.criterion.middle_value()
+ left_child_min = middle_value
+ right_child_max = middle_value
+
+ # Push right child on stack
+ builder_stack.push({
+ "start": split.pos,
+ "end": end,
+ "depth": depth + 1,
+ "parent": node_id,
+ "is_left": 0,
+ "impurity": split.impurity_right,
+ "n_constant_features": parent_record.n_constant_features,
+ "lower_bound": right_child_min,
+ "upper_bound": right_child_max,
+ })
+
+ # Push left child on stack
+ builder_stack.push({
+ "start": start,
+ "end": split.pos,
+ "depth": depth + 1,
+ "parent": node_id,
+ "is_left": 1,
+ "impurity": split.impurity_left,
+ "n_constant_features": parent_record.n_constant_features,
+ "lower_bound": left_child_min,
+ "upper_bound": left_child_max,
+ })
+
+ if depth > max_depth_seen:
+ max_depth_seen = depth
+
+ if rc >= 0:
+ rc = tree._resize_c(tree.node_count)
+
+ if rc >= 0:
+ tree.max_depth = max_depth_seen
+ if rc == -1:
+ raise MemoryError()
+
+
+cdef inline bool _compare_records(
+ const FrontierRecord& left,
+ const FrontierRecord& right,
+):
+ return left.improvement < right.improvement
+
+
+cdef inline void _add_to_frontier(
+ FrontierRecord rec,
+ vector[FrontierRecord]& frontier,
+) noexcept nogil:
+ """Adds record `rec` to the priority queue `frontier`."""
+ frontier.push_back(rec)
+ push_heap(frontier.begin(), frontier.end(), &_compare_records)
+
+
+cdef class BestFirstCausalTreeBuilder(TreeBuilder):
+ """Build a decision tree in best-first fashion.
+ The best node to expand is given by the node at the frontier that has the highest impurity improvement.
+ BestFirstCausalTreeBuilder modified for causal trees
+ Source: https://github.com/scikit-learn/scikit-learn/blob/main/sklearn/tree/_tree.pyx
+ """
+ cdef intp_t max_leaf_nodes
+ cdef intp_t min_group_samples
+
+ def __cinit__(self, Splitter splitter, intp_t min_samples_split,
+ intp_t min_samples_leaf, min_weight_leaf,
+ intp_t max_depth, intp_t max_leaf_nodes,
+ float64_t min_impurity_decrease, intp_t min_group_samples):
+ self.splitter = splitter
+ self.min_samples_split = min_samples_split
+ self.min_samples_leaf = min_samples_leaf
+ self.min_weight_leaf = min_weight_leaf
+ self.max_depth = max_depth
+ self.max_leaf_nodes = max_leaf_nodes
+ self.min_impurity_decrease = min_impurity_decrease
+ self.min_group_samples = min_group_samples
+
+ cpdef build(
+ self,
+ Tree tree,
+ object X,
+ const float64_t[:, ::1] y,
+ const float64_t[:] sample_weight=None,
+ const unsigned char[::1] missing_values_in_feature_mask=None,
+ ):
+ """Build a decision tree from the training set (X, y)."""
+
+ # check input
+ X, y, sample_weight = self._check_input(X, y, sample_weight)
+
+
+ # Parameters
+ cdef Splitter splitter = self.splitter
+ cdef intp_t max_leaf_nodes = self.max_leaf_nodes
+ cdef intp_t min_samples_leaf = self.min_samples_leaf
+ cdef float64_t min_weight_leaf = self.min_weight_leaf
+ cdef intp_t min_samples_split = self.min_samples_split
+
+ # Recursive partition (without actual recursion)
+ splitter.init(X, y, sample_weight, missing_values_in_feature_mask)
+
+ cdef vector[FrontierRecord] frontier
+ cdef FrontierRecord record
+ cdef FrontierRecord split_node_left
+ cdef FrontierRecord split_node_right
+ cdef float64_t left_child_min
+ cdef float64_t left_child_max
+ cdef float64_t right_child_min
+ cdef float64_t right_child_max
+
+ cdef intp_t n_node_samples = splitter.n_samples
+ cdef intp_t max_split_nodes = max_leaf_nodes - 1
+ cdef bint is_leaf
+ cdef intp_t max_depth_seen = -1
+ cdef int rc = 0
+ cdef Node* node
+
+ cdef ParentInfo parent_record
+ _init_parent_record(&parent_record)
+
+ # Initial capacity
+ cdef intp_t init_capacity = max_split_nodes + max_leaf_nodes
+ tree._resize(init_capacity)
+
+ with nogil:
+ # add root to frontier
+ rc = self._add_split_node(
+ splitter=splitter,
+ tree=tree,
+ start=0,
+ end=n_node_samples,
+ is_first=IS_FIRST,
+ is_left=IS_LEFT,
+ parent=NULL,
+ depth=0,
+ parent_record=&parent_record,
+ res=&split_node_left,
+ )
+ if rc >= 0:
+ _add_to_frontier(split_node_left, frontier)
+
+ while not frontier.empty():
+ pop_heap(frontier.begin(), frontier.end(), &_compare_records)
+ record = frontier.back()
+ frontier.pop_back()
+
+ node = &tree.nodes[record.node_id]
+ is_leaf = (record.is_leaf or max_split_nodes <= 0)
+
+ if is_leaf:
+ # Node is not expandable; set node as leaf
+ node.left_child = _TREE_LEAF
+ node.right_child = _TREE_LEAF
+ node.feature = _TREE_UNDEFINED
+ node.threshold = _TREE_UNDEFINED
+
+ else:
+ # Node is expandable
+
+ if (
+ not splitter.with_monotonic_cst or
+ splitter.monotonic_cst[node.feature] == 0
+ ):
+ # Split on a feature with no monotonicity constraint
+
+ # Current bounds must always be propagated to both children.
+ # If a monotonic constraint is active, bounds are used in
+ # node value clipping.
+ left_child_min = right_child_min = record.lower_bound
+ left_child_max = right_child_max = record.upper_bound
+ elif splitter.monotonic_cst[node.feature] == 1:
+ # Split on a feature with monotonic increase constraint
+ left_child_min = record.lower_bound
+ right_child_max = record.upper_bound
+
+ # Lower bound for right child and upper bound for left child
+ # are set to the same value.
+ right_child_min = record.middle_value
+ left_child_max = record.middle_value
+ else: # i.e. splitter.monotonic_cst[split.feature] == -1
+ # Split on a feature with monotonic decrease constraint
+ right_child_min = record.lower_bound
+ left_child_max = record.upper_bound
+
+ # Lower bound for left child and upper bound for right child
+ # are set to the same value.
+ left_child_min = record.middle_value
+ right_child_max = record.middle_value
+
+ # Decrement number of split nodes available
+ max_split_nodes -= 1
+
+ # Compute left split node
+ parent_record.lower_bound = left_child_min
+ parent_record.upper_bound = left_child_max
+ parent_record.impurity = record.impurity_left
+ rc = self._add_split_node(
+ splitter=splitter,
+ tree=tree,
+ start=record.start,
+ end=record.pos,
+ is_first=IS_NOT_FIRST,
+ is_left=IS_LEFT,
+ parent=node,
+ depth=record.depth + 1,
+ parent_record=&parent_record,
+ res=&split_node_left,
+ )
+ if rc == -1:
+ break
+
+ # tree.nodes may have changed
+ node = &tree.nodes[record.node_id]
+
+ # Compute right split node
+ parent_record.lower_bound = right_child_min
+ parent_record.upper_bound = right_child_max
+ parent_record.impurity = record.impurity_right
+ rc = self._add_split_node(
+ splitter=splitter,
+ tree=tree,
+ start=record.pos,
+ end=record.end,
+ is_first=IS_NOT_FIRST,
+ is_left=IS_NOT_LEFT,
+ parent=node,
+ depth=record.depth + 1,
+ parent_record=&parent_record,
+ res=&split_node_right,
+ )
+ if rc == -1:
+ break
+
+ # Add nodes to queue
+ _add_to_frontier(split_node_left, frontier)
+ _add_to_frontier(split_node_right, frontier)
+
+ if record.depth > max_depth_seen:
+ max_depth_seen = record.depth
+
+ if rc >= 0:
+ rc = tree._resize_c(tree.node_count)
+
+ if rc >= 0:
+ tree.max_depth = max_depth_seen
+
+ if rc == -1:
+ raise MemoryError()
+
+ cdef inline int _add_split_node(
+ self,
+ Splitter splitter,
+ Tree tree,
+ intp_t start,
+ intp_t end,
+ bint is_first,
+ bint is_left,
+ Node* parent,
+ intp_t depth,
+ ParentInfo* parent_record,
+ FrontierRecord* res
+ ) except -1 nogil:
+ """Adds node w/ partition ``[start, end)`` to the frontier. """
+ cdef SplitRecord split
+ cdef intp_t node_id
+ cdef intp_t n_node_samples
+ cdef float64_t weighted_n_samples = splitter.weighted_n_samples
+ cdef float64_t min_impurity_decrease = self.min_impurity_decrease
+ cdef float64_t weighted_n_node_samples
+ cdef bint is_leaf
+ cdef intp_t n_left, n_right
+ cdef float64_t imp_diff
+
+ splitter.node_reset(start, end, &weighted_n_node_samples)
+
+ # Groups statistic
+ cdef int64_t tr_count_mean
+ cdef int32_t ct_count
+ cdef int32_t groups_count
+ cdef int32_t min_size
+
+ # reset n_constant_features for this specific split before beginning split search
+ parent_record.n_constant_features = 0
+
+ if is_first:
+ parent_record.impurity = splitter.node_impurity()
+
+ ( splitter.criterion).get_group_stats(&groups_count, &tr_count_mean, &ct_count, &min_size)
+
+ n_node_samples = end - start
+ is_leaf = (depth >= self.max_depth or
+ n_node_samples < self.min_samples_split or
+ n_node_samples < 2 * self.min_samples_leaf or
+ tr_count_mean < self.min_samples_split // groups_count or
+ ct_count < self.min_samples_split // groups_count or
+ tr_count_mean < self.min_samples_leaf or
+ ct_count < self.min_samples_leaf or
+ min_size < self.min_group_samples or
+ weighted_n_node_samples < 2 * self.min_weight_leaf or parent_record.impurity <= EPSILON
+ )
+
+ if not is_leaf:
+ splitter.node_split(
+ parent_record,
+ &split
+ )
+ is_leaf = (is_leaf or split.pos >= end or
+ split.improvement + EPSILON < min_impurity_decrease)
+
+ node_id = tree._add_node(parent - tree.nodes
+ if parent != NULL
+ else _TREE_UNDEFINED,
+ is_left, is_leaf,
+ split.feature, split.threshold, parent_record.impurity,
+ n_node_samples, weighted_n_node_samples,
+ split.missing_go_to_left)
+ if node_id == INTPTR_MAX:
+ return -1
+
+ # compute values also for split nodes (might become leafs later).
+ splitter.node_value(tree.value + node_id * tree.value_stride)
+ if splitter.with_monotonic_cst:
+ splitter.clip_node_value(tree.value + node_id * tree.value_stride, parent_record.lower_bound, parent_record.upper_bound)
+
+ res.node_id = node_id
+ res.start = start
+ res.end = end
+ res.depth = depth
+ res.impurity = parent_record.impurity
+ res.lower_bound = parent_record.lower_bound
+ res.upper_bound = parent_record.upper_bound
+ res.middle_value = splitter.criterion.middle_value()
+
+ if not is_leaf:
+ # is split node
+ res.pos = split.pos
+ res.is_leaf = 0
+ res.improvement = split.improvement
+ res.impurity_left = split.impurity_left
+ res.impurity_right = split.impurity_right
+
+ else:
+ # is leaf => 0 improvement
+ res.pos = end
+ res.is_leaf = 1
+ res.improvement = 0.0
+ res.impurity_left = parent_record.impurity
+ res.impurity_right = parent_record.impurity
+
+ return 0
diff --git a/causalml/source/causalml/inference/tree/causal/_criterion.pxd b/causalml/source/causalml/inference/tree/causal/_criterion.pxd
new file mode 100644
index 0000000000000000000000000000000000000000..a99d2d99d9f3ec79186444ea785c09d41e5124a7
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/causal/_criterion.pxd
@@ -0,0 +1,86 @@
+# cython: cdivision=True
+# cython: boundscheck=False
+# cython: wraparound=False
+# cython: language_level=3
+# cython: linetrace=True
+# distutils: language = c++
+
+from libc.math cimport fabs
+from libc.math cimport isnan
+from libc.math cimport sqrt
+from libc.limits cimport INT_MAX
+from libc.string cimport memset
+from libc.string cimport memcpy
+from libc.stdio cimport printf
+
+from libcpp.vector cimport vector
+
+from .._tree._typedefs cimport int32_t, int64_t, intp_t, float64_t
+from .._tree._criterion cimport RegressionCriterion
+
+
+cdef class NodeState:
+ cdef public vector[float64_t] count_1d
+ cdef public vector[float64_t] y_sum_1d
+ cdef public vector[float64_t] y_sq_sum_1d
+ cdef public int32_t control_idx
+ cdef public int32_t control_total
+ cdef public int32_t treatment_total
+ cdef public int32_t groups_total
+ # Criterion-specific variables
+ cdef public float64_t split_metric
+ """
+ NodeState cython class tracks statistics of a control group and multiple test groups
+
+ count_1d: vector[float64_t], the number of observations for a particular group
+ y_sum_1d: vector[float64_t], the sum of y-s for a particular group
+ y_sq_sum_1d: vector[float64_t], the sum of squared y-s for a particular group
+ control_idx: int32_t, control group index
+ control_total int32_t, total number of observations for a control group
+ treatment_total int32_t, total number of observations for treatment groups
+ groups_total int32_t, total number of groups
+ split_metric: float64_t, split metric for TTest criterion
+ """
+
+ cdef int32_t reset(self, intp_t n_outputs) except -1 nogil
+ cdef int32_t update_counters(self) except -1 nogil
+ cdef int32_t copy_from_state(self, NodeState state) except -1 nogil
+ cdef int32_t increment_count(self, int32_t group_idx, float64_t value) except -1 nogil
+ cdef int32_t increment_y_sum(self, int32_t group_idx, float64_t value) except -1 nogil
+ cdef int32_t increment_y_sq_sum(self, int32_t group_idx, float64_t value) except -1 nogil
+ cdef float64_t outcome_mean(self, int32_t group_idx) noexcept nogil
+ cdef float64_t outcome_var(self, int32_t group_idx) noexcept nogil
+ cdef float64_t effect(self, int32_t treatment_idx) noexcept nogil
+
+
+cdef class NodeSplitState:
+ cdef public NodeState node
+ cdef public NodeState right
+ cdef public NodeState left
+
+ """
+ NodeSplitState cython class tracks statistics for the current node and potential left and right splits.
+
+ node: NodeState, current node statistics
+ right: NodeState, right split statistics
+ left: NodeState, left split statistics
+ """
+
+ cdef int32_t reset_nodes(self, intp_t n_outputs) except -1 nogil
+ """
+ Prepare vectors or set existing ones to zero for each NodeState.
+ """
+
+
+cdef class CausalRegressionCriterion(RegressionCriterion):
+
+ cdef public NodeSplitState state
+ cdef public float64_t groups_penalty
+
+ cdef int get_group_stats(
+ self,
+ int32_t* groups_count,
+ int64_t* tr_count_mean,
+ int32_t* ct_count,
+ int32_t* min_size_among_groups) except -1 nogil
+ cdef float64_t get_groups_penalty(self, NodeState node) noexcept nogil
\ No newline at end of file
diff --git a/causalml/source/causalml/inference/tree/causal/_criterion.pyx b/causalml/source/causalml/inference/tree/causal/_criterion.pyx
new file mode 100644
index 0000000000000000000000000000000000000000..7c49599ff64808554a9de3501f3dc0f24d4936aa
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/causal/_criterion.pyx
@@ -0,0 +1,625 @@
+# cython: cdivision=True
+# cython: boundscheck=False
+# cython: wraparound=False
+# cython: language_level=3
+# cython: linetrace=True
+
+
+cdef int32_t CONTROL_GROUP_IDX = 0
+
+
+cdef class NodeState:
+
+ def __cinit__(self):
+ self.split_metric = 1.
+ self.control_idx = CONTROL_GROUP_IDX
+ self.control_total = 0
+ self.treatment_total = 0
+ self.groups_total = 0
+
+ cdef int32_t reset(self, intp_t n_outputs) except -1 nogil:
+
+ if self.count_1d.size() == 0:
+ self.count_1d.resize(n_outputs, 0.)
+ self.y_sum_1d.resize(n_outputs, 0.)
+ self.y_sq_sum_1d.resize(n_outputs, 0.)
+ else:
+ self.count_1d.assign(n_outputs, 0.)
+ self.y_sum_1d.assign(n_outputs, 0.)
+ self.y_sq_sum_1d.assign(n_outputs, 0.)
+
+ self.update_counters()
+ return 0
+
+ cdef int32_t update_counters(self) except -1 nogil:
+
+ cdef int n_outputs = self.count_1d.size()
+
+ if n_outputs == 0:
+ return -1
+
+ self.groups_total = n_outputs
+ self.control_total = self.count_1d[self.control_idx]
+ self.treatment_total = 0
+ for k in range(n_outputs):
+ if k != self.control_idx:
+ self.treatment_total += self.count_1d[k]
+ return 0
+
+ cdef int32_t copy_from_state(self, NodeState state) except -1 nogil:
+
+ if self.count_1d.size() == 0:
+ return -1
+
+ for k in range(self.count_1d.size()):
+ self.count_1d[k] = state.count_1d[k]
+ self.y_sum_1d[k] = state.y_sum_1d[k]
+ self.y_sq_sum_1d[k] = state.y_sq_sum_1d[k]
+ self.update_counters()
+ return 0
+
+ cdef int32_t increment_count(self, int32_t group_idx, float64_t value) except -1 nogil:
+ self.count_1d[group_idx] += value
+ self.update_counters()
+ return 0
+
+ cdef int32_t increment_y_sum(self, int32_t group_idx, float64_t value) except -1 nogil:
+ self.y_sum_1d[group_idx] += value
+ return 0
+
+ cdef int32_t increment_y_sq_sum(self, int32_t group_idx, float64_t value) except -1 nogil:
+ self.y_sq_sum_1d[group_idx] += value
+ return 0
+
+ cdef float64_t outcome_mean(self, int32_t group_idx) noexcept nogil:
+ return self.y_sum_1d[group_idx] / self.count_1d[group_idx]
+
+ cdef float64_t outcome_var(self, int32_t group_idx) noexcept nogil:
+ cdef float64_t var
+ var = (self.y_sq_sum_1d[group_idx] / self.count_1d[group_idx] -
+ (self.y_sum_1d[group_idx] * self.y_sum_1d[group_idx]) / (
+ self.count_1d[group_idx] * self.count_1d[group_idx]))
+ # Clamp tiny negative variance to 0 instead of returning -1
+ var = max(var, 0.0)
+ return var
+
+ cdef float64_t effect(self, int32_t treatment_idx) noexcept nogil:
+ return (self.y_sum_1d[treatment_idx] / self.count_1d[treatment_idx] -
+ self.y_sum_1d[self.control_idx] / self.count_1d[self.control_idx])
+
+
+cdef class NodeSplitState:
+
+ def __cinit__(self, intp_t n_outputs):
+ self.node = NodeState(n_outputs)
+ self.right = NodeState(n_outputs)
+ self.left = NodeState(n_outputs)
+ self.reset_nodes(n_outputs)
+
+ cdef int32_t reset_nodes(self, intp_t n_outputs) except -1 nogil:
+ self.node.reset(n_outputs)
+ self.right.reset(n_outputs)
+ self.left.reset(n_outputs)
+ return 0
+
+
+cdef class CausalRegressionCriterion(RegressionCriterion):
+ """
+ Base class for causal tree criterion
+ """
+
+ def __cinit__(self, intp_t n_outputs, intp_t n_samples):
+ # Parent __cinit__ is automatically called
+ self.state = NodeSplitState(n_outputs)
+
+ cdef int get_group_stats(
+ self,
+ int32_t* groups_count,
+ int64_t* tr_count_mean,
+ int32_t* ct_count,
+ int32_t* min_size_among_groups,
+ ) except -1 nogil:
+
+ cdef int32_t min_size = self.state.node.count_1d[0]
+ for k in range(1, self.n_outputs):
+ min_size = self.state.node.count_1d[k] if self.state.node.count_1d[k] < min_size else min_size
+ cdef int32_t groups = self.state.node.groups_total
+
+ min_size_among_groups[0] = min_size
+ groups_count[0] = groups
+ ct_count[0] = self.state.node.count_1d[self.state.node.control_idx]
+ tr_count_mean[0] = ( ( self.state.node.treatment_total) / ( (groups - 1)) )
+
+ return 0
+
+ cdef int init(
+ self,
+ const float64_t[:, ::1] y,
+ const float64_t[:] sample_weight,
+ float64_t weighted_n_samples,
+ const intp_t[:] sample_indices,
+ intp_t start,
+ intp_t end,
+ ) except -1 nogil:
+ """Initialize the criterion.
+ This initializes the criterion at node sample_indices[start:end] and children
+ sample_indices[start:start] and sample_indices[start:end].
+
+ Notes:
+ 1) self.y[i, k] is nan if a particular observation is not in a group k, k is in range(0, n_outputs - 1).
+ 2) Control group index is fixed to 0 value.
+ 3) Impurity is averaged across the impurity vector calculated for all pairs of
+ control & treatment_i, i is in range(1, n_outputs - 1)
+ """
+ # Initialize fields
+ self.y = y
+ self.sample_weight = sample_weight
+ self.sample_indices = sample_indices
+ self.start = start
+ self.end = end
+ self.n_node_samples = end - start
+ # For compatibility with sklearn functions
+ self.weighted_n_samples = weighted_n_samples
+ self.weighted_n_node_samples = 0.
+
+ cdef intp_t i
+ cdef intp_t p
+ cdef intp_t k
+ cdef float64_t w = 1.0
+ cdef float64_t y_ik
+ cdef float64_t w_y_ik
+
+ memset(&self.sum_total[0], 0, self.n_outputs * sizeof(float64_t))
+ self.sq_sum_total = 0.
+ self.state.reset_nodes(self.n_outputs)
+
+ for p in range(start, end):
+ i = sample_indices[p]
+
+ if sample_weight is not None:
+ w = sample_weight[i]
+
+ # k is the number of groups
+ for k in range(self.n_outputs):
+ y_ik = self.y[i, k]
+
+ if not isnan(y_ik):
+ w_y_ik = w * y_ik
+ self.sum_total[k] += w_y_ik
+ self.sq_sum_total += w_y_ik * y_ik
+ self.weighted_n_node_samples += w
+
+ # Add groups statistics into node state
+ self.state.node.increment_count(k, 1.)
+ self.state.node.increment_y_sum(k, w_y_ik)
+ self.state.node.increment_y_sq_sum(k, w_y_ik * y_ik)
+
+ # Reset to pos=start
+ self.reset()
+ return 0
+
+ cdef int reset(self) except -1 nogil:
+ """Reset the criterion at pos=start."""
+ cdef intp_t n_bytes = self.n_outputs * sizeof(float64_t)
+
+ memset(&self.sum_left[0], 0, n_bytes)
+ memcpy(&self.sum_right[0], &self.sum_total[0], n_bytes)
+
+ self.state.left.reset(self.n_outputs)
+ self.state.right.copy_from_state(self.state.node)
+
+ # For compatibility with sklearn functions
+ self.weighted_n_left = 0.
+ self.weighted_n_right = self.weighted_n_node_samples
+
+ self.pos = self.start
+
+ return 0
+
+ cdef int reverse_reset(self) except -1 nogil:
+ """Reset the criterion at pos=end."""
+ cdef intp_t n_bytes = self.n_outputs * sizeof(float64_t)
+ memset(&self.sum_right[0], 0, n_bytes)
+ memcpy(&self.sum_left[0], &self.sum_total[0], n_bytes)
+
+ self.state.right.reset(self.n_outputs)
+ self.state.left.copy_from_state(self.state.node)
+
+ # For compatibility with sklearn functions
+ self.weighted_n_right = 0.0
+ self.weighted_n_left = self.weighted_n_node_samples
+
+ self.pos = self.end
+
+ return 0
+
+ cdef int update(self, intp_t new_pos) except -1 nogil:
+ """Updated statistics by moving sample_indices[pos:new_pos] to the left."""
+ cdef const float64_t[:] sample_weight = self.sample_weight
+ cdef const intp_t[:] sample_indices = self.sample_indices
+
+ cdef intp_t pos = self.pos
+ cdef intp_t end = self.end
+ cdef intp_t i
+ cdef intp_t p
+ cdef intp_t k = 0
+ cdef float64_t y_ik
+ cdef float64_t w_y_ik
+ cdef float64_t w = 1.0
+
+ """
+ Update statistics up to new_pos
+
+ Given that:
+ sum_total[x] = sum_left[x] + sum_right[x]
+ we are going to update sum_left from the direction that require the least amount of computations,
+ i.e. from pos to new_pos or from end to new_pos
+ """
+ if (new_pos - pos) <= (end - new_pos):
+ for p in range(pos, new_pos):
+ i = sample_indices[p]
+
+ if sample_weight is not None:
+ w = sample_weight[i]
+
+ for k in range(self.n_outputs):
+ y_ik = self.y[i, k]
+ if not isnan(y_ik):
+ w_y_ik = w * y_ik
+ self.sum_left[k] += w_y_ik
+ self.state.left.increment_count(k, 1.)
+ self.state.left.increment_y_sum(k, w_y_ik)
+ self.state.left.increment_y_sq_sum(k, w_y_ik * y_ik)
+
+ self.weighted_n_left += w
+ else:
+ self.reverse_reset()
+
+ for p in range(end - 1, new_pos - 1, -1):
+ i = sample_indices[p]
+
+ if sample_weight is not None:
+ w = sample_weight[i]
+
+ for k in range(self.n_outputs):
+ y_ik = self.y[i, k]
+ if not isnan(y_ik):
+ w_y_ik = w * y_ik
+ self.sum_left[k] -= w_y_ik
+ self.state.left.increment_count(k, -1.)
+ self.state.left.increment_y_sum(k, -w_y_ik)
+ self.state.left.increment_y_sq_sum(k, -w_y_ik * y_ik)
+
+ self.weighted_n_left -= w
+
+ for k in range(self.n_outputs):
+ self.state.right.count_1d[k] = self.state.node.count_1d[k] - self.state.left.count_1d[k]
+ self.state.right.y_sum_1d[k] = self.state.node.y_sum_1d[k] - self.state.left.y_sum_1d[k]
+ self.state.right.y_sq_sum_1d[k] = self.state.node.y_sq_sum_1d[k] - self.state.left.y_sq_sum_1d[k]
+
+ self.sum_right[k] = self.sum_total[k] - self.sum_left[k]
+
+ self.weighted_n_right = self.weighted_n_node_samples - self.weighted_n_left
+ self.pos = new_pos
+
+ return 0
+
+ cdef void node_value(self, float64_t * dest) noexcept nogil:
+ """Compute the node values of sample_indices[start:end] into dest."""
+ cdef intp_t k
+ for k in range(self.n_outputs):
+ dest[k] = self.state.node.outcome_mean(k)
+
+ cdef float64_t get_groups_penalty(self, NodeState node) noexcept nogil:
+ """Compute penalty for sample size differences across multiple treatment groups.
+ Penalizes imbalance of average absolute difference.
+ """
+ cdef intp_t k
+ cdef int32_t groups_total = self.n_outputs
+ cdef int32_t num_treatments = groups_total - 1
+ cdef float64_t fabs_diff_sum = 0.0
+
+ if num_treatments <= 0:
+ return 0.0
+
+ for k in range(groups_total):
+ if k == node.control_idx:
+ continue
+ fabs_diff_sum += fabs(node.count_1d[k] - node.count_1d[CONTROL_GROUP_IDX])
+
+ return self.groups_penalty * (fabs_diff_sum / num_treatments)
+
+
+
+cdef class StandardMSE(CausalRegressionCriterion):
+ """
+ Standard MSE with treatment effect estimates
+ Source: https://github.com/scikit-learn/scikit-learn/blob/main/sklearn/tree/_criterion.pyx
+ """
+
+ cdef float64_t node_impurity(self) noexcept nogil:
+ """Evaluate the impurity of the current node.
+ Evaluate the MSE criterion as impurity of the current node,
+ i.e. the impurity of sample_indices[start:end]. The smaller the impurity the
+ better.
+ """
+ cdef float64_t impurity
+ cdef intp_t k
+
+
+ impurity = self.sq_sum_total / self.n_node_samples
+ for k in range(self.n_outputs):
+ impurity -= (self.sum_total[k] / self.n_node_samples) ** 2.0
+
+ impurity += self.get_groups_penalty(self.state.node)
+
+ return impurity / self.n_outputs
+
+ cdef float64_t proxy_impurity_improvement(self) noexcept nogil:
+ """Compute a proxy of the impurity reduction.
+ This method is used to speed up the search for the best split.
+ It is a proxy quantity such that the split that maximizes this value
+ also maximizes the impurity improvement. It neglects all constant terms
+ of the impurity decrease for a given split.
+ The absolute impurity improvement is only computed by the
+ impurity_improvement method once the best split has been found.
+ The MSE proxy is derived from
+ sum_{i left}(y_i - y_pred_L)^2 + sum_{i right}(y_i - y_pred_R)^2
+ = sum(y_i^2) - n_L * mean_{i left}(y_i)^2 - n_R * mean_{i right}(y_i)^2
+ Neglecting constant terms, this gives:
+ - 1/n_L * sum_{i left}(y_i)^2 - 1/n_R * sum_{i right}(y_i)^2
+ """
+ cdef intp_t k
+ cdef float64_t proxy_impurity_left = 0.0
+ cdef float64_t proxy_impurity_right = 0.0
+ cdef float64_t penalty_left, penalty_right
+
+ penalty_left = self.get_groups_penalty(self.state.left)
+ penalty_right = self.get_groups_penalty(self.state.right)
+
+ for k in range(self.n_outputs):
+ proxy_impurity_left += self.sum_left[k] * self.sum_left[k] - penalty_left
+ proxy_impurity_right += self.sum_right[k] * self.sum_right[k] - penalty_right
+
+ return (proxy_impurity_left / self.weighted_n_left +
+ proxy_impurity_right / self.weighted_n_right)
+
+ cdef void children_impurity(
+ self,
+ float64_t * impurity_left,
+ float64_t * impurity_right
+ ) noexcept nogil:
+ """Evaluate the impurity in children nodes.
+ i.e. the impurity of the left child (sample_indices[start:pos]) and the
+ impurity the right child (sample_indices[pos:end]).
+ """
+ cdef const float64_t[:] sample_weight = self.sample_weight
+ cdef const intp_t[:] sample_indices = self.sample_indices
+ cdef intp_t pos = self.pos
+ cdef intp_t start = self.start
+
+ cdef float64_t y_ik
+
+ cdef float64_t sq_sum_left = 0.0
+ cdef float64_t sq_sum_right
+
+ cdef intp_t i
+ cdef intp_t p
+ cdef intp_t k
+ cdef float64_t w = 1.0
+
+ cdef float64_t penalty_left, penalty_right
+
+ for p in range(start, pos):
+ i = sample_indices[p]
+
+ if sample_weight is not None:
+ w = sample_weight[i]
+
+ for k in range(self.n_outputs):
+ y_ik = self.y[i, k]
+ if not isnan(y_ik):
+ sq_sum_left += w * y_ik * y_ik
+
+ sq_sum_right = self.sq_sum_total - sq_sum_left
+
+ impurity_left[0] = sq_sum_left / self.weighted_n_left
+ impurity_right[0] = sq_sum_right / self.weighted_n_right
+
+ for k in range(self.n_outputs):
+ impurity_left[0] -= (self.sum_left[k] / self.weighted_n_left) ** 2.0
+ impurity_right[0] -= (self.sum_right[k] / self.weighted_n_right) ** 2.0
+
+ impurity_left[0] += self.get_groups_penalty(self.state.left)
+ impurity_right[0] += self.get_groups_penalty(self.state.right)
+
+ impurity_left[0] /= self.n_outputs
+ impurity_right[0] /= self.n_outputs
+
+
+cdef class CausalMSE(CausalRegressionCriterion):
+ """
+ Mean squared error impurity criterion for Causal Tree
+ CausalTreeMSE = right_effect + left_effect
+ where,
+ effect = alpha * tau^2 - (1 - alpha) * (1 + train_to_est_ratio) * (VAR_tr / p + VAR_cont / (1 - p))
+ """
+
+ cdef float64_t node_impurity(self) noexcept nogil:
+ """
+ Evaluate the impurity of the current node, i.e. the impurity of sample_indices[start:end].
+ """
+
+ cdef float64_t impurity = 0.
+ cdef int32_t tr_group_idx
+ cdef float64_t node_tau
+ cdef float64_t tr_var
+ cdef float64_t ct_var = self.state.node.outcome_var(CONTROL_GROUP_IDX)
+ cdef float64_t tr_count
+ cdef float64_t ct_count = self.state.node.count_1d[CONTROL_GROUP_IDX]
+
+ for tr_group_idx in range(1, self.n_outputs):
+ node_tau = self.state.node.effect(tr_group_idx)
+ tr_var = self.state.node.outcome_var(tr_group_idx)
+ tr_count = self.state.node.count_1d[tr_group_idx]
+
+ impurity += (tr_var / tr_count + ct_var / ct_count) - node_tau * node_tau
+
+ impurity /= (self.n_outputs - 1)
+ impurity += self.get_groups_penalty(self.state.node)
+
+ return impurity
+
+ cdef void children_impurity(self, float64_t * impurity_left, float64_t * impurity_right) noexcept nogil:
+ """
+ Evaluate the impurity in children nodes, i.e. the impurity of the
+ left child (sample_indices[start:pos]) and the impurity the right child
+ (sample_indices[pos:end]).
+ """
+
+ cdef float64_t right_tr_var
+ cdef float64_t right_ct_var = self.state.right.outcome_var(CONTROL_GROUP_IDX)
+ cdef float64_t right_tr_count
+ cdef float64_t right_ct_count = self.state.right.count_1d[CONTROL_GROUP_IDX]
+ cdef float64_t left_tr_var
+ cdef float64_t left_ct_var = self.state.left.outcome_var(CONTROL_GROUP_IDX)
+ cdef float64_t left_tr_count
+ cdef float64_t left_ct_count = self.state.left.count_1d[CONTROL_GROUP_IDX]
+ cdef float64_t right_tau
+ cdef float64_t left_tau
+
+ impurity_right[0] = 0.
+ impurity_left[0] = 0.
+
+ for tr_group_idx in range(1, self.n_outputs):
+ right_tau = self.state.right.effect(tr_group_idx)
+ right_tr_var = self.state.right.outcome_var(tr_group_idx)
+ right_tr_count = self.state.right.count_1d[tr_group_idx]
+
+ left_tau = self.state.left.effect(tr_group_idx)
+ left_tr_var = self.state.left.outcome_var(tr_group_idx)
+ left_tr_count = self.state.left.count_1d[tr_group_idx]
+
+ impurity_right[0] += (right_tr_var / right_tr_count + right_ct_var / right_ct_count) - right_tau * right_tau
+ impurity_left[0] += (left_tr_var / left_tr_count + left_ct_var / left_ct_count) - left_tau * left_tau
+
+ impurity_right[0] /= (self.n_outputs - 1)
+ impurity_left[0] /= (self.n_outputs - 1)
+ impurity_right[0] += self.get_groups_penalty(self.state.right)
+ impurity_left[0] += self.get_groups_penalty(self.state.left)
+
+
+cdef class TTest(CausalRegressionCriterion):
+ """
+ TTest impurity criterion for Causal Tree based on "Su, Xiaogang, et al. (2009). Subgroup analysis via recursive partitioning."
+ """
+ cdef float64_t node_impurity(self) noexcept nogil:
+
+
+ cdef float64_t impurity = 0.
+ cdef int32_t tr_group_idx
+ cdef float64_t node_tau
+ cdef float64_t tr_var
+ cdef float64_t ct_var = self.state.node.outcome_var(CONTROL_GROUP_IDX)
+ cdef float64_t tr_count
+ cdef float64_t ct_count = self.state.node.count_1d[CONTROL_GROUP_IDX]
+ cdef float64_t denom
+
+ for tr_group_idx in range(1, self.n_outputs):
+ node_tau = self.state.node.effect(tr_group_idx)
+ tr_var = self.state.node.outcome_var(tr_group_idx)
+ tr_count = self.state.node.count_1d[tr_group_idx]
+ # T statistic of difference between treatment and control means
+ denom = sqrt(( (tr_var / tr_count) + (ct_var / ct_count)))
+ if denom > 0:
+ impurity += node_tau / denom
+
+ return impurity
+
+ cdef void children_impurity(self, float64_t * impurity_left, float64_t * impurity_right) noexcept nogil:
+ """
+ Evaluate the impurity in children nodes, i.e. the impurity of the
+ left child (sample_indices[start:pos]) and the impurity the right child
+ (sample_indices[pos:end]).
+ """
+
+ cdef int32_t tr_group_idx
+ cdef int32_t num_treatments = self.n_outputs - 1
+
+ cdef float64_t t_left_sum = 0.0
+ cdef float64_t t_right_sum = 0.0
+ cdef float64_t tdiff = 0.0
+ cdef float64_t tdiff_sq_sum = 0.0
+
+ cdef float64_t left_tau, right_tau
+ cdef float64_t left_tr_var, right_tr_var
+ cdef float64_t left_ct_var = self.state.left.outcome_var(CONTROL_GROUP_IDX)
+ cdef float64_t right_ct_var = self.state.right.outcome_var(CONTROL_GROUP_IDX)
+
+ cdef float64_t left_tr_count, right_tr_count
+ cdef float64_t left_ct_count = self.state.left.count_1d[CONTROL_GROUP_IDX]
+ cdef float64_t right_ct_count = self.state.right.count_1d[CONTROL_GROUP_IDX]
+
+ cdef float64_t denom_left, denom_right
+ cdef float64_t pooled_var_t
+ cdef float64_t inv_n_sum
+ cdef float64_t dof
+
+ impurity_left[0] = 0.0
+ impurity_right[0] = 0.0
+
+ for tr_group_idx in range(1, self.n_outputs):
+ right_tau = self.state.right.effect(tr_group_idx)
+ right_tr_var = self.state.right.outcome_var(tr_group_idx)
+ right_tr_count = self.state.right.count_1d[tr_group_idx]
+
+ left_tau = self.state.left.effect(tr_group_idx)
+ left_tr_var = self.state.left.outcome_var(tr_group_idx)
+ left_tr_count = self.state.left.count_1d[tr_group_idx]
+
+ denom_left = sqrt(left_tr_var / left_tr_count + left_ct_var / left_ct_count)
+ denom_right = sqrt(right_tr_var / right_tr_count + right_ct_var / right_ct_count)
+ if denom_left > 0.:
+ t_left_sum += left_tau / denom_left
+ if denom_right > 0.:
+ t_right_sum += right_tau / denom_right
+
+ # Per-treatment squared difference in taus between sides
+ inv_n_sum = (1.0 / right_tr_count + 1.0 / right_ct_count +
+ 1.0 / left_tr_count + 1.0 / left_ct_count)
+
+ # Pooled variance across four cells (left/right × tr/ct)
+ pooled_var_t = 0.0
+ pooled_var_t += ((right_tr_count - 1.0) * right_tr_var)
+ pooled_var_t += ((right_ct_count - 1.0) * right_ct_var)
+ pooled_var_t += ((left_tr_count - 1.0) * left_tr_var)
+ pooled_var_t += ((left_ct_count - 1.0) * left_ct_var)
+
+ # Normalize by total degrees of freedom if it is positive
+ dof = (right_tr_count - 1.0) + (right_ct_count - 1.0) + (left_tr_count - 1.0) + (left_ct_count - 1.0)
+ if dof > 0.0:
+ pooled_var_t /= dof
+
+ if pooled_var_t > 0.0 and inv_n_sum > 0.0:
+ tdiff = ((left_tau - right_tau) / (( sqrt(pooled_var_t) ) * ( sqrt(inv_n_sum) )))
+ tdiff_sq_sum += (tdiff * tdiff)
+
+ self.state.left.split_metric = (tdiff_sq_sum / num_treatments) + self.get_groups_penalty(self.state.node)
+
+ impurity_left[0] = t_left_sum / num_treatments
+ impurity_right[0] = t_right_sum / num_treatments
+
+ cdef float64_t impurity_improvement(self, float64_t impurity_parent,
+ float64_t impurity_left,
+ float64_t impurity_right) noexcept nogil:
+ return self.state.left.split_metric
+
+ cdef float64_t proxy_impurity_improvement(self) noexcept nogil:
+ """Compute a proxy of the impurity reduction. In case of t statistic - proxy_impurity_improvement
+ is the same as impurity_improvement.
+ """
+ cdef float64_t impurity_left
+ cdef float64_t impurity_right
+ self.children_impurity(&impurity_left, &impurity_right)
+
+ return self.state.left.split_metric
diff --git a/causalml/source/causalml/inference/tree/causal/_tree.py b/causalml/source/causalml/inference/tree/causal/_tree.py
new file mode 100644
index 0000000000000000000000000000000000000000..5d09e2f212ce704dc4e54efa3322f74a5e59d41b
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/causal/_tree.py
@@ -0,0 +1,273 @@
+import copy
+import numbers
+import warnings
+from math import ceil
+from typing import Union
+
+try:
+ from packaging.version import parse as Version
+except ModuleNotFoundError:
+ from distutils.version import LooseVersion as Version
+
+import numpy as np
+from scipy.sparse import issparse
+from sklearn import __version__ as sklearn_version
+from sklearn.utils import check_random_state
+from sklearn.utils.validation import _check_sample_weight, validate_data
+
+from .._tree._classes import DTYPE, DOUBLE
+from .._tree._classes import SPARSE_SPLITTERS, DENSE_SPLITTERS
+from .._tree._classes import Tree, BaseDecisionTree
+from .._tree._criterion import Criterion
+from .._tree._splitter import Splitter
+
+from ._builder import DepthFirstCausalTreeBuilder, BestFirstCausalTreeBuilder
+from ._criterion import StandardMSE, CausalMSE, TTest
+
+CAUSAL_TREES_CRITERIA = {
+ "causal_mse": CausalMSE,
+ "standard_mse": StandardMSE,
+ "t_test": TTest,
+}
+
+
+def get_check_y_params() -> dict:
+ """
+ Prepares flags for sklearn 1.6+.
+
+ Returns: check_y_params
+ """
+ check_y_params = dict(ensure_2d=False, dtype=None, ensure_all_finite=False)
+ return check_y_params
+
+
+class BaseCausalDecisionTree(BaseDecisionTree):
+ """
+ Modified base class BaseDecisionTree for causal trees
+ Source: https://github.com/scikit-learn/scikit-learn/blob/main/sklearn/tree/_classes.py
+ """
+
+ def __init__(self, min_group_samples: int, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.min_group_samples = min_group_samples
+
+ def _support_missing_values(self, X) -> bool:
+ """
+ TODO: Add support for missing values
+ See sklearn PR: ENH Adds missing value support for trees (#23595)
+ https://github.com/scikit-learn/scikit-learn/commit/6392148d80e9f14a9524c137ac5cfa04f2274d48
+ """
+ return False
+
+ def fit(
+ self,
+ X: np.ndarray,
+ y: np.ndarray,
+ sample_weight: Union[np.ndarray, None] = None,
+ check_input: bool = True,
+ X_idx_sorted="deprecated",
+ ):
+ random_state = check_random_state(self.random_state)
+
+ if self.ccp_alpha < 0.0:
+ raise ValueError("ccp_alpha must be greater than or equal to 0")
+
+ if check_input:
+ # Need to validate separately here.
+ # We can't pass multi_ouput=True because that would allow y to be csr.
+ check_X_params = dict(dtype=DTYPE, accept_sparse="csc")
+ check_y_params = get_check_y_params()
+ X, y = validate_data(
+ self, X, y, validate_separately=(check_X_params, check_y_params)
+ )
+ if issparse(X):
+ X.sort_indices()
+
+ if X.indices.dtype != np.intc or X.indptr.dtype != np.intc:
+ raise ValueError(
+ "No support for np.int64 index based " "sparse matrices"
+ )
+
+ if self.criterion not in CAUSAL_TREES_CRITERIA.keys():
+ raise ValueError(
+ f"Only {CAUSAL_TREES_CRITERIA.keys()} criteria are supported"
+ )
+
+ n_samples, self.n_features_ = X.shape
+ self.n_features_in_ = self.n_features_
+
+ y = np.atleast_1d(y)
+ expanded_class_weight = None
+
+ # n_outputs_ is the length of [y|control, y|treatment_1,..., y|treatment_{n-1}]
+ self.n_outputs_ = y.shape[1]
+
+ if getattr(y, "dtype", None) != DOUBLE or not y.flags.contiguous:
+ y = np.ascontiguousarray(y, dtype=DOUBLE)
+
+ # Check parameters
+ max_depth = np.iinfo(np.int32).max if self.max_depth is None else self.max_depth
+ max_leaf_nodes = -1 if self.max_leaf_nodes is None else self.max_leaf_nodes
+
+ if isinstance(self.min_samples_leaf, numbers.Integral):
+ if not 1 <= self.min_samples_leaf:
+ raise ValueError(
+ "min_samples_leaf must be at least 1 "
+ "or in (0, 0.5], got %s" % self.min_samples_leaf
+ )
+ min_samples_leaf = self.min_samples_leaf
+ else: # float
+ if not 0.0 < self.min_samples_leaf <= 0.5:
+ raise ValueError(
+ "min_samples_leaf must be at least 1 "
+ "or in (0, 0.5], got %s" % self.min_samples_leaf
+ )
+ min_samples_leaf = int(ceil(self.min_samples_leaf * n_samples))
+
+ if isinstance(self.min_samples_split, numbers.Integral):
+ if not 2 <= self.min_samples_split:
+ raise ValueError(
+ "min_samples_split must be an integer "
+ "greater than 1 or a float in (0.0, 1.0]; "
+ "got the integer %s" % self.min_samples_split
+ )
+ min_samples_split = self.min_samples_split
+ else: # float
+ if not 0.0 < self.min_samples_split <= 1.0:
+ raise ValueError(
+ "min_samples_split must be an integer "
+ "greater than 1 or a float in (0.0, 1.0]; "
+ "got the float %s" % self.min_samples_split
+ )
+ min_samples_split = int(ceil(self.min_samples_split * n_samples))
+ min_samples_split = max(2, min_samples_split)
+
+ min_samples_split = max(min_samples_split, 2 * min_samples_leaf)
+
+ if isinstance(self.max_features, str):
+ if self.max_features == "auto":
+ max_features = self.n_features_
+ elif self.max_features == "sqrt":
+ max_features = max(1, int(np.sqrt(self.n_features_)))
+ elif self.max_features == "log2":
+ max_features = max(1, int(np.log2(self.n_features_)))
+ else:
+ raise ValueError(
+ "Invalid value for max_features. "
+ "Allowed string values are 'auto', "
+ "'sqrt' or 'log2'."
+ )
+ elif self.max_features is None:
+ max_features = self.n_features_
+ elif isinstance(self.max_features, numbers.Integral):
+ max_features = self.max_features
+ else: # float
+ if self.max_features > 0.0:
+ max_features = max(1, int(self.max_features * self.n_features_))
+ else:
+ max_features = 0
+
+ self.max_features_ = max_features
+
+ if len(y) != n_samples:
+ raise ValueError(
+ "Number of labels=%d does not match "
+ "number of samples=%d" % (len(y), n_samples)
+ )
+ if not 0 <= self.min_weight_fraction_leaf <= 0.5:
+ raise ValueError("min_weight_fraction_leaf must in [0, 0.5]")
+ if max_depth <= 0:
+ raise ValueError("max_depth must be greater than zero. ")
+ if not (0 < max_features <= self.n_features_):
+ raise ValueError("max_features must be in (0, n_features]")
+ if not isinstance(max_leaf_nodes, numbers.Integral):
+ raise ValueError(
+ "max_leaf_nodes must be integral number but was " "%r" % max_leaf_nodes
+ )
+ if -1 < max_leaf_nodes < 2:
+ raise ValueError(
+ ("max_leaf_nodes {0} must be either None " "or larger than 1").format(
+ max_leaf_nodes
+ )
+ )
+
+ if sample_weight is not None:
+ sample_weight = _check_sample_weight(sample_weight, X, dtype=DOUBLE)
+
+ if expanded_class_weight is not None:
+ if sample_weight is not None:
+ sample_weight = sample_weight * expanded_class_weight
+ else:
+ sample_weight = expanded_class_weight
+
+ # Set min_weight_leaf from min_weight_fraction_leaf
+ if sample_weight is None:
+ min_weight_leaf = self.min_weight_fraction_leaf * n_samples
+ else:
+ min_weight_leaf = self.min_weight_fraction_leaf * np.sum(sample_weight)
+
+ if X_idx_sorted != "deprecated":
+ warnings.warn(
+ "The parameter 'X_idx_sorted' is deprecated and has no "
+ "effect. It will be removed in 1.1 (renaming of 0.26). You "
+ "can suppress this warning by not passing any value to the "
+ "'X_idx_sorted' parameter.",
+ FutureWarning,
+ )
+
+ # Build tree
+ criterion = self.criterion
+ if isinstance(criterion, str):
+ criterion = CAUSAL_TREES_CRITERIA[criterion](self.n_outputs_, n_samples)
+ criterion.groups_penalty = self.groups_penalty
+ else:
+ # Make a deepcopy in case the criterion has mutable attributes that
+ # might be shared and modified concurrently during parallel fitting
+ criterion = copy.deepcopy(criterion)
+
+ SPLITTERS = SPARSE_SPLITTERS if issparse(X) else DENSE_SPLITTERS
+
+ splitter = self.splitter
+ if not isinstance(self.splitter, Splitter):
+ splitter = SPLITTERS[self.splitter](
+ criterion,
+ self.max_features_,
+ min_samples_leaf,
+ min_weight_leaf,
+ random_state,
+ monotonic_cst=None,
+ )
+ self.tree_ = Tree(
+ self.n_features_,
+ np.array([1] * self.n_outputs_, dtype=np.intp),
+ self.n_outputs_,
+ )
+
+ # Use BestFirst if max_leaf_nodes given; use DepthFirst otherwise
+ if max_leaf_nodes < 0:
+ builder = DepthFirstCausalTreeBuilder(
+ splitter,
+ min_samples_split,
+ min_samples_leaf,
+ min_weight_leaf,
+ max_depth,
+ self.min_impurity_decrease,
+ self.min_group_samples,
+ )
+ else:
+ builder = BestFirstCausalTreeBuilder(
+ splitter,
+ min_samples_split,
+ min_samples_leaf,
+ min_weight_leaf,
+ max_depth,
+ max_leaf_nodes,
+ self.min_impurity_decrease,
+ self.min_group_samples,
+ )
+ # Treatment column is described via y cols. The first column is always a control group.
+ builder.build(self.tree_, X, y, sample_weight)
+
+ self._prune_tree()
+
+ return self
diff --git a/causalml/source/causalml/inference/tree/causal/causalforest.py b/causalml/source/causalml/inference/tree/causal/causalforest.py
new file mode 100644
index 0000000000000000000000000000000000000000..8bb7017cff830cc5c58bfe64af60a8bf5d7e666b
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/causal/causalforest.py
@@ -0,0 +1,515 @@
+from typing import Union
+
+import numpy as np
+import forestci as fci
+from joblib import Parallel, delayed
+from warnings import catch_warnings, simplefilter, warn
+
+from sklearn.exceptions import DataConversionWarning
+from sklearn.utils.validation import (
+ check_random_state,
+ _check_sample_weight,
+ validate_data,
+)
+from sklearn.utils.multiclass import type_of_target
+from sklearn import __version__ as sklearn_version
+from sklearn.ensemble._forest import DOUBLE, DTYPE, MAX_INT
+from sklearn.ensemble._forest import ForestRegressor
+from sklearn.ensemble._forest import compute_sample_weight, issparse
+from sklearn.ensemble._forest import _generate_sample_indices, _get_n_samples_bootstrap
+
+from .causaltree import CausalTreeRegressor
+from ._tree import get_check_y_params
+
+try:
+ from packaging.version import parse as Version
+except ModuleNotFoundError:
+ from distutils.version import LooseVersion as Version
+
+if Version(sklearn_version) >= Version("1.1.0"):
+ _joblib_parallel_args = dict(prefer="threads")
+else:
+ from sklearn.utils.fixes import _joblib_parallel_args
+
+ _joblib_parallel_args = _joblib_parallel_args(prefer="threads")
+
+
+def _parallel_build_trees(
+ tree,
+ forest,
+ X,
+ treatment,
+ y,
+ sample_weight,
+ tree_idx,
+ n_trees,
+ verbose=0,
+ class_weight=None,
+ n_samples_bootstrap=None,
+):
+ """
+ Private function used to fit a single tree in parallel."""
+ if verbose > 1:
+ print("building tree %d of %d" % (tree_idx + 1, n_trees))
+
+ if forest.bootstrap:
+ n_samples = X.shape[0]
+ if sample_weight is None:
+ curr_sample_weight = np.ones((n_samples,), dtype=np.float64)
+ else:
+ curr_sample_weight = sample_weight.copy()
+
+ indices = _generate_sample_indices(
+ tree.random_state, n_samples, n_samples_bootstrap
+ )
+ sample_counts = np.bincount(indices, minlength=n_samples)
+ curr_sample_weight *= sample_counts
+
+ if class_weight == "subsample":
+ with catch_warnings():
+ simplefilter("ignore", DeprecationWarning)
+ curr_sample_weight *= compute_sample_weight("auto", y, indices=indices)
+ elif class_weight == "balanced_subsample":
+ curr_sample_weight *= compute_sample_weight("balanced", y, indices=indices)
+
+ tree.fit(
+ X,
+ treatment,
+ y,
+ sample_weight=curr_sample_weight,
+ check_input=True,
+ prepare_data=False,
+ )
+ else:
+ tree.fit(
+ X,
+ treatment,
+ y,
+ sample_weight=sample_weight,
+ check_input=True,
+ prepare_data=False,
+ )
+
+ return tree
+
+
+class CausalRandomForestRegressor(ForestRegressor):
+ def __init__(
+ self,
+ n_estimators: int = 100,
+ *,
+ control_name: Union[int, str] = 0,
+ criterion: str = "causal_mse",
+ alpha: float = 0.05,
+ max_depth: int = None,
+ min_samples_split: int = 60,
+ min_samples_leaf: int = 100,
+ min_group_samples: int = 50,
+ min_weight_fraction_leaf: float = 0.0,
+ max_features: Union[int, float, str] = 1.0,
+ max_leaf_nodes: int = None,
+ min_impurity_decrease: float = float("-inf"),
+ bootstrap: bool = True,
+ oob_score: bool = False,
+ n_jobs: int = None,
+ random_state: int = None,
+ verbose: int = 0,
+ warm_start: bool = False,
+ ccp_alpha: float = 0.0,
+ groups_penalty: float = 0.5,
+ max_samples: int = None,
+ groups_cnt: bool = True,
+ groups_cnt_mode: str = "nodes",
+ ):
+ """
+ Initialize Random Forest of CausalTreeRegressors
+
+ Args:
+ n_estimators: (int, default=100)
+ Number of trees in the forest
+ control_name: (str or int)
+ Name of control group
+ criterion: ({"causal_mse", "standard_mse"}, default="causal_mse"):
+ Function to measure the quality of a split.
+ alpha: (float)
+ The confidence level alpha of the ATE estimate and ITE bootstrap estimates
+ max_depth: (int, default=None)
+ The maximum depth of the tree.
+ min_samples_split: (int or float, default=2)
+ The minimum number of samples required to split an internal node:
+ min_samples_leaf: (int or float), default=100
+ The minimum number of samples required to be at a leaf node.
+ min_weight_fraction_leaf: (float, default=0.0)
+ The minimum weighted fraction of the sum total of weights (of all
+ the input samples) required to be at a leaf node.
+ max_features: (int, float or {"auto", "sqrt", "log2"}, default=None)
+ The number of features to consider when looking for the best split
+ max_leaf_nodes: (int, default=None)
+ Grow a tree with ``max_leaf_nodes`` in best-first fashion.
+ min_impurity_decrease: (float, default=float("-inf")))
+ A node will be split if this split induces a decrease of the impurity
+ greater than or equal to this value.
+ bootstrap : (bool, default=True)
+ Whether bootstrap samples are used when building trees.
+ oob_score : bool, default=False
+ Whether to use out-of-bag samples to estimate the generalization score.
+ n_jobs : int, default=None
+ The number of jobs to run in parallel.
+ random_state : (int, RandomState instance or None, default=None)
+ Controls both the randomness of the bootstrapping of the samples used
+ when building trees (if ``bootstrap=True``) and the sampling of the
+ features to consider when looking for the best split at each node
+ (if ``max_features < n_features``).
+ verbose : (int, default=0)
+ Controls the verbosity when fitting and predicting.
+ warm_start : (bool, default=False)
+ When set to ``True``, reuse the solution of the previous call to fit
+ and add more estimators to the ensemble, otherwise, just fit a whole
+ new forest.
+ ccp_alpha : (non-negative float, default=0.0)
+ Complexity parameter used for Minimal Cost-Complexity Pruning.
+ groups_penalty: (float, default=0.5)
+ This penalty coefficient manages the node impurity increase in case of the difference between
+ treatment and control samples sizes.
+ max_samples : (int or float, default=None)
+ If bootstrap is True, the number of samples to draw from X
+ to train each base estimator.
+ groups_cnt: (bool), count treatment and control groups for each node/leaf
+ groups_cnt_mode: (str, 'nodes', 'leaves'), mode for samples counting
+ """
+ self._estimator = CausalTreeRegressor(
+ control_name=control_name,
+ criterion=criterion,
+ groups_cnt=groups_cnt,
+ groups_cnt_mode=groups_cnt_mode,
+ )
+
+ _estimator_key = (
+ "estimator"
+ if Version(sklearn_version) >= Version("1.2.0")
+ else "base_estimator"
+ )
+ _parent_args = {
+ _estimator_key: self._estimator,
+ "n_estimators": n_estimators,
+ "estimator_params": (
+ "criterion",
+ "control_name",
+ "max_depth",
+ "min_samples_split",
+ "min_weight_fraction_leaf",
+ "max_features",
+ "max_leaf_nodes",
+ "min_impurity_decrease",
+ "ccp_alpha",
+ "groups_penalty",
+ "min_samples_leaf",
+ "min_group_samples",
+ "random_state",
+ ),
+ "bootstrap": bootstrap,
+ "oob_score": oob_score,
+ "n_jobs": n_jobs,
+ "random_state": random_state,
+ "verbose": verbose,
+ "warm_start": warm_start,
+ "max_samples": max_samples,
+ }
+
+ super().__init__(**_parent_args)
+
+ self.criterion = criterion
+ self.control_name = control_name
+ self.max_depth = max_depth
+ self.min_samples_split = min_samples_split
+ self.min_samples_leaf = min_samples_leaf
+ self.min_group_samples = min_group_samples
+ self.min_weight_fraction_leaf = min_weight_fraction_leaf
+ self.max_features = max_features
+ self.max_leaf_nodes = max_leaf_nodes
+ self.min_impurity_decrease = min_impurity_decrease
+ self.ccp_alpha = ccp_alpha
+ self.groups_penalty = groups_penalty
+ self.alpha = alpha
+ self.groups_cnt = groups_cnt
+ self.groups_cnt_mode = groups_cnt_mode
+
+ def _fit(
+ self,
+ X: np.ndarray,
+ treatment: np.ndarray,
+ y: np.ndarray,
+ sample_weight: np.ndarray = None,
+ ):
+ """
+ Build a forest of trees from the training set (X, y).
+ With modified _parallel_build_trees for Causal Trees used in BaseForest.fit()
+ Source: https://github.com/scikit-learn/scikit-learn/blob/main/sklearn/ensemble/_forest.py
+
+ Parameters
+ ----------
+ X (np.ndarray): {array-like, sparse matrix} of shape (n_samples, n_features)
+ The training input samples. Internally, its dtype will be converted
+ to ``dtype=np.float32``. If a sparse matrix is provided, it will be
+ converted into a sparse ``csc_matrix``.
+
+ treatment (np.ndarray): treatment vector, includes control group
+
+ y (np.ndarray): array-like of shape (n_samples,) or (n_samples, n_outputs)
+ The target values (class labels in classification, real numbers in
+ regression).
+
+ sample_weight (np.ndarray): array-like of shape (n_samples,), default=None
+ Sample weights. If None, then samples are equally weighted. Splits
+ that would create child nodes with net zero or negative weight are
+ ignored while searching for a split in each node. In the case of
+ classification, splits are also ignored if they would result in any
+ single class carrying a negative weight in either child node.
+
+ Returns
+ -------
+ self : object
+ Fitted estimator.
+ """
+ # Validate or convert input data
+ if issparse(y):
+ raise ValueError("sparse multilabel-indicator for y is not supported.")
+ check_X_params = dict(dtype=DTYPE, accept_sparse="csc")
+ check_y_params = get_check_y_params()
+ X, y = validate_data(
+ self,
+ X,
+ y,
+ multi_output=True,
+ accept_sparse="csc",
+ validate_separately=(check_X_params, check_y_params),
+ )
+ if sample_weight is not None:
+ sample_weight = _check_sample_weight(sample_weight, X)
+
+ if issparse(X):
+ # Pre-sort indices to avoid that each individual tree of the
+ # ensemble sorts the indices.
+ X.sort_indices()
+
+ y = np.atleast_1d(y)
+ if y.ndim == 2 and y.shape[1] == 1:
+ warn(
+ "A column-vector y was passed when a 1d array was"
+ " expected. Please change the shape of y to "
+ "(n_samples,), for example using ravel().",
+ DataConversionWarning,
+ stacklevel=2,
+ )
+
+ if y.ndim == 1:
+ y = np.reshape(y, (-1, 1))
+
+ if self.criterion == "poisson":
+ if np.any(y < 0):
+ raise ValueError(
+ "Some value(s) of y are negative which is "
+ "not allowed for Poisson regression."
+ )
+ if np.sum(y) <= 0:
+ raise ValueError(
+ "Sum of y is not strictly positive which "
+ "is necessary for Poisson regression."
+ )
+ groups = np.unique(treatment).astype(int).size
+ self.n_outputs_ = groups - 1
+ self.max_outputs_ = self.n_outputs_ + groups
+ y, expanded_class_weight = self._validate_y_class_weight(y)
+
+ if getattr(y, "dtype", None) != DOUBLE or not y.flags.contiguous:
+ y = np.ascontiguousarray(y, dtype=DOUBLE)
+
+ if expanded_class_weight is not None:
+ if sample_weight is not None:
+ sample_weight = sample_weight * expanded_class_weight
+ else:
+ sample_weight = expanded_class_weight
+
+ if not self.bootstrap and self.max_samples is not None:
+ raise ValueError(
+ "`max_sample` cannot be set if `bootstrap=False`. "
+ "Either switch to `bootstrap=True` or set "
+ "`max_sample=None`."
+ )
+ elif self.bootstrap:
+ n_samples_bootstrap = _get_n_samples_bootstrap(
+ n_samples=X.shape[0], max_samples=self.max_samples
+ )
+ else:
+ n_samples_bootstrap = None
+
+ # Check parameters
+ self._validate_estimator()
+
+ if not self.bootstrap and self.oob_score:
+ raise ValueError("Out of bag estimation only available if bootstrap=True")
+
+ random_state = check_random_state(self.random_state)
+
+ if not self.warm_start or not hasattr(self, "estimators_"):
+ # Free allocated memory, if any
+ self.estimators_ = []
+
+ n_more_estimators = self.n_estimators - len(self.estimators_)
+
+ if n_more_estimators < 0:
+ raise ValueError(
+ "n_estimators=%d must be larger or equal to "
+ "len(estimators_)=%d when warm_start==True"
+ % (self.n_estimators, len(self.estimators_))
+ )
+
+ elif n_more_estimators == 0:
+ warn(
+ "Warm-start fitting without increasing n_estimators does not "
+ "fit new trees."
+ )
+ else:
+ if self.warm_start and len(self.estimators_) > 0:
+ # We draw from the random state to get the random state we
+ # would have got if we hadn't used a warm_start.
+ random_state.randint(MAX_INT, size=len(self.estimators_))
+
+ trees = [
+ self._make_estimator(append=False, random_state=random_state)
+ for _ in range(n_more_estimators)
+ ]
+ trees = Parallel(
+ n_jobs=self.n_jobs,
+ verbose=self.verbose,
+ **_joblib_parallel_args,
+ )(
+ delayed(_parallel_build_trees)(
+ tree=t,
+ forest=self,
+ X=X,
+ treatment=treatment,
+ y=y,
+ sample_weight=sample_weight,
+ tree_idx=i,
+ n_trees=len(trees),
+ verbose=self.verbose,
+ class_weight=self.class_weight,
+ n_samples_bootstrap=n_samples_bootstrap,
+ )
+ for i, t in enumerate(trees)
+ )
+
+ self.estimators_.extend(trees)
+
+ if self.oob_score:
+ y_type = type_of_target(y)
+ if y_type in ("multiclass-multioutput", "unknown"):
+ raise ValueError(
+ "The type of target cannot be used to compute OOB "
+ f"estimates. Got {y_type} while only the following are "
+ "supported: continuous, continuous-multioutput, binary, "
+ "multiclass, multilabel-indicator."
+ )
+ self._set_oob_score_and_attributes(X, y)
+
+ if hasattr(self, "classes_") and self.n_outputs_ == 1:
+ self.n_classes_ = self.n_classes_[0]
+ self.classes_ = self.classes_[0]
+
+ return self
+
+ def fit(
+ self,
+ X: np.ndarray,
+ treatment: np.ndarray,
+ y: np.ndarray,
+ sample_weight: np.ndarray = None,
+ ):
+ """
+ Fit Causal RandomForest
+ Args:
+ X: (np.ndarray), feature matrix
+ treatment: (np.ndarray), treatment vector
+ y: (np.ndarray), outcome vector
+ sample_weight: (np.ndarray), sample weights
+ Returns:
+ self
+ """
+ X, y = self._estimator._prepare_data(X=X, treatment=treatment, y=y)
+ return self._fit(X=X, treatment=treatment, y=y, sample_weight=sample_weight)
+
+ def predict(self, X: np.ndarray, with_outcomes: bool = False) -> np.ndarray:
+ """Predict individual treatment effects
+
+ Args:
+ X (np.ndarray): a feature matrix
+ with_outcomes (bool), default=False,
+ include outcomes Y_hat(X|T=0), Y_hat(X|T=1) along with individual treatment effect
+ Returns:
+ (np.ndarray): individual treatment effect (ITE), dim=(samples, groups-1)
+ or ITE with outcomes:
+ [Y_hat(X|T=0), Y_hat(X|T=1),...,Y_hat(X|T=n), ITE_1, ITE_2,...,ITE_n], dim=(samples, 2*groups-1)
+ """
+ if with_outcomes:
+ self.n_outputs_ = self.max_outputs_
+ for estimator in self.estimators_:
+ estimator._with_outcomes = True
+ y_pred = super().predict(X)
+ return y_pred
+
+ def calculate_error(
+ self,
+ X_train: np.ndarray,
+ X_test: np.ndarray,
+ inbag: np.ndarray = None,
+ calibrate: bool = True,
+ memory_constrained: bool = False,
+ memory_limit: int = None,
+ ) -> np.ndarray:
+ """
+ Calculate error bars from scikit-learn RandomForest estimators
+ Source:
+ https://github.com/scikit-learn-contrib/forest-confidence-interval
+
+ Args:
+ X_train: (np.ndarray), training subsample of feature matrix, (n_train_sample, n_features)
+ X_test: (np.ndarray), test subsample of feature matrix, (n_train_sample, n_features)
+ inbag: (ndarray, optional),
+ The inbag matrix that fit the data. If set to `None` (default) it
+ will be inferred from the forest. However, this only works for trees
+ for which bootstrapping was set to `True`. That is, if sampling was
+ done with replacement. Otherwise, users need to provide their own
+ inbag matrix.
+ calibrate: (boolean, optional)
+ Whether to apply calibration to mitigate Monte Carlo noise.
+ Some variance estimates may be negative due to Monte Carlo effects if
+ the number of trees in the forest is too small. To use calibration,
+ Default: True
+ memory_constrained: (boolean, optional)
+ Whether or not there is a restriction on memory. If False, it is
+ assumed that a ndarray of shape (n_train_sample,n_test_sample) fits
+ in main memory. Setting to True can actually provide a speedup if
+ memory_limit is tuned to the optimal range.
+ memory_limit: (int, optional)
+ An upper bound for how much memory the intermediate matrices will take
+ up in Megabytes. This must be provided if memory_constrained=True.
+
+ Returns:
+ (np.ndarray), An array with the unbiased sampling variance for a RandomForest object.
+ """
+ if self.n_outputs_ != 1:
+ raise NotImplementedError(
+ f"forestci supports n_outputs=1. n_outputs={self.n_outputs_}"
+ )
+
+ var = fci.random_forest_error(
+ self,
+ X_train,
+ X_test,
+ inbag=inbag,
+ calibrate=calibrate,
+ memory_constrained=memory_constrained,
+ memory_limit=memory_limit,
+ )
+ return var
diff --git a/causalml/source/causalml/inference/tree/causal/causaltree.py b/causalml/source/causalml/inference/tree/causal/causaltree.py
new file mode 100644
index 0000000000000000000000000000000000000000..d45ba7db70d6f07dd6aad894c10b04773e0289a6
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/causal/causaltree.py
@@ -0,0 +1,443 @@
+import logging
+from typing import Union
+
+import tqdm
+import numpy as np
+from numpy import float32 as DTYPE
+
+from pathos.pools import ProcessPool as PPool
+from scipy.stats import norm
+from sklearn.base import RegressorMixin
+from sklearn.utils import check_array
+from sklearn.utils.validation import check_is_fitted
+
+from causalml.inference.meta.utils import check_treatment_vector
+
+from ._tree import BaseCausalDecisionTree
+from ..utils import get_tree_leaves_mask, timeit
+
+logger = logging.getLogger("causalml")
+
+
+class CausalTreeRegressor(RegressorMixin, BaseCausalDecisionTree):
+ """A Causal Tree regressor class.
+ The Causal Tree is a decision tree regressor with a split criteria for treatment effects.
+ Details are available at `Athey and Imbens (2015) `_.
+ """
+
+ def __init__(
+ self,
+ *,
+ criterion: str = "causal_mse",
+ splitter: str = "best",
+ alpha: float = 0.05,
+ control_name: Union[int, str] = 0,
+ max_depth: int = None,
+ min_samples_split: Union[int, float] = 60,
+ min_weight_fraction_leaf: float = 0.0,
+ max_features: Union[int, float, str] = None,
+ max_leaf_nodes: int = None,
+ min_impurity_decrease: float = float("-inf"),
+ ccp_alpha: float = 0.0,
+ groups_penalty: float = 0.5,
+ min_group_samples: int = 50,
+ min_samples_leaf: int = 100,
+ random_state: int = None,
+ groups_cnt: bool = False,
+ groups_cnt_mode: str = "nodes",
+ ):
+ """
+ Initialize a Causal Tree
+ Args:
+ criterion: ({"causal_mse", "standard_mse"}, default="causal_mse")
+ The function to measure the quality of a split.
+ splitter: ({"best", "random"}, default="best")
+ The strategy used to choose the split at each node. Supported
+ strategies are "best" to choose the best split and "random" to choose
+ the best random split.
+ alpha: (float): the confidence level alpha of the ATE estimate and ITE bootstrap estimates
+ control_name: (str or int): name or index of control group
+ max_depth: (int, default=None)
+ The maximum depth of the tree. If None, then nodes are expanded until
+ all leaves are pure or until all leaves contain less than
+ min_samples_split samples.
+ min_samples_split: (int or float, default=2)
+ The minimum number of samples required to split an internal node:
+ - If int, then consider `min_samples_split` as the minimum number.
+ - If float, then `min_samples_split` is a fraction and
+ `ceil(min_samples_split * n_samples)` are the minimum
+ number of samples for each split.
+ min_weight_fraction_leaf: (float, default=0.0)
+ The minimum weighted fraction of the sum total of weights (of all
+ the input samples) required to be at a leaf node. Samples have
+ equal weight when sample_weight is not provided.
+ max_features: (int, float or {"auto", "sqrt", "log2"}, default=None)
+ The number of features to consider when looking for the best split:
+
+ - If int, then consider `max_features` features at each split.
+ - If float, then `max_features` is a fraction and
+ `int(max_features * n_features)` features are considered at each
+ split.
+ - If "auto", then `max_features=n_features`.
+ - If "sqrt", then `max_features=sqrt(n_features)`.
+ - If "log2", then `max_features=log2(n_features)`.
+ - If None, then `max_features=n_features`.
+ max_leaf_nodes: (int, default=None)
+ Grow a tree with ``max_leaf_nodes`` in best-first fashion.
+ Best nodes are defined as relative reduction in impurity.
+ If None then unlimited number of leaf nodes.
+ min_impurity_decrease: (float, default=float("-inf")))
+ A node will be split if this split induces a decrease of the impurity
+ greater than or equal to this value.
+ ccp_alpha: (non-negative float, default=0.0)
+ Complexity parameter used for Minimal Cost-Complexity Pruning. The
+ subtree with the largest cost complexity that is smaller than
+ ``ccp_alpha`` will be chosen. By default, no pruning is performed. See
+ :ref:`minimal_cost_complexity_pruning` for details.
+ groups_penalty: (float, default=0.5)
+ This penalty coefficient manages the node impurity increase in case of the difference between
+ treatment and control samples sizes.
+ min_group_samples: (int, default=50)
+ The minimum number of samples per each group: k treatment groups and control group.
+ min_samples_leaf: (int or float), default=100
+ The minimum number of samples required to be at a leaf node.
+ A split point at any depth will only be considered if it leaves at
+ least ``min_samples_leaf`` training samples in each of the left and
+ right branches. This may have the effect of smoothing the model,
+ especially in regression.
+
+ - If int, then consider `min_samples_leaf` as the minimum number.
+ - If float, then `min_samples_leaf` is a fraction and
+ `ceil(min_samples_leaf * n_samples)` are the minimum
+ number of samples for each node.
+ random_state: (int), RandomState instance or None, default=None
+ Used to pick randomly the `max_features` used at each split.
+ See :term:`Glossary ` for details.
+ groups_cnt: (bool), count treatment and control groups for each node/leaf
+ groups_cnt_mode: (str, 'nodes', 'leaves'), mode for samples counting
+ """
+
+ self.criterion = criterion
+ self.splitter = splitter
+ self.alpha = alpha
+ self.control_name = control_name
+ self.max_depth = max_depth
+ self.min_samples_split = min_samples_split
+ self.min_weight_fraction_leaf = min_weight_fraction_leaf
+ self.max_features = max_features
+ self.min_group_samples = min_group_samples
+ self.max_leaf_nodes = max_leaf_nodes
+ self.min_impurity_decrease = min_impurity_decrease
+ self.ccp_alpha = ccp_alpha
+ self.groups_penalty = groups_penalty
+ self.min_samples_leaf = min_samples_leaf
+ self.random_state = random_state
+
+ self._classes = {}
+ self.groups_cnt = groups_cnt
+ self.groups_cnt_mode = groups_cnt_mode
+ self._with_outcomes = False
+ self._groups_cnt = {}
+
+ super().__init__(
+ criterion=criterion,
+ splitter=splitter,
+ max_depth=max_depth,
+ min_samples_split=min_samples_split,
+ min_weight_fraction_leaf=min_weight_fraction_leaf,
+ max_features=max_features,
+ min_group_samples=min_group_samples,
+ max_leaf_nodes=max_leaf_nodes,
+ min_impurity_decrease=min_impurity_decrease,
+ ccp_alpha=ccp_alpha,
+ min_samples_leaf=min_samples_leaf,
+ random_state=random_state,
+ )
+
+ def fit(
+ self,
+ X: np.ndarray,
+ treatment: np.ndarray,
+ y: np.ndarray,
+ sample_weight: Union[np.ndarray, None] = None,
+ check_input: bool = True,
+ prepare_data: bool = True,
+ ):
+ """
+ Fit CausalTreeRegressor
+ Args:
+ X (np.ndarray): feature matrix
+ treatment (np.ndarray): treatment vector, includes control group
+ y (np.ndarray): outcome vector
+ sample_weight (np.ndarray): sample_weight, optional
+ check_input (bool, optional): default=False
+ prepare_data (bool): default=True
+ Returns:
+ self
+ """
+
+ if self.criterion == "causal_mse" and self.min_impurity_decrease != float(
+ "-inf"
+ ):
+ raise ValueError(
+ "min_impurity_decrease must be set to -inf for causal_mse criterion"
+ )
+
+ if prepare_data:
+ X, y = self._prepare_data(X=X, y=y, treatment=treatment)
+
+ super().fit(X=X, y=y, sample_weight=sample_weight, check_input=check_input)
+
+ if self.groups_cnt:
+ self._groups_cnt = self._count_groups_distribution(X=X, treatment=treatment)
+ return self
+
+ def predict(
+ self, X: np.ndarray, with_outcomes: bool = False, check_input=True
+ ) -> np.ndarray:
+ """Predict individual treatment effects
+
+ Args:
+ X (np.ndarray): a feature matrix
+ with_outcomes (bool), default=False,
+ include outcomes Y_hat(X|T=0), Y_hat(X|T=1),...,Y_hat(X|T=n)
+ along with individual treatment effects
+ check_input (bool), default=True,
+ Allow to bypass several input checking.
+ Returns:
+ (np.ndarray): individual treatment effect (ITE), dim=(samples, groups)
+ or ITE with outcomes:
+ [Y_hat(X|T=0), Y_hat(X|T=1),...,Y_hat(X|T=n), ITE_1, ITE_2,...,ITE_n], dim=(samples, 2*groups-1)
+ """
+ if check_input:
+ X = self._validate_X_predict(X, check_input)
+ y_outcomes = super().predict(X)
+ y_pred = y_outcomes[:, 1:] - y_outcomes[:, [0]]
+ need_outcomes = with_outcomes or self._with_outcomes
+ out = np.hstack([y_outcomes, y_pred]) if need_outcomes else y_pred
+ # Provides scikit-learn support for _accumulate_prediction() required for causal forests
+ if out.shape[1] == 1:
+ out = out.ravel()
+ return out
+
+ def fit_predict(
+ self,
+ X: np.ndarray,
+ treatment: np.ndarray,
+ y: np.ndarray,
+ return_ci: bool = False,
+ n_bootstraps: int = 1000,
+ bootstrap_size: int = 10000,
+ n_jobs: int = 1,
+ verbose: bool = False,
+ ) -> tuple:
+ """Fit the Causal Tree model and predict treatment effects.
+
+ Args:
+ X (np.ndarray): a feature matrix
+ treatment (np.ndarray): a treatment vector
+ y (np.array): 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
+ n_jobs (int): the number of jobs for bootstrap
+ verbose (str): whether to output progress logs
+
+ Returns:
+ (tuple):
+
+ - te (numpy.ndarray): Predictions of treatment effects.
+ - te_lower (numpy.ndarray, optional): lower bounds of treatment effects
+ - te_upper (numpy.ndarray, optional): upper bounds of treatment effects
+ """
+ self.fit(X=X, y=y, treatment=treatment)
+ te = self.predict(X=X)
+
+ if return_ci:
+ te_bootstraps = self.bootstrap_pool(
+ X=X,
+ y=y,
+ treatment=treatment,
+ n_bootstraps=n_bootstraps,
+ bootstrap_size=bootstrap_size,
+ n_jobs=n_jobs,
+ verbose=verbose,
+ )
+ te_lower = np.percentile(te_bootstraps, (self.alpha / 2) * 100, axis=0)
+ te_upper = np.percentile(te_bootstraps, (1 - self.alpha / 2) * 100, axis=0)
+ return te, te_lower, te_upper
+ else:
+ return te
+
+ def estimate_ate(
+ self, X: np.ndarray, treatment: np.ndarray, y: np.ndarray
+ ) -> tuple:
+ """Estimate the Average Treatment Effect (ATE).
+ Args:
+ X (np.ndarray): a feature matrix
+ treatment (np.array): a treatment vector
+ y (np.ndarray): an outcome vector
+ Returns:
+ tuple, The mean and confidence interval (LB, UB) of the ATE estimate.
+ """
+ dhat = self.fit_predict(X, treatment, y)
+
+ te = dhat.mean()
+ se = dhat.std() / X.shape[0]
+
+ te_lb = te - se * norm.ppf(1 - self.alpha / 2)
+ te_ub = te + se * norm.ppf(1 - self.alpha / 2)
+
+ return te, te_lb, te_ub
+
+ @timeit(exclude_kwargs=("X", "treatment", "y"))
+ def bootstrap_pool(
+ self,
+ X: np.ndarray,
+ treatment: np.ndarray,
+ y: np.ndarray,
+ n_bootstraps: int,
+ bootstrap_size: int,
+ n_jobs: int,
+ verbose: bool,
+ ):
+ """
+ Run a pool of bootstraps
+ Args:
+ X (np.ndarray): a feature matrix
+ treatment (np.ndarray): a treatment vector
+ y (np.ndarray): an outcome vector
+ n_bootstraps (int): number of bootstrap iterations
+ bootstrap_size (int): number of samples per bootstrap
+ n_jobs (int): number of processes
+ verbose (bool): whether to output progress logs
+
+ Returns:
+ (np.ndarray), bootstrap estimates
+
+ """
+
+ def _bootstrap(i: int):
+ if verbose:
+ logger.info(f"Boostrap iteration: {i}")
+ return self.bootstrap(
+ X=X, y=y, treatment=treatment, sample_size=bootstrap_size, seed=i
+ )
+
+ pool = PPool(nodes=n_jobs)
+ pool.restart(force=True)
+
+ bootstrap_estimates = np.array(
+ list(
+ tqdm.tqdm(
+ pool.imap(_bootstrap, (i for i in range(n_bootstraps))),
+ total=n_bootstraps,
+ )
+ )
+ )
+ pool.close()
+ pool.join()
+ return bootstrap_estimates
+
+ def bootstrap(
+ self,
+ X: np.ndarray,
+ treatment: np.ndarray,
+ y: np.ndarray,
+ sample_size: int,
+ seed: int,
+ ) -> np.ndarray:
+ """Runs a single bootstrap.
+
+ Fits on bootstrapped sample, then predicts on whole population.
+
+ Args:
+ X (np.ndarray): a feature matrix
+ treatment (np.ndarray): a treatment vector
+ y (np.ndarray): an outcome vector
+ sample_size (int): bootstrap sample size
+ seed: (int): bootstrap seed
+
+ Returns:
+ (np.ndarray): bootstrap predictions
+ """
+ _rnd = np.random.RandomState(seed=seed)
+ idxs = _rnd.choice(np.arange(0, X.shape[0]), size=sample_size)
+ X_b, y_b, treatment_b = X[idxs], y[idxs], treatment[idxs]
+ self.fit(X=X_b, treatment=treatment_b, y=y_b)
+ te_b = self.predict(X=X)
+ return te_b
+
+ def _prepare_data(
+ self,
+ X: np.ndarray,
+ treatment: np.ndarray,
+ y: np.ndarray,
+ ) -> tuple[np.ndarray, np.ndarray]:
+ """
+ Prepare input data with treatment info for DecisionTreeRegressor.
+ Outcome vector y transforms into y_2dim with (samples x groups) dimensions.
+ Outcomes for the control group are always placed in the first column with index 0.
+ Attribute _group2index stores mapping for y_2dim columns: ({control: 0, treatmentA: 1, treatmentB: 2, ...})
+ Args:
+ X: : (np.ndarray), feature matrix
+ treatment: : (np.ndarray), treatment vector, includes control group
+ y: : (np.ndarray), outcome vector
+ Returns: X, y (samples x groups)
+ """
+ if y.shape[0] != treatment.shape[0]:
+ raise ValueError(
+ f"The number of `treatment` and `y` rows are not equal: {y.shape[0]} {treatment.shape[0]}"
+ )
+ check_treatment_vector(treatment, self.control_name)
+ self.unique_groups = list(set(treatment))
+ self.unique_treatments = sorted(
+ [x for x in self.unique_groups if x != self.control_name]
+ )
+ self._group2index = {
+ self.control_name: 0,
+ **{treatment: i + 1 for i, treatment in enumerate(self.unique_treatments)},
+ }
+
+ X = check_array(X, dtype=DTYPE, accept_sparse="csc")
+ y = check_array(y, ensure_2d=False, dtype=None)
+ self.n_samples, self.n_features = X.shape
+
+ y_2dim = np.zeros((self.n_samples, len(self.unique_treatments) + 1))
+ for group, group_index in self._group2index.items():
+ y_2dim[:, group_index] = np.where(treatment == group, y, np.nan)
+
+ return X, y_2dim
+
+ def _count_groups_distribution(self, X: np.ndarray, treatment: np.ndarray) -> dict:
+ """
+ Count treatment, control distribution for tree nodes/leaves
+ Args:
+ X: (np.ndarray), feature matrix
+ treatment: (np.ndarray), treatment vector
+ Returns:
+ dict: treatment groups for each tree node/leaves
+ """
+ check_is_fitted(self)
+
+ self.is_leaves = get_tree_leaves_mask(self)
+ groups = np.unique(treatment)
+ groups_cnt = {
+ idx: {group: 0 for group in groups}
+ for idx in np.array(range(self.tree_.node_count))
+ }
+ node_indicators = self.tree_.decision_path(X.astype(np.float32))
+
+ for sample_id in range(X.shape[0]):
+ nodes_path = node_indicators.indices[
+ node_indicators.indptr[sample_id] : node_indicators.indptr[
+ sample_id + 1
+ ]
+ ]
+
+ if self.groups_cnt_mode == "leaves":
+ groups_cnt[nodes_path[-1]][treatment[sample_id]] += 1
+ elif self.groups_cnt_mode == "nodes":
+ for node_id in nodes_path:
+ groups_cnt[node_id][treatment[sample_id]] += 1
+ return groups_cnt
diff --git a/causalml/source/causalml/inference/tree/plot.py b/causalml/source/causalml/inference/tree/plot.py
new file mode 100644
index 0000000000000000000000000000000000000000..6e17eb8de9d14f4c8eefbe47f61becde2e482fde
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/plot.py
@@ -0,0 +1,658 @@
+"""
+Visualization functions for forest of trees-based ensemble methods for Uplift modeling on Classification
+Problem.
+"""
+
+from collections import defaultdict
+from typing import Union
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pydotplus
+import seaborn as sns
+from sklearn.tree import _tree
+from sklearn.tree._export import _MPLTreeExporter, _color_brew
+from sklearn.utils.validation import check_is_fitted
+
+from . import CausalTreeRegressor
+from .utils import get_tree_leaves_mask
+
+
+def uplift_tree_string(decisionTree, x_names):
+ """
+ Convert the tree to string for print.
+
+ Args
+ ----
+
+ decisionTree : object
+ object of DecisionTree class
+
+ x_names : list
+ List of feature names
+
+ Returns
+ -------
+ A string representation of the tree.
+ """
+
+ # Column Heading
+ dcHeadings = {}
+ for i, szY in enumerate(x_names + ["treatment_group_key"]):
+ szCol = "Column %d" % i
+ dcHeadings[szCol] = str(szY)
+
+ def toString(decisionTree, indent=""):
+ if decisionTree.results is not None: # leaf node
+ return str(decisionTree.results)
+ else:
+ szCol = "Column %s" % decisionTree.col
+ if szCol in dcHeadings:
+ szCol = dcHeadings[szCol]
+ if isinstance(decisionTree.value, int) or isinstance(
+ decisionTree.value, float
+ ):
+ decision = "%s >= %s?" % (szCol, decisionTree.value)
+ else:
+ decision = "%s == %s?" % (szCol, decisionTree.value)
+ trueBranch = (
+ indent + "yes -> " + toString(decisionTree.trueBranch, indent + "\t\t")
+ )
+ falseBranch = (
+ indent + "no -> " + toString(decisionTree.falseBranch, indent + "\t\t")
+ )
+ return decision + "\n" + trueBranch + "\n" + falseBranch
+
+ print(toString(decisionTree))
+
+
+def uplift_tree_plot(decisionTree, x_names):
+ """
+ Convert the tree to dot graph for plots.
+
+ Args
+ ----
+
+ decisionTree : object
+ object of DecisionTree class
+
+ x_names : list
+ List of feature names
+
+ Returns
+ -------
+ Dot class representing the tree graph.
+ """
+
+ # Column Heading
+ dcHeadings = {}
+ for i, szY in enumerate(x_names + ["treatment_group_key"]):
+ szCol = "Column %d" % i
+ dcHeadings[szCol] = str(szY)
+
+ dcNodes = defaultdict(list)
+ """Plots the obtained decision tree. """
+
+ def toString(
+ iSplit,
+ decisionTree,
+ bBranch,
+ szParent="null",
+ indent="",
+ indexParent=0,
+ upliftScores=list(),
+ ):
+ if decisionTree.results is not None: # leaf node
+ lsY = []
+ for tr, p in zip(decisionTree.classes_, decisionTree.results):
+ lsY.append(f"{tr}:{p:.2f}")
+ dcY = {"name": ", ".join(lsY), "parent": szParent}
+ dcSummary = decisionTree.summary
+ upliftScores += [dcSummary["matchScore"]]
+ dcNodes[iSplit].append(
+ [
+ "leaf",
+ dcY["name"],
+ szParent,
+ bBranch,
+ str(-round(float(decisionTree.summary["impurity"]), 3)),
+ dcSummary["samples"],
+ dcSummary["group_size"],
+ dcSummary["upliftScore"],
+ dcSummary["matchScore"],
+ indexParent,
+ ]
+ )
+ else:
+ szCol = "Column %s" % decisionTree.col
+ if szCol in dcHeadings:
+ szCol = dcHeadings[szCol]
+ if isinstance(decisionTree.value, int) or isinstance(
+ decisionTree.value, float
+ ):
+ decision = "%s >= %s" % (szCol, decisionTree.value)
+ else:
+ decision = "%s == %s" % (szCol, decisionTree.value)
+
+ indexOfLevel = len(dcNodes[iSplit])
+ toString(
+ iSplit + 1,
+ decisionTree.trueBranch,
+ True,
+ decision,
+ indent + "\t\t",
+ indexOfLevel,
+ upliftScores,
+ )
+ toString(
+ iSplit + 1,
+ decisionTree.falseBranch,
+ False,
+ decision,
+ indent + "\t\t",
+ indexOfLevel,
+ upliftScores,
+ )
+ dcSummary = decisionTree.summary
+ upliftScores += [dcSummary["matchScore"]]
+ dcNodes[iSplit].append(
+ [
+ iSplit + 1,
+ decision,
+ szParent,
+ bBranch,
+ str(-round(float(decisionTree.summary["impurity"]), 3)),
+ dcSummary["samples"],
+ dcSummary["group_size"],
+ dcSummary["upliftScore"],
+ dcSummary["matchScore"],
+ indexParent,
+ ]
+ )
+
+ upliftScores = list()
+ toString(0, decisionTree, None, upliftScores=upliftScores)
+
+ upliftScoreToColor = dict()
+ try:
+ # calculate colors for nodes based on uplifts
+ minUplift = min(upliftScores)
+ maxUplift = max(upliftScores)
+ upliftLevels = [
+ (uplift - minUplift) / (maxUplift - minUplift) for uplift in upliftScores
+ ] # min max scaler
+ baseUplift = float(decisionTree.summary.get("matchScore"))
+ baseUpliftLevel = (baseUplift - minUplift) / (
+ maxUplift - minUplift
+ ) # min max scaler normalization
+ white = np.array([255.0, 255.0, 255.0])
+ blue = np.array([31.0, 119.0, 180.0])
+ green = np.array([0.0, 128.0, 0.0])
+ for i, upliftLevel in enumerate(upliftLevels):
+ if upliftLevel >= baseUpliftLevel: # go blue
+ color = upliftLevel * blue + (1 - upliftLevel) * white
+ else: # go green
+ color = (1 - upliftLevel) * green + upliftLevel * white
+ color = [int(c) for c in color]
+ upliftScoreToColor[upliftScores[i]] = ("#%2x%2x%2x" % tuple(color)).replace(
+ " ", "0"
+ ) # color code
+ except Exception as e:
+ print(e)
+
+ lsDot = [
+ "digraph Tree {",
+ 'node [shape=box, style="filled, rounded", color="black", fontname=helvetica] ;',
+ "edge [fontname=helvetica] ;",
+ ]
+ i_node = 0
+ dcParent = {}
+ totalSample = int(
+ decisionTree.summary.get("samples")
+ ) # initialize the value with the total sample size at root
+ for nSplit in range(len(dcNodes.items())):
+ lsY = dcNodes[nSplit]
+ indexOfLevel = 0
+ for lsX in lsY:
+ (
+ iSplit,
+ decision,
+ szParent,
+ bBranch,
+ szImpurity,
+ szSamples,
+ szGroup,
+ upliftScore,
+ matchScore,
+ indexParent,
+ ) = lsX
+
+ sampleProportion = round(int(szSamples) * 100.0 / totalSample, 1)
+ if type(iSplit) is int:
+ szSplit = "%d-%d" % (iSplit, indexOfLevel)
+ dcParent[szSplit] = i_node
+ lsDot.append(
+ "%d [label=<%s impurity %s total_sample %s (%s%) group_sample %s "
+ "uplift score: %s uplift p_value %s "
+ 'validation uplift score %s>, fillcolor="%s"] ;'
+ % (
+ i_node,
+ decision.replace(">=", "≥").replace("?", ""),
+ szImpurity,
+ szSamples,
+ str(sampleProportion),
+ szGroup,
+ str(upliftScore[0]),
+ str(upliftScore[1]),
+ str(matchScore),
+ upliftScoreToColor.get(matchScore, "#e5813900"),
+ )
+ )
+ else:
+ lsDot.append(
+ "%d [label=< impurity %s total_sample %s (%s%) group_sample %s "
+ "uplift score: %s uplift p_value %s validation uplift score %s "
+ 'mean %s>, fillcolor="%s"] ;'
+ % (
+ i_node,
+ szImpurity,
+ szSamples,
+ str(sampleProportion),
+ szGroup,
+ str(upliftScore[0]),
+ str(upliftScore[1]),
+ str(matchScore),
+ decision,
+ upliftScoreToColor.get(matchScore, "#e5813900"),
+ )
+ )
+
+ if szParent != "null":
+ if bBranch:
+ szAngle = "45"
+ szHeadLabel = "True"
+ else:
+ szAngle = "-45"
+ szHeadLabel = "False"
+ szSplit = "%d-%d" % (nSplit, indexParent)
+ p_node = dcParent[szSplit]
+ if nSplit == 1:
+ lsDot.append(
+ '%d -> %d [labeldistance=2.5, labelangle=%s, headlabel="%s"] ;'
+ % (p_node, i_node, szAngle, szHeadLabel)
+ )
+ else:
+ lsDot.append("%d -> %d ;" % (p_node, i_node))
+ i_node += 1
+ indexOfLevel += 1
+ lsDot.append("}")
+ dot_data = "\n".join(lsDot)
+ graph = pydotplus.graph_from_dot_data(dot_data)
+ return graph
+
+
+def plot_dist_tree_leaves_values(
+ tree: CausalTreeRegressor,
+ title: str = "Leaves values distribution",
+ figsize: tuple = (5, 5),
+ fontsize: int = 12,
+) -> None:
+ """
+ Create distplot for tree leaves values
+ Args:
+ tree: (CausalTreeRegressor), Tree object
+ title: (str), plot title
+ figsize: (tuple), figure size
+ fontsize: (int), title font size
+
+ Returns: None
+
+ """
+ tree_leaves_mask = get_tree_leaves_mask(tree)
+ leaves_values = tree.tree_.value
+ treatment_effects = leaves_values[:, 1] - leaves_values[:, 0]
+ treatment_effects = treatment_effects.reshape(
+ -1,
+ )[tree_leaves_mask]
+ fig, ax = plt.subplots(figsize=figsize)
+ sns.distplot(
+ treatment_effects,
+ ax=ax,
+ )
+ plt.title(title, fontsize=fontsize)
+ plt.show()
+
+
+class _MPLCTreeExporter(_MPLTreeExporter):
+ def __init__(
+ self,
+ causal_tree: CausalTreeRegressor,
+ max_depth: int,
+ feature_names: list,
+ class_names: list,
+ label: str,
+ filled: bool,
+ impurity: bool,
+ groups_count: bool,
+ treatment_groups: tuple,
+ node_ids: bool,
+ proportion: bool,
+ rounded: bool,
+ precision: int,
+ fontsize: int,
+ ):
+ """
+ Causal Tree exporter for matplotlib
+ Source: https://github.com/scikit-learn/scikit-learn/blob/1.0.X/sklearn/tree/_export.py
+ Args:
+ causal_tree: CausalTreeRegressor
+ The causal tree to be plotted
+ max_depth: int, default=None
+ The maximum depth of the representation. If None, the tree is fully generated.
+ feature_names: list of strings, default=None
+ Names of each of the features.
+ If None, generic names will be used ("X[0]", "X[1]", ...).
+ class_names: list of str or bool, default=None
+ Names of each of the target classes in ascending numerical order.
+ Only relevant for classification and not supported for multi-output.
+ If ``True``, shows a symbolic representation of the class name.
+ label: {'all', 'root', 'none'}, default='all'
+ Whether to show informative labels for impurity, etc.
+ Options include 'all' to show at every node, 'root' to show only at
+ the top root node, or 'none' to not show at any node.
+ filled: bool, default=False
+ When set to ``True``, paint nodes to indicate extremity of node values
+ impurity: bool, default=True
+ When set to ``True``, show the impurity at each node.
+ groups_count: bool, default=True
+ Add the number of treatment and control groups
+ treatment_groups: tuple, default=(0, 1)
+ Treatment and control groups labels
+ node_ids: bool, default=False
+ When set to ``True``, show the ID number on each node.
+ proportion: bool, default=False
+ When set to ``True``, change the display of 'values' and/or 'samples'
+ to be proportions and percentages respectively.
+ rounded: bool, default=False
+ When set to ``True``, draw node boxes with rounded corners and use
+ Helvetica fonts instead of Times-Roman.
+ precision: int, default=3
+ Number of digits of precision for floating point in the values of
+ impurity, threshold and value attributes of each node.
+ fontsize: int, default=None
+ Size of text font. If None, determined automatically to fit figure.
+ """
+ super().__init__(
+ max_depth,
+ feature_names,
+ class_names,
+ label,
+ filled,
+ impurity,
+ node_ids,
+ proportion,
+ rounded,
+ precision,
+ fontsize,
+ )
+ self.causal_tree = causal_tree
+ self.groups_count = groups_count
+ self.treatment_groups = treatment_groups
+
+ def node_to_str(
+ self, tree: _tree.Tree, node_id: int, criterion: Union[str, object]
+ ) -> str:
+ """
+ Generate the node content string
+ Args:
+ tree: Tree class
+ node_id: int, Tree node id
+ criterion: str or object, split criterion
+ Returns: str, node content
+ """
+ if tree.n_outputs == 1:
+ value = tree.value[node_id][0, :]
+ else:
+ value = tree.value[node_id]
+
+ # Should labels be shown?
+ labels = (self.label == "root" and node_id == 0) or self.label == "all"
+
+ characters = self.characters
+ node_string = characters[-1]
+
+ # Write node ID
+ if self.node_ids:
+ if labels:
+ node_string += "node "
+ node_string += characters[0] + str(node_id) + characters[4]
+
+ # Write decision criteria
+ if tree.children_left[node_id] != _tree.TREE_LEAF:
+ # Always write node decision criteria, except for leaves
+ if self.feature_names is not None:
+ feature = self.feature_names[tree.feature[node_id]]
+ else:
+ feature = "X%s%s%s" % (
+ characters[1],
+ tree.feature[node_id],
+ characters[2],
+ )
+ node_string += "%s %s %s%s" % (
+ feature,
+ characters[3],
+ round(tree.threshold[node_id], self.precision),
+ characters[4],
+ )
+
+ # Write impurity
+ if self.impurity:
+ if not isinstance(criterion, str):
+ criterion = "impurity"
+ if labels:
+ node_string += "%s = " % criterion
+ node_string += (
+ str(round(tree.impurity[node_id], self.precision)) + characters[4]
+ )
+
+ # Write node sample count
+ if labels:
+ node_string += "samples = "
+ if self.proportion:
+ percent = (
+ 100.0 * tree.n_node_samples[node_id] / float(tree.n_node_samples[0])
+ )
+ node_string += str(round(percent, 1)) + "%" + characters[4]
+ else:
+ node_string += str(tree.n_node_samples[node_id]) + characters[4]
+
+ # Write the number of samples per treatment and control groups
+ if self.groups_count:
+ for group in self.treatment_groups:
+ node_string += (
+ f"Group {group} = {self.causal_tree._groups_cnt[node_id][group]} "
+ )
+ node_string += characters[4]
+
+ # Write node class distribution / regression value
+ if self.proportion and tree.n_classes[0] != 1:
+ # For classification this will show the proportion of samples
+ value = value / tree.weighted_n_node_samples[node_id]
+ if labels:
+ node_string += "value = "
+ if tree.n_classes[0] == 1:
+ # Regression
+ value_text = np.around(value, self.precision)
+ elif self.proportion:
+ # Classification
+ value_text = np.around(value, self.precision)
+ elif np.all(np.equal(np.mod(value, 1), 0)):
+ # Classification without floating-point weights
+ value_text = value.astype(int)
+ else:
+ # Classification with floating-point weights
+ value_text = np.around(value, self.precision)
+ # Strip whitespace
+ value_text = str(value_text.astype("S32")).replace("b'", "'")
+ value_text = value_text.replace("' '", ", ").replace("'", "")
+ if tree.n_classes[0] == 1 and tree.n_outputs == 1:
+ value_text = value_text.replace("[", "").replace("]", "")
+ value_text = value_text.replace("\n ", characters[4])
+ node_string += value_text + characters[4]
+
+ # Write node majority class
+ if (
+ self.class_names is not None
+ and tree.n_classes[0] != 1
+ and tree.n_outputs == 1
+ ):
+ # Only done for single-output classification trees
+ if labels:
+ node_string += "class = "
+ if self.class_names is not True:
+ class_name = self.class_names[np.argmax(value)]
+ else:
+ class_name = "y%s%s%s" % (
+ characters[1],
+ np.argmax(value),
+ characters[2],
+ )
+ node_string += class_name
+
+ # Clean up any trailing newlines
+ if node_string.endswith(characters[4]):
+ node_string = node_string[: -len(characters[4])]
+
+ return node_string + characters[5]
+
+ def get_color(self, value: np.ndarray) -> str:
+ """
+ Compute HTML color for a Tree node
+ Args:
+ value: Tree node value
+ Returns: str, html color code in #RRGGBB format
+ """
+ # Regression tree or multi-output
+ color = list(self.colors["rgb"][0])
+ alpha = float(value - self.colors["bounds"][0]) / (
+ self.colors["bounds"][1] - self.colors["bounds"][0]
+ )
+ alpha = 0 if np.isnan(alpha) else alpha
+ # Compute the color as alpha against white
+ color = [int(round(alpha * c + (1 - alpha) * 255, 0)) for c in color]
+ return "#%2x%2x%2x" % tuple(color)
+
+ def get_fill_color(self, tree: _tree.Tree, node_id: int) -> str:
+ """
+ Fetch appropriate color for node
+ Args:
+ tree: Tree class
+ node_id: int, node index
+ Returns: str
+ """
+ if "rgb" not in self.colors:
+ # Initialize colors and bounds if required
+ self.colors["rgb"] = _color_brew(tree.n_classes[0])
+ if tree.n_outputs != 1:
+ # Find max and min impurities for multi-output
+ self.colors["bounds"] = (
+ np.nanmin(-tree.impurity),
+ np.nanmax(-tree.impurity),
+ )
+ elif tree.n_classes[0] == 1 and len(np.unique(tree.value)) != 1:
+ # Find max and min values in leaf nodes for regression
+ self.colors["bounds"] = (np.nanmin(tree.value), np.nanmax(tree.value))
+ if tree.n_outputs == 1:
+ node_val = tree.value[node_id][0, :] / tree.weighted_n_node_samples[node_id]
+ if tree.n_classes[0] == 1:
+ # Regression
+ node_val = tree.value[node_id][0, :]
+ else:
+ # If multi-output color node by impurity
+ node_val = -tree.impurity[node_id]
+ return self.get_color(node_val)
+
+
+def plot_causal_tree(
+ causal_tree: CausalTreeRegressor,
+ *,
+ max_depth: int = None,
+ feature_names: list = None,
+ class_names: list = None,
+ label: str = "all",
+ filled: bool = False,
+ impurity: bool = True,
+ groups_count: bool = True,
+ treatment_groups: tuple = (0, 1),
+ node_ids: bool = False,
+ proportion: bool = False,
+ rounded: bool = False,
+ precision: int = 3,
+ ax: plt.Axes = None,
+ fontsize: int = None,
+):
+ """
+ Plot a Causal Tree.
+ Source: https://github.com/scikit-learn/scikit-learn/blob/1.0.X/sklearn/tree/_export.py
+ Args:
+ causal_tree: CausalTreeRegressor
+ The causal tree to be plotted
+ max_depth: int, default=None
+ The maximum depth of the representation. If None, the tree is fully generated.
+ feature_names: list of strings, default=None
+ Names of each of the features.
+ If None, generic names will be used ("X[0]", "X[1]", ...).
+ class_names: list of str or bool, default=None
+ Names of each of the target classes in ascending numerical order.
+ Only relevant for classification and not supported for multi-output.
+ If ``True``, shows a symbolic representation of the class name.
+ label: {'all', 'root', 'none'}, default='all'
+ Whether to show informative labels for impurity, etc.
+ Options include 'all' to show at every node, 'root' to show only at
+ the top root node, or 'none' to not show at any node.
+ filled: bool, default=False
+ When set to ``True``, paint nodes to indicate extremity of node values
+ impurity: bool, default=True
+ When set to ``True``, show the impurity at each node.
+ groups_count: bool, default=True
+ Add the number of treatment and control groups
+ treatment_groups: tuple, default=(0, 1)
+ Treatment and control groups labels
+ node_ids: bool, default=False
+ When set to ``True``, show the ID number on each node.
+ proportion: bool, default=False
+ When set to ``True``, change the display of 'values' and/or 'samples'
+ to be proportions and percentages respectively.
+ rounded: bool, default=False
+ When set to ``True``, draw node boxes with rounded corners and use
+ Helvetica fonts instead of Times-Roman.
+ precision: int, default=3
+ Number of digits of precision for floating point in the values of
+ impurity, threshold and value attributes of each node.
+ ax: matplotlib axis, default=None
+ Axes to plot to. If None, use current axis. Any previous content
+ is cleared.
+ fontsize: int, default=None
+ Size of text font. If None, determined automatically to fit figure.
+ Returns:
+
+ """
+ check_is_fitted(causal_tree)
+
+ exporter = _MPLCTreeExporter(
+ causal_tree=causal_tree,
+ max_depth=max_depth,
+ feature_names=feature_names,
+ class_names=class_names,
+ label=label,
+ filled=filled,
+ impurity=impurity,
+ groups_count=groups_count,
+ treatment_groups=treatment_groups,
+ node_ids=node_ids,
+ proportion=proportion,
+ rounded=rounded,
+ precision=precision,
+ fontsize=fontsize,
+ )
+ exporter.export(causal_tree, ax=ax)
diff --git a/causalml/source/causalml/inference/tree/uplift.pyx b/causalml/source/causalml/inference/tree/uplift.pyx
new file mode 100644
index 0000000000000000000000000000000000000000..29681b6c05eea743e205d7756da5de756a3df088
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/uplift.pyx
@@ -0,0 +1,2572 @@
+# cython: cdivision=True
+# cython: boundscheck=False
+# cython: wraparound=False
+# cython: language_level=3
+"""
+Forest of trees-based ensemble methods for Uplift modeling on Classification
+Problem. Those methods include random forests and extremely randomized trees.
+
+The module structure is the following:
+- The ``UpliftRandomForestClassifier`` base class implements different
+ variants of uplift models based on random forest, with 'fit' and 'predict'
+ method.
+- The ``UpliftTreeClassifier`` base class implements the uplift trees (without
+ Bootstrapping for random forest), this class is called within
+ ``UpliftRandomForestClassifier`` for constructing random forest.
+
+"""
+
+# Authors: Zhenyu Zhao
+# Totte Harinen
+
+import multiprocessing as mp
+from collections import defaultdict
+
+import logging
+import cython
+import numpy as np
+cimport numpy as np
+import pandas as pd
+import scipy.stats as stats
+import sklearn
+from joblib import Parallel, delayed
+from packaging import version
+from sklearn.model_selection import train_test_split
+from sklearn.utils import check_X_y, check_array, check_random_state
+import numbers
+
+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
+
+N_TYPE = np.int32
+TR_TYPE = np.int8
+Y_TYPE = np.int8
+P_TYPE = np.float64
+
+ctypedef np.int32_t N_TYPE_t
+ctypedef np.int8_t TR_TYPE_t
+ctypedef np.int8_t Y_TYPE_t
+ctypedef np.float64_t P_TYPE_t
+
+MAX_INT = np.iinfo(np.int32).max
+
+logger = logging.getLogger("causalml")
+
+cdef extern from "math.h":
+ double log(double x) nogil
+ double fabs(double x) nogil
+ double sqrt(double x) nogil
+
+@cython.cfunc
+def kl_divergence(pk: cython.float, qk: cython.float) -> cython.float:
+ '''
+ Calculate KL Divergence for binary classification.
+
+ sum(np.array(pk) * np.log(np.array(pk) / np.array(qk)))
+
+ Args
+ ----
+ pk : float
+ The probability of 1 in one distribution.
+ qk : float
+ The probability of 1 in the other distribution.
+
+ Returns
+ -------
+ S : float
+ The KL divergence.
+ '''
+
+ eps: cython.float = 1e-6
+ S: cython.float
+
+ if qk == 0.:
+ return 0.
+
+ qk = min(max(qk, eps), 1 - eps)
+
+ if pk == 0.:
+ S = -log(1 - qk)
+ elif pk == 1.:
+ S = -log(qk)
+ else:
+ S = pk * log(pk / qk) + (1 - pk) * log((1 - pk) / (1 - qk))
+
+ return S
+
+
+@cython.cfunc
+def entropyH(p: cython.float, q: cython.float=-1.) -> cython.float:
+ '''
+ Entropy
+
+ Entropy calculation for normalization.
+
+ Args
+ ----
+ p : float
+ The probability used in the entropy calculation.
+
+ q : float, optional, (default = -1.)
+ The second probability used in the entropy calculation.
+
+ Returns
+ -------
+ entropy : float
+ '''
+
+ if q == -1. and p > 0.:
+ return -p * log(p)
+ elif q > 0.:
+ return -p * log(q)
+ else:
+ return 0.
+
+
+class DecisionTree:
+ """ Tree Node Class
+
+ Tree node class to contain all the statistics of the tree node.
+
+ Parameters
+ ----------
+ classes_ : list of str
+ A list of the control and treatment group names.
+
+ col : int, optional (default = -1)
+ The column index for splitting the tree node to children nodes.
+
+ value : float, optional (default = None)
+ The value of the feature column to split the tree node to children nodes.
+
+ trueBranch : object of DecisionTree
+ The true branch tree node (feature > value).
+
+ falseBranch : object of DecisionTree
+ The false branch tree node (feature > value).
+
+ results : list of float
+ The classification probability P(Y=1|T) for each of the control and treatment groups
+ in the tree node.
+
+ summary : list of list
+ Summary statistics of the tree nodes, including impurity, sample size, uplift score, etc.
+
+ maxDiffTreatment : int
+ The treatment index generating the maximum difference between the treatment and control groups.
+
+ maxDiffSign : float
+ The sign of the maximum difference (1. or -1.).
+
+ nodeSummary : list of list
+ Summary statistics of the tree nodes [P(Y=1|T), N(T)], where y_mean stands for the target metric mean
+ and n is the sample size.
+
+ backupResults : list of float
+ The positive probabilities in each of the control and treatment groups in the parent node. The parent node
+ information is served as a backup for the children node, in case no valid statistics can be calculated from the
+ children node, the parent node information will be used in certain cases.
+
+ bestTreatment : int
+ The treatment index providing the best uplift (treatment effect).
+
+ upliftScore : list
+ The uplift score of this node: [max_Diff, p_value], where max_Diff stands for the maximum treatment effect, and
+ p_value stands for the p_value of the treatment effect.
+
+ matchScore : float
+ The uplift score by filling a trained tree with validation dataset or testing dataset.
+
+ """
+
+ def __init__(self, classes_, col=-1, value=None, trueBranch=None, falseBranch=None, results=None, summary=None,
+ maxDiffTreatment=None, maxDiffSign=1., nodeSummary=None, backupResults=None, bestTreatment=None,
+ upliftScore=None, matchScore=None):
+ self.classes_ = classes_
+ self.col = col
+ self.value = value
+ self.trueBranch = trueBranch
+ self.falseBranch = falseBranch
+ self.results = results # None for nodes, not None for leaves
+ self.summary = summary
+ # the treatment with max( |p(y|treatment) - p(y|control)| )
+ self.maxDiffTreatment = maxDiffTreatment
+ # the sign for p(y|maxDiffTreatment) - p(y|control)
+ self.maxDiffSign = maxDiffSign
+ self.nodeSummary = nodeSummary
+ self.backupResults = backupResults
+ self.bestTreatment = bestTreatment
+ self.upliftScore = upliftScore
+ # match actual treatment for validation and testing
+ self.matchScore = matchScore
+
+
+def group_uniqueCounts_to_arr(np.ndarray[TR_TYPE_t, ndim=1] treatment_idx,
+ np.ndarray[Y_TYPE_t, ndim=1] y,
+ np.ndarray[N_TYPE_t, ndim=1] out_arr):
+ '''
+ Count sample size by experiment group.
+
+ Args
+ ----
+ treatment_idx : array-like, shape = [num_samples]
+ An array containing the treatment group index for each unit.
+ Should be of type numpy.int8
+ y : array-like, shape = [num_samples]
+ An array containing the outcome of interest for each unit.
+ Should be of type numpy.int8
+ out_arr : array-like, shape = [2 * n_class]
+ An array to store the output counts, should have type numpy.int32
+
+ Returns
+ -------
+
+ No return value, but modified the out_arr to hold the negative and positive
+ outcome sample sizes for each of the control and treatment groups.
+ out_arr[2*i] is N(Y = 0, T = i) for i = 0, ..., n_class
+ out_arr[2*i+1] is N(Y = 1, T = i) for i = 0, ..., n_class
+ '''
+ cdef int out_arr_len = out_arr.shape[0]
+ cdef int n_class = out_arr_len / 2
+ cdef int num_samples = treatment_idx.shape[0]
+ cdef int yv = 0
+ cdef int tv = 0
+ cdef int i = 0
+ # first clear the output
+ for i in range(out_arr_len):
+ out_arr[i] = 0
+ # then loop through treatment_idx and y, sum the counts
+ # first sum as N(T = i) and N(Y = 1, T = i) at index (2*i, 2*i+1), and later adjust
+ for i in range(num_samples):
+ tv = treatment_idx[i]
+ # assume treatment index is in range
+ out_arr[2*tv] += 1
+ # assume y should be either 0 or 1, so this is summing
+ out_arr[2*tv + 1] += y[i]
+ # adjust the entry at index 2*i to be N(Y = 0, T = i) = N(T = i) - N(Y = 1, T = i)
+ for i in range(n_class):
+ out_arr[2*i] -= out_arr[2*i + 1]
+ # done, modified out_arr, so no need to return it
+
+def group_counts_by_divide(
+ col_vals, threshold_val, is_split_by_gt,
+ np.ndarray[TR_TYPE_t, ndim=1] treatment_idx,
+ np.ndarray[Y_TYPE_t, ndim=1] y,
+ np.ndarray[N_TYPE_t, ndim=1] out_arr):
+ '''
+ Count sample size by experiment group for the left branch,
+ after splitting col_vals by threshold_val.
+ If is_split_by_gt, the left branch is (col_vals >= threshold_val),
+ otherwise the left branch is (col_vals == threshold_val).
+
+ This aims to combine the previous divideSet_len and
+ group_uniqueCounts_to_arr into one function, so as to reduce the
+ number of intermediate objects.
+
+ Args
+ ----
+ col_vals : array-like, shape = [num_samples]
+ An array containing one column of x values.
+ threshold_val : compatible value with col_vals
+ A value for splitting col_vals.
+ If is_split_by_gt, the left branch is (col_vals >= threshold_val),
+ otherwise the left branch is (col_vals == threshold_val).
+ is_split_by_gt : bool
+ Whether to split by (col_vals >= threshold_val).
+ If False, will split by (col_vals == threshold_val).
+ treatment_idx : array-like, shape = [num_samples]
+ An array containing the treatment group index for each unit.
+ Should be of type numpy.int8
+ y : array-like, shape = [num_samples]
+ An array containing the outcome of interest for each unit.
+ Should be of type numpy.int8
+ out_arr : array-like, shape = [2 * n_class]
+ An array to store the output counts, should have type numpy.int32
+
+ Returns
+ -------
+ len_X_l: the number of samples in the left branch.
+ Also modify the out_arr to hold the negative and positive
+ outcome sample sizes for each of the control and treatment groups.
+ out_arr[2*i] is N(Y = 0, T = i) for i = 0, ..., n_class
+ out_arr[2*i+1] is N(Y = 1, T = i) for i = 0, ..., n_class
+ '''
+ cdef int out_arr_len = out_arr.shape[0]
+ cdef int n_class = out_arr_len / 2
+ cdef int num_samples = treatment_idx.shape[0]
+ cdef int yv = 0
+ cdef int tv = 0
+ cdef int i = 0
+ cdef N_TYPE_t len_X_l = 0
+ cdef np.ndarray[np.uint8_t, ndim=1, cast=True] filt
+ # first clear the output
+ for i in range(out_arr_len):
+ out_arr[i] = 0
+
+ # split
+ if is_split_by_gt:
+ filt = col_vals >= threshold_val
+ else:
+ filt = col_vals == threshold_val
+
+ # then loop through treatment_idx and y, sum the counts where filt
+ # is True, and it is the count for the left branch.
+ # Also count len_X_l in the process.
+
+ # first sum as N(T = i) and N(Y = 1, T = i) at index (2*i, 2*i+1), and later adjust
+ for i in range(num_samples):
+ if filt[i]> 0:
+ len_X_l += 1
+ tv = treatment_idx[i]
+ # assume treatment index is in range
+ out_arr[2*tv] += 1
+ # assume y should be either 0 or 1, so this is summing
+ out_arr[2*tv + 1] += y[i]
+ # adjust the entry at index 2*i to be N(Y = 0, T = i) = N(T = i) - N(Y = 1, T = i)
+ for i in range(n_class):
+ out_arr[2*i] -= out_arr[2*i + 1]
+ # done, modified out_arr
+ return len_X_l
+
+# Uplift Tree Classifier
+class UpliftTreeClassifier:
+ """ Uplift Tree Classifier for Classification Task.
+
+ A uplift tree classifier estimates the individual treatment effect by modifying the loss function in the
+ classification trees.
+
+ The uplift tree classifier is used in uplift random forest to construct the trees in the forest.
+
+ Parameters
+ ----------
+
+ evaluationFunction : string
+ Choose from one of the models: 'KL', 'ED', 'Chi', 'CTS', 'DDP', 'IT', 'CIT', 'IDDP'.
+
+ max_features: int, optional (default=None)
+ The number of features to consider when looking for the best split.
+
+ max_depth: int, optional (default=3)
+ The maximum depth of the tree.
+
+ min_samples_leaf: int, optional (default=100)
+ The minimum number of samples required to be split at a leaf node.
+
+ min_samples_treatment: int, optional (default=10)
+ The minimum number of samples required of the experiment group to be split at a leaf node.
+
+ n_reg: int, optional (default=100)
+ The regularization parameter defined in Rzepakowski et al. 2012, the weight (in terms of sample size) of the
+ parent node influence on the child node, only effective for 'KL', 'ED', 'Chi', 'CTS' methods.
+
+ early_stopping_eval_diff_scale: float, optional (default=1)
+ If train and valid uplift score diff bigger than
+ min(train_uplift_score,valid_uplift_score)/early_stopping_eval_diff_scale, stop.
+
+ control_name: string
+ The name of the control group (other experiment groups will be regarded as treatment groups).
+
+ normalization: boolean, optional (default=True)
+ The normalization factor defined in Rzepakowski et al. 2012, correcting for tests with large number of splits
+ and imbalanced treatment and control splits.
+
+ honesty: bool (default=False)
+ True if the honest approach based on "Athey, S., & Imbens, G. (2016). Recursive partitioning for heterogeneous causal effects."
+ shall be used. If 'IDDP' is used as evaluation function, this parameter is automatically set to true.
+
+ estimation_sample_size: float (default=0.5)
+ Sample size for estimating the CATE score in the leaves if honesty == True.
+
+ random_state: int, RandomState instance or None (default=None)
+ A random seed or `np.random.RandomState` to control randomness in building a tree.
+
+ """
+ def __init__(self, control_name, max_features=None, max_depth=3, min_samples_leaf=100,
+ min_samples_treatment=10, n_reg=100, early_stopping_eval_diff_scale=1, evaluationFunction='KL',
+ normalization=True, honesty=False, estimation_sample_size=0.5, random_state=None):
+ self.max_depth = max_depth
+ self.min_samples_leaf = min_samples_leaf
+ self.min_samples_treatment = min_samples_treatment
+ self.n_reg = n_reg
+ self.early_stopping_eval_diff_scale = early_stopping_eval_diff_scale
+ self.max_features = max_features
+
+ assert evaluationFunction in ['KL', 'ED', 'Chi', 'CTS', 'DDP', 'IT', 'CIT', 'IDDP'], \
+ f"evaluationFunction should be either 'KL', 'ED', 'Chi', 'CTS', 'DDP', 'IT', 'CIT', or 'IDDP' but {evaluationFunction} is passed"
+
+ if evaluationFunction == 'KL':
+ self.evaluationFunction = self.evaluate_KL
+ self.arr_eval_func = self.arr_evaluate_KL
+ elif evaluationFunction == 'ED':
+ self.evaluationFunction = self.evaluate_ED
+ self.arr_eval_func = self.arr_evaluate_ED
+ elif evaluationFunction == 'Chi':
+ self.evaluationFunction = self.evaluate_Chi
+ self.arr_eval_func = self.arr_evaluate_Chi
+ elif evaluationFunction == 'DDP':
+ self.evaluationFunction = self.evaluate_DDP
+ self.arr_eval_func = self.arr_evaluate_DDP
+ elif evaluationFunction == 'IT':
+ self.evaluationFunction = self.evaluate_IT
+ self.arr_eval_func = self.arr_evaluate_IT
+ elif evaluationFunction == 'CIT':
+ self.evaluationFunction = self.evaluate_CIT
+ self.arr_eval_func = self.arr_evaluate_CIT
+ elif evaluationFunction == 'IDDP':
+ self.evaluationFunction = self.evaluate_IDDP
+ self.arr_eval_func = self.arr_evaluate_IDDP
+ elif evaluationFunction == 'CTS':
+ self.evaluationFunction = self.evaluate_CTS
+ self.arr_eval_func = self.arr_evaluate_CTS
+ self.fitted_uplift_tree = None
+
+ assert control_name is not None and isinstance(control_name, str), \
+ f"control_group should be string but {control_name} is passed"
+ self.control_name = control_name
+ self.classes_ = [self.control_name]
+ self.n_class = 1
+ self.normalization = normalization
+ self.honesty = honesty
+ self.estimation_sample_size = estimation_sample_size
+ self.random_state = random_state
+ if evaluationFunction == 'IDDP' and self.honesty is False:
+ self.honesty = True
+
+
+ def fit(self, X, treatment, y, X_val=None, treatment_val=None, y_val=None):
+ """ Fit the uplift model.
+
+ Args
+ ----
+ X : ndarray, shape = [num_samples, num_features]
+ An ndarray of the covariates used to train the uplift model.
+
+ treatment : array-like, shape = [num_samples]
+ An array containing the treatment group for each unit.
+
+ y : array-like, shape = [num_samples]
+ An array containing the outcome of interest for each unit.
+
+ Returns
+ -------
+ self : object
+ """
+
+ self.random_state_ = check_random_state(self.random_state)
+
+ X, y = check_X_y(X, y)
+ y = (y > 0).astype(Y_TYPE) # make sure it is 0 or 1, and is int8
+ treatment = np.asarray(treatment)
+ assert len(y) == len(treatment), 'Data length must be equal for X, treatment, and y.'
+ if X_val is not None:
+ X_val, y_val = check_X_y(X_val, y_val)
+ y_val = (y_val > 0).astype(Y_TYPE) # make sure it is 0 or 1, and is int8
+ treatment_val = np.asarray(treatment_val)
+ assert len(y_val) == len(treatment_val), 'Data length must be equal for X_val, treatment_val, and y_val.'
+
+ # Get treatment group keys. self.classes_[0] is reserved for the control group.
+ treatment_groups = sorted([x for x in list(set(treatment)) if x != self.control_name])
+ self.classes_ = [self.control_name]
+ treatment_idx = np.zeros_like(treatment, dtype=TR_TYPE)
+ treatment_val_idx = None
+ if treatment_val is not None:
+ treatment_val_idx = np.zeros_like(treatment_val, dtype=TR_TYPE)
+ for i, tr in enumerate(treatment_groups, 1):
+ self.classes_.append(tr)
+ treatment_idx[treatment == tr] = i
+ if treatment_val_idx is not None:
+ treatment_val_idx[treatment_val == tr] = i
+ self.n_class = len(self.classes_)
+
+ self.feature_imp_dict = defaultdict(float)
+
+ if (self.n_class > 2) and (self.evaluationFunction in [self.evaluate_DDP, self.evaluate_IDDP, self.evaluate_IT, self.evaluate_CIT]):
+ raise ValueError("The DDP, IDDP, IT, and CIT approach can only cope with two class problems, that is two different treatment "
+ "options (e.g., control vs treatment). Please select another approach or only use a "
+ "dataset which employs two treatment options.")
+
+ if self.honesty:
+ try:
+ X, X_est, treatment_idx, treatment_idx_est, y, y_est = train_test_split(X, treatment_idx, y, stratify=np.stack([treatment_idx, y], axis=1), test_size=self.estimation_sample_size,
+ shuffle=True, random_state=self.random_state)
+ except ValueError:
+ logger.warning(f"Stratified sampling failed. Falling back to random sampling.")
+ X, X_est, treatment_idx, treatment_idx_est, y, y_est = train_test_split(X, treatment_idx, y, test_size=self.estimation_sample_size, shuffle=True,
+ random_state=self.random_state)
+
+ self.fitted_uplift_tree = self.growDecisionTreeFrom(
+ X, treatment_idx, y, X_val, treatment_val_idx, y_val,
+ max_depth=self.max_depth, early_stopping_eval_diff_scale=self.early_stopping_eval_diff_scale,
+ min_samples_leaf=self.min_samples_leaf,
+ depth=1, min_samples_treatment=self.min_samples_treatment,
+ n_reg=self.n_reg, parentNodeSummary_p=None
+ )
+
+ if self.honesty:
+ self.honestApproach(X_est, treatment_idx_est, y_est)
+
+ self.feature_importances_ = np.zeros(X.shape[1])
+ for col, imp in self.feature_imp_dict.items():
+ self.feature_importances_[col] = imp
+ self.feature_importances_ /= self.feature_importances_.sum() # normalize to add to 1
+
+ # Prune Trees
+ def prune(self, X, treatment, y, minGain=0.0001, rule='maxAbsDiff'):
+ """ Prune the uplift model.
+ Args
+ ----
+ X : ndarray, shape = [num_samples, num_features]
+ An ndarray of the covariates used to train the uplift model.
+ treatment : array-like, shape = [num_samples]
+ An array containing the treatment group for each unit.
+ y : array-like, shape = [num_samples]
+ An array containing the outcome of interest for each unit.
+ minGain : float, optional (default = 0.0001)
+ The minimum gain required to make a tree node split. The children
+ tree branches are trimmed if the actual split gain is less than
+ the minimum gain.
+ rule : string, optional (default = 'maxAbsDiff')
+ The prune rules. Supported values are 'maxAbsDiff' for optimizing
+ the maximum absolute difference, and 'bestUplift' for optimizing
+ the node-size weighted treatment effect.
+ Returns
+ -------
+ self : object
+ """
+
+ X, y = check_X_y(X, y)
+ treatment = np.asarray(treatment)
+ assert len(y) == len(treatment), 'Data length must be equal for X, treatment, and y.'
+
+ # Get treatment group keys. self.classes_[0] is reserved for the control group.
+ treatment_idx = np.zeros_like(treatment)
+ for i, tr in enumerate(self.classes_[1:], 1):
+ treatment_idx[treatment == tr] = i
+
+ self.pruneTree(X, treatment_idx, y,
+ tree=self.fitted_uplift_tree,
+ rule=rule,
+ minGain=minGain,
+ n_reg=self.n_reg,
+ parentNodeSummary=None)
+ return self
+
+ def honestApproach(self, X_est, T_est, Y_est):
+ """ Apply the honest approach based on "Athey, S., & Imbens, G. (2016). Recursive partitioning for heterogeneous causal effects."
+ Args
+ ----
+ X_est : ndarray, shape = [num_samples, num_features]
+ An ndarray of the covariates used to calculate the unbiased estimates in the leafs of the decision tree.
+ T_est : array-like, shape = [num_samples]
+ An array containing the treatment group for each unit.
+ Y_est : array-like, shape = [num_samples]
+ An array containing the outcome of interest for each unit.
+ """
+
+ self.fillTree(X_est, T_est, Y_est, self.fitted_uplift_tree)
+
+ def pruneTree(self, X, treatment_idx, y, tree, rule='maxAbsDiff', minGain=0.,
+ n_reg=0,
+ parentNodeSummary=None):
+ """Prune one single tree node in the uplift model.
+ Args
+ ----
+ X : ndarray, shape = [num_samples, num_features]
+ An ndarray of the covariates used to train the uplift model.
+ treatment_idx : array-like, shape = [num_samples]
+ An array containing the treatment group index for each unit.
+ y : array-like, shape = [num_samples]
+ An array containing the outcome of interest for each unit.
+ rule : string, optional (default = 'maxAbsDiff')
+ The prune rules. Supported values are 'maxAbsDiff' for optimizing the maximum absolute difference, and
+ 'bestUplift' for optimizing the node-size weighted treatment effect.
+ minGain : float, optional (default = 0.)
+ The minimum gain required to make a tree node split. The children tree branches are trimmed if the actual
+ split gain is less than the minimum gain.
+ n_reg: int, optional (default=0)
+ The regularization parameter defined in Rzepakowski et al. 2012, the weight (in terms of sample size) of the
+ parent node influence on the child node, only effective for 'KL', 'ED', 'Chi', 'CTS' methods.
+ parentNodeSummary : list of list, optional (default = None)
+ Node summary statistics, [P(Y=1|T), N(T)] of the parent tree node.
+ Returns
+ -------
+ self : object
+ """
+ # Current Node Summary for Validation Data Set
+ currentNodeSummary = self.tree_node_summary(
+ treatment_idx, y, min_samples_treatment=self.min_samples_treatment,
+ n_reg=n_reg, parentNodeSummary=parentNodeSummary
+ )
+ tree.nodeSummary = currentNodeSummary
+ # Divide sets for child nodes
+ if (tree.trueBranch is None) or (tree.falseBranch is None):
+ X_l, X_r, w_l, w_r, y_l, y_r = self.divideSet(X, treatment_idx, y, tree.col, tree.value)
+
+ # recursive call for each branch
+ if tree.trueBranch.results is None:
+ self.pruneTree(X_l, w_l, y_l, tree.trueBranch, rule, minGain,
+ n_reg,
+ parentNodeSummary=currentNodeSummary)
+ if tree.falseBranch.results is None:
+ self.pruneTree(X_r, w_r, y_r, tree.falseBranch, rule, minGain,
+ n_reg,
+ parentNodeSummary=currentNodeSummary)
+
+ # merge leaves (potentially)
+ if (tree.trueBranch.results is not None and
+ tree.falseBranch.results is not None):
+ if rule == 'maxAbsDiff':
+ # Current D
+ if (tree.maxDiffTreatment in currentNodeSummary and
+ self.control_name in currentNodeSummary):
+ currentScoreD = tree.maxDiffSign * (currentNodeSummary[tree.maxDiffTreatment][0]
+ - currentNodeSummary[self.control_name][0])
+ else:
+ currentScoreD = 0
+
+ # trueBranch D
+ trueNodeSummary = self.tree_node_summary(
+ w_l, y_l, min_samples_treatment=self.min_samples_treatment,
+ n_reg=n_reg, parentNodeSummary=currentNodeSummary
+ )
+ if (tree.trueBranch.maxDiffTreatment in trueNodeSummary and
+ self.control_name in trueNodeSummary):
+ trueScoreD = tree.trueBranch.maxDiffSign * (trueNodeSummary[tree.trueBranch.maxDiffTreatment][0]
+ - trueNodeSummary[self.control_name][0])
+ trueScoreD = (
+ trueScoreD
+ * (trueNodeSummary[tree.trueBranch.maxDiffTreatment][1]
+ + trueNodeSummary[self.control_name][1])
+ / (currentNodeSummary[tree.trueBranch.maxDiffTreatment][1]
+ + currentNodeSummary[self.control_name][1])
+ )
+ else:
+ trueScoreD = 0
+
+ # falseBranch D
+ falseNodeSummary = self.tree_node_summary(
+ w_r, y_r, min_samples_treatment=self.min_samples_treatment,
+ n_reg=n_reg, parentNodeSummary=currentNodeSummary
+ )
+ if (tree.falseBranch.maxDiffTreatment in falseNodeSummary and
+ self.control_name in falseNodeSummary):
+ falseScoreD = (
+ tree.falseBranch.maxDiffSign *
+ (falseNodeSummary[tree.falseBranch.maxDiffTreatment][0]
+ - falseNodeSummary[self.control_name][0])
+ )
+
+ falseScoreD = (
+ falseScoreD *
+ (falseNodeSummary[tree.falseBranch.maxDiffTreatment][1]
+ + falseNodeSummary[self.control_name][1])
+ / (currentNodeSummary[tree.falseBranch.maxDiffTreatment][1]
+ + currentNodeSummary[self.control_name][1])
+ )
+ else:
+ falseScoreD = 0
+
+ if ((trueScoreD + falseScoreD) - currentScoreD <= minGain or
+ (trueScoreD + falseScoreD < 0.)):
+ tree.trueBranch, tree.falseBranch = None, None
+ tree.results = tree.backupResults
+
+ elif rule == 'bestUplift':
+ # Current D
+ if (tree.bestTreatment in currentNodeSummary and
+ self.control_name in currentNodeSummary):
+ currentScoreD = (
+ currentNodeSummary[tree.bestTreatment][0]
+ - currentNodeSummary[self.control_name][0]
+ )
+ else:
+ currentScoreD = 0
+
+ # trueBranch D
+ trueNodeSummary = self.tree_node_summary(
+ w_l, y_l, min_samples_treatment=self.min_samples_treatment,
+ n_reg=n_reg, parentNodeSummary=currentNodeSummary
+ )
+ if (tree.trueBranch.bestTreatment in trueNodeSummary and
+ self.control_name in trueNodeSummary):
+ trueScoreD = (
+ trueNodeSummary[tree.trueBranch.bestTreatment][0]
+ - trueNodeSummary[self.control_name][0]
+ )
+ else:
+ trueScoreD = 0
+
+ # falseBranch D
+ falseNodeSummary = self.tree_node_summary(
+ w_r, y_r, min_samples_treatment=self.min_samples_treatment,
+ n_reg=n_reg, parentNodeSummary=currentNodeSummary
+ )
+ if (tree.falseBranch.bestTreatment in falseNodeSummary and
+ self.control_name in falseNodeSummary):
+ falseScoreD = (
+ falseNodeSummary[tree.falseBranch.bestTreatment][0]
+ - falseNodeSummary[self.control_name][0]
+ )
+ else:
+ falseScoreD = 0
+ gain = ((1. * len(y_l) / len(y) * trueScoreD
+ + 1. * len(y_r) / len(y) * falseScoreD)
+ - currentScoreD)
+ if gain <= minGain or (trueScoreD + falseScoreD < 0.):
+ tree.trueBranch, tree.falseBranch = None, None
+ tree.results = tree.backupResults
+ return self
+
+ def fill(self, X, treatment, y):
+ """ Fill the data into an existing tree.
+ This is a higher-level function to transform the original data inputs
+ into lower level data inputs (list of list and tree).
+
+ Args
+ ----
+ X : ndarray, shape = [num_samples, num_features]
+ An ndarray of the covariates used to train the uplift model.
+ treatment : array-like, shape = [num_samples]
+ An array containing the treatment group for each unit.
+ y : array-like, shape = [num_samples]
+ An array containing the outcome of interest for each unit.
+
+ Returns
+ -------
+ self : object
+ """
+
+ X, y = check_X_y(X, y)
+ treatment = np.asarray(treatment)
+ assert len(y) == len(treatment), 'Data length must be equal for X, treatment, and y.'
+
+ # Get treatment group keys. self.classes_[0] is reserved for the control group.
+ treatment_idx = np.zeros_like(treatment)
+ for i, tr in enumerate(self.classes_[1:], 1):
+ treatment_idx[treatment == tr] = i
+
+ self.fillTree(X, treatment_idx, y, tree=self.fitted_uplift_tree)
+ return self
+
+ def fillTree(self, X, treatment_idx, y, tree):
+ """ Fill the data into an existing tree.
+ This is a lower-level function to execute on the tree filling task.
+
+ Args
+ ----
+ X : ndarray, shape = [num_samples, num_features]
+ An ndarray of the covariates used to train the uplift model.
+ treatment_idx : array-like, shape = [num_samples]
+ An array containing the treatment group index for each unit.
+ y : array-like, shape = [num_samples]
+ An array containing the outcome of interest for each unit.
+ tree : object
+ object of DecisionTree class
+
+ Returns
+ -------
+ self : object
+ """
+ # Current Node Summary for Validation Data Set
+ currentNodeSummary = self.tree_node_summary(treatment_idx, y,
+ min_samples_treatment=0,
+ n_reg=0,
+ parentNodeSummary=None)
+ tree.nodeSummary = currentNodeSummary
+
+ # Divide sets for child nodes
+ if tree.trueBranch or tree.falseBranch:
+ X_l, X_r, w_l, w_r, y_l, y_r = self.divideSet(X, treatment_idx, y, tree.col, tree.value)
+
+ # recursive call for each branch
+ if tree.trueBranch is not None:
+ self.fillTree(X_l, w_l, y_l, tree.trueBranch)
+ if tree.falseBranch is not None:
+ self.fillTree(X_r, w_r, y_r, tree.falseBranch)
+
+ # Update Information
+
+ # matchScore
+ matchScore = (currentNodeSummary[tree.bestTreatment][0] - currentNodeSummary[0][0])
+ tree.matchScore = round(matchScore, 4)
+ tree.summary['matchScore'] = round(matchScore, 4)
+
+ # Samples, Group_size
+ tree.summary['samples'] = len(y)
+ tree.summary['group_size'] = ''
+ for treatment_group, summary in zip(self.classes_, currentNodeSummary):
+ tree.summary['group_size'] += ' ' + treatment_group + ': ' + str(summary[1])
+ # classProb
+ if tree.results is not None:
+ tree.results = self.uplift_classification_results(treatment_idx, y)
+ return self
+
+ def predict(self, X):
+ '''
+ Returns the recommended treatment group and predicted optimal
+ probability conditional on using the recommended treatment group.
+
+ Args
+ ----
+ X : ndarray, shape = [num_samples, num_features]
+ An ndarray of the covariates used to train the uplift model.
+
+ Returns
+ -------
+ pred: ndarray, shape = [num_samples, num_treatments]
+ An ndarray of predicted treatment effects across treatments.
+ '''
+
+ X = check_array(X)
+
+ pred_nodes = []
+ for i_row in range(len(X)):
+ pred_leaf, _ = self.classify(X[i_row], self.fitted_uplift_tree, dataMissing=False)
+ pred_nodes.append(pred_leaf)
+ return np.array(pred_nodes)
+
+ @staticmethod
+ def divideSet(X, treatment_idx, y, column, value):
+ '''
+ Tree node split.
+
+ Args
+ ----
+ X : ndarray, shape = [num_samples, num_features]
+ An ndarray of the covariates used to train the uplift model.
+ treatment_idx : array-like, shape = [num_samples]
+ An array containing the treatment group index for each unit.
+ y : array-like, shape = [num_samples]
+ An array containing the outcome of interest for each unit.
+ column : int
+ The column used to split the data.
+ value : float or int
+ The value in the column for splitting the data.
+
+ Returns
+ -------
+ (X_l, X_r, treatment_l, treatment_r, y_l, y_r) : list of ndarray
+ The covariates, treatments and outcomes of left node and the right node.
+ '''
+ # for int and float values
+ if np.issubdtype(value.dtype, np.number):
+ filt = X[:, column] >= value
+ else: # for strings
+ filt = X[:, column] == value
+
+ return X[filt], X[~filt], treatment_idx[filt], treatment_idx[~filt], y[filt], y[~filt]
+
+ @staticmethod
+ def divideSet_len(X, treatment_idx, y, column, value):
+ '''Tree node split.
+
+ Modified from dividedSet(), but return the len(X_l) and
+ len(X_r) instead of the split X_l and X_r, to avoid some
+ overhead, intended to be used for finding the split. After
+ finding the best splits, can split to find the X_l and X_r.
+
+ Args
+ ----
+ X : ndarray, shape = [num_samples, num_features]
+ An ndarray of the covariates used to train the uplift model.
+ treatment_idx : array-like, shape = [num_samples]
+ An array containing the treatment group index for each unit.
+ y : array-like, shape = [num_samples]
+ An array containing the outcome of interest for each unit.
+ column : int
+ The column used to split the data.
+ value : float or int
+ The value in the column for splitting the data.
+
+ Returns
+ -------
+ (len_X_l, len_X_r, treatment_l, treatment_r, y_l, y_r) : list of ndarray
+ The covariates nrows, treatments and outcomes of left node and the right node.
+
+ '''
+ # for int and float values
+ if np.issubdtype(value.dtype, np.number):
+ filt = X[:, column] >= value
+ else: # for strings
+ filt = X[:, column] == value
+
+ len_X_l = np.sum(filt)
+ return len_X_l, len(X) - len_X_l, treatment_idx[filt], treatment_idx[~filt], y[filt], y[~filt]
+
+ def group_uniqueCounts(self, treatment_idx, y):
+ '''
+ Count sample size by experiment group.
+
+ Args
+ ----
+ treatment_idx : array-like, shape = [num_samples]
+ An array containing the treatment group index for each unit.
+ y : array-like, shape = [num_samples]
+ An array containing the outcome of interest for each unit.
+
+ Returns
+ -------
+ results : list of list
+ The negative and positive outcome sample sizes for each of the control and treatment groups.
+ '''
+ results = []
+ for i in range(self.n_class):
+ filt = treatment_idx == i
+ n_pos = y[filt].sum()
+
+ # [N(Y = 0, T = 1), N(Y = 1, T = 1)]
+ results.append([filt.sum() - n_pos, n_pos])
+
+ return results
+
+ @staticmethod
+ def evaluate_KL(nodeSummary):
+ '''
+ Calculate KL Divergence as split evaluation criterion for a given node.
+
+ Args
+ ----
+ nodeSummary : list of list
+ The tree node summary statistics, [P(Y=1|T), N(T)], produced by tree_node_summary()
+ method.
+
+ Returns
+ -------
+ d_res : KL Divergence
+ '''
+ p_c = nodeSummary[0][0]
+ d_res = 0.
+ for treatment_group in nodeSummary[1:]:
+ d_res += kl_divergence(treatment_group[0], p_c)
+ return d_res
+
+ @staticmethod
+ def arr_evaluate_KL(np.ndarray[P_TYPE_t, ndim=1] node_summary_p,
+ np.ndarray[N_TYPE_t, ndim=1] node_summary_n):
+ '''
+ Calculate KL Divergence as split evaluation criterion for a given node.
+ Modified to accept new node summary format.
+
+ Args
+ ----
+ node_summary_p : array of shape [n_class]
+ Has type numpy.double.
+ The positive probabilities of each of the control
+ and treament groups of the current node, i.e. [P(Y=1|T=i)...]
+ node_summary_n : array of shape [n_class]
+ Has type numpy.int32.
+ The counts of each of the control
+ and treament groups of the current node, i.e. [N(T=i)...]
+
+ Returns
+ -------
+ d_res : KL Divergence
+ '''
+ cdef int n_class = node_summary_p.shape[0]
+ cdef P_TYPE_t p_c = node_summary_p[0]
+ cdef P_TYPE_t d_res = 0.0
+ cdef int i = 0
+ for i in range(1, n_class):
+ d_res += kl_divergence(node_summary_p[i], p_c)
+ return d_res
+
+ @staticmethod
+ def evaluate_ED(nodeSummary):
+ '''
+ Calculate Euclidean Distance as split evaluation criterion for a given node.
+
+ Args
+ ----
+ nodeSummary : dictionary
+ The tree node summary statistics, produced by tree_node_summary()
+ method.
+
+ Returns
+ -------
+ d_res : Euclidean Distance
+ '''
+ pc = nodeSummary[0][0]
+ d_res = 0
+ for treatment_group in nodeSummary[1:]:
+ d_res += 2*(treatment_group[0] - pc)**2
+ return d_res
+
+ @staticmethod
+ def arr_evaluate_ED(np.ndarray[P_TYPE_t, ndim=1] node_summary_p,
+ np.ndarray[N_TYPE_t, ndim=1] node_summary_n):
+ '''
+ Calculate Euclidean Distance as split evaluation criterion for a given node.
+
+ Args
+ ----
+ node_summary_p : array of shape [n_class]
+ Has type numpy.double.
+ The positive probabilities of each of the control
+ and treament groups of the current node, i.e. [P(Y=1|T=i)...]
+ node_summary_n : array of shape [n_class]
+ Has type numpy.int32.
+ The counts of each of the control
+ and treament groups of the current node, i.e. [N(T=i)...]
+
+ Returns
+ -------
+ d_res : Euclidean Distance
+ '''
+ cdef int n_class = node_summary_p.shape[0]
+ cdef P_TYPE_t p_c = node_summary_p[0]
+ cdef P_TYPE_t d_res = 0.0
+ cdef int i = 0
+ for i in range(1, n_class):
+ d_res += 2*(node_summary_p[i] - p_c)*(node_summary_p[i] - p_c)
+ return d_res
+
+ @staticmethod
+ def evaluate_Chi(nodeSummary):
+ '''
+ Calculate Chi-Square statistic as split evaluation criterion for a given node.
+
+ Args
+ ----
+ nodeSummary : dictionary
+ The tree node summary statistics, produced by tree_node_summary() method.
+
+ Returns
+ -------
+ d_res : Chi-Square
+ '''
+ pc = nodeSummary[0][0]
+ d_res = 0
+ for treatment_group in nodeSummary[1:]:
+ d_res += ((treatment_group[0] - pc) ** 2 / max(0.1 ** 6, pc)
+ + (treatment_group[0] - pc) ** 2 / max(0.1 ** 6, 1 - pc))
+ return d_res
+
+ @staticmethod
+ def arr_evaluate_Chi(np.ndarray[P_TYPE_t, ndim=1] node_summary_p,
+ np.ndarray[N_TYPE_t, ndim=1] node_summary_n):
+ '''
+ Calculate Chi-Square statistic as split evaluation criterion for a given node.
+
+ Args
+ ----
+ node_summary_p : array of shape [n_class]
+ Has type numpy.double.
+ The positive probabilities of each of the control
+ and treament groups of the current node, i.e. [P(Y=1|T=i)...]
+ node_summary_n : array of shape [n_class]
+ Has type numpy.int32.
+ The counts of each of the control
+ and treament groups of the current node, i.e. [N(T=i)...]
+
+ Returns
+ -------
+ d_res : Chi-Square
+ '''
+ cdef int n_class = node_summary_p.shape[0]
+ cdef P_TYPE_t p_c = node_summary_p[0]
+ cdef P_TYPE_t d_res = 0.0
+ cdef int i = 0
+ cdef P_TYPE_t max_eps_pc = max(0.1 ** 6, p_c)
+ cdef P_TYPE_t max_eps_1_pc = max(0.1 ** 6, 1 - p_c)
+ cdef P_TYPE_t diff_sq = 0.0
+ for i in range(1, n_class):
+ diff_sq = (node_summary_p[i] - p_c) * (node_summary_p[i] - p_c)
+ d_res += (diff_sq / max_eps_pc + diff_sq / max_eps_1_pc)
+ return d_res
+
+ @staticmethod
+ def evaluate_DDP(nodeSummary):
+ '''
+ Calculate Delta P as split evaluation criterion for a given node.
+
+ Args
+ ----
+ nodeSummary : list of list
+ The tree node summary statistics, [P(Y=1|T), N(T)], produced by tree_node_summary() method.
+
+ Returns
+ -------
+ d_res : Delta P
+ '''
+ pc = nodeSummary[0][0]
+ d_res = 0
+ for treatment_group in nodeSummary[1:]:
+ d_res += treatment_group[0] - pc
+ return d_res
+
+ @staticmethod
+ def arr_evaluate_DDP(np.ndarray[P_TYPE_t, ndim=1] node_summary_p,
+ np.ndarray[N_TYPE_t, ndim=1] node_summary_n):
+ '''
+ Calculate Delta P as split evaluation criterion for a given node.
+
+ Args
+ ----
+ node_summary_p : array of shape [n_class]
+ Has type numpy.double.
+ The positive probabilities of each of the control
+ and treament groups of the current node, i.e. [P(Y=1|T=i)...]
+ node_summary_n : array of shape [n_class]
+ Has type numpy.int32.
+ The counts of each of the control
+ and treament groups of the current node, i.e. [N(T=i)...]
+
+ Returns
+ -------
+ d_res : Delta P
+ '''
+ cdef int n_class = node_summary_p.shape[0]
+ cdef P_TYPE_t p_c = node_summary_p[0]
+ cdef P_TYPE_t d_res = 0.0
+ cdef int i = 0
+ for i in range(1, n_class):
+ d_res += node_summary_p[i] - p_c
+ return d_res
+
+ @staticmethod
+ def evaluate_IT(leftNodeSummary, rightNodeSummary, w_l, w_r):
+ '''
+ Calculate Squared T-Statistic as split evaluation criterion for a given node
+
+ Args
+ ----
+ leftNodeSummary : list of list
+ The left node summary statistics.
+ rightNodeSummary : list of list
+ The right node summary statistics.
+ w_l: array-like, shape = [num_samples]
+ An array containing the treatment for each unit in the left node
+ w_r: array-like, shape = [num_samples]
+ An array containing the treatment for each unit in the right node
+
+ Returns
+ -------
+ g_s : Squared T-Statistic
+ '''
+ g_s = 0
+
+ ## Control Group
+ # Sample mean in left & right child node
+ y_l_0 = leftNodeSummary[0][0]
+ y_r_0 = rightNodeSummary[0][0]
+ # Sample size left & right child node
+ n_3 = leftNodeSummary[0][1]
+ n_4 = rightNodeSummary[0][1]
+ # Sample variance in left & right child node (p*(p-1) for bernoulli)
+ s_3 = y_l_0*(1-y_l_0)
+ s_4 = y_r_0*(1-y_r_0)
+
+ for treatment_left, treatment_right in zip(leftNodeSummary[1:], rightNodeSummary[1:]):
+ ## Treatment Group
+ # Sample mean in left & right child node
+ y_l_1 = treatment_left[0]
+ y_r_1 = treatment_right[0]
+ # Sample size left & right child node
+ n_1 = treatment_left[1]
+ n_2 = treatment_right[1]
+ # Sample variance in left & right child node
+ s_1 = y_l_1*(1-y_l_1)
+ s_2 = y_r_1*(1-y_r_1)
+
+ sum_n = np.sum([n_1 - 1, n_2 - 1, n_3 - 1, n_4 - 1])
+ w_1 = (n_1 - 1) / sum_n
+ w_2 = (n_2 - 1) / sum_n
+ w_3 = (n_3 - 1) / sum_n
+ w_4 = (n_4 - 1) / sum_n
+
+ # Pooled estimator of the constant variance
+ sigma = np.sqrt(np.sum([w_1 * s_1, w_2 * s_2, w_3 * s_3, w_4 * s_4]))
+
+ # Squared t-statistic
+ g_s = np.power(((y_l_1 - y_l_0) - (y_r_1 - y_r_0)) / (sigma * np.sqrt(np.sum([1 / n_1, 1 / n_2, 1 / n_3, 1 / n_4]))), 2)
+
+ return g_s
+
+ @staticmethod
+ def arr_evaluate_IT(np.ndarray[P_TYPE_t, ndim=1] left_node_summary_p,
+ np.ndarray[N_TYPE_t, ndim=1] left_node_summary_n,
+ np.ndarray[P_TYPE_t, ndim=1] right_node_summary_p,
+ np.ndarray[N_TYPE_t, ndim=1] right_node_summary_n):
+ '''
+ Calculate Squared T-Statistic as split evaluation criterion for a given node
+
+ NOTE: n_class should be 2.
+
+ Args
+ ----
+ left_node_summary_p : array of shape [n_class]
+ Has type numpy.double.
+ The positive probabilities of each of the control
+ and treament groups of the left node, i.e. [P(Y=1|T=i)...]
+ left_node_summary_n : array of shape [n_class]
+ Has type numpy.int32.
+ The counts of each of the control
+ and treament groups of the left node, i.e. [N(T=i)...]
+ right_node_summary_p : array of shape [n_class]
+ Has type numpy.double.
+ The positive probabilities of each of the control
+ and treament groups of the right node, i.e. [P(Y=1|T=i)...]
+ right_node_summary_n : array of shape [n_class]
+ Has type numpy.int32.
+ The counts of each of the control
+ and treament groups of the right node, i.e. [N(T=i)...]
+
+ Returns
+ -------
+ g_s : Squared T-Statistic
+ '''
+ ## Control Group
+ # Sample mean in left & right child node
+ cdef P_TYPE_t y_l_0 = left_node_summary_p[0]
+ cdef P_TYPE_t y_r_0 = right_node_summary_p[0]
+ # Sample size left & right child node
+ cdef N_TYPE_t n_3 = left_node_summary_n[0]
+ cdef N_TYPE_t n_4 = right_node_summary_n[0]
+ # Sample variance in left & right child node (p*(p-1) for bernoulli)
+ cdef P_TYPE_t s_3 = y_l_0*(1-y_l_0)
+ cdef P_TYPE_t s_4 = y_r_0*(1-y_r_0)
+
+ # only one treatment, contrast with control, so no need to loop
+ ## Treatment Group
+ # Sample mean in left & right child node
+ cdef P_TYPE_t y_l_1 = left_node_summary_p[1]
+ cdef P_TYPE_t y_r_1 = right_node_summary_p[1]
+ # Sample size left & right child node
+ cdef N_TYPE_t n_1 = left_node_summary_n[1]
+ cdef N_TYPE_t n_2 = right_node_summary_n[1]
+ # Sample variance in left & right child node
+ cdef P_TYPE_t s_1 = y_l_1*(1-y_l_1)
+ cdef P_TYPE_t s_2 = y_r_1*(1-y_r_1)
+
+ cdef P_TYPE_t sum_n = (n_1 - 1) + (n_2 - 1) + (n_3 - 1) + (n_4 - 1)
+ cdef P_TYPE_t w_1 = (n_1 - 1) / sum_n
+ cdef P_TYPE_t w_2 = (n_2 - 1) / sum_n
+ cdef P_TYPE_t w_3 = (n_3 - 1) / sum_n
+ cdef P_TYPE_t w_4 = (n_4 - 1) / sum_n
+
+ # Pooled estimator of the constant variance
+ cdef P_TYPE_t sigma = sqrt(w_1 * s_1 + w_2 * s_2 + w_3 * s_3 + w_4 * s_4)
+
+ # Squared t-statistic
+ cdef P_TYPE_t g_s = ((y_l_1 - y_l_0) - (y_r_1 - y_r_0)) / (sigma * sqrt(1.0 / n_1 + 1.0 / n_2 + 1.0 / n_3 + 1.0 / n_4))
+ g_s = g_s * g_s
+
+ return g_s
+
+ @staticmethod
+ def evaluate_CIT(currentNodeSummary, leftNodeSummary, rightNodeSummary, y_l, y_r, w_l, w_r, y, w):
+ '''
+ Calculate likelihood ratio test statistic as split evaluation criterion for a given node
+ Args
+ ----
+ currentNodeSummary: list of lists
+ The parent node summary statistics
+ leftNodeSummary : list of lists
+ The left node summary statistics.
+ rightNodeSummary : list of lists
+ The right node summary statistics.
+ y_l: array-like, shape = [num_samples]
+ An array containing the outcome of interest for each unit in the left node
+ y_r: array-like, shape = [num_samples]
+ An array containing the outcome of interest for each unit in the right node
+ w_l: array-like, shape = [num_samples]
+ An array containing the treatment for each unit in the left node
+ w_r: array-like, shape = [num_samples]
+ An array containing the treatment for each unit in the right node
+ y: array-like, shape = [num_samples]
+ An array containing the outcome of interest for each unit
+ w: array-like, shape = [num_samples]
+ An array containing the treatment for each unit
+ Returns
+ -------
+ lrt : Likelihood ratio test statistic
+ '''
+ lrt = 0
+
+ # Control sample size left & right child node
+ n_l_t_0 = leftNodeSummary[0][1]
+ n_r_t_0 = rightNodeSummary[0][1]
+
+ for treatment_left, treatment_right in zip(leftNodeSummary[1:], rightNodeSummary[1:]):
+ # Treatment sample size left & right child node
+ n_l_t_1 = treatment_left[1]
+ n_r_t_1 = treatment_right[1]
+
+ # Total size of left & right node
+ n_l_t = n_l_t_1 + n_l_t_0
+ n_r_t = n_r_t_1 + n_r_t_0
+
+ # Total size of parent node
+ n_t = n_l_t + n_r_t
+
+ # Total treatment & control size in parent node
+ n_t_1 = n_l_t_1 + n_r_t_1
+ n_t_0 = n_l_t_0 + n_r_t_0
+
+ # Standard squared error of left child node
+ sse_tau_l = np.sum(np.power(y_l[w_l == 1] - treatment_left[0], 2)) + np.sum(
+ np.power(y_l[w_l == 0] - treatment_left[0], 2))
+
+ # Standard squared error of right child node
+ sse_tau_r = np.sum(np.power(y_r[w_r == 1] - treatment_right[0], 2)) + np.sum(
+ np.power(y_r[w_r == 0] - treatment_right[0], 2))
+
+ # Standard squared error of parent child node
+ sse_tau = np.sum(np.power(y[w == 1] - currentNodeSummary[1][0], 2)) + np.sum(
+ np.power(y[w == 0] - currentNodeSummary[0][0], 2))
+
+ # Maximized log-likelihood function
+ i_tau_l = - (n_l_t / 2) * np.log(n_l_t * sse_tau_l) + n_l_t_1 * np.log(n_l_t_1) + n_l_t_0 * np.log(n_l_t_0)
+ i_tau_r = - (n_r_t / 2) * np.log(n_r_t * sse_tau_r) + n_r_t_1 * np.log(n_r_t_1) + n_r_t_0 * np.log(n_r_t_0)
+ i_tau = - (n_t / 2) * np.log(n_t * sse_tau) + n_t_1 * np.log(n_t_1) + n_t_0 * np.log(n_t_0)
+
+ # Likelihood ration test statistic
+ lrt = 2 * (i_tau_l + i_tau_r - i_tau)
+
+ return lrt
+
+ @staticmethod
+ def arr_evaluate_CIT(np.ndarray[P_TYPE_t, ndim=1] cur_node_summary_p,
+ np.ndarray[N_TYPE_t, ndim=1] cur_node_summary_n,
+ np.ndarray[P_TYPE_t, ndim=1] left_node_summary_p,
+ np.ndarray[N_TYPE_t, ndim=1] left_node_summary_n,
+ np.ndarray[P_TYPE_t, ndim=1] right_node_summary_p,
+ np.ndarray[N_TYPE_t, ndim=1] right_node_summary_n):
+ '''
+ Calculate likelihood ratio test statistic as split evaluation criterion for a given node
+
+ NOTE: n_class should be 2.
+
+ Args
+ ----
+ cur_node_summary_p : array of shape [n_class]
+ Has type numpy.double.
+ The positive probabilities of each of the control
+ and treament groups of the current node, i.e. [P(Y=1|T=i)...]
+ cur_node_summary_n : array of shape [n_class]
+ Has type numpy.int32.
+ The counts of each of the control
+ and treament groups of the current node, i.e. [N(T=i)...]
+ left_node_summary_p : array of shape [n_class]
+ Has type numpy.double.
+ The positive probabilities of each of the control
+ and treament groups of the left node, i.e. [P(Y=1|T=i)...]
+ left_node_summary_n : array of shape [n_class]
+ Has type numpy.int32.
+ The counts of each of the control
+ and treament groups of the left node, i.e. [N(T=i)...]
+ right_node_summary_p : array of shape [n_class]
+ Has type numpy.double.
+ The positive probabilities of each of the control
+ and treament groups of the right node, i.e. [P(Y=1|T=i)...]
+ right_node_summary_n : array of shape [n_class]
+ Has type numpy.int32.
+ The counts of each of the control
+ and treament groups of the right node, i.e. [N(T=i)...]
+
+ Returns
+ -------
+ lrt : Likelihood ratio test statistic
+ '''
+ cdef P_TYPE_t lrt = 0.0
+
+ # since will take log of these N, so use a double type
+
+ # Control sample size left & right child node
+ cdef P_TYPE_t n_l_t_0 = left_node_summary_n[0]
+ cdef P_TYPE_t n_r_t_0 = right_node_summary_n[0]
+
+ # Treatment sample size left & right child node
+ cdef P_TYPE_t n_l_t_1 = left_node_summary_n[1]
+ cdef P_TYPE_t n_r_t_1 = right_node_summary_n[1]
+
+ # Total size of left & right node
+ cdef P_TYPE_t n_l_t = n_l_t_1 + n_l_t_0
+ cdef P_TYPE_t n_r_t = n_r_t_1 + n_r_t_0
+
+ # Total size of parent node
+ cdef P_TYPE_t n_t = n_l_t + n_r_t
+
+ # Total treatment & control size in parent node
+ cdef P_TYPE_t n_t_1 = n_l_t_1 + n_r_t_1
+ cdef P_TYPE_t n_t_0 = n_l_t_0 + n_r_t_0
+
+ # NOTE: the original code for sse_tau_l and sse_tau_r does not seem to follow the paper.
+ # sse = \sum_{i for treatment} (y_i - p_treatment)^2 + \sum_{i for control} (y_i - p_control)^2
+
+ # NOTE: since for classification, the y is either 0 or 1, we can calculate sse more simply
+ # for y being 0 or 1, sse = n*p*(1-p), but here need to calculate separately for treatment and control groups.
+
+ # Standard squared error of left child node
+ cdef P_TYPE_t sse_tau_l = n_l_t_0 * left_node_summary_p[0] * (1.0 - left_node_summary_p[0]) + n_l_t_1 * left_node_summary_p[1] * (1.0 - left_node_summary_p[1])
+
+ # Standard squared error of right child node
+ cdef P_TYPE_t sse_tau_r = n_r_t_0 * right_node_summary_p[0] * (1.0 - right_node_summary_p[0]) + n_r_t_1 * right_node_summary_p[1] * (1.0 - right_node_summary_p[1])
+
+ # Standard squared error of parent child node
+ cdef P_TYPE_t sse_tau = n_t_0 * cur_node_summary_p[0] * (1.0 - cur_node_summary_p[0]) + n_t_1 * cur_node_summary_p[1] * (1.0 - cur_node_summary_p[1])
+
+ # Maximized log-likelihood function
+ cdef P_TYPE_t i_tau_l = - (n_l_t / 2.0) * log(n_l_t * sse_tau_l) + n_l_t_1 * log(n_l_t_1) + n_l_t_0 * log(n_l_t_0)
+ cdef P_TYPE_t i_tau_r = - (n_r_t / 2.0) * log(n_r_t * sse_tau_r) + n_r_t_1 * log(n_r_t_1) + n_r_t_0 * log(n_r_t_0)
+ cdef P_TYPE_t i_tau = - (n_t / 2.0) * log(n_t * sse_tau) + n_t_1 * log(n_t_1) + n_t_0 * log(n_t_0)
+
+ # Likelihood ration test statistic
+ lrt = 2 * (i_tau_l + i_tau_r - i_tau)
+
+ return lrt
+
+ @staticmethod
+ def evaluate_IDDP(nodeSummary):
+ '''
+ Calculate Delta P as split evaluation criterion for a given node.
+
+ Args
+ ----
+ nodeSummary : dictionary
+ The tree node summary statistics, produced by tree_node_summary() method.
+ control_name : string
+ The control group name.
+ Returns
+ -------
+ d_res : Delta P
+ '''
+ pc = nodeSummary[0][0]
+ d_res = 0
+ for treatment_group in nodeSummary[1:]:
+ d_res += treatment_group[0] - pc
+ return d_res
+
+ @staticmethod
+ def arr_evaluate_IDDP(np.ndarray[P_TYPE_t, ndim=1] node_summary_p,
+ np.ndarray[N_TYPE_t, ndim=1] node_summary_n):
+ '''
+ Calculate Delta P as split evaluation criterion for a given node.
+
+ Args
+ ----
+ node_summary_p : array of shape [n_class]
+ Has type numpy.double.
+ The positive probabilities of each of the control
+ and treament groups of the current node, i.e. [P(Y=1|T=i)...]
+ node_summary_n : array of shape [n_class]
+ Has type numpy.int32.
+ The counts of each of the control
+ and treament groups of the current node, i.e. [N(T=i)...]
+
+ Returns
+ -------
+ d_res : Delta P
+ '''
+ cdef int n_class = node_summary_p.shape[0]
+ cdef P_TYPE_t p_c = node_summary_p[0]
+ cdef P_TYPE_t d_res = 0.0
+ cdef int i = 0
+ for i in range(1, n_class):
+ d_res += node_summary_p[i] - p_c
+ return d_res
+
+ @staticmethod
+ def evaluate_CTS(nodeSummary):
+ '''
+ Calculate CTS (conditional treatment selection) as split evaluation criterion for a given node.
+
+ Args
+ ----
+ nodeSummary : list of list
+ The tree node summary statistics, [P(Y=1|T), N(T)], produced by tree_node_summary() method.
+
+ Returns
+ -------
+ d_res : CTS score
+ '''
+ return -max([stat[0] for stat in nodeSummary])
+
+ @staticmethod
+ def arr_evaluate_CTS(np.ndarray[P_TYPE_t, ndim=1] node_summary_p,
+ np.ndarray[N_TYPE_t, ndim=1] node_summary_n):
+ '''
+ Calculate CTS (conditional treatment selection) as split evaluation criterion for a given node.
+
+ Args
+ ----
+ node_summary_p : array of shape [n_class]
+ Has type numpy.double.
+ The positive probabilities of each of the control
+ and treament groups of the current node, i.e. [P(Y=1|T=i)...]
+ node_summary_n : array of shape [n_class]
+ Has type numpy.int32.
+ The counts of each of the control
+ and treament groups of the current node, i.e. [N(T=i)...]
+
+ Returns
+ -------
+ d_res : CTS score
+ '''
+ # not sure why use negative for CTS, but in calculating the
+ # gain, it is adjusted back so as to maximize the gain.
+ cdef int n_class = node_summary_p.shape[0]
+ cdef P_TYPE_t d_res = node_summary_p[0]
+ cdef int i = 0
+ for i in range(1, n_class):
+ if node_summary_p[i] > d_res:
+ d_res = node_summary_p[i]
+ return -d_res
+
+ def normI(self, n_c: cython.int, n_c_left: cython.int, n_t: list, n_t_left: list, alpha: cython.float = 0.9, currentDivergence: cython.float = 0.0) -> cython.float:
+ '''
+ Normalization factor.
+
+ Args
+ ----
+ currentNodeSummary : list of list
+ The summary statistics of the current tree node, [P(Y=1|T), N(T)].
+
+ leftNodeSummary : list of list
+ The summary statistics of the left tree node, [P(Y=1|T), N(T)].
+
+ alpha : float
+ The weight used to balance different normalization parts.
+
+ Returns
+ -------
+ norm_res : float
+ Normalization factor.
+ '''
+
+ norm_res: cython.float = 0.
+ pt_a: cython.float
+ pc_a: cython.float
+
+ pt_a = 1. * np.sum(n_t_left) / (np.sum(n_t) + 0.1)
+ pc_a = 1. * n_c_left / (n_c + 0.1)
+
+ if self.evaluationFunction == self.evaluate_IDDP:
+ # Normalization Part 1
+ norm_res += (entropyH(1. * np.sum(n_t) / (np.sum(n_t) + n_c), 1. * n_c / (np.sum(n_t) + n_c)) * currentDivergence)
+ norm_res += (1. * np.sum(n_t) / (np.sum(n_t) + n_c) * entropyH(pt_a))
+
+ else:
+ # Normalization Part 1
+ norm_res += (alpha * entropyH(1. * np.sum(n_t) / (np.sum(n_t) + n_c), 1. * n_c / (np.sum(n_t) + n_c)) * kl_divergence(pt_a, pc_a))
+ # Normalization Part 2 & 3
+ for i in range(len(n_t)):
+ pt_a_i = 1. * n_t_left[i] / (n_t[i] + 0.1)
+ norm_res += ((1 - alpha) * entropyH(1. * n_t[i] / (n_t[i] + n_c), 1. * n_c / (n_t[i] + n_c)) * kl_divergence(1. * pt_a_i, pc_a))
+ norm_res += (1. * n_t[i] / (np.sum(n_t) + n_c) * entropyH(pt_a_i))
+ # Normalization Part 4
+ norm_res += 1. * n_c / (np.sum(n_t) + n_c) * entropyH(pc_a)
+
+ # Normalization Part 5
+ norm_res += 0.5
+ return norm_res
+
+ def arr_normI(self, cur_node_summary_n, left_node_summary_n,
+ alpha: cython.float = 0.9, currentDivergence: cython.float = 0.0) -> cython.float:
+ '''
+ Normalization factor.
+
+ Args
+ ----
+ cur_node_summary_n : array of shape [n_class]
+ Has type numpy.int32.
+ The counts of each of the control
+ and treament groups of the current node, i.e. [N(T=i)...]
+
+ left_node_summary_n : array of shape [n_class]
+ Has type numpy.int32.
+ The counts of each of the control
+ and treament groups of the left node, i.e. [N(T=i)...]
+
+ alpha : float
+ The weight used to balance different normalization parts.
+
+ Returns
+ -------
+ norm_res : float
+ Normalization factor.
+ '''
+ cdef N_TYPE_t[::1] cur_summary_n = cur_node_summary_n
+ cdef N_TYPE_t[::1] left_summary_n = left_node_summary_n
+ cdef int n_class = cur_summary_n.shape[0]
+ cdef int i = 0
+
+ cdef P_TYPE_t norm_res = 0.0
+ cdef P_TYPE_t n_c = cur_summary_n[0]
+ cdef P_TYPE_t n_c_left = left_summary_n[0]
+ cdef P_TYPE_t pt_a = 0.0, pt_a_i = 0.0, pc_a = 0.0, sum_n_t_left = 0.0, sum_n_t = 0.0
+
+ for i in range(1, n_class):
+ sum_n_t_left += left_summary_n[i]
+ sum_n_t += cur_summary_n[i]
+
+ pt_a = 1. * sum_n_t_left / (sum_n_t + 0.1)
+ pc_a = 1. * n_c_left / (n_c + 0.1)
+
+ if self.evaluationFunction == self.evaluate_IDDP:
+ # Normalization Part 1
+ norm_res += (entropyH(1. * sum_n_t / (sum_n_t + n_c), 1. * n_c / (sum_n_t + n_c)) * currentDivergence)
+ norm_res += (1. * sum_n_t / (sum_n_t + n_c) * entropyH(pt_a))
+
+ else:
+ # Normalization Part 1
+ norm_res += (alpha * entropyH(1. * sum_n_t / (sum_n_t + n_c), 1. * n_c / (sum_n_t + n_c)) * kl_divergence(pt_a, pc_a))
+ # Normalization Part 2 & 3
+ for i in range(1, n_class):
+ pt_a_i = 1. * left_summary_n[i] / (cur_summary_n[i] + 0.1)
+ norm_res += ((1 - alpha) * entropyH(1. * cur_summary_n[i] / (cur_summary_n[i] + n_c), 1. * n_c / (cur_summary_n[i] + n_c)) * kl_divergence(1. * pt_a_i, pc_a))
+ norm_res += (1. * cur_summary_n[i] / (sum_n_t + n_c) * entropyH(pt_a_i))
+ # Normalization Part 4
+ norm_res += 1. * n_c / (sum_n_t + n_c) * entropyH(pc_a)
+
+ # Normalization Part 5
+ norm_res += 0.5
+ return norm_res
+
+ def tree_node_summary(self, treatment_idx, y, min_samples_treatment=10, n_reg=100, parentNodeSummary=None):
+ '''
+ Tree node summary statistics.
+
+ Args
+ ----
+ treatment_idx : array-like, shape = [num_samples]
+ An array containing the treatment group index for each unit.
+ y : array-like, shape = [num_samples]
+ An array containing the outcome of interest for each unit.
+ min_samples_treatment: int, optional (default=10)
+ The minimum number of samples required of the experiment group t be split at a leaf node.
+ n_reg : int, optional (default=10)
+ The regularization parameter defined in Rzepakowski et al. 2012,
+ the weight (in terms of sample size) of the parent node influence
+ on the child node, only effective for 'KL', 'ED', 'Chi', 'CTS' methods.
+ parentNodeSummary : list of list
+ The positive probabilities and sample sizes of each of the control and treatment groups
+ in the parent node.
+
+ Returns
+ -------
+ nodeSummary : list of list
+ The positive probabilities and sample sizes of each of the control and treatment groups
+ in the current node.
+ '''
+ # counts: [[N(Y=0, T=0), N(Y=1, T=0)], [N(Y=0, T=1), N(Y=1, T=1)], ...]
+ counts = self.group_uniqueCounts(treatment_idx, y)
+
+ # nodeSummary: [[P(Y=1|T=0), N(T=0)], [P(Y=1|T=1), N(T=1)], ...]
+ nodeSummary = []
+ # Iterate the control and treatment groups
+ for i, count in enumerate(counts):
+ n_pos = count[1]
+ n = count[0] + n_pos
+ if parentNodeSummary is None:
+ p = n_pos / n if n > 0 else 0.
+ elif n > min_samples_treatment:
+ p = (n_pos + parentNodeSummary[i][0] * n_reg) / (n + n_reg)
+ else:
+ p = parentNodeSummary[i][0]
+
+ nodeSummary.append([p, n])
+
+ return nodeSummary
+
+ @staticmethod
+ def tree_node_summary_to_arr(np.ndarray[TR_TYPE_t, ndim=1] treatment_idx,
+ np.ndarray[Y_TYPE_t, ndim=1] y,
+ np.ndarray[P_TYPE_t, ndim=1] out_summary_p,
+ np.ndarray[N_TYPE_t, ndim=1] out_summary_n,
+ np.ndarray[N_TYPE_t, ndim=1] buf_count_arr,
+ np.ndarray[P_TYPE_t, ndim=1] parentNodeSummary_p,
+ int has_parent_summary,
+ min_samples_treatment=10, n_reg=100
+ ):
+ '''
+ Tree node summary statistics.
+ Modified from tree_node_summary, to use different format for the summary.
+ Instead of [[P(Y=1|T=0), N(T=0)], [P(Y=1|T=1), N(T=1)], ...],
+ use two arrays [N(T=i)...] and [P(Y=1|T=i)...].
+
+ Args
+ ----
+ treatment_idx : array-like, shape = [num_samples]
+ An array containing the treatment group index for each unit.
+ Has type numpy.int8.
+ y : array-like, shape = [num_samples]
+ An array containing the outcome of interest for each unit.
+ Has type numpy.int8.
+ out_summary_p : array of shape [n_class]
+ Has type numpy.double.
+ To be filled with the positive probabilities of each of the control
+ and treament groups of the current node.
+ out_summary_n : array of shape [n_class]
+ Has type numpy.int32.
+ To be filled with the counts of each of the control
+ and treament groups of the current node.
+ buf_count_arr : array of shape [2*n_class]
+ Has type numpy.int32.
+ To be use as temporary buffer for group_uniqueCounts_to_arr.
+ parentNodeSummary_p : array of shape [n_class]
+ The positive probabilities of each of the control and treatment groups
+ in the parent node.
+ has_parent_summary : bool as int
+ If True (non-zero), then parentNodeSummary_p is a valid parent node summary probabilities.
+ If False (0), assume no parent node summary and parentNodeSummary_p is not touched.
+ min_samples_treatment: int, optional (default=10)
+ The minimum number of samples required of the experiment group t be split at a leaf node.
+ n_reg : int, optional (default=10)
+ The regularization parameter defined in Rzepakowski et al. 2012,
+ the weight (in terms of sample size) of the parent node influence
+ on the child node, only effective for 'KL', 'ED', 'Chi', 'CTS' methods.
+
+ Returns
+ -------
+ No return values, but will modify out_summary_p and out_summary_n.
+ '''
+ # buf_count_arr: [N(Y=0, T=0), N(Y=1, T=0), N(Y=0, T=1), N(Y=1, T=1), ...]
+ group_uniqueCounts_to_arr(treatment_idx, y, buf_count_arr)
+
+ cdef int i = 0
+ cdef int n_class = buf_count_arr.shape[0] / 2
+ cdef int n = 0
+ cdef int n_pos = 0
+ cdef P_TYPE_t p = 0.0
+ cdef int n_min_sams = min_samples_treatment
+ cdef P_TYPE_t n_reg_p = n_reg
+
+ # out_summary_p: [P(Y=1|T=i)...]
+ # out_summary_n: [N(T=i) ... ]
+ if has_parent_summary == 0:
+ for i in range(n_class):
+ n_pos = buf_count_arr[2*i + 1] # N(Y=1|T=i)
+ n = buf_count_arr[2*i] + n_pos # N(Y=0|T=i) + N(Y=1|T=i) == N(T=i)
+ p = (n_pos / n) if n > 0 else 0.
+ out_summary_n[i] = n
+ out_summary_p[i] = p
+ else:
+ for i in range(n_class):
+ n_pos = buf_count_arr[2*i + 1]
+ n = buf_count_arr[2*i] + n_pos
+ if n > n_min_sams:
+ p = (n_pos + parentNodeSummary_p[i] * n_reg_p) / ( n + n_reg_p)
+ else:
+ p = parentNodeSummary_p[i]
+ out_summary_n[i] = n
+ out_summary_p[i] = p
+
+ @staticmethod
+ def tree_node_summary_from_counts(
+ np.ndarray[N_TYPE_t, ndim=1] group_count_arr,
+ np.ndarray[P_TYPE_t, ndim=1] out_summary_p,
+ np.ndarray[N_TYPE_t, ndim=1] out_summary_n,
+ np.ndarray[P_TYPE_t, ndim=1] parentNodeSummary_p,
+ int has_parent_summary,
+ min_samples_treatment=10, n_reg=100
+ ):
+ '''Tree node summary statistics.
+
+ Modified from tree_node_summary_to_arr, to use different
+ format for the summary and to calculate based on already
+ calculated group counts. Instead of [[P(Y=1|T=0), N(T=0)],
+ [P(Y=1|T=1), N(T=1)], ...], use two arrays [N(T=i)...] and
+ [P(Y=1|T=i)...].
+
+ Args
+ ----
+ group_count_arr : array of shape [2*n_class]
+ Has type numpy.int32.
+ The grounp counts, where entry 2*i is N(Y=0, T=i),
+ and entry 2*i+1 is N(Y=1, T=i).
+ out_summary_p : array of shape [n_class]
+ Has type numpy.double.
+ To be filled with the positive probabilities of each of the control
+ and treament groups of the current node.
+ out_summary_n : array of shape [n_class]
+ Has type numpy.int32.
+ To be filled with the counts of each of the control
+ and treament groups of the current node.
+ parentNodeSummary_p : array of shape [n_class]
+ The positive probabilities of each of the control and treatment groups
+ in the parent node.
+ has_parent_summary : bool as int
+ If True (non-zero), then parentNodeSummary_p is a valid parent node summary probabilities.
+ If False (0), assume no parent node summary and parentNodeSummary_p is not touched.
+ min_samples_treatment: int, optional (default=10)
+ The minimum number of samples required of the experiment group t be split at a leaf node.
+ n_reg : int, optional (default=10)
+ The regularization parameter defined in Rzepakowski et al. 2012,
+ the weight (in terms of sample size) of the parent node influence
+ on the child node, only effective for 'KL', 'ED', 'Chi', 'CTS' methods.
+
+ Returns
+ -------
+ No return values, but will modify out_summary_p and out_summary_n.
+
+ '''
+ # group_count_arr: [N(Y=0, T=0), N(Y=1, T=0), N(Y=0, T=1), N(Y=1, T=1), ...]
+ cdef int i = 0
+ cdef int n_class = group_count_arr.shape[0] / 2
+ cdef int n = 0
+ cdef int n_pos = 0
+ cdef P_TYPE_t p = 0.0
+ cdef int n_min_sams = min_samples_treatment
+ cdef P_TYPE_t n_reg_p = n_reg
+
+ # out_summary_p: [P(Y=1|T=i)...]
+ # out_summary_n: [N(T=i) ... ]
+ if has_parent_summary == 0:
+ for i in range(n_class):
+ n_pos = group_count_arr[2*i + 1] # N(Y=1|T=i)
+ n = group_count_arr[2*i] + n_pos # N(Y=0|T=i) + N(Y=1|T=i) == N(T=i)
+ p = (n_pos / n) if n > 0 else 0.
+ out_summary_n[i] = n
+ out_summary_p[i] = p
+ else:
+ for i in range(n_class):
+ n_pos = group_count_arr[2*i + 1]
+ n = group_count_arr[2*i] + n_pos
+ if n > n_min_sams:
+ p = (n_pos + parentNodeSummary_p[i] * n_reg_p) / ( n + n_reg_p)
+ else:
+ p = parentNodeSummary_p[i]
+ out_summary_n[i] = n
+ out_summary_p[i] = p
+
+ def uplift_classification_results(self, treatment_idx, y):
+ '''
+ Classification probability for each treatment in the tree node.
+
+ Args
+ ----
+ treatment_idx : array-like, shape = [num_samples]
+ An array containing the treatment group index for each unit.
+ y : array-like, shape = [num_samples]
+ An array containing the outcome of interest for each unit.
+
+ Returns
+ -------
+ res : list of list
+ The positive probabilities P(Y = 1) of each of the control and treatment groups
+ '''
+ # counts: [[N(Y=0, T=0), N(Y=1, T=0)], [N(Y=0, T=1), N(Y=1, T=1)], ...]
+ counts = self.group_uniqueCounts(treatment_idx, y)
+ res = []
+ for count in counts:
+ n_pos = count[1]
+ n = count[0] + n_pos
+ p = n_pos / n if n > 0 else 0.
+ res.append(p)
+ return res
+
+ def growDecisionTreeFrom(self, X, treatment_idx, y, X_val, treatment_val_idx, y_val,
+ early_stopping_eval_diff_scale=1, max_depth=10,
+ min_samples_leaf=100, depth=1,
+ min_samples_treatment=10, n_reg=100,
+ parentNodeSummary_p=None):
+ '''
+ Train the uplift decision tree.
+
+ Args
+ ----
+ X : ndarray, shape = [num_samples, num_features]
+ An ndarray of the covariates used to train the uplift model.
+ treatment_idx : array-like, shape = [num_samples]
+ An array containing the treatment group idx for each unit.
+ The dtype should be numpy.int8.
+ y : array-like, shape = [num_samples]
+ An array containing the outcome of interest for each unit.
+ X_val : ndarray, shape = [num_samples, num_features]
+ An ndarray of the covariates used to valid the uplift model.
+ treatment_val_idx : array-like, shape = [num_samples]
+ An array containing the validation treatment group idx for each unit.
+ y_val : array-like, shape = [num_samples]
+ An array containing the validation outcome of interest for each unit.
+ max_depth: int, optional (default=10)
+ The maximum depth of the tree.
+ min_samples_leaf: int, optional (default=100)
+ The minimum number of samples required to be split at a leaf node.
+ depth : int, optional (default = 1)
+ The current depth.
+ min_samples_treatment: int, optional (default=10)
+ The minimum number of samples required of the experiment group to be split at a leaf node.
+ n_reg: int, optional (default=10)
+ The regularization parameter defined in Rzepakowski et al. 2012,
+ the weight (in terms of sample size) of the parent node influence
+ on the child node, only effective for 'KL', 'ED', 'Chi', 'CTS' methods.
+ parentNodeSummary_p : array-like, shape [n_class]
+ Node summary probability statistics of the parent tree node.
+
+ Returns
+ -------
+ object of DecisionTree class
+ '''
+
+ if len(X) == 0:
+ return DecisionTree(classes_=self.classes_)
+
+ assert treatment_idx.dtype == TR_TYPE
+ assert y.dtype == Y_TYPE
+
+ # some temporary buffers for node summaries
+ cdef int n_class = self.n_class
+ # buffers for group counts, right can be derived from total and left
+ cdef np.ndarray[N_TYPE_t, ndim=1] left_count_arr = np.zeros(2 * self.n_class, dtype = N_TYPE)
+ cdef np.ndarray[N_TYPE_t, ndim=1] right_count_arr = np.zeros(2 * self.n_class, dtype = N_TYPE)
+ cdef np.ndarray[N_TYPE_t, ndim=1] total_count_arr = np.zeros(2 * self.n_class, dtype = N_TYPE)
+ # for X_val if any, allocate if needed below
+ cdef np.ndarray[N_TYPE_t, ndim=1] val_left_count_arr
+ cdef np.ndarray[N_TYPE_t, ndim=1] val_right_count_arr
+ cdef np.ndarray[N_TYPE_t, ndim=1] val_total_count_arr
+ # buffers for node summary
+ cdef np.ndarray[P_TYPE_t, ndim=1] cur_summary_p = np.zeros(self.n_class, dtype = P_TYPE)
+ cdef np.ndarray[N_TYPE_t, ndim=1] cur_summary_n = np.zeros(self.n_class, dtype = N_TYPE)
+ cdef np.ndarray[P_TYPE_t, ndim=1] left_summary_p = np.zeros(self.n_class, dtype = P_TYPE)
+ cdef np.ndarray[N_TYPE_t, ndim=1] left_summary_n = np.zeros(self.n_class, dtype = N_TYPE)
+ cdef np.ndarray[P_TYPE_t, ndim=1] right_summary_p = np.zeros(self.n_class, dtype = P_TYPE)
+ cdef np.ndarray[N_TYPE_t, ndim=1] right_summary_n = np.zeros(self.n_class, dtype = N_TYPE)
+ # for val left and right summary
+ cdef np.ndarray[P_TYPE_t, ndim=1] val_left_summary_p = np.zeros(self.n_class, dtype = P_TYPE)
+ cdef np.ndarray[N_TYPE_t, ndim=1] val_left_summary_n = np.zeros(self.n_class, dtype = N_TYPE)
+ cdef np.ndarray[P_TYPE_t, ndim=1] val_right_summary_p = np.zeros(self.n_class, dtype = P_TYPE)
+ cdef np.ndarray[N_TYPE_t, ndim=1] val_right_summary_n = np.zeros(self.n_class, dtype = N_TYPE)
+
+ # dummy
+ cdef int has_parent_summary = 0
+ if parentNodeSummary_p is None:
+ parent_summary_p = np.zeros(self.n_class, dtype = P_TYPE) # dummy for calling tree_node_summary_to_arr
+ has_parent_summary = 0
+ else:
+ parent_summary_p = parentNodeSummary_p
+ has_parent_summary = 1
+
+ cdef int i = 0
+
+ # preparation: fill in the total count, then for each
+ # candidate split, we calculate the count for left branch, and
+ # can derive count for right branch using the total count.
+
+ # group_count_arr: [N(Y=0, T=0), N(Y=1, T=0), N(Y=0, T=1), N(Y=1, T=1), ...]
+ group_uniqueCounts_to_arr(treatment_idx, y, total_count_arr)
+ if X_val is not None:
+ val_left_count_arr = np.zeros(2 * self.n_class, dtype = N_TYPE)
+ val_right_count_arr = np.zeros(2 * self.n_class, dtype = N_TYPE)
+ val_total_count_arr = np.zeros(2 * self.n_class, dtype = N_TYPE)
+ group_uniqueCounts_to_arr(treatment_val_idx, y_val, val_total_count_arr)
+
+ # Current node summary: [P(Y=1|T=i)...] and [N(T=i)...]
+ self.tree_node_summary_from_counts(
+ total_count_arr,
+ cur_summary_p, cur_summary_n,
+ parent_summary_p,
+ has_parent_summary,
+ min_samples_treatment=min_samples_treatment,
+ n_reg=n_reg
+ )
+
+ # to reconstruct current node summary in list of list form, so
+ # that the constructed tree follows previous format.
+
+ # Current node summary: [[P(Y=1|T=i), N(T=i)]...]
+ currentNodeSummary = []
+ for i in range(n_class):
+ currentNodeSummary.append([cur_summary_p[i], cur_summary_n[i]])
+ #
+
+ if self.evaluationFunction == self.evaluate_IT or self.evaluationFunction == self.evaluate_CIT:
+ currentScore = 0
+ else:
+ currentScore = self.arr_eval_func(cur_summary_p, cur_summary_n)
+
+ # Prune Stats:
+ cdef P_TYPE_t maxAbsDiff = 0.0
+ cdef P_TYPE_t maxDiff = -1.
+ cdef int bestTreatment = 0 # treatment index for the control group, also used in returning the tree for this node
+ cdef int suboptTreatment = 0 # treatment index for the control group
+ cdef int maxDiffTreatment = 0 # treatment index for the control group, also used in returning the tree for this node
+ maxDiffSign = 0 # also used in returning the tree for this node
+ # adapted to new current node summary format
+ cdef P_TYPE_t p_c = cur_summary_p[0]
+ cdef N_TYPE_t n_c = cur_summary_n[0]
+ cdef N_TYPE_t n_t = 0
+ cdef int i_tr = 0
+ cdef P_TYPE_t p_t = 0.0, diff = 0.0
+
+ for i_tr in range(1, n_class):
+ p_t = cur_summary_p[i_tr]
+ # P(Y=1|T=t) - P(Y=1|T=0)
+ diff = p_t - p_c
+ if fabs(diff) >= maxAbsDiff:
+ maxDiffTreatment = i_tr
+ maxDiffSign = np.sign(diff)
+ maxAbsDiff = fabs(diff)
+ if diff >= maxDiff:
+ maxDiff = diff
+ suboptTreatment = i_tr
+ if diff > 0:
+ bestTreatment = i_tr
+ if maxDiff > 0:
+ p_t = cur_summary_p[bestTreatment]
+ n_t = cur_summary_n[bestTreatment]
+ else:
+ p_t = cur_summary_p[suboptTreatment]
+ n_t = cur_summary_n[suboptTreatment]
+ p_value = (1. - stats.norm.cdf(fabs(p_c - p_t) / sqrt(p_t * (1 - p_t) / n_t + p_c * (1 - p_c) / n_c))) * 2
+ upliftScore = [maxDiff, p_value]
+
+ bestGain = 0.0
+ bestGainImp = 0.0
+ bestAttribute = None
+ # keep mostly scalar when finding best split, then get the structural value after finding the best split
+ best_col = None
+ best_value = None
+ len_X = len(X)
+ len_X_val = len(X_val) if X_val is not None else 0
+
+ c_num_percentiles = [3, 5, 10, 20, 30, 50, 70, 80, 90, 95, 97]
+ c_cat_percentiles = [10, 50, 90]
+
+ # last column is the result/target column, 2nd to the last is the treatment group
+ columnCount = X.shape[1]
+ if (self.max_features and self.max_features > 0 and self.max_features <= columnCount):
+ max_features = self.max_features
+ else:
+ max_features = columnCount
+
+ for col in list(self.random_state_.choice(a=range(columnCount), size=max_features, replace=False)):
+ columnValues = X[:, col]
+ # unique values
+ lsUnique = np.unique(columnValues)
+
+ if np.issubdtype(lsUnique.dtype, np.number):
+ is_split_by_gt = True
+ if len(lsUnique) > 10:
+ lspercentile = np.percentile(columnValues, c_num_percentiles)
+ else:
+ lspercentile = np.percentile(lsUnique, c_cat_percentiles)
+ lsUnique = np.unique(lspercentile)
+ else:
+ # to split by equality check.
+ is_split_by_gt = False
+
+ for value in lsUnique:
+ len_X_l = group_counts_by_divide(columnValues, value, is_split_by_gt, treatment_idx, y, left_count_arr)
+ len_X_r = len_X - len_X_l
+
+ # check the split validity on min_samples_leaf 372
+ if (len_X_l < min_samples_leaf or len_X_r < min_samples_leaf):
+ continue
+ # summarize notes
+ # Gain -- Entropy or Gini
+ p = float(len_X_l) / len_X
+
+ # right branch group counts can be calculated from left branch counts and total counts
+ for i in range(2 * n_class):
+ right_count_arr[i] = total_count_arr[i] - left_count_arr[i]
+
+ # left and right node summary, into the temporary buffers {left,right}_summary_{p,n}
+ self.tree_node_summary_from_counts(
+ left_count_arr,
+ left_summary_p, left_summary_n,
+ cur_summary_p,
+ 1,
+ min_samples_treatment,
+ n_reg
+ )
+
+ self.tree_node_summary_from_counts(
+ right_count_arr,
+ right_summary_p, right_summary_n,
+ cur_summary_p,
+ 1,
+ min_samples_treatment,
+ n_reg
+ )
+
+ if X_val is not None:
+ len_X_val_l = group_counts_by_divide(X_val[:, col], value, is_split_by_gt, treatment_val_idx, y_val, val_left_count_arr)
+
+ # right branch group counts can be calculated from left branch counts and total counts
+ for i in range(2 * n_class):
+ val_right_count_arr[i] = val_total_count_arr[i] - val_left_count_arr[i]
+
+ self.tree_node_summary_from_counts(
+ val_left_count_arr,
+ val_left_summary_p, val_left_summary_n,
+ cur_summary_p, # parentNodeSummary_p
+ 1 # has_parent_summary
+ )
+
+ self.tree_node_summary_from_counts(
+ val_right_count_arr,
+ val_right_summary_p, val_right_summary_n,
+ cur_summary_p, # parentNodeSummary_p
+ 1 # has_parent_summary
+ )
+
+ early_stopping_flag = False
+ for k in range(n_class):
+ if (abs(val_left_summary_p[k] - left_summary_p[k]) >
+ min(val_left_summary_p[k], left_summary_p[k])/early_stopping_eval_diff_scale or
+ abs(val_right_summary_p[k] - right_summary_p[k]) >
+ min(val_right_summary_p[k], right_summary_p[k])/early_stopping_eval_diff_scale):
+ early_stopping_flag = True
+ break
+
+ if early_stopping_flag:
+ continue
+
+ # check the split validity on min_samples_treatment
+ node_mst = min(np.min(left_summary_n), np.min(right_summary_n))
+ if node_mst < min_samples_treatment:
+ continue
+
+ # evaluate the split
+ if self.arr_eval_func == self.arr_evaluate_CTS:
+ leftScore1 = self.arr_eval_func(left_summary_p, left_summary_n)
+ rightScore2 = self.arr_eval_func(right_summary_p, right_summary_n)
+ gain = (currentScore - p * leftScore1 - (1 - p) * rightScore2)
+ gain_for_imp = (len_X * currentScore - len_X_l * leftScore1 - len_X_r * rightScore2)
+ elif self.arr_eval_func == self.arr_evaluate_DDP:
+ leftScore1 = self.arr_eval_func(left_summary_p, left_summary_n)
+ rightScore2 = self.arr_eval_func(right_summary_p, right_summary_n)
+ gain = np.abs(leftScore1 - rightScore2)
+ gain_for_imp = np.abs(len_X_l * leftScore1 - len_X_r * rightScore2)
+ elif self.arr_eval_func == self.arr_evaluate_IT:
+ gain = self.arr_eval_func(left_summary_p, left_summary_n, right_summary_p, right_summary_n)
+ gain_for_imp = gain * len_X
+ elif self.arr_eval_func == self.arr_evaluate_CIT:
+ gain = self.arr_eval_func(cur_summary_p, cur_summary_n,
+ left_summary_p, left_summary_n,
+ right_summary_p, right_summary_n)
+ gain_for_imp = gain * len_X
+ elif self.arr_eval_func == self.arr_evaluate_IDDP:
+ leftScore1 = self.arr_eval_func(left_summary_p, left_summary_n)
+ rightScore2 = self.arr_eval_func(right_summary_p, right_summary_n)
+ gain = np.abs(leftScore1 - rightScore2) - np.abs(currentScore)
+ gain_for_imp = (len_X_l * leftScore1 + len_X_r * rightScore2 - len_X * np.abs(currentScore))
+ if self.normalization:
+ # Normalize used divergence
+ currentDivergence = 2 * (gain + 1) / 3
+ norm_factor = self.arr_normI(cur_summary_n, left_summary_n, alpha=0.9, currentDivergence=currentDivergence)
+ else:
+ norm_factor = 1
+ gain = gain / norm_factor
+ else:
+ leftScore1 = self.arr_eval_func(left_summary_p, left_summary_n)
+ rightScore2 = self.arr_eval_func(right_summary_p, right_summary_n)
+ gain = (p * leftScore1 + (1 - p) * rightScore2 - currentScore)
+ gain_for_imp = (len_X_l * leftScore1 + len_X_r * rightScore2 - len_X * currentScore)
+ if self.normalization:
+ norm_factor = self.arr_normI(cur_summary_n, left_summary_n, alpha=0.9)
+ else:
+ norm_factor = 1
+ gain = gain / norm_factor
+ if (gain > bestGain and len_X_l > min_samples_leaf and len_X_r > min_samples_leaf):
+ bestGain = gain
+ bestGainImp = gain_for_imp
+ best_col = col
+ best_value = value
+
+ # after finding the best split col and value
+ if best_col is not None:
+ bestAttribute = (best_col, best_value)
+ # re-calculate the divideSet
+ X_l, X_r, w_l, w_r, y_l, y_r = self.divideSet(X, treatment_idx, y, best_col, best_value)
+ if X_val is not None:
+ X_val_l, X_val_r, w_val_l, w_val_r, y_val_l, y_val_r = self.divideSet(X_val, treatment_val_idx, y_val, best_col, best_value)
+ best_set_left = [X_l, w_l, y_l, X_val_l, w_val_l, y_val_l]
+ best_set_right = [X_r, w_r, y_r, X_val_r, w_val_r, y_val_r]
+ else:
+ best_set_left = [X_l, w_l, y_l, None, None, None]
+ best_set_right = [X_r, w_r, y_r, None, None, None]
+
+ dcY = {'impurity': '%.3f' % currentScore, 'samples': '%d' % len(X)}
+ # Add treatment size
+ dcY['group_size'] = ''
+ for i, summary in enumerate(currentNodeSummary):
+ dcY['group_size'] += ' ' + self.classes_[i] + ': ' + str(summary[1])
+ dcY['upliftScore'] = [round(upliftScore[0], 4), round(upliftScore[1], 4)]
+ dcY['matchScore'] = round(upliftScore[0], 4)
+
+ if bestGain > 0 and depth < max_depth:
+ self.feature_imp_dict[bestAttribute[0]] += bestGainImp
+ trueBranch = self.growDecisionTreeFrom(
+ *best_set_left, self.early_stopping_eval_diff_scale, max_depth, min_samples_leaf,
+ depth + 1, min_samples_treatment=min_samples_treatment,
+ n_reg=n_reg, parentNodeSummary_p=cur_summary_p
+ )
+ falseBranch = self.growDecisionTreeFrom(
+ *best_set_right, self.early_stopping_eval_diff_scale, max_depth, min_samples_leaf,
+ depth + 1, min_samples_treatment=min_samples_treatment,
+ n_reg=n_reg, parentNodeSummary_p=cur_summary_p
+ )
+
+ return DecisionTree(
+ classes_=self.classes_,
+ col=bestAttribute[0], value=bestAttribute[1],
+ trueBranch=trueBranch, falseBranch=falseBranch, summary=dcY,
+ maxDiffTreatment=maxDiffTreatment, maxDiffSign=maxDiffSign,
+ nodeSummary=currentNodeSummary,
+ backupResults=self.uplift_classification_results(treatment_idx, y),
+ bestTreatment=bestTreatment, upliftScore=upliftScore
+ )
+ else:
+ if self.evaluationFunction == self.evaluate_CTS:
+ return DecisionTree(
+ classes_=self.classes_,
+ results=self.uplift_classification_results(treatment_idx, y),
+ summary=dcY, nodeSummary=currentNodeSummary,
+ bestTreatment=bestTreatment, upliftScore=upliftScore
+ )
+ else:
+ return DecisionTree(
+ classes_=self.classes_,
+ results=self.uplift_classification_results(treatment_idx, y),
+ summary=dcY, maxDiffTreatment=maxDiffTreatment,
+ maxDiffSign=maxDiffSign, nodeSummary=currentNodeSummary,
+ bestTreatment=bestTreatment, upliftScore=upliftScore
+ )
+
+ @staticmethod
+ def classify(observations, tree, dataMissing=False):
+ '''
+ Classifies (prediction) the observations according to the tree.
+
+ Args
+ ----
+ observations : list of list
+ The internal data format for the training data (combining X, Y, treatment).
+
+ dataMissing: boolean, optional (default = False)
+ An indicator for if data are missing or not.
+
+ Returns
+ -------
+ tree.results, tree.upliftScore :
+ The results in the leaf node.
+ '''
+
+ def classifyWithoutMissingData(observations, tree):
+ '''
+ Classifies (prediction) the observations according to the tree, assuming without missing data.
+
+ Args
+ ----
+ observations : list of list
+ The internal data format for the training data (combining X, Y, treatment).
+
+ Returns
+ -------
+ tree.results, tree.upliftScore :
+ The results in the leaf node.
+ '''
+ if tree.results is not None: # leaf
+ return tree.results, tree.upliftScore
+ else:
+ v = observations[tree.col]
+ branch = None
+ if isinstance(v, numbers.Number):
+ if v >= tree.value:
+ branch = tree.trueBranch
+ else:
+ branch = tree.falseBranch
+ else:
+ if v == tree.value:
+ branch = tree.trueBranch
+ else:
+ branch = tree.falseBranch
+ return classifyWithoutMissingData(observations, branch)
+
+ def classifyWithMissingData(observations, tree):
+ '''
+ Classifies (prediction) the observations according to the tree, assuming with missing data.
+
+ Args
+ ----
+ observations : list of list
+ The internal data format for the training data (combining X, Y, treatment).
+
+ Returns
+ -------
+ tree.results, tree.upliftScore :
+ The results in the leaf node.
+ '''
+ if tree.results is not None: # leaf
+ return tree.results
+ else:
+ v = observations[tree.col]
+ if v is None:
+ tr = classifyWithMissingData(observations, tree.trueBranch)
+ fr = classifyWithMissingData(observations, tree.falseBranch)
+ tcount = sum(tr.values())
+ fcount = sum(fr.values())
+ tw = float(tcount) / (tcount + fcount)
+ fw = float(fcount) / (tcount + fcount)
+
+ # Problem description: http://blog.ludovf.net/python-collections-defaultdict/
+ result = defaultdict(int)
+ for k, v in tr.items():
+ result[k] += v * tw
+ for k, v in fr.items():
+ result[k] += v * fw
+ return dict(result)
+ else:
+ branch = None
+ if isinstance(v, numbers.Number):
+ if v >= tree.value:
+ branch = tree.trueBranch
+ else:
+ branch = tree.falseBranch
+ else:
+ if v == tree.value:
+ branch = tree.trueBranch
+ else:
+ branch = tree.falseBranch
+ return classifyWithMissingData(observations, branch)
+
+ # function body
+ if dataMissing:
+ return classifyWithMissingData(observations, tree)
+ else:
+ return classifyWithoutMissingData(observations, tree)
+
+
+# Uplift Random Forests
+class UpliftRandomForestClassifier:
+ """ Uplift Random Forest for Classification Task.
+
+ Parameters
+ ----------
+ n_estimators : integer, optional (default=10)
+ The number of trees in the uplift random forest.
+
+ evaluationFunction : string
+ Choose from one of the models: 'KL', 'ED', 'Chi', 'CTS', 'DDP', 'IT', 'CIT', 'IDDP'.
+
+ max_features: int, optional (default=10)
+ The number of features to consider when looking for the best split.
+
+ random_state: int, RandomState instance or None (default=None)
+ A random seed or `np.random.RandomState` to control randomness in building the trees and forest.
+
+ max_depth: int, optional (default=5)
+ The maximum depth of the tree.
+
+ min_samples_leaf: int, optional (default=100)
+ The minimum number of samples required to be split at a leaf node.
+
+ min_samples_treatment: int, optional (default=10)
+ The minimum number of samples required of the experiment group to be split at a leaf node.
+
+ n_reg: int, optional (default=10)
+ The regularization parameter defined in Rzepakowski et al. 2012, the
+ weight (in terms of sample size) of the parent node influence on the
+ child node, only effective for 'KL', 'ED', 'Chi', 'CTS' methods.
+
+ early_stopping_eval_diff_scale: float, optional (default=1)
+ If train and valid uplift score diff bigger than
+ min(train_uplift_score,valid_uplift_score)/early_stopping_eval_diff_scale, stop.
+
+ control_name: string
+ The name of the control group (other experiment groups will be regarded as treatment groups)
+
+ normalization: boolean, optional (default=True)
+ The normalization factor defined in Rzepakowski et al. 2012,
+ correcting for tests with large number of splits and imbalanced
+ treatment and control splits
+
+ honesty: bool (default=False)
+ True if the honest approach based on "Athey, S., & Imbens, G. (2016). Recursive partitioning for
+ heterogeneous causal effects." shall be used.
+
+ estimation_sample_size: float (default=0.5)
+ Sample size for estimating the CATE score in the leaves if honesty == True.
+
+ n_jobs: int, optional (default=-1)
+ The parallelization parameter to define how many parallel jobs need to be created.
+ This is passed on to joblib library for parallelizing uplift-tree creation and prediction.
+
+ joblib_prefer: str, optional (default="threads")
+ The preferred backend for joblib (passed as `prefer` to joblib.Parallel). See the joblib
+ documentation for valid values.
+
+ Outputs
+ ----------
+ df_res: pandas dataframe
+ A user-level results dataframe containing the estimated individual treatment effect.
+ """
+ def __init__(self,
+ control_name,
+ n_estimators=10,
+ max_features=10,
+ random_state=None,
+ max_depth=5,
+ min_samples_leaf=100,
+ min_samples_treatment=10,
+ n_reg=10,
+ early_stopping_eval_diff_scale=1,
+ evaluationFunction='KL',
+ normalization=True,
+ honesty=False,
+ estimation_sample_size=0.5,
+ n_jobs=-1,
+ joblib_prefer: str = "threads"):
+
+ """
+ Initialize the UpliftRandomForestClassifier class.
+ """
+ self.n_estimators = n_estimators
+ self.max_features = max_features
+ self.random_state = random_state
+ self.max_depth = max_depth
+ self.min_samples_leaf = min_samples_leaf
+ self.min_samples_treatment = min_samples_treatment
+ self.n_reg = n_reg
+ self.early_stopping_eval_diff_scale = early_stopping_eval_diff_scale
+ self.evaluationFunction = evaluationFunction
+ self.control_name = control_name
+ self.normalization = normalization
+ self.honesty = honesty
+ self.estimation_sample_size = estimation_sample_size
+ self.n_jobs = n_jobs
+ self.joblib_prefer = joblib_prefer
+
+ assert control_name is not None and isinstance(control_name, str), \
+ f"control_group should be string but {control_name} is passed"
+ self.control_name = control_name
+ self.classes_ = [control_name]
+ self.n_class = 1
+
+ if self.n_jobs == -1:
+ self.n_jobs = mp.cpu_count()
+
+ def fit(self, X, treatment, y, X_val=None, treatment_val=None, y_val=None):
+ """
+ Fit the UpliftRandomForestClassifier.
+
+ Args
+ ----
+ X : ndarray, shape = [num_samples, num_features]
+ An ndarray of the covariates used to train the uplift model.
+
+ treatment : array-like, shape = [num_samples]
+ An array containing the treatment group for each unit.
+
+ y : array-like, shape = [num_samples]
+ An array containing the outcome of interest for each unit.
+
+ X_val : ndarray, shape = [num_samples, num_features]
+ An ndarray of the covariates used to valid the uplift model.
+
+ treatment_val : array-like, shape = [num_samples]
+ An array containing the validation treatment group for each unit.
+
+ y_val : array-like, shape = [num_samples]
+ An array containing the validation outcome of interest for each unit.
+ """
+ random_state = check_random_state(self.random_state)
+
+ # Create forest
+ self.uplift_forest = [
+ UpliftTreeClassifier(
+ max_features=self.max_features, max_depth=self.max_depth,
+ min_samples_leaf=self.min_samples_leaf,
+ min_samples_treatment=self.min_samples_treatment,
+ n_reg=self.n_reg,
+ early_stopping_eval_diff_scale=self.early_stopping_eval_diff_scale,
+ evaluationFunction=self.evaluationFunction,
+ control_name=self.control_name,
+ normalization=self.normalization,
+ honesty=self.honesty,
+ estimation_sample_size=self.estimation_sample_size,
+ random_state=random_state.randint(MAX_INT))
+ for _ in range(self.n_estimators)
+ ]
+
+ # Get treatment group keys. self.classes_[0] is reserved for the control group.
+ treatment_groups = sorted([x for x in list(set(treatment)) if x != self.control_name])
+ self.classes_ = [self.control_name]
+ for tr in treatment_groups:
+ self.classes_.append(tr)
+ self.n_class = len(self.classes_)
+
+ self.uplift_forest = (
+ Parallel(n_jobs=self.n_jobs, prefer=self.joblib_prefer)
+ (delayed(self.bootstrap)(X, treatment, y, X_val, treatment_val, y_val, tree) for tree in self.uplift_forest)
+ )
+
+ all_importances = [tree.feature_importances_ for tree in self.uplift_forest]
+ self.feature_importances_ = np.mean(all_importances, axis=0)
+ self.feature_importances_ /= self.feature_importances_.sum() # normalize to add to 1
+
+ @staticmethod
+ def bootstrap(X, treatment, y, X_val, treatment_val, y_val, tree):
+ random_state = check_random_state(tree.random_state)
+ bt_index = random_state.choice(len(X), len(X))
+ x_train_bt = X[bt_index]
+ y_train_bt = y[bt_index]
+ treatment_train_bt = treatment[bt_index]
+
+ if X_val is None:
+ tree.fit(X=x_train_bt, treatment=treatment_train_bt, y=y_train_bt)
+ else:
+ bt_val_index = random_state.choice(len(X_val), len(X_val))
+ x_val_bt = X_val[bt_val_index]
+ y_val_bt = y_val[bt_val_index]
+ treatment_val_bt = treatment_val[bt_val_index]
+
+ tree.fit(X=x_train_bt, treatment=treatment_train_bt, y=y_train_bt, X_val=x_val_bt, treatment_val=treatment_val_bt, y_val=y_val_bt)
+ return tree
+
+ @ignore_warnings(category=FutureWarning)
+ def predict(self, X, full_output=False):
+ '''
+ Returns the recommended treatment group and predicted optimal
+ probability conditional on using the recommended treatment group.
+
+ Args
+ ----
+ X : ndarray, shape = [num_samples, num_features]
+ An ndarray of the covariates used to train the uplift model.
+
+ full_output : bool, optional (default=False)
+ Whether the UpliftTree algorithm returns upliftScores, pred_nodes
+ alongside the recommended treatment group and p_hat in the treatment group.
+
+ Returns
+ -------
+ y_pred_list : ndarray, shape = (num_samples, num_treatments])
+ An ndarray containing the predicted treatment effect of each treatment group for each sample
+
+ df_res : DataFrame, shape = [num_samples, (num_treatments * 2 + 3)]
+ If `full_output` is `True`, a DataFrame containing the predicted outcome of each treatment and
+ control group, the treatment effect of each treatment group, the treatment group with the
+ highest treatment effect, and the maximum treatment effect for each sample.
+
+ '''
+ # Make predictions with all trees and take the average
+
+ if self.n_jobs != 1:
+ y_pred_ensemble = sum(
+ Parallel(n_jobs=self.n_jobs, prefer=self.joblib_prefer)
+ (delayed(tree.predict)(X=X) for tree in self.uplift_forest)
+ ) / len(self.uplift_forest)
+ else:
+ y_pred_ensemble = sum([tree.predict(X=X) for tree in self.uplift_forest]) / len(self.uplift_forest)
+
+ # Summarize results into dataframe
+ df_res = pd.DataFrame(y_pred_ensemble, columns=self.classes_)
+ df_res['recommended_treatment'] = df_res.apply(np.argmax, axis=1)
+
+ # Calculate delta
+ delta_cols = [f'delta_{treatment_group}' for treatment_group in self.classes_[1:]]
+ for i_tr in range(1, self.n_class):
+ treatment_group = self.classes_[i_tr]
+ df_res[f'delta_{treatment_group}'] = df_res[treatment_group] - df_res[self.control_name]
+
+ df_res['max_delta'] = df_res[delta_cols].max(axis=1)
+
+ if full_output:
+ return df_res
+ else:
+ return df_res[delta_cols].values
diff --git a/causalml/source/causalml/inference/tree/utils.py b/causalml/source/causalml/inference/tree/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..fbd22efef08623bd03d8c26623b4633d49a7ea16
--- /dev/null
+++ b/causalml/source/causalml/inference/tree/utils.py
@@ -0,0 +1,359 @@
+"""
+Utility functions for uplift trees.
+"""
+
+import time
+from typing import Callable
+
+import numpy as np
+import pandas as pd
+
+
+def cat_group(dfx, kpix, n_group=10):
+ """
+ Category Reduction for Categorical Variables
+
+ Args
+ ----
+
+ dfx : dataframe
+ The inputs data dataframe.
+
+ kpix : string
+ The column of the feature.
+
+ n_group : int, optional (default = 10)
+ The number of top category values to be remained, other category values will be put into "Other".
+
+ Returns
+ -------
+ The transformed categorical feature value list.
+ """
+ if dfx[kpix].nunique() > n_group:
+ # get the top categories
+ top = dfx[kpix].isin(dfx[kpix].value_counts().index[:n_group])
+ dfx.loc[~top, kpix] = "Other"
+ return dfx[kpix].values
+ else:
+ return dfx[kpix].values
+
+
+def cat_transform(dfx, kpix, kpi1):
+ """
+ Encoding string features.
+
+ Args
+ ----
+
+ dfx : dataframe
+ The inputs data dataframe.
+
+ kpix : string
+ The column of the feature.
+
+ kpi1 : list
+ The list of feature names.
+
+ Returns
+ -------
+ dfx : DataFrame
+ The updated dataframe containing the encoded data.
+
+ kpi1 : list
+ The updated feature names containing the new dummy feature names.
+ """
+ df_dummy = pd.get_dummies(dfx[kpix].values)
+ new_col_names = ["%s_%s" % (kpix, x) for x in df_dummy.columns]
+ df_dummy.columns = new_col_names
+ dfx = pd.concat([dfx, df_dummy], axis=1)
+ for new_col in new_col_names:
+ if new_col not in kpi1:
+ kpi1.append(new_col)
+ if kpix in kpi1:
+ kpi1.remove(kpix)
+ return dfx, kpi1
+
+
+def cv_fold_index(n, i, k, random_seed=2018):
+ """
+ Encoding string features.
+
+ Args
+ ----
+
+ dfx : dataframe
+ The inputs data dataframe.
+
+ kpix : string
+ The column of the feature.
+
+ kpi1 : list
+ The list of feature names.
+
+ Returns
+ -------
+ dfx : DataFrame
+ The updated dataframe containing the encoded data.
+
+ kpi1 : list
+ The updated feature names containing the new dummy feature names.
+ """
+ np.random.seed(random_seed)
+ rlist = np.random.choice(a=range(k), size=n, replace=True)
+ fold_i_index = np.where(rlist == i)[0]
+ return fold_i_index
+
+
+# Categorize continuous variable
+def cat_continuous(x, granularity="Medium"):
+ """
+ Categorize (bin) continuous variable based on percentile.
+
+ Args
+ ----
+
+ x : list
+ Feature values.
+
+ granularity : string, optional, (default = 'Medium')
+ Control the granularity of the bins, optional values are: 'High', 'Medium', 'Low'.
+
+ Returns
+ -------
+ res : list
+ List of percentile bins for the feature value.
+ """
+ if granularity == "High":
+ lspercentile = [
+ np.percentile(x, 5),
+ np.percentile(x, 10),
+ np.percentile(x, 15),
+ np.percentile(x, 20),
+ np.percentile(x, 25),
+ np.percentile(x, 30),
+ np.percentile(x, 35),
+ np.percentile(x, 40),
+ np.percentile(x, 45),
+ np.percentile(x, 50),
+ np.percentile(x, 55),
+ np.percentile(x, 60),
+ np.percentile(x, 65),
+ np.percentile(x, 70),
+ np.percentile(x, 75),
+ np.percentile(x, 80),
+ np.percentile(x, 85),
+ np.percentile(x, 90),
+ np.percentile(x, 95),
+ np.percentile(x, 99),
+ ]
+ res = [
+ (
+ "> p90 (%s)" % (lspercentile[8])
+ if z > lspercentile[8]
+ else (
+ "<= p10 (%s)" % (lspercentile[0])
+ if z <= lspercentile[0]
+ else (
+ "<= p20 (%s)" % (lspercentile[1])
+ if z <= lspercentile[1]
+ else (
+ "<= p30 (%s)" % (lspercentile[2])
+ if z <= lspercentile[2]
+ else (
+ "<= p40 (%s)" % (lspercentile[3])
+ if z <= lspercentile[3]
+ else (
+ "<= p50 (%s)" % (lspercentile[4])
+ if z <= lspercentile[4]
+ else (
+ "<= p60 (%s)" % (lspercentile[5])
+ if z <= lspercentile[5]
+ else (
+ "<= p70 (%s)" % (lspercentile[6])
+ if z <= lspercentile[6]
+ else (
+ "<= p80 (%s)" % (lspercentile[7])
+ if z <= lspercentile[7]
+ else (
+ "<= p90 (%s)" % (lspercentile[8])
+ if z <= lspercentile[8]
+ else "> p90 (%s)"
+ % (lspercentile[8])
+ )
+ )
+ )
+ )
+ )
+ )
+ )
+ )
+ )
+ )
+ for z in x
+ ]
+ elif granularity == "Medium":
+ lspercentile = [
+ np.percentile(x, 10),
+ np.percentile(x, 20),
+ np.percentile(x, 30),
+ np.percentile(x, 40),
+ np.percentile(x, 50),
+ np.percentile(x, 60),
+ np.percentile(x, 70),
+ np.percentile(x, 80),
+ np.percentile(x, 90),
+ ]
+ res = [
+ (
+ "<= p10 (%s)" % (lspercentile[0])
+ if z <= lspercentile[0]
+ else (
+ "<= p20 (%s)" % (lspercentile[1])
+ if z <= lspercentile[1]
+ else (
+ "<= p30 (%s)" % (lspercentile[2])
+ if z <= lspercentile[2]
+ else (
+ "<= p40 (%s)" % (lspercentile[3])
+ if z <= lspercentile[3]
+ else (
+ "<= p50 (%s)" % (lspercentile[4])
+ if z <= lspercentile[4]
+ else (
+ "<= p60 (%s)" % (lspercentile[5])
+ if z <= lspercentile[5]
+ else (
+ "<= p70 (%s)" % (lspercentile[6])
+ if z <= lspercentile[6]
+ else (
+ "<= p80 (%s)" % (lspercentile[7])
+ if z <= lspercentile[7]
+ else (
+ "<= p90 (%s)" % (lspercentile[8])
+ if z <= lspercentile[8]
+ else "> p90 (%s)" % (lspercentile[8])
+ )
+ )
+ )
+ )
+ )
+ )
+ )
+ )
+ )
+ for z in x
+ ]
+ else:
+ lspercentile = [
+ np.percentile(x, 15),
+ np.percentile(x, 50),
+ np.percentile(x, 85),
+ ]
+ res = [
+ (
+ "1-Very Low"
+ if z < lspercentile[0]
+ else (
+ "2-Low"
+ if z < lspercentile[1]
+ else "3-High" if z < lspercentile[2] else "4-Very High"
+ )
+ )
+ for z in x
+ ]
+ return res
+
+
+def kpi_transform(dfx, kpi_combo, kpi_combo_new):
+ """
+ Feature transformation from continuous feature to binned features for a list of features
+
+ Args
+ ----
+
+ dfx : DataFrame
+ DataFrame containing the features.
+
+ kpi_combo : list of string
+ List of feature names to be transformed
+
+ kpi_combo_new : list of string
+ List of new feature names to be assigned to the transformed features.
+
+ Returns
+ -------
+ dfx : DataFrame
+ Updated DataFrame containing the new features.
+ """
+ for j in range(len(kpi_combo)):
+ if type(dfx[kpi_combo[j]].values[0]) is str:
+ dfx[kpi_combo_new[j]] = dfx[kpi_combo[j]].values
+ dfx[kpi_combo_new[j]] = cat_group(dfx=dfx, kpix=kpi_combo_new[j])
+ else:
+ if len(kpi_combo) > 1:
+ dfx[kpi_combo_new[j]] = cat_continuous(
+ dfx[kpi_combo[j]].values, granularity="Low"
+ )
+ else:
+ dfx[kpi_combo_new[j]] = cat_continuous(
+ dfx[kpi_combo[j]].values, granularity="High"
+ )
+ return dfx
+
+
+def get_tree_leaves_mask(tree) -> np.ndarray:
+ """
+ Get mask array for tree leaves
+ Args:
+ tree: CausalTreeRegressor
+ Tree object
+ Returns: np.ndarray
+ Mask array
+
+ """
+ n_nodes = tree.tree_.node_count
+ children_left = tree.tree_.children_left
+ children_right = tree.tree_.children_right
+
+ node_depth = np.zeros(shape=n_nodes, dtype=np.int64)
+ is_leaves = np.zeros(shape=n_nodes, dtype=bool)
+ stack = [(0, 0)]
+ while len(stack) > 0:
+ node_id, depth = stack.pop()
+ node_depth[node_id] = depth
+
+ is_split_node = children_left[node_id] != children_right[node_id]
+
+ if is_split_node:
+ stack.append((children_left[node_id], depth + 1))
+ stack.append((children_right[node_id], depth + 1))
+ else:
+ is_leaves[node_id] = True
+ return is_leaves
+
+
+def timeit(exclude_kwargs: tuple = ()) -> Callable:
+ """
+ timeit decorator
+ Args:
+ exclude_kwargs: (tuple), keyword arguments that should be excluded from display
+ Returns: Callable
+
+ """
+
+ def wrapper(f: Callable):
+ def wrapped(*args, **kw):
+ ts = time.time()
+ result = f(*args, **kw)
+ te = time.time()
+ display_kw = {k: v for k, v in kw.items() if k not in exclude_kwargs}
+ print(
+ "Function: {} Kwargs: {} Elapsed time: {:2.4f}".format(
+ f.__name__, display_kw, te - ts
+ )
+ )
+ return result
+
+ return wrapped
+
+ return wrapper
diff --git a/causalml/source/causalml/match.py b/causalml/source/causalml/match.py
new file mode 100644
index 0000000000000000000000000000000000000000..395a0777a081bad9f100385cfab724ae2151008b
--- /dev/null
+++ b/causalml/source/causalml/match.py
@@ -0,0 +1,516 @@
+import argparse
+import logging
+import sys
+
+import numpy as np
+import pandas as pd
+from sklearn.neighbors import NearestNeighbors
+from sklearn.preprocessing import StandardScaler
+from sklearn.utils import check_random_state
+
+logger = logging.getLogger("causalml")
+
+
+def smd(feature, treatment):
+ """Calculate the standard mean difference (SMD) of a feature between the
+ treatment and control groups.
+
+ The definition is available at
+ https://www.ncbi.nlm.nih.gov/pmc/articles/PMC3144483/#s11title
+
+ Args:
+ feature (pandas.Series): a column of a feature to calculate SMD for
+ treatment (pandas.Series): a column that indicate whether a row is in
+ the treatment group or not
+
+ Returns:
+ (float): The SMD of the feature
+ """
+ t = feature[treatment == 1]
+ c = feature[treatment == 0]
+ return (t.mean() - c.mean()) / np.sqrt(0.5 * (t.var() + c.var()))
+
+
+def create_table_one(data, treatment_col, features, with_std=True, with_counts=True):
+ """Report balance in input features between the treatment and control groups.
+
+ References:
+ R's tableone at CRAN: https://github.com/kaz-yos/tableone
+ Python's tableone at PyPi: https://github.com/tompollard/tableone
+
+ Args:
+ data (pandas.DataFrame): total or matched sample data
+ treatment_col (str): the column name for the treatment
+ features (list of str): the column names of features
+ with_std (bool): whether to output std together with mean values as in () format
+ with_counts (bool): whether to include a row counting the total number of samples
+
+ Returns:
+ (pandas.DataFrame): A table with the means and standard deviations in
+ the treatment and control groups, and the SMD between two groups
+ for the features.
+ """
+ t1 = pd.pivot_table(
+ data[features + [treatment_col]],
+ columns=treatment_col,
+ aggfunc=[
+ lambda x: (
+ "{:.2f} ({:.2f})".format(x.mean(), x.std())
+ if with_std
+ else "{:.2f}".format(x.mean())
+ )
+ ],
+ )
+ t1.columns = t1.columns.droplevel(level=0)
+ t1["SMD"] = data[features].apply(lambda x: smd(x, data[treatment_col])).round(4)
+
+ if with_counts:
+ n_row = pd.pivot_table(
+ data[[features[0], treatment_col]], columns=treatment_col, aggfunc=["count"]
+ )
+ n_row.columns = n_row.columns.droplevel(level=0)
+ n_row["SMD"] = ""
+ n_row.index = ["n"]
+
+ t1 = pd.concat([n_row, t1], axis=0)
+
+ t1.columns.name = ""
+ t1.columns = ["Control", "Treatment", "SMD"]
+ t1.index.name = "Variable"
+
+ return t1
+
+
+class NearestNeighborMatch:
+ """
+ Propensity score matching based on the nearest neighbor algorithm.
+
+ Attributes:
+ caliper (float): threshold to be considered as a match.
+ replace (bool): whether to match with replacement or not
+ ratio (int): ratio of control / treatment to be matched.
+ shuffle (bool): whether to shuffle the treatment group data before
+ matching
+ treatment_to_control (bool): whether to match treatment to control
+ or control to treatment
+ random_state (numpy.random.RandomState or int): RandomState or an int
+ seed
+ n_jobs (int): The number of parallel jobs to run for neighbors search.
+ None means 1 unless in a joblib.parallel_backend context. -1 means using all processors
+ """
+
+ def __init__(
+ self,
+ caliper=0.2,
+ replace=False,
+ ratio=1,
+ shuffle=True,
+ treatment_to_control=True,
+ random_state=None,
+ n_jobs=-1,
+ ):
+ """Initialize a propensity score matching model.
+
+ Args:
+ caliper (float): threshold to be considered as a match.
+ replace (bool): whether to match with replacement or not
+ ratio (int): ratio of control / treatment to be matched.
+ shuffle (bool): whether to shuffle the treatment group data before
+ matching or not
+ random_state (numpy.random.RandomState or int): RandomState or an
+ int seed
+ n_jobs (int): The number of parallel jobs to run for neighbors search.
+ None means 1 unless in a joblib.parallel_backend context. -1 means using all processors
+ """
+ self.caliper = caliper
+ self.replace = replace
+ self.ratio = ratio
+ self.shuffle = shuffle
+ self.treatment_to_control = treatment_to_control
+ self.random_state = check_random_state(random_state)
+ self.n_jobs = n_jobs
+
+ def match(self, data, treatment_col, score_cols):
+ """Find matches from the control group by matching on specified columns
+ (propensity preferred).
+
+ Args:
+ data (pandas.DataFrame): total input data
+ treatment_col (str): the column name for the treatment
+ score_cols (list): list of column names for matching (propensity
+ column should be included)
+
+ Returns:
+ (pandas.DataFrame): The subset of data consisting of matched
+ treatment and control group data.
+ """
+ assert isinstance(score_cols, list), "score_cols must be a list"
+ treatment = data.loc[data[treatment_col] == 1, score_cols]
+ control = data.loc[data[treatment_col] == 0, score_cols]
+
+ # Picks whether to use treatment or control for matching direction
+ match_from = treatment if self.treatment_to_control else control
+ match_to = control if self.treatment_to_control else treatment
+ sdcal = self.caliper * np.std(data[score_cols].values)
+
+ if self.replace:
+ scaler = StandardScaler()
+ scaler.fit(data[score_cols])
+ match_from_scaled = pd.DataFrame(
+ scaler.transform(match_from), index=match_from.index
+ )
+ match_to_scaled = pd.DataFrame(
+ scaler.transform(match_to), index=match_to.index
+ )
+
+ # SD is the same as caliper because we use a StandardScaler above
+ sdcal = self.caliper
+
+ matching_model = NearestNeighbors(
+ n_neighbors=self.ratio, n_jobs=self.n_jobs
+ )
+ matching_model.fit(match_to_scaled)
+ distances, indices = matching_model.kneighbors(match_from_scaled)
+ # distances and indices are (n_obs, self.ratio) matrices.
+ # To index easily, reshape distances, indices and treatment into
+ # the (n_obs * self.ratio, 1) matrices and data frame.
+ distances = distances.T.flatten()
+ indices = indices.T.flatten()
+ match_from_scaled = pd.concat([match_from_scaled] * self.ratio, axis=0)
+
+ cond = (distances / np.sqrt(len(score_cols))) < sdcal
+ # Deduplicate the indices of the treatment group
+ from_idx_matched = np.unique(match_from_scaled.loc[cond].index)
+ # XXX: Should we deduplicate the indices of the control group too?
+ to_idx_matched = np.array(match_to_scaled.iloc[indices[cond]].index)
+ else:
+ assert len(score_cols) == 1, (
+ "Matching on multiple columns is only supported using the "
+ "replacement method (if matching on multiple columns, set "
+ "replace=True)."
+ )
+ # unpack score_cols for the single-variable matching case
+ score_col = score_cols[0]
+
+ if self.shuffle:
+ from_indices = self.random_state.permutation(match_from.index)
+ else:
+ from_indices = match_from.index
+
+ from_idx_matched = []
+ to_idx_matched = []
+ match_to["unmatched"] = True
+
+ for from_idx in from_indices:
+ dist = np.abs(
+ match_to.loc[match_to.unmatched, score_col]
+ - match_from.loc[from_idx, score_col]
+ )
+ # Gets self.ratio lowest dists
+ to_np_idx_list = np.argpartition(dist, self.ratio)[: self.ratio]
+ to_idx_list = dist.index[to_np_idx_list]
+ for i, to_idx in enumerate(to_idx_list):
+ if dist[to_idx] <= sdcal:
+ if i == 0:
+ from_idx_matched.append(from_idx)
+ to_idx_matched.append(to_idx)
+ match_to.loc[to_idx, "unmatched"] = False
+
+ return data.loc[
+ np.concatenate([np.array(from_idx_matched), np.array(to_idx_matched)])
+ ]
+
+ def match_by_group(self, data, treatment_col, score_cols, groupby_col):
+ """Find matches from the control group stratified by groupby_col, by
+ matching on specified columns (propensity preferred).
+
+ Args:
+ data (pandas.DataFrame): total sample data
+ treatment_col (str): the column name for the treatment
+ score_cols (list): list of column names for matching (propensity
+ column should be included)
+ groupby_col (str): the column name to be used for stratification
+
+ Returns:
+ (pandas.DataFrame): The subset of data consisting of matched
+ treatment and control group data.
+ """
+ matched = data.groupby(groupby_col).apply(
+ lambda x: self.match(
+ data=x, treatment_col=treatment_col, score_cols=score_cols
+ )
+ )
+ return matched.reset_index(level=0, drop=True)
+
+
+class MatchOptimizer:
+ def __init__(
+ self,
+ treatment_col="is_treatment",
+ ps_col="pihat",
+ user_col=None,
+ matching_covariates=["pihat"],
+ max_smd=0.1,
+ max_deviation=0.1,
+ caliper_range=(0.01, 0.5),
+ max_pihat_range=(0.95, 0.999),
+ max_iter_per_param=5,
+ min_users_per_group=1000,
+ smd_cols=["pihat"],
+ dev_cols_transformations={"pihat": np.mean},
+ dev_factor=1.0,
+ verbose=True,
+ ):
+ """Finds the set of parameters that gives the best matching result.
+
+ Score = (number of features with SMD > max_smd)
+ + (sum of deviations for important variables
+ * deviation factor)
+
+ The logic behind the scoring is that we are most concerned with
+ minimizing the number of features where SMD is lower than a certain
+ threshold (max_smd). However, we would also like the matched dataset
+ not deviate too much from the original dataset, in terms of key
+ variable(s), so that we still retain a similar userbase.
+
+ Args:
+ - treatment_col (str): name of the treatment column
+ - ps_col (str): name of the propensity score column
+ - max_smd (float): maximum acceptable SMD
+ - max_deviation (float): maximum acceptable deviation for
+ important variables
+ - caliper_range (tuple): low and high bounds for caliper search
+ range
+ - max_pihat_range (tuple): low and high bounds for max pihat
+ search range
+ - max_iter_per_param (int): maximum number of search values per
+ parameters
+ - min_users_per_group (int): minimum number of users per group in
+ matched set
+ - smd_cols (list): score is more sensitive to these features
+ exceeding max_smd
+ - dev_factor (float): importance weight factor for dev_cols
+ (e.g. dev_factor=1 means a 10% deviation leads to penalty of 1
+ in score)
+ - dev_cols_transformations (dict): dict of transformations to be
+ made on dev_cols
+ - verbose (bool): boolean flag for printing statements
+
+ Returns:
+ The best matched dataset (pd.DataFrame)
+ """
+ self.treatment_col = treatment_col
+ self.ps_col = ps_col
+ self.user_col = user_col
+ self.matching_covariates = matching_covariates
+ self.max_smd = max_smd
+ self.max_deviation = max_deviation
+ self.caliper_range = np.linspace(*caliper_range, num=max_iter_per_param)
+ self.max_pihat_range = np.linspace(*max_pihat_range, num=max_iter_per_param)
+ self.max_iter_per_param = max_iter_per_param
+ self.min_users_per_group = min_users_per_group
+ self.smd_cols = smd_cols
+ self.dev_factor = dev_factor
+ self.dev_cols_transformations = dev_cols_transformations
+ self.best_params = {}
+ self.best_score = 1e7 # ideal score is 0
+ self.verbose = verbose
+ self.pass_all = False
+
+ def single_match(self, score_cols, pihat_threshold, caliper):
+ matcher = NearestNeighborMatch(caliper=caliper, replace=True)
+ df_matched = matcher.match(
+ data=self.df[self.df[self.ps_col] < pihat_threshold],
+ treatment_col=self.treatment_col,
+ score_cols=score_cols,
+ )
+ return df_matched
+
+ def check_table_one(self, tableone, matched, score_cols, pihat_threshold, caliper):
+ # check if better than past runs
+ smd_values = np.abs(tableone[tableone.index != "n"]["SMD"].astype(float))
+ num_cols_over_smd = (smd_values >= self.max_smd).sum()
+ self.cols_to_fix = (
+ smd_values[smd_values >= self.max_smd]
+ .sort_values(ascending=False)
+ .index.values
+ )
+ if self.user_col is None:
+ num_users_per_group = (
+ matched.reset_index().groupby(self.treatment_col)["index"].count().min()
+ )
+ else:
+ num_users_per_group = (
+ matched.groupby(self.treatment_col)[self.user_col].count().min()
+ )
+ deviations = [
+ np.abs(
+ self.original_stats[col]
+ / matched[matched[self.treatment_col] == 1][col].mean()
+ - 1
+ )
+ for col in self.dev_cols_transformations.keys()
+ ]
+
+ score = num_cols_over_smd
+ score += len(
+ [col for col in self.smd_cols if smd_values.loc[col] >= self.max_smd]
+ )
+ score += np.sum([dev * 10 * self.dev_factor for dev in deviations])
+
+ # check if can be considered as best score
+ if score < self.best_score and num_users_per_group > self.min_users_per_group:
+ self.best_score = score
+ self.best_params = {
+ "score_cols": score_cols.copy(),
+ "pihat": pihat_threshold,
+ "caliper": caliper,
+ }
+ self.best_matched = matched.copy()
+ if self.verbose:
+ logger.info(
+ "\tScore: {:.03f} (Best Score: {:.03f})\n".format(
+ score, self.best_score
+ )
+ )
+
+ # check if passes all criteria
+ self.pass_all = (
+ (num_users_per_group > self.min_users_per_group)
+ and (num_cols_over_smd == 0)
+ and all(dev < self.max_deviation for dev in deviations)
+ )
+
+ def match_and_check(self, score_cols, pihat_threshold, caliper):
+ if self.verbose:
+ logger.info(
+ "Preparing match for: caliper={:.03f}, "
+ "pihat_threshold={:.03f}, "
+ "score_cols={}".format(caliper, pihat_threshold, score_cols)
+ )
+ df_matched = self.single_match(
+ score_cols=score_cols, pihat_threshold=pihat_threshold, caliper=caliper
+ )
+ tableone = create_table_one(
+ df_matched, self.treatment_col, self.matching_covariates
+ )
+ self.check_table_one(tableone, df_matched, score_cols, pihat_threshold, caliper)
+
+ def search_best_match(self, df):
+ self.df = df
+
+ self.original_stats = {}
+ for col, trans in self.dev_cols_transformations.items():
+ self.original_stats[col] = trans(
+ self.df[self.df[self.treatment_col] == 1][col]
+ )
+
+ # search best max pihat
+ if self.verbose:
+ logger.info("SEARCHING FOR BEST PIHAT")
+ score_cols = [self.ps_col]
+ caliper = self.caliper_range[-1]
+ for pihat_threshold in self.max_pihat_range:
+ self.match_and_check(score_cols, pihat_threshold, caliper)
+
+ # search best score_cols
+ if self.verbose:
+ logger.info("SEARCHING FOR BEST SCORE_COLS")
+ pihat_threshold = self.best_params["pihat"]
+ caliper = self.caliper_range[int(self.caliper_range.shape[0] / 2)]
+ score_cols = [self.ps_col]
+ while not self.pass_all:
+ if len(self.cols_to_fix) == 0:
+ break
+ elif np.intersect1d(self.cols_to_fix, score_cols).shape[0] > 0:
+ break
+ else:
+ score_cols.append(self.cols_to_fix[0])
+ self.match_and_check(score_cols, pihat_threshold, caliper)
+
+ # search best caliper
+ if self.verbose:
+ logger.info("SEARCHING FOR BEST CALIPER")
+ score_cols = self.best_params["score_cols"]
+ pihat_threshold = self.best_params["pihat"]
+ for caliper in self.caliper_range:
+ self.match_and_check(score_cols, pihat_threshold, caliper)
+
+ # summarize
+ if self.verbose:
+ logger.info("\n-----\nBest params are:\n{}".format(self.best_params))
+
+ return self.best_matched
+
+
+if __name__ == "__main__":
+ from .features import load_data
+ from .propensity import ElasticNetPropensityModel
+
+ TREATMENT_COL = "treatment"
+ SCORE_COL = "score"
+ GROUPBY_COL = "group"
+
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--input-file", required=True, dest="input_file")
+ parser.add_argument("--output-file", required=True, dest="output_file")
+ parser.add_argument("--treatment-col", default=TREATMENT_COL, dest="treatment_col")
+ parser.add_argument("--groupby-col", default=GROUPBY_COL, dest="groupby_col")
+ parser.add_argument("--score-col", default=SCORE_COL, dest="score_col")
+ parser.add_argument("--feature-cols", nargs="+", required=True, dest="feature_cols")
+ parser.add_argument(
+ "--matching-cols", nargs="+", required=True, dest="matching_cols"
+ )
+ parser.add_argument("--caliper", type=float, default=0.2)
+ parser.add_argument("--replace", default=False, action="store_true")
+ parser.add_argument("--ratio", type=int, default=1)
+
+ args = parser.parse_args()
+
+ logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
+
+ logger.info("Loading data from {}".format(args.input_file))
+ df = pd.read_csv(args.input_file)
+ df[args.treatment_col] = df[args.treatment_col].astype(int)
+ logger.info("shape: {}\n{}".format(df.shape, df.head()))
+
+ pm = ElasticNetPropensityModel(random_state=42)
+ w = df[args.treatment_col].values
+ X = load_data(
+ data=df,
+ features=args.feature_cols,
+ )
+
+ logger.info("Scoring with a propensity model: {}".format(pm))
+ df[args.score_col] = pm.fit_predict(X, w)
+
+ logger.info(
+ "Balance before matching:\n{}".format(
+ create_table_one(
+ data=df, treatment_col=args.treatment_col, features=args.matching_cols
+ )
+ )
+ )
+ logger.info(
+ "Matching based on the propensity score with the nearest neighbor model"
+ )
+ psm = NearestNeighborMatch(replace=args.replace, ratio=args.ratio, random_state=42)
+ matched = psm.match_by_group(
+ data=df,
+ treatment_col=args.treatment_col,
+ score_cols=[args.score_col],
+ groupby_col=args.groupby_col,
+ )
+ logger.info("shape: {}\n{}".format(matched.shape, matched.head()))
+
+ logger.info(
+ "Balance after matching:\n{}".format(
+ create_table_one(
+ data=matched,
+ treatment_col=args.treatment_col,
+ features=args.matching_cols,
+ )
+ )
+ )
+ matched.to_csv(args.output_file, index=False)
+ logger.info("Matched data saved as {}".format(args.output_file))
diff --git a/causalml/source/causalml/metrics/__init__.py b/causalml/source/causalml/metrics/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..0bc9d187f6559d31991b1390972e4d6045efb82d
--- /dev/null
+++ b/causalml/source/causalml/metrics/__init__.py
@@ -0,0 +1,34 @@
+from .classification import roc_auc_score, logloss, classification_metrics # noqa
+from .regression import (
+ ape,
+ mape,
+ mae,
+ rmse,
+ r2_score,
+ gini,
+ smape,
+ regression_metrics,
+) # noqa
+from .visualize import (
+ plot,
+ plot_gain,
+ plot_lift,
+ plot_qini,
+ plot_tmlegain,
+ plot_tmleqini,
+) # noqa
+from .visualize import (
+ get_cumgain,
+ get_cumlift,
+ get_qini,
+ get_tmlegain,
+ get_tmleqini,
+) # noqa
+from .visualize import auuc_score, qini_score # noqa
+from .sensitivity import Sensitivity, SensitivityPlaceboTreatment # noqa
+from .sensitivity import (
+ SensitivityRandomCause,
+ SensitivityRandomReplace,
+ SensitivitySubsetData,
+ SensitivitySelectionBias,
+) # noqa
diff --git a/causalml/source/causalml/metrics/classification.py b/causalml/source/causalml/metrics/classification.py
new file mode 100644
index 0000000000000000000000000000000000000000..d7dec9da570473eab5b1fd636c0f356f2b1477f1
--- /dev/null
+++ b/causalml/source/causalml/metrics/classification.py
@@ -0,0 +1,36 @@
+import logging
+from sklearn.metrics import log_loss, roc_auc_score
+
+from .const import EPS
+from .regression import regression_metrics
+
+logger = logging.getLogger("causalml")
+
+
+def logloss(y, p):
+ """Bounded log loss error.
+ Args:
+ y (numpy.array): target
+ p (numpy.array): prediction
+ Returns:
+ bounded log loss error
+ """
+
+ p[p < EPS] = EPS
+ p[p > 1 - EPS] = 1 - EPS
+ return log_loss(y, p)
+
+
+def classification_metrics(
+ y, p, w=None, metrics={"AUC": roc_auc_score, "Log Loss": logloss}
+):
+ """Log metrics for classifiers.
+
+ Args:
+ y (numpy.array): target
+ p (numpy.array): prediction
+ w (numpy.array, optional): a treatment vector (1 or True: treatment, 0 or False: control). If given, log
+ metrics for the treatment and control group separately
+ metrics (dict, optional): a dictionary of the metric names and functions
+ """
+ regression_metrics(y=y, p=p, w=w, metrics=metrics)
diff --git a/causalml/source/causalml/metrics/const.py b/causalml/source/causalml/metrics/const.py
new file mode 100644
index 0000000000000000000000000000000000000000..abca9d3e3245709c0c41862d84211ae4d4668996
--- /dev/null
+++ b/causalml/source/causalml/metrics/const.py
@@ -0,0 +1 @@
+EPS = 1e-15
diff --git a/causalml/source/causalml/metrics/regression.py b/causalml/source/causalml/metrics/regression.py
new file mode 100644
index 0000000000000000000000000000000000000000..ccd3f28a0d821dc34cb7d8da984c3cd572fddd24
--- /dev/null
+++ b/causalml/source/causalml/metrics/regression.py
@@ -0,0 +1,126 @@
+import logging
+import numpy as np
+from sklearn.metrics import mean_squared_error as mse
+from sklearn.metrics import mean_absolute_error as mae # noqa
+from sklearn.metrics import r2_score # noqa
+
+from .const import EPS
+
+logger = logging.getLogger("causalml")
+
+
+def ape(y, p):
+ """Absolute Percentage Error (APE).
+ Args:
+ y (float): target
+ p (float): prediction
+
+ Returns:
+ e (float): APE
+ """
+
+ assert np.abs(y) > EPS
+ return np.abs(1 - p / y)
+
+
+def mape(y, p):
+ """Mean Absolute Percentage Error (MAPE).
+ Args:
+ y (numpy.array): target
+ p (numpy.array): prediction
+
+ Returns:
+ e (numpy.float64): MAPE
+ """
+
+ filt = np.abs(y) > EPS
+ return np.mean(np.abs(1 - p[filt] / y[filt]))
+
+
+def smape(y, p):
+ """Symmetric Mean Absolute Percentage Error (sMAPE).
+ Args:
+ y (numpy.array): target
+ p (numpy.array): prediction
+
+ Returns:
+ e (numpy.float64): sMAPE
+ """
+ return 2.0 * np.mean(np.abs(y - p) / (np.abs(y) + np.abs(p)))
+
+
+def rmse(y, p):
+ """Root Mean Squared Error (RMSE).
+ Args:
+ y (numpy.array): target
+ p (numpy.array): prediction
+
+ Returns:
+ e (numpy.float64): RMSE
+ """
+
+ # check and get number of samples
+ assert y.shape == p.shape
+
+ return np.sqrt(mse(y, p))
+
+
+def gini(y, p):
+ """Normalized Gini Coefficient.
+
+ Args:
+ y (numpy.array): target
+ p (numpy.array): prediction
+
+ Returns:
+ e (numpy.float64): normalized Gini coefficient
+ """
+
+ # check and get number of samples
+ assert y.shape == p.shape
+
+ n_samples = y.shape[0]
+
+ # sort rows on prediction column
+ # (from largest to smallest)
+ arr = np.array([y, p]).transpose()
+ true_order = arr[arr[:, 0].argsort()][::-1, 0]
+ pred_order = arr[arr[:, 1].argsort()][::-1, 0]
+
+ # get Lorenz curves
+ l_true = np.cumsum(true_order) / np.sum(true_order)
+ l_pred = np.cumsum(pred_order) / np.sum(pred_order)
+ l_ones = np.linspace(1 / n_samples, 1, n_samples)
+
+ # get Gini coefficients (area between curves)
+ g_true = np.sum(l_ones - l_true)
+ g_pred = np.sum(l_ones - l_pred)
+
+ # normalize to true Gini coefficient
+ return g_pred / g_true
+
+
+def regression_metrics(
+ y, p, w=None, metrics={"RMSE": rmse, "sMAPE": smape, "Gini": gini}
+):
+ """Log metrics for regressors.
+
+ Args:
+ y (numpy.array): target
+ p (numpy.array): prediction
+ w (numpy.array, optional): a treatment vector (1 or True: treatment, 0 or False: control). If given, log
+ metrics for the treatment and control group separately
+ metrics (dict, optional): a dictionary of the metric names and functions
+ """
+ assert metrics
+ assert y.shape[0] == p.shape[0]
+
+ for name, func in metrics.items():
+ if w is not None:
+ assert y.shape[0] == w.shape[0]
+ if w.dtype != bool:
+ w = w == 1
+ logger.info("{:>8s} (Control): {:10.4f}".format(name, func(y[~w], p[~w])))
+ logger.info("{:>8s} (Treatment): {:10.4f}".format(name, func(y[w], p[w])))
+ else:
+ logger.info("{:>8s}: {:10.4f}".format(name, func(y, p)))
diff --git a/causalml/source/causalml/metrics/sensitivity.py b/causalml/source/causalml/metrics/sensitivity.py
new file mode 100644
index 0000000000000000000000000000000000000000..3bd186d21b744f16209ebc3b3fe649ccb03a656f
--- /dev/null
+++ b/causalml/source/causalml/metrics/sensitivity.py
@@ -0,0 +1,607 @@
+import logging
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+from importlib import import_module
+
+logger = logging.getLogger("sensitivity")
+
+SUMMARY_COLS = ["Method", "ATE", "New ATE", "New ATE LB", "New ATE UB"]
+
+
+def one_sided(alpha, p, treatment):
+ """One sided confounding function.
+ Reference: Blackwell, Matthew. "A selection bias approach to sensitivity analysis
+ for causal effects." Political Analysis 22.2 (2014): 169-182.
+ https://www.mattblackwell.org/files/papers/causalsens.pdf
+
+ Args:
+ alpha (np.array): a confounding values vector
+ p (np.array): a propensity score vector between 0 and 1
+ treatment (np.array): a treatment vector (1 if treated, otherwise 0)
+ """
+ assert p.shape[0] == treatment.shape[0]
+ adj = alpha * (1 - p) * treatment - alpha * p * (1 - treatment)
+ return adj
+
+
+def alignment(alpha, p, treatment):
+ """Alignment confounding function.
+ Reference: Blackwell, Matthew. "A selection bias approach to sensitivity analysis
+ for causal effects." Political Analysis 22.2 (2014): 169-182.
+ https://www.mattblackwell.org/files/papers/causalsens.pdf
+
+ Args:
+ alpha (np.array): a confounding values vector
+ p (np.array): a propensity score vector between 0 and 1
+ treatment (np.array): a treatment vector (1 if treated, otherwise 0)
+ """
+
+ assert p.shape[0] == treatment.shape[0]
+ adj = alpha * (1 - p) * treatment + alpha * p * (1 - treatment)
+ return adj
+
+
+def one_sided_att(alpha, p, treatment):
+ """One sided confounding function for the average effect of the treatment among the treated units (ATT)
+
+ Reference: Blackwell, Matthew. "A selection bias approach to sensitivity analysis
+ for causal effects." Political Analysis 22.2 (2014): 169-182.
+ https://www.mattblackwell.org/files/papers/causalsens.pdf
+
+ Args:
+ alpha (np.array): a confounding values vector
+ p (np.array): a propensity score vector between 0 and 1
+ treatment (np.array): a treatment vector (1 if treated, otherwise 0)
+ """
+ assert p.shape[0] == treatment.shape[0]
+ adj = alpha * (1 - treatment)
+ return adj
+
+
+def alignment_att(alpha, p, treatment):
+ """Alignment confounding function for the average effect of the treatment among the treated units (ATT)
+
+ Reference: Blackwell, Matthew. "A selection bias approach to sensitivity analysis
+ for causal effects." Political Analysis 22.2 (2014): 169-182.
+ https://www.mattblackwell.org/files/papers/causalsens.pdf
+
+ Args:
+ alpha (np.array): a confounding values vector
+ p (np.array): a propensity score vector between 0 and 1
+ treatment (np.array): a treatment vector (1 if treated, otherwise 0)
+ """
+ assert p.shape[0] == treatment.shape[0]
+ adj = alpha * (1 - treatment)
+ return adj
+
+
+class Sensitivity:
+ """A Sensitivity Check class to support Placebo Treatment, Irrelevant Additional Confounder
+ and Subset validation refutation methods to verify causal inference.
+
+ Reference: https://github.com/microsoft/dowhy/blob/master/dowhy/causal_refuters/
+ """
+
+ def __init__(
+ self,
+ df,
+ inference_features,
+ p_col,
+ treatment_col,
+ outcome_col,
+ learner,
+ *args,
+ **kwargs,
+ ):
+ """Initialize.
+
+ Args:
+ df (pd.DataFrame): input data frame
+ inferenece_features (list of str): a list of columns that used in learner for inference
+ p_col (str): column name of propensity score
+ treatment_col (str): column name of whether in treatment of control
+ outcome_col (str): column name of outcome
+ learner (model): a model to estimate outcomes and treatment effects
+ """
+
+ self.df = df
+ self.inference_features = inference_features
+ self.p_col = p_col
+ self.treatment_col = treatment_col
+ self.outcome_col = outcome_col
+ self.learner = learner
+
+ def get_prediction(self, X, p, treatment, y):
+ """Return the treatment effects prediction.
+
+ Args:
+ X (np.matrix): a feature matrix
+ p (np.array): a propensity score vector between 0 and 1
+ treatment (np.array): a treatment vector (1 if treated, otherwise 0)
+ y (np.array): an outcome vector
+ Returns:
+ (numpy.ndarray): Predictions of treatment effects
+ """
+
+ learner = self.learner
+ try:
+ preds = learner.fit_predict(X=X, p=p, treatment=treatment, y=y).flatten()
+ except TypeError:
+ preds = learner.fit_predict(X=X, treatment=treatment, y=y).flatten()
+ return preds
+
+ def get_ate_ci(self, X, p, treatment, y):
+ """Return the confidence intervals for treatment effects prediction.
+
+ Args:
+ X (np.matrix): a feature matrix
+ p (np.array): a propensity score vector between 0 and 1
+ treatment (np.array): a treatment vector (1 if treated, otherwise 0)
+ y (np.array): an outcome vector
+ Returns:
+ (numpy.ndarray): Mean and confidence interval (LB, UB) of the ATE estimate.
+ """
+
+ try:
+ ate, ate_lower, ate_upper = self.learner.estimate_ate(
+ X=X, p=p, treatment=treatment, y=y, return_ci=True
+ )
+ except TypeError:
+ ate, ate_lower, ate_upper = self.learner.estimate_ate(
+ X=X, p=p, treatment=treatment, y=y
+ )
+ return ate[0], ate_lower[0], ate_upper[0]
+
+ @staticmethod
+ def get_class_object(method_name, *args, **kwargs):
+ """Return class object based on input method
+ Args:
+ method_name (list of str): a list of sensitivity analysis method
+ Returns:
+ (class): Sensitivy Class
+ """
+
+ method_list = [
+ "Placebo Treatment",
+ "Random Cause",
+ "Subset Data",
+ "Random Replace",
+ "Selection Bias",
+ ]
+ class_name = "Sensitivity" + method_name.replace(" ", "")
+
+ try:
+ getattr(import_module("causalml.metrics.sensitivity"), class_name)
+ return getattr(import_module("causalml.metrics.sensitivity"), class_name)
+ except AttributeError:
+ raise AttributeError(
+ "{} is not an existing method for sensitiviy analysis.".format(
+ method_name
+ )
+ + " Select one of {}".format(method_list)
+ )
+
+ def sensitivity_analysis(
+ self, methods, sample_size=None, confound="one_sided", alpha_range=None
+ ):
+ """Return the sensitivity data by different method
+
+ Args:
+ method (list of str): a list of sensitivity analysis method
+ sample_size (float, optional): ratio for subset the original data
+ confound (string, optional): the name of confouding function
+ alpha_range (np.array, optional): a parameter to pass the confounding function
+
+ Returns:
+ X (np.matrix): a feature matrix
+ p (np.array): a propensity score vector between 0 and 1
+ treatment (np.array): a treatment vector (1 if treated, otherwise 0)
+ y (np.array): an outcome vector
+ """
+ if alpha_range is None:
+ y = self.df[self.outcome_col]
+ iqr = y.quantile(0.75) - y.quantile(0.25)
+ alpha_range = np.linspace(-iqr / 2, iqr / 2, 11)
+ if 0 not in alpha_range:
+ alpha_range = np.append(alpha_range, 0)
+ else:
+ alpha_range = alpha_range
+
+ alpha_range.sort()
+
+ summary = []
+ for method in methods:
+ sens = self.get_class_object(method)
+ sens = sens(
+ self.df,
+ self.inference_features,
+ self.p_col,
+ self.treatment_col,
+ self.outcome_col,
+ self.learner,
+ sample_size=sample_size,
+ confound=confound,
+ alpha_range=alpha_range,
+ )
+
+ if method == "Subset Data":
+ method = method + "(sample size @{})".format(sample_size)
+
+ sens_df = sens.summary(method=method)
+ summary.append(sens_df.values.tolist()[0])
+
+ summary_df = pd.DataFrame(summary, columns=SUMMARY_COLS)
+
+ return summary_df
+
+ def summary(self, method):
+ """Summary report
+ Args:
+ method_name (str): sensitivity analysis method
+
+ Returns:
+ (pd.DataFrame): a summary dataframe
+ """
+ method_name = method
+
+ X = self.df[self.inference_features].values
+ p = self.df[self.p_col].values
+ treatment = self.df[self.treatment_col].values
+ y = self.df[self.outcome_col].values
+
+ preds = self.get_prediction(X, p, treatment, y)
+ ate = preds.mean()
+ ate_new, ate_new_lower, ate_new_upper = self.sensitivity_estimate()
+
+ sensitivity_summary = pd.DataFrame(
+ [method_name, ate, ate_new, ate_new_lower, ate_new_upper]
+ ).T
+ sensitivity_summary.columns = SUMMARY_COLS
+ return sensitivity_summary
+
+ def sensitivity_estimate(self):
+ raise NotImplementedError
+
+
+class SensitivityPlaceboTreatment(Sensitivity):
+ """Replaces the treatment variable with a new variable randomly generated."""
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+
+ def sensitivity_estimate(self):
+ """Summary report
+ Args:
+ return_ci (str): sensitivity analysis method
+
+ Returns:
+ (pd.DataFrame): a summary dataframe
+ """
+ num_rows = self.df.shape[0]
+
+ X = self.df[self.inference_features].values
+ p = self.df[self.p_col].values
+ treatment_new = np.random.randint(2, size=num_rows)
+ y = self.df[self.outcome_col].values
+
+ ate_new, ate_new_lower, ate_new_upper = self.get_ate_ci(X, p, treatment_new, y)
+ return ate_new, ate_new_lower, ate_new_upper
+
+
+class SensitivityRandomCause(Sensitivity):
+ """Adds an irrelevant random covariate to the dataframe."""
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+
+ def sensitivity_estimate(self):
+ num_rows = self.df.shape[0]
+ new_data = np.random.randn(num_rows)
+
+ X = self.df[self.inference_features].values
+ p = self.df[self.p_col].values
+ treatment = self.df[self.treatment_col].values
+ y = self.df[self.outcome_col].values
+ X_new = np.hstack((X, new_data.reshape((-1, 1))))
+
+ ate_new, ate_new_lower, ate_new_upper = self.get_ate_ci(X_new, p, treatment, y)
+ return ate_new, ate_new_lower, ate_new_upper
+
+
+class SensitivityRandomReplace(Sensitivity):
+ """Replaces a random covariate with an irrelevant variable."""
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ if "replaced_feature" not in kwargs:
+ replaced_feature_index = np.random.randint(len(self.inference_features))
+ self.replaced_feature = self.inference_features[replaced_feature_index]
+ else:
+ self.replaced_feature = kwargs["replaced_feature"]
+
+ def sensitivity_estimate(self):
+ """Replaces a random covariate with an irrelevant variable."""
+
+ logger.info(
+ "Replace feature {} with an random irrelevant variable".format(
+ self.replaced_feature
+ )
+ )
+ df_new = self.df.copy()
+ num_rows = self.df.shape[0]
+ df_new[self.replaced_feature] = np.random.randn(num_rows)
+
+ X_new = df_new[self.inference_features].values
+ p_new = df_new[self.p_col].values
+ treatment_new = df_new[self.treatment_col].values
+ y_new = df_new[self.outcome_col].values
+
+ ate_new, ate_new_lower, ate_new_upper = self.get_ate_ci(
+ X_new, p_new, treatment_new, y_new
+ )
+ return ate_new, ate_new_lower, ate_new_upper
+
+
+class SensitivitySubsetData(Sensitivity):
+ """Takes a random subset of size sample_size of the data."""
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.sample_size = kwargs["sample_size"]
+ assert self.sample_size is not None
+
+ def sensitivity_estimate(self):
+ df_new = self.df.sample(frac=self.sample_size).copy()
+
+ X_new = df_new[self.inference_features].values
+ p_new = df_new[self.p_col].values
+ treatment_new = df_new[self.treatment_col].values
+ y_new = df_new[self.outcome_col].values
+
+ ate_new, ate_new_lower, ate_new_upper = self.get_ate_ci(
+ X_new, p_new, treatment_new, y_new
+ )
+ return ate_new, ate_new_lower, ate_new_upper
+
+
+class SensitivitySelectionBias(Sensitivity):
+ """Reference:
+
+ [1] Blackwell, Matthew. "A selection bias approach to sensitivity analysis
+ for causal effects." Political Analysis 22.2 (2014): 169-182.
+ https://www.mattblackwell.org/files/papers/causalsens.pdf
+
+ [2] Confouding parameter alpha_range using the same range as in:
+ https://github.com/mattblackwell/causalsens/blob/master/R/causalsens.R
+
+ """
+
+ def __init__(
+ self,
+ *args,
+ confound="one_sided",
+ alpha_range=None,
+ sensitivity_features=None,
+ **kwargs,
+ ):
+ super().__init__(*args, **kwargs)
+ """Initialize.
+
+ Args:
+ confound (string): the name of confouding function
+ alpha_range (np.array): a parameter to pass the confounding function
+ sensitivity_features (list of str): ): a list of columns that to check each individual partial r-square
+ """
+
+ logger.info("Only works for linear outcome models right now. Check back soon.")
+ confounding_functions = {
+ "one_sided": one_sided,
+ "alignment": alignment,
+ "one_sided_att": one_sided_att,
+ "alignment_att": alignment_att,
+ }
+
+ try:
+ confound_func = confounding_functions[confound]
+ except KeyError:
+ raise NotImplementedError(
+ f"Confounding function, {confound} is not implemented. \
+ Use one of {confounding_functions.keys()}"
+ )
+
+ self.confound = confound_func
+
+ if sensitivity_features is None:
+ self.sensitivity_features = self.inference_features
+ else:
+ self.sensitivity_features = sensitivity_features
+
+ if alpha_range is None:
+ y = self.df[self.outcome_col]
+ iqr = y.quantile(0.75) - y.quantile(0.25)
+ self.alpha_range = np.linspace(-iqr / 2, iqr / 2, 11)
+ if 0 not in self.alpha_range:
+ self.alpha_range = np.append(self.alpha_range, 0)
+ else:
+ self.alpha_range = alpha_range
+
+ self.alpha_range.sort()
+
+ def causalsens(self):
+ alpha_range = self.alpha_range
+ confound = self.confound
+ df = self.df
+ X = df[self.inference_features].values
+ p = df[self.p_col].values
+ treatment = df[self.treatment_col].values
+ y = df[self.outcome_col].values
+
+ preds = self.get_prediction(X, p, treatment, y)
+
+ sens_df = pd.DataFrame()
+
+ sens = []
+ for a in alpha_range:
+ adj = confound(a, p, treatment)
+ preds_adj = y - adj
+ s_preds = self.get_prediction(X, p, treatment, preds_adj)
+ ate, ate_lb, ate_ub = self.get_ate_ci(X, p, treatment, preds_adj)
+
+ s_preds_residul = preds_adj - s_preds
+ rsqs = a**2 * np.var(treatment) / np.var(s_preds_residul)
+
+ sens.append([a, rsqs, ate, ate_lb, ate_ub])
+
+ sens_df = pd.DataFrame(
+ sens, columns=["alpha", "rsqs", "New ATE", "New ATE LB", "New ATE UB"]
+ )
+
+ rss = np.sum(np.square(y - preds))
+ partial_rsqs = []
+ for feature in self.sensitivity_features:
+ df_new = df.copy()
+ X_new = df_new[self.inference_features].drop(feature, axis=1).copy()
+ y_new_preds = self.get_prediction(X_new, p, treatment, y)
+ rss_new = np.sum(np.square(y - y_new_preds))
+ partial_rsqs.append(((rss_new - rss) / rss))
+
+ partial_rsqs_df = pd.DataFrame([self.sensitivity_features, partial_rsqs]).T
+ partial_rsqs_df.columns = ["feature", "partial_rsqs"]
+
+ return sens_df, partial_rsqs_df
+
+ def summary(self, method="Selection Bias"):
+ """Summary report for Selection Bias Method
+ Args:
+ method_name (str): sensitivity analysis method
+ Returns:
+ (pd.DataFrame): a summary dataframe
+ """
+
+ method_name = method
+ sensitivity_summary = self.causalsens()[0]
+ sensitivity_summary["Method"] = [
+ method_name + " (alpha@" + str(round(i, 5)) + ", with r-sqaure:"
+ for i in sensitivity_summary.alpha
+ ]
+ sensitivity_summary["Method"] = sensitivity_summary[
+ "Method"
+ ] + sensitivity_summary["rsqs"].round(5).astype(str)
+ sensitivity_summary["ATE"] = sensitivity_summary[
+ sensitivity_summary.alpha == 0
+ ]["New ATE"]
+ return sensitivity_summary[SUMMARY_COLS]
+
+ @staticmethod
+ def plot(sens_df, partial_rsqs_df=None, type="raw", ci=False, partial_rsqs=False):
+ """Plot the results of a sensitivity analysis against unmeasured
+ Args:
+ sens_df (pandas.DataFrame): a data frame output from causalsens
+ partial_rsqs_d (pandas.DataFrame) : a data frame output from causalsens including partial rsqure
+ type (str, optional): the type of plot to draw, 'raw' or 'r.squared' are supported
+ ci (bool, optional): whether plot confidence intervals
+ partial_rsqs (bool, optional): whether plot partial rsquare results
+ """
+
+ if type == "raw" and not ci:
+ fig, ax = plt.subplots()
+ y_max = round(sens_df["New ATE UB"].max() * 1.1, 4)
+ y_min = round(sens_df["New ATE LB"].min() * 0.9, 4)
+ x_max = round(sens_df.alpha.max() * 1.1, 4)
+ x_min = round(sens_df.alpha.min() * 0.9, 4)
+ plt.ylim(y_min, y_max)
+ plt.xlim(x_min, x_max)
+ ax.plot(sens_df.alpha, sens_df["New ATE"])
+ elif type == "raw" and ci:
+ fig, ax = plt.subplots()
+ y_max = round(sens_df["New ATE UB"].max() * 1.1, 4)
+ y_min = round(sens_df["New ATE LB"].min() * 0.9, 4)
+ x_max = round(sens_df.alpha.max() * 1.1, 4)
+ x_min = round(sens_df.alpha.min() * 0.9, 4)
+ plt.ylim(y_min, y_max)
+ plt.xlim(x_min, x_max)
+ ax.fill_between(
+ sens_df.alpha,
+ sens_df["New ATE LB"],
+ sens_df["New ATE UB"],
+ color="gray",
+ alpha=0.5,
+ )
+ ax.plot(sens_df.alpha, sens_df["New ATE"])
+ elif type == "r.squared" and ci:
+ fig, ax = plt.subplots()
+ y_max = round(sens_df["New ATE UB"].max() * 1.1, 4)
+ y_min = round(sens_df["New ATE LB"].min() * 0.9, 4)
+ plt.ylim(y_min, y_max)
+ ax.fill_between(
+ sens_df.rsqs,
+ sens_df["New ATE LB"],
+ sens_df["New ATE UB"],
+ color="gray",
+ alpha=0.5,
+ )
+ ax.plot(sens_df.rsqs, sens_df["New ATE"])
+ if partial_rsqs:
+ plt.scatter(
+ partial_rsqs_df.partial_rsqs,
+ list(sens_df[sens_df.alpha == 0]["New ATE"])
+ * partial_rsqs_df.shape[0],
+ marker="x",
+ color="red",
+ linewidth=10,
+ )
+ elif type == "r.squared" and not ci:
+ fig, ax = plt.subplots()
+ y_max = round(sens_df["New ATE UB"].max() * 1.1, 4)
+ y_min = round(sens_df["New ATE LB"].min() * 0.9, 4)
+ plt.ylim(y_min, y_max)
+ plt.plot(sens_df.rsqs, sens_df["New ATE"])
+ if partial_rsqs:
+ plt.scatter(
+ partial_rsqs_df.partial_rsqs,
+ list(sens_df[sens_df.alpha == 0]["New ATE"])
+ * partial_rsqs_df.shape[0],
+ marker="x",
+ color="red",
+ linewidth=10,
+ )
+
+ @staticmethod
+ def partial_rsqs_confounding(sens_df, feature_name, partial_rsqs_value, range=0.01):
+ """Check partial rsqs values of feature corresponding confounding amonunt of ATE
+ Args:
+ sens_df (pandas.DataFrame): a data frame output from causalsens
+ feature_name (str): feature name to check
+ partial_rsqs_value (float) : partial rsquare value of feature
+ range (float) : range to search from sens_df
+
+ Return: min and max value of confounding amount
+ """
+
+ rsqs_dict = []
+ for i in sens_df.rsqs:
+ if (
+ partial_rsqs_value - partial_rsqs_value * range
+ < i
+ < partial_rsqs_value + partial_rsqs_value * range
+ ):
+ rsqs_dict.append(i)
+
+ if rsqs_dict:
+ confounding_min = sens_df[sens_df.rsqs.isin(rsqs_dict)].alpha.min()
+ confounding_max = sens_df[sens_df.rsqs.isin(rsqs_dict)].alpha.max()
+ logger.info(
+ "Only works for linear outcome models right now. Check back soon."
+ )
+ logger.info(
+ "For feature {} with partial rsquare {} confounding amount with possible values: {}, {}".format(
+ feature_name, partial_rsqs_value, confounding_min, confounding_max
+ )
+ )
+ return [confounding_min, confounding_max]
+ else:
+ logger.info(
+ "Cannot find correponding rsquare value within the range for input, please edit confounding",
+ "values vector or use a larger range and try again",
+ )
diff --git a/causalml/source/causalml/metrics/visualize.py b/causalml/source/causalml/metrics/visualize.py
new file mode 100644
index 0000000000000000000000000000000000000000..10f00455554f04f67926186026df39e1ebee0726
--- /dev/null
+++ b/causalml/source/causalml/metrics/visualize.py
@@ -0,0 +1,1033 @@
+from typing import Optional
+from matplotlib import pyplot as plt
+import logging
+import numpy as np
+import pandas as pd
+import seaborn as sns
+from lightgbm import LGBMRegressor
+from ..inference.meta.tmle import TMLELearner
+
+plt.style.use("fivethirtyeight")
+sns.set_palette("Paired")
+RANDOM_COL = "Random"
+
+logger = logging.getLogger("causalml")
+
+
+def plot(
+ df,
+ kind="gain",
+ tmle=False,
+ n=100,
+ figsize=(8, 8),
+ ci=False,
+ plot_chance_level=True,
+ chance_level_kw=None,
+ ax: Optional[plt.Axes] = None,
+ *args,
+ **kwarg,
+) -> plt.Axes:
+ """Plot one of the lift/gain/Qini charts of model estimates.
+
+ A factory method for `plot_lift()`, `plot_gain()`, `plot_qini()`, `plot_tmlegain()` and `plot_tmleqini()`.
+ For details, pleas see docstrings of each function.
+
+ Args:
+ df (pandas.DataFrame): a data frame with model estimates and actual data as columns.
+ kind (str, optional): the kind of plot to draw. 'lift', 'gain', and 'qini' are supported.
+ n (int, optional): the number of samples to be used for plotting.
+ figsize (set of float, optional): the size of the figure to plot.
+ ci (bool, optional): whether to plot confidence intervals or not. Only available for `tmle=True`.
+ Default is False.
+ plot_chance_level (bool, optional): whether to plot the chance level (i.e., random) line or not.
+ Default is True.
+ chance_level_line_kw (dict, optional): the keyword arguments for the chance level line. Default is None.
+ *args: Variable length argument list.
+ **kwargs: Arbitrary keyword arguments.
+ """
+
+ if tmle:
+ catalog = {"gain": get_tmlegain, "qini": get_tmleqini}
+ else:
+ catalog = {"lift": get_cumlift, "gain": get_cumgain, "qini": get_qini}
+
+ assert (
+ kind in catalog.keys()
+ ), "{} plot is not implemented. Select one of {}".format(kind, catalog.keys())
+
+ if ax is None:
+ _, ax = plt.subplots(figsize=figsize)
+ if tmle:
+ df = catalog[kind](df, ci=ci, *args, **kwarg)
+
+ if ci:
+ model_names = [x.replace(" LB", "") for x in df.columns]
+ model_names = list(set([x.replace(" UB", "") for x in model_names]))
+
+ cmap = plt.get_cmap("tab10")
+ cindex = 0
+
+ for col in model_names:
+ lb_col = col + " LB"
+ up_col = col + " UB"
+
+ ax.plot(df.index, df[col], color=cmap(cindex))
+ ax.fill_between(
+ df.index,
+ df[lb_col],
+ df[up_col],
+ color=cmap(cindex),
+ alpha=0.25,
+ )
+ cindex += 1
+
+ ax.legend()
+ else:
+ ax = df.plot(ax=ax)
+
+ else:
+ df = catalog[kind](df, *args, **kwarg)
+
+ if (n is not None) and (n < df.shape[0]):
+ df = df.iloc[np.linspace(0, df.index[-1], n, endpoint=True).astype(int)]
+
+ ax = df.plot(ax=ax)
+
+ if plot_chance_level:
+ chance_level_line_kw = {
+ "label": RANDOM_COL,
+ "color": "k",
+ "linestyle": "--",
+ }
+
+ if chance_level_kw is not None:
+ chance_level_line_kw.update(**chance_level_kw)
+
+ ax.plot([0, df.index[-1]], [0, df.iloc[-1, 0]], **chance_level_line_kw)
+ ax.legend()
+
+ ax.set_xlabel("Population")
+ ax.set_ylabel("{}".format(kind.title()))
+ return ax
+
+
+def get_cumlift(
+ df, outcome_col="y", treatment_col="w", treatment_effect_col="tau", random_seed=42
+):
+ """Get average uplifts of model estimates in cumulative population.
+
+ If the true treatment effect is provided (e.g. in synthetic data), it's calculated
+ as the mean of the true treatment effect in each of cumulative population.
+ Otherwise, it's calculated as the difference between the mean outcomes of the
+ treatment and control groups in each of cumulative population.
+
+ For details, see Section 4.1 of Gutierrez and G{\'e}rardy (2016), `Causal Inference
+ and Uplift Modeling: A review of the literature`.
+
+ For the former, `treatment_effect_col` should be provided. For the latter, both
+ `outcome_col` and `treatment_col` should be provided.
+
+ Args:
+ df (pandas.DataFrame): a data frame with model estimates and actual data as columns
+ 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
+ random_seed (int, optional): deprecated
+
+ Returns:
+ (pandas.DataFrame): average uplifts of model estimates in cumulative population
+ """
+ assert (
+ (outcome_col in df.columns and df[outcome_col].notnull().all())
+ and (treatment_col in df.columns and df[treatment_col].notnull().all())
+ or (
+ treatment_effect_col in df.columns
+ and df[treatment_effect_col].notnull().all()
+ )
+ ), "{outcome_col} and {treatment_col}, or {treatment_effect_col} should be present without null.".format(
+ outcome_col=outcome_col,
+ treatment_col=treatment_col,
+ treatment_effect_col=treatment_effect_col,
+ )
+
+ df = df.copy()
+
+ model_names = [
+ x
+ for x in df.columns
+ if x not in [outcome_col, treatment_col, treatment_effect_col]
+ ]
+
+ lift = []
+ for i, col in enumerate(model_names):
+ sorted_df = df.sort_values(col, ascending=False).reset_index(drop=True)
+ sorted_df.index = sorted_df.index + 1
+
+ if treatment_effect_col in sorted_df.columns:
+ # When treatment_effect_col is given, use it to calculate the average treatment effects
+ # of cumulative population.
+ lift.append(sorted_df[treatment_effect_col].cumsum() / sorted_df.index)
+ else:
+ # When treatment_effect_col is not given, use outcome_col and treatment_col
+ # to calculate the average treatment_effects of cumulative population.
+ sorted_df["cumsum_tr"] = sorted_df[treatment_col].cumsum()
+ sorted_df["cumsum_ct"] = sorted_df.index.values - sorted_df["cumsum_tr"]
+ sorted_df["cumsum_y_tr"] = (
+ sorted_df[outcome_col] * sorted_df[treatment_col]
+ ).cumsum()
+ sorted_df["cumsum_y_ct"] = (
+ sorted_df[outcome_col] * (1 - sorted_df[treatment_col])
+ ).cumsum()
+
+ lift.append(
+ sorted_df["cumsum_y_tr"] / sorted_df["cumsum_tr"]
+ - sorted_df["cumsum_y_ct"] / sorted_df["cumsum_ct"]
+ )
+
+ lift = pd.concat(lift, join="inner", axis=1)
+ lift.loc[0] = np.zeros((lift.shape[1],))
+ lift = lift.sort_index().interpolate()
+
+ lift.columns = model_names
+
+ return lift
+
+
+def get_cumgain(
+ df,
+ outcome_col="y",
+ treatment_col="w",
+ treatment_effect_col="tau",
+ normalize=False,
+ random_seed=42,
+):
+ """Get cumulative gains of model estimates in population.
+
+ If the true treatment effect is provided (e.g. in synthetic data), it's calculated
+ as the cumulative gain of the true treatment effect in each population.
+ Otherwise, it's calculated as the cumulative difference between the mean outcomes
+ of the treatment and control groups in each population.
+
+ For details, see Section 4.1 of Gutierrez and G{\'e}rardy (2016), `Causal Inference
+ and Uplift Modeling: A review of the literature`.
+
+ For the former, `treatment_effect_col` should be provided. For the latter, both
+ `outcome_col` and `treatment_col` should be provided.
+
+ Args:
+ df (pandas.DataFrame): a data frame with model estimates and actual data as columns
+ 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
+ normalize (bool, optional): whether to normalize the y-axis to 1 or not
+ random_seed (int, optional): deprecated
+
+ Returns:
+ (pandas.DataFrame): cumulative gains of model estimates in population
+ """
+
+ lift = get_cumlift(df, outcome_col, treatment_col, treatment_effect_col)
+
+ # cumulative gain = cumulative lift x (# of population)
+ gain = lift.mul(lift.index.values, axis=0)
+
+ if normalize:
+ gain = gain.div(np.abs(gain.iloc[-1, :]), axis=1)
+
+ return gain
+
+
+def get_qini(
+ df,
+ outcome_col="y",
+ treatment_col="w",
+ treatment_effect_col="tau",
+ normalize=False,
+ random_seed=42,
+):
+ """Get Qini of model estimates in population.
+
+ If the true treatment effect is provided (e.g. in synthetic data), it's calculated
+ as the cumulative gain of the true treatment effect in each population.
+ Otherwise, it's calculated as the cumulative difference between the mean outcomes
+ of the treatment and control groups in each population.
+
+ For details, see Radcliffe (2007), `Using Control Group to Target on Predicted Lift:
+ Building and Assessing Uplift Models`
+
+ For the former, `treatment_effect_col` should be provided. For the latter, both
+ `outcome_col` and `treatment_col` should be provided.
+
+ Args:
+ df (pandas.DataFrame): a data frame with model estimates and actual data as columns
+ 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
+ normalize (bool, optional): whether to normalize the y-axis to 1 or not
+ random_seed (int, optional): deprecated
+
+ Returns:
+ (pandas.DataFrame): cumulative gains of model estimates in population
+ """
+ assert (
+ (outcome_col in df.columns and df[outcome_col].notnull().all())
+ and (treatment_col in df.columns and df[treatment_col].notnull().all())
+ or (
+ treatment_effect_col in df.columns
+ and df[treatment_effect_col].notnull().all()
+ )
+ ), "{outcome_col} and {treatment_col}, or {treatment_effect_col} should be present without null.".format(
+ outcome_col=outcome_col,
+ treatment_col=treatment_col,
+ treatment_effect_col=treatment_effect_col,
+ )
+
+ df = df.copy()
+
+ model_names = [
+ x
+ for x in df.columns
+ if x not in [outcome_col, treatment_col, treatment_effect_col]
+ ]
+
+ qini = []
+ for i, col in enumerate(model_names):
+ sorted_df = df.sort_values(col, ascending=False).reset_index(drop=True)
+ sorted_df.index = sorted_df.index + 1
+ sorted_df["cumsum_tr"] = sorted_df[treatment_col].cumsum()
+
+ if treatment_effect_col in sorted_df.columns:
+ # When treatment_effect_col is given, use it to calculate the average treatment effects
+ # of cumulative population.
+ l = (
+ sorted_df[treatment_effect_col].cumsum()
+ / sorted_df.index
+ * sorted_df["cumsum_tr"]
+ )
+ else:
+ # When treatment_effect_col is not given, use outcome_col and treatment_col
+ # to calculate the average treatment_effects of cumulative population.
+ sorted_df["cumsum_ct"] = sorted_df.index.values - sorted_df["cumsum_tr"]
+ sorted_df["cumsum_y_tr"] = (
+ sorted_df[outcome_col] * sorted_df[treatment_col]
+ ).cumsum()
+ sorted_df["cumsum_y_ct"] = (
+ sorted_df[outcome_col] * (1 - sorted_df[treatment_col])
+ ).cumsum()
+
+ l = (
+ sorted_df["cumsum_y_tr"]
+ - sorted_df["cumsum_y_ct"]
+ * sorted_df["cumsum_tr"]
+ / sorted_df["cumsum_ct"]
+ )
+
+ qini.append(l)
+
+ qini = pd.concat(qini, join="inner", axis=1)
+ qini.loc[0] = np.zeros((qini.shape[1],))
+ qini = qini.sort_index().interpolate()
+
+ qini.columns = model_names
+
+ if normalize:
+ qini = qini.div(np.abs(qini.iloc[-1, :]), axis=1)
+
+ return qini
+
+
+def get_tmlegain(
+ df,
+ inference_col,
+ learner=LGBMRegressor(num_leaves=64, learning_rate=0.05, n_estimators=300),
+ outcome_col="y",
+ treatment_col="w",
+ p_col="p",
+ n_segment=5,
+ cv=None,
+ ci=False,
+):
+ """Get TMLE based average uplifts of model estimates of segments.
+
+ Args:
+ df (pandas.DataFrame): a data frame with model estimates and actual data as columns
+ inferenece_col (list of str): a list of columns that used in learner for inference
+ learner (optional): a model used by TMLE to estimate the outcome
+ 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)
+ p_col (str, optional): the column name for propensity score
+ n_segment (int, optional): number of segment that TMLE will estimated for each
+ cv (sklearn.model_selection._BaseKFold, optional): sklearn CV object
+ ci (bool, optional): whether return confidence intervals for ATE or not
+ Returns:
+ (pandas.DataFrame): cumulative gains of model estimates based of TMLE
+ """
+ assert (
+ (outcome_col in df.columns and df[outcome_col].notnull().all())
+ and (treatment_col in df.columns and df[treatment_col].notnull().all())
+ or (p_col in df.columns and df[p_col].notnull().all())
+ ), "{outcome_col} and {treatment_col}, or {p_col} should be present without null.".format(
+ outcome_col=outcome_col,
+ treatment_col=treatment_col,
+ p_col=p_col,
+ )
+
+ inference_col = [x for x in inference_col if x in df.columns]
+
+ # Initialize TMLE
+ tmle = TMLELearner(learner, cv=cv)
+ ate_all, ate_all_lb, ate_all_ub = tmle.estimate_ate(
+ X=df[inference_col], p=df[p_col], treatment=df[treatment_col], y=df[outcome_col]
+ )
+
+ df = df.copy()
+ model_names = [
+ x
+ for x in df.columns
+ if x not in [outcome_col, treatment_col, p_col] + inference_col
+ ]
+
+ lift = []
+ lift_lb = []
+ lift_ub = []
+
+ for col in model_names:
+ # Create `n_segment` equal segments from sorted model estimates. Rank is used to break ties.
+ # ref: https://stackoverflow.com/a/46979206/3216742
+ segments = pd.qcut(df[col].rank(method="first"), n_segment, labels=False)
+
+ ate_model, ate_model_lb, ate_model_ub = tmle.estimate_ate(
+ X=df[inference_col],
+ p=df[p_col],
+ treatment=df[treatment_col],
+ y=df[outcome_col],
+ segment=segments,
+ )
+ lift_model = [0.0] * (n_segment + 1)
+ lift_model[n_segment] = ate_all[0]
+ for i in range(1, n_segment):
+ lift_model[i] = (
+ ate_model[0][n_segment - i] * (1 / n_segment) + lift_model[i - 1]
+ )
+ lift.append(lift_model)
+
+ if ci:
+ lift_lb_model = [0.0] * (n_segment + 1)
+ lift_lb_model[n_segment] = ate_all_lb[0]
+
+ lift_ub_model = [0.0] * (n_segment + 1)
+ lift_ub_model[n_segment] = ate_all_ub[0]
+ for i in range(1, n_segment):
+ lift_lb_model[i] = (
+ ate_model_lb[0][n_segment - i] * (1 / n_segment)
+ + lift_lb_model[i - 1]
+ )
+ lift_ub_model[i] = (
+ ate_model_ub[0][n_segment - i] * (1 / n_segment)
+ + lift_ub_model[i - 1]
+ )
+
+ lift_lb.append(lift_lb_model)
+ lift_ub.append(lift_ub_model)
+
+ lift = pd.DataFrame(lift).T
+ lift.columns = model_names
+
+ if ci:
+ lift_lb = pd.DataFrame(lift_lb).T
+ lift_lb.columns = [x + " LB" for x in model_names]
+
+ lift_ub = pd.DataFrame(lift_ub).T
+ lift_ub.columns = [x + " UB" for x in model_names]
+ lift = pd.concat([lift, lift_lb, lift_ub], axis=1)
+
+ lift.index = lift.index / n_segment
+
+ return lift
+
+
+def get_tmleqini(
+ df,
+ inference_col,
+ learner=LGBMRegressor(num_leaves=64, learning_rate=0.05, n_estimators=300),
+ outcome_col="y",
+ treatment_col="w",
+ p_col="p",
+ n_segment=5,
+ cv=None,
+ ci=False,
+ normalize=False,
+):
+ """Get TMLE based Qini of model estimates by segments.
+
+ Args:
+ df (pandas.DataFrame): a data frame with model estimates and actual data as columns
+ inferenece_col (list of str): a list of columns that used in learner for inference
+ learner(optional): a model used by TMLE to estimate the outcome
+ 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)
+ p_col (str, optional): the column name for propensity score
+ n_segment (int, optional): number of segment that TMLE will estimated for each
+ cv (sklearn.model_selection._BaseKFold, optional): sklearn CV object
+ ci (bool, optional): whether return confidence intervals for ATE or not
+ Returns:
+ (pandas.DataFrame): cumulative gains of model estimates based of TMLE
+ """
+ assert (
+ (outcome_col in df.columns and df[outcome_col].notnull().all())
+ and (treatment_col in df.columns and df[treatment_col].notnull().all())
+ or (p_col in df.columns and df[p_col].notnull().all())
+ ), "{outcome_col} and {treatment_col}, or {p_col} should be present without null.".format(
+ outcome_col=outcome_col,
+ treatment_col=treatment_col,
+ p_col=p_col,
+ )
+
+ inference_col = [x for x in inference_col if x in df.columns]
+
+ # Initialize TMLE
+ tmle = TMLELearner(learner, cv=cv)
+ ate_all, ate_all_lb, ate_all_ub = tmle.estimate_ate(
+ X=df[inference_col], p=df[p_col], treatment=df[treatment_col], y=df[outcome_col]
+ )
+
+ df = df.copy()
+ model_names = [
+ x
+ for x in df.columns
+ if x not in [outcome_col, treatment_col, p_col] + inference_col
+ ]
+
+ qini = []
+ qini_lb = []
+ qini_ub = []
+
+ for col in model_names:
+ # Create `n_segment` equal segments from sorted model estimates. Rank is used to break ties.
+ # ref: https://stackoverflow.com/a/46979206/3216742
+ segments = pd.qcut(df[col].rank(method="first"), n_segment, labels=False)
+
+ ate_model, ate_model_lb, ate_model_ub = tmle.estimate_ate(
+ X=df[inference_col],
+ p=df[p_col],
+ treatment=df[treatment_col],
+ y=df[outcome_col],
+ segment=segments,
+ )
+
+ qini_model = [0]
+ for i in range(1, n_segment):
+ n_tr = df[segments == (n_segment - i)][treatment_col].sum()
+ qini_model.append(ate_model[0][n_segment - i] * n_tr)
+
+ qini.append(qini_model)
+
+ if ci:
+ qini_lb_model = [0]
+ qini_ub_model = [0]
+ for i in range(1, n_segment):
+ n_tr = df[segments == (n_segment - i)][treatment_col].sum()
+ qini_lb_model.append(ate_model_lb[0][n_segment - i] * n_tr)
+ qini_ub_model.append(ate_model_ub[0][n_segment - i] * n_tr)
+
+ qini_lb.append(qini_lb_model)
+ qini_ub.append(qini_ub_model)
+
+ qini = pd.DataFrame(qini).T
+ qini.columns = model_names
+
+ if ci:
+ qini_lb = pd.DataFrame(qini_lb).T
+ qini_lb.columns = [x + " LB" for x in model_names]
+
+ qini_ub = pd.DataFrame(qini_ub).T
+ qini_ub.columns = [x + " UB" for x in model_names]
+ qini = pd.concat([qini, qini_lb, qini_ub], axis=1)
+
+ qini = qini.cumsum()
+ qini.loc[n_segment] = ate_all[0] * df[treatment_col].sum()
+ qini.index = np.linspace(0, 1, n_segment + 1) * df.shape[0]
+
+ return qini
+
+
+def plot_gain(
+ df,
+ outcome_col="y",
+ treatment_col="w",
+ treatment_effect_col="tau",
+ normalize=False,
+ random_seed=42,
+ n=100,
+ figsize=(8, 8),
+ ax: Optional[plt.Axes] = None,
+):
+ """Plot the cumulative gain chart (or uplift curve) of model estimates.
+
+ If the true treatment effect is provided (e.g. in synthetic data), it's calculated
+ as the cumulative gain of the true treatment effect in each population.
+ Otherwise, it's calculated as the cumulative difference between the mean outcomes
+ of the treatment and control groups in each population.
+
+ For details, see Section 4.1 of Gutierrez and G{\'e}rardy (2016), `Causal Inference
+ and Uplift Modeling: A review of the literature`.
+
+ For the former, `treatment_effect_col` should be provided. For the latter, both
+ `outcome_col` and `treatment_col` should be provided.
+
+ Args:
+ df (pandas.DataFrame): a data frame with model estimates and actual data as columns
+ 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
+ normalize (bool, optional): whether to normalize the y-axis to 1 or not
+ random_seed (int, optional): random seed for numpy.random.rand()
+ n (int, optional): the number of samples to be used for plotting
+ """
+
+ plot(
+ df,
+ kind="gain",
+ n=n,
+ figsize=figsize,
+ outcome_col=outcome_col,
+ treatment_col=treatment_col,
+ treatment_effect_col=treatment_effect_col,
+ normalize=normalize,
+ ax=ax,
+ )
+
+
+def plot_lift(
+ df,
+ outcome_col="y",
+ treatment_col="w",
+ treatment_effect_col="tau",
+ random_seed=42,
+ n=100,
+ figsize=(8, 8),
+):
+ """Plot the lift chart of model estimates in cumulative population.
+
+ If the true treatment effect is provided (e.g. in synthetic data), it's calculated
+ as the mean of the true treatment effect in each of cumulative population.
+ Otherwise, it's calculated as the difference between the mean outcomes of the
+ treatment and control groups in each of cumulative population.
+
+ For details, see Section 4.1 of Gutierrez and G{\'e}rardy (2016), `Causal Inference
+ and Uplift Modeling: A review of the literature`.
+
+ For the former, `treatment_effect_col` should be provided. For the latter, both
+ `outcome_col` and `treatment_col` should be provided.
+
+ Args:
+ df (pandas.DataFrame): a data frame with model estimates and actual data as columns
+ 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
+ random_seed (int, optional): deprecated
+ n (int, optional): the number of samples to be used for plotting
+ """
+
+ plot(
+ df,
+ kind="lift",
+ n=n,
+ figsize=figsize,
+ outcome_col=outcome_col,
+ treatment_col=treatment_col,
+ treatment_effect_col=treatment_effect_col,
+ )
+
+
+def plot_qini(
+ df,
+ outcome_col="y",
+ treatment_col="w",
+ treatment_effect_col="tau",
+ normalize=False,
+ random_seed=42,
+ n=100,
+ figsize=(8, 8),
+ ax: Optional[plt.Axes] = None,
+) -> plt.Axes:
+ """Plot the Qini chart (or uplift curve) of model estimates.
+
+ If the true treatment effect is provided (e.g. in synthetic data), it's calculated
+ as the cumulative gain of the true treatment effect in each population.
+ Otherwise, it's calculated as the cumulative difference between the mean outcomes
+ of the treatment and control groups in each population.
+
+ For details, see Radcliffe (2007), `Using Control Group to Target on Predicted Lift:
+ Building and Assessing Uplift Models`
+
+ For the former, `treatment_effect_col` should be provided. For the latter, both
+ `outcome_col` and `treatment_col` should be provided.
+
+ Args:
+ df (pandas.DataFrame): a data frame with model estimates and actual data as columns
+ 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
+ normalize (bool, optional): whether to normalize the y-axis to 1 or not
+ random_seed (int, optional): deprecated
+ n (int, optional): the number of samples to be used for plotting
+ ci (bool, optional): whether return confidence intervals for ATE or not
+ """
+
+ ax = plot(
+ df,
+ kind="qini",
+ n=n,
+ figsize=figsize,
+ outcome_col=outcome_col,
+ treatment_col=treatment_col,
+ treatment_effect_col=treatment_effect_col,
+ normalize=normalize,
+ ax=ax,
+ )
+ return ax
+
+
+def plot_tmlegain(
+ df,
+ inference_col,
+ learner=LGBMRegressor(
+ num_leaves=64, learning_rate=0.05, n_estimators=300, verbose=-1
+ ),
+ outcome_col="y",
+ treatment_col="w",
+ p_col="tau",
+ n_segment=5,
+ cv=None,
+ ci=False,
+ figsize=(8, 8),
+):
+ """Plot the lift chart based of TMLE estimation
+
+ Args:
+ df (pandas.DataFrame): a data frame with model estimates and actual data as columns
+ inferenece_col (list of str): a list of columns that used in learner for inference
+ learner (optional): a model used by TMLE to estimate the outcome
+ 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)
+ p_col (str, optional): the column name for propensity score
+ n_segment (int, optional): number of segment that TMLE will estimated for each
+ cv (sklearn.model_selection._BaseKFold, optional): sklearn CV object
+ ci (bool, optional): whether return confidence intervals for ATE or not
+ """
+
+ plot(
+ df,
+ kind="gain",
+ tmle=True,
+ figsize=figsize,
+ ci=ci,
+ learner=learner,
+ inference_col=inference_col,
+ outcome_col=outcome_col,
+ treatment_col=treatment_col,
+ p_col=p_col,
+ n_segment=n_segment,
+ cv=cv,
+ )
+
+
+def plot_tmleqini(
+ df,
+ inference_col,
+ learner=LGBMRegressor(num_leaves=64, learning_rate=0.05, n_estimators=300),
+ outcome_col="y",
+ treatment_col="w",
+ p_col="tau",
+ n_segment=5,
+ cv=None,
+ ci=False,
+ figsize=(8, 8),
+):
+ """Plot the qini chart based of TMLE estimation
+
+ Args:
+ df (pandas.DataFrame): a data frame with model estimates and actual data as columns
+ inferenece_col (list of str): a list of columns that used in learner for inference
+ learner (optional): a model used by TMLE to estimate the outcome
+ 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)
+ p_col (str, optional): the column name for propensity score
+ n_segment (int, optional): number of segment that TMLE will estimated for each
+ cv (sklearn.model_selection._BaseKFold, optional): sklearn CV object
+ ci (bool, optional): whether return confidence intervals for ATE or not
+ """
+
+ plot(
+ df,
+ kind="qini",
+ tmle=True,
+ figsize=figsize,
+ ci=ci,
+ learner=learner,
+ inference_col=inference_col,
+ outcome_col=outcome_col,
+ treatment_col=treatment_col,
+ p_col=p_col,
+ n_segment=n_segment,
+ cv=cv,
+ )
+
+
+def auuc_score(
+ df,
+ outcome_col="y",
+ treatment_col="w",
+ treatment_effect_col="tau",
+ normalize=True,
+ tmle=False,
+ *args,
+ **kwarg,
+):
+ """Calculate the AUUC (Area Under the Uplift Curve) score.
+
+ Args:
+ df (pandas.DataFrame): a data frame with model estimates and actual data as columns
+ 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
+ normalize (bool, optional): whether to normalize the y-axis to 1 or not
+
+ Returns:
+ (float): the AUUC score
+ """
+
+ if not tmle:
+ cumgain = get_cumgain(
+ df, outcome_col, treatment_col, treatment_effect_col, normalize
+ )
+ else:
+ cumgain = get_tmlegain(
+ df, outcome_col=outcome_col, treatment_col=treatment_col, *args, **kwarg
+ )
+ return cumgain.sum() / cumgain.shape[0]
+
+
+def qini_score(
+ df,
+ outcome_col="y",
+ treatment_col="w",
+ treatment_effect_col="tau",
+ normalize=True,
+ tmle=False,
+ *args,
+ **kwarg,
+):
+ """Calculate the Qini score: the area between the Qini curves of a model and random.
+
+ For details, see Radcliffe (2007), `Using Control Group to Target on Predicted Lift:
+ Building and Assessing Uplift Models`
+
+ Args:
+ df (pandas.DataFrame): a data frame with model estimates and actual data as columns
+ 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
+ normalize (bool, optional): whether to normalize the y-axis to 1 or not
+
+ Returns:
+ (float): the Qini score
+ """
+
+ if not tmle:
+ qini = get_qini(df, outcome_col, treatment_col, treatment_effect_col, normalize)
+ else:
+ qini = get_tmleqini(
+ df, outcome_col=outcome_col, treatment_col=treatment_col, *args, **kwarg
+ )
+
+ random_area = np.linspace(qini.iloc[0, 0], qini.iloc[-1, 0], qini.shape[0]).sum()
+ return (qini.sum(axis=0) - random_area) / qini.shape[0]
+
+
+def plot_ps_diagnostics(df, covariate_col, treatment_col="w", p_col="p", bal_tol=0.1):
+ """Plot covariate balances (standardized differences between the treatment and the control)
+ before and after weighting the sample using the inverse probability of treatment weights.
+
+ Args:
+ df (pandas.DataFrame): a data frame containing the covariates and treatment indicator
+ covariate_col (list of str): a list of columns that are used a covariates
+ treatment_col (str, optional): the column name for the treatment indicator (0 or 1)
+ p_col (str, optional): the column name for propensity score
+ """
+ X = df[covariate_col]
+ W = df[treatment_col]
+ PS = df[p_col]
+
+ IPTW = get_simple_iptw(W, PS)
+
+ diffs_pre = get_std_diffs(X, W, weighted=False)
+ num_unbal_pre = (np.abs(diffs_pre) > bal_tol).sum()[0]
+
+ diffs_post = get_std_diffs(X, W, IPTW, weighted=True)
+ num_unbal_post = (np.abs(diffs_post) > bal_tol).sum()[0]
+
+ diff_plot = _plot_std_diffs(
+ diffs_pre, num_unbal_pre, diffs_post, num_unbal_post, bal_tol=bal_tol
+ )
+
+ return diff_plot
+
+
+def _plot_std_diffs(diffs_pre, num_unbal_pre, diffs_post, num_unbal_post, bal_tol=0.1):
+ fig, ax1 = plt.subplots()
+
+ color = "#EA2566"
+
+ sds_pre = pd.DataFrame(
+ {"std_diff": diffs_pre[0], "covariate": diffs_pre.index, "prepost": "pre"}
+ )
+ sds_post = pd.DataFrame(
+ {"std_diff": diffs_post[0], "covariate": diffs_post.index, "prepost": "post"}
+ )
+
+ sds = pd.concat([sds_pre, sds_post], ignore_index=True)
+
+ sns.stripplot(data=sds, x="std_diff", y="covariate", hue="prepost", ax=ax1)
+
+ ax1.set_xlabel(
+ "Pre/Post Number of unbalanced covariates: {num_unbal_pre}/{num_unbal_post}".format(
+ num_unbal_pre=num_unbal_pre, num_unbal_post=num_unbal_post
+ ),
+ fontsize=14,
+ )
+ ax1.axvline(x=-bal_tol, ymin=0, ymax=1, color=color, linestyle="--", lw=2)
+ ax1.axvline(x=bal_tol, ymin=0, ymax=1, color=color, linestyle="--", lw=2)
+
+ fig.suptitle("Standardized differences in means", fontsize=16)
+
+ return fig
+
+
+def get_simple_iptw(W, propensity_score):
+ IPTW = (W / propensity_score) + (1 - W) / (1 - propensity_score)
+
+ return IPTW
+
+
+def get_std_diffs(X, W, weight=None, weighted=False, numeric_threshold=5):
+ """Calculate the inverse probability of treatment weighted standardized
+ differences in covariate means between the treatment and the control.
+ If weighting is set to 'False', calculate unweighted standardized
+ differences. Accepts only continuous and binary numerical variables.
+ """
+ cont_cols, prop_cols = _get_numeric_vars(X, threshold=numeric_threshold)
+ cols = cont_cols + prop_cols
+
+ if len(cols) == 0:
+ raise ValueError(
+ "No variable passed the test for continuous or binary variables."
+ )
+
+ treat = W == 1
+ contr = W == 0
+
+ X_1 = X.loc[treat, cols]
+ X_0 = X.loc[contr, cols]
+
+ cont_index = np.array([col in cont_cols for col in cols])
+ prop_index = np.array([col in prop_cols for col in cols])
+
+ std_diffs_cont = np.empty(sum(cont_index))
+ std_diffs_prop = np.empty(sum(prop_index))
+
+ if weighted:
+ assert (
+ weight is not None
+ ), 'weight should be provided when weighting is set to "True"'
+
+ weight_1 = weight[treat]
+ weight_0 = weight[contr]
+
+ X_1_mean, X_1_var = np.apply_along_axis(
+ lambda x: _get_wmean_wvar(x, weight_1), 0, X_1
+ )
+ X_0_mean, X_0_var = np.apply_along_axis(
+ lambda x: _get_wmean_wvar(x, weight_0), 0, X_0
+ )
+
+ elif not weighted:
+ X_1_mean, X_1_var = np.apply_along_axis(lambda x: _get_mean_var(x), 0, X_1)
+ X_0_mean, X_0_var = np.apply_along_axis(lambda x: _get_mean_var(x), 0, X_0)
+
+ X_1_mean_cont, X_1_var_cont = X_1_mean[cont_index], X_1_var[cont_index]
+ X_0_mean_cont, X_0_var_cont = X_0_mean[cont_index], X_0_var[cont_index]
+
+ std_diffs_cont = (X_1_mean_cont - X_0_mean_cont) / np.sqrt(
+ (X_1_var_cont + X_0_var_cont) / 2
+ )
+
+ X_1_mean_prop = X_1_mean[prop_index]
+ X_0_mean_prop = X_0_mean[prop_index]
+
+ std_diffs_prop = (X_1_mean_prop - X_0_mean_prop) / np.sqrt(
+ ((X_1_mean_prop * (1 - X_1_mean_prop)) + (X_0_mean_prop * (1 - X_0_mean_prop)))
+ / 2
+ )
+
+ std_diffs = np.concatenate([std_diffs_cont, std_diffs_prop], axis=0)
+ std_diffs_df = pd.DataFrame(std_diffs, index=cols)
+
+ return std_diffs_df
+
+
+def _get_numeric_vars(X, threshold=5):
+ """Attempt to determine which variables are numeric and which
+ are categorical. The threshold for a 'continuous' variable
+ is set to 5 by default.
+ """
+
+ cont = [
+ (not hasattr(X.iloc[:, i], "cat")) and (X.iloc[:, i].nunique() >= threshold)
+ for i in range(X.shape[1])
+ ]
+
+ prop = [X.iloc[:, i].nunique() == 2 for i in range(X.shape[1])]
+
+ cont_cols = list(X.loc[:, cont].columns)
+ prop_cols = list(X.loc[:, prop].columns)
+
+ dropped = set(X.columns) - set(cont_cols + prop_cols)
+
+ if dropped:
+ logger.info(
+ 'Some non-binary variables were dropped because they had fewer than {} unique values or were of the \
+ dtype "cat". The dropped variables are: {}'.format(
+ threshold, dropped
+ )
+ )
+
+ return cont_cols, prop_cols
+
+
+def _get_mean_var(X):
+ """Calculate the mean and variance of a variable."""
+ mean = X.mean()
+ var = X.var()
+
+ return [mean, var]
+
+
+def _get_wmean_wvar(X, weight):
+ """
+ Calculate the weighted mean of a variable given an arbitrary
+ sample weight. Formulas from:
+
+ Austin, Peter C., and Elizabeth A. Stuart. 2015. Moving towards Best
+ Practice When Using Inverse Probability of Treatment Weighting (IPTW)
+ Using the Propensity Score to Estimate Causal Treatment Effects in
+ Observational Studies.
+ Statistics in Medicine 34 (28): 3661 79. https://doi.org/10.1002/sim.6607.
+ """
+ weighted_mean = np.sum(weight * X) / np.sum(weight)
+ weighted_var = (
+ np.sum(weight) / (np.power(np.sum(weight), 2) - np.sum(np.power(weight, 2)))
+ ) * (np.sum(weight * np.power((X - weighted_mean), 2)))
+
+ return [weighted_mean, weighted_var]
diff --git a/causalml/source/causalml/optimize/__init__.py b/causalml/source/causalml/optimize/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..6379fc0cec0449ce73c33681889d9c61766fb719
--- /dev/null
+++ b/causalml/source/causalml/optimize/__init__.py
@@ -0,0 +1,5 @@
+from .policylearner import PolicyLearner
+from .unit_selection import CounterfactualUnitSelector
+from .utils import get_treatment_costs, get_actual_value, get_uplift_best
+from .value_optimization import CounterfactualValueEstimator
+from .pns import get_pns_bounds
diff --git a/causalml/source/causalml/optimize/pns.py b/causalml/source/causalml/optimize/pns.py
new file mode 100644
index 0000000000000000000000000000000000000000..6e62ce46b14e0afe5473fa853f695bead89d05c1
--- /dev/null
+++ b/causalml/source/causalml/optimize/pns.py
@@ -0,0 +1,75 @@
+def get_pns_bounds(data_exp, data_obs, T, Y, type="PNS"):
+ """
+ Args
+ ----
+ data_exp : DataFrame
+ Data from an experiment.
+ data_obs : DataFrame
+ Data from an observational study
+ T : str
+ Name of the binary treatment indicator
+ y : str
+ Name of the binary outcome indicator
+ type : str
+ Type of probability of causation desired. Acceptable args are:
+ - ``PNS``: Probability of necessary and sufficient causation
+ - ``PS``: Probability of sufficient causation
+ - ``PN``: Probability of necessary causation
+
+ Notes
+ -----
+ Based on Equation (24) in `Tian and Pearl (2000) `_.
+
+ To capture the counterfactual notation, we use ``1`` and ``0`` to indicate the actual and
+ counterfactual values of a variable, respectively, and we use ``do`` to indicate the effect
+ of an intervention.
+
+ The experimental and observational data are either assumed to come to the same population,
+ or from random samples of the population. If the data are from a sample, the bounds may
+ be incorrectly calculated because the relevant quantities in the Tian-Pearl equations are
+ defined e.g. as :math:`P(Y|do(T))`, not :math:`P(Y|do(T), S)` where :math:`S` corresponds to sample selection.
+ `Bareinboim and Pearl (2016) `_ discuss conditions
+ under which :math:`P(Y|do(T))` can be recovered from :math:`P(Y|do(T), S)`.
+ """
+
+ # Probabilities calculated from observational data
+ Y1 = data_obs[Y].mean()
+ T1Y0 = (
+ data_obs.loc[(data_obs[T] == 1) & (data_obs[Y] == 0)].shape[0]
+ / data_obs.shape[0]
+ )
+ T1Y1 = (
+ data_obs.loc[(data_obs[T] == 1) & (data_obs[Y] == 1)].shape[0]
+ / data_obs.shape[0]
+ )
+ T0Y0 = (
+ data_obs.loc[(data_obs[T] == 0) & (data_obs[Y] == 0)].shape[0]
+ / data_obs.shape[0]
+ )
+ T0Y1 = (
+ data_obs.loc[(data_obs[T] == 0) & (data_obs[Y] == 1)].shape[0]
+ / data_obs.shape[0]
+ )
+
+ # Probabilities calculated from experimental data
+ Y1doT1 = data_exp.loc[data_exp[T] == 1, Y].mean()
+ Y1doT0 = data_exp.loc[data_exp[T] == 0, Y].mean()
+ Y0doT0 = 1 - Y1doT0
+
+ if type == "PNS":
+ lb_args = [0, Y1doT1 - Y1doT0, Y1 - Y1doT0, Y1doT1 - Y1]
+
+ ub_args = [Y1doT1, Y0doT0, T1Y1 + T0Y0, Y1doT1 - Y1doT0 + T1Y0 + T0Y1]
+
+ if type == "PN":
+ lb_args = [0, (Y1 - Y1doT0) / T1Y1]
+ ub_args = [1, (Y0doT0 - T0Y0) / T1Y1]
+
+ if type == "PS":
+ lb_args = [0, (Y1doT1 - Y1) / T0Y0]
+ ub_args = [1, (Y1doT1 - T1Y1) / T0Y0]
+
+ lower_bound = max(lb_args)
+ upper_bound = min(ub_args)
+
+ return lower_bound, upper_bound
diff --git a/causalml/source/causalml/optimize/policylearner.py b/causalml/source/causalml/optimize/policylearner.py
new file mode 100644
index 0000000000000000000000000000000000000000..1dea008a0869adb861a0d85d089875d4c017107a
--- /dev/null
+++ b/causalml/source/causalml/optimize/policylearner.py
@@ -0,0 +1,172 @@
+import logging
+
+import numpy as np
+from causalml.propensity import compute_propensity_score
+from sklearn.ensemble import GradientBoostingRegressor, GradientBoostingClassifier
+from sklearn.model_selection import KFold
+from sklearn.tree import DecisionTreeClassifier
+
+logger = logging.getLogger("causalml")
+
+
+class PolicyLearner:
+ """
+ A Learner that learns a treatment assignment policy with observational data using doubly robust estimator of causal
+ effect for binary treatment.
+
+ Details of the policy learner are available at `Athey and Wager (2018) `_.
+
+ """
+
+ def __init__(
+ self,
+ outcome_learner=GradientBoostingRegressor(),
+ treatment_learner=GradientBoostingClassifier(),
+ policy_learner=DecisionTreeClassifier(),
+ clip_bounds=(1e-3, 1 - 1e-3),
+ n_fold=5,
+ random_state=None,
+ calibration=False,
+ ):
+ """Initialize a treatment assignment policy learner.
+
+ Args:
+ outcome_learner (optional): a regression model to estimate outcomes
+ policy_learner (optional): a classification model to estimate treatment assignment. It needs to take
+ `sample_weight` as an input argument for `fit()`
+ clip_bounds (tuple, optional): lower and upper bounds for clipping propensity scores to avoid division by
+ zero in PolicyLearner.fit()
+ 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)
+ """
+ self.model_mu = outcome_learner
+ self.model_w = treatment_learner
+ self.model_pi = policy_learner
+ self.clip_bounds = clip_bounds
+ self.cv = KFold(n_splits=n_fold, shuffle=True, random_state=random_state)
+ self.calibration = calibration
+
+ self._y_pred, self._tau_pred, self._w_pred, self._dr_score = (
+ None,
+ None,
+ None,
+ None,
+ )
+
+ def __repr__(self):
+ return (
+ "{}(model_mu={},\n"
+ "\tmodel_w={},\n"
+ "\tmodel_pi={})".format(
+ self.__class__.__name__,
+ self.model_mu.__repr__(),
+ self.model_w.__repr__(),
+ self.model_pi.__repr__(),
+ )
+ )
+
+ def _outcome_estimate(self, X, w, y):
+ self._y_pred = np.zeros(len(y))
+ self._tau_pred = np.zeros(len(y))
+
+ for train_index, test_index in self.cv.split(y):
+ X_train, X_test = X[train_index], X[test_index]
+ w_train, w_test = w[train_index], w[test_index]
+ y_train, _ = y[train_index], y[test_index]
+
+ self.model_mu.fit(
+ np.concatenate([X_train, w_train.reshape(-1, 1)], axis=1), y_train
+ )
+ self._y_pred[test_index] = self.model_mu.predict(
+ np.concatenate([X_test, w_test.reshape(-1, 1)], axis=1)
+ )
+ self._tau_pred[test_index] = self.model_mu.predict(
+ np.concatenate([X_test, np.ones((len(w_test), 1))], axis=1)
+ ) - self.model_mu.predict(
+ np.concatenate([X_test, np.zeros((len(w_test), 1))], axis=1)
+ )
+
+ def _treatment_estimate(self, X, w):
+ self._w_pred = np.zeros(len(w))
+
+ for train_index, test_index in self.cv.split(w):
+ X_train, X_test = X[train_index], X[test_index]
+ w_train, w_test = w[train_index], w[test_index]
+
+ self._w_pred[test_index], _ = compute_propensity_score(
+ X=X_train,
+ treatment=w_train,
+ X_pred=X_test,
+ treatment_pred=w_test,
+ calibrate_p=self.calibration,
+ )
+
+ self._w_pred = np.clip(
+ self._w_pred, a_min=self.clip_bounds[0], a_max=self.clip_bounds[1]
+ )
+
+ def fit(self, X, treatment, y, p=None, dhat=None):
+ """Fit the treatment assignment policy learner.
+
+ Args:
+ X (np.matrix): a feature matrix
+ treatment (np.array): a treatment vector (1 if treated, otherwise 0)
+ y (np.array): an outcome vector
+ p (optional, np.array): user provided propensity score vector between 0 and 1
+ dhat (optinal, np.array): user provided predicted treatment effect vector
+
+ Returns:
+ self: returns an instance of self.
+ """
+
+ logger.info(
+ "generating out-of-fold CV outcome estimates with {}".format(self.model_mu)
+ )
+ self._outcome_estimate(X, treatment, y)
+
+ if dhat is not None:
+ self._tau_pred = dhat
+
+ if p is None:
+ self._treatment_estimate(X, treatment)
+ else:
+ self._w_pred = np.clip(p, self.clip_bounds[0], self.clip_bounds[1])
+
+ # Doubly Robust Modification
+ self._dr_score = self._tau_pred + (treatment - self._w_pred) / self._w_pred / (
+ 1 - self._w_pred
+ ) * (y - self._y_pred)
+
+ target = self._dr_score.copy()
+ target = np.sign(target)
+
+ logger.info("training the treatment assignment model, {}".format(self.model_pi))
+ self.model_pi.fit(X, target, sample_weight=abs(self._dr_score))
+
+ return self
+
+ def predict(self, X):
+ """Predict treatment assignment that optimizes the outcome.
+
+ Args:
+ X (np.matrix): a feature matrix
+
+ Returns:
+ (numpy.ndarray): predictions of treatment assignment.
+ """
+
+ return self.model_pi.predict(X)
+
+ def predict_proba(self, X):
+ """Predict treatment assignment score that optimizes the outcome.
+
+ Args:
+ X (np.matrix): a feature matrix
+
+ Returns:
+ (numpy.ndarray): predictions of treatment assignment score.
+ """
+
+ pi_hat = self.model_pi.predict_proba(X)[:, 1]
+
+ return pi_hat
diff --git a/causalml/source/causalml/optimize/unit_selection.py b/causalml/source/causalml/optimize/unit_selection.py
new file mode 100644
index 0000000000000000000000000000000000000000..0d4232f270b85d79458e7e01b898eef13361e310
--- /dev/null
+++ b/causalml/source/causalml/optimize/unit_selection.py
@@ -0,0 +1,297 @@
+import numpy as np
+
+from sklearn.base import clone
+
+import warnings
+
+
+class CounterfactualUnitSelector:
+ """
+ A highly experimental implementation of the counterfactual unit selection
+ model proposed by Li and Pearl (2019).
+
+ Parameters
+ ----------
+ learner : object
+ The base learner used to estimate the segment probabilities.
+
+ nevertaker_payoff : float
+ The payoff from targeting a never-taker
+
+ alwaystaker_payoff : float
+ The payoff from targeting an always-taker
+
+ complier_payoff : float
+ The payoff from targeting a complier
+
+ defier_payoff : float
+ The payoff from targeting a defier
+
+ organic_conversion : float, optional (default=None)
+ The organic conversion rate in the population without an intervention.
+ If None, the organic conversion rate is obtained from tne control group.
+
+ NB: The organic conversion in the control group is not always the same
+ as the organic conversion rate without treatment.
+
+ data : DataFrame
+ A pandas DataFrame containing the features, treatment assignment
+ indicator and the outcome of interest.
+
+ treatment : string
+ A string corresponding to the name of the treatment column. The
+ assumed coding in the column is 1 for treatment and 0 for control.
+
+ outcome : string
+ A string corresponding to the name of the outcome column. The assumed
+ coding in the column is 1 for conversion and 0 for no conversion.
+
+ References
+ ----------
+ Li, Ang, and Judea Pearl. 2019. “Unit Selection Based on Counterfactual
+ Logic.” https://ftp.cs.ucla.edu/pub/stat_ser/r488.pdf.
+ """
+
+ def __init__(
+ self,
+ learner,
+ nevertaker_payoff,
+ alwaystaker_payoff,
+ complier_payoff,
+ defier_payoff,
+ organic_conversion=None,
+ ):
+ self.learner = learner
+ self.nevertaker_payoff = nevertaker_payoff
+ self.alwaystaker_payoff = alwaystaker_payoff
+ self.complier_payoff = complier_payoff
+ self.defier_payoff = defier_payoff
+ self.organic_conversion = organic_conversion
+
+ def fit(self, data, treatment, outcome):
+ """
+ Fits the class.
+ """
+
+ if self._gain_equality_check():
+ self._fit_segment_model(data, treatment, outcome)
+
+ else:
+ self._fit_segment_model(data, treatment, outcome)
+ self._fit_condprob_models(data, treatment, outcome)
+
+ def predict(self, data, treatment, outcome):
+ """
+ Predicts an individual-level payoff. If gain equality is satisfied, uses
+ the exact function; if not, uses the midpoint between bounds.
+ """
+
+ if self._gain_equality_check():
+ est_payoff = self._get_exact_benefit(data, treatment, outcome)
+
+ else:
+ est_payoff = self._obj_func_midp(data, treatment, outcome)
+
+ return est_payoff
+
+ def _gain_equality_check(self):
+ """
+ Checks if gain equality is satisfied. If so, the optimization task can
+ be simplified.
+ """
+
+ return (
+ self.complier_payoff + self.defier_payoff
+ == self.alwaystaker_payoff + self.nevertaker_payoff
+ )
+
+ @staticmethod
+ def _make_segments(data, treatment, outcome):
+ """
+ Constructs the following segments:
+
+ * AC = Pr(Y = 1, W = 1 /mid X)
+ * AD = Pr(Y = 1, W = 0 /mid X)
+ * ND = Pr(Y = 0, W = 1 /mid X)
+ * ND = Pr(Y = 0, W = 0 /mid X)
+
+ where the names of the outcomes correspond the combinations of
+ the relevant segments, eg AC = Always-taker or Complier.
+ """
+
+ segments = np.empty(data.shape[0], dtype="object")
+
+ segments[(data[treatment] == 1) & (data[outcome] == 1)] = "AC"
+ segments[(data[treatment] == 0) & (data[outcome] == 1)] = "AD"
+ segments[(data[treatment] == 1) & (data[outcome] == 0)] = "ND"
+ segments[(data[treatment] == 0) & (data[outcome] == 0)] = "NC"
+
+ return segments
+
+ def _fit_segment_model(self, data, treatment, outcome):
+ """
+ Fits a classifier for estimating the probabilities for the unit
+ segment combinations.
+ """
+
+ model = clone(self.learner)
+
+ X = data.drop([treatment, outcome], axis=1)
+ y = self._make_segments(data, treatment, outcome)
+
+ self.segment_model = model.fit(X, y)
+
+ def _fit_condprob_models(self, data, treatment, outcome):
+ """
+ Fits two classifiers to estimate conversion probabilities conditional
+ on the treatment.
+ """
+
+ trt_learner = clone(self.learner)
+ ctr_learner = clone(self.learner)
+
+ treated = data[treatment] == 1
+
+ X = data.drop([treatment, outcome], axis=1)
+ y = data[outcome]
+
+ self.trt_model = trt_learner.fit(X[treated], y[treated])
+ self.ctr_model = ctr_learner.fit(X[~treated], y[~treated])
+
+ def _get_exact_benefit(self, data, treatment, outcome):
+ """
+ Calculates the exact benefit function of Theorem 4 in Li and Pearl (2019).
+ Returns the exact benefit.
+ """
+ beta = self.complier_payoff
+ gamma = self.alwaystaker_payoff
+ theta = self.nevertaker_payoff
+
+ X = data.drop([treatment, outcome], axis=1)
+
+ segment_prob = self.segment_model.predict_proba(X)
+ segment_name = self.segment_model.classes_
+
+ benefit = (
+ (beta - theta) * segment_prob[:, segment_name == "AC"]
+ + (gamma - beta) * segment_prob[:, segment_name == "AD"]
+ + theta
+ )
+
+ return benefit
+
+ def _obj_func_midp(self, data, treatment, outcome):
+ """
+ Calculates bounds for the objective function. Returns the midpoint
+ between bounds.
+
+ Parameters
+ ----------
+ pr_y1_w1 : float
+ The probability of conversion given treatment assignment.
+
+ pr_y1_w0 : float
+ The probability of conversion given control assignment.
+
+ pr_y0_w1 : float
+ The probability of no conversion given treatment assignment
+ (1 - pr_y1_w1).
+
+ pr_y0_w0 : float
+ The probability of no conversion given control assignment
+ (1 - pr_1y_w0)
+
+ pr_y1w1_x : float
+ Probability of complier or always-taker given X.
+
+ pr_y0w0_x : float
+ Probability of complier or never-taker given X.
+
+ pr_y1w0_x : float
+ Probability of defier or always-taker given X.
+
+ pr_y0w1_x : float
+ Probability of never-taker or defier given X.
+
+ pr_y_x : float
+ Organic probability of conversion.
+ """
+
+ X = data.drop([treatment, outcome], axis=1)
+
+ beta = self.complier_payoff
+ gamma = self.alwaystaker_payoff
+ theta = self.nevertaker_payoff
+ delta = self.defier_payoff
+
+ pr_y0_w1, pr_y1_w1 = np.split(
+ self.trt_model.predict_proba(X), indices_or_sections=2, axis=1
+ )
+ pr_y0_w0, pr_y1_w0 = np.split(
+ self.ctr_model.predict_proba(X), indices_or_sections=2, axis=1
+ )
+
+ segment_prob = self.segment_model.predict_proba(X)
+ segment_name = self.segment_model.classes_
+
+ pr_y1w1_x = segment_prob[:, segment_name == "AC"]
+ pr_y0w0_x = segment_prob[:, segment_name == "NC"]
+ pr_y1w0_x = segment_prob[:, segment_name == "AD"]
+ pr_y0w1_x = segment_prob[:, segment_name == "ND"]
+
+ if self.organic_conversion is not None:
+ pr_y_x = self.organic_conversion
+
+ else:
+ pr_y_x = pr_y1_w0
+ warnings.warn(
+ "Probability of organic conversion estimated from control observations."
+ )
+
+ p1 = (beta - theta) * pr_y1_w1 + delta * pr_y1_w0 + theta * pr_y0_w0
+ p2 = gamma * pr_y1_w1 + delta * pr_y0_w1 + (beta - gamma) * pr_y0_w0
+ p3 = (
+ (gamma - delta) * pr_y1_w1
+ + delta * pr_y1_w0
+ + theta * pr_y0_w0
+ + (beta - gamma - theta + delta) * (pr_y1w1_x + pr_y0w0_x)
+ )
+ p4 = (
+ (beta - theta) * pr_y1_w1
+ - (beta - gamma - theta) * pr_y1_w0
+ + theta * pr_y0_w0
+ + (beta - gamma - theta + delta) * (pr_y1w0_x + pr_y0w1_x)
+ )
+ p5 = (gamma - delta) * pr_y1_w1 + delta * pr_y1_w0 + theta * pr_y0_w0
+ p6 = (
+ (beta - theta) * pr_y1_w1
+ - (beta - gamma - theta) * pr_y1_w0
+ + theta * pr_y0_w0
+ )
+ p7 = (
+ (gamma - delta) * pr_y1_w1
+ - (beta - gamma - theta) * pr_y1_w0
+ + theta * pr_y0_w0
+ + (beta - gamma - theta + delta) * pr_y_x
+ )
+ p8 = (
+ (beta - theta) * pr_y1_w1
+ + delta * pr_y1_w0
+ + theta * pr_y0_w0
+ - (beta - gamma - theta + delta) * pr_y_x
+ )
+
+ params_1 = np.concatenate((p1, p2, p3, p4), axis=1)
+ params_2 = np.concatenate((p5, p6, p7, p8), axis=1)
+
+ sigma = beta - gamma - theta + delta
+
+ if sigma < 0:
+ lower_bound = np.max(params_1, axis=1)
+ upper_bound = np.min(params_2, axis=1)
+
+ elif sigma > 0:
+ lower_bound = np.max(params_2, axis=1)
+ upper_bound = np.min(params_1, axis=1)
+
+ return (lower_bound + upper_bound) / 2
diff --git a/causalml/source/causalml/optimize/utils.py b/causalml/source/causalml/optimize/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..5c610b896b84e16cff75049b21a05f383011ab9c
--- /dev/null
+++ b/causalml/source/causalml/optimize/utils.py
@@ -0,0 +1,137 @@
+import numpy as np
+
+
+def get_treatment_costs(treatment, control_name, cc_dict, ic_dict):
+ """
+ Set the conversion and impression costs based on a dict of parameters.
+
+ Calculate the actual cost of targeting a user with the actual treatment
+ group using the above parameters.
+
+ Params
+ ------
+ treatment : array, shape = (num_samples, )
+ Treatment array.
+
+ control_name, str
+ Control group name as string.
+
+ cc_dict : dict
+ Dict containing the conversion cost for each treatment.
+
+ ic_dict
+ Dict containing the impression cost for each treatment.
+
+ Returns
+ -------
+ conversion_cost : ndarray, shape = (num_samples, num_treatments)
+ An array of conversion costs for each treatment.
+
+ impression_cost : ndarray, shape = (num_samples, num_treatments)
+ An array of impression costs for each treatment.
+
+ conditions : list, len = len(set(treatment))
+ A list of experimental conditions.
+ """
+
+ # Set the conversion costs of the treatments
+ conversion_cost = np.zeros((len(treatment), len(cc_dict.keys())))
+ for idx, dict_key in enumerate(cc_dict.keys()):
+ conversion_cost[:, idx] = cc_dict.get(dict_key)
+
+ # Set the impression costs of the treatments
+ impression_cost = np.zeros((len(treatment), len(ic_dict.keys())))
+ for idx, dict_key in enumerate(ic_dict.keys()):
+ impression_cost[:, idx] = ic_dict.get(dict_key)
+
+ # Get a sorted list of conditions
+ conditions = list(set(treatment))
+ conditions.remove(control_name)
+ conditions_sorted = sorted(conditions)
+ conditions_sorted.insert(0, control_name)
+
+ return conversion_cost, impression_cost, conditions_sorted
+
+
+def get_actual_value(
+ treatment,
+ observed_outcome,
+ conversion_value,
+ conditions,
+ conversion_cost,
+ impression_cost,
+):
+ """
+ Set the conversion and impression costs based on a dict of parameters.
+
+ Calculate the actual value of targeting a user with the actual treatment group
+ using the above parameters.
+
+ Params
+ ------
+ treatment : array, shape = (num_samples, )
+ Treatment array.
+
+ observed_outcome : array, shape = (num_samples, )
+ Observed outcome array, aka y.
+
+ conversion_value : array, shape = (num_samples, )
+ The value of converting a given user.
+
+ conditions : list, len = len(set(treatment))
+ List of treatment conditions.
+
+ conversion_cost : array, shape = (num_samples, num_treatment)
+ Array of conversion costs for each unit in each treatment.
+
+ impression_cost : array, shape = (num_samples, num_treatment)
+ Array of impression costs for each unit in each treatment.
+
+ Returns
+ -------
+ actual_value : array, shape = (num_samples, )
+ Array of actual values of havng a user in their actual treatment group.
+
+ conversion_value : array, shape = (num_samples, )
+ Array of payoffs from converting a user.
+ """
+
+ cost_filter = [
+ actual_group == possible_group
+ for actual_group in treatment
+ for possible_group in conditions
+ ]
+
+ conversion_cost_flat = conversion_cost.flatten()
+ actual_cc = conversion_cost_flat[cost_filter]
+ impression_cost_flat = impression_cost.flatten()
+ actual_ic = impression_cost_flat[cost_filter]
+
+ # Calculate the actual value of having a user in their actual treatment
+ actual_value = (conversion_value - actual_cc) * observed_outcome - actual_ic
+
+ return actual_value
+
+
+def get_uplift_best(cate, conditions):
+ """
+ Takes the CATE prediction from a learner, adds the control
+ outcome array and finds the name of the argmax conditon.
+
+ Params
+ ------
+ cate : array, shape = (num_samples, )
+ The conditional average treatment effect prediction.
+
+ conditions : list, len = len(set(treatment))
+
+ Returns
+ -------
+ uplift_recomm_name : array, shape = (num_samples, )
+ The experimental group recommended by the learner.
+ """
+ cate_with_control = np.c_[np.zeros(cate.shape[0]), cate]
+ uplift_best_idx = np.argmax(cate_with_control, axis=1)
+ uplift_best_name = [conditions[idx] for idx in uplift_best_idx]
+
+ return uplift_best_name
diff --git a/causalml/source/causalml/optimize/value_optimization.py b/causalml/source/causalml/optimize/value_optimization.py
new file mode 100644
index 0000000000000000000000000000000000000000..a286e49d4bdead55109ed7b4f7793f94a1f61ea9
--- /dev/null
+++ b/causalml/source/causalml/optimize/value_optimization.py
@@ -0,0 +1,118 @@
+import numpy as np
+
+
+class CounterfactualValueEstimator:
+ """
+ Args
+ ----
+ treatment : array, shape = (num_samples, )
+ An array of treatment group indicator values.
+
+ control_name : string
+ The name of the control condition as a string. Must be contained in the treatment array.
+
+ treatment_names : list, length = cate.shape[1]
+ A list of treatment group names. NB: The order of the items in the
+ list must correspond to the order in which the conditional average
+ treatment effect estimates are in cate_array.
+
+ y_proba : array, shape = (num_samples, )
+ The predicted probability of conversion using the Y ~ X model across
+ the total sample.
+
+ cate : array, shape = (num_samples, len(set(treatment)))
+ Conditional average treatment effect estimations from any model.
+
+ value : array, shape = (num_samples, )
+ Value of converting each unit.
+
+ conversion_cost : shape = (num_samples, len(set(treatment)))
+ The cost of a treatment that is triggered if a unit converts after having been in the treatment, such as a
+ promotion code.
+
+ impression_cost : shape = (num_samples, len(set(treatment)))
+ The cost of a treatment that is the same for each unit whether or not they convert, such as a cost associated
+ with a promotion channel.
+
+
+ Notes
+ -----
+ Because we get the conditional average treatment effects from
+ cate-learners relative to the control condition, we subtract the
+ cate for the unit in their actual treatment group from y_proba for that
+ unit, in order to recover the control outcome. We then add the cates
+ to the control outcome to obtain y_proba under each condition. These
+ outcomes are counterfactual because just one of them is actually
+ observed.
+ """
+
+ def __init__(
+ self,
+ treatment,
+ control_name,
+ treatment_names,
+ y_proba,
+ cate,
+ value,
+ conversion_cost,
+ impression_cost,
+ *args,
+ **kwargs,
+ ):
+ self.treatment = treatment
+ self.control_name = control_name
+ self.treatment_names = treatment_names
+ self.y_proba = y_proba
+ self.cate = cate
+ self.value = value
+ self.conversion_cost = conversion_cost
+ self.impression_cost = impression_cost
+
+ def predict_best(self):
+ """
+ Predict the best treatment group based on the highest counterfactual
+ value for a treatment.
+ """
+ self._get_counterfactuals()
+ self._get_counterfactual_values()
+ return self.best_treatment
+
+ def predict_counterfactuals(self):
+ """
+ Predict the counterfactual values for each treatment group.
+ """
+ self._get_counterfactuals()
+ self._get_counterfactual_values()
+ return self.expected_values
+
+ def _get_counterfactuals(self):
+ """
+ Get an array of counterfactual outcomes based on control outcome and
+ the array of conditional average treatment effects.
+ """
+ conditions = self.treatment_names.copy()
+ conditions.insert(0, self.control_name)
+ cates_with_control = np.c_[np.zeros(self.cate.shape[0]), self.cate]
+ cates_flat = cates_with_control.flatten()
+
+ cates_filt = [
+ actual_group == poss_group
+ for actual_group in self.treatment
+ for poss_group in conditions
+ ]
+
+ control_outcome = self.y_proba - cates_flat[cates_filt]
+ self.counterfactuals = cates_with_control + control_outcome[:, None]
+
+ def _get_counterfactual_values(self):
+ """
+ Calculate the expected value of assigning a unit to each of the
+ treatment conditions given the value of conversion and the conversion
+ and impression costs associated with the treatment.
+ """
+
+ self.expected_values = (
+ self.value[:, None] - self.conversion_cost
+ ) * self.counterfactuals - self.impression_cost
+
+ self.best_treatment = np.argmax(self.expected_values, axis=1)
diff --git a/causalml/source/causalml/propensity.py b/causalml/source/causalml/propensity.py
new file mode 100644
index 0000000000000000000000000000000000000000..4e12dcb99a03ab69f050953a5c025d7220a6b323
--- /dev/null
+++ b/causalml/source/causalml/propensity.py
@@ -0,0 +1,230 @@
+from abc import ABCMeta, abstractmethod
+import logging
+import numpy as np
+from sklearn.metrics import roc_auc_score as auc
+from sklearn.linear_model import LogisticRegressionCV
+from sklearn.model_selection import StratifiedKFold, train_test_split
+from sklearn.isotonic import IsotonicRegression
+import xgboost as xgb
+
+logger = logging.getLogger("causalml")
+
+
+class PropensityModel(metaclass=ABCMeta):
+ def __init__(self, clip_bounds=(1e-3, 1 - 1e-3), calibrate=True, **model_kwargs):
+ """
+ Args:
+ clip_bounds (tuple): lower and upper bounds for clipping propensity scores. Bounds should be implemented
+ such that: 0 < lower < upper < 1, to avoid division by zero in BaseRLearner.fit_predict() step.
+ calibrate (bool): whether calibrate the propensity score
+ model_kwargs: Keyword arguments to be passed to the underlying classification model.
+ """
+ self.clip_bounds = clip_bounds
+ self.calibrate = calibrate
+ self.model_kwargs = model_kwargs
+ self.model = self._model
+ self.calibrator = None
+
+ @property
+ @abstractmethod
+ def _model(self):
+ pass
+
+ def __repr__(self):
+ return self.model.__repr__()
+
+ def fit(self, X, y):
+ """
+ Fit a propensity model.
+
+ Args:
+ X (numpy.ndarray): a feature matrix
+ y (numpy.ndarray): a binary target vector
+ """
+ self.model.fit(X, y)
+ if self.calibrate:
+ # Fit a calibrator to the propensity scores with IsotonicRegression.
+ # Ref: https://scikit-learn.org/stable/modules/isotonic.html
+ self.calibrator = IsotonicRegression(
+ out_of_bounds="clip",
+ y_min=self.clip_bounds[0],
+ y_max=self.clip_bounds[1],
+ )
+ self.calibrator.fit(self.model.predict_proba(X)[:, 1], y)
+
+ def predict(self, X):
+ """
+ Predict propensity scores.
+
+ Args:
+ X (numpy.ndarray): a feature matrix
+
+ Returns:
+ (numpy.ndarray): Propensity scores between 0 and 1.
+ """
+ p = self.model.predict_proba(X)[:, 1]
+ if self.calibrate:
+ p = self.calibrator.transform(p)
+
+ return np.clip(p, *self.clip_bounds)
+
+ def fit_predict(self, X, y):
+ """
+ Fit a propensity model and predict propensity scores.
+
+ Args:
+ X (numpy.ndarray): a feature matrix
+ y (numpy.ndarray): a binary target vector
+
+ Returns:
+ (numpy.ndarray): Propensity scores between 0 and 1.
+ """
+ self.fit(X, y)
+ propensity_scores = self.predict(X)
+ return propensity_scores
+
+
+class LogisticRegressionPropensityModel(PropensityModel):
+ """
+ Propensity regression model based on the LogisticRegression algorithm.
+ """
+
+ @property
+ def _model(self):
+ kwargs = {
+ "penalty": "elasticnet",
+ "solver": "saga",
+ "Cs": np.logspace(1e-3, 1 - 1e-3, 4),
+ "l1_ratios": np.linspace(1e-3, 1 - 1e-3, 4),
+ "cv": StratifiedKFold(
+ n_splits=(
+ self.model_kwargs.pop("n_fold")
+ if "n_fold" in self.model_kwargs
+ else 4
+ ),
+ shuffle=True,
+ random_state=self.model_kwargs.get("random_state", 42),
+ ),
+ "random_state": 42,
+ }
+ kwargs.update(self.model_kwargs)
+
+ return LogisticRegressionCV(**kwargs)
+
+
+class ElasticNetPropensityModel(LogisticRegressionPropensityModel):
+ pass
+
+
+class GradientBoostedPropensityModel(PropensityModel):
+ """
+ Gradient boosted propensity score model with optional early stopping.
+
+ Notes
+ -----
+ Please see the xgboost documentation for more information on gradient boosting tuning parameters:
+ https://xgboost.readthedocs.io/en/latest/python/python_api.html
+ """
+
+ def __init__(
+ self,
+ early_stop=False,
+ clip_bounds=(1e-3, 1 - 1e-3),
+ calibrate=True,
+ **model_kwargs,
+ ):
+ self.early_stop = early_stop
+ super().__init__(clip_bounds, calibrate, **model_kwargs)
+
+ @property
+ def _model(self):
+ kwargs = {
+ "max_depth": 8,
+ "learning_rate": 0.1,
+ "n_estimators": 100,
+ "objective": "binary:logistic",
+ "nthread": -1,
+ "colsample_bytree": 0.8,
+ "random_state": 42,
+ }
+ kwargs.update(self.model_kwargs)
+
+ if self.early_stop:
+ kwargs.update({"early_stopping_rounds": 10})
+
+ return xgb.XGBClassifier(**kwargs)
+
+ def fit(self, X, y, stop_val_size=0.2):
+ """
+ Fit a propensity model.
+
+ Args:
+ X (numpy.ndarray): a feature matrix
+ y (numpy.ndarray): a binary target vector
+ """
+
+ if self.early_stop:
+ X_train, X_val, y_train, y_val = train_test_split(
+ X, y, test_size=stop_val_size
+ )
+
+ self.model.fit(
+ X_train,
+ y_train,
+ eval_set=[(X_val, y_val)],
+ )
+ if self.calibrate:
+ self.calibrator = IsotonicRegression(
+ out_of_bounds="clip",
+ y_min=self.clip_bounds[0],
+ y_max=self.clip_bounds[1],
+ )
+ self.calibrator.fit(self.model.predict_proba(X)[:, 1], y)
+ else:
+ super().fit(X, y)
+
+
+def compute_propensity_score(
+ X,
+ treatment,
+ p_model=None,
+ X_pred=None,
+ treatment_pred=None,
+ calibrate_p=True,
+ clip_bounds=(1e-3, 1 - 1e-3),
+):
+ """Generate propensity score if user didn't provide and optionally calibrate.
+
+ Args:
+ X (np.matrix): features for training
+ treatment (np.array or pd.Series): a treatment vector for training
+ p_model (model object, optional): a binary classifier with either a predict_proba or predict method
+ X_pred (np.matrix, optional): features for prediction
+ treatment_pred (np.array or pd.Series, optional): a treatment vector for prediciton
+ calibrate_p (bool, optional): whether calibrate the propensity score
+ clip_bounds (tuple, optional): lower and upper bounds for clipping propensity scores. Bounds should be implemented
+ such that: 0 < lower < upper < 1, to avoid division by zero in BaseRLearner.fit_predict() step.
+
+ Returns:
+ (tuple)
+ - p (numpy.ndarray): propensity score
+ - p_model (PropensityModel): either the original p_model or a trained ElasticNetPropensityModel
+ """
+ if treatment_pred is None:
+ treatment_pred = treatment.copy()
+ if p_model is None:
+ p_model = ElasticNetPropensityModel(
+ clip_bounds=clip_bounds, calibrate=calibrate_p
+ )
+
+ p_model.fit(X, treatment)
+
+ X_pred = X if X_pred is None else X_pred
+
+ try:
+ p = p_model.predict_proba(X_pred)[:, 1]
+ except AttributeError:
+ logger.info("predict_proba not available, using predict instead")
+ p = p_model.predict(X_pred)
+
+ return p, p_model
diff --git a/causalml/source/docs/Makefile b/causalml/source/docs/Makefile
new file mode 100644
index 0000000000000000000000000000000000000000..b474b1706543c8669db90bc211b9f81b81197445
--- /dev/null
+++ b/causalml/source/docs/Makefile
@@ -0,0 +1,177 @@
+# Makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line.
+SPHINXOPTS =
+SPHINXBUILD = sphinx-build
+PAPER =
+BUILDDIR = _build
+
+# User-friendly check for sphinx-build
+ifeq ($(shell which $(SPHINXBUILD) >/dev/null 2>&1; echo $$?), 1)
+$(error The '$(SPHINXBUILD)' command was not found. Make sure you have Sphinx installed, then set the SPHINXBUILD environment variable to point to the full path of the '$(SPHINXBUILD)' executable. Alternatively you can add the directory with the executable to your PATH. If you don't have Sphinx installed, grab it from http://sphinx-doc.org/)
+endif
+
+# Internal variables.
+PAPEROPT_a4 = -D latex_paper_size=a4
+PAPEROPT_letter = -D latex_paper_size=letter
+ALLSPHINXOPTS = -d $(BUILDDIR)/doctrees $(PAPEROPT_$(PAPER)) $(SPHINXOPTS) .
+# the i18n builder cannot share the environment and doctrees with the others
+I18NSPHINXOPTS = $(PAPEROPT_$(PAPER)) $(SPHINXOPTS) .
+
+.PHONY: help clean html dirhtml singlehtml pickle json htmlhelp qthelp devhelp epub latex latexpdf text man changes linkcheck doctest gettext
+
+help:
+ @echo "Please use \`make ' where is one of"
+ @echo " html to make standalone HTML files"
+ @echo " dirhtml to make HTML files named index.html in directories"
+ @echo " singlehtml to make a single large HTML file"
+ @echo " pickle to make pickle files"
+ @echo " json to make JSON files"
+ @echo " htmlhelp to make HTML files and a HTML help project"
+ @echo " qthelp to make HTML files and a qthelp project"
+ @echo " devhelp to make HTML files and a Devhelp project"
+ @echo " epub to make an epub"
+ @echo " latex to make LaTeX files, you can set PAPER=a4 or PAPER=letter"
+ @echo " latexpdf to make LaTeX files and run them through pdflatex"
+ @echo " latexpdfja to make LaTeX files and run them through platex/dvipdfmx"
+ @echo " text to make text files"
+ @echo " man to make manual pages"
+ @echo " texinfo to make Texinfo files"
+ @echo " info to make Texinfo files and run them through makeinfo"
+ @echo " gettext to make PO message catalogs"
+ @echo " changes to make an overview of all changed/added/deprecated items"
+ @echo " xml to make Docutils-native XML files"
+ @echo " pseudoxml to make pseudoxml-XML files for display purposes"
+ @echo " linkcheck to check all external links for integrity"
+ @echo " doctest to run all doctests embedded in the documentation (if enabled)"
+
+clean:
+ rm -rf $(BUILDDIR)/*
+
+html:
+ $(SPHINXBUILD) -b html $(ALLSPHINXOPTS) $(BUILDDIR)/html
+ @echo
+ @echo "Build finished. The HTML pages are in $(BUILDDIR)/html."
+
+dirhtml:
+ $(SPHINXBUILD) -b dirhtml $(ALLSPHINXOPTS) $(BUILDDIR)/dirhtml
+ @echo
+ @echo "Build finished. The HTML pages are in $(BUILDDIR)/dirhtml."
+
+singlehtml:
+ $(SPHINXBUILD) -b singlehtml $(ALLSPHINXOPTS) $(BUILDDIR)/singlehtml
+ @echo
+ @echo "Build finished. The HTML page is in $(BUILDDIR)/singlehtml."
+
+pickle:
+ $(SPHINXBUILD) -b pickle $(ALLSPHINXOPTS) $(BUILDDIR)/pickle
+ @echo
+ @echo "Build finished; now you can process the pickle files."
+
+json:
+ $(SPHINXBUILD) -b json $(ALLSPHINXOPTS) $(BUILDDIR)/json
+ @echo
+ @echo "Build finished; now you can process the JSON files."
+
+htmlhelp:
+ $(SPHINXBUILD) -b htmlhelp $(ALLSPHINXOPTS) $(BUILDDIR)/htmlhelp
+ @echo
+ @echo "Build finished; now you can run HTML Help Workshop with the" \
+ ".hhp project file in $(BUILDDIR)/htmlhelp."
+
+qthelp:
+ $(SPHINXBUILD) -b qthelp $(ALLSPHINXOPTS) $(BUILDDIR)/qthelp
+ @echo
+ @echo "Build finished; now you can run "qcollectiongenerator" with the" \
+ ".qhcp project file in $(BUILDDIR)/qthelp, like this:"
+ @echo "# qcollectiongenerator $(BUILDDIR)/qthelp/causalml.qhcp"
+ @echo "To view the help file:"
+ @echo "# assistant -collectionFile $(BUILDDIR)/qthelp/causalml.qhc"
+
+devhelp:
+ $(SPHINXBUILD) -b devhelp $(ALLSPHINXOPTS) $(BUILDDIR)/devhelp
+ @echo
+ @echo "Build finished."
+ @echo "To view the help file:"
+ @echo "# mkdir -p $$HOME/.local/share/devhelp/causalml"
+ @echo "# ln -s $(BUILDDIR)/devhelp $$HOME/.local/share/devhelp/causalml"
+ @echo "# devhelp"
+
+epub:
+ $(SPHINXBUILD) -b epub $(ALLSPHINXOPTS) $(BUILDDIR)/epub
+ @echo
+ @echo "Build finished. The epub file is in $(BUILDDIR)/epub."
+
+latex:
+ $(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(BUILDDIR)/latex
+ @echo
+ @echo "Build finished; the LaTeX files are in $(BUILDDIR)/latex."
+ @echo "Run \`make' in that directory to run these through (pdf)latex" \
+ "(use \`make latexpdf' here to do that automatically)."
+
+latexpdf:
+ $(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(BUILDDIR)/latex
+ @echo "Running LaTeX files through pdflatex..."
+ $(MAKE) -C $(BUILDDIR)/latex all-pdf
+ @echo "pdflatex finished; the PDF files are in $(BUILDDIR)/latex."
+
+latexpdfja:
+ $(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(BUILDDIR)/latex
+ @echo "Running LaTeX files through platex and dvipdfmx..."
+ $(MAKE) -C $(BUILDDIR)/latex all-pdf-ja
+ @echo "pdflatex finished; the PDF files are in $(BUILDDIR)/latex."
+
+text:
+ $(SPHINXBUILD) -b text $(ALLSPHINXOPTS) $(BUILDDIR)/text
+ @echo
+ @echo "Build finished. The text files are in $(BUILDDIR)/text."
+
+man:
+ $(SPHINXBUILD) -b man $(ALLSPHINXOPTS) $(BUILDDIR)/man
+ @echo
+ @echo "Build finished. The manual pages are in $(BUILDDIR)/man."
+
+texinfo:
+ $(SPHINXBUILD) -b texinfo $(ALLSPHINXOPTS) $(BUILDDIR)/texinfo
+ @echo
+ @echo "Build finished. The Texinfo files are in $(BUILDDIR)/texinfo."
+ @echo "Run \`make' in that directory to run these through makeinfo" \
+ "(use \`make info' here to do that automatically)."
+
+info:
+ $(SPHINXBUILD) -b texinfo $(ALLSPHINXOPTS) $(BUILDDIR)/texinfo
+ @echo "Running Texinfo files through makeinfo..."
+ make -C $(BUILDDIR)/texinfo info
+ @echo "makeinfo finished; the Info files are in $(BUILDDIR)/texinfo."
+
+gettext:
+ $(SPHINXBUILD) -b gettext $(I18NSPHINXOPTS) $(BUILDDIR)/locale
+ @echo
+ @echo "Build finished. The message catalogs are in $(BUILDDIR)/locale."
+
+changes:
+ $(SPHINXBUILD) -b changes $(ALLSPHINXOPTS) $(BUILDDIR)/changes
+ @echo
+ @echo "The overview file is in $(BUILDDIR)/changes."
+
+linkcheck:
+ $(SPHINXBUILD) -b linkcheck $(ALLSPHINXOPTS) $(BUILDDIR)/linkcheck
+ @echo
+ @echo "Link check complete; look for any errors in the above output " \
+ "or in $(BUILDDIR)/linkcheck/output.txt."
+
+doctest:
+ $(SPHINXBUILD) -b doctest $(ALLSPHINXOPTS) $(BUILDDIR)/doctest
+ @echo "Testing of doctests in the sources finished, look at the " \
+ "results in $(BUILDDIR)/doctest/output.txt."
+
+xml:
+ $(SPHINXBUILD) -b xml $(ALLSPHINXOPTS) $(BUILDDIR)/xml
+ @echo
+ @echo "Build finished. The XML files are in $(BUILDDIR)/xml."
+
+pseudoxml:
+ $(SPHINXBUILD) -b pseudoxml $(ALLSPHINXOPTS) $(BUILDDIR)/pseudoxml
+ @echo
+ @echo "Build finished. The pseudo-XML files are in $(BUILDDIR)/pseudoxml."
diff --git a/causalml/source/docs/_static/img/auuc_table_vis.png b/causalml/source/docs/_static/img/auuc_table_vis.png
new file mode 100644
index 0000000000000000000000000000000000000000..bb3770fce48a1a2add73c3bf9e176b71035d6dcc
Binary files /dev/null and b/causalml/source/docs/_static/img/auuc_table_vis.png differ
diff --git a/causalml/source/docs/_static/img/auuc_vis.png b/causalml/source/docs/_static/img/auuc_vis.png
new file mode 100644
index 0000000000000000000000000000000000000000..180b3d41582ddc63ddae56f0c2a29ee82d61755e
Binary files /dev/null and b/causalml/source/docs/_static/img/auuc_vis.png differ
diff --git a/causalml/source/docs/_static/img/counterfactual_value_optimization.png b/causalml/source/docs/_static/img/counterfactual_value_optimization.png
new file mode 100644
index 0000000000000000000000000000000000000000..197844786f667baf88df9f26117c8056e31595fe
Binary files /dev/null and b/causalml/source/docs/_static/img/counterfactual_value_optimization.png differ
diff --git a/causalml/source/docs/_static/img/logo/android-chrome-192x192.png b/causalml/source/docs/_static/img/logo/android-chrome-192x192.png
new file mode 100644
index 0000000000000000000000000000000000000000..033cd5b0a0fabe33cbce03a8e00fd08e8177cfc5
Binary files /dev/null and b/causalml/source/docs/_static/img/logo/android-chrome-192x192.png differ
diff --git a/causalml/source/docs/_static/img/logo/android-chrome-512x512.png b/causalml/source/docs/_static/img/logo/android-chrome-512x512.png
new file mode 100644
index 0000000000000000000000000000000000000000..4bde8b675a7cbd852319f7d4b76cbb7709fc7f6b
Binary files /dev/null and b/causalml/source/docs/_static/img/logo/android-chrome-512x512.png differ
diff --git a/causalml/source/docs/_static/img/logo/apple-touch-icon.png b/causalml/source/docs/_static/img/logo/apple-touch-icon.png
new file mode 100644
index 0000000000000000000000000000000000000000..12e1abd1b2e3e6143f57812f085970463be0915f
Binary files /dev/null and b/causalml/source/docs/_static/img/logo/apple-touch-icon.png differ
diff --git a/causalml/source/docs/_static/img/logo/causalml_logo.png b/causalml/source/docs/_static/img/logo/causalml_logo.png
new file mode 100644
index 0000000000000000000000000000000000000000..f99812a9c2dcad19abea8e04a7b677edc69a7b49
Binary files /dev/null and b/causalml/source/docs/_static/img/logo/causalml_logo.png differ
diff --git a/causalml/source/docs/_static/img/logo/causalml_logo.svg b/causalml/source/docs/_static/img/logo/causalml_logo.svg
new file mode 100644
index 0000000000000000000000000000000000000000..6f029805e1794ae9954bd953070ec96e337164a1
--- /dev/null
+++ b/causalml/source/docs/_static/img/logo/causalml_logo.svg
@@ -0,0 +1 @@
+
\ No newline at end of file
diff --git a/causalml/source/docs/_static/img/logo/causalml_logo_square.png b/causalml/source/docs/_static/img/logo/causalml_logo_square.png
new file mode 100644
index 0000000000000000000000000000000000000000..3f33de1b1b08d4e1c29933486edb3ca2d1e7d85f
Binary files /dev/null and b/causalml/source/docs/_static/img/logo/causalml_logo_square.png differ
diff --git a/causalml/source/docs/_static/img/logo/causalml_logo_square_transparent.png b/causalml/source/docs/_static/img/logo/causalml_logo_square_transparent.png
new file mode 100644
index 0000000000000000000000000000000000000000..f7d1c2ae1a5ed6bf1591425dd93b16666fc76045
Binary files /dev/null and b/causalml/source/docs/_static/img/logo/causalml_logo_square_transparent.png differ
diff --git a/causalml/source/docs/_static/img/logo/causalml_logo_transparent.png b/causalml/source/docs/_static/img/logo/causalml_logo_transparent.png
new file mode 100644
index 0000000000000000000000000000000000000000..c9a135cf3a43a4e674e66fbbd4c63ff8346f8f0e
Binary files /dev/null and b/causalml/source/docs/_static/img/logo/causalml_logo_transparent.png differ
diff --git a/causalml/source/docs/_static/img/logo/favicon-16x16.png b/causalml/source/docs/_static/img/logo/favicon-16x16.png
new file mode 100644
index 0000000000000000000000000000000000000000..afc805c0ce919952de64dd06f61d235320b16124
Binary files /dev/null and b/causalml/source/docs/_static/img/logo/favicon-16x16.png differ
diff --git a/causalml/source/docs/_static/img/logo/favicon-32x32.png b/causalml/source/docs/_static/img/logo/favicon-32x32.png
new file mode 100644
index 0000000000000000000000000000000000000000..64020639aa20421d1d8b1ec854968323bf94914c
Binary files /dev/null and b/causalml/source/docs/_static/img/logo/favicon-32x32.png differ
diff --git a/causalml/source/docs/_static/img/logo/favicon.ico b/causalml/source/docs/_static/img/logo/favicon.ico
new file mode 100644
index 0000000000000000000000000000000000000000..522e458727e797b344b33db689b92723e8e65274
Binary files /dev/null and b/causalml/source/docs/_static/img/logo/favicon.ico differ
diff --git a/causalml/source/docs/_static/img/meta_feature_imp_vis.png b/causalml/source/docs/_static/img/meta_feature_imp_vis.png
new file mode 100644
index 0000000000000000000000000000000000000000..51192a472e68f3df2aa4781adc032a9c105a37c4
Binary files /dev/null and b/causalml/source/docs/_static/img/meta_feature_imp_vis.png differ
diff --git a/causalml/source/docs/_static/img/meta_shap_dependence_vis.png b/causalml/source/docs/_static/img/meta_shap_dependence_vis.png
new file mode 100644
index 0000000000000000000000000000000000000000..4d59208c50bd615b6314dc196485f1e64bf239cc
Binary files /dev/null and b/causalml/source/docs/_static/img/meta_shap_dependence_vis.png differ
diff --git a/causalml/source/docs/_static/img/meta_shap_vis.png b/causalml/source/docs/_static/img/meta_shap_vis.png
new file mode 100644
index 0000000000000000000000000000000000000000..ce1c6528381d193d33a86afbe4fb9b212b671ac3
Binary files /dev/null and b/causalml/source/docs/_static/img/meta_shap_vis.png differ
diff --git a/causalml/source/docs/_static/img/sensitivity_selection_bias_r2.png b/causalml/source/docs/_static/img/sensitivity_selection_bias_r2.png
new file mode 100644
index 0000000000000000000000000000000000000000..6bb199a7246793a4d032089bfa9a9024948c79b9
Binary files /dev/null and b/causalml/source/docs/_static/img/sensitivity_selection_bias_r2.png differ
diff --git a/causalml/source/docs/_static/img/shap_vis.png b/causalml/source/docs/_static/img/shap_vis.png
new file mode 100644
index 0000000000000000000000000000000000000000..87d92f2f014baa0603dccd60942831d8668362b3
Binary files /dev/null and b/causalml/source/docs/_static/img/shap_vis.png differ
diff --git a/causalml/source/docs/_static/img/synthetic_dgp_bar_plot_multiple.png b/causalml/source/docs/_static/img/synthetic_dgp_bar_plot_multiple.png
new file mode 100644
index 0000000000000000000000000000000000000000..bc849783dcec2b2801144bab2db494eb9acb422f
Binary files /dev/null and b/causalml/source/docs/_static/img/synthetic_dgp_bar_plot_multiple.png differ
diff --git a/causalml/source/docs/_static/img/synthetic_dgp_scatter_plot.png b/causalml/source/docs/_static/img/synthetic_dgp_scatter_plot.png
new file mode 100644
index 0000000000000000000000000000000000000000..95161986488b6c4c2f649929adc54f3b0b28ae0b
--- /dev/null
+++ b/causalml/source/docs/_static/img/synthetic_dgp_scatter_plot.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f155b565b101e84659c3a9bd559a422399fc3321bd78ff861aead83235acc359
+size 182765
diff --git a/causalml/source/docs/_static/img/synthetic_dgp_scatter_plot_multiple.png b/causalml/source/docs/_static/img/synthetic_dgp_scatter_plot_multiple.png
new file mode 100644
index 0000000000000000000000000000000000000000..5b80131c10b585179d68586cfd4327c1f1d156bb
Binary files /dev/null and b/causalml/source/docs/_static/img/synthetic_dgp_scatter_plot_multiple.png differ
diff --git a/causalml/source/docs/_static/img/uplift_tree_feature_imp_vis.png b/causalml/source/docs/_static/img/uplift_tree_feature_imp_vis.png
new file mode 100644
index 0000000000000000000000000000000000000000..d2d168f7cec2fe95831d14c904ef47cfe236d71f
Binary files /dev/null and b/causalml/source/docs/_static/img/uplift_tree_feature_imp_vis.png differ
diff --git a/causalml/source/docs/_static/img/uplift_tree_vis.png b/causalml/source/docs/_static/img/uplift_tree_vis.png
new file mode 100644
index 0000000000000000000000000000000000000000..a4e5e55fda44ae4631d00416d69fbb26c6a1cfa6
--- /dev/null
+++ b/causalml/source/docs/_static/img/uplift_tree_vis.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0635de48c76bcf4083527b18e42e4e9318eae41090e6710fb4f60fb9c391f8e5
+size 149927
diff --git a/causalml/source/docs/about.rst b/causalml/source/docs/about.rst
new file mode 100644
index 0000000000000000000000000000000000000000..ebeb7d743dc8c803984c2d098eb462d3b3b33ad7
--- /dev/null
+++ b/causalml/source/docs/about.rst
@@ -0,0 +1,67 @@
+About CausalML
+===========================
+
+``CausalML`` is a Python package that provides a suite of uplift modeling and causal inference methods using machine learning algorithms based on recent research.
+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 **W** on outcome **Y** for users with observed features **X**, without strong assumptions on the model form.
+
+GitHub Repo
+-----------
+
+https://github.com/uber/causalml
+
+Mission
+-------
+
+From the CausalML `Charter `_:
+
+ 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.
+
+Contributing
+------------
+`Contributing.md `_
+
+Governance
+----------
+* `Charter `_
+* `Contributors `_
+* `Maintainers `_
+
+Intro to Causal Machine Learning
+================================
+
+What is Causal Machine Learning?
+--------------------------------
+
+Causal machine learning is a branch of machine learning that focuses on understanding the cause and effect relationships in data. It goes beyond just predicting outcomes based on patterns in the data, and tries to understand how changing one variable can affect an outcome.
+Suppose we are trying to predict a student’s test score based on how many hours they study and how much sleep they get. Traditional machine learning models would find patterns in the data, like students who study more or sleep more tend to get higher scores.
+But what if you want to know what would happen if a student studied an extra hour each day? Or slept an extra hour each night? Modeling these potential outcomes or counterfactuals is where causal machine learning comes in. It tries to understand cause-and-effect relationships - how much changing one variable (like study hours or sleep hours) will affect the outcome (the test score).
+This is useful in many fields, including economics, healthcare, and policy making, where understanding the impact of interventions is crucial.
+While traditional machine learning is great for prediction, causal machine learning helps us understand the difference in outcomes due to interventions.
+
+
+
+Difference from Traditional Machine Learning
+--------------------------------------------
+
+Traditional machine learning and causal machine learning are both powerful tools, but they serve different purposes and answer different types of questions.
+Traditional Machine Learning is primarily concerned with prediction. Given a set of input features, it learns a function from the data that can predict an outcome. It’s great at finding patterns and correlations in large datasets, but it doesn’t tell us about the cause-and-effect relationships between variables. It answers questions like “Given a patient’s symptoms, what disease are they likely to have?”
+On the other hand, Causal Machine Learning is concerned with understanding the cause-and-effect relationships between variables. It goes beyond prediction and tries to answer questions about intervention: “What will happen if we change this variable?” For example, in a medical context, it could help answer questions like “What will happen if a patient takes this medication?”
+In essence, while traditional machine learning can tell us “what is”, causal machine learning can help us understand “what if”. This makes causal machine learning particularly useful in fields where we need to make decisions based on data, such as policy making, economics, and healthcare.
+
+
+Measuring Causal Effects
+------------------------
+
+**Randomized Control Trials (RCT)** are the gold standard for causal effect measurements. Subjects are randomly exposed to a treatment and the Average Treatment Effect (ATE) is measured as the difference between the mean effects in the treatment and control groups. Random assignment removes the effect of any confounders on the treatment.
+
+If an RCT is available and the treatment effects are heterogeneous across covariates, measuring the conditional average treatment effect(CATE) can be of interest. The CATE is an estimate of the treatment effect conditioned on all available experiment covariates and confounders. We call these Heterogeneous Treatment Effects (HTEs).
+
+
+Example Use Cases
+-----------------
+
+- **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 might have multiple options to interact with its customers such as different product choices in up-sell or different 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 engagement experience.
+
diff --git a/causalml/source/docs/causalml.rst b/causalml/source/docs/causalml.rst
new file mode 100644
index 0000000000000000000000000000000000000000..5ea378c802866dd6cc29310170232a0895288f38
--- /dev/null
+++ b/causalml/source/docs/causalml.rst
@@ -0,0 +1,119 @@
+causalml package
+================
+
+Submodules
+----------
+
+causalml.inference.tree module
+------------------------------
+
+.. automodule:: causalml.inference.tree
+ :members:
+ :imported-members:
+ :undoc-members:
+ :show-inheritance:
+
+causalml.inference.meta module
+------------------------------
+
+.. automodule:: causalml.inference.meta
+ :members:
+ :imported-members:
+ :undoc-members:
+ :show-inheritance:
+
+causalml.inference.iv module
+----------------------------
+
+.. automodule:: causalml.inference.iv
+ :members:
+ :imported-members:
+ :undoc-members:
+ :show-inheritance:
+
+causalml.inference.nn module
+----------------------------
+
+.. automodule:: causalml.inference.nn
+ :members:
+ :imported-members:
+ :undoc-members:
+ :show-inheritance:
+
+causalml.inference.tf module
+----------------------------
+
+.. automodule:: causalml.inference.tf
+ :members:
+ :imported-members:
+ :undoc-members:
+ :show-inheritance:
+
+causalml.optimize module
+------------------------
+
+.. automodule:: causalml.optimize
+ :members:
+ :imported-members:
+ :undoc-members:
+ :show-inheritance:
+
+causalml.dataset module
+-----------------------
+
+.. automodule:: causalml.dataset
+ :members:
+ :imported-members:
+ :undoc-members:
+ :show-inheritance:
+
+causalml.match module
+---------------------
+
+.. automodule:: causalml.match
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+causalml.propensity module
+--------------------------
+
+.. automodule:: causalml.propensity
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+causalml.metrics module
+-----------------------
+
+.. automodule:: causalml.metrics
+ :members:
+ :imported-members:
+ :undoc-members:
+ :show-inheritance:
+
+causalml.feature_selection module
+---------------------------------
+
+.. automodule:: causalml.feature_selection
+ :members:
+ :imported-members:
+ :undoc-members:
+ :show-inheritance:
+
+causalml.features module
+------------------------
+
+.. automodule:: causalml.features
+ :members:
+ :undoc-members:
+ :show-inheritance:
+
+
+Module contents
+---------------
+
+.. automodule:: causalml
+ :members:
+ :undoc-members:
+ :show-inheritance:
diff --git a/causalml/source/docs/changelog.rst b/causalml/source/docs/changelog.rst
new file mode 100644
index 0000000000000000000000000000000000000000..ad6291b94754497e1f4e65d42ee2f9e9b7e21cae
--- /dev/null
+++ b/causalml/source/docs/changelog.rst
@@ -0,0 +1,400 @@
+.. :changelog:
+
+Changelog
+=========
+
+You can find the latest changes in the `GitHub releases `_
+
+0.15.1 (Apr 2024)
+-----------------
+* This release fixes the build failure on macOS and a few bugs in ``UpliftTreeClassifier``.
+* We have two new contributors, @lee-junseok and @IanDelbridge. Thanks for your contributions!
+
+Updates
+~~~~~~~
+* Relax ``pandas`` version requirement by @jeongyoonlee in https://github.com/uber/causalml/pull/743
+* Remove undefined variables in ``match.__main__()`` by @jeongyoonlee in https://github.com/uber/causalml/pull/749
+* Fix ``distr_plot_single_sim()`` by @jeongyoonlee in https://github.com/uber/causalml/pull/750
+* Add ``with_std``, ``with_counts`` to ``create_table_one`` by @lee-junseok in https://github.com/uber/causalml/pull/748
+* fix stratified sampling call by @IanDelbridge in https://github.com/uber/causalml/pull/756
+* 20240207 honest leaf size by @IanDelbridge in https://github.com/uber/causalml/pull/753
+* 757: add ``return_ci=True`` in sensitivity by @lee-junseok in https://github.com/uber/causalml/pull/758
+* Update sensitivity tests with more meta-learners by @jeongyoonlee in https://github.com/uber/causalml/pull/759
+* manually specify ``multiprocessing`` use ``fork`` in ``setup.py`` by @IanDelbridge in https://github.com/uber/causalml/pull/754
+
+New contributors
+~~~~~~~~~~~~~~~~
+* @lee-junseok made their first contribution in https://github.com/uber/causalml/pull/748
+* @IanDelbridge made their first contribution in https://github.com/uber/causalml/pull/756
+
+0.15.0 (Feb 2024)
+-----------------
+* In this release, we revamped documentation, cleaned up dependencies, and improved installation - in addition to the long list of bug fixes.
+* We have three new contributors, @peterloleungyau, @SuperBo, and @ZiJiaW, who submitted their first PRs to CausalML. @erikcs also contributed to @ras44's PR #729 to add the wrapper for his MAQ implementation to CausalML. Thanks for your contributions!
+
+Updates
+~~~~~~~
+* Update python-publish.yml by @jeongyoonlee in https://github.com/uber/causalml/pull/673
+* Add build.[os, tools.python] to .readthedocs.yml by @jeongyoonlee in https://github.com/uber/causalml/pull/676
+* Update notebook example with causal trees interpretation by @alexander-pv in https://github.com/uber/causalml/pull/683
+* Remove the numpy and pandas version restriction in pyproject.toml by @jeongyoonlee in https://github.com/uber/causalml/pull/681
+* Add governance documents by @jeongyoonlee in https://github.com/uber/causalml/pull/688
+* Update GOVERNANCE.md by @ras44 in https://github.com/uber/causalml/pull/691
+* Dev/governance docs to snake-case by @ras44 in https://github.com/uber/causalml/pull/693
+* Reduce sklearn dependency in causalml by @alexander-pv in https://github.com/uber/causalml/pull/686
+* Update MAINTAINERS.md by @jeongyoonlee in https://github.com/uber/causalml/pull/696
+* Modified to speed up UpliftTreeClassifier.growDecisionTreeFrom. by @peterloleungyau in https://github.com/uber/causalml/pull/695
+* Update README.md by @ras44 in https://github.com/uber/causalml/pull/698
+* Add notebook examples to docs by @jeongyoonlee in https://github.com/uber/causalml/pull/697
+* resolves change requests in #166 by @ras44 in https://github.com/uber/causalml/pull/701
+* Fix the readthedocs build error by @jeongyoonlee in https://github.com/uber/causalml/pull/702
+* Replace Stack and PriorityHeap with cpp stack/heap methods in trees by @SuperBo in https://github.com/uber/causalml/pull/700
+* Hotfix for #701 by @jeongyoonlee in https://github.com/uber/causalml/pull/705
+* Dev/699 win build fix by @ras44 in https://github.com/uber/causalml/pull/710
+* expose n_jobs for rlearner by @ZiJiaW in https://github.com/uber/causalml/pull/714
+* minimal fix to resolve #707 by @ras44 in https://github.com/uber/causalml/pull/720
+* Add Python 3.10, 3.11, 3.12 to the testing by @cclauss in https://github.com/uber/causalml/pull/454
+* Remove Python 3.12 from the build tests in python-test.yaml by @jeongyoonlee in https://github.com/uber/causalml/pull/726
+* fix plot_std_diffs, add bal_tol, condense to one plot by @ras44 in https://github.com/uber/causalml/pull/723
+* Dev/677 documentation by @ras44 in https://github.com/uber/causalml/pull/725
+* documentation updates by @ras44 in https://github.com/uber/causalml/pull/728
+* resolves #730, docs clean conda install by @ras44 in https://github.com/uber/causalml/pull/731
+* minimal wrapper of MAQ #662 by @ras44 in https://github.com/uber/causalml/pull/729
+* Temporary fix for causal trees missing values support #733 by @alexander-pv in https://github.com/uber/causalml/pull/734
+* resolves #639, credit due to Dong Liu by @ras44 in https://github.com/uber/causalml/pull/722
+
+New contributors
+~~~~~~~~~~~~~~~~
+* @peterloleungyau made their first contribution in https://github.com/uber/causalml/pull/695
+* @SuperBo made their first contribution in https://github.com/uber/causalml/pull/700
+* @ZiJiaW made their first contribution in https://github.com/uber/causalml/pull/714
+
+
+0.14.1 (Aug 2023)
+-----------------
+* This release mainly addressed installation issues and updated documentation accordingly.
+* We have 4 new contributors. @bsaunders27, @xhulianoThe1, @zpppy, and @bsaunders23. Thanks for your contributions!
+
+Updates
+~~~~~~~
+* Update the python-publish workflow file to fix the package publish Gi… by @jeongyoonlee in https://github.com/uber/causalml/pull/633
+* Update Cython dependency by @alexander-pv in https://github.com/uber/causalml/pull/640
+* Fix for builds on Mac M1 infrastructure by @bsaunders27 in https://github.com/uber/causalml/pull/641
+* code cleanups by @xhulianoThe1 in https://github.com/uber/causalml/pull/634
+* support valid error early stopping by @zpppy in https://github.com/uber/causalml/pull/614
+* fix: update to ``envs/`` conda build for precompiled M1 installs by @bsaunders27 in https://github.com/uber/causalml/pull/646
+* Installation updates to README and .github/workflows by @ras44 in https://github.com/uber/causalml/pull/637
+* fix: simulate_randomized_trial by @bsaunders23 in https://github.com/uber/causalml/pull/656
+* issue 252 by @vincewu51 in https://github.com/uber/causalml/pull/660
+* ras44/651 graph viz, resolves #651 by @ras44 in https://github.com/uber/causalml/pull/661
+* linted with black by @ras44 in https://github.com/uber/causalml/pull/663
+* Fix issue 650 by @vincewu51 in https://github.com/uber/causalml/pull/659
+* Install graphviz in the workflow builds by @jeongyoonlee in https://github.com/uber/causalml/pull/668
+* Update docs/installation.rst by @jeongyoonlee in https://github.com/uber/causalml/pull/667
+* Schedule monthly PyPI install tests by @jeongyoonlee in https://github.com/uber/causalml/pull/670
+
+New contributors
+~~~~~~~~~~~~~~~~
+* @bsaunders27 made their first contribution in https://github.com/uber/causalml/pull/641
+* @xhulianoThe1 made their first contribution in https://github.com/uber/causalml/pull/634
+* @zpppy made their first contribution in https://github.com/uber/causalml/pull/614
+* @bsaunders23 made their first contribution in https://github.com/uber/causalml/pull/656
+
+
+0.14.0 (July 2023)
+------------------
+- CausalML surpassed `2MM downloads `_ on PyPI and `4,100 stars `_ on GitHub. Thanks for choosing CausalML and supporting us on GitHub.
+- We have 7 new contributors: @darthtrevino, @ras44, @AbhishekVermaDH, @joel-mcmurry, @AlxClt, @kklein, and @volico. Thanks for your contributions!
+
+Updates
+~~~~~~~
+- Fix the readthedocs build failure by @jeongyoonlee in https://github.com/uber/causalml/pull/545
+- Add ``pyproject.toml`` with basic build dependencies for PEP518 compliance by @darthtrevino in https://github.com/uber/causalml/pull/553
+- bump ``numpy`` from 1.20.3 to 1.23.2 in ``environment-py38.yml`` #338 by @ras44 in https://github.com/uber/causalml/pull/550
+- CausalTree split criterions fix and fit optimization by @alexander-pv in https://github.com/uber/causalml/pull/557
+- fixing math notations for proper rendering by @AbhishekVermaDH in https://github.com/uber/causalml/pull/558
+- Update ``methodology.rst`` by @joel-mcmurry in https://github.com/uber/causalml/pull/568
+- Causal trees bootstrapping and ``max_leaf_nodes`` fixes with minor update by @alexander-pv in https://github.com/uber/causalml/pull/583
+- Fix #596 by @AlxClt in https://github.com/uber/causalml/pull/597
+- Add ``**kwargs`` to ``Explainer.plot_shap_values()`` by @jeongyoonlee in https://github.com/uber/causalml/pull/603
+- Make the Adam optimization optional and learning rate/epochs configurable in DragonNet by @jeongyoonlee in https://github.com/uber/causalml/pull/604
+- Fix bug in variance calculation in drivlearner. by @huigangchen in https://github.com/uber/causalml/pull/606
+- Bug Fix in Dragonnet: Adam parameter name lr depreciation by @huigangchen in https://github.com/uber/causalml/pull/617
+- Fix AttributeError in builds with ``numpy>=1.24`` and ``pandas>=2.0`` by @jeongyoonlee in https://github.com/uber/causalml/pull/631
+- Pass on ``**kwargs`` in ``plot_shap_values`` of base meta leaner by @kklein in https://github.com/uber/causalml/pull/627
+- Bump ``scipy`` from 1.4.1 to 1.10.0 by @dependabot in https://github.com/uber/causalml/pull/629
+- Feature/ttest criterion by @volico in https://github.com/uber/causalml/pull/570
+- Added Interaction Tree (IT), Causal Inference Tree (CIT), and Invariant DDP (IDDP) by @jroessler in https://github.com/uber/causalml/pull/562
+- Causal trees option to return counterfactual outcomes by @alexander-pv in https://github.com/uber/causalml/pull/623
+
+New contributors
+~~~~~~~~~~~~~~~~
+- @darthtrevino made their first contribution in https://github.com/uber/causalml/pull/553
+- @ras44 made their first contribution in https://github.com/uber/causalml/pull/550
+- @AbhishekVermaDH made their first contribution in https://github.com/uber/causalml/pull/558
+- @joel-mcmurry made their first contribution in https://github.com/uber/causalml/pull/568
+- @AlxClt made their first contribution in https://github.com/uber/causalml/pull/597
+- @kklein made their first contribution in https://github.com/uber/causalml/pull/627
+- @volico made their first contribution in https://github.com/uber/causalml/pull/570
+
+
+0.13.0 (Sep 2022)
+-----------------
+- CausalML surpassed `1MM downloads `_ on PyPI and `3,200 stars `_ on GitHub. Thanks for choosing CausalML and supporting us on GitHub.
+- We have 7 new contributors @saiwing-yeung, @lixuan12315, @aldenrogers, @vincewu51, @AlkanSte, @enzoliao, and @alexander-pv. Thanks for your contributions!
+- @alexander-pv revamped `CausalTreeRegressor` and added `CausalRandomForestRegressor` with more seamless integration with `scikit-learn`'s Cython tree module. He also added integration with `shap` for causal tree/ random forest interpretation. Please check out the `example notebook `_.
+- We dropped the support for Python 3.6 and removed its test workflow.
+
+Updates
+~~~~~~~
+- Fix typo ``(% -> $)`` by @saiwing-yeung in https://github.com/uber/causalml/pull/488
+- Add function for calculating PNS bounds by @t-tte in https://github.com/uber/causalml/pull/482
+- Fix hard coding bug by @t-tte in https://github.com/uber/causalml/pull/492
+- Update README of ``conda`` install and instruction of maintain in ``conda-forge`` by @ppstacy in https://github.com/uber/causalml/pull/485
+- Update ``examples.rst`` by @lixuan12315 in https://github.com/uber/causalml/pull/496
+- Fix incorrect ``effect_learner_objective`` in ``XGBRRegressor`` by @jeongyoonlee in https://github.com/uber/causalml/pull/504
+- Fix Filter F doesn't work with latest ``statsmodels``' F test f-value format by @paullo0106 in https://github.com/uber/causalml/pull/505
+- Exclude tests in ``setup.py`` by @aldenrogers in https://github.com/uber/causalml/pull/508
+- Enabling higher orders feature importance for F filter and LR filter by @zhenyuz0500 in https://github.com/uber/causalml/pull/509
+- Ate pretrain 0506 by @vincewu51 in https://github.com/uber/causalml/pull/511
+- Update ``methodology.rst`` by @AlkanSte in https://github.com/uber/causalml/pull/518
+- Fix the bug of incorrect result in qini for multiple models by @enzoliao in https://github.com/uber/causalml/pull/520
+- Test ``get_qini()`` by @enzoliao in https://github.com/uber/causalml/pull/523
+- Fixed typo in ``uplift_trees_with_synthetic_data.ipynb`` by @jroessler in https://github.com/uber/causalml/pull/531
+- Remove Python 3.6 test from workflows by @jeongyoonlee in https://github.com/uber/causalml/pull/535
+- Causal trees update by @alexander-pv in https://github.com/uber/causalml/pull/522
+- Causal trees interpretation example by @alexander-pv in https://github.com/uber/causalml/pull/536
+
+
+0.12.3 (Feb 2022)
+-----------------
+This patch is to release a version without the constraint for Shap to be abled to use for Conda.
+
+Updates
+~~~~~~~
+- `#483 `_ by @ppstacy: Modify the requirement version of Shap
+
+
+0.12.2 (Feb 2022)
+-----------------
+This patch includes three updates by @tonkolviktor and @heiderich as follows. We also start using `black `_, a Python formatter. Please check out the updated `contribution guideline `_ to learn how to use it.
+
+Updates
+~~~~~~~
+- `#473 `_ by @tonkolviktor: Open up the scipy dependency version
+- `#476 `_ by @heiderich: Use preferred backend for joblib instead of hard-coding it
+- `#477 `_ by @heiderich: Allow parallel prediction for UpliftRandomForestClassifier and make the joblib's preferred backend configurable
+
+
+0.12.1 (Feb 2022)
+-----------------
+This patch includes two bug fixes for UpliftRandomForestClassifier as follows:
+
+Updates
+~~~~~~~
+- `#462 `_ by @paullo0106: Use the correct treatment_idx for fillTree() when applying validation data set
+- `#468 `_ by @jeongyoonlee: Switch the joblib backend for UpliftRandomForestClassifier to threading to avoid memory copy across trees
+
+
+0.12.0 (Jan 2022)
+-----------------
+- CausalML surpassed `637K downloads `_ on PyPI and `2,500 stars `_ on Github!
+- We have 4 new community contributors, Luis (`@lgmoneda `_), Ravi (`@raviksharma `_), Louis (`@LouisHernandez17 `_) and JackRab (`@JackRab `_). Thanks for the contribution!
+- We refactored and speeded up UpliftTreeClassifier/UpliftRandomForestClassifier by 5x with Cython (`#422 `_ `#440 `_ by @jeongyoonlee)
+- We revamped our `API documentation `_, it now includes the latest methodology, references, installation, notebook examples, and graphs! (`#413 `_ by @huigangchen @t-tte @zhenyuz0500 @jeongyoonlee @paullo0106)
+- Our team gave talks at `2021 Conference on Digital Experimentation @ MIT (CODE@MIT) `_, `Causal Data Science Meeting 2021 `_, and `KDD 2021 Tutorials `_ on CausalML introduction and applications. Please take a look if you missed them! Full list of publications and talks can be found here.
+
+Updates
+~~~~~~~
+- Update documentation on Instrument Variable methods @huigangchen (`#447 `_)
+- Add benchmark simulation studies example notebook by @t-tte (`#443 `_)
+- Add sample_weight support for R-learner by @paullo0106 (`#425 `_)
+- Fix incorrect binning of numeric features in UpliftTreeClassifier by @jeongyoonlee (`#420 `_)
+- Update papers, talks, and publication info to README and refs.bib by @zhenyuz0500 (`#410 `_ `#414 `_ `#433 `_)
+- Add instruction for contributing.md doc by @jeongyoonlee (`#408 `_)
+- Fix incorrect feature importance calculation logic by @paullo0106 (`#406 `_)
+- Add parallel jobs support for NearestNeighbors search with n_jobs parameter by @paullo0106 (`#389 `_)
+- Fix bug in simulate_randomized_trial by @jroessler (`#385 `_)
+- Add GA pytest workflow by @ppstacy (`#380 `_)
+
+
+
+0.11.0 (2021-07-28)
+-------------------
+- CausalML surpassed `2K stars `_!
+- We have 3 new community contributors, Jannik (`@jroessler `_), Mohamed (`@ibraaaa `_), and Leo (`@lleiou `_). Thanks for the contribution!
+
+Major Updates
+~~~~~~~~~~~~~
+- Make tensorflow dependency optional and add python 3.9 support by @jeongyoonlee (`#343 `_)
+- Add delta-delta-p (ddp) tree inference approach by @jroessler (`#327 `_)
+- Add conda env files for Python 3.6, 3.7, and 3.8 by @jeongyoonlee (`#324 `_)
+
+Minor Updates
+~~~~~~~~~~~~~
+- Fix inconsistent feature importance calculation in uplift tree by @paullo0106 (`#372 `_)
+- Fix filter method failure with NaNs in the data issue by @manojbalaji1 (`#367 `_)
+- Add automatic package publish by @jeongyoonlee (`#354 `_)
+- Fix typo in unit_selection optimization by @jeongyoonlee (`#347 `_)
+- Fix docs build failure by @jeongyoonlee (`#335 `_)
+- Convert pandas inputs to numpy in S/T/R Learners by @jeongyoonlee (`#333 `_)
+- Require scikit-learn as a dependency of setup.py by @ibraaaa (`#325 `_)
+- Fix AttributeError when passing in Outcome and Effect learner to R-Learner by @paullo0106 (`#320 `_)
+- Fix error when there is no positive class for KL Divergence filter by @lleiou (`#311 `_)
+- Add versions to cython and numpy in setup.py for requirements.txt accordingly by @maccam912 (`#306 `_)
+
+
+
+0.10.0 (2021-02-18)
+-------------------
+- CausalML surpassed `235,000 downloads `_!
+- We have 5 new community contributors, Suraj (`@surajiyer `_), Harsh (`@HarshCasper `_), Manoj (`@manojbalaji1 `_), Matthew (`@maccam912 `_) and Václav (`@vaclavbelak `_). Thanks for the contribution!
+
+Major Updates
+~~~~~~~~~~~~~
+- Add Policy learner, DR learner, DRIV learner by @huigangchen (`#292 `_)
+- Add wrapper for CEVAE, a deep latent-variable and variational autoencoder based model by @ppstacy(`#276 `_)
+
+Minor Updates
+~~~~~~~~~~~~~
+- Add propensity_learner to R-learner by @jeongyoonlee (`#297 `_)
+- Add BaseLearner class for other meta-learners to inherit from without duplicated code by @jeongyoonlee (`#295 `_)
+- Fix installation issue for Shap>=0.38.1 by @paullo0106 (`#287 `_)
+- Fix import error for sklearn>= 0.24 by @jeongyoonlee (`#283 `_)
+- Fix KeyError issue in Filter method for certain dataset by @surajiyer (`#281 `_)
+- Fix inconsistent cumlift score calculation of multiple models by @vaclavbelak (`#273 `_)
+- Fix duplicate values handling in feature selection method by @manojbalaji1 (`#271 `_)
+- Fix the color spectrum of SHAP summary plot for feature interpretations of meta-learners by @paullo0106 (`#269 `_)
+- Add IIA and value optimization related documentation by @t-tte (`#264 `_)
+- Fix StratifiedKFold arguments for propensity score estimation by @paullo0106 (`#262 `_)
+- Refactor the code with string format argument and is to compare object types, and change methods not using bound instance to static methods by @harshcasper (`#256 `_, `#260 `_)
+
+
+
+0.9.0 (2020-10-23)
+------------------
+- CausalML won the 1st prize at the poster session in UberML'20
+- DoWhy integrated CausalML starting v0.4 (`release note `_)
+- CausalML team welcomes new project leadership, Mert Bay
+- We have 4 new community contributors, Mario Wijaya (`@mwijaya3 `_), Harry Zhao (`@deeplaunch `_), Christophe (`@ccrndn `_) and Georg Walther (`@waltherg `_). Thanks for the contribution!
+
+Major Updates
+~~~~~~~~~~~~~
+- Add feature importance and its visualization to UpliftDecisionTrees and UpliftRF by @yungmsh (`#220 `_)
+- Add feature selection example with Filter methods by @paullo0106 (`#223 `_)
+
+Minor Updates
+~~~~~~~~~~~~~
+- Implement propensity model abstraction for common interface by @waltherg (`#223 `_)
+- Fix bug in BaseSClassifier and BaseXClassifier by @yungmsh and @ppstacy (`#217
+