Files
obitools4/pkg/obikmer/enumkmer.go
T
Eric Coissac a0d8187fd6 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.
2026-11-03 09:23:37 +11:00

92 lines
2.3 KiB
Go

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