34047007@qq.com
853d01d796
- _order_pedigree: base 与同世代并列个体均按 id 确定性排序，输入/DB 行序不再参与计算 - 平似然/退化数据上同内容输入必逐位同输出（此前 golden-section/二分会被末位浮点噪声带偏到任意 h²，导致回归间歇失败）
1705 lines
72 KiB
Python
1705 lines
72 KiB
Python
"""纯 Python 单性状动物模型 BLUP（Henderson MME + 剖面 REML）。
|
||
|
||
仅依赖 numpy，无 R / 外部运行时。
|
||
|
||
系谱锚点（本系统约定）：
|
||
- 个体 = bre_tree 单株（子代），其 dam/sire = bre_germplasm 种质（base/founder，无父母）。
|
||
- 同组合单株互为全同胞（A=0.5），跨组合共享单一亲本为半同胞（A=0.25），
|
||
单株↔亲本 A=0.5 —— 全部由 A 矩阵递归自动得出，无需人工指定。
|
||
- 无父母的个体一律视为 founder（base population，彼此不相关）。
|
||
|
||
算法：
|
||
1. A⁻¹ 用 Henderson 稀疏规则构造（无需求 A 的逆）；
|
||
2. MME 求解：y = Xb + Zu + e，X=截距+固定效应虚拟列+协变量（默认仅截距），λ = σ²e/σ²a；
|
||
个体数 ≤ N_SPARSE（默认 1000）→ 稠密精确求解（solver="dense"：索引化 Z、稠密 A、
|
||
剖面 REML 黄金分割）；个体数 > N_SPARSE 且系谱存在 → 稀疏路径（solver="sparse-em"）：
|
||
Z 索引化设计矩阵（bincount）、稀疏 A⁻¹ 三元组 + 共轭梯度迭代，全程不落稠密 Z/A/V
|
||
（35k 树在常规内存内跑通）。方差分量按 n 分档：
|
||
个体数 ≤ N_EXACT（3000）→ tr(A⁻¹C²²) 逐列 CG 精确（无随机噪声），REML 解
|
||
h(λ)=g(λ)−λ=0 单调单根 → log10λ 二分 ~32 步到 1e-5（规避 EM 慢收缩）；
|
||
个体数 > N_EXACT 且观测数 m>N_SUBSAMPLE（可真实子采样）→ 随机子样本（观测
|
||
≤N_SUBSAMPLE，含祖先闭包）上精确迹二分求 REML 解——子样本 REML 为全群方差组分
|
||
一致估计（抽样 SE h² 约 ±0.04），再全群 MME 于子样本 λ* 下解 EBV（大样本 EM 慢收缩
|
||
ρ≈0.999 冻结不可用；EBV 对 λ 稳健，warning 标注）；
|
||
3. reliability：稠密路径精确（diag(C22) 取 MME 逆对角）；稀疏路径 n≤EXACT_DIAG_N 用
|
||
稠密 MME 逆精确、n 大用 Hutchinson 随机探测（reliability_approx=True，均值/PA 可靠、个体级近似）；
|
||
4. h² = σ²a/(σ²a+σ²e)；reliability = 1 - PEV/σ²a，PEV = C22[ii]·σ²e。
|
||
|
||
无系谱（所有个体均无父母）时：不背书遗传力——
|
||
h2 置 None、reliability 全部置 0、EBV 取中心化表型，并返回 warning。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import math
|
||
import random
|
||
|
||
import numpy as np
|
||
|
||
TOL = 1e-4
|
||
MAX_PROFILE_EVALS = 100
|
||
MAX_GXE_ITERS = 25 # G×E 分支坐标上升迭代上限
|
||
_FLOOR = 1e-8 # 方差下界，避免除零/负方差
|
||
H2_MIN, H2_MAX = 1e-3, 1.0 - 1e-3
|
||
N_SPARSE = 1000 # 个体数超过该阈值 → 稀疏 A⁻¹ + 共轭梯度（回避 O(n³) 稠密 MME 逆）
|
||
CG_MAX_ITER = 800 # 共轭梯度迭代上限
|
||
CG_TOL = 1e-9 # 共轭梯度相对残差收敛阈值
|
||
# Hutchinson 随机探测次数（估计 C22 对角线 → PEV → reliability）。稀疏路径 EBV/h²/σ² 为
|
||
# CG 精确解；reliability 是 MC 对角估计——同家族强相关的 C22 令其收敛慢（~1/√k），
|
||
# 误差随 k 减小但非精确，作为文档标注的近似（排序可靠、量级可信）。
|
||
HUTCH_K = 100
|
||
# 稀疏路径 reliability 分档：个体数 ≤ EXACT_DIAG_N → diag(C22) 用稠密 MME 逆精确求
|
||
# （golden 等价档，无 MC 噪声——n=20 全同胞小样本 h²→1 边界实测 Hutchinson 均值误差 ~0.08）；
|
||
# > EXACT_DIAG_N → Hutchinson 随机探测（排序可靠、量级可信，误差随个体数稀释）。
|
||
EXACT_DIAG_N = 1000
|
||
# 稀疏路径分档：精确迹二分求 λ*（逐列 CG 精确 tr(A⁻¹C²²)，无随机噪声，
|
||
# REML 解 = h(λ)=g(λ)−λ=0 单调单根，log10λ 二分 ~32 步到 1e-5，规避 EM 慢收缩 ρ≈0.987）。
|
||
# 精确迹 O(n²)（每步似然 n 次逐列 CG），仅个体数 ≤ N_EXACT 才可负担；>N_EXACT 且可真实
|
||
# 子采样观测（m>N_SUBSAMPLE）→ 随机子样本精确迹二分（观测 ≤N_SUBSAMPLE，含祖先闭包）
|
||
# 求 λ*，再全群 MME 求解。
|
||
# 大样本 EM 收缩率 ρ→0.999（100 步 λ 几乎不动，冻结在初始猜测），Hutchinson 噪声又淹没
|
||
# h(λ) 信号（~2e-4 vs 噪声 ~2e-3），全群精确迹 O(n²) 不可行；子样本 iid 抽样 REML 为全群
|
||
# 方差组分一致估计（n_sub=3000 抽样 SE h² ≈ ±0.04），EBV 对 λ 稳健（偏 25% Kendall τ>0.97）。
|
||
# N_EXACT 对齐 N_SUBSAMPLE：m>N_SUBSAMPLE 时 n≈m+系谱必超 N_EXACT → 走子样本封顶 ~70min，
|
||
# 消除 3000<n≤6000 精确迹数小时级尾（性能倒挂：局部选 3000-6000 树比全量还慢）。
|
||
N_EXACT = 3000
|
||
N_SUBSAMPLE = 3000 # 生产档子样本观测数上限（n>N_EXACT 且 m>N_SUBSAMPLE 才触发）
|
||
BISECT_LO, BISECT_HI = -4.0, 4.0 # log10(λ=Ve/Va) 二分区间（REML 解通常在其内）
|
||
BISECT_TOL = 1e-8 # 二分 log10(λ) 收敛宽度（→ λ 相对精度 ~1e-8）
|
||
MAX_BISECT_EVAL = 60 # 二分求值上限
|
||
|
||
ENGINE_VERSION = "1.7.0" # 单性状动物模型 BLUP 求解器版本（落 bre_prediction.engine_version；1.1.0: 供 mtblup 单性状方差复用；1.2.0: G×R 砧木×接穗互作；1.3.0: AR1×AR1 空间协方差 solve_spatial + 稀疏 A⁻¹/共轭梯度/Hutchinson 可靠性；1.4.0: AR1×AR1 各向异性双参数 ρ_row/ρ_col；1.5.0: 区组随机效应 _solve_block（增广/α-格子 block_no 第二随机效应）；1.6.0: 稀疏路径 EM-REML（索引化设计矩阵 + 稀疏 A⁻¹ 无稠密 A + EM-REML，35k 树不再 OOM；n≤N_EXACT 精确迹二分 λ*，n>N_EXACT 随机子样本精确迹二分 → 全群 MME 解 EBV）；1.7.0: 性能倒挂修复——N_EXACT 6000→3000 对齐 N_SUBSAMPLE，生产档仅 m>N_SUBSAMPLE 才触发；系谱闭包改 BFS 可达祖先（服务层），选中 200~3000 树不再被全量 germplasm 闭包顶进数小时精确迹尾）
|
||
|
||
|
||
def _pearson(x: list[float], y: list[float]) -> float:
|
||
n = len(x)
|
||
if n < 2:
|
||
return 0.0
|
||
mx, my = sum(x) / n, sum(y) / n
|
||
cov = sum((a - mx) * (b - my) for a, b in zip(x, y))
|
||
dx = sum((a - mx) ** 2 for a in x) ** 0.5
|
||
dy = sum((b - my) ** 2 for b in y) ** 0.5
|
||
if dx == 0 or dy == 0:
|
||
return 0.0
|
||
return cov / (dx * dy)
|
||
|
||
|
||
def _order_pedigree(base: list[int], non_base: list[tuple[int, int | None, int | None]]) -> tuple[list[int], dict[int, int]]:
|
||
"""返回拓扑序（祖先恒在子代之前，base 恒在最前）与 {个体id: 序号}。
|
||
|
||
输入顺序不参与计算：base 与同世代并列个体均按 id 确定性排序，调用方/DB 行序不再
|
||
影响 MME 变量布局——同内容输入必得逐位相同输出（平似然/退化数据上尤其关键，否则
|
||
golden-section/二分会被末位浮点噪声带偏到任意 h²）。
|
||
"""
|
||
base = sorted(base, key=str)
|
||
non_base = sorted(non_base, key=lambda it: str(it[0]))
|
||
order: list[int] = list(base)
|
||
placed: set[int] = set(base)
|
||
remaining = list(non_base)
|
||
while remaining:
|
||
progressed = False
|
||
for it in list(remaining):
|
||
i, d, s = it
|
||
if (d is None or d in placed) and (s is None or s in placed):
|
||
order.append(i)
|
||
placed.add(i)
|
||
remaining.remove(it)
|
||
progressed = True
|
||
if not progressed:
|
||
raise ValueError(
|
||
"系谱存在环或父母引用缺失，无法拓扑排序"
|
||
"（父母 id 必须是已知个体或 None）"
|
||
)
|
||
return order, {ind: idx for idx, ind in enumerate(order)}
|
||
|
||
|
||
def _build_ainv(order: list[int], n_base: int,
|
||
parent_of: dict[int, tuple[int | None, int | None]]) -> tuple[np.ndarray, np.ndarray]:
|
||
"""返回 (A, A⁻¹)。A 稠密递归 O(n²)；A⁻¹ 用 Henderson 稀疏规则。"""
|
||
n = len(order)
|
||
idx = {ind: pos for pos, ind in enumerate(order)}
|
||
A = np.eye(n)
|
||
for pos, ind in enumerate(order):
|
||
d, s = parent_of.get(ind, (None, None))
|
||
if d is None and s is None:
|
||
continue
|
||
di = idx.get(d) if d is not None else None
|
||
si = idx.get(s) if s is not None else None
|
||
row = np.zeros(n)
|
||
if di is not None:
|
||
row += A[di]
|
||
if si is not None:
|
||
row += A[si]
|
||
A[pos, :] = 0.5 * row
|
||
A[:pos, pos] = A[pos, :pos]
|
||
apq = A[di, si] if (di is not None and si is not None) else 0.0
|
||
A[pos, pos] = 1.0 + 0.5 * apq
|
||
|
||
Ainv = np.zeros((n, n))
|
||
for i in range(n_base):
|
||
Ainv[i, i] += 1.0
|
||
for ind, (d, s) in parent_of.items():
|
||
i = idx[ind]
|
||
di = idx.get(d) if d is not None else None
|
||
si = idx.get(s) if s is not None else None
|
||
fp = A[di, di] - 1.0 if di is not None else 0.0
|
||
fq = A[si, si] - 1.0 if si is not None else 0.0
|
||
t = 0.5 - 0.25 * (fp + fq)
|
||
# 极端近交（双亲 F→1）时 t→0，除零保护；1e-6 下界对正常系谱无感
|
||
inv_t = 1.0 / max(t, 1e-6)
|
||
Ainv[i, i] += inv_t
|
||
for p in (di, si):
|
||
if p is not None:
|
||
Ainv[p, p] += 0.25 * inv_t
|
||
Ainv[i, p] -= 0.5 * inv_t
|
||
Ainv[p, i] -= 0.5 * inv_t
|
||
if di is not None and si is not None:
|
||
Ainv[di, si] += 0.25 * inv_t
|
||
Ainv[si, di] += 0.25 * inv_t
|
||
return A, Ainv
|
||
|
||
|
||
def _build_ainv_sparse(order: list[int], n_base: int,
|
||
parent_of: dict[int, tuple[int | None, int | None]]
|
||
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
|
||
"""返回 (diag_a, ii, jj, vv)：A 对角（近交 F=A_ii−1 来源）；A⁻¹ 以 COO 稀疏三元组存储。
|
||
|
||
不构建稠密 A：加性关系对角用「祖先稀疏行」递推（每体一行 {祖先: A_ij}，浅系谱祖先数
|
||
小，O(n·s²)），Henderson 规则只需 A 对角（fp/fq）与双亲交叉 A[di,si]，均可在稀疏行上
|
||
读出；A⁻¹ 非零元（对角 + 2 亲本 × 对角/双向共祖先）逐项与 _build_ainv 一致。
|
||
matvec 时用 np.bincount 按 ii 累加 ainv_vv·u[jj] 即得 A⁻¹·u（重复元自动合并）。
|
||
"""
|
||
n = len(order)
|
||
idx = {ind: pos for pos, ind in enumerate(order)}
|
||
anc: list[dict[int, float]] = [None] * n
|
||
diag_a = np.zeros(n)
|
||
for pos, ind in enumerate(order):
|
||
d, s = parent_of.get(ind, (None, None))
|
||
if d is None and s is None:
|
||
anc[pos] = {pos: 1.0}
|
||
diag_a[pos] = 1.0
|
||
continue
|
||
di = idx.get(d) if d is not None else None
|
||
si = idx.get(s) if s is not None else None
|
||
row: dict[int, float] = {}
|
||
for src in (di, si):
|
||
if src is not None:
|
||
for j, v in anc[src].items():
|
||
row[j] = row.get(j, 0.0) + 0.5 * v
|
||
apq = 0.0
|
||
if di is not None and si is not None:
|
||
small, big = (anc[di], anc[si]) if len(anc[di]) < len(anc[si]) else (anc[si], anc[di])
|
||
apq = sum(v * big[j] for j, v in small.items() if j in big)
|
||
row[pos] = 1.0 + 0.5 * apq
|
||
anc[pos] = row
|
||
diag_a[pos] = row[pos]
|
||
|
||
ii: list[int] = []
|
||
jj: list[int] = []
|
||
vv: list[float] = []
|
||
|
||
def _put(i: int, j: int, v: float) -> None:
|
||
if v == 0.0:
|
||
return
|
||
ii.append(i)
|
||
jj.append(j)
|
||
vv.append(v)
|
||
|
||
for i in range(n_base):
|
||
_put(i, i, 1.0)
|
||
for ind, (d, s) in parent_of.items():
|
||
i = idx[ind]
|
||
di = idx.get(d) if d is not None else None
|
||
si = idx.get(s) if s is not None else None
|
||
fp = diag_a[di] - 1.0 if di is not None else 0.0
|
||
fq = diag_a[si] - 1.0 if si is not None else 0.0
|
||
t = 0.5 - 0.25 * (fp + fq)
|
||
# 极端近交（双亲 F→1）时 t→0，除零保护；1e-6 下界对正常系谱无感
|
||
inv_t = 1.0 / max(t, 1e-6)
|
||
_put(i, i, inv_t)
|
||
for pp in (di, si):
|
||
if pp is not None:
|
||
_put(pp, pp, 0.25 * inv_t)
|
||
_put(i, pp, -0.5 * inv_t)
|
||
_put(pp, i, -0.5 * inv_t)
|
||
if di is not None and si is not None:
|
||
_put(di, si, 0.25 * inv_t)
|
||
_put(si, di, 0.25 * inv_t)
|
||
return (diag_a, np.array(ii, dtype=np.int64), np.array(jj, dtype=np.int64),
|
||
np.array(vv, dtype=float))
|
||
|
||
|
||
def _cg_solve(matvec, rhs: np.ndarray, max_iter: int = CG_MAX_ITER,
|
||
tol: float = CG_TOL) -> tuple[np.ndarray, int, bool]:
|
||
"""共轭梯度求解 C·x = rhs（C 对称正定，仅需 matvec 回调）。
|
||
|
||
收敛判据：‖r‖ ≤ tol·‖rhs‖（相对残差）；达到 max_iter 未达标 → 返回 (解, max_iter, False)。
|
||
调用方对未收敛结果回退稠密路径（小 n）或加 warning 并采用已探明最优。
|
||
"""
|
||
n = rhs.shape[0]
|
||
x = np.zeros(n)
|
||
r = rhs.astype(float).copy()
|
||
p = r.copy()
|
||
rsold = float(r @ r)
|
||
rhs_norm = float(np.linalg.norm(rhs))
|
||
if rhs_norm == 0.0 or rsold == 0.0:
|
||
return x, 0, True
|
||
target2 = (tol * rhs_norm) ** 2
|
||
for it in range(1, max_iter + 1):
|
||
Ap = matvec(p)
|
||
pAp = float(p @ Ap)
|
||
if pAp <= 1e-300:
|
||
return x, it, False
|
||
alpha = rsold / pAp
|
||
x += alpha * p
|
||
r -= alpha * Ap
|
||
rsnew = float(r @ r)
|
||
if rsnew <= target2:
|
||
return x, it, True
|
||
p = r + (rsnew / rsold) * p
|
||
rsold = rsnew
|
||
return x, max_iter, False
|
||
|
||
|
||
def _golden_max(f, lo: float, hi: float, tol: float = TOL, max_iter: int = MAX_PROFILE_EVALS) -> tuple[float, int]:
|
||
"""黄金分割一维最大化（f 单调时收敛到边界）。返回 (极值点, 函数求值次数)。"""
|
||
invphi = (np.sqrt(5.0) - 1.0) / 2.0
|
||
a, b = lo, hi
|
||
x1 = b - invphi * (b - a)
|
||
x2 = a + invphi * (b - a)
|
||
f1, f2 = f(x1), f(x2)
|
||
n = 2
|
||
while abs(b - a) > tol and n < max_iter:
|
||
if f1 > f2:
|
||
b, x2, f2 = x2, x1, f1
|
||
x1 = b - invphi * (b - a)
|
||
f1 = f(x1)
|
||
else:
|
||
a, x1, f1 = x1, x2, f2
|
||
x2 = a + invphi * (b - a)
|
||
f2 = f(x2)
|
||
n += 1
|
||
return (a + b) / 2.0, n
|
||
|
||
|
||
def relationship_matrix(pedigree: list[dict]) -> dict:
|
||
"""计算个体间加性亲缘矩阵 A（个体数少时稠密递归），供亲缘/近交分析。
|
||
|
||
亲缘系数 r_ij = A_ij（加性相关，双列之间 0.5=亲子/全同胞、0.25=半同胞）；
|
||
近交系数 F_i = A_ii - 1。
|
||
|
||
参数:
|
||
pedigree: 同 solve() 的系谱结构（id 唯一，dam/sire 为 None 表示 founder）。
|
||
|
||
返回:
|
||
individuals: [个体id]（拓扑序：founder 在前）；
|
||
matrix: A 矩阵（len×len，6 位小数）；
|
||
inbreeding: {individual_id: F}（仅 F>0 的个体）。
|
||
"""
|
||
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))
|
||
if not seen:
|
||
raise ValueError("系谱为空")
|
||
order, _ = _order_pedigree(base, non_base)
|
||
parent_of = {i: (d, s) for (i, d, s) in non_base}
|
||
A, _ = _build_ainv(order, len(base), parent_of)
|
||
names = [str(o) for o in order]
|
||
inbreeding = {
|
||
str(o): float(round(A[i, i] - 1.0, 6))
|
||
for i, o in enumerate(order)
|
||
if A[i, i] - 1.0 > 1e-8
|
||
}
|
||
return {
|
||
"individuals": names,
|
||
"matrix": A.round(6).tolist(),
|
||
"inbreeding": inbreeding,
|
||
}
|
||
|
||
|
||
def _build_design(obs_ids: list[int], fixed: dict | None = None,
|
||
covariate: dict | None = None) -> np.ndarray:
|
||
"""构建设计阵 X：截距 + 固定效应水平虚拟列 + 协变量列。
|
||
|
||
fixed: {因子名: {individual_id: 水平}}——每因子对基准水平(排序首个)加虚拟列，
|
||
缺失水平的个体归基准（全 0）；仅 1 个水平的因子无信息，跳过。
|
||
covariate: {individual_id: float}——单一协变量列，缺失按观测均值填充；
|
||
全部同值（无变异）时无信息，跳过。
|
||
"""
|
||
m = len(obs_ids)
|
||
cols: list[np.ndarray] = [np.ones(m)]
|
||
for _fname, mapping in (fixed or {}).items():
|
||
levels = sorted({str(mapping[i]) for i in obs_ids if mapping.get(i) is not None})
|
||
if len(levels) <= 1:
|
||
continue
|
||
ref = levels[0]
|
||
for lev in levels[1:]:
|
||
col = np.array(
|
||
[1.0 if str(mapping.get(i)) == lev else 0.0 for i in obs_ids],
|
||
dtype=float,
|
||
)
|
||
cols.append(col)
|
||
if covariate:
|
||
raw = [covariate.get(i) for i in obs_ids]
|
||
obs_vals = [float(v) for v in raw if v is not None]
|
||
if obs_vals:
|
||
mean = float(np.mean(obs_vals))
|
||
filled = np.array([float(v) if v is not None else mean for v in raw])
|
||
if np.std(filled) > 1e-12:
|
||
cols.append(filled)
|
||
return np.column_stack(cols)
|
||
|
||
|
||
def solve(pedigree: list[dict], phenotypes: dict, *, tol: float = TOL,
|
||
fixed: dict | None = None, covariate: dict | None = None,
|
||
gxe: dict | None = None, record_map: dict | None = None,
|
||
block: dict | None = None,
|
||
factor_label: str = "G×E", factor_name: str = "环境") -> dict:
|
||
"""运行单性状动物模型 BLUP（可含固定效应与协变量，如站点、crop_load）。
|
||
|
||
参数:
|
||
pedigree: [{"individual": id, "dam": id|None, "sire": id|None}]；
|
||
个体 id 唯一；dam/sire 为 None 表示 founder（如种质）。
|
||
phenotypes: {记录id: float}；默认每记录=一个加性个体（树级均值）；
|
||
提供 record_map 时记录可多条映射到同一个体（如同一 clone 多年观测）。
|
||
fixed: {因子名: {记录id: 水平}}——固定效应（如 {"site": {1:"A", 2:"B"}}）；
|
||
每因子对基准水平加虚拟列，缺失个体归基准。
|
||
covariate: {记录id: float}——一个协变量列（如 crop_load，缺失按观测均值）。
|
||
gxe: {记录id: (基因型分组, 环境水平)}——提供即启用 G×E 分支：
|
||
增加第二随机效应 v（基因型×环境，协方差 I·σ²gxe），
|
||
剖面 REML 坐标上升估三方差分量（比率边界与 h² 网格一致，λ 恒正）。
|
||
block: {记录id: 区组水平}——提供即启用区组随机效应分支（不完全区组/增广/α-格子）：
|
||
增加第二随机效应 v（区组，协方差 I·σ²block），与 gxe 互斥（共用 Z₂ 槽）。
|
||
record_map: {记录id: 加性个体id}——记录→个体映射（gxe/block 分支用），默认恒等。
|
||
factor_label / factor_name: 交互随机效应文案（默认 G×E/环境；G×R 传 G×R/砧木），仅影响警告字符串。
|
||
|
||
返回:
|
||
ebv: {id: float} 全部个体（含 base 种质）的育种值；
|
||
reliability: {id: float} 个体可靠性（无系谱时为 0）；
|
||
h2: float | None 窄义遗传力（无系谱时为 None）；
|
||
sigma_a / sigma_e: float；sigma_gxe: float（G×E 分支）；sigma_block: float（区组分支）；
|
||
gxe_ratio / n_gxe / n_cross_env / n_genotypes（G×E 分支）；
|
||
block_ratio / n_blocks / block_effects（区组分支）；
|
||
n_obs / n_individuals / n_base / n_with_parents / n_fixed: int；
|
||
converged: bool；n_iter: int（稠密=剖面求值次数，稀疏=迹二分/子样本求值次数）；warning: str | None；
|
||
solver: str（"dense" 稠密精确 / "sparse-em" 稀疏 EM-REML，n>N_SPARSE 时触发）；
|
||
稀疏路径另含 cg_iter（CG 迭代数）与 reliability_approx=True（可靠性为 Hutchinson 近似）；
|
||
生产档（n>N_EXACT）方差组分为随机子样本 REML 一致估计（抽样 SE h² ≈ ±0.04），EBV 由
|
||
全群 MME 于子样本 λ 下求解——对 λ 稳健，排序可靠，h²/σ² 视为近似区间。
|
||
"""
|
||
global N_EXACT # 生产档内联递归临时抬档（try/finally 还原）；须在函数内首次读取 N_EXACT 前声明
|
||
if gxe is not None:
|
||
return _solve_gxe(pedigree, phenotypes, fixed=fixed, covariate=covariate,
|
||
gxe=gxe, record_map=record_map or {}, tol=tol,
|
||
factor_label=factor_label, factor_name=factor_name)
|
||
if block is not None:
|
||
return _solve_block(pedigree, phenotypes, fixed=fixed, covariate=covariate,
|
||
block_of=block, record_map=record_map or {}, tol=tol,
|
||
factor_label=factor_label, factor_name=factor_name)
|
||
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 = rec.get("dam")
|
||
s = rec.get("sire")
|
||
if d is None and s is None:
|
||
base.append(i)
|
||
else:
|
||
non_base.append((i, d, s))
|
||
# 表型中出现的个体若不在系谱中，自动并入 base（鲁棒）。
|
||
for pid in phenotypes:
|
||
if pid not in seen:
|
||
seen.add(pid)
|
||
base.append(pid)
|
||
if not seen:
|
||
raise ValueError("系谱与表型均为空")
|
||
|
||
order, idx = _order_pedigree(base, non_base)
|
||
n = len(order)
|
||
n_base = len(base)
|
||
parent_of: dict[int, tuple[int | None, int | None]] = {i: (d, s) for (i, d, s) in non_base}
|
||
n_with_parents = len(parent_of)
|
||
|
||
obs_ids = [ind for ind in order if ind in phenotypes]
|
||
m = len(obs_ids)
|
||
if m < 2:
|
||
raise ValueError("有效表型个体不足 2 个，无法估计方差组分")
|
||
|
||
X = _build_design(obs_ids, fixed, covariate)
|
||
p = X.shape[1]
|
||
n_fixed = p - 1
|
||
y = np.array([float(phenotypes[i]) for i in obs_ids], dtype=float)
|
||
|
||
if n_with_parents == 0:
|
||
# 无系谱：不背书遗传力。EBV 取表型残差（固定效应已校正）。
|
||
beta = np.linalg.lstsq(X, y, rcond=None)[0]
|
||
resid = y - X @ beta
|
||
ebv: dict[int, float] = {}
|
||
for k, ind in enumerate(obs_ids):
|
||
ebv[ind] = float(resid[k])
|
||
for ind in order:
|
||
ebv.setdefault(ind, 0.0)
|
||
return {
|
||
"ebv": ebv,
|
||
"reliability": {ind: 0.0 for ind in order},
|
||
"h2": None,
|
||
"sigma_a": 0.0,
|
||
"sigma_e": 0.0,
|
||
"n_obs": m,
|
||
"n_individuals": n,
|
||
"n_base": n_base,
|
||
"n_with_parents": 0,
|
||
"n_fixed": n_fixed,
|
||
"converged": True,
|
||
"n_iter": 0,
|
||
"warning": "无系谱信息（所有个体均无父母）：EBV 仅为表型残差（固定效应已校正），未估计遗传力与可靠性。",
|
||
}
|
||
|
||
XtX = X.T @ X
|
||
Xty = X.T @ y
|
||
# 索引化设计矩阵（Z 每行恰一个 1）：Zty / ZtZ 对角 / XtZ 全用 bincount 累加，
|
||
# 稀疏路径不再落稠密 Z（35k 树稠密 Z ≈ 54 GiB）。稠密分支在小 n 下重建等价 Z。
|
||
obs_idx = np.array([idx[ind] for ind in obs_ids], dtype=np.int64)
|
||
Zty = np.bincount(obs_idx, weights=y, minlength=n)
|
||
ZtZ_diag = np.bincount(obs_idx, minlength=n).astype(float)
|
||
XtZ = np.zeros((p, n))
|
||
for col in range(p):
|
||
XtZ[col] = np.bincount(obs_idx, weights=X[:, col], minlength=n)
|
||
|
||
# A7 阈值分派：个体数 > N_SPARSE 且系谱存在 → 稀疏 A⁻¹ + 索引化设计 + EM-REML
|
||
# （回避 O(n³) 稠密 MME 逆与 O(m²) 稠密 V）；可靠性 n≤EXACT_DIAG_N 用稠密 MME 逆精确，
|
||
# n 大用 Hutchinson 随机探测估计 C22 对角线。
|
||
use_sparse = n_with_parents > 0 and n > N_SPARSE
|
||
if use_sparse:
|
||
diag_a, ainv_ii, ainv_jj, ainv_vv = _build_ainv_sparse(order, n_base, parent_of)
|
||
else:
|
||
A, Ainv = _build_ainv(order, n_base, parent_of)
|
||
|
||
if use_sparse:
|
||
# ---- 稀疏路径：EM-REML（Meyer 1985 单性状动物模型）----
|
||
# Va = (û'A⁻¹û + σ²e·tr(A⁻¹C²²))/n；Ve = yPy/(m−p)；λ=Ve/Va。全程不落稠密 Z/A/V。
|
||
# n ≤ N_EXACT：tr(A⁻¹C²²) 逐列 CG 精确（无随机噪声）；REML 解 = h(λ)=g(λ)−λ=0
|
||
# 单调单根 → log10λ 二分，~32 步到 1e-5（规避 EM 慢收缩 ρ≈0.987）；
|
||
# n > N_EXACT：随机子样本（≤N_SUBSAMPLE 观测，祖先闭包）内联递归走精确迹二分
|
||
# 求 λ*，再全群 MME 于 λ* 解 EBV——子样本抽样误差即真实不确定性（warning 标注）。
|
||
ztz = ZtZ_diag
|
||
|
||
def _matvec(lam: float):
|
||
def mv(v: np.ndarray) -> np.ndarray:
|
||
b, u = v[:p], v[p:]
|
||
out1 = XtX @ b + XtZ @ u
|
||
out2 = XtZ.T @ b + ztz * u
|
||
# A⁻¹·u：bincount 按 ii 累加（重复元自动合并），比 np.add.at 快一个量级
|
||
out2 += np.bincount(ainv_ii, weights=lam * ainv_vv * u[ainv_jj], minlength=n)
|
||
return np.concatenate([out1, out2])
|
||
return mv
|
||
|
||
mme_rhs = np.concatenate([Xty, Zty])
|
||
rng = np.random.RandomState(20260804)
|
||
|
||
em_it = 0
|
||
em_ok = True
|
||
cg_ok = True
|
||
cg_iter = 0
|
||
# 生产档仅当能真实子采样观测（m>N_SUBSAMPLE）才走：m 不大时子样本=全样本，退化为
|
||
# 精确迹 + 冗余全群 MME（更慢且子样本误差 SE 无意义）——直接全 n 精确迹更省更准。
|
||
use_production = n > N_EXACT and m > N_SUBSAMPLE
|
||
if not use_production:
|
||
# ---- 精确迹二分（golden 等价档）----
|
||
# A⁻¹ 第 j 列 = {i: ainv[i,j]}；C²²e_j = MME 解 RHS=[0;e_j] 的 u 块；
|
||
# tr(A⁻¹C²²) = Σ_j (A⁻¹e_j)ᵀ(C²²e_j)。逐列 CG 精确，无 Hutchinson 噪声。
|
||
col_of: list[list[tuple[int, float]]] = [[] for _ in range(n)]
|
||
for a, b, v in zip(ainv_ii.tolist(), ainv_jj.tolist(), ainv_vv.tolist()):
|
||
col_of[b].append((a, v))
|
||
|
||
def _em_step(lam: float) -> tuple[float, float]:
|
||
mv = _matvec(lam)
|
||
sol, _ci, _ok = _cg_solve(mv, mme_rhs)
|
||
b, u = sol[:p], sol[p:]
|
||
yPy = float(y @ y - b @ Xty - u @ Zty)
|
||
Ve = max(yPy / max(m - p, 1), _FLOOR)
|
||
tr = 0.0
|
||
for j in range(n):
|
||
rhs2 = np.zeros(p + n)
|
||
rhs2[p + j] = 1.0
|
||
sol2, _i2, _o2 = _cg_solve(mv, rhs2)
|
||
uj = sol2[p:]
|
||
s = 0.0
|
||
for i, v in col_of[j]:
|
||
s += v * uj[i]
|
||
tr += s
|
||
uAu = float(np.bincount(ainv_ii, weights=ainv_vv * u[ainv_jj], minlength=n) @ u)
|
||
Va = max((uAu + Ve * tr) / max(n, 1), _FLOOR)
|
||
return Va, Ve
|
||
|
||
def _h(lv: float) -> float:
|
||
lam = 10.0 ** lv
|
||
Va, Ve = _em_step(lam)
|
||
return Ve / Va - lam
|
||
|
||
lo, hi = BISECT_LO, BISECT_HI
|
||
h_lo, h_hi = _h(lo), _h(hi)
|
||
while h_lo <= 0.0 and lo > BISECT_LO - 4.0:
|
||
lo -= 2.0
|
||
h_lo = _h(lo)
|
||
while h_hi >= 0.0 and hi < BISECT_HI + 4.0:
|
||
hi += 2.0
|
||
h_hi = _h(hi)
|
||
if h_lo > 0.0 > h_hi:
|
||
em_it = 2
|
||
while (hi - lo) > BISECT_TOL and em_it < MAX_BISECT_EVAL:
|
||
mid = 0.5 * (lo + hi)
|
||
h_mid = _h(mid)
|
||
em_it += 1
|
||
if h_mid > 0.0:
|
||
lo = mid
|
||
else:
|
||
hi = mid
|
||
Va, Ve = _em_step(10.0 ** (0.5 * (lo + hi)))
|
||
else:
|
||
# 无变号（REML 解在区间外/边界最优）：取 |h| 最小的网格点
|
||
lv_grid = np.linspace(max(lo - 4.0, -8.0), min(hi + 4.0, 8.0), 41)
|
||
h_abs = [abs(_h(lv)) for lv in lv_grid]
|
||
Va, Ve = _em_step(10.0 ** float(lv_grid[int(np.argmin(h_abs))]))
|
||
em_ok = False
|
||
else:
|
||
# ---- 生产档（n>N_EXACT 且 m>N_SUBSAMPLE）：随机子样本精确迹二分 → 全群 MME ----
|
||
# 全群精确迹 O(n²) 不可行；EM 收缩率 ρ→0.999 冻结在初始猜测，Hutchinson 噪声
|
||
# 又淹没 h(λ) 信号（~2e-4 vs 噪声 ~2e-3）。子样本 iid 抽样 REML 为全群方差组分
|
||
# 一致估计（n_sub=3000 抽样 SE h² ≈ ±0.04，EBV corr>0.998），抽样误差即真实
|
||
# 不确定性。内联递归把 N_EXACT 临时抬到 10⁹，让子样本走精确迹二分求 λ*。
|
||
chosen = set(rng.choice(obs_ids, size=min(N_SUBSAMPLE, m), replace=False).tolist())
|
||
# 祖先闭包：parent_of 以个体 id 为键、值亦为 id（与 _build_ainv_sparse 一致）；
|
||
# 勿用 idx[]/order[] 做中间层（那是位置空间，混用会把闭包加进错对象）。
|
||
q: list[int] = list(chosen)
|
||
while q:
|
||
cur = q.pop()
|
||
d, s = parent_of.get(cur, (None, None))
|
||
for pp in (d, s):
|
||
if pp is not None and pp not in chosen:
|
||
chosen.add(pp)
|
||
q.append(pp)
|
||
ped_sub = []
|
||
for i in chosen:
|
||
d, s = parent_of.get(i, (None, None))
|
||
ped_sub.append({
|
||
"individual": i,
|
||
"dam": d if d is not None else None,
|
||
"sire": s if s is not None else None,
|
||
})
|
||
phen_sub = {k: v for k, v in phenotypes.items() if k in chosen}
|
||
fixed_sub = None
|
||
if fixed:
|
||
fixed_sub = {fk: {k: v for k, v in fv.items() if k in chosen}
|
||
for fk, fv in fixed.items()}
|
||
cov_sub = None
|
||
if covariate:
|
||
cov_sub = {k: v for k, v in covariate.items() if k in chosen}
|
||
saved_exact = N_EXACT
|
||
try:
|
||
N_EXACT = 10 ** 9 # noqa: PLW0603 临时抬档（子样本个体数远小于全群）
|
||
sub = solve(ped_sub, phen_sub, fixed=fixed_sub, covariate=cov_sub)
|
||
finally:
|
||
N_EXACT = saved_exact
|
||
Va = max(float(sub["sigma_a"]), _FLOOR)
|
||
Ve = max(float(sub["sigma_e"]), _FLOOR)
|
||
em_ok = bool(sub.get("converged", False))
|
||
em_it = int(sub.get("n_iter", 0))
|
||
|
||
lam = Ve / Va
|
||
mv = _matvec(lam)
|
||
sol, cg_iter, cg_ok = _cg_solve(mv, mme_rhs)
|
||
u = sol[p:]
|
||
h2 = Va / (Va + Ve) if (Va + Ve) > 0 else None
|
||
|
||
# 可靠性：n ≤ EXACT_DIAG_N（稠密逆仍便宜）→ diag(C22) 用稠密 MME 逆精确求（golden
|
||
# 等价档，无 MC 噪声——全同胞小样本 h²→1 边界 Hutchinson 均值误差实测 ~0.08 量级，
|
||
# 见 e2e_spatial_sparse n=20 段）；n 大时退 Hutchinson 随机探测（排序可靠、量级可信）。
|
||
if n <= EXACT_DIAG_N:
|
||
Ainv_dense = np.zeros((n, n))
|
||
np.add.at(Ainv_dense, (ainv_ii, ainv_jj), ainv_vv) # 三元组含重复 (i,j)（多后代），须累加
|
||
big = np.zeros((p + n, p + n))
|
||
big[:p, :p] = XtX
|
||
big[:p, p:] = XtZ
|
||
big[p:, :p] = XtZ.T
|
||
big[p:, p:] = np.diag(ZtZ_diag) + lam * Ainv_dense
|
||
try:
|
||
big_inv = np.linalg.inv(big)
|
||
except np.linalg.LinAlgError:
|
||
big_inv = np.linalg.pinv(big)
|
||
C22d = big_inv[p:, p:]
|
||
pev = np.diag(C22d) * Ve
|
||
else:
|
||
diag_c22 = np.zeros(n)
|
||
for _ in range(HUTCH_K):
|
||
pv = rng.choice([-1.0, 1.0], size=n)
|
||
rhs2 = np.zeros(p + n)
|
||
rhs2[p:] = pv
|
||
sol2, _it2, _ok2 = _cg_solve(mv, rhs2)
|
||
diag_c22 += sol2[p:] * pv
|
||
diag_c22 /= HUTCH_K
|
||
pev = diag_c22 * Ve
|
||
|
||
ebv = {ind: float(u[idx[ind]]) for ind in order}
|
||
rel = {ind: float(np.clip(1.0 - pev[idx[ind]] / Va, 0.0, 1.0)) for ind in order}
|
||
|
||
warning = None
|
||
if use_production:
|
||
# 生产档（n>N_EXACT 且 m>N_SUBSAMPLE）：方差组分来自随机子样本
|
||
# （≤N_SUBSAMPLE 观测，含祖先闭包）的精确 REML 估计——一致估计但带抽样误差
|
||
# （n_sub=3000 SE h² ≈ ±0.04）；EBV 由全群 MME 在子样本 λ* 下求解，
|
||
# 对 λ 稳健（偏 25% 内 Kendall τ>0.97），排序可靠。
|
||
warning = (f"生产档（个体数 {n}）方差分量取自随机子样本"
|
||
f"（≤{N_SUBSAMPLE} 观测，含祖先闭包）的 REML 一致估计：存在抽样误差"
|
||
"（h² 标准误约 ±0.04），h²/σ² 报告值请视为近似区间；EBV 由全群 MME"
|
||
"在子样本 λ 下求解，排名可靠。")
|
||
elif not em_ok:
|
||
warning = ("REML 解落在 log10λ 二分区间外（似然面边界最优），"
|
||
"结果采用 |h| 最小网格点，可能不精确。")
|
||
if not cg_ok:
|
||
resid = mme_rhs - mv(sol)
|
||
rel_res = float(np.linalg.norm(resid) / (np.linalg.norm(mme_rhs) + 1e-12))
|
||
w = (f"共轭梯度 {cg_iter} 次迭代未达收敛阈值（相对残差 {rel_res:.2e}），"
|
||
"结果采用已探明最优，可靠性为 Hutchinson 近似。")
|
||
warning = f"{warning} {w}" if warning else w
|
||
return {
|
||
"ebv": ebv,
|
||
"reliability": rel,
|
||
"h2": float(h2) if h2 is not None else None,
|
||
"sigma_a": float(Va),
|
||
"sigma_e": float(Ve),
|
||
"n_obs": m,
|
||
"n_individuals": n,
|
||
"n_base": n_base,
|
||
"n_with_parents": n_with_parents,
|
||
"n_fixed": n_fixed,
|
||
"converged": em_ok and cg_ok,
|
||
"n_iter": em_it,
|
||
"solver": "sparse-em",
|
||
"cg_iter": cg_iter,
|
||
"reliability_approx": n > EXACT_DIAG_N,
|
||
"warning": warning,
|
||
}
|
||
|
||
# 稠密分支（n ≤ N_SPARSE）：重建等价 Z / ZtZ（bincount 已给各体观测计数）
|
||
Z = np.zeros((m, n))
|
||
for k, ind in enumerate(obs_ids):
|
||
Z[k, idx[ind]] = 1.0
|
||
ZtZ = np.diag(ZtZ_diag)
|
||
|
||
def _reml_ll(Va: float, Ve: float) -> float:
|
||
"""精确 REML 对数似然（直接 V 计算，O(m³)）。"""
|
||
V = (Z @ A @ Z.T) * Va + np.eye(m) * Ve
|
||
VinvX = np.linalg.solve(V, X) # p 个 RHS，1 次 LU
|
||
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)
|
||
|
||
def _mme_solve(lam: float) -> tuple[np.ndarray, np.ndarray]:
|
||
big = np.zeros((p + n, p + n))
|
||
big[:p, :p] = XtX
|
||
big[:p, p:] = XtZ
|
||
big[p:, :p] = XtZ.T
|
||
big[p:, p:] = ZtZ + lam * Ainv
|
||
rhs = np.concatenate([Xty, Zty])
|
||
try:
|
||
sol = np.linalg.solve(big, rhs)
|
||
big_inv = np.linalg.inv(big)
|
||
except np.linalg.LinAlgError:
|
||
sol = np.linalg.pinv(big) @ rhs
|
||
big_inv = np.linalg.pinv(big)
|
||
return sol, big_inv
|
||
|
||
best = {"h2": None, "ll": -np.inf, "Va": 0.0, "Ve": 0.0}
|
||
|
||
def _profile(h2: float) -> float:
|
||
lam = (1.0 - h2) / h2
|
||
sol, _ = _mme_solve(lam)
|
||
b = sol[:p]
|
||
u = sol[p:]
|
||
yPy = float(y @ y - b @ Xty - u @ Zty)
|
||
Ve = max(yPy / max(m - p, 1), _FLOOR)
|
||
Va = Ve / lam
|
||
ll = _reml_ll(Va, Ve)
|
||
if ll > best["ll"]:
|
||
best.update(h2=h2, ll=ll, Va=Va, Ve=Ve)
|
||
return ll
|
||
|
||
h2_opt, n_iter = _golden_max(_profile, H2_MIN, H2_MAX, tol=tol)
|
||
# 黄金分割端点通常已探明；再用最优 h2 精确求一次
|
||
_profile(h2_opt)
|
||
h2 = best["h2"] if best["h2"] is not None else h2_opt
|
||
Va, Ve = best["Va"], best["Ve"]
|
||
|
||
lam = (1.0 - h2) / h2
|
||
sol, big_inv = _mme_solve(lam)
|
||
u = sol[p:]
|
||
C22 = big_inv[p:, p:]
|
||
|
||
ebv = {ind: float(u[idx[ind]]) for ind in order}
|
||
pev = np.diag(C22) * Ve
|
||
rel = {ind: float(np.clip(1.0 - pev[idx[ind]] / Va, 0.0, 1.0)) for ind in order}
|
||
|
||
warning = None
|
||
if n_iter >= MAX_PROFILE_EVALS:
|
||
warning = "剖面 REML 未完全收敛（似然面极平/边界最优），结果已采用已探明最优。"
|
||
return {
|
||
"ebv": ebv,
|
||
"reliability": rel,
|
||
"h2": float(h2),
|
||
"sigma_a": float(Va),
|
||
"sigma_e": float(Ve),
|
||
"n_obs": m,
|
||
"n_individuals": n,
|
||
"n_base": n_base,
|
||
"n_with_parents": n_with_parents,
|
||
"n_fixed": n_fixed,
|
||
"converged": n_iter < MAX_PROFILE_EVALS,
|
||
"n_iter": n_iter,
|
||
"solver": "dense",
|
||
"warning": warning,
|
||
}
|
||
|
||
|
||
def solve_spatial(pedigree: list[dict], phenotypes: dict, *, coords: dict,
|
||
tol: float = TOL, fixed: dict | None = None,
|
||
covariate: dict | None = None, aniso: bool = False) -> dict:
|
||
"""AR1×AR1 空间协方差动物模型 BLUP（v1：ρ_row=ρ_col=ρ 降维；aniso=True：双参数）。
|
||
|
||
y = Xb + Zu + e，e~N(0, σ²e·R)，R_ij = ρ^(|Δrow|+|Δcol|)（行/列 AR1 之积，
|
||
即 AR1(ρ)⊗AR1(ρ)）。R⁻¹ 用纯 numpy eigh 对角化；REML 对 (h², ρ) 坐标上升黄金分割。
|
||
aniso=True 时 R_ij = ρ_row^|Δrow|·ρ_col^|Δcol|，对 (h², ρ_row, ρ_col) 坐标上升
|
||
（方法标签 method="AR1×AR1(aniso)"、落 rho_row/rho_col）；aniso=False 时输出与 v1 逐字节一致。
|
||
|
||
参数:
|
||
pedigree / phenotypes / fixed / covariate: 同 solve()。
|
||
coords: {个体id: (row_no, col_no)}——每个观测个体必须提供行/列坐标，
|
||
缺失 → ValueError（门禁：缺坐标无法建模空间残差）。
|
||
aniso: 是否启用各向异性双参数 ρ_row/ρ_col（默认 False=单 ρ 降维）。
|
||
|
||
返回:
|
||
ebv / reliability / h2 / sigma_a / sigma_e / n_obs / n_individuals /
|
||
n_base / n_with_parents / n_fixed / converged / n_iter / warning（同 solve）；
|
||
另含 rho（单参，aniso=False）或 rho_row/rho_col（aniso=True）、
|
||
spatial=True、method="AR1×AR1" 或 "AR1×AR1(aniso)"。
|
||
"""
|
||
if aniso:
|
||
return _solve_spatial_aniso(pedigree, phenotypes, coords=coords, tol=tol,
|
||
fixed=fixed, covariate=covariate)
|
||
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))
|
||
for pid in phenotypes:
|
||
if pid not in seen:
|
||
seen.add(pid)
|
||
base.append(pid)
|
||
if not seen:
|
||
raise ValueError("系谱与表型均为空")
|
||
|
||
order, idx = _order_pedigree(base, non_base)
|
||
n = len(order)
|
||
n_base = len(base)
|
||
parent_of: dict[int, tuple[int | None, int | None]] = {i: (d, s) for (i, d, s) in non_base}
|
||
n_with_parents = len(parent_of)
|
||
|
||
obs_ids = [ind for ind in order if ind in phenotypes]
|
||
m = len(obs_ids)
|
||
if m < 2:
|
||
raise ValueError("有效表型个体不足 2 个，无法估计方差组分")
|
||
missing = [i for i in obs_ids if coords.get(i) is None]
|
||
if missing:
|
||
raise ValueError(f"缺行/列坐标的个体（空间协方差需每个观测提供 row_no/col_no）: {missing[:5]}")
|
||
|
||
X = _build_design(obs_ids, fixed, covariate)
|
||
p = X.shape[1]
|
||
n_fixed = p - 1
|
||
y = np.array([float(phenotypes[i]) for i in obs_ids], dtype=float)
|
||
|
||
if n_with_parents == 0:
|
||
beta = np.linalg.lstsq(X, y, rcond=None)[0]
|
||
resid = y - X @ beta
|
||
ebv: dict[int, float] = {}
|
||
for k, ind in enumerate(obs_ids):
|
||
ebv[ind] = float(resid[k])
|
||
for ind in order:
|
||
ebv.setdefault(ind, 0.0)
|
||
return {
|
||
"ebv": ebv,
|
||
"reliability": {ind: 0.0 for ind in order},
|
||
"h2": None, "sigma_a": 0.0, "sigma_e": 0.0,
|
||
"rho": 0.0, "spatial": True, "method": "AR1×AR1",
|
||
"n_obs": m, "n_individuals": n, "n_base": n_base,
|
||
"n_with_parents": 0, "n_fixed": n_fixed,
|
||
"converged": True, "n_iter": 0,
|
||
"solver": "dense",
|
||
"warning": "无系谱信息（所有个体均无父母）：EBV 仅为表型残差，未估计遗传力与可靠性。",
|
||
}
|
||
|
||
A, Ainv = _build_ainv(order, n_base, parent_of)
|
||
Z = np.zeros((m, n))
|
||
for k, ind in enumerate(obs_ids):
|
||
Z[k, idx[ind]] = 1.0
|
||
rows = np.array([coords[i][0] for i in obs_ids], dtype=float)
|
||
cols = np.array([coords[i][1] for i in obs_ids], dtype=float)
|
||
# R_ij = ρ^(|Δrow|+|Δcol|)
|
||
dmat = np.abs(rows[:, None] - rows[None, :]) + np.abs(cols[:, None] - cols[None, :])
|
||
XtX, Xty = X.T @ X, X.T @ y
|
||
Zty = Z.T @ y
|
||
ZtZ = Z.T @ Z
|
||
XtZ = X.T @ Z
|
||
|
||
def _rinv(rho: float) -> np.ndarray:
|
||
R = np.power(rho, dmat)
|
||
evals, Q = np.linalg.eigh(R)
|
||
return Q @ np.diag(1.0 / np.maximum(evals, _FLOOR)) @ Q.T
|
||
|
||
def _mme_solve(lam: float, rho: float) -> tuple[np.ndarray, np.ndarray]:
|
||
Rinv = _rinv(rho)
|
||
XrX = X.T @ Rinv @ X
|
||
XrZ = X.T @ Rinv @ Z
|
||
ZrZ = Z.T @ Rinv @ Z
|
||
Xry = X.T @ Rinv @ y
|
||
Zry = Z.T @ Rinv @ y
|
||
big = np.zeros((p + n, p + n))
|
||
big[:p, :p] = XrX
|
||
big[:p, p:] = XrZ
|
||
big[p:, :p] = XrZ.T
|
||
big[p:, p:] = ZrZ + lam * Ainv
|
||
rhs = np.concatenate([Xry, Zry])
|
||
try:
|
||
sol = np.linalg.solve(big, rhs)
|
||
big_inv = np.linalg.inv(big)
|
||
except np.linalg.LinAlgError:
|
||
sol = np.linalg.pinv(big) @ rhs
|
||
big_inv = np.linalg.pinv(big)
|
||
return sol, big_inv
|
||
|
||
def _reml_ll(Va: float, Ve: float, rho: float) -> float:
|
||
R = np.power(rho, dmat)
|
||
V = (Z @ A @ Z.T) * Va + R * Ve
|
||
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)
|
||
|
||
RHO_LO, RHO_HI = 0.02, 0.98
|
||
best = {"h2": None, "rho": 0.0, "ll": -np.inf, "Va": 0.0, "Ve": 0.0}
|
||
h2_l, rho_l = 0.5, 0.5
|
||
|
||
def _track(h2: float, rho: float) -> float:
|
||
lam = (1.0 - h2) / h2
|
||
sol, _ = _mme_solve(lam, rho)
|
||
b, u = sol[:p], sol[p:]
|
||
resid = y - X @ b - Z @ u
|
||
yPy = float(resid @ _rinv(rho) @ resid)
|
||
Ve = max(yPy / max(m - p, 1), _FLOOR)
|
||
Va = Ve / lam
|
||
ll = _reml_ll(Va, Ve, rho)
|
||
if ll > best["ll"]:
|
||
best.update(h2=h2, rho=rho, ll=ll, Va=Va, Ve=Ve)
|
||
return ll
|
||
|
||
def _f_h2(h2: float) -> float:
|
||
return _track(h2, rho_l)
|
||
|
||
def _f_rho(rho: float) -> float:
|
||
return _track(h2_l, rho)
|
||
|
||
converged = False
|
||
n_iter = 0
|
||
prev_best = -np.inf
|
||
for _ in range(MAX_GXE_ITERS):
|
||
n_iter += 1
|
||
h2_l, _ = _golden_max(_f_h2, H2_MIN, H2_MAX, tol=1e-3)
|
||
rho_l, _ = _golden_max(_f_rho, RHO_LO, RHO_HI, tol=1e-3)
|
||
if best["ll"] - prev_best < 1e-2:
|
||
converged = True
|
||
break
|
||
prev_best = best["ll"]
|
||
_track(h2_l, rho_l)
|
||
|
||
h2 = best["h2"] if best["h2"] is not None else h2_l
|
||
Va, Ve = best["Va"], best["Ve"]
|
||
rho = best["rho"]
|
||
|
||
lam = (1.0 - h2) / h2
|
||
sol, big_inv, *_ = _mme_solve(lam, rho)
|
||
u = sol[p:]
|
||
C22 = big_inv[p:, p:]
|
||
|
||
ebv = {ind: float(u[idx[ind]]) for ind in order}
|
||
pev = np.diag(C22) * Ve
|
||
rel = {ind: float(np.clip(1.0 - pev[idx[ind]] / Va, 0.0, 1.0)) for ind in order}
|
||
|
||
warning = None
|
||
if rho <= RHO_LO * 1.05:
|
||
warning = "空间自相关 ρ 收敛到下界，残差接近独立（空间结构弱）。"
|
||
elif not converged:
|
||
warning = "剖面 REML 坐标上升未完全收敛（空间似然面极平/边界最优），结果采用已探明最优。"
|
||
return {
|
||
"ebv": ebv,
|
||
"reliability": rel,
|
||
"h2": float(h2),
|
||
"sigma_a": float(Va),
|
||
"sigma_e": float(Ve),
|
||
"rho": float(rho),
|
||
"spatial": True,
|
||
"method": "AR1×AR1",
|
||
"n_obs": m,
|
||
"n_individuals": n,
|
||
"n_base": n_base,
|
||
"n_with_parents": n_with_parents,
|
||
"n_fixed": n_fixed,
|
||
"converged": converged,
|
||
"n_iter": n_iter,
|
||
"solver": "dense",
|
||
"warning": warning,
|
||
}
|
||
|
||
|
||
def _solve_spatial_aniso(pedigree: list[dict], phenotypes: dict, *, coords: dict,
|
||
tol: float = TOL, fixed: dict | None = None,
|
||
covariate: dict | None = None) -> dict:
|
||
"""AR1×AR1 各向异性空间协方差 BLUP（aniso=True 分支，双参数 ρ_row/ρ_col）。
|
||
|
||
R_ij = ρ_row^|Δrow| · ρ_col^|Δcol|（行/列独立 AR1 之积）。REML 对 (h², ρ_row, ρ_col)
|
||
三变量坐标上升黄金分割；门禁/互斥/无系谱语义与 v1 solve_spatial 逐项一致，
|
||
仅 ρ 由单参扩展为行/列双参。方法标签 method="AR1×AR1(aniso)"、落 rho_row/rho_col。
|
||
"""
|
||
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))
|
||
for pid in phenotypes:
|
||
if pid not in seen:
|
||
seen.add(pid)
|
||
base.append(pid)
|
||
if not seen:
|
||
raise ValueError("系谱与表型均为空")
|
||
|
||
order, idx = _order_pedigree(base, non_base)
|
||
n = len(order)
|
||
n_base = len(base)
|
||
parent_of: dict[int, tuple[int | None, int | None]] = {i: (d, s) for (i, d, s) in non_base}
|
||
n_with_parents = len(parent_of)
|
||
|
||
obs_ids = [ind for ind in order if ind in phenotypes]
|
||
m = len(obs_ids)
|
||
if m < 2:
|
||
raise ValueError("有效表型个体不足 2 个，无法估计方差组分")
|
||
missing = [i for i in obs_ids if coords.get(i) is None]
|
||
if missing:
|
||
raise ValueError(f"缺行/列坐标的个体（空间协方差需每个观测提供 row_no/col_no）: {missing[:5]}")
|
||
|
||
X = _build_design(obs_ids, fixed, covariate)
|
||
p = X.shape[1]
|
||
n_fixed = p - 1
|
||
y = np.array([float(phenotypes[i]) for i in obs_ids], dtype=float)
|
||
|
||
if n_with_parents == 0:
|
||
beta = np.linalg.lstsq(X, y, rcond=None)[0]
|
||
resid = y - X @ beta
|
||
ebv: dict[int, float] = {}
|
||
for k, ind in enumerate(obs_ids):
|
||
ebv[ind] = float(resid[k])
|
||
for ind in order:
|
||
ebv.setdefault(ind, 0.0)
|
||
return {
|
||
"ebv": ebv,
|
||
"reliability": {ind: 0.0 for ind in order},
|
||
"h2": None, "sigma_a": 0.0, "sigma_e": 0.0,
|
||
"rho_row": 0.0, "rho_col": 0.0, "spatial": True, "method": "AR1×AR1(aniso)",
|
||
"n_obs": m, "n_individuals": n, "n_base": n_base,
|
||
"n_with_parents": 0, "n_fixed": n_fixed,
|
||
"converged": True, "n_iter": 0,
|
||
"solver": "dense",
|
||
"warning": "无系谱信息（所有个体均无父母）：EBV 仅为表型残差，未估计遗传力与可靠性。",
|
||
}
|
||
|
||
A, Ainv = _build_ainv(order, n_base, parent_of)
|
||
Z = np.zeros((m, n))
|
||
for k, ind in enumerate(obs_ids):
|
||
Z[k, idx[ind]] = 1.0
|
||
rows = np.array([coords[i][0] for i in obs_ids], dtype=float)
|
||
cols = np.array([coords[i][1] for i in obs_ids], dtype=float)
|
||
dmat_row = np.abs(rows[:, None] - rows[None, :])
|
||
dmat_col = np.abs(cols[:, None] - cols[None, :])
|
||
XtX, Xty = X.T @ X, X.T @ y
|
||
Zty = Z.T @ y
|
||
ZtZ = Z.T @ Z
|
||
XtZ = X.T @ Z
|
||
|
||
def _rinv(rr: float, rc: float) -> np.ndarray:
|
||
R = np.power(rr, dmat_row) * np.power(rc, dmat_col)
|
||
evals, Q = np.linalg.eigh(R)
|
||
return Q @ np.diag(1.0 / np.maximum(evals, _FLOOR)) @ Q.T
|
||
|
||
def _mme_solve(lam: float, rr: float, rc: float) -> tuple[np.ndarray, np.ndarray]:
|
||
Rinv = _rinv(rr, rc)
|
||
XrX = X.T @ Rinv @ X
|
||
XrZ = X.T @ Rinv @ Z
|
||
ZrZ = Z.T @ Rinv @ Z
|
||
Xry = X.T @ Rinv @ y
|
||
Zry = Z.T @ Rinv @ y
|
||
big = np.zeros((p + n, p + n))
|
||
big[:p, :p] = XrX
|
||
big[:p, p:] = XrZ
|
||
big[p:, :p] = XrZ.T
|
||
big[p:, p:] = ZrZ + lam * Ainv
|
||
rhs = np.concatenate([Xry, Zry])
|
||
try:
|
||
sol = np.linalg.solve(big, rhs)
|
||
big_inv = np.linalg.inv(big)
|
||
except np.linalg.LinAlgError:
|
||
sol = np.linalg.pinv(big) @ rhs
|
||
big_inv = np.linalg.pinv(big)
|
||
return sol, big_inv
|
||
|
||
def _reml_ll(Va: float, Ve: float, rr: float, rc: float) -> float:
|
||
R = np.power(rr, dmat_row) * np.power(rc, dmat_col)
|
||
V = (Z @ A @ Z.T) * Va + R * Ve
|
||
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)
|
||
|
||
RHO_LO, RHO_HI = 0.02, 0.98
|
||
best = {"h2": None, "rho_row": 0.5, "rho_col": 0.5, "ll": -np.inf,
|
||
"Va": 0.0, "Ve": 0.0}
|
||
h2_l, rho_row_l, rho_col_l = 0.5, 0.5, 0.5
|
||
|
||
def _track(h2: float, rr: float, rc: float) -> float:
|
||
lam = (1.0 - h2) / h2
|
||
sol, _ = _mme_solve(lam, rr, rc)
|
||
b, u = sol[:p], sol[p:]
|
||
resid = y - X @ b - Z @ u
|
||
yPy = float(resid @ _rinv(rr, rc) @ resid)
|
||
Ve = max(yPy / max(m - p, 1), _FLOOR)
|
||
Va = Ve / lam
|
||
ll = _reml_ll(Va, Ve, rr, rc)
|
||
if ll > best["ll"]:
|
||
best.update(h2=h2, rho_row=rr, rho_col=rc, ll=ll, Va=Va, Ve=Ve)
|
||
return ll
|
||
|
||
def _f_h2(h2: float) -> float:
|
||
return _track(h2, rho_row_l, rho_col_l)
|
||
|
||
def _f_rr(rr: float) -> float:
|
||
return _track(h2_l, rr, rho_col_l)
|
||
|
||
def _f_rc(rc: float) -> float:
|
||
return _track(h2_l, rho_row_l, rc)
|
||
|
||
converged = False
|
||
n_iter = 0
|
||
prev_best = -np.inf
|
||
for _ in range(MAX_GXE_ITERS):
|
||
n_iter += 1
|
||
h2_l, _ = _golden_max(_f_h2, H2_MIN, H2_MAX, tol=1e-3)
|
||
rho_row_l, _ = _golden_max(_f_rr, RHO_LO, RHO_HI, tol=1e-3)
|
||
rho_col_l, _ = _golden_max(_f_rc, RHO_LO, RHO_HI, tol=1e-3)
|
||
if best["ll"] - prev_best < 1e-2:
|
||
converged = True
|
||
break
|
||
prev_best = best["ll"]
|
||
_track(h2_l, rho_row_l, rho_col_l)
|
||
|
||
h2 = best["h2"] if best["h2"] is not None else h2_l
|
||
Va, Ve = best["Va"], best["Ve"]
|
||
rho_row, rho_col = best["rho_row"], best["rho_col"]
|
||
|
||
lam = (1.0 - h2) / h2
|
||
sol, big_inv, *_ = _mme_solve(lam, rho_row, rho_col)
|
||
u = sol[p:]
|
||
C22 = big_inv[p:, p:]
|
||
|
||
ebv = {ind: float(u[idx[ind]]) for ind in order}
|
||
pev = np.diag(C22) * Ve
|
||
rel = {ind: float(np.clip(1.0 - pev[idx[ind]] / Va, 0.0, 1.0)) for ind in order}
|
||
|
||
warning = None
|
||
if rho_row <= RHO_LO * 1.05 and rho_col <= RHO_LO * 1.05:
|
||
warning = "空间自相关 ρ_row/ρ_col 均收敛到下界，残差接近独立（空间结构弱）。"
|
||
elif not converged:
|
||
warning = "剖面 REML 坐标上升未完全收敛（空间似然面极平/边界最优），结果采用已探明最优。"
|
||
return {
|
||
"ebv": ebv,
|
||
"reliability": rel,
|
||
"h2": float(h2),
|
||
"sigma_a": float(Va),
|
||
"sigma_e": float(Ve),
|
||
"rho_row": float(rho_row),
|
||
"rho_col": float(rho_col),
|
||
"spatial": True,
|
||
"method": "AR1×AR1(aniso)",
|
||
"n_obs": m,
|
||
"n_individuals": n,
|
||
"n_base": n_base,
|
||
"n_with_parents": n_with_parents,
|
||
"n_fixed": n_fixed,
|
||
"converged": converged,
|
||
"n_iter": n_iter,
|
||
"solver": "dense",
|
||
"warning": warning,
|
||
}
|
||
|
||
|
||
def _solve_gxe(pedigree: list[dict], phenotypes: dict, *, fixed: dict | None,
|
||
covariate: dict | None, gxe: dict, record_map: dict, tol: float,
|
||
factor_label: str = "G×E", factor_name: str = "环境") -> dict:
|
||
"""G×E 分支：y = Xb + Z₁u + Z₂v + e，u~N(0,A·σ²a)、v~N(0,I·σ²gxe)、e~N(0,I·σ²e)。
|
||
|
||
交互随机效应槽 Z₂ 按 (基因型, 因子水平) 分组，不限定因子语义——
|
||
G×R（砧木）复用同一求解器：factor_label/factor_name 仅影响警告文案
|
||
（默认「G×E/环境」，G×R 传「G×R/砧木」），求解逻辑逐字节不变。
|
||
|
||
三方差分量用剖面 REML 估计：对 (log10 σ²a/σ²e, log10 σ²gxe/σ²e) 坐标上升黄金分割，
|
||
比率边界 λ∈[1e-3,1e3] 恒正，避免 EM-REML 边界崩塌（σ²e→0 → λ→0 → 插值）。
|
||
记录级：phenotypes/gxe/fixed/covariate 均以记录 id 为键；record_map 记录→加性个体。
|
||
"""
|
||
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))
|
||
for pid in phenotypes:
|
||
ind = record_map.get(pid, pid)
|
||
if ind not in seen:
|
||
seen.add(ind)
|
||
base.append(ind)
|
||
if not seen:
|
||
raise ValueError("系谱与表型均为空")
|
||
|
||
order, idx = _order_pedigree(base, non_base)
|
||
n = len(order)
|
||
n_base = len(base)
|
||
parent_of: dict[int, tuple[int | None, int | None]] = {i: (d, s) for (i, d, s) in non_base}
|
||
n_with_parents = len(parent_of)
|
||
|
||
rec_ids = list(phenotypes.keys())
|
||
m = len(rec_ids)
|
||
if m < 2:
|
||
raise ValueError("有效表型个体不足 2 个，无法估计方差组分")
|
||
|
||
A, Ainv = _build_ainv(order, n_base, parent_of)
|
||
X = _build_design(rec_ids, fixed, covariate)
|
||
p = X.shape[1]
|
||
n_fixed = p - 1
|
||
y = np.array([float(phenotypes[i]) for i in rec_ids], dtype=float)
|
||
|
||
Z1 = np.zeros((m, n))
|
||
for k, rid in enumerate(rec_ids):
|
||
Z1[k, idx[record_map.get(rid, rid)]] = 1.0
|
||
|
||
pairs: dict[tuple[str, str], int] = {}
|
||
pair_list: list[tuple[str, str]] = []
|
||
for rid in rec_ids:
|
||
key = (str(gxe[rid][0]), str(gxe[rid][1]))
|
||
if key not in pairs:
|
||
pairs[key] = len(pair_list)
|
||
pair_list.append(key)
|
||
q = len(pair_list)
|
||
if q < 1:
|
||
raise ValueError("无 G×E 效应水平")
|
||
Z2 = np.zeros((m, q))
|
||
for k, rid in enumerate(rec_ids):
|
||
Z2[k, pairs[(str(gxe[rid][0]), str(gxe[rid][1]))]] = 1.0
|
||
|
||
XtX, Xty = X.T @ X, X.T @ y
|
||
Z1tZ1, Z1ty = Z1.T @ Z1, Z1.T @ y
|
||
Z2tZ2, Z2ty = Z2.T @ Z2, Z2.T @ y
|
||
XtZ1, XtZ2, Z1tZ2 = X.T @ Z1, X.T @ Z2, Z1.T @ Z2
|
||
|
||
def _mme(la: float, lg: float) -> tuple[np.ndarray, np.ndarray]:
|
||
big = np.zeros((p + n + q, p + n + q))
|
||
big[:p, :p] = XtX
|
||
big[:p, p:p + n] = XtZ1
|
||
big[:p, p + n:] = XtZ2
|
||
big[p:p + n, :p] = XtZ1.T
|
||
big[p:p + n, p:p + n] = Z1tZ1 + la * Ainv
|
||
big[p:p + n, p + n:] = Z1tZ2
|
||
big[p + n:, :p] = XtZ2.T
|
||
big[p + n:, p:p + n] = Z1tZ2.T
|
||
big[p + n:, p + n:] = Z2tZ2 + lg * np.eye(q)
|
||
rhs = np.concatenate([Xty, Z1ty, Z2ty])
|
||
try:
|
||
sol = np.linalg.solve(big, rhs)
|
||
big_inv = np.linalg.inv(big)
|
||
except np.linalg.LinAlgError:
|
||
sol = np.linalg.pinv(big) @ rhs
|
||
big_inv = np.linalg.pinv(big)
|
||
return sol, big_inv
|
||
|
||
# 剖面 REML：对 (log10 σ²a/σ²e, log10 σ²gxe/σ²e) 坐标上升黄金分割最大化。
|
||
# 比率边界与单随机效应分支的 h² 网格一致（λ∈[1e-3,1e3] 恒正），
|
||
# 避免 EM-REML 边界崩塌（σ²e→0 → λ→0 → MME 退化为插值）。
|
||
rank_x = max(np.linalg.matrix_rank(X), 1)
|
||
df = max(m - rank_x, 1)
|
||
|
||
def _profile(ra: float, rg: float) -> tuple[float, float]:
|
||
"""固定比率 (ra,rg)=(σ²a/σ²e, σ²gxe/σ²e)，剖面出 σ²e，返回 (REML LL, σ²e)。"""
|
||
V0 = (Z1 @ A @ Z1.T) * ra + (Z2 @ Z2.T) * rg + np.eye(m)
|
||
VinvX = np.linalg.solve(V0, X)
|
||
Vinvy = np.linalg.solve(V0, y)
|
||
XtVinvX = X.T @ VinvX
|
||
XtVinvY = X.T @ Vinvy
|
||
yP0y = float(y @ Vinvy) - float(XtVinvY @ np.linalg.solve(XtVinvX, XtVinvY))
|
||
_, lv = np.linalg.slogdet(V0)
|
||
_, lx = np.linalg.slogdet(XtVinvX)
|
||
ve = max(yP0y / df, _FLOOR)
|
||
ll = -0.5 * (df * np.log(ve) + float(lv) + float(lx))
|
||
return ll, ve
|
||
|
||
RA_LO, RA_HI = np.log10(1e-3), np.log10(1e3)
|
||
best = {"ll": -np.inf, "ra": 1.0, "rg": 1.0, "ve": 1.0}
|
||
ra_l, rg_l = 0.0, 0.0 # 起始比率=1（等方差）
|
||
|
||
def _track(ra: float, rg: float) -> float:
|
||
ll, ve = _profile(ra, rg)
|
||
if ll > best["ll"]:
|
||
best.update(ll=ll, ra=ra, rg=rg, ve=ve)
|
||
return ll
|
||
|
||
def _f_ra(lra: float) -> float:
|
||
return _track(10.0 ** lra, 10.0 ** rg_l)
|
||
|
||
def _f_rg(lrg: float) -> float:
|
||
return _track(10.0 ** ra_l, 10.0 ** lrg)
|
||
|
||
converged = False
|
||
n_iter = 0
|
||
prev_best = -np.inf
|
||
for _ in range(MAX_GXE_ITERS):
|
||
n_iter += 1
|
||
new_ra_l, _ = _golden_max(_f_ra, RA_LO, RA_HI, tol=1e-3)
|
||
ra_l = new_ra_l
|
||
new_rg_l, _ = _golden_max(_f_rg, RA_LO, RA_HI, tol=1e-3)
|
||
rg_l = new_rg_l
|
||
# 坐标上升在似然山脊上 LL 增益趋零即收敛；阈值放宽避免斜坡上反复横跳
|
||
if best["ll"] - prev_best < 1e-2:
|
||
converged = True
|
||
break
|
||
prev_best = best["ll"]
|
||
_track(10.0 ** ra_l, 10.0 ** rg_l)
|
||
|
||
ra, rg, ve = best["ra"], best["rg"], best["ve"]
|
||
va, vg = ra * ve, rg * ve
|
||
la, lg = 1.0 / ra, 1.0 / rg
|
||
sol, big_inv = _mme(la, lg)
|
||
u, v = sol[p:p + n], sol[p + n:]
|
||
C22 = big_inv[p:p + n, p:p + n]
|
||
|
||
ebv = {ind: float(u[idx[ind]]) for ind in order}
|
||
pev = np.diag(C22) * ve
|
||
rel = {ind: float(np.clip(1.0 - pev[idx[ind]] / va, 0.0, 1.0)) for ind in order}
|
||
gxe_eff = {f"{g}|{e}": float(v[i]) for i, (g, e) in enumerate(pair_list)}
|
||
|
||
total = va + vg + ve
|
||
h2 = va / total if total > 0 else None
|
||
gxe_ratio = vg / total if total > 0 else None
|
||
|
||
geno_env: dict[str, set[str]] = {}
|
||
for rid in rec_ids:
|
||
grp, env = str(gxe[rid][0]), str(gxe[rid][1])
|
||
geno_env.setdefault(grp, set()).add(env)
|
||
n_cross_env = sum(1 for s in geno_env.values() if len(s) >= 2)
|
||
n_genotypes = len(geno_env)
|
||
|
||
warning = None
|
||
if not converged:
|
||
warning = f"剖面 REML 坐标上升未完全收敛（{factor_label} 似然面可能极平/边界最优），结果采用已探明最优。"
|
||
if n_with_parents == 0:
|
||
w = "无系谱信息：σ²a 与 σ²gxe 的遗传学解释受限（EBV 无加性依据）。"
|
||
warning = w if warning is None else f"{warning} {w}"
|
||
if vg <= _FLOOR * 1.01:
|
||
w = f"{factor_label} 方差分量收敛到 0（数据不支持基因型×{factor_name}互作）。"
|
||
warning = w if warning is None else f"{warning} {w}"
|
||
return {
|
||
"ebv": ebv,
|
||
"reliability": rel,
|
||
"h2": float(h2) if h2 is not None else None,
|
||
"sigma_a": float(va),
|
||
"sigma_e": float(ve),
|
||
"sigma_gxe": float(vg),
|
||
"gxe_ratio": float(gxe_ratio) if gxe_ratio is not None else None,
|
||
"gxe_effects": gxe_eff,
|
||
"n_gxe": q,
|
||
"n_cross_env": n_cross_env,
|
||
"n_genotypes": n_genotypes,
|
||
"n_obs": m,
|
||
"n_individuals": n,
|
||
"n_base": n_base,
|
||
"n_with_parents": n_with_parents,
|
||
"n_fixed": n_fixed,
|
||
"converged": converged,
|
||
"n_iter": n_iter,
|
||
"solver": "dense",
|
||
"warning": warning,
|
||
}
|
||
|
||
|
||
def _solve_block(pedigree: list[dict], phenotypes: dict, *, fixed: dict | None,
|
||
covariate: dict | None, block_of: dict, record_map: dict, tol: float,
|
||
factor_label: str = "区组", factor_name: str = "区组") -> dict:
|
||
"""区组随机效应分支：y = Xb + Z₁u + Z₂v + e，u~N(0,A·σ²a)、v~N(0,I·σ²block)、e~N(0,I·σ²e)。
|
||
|
||
为不完全区组设计（增广/α-格子）服务：trial_design 已把 block_no 落库到观测株，
|
||
run_ablup 把区组作第二随机效应，精度收益进入遗传评估。Z₂ 按记录所属区组水平分组
|
||
（block_of: {记录id: 区组水平}），求解器与 G×E 分支同构（剖面 REML 坐标上升估
|
||
(log10 σ²a/σ²e, log10 σ²block/σ²e) 三方差分量），仅 Z₂ 分组与输出键改名。
|
||
"""
|
||
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))
|
||
for pid in phenotypes:
|
||
ind = record_map.get(pid, pid)
|
||
if ind not in seen:
|
||
seen.add(ind)
|
||
base.append(ind)
|
||
if not seen:
|
||
raise ValueError("系谱与表型均为空")
|
||
|
||
order, idx = _order_pedigree(base, non_base)
|
||
n = len(order)
|
||
n_base = len(base)
|
||
parent_of: dict[int, tuple[int | None, int | None]] = {i: (d, s) for (i, d, s) in non_base}
|
||
n_with_parents = len(parent_of)
|
||
|
||
rec_ids = list(phenotypes.keys())
|
||
m = len(rec_ids)
|
||
if m < 2:
|
||
raise ValueError("有效表型个体不足 2 个，无法估计方差组分")
|
||
|
||
A, Ainv = _build_ainv(order, n_base, parent_of)
|
||
X = _build_design(rec_ids, fixed, covariate)
|
||
p = X.shape[1]
|
||
n_fixed = p - 1
|
||
y = np.array([float(phenotypes[i]) for i in rec_ids], dtype=float)
|
||
|
||
Z1 = np.zeros((m, n))
|
||
for k, rid in enumerate(rec_ids):
|
||
Z1[k, idx[record_map.get(rid, rid)]] = 1.0
|
||
|
||
pairs: dict[str, int] = {}
|
||
pair_list: list[str] = []
|
||
for rid in rec_ids:
|
||
key = str(block_of[rid])
|
||
if key not in pairs:
|
||
pairs[key] = len(pair_list)
|
||
pair_list.append(key)
|
||
q = len(pair_list)
|
||
if q < 1:
|
||
raise ValueError("无区组效应水平")
|
||
Z2 = np.zeros((m, q))
|
||
for k, rid in enumerate(rec_ids):
|
||
Z2[k, pairs[str(block_of[rid])]] = 1.0
|
||
|
||
XtX, Xty = X.T @ X, X.T @ y
|
||
Z1tZ1, Z1ty = Z1.T @ Z1, Z1.T @ y
|
||
Z2tZ2, Z2ty = Z2.T @ Z2, Z2.T @ y
|
||
XtZ1, XtZ2, Z1tZ2 = X.T @ Z1, X.T @ Z2, Z1.T @ Z2
|
||
|
||
def _mme(la: float, lg: float) -> tuple[np.ndarray, np.ndarray]:
|
||
big = np.zeros((p + n + q, p + n + q))
|
||
big[:p, :p] = XtX
|
||
big[:p, p:p + n] = XtZ1
|
||
big[:p, p + n:] = XtZ2
|
||
big[p:p + n, :p] = XtZ1.T
|
||
big[p:p + n, p:p + n] = Z1tZ1 + la * Ainv
|
||
big[p:p + n, p + n:] = Z1tZ2
|
||
big[p + n:, :p] = XtZ2.T
|
||
big[p + n:, p:p + n] = Z1tZ2.T
|
||
big[p + n:, p + n:] = Z2tZ2 + lg * np.eye(q)
|
||
rhs = np.concatenate([Xty, Z1ty, Z2ty])
|
||
try:
|
||
sol = np.linalg.solve(big, rhs)
|
||
big_inv = np.linalg.inv(big)
|
||
except np.linalg.LinAlgError:
|
||
sol = np.linalg.pinv(big) @ rhs
|
||
big_inv = np.linalg.pinv(big)
|
||
return sol, big_inv
|
||
|
||
rank_x = max(np.linalg.matrix_rank(X), 1)
|
||
df = max(m - rank_x, 1)
|
||
|
||
def _profile(ra: float, rb: float) -> tuple[float, float]:
|
||
"""固定比率 (ra,rb)=(σ²a/σ²e, σ²block/σ²e)，剖面出 σ²e，返回 (REML LL, σ²e)。"""
|
||
V0 = (Z1 @ A @ Z1.T) * ra + (Z2 @ Z2.T) * rb + np.eye(m)
|
||
VinvX = np.linalg.solve(V0, X)
|
||
Vinvy = np.linalg.solve(V0, y)
|
||
XtVinvX = X.T @ VinvX
|
||
XtVinvY = X.T @ Vinvy
|
||
yP0y = float(y @ Vinvy) - float(XtVinvY @ np.linalg.solve(XtVinvX, XtVinvY))
|
||
_, lv = np.linalg.slogdet(V0)
|
||
_, lx = np.linalg.slogdet(XtVinvX)
|
||
ve = max(yP0y / df, _FLOOR)
|
||
ll = -0.5 * (df * np.log(ve) + float(lv) + float(lx))
|
||
return ll, ve
|
||
|
||
RA_LO, RA_HI = np.log10(1e-3), np.log10(1e3)
|
||
best = {"ll": -np.inf, "ra": 1.0, "rb": 1.0, "ve": 1.0}
|
||
ra_l, rb_l = 0.0, 0.0 # 起始比率=1（等方差）
|
||
|
||
def _track(ra: float, rb: float) -> float:
|
||
ll, ve = _profile(ra, rb)
|
||
if ll > best["ll"]:
|
||
best.update(ll=ll, ra=ra, rb=rb, ve=ve)
|
||
return ll
|
||
|
||
def _f_ra(lra: float) -> float:
|
||
return _track(10.0 ** lra, 10.0 ** rb_l)
|
||
|
||
def _f_rb(lrb: float) -> float:
|
||
return _track(10.0 ** ra_l, 10.0 ** lrb)
|
||
|
||
converged = False
|
||
n_iter = 0
|
||
prev_best = -np.inf
|
||
for _ in range(MAX_GXE_ITERS):
|
||
n_iter += 1
|
||
new_ra_l, _ = _golden_max(_f_ra, RA_LO, RA_HI, tol=1e-3)
|
||
ra_l = new_ra_l
|
||
new_rb_l, _ = _golden_max(_f_rb, RA_LO, RA_HI, tol=1e-3)
|
||
rb_l = new_rb_l
|
||
if best["ll"] - prev_best < 1e-2:
|
||
converged = True
|
||
break
|
||
prev_best = best["ll"]
|
||
_track(10.0 ** ra_l, 10.0 ** rb_l)
|
||
|
||
ra, rb, ve = best["ra"], best["rb"], best["ve"]
|
||
va, vb = ra * ve, rb * ve
|
||
la, lb = 1.0 / ra, 1.0 / rb
|
||
sol, big_inv = _mme(la, lb)
|
||
u, v = sol[p:p + n], sol[p + n:]
|
||
C22 = big_inv[p:p + n, p:p + n]
|
||
|
||
ebv = {ind: float(u[idx[ind]]) for ind in order}
|
||
pev = np.diag(C22) * ve
|
||
rel = {ind: float(np.clip(1.0 - pev[idx[ind]] / va, 0.0, 1.0)) for ind in order}
|
||
block_eff = {f"b{pair_list[i]}": float(v[i]) for i in range(q)}
|
||
|
||
total = va + vb + ve
|
||
h2 = va / total if total > 0 else None
|
||
block_ratio = vb / total if total > 0 else None
|
||
|
||
warning = None
|
||
if not converged:
|
||
warning = f"剖面 REML 坐标上升未完全收敛（{factor_label} 似然面可能极平/边界最优），结果采用已探明最优。"
|
||
if n_with_parents == 0:
|
||
w = "无系谱信息：σ²a 与 σ²block 的遗传学解释受限（EBV 无加性依据）。"
|
||
warning = w if warning is None else f"{warning} {w}"
|
||
if vb <= _FLOOR * 1.01:
|
||
w = f"{factor_label} 方差分量收敛到 0（数据不支持{factor_name}随机效应）。"
|
||
warning = w if warning is None else f"{warning} {w}"
|
||
return {
|
||
"ebv": ebv,
|
||
"reliability": rel,
|
||
"h2": float(h2) if h2 is not None else None,
|
||
"sigma_a": float(va),
|
||
"sigma_e": float(ve),
|
||
"sigma_block": float(vb),
|
||
"block_ratio": float(block_ratio) if block_ratio is not None else None,
|
||
"block_effects": block_eff,
|
||
"n_blocks": q,
|
||
"n_obs": m,
|
||
"n_individuals": n,
|
||
"n_base": n_base,
|
||
"n_with_parents": n_with_parents,
|
||
"n_fixed": n_fixed,
|
||
"converged": converged,
|
||
"n_iter": n_iter,
|
||
"solver": "dense",
|
||
"warning": warning,
|
||
}
|
||
|
||
|
||
def kfold_cv(pedigree: list[dict], phenotypes: dict, *, k: int = 5,
|
||
fixed: dict | None = None, covariate: dict | None = None,
|
||
gxe: dict | None = None, record_map: dict | None = None,
|
||
tol: float = TOL, seed: int = 20260804) -> dict:
|
||
"""k 折交叉验证（外部验证：预测准确度，区别于 PEV 可靠性）。
|
||
|
||
折划分按**加性个体**分层：优先按父本（sire）家系分组留出（同一家系不跨训练/测试双折，
|
||
避免同家系信息泄漏导致精度高估）；无 sire 的个体自成一家系（单个体）；家系用固定种子
|
||
shuffle 后 round-robin 分配到 k 折。gxe 分支同一个体多环境记录一起留出。
|
||
|
||
每折：训练子集 re-estimate 方差组分（solve 复用）→ 留出个体 EBV 由系谱预测
|
||
（个体保留在 A 矩阵，仅掩蔽其表型）→ pearson(EBV, 观测表型均值) + RMSE。
|
||
单折失败（训练不足 / G×E 结构不可辨识）→ 该折 pearson=None 并累计 warning，不中断整体。
|
||
|
||
参数: 同 solve()；phenotypes 键为记录/个体 id，gxe 分支按 record_map 归一到个体。
|
||
|
||
返回:
|
||
k / n_total / n_individuals: int；
|
||
folds: [{fold, n_train, n_test, n_eval, pearson, rmse, error?, h2?}]；
|
||
mean_pearson / mean_rmse: 有效折均值；pooled_pearson / pooled_rmse: 合并全部
|
||
(pred, obs) 对计算；cv_accuracy: 预测准确度（=mean_pearson）；
|
||
h2: 有效折均值；warning: str | None（失败折汇总）。
|
||
"""
|
||
if gxe is not None:
|
||
ind_of = {rid: record_map.get(rid, rid) for rid in phenotypes}
|
||
else:
|
||
ind_of = {ind: ind for ind in phenotypes}
|
||
individuals = sorted({v for v in ind_of.values()})
|
||
n_total = len(phenotypes)
|
||
n_individuals = len(individuals)
|
||
if n_individuals < 2 or n_total < 2:
|
||
return {
|
||
"k": k, "n_total": n_total, "n_individuals": n_individuals,
|
||
"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))
|
||
|
||
sire_of: dict = {}
|
||
for rec in pedigree:
|
||
if rec.get("sire") is not None:
|
||
sire_of.setdefault(rec["individual"], rec["sire"])
|
||
fam_of = {ind: sire_of.get(ind, ("self", ind)) for ind in individuals}
|
||
fams = sorted({f for f in fam_of.values()})
|
||
rng = random.Random(seed)
|
||
rng.shuffle(fams)
|
||
fam_bin = {f: i % k for i, f in enumerate(fams)}
|
||
ind_bin = {ind: fam_bin[fam_of[ind]] for ind in individuals}
|
||
|
||
def _subset(bin_idx: int) -> tuple[dict, dict, set[str], set[str]]:
|
||
test_inds = {ind for ind, b in ind_bin.items() if b == bin_idx}
|
||
train_inds = set(individuals) - test_inds
|
||
train_p, test_p = {}, {}
|
||
for key, ind in ind_of.items():
|
||
if ind in train_inds:
|
||
train_p[key] = phenotypes[key]
|
||
else:
|
||
test_p[key] = phenotypes[key]
|
||
return train_p, test_p, train_inds, test_inds
|
||
|
||
folds: list[dict] = []
|
||
all_x: list[float] = []
|
||
all_y: list[float] = []
|
||
for fold in range(k):
|
||
train_p, test_p, train_inds, test_inds = _subset(fold)
|
||
if not test_inds or len(train_inds) < 2:
|
||
folds.append({
|
||
"fold": fold + 1, "n_train": len(train_p), "n_test": len(test_inds),
|
||
"n_eval": 0, "pearson": None, "rmse": None,
|
||
"error": "训练集个体不足 2 或测试集为空",
|
||
})
|
||
continue
|
||
try:
|
||
res = solve(pedigree, train_p, tol=tol, fixed=fixed, covariate=covariate,
|
||
gxe=gxe, record_map=record_map)
|
||
except Exception as e: # noqa: BLE001
|
||
folds.append({
|
||
"fold": fold + 1, "n_train": len(train_p), "n_test": len(test_inds),
|
||
"n_eval": 0, "pearson": None, "rmse": None,
|
||
"error": f"求解失败: {e!s}",
|
||
})
|
||
continue
|
||
ebv = res.get("ebv", {})
|
||
pred: dict[str, list[float]] = {}
|
||
obs: dict[str, list[float]] = {}
|
||
for key, v in test_p.items():
|
||
ind = ind_of[key]
|
||
pv = ebv.get(ind)
|
||
if pv is None:
|
||
continue
|
||
pred.setdefault(ind, []).append(float(pv))
|
||
obs.setdefault(ind, []).append(float(v))
|
||
xs = [sum(pred[i]) / len(pred[i]) for i in pred]
|
||
ys = [sum(obs[i]) / len(obs[i]) for i in obs]
|
||
pearson = _pearson(xs, ys) if len(xs) >= 2 else 0.0
|
||
rmse = (math.sqrt(sum((a - b) ** 2 for a, b in zip(xs, ys)) / len(xs))
|
||
if xs else None)
|
||
if len(xs) >= 2:
|
||
all_x.extend(xs)
|
||
all_y.extend(ys)
|
||
folds.append({
|
||
"fold": fold + 1, "n_train": len(train_p), "n_test": len(test_inds),
|
||
"n_eval": len(xs), "pearson": round(pearson, 4) if xs else None,
|
||
"rmse": round(rmse, 6) if rmse is not None else None,
|
||
"h2": float(res["h2"]) if res.get("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(_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 = f"{len(errors)} 折失败: " + "；".join(errors) if errors else None
|
||
return {
|
||
"k": k, "n_total": n_total, "n_individuals": n_individuals,
|
||
"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, "warning": warning,
|
||
}
|