rd2d 0.1.0__py3-none-any.whl

rd2d/__init__.py

@@ -0,0 +1,17 @@
+"""Local polynomial methods for boundary discontinuity designs."""
+
+from .distance import rdbw2d_distance, rdbw2d_dist, rd2d_distance, rd2d_dist
+from .location import rdbw2d, rd2d
+from .results import RD2DResult, SummaryResult, summary
+
+__all__ = [
+    "RD2DResult",
+    "SummaryResult",
+    "rdbw2d",
+    "rd2d",
+    "rdbw2d_dist",
+    "rdbw2d_distance",
+    "rd2d_dist",
+    "rd2d_distance",
+    "summary",
+]

rd2d/_utils.py

@@ -0,0 +1,272 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Iterable
+
+import numpy as np
+import pandas as pd
+from scipy import stats
+
+
+def as_1d(x, name: str) -> np.ndarray:
+    arr = np.asarray(x, dtype=float)
+    if arr.ndim != 1:
+        arr = np.ravel(arr)
+    if arr.size == 0:
+        raise ValueError(f"{name} must not be empty.")
+    return arr
+
+
+def as_2d(x, name: str, ncol: int | None = None) -> np.ndarray:
+    arr = np.asarray(x, dtype=float)
+    if arr.ndim == 1:
+        arr = arr.reshape(-1, 1)
+    if arr.ndim != 2:
+        raise ValueError(f"{name} must be a two-dimensional array.")
+    if ncol is not None and arr.shape[1] != ncol:
+        raise ValueError(f"{name} must have exactly {ncol} columns.")
+    return arr
+
+
+def check_lengths(y: np.ndarray, *arrays: np.ndarray) -> None:
+    n = len(y)
+    for arr in arrays:
+        if len(arr) != n:
+            raise ValueError("Input vectors and rows of matrix inputs must have the same length.")
+
+
+def complete_cases(*arrays: np.ndarray) -> np.ndarray:
+    mask = np.ones(len(arrays[0]), dtype=bool)
+    for arr in arrays:
+        arr = np.asarray(arr)
+        try:
+            finite = np.isfinite(arr.astype(float))
+        except (TypeError, ValueError):
+            finite = ~pd.isna(arr)
+        if arr.ndim == 1:
+            mask &= finite
+        else:
+            mask &= np.all(finite, axis=1)
+    return mask
+
+
+def normalize_binary(x, name: str) -> np.ndarray:
+    arr = np.asarray(x)
+    if arr.dtype == bool:
+        return arr.astype(bool)
+    vals = np.unique(arr[np.isfinite(arr.astype(float))].astype(float))
+    if not set(vals.tolist()).issubset({0.0, 1.0}):
+        raise ValueError(f"{name} must be logical or contain only 0 and 1.")
+    return arr.astype(float).astype(bool)
+
+
+def validate_order(value, name: str) -> int:
+    if value is None:
+        raise ValueError(f"{name} must not be None.")
+    value = float(value)
+    if not np.isfinite(value) or value < 0 or abs(value - round(value)) > np.sqrt(np.finfo(float).eps):
+        raise ValueError(f"{name} must be a nonnegative integer.")
+    return int(round(value))
+
+
+def validate_deriv(deriv: Iterable[float], p: int) -> tuple[int, int]:
+    arr = np.asarray(tuple(deriv), dtype=float)
+    if arr.shape != (2,) or np.any(~np.isfinite(arr)) or np.any(arr < 0):
+        raise ValueError("deriv must be a nonnegative integer vector of length 2.")
+    if np.any(np.abs(arr - np.round(arr)) > np.sqrt(np.finfo(float).eps)):
+        raise ValueError("deriv must be a nonnegative integer vector of length 2.")
+    out = tuple(int(v) for v in np.round(arr))
+    if sum(out) > p:
+        raise ValueError("sum(deriv) must be less than or equal to p.")
+    return out
+
+
+def kernel_weights(u: np.ndarray, kernel: str) -> np.ndarray:
+    kernel = kernel.lower()
+    if kernel in {"tri", "triangular"}:
+        return np.maximum(1.0 - np.abs(u), 0.0) * (np.abs(u) <= 1.0)
+    if kernel in {"epa", "epanechnikov"}:
+        return 0.75 * (1.0 - u**2) * (np.abs(u) <= 1.0)
+    if kernel in {"uni", "uniform"}:
+        return 0.5 * (np.abs(u) <= 1.0)
+    if kernel in {"gau", "gaussian"}:
+        return stats.norm.pdf(u)
+    raise ValueError("kernel must be one of tri, epa, uni, or gau.")
+
+
+def multi_indices_2d(p: int) -> list[tuple[int, int]]:
+    out: list[tuple[int, int]] = []
+    for deg in range(p + 1):
+        for ypow in range(deg + 1):
+            xpow = deg - ypow
+            out.append((xpow, ypow))
+    return out
+
+
+def basis_2d(centered: np.ndarray, p: int) -> np.ndarray:
+    idx = multi_indices_2d(p)
+    x1 = centered[:, 0]
+    x2 = centered[:, 1]
+    return np.column_stack([(x1**a) * (x2**b) for a, b in idx])
+
+
+def basis_1d(distance: np.ndarray, p: int) -> np.ndarray:
+    return np.column_stack([distance**j for j in range(p + 1)])
+
+
+def target_2d(p: int, deriv: tuple[int, int], tangvec: np.ndarray | None, row: int) -> np.ndarray:
+    target = np.zeros(len(multi_indices_2d(p)))
+    if tangvec is not None:
+        if p < 1:
+            raise ValueError("tangvec requires p >= 1.")
+        idx = multi_indices_2d(p)
+        target[idx.index((1, 0))] = tangvec[row, 0]
+        target[idx.index((0, 1))] = tangvec[row, 1]
+        return target
+    if deriv in multi_indices_2d(p):
+        pos = multi_indices_2d(p).index(deriv)
+        target[pos] = float(math.factorial(deriv[0]) * math.factorial(deriv[1]))
+    return target
+
+
+def target_1d(p: int, deriv: int = 0) -> np.ndarray:
+    if deriv > p:
+        raise ValueError("deriv must be less than or equal to p.")
+    target = np.zeros(p + 1)
+    target[deriv] = float(math.factorial(deriv))
+    return target
+
+
+def weighted_pinv_design(design: np.ndarray, weights: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    keep = weights > 0
+    X = design[keep, :]
+    w = weights[keep]
+    if X.shape[0] == 0:
+        raise ValueError("No observations inside the bandwidth.")
+    WX = X * w[:, None]
+    gram = X.T @ WX
+    inv_gram = np.linalg.pinv(gram, rcond=1e-12)
+    return keep, X, inv_gram
+
+
+@dataclass
+class LocalFit:
+    estimate: np.ndarray
+    se: np.ndarray
+    influence: np.ndarray
+    n_eff: int
+
+
+def local_fit_targets(
+    design_full: np.ndarray,
+    weights_full: np.ndarray,
+    outcomes_full: np.ndarray,
+    target: np.ndarray,
+    *,
+    vce: str = "hc1",
+    cluster: np.ndarray | None = None,
+) -> LocalFit:
+    outcomes_full = np.asarray(outcomes_full, dtype=float)
+    if outcomes_full.ndim == 1:
+        outcomes_full = outcomes_full.reshape(-1, 1)
+
+    keep, X, inv_gram = weighted_pinv_design(design_full, weights_full)
+    w = weights_full[keep]
+    Y = outcomes_full[keep, :]
+    beta = inv_gram @ (X.T @ (w[:, None] * Y))
+    fitted = X @ beta
+    resid = Y - fitted
+    n_eff = X.shape[0]
+    k = X.shape[1]
+
+    leverage = np.sum((X @ inv_gram) * X, axis=1) * w
+    adj = np.ones(n_eff)
+    vce = vce.lower()
+    if vce == "hc2":
+        adj = 1.0 / np.maximum(1.0 - leverage, 1e-8)
+    elif vce == "hc3":
+        adj = 1.0 / np.maximum(1.0 - leverage, 1e-8) ** 2
+
+    # Influence contribution for the requested linear functional.
+    row = target @ inv_gram
+    score_base = (X * w[:, None]) @ row
+    infl_kept = score_base[:, None] * resid * np.sqrt(adj)[:, None]
+
+    scale = 1.0
+    if vce == "hc1" and n_eff > k:
+        scale = n_eff / (n_eff - k)
+
+    if cluster is not None:
+        cluster_kept = np.asarray(cluster)[keep]
+        groups = pd.unique(cluster_kept)
+        summed = np.zeros((len(groups), outcomes_full.shape[1]))
+        for i, group in enumerate(groups):
+            summed[i, :] = np.sum(infl_kept[cluster_kept == group, :], axis=0)
+        if vce == "hc1" and len(groups) > 1 and n_eff > k:
+            scale = (len(groups) / (len(groups) - 1.0)) * ((n_eff - 1.0) / (n_eff - k))
+        cov = scale * (summed.T @ summed)
+        infl_source = summed
+    else:
+        cov = scale * (infl_kept.T @ infl_kept)
+        infl_source = np.sqrt(scale) * infl_kept
+
+    estimate = target @ beta
+    se = np.sqrt(np.maximum(np.diag(cov), 0.0))
+
+    infl_full = np.zeros((design_full.shape[0], outcomes_full.shape[1]))
+    if cluster is None:
+        infl_full[keep, :] = infl_source
+    else:
+        # For cross-evaluation covariance, keep cluster-level sums in rows
+        # matching the first occurrence of each cluster. This preserves sums
+        # without needing a second representation.
+        cluster_all = np.asarray(cluster)
+        groups = pd.unique(cluster_all)
+        infl_full = np.zeros((len(groups), outcomes_full.shape[1]))
+        kept_groups = pd.unique(np.asarray(cluster)[keep])
+        group_to_row = {g: i for i, g in enumerate(groups)}
+        for j, g in enumerate(kept_groups):
+            infl_full[group_to_row[g], :] = infl_source[j, :]
+
+    return LocalFit(estimate=np.asarray(estimate), se=se, influence=infl_full, n_eff=int(n_eff))
+
+
+def ci_columns(est: np.ndarray, se: np.ndarray, level: float, side: str) -> tuple[np.ndarray, np.ndarray]:
+    if side == "two":
+        cval = stats.norm.ppf((level + 100.0) / 200.0)
+        return est - cval * se, est + cval * se
+    cval = stats.norm.ppf(level / 100.0)
+    if side == "left":
+        return np.repeat(-np.inf, len(est)), est + cval * se
+    if side == "right":
+        return est - cval * se, np.repeat(np.inf, len(est))
+    raise ValueError("side must be two, left, or right.")
+
+
+def p_values(tvalues: np.ndarray) -> np.ndarray:
+    return 2.0 * stats.norm.sf(np.abs(tvalues))
+
+
+def bandwidth_floor(values: np.ndarray, bwcheck: int | None) -> float:
+    values = np.sort(np.asarray(values, dtype=float)[np.isfinite(values)])
+    if values.size == 0:
+        return np.nan
+    if bwcheck is None:
+        return 0.0
+    k = min(max(int(bwcheck), 1), values.size)
+    return float(values[k - 1])
+
+
+def cer_factor(n: int, p: int) -> float:
+    n = max(int(n), 1)
+    return n ** (1.0 / (2.0 * p + 4.0) - 1.0 / (p + 4.0))
+
+
+def ensure_dataframe(x: np.ndarray, columns: list[str]) -> pd.DataFrame:
+    return pd.DataFrame(x, columns=columns)
+
+
+def covariance_from_influence(influence: np.ndarray) -> np.ndarray:
+    influence = np.asarray(influence, dtype=float)
+    return influence @ influence.T