Implement core clustering logic and scoring features

This update introduces the core clustering algorithm, including parallel execution, dynamic block sizing, and SIMD-accelerated scoring methods. It also adds robust sequence preparation, k-mer analysis, and comprehensive test data fixtures for validation.

pkg/obialign/lcsband.go

@@ -40,7 +40,10 @@ const (
 type bandedLCSWorkspace struct {
 	aRev, bCode []uint32    // encoded bases (aRev is reversed), SIMD kernel only
 	rows        [3][]uint32 // antidiagonal (SIMD) or row (scalar) buffers
-	widened     int         // number of alignments that needed a second pass
+	aRev16      []uint16    // 16-bit encoded bases of the score-only kernel
+	bCode16     []uint16
+	rows16      [3][]uint16
+	widened     int // number of alignments that needed a second pass
 }
 
 var bandedLCSPool = sync.Pool{New: func() any { return new(bandedLCSWorkspace) }}
@@ -146,13 +149,7 @@ func bandedLCS(a, b []byte, maxError int, ws *bandedLCSWorkspace, useSIMD bool)
 	}
 	run := func(g int) (int, int) {
 		kmin, kmax := bandOffsets(lA, lB, g)
-		var key uint32
-		if simdOK {
-			key = bandedLCSSIMD(lA, lB, kmin, kmax, ws)
-		} else {
-			key = bandedLCSScalar(a, b, kmin, kmax, ws)
-		}
-		return int(key >> 16), int(lcsLenBase - key&0xFFFF)
+		return ws.align(a, b, kmin, kmax, simdOK)
 	}
 
 	if maxError < 0 {
@@ -173,3 +170,16 @@ func bandedLCS(a, b []byte, maxError int, ws *bandedLCSWorkspace, useSIMD bool)
 	}
 	return lcs, alilen
 }
+
+// align runs the banded DP on the offsets [kmin, kmax] (kmin <= 0 <= kmax) and
+// returns the best (lcs, alilen) of the band. The bases must have been encoded
+// with encodeBases when simdOK is set.
+func (ws *bandedLCSWorkspace) align(a, b []byte, kmin, kmax int, simdOK bool) (int, int) {
+	var key uint32
+	if simdOK {
+		key = bandedLCSSIMD(len(a), len(b), kmin, kmax, ws)
+	} else {
+		key = bandedLCSScalar(a, b, kmin, kmax, ws)
+	}
+	return int(key >> 16), int(lcsLenBase - key&0xFFFF)
+}

pkg/obialign/lcsband_nosimd.go

@@ -11,3 +11,7 @@ func encodeBases(ws *bandedLCSWorkspace, a, b []byte) {}
 func bandedLCSSIMD(lA, lB, kmin, kmax int, ws *bandedLCSWorkspace) uint32 {
 	return 0
 }
+
+func lcsScoreSIMD(a, b []byte, kmin, kmax, minLCS int, ws *bandedLCSWorkspace) int {
+	return -1
+}

pkg/obialign/lcsband_test.go

@@ -286,3 +286,75 @@ func TestFastLCSScoreWrapper(t *testing.T) {
 		t.Errorf("got (%d,%d), want (-1,-1)", s, l)
 	}
 }
+
+// TestBandedLCSAtLeast checks the exactness of the bound expressed on the
+// score: the optimum when it reaches minLCS, (-1,-1) otherwise.
+func TestBandedLCSAtLeast(t *testing.T) {
+	rng := rand.New(rand.NewSource(13))
+	var buf BandedLCSBuffer
+	for _, alpha := range testAlphabets {
+		for it := 0; it < 1500; it++ {
+			a := randomFrom(rng, alpha.letters, 1+rng.Intn(100))
+			b := mutateFrom(rng, alpha.letters, a, rng.Float64()*0.5)
+			if len(b) == 0 {
+				continue
+			}
+			os, ol := fullDP(a, b)
+			for _, minLCS := range []int{-3, 0, 1, os - 1, os, os + 1, len(b) / 2, len(b), len(b) + 1} {
+				ws, wl := os, ol
+				if os < minLCS {
+					ws, wl = -1, -1
+				}
+				for _, got := range [][2]int{func() [2]int { s, l, _ := buf.AtLeast(a, b, minLCS); return [2]int{s, l} }(),
+					func() [2]int { s, l, _ := BandedLCSAtLeast(b, a, minLCS); return [2]int{s, l} }()} {
+					if got[0] != ws || got[1] != wl {
+						t.Fatalf("%s minLCS=%d a=%q b=%q: got (%d,%d), want (%d,%d)",
+							alpha.name, minLCS, a, b, got[0], got[1], ws, wl)
+					}
+				}
+			}
+		}
+	}
+	if _, _, err := buf.AtLeast(make([]byte, MaxBandedLCSLength+1), []byte("a"), 0); err != ErrSequenceTooLong {
+		t.Errorf("too long: err = %v", err)
+	}
+}
+
+// TestBandedLCSScoreAtLeast checks the score-only entry point, with sequences
+// long enough for the early stop of the kernel to be exercised.
+func TestBandedLCSScoreAtLeast(t *testing.T) {
+	rng := rand.New(rand.NewSource(17))
+	var buf BandedLCSBuffer
+	check := func(name string, a, b []byte) {
+		os, _ := fullDP(a, b)
+		for _, minLCS := range []int{-1, 0, os - 3, os - 1, os, os + 1, os + 5, len(b) - 1, len(b), (len(b) * 9) / 10, len(b) + 1} {
+			want := os
+			if os < minLCS {
+				want = -1
+			}
+			got, err := buf.ScoreAtLeast(a, b, minLCS)
+			if err != nil || got != want {
+				t.Fatalf("%s minLCS=%d a=%q b=%q: got %d (%v), want %d", name, minLCS, a, b, got, err, want)
+			}
+		}
+	}
+	for _, alpha := range testAlphabets {
+		for it := 0; it < 800; it++ {
+			a := randomFrom(rng, alpha.letters, 1+rng.Intn(60))
+			b := mutateFrom(rng, alpha.letters, a, rng.Float64()*0.6)
+			if len(b) > 0 {
+				check(alpha.name, a, b)
+			}
+		}
+	}
+	for it := 0; it < 150; it++ {
+		a := randomFrom(rng, "ACGT", 100+rng.Intn(400))
+		b := mutateFrom(rng, "ACGT", a, rng.Float64()*0.3)
+		if it%3 == 0 {
+			b = randomFrom(rng, "ACGT", 50+rng.Intn(300))
+		}
+		if len(b) > 0 {
+			check("long", a, b)
+		}
+	}
+}

pkg/obialign/lcsmin.go

@@ -0,0 +1,102 @@
+package obialign
+
+// BandedLCSBuffer is a working memory that can be reused from one alignment to
+// the next to avoid any allocation. It must not be shared between goroutines.
+type BandedLCSBuffer struct {
+	ws bandedLCSWorkspace
+}
+
+// AtLeast is BandedLCSAtLeast on the buffer.
+func (b *BandedLCSBuffer) AtLeast(seqA, seqB []byte, minLCS int) (lcs, alilen int, err error) {
+	return atLeast(seqA, seqB, minLCS, &b.ws)
+}
+
+func atLeast(seqA, seqB []byte, minLCS int, ws *bandedLCSWorkspace) (lcs, alilen int, err error) {
+	if len(seqA) > MaxBandedLCSLength || len(seqB) > MaxBandedLCSLength {
+		return -1, -1, ErrSequenceTooLong
+	}
+	if len(seqA) < len(seqB) {
+		seqA, seqB = seqB, seqA
+	}
+	lcs, alilen = bandedLCSAtLeast(seqA, seqB, minLCS, ws, haveSIMD)
+	return lcs, alilen, nil
+}
+
+// BandedLCSAtLeast computes the LCS score of two nucleotide sequences and the
+// length of the shortest alignment achieving it, as BandedLCS does, but with
+// the bound expressed on the score: the result is the exact optimum when its
+// score is at least minLCS, and (-1, -1) otherwise.
+//
+// The DP only visits the diagonals an alignment of score minLCS can use (one
+// whose path makes at most lB-minLCS gaps on one side and lA-minLCS on the
+// other), so the cost decreases as minLCS grows. This is exact: every alignment
+// of score >= minLCS lies in that band, hence so does the optimum when it
+// reaches minLCS, and when the best path of the band is below minLCS so is
+// the global optimum.
+func BandedLCSAtLeast(seqA, seqB []byte, minLCS int) (lcs, alilen int, err error) {
+	ws := bandedLCSPool.Get().(*bandedLCSWorkspace)
+	lcs, alilen, err = atLeast(seqA, seqB, minLCS, ws)
+	bandedLCSPool.Put(ws)
+	return lcs, alilen, err
+}
+
+// bandedLCSAtLeast expects len(a) >= len(b), both within MaxBandedLCSLength.
+func bandedLCSAtLeast(a, b []byte, minLCS int, ws *bandedLCSWorkspace, useSIMD bool) (int, int) {
+	lA, lB := len(a), len(b)
+	minLCS = max(minLCS, 0)
+	if minLCS > lB {
+		return -1, -1
+	}
+	if lB == 0 {
+		return 0, lA
+	}
+	simdOK := useSIMD && allRegularIUPAC(a) && allRegularIUPAC(b)
+	kmin, kmax := -(lB - minLCS), lA-minLCS
+
+	if !simdOK {
+		lcs, alilen := ws.align(a, b, kmin, kmax, false)
+		if lcs < minLCS {
+			return -1, -1
+		}
+		return lcs, alilen
+	}
+
+	// The score alone is much cheaper than the score with the length of the
+	// alignment: the latter is only computed when the pair reaches minLCS.
+	if lcsScoreSIMD(a, b, kmin, kmax, minLCS, ws) < minLCS {
+		return -1, -1
+	}
+	encodeBases(ws, a, b)
+	return ws.align(a, b, kmin, kmax, true)
+}
+
+// ScoreAtLeast computes the LCS score of two nucleotide sequences when it is
+// at least minLCS, and returns -1 otherwise. It is much cheaper than AtLeast,
+// which also computes the length of the shortest alignment.
+func (b *BandedLCSBuffer) ScoreAtLeast(seqA, seqB []byte, minLCS int) (int, error) {
+	if len(seqA) > MaxBandedLCSLength || len(seqB) > MaxBandedLCSLength {
+		return -1, ErrSequenceTooLong
+	}
+	if len(seqA) < len(seqB) {
+		seqA, seqB = seqB, seqA
+	}
+	lA, lB := len(seqA), len(seqB)
+	minLCS = max(minLCS, 0)
+	if minLCS > lB {
+		return -1, nil
+	}
+	if lB == 0 {
+		return 0, nil
+	}
+	kmin, kmax := -(lB - minLCS), lA-minLCS
+	var lcs int
+	if haveSIMD && allRegularIUPAC(seqA) && allRegularIUPAC(seqB) {
+		lcs = lcsScoreSIMD(seqA, seqB, kmin, kmax, minLCS, &b.ws)
+	} else {
+		lcs, _ = b.ws.align(seqA, seqB, kmin, kmax, false)
+	}
+	if lcs < minLCS {
+		return -1, nil
+	}
+	return lcs, nil
+}

pkg/obialign/lcsscore_simd.go

@@ -0,0 +1,156 @@
+//go:build goexperiment.simd
+
+package obialign
+
+import "simd"
+
+func simdLanes16() int {
+	var probe simd.Uint16s
+	return probe.Len()
+}
+
+func growUint16(s []uint16, n int) []uint16 {
+	if cap(s) < n {
+		return make([]uint16, n)
+	}
+	return s[:n]
+}
+
+func encodeIUPAC16(dst []uint16, s []byte, rev bool) {
+	n := len(s)
+	clear(dst)
+	if rev {
+		for x, c := range s {
+			dst[n-1-x] = uint16(iupacMask[c])
+		}
+		return
+	}
+	for x, c := range s {
+		dst[x] = uint16(iupacMask[c])
+	}
+}
+
+// pruneEvery is the number of antidiagonals between two tests of the upper
+// bound that stops a hopeless alignment.
+const pruneEvery = 8
+
+// lcsScoreSIMD runs the same banded DP as bandedLCSSIMD on the score only,
+// with 16-bit lanes (twice the cells per instruction):
+//
+//	cell(i,j) = max(diag + match, up, left)
+//
+// A cell outside the band is 0, which never beats a real cell. The score of a
+// path ending at the cell (i, j) of an antidiagonal cannot exceed the score of
+// the cell plus min(lB-i, lA-j), the characters left to pair. Every path
+// crosses each antidiagonal inside the window, so when the best score of the
+// window plus the largest of these remainders is under minLCS the final score
+// is under minLCS: -1 is returned without computing the rest of the matrix.
+//
+// Requires lA >= lB >= 1, kmin <= 0 <= kmax, and bases with non-zero masks.
+// Returns the best score of the band, or -1 if it is proven to be under
+// minLCS.
+func lcsScoreSIMD(a, b []byte, kmin, kmax, minLCS int, ws *bandedLCSWorkspace) int {
+	lA, lB := len(a), len(b)
+	vl := simdLanes16()
+
+	ws.aRev16 = growUint16(ws.aRev16, lA+vl)
+	ws.bCode16 = growUint16(ws.bCode16, lB+vl)
+	encodeIUPAC16(ws.aRev16, a, true)
+	encodeIUPAC16(ws.bCode16, b, false)
+	aRev, bCode := ws.aRev16, ws.bCode16
+
+	size := lB + 2*vl + 2
+	for r := range ws.rows16 {
+		ws.rows16[r] = growUint16(ws.rows16[r], size)
+		clear(ws.rows16[r])
+	}
+	p2, p1, cur := ws.rows16[0], ws.rows16[1], ws.rows16[2]
+
+	// the first row and the first column hold a score of 0: nothing to set
+
+	zero := simd.BroadcastUint16s(0)
+	one := simd.BroadcastUint16s(1)
+
+	for t := 2; t <= lA+lB; t++ {
+		lo := max(1, t-lA, (t-kmax+1)>>1)
+		hi := min(lB, t-1, (t-kmin)>>1)
+
+		for i := lo; i <= hi; i += vl {
+			ca := simd.LoadUint16s(aRev[lA-t+i:])
+			cb := simd.LoadUint16s(bCode[i-1:])
+			inc := one.Masked(ca.And(cb).NotEqual(zero))
+			d := simd.LoadUint16s(p2[i-1:]).Add(inc)
+			u := simd.LoadUint16s(p1[i-1:])
+			l := simd.LoadUint16s(p1[i:])
+			d.Max(u).Max(l).Store(cur[i:])
+		}
+
+		zero.Store(cur[hi+1:])
+		if lo > 1 {
+			cur[lo-1] = 0
+		}
+
+		// A diagonal move skips an antidiagonal, but a path always visits
+		// one of two consecutive antidiagonals.
+		if t%pruneEvery == 0 && minLCS > 0 &&
+			scoreBound(p1, t-1, lA, lB, kmin, kmax) < minLCS &&
+			scoreBound(cur, t, lA, lB, kmin, kmax) < minLCS {
+			return -1
+		}
+
+		p2, p1, cur = p1, cur, p2
+	}
+	return int(p1[lB])
+}
+
+// iota16 is 0, 1, 2... to build the row index of the lanes.
+var iota16 = func() (t [64]uint16) {
+	for i := range t {
+		t[i] = uint16(i)
+	}
+	return t
+}()
+
+// scoreBound is an upper bound of the final score of the paths going through
+// the antidiagonal t, whose scores are in row. A path through the cell (i, j)
+// scores at most the score of the cell plus min(lB-i, lA-j), the characters
+// left to pair, so the bound is the maximum of this sum on the cells of the
+// window. The sum is computed by vectors on the full chunks of the window.
+func scoreBound(row []uint16, t, lA, lB, kmin, kmax int) int {
+	lo := max(1, t-lA, (t-kmax+1)>>1)
+	hi := min(lB, t-1, (t-kmin)>>1)
+
+	bound := 0
+	i := lo
+	if hi-lo >= simdLanes16() {
+		vl := simdLanes16()
+		idx := simd.LoadUint16s(iota16[:]).Add(simd.BroadcastUint16s(uint16(lo)))
+		step := simd.BroadcastUint16s(uint16(vl))
+		remB := simd.BroadcastUint16s(uint16(lB))
+		// lA-t+i is not negative on the window; the 16-bit arithmetic is modular
+		offA := simd.BroadcastUint16s(uint16(lA - t))
+		acc := simd.BroadcastUint16s(0)
+		for ; i+vl-1 <= hi; i += vl {
+			left := remB.Sub(idx).Min(idx.Add(offA))
+			acc = acc.Max(simd.LoadUint16s(row[i:]).Add(left))
+			idx = idx.Add(step)
+		}
+		var lanes [64]uint16
+		acc.Store(lanes[:])
+		for _, v := range lanes[:vl] {
+			bound = max(bound, int(v))
+		}
+	}
+	for ; i <= hi; i++ {
+		bound = max(bound, int(row[i])+min(lB-i, lA-t+i))
+	}
+
+	// the paths running along the first row or the first column have a score of 0
+	if t <= lA && t <= kmax {
+		bound = max(bound, min(lB, lA-t))
+	}
+	if t <= lB && -t >= kmin {
+		bound = max(bound, min(lB-t, lA))
+	}
+	return bound
+}