mirror of
https://github.com/metabarcoding/obitools4.git
synced 2026-11-09 09:24:41 +00:00
Implement k-mer counting and sequence handling
Introduces the SeqKmerCounter structure for parallel k-mer counting, refactors sequence reading and writing logic, and adds comprehensive test coverage for various sequence scenarios.
This commit is contained in:
1 parent
e511650f5b
commit
3f0ca993d2
6 files changed
+684
-32
No files matched your search
+33
-3
@@ -168,9 +168,7 @@ Validation :
|
|||||||
- Les 13 Tm imprimés par ecoPrimers sont retrouvés à +/- 0,05 °C près.
|
- Les 13 Tm imprimés par ecoPrimers sont retrouvés à +/- 0,05 °C près.
|
||||||
|
|
||||||
Écarts volontaires avec ecoPrimers (à confirmer) :
|
Écarts volontaires avec ecoPrimers (à confirmer) :
|
||||||
1. **273,15 au lieu de 273,0** : ecoPrimers calcule `Tm(K) - 273.0`. Les Tm de `obithermo` sont donc
|
1. **273,15 au lieu de 273,0** : ecoPrimers calcule `Tm(K) - 273.0`, `obithermo` utilise la valeur exacte 273,15 (Tm 0,15 °C plus bas ; le test tient compte du décalage). **Validé par l'utilisateur le 2026-11-03 : on garde 273,15.**
|
||||||
0,15 °C plus bas que ceux d'ecoPrimers (le test tient compte du décalage). Une constante,
|
|
||||||
facile à inverser si la compatibilité exacte des colonnes Tm prime.
|
|
||||||
2. **Owczarzy corrigé** : dans ecoPrimers, `gcContent` n'est jamais affecté (grep sur tout `src/`),
|
2. **Owczarzy corrigé** : dans ecoPrimers, `gcContent` n'est jamais affecté (grep sur tout `src/`),
|
||||||
donc `-m 2` lit de la mémoire non initialisée (`CNNParams` est sur la pile dans `main`).
|
donc `-m 2` lit de la mémoire non initialisée (`CNNParams` est sur la pile dans `main`).
|
||||||
`obithermo` utilise la fraction GC de l'amorce. Non comparable au C ; testé par invariants
|
`obithermo` utilise la fraction GC de l'amorce. Non comparable au C ; testé par invariants
|
||||||
@@ -181,3 +179,35 @@ Validation :
|
|||||||
|
|
||||||
Pour plus tard : les k-mers d'obikmer sont codés A,C,G,T = 0..3, comme les tables de `obithermo` ;
|
Pour plus tard : les k-mers d'obikmer sont codés A,C,G,T = 0..3, comme les tables de `obithermo` ;
|
||||||
une variante `TmCode(uint64, k)` évitera le décodage quand le cœur appellera Tm par paire.
|
une variante `TmCode(uint64, k)` évitera le décodage quand le cœur appellera Tm par paire.
|
||||||
|
|
||||||
|
## 9. Mode A, étape 2 : `obikmer.SeqKmerCounter` (2026-11-03)
|
||||||
|
|
||||||
|
**Constat qui change le plan** : `KmerSetGroupBuilder` crée n x P écrivains `.skm` simultanés (un set par
|
||||||
|
séquence, P >= 64). Pour 3 909 génomes cela fait > 250 000 fichiers ouverts : `KmerSetGroup` +
|
||||||
|
`QuorumAtLeast` ne convient pas pour « dans combien de séquences ce k-mer apparaît ».
|
||||||
|
|
||||||
|
**Solution** (enrichissement d'obikmer, pas de code local à obiprimers) : `pkg/obikmer/seqquorum.go` et
|
||||||
|
`seqquorum_count.go`.
|
||||||
|
- Phase 1, `AddSequence(id, seq, circular)` (concurrent) : super-kmers groupés par minimiseur, écrits avec
|
||||||
|
l'id de séquence dans **un seul** jeu de P fichiers (format `[uvarint id][uvarint len][2 bits/base]`).
|
||||||
|
Les bases non ACGTU découpent la séquence en segments : aucun k-mer ne contient d'ambiguïté (comme
|
||||||
|
le masque d'erreur d'ecoPrimers). Génome circulaire : on ajoute les k-1 premières bases.
|
||||||
|
- Phase 2, `CountAtLeast(q)` : partitions traitées en parallèle (`obidefault.ParallelWorkers()`),
|
||||||
|
paires (k-mer canonique, id) triées, séquences distinctes comptées. Résultat en mémoire, trié
|
||||||
|
(il est filtré par le quorum : 255 mots pour Teleostei, 1 208 pour Cyprinidae chez ecoPrimers).
|
||||||
|
- `SeqKmerCount.Repeated` reproduit le drapeau MULTIWORD d'ecoPrimers (k-mer présent plusieurs fois
|
||||||
|
dans une même séquence, les deux brins confondus), utile pour `-U` (pas de multi-match).
|
||||||
|
- Mémoire de la phase 2 : ~16 octets par occurrence de k-mer des partitions en cours de traitement ;
|
||||||
|
P doit croître avec la taille totale des données.
|
||||||
|
- Refactor neutre dans `skm_writer.go` / `skm_reader.go` : `appendPacked` et `unpackBases` extraits et
|
||||||
|
partagés avec le nouveau format (tests skm inchangés et verts).
|
||||||
|
|
||||||
|
Tests (`seqquorum_test.go`) : oracle force brute indépendant (`EncodeCanonicalKmer` sur chaque fenêtre),
|
||||||
|
6 configurations (k = 4, 7, 18, 31 ; linéaire/circulaire ; avec N/R ; m et P variés ; insertion
|
||||||
|
concurrente ; motifs partagés sur les deux brins), plus cas des deux brins dans une même séquence,
|
||||||
|
ambiguïtés + jonction circulaire, paramètres invalides. `go test ./pkg/obikmer/ -race` vert.
|
||||||
|
Un bug trouvé par ces tests dans ma boucle de réduction (compteur de séquences) a été corrigé.
|
||||||
|
|
||||||
|
Contrôle de bout en bout prévu (étape suivante) : `Strict primer count` d'ecoPrimers = **1 208**
|
||||||
|
(Cyprinidae) et **255** (Teleostei) avec `-e 3 -3 2 -l 30 -L 150 -c`, k = 18, quorum 0,7 ; le calcul
|
||||||
|
d'ecoPrimers est `floor(|E| x 0,7)` séquences (> longueur d'amorce).
|
||||||
@@ -0,0 +1,190 @@
|
|||||||
|
package obikmer
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bufio"
|
||||||
|
"encoding/binary"
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"sync"
|
||||||
|
)
|
||||||
|
|
||||||
|
// SeqKmerCounter counts, for every canonical k-mer, in how many distinct input
|
||||||
|
// sequences it occurs. It answers "which k-mers are present in at least q of
|
||||||
|
// the N sequences" without one k-mer set per sequence: a KmerSetGroup needs
|
||||||
|
// N*P partition files, which is out of reach for thousands of genomes, whereas
|
||||||
|
// a SeqKmerCounter uses a single set of P files whatever N is.
|
||||||
|
//
|
||||||
|
// Construction is two-phase, like KmerSetGroupBuilder:
|
||||||
|
//
|
||||||
|
// 1. AddSequence extracts super-kmers (grouped by minimizer, so a k-mer and its
|
||||||
|
// reverse complement always land in the same partition) and appends them,
|
||||||
|
// tagged with the sequence id, to one of P temporary files. It is safe to
|
||||||
|
// call concurrently.
|
||||||
|
// 2. CountAtLeast processes the partitions in parallel: each one is loaded,
|
||||||
|
// expanded to (canonical k-mer, sequence id) pairs, sorted, and reduced.
|
||||||
|
//
|
||||||
|
// Memory during phase 2 is about 16 bytes per k-mer occurrence of the largest
|
||||||
|
// partitions being processed at the same time, so P must grow with the total
|
||||||
|
// input size (total bases / P is the occurrence count of an average partition).
|
||||||
|
//
|
||||||
|
// Bases other than A, C, G, T, U break a sequence into independent runs: no
|
||||||
|
// k-mer ever contains an ambiguous base.
|
||||||
|
type SeqKmerCounter struct {
|
||||||
|
dir string
|
||||||
|
k, m, p int
|
||||||
|
|
||||||
|
parts []*seqPartWriter
|
||||||
|
|
||||||
|
mu sync.Mutex
|
||||||
|
firstErr error
|
||||||
|
sealed bool
|
||||||
|
}
|
||||||
|
|
||||||
|
type seqPartWriter struct {
|
||||||
|
mu sync.Mutex
|
||||||
|
w *bufio.Writer
|
||||||
|
file *os.File
|
||||||
|
path string
|
||||||
|
}
|
||||||
|
|
||||||
|
// SeqKmerCount is the result for one k-mer.
|
||||||
|
type SeqKmerCount struct {
|
||||||
|
Kmer uint64 // canonical k-mer
|
||||||
|
Sequences uint32 // number of distinct sequences containing it
|
||||||
|
Repeated bool // occurs more than once (either strand) in at least one sequence
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewSeqKmerCounter creates a counter whose temporary files live in a new
|
||||||
|
// directory created under parentDir (use "" for the system temporary
|
||||||
|
// directory). k is the k-mer size (2..31), m the minimizer size (<=0 for the
|
||||||
|
// default) and partitions the number of partition files P.
|
||||||
|
func NewSeqKmerCounter(parentDir string, k, m, partitions int) (*SeqKmerCounter, error) {
|
||||||
|
if k < 2 || k > 31 {
|
||||||
|
return nil, fmt.Errorf("obikmer: k must be between 2 and 31, got %d", k)
|
||||||
|
}
|
||||||
|
if partitions < 1 {
|
||||||
|
return nil, fmt.Errorf("obikmer: partitions must be >= 1, got %d", partitions)
|
||||||
|
}
|
||||||
|
if m <= 0 {
|
||||||
|
m = DefaultMinimizerSize(k)
|
||||||
|
}
|
||||||
|
if m >= k {
|
||||||
|
return nil, fmt.Errorf("obikmer: minimizer size %d must be smaller than k=%d", m, k)
|
||||||
|
}
|
||||||
|
|
||||||
|
dir, err := os.MkdirTemp(parentDir, "obikmer-seqcount-")
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
c := &SeqKmerCounter{dir: dir, k: k, m: m, p: partitions}
|
||||||
|
c.parts = make([]*seqPartWriter, partitions)
|
||||||
|
for i := range c.parts {
|
||||||
|
path := filepath.Join(dir, fmt.Sprintf("part_%04d.sqk", i))
|
||||||
|
f, err := os.Create(path)
|
||||||
|
if err != nil {
|
||||||
|
c.Close()
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
c.parts[i] = &seqPartWriter{w: bufio.NewWriterSize(f, 32*1024), file: f, path: path}
|
||||||
|
}
|
||||||
|
return c, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// K returns the k-mer size.
|
||||||
|
func (c *SeqKmerCounter) K() int { return c.k }
|
||||||
|
|
||||||
|
var unambiguousBase = func() (t [256]bool) {
|
||||||
|
for _, b := range "ACGTUacgtu" {
|
||||||
|
t[b] = true
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}()
|
||||||
|
|
||||||
|
// AddSequence registers sequence id. The same id must be used for every call
|
||||||
|
// concerning the same sequence; ids need not be contiguous. If circular is
|
||||||
|
// true the k-1 first bases are appended so that k-mers spanning the junction
|
||||||
|
// are counted. AddSequence is safe for concurrent use and must not be called
|
||||||
|
// after CountAtLeast.
|
||||||
|
func (c *SeqKmerCounter) AddSequence(id uint32, seq []byte, circular bool) {
|
||||||
|
if len(seq) < c.k {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if circular {
|
||||||
|
ext := make([]byte, 0, len(seq)+c.k-1)
|
||||||
|
ext = append(ext, seq...)
|
||||||
|
seq = append(ext, seq[:c.k-1]...)
|
||||||
|
}
|
||||||
|
|
||||||
|
var scratch []byte
|
||||||
|
start := -1
|
||||||
|
for i, b := range seq {
|
||||||
|
switch {
|
||||||
|
case unambiguousBase[b]:
|
||||||
|
if start < 0 {
|
||||||
|
start = i
|
||||||
|
}
|
||||||
|
case start >= 0:
|
||||||
|
scratch = c.addRun(id, seq[start:i], scratch)
|
||||||
|
start = -1
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if start >= 0 {
|
||||||
|
c.addRun(id, seq[start:], scratch)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// addRun writes the super-kmers of one ambiguity-free run.
|
||||||
|
func (c *SeqKmerCounter) addRun(id uint32, run []byte, scratch []byte) []byte {
|
||||||
|
if len(run) < c.k {
|
||||||
|
return scratch
|
||||||
|
}
|
||||||
|
for sk := range IterSuperKmers(run, c.k, c.m) {
|
||||||
|
scratch = binary.AppendUvarint(scratch[:0], uint64(id))
|
||||||
|
scratch = binary.AppendUvarint(scratch, uint64(len(sk.Sequence)))
|
||||||
|
scratch = appendPacked(scratch, sk.Sequence)
|
||||||
|
|
||||||
|
pw := c.parts[int(sk.Minimizer%uint64(c.p))]
|
||||||
|
pw.mu.Lock()
|
||||||
|
_, err := pw.w.Write(scratch)
|
||||||
|
pw.mu.Unlock()
|
||||||
|
if err != nil {
|
||||||
|
c.setErr(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return scratch
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *SeqKmerCounter) setErr(err error) {
|
||||||
|
c.mu.Lock()
|
||||||
|
if c.firstErr == nil {
|
||||||
|
c.firstErr = err
|
||||||
|
}
|
||||||
|
c.mu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
// seal flushes and closes the partition files; no more sequences can be added.
|
||||||
|
func (c *SeqKmerCounter) seal() error {
|
||||||
|
c.mu.Lock()
|
||||||
|
defer c.mu.Unlock()
|
||||||
|
if c.sealed {
|
||||||
|
return c.firstErr
|
||||||
|
}
|
||||||
|
c.sealed = true
|
||||||
|
for _, pw := range c.parts {
|
||||||
|
if err := pw.w.Flush(); err != nil && c.firstErr == nil {
|
||||||
|
c.firstErr = err
|
||||||
|
}
|
||||||
|
if err := pw.file.Close(); err != nil && c.firstErr == nil {
|
||||||
|
c.firstErr = err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return c.firstErr
|
||||||
|
}
|
||||||
|
|
||||||
|
// Close removes the temporary files. It is safe to call more than once.
|
||||||
|
func (c *SeqKmerCounter) Close() error {
|
||||||
|
_ = c.seal()
|
||||||
|
return os.RemoveAll(c.dir)
|
||||||
|
}
|
||||||
@@ -0,0 +1,155 @@
|
|||||||
|
package obikmer
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bufio"
|
||||||
|
"cmp"
|
||||||
|
"encoding/binary"
|
||||||
|
"io"
|
||||||
|
"os"
|
||||||
|
"slices"
|
||||||
|
"sync"
|
||||||
|
|
||||||
|
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obidefault"
|
||||||
|
)
|
||||||
|
|
||||||
|
// kmerOcc is one occurrence of a canonical k-mer in a sequence.
|
||||||
|
type kmerOcc struct {
|
||||||
|
kmer uint64
|
||||||
|
seq uint32
|
||||||
|
}
|
||||||
|
|
||||||
|
func cmpOcc(a, b kmerOcc) int {
|
||||||
|
if c := cmp.Compare(a.kmer, b.kmer); c != 0 {
|
||||||
|
return c
|
||||||
|
}
|
||||||
|
return cmp.Compare(a.seq, b.seq)
|
||||||
|
}
|
||||||
|
|
||||||
|
// CountAtLeast returns, sorted by k-mer, every canonical k-mer found in at
|
||||||
|
// least minSeq distinct sequences, with its exact sequence count.
|
||||||
|
//
|
||||||
|
// The result is held in memory: it is meant for quorum-filtered sets (the
|
||||||
|
// conserved k-mers), not for the complete k-mer spectrum. The first call
|
||||||
|
// seals the counter; AddSequence must not be called afterwards. The count
|
||||||
|
// phase uses obidefault.ParallelWorkers() goroutines over the partitions.
|
||||||
|
func (c *SeqKmerCounter) CountAtLeast(minSeq int) ([]SeqKmerCount, error) {
|
||||||
|
if err := c.seal(); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if minSeq < 1 {
|
||||||
|
minSeq = 1
|
||||||
|
}
|
||||||
|
|
||||||
|
workers := min(obidefault.ParallelWorkers(), c.p)
|
||||||
|
jobs := make(chan int)
|
||||||
|
results := make([][]SeqKmerCount, c.p)
|
||||||
|
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
var errMu sync.Mutex
|
||||||
|
var firstErr error
|
||||||
|
|
||||||
|
for range workers {
|
||||||
|
wg.Add(1)
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
var occ []kmerOcc // reused between partitions
|
||||||
|
for p := range jobs {
|
||||||
|
var err error
|
||||||
|
occ, err = c.loadPartition(p, occ[:0])
|
||||||
|
if err == nil {
|
||||||
|
results[p] = reduceOccurrences(occ, minSeq)
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
errMu.Lock()
|
||||||
|
if firstErr == nil {
|
||||||
|
firstErr = err
|
||||||
|
}
|
||||||
|
errMu.Unlock()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
for p := range c.p {
|
||||||
|
jobs <- p
|
||||||
|
}
|
||||||
|
close(jobs)
|
||||||
|
wg.Wait()
|
||||||
|
if firstErr != nil {
|
||||||
|
return nil, firstErr
|
||||||
|
}
|
||||||
|
|
||||||
|
total := 0
|
||||||
|
for _, r := range results {
|
||||||
|
total += len(r)
|
||||||
|
}
|
||||||
|
all := make([]SeqKmerCount, 0, total)
|
||||||
|
for _, r := range results {
|
||||||
|
all = append(all, r...)
|
||||||
|
}
|
||||||
|
slices.SortFunc(all, func(a, b SeqKmerCount) int { return cmp.Compare(a.Kmer, b.Kmer) })
|
||||||
|
return all, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// loadPartition appends to occ every (canonical k-mer, sequence id) of the
|
||||||
|
// partition file.
|
||||||
|
func (c *SeqKmerCounter) loadPartition(p int, occ []kmerOcc) ([]kmerOcc, error) {
|
||||||
|
f, err := os.Open(c.parts[p].path)
|
||||||
|
if err != nil {
|
||||||
|
return occ, err
|
||||||
|
}
|
||||||
|
defer f.Close()
|
||||||
|
r := bufio.NewReaderSize(f, 64*1024)
|
||||||
|
|
||||||
|
var packed, bases []byte
|
||||||
|
for {
|
||||||
|
id, err := binary.ReadUvarint(r)
|
||||||
|
if err == io.EOF {
|
||||||
|
return occ, nil
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return occ, err
|
||||||
|
}
|
||||||
|
n, err := binary.ReadUvarint(r)
|
||||||
|
if err != nil {
|
||||||
|
return occ, err
|
||||||
|
}
|
||||||
|
|
||||||
|
nb := (int(n) + 3) / 4
|
||||||
|
packed = slices.Grow(packed[:0], nb)[:nb]
|
||||||
|
if _, err := io.ReadFull(r, packed); err != nil {
|
||||||
|
return occ, err
|
||||||
|
}
|
||||||
|
bases = slices.Grow(bases[:0], int(n))[:n]
|
||||||
|
unpackBases(bases, packed)
|
||||||
|
|
||||||
|
for kmer := range IterCanonicalKmers(bases, c.k) {
|
||||||
|
occ = append(occ, kmerOcc{kmer, uint32(id)})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// reduceOccurrences sorts occ and returns the k-mers present in at least
|
||||||
|
// minSeq distinct sequences. occ is reordered in place.
|
||||||
|
func reduceOccurrences(occ []kmerOcc, minSeq int) []SeqKmerCount {
|
||||||
|
slices.SortFunc(occ, cmpOcc)
|
||||||
|
|
||||||
|
var out []SeqKmerCount
|
||||||
|
for i := 0; i < len(occ); {
|
||||||
|
kmer := occ[i].kmer
|
||||||
|
var seqs uint32
|
||||||
|
repeated := false
|
||||||
|
j := i
|
||||||
|
for ; j < len(occ) && occ[j].kmer == kmer; j++ {
|
||||||
|
if j == i || occ[j].seq != occ[j-1].seq {
|
||||||
|
seqs++
|
||||||
|
} else {
|
||||||
|
repeated = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if int(seqs) >= minSeq {
|
||||||
|
out = append(out, SeqKmerCount{Kmer: kmer, Sequences: seqs, Repeated: repeated})
|
||||||
|
}
|
||||||
|
i = j
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
@@ -0,0 +1,272 @@
|
|||||||
|
package obikmer
|
||||||
|
|
||||||
|
import (
|
||||||
|
"math/rand"
|
||||||
|
"slices"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
type bruteInfo struct {
|
||||||
|
seqs map[uint32]int // sequence id -> occurrences
|
||||||
|
repeated bool
|
||||||
|
}
|
||||||
|
|
||||||
|
// bruteForce is the oracle: every window of k unambiguous bases, circular
|
||||||
|
// junction included, encoded with EncodeCanonicalKmer.
|
||||||
|
func bruteForce(seqs [][]byte, circular []bool, k int) map[uint64]*bruteInfo {
|
||||||
|
res := map[uint64]*bruteInfo{}
|
||||||
|
for id, s := range seqs {
|
||||||
|
if len(s) < k {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if circular[id] {
|
||||||
|
s = append(slices.Clone(s), s[:k-1]...)
|
||||||
|
}
|
||||||
|
window:
|
||||||
|
for i := 0; i+k <= len(s); i++ {
|
||||||
|
w := s[i : i+k]
|
||||||
|
for _, b := range w {
|
||||||
|
if !unambiguousBase[b] {
|
||||||
|
continue window
|
||||||
|
}
|
||||||
|
}
|
||||||
|
kmer := EncodeCanonicalKmer(w, k)
|
||||||
|
info := res[kmer]
|
||||||
|
if info == nil {
|
||||||
|
info = &bruteInfo{seqs: map[uint32]int{}}
|
||||||
|
res[kmer] = info
|
||||||
|
}
|
||||||
|
info.seqs[uint32(id)]++
|
||||||
|
if info.seqs[uint32(id)] > 1 {
|
||||||
|
info.repeated = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return res
|
||||||
|
}
|
||||||
|
|
||||||
|
func randomSeqs(rng *rand.Rand, n, minLen, maxLen int, withN bool) [][]byte {
|
||||||
|
alphabet := "acgtACGT"
|
||||||
|
if withN {
|
||||||
|
alphabet = "acgtACGTacgtACGTNnR"
|
||||||
|
}
|
||||||
|
// A shared motif and its reverse complement make k-mers common to several
|
||||||
|
// sequences on both strands.
|
||||||
|
motif := make([]byte, 40)
|
||||||
|
for i := range motif {
|
||||||
|
motif[i] = "ACGT"[rng.Intn(4)]
|
||||||
|
}
|
||||||
|
rc := make([]byte, len(motif))
|
||||||
|
for i, b := range motif {
|
||||||
|
rc[len(motif)-1-i] = map[byte]byte{'A': 'T', 'C': 'G', 'G': 'C', 'T': 'A'}[b]
|
||||||
|
}
|
||||||
|
|
||||||
|
out := make([][]byte, n)
|
||||||
|
for i := range out {
|
||||||
|
l := minLen + rng.Intn(maxLen-minLen+1)
|
||||||
|
s := make([]byte, l)
|
||||||
|
for j := range s {
|
||||||
|
s[j] = alphabet[rng.Intn(len(alphabet))]
|
||||||
|
}
|
||||||
|
mode := rng.Intn(4)
|
||||||
|
if l < 60 { // too short to host the motifs
|
||||||
|
mode = 3
|
||||||
|
}
|
||||||
|
switch mode {
|
||||||
|
case 0:
|
||||||
|
copy(s[l/4:], motif)
|
||||||
|
case 1:
|
||||||
|
copy(s[l/4:], rc)
|
||||||
|
case 2: // internal repeat, both strands
|
||||||
|
copy(s[2:], motif[:20])
|
||||||
|
copy(s[l-25:], rc[:20])
|
||||||
|
}
|
||||||
|
out[i] = s
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSeqKmerCounterAgainstBruteForce(t *testing.T) {
|
||||||
|
cases := []struct {
|
||||||
|
name string
|
||||||
|
k, m, p int
|
||||||
|
withN bool
|
||||||
|
circ bool
|
||||||
|
}{
|
||||||
|
{"k7-linear", 7, 3, 8, false, false},
|
||||||
|
{"k7-circular-N", 7, 3, 5, true, true},
|
||||||
|
{"k18-linear-N", 18, 8, 16, true, false},
|
||||||
|
{"k18-circular", 18, 8, 7, false, true},
|
||||||
|
{"k31-p1", 31, 10, 1, false, true},
|
||||||
|
{"k4-small-m", 4, 1, 3, true, false},
|
||||||
|
}
|
||||||
|
for _, tc := range cases {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
rng := rand.New(rand.NewSource(int64(tc.k*100 + tc.p)))
|
||||||
|
seqs := randomSeqs(rng, 30, 5, 400, tc.withN)
|
||||||
|
circ := make([]bool, len(seqs))
|
||||||
|
for i := range circ {
|
||||||
|
circ[i] = tc.circ && i%3 != 0
|
||||||
|
}
|
||||||
|
|
||||||
|
c, err := NewSeqKmerCounter(t.TempDir(), tc.k, tc.m, tc.p)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer c.Close()
|
||||||
|
|
||||||
|
// Concurrent insertion, shuffled order, ids not contiguous.
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
for i := range seqs {
|
||||||
|
wg.Add(1)
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
c.AddSequence(uint32(i), seqs[i], circ[i])
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
wg.Wait()
|
||||||
|
|
||||||
|
want := bruteForce(seqs, circ, tc.k)
|
||||||
|
for _, q := range []int{1, 2, 3, 10} {
|
||||||
|
got, err := c.CountAtLeast(q)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !slices.IsSortedFunc(got, func(a, b SeqKmerCount) int {
|
||||||
|
switch {
|
||||||
|
case a.Kmer < b.Kmer:
|
||||||
|
return -1
|
||||||
|
case a.Kmer > b.Kmer:
|
||||||
|
return 1
|
||||||
|
}
|
||||||
|
return 0
|
||||||
|
}) {
|
||||||
|
t.Fatal("result is not sorted by k-mer")
|
||||||
|
}
|
||||||
|
|
||||||
|
n := 0
|
||||||
|
for kmer, info := range want {
|
||||||
|
if len(info.seqs) >= q {
|
||||||
|
n++
|
||||||
|
_ = kmer
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if len(got) != n {
|
||||||
|
t.Fatalf("q=%d: %d k-mers, want %d", q, len(got), n)
|
||||||
|
}
|
||||||
|
for _, g := range got {
|
||||||
|
info := want[g.Kmer]
|
||||||
|
if info == nil {
|
||||||
|
t.Fatalf("q=%d: unexpected k-mer %d", q, g.Kmer)
|
||||||
|
}
|
||||||
|
if int(g.Sequences) != len(info.seqs) || g.Repeated != info.repeated {
|
||||||
|
t.Fatalf("q=%d kmer %d: got (%d,%v) want (%d,%v)",
|
||||||
|
q, g.Kmer, g.Sequences, g.Repeated, len(info.seqs), info.repeated)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// A k-mer appearing on both strands of one sequence counts once but is Repeated.
|
||||||
|
func TestSeqKmerCounterBothStrandsOneSequence(t *testing.T) {
|
||||||
|
const k = 8
|
||||||
|
fwd := "ACGGTCAT"
|
||||||
|
rc := "ATGACCGT"
|
||||||
|
seq := []byte("TTTT" + fwd + "GGGGGGGG" + rc + "CCCC")
|
||||||
|
|
||||||
|
c, err := NewSeqKmerCounter(t.TempDir(), k, 3, 4)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer c.Close()
|
||||||
|
c.AddSequence(7, seq, false)
|
||||||
|
|
||||||
|
got, err := c.CountAtLeast(1)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
target := EncodeCanonicalKmer([]byte(fwd), k)
|
||||||
|
found := false
|
||||||
|
for _, g := range got {
|
||||||
|
if g.Kmer == target {
|
||||||
|
found = true
|
||||||
|
if g.Sequences != 1 || !g.Repeated {
|
||||||
|
t.Errorf("got %+v, want 1 sequence and Repeated", g)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !found {
|
||||||
|
t.Fatal("k-mer not found")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Ambiguous bases must never be part of a k-mer, and the circular junction
|
||||||
|
// must produce the wrap-around k-mers.
|
||||||
|
func TestSeqKmerCounterAmbiguityAndCircular(t *testing.T) {
|
||||||
|
const k = 5
|
||||||
|
c, err := NewSeqKmerCounter(t.TempDir(), k, 2, 3)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer c.Close()
|
||||||
|
|
||||||
|
c.AddSequence(0, []byte("ACGTNACGT"), false) // runs ACGT (<k) and ACGT (<k): nothing
|
||||||
|
c.AddSequence(1, []byte("AACCGGTT"), true) // 8 circular windows
|
||||||
|
|
||||||
|
got, err := c.CountAtLeast(1)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(got) == 0 {
|
||||||
|
t.Fatal("expected k-mers from the circular sequence")
|
||||||
|
}
|
||||||
|
var circ []string
|
||||||
|
for i := 0; i < 8; i++ {
|
||||||
|
var sb strings.Builder
|
||||||
|
for j := 0; j < k; j++ {
|
||||||
|
sb.WriteByte("AACCGGTT"[(i+j)%8])
|
||||||
|
}
|
||||||
|
circ = append(circ, sb.String())
|
||||||
|
}
|
||||||
|
want := map[uint64]bool{}
|
||||||
|
for _, w := range circ {
|
||||||
|
want[EncodeCanonicalKmer([]byte(w), k)] = true
|
||||||
|
}
|
||||||
|
if len(got) != len(want) {
|
||||||
|
t.Fatalf("got %d k-mers, want %d", len(got), len(want))
|
||||||
|
}
|
||||||
|
for _, g := range got {
|
||||||
|
if !want[g.Kmer] || g.Sequences != 1 {
|
||||||
|
t.Errorf("unexpected %+v", g)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSeqKmerCounterParameters(t *testing.T) {
|
||||||
|
if _, err := NewSeqKmerCounter(t.TempDir(), 1, 1, 4); err == nil {
|
||||||
|
t.Error("k=1 must fail")
|
||||||
|
}
|
||||||
|
if _, err := NewSeqKmerCounter(t.TempDir(), 32, 5, 4); err == nil {
|
||||||
|
t.Error("k=32 must fail")
|
||||||
|
}
|
||||||
|
if _, err := NewSeqKmerCounter(t.TempDir(), 10, 10, 4); err == nil {
|
||||||
|
t.Error("m>=k must fail")
|
||||||
|
}
|
||||||
|
if _, err := NewSeqKmerCounter(t.TempDir(), 10, 4, 0); err == nil {
|
||||||
|
t.Error("zero partitions must fail")
|
||||||
|
}
|
||||||
|
c, err := NewSeqKmerCounter(t.TempDir(), 10, 0, 4) // default minimizer
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := c.Close(); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := c.Close(); err != nil {
|
||||||
|
t.Fatalf("second Close must be harmless: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -10,6 +10,16 @@ import (
|
|||||||
// decode2bit maps 2-bit codes back to nucleotide bytes.
|
// decode2bit maps 2-bit codes back to nucleotide bytes.
|
||||||
var decode2bit = [4]byte{'a', 'c', 'g', 't'}
|
var decode2bit = [4]byte{'a', 'c', 'g', 't'}
|
||||||
|
|
||||||
|
// unpackBases decodes the 2-bit packed bases in packed into dst (one lowercase
|
||||||
|
// nucleotide per byte, len(dst) bases) and returns dst.
|
||||||
|
func unpackBases(dst, packed []byte) []byte {
|
||||||
|
for i := range dst {
|
||||||
|
code := (packed[i/4] >> uint(6-(i%4)*2)) & 0x03
|
||||||
|
dst[i] = decode2bit[code]
|
||||||
|
}
|
||||||
|
return dst
|
||||||
|
}
|
||||||
|
|
||||||
// SkmReader reads super-kmers from a binary .skm file.
|
// SkmReader reads super-kmers from a binary .skm file.
|
||||||
type SkmReader struct {
|
type SkmReader struct {
|
||||||
r *bufio.Reader
|
r *bufio.Reader
|
||||||
@@ -46,13 +56,7 @@ func (sr *SkmReader) Next() (SuperKmer, bool) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Decode to nucleotide bytes
|
// Decode to nucleotide bytes
|
||||||
seq := make([]byte, seqLen)
|
seq := unpackBases(make([]byte, seqLen), packed)
|
||||||
for i := 0; i < seqLen; i++ {
|
|
||||||
byteIdx := i / 4
|
|
||||||
bitPos := uint(6 - (i%4)*2)
|
|
||||||
code := (packed[byteIdx] >> bitPos) & 0x03
|
|
||||||
seq[i] = decode2bit[code]
|
|
||||||
}
|
|
||||||
|
|
||||||
return SuperKmer{
|
return SuperKmer{
|
||||||
Sequence: seq,
|
Sequence: seq,
|
||||||
|
|||||||
+23
-22
@@ -16,8 +16,9 @@ import (
|
|||||||
// Nucleotide encoding: A=00, C=01, G=10, T=11.
|
// Nucleotide encoding: A=00, C=01, G=10, T=11.
|
||||||
// The last byte is zero-padded on the low bits if len%4 != 0.
|
// The last byte is zero-padded on the low bits if len%4 != 0.
|
||||||
type SkmWriter struct {
|
type SkmWriter struct {
|
||||||
w *bufio.Writer
|
w *bufio.Writer
|
||||||
file *os.File
|
file *os.File
|
||||||
|
scratch []byte
|
||||||
}
|
}
|
||||||
|
|
||||||
// NewSkmWriter creates a new SkmWriter writing to the given file path.
|
// NewSkmWriter creates a new SkmWriter writing to the given file path.
|
||||||
@@ -32,36 +33,36 @@ func NewSkmWriter(path string) (*SkmWriter, error) {
|
|||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// appendPacked appends seq to dst packed 2 bits per base (A=00, C=01, G=10,
|
||||||
|
// T=11), four bases per byte, the last byte zero-padded on the low bits.
|
||||||
|
func appendPacked(dst, seq []byte) []byte {
|
||||||
|
for i := 0; i < len(seq); i += 4 {
|
||||||
|
var packed byte
|
||||||
|
for j := 0; j < 4; j++ {
|
||||||
|
packed <<= 2
|
||||||
|
if i+j < len(seq) {
|
||||||
|
packed |= __single_base_code__[seq[i+j]&31]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
dst = append(dst, packed)
|
||||||
|
}
|
||||||
|
return dst
|
||||||
|
}
|
||||||
|
|
||||||
// Write encodes a SuperKmer to the .skm file.
|
// Write encodes a SuperKmer to the .skm file.
|
||||||
// The sequence bytes are packed 2 bits per base.
|
// The sequence bytes are packed 2 bits per base.
|
||||||
func (sw *SkmWriter) Write(sk SuperKmer) error {
|
func (sw *SkmWriter) Write(sk SuperKmer) error {
|
||||||
seq := sk.Sequence
|
|
||||||
seqLen := uint16(len(seq))
|
|
||||||
|
|
||||||
// Write length
|
// Write length
|
||||||
var lenbuf [2]byte
|
var lenbuf [2]byte
|
||||||
binary.LittleEndian.PutUint16(lenbuf[:], seqLen)
|
binary.LittleEndian.PutUint16(lenbuf[:], uint16(len(sk.Sequence)))
|
||||||
if _, err := sw.w.Write(lenbuf[:]); err != nil {
|
if _, err := sw.w.Write(lenbuf[:]); err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
// Encode and write packed sequence (2 bits/base)
|
// Encode and write packed sequence (2 bits/base)
|
||||||
nBytes := (int(seqLen) + 3) / 4
|
sw.scratch = appendPacked(sw.scratch[:0], sk.Sequence)
|
||||||
for i := 0; i < nBytes; i++ {
|
_, err := sw.w.Write(sw.scratch)
|
||||||
var packed byte
|
return err
|
||||||
for j := 0; j < 4; j++ {
|
|
||||||
pos := i*4 + j
|
|
||||||
packed <<= 2
|
|
||||||
if pos < int(seqLen) {
|
|
||||||
packed |= __single_base_code__[seq[pos]&31]
|
|
||||||
}
|
|
||||||
}
|
|
||||||
if err := sw.w.WriteByte(packed); err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return nil
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// Close flushes buffered data and closes the underlying file.
|
// Close flushes buffered data and closes the underlying file.
|
||||||
|
|||||||
Reference in new issue
Block a user