causalkit.eda.CausalEDA#

class causalkit.eda.CausalEDA(data, ps_model=None, n_splits=5, random_state=42)[source]#

Exploratory diagnostics for causal designs with binary treatment.

The class exposes methods to:

Summarize outcome by treatment and naive mean difference.
Estimate cross-validated propensity scores and assess treatment predictability (AUC) and positivity/overlap.
Inspect covariate balance via standardized mean differences (SMD) before/after IPTW weighting; visualize with a love plot.
Inspect weight distributions and effective sample size (ESS).

Parameters:

data (Any)
ps_model (Optional[Any])
n_splits (int)
random_state (int)

__init__(data, ps_model=None, n_splits=5, random_state=42)[source]#

Parameters:

data (Any)
ps_model (Any | None)
n_splits (int)
random_state (int)

Methods

`__init__`(data[, ps_model, n_splits, ...])
`confounders_means`()	Comprehensive confounders balance assessment with means by treatment group.
`confounders_roc_auc`([ps])	Compute ROC AUC of treatment assignment vs.
`data_shape`()	Return the shape information of the causal dataset.
`fit_propensity`()	Estimate cross-validated propensity scores P(T=1\|X).
`outcome_fit`([outcome_model])	Fit a regression model to predict outcome from confounders only.
`outcome_plots`([treatment, target, bins, ...])	Plot the distribution of the outcome for every treatment on one plot, and also produce a boxplot by treatment to visualize outliers.
`outcome_stats`()	Comprehensive outcome statistics grouped by treatment.
`plot_ps_overlap`([ps])	Plot overlaid histograms of propensity scores for treated vs control.
`positivity_check`([ps, bounds])	Check overlap/positivity based on propensity score thresholds.
`treatment_features`()	Return SHAP values from the fitted propensity score model.

__init__(data, ps_model=None, n_splits=5, random_state=42)[source]#

Parameters:

data (Any)
ps_model (Any | None)
n_splits (int)
random_state (int)

data_shape()[source]#

Return the shape information of the causal dataset.

Returns a dict with: - n_rows: number of rows (observations) in the dataset - n_columns: number of columns (features) in the dataset

This provides a quick overview of the dataset dimensions for exploratory analysis and reporting purposes.

Returns:: Dictionary containing ‘n_rows’ and ‘n_columns’ keys with corresponding integer values representing the dataset dimensions.
Return type:: Dict[str, int]

Examples

>>> eda = CausalEDA(causal_data)
>>> shape_info = eda.data_shape()
>>> print(f"Dataset has {shape_info['n_rows']} rows and {shape_info['n_columns']} columns")

outcome_stats()[source]#

Comprehensive outcome statistics grouped by treatment.

Returns a DataFrame with detailed outcome statistics for each treatment group, including count, mean, std, min, various percentiles, and max. This method provides comprehensive outcome analysis and returns data in a clean DataFrame format suitable for reporting.

Returns:: DataFrame with treatment groups as index and the following columns: - count: number of observations in each group - mean: average outcome value - std: standard deviation of outcome - min: minimum outcome value - p10: 10th percentile - p25: 25th percentile (Q1) - median: 50th percentile (median) - p75: 75th percentile (Q3) - p90: 90th percentile - max: maximum outcome value
Return type:: pd.DataFrame

Examples

>>> eda = CausalEDA(causal_data)
>>> stats = eda.outcome_stats()
>>> print(stats)
        count      mean       std       min       p10       p25    median       p75       p90       max
treatment
0        3000  5.123456  2.345678  0.123456  2.345678  3.456789  5.123456  6.789012  7.890123  9.876543
1        2000  6.789012  2.456789  0.234567  3.456789  4.567890  6.789012  8.901234  9.012345  10.987654

fit_propensity()[source]#

Estimate cross-validated propensity scores P(T=1|X).

Uses a preprocessing+CatBoost classifier pipeline with stratified K-fold cross_val_predict to generate out-of-fold probabilities. CatBoost uses all available threads and handles categorical features natively. Returns a PropensityModel instance containing propensity scores and diagnostic methods.

Returns:: A PropensityModel instance with methods for: - roc_auc: ROC AUC score property - shap: SHAP values DataFrame property - ps_graph(): method to plot propensity score overlap - positivity_check(): method to check positivity/overlap
Return type:: PropensityModel

outcome_fit(outcome_model=None)[source]#

Fit a regression model to predict outcome from confounders only.

Uses a preprocessing+CatBoost regressor pipeline with K-fold cross_val_predict to generate out-of-fold predictions. CatBoost uses all available threads and handles categorical features natively. Returns an OutcomeModel instance containing predicted outcomes and diagnostic methods.

The outcome model predicts the baseline outcome from confounders only, excluding treatment. This is essential for proper causal analysis.

Parameters:: outcome_model (Optional[Any]) – Custom regression model to use. If None, uses CatBoostRegressor.
Returns:: An OutcomeModel instance with methods for: - scores: RMSE and MAE regression metrics - shap: SHAP values DataFrame property for outcome prediction
Return type:: OutcomeModel

confounders_roc_auc(ps=None)[source]#

Compute ROC AUC of treatment assignment vs. estimated propensity score.

Interpretation: Higher AUC means treatment is more predictable from confounders, indicating stronger systematic differences between groups (potential confounding). Values near 0.5 suggest random-like assignment.

Return type:: float
Parameters:: ps (ndarray | None)

positivity_check(ps=None, bounds=(0.05, 0.95))[source]#

Check overlap/positivity based on propensity score thresholds.

Returns a dict with: - bounds: (low, high) thresholds used - share_below: fraction with PS < low - share_above: fraction with PS > high - flag: heuristic boolean True if the tails collectively exceed ~2%

Return type:

Dict[str, Any]

Parameters:

ps (ndarray | None)
bounds (Tuple[float, float])

plot_ps_overlap(ps=None)[source]#

Plot overlaid histograms of propensity scores for treated vs control.

Useful to visually assess group overlap. Does not return data; it draws on the current matplotlib figure.

Parameters:: ps (ndarray | None)

confounders_means()[source]#

Comprehensive confounders balance assessment with means by treatment group.

Returns a DataFrame with detailed balance information including: - Mean values of each confounder for control group (treatment=0) - Mean values of each confounder for treated group (treatment=1) - Absolute difference between treatment groups - Standardized Mean Difference (SMD) for formal balance assessment

This method provides a comprehensive view of confounder balance by showing the actual mean values alongside the standardized differences, making it easier to understand both the magnitude and direction of imbalances.

Returns:: DataFrame with confounders as index and the following columns: - mean_t_0: mean value for control group (treatment=0) - mean_t_1: mean value for treated group (treatment=1) - abs_diff: absolute difference abs(mean_t_1 - mean_t_0) - smd: standardized mean difference (Cohen’s d)
Return type:: pd.DataFrame

Notes

SMD values > 0.1 in absolute value typically indicate meaningful imbalance. Categorical variables are automatically converted to dummy variables.

Examples

>>> eda = CausalEDA(causal_data)
>>> balance = eda.confounders_means()
>>> print(balance.head())
             mean_t_0  mean_t_1  abs_diff       smd
confounders
age              29.5      31.2      1.7     0.085
income        45000.0   47500.0   2500.0     0.125
education         0.25      0.35      0.1     0.215

outcome_plots(treatment=None, target=None, bins=30, density=True, alpha=0.5, figsize=(7, 4), sharex=True)[source]#

Plot the distribution of the outcome for every treatment on one plot, and also produce a boxplot by treatment to visualize outliers.

Parameters:

treatment (Optional[str]) – Treatment column name. Defaults to the treatment stored in the CausalEDA data.
target (Optional[str]) – Target/outcome column name. Defaults to the outcome stored in the CausalEDA data.
bins (int) – Number of bins for histograms when the outcome is numeric.
density (bool) – Whether to normalize histograms to form a density.
alpha (float) – Transparency for overlaid histograms.
figsize (tuple) – Figure size for the plots.
sharex (bool) – If True and the outcome is numeric, use the same x-limits across treatments.

Returns:

(fig_distribution, fig_boxplot)

Return type:

Tuple[matplotlib.figure.Figure, matplotlib.figure.Figure]

treatment_features()[source]#

Return SHAP values from the fitted propensity score model.

This method extracts SHAP values from the propensity score model that was trained during fit_propensity(). SHAP values show the directional contribution of each feature to treatment assignment prediction, where positive values increase treatment probability and negative values decrease it.

Returns:

For CatBoost models: DataFrame with columns ‘feature’ and ‘shap_mean’, where ‘shap_mean’ represents the mean SHAP value across all samples. Positive values indicate features that increase treatment probability, negative values indicate features that decrease treatment probability.

For sklearn models: DataFrame with columns ‘feature’ and ‘importance’ (absolute coefficient values, for backward compatibility).

Return type:

pd.DataFrame

Raises:

RuntimeError – If fit_propensity() has not been called yet, or if the fitted model does not support SHAP values extraction.

Examples

>>> eda = CausalEDA(data)
>>> ps = eda.fit_propensity()  # Must be called first
>>> shap_df = eda.treatment_features()
>>> print(shap_df.head())
   feature  shap_mean
0  age         0.45  # Positive: increases treatment prob
1  income     -0.32  # Negative: decreases treatment prob
2  education   0.12  # Positive: increases treatment prob