How It Works
The DAB pipeline runs five main stages for each outcome and seed combination. They are mostly sequential, with two branches: the held-out set splits off before matching, and --data_only stops the run after cohort construction.
Pipeline Overview
flowchart TD
IN[("Input CSV<br/>--data")] --> VAL{"validate_dataset_for_model"}
VAL -->|fails| EXIT(["SystemExit"])
VAL -->|passes| PRE["preprocess_merged_data<br/>encoding · TLFB features"]
PRE --> DROP["drop_other_outcomes<br/>keep only this run's endpoint"]
DROP --> SEED["seed set · random + numpy"]
SEED --> HOLD["holdOutTestData"]
HOLD --> HELD[/"Held-out set<br/>--heldout_size · stratified"/]
HOLD --> POOL[/"Training pool"/]
POOL --> PSM["propensityScoreMatch<br/>--split · --match · --group_size"]
PSM --> SUB["create_subsets<br/>11 cohorts · 0/1000 → 500/500 ladder"]
SUB --> ONLY{"--data_only?"}
ONLY -->|yes| SAVE["save_model_input_datasets"]
SAVE --> DONE(["done"])
ONLY -->|no| DISP{"endpointType"}
DISP -->|logical| LOG["LogisticModel"]
DISP -->|integer| NB["NegativeBinomialModel"]
DISP -->|survival| COX["CoxProportionalHazard"]
LOG & NB & COX --> FIT["per subset × 11<br/>selectFeatures → train → evaluate"]
HELD --> FIT
FIT --> SP[/"subset_predictions"/]
FIT --> HP[/"heldout_predictions"/]
FIT --> SE[/"subset_evaluations"/]
FIT --> HE[/"heldout_evaluations"/]
SP & HP & SE & HE --> LOOP(["repeat over seeds × outcomes"])
classDef gate fill:none,stroke:#d97706,stroke-width:1.5px;
classDef data fill:none,stroke:#0891b2,stroke-width:1.5px;
classDef term fill:none,stroke:#6b7280,stroke-width:1.5px;
class VAL,ONLY,DISP gate;
class HELD,POOL,SP,HP,SE,HE data;
class EXIT,DONE,LOOP term;
Every path above runs once per (outcome, seed) pair. Two edges are easy to miss in a
step-by-step reading:
Structure worth noting
The held-out set is carved out before matching, so it never passes through PSM or the
subset ladder and stays untouched until evaluation. And --data_only exits after cohort
construction, before any model is fit.
Step 1 — Validation & Preprocessing
The pipeline validates the input CSV against expected schemas before any modeling occurs. It checks:
- Required columns are present
- Target column values match the endpoint type (binary 0/1, non-negative integer, positive duration)
- No duplicate patient IDs
Feature engineering includes binary encoding of categorical variables and construction of Treatment Lifecycle Feedback (TLFB) features from weekly urine drug screening data.
Step 2 — Propensity Score Matching (PSM)
PSM builds a ladder of training cohorts whose demographic composition changes while everything else about them stays as comparable as possible. It runs in two parts: matching, then the ladder itself.
Part 1 — Matching
Every row is flagged as minority or majority by the --split column: majority
means RaceEth == 1 (Non-Hispanic White, by default --majority 1), and every
other value is minority.
The data is handed to R's MatchIt via
rpy2, which fits a probit-link GLM propensity model on the --match covariates
(age and is_female by default) and performs optimal matching at a 1:2
ratio — each minority participant is paired with two majority controls who look
like them on those covariates. Only the first --group_size (default 500)
minority participants are matched.
The result is 500 matched triples, reshaped so that each row holds one treated participant and their two controls:
| treated_row | control_row_0 | control_row_1 | |
|---|---|---|---|
| triple 1 | minority participant | majority control | majority control |
| triple 2 | minority participant | majority control | majority control |
| … | … | … | … |
Those three columns become three DataFrames of 500 rows each — [minority,
majority_A, majority_B] — still aligned row by row, so row i of all three
frames belongs to the same matched triple. That alignment is what makes the next
part work.
Part 2 — The ladder
create_subsets builds 11 cohorts. Cohort k (counting from 0) takes:
- the first
k × 50rows of the minority frame, - the majority_A rows from
k × 50onward, - all 500 rows of majority_B.
So each step swaps 50 majority_A controls out for the 50 minority participants they were matched to. Because the swap happens inside matched triples, the cohort's age and sex profile stays balanced even as its racial composition changes — which is the whole point. majority_B is never touched; it is the constant backbone present in every cohort.
Total size stays fixed at 1000 on every rung:
| Cohort | Minority | Majority | Minority share |
|---|---|---|---|
| 1 | 0 | 1000 | 0% |
| 2 | 50 | 950 | 5% |
| 3 | 100 | 900 | 10% |
| 4 | 150 | 850 | 15% |
| 5 | 200 | 800 | 20% |
| 6 | 250 | 750 | 25% |
| 7 | 300 | 700 | 30% |
| 8 | 350 | 650 | 35% |
| 9 | 400 | 600 | 40% |
| 10 | 450 | 550 | 45% |
| 11 | 500 | 500 | 50% |
Every model is trained once per cohort, so each run produces 11 sets of metrics that can be read against the composition above.
The ladder tops out at 50/50
It does not sweep to an all-minority cohort. Because majority_B is always included in full, minority representation can never exceed half the cohort. The sweep is 0% → 50%, in 5-point steps.
How the parameters interact
--group_size sets both the matched group size and the step: each cohort
totals 2 × group_size, and the step is group_size ÷ 10. The rung count
(11) is fixed in create_subsets and is not exposed on the CLI. If
group_size is not divisible by 10, the final rung falls slightly short of a
balanced cohort — --group_size 505 still steps by 50 and stops at 500
minority rows.
The held-out set
A stratified held-out evaluation set is carved out before any of this, so it
never passes through matching or the ladder. It holds a fixed majority/minority
ratio (--heldout_set_percent_majority, default 58%) so evaluation demographics
stay constant no matter how the training cohort is rebalanced. It is also sampled
with a module-local seed of 42 rather than the run seed, so the same rows are
held out across a multi-seed run.
Step 3 — Feature Selection
L1 (Lasso) regularization is applied to automatically select predictive features before training. Features with zero coefficients after regularization are dropped. This reduces dimensionality and prevents overfitting on small cohorts.
Note
If Lasso drops all features (over-regularization), the pipeline raises an error rather than silently training on no signal.
Step 4 — Model Training
The model class is automatically selected based on the endpoint type of the chosen outcome:
| Endpoint Type | Model | Library |
|---|---|---|
logical — binary |
Logistic Regression (L1) | scikit-learn |
integer — count |
Negative Binomial Regression | statsmodels |
survival — time-to-event |
Cox Proportional Hazard | lifelines |
Each model is trained independently on each PSM cohort.
Step 5 — Evaluation
Each trained model is evaluated on two sets:
- Internal test split — 25% of the PSM cohort held out during training
- Held-out set — the fixed held-out set constructed in Step 2
Both evaluations record the demographic makeup of the training cohort alongside the metrics.
Metrics by Endpoint Type
| Metric | Description |
|---|---|
| ROC-AUC | Area under the receiver operating characteristic curve |
| Precision | True positives / (true positives + false positives) |
| Recall | True positives / (true positives + false negatives) |
| Confusion Matrix | Full 2×2 breakdown |
| Metric | Description |
|---|---|
| MSE | Mean squared error |
| RMSE | Root mean squared error |
| MAE | Mean absolute error |
| Pearson r | Linear correlation between predicted and actual |
| McFadden R² | Goodness-of-fit relative to null model |
| Metric | Description |
|---|---|
| Concordance Index | C-statistic: probability model ranks a random pair correctly |
Interpreting Results
The key question is: does model performance change as the demographic composition of the training cohort changes?
- No change across ratios → the outcome measure is measurement invariant — the model generalizes across groups equally
- Performance drops as minority proportion increases → the outcome is measurement variant — the model has learned patterns specific to the majority group
This framing is based on the measurement invariance framework from Odom et al. (2025).