-
Notifications
You must be signed in to change notification settings - Fork 39
Expand file tree
/
Copy pathlocal_linear.py
More file actions
1332 lines (1198 loc) · 53.3 KB
/
Copy pathlocal_linear.py
File metadata and controls
1332 lines (1198 loc) · 53.3 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
"""
Kernels and univariate local-linear regression at a boundary.
Ships the foundational RDD infrastructure that downstream estimators compose:
- Bounded one-sided kernels (Epanechnikov, triangular, uniform) on ``[0, 1]``
suitable for boundary-point nonparametric regression.
- Closed-form kernel-moment constants ``kappa_k := int_0^1 t^k * k(t) dt`` and the
derived boundary-kernel constant ``C = (kappa_2^2 - kappa_1 kappa_3) /
(kappa_0 kappa_2 - kappa_1^2)`` that appears in the asymptotic bias of
local-linear at a boundary.
- A univariate local-linear regression fitter ``local_linear_fit`` that estimates
the conditional mean ``m(d0) := E[Y | D = d0]`` at the boundary of ``D``'s
support via kernel-weighted OLS.
This module is used by the :class:`HeterogeneousAdoptionDiD` phases:
- Phase 1a ships the kernels and fitter (this module).
- Phase 1b ships the MSE-optimal bandwidth selector
``mse_optimal_bandwidth`` / ``BandwidthResult`` (this module), a thin
wrapper over the Calonico-Cattaneo-Farrell (2018) plug-in selector
ported in ``diff_diff/_nprobust_port.py``.
- Phase 1c ships the bias-corrected local-linear fit
``bias_corrected_local_linear`` / ``BiasCorrectedFit`` (this module), a
thin wrapper over the Calonico-Cattaneo-Titiunik (2014) robust-bias
correction ported in ``diff_diff/_nprobust_port.py``. This produces
the mu-scale point estimate and CI for Equation 8 of de Chaisemartin,
Ciccia, D'Haultfoeuille & Knau (2026, arXiv:2405.04465v6); Phase 2
applies the ``(1/G) * sum(D_{g,2})`` beta-scale rescaling.
References
----------
- de Chaisemartin, C., Ciccia, D., D'Haultfoeuille, X., & Knau, F. (2026).
Difference-in-Differences Estimators When No Unit Remains Untreated.
arXiv:2405.04465v6. Section 3.1.3 defines the kernel-moment constants used
here.
- Calonico, S., Cattaneo, M. D., & Farrell, M. H. (2018). On the effect of bias
estimation on coverage accuracy in nonparametric inference. Journal of the
American Statistical Association, 113(522), 767-779.
- Fan, J., & Gijbels, I. (1996). Local Polynomial Modelling and Its
Applications. Chapman & Hall.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Callable, Dict, Optional
import numpy as np
from scipy import integrate
from diff_diff.linalg import solve_ols
__all__ = [
"BandwidthResult",
"BiasCorrectedFit",
"bias_corrected_local_linear",
"epanechnikov_kernel",
"triangular_kernel",
"uniform_kernel",
"KERNELS",
"kernel_moments",
"LocalLinearFit",
"local_linear_fit",
"mse_optimal_bandwidth",
]
# =============================================================================
# Kernel functions
# =============================================================================
#
# Each kernel is defined on [0, 1] (one-sided, for boundary estimation where
# the running variable D is supported on [0, infinity) and the evaluation point
# is at d0 = 0). Kernels return 0 outside [0, 1].
def epanechnikov_kernel(u: np.ndarray) -> np.ndarray:
"""Epanechnikov kernel on ``[0, 1]``.
``k(u) = (3/4)(1 - u^2)`` for ``u in [0, 1]``, zero elsewhere.
Parameters
----------
u : np.ndarray
Points on the scaled domain ``u = (d - d0) / h``.
Returns
-------
np.ndarray
Kernel values, same shape as ``u``.
"""
u = np.asarray(u, dtype=np.float64)
inside = (u >= 0.0) & (u <= 1.0)
return np.where(inside, 0.75 * (1.0 - u * u), 0.0)
def triangular_kernel(u: np.ndarray) -> np.ndarray:
"""Triangular kernel on ``[0, 1]``.
``k(u) = 1 - u`` for ``u in [0, 1]``, zero elsewhere.
Using the convention ``int_0^1 k(u) du = 1/2`` to match Epanechnikov's
one-sided normalization.
"""
u = np.asarray(u, dtype=np.float64)
inside = (u >= 0.0) & (u <= 1.0)
return np.where(inside, 1.0 - u, 0.0)
def uniform_kernel(u: np.ndarray) -> np.ndarray:
"""Uniform (rectangular) kernel on ``[0, 1]``.
``k(u) = 1`` for ``u in [0, 1]``, zero elsewhere.
Already normalized to ``int_0^1 k(u) du = 1`` (no factor of 1/2 like the
other two).
"""
u = np.asarray(u, dtype=np.float64)
inside = (u >= 0.0) & (u <= 1.0)
return np.where(inside, 1.0, 0.0)
KERNELS: Dict[str, Callable[[np.ndarray], np.ndarray]] = {
"epanechnikov": epanechnikov_kernel,
"triangular": triangular_kernel,
"uniform": uniform_kernel,
}
# =============================================================================
# Closed-form kernel moments
# =============================================================================
#
# For each kernel k(t) on [0, 1], the moments
#
# kappa_k := int_0^1 t^k * k(t) dt
#
# admit closed forms. These are the values from elementary integration; the
# test suite verifies each against a numerical scipy.integrate.quad call.
#
# The derived constant C from the paper's Section 3.1.3 is
#
# C = (kappa_2^2 - kappa_1 * kappa_3) / (kappa_0 * kappa_2 - kappa_1^2)
#
# and can be negative depending on kernel shape (e.g. Epanechnikov C < 0).
_CLOSED_FORM_MOMENTS: Dict[str, Dict[str, float]] = {
"epanechnikov": {
"kappa_0": 1.0 / 2.0,
"kappa_1": 3.0 / 16.0,
"kappa_2": 1.0 / 10.0,
"kappa_3": 1.0 / 16.0,
"kappa_4": 3.0 / 70.0,
},
"triangular": {
"kappa_0": 1.0 / 2.0,
"kappa_1": 1.0 / 6.0,
"kappa_2": 1.0 / 12.0,
"kappa_3": 1.0 / 20.0,
"kappa_4": 1.0 / 30.0,
},
"uniform": {
"kappa_0": 1.0,
"kappa_1": 1.0 / 2.0,
"kappa_2": 1.0 / 3.0,
"kappa_3": 1.0 / 4.0,
"kappa_4": 1.0 / 5.0,
},
}
def kernel_moments(kernel: str = "epanechnikov") -> Dict[str, float]:
"""Return kernel-moment constants used in boundary local-linear asymptotics.
The returned dict contains five raw moments ``kappa_k`` for
``k in {0, 1, 2, 3, 4}``, plus two derived constants:
- ``"C"``: the paper's boundary-kernel constant used in the asymptotic
bias term ``h^2 * C * m''(0)``. Per de Chaisemartin, Ciccia,
D'Haultfoeuille & Knau (2026, Section 3.1.3),
``C = (kappa_2^2 - kappa_1 * kappa_3) / (kappa_0 * kappa_2 - kappa_1^2)``.
- ``"kstar_L2_norm"``: the asymptotic-variance constant
``int_0^1 k*(t)^2 dt`` where
``k*(t) = (kappa_2 - kappa_1 * t) / (kappa_0 * kappa_2 - kappa_1^2) * k(t)``
is the equivalent kernel for local-linear at a boundary. Computed by
numerical integration.
Parameters
----------
kernel : str
One of ``"epanechnikov"``, ``"triangular"``, ``"uniform"``.
Returns
-------
dict of {str: float}
Keys ``kappa_0``, ``kappa_1``, ``kappa_2``, ``kappa_3``, ``kappa_4``,
``C``, ``kstar_L2_norm``.
Raises
------
ValueError
If ``kernel`` is not a recognized name.
"""
if kernel not in _CLOSED_FORM_MOMENTS:
raise ValueError(
f"Unknown kernel {kernel!r}. Expected one of " f"{sorted(_CLOSED_FORM_MOMENTS.keys())}."
)
kappas = dict(_CLOSED_FORM_MOMENTS[kernel])
k0, k1, k2, k3 = kappas["kappa_0"], kappas["kappa_1"], kappas["kappa_2"], kappas["kappa_3"]
denom = k0 * k2 - k1 * k1
C = (k2 * k2 - k1 * k3) / denom
kappas["C"] = C
kfun = KERNELS[kernel]
def _kstar_sq(t: float) -> float:
kt = kfun(np.array([t]))[0]
return ((k2 - k1 * t) / denom) ** 2 * kt * kt
val, _ = integrate.quad(_kstar_sq, 0.0, 1.0, limit=200)
kappas["kstar_L2_norm"] = float(val)
return kappas
# =============================================================================
# Local-linear regression at a boundary
# =============================================================================
@dataclass
class LocalLinearFit:
"""Result of a local-linear regression at a boundary.
Attributes
----------
intercept : float
Estimated conditional mean at the boundary, ``mu_hat_h(d0)``.
slope : float
Estimated slope of the local linear fit (coefficient on ``d - d0``).
n_effective : int
Count of observations with strictly positive kernel weight (within
``[d0, d0 + h]`` for the one-sided kernels shipped here).
bandwidth : float
Bandwidth ``h`` used.
kernel : str
Kernel name.
boundary : float
Evaluation point ``d0``.
residuals : np.ndarray, shape (n_effective,)
Residuals from the weighted OLS fit, in the order of the retained
observations.
kernel_weights : np.ndarray, shape (n_effective,)
Kernel weights ``k((d_i - d0) / h)``. These are the pre-scaled weights;
the ``1/h`` scaling cancels out of the weighted-OLS estimator (a
constant factor on all weights does not change the point estimate).
design_matrix : np.ndarray, shape (n_effective, 2)
Design matrix ``X = [1, d_i - d0]`` used in the fit. Preserved for
Phase 1c bias-correction machinery.
"""
intercept: float
slope: float
n_effective: int
bandwidth: float
kernel: str
boundary: float
residuals: np.ndarray
kernel_weights: np.ndarray
design_matrix: np.ndarray
def local_linear_fit(
d: np.ndarray,
y: np.ndarray,
bandwidth: float,
boundary: float = 0.0,
kernel: str = "epanechnikov",
weights: Optional[np.ndarray] = None,
) -> LocalLinearFit:
"""Local-linear regression of ``y`` on ``d`` at a boundary.
Fits ``y ~ a + b * (d - boundary)`` using kernel weights
``k((d - boundary) / h)`` on observations with ``d in [boundary,
boundary + h]``. Returns the intercept (the boundary estimate ``mu_hat``)
and slope.
Parameters
----------
d : np.ndarray, shape (n,)
Regressor values. For the HAD application, ``d`` is the period-2 dose
``D_{g,2}`` and the boundary is 0.
y : np.ndarray, shape (n,)
Outcome values. For the HAD application, ``y`` is the first-difference
``Delta Y_g``.
bandwidth : float
Bandwidth ``h > 0``.
boundary : float, default=0.0
Evaluation point ``d0``. Observations with ``d < d0`` are excluded
(one-sided boundary estimation).
kernel : str, default="epanechnikov"
One of ``"epanechnikov"``, ``"triangular"``, ``"uniform"``.
weights : np.ndarray or None, optional
Optional per-observation weights ``w_i >= 0`` multiplied into the
kernel weights. Useful for survey weighting; when ``None``, treated as
unit weights.
Returns
-------
LocalLinearFit
Named container with ``intercept``, ``slope``, ``n_effective``, and
diagnostics needed by downstream bias-correction phases.
Raises
------
ValueError
If ``bandwidth <= 0``, ``kernel`` is unknown, ``d`` and ``y`` differ
in length, or the bandwidth window retains fewer than 2 observations.
"""
if bandwidth <= 0.0:
raise ValueError(f"bandwidth must be positive; got {bandwidth}")
if kernel not in KERNELS:
raise ValueError(f"Unknown kernel {kernel!r}. Expected one of {sorted(KERNELS.keys())}.")
d = np.asarray(d, dtype=np.float64).ravel()
y = np.asarray(y, dtype=np.float64).ravel()
if d.shape != y.shape:
raise ValueError(f"d and y must have the same shape; got {d.shape} and {y.shape}")
# Explicit NaN / Inf validation at the API boundary so the caller gets a
# targeted error rather than a downstream failure inside the kernel or OLS.
if not np.all(np.isfinite(d)):
raise ValueError("d contains non-finite values (NaN or Inf)")
if not np.all(np.isfinite(y)):
raise ValueError("y contains non-finite values (NaN or Inf)")
if weights is None:
user_w = np.ones_like(d)
else:
user_w = np.asarray(weights, dtype=np.float64).ravel()
if user_w.shape != d.shape:
raise ValueError(
f"weights must have the same shape as d; got " f"{user_w.shape} vs {d.shape}"
)
if not np.all(np.isfinite(user_w)):
raise ValueError("weights contains non-finite values (NaN or Inf)")
if np.any(user_w < 0):
raise ValueError("weights must be nonnegative")
# Kernel weights on the scaled domain u = (d - d0) / h.
kfun = KERNELS[kernel]
u = (d - boundary) / bandwidth
k_weights_full = kfun(u)
# Compose with user weights and restrict to the bandwidth window.
combined = k_weights_full * user_w
retain = combined > 0.0
n_effective = int(retain.sum())
if n_effective < 2:
raise ValueError(
f"bandwidth window retained {n_effective} observation(s); "
f"need at least 2. Widen the bandwidth or move the boundary."
)
d_in = d[retain]
y_in = y[retain]
w_in = combined[retain]
k_in = k_weights_full[retain]
design = np.column_stack([np.ones_like(d_in), d_in - boundary])
# Weighted OLS via solve_ols. "aweight" treats the weights as analytic
# frequency weights so the unweighted-OLS formulas apply with w-scaled X.
# We only need the coefficients and residuals, not a vcov for the fit
# itself (Phase 1c will build its own bias-aware variance).
coef, residuals, _ = solve_ols(
design,
y_in,
cluster_ids=None,
return_vcov=False,
weights=w_in,
weight_type="aweight",
)
return LocalLinearFit(
intercept=float(coef[0]),
slope=float(coef[1]),
n_effective=n_effective,
bandwidth=float(bandwidth),
kernel=kernel,
boundary=float(boundary),
residuals=np.asarray(residuals, dtype=np.float64),
kernel_weights=np.asarray(k_in, dtype=np.float64),
design_matrix=np.asarray(design, dtype=np.float64),
)
# =============================================================================
# MSE-optimal bandwidth selector (Phase 1b)
# =============================================================================
#
# Public wrapper around diff_diff._nprobust_port.lpbwselect_mse_dpi. The port
# is a faithful Python translation of nprobust::lpbwselect(bwselect="mse-dpi")
# (R package nprobust 0.5.0, SHA 36e4e53); see diff_diff/_nprobust_port.py for
# source mapping.
#
# The public API here is a thin wrapper that:
# 1. Accepts the diff-diff library's kernel naming convention (full words) and
# translates to nprobust's short codes ("epa", "tri", "uni").
# 2. Converts the internal MseDpiStages dataclass to the user-facing
# BandwidthResult dataclass.
# 3. Enforces input validation consistent with local_linear_fit.
# 4. Rejects unsupported combinations (e.g. weights=) with NotImplementedError
# (nprobust has no weight support).
@dataclass
class BandwidthResult:
"""MSE-optimal bandwidth selector output plus per-stage diagnostics.
Returned by ``mse_optimal_bandwidth(..., return_diagnostics=True)``.
Mirrors the five-bandwidth + four-stage structure of
``nprobust::lpbwselect.mse.dpi``; see
``diff_diff/_nprobust_port.py`` for the source mapping.
Attributes
----------
h_mse : float
Final MSE-optimal bandwidth ``h*`` for local-linear estimation at
``boundary``. The argument to pass to
``local_linear_fit(..., bandwidth=h_mse)``.
b_mse : float
Bias-correction bandwidth. Consumed by Phase 1c for the
bias-corrected confidence interval (CCF 2018 Equation 8).
c_bw : float
Stage 1 preliminary bandwidth used as ``h.V`` in every
``lprobust.bw`` call downstream:
``C_kernel * min(sd(d), IQR(d)/1.349) * G^{-1/5}``.
Kernel constants: ``epa=2.34``, ``uni=1.843``, ``tri=2.576``.
bw_mp2 : float
Stage 2 pilot bandwidth for the ``m^{(q+1)}`` derivative estimator.
bw_mp3 : float
Stage 2 pilot bandwidth for the ``m^{(q+2)}`` derivative estimator.
stage_d1_V, stage_d1_B1, stage_d1_B2, stage_d1_R : float
Variance and bias coefficients from the first Stage-2
``lprobust.bw`` call (order ``q+1``, reading the ``(q+1)``-th
derivative). Parity-checked to 1% against R.
stage_d2_V, stage_d2_B1, stage_d2_B2, stage_d2_R : float
Same for the second Stage-2 call (order ``q+2``).
stage_b_V, stage_b_B1, stage_b_B2, stage_b_R : float
Same for the Stage-3 bias-bandwidth call (order ``q``, nu ``p+1``).
stage_h_V, stage_h_B1, stage_h_B2, stage_h_R : float
Same for the final Stage-3 main-bandwidth call (order ``p``,
nu ``deriv``).
n : int
Sample size.
kernel : str
Kernel name as user supplied it ("epanechnikov" / "triangular" /
"uniform").
boundary : float
Evaluation point ``d_0``.
"""
h_mse: float
b_mse: float
c_bw: float
bw_mp2: float
bw_mp3: float
stage_d1_V: float
stage_d1_B1: float
stage_d1_B2: float
stage_d1_R: float
stage_d2_V: float
stage_d2_B1: float
stage_d2_B2: float
stage_d2_R: float
stage_b_V: float
stage_b_B1: float
stage_b_B2: float
stage_b_R: float
stage_h_V: float
stage_h_B1: float
stage_h_B2: float
stage_h_R: float
n: int
kernel: str
boundary: float
# Mapping between diff-diff's full kernel names and nprobust's short codes.
_KERNEL_NAME_TO_NPROBUST = {
"epanechnikov": "epa",
"triangular": "tri",
"uniform": "uni",
}
def _validate_had_inputs(
d: np.ndarray,
y: np.ndarray,
boundary: float,
) -> tuple[np.ndarray, np.ndarray]:
"""HAD-scope input validation shared across Phase 1b and Phase 1c.
Coerces ``d`` and ``y`` to 1-D ``float64`` arrays and enforces the
HAD input contract used by both ``mse_optimal_bandwidth`` (Phase 1b)
and ``bias_corrected_local_linear`` (Phase 1c). Each caller handles
its own ``weights=`` rejection upfront because the phase name and
message differ; every other rule lives here.
Rules (matching Phase 1b's original inline checks verbatim):
- ``d`` and ``y`` must be shape-matched, non-empty, and finite.
- ``boundary`` must be finite.
- Doses are nonnegative (HAD support ``D_{g,2} >= 0``).
- ``boundary`` is approximately ``0`` (Design 1') or approximately
``d.min()`` (Design 1 continuous-near-d_lower); other boundaries
raise so off-support calls do not silently target a different
estimand.
- Mass-point designs at ``d.min() > 0`` with modal-min fraction
``> 2%`` raise ``NotImplementedError`` pointing to Phase 2's 2SLS
path, per the paper's Section 3.2.4.
- Design 1' plausibility heuristic: when ``boundary ~ 0``, require
``d.min() <= 0.05 * median(|d|)``; samples above that ratio are
redirected to ``boundary=float(d.min())``.
Parameters
----------
d, y : np.ndarray
Regressor (dose) and outcome.
boundary : float
Evaluation point.
Returns
-------
tuple[np.ndarray, np.ndarray]
Coerced ``(d, y)`` as 1-D float64 arrays.
Raises
------
ValueError, NotImplementedError
"""
if not np.isfinite(boundary):
raise ValueError(f"boundary must be finite; got {boundary}")
d = np.asarray(d, dtype=np.float64).ravel()
y = np.asarray(y, dtype=np.float64).ravel()
if d.shape != y.shape:
raise ValueError(f"d and y must have the same shape; got {d.shape} and {y.shape}")
if d.size == 0:
raise ValueError(
"d and y must be non-empty; the selector cannot estimate a "
"bandwidth from zero observations."
)
if not np.all(np.isfinite(d)):
raise ValueError("d contains non-finite values (NaN or Inf)")
if not np.all(np.isfinite(y)):
raise ValueError("y contains non-finite values (NaN or Inf)")
# HAD support restriction: de Chaisemartin et al. (2026) Assumption
# (dose definition in Section 2) treats ``D_{g,2}`` as the period-2
# treatment dose with ``D_{g,2} >= 0``. Negative dose values are
# outside the HAD design and would silently calibrate the selector
# against a symmetric-kernel two-sided problem while the downstream
# fitter remains one-sided. Reject front-door rather than produce a
# plausible bandwidth on a malformed input.
d_neg = d < 0.0
if np.any(d_neg):
n_neg = int(d_neg.sum())
min_neg = float(d[d_neg].min())
raise ValueError(
f"Negative dose values detected in d (n_neg={n_neg}, "
f"min={min_neg!r}). The HAD estimator (de Chaisemartin et "
f"al. 2026) requires the period-2 dose D_{{g,2}} >= 0. "
f"Nonnegative-dose data is required for both Design 1' "
f"(d_lower = 0) and Design 1 (d_lower > 0)."
)
# Boundary-applicability check (Phase 1b scope).
# The exported wrapper is scoped to the two documented HAD
# nonparametric estimands:
# - Design 1' evaluates m(0) = lim_{d down 0} E[Delta Y | D_2 <= d]
# with ``boundary = 0``.
# - Design 1 continuous-near-d_lower evaluates m(d_lower) with
# ``boundary = d_lower = min(D_2)``.
# Any other off-support boundary (interior, upper-boundary, or an
# arbitrary value between 0 and d.min()) would silently target an
# undocumented limit and is rejected.
d_min = float(d.min())
_boundary_tol = 1e-12 * max(1.0, abs(d_min), abs(boundary))
_at_zero = abs(boundary) <= _boundary_tol
_at_d_min = abs(boundary - d_min) <= _boundary_tol
if not (_at_zero or _at_d_min):
raise ValueError(
f"boundary={boundary!r} is not at a supported HAD estimand. "
f"The Phase 1b public wrapper accepts only boundary ~ 0 "
f"(Design 1') or boundary ~ d.min()={d_min!r} (Design 1 "
f"continuous-near-d_lower). Off-support values would "
f"silently target an undocumented limit. For interior or "
f"other boundary points, use "
f"diff_diff._nprobust_port.lpbwselect_mse_dpi directly and "
f"note that those paths are not separately parity-tested."
)
# Mass-point design check (paper Section 3.2.4, REGISTRY 2% rule).
# Must fire BEFORE the Design 1' support check: mass-point data is
# never appropriate for the CCF nonparametric selector regardless
# of the boundary the caller supplied. The correct remediation is
# the 2SLS sample-average path (Phase 2), not a boundary
# reclassification.
#
# The check explicitly excludes d_min ~ 0 (the Design 1'
# "untreated units present" subcase that the paper's simulations
# and the Garrett et al. (2020) application accept).
_MASS_POINT_THRESHOLD = 0.02 # REGISTRY rule: > 2% modal-min
if d_min > _boundary_tol:
eps_eq = 1e-12 * max(1.0, abs(d_min))
at_d_min_mask = np.abs(d - d_min) <= eps_eq
modal_fraction = float(np.mean(at_d_min_mask))
if modal_fraction > _MASS_POINT_THRESHOLD:
raise NotImplementedError(
f"Detected mass-point design at d.min()={d_min!r} "
f"(modal fraction {modal_fraction:.4f} > "
f"{_MASS_POINT_THRESHOLD:.2f}). Per de Chaisemartin et "
f"al. (2026) Section 3.2.4, Design 1 mass-point cases "
f"require the 2SLS sample-average estimator with "
f"instrument 1{{D_2 > d_lower}}, not the CCF "
f"nonparametric selector. That path is queued for "
f"Phase 2 (HeterogeneousAdoptionDiD). For continuous "
f"near-d_lower designs (modal fraction <= "
f"{_MASS_POINT_THRESHOLD:.2f}), this wrapper is "
f"applicable."
)
# Design 1' support check: boundary ~ 0 requires the realized
# sample minimum to be compatible with a population support
# infimum at 0. Otherwise the selector calibrates ``h_mse`` at
# an off-support limit.
#
# Rule: when boundary ~ 0 (not also at d.min()), require
# d.min() <= 5% * median(|d|). The 5% threshold is generous
# enough to accept Design 1' samples with vanishing boundary
# density (Beta(2,2): d.min/median ~ 3%) while rejecting samples
# substantially off-support (U(0.5, 1): d.min/median ~ 1.0).
# Samples just between these (e.g. U(0.05, 1), d.min/median ~ 10%)
# are directed to boundary=float(d.min()) for the continuous-
# near-d_lower path.
_DESIGN_1_PRIME_RATIO = 0.05
if _at_zero and not _at_d_min:
d_median_abs = float(np.median(np.abs(d)))
effective_threshold = _DESIGN_1_PRIME_RATIO * max(d_median_abs, 1e-12)
if d_min > effective_threshold:
raise ValueError(
f"boundary ~ 0 selected but d.min()={d_min!r} is not "
f"compatible with a Design 1' support infimum at 0 "
f"(rule: d.min() <= "
f"{_DESIGN_1_PRIME_RATIO} * median(|d|) = "
f"{effective_threshold!r}). This sample is not "
f"Design 1'. Either: (a) pass boundary=float(d.min()) "
f"for the Design 1 continuous-near-d_lower path, or "
f"(b) verify the population support actually has "
f"infimum at 0 (in which case the realized d.min() "
f"would be closer to zero relative to the data scale)."
)
return d, y
def mse_optimal_bandwidth(
d: np.ndarray,
y: np.ndarray,
boundary: float = 0.0,
kernel: str = "epanechnikov",
weights: Optional[np.ndarray] = None,
bwcheck: Optional[int] = 21,
bwregul: float = 1.0,
return_diagnostics: bool = False,
):
"""MSE-optimal bandwidth for local-linear regression at a boundary.
Port of ``nprobust::lpbwselect(bwselect="mse-dpi")`` (R package
``nprobust`` 0.5.0, SHA ``36e4e53``). Implements the Calonico,
Cattaneo, and Farrell (2018, JASA 113(522)) direct-plug-in DPI
bandwidth selector for the local-linear boundary estimator used in
Design 1' of de Chaisemartin et al. (2026) (HAD).
The three-stage algorithm produces five bandwidths (``c_bw``,
``bw_mp2``, ``bw_mp3``, ``b_mse``, ``h_mse``); see
``BandwidthResult`` for full diagnostics.
**Public API scope (Phase 1b of HAD).** This wrapper is intentionally
restricted to the HAD configuration: ``p=1`` (local-linear),
``deriv=0`` (mean regression), ``interior=False`` (boundary eval
point), ``vce="nn"`` (nearest-neighbor variance), and ``nnmatch=3``
are hard-coded. The underlying ``diff_diff._nprobust_port`` supports
additional ``vce`` modes (``hc0`` / ``hc1`` / ``hc2`` / ``hc3``),
interior evaluation, and higher polynomial orders, but those paths
are NOT parity-tested against ``nprobust`` and are deferred to Phase
1c / Phase 2. Do not rely on this public wrapper for anything
outside HAD Phase 1b; use the port directly if you need the broader
surface and accept that parity has not been separately verified.
Parameters
----------
d : np.ndarray, shape (G,)
Regressor values (the dose ``D_{g,2}`` in HAD).
y : np.ndarray, shape (G,)
Outcome values (the first-difference ``Delta Y_g`` in HAD).
boundary : float, default=0.0
Evaluation point ``d_0``. The Phase 1b wrapper accepts only
two values (within float tolerance): ``boundary = 0`` for
Design 1' or ``boundary = float(d.min())`` for Design 1
continuous-near-``d_lower``. Use the sample minimum
``d.min()`` (not a known theoretical lower bound of the
support), because the downstream selector operates on the
realized data. Any other value -- including
``boundary < d.min()``, interior points, or
``boundary > d.min()`` -- raises ``ValueError``.
kernel : str, default="epanechnikov"
One of ``"epanechnikov"``, ``"triangular"``, ``"uniform"``.
weights : np.ndarray or None, default=None
Not supported in Phase 1b (raises ``NotImplementedError``).
``nprobust::lpbwselect`` has no weight argument and thus no
parity anchor. Weighted-data support is queued for Phase 2+.
bwcheck : int or None, default=21
If set, clip ``c_bw`` (and all downstream bandwidths) below the
distance to the ``bwcheck``-th nearest neighbor of ``boundary``.
Matches ``nprobust::lpbwselect`` default.
bwregul : float, default=1.0
Bias-regularization scale used in Stage 3. ``bwregul=0`` disables
the regularization term; ``bwregul=1`` matches ``nprobust``
default.
return_diagnostics : bool, default=False
When ``False`` (default) return the final bandwidth as a
``float``. When ``True`` return a ``BandwidthResult`` with all
five stage bandwidths plus per-stage ``(V, B1, B2, R)``
diagnostics.
Returns
-------
float or BandwidthResult
The MSE-optimal bandwidth ``h*``, or (if
``return_diagnostics=True``) a ``BandwidthResult`` dataclass.
Raises
------
ValueError
Raised on: shape mismatch between ``d`` and ``y``; non-finite
values in ``d``, ``y``, or ``boundary``; unknown ``kernel``
name; ``bwcheck`` outside ``[1, len(d)]``; ``boundary`` that
is not approximately 0 or approximately ``d.min()`` (the only
two supported HAD estimands in Phase 1b); or a rank-deficient
/ under-determined pilot fit inside the DPI port (surfaced
from ``qrXXinv`` or the per-stage count guards in
``lprobust_bw``).
NotImplementedError
Raised on: ``weights=`` passed (no nprobust parity anchor);
detected Design 1 mass-point design (``d.min() > 0`` and
modal fraction at ``d.min()`` exceeds 2%, per the paper's
Section 3.2.4 redirection to the 2SLS sample-average
estimator, queued for Phase 2).
Notes
-----
The port parity-tests at 1% relative error against R
``nprobust::lpbwselect(bwselect="mse-dpi", vce="nn")`` on three
deterministic DGPs (see
``benchmarks/R/generate_nprobust_golden.R``). In practice the
agreement is at machine precision (0.0000% on the golden tests);
the 1% tolerance is the hard gate before a deviation is considered
a regression.
"""
if weights is not None:
raise NotImplementedError(
"weights= is not supported in Phase 1b of the MSE-optimal "
"bandwidth selector. nprobust::lpbwselect has no weight "
"argument, so there is no parity anchor. Weighted-data "
"support is queued for Phase 2+ (survey-design adaptation)."
)
if kernel not in _KERNEL_NAME_TO_NPROBUST:
raise ValueError(
f"Unknown kernel {kernel!r}. Expected one of "
f"{sorted(_KERNEL_NAME_TO_NPROBUST.keys())}."
)
nprobust_kernel = _KERNEL_NAME_TO_NPROBUST[kernel]
d, y = _validate_had_inputs(d, y, boundary)
# Defer heavy import to call time to avoid import-cycle risk.
from diff_diff._nprobust_port import lpbwselect_mse_dpi
stages = lpbwselect_mse_dpi(
y=y,
x=d,
cluster=None,
eval_point=float(boundary),
p=1,
q=2,
deriv=0,
kernel=nprobust_kernel,
bwcheck=bwcheck,
bwregul=bwregul,
vce="nn",
nnmatch=3,
interior=False,
)
if not return_diagnostics:
return stages.h_mse_dpi
return BandwidthResult(
h_mse=stages.h_mse_dpi,
b_mse=stages.b_mse_dpi,
c_bw=stages.c_bw,
bw_mp2=stages.bw_mp2,
bw_mp3=stages.bw_mp3,
stage_d1_V=stages.stage_d1.V,
stage_d1_B1=stages.stage_d1.B1,
stage_d1_B2=stages.stage_d1.B2,
stage_d1_R=stages.stage_d1.R,
stage_d2_V=stages.stage_d2.V,
stage_d2_B1=stages.stage_d2.B1,
stage_d2_B2=stages.stage_d2.B2,
stage_d2_R=stages.stage_d2.R,
stage_b_V=stages.stage_b.V,
stage_b_B1=stages.stage_b.B1,
stage_b_B2=stages.stage_b.B2,
stage_b_R=stages.stage_b.R,
stage_h_V=stages.stage_h.V,
stage_h_B1=stages.stage_h.B1,
stage_h_B2=stages.stage_h.B2,
stage_h_R=stages.stage_h.R,
n=int(d.shape[0]),
kernel=kernel,
boundary=float(boundary),
)
# =============================================================================
# Bias-corrected local-linear fit (Phase 1c)
# =============================================================================
#
# Public wrapper around diff_diff._nprobust_port.lprobust. The port is a
# faithful Python translation of the single-eval-point path of
# nprobust::lprobust (R package nprobust 0.5.0, SHA 36e4e53) implementing
# Calonico, Cattaneo, and Titiunik (2014) robust bias correction.
#
# This wrapper produces the mu-scale quantities of Equation 8 in de
# Chaisemartin, Ciccia, D'Haultfoeuille, and Knau (2026): a classical and
# bias-corrected point estimate plus naive and robust standard errors and
# the bias-corrected confidence interval [tau.bc +/- z_{1-alpha/2} * se.rb].
# The beta-scale rescaling in Equation 8 (divide by ``(1/G) * sum(D_{g,2})``)
# is applied by Phase 2's ``HeterogeneousAdoptionDiD.fit()``, not here.
@dataclass
class BiasCorrectedFit:
"""Bias-corrected local-linear fit at a boundary (Phase 1c).
Output of :func:`bias_corrected_local_linear`. Produces the mu-scale
quantities needed by Equation 8 of de Chaisemartin, Ciccia,
D'Haultfoeuille, and Knau (2026). Phase 2's ``HeterogeneousAdoptionDiD``
class applies the beta-scale ``(1/G) * sum(D_{g,2})`` rescaling.
Attributes
----------
estimate_classical : float
Classical point estimate ``tau.cl`` from ``nprobust::lprobust``
(local-linear boundary intercept at ``h``; no bias correction).
estimate_bias_corrected : float
Bias-corrected point estimate ``tau.bc = mu_hat + M_hat`` from the
Calonico-Cattaneo-Titiunik (2014) combined design-matrix statistic.
se_classical : float
Naive plug-in standard error.
se_robust : float
Robust standard error accounting for the additional variability
introduced by the bias-correction term (CCT 2014).
ci_low, ci_high : float
Endpoints of the bias-corrected CI: ``tau.bc +/- z_{1-alpha/2} *
se.rb``.
alpha : float
CI level (``0.05`` gives a 95% CI).
h, b : float
Main and bias-correction bandwidths actually used (post-``bwcheck``
floor).
bandwidth_source : {"auto", "user"}
``"auto"`` when the wrapper called the Phase 1b DPI selector
(``_nprobust_port.lpbwselect_mse_dpi``) internally with the
caller's ``cluster`` / ``vce`` / ``nnmatch``; ``"user"`` when the
caller passed explicit bandwidths. Auto mode then enforces
nprobust's ``rho=1`` default by setting ``b = h``; the
selector's distinct ``b_mse`` is surfaced via
``bandwidth_diagnostics`` but not applied.
bandwidth_diagnostics : BandwidthResult or None
Full Phase 1b selector output when ``bandwidth_source == "auto"``;
``None`` when the user supplied bandwidths (to avoid a redundant
selector call).
n_used : int
Observations retained in the active kernel window (``sum(ind.b)``
when ``h <= b`` and ``sum(ind.h)`` when ``h > b``; with the
``rho=1`` default the two coincide).
n_total : int
Total observations passed in (before kernel filtering).
kernel : str
Kernel name as supplied by the caller.
boundary : float
Evaluation point ``c``.
Notes
-----
``p=1``, ``q=2``, ``deriv=0`` are hard-coded for HAD Phase 1c and are
not exposed as fields. Phase 2 may surface them on the estimator-level
result class if a use case materializes.
"""
estimate_classical: float
estimate_bias_corrected: float
se_classical: float
se_robust: float
ci_low: float
ci_high: float
alpha: float
h: float
b: float
bandwidth_source: str
bandwidth_diagnostics: Optional[BandwidthResult]
n_used: int
n_total: int
kernel: str
boundary: float
influence_function: Optional[np.ndarray] = None
"""Per-observation influence function of the BIAS-CORRECTED point
estimate ``tau.bc`` (Phase 4.5 survey composition). Aligned with
the original caller-supplied ``d``/``y`` ordering; observations
outside the active kernel window have IF=0. Populated only when
``return_influence=True``; ``None`` otherwise.
Derived from ``Q_q`` + ``res_b`` so the variance self-check is
``sum(IF^2) == V_Y_bc[0, 0]`` under unclustered HC0 — matching the
bias-corrected scale of ``estimate_bias_corrected``. Using the
classical IF here would silently under-estimate survey SE by
ignoring the CCT-2014 bias-correction variance inflation."""
def bias_corrected_local_linear(
d: np.ndarray,
y: np.ndarray,
boundary: float = 0.0,
kernel: str = "epanechnikov",
h: Optional[float] = None,
b: Optional[float] = None,
alpha: float = 0.05,
vce: str = "nn",
cluster: Optional[np.ndarray] = None,
nnmatch: int = 3,
weights: Optional[np.ndarray] = None,
return_influence: bool = False,
) -> BiasCorrectedFit:
"""Bias-corrected local-linear fit with robust CI at a boundary.
Ports the single-eval-point path of ``nprobust::lprobust`` (Calonico,
Cattaneo, and Titiunik 2014) and produces the mu-scale outputs of
Equation 8 in de Chaisemartin, Ciccia, D'Haultfoeuille, and Knau
(2026). Phase 2's ``HeterogeneousAdoptionDiD`` class applies the
beta-scale rescaling to return the final estimand.
**Public API scope (Phase 1c of HAD).** Hard-coded for the HAD
configuration: ``p=1`` (local-linear), ``q=2`` (bias-correction
order), ``deriv=0`` (level), ``interior=False`` (boundary eval
point), ``bwcheck=21``, ``bwregul=1``, and ``vce="nn"`` (the only
variance mode golden-parity-tested against R in this phase). The
underlying ``diff_diff._nprobust_port.lprobust`` accepts the broader
surface (hc0/hc1/hc2/hc3, higher ``p``, interior eval), but those
paths are not separately parity-tested and remain private until
Phase 2+ ships dedicated goldens.
Parameters
----------
d : np.ndarray, shape (G,)
Regressor (dose ``D_{g,2}`` in HAD).
y : np.ndarray, shape (G,)
Outcome (first-difference ``Delta Y_g`` in HAD).
boundary : float, default=0.0
Evaluation point ``d_0``. Accepts only ``boundary ~ 0`` (Design 1')
or ``boundary ~ float(d.min())`` (Design 1 continuous-near-d_lower);
see Phase 1b's input contract.
kernel : {"epanechnikov", "triangular", "uniform"}, default="epanechnikov"
h : float or None, default=None
Main bandwidth. ``None`` auto-selects ``h`` by calling
``diff_diff._nprobust_port.lpbwselect_mse_dpi`` directly with the
supplied ``cluster``, ``vce``, and ``nnmatch`` so the selector
and the final fit use the same estimator. ``b`` is then set to
``h`` per nprobust's ``rho=1`` default; the selector's distinct
``b_mse`` is surfaced through
``BiasCorrectedFit.bandwidth_diagnostics`` for inspection but