From 0f1654baeae8c5585ee87b34e60d491d3b03996d Mon Sep 17 00:00:00 2001 From: burual <3367932896@qq.com> Date: Thu, 15 Oct 2026 21:11:20 +0800 Subject: [PATCH] Add custom CSV EHR dataset adapter for lightweight onboarding - pyhealth/datasets/synthetic_ehr.py: parse custom CSV into Patient-Visit-Event structure, compatible with 1.x-style API - tests/core/test_synthetic_ehr.py: 7 pytest unit tests - examples/synthetic_ehr/: example data generator + synthetic_ehr.csv --- .../synthetic_ehr/generate_synthetic_ehr.py | 208 ++++++++++ examples/synthetic_ehr/synthetic_ehr.csv | 24 ++ pyhealth/datasets/synthetic_ehr.py | 379 ++++++++++++++++++ tests/core/test_synthetic_ehr.py | 184 +++++++++ 4 files changed, 795 insertions(+) create mode 100644 examples/synthetic_ehr/generate_synthetic_ehr.py create mode 100644 examples/synthetic_ehr/synthetic_ehr.csv create mode 100644 pyhealth/datasets/synthetic_ehr.py create mode 100644 tests/core/test_synthetic_ehr.py diff --git a/examples/synthetic_ehr/generate_synthetic_ehr.py b/examples/synthetic_ehr/generate_synthetic_ehr.py new file mode 100644 index 000000000..c28f16a50 --- /dev/null +++ b/examples/synthetic_ehr/generate_synthetic_ehr.py @@ -0,0 +1,208 @@ +""" +生成教学用合成 EHR 数据集 (PyHealth) + +字段: + patient_id, visit_id, visit_date, age, gender, + diagnosis_code (ICD-9), procedure_code (ICD-9-CM), + drug_code (NDC), lab_glucose, lab_creatinine + +临床合理性设计: + - 每个患者被赋予一个"主诊断"(primary condition)， + 该诊断决定其用药与实验室值的倾向。 + - 糖尿病患者 (250.00): 高血糖倾向 + 更可能用二甲双胍(860975) + - 高血压/冠心病/心衰患者: 更可能用赖诺普利(003127)、阿司匹林(003127)、阿托伐他汀(551110) + - 体检患者 (V70.0): 实验室值接近正常，通常不用药 +""" + +import random +import csv +from datetime import date, timedelta + +random.seed(42) # 可复现 + +# ---- 诊断码与药物码定义 ---- +DIAGNOSES = { + "250.00": "糖尿病", + "401.9": "高血压", + "414.01": "冠心病", + "428.0": "心衰", + "530.81": "GERD", + "V70.0": "体检", +} + +DRUGS = { + "860975": "二甲双胍", # metformin + "003127": "赖诺普利", # lisinopril + "551110": "阿托伐他汀", # atorvastatin + "504580": "阿司匹林", # aspirin +} + +# ---- 每位患者的主诊断, 用于指导临床合理性 ---- +# 10 位患者: 4 糖尿病, 2 高血压, 1 冠心病, 1 心衰, 1 GERD, 1 体检 +PRIMARY_CONDITIONS = [ + "250.00", "250.00", "250.00", "250.00", # 4 糖尿病 + "401.9", "401.9", # 2 高血压 + "414.01", # 1 冠心病 + "428.0", # 1 心衰 + "530.81", # 1 GERD + "V70.0", # 1 体检 +] + + +def pick_primary_diagnosis(condition: str) -> str: + """根据主诊断返回本次就诊的诊断码, 有时会附带并发症诊断。""" + return condition + + +def pick_secondary_diagnoses(condition: str) -> list: + """返回可能伴随的诊断码列表(不含主诊断)。""" + secondary = [] + if condition == "250.00" and random.random() < 0.5: + secondary.append("401.9") # 糖尿病常合并高血压 + if condition == "414.01" and random.random() < 0.6: + secondary.append("401.9") # 冠心病常合并高血压 + if condition == "428.0" and random.random() < 0.5: + secondary.append("401.9") + if condition == "401.9" and random.random() < 0.3: + secondary.append("250.00") + return secondary + + +def pick_drugs(condition: str) -> list: + """根据主诊断选择药物, 体现临床合理性。""" + drugs = [] + if condition == "250.00": + if random.random() < 0.9: + drugs.append("860975") # 二甲双胍 + if random.random() < 0.4: + drugs.append("003127") # 赖诺普利 (合并高血压时) + elif condition == "401.9": + if random.random() < 0.9: + drugs.append("003127") # 赖诺普利 + if random.random() < 0.4: + drugs.append("551110") # 阿托伐他汀 + elif condition == "414.01": + if random.random() < 0.8: + drugs.append("003127") # 阿司匹林 + if random.random() < 0.7: + drugs.append("551110") # 阿托伐他汀 + elif condition == "428.0": + if random.random() < 0.8: + drugs.append("003127") # 赖诺普利 + elif condition == "530.81": + if random.random() < 0.5: + pass # GERD 有时不用所列药物 + # V70.0 体检通常不用药 + if not drugs: + # 保证非空, 给一个低概率安慰用药或留空由调用方处理 + if random.random() < 0.3: + drugs.append("551110") + return drugs + + +def gen_glucose(condition: str) -> float: + """根据诊断生成合理血糖值。""" + if condition == "250.00": + # 糖尿病: 偏高 (110-300) + return round(random.uniform(110, 300), 1) + else: + # 非糖尿病: 正常偏低 (70-130) + return round(random.uniform(70, 130), 1) + + +def gen_creatinine(condition: str) -> float: + """根据诊断生成合理肌酐值 (心衰/糖尿病肾功能略差)。""" + base = random.uniform(0.5, 1.2) + if condition == "428.0": + base = random.uniform(0.9, 2.5) # 心衰肾功能可能受损 + elif condition == "250.00": + base = random.uniform(0.6, 2.0) # 糖尿病肾病风险 + return round(min(base, 3.0), 2) + + +def gen_procedure(condition: str) -> str: + """简单流程码: 用诊断码简化替代 ICD-9-CM 流程 (教学用途)。""" + # 常见配套检查流程 (ICD-9-CM 示例) + procedures = { + "250.00": "89.52", # 心电图? 此处仅为示例码 + "401.9": "89.65", # 动脉血气/血压监测示例 + "414.01": "88.72", # 诊断性心导管 + "428.0": "89.52", # 心电图 + "530.81": "87.62", # 上消化道造影示例 + "V70.0": "89.01", # 常规体检问诊 + } + return procedures[condition] + + +rows = [] +visit_id_counter = 1 + +for pid_idx, condition in enumerate(PRIMARY_CONDITIONS, start=1): + patient_id = f"P{pid_idx:03d}" + age = random.randint(30, 85) + gender = "M" if pid_idx <= 5 else "F" # 前 5 位 M, 后 5 位 F + + n_visits = random.randint(1, 4) + + # 就诊日期: 在 2022-01 ~ 2023-12 内递增分布 + start = date(2022, 1, 1) + end = date(2023, 12, 31) + total_days = (end - start).days + + visit_dates = [] + # 生成 n_visits 个递增日期 + offsets = sorted(random.sample(range(total_days + 1), n_visits)) + for off in offsets: + visit_dates.append(start + timedelta(days=off)) + + for vi in range(n_visits): + visit_id = f"V{visit_id_counter:04d}" + visit_id_counter += 1 + + # 诊断码: 主诊断 + 可能的伴随诊断, 用 '|' 分隔 + diag_list = [pick_primary_diagnosis(condition)] + diag_list += pick_secondary_diagnoses(condition) + # 去重保持顺序 + seen = set() + diag_codes = [] + for d in diag_list: + if d not in seen: + seen.add(d) + diag_codes.append(d) + diagnosis_code = "|".join(diag_codes) + + drugs = pick_drugs(condition) + drug_code = "|".join(drugs) if drugs else "" + + rows.append({ + "patient_id": patient_id, + "visit_id": visit_id, + "visit_date": visit_dates[vi].isoformat(), + "age": age, + "gender": gender, + "diagnosis_code": diagnosis_code, + "procedure_code": gen_procedure(condition), + "drug_code": drug_code, + "lab_glucose": gen_glucose(condition), + "lab_creatinine": gen_creatinine(condition), + }) + +# ---- 保存 CSV ---- +fieldnames = [ + "patient_id", "visit_id", "visit_date", "age", "gender", + "diagnosis_code", "procedure_code", "drug_code", + "lab_glucose", "lab_creatinine", +] + +with open("synthetic_ehr.csv", "w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(rows) + +# ---- 打印前 10 行确认 ---- +print("生成完成: synthetic_ehr.csv") +print("总记录数:", len(rows)) +print("-" * 100) +print(",".join(fieldnames)) +for r in rows[:10]: + print(",".join(str(r[k]) for k in fieldnames)) diff --git a/examples/synthetic_ehr/synthetic_ehr.csv b/examples/synthetic_ehr/synthetic_ehr.csv new file mode 100644 index 000000000..be82929f7 --- /dev/null +++ b/examples/synthetic_ehr/synthetic_ehr.csv @@ -0,0 +1,24 @@ +patient_id,visit_id,visit_date,age,gender,diagnosis_code,procedure_code,drug_code,lab_glucose,lab_creatinine +P001,V0001,2022-01-26,70,M,250.00,89.52,860975|003127,129.5,1.36 +P002,V0002,2022-01-31,67,M,250.00|401.9,89.52,860975,246.0,1.19 +P002,V0003,2022-02-02,67,M,250.00|401.9,89.52,860975,254.2,1.19 +P002,V0004,2022-04-06,67,M,250.00|401.9,89.52,860975,129.4,1.1 +P002,V0005,2022-08-12,67,M,250.00|401.9,89.52,860975|003127,197.3,1.89 +P003,V0006,2023-01-06,35,M,250.00,89.52,860975|003127,153.3,0.71 +P003,V0007,2023-09-26,35,M,250.00|401.9,89.52,860975|003127,230.8,1.12 +P003,V0008,2023-10-06,35,M,250.00|401.9,89.52,860975,233.1,0.84 +P004,V0009,2022-06-17,76,M,250.00|401.9,89.52,,240.1,1.69 +P004,V0010,2023-04-19,76,M,250.00|401.9,89.52,860975|003127,160.9,1.92 +P005,V0011,2022-08-06,75,M,401.9,89.65,551110,84.8,0.89 +P005,V0012,2023-05-27,75,M,401.9|250.00,89.65,003127,94.0,0.65 +P005,V0013,2023-11-03,75,M,401.9,89.65,003127|551110,72.8,0.58 +P006,V0014,2023-03-09,70,F,401.9,89.65,003127,98.1,0.68 +P006,V0015,2023-11-28,70,F,401.9,89.65,551110,123.1,1.03 +P007,V0016,2022-04-25,79,F,414.01|401.9,88.72,003127,113.2,1.18 +P007,V0017,2022-10-28,79,F,414.01,88.72,003127|551110,107.5,1.09 +P007,V0018,2023-03-22,79,F,414.01|401.9,88.72,003127|551110,79.7,1.17 +P008,V0019,2023-09-06,63,F,428.0|401.9,89.52,003127,125.7,2.23 +P009,V0020,2022-03-01,49,F,530.81,87.62,,75.1,0.84 +P009,V0021,2022-09-04,49,F,530.81,87.62,,116.0,0.59 +P010,V0022,2022-09-29,60,F,V70.0,89.01,,95.4,0.65 +P010,V0023,2023-06-25,60,F,V70.0,89.01,,113.8,0.64 diff --git a/pyhealth/datasets/synthetic_ehr.py b/pyhealth/datasets/synthetic_ehr.py new file mode 100644 index 000000000..9d31de938 --- /dev/null +++ b/pyhealth/datasets/synthetic_ehr.py @@ -0,0 +1,379 @@ +""" +synthetic_ehr_dataset.py +======================= + +PyHealth 自定义 CSV 数据集的 adapter（教学用）。 + +将 ``synthetic_ehr.csv`` 解析为 PyHealth 风格的 +``Patient -> Visit -> Event`` 结构，事件类型包括： + + - diagnosis (ICD-9, ``|`` 分隔多码) + - procedure (ICD-9-CM,``|`` 分隔多码) + - prescription (NDC, ``|`` 分隔多码) + - lab (数值型: glucose / creatinine, 存为 Event 的 value) + +设计说明 +-------- +* PyHealth 2.x 中 ``Patient`` 改为 Polars DataFrame 驱动、``BaseDataset`` 改为 + YAML 配置驱动，已不存在 ``BaseEHRDataset`` 与 ``Visit`` 类。因此本 adapter + 在**尽量贴近 MIMIC3Dataset 的解析风格**（parse_basic_info + 按 patient 分组 + 构造事件）的同时，自包含地实现了 ``Patient -> Visit -> Event`` 结构，并保证： + + 1) 若安装了 PyHealth，``SyntheticEHRDataset`` 继承 ``pyhealth.datasets.BaseDataset``， + 同时保留 PyHealth 原生的 ``Patient``/``Event`` 作为底层事件对象； + 2) 若未安装 PyHealth，则回退到一个等价的轻量基类，代码仍可独立运行。 + +* lab 值既作为 ``Event``（event_type="lab"，属性含 ``lab_name`` 与 ``value``）， + 也写入 visit 级 metadata（``visit.labs`` dict），便于下游按需取用。 +""" + +from __future__ import annotations + +import csv +from datetime import datetime +from pathlib import Path +from typing import Dict, List, Optional + +# --------------------------------------------------------------------------- # +# 可选依赖：PyHealth 的 Patient / Event / BaseDataset +# 未安装时回退到本文件内定义的轻量实现，保证脚本可独立运行。 +# --------------------------------------------------------------------------- # +try: + from pyhealth.data import Event as PyHealthEvent + from pyhealth.data import Patient as PyHealthPatient + from pyhealth.datasets import BaseDataset as PyHealthBaseDataset + + _HAS_PYHEALTH = True +except Exception: # pragma: no cover - 环境未装 PyHealth + PyHealthEvent = None + PyHealthPatient = None + PyHealthBaseDataset = None + _HAS_PYHEALTH = False + + +# --------------------------------------------------------------------------- # +# 轻量事件类（无 PyHealth 时使用） +# --------------------------------------------------------------------------- # +class Event: + """单个临床事件：event_type + timestamp + 任意属性。""" + + __slots__ = ("event_type", "timestamp", "attr_dict") + + def __init__(self, event_type: str, timestamp: Optional[datetime] = None, **kwargs): + self.event_type = event_type + self.timestamp = timestamp + self.attr_dict = dict(kwargs) + + def __getattr__(self, name): + if name in self.attr_dict: + return self.attr_dict[name] + raise AttributeError(name) + + def __getitem__(self, key): + return self.attr_dict[key] + + def __repr__(self): + attrs = ", ".join(f"{k}={v!r}" for k, v in self.attr_dict.items()) + return f"Event({self.event_type!r}, {self.timestamp}, {attrs})" + + +# 统一事件类：有 PyHealth 用其原生 Event，否则用上面的轻量实现。 +_EventCls = PyHealthEvent if _HAS_PYHEALTH else Event + + +class Visit: + """一次就诊，聚合该次就诊的所有 Event。""" + + def __init__(self, visit_id: str, timestamp: datetime): + self.visit_id = visit_id + self.timestamp = timestamp + self.events: List = [] + self.metadata: Dict = {} # age, gender 等 visit 级字段 + self.labs: Dict[str, float] = {} # lab 值也可作为 visit 级 metadata + + def add_event(self, event) -> None: + self.events.append(event) + + def __repr__(self): + return f"Visit({self.visit_id!r}, {self.timestamp.date()}, {len(self.events)} events)" + + +class Patient: + """一个患者，聚合其所有 Visit。""" + + def __init__(self, patient_id: str): + self.patient_id = patient_id + self.visits: List[Visit] = [] + self.metadata: Dict = {} + + def add_visit(self, visit: Visit) -> None: + self.visits.append(visit) + + def __repr__(self): + return f"Patient({self.patient_id!r}, {len(self.visits)} visits)" + + +# --------------------------------------------------------------------------- # +# 基类（有 PyHealth 时继承其 BaseDataset，否则用轻量基类） +# --------------------------------------------------------------------------- # +if _HAS_PYHEALTH: + + class _BaseDataset(PyHealthBaseDataset): + pass + +else: + + class _BaseDataset: + """轻量基类，仅提供与 BaseDataset 同名的最小组件。""" + + def __init__(self, *args, **kwargs): + self.patients: Dict[str, Patient] = {} + self._unique_patient_ids: List[str] = [] + + @property + def unique_patient_ids(self) -> List[str]: + return list(self._unique_patient_ids) + + def get_patient(self, patient_id: str) -> Patient: + if patient_id not in self.patients: + raise KeyError(f"patient {patient_id} not found") + return self.patients[patient_id] + + def iter_patients(self): + for pid in self._unique_patient_ids: + yield self.patients[pid] + + +# --------------------------------------------------------------------------- # +# 数据集 adapter +# --------------------------------------------------------------------------- # +class SyntheticEHRDataset(_BaseDataset): + """读取 ``synthetic_ehr.csv`` 并构造 Patient -> Visit -> Event 结构。 + + Parameters + ---------- + root : str | Path + 数据集文件（CSV）所在目录；若传入的是文件路径则直接使用该文件。 + dataset_name : str, optional + 数据集名称，默认为 "synthetic_ehr"。 + """ + + # 列名到事件类型的映射（lab 两列合并为同一 event_type="lab"） + EVENT_TYPE_MAP = { + "diagnosis_code": "diagnosis", + "procedure_code": "procedure", + "drug_code": "prescription", + "lab_glucose": "lab", + "lab_creatinine": "lab", + } + + LAB_NAMES = {"lab_glucose": "glucose", "lab_creatinine": "creatinine"} + + def __init__( + self, + root: str | Path, + dataset_name: str = "synthetic_ehr", + **kwargs, + ): + self.root = Path(root) + # 兼容：root 既可指向目录也可指向 csv 文件 + self.csv_path = self.root if self.root.is_file() else self.root / "synthetic_ehr.csv" + if not self.csv_path.exists(): + raise FileNotFoundError(f"未找到数据文件: {self.csv_path}") + + self.dataset_name = dataset_name + + if _HAS_PYHEALTH: + # 用 PyHealth BaseDataset 的最小初始化（不触发其 YAML 加载逻辑）。 + # 我们自行构造 patient 树，因此直接跳过 BaseDataset 的 load_data 机制。 + PyHealthBaseDataset.__init__( + self, + root=str(self.root), + tables=[], # 不使用原生 table 加载 + dataset_name=dataset_name, + **kwargs, + ) + else: + super().__init__(root=str(self.root), dataset_name=dataset_name, **kwargs) + + # 核心：解析 CSV 并构造 Patient -> Visit -> Event + self.patients: Dict[str, Patient] = {} + self._unique_patient_ids: List[str] = [] + self.parse_basic_info() + + self._print_stats() + + # ------------------------------------------------------------------ # + # 解析主流程 + # ------------------------------------------------------------------ # + def parse_basic_info(self) -> None: + """读取 CSV，按 patient_id 分组，构造 Patient -> Visit -> Event。""" + rows = self._read_rows() + + for row in rows: + pid = row.get("patient_id", "").strip() + if not pid: + continue # 空 patient_id 跳过 + + visit = self._build_visit(row) + if visit is None: + continue # 无效行（如日期无法解析）跳过 + + patient = self.patients.get(pid) + if patient is None: + patient = Patient(pid) + patient.metadata["gender"] = self._safe(row, "gender") + patient.metadata["age"] = self._to_int(row.get("age")) + self.patients[pid] = patient + self._unique_patient_ids.append(pid) + + patient.add_visit(visit) + + # 每个患者的 visit 按时间排序 + for patient in self.patients.values(): + patient.visits.sort(key=lambda v: v.timestamp) + + def _read_rows(self) -> List[dict]: + """读取 CSV 为 dict 列表（空行/空字段按原样保留，由调用方处理）。""" + with open(self.csv_path, "r", newline="", encoding="utf-8-sig") as f: + reader = csv.DictReader(f) + return [row for row in reader] + + def _build_visit(self, row: dict) -> Optional[Visit]: + """将一行 CSV 解析为一次 Visit，含 diagnosis/procedure/prescription/lab 事件。""" + visit_id = self._safe(row, "visit_id") + ts = self._parse_date(row.get("visit_date")) + if ts is None: + return None + + visit = Visit(visit_id, ts) + visit.metadata["age"] = self._to_int(row.get("age")) + visit.metadata["gender"] = self._safe(row, "gender") + + for col, event_type in self.EVENT_TYPE_MAP.items(): + if event_type == "lab": + # lab 两列分别作为一个 Event，value 存数值 + self._add_lab_event(visit, col, row.get(col), ts) + else: + self._add_code_events(visit, col, event_type, row.get(col), ts) + + return visit + + def _add_code_events(self, visit: Visit, col: str, event_type: str, + raw: Optional[str], ts: datetime) -> None: + """拆分 ``|`` 分隔的多码，逐码生成一个 Event；空值跳过。""" + if raw is None: + return + codes = [c.strip() for c in str(raw).split("|") if c.strip()] + for code in codes: + event = _EventCls(event_type, timestamp=ts, code=code) + visit.add_event(event) + + def _add_lab_event(self, visit: Visit, col: str, raw: Optional[str], + ts: datetime) -> None: + """lab 值作为数值型 Event，同时写入 visit.labs；空值/非法值跳过。""" + val = self._to_float(raw) + if val is None: + return + lab_name = self.LAB_NAMES[col] + event = _EventCls("lab", timestamp=ts, lab_name=lab_name, value=val) + visit.add_event(event) + visit.labs[lab_name] = val + + # ------------------------------------------------------------------ # + # 工具函数 + # ------------------------------------------------------------------ # + @staticmethod + def _parse_date(raw) -> Optional[datetime]: + if raw is None: + return None + raw = str(raw).strip() + if not raw: + return None + try: + return datetime.strptime(raw, "%Y-%m-%d") + except ValueError: + return None + + @staticmethod + def _safe(row: dict, key: str) -> str: + v = row.get(key) + return "" if v is None else str(v).strip() + + @staticmethod + def _to_int(raw) -> Optional[int]: + try: + return int(float(str(raw).strip())) + except (TypeError, ValueError): + return None + + @staticmethod + def _to_float(raw) -> Optional[float]: + try: + return float(str(raw).strip()) + except (TypeError, ValueError): + return None + + # ------------------------------------------------------------------ # + # 统计与查询 + # ------------------------------------------------------------------ # + def _print_stats(self) -> None: + """打印数据集统计：患者数、就诊总数、诊断码种类数、药物码种类数。""" + n_patients = len(self._unique_patient_ids) + n_visits = sum(len(p.visits) for p in self.patients.values()) + + diag_codes: set = set() + drug_codes: set = set() + for p in self.patients.values(): + for v in p.visits: + for e in v.events: + if e.event_type == "diagnosis": + diag_codes.add(e.code) + elif e.event_type == "prescription": + drug_codes.add(e.code) + + print("=" * 50) + print(f"[SyntheticEHRDataset] {self.dataset_name}") + print(f" 患者数: {n_patients}") + print(f" 就诊总数: {n_visits}") + print(f" 诊断码种类数: {len(diag_codes)}") + print(f" 药物码种类数: {len(drug_codes)}") + print("=" * 50) + + def stats(self) -> None: + """兼容 BaseDataset.stats() 接口。""" + self._print_stats() + + def describe_patient(self, patient_id: str) -> str: + """返回单个患者结构的人类可读文本，便于测试块打印。""" + p = self.get_patient(patient_id) + lines = [f"患者 {p.patient_id} (gender={p.metadata.get('gender')}, " + f"age={p.metadata.get('age')}, visits={len(p.visits)})"] + for v in p.visits: + lines.append(f" ├─ {v}") + for e in v.events: + if e.event_type == "lab": + lines.append(f" │ └─ {e.event_type:12s} {e.lab_name:<10s} = {e.value}") + else: + lines.append(f" │ └─ {e.event_type:12s} {e.code}") + return "\n".join(lines) + + +# --------------------------------------------------------------------------- # +# 测试块 +# --------------------------------------------------------------------------- # +if __name__ == "__main__": + import sys + + # 兼容：既可从命令行传路径，也可默认当前目录 + root = sys.argv[1] if len(sys.argv) > 1 else "." + + ds = SyntheticEHRDataset(root=root) + + print() + print("患者 ID 列表:", ds.unique_patient_ids) + print() + + # 打印第一个患者的结构 + first_pid = ds.unique_patient_ids[0] + print(ds.describe_patient(first_pid)) diff --git a/tests/core/test_synthetic_ehr.py b/tests/core/test_synthetic_ehr.py new file mode 100644 index 000000000..6e9435eb8 --- /dev/null +++ b/tests/core/test_synthetic_ehr.py @@ -0,0 +1,184 @@ +""" +tests/core/test_synthetic_ehr.py +================================ + +:class:`pyhealth.datasets.synthetic_ehr.SyntheticEHRDataset` 的单元测试。 + +被测结构回顾（见 ``pyhealth/datasets/synthetic_ehr.py``）： + + SyntheticEHRDataset(root) # root 为目录或 csv 文件路径 + .patients # {patient_id: Patient} + Patient.patient_id / Patient.metadata # gender, age + Patient.visits # [Visit, ...]（按 timestamp 升序） + Visit.visit_id / Visit.timestamp / Visit.events / Visit.labs + Event.event_type / Event.code（或 Event.lab_name + Event.value） + +测试数据用标准库 ``csv`` 在 pytest 的 ``tmp_path`` 下现场生成， +不依赖任何仓库内的固定数据文件。 +""" + +import csv +import os +import sys + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..")) + +from pyhealth.datasets.synthetic_ehr import SyntheticEHRDataset + + +# --------------------------------------------------------------------------- # +# 公共工具 +# --------------------------------------------------------------------------- # +COLUMNS = [ + "patient_id", + "visit_id", + "visit_date", + "age", + "gender", + "diagnosis_code", + "procedure_code", + "drug_code", + "lab_glucose", + "lab_creatinine", +] + + +def _row(pid, vid, date, age=60, gender="M", diag="250.00", + proc="89.52", drug="860975", glu=120.0, cre=1.0): + """构造一条完整的就诊记录。""" + return { + "patient_id": pid, + "visit_id": vid, + "visit_date": date, + "age": age, + "gender": gender, + "diagnosis_code": diag, + "procedure_code": proc, + "drug_code": drug, + "lab_glucose": glu, + "lab_creatinine": cre, + } + + +def _write_csv(tmp_path, rows): + """把若干行记录写入 ``tmp_path/synthetic_ehr.csv``。""" + path = os.path.join(str(tmp_path), "synthetic_ehr.csv") + with open(path, "w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=COLUMNS) + writer.writeheader() + writer.writerows(rows) + return path + + +def _load(tmp_path): + """按目录路径加载数据集（被测代码会自动拼接 synthetic_ehr.csv）。""" + return SyntheticEHRDataset(root=str(tmp_path)) + + +def _events_of_type(visit, event_type): + """返回某个 visit 中指定 event_type 的事件列表。""" + return [e for e in visit.events if e.event_type == event_type] + + +# --------------------------------------------------------------------------- # +# 测试用例 +# --------------------------------------------------------------------------- # +def test_load_normal_csv(tmp_path): + """构造 3 患者 × 各 2 次就诊，验证患者数=3、就诊总数=6。""" + rows = [] + for i in range(1, 4): + pid = f"P{i:03d}" + for j in range(1, 3): + rows.append(_row(pid, f"V{i}{j}", f"2022-0{i}-0{j}")) + + _write_csv(tmp_path, rows) + ds = _load(tmp_path) + + assert len(ds.patients) == 3 + total_visits = sum(len(p.visits) for p in ds.patients.values()) + assert total_visits == 6 + + +def test_patient_metadata(tmp_path): + """验证 patient.metadata 正确存储 age 与 gender。""" + rows = [ + _row("P001", "V001", "2022-01-01", age=72, gender="M"), + _row("P002", "V001", "2022-01-02", age=58, gender="F"), + ] + _write_csv(tmp_path, rows) + ds = _load(tmp_path) + + p1 = ds.get_patient("P001") + p2 = ds.get_patient("P002") + assert p1.metadata["age"] == 72 + assert p1.metadata["gender"] == "M" + assert p2.metadata["age"] == 58 + assert p2.metadata["gender"] == "F" + + +def test_visit_timestamp_sorting(tmp_path): + """同一患者 3 次就诊日期乱序输入，验证 visit 按日期升序排列。""" + rows = [ + _row("P001", "V003", "2022-03-01"), + _row("P001", "V001", "2022-01-01"), + _row("P001", "V002", "2022-02-01"), + ] + _write_csv(tmp_path, rows) + ds = _load(tmp_path) + + visits = ds.get_patient("P001").visits + dates = [v.timestamp.strftime("%Y-%m-%d") for v in visits] + assert dates == ["2022-01-01", "2022-02-01", "2022-03-01"] + + +def test_event_types_present(tmp_path): + """验证一次 visit 中 diagnosis / procedure / prescription 三类事件均被生成。""" + rows = [_row("P001", "V001", "2022-01-01")] + _write_csv(tmp_path, rows) + ds = _load(tmp_path) + + visit = ds.get_patient("P001").visits[0] + types = {e.event_type for e in visit.events} + assert {"diagnosis", "procedure", "prescription"} <= types + + +def test_lab_values(tmp_path): + """验证 lab 事件携带正确的 lab_name 与数值 value，并写入 visit.labs。""" + rows = [_row("P001", "V001", "2022-01-01", glu=185.5, cre=1.2)] + _write_csv(tmp_path, rows) + ds = _load(tmp_path) + + visit = ds.get_patient("P001").visits[0] + lab_events = _events_of_type(visit, "lab") + lab_map = {e.lab_name: e.value for e in lab_events} + + assert lab_map["glucose"] == 185.5 + assert lab_map["creatinine"] == 1.2 + # 同时确认 visit 级 labs 元数据 + assert visit.labs["glucose"] == 185.5 + assert visit.labs["creatinine"] == 1.2 + + +def test_empty_csv(tmp_path): + """空 CSV（仅表头）加载不崩溃，患者数与就诊数均为 0。""" + _write_csv(tmp_path, []) # 只有表头，无数据行 + ds = _load(tmp_path) + + assert len(ds.patients) == 0 + assert ds.unique_patient_ids == [] + + +def test_multiple_visits_same_patient(tmp_path): + """同一患者 3 次就诊，验证均归属同一 patient 且 visit_id 正确。""" + rows = [ + _row("P001", "V001", "2022-01-01"), + _row("P001", "V002", "2022-04-01"), + _row("P001", "V003", "2022-07-01"), + ] + _write_csv(tmp_path, rows) + ds = _load(tmp_path) + + assert len(ds.patients) == 1 + p = ds.get_patient("P001") + assert len(p.visits) == 3 + assert [v.visit_id for v in p.visits] == ["V001", "V002", "V003"]