### Setup Environment and Reproducibility Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/causal_effect.ipynb Initializes the Python environment by loading extensions, setting autoreload, importing necessary libraries, and configuring the device (GPU or CPU). It also sets random seeds for reproducibility. ```python %load_ext autoreload %autoreload 2 import warnings import numpy as np import torch import random import sys import pandas as pd sys.path.append(".."") warnings.filterwarnings('ignore') # ignore warnings device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # Set seeds for reproducibility seed = 82718 random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) ``` -------------------------------- ### Setup CausalPFN Environment Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/calibration.ipynb Initializes the environment for CausalPFN, including loading necessary libraries, setting the device (GPU or CPU), and configuring random seeds for reproducibility. ```python %load_ext autoreload %autoreload 2 import warnings import numpy as np import torch import random import sys sys.path.append("..\n") warnings.filterwarnings('ignore') # ignore warnings device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # Set seeds for reproducibility seed = 82718 random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) ``` -------------------------------- ### Install CausalPFN via PyPI Source: https://github.com/vdblm/causalpfn/blob/main/README.md Install the CausalPFN library using pip. Ensure you have Python 3.10+ and PyTorch 2.3+ installed. ```bash pip install causalpfn ``` -------------------------------- ### Install R Baseline Dependencies Source: https://github.com/vdblm/causalpfn/blob/main/REPRODUCE.md Installs necessary R packages for baseline methods within the R console. ```r install.packages("remotes") remotes::install_version("Matrix", version = "1.5-4") install.packages("grf") install.packages("bartCause") ``` -------------------------------- ### Initialize Environment and Seeds Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/hillstrom_marketing.ipynb Sets up the environment, imports necessary libraries, and initializes random seeds for reproducibility. ```python %load_ext autoreload %autoreload 2 import warnings import numpy as np import torch import random import sys import pandas as pd sys.path.append("..") warnings.filterwarnings('ignore') # ignore warnings device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # Set seeds for reproducibility seed = 82718 random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) ``` -------------------------------- ### Initialize environment and dependencies Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/qini.ipynb Configures the runtime environment, imports necessary libraries, sets random seeds for reproducibility, and prepares the output directory. ```python %load_ext autoreload %autoreload 2 import warnings import numpy as np import torch import random from dotenv import load_dotenv import os import sys from utils import * sys.path.append("..") load_dotenv(override=True) dataset_patterns = dataset_patterns.split(",") method_patterns = method_patterns.split(",") warnings.filterwarnings('ignore') # ignore warnings device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # Set seeds for reproducibility seed = 82718 random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) # make the qini directory if it doesn't exist qini_dir = os.path.join(os.environ["OUTPUT_DIR"], "qini") os.makedirs(qini_dir, exist_ok=True) ``` -------------------------------- ### Initialize Environment and Seeds Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/ablation_dataset_size.ipynb Load environment variables, configure device settings, and set random seeds for reproducibility. ```python %load_ext autoreload %autoreload 2 import warnings import numpy as np import torch import random from dotenv import load_dotenv import os import sys from utils import * sys.path.append("..") load_dotenv(override=True) warnings.filterwarnings("ignore") # ignore warnings device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # Set seeds for reproducibility seed = 82718 random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) ``` -------------------------------- ### Initialize Environment and Load Libraries Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/causal_effect_full.ipynb Initializes the Python environment by loading necessary libraries, setting up the device (GPU or CPU), and configuring warnings. It also sets random seeds for reproducibility. ```python %load_ext autoreload %autoreload 2 import warnings import numpy as np import torch import random from dotenv import load_dotenv import os import sys from utils import * sys.path.append("..") load_dotenv(override=True) warnings.filterwarnings('ignore') # ignore warnings device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") dataset_patterns = dataset_patterns.split(",") method_patterns = method_patterns.split(",") # Set seeds for reproducibility seed = 82718 random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) ``` -------------------------------- ### Load and configure benchmark datasets Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/qini.ipynb Iterates through defined patterns to instantiate specific benchmark datasets with appropriate parameters. ```python %autoreload 2 import re from benchmarks import LentaDataset, CriteoDataset, HillstromDataset, X5Dataset, MegafonDataset # this might take some time to run & load datasets = {} for dataset_pattern in dataset_patterns: if check_match("Hill (1)", dataset_pattern): print("Hit Hill (1)!") datasets["Hill (1)"] = HillstromDataset( seed=seed, n_folds=5, outcome_col="visit", control_arm = "No E-Mail", treatment_arm= "Womens E-Mail", ) if check_match("Hill (2)", dataset_pattern): print("Hit Hill (2)!") datasets["Hill (2)"] = HillstromDataset( seed=seed, n_folds=5, outcome_col="visit", control_arm = "No E-Mail", treatment_arm= "Mens E-Mail", ) if check_match("Criteo RCT", dataset_pattern): print("Hit Criteo RCT!") datasets["Criteo RCT"] = CriteoDataset( seed=seed, n_folds=5, outcome_col="visit", treatment_col="treatment", ) if check_match("Lenta RCT", dataset_pattern): print("Hit Lenta RCT!") datasets["Lenta RCT"] = LentaDataset( seed=seed, n_folds=5, ) if check_match("X5 RCT", dataset_pattern): print("Hit X5 RCT!") datasets["X5 RCT"] = X5Dataset( seed=seed, n_folds=5, ) if check_match("Megafon RCT", dataset_pattern): print("Hit Megafon RCT!") datasets["Megafon RCT"] = MegafonDataset( seed=seed, n_folds=5, ) if check_match("Criteo RCT (Sub)", dataset_pattern): print("Hit Criteo RCT (Sub)!") datasets["Criteo RCT (Sub)"] = CriteoDataset( seed=seed, n_folds=5, outcome_col="visit", treatment_col="treatment", subsample_max_rows=50_000, ) if check_match("X5 RCT (Sub)", dataset_pattern): print("Hit X5 RCT (Sub)!") datasets["X5 RCT (Sub)"] = X5Dataset( seed=seed, n_folds=5, subsample_max_rows=50_000, ) if check_match("Lenta RCT (Sub)", dataset_pattern): print("Hit Lenta RCT (Sub)!") datasets["Lenta RCT (Sub)"] = LentaDataset( seed=seed, n_folds=5, subsample_max_rows=50_000, ) if check_match("Megafon RCT (Sub)", dataset_pattern): print("Hit Megafon RCT (Sub)!") datasets["Megafon RCT (Sub)"] = MegafonDataset( seed=seed, n_folds=5, subsample_max_rows=50_000, ) ``` -------------------------------- ### CausalPFN Custom Model Path and Configuration Source: https://context7.com/vdblm/causalpfn/llms.txt Demonstrates initializing CATEEstimator with custom model paths for loading fine-tuned models or different checkpoints. Advanced configuration options are available. ```python import numpy as np import torch from causalpfn import CATEEstimator device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") ``` -------------------------------- ### Initialize CATEEstimator with Custom Configuration Source: https://context7.com/vdblm/causalpfn/llms.txt Initialize CATEEstimator with specific settings for model path, caching, context/query lengths, and calibration parameters. This is useful for fine-tuning the model's behavior and resource usage. ```python import numpy as np import torch from causalpfn import CATEEstimator device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") cate_estimator = CATEEstimator( device=device, # Model source (default: "vdblm/causalpfn" from Hugging Face) model_path="vdblm/causalpfn", # Cache directory for downloaded models cache_dir="~/.cache/causalpfn", # Memory management for large datasets max_context_length=4096, # Maximum context samples per query batch max_query_length=4096, # Maximum query samples per batch num_neighbours=1024, # Nearest neighbors for stratification # Calibration settings calibrate=True, n_folds=3, # Cross-validation folds for calibration calibrate_T_min=0.001, # Minimum temperature search bound calibrate_T_max=10.0, # Maximum temperature search bound calibrate_T_size=500, # Number of temperature values to search calibrate_T_batch_size=50, # Uncertainty quantification settings ICE_n_bins=10, # Bins for integrated coverage error ICE_n_samples=1000, # Samples for ICE calculation verbose=True, ) # Generate sample data np.random.seed(42) n = 5000 X = np.random.normal(0, 1, size=(n, 10)).astype(np.float32) T = np.random.binomial(1, 0.5, size=n).astype(np.float32) Y = X[:, 0] + T * (0.5 + X[:, 1]) + np.random.normal(0, 0.1, size=n).astype(np.float32) # Fit and estimate cate_estimator.fit(X, T, Y) cate = cate_estimator.estimate_cate(X[:100]) print(f"Estimated CATE shape: {cate.shape}") print(f"Calibrated temperature: {cate_estimator.temperature:.4f}") ``` -------------------------------- ### Create and Activate Conda Environment Source: https://github.com/vdblm/causalpfn/blob/main/REPRODUCE.md Initializes the environment required for running baselines in R and CATENet. ```bash conda env create -f env_reproduce.yaml conda activate reproduce ``` -------------------------------- ### Configure Benchmark Parameters Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/causal_effect_full.ipynb Sets up parameters for running causal inference benchmarks, such as dataset and method patterns, and whether to replace existing runs. ```python # Parameters dataset_patterns = "*" # all of the datasets that are being run method_patterns = "*" # all of the methods to run index = None replace_runs = False # whether to replace existing runs ``` -------------------------------- ### Load Datasets and Baselines for Comparison Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/causal_effect.ipynb Loads datasets and required functions for causal inference benchmarks. Initializes various baseline models (e.g., X-Learner, S-Learner, T-Learner) with and without hyperparameter optimization (HPO). ```python # Load datasets and required functions %autoreload 2 from benchmarks import IHDPDataset, ACIC2016Dataset from benchmarks import RealCauseLalondeCPSDataset, RealCauseLalondePSIDDataset import time from causalpfn import ATEEstimator, CATEEstimator from benchmarks.base import CATE_Dataset, ATE_Dataset from benchmarks.baselines import ( TLearnerBaseline, SLearnerBaseline, XLearnerBaseline, BaselineModel ) from causalpfn.evaluation import calculate_pehe from tqdm import tqdm # Get different realizations for each dataset (only the first two realization - you can change `n_tables`) datasets = { "IHDP": IHDPDataset(n_tables=2), "ACIC 2016": ACIC2016Dataset(n_tables=2), "RealCause Lalonde CPS": RealCauseLalondeCPSDataset(n_tables=2), "RealCause Lalonde PSID": RealCauseLalondePSIDDataset(n_tables=2), } # get all of the baselines to compare with (not exhaustive -- feel free to comment out some) baselines = { "X-Learner (no HPO)": XLearnerBaseline(hpo=False), "S-Learner (no HPO)": SLearnerBaseline(hpo=False), "T-Learner (no HPO)": TLearnerBaseline(hpo=False), ################################################ # Ucomment the following lines to run with HPO # ################################################ # "X-Learner (HPO)": XLearnerBaseline(hpo=True), # "S-Learner (HPO)": SLearnerBaseline(hpo=True), # "T-Learner (HPO)": TLearnerBaseline(hpo=True), } # Initialize results DataFrame results = pd.DataFrame(columns=["dataset", "realization", "method", "ate_rel_err", "cate_pehe", "ate_time", "cate_time"]) ``` -------------------------------- ### Complete Uplift Modeling Pipeline Source: https://context7.com/vdblm/causalpfn/llms.txt Demonstrates a full uplift modeling pipeline using CATEEstimator for marketing applications. It includes simulating customer data, training the model, predicting uplift, and evaluating a targeting strategy. ```python import numpy as np import torch from causalpfn import CATEEstimator device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # Simulate marketing campaign data np.random.seed(42) n_customers = 10000 # Customer features: age, income, engagement_score, tenure, purchase_history X = np.column_stack([ np.random.normal(40, 15, n_customers), # age np.random.lognormal(10.5, 0.5, n_customers), # income np.random.beta(2, 5, n_customers) * 100, # engagement_score np.random.exponential(24, n_customers), # tenure_months np.random.poisson(5, n_customers), # past_purchases ]).astype(np.float32) # True heterogeneous treatment effect (campaign effectiveness varies by customer) true_uplift = ( 0.1 * (X[:, 2] / 100) + # Higher engagement = more responsive 0.05 * np.log(X[:, 4] + 1) - # More past purchases = more responsive 0.02 * (X[:, 0] - 40) / 15 # Younger customers slightly more responsive ).astype(np.float32) # Treatment assignment and outcomes (conversion rate) T = np.random.binomial(1, 0.5, n_customers).astype(np.float32) base_conversion = 0.05 + 0.1 * (X[:, 2] / 100) Y = (base_conversion + T * true_uplift + np.random.normal(0, 0.02, n_customers)).astype(np.float32) Y = np.clip(Y, 0, 1) # Ensure valid probability range # Train/test split train_size = int(0.7 * n_customers) train_idx = np.random.choice(n_customers, train_size, replace=False) test_idx = np.setdiff1d(np.arange(n_customers), train_idx) # Fit uplift model uplift_model = CATEEstimator(device=device, verbose=True) uplift_model.fit(X[train_idx], T[train_idx], Y[train_idx]) # Predict uplift for test customers predicted_uplift = uplift_model.estimate_cate(X[test_idx]) # Targeting strategy: treat top 20% by predicted uplift n_to_treat = int(0.2 * len(test_idx)) top_indices = np.argsort(predicted_uplift)[-n_to_treat:] # Evaluate targeting performance mean_uplift_targeted = true_uplift[test_idx][top_indices].mean() mean_uplift_random = true_uplift[test_idx].mean() print(f"Mean uplift (targeted top 20%): {mean_uplift_targeted:.4f}") print(f"Mean uplift (random selection): {mean_uplift_random:.4f}") print(f"Targeting lift: {mean_uplift_targeted / mean_uplift_random:.2f}x") ``` -------------------------------- ### Estimate CATE with CATEEstimator Source: https://context7.com/vdblm/causalpfn/llms.txt Demonstrates initializing the CATEEstimator, fitting it on synthetic observational data, and calculating the PEHE metric. ```python import numpy as np import torch from causalpfn import CATEEstimator # Initialize device device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # Generate synthetic observational data np.random.seed(42) n, d = 10000, 5 X = np.random.normal(0, 1, size=(n, d)).astype(np.float32) # True CATE function: tau(x) = sin(x_0) + 0.5 * x_1 true_cate = np.sin(X[:, 0]) + 0.5 * X[:, 1] # Generate treatment assignment and outcomes T = np.random.binomial(1, p=0.5, size=n).astype(np.float32) Y0 = X[:, 0] - X[:, 1] + np.random.normal(0, 0.1, size=n).astype(np.float32) Y1 = Y0 + true_cate Y = Y0 * (1 - T) + Y1 * T # Split into train/test train_idx = np.random.choice(n, size=int(0.7 * n), replace=False) test_idx = np.setdiff1d(np.arange(n), train_idx) X_train, X_test = X[train_idx], X[test_idx] T_train, Y_train = T[train_idx], Y[train_idx] true_cate_test = true_cate[test_idx] # Initialize and fit the CATE estimator cate_estimator = CATEEstimator( device=device, verbose=True, calibrate=False, # Set True for calibrated confidence intervals max_context_length=4096, max_query_length=4096, ) cate_estimator.fit(X_train, T_train, Y_train) # Estimate CATE for test data cate_predictions = cate_estimator.estimate_cate(X_test) # Calculate PEHE (Precision in Estimation of Heterogeneous Effect) pehe = np.sqrt(np.mean((cate_predictions - true_cate_test) ** 2)) print(f"PEHE: {pehe:.4f}") # Output: PEHE: 0.1523 ``` -------------------------------- ### Configure Benchmarking Parameters Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/ablation_dataset_size.ipynb Define global variables for test size and run management. ```python # Parameters test_size = 100 num_tables_per_setting = 50 replace_runs = False # whether to replace existing runs ``` -------------------------------- ### Import Baseline Models Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/causal_effect_full.ipynb Imports various baseline models for causal inference benchmarking, categorized into Base, EconML, and CATE Net implementations. ```python # Baselines (Base) from benchmarks.baselines import BaselineModel # Baselines (EconML) from benchmarks.baselines import ( ForestDMLBaseline, TLearnerBaseline, SLearnerBaseline, XLearnerBaseline, DALearnerBaseline, XLearnerBaseline, ForestDRLearnerBaseline, ) # Baselines (CATE Net) from benchmarks.baselines import ( TarNetBaseline, DragonNetBaseline, RANetBaseline, SNetBaseline, FlexTENetBaseline, XNetBaseline, ) ``` -------------------------------- ### Run CausalPFN and X-Learner Estimation Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/hillstrom_marketing.ipynb Executes CATE estimation using 5-fold honest splitting to compare CausalPFN and X-Learner performance. ```python %autoreload 2 from causalpfn.causal_estimator import CATEEstimator from benchmarks.base import Qini_Dataset from tqdm import tqdm from benchmarks.baselines import XLearnerBaseline datasets = [hill1, hill2] results = [ pd.DataFrame(columns=['y', 't', 'effect_causalpfn', 'effect_xlearner']), pd.DataFrame(columns=['y', 't', 'effect_causalpfn', 'effect_xlearner']), ] pbar = tqdm( total = sum([len(dset) for dset in datasets]), desc = "Running CausalPFN and X-Learner", ) for dset_idx in range(len(datasets)): dset: HillstromDataset = datasets[dset_idx] for fold_idx in range(len(dset)): # fit on the rest of the data and get estimates for the current fold (honest splitting) qini_data: Qini_Dataset = dset[fold_idx] train_size = qini_data.X_train.shape[0] test_size = qini_data.X_test.shape[0] # run causalpfn causalpfn = CATEEstimator( device=device, max_context_length=train_size, # set to a large value to fit everything in the context ) causalpfn.fit(X=qini_data.X_train, y=qini_data.y_train, t=qini_data.t_train) causalpfn_tau_hat = causalpfn.estimate_cate(X=qini_data.X_test) # run an x-learner xlearner = XLearnerBaseline(hpo=False) xlearner_tau_hat = xlearner.estimate_cate( X_train=qini_data.X_train, y_train=qini_data.y_train, t_train=qini_data.t_train, X_test=qini_data.X_test, ) # add the fold results[dset_idx] = pd.concat([results[dset_idx], pd.DataFrame({ 'y': qini_data.y_test, 't': qini_data.t_test, 'effect_causalpfn': causalpfn_tau_hat, 'effect_xlearner': xlearner_tau_hat, })], ignore_index=True) pbar.update(1) pbar.close() ``` -------------------------------- ### Train and Estimate CATE with CausalPFN Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/qini.ipynb Implements the training and CATE estimation process using the CausalPFN model. It iterates through datasets, initializes the CATEEstimator, fits the model on training data, and estimates CATE on test data. Results, including estimated effects and time spent, are saved. ```python from src.causalpfn import CATEEstimator from benchmarks.base import Qini_Dataset import time from tqdm import tqdm pbar = tqdm( total=sum([len(dataset) for dataset in datasets.values()]), desc="CausalPFN", ) for dataset_name, dataset in datasets.items(): pbar.set_postfix(dataset=dataset_name) for i in range(len(dataset)): # dataset_name: str, method_name: str, all_method_patterns: list, all_datasets_patterns: list, idx: int, artifact_dir: str, replace: bool = False with result_saver( dataset_name=dataset_name, method_name="CausalPFN", all_method_patterns=method_patterns, all_datasets_patterns=dataset_patterns, idx=i, artifact_dir=qini_dir, replace=replace_runs, ) as result: if result is not None: qini_data: Qini_Dataset = dataset[i] time_start = time.time() cate_estimator = CATEEstimator( device=device, max_context_length=max_context_length, ) cate_estimator.fit(X=qini_data.X_train, y=qini_data.y_train, t=qini_data.t_train) estimated_tau = cate_estimator.estimate_cate(X=qini_data.X_test) time_spent = time.time() - time_start result["estimated_effect"] = estimated_tau result["t"] = qini_data.t_test result["y"] = qini_data.y_test result["time_spent"] = time_spent / ((len(qini_data.X_test) + len(qini_data.X_train)) / 1000) pbar.update(1) pbar.close() ``` -------------------------------- ### Initialize Causal Inference Baselines Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/causal_effect_full.ipynb Defines a dictionary of various causal inference baseline models, including T-Learner, S-Learner, X-Learner, DA-Learner, Forest DR, Forest DML, DragonNet, TarNet, SNet, FlexTENet, XNET, RA Net, GRF, BART, and IPW. Each baseline can be configured with or without Hyperparameter Optimization (HPO). ```python from benchmarks.baselines import GRFBaseline, BartBaseline, IPWBaseline baselines = { "T Learner (no HPO)": TLearnerBaseline(hpo=False), "T Learner (HPO)": TLearnerBaseline(hpo=True), "S Learner (no HPO)": SLearnerBaseline(hpo=False), "S Learner (HPO)": SLearnerBaseline(hpo=True), "X Learner (no HPO)": XLearnerBaseline(hpo=False), "X Learner (HPO)": XLearnerBaseline(hpo=True), "DA Learner (no HPO)": DALearnerBaseline(hpo=False), "DA Learner (HPO)": DALearnerBaseline(hpo=True), "Forest DR Learner (no HPO)": ForestDRLearnerBaseline(hpo=False), "Forest DR Learner (HPO)": ForestDRLearnerBaseline(hpo=True), "Forest DML (no HPO)": ForestDMLBaseline(hpo=False), "Forest DML (HPO)": ForestDMLBaseline(hpo=True), "DragonNet (no HPO)": DragonNetBaseline(hpo=False), "DragonNet (HPO)": DragonNetBaseline(hpo=True), "TarNet (no HPO)": TarNetBaseline(hpo=False), "TarNet (HPO)": TarNetBaseline(hpo=True), "SNet (no HPO)": SNetBaseline(hpo=False), "SNet (HPO)": SNetBaseline(hpo=True), "FlexTENet (no HPO)": FlexTENetBaseline(hpo=False, batch_size=256), "FlexTENet (HPO)": FlexTENetBaseline(hpo=True, batch_size=256), "XNET (no HPO)": XNetBaseline(hpo=False, batch_size=256), "XNET (HPO)": XNetBaseline(hpo=True, batch_size=256), "RA Net (no HPO)": RANetBaseline(hpo=False), "RA Net (HPO)": RANetBaseline(hpo=True), "GRF (no HPO)": GRFBaseline(hpo=False), "GRF (HPO)": GRFBaseline(hpo=True), "BART": BartBaseline(hpo=False), "IPW (no HPO)": IPWBaseline(hpo=False), "IPW (HPO)": IPWBaseline(hpo=True), } ``` -------------------------------- ### Import Baseline Models Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/ablation_dataset_size.ipynb Imports for baseline model evaluation. ```python from src.causalpfn import calculate_pehe from tqdm import tqdm # Baselines (Base) from benchmarks.baselines import BaselineModel ``` -------------------------------- ### CATEEstimator - Fit and Estimate CATE Source: https://context7.com/vdblm/causalpfn/llms.txt Demonstrates how to initialize, fit, and use the CATEEstimator to predict Conditional Average Treatment Effects (CATE) on new data. ```APIDOC ## CATEEstimator The `CATEEstimator` class estimates Conditional Average Treatment Effects (CATE), which represent the heterogeneous treatment effect for individual observations. It automatically downloads the pre-trained CausalPFN model from Hugging Face Hub, fits on observational data with covariates, treatments, and outcomes, and provides point estimates along with calibrated confidence intervals for uncertainty quantification. ### Method `fit(X, T, Y)` ### Parameters - **X** (np.ndarray) - Required - Covariate data. - **T** (np.ndarray) - Required - Treatment assignment data. - **Y** (np.ndarray) - Required - Outcome data. ### Method `estimate_cate(X)` ### Parameters - **X** (np.ndarray) - Required - Covariate data for prediction. ### Request Example ```python import numpy as np import torch from causalpfn import CATEEstimator # Initialize device device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # Generate synthetic observational data np.random.seed(42) n, d = 10000, 5 X = np.random.normal(0, 1, size=(n, d)).astype(np.float32) # True CATE function: tau(x) = sin(x_0) + 0.5 * x_1 true_cate = np.sin(X[:, 0]) + 0.5 * X[:, 1] # Generate treatment assignment and outcomes T = np.random.binomial(1, p=0.5, size=n).astype(np.float32) Y0 = X[:, 0] - X[:, 1] + np.random.normal(0, 0.1, size=n).astype(np.float32) Y1 = Y0 + true_cate Y = Y0 * (1 - T) + Y1 * T # Split into train/test train_idx = np.random.choice(n, size=int(0.7 * n), replace=False) test_idx = np.setdiff1d(np.arange(n), train_idx) X_train, X_test = X[train_idx], X[test_idx] T_train, Y_train = T[train_idx], Y[train_idx] true_cate_test = true_cate[test_idx] # Initialize and fit the CATE estimator cate_estimator = CATEEstimator( device=device, verbose=True, calibrate=False, # Set True for calibrated confidence intervals max_context_length=4096, max_query_length=4096, ) cate_estimator.fit(X_train, T_train, Y_train) # Estimate CATE for test data cate_predictions = cate_estimator.estimate_cate(X_test) # Calculate PEHE (Precision in Estimation of Heterogeneous Effect) pehe = np.sqrt(np.mean((cate_predictions - true_cate_test) ** 2)) print(f"PEHE: {pehe:.4f}") ``` ### Response Example ```json { "example": "PEHE: 0.1523" } ``` ``` -------------------------------- ### Prepare Causal Inference Results for Visualization Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/causal_effect_full.ipynb Maps internal method names to display-friendly labels for plotting performance benchmarks. ```python import matplotlib.patches as mpatches df = results_df.copy() # replace method names for better visualization method_name_map = { "CausalPFN": "CausalPFN", "T Learner (HPO)": "T-Learner", "S Learner (HPO)": "S-Learner", "X Learner (HPO)": "X-Learner", "DA Learner (HPO)": "DA-Learner", "Forest DR Learner (HPO)": "DR-Learner", "Forest DML (HPO)": "Forest DML", "DragonNet (HPO)": "DragonNet", "TarNet (HPO)": "TarNet", "RA Net (HPO)": "RA-Net", "GRF (HPO)": "GRF", "BART": "BART", } df["method"] = df["method"].map(method_name_map) ``` -------------------------------- ### Import Matplotlib Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/ablation_dataset_size.ipynb Standard import for plotting library. ```python import matplotlib.pyplot as plt ``` -------------------------------- ### Run EconML Baselines Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/ablation_dataset_size.ipynb Executes various EconML learners and saves results using a result_saver context manager. ```python from benchmarks.baselines import ( TLearnerBaseline, SLearnerBaseline, XLearnerBaseline, DALearnerBaseline, ) baselines = { "T Learner (no HPO)": TLearnerBaseline(hpo=False), "S Learner (no HPO)": SLearnerBaseline(hpo=False), "X Learner (no HPO)": XLearnerBaseline(hpo=False), "DA Learner (no HPO)": DALearnerBaseline(hpo=False), } for datasets, path in zip( [sample_size_datasets, feature_size_datasets], [sample_size_ablation_path, feature_size_ablation_path], ): pbar = tqdm( total=sum([len(dataset) * len(baselines) for dataset in datasets.values()]), desc="Baselines", ) for dataset_name, dataset in datasets.items(): for baseline_name, baseline in baselines.items(): pbar.set_postfix(dataset=dataset_name, baseline=baseline_name) for i in range(len(dataset)): with result_saver( dataset_name=dataset_name, method_name=baseline_name, all_method_patterns="*", all_datasets_patterns="*", idx=i, artifact_dir=path, replace=replace_runs, ) as result: if result is not None: baseline: BaselineModel cate_dset, ate_dset = dataset[i] # CATE cate_pred = baseline.estimate_cate( X_train=cate_dset.X_train, t_train=cate_dset.t_train, y_train=cate_dset.y_train, X_test=cate_dset.X_test, ) pehe = calculate_pehe(cate_true=cate_dset.true_cate, cate_pred=cate_pred) result["pehe"] = pehe pbar.update(1) pbar.close() ``` -------------------------------- ### Iterate through datasets and realizations for CausalPFN evaluation Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/causal_effect.ipynb Iterates through datasets to compute ATE and CATE using CausalPFN and baseline models, tracking progress with tqdm. ```python pbar = tqdm( total=sum(len(dataset) * (1 + len(baselines)) for dataset in datasets.values()), desc="Processing datasets", ) for dataset_name, dataset in datasets.items(): for realization_idx in range(len(dataset)): pbar.set_postfix({"dataset": dataset_name, "method": "CausalPFN"}) res = dataset[realization_idx] cate_dset: CATE_Dataset = res[0] ate_dset: ATE_Dataset = res[1] # run CausalPFN estimator for ATE start_time = time.time() causalpfn_ate = ATEEstimator( device=device, ) causalpfn_ate.fit(ate_dset.X, ate_dset.t, ate_dset.y) true_ate = ate_dset.true_ate causalpfn_ate_hat = causalpfn_ate.estimate_ate() causalpfn_rel_error = abs(causalpfn_ate_hat - true_ate) / abs(true_ate) ate_time = time.time() - start_time # run CausalPFN estimator for CATE start_time = time.time() causalpfn_cate = CATEEstimator( device=device, ) causalpfn_cate.fit(cate_dset.X_train, cate_dset.t_train, cate_dset.y_train) causalpfn_cate_hat = causalpfn_cate.estimate_cate(cate_dset.X_test) cate_pehe = calculate_pehe(cate_dset.true_cate, causalpfn_cate_hat) cate_time = time.time() - start_time # add results for CausalPFN row = dict( dataset=dataset_name, realization=realization_idx, method="CausalPFN", ate_rel_err=round(causalpfn_rel_error, 2), cate_pehe=round(cate_pehe, 2), ate_time=round(ate_time / (ate_dset.X.shape[0] + ate_dset.X.shape[0]) * 100, 2), cate_time=round(cate_time / (cate_dset.X_train.shape[0] + cate_dset.X_test.shape[0]) * 100, 2), ) pbar.update(1) results = pd.concat([results, pd.DataFrame([row])], ignore_index=True) for method_name, baseline in baselines.items(): pbar.set_postfix({"dataset": dataset_name, "method": method_name}) baseline: BaselineModel # run baseline estimator for ATE start_time = time.time() ate_pred = baseline.estimate_ate(X=ate_dset.X, t=ate_dset.t, y=ate_dset.y) rel_err = np.abs(ate_pred - true_ate) / np.abs(true_ate) ate_time = time.time() - start_time # run baseline estimator for CATE start_time = time.time() cate_pred = baseline.estimate_cate(X_train=cate_dset.X_train, t_train=cate_dset.t_train, y_train=cate_dset.y_train, X_test=cate_dset.X_test) cate_pehe = calculate_pehe(cate_dset.true_cate, cate_pred) cate_time = time.time() - start_time # add results for baseline row = dict( dataset=dataset_name, realization=realization_idx, method=method_name, ate_rel_err=round(rel_err, 2), cate_pehe=round(cate_pehe, 2), ate_time=round(ate_time / (ate_dset.X.shape[0] + ate_dset.X.shape[0]) * 100, 2), cate_time=round(cate_time / (cate_dset.X_train.shape[0] + cate_dset.X_test.shape[0]) * 100, 2), ) pbar.update(1) results = pd.concat([results, pd.DataFrame([row])], ignore_index=True) ``` -------------------------------- ### Define evaluation parameters Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/qini.ipynb Sets the configuration variables for dataset selection, method patterns, and execution constraints. ```python # Parameters dataset_patterns = "Hill (1),Hill (2),*(Sub)*" # ,*(Sub)* method_patterns = "CausalPFN,T Learner,S Learner,X Learner,DA Learner,Forest DR Learner" max_context_length = 50_000 replace_runs = False # whether to replace existing runs ``` -------------------------------- ### Configure Output Directory Environment Variable Source: https://github.com/vdblm/causalpfn/blob/main/REPRODUCE.md Sets the output directory for logging and saving results using python-dotenv. ```bash # the directory where everything will be logged and saved with git ignoring it dotenv set OUTPUT_DIR ``` -------------------------------- ### Summarize estimation metrics Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/causal_effect.ipynb Aggregates ATE and CATE estimation times by calculating the mean over realizations using a pivot table. ```python # summarize all of the ATE and CATE estimation times by averaging over realizations time_spent_df = ( results.pivot_table( index="method", # rows: one per method columns="dataset", # multi‐columns: first level will be dataset values=["ate_time", "cate_time"], # the values to aggregate aggfunc="mean", # take the mean over realizations ) .swaplevel(0, 1, axis=1) .sort_index(axis=1, level=0) ) ``` -------------------------------- ### Progress Bar Output Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/causal_effect_full.ipynb Displays the progress of the baseline benchmarking process, showing the completion percentage, iteration count, and current dataset and baseline being processed. ```text Baselines: 100%|██████████| 8990/8990 [00:02<00:00, 4393.65it/s, baseline=IPW (HPO), dataset=RealCause Lalonde PSID] ``` -------------------------------- ### Execute Ablation Studies Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/ablation_dataset_size.ipynb Generate datasets for varying sample sizes and feature dimensions to evaluate model performance. ```python # the effect of sample size on the performance of the methods sample_size_ablation_path = os.path.join(os.environ["OUTPUT_DIR"], "ablation/sample_size/") os.makedirs(sample_size_ablation_path, exist_ok=True) NUM_SAMPLES = [10, 20, 50, 100, 200, 500, 1000, 5000, 10000] X_DIM = 10 sample_size_datasets = { f"polynomial_{sample_size}": get_datasets( sample_size=sample_size, x_dim=X_DIM, test_size=test_size, n_tables=num_tables_per_setting, ) for sample_size in NUM_SAMPLES } # the effect of feature size on the performance of the methods feature_size_ablation_path = os.path.join(os.environ["OUTPUT_DIR"], "ablation/feature_size/") os.makedirs(feature_size_ablation_path, exist_ok=True) NUM_SAMPLES = 1000 X_DIMS = [1, 5, 10, 20, 50, 100, 200, 500, 1000] feature_size_datasets = { f"polynomial_features_{x_dim}": get_datasets( sample_size=NUM_SAMPLES, x_dim=x_dim, test_size=test_size, n_tables=num_tables_per_setting, ) for x_dim in X_DIMS } ``` -------------------------------- ### Import Baseline Model Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/qini.ipynb Imports the BaselineModel class from the benchmarks.baselines module, typically used for comparison against other causal inference methods. ```python # Baselines (Base) from benchmarks.baselines import BaselineModel ``` -------------------------------- ### Run CausalPFN Evaluation Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/ablation_dataset_size.ipynb Iterate through datasets to calculate PEHE metrics using the CATEEstimator. ```python from src.causalpfn import CATEEstimator, calculate_pehe from tqdm import tqdm for datasets, path in zip( [sample_size_datasets, feature_size_datasets], [sample_size_ablation_path, feature_size_ablation_path], ): pbar = tqdm( total=sum([len(dataset) for dataset in datasets.values()]), desc="CausalPFN", ) for dataset_name, dataset in datasets.items(): pbar.set_postfix(dataset=dataset_name) for i in range(len(dataset)): with result_saver( dataset_name=dataset_name, method_name="CausalPFN", all_method_patterns="*", all_datasets_patterns="*", idx=i, artifact_dir=path, replace=replace_runs, ) as result: if result is not None: cate_dset, ate_dset = dataset[i] # CATE cate_estimator = CATEEstimator( device=device, ) cate_estimator.fit(cate_dset.X_train, cate_dset.t_train, cate_dset.y_train) estimated_cate = cate_estimator.estimate_cate(X=cate_dset.X_test) pehe = calculate_pehe(cate_pred=estimated_cate, cate_true=cate_dset.true_cate) result["pehe"] = pehe pbar.update(1) pbar.close() ``` -------------------------------- ### Run Causal Inference Baselines Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/qini.ipynb Iterates through datasets and baselines to estimate conditional average treatment effects (CATE). Requires pre-defined datasets, baselines, and result saving mechanisms. ```python from benchmarks.baselines import ( TLearnerBaseline, SLearnerBaseline, XLearnerBaseline, DALearnerBaseline, XLearnerBaseline, ForestDRLearnerBaseline, ) baselines = { "T Learner": TLearnerBaseline(hpo=False), "S Learner": SLearnerBaseline(hpo=False), "X Learner": XLearnerBaseline(hpo=False), "DA Learner": DALearnerBaseline(hpo=False), "Forest DR Learner": ForestDRLearnerBaseline(hpo=False), } pbar = tqdm( total=sum([len(dataset) * len(baselines) for dataset in datasets.values()]), desc="Baselines", ) for dataset_name, dataset in datasets.items(): for baseline_name, baseline in baselines.items(): pbar.set_postfix(dataset=dataset_name, baseline=baseline_name) for i in range(len(dataset)): with result_saver( dataset_name=dataset_name, method_name=baseline_name, all_method_patterns=method_patterns, all_datasets_patterns=dataset_patterns, idx=i, artifact_dir=qini_dir, replace=replace_runs, ) as result: if result is not None: baseline: BaselineModel qini_data: Qini_Dataset = dataset[i] time_start = time.time() estimated_tau = baseline.estimate_cate( X_train=qini_data.X_train, t_train=qini_data.t_train, y_train=qini_data.y_train, X_test=qini_data.X_test, ) time_spent = time.time() - time_start result["estimated_effect"] = estimated_tau result["t"] = qini_data.t_test result["y"] = qini_data.y_test result["time_spent"] = time_spent / ((len(qini_data.X_test) + len(qini_data.X_train)) / 1000) pbar.update(1) pbar.close() ``` -------------------------------- ### Dataset Generation Helper Source: https://github.com/vdblm/causalpfn/blob/main/notebooks/ablation_dataset_size.ipynb Define a function to load polynomial datasets with specified sample and feature dimensions. ```python # Load the polynomial dataset and required functions %autoreload 2 from benchmarks import PolynomialDataset def get_datasets( sample_size: int, x_dim: int, test_size: int = test_size, n_tables: int = num_tables_per_setting, ): effective_sample_size = sample_size + test_size test_ratio = test_size / effective_sample_size datasets = PolynomialDataset( n_tables=n_tables, n_samples=effective_sample_size, test_ratio=test_ratio, x_dim_dist=lambda: x_dim, seed=42, ) return datasets ```