Plotting

plotting

plot_outcome_proba_dist

plot_outcome_proba_dist(df: DataFrame, outcome_proba_col: str, treatment_col: str, xlabel: str = 'Predicted Outcome Probability', title: str = 'Outcome Probability Distribution', bin_edges: ndarray = None, normalize: bool = False, fig: Figure = None, ax: Axes = None, figsize: tuple = (10, 6))

Plot a predicted-outcome probability distribution for treatment vs. control groups. E.g., if 'outcome_proba_col' stores model-predicted probabilities.

Source code in CausalEstimate/vis/plotting.py

def plot_outcome_proba_dist(
    df: pd.DataFrame,
    outcome_proba_col: str,
    treatment_col: str,
    xlabel: str = "Predicted Outcome Probability",
    title: str = "Outcome Probability Distribution",
    bin_edges: np.ndarray = None,
    normalize: bool = False,
    fig: plt.Figure = None,
    ax: plt.Axes = None,
    figsize: tuple = (10, 6),
):
    """
    Plot a predicted-outcome probability distribution for treatment vs. control groups.
    E.g., if 'outcome_proba_col' stores model-predicted probabilities.
    """
    return plot_hist_by_groups(
        df=df,
        value_col=outcome_proba_col,
        group_col=treatment_col,
        group_values=(0, 1),
        group_labels=("Control", "Treatment"),
        bin_edges=bin_edges,
        normalize=normalize,
        xlabel=xlabel,
        title=title,
        fig=fig,
        ax=ax,
        figsize=figsize,
    )

plot_propensity_score_dist

plot_propensity_score_dist(df: DataFrame, ps_col: str, treatment_col: str, xlabel: str = 'Propensity Score', title: str = 'Propensity Score Distribution', bin_edges: ndarray = None, normalize: bool = False, fig: Figure = None, ax: Axes = None, figsize: tuple = (10, 6))

Plot a propensity score distribution for treatment and control groups.

Source code in CausalEstimate/vis/plotting.py

def plot_propensity_score_dist(
    df: pd.DataFrame,
    ps_col: str,
    treatment_col: str,
    xlabel: str = "Propensity Score",
    title: str = "Propensity Score Distribution",
    bin_edges: np.ndarray = None,
    normalize: bool = False,
    fig: plt.Figure = None,
    ax: plt.Axes = None,
    figsize: tuple = (10, 6),
):
    """
    Plot a propensity score distribution for treatment and control groups.
    """
    return plot_hist_by_groups(
        df=df,
        value_col=ps_col,
        group_col=treatment_col,
        group_values=(0, 1),
        group_labels=("Control", "Treatment"),
        bin_edges=bin_edges,
        normalize=normalize,
        xlabel=xlabel,
        title=title,
        fig=fig,
        ax=ax,
        figsize=figsize,
    )

plot_hist_by_groups

plot_hist_by_groups(df: DataFrame, value_col: str, group_col: str, group_values=(0, 1), group_labels=('Group 0', 'Group 1'), bin_edges=None, normalize: bool = False, xlabel: str = None, title: str = None, alpha: float = 0.5, colors=('#1F77B4', '#D62728'), fig: Figure = None, ax: Axes = None, figsize: tuple = (10, 6)) -> Tuple[plt.Figure, plt.Axes]

A generic helper that plots a histogram of 'value_col' for two groups defined by 'group_col', e.g. group_col=0 vs. group_col=1.

Parameters:

Name	Type	Description	Default
df	DataFrame	DataFrame containing the data.	required
value_col	str	The column whose distribution we want to plot.	required
group_col	str	The column that indicates group membership.	required
group_values	tuple	The two distinct values used to split the DataFrame.	(0, 1)
group_labels	tuple	Labels for legend (e.g. "Control", "Treatment").	('Group 0', 'Group 1')
bin_edges	array	The bin edges for histogram. If None, defaults to 50 bins from 0..1	None
normalize	bool	Whether to normalize the histogram (density=True).	False
xlabel	str	X-axis label.	None
title	str	Plot title.	None
alpha	float	Transparency for the histogram overlay.	0.5
colors	tuple	Colors for the two histograms.	('#1F77B4', '#D62728')
fig, ax		If provided, plot into them; otherwise create new figure/axes.	required
figsize	tuple	Size of figure if we create a new one.	(10, 6)

Returns:

Type	Description
Tuple[Figure, Axes]	(fig, ax)

Source code in CausalEstimate/vis/plotting.py

def plot_hist_by_groups(
    df: pd.DataFrame,
    value_col: str,
    group_col: str,
    group_values=(0, 1),
    group_labels=("Group 0", "Group 1"),
    bin_edges=None,
    normalize: bool = False,
    xlabel: str = None,
    title: str = None,
    alpha: float = 0.5,
    colors=("#1F77B4", "#D62728"),
    fig: plt.Figure = None,
    ax: plt.Axes = None,
    figsize: tuple = (10, 6),
) -> Tuple[plt.Figure, plt.Axes]:
    """
    A generic helper that plots a histogram of 'value_col' for two groups
    defined by 'group_col', e.g. group_col=0 vs. group_col=1.

    Args:
        df (pd.DataFrame): DataFrame containing the data.
        value_col (str): The column whose distribution we want to plot.
        group_col (str): The column that indicates group membership.
        group_values (tuple): The two distinct values used to split the DataFrame.
        group_labels (tuple): Labels for legend (e.g. "Control", "Treatment").
        bin_edges (array): The bin edges for histogram. If None, defaults to 50 bins from 0..1
        normalize (bool): Whether to normalize the histogram (density=True).
        xlabel (str): X-axis label.
        title (str): Plot title.
        alpha (float): Transparency for the histogram overlay.
        colors (tuple): Colors for the two histograms.
        fig, ax: If provided, plot into them; otherwise create new figure/axes.
        figsize (tuple): Size of figure if we create a new one.

    Returns:
        (fig, ax)
    """
    # create or reuse figure/axes
    if fig is None and ax is None:
        fig, ax = plt.subplots(figsize=figsize)
    elif ax is None:
        ax = fig.add_subplot(111)
    elif fig is None:
        raise ValueError("fig and ax cannot both be None")

    # default bins
    if bin_edges is None:
        bin_edges = np.linspace(0, 1, 51)  # 50 bins in [0,1]

    # group 0
    mask0 = df[group_col] == group_values[0]
    ax.hist(
        df.loc[mask0, value_col],
        bins=bin_edges,
        alpha=alpha,
        label=group_labels[0],
        color=colors[0],
        density=normalize,
    )

    # group 1
    mask1 = df[group_col] == group_values[1]
    ax.hist(
        df.loc[mask1, value_col],
        bins=bin_edges,
        alpha=alpha,
        label=group_labels[1],
        color=colors[1],
        density=normalize,
    )

    ax.set_xlabel(xlabel if xlabel else value_col)
    ax.set_ylabel("Count" if not normalize else "Density")
    if title:
        ax.set_title(title)
    ax.legend()

    return fig, ax

plot_calibration

plot_calibration(df: DataFrame, proba_col: str = 'probas', target_col: str = 'targets', df2: Optional[DataFrame] = None, n_bins: int = 10, strategy: str = 'uniform', labels: Tuple[str, str] = ('Model', 'Comparison'), include_brier: bool = True, include_counts: bool = False, include_ideal: bool = True, markers: Tuple[str, str] = ('o', 's'), colors: Tuple[str, str] = ('b', 'r'), alpha: float = 0.7, xlabel: str = 'Mean Predicted Probability', ylabel: str = 'Fraction of Positives', title: str = 'Calibration Plot', fig: Optional[Figure] = None, ax: Optional[Axes] = None, figsize: Tuple[int, int] = (10, 6), pos_label: Union[int, str] = 1) -> Tuple[plt.Figure, plt.Axes]

Plot calibration curves for one or two datasets.

Parameters:

Name	Type	Description	Default
df	DataFrame	DataFrame containing true labels and probability predictions.	required
proba_col	str	Column name for predicted probabilities.	'probas'
target_col	str	Column name for true binary labels (0/1).	'targets'
df2	DataFrame	Optional second DataFrame for comparison.	None
n_bins	int	Number of bins for calibration curve.	10
strategy	str	Binning strategy, 'uniform' creates equal-width bins, 'quantile' creates equal-populated bins.	'uniform'
labels	tuple	Labels for each dataset (shown in legend with optional Brier scores).	('Model', 'Comparison')
include_brier	bool	Whether to include Brier scores in the legend.	True
include_counts	bool	Whether to display counts in each bin as text.	False
include_ideal	bool	Whether to plot the ideal diagonal line.	True
markers	tuple	Marker styles for the two curves.	('o', 's')
colors	tuple	Colors for the two curves.	('b', 'r')
alpha	float	Transparency for markers.	0.7
xlabel	str	X-axis label.	'Mean Predicted Probability'
ylabel	str	Y-axis label.	'Fraction of Positives'
title	str	Plot title.	'Calibration Plot'
fig, ax		If provided, plot into them; otherwise create new figure/axes.	required
figsize	tuple	Size of figure if we create a new one.	(10, 6)
pos_label	Union[int, str]	Label of the positive class for brier score calculation.	1

Returns:

Type	Description
(fig, ax)	Figure and axes objects

Source code in CausalEstimate/vis/plotting.py

def plot_calibration(
    df: pd.DataFrame,
    proba_col: str = "probas",
    target_col: str = "targets",
    df2: Optional[pd.DataFrame] = None,
    n_bins: int = 10,
    strategy: str = "uniform",
    labels: Tuple[str, str] = ("Model", "Comparison"),
    include_brier: bool = True,
    include_counts: bool = False,
    include_ideal: bool = True,
    markers: Tuple[str, str] = ("o", "s"),
    colors: Tuple[str, str] = ("b", "r"),
    alpha: float = 0.7,
    xlabel: str = "Mean Predicted Probability",
    ylabel: str = "Fraction of Positives",
    title: str = "Calibration Plot",
    fig: Optional[plt.Figure] = None,
    ax: Optional[plt.Axes] = None,
    figsize: Tuple[int, int] = (10, 6),
    pos_label: Union[int, str] = 1,
) -> Tuple[plt.Figure, plt.Axes]:
    """
    Plot calibration curves for one or two datasets.

    Args:
        df (pd.DataFrame): DataFrame containing true labels and probability predictions.
        proba_col (str): Column name for predicted probabilities.
        target_col (str): Column name for true binary labels (0/1).
        df2 (pd.DataFrame, optional): Optional second DataFrame for comparison.
        n_bins (int): Number of bins for calibration curve.
        strategy (str): Binning strategy, 'uniform' creates equal-width bins,
                       'quantile' creates equal-populated bins.
        labels (tuple): Labels for each dataset (shown in legend with optional Brier scores).
        include_brier (bool): Whether to include Brier scores in the legend.
        include_counts (bool): Whether to display counts in each bin as text.
        include_ideal (bool): Whether to plot the ideal diagonal line.
        markers (tuple): Marker styles for the two curves.
        colors (tuple): Colors for the two curves.
        alpha (float): Transparency for markers.
        xlabel (str): X-axis label.
        ylabel (str): Y-axis label.
        title (str): Plot title.
        fig, ax: If provided, plot into them; otherwise create new figure/axes.
        figsize (tuple): Size of figure if we create a new one.
        pos_label: Label of the positive class for brier score calculation.

    Returns:
        (fig, ax): Figure and axes objects
    """
    # Create or reuse figure/axes
    if fig is None and ax is None:
        fig, ax = plt.subplots(figsize=figsize)
    elif ax is None:
        ax = fig.add_subplot(111)
    elif fig is None:
        fig = ax.figure

    # Process datasets
    datasets = [df]
    if df2 is not None:
        datasets.append(df2)

    results = []
    for i, dataset in enumerate(datasets):
        result = _process_single_dataset(
            df=dataset,
            target_col=target_col,
            proba_col=proba_col,
            n_bins=n_bins,
            strategy=strategy,
            include_brier=include_brier,
            pos_label=pos_label,
        )
        results.append(result)

    # Plot each dataset
    for i, result in enumerate(results):
        # Create label with optional Brier score
        label = labels[i]
        if include_brier and result["brier_score"] is not None:
            label = f"{label} (Brier = {result['brier_score']:.3f})"

        # Plot calibration curve
        ax.plot(
            result["prob_pred"],
            result["prob_true"],
            marker=markers[i],
            color=colors[i],
            alpha=alpha,
            label=label,
        )

        # Add count annotations if requested
        if include_counts:
            for x, y, count in zip(
                result["prob_pred"], result["prob_true"], result["bin_counts"]
            ):
                if count > 0:  # Only annotate if there are points in the bin
                    ax.annotate(
                        f"{count}",
                        (x, y),
                        textcoords="offset points",
                        xytext=(0, 5),
                        ha="center",
                        fontsize=8,
                    )

    # Plot ideal diagonal if requested
    if include_ideal:
        ax.plot([0, 1], [0, 1], linestyle="--", color="gray", alpha=0.7, label="Ideal")

    # Set labels and title
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.set_title(title)
    ax.grid(True, alpha=0.3)
    ax.legend()

    return fig, ax

plot_calibration_comparison

plot_calibration_comparison(df1: DataFrame, df2: DataFrame, target_col: str = 'targets', proba_col: str = 'probas', n_bins: int = 10, strategy: str = 'uniform', labels: Tuple[str, str] = ('Before', 'After'), fig: Optional[Figure] = None, ax: Optional[Axes] = None, figsize: Tuple[int, int] = (10, 6), **kwargs) -> Tuple[plt.Figure, plt.Axes]

Plot calibration curves for two datasets on the same axes. A convenience wrapper around plot_calibration.

Parameters:

Name	Type	Description	Default
df1, df2	DataFrame	DataFrames containing true labels and probability predictions.	required
target_col	str	Column name for true binary labels (0/1).	'targets'
proba_col	str	Column name for predicted probabilities.	'probas'
n_bins	int	Number of bins for calibration curve.	10
strategy	str	Binning strategy, 'uniform' creates equal-width bins, 'quantile' creates equal-populated bins.	'uniform'
labels	tuple	Labels for each dataset (shown in legend with Brier scores).	('Before', 'After')
fig, ax		If provided, plot into them; otherwise create new figure/axes.	required
figsize	tuple	Size of figure if we create a new one.	(10, 6)
**kwargs		Additional arguments passed to plot_calibration	{}

Returns:

Type	Description
(fig, ax)	Figure and axes objects

Source code in CausalEstimate/vis/plotting.py

def plot_calibration_comparison(
    df1: pd.DataFrame,
    df2: pd.DataFrame,
    target_col: str = "targets",
    proba_col: str = "probas",
    n_bins: int = 10,
    strategy: str = "uniform",
    labels: Tuple[str, str] = ("Before", "After"),
    fig: Optional[plt.Figure] = None,
    ax: Optional[plt.Axes] = None,
    figsize: Tuple[int, int] = (10, 6),
    **kwargs,
) -> Tuple[plt.Figure, plt.Axes]:
    """
    Plot calibration curves for two datasets on the same axes.
    A convenience wrapper around plot_calibration.

    Args:
        df1, df2 (pd.DataFrame): DataFrames containing true labels and probability predictions.
        target_col (str): Column name for true binary labels (0/1).
        proba_col (str): Column name for predicted probabilities.
        n_bins (int): Number of bins for calibration curve.
        strategy (str): Binning strategy, 'uniform' creates equal-width bins,
                       'quantile' creates equal-populated bins.
        labels (tuple): Labels for each dataset (shown in legend with Brier scores).
        fig, ax: If provided, plot into them; otherwise create new figure/axes.
        figsize (tuple): Size of figure if we create a new one.
        **kwargs: Additional arguments passed to plot_calibration

    Returns:
        (fig, ax): Figure and axes objects
    """
    return plot_calibration(
        df=df1,
        df2=df2,
        proba_col=proba_col,
        target_col=target_col,
        n_bins=n_bins,
        strategy=strategy,
        labels=labels,
        fig=fig,
        ax=ax,
        figsize=figsize,
        **kwargs,
    )

plot_weight_dist

plot_weight_dist(df: DataFrame, ps_col: str = PS_COL, treatment_col: str = TREATMENT_COL, weight_type: str = 'ATE', clip_percentile: float = 1, bin_edges: ndarray = None, normalize: bool = False, xlabel: str = 'IPW weight', title: str = 'IPW Weight Distribution', fig: Figure = None, ax: Axes = None, figsize: tuple = (10, 6)) -> Tuple[plt.Figure, plt.Axes]

Plot the distribution of IPW weights for treatment vs. control groups.

Weights are computed with compute_ipw_weights (raises on propensity scores of exactly 0 or 1). Default bins span the observed weight range.

Parameters:

Name	Type	Description	Default
df	DataFrame	DataFrame with treatment and propensity score columns.	required
ps_col	str	Name of the propensity score column.	PS_COL
treatment_col	str	Name of the treatment status column (1 treated, 0 control).	TREATMENT_COL
weight_type	str	"ATE" or "ATT".	'ATE'
clip_percentile	float	Upper-tail clipping passed to compute_ipw_weights.	1
bin_edges	ndarray	Histogram bin edges; defaults to 50 bins over the weight range.	None
normalize	bool	Whether to normalize the histograms (density=True).	False
xlabel, title, fig, ax, figsize		As in plot_hist_by_groups.	required

Returns:

Type	Description
Tuple[Figure, Axes]	(fig, ax)

Source code in CausalEstimate/vis/plotting.py

def plot_weight_dist(
    df: pd.DataFrame,
    ps_col: str = PS_COL,
    treatment_col: str = TREATMENT_COL,
    weight_type: str = "ATE",
    clip_percentile: float = 1,
    bin_edges: np.ndarray = None,
    normalize: bool = False,
    xlabel: str = "IPW weight",
    title: str = "IPW Weight Distribution",
    fig: plt.Figure = None,
    ax: plt.Axes = None,
    figsize: tuple = (10, 6),
) -> Tuple[plt.Figure, plt.Axes]:
    """
    Plot the distribution of IPW weights for treatment vs. control groups.

    Weights are computed with compute_ipw_weights (raises on propensity scores
    of exactly 0 or 1). Default bins span the observed weight range.

    Args:
        df: DataFrame with treatment and propensity score columns.
        ps_col: Name of the propensity score column.
        treatment_col: Name of the treatment status column (1 treated, 0 control).
        weight_type: "ATE" or "ATT".
        clip_percentile: Upper-tail clipping passed to compute_ipw_weights.
        bin_edges: Histogram bin edges; defaults to 50 bins over the weight range.
        normalize: Whether to normalize the histograms (density=True).
        xlabel, title, fig, ax, figsize: As in plot_hist_by_groups.

    Returns:
        (fig, ax)
    """
    weights = compute_ipw_weights(
        df[treatment_col].to_numpy(),
        df[ps_col].to_numpy(),
        weight_type=weight_type,
        clip_percentile=clip_percentile,
    )
    if bin_edges is None:
        # span to the 99.5th percentile so a few extreme weights don't
        # squash the bulk of the distribution into a sliver
        w_min, w_max = weights.min(), np.percentile(weights, 99.5)
        if w_min == w_max:  # constant weights: avoid zero-width bins
            w_min, w_max = w_min - 0.5, w_max + 0.5
        bin_edges = np.linspace(w_min, w_max, 51)
    plot_df = pd.DataFrame(
        {"_weight": weights, treatment_col: df[treatment_col].to_numpy()}
    )
    return plot_hist_by_groups(
        df=plot_df,
        value_col="_weight",
        group_col=treatment_col,
        group_values=(0, 1),
        group_labels=("Control", "Treatment"),
        bin_edges=bin_edges,
        normalize=normalize,
        xlabel=xlabel,
        title=title,
        fig=fig,
        ax=ax,
        figsize=figsize,
    )

plot_love

plot_love(balance_table: DataFrame, threshold: float = 0.1, fig: Figure = None, ax: Axes = None, figsize: tuple = None) -> Tuple[plt.Figure, plt.Axes]

Love plot of covariate balance from a compute_balance_table result.

Shows |SMD| per covariate before (open circles) and after (filled circles) IPW weighting, connected per covariate, with a dashed line at the balance threshold. Covariates are sorted by unweighted |SMD| (largest at the top); rows with undefined (NaN) SMDs are dropped with a warning.

Parameters:

Name	Type	Description	Default
balance_table	DataFrame	Output of CausalEstimate.diagnostics.compute_balance_table.	required
threshold	float	|SMD| bound drawn as the balance reference line.	0.1
fig, ax		As in plot_hist_by_groups.	required
figsize	tuple	Defaults to a height scaled to the number of covariates.	None

Returns:

Type	Description
Tuple[Figure, Axes]	(fig, ax)

Source code in CausalEstimate/vis/plotting.py

def plot_love(
    balance_table: pd.DataFrame,
    threshold: float = 0.1,
    fig: plt.Figure = None,
    ax: plt.Axes = None,
    figsize: tuple = None,
) -> Tuple[plt.Figure, plt.Axes]:
    """
    Love plot of covariate balance from a compute_balance_table result.

    Shows |SMD| per covariate before (open circles) and after (filled circles)
    IPW weighting, connected per covariate, with a dashed line at the balance
    threshold. Covariates are sorted by unweighted |SMD| (largest at the top);
    rows with undefined (NaN) SMDs are dropped with a warning.

    Args:
        balance_table: Output of CausalEstimate.diagnostics.compute_balance_table.
        threshold: |SMD| bound drawn as the balance reference line.
        fig, ax: As in plot_hist_by_groups.
        figsize: Defaults to a height scaled to the number of covariates.

    Returns:
        (fig, ax)
    """
    table = balance_table[[SMD_UNWEIGHTED_COL, SMD_WEIGHTED_COL]].abs()
    if table.isna().any(axis=None):
        warnings.warn(
            "Dropping covariates with undefined (NaN) SMD from the Love plot.",
            RuntimeWarning,
            stacklevel=2,
        )
        table = table.dropna()
    if len(table) == 0:
        raise ValueError("No covariates with a defined SMD to plot.")
    table = table.sort_values(SMD_UNWEIGHTED_COL)

    if figsize is None:
        figsize = (7, max(2.2, 0.5 * len(table) + 1.4))
    if fig is None and ax is None:
        fig, ax = plt.subplots(figsize=figsize)
    elif ax is None:
        ax = fig.add_subplot(111)
    elif fig is None:
        raise ValueError("fig and ax cannot both be None")

    y = np.arange(len(table))
    for yi, (u, w) in enumerate(
        zip(table[SMD_UNWEIGHTED_COL], table[SMD_WEIGHTED_COL])
    ):
        ax.plot([u, w], [yi, yi], color="0.8", linewidth=1.5, zorder=1)
    ax.scatter(
        table[SMD_UNWEIGHTED_COL],
        y,
        s=55,
        facecolors="white",
        edgecolors="#1F77B4",
        linewidths=1.8,
        label="Unweighted",
        zorder=2,
    )
    ax.scatter(
        table[SMD_WEIGHTED_COL],
        y,
        s=55,
        color="#D62728",
        label="Weighted",
        zorder=2,
    )
    ax.axvline(threshold, linestyle="--", color="0.6", linewidth=1)
    ax.set_yticks(y)
    ax.set_yticklabels(table.index)
    ax.set_xlim(left=0)
    ax.set_ylim(-0.6, len(table) - 0.4)
    ax.xaxis.grid(True, color="0.9", linewidth=0.8)
    ax.set_axisbelow(True)
    for side in ("top", "right"):
        ax.spines[side].set_visible(False)
    ax.set_xlabel("|Standardized mean difference|")
    ax.set_title("Covariate Balance (Love Plot)")
    ax.legend(frameon=False, loc="lower right")
    return fig, ax

plot_ps_boxplot

plot_ps_boxplot(df: DataFrame, ps_col: str = PS_COL, treatment_col: str = TREATMENT_COL, weight_type: str = 'ATE', clip_percentile: float = 1, fig: Figure = None, ax: Axes = None, figsize: tuple = (8, 5)) -> Tuple[plt.Figure, plt.Axes]

Boxplots of the propensity score by treatment group, before and after IPW weighting.

The unweighted boxes show the raw overlap between arms; the weighted boxes use IPW-weighted quantiles, so under good weighting the treated and control boxes should nearly coincide. Whiskers extend to the most extreme point within 1.5 IQR of the box; outliers are not drawn.

Parameters:

Name	Type	Description	Default
df	DataFrame	DataFrame with treatment and propensity score columns.	required
ps_col	str	Name of the propensity score column.	PS_COL
treatment_col	str	Name of the treatment status column (1 treated, 0 control).	TREATMENT_COL
weight_type	str	"ATE" or "ATT", passed to compute_ipw_weights.	'ATE'
clip_percentile	float	Upper-tail clipping passed to compute_ipw_weights.	1
fig, ax, figsize		As in plot_hist_by_groups.	required

Returns:

Type	Description
Tuple[Figure, Axes]	(fig, ax)

Source code in CausalEstimate/vis/plotting.py

def plot_ps_boxplot(
    df: pd.DataFrame,
    ps_col: str = PS_COL,
    treatment_col: str = TREATMENT_COL,
    weight_type: str = "ATE",
    clip_percentile: float = 1,
    fig: plt.Figure = None,
    ax: plt.Axes = None,
    figsize: tuple = (8, 5),
) -> Tuple[plt.Figure, plt.Axes]:
    """
    Boxplots of the propensity score by treatment group, before and after
    IPW weighting.

    The unweighted boxes show the raw overlap between arms; the weighted boxes
    use IPW-weighted quantiles, so under good weighting the treated and control
    boxes should nearly coincide. Whiskers extend to the most extreme point
    within 1.5 IQR of the box; outliers are not drawn.

    Args:
        df: DataFrame with treatment and propensity score columns.
        ps_col: Name of the propensity score column.
        treatment_col: Name of the treatment status column (1 treated, 0 control).
        weight_type: "ATE" or "ATT", passed to compute_ipw_weights.
        clip_percentile: Upper-tail clipping passed to compute_ipw_weights.
        fig, ax, figsize: As in plot_hist_by_groups.

    Returns:
        (fig, ax)
    """
    A = df[treatment_col].to_numpy()
    ps = df[ps_col].to_numpy(dtype=float)
    weights = compute_ipw_weights(
        A, ps, weight_type=weight_type, clip_percentile=clip_percentile
    )
    if fig is None and ax is None:
        fig, ax = plt.subplots(figsize=figsize)
    elif ax is None:
        ax = fig.add_subplot(111)
    elif fig is None:
        raise ValueError("fig and ax cannot both be None")

    ctrl, trt = A == 0, A == 1
    ones = np.ones_like(ps)
    stats = [
        _box_stats(ps[ctrl], ones[ctrl], "Control"),
        _box_stats(ps[trt], ones[trt], "Treatment"),
        _box_stats(ps[ctrl], weights[ctrl], "Control"),
        _box_stats(ps[trt], weights[trt], "Treatment"),
    ]
    positions = [0, 1, 2.5, 3.5]
    colors = ["#1F77B4", "#D62728", "#1F77B4", "#D62728"]
    bp = ax.bxp(
        stats,
        positions=positions,
        widths=0.7,
        patch_artist=True,
        showfliers=False,
        medianprops={"color": "black"},
    )
    for patch, color in zip(bp["boxes"], colors):
        patch.set_facecolor(color)
        patch.set_alpha(0.5)
    ax.text(0.5, 1.02, "Unweighted", ha="center", transform=ax.get_xaxis_transform())
    ax.text(3.0, 1.02, "Weighted", ha="center", transform=ax.get_xaxis_transform())
    ax.set_ylabel("Propensity Score")
    ax.set_ylim(0, 1)
    ax.yaxis.grid(True, color="0.9", linewidth=0.8)
    ax.set_axisbelow(True)
    for side in ("top", "right"):
        ax.spines[side].set_visible(False)
    ax.set_title("Propensity Score by Treatment Group", pad=18)
    return fig, ax

plot_zipper

plot_zipper(truth: Union[float, ndarray], lower: ndarray, upper: ndarray, fig: Figure = None, ax: Axes = None, figsize: tuple = (7, 6)) -> Tuple[plt.Figure, plt.Axes]

Zipper plot of confidence-interval coverage across simulation replicates.

Each replicate's interval is drawn as a horizontal segment, sorted by its midpoint, and colored by whether it covers the truth. The empirical coverage is shown in the legend. Intervals are drawn relative to the truth, so truth may be a single value or one value per replicate.

Parameters:

Name	Type	Description	Default
truth	Union[float, ndarray]	True effect, scalar or array of length len(lower).	required
lower, upper		Interval bounds, one entry per replicate.	required
fig, ax, figsize		As in plot_hist_by_groups.	required

Returns:

Type	Description
Tuple[Figure, Axes]	(fig, ax)

Source code in CausalEstimate/vis/plotting.py

def plot_zipper(
    truth: Union[float, np.ndarray],
    lower: np.ndarray,
    upper: np.ndarray,
    fig: plt.Figure = None,
    ax: plt.Axes = None,
    figsize: tuple = (7, 6),
) -> Tuple[plt.Figure, plt.Axes]:
    """
    Zipper plot of confidence-interval coverage across simulation replicates.

    Each replicate's interval is drawn as a horizontal segment, sorted by its
    midpoint, and colored by whether it covers the truth. The empirical
    coverage is shown in the legend. Intervals are drawn relative to the
    truth, so `truth` may be a single value or one value per replicate.

    Args:
        truth: True effect, scalar or array of length len(lower).
        lower, upper: Interval bounds, one entry per replicate.
        fig, ax, figsize: As in plot_hist_by_groups.

    Returns:
        (fig, ax)
    """
    lower = np.asarray(lower, dtype=float)
    upper = np.asarray(upper, dtype=float)
    truth = np.broadcast_to(np.asarray(truth, dtype=float), lower.shape)
    if lower.shape != upper.shape or lower.ndim != 1 or lower.size == 0:
        raise ValueError(
            "lower and upper must be non-empty 1-D arrays of equal length."
        )
    if np.any(lower > upper):
        raise ValueError("lower must not exceed upper.")
    lo, hi = lower - truth, upper - truth
    order = np.argsort((lo + hi) / 2)
    lo, hi = lo[order], hi[order]
    covers = (lo <= 0) & (hi >= 0)
    coverage = covers.mean()

    if fig is None and ax is None:
        fig, ax = plt.subplots(figsize=figsize)
    elif ax is None:
        ax = fig.add_subplot(111)
    elif fig is None:
        raise ValueError("fig and ax cannot both be None")

    y = np.arange(len(lo))
    for mask, color, label in (
        (covers, "#1F77B4", f"Covers truth ({coverage:.1%})"),
        (~covers, "#D62728", f"Misses truth ({1 - coverage:.1%})"),
    ):
        if mask.any():
            ax.hlines(
                y[mask], lo[mask], hi[mask], color=color, linewidth=1.2, label=label
            )
    ax.axvline(0, color="black", linewidth=1)
    ax.set_ylim(-1, len(lo))
    ax.set_yticks([])
    ax.set_ylabel("Replicates (sorted by interval midpoint)")
    ax.set_xlabel("Interval relative to truth")
    ax.xaxis.grid(True, color="0.9", linewidth=0.8)
    ax.set_axisbelow(True)
    for side in ("top", "right", "left"):
        ax.spines[side].set_visible(False)
    ax.set_title("Confidence Interval Coverage (Zipper Plot)")
    ax.legend(frameon=False, loc="upper left")
    return fig, ax