第 11 章 · 读取 XPT / SAS7BDAT 与临床数据结构

import pandas as pd
import pyreadstat

# ---- 方式 A：pandas（简单，但不带变量标签）----
df = pd.read_sas("data/raw/dm.xpt", format="xport", encoding="utf-8")

# ---- 方式 B：pyreadstat（推荐，带标签 + 值标签）----
df, meta = pyreadstat.read_xport("data/raw/dm.xpt")
df, meta = pyreadstat.read_sas7bdat("adsl.sas7bdat")

# 读永久数据集的目录（自定义 format）
df, meta = pyreadstat.read_sas7bdat("adsl.sas7bdat", catalog_file="formats.sas7bcat")

df["TRTSDT"].dtype            # dtype('O')  —— 字符串！
df["TRTSDT"].head(3)
# 0    2014-01-02
# 1    2012-08-05

df["TRTSDT_DT"] = pd.to_datetime(df["TRTSDT"], errors="coerce")
df["TRTSDT_DT"].dtype         # datetime64[ns]

# 数值型变量（如 TRTDUR / AGE）如果含缺失，也会变成 float64
df.dtypes.value_counts()
# float64    28
# object     18
# int64       2

df["AGE"].head(3)     # 63.0, 64.0, 71.0

# 需要整数显示时
df["AGE"] = df["AGE"].astype("Int64")            # 大写 I = 可空整型
df["AGE"].head(3)                                 # 63, 64, 71

# 或者只在输出时报表格时格式化
print(df["AGE"].describe())

# 利用命名规律写出"通用域迭代器"
BDS_PATTERNS = {
    "subject":   "USUBJID",
    "param":     lambda c: [x for x in c if x.startswith("PARAM")],
    "result":    lambda c: [x for x in c if x in ("AVAL", "AVALC")],
    "baseline":  lambda c: [x for x in c if x in ("BASE", "ABLFL")],
    "change":    lambda c: [x for x in c if x in ("CHG", "PCHG")],
    "visit":     lambda c: [x for x in c if x.startswith("VISIT") or x == "AVISIT"],
    "flag":      lambda c: [x for x in c if x.endswith("FL")],
    "timing":    lambda c: [x for x in c if x.endswith(("DT", "DTM", "DY"))],
}

USUBJID_STRUCT = {
    "01":   "研究编号后缀（CDISCPILOT01 → 01）",
    "701":  "中心编号（SITEID）",
    "1015": "受试者编号（SUBJID）",
}

# 拆解 USUBJID（用正则更安全，因为分隔符可能变化）
import re
m = re.match(r"^(?P<study>[^-]+)-(?P<site>[^-]+)-(?P<subj>.+)$", "01-701-1015")
m.groupdict()   # {'study': '01', 'site': '701', 'subj': '1015'}

# 或直接 split（结构固定时更简单）
df[["STUDY", "SITEID2", "SUBJID2"]] = df["USUBJID"].str.split("-", expand=True)

"""clinic/io.py —— 临床数据读取工具。"""
from __future__ import annotations

import json
import logging
from pathlib import Path

import pandas as pd
import pyreadstat

log = logging.getLogger(__name__)

# CDISC 里常见的"伪缺失"标记
NA_TOKENS = {"", " ", "NA", "N/A", "NULL", "null", ".", "UNKNOWN", "Unknown"}


def read_xpt(path: str | Path, keep_labels: bool = True,
             as_datetime: bool = False) -> pd.DataFrame:
    """读取 SAS XPORT (.xpt)，可选保留变量标签。

    Parameters
    ----------
    path : 文件路径
    keep_labels : 是否把变量标签存到 df.attrs['labels']
    as_datetime : 是否自动把 --DTC / *DT 结尾的字符列转成 datetime
                  （默认 False，避免误转部分日期）
    """
    path = Path(path)
    df, meta = pyreadstat.read_xport(str(path))

    if keep_labels:
        df.attrs["labels"] = meta.column_names_to_labels
        df.attrs["source"] = str(path)

    if as_datetime:
        for c in df.columns:
            if c.endswith(("DTC", "DT")) and df[c].dtype == "object":
                df[c] = pd.to_datetime(df[c], errors="coerce")

    log.info("读取 %-16s %7d 行 x %2d 列", path.name, *df.shape)
    return df


def read_sas7bdat(path, catalog=None, keep_labels=True):
    """读取 SAS 数据集（.sas7bdat）。"""
    df, meta = pyreadstat.read_sas7bdat(str(path), catalog_file=catalog)
    if keep_labels:
        df.attrs["labels"] = meta.column_names_to_labels
        df.attrs["value_labels"] = getattr(meta, "variable_value_labels", {}) or {}
    return df


def read_any(path, **kw) -> pd.DataFrame:
    """按扩展名自动选择读取方式。"""
    p = Path(path)
    suffix = p.suffix.lower()
    if suffix == ".xpt":
        return read_xpt(p, **kw)
    if suffix == ".sas7bdat":
        return read_sas7bdat(p, **kw)
    if suffix == ".csv":
        return pd.read_csv(p, na_values=list(NA_TOKENS), keep_default_na=True)
    if suffix in (".xlsx", ".xls"):
        return pd.read_excel(p)
    if suffix in (".parquet", ".pq"):
        return pd.read_parquet(p)
    raise ValueError(f"不支持的文件类型: {suffix}")


def load_domains(data_dir, domains: list[str], ext: str = "auto") -> dict[str, pd.DataFrame]:
    """批量加载多个域；单个失败不影响其他（见第 04 章 的容错模式）。"""
    data_dir = Path(data_dir)
    out: dict[str, pd.DataFrame] = {}
    for dom in domains:
        # 按优先级找文件：csv -> xpt -> sas7bdat
        for suffix in ([".csv"] if ext == "csv" else
                       [".xpt"] if ext == "xpt" else
                       [".csv", ".xpt", ".sas7bdat"]):
            f = data_dir / f"{dom}{suffix}"
            if f.exists():
                try:
                    out[dom] = read_any(f)
                    break
                except Exception as e:  # noqa: BLE001
                    log.error("读取 %s 失败: %s", f.name, e)
        else:
            log.warning("未找到 %s 的数据文件（已尝试 csv/xpt/sas7bdat）", dom)
    return out


def get_labels(df: pd.DataFrame) -> dict[str, str]:
    """取变量标签；没有则返回 {变量: 变量} 的恒等映射。"""
    return df.attrs.get("labels") or {c: c for c in df.columns}

from clinic.io import load_domains, read_xpt, get_labels

# 一次性加载多个域
src = load_domains("data/raw", ["dm", "ae", "adsl", "adae", "ex", "ds"])

# 或者从 CSV 样本（更快，无需下载）
src = load_domains("data/samples", ["dm", "ae", "adsl", "adae"], ext="csv")

adsl = src["adsl"]
labels = get_labels(adsl)
print(labels["AGE"])    # 'Age'  —— SAS 的 label 被保留了

# 快速验证这个差异
dm = read_xpt("data/raw/dm.xpt")
adsl = read_xpt("data/raw/adsl.xpt")

only_dm = set(dm["USUBJID"]) - set(adsl["USUBJID"])
print(f"DM {len(dm)} 人，ADSL {len(adsl)} 人，只在 DM 中的 {len(only_dm)} 人")
# DM 306 人，ADSL 254 人，只在 DM 中的 52 人

# 看看这些人为什么不在 ADSL 里（通常是筛选失败）
print(dm.loc[dm["USUBJID"].isin(only_dm),
             ["USUBJID", "ARMCD", "ARM", "RFSTDTC"]].head())

# 查看 DataFrame 的真实内存占用（含字符串）
df.memory_usage(deep=True).sum() / 1024**2      # MB

# 优化技巧 1：把低基数的字符串列转成 category（临床数据里效果极好）
for c in ["DOMAIN", "SEX", "RACE", "AESEV", "AESER", "TRT01P"]:
    if c in df.columns:
        df[c] = df[c].astype("category")
# 通常能省 50-80% 内存

# 优化技巧 2：数值降位（如果范围允许）
df["AESEQ"] = df["AESEQ"].astype("int16")

# 优化技巧 3：大文件用分块读取（对应 SAS 的 data _null_ 扫描）
for chunk in pd.read_csv("huge.csv", chunksize=100_000):
    process(chunk)

# 优化技巧 4：Parquet 代替 CSV（体积小 3-5 倍，读取快 5-10 倍）
df.to_parquet("adsl.parquet", index=False)
df = pd.read_parquet("adsl.parquet")