Skip to content

Vectorise edge_effect using tskit 1.0 mutation arrays - #193

Merged
jeromekelleher merged 2 commits into
tskit-dev:mainfrom
jeromekelleher:perf
Sep 2, 2026
Merged

Vectorise edge_effect using tskit 1.0 mutation arrays#193
jeromekelleher merged 2 commits into
tskit-dev:mainfrom
jeromekelleher:perf

Conversation

@jeromekelleher

Copy link
Copy Markdown
Member

Replace the Python loop over causal sites and their mutations with an array-based computation using ts.mutations_edge,
ts.mutations_inherited_state and ts.mutations_derived_state. Each causal site owns a contiguous run of mutation IDs, so the per-site boundaries come from a single searchsorted, the causal allele state changes are array comparisons, and the effects are accumulated with np.bincount. This is around 20 times faster: 1.44s -> 0.07s for 168k trait rows over a tree sequence with 338k edges.

Extract the trait dataframe validation shared by genetic_value and edge_effect into _check_trait_df, and add a site_id bounds check there. Both functions now raise a ValueError for a site_id outside [0, ts.num_sites), where a negative value previously wrapped around to the end of the site table.

Expand the edge_effect tests with a state-based reference implementation that is independent of the mutation walking approach, run over the tests/data.py tree sequences, all_trees_ts(2..5) with recurrent and back mutations, and simulations with and without recombination. Cross check against genetic_value(level="node") using the single tree identity edge_effect[e] == value[child(e)] - value[parent(e)], and cover multiple traits, shared causal sites, sites without mutations, stacked mutations, multi-character alleles, isolated nodes and multi-root forests.

Closes #190

Replace the Python loop over causal sites and their mutations with an
array-based computation using ts.mutations_edge,
ts.mutations_inherited_state and ts.mutations_derived_state. Each causal
site owns a contiguous run of mutation IDs, so the per-site boundaries
come from a single searchsorted, the causal allele state changes are
array comparisons, and the effects are accumulated with np.bincount.
This is around 20 times faster: 1.44s -> 0.07s for 168k trait rows over
a tree sequence with 338k edges.

Extract the trait dataframe validation shared by genetic_value and
edge_effect into _check_trait_df, and add a site_id bounds check there.
Both functions now raise a ValueError for a site_id outside
[0, ts.num_sites), where a negative value previously wrapped around to
the end of the site table.

Expand the edge_effect tests with a state-based reference implementation
that is independent of the mutation walking approach, run over the
tests/data.py tree sequences, all_trees_ts(2..5) with recurrent and back
mutations, and simulations with and without recombination. Cross check
against genetic_value(level="node") using the single tree identity
edge_effect[e] == value[child(e)] - value[parent(e)], and cover multiple
traits, shared causal sites, sites without mutations, stacked mutations,
multi-character alleles, isolated nodes and multi-root forests.
@codecov

codecov Bot commented Sep 2, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 100.00%. Comparing base (c1d5dba) to head (6c9d01f).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files
@@            Coverage Diff            @@
##              main      #193   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files            9         9           
  Lines          417       417           
  Branches        56        50    -6     
=========================================
  Hits           417       417           
Flag Coverage Δ
python-tests 100.00% <100.00%> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

tskit provides the state that existed at the site before a mutation
occurred, so there is no need to look up the parent mutation's derived
state and fall back to the site's ancestral state by hand.
@jeromekelleher

Copy link
Copy Markdown
Member Author

Vectorised numpy version is about 20X faster than Python loops, so good enough for now. Also use some tskit 1.0 infrastructure to simplify the code.

@jeromekelleher
jeromekelleher merged commit 31fec9a into tskit-dev:main Sep 2, 2026
11 checks passed
@jeromekelleher
jeromekelleher deleted the perf branch September 2, 2026 13:28
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.

Implement numba algorithm to speed-up genetic_value.py

1 participant