Refactor sequence processing for parallel loading and memory management

This commit introduces significant changes to the sequence processing pipeline, implementing parallel loading, k-mer budget control, explicit memory management, and refined logic for pair evaluation and resolution.
This commit is contained in:
Eric Coissac committed 2026-11-03 22:53:54 +11:00
1 parent d6dd39f878
commit 8cc02aa9c1
28 files changed
+1673 -382

No files matched your search

+47 -8
View File
@@ -24,9 +24,12 @@ import (
// 2. CountAtLeast processes the partitions in parallel: each one is loaded,
// expanded to (canonical k-mer, sequence id) pairs, sorted, and reduced.
//
// Memory during phase 2 is about 16 bytes per k-mer occurrence of the largest
// partitions being processed at the same time, so P must grow with the total
// input size (total bases / P is the occurrence count of an average partition).
// Memory during phase 2 is 16 bytes per k-mer occurrence of the partitions
// being processed at the same time. Their number is bounded by an occurrence
// budget (see SetOccurrenceBudget), so that the memory does not depend on how
// the k-mers happen to be spread over the partitions; P only has to be large
// enough for the largest partition to be small against the budget (total bases
// / P is the occurrence count of an average partition, see DefaultPartitions).
//
// Bases other than A, C, G, T, U break a sequence into independent runs: no
// k-mer ever contains an ambiguous base.
@@ -39,13 +42,16 @@ type SeqKmerCounter struct {
mu sync.Mutex
firstErr error
sealed bool
budget int // k-mer occurrences held at once in phase 2
}
type seqPartWriter struct {
mu sync.Mutex
w *bufio.Writer
file *os.File
path string
mu sync.Mutex
w *bufio.Writer
file *os.File
path string
kmers int // k-mer occurrences written, so that the partition is loaded in an array of the exact size
}
// SeqKmerCount is the result for one k-mer.
@@ -78,7 +84,7 @@ func NewSeqKmerCounter(parentDir string, k, m, partitions int) (*SeqKmerCounter,
return nil, err
}
c := &SeqKmerCounter{dir: dir, k: k, m: m, p: partitions}
c := &SeqKmerCounter{dir: dir, k: k, m: m, p: partitions, budget: DefaultOccurrenceBudget}
c.parts = make([]*seqPartWriter, partitions)
for i := range c.parts {
path := filepath.Join(dir, fmt.Sprintf("part_%04d.sqk", i))
@@ -92,6 +98,38 @@ func NewSeqKmerCounter(parentDir string, k, m, partitions int) (*SeqKmerCounter,
return c, nil
}
// DefaultOccurrenceBudget is the number of k-mer occurrences phase 2 holds at
// once by default: 32 million, 512 MB. A partition larger than the budget is
// processed alone. Measured on 16 cores, 526 million occurrences to count: with
// 32M the count is as fast as without any limit; 16M make it 15% slower, 8M
// 60% slower.
const DefaultOccurrenceBudget = 32 << 20
// SetOccurrenceBudget bounds the k-mer occurrences (16 bytes each) that
// CountAtLeast holds in memory at once. It must be called before CountAtLeast.
func (c *SeqKmerCounter) SetOccurrenceBudget(occurrences int) {
if occurrences > 0 {
c.budget = occurrences
}
}
// partitionOccurrences is the occurrence count aimed at for an average
// partition by DefaultPartitions (a partition of 1M occurrences is 16 MB).
const partitionOccurrences = 1 << 20
// DefaultPartitions returns a partition count suited to a total of kmers
// k-mer occurrences: 256 at least, and more, as a power of two up to 1024, when
// the average partition would otherwise exceed about a million occurrences.
// The largest partitions hold several times the average. The limit keeps the
// number of simultaneously open files reasonable.
func DefaultPartitions(kmers int) int {
p := 256
for p < 1024 && kmers/p > partitionOccurrences {
p *= 2
}
return p
}
// K returns the k-mer size.
func (c *SeqKmerCounter) K() int { return c.k }
@@ -148,6 +186,7 @@ func (c *SeqKmerCounter) addRun(id uint32, run []byte, scratch []byte) []byte {
pw := c.parts[int(sk.Minimizer%uint64(c.p))]
pw.mu.Lock()
_, err := pw.w.Write(scratch)
pw.kmers += len(sk.Sequence) - c.k + 1
pw.mu.Unlock()
if err != nil {
c.setErr(err)
+118
View File
@@ -0,0 +1,118 @@
package obikmer
import (
"math/rand"
"reflect"
"sync"
"testing"
"time"
)
// The occurrence budget only bounds the memory: any budget, down to one
// occurrence (every partition then goes alone), gives the same counts.
func TestSeqKmerCounterBudgetDoesNotChangeTheCounts(t *testing.T) {
rng := rand.New(rand.NewSource(7))
seqs := randomSeqs(rng, 40, 20, 500, true)
build := func(budget int) []SeqKmerCount {
c, err := NewSeqKmerCounter(t.TempDir(), 9, 4, 11)
if err != nil {
t.Fatal(err)
}
defer c.Close()
if budget > 0 {
c.SetOccurrenceBudget(budget)
}
for i, s := range seqs {
c.AddSequence(uint32(i), s, i%2 == 0)
}
total := 0
for _, pw := range c.parts {
total += pw.kmers
}
if total == 0 {
t.Fatal("no occurrence counted")
}
got, err := c.CountAtLeast(2)
if err != nil {
t.Fatal(err)
}
return got
}
want := build(0)
if len(want) == 0 {
t.Fatal("the test needs k-mers shared by several sequences")
}
for _, budget := range []int{1, 100, 5000} {
if got := build(budget); !reflect.DeepEqual(got, want) {
t.Errorf("budget %d changes the counts", budget)
}
}
}
// The occurrences announced for a partition are those that its file holds.
func TestSeqKmerCounterPartitionOccurrenceCount(t *testing.T) {
rng := rand.New(rand.NewSource(3))
seqs := randomSeqs(rng, 25, 10, 300, true)
c, err := NewSeqKmerCounter(t.TempDir(), 8, 3, 5)
if err != nil {
t.Fatal(err)
}
defer c.Close()
for i, s := range seqs {
c.AddSequence(uint32(i), s, i%3 == 0)
}
if err := c.seal(); err != nil {
t.Fatal(err)
}
for p, pw := range c.parts {
occ, err := c.loadPartition(p, nil)
if err != nil {
t.Fatal(err)
}
if len(occ) != pw.kmers {
t.Errorf("partition %d: %d occurrences loaded, %d counted when written", p, len(occ), pw.kmers)
}
}
}
func TestOccurrenceBudget(t *testing.T) {
b := newOccurrenceBudget(10)
if got := b.acquire(4); got != 4 {
t.Fatalf("acquire(4) = %d", got)
}
// 7 does not fit in the 6 left: it must wait for the release
done := make(chan int)
go func() { done <- b.acquire(7) }()
select {
case <-done:
t.Fatal("acquire(7) must wait")
case <-time.After(50 * time.Millisecond):
}
b.release(4)
if got := <-done; got != 7 {
t.Fatalf("acquire(7) = %d", got)
}
b.release(7)
// a request above the whole budget takes all of it, once everything is free
if got := b.acquire(1000); got != 10 {
t.Fatalf("an oversized request must take the whole budget, got %d", got)
}
var wg sync.WaitGroup
wg.Add(1)
go func() { defer wg.Done(); b.release(b.acquire(3)) }()
time.Sleep(20 * time.Millisecond)
b.release(10)
wg.Wait()
}
func TestDefaultPartitions(t *testing.T) {
for _, tc := range []struct{ kmers, want int }{
{0, 256}, {5_000_000, 256}, {256 << 20, 256}, {257 << 20, 512}, {526_000_000, 512},
{600 << 20, 1024}, {1 << 40, 1024},
} {
if got := DefaultPartitions(tc.kmers); got != tc.want {
t.Errorf("DefaultPartitions(%d) = %d, want %d", tc.kmers, got, tc.want)
}
}
}
+54 -2
View File
@@ -41,6 +41,16 @@ func (c *SeqKmerCounter) CountAtLeast(minSeq int) ([]SeqKmerCount, error) {
}
workers := min(obidefault.ParallelWorkers(), c.p)
// the largest partitions first: they are the hardest to fit in the budget
// and the last ones to finish, if they come late
order := make([]int, c.p)
for i := range order {
order[i] = i
}
slices.SortStableFunc(order, func(a, b int) int { return cmp.Compare(c.parts[b].kmers, c.parts[a].kmers) })
budget := newOccurrenceBudget(c.budget)
jobs := make(chan int)
results := make([][]SeqKmerCount, c.p)
@@ -52,13 +62,19 @@ func (c *SeqKmerCounter) CountAtLeast(minSeq int) ([]SeqKmerCount, error) {
wg.Add(1)
go func() {
defer wg.Done()
var occ []kmerOcc // reused between partitions
var occ []kmerOcc // reused between partitions of similar sizes
for p := range jobs {
n := c.parts[p].kmers
held := budget.acquire(n)
if cap(occ) < n || cap(occ) > 2*n+1024 {
occ = make([]kmerOcc, 0, n)
}
var err error
occ, err = c.loadPartition(p, occ[:0])
if err == nil {
results[p] = reduceOccurrences(occ, minSeq)
}
budget.release(held)
if err != nil {
errMu.Lock()
if firstErr == nil {
@@ -69,7 +85,7 @@ func (c *SeqKmerCounter) CountAtLeast(minSeq int) ([]SeqKmerCount, error) {
}
}()
}
for p := range c.p {
for _, p := range order {
jobs <- p
}
close(jobs)
@@ -153,3 +169,39 @@ func reduceOccurrences(occ []kmerOcc, minSeq int) []SeqKmerCount {
}
return out
}
// occurrenceBudget bounds the k-mer occurrences held in memory at once by the
// goroutines that load the partitions.
type occurrenceBudget struct {
mu sync.Mutex
freed *sync.Cond
total int
free int
}
func newOccurrenceBudget(total int) *occurrenceBudget {
b := &occurrenceBudget{total: total, free: total}
b.freed = sync.NewCond(&b.mu)
return b
}
// acquire waits until n occurrences fit in the budget and takes them; a request
// above the whole budget waits for the budget to be free and takes all of it.
// It returns what was taken, to give back with release.
func (b *occurrenceBudget) acquire(n int) int {
n = min(n, b.total)
b.mu.Lock()
for b.free < n {
b.freed.Wait()
}
b.free -= n
b.mu.Unlock()
return n
}
func (b *occurrenceBudget) release(n int) {
b.mu.Lock()
b.free += n
b.mu.Unlock()
b.freed.Broadcast()
}
+159 -98
View File
@@ -1,23 +1,13 @@
package obiprimer
import (
"cmp"
"math"
"slices"
"sync"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obidefault"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obithermo"
)
// ampRecord is an amplicon of a retained pair, in one sequence.
type ampRecord struct {
hash uint64 // the barcode: see ampHash
pair int32
taxon int32 // taxon of the sequence at the chosen rank, -1 when it has none
role Role
}
// refHit is an amplicon of a retained pair on the reference sequence.
type refHit struct {
up, down seqMatch
@@ -48,47 +38,142 @@ func ampHash(raw []byte, from, to int, reverse bool) uint64 {
return h
}
// amplicons holds what the second phase gathers on the retained pairs.
type amplicons struct {
parts [][]ampRecord // one slice per worker, each sorted by pair
minTm1 []float32 // by pair: lowest Tm of the first primer over the example sites
minTm2 []float32
refs map[int32]refHit
}
// eachAmplicon runs forEachOccurrence over the sequences of the corpus (the
// examples only, or all of them) in parallel, and calls visit for every
// amplicon of a pair of index. Every worker gets its own visit function, from
// newVisit, and done is called once the worker has no more sequence.
//
// Without a plan the sequences are handed to the workers as they get free, and
// the assignment is recorded in record. With one, every worker scans the
// sequences that the plan gave it.
func (ps *primerSet) eachAmplicon(c *Corpus, p PairParams, index map[pairKey]int32, examplesOnly bool, plan *ampPlan, record []uint16,
newVisit func(worker int) (visit func(s int, raw []byte, o occurrence, pair int32), done func())) {
// pair returns the amplicons of one retained pair, as one slice per worker.
func (a *amplicons) pair(i int32) [][]ampRecord {
out := make([][]ampRecord, 0, len(a.parts))
byPair := func(r ampRecord, p int32) int { return cmp.Compare(r.pair, p) }
for _, part := range a.parts {
lo, _ := slices.BinarySearchFunc(part, i, byPair)
hi, _ := slices.BinarySearchFunc(part, i+1, byPair)
if hi > lo {
out = append(out, part[lo:hi])
}
workers := obidefault.ParallelWorkers()
// only the matches of the primers of the pairs of index can make an amplicon
// of these pairs, and dropping the others spares most of the comparisons
wanted := make([]bool, len(ps.primers))
for k := range index {
wanted[k.p1], wanted[k.p2] = true, true
}
return out
type scratch struct {
ms []seqMatch
bases []byte
}
scan := func(sc *scratch, visit func(int, []byte, occurrence, int32), s int) {
if examplesOnly && c.Roles[s] != Example {
return
}
ms := sc.ms[:0]
for _, m := range ps.perSeq[s] {
if wanted[m.primer] {
ms = append(ms, m)
}
}
sc.ms = ms
if len(ms) < 2 {
return // an amplicon needs two primer sites
}
var raw []byte
raw, sc.bases = c.bases(s, sc.bases)
ps.forEachOccurrenceIn(raw, ms, p, func(o occurrence) {
if pi, ok := index[o.key]; ok {
visit(s, raw, o, pi)
}
})
}
var wg sync.WaitGroup
if plan != nil {
mine := make([][]int, workers)
for s, w := range plan.owner {
mine[w] = append(mine[w], s)
}
for w := range workers {
wg.Add(1)
go func() {
defer wg.Done()
visit, done := newVisit(w)
var sc scratch
for _, s := range mine[w] {
scan(&sc, visit, s)
}
if done != nil {
done()
}
}()
}
wg.Wait()
return
}
jobs := make(chan int)
for w := range workers {
wg.Add(1)
go func() {
defer wg.Done()
visit, done := newVisit(w)
var sc scratch
for s := range jobs {
record[s] = uint16(w)
scan(&sc, visit, s)
}
if done != nil {
done()
}
}()
}
for s := range c.Roles {
jobs <- s
}
close(jobs)
wg.Wait()
}
// collectOptions selects what collectAmplicons gathers.
// planAmplicons counts the amplicons of every pair of index in the sequences
// of the corpus (all of them, or the examples only), by worker.
func (ps *primerSet) planAmplicons(c *Corpus, p PairParams, index map[pairKey]int32, examplesOnly bool) *ampPlan {
workers := obidefault.ParallelWorkers()
plan := &ampPlan{owner: make([]uint16, c.Len()), count: make([][]uint32, workers)}
for w := range plan.count {
plan.count[w] = make([]uint32, len(index))
}
ps.eachAmplicon(c, p, index, examplesOnly, nil, plan.owner, func(w int) (func(int, []byte, occurrence, int32), func()) {
row := plan.count[w]
return func(_ int, _ []byte, _ occurrence, pi int32) { row[pi]++ }, nil
})
return plan
}
// ampSites is what the gathering of the amplicons of the targets learns besides
// the records.
type ampSites struct {
minTm1 []float32 // by pair: lowest Tm of the first primer over the example sites
minTm2 []float32
refs map[int32]refHit // by pair: the amplicon on the reference sequence
}
// collectOptions selects what collectAmplicons gathers besides the records.
type collectOptions struct {
examplesOnly bool // skip the counterexamples and the other sequences
model *obithermo.Model // melting temperatures of the example sites (nil: none)
reference bool // record the amplicons on the reference sequence
}
// collectAmplicons gathers, in parallel, one record per amplicon of a retained
// pair, in the sequences of the corpus (all of them, or the examples only), and
// the lowest melting temperatures over the example sites. The records are not
// copied nor regrouped: each worker keeps its own, sorted by pair.
func (ps *primerSet) collectAmplicons(c *Corpus, taxa *Taxa, p PairParams, index map[pairKey]int32, opt collectOptions) *amplicons {
// collectAmplicons writes into store, in parallel, one record per amplicon of
// the pairs of index (whose values are the pair numbers in the store) and
// returns the lowest melting temperatures over the example sites. The store
// was sized from the plan, which says which worker scans which sequence; the
// batch only holds the pairs of index, the plan more.
func (ps *primerSet) collectAmplicons(c *Corpus, taxa *Taxa, p PairParams, index map[pairKey]int32,
plan *ampPlan, store *ampStore, opt collectOptions) *ampSites {
model := opt.model
lp := p.PrimerLength
workers := obidefault.ParallelWorkers()
npairs := len(index)
res := &amplicons{
parts: make([][]ampRecord, workers),
res := &ampSites{
minTm1: make([]float32, npairs),
minTm2: make([]float32, npairs),
refs: map[int32]refHit{}, // written by the single worker that reads the reference
@@ -106,63 +191,43 @@ func (ps *primerSet) collectAmplicons(c *Corpus, taxa *Taxa, p PairParams, index
oriented[i][1] = []byte(orientedWord(pr.Word, lp, true))
}
jobs := make(chan int)
var wg sync.WaitGroup
for w := range workers {
wg.Add(1)
go func() {
defer wg.Done()
site1, site2 := make([]byte, lp), make([]byte, lp)
lowTm1, lowTm2 := map[int32]float32{}, map[int32]float32{}
for s := range jobs {
raw := c.Seqs[s].Sequence()
n := len(raw)
role := c.Roles[s]
if opt.examplesOnly && role != Example {
continue
ps.eachAmplicon(c, p, index, opt.examplesOnly, plan, nil, func(w int) (func(int, []byte, occurrence, int32), func()) {
site1, site2 := make([]byte, lp), make([]byte, lp)
lowTm1, lowTm2 := map[int32]float32{}, map[int32]float32{}
visit := func(s int, raw []byte, o occurrence, pi int32) {
n := len(raw)
role := c.Roles[s]
if opt.reference && p.HasReference && s == p.Reference {
if old, seen := res.refs[pi]; !seen || o.up.pos() < old.up.pos() {
res.refs[pi] = refHit{o.up, o.down, o.swapped}
}
ps.forEachOccurrence(c, s, p, func(o occurrence) {
pi, ok := index[o.key]
if !ok {
return
}
if opt.reference && p.HasReference && s == p.Reference {
if old, seen := res.refs[pi]; !seen || o.up.pos < old.up.pos {
res.refs[pi] = refHit{o.up, o.down, o.swapped}
}
}
res.parts[w] = append(res.parts[w], ampRecord{
hash: ampHash(raw, int(o.up.pos)+lp, int(o.down.pos), o.swapped),
pair: pi,
taxon: taxa.Of[s],
role: role,
})
if role != Example || model == nil {
return
}
fp := oriented[o.up.primer][b2i(!o.up.reverse)]
rp := oriented[o.down.primer][b2i(o.down.reverse)]
tmF, tmR := inf, inf
if windowAsIs(site1, raw, int(o.up.pos), n, p.Circular) {
if t, err := model.Tm(fp, site1); err == nil && !math.IsNaN(t) {
tmF = float32(t)
}
}
if windowRevComp(site2, raw, int(o.down.pos), n, p.Circular) {
if t, err := model.Tm(rp, site2); err == nil && !math.IsNaN(t) {
tmR = float32(t)
}
}
if o.swapped {
tmF, tmR = tmR, tmF // slot 1 is the downstream primer
}
lowTm1[pi] = min(lowOr(lowTm1, pi, inf), tmF)
lowTm2[pi] = min(lowOr(lowTm2, pi, inf), tmR)
})
}
store.add(w, pi, ampHash(raw, int(o.up.pos())+lp, int(o.down.pos()), o.swapped), packTaxon(taxa.Of[s], role))
slices.SortFunc(res.parts[w], func(a, b ampRecord) int { return cmp.Compare(a.pair, b.pair) })
if role != Example || model == nil {
return
}
fp := oriented[o.up.primer][b2i(!o.up.reverse())]
rp := oriented[o.down.primer][b2i(o.down.reverse())]
tmF, tmR := inf, inf
if windowAsIs(site1, raw, int(o.up.pos()), n, p.Circular) {
if t, err := model.Tm(fp, site1); err == nil && !math.IsNaN(t) {
tmF = float32(t)
}
}
if windowRevComp(site2, raw, int(o.down.pos()), n, p.Circular) {
if t, err := model.Tm(rp, site2); err == nil && !math.IsNaN(t) {
tmR = float32(t)
}
}
if o.swapped {
tmF, tmR = tmR, tmF // slot 1 is the downstream primer
}
lowTm1[pi] = min(lowOr(lowTm1, pi, inf), tmF)
lowTm2[pi] = min(lowOr(lowTm2, pi, inf), tmR)
}
done := func() {
tmMu.Lock()
for pi, t := range lowTm1 {
res.minTm1[pi] = min(res.minTm1[pi], t)
@@ -171,13 +236,9 @@ func (ps *primerSet) collectAmplicons(c *Corpus, taxa *Taxa, p PairParams, index
res.minTm2[pi] = min(res.minTm2[pi], t)
}
tmMu.Unlock()
}()
}
for s := range c.Seqs {
jobs <- s
}
close(jobs)
wg.Wait()
}
return visit, done
})
return res
}
+132
View File
@@ -0,0 +1,132 @@
package obiprimer
import "fmt"
// ampRecord is the logical content of the record of one amplicon: its barcode
// (see ampHash), the taxon of its sequence at the chosen rank (-1 when it has
// none) and the role of the sequence. It is stored packed in an ampStore.
type ampRecord struct {
hash uint64
taxon int32
role Role
}
const (
roleBits = 2
roleMask = 1<<roleBits - 1
// maxTaxa is the number of taxa a record can carry: the taxon, shifted by
// one so that "none" is 0, shares a 32 bits word with the role.
maxTaxa = 1<<(32-roleBits) - 2
)
func packTaxon(taxon int32, role Role) uint32 { return uint32(taxon+1)<<roleBits | uint32(role) }
func unpackTaxon(t uint32) (int32, Role) { return int32(t>>roleBits) - 1, Role(t & roleMask) }
// ampSlice is the amplicons of one pair: 12 bytes per amplicon, kept as two
// parallel arrays (a struct holding a uint64 and a uint32 would take 16).
type ampSlice struct {
hash []uint64
tax []uint32 // packTaxon
}
// ampPlan says how many amplicons each worker finds for each pair, and which
// worker scans which sequence: the write pass replays the counting pass with
// the very same assignment, so that every worker knows where to write its
// records without sharing any counter with the others.
type ampPlan struct {
owner []uint16 // owner[s]: the worker that scans sequence s
count [][]uint32 // count[w][pair]: amplicons of the pair in the sequences of worker w
}
// total returns the number of amplicons of every pair.
func (pl *ampPlan) total() []int {
t := make([]int, len(pl.count[0]))
for _, row := range pl.count {
for b, n := range row {
t[b] += int(n)
}
}
return t
}
// ampStore holds the amplicons of a batch of pairs, grouped by pair: for each
// pair, two slices of the exact size, since their number is known before they
// are written (see planAmplicons). Nothing grows by doubling, nothing is copied
// nor sorted afterwards. The records of a pair are those of worker 0, then those
// of worker 1, and so on.
//
// The pairs have slices of their own, rather than slices of two huge arrays: a
// few kB each find room in the free memory, where blocks of tens of MB make the
// process map new memory (and, on some systems, count the freed one still).
type ampStore struct {
plan *ampPlan
lo int // first pair of the batch in the plan
pairs []ampSlice
cursor [][]int // cursor[w][b]: next free slot of worker w for pair b
}
// newAmpStore allocates the store of the pairs lo..hi-1 of the plan.
func newAmpStore(plan *ampPlan, lo, hi int) *ampStore {
n := hi - lo
s := &ampStore{plan: plan, lo: lo, pairs: make([]ampSlice, n), cursor: make([][]int, len(plan.count))}
for w := range s.cursor {
s.cursor[w] = make([]int, n)
}
for b := range n {
pos := 0
for w, row := range plan.count {
s.cursor[w][b] = pos
pos += int(row[lo+b])
}
if pos > 0 {
s.pairs[b] = ampSlice{make([]uint64, pos), make([]uint32, pos)}
}
}
return s
}
// add writes one record of pair b (numbered in the batch), found by worker w.
func (s *ampStore) add(w int, b int32, hash uint64, taxon uint32) {
i := s.cursor[w][b]
s.cursor[w][b] = i + 1
s.pairs[b].hash[i], s.pairs[b].tax[i] = hash, taxon
}
// check panics if a worker did not write exactly the records counted for it:
// the counting and write passes must enumerate the same amplicons.
func (s *ampStore) check() {
for b := range s.pairs {
pos := 0
for w, row := range s.plan.count {
pos += int(row[s.lo+b])
if s.cursor[w][b] != pos {
panic(fmt.Sprintf("obiprimer: pair %d: worker %d wrote %d amplicons, %d were counted",
s.lo+b, w, s.cursor[w][b]-(pos-int(row[s.lo+b])), row[s.lo+b]))
}
}
}
}
// pair returns the amplicons of pair b (numbered in the batch).
func (s *ampStore) pair(b int) ampSlice { return s.pairs[b] }
// ampBatches cuts the pairs 0..len(count)-1 in consecutive ranges [lo,hi)
// whose amplicons together do not exceed budget records (a pair alone above the
// budget makes a batch of its own). A non positive budget gives one batch.
func ampBatches(count []int, budget int) [][2]int {
var out [][2]int
lo, sum := 0, 0
for i, n := range count {
if budget > 0 && i > lo && sum+n > budget {
out = append(out, [2]int{lo, i})
lo, sum = i, 0
}
sum += n
}
if lo < len(count) {
out = append(out, [2]int{lo, len(count)})
}
return out
}
+60 -65
View File
@@ -1,16 +1,15 @@
package obiprimer
import (
"cmp"
"fmt"
"slices"
"sync"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obidefault"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiiter"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiseq"
)
// maxSeqLen bounds the length of a sequence: the position of a primer site and
// its strand share 32 bits (see seqMatch).
const maxSeqLen = 1 << 31
// Counts gives the number of sequences by role.
type Counts struct {
Examples int
@@ -26,8 +25,11 @@ type Counts struct {
// The primer search needs several passes over the sequences (words are
// searched approximately in every one of them), so the corpus lives in memory.
type Corpus struct {
Seqs []*obiseq.BioSequence
Roles []Role // Roles[i] is the role of Seqs[i]
Seqs []*obiseq.BioSequence // nil entries once the corpus is packed
Roles []Role // Roles[i] is the role of Seqs[i]
packed []packedSeq // see Pack
ids []string
}
// Counts returns the number of sequences by role.
@@ -46,68 +48,12 @@ func (c *Corpus) Counts() Counts {
return n
}
type corpusEntry struct {
batch, index int
seq *obiseq.BioSequence
role Role
}
// LoadCorpus reads every sequence of seqs, once, and keeps all of them with the
// role given by the selection. The order of the input is preserved whatever
// the number of reader goroutines.
func LoadCorpus(seqs obiiter.IBioSequence, sel Selection) (*Corpus, error) {
var mu sync.Mutex
var entries []corpusEntry
consumer := func(it obiiter.IBioSequence, wg *sync.WaitGroup) {
defer wg.Done()
var local []corpusEntry
for it.Next() {
batch := it.Get()
for i, seq := range batch.Slice() {
local = append(local, corpusEntry{batch.Order(), i, seq, sel.Role(seq)})
}
}
mu.Lock()
entries = append(entries, local...)
mu.Unlock()
}
var wg sync.WaitGroup
for i := 1; i < obidefault.ParallelWorkers(); i++ {
wg.Add(1)
go consumer(seqs.Split(), &wg)
}
wg.Add(1)
go consumer(seqs, &wg)
wg.Wait()
slices.SortFunc(entries, func(a, b corpusEntry) int {
if c := cmp.Compare(a.batch, b.batch); c != 0 {
return c
}
return cmp.Compare(a.index, b.index)
})
c := &Corpus{
Seqs: make([]*obiseq.BioSequence, len(entries)),
Roles: make([]Role, len(entries)),
}
for i, e := range entries {
c.Seqs[i], c.Roles[i] = e.seq, e.role
}
if c.Counts().Examples == 0 {
return nil, fmt.Errorf("obiprimer: no example sequence selected")
}
return c, nil
}
// IndexOfExample returns the index of the example sequence with the given
// identifier. It is an error if there is none, or if the sequence is a
// counterexample.
func (c *Corpus) IndexOfExample(id string) (int, error) {
for i, s := range c.Seqs {
if s.Id() == id {
for i := range c.Roles {
if c.ID(i) == id {
if c.Roles[i] != Example {
return -1, fmt.Errorf("obiprimer: reference sequence %q is a counterexample, it must be an example", id)
}
@@ -116,3 +62,52 @@ func (c *Corpus) IndexOfExample(id string) (int, error) {
}
return -1, fmt.Errorf("obiprimer: reference sequence %q is not among the selected example sequences", id)
}
// Len returns the number of sequences.
func (c *Corpus) Len() int { return len(c.Roles) }
// SeqLen returns the length of sequence i.
func (c *Corpus) SeqLen(i int) int {
if c.Seqs[i] != nil {
return c.Seqs[i].Len()
}
return c.packed[i].n
}
// ID returns the identifier of sequence i.
func (c *Corpus) ID(i int) string {
if c.Seqs[i] != nil {
return c.Seqs[i].Id()
}
return c.ids[i]
}
// bases returns the bases of sequence i, lower case. A packed sequence is
// unpacked into buf, grown if need be, which the caller keeps for the next call
// (the second result); an unpacked one is returned as it is. The result must
// not be modified, and is valid until the next call with the same buffer.
func (c *Corpus) bases(i int, buf []byte) (raw, newBuf []byte) {
if c.Seqs[i] != nil {
return c.Seqs[i].Sequence(), buf
}
buf = c.packed[i].unpack(buf)
return buf, buf
}
// Pack replaces the sequences by their packed form, on 2 bits per base, and
// frees the sequence objects (their annotations, their slack): the primer search
// only needs the bases, the identifiers, and the taxa, which must have been
// assigned before (AssignTaxa). A packed corpus holds about a quarter of the
// memory of the bases themselves.
func (c *Corpus) Pack() {
if c.packed != nil {
return
}
c.packed = make([]packedSeq, len(c.Seqs))
c.ids = make([]string, len(c.Seqs))
parallelFor(len(c.Seqs), func(_, i int) {
c.packed[i] = packBases(c.Seqs[i].Sequence())
c.ids[i] = c.Seqs[i].Id()
c.Seqs[i] = nil
})
}
+137
View File
@@ -0,0 +1,137 @@
package obiprimer
import (
"cmp"
"fmt"
"slices"
"sync"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obidefault"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiiter"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiseq"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obitax"
)
// corpusEntry is a sequence read by one of the loading goroutines: either the
// sequence itself, or, when it is packed as it is read, what is kept of it.
type corpusEntry struct {
batch, index int
role Role
seq *obiseq.BioSequence // unpacked loading
packed packedSeq // packed loading
id string
taxon *obitax.Taxon // at the chosen rank, nil when there is none
}
// loadEntries reads every sequence of seqs, once, on all the cores, and returns
// the entries in the order of the input, whatever the number of reader
// goroutines. keep fills the entry of a sequence (besides its position and its
// role); it runs in the reading goroutines.
func loadEntries(seqs obiiter.IBioSequence, sel Selection, keep func(*obiseq.BioSequence, *corpusEntry)) []corpusEntry {
var mu sync.Mutex
var entries []corpusEntry
consumer := func(it obiiter.IBioSequence, wg *sync.WaitGroup) {
defer wg.Done()
var local []corpusEntry
for it.Next() {
batch := it.Get()
for i, seq := range batch.Slice() {
e := corpusEntry{batch: batch.Order(), index: i, role: sel.Role(seq)}
keep(seq, &e)
local = append(local, e)
}
}
mu.Lock()
entries = append(entries, local...)
mu.Unlock()
}
var wg sync.WaitGroup
for i := 1; i < obidefault.ParallelWorkers(); i++ {
wg.Add(1)
go consumer(seqs.Split(), &wg)
}
wg.Add(1)
go consumer(seqs, &wg)
wg.Wait()
slices.SortFunc(entries, func(a, b corpusEntry) int {
if c := cmp.Compare(a.batch, b.batch); c != 0 {
return c
}
return cmp.Compare(a.index, b.index)
})
return entries
}
func checkSeqLen(id string, n int) error {
if n >= maxSeqLen {
return fmt.Errorf("obiprimer: sequence %q is too long (%d bases, the limit is %d)", id, n, maxSeqLen-1)
}
return nil
}
// LoadCorpus reads every sequence of seqs, once, and keeps all of them with the
// role given by the selection. The order of the input is preserved whatever
// the number of reader goroutines.
func LoadCorpus(seqs obiiter.IBioSequence, sel Selection) (*Corpus, error) {
entries := loadEntries(seqs, sel, func(seq *obiseq.BioSequence, e *corpusEntry) { e.seq = seq })
c := &Corpus{
Seqs: make([]*obiseq.BioSequence, len(entries)),
Roles: make([]Role, len(entries)),
}
for i, e := range entries {
if err := checkSeqLen(e.seq.Id(), e.seq.Len()); err != nil {
return nil, err
}
c.Seqs[i], c.Roles[i] = e.seq, e.role
}
if c.Counts().Examples == 0 {
return nil, fmt.Errorf("obiprimer: no example sequence selected")
}
return c, nil
}
// LoadPackedCorpus is LoadCorpus followed by Pack and AssignTaxa, without ever
// holding the sequences themselves: each one is packed, and its taxon at rank
// looked up, as it is read, and then dropped. The corpus of a large database
// would otherwise first take, in memory, the sequences with all their
// annotations, four times the size of their bases and more.
func LoadPackedCorpus(seqs obiiter.IBioSequence, sel Selection, taxonomy *obitax.Taxonomy, rank string) (*Corpus, *Taxa, error) {
if err := checkRank(taxonomy, rank); err != nil {
return nil, nil, err
}
entries := loadEntries(seqs, sel, func(seq *obiseq.BioSequence, e *corpusEntry) {
e.packed = packBases(seq.Sequence())
e.id = seq.Id()
e.taxon = seq.Taxon(taxonomy).TaxonAtRank(rank)
})
c := &Corpus{
Seqs: make([]*obiseq.BioSequence, len(entries)),
Roles: make([]Role, len(entries)),
packed: make([]packedSeq, len(entries)),
ids: make([]string, len(entries)),
}
taxons := make([]*obitax.Taxon, len(entries))
for i := range entries {
e := &entries[i]
if err := checkSeqLen(e.id, e.packed.n); err != nil {
return nil, nil, err
}
c.Roles[i], c.packed[i], c.ids[i], taxons[i] = e.role, e.packed, e.id, e.taxon
*e = corpusEntry{} // frees the entry as it goes
}
if c.Counts().Examples == 0 {
return nil, nil, fmt.Errorf("obiprimer: no example sequence selected")
}
taxa, err := buildTaxa(rank, taxons, c.Roles)
if err != nil {
return nil, nil, err
}
return c, taxa, nil
}
+40 -18
View File
@@ -8,7 +8,6 @@ import (
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obidefault"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obikmer"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiseq"
)
// Match is one approximate occurrence of a word in a corpus sequence.
@@ -54,15 +53,18 @@ func (p MatchParams) validate() error {
return nil
}
// hit is a Match before its sequence index is known to be final.
// hit is a Match before its sequence is known: the sequence it belongs to is
// the one being scanned.
type hit struct {
word int
m Match
word uint32
pos uint32
errors uint8
reverse bool
}
// scanFunc finds every occurrence of every word, on both strands, in one
// sequence and reports them through emit.
type scanFunc func(seq *obiseq.BioSequence, emit func(word, pos, errors int, reverse bool))
// sequence (its bases) and reports them through emit.
type scanFunc func(raw []byte, emit func(word, pos, errors int, reverse bool))
// scannerMaker returns an independent scanFunc; each worker gets its own so
// that scanners can keep private buffers.
@@ -83,8 +85,13 @@ func wordStrings(words []uint64, k int) (fwd, rev []string) {
// matchCorpus runs the scanners over every sequence of the corpus in parallel
// and assembles the per-word results. The output does not depend on the number
// of workers.
//
// The occurrences of a sequence are kept in a slice of their exact size, then
// counted per word, so that the occurrences of all the words are laid out in
// one allocation of the exact size too: nothing grows by doubling, nor is held
// twice for long.
func matchCorpus(c *Corpus, words []uint64, mk scannerMaker) []WordMatches {
perSeq := make([][]hit, len(c.Seqs))
perSeq := make([][]hit, c.Len())
jobs := make(chan int)
var wg sync.WaitGroup
@@ -93,31 +100,46 @@ func matchCorpus(c *Corpus, words []uint64, mk scannerMaker) []WordMatches {
go func() {
defer wg.Done()
scan := mk()
var buf []hit
var raw, bases []byte
for i := range jobs {
var hits []hit
scan(c.Seqs[i], func(word, pos, errors int, reverse bool) {
hits = append(hits, hit{word, Match{uint32(i), uint32(pos), uint8(errors), reverse}})
buf = buf[:0]
raw, bases = c.bases(i, bases)
scan(raw, func(word, pos, errors int, reverse bool) {
buf = append(buf, hit{uint32(word), uint32(pos), uint8(errors), reverse})
})
perSeq[i] = hits
if len(buf) > 0 {
perSeq[i] = slices.Clone(buf)
}
}
}()
}
for i := range c.Seqs {
for i := range c.Roles {
jobs <- i
}
close(jobs)
wg.Wait()
res := make([]WordMatches, len(words))
for i, w := range words {
res[i].Word = w
}
count := make([]int, len(words))
for _, hits := range perSeq {
for _, h := range hits {
res[h.word].Matches = append(res[h.word].Matches, h.m)
count[h.word]++
}
}
for i := range res {
for i, w := range words {
res[i].Word = w
if count[i] > 0 {
res[i].Matches = make([]Match, 0, count[i])
}
}
for s, hits := range perSeq {
for _, h := range hits {
res[h.word].Matches = append(res[h.word].Matches, Match{uint32(s), h.pos, h.errors, h.reverse})
}
perSeq[s] = nil // the hits of the sequence are no longer needed
}
parallelFor(len(res), func(_, i int) {
slices.SortFunc(res[i].Matches, func(a, b Match) int {
if c := cmp.Compare(a.Seq, b.Seq); c != 0 {
return c
@@ -133,7 +155,7 @@ func matchCorpus(c *Corpus, words []uint64, mk scannerMaker) []WordMatches {
}
return 0
})
}
})
return res
}
+1 -4
View File
@@ -1,7 +1,5 @@
package obiprimer
import "git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiseq"
// baseCode maps a nucleotide to 0..3 (A,C,G,T; U is T) and everything else to 4,
// which never equals a word base.
var baseCode = func() (t [256]byte) {
@@ -39,8 +37,7 @@ func matchNaive(c *Corpus, words []uint64, p MatchParams) ([]WordMatches, error)
mk := func() scanFunc {
var codes []byte
return func(seq *obiseq.BioSequence, emit func(word, pos, errors int, reverse bool)) {
raw := seq.Sequence()
return func(raw []byte, emit func(word, pos, errors int, reverse bool)) {
n := len(raw)
if n < p.K {
return
+1 -4
View File
@@ -3,8 +3,6 @@ package obiprimer
import (
"math/bits"
"slices"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiseq"
)
// seedIndex maps the value of one half of a window to the patterns that are
@@ -136,8 +134,7 @@ func MatchWords(c *Corpus, words []uint64, p MatchParams) ([]WordMatches, error)
maxErr := p.MaxErrors
mk := func() scanFunc {
return func(seq *obiseq.BioSequence, emit func(word, pos, errors int, reverse bool)) {
raw := seq.Sequence()
return func(raw []byte, emit func(word, pos, errors int, reverse bool)) {
n := len(raw)
if n < k {
return
+80
View File
@@ -0,0 +1,80 @@
package obiprimer
import (
"encoding/binary"
"slices"
)
// packedSeq holds the bases of a sequence on 2 bits each, a quarter of the
// bytes of the sequence itself. The bytes that are not a, c, g or t (N, the
// IUPAC codes, u...) are kept apart in runs, so that unpacking gives back the
// sequence byte for byte.
type packedSeq struct {
n int
bits []byte // 4 bases per byte, the first in the low bits
exc []excRun // sorted by position
}
// excRun is a run of length identical bytes b that are not a, c, g or t,
// starting at pos.
type excRun struct {
pos, length uint32
b byte
}
// packCode maps a, c, g, t to 0..3 and every other byte to 4.
var packCode = func() (t [256]byte) {
for i := range t {
t[i] = 4
}
for i, c := range "acgt" {
t[c] = byte(i)
}
return
}()
// unpackWord gives, for a byte of 4 packed bases, the 4 bytes they stand for.
var unpackWord = func() (t [256]uint32) {
for b := range t {
for k := range 4 {
t[b] |= uint32("acgt"[(b>>(2*k))&3]) << (8 * k)
}
}
return
}()
// packBases packs raw (lower case).
func packBases(raw []byte) packedSeq {
p := packedSeq{n: len(raw), bits: make([]byte, (len(raw)+3)/4)}
for i, b := range raw {
code := packCode[b]
if code > 3 {
if k := len(p.exc) - 1; k >= 0 && p.exc[k].b == b && int(p.exc[k].pos)+int(p.exc[k].length) == i {
p.exc[k].length++
} else {
p.exc = append(p.exc, excRun{uint32(i), 1, b})
}
code = 0
}
p.bits[i>>2] |= code << (2 * (i & 3))
}
p.exc = slices.Clone(p.exc) // the exact size, nil when there is none
return p
}
// unpack writes the sequence into dst, which it grows if need be, and returns it.
func (p packedSeq) unpack(dst []byte) []byte {
dst = slices.Grow(dst[:0], p.n)[:p.n]
for i := range p.n / 4 {
binary.LittleEndian.PutUint32(dst[4*i:], unpackWord[p.bits[i]])
}
for t := p.n &^ 3; t < p.n; t++ {
dst[t] = "acgt"[(p.bits[t>>2]>>(2*(t&3)))&3]
}
for _, e := range p.exc {
for k := range e.length {
dst[e.pos+k] = e.b
}
}
return dst
}
+109
View File
@@ -0,0 +1,109 @@
package obiprimer
import (
"bytes"
"math/rand"
"reflect"
"strings"
"testing"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiseq"
)
func TestPackedSeqRoundTrip(t *testing.T) {
rng := rand.New(rand.NewSource(1))
alphabet := []byte("acgtacgtacgtacgtnnrykmu-.")
for trial := 0; trial < 500; trial++ {
n := rng.Intn(70) // all the lengths modulo 4, the empty sequence included
raw := make([]byte, n)
for i := range raw {
raw[i] = alphabet[rng.Intn(len(alphabet))]
}
if trial%5 == 0 && n > 10 { // a long run of N, and one at each end
copy(raw[2:9], "nnnnnnn")
raw[0], raw[n-1] = 'n', 'n'
}
p := packBases(raw)
buf := make([]byte, 3, 8) // dirty, too small or too big: whatever
if got := p.unpack(buf); !bytes.Equal(got, raw) {
t.Fatalf("trial %d: %q unpacked as %q", trial, raw, got)
}
if len(p.bits) != (n+3)/4 {
t.Fatalf("trial %d: %d bytes for %d bases", trial, len(p.bits), n)
}
}
if p := packBases([]byte("acgtacgt")); p.exc != nil {
t.Errorf("no exception expected: %v", p.exc)
}
}
// Packing the corpus must change nothing to the search, ambiguous bases and
// reference included.
func TestPackedCorpusGivesTheSameSearch(t *testing.T) {
build := func() synthetic {
s := newSynthetic(10, -1, 7)
rng := rand.New(rand.NewSource(3))
for k, seq := range s.corpus.Seqs { // sprinkle ambiguous bases, some inside the sites
raw := append([]byte(nil), seq.Sequence()...)
for range 3 { // in the flanks, and in the insert past the first bases
at := rng.Intn(len(raw))
if at >= s.pos-2 && at < s.pos+18+2 || at >= s.pos+18+40-2 && at < s.pos+18+40+18+2 {
continue
}
raw[at] = "nrykm"[rng.Intn(5)]
}
if k < 2 { // one mismatch inside a site, which the search tolerates
raw[s.pos+5] = 'n'
}
s.corpus.Seqs[k] = obiseq.NewBioSequence(seq.Id(), raw, "")
}
return s
}
run := func(pack bool, circular bool) []PairStats {
s := build()
if pack {
s.corpus.Pack()
for i := range s.corpus.Seqs {
if s.corpus.Seqs[i] != nil {
t.Fatal("Pack must free the sequences")
}
}
}
p := s.params()
p.Reference = "s3"
p.ThreePrime = 3
p.Circular = circular
p.MinSpecificity = 0
pairs, _, err := Search(s.corpus, s.taxa, p)
if err != nil {
t.Fatal(err)
}
return pairs
}
for _, circular := range []bool{false, true} {
want := run(false, circular)
if len(want) == 0 {
t.Fatal("the test needs pairs")
}
if got := run(true, circular); !reflect.DeepEqual(got, want) {
t.Errorf("circular=%v: packing changes the result (%d pairs, want %d)", circular, len(got), len(want))
}
}
}
func TestPackedCorpusAccessors(t *testing.T) {
c := &Corpus{
Seqs: []*obiseq.BioSequence{obiseq.NewBioSequence("a", []byte(strings.ToLower("ACGTNNAC")), ""), obiseq.NewBioSequence("b", []byte("gg"), "")},
Roles: []Role{Example, Other},
}
c.Pack()
if c.Len() != 2 || c.SeqLen(0) != 8 || c.SeqLen(1) != 2 || c.ID(1) != "b" {
t.Errorf("accessors: %d %d %d %q", c.Len(), c.SeqLen(0), c.SeqLen(1), c.ID(1))
}
if raw, _ := c.bases(0, nil); string(raw) != "acgtnnac" {
t.Errorf("bases: %q", raw)
}
if i, err := c.IndexOfExample("a"); err != nil || i != 0 {
t.Errorf("IndexOfExample: %d %v", i, err)
}
}
+66 -26
View File
@@ -24,6 +24,8 @@ type PairParams struct {
Circular bool // -c
Reference int // -R: corpus index of the reference sequence (used when HasReference)
HasReference bool
AmpBudget int // amplicon records held at once in phase 2 (0 = DefaultAmpBudget)
Release func() // see SearchParams.Release
}
func (p PairParams) validate() error {
@@ -51,13 +53,20 @@ type Primer struct {
Good bool
}
// seqMatch is one occurrence of a primer in one sequence.
// seqMatch is one occurrence of a primer in one sequence, on 8 bytes: there
// are millions of them.
type seqMatch struct {
primer int32
pos uint32
reverse bool
primer int32
place uint32 // position << 1, plus 1 when the primer is found reversed
}
func newSeqMatch(primer int32, pos uint32, reverse bool) seqMatch {
return seqMatch{primer, pos<<1 | uint32(b2i(reverse))}
}
func (m seqMatch) pos() uint32 { return m.place >> 1 }
func (m seqMatch) reverse() bool { return m.place&1 == 1 }
// primerSet is the conserved words and, for each corpus sequence, their
// occurrences sorted by position.
type primerSet struct {
@@ -76,21 +85,41 @@ func selectPrimers(c *Corpus, wm []WordMatches, p PairParams) *primerSet {
inQ := quorumCount(counts.Examples, p.Sensitivity)
outQ := quorumCount(counts.CounterExamples, p.FalsePositive)
ps := &primerSet{perSeq: make([][]seqMatch, len(c.Seqs))}
for _, w := range wm {
ex, co := w.SequencesMatched(c)
if ex < inQ {
ex, co := make([]int, len(wm)), make([]int, len(wm))
parallelFor(len(wm), func(_, i int) { ex[i], co[i] = wm[i].SequencesMatched(c) })
// the occurrences of the words kept, by sequence, each in a slice of the
// exact size (and not in one big array: a few hundred bytes to some kB each
// find room in the free memory, which a single block of tens of MB does not)
ps := &primerSet{perSeq: make([][]seqMatch, c.Len())}
size := make([]int, c.Len())
for i, w := range wm {
if ex[i] < inQ {
continue
}
id := int32(len(ps.primers))
ps.primers = append(ps.primers, Primer{w.Word, ex, co, co <= outQ})
ps.primers = append(ps.primers, Primer{w.Word, ex[i], co[i], co[i] <= outQ})
for _, m := range w.Matches {
ps.perSeq[m.Seq] = append(ps.perSeq[m.Seq], seqMatch{id, m.Pos, m.Reverse})
size[m.Seq]++
}
}
for _, ms := range ps.perSeq {
slices.SortStableFunc(ms, func(a, b seqMatch) int { return cmp.Compare(a.pos, b.pos) })
for s, n := range size {
if n > 0 {
ps.perSeq[s] = make([]seqMatch, 0, n)
}
}
id := int32(0)
for i, w := range wm {
if ex[i] < inQ {
continue
}
for _, m := range w.Matches {
ps.perSeq[m.Seq] = append(ps.perSeq[m.Seq], newSeqMatch(id, m.Pos, m.Reverse))
}
id++
}
parallelFor(c.Len(), func(_, s int) {
slices.SortStableFunc(ps.perSeq[s], func(a, b seqMatch) int { return cmp.Compare(a.pos(), b.pos()) })
})
return ps
}
@@ -111,18 +140,28 @@ type occurrence struct {
}
// forEachOccurrence calls fn for every amplicon of sequence s, applying the
// length, good-primer and 3' constraints of ecoPrimers.
func (ps *primerSet) forEachOccurrence(c *Corpus, s int, p PairParams, fn func(occurrence)) {
// length, good-primer and 3' constraints of ecoPrimers. buf is the buffer in
// which the sequence is unpacked, if the corpus is packed.
func (ps *primerSet) forEachOccurrence(c *Corpus, s int, buf []byte, p PairParams, fn func(occurrence)) []byte {
ms := ps.perSeq[s]
lp := p.PrimerLength
var raw []byte
if p.ThreePrime > 0 {
raw = c.Seqs[s].Sequence()
if p.ThreePrime > 0 && len(ms) > 1 {
raw, buf = c.bases(s, buf)
}
ps.forEachOccurrenceIn(raw, ms, p, fn)
return buf
}
// forEachOccurrenceIn is forEachOccurrence over ms, the matches of a sequence
// sorted by position (all of them, or only those of the primers of interest,
// which yields the amplicons between these primers and none other), and raw its
// bases, needed when the 3' bases must match exactly.
func (ps *primerSet) forEachOccurrenceIn(raw []byte, ms []seqMatch, p PairParams, fn func(occurrence)) {
lp := p.PrimerLength
for i := range ms {
for j := i + 1; j < len(ms); j++ {
dpos := int(ms[j].pos) - int(ms[i].pos)
dpos := int(ms[j].pos()) - int(ms[i].pos())
if dpos <= lp {
continue
}
@@ -137,7 +176,7 @@ func (ps *primerSet) forEachOccurrence(c *Corpus, s int, p PairParams, fn func(o
// The upstream match is read as is (primer 1 of the pair when its word
// is smaller), the downstream one as the reverse primer.
key := pairKey{up.primer, down.primer, !up.reverse, down.reverse}
key := pairKey{up.primer, down.primer, !up.reverse(), down.reverse()}
swapped := !(ps.primers[key.p2].Word > ps.primers[key.p1].Word)
if swapped {
key = pairKey{key.p2, key.p1, key.a2, key.a1}
@@ -172,23 +211,23 @@ func (ps *primerSet) threePrimeMatches(raw []byte, up, down seqMatch, p PairPara
for t := lp - p.ThreePrime; t < lp; t++ {
// forward primer: the word when found direct, its reverse complement otherwise
var f byte
if !up.reverse {
if !up.reverse() {
f = wordBase(fw, lp, t)
} else {
f = 3 - wordBase(fw, lp, lp-1-t)
}
if s := at(int(up.pos) + t); s > 3 || s != f {
if s := at(int(up.pos()) + t); s > 3 || s != f {
return false
}
// reverse primer: the word when found reverse, its reverse complement otherwise
var r byte
if down.reverse {
if down.reverse() {
r = wordBase(rv, lp, t)
} else {
r = 3 - wordBase(rv, lp, lp-1-t)
}
if s := at(int(down.pos) + lp - 1 - t); s > 3 || 3-s != r {
if s := at(int(down.pos()) + lp - 1 - t); s > 3 || 3-s != r {
return false
}
}
@@ -220,12 +259,13 @@ func (ps *primerSet) countPairs(c *Corpus, p PairParams) map[pairKey]*pairAcc {
go func() {
defer wg.Done()
acc := partial[w]
var buf []byte
for s := range jobs {
if c.Roles[s] == Other {
continue // only the quorums are counted here
}
isExample := c.Roles[s] == Example
ps.forEachOccurrence(c, s, p, func(o occurrence) {
buf = ps.forEachOccurrence(c, s, buf, p, func(o occurrence) {
a := acc[o.key]
if a == nil {
a = &pairAcc{minLen: ^uint32(0)}
@@ -250,7 +290,7 @@ func (ps *primerSet) countPairs(c *Corpus, p PairParams) map[pairKey]*pairAcc {
}
}()
}
for s := range c.Seqs {
for s := range c.Roles {
jobs <- s
}
close(jobs)
+78
View File
@@ -0,0 +1,78 @@
package obiprimer
import (
"math/rand"
"reflect"
"strings"
"testing"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiseq"
)
func TestAmpBatches(t *testing.T) {
cases := []struct {
count []int
budget int
want [][2]int
}{
{[]int{3, 3, 3}, 0, [][2]int{{0, 3}}},
{[]int{3, 3, 3}, 100, [][2]int{{0, 3}}},
{[]int{3, 3, 3}, 6, [][2]int{{0, 2}, {2, 3}}},
{[]int{3, 3, 3}, 1, [][2]int{{0, 1}, {1, 2}, {2, 3}}}, // a pair above the budget goes alone
{[]int{9, 1, 1}, 5, [][2]int{{0, 1}, {1, 3}}},
{nil, 5, nil},
{[]int{0, 0}, 5, [][2]int{{0, 2}}},
}
for _, c := range cases {
if got := ampBatches(c.count, c.budget); !reflect.DeepEqual(got, c.want) {
t.Errorf("ampBatches(%v, %d) = %v, want %v", c.count, c.budget, got, c.want)
}
}
}
// The batches of phase 2 only bound the memory: whatever their size, the pairs,
// their indices and their sites on the reference are the same.
func TestAmpBudgetDoesNotChangeTheResult(t *testing.T) {
s := newSynthetic(10, -1, 5)
// non-targets sharing barcodes with targets, and counterexamples
for k, role := range []Role{Other, Other, CounterExample, Other, CounterExample} {
cp := append([]byte(nil), s.corpus.Seqs[k*2].Sequence()...)
if role == CounterExample {
for i := s.pos; i < s.pos+10; i++ { // no forward site: not amplified
cp[i] = "cgta"[strings.IndexByte("acgt", cp[i])]
}
}
s.corpus.Seqs = append(s.corpus.Seqs, obiseq.NewBioSequence(strings.Repeat("x", k+1), cp, ""))
s.corpus.Roles = append(s.corpus.Roles, role)
s.taxa.Of = append(s.taxa.Of, int32(100+k%3))
}
// a few unrelated sequences, with a taxon of their own or none
rng := rand.New(rand.NewSource(9))
for k := 0; k < 3; k++ {
s.corpus.Seqs = append(s.corpus.Seqs, obiseq.NewBioSequence("u"+strings.Repeat("y", k), []byte(strings.ToLower(randDNA(rng, 300))), ""))
s.corpus.Roles = append(s.corpus.Roles, Other)
s.taxa.Of = append(s.taxa.Of, int32(k-1))
}
run := func(budget int) []PairStats {
p := s.params()
p.MinSpecificity = 0 // keep every pair: the stress case of the batches
p.FalsePositive = 1
p.Reference = "s3"
p.AmpBudget = budget
pairs, info, err := Search(s.corpus, s.taxa, p)
if err != nil {
t.Fatal(err)
}
if info.Pair.Candidates < 2 {
t.Fatalf("budget %d: only %d pairs reach phase 2b, the test needs several", budget, info.Pair.Candidates)
}
return pairs
}
want := run(0)
for _, budget := range []int{1, 7, 50, 1000} {
if got := run(budget); !reflect.DeepEqual(got, want) {
t.Errorf("budget %d changes the result:\n got %d pairs\nwant %d pairs", budget, len(got), len(want))
}
}
}
+5 -47
View File
@@ -5,9 +5,7 @@ import (
"fmt"
"math"
"slices"
"sync"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obidefault"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obithermo"
)
@@ -69,10 +67,11 @@ func FindPairs(c *Corpus, taxa *Taxa, wm []WordMatches, p PairParams, model *obi
if err := p.validate(); err != nil {
return nil, info, err
}
if taxa == nil || len(taxa.Of) != len(c.Seqs) {
if taxa == nil || len(taxa.Of) != c.Len() {
return nil, info, fmt.Errorf("obiprimer: taxa do not match the corpus")
}
ps := selectPrimers(c, wm, p)
wm = nil // the occurrences are in ps now
info.Primers = len(ps.primers)
for _, pr := range ps.primers {
if pr.Good {
@@ -104,63 +103,22 @@ func FindPairs(c *Corpus, taxa *Taxa, wm []WordMatches, p PairParams, model *obi
return cmp.Or(cmp.Compare(a.p1, b.p1), cmp.Compare(a.p2, b.p2),
boolCmp(a.a1, b.a1), boolCmp(a.a2, b.a2))
})
index := make(map[pairKey]int32, len(kept))
for i, k := range kept {
index[k] = int32(i)
}
// Phase 2a: the amplicons of the target sequences, which give the melting
// temperatures, the reference sites, and an upper bound of Bs. It is cheap
// and rejects most pairs on the specificity threshold.
first := ps.collectAmplicons(c, taxa, p, index, collectOptions{examplesOnly: true, model: model, reference: true})
first := ps.boundPairs(c, taxa, p, kept, model)
var survivors []int // indices in kept
for i := range kept {
targets, identified := lenientIdentified(first.pair(int32(i))...)
if targets > 0 && float32(identified)/float32(targets) >= float32(p.MinSpecificity) {
if first.targets[i] > 0 && float32(first.identified[i])/float32(first.targets[i]) >= float32(p.MinSpecificity) {
survivors = append(survivors, i)
}
}
first.parts = nil
info.Candidates = len(survivors)
// Phase 2b: for the survivors only, the amplicons of every sequence, whatever
// its role: they decide which target taxa are identified.
index2 := make(map[pairKey]int32, len(survivors))
for j, i := range survivors {
index2[kept[i]] = int32(j)
}
amps := ps.collectAmplicons(c, taxa, p, index2, collectOptions{})
results := make([]*PairStats, len(survivors))
jobs := make(chan int)
var wg sync.WaitGroup
for range obidefault.ParallelWorkers() {
wg.Add(1)
go func() {
defer wg.Done()
for j := range jobs {
i := survivors[j]
k := kept[i]
st, ok := ps.evaluate(taxa, counts, p, k, acc[k], resolve(amps.pair(int32(j))...),
first.minTm1[i], first.minTm2[i], model)
if !ok {
continue
}
if h, found := first.refs[int32(i)]; found {
// evaluate orders the primers by string; the amplicon is read from the first one
s1 := orientedWord(ps.primers[k.p1].Word, p.PrimerLength, k.a1)
s2 := orientedWord(ps.primers[k.p2].Word, p.PrimerLength, k.a2)
st.RefLocation, st.RefAmplicon = referenceSites(c.Seqs[p.Reference].Sequence(), h, p, s2 < s1)
}
results[j] = &st
}
}()
}
for j := range survivors {
jobs <- j
}
close(jobs)
wg.Wait()
results := ps.evaluatePairs(c, taxa, p, kept, survivors, acc, first, model)
var out []PairStats
for _, st := range results {
+160
View File
@@ -0,0 +1,160 @@
package obiprimer
import (
"sync"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obidefault"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obithermo"
)
// DefaultAmpBudget is the number of amplicon records (12 bytes each) phase 2
// holds at once: 20 million, 240 MB. The pairs are processed in batches whose
// amplicons fit in it. With the usual specificity threshold a search needs one
// batch; with none (-T 0) no pair is rejected early and there are many more
// amplicons than any memory could hold.
const DefaultAmpBudget = 20_000_000
func (p PairParams) release() {
if p.Release != nil {
p.Release()
}
}
func (p PairParams) ampBudget() int {
if p.AmpBudget > 0 {
return p.AmpBudget
}
return DefaultAmpBudget
}
// batchIndex maps the pairs kept[lo:hi] to their number in the batch.
func batchIndex(kept []pairKey, lo, hi int) map[pairKey]int32 {
index := make(map[pairKey]int32, hi-lo)
for i := lo; i < hi; i++ {
index[kept[i]] = int32(i - lo)
}
return index
}
// parallelFor calls fn(i) for i in [0,n) on obidefault.ParallelWorkers() goroutines.
func parallelFor(n int, fn func(worker, i int)) {
jobs := make(chan int)
var wg sync.WaitGroup
for w := range obidefault.ParallelWorkers() {
wg.Add(1)
go func() {
defer wg.Done()
for i := range jobs {
fn(w, i)
}
}()
}
for i := 0; i < n; i++ {
jobs <- i
}
close(jobs)
wg.Wait()
}
// bound is what phase 2a learns about each retained pair (indexed like kept).
type bound struct {
targets []int // target taxa amplified
identified []int // upper bound of the target taxa identified (lenientIdentified)
minTm1 []float32
minTm2 []float32
refs map[int32]refHit
}
// boundPairs is phase 2a: from the amplicons of the target sequences alone, it
// gives each retained pair the upper bound of its Bs, the lowest melting
// temperatures and its site on the reference. The pairs go by batches whose
// amplicons fit in the budget; they are counted first, so that nothing is resized.
func (ps *primerSet) boundPairs(c *Corpus, taxa *Taxa, p PairParams, kept []pairKey,
model *obithermo.Model) *bound {
plan := ps.planAmplicons(c, p, batchIndex(kept, 0, len(kept)), true)
res := &bound{
targets: make([]int, len(kept)),
identified: make([]int, len(kept)),
minTm1: make([]float32, len(kept)),
minTm2: make([]float32, len(kept)),
refs: map[int32]refHit{},
}
resolvers := make([]*resolver, obidefault.ParallelWorkers())
for i := range resolvers {
resolvers[i] = newResolver()
}
for _, b := range ampBatches(plan.total(), p.ampBudget()) {
lo, hi := b[0], b[1]
p.release() // the previous batch, or the stage before, is garbage
store := newAmpStore(plan, lo, hi)
sites := ps.collectAmplicons(c, taxa, p, batchIndex(kept, lo, hi), plan, store,
collectOptions{examplesOnly: true, model: model, reference: true})
store.check()
parallelFor(hi-lo, func(w, i int) {
res.targets[lo+i], res.identified[lo+i] = resolvers[w].lenientIdentified(store.pair(i))
})
copy(res.minTm1[lo:hi], sites.minTm1)
copy(res.minTm2[lo:hi], sites.minTm2)
for i, h := range sites.refs {
res.refs[int32(lo)+i] = h
}
}
return res
}
// evaluatePairs is phase 2b: for the surviving pairs (indices in kept), it
// gathers the amplicons of every sequence, whatever its role, by batches whose
// amplicons fit in the budget, and evaluates the pairs. The result is indexed
// like survivors; a nil entry is a pair dropped by the specificity threshold.
func (ps *primerSet) evaluatePairs(c *Corpus, taxa *Taxa, p PairParams, kept []pairKey, survivors []int,
acc map[pairKey]*pairAcc, first *bound, model *obithermo.Model) []*PairStats {
results := make([]*PairStats, len(survivors))
if len(survivors) == 0 {
return results
}
keys := make([]pairKey, len(survivors))
for j, i := range survivors {
keys[j] = kept[i]
}
// the amplicons of every sequence are counted first: the store is allocated
// at its exact size, and the batches are cut on the exact counts
plan := ps.planAmplicons(c, p, batchIndex(keys, 0, len(keys)), false)
counts := c.Counts()
var refBases []byte
if p.HasReference {
refBases, _ = c.bases(p.Reference, nil)
}
resolvers := make([]*resolver, obidefault.ParallelWorkers())
for i := range resolvers {
resolvers[i] = newResolver()
}
for _, b := range ampBatches(plan.total(), p.ampBudget()) {
lo, hi := b[0], b[1]
p.release()
store := newAmpStore(plan, lo, hi)
ps.collectAmplicons(c, taxa, p, batchIndex(keys, lo, hi), plan, store, collectOptions{})
store.check()
parallelFor(hi-lo, func(w, n int) {
j := lo + n
i := survivors[j]
k := keys[j]
st, ok := ps.evaluate(taxa, counts, p, k, acc[k], resolvers[w].resolve(store.pair(n)),
first.minTm1[i], first.minTm2[i], model)
if !ok {
return
}
if h, found := first.refs[int32(i)]; found {
// evaluate orders the primers by string; the amplicon is read from the first one
s1 := orientedWord(ps.primers[k.p1].Word, p.PrimerLength, k.a1)
s2 := orientedWord(ps.primers[k.p2].Word, p.PrimerLength, k.a2)
st.RefLocation, st.RefAmplicon = referenceSites(refBases, h, p, s2 < s1)
}
results[j] = &st
})
}
return results
}
+2 -2
View File
@@ -34,8 +34,8 @@ func referenceSites(raw []byte, h refHit, p PairParams, swappedByOrder bool) (lo
// the exchange due to the string order flips it again.
firstIsUpstream := h.swapped == swappedByOrder
location = "join(" + span(int(h.up.pos), lp, n, p.Circular) + "," + span(int(h.down.pos), lp, n, p.Circular) + ")"
region := []byte(strings.ToLower(string(raw[int(h.up.pos)+lp : int(h.down.pos)])))
location = "join(" + span(int(h.up.pos()), lp, n, p.Circular) + "," + span(int(h.down.pos()), lp, n, p.Circular) + ")"
region := []byte(strings.ToLower(string(raw[int(h.up.pos())+lp : int(h.down.pos())])))
if firstIsUpstream {
return location, string(region)
}
+78 -57
View File
@@ -43,49 +43,67 @@ type resolution struct {
discriminant int // of which lead to a single taxon
}
// resolve computes the resolution of one pair from its amplicons. The slices
// need not be sorted; they are only read.
func resolve(parts ...[]ampRecord) resolution {
// resolver computes the indices of pairs one after the other, reusing its maps
// (a pair needs a few thousand entries, and there are thousands of pairs).
type resolver struct {
state map[uint64]barcodeState
targets map[int32]struct{}
counters map[int32]struct{}
bad map[int32]struct{}
owner map[uint64]int32
}
func newResolver() *resolver {
return &resolver{
state: map[uint64]barcodeState{},
targets: map[int32]struct{}{},
counters: map[int32]struct{}{},
bad: map[int32]struct{}{},
owner: map[uint64]int32{},
}
}
// resolve computes the resolution of one pair from its amplicons, in any order.
func resolve(a ampSlice) resolution { return newResolver().resolve(a) }
func (rs *resolver) resolve(a ampSlice) resolution {
var res resolution
clear(rs.state)
clear(rs.targets)
clear(rs.counters)
clear(rs.bad)
state := map[uint64]barcodeState{}
targets := map[int32]struct{}{}
counters := map[int32]struct{}{}
for i, h := range a.hash {
taxon, role := unpackTaxon(a.tax[i])
st := rs.state[h]
switch {
case taxon < 0:
st.flags |= barcodeAnonymous
case st.taxon == 0:
st.taxon = taxon + 1
case st.taxon != taxon+1:
st.flags |= barcodeAmbiguous
}
if taxon >= 0 {
switch role {
case Example:
st.flags |= barcodeTarget
rs.targets[taxon] = struct{}{}
case CounterExample:
rs.counters[taxon] = struct{}{}
}
}
rs.state[h] = st
}
for _, part := range parts {
for _, r := range part {
st := state[r.hash]
switch {
case r.taxon < 0:
st.flags |= barcodeAnonymous
case st.taxon == 0:
st.taxon = r.taxon + 1
case st.taxon != r.taxon+1:
st.flags |= barcodeAmbiguous
}
if r.taxon >= 0 {
switch r.role {
case Example:
st.flags |= barcodeTarget
targets[r.taxon] = struct{}{}
case CounterExample:
counters[r.taxon] = struct{}{}
}
}
state[r.hash] = st
for i, h := range a.hash {
taxon, role := unpackTaxon(a.tax[i])
if role == Example && taxon >= 0 && rs.state[h].ambiguous() {
rs.bad[taxon] = struct{}{}
}
}
bad := map[int32]struct{}{}
for _, part := range parts {
for _, r := range part {
if r.role == Example && r.taxon >= 0 && state[r.hash].ambiguous() {
bad[r.taxon] = struct{}{}
}
}
}
for _, st := range state {
for _, st := range rs.state {
if st.flags&barcodeTarget != 0 {
res.barcodes++
if !st.ambiguous() {
@@ -93,8 +111,8 @@ func resolve(parts ...[]ampRecord) resolution {
}
}
}
res.targetTaxa, res.counterTaxa = len(targets), len(counters)
res.identified = len(targets) - len(bad)
res.targetTaxa, res.counterTaxa = len(rs.targets), len(rs.counters)
res.identified = len(rs.targets) - len(rs.bad)
return res
}
@@ -107,27 +125,30 @@ func resolve(parts ...[]ampRecord) resolution {
// counted here too. So identified <= the second result, with the same number of
// target taxa, which lets a pair be rejected on the specificity threshold without
// gathering the amplicons of the non-target sequences.
func lenientIdentified(parts ...[]ampRecord) (targetTaxa, identified int) {
owner := map[uint64]int32{} // barcode -> its taxon, or -1 when shared
targets := map[int32]struct{}{}
for _, part := range parts {
for _, r := range part {
if r.role != Example || r.taxon < 0 {
continue
}
targets[r.taxon] = struct{}{}
if prev, seen := owner[r.hash]; !seen {
owner[r.hash] = r.taxon
} else if prev != r.taxon {
owner[r.hash] = -1
}
func lenientIdentified(a ampSlice) (targetTaxa, identified int) {
return newResolver().lenientIdentified(a)
}
func (rs *resolver) lenientIdentified(a ampSlice) (targetTaxa, identified int) {
clear(rs.owner) // barcode -> its taxon, or -1 when shared
clear(rs.targets)
for i, h := range a.hash {
taxon, role := unpackTaxon(a.tax[i])
if role != Example || taxon < 0 {
continue
}
rs.targets[taxon] = struct{}{}
if prev, seen := rs.owner[h]; !seen {
rs.owner[h] = taxon
} else if prev != taxon {
rs.owner[h] = -1
}
}
alone := map[int32]struct{}{}
for _, t := range owner {
clear(rs.bad) // here: the taxa owning a barcode
for _, t := range rs.owner {
if t >= 0 {
alone[t] = struct{}{}
rs.bad[t] = struct{}{}
}
}
return len(targets), len(alone)
return len(rs.targets), len(rs.bad)
}
+31 -19
View File
@@ -6,6 +6,18 @@ import (
)
// rec builds an amplicon record: barcode, taxon (-1: none at the rank), role.
// pack builds an ampSlice from logical records.
func pack(parts ...[]ampRecord) ampSlice {
var a ampSlice
for _, part := range parts {
for _, r := range part {
a.hash = append(a.hash, r.hash)
a.tax = append(a.tax, packTaxon(r.taxon, r.role))
}
}
return a
}
func rec(barcode uint64, taxon int32, role Role) ampRecord {
return ampRecord{hash: barcode, taxon: taxon, role: role}
}
@@ -15,10 +27,10 @@ func TestResolutionUserExample(t *testing.T) {
// nor B is identified, but A1 and B1 lead to one species each.
const A, B = 0, 1
const A1, A2, B1 = 11, 12, 21
res := resolve([]ampRecord{
res := resolve(pack([]ampRecord{
rec(A1, A, Example), rec(A2, A, Example),
rec(A2, B, Example), rec(B1, B, Example),
})
}))
want := resolution{targetTaxa: 2, identified: 0, barcodes: 3, discriminant: 2}
if res != want {
t.Fatalf("got %+v, want %+v", res, want)
@@ -26,7 +38,7 @@ func TestResolutionUserExample(t *testing.T) {
}
func TestResolutionAllUnique(t *testing.T) {
res := resolve([]ampRecord{rec(1, 0, Example), rec(2, 1, Example), rec(3, 2, Example)})
res := resolve(pack([]ampRecord{rec(1, 0, Example), rec(2, 1, Example), rec(3, 2, Example)}))
if res.identified != 3 || res.discriminant != 3 || res.barcodes != 3 || res.targetTaxa != 3 {
t.Errorf("all barcodes distinct: %+v", res)
}
@@ -35,10 +47,10 @@ func TestResolutionAllUnique(t *testing.T) {
func TestResolutionSameTaxonSharesItsBarcodeFreely(t *testing.T) {
// several sequences of one species with the same barcode are not ambiguous,
// and several barcodes for one species are fine when none is shared
res := resolve([]ampRecord{
res := resolve(pack([]ampRecord{
rec(1, 0, Example), rec(1, 0, Example), rec(1, 0, Example),
rec(2, 0, Example),
})
}))
if res.identified != 1 || res.barcodes != 2 || res.discriminant != 2 {
t.Errorf("%+v", res)
}
@@ -48,12 +60,12 @@ func TestResolutionSameTaxonSharesItsBarcodeFreely(t *testing.T) {
// non-target, or with a counterexample, does not identify it.
func TestResolutionNonTargetsCount(t *testing.T) {
for _, role := range []Role{Other, CounterExample} {
res := resolve([]ampRecord{
res := resolve(pack([]ampRecord{
rec(1, 0, Example), // unique among the targets...
rec(1, 5, role), // ...but shared with a non-target taxon
rec(2, 1, Example), // identified
rec(3, 6, role), // a barcode of non-targets only
})
}))
if res.targetTaxa != 2 || res.identified != 1 {
t.Errorf("role %d: %+v, want 2 target taxa and 1 identified", role, res)
}
@@ -65,10 +77,10 @@ func TestResolutionNonTargetsCount(t *testing.T) {
}
func TestResolutionAmplifiedTaxaCounts(t *testing.T) {
res := resolve([]ampRecord{
res := resolve(pack([]ampRecord{
rec(1, 0, Example), rec(2, 1, CounterExample), rec(3, 1, CounterExample),
rec(4, 2, CounterExample), rec(5, 3, Other),
})
}))
if res.targetTaxa != 1 || res.counterTaxa != 2 {
t.Errorf("target taxa %d (want 1), counterexample taxa %d (want 2: others are not counted)",
res.targetTaxa, res.counterTaxa)
@@ -78,20 +90,20 @@ func TestResolutionAmplifiedTaxaCounts(t *testing.T) {
// A sequence with no taxon at the rank is an anonymous taxon, distinct from all
// the others: a barcode shared with it is ambiguous.
func TestResolutionAnonymousTaxon(t *testing.T) {
res := resolve([]ampRecord{
res := resolve(pack([]ampRecord{
rec(1, 0, Example), rec(1, -1, Other), // shared with an unknown: ambiguous
rec(2, 1, Example),
})
}))
if res.identified != 1 || res.barcodes != 2 || res.discriminant != 1 {
t.Errorf("target sharing with an anonymous sequence: %+v", res)
}
// two anonymous sequences with the same barcode never count: no target carries it
res = resolve([]ampRecord{rec(7, -1, Other), rec(7, -1, Other), rec(8, 0, Example)})
res = resolve(pack([]ampRecord{rec(7, -1, Other), rec(7, -1, Other), rec(8, 0, Example)}))
if res.barcodes != 1 || res.discriminant != 1 || res.identified != 1 {
t.Errorf("anonymous barcodes alone: %+v", res)
}
// a target with no taxon is not a target taxon, and is anonymous too
res = resolve([]ampRecord{rec(1, -1, Example), rec(1, 0, Example), rec(2, 0, Example)})
res = resolve(pack([]ampRecord{rec(1, -1, Example), rec(1, 0, Example), rec(2, 0, Example)}))
if res.targetTaxa != 1 || res.identified != 0 {
t.Errorf("an anonymous target sequence sharing a barcode must make the taxon ambiguous: %+v", res)
}
@@ -101,10 +113,10 @@ func TestResolutionSeveralPartsAndOrder(t *testing.T) {
// the amplicons come as one slice per worker, in any order
a := []ampRecord{rec(1, 0, Example), rec(2, 1, Example)}
b := []ampRecord{rec(2, 2, Other), rec(3, 1, Example)}
if res := resolve(a, b); res.targetTaxa != 2 || res.identified != 1 {
if res := resolve(pack(a, b)); res.targetTaxa != 2 || res.identified != 1 {
t.Errorf("taxon 1 shares barcode 2 with taxon 2: %+v", res)
}
if res := resolve(b, a); res.targetTaxa != 2 || res.identified != 1 {
if res := resolve(pack(b, a)); res.targetTaxa != 2 || res.identified != 1 {
t.Errorf("order must not matter: %+v", res)
}
}
@@ -117,10 +129,10 @@ func TestLenientIdentifiedUserExample(t *testing.T) {
rec(11, 0, Example), rec(12, 0, Example),
rec(12, 1, Example), rec(21, 1, Example),
}
if targets, identified := lenientIdentified(recs); targets != 2 || identified != 2 {
if targets, identified := lenientIdentified(pack(recs)); targets != 2 || identified != 2 {
t.Errorf("lenient rule: %d of %d, want 2 of 2", identified, targets)
}
if res := resolve(recs); res.identified != 0 {
if res := resolve(pack(recs)); res.identified != 0 {
t.Errorf("strict rule: %d identified, want 0", res.identified)
}
}
@@ -138,8 +150,8 @@ func TestLenientIdentifiedIsAnUpperBound(t *testing.T) {
}
// split in a few parts, as the workers do
cut := rng.Intn(len(recs) + 1)
res := resolve(recs[:cut], recs[cut:])
targets, lenient := lenientIdentified(recs[:cut], recs[cut:])
res := resolve(pack(recs[:cut], recs[cut:]))
targets, lenient := lenientIdentified(pack(recs[:cut], recs[cut:]))
if targets != res.targetTaxa {
t.Fatalf("trial %d: %d target taxa for the bound, %d for resolve", trial, targets, res.targetTaxa)
}
+14 -1
View File
@@ -22,6 +22,14 @@ type SearchParams struct {
NoMultiMatch bool // -U: discard the primers matching several times in a sequence
Reference string // -R: identifier of an example sequence on which the amplicons are located ("" = none)
// Release, if not nil, is called when a stage has just freed a lot of memory
// and the next one is about to take some: the command makes the runtime hand
// it back to the system then, which the garbage collector would only do
// slowly, and the peak of memory would be the sum of the stages.
Release func()
AmpBudget int // amplicon records held at once by the pair evaluation (0 = DefaultAmpBudget)
Partitions int // partitions of the k-mer counter (default 256)
TmpDir string
@@ -33,7 +41,7 @@ func DefaultSearchParams() SearchParams {
return SearchParams{
PrimerLength: 18, MaxErrors: 3, MinLength: 0, MaxLength: 1000,
StrictQuorum: 0.7, Sensitivity: 0.9, FalsePositive: 0.1, MinSpecificity: 0.6,
Partitions: 256, Model: obithermo.DefaultModel(),
Model: obithermo.DefaultModel(),
}
}
@@ -82,6 +90,9 @@ func Search(c *Corpus, taxa *Taxa, p SearchParams) ([]PairStats, *SearchInfo, er
for i, w := range sw.Words {
words[i] = w.Kmer
}
if p.Release != nil {
p.Release() // the strict word search has just freed its occurrences
}
wm, err := MatchWords(c, words, MatchParams{K: p.PrimerLength, MaxErrors: p.MaxErrors, Circular: p.Circular})
if err != nil {
return nil, info, err
@@ -104,6 +115,8 @@ func Search(c *Corpus, taxa *Taxa, p SearchParams) ([]PairStats, *SearchInfo, er
Circular: p.Circular,
Reference: refIndex,
HasReference: refIndex >= 0,
AmpBudget: p.AmpBudget,
Release: p.Release,
}, p.Model)
info.Pair = pinfo
return pairs, info, err
+30 -17
View File
@@ -2,6 +2,7 @@ package obiprimer
import (
"fmt"
"slices"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obitax"
)
@@ -23,42 +24,54 @@ type Taxa struct {
// Name returns the taxon with the given index, for reporting.
func (t *Taxa) Name(i int32) *obitax.Taxon { return t.nodes[i] }
// AssignTaxa determines the taxon of every sequence of the corpus at rank.
func (c *Corpus) AssignTaxa(taxonomy *obitax.Taxonomy, rank string) (*Taxa, error) {
// checkRank reports an error if the taxonomy has no such rank.
func checkRank(taxonomy *obitax.Taxonomy, rank string) error {
if taxonomy == nil {
return nil, fmt.Errorf("obiprimer: a taxonomy is required to compute taxon statistics")
return fmt.Errorf("obiprimer: a taxonomy is required to compute taxon statistics")
}
ranks := taxonomy.RankList()
known := false
for _, r := range ranks {
if r == rank {
known = true
break
}
}
if !known {
return nil, fmt.Errorf("obiprimer: unknown taxonomic rank %q", rank)
if !slices.Contains(taxonomy.RankList(), rank) {
return fmt.Errorf("obiprimer: unknown taxonomic rank %q", rank)
}
return nil
}
t := &Taxa{Rank: rank, Of: make([]int32, len(c.Seqs))}
// AssignTaxa determines the taxon of every sequence of the corpus at rank. The
// sequences must not be packed (see LoadPackedCorpus).
func (c *Corpus) AssignTaxa(taxonomy *obitax.Taxonomy, rank string) (*Taxa, error) {
if err := checkRank(taxonomy, rank); err != nil {
return nil, err
}
taxons := make([]*obitax.Taxon, len(c.Seqs))
for i, seq := range c.Seqs {
taxons[i] = seq.Taxon(taxonomy).TaxonAtRank(rank)
}
return buildTaxa(rank, taxons, c.Roles)
}
// buildTaxa gives the sequences, whose taxon at rank is taxons[i] (nil when
// there is none), their dense taxon index.
func buildTaxa(rank string, taxons []*obitax.Taxon, roles []Role) (*Taxa, error) {
t := &Taxa{Rank: rank, Of: make([]int32, len(taxons))}
index := map[*obitax.TaxNode]int32{}
exSeen := map[int32]struct{}{}
coSeen := map[int32]struct{}{}
for i, seq := range c.Seqs {
for i, taxon := range taxons {
t.Of[i] = -1
taxon := seq.Taxon(taxonomy).TaxonAtRank(rank)
if taxon == nil {
continue
}
id, ok := index[taxon.Node]
if !ok {
if len(t.nodes) >= maxTaxa {
return nil, fmt.Errorf("obiprimer: more than %d taxa at rank %q", maxTaxa, rank)
}
id = int32(len(t.nodes))
index[taxon.Node] = id
t.nodes = append(t.nodes, taxon)
}
t.Of[i] = id
switch c.Roles[i] {
switch roles[i] {
case Example:
exSeen[id] = struct{}{}
case CounterExample:
+16 -5
View File
@@ -14,7 +14,7 @@ type WordSearch struct {
PrimerLength int // word size (ecoPrimers -O, 18 by default; 2..31)
Quorum float64 // fraction of example sequences that must contain the word (ecoPrimers -q, 0.7)
Circular bool // sequences are circular (ecoPrimers -c)
Partitions int // partition files of the counter (default 256)
Partitions int // partition files of the counter (0 = chosen from the size of the examples, see obikmer.DefaultPartitions)
TmpDir string // parent of the temporary directory ("" = system default)
}
@@ -48,7 +48,15 @@ func FindStrictWords(c *Corpus, cfg WordSearch) (*StrictWords, error) {
return nil, fmt.Errorf("obiprimer: quorum must be in ]0,1], got %g", cfg.Quorum)
}
if cfg.Partitions <= 0 {
cfg.Partitions = 256
// the k-mer occurrences to count: those of the examples, which are the
// only sequences indexed
total := 0
for i, r := range c.Roles {
if r == Example && c.SeqLen(i) > cfg.PrimerLength {
total += c.SeqLen(i)
}
}
cfg.Partitions = obikmer.DefaultPartitions(total)
}
counter, err := obikmer.NewSeqKmerCounter(cfg.TmpDir, cfg.PrimerLength, 0, cfg.Partitions)
@@ -63,14 +71,17 @@ func FindStrictWords(c *Corpus, cfg WordSearch) (*StrictWords, error) {
wg.Add(1)
go func() {
defer wg.Done()
var buf []byte
for i := range jobs {
if seq := c.Seqs[i]; c.Roles[i] == Example && seq.Len() > cfg.PrimerLength {
counter.AddSequence(uint32(i), seq.Sequence(), cfg.Circular)
if c.Roles[i] == Example && c.SeqLen(i) > cfg.PrimerLength {
var raw []byte
raw, buf = c.bases(i, buf)
counter.AddSequence(uint32(i), raw, cfg.Circular)
}
}
}()
}
for i := range c.Seqs {
for i := range c.Roles {
jobs <- i
}
close(jobs)
+50
View File
@@ -3,9 +3,11 @@ package obiprimers
import (
"os"
"path/filepath"
"reflect"
"strings"
"testing"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obidefault"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiiter"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiprimer"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiseq"
@@ -195,3 +197,51 @@ func TestTaxonomicPathIsDroppedAfterExtraction(t *testing.T) {
t.Errorf("stdin: %d sequences, %d with a path left, %d with a taxid", total, withPath, withTaxid)
}
}
// LoadPackedCorpus must give the same corpus and taxa as LoadCorpus followed by
// AssignTaxa, whatever the number of reading goroutines.
func TestLoadPackedCorpusEqualsLoadThenAssign(t *testing.T) {
path := writeTemp(t, embeddedFasta)
defer obidefault.SetMaxCPU(obidefault.MaxCPU())
for _, workers := range []int{1, 4} {
obidefault.SetMaxCPU(workers)
sequences, taxonomy, err := readSequencesAndTaxonomy([]string{path}, nil)
if err != nil {
t.Fatal(err)
}
restrict, _ := obiprimer.BelongingTo(taxonomy, []string{"5"})
sel := obiprimer.Selection{Restrict: restrict}
want, err := obiprimer.LoadCorpus(sequences, sel)
if err != nil {
t.Fatal(err)
}
wantTaxa, err := want.AssignTaxa(taxonomy, "species")
if err != nil {
t.Fatal(err)
}
sequences, _, err = readSequencesAndTaxonomy([]string{path}, taxonomy)
if err != nil {
t.Fatal(err)
}
got, gotTaxa, err := obiprimer.LoadPackedCorpus(sequences, sel, taxonomy, "species")
if err != nil {
t.Fatal(err)
}
if got.Counts() != want.Counts() || got.Len() != want.Len() {
t.Errorf("workers %d: counts %+v / %+v", workers, got.Counts(), want.Counts())
}
for i := range want.Roles {
if got.Roles[i] != want.Roles[i] || got.ID(i) != want.ID(i) || got.SeqLen(i) != want.SeqLen(i) || got.Seqs[i] != nil {
t.Errorf("workers %d, sequence %d differs", workers, i)
}
}
if !reflect.DeepEqual(gotTaxa.Of, wantTaxa.Of) || gotTaxa.ExampleTaxa != wantTaxa.ExampleTaxa {
t.Errorf("workers %d: taxa %v / %v", workers, gotTaxa.Of, wantTaxa.Of)
}
}
if _, _, err := obiprimer.LoadPackedCorpus(obiiter.NilIBioSequence, obiprimer.Selection{}, nil, "species"); err == nil {
t.Error("a nil taxonomy must be refused")
}
}
+45
View File
@@ -0,0 +1,45 @@
package obiprimers
import (
"os"
"runtime/debug"
"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obidefault"
)
// gcPercent is the garbage collector target of the command: the heap may grow
// by that percentage over the live data before the next collection (the Go
// default is 100). Reading the sequences and the search allocate large buffers
// and tables that die as soon as a stage is over, and with the default the
// process holds up to twice what it needs at its peak. Peak of memory and CPU
// time against the default (see OBIPRIMERS_NOTES.md, section 26): 7822
// vertebrate genomes, 50 gives -8% and +0%, 25 -13% and +8%; 18061 genomes
// (593 Mb), 50 gives -20% and +3%, 25 -29% and +8.5%.
const gcPercent = 50
// batchSizeMax is the default of the --batch-size-max option for this command.
// The reading goroutines hand the sequences over by batches, and a batch of 2000
// genomes of 16 kb (the general default) is 33 Mb of bases once parsed, times
// the number of workers and the batches waiting between the stages. Genomes are
// few and long: 100 per batch keeps all the cores busy.
const batchSizeMax = 100
// SetReadingDefaults gives the options shared by the OBITools their defaults for
// this command; it must be called before the options are parsed, which still
// override them.
func SetReadingDefaults() {
*obidefault.BatchSizeMaxPtr() = batchSizeMax
}
// tuneMemory sets the collector target. A user who set GOGC keeps it.
func tuneMemory() {
if os.Getenv("GOGC") == "" {
debug.SetGCPercent(gcPercent)
}
}
// releaseMemory collects the garbage and gives the free memory back to the
// system. Without it the memory freed by a stage stays in the process (the
// runtime returns it slowly), and the next stage takes new memory on top of it:
// the peak is then the sum of the stages instead of the largest.
func releaseMemory() { debug.FreeOSMemory() }
+5 -5
View File
@@ -94,6 +94,7 @@ func readSequencesAndTaxonomy(args []string, given *obitax.Taxonomy) (obiiter.IB
// returns them with the taxonomy to use, set as the default taxonomy: the one
// of the -t option, or else the one embedded in the sequences.
func CLIReadSequencesAndTaxonomy(args []string) (obiiter.IBioSequence, *obitax.Taxonomy, error) {
tuneMemory()
given := obitax.DefaultTaxonomy()
sequences, taxonomy, err := readSequencesAndTaxonomy(args, given)
if err != nil {
@@ -130,11 +131,8 @@ func CLIPrimerPairs(sequences obiiter.IBioSequence, taxonomy *obitax.Taxonomy) (
}
log.Infof("Reading the sequences")
corpus, err := obiprimer.LoadCorpus(sequences, sel)
if err != nil {
return nil, err
}
taxa, err := corpus.AssignTaxa(taxonomy, CLIRank())
// the sequences are packed, and their taxa looked up, as they are read
corpus, taxa, err := obiprimer.LoadPackedCorpus(sequences, sel, taxonomy, CLIRank())
if err != nil {
return nil, err
}
@@ -144,6 +142,8 @@ func CLIPrimerPairs(sequences obiiter.IBioSequence, taxonomy *obitax.Taxonomy) (
counts.Examples, taxa.ExampleTaxa, CLIRank(),
counts.CounterExamples, taxa.CounterExamples, CLIRank(), counts.Others)
releaseMemory() // the sequences were read through large transient buffers
pairs, info, err := obiprimer.Search(corpus, taxa, params)
if info != nil {
log.Infof("%d strict words (found without error in at least %d examples)", info.StrictWords, info.MinStrict)
+5 -3
View File
@@ -40,7 +40,7 @@ var _saltMethod = "santalucia"
var _salt = obithermo.DefaultSalt
var _primerConc = obithermo.DefaultPrimerConc
var _partitions = 256
var _partitions = 0
// OptionSet registers the options of obiprimers.
func OptionSet(options *getoptions.GetOpt) {
@@ -148,8 +148,9 @@ func OptionSet(options *getoptions.GetOpt) {
options.IntVar(&_partitions, "partitions", _partitions,
options.ArgName("COUNT"),
options.Description("Number of temporary partition files used to count the conserved words. "+
"Increase it for very large inputs (memory per partition grows with data size / COUNT)."))
options.Description("Number of temporary partition files used to count the conserved words "+
"(0 = chosen from the size of the examples: 256 to 1024). The memory of the count is bounded "+
"whatever the number; more files only make the largest partitions smaller."))
}
// CLIRestrictTaxa returns the taxids restricting the examples.
@@ -198,5 +199,6 @@ func CLISearchParams() (obiprimer.SearchParams, error) {
Reference: _reference,
Partitions: _partitions,
Model: model,
Release: releaseMemory,
}, nil
}