59 lines
3.2 KiB
Python
59 lines
3.2 KiB
Python
"""对拍：python -m ml.eval.check_parity [--kin ml/assets/kin.json] [--data ml/data/parity.json]
|
||
检查 PyTorch 版 fk / orient / 组对间隙 / 83 维特征 与 JS 一致，以及间隙梯度与有限差分一致。"""
|
||
import argparse, json
|
||
import torch
|
||
from ml.common.kin_torch import Kin
|
||
from ml.common.features import features, GAP_CAP
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument('--kin', default=None); ap.add_argument('--data', default='ml/data/parity.json')
|
||
a = ap.parse_args()
|
||
kin = Kin(a.kin) if a.kin else Kin(dtype=torch.float64)
|
||
D = json.load(open(a.data, encoding='utf-8'))
|
||
assert D['fp'] == kin.fp, f"指纹不符：parity {D['fp']} vs kin {kin.fp}（重跑 export_kin.cjs）"
|
||
R = D['recs']; f64 = torch.float64
|
||
q = torch.tensor([r['q'] for r in R], dtype=f64)
|
||
qp = torch.tensor([r['qPrev'] for r in R], dtype=f64)
|
||
T = torch.tensor([r['T'] for r in R], dtype=f64)
|
||
Rt = torch.tensor([r['R'] if r['R'] else [0.] * 9 for r in R], dtype=f64)
|
||
pre = torch.zeros(len(R), 3, 6, dtype=f64); pl = torch.zeros(len(R), dtype=torch.long)
|
||
for n, r in enumerate(R):
|
||
for k, d in enumerate(r['prefix']): pre[n, k] = torch.tensor(d, dtype=f64)
|
||
pl[n] = len(r['prefix'])
|
||
sub = torch.tensor([r['sub'] if r['sub'] else [0.] * 6 for r in R], dtype=f64)
|
||
sf = torch.tensor([1. if r['sub'] else 0. for r in R], dtype=f64)
|
||
kin = Kin(a.kin, dtype=f64) if a.kin else kin
|
||
xf = kin.cum_xf(q)
|
||
e_fk = (kin.fk(q, xf) - torch.tensor([r['fk'] for r in R], dtype=f64)).abs().max().item()
|
||
e_or = (kin.orient(q, xf).reshape(-1, 9) - torch.tensor([r['orient'] for r in R], dtype=f64)).abs().max().item()
|
||
G = kin.gaps(q, GAP_CAP, xf)
|
||
gm_js = torch.tensor([r['gm'] for r in R], dtype=f64)
|
||
e_gm = (G['gm'] - gm_js.clamp(max=GAP_CAP)).abs().max().item()
|
||
col_js = gm_js.min(-1).values <= 0; col_py = G['gm'].min(-1).values <= 0
|
||
f = features(kin, q, qp, T, Rt, pre, pl, sub, sf, gm=G['gm'], xf=xf)
|
||
e_f = (f - torch.tensor([r['feat'] for r in R], dtype=f64)).abs().max().item()
|
||
# 梯度：最小间隙对 q 的梯度 vs 有限差分（取 50 个无干涉且间隙 < 20 的样本）
|
||
sel = ((G['gm'].min(-1).values > 0.5) & (G['gm'].min(-1).values < 20)).nonzero().flatten()[:50]
|
||
rel = []
|
||
if len(sel):
|
||
qs = q[sel].clone().requires_grad_(True)
|
||
g = kin.gaps(qs, GAP_CAP)['gm'].min(-1).values
|
||
gr, = torch.autograd.grad(g.sum(), qs)
|
||
h = 1e-4
|
||
for j in range(6):
|
||
dq = torch.zeros_like(qs); dq[:, j] = h
|
||
fd = (kin.gaps(qs.detach() + dq, GAP_CAP)['gm'].min(-1).values - kin.gaps(qs.detach() - dq, GAP_CAP)['gm'].min(-1).values) / (2 * h)
|
||
rel.append(((gr[:, j] - fd).abs() / (fd.abs() + 1e-3)).max().item())
|
||
res = {'n': len(R), 'fk_max_abs_mm': e_fk, 'orient_max_abs': e_or, 'gap_max_abs_mm': e_gm,
|
||
'collide_bool_mismatch': int((col_js != col_py).sum()), 'feat_max_abs': e_f,
|
||
'grad_rel_err_max': max(rel) if rel else None}
|
||
ok = e_fk < 1e-6 and e_or < 1e-9 and e_gm < 1e-4 and res['collide_bool_mismatch'] == 0 and e_f < 1e-4 and (not rel or max(rel) < 1e-3)
|
||
res['PASS'] = bool(ok)
|
||
print(json.dumps(res, indent=1))
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|