Feature: Implement constrained k-mer generation and clique search

This commit introduces core functionality for oligonucleotide family tagging, constrained k-mer generation, and clique finding. It also includes updates to the Go language version and refactors the logic for k-mer dissimilarity metrics.

pkg/obikmer/enumkmer.go

@@ -8,18 +8,32 @@ import "iter"
 // a negative value means unlimited, 0 forbids the base altogether
 // (so MaxCount[2] = 0 restricts the alphabet to A, C, T).
 //
+// MaxGC is the maximum number of C or G bases; a negative value means
+// unlimited.
+//
 // MaxRun is the maximum number of consecutive identical bases;
-// a value <= 0 means unlimited.
+// a value <= 0 means unlimited. MinRun requires at least one run of that
+// length; a value <= 1 means no requirement.
+//
+// Accepted and Rejected are optional per-position sets of bases (4 bits masks,
+// bit b set when the base of code b is in the set, see ParseIUPACMasks).
+// When set, their length must be k. A k-mer is accepted only if each of its
+// bases belongs to the Accepted set of its position, and is refused if every
+// one of its bases belongs to the Rejected set of its position.
 //
 // The zero value forbids every base: use NoKmerConstraints as a starting point.
 type KmerConstraints struct {
 	MaxCount [4]int
+	MaxGC    int
 	MaxRun   int
+	MinRun   int
+	Accepted []uint8
+	Rejected []uint8
 }
 
 // NoKmerConstraints returns constraints accepting every k-mer.
 func NoKmerConstraints() KmerConstraints {
-	return KmerConstraints{MaxCount: [4]int{-1, -1, -1, -1}}
+	return KmerConstraints{MaxCount: [4]int{-1, -1, -1, -1}, MaxGC: -1}
 }
 
 // IterConstrainedKmers returns an iterator over all k-mers of size k
@@ -27,11 +41,14 @@ func NoKmerConstraints() KmerConstraints {
 // 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.
+// constraint is violated or cannot be met any more (MaxCount, MaxGC, MaxRun
+// and Accepted are monotone: a violating prefix cannot be repaired by
+// extension; MinRun is checked against the positions left), so the cost is
+// proportional to the number of explored prefixes, not to 4^k. Rejected is
+// checked on complete k-mers.
 //
-// k must be between 1 and 31, otherwise the iterator is empty.
+// k must be between 1 and 31, otherwise the iterator is empty. It is also
+// empty if Accepted or Rejected is set with a length different from k.
 //
 // Example:
 //
@@ -41,7 +58,7 @@ func NoKmerConstraints() KmerConstraints {
 //	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 {
+		if k < 1 || k > 31 || !c.validPositions(k) {
 			return
 		}
 
@@ -59,16 +76,27 @@ func IterConstrainedKmers(k int, c KmerConstraints) iter.Seq[uint64] {
 		}
 
 		var counts [4]int
+		gc := 0
 
-		var walk func(depth int, kmer uint64, last, run int) bool
-		walk = func(depth int, kmer uint64, last, run int) bool {
+		var walk func(depth int, kmer uint64, last, run, longest int) bool
+		walk = func(depth int, kmer uint64, last, run, longest int) bool {
 			if depth == k {
+				if c.Rejected != nil && matchesAllMasks(kmer, k, c.Rejected) {
+					return true
+				}
 				return yield(kmer)
 			}
 			for b := range 4 {
 				if c.MaxCount[b] >= 0 && counts[b] >= c.MaxCount[b] {
 					continue
 				}
+				if c.Accepted != nil && c.Accepted[depth]>>uint(b)&1 == 0 {
+					continue
+				}
+				isGC := b == 1 || b == 2
+				if isGC && c.MaxGC >= 0 && gc >= c.MaxGC {
+					continue
+				}
 				r := 1
 				if b == last {
 					r = run + 1
@@ -76,9 +104,19 @@ func IterConstrainedKmers(k int, c KmerConstraints) iter.Seq[uint64] {
 				if c.MaxRun > 0 && r > c.MaxRun {
 					continue
 				}
+				l := max(longest, r)
+				if c.MinRun > 1 && l < c.MinRun && r+(k-depth-1) < c.MinRun {
+					continue
+				}
 				counts[b]++
-				ok := walk(depth+1, kmer<<2|uint64(b), b, r)
+				if isGC {
+					gc++
+				}
+				ok := walk(depth+1, kmer<<2|uint64(b), b, r, l)
 				counts[b]--
+				if isGC {
+					gc--
+				}
 				if !ok {
 					return false
 				}
@@ -86,6 +124,6 @@ func IterConstrainedKmers(k int, c KmerConstraints) iter.Seq[uint64] {
 			return true
 		}
 
-		walk(0, 0, -1, 0)
+		walk(0, 0, -1, 0, 0)
 	}
 }

pkg/obikmer/kmer_constraints.go

@@ -0,0 +1,141 @@
+package obikmer
+
+import (
+	"fmt"
+
+	"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiseq"
+)
+
+// ParseIUPACMasks converts an IUPAC nucleotide pattern into one 4 bits mask per
+// position (bit 0 is A, bit 1 is C, bit 2 is G, bit 3 is T), the format used by
+// KmerConstraints.Accepted and Rejected.
+func ParseIUPACMasks(pattern string) ([]uint8, error) {
+	masks := make([]uint8, len(pattern))
+	for i := 0; i < len(pattern); i++ {
+		m := obiseq.IUPACNucMask(pattern[i])
+		if m == 0 {
+			return nil, fmt.Errorf("invalid IUPAC nucleotide code %q at position %d of pattern %q",
+				pattern[i], i+1, pattern)
+		}
+		masks[i] = m
+	}
+	return masks, nil
+}
+
+// RestrictAlphabet forbids every base that is not denoted by one of the IUPAC
+// codes of alphabet. For instance "ACT" (or "H") forbids G.
+func (c *KmerConstraints) RestrictAlphabet(alphabet string) error {
+	masks, err := ParseIUPACMasks(alphabet)
+	if err != nil {
+		return err
+	}
+	var allowed uint8
+	for _, m := range masks {
+		allowed |= m
+	}
+	if allowed == 0 {
+		return fmt.Errorf("empty alphabet")
+	}
+	for b := range 4 {
+		if allowed>>uint(b)&1 == 0 {
+			c.MaxCount[b] = 0
+		}
+	}
+	return nil
+}
+
+// validPositions tells whether Accepted and Rejected, when set, have length k.
+func (c *KmerConstraints) validPositions(k int) bool {
+	return (c.Accepted == nil || len(c.Accepted) == k) &&
+		(c.Rejected == nil || len(c.Rejected) == k)
+}
+
+// matchesAllMasks tells whether every base of the k-mer belongs to the set
+// given for its position.
+func matchesAllMasks(kmer uint64, k int, masks []uint8) bool {
+	for i := k - 1; i >= 0; i-- {
+		if masks[i]>>(kmer&3)&1 == 0 {
+			return false
+		}
+		kmer >>= 2
+	}
+	return true
+}
+
+// Accepts tells whether a k-mer of size k satisfies every constraint.
+// It is the predicate whose solutions IterConstrainedKmers enumerates, to be
+// used to filter k-mers coming from elsewhere.
+func (c *KmerConstraints) Accepts(kmer uint64, k int) bool {
+	if k < 1 || k > 31 || !c.validPositions(k) {
+		return false
+	}
+
+	var counts [4]int
+	gc, run, longest, last := 0, 0, 0, -1
+	for i := k - 1; i >= 0; i-- {
+		b := int(kmer>>(2*uint(i))) & 3
+		if c.Accepted != nil && c.Accepted[k-1-i]>>uint(b)&1 == 0 {
+			return false
+		}
+		counts[b]++
+		if c.MaxCount[b] >= 0 && counts[b] > c.MaxCount[b] {
+			return false
+		}
+		if b == 1 || b == 2 {
+			gc++
+			if c.MaxGC >= 0 && gc > c.MaxGC {
+				return false
+			}
+		}
+		if b == last {
+			run++
+		} else {
+			run = 1
+		}
+		last = b
+		longest = max(longest, run)
+		if c.MaxRun > 0 && run > c.MaxRun {
+			return false
+		}
+	}
+
+	if c.MinRun > 1 && longest < c.MinRun {
+		return false
+	}
+	if c.Rejected != nil && matchesAllMasks(kmer, k, c.Rejected) {
+		return false
+	}
+	return true
+}
+
+// HammingBound returns the sphere-packing (Hamming) upper bound on the number
+// of words of length k over an alphabet of alphabetSize letters that can be
+// pairwise at Hamming distance >= d. No larger family exists; a family of
+// that size may not exist either.
+func HammingBound(alphabetSize, k, d int) float64 {
+	if d <= 1 {
+		total := 1.0
+		for i := 0; i < k; i++ {
+			total *= float64(alphabetSize)
+		}
+		return total
+	}
+	if d%2 == 0 { // A(k, d) <= A(k-1, d-1)
+		k--
+		d--
+	}
+	if k < 1 {
+		return 1
+	}
+	t := (d - 1) / 2
+	total, ball, binom, pow := 1.0, 1.0, 1.0, 1.0
+	for i := 0; i < k; i++ {
+		total *= float64(alphabetSize)
+	}
+	for i := 1; i <= t && i <= k; i++ {
+		binom = binom * float64(k-i+1) / float64(i)
+		pow *= float64(alphabetSize - 1)
+		ball += binom * pow
+	}
+	return total / ball
+}

pkg/obikmer/kmer_constraints_test.go

@@ -0,0 +1,123 @@
+package obikmer
+
+import (
+	"math"
+	"math/rand"
+	"testing"
+)
+
+// IterConstrainedKmers must enumerate exactly the k-mers accepted by Accepts.
+func TestIterConstrainedKmersMatchesAccepts(t *testing.T) {
+	rng := rand.New(rand.NewSource(10))
+	patterns := []string{"", "nnnnnn", "yrnnnn", "cnnnnn", "annnnt"}
+
+	for trial := 0; trial < 60; trial++ {
+		k := 6
+		c := NoKmerConstraints()
+		for b := range 4 {
+			if rng.Intn(3) == 0 {
+				c.MaxCount[b] = rng.Intn(4)
+			}
+		}
+		if rng.Intn(2) == 0 {
+			c.MaxGC = rng.Intn(5)
+		}
+		if rng.Intn(2) == 0 {
+			c.MaxRun = 1 + rng.Intn(3)
+		}
+		if rng.Intn(2) == 0 {
+			c.MinRun = 2 + rng.Intn(3)
+		}
+		if p := patterns[rng.Intn(len(patterns))]; p != "" {
+			c.Accepted, _ = ParseIUPACMasks(p)
+		}
+		if p := patterns[rng.Intn(len(patterns))]; p != "" && rng.Intn(2) == 0 {
+			c.Rejected, _ = ParseIUPACMasks(p)
+		}
+
+		var want []uint64
+		for km := uint64(0); km < 1<<(2*uint(k)); km++ {
+			if c.Accepts(km, k) {
+				want = append(want, km)
+			}
+		}
+		got := collectConstrained(k, c)
+		if len(got) != len(want) {
+			t.Fatalf("trial %d %+v: got %d k-mers, want %d", trial, c, len(got), len(want))
+		}
+		for i := range got {
+			if got[i] != want[i] {
+				t.Fatalf("trial %d %+v: mismatch at %d: %d != %d", trial, c, i, got[i], want[i])
+			}
+		}
+	}
+}
+
+func TestConstraintsGCAndRuns(t *testing.T) {
+	enc := func(s string) uint64 { return EncodeKmer([]byte(s), len(s)) }
+
+	c := NoKmerConstraints()
+	c.MaxGC = 2
+	if !c.Accepts(enc("ACGTAA"), 6) || c.Accepts(enc("ACGCAA"), 6) {
+		t.Fatal("MaxGC not enforced")
+	}
+
+	c = NoKmerConstraints()
+	c.MinRun = 3
+	if !c.Accepts(enc("ACCCGT"), 6) || c.Accepts(enc("ACCGGT"), 6) {
+		t.Fatal("MinRun not enforced")
+	}
+
+	c = NoKmerConstraints()
+	c.Rejected, _ = ParseIUPACMasks("cnnnnn") // no C in first position
+	if c.Accepts(enc("CAAAAA"), 6) || !c.Accepts(enc("AAAAAA"), 6) {
+		t.Fatal("Rejected not enforced")
+	}
+}
+
+func TestRestrictAlphabetAndPatterns(t *testing.T) {
+	c := NoKmerConstraints()
+	if err := c.RestrictAlphabet("ACT"); err != nil {
+		t.Fatal(err)
+	}
+	if got := len(collectConstrained(4, c)); got != 81 {
+		t.Fatalf("expected 3^4 = 81 k-mers, got %d", got)
+	}
+	c = NoKmerConstraints()
+	if err := c.RestrictAlphabet("H"); err != nil { // H = A, C or T
+		t.Fatal(err)
+	}
+	if c.MaxCount != [4]int{-1, -1, 0, -1} {
+		t.Fatalf("unexpected counts %v", c.MaxCount)
+	}
+	if err := c.RestrictAlphabet("AX"); err == nil {
+		t.Fatal("invalid IUPAC code must be refused")
+	}
+	if _, err := ParseIUPACMasks("acgz"); err == nil {
+		t.Fatal("invalid IUPAC code must be refused")
+	}
+	// wrong pattern length
+	c = NoKmerConstraints()
+	c.Accepted, _ = ParseIUPACMasks("nnn")
+	if got := collectConstrained(4, c); len(got) != 0 {
+		t.Fatalf("pattern of wrong length must give nothing, got %d", len(got))
+	}
+}
+
+func TestHammingBound(t *testing.T) {
+	cases := []struct {
+		q, k, d int
+		want    float64
+	}{
+		{4, 7, 3, 16384.0 / 22},
+		{4, 8, 4, 16384.0 / 22}, // d even: reduces to (7, 3)
+		{4, 9, 5, 262144.0 / 352},
+		{3, 9, 3, 19683.0 / 19},
+		{4, 5, 1, 1024},
+	}
+	for _, c := range cases {
+		if got := HammingBound(c.q, c.k, c.d); math.Abs(got-c.want) > 1e-6 {
+			t.Errorf("HammingBound(%d,%d,%d) = %v, want %v", c.q, c.k, c.d, got, c.want)
+		}
+	}
+}

pkg/obikmer/kmer_dissimilarity.go

@@ -0,0 +1,84 @@
+package obikmer
+
+import (
+	"fmt"
+	"strings"
+	"sync"
+)
+
+// KmerMetric selects the distance used to compare two k-mers.
+type KmerMetric int
+
+const (
+	// HammingMetric counts the positions at which two k-mers differ.
+	HammingMetric KmerMetric = iota
+	// LCSMetric is the alignment length minus the LCS length (see KmerLCS).
+	LCSMetric
+)
+
+// ParseKmerMetric converts "hamming" or "lcs" (case insensitive) to a KmerMetric.
+func ParseKmerMetric(name string) (KmerMetric, error) {
+	switch strings.ToLower(name) {
+	case "hamming":
+		return HammingMetric, nil
+	case "lcs":
+		return LCSMetric, nil
+	}
+	return HammingMetric, fmt.Errorf("unknown distance type %q (expected hamming or lcs)", name)
+}
+
+func (m KmerMetric) String() string {
+	if m == LCSMetric {
+		return "lcs"
+	}
+	return "hamming"
+}
+
+type lcsScratch struct {
+	partners []uint64
+	results  []KmerLCSResult
+}
+
+// NewKmerDissimilarityFilter describes the graph whose vertices are kmers[i]
+// (k-mers of size k, 1 <= k <= 31) and where two vertices are linked when
+// their distance is at least dmin. The returned function has the signature of
+// obigraph.NeighborFilter: it appends to out[:0] the candidates (indices in
+// kmers) linked to v. It is safe for concurrent use.
+func NewKmerDissimilarityFilter(kmers []uint64, k int, metric KmerMetric, dmin int) func(v int32, candidates, out []int32) []int32 {
+	if metric == HammingMetric {
+		return func(v int32, candidates, out []int32) []int32 {
+			q := kmers[v]
+			out = out[:0]
+			for _, c := range candidates {
+				if HammingDistance(q, kmers[c], k) >= dmin {
+					out = append(out, c)
+				}
+			}
+			return out
+		}
+	}
+
+	pool := sync.Pool{New: func() any { return new(lcsScratch) }}
+	return func(v int32, candidates, out []int32) []int32 {
+		s := pool.Get().(*lcsScratch)
+		defer pool.Put(s)
+
+		s.partners = s.partners[:0]
+		for _, c := range candidates {
+			s.partners = append(s.partners, kmers[c])
+		}
+		if cap(s.results) < len(candidates) {
+			s.results = make([]KmerLCSResult, len(candidates))
+		}
+		s.results = s.results[:len(candidates)]
+		KmerLCSBatch(kmers[v], s.partners, k, s.results)
+
+		out = out[:0]
+		for i, c := range candidates {
+			if s.results[i].Distance() >= dmin {
+				out = append(out, c)
+			}
+		}
+		return out
+	}
+}

pkg/obikmer/kmer_dissimilarity_test.go

@@ -0,0 +1,113 @@
+package obikmer_test
+
+import (
+	"context"
+	"math/rand"
+	"testing"
+	"time"
+
+	"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obigraph"
+	"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obikmer"
+)
+
+func allKmers(k int) []uint64 {
+	kmers := make([]uint64, 0, 1<<(2*uint(k)))
+	for km := range obikmer.IterConstrainedKmers(k, obikmer.NoKmerConstraints()) {
+		kmers = append(kmers, km)
+	}
+	return kmers
+}
+
+func TestDissimilarityFilterMatchesDistances(t *testing.T) {
+	const k = 7
+	rng := rand.New(rand.NewSource(1))
+	kmers := make([]uint64, 300)
+	for i := range kmers {
+		kmers[i] = rng.Uint64() & (1<<(2*k) - 1)
+	}
+	candidates := make([]int32, 0, len(kmers))
+	for i := 1; i < len(kmers); i++ {
+		candidates = append(candidates, int32(i))
+	}
+
+	for _, metric := range []obikmer.KmerMetric{obikmer.HammingMetric, obikmer.LCSMetric} {
+		for _, dmin := range []int{1, 3, 5} {
+			filter := obikmer.NewKmerDissimilarityFilter(kmers, k, metric, dmin)
+			got := filter(0, candidates, nil)
+
+			var want []int32
+			for _, c := range candidates {
+				var d int
+				if metric == obikmer.HammingMetric {
+					d = obikmer.HammingDistance(kmers[0], kmers[c], k)
+				} else {
+					d = obikmer.KmerLCS(kmers[0], kmers[c], k).Distance()
+				}
+				if d >= dmin {
+					want = append(want, c)
+				}
+			}
+			if len(got) != len(want) {
+				t.Fatalf("%v d>=%d: got %d vertices, want %d", metric, dmin, len(got), len(want))
+			}
+			for i := range got {
+				if got[i] != want[i] {
+					t.Fatalf("%v d>=%d: mismatch at %d", metric, dmin, i)
+				}
+			}
+		}
+	}
+}
+
+// A family of 5-mers pairwise at Hamming distance >= 3, built with FindClique.
+func TestFindClique_HammingFamily(t *testing.T) {
+	const k, dmin, target = 5, 3, 32
+	kmers := allKmers(k)
+	filter := obikmer.NewKmerDissimilarityFilter(kmers, k, obikmer.HammingMetric, dmin)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+	r := obigraph.FindClique(ctx, len(kmers), filter, target, obigraph.CliqueOptions{Workers: 2})
+	if !r.Reached {
+		t.Fatalf("family of %d not found (best %d, exhausted %v)", target, len(r.Clique), r.Exhausted)
+	}
+	for i, a := range r.Clique {
+		for _, b := range r.Clique[i+1:] {
+			if obikmer.HammingDistance(kmers[a], kmers[b], k) < dmin {
+				t.Fatalf("%d and %d are too close", kmers[a], kmers[b])
+			}
+		}
+	}
+}
+
+func TestFindClique_LCSFamily(t *testing.T) {
+	const k, dmin, target = 6, 3, 20
+	kmers := allKmers(k)
+	filter := obikmer.NewKmerDissimilarityFilter(kmers, k, obikmer.LCSMetric, dmin)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+	r := obigraph.FindClique(ctx, len(kmers), filter, target, obigraph.CliqueOptions{Workers: 2})
+	if !r.Reached {
+		t.Fatalf("family of %d not found (best %d, exhausted %v)", target, len(r.Clique), r.Exhausted)
+	}
+	for i, a := range r.Clique {
+		for _, b := range r.Clique[i+1:] {
+			if obikmer.KmerLCS(kmers[a], kmers[b], k).Distance() < dmin {
+				t.Fatalf("%d and %d are too close", kmers[a], kmers[b])
+			}
+		}
+	}
+}
+
+func TestParseKmerMetric(t *testing.T) {
+	if m, err := obikmer.ParseKmerMetric("LCS"); err != nil || m != obikmer.LCSMetric {
+		t.Fatal("lcs not parsed")
+	}
+	if m, err := obikmer.ParseKmerMetric("hamming"); err != nil || m != obikmer.HammingMetric {
+		t.Fatal("hamming not parsed")
+	}
+	if _, err := obikmer.ParseKmerMetric("levenshtein"); err == nil {
+		t.Fatal("unknown metric must be refused")
+	}
+}