sankar-raul/ICD-10-code-predictor-env

1"""Simulate incremental learning from diagnoses.csv.2 3The learner is intentionally simple and deterministic:4- processes visits in chronological order5- predicts primary ICD-10 from aggregate frequency signals6- updates itself after each labeled visit (online learning)7"""8 9from __future__ import annotations10 11import argparse12import csv13import math14import pickle15from collections import Counter, defaultdict16from dataclasses import dataclass17from datetime import date18from pathlib import Path19 20 21@dataclass(frozen=True)22class Visit:23    patient_id: str24    visit_date: date25    visit_type: str26    primary_icd10: str27    secondary_icd10s: tuple[str, ...]28    provider_specialty: str29 30 31def _parse_date(value: str) -> date:32    parts = value.split("-")33    return date(int(parts[0]), int(parts[1]), int(parts[2]))34 35 36def _split_codes(value: str) -> tuple[str, ...]:37    if not value:38        return ()39    return tuple(code.strip().upper() for code in value.split("|") if code.strip())40 41 42def load_visits(csv_path: Path) -> list[Visit]:43    visits: list[Visit] = []44    with csv_path.open("r", encoding="utf-8", newline="") as f:45        reader = csv.DictReader(f)46        for row in reader:47            visits.append(48                Visit(49                    patient_id=row["patient_id"],50                    visit_date=_parse_date(row["visit_date"]),51                    visit_type=row["visit_type"].strip().lower(),52                    primary_icd10=row["primary_icd10"].strip().upper(),53                    secondary_icd10s=_split_codes(row["secondary_icd10s"].strip()),54                    provider_specialty=row["provider_specialty"].strip().lower(),55                )56            )57    visits.sort(key=lambda v: (v.visit_date, v.patient_id))58    return visits59 60 61class OnlineDiagnosisLearner:62    """Online Naive Bayes learner with contextual and patient-history features."""63 64    def __init__(self) -> None:65        self.global_counts: Counter[str] = Counter()66        self.by_visit_type: dict[str, Counter[str]] = defaultdict(Counter)67        self.by_visit_type_total: Counter[str] = Counter()68        self.by_specialty: dict[str, Counter[str]] = defaultdict(Counter)69        self.by_specialty_total: Counter[str] = Counter()70        self.by_secondary: dict[str, Counter[str]] = defaultdict(Counter)71        self.by_secondary_total: Counter[str] = Counter()72        self.transitions: dict[str, Counter[str]] = defaultdict(Counter)73        self.transition_total: Counter[str] = Counter()74        self.patient_counts: dict[str, Counter[str]] = defaultdict(Counter)75        self.patient_total: Counter[str] = Counter()76        self.by_visit_specialty: dict[tuple[str, str], Counter[str]] = defaultdict(Counter)77        self.by_visit_specialty_total: Counter[tuple[str, str]] = Counter()78        self.by_exact_context: dict[tuple[str, str, tuple[str, ...]], Counter[str]] = defaultdict(Counter)79        self.by_exact_context_total: Counter[tuple[str, str, tuple[str, ...]]] = Counter()80        self.last_primary_by_patient: dict[str, str] = {}81        self.total_examples = 082 83        # Weights tuned for this deterministic simulation.84        self.w_prior = 0.785        self.w_visit_type = 1.586        self.w_specialty = 1.487        self.w_visit_specialty = 1.888        self.w_secondary = 0.889        self.w_transition = 1.290        self.w_patient = 2.291        self.w_exact_context = 2.092        self.alpha = 0.193 94    def predict(self, visit: Visit) -> str:95        if not self.global_counts:96            # Before any training data exists.97            return "I10"98 99        alpha = self.alpha100        classes = list(self.global_counts.keys())101        num_classes = len(classes)102        visit_specialty = (visit.visit_type, visit.provider_specialty)103        secondary_tuple = tuple(sorted(visit.secondary_icd10s))104        exact_context = (visit.visit_type, visit.provider_specialty, secondary_tuple)105        previous = self.last_primary_by_patient.get(visit.patient_id)106 107        scores: dict[str, float] = {}108        for icd in classes:109            # Prior prevalence.110            prior = math.log(111                (self.global_counts[icd] + alpha)112                / (self.total_examples + alpha * num_classes)113            )114            score = self.w_prior * prior115 116            # Visit-type likelihood.117            vt_total = self.by_visit_type_total[visit.visit_type]118            if vt_total:119                score += self.w_visit_type * math.log(120                    (self.by_visit_type[visit.visit_type][icd] + alpha)121                    / (vt_total + alpha * num_classes)122                )123 124            # Specialty likelihood.125            sp_total = self.by_specialty_total[visit.provider_specialty]126            if sp_total:127                score += self.w_specialty * math.log(128                    (self.by_specialty[visit.provider_specialty][icd] + alpha)129                    / (sp_total + alpha * num_classes)130                )131 132            # Joint visit-type + specialty likelihood.133            vs_total = self.by_visit_specialty_total[visit_specialty]134            if vs_total:135                score += self.w_visit_specialty * math.log(136                    (self.by_visit_specialty[visit_specialty][icd] + alpha)137                    / (vs_total + alpha * num_classes)138                )139 140            # Secondary-code signals.141            for sec in visit.secondary_icd10s:142                sec_total = self.by_secondary_total[sec]143                if sec_total:144                    score += self.w_secondary * math.log(145                        (self.by_secondary[sec][icd] + alpha)146                        / (sec_total + alpha * num_classes)147                    )148 149            # Patient transition signal.150            previous = self.last_primary_by_patient.get(visit.patient_id)151            if previous:152                tr_total = self.transition_total[previous]153                if tr_total:154                    score += self.w_transition * math.log(155                        (self.transitions[previous][icd] + alpha)156                        / (tr_total + alpha * num_classes)157                    )158 159            # Patient-specific distribution.160            patient_total = self.patient_total[visit.patient_id]161            if patient_total:162                score += self.w_patient * math.log(163                    (self.patient_counts[visit.patient_id][icd] + alpha)164                    / (patient_total + alpha * num_classes)165                )166 167            # Fully combined context distribution.168            exact_total = self.by_exact_context_total[exact_context]169            if exact_total:170                score += self.w_exact_context * math.log(171                    (self.by_exact_context[exact_context][icd] + alpha)172                    / (exact_total + alpha * num_classes)173                )174 175            scores[icd] = score176 177        # Deterministic tie-break: higher score, then lexicographically smaller ICD.178        return min(scores, key=lambda icd: (-scores[icd], icd))179 180    def update(self, visit: Visit) -> None:181        gold = visit.primary_icd10182        previous = self.last_primary_by_patient.get(visit.patient_id)183 184        self.global_counts[gold] += 1185        self.total_examples += 1186 187        self.by_visit_type[visit.visit_type][gold] += 1188        self.by_visit_type_total[visit.visit_type] += 1189 190        self.by_specialty[visit.provider_specialty][gold] += 1191        self.by_specialty_total[visit.provider_specialty] += 1192 193        visit_specialty = (visit.visit_type, visit.provider_specialty)194        self.by_visit_specialty[visit_specialty][gold] += 1195        self.by_visit_specialty_total[visit_specialty] += 1196 197        for sec in visit.secondary_icd10s:198            self.by_secondary[sec][gold] += 1199            self.by_secondary_total[sec] += 1200 201        secondary_tuple = tuple(sorted(visit.secondary_icd10s))202        exact_context = (visit.visit_type, visit.provider_specialty, secondary_tuple)203        self.by_exact_context[exact_context][gold] += 1204        self.by_exact_context_total[exact_context] += 1205 206        if previous:207            self.transitions[previous][gold] += 1208            self.transition_total[previous] += 1209 210        self.patient_counts[visit.patient_id][gold] += 1211        self.patient_total[visit.patient_id] += 1212        self.last_primary_by_patient[visit.patient_id] = gold213 214 215def simulate_learning(visits: list[Visit], warmup: int) -> dict[str, float | int]:216    learner = OnlineDiagnosisLearner()217    correct = 0218    evaluated = 0219 220    for i, visit in enumerate(visits):221        if i < warmup:222            learner.update(visit)223            continue224 225        predicted = learner.predict(visit)226        if predicted == visit.primary_icd10:227            correct += 1228        evaluated += 1229        learner.update(visit)230 231    accuracy = (correct / evaluated) if evaluated else 0.0232    return {233        "rows_total": len(visits),234        "rows_warmup": warmup,235        "rows_evaluated": evaluated,236        "correct": correct,237        "accuracy": round(accuracy, 4),238    }239 240 241def simulate_learning_with_learner(242    learner: OnlineDiagnosisLearner,243    visits: list[Visit],244    warmup: int,245    print_each_episode: bool = False,246) -> dict[str, float | int]:247    correct = 0248    evaluated = 0249 250    for i, visit in enumerate(visits):251        if i < warmup:252            learner.update(visit)253            continue254 255        predicted = learner.predict(visit)256        if predicted == visit.primary_icd10:257            correct += 1258        evaluated += 1259        if print_each_episode:260            running_accuracy = correct / evaluated261            print(262                f"episode={evaluated} index={i} predicted={predicted} "263                f"gold={visit.primary_icd10} accuracy={running_accuracy:.4f}"264            )265        learner.update(visit)266 267    accuracy = (correct / evaluated) if evaluated else 0.0268    return {269        "rows_total": len(visits),270        "rows_warmup": warmup,271        "rows_evaluated": evaluated,272        "correct": correct,273        "accuracy": round(accuracy, 4),274    }275 276 277def save_learner(learner: OnlineDiagnosisLearner, pkl_path: Path) -> None:278    pkl_path.parent.mkdir(parents=True, exist_ok=True)279    with pkl_path.open("wb") as f:280        pickle.dump(learner, f)281 282 283def load_learner(pkl_path: Path) -> OnlineDiagnosisLearner:284    with pkl_path.open("rb") as f:285        learner = pickle.load(f)286    if not isinstance(learner, OnlineDiagnosisLearner):287        raise TypeError(f"Pickle at {pkl_path} does not contain OnlineDiagnosisLearner.")288    return learner289 290 291def main() -> None:292    parser = argparse.ArgumentParser()293    parser.add_argument(294        "--csv",295        default="data/diagnoses.csv",296        help="Path to diagnoses CSV dataset.",297    )298    parser.add_argument(299        "--warmup",300        type=int,301        default=500,302        help="Number of earliest rows used only for initial learning.",303    )304    parser.add_argument(305        "--save-pkl",306        default="outputs/diagnosis_learner.pkl",307        help="Where to save learner state as pickle after simulation.",308    )309    parser.add_argument(310        "--load-pkl",311        default=None,312        help="Optional path to an existing learner pickle to continue from.",313    )314    parser.add_argument(315        "--print-each-episode",316        action="store_true",317        help="Print running accuracy after every evaluated episode.",318    )319    args = parser.parse_args()320 321    csv_path = Path(args.csv)322    visits = load_visits(csv_path)323    warmup = max(0, min(args.warmup, len(visits)))324    if args.load_pkl:325        learner = load_learner(Path(args.load_pkl))326    else:327        learner = OnlineDiagnosisLearner()328 329    metrics = simulate_learning_with_learner(330        learner,331        visits,332        warmup=warmup,333        print_each_episode=args.print_each_episode,334    )335    save_path = Path(args.save_pkl)336    save_learner(learner, save_path)337 338    print("Learning simulation complete")339    print(f"dataset: {csv_path}")340    print(f"rows_total: {metrics['rows_total']}")341    print(f"rows_warmup: {metrics['rows_warmup']}")342    print(f"rows_evaluated: {metrics['rows_evaluated']}")343    print(f"correct: {metrics['correct']}")344    print(f"accuracy: {metrics['accuracy']:.4f}")345    print(f"saved_learner_pkl: {save_path}")346 347 348if __name__ == "__main__":349    main()350