122 lines
5.5 KiB
Python
122 lines
5.5 KiB
Python
"""断链焊接第二段 e2e：配合力纯Python化 + 选择指数(EBV/h²落库+决选) + 亲缘近交。
|
|
|
|
用法（需在 backend 目录下，PYTHONPATH=.）：
|
|
ENVIRONMENT=dev C:/ai/miniconda3/envs/dpb/python.exe -X utf8 -m tests._e2e_weld2
|
|
|
|
验证通过后删除脚本产生的数据（ABLUP/配合力/选择指数/决选 记录）。
|
|
"""
|
|
import asyncio
|
|
import os
|
|
|
|
os.environ.setdefault("ENVIRONMENT", "dev")
|
|
|
|
from sqlalchemy import delete, select # noqa: E402
|
|
|
|
from app.api.v1.module_bre.statistics.service import StatisticsService # noqa: E402
|
|
from app.core.base_schema import AuthSchema, CoreUserSchema # noqa: E402
|
|
from app.core.database import create_async_engine_and_session # noqa: E402
|
|
|
|
from app.api.v1.module_system.user.model import UserModel # noqa: E402 (注册 mapper，避免关系解析失败)
|
|
from app.api.v1.module_bre.selection_result.model import SelectionResultModel # noqa: E402
|
|
from app.api.v1.module_bre.statistics.model import ( # noqa: E402
|
|
CombiningAbilityModel,
|
|
PredictionModel,
|
|
PredictionValueModel,
|
|
SelectionIndexModel,
|
|
StatisticsJobModel,
|
|
)
|
|
|
|
AUTH = AuthSchema(user=CoreUserSchema(id=1, is_superuser=True))
|
|
TRAIT_ID = 29
|
|
TRAIT_CODE = "avg_fruit_weight"
|
|
WEIGHTS = {"avg_fruit_weight": 1.0, "max_fruit_weight": 0.8}
|
|
|
|
created = {"pred": None, "ca": None, "si": None, "job_ids": []}
|
|
|
|
|
|
async def main() -> None:
|
|
_engine, session_factory = create_async_engine_and_session()
|
|
try:
|
|
await _run(session_factory)
|
|
print("E2E 全部通过")
|
|
finally:
|
|
await _engine.dispose()
|
|
await cleanup(session_factory)
|
|
|
|
|
|
async def _run(session_factory) -> None:
|
|
async with session_factory() as db:
|
|
svc = StatisticsService(AUTH, db)
|
|
|
|
# 1. ABLUP（含 h²/可靠性，纯 Python 引擎）
|
|
pred_id = await svc.run_ablup(TRAIT_ID, TRAIT_CODE)
|
|
created["pred"] = pred_id
|
|
job = await db.execute(
|
|
select(StatisticsJobModel).where(StatisticsJobModel.result_ref == pred_id, StatisticsJobModel.job_type == "ABLUP")
|
|
)
|
|
created["job_ids"] += [j.id for j in job.scalars().all()]
|
|
preds = await svc.list_predictions()
|
|
p = next(x for x in preds if x.id == pred_id)
|
|
assert p.heritability is not None, "ABLUP 应写入 h²"
|
|
assert 0.0 <= float(p.heritability) <= 1.0, f"h² 越界: {p.heritability}"
|
|
ranking = await svc.ebv_ranking(pred_id)
|
|
assert ranking, "EBV 排行应有数据"
|
|
rels = [float(r.reliability) for r in ranking if r.reliability is not None]
|
|
assert rels and all(0.0 <= r <= 1.0 for r in rels), f"可靠性越界: {rels}"
|
|
print(f"[1] ABLUP ok pred_id={pred_id} h2={p.heritability} acc={p.accuracy} n={p.train_n} ebv_rows={len(ranking)}")
|
|
|
|
# 2. 配合力 GCA/SCA（纯 Python 引擎）
|
|
ca_id = await svc.run_combining(TRAIT_ID, TRAIT_CODE)
|
|
created["ca"] = ca_id
|
|
job = await db.execute(
|
|
select(StatisticsJobModel).where(StatisticsJobModel.result_ref == ca_id, StatisticsJobModel.job_type == "COMBINING")
|
|
)
|
|
created["job_ids"] += [j.id for j in job.scalars().all()]
|
|
ca = await svc.combining_detail(ca_id)
|
|
assert ca.gca_json and ca.sca_json, "GCA/SCA 应有结果"
|
|
print(f"[2] COMBINING ok ca_id={ca_id} parents={len(ca.gca_json)} combos={len(ca.sca_json)}")
|
|
|
|
# 3. 选择指数（EBV 接入 + h² 加权 + 落库）
|
|
si = await svc.selection_index(weights=WEIGHTS, batch_ids={TRAIT_CODE: pred_id}, use_h2=True, top_n=5)
|
|
created["si"] = si["id"]
|
|
assert si["top"], "选择指数应有排名"
|
|
assert si["heritabilities"][TRAIT_CODE] is not None, "应带入批次 h²"
|
|
top = await svc.list_index_batches()
|
|
assert any(x.id == si["id"] for x in top), "选择指数批次应落库可查"
|
|
print(f"[3] SELECTION-INDEX ok si_id={si['id']} method={si['method']} h2={si['heritabilities']} top={len(si['top'])}")
|
|
|
|
# 4. 写入决选
|
|
applied = await svc.apply_selection(si["id"], top_n=3)
|
|
assert applied["written"] >= 1, "决选应写入"
|
|
print(f"[4] APPLY ok written={applied['written']} year={applied['selection_year']}")
|
|
|
|
# 5. 亲缘/近交
|
|
kin = await svc.kinship_matrix(threshold=0.25)
|
|
assert kin["n"] > 0 and kin["individuals"], "亲缘分析应有结果"
|
|
print(f"[5] KINSHIP ok n={kin['n']} flags={len(kin['flags'])} F>0={len(kin['inbreeding'])} matrix={'yes' if kin['matrix'] else 'no'}")
|
|
if kin["flags"]:
|
|
print(" top flags:", kin["flags"][:3])
|
|
|
|
|
|
async def cleanup(session_factory) -> None:
|
|
async with session_factory() as db:
|
|
si_id = created["si"]
|
|
if si_id:
|
|
await db.execute(
|
|
delete(SelectionResultModel).where(SelectionResultModel.reason.like(f"选择指数批次#{si_id}%"))
|
|
)
|
|
await db.execute(delete(SelectionIndexModel).where(SelectionIndexModel.id == si_id))
|
|
if created["ca"]:
|
|
await db.execute(delete(CombiningAbilityModel).where(CombiningAbilityModel.id == created["ca"]))
|
|
if created["pred"]:
|
|
await db.execute(delete(PredictionValueModel).where(PredictionValueModel.prediction_id == created["pred"]))
|
|
await db.execute(delete(PredictionModel).where(PredictionModel.id == created["pred"]))
|
|
if created["job_ids"]:
|
|
await db.execute(delete(StatisticsJobModel).where(StatisticsJobModel.id.in_(created["job_ids"])))
|
|
await db.commit()
|
|
print("已清理 e2e 产生的数据")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|