diff-diff 2.1.8__tar.gz → 2.2.0__tar.gz

{diff_diff-2.1.8 → diff_diff-2.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diff-diff
-Version: 2.1.8
+Version: 2.2.0
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Science/Research
 Classifier: Operating System :: OS Independent
@@ -1302,6 +1302,7 @@ trop = TROP(
 
 ```python
 TROP(
+    method='twostep',           # Estimation method: 'twostep' (default) or 'joint'
     lambda_time_grid=None,      # Time decay grid (default: [0, 0.1, 0.5, 1, 2, 5])
     lambda_unit_grid=None,      # Unit distance grid (default: [0, 0.1, 0.5, 1, 2, 5])
     lambda_nn_grid=None,        # Nuclear norm grid (default: [0, 0.01, 0.1, 1, 10])
@@ -1314,6 +1315,10 @@ TROP(
 )
 ```
 
+**Estimation methods:**
+- `'twostep'` (default): Per-observation model fitting following Algorithm 2 of the paper. Computes observation-specific weights and fits a model for each treated observation, then averages the individual treatment effects. More flexible but computationally intensive.
+- `'joint'`: Joint weighted least squares optimization. Estimates a single scalar treatment effect τ along with fixed effects and optional low-rank factor adjustment. Faster but assumes homogeneous treatment effects.
+
 **Convenience function:**
 
 ```python

{diff_diff-2.1.8 → diff_diff-2.2.0}/README.md

@@ -1267,6 +1267,7 @@ trop = TROP(
 
 ```python
 TROP(
+    method='twostep',           # Estimation method: 'twostep' (default) or 'joint'
     lambda_time_grid=None,      # Time decay grid (default: [0, 0.1, 0.5, 1, 2, 5])
     lambda_unit_grid=None,      # Unit distance grid (default: [0, 0.1, 0.5, 1, 2, 5])
     lambda_nn_grid=None,        # Nuclear norm grid (default: [0, 0.01, 0.1, 1, 10])
@@ -1279,6 +1280,10 @@ TROP(
 )
 ```
 
+**Estimation methods:**
+- `'twostep'` (default): Per-observation model fitting following Algorithm 2 of the paper. Computes observation-specific weights and fits a model for each treated observation, then averages the individual treatment effects. More flexible but computationally intensive.
+- `'joint'`: Joint weighted least squares optimization. Estimates a single scalar treatment effect τ along with fixed effects and optional low-rank factor adjustment. Faster but assumes homogeneous treatment effects.
+
 **Convenience function:**
 
 ```python

{diff_diff-2.1.8 → diff_diff-2.2.0}/diff_diff/__init__.py

@@ -136,7 +136,7 @@ from diff_diff.datasets import (
     load_mpdta,
 )
 
-__version__ = "2.1.8"
+__version__ = "2.2.0"
 __all__ = [
     # Estimators
     "DifferenceInDifferences",

{diff_diff-2.1.8 → diff_diff-2.2.0}/diff_diff/_backend.py

@@ -23,10 +23,13 @@ try:
         project_simplex as _rust_project_simplex,
         solve_ols as _rust_solve_ols,
         compute_robust_vcov as _rust_compute_robust_vcov,
-        # TROP estimator acceleration
+        # TROP estimator acceleration (twostep method)
         compute_unit_distance_matrix as _rust_unit_distance_matrix,
         loocv_grid_search as _rust_loocv_grid_search,
         bootstrap_trop_variance as _rust_bootstrap_trop_variance,
+        # TROP estimator acceleration (joint method)
+        loocv_grid_search_joint as _rust_loocv_grid_search_joint,
+        bootstrap_trop_variance_joint as _rust_bootstrap_trop_variance_joint,
     )
     _rust_available = True
 except ImportError:
@@ -36,10 +39,13 @@ except ImportError:
     _rust_project_simplex = None
     _rust_solve_ols = None
     _rust_compute_robust_vcov = None
-    # TROP estimator acceleration
+    # TROP estimator acceleration (twostep method)
     _rust_unit_distance_matrix = None
     _rust_loocv_grid_search = None
     _rust_bootstrap_trop_variance = None
+    # TROP estimator acceleration (joint method)
+    _rust_loocv_grid_search_joint = None
+    _rust_bootstrap_trop_variance_joint = None
 
 # Determine final backend based on environment variable and availability
 if _backend_env == 'python':
@@ -50,10 +56,13 @@ if _backend_env == 'python':
     _rust_project_simplex = None
     _rust_solve_ols = None
     _rust_compute_robust_vcov = None
-    # TROP estimator acceleration
+    # TROP estimator acceleration (twostep method)
     _rust_unit_distance_matrix = None
     _rust_loocv_grid_search = None
     _rust_bootstrap_trop_variance = None
+    # TROP estimator acceleration (joint method)
+    _rust_loocv_grid_search_joint = None
+    _rust_bootstrap_trop_variance_joint = None
 elif _backend_env == 'rust':
     # Force Rust mode - fail if not available
     if not _rust_available:
@@ -73,8 +82,11 @@ __all__ = [
     '_rust_project_simplex',
     '_rust_solve_ols',
     '_rust_compute_robust_vcov',
-    # TROP estimator acceleration
+    # TROP estimator acceleration (twostep method)
     '_rust_unit_distance_matrix',
     '_rust_loocv_grid_search',
     '_rust_bootstrap_trop_variance',
+    # TROP estimator acceleration (joint method)
+    '_rust_loocv_grid_search_joint',
+    '_rust_bootstrap_trop_variance_joint',
 ]

{diff_diff-2.1.8 → diff_diff-2.2.0}/diff_diff/linalg.py

@@ -251,10 +251,10 @@ def _solve_ols_rust(
     cluster_ids: Optional[np.ndarray] = None,
     return_vcov: bool = True,
     return_fitted: bool = False,
-) -> Union[
+) -> Optional[Union[
     Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]],
     Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]],
-]:
+]]:
     """
     Rust backend implementation of solve_ols for full-rank matrices.
 
@@ -296,15 +296,30 @@ def _solve_ols_rust(
         Fitted values if return_fitted=True.
     vcov : np.ndarray, optional
         Variance-covariance matrix if return_vcov=True.
+    None
+        If Rust backend detects numerical instability and caller should
+        fall back to Python backend.
     """
     # Convert cluster_ids to int64 for Rust (handles string/categorical IDs)
     if cluster_ids is not None:
         cluster_ids = _factorize_cluster_ids(cluster_ids)
 
-    # Call Rust backend
-    coefficients, residuals, vcov = _rust_solve_ols(
-        X, y, cluster_ids=cluster_ids, return_vcov=return_vcov
-    )
+    # Call Rust backend with fallback on numerical instability
+    try:
+        coefficients, residuals, vcov = _rust_solve_ols(
+            X, y, cluster_ids=cluster_ids, return_vcov=return_vcov
+        )
+    except ValueError as e:
+        error_msg = str(e).lower()
+        if "numerically unstable" in error_msg or "singular" in error_msg:
+            warnings.warn(
+                f"Rust backend detected numerical instability: {e}. "
+                "Falling back to Python backend.",
+                UserWarning,
+                stacklevel=3,
+            )
+            return None  # Signal caller to use Python fallback
+        raise
 
     # Convert to numpy arrays
     coefficients = np.asarray(coefficients)
@@ -468,12 +483,15 @@ def solve_ols(
     # This saves O(nk²) QR overhead but won't detect rank-deficient matrices
     if skip_rank_check:
         if HAS_RUST_BACKEND and _rust_solve_ols is not None:
-            return _solve_ols_rust(
+            result = _solve_ols_rust(
                 X, y,
                 cluster_ids=cluster_ids,
                 return_vcov=return_vcov,
                 return_fitted=return_fitted,
             )
+            if result is not None:
+                return result
+            # Fall through to NumPy on numerical instability
         # Fall through to Python without rank check (user guarantees full rank)
         return _solve_ols_numpy(
             X, y,
@@ -499,6 +517,7 @@ def solve_ols(
     # Routing strategy:
     # - Full-rank + Rust available → fast Rust backend (SVD-based solve)
     # - Rank-deficient → Python backend (proper NA handling, valid SEs)
+    # - Rust numerical instability → Python fallback (via None return)
     # - No Rust → Python backend (works for all cases)
     if HAS_RUST_BACKEND and _rust_solve_ols is not None and not is_rank_deficient:
         result = _solve_ols_rust(
@@ -508,6 +527,19 @@ def solve_ols(
             return_fitted=return_fitted,
         )
 
+        # Check for None: Rust backend detected numerical instability and
+        # signaled us to fall back to Python backend
+        if result is None:
+            return _solve_ols_numpy(
+                X, y,
+                cluster_ids=cluster_ids,
+                return_vcov=return_vcov,
+                return_fitted=return_fitted,
+                rank_deficient_action=rank_deficient_action,
+                column_names=column_names,
+                _precomputed_rank_info=None,  # Force fresh rank detection
+            )
+
         # Check for NaN vcov: Rust SVD may detect rank-deficiency that QR missed
         # for ill-conditioned matrices (QR and SVD have different numerical properties).
         # When this happens, fall back to Python's R-style handling.
@@ -732,7 +764,7 @@ def compute_robust_vcov(
         try:
             return _rust_compute_robust_vcov(X, residuals, cluster_ids_int)
         except ValueError as e:
-            # Translate Rust LAPACK errors to consistent Python error messages
+            # Translate Rust errors to consistent Python error messages or fallback
             error_msg = str(e)
             if "Matrix inversion failed" in error_msg:
                 raise ValueError(
@@ -740,6 +772,15 @@ def compute_robust_vcov(
                     "This indicates perfect multicollinearity. Check your fixed effects "
                     "and covariates for linear dependencies."
                 ) from e
+            if "numerically unstable" in error_msg.lower():
+                # Fall back to NumPy on numerical instability (with warning)
+                warnings.warn(
+                    f"Rust backend detected numerical instability: {e}. "
+                    "Falling back to Python backend for variance computation.",
+                    UserWarning,
+                    stacklevel=2,
+                )
+                return _compute_robust_vcov_numpy(X, residuals, cluster_ids)
             raise
 
     # Fallback to NumPy implementation