diff --git a/OBIPRIMERS_NOTES.md b/OBIPRIMERS_NOTES.md index bec957a..9b39c4d 100644 --- a/OBIPRIMERS_NOTES.md +++ b/OBIPRIMERS_NOTES.md @@ -27,7 +27,7 @@ ce qu'elles comparent (paires, drapeaux good, séquences et taxons amplifiés, B **Bs et `identifiés` ne sont plus comparables** à ceux d'ecoPrimers ; ils sont validés autrement (section 25). Sections en attente de décision de l'utilisateur : voie B (21.4), statistique de Tm (22.3-a), nombre de partitions -automatique, mémoire, tests shell, documentation, test `-e 0` (ecoPrimers n'avait pas terminé). +automatique (fait, section 26), mémoire (fait, section 26), tests shell, documentation, test `-e 0` (ecoPrimers n'avait pas terminé). Sources de référence : - Article : `~/Sync/travail/__MOI__/ecoprimers/Documents/Riaz et al. 2011 ... .pdf` (NAR 39:e145) @@ -1205,3 +1205,150 @@ Cyprinidae 3 448 paires, 0 absente, 0 champ différent ; Teleostei 407, 0 absent (l'artefact du `N`, section 12) ; `-U` 1 366, 0 absente, 0 différente ; contre-exemples 173, 0 absente, 0 différente. obiprimers produit en plus 22 251 (Cyprinidae) et 2 389 (Teleostei) paires qu'ecoPrimers avait écartées avec son propre Bs. + +## 26. Optimisation de la mémoire (2026-11-03) + +Mission : réduire la mémoire de `obiprimers` **sans changer un octet de la sortie**, en procédant par la mesure. +Contraintes respectées : tout reste en mémoire (l'entrée standard marche), un seul calcul par indice +(`resolution.go` garde ses définitions), aucun paquet partagé modifié (`obiformats`, `obiiter`, `obiseq`, +`obitax`, `obidefault` : `git diff` vide ; seuls `pkg/obiprimer`, `pkg/obikmer/seqquorum*.go`, +`pkg/obitools/obiprimers` et `cmd/obitools/obiprimers` changent). Pas de commit. + +### 26.1 Méthode, et pièges + +- **Garde-fou** : binaire de référence construit avant toute modification ; six CSV de référence (Cyprinidae + `-r 7953`, Teleostei `-r 32443`, `-r 7953 -i 1`, `-U`, `-R NC_083213`, entrée standard) comparés par `cmp` + après **chaque** étape : identiques à chaque fois. Idem pour `-T 0` (25 699 paires), `-t` (taxonomie NCBI) et le + test d'échelle (Metazoa) comparés au binaire de référence. +- **Mesure** : exécutions entrelacées, médiane (5 à 9 exécutions ; le pic RSS d'une même commande varie de ±100 Mo + avant, de ±10 Mo après). Le **temps CPU utilisateur est trompeur** : 40 % des échantillons d'un profil CPU sont du + `runtime.usleep` des P inactifs. Pour juger le vrai coût d'un changement, comparer aussi avec `--max-cpu 2`. +- **Le RSS n'est pas le tas** (observation, pas hypothèse : mesures croisées `ps` + `runtime.MemStats` toutes les + 60 ms). Sur cette machine le RSS ne descend jamais pendant l'exécution et suit `Sys`, même quand le runtime + annonce avoir *relâché* 258 Mo (`HeapReleased`). Une cause plausible, **non vérifiée** : macOS continue à compter + les pages rendues par `madvise` tant que le système ne les reprend pas. Conséquence pratique, vérifiée par + l'expérience du §26.3 (étapes 5 et 6) : ce qui fait monter le RSS, ce sont les très gros tableaux contigus + alloués successivement, qui n'entrent pas dans les trous du tas et font mapper de la mémoire neuve. +- Les profils de tas par étape (tas après GC forcé, `inuse_space`) sont faits sur une copie du dépôt dans `/tmp` + avec des marqueurs temporaires (`MemStage`) ; rien n'en reste dans le code. + +### 26.2 Rapport de profil, avant toute modification (Cyprinidae, 7 822 séquences) + +Tas vivant après GC, par étape, et ce qui le compose : + +| Étape | vivant | composition | +|---|---|---| +| lecture terminée | 178 Mo | 165 Mo d'octets de séquence (131 Mo utiles : 26 % de marge laissée par les `append`) + attributs | +| appariement approché | 274 Mo | + 97 Mo de `[]WordMatches` (82 Mo utiles) | +| `selectPrimers` | 283 Mo | `perSeq` 109 Mo (82 utiles), construit par `append` (238 Mo alloués), tri stable en série : **0,74 s** | +| phase 2a | 511 Mo | + **214 Mo** d'`ampRecord` (8,1 millions × 24 o, avec la marge des `append`) | +| phase 2b | 517 Mo | + **220 Mo** (9,2 millions × 24 o) ; 2,17 Go alloués au total par ces `append` | +| fin | pic du tas 1 005 Mo | environ 2 × le vivant (`GOGC=100`) | + +Les 1,1 Go « inexpliqués » étaient donc : (a) les enregistrements d'amplicons, 24 octets, marge des `append`, +copiés par pair puis triés ; (b) les occurrences stockées deux fois pendant `selectPrimers` (`WordMatches` puis +`perSeq`) avec des `append` en doublement ; (c) le retard du ramasse-miettes (le tas monte à ~2 × le vivant) ; +(d) la mémoire libérée par une étape qui n'est pas réutilisée par la suivante (§26.1). **Hypothèses réfutées** : +la lecture n'est pas le pic à cette taille (tas ≤ 330 Mo) ; `SeqKmerCounter` n'y est pour rien (il n'indexe que +les 325 exemples, 5 Mo) ; `GOGC=25` ne double pas le temps utilisateur ici (+7 % sur la référence, +5 % avec les +changements ci-dessous). + +### 26.3 Optimisations, dans l'ordre où elles ont été faites + +Médianes entrelacées (RSS en Mo, temps réel, temps utilisateur) ; « avant » est le binaire de l'étape précédente +mesuré **dans la même série** (les valeurs absolues changent d'une série à l'autre, d'où les deux colonnes). + +| # | Changement | RSS avant → après | réel | utilisateur | +|---|---|---|---|---| +| 1 | **Amplicons** (`ampstore.go`, `amplicons.go`, `pairs_phase2.go`) : enregistrement de **12 octets** (deux tableaux parallèles `hash`/`taxon+rôle` au lieu de 24 octets), taille exacte connue avant l'écriture (passe de comptage par worker, puis écriture sans compteur partagé), ni copie ni tri ; **lots bornés** de paires (budget `DefaultAmpBudget` = 20 millions d'enregistrements, 240 Mo) ; maps de `resolve` réutilisées | 1 402 → 866 | 3,76 → 2,77 s | 18,9 → 13,5 s | +| 2 | **Occurrences** (`match.go`, `pairs.go`) : `hit` de 12 octets, regroupement par mot en tranches de taille exacte, `seqMatch` de **8 octets**, tris parallélisés | 949 → 704 | 2,80 → 2,04 s | 13,5 → 13,2 s | +| 3 | **Séquences à 2 bits par base** (`packed.go`) : plages d'exceptions pour tout octet ≠ a/c/g/t (restitution exacte), décodage dans un tampon par worker ; 178 → 42 Mo vivants | 705 → 628 | = | 13,4 → 13,8 s | +| 4 | **Compactage à la lecture** (`LoadPackedCorpus`, `corpus_load.go`) : chaque séquence est compactée et son taxon cherché par les goroutines de lecture, l'objet `BioSequence` n'est jamais gardé | petit jeu : sans effet mesurable ; **échelle : 1 999 → 1 390** | = | = | +| 5 | **Crochet `Release`** (`SearchParams.Release`, branché sur `debug.FreeOSMemory`) aux frontières d'étapes ; `debug.SetGCPercent` (étape 7) | 660 → 532 | = | = | +| 6 | **Tranches exactes par mot / par séquence / par paire** au lieu d'un seul tableau de 55 à 110 Mo | 546 → 365 | 2,12 → 2,19 s | 13,6 → 13,8 s | +| 7 | **Ramasse-miettes** : `SetGCPercent` dès le début de la commande (`memory.go`), sauf si `GOGC` est fixé par l'utilisateur | 370 → 305 avec 25 ; 292 avec 50 (autre série) | = | 13,9 → 14,4 s avec 25 | + +**Étape 1, un faux départ instructif** : la première version (comptage exact, écriture dans un tableau partagé, +compteurs atomiques) faisait gagner 15 % de RSS mais coûtait **+18 % de CPU** (la passe de comptage énumère +deux fois les 7 822 séquences). Deux corrections : des curseurs privés par worker (le plan de comptage fixe +aussi quel worker lit quelle séquence, la passe d'écriture le rejoue), et surtout ne garder pour ces deux passes +que les occurrences **des amorces qui figurent dans les paires concernées** (613 amorces sur 1 200 pour les +survivants de la phase 2b) : l'énumération des paires d'occurrences est quadratique en densité. Résultat à +`--max-cpu 2` : réel 10,35 → 7,33 s et CPU 18,8 → 13,3 s **par rapport à la référence**. Même résultat par +construction : une paire dont une des deux amorces est absente de la liste ne peut pas être produite. + +**Étape 5, le choix du pourcentage** : `GOGC=50` donne l'essentiel du gain pour presque rien de CPU ; `25` retire +encore 150 Mo à l'échelle pour +5 % de CPU. Retenu : 50 (`gcPercent` dans `memory.go`). + +### 26.4 Résultat final (médianes entrelacées ; binaire de référence → binaire final) + +| Cas | RSS | temps réel | temps utilisateur | +|---|---|---|---| +| Cyprinidae, fichier (9 exécutions) | **1 376 → 292 Mo** (−79 %) | 3,77 → 2,15 s | 19,2 → 14,2 s | +| Cyprinidae, entrée standard (7) | 1 365 → 369 Mo | 2,91 → 1,33 s | 17,2 → 12,7 s | +| Teleostei `-r 32443` (5) | 1 047 → 378 Mo | 3,41 → 2,58 s | 21,6 → 20,8 s | +| Cyprinidae `-T 0`, aucune paire élaguée (3) | **5 883 → 409 Mo** | 6,06 → 4,82 s | 46,3 → 46,1 s | +| avec `-t` (taxonomie NCBI complète) (3) | 2 664 → 1 273 Mo | 10,2 → 8,7 s | 32,7 → 29,5 s | +| échelle : 18 061 génomes, 593 Mb, `-r 33208` (3) | 3 189 → 1 370 Mo | 9,06 → 8,95 s | 64,5 → 68,2 s | + +Pour comparaison, ecoPrimers (C) : 1,86 Go (Cyprinidae), 2,11 Go (Teleostei). Les tests : `go test ./pkg/obiprimer/ +./pkg/obitools/obiprimers/ ./pkg/obithermo/ ./pkg/obikmer/ -race -count=1` verts ; tests opt-in sur les données +réelles (`OBIPRIMERS_TUTORIAL`) verts (65 s) : mêmes paires qu'ecoPrimers, 0 absente, 0 champ différent. + +Tests ajoutés : `TestAmpBatches` et `TestAmpBudgetDoesNotChangeTheResult` (tout budget, jusqu'à un enregistrement, +donne les mêmes paires, indices et sites de référence) ; `TestPackedSeqRoundTrip` (toutes les longueurs modulo 4, +N, IUPAC, `u`, plages), `TestPackedCorpusGivesTheSameSearch` (recherche sur corpus compacté ou non, avec +ambiguïtés, référence, 3', circulaire) ; `TestLoadPackedCorpusEqualsLoadThenAssign` ; pour le compteur de k-mers : +budget sans effet sur les comptes, décompte des occurrences par partition, budget d'occurrences, +`DefaultPartitions`. + +### 26.5 Le test d'échelle (jamais fait avant) et le compteur de k-mers + +Sur 593 Mb d'un seul tenant (16 027 exemples), la mémoire n'est plus dominée par la phase 2 (aucune paire n'est +trouvée avec les Metazoa pour cibles : 18-mers conservés chez 70 % des cibles = 0) mais par deux autres postes, +que le petit jeu ne montrait pas : + +1. **Le comptage des mots stricts** (`SeqKmerCounter`, 16 octets par occurrence de k-mer, 526 millions + d'occurrences). Défauts trouvés : le tampon réutilisé par chaque worker gardait la taille de la plus grosse + partition traitée (jusqu'à 106 Mo) ; la plus grosse partition fait 6 fois la médiane. Corrigé + (`seqquorum*.go`) : les occurrences de chaque partition sont comptées à l'écriture et la partition est chargée + dans un tableau de taille exacte ; un **budget d'occurrences** (`DefaultOccurrenceBudget` = 32 millions, 512 Mo) + borne ce que les workers chargent ensemble, partitions les plus grosses d'abord (une partition plus grosse que le + budget passe seule) ; **nombre de partitions automatique** (`DefaultPartitions` : 256, jusqu'à 1 024 quand la + partition moyenne dépasserait ~1 million d'occurrences ; `--partitions 0`, nouveau défaut). La mémoire du comptage + ne dépend plus de la répartition des k-mers. Le budget est un compromis temps/mémoire, **mesuré** (16 cœurs, + RSS total / réel) : 8 M → 1 348 Mo / 14,4 s ; 16 M → 1 621 Mo / 10,2 s ; **32 M → 1 838 Mo / 8,75 s** ; sans + limite → 2 157 Mo / 8,70 s. 32 M ne coûte rien en temps. +2. **La lecture** : les lots de 2 000 séquences (défaut commun des obitools) sont de ~33 Mo de bases chacun ; + avec 16 workers et les lots qui attendent entre les étapes, 793 Mo restaient vivants pendant la lecture. + `--batch-size-max` à 100 par défaut pour cette commande (`SetReadingDefaults`, appelée avant l'analyse des + options, qui l'emportent) : à l'échelle 1 819 → 1 190 Mo (−35 %, budget de k-mers à 32 M), temps réel +2,5 %, CPU +7 % ; + petit jeu : 306 → 274 Mo, +4 % de CPU. `--batch-mem` (existant) + a aussi été évalué : 4 M → 2 312 → 1 582 Mo (+6 % CPU), mais le plafond en nombre de séquences est plus efficace. + **À signaler** : le texte d'aide commun de `--batch-size-max` (paquet `obioptions`, partagé, non modifié) dit + « default 2000 » alors que le défaut affiché par `getoptions` est maintenant 100 pour `obiprimers`. + +### 26.6 Pistes écartées, et pourquoi + +- **Clé `uint64` pour les tables de paires** (au lieu de la structure de 12 octets) : aucun gain mesurable + (22,8 → 22,1 s de CPU à 2 cœurs, bruit). Annulé. +- **`GOGC=25` par défaut** : 5 % de CPU de plus que 50 pour 150 Mo de moins à l'échelle, rien sur le petit jeu. + `GOGC` fixé par l'utilisateur est respecté, donc 25 reste à portée de main. +- **Budget d'amplicons plus petit** (`DefaultAmpBudget`) : non réglé finement. 20 millions d'enregistrements + (240 Mo) tiennent en un lot pour les cas usuels ; `-T 0` passe en plusieurs lots sans perte de temps mesurable. +- **Relire le fichier** : exclu (entrée standard). + +### 26.7 Ce qui reste, et décisions à prendre + +- **Entrée standard** (369 Mo au lieu de 292) : les séquences sont gardées comme objets pendant l'extraction de la + taxonomie embarquée (`embeddedTaxonomyInMemory`), avant d'être compactées. Les compacter dès la lecture + demanderait de scinder la sélection (qui a besoin de la taxonomie) : non fait. +- **Taxonomie NCBI (`-t`)** : ~1,07 Go à elle seule (chargeur `obitax`, partagé, hors périmètre). La taxonomie + embarquée (`obiannotate --taxonomic-path`) l'évite. À dire dans l'aide : non fait. +- **Budget du compteur de k-mers** (32 M occurrences) : décision de l'utilisateur si la mémoire prime sur le + temps (16 M : −12 % de RSS à l'échelle pour +16 % de temps). +- `Corpus.Pack()` n'est plus appelé que par les tests (et utile à qui bâtit un corpus à la main) ; l'ancien chemin + `LoadCorpus` + `AssignTaxa` reste pour les tests. À supprimer si l'on veut un seul chemin. +- Le crochet `Release` appelle une collecte complète à chaque frontière d'étape (4 appels dans le cas usuel : après + la lecture, avant l'appariement, avant chaque lot de 2a et de 2b) ; sans effet mesurable sur le temps. + diff --git a/cmd/obitools/obiprimers/main.go b/cmd/obitools/obiprimers/main.go index c78eb29..c8fa6c3 100644 --- a/cmd/obitools/obiprimers/main.go +++ b/cmd/obitools/obiprimers/main.go @@ -12,6 +12,8 @@ import ( ) func main() { + obiprimers.SetReadingDefaults() + optionParser := obioptions.GenerateOptionParser( "obiprimers", "designs the primer pairs amplifying a set of example sequences and evaluates the DNA barcode they delimit. "+ diff --git a/pkg/obikmer/seqquorum.go b/pkg/obikmer/seqquorum.go index 3471eaa..3b82da4 100644 --- a/pkg/obikmer/seqquorum.go +++ b/pkg/obikmer/seqquorum.go @@ -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) diff --git a/pkg/obikmer/seqquorum_budget_test.go b/pkg/obikmer/seqquorum_budget_test.go new file mode 100644 index 0000000..debe287 --- /dev/null +++ b/pkg/obikmer/seqquorum_budget_test.go @@ -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) + } + } +} diff --git a/pkg/obikmer/seqquorum_count.go b/pkg/obikmer/seqquorum_count.go index 9af0cd6..7b216c7 100644 --- a/pkg/obikmer/seqquorum_count.go +++ b/pkg/obikmer/seqquorum_count.go @@ -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() +} diff --git a/pkg/obiprimer/amplicons.go b/pkg/obiprimer/amplicons.go index 828bd8a..d547089 100644 --- a/pkg/obiprimer/amplicons.go +++ b/pkg/obiprimer/amplicons.go @@ -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 } diff --git a/pkg/obiprimer/ampstore.go b/pkg/obiprimer/ampstore.go new file mode 100644 index 0000000..416d4ff --- /dev/null +++ b/pkg/obiprimer/ampstore.go @@ -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, 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 := &Store{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 +} diff --git a/pkg/obiprimer/corpus.go b/pkg/obiprimer/corpus.go index c64ea34..bccc7ec 100644 --- a/pkg/obiprimer/corpus.go +++ b/pkg/obiprimer/corpus.go @@ -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 + }) +} diff --git a/pkg/obiprimer/corpus_load.go b/pkg/obiprimer/corpus_load.go new file mode 100644 index 0000000..f3d2848 --- /dev/null +++ b/pkg/obiprimer/corpus_load.go @@ -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 +} diff --git a/pkg/obiprimer/match.go b/pkg/obiprimer/match.go index 8f4216b..0e9af27 100644 --- a/pkg/obiprimer/match.go +++ b/pkg/obiprimer/match.go @@ -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 } diff --git a/pkg/obiprimer/match_naive.go b/pkg/obiprimer/match_naive.go index 54e9567..a815c94 100644 --- a/pkg/obiprimer/match_naive.go +++ b/pkg/obiprimer/match_naive.go @@ -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 diff --git a/pkg/obiprimer/match_seed.go b/pkg/obiprimer/match_seed.go index d720088..a63020d 100644 --- a/pkg/obiprimer/match_seed.go +++ b/pkg/obiprimer/match_seed.go @@ -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 diff --git a/pkg/obiprimer/packed.go b/pkg/obiprimer/packed.go new file mode 100644 index 0000000..431be04 --- /dev/null +++ b/pkg/obiprimer/packed.go @@ -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 +} diff --git a/pkg/obiprimer/packed_test.go b/pkg/obiprimer/packed_test.go new file mode 100644 index 0000000..50a7858 --- /dev/null +++ b/pkg/obiprimer/packed_test.go @@ -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) + } +} diff --git a/pkg/obiprimer/pairs.go b/pkg/obiprimer/pairs.go index 8457285..a4f04f8 100644 --- a/pkg/obiprimer/pairs.go +++ b/pkg/obiprimer/pairs.go @@ -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) diff --git a/pkg/obiprimer/pairs_batches_test.go b/pkg/obiprimer/pairs_batches_test.go new file mode 100644 index 0000000..39bc27f --- /dev/null +++ b/pkg/obiprimer/pairs_batches_test.go @@ -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)) + } + } +} diff --git a/pkg/obiprimer/pairs_eval.go b/pkg/obiprimer/pairs_eval.go index aee5b5d..e5cc823 100644 --- a/pkg/obiprimer/pairs_eval.go +++ b/pkg/obiprimer/pairs_eval.go @@ -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 { diff --git a/pkg/obiprimer/pairs_phase2.go b/pkg/obiprimer/pairs_phase2.go new file mode 100644 index 0000000..25d930b --- /dev/null +++ b/pkg/obiprimer/pairs_phase2.go @@ -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 +} diff --git a/pkg/obiprimer/reference.go b/pkg/obiprimer/reference.go index a2f6b31..835abc0 100644 --- a/pkg/obiprimer/reference.go +++ b/pkg/obiprimer/reference.go @@ -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) } diff --git a/pkg/obiprimer/resolution.go b/pkg/obiprimer/resolution.go index 86d0c89..25858d8 100644 --- a/pkg/obiprimer/resolution.go +++ b/pkg/obiprimer/resolution.go @@ -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) } diff --git a/pkg/obiprimer/resolution_test.go b/pkg/obiprimer/resolution_test.go index 4f625c9..37c6ce0 100644 --- a/pkg/obiprimer/resolution_test.go +++ b/pkg/obiprimer/resolution_test.go @@ -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) } diff --git a/pkg/obiprimer/search.go b/pkg/obiprimer/search.go index f4a5a8f..9ef249b 100644 --- a/pkg/obiprimer/search.go +++ b/pkg/obiprimer/search.go @@ -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 diff --git a/pkg/obiprimer/taxa.go b/pkg/obiprimer/taxa.go index 89698b0..92f0351 100644 --- a/pkg/obiprimer/taxa.go +++ b/pkg/obiprimer/taxa.go @@ -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: diff --git a/pkg/obiprimer/words.go b/pkg/obiprimer/words.go index d492039..868126f 100644 --- a/pkg/obiprimer/words.go +++ b/pkg/obiprimer/words.go @@ -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) diff --git a/pkg/obitools/obiprimers/embedded_test.go b/pkg/obitools/obiprimers/embedded_test.go index cef413e..c1f6438 100644 --- a/pkg/obitools/obiprimers/embedded_test.go +++ b/pkg/obitools/obiprimers/embedded_test.go @@ -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") + } +} diff --git a/pkg/obitools/obiprimers/memory.go b/pkg/obitools/obiprimers/memory.go new file mode 100644 index 0000000..04499e5 --- /dev/null +++ b/pkg/obitools/obiprimers/memory.go @@ -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() } diff --git a/pkg/obitools/obiprimers/obiprimers.go b/pkg/obitools/obiprimers/obiprimers.go index 59e6dc1..b1c27a0 100644 --- a/pkg/obitools/obiprimers/obiprimers.go +++ b/pkg/obitools/obiprimers/obiprimers.go @@ -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) diff --git a/pkg/obitools/obiprimers/options.go b/pkg/obitools/obiprimers/options.go index 13d37f4..85b2078 100644 --- a/pkg/obitools/obiprimers/options.go +++ b/pkg/obitools/obiprimers/options.go @@ -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 }