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.

pkg/obikmer/seqquorum.go

@@ -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)
+}

pkg/obikmer/seqquorum_count.go

@@ -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
+}

pkg/obikmer/seqquorum_test.go

@@ -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)
+	}
+}

pkg/obikmer/skm_reader.go

@@ -10,6 +10,16 @@ import (
 // decode2bit maps 2-bit codes back to nucleotide bytes.
 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.
 type SkmReader struct {
 	r    *bufio.Reader
@@ -46,13 +56,7 @@ func (sr *SkmReader) Next() (SuperKmer, bool) {
 	}
 
 	// Decode to nucleotide bytes
-	seq := make([]byte, seqLen)
-	for i := 0; i < seqLen; i++ {
-		byteIdx := i / 4
-		bitPos := uint(6 - (i%4)*2)
-		code := (packed[byteIdx] >> bitPos) & 0x03
-		seq[i] = decode2bit[code]
-	}
+	seq := unpackBases(make([]byte, seqLen), packed)
 
 	return SuperKmer{
 		Sequence: seq,

pkg/obikmer/skm_writer.go

@@ -16,8 +16,9 @@ import (
 // Nucleotide encoding: A=00, C=01, G=10, T=11.
 // The last byte is zero-padded on the low bits if len%4 != 0.
 type SkmWriter struct {
-	w    *bufio.Writer
-	file *os.File
+	w       *bufio.Writer
+	file    *os.File
+	scratch []byte
 }
 
 // NewSkmWriter creates a new SkmWriter writing to the given file path.
@@ -32,36 +33,36 @@ func NewSkmWriter(path string) (*SkmWriter, error) {
 	}, 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.
 // The sequence bytes are packed 2 bits per base.
 func (sw *SkmWriter) Write(sk SuperKmer) error {
-	seq := sk.Sequence
-	seqLen := uint16(len(seq))
-
 	// Write length
 	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 {
 		return err
 	}
 
 	// Encode and write packed sequence (2 bits/base)
-	nBytes := (int(seqLen) + 3) / 4
-	for i := 0; i < nBytes; i++ {
-		var packed byte
-		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
+	sw.scratch = appendPacked(sw.scratch[:0], sk.Sequence)
+	_, err := sw.w.Write(sw.scratch)
+	return err
 }
 
 // Close flushes buffered data and closes the underlying file.