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, ) }