Add phylogenetic export support for TNT, PhyG, and IQ-TREE

Extends the phylogenetic pipeline with exporters for TNT, PhyG, and IQ-TREE that generate executable scripts, cost matrices, and recoded alignments. Adds internal helpers for state indexing, floating-point matrix scaling, and Floyd-Warshall metric closure to satisfy external tool constraints. Adjusts visibility modifiers for sibling iterators and entropy annex structs, and updates design documentation for evolutionary distance metrics.
This commit is contained in:
Eric Coissac
2026-08-28 21:32:32 +02:00
parent 1536217ac1
commit 0b40d2d0da
12 changed files with 1057 additions and 41 deletions
+86
View File
@@ -2182,6 +2182,92 @@ Covered by `iqtree::tests::iqtree_min_freq_folds_rare_states_into_missing`
from the written `_iqtree_states.csv` and `A`/`C` still present). Full from the written `_iqtree_states.csv` and `A`/`C` still present). Full
workspace `cargo test` green. workspace `cargo test` green.
## `--distance` unification: SNP corrections as first-class metrics (discussion, 2026-08-28)
**Decided, not yet implemented.** `--metric` (renamed `--distance` — several of
its existing values, e.g. Bray-Curtis, aren't metrics in the strict sense,
`--metric` was a misnomer) gains a family of `snp-*` values computed from the
central-position SNP pipeline, routed internally to the sibling-annex
machinery (`PairwiseTally`, `obikphylo::siblings::algorithms::pairwise`)
instead of `cache.distance(...)`'s existing per-layer traversal — a different
code path behind the same CLI surface, not just another branch of one
formula function.
**Why unify at the CLI level despite the implementation split**: phylogenetically
a SNP-corrected distance is a distance like any other — NJ/UPGMA are agnostic
to how the matrix was produced, so exposing it as a special-cased subcommand
instead of a `--distance` value would misrepresent its role. The
implementation divergence (sibling-annex-based vs. plain index scan) is real
but belongs at the routing layer, invisible to the CLI's own vocabulary.
**`--subsample` becomes optional for `snp-*` distances** (it stays mandatory
for `--sankoff`/`--pseudo-alignment`, unrelated commands): absent means
exhaustive, achieved for free by reusing `sample_index`'s existing
proportional-per-layer-quota mechanism with `n` set to the index-wide total
non-monomorphic-minorant count (already available from the sibling-annex
stats) — every layer's quota then equals its own full count, giving Bernoulli
`p = 1` everywhere, i.e. every eligible family is drawn. No second,
exhaustive-only driver needed. Present means sampled, exactly as `--sankoff`
already behaves.
**One shared tally, many derived formulas.** `PairwiseTally`'s `subst[4][4]`
per-pair substitution counts (plus marginal base frequencies derived from it)
are the sufficient statistic for every closed-form correction below — each
is a small pure function `PairwiseTally -> Array2<f64>`, at the same level as
the already-implemented `raw_snp_distance`/`base_pair_tally`/
`cardinality_tally`. No new full scan per formula, whether the tally itself
was built exhaustively or from a subsample.
**`--raw-snp-counts` stays a separate, unrelated flag** — same underlying
tally, but a diagnostic (`n_snp`/`n_shared`/`n_eligible` per genome pair, one
row per pair) rather than a distance value, and its long-table shape doesn't
fold into a single N×N matrix the way a distance does. No change to its
existing CSV format.
### `snp-*` distance catalog
All closed-form (method-of-moments / direct formula), none requiring
per-pair or per-tree maximum-likelihood fitting — that excludes HKY85's
*tree*-ML usage but not its *pairwise* estimator, which is closed-form like
F84/TN93 and is included below. `snp-` prefix on every CLI value.
| value | corrects for | inputs beyond raw counts |
|---|---|---|
| `snp-raw` | nothing (uncorrected p-distance) | — |
| `snp-jc` (Jukes-Cantor, JC69) | multiple substitutions per site | — |
| `snp-k2p` (Kimura 2-parameter, K80) | + transition/transversion rate bias | ts/tv split |
| `snp-k81` (Kimura 3-parameter, K3ST) | + splits transversions into 2 categories | ts/tv split, by category |
| `snp-f81` (Felsenstein 81) | + unequal base frequencies (no ts/tv split) | empirical base freqs |
| `snp-tajima-nei` (Tajima-Nei 1984) | same goal as F81 (equal-input model), different formula, better small-sample behavior | empirical base freqs |
| `snp-t92` (Tamura 3-parameter) | K2P + GC-content bias | ts/tv split, GC content |
| `snp-f84` (Felsenstein 84) | full empirical base freqs + single ts/tv rate | empirical base freqs, ts/tv split |
| `snp-hky85` (Hasegawa-Kishino-Yano, pairwise estimator) | same inputs as F84, different formula | empirical base freqs, ts/tv split |
| `snp-tn93` (Tamura-Nei) | full empirical base freqs + separate purine/pyrimidine transition rates + transversion rate | empirical base freqs, purine-ts/pyrimidine-ts/tv split |
| `snp-logdet` (LogDet / paralinear) | no shared-model or stationarity assumption at all — general divergence-matrix determinant | full empirical 4×4 divergence matrix (already `subst[4][4]`) |
| `snp-tv` (transversions-only p-distance) | diagnostic/deep-divergence variant — drops transitions entirely (they saturate first) | tv-only counts |
**`` rate-heterogeneity modifier, applicable to `snp-jc`, `snp-k2p`,
`snp-k81`, `snp-t92`, `snp-f84`, `snp-hky85`, `snp-tn93`** (not `snp-raw`,
nothing to correct; not `snp-logdet`, no standard gamma formulation) — same
formula as the base correction, weighted by a shape parameter `α` supplied
by the user (`--gamma-shape <alpha>`), not estimated by ML. A modifier on
existing values, not a separate enum arm per distance.
### Output format: PHYLIP-relaxed by default for the distance matrix
**Decided, not yet implemented.** The primary distance-matrix output
(`_dist.csv` today) gains multiple formats: **PHYLIP-relaxed becomes the
default** (widely read by external NJ tools — PHYLIP `neighbor`, FastME,
T-REX, SplitsTree — relaxed rather than strict to avoid the 10-character
label truncation, since genome labels here routinely exceed it), a `--csv`
flag opts back into the current CSV format, PHYLIP-strict is a possible
future addition (not now). This changes the *default* output of every
existing `--distance` value (jaccard, hamming, bray-curtis, ...), not just
the new `snp-*` ones — accepted explicitly (pre-release, single developer
user, no external consumers to break). Scoped to the distance matrix only:
`--shared-kmers` and `--raw-snp-counts` are counts, not distances, and keep
their existing CSV-only format.
## References ## References
The Mash mutation-rate model this discussion contrasts with: The Mash mutation-rate model this discussion contrasts with:
+68 -5
View File
@@ -40,11 +40,12 @@ impl From<MetricArg> for DistanceMetric {
/// annex diagnostics (`--sibling-stats`, `--sibling-hist`), entropy /// annex diagnostics (`--sibling-stats`, `--sibling-hist`), entropy
/// reporting (`--shannon`), SNP pseudo-alignment sampling /// reporting (`--shannon`), SNP pseudo-alignment sampling
/// (`--pseudo-alignment`, `--subsample`, `--free-loss`, `--no-ambiguity`, /// (`--pseudo-alignment`, `--subsample`, `--free-loss`, `--no-ambiguity`,
/// `--entropy`/`--entropy-sd`) and Sankoff cost-matrix calibration /// `--entropy`/`--entropy-sd`), Sankoff cost-matrix calibration
/// (`--sankoff`, `--sankoff-ratio-ceiling`) — everything else sibling-annex-based /// (`--sankoff`, `--sankoff-ratio-ceiling`) and its TNT/PhyG/IQ-TREE exports
/// (`--tnt`/`--phyg`/`--iqtree`, raw SNP distance, family overlap, ...) /// (`--tnt`, `--phyg`, `--iqtree`/`--iqtree-min-freq`,
/// stays in `obikmer` until the rest of `obikphylo::siblings` is /// `--sankoff-cost-scale`) — everything else sibling-annex-based (raw SNP
/// reconnected (see the project memory on this). /// distance, family overlap, ...) stays in `obikmer` until the rest of
/// `obikphylo::siblings` is reconnected (see the project memory on this).
#[derive(Args)] #[derive(Args)]
pub struct PhyloArgs { pub struct PhyloArgs {
/// Index directory /// Index directory
@@ -142,6 +143,68 @@ pub struct PhyloArgs {
#[arg(long, default_value = "0.5")] #[arg(long, default_value = "0.5")]
pub sankoff_ratio_ceiling: f64, pub sankoff_ratio_ceiling: f64,
/// Also write <prefix>_sankoff.tnt, a ready-to-run TNT script (`proc
/// <file>;`) for the same matrix/alignment `--sankoff` computes —
/// recoded to TNT's default xread alphabet (0-9A-F only; TNT rejects
/// the wider IUPAC set `--sankoff`'s own output uses unless `nstates
/// dna` is set, which imposes TNT's own incompatible DNA encoding
/// instead) with integer-scaled costs (TNT's smatrix/cost commands
/// reject decimals). Implies `--sankoff`.
#[arg(long)]
pub tnt: bool,
/// Also write <prefix>_sankoff.tcm and <prefix>_sankoff.pg, a
/// custom-alphabet cost matrix and a ready-to-run PhyG script (`read`/
/// `search`/`report`) for the same matrix/alignment `--sankoff`
/// computes. Reuses `--sankoff`'s own `_sankoff.fasta` directly — PhyG's
/// `tcm:` alphabet is read from the matrix file itself, so the IUPAC+`0`
/// alphabet needs no recoding here, unlike `--tnt`. Implies `--sankoff`.
#[arg(long)]
pub phyg: bool,
/// Also write <prefix>_iqtree.model and <prefix>_iqtree.fasta, a
/// custom-model file and a matching recoded alignment for genuine
/// maximum-likelihood inference with IQ-TREE (`iqtree3 -s ...
/// --seqtype MORPH -m ...+ASC`) — real branch lengths, unlike
/// `--tnt`/`--phyg`'s parsimony step counts. The model is the
/// reversible `Q(i,j) = R(i,j)·π_j` construction: `R` (exchangeability,
/// symmetric) recovered from the same calibrated cost matrix
/// `--sankoff` computes, `π` the real empirical state frequencies
/// counted from the alignment. Only the states that actually occur in
/// this alignment are kept, compactly renumbered (IQ-TREE infers its
/// state count from the alignment itself, and a gap in the numbering
/// would silently misalign the model file). Implies `--sankoff`.
#[arg(long)]
pub iqtree: bool,
/// Under `--iqtree --free-loss`, also recode to `?` (the same
/// missing-data treatment as `-`) any state whose empirical frequency
/// in the alignment falls below this threshold — not just genuinely
/// absent calls. States encoding 3 or 4 simultaneously-observed central
/// bases (IUPAC `V`/`H`/`K`.../`N` for 3, `N` for 4) are rare by
/// construction and often land in exactly this low-frequency range —
/// more likely assembly/detection noise than a genuine, widely-preserved
/// multi-way polymorphism, the same "sampling failure, not true signal"
/// reasoning `--free-loss` already applies to absence. No effect
/// without `--free-loss` (there is no missing-data symbol to recode to
/// otherwise). `<prefix>_iqtree_states.csv` reports the frequency
/// actually used to decide.
#[arg(long, default_value = "0.001")]
pub iqtree_min_freq: f64,
/// Scale factor applied before rounding real-valued costs to the
/// integers both `--tnt`'s smatrix/cost commands and `--phyg`'s `tcm:`
/// matrix require. Keep this small: the total tree score is this scale
/// times the sum of per-character costs across every character, and
/// there are hints in TNT's own manual that at least some of its
/// internal accumulators are 32-bit — a large scale risks a silent
/// integer overflow (undetectable, not just a crash) far more costly
/// than the resolution a bigger factor would buy. Shared between `--tnt`
/// and `--phyg` rather than split into two flags: both scale the same
/// calibrated matrix for the same reason (integer-only cost commands).
#[arg(long, default_value = "100")]
pub sankoff_cost_scale: f64,
/// Distance metric to compute /// Distance metric to compute
#[arg(long, value_enum, default_value = "jaccard")] #[arg(long, value_enum, default_value = "jaccard")]
pub metric: MetricArg, pub metric: MetricArg,
+502
View File
@@ -0,0 +1,502 @@
use std::io::{BufWriter, Write};
use std::path::PathBuf;
use obifastwrite::{JsonVal, write_record};
use obikphylo::siblings::SnpAlignment;
use tracing::info;
use super::sankoff::{STATE_SYMBOL, state_index_table};
// ── Sankoff-calibrated data → IQ-TREE custom ML model + recoded alignment ──
//
// Not itself a Sankoff computation — IQ-TREE does maximum likelihood, not
// parsimony. Only the *source data* is shared with `--tnt`/`--phyg` (the
// calibrated 16-state cost matrix, the pseudo-alignment); the operation
// performed on it here is different, hence no "sankoff" in these names,
// unlike `tnt::write_sankoff_tnt`/`phyg::write_sankoff_phyg`.
//
// The mechanism: pass a **file path** directly as `-m`, containing (as
// whitespace/newline-separated numbers) the lower-triangular exchangeability
// matrix `R` (`k(k-1)/2` values, PAML row-major order) immediately followed
// by the `k` state frequencies `π` on the same stream —
// `ModelMarkov::readRates`/`readStateFreq` read them in that exact order, no
// header, no separator required.
//
// `R` is recovered from the calibrated Sankoff cost matrix via
// `R(a,b) = exp(-cost(a,b))` (the cost is `-ln(rate)`), symmetric by
// construction (the underlying tally never captured direction). `π` is the
// real, empirical, non-uniform marginal frequency of each state across the
// whole alignment. IQ-TREE reconstructs the (generally asymmetric) rate
// matrix internally as `Q(i,j) = R(i,j)·π_j` — reversible for *any* `π`, not
// just uniform, because `R` is symmetric.
//
// IQ-TREE infers its state count from the highest-ordinal symbol actually
// present in the alignment, not from a declared count. So states that never
// occur anywhere in this particular alignment are dropped, and the
// survivors are renumbered compactly (`0..k-1`, order preserved) rather than
// leaving gaps that would silently misalign every value IQ-TREE reads. Both
// the model and the alignment must agree on this same renumbering, so it's
// computed once (`CompactAlphabet`) and shared between them.
const IQTREE_STATE_SYMBOL: [char; 16] = [
'0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F',
];
struct CompactAlphabet {
/// Canonical (0..16) state index -> compact index, for states that occur.
old_to_compact: [Option<u8>; 16],
/// Compact index -> canonical state index, order-preserving.
compact_to_old: Vec<u8>,
/// Empirical frequency of each compact-indexed state (sums to 1).
freq: Vec<f64>,
}
impl CompactAlphabet {
fn k(&self) -> usize {
self.compact_to_old.len()
}
}
/// Under `--free-loss`, non-detection (`-`) becomes IQ-TREE's own missing
/// symbol (`?`) — ignored when IQ-TREE checks a site's constancy for
/// `+ASC`. A family kept as "variable" by the sampling (`family_size() >=
/// 2`, a whole-annex property, oblivious to any one column's actual calls)
/// can still turn constant *among the genomes that actually have data* once
/// the non-detected ones are excluded from that check — the same failure
/// mode `--exclude-genome` already had to account for, just triggered by
/// hiding cells instead of dropping whole rows. Same remedy: rescan columns
/// treating `-` as ignored, drop any where the remaining calls agree on a
/// single state. Parsimony (`--tnt`/`--phyg`) has no no-invariant-site
/// requirement, so this only runs on IQ-TREE's own copy of the alignment,
/// never mutating the one the caller also hands to those two exports.
fn drop_ascertainment_noninformative(alignment: &SnpAlignment) -> SnpAlignment {
let n_sites = alignment.sequences.first().map(|s| s.len()).unwrap_or(0);
let keep: Vec<bool> = (0..n_sites)
.map(|site| {
let mut first: Option<u8> = None;
for seq in &alignment.sequences {
let b = seq[site];
if b == b'-' {
continue;
}
match first {
None => first = Some(b),
Some(f) if f != b => return true,
_ => {}
}
}
false // all calls missing, or all calls agree — non-informative
})
.collect();
let sequences = alignment
.sequences
.iter()
.map(|seq| {
seq.iter()
.zip(keep.iter())
.filter(|&(_, &k)| k)
.map(|(&b, _)| b)
.collect()
})
.collect();
SnpAlignment { sequences, genome_indices: alignment.genome_indices.clone() }
}
/// Recode every occurrence of a byte in `symbols` to `-` — the same
/// "absent" byte `drop_ascertainment_noninformative`/`compact_alphabet`
/// already treat specially under `--free-loss` (recoded to `?` further
/// downstream). Used by `--iqtree-min-freq` to fold rare, likely-noisy
/// states into the missing-data treatment before a second
/// `compact_alphabet` pass, without duplicating that treatment's logic.
fn recode_symbols_as_absent(alignment: &SnpAlignment, symbols: &[u8]) -> SnpAlignment {
let sequences = alignment
.sequences
.iter()
.map(|seq| {
seq.iter()
.map(|&b| if symbols.contains(&b) { b'-' } else { b })
.collect()
})
.collect();
SnpAlignment { sequences, genome_indices: alignment.genome_indices.clone() }
}
fn compact_alphabet(alignment: &SnpAlignment, free_loss: bool) -> CompactAlphabet {
let iupac_to_state = state_index_table();
let mut occurs = [false; 16];
let mut counts = [0u64; 16];
for seq in &alignment.sequences {
for &b in seq {
if free_loss && b == b'-' {
// `?`: IQ-TREE's own missing-data symbol for `--seqtype
// MORPH`, marginalised by Felsenstein pruning — not a
// numbered state, so excluded from `occurs`/`counts` and
// from the compact alphabet built below.
continue;
}
let b = if b == b'-' { b'0' } else { b };
let state = iupac_to_state[b as usize] as usize;
occurs[state] = true;
counts[state] += 1;
}
}
let mut old_to_compact: [Option<u8>; 16] = [None; 16];
let mut compact_to_old: Vec<u8> = Vec::new();
for old in 0..16 {
if occurs[old] {
old_to_compact[old] = Some(compact_to_old.len() as u8);
compact_to_old.push(old as u8);
}
}
let total: u64 = compact_to_old.iter().map(|&old| counts[old as usize]).sum();
let freq: Vec<f64> = compact_to_old
.iter()
.map(|&old| counts[old as usize] as f64 / total as f64)
.collect();
CompactAlphabet {
old_to_compact,
compact_to_old,
freq,
}
}
/// Write `<prefix>_iqtree_states.csv`: the mapping from IQ-TREE's own
/// compact state symbols (`0..9A-F`, what actually appears in
/// `_iqtree.fasta`/`_iqtree.model`) back to the canonical 16-state
/// alphabet (`STATE_SYMBOL` — the same one `_sankoff_matrix.csv` is
/// indexed by), plus each state's empirical frequency at full precision
/// (`_iqtree.model`'s own frequency line is truncated to 6 decimals).
/// Without this file, a compact index in `_iqtree.model`'s `R`/`π` output
/// (e.g. "state 0 has zero exchangeability with everything else") can't be
/// traced back to which real state that is.
fn write_iqtree_states_csv(alphabet: &CompactAlphabet, output: &Option<PathBuf>) -> String {
let path = output
.as_ref()
.map(|p| format!("{}_iqtree_states.csv", p.display()))
.unwrap_or_else(|| "iqtree_states.csv".into());
let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
eprintln!("error creating {path}: {e}");
std::process::exit(1);
}));
writeln!(f, "iqtree_symbol,canonical_symbol,frequency").unwrap();
for (compact, &old) in alphabet.compact_to_old.iter().enumerate() {
writeln!(
f,
"{},{},{}",
IQTREE_STATE_SYMBOL[compact], STATE_SYMBOL[old as usize], alphabet.freq[compact]
)
.unwrap();
}
path
}
/// Write the `R` (exchangeability) + `π` (frequencies) model file IQ-TREE's
/// `-m <file>+ASC` reads. Returns the path, so the caller can print a
/// single combined "how to run this" message once the alignment is also
/// written. The one bit of real computation this whole adapter does:
/// `R(a,b) = exp(-cost(a,b))`, recovering the exchangeability rate a
/// calibrated Sankoff parsimony cost implies for a continuous-time model —
/// a one-line inversion of the cost matrix's own `-ln(rate)` construction,
/// not a new estimate.
fn write_iqtree_model(
matrix: &[[f64; 16]; 16],
alphabet: &CompactAlphabet,
output: &Option<PathBuf>,
) -> String {
let rate = |old_i: u8, old_j: u8| (-matrix[old_i as usize][old_j as usize]).exp();
let model_path = output
.as_ref()
.map(|p| format!("{}_iqtree.model", p.display()))
.unwrap_or_else(|| "iqtree.model".into());
let mut f = BufWriter::new(std::fs::File::create(&model_path).unwrap_or_else(|e| {
eprintln!("error creating {model_path}: {e}");
std::process::exit(1);
}));
for i in 1..alphabet.k() {
let row: Vec<String> = (0..i)
.map(|j| {
format!(
"{:.6}",
rate(alphabet.compact_to_old[i], alphabet.compact_to_old[j])
)
})
.collect();
writeln!(f, "{}", row.join(" ")).unwrap();
}
writeln!(
f,
"{}",
alphabet
.freq
.iter()
.map(|p| format!("{p:.6}"))
.collect::<Vec<_>>()
.join(" ")
)
.unwrap();
info!(
"IQ-TREE model file → {model_path} ({} of 16 states present in the alignment)",
alphabet.k()
);
model_path
}
/// Write the pseudo-alignment recoded to the same compact `0..k-1` alphabet
/// as `write_iqtree_model`'s matrix — not `--sankoff`'s own IUPAC alphabet,
/// since IQ-TREE needs the symbol ordinal itself to match the surviving
/// state count (see this module's own doc comment on state inference).
fn write_iqtree_alignment(
alignment: &SnpAlignment,
labels: &[String],
alphabet: &CompactAlphabet,
output: &Option<PathBuf>,
free_loss: bool,
) -> (String, usize) {
let iupac_to_state = state_index_table();
let fasta_path = output
.as_ref()
.map(|p| format!("{}_iqtree.fasta", p.display()))
.unwrap_or_else(|| "iqtree.fasta".into());
let mut f = BufWriter::new(std::fs::File::create(&fasta_path).unwrap_or_else(|e| {
eprintln!("error creating {fasta_path}: {e}");
std::process::exit(1);
}));
let n_sites = alignment.sequences.first().map(|s| s.len()).unwrap_or(0);
for (&g, seq) in alignment.genome_indices.iter().zip(alignment.sequences.iter()) {
let recoded: Vec<u8> = seq
.iter()
.map(|&b| {
if free_loss && b == b'-' {
return b'?';
}
let b = if b == b'-' { b'0' } else { b };
let old = iupac_to_state[b as usize] as usize;
let compact = alphabet.old_to_compact[old]
.expect("state occurs in the alignment, so it must have a compact index");
IQTREE_STATE_SYMBOL[compact as usize] as u8
})
.collect();
write_record(
&recoded,
&labels[g],
&[("n_sites", JsonVal::Num(n_sites as u64))],
&mut f,
)
.unwrap_or_else(|e| {
eprintln!("error writing {fasta_path}: {e}");
std::process::exit(1);
});
}
(fasta_path, n_sites)
}
pub(super) fn write_iqtree(
matrix: &[[f64; 16]; 16],
alignment: &SnpAlignment,
labels: &[String],
output: &Option<PathBuf>,
free_loss: bool,
min_freq: f64,
) {
let filtered;
let alignment = if free_loss {
let before = alignment.sequences.first().map(|s| s.len()).unwrap_or(0);
filtered = drop_ascertainment_noninformative(alignment);
let after = filtered.sequences.first().map(|s| s.len()).unwrap_or(0);
if after != before {
info!(
"--free-loss: {before} → {after} sites (dropped columns non-informative once `-` \
is treated as missing — required for +ASC)"
);
}
&filtered
} else {
alignment
};
let mut alphabet = compact_alphabet(alignment, free_loss);
// `--iqtree-min-freq`: fold rare (likely-noisy) states into the same
// missing-data treatment `-` already gets under `--free-loss`, then
// recompute the alphabet on the further-filtered alignment.
let refiltered;
let alignment = if free_loss {
let low_freq_symbols: Vec<u8> = alphabet
.compact_to_old
.iter()
.zip(alphabet.freq.iter())
.filter(|&(_, &f)| f < min_freq)
.map(|(&old, _)| STATE_SYMBOL[old as usize] as u8)
.collect();
if low_freq_symbols.is_empty() {
alignment
} else {
let recoded = recode_symbols_as_absent(alignment, &low_freq_symbols);
let before = recoded.sequences.first().map(|s| s.len()).unwrap_or(0);
refiltered = drop_ascertainment_noninformative(&recoded);
let after = refiltered.sequences.first().map(|s| s.len()).unwrap_or(0);
info!(
"--iqtree-min-freq {min_freq}: {} rare state(s) ({}) recoded as missing, {before} → {after} sites",
low_freq_symbols.len(),
low_freq_symbols
.iter()
.map(|&b| b as char)
.collect::<String>(),
);
alphabet = compact_alphabet(&refiltered, free_loss);
&refiltered
}
} else {
alignment
};
let states_path = write_iqtree_states_csv(&alphabet, output);
let model_path = write_iqtree_model(matrix, &alphabet, output);
let (fasta_path, n_sites) =
write_iqtree_alignment(alignment, labels, &alphabet, output, free_loss);
let prefix_name = output
.as_ref()
.and_then(|p| p.file_name())
.map(|n| format!("{}_iqtree", n.to_string_lossy()))
.unwrap_or_else(|| "iqtree".into());
info!(
"IQ-TREE alignment → {fasta_path} ({n_sites} sites, {} states)\n\
IQ-TREE state mapping → {states_path}\n\
Run with:\n \
iqtree3 -s {fasta_path} --seqtype MORPH -m {model_path}+ASC --prefix {prefix_name} -T AUTO\n\
\n\
options -alrt 1000 -B 1000 can be added to evaluate robustness of the tree",
alphabet.k()
);
}
#[cfg(test)]
mod tests {
use super::*;
fn alignment(sequences: Vec<Vec<u8>>) -> SnpAlignment {
let genome_indices = (0..sequences.len()).collect();
SnpAlignment { sequences, genome_indices }
}
#[test]
fn free_loss_excludes_absent_state_and_freq_sums_to_one() {
// 3 genomes, 2 sites. Site 0: g1='A', g2='C', g3='-' (absent).
// Site 1: g1='-', g2='-', g3='G'. Under free_loss, every '-' must
// be excluded from the frequency count entirely (not folded into
// state 0).
let alignment = alignment(vec![vec![b'A', b'-'], vec![b'C', b'-'], vec![b'-', b'G']]);
let alphabet = compact_alphabet(&alignment, true);
assert!(
!alphabet.compact_to_old.contains(&0),
"state 0 (absent) must not appear in the compact alphabet under --free-loss, got {:?}",
alphabet.compact_to_old
);
let sum: f64 = alphabet.freq.iter().sum();
assert!(
(sum - 1.0).abs() < 1e-9,
"frequencies must sum to 1, got {sum} ({:?})",
alphabet.freq
);
assert_eq!(alphabet.k(), 3, "A, C, G — 3 real states, `-` excluded");
}
#[test]
fn without_free_loss_absent_state_is_counted_normally() {
let alignment = alignment(vec![vec![b'A', b'-'], vec![b'C', b'-'], vec![b'-', b'G']]);
let alphabet = compact_alphabet(&alignment, false);
assert!(
alphabet.compact_to_old.contains(&0),
"state 0 (absent, recoded from '-') must be counted when --free-loss is off"
);
let sum: f64 = alphabet.freq.iter().sum();
assert!(
(sum - 1.0).abs() < 1e-9,
"frequencies must sum to 1, got {sum} ({:?})",
alphabet.freq
);
}
#[test]
fn states_csv_maps_compact_symbols_back_to_canonical_ones() {
// 'A' (state 1) and 'G' (state 4) occur, '-' (state 0) excluded by
// --free-loss — compact index 0 -> 'A', compact index 1 -> 'G'.
let alignment = alignment(vec![vec![b'A', b'-'], vec![b'-', b'G']]);
let alphabet = compact_alphabet(&alignment, true);
let output = Some(
std::env::temp_dir().join(format!("obikmer2_test_iqtree_states_{}", std::process::id())),
);
let path = write_iqtree_states_csv(&alphabet, &output);
let csv = std::fs::read_to_string(&path).unwrap();
std::fs::remove_file(&path).ok();
let mut lines = csv.lines();
assert_eq!(
lines.next(),
Some("iqtree_symbol,canonical_symbol,frequency")
);
assert_eq!(lines.next(), Some("0,A,0.5"));
assert_eq!(lines.next(), Some("1,G,0.5"));
assert!(lines.next().is_none());
}
#[test]
fn iqtree_min_freq_folds_rare_states_into_missing() {
// 20 common A/C sites (60 calls total across 3 genomes) plus one
// site where genome 0 carries the rare ambiguity state `M` and
// genome 1 carries `A` (kept informative by the first
// ascertainment filter: two distinct non-`-` calls) — `M` ends up
// at 1/62, well below the 0.05 threshold used here.
let mut sequences: Vec<Vec<u8>> = vec![Vec::new(); 3];
for i in 0..20 {
let (a, b, c) = if i % 2 == 0 {
(b'A', b'C', b'A')
} else {
(b'C', b'A', b'C')
};
sequences[0].push(a);
sequences[1].push(b);
sequences[2].push(c);
}
sequences[0].push(b'M');
sequences[1].push(b'A');
sequences[2].push(b'-');
let alignment = alignment(sequences);
let labels = vec!["g1".to_string(), "g2".to_string(), "g3".to_string()];
let matrix = [[0.0f64; 16]; 16];
let prefix = std::env::temp_dir().join(format!(
"obikmer2_test_iqtree_minfreq_{}",
std::process::id()
));
let output = Some(prefix.clone());
write_iqtree(&matrix, &alignment, &labels, &output, true, 0.05);
let states_path = format!("{}_iqtree_states.csv", prefix.display());
let csv = std::fs::read_to_string(&states_path).unwrap();
assert!(
!csv.contains(",M,"),
"M (freq ~1/62) must be folded into missing under --iqtree-min-freq 0.05, got:\n{csv}"
);
assert!(
csv.contains(",A,") && csv.contains(",C,"),
"A/C must survive (well above threshold), got:\n{csv}"
);
for suffix in ["_iqtree_states.csv", "_iqtree.model", "_iqtree.fasta"] {
std::fs::remove_file(format!("{}{suffix}", prefix.display())).ok();
}
}
}
+32 -2
View File
@@ -1,5 +1,8 @@
mod args; mod args;
mod iqtree;
mod phyg;
mod sankoff; mod sankoff;
mod tnt;
use std::io::{self, BufWriter, Write}; use std::io::{self, BufWriter, Write};
use std::sync::Arc; use std::sync::Arc;
@@ -14,7 +17,10 @@ use obikphylo::{Metrics, neighbor_joining, upgma};
use obisys::{Reporter, Stage}; use obisys::{Reporter, Stage};
use tracing::info; use tracing::info;
use iqtree::write_iqtree;
use phyg::write_sankoff_phyg;
use sankoff::{write_sankoff_alignment_fasta, write_sankoff_matrix_csv, write_sankoff_params}; use sankoff::{write_sankoff_alignment_fasta, write_sankoff_matrix_csv, write_sankoff_params};
use tnt::write_sankoff_tnt;
pub use args::PhyloArgs; pub use args::PhyloArgs;
@@ -201,8 +207,8 @@ pub fn run(args: PhyloArgs) {
info!("pseudo-alignment ({n_sites} site(s), {} genome(s)) → {path}", alignment.genome_indices.len()); info!("pseudo-alignment ({n_sites} site(s), {} genome(s)) → {path}", alignment.genome_indices.len());
} }
// ── Sankoff cost-matrix calibration (`--sankoff`) ─────────────────────────── // ── Sankoff cost-matrix calibration (`--sankoff`, `--tnt`, `--phyg`, `--iqtree`) ──
if args.sankoff { if args.sankoff || args.tnt || args.phyg || args.iqtree {
let Some(subsample_n) = args.subsample else { let Some(subsample_n) = args.subsample else {
eprintln!("error: --sankoff requires --subsample <N>"); eprintln!("error: --sankoff requires --subsample <N>");
std::process::exit(1); std::process::exit(1);
@@ -239,6 +245,30 @@ pub fn run(args: PhyloArgs) {
&args.output, &args.output,
); );
write_sankoff_alignment_fasta(&bundle.alignment, &labels, &args.output, args.free_loss); write_sankoff_alignment_fasta(&bundle.alignment, &labels, &args.output, args.free_loss);
if args.tnt {
write_sankoff_tnt(
&matrix,
&bundle.alignment,
&labels,
&args.output,
args.sankoff_cost_scale,
args.free_loss,
);
}
if args.phyg {
write_sankoff_phyg(&matrix, &args.output, args.sankoff_cost_scale);
}
if args.iqtree {
write_iqtree(
&matrix,
&bundle.alignment,
&labels,
&args.output,
args.free_loss,
args.iqtree_min_freq,
);
}
} }
info!("computing {:?} distances for {} genome(s)", args.metric, n); info!("computing {:?} distances for {} genome(s)", args.metric, n);
+80
View File
@@ -0,0 +1,80 @@
use std::io::{BufWriter, Write};
use std::path::PathBuf;
use tracing::info;
use super::sankoff::{STATE_SYMBOL, scaled_metric_matrix};
// ── Sankoff cost matrix → PhyG custom-alphabet TCM + ready-to-run script ────
//
// PhyG's `tcm:STRING` format needs no alphabet recoding, unlike `--tnt`:
// its parser reads the alphabet straight from the tcm file's own first
// line, so `--sankoff`'s own `_sankoff.fasta` (already IUPAC+`0`) is reused
// as-is via `prefasta:`. PhyG auto-adds its own indel/gap state as an
// (n+1)-th row/column of the tcm — inert here since the alignment already
// encodes absence as an ordinary state (`0`), never as `-` (see
// `sankoff::write_sankoff_alignment_fasta`'s own comment on why). The gap
// row/column below reuses `matrix[i][0]`/`matrix[0][j]` (cost to/from `∅`)
// as the closest principled value for a state that, in practice, is never
// actually triggered.
pub(super) fn write_sankoff_phyg(matrix: &[[f64; 16]; 16], output: &Option<PathBuf>, cost_scale: f64) {
let scaled_matrix = scaled_metric_matrix(matrix, cost_scale);
let basename = |suffix: &str| -> String {
output.as_ref()
.and_then(|p| p.file_name())
.map(|n| format!("{}{suffix}", n.to_string_lossy()))
.unwrap_or_else(|| format!("sankoff{suffix}"))
};
let full_path = |suffix: &str| -> String {
output.as_ref()
.map(|p| format!("{}{suffix}", p.display()))
.unwrap_or_else(|| format!("sankoff{suffix}"))
};
let tcm_path = full_path("_sankoff.tcm");
let mut f = BufWriter::new(std::fs::File::create(&tcm_path).unwrap_or_else(|e| {
eprintln!("error creating {tcm_path}: {e}");
std::process::exit(1);
}));
let alphabet_line = STATE_SYMBOL.iter().map(|c| c.to_string()).collect::<Vec<_>>().join(" ");
writeln!(f, "{alphabet_line}").unwrap();
for i in 0..16 {
let mut row: Vec<i64> = (0..16).map(|j| scaled_matrix[i][j]).collect();
row.push(scaled_matrix[i][0]); // gap column: same cost as to/from ∅
writeln!(f, "{}", row.iter().map(|v| v.to_string()).collect::<Vec<_>>().join(" ")).unwrap();
}
let mut gap_row: Vec<i64> = (0..16).map(|j| scaled_matrix[0][j]).collect();
gap_row.push(0);
writeln!(f, "{}", gap_row.iter().map(|v| v.to_string()).collect::<Vec<_>>().join(" ")).unwrap();
info!("PhyG TCM → {tcm_path}");
let pg_path = full_path("_sankoff.pg");
let mut f = BufWriter::new(std::fs::File::create(&pg_path).unwrap_or_else(|e| {
eprintln!("error creating {pg_path}: {e}");
std::process::exit(1);
}));
let fasta_name = basename("_sankoff.fasta");
let tcm_name = basename("_sankoff.tcm");
let tre_name = basename("_sankoff.tre");
writeln!(f, "read(prefasta:\"{fasta_name}\", tcm:\"{tcm_name}\")").unwrap();
writeln!(f, "search(seconds:300, instances:4)").unwrap();
writeln!(f, "report(\"{tre_name}\", graphs, newick, overwrite)").unwrap();
let pg_dir = std::path::Path::new(&pg_path).parent()
.filter(|d| !d.as_os_str().is_empty())
.map(|d| d.display().to_string())
.unwrap_or_else(|| ".".into());
let pg_name = std::path::Path::new(&pg_path).file_name()
.map(|n| n.to_string_lossy().into_owned())
.unwrap_or_else(|| pg_path.clone());
info!(
"PhyG script → {pg_path} (costs scaled x{cost_scale:.0}, runs a default 300s/4-instance \
search and writes trees to {tre_name})\n\
Run it with:\n \
cd {pg_dir} && phyg {pg_name}\n\
(`phyg` must run from that directory — `read()`/`report()` in the script use relative \
file names)"
);
}
+47
View File
@@ -66,6 +66,17 @@ pub(super) const STATE_SYMBOL: [char; 16] = [
'0', 'A', 'C', 'M', 'G', 'R', 'S', 'V', 'T', 'W', 'Y', 'H', 'K', 'D', 'B', 'N', '0', 'A', 'C', 'M', 'G', 'R', 'S', 'V', 'T', 'W', 'Y', 'H', 'K', 'D', 'B', 'N',
]; ];
/// `STATE_SYMBOL` byte -> state index (0..16), for adapters that need to
/// translate an alignment written in this alphabet into their own. Shared
/// rather than rebuilt per adapter (`tnt`, `iqtree`).
pub(super) fn state_index_table() -> [u8; 128] {
let mut table = [0u8; 128];
for (state, &sym) in STATE_SYMBOL.iter().enumerate() {
table[sym as usize] = state as u8;
}
table
}
pub(super) fn write_sankoff_matrix_csv(matrix: &[[f64; 16]; 16], output: &Option<PathBuf>) { pub(super) fn write_sankoff_matrix_csv(matrix: &[[f64; 16]; 16], output: &Option<PathBuf>) {
let path = output.as_ref() let path = output.as_ref()
.map(|p| format!("{}_sankoff_matrix.csv", p.display())) .map(|p| format!("{}_sankoff_matrix.csv", p.display()))
@@ -166,3 +177,39 @@ pub(super) fn write_sankoff_params(
}); });
info!("Sankoff calibration parameters → {path}"); info!("Sankoff calibration parameters → {path}");
} }
/// Scale `matrix` by `cost_scale` and round to integers (TNT's smatrix/cost
/// and PhyG's `tcm:` commands both reject decimals), then take the *metric
/// closure* of the result (Floyd-Warshall over the 16 states again, on the
/// now-integer values).
///
/// `pairwise_cost_matrix`'s row-normalise-then-`-ln` construction gives no
/// guarantee of being a metric (unlike a cost graph closed by shortest path
/// by construction) — so this closure isn't only needed to correct
/// integer-rounding artifacts (two real costs of `1.734` each round to
/// `173`, summing to `346`, while their own real sum `3.468` rounds to
/// `347` — TNT then reports "triangle inequality violated ... Fixed" and
/// silently substitutes its own corrected value), it may also be the only
/// thing making the *real-valued* matrix a metric in the first place.
/// Re-closing after rounding makes both corrections explicit and
/// reproducible here instead, rather than left implicit and
/// tool-version-dependent.
pub(super) fn scaled_metric_matrix(matrix: &[[f64; 16]; 16], cost_scale: f64) -> [[i64; 16]; 16] {
let mut m = [[0i64; 16]; 16];
for i in 0..16 {
for j in 0..16 {
m[i][j] = (matrix[i][j] * cost_scale).round() as i64;
}
}
for k in 0..16 {
for i in 0..16 {
for j in 0..16 {
let via = m[i][k] + m[k][j];
if via < m[i][j] {
m[i][j] = via;
}
}
}
}
m
}
+184
View File
@@ -0,0 +1,184 @@
use std::io::{BufWriter, Write};
use std::path::PathBuf;
use obikphylo::siblings::SnpAlignment;
use tracing::info;
use super::sankoff::{scaled_metric_matrix, state_index_table};
// ── Sankoff cost matrix + alignment → ready-to-run TNT script ──────────────
//
// TNT's *default* xread reader only accepts its own 0-9A-F alphabet (see
// its manual: "up to 16 states are allowed by xread, using symbols 0-9 ...
// and A-F") — the wider IUPAC set `STATE_SYMBOL` uses is rejected as an
// "alien symbol" unless `nstates dna` is set, which imposes TNT's own fixed
// DNA encoding instead, incompatible with a custom smatrix. And TNT's
// `smatrix`/`cost` commands reject decimal costs ("found symbol . when
// reading transformation costs"). So: recode to TNT's alphabet and
// integer-scale the costs here, at this adapter's boundary, rather than
// degrading the project's own canonical (IUPAC, real-valued) output.
const TNT_STATE_SYMBOL: [char; 16] = [
'0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F',
];
pub(super) fn write_sankoff_tnt(
matrix: &[[f64; 16]; 16],
alignment: &SnpAlignment,
labels: &[String],
output: &Option<PathBuf>,
cost_scale: f64,
free_loss: bool,
) {
let path = output
.as_ref()
.map(|p| format!("{}_sankoff.tnt", p.display()))
.unwrap_or_else(|| "sankoff.tnt".into());
let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| {
eprintln!("error creating {path}: {e}");
std::process::exit(1);
}));
// IUPAC-ish symbol -> bitmask, to translate the alignment (which uses
// `STATE_SYMBOL`, `-` already normalised to `0` by `sankoff_bundle`
// callers) into TNT's alphabet without re-deriving state indices.
let iupac_to_state = state_index_table();
let n_sites = alignment.sequences.first().map(|s| s.len()).unwrap_or(0);
let kept_labels: Vec<&String> = alignment.genome_indices.iter().map(|&g| &labels[g]).collect();
writeln!(f, "xread").unwrap();
writeln!(f, "mxram 16000;").unwrap();
writeln!(f, "taxname =;").unwrap();
writeln!(f, "taxname +50;").unwrap();
writeln!(
f,
"'obikmer central-position SNP families, calibrated Sankoff 16-state encoding'"
)
.unwrap();
writeln!(f, "{n_sites} {}", kept_labels.len()).unwrap();
for (label, seq) in kept_labels.iter().zip(alignment.sequences.iter()) {
write!(f, "{label} ").unwrap();
for &b in seq {
if free_loss && b == b'-' {
// `?`: TNT's own missing-data symbol, read directly, not
// routed through `TNT_STATE_SYMBOL` (there is no state for
// it) — see `write_sankoff_alignment_fasta`'s doc comment.
write!(f, "?").unwrap();
continue;
}
let b = if b == b'-' { b'0' } else { b };
let state = iupac_to_state[b as usize];
write!(f, "{}", TNT_STATE_SYMBOL[state as usize]).unwrap();
}
writeln!(f).unwrap();
}
writeln!(f, ";\n").unwrap();
let scaled_matrix = scaled_metric_matrix(matrix, cost_scale);
writeln!(f, "smatrix =0 (family16)").unwrap();
for i in 0..16 {
for j in (i + 1)..16 {
writeln!(
f,
"{}/{} {}",
TNT_STATE_SYMBOL[i], TNT_STATE_SYMBOL[j], scaled_matrix[i][j]
)
.unwrap();
}
}
writeln!(f, ";\n").unwrap();
writeln!(f, "ccode ( 0.{} ;", n_sites - 1).unwrap();
writeln!(f, "smatrix +0 0.{} ;", n_sites - 1).unwrap();
writeln!(f).unwrap();
// Basename only (not the full `path`/`output` prefix): TNT's natural
// workflow is to `cd` into the output directory before `proc`-ing the
// script, and an absolute path here would break if that directory is
// later moved or copied elsewhere.
let tre_name = output
.as_ref()
.and_then(|p| p.file_name())
.map(|n| format!("{}_sankoff.tre", n.to_string_lossy()))
.unwrap_or_else(|| "sankoff.tre".into());
// TNT's plain command parser has no comment syntax of its own — `/* */`
// and `[ ]` are only recognised inside the (separately-enabled) macro
// scripting language, and fail with "No command!" here otherwise
// (verified against this file with the local TNT binary). `quote` is
// the closest working equivalent: it prints free text and does not
// otherwise affect parsing, so it doubles as an explanation of the
// defaults below when the script is run. `;` ends a `quote` block like
// any other TNT command, so the text itself must avoid semicolons.
writeln!(f, "quote").unwrap();
writeln!(
f,
"Default search below (edit or delete this block to run your own strategy):"
)
.unwrap();
writeln!(
f,
" hold N : size of TNT's tree buffer (how many equally-parsimonious"
)
.unwrap();
writeln!(
f,
" trees it keeps in memory at once), 20 is a small, fast"
)
.unwrap();
writeln!(
f,
" default, raise it if mult reports it had to drop trees."
)
.unwrap();
writeln!(
f,
" mult : traditional search (random addition sequences followed by"
)
.unwrap();
writeln!(
f,
" TBR branch-swapping, TNT's own default replication count),"
)
.unwrap();
writeln!(
f,
" a reasonable first-pass strategy on this data's memory"
)
.unwrap();
writeln!(
f,
" footprint, xmult's ratchet/drift/tree-fusion buffers ran"
)
.unwrap();
writeln!(
f,
" this out of RAM at TNT's default mxram on this dataset."
)
.unwrap();
writeln!(
f,
" export - F : write the trees held in the buffer to file F, in"
)
.unwrap();
writeln!(
f,
" TNT/Hennig86 format ('-' means trees, as opposed to data)."
)
.unwrap();
writeln!(f, ";").unwrap();
writeln!(f, "hold 20;").unwrap();
writeln!(f, "mult;").unwrap();
writeln!(f, "export - {tre_name};").unwrap();
info!(
"TNT script → {path} (costs scaled x{cost_scale:.0}, runs a default `hold 20; mult;` \
search and writes trees to {tre_name} in TNT's working directory — edit the trailing \
comment block in the script to change this)\n\
Run it with:\n \
printf 'proc {path};\\nquit;\\n' | tnt\n\
(or start `tnt` interactively and type `proc {path};`)"
);
}
+11 -5
View File
@@ -27,20 +27,21 @@ mod subsample;
use obikidxcache::index_cache::IndexCache; use obikidxcache::index_cache::IndexCache;
pub use cardcomp::{cardinality_transition_probs, composition_transition_probs, pairwise_cost_matrix};
pub use alignment::SnpAlignment; pub use alignment::SnpAlignment;
pub use cardcomp::{
cardinality_transition_probs, composition_transition_probs, pairwise_cost_matrix,
};
pub use pairwise::{BasePairTally, CardinalityTally, RawSnpDistanceOutput}; pub use pairwise::{BasePairTally, CardinalityTally, RawSnpDistanceOutput};
pub use sankoff::SankoffBundle; pub use sankoff::SankoffBundle;
pub use stats::SiblingAnnexStats; pub use stats::SiblingAnnexStats;
pub use subsample::EntropyBias; pub use subsample::{EntropyBias, SurvivingFamily};
pub(crate) use alignment::snp_pseudo_alignment; pub(crate) use alignment::snp_pseudo_alignment;
pub(crate) use annex::build_layer_sibling_annex; pub(crate) use annex::build_layer_sibling_annex;
pub(crate) use entropy::{ensure_layer_entropy_annex, family_entropy, family_entropy_4, iter_full_entropy}; pub(crate) use entropy::{ensure_layer_entropy_annex, family_entropy, family_entropy_4};
pub(crate) use family_scan::{Selection, scan_layer_families}; pub(crate) use family_scan::{Selection, scan_layer_families};
pub(crate) use sankoff::sankoff_bundle; pub(crate) use sankoff::sankoff_bundle;
pub(crate) use stats::{sibling_annex_stats, sibling_family_size_histogram}; pub(crate) use stats::{sibling_annex_stats, sibling_family_size_histogram};
pub(crate) use subsample::sample_index;
/// Whether every layer number in `cache` fits in a `FamilyMask` field /// Whether every layer number in `cache` fits in a `FamilyMask` field
/// (`< 7`) — a single fact about the whole index, decided once here so /// (`< 7`) — a single fact about the whole index, decided once here so
@@ -51,5 +52,10 @@ pub(crate) use subsample::sample_index;
/// interprets that field (`SiblingExt::build_sibling_annex`, /// interprets that field (`SiblingExt::build_sibling_annex`,
/// `SiblingBuilder::ensure_entropy_annexes`), so they can never disagree. /// `SiblingBuilder::ensure_entropy_annexes`), so they can never disagree.
pub(crate) fn is_fast_mode(cache: &IndexCache) -> bool { pub(crate) fn is_fast_mode(cache: &IndexCache) -> bool {
cache.partitions().filter_map(|p| cache.n_layer(p)).max().unwrap_or(0) <= 7 cache
.partitions()
.filter_map(|p| cache.n_layer(p))
.max()
.unwrap_or(0)
<= 7
} }
@@ -71,7 +71,7 @@ const MAX_TOPUP_ROUNDS: usize = 5;
/// outlive that callback, but a `SurvivingFamily` is meant to be handed to /// outlive that callback, but a `SurvivingFamily` is meant to be handed to
/// [`sample_index`]'s caller as part of a whole layer's batch, well after /// [`sample_index`]'s caller as part of a whole layer's batch, well after
/// that callback returns. /// that callback returns.
pub(crate) struct SurvivingFamily { pub struct SurvivingFamily {
pub family_idx: usize, pub family_idx: usize,
pub mask: FamilyMask, pub mask: FamilyMask,
pub genome_mask: Vec<u8>, pub genome_mask: Vec<u8>,
@@ -162,7 +162,10 @@ pub(crate) fn sample_index(
} }
} }
let total_eligible: u64 = layers.iter().map(|(_, _, e, _)| e.view().count_ones()).sum(); let total_eligible: u64 = layers
.iter()
.map(|(_, _, e, _)| e.view().count_ones())
.sum();
if total_eligible == 0 { if total_eligible == 0 {
return Ok(0); return Ok(0);
} }
@@ -272,7 +275,9 @@ fn sample_layer(
Some(_) => Box::new(circular_entropy_values(&layer_dir, round_start)?), Some(_) => Box::new(circular_entropy_values(&layer_dir, round_start)?),
None => Box::new(std::iter::repeat(0.0f32)), // never read: kernel forced to 1.0 below None => Box::new(std::iter::repeat(0.0f32)), // never read: kernel forced to 1.0 below
}; };
let candidates = (round_start..n_minorants).chain(0..round_start).zip(entropies); let candidates = (round_start..n_minorants)
.chain(0..round_start)
.zip(entropies);
let mut accepted: HashSet<usize> = HashSet::new(); let mut accepted: HashSet<usize> = HashSet::new();
for (family_idx, entropy) in candidates { for (family_idx, entropy) in candidates {
+17 -8
View File
@@ -35,41 +35,47 @@ const SENTINEL: f32 = -1.0;
pub(crate) const ENTROPY_ANNEX_FILE_NAME: &str = "entropy.pent"; pub(crate) const ENTROPY_ANNEX_FILE_NAME: &str = "entropy.pent";
pub(crate) struct EntropyAnnex { pub struct EntropyAnnex {
mmap: Mmap, mmap: Mmap,
n: usize, n: usize,
} }
impl EntropyAnnex { impl EntropyAnnex {
pub(crate) fn open(path: &Path) -> io::Result<Self> { pub fn open(path: &Path) -> io::Result<Self> {
let mmap = unsafe { Mmap::map(&File::open(path)?)? }; let mmap = unsafe { Mmap::map(&File::open(path)?)? };
if mmap.len() < HEADER_SIZE { if mmap.len() < HEADER_SIZE {
return Err(io::Error::new(io::ErrorKind::InvalidData, "PENT file too short")); return Err(io::Error::new(
io::ErrorKind::InvalidData,
"PENT file too short",
));
} }
if mmap[0..4] != MAGIC { if mmap[0..4] != MAGIC {
return Err(io::Error::new(io::ErrorKind::InvalidData, "bad PENT magic")); return Err(io::Error::new(io::ErrorKind::InvalidData, "bad PENT magic"));
} }
let n = u64::from_le_bytes(mmap[8..16].try_into().unwrap()) as usize; let n = u64::from_le_bytes(mmap[8..16].try_into().unwrap()) as usize;
if mmap.len() < HEADER_SIZE + n * 4 { if mmap.len() < HEADER_SIZE + n * 4 {
return Err(io::Error::new(io::ErrorKind::InvalidData, "PENT file truncated")); return Err(io::Error::new(
io::ErrorKind::InvalidData,
"PENT file truncated",
));
} }
Ok(Self { mmap, n }) Ok(Self { mmap, n })
} }
pub(crate) fn len(&self) -> usize { pub fn len(&self) -> usize {
self.n self.n
} }
/// `None` for the sentinel (not yet computed — shouldn't happen against /// `None` for the sentinel (not yet computed — shouldn't happen against
/// a fully-built annex). /// a fully-built annex).
pub(crate) fn get(&self, idx: usize) -> Option<f32> { pub fn get(&self, idx: usize) -> Option<f32> {
let off = HEADER_SIZE + idx * 4; let off = HEADER_SIZE + idx * 4;
let v = f32::from_le_bytes(self.mmap[off..off + 4].try_into().unwrap()); let v = f32::from_le_bytes(self.mmap[off..off + 4].try_into().unwrap());
(v >= 0.0).then_some(v) (v >= 0.0).then_some(v)
} }
} }
pub(crate) struct EntropyAnnexBuilder { pub struct EntropyAnnexBuilder {
mmap: MmapMut, mmap: MmapMut,
} }
@@ -79,7 +85,10 @@ impl EntropyAnnexBuilder {
pub(crate) fn new(n: usize, path: &Path) -> io::Result<Self> { pub(crate) fn new(n: usize, path: &Path) -> io::Result<Self> {
let file_size = HEADER_SIZE + n * 4; let file_size = HEADER_SIZE + n * 4;
let file = OpenOptions::new() let file = OpenOptions::new()
.read(true).write(true).create(true).truncate(true) .read(true)
.write(true)
.create(true)
.truncate(true)
.open(path)?; .open(path)?;
file.set_len(file_size as u64)?; file.set_len(file_size as u64)?;
let mut mmap = unsafe { MmapMut::map_mut(&file)? }; let mut mmap = unsafe { MmapMut::map_mut(&file)? };
+14 -14
View File
@@ -21,8 +21,10 @@
//! | batch | [`SiblingBatchIter`] | [`MinorantBatchIter`] | //! | batch | [`SiblingBatchIter`] | [`MinorantBatchIter`] |
//! //!
//! No separate "enumerate" variant (unlike `KmerIter`/`enumerate_kmers`): //! No separate "enumerate" variant (unlike `KmerIter`/`enumerate_kmers`):
//! [`SiblingEntry`] already carries `order` for free, since pairing with //! pairing with the annex is a plain positional `zip` internal to
//! the annex requires it anyway. //! [`SiblingIter`] (a running counter indexing straight into the annex's
//! mmap) — nothing a consumer needs to see, so [`SiblingEntry`] doesn't
//! carry it.
use std::sync::Arc; use std::sync::Arc;
@@ -31,19 +33,17 @@ use obikseq::CanonicalKmer;
use super::{FamilyMask, SiblingAnnex}; use super::{FamilyMask, SiblingAnnex};
/// One layer entry: a k-mer's position in the layer's iteration order (the /// One layer entry: a k-mer and its family mask, paired positionally with
/// same index the sibling annex is keyed on — not an MPHF slot), the k-mer /// the sibling annex (see [`SiblingIter`]'s own docs for how).
/// itself, and its family mask.
#[derive(Debug, Clone, Copy)] #[derive(Debug, Clone, Copy)]
pub(crate) struct SiblingEntry { pub struct SiblingEntry {
pub order: usize,
pub kmer: CanonicalKmer, pub kmer: CanonicalKmer,
pub mask: FamilyMask, pub mask: FamilyMask,
} }
/// Streams `(order, kmer, mask)` triples for one layer, in iteration order. /// Streams `(kmer, mask)` pairs for one layer, in iteration order.
/// Produced by [`SiblingLayerExt::iter_siblings`]. /// Produced by [`SiblingLayerExt::iter_siblings`].
pub(crate) struct SiblingIter { pub struct SiblingIter {
kmers: KmerIter, kmers: KmerIter,
annex: Arc<SiblingAnnex>, annex: Arc<SiblingAnnex>,
order: usize, order: usize,
@@ -61,7 +61,7 @@ impl Iterator for SiblingIter {
// docs) — shouldn't happen against a fully-built annex, but // docs) — shouldn't happen against a fully-built annex, but
// skip rather than misalign the two streams if it does. // skip rather than misalign the two streams if it does.
if let Some(mask) = self.annex.get(order) { if let Some(mask) = self.annex.get(order) {
return Some(SiblingEntry { order, kmer, mask }); return Some(SiblingEntry { kmer, mask });
} }
} }
} }
@@ -69,7 +69,7 @@ impl Iterator for SiblingIter {
/// Batches of [`SiblingIter`]'s entries, `batch_size` at a time — the last /// Batches of [`SiblingIter`]'s entries, `batch_size` at a time — the last
/// batch may be shorter. Produced by [`SiblingLayerExt::iter_siblings_batch`]. /// batch may be shorter. Produced by [`SiblingLayerExt::iter_siblings_batch`].
pub(crate) struct SiblingBatchIter { pub struct SiblingBatchIter {
inner: SiblingIter, inner: SiblingIter,
batch_size: usize, batch_size: usize,
} }
@@ -85,7 +85,7 @@ impl Iterator for SiblingBatchIter {
/// Like [`SiblingIter`], filtered to the minorant of each family — the /// Like [`SiblingIter`], filtered to the minorant of each family — the
/// common case, since a family is tallied once, at its minorant. Produced /// common case, since a family is tallied once, at its minorant. Produced
/// by [`SiblingLayerExt::iter_minorants`]. /// by [`SiblingLayerExt::iter_minorants`].
pub(crate) struct MinorantIter { pub struct MinorantIter {
inner: SiblingIter, inner: SiblingIter,
} }
@@ -100,7 +100,7 @@ impl Iterator for MinorantIter {
/// Batches of [`MinorantIter`]'s entries, `batch_size` at a time — the last /// Batches of [`MinorantIter`]'s entries, `batch_size` at a time — the last
/// batch may be shorter. Produced by /// batch may be shorter. Produced by
/// [`SiblingLayerExt::iter_minorants_batch`]. /// [`SiblingLayerExt::iter_minorants_batch`].
pub(crate) struct MinorantBatchIter { pub struct MinorantBatchIter {
inner: MinorantIter, inner: MinorantIter,
batch_size: usize, batch_size: usize,
} }
@@ -131,7 +131,7 @@ fn collect_batch<I: Iterator>(inner: &mut I, batch_size: usize) -> Option<Vec<I:
/// (`LayerData`) rather than implemented once per matrix kind: kmer /// (`LayerData`) rather than implemented once per matrix kind: kmer
/// iteration doesn't depend on the data payload, and `TypedLayer<D>` already /// iteration doesn't depend on the data payload, and `TypedLayer<D>` already
/// carries `iter_kmers`/`hash_batch`/`fill_sub_matrix_carries` for every `D`. /// carries `iter_kmers`/`hash_batch`/`fill_sub_matrix_carries` for every `D`.
pub(crate) trait SiblingLayerExt { pub trait SiblingLayerExt {
/// Zip this layer's k-mers with their sibling-annex entry, in iteration /// Zip this layer's k-mers with their sibling-annex entry, in iteration
/// order. `annex` must have been built from this same layer (its length /// order. `annex` must have been built from this same layer (its length
/// must match the layer's k-mer count). /// must match the layer's k-mer count).
+8 -4
View File
@@ -20,14 +20,18 @@
//! reverse. //! reverse.
pub mod algorithms; pub mod algorithms;
pub mod extensions;
mod entropy_annex; mod entropy_annex;
pub mod extensions;
mod helpers; mod helpers;
mod iter; pub mod iter;
mod siblingannex; mod siblingannex;
pub(crate) use entropy_annex::{ENTROPY_ANNEX_FILE_NAME, EntropyAnnex, EntropyAnnexBuilder}; pub(crate) use entropy_annex::ENTROPY_ANNEX_FILE_NAME;
pub(crate) use siblingannex::{FamilyMask, SiblingAnnex, SiblingAnnexBuilder}; pub use entropy_annex::{EntropyAnnex, EntropyAnnexBuilder};
pub use iter::{
MinorantBatchIter, MinorantIter, SiblingBatchIter, SiblingEntry, SiblingIter, SiblingLayerExt,
};
pub use siblingannex::{FamilyMask, SiblingAnnex, SiblingAnnexBuilder};
pub use algorithms::{ pub use algorithms::{
BasePairTally, CardinalityTally, EntropyBias, RawSnpDistanceOutput, SankoffBundle, BasePairTally, CardinalityTally, EntropyBias, RawSnpDistanceOutput, SankoffBundle,