Implement k-mer distance metrics and constrained generation

Adds support for Hamming distance, constrained k-mer generation, and optimized Longest Common Subsequence (LCS) calculations using scalar, batch, and SIMD methods.

This includes new methods for k-mer distance calculation, structures for k-mer constraints, and SIMD-based batch processing for LCS to improve performance and correctness.
This commit is contained in:
Eric Coissac committed 2026-11-03 09:23:37 +11:00
1 parent 0012a55cad
commit a0d8187fd6
11 files changed
+601 -19

No files matched your search

-9
View File
@@ -5,7 +5,6 @@ import (
"container/heap" "container/heap"
"fmt" "fmt"
"math" "math"
"math/bits"
"os" "os"
"slices" "slices"
@@ -610,14 +609,6 @@ func (graph *DeBruijnGraph) WriteGml(filename string) error {
return err 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 type UInt64Heap []uint64
func (h UInt64Heap) Len() int { return len(h) } func (h UInt64Heap) Len() int { return len(h) }
+91
View File
@@ -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)
}
}
+85
View File
@@ -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)
}
}
+104
View File
@@ -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)
}
}
+112
View File
@@ -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) }
+62
View File
@@ -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,
}
}
}
}
@@ -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) }
+29
View File
@@ -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)
}
}
}
}
+7
View File
@@ -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)
}
+80
View File
@@ -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,
}
}
}
}
-10
View File
@@ -93,11 +93,6 @@ func CLIPrintTaxon() bool {
_outputTaxon = true _outputTaxon = true
} }
if i := obiutils.LookFor(_keepOnly, "scientific_name"); i >= 0 {
_keepOnly = obiutils.RemoveIndex(_keepOnly, i)
_outputTaxon = true
}
return _outputTaxon return _outputTaxon
} }
func CLIPrintQuality() bool { func CLIPrintQuality() bool {
@@ -131,11 +126,6 @@ func CLIToBeKeptAttributes() []string {
_outputTaxon = true _outputTaxon = true
} }
if i := obiutils.LookFor(_keepOnly, "scientific_name"); i >= 0 {
_keepOnly = obiutils.RemoveIndex(_keepOnly, i)
_outputTaxon = true
}
return _keepOnly return _keepOnly
} }