[gregorgorjanc]

Continuing work from jeromekelleher#1

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 100.00%. Comparing base (7b08dbe) to head (a6847c9).

Additional details and impacted files

@@            Coverage Diff            @@
##              main      #189   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files            8         8           
  Lines          348       394   +46     
  Branches        43        54   +11     
=========================================
+ Hits           348       394   +46     

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.

[jeromekelleher]

Generally LGTM. Some performance issues that we perhaps want to leave for a follow-up, but let's try and get the nice docs and the overall framework merged

[jeromekelleher]

Why is this pragma no cover? This seems like a fundamental operation that we'd be testing thoroughly

[gregorgorjanc]

Hmm, this was in from before when we did this in Oxford. I have now removed that and added a test on my end.

[jeromekelleher]

This is quite an inefficient implementation. We need to either implement using numpy operations or numba JIT compile it. I'd prefer numpy if possible (ask your favourite LLM for help).

[jeromekelleher]

The docs build failure looks like fairly basic stuff. You just need to build locally and fixup.

/home/runner/work/tstrait/tstrait/tstrait/genetic_value.py:docstring of tstrait.genetic_value.genetic_value:27: ERROR: Unexpected indentation.
330
/home/runner/work/tstrait/tstrait/tstrait/genetic_value.py:docstring of tstrait.genetic_value.genetic_value:28: WARNING: Block quote ends without a blank line; unexpected unindent.

[gregorgorjanc]

Looked into improving the _accumulate_edge_values() based on #189 (comment)

Using this script

"""
Benchmark the Python-loop and NumPy implementations of edge-value accumulation.

Run from the repository root with:

    python scripts/benchmark_accumulate_edge_values.py
"""

import argparse
import platform
import statistics
import timeit

import numpy as np
import tskit


def accumulate_edge_values_python(nodes_genetic_value, nodes_edge, num_nodes, num_edges):
    """Original implementation using a Python loop."""
    edges_genetic_value = np.zeros(num_edges)
    for u in range(num_nodes):
        edge = nodes_edge[u]
        if edge != tskit.NULL:
            edges_genetic_value[edge] += nodes_genetic_value[u]
    return edges_genetic_value


def accumulate_edge_values_numpy(nodes_genetic_value, nodes_edge, num_nodes, num_edges):
    """Proposed implementation using NumPy operations."""
    nodes_edge = nodes_edge[:num_nodes]
    nodes_genetic_value = nodes_genetic_value[:num_nodes]
    has_edge = nodes_edge != tskit.NULL
    return np.bincount(
        nodes_edge[has_edge],
        weights=nodes_genetic_value[has_edge],
        minlength=num_edges,
    )


def test_equivalence():
    nodes_genetic_value = np.array([10, 2, 0, 0, -3, 3, 5, 20, -1])
    node_edges = np.array(
        [tskit.NULL, 0, tskit.NULL, 1, 1, 1, 2, tskit.NULL, 2], dtype=int
    )

    observed = accumulate_edge_values_python(
        nodes_genetic_value=nodes_genetic_value,
        nodes_edge=node_edges,
        num_nodes=len(nodes_genetic_value),
        num_edges=3,
    )

    np.testing.assert_array_equal(observed, [2, 0, 4])

    observed2 = accumulate_edge_values_numpy(
        nodes_genetic_value=nodes_genetic_value,
        nodes_edge=node_edges,
        num_nodes=len(nodes_genetic_value),
        num_edges=3,
    )

    np.testing.assert_array_equal(observed, observed2)


def make_data(num_nodes, rng):
    """
    Generate inputs with unique edge IDs and some null incoming edges.

    A node has at most one incoming edge in a tree. Root nodes and nodes absent
    from the current tree have tskit.NULL in Tree.edge_array.
    """
    num_edges = max(1, 2 * num_nodes)
    num_attached_nodes = int(0.8 * num_nodes)
    nodes_edge = np.full(num_nodes, tskit.NULL, dtype=np.int32)
    nodes_edge[:num_attached_nodes] = rng.choice(
        num_edges, size=num_attached_nodes, replace=False
    )
    rng.shuffle(nodes_edge)
    nodes_genetic_value = rng.normal(size=num_nodes)
    return nodes_genetic_value, nodes_edge, num_edges


def time_implementation(function, args, number, repeat):
    timings = timeit.repeat(lambda: function(*args), number=number, repeat=repeat)
    return statistics.median(timings) / number


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--sizes",
        nargs="+",
        type=int,
        default=[10, 100, 1_000, 10_000, 100_000, 1_000_000],
        help="Numbers of nodes to benchmark.",
    )
    parser.add_argument(
        "--repeat",
        type=int,
        default=5,
        help="Number of independent timing repetitions.",
    )
    parser.add_argument(
        "--target-nodes",
        type=int,
        default=2_000_000,
        help="Approximate number of node visits timed per repetition.",
    )
    args = parser.parse_args()

    print(f"Python {platform.python_version()}")
    print(f"NumPy {np.__version__}")
    print()
    print(
        f"{'nodes':>10} {'calls':>8} {'Python loop':>14} {'NumPy':>14} {'speedup':>10}"
    )
    print("-" * 62)

    rng = np.random.default_rng(12345)
    for num_nodes in args.sizes:
        nodes_genetic_value, nodes_edge, num_edges = make_data(num_nodes, rng)
        function_args = (
            nodes_genetic_value,
            nodes_edge,
            num_nodes,
            num_edges,
        )

        expected = accumulate_edge_values_python(*function_args)
        observed = accumulate_edge_values_numpy(*function_args)
        np.testing.assert_allclose(observed, expected)

        number = max(1, args.target_nodes // num_nodes)
        python_time = time_implementation(
            accumulate_edge_values_python, function_args, number, args.repeat
        )
        numpy_time = time_implementation(
            accumulate_edge_values_numpy, function_args, number, args.repeat
        )
        speedup = python_time / numpy_time

        print(
            f"{num_nodes:10,d} {number:8,d} "
            f"{python_time * 1e6:11,.1f} us "
            f"{numpy_time * 1e6:11,.1f} us "
            f"{speedup:9.1f}x"
        )


if __name__ == "__main__":
    test_equivalence()
    main()

I get

❯ mamba run -n arg_work \
  python scripts/benchmark_accumulate_edge_values.py
Python 3.13.5
NumPy 2.2.6

     nodes    calls    Python loop          NumPy    speedup
--------------------------------------------------------------
        10  200,000         1.9 us         1.7 us       1.1x
       100   20,000        17.1 us         2.0 us       8.5x
     1,000    2,000       183.5 us         6.3 us      29.4x
    10,000      200     1,815.4 us        47.4 us      38.3x
   100,000       20    19,049.1 us       566.1 us      33.7x
 1,000,000        2   204,744.7 us     8,503.9 us      24.1x

Nice;)

[jeromekelleher]

Great stuff, thanks @gregorgorjanc. I think we should merge. I can follow up with a few updates, like the perf and bugfix for #191

[gregorgorjanc]

I have now squashed the commits so this can be merged in.