Skip to content

chunchiehfan/bigcbps

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

1 Commit
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

bigcbps

Scalable covariate-balancing propensity scores (CBPS) and IPW/AIPW estimators in R, designed to run a full causal-inference pipeline on 5M-row datasets within an 8 GB RAM ceiling on a single node.

Why it exists: valid causal inference at biobank and national-registry scale. Standard R implementations of covariate-balancing propensity scores struggle well before the million-row mark, and are often infeasible inside the hardware-constrained, locked-down environments that hold sensitive registry and electronic-health-record data. bigcbps streams over file-backed matrices to estimate treatment effects efficiently at the tens-of-millions-of-rows scale under a fixed memory budget, making covariate-balancing causal analysis practical for ultra-large EHR and national-registry cohorts.

The primary estimand is the marginal ATE E[Y(1) − Y(0)] via g-computation on a weighted outcome model, with stacked-sandwich + delta-method standard errors (IPW) or empirical influence-function variance (AIPW). The conditional treatment coefficient is a secondary output; for logit / NB / Cox it is not the marginal ATE.

For users who want matched samples instead of weighting, the package also ships a propensity-matching path (estimand: ATT) — see §2.d.

Status v1 + AIPW (Phases 0–9 DONE) + propensity matching (ATT) + parquet UX + auto-encoding + balance reporting + rank-deficient-covariate auto-drop + degenerate-input guards
Tests 248 test blocks (4,756 assertions): 0 FAIL / 29 env-gated skips; parity oracles: CBPS, WeightIt, MatchIt, sandwich::vcovCL, marginaleffects, survival::cch/coxph
Stress test N = 10M (5M per arm), 4 cells: 12 min total, peak RSS 2.5–2.7 GB (8 GB ceiling)
Validation MC coverage + scale grid: docs/validation/

Contents

  1. Installation
  2. Tutorial
  3. Conceptual framework
  4. Implementation schema
  5. Benchmark results
  6. Pointers

1. Installation

Requirements: R ≥ 4.2 and a C++ compiler (Xcode CLT on macOS, build-essential on Ubuntu — needed by bigmemory for the file-backed matrix layer).

One command

# install.packages("remotes")  # if you don't already have it

remotes::install_github(
  "chunchiehfan/bigcbps",
  dependencies   = TRUE,         # installs Imports + Suggests in one shot
  build_vignettes = TRUE
)

dependencies = TRUE pulls in:

  • Core (Imports — required by the package itself): bigmemory, biglm, arrow, RhpcBLASctl, generics. These are the streaming filebacked-matrix, parquet, and BLAS-pinning machinery.
  • Optional (Suggests — used by benchmarks, alternative SEs, doubly-robust paths, and the vignette): CBPS, WeightIt, MatchIt, PSweight, SuperLearner, MASS, survival, sandwich, survey, marginaleffects, cobalt, testthat, knitr, rmarkdown.

You don't need to install anything else by hand. The one command above is the whole setup.

From a local clone

If you already cloned the repo (or this is a private build):

remotes::install_local("bigcbps", dependencies = TRUE, build_vignettes = TRUE)
# or, equivalently from the shell at repo root:
#   R -e 'remotes::install_deps("bigcbps", dependencies = TRUE)'
#   R CMD INSTALL bigcbps

2. Tutorial

2.a Preparing the data

In practice your dataset will be on disk as parquet. To keep this tutorial reproducible we'll simulate one in memory, write it out as parquet, and then inspect it the same way you'd inspect a real file you've just received.

library(bigcbps)
library(arrow)
set.seed(1)

n <- 5000
df <- data.frame(
  trt      = rbinom(n, 1, plogis(rnorm(n))),               # treatment 0/1
  age      = rnorm(n, 50, 10),                              # continuous
  sex      = sample(c("F", "M"), n, replace = TRUE),        # binary char
  smoker   = sample(c(TRUE, FALSE), n, replace = TRUE),     # logical
  diabetes = sample(c("none", "type1", "type2"), n, replace = TRUE),  # 3-level char
  y        = rbinom(n, 1, 0.3)                              # binary outcome
)
write_parquet(df, "cohort.parquet")

Before writing your formula, run preview_parquet() to see what's in the file and whether each column will be usable by the pipeline:

info <- preview_parquet("cohort.parquet")

info$n_rows
#> [1] 5000

info$columns
#>       name   type bigcbps_usable
#> 1      trt  int32           TRUE
#> 2      age double           TRUE
#> 3      sex string           TRUE     <- character; auto-encoded to 0/1
#> 4   smoker   bool           TRUE
#> 5 diabetes string           TRUE     <- 3 levels; auto-encoded to 2 dummies
#> 6        y  int32           TRUE

info$head    # first 6 rows for a sanity check

bigcbps_usable = TRUE means the column will enter the pipeline either directly (numeric/bool) or after auto-encoding (character/factor — binary becomes 0/1, multi-level becomes k-1 dummy columns named <orig>_<level>). A FALSE flag means the column is a date, a timestamp, or has more than 20 unique values; bigcbps() will error with a specific actionable message if you reference it in the formula.

2.b Running the analysis

One call: parse the formula, stream the parquet into a file-backed matrix, fit CBPS, build IPW weights, fit the outcome model, compute the sandwich + delta-method SE, marginalize via g-computation.

res <- bigcbps(trt ~ age + sex + smoker + diabetes,
               outcome = "y",
               data    = "cohort.parquet",
               family  = "logit")

summary(res) returns everything at once — point estimate + CI, ESS, every gate's status, CBPS convergence diagnostics, identifying assumptions, and covariate balance:

summary(res)
#> bigcbps_result (family = logit, estimator = ipw)
#> Marginal ATE:     0.0421  (SE = 0.0227, 95% CI = [-0.0024, 0.0866])
#> Conditional coef: 0.384   (SE = 0.194,  CI = [0.003, 0.765])
#> ESS (post-trim):  0.871 of N
#>
#> Gates:
#>   cbps_cond    value=1.06       threshold=1e+10   PASS
#>   ess          value=0.871      threshold=0.5     PASS
#>   weight_tail  value=4.92       threshold=1000    PASS
#>   overlap      value=0.001      threshold=0.01    PASS
#>
#> Diagnostics:
#>   cbps_iter            3
#>   cbps_cond_num        1.062
#>   cbps_converged       TRUE
#>
#> Assumptions (v1):
#>   - consistency: Y = Y(T)
#>   - no unmeasured confounding: {Y(0), Y(1)} ⊥ T | X
#>   - overlap: 0 < P(T=1 | X) < 1
#>
#> Covariate balance (max |SMD| per covariate, post-weighting):
#>   All 3 covariates have |SMD| < 0.10 post-weighting.
#>
#>   Top 5 by |SMD| post-weighting:
#>    covariate original level smd_pre  smd_post abs_smd_post
#>          age      age        0.058 -0.0165       0.0165
#>          sex      sex        0.049 -0.0114       0.0114
#>       smoker   smoker        0.108 -0.0112       0.0112

Reading the output block by block:

  • Marginal ATE + CI — the primary causal estimand E[Y(1) − Y(0)], with a population-level influence-function SE built on the stacked M-estimation sandwich (it includes both the coefficient-estimation uncertainty and the covariate-averaging term; the fixed-X delta-method SE that marginaleffects reports is also available as res$marginal$se_fixed_x). This is what you report.
  • Conditional coef — the treatment coefficient inside the weighted outcome model (log-OR for logit, log-RR for NB, log-HR for Cox). It is not the marginal ATE on a non-identity link (non-collapsibility) and is shown as a secondary diagnostic only.
  • ESS (effective sample size) — fraction of N effectively contributing after percentile trimming and stabilization. Below 50% triggers the ess gate.
  • Gates — five identifiability checks (Hessian conditioning, ESS, weight tail, propensity overlap, NB θ stability). By default each warns and continues when its threshold is crossed, recording the failure in res$gate_results and showing a !! VALIDITY GATE FAILURE banner on print(res); arm a hard stop with gates = list(<gate> = "halt").
  • Diagnostics — CBPS convergence. cbps_cond_num (Hessian condition number κ) below ~1e3 means the propensity model is numerically well-behaved.
  • Assumptions — the three causal assumptions listed for the reader's reference. Gates check overlap directly; consistency and no-unmeasured- confounding are domain claims you make about your design.
  • Covariate balance — post-weighting standardized mean differences per covariate. |SMD| < 0.10 is the conventional "balanced" threshold; the printout shows the count above and the top 5 by |SMD|.

On messy real-world data (redundant or near-collinear covariates, missing cells), rank-deficient columns are dropped automatically and reported here as a Dropped covariates (rank-deficient): … line — see §2.f.

print(res) returns just the four-line headline if you want it brief. For programmatic access into reports (knitr::kable, gt, flextable, or any tidy/broom workflow):

tidy(res)        # one row per term: marginal_ate (+ conditional_coef)
glance(res)      # one-row diagnostics: family, estimator, n_obs, ess, cbps_iter
balance(res)     # full SMD table per covariate (aggregate = TRUE for per-variable)
coef(res); confint(res, level = 0.95)

Survival outcomes (Cox)

For a time-to-event outcome, the formula stays the same — the outcome argument becomes a list(time = ..., event = ...), the family is "cox", and horizon is the t* at which to evaluate the marginal survival difference:

# df_surv has the same covariates plus fup (follow-up time) and
# event (1 = event observed, 0 = censored).
write_parquet(df_surv, "cohort_surv.parquet")

res_cox <- bigcbps(trt ~ age + sex + smoker,
                   outcome = list(time = "fup", event = "event"),
                   data    = "cohort_surv.parquet",
                   family  = "cox",
                   horizon = 180)
summary(res_cox)
#> bigcbps_result (family = cox, estimator = ipw)
#> Marginal ATE:     0.00559  (SE = 0.0111, 95% CI = [-0.0161, 0.0273])
#> Conditional coef: -0.357   (SE = 0.159,  CI = [-0.669, -0.045])
#> ESS (post-trim):  0.990 of N
#> Horizon: 180
#>
#> Gates:
#>   cbps_cond    value=7.48e4    threshold=1e+10   PASS
#>   ess          value=0.990     threshold=0.5     PASS
#>   weight_tail  value=1.03      threshold=1000    PASS
#>   overlap      value=0.000     threshold=0.01    PASS
#>
#> Diagnostics:
#>   cbps_converged       TRUE
#>   cbps_iter            2
#>   cox_n_cc             3924
#>
#> Assumptions (v1):
#>   - SUTVA, ignorability, positivity
#>   - Proportional hazards holds conditional on X for the case-cohort
#>     sub-sample; censoring is non-informative.
#>
#> Covariate balance (max |SMD| per covariate, post-weighting):
#>   All 3 covariates have |SMD| < 0.10 post-weighting.

Three Cox-specific differences from the logit case:

  • "Marginal ATE" is the marginal survival-probability difference at t* = 180 (e.g. P(survive past 180 | T=1) − P(survive past 180 | T=0)), not a log hazard ratio. The HR is intentionally not the primary estimand — it is non-collapsible and not directly interpretable as a population-level contrast.
  • "Conditional coef" is the log-HR from the weighted Cox model. It is reported as a secondary diagnostic only.
  • cox_n_cc is the case-cohort sub-sample size used (Lin–Ying Horvitz-Thompson weighting: cases weight 1, non-case subcohort members 1/α; default 1% subcohort + all events, configurable via subcohort_frac). The Cox path is the only one that materializes rows into RAM — capped at ~200 k via Horvitz-Thompson weighted sub-sampling.

For negative-binomial counts use family = "nb" with outcome = "y_count"; for the doubly-robust variant of any family add estimator = "aipw".

2.c Step by step (and the FBM)

bigcbps() is a wrapper. Under the hood it calls a few exported steps, each of which can be invoked, repeated, or substituted on its own. The key intermediate is a file-backed matrix (FBM) — a numeric matrix stored on disk via bigmemory so it stays out of RAM. Ingest the parquet into an FBM once, then run the pipeline as many times as you want without re-ingesting:

# 1. Ingest parquet -> FBM on disk (the slow step at large N)
sim <- ingest_to_fbm(
  path           = "cohort.parquet",
  outdir         = "fbm_cache/cohort",      # writes cohort.bk + cohort.rds
  treatment_col  = "trt",
  outcome_spec   = list(type = "binary", y_col = "y"),
  covariate_cols = c("age", "sex", "smoker", "diabetes")
)
# sim$fbm  -- handle into the .bk file on disk
# sim$meta -- column indices + auto-encoding map

# 2. Run the full pipeline on that FBM
res_ipw <- run_pipeline(sim$fbm, sim$meta,
                        spec = list(family = "logit", estimator = "ipw"))

# 3. Same FBM, different spec -- no re-ingest needed
res_aipw <- run_pipeline(sim$fbm, sim$meta,
                         spec = list(family = "logit", estimator = "aipw"))

# 4. New R session later? Re-attach the FBM without re-ingesting:
sim <- attach_fbm("fbm_cache/cohort")
res  <- run_pipeline(sim$fbm, sim$meta, spec = list(family = "logit"))

At N ≥ 1M, the FBM-cache pattern saves the parquet → FBM conversion time (tens of seconds to a few minutes per re-ingest), so this is the right workflow for iterating on different model specs against a fixed dataset.

You can also drop a step deeper and call the streaming building blocks directly. Useful when you want to insert your own outcome model, recompute balance against a non-default weighting, or experiment with one phase in isolation:

# Streaming logistic warm-start (one FBM sweep) -> initial beta
beta0 <- streaming_biglogit_propensity(sim$fbm, sim$meta)$beta

# Streaming CBPS-ATE Fisher scoring -> fitted beta, propensities e_hat, kappa
cbps  <- streaming_cbps_ate(sim$fbm, sim$meta, beta0 = beta0)

# IPW weights with percentile trimming + stabilization
w     <- build_ipw_weights(sim$fbm, sim$meta, ehat = cbps$ehat)

# Per-covariate SMD pre and post weighting
streaming_smd(sim$fbm, sim$meta, weights = w$w_vec, ehat = cbps$ehat)$smd

For survival outcomes, the outcome_spec becomes a 2-column list and spec needs a horizon t* at which to evaluate the marginal survival-difference:

sim_surv <- ingest_to_fbm(
  path = "cohort_surv.parquet", outdir = "fbm_cache/surv",
  treatment_col  = "trt",
  outcome_spec   = list(type = "survival",
                         time_col = "fup", event_col = "event"),
  covariate_cols = c("age", "sex", "smoker")
)
res_cox <- run_pipeline(sim_surv$fbm, sim_surv$meta,
                        spec = list(family = "cox", horizon = 180))

The bigcbps() formula wrapper supports all of these data inputs and outcome families:

Input How to pass to bigcbps()
In-memory data.frame data = df
Parquet file or directory data = "cohort.parquet"
Pre-ingested FBM (cached) data = attach_fbm("fbm_cache/cohort")
Outcome family How to specify
Binary outcome = "y", family = "logit"
Count (negative binomial) outcome = "y_count", family = "nb"
Survival (marginal survival diff at t*) outcome = list(time = "fup", event = "event"), family = "cox", horizon = 180
Doubly-robust variant of any of the above add estimator = "aipw"

2.d Matched samples instead of weighting (ATT)

Everything above answers the ATE ("what if everyone vs no one were treated?") by weighting. Some analyses instead call for a matched sample — for interpretability, audit, or journal convention. The matching path answers the ATT ("what was the effect on those actually treated?") and reuses the same propensity machinery, so it scales to the same data sizes: matching happens on the propensity score itself, entirely in RAM, in seconds even at 5M rows.

The flow is four calls after ingest:

sim <- attach_fbm("fbm_cache/cohort")          # or ingest_to_fbm(...) once

# 1. Propensity scores (same streaming CBPS as the weighting path)
fit_p <- streaming_cbps_ate(sim$fbm, sim$meta)
Tr    <- as.integer(sim$fbm[, sim$meta$treatment_col])

# 2. Match: 1 nearest control per treated, no reuse, caliper 0.2 SD
m <- match_propensity(fit_p$ehat, Tr,
                      ratio   = 1,      # controls per treated
                      replace = FALSE,  # each control used at most once
                      caliper = 0.2)    # max distance, in SDs of logit(e)
m$info
#> $n_treated               2412   <- treated units in the data
#> $n_matched_treated       2350   <- got at least one control
#> $n_discarded_by_caliper    62   <- no control within the caliper
#> ...

# 3. Check balance on the matched sample (same SMD tool, match weights)
bal <- streaming_smd(sim$fbm, sim$meta,
                     weights = m$w_match, ehat = fit_p$ehat)
bal$smd                                       # smd_w = post-matching SMD

# 4. Outcome model on the matched sample + the ATT
mo  <- matched_outcome(sim$fbm, sim$meta, m, family = "logit")
att <- marginal_att(mo, sim$fbm, sim$meta, m)
att$tau; att$se; att$ci                       # the ATT and its 95% CI
att$conditional                               # log-OR + cluster-robust SE

Reading the results:

  • att$tau is the ATT — the average effect among the matched treated, by g-computation over their covariates. Its SE is a cluster-level influence function: outcome-model uncertainty enters through a cluster-robust-by-matched-set sandwich (the MatchIt-recommended convention), and the covariate-averaging term is included.
  • mo$vcov_cluster is that clustered covariance for the outcome-model coefficients; att$conditional reports the treatment coefficient with it.
  • m$w_match are MatchIt-style match weights (treated 1, control = times-used/ratio, unmatched 0) — they plug directly into streaming_smd or any weighted fit of your own.

Knobs on match_propensity:

Argument Meaning Default
ratio controls matched per treated (1:k) 1
replace controls reusable across treated? FALSE
caliper max match distance in SDs of logit(ê); NULL = none 0.2

All three outcome families work: family = "logit", "nb", or "cox" (Cox needs horizon = t* in the marginal_att call, and reports the treated-average survival difference at t*).

Three caveats (also in vignette("matching")):

  1. The estimand is the ATT, not the ATE — and with a caliper it narrows further to the ATT among matchable treated. Check m$info$n_discarded_by_caliper; a warning fires automatically above 5%.
  2. SEs treat the match as fixed. Propensity-estimation uncertainty does not propagate through the (non-smooth) matching step — the standard convention in applied practice.
  3. With replace = TRUE, cluster-by-set SEs ignore the positive correlation a reused control induces across sets and tend to understate variance when controls are scarce (Austin & Cafri 2020).

2.e Function reference

User-facing entry points (recommended):

Function What it does
bigcbps() One-call formula wrapper: ingest → CBPS → IPW → outcome → ATE
preview_parquet() Inspect a parquet file before writing your formula
balance() Per-covariate SMD; aggregate = TRUE for per-original-variable summary
ingest_to_fbm() Stream parquet → FBM on disk; auto-encodes categoricals
attach_fbm() Re-attach a previously-written FBM from .bk + .rds
run_pipeline() Lower-level orchestrator on a pre-built FBM + meta list
sim_cbps_truth() Known-truth fixture for testing / benchmarking

S3 methods on bigcbps_result:

Method What it does
print() Four-line headline: marginal ATE + CI + conditional + ESS
summary() Full diagnostics: gates, convergence, balance, assumptions
tidy() One row per term (broom-compatible)
glance() One-row model-level diagnostics
coef() Named numeric of point estimates
confint() 2-column CI matrix

Streaming building blocks (swap or repeat individual phases):

Function What it does
streaming_biglogit_propensity() Streaming logistic warm-start; one FBM sweep
streaming_cbps_ate() Streaming CBPS-ATE Fisher scoring; β, ê, p×p Hessian, κ; auto-drops rank-deficient covariates (drop_collinear, default on)
build_ipw_weights() ATE weights with percentile trimming + stabilization
streaming_smd() Per-covariate SMD pre/post + overlap gate
weighted_biglogit() Streaming weighted logistic GLM
weighted_bignb() Streaming weighted NB log-link GLM with profile MLE for θ
weighted_bigcox() Cox PH on case-cohort sample; coxph(robust = TRUE)
stacked_sandwich() Stacked M-estimation sandwich; one FBM sweep, DSYRK contract
marginal_ate_logit() / marginal_ate_nb() / marginal_ate_cox() G-computation + delta-method SE
aipw_logit() / aipw_nb() / aipw_cox() AIPW doubly-robust paths
case_cohort_sample() Lin–Ying case-cohort design with HT weights (Cox only)

Matching path (ATT — see §2.d):

Function What it does
match_propensity() Nearest-neighbor matching on logit(ê); 1:k, with/without replacement, caliper; all-in-RAM
matched_outcome() Outcome model (logit/NB/Cox) on the matched sample with cluster-robust-by-set SEs
marginal_att() ATT by g-computation over matched treated, cluster-level IF SE; Cox = survival diff at t*
pin_blas_single_thread() Pin BLAS to 1 thread inside parallel workers
fire_gate() Internal gate evaluator (warn-and-continue by default; "halt"/TRUE or "off"/FALSE via gates)

For details on FBM persistence at N ≥ 1M, categorical-encoding overrides, and per-gate overrides, see ?bigcbps and the package vignette (vignette("end-to-end", package = "bigcbps")).

2.f Robustness on real-world data

Real cohorts are messier than simulations — redundant or near-collinear covariates, the occasional missing cell, an outcome that would overflow a naive exp(Xβ). bigcbps is built to degrade gracefully rather than crash opaquely on any of these.

Rank-deficient covariates are dropped automatically (like lm/glm). Before fitting, the CBPS step streams the covariate Gram matrix and drops any column that is constant or a linear combination of columns already kept — the same aliasing behavior lm()/glm() apply, but computed without materializing the design matrix. This is what lets the pipeline run on a large, partly- redundant covariate set that would otherwise produce a singular Hessian. The dropped columns are reported on print(res) / summary(res) and recorded in res$dropped_covariates:

Dropped covariates (rank-deficient): region_west (collinear), visits_dup (zero variance)
bigcbps_result (family = nb, estimator = ipw)
Marginal ATE:  ...

Pass drop_collinear = FALSE (through bigcbps(...) or run_pipeline()'s spec) to keep every column instead; the cbps_cond gate will then flag the singularity (κ > 1e10) rather than silently dropping it. Because κ is computed on the internally standardized scale (see §3), a raw calendar-year or days-count column will not trip that gate — a cbps_cond failure means genuine near-collinearity, not mismatched units.

Degenerate inputs fail loud and early, never silently. A read-only validation sweep runs before any fitting and stops with an actionable, column-named error if a covariate contains missing/non-finite values, or if the treatment has only one arm (the ATE is then not estimable). Extreme linear predictors are clamped before exp() so a large covariate scale or count cannot overflow to NaN. As a final backstop, if any marginal estimate or SE comes back non-finite the run emits a warning and records it — the pipeline never hands back a silent NaN.

These guards are parity-preserving: they fire only on degenerate input and leave every legitimate fit bit-identical.


3. Conceptual framework

The estimand

The primary target is the marginal average treatment effect (ATE): τ = E[Y(1) − Y(0)], where Y(t) is the counterfactual outcome under treatment assignment T = t. Identification requires: (1) consistency (Y = Y(T)), (2) no unmeasured confounding ({Y(0),Y(1)} ⊥ T | X), and (3) overlap (0 < P(T=1 | X) < 1 almost surely).

CBPS-IPW

Imai & Ratkovic's covariate-balancing propensity score fits the propensity model not by maximum likelihood (which optimizes predictive accuracy) but by exactly satisfying the moment conditions E[(T − ê(X)) · X] = 0. This produces weights with strictly better finite-sample covariate balance than naive logistic-PS IPW. The resulting IPW estimator τ̂ = Σ w_i T_i Y_i / Σ w_i T_i − Σ w_i (1−T_i) Y_i / Σ w_i (1−T_i) (with stabilized, percentile-trimmed weights) is then fed into the outcome model for g-computation.

AIPW (doubly robust)

The augmented IPW estimator is consistent if either the propensity model or the outcome model is correctly specified — doubly robust in the Robins-Rotnitzky-Zhao sense. The AIPW path (estimator = "aipw") uses an empirical influence-function variance estimator. Phases 0–8 (IPW path) must be complete and validated before the AIPW path is invoked.

Streaming architecture (why this scales)

The CBPS Fisher scoring and the sandwich meat matrix both reduce to sufficient statistics of fixed size: a p × p Hessian and a p-vector score, accumulated as the file-backed matrix is swept in 250k–500k row chunks. The package never materializes the full design matrix in RAM. The two critical BLAS contracts:

  • CBPS Hessian: crossprod(X * sqrt(w)) (DSYRK). Never t(X) %*% diag(w) — that intermediate is p × c and blows the RAM budget by a factor of c/p.
  • Sandwich meat: crossprod(Psi) (DSYRK). Never Psi %*% t(Psi) — at c = 500k that would be a 2 TB matrix.

Marginal vs conditional

The primary output is the marginal ATE via g-computation (predict counterfactual outcomes under T=0 and T=1, average, take the difference). The conditional treatment coefficient (log-OR for logit, log-RR for NB, log-HR for Cox) is reported as a secondary and is not the marginal ATE due to non-collapsibility on non-identity link scales.

Sandwich SE

The CBPS and outcome-model estimating equations are stacked into ψ_i(θ); the variance is V = A⁻¹ B (A⁻¹)' with A := −E[∂ψ/∂θ] and B := E[ψ ψ']. Delta-method propagation through the marginalization functional yields the final SE on the marginal ATE.

Validity gates (warn-and-continue by default; halt opt-in)

Five gates protect identifiability. By default each gate warns and continues when its threshold is crossed — the warning carries the diagnostic detail and suggested fixes, the failure is recorded in res$gate_results, and print(res) opens with a !! VALIDITY GATE FAILURE banner. Arm a gate as a hard stop with gates = list(<gate> = "halt") (or TRUE); silence it with gates = list(<gate> = FALSE) (override is still logged):

Gate Threshold Stage
cbps_cond Hessian condition number κ > 1e10 CBPS scoring
κ is computed on the internally standardized scale, so it is invariant to covariate units (a raw calendar-year or days-count column will not trip it) and fires only on genuine near-collinearity. On failure the message reports the smallest singular values and the covariates implicated.
ess Post-trim ESS < 50% of N Weight construction
weight_tail max(w) / median(w) > 1000 Weight construction
overlap > 1% of rows with ê ≤ 0.01 or ê ≥ 0.99 Balance check
nb_theta NB θ oscillation > 10% after iter 10 NB outcome model

4. Implementation schema

Pipeline

parquet / data.frame
    │
    ▼
ingest_to_fbm()  ─────────►  FBM on disk + meta sidecar
    │
    ▼
streaming logistic warm-start  ──►  initial β
    │
    ▼
streaming CBPS Fisher scoring  ──►  fitted β + p×p Hessian
    │  [rank-deficient covariates auto-dropped first (drop_collinear)]
    │
    ▼
build_ipw_weights()  ──────────►  weights column in FBM
    │  [Validity gates (warn by default; halt opt-in): ESS, overlap, weight tail, CBPS condition]
    ▼
streaming outcome model (logit / NB / Cox)
    │
    ▼
stacked sandwich SE  (one additional FBM sweep, DSYRK contract)
    │
    ▼
marginal ATE via g-computation + delta-method SE
    │
    ▼
bigcbps_result   (S3: marginal, conditional, balance,
                  ehat, diagnostics, gate_results, encoding)

(See §2.e Function reference for the function tables.)

File layout

bigcbps/
├── R/
│   ├── bigcbps.R           # formula wrapper (bigcbps())
│   ├── run_pipeline.R      # end-to-end orchestrator
│   ├── ingest.R            # ingest_to_fbm(), attach_fbm()
│   ├── encode.R            # categorical auto-encoding
│   ├── preview.R           # preview_parquet()
│   ├── propensity_logistic.R  # streaming warm-start logistic
│   ├── propensity_cbps.R   # streaming CBPS Fisher scoring
│   ├── weights.R           # build_ipw_weights()
│   ├── balance.R           # balance(), streaming_smd()
│   ├── outcome_logit.R     # weighted_biglogit()
│   ├── outcome_nb.R        # weighted_bignb()
│   ├── outcome_cox.R       # weighted_bigcox(), case_cohort_sample()
│   ├── sandwich.R          # stacked_sandwich()
│   ├── marginal_ate.R      # marginal_ate_logit/nb/cox()
│   ├── aipw.R              # aipw_logit/nb/cox()
│   ├── tidiers.R           # tidy/glance/coef/confint methods
│   ├── bigcbps_result.R    # print/summary methods
│   ├── gates.R             # fire_gate()
│   ├── dgp.R               # sim_cbps_truth()
│   └── utils.R             # pin_blas_single_thread(), helpers
├── inst/
│   ├── benchmarks/
│   │   └── realistic/      # 10-rep × N=75k head-to-head benchmark
│   └── simdata/            # pre-generated fixture RDS files
└── tests/testthat/
    ├── helper-fixtures.R   # Tier 0/1/2 fixture builder (N=1k/10k/100k)
    ├── test-parity-*.R     # numerical parity vs reference packages
    ├── test-gates.R        # adversarial gate-halt fixtures
    ├── test-memory-budget.R  # peak RAM contracts (bench-gated)
    ├── test-recovery.R     # 200-rep MC coverage (bench-gated)
    └── test-*.R            # unit tests per module

5. Benchmark results

Head-to-head: 10 reps × N = 75k (realistic three-arm DGP)

Design: Three-arm DGP with continuous and binary covariates (p = 15). Pairwise comparisons: A (G1 vs G0, true RR = 1.10) and B (G2 vs G0, true RR = 1.05). 10 Monte Carlo reps, N = 75k each, 9 estimators × 2 families (logit and Cox) × 2 comparisons = 36 cells.

Comparison A / logit (headline cell):

Estimator mean(τ̂) ± MC-SE bias ± MC-SE mean(SE) SE cal. ratio Coverage wall (s) peak (MiB)
bigcbps_ipw 0.0149 (± 0.0007) −0.0001 (± 0.0007) 0.0033 1.55 1.00 1.85 281
bigcbps_aipw 0.0150 (± 0.0007) −0.0000 (± 0.0007) 0.0033 1.51 1.00 1.13 287
cbps 0.0150 (± 0.0007) −0.0000 (± 0.0007) 0.5020 228.0 † 1.00 65.85 511
weightit 0.0150 (± 0.0007) −0.0000 (± 0.0007) 0.0022 1.02 1.00 27.02 433
matchit 0.0151 (± 0.0007) +0.0001 (± 0.0007) 0.0033 1.53 1.00 7.71 464
psweight 0.0150 (± 0.0007) −0.0000 (± 0.0007) 0.0033 1.50 1.00 12.58 426
naive_ipw 0.0150 (± 0.0007) −0.0000 (± 0.0007) 0.0022 1.02 1.00 3.41 437
svyglm 0.0150 (± 0.0007) −0.0000 (± 0.0007) 0.0033 1.51 1.00 6.09 440
aipw_pkg 0.0150 (± 0.0007) −0.0000 (± 0.0007) 0.0033 1.50 1.00 8.35 398

† The cbps row shows mean(SE) ≈ 0.50 (calibration ratio 228×). This was a known wrapper bug in est_cbps that inflated the reported SE by ~228×; the point estimate is correct. The bug has since been fixed (commit a1ab36e).

Reading the table: all estimators recover the true RD with near-zero bias and 100% coverage. bigcbps_ipw is the fastest (1.85 s) and the most memory-efficient (281 MiB) by a wide margin. The full 4-table report (two comparisons × two families) is at inst/benchmarks/realistic/results/REPORT_20260523.md.

Figures (Comparison A, N = 75k):

Forest plot, Comparison A (logit): every estimator recovers the true risk difference with near-zero bias; bigcbps_ipw and bigcbps_aipw are the fastest.

Love plot, Comparison A: all propensity-score estimators bring the maximum standardized mean difference below 0.10 after weighting.

Median wall time per estimator (logit and Cox combined): bigcbps is 4 to 50 times faster than the reference packages at N = 75k.

Stress test: N = 10M, p = 15

Single run with a binary treatment DGP (sim_cbps_truth, N = 5M per arm), all 4 cells: family (logit / Cox) × estimator (IPW / AIPW).

family estimator wall (s) peak RSS (MB) τ̂ SE true τ bias CI [95%]
logit ipw 232.5 2747.5 0.035410 0.000240 0.0356 −0.0002 [0.0349, 0.0359] COVERS
logit aipw 219.0 2659.0 0.035256 0.000259 0.0356 −0.0004 [0.0347, 0.0358] COVERS
cox ipw 143.8 2744.2 −0.097847 0.002229 NA NA [−0.1022, −0.0935]
cox aipw 140.7 2510.5 −0.094342 0.002750 NA NA [−0.0997, −0.0890]

Memory verdict: all 4 cells stay well under the 8 GB ceiling (peak RSS 2.5–2.7 GB, ~34% of the limit). The Cox path previously hit 14.9 GB before a case_cohort_sample() row-budget fix capped materialization at 200k rows via Horvitz-Thompson weighted sub-sampling.

Unbiasedness: logit CIs cover the known truth for both IPW and AIPW (bias < 1.4 SEs). Cox truth is unavailable (sim_cbps_truth returns tau_marginal = NA for survival — marginal RMST is a Phase 9 deferred estimand); Cox point estimates and SEs are numerically stable.

Backend: these numbers are from the bigmemory backend (replacing bigstatsr). Every τ̂ and SE is bit-identical to the prior bigstatsr run; wall time is +0-3 % and peak RSS is 5-12 % lower.

Wall time: 12 min total for all 4 cells at N = 10M (BLAS pinned to 1 thread). Full report at inst/benchmarks/realistic/results/STRESS_REPORT_20260527.md.


6. Pointers

Vignettes vignette("end-to-end", package = "bigcbps"), vignette("matching")
Benchmark code inst/benchmarks/realistic/
GitHub https://github.com/chunchiehfan/bigcbps
License MIT, see LICENSE

About

Scalable covariate-balancing propensity scores (CBPS) and IPW/AIPW causal-inference estimators in R: a full pipeline on tens of millions of rows within a fixed memory budget, for EHR and national-registry cohorts. Benchmarked against reference packages.

Topics

Resources

License

Stars

2 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages