diff --git a/.idea/misc.xml b/.idea/misc.xml
index 1ffff1f..99bc5d6 100644
--- a/.idea/misc.xml
+++ b/.idea/misc.xml
@@ -3,6 +3,7 @@
+
+
pyrdd vs rdrobust — Runtime (log scale)
+Runtime (log scale)
0.1ms
1s
10s
-Senate (n≈1,390)
+Senate (n=1,390)
Manual bandwidth
-6.2ms
-0.47ms (13×)
+5.4ms
+0.48ms (11×)
MSE bandwidth
-28ms
-1.6ms (18×)
+25ms
+1.7ms (15×)
MSE bandwidth + clustered SE
-33ms
-2.2ms (15×)
+29ms
+2.0ms (15×)
MSE bandwidth + clustered SE + covariates
-70ms
-3.0ms (24×)
-Synthetic (n=20,000)
+67ms
+2.8ms (24×)
+Synthetic (n=20,000)
Manual bandwidth
-59ms
-3.0ms (20×)
+50ms
+2.8ms (18×)
MSE bandwidth
-283ms
-8.8ms (32×)
+245ms
+8.7ms (28×)
MSE bandwidth + clustered SE
-302ms
-12ms (25×)
+267ms
+12ms (23×)
MSE bandwidth + clustered SE + covariates
-2.6s
-18ms (143×)
+2.5s
+17ms (145×)
rdrobust
'
- f"pyrdd vs rdrobust \u2014 Runtime (log scale) "
+ f"Runtime (log scale)"
)
# Gridlines
@@ -358,7 +261,7 @@ def x_pos(ms: float) -> float:
lines.append(
f''
+ f'font-weight="600" fill="#1f2937">'
f"{_escape_xml(r.group)} "
)
diff --git a/pyrdd/bandwidth.py b/pyrdd/bandwidth.py
index 685690d..1d0f8c1 100644
--- a/pyrdd/bandwidth.py
+++ b/pyrdd/bandwidth.py
@@ -34,7 +34,12 @@
from pyrdd.internals.data import Data, EstimationSample
from pyrdd.internals.fit import PolynomialFit
-from pyrdd.internals.misc import gram_inverse, rdrobust_res, rdrobust_vce
+from pyrdd.internals.misc import (
+ gram_inverse,
+ hc_residuals,
+ nn_residuals,
+ variance_covariance_matrix,
+)
from pyrdd.internals.types import (
BandwidthSelector,
Kernel,
@@ -52,11 +57,6 @@
}
-# ===================================================================
-# Data containers
-# ===================================================================
-
-
@dataclass(frozen=True, slots=True)
class _BandwidthComponents:
"""Variance, bias, regularisation, and rate exponent for one side.
@@ -86,17 +86,12 @@ class _VarianceFit:
rate: float # 1/(2p+3)
-# ===================================================================
-# Preliminary cache (one per side, reused across all plug-in rounds)
-# ===================================================================
-
-
@dataclass(slots=True)
class _PreliminaryCache:
"""Pre-computed quantities at the preliminary (order-0) bandwidth.
The variance window is identical across all calls to
- ``_compute_bandwidth_components`` because they all use the same
+ `_compute_bandwidth_components` because they all use the same
preliminary bandwidth. This cache stores the estimation sample,
Vandermonde columns, NN residuals, and memoised variance-side fits.
"""
@@ -160,23 +155,16 @@ def _build_preliminary_cache(
)
# NN residuals are independent of polynomial order — compute once.
- # CCT14a, Section 5: fixed-matches NN variance estimator.
- nn_residuals: np.ndarray | None = None
if vce == "nn":
- nn_residuals = rdrobust_res(
- sample.running,
- sample.outcome,
- sample.fuzzy,
- sample.covariates,
- 0,
- 0,
- "nn",
- nnmatch,
- sample.duplicate_count,
- sample.duplicate_index,
- 0,
+ residuals = nn_residuals(
+ running=sample.running,
+ responses=sample.responses,
+ duplicate_count=sample.duplicate_count,
+ duplicate_index=sample.duplicate_index,
+ matches=nnmatch,
)
-
+ else:
+ residuals = None
return _PreliminaryCache(
sample=sample,
responses=responses,
@@ -185,7 +173,7 @@ def _build_preliminary_cache(
preliminary_bandwidth=preliminary_bandwidth,
n_fuzzy_cols=n_fuzzy_cols,
n_covariate_cols=n_covariate_cols,
- nn_residuals=nn_residuals,
+ nn_residuals=residuals,
data=data,
vce=vce,
nnmatch=nnmatch,
@@ -194,11 +182,6 @@ def _build_preliminary_cache(
)
-# ===================================================================
-# Influence weights (delta method for fuzzy/covariate designs)
-# ===================================================================
-
-
def _compute_influence_weights(
cache: _PreliminaryCache,
fit: PolynomialFit,
@@ -282,11 +265,6 @@ def _compute_influence_weights(
return s
-# ===================================================================
-# Variance-side computation (cached by poly_order, deriv_order)
-# ===================================================================
-
-
def _compute_variance_fit(
cache: _PreliminaryCache,
poly_order: int,
@@ -316,68 +294,50 @@ def _compute_variance_fit(
# Residuals: NN (CCT14a, Section 5) or HC variants
if cache.vce == "nn":
assert cache.nn_residuals is not None
- residuals = cache.nn_residuals
+ res = cache.nn_residuals
else:
fitted_values = regressors @ beta
- hat_diagonal: np.ndarray | int = 0
+ hat_diagonal: np.ndarray | None = None
if cache.vce in ("hc2", "hc3"):
hat_diagonal = np.sum(
(regressors @ gram_inv) * (regressors * weights),
axis=1,
)
- residuals = rdrobust_res(
- sample.running,
- sample.outcome,
- sample.fuzzy,
- sample.covariates,
- fitted_values,
- hat_diagonal,
- cache.vce,
- cache.nnmatch,
- sample.duplicate_count,
- sample.duplicate_index,
- p + 1,
+ res = hc_residuals(
+ responses=cache.responses,
+ fitted_values=fitted_values,
+ hat_diagonal=hat_diagonal,
+ vce=cache.vce,
+ n_params=p + 1,
)
# Sandwich VCE: V̂_ν = ê_ν' G⁻¹ Σ̂ G⁻¹ ê_ν
# CCT14a, Section 5, eq. 5.1 (NN) / eq. 5.3 (HC)
weighted_regressors = regressors * weights
- sandwich_meat = rdrobust_vce(
- cache.n_fuzzy_cols + cache.n_covariate_cols,
- influence_weights,
- weighted_regressors,
- residuals,
- sample.cluster,
+ sandwich_meat = variance_covariance_matrix(
+ influence_weights=influence_weights,
+ weighted_design_matrix=weighted_regressors,
+ residuals=res,
+ clusters=sample.cluster,
)
V_nu = float((gram_inv @ sandwich_meat @ gram_inv)[nu, nu])
# Bias constant: ê_ν' H_p (G⁻¹ R'W u^{p+1})
# CCT14a, Lemma A.1(B) — the "leakage" of the (p+1)th power
- leakage = weighted_regressors.T @ (
- (cache.centered_running / h).reshape(-1, 1) ** (p + 1)
- )
+ leakage = weighted_regressors.T @ (cache.centered_running / h) ** (p + 1)
bandwidth_powers = np.array([h**j for j in range(p + 1)])
- bias_constant = float((bandwidth_powers * (gram_inv @ leakage).ravel())[nu])
+ bias_constant = float((bandwidth_powers * (gram_inv @ leakage))[nu])
# Integrated variance: (2ν+1) h^{2ν+1} V̂_ν (CCF20, eq. 2.1)
variance_component = (2 * nu + 1) * h ** (2 * nu + 1) * V_nu
-
- # Rate exponent: 1/(2p+3) (CCF20, eq. 2.2)
- rate = 1 / (2 * p + 3)
-
return _VarianceFit(
influence_weights=influence_weights,
variance_component=variance_component,
bias_constant=bias_constant,
- rate=rate,
+ rate=1 / (2 * p + 3), # Rate exponent: 1/(2p+3) (CCF20, eq. 2.2)
)
-# ===================================================================
-# Bandwidth components (variance from cache, bias from fresh sample)
-# ===================================================================
-
-
def _compute_bandwidth_components(
cache: _PreliminaryCache,
cutoff: float,
@@ -390,7 +350,7 @@ def _compute_bandwidth_components(
"""Compute (V̂, B̂, R̂) for one side at given orders and bias bandwidth.
The variance side is looked up from the cache; only the bias window
- requires a fresh EstimationSample at ``bias_bandwidth``.
+ requires a fresh EstimationSample at `bias_bandwidth`.
Implements the plug-in estimates for the MSE formula in CCF20, eq. 2.1-2.2,
with regularisation from CCT14b, Supplemental Material Section S.7.2.
@@ -423,38 +383,38 @@ def _compute_bandwidth_components(
# Regularisation term (CCT14b, Section S.7.2)
regularisation = 0.0
if scaleregul > 0:
- hat_diagonal: np.ndarray | int = 0
- fitted_values: np.ndarray | int = 0
- if cache.vce in ("hc0", "hc1", "hc2", "hc3"):
+ if cache.vce == "nn":
+ bias_res = nn_residuals(
+ bias_sample.running,
+ bias_sample.responses,
+ bias_sample.duplicate_count,
+ bias_sample.duplicate_index,
+ cache.nnmatch,
+ )
+ else:
fitted_values = bias_fit.regressors @ bias_fit.beta
+ hat_diagonal: np.ndarray | None = None
if cache.vce in ("hc2", "hc3"):
hat_diagonal = np.sum(
(bias_fit.regressors @ bias_fit.gram_inv)
- * (bias_fit.regressors * bias_sample.main_weights.reshape(-1, 1)),
+ * (bias_fit.regressors * bias_sample.main_weights[:, None]),
axis=1,
)
+ bias_res = hc_residuals(
+ bias_sample.responses,
+ fitted_values,
+ hat_diagonal,
+ cache.vce,
+ bias_poly_order + 1,
+ )
- bias_residuals = rdrobust_res(
- bias_sample.running,
- bias_sample.outcome,
- bias_sample.fuzzy,
- bias_sample.covariates,
- fitted_values,
- hat_diagonal,
- cache.vce,
- cache.nnmatch,
- bias_sample.duplicate_count,
- bias_sample.duplicate_index,
- bias_poly_order + 1,
- )
weighted_bias_regressors = (
- bias_fit.regressors * bias_sample.main_weights.reshape(-1, 1)
+ bias_fit.regressors * bias_sample.main_weights[:, None]
)
- bias_sandwich_meat = rdrobust_vce(
- cache.n_fuzzy_cols + cache.n_covariate_cols,
+ bias_sandwich_meat = variance_covariance_matrix(
influence_weights,
weighted_bias_regressors,
- bias_residuals,
+ bias_res,
bias_sample.cluster,
)
bias_variance = (bias_fit.gram_inv @ bias_sandwich_meat @ bias_fit.gram_inv)[
@@ -463,11 +423,7 @@ def _compute_bandwidth_components(
regularisation = 3 * variance_fit.bias_constant**2 * bias_variance
# Bias estimate: B̂ = √(2(p+1-ν)) · B̂_const · s'β̂_{q+1}
- try:
- bias_beta = bias_fit.beta[-1, :]
- except IndexError:
- bias_beta = bias_fit.beta[-1]
-
+ bias_beta = bias_fit.beta[-1, :]
weighted_bias_coefficient = (
float(np.dot(influence_weights, bias_beta))
if np.ndim(influence_weights) > 0
@@ -522,11 +478,6 @@ def _compute_components_both_sides(
)
-# ===================================================================
-# Combining bandwidth components into a bandwidth
-# ===================================================================
-
-
def _combine_two(
comp: _BandwidthComponents,
scaleregul: float,
@@ -738,11 +689,6 @@ def _select_bandwidth_two(
return main_bw_l, main_bw_r, bias_bw_l, bias_bw_r
-# ===================================================================
-# Result container
-# ===================================================================
-
-
@dataclass(frozen=True, slots=True)
class BandwidthSelectionResult:
"""Selected bandwidths."""
@@ -758,11 +704,6 @@ class BandwidthSelectionResult:
vce: VarianceCovarianceEstimator
-# ===================================================================
-# Internal entry point (called by discontinuity.py with pre-split data)
-# ===================================================================
-
-
def _select_bandwidth(
data_l: Data,
data_r: Data,
@@ -978,7 +919,6 @@ def _select_bandwidth(
"msecomb2",
"cercomb1",
"cercomb2",
- "",
)
or compute_all
)
@@ -1076,7 +1016,6 @@ def _select_bandwidth(
selector_map: dict[str, tuple[float, float, float, float]] = {
"mserd": (h_mserd, h_mserd, b_mserd, b_mserd),
- "": (h_mserd, h_mserd, b_mserd, b_mserd),
"msetwo": (h_msetwo_l, h_msetwo_r, b_msetwo_l, b_msetwo_r),
"msesum": (h_msesum, h_msesum, b_msesum, b_msesum),
"msecomb1": (h_msecomb1, h_msecomb1, b_msecomb1, b_msecomb1),
@@ -1123,11 +1062,6 @@ def _select_bandwidth(
)
-# ===================================================================
-# Public entry point
-# ===================================================================
-
-
def select_bandwidth(
y: np.ndarray,
x: np.ndarray,
@@ -1168,11 +1102,11 @@ def select_bandwidth(
RD cutoff (default 0).
p, q, deriv : int
Polynomial orders and derivative order.
- kernel : ``"triangular"`` | ``"epanechnikov"`` | ``"uniform"``
+ kernel : `"triangular"` | `"epanechnikov"` | `"uniform"`
Kernel function.
bwselect : str
Bandwidth selection procedure.
- vce : ``"nn"`` | ``"hc0"`` | ``"hc1"`` | ``"hc2"`` | ``"hc3"``
+ vce : `"nn"` | `"hc0"` | `"hc1"` | `"hc2"` | `"hc3"`
Variance-covariance estimator.
nnmatch : int
Minimum neighbours for NN variance.
@@ -1180,7 +1114,7 @@ def select_bandwidth(
Regularisation scaling.
sharpbw : bool
Use sharp-RD bandwidth in fuzzy design.
- masspoints : ``"off"`` | ``"check"`` | ``"adjust"``
+ masspoints : `"off"` | `"check"` | `"adjust"`
Mass-point handling.
bwcheck : int or None
Minimum unique obs in bandwidth.
diff --git a/pyrdd/discontinuity.py b/pyrdd/discontinuity.py
index 52b8534..f8721a2 100644
--- a/pyrdd/discontinuity.py
+++ b/pyrdd/discontinuity.py
@@ -1,3 +1,20 @@
+"""
+Local polynomial RD estimation with robust bias-corrected inference.
+
+Implements the estimators and confidence intervals from:
+
+- CCT14: Calonico, Cattaneo, Titiunik (2014). Robust Nonparametric
+ Confidence Intervals for Regression-Discontinuity Designs.
+ Econometrica 82(6): 2295-2326.
+
+- CCF20: Calonico, Cattaneo, Farrell (2020). Optimal Bandwidth Choice
+ for Robust Bias-Corrected Inference in RD Designs.
+ Econometrics Journal 23(2): 192-210.
+
+- CCFT19: Calonico, Cattaneo, Farrell, Titiunik (2019). Regression
+ Discontinuity Designs Using Covariates. REStat 101(3): 442-451.
+"""
+
import math
import warnings
from dataclasses import dataclass
@@ -8,7 +25,12 @@
from pyrdd.bandwidth import _select_bandwidth
from pyrdd.internals.data import Data, EstimationSample
from pyrdd.internals.fit import PolynomialFitBiascorrection
-from pyrdd.internals.misc import check_bandwidth, rdrobust_res, rdrobust_vce
+from pyrdd.internals.misc import (
+ check_bandwidth,
+ hc_residuals,
+ nn_residuals,
+ variance_covariance_matrix,
+)
from pyrdd.internals.result import (
Bandwidths,
Estimate,
@@ -28,7 +50,7 @@
@dataclass(kw_only=True, frozen=True, slots=True)
class _CovariateAdjustment:
- """Results of partialling out covariates."""
+ """Frisch-Waugh-Lovell covariate projection results. See CCFT19, Section 2."""
influence_weights: np.ndarray
first_stage_projection: np.ndarray | None
@@ -47,25 +69,18 @@ class _TreatmentEffect:
bias_right: float
-@dataclass(kw_only=True, frozen=True, slots=True)
-class _VarianceEstimate:
- """Variance-covariance matrices and standard errors."""
-
- conventional_left: np.ndarray
- conventional_right: np.ndarray
- robust_left: np.ndarray
- robust_right: np.ndarray
- se_conventional: float
- se_robust: float
-
-
def _apply_fuzzy_correction(
reduced_form_conventional: float,
first_stage_conventional: float,
reduced_form_corrected: float,
first_stage_corrected: float,
) -> tuple[float, float, np.ndarray]:
- """Wald ratio and delta-method bias correction."""
+ """Wald ratio and delta-method bias correction for fuzzy RD.
+
+ The conventional estimate is the ratio of reduced-form and first-stage
+ discontinuities. The bias-corrected estimate subtracts the first-order
+ delta-method expansion of the bias. See CCT14, Section 3, eq. 3.1.
+ """
conventional = reduced_form_conventional / first_stage_conventional
influence_weights = np.array(
[
@@ -90,7 +105,11 @@ def _extract_reduced_form(
deriv_factorial: int,
adjustment: _CovariateAdjustment | None = None,
) -> float:
- """Extract the reduced-form discontinuity from a beta matrix."""
+ """Extract the reduced-form discontinuity from a coefficient matrix.
+
+ Returns `scalepar * deriv! * beta[deriv, 0]` without covariates,
+ or the FWL-adjusted version with covariates (CCFT19). See CCT14, Appendix A.1.
+ """
if adjustment is None:
return float(scalepar * deriv_factorial * beta[deriv, 0])
return float(
@@ -107,7 +126,7 @@ def _extract_first_stage(
deriv_factorial: int,
adjustment: _CovariateAdjustment | None = None,
) -> float:
- """Extract the first-stage discontinuity from a beta matrix."""
+ """Extract the first-stage discontinuity for fuzzy designs. See CCT14, Section 3."""
if adjustment is None or adjustment.first_stage_projection is None:
return float(deriv_factorial * beta[deriv, 1])
return float(
@@ -125,7 +144,11 @@ def _compute_side_bias(
fuzzy_influence: np.ndarray | None = None,
adjustment: _CovariateAdjustment | None = None,
) -> float:
- """Compute the bias for one side: conventional minus corrected."""
+ """Per-side bias: difference between conventional and bias-corrected estimates.
+
+ For fuzzy designs, the bias is projected through the delta-method influence weights.
+ See CCT14, Theorem A.1.
+ """
reduced_form_conventional = _extract_reduced_form(
fit.beta_conventional,
deriv,
@@ -142,6 +165,7 @@ def _compute_side_bias(
)
if fuzzy_influence is None:
return reduced_form_conventional - reduced_form_corrected
+
first_stage_conventional = _extract_first_stage(
fit.beta_conventional,
deriv,
@@ -170,9 +194,9 @@ def _compute_covariate_adjustment(
sample_r: EstimationSample,
n_fuzzy_cols: int,
n_covariate_cols: int,
- drop_collinear: int,
) -> _CovariateAdjustment:
- """Partial out covariates and compute influence weights."""
+ """Partial out covariates via Frisch-Waugh-Lovell and compute influence weights.
+ See CCFT19, Section 2."""
covariate_columns = np.arange(
1 + n_fuzzy_cols,
max(2 + n_fuzzy_cols + n_covariate_cols - 1, 2 + n_fuzzy_cols),
@@ -184,16 +208,13 @@ def _compute_covariate_adjustment(
sample_r,
n_fuzzy_cols,
covariate_columns,
- drop_collinear,
)
influence_weights = np.hstack([1.0, -projection[:, 0]])
- first_stage_projection: np.ndarray | None = None
- if n_fuzzy_cols > 0:
- first_stage_projection = np.append(1.0, -projection[:, 1])
-
return _CovariateAdjustment(
influence_weights=influence_weights,
- first_stage_projection=first_stage_projection,
+ first_stage_projection=np.append(1.0, -projection[:, 1])
+ if n_fuzzy_cols > 0
+ else None,
covariate_columns=covariate_columns,
projection=projection,
)
@@ -206,51 +227,48 @@ def _project_out_covariates(
sample_r: EstimationSample,
n_fuzzy_cols: int,
covariate_columns: np.ndarray,
- drop_collinear: int,
) -> np.ndarray:
- """Compute the covariate partialling-out coefficients."""
+ """Compute the FWL covariate projection coefficients, pooled across sides.
+
+ Solves `projection = (Z'M Z)^{-1} (Z'M Y)` where `M` partials
+ out the polynomial basis. See CCFT19, Section 2, Theorem 1.
+ """
main_weights_l = sample_l.main_weights.reshape(-1, 1)
main_weights_r = sample_r.main_weights.reshape(-1, 1)
responses_l = sample_l.responses
responses_r = sample_r.responses
- # Weighted cross-products of regressors with responses
- weighted_projection_l = (fit_l.regressor_p * main_weights_l).T @ responses_l
- weighted_projection_r = (fit_r.regressor_p * main_weights_r).T @ responses_r
+ # Weighted cross-products of regressors with all response columns
+ regressor_cross_response_l = (fit_l.regressor_p * main_weights_l).T @ responses_l
+ regressor_cross_response_r = (fit_r.regressor_p * main_weights_r).T @ responses_r
- # Weighted cross-products of covariates with responses
- covariates_cross_responses_l = (
- sample_l.covariates * main_weights_l
- ).T @ responses_l
- covariates_cross_responses_r = (
- sample_r.covariates * main_weights_r
- ).T @ responses_r
+ # Weighted cross-products of covariates with all response columns
+ covariate_cross_response_l = (sample_l.covariates * main_weights_l).T @ responses_l
+ covariate_cross_response_r = (sample_r.covariates * main_weights_r).T @ responses_r
- # Adjustment for the polynomial projection
- projection_adjustment_l = weighted_projection_l[:, covariate_columns].T @ (
- fit_l.gram_inv_p @ weighted_projection_l
+ # Remove the polynomial component from the covariate cross-products
+ projection_adjustment_l = regressor_cross_response_l[:, covariate_columns].T @ (
+ fit_l.gram_inv_p @ regressor_cross_response_l
)
- projection_adjustment_r = weighted_projection_r[:, covariate_columns].T @ (
- fit_r.gram_inv_p @ weighted_projection_r
+ projection_adjustment_r = regressor_cross_response_r[:, covariate_columns].T @ (
+ fit_r.gram_inv_p @ regressor_cross_response_r
)
- # Partial out the polynomial fit from the covariate cross-products
+ # Partialled-out covariate cross-products, pooled across sides
covariates_cross_covariates = (
- covariates_cross_responses_l[:, covariate_columns]
+ covariate_cross_response_l[:, covariate_columns]
- projection_adjustment_l[:, covariate_columns]
- + covariates_cross_responses_r[:, covariate_columns]
+ + covariate_cross_response_r[:, covariate_columns]
- projection_adjustment_r[:, covariate_columns]
)
covariates_cross_outcome = (
- covariates_cross_responses_l[:, : (1 + n_fuzzy_cols)]
+ covariate_cross_response_l[:, : (1 + n_fuzzy_cols)]
- projection_adjustment_l[:, : (1 + n_fuzzy_cols)]
- + covariates_cross_responses_r[:, : (1 + n_fuzzy_cols)]
+ + covariate_cross_response_r[:, : (1 + n_fuzzy_cols)]
- projection_adjustment_r[:, : (1 + n_fuzzy_cols)]
)
- if drop_collinear:
- return np.linalg.pinv(covariates_cross_covariates) @ covariates_cross_outcome
- return np.linalg.inv(covariates_cross_covariates) @ covariates_cross_outcome
+ return np.linalg.solve(covariates_cross_covariates, covariates_cross_outcome)
def _estimate_treatment_effect(
@@ -260,9 +278,15 @@ def _estimate_treatment_effect(
sample_r: EstimationSample,
deriv: int,
scalepar: float,
- drop_collinear: int,
) -> _TreatmentEffect:
- """Compute conventional and bias-corrected treatment effects."""
+ """Compute conventional and bias-corrected treatment effects.
+
+ For sharp RD, the treatment effect is the difference in polynomial
+ intercepts across the cutoff (CCT14, Section 2). For fuzzy RD,
+ it is the Wald ratio of reduced-form and first-stage discontinuities
+ (CCT14, Section 3). Covariates are partialled out via FWL before
+ computing either (CCFT19, Section 2).
+ """
deriv_factorial = math.factorial(deriv)
beta_conventional = fit_r.beta_conventional - fit_l.beta_conventional
beta_corrected = fit_r.beta_corrected - fit_l.beta_corrected
@@ -270,7 +294,6 @@ def _estimate_treatment_effect(
has_fuzzy = sample_l.fuzzy is not None
n_fuzzy_cols = 1 if has_fuzzy else 0
- # Covariate adjustment (if applicable)
adjustment: _CovariateAdjustment | None = None
if sample_l.covariates is not None:
adjustment = _compute_covariate_adjustment(
@@ -280,10 +303,8 @@ def _estimate_treatment_effect(
sample_r,
n_fuzzy_cols,
sample_l.covariates.shape[1],
- drop_collinear,
)
- # Reduced-form discontinuities
reduced_form_conventional = _extract_reduced_form(
beta_conventional,
deriv,
@@ -299,7 +320,6 @@ def _estimate_treatment_effect(
adjustment,
)
- # Treatment effect (with optional fuzzy correction)
influence_weights: np.ndarray | float
fuzzy_influence: np.ndarray | None = None
@@ -330,6 +350,7 @@ def _estimate_treatment_effect(
first_stage_corrected,
)
)
+ # Combine fuzzy delta-method weights with covariate projection
if adjustment is not None:
influence_weights = np.hstack(
[
@@ -342,7 +363,6 @@ def _estimate_treatment_effect(
else:
influence_weights = fuzzy_influence
- # Per-side biases
bias_left = _compute_side_bias(
fit_l,
deriv,
@@ -368,7 +388,7 @@ def _estimate_treatment_effect(
)
-def estimate(
+def fit(
y: np.ndarray,
x: np.ndarray,
c: float,
@@ -391,14 +411,92 @@ def estimate(
scalepar: float = 1.0,
scaleregul: float = 1.0,
sharpbw: bool = False,
- all: bool = True,
subset: np.ndarray | None = None,
masspoints: MassPointsCheck = "adjust",
bwcheck: int | None = None,
bwrestrict: bool = True,
stdvars: bool = False,
) -> Result:
- # Check for errors in the INPUT
+ """Local polynomial RD estimation with robust bias-corrected inference.
+
+ Implements the estimator and confidence intervals from Calonico,
+ Cattaneo, and Titiunik (2014, Econometrica).
+
+ Parameters
+ ----------
+ y : array
+ Outcome variable.
+ x : array
+ Running variable.
+ c : float
+ RD cutoff.
+ fuzzy : array or None
+ Treatment variable for fuzzy RD designs. When provided, the
+ treatment effect is estimated as the Wald ratio of the
+ reduced-form and first-stage discontinuities.
+ deriv : int
+ Derivative order of the treatment effect (0 = level, 1 = kink).
+ p : int
+ Polynomial order for point estimation (default 1 = local linear).
+ q : int
+ Polynomial order for bias correction (default 2 = local quadratic).
+ h : float or (float, float) or None
+ Main bandwidth. If a tuple, specifies (left, right) separately.
+ If None, selected automatically via `bwselect`.
+ b : float or (float, float) or None
+ Bias bandwidth. If None, set equal to `h`.
+ rho : float or None
+ If provided, sets `b = h / rho`, overriding `b`.
+ covs : array or None
+ Covariate matrix for adjustment via Frisch-Waugh-Lovell (CCFT19).
+ covs_drop : bool
+ Drop collinear covariates via pivoted QR.
+ kernel : `"triangular"` | `"epanechnikov"` | `"uniform"`
+ Kernel function for local polynomial weights.
+ weights : array or None
+ Observation weights.
+ bwselect : str
+ Bandwidth selector: `"mserd"`, `"msetwo"`, `"msesum"`,
+ `"cerrd"`, `"certwo"`, `"cersum"`, or combination variants.
+ vce : `"nn"` | `"hc0"` | `"hc1"` | `"hc2"` | `"hc3"`
+ Variance-covariance estimator. `"nn"` uses nearest-neighbour
+ matching (CCT14, eq. 5.1); `"hc0"`--`"hc3"` use
+ heteroskedasticity-consistent plug-in residuals (eq. 5.3).
+ cluster : array or None
+ Cluster variable for clustered standard errors.
+ nnmatch : int
+ Minimum number of neighbours for NN variance estimation.
+ level : float
+ Confidence level in percent (default 95).
+ scalepar : float
+ Scaling factor for the treatment effect.
+ scaleregul : float
+ Regularisation scaling for bandwidth selection.
+ sharpbw : bool
+ Use sharp-RD bandwidth formula in fuzzy designs.
+ subset : array or None
+ Boolean array selecting a subset of observations.
+ masspoints : `"off"` | `"check"` | `"adjust"`
+ Mass-point handling for the running variable.
+ bwcheck : int or None
+ Minimum number of unique observations within the bandwidth.
+ bwrestrict : bool
+ Restrict bandwidth to the range of the data.
+ stdvars : bool
+ Standardise variables before bandwidth selection.
+
+ Returns
+ -------
+ Result
+ Frozen dataclass containing:
+
+ - `conventional`, `bias_corrected`, `robust`: `Estimate`
+ objects with coefficient, standard error, t-statistic, p-value,
+ and confidence interval.
+ - `bandwidths`: selected (h_l, h_r) and (b_l, b_r).
+ - `sample_sizes`: total, effective, and unique counts per side.
+ - `settings`: estimation parameters.
+ """
if q < p:
raise ValueError(
"Polynomial order for bias correction must not be smaller than polynomial order"
@@ -409,10 +507,6 @@ def estimate(
raise ValueError("`level` should be set between 0 and 100")
if rho is not None and rho < 0:
raise ValueError("`rho` should be greater than 0")
- if h is not None:
- h_l, h_r = check_bandwidth(h)
- if b is not None:
- b_l, b_r = check_bandwidth(b)
data = Data.from_raw(
outcome=y,
@@ -449,15 +543,18 @@ def estimate(
M_l = N_l
M_r = N_r
- ######### Calculate bandwidth
- if h is None and N < 20:
+ # Bandwidth selection
+ if h is not None:
+ h_l, h_r = check_bandwidth(h)
+ b_l, b_r = check_bandwidth(b) if b is not None else (h_l, h_r)
+ elif N < 20:
warnings.warn(
"Not enough observations to perform bandwidth calculations. "
"Estimates computed using entire sample."
)
h_l = h_r = b_l = b_r = float(max(range_l, range_r))
- elif h is None:
- bw_result = _select_bandwidth(
+ else:
+ bandwidth = _select_bandwidth(
data_l=data_left,
data_r=data_right,
cutoff=c,
@@ -476,16 +573,14 @@ def estimate(
stdvars=stdvars,
drop_collinear=covs_drop,
)
- h_l, h_r = bw_result.h
- b_l, b_r = bw_result.b
- if rho is not None:
- b_l = h_l / rho
- b_r = h_r / rho
- elif b is None:
- b_l = h_l
- b_r = h_r
-
- # Get estimation samples to the left and right of cutoff
+ h_l, h_r = bandwidth.h
+ b_l, b_r = bandwidth.b
+
+ if rho is not None and N >= 20:
+ b_l = h_l / rho
+ b_r = h_r / rho
+
+ # Estimation samples and local polynomial fits
sample_left = EstimationSample.from_data(
data=data_left,
cutoff=c,
@@ -502,13 +597,23 @@ def estimate(
kernel=kernel,
vce=vce,
)
- # Estimate local polynomials
+
fit_left = PolynomialFitBiascorrection.fit(
- sample=sample_left, cutoff=c, p=p, q=q, bandwidth=h_l
+ sample=sample_left,
+ cutoff=c,
+ p=p,
+ q=q,
+ bandwidth=h_l,
)
fit_right = PolynomialFitBiascorrection.fit(
- sample=sample_right, cutoff=c, p=p, q=q, bandwidth=h_r
+ sample=sample_right,
+ cutoff=c,
+ p=p,
+ q=q,
+ bandwidth=h_r,
)
+
+ # Treatment effect
treatment_effect = _estimate_treatment_effect(
fit_l=fit_left,
fit_r=fit_right,
@@ -516,125 +621,145 @@ def estimate(
sample_r=sample_right,
deriv=deriv,
scalepar=scalepar,
- drop_collinear=covs_drop,
)
- # ####### Computing variance-covariance matrix.
- # Residuals
- hii_l = hii_r = fitted_p_l = fitted_p_r = fitted_q_l = fitted_q_r = 0
- if vce in ("hc0", "hc1", "hc2", "hc3"):
+
+ # Residuals for variance estimation
+ # See CCT14, Section 5 for both NN and HC approaches.
+ responses_l = sample_left.responses
+ responses_r = sample_right.responses
+
+ if vce == "nn":
+ # NN residuals are order-independent — same for conventional and robust
+ res_h_l = nn_residuals(
+ running=sample_left.running,
+ responses=responses_l,
+ duplicate_count=sample_left.duplicate_count,
+ duplicate_index=sample_left.duplicate_index,
+ matches=nnmatch,
+ )
+ res_h_r = nn_residuals(
+ running=sample_right.running,
+ responses=responses_r,
+ duplicate_count=sample_right.duplicate_count,
+ duplicate_index=sample_right.duplicate_index,
+ matches=nnmatch,
+ )
+ res_b_l = res_h_l
+ res_b_r = res_h_r
+ else:
+ # HC residuals: order p for conventional, order q for robust
fitted_p_l = fit_left.regressor_p @ fit_left.beta_conventional
fitted_p_r = fit_right.regressor_p @ fit_right.beta_conventional
fitted_q_l = fit_left.regressor_q @ fit_left.beta_bias
fitted_q_r = fit_right.regressor_q @ fit_right.beta_bias
+
+ hat_diagonal_l = hat_diagonal_r = None
if vce in ("hc2", "hc3"):
- hii_l = np.sum(
+ hat_diagonal_l = np.sum(
(fit_left.regressor_p @ fit_left.gram_inv_p)
- * (fit_left.regressor_p * sample_left.main_weights.reshape(-1, 1)),
+ * (fit_left.regressor_p * sample_left.main_weights[:, None]),
axis=1,
)
- hii_r = np.sum(
+ hat_diagonal_r = np.sum(
(fit_right.regressor_p @ fit_right.gram_inv_p)
- * (fit_right.regressor_p * sample_right.main_weights.reshape(-1, 1)),
+ * (fit_right.regressor_p * sample_right.main_weights[:, None]),
axis=1,
)
- res_h_l = rdrobust_res(
- sample_left.running,
- sample_left.outcome,
- sample_left.fuzzy,
- sample_left.covariates,
- fitted_p_l,
- hii_l,
- vce,
- nnmatch,
- sample_left.duplicate_count,
- sample_left.duplicate_index,
- p + 1,
- )
- res_h_r = rdrobust_res(
- sample_right.running,
- sample_right.outcome,
- sample_right.fuzzy,
- sample_right.covariates,
- fitted_p_r,
- hii_r,
- vce,
- nnmatch,
- sample_right.duplicate_count,
- sample_right.duplicate_index,
- p + 1,
- )
-
- if vce == "nn":
- res_b_l = res_h_l
- res_b_r = res_h_r
- else:
- res_b_l = rdrobust_res(
- sample_left.running,
- sample_left.outcome,
- sample_left.fuzzy,
- sample_left.covariates,
- fitted_q_l,
- hii_l,
- vce,
- nnmatch,
- sample_left.duplicate_count,
- sample_left.duplicate_index,
- q + 1,
+ res_h_l = hc_residuals(
+ responses=responses_l,
+ fitted_values=fitted_p_l,
+ hat_diagonal=hat_diagonal_l,
+ vce=vce,
+ n_params=p + 1,
+ )
+ res_h_r = hc_residuals(
+ responses=responses_r,
+ fitted_values=fitted_p_r,
+ hat_diagonal=hat_diagonal_r,
+ vce=vce,
+ n_params=p + 1,
)
- res_b_r = rdrobust_res(
- sample_right.running,
- sample_right.outcome,
- sample_right.fuzzy,
- sample_right.covariates,
- fitted_q_r,
- hii_r,
- vce,
- nnmatch,
- sample_right.duplicate_count,
- sample_right.duplicate_index,
- q + 1,
+ res_b_l = hc_residuals(
+ responses=responses_l,
+ fitted_values=fitted_q_l,
+ hat_diagonal=hat_diagonal_l,
+ vce=vce,
+ n_params=q + 1,
+ )
+ res_b_r = hc_residuals(
+ responses=responses_r,
+ fitted_values=fitted_q_r,
+ hat_diagonal=hat_diagonal_r,
+ vce=vce,
+ n_params=q + 1,
)
# Sandwich variance matrices
- n_extra = (1 if sample_left.fuzzy is not None else 0) + (
- sample_left.covariates.shape[1] if sample_left.covariates is not None else 0
- )
- weights_l = fit_left.regressor_p * sample_left.main_weights.reshape(-1, 1)
- weights_r = fit_right.regressor_p * sample_right.main_weights.reshape(-1, 1)
- s_Y = treatment_effect.influence_weights
-
- V_cl_l = (
+ # Conventional: order-p weights and residuals.
+ # Robust: bias-correction projection and order-q residuals,
+ # accounting for the variability of bias estimation (CCT14, Theorem A.2).
+ variance_conventional_l = (
fit_left.gram_inv_p
- @ rdrobust_vce(n_extra, s_Y, weights_l, res_h_l, sample_left.cluster)
+ @ variance_covariance_matrix(
+ influence_weights=treatment_effect.influence_weights,
+ weighted_design_matrix=fit_left.regressor_p
+ * sample_left.main_weights[:, None],
+ residuals=res_h_l,
+ clusters=sample_left.cluster,
+ )
@ fit_left.gram_inv_p
)
- V_cl_r = (
+ variance_conventional_r = (
fit_right.gram_inv_p
- @ rdrobust_vce(n_extra, s_Y, weights_r, res_h_r, sample_right.cluster)
+ @ variance_covariance_matrix(
+ influence_weights=treatment_effect.influence_weights,
+ weighted_design_matrix=fit_right.regressor_p
+ * sample_right.main_weights[:, None],
+ residuals=res_h_r,
+ clusters=sample_right.cluster,
+ )
@ fit_right.gram_inv_p
)
- V_rb_l = (
+
+ variance_robust_l = (
fit_left.gram_inv_p
- @ rdrobust_vce(
- n_extra, s_Y, fit_left.bias_correction, res_b_l, sample_left.cluster
+ @ variance_covariance_matrix(
+ influence_weights=treatment_effect.influence_weights,
+ weighted_design_matrix=fit_left.bias_correction,
+ residuals=res_b_l,
+ clusters=sample_left.cluster,
)
@ fit_left.gram_inv_p
)
- V_rb_r = (
+ variance_robust_r = (
fit_right.gram_inv_p
- @ rdrobust_vce(
- n_extra, s_Y, fit_right.bias_correction, res_b_r, sample_right.cluster
+ @ variance_covariance_matrix(
+ influence_weights=treatment_effect.influence_weights,
+ weighted_design_matrix=fit_right.bias_correction,
+ residuals=res_b_r,
+ clusters=sample_right.cluster,
)
@ fit_right.gram_inv_p
)
+ # Standard errors and confidence intervals
deriv_factorial = math.factorial(deriv)
scale = scalepar**2 * deriv_factorial**2
- se_conventional = float(np.sqrt(scale * (V_cl_l + V_cl_r)[deriv, deriv]))
- se_robust = float(np.sqrt(scale * (V_rb_l + V_rb_r)[deriv, deriv]))
+ se_conventional = float(
+ np.sqrt(
+ scale * (variance_conventional_l + variance_conventional_r)[deriv, deriv]
+ )
+ )
+ se_robust = float(
+ np.sqrt(scale * (variance_robust_l + variance_robust_r)[deriv, deriv])
+ )
critical_value = float(-scipy.stats.norm.ppf(abs((1 - level / 100) / 2)))
+ # Collect three inference methods (CCT14, Remarks 6-7):
+ # conventional: point estimate from order p, conventional SE
+ # bias_corrected: bias-corrected estimate, conventional SE
+ # robust: bias-corrected estimate, robust SE (recommended)
return Result(
_estimates=(
Estimate.from_estimate(
@@ -674,6 +799,6 @@ def estimate(
),
bias=(treatment_effect.bias_left, treatment_effect.bias_right),
beta=(fit_left.beta_conventional, fit_right.beta_conventional),
- V_conventional=(V_cl_l, V_cl_r),
- V_robust=(V_rb_l, V_rb_r),
+ V_conventional=(variance_conventional_l, variance_conventional_r),
+ V_robust=(variance_robust_l, variance_robust_r),
)
diff --git a/pyrdd/internals/misc.py b/pyrdd/internals/misc.py
index bca0bca..bcd376b 100644
--- a/pyrdd/internals/misc.py
+++ b/pyrdd/internals/misc.py
@@ -4,7 +4,7 @@
import scipy.linalg
from numba import njit
-from pyrdd.internals.types import Bandwidth, Kernel
+from pyrdd.internals.types import Bandwidth, Kernel, VarianceCovarianceEstimator
def drop_collinear_variables(z: np.ndarray, tol: float = 1e-5) -> np.ndarray:
@@ -119,144 +119,113 @@ def _nn_bounds(
return lo, hi
-def rdrobust_res(
- X: np.ndarray,
- y: np.ndarray,
- T: np.ndarray | None,
- Z: np.ndarray | None,
- m: np.ndarray | int,
- hii: np.ndarray | int,
- vce: str,
+def _stack_responses(
+ outcome: np.ndarray,
+ fuzzy: np.ndarray | None,
+ covariates: np.ndarray | None,
+) -> np.ndarray:
+ """Stack outcome, fuzzy, and covariates into a response matrix."""
+ columns: list[np.ndarray] = [outcome]
+ if fuzzy is not None:
+ columns.append(fuzzy)
+ if covariates is not None:
+ columns.append(covariates)
+ return np.column_stack(columns)
+
+
+def nn_residuals(
+ running: np.ndarray,
+ responses: np.ndarray,
+ duplicate_count: np.ndarray,
+ duplicate_index: np.ndarray,
matches: int,
- dups: np.ndarray,
- dupsid: np.ndarray,
- d: int,
) -> np.ndarray:
- """Compute residuals for rdrobust variance estimation.
+ """Nearest-neighbour matching residuals (CCT14a, Section 5, eq. 5.1)."""
+ n = running.shape[0]
+ lo, hi = _nn_bounds(running, duplicate_count, duplicate_index, matches, n)
- Assumes all arrays have consistent shapes:
- - X, y, T, dups, dupsid, hii are 1-D
- - Z is 2-D (n, dZ) or None
- - m (fitted values) is 2-D (n, cols) or 0
+ cumulative = np.cumsum(responses, axis=0)
+ window_sum = cumulative[hi]
+ mask = lo > 0
+ window_sum[mask] -= cumulative[lo[mask] - 1]
- Parameters
- ----------
- X : (n,) — Running variable (sorted).
- y : (n,) — Outcome variable.
- T : (n,) or None — Treatment variable (fuzzy designs).
- Z : (n, dZ) or None — Covariate matrix.
- m : (n, cols) or 0 — Predicted values (HC path only).
- hii : (n,) or 0 — Hat matrix diagonal (HC2/HC3 only).
- vce : str — Variance estimator type.
- matches : int — Minimum number of NN matches.
- dups, dupsid : (n,) int arrays — Duplicate structure (NN path only).
- d : int — Number of estimated parameters (HC1 correction).
+ neighbour_sum = window_sum - responses
+ window_size = (hi - lo).astype(np.float64)
+ scale = np.sqrt(window_size / (window_size + 1.0))
- Returns
- -------
- res : (n, 1+dT+dZ) — Residual matrix.
- """
- n = len(y)
- dT = int(T is not None)
- dZ = Z.shape[1] if Z is not None else 0
-
- if vce != "nn":
- assert isinstance(m, np.ndarray), "HC path requires fitted values"
- res = np.empty((n, 1 + dT + dZ))
- if vce == "hc0":
- w: np.ndarray | float = 1.0
- elif vce == "hc1":
- w = np.sqrt(n / (n - d))
- elif vce == "hc2":
- assert isinstance(hii, np.ndarray), "HC2/HC3 requires hat diagonal"
- w = np.sqrt(1.0 / (1.0 - hii))
- else: # hc3
- assert isinstance(hii, np.ndarray), "HC2/HC3 requires hat diagonal"
- w = 1.0 / (1.0 - hii)
- res[:, 0] = w * (y - m[:, 0]) # ty: ignore[invalid-argument-type]
- if T is not None:
- res[:, 1] = w * (T - m[:, 1]) # ty: ignore[invalid-argument-type]
- if Z is not None:
- for i in range(dZ):
- res[:, 1 + dT + i] = w * (Z[:, i] - m[:, 1 + dT + i]) # ty: ignore[invalid-argument-type]
- return res
-
- # ── NN path ─────────────────────────────────────────────────────
-
- # Step 1: compute window boundaries via numba-jitted loop
- lo, hi = _nn_bounds(X, dups, dupsid, matches, n)
-
- # Step 2: stack all response columns and compute window sums via cumsum
- cols: list[np.ndarray] = [y]
- if T is not None:
- cols.append(T)
- if Z is not None:
- cols.extend(Z[:, i] for i in range(dZ))
- D = np.column_stack(cols) # (n, 1+dT+dZ)
-
- cumD = np.cumsum(D, axis=0)
- window_sum = cumD[hi]
- mask = lo > 0
- window_sum[mask] -= cumD[lo[mask] - 1]
+ return scale[:, None] * (responses - neighbour_sum / window_size[:, None])
- neighbour_sum = window_sum - D
- Ji = (hi - lo).astype(np.float64) # window_size - 1
- scale = np.sqrt(Ji / (Ji + 1.0))
- return scale[:, None] * (D - neighbour_sum / Ji[:, None])
+def hc_residuals(
+ responses: np.ndarray,
+ fitted_values: np.ndarray,
+ hat_diagonal: np.ndarray | None,
+ vce: VarianceCovarianceEstimator,
+ n_params: int,
+) -> np.ndarray:
+ """Heteroskedasticity-consistent residuals (CCT14, Section 5, eq. 5.3)."""
+ raw = responses - fitted_values
+ if vce == "hc0":
+ return raw
+ elif vce == "hc1":
+ n = responses.shape[0]
+ return np.sqrt(n / (n - n_params)) * raw
+ elif vce == "hc2":
+ if hat_diagonal is None:
+ raise ValueError("HC2 requires hat_diagonal")
+ return raw / np.sqrt(1.0 - hat_diagonal)[:, None]
+ elif vce == "hc3":
+ if hat_diagonal is None:
+ raise ValueError("HC3 requires hat_diagonal")
+ return raw / (1.0 - hat_diagonal)[:, None]
+ else:
+ raise ValueError(f"Unknown VCE type: {vce}")
-def rdrobust_vce(
- d: int,
- s: np.ndarray | float,
- RX: np.ndarray,
- res: np.ndarray,
- C: np.ndarray | None,
+def variance_covariance_matrix(
+ influence_weights: np.ndarray | float,
+ weighted_design_matrix: np.ndarray,
+ residuals: np.ndarray,
+ clusters: np.ndarray | None,
) -> np.ndarray:
"""Sandwich-form variance-covariance matrix estimator.
+ Computes the meat of the sandwich: A'A (unclustered) or the
+ clustered equivalent with small-sample correction.
+ See CCT14a, Section 5.
+
Parameters
----------
- d : int — Number of extra response columns (fuzzy + covariates).
- s : array or float — Influence weights.
- RX : (n, k) — Weighted design matrix.
- res : (n, 1+d) — Residual matrix.
- C : (n,) or None — Cluster IDs.
+ influence_weights : np.ndarray or float
+ Delta-method weights for projecting multivariate residuals.
+ weighted_design_matrix : np.ndarray of shape (n, k)
+ Weighted design matrix.
+ residuals : np.ndarray of shape (n, n_responses)
+ Residual matrix.
+ clusters : np.ndarray of shape (n,) or None
+ Cluster IDs.
Returns
-------
- M : (k, k) — Meat of the sandwich estimator.
+ np.ndarray of shape (k, k)
+ Meat of the sandwich estimator.
"""
- k = RX.shape[1]
-
- if C is None:
- if d == 0:
- A = RX * res
- return A.T @ A
- s_flat = np.asarray(s).ravel()
- e = res @ s_flat # (n,)
- A = RX * e[:, None] # (n, k)
- return A.T @ A
-
- # ── Cluster branch ──────────────────────────────────────────────
- try:
- n = len(C)
- except TypeError:
- n = 1
-
- if d == 0:
- A = RX * res
- else:
- s_flat = np.asarray(s).ravel()
- e = res @ s_flat
- A = RX * e[:, None]
+ n_regressors = weighted_design_matrix.shape[1]
+ projected_residuals = residuals @ np.atleast_1d(influence_weights)
+ score = weighted_design_matrix * projected_residuals[:, None]
- _, inverse = np.unique(C, return_inverse=True)
- g = int(inverse.max()) + 1
- w = ((n - 1) / (n - k)) * (g / (g - 1))
+ if clusters is None:
+ return score.T @ score
- cluster_sums = np.column_stack(
- [np.bincount(inverse, weights=A[:, j], minlength=g) for j in range(k)]
- ) # (g, k)
+ n = clusters.shape[0]
+ _, inverse = np.unique(clusters, return_inverse=True)
+ n_clusters = int(inverse.max()) + 1
+ correction = ((n - 1) / (n - n_regressors)) * (n_clusters / (n_clusters - 1))
- return w * (cluster_sums.T @ cluster_sums)
+ cluster_sums = np.column_stack(
+ [
+ np.bincount(inverse, weights=score[:, j], minlength=n_clusters)
+ for j in range(n_regressors)
+ ]
+ )
+ return correction * (cluster_sums.T @ cluster_sums)
diff --git a/tests/test_discontinuity_estimation.py b/tests/test_discontinuity_estimation.py
index 1581ca6..ce987bd 100644
--- a/tests/test_discontinuity_estimation.py
+++ b/tests/test_discontinuity_estimation.py
@@ -6,7 +6,7 @@
import pytest
import rdrobust
-from pyrdd.discontinuity import estimate
+from pyrdd.discontinuity import fit
from pyrdd.internals.result import Result
from tests.conftest import TOLERANCE, requires_numpy_lt2
@@ -30,7 +30,7 @@ class TestSharpRD:
def test_full_output(self, sharp_rd_data: tuple[np.ndarray, np.ndarray]) -> None:
x, y = sharp_rd_data
orig = rdrobust.rdrobust(y=y, x=x, c=0, h=0.5, kernel="triangular", vce="nn")
- new = estimate(y=y, x=x, c=0, h=0.5, kernel="triangular", vce="nn")
+ new = fit(y=y, x=x, c=0, h=0.5, kernel="triangular", vce="nn")
_compare(orig, new)
assert new.conventional.t_statistic == pytest.approx(
expected=orig.t.iloc[0].item(), abs=TOLERANCE
@@ -60,14 +60,14 @@ def test_full_output(self, sharp_rd_data: tuple[np.ndarray, np.ndarray]) -> None
def test_sample_sizes(self, sharp_rd_data: tuple[np.ndarray, np.ndarray]) -> None:
x, y = sharp_rd_data
orig = rdrobust.rdrobust(y=y, x=x, c=0, h=0.5, kernel="triangular", vce="nn")
- new = estimate(y=y, x=x, c=0, h=0.5, kernel="triangular", vce="nn")
+ new = fit(y=y, x=x, c=0, h=0.5, kernel="triangular", vce="nn")
assert new.sample_sizes.n == tuple(orig.N)
assert new.sample_sizes.n_effective == tuple(int(v) for v in orig.N_h)
def test_bandwidths(self, sharp_rd_data: tuple[np.ndarray, np.ndarray]) -> None:
x, y = sharp_rd_data
orig = rdrobust.rdrobust(y=y, x=x, c=0, h=0.5, kernel="triangular", vce="nn")
- new = estimate(y=y, x=x, c=0, h=0.5, kernel="triangular", vce="nn")
+ new = fit(y=y, x=x, c=0, h=0.5, kernel="triangular", vce="nn")
assert new.bandwidths.h == pytest.approx(
expected=(orig.bws.iloc[0, 0], orig.bws.iloc[0, 1]), abs=TOLERANCE
)
@@ -80,7 +80,7 @@ class TestHC1Variance:
def test_hc1(self, sharp_rd_data: tuple[np.ndarray, np.ndarray]) -> None:
x, y = sharp_rd_data
orig = rdrobust.rdrobust(y=y, x=x, c=0, h=0.5, kernel="triangular", vce="hc1")
- new = estimate(y=y, x=x, c=0, h=0.5, kernel="triangular", vce="hc1")
+ new = fit(y=y, x=x, c=0, h=0.5, kernel="triangular", vce="hc1")
_compare(orig, new)
@@ -90,7 +90,7 @@ def test_separate_bw(self, sharp_rd_data: tuple[np.ndarray, np.ndarray]) -> None
orig = rdrobust.rdrobust(
y=y, x=x, c=0, h=0.5, b=0.7, kernel="triangular", vce="nn"
)
- new = estimate(y=y, x=x, c=0, h=0.5, b=0.7, kernel="triangular", vce="nn")
+ new = fit(y=y, x=x, c=0, h=0.5, b=0.7, kernel="triangular", vce="nn")
_compare(orig, new)
@@ -112,7 +112,7 @@ def test_covariates_fixed_bw(
kernel="triangular",
vce="nn",
)
- new = estimate(
+ new = fit(
y=y,
x=x,
c=0,
@@ -141,7 +141,7 @@ def test_single_covariate(
kernel="triangular",
vce="nn",
)
- new = estimate(
+ new = fit(
y=y,
x=x,
c=0,
@@ -169,7 +169,7 @@ def test_covariates_hc1(
kernel="triangular",
vce="hc1",
)
- new = estimate(
+ new = fit(
y=y,
x=x,
c=0,
diff --git a/tests/test_senate.py b/tests/test_senate.py
index 63fe916..4d3ab5b 100644
--- a/tests/test_senate.py
+++ b/tests/test_senate.py
@@ -7,7 +7,7 @@
import rdrobust
from pyrdd.bandwidth import BandwidthSelectionResult, select_bandwidth
-from pyrdd.discontinuity import estimate
+from pyrdd.discontinuity import fit
from pyrdd.internals.result import Result
from pyrdd.internals.types import BandwidthSelector
from tests.conftest import ALL_SELECTORS, TOLERANCE, requires_numpy_lt2
@@ -40,19 +40,19 @@ def _compare_bw(
class TestSenateDefault:
def test_default(self, senate_vote: np.ndarray, senate_margin: np.ndarray) -> None:
orig = rdrobust.rdrobust(y=senate_vote, x=senate_margin)
- new = estimate(y=senate_vote, x=senate_margin, c=0.0)
+ new = fit(y=senate_vote, x=senate_margin, c=0.0)
_compare(orig, new)
def test_p_value(self, senate_vote: np.ndarray, senate_margin: np.ndarray) -> None:
orig = rdrobust.rdrobust(y=senate_vote, x=senate_margin)
- new = estimate(y=senate_vote, x=senate_margin, c=0.0)
+ new = fit(y=senate_vote, x=senate_margin, c=0.0)
assert abs(orig.pv.iloc[2].item() - new.robust.p_value) < TOLERANCE
def test_confidence_interval(
self, senate_vote: np.ndarray, senate_margin: np.ndarray
) -> None:
orig = rdrobust.rdrobust(y=senate_vote, x=senate_margin)
- new = estimate(y=senate_vote, x=senate_margin, c=0.0)
+ new = fit(y=senate_vote, x=senate_margin, c=0.0)
assert abs(orig.ci.iloc[2, 0].item() - new.robust.ci[0]) < TOLERANCE
assert abs(orig.ci.iloc[2, 1].item() - new.robust.ci[1]) < TOLERANCE
@@ -60,14 +60,14 @@ def test_sample_sizes(
self, senate_vote: np.ndarray, senate_margin: np.ndarray
) -> None:
orig = rdrobust.rdrobust(y=senate_vote, x=senate_margin)
- new = estimate(y=senate_vote, x=senate_margin, c=0.0)
+ new = fit(y=senate_vote, x=senate_margin, c=0.0)
assert new.sample_sizes.n == tuple(orig.N)
def test_bandwidths(
self, senate_vote: np.ndarray, senate_margin: np.ndarray
) -> None:
orig = rdrobust.rdrobust(y=senate_vote, x=senate_margin)
- new = estimate(y=senate_vote, x=senate_margin, c=0.0)
+ new = fit(y=senate_vote, x=senate_margin, c=0.0)
assert abs(orig.bws.iloc[0, 0] - new.bandwidths.h[0]) < TOLERANCE
assert abs(orig.bws.iloc[0, 1] - new.bandwidths.h[1]) < TOLERANCE
assert abs(orig.bws.iloc[1, 0] - new.bandwidths.b[0]) < TOLERANCE
@@ -79,7 +79,7 @@ def test_manual_bandwidth(
self, senate_vote: np.ndarray, senate_margin: np.ndarray
) -> None:
orig = rdrobust.rdrobust(y=senate_vote, x=senate_margin, h=16.79369, b=27.43745)
- new = estimate(y=senate_vote, x=senate_margin, c=0.0, h=16.79369, b=27.43745)
+ new = fit(y=senate_vote, x=senate_margin, c=0.0, h=16.79369, b=27.43745)
_compare(orig, new)
@@ -88,7 +88,7 @@ def test_asymmetric(
self, senate_vote: np.ndarray, senate_margin: np.ndarray
) -> None:
orig = rdrobust.rdrobust(y=senate_vote, x=senate_margin, h=(12, 15), b=(18, 20))
- new = estimate(y=senate_vote, x=senate_margin, c=0.0, h=(12, 15), b=(18, 20))
+ new = fit(y=senate_vote, x=senate_margin, c=0.0, h=(12, 15), b=(18, 20))
_compare(orig, new)
@@ -102,9 +102,7 @@ def test_cluster_nn(
orig = rdrobust.rdrobust(
y=senate_vote, x=senate_margin, vce="nn", cluster=senate_state
)
- new = estimate(
- y=senate_vote, x=senate_margin, c=0.0, vce="nn", cluster=senate_state
- )
+ new = fit(y=senate_vote, x=senate_margin, c=0.0, vce="nn", cluster=senate_state)
_compare(orig, new)
@@ -122,7 +120,7 @@ def test_uniform_hc1_cluster(
vce="hc1",
cluster=senate_state,
)
- new = estimate(
+ new = fit(
y=senate_vote,
x=senate_margin,
c=0.0,
@@ -138,9 +136,7 @@ def test_certwo_hc3(
orig = rdrobust.rdrobust(
y=senate_vote, x=senate_margin, bwselect="certwo", vce="hc3"
)
- new = estimate(
- y=senate_vote, x=senate_margin, c=0.0, bwselect="certwo", vce="hc3"
- )
+ new = fit(y=senate_vote, x=senate_margin, c=0.0, bwselect="certwo", vce="hc3")
_compare(orig, new)
@@ -157,7 +153,7 @@ def test_balance(
) -> None:
z = senate_data[covariate_name].astype(float)
orig = rdrobust.rdrobust(y=z, x=senate_margin)
- new = estimate(y=z, x=senate_margin, c=0.0)
+ new = fit(y=z, x=senate_margin, c=0.0)
_compare(orig, new)
@@ -216,7 +212,7 @@ def test_covariates_fixed_bw(
orig = rdrobust.rdrobust(
y=senate_vote, x=senate_margin, covs=senate_covariates, h=17.75, b=28.03
)
- new = estimate(
+ new = fit(
y=senate_vote,
x=senate_margin,
c=0.0,
@@ -239,7 +235,7 @@ def test_covariates_asymmetric_bw(
h=(12, 18),
b=(20, 25),
)
- new = estimate(
+ new = fit(
y=senate_vote,
x=senate_margin,
c=0.0,
@@ -259,9 +255,7 @@ def test_single_covariate(
orig = rdrobust.rdrobust(
y=senate_vote, x=senate_margin, covs=covs, h=17.75, b=28.03
)
- new = estimate(
- y=senate_vote, x=senate_margin, c=0.0, covs=covs, h=17.75, b=28.03
- )
+ new = fit(y=senate_vote, x=senate_margin, c=0.0, covs=covs, h=17.75, b=28.03)
_compare(orig, new)
def test_covariates_with_clustering(
@@ -279,7 +273,7 @@ def test_covariates_with_clustering(
h=17.75,
b=28.03,
)
- new = estimate(
+ new = fit(
y=senate_vote,
x=senate_margin,
c=0.0,