mirror of
https://github.com/metabarcoding/obitools4.git
synced 2026-11-10 01:24:41 +00:00
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
1 parent
d6dd39f878
commit
8cc02aa9c1
28 files changed
+1673
-382
No files matched your search
+159
-98
@@ -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 := &Plan{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 := &licons{
|
||||
parts: make([][]ampRecord, workers),
|
||||
res := &Sites{
|
||||
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
|
||||
}
|
||||
|
||||
|
||||
Reference in new issue
Block a user