Refactor MatchWords indexing and introduce seed indexing

Corrects the indexing strategy for seed searching, improving performance and introducing a new seed index structure for efficient pattern lookup.

This commit updates the core matching logic by shifting the indexing strategy to use two halves of k/2 bases, resulting in significant speed improvements for sequence matching. It also introduces a new `seedIndex` structure for fast pattern retrieval and implements comprehensive performance testing logic for matching functions. Future features, such as bit-parallel search, are deferred.

pkg/obiprimer/match_apat.go

@@ -1,87 +0,0 @@
-package obiprimer
-
-import (
-	"fmt"
-
-	"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiapat"
-	"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiseq"
-)
-
-// apatSafe returns seq when it only holds a, c, g, t; otherwise a copy where u
-// becomes t and every other symbol becomes z. obiapat reads IUPAC codes in the
-// sequence (an n matches every base); z is the symbol it matches with nothing,
-// which gives back the ecoPrimers rule "an ambiguous base is an error".
-func apatSafe(seq *obiseq.BioSequence) *obiseq.BioSequence {
-	raw := seq.Sequence()
-	clean := true
-	for _, b := range raw {
-		if b != 'a' && b != 'c' && b != 'g' && b != 't' {
-			clean = false
-			break
-		}
-	}
-	if clean {
-		return seq
-	}
-	out := make([]byte, len(raw))
-	for i, b := range raw {
-		switch b {
-		case 'a', 'c', 'g', 't':
-			out[i] = b
-		case 'A', 'C', 'G', 'T':
-			out[i] = b + 'a' - 'A'
-		case 'u', 'U':
-			out[i] = 't'
-		default:
-			out[i] = 'z'
-		}
-	}
-	return obiseq.NewBioSequence(seq.Id(), out, "")
-}
-
-// matchApat searches the words with obiapat, the bit-parallel Baeza-Yates /
-// Manber matcher of OBITools (the engine of ecoPrimers and obipcr).
-func matchApat(c *Corpus, words []uint64, p MatchParams) ([]WordMatches, error) {
-	if err := p.validate(); err != nil {
-		return nil, err
-	}
-	fwd, rev := wordStrings(words, p.K)
-	pats := make([][2]obiapat.ApatPattern, len(words))
-	for i := range words {
-		var err error
-		if pats[i][0], err = obiapat.MakeApatPattern(fwd[i], p.MaxErrors, false); err != nil {
-			return nil, fmt.Errorf("obiprimer: pattern %s: %w", fwd[i], err)
-		}
-		if pats[i][1], err = obiapat.MakeApatPattern(rev[i], p.MaxErrors, false); err != nil {
-			return nil, fmt.Errorf("obiprimer: pattern %s: %w", rev[i], err)
-		}
-	}
-
-	mk := func() scanFunc {
-		var recycle []obiapat.ApatSequence
-		return func(seq *obiseq.BioSequence, emit func(word, pos, errors int, reverse bool)) {
-			n := seq.Len()
-			if n < p.K {
-				return
-			}
-			aseq, err := obiapat.MakeApatSequence(apatSafe(seq), p.Circular, recycle...)
-			if err != nil {
-				panic(fmt.Sprintf("obiprimer: cannot encode sequence %s: %v", seq.Id(), err))
-			}
-			recycle = append(recycle[:0], aseq)
-
-			for w := range pats {
-				for strand := range 2 {
-					for _, m := range pats[w][strand].FindAllIndex(aseq, 0, -1) {
-						// On circular sequences obiapat also scans a copy of the
-						// start: those hits repeat the ones at start-n.
-						if m[0] >= 0 && m[0] < n {
-							emit(w, m[0], m[2], strand == 1)
-						}
-					}
-				}
-			}
-		}
-	}
-	return matchCorpus(c, words, mk), nil
-}

pkg/obiprimer/match_seed.go

@@ -0,0 +1,214 @@
+package obiprimer
+
+import (
+	"math/bits"
+	"slices"
+
+	"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiseq"
+)
+
+// seedIndex maps the value of one half of a window to the patterns that are
+// candidates for it: those whose same half is within d substitutions of it.
+//
+// Pigeonhole: a window with at most e = 2d+1 errors against a pattern has at
+// most d errors in at least one of the two halves of the pattern, so the
+// windows to verify are exactly those found through one of the two indexes.
+// Indexing the half of every pattern together with its neighbours at distance
+// d (1 + 3L values for d = 1, 1 + 3L + 9L(L-1)/2 for d = 2, L = 9 bases) costs
+// a few tens of entries per pattern, and a window of a random text finds about
+// patterns x neighbours / 4^L candidates, far fewer than the patterns.
+type seedIndex struct {
+	shift   uint     // slot = seed >> shift when hashed, else seed
+	hashed  bool     // the seed space is too large for a direct table
+	offsets []uint32 // CSR: ids[offsets[s]:offsets[s+1]] are the candidates of slot s
+	ids     []int32  // pattern ids (2*word + strand)
+}
+
+const maxDirectSeedBits = 22
+
+func (x *seedIndex) slot(seed uint64) uint32 {
+	if !x.hashed {
+		return uint32(seed)
+	}
+	return uint32((seed * 0x9E3779B97F4A7C15) >> x.shift)
+}
+
+func (x *seedIndex) candidates(seed uint64) []int32 {
+	s := x.slot(seed)
+	return x.ids[x.offsets[s]:x.offsets[s+1]]
+}
+
+// neighbours calls fn for v and every value within d substitutions of it (each
+// once), v being L bases packed on 2 bits.
+func neighbours(v uint64, L, d int, fn func(uint64)) {
+	var rec func(v uint64, from, left int)
+	rec = func(v uint64, from, left int) {
+		fn(v)
+		if left == 0 {
+			return
+		}
+		for pos := from; pos < L; pos++ {
+			for alt := uint64(1); alt <= 3; alt++ {
+				rec(v^(alt<<(2*uint(pos))), pos+1, left-1)
+			}
+		}
+	}
+	rec(v, 0, d)
+}
+
+// newSeedIndex indexes, for every pattern, the half of the given width (in
+// bases) selected by (code >> shiftBases*2) & mask, with its neighbours at
+// distance d.
+func newSeedIndex(patterns []uint64, shiftBases, width, d int) *seedIndex {
+	x := &seedIndex{}
+	bitsNeeded := 2 * width
+	if bitsNeeded > maxDirectSeedBits {
+		x.hashed, bitsNeeded = true, maxDirectSeedBits
+		x.shift = uint(64 - maxDirectSeedBits)
+	}
+	nslots := 1 << uint(bitsNeeded)
+	halfMask := uint64(1)<<(2*uint(width)) - 1
+
+	type entry struct {
+		slot uint32
+		id   int32
+	}
+	var entries []entry
+	var scratch []uint32
+	for id, pat := range patterns {
+		half := (pat >> (2 * uint(shiftBases))) & halfMask
+		scratch = scratch[:0]
+		neighbours(half, width, d, func(v uint64) { scratch = append(scratch, x.slot(v)) })
+		slices.Sort(scratch)
+		scratch = slices.Compact(scratch) // hashing may merge neighbours
+		for _, s := range scratch {
+			entries = append(entries, entry{s, int32(id)})
+		}
+	}
+
+	x.offsets = make([]uint32, nslots+1)
+	for _, e := range entries {
+		x.offsets[e.slot+1]++
+	}
+	for i := 1; i <= nslots; i++ {
+		x.offsets[i] += x.offsets[i-1]
+	}
+	x.ids = make([]int32, len(entries))
+	fill := slices.Clone(x.offsets[:nslots])
+	for _, e := range entries {
+		x.ids[fill[e.slot]] = e.id
+		fill[e.slot]++
+	}
+	return x
+}
+
+// MatchWords finds every approximate occurrence, on both strands, of the given
+// canonical words in the sequences of the corpus, with the rules of MatchParams
+// (substitutions only, ambiguous bases count as errors, circular junction).
+//
+// It uses the pigeonhole seed indexes above. It returns exactly what the
+// reference matchNaive returns, and what ecoPrimers' fuzzy search finds, in a
+// small fraction of the time: on the Teleostei tutorial data (255 words, 3909
+// sequences, 65 Mb, e = 3) 0.42 s on one core, against 113 s for the
+// obiapat (agrep) search it replaces.
+func MatchWords(c *Corpus, words []uint64, p MatchParams) ([]WordMatches, error) {
+	if err := p.validate(); err != nil {
+		return nil, err
+	}
+	k := p.K
+	d := p.MaxErrors / 2
+
+	// pattern 2w is the word, 2w+1 its reverse complement
+	patterns := make([]uint64, 2*len(words))
+	for i, w := range words {
+		patterns[2*i] = codeOf(wordCodes(w, k, false), k)
+		patterns[2*i+1] = codeOf(wordCodes(w, k, true), k)
+	}
+
+	right := k / 2 // the right half holds the low bits
+	left := k - right
+	leftIdx := newSeedIndex(patterns, right, left, d)
+	rightIdx := newSeedIndex(patterns, 0, right, d)
+
+	rightMask := uint64(1)<<(2*uint(right)) - 1
+	windowMask := uint64(1)<<(2*uint(k)) - 1
+	lowBits := uint64(0x5555555555555555) & windowMask
+	maxErr := p.MaxErrors
+
+	mk := func() scanFunc {
+		return func(seq *obiseq.BioSequence, emit func(word, pos, errors int, reverse bool)) {
+			raw := seq.Sequence()
+			n := len(raw)
+			if n < k {
+				return
+			}
+			total := n // number of bases fed to the rolling window
+			windows := n - k + 1
+			if p.Circular {
+				total, windows = n+k-1, n
+			}
+
+			var code, bad uint64 // window, and its non-ACGT positions (bit 2j)
+			verify := func(pos int, id int32) {
+				pat := patterns[id]
+				x := code ^ pat
+				if errs := bits.OnesCount64(((x | x>>1) & lowBits) | bad); errs <= maxErr {
+					emit(int(id)>>1, pos, errs, id&1 == 1)
+				}
+			}
+
+			for i := 0; i < total; i++ {
+				j := i
+				if j >= n {
+					j -= n
+				}
+				v := baseCode[raw[j]]
+				isBad := uint64(0)
+				if v > 3 {
+					v, isBad = 0, 1
+				}
+				code = (code<<2 | uint64(v)) & windowMask
+				bad = (bad<<2 | isBad) & windowMask
+
+				if i < k-1 {
+					continue
+				}
+				pos := i - k + 1
+				if pos >= windows {
+					break
+				}
+				leftCands := leftIdx.candidates(code >> (2 * uint(right)))
+				for _, id := range leftCands {
+					verify(pos, id)
+				}
+				for _, id := range rightIdx.candidates(code & rightMask) {
+					// ids are sorted: skip the ones the left index has just verified
+					if _, done := slices.BinarySearch(leftCands, id); !done {
+						verify(pos, id)
+					}
+				}
+			}
+		}
+	}
+	return matchCorpus(c, words, mk), nil
+}
+
+// wordCodes returns the 2-bit codes (A=0..T=3) of the word or of its reverse
+// complement, 5' to 3'.
+func wordCodes(word uint64, k int, reverse bool) []byte {
+	fwd, rev := wordStrings([]uint64{word}, k)
+	s := fwd[0]
+	if reverse {
+		s = rev[0]
+	}
+	return codesOf(s)
+}
+
+// codeOf packs codes (each 0..3) on 2 bits, first base in the high bits.
+func codeOf(codes []byte, k int) uint64 {
+	var v uint64
+	for _, c := range codes[:k] {
+		v = v<<2 | uint64(c)
+	}
+	return v
+}

pkg/obiprimer/match_test.go

@@ -71,8 +71,8 @@ func TestMatchParamsValidation(t *testing.T) {
 		if _, err := matchNaive(c, nil, p); err == nil {
 			t.Errorf("%+v must be rejected", p)
 		}
-		if _, err := matchApat(c, nil, p); err == nil {
-			t.Errorf("%+v must be rejected by obiapat variant", p)
+		if _, err := MatchWords(c, nil, p); err == nil {
+			t.Errorf("%+v must be rejected by MatchWords", p)
 		}
 	}
 }
@@ -130,7 +130,7 @@ func TestMatchersAgreeWithReference(t *testing.T) {
 	for _, tc := range []struct {
 		k, e     int
 		circular bool
-	}{{18, 0, false}, {18, 1, true}, {18, 3, false}, {18, 3, true}, {12, 2, true}, {25, 4, false}} {
+	}{{18, 0, false}, {18, 1, true}, {18, 3, false}, {18, 3, true}, {12, 2, true}, {25, 4, false}, {18, 2, true}, {18, 5, false}, {31, 3, true}, {24, 3, false}, {9, 2, true}, {18, 6, true}, {18, 7, false}, {20, 8, true}, {12, 11, false}} {
 		t.Run(fmt.Sprintf("k%d-e%d-circ%v", tc.k, tc.e, tc.circular), func(t *testing.T) {
 			rng := rand.New(rand.NewSource(int64(tc.k*100 + tc.e)))
 			var wordStr []string
@@ -155,13 +155,13 @@ func TestMatchersAgreeWithReference(t *testing.T) {
 				t.Fatalf("test data too poor: only %d reference matches", total)
 			}
 
-			got, err := matchApat(c, words, p)
+			got, err := MatchWords(c, words, p)
 			if err != nil {
 				t.Fatal(err)
 			}
 			for i := range want {
 				if !reflect.DeepEqual(got[i].Matches, want[i].Matches) {
-					t.Fatalf("word %d: obiapat gives %v\nreference gives %v", i, got[i].Matches, want[i].Matches)
+					t.Fatalf("word %d: MatchWords gives %v\nreference gives %v", i, got[i].Matches, want[i].Matches)
 				}
 			}
 		})
@@ -200,7 +200,7 @@ func readEcoCounts(t *testing.T, path string) map[string]int {
 // strict word test). For every strict word, ecoPrimers logs the number of
 // example sequences in which the fuzzy search (-e 3, both strands, circular)
 // finds it; testdata/ecoprimers_*.words.tsv holds those numbers.
-func TestMatchApatAgainstEcoPrimers(t *testing.T) {
+func TestMatchWordsAgainstEcoPrimers(t *testing.T) {
 	dir := os.Getenv("OBIPRIMERS_TUTORIAL")
 	if dir == "" {
 		t.Skip("OBIPRIMERS_TUTORIAL not set")
@@ -244,7 +244,7 @@ func TestMatchApatAgainstEcoPrimers(t *testing.T) {
 			runtime.GC()
 			runtime.ReadMemStats(&before)
 			start := time.Now()
-			res, err := matchApat(corpus, words, MatchParams{K: 18, MaxErrors: 3, Circular: true})
+			res, err := MatchWords(corpus, words, MatchParams{K: 18, MaxErrors: 3, Circular: true})
 			elapsed := time.Since(start)
 			runtime.ReadMemStats(&after)
 			if err != nil {
@@ -258,10 +258,29 @@ func TestMatchApatAgainstEcoPrimers(t *testing.T) {
 			for _, s := range corpus.Seqs {
 				bases += s.Len()
 			}
-			t.Logf("obiapat: %d words x %d sequences (%.1f Mb) e=3 in %s (%.0f Mb/s per word-strand pass); %d matches; heap +%.0f MB, allocated %.0f MB",
-				len(words), len(corpus.Seqs), float64(bases)/1e6, elapsed.Round(time.Millisecond),
-				float64(bases)*float64(len(words))*2/1e6/elapsed.Seconds(), nm,
-				float64(after.HeapAlloc-before.HeapAlloc)/1e6, float64(after.TotalAlloc-before.TotalAlloc)/1e6)
+			t.Logf("%d words x %d sequences (%.1f Mb) e=3 in %s; %d matches; allocated %.0f MB",
+				len(words), len(corpus.Seqs), float64(bases)/1e6, elapsed.Round(time.Millisecond), nm,
+				float64(after.TotalAlloc-before.TotalAlloc)/1e6)
+
+			// On real sequences (with their ambiguity codes), the first ones must give
+			// exactly what the reference implementation gives.
+			sub := &Corpus{Seqs: corpus.Seqs[:40], Roles: corpus.Roles[:40]}
+			wantSub, err := matchNaive(sub, words, MatchParams{K: 18, MaxErrors: 3, Circular: true})
+			if err != nil {
+				t.Fatal(err)
+			}
+			gotSub, err := MatchWords(sub, words, MatchParams{K: 18, MaxErrors: 3, Circular: true})
+			if err != nil {
+				t.Fatal(err)
+			}
+			subMatches := 0
+			for i := range wantSub {
+				subMatches += len(wantSub[i].Matches)
+				if !reflect.DeepEqual(gotSub[i].Matches, wantSub[i].Matches) {
+					t.Fatalf("word %d differs from the reference on the first 40 sequences", i)
+				}
+			}
+			t.Logf("first 40 sequences identical to the reference (%d matches)", subMatches)
 
 			eco := readEcoCounts(t, filepath.Join("testdata", "ecoprimers_"+tc.name+".words.tsv"))
 			compared, differ := 0, 0

pkg/obiprimer/pairs_test.go

@@ -136,7 +136,7 @@ func TestPairsAgainstEcoPrimers(t *testing.T) {
 			for i, w := range sw.Words {
 				words[i] = w.Kmer
 			}
-			wm, err := matchApat(corpus, words, MatchParams{K: 18, MaxErrors: 3, Circular: true})
+			wm, err := MatchWords(corpus, words, MatchParams{K: 18, MaxErrors: 3, Circular: true})
 			if err != nil {
 				t.Fatal(err)
 			}