File size: 3,282 Bytes
017b9f5 9401a93 017b9f5 9401a93 017b9f5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | import argparse
import json
import numpy as np
import pandas as pd
from inputs import prepare_inputs
parser = argparse.ArgumentParser()
parser.add_argument('--t', type=float, default=0.05)
parser.add_argument('--frac', type=float, default=0.1)
parser.add_argument('--seed', type=int, default=0)
parser.add_argument('--output', default='hard_split.json')
args = parser.parse_args()
K, MAXDF = 5, 2000
LUT = np.full(256, 20, dtype=np.int64)
for i, c in enumerate('ACDEFGHIKLMNPQRSTVWY'):
LUT[ord(c)] = i
def kmers(seq):
codes = LUT[np.frombuffer(seq.encode('ascii', 'ignore'), dtype=np.uint8)]
n = len(codes) - K + 1
if n <= 0:
return np.empty(0, dtype=np.int32)
ids = np.zeros(n, dtype=np.int64)
for j in range(K):
ids = ids * 21 + codes[j:j + n]
return np.unique(ids).astype(np.int32)
train = pd.read_parquet('../data/train.parquet').drop_duplicates(['sequence', 'ec_number'])
seqs = train.groupby('sequence', sort=False).size().index.tolist()
N = len(seqs)
parts = [kmers(s) for s in seqs]
sizes = np.array([len(p) for p in parts])
kmer_all = np.concatenate(parts)
del parts
seq_all = np.repeat(np.arange(N, dtype=np.int32), sizes)
order = np.argsort(kmer_all, kind='stable')
kmer_sorted = kmer_all[order]
seq_sorted = seq_all[order]
del kmer_all, seq_all, order
print('index built for', N, 'sequences')
def shared_counts(seq):
q = kmers(seq)
if len(q) == 0:
return np.zeros(N, dtype=np.int64), 1
lo = np.searchsorted(kmer_sorted, q, 'left')
cnt = np.searchsorted(kmer_sorted, q, 'right') - lo
keep = cnt <= MAXDF
lo, cnt = lo[keep], cnt[keep]
offs = np.cumsum(cnt) - cnt
idx = np.arange(cnt.sum()) - np.repeat(offs, cnt) + np.repeat(lo, cnt)
return np.bincount(seq_sorted[idx], minlength=N), len(q)
def nearest(seq, blocked=None):
shared, nq = shared_counts(seq)
if blocked is not None:
shared[blocked] = 0
return shared.max() / nq
rng = np.random.RandomState(args.seed)
target = int(args.frac * N)
in_hold = np.zeros(N, dtype=bool)
for s in rng.permutation(N):
if in_hold[s]:
continue
shared, nq = shared_counts(seqs[s])
group = (shared >= args.t * nq) & ~in_hold
group[s] = True
if in_hold.sum() + group.sum() > 1.2 * target:
continue
in_hold |= group
if in_hold.sum() >= target:
break
print('holdout sequences:', int(in_hold.sum()), 'of', N)
val_seqs = prepare_inputs(pd.read_parquet('../data/val_inputs.parquet')).sequence.tolist()
hold_ids = np.where(in_hold)[0]
sample = rng.choice(hold_ids, min(3000, len(hold_ids)), replace=False)
random_ids = rng.choice(N, 3000, replace=False)
rows = []
for name, sims in (
('validation to full train', [nearest(s) for s in val_seqs]),
('new holdout to remaining train', [nearest(seqs[i], in_hold) for i in sample]),
('random holdout to remaining train', [nearest(seqs[i], np.arange(N) == i) for i in random_ids])):
p = np.percentile(sims, [10, 25, 50, 75, 90])
rows.append({'set': name, 'p10': p[0], 'p25': p[1], 'p50': p[2], 'p75': p[3], 'p90': p[4]})
print(pd.DataFrame(rows).round(3).to_string(index=False))
json.dump({'t': args.t, 'holdout': [int(i) for i in hold_ids]}, open(args.output, 'w'))
print('saved', args.output)
|