[ Web Proxy ]
URL:
Viewing: https://raw.githubusercontent.com/UtilityTools/diff-diff/main/diff_diff/triple_diff.py [Back]  [Original]

"""
Triple Difference (DDD) estimators.

Implements the methodology from Ortiz-Villavicencio & Sant'Anna (2025)
"Better Understanding Triple Differences Estimators" for causal inference
when treatment requires satisfying two criteria:
1. Belonging to a treated group (e.g., a state with a policy)
2. Being in an eligible partition (e.g., women, low-income, etc.)

This module provides regression adjustment, inverse probability weighting,
and doubly robust estimators that correctly handle covariate adjustment,
unlike naive implementations. Standard errors use the efficient influence
function: SE = std(IF) / sqrt(n), which is inherently heteroskedasticity-
robust. Cluster-robust SEs are available via the ``cluster`` parameter.

The DDD is computed via three pairwise DiD comparisons matching R's
``triplediff::ddd()`` package (panel=FALSE mode).

Reference:
    Ortiz-Villavicencio, M., & Sant'Anna, P. H. C. (2025).
    Better Understanding Triple Differences Estimators.
    arXiv:2505.09942.
"""

import warnings
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import pandas as pd

from diff_diff.linalg import solve_logit, solve_ols
from diff_diff.results import _format_survey_block, _get_significance_stars
from diff_diff.utils import safe_inference

_MIN_CELL_SIZE = 10

# =============================================================================
# Results Classes
# =============================================================================


@dataclass
class TripleDifferenceResults:
    """
    Results from Triple Difference (DDD) estimation.

    Provides access to the estimated average treatment effect on the treated
    (ATT), standard errors, confidence intervals, and diagnostic information.

    Attributes
    ----------
    att : float
        Average Treatment effect on the Treated (ATT).
        This is the effect on units in the treated group (G=1) and eligible
        partition (P=1) after treatment (T=1).
    se : float
        Standard error of the ATT estimate.
    t_stat : float
        T-statistic for the ATT estimate.
    p_value : float
        P-value for the null hypothesis that ATT = 0.
    conf_int : tuple[float, float]
        Confidence interval for the ATT.
    n_obs : int
        Total number of observations used in estimation.
    n_treated_eligible : int
        Number of observations in treated group and eligible partition.
    n_treated_ineligible : int
        Number of observations in treated group and ineligible partition.
    n_control_eligible : int
        Number of observations in control group and eligible partition.
    n_control_ineligible : int
        Number of observations in control group and ineligible partition.
    estimation_method : str
        Estimation method used: "dr" (doubly robust), "reg" (regression
        adjustment), or "ipw" (inverse probability weighting).
    alpha : float
        Significance level used for confidence intervals.
    """

    att: float
    se: float
    t_stat: float
    p_value: float
    conf_int: Tuple[float, float]
    n_obs: int
    n_treated_eligible: int
    n_treated_ineligible: int
    n_control_eligible: int
    n_control_ineligible: int
    estimation_method: str
    alpha: float = 0.05
    # Group means for diagnostics
    group_means: Optional[Dict[str, float]] = field(default=None)
    # Propensity score diagnostics (for IPW/DR)
    pscore_stats: Optional[Dict[str, float]] = field(default=None)
    # Regression diagnostics
    r_squared: Optional[float] = field(default=None)
    # Covariate balance statistics
    covariate_balance: Optional[pd.DataFrame] = field(default=None, repr=False)
    # Inference details
    inference_method: str = field(default="analytical")
    n_bootstrap: Optional[int] = field(default=None)
    n_clusters: Optional[int] = field(default=None)
    # Survey design metadata (SurveyMetadata instance from diff_diff.survey)
    survey_metadata: Optional[Any] = field(default=None)
    # EPV diagnostics per subgroup comparison
    epv_diagnostics: Optional[Dict[int, Dict[str, Any]]] = field(
        default=None, repr=False
    )
    epv_threshold: float = 10
    pscore_fallback: str = "error"

    def __repr__(self) -> str:
        """Concise string representation."""
        return (
            f"TripleDifferenceResults(ATT={self.att:.4f}{self.significance_stars}, "
            f"SE={self.se:.4f}, p={self.p_value:.4f}, method={self.estimation_method})"
        )

    def summary(self, alpha: Optional[float] = None) -> str:
        """
        Generate a formatted summary of the estimation results.

        Parameters
        ----------
        alpha : float, optional
            Significance level for confidence intervals. Defaults to the
            alpha used during estimation.

        Returns
        -------
        str
            Formatted summary table.
        """
        alpha = alpha or self.alpha
        conf_level = int((1 - alpha) * 100)

        lines = [
            "=" * 75,
            "Triple Difference (DDD) Estimation Results".center(75),
            "=" * 75,
            "",
            f"{'Estimation method:':15}",
            f"{'Total observations:':15}",
            "",
            "Sample Composition by Cell:",
            f"  {'Treated group, Eligible:':15}",
            f"  {'Treated group, Ineligible:':15}",
            f"  {'Control group, Eligible:':15}",
            f"  {'Control group, Ineligible:':15}",
        ]

        if self.r_squared is not None:
            lines.append(f"{'R-squared:':15.4f}")

        # Add survey design info
        if self.survey_metadata is not None:
            sm = self.survey_metadata
            lines.extend(_format_survey_block(sm, 75))

        if self.inference_method != "analytical":
            lines.append(f"{'Inference method:':15}")
            if self.n_bootstrap is not None:
                lines.append(f"{'Bootstrap replications:':15}")
        if self.n_clusters is not None:
            lines.append(f"{'Number of clusters:':15}")

        lines.extend(
            [
                "",
                "-" * 75,
                f"{'Parameter':12} {'Std. Err.':>12} {'t-stat':>10} {'P>|t|':>10} {'':>5}",
                "-" * 75,
                f"{'ATT':12.4f} {self.se:>12.4f} {self.t_stat:>10.3f} {self.p_value:>10.4f} {self.significance_stars:>5}",
                "-" * 75,
                "",
                f"{conf_level}% Confidence Interval: [{self.conf_int[0]:.4f}, {self.conf_int[1]:.4f}]",
            ]
        )

        # EPV diagnostics block (if any subgroup has low EPV)
        if self.epv_diagnostics:
            low_epv = {k: v for k, v in self.epv_diagnostics.items() if v.get("is_low")}
            if low_epv:
                n_affected = len(low_epv)
                n_total = len(self.epv_diagnostics)
                min_entry = min(low_epv.values(), key=lambda v: v["epv"])
                lines.extend(
                    [
                        "",
                        "-" * 75,
                        "EPV Diagnostics".center(75),
                        "-" * 75,
                        f"WARNING: Low Events Per Variable (EPV) in "
                        f"{n_affected} of {n_total} subgroup comparison(s).",
                        f"Minimum EPV: {min_entry['epv']:.1f}. "
                        f"Threshold: {self.epv_threshold:.0f}.",
                        "Consider: estimation_method='reg' or fewer covariates.",
                        "Call results.epv_summary() for details.",
                        "-" * 75,
                    ]
                )

        # Show group means if available
        if self.group_means:
            lines.extend(
                [
                    "",
                    "-" * 75,
                    "Cell Means (Y):",
                    "-" * 75,
                ]
            )
            for cell, mean in self.group_means.items():
                lines.append(f"  {cell:12.4f}")

        # Show propensity score diagnostics if available
        if self.pscore_stats:
            lines.extend(
                [
                    "",
                    "-" * 75,
                    "Propensity Score Diagnostics:",
                    "-" * 75,
                ]
            )
            for stat, value in self.pscore_stats.items():
                lines.append(f"  {stat:12.4f}")

        lines.extend(
            [
                "",
                "Signif. codes: '***' 0.001, '**' 0.01, '*' 0.05, '.' 0.1",
                "=" * 75,
            ]
        )

        return "\n".join(lines)

    def print_summary(self, alpha: Optional[float] = None) -> None:
        """Print the summary to stdout."""
        print(self.summary(alpha))

    def to_dict(self) -> Dict[str, Any]:
        """
        Convert results to a dictionary.

        Returns
        -------
        Dict[str, Any]
            Dictionary containing all estimation results.
        """
        result = {
            "att": self.att,
            "se": self.se,
            "t_stat": self.t_stat,
            "p_value": self.p_value,
            "conf_int_lower": self.conf_int[0],
            "conf_int_upper": self.conf_int[1],
            "n_obs": self.n_obs,
            "n_treated_eligible": self.n_treated_eligible,
            "n_treated_ineligible": self.n_treated_ineligible,
            "n_control_eligible": self.n_control_eligible,
            "n_control_ineligible": self.n_control_ineligible,
            "estimation_method": self.estimation_method,
            "inference_method": self.inference_method,
        }
        if self.r_squared is not None:
            result["r_squared"] = self.r_squared
        if self.n_bootstrap is not None:
            result["n_bootstrap"] = self.n_bootstrap
        if self.n_clusters is not None:
            result["n_clusters"] = self.n_clusters
        if self.survey_metadata is not None:
            sm = self.survey_metadata
            result["weight_type"] = sm.weight_type
            result["effective_n"] = sm.effective_n
            result["design_effect"] = sm.design_effect
            result["sum_weights"] = sm.sum_weights
            result["n_strata"] = sm.n_strata
            result["n_psu"] = sm.n_psu
            result["df_survey"] = sm.df_survey
        return result

    def to_dataframe(self) -> pd.DataFrame:
        """
        Convert results to a pandas DataFrame.

        Returns
        -------
        pd.DataFrame
            DataFrame with estimation results.
        """
        return pd.DataFrame([self.to_dict()])

    @property
    def is_significant(self) -> bool:
        """Check if the ATT is statistically significant at the alpha level."""
        return bool(self.p_value < self.alpha)

    @property
    def significance_stars(self) -> str:
        """Return significance stars based on p-value."""
        return _get_significance_stars(self.p_value)

    def epv_summary(self, show_all: bool = False) -> pd.DataFrame:
        """
        Return per-subgroup EPV diagnostics as a DataFrame.

        Parameters
        ----------
        show_all : bool, default False
            If False, only show subgroups with low EPV. If True, show all.

        Returns
        -------
        pd.DataFrame
            Columns: subgroup, epv, n_events, n_params, is_low.
        """
        if not self.epv_diagnostics:
            return pd.DataFrame(
                columns=["subgroup", "epv", "n_events", "n_params", "is_low"]
            )
        rows = []
        for sg, diag in sorted(self.epv_diagnostics.items()):
            if show_all or diag.get("is_low", False):
                rows.append(
                    {
                        "subgroup": sg,
                        "epv": diag.get("epv"),
                        "n_events": diag.get("n_events"),
                        "n_params": diag.get("k"),
                        "is_low": diag.get("is_low", False),
                    }
                )
        cols = ["subgroup", "epv", "n_events", "n_params", "is_low"]
        return pd.DataFrame(rows, columns=cols) if rows else pd.DataFrame(columns=cols)


# =============================================================================
# Helper Functions
# =============================================================================


# =============================================================================
# Main Estimator Class
# =============================================================================


class TripleDifference:
    """
    Triple Difference (DDD) estimator.

    Estimates the Average Treatment effect on the Treated (ATT) when treatment
    requires satisfying two criteria: belonging to a treated group AND being
    in an eligible partition of the population. The DDD design was popularized
    by Gruber (1994) [2]_.

    This implementation follows Ortiz-Villavicencio & Sant'Anna (2025) [1]_,
    which shows that naive DDD implementations (difference of two DiDs,
    three-way fixed effects) are invalid when covariates are needed for
    identification.

    Parameters
    ----------
    estimation_method : str, default="dr"
        Estimation method to use:

        - "dr": Doubly robust (recommended). Consistent if either the outcome
          model or propensity score model is correctly specified.
        - "reg": Regression adjustment (outcome regression).
        - "ipw": Inverse probability weighting.
    robust : bool, default=True
        Whether to use heteroskedasticity-robust standard errors.
        Note: influence function-based SEs are inherently robust to
        heteroskedasticity, so this parameter has no effect. Retained
        for API compatibility.
    cluster : str, optional
        Column name for cluster-robust standard errors. When provided,
        SEs are computed using the Liang-Zeger cluster-robust variance
        estimator on the influence function.
    alpha : float, default=0.05
        Significance level for confidence intervals.
    pscore_trim : float, default=0.01
        Trimming threshold for propensity scores. Scores below this value
        or above (1 - pscore_trim) are clipped to avoid extreme weights.
    rank_deficient_action : str, default="warn"
        Action when design matrix is rank-deficient (linearly dependent columns):

        - "warn": Issue warning and drop linearly dependent columns (default)
        - "error": Raise ValueError
        - "silent": Drop columns silently without warning
    epv_threshold : float, default=10
        Events Per Variable threshold for propensity score logit.
        When the ratio of minority-class observations to predictor
        variables (excluding intercept) falls below this value, a
        warning is emitted (or ``ValueError`` raised if
        ``rank_deficient_action="error"``). Based on Peduzzi et al.
        (1996). Only applies to IPW and DR estimation methods.
    pscore_fallback : str, default="error"
        Action when propensity score estimation fails:

        - "error": Raise the exception (default)
        - "unconditional": Fall back to unconditional propensity with
          a warning. For IPW, drops all covariates. For DR, the
          propensity model becomes unconditional but outcome regression
          still uses covariates.

        When ``rank_deficient_action="error"``, errors are always
        re-raised regardless of this setting.

    Attributes
    ----------
    results_ : TripleDifferenceResults
        Estimation results after calling fit().
    is_fitted_ : bool
        Whether the model has been fitted.

    Examples
    --------
    Basic usage with a DataFrame:

    >>> import pandas as pd
    >>> from diff_diff import TripleDifference
    >>>
    >>> # Data where treatment affects women (partition=1) in states
    >>> # that enacted a policy (group=1)
    >>> data = pd.DataFrame({
    ...     'outcome': [...],
    ...     'group': [1, 1, 0, 0, ...],      # 1=policy state, 0=control state
    ...     'partition': [1, 0, 1, 0, ...],  # 1=women, 0=men
    ...     'post': [0, 0, 1, 1, ...],       # 1=post-treatment period
    ... })
    >>>
    >>> # Fit using doubly robust estimation
    >>> ddd = TripleDifference(estimation_method="dr")
    >>> results = ddd.fit(
    ...     data,
    ...     outcome='outcome',
    ...     group='group',
    ...     partition='partition',
    ...     time='post'
    ... )
    >>> print(results.att)  # ATT estimate

    With covariates (properly handled unlike naive DDD):

    >>> results = ddd.fit(
    ...     data,
    ...     outcome='outcome',
    ...     group='group',
    ...     partition='partition',
    ...     time='post',
    ...     covariates=['age', 'income']
    ... )

    Notes
    -----
    The DDD estimator is appropriate when:

    1. Treatment affects only units satisfying BOTH criteria:
       - Belonging to a treated group (G=1), e.g., states with a policy
       - Being in an eligible partition (P=1), e.g., women, low-income

    2. The DDD parallel trends assumption holds: the differential trend
       between eligible and ineligible partitions would have been the same
       across treated and control groups, absent treatment.

    This is weaker than requiring separate parallel trends for two DiDs,
    as biases can cancel out in the differencing.

    References
    ----------
    .. [1] Ortiz-Villavicencio, M., & Sant'Anna, P. H. C. (2025).
           Better Understanding Triple Differences Estimators.
           arXiv:2505.09942.

    .. [2] Gruber, J. (1994). The incidence of mandated maternity benefits.
           American Economic Review, 84(3), 622-641.
    """

    def __init__(
        self,
        estimation_method: str = "dr",
        robust: bool = True,
        cluster: Optional[str] = None,
        alpha: float = 0.05,
        pscore_trim: float = 0.01,
        rank_deficient_action: str = "warn",
        epv_threshold: float = 10,
        pscore_fallback: str = "error",
    ):
        if estimation_method not in ("dr", "reg", "ipw"):
            raise ValueError(
                f"estimation_method must be 'dr', 'reg', or 'ipw', " f"got '{estimation_method}'"
            )
        if rank_deficient_action not in ["warn", "error", "silent"]:
            raise ValueError(
                f"rank_deficient_action must be 'warn', 'error', or 'silent', "
                f"got '{rank_deficient_action}'"
            )
        if epv_threshold  TripleDifferenceResults:
        """
        Fit the Triple Difference model.

        Parameters
        ----------
        data : pd.DataFrame
            DataFrame containing all variables.
        outcome : str
            Name of the outcome variable column.
        group : str
            Name of the group indicator column (0/1).
            1 = treated group (e.g., states that enacted policy).
            0 = control group.
        partition : str
            Name of the partition/eligibility indicator column (0/1).
            1 = eligible partition (e.g., women, targeted demographic).
            0 = ineligible partition.
        time : str
            Name of the time period indicator column (0/1).
            1 = post-treatment period.
            0 = pre-treatment period.
        covariates : list of str, optional
            List of covariate column names to adjust for.
            These are properly incorporated using the selected estimation
            method (unlike naive DDD implementations).
        survey_design : SurveyDesign, optional
            Survey design specification for complex survey data. When
            provided, uses survey weights for estimation and Taylor Series
            Linearization (TSL) for variance estimation. Supported with
            all estimation methods ("reg", "ipw", "dr").

        Returns
        -------
        TripleDifferenceResults
            Object containing estimation results.

        Raises
        ------
        ValueError
            If required columns are missing or data validation fails.
        NotImplementedError
            If survey_design is used with wild_bootstrap inference.
        """
        # Reset replicate state from any previous fit
        self._replicate_n_valid = None

        # Resolve survey design if provided
        from diff_diff.survey import (
            _inject_cluster_as_psu,
            _resolve_effective_cluster,
            _resolve_survey_for_fit,
            compute_survey_metadata,
        )

        resolved_survey, survey_weights, survey_weight_type, survey_metadata = (
            _resolve_survey_for_fit(survey_design, data, "analytical")
        )

        if resolved_survey is not None and resolved_survey.weight_type != "pweight":
            raise ValueError(
                f"TripleDifference survey support requires weight_type='pweight', "
                f"got '{resolved_survey.weight_type}'. The survey variance math "
                f"assumes probability weights (pweight)."
            )

        # Validate inputs
        self._validate_data(data, outcome, group, partition, time, covariates)

        # Extract data
        y = data[outcome].values.astype(float)
        G = data[group].values.astype(float)
        P = data[partition].values.astype(float)
        T = data[time].values.astype(float)

        # Store cluster IDs for SE computation
        self._cluster_ids = data[self.cluster].values if self.cluster is not None else None
        if self._cluster_ids is not None and np.any(pd.isna(data[self.cluster])):
            raise ValueError(f"Cluster column '{self.cluster}' contains missing values")

        # Resolve effective cluster and inject cluster-as-PSU for survey variance
        if resolved_survey is not None:
            effective_cluster_ids = _resolve_effective_cluster(
                resolved_survey, self._cluster_ids, self.cluster
            )
            if effective_cluster_ids is not None:
                resolved_survey = _inject_cluster_as_psu(resolved_survey, effective_cluster_ids)
                if resolved_survey.psu is not None and survey_metadata is not None:
                    raw_w = (
                        data[survey_design.weights].values.astype(np.float64)
                        if survey_design.weights
                        else np.ones(len(data), dtype=np.float64)
                    )
                    survey_metadata = compute_survey_metadata(resolved_survey, raw_w)

        # Get covariates if specified
        X = None
        if covariates:
            X = data[covariates].values.astype(float)
            if np.any(np.isnan(X)):
                raise ValueError("Covariates contain missing values")

        # Count observations in each cell
        n_obs = len(y)
        n_treated_eligible = int(np.sum((G == 1) & (P == 1)))
        n_treated_ineligible = int(np.sum((G == 1) & (P == 0)))
        n_control_eligible = int(np.sum((G == 0) & (P == 1)))
        n_control_ineligible = int(np.sum((G == 0) & (P == 0)))

        # Compute cell means for diagnostics
        group_means = self._compute_cell_means(y, G, P, T, weights=survey_weights)

        # Estimate ATT based on method
        if self.estimation_method == "reg":
            att, se, r_squared, pscore_stats, epv_diag = self._regression_adjustment(
                y,
                G,
                P,
                T,
                X,
                survey_weights=survey_weights,
                resolved_survey=resolved_survey,
            )
        elif self.estimation_method == "ipw":
            att, se, r_squared, pscore_stats, epv_diag = self._ipw_estimation(
                y,
                G,
                P,
                T,
                X,
                survey_weights=survey_weights,
                resolved_survey=resolved_survey,
            )
        else:  # doubly robust
            att, se, r_squared, pscore_stats, epv_diag = self._doubly_robust(
                y,
                G,
                P,
                T,
                X,
                survey_weights=survey_weights,
                resolved_survey=resolved_survey,
            )

        # Compute inference
        # When survey design is active, use survey df (n_PSU - n_strata)
        if survey_metadata is not None and survey_metadata.df_survey is not None:
            df = survey_metadata.df_survey
            # Override with effective replicate df only when replicates were dropped
            if (hasattr(self, '_replicate_n_valid') and self._replicate_n_valid is not None
                    and resolved_survey is not None
                    and self._replicate_n_valid < resolved_survey.n_replicates):
                df = self._replicate_n_valid - 1
                survey_metadata.df_survey = self._replicate_n_valid - 1
            # df 

Web Proxy Viewer  |  New URL  |  Original Page