"""

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)

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:

"""

alpha = alpha or self.alpha

conf_level = int((1 - alpha) * 100)

lines = [

"=" * 75,

"Triple Difference (DDD) Estimation Results".center(75),

"=" * 75,

"",

f"{'Estimation method:':<30} {self.estimation_method:>15}",

f"{'Total observations:':<30} {self.n_obs:>15}",

"",

"Sample Composition by Cell:",

f" {'Treated group, Eligible:':<28} {self.n_treated_eligible:>15}",

f" {'Treated group, Ineligible:':<28} {self.n_treated_ineligible:>15}",

f" {'Control group, Eligible:':<28} {self.n_control_eligible:>15}",

f" {'Control group, Ineligible:':<28} {self.n_control_ineligible:>15}",

]

if self.r_squared is not None:

lines.append(f"{'R-squared:':<30} {self.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:':<30} {self.inference_method:>15}")

if self.n_bootstrap is not None:

lines.append(f"{'Bootstrap replications:':<30} {self.n_bootstrap:>15}")

if self.n_clusters is not None:

lines.append(f"{'Number of clusters:':<30} {self.n_clusters:>15}")

lines.extend(

[

"",

"-" * 75,

f"{'Parameter':<15} {'Estimate':>12} {'Std. Err.':>12} {'t-stat':>10} {'P>|t|':>10} {'':>5}",

"-" * 75,

f"{'ATT':<15} {self.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}]",

]

)

# 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:<35} {mean:>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:<35} {value:>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]:

"""

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:

"""

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."""

"""

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",

):

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}'"

)

self.estimation_method = estimation_method

self.robust = robust

self.cluster = cluster

self.alpha = alpha

self.pscore_trim = pscore_trim

self.rank_deficient_action = rank_deficient_action

self.is_fitted_ = False

self.results_: Optional[TripleDifferenceResults] = None

def fit(

self,

data: pd.DataFrame,

outcome: str,

group: str,

partition: str,

time: str,

covariates: Optional[List[str]] = None,

survey_design=None,

) -> TripleDifferenceResults:

"""

# 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 = 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 = 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 = 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 <= 0 means insufficient rank for t-based inference

if df is not None and df <= 0:

df = 0 # Forces NaN from t-distribution

elif (resolved_survey is not None

and hasattr(resolved_survey, 'uses_replicate_variance')

and resolved_survey.uses_replicate_variance):

# Replicate design with undefined df (rank <= 1) — NaN inference

df = 0 # Forces NaN from t-distribution

else:

df = n_obs - 8 # Approximate df (8 cell means)

if covariates:

df -= len(covariates)

df = max(df, 1)

t_stat, p_value, conf_int = safe_inference(att, se, alpha=self.alpha, df=df)

# Get number of clusters if clustering

n_clusters = None

if self.cluster is not None:

n_clusters = data[self.cluster].nunique()

# Create results object

self.results_ = TripleDifferenceResults(

att=att,

se=se,

t_stat=t_stat,

p_value=p_value,

conf_int=conf_int,

n_obs=n_obs,

n_treated_eligible=n_treated_eligible,

n_treated_ineligible=n_treated_ineligible,

n_control_eligible=n_control_eligible,

n_control_ineligible=n_control_ineligible,

estimation_method=self.estimation_method,

alpha=self.alpha,

group_means=group_means,

pscore_stats=pscore_stats,

r_squared=r_squared,

inference_method="analytical",

n_clusters=n_clusters,

survey_metadata=survey_metadata,

)

self.is_fitted_ = True

return self.results_

def _validate_data(

self,

data: pd.DataFrame,

outcome: str,

group: str,

partition: str,

time: str,

covariates: Optional[List[str]] = None,

) -> None:

"""Validate input data."""

if not isinstance(data, pd.DataFrame):

raise TypeError("data must be a pandas DataFrame")

# Check required columns exist

required_cols = [outcome, group, partition, time]

if covariates:

required_cols.extend(covariates)

if self.cluster is not None:

required_cols.append(self.cluster)

missing_cols = [col for col in required_cols if col not in data.columns]

if missing_cols:

raise ValueError(f"Missing columns in data: {missing_cols}")

# Check for missing values in required columns

for col in [outcome, group, partition, time]:

if data[col].isna().any():

raise ValueError(f"Column '{col}' contains missing values")

# Validate binary variables

for col, name in [(group, "group"), (partition, "partition"), (time, "time")]:

unique_vals = set(data[col].unique())

if not unique_vals.issubset({0, 1, 0.0, 1.0}):

raise ValueError(

f"'{name}' column must be binary (0/1), " f"got values: {sorted(unique_vals)}"

)

if len(unique_vals) < 2:

raise ValueError(f"'{name}' column must have both 0 and 1 values")

# Check we have observations in all cells

G = data[group].values

P = data[partition].values

T = data[time].values

cells = [

((G == 1) & (P == 1) & (T == 0), "treated, eligible, pre"),

((G == 1) & (P == 1) & (T == 1), "treated, eligible, post"),

((G == 1) & (P == 0) & (T == 0), "treated, ineligible, pre"),

((G == 1) & (P == 0) & (T == 1), "treated, ineligible, post"),

((G == 0) & (P == 1) & (T == 0), "control, eligible, pre"),

((G == 0) & (P == 1) & (T == 1), "control, eligible, post"),

((G == 0) & (P == 0) & (T == 0), "control, ineligible, pre"),

((G == 0) & (P == 0) & (T == 1), "control, ineligible, post"),

]

for mask, cell_name in cells:

n_cell = int(np.sum(mask))

if n_cell == 0:

raise ValueError(

f"No observations in cell: {cell_name}. "

"DDD requires observations in all 8 cells."

)

elif n_cell < _MIN_CELL_SIZE:

warnings.warn(

f"Low observation count ({n_cell}) in cell: {cell_name}. "

f"Estimates may be unreliable with fewer than "

f"{_MIN_CELL_SIZE} observations per cell.",

UserWarning,

stacklevel=2,

)

def _compute_cell_means(

self,

y: np.ndarray,

G: np.ndarray,

P: np.ndarray,

T: np.ndarray,

weights: Optional[np.ndarray] = None,

) -> Dict[str, float]:

"""Compute mean outcomes for each of the 8 DDD cells."""

means = {}

for g_val, g_name in [(1, "Treated"), (0, "Control")]:

for p_val, p_name in [(1, "Eligible"), (0, "Ineligible")]:

for t_val, t_name in [(0, "Pre"), (1, "Post")]:

mask = (G == g_val) & (P == p_val) & (T == t_val)

cell_name = f"{g_name}, {p_name}, {t_name}"

if weights is not None:

w_cell = weights[mask]

if np.sum(w_cell) <= 0:

raise ValueError(

f"Cell '{cell_name}' has zero effective survey "

f"weight. Cannot compute weighted cell mean. "

f"Check subpopulation/domain definition."

)

means[cell_name] = float(np.average(y[mask], weights=w_cell))

else:

means[cell_name] = float(np.mean(y[mask]))

return means

# =========================================================================

# Three-DiD Decomposition (matches R's triplediff::ddd())

# =========================================================================

#

# The DDD is decomposed into three pairwise DiD comparisons:

# DiD_3: subgroup 3 (G=1,P=0) vs subgroup 4 (G=1,P=1)

# DiD_2: subgroup 2 (G=0,P=1) vs subgroup 4 (G=1,P=1)

# DiD_1: subgroup 1 (G=0,P=0) vs subgroup 4 (G=1,P=1)

#

# DDD = DiD_3 + DiD_2 - DiD_1

#

# Each DiD uses the selected estimation method (DR, IPW, or RA).

# SE is computed from the combined influence function:

# inf = w3*inf_3 + w2*inf_2 - w1*inf_1

# SE = std(inf, ddof=1) / sqrt(n)

#

# Reference: Ortiz-Villavicencio & Sant'Anna (2025), implemented in

# R's triplediff::ddd() with panel=FALSE (repeated cross-section).

# =========================================================================

def _regression_adjustment(

self,

y: np.ndarray,

G: np.ndarray,

P: np.ndarray,

T: np.ndarray,

X: Optional[np.ndarray],

survey_weights: Optional[np.ndarray] = None,

resolved_survey=None,

) -> Tuple[float, float, Optional[float], Optional[Dict[str, float]]]:

"""

return self._estimate_ddd_decomposition(

y,

G,

P,

T,

X,

survey_weights=survey_weights,

resolved_survey=resolved_survey,

)

def _ipw_estimation(

self,

y: np.ndarray,

G: np.ndarray,

P: np.ndarray,

T: np.ndarray,

X: Optional[np.ndarray],

survey_weights: Optional[np.ndarray] = None,

resolved_survey=None,

) -> Tuple[float, float, Optional[float], Optional[Dict[str, float]]]:

"""

return self._estimate_ddd_decomposition(

y,

G,

P,

T,

X,

survey_weights=survey_weights,

resolved_survey=resolved_survey,

)

def _doubly_robust(

self,

y: np.ndarray,

G: np.ndarray,

P: np.ndarray,

T: np.ndarray,

X: Optional[np.ndarray],

survey_weights: Optional[np.ndarray] = None,

resolved_survey=None,

) -> Tuple[float, float, Optional[float], Optional[Dict[str, float]]]:

"""

return self._estimate_ddd_decomposition(

y,

G,

P,

T,

X,

survey_weights=survey_weights,

resolved_survey=resolved_survey,

)

def _estimate_ddd_decomposition(

self,

y: np.ndarray,

G: np.ndarray,

P: np.ndarray,

T: np.ndarray,

X: Optional[np.ndarray],

survey_weights: Optional[np.ndarray] = None,

resolved_survey=None,

) -> Tuple[float, float, Optional[float], Optional[Dict[str, float]]]:

"""

n = len(y)

est_method = self.estimation_method

# Assign subgroups following R convention:

# 4: G=1, P=1 (treated, eligible - reference/"treated")

# 3: G=1, P=0 (treated, ineligible)

# 2: G=0, P=1 (control, eligible)

# 1: G=0, P=0 (control, ineligible)

subgroup = np.zeros(n, dtype=int)

subgroup[(G == 1) & (P == 1)] = 4

subgroup[(G == 1) & (P == 0)] = 3

subgroup[(G == 0) & (P == 1)] = 2

subgroup[(G == 0) & (P == 0)] = 1

post = T.astype(float)

# Covariate matrix (always includes intercept)

if X is not None and X.shape[1] > 0:

covX = np.column_stack([np.ones(n), X])

has_covariates = True

else:

covX = np.ones((n, 1))

has_covariates = False

# Three DiD comparisons: j vs 4 for j in {3, 2, 1}

did_results = {}

pscore_stats = None

all_pscores = {} # Collect pscores for diagnostics

overlap_issues = [] # Collect overlap diagnostics across comparisons

any_nonfinite_if = False

with np.errstate(divide="ignore", invalid="ignore", over="ignore"):

for j in [3, 2, 1]:

mask = (subgroup == j) | (subgroup == 4)

y_sub = y[mask]

post_sub = post[mask]

sg_sub = subgroup[mask]

covX_sub = covX[mask]

n_sub = len(y_sub)

PA4 = (sg_sub == 4).astype(float)

PAa = (sg_sub == j).astype(float)

# Subset survey weights for this comparison (needed for logit)

w_sub = survey_weights[mask] if survey_weights is not None else None

# --- Propensity scores ---

if est_method == "reg":

# RA: no propensity scores needed

pscore_sub = np.ones(n_sub)

hessian = None

elif has_covariates:

# Logistic regression: P(subgroup=4 | X) within {j, 4}

ps_estimated = True

try:

_, pscore_sub = solve_logit(

covX_sub[:, 1:],

PA4,

rank_deficient_action=self.rank_deficient_action,

weights=w_sub,

)

except Exception:

if self.rank_deficient_action == "error":

raise

pscore_sub = np.full(n_sub, np.mean(PA4))

ps_estimated = False

warnings.warn(

f"Propensity score estimation failed for subgroup "

f"{j} vs 4; using unconditional probability. "

f"SEs may be less efficient.",

UserWarning,

stacklevel=3,

)

pscore_sub = np.clip(pscore_sub, self.pscore_trim, 1 - self.pscore_trim)

all_pscores[j] = pscore_sub

# Check overlap: count obs at trim bounds

# (1e-10 tolerance for floating-point after np.clip)

n_trimmed = int(

np.sum(

(pscore_sub <= self.pscore_trim + 1e-10)

| (pscore_sub >= 1 - self.pscore_trim - 1e-10)

)

)

frac_trimmed = n_trimmed / len(pscore_sub)

if frac_trimmed > 0.05:

overlap_issues.append((j, frac_trimmed))

# Hessian only when PS was actually estimated

if ps_estimated:

W_ps = pscore_sub * (1 - pscore_sub)

if w_sub is not None:

W_ps = W_ps * w_sub

try:

XWX = covX_sub.T @ (W_ps[:, None] * covX_sub)

hessian = np.linalg.inv(XWX) * n_sub

except np.linalg.LinAlgError:

hessian = np.linalg.pinv(XWX) * n_sub

else:

hessian = None

else:

# No covariates: unconditional probability

pscore_sub = np.full(n_sub, np.mean(PA4))

pscore_sub = np.clip(pscore_sub, self.pscore_trim, 1 - self.pscore_trim)

# Check overlap (same logic as covariate branch)

n_trimmed = int(

np.sum(

(pscore_sub <= self.pscore_trim + 1e-10)

| (pscore_sub >= 1 - self.pscore_trim - 1e-10)

)

)

frac_trimmed = n_trimmed / len(pscore_sub)

if frac_trimmed > 0.05:

overlap_issues.append((j, frac_trimmed))

hessian = None

# --- Outcome regression ---

if est_method == "ipw":

# IPW: no outcome regression

or_ctrl_pre = np.zeros(n_sub)

or_ctrl_post = np.zeros(n_sub)

or_trt_pre = np.zeros(n_sub)

or_trt_post = np.zeros(n_sub)

else:

# Fit separate OLS per subgroup-time cell, predict for all

or_ctrl_pre = self._fit_predict_mu(

y_sub,

covX_sub,

sg_sub == j,

post_sub == 0,

n_sub,

weights=w_sub,