409 lines
18 KiB
Python
409 lines
18 KiB
Python
"""GWAS 全基因组关联分析：GLM+PC（群体结构主成分校正）/ EMMAX / ssGWAS，纯 numpy 无 R。
|
||
|
||
模型：逐标记单标记回归 y ~ μ + PC1..PCk + marker，加性效应=标记回归系数。
|
||
- 群体结构：标记剂量矩阵(缺失按列均值填)中心化后 SVD，取前 n_pc 个主成分得分作协变量。
|
||
- 显著性：t² ~ F(1, df)（df = n − n_pc − 2），p 值复用 fdist.f_pvalue（正则不完全 beta）。
|
||
- 多重检验：Bonferroni 阈值 + Benjamini-Hochberg FDR q 值（numpy argsort 实现）。
|
||
- QTL 定位：显著(Bonf)标记按染色体/位置排序，相邻间距 < qtl_window 合并为同一 QTL，
|
||
簇内最小 p 为峰标记；单个显著标记也落 QTL（首版兜底）。
|
||
|
||
method 三路径共用 _prepare 数据装配与 _finalize 后处理：
|
||
- gwas：GLM+PC 逐标记单标记回归（默认）；
|
||
- emmax：EMMAX 混合模型（零模型方差分量 → 固定 V 逐标记 GLS）；
|
||
- ssgwas：单步 GWAS（GBLUP 拟合全样本 GEBV → MEM 对角近似反推标记效应，Wang et al. 2012）。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import math
|
||
|
||
import numpy as np
|
||
|
||
from scripts.breeding_stats import fdist, genomic
|
||
|
||
ENGINE_VERSION = "1.2.0" # 1.2.0: ssGWAS（GBLUP GEBV → MEM 标记效应反推）
|
||
|
||
|
||
def _chrom_key(c: str) -> tuple[int, str]:
|
||
"""染色体排序键：数字开头按数字，否则按字符串（'Chr2' < 'Chr10' 需求）。"""
|
||
if c is None:
|
||
return (10 ** 9, "")
|
||
s = str(c).strip()
|
||
i = 0
|
||
while i < len(s) and s[i].isdigit():
|
||
i += 1
|
||
if i and i == len(s):
|
||
return (int(s), "")
|
||
if i:
|
||
return (int(s[:i]), s[i:])
|
||
return (10 ** 9, s)
|
||
|
||
|
||
def _bh_qvalues(pvals: np.ndarray) -> np.ndarray:
|
||
"""Benjamini-Hochberg FDR q 值（p 升序：q_i = p_i·m/i，自底向上保单调）。"""
|
||
m = len(pvals)
|
||
if m == 0:
|
||
return np.array([], dtype=float)
|
||
order = np.argsort(pvals)
|
||
q = np.empty(m)
|
||
running = 1.0
|
||
for rank in range(m - 1, -1, -1):
|
||
running = min(running, float(pvals[order[rank]]) * m / (rank + 1))
|
||
q[order[rank]] = running
|
||
return np.clip(q, 0.0, 1.0)
|
||
|
||
|
||
def _prepare(phenos: dict[str, float], dosage: dict[str, dict[str, float]],
|
||
markers_meta: dict[str, dict[str, object]], maf_min: float, n_pc: int,
|
||
) -> tuple[list[str], np.ndarray, np.ndarray, np.ndarray, int, int,
|
||
list[str], np.ndarray, np.ndarray, int]:
|
||
"""共享数据准备（GLM/EMMAX 两路径用同一份 M/MAF/PC）：交集样本、剂量矩阵、
|
||
MAF 过滤+均值填补、前 n_pc 主成分得分。返回 (common, M, p, maf, m, m_keep,
|
||
markers_k, pc_scores, y, n)；p/maf/markers_k 均按 keep 过滤对齐。"""
|
||
common = sorted(set(phenos) & set(dosage))
|
||
if len(common) < 2:
|
||
raise ValueError("表型与基因型交集样本不足 2，无法关联分析")
|
||
if len(common) < n_pc + 3:
|
||
raise ValueError(f"样本数 {len(common)} < n_pc+3（{n_pc + 3}），PC 校正需更多样本")
|
||
|
||
markers = [mk for mk in markers_meta if mk in dosage[common[0]]]
|
||
if not markers:
|
||
raise ValueError("无有效标记（需在 markers_meta 与 dosage 中均存在）")
|
||
n, m = len(common), len(markers)
|
||
M = np.full((n, m), np.nan)
|
||
for i, ind in enumerate(common):
|
||
g = dosage[ind]
|
||
for j, mk in enumerate(markers):
|
||
v = g.get(mk)
|
||
if v is None or (isinstance(v, float) and math.isnan(v)):
|
||
continue
|
||
M[i, j] = float(v)
|
||
|
||
# MAF 过滤 + 均值填补
|
||
p = np.zeros(m)
|
||
for j in range(m):
|
||
col = M[:, j]
|
||
valid = col[~np.isnan(col)]
|
||
p[j] = float(np.mean(valid)) / 2.0 if len(valid) else 0.5
|
||
M[:, j] = np.where(np.isnan(M[:, j]), 2.0 * p[j], M[:, j])
|
||
maf = np.minimum(p, 1.0 - p)
|
||
keep = maf >= maf_min
|
||
if int(keep.sum()) == 0:
|
||
raise ValueError(f"MAF≥{maf_min} 的标记为 0，请降低 maf_min 或检查标记等位频率")
|
||
M = M[:, keep]
|
||
p = p[keep]
|
||
maf = maf[keep]
|
||
m_keep = int(keep.sum())
|
||
markers_k = [mk for mk, k in zip(markers, keep) if k]
|
||
|
||
# 群体结构：中心化标记矩阵 SVD → 前 n_pc 个主成分得分
|
||
M_c = M - M.mean(axis=0, keepdims=True)
|
||
try:
|
||
U, s, _ = np.linalg.svd(M_c, full_matrices=False)
|
||
except np.linalg.LinAlgError:
|
||
raise ValueError("SVD 分解失败（标记矩阵病态，请检查标记共线性）") from None
|
||
pc_scores = U[:, :n_pc] * s[:n_pc] if n_pc <= len(s) else U[:, :len(s)] * s
|
||
|
||
y = np.array([float(phenos[i]) for i in common], dtype=float)
|
||
return common, M, p, maf, m, m_keep, markers_k, pc_scores, y, n
|
||
|
||
|
||
def _finalize(snps: list[dict], m_keep: int, sig_level: float,
|
||
qtl_window: int) -> tuple[float, list[dict]]:
|
||
"""多重检验（Bonferroni 阈值 + BH-FDR q）+ QTL 定位（同染色体相邻间距合并）。
|
||
|
||
GLM 与 EMMAX 共用的后处理；GLM 路径输出逐字节不变。
|
||
原地给 snps 逐项补 q/sig_bonf/sig_fdr，返回 (threshold_bonf, qtls)。
|
||
"""
|
||
threshold_bonf = sig_level / m_keep
|
||
pvals = np.array([s["p"] for s in snps])
|
||
qvals = _bh_qvalues(pvals)
|
||
for s, q in zip(snps, qvals):
|
||
s["q"] = round(float(q), 6)
|
||
s["sig_bonf"] = bool(s["p"] < threshold_bonf)
|
||
s["sig_fdr"] = bool(q < sig_level)
|
||
|
||
qtls: list[dict] = []
|
||
sigs = [s for s in snps if s["sig_bonf"]]
|
||
sigs.sort(key=lambda s: (_chrom_key(s["chromosome"]), s["position"]))
|
||
for c in _cluster(sigs, qtl_window):
|
||
peak = min(c, key=lambda s: s["p"])
|
||
qtls.append({
|
||
"qtl": len(qtls) + 1,
|
||
"chromosome": peak["chromosome"],
|
||
"start_bp": c[0]["position"],
|
||
"end_bp": c[-1]["position"],
|
||
"peak_marker": peak["marker"],
|
||
"peak_p": peak["p"],
|
||
"n_markers": len(c),
|
||
"effect": peak["effect"],
|
||
})
|
||
return threshold_bonf, qtls
|
||
|
||
|
||
def run_gwas(phenos: dict[str, float], dosage: dict[str, dict[str, float]],
|
||
markers_meta: dict[str, dict[str, object]], *,
|
||
maf_min: float = 0.05, n_pc: int = 3, sig_level: float = 0.05,
|
||
qtl_window: int = 1_000_000, method: str = "gwas") -> dict:
|
||
"""GWAS 关联扫描：GLM+PC（默认）/ EMMAX 混合模型 / ssGWAS 单步 + QTL 定位。
|
||
|
||
phenos: {个体: 表型}；dosage: {个体: {标记: 剂量0/1/2}}；
|
||
markers_meta: {标记: {chromosome, position}}。个体键约定 f"t{tid}"（树为个体）。
|
||
method: gwas=GLM+PC 单标记回归；emmax=EMMAX（零模型方差分量→固定 V 逐标记 GLS）；
|
||
ssgwas=GBLUP 拟合 GEBV 后 MEM 对角近似反推标记效应（Wang et al. 2012 单步 GWAS）。
|
||
"""
|
||
method = (method or "gwas").lower()
|
||
if method not in ("gwas", "ssgwas", "emmax"):
|
||
raise ValueError("method 取值不合法（可用: gwas / emmax / ssgwas）")
|
||
n_pc = max(1, int(n_pc or 3))
|
||
if n_pc > 10:
|
||
raise ValueError(f"n_pc 过大（≤10），当前 {n_pc}")
|
||
|
||
common, M, p, maf, m, m_keep, markers_k, pc_scores, y, n = _prepare(
|
||
phenos, dosage, markers_meta, maf_min, n_pc)
|
||
if method == "emmax":
|
||
return _emmax(M, p, maf, m_keep, markers_k, pc_scores, y, n, m, markers_meta,
|
||
maf_min, n_pc, sig_level, qtl_window)
|
||
if method == "ssgwas":
|
||
return _ssgwas(M, p, maf, m_keep, markers_k, y, n, m, markers_meta,
|
||
maf_min, sig_level, qtl_window, common, phenos, dosage)
|
||
|
||
df = n - (n_pc + 2)
|
||
snps: list[dict] = []
|
||
X_cov = np.column_stack([np.ones(n), pc_scores[:, :min(n_pc, pc_scores.shape[1])]])
|
||
for j, mk in enumerate(markers_k):
|
||
X = np.column_stack([X_cov, M[:, j]])
|
||
try:
|
||
beta, *_ = np.linalg.lstsq(X, y, rcond=None)
|
||
except np.linalg.LinAlgError:
|
||
continue
|
||
resid = y - X @ beta
|
||
sigma2 = float(resid @ resid) / df
|
||
XtX = X.T @ X
|
||
# 微小岭回归兜底：标记与 PC 共线时 XtX 奇异，inv 返回病态（可致负对角元 → sqrt 域错）。
|
||
# reg=1e-8·max(diag) 令共线方向 SE 大而有限 → t 小 → p≈1（效应被 PC 吸收，正确非显著）。
|
||
reg = 1e-8 * float(np.max(np.diag(XtX)))
|
||
XtX_inv = np.linalg.inv(XtX + reg * np.eye(XtX.shape[0]))
|
||
se = math.sqrt(sigma2 * float(XtX_inv[-1, -1])) if sigma2 > 0 else 0.0
|
||
if not math.isfinite(se) or se < 1e-10:
|
||
t, pval = 0.0, 1.0
|
||
else:
|
||
t = float(beta[-1]) / se
|
||
pval = float(fdist.f_pvalue(t * t, 1, df))
|
||
meta = markers_meta.get(mk) or {}
|
||
snps.append({
|
||
"marker": mk,
|
||
"chromosome": (meta.get("chromosome") if meta.get("chromosome") is not None else "NA"),
|
||
"position": int(meta.get("position")) if meta.get("position") is not None else 0,
|
||
"maf": round(float(maf[j]), 6),
|
||
"effect": round(float(beta[-1]), 6),
|
||
"se": round(se, 6),
|
||
"t": round(t, 6),
|
||
"p": pval,
|
||
"neg_log10p": -math.log10(pval) if pval > 0 else 300.0,
|
||
})
|
||
if not snps:
|
||
raise ValueError("标记回归全部失败（数据奇异）")
|
||
|
||
threshold_bonf, qtls = _finalize(snps, m_keep, sig_level, qtl_window)
|
||
return {
|
||
"snps": snps,
|
||
"qtls": qtls,
|
||
"n_individuals": n,
|
||
"n_markers": m,
|
||
"m_after_maf": m_keep,
|
||
"maf_min": maf_min,
|
||
"n_pc": n_pc,
|
||
"sig_level": sig_level,
|
||
"threshold_bonf": round(threshold_bonf, 10),
|
||
"n_sig_bonf": int(sum(s["sig_bonf"] for s in snps)),
|
||
"n_sig_fdr": int(sum(s["sig_fdr"] for s in snps)),
|
||
"n_qtl": len(qtls),
|
||
"engine_version": ENGINE_VERSION,
|
||
}
|
||
|
||
|
||
def _ssgwas(M: np.ndarray, p: np.ndarray, maf: np.ndarray, m_keep: int,
|
||
markers_k: list[str], y: np.ndarray, n: int, m: int,
|
||
markers_meta: dict[str, dict[str, object]], maf_min: float,
|
||
sig_level: float, qtl_window: int, common: list[str],
|
||
phenos: dict[str, float], dosage: dict[str, dict[str, float]]) -> dict:
|
||
"""ssGWAS 单步 GWAS（Wang et al. 2012 MEM 对角近似）。
|
||
|
||
- ① VanRaden G（build_g_matrix **method=1**：G=ZZ'/k，diag≈1，h² 性状尺度）→ GBLUP
|
||
（solve_gblup）拟合全样本 GEBV û（common=表型∩基因型交集，全部有基因型+表型）；
|
||
- ② 标记效应反推（MEM 对角近似，Wang et al. 2012）：ê_j = (Z_j' G⁻¹ û) / k、
|
||
Var(ê_j) = σ²u / k，k = 2Σp_j(1−p_j)（method-1 G 的 VanRaden 缩放分母）；
|
||
Wald t=ê/SE、df=n−2 → fdist.f_pvalue；
|
||
- ③ 显著判定/Bonferroni/BH-q/QTL 聚类与 GLM/EMMAX 共用 _finalize（逐字节同后处理）。
|
||
|
||
注意：G 必须用 method=1（G=ZZ'/k，对角≈1）。method=2（ZDZ'）对角≈标记数 m，
|
||
令 h² 与 σ²u 落在 G 尺度上，与 Wang 公式的性状尺度失配（效应被整体压低）。
|
||
"""
|
||
G, _gmeta = genomic.build_g_matrix(
|
||
{i: dosage[i] for i in common}, markers_k, maf_min=maf_min, method=1, blend=0.02)
|
||
res = genomic.solve_gblup(common, phenos, G)
|
||
if not res.get("converged") or not res.get("sigma_a") or float(res["sigma_a"]) <= 0:
|
||
raise ValueError("ssGWAS：GBLUP 拟合失败或加性方差为 0，无法反推标记效应")
|
||
va = float(res["sigma_a"])
|
||
uhat = np.array([float(res["ebv"][i]) for i in common])
|
||
try:
|
||
Ginv = np.linalg.inv(G)
|
||
except np.linalg.LinAlgError:
|
||
Ginv = np.linalg.pinv(G)
|
||
Z = M - 2.0 * p[None, :]
|
||
k = 2.0 * float(np.sum(p * (1.0 - p)))
|
||
if k <= 0:
|
||
raise ValueError("ssGWAS：标记等位频率总和为 0，无法反推标记效应")
|
||
Ginvu = Ginv @ uhat
|
||
df = n - 2
|
||
var_effect = va / k
|
||
|
||
snps: list[dict] = []
|
||
for j, mk in enumerate(markers_k):
|
||
effect = float(Z[:, j] @ Ginvu) / k
|
||
se = math.sqrt(var_effect) if var_effect > 0 else 0.0
|
||
if not math.isfinite(se) or se < 1e-12:
|
||
t, pval = 0.0, 1.0
|
||
else:
|
||
t = effect / se
|
||
pval = float(fdist.f_pvalue(t * t, 1, df))
|
||
meta = markers_meta.get(mk) or {}
|
||
snps.append({
|
||
"marker": mk,
|
||
"chromosome": (meta.get("chromosome") if meta.get("chromosome") is not None else "NA"),
|
||
"position": int(meta.get("position")) if meta.get("position") is not None else 0,
|
||
"maf": round(float(maf[j]), 6),
|
||
"effect": round(effect, 6),
|
||
"se": round(se, 6),
|
||
"t": round(t, 6),
|
||
"p": pval,
|
||
"neg_log10p": -math.log10(pval) if pval > 0 else 300.0,
|
||
})
|
||
if not snps:
|
||
raise ValueError("ssGWAS：标记效应反推全部失败（数据奇异）")
|
||
|
||
threshold_bonf, qtls = _finalize(snps, m_keep, sig_level, qtl_window)
|
||
return {
|
||
"snps": snps,
|
||
"qtls": qtls,
|
||
"n_individuals": n,
|
||
"n_markers": m,
|
||
"m_after_maf": m_keep,
|
||
"maf_min": maf_min,
|
||
"n_pc": 0,
|
||
"sig_level": sig_level,
|
||
"threshold_bonf": round(threshold_bonf, 10),
|
||
"n_sig_bonf": int(sum(s["sig_bonf"] for s in snps)),
|
||
"n_sig_fdr": int(sum(s["sig_fdr"] for s in snps)),
|
||
"n_qtl": len(qtls),
|
||
"method": "ssgwas",
|
||
"h2": round(float(res["h2"]), 4) if res.get("h2") is not None else None,
|
||
"sigma_a": round(va, 6),
|
||
"sigma_e": round(float(res["sigma_e"]), 6) if res.get("sigma_e") is not None else None,
|
||
"engine_version": ENGINE_VERSION,
|
||
}
|
||
|
||
|
||
def _emmax(M: np.ndarray, p: np.ndarray, maf: np.ndarray, m_keep: int,
|
||
markers_k: list[str], pc_scores: np.ndarray, y: np.ndarray, n: int, m: int,
|
||
markers_meta: dict[str, dict[str, object]], maf_min: float, n_pc: int,
|
||
sig_level: float, qtl_window: int) -> dict:
|
||
"""EMMAX 混合模型关联（Kang et al. 2010 常规化）：零模型方差分量 → 固定 V 逐标记 GLS。
|
||
|
||
- 零模型 y = Xβ + u + e，u~N(0, σg²·K)，K=VanRaden GRM（M 中心化，blend 岭正定）；
|
||
方差分量复用 genomic._profile_solve（剖面 REML，ZGZ=K、Ginv=K⁻¹、z=全样本）。
|
||
- 固定 V=σg²K+σe²I：纯 numpy eigh 对角化 K → W⁻½=Q(λΛ+I)⁻½Q'，旋转 y*/X*/标记列后 OLS；
|
||
t²~F(1, df)，p 复用 fdist.f_pvalue，显著判定/BH-q/QTL 与 GLM 共用 _finalize。
|
||
"""
|
||
Z = M - 2.0 * p[None, :]
|
||
denom = 2.0 * float(np.sum(p * (1.0 - p)))
|
||
K = (Z @ Z.T) / denom if denom > 0 else np.eye(n)
|
||
blend = 0.02
|
||
K = (1.0 - blend) * K + blend * float(np.mean(np.diag(K))) * np.eye(n)
|
||
X = np.column_stack([np.ones(n), pc_scores[:, :min(n_pc, pc_scores.shape[1])]])
|
||
try:
|
||
Kinv = np.linalg.inv(K)
|
||
except np.linalg.LinAlgError:
|
||
Kinv = np.linalg.pinv(K)
|
||
null = genomic._profile_solve(K, X, y, Kinv, n, np.arange(n))
|
||
va = float(null["sigma_a"] or 0.0)
|
||
ve = float(null["sigma_e"] or 0.0)
|
||
h2 = float(null["h2"]) if null.get("h2") is not None else None
|
||
if va <= 0 or ve <= 0:
|
||
va, ve = max(va, 1e-6), max(ve, 1e-6)
|
||
lam = va / ve
|
||
|
||
evals, Q = np.linalg.eigh(K)
|
||
d = np.maximum(lam * evals + 1.0, 1e-12)
|
||
Winv_half = Q @ np.diag(1.0 / np.sqrt(d)) @ Q.T
|
||
ystar = Winv_half @ y
|
||
Xstar = Winv_half @ X
|
||
df = n - (n_pc + 2)
|
||
|
||
snps: list[dict] = []
|
||
for j, mk in enumerate(markers_k):
|
||
mstar = Winv_half @ M[:, j]
|
||
Xj = np.column_stack([Xstar, mstar])
|
||
try:
|
||
beta, *_ = np.linalg.lstsq(Xj, ystar, rcond=None)
|
||
except np.linalg.LinAlgError:
|
||
continue
|
||
resid = ystar - Xj @ beta
|
||
sigma2 = float(resid @ resid) / df
|
||
XtX = Xj.T @ Xj
|
||
reg = 1e-8 * float(np.max(np.diag(XtX)))
|
||
XtX_inv = np.linalg.inv(XtX + reg * np.eye(XtX.shape[0]))
|
||
se = math.sqrt(sigma2 * float(XtX_inv[-1, -1])) if sigma2 > 0 else 0.0
|
||
if not math.isfinite(se) or se < 1e-10:
|
||
t, pval = 0.0, 1.0
|
||
else:
|
||
t = float(beta[-1]) / se
|
||
pval = float(fdist.f_pvalue(t * t, 1, df))
|
||
meta = markers_meta.get(mk) or {}
|
||
snps.append({
|
||
"marker": mk,
|
||
"chromosome": (meta.get("chromosome") if meta.get("chromosome") is not None else "NA"),
|
||
"position": int(meta.get("position")) if meta.get("position") is not None else 0,
|
||
"maf": round(float(maf[j]), 6),
|
||
"effect": round(float(beta[-1]), 6),
|
||
"se": round(se, 6),
|
||
"t": round(t, 6),
|
||
"p": pval,
|
||
"neg_log10p": -math.log10(pval) if pval > 0 else 300.0,
|
||
})
|
||
if not snps:
|
||
raise ValueError("标记回归全部失败（数据奇异）")
|
||
|
||
threshold_bonf, qtls = _finalize(snps, m_keep, sig_level, qtl_window)
|
||
return {
|
||
"snps": snps,
|
||
"qtls": qtls,
|
||
"n_individuals": n,
|
||
"n_markers": m,
|
||
"m_after_maf": m_keep,
|
||
"maf_min": maf_min,
|
||
"n_pc": n_pc,
|
||
"sig_level": sig_level,
|
||
"threshold_bonf": round(threshold_bonf, 10),
|
||
"n_sig_bonf": int(sum(s["sig_bonf"] for s in snps)),
|
||
"n_sig_fdr": int(sum(s["sig_fdr"] for s in snps)),
|
||
"n_qtl": len(qtls),
|
||
"method": "emmax",
|
||
"h2": round(h2, 4) if h2 is not None else None,
|
||
"sigma_a": round(va, 6),
|
||
"sigma_e": round(ve, 6),
|
||
"engine_version": ENGINE_VERSION,
|
||
}
|
||
|
||
|
||
def _cluster(sigs: list[dict], qtl_window: int) -> list[list[dict]]:
|
||
"""显著标记按同染色体相邻间距合并簇（间距 ≤ qtl_window 同簇）。"""
|
||
groups: list[list[dict]] = []
|
||
for s in sigs:
|
||
if groups and groups[-1][-1]["chromosome"] == s["chromosome"] and (
|
||
s["position"] - groups[-1][-1]["position"] <= qtl_window):
|
||
groups[-1].append(s)
|
||
else:
|
||
groups.append([s])
|
||
return groups
|