Skip to content

Large scale perf - #194

Open
jeromekelleher wants to merge 13 commits into
tskit-dev:mainfrom
jeromekelleher:large-scale-perf
Open

Large scale perf#194
jeromekelleher wants to merge 13 commits into
tskit-dev:mainfrom
jeromekelleher:large-scale-perf

Conversation

@jeromekelleher

Copy link
Copy Markdown
Member

Currently performance is very poor on large tree sequences with large numbers of causal sites. This PR changes the core algorithm from a tree-by-tree traversal approach to a whole-ARG traversal.

@daikitag would you mind trying this out to see if it helps in your case? There may be some memory issues, so do report back on what you see.

The genetic value computation is slow for a trait with a large number of
weakly causal sites, so add a benchmark that measures how it scales.

The script simulates 100k sample nodes over 5Mb, caching the tree
sequence, and times genetic_value for 1, 100, 1000, 10,000 and 100,000
causal sites at each of the individual, node and edge levels. The
mutation rate is raised to ten times the human rate so that there are
enough sites in a genome short enough to simulate quickly; this leaves
the allele frequency spectrum, and so the weakly causal regime, intact.
sim_trait is timed separately because it has a per-site Python loop of
its own, the numba kernels are compiled by an untimed warm up call, and
--max-seconds gives up on the larger points of the grid once a single
call goes over budget.

On a tree sequence with 217k nodes the cost per causal site is flat from
1000 sites upwards, at 444us for individuals, 384us for nodes and 1807us
for edges, so 100,000 causal sites take 44s, 38s and 181s respectively.
Holding the number of causal sites fixed and varying the sample size,
the cost per site normalised by num_nodes is constant across a sixteen
fold range, so the computation is O(num_causal * num_nodes): every
causal site pays several passes over every node in the tree sequence.
The median causal site is carried by 0.4% of the nodes.
Replace the tree by tree loop in genetic_value with one descent over the
ARG that accumulates every causal site at once, using the tskit child
index to find the child edges of a node and carrying a range of causal
site indexes down so that a subtree spanning no causal site is pruned.
The causal allele state change of a mutation telescopes down a path, so
values are additive from root to node and nested and back mutations need
no special handling; the has_mutation pruning is gone, and so is the
separate accumulation pass for edges, which now fall out of the descent
by crediting the edge that a value arrives through.

The value carried down is a sum over the sites in the current range, so
when an edge narrows that range it is not the value the child should
inherit: the mutations above at the sites that dropped out do not apply
below. The path back to the root is therefore kept, and the value is
recomputed over the narrowed range when that happens. It is rare, a
hundredth of a percent of descents on a large tree sequence.

This is slower than what it replaces and is committed as a checkpoint
only. On 100k samples at level="node", 1000 causal sites go from 0.436s
to 14.4s and 10,000 from 3.761s to 51.7s. Cost is independent of allele
frequency as intended, rare and uniform causal sites taking the same
time, but the per visit constant is around 40 times worse than a
prototype that only counted visits, once the value lookups, the per edge
mutation searches and the output writes are in the loop.

Add naive_genetic_value, the tree by tree implementation kept as a
reference oracle, and TestGeneticValueReference comparing against it over
the tests/data.py tree sequences, all_trees_ts(2..5) with recurrent and
back mutations, and simulations with and without recombination, at each
of the individual, node and edge levels and with multiple traits, along
with isolated samples, multiple roots and mutations above a root. Give
the benchmark a causal site selection axis, since drawing sites
uniformly is dominated by the common variants in the tail of the
frequency spectrum and the rare ones behave quite differently.
Replace the range carrying descent with a single sweep of the nodes from
the past to the present. Each node conceptually holds a set of (causal
site, effect) tuples, seeded at the node of each causal mutation. A
parent is always older than its child, so time order is a topological
order and every node is reached after all of its ancestors: for each
tuple the outbound edges spanning its causal site are found, and the
tuple is credited to the edge and to the node at the far end before
being passed on to it. Nodes are credited when a tuple arrives rather
than when they are swept, so a node that is never a parent never holds a
tuple, which is about half of the work on a large tree sequence.

Each tuple carries a single causal site rather than a range, so there is
nothing to narrow and none of the path recomputation the descent needed.
Tuples live in an arena threaded as a per node linked list, with a free
list, since every tuple of a node is dead once that node has been swept
and only those in flight need to be held.

The causal allele being the ancestral state needs no separate treatment
either: there is no mutation to seed from, so the roots are seeded
instead and the effect reaches exactly the nodes of the tree. Mutations
above a root need nothing at all, since seeding at the node and pushing
down is already right.

Cost is the number of nodes that carry a causal allele, rather than the
size of the tree sequence, so the gain depends entirely on allele
frequency. On 100k samples at level="node", against the tree by tree
implementation, causal sites drawn from those below a frequency of 0.001
go from 0.436s to 0.060s at 1000 sites, 3.761s to 0.099s at 10,000, and
38.374s to 0.665s at 100,000. Causal sites drawn uniformly over all
sites are around 2.8 times slower throughout, because the common
variants in the tail of the frequency spectrum touch 7.4% of the nodes
on average where the median site touches 0.39%, and a dense pass over
every node is sequential where following carriers is not.

Retarget the jit tests at the new kernel, dropping the binary search
helpers that only existed for the range searches.
Replace the hand rolled array of linked lists with a typed list per node,
holding indexes into the seed arrays. The causal site and effect size of
a seed are fixed for the whole sweep, so an item is one integer and is
looked up rather than copied. A node's list is made when something first
reaches it, since most nodes are never reached when the causal alleles
are rare, and dropping the reference once the node has been swept hands
the storage straight back. This removes the arena, the free list and the
doubling block that was written out twice.

All figures below are level="node" on a tree sequence of 100,000 samples,
217,091 nodes, 285,979 edges, 22,523 trees and 235,458 sites, taking
causal sites either uniformly over all sites or from those with an allele
frequency below 0.001.

Against the array of linked lists it replaces, per causal site count:

                     1,000    10,000   100,000
    uniform  typed   0.656s    5.615s    54.8s
             manual  1.121s   10.427s   109.7s
                      1.71x     1.86x    2.00x
    rare     typed   0.076s    0.163s   0.558s
             manual  0.060s    0.099s   0.665s
                      0.79x     0.61x    1.19x

It is faster wherever the structure is under any pressure, and slower
only where the lists are a few items long and the sweep is short enough
that making them is a noticeable part of it, a difference of tens of
milliseconds. Peak memory at 100,000 uniformly drawn causal sites falls
from 4.57GB to 1.22GB, because each node's storage is released as the
sweep passes it rather than being held to the high water mark with up to
twice as much again in slack.

Against the tree by tree implementation this branch started from, which
took 0.436s, 3.761s and 38.374s for 1000, 10,000 and 100,000 causal
sites, and whose cost comes from passes over every node rather than from
the number of carriers:

                     1,000    10,000   100,000
    rare              5.7x     23.1x     68.8x
    uniform          0.66x     0.67x     0.70x

The remaining loss on uniformly drawn causal sites is the common variants
in the tail of the frequency spectrum, which touch 7.4% of the nodes on
average where the median site touches 0.39%. Following carriers is random
access at around 50ns each, while a dense pass over every node is
sequential at under a nanosecond, so the two cross over at a carrier
fraction of one or two percent.
@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 (31fec9a) to head (aca6627).

Additional details and impacted files
@@            Coverage Diff            @@
##              main      #194   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files            9         9           
  Lines          417       510   +93     
  Branches        50        63   +13     
=========================================
+ Hits           417       510   +93     
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.

The tree sequence the ARG sweep was written against takes nine minutes
for a full grid and 48 seconds for its longest single call, which is too
slow to work against, and the benchmark measured wall time and nothing
else. The peak memory in the last commit message and the carrier
fractions in the one before it were both produced out of tree.

Add a --preset flag and default it to a simulation of 30,000 samples
over 1Mb, which has 63,287 nodes against 217,091 and runs the grid in
thirty seconds. What makes it a fair substitute is not the timings but
the distribution underneath them, since the sweep costs the number of
nodes that carry a causal allele: the mean fraction of the nodes that a
causal site reaches is 7.19% against 5.99%, and the median 0.406%
against 0.184%, which is as close as 150 sampled sites from a
distribution with that tail can show. Both presets then have a flat per
call floor, a uniform cost per causal site that falls to a plateau from
1000 sites upwards, a rare cost still falling at the top of the grid, a
uniform to rare ratio growing to two orders of magnitude, and parity
between the three levels. The plateau is 160us against 483us per causal
site, a factor of 3.0 on a node count ratio of 3.4. Sixteen times
smaller was tried and is measurably less faithful, and the grid stops at
10,000 causal sites because raising the mutation rate to lift the
ceiling would take sites per node from 0.68 to 2.7 against the large
preset's 1.08 and inflate the setup floor.

Measure four things besides the wall clock. --phases times the check,
the setup, the kernel and the dataframe separately, without which a flat
setup cost and a kernel that grows with the causal sites cannot be told
apart; setup goes from 8ms to 17ms across the whole small grid and the
dataframe is 1ms, so at 10,000 causal sites the kernel is 98% of the
call. --counters runs a counting only copy of the sweep, since perf
cannot see inside the kernel. --structure reports the carrier fractions
above. --memory reports peak RSS, resetting VmHWM through clear_refs
because it never falls on its own.

The counters say three things that were assumed otherwise. The scan of a
node's out edges is not where the waste is, because 94% of the trips
find an edge spanning the seed's causal site. Uniform selection costs
178 times as many scans as rare at 10,000 causal sites but only 46 times
the time, because rare costs 63ns a scan against 28ns, which is the
random access against sequential difference measured directly rather
than inferred. And the kernel has a floor in the number of nodes: it
builds a pending entry for every node and sweeps every node whether or
not anything reached it, so one rare causal site, ten scans, still takes
2ms.

Add profile_genetic_value.py, which profiles one cell of the grid with
cProfile for the setup and sets up a cell to run just the kernel under
perf, printing the commands. Two things had to be established for that
to be worth anything. Source lines inside the kernel are not available,
because this llvmlite has no LLVM PerfJITEventListener and so nothing
writes a perf map or a jitdump; what perf does give is the split between
the kernel, the numba runtime with its typed list functions named, the
interpreter and compilation. And NUMBA_ENABLE_PROFILING=1, the
documented way to profile numba, is the wrong thing to use here: it
would only help through the listener that is missing, and it defaults
NUMBA_DEBUGINFO to 1, which measured 2.62s a call against 1.61s. perf
finds the JIT mappings by itself.

Record that _GeneticValue takes the _root_runs branch, a Python loop
over every tree costing 7ms on the small preset and 38ms on the large
one. It is not the edge case it looks like: drawing 10,000 sites
uniformly makes it near certain that one of them has the ancestral state
as its causal allele, and it fires at 10,000 and 100,000 causal sites on
both presets, so on the large one it is over a third of the setup.

Check in the small preset grid to diff against. The counts in it are
identical between runs and are the part worth treating as a regression
test; the timings are specific to the machine they were taken on.
The three levels differ only in which slot a contribution is credited to,
and the kernel already discards a negative slot, so individuals need no
more than a node to individual mapping: nodes_individual is exactly that
array, and tskit.NULL for a node belonging to no individual is already
negative. That removes the separate O(num_nodes) accumulation pass the
individual level ran after the descent, one per trait, and with it
_accumulate_individual_values.

Accumulate into the caller's output array rather than allocating one per
trait, so that a second implementation can write into the same table and
combining the two is nothing at all.

Split the causal mutation expansion in two. A descent of the trees needs
every mutation at a causal site, because a mutation blocks the
inheritance of the allele above it whatever it changes the state to,
where a sum of state changes needs only those that change the state.
_row_mutations returns all of them and _causal_mutations filters, which
is what the push down seeds and edge_effect keep using. Keep the row each
seed came from for the same reason: a subset of the rows can then be
selected without recomputing the seeds.

No change in what is computed. Verified against the tree by tree
implementation at 31fec9a on 500 causal sites over a 30,000 sample tree
sequence, at each level and for one and three traits, agreeing to 3e-14.
A descent of a tree needs the tree, so add the incremental construction
that a second implementation of genetic values will run over: apply the
edges leaving and then entering at each tree of tskit's TreeIndex,
maintaining the quintuply linked encoding along with the edge above each
node. Attaching and detaching a child are both constant time, so a full
pass costs one insertion and one removal per edge: 2.25ms over the 75,586
edges of a 30,000 sample tree sequence and 12.35ms over the 285,979 edges
of a 100,000 sample one, which is less than the setup that the push down
already pays on the same data.

There is no virtual root. The descent starts at the mutations of a causal
site and needs the roots only when the ancestral state is the causal
allele, which is 0.04% of the sites of the smaller tree sequence and
0.07% of the larger. Maintaining the virtual root's children costs
something on every edge, and knowing which nodes belong in that list
means tracking the samples below every node; walking up from each sample
and taking the top of the path finds the same roots on demand, and marking
the nodes already walked through keeps the whole thing to one visit per
node however many samples there are.

That leaves one difference from tskit: it threads the roots together as
children of the virtual root, so they are siblings of each other there
and have none here. Nothing else differs, and the descent never walks the
siblings of a root, so the test asserts the sibling arrays match at every
node that has a parent and are null at the roots, along with the parent,
left child and edge arrays and the root set itself, tree by tree over
every fixture in tests.data, all_trees_ts(2..5), comb trees, multiple
roots, isolated samples and a tree sequence with no nodes at all.
The push down of the ARG costs the nodes that carry a causal allele, at
around 28ns each because following carriers is random access. A descent
of the trees costs the same nodes at around 20ns, since a child is the
next link of a list rather than a search of the out edges of a node, and
it needs no typed list per node to hold what is in flight. Against that
it has to build the trees, which the push down never does.

So neither wins everywhere, and which one a causal site should take
follows from its allele frequency: the tree pass is a fixed cost that
common sites amortise and rare ones do not. Add the descent as a second
implementation and route each row to one or the other on allele_freq,
which sim_trait already returns and _check_trait_df now keeps when it is
there. A trait dataframe assembled by hand has no such column and takes
the push down for everything, which is what it did before.

Measured on 30,000 samples over 1Mb at level="node", end to end:

                     1,000    10,000
    uniform  before   195ms    1624ms
             after    129ms    1080ms
                      1.51x     1.50x
    rare     before  16.3ms    31.0ms
             after   16.6ms    31.1ms

Rare sites are below the threshold and take the same path as before, so
the gain is on the common ones, which is where the time was: sites above
a frequency of 0.03 are under a third of the sites and 94 to 97% of the
time.

The descent walks the rows in step with the trees, which works because
_check_trait_df already requires the rows to be sorted by site. A row's
mutations are marked with the row's own index rather than into an array
that has to be cleared, so a descent costs the nodes it reaches and
nothing per node of the tree sequence. All of the traits share the pass,
where the push down runs once per trait.

Two things the descent has to do that the push down does not. It takes
every mutation at a causal site, not only the state changing ones, since
a mutation blocks the inheritance of the allele above it whatever it
changes the state to. And where the causal allele is the ancestral state
it seeds the roots, skipping any root that carries a mutation: the push
down seeds every root and lets the root's own mutation cancel it, and the
descent has no such cancellation, so seeding an ancestral root that
carries the allele would count it twice. Removing that skip fails
test_mutation_on_root and two of the reference comparisons.

Run every comparison against the tree by tree reference through both
implementations and a mix of the two, and add tests for the cases where
they are most likely to part company: a mutation on a root under both
causal alleles, a causal site carrying no mutations at all, a site
exactly on a breakpoint, and that a middling threshold really does divide
the rows rather than quietly sending them all the same way.
Compare the two implementations over the whole benchmark grid, on both
presets, at all three levels, for one and three traits, and for causal
sites drawn both uniformly and from the rare ones. A threshold of zero,
sending everything to the descent of the trees, was fastest or tied in
every cell:

                      uniform            rare
                    1,000  10,000   1,000  10,000
    small  descent   122ms  1058ms    15ms    23ms
           push down 196ms  1619ms    16ms    35ms
    large  descent   427ms  3709ms    71ms    90ms
           push down 592ms  4917ms    68ms   108ms

The push down led only where the causal sites are few and rare, so the
tree pass is not amortised, and there it led by a millisecond or two on a
call taking a few. With three traits the gap widens the other way, to
2.27x on rare causal sites, since the push down sweeps once per trait
where the descent takes all of them in one pass.

So set the threshold to zero. Every allele frequency is at or above zero,
so there is nothing to compare against and nothing to look up, which is
what lets the default apply to a trait dataframe assembled by hand as
well as to one from sim_trait. The push down is kept, and --threshold inf
still runs it, only so that the two can go on being compared.

Against the tree by tree implementation this branch started from, at
level="node" on 100,000 samples, the whole of it is now ahead:

                     1,000    10,000   100,000
    tree by tree     0.436s    3.761s    38.374s
    rare             0.072s    0.091s     0.216s
                       6.1x      41x       178x
    uniform          0.443s    3.848s    32.589s
                      0.98x     0.98x      1.18x

which closes the regression on uniformly drawn causal sites that the push
down left behind at 0.70x.

Report the work the descent does rather than the work the push down does,
since that is the code that runs. The descent kernel returns the number
of nodes it visited, so there is no second copy of the loop to keep in
step with the first, and the count is the thing its run time is
proportional to: seconds over visits is around 20ns once there are enough
causal sites to amortise the tree pass. Visits per row as a fraction of
the nodes comes to 8.4% for uniformly drawn causal sites and 0.02% for
rare ones, against the 7.2% and 0.4% that --structure measures from the
other end.

The _root_runs branch is only taken by rows that go to the push down, so
at the default threshold it never runs; say so rather than reporting that
any ancestral causal allele triggers it.
@daikitag

daikitag commented Sep 3, 2026

Copy link
Copy Markdown
Collaborator
image
import tskit
import tstrait

ts = tskit.load("/data/smew01/not-backed-up/scratch/tagami/ukb-simulation/output/slim_mu_1.7e-10_demography_relate_OutOfAfrica_4J17_ceu_slim_seed_100_chr_1q/slim_allele.tree")
print(ts)

print("tstrait version:", tstrait.__version__)

model = tstrait.trait_model(distribution="normal", mean=0, var=1)
sim_result = tstrait.sim_phenotype(
    ts=ts, num_causal=ts.num_sites, model=model, h2=0.3, random_seed=1
)
print("phenotype:\n")
print(sim_result.phenotype.head())
print("trait:\n")
print(sim_result.trait.head())

I ran the above codes by using the tree sequence indicated in the image.

tstrait version 0.1.0 (previous one):
CPU Utilized: 07:22:02
CPU Efficiency: 9.97% of 3-01:55:00 core-walltime
Job Wall-clock time: 07:23:30
Memory Utilized: 28.07 GB
Memory Efficiency: 18.71% of 150.00 GB

pull request tstrait version:
CPU Utilized: 00:45:06
CPU Efficiency: 9.96% of 07:32:50 core-walltime
Job Wall-clock time: 00:45:17
Memory Utilized: 56.15 GB
Memory Efficiency: 37.43% of 150.00 GB

There is a significant improvement in the computational speed in the current codes compared with the previous tstrait.

@jeromekelleher

@daikitag

daikitag commented Sep 3, 2026

Copy link
Copy Markdown
Collaborator

The seed is also reflected well, and the two simulations produced identical phenotype and trait dataframe.

The push down holds the seeds still in flight, which grows with the
number of causal sites; the descent holds a fixed handful of arrays the
length of the nodes whatever the trait looks like. At 100,000 uniformly
drawn causal sites on 100,000 samples, that is 0.01GB over the baseline
against 0.89GB, and a peak of 0.42GB against 1.30GB.
@jeromekelleher

Copy link
Copy Markdown
Member Author

Excellent, thanks @daikitag. I'm working on further improvements, which should bring it down more.

Two implementations of the genetic value computation have been carried
side by side so that they could be compared. The comparison is done, so
delete the one that lost, along with the threshold that chose between
them and everything built only to feed it.

Measured over the whole benchmark grid, on both tree sequences, at each
level, for one and three traits, and for causal sites drawn both
uniformly and from the rare ones, the push down was slower everywhere the
cost mattered, by 1.33 to 1.70 times on uniformly drawn causal sites and
up to 2.27 times on rare ones with three traits, where it swept once per
trait against a single pass for all of them:

                      uniform            rare
                    1,000  10,000   1,000  10,000
    small  descent   122ms  1058ms    15ms    23ms
           push down 196ms  1619ms    16ms    35ms
    large  descent   427ms  3709ms    71ms    90ms
           push down 592ms  4917ms    68ms   108ms

It was ahead in ten of the thirty six cells, every one of them rare
causal sites at a thousand or fewer, where the tree building the descent
starts with has too little to amortise it against. In none of them was it
ahead by more than a few milliseconds on a call taking a few. The note
left in jit.py says as much, since a reader can check it against the same
benchmark.

Deleting it takes the setup with it, which was built on every call
whether the push down ran or not: an argsort over the nodes, two
searchsorted over the edge table, the tskit child index, and the seed
arrays. _GeneticValue setup on 100,000 samples goes from 98ms to 42ms,
almost all of what is left being tskit's jitwrap, and every cell of the
small preset grid gets between 1.01 and 1.41 times faster with no change
to the kernel at all.

_root_runs goes too. It was a Python loop over every tree, finding the
roots for causal sites whose ancestral state is the causal allele, and
only the push down needed it: the descent is already standing in the tree
and asks _tree_roots. The threshold argument, the optional allele_freq
column that it read, and the row selection the descent kernel took all go
with it, so _check_trait_df is back to the four required columns and
sim_trait's allele_freq is once again only an output.

The thirty five topology tests carry over unchanged by repointing the one
fixture they share; the jit and nojit parametrisation survives, though
nojit now interprets only the descent's own loop, since the tree building
kernels it calls are compiled. Those are covered instead by TestTreeState
against tskit.Tree. The tests that existed to compare the two
implementations go, and the reference comparison against the tree by tree
oracle stops looping over thresholds.

Verified against that same tree by tree implementation, recovered from
31fec9a, on both tree sequences at all three levels for one and three
traits.
@jeromekelleher

Copy link
Copy Markdown
Member Author

I've changed the algorithm, which should be a bit better. Can you try it again please @daikitag ?

Coverage of jit.py fell to 62% when the push down went. It had called no
other kernel, so running it through py_func interpreted everything it did;
the descent calls five kernels to build the trees, and a numba function
calls whatever the name is bound to in the module, so those five stayed
compiled and the coverage of their bodies went with them.

Swap the whole set for their py_func together when the nojit parameter
asks for it, so that a kernel is interpreted the whole way down rather
than only in its own loop, and give the tree state test the same
parameter as the rest. That is what the module docstring has always
claimed the convention is; correct the caveat that said otherwise.

jit.py is back to 100%, and so is every other module, on branches as well
as lines. Running only the nojit half covers jit.py completely and only
the jit half covers 16% of it, so the measurement is coming from the
interpreted path rather than from somewhere incidental.
@jeromekelleher

Copy link
Copy Markdown
Member Author

I think we can merge this, assuming that it does in fact lead to performance improvement over the previous version.

@daikitag

daikitag commented Sep 3, 2026

Copy link
Copy Markdown
Collaborator

CPU Utilized: 02:03:34
CPU Efficiency: 9.96% of 20:40:20 core-walltime
Job Wall-clock time: 02:04:02
Memory Utilized: 28.46 GB
Memory Efficiency: 18.97% of 150.00 GB

The new codes took a slightly longer time to simulate traits but with a lower memory.

@jeromekelleher

@jeromekelleher

Copy link
Copy Markdown
Member Author

Huh, 124m vs 45 minutes is quite a big difference. That's completely against what my local benchmarks were saying. Always good to try on the real data!

I'm going to try parallelising the current algorithm though. If that scales over multiple threads, it's a much simpler algorithm with predictable memory usage so could still win.

Thanks @daikitag - I'll ping you again when there's something to try out.

@daikitag

daikitag commented Sep 3, 2026

Copy link
Copy Markdown
Collaborator

Thank you for modifying the codes, Jerome! There is still a substantial improvement in the current codes compared with the previous codes, so it will definitely improve the usability of the tstrait package on large tree sequences.

The rows of the trait dataframe are independent of each other, so add a
num_threads argument that gives each of that many worker threads a
contiguous range of them. It defaults to 0, which does the work on the
calling thread with no pool created, matching what tskit's
divergence_matrix and its neighbours mean by the same argument.

Releasing the GIL is what makes it worth anything: a numba kernel holds
it for its whole execution by default, so the threads would have taken
turns. With nogil=True on the descent, four of them report a busy to wall
ratio of 3.8 out of 4.

A range needs nothing from any other. The kernel takes row_start and
row_stop rather than a slice, so nothing is copied or rebased: pair_offset
is indexed by absolute row, and the arrays a thread only reads are built
once and shared. Each range accumulates into its own table, since the
kernel adds into it and that is not atomic, and the tables are summed at
the end, which is where the answer stops being bit for bit what one
thread would have produced. The rows are marked with their own index, and
two ranges cannot collide.

Minimum of two replicates at level="node" on four cores:

    preset selection num_causal  1 thread     2     3     4
    small  uniform       10,000     980ms 1.94x 2.86x 3.41x
    small  uniform          300      41ms 1.36x 1.66x 1.69x
    small  rare          10,000      17ms 0.94x 0.91x 0.89x
    large  uniform       10,000    3338ms 1.47x 1.64x 1.64x
    large  rare          10,000      76ms 1.01x 1.00x 0.99x

Two things in there are worth saying out loud, since neither is what one
would assume. A bigger trait parallelises and a bigger tree sequence does
not: the kernel holds around eleven arrays the length of the nodes per
thread, 49 bytes a node, which is 3.1MB on 30,000 samples so four threads
sit inside a 16MB L3, and 10.6MB on 100,000 samples so four want 42MB and
are held up by memory bandwidth instead. Going from 10,000 causal sites
to 100,000 on the larger one moves four threads from 1.50x only to 1.61x,
so this is cache capacity rather than the division of work. And threads
cost more than they save on a rare or a small trait, because a thread
walks the whole tree sequence whatever range it takes: there is no
seeking to the first tree a range wants, and that pass is nothing against
a second of descending and everything against seventeen milliseconds.

Equal row counts turned out to need no cost model behind them. They were
the thing most likely to want one, since rows differ enormously in how
many nodes they reach, but four ranges came out within 1.11x of each other
on time and 1.18x on nodes visited.

The sequential path is unchanged, within 0.96x to 1.09x of the committed
baseline over five replicates.
@jeromekelleher

Copy link
Copy Markdown
Member Author

I've just added a num_threads argument to genetic_value @daikitag. Can you update your script to call this please? (sim_phenotype hasn't been wired up).

It seems to be cache-sensitive, so it would be good to try num_threads=4 and num_threads=8 separately to see how well it scales with such a large ARG.

@daikitag

daikitag commented Sep 4, 2026

Copy link
Copy Markdown
Collaborator

num_threads=0:
State: COMPLETED (exit code 0)
Nodes: 1
Cores per node: 8
CPU Utilized: 02:08:44
CPU Efficiency: 12.45% of 17:13:52 core-walltime
Job Wall-clock time: 02:09:14
Memory Utilized: 28.44 GB
Memory Efficiency: 28.44% of 100.00 GB

num_threads=4:
State: COMPLETED (exit code 0)
Nodes: 1
Cores per node: 4 (I only requested 4 cores)
CPU Utilized: 02:34:33
CPU Efficiency: 65.05% of 03:57:36 core-walltime
Job Wall-clock time: 00:59:24
Memory Utilized: 38.03 GB
Memory Efficiency: 38.03% of 100.00 GB

num_threads=8:
State: COMPLETED (exit code 0)
Nodes: 1
Cores per node: 8
CPU Utilized: 03:37:45
CPU Efficiency: 51.52% of 07:02:40 core-walltime
Job Wall-clock time: 00:52:50
Memory Utilized: 50.81 GB
Memory Efficiency: 50.81% of 100.00 GB

There was a slight improvement in computational time when we increase the number of threads.

@jeromekelleher

@jeromekelleher

Copy link
Copy Markdown
Member Author

Hmm, well we halved the time with 4 threads, but it's still slower than the arg algorithm. Seems like it's better to revert to the arg algorithm after all.

Can you tszip that tree sequence file and send it to me on slack please? My benchmarks aren't capturing your use case.

@daikitag

daikitag commented Sep 4, 2026

Copy link
Copy Markdown
Collaborator

I tried to zip the tree sequence file, but it was too big to download on my computer.
I also did the simulation on a smaller tree sequence file with 600K individuals, but there was a negligible difference (11 mins for 4 and 8 threads).
Maybe we can try using tstrait with a large number of causal sites on the French Canadian dataset? @jeromekelleher

@jeromekelleher

Copy link
Copy Markdown
Member Author

Why can't you run tszip on the server?

@jeromekelleher

Copy link
Copy Markdown
Member Author

Just to clarify here @daikitag - are you using tstrait end-to-end to simulate the traits and compute the genetic values, or are you computing genetic values based on traits simulated by SLiM? On slack you said you loaded the effect size dataframe and used tstrait.genetic_value, but above you used sim_phenotype?

This will matter as the performance here depends a lot on the allele frequency, which could differ significantly.

@daikitag

daikitag commented Sep 8, 2026

Copy link
Copy Markdown
Collaborator

@jeromekelleher
Thank you for your message. For the performance test, I simply used the following codes:

model = tstrait.trait_model(distribution="normal", mean=0, var=1)
sim_result = tstrait.sim_phenotype(
    ts=ts, num_causal=ts.num_sites, model=model, h2=0.3, random_seed=1
)

and I did not use SLiM information to simulate the traits.
However, this tree sequence is a raw output of SLiM and it is not recapitated.

@jeromekelleher

Copy link
Copy Markdown
Member Author

Can you profile genetic_value on the real example that you're using please - it's not just about scale, it's about having realistic allele frequency patterns.

@daikitag

daikitag commented Sep 9, 2026

Copy link
Copy Markdown
Collaborator

I used the following codes to measure the times.

import pandas as pd
import tskit
import tstrait
import sys

import time

folder = "/data/smew01/not-backed-up/scratch/tagami/ukb-simulation/output/slim_mu_1.7e-10_demography_relate_OutOfAfrica_
4J17_ceu_slim_seed_100_chr_1q/"

ts = tskit.load(f"{folder}subset_nucleotide_yri_10000_ceu_500000_chb_10000_jpt_10000_individualseed_10_recap_seed_99.tree")
print(ts)

print("tstrait version:", tstrait.__version__)

trait_df = pd.read_csv(f"{folder}mutation_yri_10000_ceu_500000_chb_10000_jpt_10000_individualseed_10_recap_seed_99_prop_0.779_genetic_seed_15.csv")

num_threads = int(sys.argv[1])

start = time.time()
genetic_df = tstrait.genetic_value(ts, trait_df, num_threads=num_threads)
end = time.time()

print("Number of threads:", num_threads)
print("Time Elapsed (Seconds):", end - start)
print(genetic_df.head())

num_threads=0: 425.34 seconds
num_threads=4: 230.50 seconds
num_threads=8: 205.65 seconds

@jeromekelleher

@daikitag

daikitag commented Sep 9, 2026

Copy link
Copy Markdown
Collaborator

This is the summary of the tree sequence data that is used:

image

@jeromekelleher

Copy link
Copy Markdown
Member Author

Ok, great. We're getting a reasonable return from threads here, so I think we'll commit this and tag a release.

Thanks for your help @daikitag !

@daikitag

daikitag commented Sep 9, 2026

Copy link
Copy Markdown
Collaborator

Thank you very much for modifying the algorithm @jeromekelleher !

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.

2 participants