diff-diff 2.0.0__tar.gz → 2.0.2__tar.gz

{diff_diff-2.0.0 → diff_diff-2.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diff-diff
-Version: 2.0.0
+Version: 2.0.2
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Science/Research
 Classifier: Operating System :: OS Independent

{diff_diff-2.0.0 → diff_diff-2.0.2}/diff_diff/__init__.py

@@ -113,7 +113,7 @@ from diff_diff.visualization import (
     plot_sensitivity,
 )
 
-__version__ = "2.0.0"
+__version__ = "2.0.2"
 __all__ = [
     # Estimators
     "DifferenceInDifferences",

{diff_diff-2.0.0 → diff_diff-2.0.2}/diff_diff/bacon.py

@@ -17,6 +17,8 @@ from typing import Any, Dict, List, Optional, Tuple
 import numpy as np
 import pandas as pd
 
+from diff_diff.utils import within_transform as _within_transform_util
+
 
 @dataclass
 class Comparison2x2:
@@ -573,66 +575,16 @@ class BaconDecomposition:
         treat_col: str = '__bacon_treated_internal__',
     ) -> float:
         """Compute TWFE estimate using within-transformation."""
-        # Demean by unit and time
-        y = df[outcome].values
-        d = df[treat_col].astype(float).values
-
-        # Create unit and time dummies for demeaning
-        units = df[unit].values
-        times = df[time].values
-
-        # Unit means
-        unit_map = {u: i for i, u in enumerate(df[unit].unique())}
-        unit_idx = np.array([unit_map[u] for u in units])
-        n_units = len(unit_map)
-
-        # Time means
-        time_map = {t: i for i, t in enumerate(df[time].unique())}
-        time_idx = np.array([time_map[t] for t in times])
-        n_times = len(time_map)
-
-        # Compute means
-        y_unit_mean = np.zeros(n_units)
-        d_unit_mean = np.zeros(n_units)
-        unit_counts = np.zeros(n_units)
-
-        for i in range(len(y)):
-            u = unit_idx[i]
-            y_unit_mean[u] += y[i]
-            d_unit_mean[u] += d[i]
-            unit_counts[u] += 1
-
-        y_unit_mean /= np.maximum(unit_counts, 1)
-        d_unit_mean /= np.maximum(unit_counts, 1)
-
-        y_time_mean = np.zeros(n_times)
-        d_time_mean = np.zeros(n_times)
-        time_counts = np.zeros(n_times)
-
-        for i in range(len(y)):
-            t = time_idx[i]
-            y_time_mean[t] += y[i]
-            d_time_mean[t] += d[i]
-            time_counts[t] += 1
-
-        y_time_mean /= np.maximum(time_counts, 1)
-        d_time_mean /= np.maximum(time_counts, 1)
-
-        # Overall mean
-        y_mean = np.mean(y)
-        d_mean = np.mean(d)
-
-        # Within transformation: y_it - y_i - y_t + y
-        y_within = np.zeros(len(y))
-        d_within = np.zeros(len(d))
-
-        for i in range(len(y)):
-            u = unit_idx[i]
-            t = time_idx[i]
-            y_within[i] = y[i] - y_unit_mean[u] - y_time_mean[t] + y_mean
-            d_within[i] = d[i] - d_unit_mean[u] - d_time_mean[t] + d_mean
-
-        # OLS on demeaned data
+        # Apply two-way within transformation
+        df_dm = _within_transform_util(
+            df, [outcome, treat_col], unit, time, suffix="_within"
+        )
+
+        # Extract within-transformed values
+        y_within = df_dm[f"{outcome}_within"].values
+        d_within = df_dm[f"{treat_col}_within"].values
+
+        # OLS on demeaned data: beta = sum(d * y) / sum(d^2)
         d_var = np.sum(d_within ** 2)
         if d_var > 0:
             beta = np.sum(d_within * y_within) / d_var

{diff_diff-2.0.0 → diff_diff-2.0.2}/diff_diff/estimators.py

@@ -23,6 +23,7 @@ from diff_diff.utils import (
     WildBootstrapResults,
     compute_confidence_interval,
     compute_p_value,
+    demean_by_group,
     validate_binary,
     wild_bootstrap_se,
 )
@@ -227,10 +228,10 @@ class DifferenceInDifferences:
             # unit-invariant, so demeaning them would create multicollinearity
             vars_to_demean = [outcome] + (covariates or [])
             for ab_var in absorb:
-                n_absorbed_effects += working_data[ab_var].nunique() - 1
-                for var in vars_to_demean:
-                    group_means = working_data.groupby(ab_var)[var].transform("mean")
-                    working_data[var] = working_data[var] - group_means
+                working_data, n_fe = demean_by_group(
+                    working_data, vars_to_demean, ab_var, inplace=True
+                )
+                n_absorbed_effects += n_fe
                 absorbed_vars.append(ab_var)
 
         # Extract variables (may be demeaned if absorb was used)
@@ -828,10 +829,10 @@ class MultiPeriodDiD(DifferenceInDifferences):
         if absorb:
             vars_to_demean = [outcome] + (covariates or [])
             for ab_var in absorb:
-                n_absorbed_effects += working_data[ab_var].nunique() - 1
-                for var in vars_to_demean:
-                    group_means = working_data.groupby(ab_var)[var].transform("mean")
-                    working_data[var] = working_data[var] - group_means
+                working_data, n_fe = demean_by_group(
+                    working_data, vars_to_demean, ab_var, inplace=True
+                )
+                n_absorbed_effects += n_fe
 
         # Extract outcome and treatment
         y = working_data[outcome].values.astype(float)

{diff_diff-2.0.0 → diff_diff-2.0.2}/diff_diff/staggered.py

@@ -1479,29 +1479,36 @@ class CallawaySantAnna:
 
         Standard errors are computed using influence function aggregation,
         which properly accounts for covariances across (g,t) pairs due to
-        shared control units. This matches R's `did` package approach.
+        shared control units. This includes the wif (weight influence function)
+        adjustment from R's `did` package that accounts for uncertainty in
+        estimating the group-size weights.
         """
         effects = []
         weights_list = []
         gt_pairs = []
+        groups_for_gt = []
 
         for (g, t), data in group_time_effects.items():
             effects.append(data['effect'])
             weights_list.append(data['n_treated'])
             gt_pairs.append((g, t))
+            groups_for_gt.append(g)
 
         effects = np.array(effects)
         weights = np.array(weights_list, dtype=float)
+        groups_for_gt = np.array(groups_for_gt)
 
         # Normalize weights
-        weights = weights / np.sum(weights)
+        total_weight = np.sum(weights)
+        weights_norm = weights / total_weight
 
         # Weighted average
-        overall_att = np.sum(weights * effects)
+        overall_att = np.sum(weights_norm * effects)
 
-        # Compute SE using influence function aggregation
-        overall_se = self._compute_aggregated_se(
-            gt_pairs, weights, influence_func_info
+        # Compute SE using influence function aggregation with wif adjustment
+        overall_se = self._compute_aggregated_se_with_wif(
+            gt_pairs, weights_norm, effects, groups_for_gt,
+            influence_func_info, df, unit
         )
 
         return overall_att, overall_se
@@ -1566,6 +1573,150 @@ class CallawaySantAnna:
         variance = np.sum(psi_overall ** 2)
         return np.sqrt(variance)
 
+    def _compute_aggregated_se_with_wif(
+        self,
+        gt_pairs: List[Tuple[Any, Any]],
+        weights: np.ndarray,
+        effects: np.ndarray,
+        groups_for_gt: np.ndarray,
+        influence_func_info: Dict,
+        df: pd.DataFrame,
+        unit: str,
+    ) -> float:
+        """
+        Compute SE with weight influence function (wif) adjustment.
+
+        This matches R's `did` package approach for "simple" aggregation,
+        which accounts for uncertainty in estimating group-size weights.
+
+        The wif adjustment adds variance due to the fact that aggregation
+        weights w_g = n_g / N depend on estimated group sizes.
+
+        Formula (matching R's did::aggte):
+            agg_inf_i = Σ_k w_k × inf_i_k + wif_i × ATT_k
+            se = sqrt(mean(agg_inf^2) / n)
+
+        where:
+        - k indexes "keepers" (post-treatment (g,t) pairs)
+        - w_k = pg[k] / sum(pg[keepers]) where pg = n_g / n_all
+        - wif captures how unit i influences the weight estimation
+        """
+        if not influence_func_info:
+            return 0.0
+
+        # Build unit index mapping
+        all_units = set()
+        for (g, t) in gt_pairs:
+            if (g, t) in influence_func_info:
+                info = influence_func_info[(g, t)]
+                all_units.update(info['treated_units'])
+                all_units.update(info['control_units'])
+
+        if not all_units:
+            return 0.0
+
+        all_units = sorted(all_units)
+        n_units = len(all_units)
+        unit_to_idx = {u: i for i, u in enumerate(all_units)}
+
+        # Get unique groups and their information
+        unique_groups = sorted(set(groups_for_gt))
+        group_to_idx = {g: i for i, g in enumerate(unique_groups)}
+
+        # Compute group-level probabilities matching R's formula:
+        # pg[g] = n_g / n_all (fraction of ALL units in group g)
+        # This differs from our old formula which used n_g / total_treated
+        group_sizes = {}
+        for g in unique_groups:
+            treated_in_g = df[df['first_treat'] == g][unit].nunique()
+            group_sizes[g] = treated_in_g
+
+        # pg indexed by group
+        pg_by_group = np.array([group_sizes[g] / n_units for g in unique_groups])
+
+        # pg indexed by keeper (each (g,t) pair gets its group's pg)
+        # This matches R's: pg <- pgg[match(group, originalglist)]
+        pg_keepers = np.array([pg_by_group[group_to_idx[g]] for g in groups_for_gt])
+        sum_pg_keepers = np.sum(pg_keepers)
+
+        # Standard aggregated influence (without wif)
+        psi_standard = np.zeros(n_units)
+
+        for j, (g, t) in enumerate(gt_pairs):
+            if (g, t) not in influence_func_info:
+                continue
+
+            info = influence_func_info[(g, t)]
+            w = weights[j]
+
+            for i, uid in enumerate(info['treated_units']):
+                idx = unit_to_idx[uid]
+                psi_standard[idx] += w * info['treated_inf'][i]
+
+            for i, uid in enumerate(info['control_units']):
+                idx = unit_to_idx[uid]
+                psi_standard[idx] += w * info['control_inf'][i]
+
+        # Build unit-group membership indicator
+        unit_groups = {}
+        for uid in all_units:
+            unit_first_treat = df[df[unit] == uid]['first_treat'].iloc[0]
+            if unit_first_treat in unique_groups:
+                unit_groups[uid] = unit_first_treat
+            else:
+                unit_groups[uid] = None  # Never-treated or other
+
+        # Compute wif using R's exact formula (iterate over keepers, not groups)
+        # R's wif function:
+        #   if1[i,k] = (indicator(G_i == group_k) - pg[k]) / sum(pg[keepers])
+        #   if2[i,k] = indicator_sum[i] * pg[k] / sum(pg[keepers])^2
+        #   wif[i,k] = if1[i,k] - if2[i,k]
+        #
+        # Then: wif_contrib[i] = sum_k(wif[i,k] * att[k])
+
+        n_keepers = len(gt_pairs)
+        wif_contrib = np.zeros(n_units)
+
+        # Pre-compute indicator_sum for each unit
+        # indicator_sum[i] = sum_k(indicator(G_i == group_k) - pg[k])
+        indicator_sum = np.zeros(n_units)
+        for j, g in enumerate(groups_for_gt):
+            pg_k = pg_keepers[j]
+            for uid in all_units:
+                i = unit_to_idx[uid]
+                unit_g = unit_groups[uid]
+                indicator = 1.0 if unit_g == g else 0.0
+                indicator_sum[i] += (indicator - pg_k)
+
+        # Compute wif contribution for each keeper
+        for j, (g, t) in enumerate(gt_pairs):
+            pg_k = pg_keepers[j]
+            att_k = effects[j]
+
+            for uid in all_units:
+                i = unit_to_idx[uid]
+                unit_g = unit_groups[uid]
+                indicator = 1.0 if unit_g == g else 0.0
+
+                # R's formula for wif
+                if1_ik = (indicator - pg_k) / sum_pg_keepers
+                if2_ik = indicator_sum[i] * pg_k / (sum_pg_keepers ** 2)
+                wif_ik = if1_ik - if2_ik
+
+                # Add contribution: wif[i,k] * att[k]
+                wif_contrib[i] += wif_ik * att_k
+
+        # Scale by 1/n_units to match R's getSE formula: sqrt(mean(IF^2)/n)
+        psi_wif = wif_contrib / n_units
+
+        # Combine standard and wif terms
+        psi_total = psi_standard + psi_wif
+
+        # Compute variance and SE
+        # R's formula: sqrt(mean(IF^2) / n) = sqrt(sum(IF^2) / n^2)
+        variance = np.sum(psi_total ** 2)
+        return np.sqrt(variance)
+
     def _aggregate_event_study(
         self,
         group_time_effects: Dict,

{diff_diff-2.0.0 → diff_diff-2.0.2}/diff_diff/sun_abraham.py

@@ -21,6 +21,7 @@ from diff_diff.results import _get_significance_stars
 from diff_diff.utils import (
     compute_confidence_interval,
     compute_p_value,
+    within_transform as _within_transform_util,
 )
 
 
@@ -789,28 +790,7 @@ class SunAbraham:
 
         y_it - y_i. - y_.t + y_..
         """
-        df = df.copy()
-
-        # Build all demeaned columns at once to avoid fragmentation
-        demeaned_data = {}
-        for var in variables:
-            # Unit means
-            unit_means = df.groupby(unit)[var].transform("mean")
-            # Time means
-            time_means = df.groupby(time)[var].transform("mean")
-            # Grand mean
-            grand_mean = df[var].mean()
-
-            # Within transformation
-            demeaned_data[f"{var}_dm"] = (
-                df[var] - unit_means - time_means + grand_mean
-            ).values
-
-        # Add all demeaned columns at once
-        demeaned_df = pd.DataFrame(demeaned_data, index=df.index)
-        df = pd.concat([df, demeaned_df], axis=1)
-
-        return df
+        return _within_transform_util(df, variables, unit, time, suffix="_dm")
 
     def _compute_iw_effects(
         self,

{diff_diff-2.0.0 → diff_diff-2.0.2}/diff_diff/twfe.py

@@ -17,6 +17,7 @@ from diff_diff.results import DiDResults
 from diff_diff.utils import (
     compute_confidence_interval,
     compute_p_value,
+    within_transform as _within_transform_util,
 )
 
 
@@ -211,25 +212,8 @@ class TwoWayFixedEffects(DifferenceInDifferences):
         pd.DataFrame
             Data with demeaned variables.
         """
-        data = data.copy()
         variables = [outcome] + (covariates or [])
-
-        # Cache groupby objects for efficiency (avoids re-computing group indexes)
-        unit_grouper = data.groupby(unit, sort=False)
-        time_grouper = data.groupby(time, sort=False)
-
-        for var in variables:
-            # Unit means (using cached grouper)
-            unit_means = unit_grouper[var].transform("mean")
-            # Time means (using cached grouper)
-            time_means = time_grouper[var].transform("mean")
-            # Grand mean
-            grand_mean = data[var].mean()
-
-            # Within transformation
-            data[f"{var}_demeaned"] = data[var] - unit_means - time_means + grand_mean
-
-        return data
+        return _within_transform_util(data, variables, unit, time, suffix="_demeaned")
 
     def _check_staggered_treatment(
         self,

{diff_diff-2.0.0 → diff_diff-2.0.2}/diff_diff/utils.py

@@ -1342,3 +1342,140 @@ def compute_placebo_effects(
         placebo_effects.append(placebo_tau)
 
     return np.asarray(placebo_effects)
+
+
+def demean_by_group(
+    data: pd.DataFrame,
+    variables: List[str],
+    group_var: str,
+    inplace: bool = False,
+    suffix: str = "",
+) -> Tuple[pd.DataFrame, int]:
+    """
+    Demean variables by a grouping variable (one-way within transformation).
+
+    For each variable, computes: x_ig - mean(x_g) where g is the group.
+
+    Parameters
+    ----------
+    data : pd.DataFrame
+        DataFrame containing the variables to demean.
+    variables : list of str
+        Column names to demean.
+    group_var : str
+        Column name for the grouping variable.
+    inplace : bool, default False
+        If True, modifies the original columns. If False, leaves original
+        columns unchanged (demeaning is still applied to return value).
+    suffix : str, default ""
+        Suffix to add to demeaned column names (only used when inplace=False
+        and you want to keep both original and demeaned columns).
+
+    Returns
+    -------
+    data : pd.DataFrame
+        DataFrame with demeaned variables.
+    n_effects : int
+        Number of absorbed fixed effects (nunique - 1).
+
+    Examples
+    --------
+    >>> df, n_fe = demean_by_group(df, ['y', 'x1', 'x2'], 'unit')
+    >>> # df['y'], df['x1'], df['x2'] are now demeaned by unit
+    """
+    if not inplace:
+        data = data.copy()
+
+    # Count fixed effects (categories - 1 for identification)
+    n_effects = data[group_var].nunique() - 1
+
+    # Cache the groupby object for efficiency
+    grouper = data.groupby(group_var, sort=False)
+
+    for var in variables:
+        col_name = var if not suffix else f"{var}{suffix}"
+        group_means = grouper[var].transform("mean")
+        data[col_name] = data[var] - group_means
+
+    return data, n_effects
+
+
+def within_transform(
+    data: pd.DataFrame,
+    variables: List[str],
+    unit: str,
+    time: str,
+    inplace: bool = False,
+    suffix: str = "_demeaned",
+) -> pd.DataFrame:
+    """
+    Apply two-way within transformation to remove unit and time fixed effects.
+
+    Computes: y_it - y_i. - y_.t + y_.. for each variable.
+
+    This is the standard fixed effects transformation for panel data that
+    removes both unit-specific and time-specific effects.
+
+    Parameters
+    ----------
+    data : pd.DataFrame
+        Panel data containing the variables to transform.
+    variables : list of str
+        Column names to transform.
+    unit : str
+        Column name for unit identifier.
+    time : str
+        Column name for time period identifier.
+    inplace : bool, default False
+        If True, modifies the original columns. If False, creates new columns
+        with the specified suffix.
+    suffix : str, default "_demeaned"
+        Suffix for new column names when inplace=False.
+
+    Returns
+    -------
+    pd.DataFrame
+        DataFrame with within-transformed variables.
+
+    Notes
+    -----
+    The within transformation removes variation that is constant within units
+    (unit fixed effects) and constant within time periods (time fixed effects).
+    The resulting estimates are equivalent to including unit and time dummies
+    but is computationally more efficient for large panels.
+
+    Examples
+    --------
+    >>> df = within_transform(df, ['y', 'x'], 'unit_id', 'year')
+    >>> # df now has 'y_demeaned' and 'x_demeaned' columns
+    """
+    if not inplace:
+        data = data.copy()
+
+    # Cache groupby objects for efficiency
+    unit_grouper = data.groupby(unit, sort=False)
+    time_grouper = data.groupby(time, sort=False)
+
+    if inplace:
+        # Modify columns in place
+        for var in variables:
+            unit_means = unit_grouper[var].transform("mean")
+            time_means = time_grouper[var].transform("mean")
+            grand_mean = data[var].mean()
+            data[var] = data[var] - unit_means - time_means + grand_mean
+    else:
+        # Build all demeaned columns at once to avoid DataFrame fragmentation
+        demeaned_data = {}
+        for var in variables:
+            unit_means = unit_grouper[var].transform("mean")
+            time_means = time_grouper[var].transform("mean")
+            grand_mean = data[var].mean()
+            demeaned_data[f"{var}{suffix}"] = (
+                data[var] - unit_means - time_means + grand_mean
+            ).values
+
+        # Add all columns at once
+        demeaned_df = pd.DataFrame(demeaned_data, index=data.index)
+        data = pd.concat([data, demeaned_df], axis=1)
+
+    return data

{diff_diff-2.0.0 → diff_diff-2.0.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "diff-diff"
-version = "2.0.0"
+version = "2.0.2"
 description = "A library for Difference-in-Differences causal inference analysis"
 readme = "README.md"
 license = "MIT"

{diff_diff-2.0.0 → diff_diff-2.0.2}/rust/Cargo.lock

@@ -131,9 +131,9 @@ checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
 
 [[package]]
 name = "chrono"
-version = "0.4.42"
+version = "0.4.43"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "145052bdd345b87320e369255277e3fb5152762ad123a901ef5c262dd38fe8d2"
+checksum = "fac4744fb15ae8337dc853fee7fb3f4e48c0fbaa23d0afe49c447b4fab126118"
 dependencies = [
  "iana-time-zone",
  "js-sys",
@@ -709,9 +709,9 @@ checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2"
 
 [[package]]
 name = "js-sys"
-version = "0.3.83"
+version = "0.3.85"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "464a3709c7f55f1f721e5389aa6ea4e3bc6aba669353300af094b29ffbdde1d8"
+checksum = "8c942ebf8e95485ca0d52d97da7c5a2c387d0e7f0ba4c35e93bfcaee045955b3"
 dependencies = [
  "once_cell",
  "wasm-bindgen",
@@ -1895,9 +1895,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen"
-version = "0.2.106"
+version = "0.2.108"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0d759f433fa64a2d763d1340820e46e111a7a5ab75f993d1852d70b03dbb80fd"
+checksum = "64024a30ec1e37399cf85a7ffefebdb72205ca1c972291c51512360d90bd8566"
 dependencies = [
  "cfg-if",
  "once_cell",
@@ -1908,9 +1908,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-macro"
-version = "0.2.106"
+version = "0.2.108"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "48cb0d2638f8baedbc542ed444afc0644a29166f1595371af4fecf8ce1e7eeb3"
+checksum = "008b239d9c740232e71bd39e8ef6429d27097518b6b30bdf9086833bd5b6d608"
 dependencies = [
  "quote",
  "wasm-bindgen-macro-support",
@@ -1918,9 +1918,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-macro-support"
-version = "0.2.106"
+version = "0.2.108"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "cefb59d5cd5f92d9dcf80e4683949f15ca4b511f4ac0a6e14d4e1ac60c6ecd40"
+checksum = "5256bae2d58f54820e6490f9839c49780dff84c65aeab9e772f15d5f0e913a55"
 dependencies = [
  "bumpalo",
  "proc-macro2",
@@ -1931,9 +1931,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-shared"
-version = "0.2.106"
+version = "0.2.108"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "cbc538057e648b67f72a982e708d485b2efa771e1ac05fec311f9f63e5800db4"
+checksum = "1f01b580c9ac74c8d8f0c0e4afb04eeef2acf145458e52c03845ee9cd23e3d12"
 dependencies = [
  "unicode-ident",
 ]
@@ -2391,6 +2391,6 @@ dependencies = [
 
 [[package]]
 name = "zmij"
-version = "1.0.13"
+version = "1.0.14"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ac93432f5b761b22864c774aac244fa5c0fd877678a4c37ebf6cf42208f9c9ec"
+checksum = "bd8f3f50b848df28f887acb68e41201b5aea6bc8a8dacc00fb40635ff9a72fea"