| """Sparse k-mer index: TF-IDF-weighted, L2-normalised presence vectors for cosine search.""" |
| import numpy as np |
| import scipy.sparse as sp |
|
|
| AA = 'ACDEFGHIKLMNPQRSTVWY' |
| LUT = np.full(256, 20, dtype=np.int64) |
| for i, c in enumerate(AA): |
| LUT[ord(c)] = i |
| LUT[ord(c.lower())] = i |
|
|
|
|
| def kmer_ids(seq, k): |
| a = LUT[np.frombuffer(seq.encode('ascii', 'replace'), dtype=np.uint8)] |
| if len(a) < k: |
| return np.empty(0, dtype=np.int64) |
| ids = np.zeros(len(a) - k + 1, dtype=np.int64) |
| for j in range(k): |
| ids = ids * 21 + a[j:len(a) - k + 1 + j] |
| return np.unique(ids) |
|
|
|
|
| def presence_matrix(seqs, k): |
| rows, cols = [], [] |
| for i, s in enumerate(seqs): |
| ids = kmer_ids(s, k) |
| rows.append(np.full(len(ids), i, dtype=np.int32)) |
| cols.append(ids) |
| rows = np.concatenate(rows) |
| cols = np.concatenate(cols) |
| data = np.ones(len(rows), dtype=np.float32) |
| return sp.csr_matrix((data, (rows, cols)), shape=(len(seqs), 21 ** k)) |
|
|
|
|
| def weight(m, idf): |
| m = m.multiply(idf[None, :]).tocsr().astype(np.float32) |
| norms = np.sqrt(np.asarray(m.multiply(m).sum(axis=1)).ravel()) |
| norms[norms == 0] = 1 |
| return sp.diags(1 / norms).dot(m).tocsr().astype(np.float32) |
|
|
|
|
| class KmerIndex: |
| def __init__(self, seqs, k=5): |
| self.k = k |
| m = presence_matrix(seqs, k) |
| df = np.bincount(m.indices, minlength=m.shape[1]).astype(np.float32) |
| self.idf = np.log((1 + m.shape[0]) / (1 + df)).astype(np.float32) |
| self.T = weight(m, self.idf).T.tocsr() |
|
|
| def topk(self, seqs, k=50, chunk=256): |
| """Return (indices, scores) arrays of shape (len(seqs), k).""" |
| n = len(seqs) |
| idx = np.full((n, k), -1, dtype=np.int64) |
| sc = np.zeros((n, k), dtype=np.float32) |
| for start in range(0, n, chunk): |
| q = weight(presence_matrix(seqs[start:start + chunk], self.k), self.idf) |
| s = (q @ self.T).tocsr() |
| for r in range(s.shape[0]): |
| lo, hi = s.indptr[r], s.indptr[r + 1] |
| if hi == lo: |
| continue |
| d, c = s.data[lo:hi], s.indices[lo:hi] |
| m = min(k, hi - lo) |
| part = np.argpartition(-d, m - 1)[:m] |
| part = part[np.argsort(-d[part])] |
| idx[start + r, :m] = c[part] |
| sc[start + r, :m] = d[part] |
| return idx, sc |
|
|
