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"
|
||||
"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) }
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
|
||||
Reference in new issue
Block a user