diff-diff 1.4.0__tar.gz → 2.0.1__tar.gz

{diff_diff-1.4.0 → diff_diff-2.0.1}/PKG-INFO

@@ -1,14 +1,6 @@
 Metadata-Version: 2.4
 Name: diff-diff
-Version: 1.4.0
-Summary: A library for Difference-in-Differences causal inference analysis
-Author: diff-diff contributors
-License-Expression: MIT
-Project-URL: Homepage, https://github.com/igerber/diff-diff
-Project-URL: Documentation, https://diff-diff.readthedocs.io
-Project-URL: Repository, https://github.com/igerber/diff-diff
-Project-URL: Issues, https://github.com/igerber/diff-diff/issues
-Keywords: causal-inference,difference-in-differences,econometrics,statistics,treatment-effects
+Version: 2.0.1
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Science/Research
 Classifier: Operating System :: OS Independent
@@ -18,20 +10,28 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: Scientific/Engineering :: Mathematics
-Requires-Python: >=3.9
-Description-Content-Type: text/markdown
 Requires-Dist: numpy>=1.20.0
 Requires-Dist: pandas>=1.3.0
 Requires-Dist: scipy>=1.7.0
+Requires-Dist: pytest>=7.0 ; extra == 'dev'
+Requires-Dist: pytest-cov>=4.0 ; extra == 'dev'
+Requires-Dist: black>=23.0 ; extra == 'dev'
+Requires-Dist: ruff>=0.1.0 ; extra == 'dev'
+Requires-Dist: mypy>=1.0 ; extra == 'dev'
+Requires-Dist: sphinx>=6.0 ; extra == 'docs'
+Requires-Dist: sphinx-rtd-theme>=1.0 ; extra == 'docs'
 Provides-Extra: dev
-Requires-Dist: pytest>=7.0; extra == "dev"
-Requires-Dist: pytest-cov>=4.0; extra == "dev"
-Requires-Dist: black>=23.0; extra == "dev"
-Requires-Dist: ruff>=0.1.0; extra == "dev"
-Requires-Dist: mypy>=1.0; extra == "dev"
 Provides-Extra: docs
-Requires-Dist: sphinx>=6.0; extra == "docs"
-Requires-Dist: sphinx-rtd-theme>=1.0; extra == "docs"
+Summary: A library for Difference-in-Differences causal inference analysis
+Keywords: causal-inference,difference-in-differences,econometrics,statistics,treatment-effects
+Author: diff-diff contributors
+License-Expression: MIT
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
+Project-URL: Documentation, https://diff-diff.readthedocs.io
+Project-URL: Homepage, https://github.com/igerber/diff-diff
+Project-URL: Issues, https://github.com/igerber/diff-diff/issues
+Project-URL: Repository, https://github.com/igerber/diff-diff
 
 # diff-diff
 
@@ -2253,3 +2253,4 @@ The `HonestDiD` module implements sensitivity analysis methods for relaxing the
 ## License
 
 MIT License
+

{diff_diff-1.4.0 → diff_diff-2.0.1}/diff_diff/__init__.py

@@ -5,6 +5,16 @@ This library provides sklearn-like estimators for causal inference
 using the difference-in-differences methodology.
 """
 
+# Import backend detection from dedicated module (avoids circular imports)
+from diff_diff._backend import (
+    HAS_RUST_BACKEND,
+    _rust_bootstrap_weights,
+    _rust_compute_robust_vcov,
+    _rust_project_simplex,
+    _rust_solve_ols,
+    _rust_synthetic_weights,
+)
+
 from diff_diff.bacon import (
     BaconDecomposition,
     BaconDecompositionResults,
@@ -103,7 +113,7 @@ from diff_diff.visualization import (
     plot_sensitivity,
 )
 
-__version__ = "1.4.0"
+__version__ = "2.0.1"
 __all__ = [
     # Estimators
     "DifferenceInDifferences",
@@ -187,4 +197,6 @@ __all__ = [
     "compute_pretrends_power",
     "compute_mdv",
     "plot_pretrends_power",
+    # Rust backend
+    "HAS_RUST_BACKEND",
 ]

diff_diff-2.0.1/diff_diff/_backend.py

@@ -0,0 +1,64 @@
+"""
+Backend detection and configuration for diff-diff.
+
+This module handles:
+1. Detection of optional Rust backend
+2. Environment variable configuration (DIFF_DIFF_BACKEND)
+3. Exports HAS_RUST_BACKEND and Rust function references
+
+Other modules should import from here to avoid circular imports with __init__.py.
+"""
+
+import os
+
+# Check for backend override via environment variable
+# DIFF_DIFF_BACKEND can be: 'auto' (default), 'python', or 'rust'
+_backend_env = os.environ.get('DIFF_DIFF_BACKEND', 'auto').lower()
+
+# Try to import Rust backend for accelerated operations
+try:
+    from diff_diff._rust_backend import (
+        generate_bootstrap_weights_batch as _rust_bootstrap_weights,
+        compute_synthetic_weights as _rust_synthetic_weights,
+        project_simplex as _rust_project_simplex,
+        solve_ols as _rust_solve_ols,
+        compute_robust_vcov as _rust_compute_robust_vcov,
+    )
+    _rust_available = True
+except ImportError:
+    _rust_available = False
+    _rust_bootstrap_weights = None
+    _rust_synthetic_weights = None
+    _rust_project_simplex = None
+    _rust_solve_ols = None
+    _rust_compute_robust_vcov = None
+
+# Determine final backend based on environment variable and availability
+if _backend_env == 'python':
+    # Force pure Python mode - disable Rust even if available
+    HAS_RUST_BACKEND = False
+    _rust_bootstrap_weights = None
+    _rust_synthetic_weights = None
+    _rust_project_simplex = None
+    _rust_solve_ols = None
+    _rust_compute_robust_vcov = None
+elif _backend_env == 'rust':
+    # Force Rust mode - fail if not available
+    if not _rust_available:
+        raise ImportError(
+            "DIFF_DIFF_BACKEND=rust but Rust backend is not available. "
+            "Install with: pip install diff-diff[rust]"
+        )
+    HAS_RUST_BACKEND = True
+else:
+    # Auto mode - use Rust if available
+    HAS_RUST_BACKEND = _rust_available
+
+__all__ = [
+    'HAS_RUST_BACKEND',
+    '_rust_bootstrap_weights',
+    '_rust_synthetic_weights',
+    '_rust_project_simplex',
+    '_rust_solve_ols',
+    '_rust_compute_robust_vcov',
+]

{diff_diff-1.4.0 → diff_diff-2.0.1}/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-1.4.0 → diff_diff-2.0.1}/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-1.4.0 → diff_diff-2.0.1}/diff_diff/linalg.py

@@ -1,15 +1,17 @@
 """
 Unified linear algebra backend for diff-diff.
 
-This module provides optimized OLS and variance estimation that can be
-swapped to a compiled backend (Rust/C++) for maximum performance.
+This module provides optimized OLS and variance estimation with an optional
+Rust backend for maximum performance.
 
 The key optimizations are:
 1. scipy.linalg.lstsq with 'gelsy' driver (QR-based, faster than SVD)
 2. Vectorized cluster-robust SE via groupby (eliminates O(n*clusters) loop)
 3. Single interface for all estimators (reduces code duplication)
+4. Optional Rust backend for additional speedup (when available)
 
-Future: This module can be extended with a Rust backend for additional speedup.
+The Rust backend is automatically used when available, with transparent
+fallback to NumPy/SciPy implementations.
 """
 
 from typing import Optional, Tuple, Union
@@ -18,6 +20,13 @@ import numpy as np
 import pandas as pd
 from scipy.linalg import lstsq as scipy_lstsq
 
+# Import Rust backend if available (from _backend to avoid circular imports)
+from diff_diff._backend import (
+    HAS_RUST_BACKEND,
+    _rust_compute_robust_vcov,
+    _rust_solve_ols,
+)
+
 
 def solve_ols(
     X: np.ndarray,
@@ -119,6 +128,87 @@ def solve_ols(
                 "Clean your data or set check_finite=False to skip this check."
             )
 
+    # Use Rust backend if available
+    # Note: Fall back to NumPy if check_finite=False since Rust's LAPACK
+    # doesn't support non-finite values
+    if HAS_RUST_BACKEND and check_finite:
+        # Ensure contiguous arrays for Rust
+        X = np.ascontiguousarray(X, dtype=np.float64)
+        y = np.ascontiguousarray(y, dtype=np.float64)
+
+        # Convert cluster_ids to int64 for Rust (if provided)
+        cluster_ids_int = None
+        if cluster_ids is not None:
+            cluster_ids_int = pd.factorize(cluster_ids)[0].astype(np.int64)
+
+        try:
+            coefficients, residuals, vcov = _rust_solve_ols(
+                X, y, cluster_ids_int, return_vcov
+            )
+        except ValueError as e:
+            # Translate Rust LAPACK errors to consistent Python error messages
+            error_msg = str(e)
+            if "Matrix inversion failed" in error_msg or "Least squares failed" in error_msg:
+                raise ValueError(
+                    "Design matrix is rank-deficient (singular X'X matrix). "
+                    "This indicates perfect multicollinearity. Check your fixed effects "
+                    "and covariates for linear dependencies."
+                ) from e
+            raise
+
+        if return_fitted:
+            fitted = X @ coefficients
+            return coefficients, residuals, fitted, vcov
+        else:
+            return coefficients, residuals, vcov
+
+    # Fallback to NumPy/SciPy implementation
+    return _solve_ols_numpy(
+        X, y, cluster_ids=cluster_ids, return_vcov=return_vcov, return_fitted=return_fitted
+    )
+
+
+def _solve_ols_numpy(
+    X: np.ndarray,
+    y: np.ndarray,
+    *,
+    cluster_ids: Optional[np.ndarray] = None,
+    return_vcov: bool = True,
+    return_fitted: bool = False,
+) -> Union[
+    Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]],
+    Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]],
+]:
+    """
+    NumPy/SciPy fallback implementation of solve_ols.
+
+    Uses scipy.linalg.lstsq with 'gelsy' driver (QR with column pivoting)
+    for fast and stable least squares solving.
+
+    Parameters
+    ----------
+    X : np.ndarray
+        Design matrix of shape (n, k).
+    y : np.ndarray
+        Response vector of shape (n,).
+    cluster_ids : np.ndarray, optional
+        Cluster identifiers for cluster-robust SEs.
+    return_vcov : bool
+        Whether to compute variance-covariance matrix.
+    return_fitted : bool
+        Whether to return fitted values.
+
+    Returns
+    -------
+    coefficients : np.ndarray
+        OLS coefficients of shape (k,).
+    residuals : np.ndarray
+        Residuals of shape (n,).
+    fitted : np.ndarray, optional
+        Fitted values if return_fitted=True.
+    vcov : np.ndarray, optional
+        Variance-covariance matrix if return_vcov=True.
+    """
     # Solve OLS using scipy's optimized solver
     # 'gelsy' uses QR with column pivoting, faster than default 'gelsd' (SVD)
     # Note: gelsy doesn't reliably report rank, so we don't check for deficiency
@@ -131,7 +221,7 @@ def solve_ols(
     # Compute variance-covariance matrix if requested
     vcov = None
     if return_vcov:
-        vcov = compute_robust_vcov(X, residuals, cluster_ids)
+        vcov = _compute_robust_vcov_numpy(X, residuals, cluster_ids)
 
     if return_fitted:
         return coefficients, residuals, fitted, vcov
@@ -176,6 +266,63 @@ def compute_robust_vcov(
     The cluster-robust computation is vectorized using pandas groupby,
     which is much faster than a Python loop over clusters.
     """
+    # Use Rust backend if available
+    if HAS_RUST_BACKEND:
+        X = np.ascontiguousarray(X, dtype=np.float64)
+        residuals = np.ascontiguousarray(residuals, dtype=np.float64)
+
+        cluster_ids_int = None
+        if cluster_ids is not None:
+            cluster_ids_int = pd.factorize(cluster_ids)[0].astype(np.int64)
+
+        try:
+            return _rust_compute_robust_vcov(X, residuals, cluster_ids_int)
+        except ValueError as e:
+            # Translate Rust LAPACK errors to consistent Python error messages
+            error_msg = str(e)
+            if "Matrix inversion failed" in error_msg:
+                raise ValueError(
+                    "Design matrix is rank-deficient (singular X'X matrix). "
+                    "This indicates perfect multicollinearity. Check your fixed effects "
+                    "and covariates for linear dependencies."
+                ) from e
+            raise
+
+    # Fallback to NumPy implementation
+    return _compute_robust_vcov_numpy(X, residuals, cluster_ids)
+
+
+def _compute_robust_vcov_numpy(
+    X: np.ndarray,
+    residuals: np.ndarray,
+    cluster_ids: Optional[np.ndarray] = None,
+) -> np.ndarray:
+    """
+    NumPy fallback implementation of compute_robust_vcov.
+
+    Computes HC1 (heteroskedasticity-robust) or cluster-robust variance-covariance
+    matrix using the sandwich estimator.
+
+    Parameters
+    ----------
+    X : np.ndarray
+        Design matrix of shape (n, k).
+    residuals : np.ndarray
+        OLS residuals of shape (n,).
+    cluster_ids : np.ndarray, optional
+        Cluster identifiers. If None, uses HC1. If provided, uses
+        cluster-robust with G/(G-1) small-sample adjustment.
+
+    Returns
+    -------
+    vcov : np.ndarray
+        Variance-covariance matrix of shape (k, k).
+
+    Notes
+    -----
+    Uses vectorized groupby aggregation for cluster-robust SEs to avoid
+    the O(n * G) loop that would be required with explicit iteration.
+    """
     n, k = X.shape
     XtX = X.T @ X