From 34453804fe4a5195ca136932d299b4c5d177cb54 Mon Sep 17 00:00:00 2001 From: Paula Ruiz-Rodriguez <50167687+Paururo@users.noreply.github.com> Date: Sat, 17 Oct 2026 20:23:15 +0200 Subject: [PATCH] perf: keep annotations and gene lookups compact, so memory stops growing megabytes per genome Each genome's GFF was held as a String per name in up to six maps, and its gene-to-cluster lookup as a map of every name, copied again for the corrected matrix: 3.4 MB per genome with minimal GFFs and 6.2 MB with PGAP-style ones. Names now live once per genome in one buffer with a sorted hash table, the lookup is one cluster per name corrected in place, matrix cells are one string per cluster, and the block and context tables are built in chunks bounded by genome-cluster pairs, the contexts now in parallel. On 800 synthetic genomes peak memory goes from 2.8 to 0.56 GB (5.5 to 1.0 GB with PGAP-style GFFs) and run time from 33 to 22 seconds, with every output byte for byte the same. --- CHANGELOG.md | 17 + benchmarks/README.md | 6 + benchmarks/search_scaling.py | 46 +- benchmarks/search_scaling_metrics.tsv | 10 +- benchmarks/search_scaling_rich_metrics.tsv | 5 + docs/how-it-works/algorithm.md | 13 +- docs/how-it-works/limits.md | 42 +- src/parser.rs | 83 ++- src/pipeline.rs | 130 +++-- src/pipeline/evolution.rs | 459 ++++++++------- src/pipeline/output.rs | 2 +- src/pipeline/reference_audit.rs | 12 +- src/pipeline/scale_harness.rs | 2 +- src/pipeline/tests.rs | 11 +- src/rescue.rs | 32 +- src/rescue/genes.rs | 19 +- src/rescue/gff.rs | 629 ++++++++++++++++----- src/rescue/gff/tests.rs | 474 ++++++++++++++++ src/rescue/synteny.rs | 18 +- src/rescue/tests.rs | 282 +++++---- 20 files changed, 1715 insertions(+), 577 deletions(-) create mode 100644 benchmarks/search_scaling_rich_metrics.tsv create mode 100644 src/rescue/gff/tests.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 1375bf3..4736b05 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,23 @@ and Recore uses semantic versioning. ### Fixed +- Memory no longer grows by megabytes with every genome. Each genome's + annotation was held as a String per name in up to six maps, and its + gene-to-cluster lookup as a map of every name, copied again for the corrected + matrix: 3.4 MB per genome with minimal GFFs and 6.2 MB with PGAP-style ones + (six names per CDS), so 5,000 genomes would have needed about 17 and 32 GB. + Each name is now stored once, in one buffer per genome with a sorted table + of hashes; the lookup keeps one cluster per name and is corrected in place; + matrix cells are one string per cluster; and `cluster_synteny_blocks.tsv` and + `cluster_locus_context.tsv` are built in chunks bounded by genome-cluster + pairs instead of by clusters, the locus contexts now in parallel too. On 800 + synthetic genomes peak memory went from 2.8 to 0.56 GB (5.5 to 1.0 GB with + PGAP-style GFFs) and run time from 33 to 22 seconds. Every output is byte for + byte the same. +- When two corrections in one genome named the same gene ID, the cluster its + neighbours took in `cluster_locus_context.tsv` and + `cluster_synteny_blocks.tsv` depended on hash-map order and could change from + run to run. The first cluster in matrix order now wins. - Large runs are faster and use less memory, with identical outputs. `cluster_synteny_blocks.tsv` held every cluster's blocks before writing the first and was built in one thread: on 200 genomes of 4,000 genes, half the run diff --git a/benchmarks/README.md b/benchmarks/README.md index 6f75d30..4169dbe 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -149,6 +149,12 @@ equal the false absences planted). The datasets are generated afresh and removed afterwards unless `--keep-workdir` is given. The numbers belong to the machine they ran on, which the table records; the shape is what carries over. +The GFFs carry an ID and a locus tag. `--rich-gff` writes PGAP-style records +instead (a name, a gene symbol, an old locus tag, a protein ID, +cross-references, a product and a long note), six names per CDS, which is what +an annotation index has to hold for most real assemblies; that run goes to +`benchmarks/search_scaling_rich_metrics.tsv`. + ## AMAP performance benchmark `amap_performance_benchmark.py` profiles Recore on the full local AMAP diff --git a/benchmarks/search_scaling.py b/benchmarks/search_scaling.py index edc52d7..c0e984b 100644 --- a/benchmarks/search_scaling.py +++ b/benchmarks/search_scaling.py @@ -33,6 +33,7 @@ REPO_ROOT = Path(__file__).resolve().parents[1] DEFAULT_METRICS = REPO_ROOT / "benchmarks" / "search_scaling_metrics.tsv" +DEFAULT_RICH_METRICS = REPO_ROOT / "benchmarks" / "search_scaling_rich_metrics.tsv" BASES = np.frombuffer(b"ACGT", dtype=np.uint8) STOPS = {b"TAA", b"TAG", b"TGA"} @@ -44,7 +45,33 @@ def random_orf(rng: np.random.Generator, codons: int) -> bytes: return b"ATG" + b"".join(sense[:codons]) + b"TAA" -def build_dataset(out: Path, n_genomes: int, n_genes: int, gaps: int, seed: int) -> dict: +PRODUCTS = [ + "DNA-directed RNA polymerase subunit beta", + "ATP-dependent Clp protease proteolytic subunit", + "PE family protein", + "hypothetical protein", + "ABC transporter ATP-binding protein", + "transcriptional regulator, TetR family", +] + + +def gff_attributes(genome: str, i: int, rich: bool) -> str: + """The CDS attributes: an ID and a locus tag, or a PGAP-style record with a + name, a product, a long note, cross-references and a protein ID, which is + what most real annotations carry and what an annotation index holds.""" + tag = f"{genome}_{i:05d}" + if not rich: + return f"ID={tag};locus_tag={tag}" + return ( + f"ID=cds-{tag};Parent=gene-{tag};Name=WP_{i:09d}.1;gene=gen{i:05d};locus_tag={tag};" + f"old_locus_tag=Rv{i:04d};protein_id=WP_{i:09d}.1;Dbxref=GenBank:WP_{i:09d}.1,GeneID:{1000000 + i};" + f"product={PRODUCTS[i % len(PRODUCTS)]};" + "Note=Derived by automated computational analysis using gene prediction method%3A Protein Homology. " + "GO_function%3A GO%3A0003677 - DNA binding [Evidence IEA]" + ) + + +def build_dataset(out: Path, n_genomes: int, n_genes: int, gaps: int, seed: int, rich_gff: bool = False) -> dict: """Write pa.csv, pan_ref.fa, genomes/ and annotations/ under `out`.""" rng = np.random.default_rng(seed) genes = [random_orf(rng, int(rng.integers(250, 500))) for _ in range(n_genes)] @@ -82,7 +109,7 @@ def build_dataset(out: Path, n_genomes: int, n_genes: int, gaps: int, seed: int) if i not in absent[g]: gff_lines.append( f"{genome}_1\tsynthetic\tCDS\t{pos + 1}\t{pos + len(gene)}\t.\t+\t0\t" - f"ID={genome}_{i:05d};locus_tag={genome}_{i:05d}" + + gff_attributes(genome, i, rich_gff) ) pos += len(gene) contig = bytearray(b"".join(seq_parts)) @@ -161,11 +188,19 @@ def main() -> int: parser.add_argument("--gaps", type=int, default=80, help="matrix gaps per genome, half of them false") parser.add_argument("--threads", type=int, default=8) parser.add_argument("--seed", type=int, default=1) + parser.add_argument( + "--rich-gff", + action="store_true", + help="PGAP-style annotations (names, products, notes, cross-references) instead of an ID and a locus tag", + ) parser.add_argument("--minimap2") parser.add_argument("--recore", help="recore binary (default: target/release/recore, built if missing)") parser.add_argument("--workdir", help="where datasets and outputs go (default: a temporary directory)") parser.add_argument("--keep-workdir", action="store_true") - parser.add_argument("--metrics", default=str(DEFAULT_METRICS)) + parser.add_argument( + "--metrics", + help="output TSV (default: search_scaling_metrics.tsv, or search_scaling_rich_metrics.tsv with --rich-gff)", + ) args = parser.parse_args() minimap2 = find_minimap2(args.minimap2) @@ -187,7 +222,7 @@ def main() -> int: if data.exists(): shutil.rmtree(data) t0 = time.monotonic() - info = build_dataset(data, n, args.genes, args.gaps, args.seed) + info = build_dataset(data, n, args.genes, args.gaps, args.seed, args.rich_gff) print(f"{n} genomes: dataset in {time.monotonic() - t0:.0f}s ({info['genome_mb']} Mb each)", flush=True) cache = workdir / f"cache_{n}" if cache.exists(): @@ -199,6 +234,7 @@ def main() -> int: "genomes": n, "genome_mb": info["genome_mb"], "genes_per_genome": args.genes, + "annotation": "rich" if args.rich_gff else "minimal", "gaps_searched": plain["searched"], "false_absences": info["false_absences"], "rescued": plain["rescued"], @@ -223,7 +259,7 @@ def main() -> int: if not args.keep_workdir and not args.workdir: shutil.rmtree(workdir, ignore_errors=True) - metrics = Path(args.metrics) + metrics = Path(args.metrics) if args.metrics else DEFAULT_RICH_METRICS if args.rich_gff else DEFAULT_METRICS metrics.parent.mkdir(parents=True, exist_ok=True) with open(metrics, "w", newline="") as fh: writer = csv.DictWriter(fh, fieldnames=list(rows[0].keys()), delimiter="\t") diff --git a/benchmarks/search_scaling_metrics.tsv b/benchmarks/search_scaling_metrics.tsv index 64941c3..e0bff16 100644 --- a/benchmarks/search_scaling_metrics.tsv +++ b/benchmarks/search_scaling_metrics.tsv @@ -1,5 +1,5 @@ -genomes genome_mb genes_per_genome gaps_searched false_absences rescued threads seconds_no_cache seconds_cold_cache seconds_warm_cache seconds_per_genome peak_rss_mb recore minimap2 machine -100 4.91 4000 8000 4000 3999 8 3.4 3.4 2.4 0.034 461 recore 0.2.0 2.31-r1302 arm64 Darwin 14 cpus -200 4.91 4000 16000 8000 8000 8 7.1 7.9 6.1 0.035 762 recore 0.2.0 2.31-r1302 arm64 Darwin 14 cpus -400 4.91 4000 32000 16000 16000 8 18.0 18.5 13.0 0.045 1456 recore 0.2.0 2.31-r1302 arm64 Darwin 14 cpus -800 4.91 4000 64000 32000 32000 8 48.1 50.1 38.1 0.06 2570 recore 0.2.0 2.31-r1302 arm64 Darwin 14 cpus +genomes genome_mb genes_per_genome annotation gaps_searched false_absences rescued threads seconds_no_cache seconds_cold_cache seconds_warm_cache seconds_per_genome peak_rss_mb recore minimap2 machine +100 4.91 4000 minimal 8000 4000 3999 8 2.4 2.5 1.3 0.024 236 recore 0.2.0 2.31-r1302 arm64 Darwin 14 cpus +200 4.91 4000 minimal 16000 8000 8000 8 4.7 5.1 2.8 0.024 300 recore 0.2.0 2.31-r1302 arm64 Darwin 14 cpus +400 4.91 4000 minimal 32000 16000 16000 8 10.7 10.5 6.3 0.027 377 recore 0.2.0 2.31-r1302 arm64 Darwin 14 cpus +800 4.91 4000 minimal 64000 32000 32000 8 21.4 22.1 13.5 0.027 556 recore 0.2.0 2.31-r1302 arm64 Darwin 14 cpus diff --git a/benchmarks/search_scaling_rich_metrics.tsv b/benchmarks/search_scaling_rich_metrics.tsv new file mode 100644 index 0000000..72cbc65 --- /dev/null +++ b/benchmarks/search_scaling_rich_metrics.tsv @@ -0,0 +1,5 @@ +genomes genome_mb genes_per_genome annotation gaps_searched false_absences rescued threads seconds_no_cache seconds_cold_cache seconds_warm_cache seconds_per_genome peak_rss_mb recore minimap2 machine +100 4.91 4000 rich 8000 4000 3999 8 2.2 2.2 1.1 0.022 322 recore 0.2.0 2.31-r1302 arm64 Darwin 14 cpus +200 4.91 4000 rich 16000 8000 8000 8 4.2 4.3 2.3 0.021 430 recore 0.2.0 2.31-r1302 arm64 Darwin 14 cpus +400 4.91 4000 rich 32000 16000 16000 8 8.8 9.1 5.0 0.022 637 recore 0.2.0 2.31-r1302 arm64 Darwin 14 cpus +800 4.91 4000 rich 64000 32000 32000 8 20.7 22.0 12.7 0.026 1053 recore 0.2.0 2.31-r1302 arm64 Darwin 14 cpus diff --git a/docs/how-it-works/algorithm.md b/docs/how-it-works/algorithm.md index bbcfa5a..5b21edd 100644 --- a/docs/how-it-works/algorithm.md +++ b/docs/how-it-works/algorithm.md @@ -83,7 +83,11 @@ For each genome with a GFF file, builds a `GffIndex`: 1. Genes sorted by (contig, start position) 2. Per-contig range index for O(log n) overlap queries via binary search (`partition_point`) -3. Gene ID index for O(1) lookup by name +3. Every name of every CDS (its gene ID, then its aliases) stored once, back to back in one buffer, with a sorted table of name hashes for lookup by name +4. Products and notes read once to recognise mobile elements, then dropped + +A CDS costs 32 bytes plus its names, so the index grows with a genome's gene +count and the length of its names, not with how many maps refer to them. The index supports: - `overlapping(contig, start, end)` → all CDS features overlapping a query region @@ -94,7 +98,7 @@ The index supports: ## Phase 4: Gene-to-Cluster Maps -For each genome involved in the search, builds a HashMap from gene ID → ClusterIdx. This maps every gene ID found in the CSV cells to its Panaroo cluster assignment. +For each genome with an annotation, builds a map from gene ID → ClusterIdx. This maps every gene ID found in the CSV cells to its Panaroo cluster assignment. Panaroo appends suffixes to gene IDs that don't exist in the original GFF: - `_pseudo` for pseudogene annotations @@ -102,6 +106,11 @@ Panaroo appends suffixes to gene IDs that don't exist in the original GFF: The map indexes both the original ID and stripped versions to maximize GFF matching success. +The map is then read into one cluster for each name of the genome's +`GffIndex` (4 bytes each, a `GeneClusters`) and dropped. Every later lookup is +for a name the annotation uses, so this answers each one as the map did. The +genomes are done in parallel. + --- ## Phase 5: Search (minimap2) diff --git a/docs/how-it-works/limits.md b/docs/how-it-works/limits.md index f2dc93e..f414fb9 100644 --- a/docs/how-it-works/limits.md +++ b/docs/how-it-works/limits.md @@ -108,28 +108,36 @@ Measured with [`benchmarks/search_scaling.py`](https://github.com/PathoGenOmics- on an Apple M4 Pro (14 cores, 24 GB) with 8 threads and minimap2 2.31: synthetic MTBC-sized genomes (one 4.9 Mb contig and 4,000 genes each), 80 matrix gaps per genome of which 40 are false absences, GFFs and `--synteny`. +The GFFs give each CDS an ID and a locus tag; the last two columns are the +same runs with PGAP-style GFFs (`--rich-gff`), six names per CDS plus a +product and a note. -| Genomes | Gaps searched | Run time | Warm `--search-cache` | Peak memory | -|---------|---------------|----------|-----------------------|-------------| -| 100 | 8,000 | 3.4 s | 2.4 s | 0.5 GB | -| 200 | 16,000 | 7.1 s | 6.1 s | 0.8 GB | -| 400 | 32,000 | 18.0 s | 13.0 s | 1.5 GB | -| 800 | 64,000 | 48.1 s | 38.1 s | 2.6 GB | +| Genomes | Gaps searched | Run time | Warm `--search-cache` | Peak memory | PGAP-style GFFs: run time | PGAP-style GFFs: peak memory | +|---------|---------------|----------|-----------------------|-------------|---------------------------|------------------------------| +| 100 | 8,000 | 2.4 s | 1.3 s | 0.24 GB | 2.2 s | 0.32 GB | +| 200 | 16,000 | 4.7 s | 2.8 s | 0.30 GB | 4.2 s | 0.43 GB | +| 400 | 32,000 | 10.7 s | 6.3 s | 0.38 GB | 8.8 s | 0.64 GB | +| 800 | 64,000 | 21.4 s | 13.5 s | 0.56 GB | 20.7 s | 1.05 GB | Every planted false absence was recovered, except one at 100 genomes that synteny rejected because the random deletions had changed its neighbourhood. -Memory grows by about 3 MB per genome, mostly the annotations synteny keeps for -every genome. Time grows a little faster than the genome count, from 0.03 s -per genome at 100 genomes to 0.06 s at 800, and the minimap2 search is the -smaller part of it (10 of the 48 seconds at 800 genomes). The rest is synteny, -and the audit tables in `cluster_synteny_blocks.tsv`, whose size follows the -number of distinct gene neighbourhoods. These synthetic genomes are a worst -case for that: every gap reshuffles the context of ten neighbours, and at 800 -genomes the table reached 2.3 GB. Assemblies with conserved synteny write far -less. - -`--search-cache` removes the search from a rerun (10 seconds to 1 at 800 +Memory grows by about 0.5 MB per genome, or 1 MB with PGAP-style GFFs, nearly +all of it the annotation synteny reads for every genome: 5,000 genomes like +these would need about 2.5 GB, or 5.5 GB with PGAP-style GFFs. Each name is +stored once per genome and the gene-to-cluster lookup keeps one cluster per +name; when every name was a String in several maps, memory grew by 3.4 MB per +genome (6.2 MB with PGAP-style GFFs). Time grows with the genome count, about +0.027 s per genome, and the minimap2 search is about 9 of the 21 seconds at 800 +genomes. The rest is synteny, and the audit tables in +`cluster_synteny_blocks.tsv`, whose size follows the number of distinct gene +neighbourhoods. These synthetic genomes are a worst case for that: every gap +reshuffles the context of ten neighbours, and at 800 genomes the table reached +2.3 GB. Assemblies with conserved synteny write far less, and the table is +written a bounded chunk of clusters at a time, so its size does not show in the +memory. + +`--search-cache` removes the search from a rerun (21 seconds to 13.5 at 800 genomes), so it pays off most where the search is the expensive part: many gaps per genome (`--min-prevalence 0`), `--multi-allele`, or larger genomes. The cache directory can be deleted at any time; the next run simply searches again. diff --git a/src/parser.rs b/src/parser.rs index ff58710..150cc4d 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -8,6 +8,39 @@ use crate::types::{ClusterIdx, GenomeIdx}; // ── Presence/absence matrix (encapsulated) ── +/// One cluster's cells, back to back in one string. A String per cell cost 24 +/// bytes plus an allocation, which with gene IDs a dozen characters long was +/// most of what the matrix held for each genome. +struct CellRow { + text: Box, + /// Where each genome's cell ends in `text`; each starts where the one + /// before ends. + ends: Box<[u32]>, +} + +impl CellRow { + fn new<'a>(cells: impl IntoIterator) -> Self { + let mut text = String::new(); + let mut ends = Vec::new(); + for cell in cells { + text.push_str(cell); + ends.push(u32::try_from(text.len()).expect("a matrix row under 4 GB")); + } + CellRow { + text: text.into_boxed_str(), + ends: ends.into_boxed_slice(), + } + } + + fn get(&self, gi: usize) -> &str { + let start = match gi { + 0 => 0, + _ => self.ends[gi - 1] as usize, + }; + &self.text[start..self.ends[gi] as usize] + } +} + /// Presence/absence matrix from Panaroo/Roary output. /// /// Fields are private to enforce consistent access and prevent @@ -17,7 +50,7 @@ pub struct PresenceAbsence { genomes: Vec, present_genomes: Vec>, present_bits: Vec>, - cell_values: Vec>, + cell_values: Vec, header: Vec, metadata: Vec>, genome_col_indices: Vec, @@ -80,7 +113,7 @@ impl PresenceAbsence { #[inline] pub fn cell(&self, ci: ClusterIdx, gi: GenomeIdx) -> &str { - &self.cell_values[ci.val()][gi.val()] + self.cell_values[ci.val()].get(gi.val()) } // ── CSV structure (for output reconstruction) ── @@ -162,6 +195,10 @@ impl PresenceAbsence { let genome_col_indices: Vec = (0..n_ge).collect(); let mut header = vec!["Gene".to_string()]; header.extend(genomes.iter().cloned()); + let cell_values = cell_values + .iter() + .map(|row| CellRow::new(row.iter().map(String::as_str))) + .collect(); PresenceAbsence { clusters, genomes, @@ -271,16 +308,14 @@ pub fn parse_presence_absence(path: &Path) -> Result { let mut present_row = Vec::new(); let mut bit_row = vec![false; genomes.len()]; - let mut row_values = Vec::with_capacity(genomes.len()); - - for (gi, &col_idx) in genome_indices.iter().enumerate() { - let val = record.get(col_idx).unwrap_or("").trim().to_string(); + let row_values = CellRow::new(genome_indices.iter().enumerate().map(|(gi, &col_idx)| { + let val = record.get(col_idx).unwrap_or("").trim(); if !val.is_empty() { present_row.push(GenomeIdx(gi)); bit_row[gi] = true; } - row_values.push(val); - } + val + })); let meta_values: Vec = meta_col_indices .iter() @@ -394,15 +429,14 @@ fn parse_pirate_gene_families(path: &Path) -> Result { let field = |col: Option| col.and_then(|i| record.get(i)).unwrap_or("").to_string(); let mut present_row = Vec::new(); let mut bit_row = vec![false; genomes.len()]; - let mut row_values = Vec::with_capacity(genomes.len()); - for (gi, &col) in genome_cols.iter().enumerate() { - let value = record.get(col).unwrap_or("").trim().to_string(); + let row_values = CellRow::new(genome_cols.iter().enumerate().map(|(gi, &col)| { + let value = record.get(col).unwrap_or("").trim(); if !value.is_empty() { present_row.push(GenomeIdx(gi)); bit_row[gi] = true; } - row_values.push(value); - } + value + })); metadata.push(vec![ field(Some(family_col)), field(gene_name_col), @@ -1373,13 +1407,30 @@ mod tests { Accessory Order,QC,Min group size nuc,Max group size nuc,\ Avg group size nuc,g1,g2,g3\n\ gA,gA,prod,3,3,1.0,0,0,0,,100,100,100,id1,id2,id3\n\ - gB,gB,prod,2,2,1.0,0,0,0,,200,200,200,id4,,id5\n", + gB,gB,prod,2,2,1.0,0,0,0,,200,200,200,id4,,id5\n\ + gC,gC,prod,1,2,2.0,0,0,0,,300,300,300,, id6;id7 ,\n", ); let pa = parse_presence_absence(csv.path()).unwrap(); - assert_eq!(pa.n_clusters(), 2); + assert_eq!(pa.n_clusters(), 3); assert_eq!(pa.n_genomes(), 3); assert_eq!(pa.core_size(), 1); - assert_eq!(pa.gaps().len(), 1); + assert_eq!(pa.gaps().len(), 3); + // Every cell comes back as written, trimmed, empty ones included. + let cells: Vec> = (0..3) + .map(|ci| { + (0..3) + .map(|gi| pa.cell(ClusterIdx(ci), GenomeIdx(gi))) + .collect() + }) + .collect(); + assert_eq!( + cells, + vec![ + vec!["id1", "id2", "id3"], + vec!["id4", "", "id5"], + vec!["", "id6;id7", ""], + ] + ); } #[test] diff --git a/src/pipeline.rs b/src/pipeline.rs index f1074df..18c1526 100644 --- a/src/pipeline.rs +++ b/src/pipeline.rs @@ -15,10 +15,11 @@ use crate::parser::{ }; use crate::relationships::{write_cluster_components, write_cluster_relationships}; use crate::rescue::{ - build_gene_to_cluster_with_gff_aliases, cell_gene_id_variants, check_internal_stops_in_cds, + build_gene_clusters, cell_gene_id_variants, check_internal_stops_in_cds, check_internal_stops_in_hit, check_internal_stops_in_span, classify_hits, - panaroo_gene_id_suffix_variants, resolve_conflicts, validate_synteny_with_cache, GffGene, - GffIndex, PseudogeneReason, RescueEvent, RescueType, SyntenyEvidence, SyntenyReferenceCache, + panaroo_gene_id_suffix_variants, resolve_conflicts, validate_synteny_with_cache, GeneClusters, + GffGene, GffIndex, GffRecord, PseudogeneReason, RescueEvent, RescueType, SyntenyEvidence, + SyntenyReferenceCache, }; use crate::search::{batch_search, Hit}; use crate::types::{ClusterIdx, GenomeIdx}; @@ -393,23 +394,25 @@ pub fn run(args: &Args) -> Result<()> { HashMap::new() }; - // ── 5. Build gene_to_cluster maps ONCE ── - let gene_to_cluster_maps: HashMap> = { - let mut maps = HashMap::new(); - for (gi, _) in &searchable_genomes { - maps.entry(*gi).or_insert_with(|| { - build_gene_to_cluster_with_gff_aliases(&pa, *gi, gff_indices.get(gi)) - }); - } - if args.synteny || !gff_indices.is_empty() { - for &gi in gff_indices.keys() { - maps.entry(gi).or_insert_with(|| { - build_gene_to_cluster_with_gff_aliases(&pa, gi, gff_indices.get(&gi)) - }); - } - } - maps - }; + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(args.threads) + .build() + .context("Failed to build rayon thread pool")?; + + // ── 5. Build the gene-to-cluster lookups ONCE ── + // Every lookup is for a name a genome's annotation uses, so a genome + // without one needs none. Each genome's map of names is built in + // parallel, read into its index's names and dropped; see GeneClusters. + let mut gene_to_cluster_maps: HashMap = pool.install(|| { + use rayon::prelude::*; + let genomes: Vec<(&GenomeIdx, &GffIndex)> = gff_indices.iter().collect(); + genomes + .into_par_iter() + .map(|(&gi, gff)| (gi, build_gene_clusters(&pa, gi, gff))) + .collect::>() + .into_iter() + .collect() + }); // ── 6. Search each gap genome for its absent clusters ── // @@ -502,17 +505,12 @@ pub fn run(args: &Args) -> Result<()> { }; // ── 7. Classify hits (parallel) ── - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(args.threads) - .build() - .context("Failed to build rayon thread pool")?; - let classify_inputs: Vec<(ClusterIdx, GenomeIdx, &[crate::search::Hit])> = hits_map .iter() .map(|((gi, ci), hits)| (*ci, *gi, hits.as_slice())) .collect(); - let empty_map: HashMap = HashMap::new(); + let no_clusters = GeneClusters::none(); let rescues: Vec = pool.install(|| { use rayon::prelude::*; @@ -520,7 +518,7 @@ pub fn run(args: &Args) -> Result<()> { .par_iter() .filter_map(|&(ci, gi, hits)| { let gff_index = gff_indices.get(&gi); - let g2c = gene_to_cluster_maps.get(&gi).unwrap_or(&empty_map); + let g2c = gene_to_cluster_maps.get(&gi).unwrap_or(&no_clusters); classify_hits(hits, ci, gi, &pa, gff_index, g2c, args.min_coverage) }) .collect() @@ -921,27 +919,54 @@ pub fn run(args: &Args) -> Result<()> { write_reference_cluster_events(&reference_cluster_events, &pa, &args.output)?; } } + write_multistate_matrix(MultistateMatrixInputs { + pa: &pa, + accepted_rescue_map: &all_rescue_map, + nsc_rescue_map: &matrix_rescue_map, + synteny_rejected: &synteny_rejected, + paralog_rejected: ¶log_rejected, + gene_to_cluster_maps: &gene_to_cluster_maps, + output_dir: &args.output, + nsc_result: &nsc_result, + coverage: &coverage, + })?; // Contexts and synteny blocks are read off the corrected matrix: a CDS this // run rescued belongs to its cluster like any other neighbour. Read off the // input matrix, the falsely absent gene was a CDS with no cluster, it was // skipped, and the window slid along by one. Every single-copy neighbour of // a false absence then had two contexts and was reported as a duplication // or paralog expansion: 12 of the tutorial's 24 clusters, DnaA included. - let mut corrected_gene_to_cluster_maps = gene_to_cluster_maps.clone(); - for rescue in matrix_rescue_map.values() { - if let Some(gene_id) = rescue.annotation_match.as_deref() { - corrected_gene_to_cluster_maps - .entry(rescue.genome_idx) - .or_default() - .entry(gene_id.to_string()) - .or_insert(rescue.cluster_idx); + // + // The lookups are corrected in place, so this comes after everything that + // reads the input matrix's clusters (the multistate matrix last). The + // corrections used to go into a copy of every map, filled in the order of a + // hash map; they are taken in cluster order, so when two corrections name + // one gene ID the same one wins from run to run. + let mut corrections: Vec<(GenomeIdx, ClusterIdx, &str)> = matrix_rescue_map + .values() + .filter_map(|rescue| { + rescue + .annotation_match + .as_deref() + .map(|gene_id| (rescue.genome_idx, rescue.cluster_idx, gene_id)) + }) + .collect(); + corrections.sort_unstable(); + for genome_corrections in corrections.chunk_by(|a, b| a.0 == b.0) { + if let Some(g2c) = gene_to_cluster_maps.get_mut(&genome_corrections[0].0) { + let names: Vec<(&str, ClusterIdx)> = genome_corrections + .iter() + .map(|&(_, ci, gene_id)| (gene_id, ci)) + .collect(); + g2c.fill(&names); } } let context_summaries = if !gff_indices.is_empty() { let (context_rows, summaries) = write_cluster_locus_contexts( &pa, &gff_indices, - &corrected_gene_to_cluster_maps, + &gene_to_cluster_maps, + args.threads, &args.output, )?; let synteny_reference_gi = args @@ -951,7 +976,7 @@ pub fn run(args: &Args) -> Result<()> { let block_rows = write_cluster_synteny_blocks( &pa, &gff_indices, - &corrected_gene_to_cluster_maps, + &gene_to_cluster_maps, synteny_reference_gi, args.threads, &args.output, @@ -967,17 +992,6 @@ pub fn run(args: &Args) -> Result<()> { Vec::new() }; let corrected_counts = write_corrected_csv(&pa, &matrix_rescue_map, &args.output, &nsc_result)?; - write_multistate_matrix(MultistateMatrixInputs { - pa: &pa, - accepted_rescue_map: &all_rescue_map, - nsc_rescue_map: &matrix_rescue_map, - synteny_rejected: &synteny_rejected, - paralog_rejected: ¶log_rejected, - gene_to_cluster_maps: &gene_to_cluster_maps, - output_dir: &args.output, - nsc_result: &nsc_result, - coverage: &coverage, - })?; let n_relationships = write_cluster_relationships(&pa, &rescues, &nsc_result, &args.output)?; if n_relationships > 0 { info!("Cluster relationships: {} rows", n_relationships); @@ -1335,7 +1349,7 @@ fn clean_tsv_field(value: &str) -> String { .collect() } -fn write_tsv_fields(w: &mut BufWriter, fields: &[String]) -> Result<()> { +fn write_tsv_fields(w: &mut impl Write, fields: &[String]) -> Result<()> { let cleaned: Vec = fields.iter().map(|field| clean_tsv_field(field)).collect(); writeln!(w, "{}", cleaned.join("\t"))?; Ok(()) @@ -1425,7 +1439,7 @@ fn scan_genome_for_internal_stops( && rescue.start < gene.end }); let stops = if let Some(gene) = annotated_cds { - check_internal_stops_in_cds(&gene.contig, gene.start, gene.end, gene.strand, &contigs) + check_internal_stops_in_cds(gene.contig, gene.start, gene.end, gene.strand, &contigs) } else if let Some(query_start) = rescue.query_start { // No annotation, but one alignment: its start on the representative // fixes the reading frame. @@ -1530,11 +1544,11 @@ fn load_genome_gff( ))); } } - let genes: Vec<(GffGene, Vec, Vec)> = raw_genes + let genes: Vec<(GffRecord, Vec, Vec)> = raw_genes .into_iter() .map(|(contig, start, end, strand, gene_id, aliases, labels)| { ( - GffGene { + GffRecord { contig, start, end, @@ -1553,13 +1567,13 @@ fn load_genome_gff( }) } -fn gff_gene_for_id<'a>(gff: &'a GffIndex, gene_id: &str) -> Option<&'a GffGene> { +fn gff_gene_for_id<'a>(gff: &'a GffIndex, gene_id: &str) -> Option> { cell_gene_id_variants(gene_id) .into_iter() .find_map(|id| gff.get_by_id(id)) } -fn gff_gene_for_cell<'a>(gff: &'a GffIndex, cell: &str) -> Option<&'a GffGene> { +fn gff_gene_for_cell<'a>(gff: &'a GffIndex, cell: &str) -> Option> { cell.split(';') .map(str::trim) .filter(|id| !id.is_empty()) @@ -1569,7 +1583,7 @@ fn gff_gene_for_cell<'a>(gff: &'a GffIndex, cell: &str) -> Option<&'a GffGene> { fn annotation_cluster_name( rescue: &RescueEvent, pa: &crate::parser::PresenceAbsence, - gene_to_cluster_maps: &HashMap>, + gene_to_cluster_maps: &HashMap, ) -> String { let Some(annotation_match) = rescue.annotation_match.as_deref() else { return String::new(); @@ -1578,7 +1592,7 @@ fn annotation_cluster_name( return String::new(); }; g2c.get(annotation_match) - .map(|ci| pa.cluster_name(*ci).to_string()) + .map(|ci| pa.cluster_name(ci).to_string()) .unwrap_or_default() } @@ -1711,7 +1725,7 @@ fn matrix_action_with_quality( fn accepted_rescue_evidence_state( rescue: &RescueEvent, pa: &crate::parser::PresenceAbsence, - gene_to_cluster_maps: &HashMap>, + gene_to_cluster_maps: &HashMap, ) -> &'static str { let query_cluster = pa.cluster_name(rescue.cluster_idx); let annotation_cluster = annotation_cluster_name(rescue, pa, gene_to_cluster_maps); @@ -2001,7 +2015,7 @@ struct OrthologyEvidenceInputs<'a> { synteny_rejected: &'a [RescueEvent], paralog_rejected: &'a [RescueEvent], pa: &'a crate::parser::PresenceAbsence, - gene_to_cluster_maps: &'a HashMap>, + gene_to_cluster_maps: &'a HashMap>, cluster_prevalences: &'a [f64], cluster_seq_lengths: &'a [usize], matrix_policy: MatrixPolicy, diff --git a/src/pipeline/evolution.rs b/src/pipeline/evolution.rs index 3b80fd6..8e003db 100644 --- a/src/pipeline/evolution.rs +++ b/src/pipeline/evolution.rs @@ -732,31 +732,6 @@ pub(super) fn corrected_cluster_present( || nsc_presence.contains(&(ci.val(), gi.val())) } -pub(super) fn mobile_kind_from_labels(labels: &[&str]) -> Option<(&'static str, String)> { - let joined = labels.join("|"); - let normalized = joined - .to_ascii_lowercase() - .replace("%20", " ") - .replace('_', " "); - let kind = if normalized.contains("is6110") { - "IS6110" - } else if normalized.contains("is1081") { - "IS1081" - } else if normalized.contains("insertion sequence") - || normalized.contains("transposase") - || normalized.contains("is element") - { - "insertion_sequence" - } else if normalized.contains("integrase") || normalized.contains("phage") { - "phage_or_integrase" - } else if normalized.contains("recombinase") { - "recombinase" - } else { - return None; - }; - Some((kind, joined)) -} - pub(super) struct MobileElementSearchInput<'a> { pub(super) path: &'a Path, pub(super) genome_files: &'a HashMap, @@ -909,7 +884,7 @@ pub(super) fn nearest_mobile_element_for_rescue( let Some(gff) = gff_indices.get(&rescue.genome_idx) else { return best; }; - for gene in &gff.genes { + for gene in gff.genes() { if gene.contig != rescue.contig { continue; } @@ -917,10 +892,12 @@ pub(super) fn nearest_mobile_element_for_rescue( if distance > mobile_window { continue; } - let labels = gff.labels_for_gene(gene); - let Some((kind, label)) = mobile_kind_from_labels(&labels) else { + // Worked out from the gene's names and descriptions when the GFF was + // read; see GffIndex::mobile_annotation. + let Some((kind, label)) = gff.mobile_annotation(&gene) else { continue; }; + let label = label.to_string(); let candidate = MobileElementEvidence { genome_name: pa.genome_name(rescue.genome_idx).to_string(), source: "annotation", @@ -1994,14 +1971,13 @@ pub(super) struct ClusterContextSummary { } pub(super) fn context_signature( - gff: &GffIndex, - genes: &[&GffGene], - gene_to_cluster: &HashMap, + genes: &[GffGene], + gene_to_cluster: &GeneClusters, pa: &crate::parser::PresenceAbsence, ) -> String { let names: Vec<&str> = genes .iter() - .filter_map(|gene| gff.cluster_for_gene(gene, gene_to_cluster)) + .filter_map(|gene| gene_to_cluster.cluster_for_gene(gene)) .map(|ci| pa.cluster_name(ci)) .collect(); if names.is_empty() { @@ -2047,11 +2023,11 @@ pub(super) fn context_interpretation( fn synteny_gene_record( gff: &GffIndex, gene: &GffGene, - gene_to_cluster: &HashMap, + gene_to_cluster: &GeneClusters, pa: &crate::parser::PresenceAbsence, ) -> SyntenyGeneRecord { - let cluster = gff - .cluster_for_gene(gene, gene_to_cluster) + let cluster = gene_to_cluster + .cluster_for_gene(gene) .map(|ci| pa.cluster_name(ci).to_string()) .unwrap_or_else(|| "NA".to_string()); let display_gene_id = gff @@ -2059,11 +2035,11 @@ fn synteny_gene_record( .into_iter() .map(display_reference_gene_id) .find(|id| id.starts_with("Rv")) - .unwrap_or_else(|| display_reference_gene_id(&gene.gene_id)); + .unwrap_or_else(|| display_reference_gene_id(gene.gene_id)); SyntenyGeneRecord { cluster, gene_id: display_gene_id, - contig: gene.contig.clone(), + contig: gene.contig.to_string(), start: gene.start + 1, end: gene.end, strand: gene.strand, @@ -2093,10 +2069,10 @@ fn synteny_gap_record(contig: &str, start_0based: u64, end_0based: u64) -> Synte fn ordered_synteny_gene_records( gff: &GffIndex, - upstream: &[&GffGene], + upstream: &[GffGene], center: Option<&GffGene>, - downstream: &[&GffGene], - gene_to_cluster: &HashMap, + downstream: &[GffGene], + gene_to_cluster: &GeneClusters, pa: &crate::parser::PresenceAbsence, ) -> Vec { let mut records = @@ -2158,32 +2134,40 @@ fn synteny_gene_identity_signature(records: &[SyntenyGeneRecord]) -> String { fn build_reference_boundary_index( gff: &GffIndex, - gene_to_cluster: &HashMap, + gene_to_cluster: &GeneClusters, pa: &crate::parser::PresenceAbsence, n_neighbors: usize, ) -> ReferenceBoundaryIndex { let mut index = ReferenceBoundaryIndex::default(); - if gff.genes.is_empty() { + let all: Vec = gff.genes().collect(); + if all.is_empty() { return index; } let mut range_start = 0usize; - while range_start < gff.genes.len() { - let contig = gff.genes[range_start].contig.clone(); + while range_start < all.len() { + let contig = all[range_start].contig.to_string(); let mut range_end = range_start + 1; - while range_end < gff.genes.len() && gff.genes[range_end].contig == contig { + while range_end < all.len() && all[range_end].contig == contig { range_end += 1; } - let slice = &gff.genes[range_start..range_end]; + let slice = &all[range_start..range_end]; for (gene_idx, gene) in slice.iter().enumerate() { - let Some(focal_ci) = gff.cluster_for_gene(gene, gene_to_cluster) else { + let Some(focal_ci) = gene_to_cluster.cluster_for_gene(gene) else { continue; }; index.by_focal.entry(focal_ci).or_insert_with(|| { - let upstream: Vec<&GffGene> = - slice[..gene_idx].iter().rev().take(n_neighbors).collect(); - let downstream: Vec<&GffGene> = - slice[gene_idx + 1..].iter().take(n_neighbors).collect(); + let upstream: Vec = slice[..gene_idx] + .iter() + .rev() + .take(n_neighbors) + .copied() + .collect(); + let downstream: Vec = slice[gene_idx + 1..] + .iter() + .take(n_neighbors) + .copied() + .collect(); ordered_synteny_gene_records( gff, &upstream, @@ -2195,9 +2179,17 @@ fn build_reference_boundary_index( }); } for boundary in 0..=slice.len() { - let upstream: Vec<&GffGene> = - slice[..boundary].iter().rev().take(n_neighbors).collect(); - let downstream: Vec<&GffGene> = slice[boundary..].iter().take(n_neighbors).collect(); + let upstream: Vec = slice[..boundary] + .iter() + .rev() + .take(n_neighbors) + .copied() + .collect(); + let downstream: Vec = slice[boundary..] + .iter() + .take(n_neighbors) + .copied() + .collect(); let gap_start = if boundary > 0 { slice[boundary - 1].end } else { @@ -2221,7 +2213,7 @@ fn build_reference_boundary_index( let mut cluster_ids: Vec = upstream .iter() .chain(downstream.iter()) - .filter_map(|gene| gff.cluster_for_gene(gene, gene_to_cluster)) + .filter_map(|gene| gene_to_cluster.cluster_for_gene(gene)) .collect(); cluster_ids.sort_unstable(); cluster_ids.dedup(); @@ -2288,101 +2280,157 @@ fn best_reference_context( index.contexts[best_idx].genes.clone() } -pub(super) fn write_cluster_locus_contexts( +/// Runs of consecutive clusters to build in parallel and write in order. A run +/// is cut by the genomes its clusters are present in, not by a count of +/// clusters, so what one holds does not grow with the number of genomes: a +/// cluster for every thread at least, then clusters until the run's +/// genome-cluster pairs reach the budget. +fn cluster_chunks( + pa: &crate::parser::PresenceAbsence, + threads: usize, +) -> Vec> { + const PAIRS_PER_CHUNK: usize = 32 * 1024; + let mut chunks = Vec::new(); + let mut chunk_start = 0usize; + while chunk_start < pa.n_clusters() { + let mut chunk_end = chunk_start; + let mut pairs = 0usize; + while chunk_end < pa.n_clusters() + && (chunk_end - chunk_start < threads || pairs < PAIRS_PER_CHUNK) + { + pairs += pa.present_count(ClusterIdx(chunk_end)); + chunk_end += 1; + } + chunks.push(chunk_start..chunk_end); + chunk_start = chunk_end; + } + chunks +} + +/// One cluster's locus contexts: its summary, and its context rows ranked by +/// carriers. +fn cluster_locus_contexts( + ci: ClusterIdx, pa: &crate::parser::PresenceAbsence, gff_indices: &HashMap, - gene_to_cluster_maps: &HashMap>, - output_dir: &std::path::Path, -) -> Result<(usize, Vec)> { - let mut context_rows: BTreeMap<(usize, String, String), ContextAccumulator> = BTreeMap::new(); - let mut summaries = Vec::new(); + gene_to_cluster_maps: &HashMap, +) -> ( + ClusterContextSummary, + Vec<(String, String, ContextAccumulator)>, +) { + let mut context_rows: BTreeMap<(String, String), ContextAccumulator> = BTreeMap::new(); + let mut n_mapped = 0usize; + let mut n_unresolved = 0usize; - for ci_val in 0..pa.n_clusters() { - let ci = ClusterIdx(ci_val); - let mut n_mapped = 0usize; - let mut n_unresolved = 0usize; - let mut context_counts: HashMap<(String, String), usize> = HashMap::new(); + let mut genomes = pa.present_genomes(ci).to_vec(); + genomes.sort_by_key(|gi| gi.val()); - let mut genomes = pa.present_genomes(ci).to_vec(); - genomes.sort_by_key(|gi| gi.val()); + for gi in genomes { + let Some(gff) = gff_indices.get(&gi) else { + continue; + }; + let Some(g2c) = gene_to_cluster_maps.get(&gi) else { + continue; + }; + let cell = pa.cell(ci, gi); + let Some(gene) = gff_gene_for_cell(gff, cell) else { + n_unresolved += 1; + continue; + }; - for gi in genomes { - let Some(gff) = gff_indices.get(&gi) else { - continue; - }; - let Some(g2c) = gene_to_cluster_maps.get(&gi) else { - continue; - }; - let cell = pa.cell(ci, gi); - let Some(gene) = gff_gene_for_cell(gff, cell) else { - n_unresolved += 1; - continue; - }; + let (upstream, downstream) = + gff.context_window(gene.contig, gene.start, gene.end, Some(gene.gene_id), 2); + let upstream_sig = context_signature(&upstream, g2c, pa); + let downstream_sig = context_signature(&downstream, g2c, pa); + n_mapped += 1; - let (upstream, downstream) = gff.context_window( - &gene.contig, - gene.start, - gene.end, - Some(gene.gene_id.as_str()), - 2, - ); - let upstream_sig = context_signature(gff, &upstream, g2c, pa); - let downstream_sig = context_signature(gff, &downstream, g2c, pa); - let key = (upstream_sig.clone(), downstream_sig.clone()); - *context_counts.entry(key).or_insert(0) += 1; - n_mapped += 1; - - let row_key = (ci.val(), upstream_sig, downstream_sig); - let row = context_rows.entry(row_key).or_default(); - row.count += 1; - if row.examples.len() < 5 { - row.examples - .push(format!("{}:{}", pa.genome_name(gi), gene.gene_id)); - } + let row = context_rows + .entry((upstream_sig, downstream_sig)) + .or_default(); + row.count += 1; + if row.examples.len() < 5 { + row.examples + .push(format!("{}:{}", pa.genome_name(gi), gene.gene_id)); } - - let n_unique_contexts = context_counts.len(); - let dominant_count = context_counts.values().copied().max().unwrap_or(0); - let dominant_fraction = if n_mapped > 0 { - dominant_count as f64 / n_mapped as f64 - } else { - 0.0 - }; - // Judged by neighbourhood, not by exact signature. A context that shares - // even one neighbouring cluster with the dominant one is the same locus - // with something gained, lost or unannotated next to it: a pseudogene, - // a missing CDS or a true loss beside a single-copy gene changes its - // signature without making it a second copy. Only a context with no - // neighbour in common is evidence of another locus. Signature - // differences alone used to label single-copy genes as paralogs. - let dominant = context_counts - .iter() - .max_by(|a, b| a.1.cmp(b.1).then_with(|| b.0.cmp(a.0))) - .map(|(key, _)| key.clone()); - let compatible: usize = context_counts - .iter() - .filter(|(key, _)| dominant.as_ref().is_none_or(|d| shares_neighbour(d, key))) - .map(|(_, count)| *count) - .sum(); - let compatible_fraction = if n_mapped > 0 { - compatible as f64 / n_mapped as f64 - } else { - 0.0 - }; - let interpretation = - context_interpretation(n_mapped, n_unique_contexts, compatible_fraction); - summaries.push(ClusterContextSummary { - cluster_idx: ci, - n_present: pa.present_count(ci), - n_mapped, - n_unresolved, - n_unique_contexts, - dominant_count, - dominant_fraction, - interpretation, - }); } + let n_unique_contexts = context_rows.len(); + let dominant_count = context_rows + .values() + .map(|row| row.count) + .max() + .unwrap_or(0); + let dominant_fraction = if n_mapped > 0 { + dominant_count as f64 / n_mapped as f64 + } else { + 0.0 + }; + // Judged by neighbourhood, not by exact signature. A context that shares + // even one neighbouring cluster with the dominant one is the same locus + // with something gained, lost or unannotated next to it: a pseudogene, + // a missing CDS or a true loss beside a single-copy gene changes its + // signature without making it a second copy. Only a context with no + // neighbour in common is evidence of another locus. Signature + // differences alone used to label single-copy genes as paralogs. + let dominant = context_rows + .iter() + .max_by(|a, b| a.1.count.cmp(&b.1.count).then_with(|| b.0.cmp(a.0))) + .map(|(key, _)| key.clone()); + let compatible: usize = context_rows + .iter() + .filter(|(key, _)| dominant.as_ref().is_none_or(|d| shares_neighbour(d, key))) + .map(|(_, row)| row.count) + .sum(); + let compatible_fraction = if n_mapped > 0 { + compatible as f64 / n_mapped as f64 + } else { + 0.0 + }; + let interpretation = context_interpretation(n_mapped, n_unique_contexts, compatible_fraction); + let summary = ClusterContextSummary { + cluster_idx: ci, + n_present: pa.present_count(ci), + n_mapped, + n_unresolved, + n_unique_contexts, + dominant_count, + dominant_fraction, + interpretation, + }; + + let mut rows: Vec<(String, String, ContextAccumulator)> = context_rows + .into_iter() + .map(|((upstream, downstream), acc)| (upstream, downstream, acc)) + .collect(); + rows.sort_by(|a, b| { + b.2.count + .cmp(&a.2.count) + .then(a.0.cmp(&b.0)) + .then(a.1.cmp(&b.1)) + }); + (summary, rows) +} + +/// Each cluster's neighbourhoods, in `cluster_locus_context.tsv`, and a summary +/// per cluster, in `cluster_locus_summary.tsv`. +/// +/// One cluster at a time in one thread, this was a sixth of a run on 400 +/// genomes and grew with every genome. Clusters do not depend on each other, +/// so they are built in parallel a run at a time and written in cluster order, +/// the same rows in the same order. +pub(super) fn write_cluster_locus_contexts( + pa: &crate::parser::PresenceAbsence, + gff_indices: &HashMap, + gene_to_cluster_maps: &HashMap, + threads: usize, + output_dir: &std::path::Path, +) -> Result<(usize, Vec)> { + let threads = threads.max(1); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .context("Failed to build the locus context thread pool")?; + let context_path = output_dir.join("cluster_locus_context.tsv"); let context_file = fs::File::create(&context_path)?; let mut context_w = BufWriter::new(context_file); @@ -2390,37 +2438,41 @@ pub(super) fn write_cluster_locus_contexts( context_w, "gene_cluster\tcontext_rank\tn_genomes\tfraction_of_mapped_contexts\tupstream_context\tdownstream_context\texample_genomes" )?; + let mut summaries = Vec::with_capacity(pa.n_clusters()); let mut written_context_rows = 0usize; - let mut rows_by_cluster: BTreeMap> = - BTreeMap::new(); - for ((ci, upstream, downstream), acc) in context_rows { - rows_by_cluster - .entry(ci) - .or_default() - .push((upstream, downstream, acc)); - } - for (ci_val, mut rows) in rows_by_cluster { - rows.sort_by(|a, b| { - b.2.count - .cmp(&a.2.count) - .then(a.0.cmp(&b.0)) - .then(a.1.cmp(&b.1)) + for chunk in cluster_chunks(pa, threads) { + let built: Vec<_> = pool.install(|| { + use rayon::prelude::*; + chunk + .into_par_iter() + .map(|ci_val| { + cluster_locus_contexts( + ClusterIdx(ci_val), + pa, + gff_indices, + gene_to_cluster_maps, + ) + }) + .collect() }); - let mapped = summaries[ci_val].n_mapped.max(1); - for (rank, (upstream, downstream, acc)) in rows.into_iter().enumerate() { - write_tsv_fields( - &mut context_w, - &[ - pa.cluster_name(ClusterIdx(ci_val)).to_string(), - (rank + 1).to_string(), - acc.count.to_string(), - format!("{:.3}", acc.count as f64 / mapped as f64), - upstream, - downstream, - acc.examples.join(";"), - ], - )?; - written_context_rows += 1; + for (summary, rows) in built { + let mapped = summary.n_mapped.max(1); + for (rank, (upstream, downstream, acc)) in rows.into_iter().enumerate() { + write_tsv_fields( + &mut context_w, + &[ + pa.cluster_name(summary.cluster_idx).to_string(), + (rank + 1).to_string(), + acc.count.to_string(), + format!("{:.3}", acc.count as f64 / mapped as f64), + upstream, + downstream, + acc.examples.join(";"), + ], + )?; + written_context_rows += 1; + } + summaries.push(summary); } } context_w.flush()?; @@ -2464,11 +2516,16 @@ pub(super) fn write_cluster_locus_contexts( /// in one thread: on 200 genomes of 4,000 genes that held 1.2 GB and took half /// the run. Clusters do not depend on each other, so they are now built in /// parallel a chunk at a time and written in cluster order, the same rows in -/// the same order, holding one chunk's blocks at a time. +/// the same order, holding one chunk's rows at a time. +/// +/// A block is a genome's neighbourhood, so where annotations share no gene IDs +/// every genome is a block of its own: chunks of 256 clusters held 200 MB of +/// rows on 400 genomes, more with every genome added. Chunks are now cut by +/// genome-cluster pairs; see cluster_chunks. pub(super) fn write_cluster_synteny_blocks( pa: &crate::parser::PresenceAbsence, gff_indices: &HashMap, - gene_to_cluster_maps: &HashMap>, + gene_to_cluster_maps: &HashMap, reference_gi: Option, threads: usize, output_dir: &std::path::Path, @@ -2487,19 +2544,18 @@ pub(super) fn write_cluster_synteny_blocks( "gene_cluster\tblock_rank\tn_genomes\tfraction_of_present\tupstream_block_5\tdownstream_block_5\tblock_interpretation\texample_genomes\tcarrier_genomes\trepresentative_genome\trepresentative_gene\trepresentative_contig\trepresentative_start\trepresentative_end\trepresentative_strand\trepresentative_length_bp\trepresentative_block_genes\treference_block_genes" )?; - const CLUSTERS_PER_CHUNK: usize = 256; + let threads = threads.max(1); let pool = rayon::ThreadPoolBuilder::new() - .num_threads(threads.max(1)) + .num_threads(threads) .build() .context("Failed to build the synteny block thread pool")?; - let clusters: Vec = (0..pa.n_clusters()).collect(); let mut written = 0usize; - for chunk in clusters.chunks(CLUSTERS_PER_CHUNK) { - let rows: Vec>> = pool.install(|| { + for chunk in cluster_chunks(pa, threads) { + let rows: Vec, usize)>> = pool.install(|| { use rayon::prelude::*; chunk - .par_iter() - .map(|&ci_val| { + .into_par_iter() + .map(|ci_val| { cluster_synteny_block_rows( ClusterIdx(ci_val), pa, @@ -2512,25 +2568,25 @@ pub(super) fn write_cluster_synteny_blocks( .collect() }); for cluster_rows in rows { - for fields in cluster_rows { - write_tsv_fields(&mut w, &fields)?; - written += 1; - } + let (lines, n_rows) = cluster_rows?; + w.write_all(&lines)?; + written += n_rows; } } w.flush()?; Ok(written) } -/// One cluster's synteny block rows, ranked by carriers. +/// One cluster's synteny block rows, ranked by carriers, as TSV lines, and how +/// many there are. fn cluster_synteny_block_rows( ci: ClusterIdx, pa: &crate::parser::PresenceAbsence, gff_indices: &HashMap, - gene_to_cluster_maps: &HashMap>, + gene_to_cluster_maps: &HashMap, reference_boundary_index: &ReferenceBoundaryIndex, cluster_index: &HashMap, -) -> Vec> { +) -> Result<(Vec, usize)> { let mut blocks: BTreeMap = BTreeMap::new(); let mut genomes = pa.present_genomes(ci).to_vec(); genomes.sort_by_key(|gi| gi.val()); @@ -2545,17 +2601,12 @@ fn cluster_synteny_block_rows( continue; }; - let (upstream, downstream) = gff.context_window( - &gene.contig, - gene.start, - gene.end, - Some(gene.gene_id.as_str()), - 5, - ); - let upstream_sig = context_signature(gff, &upstream, g2c, pa); - let downstream_sig = context_signature(gff, &downstream, g2c, pa); + let (upstream, downstream) = + gff.context_window(gene.contig, gene.start, gene.end, Some(gene.gene_id), 5); + let upstream_sig = context_signature(&upstream, g2c, pa); + let downstream_sig = context_signature(&downstream, g2c, pa); let representative_genes = - ordered_synteny_gene_records(gff, &upstream, Some(gene), &downstream, g2c, pa); + ordered_synteny_gene_records(gff, &upstream, Some(&gene), &downstream, g2c, pa); let block_signature = synteny_gene_identity_signature(&representative_genes); let row = blocks.entry(block_signature).or_default(); row.count += 1; @@ -2571,8 +2622,8 @@ fn cluster_synteny_block_rows( if row.representative.is_none() { row.representative = Some(SyntenyBlockRepresentative { genome: pa.genome_name(gi).to_string(), - gene_id: gene.gene_id.clone(), - contig: gene.contig.clone(), + gene_id: gene.gene_id.to_string(), + contig: gene.contig.to_string(), start: gene.start + 1, end: gene.end, strand: gene.strand, @@ -2586,7 +2637,9 @@ fn cluster_synteny_block_rows( rows.sort_by(|a, b| b.1.count.cmp(&a.1.count).then(a.0.cmp(&b.0))); let present = pa.present_count(ci).max(1); let n_blocks = rows.len(); - rows.into_iter() + // Each row is formatted and written as it is made. + let row_fields = rows + .into_iter() .enumerate() .map(|(rank, (_signature, acc))| { let fraction = acc.count as f64 / present as f64; @@ -2641,6 +2694,10 @@ fn cluster_synteny_block_rows( .unwrap_or_default(), encode_synteny_gene_records(&reference_genes), ] - }) - .collect() + }); + let mut lines = Vec::new(); + for fields in row_fields { + write_tsv_fields(&mut lines, &fields)?; + } + Ok((lines, n_blocks)) } diff --git a/src/pipeline/output.rs b/src/pipeline/output.rs index cd8d9c8..3697045 100644 --- a/src/pipeline/output.rs +++ b/src/pipeline/output.rs @@ -275,7 +275,7 @@ pub(super) struct MultistateMatrixInputs<'a> { pub(super) nsc_rescue_map: &'a HashMap<(ClusterIdx, GenomeIdx), &'a RescueEvent>, pub(super) synteny_rejected: &'a [RescueEvent], pub(super) paralog_rejected: &'a [RescueEvent], - pub(super) gene_to_cluster_maps: &'a HashMap>, + pub(super) gene_to_cluster_maps: &'a HashMap>, pub(super) output_dir: &'a std::path::Path, pub(super) nsc_result: &'a Option, pub(super) coverage: &'a SearchCoverage, diff --git a/src/pipeline/reference_audit.rs b/src/pipeline/reference_audit.rs index ea7a366..b8cbbc7 100644 --- a/src/pipeline/reference_audit.rs +++ b/src/pipeline/reference_audit.rs @@ -223,17 +223,17 @@ pub(super) fn best_overlap_gene_id(gff: &GffIndex, hit: &Hit) -> Option return None; } let reciprocal = (overlap / hit_len).min(overlap / gene_len); - (reciprocal > 0.5).then(|| (gene.gene_id.clone(), reciprocal)) + (reciprocal > 0.5).then_some((gene.gene_id, reciprocal)) }) .max_by(|a, b| a.1.total_cmp(&b.1)) - .map(|(gene_id, _)| gene_id) + .map(|(gene_id, _)| gene_id.to_string()) } pub(super) fn reference_hit_shadow_candidates( hits: &HashMap<(GenomeIdx, ClusterIdx), Vec>, pa: &crate::parser::PresenceAbsence, gff_indices: &HashMap, - gene_to_cluster_maps: &HashMap>, + gene_to_cluster_maps: &HashMap, ) -> Vec { let mut candidates = Vec::new(); for (&(gi, ci), cluster_hits) in hits { @@ -247,7 +247,7 @@ pub(super) fn reference_hit_shadow_candidates( let annotation_cluster = annotation_match .as_deref() .and_then(|gene_id| g2c.and_then(|map| map.get(gene_id))) - .map(|other_ci| pa.cluster_name(*other_ci).to_string()); + .map(|other_ci| pa.cluster_name(other_ci).to_string()); let query_cluster = pa.cluster_name(ci); let rescue_type = if let Some(other_cluster) = annotation_cluster { if other_cluster != query_cluster { @@ -699,7 +699,7 @@ pub(super) fn reference_event_score(event: &ReferenceClusterEvent) -> f64 { pub(super) struct ReferenceEventRow<'a> { reference_name: &'a str, rescue: &'a RescueEvent, - gene_to_cluster_maps: &'a HashMap>, + gene_to_cluster_maps: &'a HashMap>, cluster_prevalences: &'a [f64], cluster_seq_lengths: &'a [usize], matrix_policy: MatrixPolicy, @@ -716,7 +716,7 @@ pub(super) struct ReferenceAuditInputs<'a> { pub(super) reference_candidates: &'a [RescueEvent], pub(super) reference_cluster_events: &'a [ReferenceClusterEvent], pub(super) pa: &'a crate::parser::PresenceAbsence, - pub(super) gene_to_cluster_maps: &'a HashMap>, + pub(super) gene_to_cluster_maps: &'a HashMap>, pub(super) cluster_prevalences: &'a [f64], pub(super) cluster_seq_lengths: &'a [usize], pub(super) matrix_policy: MatrixPolicy, diff --git a/src/pipeline/scale_harness.rs b/src/pipeline/scale_harness.rs index d55a576..bac4ede 100644 --- a/src/pipeline/scale_harness.rs +++ b/src/pipeline/scale_harness.rs @@ -992,7 +992,7 @@ fn scale_report() { }; // Evidence rows the way the pipeline builds them, so the per-cell // evidence panel and the kept-out cells are exercised at this size too. - let no_annotations: HashMap> = HashMap::new(); + let no_annotations: HashMap = HashMap::new(); let prevalences: Vec = (0..pa.n_clusters()) .map(|ci| pa.prevalence(ClusterIdx(ci))) .collect(); diff --git a/src/pipeline/tests.rs b/src/pipeline/tests.rs index 5e63d1e..cfdb63a 100644 --- a/src/pipeline/tests.rs +++ b/src/pipeline/tests.rs @@ -1737,7 +1737,7 @@ fn reference_hit_shadow_candidates_capture_partial_annotated_paralog() { vec!["L4_AS00146".into()], vec![vec!["".into()], vec!["rv_duf222".into()]], ); - let gff = GffIndex::build(vec![GffGene { + let gff = GffIndex::build(vec![GffRecord { contig: "chr".into(), start: 0, end: 954, @@ -1749,7 +1749,10 @@ fn reference_hit_shadow_candidates_capture_partial_annotated_paralog() { let mut gene_map = HashMap::new(); gene_map.insert("rv_duf222".to_string(), ClusterIdx(1)); let mut gene_to_cluster_maps = HashMap::new(); - gene_to_cluster_maps.insert(GenomeIdx(0), gene_map); + gene_to_cluster_maps.insert( + GenomeIdx(0), + GeneClusters::from_map(&gff_indices[&GenomeIdx(0)], &gene_map), + ); let mut hits = HashMap::new(); hits.insert( (GenomeIdx(0), ClusterIdx(0)), @@ -1932,7 +1935,7 @@ fn mobile_element_detection_uses_gff_aliases() { vec![vec!["".into()]], ); let gff = GffIndex::build_with_aliases(vec![( - GffGene { + GffRecord { contig: "chr".into(), start: 150, end: 300, @@ -1982,7 +1985,7 @@ fn mobile_element_sequence_hits_are_name_agnostic_and_prioritized() { gff_indices.insert( GenomeIdx(0), GffIndex::build_with_aliases(vec![( - GffGene { + GffRecord { contig: "chr".into(), start: 90, end: 210, diff --git a/src/rescue.rs b/src/rescue.rs index 58d8b9b..37477ca 100644 --- a/src/rescue.rs +++ b/src/rescue.rs @@ -11,9 +11,9 @@ mod synteny; mod types; pub use self::conflicts::resolve_conflicts; -pub use self::genes::build_gene_to_cluster_with_gff_aliases; +pub use self::genes::build_gene_clusters; pub(crate) use self::genes::{cell_gene_id_variants, panaroo_gene_id_suffix_variants}; -pub use self::gff::{GffGene, GffIndex}; +pub use self::gff::{GeneClusters, GffGene, GffIndex, GffRecord}; pub(crate) use self::orf::reverse_complement; pub use self::orf::{ check_internal_stops_in_cds, check_internal_stops_in_hit, check_internal_stops_in_span, @@ -25,7 +25,7 @@ pub use self::types::{ }; #[cfg(test)] -use self::genes::build_gene_to_cluster; +use self::genes::{build_gene_to_cluster, build_gene_to_cluster_with_gff_aliases}; #[cfg(test)] use self::orf::count_internal_stops; #[cfg(test)] @@ -497,7 +497,7 @@ fn best_within_contig_chain(hits: &[&Hit]) -> Option { fn compatible_single_cds<'a>( chain: &FragmentChain, gff_index: Option<&'a GffIndex>, -) -> Option<&'a GffGene> { +) -> Option> { let gff = gff_index?; let overlapping = gff.overlapping(&chain.contig, chain.start, chain.end); let span_len = chain.end.saturating_sub(chain.start) as f64; @@ -539,10 +539,10 @@ fn classify_matched_gene( coverage: f64, ci: ClusterIdx, pa: &PresenceAbsence, - gene_to_cluster: &HashMap, + gene_to_cluster: &GeneClusters, ) -> (RescueType, f64) { let base_confidence = identity / 100.0 * coverage / 100.0; - if let Some(&other_ci) = gene_to_cluster.get(matched_gene_id) { + if let Some(other_ci) = gene_to_cluster.get(matched_gene_id) { if other_ci != ci { if identity >= 95.0 { return ( @@ -636,7 +636,7 @@ fn detect_start_variant_from_span( } let overlap = end.min(gene.end).saturating_sub(start.max(gene.start)) as f64; if hit_len > 0.0 && overlap / hit_len > 0.8 { - return Some((start_diff, gene.gene_id.clone())); + return Some((start_diff, gene.gene_id.to_string())); } } @@ -667,7 +667,7 @@ pub fn classify_hits( gi: GenomeIdx, pa: &PresenceAbsence, gff_index: Option<&GffIndex>, - gene_to_cluster: &HashMap, + gene_to_cluster: &GeneClusters, min_coverage: f64, ) -> Option { if hits.is_empty() { @@ -863,14 +863,14 @@ pub fn classify_hits( chain.evidence.note = Some("same_cds_alignment_split".to_string()); let (rescue_type, annotation_match, confidence) = { let (rescue_type, confidence) = classify_matched_gene( - &gene.gene_id, + gene.gene_id, chain.avg_identity, chain.coverage, ci, pa, gene_to_cluster, ); - (rescue_type, Some(gene.gene_id.clone()), confidence) + (rescue_type, Some(gene.gene_id.to_string()), confidence) }; let (final_type, final_ann, final_conf) = if matches!( @@ -1081,7 +1081,7 @@ fn classify_single_hit( ci: ClusterIdx, pa: &PresenceAbsence, gff_index: Option<&GffIndex>, - gene_to_cluster: &HashMap, + gene_to_cluster: &GeneClusters, ) -> (RescueType, Option, f64) { let base_confidence = hit.identity() / 100.0 * hit.query_coverage() / 100.0; @@ -1118,12 +1118,12 @@ fn classify_single_hit( let Some((matched_gene, _recip)) = best_gene else { return ( RescueType::Dropped, - Some(overlapping[0].gene_id.clone()), + Some(overlapping[0].gene_id.to_string()), base_confidence, ); }; - if let Some(&other_ci) = gene_to_cluster.get(&matched_gene.gene_id) { + if let Some(other_ci) = gene_to_cluster.get(matched_gene.gene_id) { if other_ci != ci { // Gene is annotated but assigned to a different Panaroo cluster. // High identity (≥95%) → SPLIT (same gene, mis-clustered by Panaroo). @@ -1133,7 +1133,7 @@ fn classify_single_hit( RescueType::Split { other_cluster: pa.cluster_name(other_ci).to_string(), }, - Some(matched_gene.gene_id.clone()), + Some(matched_gene.gene_id.to_string()), base_confidence * 1.1, ); } else { @@ -1141,7 +1141,7 @@ fn classify_single_hit( RescueType::Paralog { other_cluster: pa.cluster_name(other_ci).to_string(), }, - Some(matched_gene.gene_id.clone()), + Some(matched_gene.gene_id.to_string()), base_confidence * 0.7, // lower confidence: paralog, not real gene ); } @@ -1150,7 +1150,7 @@ fn classify_single_hit( ( RescueType::Dropped, - Some(matched_gene.gene_id.clone()), + Some(matched_gene.gene_id.to_string()), base_confidence, ) } diff --git a/src/rescue/genes.rs b/src/rescue/genes.rs index 64cf8c7..7455c6e 100644 --- a/src/rescue/genes.rs +++ b/src/rescue/genes.rs @@ -1,4 +1,4 @@ -use super::gff::GffIndex; +use super::gff::{GeneClusters, GffIndex}; use crate::parser::PresenceAbsence; use crate::types::{ClusterIdx, GenomeIdx}; use ahash::{HashMap, HashMapExt}; @@ -98,6 +98,19 @@ fn insert_gene_cluster_alias(map: &mut HashMap, gene_id: &st } } +/// The cluster of every name of one genome's annotation, read off the map +/// below, which is built for the moment and dropped; see GeneClusters. +pub fn build_gene_clusters<'a>( + pa: &PresenceAbsence, + gi: GenomeIdx, + gff: &'a GffIndex, +) -> GeneClusters<'a> { + GeneClusters::from_map( + gff, + &build_gene_to_cluster_with_gff_aliases(pa, gi, Some(gff)), + ) +} + /// Build a gene_id to ClusterIdx lookup and conservatively add GFF aliases. /// /// Panaroo cells can contain a gene name while the GFF uses a locus_tag (or the @@ -114,8 +127,8 @@ pub fn build_gene_to_cluster_with_gff_aliases( return map; }; - for gene in &gff.genes { - let labels = gff.aliases_for_gene(gene); + for gene in gff.genes() { + let labels = gff.aliases_for_gene(&gene); let mut candidates: Vec = labels .iter() .filter_map(|label| lookup_gene_cluster_alias(&map, label)) diff --git a/src/rescue/gff.rs b/src/rescue/gff.rs index 0b237db..6a9f33a 100644 --- a/src/rescue/gff.rs +++ b/src/rescue/gff.rs @@ -1,170 +1,411 @@ use crate::types::ClusterIdx; use ahash::{HashMap, HashMapExt}; +/// A CDS as read from a GFF: the owned record an index is built from. #[derive(Debug, Clone)] -pub struct GffGene { +pub struct GffRecord { pub contig: String, pub start: u64, pub end: u64, - #[allow(dead_code)] pub strand: char, pub gene_id: String, } +/// One CDS of an index: a view into the index's own storage, cheap to copy. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct GffGene<'a> { + pub contig: &'a str, + pub start: u64, + pub end: u64, + pub strand: char, + pub gene_id: &'a str, + /// The gene's position in its index. + pub(crate) index: u32, +} + +/// One CDS, 32 bytes. Its names (the gene ID, then its aliases) are the +/// `n_names` names from position `first_name` on. +#[derive(Debug, Clone, Copy)] +struct Record { + start: u64, + end: u64, + first_name: u32, + /// The record whose aliases answer for this one: the last record with the + /// same gene ID, which is whose aliases a map keyed by gene ID kept. + alias_owner: u32, + contig: u32, + n_names: u16, + strand: u8, +} + +/// The CDS annotation of one genome, sorted by contig and start. +/// +/// Every name of every CDS (its gene ID, then its aliases) is stored once, +/// back to back in one buffer, and found through one sorted table of hashes. +/// The index used to hold each name as its own String in up to six places +/// (the gene, an ID map, an alias map and a label map, keys and values), plus +/// every product and note: on PGAP-style annotations that was about 4.7 MB per +/// genome of 4,000 genes, which is what made memory grow with every genome a +/// run added. Products and notes are only read to recognise mobile elements, +/// so that is worked out once, here, and only the few mobile genes keep their +/// labels. pub struct GffIndex { - pub genes: Vec, - contig_ranges: HashMap, - /// Largest gene end seen so far within each contig, aligned with `genes`. - /// Genes are sorted by start, and their ends are not monotone once one CDS - /// contains another, so a binary search on `end` itself is wrong: it - /// skipped a CDS that contained a nested one, and the annotated locus read - /// as having no CDS. The running maximum is monotone. + text: String, + /// Where each name ends in `text`; each starts where the one before ends. + name_ends: Vec, + records: Vec, + contig_names: Vec>, + contig_ids: HashMap, u32>, + /// Records of each contig, as a range, by contig id. + contig_ranges: Vec<(u32, u32)>, + /// Largest gene end seen so far within each contig, aligned with the + /// records. Genes are sorted by start, and their ends are not monotone once + /// one CDS contains another, so a binary search on `end` itself is wrong: + /// it skipped a CDS that contained a nested one, and the annotated locus + /// read as having no CDS. The running maximum is monotone. prefix_max_end: Vec, - /// Longest gene on each contig, which bounds how far back an upstream - /// search has to look. - max_gene_len: HashMap, - gene_id_index: HashMap, - gene_aliases: HashMap>, - gene_labels: HashMap>, + /// Longest gene on each contig, by contig id, which bounds how far back an + /// upstream search has to look. + max_gene_len: Vec, + /// (hash of a name, position of the name that answers for it), one entry + /// per distinct name, sorted. + name_index: Vec<(u32, u32)>, + /// Mobile-element genes by alias owner: the kind and the joined labels. + mobile: Vec<(u32, &'static str, Box)>, + hasher: ahash::RandomState, +} + +/// What a gene's names and descriptions say about it being a mobile element: +/// the kind, and the labels joined as they were read. +fn mobile_kind_from_labels(labels: &[&str]) -> Option<(&'static str, String)> { + let joined = labels.join("|"); + let normalized = joined + .to_ascii_lowercase() + .replace("%20", " ") + .replace('_', " "); + let kind = if normalized.contains("is6110") { + "IS6110" + } else if normalized.contains("is1081") { + "IS1081" + } else if normalized.contains("insertion sequence") + || normalized.contains("transposase") + || normalized.contains("is element") + { + "insertion_sequence" + } else if normalized.contains("integrase") || normalized.contains("phage") { + "phage_or_integrase" + } else if normalized.contains("recombinase") { + "recombinase" + } else { + return None; + }; + Some((kind, joined)) +} + +fn name_hasher() -> ahash::RandomState { + ahash::RandomState::with_seeds( + 0x243f_6a88_85a3_08d3, + 0x1319_8a2e_0370_7344, + 0xa409_3822_299f_31d0, + 0x082e_fa98_ec4e_6c89, + ) +} + +/// The name at `pos` of a name buffer. +fn name_at<'t>(text: &'t str, name_ends: &[u32], pos: u32) -> &'t str { + let start = match pos { + 0 => 0, + _ => name_ends[pos as usize - 1] as usize, + }; + &text[start..name_ends[pos as usize] as usize] } impl GffIndex { #[allow(dead_code)] - pub fn build(mut genes: Vec) -> Self { + pub fn build(records: Vec) -> Self { Self::build_with_aliases( - genes - .drain(..) - .map(|gene| (gene, Vec::new())) - .collect::>(), + records + .into_iter() + .map(|record| (record, Vec::new())) + .collect(), ) } - pub fn build_with_aliases(mut records: Vec<(GffGene, Vec)>) -> Self { - let records = records - .drain(..) - .map(|(gene, aliases)| (gene, aliases.clone(), aliases)) - .collect(); - Self::build_with_aliases_and_labels(records) + pub fn build_with_aliases(records: Vec<(GffRecord, Vec)>) -> Self { + Self::build_with_aliases_and_labels( + records + .into_iter() + .map(|(record, aliases)| { + let labels = aliases.clone(); + (record, aliases, labels) + }) + .collect(), + ) } pub fn build_with_aliases_and_labels( - mut records: Vec<(GffGene, Vec, Vec)>, + mut records: Vec<(GffRecord, Vec, Vec)>, ) -> Self { + // Stable, by contig name then start: the order every lookup and every + // output that walks a contig relies on. records.sort_by(|a, b| a.0.contig.cmp(&b.0.contig).then(a.0.start.cmp(&b.0.start))); - let genes: Vec = records.iter().map(|(gene, _, _)| gene.clone()).collect(); - - let mut contig_ranges = HashMap::new(); - if !genes.is_empty() { - let mut current_contig = &genes[0].contig; - let mut range_start = 0; - for (i, g) in genes.iter().enumerate() { - if &g.contig != current_contig { - contig_ranges.insert(current_contig.clone(), (range_start, i)); - current_contig = &g.contig; - range_start = i; - } - } - contig_ranges.insert(current_contig.clone(), (range_start, genes.len())); - } + let mut text = String::new(); + let mut name_ends: Vec = Vec::new(); + let mut contig_names: Vec> = Vec::new(); + let mut contig_ranges: Vec<(u32, u32)> = Vec::new(); + let mut max_gene_len: Vec = Vec::new(); + let mut out: Vec = Vec::with_capacity(records.len()); + let mut prefix_max_end: Vec = Vec::with_capacity(records.len()); + // Labels of every record whose gene ID is not empty, until the mobile + // genes have been picked out. + let mut labels_by_record: Vec>> = Vec::with_capacity(records.len()); - let mut prefix_max_end: Vec = Vec::with_capacity(genes.len()); - let mut max_gene_len: HashMap = HashMap::new(); - for (i, g) in genes.iter().enumerate() { - let new_contig = i == 0 || genes[i - 1].contig != g.contig; + for (i, (gene, aliases, labels)) in records.into_iter().enumerate() { + // The genes of one contig are contiguous after the sort, so a + // contig is new exactly when it is not the one just seen. + let contig = match contig_names.last() { + Some(last) if **last == *gene.contig => contig_names.len() as u32 - 1, + _ => { + contig_names.push(gene.contig.into_boxed_str()); + contig_ranges.push((i as u32, i as u32)); + max_gene_len.push(0); + contig_names.len() as u32 - 1 + } + }; + contig_ranges[contig as usize].1 = i as u32 + 1; + let new_contig = contig_ranges[contig as usize].0 == i as u32; let running = if new_contig { - g.end + gene.end } else { - prefix_max_end[i - 1].max(g.end) + prefix_max_end[i - 1].max(gene.end) }; prefix_max_end.push(running); - let len = g.end.saturating_sub(g.start); - let longest = max_gene_len.entry(g.contig.clone()).or_insert(0); - *longest = (*longest).max(len); - } + let len = gene.end.saturating_sub(gene.start); + max_gene_len[contig as usize] = max_gene_len[contig as usize].max(len); - let mut gene_id_index = HashMap::new(); - let mut gene_aliases = HashMap::new(); - let mut gene_labels = HashMap::new(); - for (i, (g, aliases, labels)) in records.iter().enumerate() { - if !g.gene_id.is_empty() { - gene_id_index.insert(g.gene_id.clone(), i); - let mut all_aliases = vec![g.gene_id.clone()]; - for alias in aliases { - if !alias.is_empty() && !all_aliases.iter().any(|x| x == alias) { - gene_id_index.entry(alias.clone()).or_insert(i); - all_aliases.push(alias.clone()); + // Names: the gene ID, then each alias once, as the alias map held + // them. A record with no gene ID has none. + let first_name = name_ends.len() as u32; + let mut n_names = 0u16; + let mut stored_labels = None; + if !gene.gene_id.is_empty() { + let mut all_aliases: Vec<&str> = vec![gene.gene_id.as_str()]; + for alias in &aliases { + if !alias.is_empty() && !all_aliases.contains(&alias.as_str()) { + all_aliases.push(alias); } } - let mut all_labels = all_aliases.clone(); + let mut all_labels: Vec = + all_aliases.iter().map(|s| s.to_string()).collect(); for label in labels { - if !label.is_empty() && !all_labels.iter().any(|x| x == label) { - all_labels.push(label.clone()); + if !label.is_empty() && !all_labels.contains(&label) { + all_labels.push(label); } } - gene_aliases.insert(g.gene_id.clone(), all_aliases); - gene_labels.insert(g.gene_id.clone(), all_labels); + stored_labels = Some(all_labels); + all_aliases.truncate(u16::MAX as usize); + for name in &all_aliases { + text.push_str(name); + name_ends.push(u32::try_from(text.len()).expect("annotation names under 4 GB")); + } + n_names = all_aliases.len() as u16; } + labels_by_record.push(stored_labels); + out.push(Record { + start: gene.start, + end: gene.end, + first_name, + alias_owner: i as u32, + contig, + n_names, + strand: u8::try_from(gene.strand).unwrap_or(b'+'), + }); } + // Which name answers for each spelling, exactly as the maps did: a gene + // ID claims its name (a later record with the same ID takes it over), + // an alias only takes a name nobody has claimed yet. And whose aliases + // and labels a record answers with: the last record with its gene ID. + let hasher = name_hasher(); + let mut claims: HashMap<&str, u32> = HashMap::new(); + let mut owner_of_id: HashMap<&str, u32> = HashMap::new(); + for (i, record) in out.iter().enumerate() { + if record.n_names == 0 { + continue; + } + let gene_id = name_at(&text, &name_ends, record.first_name); + claims.insert(gene_id, record.first_name); + owner_of_id.insert(gene_id, i as u32); + for pos in record.first_name + 1..record.first_name + u32::from(record.n_names) { + claims.entry(name_at(&text, &name_ends, pos)).or_insert(pos); + } + } + for record in out.iter_mut() { + if record.n_names > 0 { + record.alias_owner = owner_of_id[name_at(&text, &name_ends, record.first_name)]; + } + } + let mut name_index: Vec<(u32, u32)> = claims + .into_iter() + .map(|(name, pos)| (hasher.hash_one(name) as u32, pos)) + .collect(); + name_index.sort_unstable(); + drop(owner_of_id); + + // The mobile kind of each alias owner, from the labels it kept. + let mut mobile: Vec<(u32, &'static str, Box)> = Vec::new(); + for (i, labels) in labels_by_record.into_iter().enumerate() { + let Some(labels) = labels else { + continue; + }; + if out[i].alias_owner != i as u32 { + continue; + } + let refs: Vec<&str> = labels.iter().map(String::as_str).collect(); + if let Some((kind, joined)) = mobile_kind_from_labels(&refs) { + mobile.push((i as u32, kind, joined.into_boxed_str())); + } + } + + let contig_ids = contig_names + .iter() + .enumerate() + .map(|(id, name)| (name.clone(), id as u32)) + .collect(); + text.shrink_to_fit(); + name_ends.shrink_to_fit(); GffIndex { - genes, + text, + name_ends, + records: out, + contig_names, + contig_ids, contig_ranges, prefix_max_end, max_gene_len, - gene_id_index, - gene_aliases, - gene_labels, + name_index, + mobile, + hasher, } } - pub fn get_by_id(&self, gene_id: &str) -> Option<&GffGene> { - self.gene_id_index.get(gene_id).map(|&i| &self.genes[i]) + fn name(&self, pos: u32) -> &str { + name_at(&self.text, &self.name_ends, pos) } - pub fn aliases_for_gene<'a>(&'a self, gene: &'a GffGene) -> Vec<&'a str> { - self.gene_aliases - .get(&gene.gene_id) - .map(|aliases| aliases.iter().map(String::as_str).collect()) - .unwrap_or_else(|| vec![gene.gene_id.as_str()]) + fn view(&self, i: usize) -> GffGene<'_> { + let record = &self.records[i]; + GffGene { + contig: &self.contig_names[record.contig as usize], + start: record.start, + end: record.end, + strand: record.strand as char, + gene_id: if record.n_names == 0 { + "" + } else { + self.name(record.first_name) + }, + index: i as u32, + } } - pub fn labels_for_gene<'a>(&'a self, gene: &'a GffGene) -> Vec<&'a str> { - self.gene_labels - .get(&gene.gene_id) - .map(|labels| labels.iter().map(String::as_str).collect()) - .unwrap_or_else(|| self.aliases_for_gene(gene)) + /// Every CDS in contig and start order. + pub fn genes(&self) -> impl Iterator> + '_ { + (0..self.records.len()).map(|i| self.view(i)) } - pub fn cluster_for_gene( - &self, - gene: &GffGene, - gene_to_cluster: &HashMap, - ) -> Option { - gene_to_cluster.get(&gene.gene_id).copied().or_else(|| { - self.gene_aliases - .get(&gene.gene_id) - .and_then(|aliases| aliases.iter().find_map(|alias| gene_to_cluster.get(alias))) - .copied() - }) + /// The records of one contig, as a range of positions. + fn contig_range(&self, contig: &str) -> Option<(usize, usize)> { + let id = *self.contig_ids.get(contig)?; + let (start, end) = self.contig_ranges[id as usize]; + Some((start as usize, end as usize)) + } + + /// Position of the name that answers for `name`, if the annotation uses it. + fn name_position(&self, name: &str) -> Option { + let hash = self.hasher.hash_one(name) as u32; + let from = self.name_index.partition_point(|entry| entry.0 < hash); + self.name_index[from..] + .iter() + .take_while(|entry| entry.0 == hash) + .map(|entry| entry.1) + .find(|&pos| self.name(pos) == name) + } + + /// The record a name belongs to. Records hold their names in order, so it + /// is the last record whose names start at or before it. + fn record_of_name(&self, pos: u32) -> usize { + self.records + .partition_point(|record| record.first_name <= pos) + - 1 + } + + pub fn get_by_id(&self, gene_id: &str) -> Option> { + self.name_position(gene_id) + .map(|pos| self.view(self.record_of_name(pos))) } - pub fn overlapping(&self, contig: &str, query_start: u64, query_end: u64) -> Vec<&GffGene> { - let Some(&(range_start, range_end)) = self.contig_ranges.get(contig) else { + /// The record behind a gene this index handed out. The lookups below take + /// the gene's position, not its ID, so a gene from another index would + /// read the wrong record. + fn record_of(&self, gene: &GffGene<'_>) -> &Record { + let record = &self.records[gene.index as usize]; + debug_assert!( + record.start == gene.start && record.end == gene.end, + "gene {} is not from this index", + gene.gene_id + ); + record + } + + /// The names that answer for a gene: its alias owner's, gene ID first. + fn alias_positions(&self, record: &Record) -> std::ops::Range { + let owner = &self.records[record.alias_owner as usize]; + owner.first_name..owner.first_name + u32::from(owner.n_names) + } + + /// The gene ID and every alias the gene answers to, gene ID first. + pub fn aliases_for_gene<'a>(&'a self, gene: &GffGene<'a>) -> Vec<&'a str> { + let record = self.record_of(gene); + if record.n_names == 0 { + return vec![gene.gene_id]; + } + self.alias_positions(record) + .map(|pos| self.name(pos)) + .collect() + } + + /// What the gene's names and descriptions say about it being a mobile + /// element: the kind and the joined labels. + pub fn mobile_annotation(&self, gene: &GffGene<'_>) -> Option<(&'static str, &str)> { + let record = self.record_of(gene); + if record.n_names == 0 { + return None; + } + let owner = record.alias_owner; + let at = self.mobile.partition_point(|entry| entry.0 < owner); + self.mobile + .get(at) + .filter(|entry| entry.0 == owner) + .map(|entry| (entry.1, &*entry.2)) + } + + pub fn overlapping(&self, contig: &str, query_start: u64, query_end: u64) -> Vec> { + let Some((range_start, range_end)) = self.contig_range(contig) else { return Vec::new(); }; - - let slice = &self.genes[range_start..range_end]; let running_end = &self.prefix_max_end[range_start..range_end]; - let scan_start = running_end.partition_point(|&end| end <= query_start); + let scan_start = range_start + running_end.partition_point(|&end| end <= query_start); let mut results = Vec::new(); - for g in &slice[scan_start..] { - if g.start >= query_end { + for i in scan_start..range_end { + let record = &self.records[i]; + if record.start >= query_end { break; } - if g.end > query_start { - results.push(g); + if record.end > query_start { + results.push(self.view(i)); } } results @@ -175,37 +416,36 @@ impl GffIndex { contig: &str, start: u64, end: u64, - ) -> (Option<&GffGene>, Option<&GffGene>) { - let Some(&(range_start, range_end)) = self.contig_ranges.get(contig) else { + ) -> (Option>, Option>) { + let Some((range_start, range_end)) = self.contig_range(contig) else { return (None, None); }; - - let slice = &self.genes[range_start..range_end]; + let slice = &self.records[range_start..range_end]; // The first gene starting at or after `end`: starts are sorted. let down = slice.partition_point(|g| g.start < end); - let downstream = slice.get(down); + let downstream = (down < slice.len()).then(|| self.view(range_start + down)); // The gene ending closest before `start`. Ends are not sorted when // genes nest, so look back from the last gene starting before `start`, // stopping once no earlier gene can end later than the best so far. - let longest = self.max_gene_len.get(contig).copied().unwrap_or(0); - let mut upstream: Option<&GffGene> = None; - for g in slice[..slice.partition_point(|g| g.start < start)] - .iter() - .rev() - { + let contig_id = self.records[range_start].contig as usize; + let longest = self.max_gene_len[contig_id]; + let mut upstream: Option = None; + let before = slice.partition_point(|g| g.start < start); + for k in (0..before).rev() { + let g = &slice[k]; if let Some(best) = upstream { - if g.start.saturating_add(longest) < best.end { + if g.start.saturating_add(longest) < slice[best].end { break; } } - if g.end <= start && upstream.is_none_or(|best| g.end > best.end) { - upstream = Some(g); + if g.end <= start && upstream.is_none_or(|best| g.end > slice[best].end) { + upstream = Some(k); } } - (upstream, downstream) + (upstream.map(|k| self.view(range_start + k)), downstream) } pub fn context_window( @@ -215,17 +455,18 @@ impl GffIndex { end: u64, center_gene_id: Option<&str>, n_neighbors: usize, - ) -> (Vec<&GffGene>, Vec<&GffGene>) { - let Some(&(range_start, range_end)) = self.contig_ranges.get(contig) else { + ) -> (Vec>, Vec>) { + let Some((range_start, range_end)) = self.contig_range(contig) else { return (Vec::new(), Vec::new()); }; - let slice = &self.genes[range_start..range_end]; + let slice = &self.records[range_start..range_end]; if slice.is_empty() { return (Vec::new(), Vec::new()); } let center_pos = center_gene_id - .and_then(|gene_id| self.gene_id_index.get(gene_id).copied()) + .and_then(|gene_id| self.name_position(gene_id)) + .map(|pos| self.record_of_name(pos)) .and_then(|idx| { if idx >= range_start && idx < range_end { Some(idx - range_start) @@ -250,13 +491,12 @@ impl GffIndex { .map(|(i, _)| i) }); + let views = |positions: &mut dyn Iterator| -> Vec> { + positions.map(|k| self.view(range_start + k)).collect() + }; if let Some(center_pos) = center_pos { - let upstream = slice[..center_pos].iter().rev().take(n_neighbors).collect(); - let downstream = slice - .iter() - .skip(center_pos + 1) - .take(n_neighbors) - .collect(); + let upstream = views(&mut (0..center_pos).rev().take(n_neighbors)); + let downstream = views(&mut (center_pos + 1..slice.len()).take(n_neighbors)); return (upstream, downstream); } @@ -265,8 +505,135 @@ impl GffIndex { let gene_mid = g.start.saturating_add(g.end.saturating_sub(g.start) / 2); gene_mid < midpoint }); - let upstream = slice[..pos].iter().rev().take(n_neighbors).collect(); - let downstream = slice[pos..].iter().take(n_neighbors).collect(); + let upstream = views(&mut (0..pos).rev().take(n_neighbors)); + let downstream = views(&mut (pos..slice.len()).take(n_neighbors)); (upstream, downstream) } + + /// Bytes this index holds on the heap, roughly. + #[cfg(test)] + pub(crate) fn heap_bytes(&self) -> usize { + self.text.capacity() + + self.name_ends.capacity() * 4 + + self.records.capacity() * std::mem::size_of::() + + self.prefix_max_end.capacity() * 8 + + self.name_index.capacity() * 8 + + self.contig_ranges.capacity() * 8 + + self.max_gene_len.capacity() * 8 + + self + .contig_names + .iter() + .map(|name| 2 * (name.len() + 16) + 4) + .sum::() + + self + .mobile + .iter() + .map(|m| m.2.len() + std::mem::size_of::<(u32, &str, Box)>()) + .sum::() + } +} + +const NO_CLUSTER: u32 = u32::MAX; + +fn pack_cluster(ci: ClusterIdx) -> u32 { + match u32::try_from(ci.val()) { + Ok(packed) if packed != NO_CLUSTER => packed, + _ => panic!("cluster index {} does not fit in 32 bits", ci.val()), + } +} + +fn unpack_cluster(packed: u32) -> Option { + (packed != NO_CLUSTER).then_some(ClusterIdx(packed as usize)) +} + +/// The cluster each name of one genome's annotation belongs to. +/// +/// This was a map from every name (matrix cell IDs and their variants, gene +/// IDs, aliases) to its cluster, a String and a hash slot each: about 1.1 MB +/// per genome on PGAP-style annotations, held twice once the corrected copy +/// was made. Every lookup a run makes is for a name the annotation uses, so one +/// cluster for each name of the index answers it as the map did, in 4 bytes. +pub struct GeneClusters<'a> { + gff: Option<&'a GffIndex>, + /// Aligned with the index's names. + by_name: Vec, + /// What the map held for the empty name, which a gene with no ID asks for. + empty_name: u32, +} + +impl<'a> GeneClusters<'a> { + /// For a genome with no annotation, where nothing has a cluster. + pub fn none() -> Self { + GeneClusters { + gff: None, + by_name: Vec::new(), + empty_name: NO_CLUSTER, + } + } + + /// What `map` says about every name of `gff`. + pub fn from_map(gff: &'a GffIndex, map: &HashMap) -> Self { + let pack = |ci: Option<&ClusterIdx>| ci.map_or(NO_CLUSTER, |ci| pack_cluster(*ci)); + GeneClusters { + gff: Some(gff), + by_name: (0..gff.name_ends.len() as u32) + .map(|pos| pack(map.get(gff.name(pos)))) + .collect(), + empty_name: pack(map.get("")), + } + } + + fn at(&self, pos: u32) -> Option { + unpack_cluster(self.by_name[pos as usize]) + } + + /// The cluster of a name the annotation uses: what the map held for it. + /// A name the annotation does not use has none. + pub fn get(&self, name: &str) -> Option { + if name.is_empty() { + return unpack_cluster(self.empty_name); + } + let gff = self.gff?; + gff.name_position(name).and_then(|pos| self.at(pos)) + } + + /// The gene's cluster: its ID's, or else that of the first of its names + /// that has one. + pub fn cluster_for_gene(&self, gene: &GffGene<'_>) -> Option { + let gff = self.gff?; + let record = gff.record_of(gene); + if record.n_names == 0 { + return unpack_cluster(self.empty_name); + } + self.at(record.first_name) + .or_else(|| gff.alias_positions(record).find_map(|pos| self.at(pos))) + } + + /// Gives each correction's cluster to every name spelled like its gene ID + /// that has none yet, as adding the gene ID to the map did when it was not + /// already there. The first correction of a name wins. + pub fn fill(&mut self, corrections: &[(&str, ClusterIdx)]) { + let Some(gff) = self.gff else { + return; + }; + let mut first: HashMap<&str, u32> = HashMap::new(); + for &(name, ci) in corrections { + first.entry(name).or_insert(pack_cluster(ci)); + } + if let Some(&ci) = first.get("") { + if self.empty_name == NO_CLUSTER { + self.empty_name = ci; + } + } + for (pos, slot) in self.by_name.iter_mut().enumerate() { + if *slot == NO_CLUSTER { + if let Some(&ci) = first.get(gff.name(pos as u32)) { + *slot = ci; + } + } + } + } } + +#[cfg(test)] +mod tests; diff --git a/src/rescue/gff/tests.rs b/src/rescue/gff/tests.rs new file mode 100644 index 0000000..3c728f5 --- /dev/null +++ b/src/rescue/gff/tests.rs @@ -0,0 +1,474 @@ +use super::*; + +/// The index as it was before it was made compact: every name a String, one +/// map from names to genes and one map each for aliases and labels keyed by +/// gene ID. The compact index has to answer every question exactly as this did. +mod model { + use crate::types::ClusterIdx; + use ahash::{HashMap, HashMapExt}; + + #[derive(Debug, Clone)] + pub struct Gene { + pub contig: String, + pub start: u64, + pub end: u64, + pub strand: char, + pub gene_id: String, + } + + pub struct Index { + pub genes: Vec, + contig_ranges: HashMap, + prefix_max_end: Vec, + max_gene_len: HashMap, + pub gene_id_index: HashMap, + gene_aliases: HashMap>, + gene_labels: HashMap>, + } + + impl Index { + pub fn build(mut records: Vec<(Gene, Vec, Vec)>) -> Self { + records.sort_by(|a, b| a.0.contig.cmp(&b.0.contig).then(a.0.start.cmp(&b.0.start))); + let genes: Vec = records.iter().map(|(gene, _, _)| gene.clone()).collect(); + let mut contig_ranges = HashMap::new(); + if !genes.is_empty() { + let mut current = &genes[0].contig; + let mut range_start = 0; + for (i, g) in genes.iter().enumerate() { + if &g.contig != current { + contig_ranges.insert(current.clone(), (range_start, i)); + current = &g.contig; + range_start = i; + } + } + contig_ranges.insert(current.clone(), (range_start, genes.len())); + } + let mut prefix_max_end: Vec = Vec::new(); + let mut max_gene_len: HashMap = HashMap::new(); + for (i, g) in genes.iter().enumerate() { + let new_contig = i == 0 || genes[i - 1].contig != g.contig; + let running = if new_contig { + g.end + } else { + prefix_max_end[i - 1].max(g.end) + }; + prefix_max_end.push(running); + let longest = max_gene_len.entry(g.contig.clone()).or_insert(0); + *longest = (*longest).max(g.end.saturating_sub(g.start)); + } + let mut gene_id_index = HashMap::new(); + let mut gene_aliases = HashMap::new(); + let mut gene_labels = HashMap::new(); + for (i, (g, aliases, labels)) in records.iter().enumerate() { + if !g.gene_id.is_empty() { + gene_id_index.insert(g.gene_id.clone(), i); + let mut all_aliases = vec![g.gene_id.clone()]; + for alias in aliases { + if !alias.is_empty() && !all_aliases.contains(alias) { + gene_id_index.entry(alias.clone()).or_insert(i); + all_aliases.push(alias.clone()); + } + } + let mut all_labels = all_aliases.clone(); + for label in labels { + if !label.is_empty() && !all_labels.contains(label) { + all_labels.push(label.clone()); + } + } + gene_aliases.insert(g.gene_id.clone(), all_aliases); + gene_labels.insert(g.gene_id.clone(), all_labels); + } + } + Index { + genes, + contig_ranges, + prefix_max_end, + max_gene_len, + gene_id_index, + gene_aliases, + gene_labels, + } + } + + pub fn get_by_id(&self, gene_id: &str) -> Option<&Gene> { + self.gene_id_index.get(gene_id).map(|&i| &self.genes[i]) + } + + pub fn aliases_for_gene<'a>(&'a self, gene: &'a Gene) -> Vec<&'a str> { + self.gene_aliases + .get(&gene.gene_id) + .map(|aliases| aliases.iter().map(String::as_str).collect()) + .unwrap_or_else(|| vec![gene.gene_id.as_str()]) + } + + pub fn labels_for_gene<'a>(&'a self, gene: &'a Gene) -> Vec<&'a str> { + self.gene_labels + .get(&gene.gene_id) + .map(|labels| labels.iter().map(String::as_str).collect()) + .unwrap_or_else(|| self.aliases_for_gene(gene)) + } + + pub fn cluster_for_gene( + &self, + gene: &Gene, + gene_to_cluster: &HashMap, + ) -> Option { + gene_to_cluster.get(&gene.gene_id).copied().or_else(|| { + self.gene_aliases + .get(&gene.gene_id) + .and_then(|aliases| aliases.iter().find_map(|alias| gene_to_cluster.get(alias))) + .copied() + }) + } + + pub fn overlapping(&self, contig: &str, query_start: u64, query_end: u64) -> Vec<&Gene> { + let Some(&(range_start, range_end)) = self.contig_ranges.get(contig) else { + return Vec::new(); + }; + let slice = &self.genes[range_start..range_end]; + let running_end = &self.prefix_max_end[range_start..range_end]; + let scan_start = running_end.partition_point(|&end| end <= query_start); + let mut results = Vec::new(); + for g in &slice[scan_start..] { + if g.start >= query_end { + break; + } + if g.end > query_start { + results.push(g); + } + } + results + } + + pub fn flanking( + &self, + contig: &str, + start: u64, + end: u64, + ) -> (Option<&Gene>, Option<&Gene>) { + let Some(&(range_start, range_end)) = self.contig_ranges.get(contig) else { + return (None, None); + }; + let slice = &self.genes[range_start..range_end]; + let down = slice.partition_point(|g| g.start < end); + let downstream = slice.get(down); + let longest = self.max_gene_len.get(contig).copied().unwrap_or(0); + let mut upstream: Option<&Gene> = None; + for g in slice[..slice.partition_point(|g| g.start < start)] + .iter() + .rev() + { + if let Some(best) = upstream { + if g.start.saturating_add(longest) < best.end { + break; + } + } + if g.end <= start && upstream.is_none_or(|best| g.end > best.end) { + upstream = Some(g); + } + } + (upstream, downstream) + } + + pub fn context_window( + &self, + contig: &str, + start: u64, + end: u64, + center_gene_id: Option<&str>, + n_neighbors: usize, + ) -> (Vec<&Gene>, Vec<&Gene>) { + let Some(&(range_start, range_end)) = self.contig_ranges.get(contig) else { + return (Vec::new(), Vec::new()); + }; + let slice = &self.genes[range_start..range_end]; + if slice.is_empty() { + return (Vec::new(), Vec::new()); + } + let center_pos = center_gene_id + .and_then(|gene_id| self.gene_id_index.get(gene_id).copied()) + .and_then(|idx| { + if idx >= range_start && idx < range_end { + Some(idx - range_start) + } else { + None + } + }) + .or_else(|| { + slice + .iter() + .enumerate() + .filter_map(|(i, gene)| { + let overlap = gene.end.min(end).saturating_sub(gene.start.max(start)); + let query_len = end.saturating_sub(start); + let gene_len = gene.end.saturating_sub(gene.start); + let substantial = overlap > 0 + && ((query_len > 0 && overlap as f64 / query_len as f64 >= 0.5) + || (gene_len > 0 && overlap as f64 / gene_len as f64 >= 0.5)); + substantial.then_some((i, overlap)) + }) + .max_by_key(|(_, overlap)| *overlap) + .map(|(i, _)| i) + }); + if let Some(center_pos) = center_pos { + let upstream = slice[..center_pos].iter().rev().take(n_neighbors).collect(); + let downstream = slice + .iter() + .skip(center_pos + 1) + .take(n_neighbors) + .collect(); + return (upstream, downstream); + } + let midpoint = start.saturating_add(end.saturating_sub(start) / 2); + let pos = slice.partition_point(|g| { + let gene_mid = g.start.saturating_add(g.end.saturating_sub(g.start) / 2); + gene_mid < midpoint + }); + let upstream = slice[..pos].iter().rev().take(n_neighbors).collect(); + let downstream = slice[pos..].iter().take(n_neighbors).collect(); + (upstream, downstream) + } + } +} + +/// xorshift64*: reproducible draws without a dependency. +struct Draws(u64); + +impl Draws { + fn next(&mut self) -> u64 { + self.0 ^= self.0 >> 12; + self.0 ^= self.0 << 25; + self.0 ^= self.0 >> 27; + self.0.wrapping_mul(0x2545_f491_4f6c_dd1d) + } + + fn below(&mut self, n: u64) -> u64 { + self.next() % n + } + + fn pick<'a>(&mut self, from: &[&'a str]) -> &'a str { + from[self.below(from.len() as u64) as usize] + } +} + +type Key = (String, u64, u64, char, String); + +fn key_of(gene: &GffGene<'_>) -> Key { + ( + gene.contig.to_string(), + gene.start, + gene.end, + gene.strand, + gene.gene_id.to_string(), + ) +} + +fn model_key(gene: &model::Gene) -> Key { + ( + gene.contig.clone(), + gene.start, + gene.end, + gene.strand, + gene.gene_id.clone(), + ) +} + +const GENE_IDS: &[&str] = &["", "g1", "g2", "g3", "g4", "g5", "g6", "g7", "g8"]; +const ALIASES: &[&str] = &[ + "", "g1", "g2", "g3", "g4", "g5", "g6", "g7", "g8", "a1", "a2", "a3", "a4", "a5", "a6", +]; +const LABELS: &[&str] = &[ + "", + "transposase IS6110", + "integrase", + "hypothetical protein", + "phage tail protein", + "site-specific recombinase", + "insertion_sequence element", +]; +const CONTIGS: &[&str] = &["c1", "c2", "c3"]; + +/// The compact index and its per-name clusters answer every question the old +/// index and the old map of names answered, the same way: names claimed by +/// two genes, aliases that are another gene's ID, genes with no ID, nested +/// genes, corrections that name one gene ID twice. +#[test] +fn compact_index_answers_like_the_maps_it_replaced() { + let mut draws = Draws(0x9e37_79b9_7f4a_7c15); + for trial in 0..400 { + let n_records = draws.below(40) as usize; + let mut records = Vec::new(); + let mut model_records = Vec::new(); + for _ in 0..n_records { + let contig = draws.pick(CONTIGS).to_string(); + let start = draws.below(2000); + let end = start + draws.below(600); + let strand = if draws.below(2) == 0 { '+' } else { '-' }; + let gene_id = draws.pick(GENE_IDS).to_string(); + let aliases: Vec = (0..draws.below(5)) + .map(|_| draws.pick(ALIASES).to_string()) + .collect(); + let labels: Vec = (0..draws.below(3)) + .map(|_| draws.pick(LABELS).to_string()) + .collect(); + records.push(( + GffRecord { + contig: contig.clone(), + start, + end, + strand, + gene_id: gene_id.clone(), + }, + aliases.clone(), + labels.clone(), + )); + model_records.push(( + model::Gene { + contig, + start, + end, + strand, + gene_id, + }, + aliases, + labels, + )); + } + let index = GffIndex::build_with_aliases_and_labels(records); + let old = model::Index::build(model_records); + + let genes: Vec = index.genes().collect(); + let keys: Vec = genes.iter().map(key_of).collect(); + let model_keys: Vec = old.genes.iter().map(model_key).collect(); + assert_eq!(keys, model_keys, "trial {trial}: gene order"); + + for name in ALIASES.iter().chain(["zz", "c1"].iter()) { + assert_eq!( + index.get_by_id(name).as_ref().map(key_of), + old.get_by_id(name).map(model_key), + "trial {trial}: get_by_id({name:?})" + ); + } + for (gene, model_gene) in genes.iter().zip(&old.genes) { + assert_eq!( + index.aliases_for_gene(gene), + old.aliases_for_gene(model_gene), + "trial {trial}: aliases of {model_gene:?}" + ); + let expected = mobile_kind_from_labels(&old.labels_for_gene(model_gene)); + assert_eq!( + index.mobile_annotation(gene), + expected + .as_ref() + .map(|(kind, label)| (*kind, label.as_str())), + "trial {trial}: mobile annotation of {model_gene:?}" + ); + } + + // A map from names to clusters: some of the annotation's names, some + // it does not use, and sometimes the empty name. + let mut map: HashMap = HashMap::new(); + for name in ALIASES.iter().chain(["x1", "x2"].iter()) { + if draws.below(3) == 0 { + map.insert(name.to_string(), ClusterIdx(draws.below(5) as usize)); + } + } + let mut clusters = GeneClusters::from_map(&index, &map); + let check = |clusters: &GeneClusters, map: &HashMap, stage: &str| { + for name in old.gene_id_index.keys().map(String::as_str).chain([""]) { + assert_eq!( + clusters.get(name), + map.get(name).copied(), + "trial {trial} {stage}: cluster of {name:?}" + ); + } + for (gene, model_gene) in genes.iter().zip(&old.genes) { + assert_eq!( + clusters.cluster_for_gene(gene), + old.cluster_for_gene(model_gene, map), + "trial {trial} {stage}: cluster of {model_gene:?}" + ); + } + }; + check(&clusters, &map, "input"); + + // Corrections go in as entries added only when the name is absent. + let corrections: Vec<(&str, ClusterIdx)> = (0..draws.below(4)) + .map(|_| (draws.pick(GENE_IDS), ClusterIdx(draws.below(5) as usize))) + .collect(); + let mut corrected = map.clone(); + for &(name, ci) in &corrections { + corrected.entry(name.to_string()).or_insert(ci); + } + clusters.fill(&corrections); + check(&clusters, &corrected, "corrected"); + + for _ in 0..20 { + let contig = draws.pick(&["c1", "c2", "c3", "c4"]); + let start = draws.below(2400); + let end = start + draws.below(700); + let as_keys = |genes: Vec| genes.iter().map(key_of).collect::>(); + let as_model_keys = + |genes: Vec<&model::Gene>| genes.into_iter().map(model_key).collect::>(); + assert_eq!( + as_keys(index.overlapping(contig, start, end)), + as_model_keys(old.overlapping(contig, start, end)), + "trial {trial}: overlapping {contig}:{start}-{end}" + ); + let (up, down) = index.flanking(contig, start, end); + let (model_up, model_down) = old.flanking(contig, start, end); + assert_eq!(up.as_ref().map(key_of), model_up.map(model_key)); + assert_eq!(down.as_ref().map(key_of), model_down.map(model_key)); + let center = if draws.below(2) == 0 { + None + } else { + Some(draws.pick(ALIASES)) + }; + let n = draws.below(4) as usize; + let (up, down) = index.context_window(contig, start, end, center, n); + let (model_up, model_down) = old.context_window(contig, start, end, center, n); + assert_eq!( + (as_keys(up), as_keys(down)), + (as_model_keys(model_up), as_model_keys(model_down)), + "trial {trial}: context of {contig}:{start}-{end} around {center:?}" + ); + } + } +} + +/// What the index costs per gene on a PGAP-style annotation, where every CDS +/// has six names: under 250 bytes, against about 1,200 when each name was a +/// String in several maps. A guard against storing names per gene again. +#[test] +fn index_stays_compact_on_pgap_style_annotation() { + let n_genes = 4000u64; + let records: Vec<(GffRecord, Vec, Vec)> = (0..n_genes) + .map(|i| { + let locus = format!("G00000_{i:05}"); + ( + GffRecord { + contig: "G00000_1".into(), + start: i * 1200, + end: i * 1200 + 1000, + strand: '+', + gene_id: locus.clone(), + }, + vec![ + format!("cds-{locus}"), + locus, + format!("gen{i:05}"), + format!("WP_{i:09}.1"), + format!("Rv{i:04}"), + format!("{}", 1_000_000 + i), + ], + vec!["DNA-directed RNA polymerase subunit beta".into()], + ) + }) + .collect(); + let index = GffIndex::build_with_aliases_and_labels(records); + let per_gene = index.heap_bytes() as u64 / n_genes; + assert!(per_gene < 250, "{per_gene} bytes per gene"); + let map: HashMap = HashMap::new(); + let clusters = GeneClusters::from_map(&index, &map); + assert_eq!(clusters.by_name.len() as u64, 6 * n_genes); +} diff --git a/src/rescue/synteny.rs b/src/rescue/synteny.rs index 91a9bb8..9dd144f 100644 --- a/src/rescue/synteny.rs +++ b/src/rescue/synteny.rs @@ -1,5 +1,5 @@ use super::genes::cell_gene_id_variants; -use super::gff::GffIndex; +use super::gff::{GeneClusters, GffIndex}; use super::types::{RescueEvent, SyntenyEvidence}; use crate::parser::PresenceAbsence; use crate::types::{ClusterIdx, GenomeIdx}; @@ -146,7 +146,7 @@ fn context_score(target: &SyntenyContext, reference: &SyntenyContext) -> (f64, u fn synteny_context( gff: &GffIndex, - gene_to_cluster: &HashMap, + gene_to_cluster: &GeneClusters, contig: &str, start: u64, end: u64, @@ -157,12 +157,12 @@ fn synteny_context( let (upstream, downstream) = gff.context_window(contig, start, end, center_gene_id, n_neighbors); for gene in upstream { - if let Some(cluster) = gff.cluster_for_gene(gene, gene_to_cluster) { + if let Some(cluster) = gene_to_cluster.cluster_for_gene(&gene) { context.upstream.push(cluster); } } for gene in downstream { - if let Some(cluster) = gff.cluster_for_gene(gene, gene_to_cluster) { + if let Some(cluster) = gene_to_cluster.cluster_for_gene(&gene) { context.downstream.push(cluster); } } @@ -281,7 +281,7 @@ fn reference_contexts_for_cluster<'a>( n_neighbors: usize, pa: &PresenceAbsence, gff_indices: &HashMap, - gene_to_cluster_maps: &HashMap>, + gene_to_cluster_maps: &HashMap, verbose: bool, ) -> &'a [ReferenceSyntenyContext] { cache @@ -328,10 +328,10 @@ fn reference_contexts_for_cluster<'a>( let ref_context = synteny_context( ref_gff, ref_g2c, - &ref_gene.contig, + ref_gene.contig, ref_gene.start, ref_gene.end, - Some(ref_gene.gene_id.as_str()), + Some(ref_gene.gene_id), n_neighbors, ); @@ -353,7 +353,7 @@ pub(super) fn validate_synteny( rescue: &mut RescueEvent, pa: &PresenceAbsence, gff_indices: &HashMap, - gene_to_cluster_maps: &HashMap>, + gene_to_cluster_maps: &HashMap, verbose: bool, ) { let mut cache = SyntenyReferenceCache::default(); @@ -371,7 +371,7 @@ pub(crate) fn validate_synteny_with_cache( rescue: &mut RescueEvent, pa: &PresenceAbsence, gff_indices: &HashMap, - gene_to_cluster_maps: &HashMap>, + gene_to_cluster_maps: &HashMap, cache: &mut SyntenyReferenceCache, verbose: bool, ) { diff --git a/src/rescue/tests.rs b/src/rescue/tests.rs index 1004dc0..78c1cb0 100644 --- a/src/rescue/tests.rs +++ b/src/rescue/tests.rs @@ -117,21 +117,21 @@ fn test_merged_coverage_same_length_different_alleles() { #[test] fn test_gff_index_overlap() { let genes = vec![ - GffGene { + GffRecord { contig: "chr1".into(), start: 100, end: 500, strand: '+', gene_id: "g1".into(), }, - GffGene { + GffRecord { contig: "chr1".into(), start: 600, end: 900, strand: '-', gene_id: "g2".into(), }, - GffGene { + GffRecord { contig: "chr2".into(), start: 50, end: 200, @@ -150,21 +150,21 @@ fn test_gff_index_overlap() { #[test] fn test_gff_flanking() { let genes = vec![ - GffGene { + GffRecord { contig: "c".into(), start: 100, end: 200, strand: '+', gene_id: "g1".into(), }, - GffGene { + GffRecord { contig: "c".into(), start: 300, end: 400, strand: '+', gene_id: "g2".into(), }, - GffGene { + GffRecord { contig: "c".into(), start: 500, end: 600, @@ -204,35 +204,35 @@ fn test_synteny_context_handles_overlapping_boundary_genes() { ); let target_gff = GffIndex::build(vec![ - GffGene { + GffRecord { contig: "chr".into(), start: 0, end: 100, strand: '+', gene_id: "target_outer_up".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 100, end: 200, strand: '+', gene_id: "target_up".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 197, end: 260, strand: '+', gene_id: "target_center".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 257, end: 320, strand: '+', gene_id: "target_down".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 320, end: 400, @@ -241,35 +241,35 @@ fn test_synteny_context_handles_overlapping_boundary_genes() { }, ]); let ref_gff = GffIndex::build(vec![ - GffGene { + GffRecord { contig: "chr".into(), start: 0, end: 100, strand: '+', gene_id: "ref_outer_up".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 100, end: 200, strand: '+', gene_id: "ref_up".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 197, end: 260, strand: '+', gene_id: "ref_center".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 257, end: 320, strand: '+', gene_id: "ref_down".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 320, end: 400, @@ -283,8 +283,20 @@ fn test_synteny_context_handles_overlapping_boundary_genes() { gff_indices.insert(GenomeIdx(1), ref_gff); let mut g2c = HashMap::new(); - g2c.insert(GenomeIdx(0), build_gene_to_cluster(&pa, GenomeIdx(0))); - g2c.insert(GenomeIdx(1), build_gene_to_cluster(&pa, GenomeIdx(1))); + g2c.insert( + GenomeIdx(0), + GeneClusters::from_map( + &gff_indices[&GenomeIdx(0)], + &build_gene_to_cluster(&pa, GenomeIdx(0)), + ), + ); + g2c.insert( + GenomeIdx(1), + GeneClusters::from_map( + &gff_indices[&GenomeIdx(1)], + &build_gene_to_cluster(&pa, GenomeIdx(1)), + ), + ); let mut rescue = RescueEvent { cluster_idx: ClusterIdx(0), @@ -325,21 +337,21 @@ fn test_synteny_context_rejects_different_locus_paralog() { ); let target_gff = GffIndex::build(vec![ - GffGene { + GffRecord { contig: "chr".into(), start: 100, end: 200, strand: '+', gene_id: "target_up_gene".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 200, end: 300, strand: '+', gene_id: "target_center".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 300, end: 400, @@ -348,21 +360,21 @@ fn test_synteny_context_rejects_different_locus_paralog() { }, ]); let ref_gff = GffIndex::build(vec![ - GffGene { + GffRecord { contig: "chr".into(), start: 100, end: 200, strand: '+', gene_id: "ref_up_other".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 200, end: 300, strand: '+', gene_id: "ref_center".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 300, end: 400, @@ -376,8 +388,20 @@ fn test_synteny_context_rejects_different_locus_paralog() { gff_indices.insert(GenomeIdx(1), ref_gff); let mut g2c = HashMap::new(); - g2c.insert(GenomeIdx(0), build_gene_to_cluster(&pa, GenomeIdx(0))); - g2c.insert(GenomeIdx(1), build_gene_to_cluster(&pa, GenomeIdx(1))); + g2c.insert( + GenomeIdx(0), + GeneClusters::from_map( + &gff_indices[&GenomeIdx(0)], + &build_gene_to_cluster(&pa, GenomeIdx(0)), + ), + ); + g2c.insert( + GenomeIdx(1), + GeneClusters::from_map( + &gff_indices[&GenomeIdx(1)], + &build_gene_to_cluster(&pa, GenomeIdx(1)), + ), + ); let mut rescue = RescueEvent { cluster_idx: ClusterIdx(0), @@ -422,21 +446,21 @@ fn test_synteny_without_any_reference_annotation_is_not_assessed() { ); let target_gff = GffIndex::build(vec![ - GffGene { + GffRecord { contig: "chr".into(), start: 100, end: 200, strand: '+', gene_id: "target_up_gene".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 200, end: 300, strand: '+', gene_id: "target_center".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 300, end: 400, @@ -445,21 +469,21 @@ fn test_synteny_without_any_reference_annotation_is_not_assessed() { }, ]); let ref_gff = GffIndex::build(vec![ - GffGene { + GffRecord { contig: "chr".into(), start: 100, end: 200, strand: '+', gene_id: "ref_up_other".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 200, end: 300, strand: '+', gene_id: "ref_center".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 300, end: 400, @@ -472,9 +496,15 @@ fn test_synteny_without_any_reference_annotation_is_not_assessed() { gff_indices.insert(GenomeIdx(0), target_gff); drop(ref_gff); + // Genome 1 has no annotation, so it has no lookups either. let mut g2c = HashMap::new(); - g2c.insert(GenomeIdx(0), build_gene_to_cluster(&pa, GenomeIdx(0))); - g2c.insert(GenomeIdx(1), build_gene_to_cluster(&pa, GenomeIdx(1))); + g2c.insert( + GenomeIdx(0), + GeneClusters::from_map( + &gff_indices[&GenomeIdx(0)], + &build_gene_to_cluster(&pa, GenomeIdx(0)), + ), + ); let mut rescue = RescueEvent { cluster_idx: ClusterIdx(0), @@ -576,21 +606,21 @@ fn test_synteny_uses_all_present_genomes_deterministically() { fn genes(up: &str, center: &str, down: &str) -> GffIndex { GffIndex::build(vec![ - GffGene { + GffRecord { contig: "chr".into(), start: 100, end: 200, strand: '+', gene_id: up.into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 200, end: 300, strand: '+', gene_id: center.into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 300, end: 400, @@ -617,7 +647,13 @@ fn test_synteny_uses_all_present_genomes_deterministically() { let mut g2c = HashMap::new(); for gi in 0..8 { - g2c.insert(GenomeIdx(gi), build_gene_to_cluster(&pa, GenomeIdx(gi))); + g2c.insert( + GenomeIdx(gi), + GeneClusters::from_map( + &gff_indices[&GenomeIdx(gi)], + &build_gene_to_cluster(&pa, GenomeIdx(gi)), + ), + ); } let mut rescue = RescueEvent { @@ -663,21 +699,21 @@ fn test_synteny_uses_all_gene_ids_in_multicopy_reference_cell() { ); let target_gff = GffIndex::build(vec![ - GffGene { + GffRecord { contig: "chr".into(), start: 100, end: 200, strand: '+', gene_id: "target_up".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 200, end: 300, strand: '+', gene_id: "candidate".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 300, end: 400, @@ -686,42 +722,42 @@ fn test_synteny_uses_all_gene_ids_in_multicopy_reference_cell() { }, ]); let ref_gff = GffIndex::build(vec![ - GffGene { + GffRecord { contig: "chr".into(), start: 100, end: 200, strand: '+', gene_id: "ref_bad_up".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 200, end: 300, strand: '+', gene_id: "ref_bad".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 300, end: 400, strand: '+', gene_id: "ref_bad_down".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 1000, end: 1100, strand: '+', gene_id: "ref_good_up".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 1100, end: 1200, strand: '+', gene_id: "ref_good".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 1200, end: 1300, @@ -735,8 +771,20 @@ fn test_synteny_uses_all_gene_ids_in_multicopy_reference_cell() { gff_indices.insert(GenomeIdx(1), ref_gff); let mut g2c = HashMap::new(); - g2c.insert(GenomeIdx(0), build_gene_to_cluster(&pa, GenomeIdx(0))); - g2c.insert(GenomeIdx(1), build_gene_to_cluster(&pa, GenomeIdx(1))); + g2c.insert( + GenomeIdx(0), + GeneClusters::from_map( + &gff_indices[&GenomeIdx(0)], + &build_gene_to_cluster(&pa, GenomeIdx(0)), + ), + ); + g2c.insert( + GenomeIdx(1), + GeneClusters::from_map( + &gff_indices[&GenomeIdx(1)], + &build_gene_to_cluster(&pa, GenomeIdx(1)), + ), + ); let mut rescue = RescueEvent { cluster_idx: ClusterIdx(0), @@ -780,21 +828,21 @@ fn test_missing_cds_preferred_over_annotated_paralog_and_validated_by_synteny() ); let target_gff = GffIndex::build(vec![ - GffGene { + GffRecord { contig: "chr".into(), start: 100, end: 200, strand: '+', gene_id: "target_up".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 300, end: 400, strand: '+', gene_id: "target_down".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 1000, end: 1100, @@ -803,21 +851,21 @@ fn test_missing_cds_preferred_over_annotated_paralog_and_validated_by_synteny() }, ]); let ref1_gff = GffIndex::build(vec![ - GffGene { + GffRecord { contig: "chr".into(), start: 100, end: 200, strand: '+', gene_id: "ref_up1".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 200, end: 300, strand: '+', gene_id: "ref_center1".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 300, end: 400, @@ -826,21 +874,21 @@ fn test_missing_cds_preferred_over_annotated_paralog_and_validated_by_synteny() }, ]); let ref2_gff = GffIndex::build(vec![ - GffGene { + GffRecord { contig: "chr".into(), start: 100, end: 200, strand: '+', gene_id: "ref_up2".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 200, end: 300, strand: '+', gene_id: "ref_center2".into(), }, - GffGene { + GffRecord { contig: "chr".into(), start: 300, end: 400, @@ -856,7 +904,13 @@ fn test_missing_cds_preferred_over_annotated_paralog_and_validated_by_synteny() let mut g2c = HashMap::new(); for gi in 0..3 { - g2c.insert(GenomeIdx(gi), build_gene_to_cluster(&pa, GenomeIdx(gi))); + g2c.insert( + GenomeIdx(gi), + GeneClusters::from_map( + &gff_indices[&GenomeIdx(gi)], + &build_gene_to_cluster(&pa, GenomeIdx(gi)), + ), + ); } let annotated_paralog = Hit { @@ -1043,7 +1097,7 @@ fn test_synteny_recurrence_requires_multiple_genomes_for_large_reference_sets() #[test] fn test_detect_start_variant_with_gff() { - let genes = vec![GffGene { + let genes = vec![GffRecord { contig: "c".into(), start: 0, end: 1100, @@ -1076,7 +1130,7 @@ fn test_detect_start_variant_with_gff() { /// start variant. The old rule compared lengths and called this +50bp. #[test] fn test_detect_start_variant_projects_unaligned_query_ends() { - let genes = vec![GffGene { + let genes = vec![GffRecord { contig: "c".into(), start: 0, end: 1050, @@ -1109,7 +1163,7 @@ fn test_detect_start_variant_projects_unaligned_query_ends() { /// that reads through its stop used to be reported this way. #[test] fn test_detect_start_variant_ignores_stop_side_differences() { - let idx = GffIndex::build(vec![GffGene { + let idx = GffIndex::build(vec![GffRecord { contig: "c".into(), start: 0, end: 1100, @@ -1123,7 +1177,7 @@ fn test_detect_start_variant_ignores_stop_side_differences() { /// On the minus strand the start is the high coordinate. #[test] fn test_detect_start_variant_minus_strand_start_is_the_high_end() { - let idx = GffIndex::build(vec![GffGene { + let idx = GffIndex::build(vec![GffRecord { contig: "c".into(), start: 0, end: 1100, @@ -1311,7 +1365,6 @@ fn test_fragmented_span() { vec!["genome1".into()], vec![vec!["".into()]], ); - let g2c: HashMap = HashMap::new(); let result = classify_hits( &[h1, h2], @@ -1319,7 +1372,7 @@ fn test_fragmented_span() { GenomeIdx(0), &pa, None, - &g2c, + &GeneClusters::none(), 80.0, ); assert!(result.is_some()); @@ -1366,7 +1419,7 @@ fn test_best_fragment_chain_tie_breaks_by_locus() { fn test_classify_split_vs_paralog() { // Setup: cluster "target_cluster" absent in genome, hit lands on a gene // assigned to "other_cluster". Identity determines SPLIT vs PARALOG. - let gff_genes = vec![GffGene { + let gff_genes = vec![GffRecord { contig: "c".into(), start: 0, end: 1000, @@ -1406,7 +1459,7 @@ fn test_classify_split_vs_paralog() { GenomeIdx(0), &pa, Some(&gff), - &g2c, + &GeneClusters::from_map(&gff, &g2c), 80.0, ); let r = result.unwrap(); @@ -1437,7 +1490,7 @@ fn test_classify_split_vs_paralog() { GenomeIdx(0), &pa, Some(&gff), - &g2c, + &GeneClusters::from_map(&gff, &g2c), 80.0, ); let r = result.unwrap(); @@ -1471,8 +1524,15 @@ fn test_classify_dropped_no_gff() { vec!["genome1".into()], vec![vec!["".into()]], ); - let g2c: HashMap = HashMap::new(); - let result = classify_hits(&[hit], ClusterIdx(0), GenomeIdx(0), &pa, None, &g2c, 80.0); + let result = classify_hits( + &[hit], + ClusterIdx(0), + GenomeIdx(0), + &pa, + None, + &GeneClusters::none(), + 80.0, + ); let r = result.unwrap(); assert!(matches!(r.rescue_type, RescueType::Dropped)); assert!(r.confidence > 0.9); @@ -1515,14 +1575,13 @@ fn test_classify_prefers_best_evidence_over_longest_allele() { vec!["genome1".into()], vec![vec!["".into()]], ); - let g2c: HashMap = HashMap::new(); let result = classify_hits( &[long_lower_score, short_better_score], ClusterIdx(0), GenomeIdx(0), &pa, None, - &g2c, + &GeneClusters::none(), 80.0, ) .unwrap(); @@ -1535,7 +1594,7 @@ fn test_classify_prefers_best_evidence_over_longest_allele() { #[test] fn test_classify_missing_cds_no_overlapping_genes() { // Full-coverage hit, GFF exists but no gene overlaps the hit region → MISSING_CDS - let gff_genes = vec![GffGene { + let gff_genes = vec![GffRecord { contig: "c".into(), start: 5000, end: 6000, @@ -1570,7 +1629,7 @@ fn test_classify_missing_cds_no_overlapping_genes() { GenomeIdx(0), &pa, Some(&gff), - &g2c, + &GeneClusters::from_map(&gff, &g2c), 80.0, ); let r = result.unwrap(); @@ -1580,7 +1639,7 @@ fn test_classify_missing_cds_no_overlapping_genes() { #[test] fn test_classify_start_variant_with_gff() { // Hit with high identity overlapping a similar-length GFF gene → START_VARIANT - let gff_genes = vec![GffGene { + let gff_genes = vec![GffRecord { contig: "c".into(), start: 0, end: 1100, @@ -1617,7 +1676,7 @@ fn test_classify_start_variant_with_gff() { GenomeIdx(0), &pa, Some(&gff), - &g2c, + &GeneClusters::from_map(&gff, &g2c), 80.0, ); let r = result.unwrap(); @@ -1632,8 +1691,15 @@ fn test_classify_start_variant_with_gff() { fn test_classify_returns_none_for_empty_hits() { let pa = PresenceAbsence::new_for_test(vec!["c1".into()], vec!["g1".into()], vec![vec!["".into()]]); - let g2c: HashMap = HashMap::new(); - let result = classify_hits(&[], ClusterIdx(0), GenomeIdx(0), &pa, None, &g2c, 80.0); + let result = classify_hits( + &[], + ClusterIdx(0), + GenomeIdx(0), + &pa, + None, + &GeneClusters::none(), + 80.0, + ); assert!(result.is_none()); } @@ -1657,8 +1723,15 @@ fn test_classify_returns_none_below_coverage() { }; let pa = PresenceAbsence::new_for_test(vec!["c1".into()], vec!["g1".into()], vec![vec!["".into()]]); - let g2c: HashMap = HashMap::new(); - let result = classify_hits(&[hit], ClusterIdx(0), GenomeIdx(0), &pa, None, &g2c, 80.0); + let result = classify_hits( + &[hit], + ClusterIdx(0), + GenomeIdx(0), + &pa, + None, + &GeneClusters::none(), + 80.0, + ); assert!(result.is_none()); } @@ -1709,7 +1782,7 @@ fn test_build_gene_to_cluster_adds_unambiguous_gff_aliases() { vec![vec!["commonName".into()]], ); let gff = GffIndex::build_with_aliases(vec![( - GffGene { + GffRecord { contig: "ctg".into(), start: 0, end: 100, @@ -1734,7 +1807,7 @@ fn test_descriptive_gff_labels_do_not_create_cluster_aliases() { ); let gff = GffIndex::build_with_aliases_and_labels(vec![ ( - GffGene { + GffRecord { contig: "ctg".into(), start: 0, end: 100, @@ -1745,7 +1818,7 @@ fn test_descriptive_gff_labels_do_not_create_cluster_aliases() { vec!["commonName".into(), "hypothetical%20protein".into()], ), ( - GffGene { + GffRecord { contig: "ctg".into(), start: 200, end: 300, @@ -1771,7 +1844,7 @@ fn test_classify_hit_uses_gff_alias_for_split_cluster() { vec![vec!["".into()], vec!["commonName".into()]], ); let gff = GffIndex::build_with_aliases(vec![( - GffGene { + GffRecord { contig: "ctg".into(), start: 100, end: 1100, @@ -1803,7 +1876,7 @@ fn test_classify_hit_uses_gff_alias_for_split_cluster() { GenomeIdx(0), &pa, Some(&gff), - &g2c, + &GeneClusters::from_map(&gff, &g2c), 80.0, ) .unwrap(); @@ -2047,9 +2120,16 @@ fn test_classify_partial_deletion_pseudogene() { }; let pa = PresenceAbsence::new_for_test(vec!["c1".into()], vec!["g1".into()], vec![vec!["".into()]]); - let g2c: HashMap = HashMap::new(); // min_coverage = 80% so 70% coverage won't pass Path 1 - let result = classify_hits(&[hit], ClusterIdx(0), GenomeIdx(0), &pa, None, &g2c, 80.0); + let result = classify_hits( + &[hit], + ClusterIdx(0), + GenomeIdx(0), + &pa, + None, + &GeneClusters::none(), + 80.0, + ); let r = result.unwrap(); match &r.rescue_type { RescueType::Pseudogene { @@ -2098,14 +2178,13 @@ fn test_alignment_split_for_small_colinear_gaps() { vec!["genome1".into()], vec![vec!["".into()]], ); - let g2c: HashMap = HashMap::new(); let result = classify_hits( &[h1, h2], ClusterIdx(0), GenomeIdx(0), &pa, None, - &g2c, + &GeneClusters::none(), 80.0, ) .unwrap(); @@ -2152,7 +2231,7 @@ fn test_fragment_chain_inside_same_cds_is_not_fragmented() { alignment_length: 400, mapq: 60, }; - let gff = GffIndex::build(vec![GffGene { + let gff = GffIndex::build(vec![GffRecord { contig: "c".into(), start: 100, end: 1100, @@ -2171,7 +2250,7 @@ fn test_fragment_chain_inside_same_cds_is_not_fragmented() { GenomeIdx(0), &pa, Some(&gff), - &g2c, + &GeneClusters::from_map(&gff, &g2c), 80.0, ) .unwrap(); @@ -2217,14 +2296,13 @@ fn test_non_colinear_fragments_are_not_rescued() { vec!["genome1".into()], vec![vec!["".into()]], ); - let g2c: HashMap = HashMap::new(); let result = classify_hits( &[h1, h2], ClusterIdx(0), GenomeIdx(0), &pa, None, - &g2c, + &GeneClusters::none(), 80.0, ); assert!(result.is_none()); @@ -2234,7 +2312,7 @@ fn test_non_colinear_fragments_are_not_rescued() { /// contains another. The binary search on ends skipped the containing CDS. #[test] fn test_gff_index_finds_a_cds_that_contains_a_nested_one() { - let gene = |start, end, id: &str| GffGene { + let gene = |start, end, id: &str| GffRecord { contig: "c".into(), start, end, @@ -2249,17 +2327,17 @@ fn test_gff_index_finds_a_cds_that_contains_a_nested_one() { let ids: Vec<&str> = idx .overlapping("c", 3000, 4000) .iter() - .map(|g| g.gene_id.as_str()) + .map(|g| g.gene_id) .collect(); assert_eq!(ids, vec!["A"]); let (up, down) = idx.flanking("c", 5040, 5090); assert_eq!( - up.map(|g| g.gene_id.as_str()), + up.map(|g| g.gene_id), Some("A"), "nearest end is A's, not B's" ); - assert_eq!(down.map(|g| g.gene_id.as_str()), Some("C")); + assert_eq!(down.map(|g| g.gene_id), Some("C")); } /// The two halves of a gene broken by an insertion abut exactly on the query @@ -2311,7 +2389,7 @@ fn test_gene_split_across_contigs_is_fragmented() { GenomeIdx(0), &pa, None, - &HashMap::new(), + &GeneClusters::none(), 80.0, ) .unwrap(); @@ -2335,7 +2413,7 @@ fn test_unclaimed_annotated_cds_beats_an_unannotated_copy() { vec!["g".into(), "r".into()], vec![vec!["".into(), "r1".into()]], ); - let gff = GffIndex::build(vec![GffGene { + let gff = GffIndex::build(vec![GffRecord { contig: "chr".into(), start: 1000, end: 2000, @@ -2362,7 +2440,7 @@ fn test_unclaimed_annotated_cds_beats_an_unannotated_copy() { GenomeIdx(0), &pa, Some(&gff), - &HashMap::new(), + &GeneClusters::from_map(&gff, &HashMap::new()), 80.0, ) .unwrap(); @@ -2542,7 +2620,7 @@ fn piece( fn classify_pieces(hits: &[Hit]) -> Option { // An annotation exists but nothing in it overlaps either piece. - let gff = GffIndex::build(vec![GffGene { + let gff = GffIndex::build(vec![GffRecord { contig: "chr".into(), start: 9000, end: 9900, @@ -2561,7 +2639,7 @@ fn classify_pieces(hits: &[Hit]) -> Option { GenomeIdx(0), &pa, Some(&gff), - &g2c, + &GeneClusters::from_map(&gff, &g2c), 80.0, ) }