| 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) |
|
|
