"""纯 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, # 消除 3000N_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: 序号}。""" 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, }