Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

perf: speedup parsimony ancestral reconstruction #292

Merged
merged 6 commits into from
Dec 6, 2024
Merged

Conversation

ivan-aksamentov
Copy link
Member

@ivan-aksamentov ivan-aksamentov commented Dec 6, 2024

  • Simplify speedup and consensus decision in backward pass of the ancestral Fitch method: removes izip and ManyZip pythonisms, flattens and simplifies the loop
  • Introduce BitSet128 data structure: a fixed-size 128-bit integer set with fast constant time lookups. It can store up to 128 chars, which should be enough for our char-based alphabets. For larger alphabet we might have a separate implementation.
  • Replace BTreeSet-based StateSet implementation with BitSet128
  • Replace sets in Alphabet with BitSet128 (which is now the same as StateSet)

These optimizations produce 2x speedup when running ancestral parsimony reconstruction on mpox-500 dataset (and a little more when combined with #290 and #291):

Before:

$ cargo -q build --release --target-dir=/workdir/.build/docker --bin=treetime
$ hyperfine --warmup 1 --show-output '/workdir/.build/docker/release/treetime ancestral --method-anc=marginal --dense=false --tree=data/mpox/clade-ii/500/tree.nwk --outdir=tmp/smoke-tests/ancestral/marginal/mpox/clade-ii/500 data/mpox/clade-ii/500/aln.fasta.xz'
Benchmark 1: /workdir/.build/docker/release/treetime ancestral --method-anc=marginal --dense=false --tree=data/mpox/clade-ii/500/tree.nwk --outdir=tmp/smoke-tests/ancestral/marginal/mpox/clade-ii/500 data/mpox/clade-ii/500/aln.fasta.xz
  Time (mean ± σ):      4.012 s ±  0.151 s    [User: 13.367 s, System: 0.548 s]
  Range (min … max):    3.915 s …  4.379 s    10 runs

After:

$ cargo -q build --release --target-dir=/workdir/.build/docker --bin=treetime
$ hyperfine --warmup 1 --show-output '/workdir/.build/docker/release/treetime ancestral --method-anc=parsimony --dense=false --tree=data/mpox/clade-ii/500/tree.nwk --outdir=tmp/smoke-tests/ancestral/marginal/mpox/clade-ii/500 data/mpox/clade-ii/500/aln.fasta.xz'
Benchmark 1: /workdir/.build/docker/release/treetime ancestral --method-anc=parsimony --dense=false --tree=data/mpox/clade-ii/500/tree.nwk --outdir=tmp/smoke-tests/ancestral/marginal/mpox/clade-ii/500 data/mpox/clade-ii/500/aln.fasta.xz
  Time (mean ± σ):      1.852 s ±  0.095 s    [User: 4.591 s, System: 0.523 s]
  Range (min … max):    1.768 s …  2.055 s    10 runs

One test is failing. We need to figure out ambiguous character reconstruction once and for all. But it will happen in next PR.

@ivan-aksamentov ivan-aksamentov merged commit 495800a into rust Dec 6, 2024
13 of 14 checks passed
@ivan-aksamentov ivan-aksamentov deleted the perf/fitch branch December 6, 2024 01:41
ivan-aksamentov added a commit that referenced this pull request Dec 6, 2024
Followup of: #292

The transition from character sets to bit sets introduced some confusion in handling of ambiguous characters. Ambiguous characters correspond to multiple canonical characters, so they needed explicit disambiguation and disambiguation when using sets. With profile maps (arrays of floats) and with bitsets (~arrays of bits) this happens transparently: multiple "enabled" bits denote canonical components of an ambiguous character.

Here I introduce precalculated `char_to_set` and `set_to_char` maps, which, allow to resolve particular combinations of bits from and to characters. This is similar to how profile maps operate, but with bits (0s and 1s) rather than floating point values.

This allows the remaining test for parsimony reconstruction to pass.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant