ZeBRA_open

ZeBRA_open is a research-grade implementation of a retrospective risk modeling framework based on the ZeBRA architecture. It is designed for methodological research in longitudinal health record modeling, temporally correct cohort construction, odds-ratio embeddings, and stacked prediction.

Research Use Only
This repository is not a medical device and must not be used for diagnosis, treatment, or clinical decision making. See DISCLAIMER.md and CLINICAL_USE_NOTICE.md.

---

Installation

From the ZeBRA_open repository root:

python3 -m pip install -e ./python --no-cache-dir

Alternatively:

cd python
python3 -m pip install -e . --no-cache-dir

Verify CLI availability:

zebra-generate -h
zebra-cohort -h
zebra-train -h
zebra-eval -h
zebra-score -h

---

Quickstart (CLI)

1) Generate Synthetic Cohort

Basic generation (DX-only JSONL cohort):

zebra-generate   --out_dir data/generated/j8411_seed0   --n_patients 10000   --target J84.11   --case_frac 0.10   --seed 0   --start_date 2010-01-01   --end_date 2024-12-31

Generator controls that mirror the standalone scripts:

zebra-generate   --out_dir data/generated/j8411_seed0_v2   --n_patients 10000   --target J84.11   --case_frac 0.10   --seed 1   --start_date 2010-01-01   --end_date 2024-12-31   --min_encounters_per_year 2   --max_encounters_per_year 14   --min_problem_list 2   --max_problem_list 18   --min_days_to_target 30   --max_days_to_target 2000   --target_min_occ 1   --target_max_occ 5   --target_last_third   --missingness_encounter_drop 0.05   --missingness_code_drop 0.10

Optional: enable the private LLM enrichment pack (OpenAI Responses API). This is disabled by default and must be explicitly turned on:

export OPENAI_API_KEY="..."
zebra-generate   --out_dir data/generated/j8411_seed0_llm   --n_patients 2000   --target J84.11   --case_frac 0.10   --seed 2   --start_date 2010-01-01   --end_date 2024-12-31   --enable_private_llm_pack   --openai_model gpt-5-nano   --openai_timeout_s 45   --openai_max_output_tokens 1200

Outputs:

• patients.jsonl
• cohort_manifest.json

Notes:

• The generator simulates utilization (encounters per year), per-patient observation span, and per-encounter code sampling.
• Missingness can be applied both at the encounter level (dropping encounters) and within encounters (dropping codes).
• The target code is inserted for cases with controls on placement and repetition.

---

2) Summarize Cohort

zebra-cohort   --patients data/generated/j8411_seed0/patients.jsonl   --target_prefix J84.11   --out results/cohort_summary_j8411.csv   --aggregate_json results/cohort_summary_j8411_agg.json

---

3) Train Retrospective Model

zebra-train   --patients data/generated/j8411_seed0/patients.jsonl   --target J84.11   --out results/zebra_j8411_v1   --observation_days 730   --horizon_days 28   --prediction_days 365   --confidence_days 365   --or_infer_frac 0.8   --val_frac 0.34   --random_state 0

Primary artifact:

results/zebra_j8411_v1/zebra_retro_model.joblib

---

4) Retrospective Evaluation (AUC + Calibration)

zebra-eval   --model results/zebra_j8411_v1/zebra_retro_model.joblib   --patients data/generated/j8411_seed0/patients.jsonl   --out results/zebra_j8411_v1/eval.csv   --summary results/zebra_j8411_v1/eval_summary.json   --calibration_bins 10   --calibration_csv results/zebra_j8411_v1/calibration.csv   --calibration_plot results/zebra_j8411_v1/calibration.png

---

5) Prospective Scoring

zebra-score   --model results/zebra_j8411_v1/zebra_retro_model.joblib   --patients data/generated/j8411_seed0/patients.jsonl   --out results/zebra_j8411_v1/scores.csv   --as_of 2023-12-31

If --as_of is omitted, scoring uses each patient's last observed event date.

---

Python API Quickstart

Generate Cohort

from zebra_open.generator import (
    GenParams,
    generate_dx_cohort,
    write_patients_jsonl,
    write_manifest,
    parse_yyyy_mm_dd,
)

params = GenParams(
    n_patients=10000,
    case_frac=0.10,
    target_code="J84.11",
    start_date=parse_yyyy_mm_dd("2010-01-01"),
    end_date=parse_yyyy_mm_dd("2024-12-31"),
    seed=0,
    output="data/generated/j8411_seed0",
    jsonl=True,
    target_min_occurrences=1,
    target_max_occurrences=4,
    target_last_third=True,
    missingness_encounter_drop=0.05,
    missingness_code_drop=0.10,
)

patients, manifest = generate_dx_cohort(
    params,
    enable_openai_pack=False,
)

write_patients_jsonl(patients, "data/generated/j8411_seed0/patients.jsonl")
write_manifest(manifest, "data/generated/j8411_seed0/cohort_manifest.json")

If you want the OpenAI enrichment pack:

patients, manifest = generate_dx_cohort(
    params,
    enable_openai_pack=True,
    openai_model="gpt-5-nano",
    openai_timeout_s=45,
    openai_max_output_tokens=1200,
)

---

Summarize Cohort

from zebra_open.cohort import summarize_file, summarize_aggregate

df = summarize_file(
    "data/generated/j8411_seed0/patients.jsonl",
    target_prefix="J84.11",
    out_path="results/cohort_summary_j8411.csv",
)

print(summarize_aggregate(df))

---

Train Model

from zebra_open import ZebraModel

model = ZebraModel.from_config(
    observation_days=730,
    horizon_days=28,
    prediction_days=365,
    confidence_days=365,
    or_infer_frac=0.8,
    val_frac=0.34,
    random_state=0,
)

model.fit(
    patients_json="data/generated/j8411_seed0/patients.jsonl",
    target_codes=["J84.11"],
    out_dir="results/zebra_j8411_v1",
)

---

Input Format

Accepted formats:

• JSONL (one patient per line)
• JSON (list of patient objects)

Each patient must include one of:

patient_id
pid
id

DX records must follow:

{
  "date": "MM/DD/YYYY",
  "code": "J84.11"
}

---

Detailed CLI Option Reference

zebra-generate

Option	Description
--out_dir	Output directory
--n_patients	Total patients
--target	ICD code
--case_frac	Fraction of cases
--seed	Random seed
--start_date	YYYY-MM-DD
--end_date	YYYY-MM-DD
--min_encounters_per_year	Utilization lower bound
--max_encounters_per_year	Utilization upper bound
--min_problem_list	Min chronic codes
--max_problem_list	Max chronic codes
--min_days_to_target	Target placement min (days from start)
--max_days_to_target	Target placement max (days from start)
--target_min_occ	Minimum number of target insertions for cases
--target_max_occ	Maximum number of target insertions for cases
--target_last_third	If set, force target occurrences into the final third of the timeline
--missingness_encounter_drop	Probability of dropping an encounter (0 to 1)
--missingness_code_drop	Probability of dropping a code within an encounter (0 to 1)
--enable_private_llm_pack	Enable OpenAI enrichment pack (disabled by default)
--openai_model	OpenAI model name (used only with private pack)
--openai_timeout_s	OpenAI request timeout seconds
--openai_max_output_tokens	Cap on OpenAI output tokens

---

zebra-cohort

Option	Description
--patients	Patient file
--target_prefix	ICD prefix match
--out	Summary CSV/Parquet
--aggregate_json	Aggregate stats JSON

---

zebra-train

Option	Description
--patients	Input patient file
--target	ICD code(s)
--out	Output directory
--observation_days	Feature lookback window
--horizon_days	Gap before prediction window
--prediction_days	Event window defining case
--confidence_days	Required control follow-up
--or_infer_frac	OR embedding split
--val_frac	Validation split
--random_state	Random seed

---

zebra-eval

Option	Description
--model	Path to model
--patients	Patient file
--out	Eval CSV
--summary	JSON summary
--calibration_bins	Number of bins
--calibration_csv	Calibration CSV
--calibration_plot	Calibration PNG
--target	Override target
--max_controls	Control cap

---

zebra-score

Option	Description
--model	Path to model
--patients	Patient file
--out	Scores CSV
--as_of	YYYY-MM-DD cut date

---

Notes

• Training and evaluation use identical temporal window logic.
• Prospective scoring performs feature construction at an as-of time without labeling.
• See METHODS.md for full algorithmic details.