checkpoint: real-scale 数据链 + 统计引擎 + ABLUP 稀疏性能修复
init 后首次落盘，累计工作： - 数据侧：simulate_breeding_data real-scale 3.5万树重建 + 分层观测 - 引擎侧：ABLUP 稀疏生产档 EM-REML（blup 1.6.0/1.7.0）；性能倒挂修复—— 系谱闭包收口为 BFS 可达祖先 + N_EXACT 3000 对齐 N_SUBSAMPLE（选中 200~3000 树 不再顶进数小时精确迹尾），A/B 数值等价实证 + 守卫探针 - 前端：观测过滤 + 表单/HTTP 工具链完善 - 文档：业务链/观测梳理 + 引擎实施记录 + 规划对照总表 - gitignore：排除 Temp/调试脚本/一次性验证输出
This commit is contained in:
34047007@qq.com
@@ -9,15 +9,21 @@
|
||||
- 无父母的个体一律视为 founder（base population，彼此不相关）。
|
||||
|
||||
算法：
|
||||
1. A 矩阵（稠密，O(n²) 递归）→ 各体自交系数 F；
|
||||
2. A⁻¹ 用 Henderson 稀疏规则构造（无需求 A 的逆）；
|
||||
3. MME 求解：y = Xb + Zu + e，X=截距+固定效应虚拟列+协变量（默认仅截距），λ = σ²e/σ²a；
|
||||
默认稠密精确求解；个体数 > N_SPARSE（默认 1000）且系谱存在时切稀疏 A⁻¹ + 共轭梯度
|
||||
迭代（solve 的 solver="sparse-cg"）：EBV/h²/σ² 为 CG 精确解，reliability 用
|
||||
Hutchinson 随机探测估计 C22 对角线（reliability_approx=True，均值/PA 可靠、个体级近似）；
|
||||
4. 剖面 REML：对每个候选 h² 固定 λ，用 MME 解的 EM 条件式估计 σ²e、σ²a，
|
||||
再用精确 REML 对数似然评分，黄金分割一维最大化（稳健、可收敛到边界）；
|
||||
5. h² = σ²a/(σ²a+σ²e)；reliability = 1 - PEV/σ²a，PEV = C22[ii]·σ²e。
|
||||
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。
|
||||
@@ -41,8 +47,27 @@ CG_TOL = 1e-9 # 共轭梯度相对残差收敛阈值
|
||||
# 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.5.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 第二随机效应））
|
||||
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:
|
||||
@@ -129,30 +154,37 @@ def _build_ainv(order: list[int], n_base: int,
|
||||
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]:
|
||||
"""返回 (A, ii, jj, vv)：A 稠密（供 REML 似然与近交系数）；A⁻¹ 以 COO 稀疏三元组存储。
|
||||
"""返回 (diag_a, ii, jj, vv)：A 对角（近交 F=A_ii−1 来源）；A⁻¹ 以 COO 稀疏三元组存储。
|
||||
|
||||
Henderson 规则与 _build_ainv 逐项一致，但 A⁻¹ 不落稠密矩阵，只收集非零元
|
||||
（对角 + 2 亲本 × (对角/双向共祖先)），供共轭梯度 matvec 使用：
|
||||
(Ainv·u)[i] = Σ_j Ainv_ij·u_j，np.add.at 按三元组累加即可。
|
||||
不构建稠密 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)}
|
||||
A = np.eye(n)
|
||||
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 = 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
|
||||
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] = []
|
||||
@@ -171,8 +203,8 @@ def _build_ainv_sparse(order: list[int], n_base: int,
|
||||
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
|
||||
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)
|
||||
@@ -185,7 +217,7 @@ def _build_ainv_sparse(order: list[int], n_base: int,
|
||||
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 (A, np.array(ii, dtype=np.int64), np.array(jj, dtype=np.int64),
|
||||
return (diag_a, np.array(ii, dtype=np.int64), np.array(jj, dtype=np.int64),
|
||||
np.array(vv, dtype=float))
|
||||
|
||||
|
||||
@@ -351,11 +383,13 @@ def solve(pedigree: list[dict], phenotypes: dict, *, tol: float = TOL,
|
||||
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（剖面求值次数/EM 迭代数）；warning: str | None；
|
||||
solver: str（"dense" 稠密精确 / "sparse-cg" 稀疏共轭梯度，n>N_SPARSE 时触发）；
|
||||
稀疏路径另含 cg_iter（CG 迭代数）与 reliability_approx=True（可靠性为 Hutchinson 近似，
|
||||
EBV/h²/σ² 仍为精确解）。
|
||||
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,
|
||||
@@ -427,22 +461,237 @@ def solve(pedigree: list[dict], phenotypes: dict, *, tol: float = TOL,
|
||||
"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
|
||||
XtX = X.T @ X
|
||||
Xty = X.T @ y
|
||||
Zty = Z.T @ y
|
||||
ZtZ = Z.T @ Z
|
||||
XtZ = X.T @ Z
|
||||
|
||||
# A7 阈值分派：个体数 > N_SPARSE 且系谱存在 → 稀疏 A⁻¹ + 共轭梯度迭代求解
|
||||
# （回避 O(n³) 稠密 MME 逆）；可靠性用 Hutchinson 随机探测估计 C22 对角线。
|
||||
use_sparse = n_with_parents > 0 and n > N_SPARSE
|
||||
if use_sparse:
|
||||
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)
|
||||
ZtZ = np.diag(ZtZ_diag)
|
||||
|
||||
def _reml_ll(Va: float, Ve: float) -> float:
|
||||
"""精确 REML 对数似然（直接 V 计算，O(m³)）。"""
|
||||
@@ -456,85 +705,6 @@ def solve(pedigree: list[dict], phenotypes: dict, *, tol: float = TOL,
|
||||
_, lx = np.linalg.slogdet(XtVinvX)
|
||||
return -0.5 * (float(lv) + float(lx) + yPy)
|
||||
|
||||
if use_sparse:
|
||||
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
|
||||
np.add.at(out2, ainv_ii, lam * ainv_vv * u[ainv_jj])
|
||||
return np.concatenate([out1, out2])
|
||||
return mv
|
||||
|
||||
mme_rhs = np.concatenate([Xty, Zty])
|
||||
|
||||
best = {"h2": None, "ll": -np.inf, "Va": 0.0, "Ve": 0.0}
|
||||
|
||||
def _profile(h2: float) -> float:
|
||||
lam = (1.0 - h2) / h2
|
||||
sol, _it, _ok = _cg_solve(_matvec(lam), mme_rhs)
|
||||
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)
|
||||
_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
|
||||
mv = _matvec(lam)
|
||||
sol, cg_iter, cg_ok = _cg_solve(mv, mme_rhs)
|
||||
u = sol[p:]
|
||||
# Hutchinson：E[P⊙(C22·P)] = diag(C22)，C22·P 由 CG 解 big·x=[0;P] 取 x[p:] 得到
|
||||
rng = np.random.RandomState(20260804)
|
||||
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
|
||||
|
||||
ebv = {ind: float(u[idx[ind]]) for ind in order}
|
||||
pev = 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 未完全收敛（似然面极平/边界最优），结果已采用已探明最优。"
|
||||
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),
|
||||
"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 and cg_ok,
|
||||
"n_iter": n_iter,
|
||||
"solver": "sparse-cg",
|
||||
"cg_iter": cg_iter,
|
||||
"reliability_approx": True,
|
||||
"warning": warning,
|
||||
}
|
||||
|
||||
def _mme_solve(lam: float) -> tuple[np.ndarray, np.ndarray]:
|
||||
big = np.zeros((p + n, p + n))
|
||||
big[:p, :p] = XtX
|
||||
|
||||
Reference in New Issue
Block a user