1528 lines
60 KiB
Python
1528 lines
60 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 矩阵（稠密，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。
|
||
|
||
无系谱（所有个体均无父母）时：不背书遗传力——
|
||
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
|
||
|
||
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 第二随机效应））
|
||
|
||
|
||
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]:
|
||
"""返回 (A, ii, jj, vv)：A 稠密（供 REML 似然与近交系数）；A⁻¹ 以 COO 稀疏三元组存储。
|
||
|
||
Henderson 规则与 _build_ainv 逐项一致，但 A⁻¹ 不落稠密矩阵，只收集非零元
|
||
（对角 + 2 亲本 × (对角/双向共祖先)），供共轭梯度 matvec 使用：
|
||
(Ainv·u)[i] = Σ_j Ainv_ij·u_j，np.add.at 按三元组累加即可。
|
||
"""
|
||
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
|
||
|
||
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 = 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)
|
||
_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 (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（剖面求值次数/EM 迭代数）；warning: str | None；
|
||
solver: str（"dense" 稠密精确 / "sparse-cg" 稀疏共轭梯度，n>N_SPARSE 时触发）；
|
||
稀疏路径另含 cg_iter（CG 迭代数）与 reliability_approx=True（可靠性为 Hutchinson 近似，
|
||
EBV/h²/σ² 仍为精确解）。
|
||
"""
|
||
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 仅为表型残差（固定效应已校正），未估计遗传力与可靠性。",
|
||
}
|
||
|
||
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)
|
||
|
||
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)
|
||
|
||
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
|
||
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,
|
||
}
|