385 lines
16 KiB
Python
385 lines
16 KiB
Python
"""动物模型 MT-BLUP（多性状 BLUP / 遗传相关估计）。
|
||
|
||
仅依赖 numpy，复用 blup 的系谱 A/A⁻¹、单性状剖面 REML 与黄金分割一维最大化。
|
||
|
||
两条路径：
|
||
- solve_bivariate：成对双性状（降维：Va/Ve 取单性状 REML，仅对 ρ 做 1-D 剖面最大化）；
|
||
- solve_multi：全多变量 EM-REML（一次估计完整 G0⊗A，替代逐对，供 Smith-Hazel 指数）。
|
||
|
||
模型（成对双性状）：
|
||
y = [y1; y2]，Var(u) = G0 ⊗ A，G0 = [[Va1, ρ·√(Va1·Va2)], [·, Va2]]，
|
||
Var(e) = diag(Ve1·I, Ve2·I)。
|
||
|
||
降维策略（避免全多变量 REML 的收敛与识别风险）：
|
||
Va1/Ve1、Va2/Ve2 直接取各性状单性状 blup.solve 的 REML 估计（= 既有 herit 来源），
|
||
仅对遗传相关 ρ ∈ [-0.99, 0.99] 做粗网格预扫 + 黄金分割最大化精确 REML 对数似然
|
||
（V = Z(G0⊗A)Z' + R，slogdet + y'Py，与 _solve_gxe 模板一致）。
|
||
|
||
输出 σa12 = ρ·√(Va1·Va2)、r_g = ρ，供 Smith-Hazel 指数 G 矩阵非对角替换 Calo 近似。
|
||
单对不收敛 / 单性状求解失败 → converged=False + warning（调用方回退 Calo 并记审计）。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import math
|
||
|
||
import numpy as np
|
||
|
||
from scripts.breeding_stats import blup
|
||
|
||
ENGINE_VERSION = "1.1.0" # 1.1.0: 全多变量 EM-REML solve_multi（G0⊗A 整体估计）
|
||
_RHO_LO, _RHO_HI = -0.99, 0.99
|
||
_GRID_N = 21 # 粗网格预扫点数（防多峰取错局部峰）
|
||
|
||
_PAIR_KEYS = ["code1", "code2", "va1", "va2", "va12", "r_g", "ve1", "ve2",
|
||
"n_common", "n1", "n2", "n_individuals", "n_iter", "converged", "warning"]
|
||
|
||
|
||
def _pedigree_partition(pedigree: list[dict]):
|
||
"""系谱 -> (base, non_base, seen)。与 blup.solve 相同解析。"""
|
||
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:
|
||
raise ValueError(f"系谱个体重复: {i}")
|
||
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))
|
||
return base, non_base, seen
|
||
|
||
|
||
def _fail(code1: str, code2: str, msg: str) -> dict:
|
||
return {"code1": code1, "code2": code2, "va1": None, "va2": None, "va12": None,
|
||
"r_g": None, "ve1": None, "ve2": None, "n_common": 0, "n1": 0, "n2": 0,
|
||
"n_individuals": 0, "n_iter": 0, "converged": False, "warning": msg}
|
||
|
||
|
||
def solve_bivariate(pedigree: list[dict], phenos: dict[str, dict],
|
||
code1: str, code2: str, *, tol: float = blup.TOL) -> dict:
|
||
"""成对双性状 BLUP：估计 Va1/Va2/Va12（遗传相关 ρ）。
|
||
|
||
phenos: {code: {个体id: 均值}}（个体 id 空间与系谱一致，如 "t{tree_id}"）。
|
||
"""
|
||
p1 = phenos.get(code1, {})
|
||
p2 = phenos.get(code2, {})
|
||
n1, n2 = len(p1), len(p2)
|
||
if n1 < 2 or n2 < 2:
|
||
return _fail(code1, code2, "任一性状表型个体不足 2 个，无法估计方差组分")
|
||
# 单性状方差分量（REML）——与 run_ablup 同一求解器，保证 herit 口径一致
|
||
try:
|
||
r1 = blup.solve(pedigree, p1, tol=tol)
|
||
r2 = blup.solve(pedigree, p2, tol=tol)
|
||
except Exception as e: # noqa: BLE001
|
||
return _fail(code1, code2, f"单性状求解失败: {e!s}")
|
||
va1, ve1 = float(r1["sigma_a"]), float(r1["sigma_e"])
|
||
va2, ve2 = float(r2["sigma_a"]), float(r2["sigma_e"])
|
||
if va1 <= 0 or va2 <= 0:
|
||
return _fail(code1, code2, "单性状加性方差为 0（数据不支持遗传方差），无法估计遗传相关")
|
||
|
||
base, non_base, seen = _pedigree_partition(pedigree)
|
||
for pid in list(p1) + list(p2):
|
||
if pid not in seen:
|
||
seen.add(pid)
|
||
base.append(pid)
|
||
if not seen:
|
||
return _fail(code1, code2, "系谱与表型均为空")
|
||
order, idx = blup._order_pedigree(base, non_base)
|
||
n = len(order)
|
||
obs1 = [ind for ind in order if ind in p1]
|
||
obs2 = [ind for ind in order if ind in p2]
|
||
m1, m2 = len(obs1), len(obs2)
|
||
ix1 = [idx[i] for i in obs1]
|
||
ix2 = [idx[i] for i in obs2]
|
||
|
||
y1 = np.array([float(p1[i]) for i in obs1], dtype=float)
|
||
y2 = np.array([float(p2[i]) for i in obs2], dtype=float)
|
||
y = np.concatenate([y1, y2])
|
||
X = np.zeros((m1 + m2, 2))
|
||
X[:m1, 0] = 1.0
|
||
X[m1:, 1] = 1.0
|
||
|
||
A = blup._build_ainv(order, len(base), {i: (d, s) for (i, d, s) in non_base})[0]
|
||
A11 = A[np.ix_(ix1, ix1)]
|
||
A12 = A[np.ix_(ix1, ix2)]
|
||
A22 = A[np.ix_(ix2, ix2)]
|
||
cross = float(np.sqrt(va1 * va2))
|
||
|
||
best = {"ll": -np.inf, "rho": 0.0}
|
||
n_eval = 0
|
||
|
||
def _ll(rho: float) -> float:
|
||
nonlocal n_eval
|
||
n_eval += 1
|
||
sa12 = rho * cross
|
||
V = np.zeros((m1 + m2, m1 + m2))
|
||
V[:m1, :m1] = va1 * A11 + ve1 * np.eye(m1)
|
||
V[:m1, m1:] = sa12 * A12
|
||
V[m1:, :m1] = sa12 * A12.T
|
||
V[m1:, m1:] = va2 * A22 + ve2 * np.eye(m2)
|
||
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)
|
||
ll = -0.5 * (float(lv) + float(lx) + yPy)
|
||
except np.linalg.LinAlgError:
|
||
return -np.inf
|
||
if ll > best["ll"]:
|
||
best.update(ll=ll, rho=rho)
|
||
return ll
|
||
|
||
# 粗网格预扫 + 邻域黄金分割（防多峰取错局部峰）
|
||
grid = np.linspace(_RHO_LO, _RHO_HI, _GRID_N)
|
||
vals = [_ll(r) for r in grid]
|
||
best_i = int(np.argmax(vals))
|
||
lo = grid[max(best_i - 1, 0)]
|
||
hi = grid[min(best_i + 1, _GRID_N - 1)]
|
||
rho_opt, n_gold = blup._golden_max(_ll, lo, hi, tol=1e-4, max_iter=80)
|
||
if _ll(rho_opt) > best["ll"]:
|
||
best.update(ll=_ll(rho_opt), rho=rho_opt)
|
||
rho = best["rho"]
|
||
n_iter = _GRID_N + n_gold
|
||
|
||
warnings: list[str] = []
|
||
if abs(rho) >= _RHO_HI - 1e-6:
|
||
warnings.append(f"遗传相关 ρ 达边界 {rho:.3f}（接近完全{'正' if rho > 0 else '负'}相关，估计需谨慎）")
|
||
if n_iter >= blup.MAX_PROFILE_EVALS:
|
||
warnings.append("剖面 REML 未完全收敛（似然面极平/边界最优），结果采用已探明最优。")
|
||
|
||
common = set(obs1) & set(obs2)
|
||
return {
|
||
"code1": code1, "code2": code2,
|
||
"va1": float(va1), "va2": float(va2),
|
||
"va12": float(rho * cross),
|
||
"r_g": float(rho),
|
||
"ve1": float(ve1), "ve2": float(ve2),
|
||
"n_common": len(common), "n1": n1, "n2": n2,
|
||
"n_individuals": n, "n_iter": n_iter,
|
||
"converged": n_iter < blup.MAX_PROFILE_EVALS,
|
||
"warning": ("；".join(warnings) if warnings else None),
|
||
}
|
||
|
||
|
||
def _fail_multi(codes: list[str], msg: str) -> dict:
|
||
return {"codes": codes, "G0": None, "G0_inv": None, "r_g": None, "ve": None,
|
||
"h2": None, "n_obs": None, "n_individuals": 0, "n_iter": 0,
|
||
"converged": False, "warning": msg, "engine_version": ENGINE_VERSION}
|
||
|
||
|
||
def solve_multi(pedigree: list[dict], phenos: dict[str, dict],
|
||
codes: list[str], *, tol: float = blup.TOL,
|
||
max_iter: int = 200, tol_em: float = 1e-4) -> dict:
|
||
"""全多变量 MT-BLUP（EM-REML 估计完整 G0⊗A，替代逐对 bivariate）。
|
||
|
||
y = Xb + Zu + e：u 按性状主序堆叠 [u_1;...;u_m]（各 u_k 为全部 n 个体的育种值），
|
||
Var(u) = G0⊗A（G0 m×m 遗传协方差，半正定），Var(e) = diag(ve_1·I, ..., ve_m·I)。
|
||
|
||
EM-REML（Meyer 1985）：
|
||
- 初始 G0 = diag(各性状单性状 blup.solve REML va)、ve = 单性状 ve；
|
||
- 迭代：MME(C·s=r) 求 β/u + C⁻¹ → 遗传 G0_ij = (û_i'A⁻¹û_j + tr(A⁻¹C_ij))/n、
|
||
残差 ve_i = (e_i'e_i + Σ_{k∈obs_i} C_ii[k,k]) / n_i；
|
||
- G0 每次特征值截断投影半正定 + 对称化；|ΔREML LL| < tol_em 收敛；
|
||
- Aitken 逐元素外推加速（仅当外推点 REML LL 更高才接受，否则回退 EM 步）。
|
||
|
||
phenos: {code: {个体id: 均值}}（个体 id 空间与系谱一致，如 "t{tree_id}"）。
|
||
输出 G0 全元素 + G0_inv（Smith-Hazel 指数） + r_g 相关矩阵 + ve + h2。
|
||
"""
|
||
if not codes or len(codes) < 2:
|
||
return _fail_multi(list(codes or []), "全 MT-BLUP 需至少 2 个性状")
|
||
codes = list(codes)
|
||
m = len(codes)
|
||
base, non_base, seen = _pedigree_partition(pedigree)
|
||
for code in codes:
|
||
for pid in phenos.get(code, {}):
|
||
if pid not in seen:
|
||
seen.add(pid)
|
||
base.append(pid)
|
||
if not seen:
|
||
return _fail_multi(codes, "系谱与表型均为空")
|
||
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)
|
||
|
||
Xlist: list[np.ndarray] = []
|
||
Zlist: list[np.ndarray] = []
|
||
ylist: list[np.ndarray] = []
|
||
obslist: list[list[int]] = []
|
||
va0: list[float] = []
|
||
ve0: list[float] = []
|
||
for code in codes:
|
||
pk = phenos.get(code, {})
|
||
obs_k = [ind for ind in order if ind in pk]
|
||
if len(obs_k) < 2:
|
||
return _fail_multi(codes, f"性状 {code} 表型个体不足 2，无法估计方差组分")
|
||
obslist.append(obs_k)
|
||
ylist.append(np.array([float(pk[i]) for i in obs_k], dtype=float))
|
||
Zk = np.zeros((len(obs_k), n))
|
||
for kk, ind in enumerate(obs_k):
|
||
Zk[kk, idx[ind]] = 1.0
|
||
Zlist.append(Zk)
|
||
Xlist.append(np.ones((len(obs_k), 1)))
|
||
rk = blup.solve(pedigree, pk, tol=tol)
|
||
va0.append(float(rk["sigma_a"]))
|
||
ve0.append(float(rk["sigma_e"]))
|
||
if any(v <= 0 for v in va0):
|
||
return _fail_multi(codes, "有性状加性方差为 0，全多变量 REML 无法启动（请回退成对 bivariate）")
|
||
|
||
n_obs_list = [len(o) for o in obslist]
|
||
N = sum(n_obs_list)
|
||
Xbig = np.zeros((N, m))
|
||
Zbig = np.zeros((N, m * n))
|
||
yvec = np.zeros(N)
|
||
off = 0
|
||
for k in range(m):
|
||
mk = n_obs_list[k]
|
||
Xbig[off:off + mk, k] = 1.0
|
||
Zbig[off:off + mk, k * n:(k + 1) * n] = Zlist[k]
|
||
yvec[off:off + mk] = ylist[k]
|
||
off += mk
|
||
|
||
def _reml_ll(G0c: np.ndarray, vec: np.ndarray) -> float:
|
||
ZAZ = np.zeros((N, N))
|
||
for i in range(m):
|
||
for j in range(m):
|
||
offi = sum(n_obs_list[:i])
|
||
offj = sum(n_obs_list[:j])
|
||
ZAZ[offi:offi + n_obs_list[i], offj:offj + n_obs_list[j]] = (
|
||
G0c[i, j] * (Zlist[i] @ A @ Zlist[j].T))
|
||
V = ZAZ + np.diag(np.repeat(vec, n_obs_list))
|
||
try:
|
||
VinvX = np.linalg.solve(V, Xbig)
|
||
Vinvy = np.linalg.solve(V, yvec)
|
||
XtVinvX = Xbig.T @ VinvX
|
||
XtVinvY = Xbig.T @ Vinvy
|
||
yPy = float(yvec @ 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 _project_psd(mat: np.ndarray) -> tuple[np.ndarray, bool]:
|
||
mat = 0.5 * (mat + mat.T)
|
||
w, V = np.linalg.eigh(mat)
|
||
floor = 1e-6 * max(float(w[-1]), 1e-6)
|
||
clipped = bool(w[0] <= floor)
|
||
return (V * np.maximum(w, floor)) @ V.T, clipped
|
||
|
||
G0 = np.diag(va0)
|
||
ve = np.array(ve0, dtype=float)
|
||
warnings: list[str] = []
|
||
converged = False
|
||
ll = -np.inf
|
||
prev_ll: float | None = None
|
||
ll_gain = 1.0
|
||
it = 0
|
||
theta_prev2: np.ndarray | None = None
|
||
theta_prev1: np.ndarray | None = None
|
||
|
||
def _em_step(par_g0: np.ndarray, par_ve: np.ndarray) -> tuple[np.ndarray, np.ndarray, float]:
|
||
Rinv_diag = np.repeat(1.0 / par_ve, n_obs_list)
|
||
G0inv = np.linalg.inv(par_g0)
|
||
C = np.zeros((m + m * n, m + m * n))
|
||
C[:m, :m] = Xbig.T @ (Rinv_diag[:, None] * Xbig)
|
||
XtRinvZ = Xbig.T @ (Rinv_diag[:, None] * Zbig)
|
||
C[:m, m:] = XtRinvZ
|
||
C[m:, :m] = XtRinvZ.T
|
||
C[m:, m:] = Zbig.T @ (Rinv_diag[:, None] * Zbig) + np.kron(G0inv, Ainv)
|
||
rhs = np.concatenate([Xbig.T @ (Rinv_diag * yvec), Zbig.T @ (Rinv_diag * yvec)])
|
||
try:
|
||
sol = np.linalg.solve(C, rhs)
|
||
Cinv = np.linalg.inv(C)
|
||
except np.linalg.LinAlgError:
|
||
sol = np.linalg.pinv(C) @ rhs
|
||
Cinv = np.linalg.pinv(C)
|
||
u = sol[m:]
|
||
Cuu = Cinv[m:, m:]
|
||
|
||
G0n = np.zeros((m, m))
|
||
for i in range(m):
|
||
ui = u[i * n:(i + 1) * n]
|
||
for j in range(m):
|
||
uj = u[j * n:(j + 1) * n]
|
||
Cij = Cuu[i * n:(i + 1) * n, j * n:(j + 1) * n]
|
||
G0n[i, j] = (float(ui @ Ainv @ uj) + float(np.sum(Ainv * Cij))) / n
|
||
G0n, proj = _project_psd(G0n)
|
||
if proj:
|
||
warnings.append("遗传协方差 G0 非半正定，已特征值截断投影到最近半正定阵")
|
||
|
||
ve_n = np.zeros(m)
|
||
for k in range(m):
|
||
e_k = ylist[k] - float(sol[k]) - Zlist[k] @ u[k * n:(k + 1) * n]
|
||
obs_pos = np.array([idx[i] for i in obslist[k]])
|
||
# tr(PEV_e_k) = tr(W_k C⁻¹_kk W_k')，W_k=[X_k,Z_k]（X 为截距列、Z 为逐个体指示）
|
||
# = n_k·C⁻¹[βk,βk] + 2·Σ_obs C⁻¹[βk,u_k] + Σ_obs C⁻¹[u_k,u_k]
|
||
block_ku = Cinv[k, m + k * n: m + (k + 1) * n]
|
||
block_uu_k = Cinv[m + k * n: m + (k + 1) * n, m + k * n: m + (k + 1) * n]
|
||
trace_r = (n_obs_list[k] * float(Cinv[k, k])
|
||
+ 2.0 * float(np.sum(block_ku[obs_pos]))
|
||
+ float(np.sum(block_uu_k[obs_pos, obs_pos])))
|
||
ve_n[k] = (float(e_k @ e_k) + trace_r) / n_obs_list[k]
|
||
ve_n = np.clip(ve_n, blup._FLOOR, None)
|
||
return G0n, ve_n, _reml_ll(G0n, ve_n)
|
||
|
||
for it in range(1, max_iter + 1):
|
||
G0new, ve_new, new_ll = _em_step(G0, ve)
|
||
if new_ll == -np.inf:
|
||
continue
|
||
theta_new = np.concatenate([G0new.ravel(), ve_new])
|
||
# Aitken 加速：逐元素外推 θ∞≈θₜ+d₁/(1−ρ)，仅当外推点 REML LL 更高才接受，
|
||
# 缓解 EM 在高相关多性状下收敛缓慢（trait 方差先塌缩后缓慢恢复的困境）。
|
||
if (it >= 4 and theta_prev2 is not None and theta_prev1 is not None):
|
||
d1 = theta_new - theta_prev1
|
||
d2 = theta_prev1 - theta_prev2
|
||
denom = np.where(np.abs(d2) > 1e-12, d2, 1.0)
|
||
rho = np.where(np.abs(d2) > 1e-12, d1 / denom, 0.0)
|
||
rho = np.clip(rho, -0.999, 0.999)
|
||
theta_acc = theta_new + d1 / (1.0 - rho)
|
||
G0a, _ = _project_psd(0.5 * (theta_acc[:m * m].reshape(m, m)
|
||
+ theta_acc[:m * m].reshape(m, m).T))
|
||
vea = np.clip(theta_acc[m * m:], blup._FLOOR, None)
|
||
ll_acc = _reml_ll(G0a, vea)
|
||
if ll_acc > new_ll and ll_acc != -np.inf:
|
||
G0new, ve_new, new_ll = G0a, vea, ll_acc
|
||
theta_prev2, theta_prev1 = theta_prev1, theta_new
|
||
ll_gain = abs(new_ll - prev_ll) if prev_ll is not None else 1.0
|
||
G0, ve = G0new, ve_new
|
||
ll = new_ll
|
||
prev_ll = new_ll
|
||
if ll_gain < tol_em:
|
||
converged = True
|
||
break
|
||
if not converged:
|
||
warnings.append(f"EM-REML {max_iter} 次迭代未完全收敛（|ΔLL|={ll_gain:.2e}），结果采用已探明最优")
|
||
|
||
r_g = np.zeros((m, m))
|
||
for i in range(m):
|
||
for j in range(m):
|
||
if G0[i, i] > 0 and G0[j, j] > 0:
|
||
r_g[i, j] = G0[i, j] / math.sqrt(G0[i, i] * G0[j, j])
|
||
h2 = [G0[k, k] / (G0[k, k] + ve[k]) if (G0[k, k] + ve[k]) > 0 else None
|
||
for k in range(m)]
|
||
G0i = np.linalg.inv(G0)
|
||
return {
|
||
"codes": codes,
|
||
"G0": [[round(float(G0[i, j]), 6) for j in range(m)] for i in range(m)],
|
||
"G0_inv": [[round(float(G0i[i, j]), 6) for j in range(m)] for i in range(m)],
|
||
"r_g": [[round(float(r_g[i, j]), 4) for j in range(m)] for i in range(m)],
|
||
"ve": [round(float(v), 6) for v in ve],
|
||
"h2": [round(float(h), 4) if h is not None else None for h in h2],
|
||
"n_obs": n_obs_list,
|
||
"n_individuals": n,
|
||
"n_iter": it,
|
||
"converged": converged,
|
||
"warning": "；".join(warnings) if warnings else None,
|
||
"engine_version": ENGINE_VERSION,
|
||
}
|