train: loop variable typing (mypy)

This commit is contained in:
kemal committed 2026-09-29 19:54:24 +01:00
1 parent 7cfb0fd962
commit adad0efb08
1 file changed
+4 -4
+4 -4
View File
@@ -106,10 +106,10 @@ def fit_block(theta: F32, r: F32, use: npt.NDArray[np.bool_], haploid: npt.NDArr
# a locus with a single cluster is homozygous: an AB-labelled lone cluster moves to AA or BB by position
lone_ab = keep[:, 1] & ~keep[:, 0] & ~keep[:, 2]
if lone_ab.any():
for i in np.flatnonzero(lone_ab).tolist():
hom = 0 if m[i, 1] < 0.5 else 2
assign[i, assign[i] == 1] = hom
keep[i, 1], keep[i, hom] = False, True
for j in np.flatnonzero(lone_ab).tolist():
hom = 0 if m[j, 1] < 0.5 else 2
assign[j, assign[j] == 1] = hom
keep[j, 1], keep[j, hom] = False, True
counts = np.stack([(assign == k).sum(axis=1) for k in range(3)], axis=1)
tm = np.full((L, 3), np.nan)
td = np.full((L, 3), np.nan)