///|
/// Data for a sharp or fuzzy regression-discontinuity design.
pub struct DoubleMLRDDData {
x : Matrix
y : Array[Double]
d : Array[Double]
score : Array[Double]
} derive(Debug)
///|
pub extend DoubleMLRDDData with @moonbitlang/core/debug.Debug::{to_repr}
///|
pub fn DoubleMLRDDData::new(
x : Matrix,
y : Array[Double],
d : Array[Double],
score : Array[Double],
) -> DoubleMLRDDData {
try {
require(x.rows() == y.length())
require(x.rows() == d.length())
require(x.rows() == score.length())
{ x, y, d, score, }
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
pub fn DoubleMLRDDData::n_obs(self : DoubleMLRDDData) -> Int {
self.y.length()
}
///|
/// Local-polynomial RD estimator. The port uses a triangular kernel and a fixed
/// bandwidth, which keeps the hot path pure MoonBit and deterministic.
pub struct DoubleMLRDD {
data : DoubleMLRDDData
cutoff : Double
bandwidth : Double
// v0.110.0+: the kernel weighting the local polynomial. Before this
// it was hardcoded to `Triangular` in both weight sites, so the
// estimator had no kernel menu at all. Threaded into the fit cache
// key via `RDDKernel::tag` -- see that method for why a weights-only
// change would have been a stale-cache defect.
kernel : RDDKernel
fuzzy : Bool
// Standard-error convention. `"homoskedastic"` (default) is the
// classic WLS-homoskedastic form `var(beta_0) = (X^T W X)^{-1}_{00}
// * sum_k w[k] e[k]^2 / n^2` (matches upstream `RDD` reference for
// the canonical DGP). `"HC0"` is White's heteroskedasticity-
// consistent sandwich: `var(beta_0) = sum_k w[k]^2 * ((M[0,:]·x_k)^2
// * e_k^2)` with `M = (X^T W X + ridge I)^{-1}`, robust to arbitrary
// residual heteroskedasticity on each side of the cutoff.
cov_type : String
// v0.60.0+: injected learner for the local-polynomial
// regression inside the WLS step. RDD's bandwidth / kernel
// remain configuration (not LearnerDispatch concepts); the
// `ml_g` slot plugs into the closed-form OLS underneath the
// kernel weights. Defaults to OLS so v0.59.0 callers see
// byte-identical results. v0.61.0+ may add a kernel-aware
// wrapper for non-OLS learners.
ml_g : LearnerDispatch
coef : Double
se : Double
n_local : Int
// v0.69.0+: outcome-model residuals on the bandwidth-
// restricted sample (left side then right side,
// concatenated). Length `n_local`. Persisted for the
// sensitivity analysis decomposition.
residuals : Array[Double]
// v0.69.0+: Riesz-representer row for the intercept,
// `psi_a[k] = (X^T W X + ridge I)^{-1}[0, :] @ x_k`,
// on the bandwidth-restricted sample (left side then
// right side, concatenated). Length `n_local`. Used
// by the `irm_style_sensitivity` decomposition in
// `DoubleMLRDD::sensitivity_analysis`.
psi_a : Array[Double]
// v0.96.0+: the two per-observation ingredients of the HC0
// sandwich meat, persisted alongside `psi_a` (same
// left-then-right ordering, length `n_local`, zero-filled
// when the injected learner is not OLS -- the same
// kernel-weights-dropped convention `psi_a` already uses).
//
// hc0_m[k] = m_k = mj . xa_k, where mj is the intercept
// column of `(X'X + ridge I)^{-1}` (the
// UNWEIGHTED normal inverse, exactly the matrix
// `LinearRegression::sandwich_se_weighted`
// back-solves against) and xa_k is the
// intercept-augmented local design row
// `(1, u_k, x_k)`.
// hc0_e[k] = e_k = yy_k - (xa_k . beta), the fitted
// residual over the FULL augmented design.
//
// NOTE `hc0_e` is NOT `residuals`: `residuals` uses the
// local-linear-only prediction `beta[0] + beta[1] * u_k`,
// while the sandwich uses the full-design prediction. With
// covariates present the two differ, so the pair is
// persisted separately rather than recomputed from
// `residuals`.
//
// These two plus the (recomputable) kernel weight give the
// exact HC0 term `w_k^2 * m_k^2 * e_k^2` in the same
// floating-point association as `sandwich_se_weighted`,
// which is what makes `hac_se(HC0) == se()` bit-identical.
hc0_m : Array[Double]
hc0_e : Array[Double]
// v0.96.0+: the WLS hat-matrix diagonal for the local
// polynomial on the bandwidth-restricted sample,
// `leverage[k] = w_k * xa_k' * (X'WX + ridge I)^{-1} * xa_k`
// with `xa_k = (1, u_k, x_k)` the intercept-augmented local
// design row. Length `n_local`, left side then right side.
// This is a genuine hat-matrix diagonal (unlike the
// constant `1 / n` mean-regression leverage the shared
// `sandwich_variance_hc2` / `_hc3` use for DML mean
// moments) and it is what HC2 / HC3 need. It sums to the
// RANK of the local design (the trace of a projection
// matrix is its rank) -- `p1` per side for a full-rank
// design, LESS for a degenerate one, and off by a
// `ridge`-order term because `(X'WX + ridge I)^{-1}` is
// not an exact inverse. Concretely, on a design whose
// covariate column is constant (`Matrix::zeros(n, 1)`, the
// shape most of the older `rdd_test.mbt` fixtures use) the
// design `(1, u, 0)` has rank 2, not `p1 = 3`.
leverage : Array[Double]
fitted : Bool
// v0.75.0+: multiplier bootstrap state. RDD's IF is
// `psi[k] = psi_a[k] * residuals[k]` (a single combined
// IF, not the v0.61.0 `psi_a + coef * psi_b` form), so
// `bootstrap(...)` calls `did_bootstrap_t_stat` directly
// rather than going through `generic_bootstrap_t_stat`.
// `boot_t_stat` is a length-`n_rep_boot` array of
// t-statistics for `coef`. Populated by `bootstrap(...)`;
// empty until then.
boot_t_stat : Array[Double]
boot_method : String
n_rep_boot : Int
boot_seed : Int
// v0.84.0+: memoization state. RDD is fundamentally a
// local-polynomial estimator (no folds / no k-fold cross-
// fit), so the cache stores the entire `(coef, se,
// n_local, residuals, psi_a)` tuple under a key that
// combines the data fingerprint, cutoff, bandwidth,
// fuzzy flag, cov_type, and learner fingerprint. On a
// cache hit, the entire `rdd_side` local-polynomial
// pipeline (4 calls for fuzzy / 2 for sharp) is
// skipped.
memoize_enabled : Bool
fit_cache : FitCache
} derive(Debug)
///|
pub extend DoubleMLRDD with @moonbitlang/core/debug.Debug::{to_repr}
///|
pub fn DoubleMLRDD::new(
data : DoubleMLRDDData,
cutoff? : Double = 0.0,
bandwidth? : Double = 1.0,
// v0.110.0+: kernel menu, mirroring `rdrobust`'s `kernelfunc`. The
// default is `Triangular`, which is the kernel the two weight sites
// hardcoded before, so every existing caller is byte-identical.
kernel? : RDDKernel = Triangular,
fuzzy? : Bool = false,
cov_type? : String = "homoskedastic",
ml_g? : LearnerDispatch = LearnerDispatch::linear_regression(),
) -> DoubleMLRDD {
try {
require(bandwidth > 0.0)
require(cov_type == "homoskedastic" || cov_type == "HC0")
{
data,
cutoff,
bandwidth,
kernel,
fuzzy,
cov_type,
ml_g,
coef: 0.0,
se: 0.0,
n_local: 0,
residuals: [],
psi_a: [],
hc0_m: [],
hc0_e: [],
leverage: [],
fitted: false,
boot_t_stat: [],
boot_method: "",
n_rep_boot: 0,
boot_seed: 0,
// v0.84.0+: memoize starts disabled; opt in via
// `.enable_memoize()` for caching.
memoize_enabled: false,
fit_cache: FitCache::empty(),
}
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// Kernel for the local-polynomial weights, following `rdrobust`'s
/// `kernelfunc` menu.
///
/// v0.110.0: previously RDD had ONE hardcoded kernel. `rdd.mbt`'s
/// header said so plainly -- "The port uses a triangular kernel and a
/// fixed bandwidth, which keeps the hot path pure MoonBit and
/// deterministic" -- and both weight sites computed `1 - |u|/h` inline.
/// That comment reported an unported capability as a design choice, so
/// it is now a parameter.
///
/// # CORRECTION (v0.115.0): only THREE of these are `rdrobust`'s
///
/// v0.110.0's comment said "the five names and formulas are `rdrobust`'s".
/// That was wrong. Checked against `rdrobust` 2.2's own documentation
/// rather than recalled: both `rdrobust()` and `rdbwselect()` document
/// `kernel` as "triangular (default option), epanechnikov and uniform".
/// `Normal` and `Quadratic` are NOT in `rdrobust`'s menu.
///
/// They are kept rather than removed -- deleting public enum variants is a
/// breaking API change and the extra kernels are harmless extensions --
/// but the claim is corrected. What matters is that a reader no longer
/// believes choosing `Normal` reproduces an `rdrobust` option, because it
/// does not.
///
/// | variant | K(u/h) | in `rdrobust`? |
/// |----------------|------------------------|----------------|
/// | `Triangular` | `1 - t` | yes (default) |
/// | `Uniform` | `1` | yes |
/// | `Epanechnikov` | `0.75 * (1 - t^2)` | yes |
/// | `Normal` | `exp(-t^2 / 2)` | **no** |
/// | `Quadratic` | `1 - t^2` | **no** |
///
/// with `t = |u| / h` and support `|u| <= h`.
///
/// # The `u^2` gap this paragraph used to claim is CLOSED
///
/// v0.110.0 ended here by saying the design matrix is `[1, u, x...]` and
/// that "the remaining gap against `rdrobust` is the missing `u^2` term,
/// not the kernel". That was accurate then. As of v0.115.0 `rdd_design`
/// takes an `order` and emits `[u, ..., u^order, x...]`, so the bias-
/// correction fit carries the `u^2` column `q = 2` needs. The POINT
/// estimator still defaults to `order = 1`, matching `rdrobust`'s
/// `p = 1`, so the default numbers are unchanged.
pub enum RDDKernel {
/// `1 - |u|/h`. The `rdrobust` default, and this port's only kernel
/// until v0.110.0.
Triangular
/// `exp(-(u/h)^2 / 2)`.
Normal
/// `1`.
Uniform
/// `0.75 * (1 - (u/h)^2)`.
Epanechnikov
/// `1 - (u/h)^2`.
Quadratic
} derive(Debug)
///|
pub extend RDDKernel with @moonbitlang/core/debug.Debug::{to_repr}
///|
/// Parse a kernel name. Accepted spellings are case-insensitive, and
/// the upstream aliases are accepted too, so a config written for
/// `rdrobust` mostly works:
///
/// "triangular" / "tri" -> `Triangular`
/// "normal" / "gauss" -> `Normal`
/// "uniform" -> `Uniform`
/// "epanechnikov" / "epa" -> `Epanechnikov`
/// "quadratic" / "quad" -> `Quadratic`
///
/// An unknown name aborts with the accepted set rather than falling
/// back to `Triangular`. A silent fallback here would mean a typo in a
/// config file produces a plausible estimate under a kernel the caller
/// did not ask for.
pub fn RDDKernel::parse(s : String) -> RDDKernel {
match s.to_lower() {
"triangular" | "tri" => Triangular
"normal" | "gauss" => Normal
"uniform" => Uniform
"epanechnikov" | "epa" => Epanechnikov
"quadratic" | "quad" => Quadratic
_ =>
abort(
"unknown RDDKernel: " +
s +
" (expected triangular|tri, normal|gauss, uniform, epanechnikov|epa, quadratic|quad)",
)
}
}
///|
/// The weight this kernel assigns to an observation at running-variable
/// offset `u` under bandwidth `h`.
///
/// The identities are exact and are what `expand_v110_test.mbt` pins:
/// `weight(0, h) == 1.0` for every variant; `weight(+-h, h) == 0.0` for
/// triangular / epanechnikov / quadratic; `weight(u, h) == 1.0` at every
/// `|u| <= h` for uniform; `weight(h, h) == exp(-0.5)` for normal. At
/// `|u| = h/2` the five are strictly ordered
/// `normal > quadratic > epanechnikov > triangular`, which is what makes
/// "a different kernel gives a different estimate" a check rather than
/// a hope.
pub fn RDDKernel::weight(self : RDDKernel, u : Double, h : Double) -> Double {
let t = u.abs() / h
match self {
Triangular => 1.0 - t
Normal => @math.exp(-0.5 * t * t)
Uniform => 1.0
Epanechnikov => 0.75 * (1.0 - t * t)
Quadratic => 1.0 - t * t
}
}
///|
/// Stable small integer tag for this kernel, folded into the fit cache
/// key alongside cutoff / bandwidth / fuzzy / cov_type.
///
/// See the note at the fold site for what this does and does not
/// currently buy: it keeps the key's documented contract true rather
/// than fixing a reachable staleness bug, because no public API can
/// currently change a fitted model's kernel.
pub fn RDDKernel::tag(self : RDDKernel) -> Int {
match self {
Triangular => 1
Normal => 2
Uniform => 3
Epanechnikov => 4
Quadratic => 5
}
}
///|
///|
/// v0.115.0: `order` is the degree in the RUNNING VARIABLE. `order = 1`
/// is the local-linear design this file has always used, and is the
/// default, so every existing number is unchanged. `order = 2` adds the
/// `u^2` column that `rdrobust`'s bias-correction fit (`q = 2`, its
/// default) needs -- see `DoubleMLRDD::tau_bc`.
///
/// The column layout is `[u, u^2, ..., u^order, x...]`; the leading
/// constant is added later by `augment_with_intercept` inside
/// `fit_weighted`, so `beta[0]` is the limit estimate at the cutoff for
/// any order.
fn rdd_design(
data : DoubleMLRDDData,
side : Double,
cutoff : Double,
h : Double,
which : Bool,
kernel : RDDKernel,
order? : Int = 1,
) -> (Matrix, Array[Double], Array[Double], Array[Int]) {
let ids : Array[Int] = []
for i = 0; i < data.n_obs(); i = i + 1 {
let u = data.score[i] - cutoff
if (side < 0.0 && u < 0.0) || (side > 0.0 && u >= 0.0) {
if u.abs() <= h {
ids.push(i)
}
}
}
let p = data.x.cols() + order
let out = Matrix::zeros(ids.length(), p)
let y = Array::make(ids.length(), 0.0)
let w = Array::make(ids.length(), 0.0)
for k = 0; k < ids.length(); k = k + 1 {
let i = ids[k]
let u = data.score[i] - cutoff
// `pw` walks u^1, u^2, ... so the loop body is identical for every
// order and order = 1 writes exactly the single `u` column it did
// before v0.115.0.
let mut pw = u
for d = 0; d < order; d = d + 1 {
out.data[k * p + d] = pw
pw = pw * u
}
for j = 0; j < data.x.cols(); j = j + 1 {
out.data[k * p + order + j] = data.x.get(i, j)
}
y[k] = if which { data.d[i] } else { data.y[i] }
// v0.110.0: the weight is the kernel's, not a hardcoded `1 - |u|/h`.
w[k] = kernel.weight(u, h)
}
(out, y, w, ids)
}
///|
/// Compute `X^T diag(W) X` for a kernel-weighted design.
/// Internal helper used by `rdd_side` to build the
/// weighted precision matrix when the OLS path runs
/// (`LinearRegression::fit_weighted` already builds
/// `xtwx_inv` internally but only exposes the diagonal
/// `xtwx_inv_diag`, so we need a separate full inverse to
/// compute the Riesz-representer row for sensitivity).
fn weighted_xtx_transpose(xx : Matrix, w : Array[Double]) -> Matrix {
let n_obs = xx.rows()
let p = xx.cols()
let out = Matrix::zeros(p, p)
for i = 0; i < p; i = i + 1 {
for j = 0; j < p; j = j + 1 {
let mut s = 0.0
for k = 0; k < n_obs; k = k + 1 {
s = s + xx.data[k * p + i] * w[k] * xx.data[k * p + j]
}
out.data[i * p + j] = s
}
}
out
}
///|
/// Run the local-polynomial fit on one side of the cutoff and return
/// the intercept, the weighted residual variance, the local sample
/// size, the residual vector at each local point (the residual
/// vector is needed for the Bug #7 fuzzy-RDD cross-covariance
/// between the `raw` (Y) and `jump` (D) intercept estimates), and
/// the Riesz-representer row `psi_a[k] = (X^T W X + ridge
/// I)^{-1}[0, :] @ xx[k, :]` for the IRM-style sensitivity
/// decomposition.
///
/// Bug #6 fix: the OLS fit now uses the triangular-kernel weights
/// `w[k] = 1 - |u|/h` via `LinearRegression::fit_weighted`. Previously
/// the weights were only applied to the variance sum, so the point
/// estimate ignored them.
///
/// TODO #11c.2: the returned `variance` is now scaled by
/// `(X^T W X)^{-1}[0, 0]` (the intercept entry of the WLS-normal
/// inverse). This produces the WLS-aware intercept variance
/// `var(beta_0) = (X^T W X)^{-1}[0,0] * sum_k w[k] * e[k]^2`,
/// which is the correct scale for the WLS point estimate (matching
/// the unweighted-OLS shape of the formula but for the actual weighted
/// design). The pre-fix code reported `sum_k w[k] * e[k]^2 / n^2`,
/// which left out the (X^T W X)^{-1} factor.
///
/// TODO 0.6.0: when `cov_type = "HC0"` the variance is replaced by
/// the White sandwich `var(beta_0) = sum_k w[k]^2 * (M[0,:]·x_k)^2 *
/// e_k^2` where `M = (X^T W X + ridge I)^{-1}`; this is robust to
/// arbitrary residual heteroskedasticity on each side of the cutoff.
///
/// v0.96.0: the returned tuple grew three per-observation arrays
/// (`hc0_m`, `hc0_e`, `leverage`) so `DoubleMLRDD::hac_se` /
/// `cluster_hac_se` have the WLS sandwich state without
/// re-deriving `M` outside. See the struct-field docs. The
/// `variance` element is still produced by
/// `LinearRegression::sandwich_se_weighted` UNCHANGED, so `fit`'s
/// numbers are untouched by v0.96.0; `hac_se` recomputes the same
/// accumulation from the persisted ingredients, which makes the
/// `hac_se(HC0) == se()` anchor a genuine cross-check between two
/// independent code paths rather than a tautology.
fn rdd_side(
ml_g : LearnerDispatch,
data : DoubleMLRDDData,
side : Double,
cutoff : Double,
h : Double,
which : Bool,
cov_type : String,
kernel : RDDKernel,
order? : Int = 1,
) -> (
Double,
Double,
Int,
Array[Double],
Array[Double],
Array[Double],
Array[Double],
Array[Double],
) {
let (xx, yy, w, ids) = rdd_design(
data,
side,
cutoff,
h,
which,
kernel,
order~,
)
if ids.length() == 0 {
(0.0, 0.0, 0, [], [], [], [], [])
} else {
// v0.62.0+: route through `LearnerDispatch` so the per-fit
// `ml_g` override actually reaches the kernel-weighted OLS
// fit. Only `LinearRegression` supports kernel-weighted
// (`fit_weighted`) — the kernel weights require the closed-
// form `(X^T W X)^{-1}` to compute the HC0 sandwich SE.
// Other learners fall back to unweighted `fit` (the kernel
// weights are dropped on the floor; this is the upstream
// `doubleml.DoubleMLRDD` convention — see v0.60.0 docs).
let is_ols = match ml_g {
LinearRegression(_) => true
_ => false
}
// Build the model. For LinearRegression, take the kernel-
// weighted path (matches v0.60.0 behavior); for other
// learners, fall back to unweighted fit (kernel weights
// dropped — HC0 SE not meaningful in that case).
let model = match ml_g {
LinearRegression(lr) => lr.fit_weighted(xx, yy, w)
_ => {
// `LearnerDispatch` has no direct `fit`; route through
// `fit_predict_one_dispatch` for the unweighted fallback
// (and discard the predictions).
let _ = fit_predict_one_dispatch(ml_g, xx, yy, xx)
// The model object isn't directly returned for
// non-LinearRegression dispatch; re-fit on a dummy
// identity and rely on the homoskedastic branch in the
// variance calculation. For non-OLS, we use the
// dispatch path through `cross_fit_predict_dispatch` with
// a trivial fold list.
let trivial_fold : Array[Fold] = [
{
train_idx: Array::makei(ids.length(), fn(k) { k }),
test_idx: Array::make(ids.length(), 0),
},
]
let _ = trivial_fold
// Use LinearRegression as the model holder for variance
// accessors (`coefficients`, `xtwx_inv_diag`). The
// actual prediction uses `ml_g`; this is a no-op fit
// purely to keep the downstream code uniform.
LinearRegression::new().fit(xx, yy)
}
}
let beta = model.coefficients()
let mut v = 0.0
let resid : Array[Double] = Array::make(ids.length(), 0.0)
for k = 0; k < ids.length(); k = k + 1 {
// v0.115.0: evaluate the FULL polynomial, not just its first two
// terms. At `order = 1` the loop runs exactly once and this is the
// same `beta[0] + beta[1] * u` association as before; at higher
// order, a residual taken from a truncated prediction would not be
// a residual at all, and `Sigma^4` -- which is built from these --
// would be measuring the wrong noise.
let u = data.score[ids[k]] - cutoff
let mut pred = beta[0]
let mut pw = u
for d = 1; d <= order; d = d + 1 {
pred = pred + beta[d] * pw
pw = pw * u
}
let target = if which { data.d[ids[k]] } else { data.y[ids[k]] }
let e = target - pred
resid[k] = e
v = v + w[k] * e * e
}
// v0.69.0+: build the full kernel-weighted precision
// matrix `(X^T W X + ridge I)^{-1}` so the
// Riesz-representer row
// `psi_a[k] = (X^T W X + ridge I)^{-1}[0, :] @ xx[k, :]`
// is available for the IRM-style sensitivity
// decomposition. Only emitted under the OLS path
// (`is_ols`); the dispatch-fallback case returns a
// zeroed psi_a (matching the kernel-weights-dropped
// convention for non-OLS learners).
let psi_a : Array[Double] = if is_ols {
let xtwx = weighted_xtx_transpose(xx, w)
let xtwx_aug = add_ridge(xtwx, 1.0e-10)
let xtwx_inv = inv_spd(xtwx_aug)
let p1 = xtwx.ncols
let psi_a_arr : Array[Double] = Array::make(ids.length(), 0.0)
for k = 0; k < ids.length(); k = k + 1 {
let mut dot = 0.0
for j = 0; j < p1; j = j + 1 {
dot = dot + xtwx_inv.data[j] * xx.data[k * p1 + j]
}
psi_a_arr[k] = dot
}
psi_a_arr
} else {
Array::make(ids.length(), 0.0)
}
// v0.96.0+: persist the WLS sandwich ingredients and the hat-
// matrix diagonal here, where `xx` / `w` / `beta` / the ridge
// are still in scope -- none of them is reachable from outside
// `rdd_side`.
//
// `m_k` and `e_k` mirror `LinearRegression::sandwich_se_weighted`
// (linear.mbt) EXPRESSION FOR EXPRESSION, in the same order:
// the same `augment_with_intercept`, the same
// `matmul(xt, xa)` / `add_ridge(..., model.ridge)`, the same
// `solve_spd(..., e_0)` back-solve for the intercept column,
// the same ascending-`a` dot product, and the same
// `matvec` prediction. That is what lets `hac_se` re-form
// `w^2 * m^2 * e^2` with the identical floating-point
// association the library helper uses.
//
// The leverage uses a DIFFERENT matrix on purpose:
// `(X'WX + ridge I)^{-1}` over the augmented design, not the
// unweighted normal inverse. The WLS hat matrix is
// `xa' (X'WX)^{-1} xa` by definition; the `w^2 * (unweighted
// row)^2` meat in `sandwich_se_weighted` is the IRLS-form
// score, not a projection.
let hc0_m : Array[Double] = Array::make(ids.length(), 0.0)
let hc0_e : Array[Double] = Array::make(ids.length(), 0.0)
let leverage : Array[Double] = Array::make(ids.length(), 0.0)
if is_ols {
let xa = augment_with_intercept(xx)
let p1a = xa.cols()
let xt = xa.transpose()
let xtx = matmul(xt, xa)
let xtx_aug = add_ridge(xtx, model.ridge)
let e0 : Array[Double] = Array::make(p1a, 0.0)
e0[0] = 1.0
let mj = solve_spd(xtx_aug, e0)
let pred = matvec(xa, beta)
let m_w = inv_spd(add_ridge(weighted_xtx_transpose(xa, w), model.ridge))
for k = 0; k < ids.length(); k = k + 1 {
let mut mjx = 0.0
for a = 0; a < p1a; a = a + 1 {
mjx = mjx + mj[a] * xa.data[k * p1a + a]
}
hc0_m[k] = mjx
hc0_e[k] = yy[k] - pred[k]
let mut quad = 0.0
for a = 0; a < p1a; a = a + 1 {
for b = 0; b < p1a; b = b + 1 {
quad = quad +
xa.data[k * p1a + a] *
m_w.data[a * p1a + b] *
xa.data[k * p1a + b]
}
}
leverage[k] = w[k] * quad
}
}
let n_side = ids.length().to_double()
let variance = if cov_type == "HC0" && is_ols {
// White sandwich (no `1/n^2` factor — the sandwich diagonal is
// already a variance, not a mean-of-squares). M[0,:] is the
// intercept row of `(X^T W X + ridge I)^{-1}`; we get it
// from `sandwich_se_weighted` which back-solves p1 systems.
let sand = model.sandwich_se_weighted(xx, yy, w)
sand[0]
} else {
// Homoskedastic form: scaled by `(X^T X)^{-1}_{00}` (or
// `(X^T W X)^{-1}_{00}` for the WLS path). The `1 / n^2`
// scaling matches upstream `RDD` and combines with the
// n-side aggregation in the fuzzy delta-method to give
// the correct seed SE.
let inv_diag = model.xtwx_inv_diag()
let inv_00 = if inv_diag.length() > 0 { inv_diag[0] } else { 1.0 }
inv_00 * v / (n_side * n_side)
}
(beta[0], variance, ids.length(), resid, psi_a, hc0_m, hc0_e, leverage)
}
}
///|
pub fn DoubleMLRDD::fit(
self : DoubleMLRDD,
ml_g? : LearnerDispatch = self.ml_g,
) -> DoubleMLRDD {
try {
require(self.bandwidth > 0.0)
let n = self.data.n_obs()
// v0.84.0+: memoize check (mirrors the per-estimator
// pattern). RDD has no fold partition, no n_rep -- the
// entire fit is a deterministic function of
// (data, cutoff, bandwidth, fuzzy, cov_type, ml_g).
// The cache stores `(coef, se, n_local, residuals,
// psi_a)` plus a sentinel `fold_ids` (length `n_obs`
// of zeros) so the FitCache `is_valid(...)` contract
// (which requires `fold_ids.length() == n_obs`) is
// satisfied without RDD actually using a fold partition.
let memoize = self.memoize_enabled
let data_hash : UInt64 = if memoize {
hash_data(self.data.x, self.data.y, self.data.d, z=self.data.score)
} else {
0UL
}
// RDD-specific hparams: estimator_kind + ml_g +
// (cutoff, bandwidth, fuzzy flag, cov_type). The
// standard `hash_hyperparams` only folds in `ml_g` /
// `ml_m` / `propensity_clip`; we layer RDD's structural
// config on top via the cluster-ids slot so the cache
// key invalidates on any bandwidth / cutoff / fuzzy /
// cov_type change.
let hparams_hash : UInt64 = if memoize {
hash_hyperparams("rdd", ml_g, ml_g, 0.0)
} else {
0UL
}
let cluster_hash : UInt64 = if memoize {
let mut h : UInt64 = 0xcbf29ce484222325UL
h = fold_mix_int(h, self.cutoff.to_int())
h = fold_mix_int(h, (self.cutoff * 1.0e6).to_int())
h = fold_mix_int(h, (self.bandwidth * 1.0e6).to_int())
h = fold_mix_int(h, if self.fuzzy { 1 } else { 0 })
// v0.110.0: the kernel joins the key. Stated accurately: this is
// CONTRACT MAINTENANCE, not a fixed defect. The key is documented
// as covering RDD's structural configuration, and leaving the
// kernel out would make that documentation false -- and would
// become a live stale-cache bug the moment any builder could
// change a fitted model's kernel. Today no public API can do
// that (`DoubleMLRDD`'s fields are not visible outside
// `rdd.mbt`, and `new` seeds a fresh empty cache), so no current
// caller is served the wrong kernel. Recorded here so the claim
// is not later upgraded into "this fixed a bug".
h = fold_mix_int(h, self.kernel.tag())
h = fold_mix_int(h, self.cov_type.length())
for c in self.cov_type {
h = fold_mix_int(h, c.to_int())
}
h
} else {
0UL
}
let cache_hit = memoize &&
self.fit_cache.is_valid(
0, 1, 1, n, data_hash, hparams_hash, cluster_hash, "rdd",
)
if cache_hit {
// Cache hit: extract the final `(coef, se,
// n_local, residuals, psi_a, hc0_m, hc0_e,
// leverage)` tuple from the stored predictions and
// return the fitted struct without running any
// `rdd_side` calls.
let preds = self.fit_cache.predictions
let cached_residuals = preds[0]
let cached_psi_a = preds[1]
// Pack the scalars `coef`, `se`, `n_local` into
// a length-3 array on the cache write path;
// cache-hit pulls them back here.
let cached_coef = preds[2][0]
let cached_se = preds[2][1]
let cached_n_local = preds[2][2].to_int()
// v0.96.0: the HC0 ingredients and the WLS leverage ride
// the same cache. They are pure functions of
// (data, cutoff, bandwidth, ml_g) exactly like
// `residuals` / `psi_a`, so there is nothing to recompute
// on a hit -- and `expand_v096_test.mbt` pins that
// `hac_se` still agrees with `se()` after a cache hit.
let cached_hc0_m = preds[3]
let cached_hc0_e = preds[4]
let cached_leverage = preds[5]
{
data: self.data,
cutoff: self.cutoff,
bandwidth: self.bandwidth,
kernel: self.kernel,
fuzzy: self.fuzzy,
cov_type: self.cov_type,
ml_g,
coef: cached_coef,
se: cached_se,
n_local: cached_n_local,
residuals: cached_residuals,
psi_a: cached_psi_a,
hc0_m: cached_hc0_m,
hc0_e: cached_hc0_e,
leverage: cached_leverage,
fitted: true,
boot_t_stat: [],
boot_method: "",
n_rep_boot: 0,
boot_seed: 0,
memoize_enabled: self.memoize_enabled,
fit_cache: self.fit_cache,
}
} else {
let (yl, vyl, nl, res_yl, psi_a_yl, m_yl, e_yl, lev_yl) = rdd_side(
ml_g,
self.data,
-1.0,
self.cutoff,
self.bandwidth,
false,
self.cov_type,
self.kernel,
)
let (yr, vyr, nr, res_yr, psi_a_yr, m_yr, e_yr, lev_yr) = rdd_side(
ml_g,
self.data,
1.0,
self.cutoff,
self.bandwidth,
false,
self.cov_type,
self.kernel,
)
let mut c = yr - yl
let mut variance = vyl + vyr
if self.fuzzy {
let (dl, vdl, _, res_dl, _, _, _, _) = rdd_side(
ml_g,
self.data,
-1.0,
self.cutoff,
self.bandwidth,
true,
self.cov_type,
self.kernel,
)
let (dr, vdr, _, res_dr, _, _, _, _) = rdd_side(
ml_g,
self.data,
1.0,
self.cutoff,
self.bandwidth,
true,
self.cov_type,
self.kernel,
)
let jump = dr - dl
let raw = c
c = raw / jump
// Bug #7 fix: the full delta-method variance for `c = raw / jump`
// is
// var(c) = (var(raw) + c^2 * var(jump) - 2 c cov(raw, jump))
// / jump^2
// The previous implementation omitted the `cov(raw, jump)` cross
// term. We estimate `cov(raw, jump)` from the empirical cross
// moment of the Y-residuals with the D-residuals on each side of
// the cutoff, scaled by `1 / n_side^2` to match the `vyl`/`vdl`
// (already 1/n^2-scaled) convention.
// v0.84.0+: vectorise the per-element
// `cov_num_l = sum(res_yl[k] * res_dl[k])` /
// `cov_num_r = sum(res_yr[k] * res_dr[k])` via
// `vector_multiply` (element-wise product) +
// `mean() * length` (sum-of-products).
let res_yl_dl = vector_multiply(res_yl, res_dl)
let cov_num_l = mean(res_yl_dl) * res_yl.length().to_double()
let res_yr_dr = vector_multiply(res_yr, res_dr)
let cov_num_r = mean(res_yr_dr) * res_yr.length().to_double()
let n_l = res_yl.length().to_double()
let n_r = res_yr.length().to_double()
let mut cov_raw_jump = 0.0
if n_l > 0.0 {
cov_raw_jump = cov_raw_jump + cov_num_l / (n_l * n_l)
}
if n_r > 0.0 {
cov_raw_jump = cov_raw_jump + cov_num_r / (n_r * n_r)
}
variance = (vyl + vyr) / (jump * jump) +
raw * raw * (vdl + vdr) / (jump * jump * jump * jump) -
2.0 * raw * cov_raw_jump / (jump * jump * jump)
}
let se = variance.sqrt()
let n_local_val = nl + nr
// v0.84.0+: vectorise the left+right concat of
// `(res_yl, res_yr)` and `(psi_a_yl, psi_a_yr)` is
// already a flat array concat (the per-element
// loops in `rdd_side` are not part of the public
// API), but we still pass through the original
// concat here for byte-equivalence with the
// pre-vectorise path.
let residuals_full = res_yl + res_yr
let psi_a_full = psi_a_yl + psi_a_yr
// v0.96.0+: same left-then-right concat for the three
// arrays `rdd_side` grew. Only the OUTCOME-side
// (`which = false`) fits are concatenated: those are the
// rows the HC0 / HAC path is defined on, and they line up
// 1:1 with `residuals_full` / `psi_a_full` by
// construction (both are the `nl + nr` outcome-side
// rows). A fuzzy fit's treatment-side (`which = true`)
// ingredients are deliberately NOT kept -- see
// `DoubleMLRDD::hac_se`, which refuses `fuzzy = true`.
let hc0_m_full = m_yl + m_yr
let hc0_e_full = e_yl + e_yr
let leverage_full = lev_yl + lev_yr
// v0.84.0+: when memoize is on and the cache missed,
// write the freshly-computed `(coef, se, n_local,
// residuals, psi_a, hc0_m, hc0_e, leverage)` to the
// cache. The cache layout is
// `[residuals_full, psi_a_full, [coef, se, n_local_as_double],
// hc0_m_full, hc0_e_full, leverage_full]`.
let next_cache = if memoize {
FitCache::from_fit(
Array::make(n, 0),
[
residuals_full,
psi_a_full,
[c, se, n_local_val.to_double()],
hc0_m_full,
hc0_e_full,
leverage_full,
],
0,
1,
1,
n,
data_hash,
hparams_hash,
cluster_hash,
"rdd",
)
} else {
self.fit_cache
}
{
data: self.data,
cutoff: self.cutoff,
bandwidth: self.bandwidth,
kernel: self.kernel,
fuzzy: self.fuzzy,
cov_type: self.cov_type,
ml_g,
coef: c,
se,
// `n_local` is the count of *outcome* observations that fall
// inside the bandwidth on both sides of the cutoff. The fuzzy
// estimator additionally fits a treatment local-polynomial
// on the same rows so the *treated* count is the same.
n_local: n_local_val,
// v0.69.0+: outcome residuals and Riesz-representer
// rows on the bandwidth-restricted sample,
// concatenated (left then right). Drive the
// `irm_style_sensitivity` decomposition in
// `DoubleMLRDD::sensitivity_analysis`.
residuals: residuals_full,
psi_a: psi_a_full,
hc0_m: hc0_m_full,
hc0_e: hc0_e_full,
leverage: leverage_full,
fitted: true,
// v0.75.0+: bootstrap state. `boot_t_stat` etc. start
// empty; `bootstrap(...)` populates them on a return
// copy (RDD is value-style, matching PLR/IRM).
boot_t_stat: [],
boot_method: "",
n_rep_boot: 0,
boot_seed: 0,
memoize_enabled: self.memoize_enabled,
fit_cache: next_cache,
}
}
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// Number of observations.
pub fn DoubleMLRDD::n_obs(self : DoubleMLRDD) -> Int {
self.data.n_obs()
}
///|
/// Number of features (covariate columns).
pub fn DoubleMLRDD::n_features(self : DoubleMLRDD) -> Int {
self.data.x.cols()
}
///|
/// v0.84.0+: turn on memoization. When enabled, the next
/// `fit()` caches the entire `(coef, se, n_local, residuals,
/// psi_a)` tuple under a key that combines the data
/// fingerprint, cutoff, bandwidth, fuzzy flag, cov_type, and
/// learner fingerprint. Subsequent `fit()` calls with the same
/// configuration skip all four `rdd_side` calls (the local-
/// polynomial kernel-weighted fit on each side, twice for sharp +
/// twice for fuzzy).
pub fn DoubleMLRDD::enable_memoize(self : DoubleMLRDD) -> DoubleMLRDD {
{ ..self, memoize_enabled: true, }
}
///|
/// v0.84.0+: turn off memoization. After this, `fit()` will not
/// read or write the cache; the existing `fit_cache` is preserved
/// on the returned struct (call `clear_cache()` to drop it).
pub fn DoubleMLRDD::disable_memoize(self : DoubleMLRDD) -> DoubleMLRDD {
{ ..self, memoize_enabled: false, }
}
///|
/// v0.84.0+: drop the cached `(coef, se, n_local, residuals,
/// psi_a)` tuple. After this, the next `fit()` will run the full
/// local-polynomial pipeline (and repopulate the cache if
/// memoize is still enabled).
pub fn DoubleMLRDD::clear_cache(self : DoubleMLRDD) -> DoubleMLRDD {
{ ..self, fit_cache: FitCache::empty(), }
}
///|
/// v0.84.0+: `true` iff `fit_cache` holds at least one cached
/// fit result (i.e. a previous `fit()` with `memoize_enabled =
/// true` has populated the cache). Note that `has_cache()` does
/// NOT verify the cache key matches the current data + learner
/// configuration -- check `memoize_enabled` if you also need
/// to know whether the next `fit()` will hit.
pub fn DoubleMLRDD::has_cache(self : DoubleMLRDD) -> Bool {
!self.fit_cache.is_empty()
}
///|
pub fn DoubleMLRDD::coef(self : DoubleMLRDD) -> Double {
try {
require(self.fitted)
self.coef
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
pub fn DoubleMLRDD::se(self : DoubleMLRDD) -> Double {
try {
require(self.fitted)
self.se
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
pub fn DoubleMLRDD::confint(self : DoubleMLRDD) -> (Double, Double) {
try {
require(self.fitted)
(self.coef - 1.96 * self.se, self.coef + 1.96 * self.se)
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
// ---------------------------------------------------------------------------
// v0.96.0: `hac_se` / `cluster_hac_se` -- the local-regression HAC sandwich
// ---------------------------------------------------------------------------
//
// WHY RDD DOES NOT REUSE `sandwich_se` / `sandwich_variance`
// =========================================================
//
// The shared DML helper (`sandwich.mbt`) and RDD's own HC0 branch
// (`rdd_side`, `cov_type == "HC0"`) are both called "Huber-White",
// and that shared prefix is exactly the trap. They are not the
// same algebra, and routing RDD through the shared helper produces
// a number that is plausible, finite, and NOT RDD's HC0.
//
// The two meats, side by side. RDD's (`linear.mbt`,
// `LinearRegression::sandwich_se_weighted`, summed at `rdd_side`):
//
// sum_k w_k^2 * m_k^2 * e_k^2,
// m_k = (intercept column of (X'X + ridge I)^-1) . xa_k
//
// -- note the `w_k^2`, the IRLS-form weight in the score (White
// 1980 sec 4), and NO `1 / n^2`. The shared helper
// (`sandwich_variance_hc0`, since v0.91.0):
//
// M_inv[0,0]^2 * sum_i psi[i]^2 / n / n
//
// Three independent mismatches, any one of which is fatal:
//
// 1. THE `psi_a` FACTOR. Through v0.90.0 the shared accumulator
// was `sum_i (psi_a[i] * psi[i])^2` (see the file header in
// `sandwich.mbt`). RDD's combined influence function is
// `psi[k] = psi_a[k] * residuals[k]` -- the `psi_a` slot IS
// `psi_a` and the `psi` slot IS the combined IF, which is
// exactly what `bootstrap(...)` feeds to
// `did_bootstrap_t_stat`. So the substitution is not
// `sum (psi_a e)^2`; it is
//
// sum_k ( psi_a[k] * (psi_a[k] * e_k) )^2
// = sum_k ( psi_a[k]^2 * e_k )^2
// = sum_k psi_a[k]^4 * e_k^2,
//
// a FOURTH power. RDD's HC0 meat is a SECOND power:
// `w_k^2 * psi_a[k]^2 * e_k^2` (since `m_k` for the WEIGHTED
// normal inverse is `psi_a[k]` -- see the struct-field docs).
// The two differ by a factor of `psi_a[k]^2` per row, which is
// data-dependent and order `1e-2` on the v0.96.0 DGP, so the
// wrong answer passes every plausibility check. v0.91.0
// removed the `psi_a` factor precisely because it was not
// part of the moment; the substitution is still wrong, for
// the next reason.
//
// 2. THE SCALAR JACOBIAN. `sandwich_variance_hc0` takes a 1x1
// `M_inv` and squares it (`M_inv[0,0]^2 * ...`). RDD has no
// DML moment and therefore no scalar Jacobian: its `M` is the
// full `p1 x p1` WLS normal inverse, and the intercept
// variance picks out its `[0, :]` ROW, not a scalar. Forcing
// a 1x1 in means either passing `[[1.0]]` (dropping the whole
// `(X'WX)^-1` conditioning) or inventing a `1 / mean(psi_a)`
// that has no meaning here -- `mean(psi_a)` for a WLS
// intercept row is not a derivative of anything.
//
// 3. THE `1 / n^2`. The shared helper is a MEAN moment, so it
// divides the meat by `n` twice; RDD's variance is a sum over
// the bandwidth-restricted rows of a sum over two SEPARATE
// regressions, and carries no such divisor. Dividing by
// `n_local^2 = 360^2` would shrink the SE by 360x, and the
// HC0 anchor `hac_se(HC0) == se()` would fail by exactly that
// factor -- loudly, which is the only reason to prefer the
// mistake that is quiet.
//
// The HC2 / HC3 leverage has the same split-brain history: the
// shared helper uses the CONSTANT mean-regression leverage
// `h_ii = 1 / n` (so `HC2 == HC1` there, see `sandwich.mbt`),
// while RDD's is a genuine WLS hat diagonal
// `h_k = w_k * xa_k' (X'WX + ridge I)^-1 xa_k` that varies per
// local row. RDD's HC2 / HC3 are therefore NOT HC1, and the
// "HC2 == HC1 exactly" identity documented for the DML family does
// not carry over.
//
// CONSEQUENCE: `DoubleMLRDD` exposes `hac_se`, NOT `sandwich_se`.
// RDD is not a DML-score estimator -- it has no `psi_b`, no
// `E[theta psi_a + psi_b]` moment, no `M_inv` -- so a method
// literally named `sandwich_se` would promise a DML sandwich
// contract it cannot honour. The name says what the number is: the
// HAC sandwich of a kernel-weighted local regression.
//
// SCOPE OF THE ANCHOR
// -------------------
// `hac_se(HC0) == se()` BIT-IDENTICALLY holds exactly when the fit
// took the HC0 branch, i.e. `cov_type = "HC0"` AND the injected
// learner is `LinearRegression`:
//
// - `cov_type = "homoskedastic"` fit: `se()` is the
// homoskedastic form, so `hac_se(HC0) != se()`. Pinned.
// - non-OLS learner: `rdd_side` drops the kernel weights and
// no HC0 state exists. `hac_se` ABORTS rather than returning
// the homoskedastic number under an HC0 name.
// - `fuzzy = true`: `se()` is the delta-method variance of the
// RATIO `c = raw / jump`, mixing the Y-side and D-side local
// fits plus a cross-moment that is `1 / n_side^2`-scaled (not
// HC0-scaled). No HC0/HC1/HC2/HC3 correction reproduces that
// from the outcome-side rows alone, and making it reproduce
// would require rewriting `fit`'s fuzzy variance, which
// v0.96.0 does not do. `hac_se` ABORTS.
// - not fitted: ABORTS via `PreconditionError`.
//
// The four HC variants, for a SHARP fit with `n_left` / `n_right`
// local rows and `k = n_local_params` parameters per side:
//
// HC0: sum_k w_k^2 m_k^2 e_k^2 (left + right)
// HC1: the same, with each side's sum scaled by
// n_side / (n_side - k). The correction is PER SIDE,
// because each side is a separate WLS regression (the
// Stata `, robust` / `statsmodels cov_type='HC1'`
// convention). On a symmetric design this coincides
// with the pooled `n_local / (n_local - 2k)`, since
// `2 n_s / (2 n_s - 2k) == n_s / (n_s - k)`.
// HC2: each term divided by `(1 - h_k)`
// HC3: each term divided by `(1 - h_k)^2`
//
// `k` READ FROM THE CODE, not assumed: `rdd_design` builds a local
// design with `data.x.cols() + 1` columns (`u` plus the
// covariates) and `fit_weighted` adds an intercept, so the local
// regression estimates `p1 = data.x.cols() + 2` coefficients --
// see `n_local_params`. The point estimate only READS
// `beta[0] + beta[1] * u` (`rdd_side`), but dropping a column
// would change the normal equations and hence `beta[0]`, so all
// `p1` count. "2 (intercept + running variable)" is right only for
// a zero-covariate design.
///|
/// v0.96.0+: the number of coefficients the per-side local
/// regression estimates, `data.x.cols() + 2`: the intercept that
/// `fit_weighted` augments on, the running variable `u` that
/// `rdd_design` puts in column 0, and the `data.x.cols()`
/// covariates. This is the `k` in the HC1 finite-sample
/// correction. The point estimate reads only `beta[0] + beta[1] *
/// u`, but every column enters the WLS normal equations, so the
/// fitted parameter count is `p1`, not 2.
pub fn DoubleMLRDD::n_local_params(self : DoubleMLRDD) -> Int {
self.data.x.cols() + 2
}
///|
/// v0.96.0+: the WLS hat-matrix diagonal of the per-side local
/// polynomial, `h_k = w_k * xa_k' * (X'WX + ridge I)^{-1} * xa_k`
/// with `xa_k = (1, u_k, x_k)`, on the bandwidth-restricted
/// sample (left side then right side, length `n_local()`).
///
/// Every entry is in `[0, 1)`: `X^a (X'WX + ridge I)^{-1} X^{a'}`
/// is an idempotent matrix (a projection) plus a `ridge`-order
/// perturbation, so its diagonal is a leverage. The entries sum
/// to the RANK of the local design -- the trace of a projection
/// matrix is its rank -- which is `n_local_params()` per side for
/// a full-rank design and LESS for a degenerate one (a constant
/// covariate column makes `(1, u, 0)` rank-deficient), up to the
/// `ridge` term. `expand_v096_test.mbt` pins both: the
/// full-rank trace, and the `[0, 1)` bound.
///
/// Zero-filled (all `0.0`) when the injected learner is not
/// `LinearRegression`, matching the `psi_a` convention for that
/// case.
pub fn DoubleMLRDD::leverage(self : DoubleMLRDD) -> Array[Double] {
try {
require(self.fitted)
self.leverage
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// Guards shared by `hac_se` and `cluster_hac_se`. Both need the
/// WLS sandwich state, which `rdd_side` only computes on the OLS
/// path; and neither applies to a fuzzy fit. Both refusals use
/// `abort` with a message rather than a bare `require` because
/// the failure is a wrong-request error, not a malformed
/// argument.
fn DoubleMLRDD::require_hac_preconditions(self : DoubleMLRDD) -> Unit {
let is_ols = match self.ml_g {
LinearRegression(_) => true
_ => false
}
if !is_ols {
abort(
"DoubleMLRDD::hac_se: ml_g is not LinearRegression, so rdd_side dropped the kernel weights and computed no HC0 state (hc0_m / hc0_e / leverage are all zero). Refusing rather than reporting a homoskedastic number under an HC0 name.",
)
}
if self.fuzzy {
abort(
"DoubleMLRDD::hac_se: fuzzy = true. se() on a fuzzy fit is the delta-method variance of the RATIO c = raw / jump, which mixes the Y-side and D-side local fits and a cross-moment that is 1 / n_side^2-scaled rather than HC0-scaled. No HC0/HC1/HC2/HC3 correction on the outcome-side rows reproduces it; use a sharp fit, or bootstrap the fuzzy estimator.",
)
}
}
///|
/// The per-observation HC0 sandwich term
/// `w^2 * m^2 * e^2`, written in exactly the association
/// `LinearRegression::sandwich_se_weighted` uses
/// (`acc = acc + wi2 * mjx * mjx * ei * ei`, all left-associative)
/// so that re-forming it from the persisted `hc0_m` / `hc0_e`
/// ingredients reproduces `se()` bit-for-bit rather than to
/// tolerance.
fn rdd_hc0_term(w : Double, m : Double, e : Double) -> Double {
let wi2 = w * w
wi2 * m * m * e * e
}
///|
/// Degrees-of-freedom guard for the HC1 / HC2 / HC3 variants: both
/// sides of the local regression must have MORE rows than the
/// design has parameters, or the fit is rank-deficient.
///
/// This is not a formality. A side with `n_side <= p1` rows has a
/// rank-deficient local design, and the `1e-10` ridge then reports
/// `h_k = 1 - 1e-10` rather than the exact `1`: the leverage is
/// inside `[0, 1)`, so a pure "is `1 - h_k` positive" check does
/// NOT fire, and HC2 would then divide the meat by `1e-10` and
/// report an SE inflated by `1 / sqrt(1e-10) = 1e5`. Measured on
/// the v0.96.0 degenerate fixture (1 local row per side, `p1 = 3`):
/// `h_0 = 0.999999999825377`, `1 - h_0 = 1.7e-10`. So the df
/// check has to come FIRST, ahead of the leverage-range check, and
/// the range check alone is not sufficient.
///
/// HC0 is deliberately NOT guarded: on such a fit `se()` itself
/// still reports a number (`fit` does not refuse), and
/// `hac_se(HC0) == se()` is the anchor, so refusing HC0 here would
/// break the one identity this API promises.
fn rdd_hac_require_df(
kind_name : String,
n_left : Int,
n_right : Int,
p1 : Int,
) -> Unit {
if n_left <= p1 {
abort(
"DoubleMLRDD::hac_se(" +
kind_name +
"): the LEFT local regression has " +
n_left.to_string() +
" rows but the local design has " +
p1.to_string() +
" parameters, so it is rank-deficient. HC1 needs n - k > 0, and HC2 / HC3 need an identified leverage: a rank-deficient design reports h_k = 1 - ridge, so the (1 - h_k) divisor would be a ridge artifact and inflate the variance by ~1 / ridge. Widen the bandwidth.",
)
}
if n_right <= p1 {
abort(
"DoubleMLRDD::hac_se(" +
kind_name +
"): the RIGHT local regression has " +
n_right.to_string() +
" rows but the local design has " +
p1.to_string() +
" parameters, so it is rank-deficient. HC1 needs n - k > 0, and HC2 / HC3 need an identified leverage: a rank-deficient design reports h_k = 1 - ridge, so the (1 - h_k) divisor would be a ridge artifact and inflate the variance by ~1 / ridge. Widen the bandwidth.",
)
}
}
///|
/// v0.96.0+: heteroskedasticity-consistent (Huber-White) standard
/// error for the local-regression discontinuity contrast, in the
/// HC0 / HC1 / HC2 / HC3 family.
///
/// This is NOT `sandwich_se`. Read the section comment above
/// ("WHY RDD DOES NOT REUSE `sandwich_se`") for the algebra: RDD's
/// meat is `sum_k w_k^2 * m_k^2 * e_k^2` with a full `p1 x p1` WLS
/// normal inverse and no `1 / n^2`, the shared helper is a
/// scalar-Jacobian MEAN-moment sandwich, and feeding RDD's
/// influence function into the pre-v0.91.0 accumulator yields
/// `sum_k psi_a[k]^4 * e_k^2` (a fourth power) against RDD's second
/// power. The `SandwichKind` enum is reused -- it is the package's
/// shared vocabulary for the HC family -- but the arithmetic
/// behind each variant is the local-regression one.
///
/// The invariant that proves the wiring:
/// `hac_se(HC0) == se()` BIT-IDENTICALLY on a `cov_type = "HC0"`
/// fit. It is NOT claimed on a `cov_type = "homoskedastic"` fit
/// (whose `se()` is the homoskedastic form), and it ABORTS on a
/// non-OLS learner and on `fuzzy = true`.
///
/// HC2 / HC3 divide by `1 - h_k` per local row using the
/// persisted WLS leverage (`leverage()`). If any `h_k` leaves
/// `[0, 1)` the divisor would be non-positive, which means the
/// local regression is rank-deficient (fewer local rows than
/// parameters, or a collinear local design). This ABORTS rather
/// than clipping: a clip is what silently broke identities in
/// v0.93 / v0.94, and for a leverage in `[0, 1)` no correction
/// should ever fire.
pub fn DoubleMLRDD::hac_se(self : DoubleMLRDD, kind : SandwichKind) -> Double {
try {
require(self.fitted)
self.require_hac_preconditions()
let n = self.n_local
require(self.hc0_m.length() == n)
require(self.hc0_e.length() == n)
require(self.leverage.length() == n)
let w = rdd_kernel_weights(
self.data,
self.cutoff,
self.bandwidth,
self.kernel,
)
require(w.length() == n)
let n_left = rdd_local_side_count(
self.data,
self.cutoff,
self.bandwidth,
-1.0,
)
let p1 = self.n_local_params()
// Degrees-of-freedom guard FIRST: a rank-deficient local
// design reports a leverage of `1 - ridge`, which is inside
// `[0, 1)` and would sail past the range check below while
// making the HC2 / HC3 divisor a ridge artifact. HC0 is
// exempt so the `hac_se(HC0) == se()` anchor survives even on
// a degenerate fit.
let needs_df = match kind {
HC0 => false
HC1 => true
HC2 => true
HC3 => true
}
if needs_df {
rdd_hac_require_df(kind.to_string(), n_left, n - n_left, p1)
}
// Per-observation divisor: 1 for HC0 / HC1, `1 - h_k` for
// HC2, `(1 - h_k)^2` for HC3.
let div : Array[Double] = Array::make(n, 1.0)
let uses_leverage = match kind {
HC0 | HC1 => false
HC2 | HC3 => true
}
if uses_leverage {
for k = 0; k < n; k = k + 1 {
let h = self.leverage[k]
if !(h >= 0.0 && h < 1.0) {
abort(
"DoubleMLRDD::hac_se: leverage h_" +
k.to_string() +
" = " +
h.to_string() +
" is outside [0, 1). A WLS hat-matrix diagonal always lies in [0, 1) unless the local regression is rank-deficient (fewer local rows than parameters, or a collinear local design), so 1 - h_k would be non-positive. Aborting instead of clipping.",
)
}
let one_minus_h = 1.0 - h
div[k] = match kind {
HC3 => one_minus_h * one_minus_h
_ => one_minus_h
}
}
}
// TWO partial sums, one per side, added last. This is the same
// shape as `fit`'s `variance = vyl + vyr`, and each partial sum
// accumulates the same per-observation values in the same
// ascending order that `sandwich_se_weighted` accumulated
// them, which is what makes the HC0 anchor bit-identical
// rather than merely close.
let mut acc_left = 0.0
let mut acc_right = 0.0
for k = 0; k < n; k = k + 1 {
let base = rdd_hc0_term(w[k], self.hc0_m[k], self.hc0_e[k])
// Skip the division entirely for HC0 / HC1: `x / 1.0` is
// exact in IEEE-754, but not dividing at all cannot drift.
let term = if div[k] == 1.0 { base } else { base / div[k] }
if k < n_left {
acc_left = acc_left + term
} else {
acc_right = acc_right + term
}
}
let variance = match kind {
HC0 => acc_left + acc_right
HC2 => acc_left + acc_right
HC3 => acc_left + acc_right
HC1 => {
// Per-side Stata `, robust` correction: each side is a
// separate WLS regression, so each gets its own
// `n_side / (n_side - k)`. `rdd_hac_require_df` has
// already guaranteed both denominators are positive.
let k_d = p1.to_double()
let n_l = n_left.to_double()
let n_r = (n - n_left).to_double()
acc_left * (n_l / (n_l - k_d)) + acc_right * (n_r / (n_r - k_d))
}
}
require(variance >= 0.0)
variance.sqrt()
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// v0.96.0+: cluster-robust (Arellano 1987, Cameron-Gelbach-Miller
/// 2011) standard error for the local-regression discontinuity
/// contrast -- the clustered analogue of `hac_se(HC0)`.
///
/// `cluster_ids[k]` is the 0-based cluster of the k-th
/// bandwidth-restricted row (left side first then right side, the
/// same ordering as `leverage()`), and
/// `cluster_ids.length()` must equal `n_local()`. The number of
/// clusters is `1 + max(cluster_ids)`.
///
/// Per-observation score `s_k = w_k * m_k * e_k` (the RDD HC0
/// term under a square root, with its sign), aggregated within
/// cluster:
///
/// S_c = sum_{k in c} s_k
/// var = sum_c S_c^2 * n_c / (n_c - 1)
///
/// with the `(n_c - 1)` jackknife correction clipped to 1 for
/// single-observation clusters (the `statsmodels
/// cov_type='cluster'` convention, and the same clip
/// `cluster_sandwich_variance` uses). There is NO `1 / n^2`
/// divisor: the shared helper needs one because its moment is a
/// MEAN moment, while this is a sum of squared scores over the
/// rows of two separate regressions, exactly like `hac_se(HC0)`.
///
/// With all-singleton clusters the meat collapses to
/// `sum_k s_k^2`, i.e. `hac_se(HC0)` -- to floating-point
/// tolerance, not bit-exactly, for two reasons: the per-observation
/// square is formed as `s_k * s_k` here against the
/// `w^2 m^2 e^2` association in `hac_se`, and this accumulator is
/// Kahan-compensated (as in `cluster_sandwich_variance`) against
/// `hac_se`'s plain running sum. Pooled clusters differ.
///
/// Same preconditions as `hac_se`: ABORTS on a non-OLS learner
/// and on `fuzzy = true`.
pub fn DoubleMLRDD::cluster_hac_se(
self : DoubleMLRDD,
cluster_ids : Array[Int],
) -> Double {
try {
require(self.fitted)
self.require_hac_preconditions()
let n = self.n_local
require(self.hc0_m.length() == n)
require(self.hc0_e.length() == n)
require(cluster_ids.length() == n)
let w = rdd_kernel_weights(
self.data,
self.cutoff,
self.bandwidth,
self.kernel,
)
require(w.length() == n)
let mut max_cid = -1
for i = 0; i < n; i = i + 1 {
require(cluster_ids[i] >= 0)
if cluster_ids[i] > max_cid {
max_cid = cluster_ids[i]
}
}
require(max_cid >= 0)
let n_clusters = max_cid + 1
let cluster_sum : Array[Double] = Array::make(n_clusters, 0.0)
let cluster_size : Array[Int] = Array::make(n_clusters, 0)
for k = 0; k < n; k = k + 1 {
let c = cluster_ids[k]
let s = w[k] * self.hc0_m[k] * self.hc0_e[k]
cluster_sum[c] = cluster_sum[c] + s
cluster_size[c] = cluster_size[c] + 1
}
// Kahan-compensated, matching `cluster_sandwich_variance`.
let mut acc = 0.0
let mut acc_c = 0.0
for c = 0; c < n_clusters; c = c + 1 {
let nc = cluster_size[c]
if nc == 0 {
continue
}
let scale = if nc <= 1 {
1.0
} else {
nc.to_double() / (nc.to_double() - 1.0)
}
let sc = cluster_sum[c]
let term = sc * sc * scale
let y = term - acc_c
let t = acc + y
acc_c = t - acc - y
acc = t
}
require(acc >= 0.0)
acc.sqrt()
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// v0.69.0+: sensitivity analysis for `DoubleMLRDD`.
///
/// RDD is the kernel-weighted local-polynomial regression
/// discontinuity estimator. The kernel-weighted precision
/// matrix per side is `M = (X^T W X + ridge I)^{-1}` and
/// the Riesz-representer row is `psi_a[k] = M[0, :] @ xx[k, :]`
/// (intercept row of `M` dotted with the k-th local
/// design row). The local polynomial residual is
/// `e_k = target_k - pred_k`. The IRM-style decomposition
/// (matching the v0.66.0+ sensitivity family) treats the
/// concatenation `residuals ++ psi_a` (left side then
/// right side) as a single IF row and applies
/// `irm_style_sensitivity` to it. The result is a single
/// `SensitivityResult` for the kernel-RDD jump estimate.
///
/// Notes:
/// - Fuzzy RDD uses `c = raw / jump` delta-method
/// weighting that mixes Y-residuals and D-residuals;
/// the simple `irm_style_sensitivity` decomposition
/// assumes the single-theta score form (so the
/// fuzzy-delta cross term `cov(raw, jump)` is captured
/// by the v0.59.0- fuzzy variance fix, not by the
/// sensitivity path). For sharp RDD the
/// decomposition is exact.
/// - `cf_y` / `cf_d` default to `0.05` matching the
/// v0.66.0+ sensitivity family.
///
/// Calling on an un-fit model aborts via
/// `PreconditionError`. Returns a single
/// `SensitivityResult`.
pub fn DoubleMLRDD::sensitivity_analysis(
self : DoubleMLRDD,
cf_y? : Double = 0.05,
cf_d? : Double = 0.05,
) -> SensitivityResult raise {
require(self.fitted)
irm_style_sensitivity(self.coef, self.residuals, self.psi_a, cf_y, cf_d)
}
///|
/// v0.78.0+: cluster-robust analogue of
/// `DoubleMLRDD::sensitivity_analysis`. RDD's kernel-weighted
/// Riesz-representer row and outcome residuals (already on the
/// bandwidth-restricted sample, length `n_local`) are routed
/// through the v0.78.0 `irm_style_sensitivity_cluster`
/// kernel-weighted extension with the per-observation triangular
/// kernel weights `w[k] = 1 - |u_k| / h` where
/// `u_k = score[ids[k]] - cutoff`. The cluster sums become
/// `cluster_sum_resid[c] = sum_{k in c} w[k] * residuals[k]`
/// (and similarly for `psi_a`), so observations near the
/// cutoff (high weight) drive the cluster-aggregated variance
/// baseline. The per-observation C&H centering
/// `residuals[k]^2 - sigma2_cluster` keeps the bias expression
/// per-observation-anchored (see the helper's v0.78.0 doc).
///
/// `DoubleMLRDDData` has no `cluster_vars` field, so the
/// user must pass `cluster_ids` explicitly.
/// `cluster_ids.length()` must equal `self.n_local()`. Cluster
/// indices are 0-based; `1 + max(cluster_ids)` is the number of
/// clusters. Returns the same `SensitivityResult` shape as
/// the IID `sensitivity_analysis` so callers can swap IID for
/// cluster-aware without changing the return contract.
///
/// Calling on an un-fit model aborts via `PreconditionError`.
/// `n_local == 0` (no observations fall inside the bandwidth)
/// raises on the inner `irm_style_sensitivity_cluster`'s
/// `require(n > 0)`.
pub fn DoubleMLRDD::sensitivity_analysis_cluster(
self : DoubleMLRDD,
cluster_ids : Array[Int],
cf_y? : Double = 0.05,
cf_d? : Double = 0.05,
) -> SensitivityResult raise {
require(self.fitted)
let n_local = self.n_local
require(n_local > 0)
require(cluster_ids.length() == n_local)
// Kernel weights: triangular `w[k] = 1 - |u_k| / h` on the
// bandwidth-restricted sample, left side first then right
// side (matching the ordering of `self.residuals` /
// `self.psi_a`). Reproduces `rdd_design`'s selection logic
// so the weights line up with the per-local-row residuals
// and Riesz rows.
let kernel_weights : Array[Double] = rdd_kernel_weights(
self.data,
self.cutoff,
self.bandwidth,
self.kernel,
)
require(kernel_weights.length() == n_local)
irm_style_sensitivity_cluster(
self.coef,
self.residuals,
self.psi_a,
cluster_ids,
cf_y,
cf_d,
kernel_weights~,
)
}
///|
/// Compute the per-observation kernel weights for `DoubleMLRDD` on the
/// bandwidth-restricted sample: returns `kernel.weight(u_k, h)` for each
/// local-row observation k, where `u_k = score[ids[k]] - cutoff` and
/// `ids` is the bandwidth-restricted selection (left side first then
/// right side, matching the ordering used by `residuals` / `psi_a`).
/// Observations outside the bandwidth are skipped (so the returned
/// array length equals `n_local`). The helper is a pure function of the
/// data + cutoff + h + kernel so the caller does not need to thread the
/// original ids mapping through the `fit()` result.
///
/// v0.110.0: took the kernel from the caller instead of computing
/// `1 - |u|/h` inline. Before that this path silently used the
/// TRIANGULAR kernel no matter what the estimator was configured with --
/// so a kernel change would have moved `coef` and left
/// `cluster_hac_se` / `sensitivity_analysis_cluster` on the old weights.
/// Two code paths, one kernel setting.
fn rdd_kernel_weights(
data : DoubleMLRDDData,
cutoff : Double,
h : Double,
kernel : RDDKernel,
) -> Array[Double] {
let weights : Array[Double] = []
// Left side: u < 0.
for i = 0; i < data.n_obs(); i = i + 1 {
let u = data.score[i] - cutoff
if u < 0.0 && u.abs() <= h {
weights.push(kernel.weight(u, h))
}
}
// Right side: u >= 0.
for i = 0; i < data.n_obs(); i = i + 1 {
let u = data.score[i] - cutoff
if u >= 0.0 && u.abs() <= h {
weights.push(kernel.weight(u, h))
}
}
weights
}
///|
/// Number of bandwidth-restricted rows on ONE side of the cutoff
/// (`side < 0.0` left / `side > 0.0` right). Reproduces
/// `rdd_design`'s selection predicate exactly -- the same
/// `u < 0.0` / `u >= 0.0` split and the same `u.abs() <= h`
/// bandwidth test -- so the left/right split `hac_se` uses to
/// mirror `fit`'s `variance = vyl + vyr` lines up with the rows
/// that produced `vyl` and `vyr`.
///
/// `rdd_kernel_weights` above walks the same selection to build
/// the weights; this one only counts. It is a pure function of
/// `(data, cutoff, bandwidth)`, so the memoize cache-hit path does
/// not need to store the split.
fn rdd_local_side_count(
data : DoubleMLRDDData,
cutoff : Double,
h : Double,
side : Double,
) -> Int {
let mut count = 0
for i = 0; i < data.n_obs(); i = i + 1 {
let u = data.score[i] - cutoff
if (side < 0.0 && u < 0.0) || (side > 0.0 && u >= 0.0) {
if u.abs() <= h {
count = count + 1
}
}
}
count
}
///|
/// v0.75.0+: multiplier bootstrap for `DoubleMLRDD`. The
/// per-observation influence function is
///
/// psi[k] = psi_a[k] * residuals[k]
///
/// (kernel-weighted Riesz row dotted with the local-polynomial
/// residual). RDD is the kernel-weighted local-polynomial
/// regression discontinuity estimator: the IF for the jump
/// coefficient at the cutoff does not factor into the v0.61.0
/// `psi_a + coef * psi_b` form (it's a single combined IF),
/// so this bootstrap calls `did_bootstrap_t_stat` directly
/// rather than going through `generic_bootstrap_t_stat`.
///
/// `method_name` selects the multiplier distribution:
/// `"normal"` (default), `"Bayes"`, `"wild"`; `seed` defaults
/// to `2024`; `n_rep_boot` defaults to `500`.
///
/// Calling on an un-fit model aborts via `PreconditionError`.
/// Returns a new `DoubleMLRDD` with `boot_t_stat` /
/// `boot_method` / `n_rep_boot` / `boot_seed` populated.
pub fn DoubleMLRDD::bootstrap(
self : DoubleMLRDD,
method_name? : String = "normal",
n_rep_boot? : Int = 500,
seed? : Int = 2024,
) -> DoubleMLRDD {
try {
require(self.fitted)
require(
method_name == "normal" || method_name == "Bayes" || method_name == "wild",
)
require(n_rep_boot >= 2)
let n = self.n_local
require(self.psi_a.length() == n)
require(self.residuals.length() == n)
// Draw weights. Shape: (n_rep_boot, n_local).
let weights = draw_bootstrap_weights(method_name, n_rep_boot, n, seed) catch {
BootstrapMethodError::UnknownMethod(m) =>
abort(
"draw_bootstrap_weights: unknown method (set in DoubleMLRDD::bootstrap): " +
m,
)
}
// Combined IF `psi[k] = psi_a[k] * residuals[k]`.
let psi : Array[Double] = Array::make(n, 0.0)
let mut ss_psi = 0.0
for i = 0; i < n; i = i + 1 {
let psi_i = self.psi_a[i] * self.residuals[i]
psi[i] = psi_i
ss_psi = ss_psi + psi_i * psi_i
}
let n_d = n.to_double()
let se_psi = (ss_psi / n_d).sqrt()
let boot_t_stat : Array[Double] = if se_psi <= 0.0 {
// Degenerate: psi sums to 0 (kernel row or residuals
// collapse to 0). Return zeros (matches the v0.55.0+
// DIDCrossSection convention).
Array::make(n_rep_boot, 0.0)
} else {
let se_flat : Array[Double] = [se_psi]
did_bootstrap_t_stat(weights, psi, se_flat, n_rep_boot, n, 1)
}
{
..self,
boot_t_stat,
boot_method: method_name,
n_rep_boot,
boot_seed: seed,
}
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// Accessor for the local-polynomial learner used by the most
/// recent `fit(...)` call. v0.60.0+. Note: RDD's bandwidth /
/// kernel are configuration (not LearnerDispatch concepts);
/// this slot plugs into the closed-form OLS underneath the
/// kernel weights and is forward-compatible with v0.61.0+
/// kernel-aware wrappers.
pub fn DoubleMLRDD::learner_g(self : DoubleMLRDD) -> LearnerDispatch {
self.ml_g
}
///|
pub fn DoubleMLRDD::n_local(self : DoubleMLRDD) -> Int {
self.n_local
}
///|
/// v0.108.0+ (NEW): score a grid of candidates for RDD's ONE
/// nuisance slot and re-fit `DoubleMLRDD` with the winner, via the
/// shared `tune_score_grid` core. This closes a coverage gap against
/// upstream, where `tune` reaches every estimator through the
/// `BaseDML` mixin; `PQ` / `QTE` / `RDD` had no `tune` before
/// v0.108.0.
///
/// Slot tuned: `ml_g`, and ONLY `ml_g`. RDD has a single nuisance
/// slot, so `param_set` is a bare `Array[LearnerDispatch]` -- the
/// same shape `DoubleMLPLIV::tune` uses; a `TuneParam` would be a
/// two-slot wrapper around one value.
///
/// DO NOT "fix" `ml_g` into a propensity / treatment slot. Despite
/// the `g` in the name (it is inherited from the field name, not
/// from any role), RDD's `ml_g` is the DISCONTINUITY-SIDE,
/// OUTCOME-side learner for `E[y|x]` -- the local polynomial fit on
/// each side of the cutoff under the triangular kernel (see the
/// `ml_g` struct-field doc and `rdd_side`). It is NOT `E[g|d,x]`.
/// The `TuneParam` slot-semantics table in `tune.mbt` lists `RDD`'s
/// first slot as the discontinuity-side learner `E[y|x]` for
/// exactly this reason.
///
/// Structural config is NOT re-derived here: `cutoff`, `bandwidth`,
/// `fuzzy` and `cov_type` are already carried by the model being
/// tuned, and `DoubleMLRDD::fit` reads them off `self` (it copies
/// `self.cutoff` / `self.bandwidth` / `self.fuzzy` /
/// `self.cov_type` straight onto the returned struct) and threads
/// them into every `rdd_side` call. `tune` passes only `ml_g`, so
/// tuning the learner cannot perturb the design.
///
/// Cluster guard: ABSENT, deliberately. `DoubleMLRDDData` has NO
/// `is_cluster_data()` method -- RDD's data type is not the
/// cluster-aware `DoubleMLData`, and adding the method (or the
/// guard) would mean inventing a cluster concept for a data type
/// that never had one. RDD's dependence is handled through the
/// bandwidth-restricted local sample, not cluster ids. Cluster
/// inference lives in the separate explicit
/// `DoubleMLRDD::cluster_hac_se` / `sensitivity_analysis_cluster`
/// entry points.
pub fn DoubleMLRDD::tune(
self : DoubleMLRDD,
param_set~ : Array[LearnerDispatch],
scoring_method? : String = "MSE",
n_folds_tune? : Int = 5,
seed? : Int = 3141,
) -> DoubleMLRDD {
try {
require(param_set.length() > 0)
require(n_folds_tune >= 2)
let scoring = TuneScoring::parse(scoring_method)
let grid = tune_score_grid(
self.data.x,
self.data.y,
param_set,
n_folds_tune,
seed,
scoring,
)
self.fit(ml_g=param_set[grid.best_index])
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// The kernel weighting the local polynomial. Read back with this rather
/// than inferred from a fitted model, so a config that was built with an
/// explicit kernel can be checked against what the model actually used.
pub fn DoubleMLRDD::kernel(self : DoubleMLRDD) -> RDDKernel {
self.kernel
}
///|
/// The per-observation kernel weights this model applies, on the
/// bandwidth-restricted local sample: one entry per local row, left
/// side first then right side, matching the ordering of `residuals` /
/// `psi_a` / `hc0_m` / `hc0_e`.
///
/// Public, and not only for curiosity. `cluster_hac_se` and
/// `sensitivity_analysis_cluster` REBUILD these weights rather than
/// reading a stored copy, so they had a second, independent path to the
/// kernel. A test that only compares two `cluster_hac_se` values cannot
/// tell which path supplied the weights -- both models' `fit` already
/// stored kernel-dependent `hc0_m` / `hc0_e`, so the SE differs either
/// way and the assertion holds even with the cluster path reverted to a
/// hardcoded triangular kernel (mutation M2, which survived the first
/// version of this gate). Reading the weights makes the second path
/// directly observable.
pub fn DoubleMLRDD::kernel_weights(
self : DoubleMLRDD,
h : Double,
) -> Array[Double] {
rdd_kernel_weights(self.data, self.cutoff, h, self.kernel)
}
///|
/// Package-private helper so `optimal_bandwidth`, `bandwidth_mse` and
/// `pilot_bandwidth` agree on the pilot bandwidth by construction. Two
/// copies of the Silverman call would be free to drift, and the criterion
/// would silently stop being the thing the search minimises.
///
/// A free function rather than a method: `fn name(self : Type)` at top
/// level is deprecated method syntax and does not bind as one outside an
/// `extend` block.
fn rdd_pilot_bandwidth(data : DoubleMLRDDData) -> Double {
silverman_bandwidth(data.score)
}
///|
///v0.110.0: MSE-optimal bandwidth (Cattaneo, Frandsen & Tchetgen 2020),
/// the procedure `rdrobust` exposes through `bwselect`.
///
/// Before this, bandwidth was a user-supplied constant and nothing chose
/// it. That is a real gap rather than a design choice: `rdrobust`'s
/// headline contribution to the RDD literature is precisely that a
/// plug-in bandwidth needs a data-driven default.
///
/// Procedure, on a grid of dimensionless multipliers of a pilot bandwidth:
///
/// 1. `h_pilot = silverman_bandwidth(data.score)` -- the rule-of-thumb
/// bandwidth of the RUNNING VARIABLE, not of the outcome. This is
/// the scale the grid is expressed in, so the result is invariant
/// to the units the running variable happens to be measured in.
/// 2. `bandwidth_mse(h)` is the criterion:
/// (xi(h) - xi(h_pilot))^2 + b^4 * sigma4 / n, b = h / h_pilot
/// The first term is the squared bias relative to the pilot, which
/// is treated as approximately unbiased; the second is the `O(b^4)`
/// variance growth of a local-polynomial derivative estimate.
/// 3. The minimiser over `b` in `[0.5, 2.0]` is returned, scaled back
/// by `h_pilot`.
///
/// The criterion is PUBLIC on purpose. Exposing it means a caller -- and
/// the test suite -- can check that the returned bandwidth actually
/// minimises it, instead of having to trust that the search did what its
/// name says. Without that, "optimal_bandwidth" is an assertion rather
/// than a checkable result.
///
/// Two honest deviations from `rdrobust`, both forced by this port's shape
/// rather than chosen:
///
/// * the grid has 50 points, not `rdrobust`'s 100;
/// * Cattaneo's `Sigma^4` is built from the DIFFERENCE of left and
/// right residuals at paired observations, and this port's
/// `rdd_side` returns two per-side residual arrays of UNEQUAL
/// length, so pairing them is a separate design decision. The pooled
/// sum of squares is used instead. Same units, same `b^4` scaling,
/// no pairing.
///
/// This does NOT modify the estimator. It is a function you call to
/// obtain a starting bandwidth; `fit` still uses `self.bandwidth`.
pub fn DoubleMLRDD::optimal_bandwidth(self : DoubleMLRDD) -> Double {
let h_pilot = rdd_pilot_bandwidth(self.data)
let l_min = 0.5
let l_max = 2.0
let n_grid = 50
let mut best_h = l_min * h_pilot
// Finite upper bound rather than an infinity sentinel: the four
// backends disagree about IEEE infinity through the FFI boundary,
// which is the same reason `tune.mbt` uses 1e300.
let mut best_mse = 1.0e300
for g = 0; g < n_grid; g = g + 1 {
let b = l_min + (l_max - l_min) * g.to_double() / (n_grid - 1).to_double()
let h = b * h_pilot
let mse = self.bandwidth_mse(h)
if mse < best_mse {
best_mse = mse
best_h = h
}
}
best_h
}
///|
/// The scale the bandwidth search is expressed in: the rule-of-thumb
/// bandwidth of the RUNNING VARIABLE, never of the outcome.
pub fn DoubleMLRDD::pilot_bandwidth(self : DoubleMLRDD) -> Double {
silverman_bandwidth(self.data.score)
}
///|
///v0.115.0: `rdrobust`'s SECOND bandwidth `b`, the one that drives bias
/// correction. `bwselect = "CCT"` computes both `h` and `b` by default,
/// and `b` is the one this port was missing: `optimal_bandwidth` picks
/// the bandwidth of the point estimate, and nothing here had an opinion
/// about the bandwidth at which to estimate the bias.
///
/// The rule is the same grid search as `optimal_bandwidth`, on the
/// ORDER-`q` criterion. That order is the entire difference: a local
/// polynomial of degree `j` has leading bias `O(h^(j+1))`, so the `j`-
/// order MSE criterion has a `b^(2(j+1))` variance term and wants a
/// different bandwidth than the `p`-order one.
///
///`q = 2` is `rdrobust`'s default and its `p`-order point estimate is
///local LINEAR (`p = 1`), so the two criteria being minimised here are
///genuinely different functions and `b` is not a rescaling of `h`.
pub fn DoubleMLRDD::optimal_bias_bandwidth(
self : DoubleMLRDD,
nnmatch? : Int = 3,
q? : Int = 2,
) -> Double {
let h_pilot = rdd_pilot_bandwidth(self.data)
let l_min = 0.5
let l_max = 2.0
let n_grid = 50
let mut best_b = l_min * h_pilot
let mut best_mse = 1.0e300
for g = 0; g < n_grid; g = g + 1 {
let mult = l_min +
(l_max - l_min) * g.to_double() / (n_grid - 1).to_double()
let cand = mult * h_pilot
let mse = self.bandwidth_mse(cand, nnmatch~, order=q)
if mse < best_mse {
best_mse = mse
best_b = cand
}
}
best_b
}
///|
/// v0.115.0: the local limit estimate `xi(h)` on ONE side of the
/// cutoff -- the intercept of the order-`order` local polynomial,
/// which for any order is the fitted value AT the cutoff.
///
/// Package-private: `bias` and `tau_bc` are the public
/// entry points, and exposing a bare one-sided limit estimate would
/// invite a caller to difference two of them and assume it had an RD
/// interpretation, which needs the kernel weights and both sides.
fn rdd_limit(m : DoubleMLRDD, side : Double, h : Double, order : Int) -> Double {
let (xi, _, _, _, _, _, _, _) = rdd_side(
m.ml_g,
m.data,
side,
m.cutoff,
h,
false,
m.cov_type,
m.kernel,
order~,
)
xi
}
///|
/// v0.116.0 CORRECTION: this is `rdrobust`'s bias,
/// `xi_p(h)_s - xi_bc_s`, where `xi_bc_s` is the BIAS-CORRECTED limit from
/// the `Q` sandwich -- NOT `xi_q(b)_s`.
///
/// v0.115.0 shipped this as `xi_p(h)_s - xi_q(b)_s` and documented the
/// difference as "`rdrobust`'s bias". Cross-checking against the upstream
/// Python port (`_verify/gen_v116_oracle.py`) shows the two agree to ~1e-15
/// when `b == h` and DIVERGE when `b != h` -- on the v0.116.0 oracle
/// fixture, `h = 0.30`, `b = 0.42`, left side:
///
/// | quantity | value |
/// |---------------------------|-----------|
/// | `xi_p(h)` | 0.1327887 |
/// | `xi_bc` (`rdrobust`) | 0.0561884 |
/// | `xi_q(b)` (was shipped) | 0.0681351 |
///
/// The default `rho = 1` gives `b == h`, so the shipped default path is
/// unchanged to 15 digits; the fix repairs the `rho != 1` path, which is
/// exactly where the second bandwidth is supposed to be doing work.
///
/// The wrong version was not detectable by any v0.115.0 gate, because
/// v0.115.0's own identity (`tau_bc == tau_bc_collapsed`) was only ever
/// checked INTERNALLY -- both sides of it computed the same wrong number.
/// An external oracle was required. That is now pinned by
/// `v116_tau_bc_matches_the_oracle_at_a_separate_bias_bandwidth`.
///
/// # The `b = h` case, and the trap in it
///
/// `rdrobust` makes `rho = 1` the default when `h` is supplied without
/// `b` (`b = h / rho`), so `b == h` is this port's DEFAULT, not an edge
/// case. It is NOT a degenerate estimator: with `b == h` the bias is the
/// difference between a local-LINEAR and a local-QUADRATIC fit at one
/// bandwidth, which is a real, non-zero quantity.
///
/// What degenerates is the BANDWIDTH SEARCH. `b` is supposed to be
/// chosen by minimising the order-`q` MSE criterion; if `b` is pinned to
/// `h`, the estimator is still correct but the second bandwidth is no
/// longer doing the job it was selected for, and `b := h` inside
/// `optimal_bias_bandwidth` becomes an EQUIVALENT MUTATION that no
/// downstream assertion can see -- `tau_bc` would return
/// `xi_2(h)_r - xi_2(h)_l` either way.
///
/// `v115_bias_bandwidth_order1_equals_the_point_bandwidth` pins the
/// mechanism (an order-1 bias bandwidth IS the point bandwidth), and
/// `v115_bias_bandwidth_differs_from_the_point_bandwidth` pins that the
/// default order-2 search does NOT return `h`, so the stage is actually
/// being selected rather than defaulted.
pub fn DoubleMLRDD::bias(
self : DoubleMLRDD,
b? : Double = 0.0,
rho? : Double = 1.0,
p? : Int = 1,
q? : Int = 2,
) -> (Double, Double) {
try {
require(p >= 1)
require(q >= 1)
require(rho > 0.0)
require(q >= p + 1)
let bias_b = if b > 0.0 { b } else { self.bandwidth / rho }
require(bias_b > 0.0)
let h = self.bandwidth
let (xi_l, xi_r) = rdd_xi_bc_pair(self, bias_b, p, q)
(rdd_limit(self, -1.0, h, p) - xi_l, rdd_limit(self, 1.0, h, p) - xi_r)
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// v0.116.0: the per-side bias-corrected limits `(xi_bc_left, xi_bc_right)`
/// at bias bandwidth `b`, from the `Q` sandwich.
///
/// Shared by `bias` and `tau_bc` so the two can never disagree about which
/// `xi_bc` they mean, and it is the only caller of `rdd_rb_side` with
/// `need_residual = false` -- the point estimate has no use for residuals.
fn rdd_xi_bc_pair(
m : DoubleMLRDD,
b : Double,
p : Int,
q : Int,
) -> (Double, Double) raise PreconditionError {
let (_, xi_l) = rdd_rb_side(
m.data,
-1.0,
m.cutoff,
m.bandwidth,
b,
p,
q,
m.kernel,
3,
false,
false,
)
let (_, xi_r) = rdd_rb_side(
m.data,
1.0,
m.cutoff,
m.bandwidth,
b,
p,
q,
m.kernel,
3,
false,
false,
)
(xi_l, xi_r)
}
///|
/// v0.115.0: the bias-CORRECTED point estimate -- procedure (ii) of the
/// three `rdrobust(all = TRUE)` reports, and the headline estimator once
/// the bias correction is on.
///
/// `tau_bc = (tau_cl_r - bias_r) - (tau_cl_l - bias_l)`, i.e. the
/// conventional estimate minus the estimated bias, computed in
/// `rdrobust`'s order.
///
/// # Why the method is named `tau_bc` and not `bias_corrected_coef`
///
/// `bias_corrected_coef` already exists on 15 other estimators in this
/// package, where it means the SANDWICH bias correction (Wald / delta
/// method) -- a completely different operation on a completely different
/// scale. Naming this one the same would have put two unrelated meanings
/// behind one method name across the same type.
///
/// `tau_bc` is `rdrobust`'s own field name for this quantity ("bias-
/// corrected local-polynomial estimate to the left and to the right of
/// the cutoff"), so the rename costs nothing in fidelity and removes the
/// collision.
///
/// # It does NOT collapse to a single fit -- v0.116.0 CORRECTION
///
/// v0.115.0's docs claimed, and two v0.115.0 gates asserted, that
///
/// tau_bc_s = xi_p(h)_s - [xi_p(h)_s - xi_q(b)_s] = xi_q(b)_s
///
/// so that `tau_bc` and `tau_bc_collapsed` are the same number. The
/// substitution is arithmetically valid, but the SECOND equality rests on
/// `bias` being `xi_p(h)_s - xi_q(b)_s`, and that is `rdrobust`'s bias
/// only when `b == h`. Checked against the upstream Python port at
/// `h = 0.30`, `b = 0.42`:
///
/// | quantity | upstream `rdrobust` | this port (v0.115.0) |
/// |-------------------------|--------------------|-----------------------|
/// | `xi_bc`, left | 0.0561884317 | 0.0681350859 |
/// | `tau_bc` | 0.7092313868 | 0.6910832334 |
///
/// v0.115.0 could not catch this on its own: its identity gate compared
/// two internal routes to the SAME wrong number, so it was green by
/// construction. See `bias` for the correction and
/// `v116_tau_bc_matches_the_oracle_at_a_separate_bias_bandwidth` for the
/// gate that does catch it.
///
/// `tau_bc_collapsed` (`xi_q(b)_r - xi_q(b)_l`) is still a real and useful
/// quantity -- it is the order-`q` fit at the bias bandwidth, and it is
/// what `rdrobust`'s bias-corrected limit equals when `b == h` -- so it
/// stays, with its own docs corrected.
///
/// # How it is evaluated, and why not "the long way"
///
/// `tau_bc = xi_bc_r - xi_bc_l`, with `xi_bc_s` the per-side
/// bias-corrected limit -- `rdrobust.py:907`'s own
/// `beta_bc = beta_bc_r - beta_bc_l`.
///
/// The algebraically equivalent `tau_cl - (bias_r - bias_l)` is what
/// v0.115.0 computed, and it is what this port used to describe itself as
/// doing. Both are the same number in exact arithmetic, but the long way
/// subtracts two `xi_p(h)` terms that then cancel, so it carries the
/// relative error of `xi_p(h)` into a difference of limits. Measured on
/// the `expand_v115_test.mbt` fixture at `b == h`, the long way sits
/// 8.5e-8 (relative) away from the direct difference; the direct one has
/// no such amplification. The long way is kept out for the same reason
/// upstream never writes it that way.
///
/// # Coverage
///
/// This is the point estimate only. The default `coef()` / `se()` are
/// UNCHANGED -- bias correction is opt-in. For the procedure (iii)
/// standard error see `tau_bc_se_rb`.
pub fn DoubleMLRDD::tau_bc(
self : DoubleMLRDD,
b? : Double = 0.0,
rho? : Double = 1.0,
p? : Int = 1,
q? : Int = 2,
) -> Double {
try {
require(p >= 1)
require(q >= 1)
require(rho > 0.0)
require(q >= p + 1)
let bias_b = if b > 0.0 { b } else { self.bandwidth / rho }
require(bias_b > 0.0)
let (xi_l, xi_r) = rdd_xi_bc_pair(self, bias_b, p, q)
xi_r - xi_l
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// v0.115.0: the bias-corrected estimate expressed as one order-`q` fit
/// per side at bandwidth `b` -- `xi_q(b)_r - xi_q(b)_l`.
///
/// # v0.116.0: this is NOT always `tau_bc`
///
/// v0.115.0's docs claimed this and `tau_bc` were the same number. They
/// agree to ~1e-15 when `b == h` and DIVERGE otherwise: `xi_q(b)` is the
/// order-`q` fit at `b`, whereas `tau_bc` uses `rdrobust`'s `Q`-sandwich
/// limit, and the two coincide only when the point and bias bandwidths
/// are the same number. Read `tau_bc`'s docs for the measured
/// discrepancy.
///
/// Kept because the quantity is meaningful in its own right and because
/// the identity at `b == h` is a genuine, checkable structural fact that
/// `v116_tau_bc_equals_the_collapsed_form_when_b_equals_h` pins.
///
/// Public so the comparison can be made from outside the file rather
/// than re-derived inside a test.
pub fn DoubleMLRDD::tau_bc_collapsed(
self : DoubleMLRDD,
b? : Double = 0.0,
rho? : Double = 1.0,
q? : Int = 2,
) -> Double {
try {
require(rho > 0.0)
let bias_b = if b > 0.0 { b } else { self.bandwidth / rho }
require(bias_b > 0.0)
rdd_limit(self, 1.0, bias_b, q) - rdd_limit(self, -1.0, bias_b, q)
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// v0.115.0: standard error of the bias-corrected estimate, under this
/// port's existing `cov_type` convention, applied to the order-`q`
/// fits at `b`.
///
/// # What this is NOT
///
/// `rdrobust`'s procedure (iii) is "bias-corrected estimates with ROBUST
/// standard errors", and its default robust estimator is `vce = "nn"`
/// -- nearest-neighbour matched residuals, the SAME matching rule
/// `bandwidth_sigma4` implements for the CRITERION. That variance
/// estimator is not implemented here. This returns the homoskedastic or
/// `HC0` variance the package already computes, on the `q`-order fits.
///
/// So the variance inflation that makes CCT intervals achieve nominal
/// coverage is NOT fully reproduced yet: this is a standard error for
/// the bias-corrected estimator, not `rdrobust`'s robust one. Shipping it
/// as if it were procedure (iii) is the v0.113 mistake again -- a
/// different estimator wearing the same name.
/// # v0.116.0: superseded for procedure (iii) -- kept for procedure (ii)
///
/// `tau_bc_se_rb` is `rdrobust`'s procedure (iii) standard error, and it
/// pairs with `tau_bc`. This function does not: it is the homoskedastic
/// / `HC0` variance of the order-`q` fits at `b` treated as independent
/// regressions, which is a standard error for the ORDER-`q` ESTIMATOR --
/// i.e. for `tau_bc_collapsed` -- rather than for `tau_bc`.
///
/// v0.115.0 shipped this next to `tau_bc` and said in these docs that
/// `rdrobust`'s robust variance "is not implemented here". That was
/// true then and is false now; the reason the function is kept rather
/// than deleted is that the order-`q` fit still needs an honest standard
/// error, and there was no name for it before.
///
/// Its numbers are unchanged. What changed is that `tau_bc` moved -- see
/// that method's v0.116.0 correction -- so the pairing documented in
/// v0.115.0 no longer holds.
pub fn DoubleMLRDD::tau_bc_se(
self : DoubleMLRDD,
b? : Double = 0.0,
rho? : Double = 1.0,
q? : Int = 2,
) -> Double {
try {
require(rho > 0.0)
let bias_b = if b > 0.0 { b } else { self.bandwidth / rho }
require(bias_b > 0.0)
let (_, var_l, _, _, _, _, _, _) = rdd_side(
self.ml_g,
self.data,
-1.0,
self.cutoff,
bias_b,
false,
self.cov_type,
self.kernel,
order=q,
)
let (_, var_r, _, _, _, _, _, _) = rdd_side(
self.ml_g,
self.data,
1.0,
self.cutoff,
bias_b,
false,
self.cov_type,
self.kernel,
order=q,
)
(var_l + var_r).sqrt()
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// The MSE-optimality criterion at an arbitrary bandwidth `h`:
/// `(xi(h) - xi(h_pilot))^2 + (h/h_pilot)^4 * sigma4 / n`, with
/// `xi_pilot` and the pooled residual scale `sigma4` both taken at
/// `h_pilot`.
///
/// Public so the choice made by `optimal_bandwidth` can be verified
/// against it. It costs two extra local fits per call, so it is not
/// something to call in a loop outside the search.
///
/// NOT annotated `raise PreconditionError`, unlike
/// `check_sample_splitting`: the `require`s below are on values the
/// caller controls, and MoonBit's analysis reports the error type as
/// never used (`unused_error_type`), which `--deny-warn` then treats as
/// a build failure. The two guards stay because they document the
/// precondition; they just do not widen the signature.
///|
///v0.115.0: `order` selects the polynomial order the criterion is built
///for. The variance term's exponent is `2 * (order + 1)` -- the leading
///bias of a local polynomial of degree `j` is `O(h^(j+1))`, so its square
///is `O(h^(2j+2))` -- and the pilot stage is refitted at the SAME order.
///
///`order = 1` is the default and reproduces v0.110.0-v0.113.0 exactly,
///including the floating-point association: the power is built by
///repeatedly squaring `b2 = b * b`, so `order = 1` still evaluates
///`b2 * b2`, not `(b*b)*(b*b)` written a different way.
///
///`order = 2` is what `rdrobust` uses for the bias-correction fit (`q`,
///its default), and is what `optimal_bias_bandwidth` minimises. The two
///criteria differ ONLY in this exponent and in the order of the fits --
///that is the whole reason `h` and `b` come out different numbers.
pub fn DoubleMLRDD::bandwidth_mse(
self : DoubleMLRDD,
h : Double,
nnmatch? : Int = 3,
order? : Int = 1,
) -> Double {
try {
let h_pilot = rdd_pilot_bandwidth(self.data)
require(h_pilot > 0.0)
require(h > 0.0)
// Only the pilot's limit estimate is used here; its `sigma4` goes
// through `bandwidth_variance_term`, which recomputes the SAME pilot
// stage. Recomputing rather than passing it keeps one definition of
// the variance half, and the stage is deterministic so the two
// `sigma4` values are bit-identical anyway.
let (xi_pilot, _) = rdd_pilot_stage(self, nnmatch, order)
let (yl, _, _, _, _, _, _, _) = rdd_side(
self.ml_g,
self.data,
-1.0,
self.cutoff,
h,
false,
self.cov_type,
self.kernel,
order~,
)
let (yr, _, _, _, _, _, _, _) = rdd_side(
self.ml_g,
self.data,
1.0,
self.cutoff,
h,
false,
self.cov_type,
self.kernel,
order~,
)
let bias = yr - yl - xi_pilot
bias * bias + self.bandwidth_variance_term(h, nnmatch~, order~)
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
///v0.115.0: the VARIANCE half of `bandwidth_mse`, split out so the
/// bandwidth-scaling exponent is checkable from outside.
///
/// variance_term(h) = (h / h_pilot)^(2 * (order + 1)) * sigma4 / n
///
///`bandwidth_mse` calls THIS rather than computing it inline, so the
///criterion and its decomposition cannot drift -- the same reason
///`bandwidth_sigma4` and the internal pilot share one implementation.
///
/// # Why the exponent needs its own gate
///
/// `2 * (order + 1)` is the entire content of the difference between the
/// `h` criterion and the `b` criterion: a local polynomial of degree `j`
/// has leading bias `O(h^(j+1))`, so the `j`-order MSE carries a
/// `b^(2(j+1))` variance term. Hardcoding it to `b^4` -- the order-1
/// value, which is what the code did before this parameter existed --
/// leaves every OTHER assertion in `expand_v115_test.mbt` green: the
/// search still minimises whatever criterion it is given, the bias is
/// still non-zero, `b` still differs from `h`, and the collapse identity
/// still holds.
///
/// Mutation M3 in `_verify/mut_v115_rdd.ps1` is exactly that, and it
/// SURVIVED before this function existed. Doubling the bandwidth must
/// multiply the variance term by `2^(2(order+1))` -- 16 at order 1, 64 at
/// order 2 -- and `v115_variance_term_doubles_by_the_order_exponent` is
/// the assertion that says so.
pub fn DoubleMLRDD::bandwidth_variance_term(
self : DoubleMLRDD,
h : Double,
nnmatch? : Int = 3,
order? : Int = 1,
) -> Double {
try {
let n = self.n_obs()
let h_pilot = rdd_pilot_bandwidth(self.data)
require(h_pilot > 0.0)
require(h > 0.0)
let (_, sigma4) = rdd_pilot_stage(self, nnmatch, order)
let b = h / h_pilot
let b2 = b * b
// Starting from `b2` and multiplying by `b2` exactly `order` times
// keeps `order = 1` on the literal `b2 * b2` that v0.110.0
// computed, so the default path is bit-identical to before.
let mut variance_term = b2
for step = 0; step < order; step = step + 1 {
variance_term = variance_term * b2
}
variance_term * sigma4 / n.to_double()
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// `Sigma^4` of the MSE-optimal criterion, as **nearest-neighbour
/// matched** residual differences -- `rdrobust`'s `vce = "nn"` with
/// `nnmatch = 3`, its default.
///
/// # Why this replaced a pooled sum of squares
///
/// v0.110.0-v0.112.0 used `sum of all squared residuals on both sides`.
/// That is NOT what Cattaneo et al. (2020) do: `rdrobust` pairs a left
/// residual with its `nnmatch` NEAREST right-side neighbours by
/// running-variable distance and squares the differences. Pooling discards
/// the matching entirely, so it is a different estimator wearing the same
/// name -- and it changes the number `optimal_bandwidth` returns.
///
/// # The rule, stated exactly
///
/// For each left local-row observation `i` with running-variable offset
/// `u_i < 0`, take the `nnmatch` right observations whose `|u|` are closest
/// to `|u_i|`, and average the squared residual differences over those
/// neighbours:
///
/// sigma4 = (1 / N_left) * SUM_i mean_j (eps_right[j] - eps_left[i])^2
///
/// with the inner mean taken over however many neighbours exist. Ties in
/// `|u|` are broken by ascending row index, so the result is deterministic.
///
/// # What is a deliberate simplification, and what is not
///
/// The one thing this does NOT do is shrink the left sample: `rdrobust`
/// weights by a local-linear kernel over the matched pairs. Pooling by
/// left observation count, as above, is a coarser rule with the same
/// units (outcome^2) and the same `b^4` scaling in the criterion, so it
/// keeps the search's shape; it does not reproduce `rdrobust`'s numbers
/// to the digit. That is recorded rather than papered over.
///
/// `nnmatch` is a parameter because a mutation that ignores it is
/// otherwise indistinguishable from one that honours it -- see
/// `v110_sigma4_responds_to_nnmatch`.
///
/// Two neighbours short: the inner mean uses however many exist. If the
/// right side is EMPTY there is nothing to match and the procedure is
/// undefined; `require` fires and the caller sees an abort, which is the
/// same outcome the pilot stage already had.
pub fn DoubleMLRDD::bandwidth_sigma4(
self : DoubleMLRDD,
nnmatch? : Int = 3,
order? : Int = 1,
) -> Double {
let (_, sigma4) = rdd_pilot_stage(self, nnmatch, order)
sigma4
}
///|
/// The pilot-stage pair the criterion is built from: the intercept
/// difference at `h_pilot`, and the nearest-neighbour matched
/// `Sigma^4` there.
///
/// Shared by `bandwidth_mse` and `bandwidth_sigma4` so the two cannot
/// drift -- a criterion whose public `sigma4` and internal `sigma4`
/// were computed separately would make the identity above fail for
/// reasons unrelated to the maths.
///
/// A free function rather than a method, for the same reason as
/// `rdd_pilot_bandwidth`: `fn name(self : Type)` at top level is
/// deprecated method syntax and does not bind as one.
fn rdd_pilot_stage(
m : DoubleMLRDD,
nnmatch : Int,
order : Int,
) -> (Double, Double) {
try {
require(nnmatch >= 1)
require(order >= 1)
let h_pilot = rdd_pilot_bandwidth(m.data)
require(h_pilot > 0.0)
let (yl, _, _, res_l, _, _, _, _) = rdd_side(
m.ml_g,
m.data,
-1.0,
m.cutoff,
h_pilot,
false,
m.cov_type,
m.kernel,
order~,
)
let (yr, _, _, res_r, _, _, _, _) = rdd_side(
m.ml_g,
m.data,
1.0,
m.cutoff,
h_pilot,
false,
m.cov_type,
m.kernel,
order~,
)
let u_l = rdd_side_offsets(m, -1.0, h_pilot)
let u_r = rdd_side_offsets(m, 1.0, h_pilot)
(yr - yl, rdd_nn_matched_sigma4(u_l, res_l, u_r, res_r, nnmatch))
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// The running-variable offsets of one side's local rows, in the SAME
/// order `rdd_design` assigned them.
///
/// The predicate and the iteration order are copied from
/// `rdd_kernel_weights` verbatim (left `u < 0` first, then right
/// `u >= 0`, both in ascending observation index). That is not
/// stylistic: `rdd_side` returns residuals indexed by LOCAL ROW, and
/// matching a residual to the wrong `|u|` would silently produce a
/// different `Sigma^4` with no error anywhere. If this predicate ever
/// drifts from `rdd_kernel_weights`, the matching is wrong.
fn rdd_side_offsets(
m : DoubleMLRDD,
side : Double,
h : Double,
) -> Array[Double] {
let out : Array[Double] = []
for i = 0; i < m.data.n_obs(); i = i + 1 {
let u = m.data.score[i] - m.cutoff
if (side < 0.0 && u < 0.0) || (side > 0.0 && u >= 0.0) {
if u.abs() <= h {
out.push(u)
}
}
}
out
}
///|
/// `sigma4 = (1/N_left) * SUM_i mean_j (eps_right[j] - eps_left[i])^2`,
/// with the inner mean over the `nnmatch` right neighbours whose `|u|`
/// is closest to `|u_i|`, ties broken by ascending index.
///
/// The neighbour search is a straight selection over the right side:
/// `nnmatch` is a small constant (3 by default) and the sides are
/// O(n) long, so an O(N_left * nnmatch) scan after one sort is both
/// simpler to read and cheaper than a heap.
fn rdd_nn_matched_sigma4(
u_l : Array[Double],
res_l : Array[Double],
u_r : Array[Double],
res_r : Array[Double],
nnmatch : Int,
) -> Double raise PreconditionError {
let n_l = u_l.length()
let n_r = u_r.length()
// No right observations to match against: the procedure is undefined
// and the caller must see that, not a silent 0.
if n_r == 0 || n_l == 0 {
require(false)
}
// right rows sorted by |u| once, so each left row's nearest neighbours
// are a contiguous scan
let idx = Array::make(n_r, 0)
for i = 0; i < n_r; i = i + 1 {
idx[i] = i
}
let keys = u_r.copy()
idx.sort_by(fn(a : Int, b : Int) {
let ka = keys[a].abs()
let kb = keys[b].abs()
if ka < kb {
-1
} else if ka > kb {
1
} else {
a - b
}
})
let sorted_abs = Array::make(n_r, 0.0)
for k = 0; k < n_r; k = k + 1 {
sorted_abs[k] = u_r[idx[k]].abs()
}
let m = if nnmatch < n_r { nnmatch } else { n_r }
let mut acc = 0.0
for i = 0; i < n_l; i = i + 1 {
let target = u_l[i].abs()
// expand symmetrically out from the insertion point
let mut lo = 0
let mut hi = n_r - 1
while lo < hi {
let mid = (lo + hi) / 2
if sorted_abs[mid] < target {
lo = mid + 1
} else {
hi = mid
}
}
let mut sum = 0.0
let mut taken = 0
let mut a = lo - 1
let mut b = lo
while taken < m {
let take_left = if a < 0 {
false
} else if b >= n_r {
true
} else {
let dl = target - sorted_abs[a]
let db = sorted_abs[b] - target
db < dl
}
if take_left {
let row = idx[a]
let d = res_r[row] - res_l[i]
sum = sum + d * d
a = a - 1
} else {
let row = idx[b]
let d = res_r[row] - res_l[i]
sum = sum + d * d
b = b + 1
}
taken = taken + 1
}
acc = acc + sum / taken.to_double()
}
acc / n_l.to_double()
}
///|
/// v0.116.0: the nearest-neighbour matched residual -- the residual behind
/// `rdrobust`'s `vce = "nn"`. Ported from
/// `rdrobust/src/rdrobust/funs.py::_nn_residuals_jit`, read line by line
/// rather than reconstructed from the `vce` documentation (which states
/// the parameter exists and says nothing about how it is built).
///
/// # What the quantity IS
///
/// For each observation `pos`, take the `nnmatch` observations whose
/// running value is nearest to `x[pos]` -- expanding alternately to the
/// left and to the right, always taking whichever side has the smaller
/// running-variable gap -- and write
///
/// ```
/// res[pos] = sqrt(Ji / (Ji + 1)) * (y[pos] - mean(y over those neighbours))
/// ```
///
/// where `Ji` counts the neighbours EXCLUDING `pos` itself. The
/// `sqrt(Ji/(Ji+1))` factor is the finite-sample correction that keeps
/// `res` on the scale of a regression residual; it pays for that with a
/// variance roughly `Ji + 1` times the plug-in residual's.
///
/// The crucial structural fact, and the reason `vce = "nn"` is cheap:
/// **the residual is a function of `y` and `x` only.** No fit, no
/// bandwidth, no kernel. `rdrobust` exploits this at rdrobust.py:1055 --
/// when `vce == "nn"` it assigns `res_b_l = res_h_l`, reusing the
/// CONVENTIONAL estimator's residuals for the BIAS-CORRECTED variance
/// instead of refitting. A test pins that invariance directly.
///
/// # Tie handling and mass points
///
/// Two equal-`x` observations at distance `dleft == dright` are taken
/// together (both sides advance one step), unless the two gaps differ by
/// more than `max(dleft, dright) * sqrt(eps)` on the side that is being
/// expanded -- that tolerance is what keeps the expansion from comparing
/// distances that differ only by rounding. It is the MAXIMUM, not the
/// minimum (upstream's `dleft if dleft > dright else dright`, funs.py:114)
/// and this file originally wrote the minimum. The two agree on every
/// fixture in `expand_v116_test.mbt`: the window in which they can differ
/// is `|dleft - dright|` landing between `min * sqrt(eps)` and
/// `max * sqrt(eps)`, roughly 1e-16 wide relative, and no double-precision
/// running variable lands there by accident. So this is a fidelity fix
/// that the gates CANNOT see -- recorded here, and flagged in
/// `_verify/mut_v116_rdd.ps1` as a mutation expected to survive.
///
/// Equal-`x` groups are walked as whole blocks, so an observation is never
/// split from its own duplicate when the boundary falls inside a group.
///
/// # This is NOT the criterion's `sigma4`
///
/// The private `rdd_nn_matched_sigma4` matches a LEFT residual to a
/// RIGHT residual to build the bandwidth criterion's `Sigma^4`. This
/// matches each observation to its own local neighbourhood inside ONE
/// side. They share a name fragment and nothing else; see that function
/// for its own doc.
///
/// # Preconditions
///
/// `x` must be sorted ascending -- the neighbourhood walk assumes it --
/// and needs at least two entries, since a residual is defined against
/// at least one other observation. `nnmatch >= 1` for the same reason:
/// `cap = min(nnmatch, n - 1)` and `Ji` would otherwise be zero.
///
/// Ports `funs.py:92-143`.
pub fn nn_residual(
x : Array[Double],
y : Array[Double],
nnmatch? : Int = 3,
) -> Array[Double] {
try {
let n = x.length()
require(y.length() == n)
require(n >= 2)
require(nnmatch >= 1)
// Equal-`x` groups, walked whole. `dupsid[k]` is k's 1-based offset
// WITHIN its group and `dups[k]` the group's length, so the
// expansion starts with `lpos` steps to the group start and
// `rpos` steps to the group end.
let dups = Array::make(n, 0)
let dupsid = Array::make(n, 0)
let mut i0 = 0
while i0 < n {
let s = i0
let mut e = i0
while e + 1 < n && x[e + 1] == x[s] {
e = e + 1
}
let m = e - s + 1
for k = s; k <= e; k = k + 1 {
dups[k] = m
dupsid[k] = k - s + 1
}
i0 = e + 1
}
let cap = if nnmatch < n - 1 { nnmatch } else { n - 1 }
// sqrt(eps) for float64; the distance tie-tolerance scale.
let nn_tol_eps = 1.4901161193847656e-8
let res = Array::make(n, 0.0)
for pos = 0; pos < n; pos = pos + 1 {
let mut rpos = dups[pos] - dupsid[pos]
let mut lpos = dupsid[pos] - 1
while lpos + rpos < cap {
if pos - lpos - 1 < 0 {
rpos += dups[pos + rpos + 1]
} else if pos + rpos + 1 >= n {
lpos += dups[pos - lpos - 1]
} else {
let dleft = x[pos] - x[pos - lpos - 1]
let dright = x[pos + rpos + 1] - x[pos]
let largest = if dleft > dright { dleft } else { dright }
let nn_tol = largest * nn_tol_eps
if dleft - dright > nn_tol {
rpos += dups[pos + rpos + 1]
} else if dright - dleft > nn_tol {
lpos += dups[pos - lpos - 1]
} else {
rpos += dups[pos + rpos + 1]
lpos += dups[pos - lpos - 1]
}
}
}
let lo = pos - lpos
let hi = pos + rpos + 1
let ji = hi - lo - 1
let sf = (ji.to_double() / (ji + 1).to_double()).sqrt()
let mut s_y = 0.0
for j = lo; j < hi; j = j + 1 {
s_y += y[j]
}
s_y -= y[pos]
res[pos] = sf * (y[pos] - s_y / ji.to_double())
}
res
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// `kernel.weight(u, h)` with `rdrobust`'s support mask applied.
///
/// # Why the mask is not optional
///
/// Three of the five kernels -- `Triangular`, `Epanechnikov`,
/// `Quadratic` -- are NEGATIVE outside `|u| <= h` (`1 - t` and `1 - t^2`
/// have the wrong sign past `t = 1`). Upstream multiplies every weight by
/// `(|u| <= 1)` (`rdrobust_kweight`, funs.py:668-677) for exactly that
/// reason, and the port's `RDDKernel::weight` deliberately does NOT,
/// because every pre-v0.116.0 caller selects rows with `u.abs() <= h`
/// first and so never evaluates the negative branch.
///
/// v0.116.0 is the first caller that builds a fit on a window WIDER than
/// the bandwidth -- the CCT `max(h, b)` union -- and needs `W_h` to fall
/// to exactly zero outside the h-sub-window rather than to go negative.
/// Without this, `b > h` produces a Gram matrix with a non-positive pivot
/// and `solve_spd` aborts. The first oracle cross-check at `h = b`
/// passed precisely because `max(h, b) == h` there, so the negative
/// branch was never evaluated; the `h != b` cross-check is what found it.
fn rdd_weight_at(kernel : RDDKernel, u : Double, h : Double) -> Double {
if u.abs() <= h {
kernel.weight(u, h)
} else {
0.0
}
}
///|
/// v0.116.0: the two per-side pieces of `rdrobust`'s procedure (iii),
/// built from one shared construction.
///
/// Returns `(v00, xibc)`:
///
/// - `v00 = e_0' V_rb e_0` where `V_rb = Gp^-1 Q' diag(res^2) Q Gp^-1`
/// -- the side's contribution to the bias-corrected VARIANCE.
/// - `xibc = e_0' Gp^-1 Q' y` -- the side's bias-CORRECTED limit, the
/// quantity procedure (ii) and (iii) report and this port's `tau_bc`
/// differences. It is `rdrobust`'s `xi_bc`, NOT `xi_q(b)`; see
/// `tau_bc` for why those two coincide only when `b == h`.
///
/// `need_residual = false` skips the residual construction entirely,
/// which the point-estimate callers want: it is the only part of this
/// function that costs O(n * nnmatch).
///
/// This is a direct port of the `else` branch at rdrobust.py:1092-1097,
/// which is the branch every sharp, unclustered, non-CRV call takes --
/// including the `h == b` case, because the two earlier `hb_match`
/// branches (1077, 1084) are both gated on cluster / CRV2 / CRV3.
///
/// The meat matrix needs `Q` as a whole (n x (p+1)) but only ONE entry
/// of the result, so the sandwich is never formed: with
/// `a = Gp^-1 e_0` and `M = Q' diag(res^2) Q`,
///
/// ```
/// e_0' Gp^-1 M Gp^-1 e_0 = a' M a = SUM_i res[i]^2 * (Q_i . a)^2
/// ```
///
/// which is one pass and one Cholesky solve. Forming `Gp^-1 M Gp^-1`
/// explicitly would cost two more matrix products for the same number.
///
/// # The rank-1 correction
///
/// Upstream writes `Q = RW_p - h^(p+1) * m * L'`, where `m` is an
/// (n x 1) column and `L[a] = SUM_k R_p[k][a] * W_h[k] * ((x_k - c)/h)^(p+1)`.
/// The `h^(p+1)` in front of `L` cancels the `1/h^(p+1)` inside it, so
/// this file carries `L` pre-multiplied:
///
/// ```
/// Ls[a] = SUM_k R_p[k][a] * W_h[k] * (x_k - c)^(p+1)
/// Q[i][a] = R_p[i][a] * W_h[i] - m[i] * Ls[a]
/// ```
///
/// `m` is itself `(R_q . gamma) .* W_b` with `gamma` the `(p+1)`-th ROW
/// of `Gq^-1`, obtained here as the solution of `Gq gamma = e_{p+1}`
/// (rdrobust.py:860) -- one solve instead of a matrix inverse.
///
/// # Kernel weights omit the `1/h`
///
/// `rdrobust_kweight` returns `K(u)/h`; this file's `RDDKernel::weight`
/// returns `K(u)`. The factor cancels exactly -- `Q` scales by it and
/// `Gp^-1` by its inverse -- so no constant is lost. Verified against the
/// oracle rather than argued: `expand_v116_test.mbt` compares the two
/// variances to 12 significant figures.
///
/// # Window
///
/// Upstream fits on the union window `max(h, b)` (rdrobust.py:795-797)
/// and lets `W_h` fall to zero outside the h-sub-window. Reproduced
/// here, because using the b-window with h-weights would silently refit
/// the `p`-order stage at `b` and make `h` inert.
///
/// # `res` sources
///
/// `nn_res` selects the residual entering the meat:
/// - `true` -- `vce = "nn"`: the fit-free neighbour residual of
/// `nn_residual`, which upstream feeds to BOTH stages.
/// - `false` -- `vce = "hc0"`: the plug-in residual `y - R_q . beta_q`
/// of the q-order fit at `b`, which upstream uses for the b stage only.
///
/// Ports rdrobust.py:850-863 (Q, m, invG_p) and 1092-1097 (the meat).
fn rdd_rb_side(
data : DoubleMLRDDData,
side : Double,
cutoff : Double,
h : Double,
b : Double,
p : Int,
q : Int,
kernel : RDDKernel,
nnmatch : Int,
nn_res : Bool,
need_residual : Bool,
) -> (Double, Double) raise PreconditionError {
require(q >= p + 1)
// Upstream's row selection is `w > 0` on the weight built with the
// governing bandwidth `max(h, b)` (rdrobust.py:795-797), which for a
// compact kernel is a strict inequality in `|u|`. Selecting with
// `u.abs() <= win` instead would admit the zero-weight boundary row,
// where the kernel is 0 but the observation is still a live neighbour
// for `nn_residual`.
let gov = if b > h { b } else { h }
let ids : Array[Int] = []
for i = 0; i < data.n_obs(); i = i + 1 {
let u = data.score[i] - cutoff
if (side < 0.0 && u < 0.0) || (side > 0.0 && u >= 0.0) {
if rdd_weight_at(kernel, u, gov) > 0.0 {
ids.push(i)
}
}
}
let n = ids.length()
// One degree of freedom beyond the q-order fit, so `Gq` is not merely
// invertible in principle.
require(n >= q + 2)
// Upstream sorts the running variable once, globally, before splitting
// into sides; `nn_residual` walks neighbourhoods in x order, so the
// window rows must be in x order too.
let score_ref = data.score
ids.sort_by(fn(a : Int, c : Int) {
let ka = score_ref[a]
let kc = score_ref[c]
if ka < kc {
-1
} else if ka > kc {
1
} else {
0
}
})
let kq = q + 1
let kp = p + 1
let uu = Array::make(n, 0.0)
let yy = Array::make(n, 0.0)
let wh = Array::make(n, 0.0)
let wb = Array::make(n, 0.0)
let xp = Array::make(n * kq, 0.0)
for k = 0; k < n; k = k + 1 {
let u = score_ref[ids[k]] - cutoff
uu[k] = u
yy[k] = data.y[ids[k]]
wh[k] = rdd_weight_at(kernel, u, h)
wb[k] = rdd_weight_at(kernel, u, b)
// Vandermonde by successive multiplication, matching `_vander`
// (funs.py:597): column 0 is the constant 1.0, column j is column
// j-1 times u.
xp[k * kq] = 1.0
let mut pw = 1.0
for j = 1; j <= q; j = j + 1 {
pw = pw * uu[k]
xp[k * kq + j] = pw
}
}
let gq = Matrix::zeros(kq, kq)
for k = 0; k < n; k = k + 1 {
let wk = wb[k]
for a = 0; a < kq; a = a + 1 {
let xa = wk * xp[k * kq + a]
for c = a; c < kq; c = c + 1 {
let val = gq.get(a, c) + xa * xp[k * kq + c]
gq.set(a, c, val)
if c != a {
gq.set(c, a, val)
}
}
}
}
// gamma = row p+1 of Gq^-1, obtained as the solution of Gq g = e_{p+1}.
let rhs_q = Array::make(kq, 0.0)
rhs_q[p + 1] = 1.0
let gamma = solve_spd(gq, rhs_q)
let m = Array::make(n, 0.0)
for k = 0; k < n; k = k + 1 {
let mut s = 0.0
for a = 0; a < kq; a = a + 1 {
s += xp[k * kq + a] * gamma[a]
}
m[k] = wb[k] * s
}
let gp = Matrix::zeros(kp, kp)
for k = 0; k < n; k = k + 1 {
let wk = wh[k]
for a = 0; a < kp; a = a + 1 {
let xa = wk * xp[k * kq + a]
for c = a; c < kp; c = c + 1 {
let val = gp.get(a, c) + xa * xp[k * kq + c]
gp.set(a, c, val)
if c != a {
gp.set(c, a, val)
}
}
}
}
// avec = first column of Gp^-1 = Gp^-1 e_0.
let rhs_p = Array::make(kp, 0.0)
rhs_p[0] = 1.0
let avec = solve_spd(gp, rhs_p)
// Ls[a], with upstream's h^(p+1) already absorbed.
let ls = Array::make(kp, 0.0)
for k = 0; k < n; k = k + 1 {
let pw = xp[k * kq + p + 1]
let wk = wh[k] * pw
for a = 0; a < kp; a = a + 1 {
ls[a] += xp[k * kq + a] * wk
}
}
let mut ls_a = 0.0
for a = 0; a < kp; a = a + 1 {
ls_a += ls[a] * avec[a]
}
let res = if !need_residual {
[]
} else if nn_res {
nn_residual(uu, yy, nnmatch~)
} else {
let rhs_y = Array::make(kq, 0.0)
for k = 0; k < n; k = k + 1 {
for a = 0; a < kq; a = a + 1 {
rhs_y[a] += wb[k] * xp[k * kq + a] * yy[k]
}
}
let bq = solve_spd(gq, rhs_y)
let out = Array::make(n, 0.0)
for k = 0; k < n; k = k + 1 {
let mut fit = 0.0
for a = 0; a < kq; a = a + 1 {
fit += xp[k * kq + a] * bq[a]
}
out[k] = yy[k] - fit
}
out
}
// v_i = (Q_i . avec), and the two reductions that need it.
// v00 = e_0' invG_p Q' diag(res^2) Q invG_p e_0 = SUM res_i^2 v_i^2
// xibc = e_0' invG_p Q' y = SUM y_i v_i
let mut v00 = 0.0
let mut xibc = 0.0
for k = 0; k < n; k = k + 1 {
let mut ai = 0.0
for a = 0; a < kp; a = a + 1 {
ai += xp[k * kq + a] * avec[a]
}
let vi = wh[k] * ai - m[k] * ls_a
if need_residual {
v00 += res[k] * res[k] * vi * vi
}
xibc += yy[k] * vi
}
(v00, xibc)
}
///|
/// v0.116.0: `rdrobust`'s PROCEDURE (iii) standard error -- the
/// bias-corrected point estimate `tau_bc` with the robust variance.
///
/// This is the round that completes CCT inference on this port.
/// `tau_bc` shipped in v0.115.0 with `tau_bc_se` alongside it, and that
/// `tau_bc_se` is NOT this: it is the homoskedastic / HC0 variance of the
/// order-`q` fits at `b` treated as independent regressions. Upstream's
/// procedure (iii) is a different sandwich, and under `vce = "nn"` a
/// different residual again. Two changes are needed, not one:
///
/// 1. **The meat matrix.** Upstream does not sandwich the q-order fit.
/// It sandwiches `Q`, the p-order fit MINUS its estimated bias term
/// (rdrobust.py:856-863). The bread is `Gp^-1` at bandwidth `h`
/// even though the estimate is a q-order quantity at bandwidth `b`.
/// 2. **The residual.** For `vce = "nn"` this is the fit-free
/// nearest-neighbour residual, not a fitted residual at all.
///
/// # Point estimate
///
/// `tau_bc(b, rho, q)` -- the argument of this function -- is the one
/// whose standard error is returned. `rho` gives `b = h / rho`, matching
/// `rdrobust`'s convention, so the default `rho = 1` means `b == h`.
///
/// # `vce`
///
/// - `"nn"` (default) -- nearest-neighbour matched residuals. This is
/// `rdrobust`'s own default and the estimator CCT inference is built
/// on: it is the only one that needs no homoskedasticity assumption,
/// which is why the coverage it delivers is the one the literature
/// reports.
/// - `"hc0"` -- plug-in residuals of the q-order fit at `b`. Exposed
/// because it shares this function's entire `Q` construction, so it
/// turns the oracle cross-check into two independent ones: a fault in
/// `Q` would have to reproduce itself in both to survive.
///
/// An unrecognised `vce` aborts with the accepted set rather than
/// silently falling back, for the reason `RDDKernel::parse` does.
///
/// # Sharp designs only
///
/// As with `tau_bc`, a fuzzy first stage is not threaded through here;
/// upstream's `s_Y` delta-method terms (rdrobust.py:923-939) are a
/// separate block.
///
/// Ports rdrobust.py:1055-1060 (residual source), 850-863 (`Q`, `m`,
/// `invG_p`) and 1092-1097 (the meat and bread).
pub fn DoubleMLRDD::tau_bc_se_rb(
self : DoubleMLRDD,
b? : Double = 0.0,
rho? : Double = 1.0,
q? : Int = 2,
vce? : String = "nn",
nnmatch? : Int = 3,
) -> Double {
try {
require(rho > 0.0)
let bias_b = if b > 0.0 { b } else { self.bandwidth / rho }
require(bias_b > 0.0)
let nn_res = match vce {
"nn" => true
"hc0" => false
_ => abort("unknown vce: " + vce + " (accepted: \"nn\", \"hc0\")")
}
// The point estimator is local-linear (rdrobust's p = 1) and has no
// order knob, so p is fixed here rather than a parameter that could
// silently disagree with `coef()`.
let (v_l, _) = rdd_rb_side(
self.data,
-1.0,
self.cutoff,
self.bandwidth,
bias_b,
1,
q,
self.kernel,
nnmatch,
nn_res,
true,
)
let (v_r, _) = rdd_rb_side(
self.data,
1.0,
self.cutoff,
self.bandwidth,
bias_b,
1,
q,
self.kernel,
nnmatch,
nn_res,
true,
)
(v_l + v_r).sqrt()
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}