mirror of
https://github.com/metabarcoding/obitools4.git
synced 2026-11-09 09:24:41 +00:00
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:
1 parent
0012a55cad
commit
a0d8187fd6
11 files changed
+601
-19
No files matched your search
@@ -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) }
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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) }
|
||||||
@@ -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) }
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
@@ -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,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user