obitools4/pkg/obiprimer/pairs_test.go

package obiprimer

import (
	"bufio"
	"fmt"
	"math"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"testing"
	"time"

	"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obiformats"
	"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obikmer"
	"git.metabarcoding.org/obitools/obitools4/obitools4/pkg/obithermo"
)

// ecoRow is one result line of ecoPrimers.
type ecoRow struct {
	p1, p2         string
	tm1, mtm1      float64
	tm2, mtm2      float64
	gc1, gc2       int
	good           string
	in, out        int
	yule           float64
	intaxa, outtax int
	bc             float64
	well           int
	bs             float64
	mind, maxd     int
	avg            float64
}

func readEcoRows(t *testing.T, path string) []ecoRow {
	t.Helper()
	f, err := os.Open(path)
	if err != nil {
		t.Fatal(err)
	}
	defer f.Close()
	var rows []ecoRow
	sc := bufio.NewScanner(f)
	sc.Buffer(make([]byte, 1<<20), 1<<20)
	for sc.Scan() {
		line := sc.Text()
		if strings.HasPrefix(line, "#") || strings.TrimSpace(line) == "" {
			continue
		}
		f := strings.Split(line, "\t")
		fl := func(i int) float64 {
			v, err := strconv.ParseFloat(strings.TrimSpace(f[i]), 64)
			if err != nil {
				t.Fatalf("field %d of %q: %v", i, line, err)
			}
			return v
		}
		in := func(i int) int { return int(fl(i)) }
		rows = append(rows, ecoRow{
			p1: f[1], p2: f[2], tm1: fl(3), mtm1: fl(4), tm2: fl(5), mtm2: fl(6),
			gc1: in(7), gc2: in(8), good: f[9], in: in(10), out: in(11), yule: fl(12),
			intaxa: in(13), outtax: in(14), bc: fl(15), well: in(16), bs: fl(17),
			mind: in(18), maxd: in(19), avg: fl(20),
		})
	}
	if err := sc.Err(); err != nil {
		t.Fatal(err)
	}
	return rows
}

// Opt-in end-to-end comparison with the complete ecoPrimers result tables of
// the tutorial: same parameters (-e 3 -3 2 -l 30 -L 150 -c, -r taxid).
//
// Bs and the number of identified taxa are NOT compared: obiprimers defines
// them differently (a taxon is identified only if all its barcodes are its own,
// and every amplified sequence counts, targets or not), so they are lower. The
// specificity threshold is switched off so that the pairs ecoPrimers dropped
// with its own Bs are still produced, and only the pairs of ecoPrimers are
// checked. What is checked: that every pair of ecoPrimers is found, with the
// same good flags, sequence and taxon counts, Yule, Bc, lengths, GC, Tm and
// lowest Tm; and that identified taxa never exceed those of ecoPrimers' more
// lenient rule on the same amplified sequences.
//
// Cyprinidae is reproduced exactly on these columns for its 3448 pairs.
// Teleostei finds the same 407 pairs, and 6 of them count one
// sequence less than ecoPrimers (NC_033548, whose downstream site is
// NGGGGAGGCAAGTCGTAA). The 3' check of ecoPrimers reads an N through the
// shift register of extractSite, where the code 4 of N spills on the next base
// and ~4&3 reads as T, so an N can be accepted by accident. Emulating that
// register exactly in the 3' check was verified to reproduce both tables on
// every column; it is not kept: an ambiguous base cannot be a verified 3' match.
func TestPairsAgainstEcoPrimers(t *testing.T) {
	dir := os.Getenv("OBIPRIMERS_TUTORIAL")
	if dir == "" {
		t.Skip("OBIPRIMERS_TUTORIAL not set")
	}
	taxonomy, err := obiformats.LoadTaxonomy(filepath.Join(dir, "ncbitaxo_20250211.tgz"), false, false)
	if err != nil {
		t.Fatal(err)
	}

	for _, tc := range []struct {
		name, taxid string
		species     int
		divergent   int
	}{{"Cyprinidae", "7953", 316, 0}, {"Teleostei", "32443", 3876, 6}} {
		t.Run(tc.name, func(t *testing.T) {
			var stages []string
			tick := time.Now()
			lap := func(name string) {
				stages = append(stages, fmt.Sprintf("%s %s", name, time.Since(tick).Round(time.Millisecond)))
				tick = time.Now()
			}
			restrict, err := BelongingTo(taxonomy, []string{tc.taxid})
			if err != nil {
				t.Fatal(err)
			}
			seqs, err := obiformats.ReadSequencesFromFile(filepath.Join(dir, "mito.vert.fasta.gz"))
			if err != nil {
				t.Fatal(err)
			}
			corpus, err := LoadCorpus(seqs, Selection{Restrict: restrict})
			if err != nil {
				t.Fatal(err)
			}
			lap("read+select")
			taxa, err := corpus.AssignTaxa(taxonomy, "species")
			if err != nil {
				t.Fatal(err)
			}
			lap("taxa")
			if taxa.ExampleTaxa != tc.species {
				t.Errorf("%d example species, ecoPrimers reports %d", taxa.ExampleTaxa, tc.species)
			}

			sw, err := FindStrictWords(corpus, WordSearch{PrimerLength: 18, Quorum: 0.7, Circular: true, TmpDir: t.TempDir()})
			if err != nil {
				t.Fatal(err)
			}
			lap("strict words")
			words := make([]uint64, len(sw.Words))
			for i, w := range sw.Words {
				words[i] = w.Kmer
			}
			wm, err := MatchWords(corpus, words, MatchParams{K: 18, MaxErrors: 3, Circular: true})
			if err != nil {
				t.Fatal(err)
			}
			lap("approximate search")
			got, _, err := FindPairs(corpus, taxa, wm, PairParams{
				PrimerLength: 18, MinLength: 30, MaxLength: 150, ThreePrime: 2,
				Sensitivity: 0.9, FalsePositive: 0.1, MinSpecificity: 0, Circular: true,
			}, obithermo.DefaultModel())
			if err != nil {
				t.Fatal(err)
			}

			lap("pairs+indices+Tm")
			var ms runtime.MemStats
			runtime.ReadMemStats(&ms)
			t.Logf("stages: %s; Go heap in use at the end %.0f MB (includes the taxonomy)", strings.Join(stages, ", "), float64(ms.HeapAlloc)/1e6)

			want := readEcoRows(t, filepath.Join(dir, tc.name+".ecoprimers"))
			t.Logf("%d pairs, ecoPrimers %d", len(got), len(want))

			byKey := map[string]PairStats{}
			for _, g := range got {
				byKey[g.Primer1+"/"+g.Primer2] = g
			}
			bad := map[string]int{}
			missing, divergent := 0, 0
			note := func(field string, w ecoRow, format string, args ...any) {
				if bad[field]++; bad[field] <= 3 {
					t.Errorf("%s/%s %s: "+format, append([]any{w.p1, w.p2, field}, args...)...)
				}
			}
			near := func(a, b, tol float64) bool { return math.Abs(a-b) <= tol }
			for _, w := range want {
				g, ok := byKey[w.p1+"/"+w.p2]
				if !ok {
					if missing++; missing <= 5 {
						t.Errorf("ecoPrimers pair %s/%s not found", w.p1, w.p2)
					}
					continue
				}
				delete(byKey, w.p1+"/"+w.p2)
				good := string("bG"[b2i(g.Good1)]) + string("bG"[b2i(g.Good2)])
				if good != w.good {
					note("good", w, "got %s want %s", good, w.good)
				}
				// Known divergence, see the comment of the test: ecoPrimers accepts one
				// amplicon (NC_033548) whose 3' primer site ends with an N. We reject it,
				// so the counts can be exactly one lower and nothing else may differ.
				lower := func(got, want int) bool { return got == want || got == want-1 }
				diverges := g.Examples != w.in || g.ExampleTaxa != w.intaxa
				if diverges {
					divergent++
				}
				if !lower(g.Examples, w.in) || g.CounterExamples != w.out {
					note("sequences", w, "got %d/%d want %d/%d", g.Examples, g.CounterExamples, w.in, w.out)
				}
				if !lower(g.ExampleTaxa, w.intaxa) || g.CounterTaxa != w.outtax {
					note("taxa", w, "got %d/%d want %d/%d", g.ExampleTaxa, g.CounterTaxa, w.intaxa, w.outtax)
				}
				if !diverges && g.Identified > w.well {
					note("identified", w, "got %d, more than the %d of ecoPrimers' lenient rule", g.Identified, w.well)
				}
				tol := 0.0006
				if diverges {
					tol = 0.0021
				}
				if !near(float64(g.Yule), w.yule, tol) || !near(float64(g.Bc), w.bc, tol) {
					note("indices", w, "yule/bc got %.4f/%.4f want %.3f/%.3f", g.Yule, g.Bc, w.yule, w.bc)
				}
				if g.MinLen != w.mind || g.MaxLen != w.maxd || !near(g.MeanLen, w.avg, 0.006) {
					note("lengths", w, "got %d/%d/%.2f want %d/%d/%.2f", g.MinLen, g.MaxLen, g.MeanLen, w.mind, w.maxd, w.avg)
				}
				if g.GC1 != w.gc1 || g.GC2 != w.gc2 {
					note("gc", w, "got %d/%d want %d/%d", g.GC1, g.GC2, w.gc1, w.gc2)
				}
				// ecoPrimers prints Kelvin-273.0; obithermo uses 273.15
				if !near(g.Tm1+0.15, w.tm1, 0.051) || !near(g.Tm2+0.15, w.tm2, 0.051) {
					note("tm", w, "got %.2f/%.2f want %.1f/%.1f", g.Tm1+0.15, g.Tm2+0.15, w.tm1, w.tm2)
				}
				if !near(g.MinTm1+0.15, w.mtm1, 0.051) || !near(g.MinTm2+0.15, w.mtm2, 0.051) {
					note("mintm", w, "got %.2f/%.2f want %.1f/%.1f", g.MinTm1+0.15, g.MinTm2+0.15, w.mtm1, w.mtm2)
				}
			}
			t.Logf("missing from ours: %d; pairs of ours that ecoPrimers dropped with its Bs: %d; mismatching fields: %v; pairs one sequence below ecoPrimers: %d",
				missing, len(byKey), bad, divergent)
			if divergent != tc.divergent {
				t.Errorf("%d pairs differ from ecoPrimers by one sequence, expected %d", divergent, tc.divergent)
			}
		})
	}
	_ = obikmer.DecodeKmer
}