diff --git a/pkg/obikmer/debruijn.go b/pkg/obikmer/debruijn.go index a707ca7..c2e8a9e 100644 --- a/pkg/obikmer/debruijn.go +++ b/pkg/obikmer/debruijn.go @@ -5,7 +5,6 @@ import ( "container/heap" "fmt" "math" - "math/bits" "os" "slices" @@ -610,14 +609,6 @@ func (graph *DeBruijnGraph) WriteGml(filename string) error { return err } -// Calculating the hamming distance between two k-mers. -func (g *DeBruijnGraph) HammingDistance(kmer1, kmer2 uint64) int { - ident := ^((kmer1 & kmer2) | (^kmer1 & ^kmer2)) - ident |= (ident >> 1) - ident &= 0x5555555555555555 & g.kmermask - return bits.OnesCount64(ident) -} - type UInt64Heap []uint64 func (h UInt64Heap) Len() int { return len(h) } diff --git a/pkg/obikmer/enumkmer.go b/pkg/obikmer/enumkmer.go new file mode 100644 index 0000000..678bf08 --- /dev/null +++ b/pkg/obikmer/enumkmer.go @@ -0,0 +1,91 @@ +package obikmer + +import "iter" + +// KmerConstraints describes lexicographic properties a k-mer must satisfy. +// +// MaxCount is indexed by the 2-bit base code (A=0, C=1, G=2, T=3): +// a negative value means unlimited, 0 forbids the base altogether +// (so MaxCount[2] = 0 restricts the alphabet to A, C, T). +// +// MaxRun is the maximum number of consecutive identical bases; +// a value <= 0 means unlimited. +// +// The zero value forbids every base: use NoKmerConstraints as a starting point. +type KmerConstraints struct { + MaxCount [4]int + MaxRun int +} + +// NoKmerConstraints returns constraints accepting every k-mer. +func NoKmerConstraints() KmerConstraints { + return KmerConstraints{MaxCount: [4]int{-1, -1, -1, -1}} +} + +// IterConstrainedKmers returns an iterator over all k-mers of size k +// satisfying the constraints c, encoded as in EncodeKmer, in increasing +// numerical order. +// +// K-mers are built base by base and a branch is abandoned as soon as a +// constraint is violated (both MaxCount and MaxRun are monotone: a violating +// prefix cannot be repaired by extension), so the cost is proportional to the +// number of explored valid prefixes, not to 4^k. +// +// k must be between 1 and 31, otherwise the iterator is empty. +// +// Example: +// +// c := NoKmerConstraints() +// c.MaxCount[2] = 0 // no G +// c.MaxRun = 3 +// for kmer := range IterConstrainedKmers(8, c) { ... } +func IterConstrainedKmers(k int, c KmerConstraints) iter.Seq[uint64] { + return func(yield func(uint64) bool) { + if k < 1 || k > 31 { + return + } + + // Nothing to enumerate if the per-base caps cannot even sum to k. + capacity := 0 + for _, m := range c.MaxCount { + if m < 0 || m >= k { + capacity += k + } else { + capacity += m + } + } + if capacity < k { + return + } + + var counts [4]int + + var walk func(depth int, kmer uint64, last, run int) bool + walk = func(depth int, kmer uint64, last, run int) bool { + if depth == k { + return yield(kmer) + } + for b := range 4 { + if c.MaxCount[b] >= 0 && counts[b] >= c.MaxCount[b] { + continue + } + r := 1 + if b == last { + r = run + 1 + } + if c.MaxRun > 0 && r > c.MaxRun { + continue + } + counts[b]++ + ok := walk(depth+1, kmer<<2|uint64(b), b, r) + counts[b]-- + if !ok { + return false + } + } + return true + } + + walk(0, 0, -1, 0) + } +} diff --git a/pkg/obikmer/enumkmer_test.go b/pkg/obikmer/enumkmer_test.go new file mode 100644 index 0000000..d2d528d --- /dev/null +++ b/pkg/obikmer/enumkmer_test.go @@ -0,0 +1,85 @@ +package obikmer + +import ( + "strings" + "testing" +) + +func collectConstrained(k int, c KmerConstraints) []uint64 { + var out []uint64 + for km := range IterConstrainedKmers(k, c) { + out = append(out, km) + } + return out +} + +func TestIterConstrainedKmersUnconstrained(t *testing.T) { + got := collectConstrained(4, NoKmerConstraints()) + if len(got) != 256 { + t.Fatalf("expected 256 k-mers, got %d", len(got)) + } + for i, km := range got { + if km != uint64(i) { + t.Fatalf("not in increasing order at %d: %d", i, km) + } + } +} + +func TestIterConstrainedKmersNoG(t *testing.T) { + c := NoKmerConstraints() + c.MaxCount[2] = 0 + got := collectConstrained(5, c) + if len(got) != 243 { + t.Fatalf("expected 3^5 = 243 k-mers, got %d", len(got)) + } + for _, km := range got { + if strings.Contains(string(DecodeKmer(km, 5, nil)), "g") { + t.Fatalf("k-mer %d contains G", km) + } + } +} + +func TestIterConstrainedKmersMaxRun(t *testing.T) { + c := NoKmerConstraints() + c.MaxRun = 2 + // a(n) = 3*(a(n-1)+a(n-2)), a(1)=4, a(2)=16 + want := map[int]int{3: 60, 4: 228} + for k, n := range want { + if got := len(collectConstrained(k, c)); got != n { + t.Fatalf("k=%d: expected %d, got %d", k, n, got) + } + } +} + +func TestIterConstrainedKmersMaxCount(t *testing.T) { + c := NoKmerConstraints() + c.MaxCount[0] = 1 // at most one A + for _, km := range collectConstrained(6, c) { + if strings.Count(string(DecodeKmer(km, 6, nil)), "a") > 1 { + t.Fatalf("k-mer %d has more than one A", km) + } + } +} + +func TestIterConstrainedKmersImpossible(t *testing.T) { + c := KmerConstraints{MaxCount: [4]int{1, 1, 1, 0}} + if got := collectConstrained(4, c); len(got) != 0 { + t.Fatalf("expected no k-mer, got %d", len(got)) + } + if got := collectConstrained(0, NoKmerConstraints()); len(got) != 0 { + t.Fatalf("expected no k-mer for k=0") + } +} + +func TestIterConstrainedKmersEarlyStop(t *testing.T) { + n := 0 + for range IterConstrainedKmers(10, NoKmerConstraints()) { + n++ + if n == 5 { + break + } + } + if n != 5 { + t.Fatalf("early stop failed: %d", n) + } +} diff --git a/pkg/obikmer/kmer_distance.go b/pkg/obikmer/kmer_distance.go new file mode 100644 index 0000000..8d17644 --- /dev/null +++ b/pkg/obikmer/kmer_distance.go @@ -0,0 +1,104 @@ +package obikmer + +import "math/bits" + +// KmerLCSResult is the result of a k-mer LCS alignment: the length of the +// Longest Common Subsequence and the length of the shortest alignment +// achieving it (same semantics as obialign.FastLCSScore). +type KmerLCSResult struct { + LCS int + AlignLen int +} + +// Distance returns AlignLen - LCS, the LCS-mode edit distance. +func (r KmerLCSResult) Distance() int { + return r.AlignLen - r.LCS +} + +// HammingDistance returns the number of positions at which the two k-mers +// differ. Both k-mers are encoded with 2 bits per base (see EncodeKmer); +// bits above the 2*k lowest ones are ignored. k must be between 0 and 32. +func HammingDistance(a, b uint64, k int) int { + diff := a ^ b + diff = (diff | diff>>1) & 0x5555555555555555 + diff &= ^uint64(0) >> (64 - 2*k) + return bits.OnesCount64(diff) +} + +// KmerLCS aligns two k-mers of the same size in LCS mode (match +1, mismatch +// and indel 0) and returns the LCS length together with the length of the +// shortest alignment among those reaching it, so that the distance is +// AlignLen - LCS. +// +// It is the scalar reference implementation, straightforward on purpose: +// it is the oracle against which the vectorised versions are tested. +// It returns {-1, -1} if k is not between 1 and 31. +func KmerLCS(a, b uint64, k int) KmerLCSResult { + if k < 1 || k > 31 { + return KmerLCSResult{-1, -1} + } + + var ba, bb [32]uint8 + for i := k - 1; i >= 0; i-- { + ba[i] = uint8(a & 3) + bb[i] = uint8(b & 3) + a >>= 2 + b >>= 2 + } + + // One DP row: best (score, alignment length) per column. + var score, length [32]int + for j := 0; j <= k; j++ { + length[j] = j + } + + better := func(s1, l1, s2, l2 int) bool { + return s1 > s2 || (s1 == s2 && l1 < l2) + } + + for i := 1; i <= k; i++ { + diagS, diagL := score[0], length[0] + score[0], length[0] = 0, i + + for j := 1; j <= k; j++ { + upS, upL := score[j], length[j] + + bestS, bestL := diagS, diagL+1 + if ba[j-1] == bb[i-1] { + bestS++ + } + if better(upS, upL+1, bestS, bestL) { + bestS, bestL = upS, upL+1 + } + if better(score[j-1], length[j-1]+1, bestS, bestL) { + bestS, bestL = score[j-1], length[j-1]+1 + } + + diagS, diagL = upS, upL + score[j], length[j] = bestS, bestL + } + } + + return KmerLCSResult{score[k], length[k]} +} + +// KmerLCSBatch computes KmerLCS(query, partners[i], k) for every partner and +// stores it in out[i]. out must be at least as long as partners. +// +// When built with GOEXPERIMENT=simd the partners are aligned several at a +// time, one per SIMD lane; otherwise the scalar reference is used. +// It returns false, leaving out untouched, if k is not between 1 and 31 or +// if out is too short. +func KmerLCSBatch(query uint64, partners []uint64, k int, out []KmerLCSResult) bool { + if k < 1 || k > 31 || len(out) < len(partners) { + return false + } + kmerLCSBatch(query, partners, k, out) + return true +} + +func kmerLCSBatchScalar(query uint64, partners []uint64, k int, out []KmerLCSResult) { + for i, p := range partners { + out[i] = KmerLCS(query, p, k) + } +} diff --git a/pkg/obikmer/kmer_distance_test.go b/pkg/obikmer/kmer_distance_test.go new file mode 100644 index 0000000..a46835a --- /dev/null +++ b/pkg/obikmer/kmer_distance_test.go @@ -0,0 +1,112 @@ +package obikmer + +import ( + "math/rand" + "testing" +) + +func randomKmer(rng *rand.Rand, k int) uint64 { + return rng.Uint64() & (uint64(1)<<(2*uint(k)) - 1) +} + +func TestHammingDistance(t *testing.T) { + rng := rand.New(rand.NewSource(1)) + for _, k := range []int{1, 2, 7, 21, 31} { + for n := 0; n < 200; n++ { + a, b := randomKmer(rng, k), randomKmer(rng, k) + sa, sb := DecodeKmer(a, k, nil), DecodeKmer(b, k, nil) + want := 0 + for i := range sa { + if sa[i] != sb[i] { + want++ + } + } + if got := HammingDistance(a, b, k); got != want { + t.Fatalf("k=%d %s/%s: got %d, want %d", k, sa, sb, got, want) + } + } + } +} + +func TestHammingDistanceIgnoresHighBits(t *testing.T) { + if got := HammingDistance(0xFF00, 0x0000, 4); got != 0 { + t.Fatalf("bits above 2k must be ignored, got %d", got) + } +} + +func TestKmerLCSKnown(t *testing.T) { + enc := func(s string) uint64 { return EncodeKmer([]byte(s), len(s)) } + cases := []struct { + a, b string + want KmerLCSResult + }{ + {"ACGT", "ACGT", KmerLCSResult{4, 4}}, + {"AAAA", "CCCC", KmerLCSResult{0, 4}}, + {"ACGT", "ACGA", KmerLCSResult{3, 4}}, + {"ACGT", "CGTA", KmerLCSResult{3, 5}}, // one shift: 1 indel on each side + } + for _, c := range cases { + if got := KmerLCS(enc(c.a), enc(c.b), 4); got != c.want { + t.Errorf("%s/%s: got %+v, want %+v", c.a, c.b, got, c.want) + } + } + if got := KmerLCS(0, 0, 32); got.LCS != -1 { + t.Errorf("k=32 must be rejected, got %+v", got) + } +} + +// The batch implementation (SIMD when built with GOEXPERIMENT=simd) must agree +// with the scalar reference for all k and for batch sizes around the lane count. +func TestKmerLCSBatchMatchesReference(t *testing.T) { + rng := rand.New(rand.NewSource(2)) + for _, k := range []int{1, 2, 3, 5, 8, 16, 21, 30, 31} { + for _, n := range []int{0, 1, 3, 7, 8, 9, 16, 17, 100} { + query := randomKmer(rng, k) + partners := make([]uint64, n) + for i := range partners { + partners[i] = randomKmer(rng, k) + } + if n > 2 { + partners[0] = query // identical + } + out := make([]KmerLCSResult, n) + if !KmerLCSBatch(query, partners, k, out) { + t.Fatalf("k=%d n=%d: batch refused", k, n) + } + for i, p := range partners { + if want := KmerLCS(query, p, k); out[i] != want { + t.Fatalf("k=%d n=%d partner %d: got %+v, want %+v", k, n, i, out[i], want) + } + } + } + } +} + +func TestKmerLCSBatchRejectsBadInput(t *testing.T) { + if KmerLCSBatch(0, []uint64{1}, 0, make([]KmerLCSResult, 1)) { + t.Fatal("k=0 must be rejected") + } + if KmerLCSBatch(0, []uint64{1, 2}, 4, make([]KmerLCSResult, 1)) { + t.Fatal("short output must be rejected") + } +} + +func benchBatch(b *testing.B, k int, f func(uint64, []uint64, int, []KmerLCSResult)) { + rng := rand.New(rand.NewSource(3)) + partners := make([]uint64, 1024) + for i := range partners { + partners[i] = randomKmer(rng, k) + } + out := make([]KmerLCSResult, len(partners)) + query := randomKmer(rng, k) + b.ResetTimer() + for b.Loop() { + f(query, partners, k, out) + } + b.ReportMetric(float64(b.Elapsed().Nanoseconds())/float64(b.N*len(partners)), "ns/pair") +} + +func BenchmarkKmerLCSScalar21(b *testing.B) { benchBatch(b, 21, kmerLCSBatchScalar) } +func BenchmarkKmerLCSBatch21(b *testing.B) { benchBatch(b, 21, kmerLCSBatch) } +func BenchmarkKmerLCSScalar11(b *testing.B) { benchBatch(b, 11, kmerLCSBatchScalar) } +func BenchmarkKmerLCSBatch11(b *testing.B) { benchBatch(b, 11, kmerLCSBatch) } diff --git a/pkg/obikmer/kmer_lcs_archsimd_arm64.go b/pkg/obikmer/kmer_lcs_archsimd_arm64.go new file mode 100644 index 0000000..3438d24 --- /dev/null +++ b/pkg/obikmer/kmer_lcs_archsimd_arm64.go @@ -0,0 +1,62 @@ +//go:build goexperiment.simd && arm64 + +package obikmer + +import "simd/archsimd" + +// neonLanes is the number of 16-bit lanes of a NEON 128-bit register. +const neonLanes = 8 + +// kmerLCSBatchNEON is kmerLCSBatch written directly on archsimd NEON types +// (see kmer_lcs_simd.go for the cell encoding and recurrence). +func kmerLCSBatchNEON(query uint64, partners []uint64, k int, out []KmerLCSResult) { + var qbase [32]uint16 + for i := 0; i < k; i++ { + qbase[i] = uint16(query>>(2*uint(k-1-i))) & 3 + } + + pbase := make([]uint16, k*neonLanes) + row := make([]uint16, (k+1)*neonLanes) + one := archsimd.BroadcastUint16x8(1) + bonus := archsimd.BroadcastUint16x8(lcsScoreUnit) + + for start := 0; start < len(partners); start += neonLanes { + chunk := partners[start:min(start+neonLanes, len(partners))] + + for lane, p := range chunk { + for j := 0; j < k; j++ { + pbase[j*neonLanes+lane] = uint16(p>>(2*uint(k-1-j))) & 3 + } + } + + for j := 0; j <= k; j++ { + archsimd.BroadcastUint16x8(uint16(lcsLenBase - j)).Store(row[j*neonLanes:]) + } + + for i := 1; i <= k; i++ { + qv := archsimd.BroadcastUint16x8(qbase[i-1]) + diag := archsimd.LoadUint16x8(row) + left := archsimd.BroadcastUint16x8(uint16(lcsLenBase - i)) + left.Store(row) + + for j := 1; j <= k; j++ { + up := archsimd.LoadUint16x8(row[j*neonLanes:]) + pv := archsimd.LoadUint16x8(pbase[(j-1)*neonLanes:]) + match := qv.Equal(pv) + best := diag.Add(bonus.Masked(match)).Max(up).Max(left).Sub(one) + best.Store(row[j*neonLanes:]) + diag = up + left = best + } + } + + last := row[k*neonLanes:] + for lane := range chunk { + v := int(last[lane]) + out[start+lane] = KmerLCSResult{ + LCS: v / lcsScoreUnit, + AlignLen: lcsLenBase - v%lcsScoreUnit, + } + } + } +} diff --git a/pkg/obikmer/kmer_lcs_archsimd_arm64_test.go b/pkg/obikmer/kmer_lcs_archsimd_arm64_test.go new file mode 100644 index 0000000..1210b94 --- /dev/null +++ b/pkg/obikmer/kmer_lcs_archsimd_arm64_test.go @@ -0,0 +1,31 @@ +//go:build goexperiment.simd && arm64 + +package obikmer + +import ( + "math/rand" + "testing" +) + +func TestKmerLCSBatchNEONMatchesReference(t *testing.T) { + rng := rand.New(rand.NewSource(5)) + for _, k := range []int{1, 2, 5, 11, 21, 31} { + for _, n := range []int{0, 1, 7, 8, 9, 100} { + query := randomKmer(rng, k) + partners := make([]uint64, n) + for i := range partners { + partners[i] = randomKmer(rng, k) + } + out := make([]KmerLCSResult, n) + kmerLCSBatchNEON(query, partners, k, out) + for i, p := range partners { + if want := KmerLCS(query, p, k); out[i] != want { + t.Fatalf("k=%d n=%d partner %d: got %+v, want %+v", k, n, i, out[i], want) + } + } + } + } +} + +func BenchmarkKmerLCSNEON21(b *testing.B) { benchBatch(b, 21, kmerLCSBatchNEON) } +func BenchmarkKmerLCSNEON11(b *testing.B) { benchBatch(b, 11, kmerLCSBatchNEON) } diff --git a/pkg/obikmer/kmer_lcs_obialign_test.go b/pkg/obikmer/kmer_lcs_obialign_test.go new file mode 100644 index 0000000..3a5eba6 --- /dev/null +++ b/pkg/obikmer/kmer_lcs_obialign_test.go @@ -0,0 +1,29 @@ +package obikmer_test + +import ( + "math/rand" + "testing" + + "git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obialign" + "git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obikmer" + "git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiseq" +) + +// KmerLCS must reproduce obialign.FastLCSScore on the decoded sequences. +func TestKmerLCSMatchesFastLCSScore(t *testing.T) { + rng := rand.New(rand.NewSource(4)) + for _, k := range []int{1, 2, 4, 9, 15, 21, 31} { + mask := uint64(1)<<(2*uint(k)) - 1 + for n := 0; n < 300; n++ { + a, b := rng.Uint64()&mask, rng.Uint64()&mask + sa := obiseq.NewBioSequence("a", obikmer.DecodeKmer(a, k, nil), "") + sb := obiseq.NewBioSequence("b", obikmer.DecodeKmer(b, k, nil), "") + lcs, alen := obialign.FastLCSScore(sa, sb, -1, nil) + got := obikmer.KmerLCS(a, b, k) + if got.LCS != lcs || got.AlignLen != alen { + t.Fatalf("k=%d %s/%s: KmerLCS=%+v, FastLCSScore=(%d,%d)", + k, sa.String(), sb.String(), got, lcs, alen) + } + } + } +} diff --git a/pkg/obikmer/kmer_lcs_scalar.go b/pkg/obikmer/kmer_lcs_scalar.go new file mode 100644 index 0000000..831a131 --- /dev/null +++ b/pkg/obikmer/kmer_lcs_scalar.go @@ -0,0 +1,7 @@ +//go:build !goexperiment.simd + +package obikmer + +func kmerLCSBatch(query uint64, partners []uint64, k int, out []KmerLCSResult) { + kmerLCSBatchScalar(query, partners, k, out) +} diff --git a/pkg/obikmer/kmer_lcs_simd.go b/pkg/obikmer/kmer_lcs_simd.go new file mode 100644 index 0000000..6127c62 --- /dev/null +++ b/pkg/obikmer/kmer_lcs_simd.go @@ -0,0 +1,80 @@ +//go:build goexperiment.simd + +package obikmer + +import "simd" + +// Each lane of a vector holds the state of one (query, partner) alignment as +// +// score<<6 | (63 - alignment length) +// +// so that taking the maximum of two cells ranks them by score first and, for +// equal scores, by shortest alignment. With k <= 31 the score is <= 31 and the +// alignment length <= 62, hence the value fits in 11 bits and the low field +// is never exhausted. The three DP moves then all reduce to +// +// cell(i,j) = max(diag + 64*match, up, left) - 1 +// +// (every move lengthens the alignment by one; only a diagonal match adds 1 +// to the score). +const ( + lcsScoreUnit = 64 + lcsLenBase = 63 +) + +func kmerLCSBatch(query uint64, partners []uint64, k int, out []KmerLCSResult) { + var probe simd.Uint16s + lanes := probe.Len() + + var qbase [32]uint16 + for i := 0; i < k; i++ { + qbase[i] = uint16(query>>(2*uint(k-1-i))) & 3 + } + + // partner bases, transposed: pbase[j*lanes+lane] is base j of the partner in `lane` + pbase := make([]uint16, k*lanes) + // one DP row, one vector per column + row := make([]uint16, (k+1)*lanes) + one := simd.BroadcastUint16s(1) + bonus := simd.BroadcastUint16s(lcsScoreUnit) + + for start := 0; start < len(partners); start += lanes { + chunk := partners[start:min(start+lanes, len(partners))] + + for lane, p := range chunk { + for j := 0; j < k; j++ { + pbase[j*lanes+lane] = uint16(p>>(2*uint(k-1-j))) & 3 + } + } + + for j := 0; j <= k; j++ { + simd.BroadcastUint16s(uint16(lcsLenBase - j)).Store(row[j*lanes:]) + } + + for i := 1; i <= k; i++ { + qv := simd.BroadcastUint16s(qbase[i-1]) + diag := simd.LoadUint16s(row) + left := simd.BroadcastUint16s(uint16(lcsLenBase - i)) + left.Store(row) + + for j := 1; j <= k; j++ { + up := simd.LoadUint16s(row[j*lanes:]) + pv := simd.LoadUint16s(pbase[(j-1)*lanes:]) + match := qv.Equal(pv) + best := diag.Add(bonus.Masked(match)).Max(up).Max(left).Sub(one) + best.Store(row[j*lanes:]) + diag = up + left = best + } + } + + last := row[k*lanes:] + for lane := range chunk { + v := int(last[lane]) + out[start+lane] = KmerLCSResult{ + LCS: v / lcsScoreUnit, + AlignLen: lcsLenBase - v%lcsScoreUnit, + } + } + } +} diff --git a/pkg/obitools/obicsv/options.go b/pkg/obitools/obicsv/options.go index 1928fad..395f446 100644 --- a/pkg/obitools/obicsv/options.go +++ b/pkg/obitools/obicsv/options.go @@ -93,11 +93,6 @@ func CLIPrintTaxon() bool { _outputTaxon = true } - if i := obiutils.LookFor(_keepOnly, "scientific_name"); i >= 0 { - _keepOnly = obiutils.RemoveIndex(_keepOnly, i) - _outputTaxon = true - } - return _outputTaxon } func CLIPrintQuality() bool { @@ -131,11 +126,6 @@ func CLIToBeKeptAttributes() []string { _outputTaxon = true } - if i := obiutils.LookFor(_keepOnly, "scientific_name"); i >= 0 { - _keepOnly = obiutils.RemoveIndex(_keepOnly, i) - _outputTaxon = true - } - return _keepOnly }