zebra-open 0.1.0__py3-none-any.whl

zebra_open/__init__.py

@@ -0,0 +1,33 @@
+"""zebra_open
+
+Research-grade reference implementation of a ZeBRA-style retrospective risk modeling framework.
+
+The CLI entry points are the recommended interface. Core functions are also
+importable for scripting and integration.
+"""
+
+from .generator import AgeSpec, GenParams, generate_dx_cohort, write_manifest, write_patients_jsonl
+from .cohort import summarize_file, summarize_patients, summarize_aggregate
+from .retro import (
+    train_model_from_files,
+    eval_retrospective_from_files,
+    score_patients_from_files,
+    load_model,
+)
+from .model import ZebraModel
+
+__all__ = [
+    "AgeSpec",
+    "GenParams",
+    "generate_dx_cohort",
+    "write_patients_jsonl",
+    "write_manifest",
+    "summarize_file",
+    "summarize_patients",
+    "summarize_aggregate",
+    "train_model_from_files",
+    "eval_retrospective_from_files",
+    "score_patients_from_files",
+    "load_model",
+    "ZebraModel",
+]

zebra_open/cli.py

@@ -0,0 +1,225 @@
+#!/usr/bin/env python3
+"""zebra_open.cli
+
+Console entry points:
+  zebra-generate : generate a synthetic DX-only cohort (research use)
+  zebra-cohort   : summarize cohort (QC)
+  zebra-train    : train a retrospective ZeBRA model
+  zebra-eval     : retrospective evaluation (AUC + optional calibration)
+  zebra-score    : prospective-style scoring (no labels)
+"""
+
+import argparse
+import json
+import os
+from pathlib import Path
+
+from .model import ZebraModel
+from . import retro
+from .generator import GenParams, parse_yyyy_mm_dd, is_valid_icd10cm, generate_dx_cohort, write_patients_jsonl, write_manifest
+from .cohort import summarize_file, summarize_aggregate
+
+
+def main_generate():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--out_dir", required=True, help="Output directory (will be created).")
+    ap.add_argument("--n_patients", type=int, default=10_000)
+    ap.add_argument("--target", required=True, help="Target ICD-10-CM code (e.g., J84.11).")
+    ap.add_argument("--case_frac", type=float, default=0.10)
+    ap.add_argument("--seed", type=int, default=13)
+    ap.add_argument("--start_date", default="2018-01-01", help="YYYY-MM-DD")
+    ap.add_argument("--end_date", default="2025-12-31", help="YYYY-MM-DD")
+
+    # Utilization controls (encounters/year)
+    ap.add_argument("--min_encounters_per_year", type=int, default=2)
+    ap.add_argument("--max_encounters_per_year", type=int, default=12)
+
+    # Chronic/problem list controls
+    ap.add_argument("--min_problem_list", type=int, default=0)
+    ap.add_argument("--max_problem_list", type=int, default=5)
+
+    # Target placement window (days from patient_start)
+    ap.add_argument("--min_days_to_target", type=int, default=30)
+    ap.add_argument("--max_days_to_target", type=int, default=2000)
+
+
+    # Target insertion controls
+    ap.add_argument("--target_min_occ", type=int, default=1)
+    ap.add_argument("--target_max_occ", type=int, default=3)
+    ap.add_argument("--target_last_third", action="store_true")
+
+    # Missingness controls
+    ap.add_argument("--missingness_encounter_drop", type=float, default=0.04)
+    ap.add_argument("--missingness_code_drop", type=float, default=0.06)
+
+    # Optional OpenAI enrichment
+    ap.add_argument("--enable_private_llm_pack", action="store_true")
+    ap.add_argument("--openai_model", default="gpt-5-nano")
+    ap.add_argument("--openai_timeout_s", type=int, default=30)
+    ap.add_argument("--openai_max_output_tokens", type=int, default=None)
+
+    args = ap.parse_args()
+
+    outdir = Path(args.out_dir)
+    outdir.mkdir(parents=True, exist_ok=True)
+
+    from .generator import GenParams, parse_yyyy_mm_dd, generate_dx_cohort, write_patients_jsonl, write_manifest, is_valid_icd10cm
+
+    target = str(args.target).strip().upper()
+    if not is_valid_icd10cm(target):
+        raise ValueError(f"--target {target!r} does not look like a valid ICD-10-CM code string.")
+
+    params = GenParams(
+        n_patients=int(args.n_patients),
+        case_frac=float(args.case_frac),
+        target_code=target,
+        start_date=parse_yyyy_mm_dd(args.start_date),
+        end_date=parse_yyyy_mm_dd(args.end_date),
+        seed=int(args.seed),
+        output=str(outdir / "patients.jsonl"),
+        jsonl=True,
+        target_min_occurrences=int(args.target_min_occ),
+        target_max_occurrences=int(args.target_max_occ),
+        target_last_third=bool(args.target_last_third),
+        missingness_encounter_drop=float(args.missingness_encounter_drop),
+        missingness_code_drop=float(args.missingness_code_drop),
+        min_encounters_per_year=int(args.min_encounters_per_year),
+        max_encounters_per_year=int(args.max_encounters_per_year),
+        min_problem_list=int(args.min_problem_list),
+        max_problem_list=int(args.max_problem_list),
+        min_days_to_target=int(args.min_days_to_target),
+        max_days_to_target=int(args.max_days_to_target),
+    )
+
+    patients, manifest = generate_dx_cohort(
+        params,
+        enable_openai_pack=bool(args.enable_private_llm_pack),
+        openai_model=str(args.openai_model),
+        openai_timeout_s=int(args.openai_timeout_s),
+        openai_max_output_tokens=(int(args.openai_max_output_tokens) if args.openai_max_output_tokens is not None else None),
+    )
+
+    patients_path = outdir / "patients.jsonl"
+    manifest_path = outdir / "cohort_manifest.json"
+
+    write_patients_jsonl(patients, str(patients_path))
+    write_manifest(manifest, str(manifest_path))
+
+    print(str(patients_path))
+    print(str(manifest_path))
+    return 0
+
+def main_cohort():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--patients", required=True, help="Patients JSON or JSONL path.")
+    ap.add_argument("--target_prefix", required=True, help="Prefix to match (e.g., J84.11).")
+    ap.add_argument("--out", required=True, help="Output CSV (or .parquet) summary file.")
+    ap.add_argument("--aggregate_json", default=None, help="Optional aggregate summary JSON output.")
+    args = ap.parse_args()
+
+    df = summarize_file(args.patients, target_prefix=args.target_prefix, out_path=args.out)
+    agg = summarize_aggregate(df)
+
+    if args.aggregate_json:
+        with open(args.aggregate_json, "w", encoding="utf-8") as f:
+            json.dump(agg, f, indent=2)
+
+    print(args.out)
+    if args.aggregate_json:
+        print(args.aggregate_json)
+    return 0
+
+
+def main_train():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--patients", required=True)
+    ap.add_argument("--target", required=True, help="Comma-separated target codes or single code.")
+    ap.add_argument("--out", required=True, help="Output directory.")
+    ap.add_argument("--observation_days", type=int, default=365)
+    ap.add_argument("--horizon_days", type=int, default=28)
+    ap.add_argument("--prediction_days", type=int, default=365)
+    ap.add_argument("--confidence_days", type=int, default=365)
+    ap.add_argument("--or_infer_frac", type=float, default=0.60)
+    ap.add_argument("--val_frac", type=float, default=0.34)
+    ap.add_argument("--random_state", type=int, default=0)
+    args = ap.parse_args()
+
+    outdir = Path(args.out)
+    outdir.mkdir(parents=True, exist_ok=True)
+
+    targets = [t.strip() for t in args.target.split(",") if t.strip()]
+
+    m = ZebraModel.from_config(
+        observation_days=args.observation_days,
+        horizon_days=args.horizon_days,
+        prediction_days=args.prediction_days,
+        confidence_days=args.confidence_days,
+        or_infer_frac=args.or_infer_frac,
+        val_frac=args.val_frac,
+        random_state=args.random_state,
+    )
+    m.fit(patients_json=args.patients, target_codes=targets, out_dir=str(outdir))
+    return 0
+
+
+def main_score():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", required=True, help="Path to zebra_retro_model.joblib.")
+    ap.add_argument("--patients", required=True, help="Path to patients JSON or JSONL.")
+    ap.add_argument("--out", required=True, help="Output CSV path.")
+    ap.add_argument("--as_of", default=None, help="Optional YYYY-MM-DD as-of date for all patients.")
+    args = ap.parse_args()
+
+    m = ZebraModel.load(args.model)
+    m.predict_proba(args.patients, as_of=args.as_of, out_csv=args.out)
+    return 0
+
+
+def main_eval():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", required=True, help="Path to zebra_retro_model.joblib.")
+    ap.add_argument("--patients", required=True, help="Path to patients JSON or JSONL to evaluate.")
+    ap.add_argument("--out", required=True, help="Output CSV path with columns patient_id,t0,y,p.")
+    ap.add_argument("--summary", default=None, help="Optional JSON summary path (includes AUC when computable).")
+    ap.add_argument("--calibration_bins", type=int, default=10, help="Number of bins for calibration curve.")
+    ap.add_argument("--calibration_csv", default=None, help="Optional CSV path for binned calibration stats.")
+    ap.add_argument("--calibration_plot", default=None, help="Optional PNG path for a reliability diagram.")
+    ap.add_argument("--target", default=None, help="Optional target code(s) override. Comma-separated or single.")
+    ap.add_argument("--max_controls", type=int, default=None, help="Optional control downsampling for evaluation.")
+    args = ap.parse_args()
+
+    target_codes = None
+    if args.target:
+        target_codes = [t.strip() for t in args.target.split(",") if t.strip()]
+
+    # If --out is a directory, write default artifacts inside it.
+    out_csv = args.out
+    out_is_dir = out_csv.endswith(os.sep) or (os.path.exists(out_csv) and os.path.isdir(out_csv))
+    if out_is_dir:
+        os.makedirs(out_csv, exist_ok=True)
+        out_csv = os.path.join(out_csv, 'eval.csv')
+
+    summary_path = args.summary
+    if summary_path is None and out_is_dir:
+        summary_path = os.path.join(args.out, 'eval_summary.json')
+
+    cal_csv = args.calibration_csv
+    cal_plot = args.calibration_plot
+    if out_is_dir:
+        if cal_csv is None:
+            cal_csv = os.path.join(args.out, 'calibration.csv')
+        if cal_plot is None:
+            cal_plot = os.path.join(args.out, 'calibration.png')
+
+    retro.eval_retrospective_from_files(
+        model_path=args.model,
+        patients_json=args.patients,
+        target_codes=target_codes,
+        out_csv=out_csv,
+        out_summary_json=summary_path,
+        calibration_bins=int(args.calibration_bins),
+        calibration_csv=cal_csv,
+        calibration_plot=cal_plot,
+        max_controls=args.max_controls,
+    )
+    return 0

zebra_open/cohort.py

@@ -0,0 +1,174 @@
+#!/usr/bin/env python3
+"""
+zebra_open.cohort
+
+Cohort summarization utilities.
+
+This module is derived from code/cohort_visualizer.py and provides reusable
+functions to summarize patient DX history and target-prefix presence.
+
+It does not perform ZeBRA retrospective window labeling. It is intended for
+data sanity checks, dataset characterization, and quick QC before model runs.
+"""
+
+from __future__ import annotations
+
+from datetime import date, datetime
+from typing import Any, Dict, Iterable, List, Optional, Tuple
+import json
+
+import pandas as pd
+
+DATE_FMT = "%m/%d/%Y"
+
+
+def parse_mmddyyyy(s: str) -> date:
+    return datetime.strptime(s, DATE_FMT).date()
+
+
+def norm_code(code: str) -> str:
+    return (code or "").strip().upper()
+
+
+def code_matches_prefix(code: str, target_prefix: str) -> bool:
+    c = norm_code(code)
+    t = norm_code(target_prefix)
+    return c.startswith(t) if t else False
+
+
+def load_patients(path: str) -> List[Dict[str, Any]]:
+    with open(path, "r", encoding="utf-8") as f:
+        first = f.read(1)
+        f.seek(0)
+        if first == "[":
+            return json.load(f)
+        out = []
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            out.append(json.loads(line))
+        return out
+
+
+def extract_dx_records(p: Dict[str, Any]) -> List[Tuple[date, str]]:
+    dx = p.get("DX_record", []) or []
+    out: List[Tuple[date, str]] = []
+    for r in dx:
+        d = r.get("date")
+        c = r.get("code")
+        if not d or not c:
+            continue
+        try:
+            out.append((parse_mmddyyyy(d), norm_code(c)))
+        except Exception:
+            continue
+    out.sort(key=lambda t: t[0])
+    return out
+
+
+def summarize_patient(p: Dict[str, Any], target_prefix: str) -> Dict[str, Any]:
+    pid = p.get("patient_id", p.get("pid", p.get("id", None)))
+    pid = "" if pid is None else str(pid)
+
+    dx = extract_dx_records(p)
+    n_dx = len(dx)
+    uniq = len(set(c for _, c in dx))
+
+    if n_dx == 0:
+        return {
+            "patient_id": pid,
+            "target_prefix": target_prefix,
+            "target": 0,
+            "n_dx_records": 0,
+            "n_unique_dx_codes": 0,
+            "first_observed_date": None,
+            "last_observed_date": None,
+            "duration_days": None,
+            "first_target_date": None,
+            "days_before_first_target": None,
+            "days_after_first_target": None,
+        }
+
+    first_d = dx[0][0]
+    last_d = dx[-1][0]
+    duration_days = (last_d - first_d).days
+
+    first_t: Optional[date] = None
+    for d, c in dx:
+        if code_matches_prefix(c, target_prefix):
+            first_t = d
+            break
+
+    is_case = 1 if first_t is not None else 0
+
+    if first_t is None:
+        return {
+            "patient_id": pid,
+            "target_prefix": target_prefix,
+            "target": 0,
+            "n_dx_records": n_dx,
+            "n_unique_dx_codes": uniq,
+            "first_observed_date": first_d.strftime(DATE_FMT),
+            "last_observed_date": last_d.strftime(DATE_FMT),
+            "duration_days": duration_days,
+            "first_target_date": None,
+            "days_before_first_target": None,
+            "days_after_first_target": None,
+        }
+
+    return {
+        "patient_id": pid,
+        "target_prefix": target_prefix,
+        "target": is_case,
+        "n_dx_records": n_dx,
+        "n_unique_dx_codes": uniq,
+        "first_observed_date": first_d.strftime(DATE_FMT),
+        "last_observed_date": last_d.strftime(DATE_FMT),
+        "duration_days": duration_days,
+        "first_target_date": first_t.strftime(DATE_FMT),
+        "days_before_first_target": (first_t - first_d).days,
+        "days_after_first_target": (last_d - first_t).days,
+    }
+
+
+def summarize_patients(patients: Iterable[Dict[str, Any]], target_prefix: str) -> pd.DataFrame:
+    rows = [summarize_patient(p, target_prefix=target_prefix) for p in patients]
+    return pd.DataFrame(rows)
+
+
+def summarize_file(in_path: str, target_prefix: str, out_path: Optional[str] = None) -> pd.DataFrame:
+    pats = load_patients(in_path)
+    df = summarize_patients(pats, target_prefix=target_prefix)
+    if out_path:
+        if out_path.endswith(".parquet"):
+            df.to_parquet(out_path, index=False)
+        else:
+            df.to_csv(out_path, index=False)
+    return df
+
+
+def summarize_aggregate(df: pd.DataFrame) -> Dict[str, Any]:
+    # Basic aggregate stats for quick QC.
+    n = int(df.shape[0])
+    n_cases = int(df["target"].sum()) if "target" in df.columns else 0
+    n_controls = int(n - n_cases)
+    out: Dict[str, Any] = {"n": n, "cases": n_cases, "controls": n_controls}
+
+    if "duration_days" in df.columns:
+        d = pd.to_numeric(df["duration_days"], errors="coerce")
+        out["duration_days"] = {
+            "n_finite": int(d.notna().sum()),
+            "min": float(d.min()) if d.notna().any() else None,
+            "median": float(d.median()) if d.notna().any() else None,
+            "max": float(d.max()) if d.notna().any() else None,
+        }
+
+    if "n_dx_records" in df.columns:
+        x = pd.to_numeric(df["n_dx_records"], errors="coerce")
+        out["n_dx_records"] = {
+            "min": float(x.min()) if x.notna().any() else None,
+            "median": float(x.median()) if x.notna().any() else None,
+            "max": float(x.max()) if x.notna().any() else None,
+        }
+    return out

zebra_open/examples/README.md

@@ -0,0 +1,36 @@
+# ZeBRA Open examples
+
+This directory is distributed with the `zebra-open` Python package.
+
+## Quickstart notebook
+
+`quickstart.ipynb` demonstrates an end-to-end ZeBRA Open workflow using only synthetic data:
+
+1. generate a longitudinal synthetic cohort;
+2. inspect the generated cohort;
+3. train a retrospective ZeBRA model;
+4. evaluate discrimination and calibration; and
+5. score patients with the fitted model.
+
+No clinical or private data are required.
+
+When working from a clone of the repository, install the package first:
+
+```bash
+python3 -m pip install -e ./python
+```
+
+After a PyPI release, the corresponding command is:
+
+```bash
+python3 -m pip install zebra-open
+```
+
+The installed notebook can also be located programmatically:
+
+```python
+from importlib.resources import files
+
+notebook = files("zebra_open").joinpath("examples", "quickstart.ipynb")
+print(notebook)
+```