721 lines
33 KiB
Python
721 lines
33 KiB
Python
"""基因组选择：VanRaden G 矩阵 + GBLUP + ssGBLUP（仅 numpy，无 R）。
|
||
|
||
- build_g_matrix：dosage(0/1/2) 样本×标记矩阵 → VanRaden G（method1=ZZ'/2Σp(1-p)、
|
||
method2=ZDZ'，D=1/[2p(1-p)]），缺格按列均值(2p)填充；MAF 过滤 + blend 岭回归保证可逆。
|
||
- solve_gblup：仅基因型个体入模型，MME 用 G⁻¹ 替代 A⁻¹；方差组分用剖面 REML
|
||
（黄金分割最大化精确 REML 对数似然，与 blup._solve_gxe 同模板）。
|
||
- solve_ssgblup：单步法 H⁻¹ = A⁻¹ + [[0,0],[0,G⁻¹ − A22⁻¹]]（Aguilar et al. 2010；
|
||
A22 = 基因型个体在 A 中的子阵），基因型+非基因型个体都出 EBV。
|
||
|
||
可靠性 = 1 − PEV/σ²a，PEV = C22[ii]·σ²e（C 为 MME 系数矩阵逆的加性块）。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import math
|
||
import random
|
||
|
||
import numpy as np
|
||
|
||
from scripts.breeding_stats import blup
|
||
|
||
ENGINE_VERSION = "1.4.0" # 1.4.0: 非线性基因组预测 rrBLUP + BayesB; 1.3.0: kfold split=family; 1.2.0: SSR 多等位 + EMMAX GRM 复用
|
||
_RIDGE = 1e-6 # A22⁻¹ 病态加小岭
|
||
|
||
_GBLUP_KEYS = ["ebv", "reliability", "h2", "sigma_a", "sigma_e",
|
||
"n_obs", "n_individuals", "n_markers", "converged", "warning"]
|
||
|
||
|
||
def build_g_matrix(genotypes: dict[str, dict[str, float]],
|
||
markers: list[str], *, maf_min: float = 0.05,
|
||
method: int = 2, blend: float = 0.02):
|
||
"""VanRaden G 矩阵。
|
||
|
||
genotypes: {sample: {marker: dosage(0/1/2)}}；markers: 有序标记列表。
|
||
返回 (G 按 sorted(genotypes) 排序, meta{G 统计量})。
|
||
"""
|
||
samples = sorted(genotypes)
|
||
n = len(samples)
|
||
if n < 2 or not markers:
|
||
raise ValueError("基因型样本或标记数不足（需 n≥2、m≥1）")
|
||
m = len(markers)
|
||
M = np.full((n, m), np.nan)
|
||
n_missing = 0
|
||
for i, s in enumerate(samples):
|
||
g = genotypes[s]
|
||
for j, mk in enumerate(markers):
|
||
v = g.get(mk)
|
||
if v is None or (isinstance(v, float) and math.isnan(v)):
|
||
n_missing += 1
|
||
else:
|
||
M[i, j] = float(v)
|
||
p = np.zeros(m)
|
||
for j in range(m):
|
||
col = M[:, j]
|
||
valid = col[~np.isnan(col)]
|
||
if len(valid) == 0:
|
||
p[j] = 0.5
|
||
else:
|
||
p[j] = float(np.mean(valid)) / 2.0
|
||
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]
|
||
m_keep = int(keep.sum())
|
||
Z = M - 2.0 * p[None, :]
|
||
if method == 1:
|
||
denom = 2.0 * float(np.sum(p * (1.0 - p)))
|
||
G = (Z @ Z.T) / denom if denom > 0 else np.eye(n)
|
||
else:
|
||
D = 1.0 / (2.0 * p * (1.0 - p))
|
||
G = Z @ np.diag(D) @ Z.T
|
||
if blend and blend > 0:
|
||
diag_mean = float(np.mean(np.diag(G)))
|
||
G = (1.0 - blend) * G + blend * diag_mean * np.eye(n)
|
||
off = G[np.triu_indices(n, 1)] if n > 1 else np.array([0.0])
|
||
return G, {
|
||
"n_individuals": n,
|
||
"n_markers": m,
|
||
"m_after_maf": m_keep,
|
||
"n_missing": n_missing,
|
||
"maf_min": maf_min,
|
||
"method": method,
|
||
"blend": blend,
|
||
"diag_mean": round(float(np.mean(np.diag(G))), 6),
|
||
"diag_range": [round(float(np.min(np.diag(G))), 6), round(float(np.max(np.diag(G))), 6)],
|
||
"offdiag_range": [round(float(np.min(off)), 6), round(float(np.max(off)), 6)],
|
||
}
|
||
|
||
|
||
def _scale_g_to_a22(G: np.ndarray, A22: np.ndarray) -> np.ndarray:
|
||
"""把基因关系阵 G 缩放到系谱关系阵 A22 的尺度（Christensen et al. 调整）。
|
||
|
||
VanRaden method=2 的 G 对角线≈标记数 m（远大于 A22 的 ≈1），若直接组 H，
|
||
基因型块会被 G 主导、h²/EBV 尺度偏置。用两参数对齐 G*=a·G+b 使
|
||
对角均值/非对角均值与 A22 一致；G 退化（对角≈非对角）时跳过并原样返回。
|
||
"""
|
||
n = len(A22)
|
||
if n < 2:
|
||
return G
|
||
diag_g = np.mean(np.diag(G))
|
||
diag_a = np.mean(np.diag(A22))
|
||
off_g = float(np.mean(G[np.triu_indices(n, 1)]))
|
||
off_a = float(np.mean(A22[np.triu_indices(n, 1)]))
|
||
denom = diag_g - off_g
|
||
if abs(denom) < 1e-12:
|
||
return G
|
||
a = (diag_a - off_a) / denom
|
||
b = off_a - a * off_g
|
||
return a * G + b
|
||
|
||
|
||
def _reml_ll(ratio: float, ZGZ: np.ndarray, X: np.ndarray, y: np.ndarray) -> float:
|
||
"""V = ZGZ·ratio + I（Ve=1 尺度），精确 REML 对数似然（slogdet + y'Py）。"""
|
||
m = len(y)
|
||
V = ZGZ * ratio + np.eye(m)
|
||
try:
|
||
VinvX = np.linalg.solve(V, X)
|
||
Vinvy = np.linalg.solve(V, y)
|
||
XtVinvX = X.T @ VinvX
|
||
XtVinvY = X.T @ Vinvy
|
||
yPy = float(y @ Vinvy) - float(XtVinvY @ np.linalg.solve(XtVinvX, XtVinvY))
|
||
_, lv = np.linalg.slogdet(V)
|
||
_, lx = np.linalg.slogdet(XtVinvX)
|
||
return -0.5 * (float(lv) + float(lx) + yPy)
|
||
except np.linalg.LinAlgError:
|
||
return -np.inf
|
||
|
||
|
||
def _profile_ve(V: np.ndarray, X: np.ndarray, y: np.ndarray, df: int) -> float:
|
||
"""REML 剖面 σ²e = y'Py/df（V 固定尺度下）。"""
|
||
VinvX = np.linalg.solve(V, X)
|
||
Vinvy = np.linalg.solve(V, y)
|
||
XtVinvX = X.T @ VinvX
|
||
XtVinvY = X.T @ Vinvy
|
||
yPy = float(y @ Vinvy) - float(XtVinvY @ np.linalg.solve(XtVinvX, XtVinvY))
|
||
return yPy / df if df > 0 else float("nan")
|
||
|
||
|
||
def _profile_solve(ZGZ: np.ndarray, X: np.ndarray, y: np.ndarray,
|
||
Ginv: np.ndarray, n_all: int, z: np.ndarray):
|
||
"""共享剖面 REML + MME 求解。返回 (ebv, rel, h2, sigma_a, sigma_e, n_iter)。"""
|
||
m = len(y)
|
||
df = m - X.shape[1]
|
||
# 剖面参数为 log10(σ²a/σ²e)。h²=ratio/(1+ratio)，故 h²∈[H2_MIN,H2_MAX]
|
||
# 对应 ratio∈[H2_MIN/(1−H2_MIN), H2_MAX/(1−H2_MAX)]（直接用 log10(H2_MAX)≈−0.0004
|
||
# 会把 σ²a/σ²e 封顶 0.999，令 h² 上限只有 0.5）。
|
||
lo, hi = math.log10(blup.H2_MIN / (1.0 - blup.H2_MIN)), math.log10(blup.H2_MAX / (1.0 - blup.H2_MAX))
|
||
|
||
def _ll(r: float) -> float:
|
||
return _reml_ll(10.0 ** r, ZGZ, X, y)
|
||
|
||
r_opt, n_iter = blup._golden_max(_ll, lo, hi, tol=1e-4, max_iter=blup.MAX_PROFILE_EVALS)
|
||
ratio = 10.0 ** r_opt
|
||
V = ZGZ * ratio + np.eye(m)
|
||
ve = _profile_ve(V, X, y, df)
|
||
va = ratio * ve if not math.isnan(ve) else 1.0
|
||
lam = ve / va if va > 0 else 1.0
|
||
|
||
Zmat = np.zeros((n_all, m))
|
||
Zmat[z, np.arange(m)] = 1.0
|
||
C = np.block([
|
||
[X.T @ X, X.T @ Zmat.T],
|
||
[Zmat @ X, Zmat @ Zmat.T + lam * Ginv],
|
||
])
|
||
rhs = np.concatenate([X.T @ y, Zmat @ y])
|
||
try:
|
||
sol = np.linalg.solve(C, rhs)
|
||
except np.linalg.LinAlgError:
|
||
return {"ebv": {}, "reliability": {}, "h2": None, "sigma_a": None,
|
||
"sigma_e": None, "n_iter": n_iter, "converged": False,
|
||
"warning": "MME 系数矩阵奇异（基因型共线/样本不足），求解失败"}
|
||
C_inv = np.linalg.inv(C)
|
||
C22inv = C_inv[1:, 1:]
|
||
rel = np.clip(1.0 - np.diag(C22inv) * ve / va, 0.0, 1.0) if va > 0 else np.zeros(n_all)
|
||
h2 = va / (va + ve) if (va + ve) > 0 else None
|
||
return {"ebv": sol[1:], "reliability": rel, "h2": h2, "sigma_a": float(va),
|
||
"sigma_e": float(ve) if not math.isnan(ve) else None,
|
||
"n_iter": n_iter, "converged": True, "warning": None}
|
||
|
||
|
||
def solve_gblup(individuals: list[str], phenos: dict[str, float], G: np.ndarray,
|
||
*, tol: float = blup.TOL):
|
||
"""GBLUP：仅基因型个体。individuals 为 G 行序（个体 id）；phenos 子集有观测。"""
|
||
n = len(individuals)
|
||
obs = [i for i in individuals if i in phenos]
|
||
m = len(obs)
|
||
if m < 3:
|
||
return {"ebv": {}, "reliability": {}, "h2": None, "sigma_a": None, "sigma_e": None,
|
||
"n_obs": m, "n_individuals": n, "n_markers": None,
|
||
"converged": False, "warning": "表型观测不足 3，无法估计方差组分"}
|
||
idx = {i: k for k, i in enumerate(individuals)}
|
||
z = np.array([idx[i] for i in obs])
|
||
y = np.array([float(phenos[i]) for i in obs])
|
||
X = np.ones((m, 1))
|
||
ZGZ = G[np.ix_(z, z)]
|
||
try:
|
||
Ginv = np.linalg.inv(G)
|
||
except np.linalg.LinAlgError:
|
||
Ginv = np.linalg.pinv(G)
|
||
res = _profile_solve(ZGZ, X, y, Ginv, n, z)
|
||
out = dict(res)
|
||
ebv = {individuals[i]: float(out["ebv"][i]) for i in range(n)}
|
||
rel = {individuals[i]: float(out["reliability"][i]) for i in range(n)}
|
||
out.update(ebv=ebv, reliability=rel, n_obs=m, n_individuals=n,
|
||
n_markers=None if not isinstance(G, np.ndarray) else G.shape[0])
|
||
if out["warning"] is None and out.get("h2") is not None and (out["h2"] >= blup.H2_MAX or out["h2"] <= blup.H2_MIN):
|
||
out["warning"] = "遗传力估计逼近边界，剖面 REML 未完全收敛（样本小/信噪比低）"
|
||
return out
|
||
|
||
|
||
def _build_h(pedigree: list[dict], genotyped: list[str], G: np.ndarray,
|
||
*, ridge: float = _RIDGE, extra_base: list[str] | None = None):
|
||
"""ssGBLUP 合并关系阵（H 协方差 + H⁻¹ 仅进 MME，Aguilar et al. 2010）。
|
||
|
||
返回 (order, idx, H, Hinv)。系谱中缺失的基因型个体按 base 补入（与 phenos 兜底一致）。
|
||
"""
|
||
base: list[int] = []
|
||
non_base: list[tuple[int, int | None, int | None]] = []
|
||
seen: set[int] = set()
|
||
for rec in pedigree:
|
||
i = rec["individual"]
|
||
if i in seen:
|
||
continue
|
||
seen.add(i)
|
||
d, s = rec.get("dam"), rec.get("sire")
|
||
if d is None and s is None:
|
||
base.append(i)
|
||
else:
|
||
non_base.append((i, d, s))
|
||
if extra_base:
|
||
for g in extra_base:
|
||
if g not in seen:
|
||
seen.add(g)
|
||
base.append(g)
|
||
for g in genotyped: # 无表型/无系谱基因型树也按 base 补入 H（不再 raise）
|
||
if g not in seen:
|
||
seen.add(g)
|
||
base.append(g)
|
||
if not seen:
|
||
raise ValueError("系谱与基因型均为空")
|
||
order, idx = blup._order_pedigree(base, non_base)
|
||
n = len(order)
|
||
parent_of = {i: (d, s) for (i, d, s) in non_base}
|
||
A, Ainv = blup._build_ainv(order, len(base), parent_of)
|
||
|
||
gen_indices = [idx[g] for g in genotyped]
|
||
if len(gen_indices) < 2:
|
||
raise ValueError("基因型个体不足 2，无法 ssGBLUP")
|
||
A22 = A[np.ix_(gen_indices, gen_indices)]
|
||
G = _scale_g_to_a22(G, A22)
|
||
try:
|
||
A22inv = np.linalg.inv(A22 + ridge * np.eye(len(gen_indices)))
|
||
Ginv = np.linalg.inv(G)
|
||
except np.linalg.LinAlgError:
|
||
A22inv = np.linalg.pinv(A22)
|
||
Ginv = np.linalg.pinv(G)
|
||
Hinv = Ainv.copy()
|
||
Hinv[np.ix_(gen_indices, gen_indices)] += Ginv - A22inv
|
||
|
||
# H 合并关系阵（协方差；H⁻¹=A⁻¹+[[0,0],[0,G⁻¹−A22⁻¹]]，Aguilar et al. 2010）
|
||
# H = A + [[A12·A22⁻¹·D·A22⁻¹·A21, A12·A22⁻¹·D], [D·A22⁻¹·A21, D]]，D=G−A22
|
||
gen_set = set(gen_indices)
|
||
non_gen = [i for i in range(n) if i not in gen_set]
|
||
if non_gen:
|
||
D = G - A22
|
||
L = A[np.ix_(non_gen, gen_indices)] @ A22inv
|
||
LD = L @ D
|
||
H = A.copy()
|
||
H[np.ix_(non_gen, non_gen)] += LD @ L.T
|
||
H[np.ix_(non_gen, gen_indices)] += LD
|
||
H[np.ix_(gen_indices, non_gen)] += LD.T
|
||
H[np.ix_(gen_indices, gen_indices)] = G
|
||
else:
|
||
H = G
|
||
return order, idx, H, Hinv
|
||
|
||
|
||
def solve_ssgblup(pedigree: list[dict], phenos: dict[str, float],
|
||
G: np.ndarray, genotyped: list[str], *, tol: float = blup.TOL,
|
||
ridge: float = _RIDGE):
|
||
"""ssGBLUP：单步法。genotyped 为有序列表（与 G 行对应）。全系谱个体出 EBV。"""
|
||
if not pedigree and not phenos:
|
||
return {"ebv": {}, "reliability": {}, "h2": None, "sigma_a": None, "sigma_e": None,
|
||
"n_obs": 0, "n_individuals": 0, "n_markers": None,
|
||
"converged": False, "warning": "系谱与表型均为空"}
|
||
try:
|
||
order, idx, H, Hinv = _build_h(
|
||
pedigree, genotyped, G, ridge=ridge, extra_base=list(phenos))
|
||
except ValueError as e:
|
||
return {"ebv": {}, "reliability": {}, "h2": None, "sigma_a": None, "sigma_e": None,
|
||
"n_obs": 0, "n_individuals": len({r["individual"] for r in pedigree}),
|
||
"n_markers": None, "converged": False, "warning": str(e)}
|
||
n = len(order)
|
||
|
||
obs = [i for i in order if i in phenos]
|
||
m = len(obs)
|
||
if m < 3:
|
||
return {"ebv": {}, "reliability": {}, "h2": None, "sigma_a": None, "sigma_e": None,
|
||
"n_obs": m, "n_individuals": n, "n_markers": None,
|
||
"converged": False, "warning": "表型观测不足 3，无法估计方差组分"}
|
||
z = np.array([idx[i] for i in obs])
|
||
y = np.array([float(phenos[i]) for i in obs])
|
||
X = np.ones((m, 1))
|
||
ZGZ = H[np.ix_(z, z)] # V 用 H 协方差块（H⁻¹ 仅进 MME 系数矩阵）
|
||
res = _profile_solve(ZGZ, X, y, Hinv, n, z)
|
||
out = dict(res)
|
||
ebv = {order[i]: float(out["ebv"][i]) for i in range(n)}
|
||
rel = {order[i]: float(out["reliability"][i]) for i in range(n)}
|
||
out.update(ebv=ebv, reliability=rel, n_obs=m, n_individuals=n, n_markers=len(genotyped))
|
||
if out["warning"] is None and out.get("h2") is not None and (out["h2"] >= blup.H2_MAX or out["h2"] <= blup.H2_MIN):
|
||
out["warning"] = "遗传力估计逼近边界，剖面 REML 未完全收敛"
|
||
return out
|
||
|
||
|
||
def _build_dosage_matrix(genotypes: dict[str, dict[str, float]],
|
||
markers: list[str], *, maf_min: float = 0.05):
|
||
"""dosage dict → 有序样本×标记剂量矩阵（与 build_g_matrix 同 MAF 过滤口径）。
|
||
|
||
返回 (X[n×q], 保留标记列表, 等位频率 p[q])。X 已按列均值(≈2p)填充缺失。
|
||
rrBLUP/BayesB 直接消费 X（Xc = X − 2p 即 GBLUP 的 Z）。
|
||
"""
|
||
samples = sorted(genotypes)
|
||
n = len(samples)
|
||
m = len(markers)
|
||
M = np.full((n, m), np.nan)
|
||
for i, s in enumerate(samples):
|
||
g = genotypes[s]
|
||
for j, mk in enumerate(markers):
|
||
v = g.get(mk)
|
||
if v is not None and not (isinstance(v, float) and math.isnan(v)):
|
||
M[i, j] = float(v)
|
||
p = np.zeros(m)
|
||
for j in range(m):
|
||
col = M[:, j]
|
||
valid = col[~np.isnan(col)]
|
||
p[j] = 0.5 if len(valid) == 0 else float(np.mean(valid)) / 2.0
|
||
M[:, j] = np.where(np.isnan(M[:, j]), 2.0 * p[j], M[:, j])
|
||
keep = np.minimum(p, 1.0 - p) >= maf_min
|
||
if int(keep.sum()) == 0:
|
||
raise ValueError(f"MAF≥{maf_min} 的标记为 0，请降低 maf_min 或检查标记等位频率")
|
||
return M[:, keep], [mk for mk, k in zip(markers, keep) if k], p[keep]
|
||
|
||
|
||
def solve_rrblup(individuals: list[str], phenos: dict[str, float],
|
||
X: np.ndarray, marker_order: list[str], p: np.ndarray,
|
||
*, tol: float = blup.TOL):
|
||
"""RR-BLUP（岭回归逐标记）：y = μ + Xu + e，u=(Xc'Xc+λI)⁻¹Xc'yc。
|
||
|
||
与 GBLUP 数学对偶（G = XcXc'/c 时 EBV 一致）：方差组分仍用剖面 REML
|
||
（_profile_solve 同一求解器）估 σa²/σe² → σu²=σa²/c（c=2Σp(1-p)）→
|
||
λ=σe²/σu²。输出 EBV（基因型个体，含表型截距）+ marker effects（GBLUP
|
||
无的解释能力）+ 与 GBLUP 同构的可靠性（PEV 近似）。
|
||
"""
|
||
n = len(individuals)
|
||
obs = [i for i in individuals if i in phenos]
|
||
m = len(obs)
|
||
if m < 3:
|
||
return {"ebv": {}, "reliability": {}, "h2": None, "sigma_a": None, "sigma_e": None,
|
||
"n_obs": m, "n_individuals": n, "n_markers": None, "marker_effects": {},
|
||
"converged": False, "warning": "表型观测不足 3，无法估计方差组分"}
|
||
idx = {i: k for k, i in enumerate(individuals)}
|
||
z = np.array([idx[i] for i in obs])
|
||
y = np.array([float(phenos[i]) for i in obs])
|
||
Xc = X - 2.0 * p[None, :] # 中心化剂量 = GBLUP 的 Z
|
||
q = Xc.shape[1]
|
||
c = 2.0 * float(np.sum(p * (1.0 - p)))
|
||
if c <= 0:
|
||
c = 1.0
|
||
G = Xc @ Xc.T / c
|
||
try:
|
||
Ginv = np.linalg.inv(G)
|
||
except np.linalg.LinAlgError:
|
||
Ginv = np.linalg.pinv(G)
|
||
res = _profile_solve(G[np.ix_(z, z)], np.ones((m, 1)), y, Ginv, n, z)
|
||
va, ve = res.get("sigma_a"), res.get("sigma_e")
|
||
if not res.get("converged") or va is None or ve is None or va <= 0 or ve <= 0:
|
||
return {"ebv": {}, "reliability": {}, "h2": res.get("h2"), "sigma_a": va,
|
||
"sigma_e": ve, "n_obs": m, "n_individuals": n, "n_markers": q,
|
||
"marker_effects": {},
|
||
"converged": False, "warning": res.get("warning") or "方差组分估计失败，rrBLUP 未收敛"}
|
||
lam = ve / (va / c) if va > 0 else 1.0 # λ = σe²/σu² = σe²·c/σa²
|
||
yc = y - float(y.mean())
|
||
Xo = Xc[z]
|
||
try:
|
||
u = np.linalg.solve(Xo.T @ Xo + lam * np.eye(q), Xo.T @ yc)
|
||
except np.linalg.LinAlgError:
|
||
u = np.linalg.pinv(Xo.T @ Xo + lam * np.eye(q)) @ (Xo.T @ yc)
|
||
mu = float(y.mean())
|
||
ebv = {individuals[i]: float(Xc[i] @ u + mu) for i in range(n)}
|
||
rel = res["reliability"]
|
||
out = {"ebv": ebv, "reliability": {individuals[i]: float(rel[i]) for i in range(n)},
|
||
"h2": res["h2"], "sigma_a": float(va), "sigma_e": float(ve),
|
||
"n_obs": m, "n_individuals": n, "n_markers": q,
|
||
"marker_effects": {marker_order[j]: float(u[j]) for j in range(q)},
|
||
"converged": True, "warning": res.get("warning")}
|
||
return out
|
||
|
||
|
||
def _bayesb_gibbs(Xo: np.ndarray, yc: np.ndarray, *, n_iter: int, burn_in: int,
|
||
seed: int, pi: float, df_s: float, scale0: float,
|
||
df_e: float, scale_e0: float) -> tuple[np.ndarray, float, float, int]:
|
||
"""BayesB 逐标记 Gibbs 核心：给定中心化剂量观测子阵 Xo 与中心化表型 yc。
|
||
|
||
返回 (u_hat 后验均值, σe² 末段, SSE 末段, 收敛样本数 n_save)。
|
||
numpy 固定 seed（default_rng）保证同数据同参数逐位可复现。
|
||
"""
|
||
q = Xo.shape[1]
|
||
m = len(yc)
|
||
rng = np.random.default_rng(seed)
|
||
u = np.zeros(q)
|
||
e = yc.astype(float).copy()
|
||
sigma_e2 = float(np.var(yc))
|
||
if sigma_e2 <= 0:
|
||
sigma_e2 = 1.0
|
||
u_acc = np.zeros(q)
|
||
n_save = 0
|
||
last_sse = float("inf")
|
||
for it in range(n_iter):
|
||
for j in range(q):
|
||
e_star = e + Xo[:, j] * u[j]
|
||
xjxj = float(Xo[:, j] @ Xo[:, j])
|
||
xjy = float(Xo[:, j] @ e_star)
|
||
if rng.uniform() > pi: # 1−π 概率标记入模型
|
||
s2 = (df_s * scale0 + u[j] * u[j]) / (rng.chisquare(df_s + 1))
|
||
if s2 <= 0 or xjxj <= 0:
|
||
u[j] = 0.0
|
||
else:
|
||
denom = xjxj + sigma_e2 / s2
|
||
u[j] = rng.normal(xjy / denom, math.sqrt(sigma_e2 / denom))
|
||
else:
|
||
u[j] = 0.0
|
||
e = e_star - Xo[:, j] * u[j]
|
||
sse = float(e @ e)
|
||
sigma_e2 = (df_e * scale_e0 + sse) / rng.chisquare(df_e + m)
|
||
if sigma_e2 <= 0:
|
||
sigma_e2 = 1e-8
|
||
if it >= burn_in:
|
||
u_acc += u
|
||
n_save += 1
|
||
last_sse = sse
|
||
if n_save == 0:
|
||
raise ValueError("后验样本为空（burn_in 覆盖全部迭代）")
|
||
return u_acc / n_save, sigma_e2, last_sse, n_save
|
||
|
||
|
||
def solve_bayesb(individuals: list[str], phenos: dict[str, float],
|
||
X: np.ndarray, marker_order: list[str], p: np.ndarray,
|
||
*, n_iter: int = 2000, burn_in: int = 500, seed: int = 20260805,
|
||
pi: float = 0.95, df_s: float = 5.0, scale0: float = 0.001,
|
||
df_e: float = 5.0, scale_e0: float = 1.0,
|
||
tol: float = blup.TOL):
|
||
"""BayesB（Meuwissen et al. 2001）：贝叶斯可变选择，逐标记 Gibbs 抽样。
|
||
|
||
模型 y = μ + Σ Xc_j·u_j + e；每标记 δ_j ~ Bernoulli(1−π) 概率入模型，
|
||
u_j | δ=1 ~ N(0, σ²j)、σ²j ~ inv-χ²(df_s, scale)；σe² ~ inv-χ²(df_e, scale_e0)。
|
||
固定 seed（numpy default_rng）保证可复现（MLOps 铁律：同数据同参数重跑逐位一致）。
|
||
输出 EBV（后验均值，基因型个体，含表型截距）+ marker effects + 收敛诊断。
|
||
可靠性未单独估计（后验方差才是正确不确定性度量），按 0 保守处理（note 说明）。
|
||
"""
|
||
n = len(individuals)
|
||
obs = [i for i in individuals if i in phenos]
|
||
m = len(obs)
|
||
if m < 3:
|
||
return {"ebv": {}, "reliability": {}, "h2": None, "sigma_a": None, "sigma_e": None,
|
||
"n_obs": m, "n_individuals": n, "n_markers": None, "marker_effects": {},
|
||
"converged": False, "warning": "表型观测不足 3，无法估计方差组分"}
|
||
if n_iter <= burn_in:
|
||
return {"ebv": {}, "reliability": {}, "h2": None, "sigma_a": None, "sigma_e": None,
|
||
"n_obs": m, "n_individuals": n, "n_markers": None, "marker_effects": {},
|
||
"converged": False, "warning": "n_iter 必须大于 burn_in"}
|
||
idx = {i: k for k, i in enumerate(individuals)}
|
||
z = np.array([idx[i] for i in obs])
|
||
y = np.array([float(phenos[i]) for i in obs])
|
||
Xc = X - 2.0 * p[None, :]
|
||
q = Xc.shape[1]
|
||
yc = y - float(y.mean())
|
||
Xo = Xc[z]
|
||
|
||
try:
|
||
u_hat, sigma_e2, sse, n_save = _bayesb_gibbs(
|
||
Xo, yc, n_iter=n_iter, burn_in=burn_in, seed=seed, pi=pi,
|
||
df_s=df_s, scale0=scale0, df_e=df_e, scale_e0=scale_e0)
|
||
except Exception as exc: # noqa: BLE001
|
||
return {"ebv": {}, "reliability": {}, "h2": None, "sigma_a": None, "sigma_e": None,
|
||
"n_obs": m, "n_individuals": n, "n_markers": q, "marker_effects": {},
|
||
"converged": False, "warning": f"BayesB 抽样异常: {exc!s}"}
|
||
mu = float(yc.mean())
|
||
ebv = {individuals[i]: float(Xc[i] @ u_hat + mu) for i in range(n)}
|
||
note = (f"BayesB π={pi}·{n_iter}iter·{burn_in}burn-in·seed={seed}（后验均值，"
|
||
f"收敛样本 {n_save}；σe² 末段 SSE={sse:.3f}）")
|
||
return {"ebv": ebv,
|
||
"reliability": {individuals[i]: 0.0 for i in range(n)},
|
||
"h2": None, "sigma_a": None, "sigma_e": float(sigma_e2),
|
||
"n_obs": m, "n_individuals": n, "n_markers": q,
|
||
"marker_effects": {marker_order[j]: float(u_hat[j]) for j in range(q)},
|
||
"converged": True, "warning": note}
|
||
|
||
|
||
def kfold_cv_genomic(phenos: dict[str, float], G: np.ndarray, genotyped: list[str],
|
||
*, k: int = 5, seed: int = 20260804, method: str = "gblup",
|
||
pedigree: list[dict] | None = None,
|
||
ridge: float = _RIDGE,
|
||
split: str = "random",
|
||
family_of: dict[str, str] | None = None,
|
||
genotypes: dict[str, dict[str, float]] | None = None,
|
||
markers: list[str] | None = None,
|
||
n_iter: int = 1000, burn_in: int = 300, pi: float = 0.95,
|
||
df_s: float = 5.0, scale0: float = 0.001,
|
||
df_e: float = 5.0, scale_e0: float = 1.0) -> dict:
|
||
"""GS k 折交叉验证（实证预测准确度：留出 GEBV vs 表型 pearson/RMSE）。
|
||
|
||
与 blup.kfold_cv 区别：关系矩阵为 VanRaden G（GBLUP）或单步 H（ssGBLUP），
|
||
折划分默认按**基因型个体随机分层**（固定种子；GS 同世代样本无系谱家系树，
|
||
随机分层即标准做法，区别于 ABLUP 的 sire 家系留出）。若样本确有家系结构
|
||
（同父半同胞），随机分层会把同家系亲缘个体摊入 train/test → 泄漏亲缘、
|
||
高估准确度，此时应传 `split="family"` + `family_of` 走**家系阻塞折**
|
||
（同家系个体同折，照 blup.kfold_cv 的 sire 家系留出模式）；无家系个体自成一族。
|
||
家系数 < 折数 k 时折叠为 n_families 折并附 warning。
|
||
|
||
method 分派：
|
||
- `gblup`：VanRaden G（折内 _profile_solve，掩蔽留出经 G⁻¹ 交叉关系预测）。
|
||
- `ssgblup`：单步 H（需 pedigree）。
|
||
- `rrblup`：岭回归逐标记（需 genotypes+markers 剂量矩阵；折内 REML 定 λ，
|
||
EBV=Xc·u+μ 直接由标记基因型预测，不经关系矩阵逆）。
|
||
- `bayesb`：贝叶斯可变选择（需 genotypes+markers；折内 Gibbs，固定 seed 可复现）。
|
||
rrblup/bayesb 缺 genotypes/markers 时返回含 warning 的空结果。
|
||
|
||
掩蔽留出：个体保留在 G/H 中，仅掩蔽其表型（_profile_solve 的 MME 非观测个体
|
||
经 G⁻¹/H⁻¹ 交叉关系预测）。G/H/Hinv 只建一次，逐折只换 z_train 子阵。
|
||
剂量矩阵（Xc）同样只建一次，逐折只取观测行。
|
||
|
||
输出 dict 与 blup.kfold_cv 逐键对齐 + `n_families`。cand=genotyped∩phenos（可验证个体）。
|
||
"""
|
||
method = (method or "gblup").lower()
|
||
if method == "ssgblup":
|
||
order, idx, H, Hinv = _build_h(
|
||
pedigree or [], genotyped, G, ridge=ridge, extra_base=list(phenos))
|
||
individuals = order
|
||
zmat_full = H
|
||
relmat_inv = Hinv
|
||
elif method in ("rrblup", "bayesb"):
|
||
if genotypes is None or markers is None:
|
||
return {
|
||
"k": k, "n_total": 0, "n_individuals": 0,
|
||
"folds": [], "mean_pearson": None, "mean_rmse": None,
|
||
"pooled_pearson": None, "pooled_rmse": None, "cv_accuracy": None,
|
||
"h2": None, "n_families": None,
|
||
"warning": f"{method} 需要 genotypes+markers 剂量矩阵，无法交叉验证",
|
||
}
|
||
X_full, marker_order, p_full = _build_dosage_matrix(genotypes, markers)
|
||
Xc_full = X_full - 2.0 * p_full[None, :]
|
||
marker_c = 2.0 * float(np.sum(p_full * (1.0 - p_full)))
|
||
if marker_c <= 0:
|
||
marker_c = 1.0
|
||
individuals = sorted(genotypes)
|
||
idx = {i: k for k, i in enumerate(individuals)}
|
||
# _profile_solve 需全 n×n Ginv（MME 随机效应块为 Zmat@Zmat.T + λ·Ginv_full），
|
||
# 折内只取 zmat_full[np.ix_(z,z)] 观测子阵（= XoXo'/c），与 GBLUP 分支同构。
|
||
zmat_full = Xc_full @ Xc_full.T / marker_c
|
||
try:
|
||
relmat_inv = np.linalg.inv(zmat_full)
|
||
except np.linalg.LinAlgError:
|
||
relmat_inv = np.linalg.pinv(zmat_full)
|
||
else:
|
||
individuals = list(genotyped)
|
||
idx = {i: k for k, i in enumerate(individuals)}
|
||
zmat_full = G
|
||
try:
|
||
relmat_inv = np.linalg.inv(G)
|
||
except np.linalg.LinAlgError:
|
||
relmat_inv = np.linalg.pinv(G)
|
||
n = len(individuals)
|
||
cand = sorted([i for i in genotyped if i in phenos])
|
||
n_total = len(cand)
|
||
if n_total < 2:
|
||
return {
|
||
"k": k, "n_total": n_total, "n_individuals": n,
|
||
"folds": [], "mean_pearson": None, "mean_rmse": None,
|
||
"pooled_pearson": None, "pooled_rmse": None, "cv_accuracy": None,
|
||
"h2": None, "warning": "可验证个体（有基因型+表型）不足 2，无法交叉验证",
|
||
}
|
||
k = max(2, min(int(k or 5), 10))
|
||
rng = random.Random(seed)
|
||
n_families: int | None = None
|
||
fam_warning: str | None = None
|
||
split = (split or "random").lower()
|
||
if split == "family":
|
||
fam_of = {ind: (family_of or {}).get(ind, ("self", ind)) for ind in cand}
|
||
fams = sorted({f for f in fam_of.values()})
|
||
n_families = len(fams)
|
||
if n_families < 2:
|
||
return {
|
||
"k": k, "n_total": n_total, "n_individuals": n,
|
||
"folds": [], "mean_pearson": None, "mean_rmse": None,
|
||
"pooled_pearson": None, "pooled_rmse": None, "cv_accuracy": None,
|
||
"h2": None, "n_families": n_families,
|
||
"warning": "可验证个体仅构成 1 个家系，无法家系阻塞交叉验证（请改用随机分层）",
|
||
}
|
||
if n_families < k:
|
||
k = n_families
|
||
fam_warning = f"家系数 {n_families} < 请求折数，折叠为 {k} 折（family-blocked）"
|
||
rng.shuffle(fams)
|
||
fam_bin = {f: i % k for i, f in enumerate(fams)}
|
||
bin_of = {c: fam_bin[fam_of[c]] for c in cand}
|
||
else:
|
||
order_shuf = list(cand)
|
||
rng.shuffle(order_shuf)
|
||
bin_of = {c: i % k for i, c in enumerate(order_shuf)}
|
||
|
||
folds: list[dict] = []
|
||
all_x: list[float] = []
|
||
all_y: list[float] = []
|
||
for fold in range(k):
|
||
test_set = {c for c, b in bin_of.items() if b == fold}
|
||
train_phenos = {i: v for i, v in phenos.items() if i not in test_set}
|
||
train_obs = [i for i in individuals if i in train_phenos]
|
||
if not test_set or len(train_obs) < 2:
|
||
folds.append({
|
||
"fold": fold + 1, "n_train": len(train_obs), "n_test": len(test_set),
|
||
"n_eval": 0, "pearson": None, "rmse": None,
|
||
"error": "训练集个体不足 2 或测试集为空",
|
||
})
|
||
continue
|
||
z = np.array([idx[i] for i in train_obs])
|
||
y = np.array([float(train_phenos[i]) for i in train_obs])
|
||
ebv: dict[str, float] | None = None
|
||
fold_h2: float | None = None
|
||
res: dict | None = None
|
||
try:
|
||
if method in ("rrblup", "bayesb"):
|
||
Xo = Xc_full[z]
|
||
yc = y - float(y.mean())
|
||
if method == "rrblup":
|
||
res = _profile_solve(zmat_full[np.ix_(z, z)], np.ones((len(train_obs), 1)),
|
||
y, relmat_inv, n, z)
|
||
va, ve = res.get("sigma_a"), res.get("sigma_e")
|
||
if (res.get("converged") and va is not None and ve is not None
|
||
and va > 0 and ve > 0):
|
||
lam = ve / (va / marker_c)
|
||
q = Xc_full.shape[1]
|
||
try:
|
||
u = np.linalg.solve(Xo.T @ Xo + lam * np.eye(q), Xo.T @ yc)
|
||
except np.linalg.LinAlgError:
|
||
u = np.linalg.pinv(Xo.T @ Xo + lam * np.eye(q)) @ (Xo.T @ yc)
|
||
mu = float(y.mean())
|
||
ebv = np.array([float(Xc_full[i] @ u + mu) for i in range(n)])
|
||
fold_h2 = res.get("h2")
|
||
else:
|
||
u_hat, _sigma_e2, _sse, _n_save = _bayesb_gibbs(
|
||
Xo, yc, n_iter=n_iter, burn_in=burn_in, seed=seed, pi=pi,
|
||
df_s=df_s, scale0=scale0, df_e=df_e, scale_e0=scale_e0)
|
||
mu = float(y.mean())
|
||
ebv = np.array([float(Xc_full[i] @ u_hat + mu) for i in range(n)])
|
||
res = {"h2": None, "converged": True}
|
||
else:
|
||
X = np.ones((len(train_obs), 1))
|
||
res = _profile_solve(zmat_full[np.ix_(z, z)], X, y, relmat_inv, n, z)
|
||
if res.get("converged"):
|
||
ebv = res["ebv"]
|
||
fold_h2 = res.get("h2")
|
||
except Exception as e: # noqa: BLE001
|
||
folds.append({
|
||
"fold": fold + 1, "n_train": len(train_obs), "n_test": len(test_set),
|
||
"n_eval": 0, "pearson": None, "rmse": None,
|
||
"error": f"求解失败: {e!s}",
|
||
})
|
||
continue
|
||
if not res or not res.get("converged") or ebv is None:
|
||
folds.append({
|
||
"fold": fold + 1, "n_train": len(train_obs), "n_test": len(test_set),
|
||
"n_eval": 0, "pearson": None, "rmse": None,
|
||
"error": (res or {}).get("warning") or "求解失败",
|
||
})
|
||
continue
|
||
xs = [float(ebv[idx[c]]) for c in sorted(test_set)]
|
||
ys = [float(phenos[c]) for c in sorted(test_set)]
|
||
if len(xs) < 2:
|
||
folds.append({
|
||
"fold": fold + 1, "n_train": len(train_obs), "n_test": len(test_set),
|
||
"n_eval": len(xs), "pearson": None, "rmse": None,
|
||
"error": "测试集有效个体不足 2，无法计算 pearson",
|
||
})
|
||
continue
|
||
pearson = blup._pearson(list(xs), list(ys))
|
||
rmse = math.sqrt(sum((a - b) ** 2 for a, b in zip(xs, ys)) / len(xs))
|
||
all_x.extend(xs)
|
||
all_y.extend(ys)
|
||
folds.append({
|
||
"fold": fold + 1, "n_train": len(train_obs), "n_test": len(test_set),
|
||
"n_eval": len(xs), "pearson": round(pearson, 4),
|
||
"rmse": round(rmse, 6),
|
||
"h2": float(fold_h2) if fold_h2 is not None else None,
|
||
})
|
||
|
||
valid = [f for f in folds if f.get("pearson") is not None]
|
||
mean_pearson = round(sum(f["pearson"] for f in valid) / len(valid), 4) if valid else None
|
||
rmse_vals = [f["rmse"] for f in valid if f.get("rmse") is not None]
|
||
mean_rmse = round(sum(rmse_vals) / len(rmse_vals), 6) if rmse_vals else None
|
||
pooled_pearson = round(blup._pearson(all_x, all_y), 4) if len(all_x) >= 2 else None
|
||
pooled_rmse = (round(math.sqrt(sum((a - b) ** 2 for a, b in zip(all_x, all_y)) / len(all_x)), 6)
|
||
if len(all_x) else None)
|
||
h2_vals = [f["h2"] for f in valid if f.get("h2") is not None]
|
||
h2_mean = round(sum(h2_vals) / len(h2_vals), 4) if h2_vals else None
|
||
errors = [f.get("error") for f in folds if f.get("error")]
|
||
warning_parts = []
|
||
if fam_warning:
|
||
warning_parts.append(fam_warning)
|
||
if errors:
|
||
warning_parts.append(f"{len(errors)} 折失败: " + "；".join(errors))
|
||
warning = "；".join(warning_parts) if warning_parts else None
|
||
return {
|
||
"k": k, "n_total": n_total, "n_individuals": n,
|
||
"folds": folds,
|
||
"mean_pearson": mean_pearson, "mean_rmse": mean_rmse,
|
||
"pooled_pearson": pooled_pearson, "pooled_rmse": pooled_rmse,
|
||
"cv_accuracy": mean_pearson, "h2": h2_mean,
|
||
"n_families": n_families, "warning": warning,
|
||
}
|