test_tree_stats.py

# MIT License
#
# Copyright (c) 2018-2024 Tskit Developers
# Copyright (C) 2016 University of Oxford
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
"""
Test cases for generalized statistic computation.
"""
import collections
import contextlib
import functools
import io
import itertools
import random

import msprime
import numpy as np
import numpy.testing as nt
import pytest

import tests.test_wright_fisher as wf
import tests.tsutil as tsutil
import tskit
import tskit.exceptions as exceptions

np.random.seed(5)


def cached_np(func):
    """
    Decorator to speed up functions that take numpy arrays as positional
    arguments that get called a lot with the same arguments.

    # See https://github.com/tskit-dev/tskit/issues/1856 for more info
    """
    cache = {}

    def f(*args):
        nonlocal cache  # noqa: F824
        key = tuple(x.tobytes() for x in args)
        if key not in cache:
            cache[key] = func(*args)
        return cache[key]

    return f


def subset_combos(*args, p=0.5, min_tests=3):
    # We have too many tests, combinatorially; so we will run a random subset
    # of them, using this function, below. If we don't set a seed, a different
    # random set is run each time. Ensures that at least min_tests are run.
    # Uncomment this line to run all tests (takes about an hour):
    # p = 1.0
    num_tests = 0
    skipped_tests = []
    # total_tests = 0
    for x in itertools.product(*args):
        # total_tests = total_tests + 1
        if np.random.uniform() < p:
            num_tests += 1
            yield x
        elif len(skipped_tests) < min_tests:
            skipped_tests.append(x)
        elif np.random.uniform() < 0.1:
            skipped_tests[np.random.randint(min_tests)] = x
    while num_tests < min_tests:
        yield skipped_tests.pop()
        num_tests += 1
    # print("tests", num_tests)
    assert num_tests >= min_tests


def path_length(tr, x, y):
    L = 0
    if x >= 0 and y >= 0:
        mrca = tr.mrca(x, y)
    else:
        mrca = -1
    for u in x, y:
        while u != mrca:
            L += tr.branch_length(u)
            u = tr.parent(u)
    return L


@contextlib.contextmanager
def suppress_division_by_zero_warning():
    with np.errstate(invalid="ignore", divide="ignore"):
        yield


##############################
# Branch general stat algorithms
##############################


def windowed_tree_stat(ts, stat, windows, span_normalise=True):
    shape = list(stat.shape)
    shape[0] = len(windows) - 1
    A = np.zeros(shape)

    tree_breakpoints = np.array(list(ts.breakpoints()))
    tree_index = 0
    for j in range(len(windows) - 1):
        w_left = windows[j]
        w_right = windows[j + 1]
        while True:
            t_left = tree_breakpoints[tree_index]
            t_right = tree_breakpoints[tree_index + 1]
            left = max(t_left, w_left)
            right = min(t_right, w_right)
            weight = max(0.0, (right - left) / (t_right - t_left))
            A[j] += stat[tree_index] * weight
            assert left != right
            if t_right <= w_right:
                tree_index += 1
                # TODO This is inelegant - should include this in the case below
                if t_right == w_right:
                    break
            else:
                break
    if span_normalise:
        # re-normalize by window lengths
        window_lengths = np.diff(windows)
        for j in range(len(windows) - 1):
            A[j] /= window_lengths[j]
    return A


def naive_branch_general_stat(
    ts, w, f, windows=None, time_windows=None, polarised=False, span_normalise=True
):
    # NOTE: does not behave correctly for unpolarised stats
    # with non-ancestral material.
    if windows is None:
        windows = [0.0, ts.sequence_length]
    drop_time_windows = time_windows is None
    if time_windows is None:
        time_windows = [0.0, np.inf]
    else:
        if time_windows[0] != 0:
            time_windows = [0] + time_windows
    n, k = w.shape
    tw = len(time_windows) - 1
    # hack to determine m
    m = len(f(w[0]))
    total = np.sum(w, axis=0)

    sigma = np.zeros((ts.num_trees, tw, m))
    for j, upper_time in enumerate(time_windows[1:]):
        if np.isfinite(upper_time):
            decap_ts = ts.decapitate(upper_time)
        else:
            decap_ts = ts
        assert np.all(list(ts.samples()) == list(decap_ts.samples()))
        for tree in decap_ts.trees():
            x = np.zeros((decap_ts.num_nodes, k))
            x[decap_ts.samples()] = w
            for u in tree.nodes(order="postorder"):
                for v in tree.children(u):
                    x[u] += x[v]
            if polarised:
                s = sum(tree.branch_length(u) * f(x[u]) for u in tree.nodes())
            else:
                s = sum(
                    tree.branch_length(u) * (f(x[u]) + f(total - x[u]))
                    for u in tree.nodes()
                )
            sigma[tree.index, j, :] = s * tree.span
    for j in range(1, tw):
        sigma[:, j, :] = sigma[:, j, :] - sigma[:, j - 1, :]
    if isinstance(windows, str) and windows == "trees":
        # need to average across the windows
        if span_normalise:
            for j, tree in enumerate(ts.trees()):
                sigma[j] /= tree.span
        out = sigma
    else:
        out = windowed_tree_stat(ts, sigma, windows, span_normalise=span_normalise)
    if drop_time_windows:
        assert out.ndim == 3
        out = out[:, 0]
    return out


def branch_general_stat(
    ts,
    sample_weights,
    summary_func,
    windows=None,
    time_windows=None,
    polarised=False,
    span_normalise=True,
):
    """
    Efficient implementation of the algorithm used as the basis for the
    underlying C version.
    """
    n, state_dim = sample_weights.shape
    windows = ts.parse_windows(windows)
    drop_time_windows = time_windows is None
    time_windows = ts.parse_time_windows(time_windows)
    num_windows = windows.shape[0] - 1
    num_time_windows = time_windows.shape[0] - 1

    # Determine result_dim
    result_dim = len(summary_func(sample_weights[0]))
    result = np.zeros((num_windows, num_time_windows, result_dim))
    state = np.zeros((ts.num_nodes, state_dim))
    state[ts.samples()] = sample_weights
    total_weight = np.sum(sample_weights, axis=0)

    time = ts.tables.nodes.time
    parent = np.zeros(ts.num_nodes, dtype=np.int32) - 1
    branch_length = np.zeros((num_time_windows, ts.num_nodes))
    # The value of summary_func(u) for every node.
    summary = np.zeros((ts.num_nodes, result_dim))
    # The result for the current tree *not* weighted by span.
    running_sum = np.zeros((num_time_windows, result_dim))

    def polarised_summary(u):
        s = summary_func(state[u])
        if not polarised:
            s += summary_func(total_weight - state[u])
        return s

    for u in range(ts.num_nodes):
        summary[u] = polarised_summary(u)

    window_index = 0

    def update_sum(u, sign):
        time_window_index = 0
        if parent[u] != -1:
            while (
                time_window_index < num_time_windows
                and time_windows[time_window_index] < time[parent[u]]
            ):
                running_sum[time_window_index] += sign * (
                    branch_length[time_window_index, u] * summary[u]
                )
                time_window_index += 1

    for (t_left, t_right), edges_out, edges_in in ts.edge_diffs():
        for edge in edges_out:
            u = edge.child
            update_sum(u, sign=-1)
            u = edge.parent
            while u != -1:
                update_sum(u, sign=-1)
                state[u] -= state[edge.child]
                summary[u] = polarised_summary(u)
                update_sum(u, sign=+1)
                u = parent[u]
            parent[edge.child] = -1
            for tw in range(num_time_windows):
                branch_length[tw, edge.child] = 0

        for edge in edges_in:
            parent[edge.child] = edge.parent
            for tw in range(num_time_windows):
                branch_length[tw, edge.child] = min(
                    time[edge.parent], time_windows[tw + 1]
                ) - max(time[edge.child], time_windows[tw])
            u = edge.child
            update_sum(u, sign=+1)
            u = edge.parent
            while u != -1:
                update_sum(u, sign=-1)
                state[u] += state[edge.child]
                summary[u] = polarised_summary(u)
                update_sum(u, sign=+1)
                u = parent[u]

        # Update the windows
        assert window_index < num_windows
        while windows[window_index] < t_right:
            w_left = windows[window_index]
            w_right = windows[window_index + 1]
            left = max(t_left, w_left)
            right = min(t_right, w_right)
            span = right - left
            assert span > 0
            time_window_index = 0
            while time_window_index < num_time_windows:
                result[window_index, time_window_index] += (
                    running_sum[time_window_index] * span
                )
                time_window_index += 1
            if w_right <= t_right:
                window_index += 1
            else:
                # This interval crosses a tree boundary, so we update it again in the
                # for the next tree
                break

    # print("window_index:", window_index, windows.shape)
    assert window_index == windows.shape[0] - 1
    if drop_time_windows:
        assert result.ndim == 3
        result = result[:, 0]
    if span_normalise:
        for j in range(num_windows):
            result[j] /= windows[j + 1] - windows[j]
    return result


##############################
# Site general stat algorithms
##############################


def windowed_sitewise_stat(ts, sigma, windows, span_normalise=True):
    M = sigma.shape[1]
    A = np.zeros((len(windows) - 1, M))
    window = 0
    for site in ts.sites():
        while windows[window + 1] <= site.position:
            window += 1
        assert windows[window] <= site.position < windows[window + 1]
        A[window] += sigma[site.id]
    if span_normalise:
        diff = np.zeros((A.shape[0], 1))
        diff[:, 0] = np.diff(windows).T
        A /= diff
    return A


def naive_site_general_stat(
    ts, W, f, windows=None, polarised=False, span_normalise=True
):
    n, K = W.shape
    # Hack to determine M
    M = len(f(W[0]))
    sigma = np.zeros((ts.num_sites, M))
    for tree in ts.trees():
        X = np.zeros((ts.num_nodes, K))
        X[ts.samples()] = W
        for u in tree.nodes(order="postorder"):
            for v in tree.children(u):
                X[u] += X[v]
        for site in tree.sites():
            state_map = collections.defaultdict(functools.partial(np.zeros, K))
            state_map[site.ancestral_state] = sum(X[root] for root in tree.roots)
            for mutation in site.mutations:
                state_map[mutation.derived_state] += X[mutation.node]
                if mutation.parent != tskit.NULL:
                    parent = site.mutations[mutation.parent - site.mutations[0].id]
                    state_map[parent.derived_state] -= X[mutation.node]
                else:
                    state_map[site.ancestral_state] -= X[mutation.node]
            if polarised:
                del state_map[site.ancestral_state]
            sigma[site.id] += sum(map(f, state_map.values()))
    return windowed_sitewise_stat(
        ts, sigma, ts.parse_windows(windows), span_normalise=span_normalise
    )


def site_general_stat(
    ts,
    sample_weights,
    summary_func,
    windows=None,
    time_windows=None,
    polarised=False,
    span_normalise=True,
):
    """
    Problem: 'sites' is different that the other windowing options
    because if we output by site we don't want to normalize by length of the window.
    Solution: we pass an argument "normalize", to the windowing function.
    """
    assert time_windows is None
    windows = ts.parse_windows(windows)
    num_windows = windows.shape[0] - 1
    n, state_dim = sample_weights.shape
    # Determine result_dim
    (result_dim,) = summary_func(sample_weights[0]).shape
    result = np.zeros((num_windows, result_dim))
    state = np.zeros((ts.num_nodes, state_dim))
    state[ts.samples()] = sample_weights
    total_weight = np.sum(sample_weights, axis=0)

    site_index = 0
    mutation_index = 0
    window_index = 0
    sites = ts.tables.sites
    mutations = ts.tables.mutations
    parent = np.zeros(ts.num_nodes, dtype=np.int32) - 1
    for (left, right), edges_out, edges_in in ts.edge_diffs():
        for edge in edges_out:
            u = edge.parent
            while u != -1:
                state[u] -= state[edge.child]
                u = parent[u]
            parent[edge.child] = -1
        for edge in edges_in:
            parent[edge.child] = edge.parent
            u = edge.parent
            while u != -1:
                state[u] += state[edge.child]
                u = parent[u]
        while site_index < len(sites) and sites.position[site_index] < right:
            assert left <= sites.position[site_index]
            ancestral_state = sites[site_index].ancestral_state
            allele_state = collections.defaultdict(
                functools.partial(np.zeros, state_dim)
            )
            allele_state[ancestral_state][:] = total_weight
            while (
                mutation_index < len(mutations)
                and mutations[mutation_index].site == site_index
            ):
                mutation = mutations[mutation_index]
                allele_state[mutation.derived_state] += state[mutation.node]
                if mutation.parent != -1:
                    parent_allele = mutations[mutation.parent].derived_state
                    allele_state[parent_allele] -= state[mutation.node]
                else:
                    allele_state[ancestral_state] -= state[mutation.node]
                mutation_index += 1
            if polarised:
                del allele_state[ancestral_state]

            pos = sites.position[site_index]
            while windows[window_index + 1] <= pos:
                window_index += 1
            assert windows[window_index] <= pos < windows[window_index + 1]
            site_result = result[window_index]

            for _allele, value in allele_state.items():
                site_result += summary_func(value)
            site_index += 1
    if span_normalise:
        for j in range(num_windows):
            span = windows[j + 1] - windows[j]
            result[j] /= span
    return result


##############################
# Node general stat algorithms
##############################


def naive_node_general_stat(
    ts, W, f, windows=None, polarised=False, span_normalise=True
):
    windows = ts.parse_windows(windows)
    n, K = W.shape
    M = f(W[0]).shape[0]
    total = np.sum(W, axis=0)
    sigma = np.zeros((ts.num_trees, ts.num_nodes, M))
    for tree in ts.trees():
        X = np.zeros((ts.num_nodes, K))
        X[ts.samples()] = W
        for u in tree.nodes(order="postorder"):
            for v in tree.children(u):
                X[u] += X[v]
        s = np.zeros((ts.num_nodes, M))
        for u in range(ts.num_nodes):
            s[u] = f(X[u])
            if not polarised:
                s[u] += f(total - X[u])
        sigma[tree.index] = s * tree.span
    return windowed_tree_stat(ts, sigma, windows, span_normalise=span_normalise)


def node_general_stat(
    ts,
    sample_weights,
    summary_func,
    windows=None,
    time_windows=None,
    polarised=False,
    span_normalise=True,
):
    """
    Efficient implementation of the algorithm used as the basis for the
    underlying C version.
    """
    assert time_windows is None
    n, state_dim = sample_weights.shape
    windows = ts.parse_windows(windows)
    num_windows = windows.shape[0] - 1
    result_dim = summary_func(sample_weights[0]).shape[0]
    result = np.zeros((num_windows, ts.num_nodes, result_dim))
    state = np.zeros((ts.num_nodes, state_dim))
    state[ts.samples()] = sample_weights
    total_weight = np.sum(sample_weights, axis=0)

    def node_summary(u):
        s = summary_func(state[u])
        if not polarised:
            s += summary_func(total_weight - state[u])
        return s

    window_index = 0
    parent = np.zeros(ts.num_nodes, dtype=np.int32) - 1
    # contains summary_func(state[u]) for each node
    current_values = np.zeros((ts.num_nodes, result_dim))
    for u in range(ts.num_nodes):
        current_values[u] = node_summary(u)
    # contains the location of the last time we updated the output for a node.
    last_update = np.zeros((ts.num_nodes, 1))
    for (t_left, t_right), edges_out, edges_in in ts.edge_diffs():
        for edge in edges_out:
            u = edge.child
            v = edge.parent
            while v != -1:
                result[window_index, v] += (t_left - last_update[v]) * current_values[v]
                last_update[v] = t_left
                state[v] -= state[u]
                current_values[v] = node_summary(v)
                v = parent[v]
            parent[u] = -1

        for edge in edges_in:
            u = edge.child
            v = edge.parent
            parent[u] = v
            while v != -1:
                result[window_index, v] += (t_left - last_update[v]) * current_values[v]
                last_update[v] = t_left
                state[v] += state[u]
                current_values[v] = node_summary(v)
                v = parent[v]

        # Update the windows
        while window_index < num_windows and windows[window_index + 1] <= t_right:
            w_right = windows[window_index + 1]
            # Flush the contribution of all nodes to the current window.
            for u in range(ts.num_nodes):
                result[window_index, u] += (w_right - last_update[u]) * current_values[
                    u
                ]
                last_update[u] = w_right
            window_index += 1

    assert window_index == windows.shape[0] - 1
    if span_normalise:
        for j in range(num_windows):
            result[j] /= windows[j + 1] - windows[j]
    return result


def general_stat(
    ts,
    sample_weights,
    summary_func,
    windows=None,
    time_windows=None,
    polarised=False,
    mode="site",
    span_normalise=True,
):
    """
    General iterface for algorithms above. Directly corresponds to the interface
    for TreeSequence.general_stat.
    """
    method_map = {
        "site": site_general_stat,
        "node": node_general_stat,
        "branch": branch_general_stat,
    }
    return method_map[mode](
        ts,
        sample_weights,
        summary_func,
        windows=windows,
        time_windows=time_windows,
        polarised=polarised,
        span_normalise=span_normalise,
    )


def upper_tri_to_matrix(x):
    """
    Given x, a vector of entries of the upper triangle of a matrix
    in row-major order, including the diagonal, return the corresponding matrix.
    """
    # n^2 + n = 2 u => n = (-1 + sqrt(1 + 8*u))/2
    n = int((np.sqrt(1 + 8 * len(x)) - 1) / 2.0)
    out = np.ones((n, n))
    k = 0
    for i in range(n):
        for j in range(i, n):
            out[i, j] = out[j, i] = x[k]
            k += 1
    return out


##################################
# Test cases
##################################


class StatsTestCase:
    """
    Provides convenience functions.
    """

    def assertListAlmostEqual(self, x, y):
        assert len(x) == len(y)
        for a, b in zip(x, y):
            self.assertAlmostEqual(a, b)

    def assertArrayEqual(self, x, y):
        nt.assert_equal(x, y)

    def assertArrayAlmostEqual(self, x, y, atol=1e-6, rtol=1e-7):
        nt.assert_allclose(x, y, atol=atol, rtol=rtol)

    def identity_f(self, ts):
        return lambda x: x * (x < ts.num_samples)

    def cumsum_f(self, ts):
        return lambda x: np.cumsum(x) * (x < ts.num_samples)

    def sum_f(self, ts, k=1):
        return lambda x: np.array([sum(x) * (sum(x) < 2 * ts.num_samples)] * k)


class TopologyExamplesMixin:
    """
    Defines a set of test cases on different example tree sequence topologies.
    Derived classes need to define a 'verify' function which will perform the
    actual tests.
    """

    def test_single_tree_sequence_length(self):
        ts = msprime.simulate(6, length=10, random_seed=1)
        self.verify(ts)

    def test_single_tree_multiple_roots(self):
        ts = msprime.simulate(8, random_seed=1, end_time=0.5)
        assert ts.first().num_roots > 1
        self.verify(ts)

    def test_many_trees(self, ts_4_recomb_fixture):
        ts = ts_4_recomb_fixture
        assert ts.num_trees > 2
        self.verify(ts)

    # @pytest.mark.skip(reason="Skipping short sequence length test")
    def test_short_sequence_length(self):
        ts = msprime.simulate(6, length=0.5, recombination_rate=2, random_seed=1)
        assert ts.num_trees > 2
        self.verify(ts)

    @pytest.mark.slow
    def test_wright_fisher_unsimplified(self, wf_sim_fixture):
        self.verify(wf_sim_fixture["unsimplified"])

    @pytest.mark.slow
    def test_wright_fisher_initial_generation(self, wf_sim_fixture):
        self.verify(wf_sim_fixture["initial_generation"])

    def test_wright_fisher_initial_generation_no_deep_history(self, wf_sim_fixture):
        self.verify(wf_sim_fixture["no_deep_history"])

    def test_wright_fisher_unsimplified_multiple_roots(self, wf_sim_fixture):
        self.verify(wf_sim_fixture["unsimplified_multi_roots"])

    def test_wright_fisher_simplified(self, wf_sim_fixture):
        self.verify(wf_sim_fixture["simplified"])

    def test_wright_fisher_simplified_multiple_roots(self, wf_sim_fixture):
        self.verify(wf_sim_fixture["simplified_multi_roots"])

    def test_empty_ts(self):
        tables = tskit.TableCollection(1.0)
        tables.nodes.add_row(1, 0)
        tables.nodes.add_row(1, 0)
        tables.nodes.add_row(1, 0)
        tables.nodes.add_row(1, 0)
        ts = tables.tree_sequence()
        self.verify(ts)

    def test_non_sample_ancestry(self):
        tables = tskit.TableCollection(1.0)
        tables.nodes.add_row(1, 0)
        tables.nodes.add_row(1, 0)
        tables.nodes.add_row(0, 1)
        tables.nodes.add_row(0, 0)  # 3 is a leaf but not a sample.
        # Make sure we have 4 samples for the tests.
        tables.nodes.add_row(1, 1)
        tables.nodes.add_row(1, 1)
        tables.edges.add_row(0, 1, 2, 0)
        tables.edges.add_row(0, 1, 2, 1)
        tables.edges.add_row(0, 1, 4, 3)
        ts = tables.tree_sequence()
        self.verify(ts)


# Fixtures for commonly used simulations in test_tree_stats.py
# Naming convention: ts_{num_samples}_{features}_fixture
# Features: mut (mutations), recomb (recombination), highmut/highrecomb (high rates)


@pytest.fixture(scope="session")
def ts_6_fixture():
    """Basic 6-sample tree sequence, no mutations or recombination."""
    return msprime.simulate(6, random_seed=1)


@pytest.fixture(scope="session")
def ts_10_recomb_fixture():
    """10-sample tree sequence with recombination (used 3+ times)."""
    return msprime.simulate(10, recombination_rate=1, random_seed=2)


@pytest.fixture(scope="session")
def ts_10_mut_fixture():
    """10-sample tree sequence with mutations (used 10 times)."""
    return msprime.simulate(10, mutation_rate=1, random_seed=1)


@pytest.fixture(scope="session")
def ts_10_mut_recomb_fixture():
    """10-sample tree sequence with mutations and recombination (used 5+ times)."""
    return msprime.simulate(10, mutation_rate=1, recombination_rate=2, random_seed=1)


@pytest.fixture(scope="session")
def ts_4_recomb_fixture():
    """4-sample tree sequence with recombination (used 4+ times)."""
    return msprime.simulate(4, recombination_rate=1, random_seed=2)


@pytest.fixture(scope="session")
def ts_12_highrecomb_fixture():
    """12-sample tree sequence with high recombination (used 4+ times)."""
    return msprime.simulate(12, recombination_rate=3, random_seed=2)


@pytest.fixture(scope="session")
def ts_44_recomb_fixture():
    """44-sample tree sequence with recombination (used 2 times)."""
    return msprime.simulate(44, recombination_rate=1, random_seed=2)


@pytest.fixture(scope="session")
def ts_ancestry_10_fixture():
    """Standard ancestry simulation for 10 samples."""
    return msprime.sim_ancestry(10, random_seed=1, sequence_length=10)


@pytest.fixture(scope="session")
def ts_6_length_factory_fixture():
    """Factory fixture for 6-sample tree sequences with variable length."""

    def _make_ts(length):
        return msprime.simulate(
            6, length=length, recombination_rate=2, mutation_rate=1, random_seed=1
        )

    return _make_ts


# Wright-Fisher simulation fixtures
@pytest.fixture(scope="session")
def wf_sim_fixture():
    """Common Wright-Fisher simulations used across test classes."""
    # Pre-compute all common WF simulations
    simulations = {}

    # Used in TopologyExamplesMixin tests
    tables = wf.wf_sim(
        4, 5, seed=1, deep_history=True, initial_generation_samples=False, num_loci=5
    )
    tables.sort()
    simulations["unsimplified"] = tables.tree_sequence()

    tables = wf.wf_sim(
        6, 5, seed=3, deep_history=True, initial_generation_samples=True, num_loci=2
    )
    tables.sort()
    tables.simplify()
    simulations["initial_generation"] = tables.tree_sequence()

    tables = wf.wf_sim(
        6, 15, seed=202, deep_history=False, initial_generation_samples=True, num_loci=5
    )
    tables.sort()
    tables.simplify()
    simulations["no_deep_history"] = tables.tree_sequence()

    tables = wf.wf_sim(
        6, 5, seed=1, deep_history=False, initial_generation_samples=False, num_loci=4
    )
    tables.sort()
    simulations["unsimplified_multi_roots"] = tables.tree_sequence()

    tables = wf.wf_sim(
        5, 8, seed=1, deep_history=True, initial_generation_samples=False, num_loci=5
    )
    tables.sort()
    simulations["simplified"] = tables.tree_sequence().simplify()

    tables = wf.wf_sim(
        6, 8, seed=1, deep_history=False, initial_generation_samples=False, num_loci=3
    )
    tables.sort()
    simulations["simplified_multi_roots"] = tables.tree_sequence().simplify()

    return simulations


@pytest.fixture(scope="session")
def wf_mut_sim_fixture():
    """Wright-Fisher simulations with mutations for MutatedTopologyExamplesMixin."""
    simulations = {}

    # With mutations for site-based tests
    tables = wf.wf_sim(
        4, 5, seed=1, deep_history=True, initial_generation_samples=False, num_loci=10
    )
    tables.sort()
    ts = msprime.mutate(tables.tree_sequence(), rate=0.05, random_seed=234)
    simulations["unsimplified"] = ts

    tables = wf.wf_sim(
        6, 5, seed=3, deep_history=True, initial_generation_samples=True, num_loci=2
    )
    tables.sort()
    tables.simplify()
    ts = msprime.mutate(tables.tree_sequence(), rate=0.08, random_seed=2)
    simulations["initial_generation"] = ts

    tables = wf.wf_sim(
        7, 15, seed=202, deep_history=False, initial_generation_samples=True, num_loci=5
    )
    tables.sort()
    tables.simplify()
    ts = msprime.mutate(tables.tree_sequence(), rate=0.1, random_seed=3)
    simulations["no_deep_history"] = ts

    tables = wf.wf_sim(
        8, 15, seed=1, deep_history=False, initial_generation_samples=False, num_loci=20
    )
    tables.sort()
    ts = msprime.mutate(tables.tree_sequence(), rate=0.01, random_seed=2)
    simulations["unsimplified_multi_roots"] = ts

    tables = wf.wf_sim(
        9, 10, seed=1, deep_history=True, initial_generation_samples=False, num_loci=5
    )
    tables.sort()
    ts = tables.tree_sequence().simplify()
    ts = tsutil.jukes_cantor(ts, 10, 0.01, seed=1)
    simulations["simplified"] = ts

    return simulations


class MutatedTopologyExamplesMixin:
    """
    Defines a set of test cases on different example tree sequence topologies.
    Derived classes need to define a 'verify' function which will perform the
    actual tests.
    """

    def test_single_tree_no_sites(self, ts_6_fixture):
        ts = ts_6_fixture
        assert ts.num_sites == 0
        self.verify(ts)

    def test_ghost_allele(self):
        tables = tskit.TableCollection(1)
        tables.nodes.add_row(flags=1, time=0)
        tables.nodes.add_row(flags=1, time=0)
        tables.nodes.add_row(flags=0, time=1)
        tables.edges.add_row(0, 1, 2, 0)
        tables.edges.add_row(0, 1, 2, 1)
        tables.sites.add_row(position=0.5, ancestral_state="A")
        # Make sure there's 4 samples
        tables.nodes.add_row(flags=1, time=0)
        tables.nodes.add_row(flags=1, time=0)
        # The ghost mutation that's never seen in the genotypes
        tables.mutations.add_row(site=0, node=0, derived_state="T")
        tables.mutations.add_row(site=0, node=0, derived_state="G", parent=0)
        ts = tables.tree_sequence()
        self.verify(ts)

    def test_ghost_allele_all_ancestral(self):
        tables = tskit.TableCollection(1)
        tables.nodes.add_row(flags=1, time=0)
        tables.nodes.add_row(flags=1, time=0)
        tables.nodes.add_row(flags=0, time=1)
        # Make sure there's 4 samples
        tables.nodes.add_row(flags=1, time=0)
        tables.nodes.add_row(flags=1, time=0)
        tables.edges.add_row(0, 1, 2, 0)
        tables.edges.add_row(0, 1, 2, 1)
        tables.sites.add_row(position=0.5, ancestral_state="A")
        tables.mutations.add_row(site=0, node=0, derived_state="T")
        # Mutate back to the ancestral state so that all genotypes are zero
        tables.mutations.add_row(site=0, node=0, derived_state="A", parent=0)
        ts = tables.tree_sequence()
        self.verify(ts)

    def test_non_sample_ancestry(self):
        # 2.00┊       5   ┊
        #     ┊    ┏━━┻━┓ ┊
        # 1.00┊    4    ┃ ┊
        #     ┊ ┏━┳┻┳━┓ ┃ ┊
        # 0.00┊ 0 1 2 3 6 ┊
        #    0.00        1.00
        tables = tskit.TableCollection(1)
        # Four sample nodes
        for j in range(4):
            tables.nodes.add_row(flags=1, time=0)
            tables.edges.add_row(0, 1, 4, j)
        # Their MRCA, 4, joins to older ancestor 5
        tables.nodes.add_row(flags=0, time=1)
        tables.nodes.add_row(flags=0, time=2)
        tables.edges.add_row(0, 1, 5, 4)
        # Which has non-sample leaf at time 0
        tables.nodes.add_row(flags=0, time=0)
        tables.edges.add_row(0, 1, 5, 6)
        # Two sites with mutations. One over the MRCA of the
        # samples so it's fixed at 1 and one over the non sample
        # leaf so that samples are fixed at zero.
        tables.sites.add_row(position=0.25, ancestral_state="0")
        tables.sites.add_row(position=0.5, ancestral_state="0")
        tables.mutations.add_row(site=0, node=4, derived_state="1")
        tables.mutations.add_row(site=1, node=6, derived_state="1")
        ts = tables.tree_sequence()
        self.verify(ts)

    def test_single_tree_infinite_sites(self, ts_10_mut_fixture):
        ts = ts_10_mut_fixture
        assert ts.num_sites > 0
        self.verify(ts)

    def test_single_tree_sites_no_mutations(self, ts_6_fixture):
        ts = ts_6_fixture
        tables = ts.dump_tables()
        tables.sites.add_row(0.1, "a")
        tables.sites.add_row(0.2, "aaa")
        self.verify(tables.tree_sequence())

    @pytest.mark.slow
    def test_single_tree_jukes_cantor(self, ts_10_mut_fixture):
        ts = ts_10_mut_fixture
        ts = tsutil.jukes_cantor(ts, 20, 1, seed=10)
        self.verify(ts)

    def test_single_tree_single_site_many_silent(self, ts_6_fixture):
        ts = ts_6_fixture
        ts = tsutil.jukes_cantor(ts, 1, 20, seed=10)
        self.verify(ts)

    def test_single_tree_multichar_mutations(self, ts_10_mut_fixture):
        ts = ts_10_mut_fixture
        ts = tsutil.insert_multichar_mutations(ts)
        self.verify(ts)

    def test_many_trees_infinite_sites(self, ts_10_mut_recomb_fixture):
        ts = ts_10_mut_recomb_fixture
        assert ts.num_sites > 0
        assert ts.num_trees > 2
        self.verify(ts)

    @pytest.mark.slow
    def test_many_trees_sequence_length_infinite_sites(
        self, ts_6_length_factory_fixture
    ):
        for L in [0.5, 1.5, 3.3333]:
            ts = ts_6_length_factory_fixture(L)
            self.verify(ts)

    def test_wright_fisher_unsimplified(self, wf_mut_sim_fixture):
        ts = wf_mut_sim_fixture["unsimplified"]
        assert ts.num_sites > 0
        self.verify(ts)

    def test_wright_fisher_initial_generation(self, wf_mut_sim_fixture):
        ts = wf_mut_sim_fixture["initial_generation"]
        assert ts.num_sites > 0
        self.verify(ts)

    def test_wright_fisher_initial_generation_no_deep_history(self, wf_mut_sim_fixture):
        ts = wf_mut_sim_fixture["no_deep_history"]
        assert ts.num_sites > 0
        self.verify(ts)

    def test_wright_fisher_unsimplified_multiple_roots(self, wf_mut_sim_fixture):
        ts = wf_mut_sim_fixture["unsimplified_multi_roots"]
        assert ts.num_sites > 0
        self.verify(ts)

    def test_wright_fisher_simplified(self, wf_mut_sim_fixture):
        ts = wf_mut_sim_fixture["simplified"]
        assert ts.num_sites > 0
        self.verify(ts)

    def test_empty_ts(self):
        tables = tskit.TableCollection(1.0)
        for _ in range(10):
            tables.nodes.add_row(tskit.NODE_IS_SAMPLE, 0)
        ts = tables.tree_sequence()
        self.verify(ts)


def example_sample_sets(ts, min_size=1):
    """
    Generate a series of example sample sets from the specfied tree sequence. The
    number of sample sets returned in each example must be at least min_size
    """
    samples = ts.samples()
    np.random.shuffle(samples)
    splits = np.array_split(samples, min_size)
    yield splits
    yield [[s] for s in samples]
    if min_size == 1:
        yield [samples[:1]]
    if ts.num_samples > 2 and min_size <= 2:
        yield [samples[:2], samples[2:]]
    if ts.num_samples > 4 and min_size <= 2:
        yield [samples[:2], samples[2:4]]
    if ts.num_samples > 7 and min_size <= 4:
        yield [samples[:2], samples[2:4], samples[4:6], samples[6:]]
    if ts.num_samples > 8 and min_size <= 4:
        yield [samples[:2], samples[2:4], samples[4:6], samples[6:8]]


def example_sample_set_index_pairs(sample_sets):
    assert len(sample_sets) >= 2
    yield [(0, 1)]
    yield [(1, 0), (0, 1)]
    if len(sample_sets) > 2:
        yield [(0, 1), (1, 2), (0, 2)]


def example_sample_set_index_triples(sample_sets):
    assert len(sample_sets) >= 3
    yield [(0, 1, 2)]
    yield [(0, 2, 1), (2, 1, 0)]
    if len(sample_sets) > 3:
        yield [(3, 0, 1), (0, 2, 3), (1, 2, 3)]


def example_sample_set_index_quads(sample_sets):
    assert len(sample_sets) >= 4
    yield [(0, 1, 2, 3)]
    yield [(0, 1, 2, 3), (3, 2, 1, 0)]
    yield [(0, 1, 2, 3), (3, 2, 1, 0), (1, 2, 3, 0)]


def example_windows(ts):
    """
    Generate a series of example windows for the specified tree sequence.
    """
    L = ts.sequence_length
    yield "sites"
    yield "trees"
    yield [0, L]
    yield ts.breakpoints(as_array=True)
    yield np.linspace(0, L, num=10)


class WeightStatsMixin:
    """
    Implements the verify method and dispatches it to verify_weighted_stat
    for a representative set of sample sets and windows.
    """

    def example_weights(self, ts, min_size=1):
        """
        Generate a series of example weights from the specfied tree sequence.
        """
        np.random.seed(46)
        # Reduced to 3 essential weight matrices for performance

        # 1. Simple weights with variance (k=min_size)
        k = min_size
        W = np.ones((ts.num_samples, k))
        # Ensure positive variance: different values for different samples
        for i in range(ts.num_samples):
            W[i, :] = 1.0 + i * 0.1
        yield W

        # 2. Exponential weights with k=min_size+1 (medium complexity)
        k = min_size + 1
        W = np.zeros((ts.num_samples, k))
        for j in range(k):
            W[:, j] = np.random.exponential(1, ts.num_samples)
        yield W

        # 3. Mixed weights with larger k (complex case)
        # Only test larger k if samples allow it and keep it reasonable
        k = min(min_size + 3, ts.num_samples)  # Reduced from +5 to +3
        if k > min_size + 1:
            W = np.zeros((ts.num_samples, k))
            # First column: linear gradient for variance
            W[:, 0] = np.linspace(0.5, 1.5, ts.num_samples)
            # Remaining columns: exponential
            for j in range(1, k):
                W[:, j] = np.random.exponential(0.8, ts.num_samples)
            yield W

    def transform_weights(self, W):
        """
        Specific methods will need to transform weights
        before passing them to general_stat.
        """
        return W

    def verify(self, ts):
        for W, windows in subset_combos(
            self.example_weights(ts), example_windows(ts), p=0.1
        ):
            self.verify_weighted_stat(ts, W, windows=windows)

    def verify_definition(self, ts, W, windows, summary_func, ts_method, definition):
        # general_stat will need an extra column for p
        gW = self.transform_weights(W)

        def wrapped_summary_func(x):
            with suppress_division_by_zero_warning():
                return summary_func(x)

        # Determine output_dim of the function
        M = len(wrapped_summary_func(gW[0]))
        for sn in [True, False]:
            sigma1 = ts.general_stat(
                gW, wrapped_summary_func, M, windows, mode=self.mode, span_normalise=sn
            )
            sigma2 = general_stat(
                ts, gW, wrapped_summary_func, windows, mode=self.mode, span_normalise=sn
            )
            sigma3 = ts_method(W, windows=windows, mode=self.mode, span_normalise=sn)
            sigma4 = definition(
                ts, W, windows=windows, mode=self.mode, span_normalise=sn
            )

            assert sigma1.shape == sigma2.shape
            assert sigma1.shape == sigma3.shape
            assert sigma1.shape == sigma4.shape
            self.assertArrayAlmostEqual(sigma1, sigma2)
            self.assertArrayAlmostEqual(sigma1, sigma3)
            self.assertArrayAlmostEqual(sigma1, sigma4)


class SampleSetStatsMixin:
    """
    Implements the verify method and dispatches it to verify_sample_sets
    for a representative set of sample sets and windows.
    """

    def verify(self, ts):
        for sample_sets, windows in subset_combos(
            example_sample_sets(ts),
            example_windows(ts),
            p=0.2,
        ):
            self.verify_sample_sets(ts, sample_sets, windows=windows)

    def verify_definition(
        self, ts, sample_sets, windows, summary_func, ts_method, definition
    ):
        W = np.array([[u in A for A in sample_sets] for u in ts.samples()], dtype=float)

        def wrapped_summary_func(x):
            with suppress_division_by_zero_warning():
                return summary_func(x)

        for sn in [True, False]:
            # Determine output_dim of the function
            M = len(summary_func(W[0]))
            sigma1 = ts.general_stat(
                W, wrapped_summary_func, M, windows, mode=self.mode, span_normalise=sn
            )
            sigma2 = general_stat(
                ts, W, wrapped_summary_func, windows, mode=self.mode, span_normalise=sn
            )
            sigma3 = ts_method(
                sample_sets, windows=windows, mode=self.mode, span_normalise=sn
            )
            sigma4 = definition(
                ts, sample_sets, windows=windows, mode=self.mode, span_normalise=sn
            )

            assert sigma1.shape == sigma2.shape
            assert sigma1.shape == sigma3.shape
            assert sigma1.shape == sigma4.shape
            self.assertArrayAlmostEqual(sigma1, sigma2)
            self.assertArrayAlmostEqual(sigma1, sigma3)
            self.assertArrayAlmostEqual(sigma1, sigma4)


class KWaySampleSetStatsMixin(SampleSetStatsMixin):
    """
    Defines the verify definition method, which comparse the results from
    several different ways of defining and computing the same statistic.
    """

    def verify_definition(
        self, ts, sample_sets, indexes, windows, summary_func, ts_method, definition
    ):
        def wrapped_summary_func(x):
            with suppress_division_by_zero_warning():
                return summary_func(x)

        W = np.array([[u in A for A in sample_sets] for u in ts.samples()], dtype=float)
        # Determine output_dim of the function
        M = len(wrapped_summary_func(W[0]))
        sigma1 = ts.general_stat(W, wrapped_summary_func, M, windows, mode=self.mode)
        sigma2 = general_stat(ts, W, wrapped_summary_func, windows, mode=self.mode)
        sigma3 = ts_method(
            sample_sets, indexes=indexes, windows=windows, mode=self.mode
        )
        sigma4 = definition(
            ts, sample_sets, indexes=indexes, windows=windows, mode=self.mode
        )

        assert sigma1.shape == sigma2.shape
        assert sigma1.shape == sigma3.shape
        assert sigma1.shape == sigma4.shape
        self.assertArrayAlmostEqual(sigma1, sigma2)
        self.assertArrayAlmostEqual(sigma1, sigma3)
        self.assertArrayAlmostEqual(sigma1, sigma4)


class TwoWaySampleSetStatsMixin(KWaySampleSetStatsMixin):
    """
    Implements the verify method and dispatches it to verify_sample_sets_indexes,
    which gives a representative sample of sample set indexes.
    """

    def verify(self, ts):
        for sample_sets, windows in subset_combos(
            example_sample_sets(ts, min_size=2),
            example_windows(ts),
            p=0.1,
        ):
            for indexes in example_sample_set_index_pairs(sample_sets):
                self.verify_sample_sets_indexes(ts, sample_sets, indexes, windows)


class ThreeWaySampleSetStatsMixin(KWaySampleSetStatsMixin):
    """
    Implements the verify method and dispatches it to verify_sample_sets_indexes,
    which gives a representative sample of sample set indexes.
    """

    def verify(self, ts):
        for sample_sets, windows in subset_combos(
            example_sample_sets(ts, min_size=3),
            example_windows(ts),
            p=0.1,
        ):
            for indexes in example_sample_set_index_triples(sample_sets):
                self.verify_sample_sets_indexes(ts, sample_sets, indexes, windows)


class FourWaySampleSetStatsMixin(KWaySampleSetStatsMixin):
    """
    Implements the verify method and dispatches it to verify_sample_sets_indexes,
    which gives a representative sample of sample set indexes.
    """

    def verify(self, ts):
        for sample_sets, windows in subset_combos(
            example_sample_sets(ts, min_size=4),
            example_windows(ts),
            p=0.1,
        ):
            for indexes in example_sample_set_index_quads(sample_sets):
                self.verify_sample_sets_indexes(ts, sample_sets, indexes, windows)


############################################
# Diversity
############################################


def site_diversity(ts, sample_sets, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(sample_sets)))
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        haps = ts.genotype_matrix(isolated_as_missing=False).T
        site_positions = [x.position for x in ts.sites()]
        for i, X in enumerate(sample_sets):
            S = 0
            site_in_window = False
            denom = np.float64(len(X) * (len(X) - 1))
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    site_in_window = True
                    for x in X:
                        for y in set(X) - {x}:
                            x_index = np.where(samples == x)[0][0]
                            y_index = np.where(samples == y)[0][0]
                            if haps[x_index][k] != haps[y_index][k]:
                                # x|y
                                S += 1
            if site_in_window:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def branch_diversity(ts, sample_sets, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(sample_sets)))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, X in enumerate(sample_sets):
            S = 0
            denom = np.float64(len(X) * (len(X) - 1))
            has_trees = False
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                if tr.total_branch_length > 0:
                    has_trees = True
                SS = 0
                for x in X:
                    for y in set(X) - {x}:
                        SS += path_length(tr, x, y)
                S += SS * (min(end, tr.interval.right) - max(begin, tr.interval.left))
            if has_trees:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def node_diversity(ts, sample_sets, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    K = len(sample_sets)
    out = np.zeros((len(windows) - 1, ts.num_nodes, K))
    for k in range(K):
        X = sample_sets[k]
        for j in range(len(windows) - 1):
            begin = windows[j]
            end = windows[j + 1]
            tX = len(X)
            denom = np.float64(len(X) * (len(X) - 1))
            S = np.zeros(ts.num_nodes)
            for tr in ts.trees(tracked_samples=X):
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in tr.nodes():
                    # count number of pairwise paths going through u
                    n = tr.num_tracked_samples(u)
                    SS[u] += 2 * n * (tX - n)
                S += SS * (min(end, tr.interval.right) - max(begin, tr.interval.left))
            with suppress_division_by_zero_warning():
                out[j, :, k] = S / denom
            if span_normalise:
                out[j, :, k] /= end - begin
    return out


def diversity(ts, sample_sets, windows=None, mode="site", span_normalise=True):
    """
    Computes average pairwise diversity between two random choices from x
    over the window specified.
    """
    method_map = {
        "site": site_diversity,
        "node": node_diversity,
        "branch": branch_diversity,
    }
    return method_map[mode](
        ts, sample_sets, windows=windows, span_normalise=span_normalise
    )


class TestDiversity(StatsTestCase, SampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_sample_sets(self, ts, sample_sets, windows):
        n = np.array([len(x) for x in sample_sets])

        def f(x):
            with np.errstate(invalid="ignore", divide="ignore"):
                return x * (n - x) / (n * (n - 1))

        self.verify_definition(ts, sample_sets, windows, f, ts.diversity, diversity)


class TestBranchDiversity(TestDiversity, TopologyExamplesMixin):
    mode = "branch"


class TestNodeDiversity(TestDiversity, TopologyExamplesMixin):
    mode = "node"


class TestSiteDiversity(TestDiversity, MutatedTopologyExamplesMixin):
    mode = "site"


############################################
# Segregating sites
############################################


def site_segregating_sites(ts, sample_sets, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(sample_sets)))
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        haps = ts.genotype_matrix(isolated_as_missing=False)
        site_positions = [x.position for x in ts.sites()]
        for i, X in enumerate(sample_sets):
            set_X = set(X)
            X_index = np.where(np.fromiter((s in set_X for s in samples), dtype=bool))[
                0
            ]
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    num_alleles = len(set(haps[k, X_index]))
                    out[j][i] += num_alleles - 1
            if span_normalise:
                out[j][i] /= end - begin
    return out


def branch_segregating_sites(ts, sample_sets, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(sample_sets)))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, X in enumerate(sample_sets):
            tX = len(X)
            for tr in ts.trees(tracked_samples=X):
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                SS = 0
                for u in tr.nodes():
                    nX = tr.num_tracked_samples(u)
                    if nX > 0 and nX < tX:
                        SS += tr.branch_length(u)
                out[j][i] += SS * (
                    min(end, tr.interval.right) - max(begin, tr.interval.left)
                )
            if span_normalise:
                out[j][i] /= end - begin
    return out


def node_segregating_sites(ts, sample_sets, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    K = len(sample_sets)
    out = np.zeros((len(windows) - 1, ts.num_nodes, K))
    for k in range(K):
        X = sample_sets[k]
        for j in range(len(windows) - 1):
            begin = windows[j]
            end = windows[j + 1]
            tX = len(X)
            S = np.zeros(ts.num_nodes)
            for tr in ts.trees(tracked_samples=X):
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in tr.nodes():
                    nX = tr.num_tracked_samples(u)
                    SS[u] = (nX > 0) and (nX < tX)
                S += SS * (min(end, tr.interval.right) - max(begin, tr.interval.left))
            out[j, :, k] = S
            if span_normalise:
                out[j, :, k] /= end - begin
    return out


def segregating_sites(ts, sample_sets, windows=None, mode="site", span_normalise=True):
    """
    Computes the density of segregating sites over the window specified.
    """
    method_map = {
        "site": site_segregating_sites,
        "node": node_segregating_sites,
        "branch": branch_segregating_sites,
    }
    return method_map[mode](
        ts, sample_sets, windows=windows, span_normalise=span_normalise
    )


class TestSegregatingSites(StatsTestCase, SampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_sample_sets(self, ts, sample_sets, windows):
        n = np.array([len(x) for x in sample_sets])

        # this works because sum_{i=1}^k (1-p_i) = k-1
        def f(x):
            return (x > 0) * (1 - x / n)

        self.verify_definition(
            ts, sample_sets, windows, f, ts.segregating_sites, segregating_sites
        )


class TestBranchSegregatingSites(TestSegregatingSites, TopologyExamplesMixin):
    mode = "branch"


class TestNodeSegregatingSites(TestSegregatingSites, TopologyExamplesMixin):
    mode = "node"


class TestSiteSegregatingSites(TestSegregatingSites, MutatedTopologyExamplesMixin):
    mode = "site"


class TestBranchSegregatingSitesProperties(StatsTestCase, TopologyExamplesMixin):
    def verify(self, ts):
        windows = ts.breakpoints(as_array=True)
        # If we split by tree, this should always be equal to the total
        # branch length. The definition of total_branch_length here is slightly
        # tricky: it's the sum of all branch lengths that subtend between 0
        # and n samples. This differs from the built-in total_branch_length
        # function, which just sums that total branch length reachable from
        # roots.
        tbl_tree = [
            sum(
                tree.branch_length(u)
                for u in tree.nodes()
                if 0 < tree.num_samples(u) < ts.num_samples
            )
            for tree in ts.trees()
        ]
        # We must span_normalise, because these values are always weighted
        # by the span, so we're effectively cancelling out this contribution
        tbl = ts.segregating_sites([ts.samples()], windows=windows, mode="branch")
        tbl = tbl.reshape(tbl.shape[:-1])
        self.assertArrayAlmostEqual(tbl_tree, tbl)


############################################
# Tajima's D
############################################


def site_tajimas_d(ts, sample_sets, windows=None):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(sample_sets)))
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        haps = ts.genotype_matrix(isolated_as_missing=False)
        site_positions = [x.position for x in ts.sites()]
        n = np.array([len(X) for X in sample_sets])
        for i, X in enumerate(sample_sets):
            nn = n[i]
            S = 0
            T = 0
            set_X = set(X)
            X_index = np.where(np.fromiter((s in set_X for s in samples), dtype=bool))[
                0
            ]
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    hX = haps[k, X_index]
                    alleles = set(hX)
                    num_alleles = len(alleles)
                    n_alleles = [np.sum(hX == a) for a in alleles]
                    S += num_alleles - 1
                    for k in n_alleles:
                        with suppress_division_by_zero_warning():
                            T += k * (nn - k) / (nn * (nn - 1))
            with suppress_division_by_zero_warning():
                a1 = np.sum(1 / np.arange(1, nn))  # this is h in the main version
                a2 = np.sum(1 / np.arange(1, nn) ** 2)  # this is g
                b1 = (nn + 1) / (3 * (nn - 1))
                b2 = 2 * (nn**2 + nn + 3) / (9 * nn * (nn - 1))
                c1 = b1 - 1 / a1
                c2 = b2 - (nn + 2) / (a1 * nn) + a2 / a1**2
                e1 = c1 / a1  # this is a
                e2 = c2 / (a1**2 + a2)  # this is b
                out[j][i] = (T - S / a1) / np.sqrt(e1 * S + e2 * S * (S - 1))
    return out


def tajimas_d(ts, sample_sets, windows=None, mode="site", span_normalise=True):
    method_map = {"site": site_tajimas_d}
    return method_map[mode](
        ts, sample_sets, windows=windows, span_normalise=span_normalise
    )


class TestTajimasD(StatsTestCase, SampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify(self, ts):
        # only check per-site
        for sample_sets in example_sample_sets(ts, min_size=1):
            self.verify_persite_tajimas_d(ts, sample_sets)

    def get_windows(self, ts):
        yield "sites"
        yield [0, ts.sequence_length]
        yield np.arange(0, 1.1, 0.1) * ts.sequence_length

    def verify_persite_tajimas_d(self, ts, sample_sets):
        for windows in self.get_windows(ts):
            sigma1 = ts.Tajimas_D(sample_sets, windows=windows, mode=self.mode)
            sigma2 = site_tajimas_d(ts, sample_sets, windows=windows)
            assert sigma1.shape == sigma2.shape
            self.assertArrayAlmostEqual(sigma1, sigma2)


class TestSiteTajimasD(TestTajimasD, MutatedTopologyExamplesMixin):
    mode = "site"


############################################
# Y1
############################################


def branch_Y1(ts, sample_sets, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(sample_sets)))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, X in enumerate(sample_sets):
            S = 0
            denom = np.float64(len(X) * (len(X) - 1) * (len(X) - 2))
            has_trees = False
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                if tr.total_branch_length > 0:
                    has_trees = True
                this_length = min(end, tr.interval.right) - max(begin, tr.interval.left)
                for x in X:
                    for y in set(X) - {x}:
                        for z in set(X) - {x, y}:
                            xy_mrca = tr.mrca(x, y)
                            xz_mrca = tr.mrca(x, z)
                            yz_mrca = tr.mrca(y, z)
                            if xy_mrca == xz_mrca:
                                #   /\
                                #  / /\
                                # x y  z
                                S += path_length(tr, x, yz_mrca) * this_length
                            elif xy_mrca == yz_mrca:
                                #   /\
                                #  / /\
                                # y x  z
                                S += path_length(tr, x, xz_mrca) * this_length
                            elif xz_mrca == yz_mrca:
                                #   /\
                                #  / /\
                                # z x  y
                                S += path_length(tr, x, xy_mrca) * this_length
            if has_trees:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def site_Y1(ts, sample_sets, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(sample_sets)))
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        haps = ts.genotype_matrix(isolated_as_missing=False).T
        site_positions = [x.position for x in ts.sites()]
        for i, X in enumerate(sample_sets):
            S = 0
            denom = np.float64(len(X) * (len(X) - 1) * (len(X) - 2))
            site_in_window = False
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    site_in_window = True
                    for x in X:
                        x_index = np.where(samples == x)[0][0]
                        for y in set(X) - {x}:
                            y_index = np.where(samples == y)[0][0]
                            for z in set(X) - {x, y}:
                                z_index = np.where(samples == z)[0][0]
                                condition = (
                                    haps[x_index, k] != haps[y_index, k]
                                    and haps[x_index, k] != haps[z_index, k]
                                )
                                if condition:
                                    # x|yz
                                    S += 1
            if site_in_window:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def node_Y1(ts, sample_sets, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    K = len(sample_sets)
    out = np.zeros((len(windows) - 1, ts.num_nodes, K))
    for k in range(K):
        X = sample_sets[k]
        for j in range(len(windows) - 1):
            begin = windows[j]
            end = windows[j + 1]
            tX = len(X)
            denom = np.float64(tX * (tX - 1) * (tX - 2))
            S = np.zeros(ts.num_nodes)
            for tr in ts.trees(tracked_samples=X):
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in tr.nodes():
                    # count number of paths above a but not b,c
                    n = tr.num_tracked_samples(u)
                    SS[u] += n * (tX - n) * (tX - n - 1) + (tX - n) * n * (n - 1)
                S += SS * (min(end, tr.interval.right) - max(begin, tr.interval.left))
            with suppress_division_by_zero_warning():
                out[j, :, k] = S / denom
            if span_normalise:
                out[j, :, k] /= end - begin
    return out


def Y1(ts, sample_sets, windows=None, mode="site", span_normalise=True):
    windows = ts.parse_windows(windows)
    method_map = {"site": site_Y1, "node": node_Y1, "branch": branch_Y1}
    return method_map[mode](
        ts, sample_sets, windows=windows, span_normalise=span_normalise
    )


class TestY1(StatsTestCase, SampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_sample_sets(self, ts, sample_sets, windows):
        n = np.array([len(x) for x in sample_sets])
        denom = n * (n - 1) * (n - 2)

        def f(x):
            with np.errstate(invalid="ignore", divide="ignore"):
                return x * (n - x) * (n - x - 1) / denom

        self.verify_definition(ts, sample_sets, windows, f, ts.Y1, Y1)


class TestBranchY1(TestY1, TopologyExamplesMixin):
    mode = "branch"


class TestNodeY1(TestY1, TopologyExamplesMixin):
    mode = "node"


class TestSiteY1(TestY1, MutatedTopologyExamplesMixin):
    mode = "site"


############################################
# Divergence
############################################


def site_divergence(ts, sample_sets, indexes, windows=None, span_normalise=True):
    out = np.zeros((len(windows) - 1, len(indexes)))
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        haps = ts.genotype_matrix(isolated_as_missing=False).T
        site_positions = [x.position for x in ts.sites()]
        for i, (ix, iy) in enumerate(indexes):
            X = sample_sets[ix]
            Y = sample_sets[iy]
            denom = np.float64(len(X) * len(Y))
            site_in_window = False
            S = 0
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    site_in_window = True
                    for x in X:
                        x_index = np.where(samples == x)[0][0]
                        for y in Y:
                            y_index = np.where(samples == y)[0][0]
                            if haps[x_index][k] != haps[y_index][k]:
                                # x|y
                                S += 1
            if site_in_window:
                with np.errstate(invalid="ignore", divide="ignore"):
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def branch_divergence(ts, sample_sets, indexes, windows=None, span_normalise=True):
    out = np.zeros((len(windows) - 1, len(indexes)))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, (ix, iy) in enumerate(indexes):
            X = sample_sets[ix]
            Y = sample_sets[iy]
            denom = np.float64(len(X) * len(Y))
            has_trees = False
            S = 0
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                if tr.total_branch_length > 0:
                    has_trees = True
                SS = 0
                for x in X:
                    for y in Y:
                        SS += path_length(tr, x, y)
                S += SS * (min(end, tr.interval.right) - max(begin, tr.interval.left))
            if has_trees:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def node_divergence(ts, sample_sets, indexes, windows=None, span_normalise=True):
    out = np.zeros((len(windows) - 1, ts.num_nodes, len(indexes)))
    for i, (ix, iy) in enumerate(indexes):
        X = sample_sets[ix]
        Y = sample_sets[iy]
        tX = len(X)
        tY = len(Y)
        denom = np.float64(len(X) * len(Y))
        for j in range(len(windows) - 1):
            begin = windows[j]
            end = windows[j + 1]
            S = np.zeros(ts.num_nodes)
            for t1, t2 in zip(ts.trees(tracked_samples=X), ts.trees(tracked_samples=Y)):
                if t1.interval.right <= begin:
                    continue
                if t1.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in t1.nodes():
                    # count number of pairwise paths going through u
                    nX = t1.num_tracked_samples(u)
                    nY = t2.num_tracked_samples(u)
                    SS[u] += nX * (tY - nY) + (tX - nX) * nY
                S += SS * (min(end, t1.interval.right) - max(begin, t1.interval.left))
            with suppress_division_by_zero_warning():
                out[j, :, i] = S / denom
            if span_normalise:
                out[j, :, i] /= end - begin
    return out


def divergence(
    ts, sample_sets, indexes=None, windows=None, mode="site", span_normalise=True
):
    """
    Computes average pairwise divergence between two random choices from x
    over the window specified.
    """
    windows = ts.parse_windows(windows)
    if indexes is None:
        indexes = [(0, 1)]
    method_map = {
        "site": site_divergence,
        "node": node_divergence,
        "branch": branch_divergence,
    }
    return method_map[mode](
        ts, sample_sets, indexes=indexes, windows=windows, span_normalise=span_normalise
    )


class TestDivergence(StatsTestCase, TwoWaySampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_sample_sets_indexes(self, ts, sample_sets, indexes, windows):
        n = np.array([len(x) for x in sample_sets])

        denom = np.array([n[i] * (n[j] - (i == j)) for i, j in indexes])

        def f(x):
            numer = np.array([(x[i] * (n[j] - x[j])) for i, j in indexes])
            return numer / denom

        self.verify_definition(
            ts, sample_sets, indexes, windows, f, ts.divergence, divergence
        )


class TestBranchDivergence(TestDivergence, TopologyExamplesMixin):
    mode = "branch"


class TestNodeDivergence(TestDivergence, TopologyExamplesMixin):
    mode = "node"


class TestSiteDivergence(TestDivergence, MutatedTopologyExamplesMixin):
    mode = "site"


############################################
# Genetic relatedness
############################################


def site_genetic_relatedness(
    ts,
    sample_sets,
    indexes,
    windows=None,
    span_normalise=True,
    polarised=True,
    proportion=True,
    centre=True,
):
    if windows is None:
        windows = [0.0, ts.sequence_length]
    out = np.zeros((len(windows) - 1, len(indexes)))
    all_samples = np.array(list({u for s in sample_sets for u in s}))
    denom = np.ones(len(windows))
    if proportion:
        denom = ts.segregating_sites(
            sample_sets=all_samples,
            windows=windows,
            mode="site",
            span_normalise=span_normalise,
        )
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for vv in zip(
            *[
                ts.variants(left=begin, right=end, samples=x, isolated_as_missing=False)
                for x in sample_sets
            ]
        ):
            ancestral_state = vv[0].site.ancestral_state
            alleles = vv[0].alleles
            ff = [v.frequencies() for v in vv]
            for a in alleles:
                mean_f = sum([f[a] for f in ff]) / len(ff)
                for i, (ix, iy) in enumerate(indexes):
                    fx = ff[ix][a]
                    fy = ff[iy][a]
                    if not (polarised and a == ancestral_state):
                        if centre:
                            out[j][i] += (fx - mean_f) * (fy - mean_f)
                        else:
                            out[j][i] += fx * fy
        for i in range(len(indexes)):
            with np.errstate(invalid="ignore", divide="ignore"):
                out[j][i] /= denom[j]
            if span_normalise:
                out[j][i] /= end - begin
    return out


def branch_genetic_relatedness(
    ts,
    sample_sets,
    indexes,
    windows=None,
    span_normalise=True,
    polarised=True,
    proportion=True,
    centre=True,
):
    if windows is None:
        windows = [0.0, ts.sequence_length]
    out = np.zeros((len(windows) - 1, len(indexes)))
    all_samples = np.array(list({u for s in sample_sets for u in s}))
    denom = np.ones(len(windows))
    if proportion:
        denom = ts.segregating_sites(
            sample_sets=all_samples,
            windows=windows,
            mode="branch",
            span_normalise=span_normalise,
        )
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for tr in ts.trees():
            if tr.interval.right <= begin:
                continue
            if tr.interval.left >= end:
                break
            span = min(end, tr.interval.right) - max(begin, tr.interval.left)
            # iterating over tr.nodes will miss nodes unreachable from samples
            for v in range(ts.num_nodes):
                area = tr.branch_length(v) * span
                freqs = [
                    sum([tr.is_descendant(u, v) for u in x]) / len(x)
                    for x in sample_sets
                ]
                mean_freq = sum(freqs) / len(freqs)
                for i, (ix, iy) in enumerate(indexes):
                    fx = freqs[ix]
                    fy = freqs[iy]
                    if centre:
                        out[j][i] += area * (fx - mean_freq) * (fy - mean_freq)
                        if not polarised:
                            out[j][i] += (
                                area
                                * (1 - fx - (1 - mean_freq))
                                * (1 - fy - (1 - mean_freq))
                            )
                    else:
                        out[j][i] += area * fx * fy
                        if not polarised:
                            out[j][i] += area * (1 - fx) * (1 - fy)
        for i in range(len(indexes)):
            with np.errstate(invalid="ignore", divide="ignore"):
                out[j][i] /= denom[j]
            if span_normalise:
                out[j][i] /= end - begin
    return out


def node_genetic_relatedness(
    ts,
    sample_sets,
    indexes,
    windows=None,
    span_normalise=True,
    proportion=True,
    centre=True,
    polarised=True,
):
    if windows is None:
        windows = [0.0, ts.sequence_length]
    out = np.zeros((len(windows) - 1, ts.num_nodes, len(indexes)))
    all_samples = np.array(list({u for s in sample_sets for u in s}))
    denom = np.ones((len(windows), ts.num_nodes))
    if proportion:
        denom = ts.segregating_sites(
            sample_sets=all_samples,
            windows=windows,
            mode="node",
            span_normalise=span_normalise,
        )
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for tr in ts.trees():
            if tr.interval.right <= begin:
                continue
            if tr.interval.left >= end:
                break
            span = min(end, tr.interval.right) - max(begin, tr.interval.left)
            for v in range(ts.num_nodes):
                freqs = [
                    sum([tr.is_descendant(u, v) for u in x]) / len(x)
                    for x in sample_sets
                ]
                mean_freq = sum(freqs) / len(freqs)
                for i, (ix, iy) in enumerate(indexes):
                    fx = freqs[ix]
                    fy = freqs[iy]
                    if centre:
                        out[j][v][i] += span * (fx - mean_freq) * (fy - mean_freq)
                        if not polarised:
                            out[j][v][i] += (
                                span
                                * (1 - fx - (1 - mean_freq))
                                * (1 - fy - (1 - mean_freq))
                            )
                    else:
                        out[j][v][i] += span * fx * fy
                        if not polarised:
                            out[j][v][i] += span * (1 - fx) * (1 - fy)
        for i in range(len(indexes)):
            for v in ts.nodes():
                iV = v.id
                with np.errstate(invalid="ignore", divide="ignore"):
                    out[j, iV, i] /= denom[j, iV]
                if span_normalise:
                    out[j, iV, i] /= end - begin
    return out


def genetic_relatedness(
    ts,
    sample_sets,
    indexes=None,
    windows=None,
    mode="site",
    span_normalise=True,
    proportion=True,
    centre=True,
    polarised=True,
):
    """
    Computes genetic relatedness between two random choices from x
    over the window specified.
    """
    windows = ts.parse_windows(windows)
    if indexes is None:
        indexes = [(0, 1)]
    method_map = {
        "site": site_genetic_relatedness,
        "node": node_genetic_relatedness,
        "branch": branch_genetic_relatedness,
    }
    return method_map[mode](
        ts,
        sample_sets,
        indexes=indexes,
        windows=windows,
        span_normalise=span_normalise,
        polarised=polarised,
        proportion=proportion,
        centre=centre,
    )


class TestGeneticRelatedness(StatsTestCase, TwoWaySampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_definition(
        self,
        ts,
        sample_sets,
        indexes,
        windows,
        summary_func,
        ts_method,
        definition,
        proportion,
        polarised=True,
        centre=True,
    ):
        def wrapped_summary_func(x):
            with suppress_division_by_zero_warning():
                return summary_func(x)

        W = np.array([[u in A for A in sample_sets] for u in ts.samples()], dtype=float)
        # Determine output_dim of the function
        M = len(wrapped_summary_func(W[0]))
        denom = 1
        if proportion:
            all_samples = list({u for s in sample_sets for u in s})
            denom = ts.segregating_sites(
                sample_sets=[all_samples], windows=windows, mode=self.mode
            )

        with np.errstate(divide="ignore", invalid="ignore"):
            sigma1 = (
                ts.general_stat(
                    W,
                    wrapped_summary_func,
                    M,
                    windows,
                    mode=self.mode,
                    strict=centre,
                    polarised=polarised,
                )
                / denom
            )
        sigma2 = ts_method(
            sample_sets,
            indexes=indexes,
            windows=windows,
            mode=self.mode,
            proportion=proportion,
            centre=centre,
            polarised=polarised,
        )
        sigma3 = definition(
            ts,
            sample_sets,
            indexes=indexes,
            windows=windows,
            mode=self.mode,
            proportion=proportion,
            centre=centre,
            polarised=polarised,
        )
        assert sigma1.shape == sigma2.shape
        assert sigma1.shape == sigma3.shape
        self.assertArrayAlmostEqual(sigma1, sigma2)
        self.assertArrayAlmostEqual(sigma1, sigma3)

    def verify_sample_sets_indexes(self, ts, sample_sets, indexes, windows):
        n = np.array([len(x) for x in sample_sets])

        def f_noncentred(x):
            p = x / n
            return np.array([p[i] * p[j] for i, j in indexes])

        def f_centred(x):
            p = x / n
            mp = np.mean(p)
            return np.array([(p[i] - mp) * (p[j] - mp) for i, j in indexes])

        for proportion in [True, False]:
            self.verify_definition(
                ts,
                sample_sets,
                indexes,
                windows,
                f_centred,
                ts.genetic_relatedness,
                genetic_relatedness,
                proportion,
            )

        for centre, polarised in [
            (True, True),
            (False, True),
            (True, False),
            (False, False),
        ]:
            f = f_centred if centre else f_noncentred
            self.verify_definition(
                ts,
                sample_sets,
                indexes,
                windows,
                f,
                ts.genetic_relatedness,
                genetic_relatedness,
                proportion=False,
                centre=centre,
                polarised=polarised,
            )

    @pytest.mark.parametrize("proportion", [None, True, False])
    def test_shapes(self, proportion):
        # exclude this test in the parent class
        if self.mode is None:
            return
        ts = msprime.sim_ancestry(
            8,
            random_seed=1,
            end_time=10,
            sequence_length=10,
            population_size=10,
            recombination_rate=0.02,
        )
        ts = msprime.sim_mutations(ts, rate=0.01, random_seed=2)
        x = ts.genetic_relatedness(
            sample_sets=[[0, 1, 2], [3]],
            indexes=None,
            windows=None,
            mode=self.mode,
            proportion=proportion,
        )
        if self.mode == "node":
            assert x.shape == (ts.num_nodes,)
        else:
            assert x.shape == ()
        x = ts.genetic_relatedness(
            sample_sets=[[0, 1, 2], [3]],
            indexes=[(0, 1)],
            windows=None,
            mode=self.mode,
            proportion=proportion,
        )
        if self.mode == "node":
            assert x.shape == (ts.num_nodes, 1)
        else:
            assert x.shape == (1,)
        x = ts.genetic_relatedness(
            sample_sets=[[0, 1, 2], [3]],
            indexes=[(0, 1)],
            windows=[0, 10],
            mode=self.mode,
            proportion=proportion,
        )
        if self.mode == "node":
            assert x.shape == (1, ts.num_nodes, 1)
        else:
            assert x.shape == (1, 1)
        x = ts.genetic_relatedness(
            sample_sets=[[0, 1, 2], [3]],
            indexes=[(0, 1)],
            windows=[0, 5, 10],
            mode=self.mode,
            proportion=proportion,
        )
        if self.mode == "node":
            assert x.shape == (2, ts.num_nodes, 1)
        else:
            assert x.shape == (2, 1)
        x = ts.genetic_relatedness(
            sample_sets=[[0, 1, 2], [3]],
            indexes=None,
            windows=[0, 5, 10],
            mode=self.mode,
            proportion=proportion,
        )
        if self.mode == "node":
            assert x.shape == (2, ts.num_nodes)
        else:
            assert x.shape == (2,)
        x = ts.genetic_relatedness(
            sample_sets=[[0, 1, 2], [3], [4, 5]],
            indexes=[(0, 1), (1, 2)],
            windows=[0, 5, 9, 10],
            mode=self.mode,
            proportion=proportion,
        )
        if self.mode == "node":
            assert x.shape == (3, ts.num_nodes, 2)
        else:
            assert x.shape == (3, 2)
        x = ts.genetic_relatedness(
            sample_sets=[[0, 1, 2], [3], [4, 5]],
            indexes=[(0, 1), (1, 2)],
            windows=None,
            mode=self.mode,
            proportion=proportion,
        )
        if self.mode == "node":
            assert x.shape == (ts.num_nodes, 2)
        else:
            assert x.shape == (2,)


class TestBranchGeneticRelatedness(TestGeneticRelatedness, TopologyExamplesMixin):
    mode = "branch"

    @pytest.mark.parametrize("polarised", [True, False])
    def test_simple_tree_noncentred(self, polarised):
        # 2.00┊   4   ┊
        #     ┊ ┏━┻┓  ┊
        # 1.00┊ ┃  3  ┊
        #     ┊ ┃ ┏┻┓ ┊
        # 0.00┊ 0 1 2 ┊
        #     0       1
        ts = tskit.Tree.generate_balanced(3).tree_sequence
        indexes = [(0, 0), (0, 1), (1, 1), (1, 2), (2, 2)]
        sample_sets = [[0], [1], [2]]
        if polarised:
            A = np.array(
                [
                    2,  # (0, 0)
                    0,  # (0, 1)
                    2,  # (1, 1)
                    1,  # (1, 2),
                    2,  # (2, 2)
                ]
            )
        else:
            A = np.array(
                [
                    (2 + 3),  # (0, 0)
                    (0 + 1),  # (0, 1)
                    (2 + 3),  # (1, 1)
                    (1 + 2),  # (1, 2),
                    (2 + 3),  # (2, 2)
                ]
            )
        B = branch_genetic_relatedness(
            ts,
            sample_sets=sample_sets,
            indexes=indexes,
            polarised=polarised,
            proportion=False,
            centre=False,
        ).squeeze()
        C = ts.genetic_relatedness(
            sample_sets=sample_sets,
            indexes=indexes,
            mode="branch",
            polarised=polarised,
            proportion=False,
            centre=False,
        ).squeeze()
        self.assertArrayAlmostEqual(A, B)
        self.assertArrayAlmostEqual(A, C)


class TestNodeGeneticRelatedness(TestGeneticRelatedness, TopologyExamplesMixin):
    mode = "node"


class TestSiteGeneticRelatedness(TestGeneticRelatedness, MutatedTopologyExamplesMixin):
    mode = "site"

    def test_match_K_c0(self):
        # This test checks that ts.genetic_relatedness() matches K_c0
        # from Speed & Balding (2014) https://www.nature.com/articles/nrg3821
        ts = msprime.simulate(
            10, mutation_rate=0.01, length=100, recombination_rate=0.01, random_seed=23
        )
        samples = [u for u in ts.samples()]
        sample_sets = [[0, 1], [2, 3], [4, 5]]
        all_samples = list({u for s in sample_sets for u in s})
        sample_ind = [samples.index(x) for x in all_samples]
        indexes = [(0, 0), (0, 1), (0, 2), (1, 1), (1, 2), (2, 2)]
        A = ts.genetic_relatedness(
            sample_sets, indexes=indexes, mode="site", span_normalise=False
        )
        # Genotype covariance as in Speed and Balding
        G = ts.genotype_matrix().T
        G = G[sample_ind]
        G_centered = G - G.mean(axis=0)
        B = np.zeros(len(indexes))
        for i, (ix, iy) in enumerate(indexes):
            x1 = sample_sets[ix][0]
            x2 = sample_sets[ix][1]
            y1 = sample_sets[iy][0]
            y2 = sample_sets[iy][1]
            B[i] = (
                (G_centered[x1] + G_centered[x2])
                @ (G_centered[y1] + G_centered[y2])
                / ts.segregating_sites(sample_sets=all_samples, span_normalise=False)
            )
        self.assertArrayAlmostEqual(4 * A, B)


############################################
# Genetic relatedness weighted
############################################


def genetic_relatedness_matrix(
    ts, sample_sets, windows=None, mode="site", polarised=True, centre=True
):
    n = len(sample_sets)
    indexes = [
        (n1, n2) for n1, n2 in itertools.combinations_with_replacement(range(n), 2)
    ]
    if windows is None:
        if mode == "node":
            n_nodes = ts.num_nodes
            K = np.zeros((n_nodes, n, n))
            out = ts.genetic_relatedness(
                sample_sets,
                indexes,
                mode=mode,
                proportion=False,
                span_normalise=True,
                polarised=polarised,
                centre=centre,
            )
            for node in range(n_nodes):
                this_K = np.zeros((n, n))
                this_K[np.triu_indices(n)] = out[node, :]
                this_K = this_K + np.triu(this_K, 1).transpose()
                K[node, :, :] = this_K
        else:
            K = np.zeros((n, n))
            K[np.triu_indices(n)] = ts.genetic_relatedness(
                sample_sets,
                indexes,
                mode=mode,
                proportion=False,
                span_normalise=True,
                centre=centre,
                polarised=polarised,
            )
            K = K + np.triu(K, 1).transpose()
    else:
        windows = ts.parse_windows(windows)
        n_windows = len(windows) - 1
        out = ts.genetic_relatedness(
            sample_sets,
            indexes,
            mode=mode,
            windows=windows,
            proportion=False,
            span_normalise=True,
            polarised=polarised,
            centre=centre,
        )
        if mode == "node":
            n_nodes = ts.num_nodes
            K = np.zeros((n_windows, n_nodes, n, n))
            for win in range(n_windows):
                for node in range(n_nodes):
                    K_this = np.zeros((n, n))
                    K_this[np.triu_indices(n)] = out[win, node, :]
                    K_this = K_this + np.triu(K_this, 1).transpose()
                    K[win, node, :, :] = K_this
        else:
            K = np.zeros((n_windows, n, n))
            for win in range(n_windows):
                K_this = np.zeros((n, n))
                K_this[np.triu_indices(n)] = out[win, :]
                K_this = K_this + np.triu(K_this, 1).transpose()
                K[win, :, :] = K_this
    return K


def genetic_relatedness_weighted(
    ts, W, indexes, windows=None, mode="site", polarised=True, centre=True
):
    if centre:
        W_mean = W.mean(axis=0)
        W = W - W_mean
    sample_sets = [[u] for u in ts.samples()]
    K = genetic_relatedness_matrix(
        ts, sample_sets, windows=windows, mode=mode, centre=centre, polarised=polarised
    )
    n_indexes = len(indexes)
    n_nodes = ts.num_nodes
    if windows is None:
        if mode == "node":
            out = np.zeros((n_nodes, n_indexes))
        else:
            out = np.zeros(n_indexes)
    else:
        windows = ts.parse_windows(windows)
        n_windows = len(windows) - 1
        if mode == "node":
            out = np.zeros((n_windows, n_nodes, n_indexes))
        else:
            out = np.zeros((n_windows, n_indexes))
    for pair in range(n_indexes):
        i1 = indexes[pair][0]
        i2 = indexes[pair][1]
        if windows is None:
            if mode == "node":
                for node in range(n_nodes):
                    this_K = K[node, :, :]
                    out[node, pair] = W[:, i1] @ this_K @ W[:, i2]
            else:
                out[pair] = W[:, i1] @ K @ W[:, i2]
        else:
            for win in range(n_windows):
                if mode == "node":
                    for node in range(n_nodes):
                        this_K = K[win, node, :, :]
                        out[win, node, pair] = W[:, i1] @ this_K @ W[:, i2]
                else:
                    this_K = K[win, :, :]
                    out[win, pair] = W[:, i1] @ this_K @ W[:, i2]
    return out


def example_index_pairs(weights):
    assert weights.shape[1] >= 2
    yield [(0, 1)]
    yield [(1, 0), (0, 1)]
    if weights.shape[1] > 2:
        yield [(0, 1), (1, 2), (0, 2)]


class TestGeneticRelatednessWeighted(StatsTestCase, WeightStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_definition(
        self,
        ts,
        W,
        indexes,
        windows,
        summary_func,
        ts_method,
        definition,
        polarised=True,
        centre=True,
    ):
        # Determine output_dim of the function
        M = len(indexes)

        sigma1 = ts.general_stat(
            W,
            summary_func,
            M,
            windows,
            mode=self.mode,
            span_normalise=True,
            strict=centre,
            polarised=polarised,
        )
        sigma2 = general_stat(
            ts,
            W,
            summary_func,
            windows,
            mode=self.mode,
            span_normalise=True,
            polarised=polarised,
        )

        sigma3 = ts_method(
            W,
            indexes=indexes,
            windows=windows,
            mode=self.mode,
            polarised=polarised,
            centre=centre,
        )
        sigma4 = definition(
            ts,
            W,
            indexes=indexes,
            windows=windows,
            mode=self.mode,
            polarised=polarised,
            centre=centre,
        )
        assert sigma1.shape == sigma2.shape
        assert sigma1.shape == sigma3.shape
        assert sigma1.shape == sigma4.shape
        self.assertArrayAlmostEqual(sigma1, sigma2)
        self.assertArrayAlmostEqual(sigma1, sigma3)
        self.assertArrayAlmostEqual(sigma1, sigma4)

    def verify(self, ts):
        for W, windows in subset_combos(
            self.example_weights(ts, min_size=2), example_windows(ts), p=0.1
        ):
            for indexes in example_index_pairs(W):
                self.verify_weighted_stat(ts, W, indexes, windows)

    def verify_weighted_stat(self, ts, W, indexes, windows):
        n = W.shape[0]
        K = W.shape[1]
        WW = np.column_stack([W, np.ones(n) / n])
        W_sum = WW.sum(axis=0)

        def f_noncentred(x):
            return np.array([x[i] * x[j] for i, j in indexes])

        def f_centred(x):
            pn = x[K]
            return np.array(
                [(x[i] - W_sum[i] * pn) * (x[j] - W_sum[j] * pn) for i, j in indexes]
            )

        for centre, polarised in [
            (True, True),
            (False, True),
            (True, False),
            (False, False),
        ]:
            f = f_centred if centre else f_noncentred
            self.verify_definition(
                ts,
                WW,
                indexes,
                windows,
                f,
                ts.genetic_relatedness_weighted,
                genetic_relatedness_weighted,
                centre=centre,
                polarised=polarised,
            )


class TestBranchGeneticRelatednessWeighted(
    TestGeneticRelatednessWeighted, TopologyExamplesMixin
):
    mode = "branch"


class TestNodeGeneticRelatednessWeighted(
    TestGeneticRelatednessWeighted, TopologyExamplesMixin
):
    mode = "node"


class TestSiteGeneticRelatednessWeighted(
    TestGeneticRelatednessWeighted, MutatedTopologyExamplesMixin
):
    mode = "site"


# NOTE: these classes don't follow the same (anti)-patterns as used elsewhere as they
# were added in several years afterwards.


class TestGeneticRelatednessWeightedSimpleExamples:
    # Values verified against the simple implementations above
    site_value = 22.24
    branch_value = 29.44

    def fixture(self):
        ts = tskit.Tree.generate_balanced(5).tree_sequence
        # Abitrary weights that give non-zero results
        W = np.zeros((ts.num_samples, 2))
        W[0, :] = 1
        W[1, :] = 2
        return tsutil.insert_branch_sites(ts), W

    def test_no_arguments_site(self):
        ts, W = self.fixture()
        X = ts.genetic_relatedness_weighted(W, mode="site")
        assert X.shape == tuple()
        nt.assert_almost_equal(X, self.site_value)

    def test_windows_site(self):
        ts, W = self.fixture()
        X = ts.genetic_relatedness_weighted(W, mode="site", windows=[0, 1 - 1e-12, 1])
        assert X.shape == (2,)
        nt.assert_almost_equal(X[0], self.site_value)
        nt.assert_almost_equal(X[1], 0)

    def test_no_arguments_branch(self):
        ts, W = self.fixture()
        X = ts.genetic_relatedness_weighted(W, mode="branch")
        assert X.shape == tuple()
        nt.assert_almost_equal(X, self.branch_value)

    def test_windows_branch(self):
        ts, W = self.fixture()
        X = ts.genetic_relatedness_weighted(W, mode="branch", windows=[0, 0.5, 1])
        assert X.shape == (2,)
        nt.assert_almost_equal(X, self.branch_value)

    def test_indexes_1D(self):
        ts, W = self.fixture()
        indexes = [0, 1]
        X = ts.genetic_relatedness_weighted(W, indexes, mode="branch")
        assert X.shape == tuple()
        nt.assert_almost_equal(X, self.branch_value)

    def test_indexes_2D(self):
        ts, W = self.fixture()
        indexes = [[0, 1]]
        X = ts.genetic_relatedness_weighted(W, indexes, mode="branch")
        assert X.shape == (1,)
        nt.assert_almost_equal(X, self.branch_value)

    def test_indexes_2D_windows(self):
        ts, W = self.fixture()
        indexes = [[0, 1], [0, 1]]
        X = ts.genetic_relatedness_weighted(
            W, indexes, windows=[0, 0.5, 1], mode="branch"
        )
        assert X.shape == (2, 2)
        nt.assert_almost_equal(X, self.branch_value)


class TestGeneticRelatednessWeightedErrors:
    def ts(self):
        return tskit.Tree.generate_balanced(3).tree_sequence

    @pytest.mark.parametrize("W", [[0], np.array([0]), np.zeros(100)])
    def test_bad_weight_size(self, W):
        with pytest.raises(ValueError, match="First trait dimension"):
            self.ts().genetic_relatedness_weighted(W)

    @pytest.mark.parametrize("cols", [1, 3])
    def test_no_indexes_with_non_2_cols(self, cols):
        ts = self.ts()
        W = np.zeros((ts.num_samples, cols))
        with pytest.raises(ValueError, match="Must specify indexes"):
            ts.genetic_relatedness_weighted(W)

    @pytest.mark.parametrize("indexes", [[], [[0]], [[0, 0, 0]], [[[0], [0], [0]]]])
    def test_bad_index_shapes(self, indexes):
        ts = self.ts()
        W = np.zeros((ts.num_samples, 2))
        with pytest.raises(ValueError, match="Indexes must be convertable to a 2D"):
            ts.genetic_relatedness_weighted(W, indexes=indexes)


############################################
# Fst
############################################


def single_site_Fst(ts, sample_sets, indexes):
    """
    Compute single-site Fst, which between two groups with frequencies p and q is
      1 - 2 * (p (1-p) + q(1-q)) / ( p(1-p) + q(1-q) + p(1-q) + q(1-p) )
    or in the multiallelic case, replacing p(1-p) with the sum over alleles of p(1-p),
    and adjusted for sampling without replacement.
    """
    # TODO: what to do in this case?
    if ts.num_sites == 0:
        out = np.array([np.repeat(np.nan, len(indexes))])
        return out
    out = np.zeros((ts.num_sites, len(indexes)))
    samples = ts.samples()
    # TODO deal with missing data properly.
    for j, v in enumerate(ts.variants(isolated_as_missing=False)):
        for i, (ix, iy) in enumerate(indexes):
            g = v.genotypes
            X = sample_sets[ix]
            Y = sample_sets[iy]
            gX = [a for k, a in zip(samples, g) if k in X]
            gY = [a for k, a in zip(samples, g) if k in Y]
            nX = len(X)
            nY = len(Y)
            dX = dY = dXY = 0
            for a in set(g):
                fX = np.sum(gX == a)
                fY = np.sum(gY == a)
                with suppress_division_by_zero_warning():
                    dX += fX * (nX - fX) / (nX * (nX - 1))
                    dY += fY * (nY - fY) / (nY * (nY - 1))
                    dXY += (fX * (nY - fY) + (nX - fX) * fY) / (2 * nX * nY)
            with suppress_division_by_zero_warning():
                out[j][i] = 1 - 2 * (dX + dY) / (dX + dY + 2 * dXY)
    return out


class TestFst(StatsTestCase, TwoWaySampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify(self, ts):
        # only check per-site
        for sample_sets in example_sample_sets(ts, min_size=2):
            for indexes in example_sample_set_index_pairs(sample_sets):
                self.verify_persite_Fst(ts, sample_sets, indexes)

    def verify_persite_Fst(self, ts, sample_sets, indexes):
        sigma1 = ts.Fst(
            sample_sets,
            indexes=indexes,
            windows="sites",
            mode=self.mode,
            span_normalise=False,
        )
        sigma2 = single_site_Fst(ts, sample_sets, indexes)
        assert sigma1.shape == sigma2.shape
        self.assertArrayAlmostEqual(sigma1, sigma2)


class FstInterfaceMixin:
    def test_interface(self):
        ts = msprime.simulate(10, mutation_rate=0.0)
        sample_sets = [[0, 1, 2], [6, 7], [4]]
        with pytest.raises(ValueError):
            ts.Fst(sample_sets, mode=self.mode)
        with pytest.raises(ValueError):
            ts.Fst(sample_sets, indexes=[(0, 1, 2), (3, 4, 5)], mode=self.mode)
        with pytest.raises(tskit.LibraryError):
            ts.Fst(sample_sets, indexes=[(0, 1), (0, 20)])
        sigma1 = ts.Fst(sample_sets, indexes=[(0, 1)], mode=self.mode)
        sigma2 = ts.Fst(sample_sets, indexes=[(0, 1), (0, 2), (1, 2)], mode=self.mode)
        self.assertArrayAlmostEqual(sigma1[..., 0], sigma2[..., 0])


class TestSiteFst(TestFst, MutatedTopologyExamplesMixin, FstInterfaceMixin):
    mode = "site"


# Since Fst is defined using diversity and divergence, we don't seriously
# test it for correctness for node and branch, and only test the interface.


class TestNodeFst(StatsTestCase, FstInterfaceMixin):
    mode = "node"


class TestBranchFst(StatsTestCase, FstInterfaceMixin):
    mode = "node"


############################################
# Y2
############################################


def branch_Y2(ts, sample_sets, indexes, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(indexes)))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, (ix, iy) in enumerate(indexes):
            X = sample_sets[ix]
            Y = sample_sets[iy]
            denom = np.float64(len(X) * len(Y) * (len(Y) - 1))
            has_trees = False
            S = 0
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                if tr.total_branch_length > 0:
                    has_trees = True
                this_length = min(end, tr.interval.right) - max(begin, tr.interval.left)
                for x in X:
                    for y in Y:
                        for z in set(Y) - {y}:
                            xy_mrca = tr.mrca(x, y)
                            xz_mrca = tr.mrca(x, z)
                            yz_mrca = tr.mrca(y, z)
                            if xy_mrca == xz_mrca:
                                #   /\
                                #  / /\
                                # x y  z
                                S += path_length(tr, x, yz_mrca) * this_length
                            elif xy_mrca == yz_mrca:
                                #   /\
                                #  / /\
                                # y x  z
                                S += path_length(tr, x, xz_mrca) * this_length
                            elif xz_mrca == yz_mrca:
                                #   /\
                                #  / /\
                                # z x  y
                                S += path_length(tr, x, xy_mrca) * this_length
            if has_trees:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def site_Y2(ts, sample_sets, indexes, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    samples = ts.samples()
    out = np.zeros((len(windows) - 1, len(indexes)))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        haps = ts.genotype_matrix(isolated_as_missing=False).T
        site_positions = [x.position for x in ts.sites()]
        for i, (ix, iy) in enumerate(indexes):
            X = sample_sets[ix]
            Y = sample_sets[iy]
            denom = np.float64(len(X) * len(Y) * (len(Y) - 1))
            S = 0
            site_in_window = False
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    site_in_window = True
                    for x in X:
                        x_index = np.where(samples == x)[0][0]
                        for y in Y:
                            y_index = np.where(samples == y)[0][0]
                            for z in set(Y) - {y}:
                                z_index = np.where(samples == z)[0][0]
                                condition = (
                                    haps[x_index, k] != haps[y_index, k]
                                    and haps[x_index, k] != haps[z_index, k]
                                )
                                if condition:
                                    # x|yz
                                    S += 1
            if site_in_window:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def node_Y2(ts, sample_sets, indexes, windows=None, span_normalise=True):
    out = np.zeros((len(windows) - 1, ts.num_nodes, len(indexes)))
    for i, (ix, iy) in enumerate(indexes):
        X = sample_sets[ix]
        Y = sample_sets[iy]
        tX = len(X)
        tY = len(Y)
        denom = np.float64(tX * tY * (tY - 1))
        for j in range(len(windows) - 1):
            begin = windows[j]
            end = windows[j + 1]
            S = np.zeros(ts.num_nodes)
            for t1, t2 in zip(ts.trees(tracked_samples=X), ts.trees(tracked_samples=Y)):
                if t1.interval.right <= begin:
                    continue
                if t1.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in t1.nodes():
                    # count number of pairwise paths going through u
                    nX = t1.num_tracked_samples(u)
                    nY = t2.num_tracked_samples(u)
                    SS[u] += nX * (tY - nY) * (tY - nY - 1) + (tX - nX) * nY * (nY - 1)
                S += SS * (min(end, t1.interval.right) - max(begin, t1.interval.left))
            with suppress_division_by_zero_warning():
                out[j, :, i] = S / denom
            if span_normalise:
                out[j, :, i] /= end - begin
    return out


def Y2(ts, sample_sets, indexes=None, windows=None, mode="site", span_normalise=True):
    windows = ts.parse_windows(windows)
    if indexes is None:
        indexes = [(0, 1)]
    method_map = {"site": site_Y2, "node": node_Y2, "branch": branch_Y2}
    return method_map[mode](
        ts, sample_sets, indexes=indexes, windows=windows, span_normalise=span_normalise
    )


class TestY2(StatsTestCase, TwoWaySampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_sample_sets_indexes(self, ts, sample_sets, indexes, windows):
        n = np.array([len(x) for x in sample_sets])

        denom = np.array([n[i] * n[j] * (n[j] - 1) for i, j in indexes])

        def f(x):
            numer = np.array(
                [(x[i] * (n[j] - x[j]) * (n[j] - x[j] - 1)) for i, j in indexes]
            )
            return numer / denom

        self.verify_definition(ts, sample_sets, indexes, windows, f, ts.Y2, Y2)


class TestBranchY2(TestY2, TopologyExamplesMixin):
    mode = "branch"


class TestNodeY2(TestY2, TopologyExamplesMixin):
    mode = "node"


class TestSiteY2(TestY2, MutatedTopologyExamplesMixin):
    mode = "site"


############################################
# Y3
############################################


def branch_Y3(ts, sample_sets, indexes, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(indexes)))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, (ix, iy, iz) in enumerate(indexes):
            S = 0
            X = sample_sets[ix]
            Y = sample_sets[iy]
            Z = sample_sets[iz]
            denom = np.float64(len(X) * len(Y) * len(Z))
            has_trees = False
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                if tr.total_branch_length > 0:
                    has_trees = True
                this_length = min(end, tr.interval.right) - max(begin, tr.interval.left)
                for x in X:
                    for y in Y:
                        for z in Z:
                            xy_mrca = tr.mrca(x, y)
                            xz_mrca = tr.mrca(x, z)
                            yz_mrca = tr.mrca(y, z)
                            if xy_mrca == xz_mrca:
                                #   /\
                                #  / /\
                                # x y  z
                                S += path_length(tr, x, yz_mrca) * this_length
                            elif xy_mrca == yz_mrca:
                                #   /\
                                #  / /\
                                # y x  z
                                S += path_length(tr, x, xz_mrca) * this_length
                            elif xz_mrca == yz_mrca:
                                #   /\
                                #  / /\
                                # z x  y
                                S += path_length(tr, x, xy_mrca) * this_length
            if has_trees:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def site_Y3(ts, sample_sets, indexes, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(indexes)))
    haps = ts.genotype_matrix(isolated_as_missing=False).T
    site_positions = ts.tables.sites.position
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, (ix, iy, iz) in enumerate(indexes):
            X = sample_sets[ix]
            Y = sample_sets[iy]
            Z = sample_sets[iz]
            denom = np.float64(len(X) * len(Y) * len(Z))
            S = 0
            site_in_window = False
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    site_in_window = True
                    for x in X:
                        x_index = np.where(samples == x)[0][0]
                        for y in Y:
                            y_index = np.where(samples == y)[0][0]
                            for z in Z:
                                z_index = np.where(samples == z)[0][0]
                                if (haps[x_index][k] != haps[y_index][k]) and (
                                    haps[x_index][k] != haps[z_index][k]
                                ):
                                    # x|yz
                                    with suppress_division_by_zero_warning():
                                        S += 1
            if site_in_window:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def node_Y3(ts, sample_sets, indexes, windows=None, span_normalise=True):
    out = np.zeros((len(windows) - 1, ts.num_nodes, len(indexes)))
    for i, (ix, iy, iz) in enumerate(indexes):
        X = sample_sets[ix]
        Y = sample_sets[iy]
        Z = sample_sets[iz]
        tX = len(X)
        tY = len(Y)
        tZ = len(Z)
        denom = np.float64(tX * tY * tZ)
        for j in range(len(windows) - 1):
            begin = windows[j]
            end = windows[j + 1]
            S = np.zeros(ts.num_nodes)
            for t1, t2, t3 in zip(
                ts.trees(tracked_samples=X),
                ts.trees(tracked_samples=Y),
                ts.trees(tracked_samples=Z),
            ):
                if t1.interval.right <= begin:
                    continue
                if t1.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in t1.nodes():
                    # count number of pairwise paths going through u
                    nX = t1.num_tracked_samples(u)
                    nY = t2.num_tracked_samples(u)
                    nZ = t3.num_tracked_samples(u)
                    SS[u] += nX * (tY - nY) * (tZ - nZ) + (tX - nX) * nY * nZ
                S += SS * (min(end, t1.interval.right) - max(begin, t1.interval.left))
            with suppress_division_by_zero_warning():
                out[j, :, i] = S / denom
            if span_normalise:
                out[j, :, i] /= end - begin
    return out


def Y3(ts, sample_sets, indexes=None, windows=None, mode="site", span_normalise=True):
    windows = ts.parse_windows(windows)
    if indexes is None:
        indexes = [(0, 1, 2)]
    method_map = {"site": site_Y3, "node": node_Y3, "branch": branch_Y3}
    return method_map[mode](
        ts, sample_sets, indexes=indexes, windows=windows, span_normalise=span_normalise
    )


class TestY3(StatsTestCase, ThreeWaySampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_sample_sets_indexes(self, ts, sample_sets, indexes, windows):
        n = np.array([len(x) for x in sample_sets])
        denom = np.array([n[i] * n[j] * n[k] for i, j, k in indexes])

        def f(x):
            numer = np.array(
                [x[i] * (n[j] - x[j]) * (n[k] - x[k]) for i, j, k in indexes]
            )
            return numer / denom

        self.verify_definition(ts, sample_sets, indexes, windows, f, ts.Y3, Y3)


class TestBranchY3(TestY3, TopologyExamplesMixin):
    mode = "branch"


class TestNodeY3(TestY3, TopologyExamplesMixin):
    mode = "node"


class TestSiteY3(TestY3, MutatedTopologyExamplesMixin):
    mode = "site"


############################################
# f2
############################################


def branch_f2(ts, sample_sets, indexes, windows=None, span_normalise=True):
    # this is f4(A,B;A,B) but drawing distinct samples from A and B
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(indexes)))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, (ia, ib) in enumerate(indexes):
            A = sample_sets[ia]
            B = sample_sets[ib]
            denom = np.float64(len(A) * (len(A) - 1) * len(B) * (len(B) - 1))
            has_trees = False
            S = 0
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                if tr.total_branch_length > 0:
                    has_trees = True
                this_length = min(end, tr.interval.right) - max(begin, tr.interval.left)
                SS = 0
                for a in A:
                    for b in B:
                        for c in set(A) - {a}:
                            for d in set(B) - {b}:
                                with suppress_division_by_zero_warning():
                                    SS += path_length(tr, tr.mrca(a, c), tr.mrca(b, d))
                                    SS -= path_length(tr, tr.mrca(a, d), tr.mrca(b, c))
                S += SS * this_length
            if has_trees:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def site_f2(ts, sample_sets, indexes, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(indexes)))
    samples = ts.samples()
    haps = ts.genotype_matrix(isolated_as_missing=False).T
    site_positions = ts.tables.sites.position
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, (iA, iB) in enumerate(indexes):
            A = sample_sets[iA]
            B = sample_sets[iB]
            denom = np.float64(len(A) * (len(A) - 1) * len(B) * (len(B) - 1))
            S = 0
            site_in_window = False
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    site_in_window = True
                    for a in A:
                        a_index = np.where(samples == a)[0][0]
                        for b in B:
                            b_index = np.where(samples == b)[0][0]
                            for c in set(A) - {a}:
                                c_index = np.where(samples == c)[0][0]
                                for d in set(B) - {b}:
                                    d_index = np.where(samples == d)[0][0]
                                    if (
                                        (haps[a_index][k] == haps[c_index][k])
                                        and (haps[a_index][k] != haps[d_index][k])
                                        and (haps[a_index][k] != haps[b_index][k])
                                    ):
                                        # ac|bd
                                        S += 1
                                    elif (
                                        (haps[a_index][k] == haps[d_index][k])
                                        and (haps[a_index][k] != haps[c_index][k])
                                        and (haps[a_index][k] != haps[b_index][k])
                                    ):
                                        # ad|bc
                                        S -= 1
            if site_in_window:
                with np.errstate(invalid="ignore", divide="ignore"):
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def node_f2(ts, sample_sets, indexes, windows=None, span_normalise=True):
    out = np.zeros((len(windows) - 1, ts.num_nodes, len(indexes)))
    for i, (ia, ib) in enumerate(indexes):
        A = sample_sets[ia]
        B = sample_sets[ib]
        tA = len(A)
        tB = len(B)
        denom = np.float64(tA * (tA - 1) * tB * (tB - 1))
        for j in range(len(windows) - 1):
            begin = windows[j]
            end = windows[j + 1]
            S = np.zeros(ts.num_nodes)
            for t1, t2 in zip(ts.trees(tracked_samples=A), ts.trees(tracked_samples=B)):
                if t1.interval.right <= begin:
                    continue
                if t1.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in t1.nodes():
                    # count number of pairwise paths going through u
                    nA = t1.num_tracked_samples(u)
                    nB = t2.num_tracked_samples(u)
                    # xy|uv - xv|uy with x,y in A, u, v in B
                    SS[u] += nA * (nA - 1) * (tB - nB) * (tB - nB - 1) + (tA - nA) * (
                        tA - nA - 1
                    ) * nB * (nB - 1)
                    SS[u] -= 2 * nA * nB * (tA - nA) * (tB - nB)
                S += SS * (min(end, t1.interval.right) - max(begin, t1.interval.left))
            with suppress_division_by_zero_warning():
                out[j, :, i] = S / denom
            if span_normalise:
                out[j, :, i] /= end - begin
    return out


def f2(ts, sample_sets, indexes=None, windows=None, mode="site", span_normalise=True):
    """
    Patterson's f2 statistic definitions.
    """
    windows = ts.parse_windows(windows)
    if indexes is None:
        indexes = [(0, 1)]
    method_map = {"site": site_f2, "node": node_f2, "branch": branch_f2}
    return method_map[mode](
        ts, sample_sets, indexes=indexes, windows=windows, span_normalise=span_normalise
    )


class Testf2(StatsTestCase, TwoWaySampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_sample_sets_indexes(self, ts, sample_sets, indexes, windows):
        n = np.array([len(x) for x in sample_sets])

        denom = np.array([n[i] * (n[i] - 1) * n[j] * (n[j] - 1) for i, j in indexes])

        def f(x):
            numer = np.array(
                [
                    x[i] * (x[i] - 1) * (n[j] - x[j]) * (n[j] - x[j] - 1)
                    - x[i] * (n[i] - x[i]) * (n[j] - x[j]) * x[j]
                    for i, j in indexes
                ]
            )
            return numer / denom

        self.verify_definition(ts, sample_sets, indexes, windows, f, ts.f2, f2)


class TestBranchf2(Testf2, TopologyExamplesMixin):
    mode = "branch"


class TestNodef2(Testf2, TopologyExamplesMixin):
    mode = "node"


class TestSitef2(Testf2, MutatedTopologyExamplesMixin):
    mode = "site"


############################################
# f3
############################################


def branch_f3(ts, sample_sets, indexes, windows=None, span_normalise=True):
    # this is f4(A,B;A,C) but drawing distinct samples from A
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(indexes)))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, (ia, ib, ic) in enumerate(indexes):
            A = sample_sets[ia]
            B = sample_sets[ib]
            C = sample_sets[ic]
            denom = np.float64(len(A) * (len(A) - 1) * len(B) * len(C))
            has_trees = False
            S = 0
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                if tr.total_branch_length > 0:
                    has_trees = True
                this_length = min(end, tr.interval.right) - max(begin, tr.interval.left)
                SS = 0
                for a in A:
                    for b in B:
                        for c in set(A) - {a}:
                            for d in C:
                                SS += path_length(tr, tr.mrca(a, c), tr.mrca(b, d))
                                SS -= path_length(tr, tr.mrca(a, d), tr.mrca(b, c))
                S += SS * this_length
            if has_trees:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def site_f3(ts, sample_sets, indexes, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(indexes)))
    samples = ts.samples()
    haps = ts.genotype_matrix(isolated_as_missing=False).T
    site_positions = ts.tables.sites.position
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, (iA, iB, iC) in enumerate(indexes):
            A = sample_sets[iA]
            B = sample_sets[iB]
            C = sample_sets[iC]
            denom = np.float64(len(A) * (len(A) - 1) * len(B) * len(C))
            S = 0
            site_in_window = False
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    site_in_window = True
                    for a in A:
                        a_index = np.where(samples == a)[0][0]
                        for b in B:
                            b_index = np.where(samples == b)[0][0]
                            for c in set(A) - {a}:
                                c_index = np.where(samples == c)[0][0]
                                for d in C:
                                    d_index = np.where(samples == d)[0][0]
                                    if (
                                        (haps[a_index][k] == haps[c_index][k])
                                        and (haps[a_index][k] != haps[d_index][k])
                                        and (haps[a_index][k] != haps[b_index][k])
                                    ):
                                        # ac|bd
                                        S += 1
                                    elif (
                                        (haps[a_index][k] == haps[d_index][k])
                                        and (haps[a_index][k] != haps[c_index][k])
                                        and (haps[a_index][k] != haps[b_index][k])
                                    ):
                                        # ad|bc
                                        S -= 1
            if site_in_window:
                with np.errstate(invalid="ignore", divide="ignore"):
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def node_f3(ts, sample_sets, indexes, windows=None, span_normalise=True):
    out = np.zeros((len(windows) - 1, ts.num_nodes, len(indexes)))
    for i, (iA, iB, iC) in enumerate(indexes):
        A = sample_sets[iA]
        B = sample_sets[iB]
        C = sample_sets[iC]
        tA = len(A)
        tB = len(B)
        tC = len(C)
        denom = np.float64(tA * (tA - 1) * tB * tC)
        for j in range(len(windows) - 1):
            begin = windows[j]
            end = windows[j + 1]
            S = np.zeros(ts.num_nodes)
            for t1, t2, t3 in zip(
                ts.trees(tracked_samples=A),
                ts.trees(tracked_samples=B),
                ts.trees(tracked_samples=C),
            ):
                if t1.interval.right <= begin:
                    continue
                if t1.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in t1.nodes():
                    # count number of pairwise paths going through u
                    nA = t1.num_tracked_samples(u)
                    nB = t2.num_tracked_samples(u)
                    nC = t3.num_tracked_samples(u)
                    # xy|uv - xv|uy with x,y in A, u in B and v in C
                    SS[u] += (
                        nA * (nA - 1) * (tB - nB) * (tC - nC)
                        + (tA - nA) * (tA - nA - 1) * nB * nC
                    )
                    SS[u] -= (
                        nA * nC * (tA - nA) * (tB - nB)
                        + (tA - nA) * (tC - nC) * nA * nB
                    )
                S += SS * (min(end, t1.interval.right) - max(begin, t1.interval.left))
            with suppress_division_by_zero_warning():
                out[j, :, i] = S / denom
            if span_normalise:
                out[j, :, i] /= end - begin
    return out


def f3(ts, sample_sets, indexes=None, windows=None, mode="site", span_normalise=True):
    """
    Patterson's f3 statistic definitions.
    """
    windows = ts.parse_windows(windows)
    if indexes is None:
        indexes = [(0, 1, 2)]
    method_map = {"site": site_f3, "node": node_f3, "branch": branch_f3}
    return method_map[mode](
        ts, sample_sets, indexes=indexes, windows=windows, span_normalise=span_normalise
    )


class Testf3(StatsTestCase, ThreeWaySampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_sample_sets_indexes(self, ts, sample_sets, indexes, windows):
        n = np.array([len(x) for x in sample_sets])
        denom = np.array([n[i] * (n[i] - 1) * n[j] * n[k] for i, j, k in indexes])

        def f(x):
            numer = np.array(
                [
                    x[i] * (x[i] - 1) * (n[j] - x[j]) * (n[k] - x[k])
                    - x[i] * (n[i] - x[i]) * (n[j] - x[j]) * x[k]
                    for i, j, k in indexes
                ]
            )
            return numer / denom

        self.verify_definition(ts, sample_sets, indexes, windows, f, ts.f3, f3)


class TestBranchf3(Testf3, TopologyExamplesMixin):
    mode = "branch"


class TestNodef3(Testf3, TopologyExamplesMixin):
    mode = "node"


class TestSitef3(Testf3, MutatedTopologyExamplesMixin):
    mode = "site"


############################################
# f4
############################################


def branch_f4(
    ts, sample_sets, indexes, windows=None, time_windows=None, span_normalise=True
):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, len(indexes)))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, (iA, iB, iC, iD) in enumerate(indexes):
            A = sample_sets[iA]
            B = sample_sets[iB]
            C = sample_sets[iC]
            D = sample_sets[iD]
            denom = np.float64(len(A) * len(B) * len(C) * len(D))
            has_trees = False
            S = 0
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                if tr.total_branch_length > 0:
                    has_trees = True
                this_length = min(end, tr.interval.right) - max(begin, tr.interval.left)
                SS = 0
                for a in A:
                    for b in B:
                        for c in C:
                            for d in D:
                                with suppress_division_by_zero_warning():
                                    SS += path_length(tr, tr.mrca(a, c), tr.mrca(b, d))
                                    SS -= path_length(tr, tr.mrca(a, d), tr.mrca(b, c))
                S += SS * this_length
            if has_trees:
                with suppress_division_by_zero_warning():
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def site_f4(ts, sample_sets, indexes, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    samples = ts.samples()
    haps = ts.genotype_matrix(isolated_as_missing=False).T
    site_positions = ts.tables.sites.position
    out = np.zeros((len(windows) - 1, len(indexes)))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i, (iA, iB, iC, iD) in enumerate(indexes):
            A = sample_sets[iA]
            B = sample_sets[iB]
            C = sample_sets[iC]
            D = sample_sets[iD]
            denom = np.float64(len(A) * len(B) * len(C) * len(D))
            S = 0
            site_in_window = False
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    site_in_window = True
                    for a in A:
                        a_index = np.where(samples == a)[0][0]
                        for b in B:
                            b_index = np.where(samples == b)[0][0]
                            for c in C:
                                c_index = np.where(samples == c)[0][0]
                                for d in D:
                                    d_index = np.where(samples == d)[0][0]
                                    if (
                                        (haps[a_index][k] == haps[c_index][k])
                                        and (haps[a_index][k] != haps[d_index][k])
                                        and (haps[a_index][k] != haps[b_index][k])
                                    ):
                                        # ac|bd
                                        S += 1
                                    elif (
                                        (haps[a_index][k] == haps[d_index][k])
                                        and (haps[a_index][k] != haps[c_index][k])
                                        and (haps[a_index][k] != haps[b_index][k])
                                    ):
                                        # ad|bc
                                        S -= 1
            if site_in_window:
                with np.errstate(invalid="ignore", divide="ignore"):
                    out[j][i] = S / denom
                if span_normalise:
                    out[j][i] /= end - begin
    return out


def node_f4(ts, sample_sets, indexes, windows=None, span_normalise=True):
    windows = ts.parse_windows(windows)
    out = np.zeros((len(windows) - 1, ts.num_nodes, len(indexes)))
    for i, (iA, iB, iC, iD) in enumerate(indexes):
        A = sample_sets[iA]
        B = sample_sets[iB]
        C = sample_sets[iC]
        D = sample_sets[iD]
        tA = len(A)
        tB = len(B)
        tC = len(C)
        tD = len(D)
        denom = np.float64(tA * tB * tC * tD)
        for j in range(len(windows) - 1):
            begin = windows[j]
            end = windows[j + 1]
            S = np.zeros(ts.num_nodes)
            for t1, t2, t3, t4 in zip(
                ts.trees(tracked_samples=A),
                ts.trees(tracked_samples=B),
                ts.trees(tracked_samples=C),
                ts.trees(tracked_samples=D),
            ):
                if t1.interval.right <= begin:
                    continue
                if t1.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in t1.nodes():
                    # count number of pairwise paths going through u
                    nA = t1.num_tracked_samples(u)
                    nB = t2.num_tracked_samples(u)
                    nC = t3.num_tracked_samples(u)
                    nD = t4.num_tracked_samples(u)
                    # ac|bd - ad|bc
                    SS[u] += (
                        nA * nC * (tB - nB) * (tD - nD)
                        + (tA - nA) * (tC - nC) * nB * nD
                    )
                    SS[u] -= (
                        nA * nD * (tB - nB) * (tC - nC)
                        + (tA - nA) * (tD - nD) * nB * nC
                    )
                S += SS * (min(end, t1.interval.right) - max(begin, t1.interval.left))
            with suppress_division_by_zero_warning():
                out[j, :, i] = S / denom
            if span_normalise:
                out[j, :, i] /= end - begin
    return out


def f4(
    ts,
    sample_sets,
    indexes=None,
    windows=None,
    time_windows=None,
    mode="site",
    span_normalise=True,
):
    """
    Patterson's f4 statistic definitions.
    """
    if indexes is None:
        indexes = [(0, 1, 2, 3)]
    method_map = {"site": site_f4, "node": node_f4, "branch": branch_f4}
    return method_map[mode](
        ts, sample_sets, indexes=indexes, windows=windows, span_normalise=span_normalise
    )


class Testf4(StatsTestCase, FourWaySampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_sample_sets_indexes(self, ts, sample_sets, indexes, windows):
        n = np.array([len(x) for x in sample_sets])
        denom = np.array([n[i] * n[j] * n[k] * n[l] for i, j, k, l in indexes])

        def f(x):
            numer = np.array(
                [
                    x[i] * x[k] * (n[j] - x[j]) * (n[l] - x[l])
                    - x[i] * x[l] * (n[j] - x[j]) * (n[k] - x[k])
                    for i, j, k, l in indexes
                ]
            )
            return numer / denom

        self.verify_definition(ts, sample_sets, indexes, windows, f, ts.f4, f4)


class TestBranchf4(Testf4, TopologyExamplesMixin):
    mode = "branch"


class TestNodef4(Testf4, TopologyExamplesMixin):
    mode = "node"


class TestSitef4(Testf4, MutatedTopologyExamplesMixin):
    mode = "site"


############################################
# Allele frequency spectrum
############################################


def fold(x, dims):
    """
    Folds the specified coordinates.
    """
    x = np.array(x, dtype=int)
    dims = np.array(dims, dtype=int)
    k = len(dims)
    n = np.sum(dims - 1) / 2
    s = np.sum(x)
    while s == n and k > 0:
        k -= 1
        assert k >= 0
        n -= (dims[k] - 1) / 2
        s -= x[k]
    if s > n:
        x = dims - 1 - x
    assert np.all(x >= 0)
    return tuple(x)


def foldit(A):
    B = np.zeros(A.shape)
    dims = A.shape
    inds = [range(k) for k in dims]
    for ij in itertools.product(*inds):
        nij = fold(ij, dims)
        B[nij] += A[ij]
    return B


def fold_windowed(X):
    Y = np.zeros(X.shape)
    for k in range(X.shape[0]):
        Y[k] = foldit(X[k])
    return Y


class TestFold:
    """
    Tests for the fold operation used in the AFS.
    """

    def test_examples(self):
        A = np.arange(12)
        Af = np.array(
            [11.0, 11.0, 11.0, 11.0, 11.0, 11.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        )

        assert np.all(foldit(A) == Af)

        B = A.copy().reshape(3, 4)
        Bf = np.array(
            [[11.0, 11.0, 11.0, 0.0], [11.0, 11.0, 0.0, 0.0], [11.0, 0.0, 0.0, 0.0]]
        )
        assert np.all(foldit(B) == Bf)

        C = A.copy().reshape(3, 2, 2)
        Cf = np.array(
            [
                [[11.0, 11.0], [11.0, 11.0]],
                [[11.0, 11.0], [0.0, 0.0]],
                [[0.0, 0.0], [0.0, 0.0]],
            ]
        )
        assert np.all(foldit(C) == Cf)

        D = np.arange(9).reshape((3, 3))
        Df = np.array([[8.0, 8.0, 8.0], [8.0, 4.0, 0.0], [0.0, 0.0, 0.0]])
        assert np.all(foldit(D) == Df)

        E = np.arange(9)
        Ef = np.array([8.0, 8.0, 8.0, 8.0, 4.0, 0.0, 0.0, 0.0, 0.0])
        assert np.all(foldit(E) == Ef)

    def test_branch_folded(self, ts_ancestry_10_fixture):
        ts = ts_ancestry_10_fixture
        folded = ts.allele_frequency_spectrum(
            windows=[0, 5, 8, 9, 10], mode="branch", polarised=False
        )
        unfolded = ts.allele_frequency_spectrum(
            windows=[0, 5, 8, 9, 10], mode="branch", polarised=True
        )
        assert np.allclose(fold_windowed(unfolded), folded)

    def test_site_folded(self, ts_ancestry_10_fixture):
        ts = ts_ancestry_10_fixture
        ts = msprime.sim_mutations(ts, rate=1, random_seed=1, discrete_genome=False)
        for s in ts.sites():
            assert len(s.mutations) == 1
        folded = ts.allele_frequency_spectrum(
            windows=[0, 5, 8, 9, 10], mode="site", polarised=False, span_normalise=False
        )
        unfolded = ts.allele_frequency_spectrum(
            windows=[0, 5, 8, 9, 10], mode="site", polarised=True, span_normalise=False
        )
        assert np.allclose(fold_windowed(unfolded), folded)


def naive_site_allele_frequency_spectrum(
    ts, sample_sets, windows=None, polarised=False, span_normalise=True
):
    """
    The joint allele frequency spectrum for sites.
    """
    windows = ts.parse_windows(windows)
    num_windows = len(windows) - 1
    out_dim = [1 + len(sample_set) for sample_set in sample_sets]
    out = np.zeros([num_windows] + out_dim)
    G = ts.genotype_matrix(isolated_as_missing=False)
    samples = ts.samples()
    # Indexes of the samples within the sample sets into the samples array.
    sample_set_indexes = [
        np.array([np.where(x == samples)[0][0] for x in sample_set])
        for sample_set in sample_sets
    ]
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for site in ts.sites():
            S = np.zeros(out_dim)
            if begin <= site.position < end:
                g = G[site.id]
                alleles = np.unique(g)

                # Any site monomorphic across all samples does not contribute
                if len(alleles) == 1:
                    continue

                # For each allele, count the number present in each sample set.
                count = {
                    allele: np.zeros(len(sample_sets), dtype=int) for allele in alleles
                }
                for k, sample_set in enumerate(sample_set_indexes):
                    allele_counts = zip(*np.unique(g[sample_set], return_counts=True))
                    for allele, c in allele_counts:
                        count[allele][k] = c
                increment = 0.5
                if polarised:
                    increment = 1
                    # Remove the contribution of the ancestral state
                    if 0 in count:
                        del count[0]
                for allele_count in count.values():
                    x = tuple(allele_count)
                    if not polarised:
                        x = fold(x, out_dim)
                    S[x] += increment
            if span_normalise:
                S /= end - begin
            out[j, :] += S
    return out


def naive_branch_allele_frequency_spectrum(
    ts, sample_sets, windows=None, polarised=False, span_normalise=True
):
    """
    The joint allele frequency spectrum for branches.
    """
    windows = ts.parse_windows(windows)
    num_windows = len(windows) - 1
    out_dim = [1 + len(sample_set) for sample_set in sample_sets]
    out = np.zeros([num_windows] + out_dim)
    for j in range(num_windows):
        begin = windows[j]
        end = windows[j + 1]
        S = np.zeros(out_dim)
        trees = [
            next(ts.trees(tracked_samples=sample_set)) for sample_set in sample_sets
        ]
        t = trees[0]
        while True:
            tr_len = min(end, t.interval.right) - max(begin, t.interval.left)
            if tr_len > 0:
                for node in t.nodes():
                    if 0 < t.num_samples(node) < ts.num_samples:
                        x = [tree.num_tracked_samples(node) for tree in trees]
                        # Note x must be a tuple for indexing to work
                        if not polarised:
                            x = fold(x, out_dim)
                        S[tuple(x)] += t.branch_length(node) * tr_len

            # Advance the trees
            more = [tree.next() for tree in trees]
            assert len(set(more)) == 1
            if not more[0]:
                break
        if span_normalise:
            S /= end - begin
        out[j, :] = S
    return out


def naive_allele_frequency_spectrum(
    ts, sample_sets, windows=None, polarised=False, mode="site", span_normalise=True
):
    """
    Naive definition of the generalised site frequency spectrum.
    """
    method_map = {
        "site": naive_site_allele_frequency_spectrum,
        "branch": naive_branch_allele_frequency_spectrum,
    }
    return method_map[mode](
        ts,
        sample_sets,
        windows=windows,
        polarised=polarised,
        span_normalise=span_normalise,
    )


def branch_allele_frequency_spectrum(
    ts, sample_sets, windows, polarised=False, span_normalise=True
):
    """
    Efficient implementation of the algorithm used as the basis for the
    underlying C version.
    """
    num_sample_sets = len(sample_sets)
    windows = ts.parse_windows(windows)
    num_windows = windows.shape[0] - 1
    out_dim = [1 + len(sample_set) for sample_set in sample_sets]
    time = ts.tables.nodes.time

    result = np.zeros([num_windows] + out_dim)
    # Number of nodes in sample_set j ancestral to each node u.
    count = np.zeros((ts.num_nodes, num_sample_sets + 1), dtype=np.uint32)
    for j in range(num_sample_sets):
        count[sample_sets[j], j] = 1
    # The last column counts across all samples
    count[ts.samples(), -1] = 1
    # contains the location of the last time we updated the output for a node.
    last_update = np.zeros(ts.num_nodes)
    window_index = 0
    parent = np.zeros(ts.num_nodes, dtype=np.int32) - 1
    branch_length = np.zeros(ts.num_nodes)
    tree_index = 0

    def update_result(window_index, u, right):
        if 0 < count[u, -1] < ts.num_samples:
            x = (right - last_update[u]) * branch_length[u]
            c = count[u, :num_sample_sets]
            if not polarised:
                c = fold(c, out_dim)
            index = tuple([window_index] + list(c))
            result[index] += x
        last_update[u] = right

    for (t_left, t_right), edges_out, edges_in in ts.edge_diffs():
        for edge in edges_out:
            u = edge.child
            v = edge.parent
            update_result(window_index, u, t_left)
            while v != -1:
                update_result(window_index, v, t_left)
                count[v] -= count[u]
                v = parent[v]
            parent[u] = -1
            branch_length[u] = 0

        for edge in edges_in:
            u = edge.child
            v = edge.parent
            parent[u] = v
            branch_length[u] = time[v] - time[u]
            while v != -1:
                update_result(window_index, v, t_left)
                count[v] += count[u]
                v = parent[v]

        # Update the windows
        while window_index < num_windows and windows[window_index + 1] <= t_right:
            w_right = windows[window_index + 1]
            # This seems like a bad idea as we incur a O(N) cost for each window,
            # where N is the number of nodes.  It might be hard to do much better
            # though, since we can't help but incur O(|sample_set|) cost at each window
            # which we'll assume is O(n), and for large n, N isn't much larger than n.
            # For K > 1 dimensions, the cost of the scan through the nodes is much
            # less than the O(n^K) required to copy (if n is large and K is small).
            # We could keep track of the roots and do a tree traversal, bringing this
            # down to O(n), but this adds a lot of complexity and memory and I'm
            # fairly confident would be slower overall. We could keep a set of
            # non-zero branches, but this would add a O(log n) cost to each edge
            # insertion and removal and a lot of complexity to the C implementation.
            for u in range(ts.num_nodes):
                update_result(window_index, u, w_right)
            window_index += 1
        tree_index += 1

    assert window_index == windows.shape[0] - 1
    if span_normalise:
        for j in range(num_windows):
            result[j] /= windows[j + 1] - windows[j]
    return result


def site_allele_frequency_spectrum(
    ts, sample_sets, windows, polarised=False, span_normalise=True
):
    """
    Efficient implementation of the algorithm used as the basis for the
    underlying C version.
    """
    windows = ts.parse_windows(windows)
    num_windows = windows.shape[0] - 1
    out_dim = [1 + len(sample_set) for sample_set in sample_sets]

    result = np.zeros([num_windows] + out_dim)
    # Add an extra sample set to count across all samples
    sample_sets = list(sample_sets) + [ts.samples()]
    # Number of nodes in sample_set j ancestral to each node u.
    count = np.zeros((ts.num_nodes, len(sample_sets)), dtype=np.uint32)
    for j in range(len(sample_sets)):
        count[sample_sets[j], j] = 1

    site_index = 0
    mutation_index = 0
    window_index = 0
    sites = ts.tables.sites
    mutations = ts.tables.mutations
    parent = np.zeros(ts.num_nodes, dtype=np.int32) - 1
    for (t_left, t_right), edges_out, edges_in in ts.edge_diffs():
        for edge in edges_out:
            u = edge.child
            v = edge.parent
            while v != -1:
                count[v] -= count[u]
                v = parent[v]
            parent[u] = -1

        for edge in edges_in:
            u = edge.child
            v = edge.parent
            parent[u] = v
            while v != -1:
                count[v] += count[u]
                v = parent[v]

        while site_index < len(sites) and sites.position[site_index] < t_right:
            assert t_left <= sites.position[site_index]
            ancestral_state = sites[site_index].ancestral_state
            allele_count = collections.defaultdict(
                functools.partial(np.zeros, len(sample_sets), dtype=int)
            )
            allele_count[ancestral_state][:] = [
                len(sample_set) for sample_set in sample_sets
            ]
            while (
                mutation_index < len(mutations)
                and mutations[mutation_index].site == site_index
            ):
                mutation = mutations[mutation_index]
                allele_count[mutation.derived_state] += count[mutation.node]
                if mutation.parent != -1:
                    parent_allele = mutations[mutation.parent].derived_state
                    allele_count[parent_allele] -= count[mutation.node]
                else:
                    allele_count[ancestral_state] -= count[mutation.node]
                mutation_index += 1

            pos = sites.position[site_index]
            while windows[window_index + 1] <= pos:
                window_index += 1
            assert windows[window_index] <= pos < windows[window_index + 1]
            site_result = result[window_index]

            for allele, c in dict(allele_count).items():
                # Any allele monomorphic across all samples does not
                # contribute to the AFS
                if 0 == c[-1] or c[-1] == ts.num_samples:
                    del allele_count[allele]
            if polarised and ancestral_state in allele_count:
                del allele_count[ancestral_state]

            increment = 1 if polarised else 0.5
            for _allele, c in allele_count.items():
                x = tuple(c[:-1])
                if not polarised:
                    x = fold(x, out_dim)
                site_result[x] += increment
            site_index += 1

    if span_normalise:
        for j in range(num_windows):
            span = windows[j + 1] - windows[j]
            result[j] /= span
    return result


def allele_frequency_spectrum(
    ts, sample_sets, windows=None, polarised=False, mode="site", span_normalise=True
):
    """
    Generalised site frequency spectrum.
    """
    method_map = {
        "site": site_allele_frequency_spectrum,
        "branch": branch_allele_frequency_spectrum,
    }
    return method_map[mode](
        ts,
        sample_sets,
        windows=windows,
        polarised=polarised,
        span_normalise=span_normalise,
    )


class TestAlleleFrequencySpectrum(StatsTestCase, SampleSetStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def verify_single_sample_set(self, ts):
        L = ts.sequence_length
        samples = ts.samples()
        a1 = ts.allele_frequency_spectrum(mode=self.mode)
        a2 = ts.allele_frequency_spectrum([samples], mode=self.mode)
        self.assertArrayEqual(a1, a2)
        for windows in [None, (0, L), (0, L / 2, L)]:
            a1 = ts.allele_frequency_spectrum(mode=self.mode, windows=windows)
            a2 = ts.allele_frequency_spectrum(
                [samples], mode=self.mode, windows=windows
            )
            self.assertArrayEqual(a1, a2)
        for polarised in [True, False]:
            a1 = ts.allele_frequency_spectrum(mode=self.mode, polarised=polarised)
            a2 = ts.allele_frequency_spectrum(
                [samples], mode=self.mode, polarised=polarised
            )
            self.assertArrayEqual(a1, a2)
        for span_normalise in [True, False]:
            a1 = ts.allele_frequency_spectrum(
                mode=self.mode, span_normalise=span_normalise
            )
            a2 = ts.allele_frequency_spectrum(
                [samples], mode=self.mode, span_normalise=span_normalise
            )
            self.assertArrayEqual(a1, a2)

    def verify_sample_sets(self, ts, sample_sets, windows):
        # print(ts.genotype_matrix())
        # print(ts.draw_text())
        # print("sample_sets = ", sample_sets)
        windows = ts.parse_windows(windows)
        for span_normalise, polarised in itertools.product(
            [True, False], [True, False]
        ):
            sfs1 = naive_allele_frequency_spectrum(
                ts,
                sample_sets,
                windows,
                mode=self.mode,
                polarised=polarised,
                span_normalise=span_normalise,
            )
            sfs2 = allele_frequency_spectrum(
                ts,
                sample_sets,
                windows,
                mode=self.mode,
                polarised=polarised,
                span_normalise=span_normalise,
            )
            sfs3 = ts.allele_frequency_spectrum(
                sample_sets,
                windows,
                mode=self.mode,
                polarised=polarised,
                span_normalise=span_normalise,
            )
            assert sfs1.shape[0] == len(windows) - 1
            assert len(sfs1.shape) == len(sample_sets) + 1
            for j, sample_set in enumerate(sample_sets):
                n = 1 + len(sample_set)
                assert sfs1.shape[j + 1] == n

            assert len(sfs1.shape) == len(sample_sets) + 1
            assert sfs1.shape == sfs2.shape
            assert sfs1.shape == sfs3.shape
            if not np.allclose(sfs1, sfs3):
                print()
                print("sample sets", sample_sets)
                print("simple", sfs1)
                print("effic ", sfs2)
                print("ts    ", sfs3)
            self.assertArrayAlmostEqual(sfs1, sfs2)
            self.assertArrayAlmostEqual(sfs1, sfs3)


class TestBranchAlleleFrequencySpectrum(
    TestAlleleFrequencySpectrum, TopologyExamplesMixin
):
    mode = "branch"

    def test_simple_example(self):
        ts = msprime.simulate(6, recombination_rate=0.1, random_seed=1)
        self.verify_single_sample_set(ts)

        self.verify_sample_sets(ts, [range(6)], [0, 1])
        self.verify_sample_sets(ts, [[0, 1]], [0, 1])
        self.verify_sample_sets(ts, [[0, 1], [2, 3]], [0, 1])
        self.verify_sample_sets(ts, [[0, 1, 2, 3, 4, 5]], [0, 1])
        self.verify_sample_sets(ts, [[0, 1, 2], [3, 4, 5]], [0, 1])
        self.verify_sample_sets(ts, [[0, 1], [2, 3], [4, 5]], [0, 1])


class TestSiteAlleleFrequencySpectrum(
    TestAlleleFrequencySpectrum, MutatedTopologyExamplesMixin
):
    mode = "site"

    def test_simple_example(self):
        ts = msprime.simulate(6, mutation_rate=0.2, random_seed=1)
        self.verify_single_sample_set(ts)

        self.verify_sample_sets(ts, [[0]], [0, 1])
        self.verify_sample_sets(ts, [[0, 1, 2, 3, 4, 5]], [0, 1])
        self.verify_sample_sets(ts, [[0, 1, 2], [3, 4, 5]], [0, 1])
        self.verify_sample_sets(ts, [[0, 1], [2, 3], [4, 5]], [0, 1])


class TestBranchAlleleFrequencySpectrumProperties(StatsTestCase, TopologyExamplesMixin):
    def verify(self, ts):
        # If we split by tree, the sum of the AFS should be equal to the
        # tree total branch length in each window
        windows = ts.breakpoints(as_array=True)
        S = ts.samples()
        examples = [
            [S],
            [S[:1]],
            [S[:-1]],
            [S[:1], S[1:]],
            [S[:1], S[:-1]],
        ]
        if len(S) > 2:
            examples += [[S[:1], S[2:], S[:3]]]
        # This is the same definition that we use for segregating_sites
        tbl = [
            sum(
                tree.branch_length(u)
                for u in tree.nodes()
                if 0 < tree.num_samples(u) < ts.num_samples
            )
            for tree in ts.trees()
        ]
        for polarised in [True, False]:
            for sample_sets in examples:
                afs = ts.allele_frequency_spectrum(
                    sample_sets,
                    windows=windows,
                    mode="branch",
                    polarised=polarised,
                    span_normalise=True,
                )
                afs_sum = [np.sum(window) for window in afs]
                self.assertArrayAlmostEqual(afs_sum, tbl)


############################################
# End of specific stats tests.
############################################


class TestWindowedTreeStat(StatsTestCase):
    """
    Tests that the treewise windowing function defined here has the correct
    behaviour.
    """

    # TODO add more tests here covering the various windowing possibilities.
    def get_tree_sequence(self):
        ts = msprime.simulate(10, recombination_rate=2, random_seed=1)
        assert ts.num_trees > 3
        return ts

    def test_all_trees(self):
        ts = self.get_tree_sequence()
        A1 = np.ones((ts.num_trees, 1))
        windows = np.array(list(ts.breakpoints()))
        A2 = windowed_tree_stat(ts, A1, windows)
        # print("breakpoints = ", windows)
        # print(A2)
        assert A1.shape == A2.shape
        # JK: I don't understand what we're computing here, this normalisation
        # seems pretty weird.
        # for tree in ts.trees():
        #     self.assertAlmostEqual(A2[tree.index, 0], tree.span / ts.sequence_length)

    def test_single_interval(self):
        ts = self.get_tree_sequence()
        A1 = np.ones((ts.num_trees, 1))
        windows = np.array([0, ts.sequence_length])
        A2 = windowed_tree_stat(ts, A1, windows)
        assert A2.shape == (1, 1)
        # TODO: Test output


class TestSampleSets(StatsTestCase):
    """
    Tests that passing sample sets in various ways gets interpreted correctly.
    """

    def get_example_ts(self, ts_10_mut_recomb_fixture):
        ts = ts_10_mut_recomb_fixture
        assert ts.num_mutations > 0
        return ts

    def test_duplicate_samples(self, ts_10_mut_recomb_fixture):
        ts = self.get_example_ts(ts_10_mut_recomb_fixture)
        for bad_set in [[1, 1], [1, 2, 1], list(range(10)) + [9]]:
            with pytest.raises(exceptions.LibraryError):
                ts.diversity([bad_set])
            with pytest.raises(exceptions.LibraryError):
                ts.divergence([[0, 1], bad_set])
            with pytest.raises(ValueError):
                ts.sample_count_stat([bad_set], self.identity_f(ts), 1)

    def test_empty_sample_set(self, ts_10_mut_recomb_fixture):
        ts = self.get_example_ts(ts_10_mut_recomb_fixture)
        with pytest.raises(ValueError):
            ts.diversity([[]])
        for bad_sample_sets in [[[], []], [[1], []], [[1, 2], [1], []]]:
            with pytest.raises(ValueError):
                ts.diversity(bad_sample_sets)
            with pytest.raises(ValueError):
                ts.divergence(bad_sample_sets)
            with pytest.raises(ValueError):
                ts.sample_count_stat(bad_sample_sets, self.identity_f(ts), 1)

    def test_non_samples(self, ts_10_mut_recomb_fixture):
        ts = self.get_example_ts(ts_10_mut_recomb_fixture)
        with pytest.raises(exceptions.LibraryError):
            ts.diversity([[ts.num_samples]])

        with pytest.raises(exceptions.LibraryError):
            ts.divergence([[ts.num_samples], [1, 2]])

        with pytest.raises(ValueError):
            ts.sample_count_stat([[ts.num_samples]], self.identity_f(ts), 1)

    def test_span_normalise(self, ts_10_mut_recomb_fixture):
        np.random.seed(92)
        ts = self.get_example_ts(ts_10_mut_recomb_fixture)
        sample_sets = [[0, 1], [2, 3, 4], [5, 6]]
        windows = ts.sequence_length * np.random.uniform(size=10)
        windows.sort()
        windows[0] = 0.0
        windows[-1] = ts.sequence_length
        n = np.array([len(u) for u in sample_sets])

        def f(x):
            return x * (x < n)

        # Determine output_dim of the function
        for mode in ("site", "branch", "node"):
            sigma1 = ts.sample_count_stat(sample_sets, f, 3, windows=windows, mode=mode)
            sigma2 = ts.sample_count_stat(
                sample_sets, f, 3, windows=windows, mode=mode, span_normalise=True
            )
            sigma3 = ts.sample_count_stat(
                sample_sets, f, 3, windows=windows, mode=mode, span_normalise=False
            )
            denom = np.diff(windows)[:, np.newaxis]
            if mode == "node":
                denom = np.diff(windows)[:, np.newaxis, np.newaxis]

            assert sigma1.shape == sigma2.shape
            assert sigma1.shape == sigma3.shape
            self.assertArrayAlmostEqual(sigma1, sigma2)
            self.assertArrayAlmostEqual(sigma1, sigma3 / denom)


class TestSampleSetIndexes(StatsTestCase):
    """
    Tests that we get the correct behaviour from the indexes argument to
    k-way stats functions.
    """

    def get_example_ts(self, ts_10_mut_fixture):
        ts = ts_10_mut_fixture
        assert ts.num_mutations > 0
        return ts

    def test_2_way_default(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        sample_sets = np.array_split(ts.samples(), 2)
        S1 = ts.divergence(sample_sets)
        S2 = divergence(ts, sample_sets)[0, 0]
        S3 = ts.divergence(sample_sets, [0, 1])
        assert S1.shape == S2.shape
        self.assertArrayAlmostEqual(S1, S2)
        self.assertArrayAlmostEqual(S1, S3)
        sample_sets = np.array_split(ts.samples(), 3)
        with pytest.raises(ValueError):
            _ = ts.divergence(sample_sets)
        with pytest.raises(ValueError):
            _ = ts.divergence([sample_sets[0]])

    def test_3_way_default(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        sample_sets = np.array_split(ts.samples(), 3)
        S1 = ts.f3(sample_sets)
        S2 = f3(ts, sample_sets)[0, 0]
        S3 = ts.f3(sample_sets, [0, 1, 2])
        assert S1.shape == S2.shape
        self.assertArrayAlmostEqual(S1, S2)
        self.assertArrayAlmostEqual(S1, S3)
        sample_sets = np.array_split(ts.samples(), 4)
        with pytest.raises(ValueError):
            _ = ts.f3(sample_sets)

    def test_4_way_default(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        sample_sets = np.array_split(ts.samples(), 4)
        S1 = ts.f4(sample_sets)
        S2 = f4(ts, sample_sets)
        S3 = ts.f4(sample_sets, [0, 1, 2, 3])
        assert S1.shape == S3.shape
        self.assertArrayAlmostEqual(S1, S2)
        self.assertArrayAlmostEqual(S1, S3)
        sample_sets = np.array_split(ts.samples(), 5)
        with pytest.raises(ValueError):
            _ = ts.f4(sample_sets)

    def test_2_way_combinations(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        sample_sets = np.array_split(ts.samples(), 4)
        pairs = list(itertools.combinations(range(4), 2))
        for k in range(1, len(pairs)):
            S1 = ts.divergence(sample_sets, pairs[:k])
            S2 = divergence(ts, sample_sets, pairs[:k])[0]
            assert S1.shape[-1] == k
            assert S1.shape == S2.shape
            self.assertArrayAlmostEqual(S1, S2)

    def test_3_way_combinations(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        sample_sets = np.array_split(ts.samples(), 5)
        triples = list(itertools.combinations(range(5), 3))
        for k in range(1, len(triples)):
            S1 = ts.Y3(sample_sets, triples[:k])
            S2 = Y3(ts, sample_sets, triples[:k])[0]
            assert S1.shape[-1] == k
            assert S1.shape == S2.shape
            self.assertArrayAlmostEqual(S1, S2)

    def test_4_way_combinations(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        sample_sets = np.array_split(ts.samples(), 5)
        quads = list(itertools.combinations(range(5), 4))
        for k in range(1, len(quads)):
            S1 = ts.f4(sample_sets, quads[:k], windows=[0, ts.sequence_length])
            S2 = f4(ts, sample_sets, quads[:k])
            assert S1.shape[-1] == k
            assert S2.shape == S2.shape
            self.assertArrayAlmostEqual(S1, S2)

    def test_errors(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        sample_sets = np.array_split(ts.samples(), 2)
        with pytest.raises(ValueError):
            ts.divergence(sample_sets, indexes=[])
        with pytest.raises(ValueError):
            ts.divergence(sample_sets, indexes=[(1, 1, 1)])
        with pytest.raises(exceptions.LibraryError):
            ts.divergence(sample_sets, indexes=[(1, 2)])


class TestGeneralStatInterface(StatsTestCase):
    """
    Tests for the basic interface for general_stats.
    """

    def get_tree_sequence(self):
        ts = msprime.simulate(10, recombination_rate=2, mutation_rate=2, random_seed=1)
        return ts

    def test_function_cannot_update_state(self):
        ts = self.get_tree_sequence()

        def f(x):
            out = x.copy()
            x[:] = 0.0
            return out

        def g(x):
            return x

        x = ts.sample_count_stat(
            [ts.samples()],
            f,
            output_dim=1,
            strict=False,
            mode="node",
            span_normalise=False,
        )
        y = ts.sample_count_stat(
            [ts.samples()],
            g,
            output_dim=1,
            strict=False,
            mode="node",
            span_normalise=False,
        )
        self.assertArrayEqual(x, y)

    def test_default_mode(self, ts_10_recomb_fixture):
        ts = ts_10_recomb_fixture
        W = np.ones((ts.num_samples, 2))
        sigma1 = ts.general_stat(W, self.identity_f(ts), W.shape[1])
        sigma2 = ts.general_stat(W, self.identity_f(ts), W.shape[1], mode="site")
        self.assertArrayEqual(sigma1, sigma2)

    def test_bad_mode(self, ts_10_recomb_fixture):
        ts = ts_10_recomb_fixture
        W = np.ones((ts.num_samples, 2))
        for bad_mode in ["", "MODE", "x" * 8192]:
            with pytest.raises(ValueError):
                ts.general_stat(W, self.identity_f(ts), W.shape[1], mode=bad_mode)

    def test_bad_window_strings(self):
        ts = self.get_tree_sequence()
        with pytest.raises(ValueError):
            ts.diversity([ts.samples()], mode="site", windows="abc")
        with pytest.raises(ValueError):
            ts.diversity([ts.samples()], mode="site", windows="")
        with pytest.raises(ValueError):
            ts.diversity([ts.samples()], mode="tree", windows="abc")

    def test_bad_summary_function(self):
        ts = self.get_tree_sequence()
        W = np.ones((ts.num_samples, 3))
        with pytest.raises(ValueError):
            ts.general_stat(W, lambda x: x, 3, windows="sites")
        with pytest.raises(ValueError):
            ts.general_stat(W, lambda x: np.array([1.0]), 1, windows="sites")

    def test_nonnumpy_summary_function(self):
        ts = self.get_tree_sequence()
        W = np.ones((ts.num_samples, 3))
        sigma1 = ts.general_stat(W, lambda x: [0.0], 1)
        sigma2 = ts.general_stat(W, lambda x: np.array([0.0]), 1)
        self.assertArrayEqual(sigma1, sigma2)


class TestGeneralBranchStats(StatsTestCase):
    """
    Tests for general branch stats (using functions and arbitrary weights)
    """

    def compare_general_stat(self, ts, W, f, windows=None, polarised=False):
        # Determine output_dim of the function
        M = len(f(W[0]))
        sigma1 = naive_branch_general_stat(ts, W, f, windows, polarised=polarised)
        sigma2 = ts.general_stat(W, f, M, windows, polarised=polarised, mode="branch")
        sigma3 = branch_general_stat(ts, W, f, windows, polarised=polarised)
        assert sigma1.shape == sigma2.shape
        assert sigma1.shape == sigma3.shape
        self.assertArrayAlmostEqual(sigma1, sigma2)
        self.assertArrayAlmostEqual(sigma1, sigma3)
        return sigma1

    def test_simple_identity_f_w_zeros(self, ts_12_highrecomb_fixture):
        ts = ts_12_highrecomb_fixture
        W = np.zeros((ts.num_samples, 3))
        for polarised in [True, False]:
            sigma = self.compare_general_stat(
                ts, W, self.identity_f(ts), windows="trees", polarised=polarised
            )
            assert sigma.shape == (ts.num_trees, W.shape[1])
            assert np.all(sigma == 0)

    def test_simple_identity_f_w_ones(self, ts_10_recomb_fixture):
        ts = ts_10_recomb_fixture
        W = np.ones((ts.num_samples, 2))
        sigma = self.compare_general_stat(
            ts, W, self.identity_f(ts), windows="trees", polarised=True
        )
        assert sigma.shape == (ts.num_trees, W.shape[1])
        # A W of 1 for every node and identity f counts the samples in the subtree
        # if polarised is True.
        for tree in ts.trees():
            s = sum(tree.num_samples(u) * tree.branch_length(u) for u in tree.nodes())
            assert np.allclose(sigma[tree.index], s)

    def test_simple_cumsum_f_w_ones(self):
        ts = msprime.simulate(13, recombination_rate=1, random_seed=2)
        W = np.ones((ts.num_samples, 8))
        for polarised in [True, False]:
            sigma = self.compare_general_stat(
                ts, W, self.cumsum_f(ts), windows="trees", polarised=polarised
            )
            assert sigma.shape == (ts.num_trees, W.shape[1])

    def test_simple_cumsum_f_w_ones_many_windows(self):
        ts = msprime.simulate(15, recombination_rate=3, random_seed=3)
        assert ts.num_trees > 3
        windows = np.linspace(0, ts.sequence_length, num=ts.num_trees * 10)
        W = np.ones((ts.num_samples, 3))
        sigma = self.compare_general_stat(ts, W, self.cumsum_f(ts), windows=windows)
        assert sigma.shape == (windows.shape[0] - 1, W.shape[1])

    def test_windows_equal_to_ts_breakpoints(self):
        ts = msprime.simulate(14, recombination_rate=1, random_seed=2)
        W = np.ones((ts.num_samples, 1))
        for polarised in [True, False]:
            sigma_no_windows = self.compare_general_stat(
                ts, W, self.cumsum_f(ts), windows="trees", polarised=polarised
            )
            assert sigma_no_windows.shape == (ts.num_trees, W.shape[1])
            sigma_windows = self.compare_general_stat(
                ts,
                W,
                self.cumsum_f(ts),
                windows=ts.breakpoints(as_array=True),
                polarised=polarised,
            )
            assert sigma_windows.shape == sigma_no_windows.shape
            assert np.allclose(sigma_windows.shape, sigma_no_windows.shape)

    def test_single_tree_windows(self):
        ts = msprime.simulate(15, random_seed=2, length=100)
        W = np.ones((ts.num_samples, 2))
        f = self.sum_f(ts)
        # for num_windows in range(1, 10):
        for num_windows in [2]:
            windows = np.linspace(0, ts.sequence_length, num=num_windows + 1)
            sigma = self.compare_general_stat(ts, W, f, windows)
            assert sigma.shape == (num_windows, 1)

    def test_simple_identity_f_w_zeros_windows(self, ts_12_highrecomb_fixture):
        ts = ts_12_highrecomb_fixture
        W = np.zeros((ts.num_samples, 3))
        f = self.identity_f(ts)
        windows = np.linspace(0, ts.sequence_length, num=11)
        for polarised in [True, False]:
            sigma = self.compare_general_stat(ts, W, f, windows, polarised=polarised)
            assert sigma.shape == (10, W.shape[1])
            assert np.all(sigma == 0)

    def test_nonstrict_nonancestral_material(self):
        # 0 is a sample, 1 is not
        #
        # 2.00┊  2  ┊ 2 ┊  2  ┊
        #     ┊ ┏┻┓ ┊ ┃ ┊ ┏┻┓ ┊
        # 1.00┊ ┃ 1 ┊ 1 ┊ ┃ 1 ┊
        #     ┊ ┃   ┊ ┃ ┊ ┃   ┊
        # 0.00┊ 0   ┊ 0 ┊ 0   ┊
        #     0     1   2     3

        tables = tskit.TableCollection(sequence_length=3)

        node_times = [0, 1, 2]
        samples = [0]
        for n, t in enumerate(node_times):
            tables.nodes.add_row(
                time=t, flags=tskit.NODE_IS_SAMPLE if n in samples else 0
            )

        # p, c, l, r
        edges = [
            (1, 0, 1, 2),
            (2, 0, 0, 1),
            (2, 0, 2, 3),
            (2, 1, 0, 3),
        ]
        for p, c, l, r in edges:
            tables.edges.add_row(parent=p, child=c, left=l, right=r)

        # this makes it so 'site' mode counts branches
        for x in range(int(tables.sequence_length)):
            for n in range(tables.nodes.num_rows - 1):
                offset = n / tables.nodes.num_rows
                s = tables.sites.add_row(position=x + offset, ancestral_state="0")
                tables.mutations.add_row(site=s, node=n, derived_state="1")

        ts = tables.tree_sequence()

        def f(x):
            return x

        for polarised, mode, answer in [
            (True, "branch", 6),
            (True, "site", 4),
            (False, "branch", 8),
            (False, "site", 6),
        ]:
            (stat,) = ts.sample_count_stat(
                [[0]],
                f,
                1,
                strict=False,
                span_normalise=False,
                polarised=polarised,
                mode=mode,
            )
            assert stat == answer


class TestGeneralSiteStats(StatsTestCase):
    """
    Tests for general site stats (using functions and arbitrary weights)
    """

    def compare_general_stat(self, ts, W, f, windows=None, polarised=False):
        # Determine output_dim of the function
        M = len(f(W[0]))
        sigma1 = naive_site_general_stat(ts, W, f, windows, polarised=polarised)
        sigma2 = ts.general_stat(W, f, M, windows, polarised=polarised, mode="site")
        sigma3 = site_general_stat(ts, W, f, windows, polarised=polarised)
        assert sigma1.shape == sigma2.shape
        assert sigma1.shape == sigma3.shape
        self.assertArrayAlmostEqual(sigma1, sigma2)
        self.assertArrayAlmostEqual(sigma1, sigma3)
        return sigma1

    def test_identity_f_W_0_multiple_alleles(self):
        ts = msprime.simulate(20, recombination_rate=0, random_seed=2)
        ts = tsutil.jukes_cantor(ts, 20, 1, seed=10)
        W = np.zeros((ts.num_samples, 3))
        for polarised in [True, False]:
            sigma = self.compare_general_stat(
                ts, W, self.identity_f(ts), windows="sites", polarised=polarised
            )
            assert sigma.shape == (ts.num_sites, W.shape[1])
            assert np.all(sigma == 0)

    def test_identity_f_W_0_multiple_alleles_windows(self):
        ts = msprime.simulate(34, recombination_rate=0, random_seed=2)
        ts = tsutil.jukes_cantor(ts, 20, 1, seed=10)
        W = np.zeros((ts.num_samples, 3))
        windows = np.linspace(0, 1, num=11)
        for polarised in [True, False]:
            sigma = self.compare_general_stat(
                ts, W, self.identity_f(ts), windows=windows, polarised=polarised
            )
            assert sigma.shape == (windows.shape[0] - 1, W.shape[1])
            assert np.all(sigma == 0)

    def test_cumsum_f_W_1_multiple_alleles(self):
        ts = msprime.simulate(3, recombination_rate=2, random_seed=2)
        ts = tsutil.jukes_cantor(ts, 20, 1, seed=10)
        W = np.ones((ts.num_samples, 3))
        for polarised in [True, False]:
            sigma = self.compare_general_stat(
                ts, W, self.cumsum_f(ts), windows="sites", polarised=polarised
            )
            assert sigma.shape == (ts.num_sites, W.shape[1])

    def test_cumsum_f_W_1_two_alleles(self):
        ts = msprime.simulate(33, recombination_rate=1, mutation_rate=2, random_seed=1)
        W = np.ones((ts.num_samples, 5))
        for polarised in [True, False]:
            sigma = self.compare_general_stat(
                ts, W, self.cumsum_f(ts), windows="sites", polarised=polarised
            )
            assert sigma.shape == (ts.num_sites, W.shape[1])


class TestGeneralNodeStats(StatsTestCase):
    """
    Tests for general node stats (using functions and arbitrary weights)
    """

    def compare_general_stat(self, ts, W, f, windows=None, polarised=False):
        # Determine output_dim of the function
        M = len(f(W[0]))
        sigma1 = naive_node_general_stat(ts, W, f, windows, polarised=polarised)
        sigma2 = ts.general_stat(W, f, M, windows, polarised=polarised, mode="node")
        sigma3 = node_general_stat(ts, W, f, windows, polarised=polarised)
        assert sigma1.shape == sigma2.shape
        assert sigma1.shape == sigma3.shape
        self.assertArrayAlmostEqual(sigma1, sigma2)
        self.assertArrayAlmostEqual(sigma1, sigma3)
        return sigma1

    def test_simple_sum_f_w_zeros(self, ts_12_highrecomb_fixture):
        ts = ts_12_highrecomb_fixture
        W = np.zeros((ts.num_samples, 3))
        for polarised in [True, False]:
            sigma = self.compare_general_stat(
                ts, W, self.identity_f(ts), windows="trees", polarised=polarised
            )
            assert sigma.shape == (ts.num_trees, ts.num_nodes, 3)
            assert np.all(sigma == 0)

    def test_simple_sum_f_w_ones(self, ts_44_recomb_fixture):
        ts = ts_44_recomb_fixture
        W = np.ones((ts.num_samples, 2))
        f = self.sum_f(ts)
        sigma = self.compare_general_stat(ts, W, f, windows="trees", polarised=True)
        assert sigma.shape == (ts.num_trees, ts.num_nodes, 1)
        # Drop the last dimension
        sigma = sigma.reshape((ts.num_trees, ts.num_nodes))
        # A W of 1 for every node and f(x)=sum(x) counts the samples in the subtree
        # times 2 if polarised is True.
        for tree in ts.trees():
            s = np.array(
                [
                    tree.num_samples(u) if tree.num_samples(u) < ts.num_samples else 0
                    for u in range(ts.num_nodes)
                ]
            )
            self.assertArrayAlmostEqual(sigma[tree.index], 2 * s)

    def test_simple_sum_f_w_ones_notstrict(self, ts_44_recomb_fixture):
        ts = ts_44_recomb_fixture
        W = np.ones((ts.num_samples, 2))
        sigma = ts.general_stat(
            W,
            lambda x: np.array([np.sum(x)]),
            1,
            windows="trees",
            polarised=True,
            mode="node",
            strict=False,
        )
        assert sigma.shape == (ts.num_trees, ts.num_nodes, 1)
        # Drop the last dimension
        sigma = sigma.reshape((ts.num_trees, ts.num_nodes))
        # A W of 1 for every node and f(x)=sum(x) counts the samples in the subtree
        # times 2 if polarised is True.
        for tree in ts.trees():
            s = np.array([tree.num_samples(u) for u in range(ts.num_nodes)])
            self.assertArrayAlmostEqual(sigma[tree.index], 2 * s)

    def test_small_tree_windows_polarised(self):
        ts = msprime.simulate(4, recombination_rate=0.5, random_seed=2)
        assert ts.num_trees > 1
        W = np.ones((ts.num_samples, 1))
        sigma = self.compare_general_stat(
            ts,
            W,
            self.cumsum_f(ts),
            windows=ts.breakpoints(as_array=True),
            polarised=True,
        )
        assert sigma.shape == (ts.num_trees, ts.num_nodes, 1)

    def test_one_window_polarised(self, ts_4_recomb_fixture):
        ts = ts_4_recomb_fixture
        W = np.ones((ts.num_samples, 1))
        sigma = self.compare_general_stat(
            ts, W, self.cumsum_f(ts), windows=[0, ts.sequence_length], polarised=True
        )
        assert sigma.shape == (1, ts.num_nodes, W.shape[1])

    def test_one_window_unpolarised(self, ts_4_recomb_fixture):
        ts = ts_4_recomb_fixture
        W = np.ones((ts.num_samples, 2))
        sigma = self.compare_general_stat(
            ts, W, self.cumsum_f(ts), windows=[0, ts.sequence_length], polarised=False
        )
        assert sigma.shape == (1, ts.num_nodes, 2)

    def test_many_windows(self):
        ts = msprime.simulate(24, recombination_rate=3, random_seed=2)
        W = np.ones((ts.num_samples, 3))
        for k in [1, ts.num_trees // 2, ts.num_trees, ts.num_trees * 2]:
            windows = np.linspace(0, 1, num=k + 1)
            for polarised in [True]:
                sigma = self.compare_general_stat(
                    ts, W, self.cumsum_f(ts), windows=windows, polarised=polarised
                )
            assert sigma.shape == (k, ts.num_nodes, 3)

    def test_one_tree(self):
        ts = msprime.simulate(10, random_seed=3)
        W = np.ones((ts.num_samples, 2))
        f = self.sum_f(ts, k=2)
        sigma = self.compare_general_stat(ts, W, f, windows=[0, 1], polarised=True)
        assert sigma.shape == (1, ts.num_nodes, 2)
        # A W of 1 for every node and f(x)=sum(x) counts the samples in the subtree
        # times 2 if polarised is True.
        tree = ts.first()
        s = np.array(
            [
                tree.num_samples(u) if tree.num_samples(u) < ts.num_samples else 0
                for u in range(ts.num_nodes)
            ]
        )
        self.assertArrayAlmostEqual(sigma[tree.index, :, 0], 2 * s)
        self.assertArrayAlmostEqual(sigma[tree.index, :, 1], 2 * s)


##############################
# Trait covariance
##############################


@cached_np
def covsq(x, y):
    cov = np.dot(x - np.mean(x), y - np.mean(y)) / (len(x) - 1)
    return cov * cov


@cached_np
def corsq(x, y):
    vx = covsq(x, x)
    vy = covsq(y, y)
    # sqrt is because vx and vy are *squared* variances
    return covsq(x, y) / np.sqrt(vx * vy)


def site_trait_covariance(ts, W, windows=None, span_normalise=True):
    """
    For each site, computes the covariance between the columns of W and the genotypes.
    """
    windows = ts.parse_windows(windows)
    n, K = W.shape
    assert n == ts.num_samples
    out = np.zeros((len(windows) - 1, K))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        haps = ts.genotype_matrix(isolated_as_missing=False)
        site_positions = [x.position for x in ts.sites()]
        for i in range(K):
            w = W[:, i].copy()
            w -= np.mean(w)
            S = 0
            site_in_window = False
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    site_in_window = True
                    hX = haps[k]
                    alleles = set(hX)
                    for a in alleles:
                        S += covsq(w, hX == a) / 2
            if site_in_window:
                out[j, i] = S
                if span_normalise:
                    out[j, i] /= end - begin
    return out


def branch_trait_covariance(ts, W, windows=None, span_normalise=True):
    """
    For each branch, computes the covariance between the columns of W and the split
    induced by the branch, multiplied by the length of the branch.
    """
    windows = ts.parse_windows(windows)
    n, K = W.shape
    assert n == ts.num_samples
    out = np.zeros((len(windows) - 1, K))
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i in range(K):
            w = W[:, i].copy()
            w -= np.mean(w)
            S = 0
            has_trees = False
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                if tr.total_branch_length > 0:
                    has_trees = True
                SS = 0
                for u in range(ts.num_nodes):
                    tree_samples = set(tr.samples(u))
                    below = np.fromiter(
                        (s in tree_samples for s in samples), dtype=bool
                    )
                    branch_length = tr.branch_length(u)
                    SS += covsq(w, below) * branch_length
                S += SS * (min(end, tr.interval.right) - max(begin, tr.interval.left))
            if has_trees:
                out[j, i] = S
                if span_normalise:
                    out[j, i] /= end - begin
    return out


def node_trait_covariance(ts, W, windows=None, span_normalise=True):
    """
    For each node, computes the covariance between the columns of W and the split
    induced by above/below the node.
    """
    windows = ts.parse_windows(windows)
    n, K = W.shape
    assert n == ts.num_samples
    out = np.zeros((len(windows) - 1, ts.num_nodes, K))
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i in range(K):
            w = W[:, i].copy()
            w -= np.mean(w)
            S = np.zeros(ts.num_nodes)
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in range(ts.num_nodes):
                    tree_samples = set(tr.samples(u))
                    below = np.fromiter(
                        (s in tree_samples for s in samples), dtype=bool
                    )
                    SS[u] += covsq(w, below)
                S += SS * (min(end, tr.interval.right) - max(begin, tr.interval.left))
            out[j, :, i] = S
            if span_normalise:
                out[j, :, i] /= end - begin
    return out


def trait_covariance(ts, W, windows=None, mode="site", span_normalise=True):
    method_map = {
        "site": site_trait_covariance,
        "node": node_trait_covariance,
        "branch": branch_trait_covariance,
    }
    return method_map[mode](ts, W, windows=windows, span_normalise=span_normalise)


class TestTraitCovariance(StatsTestCase, WeightStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def get_example_ts(self, ts_10_mut_recomb_fixture):
        ts = ts_10_mut_recomb_fixture
        assert ts.num_mutations > 0
        return ts

    def transform_weights(self, W):
        """
        Need centered weights to compare to general stats.
        """
        W -= np.mean(W, axis=0)
        return W

    def verify_weighted_stat(self, ts, W, windows):
        n = W.shape[0]

        def f(x):
            return (x**2) / (2 * (n - 1) * (n - 1))

        self.verify_definition(ts, W, windows, f, ts.trait_covariance, trait_covariance)

    def verify_interface(self, ts, ts_method):
        W = np.array([np.arange(ts.num_samples)]).T
        sigma1 = ts_method(W, mode=self.mode)
        sigma2 = ts_method(W, windows=None, mode=self.mode)
        sigma3 = ts_method(W, windows=[0.0, ts.sequence_length], mode=self.mode)
        assert sigma1.shape == sigma2.shape
        self.assertArrayAlmostEqual(sigma1, sigma2)
        self.assertArrayAlmostEqual(sigma1, sigma3[0])

    def verify_centering(self, ts, method, ts_method):
        # Since weights are mean-centered, adding a constant shouldn't change anything.
        # ts is already passed as parameter, no need to call get_example_ts()
        for W, windows in subset_combos(
            self.example_weights(ts), example_windows(ts), p=0.1
        ):
            shift = np.arange(1, W.shape[1] + 1)
            sigma1 = ts_method(W, windows=windows, mode=self.mode)
            sigma2 = ts_method(W + shift, windows=windows, mode=self.mode)
            sigma3 = method(ts, W, windows=windows, mode=self.mode)
            sigma4 = method(ts, W + shift, windows=windows, mode=self.mode)
            assert sigma1.shape == sigma2.shape
            assert sigma1.shape == sigma3.shape
            assert sigma1.shape == sigma4.shape
            self.assertArrayAlmostEqual(sigma1, sigma2)
            self.assertArrayAlmostEqual(sigma1, sigma3)
            self.assertArrayAlmostEqual(sigma1, sigma4)


class TraitCovarianceMixin:
    def test_interface(self, ts_10_mut_recomb_fixture):
        ts = self.get_example_ts(ts_10_mut_recomb_fixture)
        self.verify_interface(ts, ts.trait_covariance)

    def test_normalisation(self, ts_10_mut_recomb_fixture):
        ts = self.get_example_ts(ts_10_mut_recomb_fixture)
        self.verify_centering(ts, trait_covariance, ts.trait_covariance)

    def test_errors(self, ts_10_mut_recomb_fixture):
        ts = self.get_example_ts(ts_10_mut_recomb_fixture)
        W = np.ones((ts.num_samples, 2))
        # W must have the right number of rows
        with pytest.raises(ValueError):
            ts.trait_correlation(W[1:, :])


@pytest.mark.slow
class TestBranchTraitCovariance(
    TestTraitCovariance, TopologyExamplesMixin, TraitCovarianceMixin
):
    mode = "branch"


@pytest.mark.slow
class TestNodeTraitCovariance(
    TestTraitCovariance, TopologyExamplesMixin, TraitCovarianceMixin
):
    mode = "node"


class TestSiteTraitCovariance(
    TestTraitCovariance, MutatedTopologyExamplesMixin, TraitCovarianceMixin
):
    mode = "site"


##############################
# Trait correlation
##############################


def site_trait_correlation(ts, W, windows=None, span_normalise=True):
    """
    For each site, computes the correlation between the columns of W and the genotypes.
    """
    windows = ts.parse_windows(windows)
    n, K = W.shape
    assert n == ts.num_samples
    out = np.zeros((len(windows) - 1, K))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        haps = ts.genotype_matrix(isolated_as_missing=False)
        site_positions = [x.position for x in ts.sites()]
        for i in range(K):
            w = W[:, i].copy()
            w -= np.mean(w)
            w /= np.std(w) * np.sqrt(len(w) / (len(w) - 1))
            S = 0
            site_in_window = False
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    site_in_window = True
                    hX = haps[k]
                    alleles = set(hX)
                    for a in alleles:
                        p = np.mean(hX == a)
                        if p > 0 and p < 1:
                            # S += sum(w[hX == a])**2 / (2 * (p * (1 - p)))
                            S += corsq(w, hX == a) / 2
            if site_in_window:
                out[j, i] = S
                if span_normalise:
                    out[j, i] /= end - begin
    return out


def branch_trait_correlation(ts, W, windows=None, span_normalise=True):
    """
    For each branch, computes the correlation between the columns of W and the split
    induced by the branch, multiplied by the length of the branch.
    """
    windows = ts.parse_windows(windows)
    n, K = W.shape
    assert n == ts.num_samples
    out = np.zeros((len(windows) - 1, K))
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i in range(K):
            w = W[:, i].copy()
            w -= np.mean(w)
            w /= np.std(w) * np.sqrt(len(w) / (len(w) - 1))
            S = 0
            has_trees = False
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                if tr.total_branch_length > 0:
                    has_trees = True
                SS = 0
                for u in range(ts.num_nodes):
                    tree_samples = set(tr.samples(u))
                    below = np.fromiter(
                        (s in tree_samples for s in samples), dtype=bool
                    )
                    p = np.mean(below)
                    if p > 0 and p < 1:
                        branch_length = tr.branch_length(u)
                        # SS += ((sum(w[below])**2 +
                        #         sum(w[np.logical_not(below)])**2) * branch_length
                        #        / (2 * (p * (1 - p))))
                        SS += corsq(w, below) * branch_length
                S += SS * (min(end, tr.interval.right) - max(begin, tr.interval.left))
            if has_trees:
                out[j, i] = S
                if span_normalise:
                    out[j, i] /= end - begin
    return out


def node_trait_correlation(ts, W, windows=None, span_normalise=True):
    """
    For each node, computes the correlation between the columns of W and the split
    induced by above/below the node.
    """
    windows = ts.parse_windows(windows)
    n, K = W.shape
    assert n == ts.num_samples
    out = np.zeros((len(windows) - 1, ts.num_nodes, K))
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i in range(K):
            w = W[:, i].copy()
            w -= np.mean(w)
            w /= np.std(w) * np.sqrt(len(w) / (len(w) - 1))
            S = np.zeros(ts.num_nodes)
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in range(ts.num_nodes):
                    tree_samples = set(tr.samples(u))
                    below = np.fromiter(
                        (s in tree_samples for s in samples), dtype=bool
                    )
                    p = np.mean(below)
                    if p > 0 and p < 1:
                        # SS[u] += sum(w[below])**2 / 2
                        # SS[u] += sum(w[np.logical_not(below)])**2 / 2
                        # SS[u] /= (p * (1 - p))
                        SS[u] += corsq(w, below)
                S += SS * (min(end, tr.interval.right) - max(begin, tr.interval.left))
            out[j, :, i] = S
            if span_normalise:
                out[j, :, i] /= end - begin
    return out


def trait_correlation(ts, W, windows=None, mode="site", span_normalise=True):
    method_map = {
        "site": site_trait_correlation,
        "node": node_trait_correlation,
        "branch": branch_trait_correlation,
    }
    return method_map[mode](ts, W, windows=windows, span_normalise=span_normalise)


class TestTraitCorrelation(TestTraitCovariance):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def transform_weights(self, W):
        """
        Need standardised weights to compare to general stats,
        and also an extra column to compute allele frequencies.
        """
        W -= np.mean(W, axis=0)
        n = W.shape[0]
        with suppress_division_by_zero_warning():
            W /= np.std(W, axis=0) * np.sqrt(n / (n - 1))
        return np.column_stack((W, np.ones(W.shape[0]) / W.shape[0]))

    def verify_weighted_stat(self, ts, W, windows):
        n = W.shape[0]

        def f(x):
            p = x[-1]
            if p > 0 and p < 1:
                return (x[:-1] ** 2) / (2 * (p * (1 - p)) * n * (n - 1))
            else:
                return x[:-1] * 0.0

        self.verify_definition(
            ts, W, windows, f, ts.trait_correlation, trait_correlation
        )

    def test_errors(self, ts_10_mut_recomb_fixture):
        ts = self.get_example_ts(ts_10_mut_recomb_fixture)
        # columns of W must have positive SD
        W = np.ones((ts.num_samples, 2))
        with pytest.raises(ValueError):
            ts.trait_correlation(W)
        # W must have the right number of rows
        with pytest.raises(ValueError):
            ts.trait_correlation(W[1:, :])

    def verify_standardising(self, ts, method, ts_method):
        """
        Since weights are standardised, multiplying by a constant shouldn't
        change anything.
        """
        for W, windows in subset_combos(
            self.example_weights(ts), example_windows(ts), p=0.1
        ):
            scale = np.arange(1, W.shape[1] + 1)
            sigma1 = ts_method(W, windows=windows, mode=self.mode)
            sigma2 = ts_method(W * scale, windows=windows, mode=self.mode)
            sigma3 = method(ts, W, windows=windows, mode=self.mode)
            sigma4 = method(ts, W * scale, windows=windows, mode=self.mode)
            assert sigma1.shape == sigma2.shape
            self.assertArrayAlmostEqual(sigma1, sigma2)
            self.assertArrayAlmostEqual(sigma1, sigma3)
            self.assertArrayAlmostEqual(sigma1, sigma4)


class TraitCorrelationMixin:
    def test_interface(self, ts_10_mut_recomb_fixture):
        ts = self.get_example_ts(ts_10_mut_recomb_fixture)
        self.verify_interface(ts, ts.trait_correlation)

    def test_normalisation(self, ts_10_mut_recomb_fixture):
        ts = self.get_example_ts(ts_10_mut_recomb_fixture)
        self.verify_centering(ts, trait_correlation, ts.trait_correlation)
        self.verify_standardising(ts, trait_correlation, ts.trait_correlation)


@pytest.mark.slow
class TestBranchTraitCorrelation(
    TestTraitCorrelation, TopologyExamplesMixin, TraitCorrelationMixin
):
    mode = "branch"


@pytest.mark.slow
class TestNodeTraitCorrelation(
    TestTraitCorrelation, TopologyExamplesMixin, TraitCorrelationMixin
):
    mode = "node"


class TestSiteTraitCorrelation(
    TestTraitCorrelation, MutatedTopologyExamplesMixin, TraitCorrelationMixin
):
    mode = "site"


##############################
# Trait linear_model
##############################

# Quick hack to speed up the tests a bit. We're running this linear_model
# function a gazillion times with similar arguments, so it's worth
# caching the results. We could use functools.lru_cache, but the arguments
# are numpy arrays and so we have to do something different. This
# reduces runtime from ~40 seconds to ~7 seconds on the
# TestNodeTraitLinearModel.
# See https://github.com/tskit-dev/tskit/issues/1856 for more info
_lm_cache = {}


@cached_np
def linear_model(y, x, z):
    key = (y.tobytes(), x.tobytes(), z.tobytes())
    if key not in _lm_cache:
        _lm_cache[key] = _linear_model(y, x, z)
    return _lm_cache[key]


def _linear_model(y, x, z):
    """
    Returns the squared coefficient of x in the least-squares linear model
    :   y ~ x + z
    where x and y are vectors and z is a matrix.
    Note that if z is None then the output is
      cor(x, y) * sd(y) / sd(x) = cov(x, y) / (sd(x) ** 2) .
    """
    # add the constant vector to z
    if z is None:
        z = np.ones((len(x), 1))
    else:
        xz = np.column_stack([z, np.ones((len(x), 1))])
        if np.linalg.matrix_rank(xz) == xz.shape[1]:
            z = xz
    xz = np.column_stack([x, z])
    # check if y is sufficiently independent of the subspace spanned by xz
    Pz = np.matmul(z, np.linalg.pinv(z))
    Py = np.matmul(Pz, y)
    denom = np.sum((y - Py) ** 2)
    if np.linalg.matrix_rank(xz) < xz.shape[1] or denom < 1e-8:
        return 0.0
    else:
        coefs, _, _, _ = np.linalg.lstsq(xz, y, rcond=None)
        return coefs[0] * coefs[0]


def site_trait_linear_model(ts, W, Z, windows=None, span_normalise=True):
    """
    For each site, and for each trait w (column of W), computes the coefficient
    of site in the linear model:
      w ~ site + Z
    """
    windows = ts.parse_windows(windows)
    n, K = W.shape
    assert n == ts.num_samples
    out = np.zeros((len(windows) - 1, K))
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        haps = ts.genotype_matrix(isolated_as_missing=False)
        site_positions = [x.position for x in ts.sites()]
        for i in range(K):
            w = W[:, i]
            S = 0
            site_in_window = False
            for k in range(ts.num_sites):
                if (site_positions[k] >= begin) and (site_positions[k] < end):
                    site_in_window = True
                    hX = haps[k]
                    alleles = set(hX)
                    for a in alleles:
                        p = np.mean(hX == a)
                        if p > 0 and p < 1:
                            S += linear_model(w, hX == a, Z) / 2
            if site_in_window:
                out[j, i] = S
                if span_normalise:
                    out[j, i] /= end - begin
    return out


def branch_trait_linear_model(ts, W, Z, windows=None, span_normalise=True):
    """
    For each branch, fits the linear_model of each column of W onto the split
    induced by the branch and the covariates Z, multiplied by the length of the branch,
    returning the squared coefficient of the column of W.
    """
    windows = ts.parse_windows(windows)
    n, K = W.shape
    assert n == ts.num_samples
    out = np.zeros((len(windows) - 1, K))
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i in range(K):
            w = W[:, i]
            S = 0
            has_trees = False
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                if tr.total_branch_length > 0:
                    has_trees = True
                SS = 0
                for u in range(ts.num_nodes):
                    tree_samples = set(tr.samples(u))
                    below = np.fromiter(
                        (s in tree_samples for s in samples), dtype=bool
                    )
                    branch_length = tr.branch_length(u)
                    SS += linear_model(w, below, Z) * branch_length
                S += SS * (min(end, tr.interval.right) - max(begin, tr.interval.left))
            if has_trees:
                out[j, i] = S
                if span_normalise:
                    out[j, i] /= end - begin
    return out


def node_trait_linear_model(ts, W, Z, windows=None, span_normalise=True):
    """
    For each node, fits the linear model of each columns of W on the split
    induced by above/below the node and the covariates Z, returning the squared
    coefficient of the column of W.
    """
    windows = ts.parse_windows(windows)
    n, K = W.shape
    assert n == ts.num_samples
    out = np.zeros((len(windows) - 1, ts.num_nodes, K))
    samples = ts.samples()
    for j in range(len(windows) - 1):
        begin = windows[j]
        end = windows[j + 1]
        for i in range(K):
            w = W[:, i]
            S = np.zeros(ts.num_nodes)
            for tr in ts.trees():
                if tr.interval.right <= begin:
                    continue
                if tr.interval.left >= end:
                    break
                SS = np.zeros(ts.num_nodes)
                for u in range(ts.num_nodes):
                    tree_samples = set(tr.samples(u))
                    below = np.fromiter(
                        (s in tree_samples for s in samples), dtype=bool
                    )
                    SS[u] += linear_model(w, below, Z)
                S += SS * (min(end, tr.interval.right) - max(begin, tr.interval.left))
            out[j, :, i] = S
            if span_normalise:
                out[j, :, i] /= end - begin
    return out


def trait_linear_model(ts, W, Z, windows=None, mode="site", span_normalise=True):
    method_map = {
        "site": site_trait_linear_model,
        "node": node_trait_linear_model,
        "branch": branch_trait_linear_model,
    }
    return method_map[mode](ts, W, Z, windows=windows, span_normalise=span_normalise)


class TestTraitLinearModel(StatsTestCase, WeightStatsMixin):
    # Derived classes define this to get a specific stats mode.
    mode = None

    def get_example_ts(self, ts_10_mut_recomb_fixture):
        ts = ts_10_mut_recomb_fixture
        assert ts.num_mutations > 0
        return ts

    def example_covariates(self, ts, k_values=None):
        if k_values is None:
            k_values = [2]  # Default to [2] to maintain current optimization

        np.random.seed(999)
        N = ts.num_samples

        for k in k_values:
            k = min(k, ts.num_samples)

            # Uniform covariates
            Z = np.ones((N, k))
            Z[1, :] = np.arange(k, 2 * k)
            yield Z

            # Include one normal case for test coverage
            if N >= 6:  # Only for larger samples to reduce computations
                Z_normal = np.ones((N, k))
                for j in range(k):
                    Z_normal[:, j] = np.random.normal(0, 1, N)
                yield Z_normal

    def transform_weights(self, W, Z):
        n = W.shape[0]
        return np.column_stack([W, Z, np.ones((n, 1))])

    def transform_covariates(self, Z):
        tZ = np.column_stack([Z, np.ones((Z.shape[0], 1))])
        if np.linalg.matrix_rank(tZ) == tZ.shape[1]:
            Z = tZ
        assert np.linalg.matrix_rank(Z) == Z.shape[1]
        K = np.linalg.cholesky(np.matmul(Z.T, Z)).T
        Z = np.matmul(Z, np.linalg.inv(K))
        return Z

    def verify(self, ts):
        for W, Z, windows in subset_combos(
            self.example_weights(ts),
            self.example_covariates(ts, k_values=[2]),
            example_windows(ts),
            p=0.02,  # Reduced from 0.04 for performance
        ):
            self.verify_trait_linear_model(ts, W, Z, windows=windows)

    def verify_trait_linear_model(self, ts, W, Z, windows):
        n, result_dim = W.shape
        tZ = self.transform_covariates(Z)
        n, k = tZ.shape
        V = np.matmul(W.T, tZ)

        def f(x):
            m = x[-1]
            a = np.zeros(result_dim)
            for i in range(result_dim):
                # print("i=", i, "result_dim=", result_dim, "m=", m, "x=", x)
                # print("V=", V)
                if m > 0 and m < ts.num_samples:
                    v = V[i, :]
                    a[i] = x[i]
                    denom = m
                    for j in range(k):
                        xx = x[result_dim + j]
                        a[i] -= xx * v[j]
                        denom -= xx * xx
                    if abs(denom) < 1e-8:
                        a[i] = 0.0
                    else:
                        a[i] /= denom
                else:
                    a[i] = 0.0
            # print("out", a*a/2)
            return a * a / 2

        # general_stat will need Z added, and an extra column for m
        gW = self.transform_weights(W, tZ)

        def wrapped_summary_func(x):
            with suppress_division_by_zero_warning():
                return f(x)

        # Determine output_dim of the function
        M = len(wrapped_summary_func(gW[0]))
        for sn in [True, False]:
            sigma1 = ts.general_stat(
                gW, wrapped_summary_func, M, windows, mode=self.mode, span_normalise=sn
            )
            sigma2 = general_stat(
                ts, gW, wrapped_summary_func, windows, mode=self.mode, span_normalise=sn
            )
            sigma3 = ts.trait_linear_model(
                W, Z, windows=windows, mode=self.mode, span_normalise=sn
            )
            sigma4 = trait_linear_model(
                ts, W, Z, windows=windows, mode=self.mode, span_normalise=sn
            )

            assert sigma1.shape == sigma2.shape
            assert sigma1.shape == sigma3.shape
            assert sigma1.shape == sigma4.shape
            self.assertArrayAlmostEqual(sigma1, sigma2)
            self.assertArrayAlmostEqual(sigma1, sigma3)
            self.assertArrayAlmostEqual(sigma1, sigma4)


class TraitLinearModelMixin:
    def test_interface(self, ts_10_mut_recomb_fixture):
        ts = self.get_example_ts(ts_10_mut_recomb_fixture)
        W = np.array([np.arange(ts.num_samples)]).T
        Z = np.ones((ts.num_samples, 1))
        sigma1 = ts.trait_linear_model(W, Z=Z, mode=self.mode)
        sigma2 = ts.trait_linear_model(W, Z=Z, windows=None, mode=self.mode)
        sigma3 = ts.trait_linear_model(
            W, Z=Z, windows=[0.0, ts.sequence_length], mode=self.mode
        )
        sigma4 = ts.trait_linear_model(
            W, Z=None, windows=[0.0, ts.sequence_length], mode=self.mode
        )
        assert sigma1.shape == sigma2.shape
        assert sigma3.shape[0] == 1
        assert sigma1.shape == sigma3.shape[1:]
        assert sigma1.shape == sigma4.shape[1:]
        self.assertArrayAlmostEqual(sigma1, sigma2)
        self.assertArrayAlmostEqual(sigma1, sigma3[0])
        self.assertArrayAlmostEqual(sigma1, sigma4[0])

    def test_errors(self, ts_10_mut_recomb_fixture):
        ts = self.get_example_ts(ts_10_mut_recomb_fixture)
        W = np.array([np.arange(ts.num_samples)]).T
        Z = np.ones((ts.num_samples, 1))
        # singular covariates
        with pytest.raises(ValueError):
            ts.trait_linear_model(
                W,
                np.ones((ts.num_samples, 2)),
                mode=self.mode,
            )
        # wrong dimensions of W
        with pytest.raises(ValueError):
            ts.trait_linear_model(W[1:, :], Z, mode=self.mode)
        # wrong dimensions of Z
        with pytest.raises(ValueError):
            ts.trait_linear_model(W, Z[1:, :], mode=self.mode)

    def test_deprecation(self, ts_10_mut_recomb_fixture):
        ts = self.get_example_ts(ts_10_mut_recomb_fixture)
        W = np.array([np.arange(ts.num_samples)]).T
        Z = np.ones((ts.num_samples, 1))
        with pytest.warns(FutureWarning):
            ts.trait_regression(W, Z=Z, mode=self.mode)


@pytest.mark.slow
class TestBranchTraitLinearModel(
    TestTraitLinearModel, TopologyExamplesMixin, TraitLinearModelMixin
):
    mode = "branch"


@pytest.mark.slow
class TestNodeTraitLinearModel(
    TestTraitLinearModel, TopologyExamplesMixin, TraitLinearModelMixin
):
    mode = "node"


class TestSiteTraitLinearModel(
    TestTraitLinearModel, MutatedTopologyExamplesMixin, TraitLinearModelMixin
):
    mode = "site"


##############################
# Sample set statistics
##############################


@pytest.mark.skip(reason="Broken - need to port tests")
class SampleSetStatTestCase(StatsTestCase):
    """
    Provides checks for testing of sample set-based statistics.  Actual testing
    is done by derived classes, which should have attributes `stat_type` and `rng`.
    This works by using parallel structure between different statistic "modes",
    in tree sequence methods (with stat_type=X) and python stat calculators as
    implemented here (with StatCalculator.X).
    """

    random_seed = 123456

    def compare_sfs(self, ts, tree_fn, sample_sets, tsc_fn):
        for sample_set in sample_sets:
            windows = [
                k * ts.sequence_length / 20
                for k in [0] + sorted(self.rng.sample(range(1, 20), 4)) + [20]
            ]
            win_args = [
                {"begin": windows[i], "end": windows[i + 1]}
                for i in range(len(windows) - 1)
            ]
            tree_vals = [tree_fn(sample_set, **b) for b in win_args]

            tsc_vals = tsc_fn(sample_set, windows)
            assert len(tsc_vals) == len(windows) - 1
            for i in range(len(windows) - 1):
                self.assertListAlmostEqual(tsc_vals[i], tree_vals[i])

    def check_sfs_interface(self, ts):
        samples = ts.samples()

        # empty sample sets will raise an error
        with pytest.raises(ValueError):
            ts.site_frequency_spectrum([], self.stat_type)
        # sample_sets must be lists without repeated elements
        with pytest.raises(ValueError):
            ts.site_frequency_spectrum(
                [samples[2], samples[2]],
                self.stat_type,
            )
        # and must all be samples
        with pytest.raises(ValueError):
            ts.site_frequency_spectrum(
                [samples[0], max(samples) + 1],
                self.stat_type,
            )
        # windows must start at 0.0, be increasing, and extend to the end
        with pytest.raises(ValueError):
            ts.site_frequency_spectrum(
                samples[0:2],
                [0.1, ts.sequence_length],
                self.stat_type,
            )
        with pytest.raises(ValueError):
            ts.site_frequency_spectrum(
                samples[0:2],
                [0.0, 0.8 * ts.sequence_length],
                self.stat_type,
            )
        with pytest.raises(ValueError):
            ts.site_frequency_spectrum(
                samples[0:2],
                [
                    0.0,
                    0.8 * ts.sequence_length,
                    0.4 * ts.sequence_length,
                    ts.sequence_length,
                ],
                self.stat_type,
            )

    def check_sfs(self, ts):
        # check site frequency spectrum
        self.check_sfs_interface(ts)
        A = [
            self.rng.sample(list(ts.samples()), 2),
            self.rng.sample(list(ts.samples()), 4),
            self.rng.sample(list(ts.samples()), 8),
            self.rng.sample(list(ts.samples()), 10),
            self.rng.sample(list(ts.samples()), 12),
        ]
        py_tsc = self.py_stat_class(ts)

        self.compare_sfs(
            ts, py_tsc.site_frequency_spectrum, A, ts.site_frequency_spectrum
        )


class BranchSampleSetStatsTestCase(SampleSetStatTestCase):
    """
    Tests of branch statistic computation with sample sets,
    mostly running the checks in SampleSetStatTestCase.
    """

    def setUp(self):
        self.rng = random.Random(self.random_seed)
        self.stat_type = "branch"

    def get_ts(self):
        for N in [12, 15, 20]:
            yield msprime.simulate(
                N, random_seed=self.random_seed, recombination_rate=10
            )

    @pytest.mark.skip(reason="Skipping SFS.")
    def test_sfs_interface(self):
        ts = msprime.simulate(10)
        tsc = tskit.BranchStatCalculator(ts)

        # Duplicated samples raise an error
        with pytest.raises(ValueError):
            tsc.site_frequency_spectrum([1, 1])
        with pytest.raises(ValueError):
            tsc.site_frequency_spectrum([])
        with pytest.raises(ValueError):
            tsc.site_frequency_spectrum([0, 11])
        # Check for bad windows
        for bad_start in [-1, 1, 1e-7]:
            with pytest.raises(ValueError):
                tsc.site_frequency_spectrum(
                    [1, 2],
                    [bad_start, ts.sequence_length],
                )
        for bad_end in [0, ts.sequence_length - 1, ts.sequence_length + 1]:
            with pytest.raises(ValueError):
                tsc.site_frequency_spectrum([1, 2], [0, bad_end])
        # Windows must be increasing.
        with pytest.raises(ValueError):
            tsc.site_frequency_spectrum([1, 2], [0, 1, 1])

    @pytest.mark.skip(reason="No SFS.")
    def test_branch_sfs(self):
        for ts in self.get_ts():
            self.check_sfs(ts)


class SpecificTreesTestCase(StatsTestCase):
    """
    Some particular cases, that are easy to see and debug.
    """

    seed = 21

    def test_case_1(self):
        # With mutations:
        #
        # 1.0          6
        # 0.7         / \                                    5
        #            /   X                                  / \
        # 0.5       X     4                4               /   4
        #          /     / \              / \             /   X X
        # 0.4     X     X   \            X   3           X   /   \
        #        /     /     X          /   / X         /   /     \
        # 0.0   0     1       2        1   0   2       0   1       2
        #          (0.0, 0.2),        (0.2, 0.8),       (0.8, 1.0)
        #
        # haplotypes:
        # site:  0   1 2  3 4         5    6              7   8 9
        # 0:     0   1 0  1 0         0    0              1   0 0
        # 1:     1   0 0  0 1         1    0              0   1 0
        # 2:     1   0 1  0 0         0    1              0   0 1
        branch_true_diversity_01 = 2 * (
            1 * (0.2 - 0) + 0.5 * (0.8 - 0.2) + 0.7 * (1.0 - 0.8)
        )
        branch_true_diversity_02 = 2 * (
            1 * (0.2 - 0) + 0.4 * (0.8 - 0.2) + 0.7 * (1.0 - 0.8)
        )
        branch_true_diversity_12 = 2 * (
            0.5 * (0.2 - 0) + 0.5 * (0.8 - 0.2) + 0.5 * (1.0 - 0.8)
        )
        branch_true_Y = 0.2 * (1 + 0.5) + 0.6 * (0.4) + 0.2 * (0.7 + 0.2)
        site_true_Y = 3 + 0 + 1
        node_true_diversity_012 = (
            np.array(
                [
                    0.2 * np.array([2, 2, 2, 0, 2, 0, 0])
                    + 0.6 * np.array([2, 2, 2, 2, 0, 0, 0])
                    + 0.2 * np.array([2, 2, 2, 0, 2, 0, 0])
                ]
            )
            / 3
        )
        node_true_divergence_0_12 = (
            np.array(
                [
                    0.2 * np.array([2, 1, 1, 0, 2, 0, 0])
                    + 0.6 * np.array([2, 1, 1, 1, 0, 0, 0])
                    + 0.2 * np.array([2, 1, 1, 0, 2, 0, 0])
                ]
            )
            / 2
        )
        haplotypes = np.array(
            [
                [0, 1, 1],
                [1, 0, 0],
                [0, 0, 1],
                [1, 0, 0],
                [0, 1, 0],
                [0, 1, 0],
                [0, 0, 1],
                [1, 0, 0],
                [0, 1, 0],
                [0, 0, 1],
            ]
        )
        traits = np.array([[1, 2, 3, 0], [-5, 0, 1, 1], [3, 4, 1.2, 2]])
        # nb: verified the following with R
        true_cov = (
            np.cov(haplotypes, traits.T)[: haplotypes.shape[0], haplotypes.shape[0] :]
            ** 2
        )
        true_cor = (
            np.corrcoef(haplotypes, traits.T)[
                : haplotypes.shape[0], haplotypes.shape[0] :
            ]
            ** 2
        )
        cov02 = np.cov(np.array([1, 0, 1]), traits.T)[:1, 1:] ** 2
        true_branch_cov = (
            true_cov[1, :] * 1.0 * 0.2
            + true_cov[4, :] * 0.5 * 0.2  # branch 0, tree 0
            + true_cov[2, :] * 0.5 * 0.2  # branch 1, tree 0
            + true_cov[0, :] * 0.5 * 0.2  # branch 2, tree 0
            + true_cov[1, :] * 0.4 * 0.6  # branch 4, tree 0
            + true_cov[4, :] * 0.5 * 0.6  # branch 0, tree 1
            + true_cov[2, :] * 0.4 * 0.6  # branch 1, tree 1
            + cov02 * 0.1 * 0.6  # branch 2, tree 1
            + true_cov[1, :] * 0.7 * 0.2  # branch 3, tree 1
            + true_cov[4, :] * 0.5 * 0.2  # branch 0, tree 2
            + true_cov[2, :] * 0.5 * 0.2  # branch 1, tree 2
            + true_cov[0, :] * 0.2 * 0.2  # branch 2, tree 2
        )  # branch 4, tree 2
        cor02 = np.corrcoef(np.array([1, 0, 1]), traits.T)[:1, 1:] ** 2
        true_branch_cor = (
            true_cor[1, :] * 1.0 * 0.2
            + true_cor[4, :] * 0.5 * 0.2  # branch 0, tree 0
            + true_cor[2, :] * 0.5 * 0.2  # branch 1, tree 0
            + true_cor[0, :] * 0.5 * 0.2  # branch 2, tree 0
            + true_cor[1, :] * 0.4 * 0.6  # branch 4, tree 0
            + true_cor[4, :] * 0.5 * 0.6  # branch 0, tree 1
            + true_cor[2, :] * 0.4 * 0.6  # branch 1, tree 1
            + cor02 * 0.1 * 0.6  # branch 2, tree 1
            + true_cor[1, :] * 0.7 * 0.2  # branch 3, tree 1
            + true_cor[4, :] * 0.5 * 0.2  # branch 0, tree 2
            + true_cor[2, :] * 0.5 * 0.2  # branch 1, tree 2
            + true_cor[0, :] * 0.2 * 0.2  # branch 2, tree 2
        )  # branch 4, tree 2

        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       0           0.4
        4       0           0.5
        5       0           0.7
        6       0           1.0
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.2     0.8     3       0,2
        0.0     0.2     4       1,2
        0.2     0.8     4       1,3
        0.8     1.0     4       1,2
        0.8     1.0     5       0,4
        0.0     0.2     6       0,4
        """
        )
        sites = io.StringIO(
            """\
        id  position    ancestral_state
        0   0.05        0
        1   0.1         0
        2   0.11        0
        3   0.15        0
        4   0.151       0
        5   0.3         0
        6   0.6         0
        7   0.9         0
        8   0.95        0
        9   0.951       0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    derived_state
        0       4       1
        1       0       1
        2       2       1
        3       0       1
        4       1       1
        5       1       1
        6       2       1
        7       0       1
        8       1       1
        9       2       1
        """
        )
        ts = tskit.load_text(
            nodes=nodes, edges=edges, sites=sites, mutations=mutations, strict=False
        )

        # diversity between 0 and 1
        A = [[0], [1]]
        n = [len(a) for a in A]

        def f(x):
            return np.array(
                [float(x[0] * (n[1] - x[1]) + (n[0] - x[0]) * x[1]) / (2 * n[0] * n[1])]
            )

        # tree lengths:
        mode = "branch"
        self.assertAlmostEqual(
            divergence(ts, [[0], [1]], [(0, 1)], mode=mode), branch_true_diversity_01
        )
        self.assertAlmostEqual(
            ts.divergence([[0], [1]], [(0, 1)], mode=mode), branch_true_diversity_01
        )
        self.assertAlmostEqual(
            ts.sample_count_stat(A, f, 1, mode=mode)[0], branch_true_diversity_01
        )
        self.assertAlmostEqual(
            ts.diversity([[0, 1]], mode=mode)[0], branch_true_diversity_01
        )

        # mean diversity between [0, 1] and [0, 2]:
        branch_true_mean_diversity = (
            0
            + branch_true_diversity_02
            + branch_true_diversity_01
            + branch_true_diversity_12
        ) / 4
        A = [[0, 1], [0, 2]]
        n = [len(a) for a in A]

        def f(x):
            return np.array([float(x[0] * (n[1] - x[1]) + (n[0] - x[0]) * x[1]) / 8.0])

        # tree lengths:
        self.assertAlmostEqual(
            divergence(ts, [A[0], A[1]], [(0, 1)], mode=mode),
            branch_true_mean_diversity,
        )
        self.assertAlmostEqual(
            ts.divergence([A[0], A[1]], [(0, 1)], mode=mode), branch_true_mean_diversity
        )
        self.assertAlmostEqual(
            ts.sample_count_stat(A, f, 1, mode=mode)[0], branch_true_mean_diversity
        )

        # Y-statistic for (0/12)
        A = [[0], [1, 2]]

        def f(x):
            return np.array(
                [
                    float(
                        ((x[0] == 1) and (x[1] == 0)) or ((x[0] == 0) and (x[1] == 2))
                    )
                    / 2.0
                ]
            )

        # tree lengths:
        bts_Y = ts.Y3([[0], [1], [2]], mode=mode)
        py_bsc_Y = Y3(ts, [[0], [1], [2]], [(0, 1, 2)], windows=[0.0, 1.0], mode=mode)
        self.assertArrayAlmostEqual(bts_Y, branch_true_Y)
        self.assertArrayAlmostEqual(py_bsc_Y, branch_true_Y)
        self.assertArrayAlmostEqual(
            ts.sample_count_stat(A, f, 1, mode=mode)[0], branch_true_Y
        )

        mode = "site"
        # sites, Y:
        sts_Y = ts.Y3([[0], [1], [2]], mode=mode)
        py_ssc_Y = Y3(ts, [[0], [1], [2]], [(0, 1, 2)], windows=[0.0, 1.0], mode=mode)
        self.assertArrayAlmostEqual(sts_Y, site_true_Y)
        self.assertArrayAlmostEqual(py_ssc_Y, site_true_Y)
        self.assertArrayAlmostEqual(
            ts.sample_count_stat(A, f, 1, mode=mode)[0], site_true_Y
        )

        A = [[0, 1, 2]]
        n = 3
        W = np.array([[u in A[0]] for u in ts.samples()], dtype=float)

        def f(x):
            return np.array([x[0] * (n - x[0]) / (n * (n - 1))])

        mode = "node"
        # nodes, diversity in [0,1,2]
        nodes_div_012 = ts.diversity([[0, 1, 2]], mode=mode).reshape((1, 7))
        py_nodes_div_012 = diversity(ts, [[0, 1, 2]], mode=mode).reshape((1, 7))
        py_general_nodes_div_012 = general_stat(ts, W, f, mode=mode).reshape((1, 7))
        self.assertArrayAlmostEqual(py_nodes_div_012, node_true_diversity_012)
        self.assertArrayAlmostEqual(py_general_nodes_div_012, node_true_diversity_012)
        self.assertArrayAlmostEqual(nodes_div_012, node_true_diversity_012)

        # nodes, divergence [0] to [1,2]
        nodes_div_0_12 = ts.divergence([[0], [1, 2]], mode=mode).reshape((1, 7))
        py_nodes_div_0_12 = divergence(ts, [[0], [1, 2]], mode=mode).reshape((1, 7))
        self.assertArrayAlmostEqual(nodes_div_0_12, node_true_divergence_0_12)
        self.assertArrayAlmostEqual(py_nodes_div_0_12, node_true_divergence_0_12)

        # covariance and correlation
        ts_sitewise_cov = ts.trait_covariance(
            traits, mode="site", windows="sites", span_normalise=False
        )
        py_sitewise_cov = site_trait_covariance(
            ts, traits, windows="sites", span_normalise=False
        )
        self.assertArrayAlmostEqual(py_sitewise_cov, true_cov)
        self.assertArrayAlmostEqual(ts_sitewise_cov, true_cov)
        ts_sitewise_cor = ts.trait_correlation(
            traits, mode="site", windows="sites", span_normalise=False
        )
        py_sitewise_cor = site_trait_correlation(
            ts, traits, windows="sites", span_normalise=False
        )
        self.assertArrayAlmostEqual(py_sitewise_cor, true_cor)
        self.assertArrayAlmostEqual(ts_sitewise_cor, true_cor)
        # mean
        ts_mean_cov = ts.trait_covariance(
            traits, mode="site", windows=[0, ts.sequence_length]
        )
        py_mean_cov = site_trait_covariance(ts, traits)
        self.assertArrayAlmostEqual(ts_mean_cov, np.array([np.sum(true_cov, axis=0)]))
        self.assertArrayAlmostEqual(ts_mean_cov, py_mean_cov)
        ts_mean_cor = ts.trait_correlation(
            traits, mode="site", windows=[0, ts.sequence_length]
        )
        py_mean_cor = site_trait_correlation(ts, traits)
        self.assertArrayAlmostEqual(ts_mean_cor, np.array([np.sum(true_cor, axis=0)]))
        self.assertArrayAlmostEqual(ts_mean_cor, py_mean_cor)
        # mode = 'branch'
        ts_mean_cov = ts.trait_covariance(
            traits, mode="branch", windows=[0, ts.sequence_length]
        )
        py_mean_cov = branch_trait_covariance(ts, traits)
        self.assertArrayAlmostEqual(ts_mean_cov, true_branch_cov)
        self.assertArrayAlmostEqual(ts_mean_cov, py_mean_cov)
        ts_mean_cor = ts.trait_correlation(
            traits, mode="branch", windows=[0, ts.sequence_length]
        )
        py_mean_cor = branch_trait_correlation(ts, traits)
        self.assertArrayAlmostEqual(ts_mean_cor, true_branch_cor)
        self.assertArrayAlmostEqual(ts_mean_cor, py_mean_cor)

        # trait linear_model:
        # r = cor * sd(y) / sd(x) = cov / var(x)
        # geno_var = allele_freqs * (1 - allele_freqs) * (3 / (3 - 1))
        geno_var = np.var(haplotypes, axis=1) * (3 / (3 - 1))
        trait_var = np.var(traits, axis=0) * (3 / (3 - 1))
        py_r = trait_linear_model(
            ts, traits, None, mode="site", windows="sites", span_normalise=False
        )
        ts_r = ts.trait_linear_model(
            traits, None, mode="site", windows="sites", span_normalise=False
        )
        self.assertArrayAlmostEqual(py_r, ts_r)
        self.assertArrayAlmostEqual(true_cov, py_r * (geno_var[:, np.newaxis] ** 2))
        self.assertArrayAlmostEqual(
            true_cor, ts_r * geno_var[:, np.newaxis] / trait_var
        )

    def test_case_odds_and_ends(self):
        # Tests having (a) the first site after the first window, and
        # (b) no samples having the ancestral state.
        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       0           0.5
        3       0           1.0
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.0     0.5     2       0,1
        0.5     1.0     3       0,1
        """
        )
        sites = io.StringIO(
            """\
        id  position    ancestral_state
        0   0.65        0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    derived_state   parent
        0       0       1               -1
        0       1       2               -1
        """
        )
        ts = tskit.load_text(
            nodes=nodes, edges=edges, sites=sites, mutations=mutations, strict=False
        )

        mode = "site"
        py_div = divergence(
            ts, [[0], [1]], indexes=[(0, 1)], windows=[0.0, 0.5, 1.0], mode=mode
        )
        div = ts.divergence(
            [[0], [1]], indexes=[(0, 1)], windows=[0.0, 0.5, 1.0], mode=mode
        )
        self.assertArrayEqual(py_div, div)

    def test_case_four_taxa(self):
        #
        # 1.0          7
        # 0.7         / \                                    6
        #            /   \                                  / \
        # 0.5       /     5              5                 /   5
        #          /     / \            / \__             /   / \
        # 0.4     /     8   \          8     4           /   8   \
        #        /     / \   \        / \   / \         /   / \   \
        # 0.0   0     1   3   2      1   3 0   2       0   1   3   2
        #          (0.0, 0.2),        (0.2, 0.8),       (0.8, 2.5)

        # f4(0, 1, 2, 3): (0 -> 1)(2 -> 3)
        branch_true_f4_0123 = (0.1 * 0.2 + (0.1 + 0.1) * 0.6 + 0.1 * 1.7) / 2.5
        windows = [0.0, 0.4, 2.5]
        branch_true_f4_0123_windowed = np.array(
            [
                (0.1 * 0.2 + (0.1 + 0.1) * 0.2) / 0.4,
                ((0.1 + 0.1) * 0.4 + 0.1 * 1.7) / 2.1,
            ]
        )
        # f4(0, 3, 2, 1): (0 -> 3)(2 -> 1)
        branch_true_f4_0321 = (0.1 * 0.2 + (0.1 + 0.1) * 0.6 + 0.1 * 1.7) / 2.5
        # f2([0,2], [1,3]) = (1/2) (f4(0,1,2,3) + f4(0,3,2,1))
        branch_true_f2_02_13 = (branch_true_f4_0123 + branch_true_f4_0321) / 2
        # diversity([0,1,2,3])
        branch_true_diversity_windowed = (2 / 6) * np.array(
            [
                [
                    (
                        0.2 * (1 + 1 + 1 + 0.5 + 0.4 + 0.5)
                        + (0.4 - 0.2) * (0.5 + 0.4 + 0.5 + 0.5 + 0.4 + 0.5)
                    )
                    / 0.4
                ],
                [
                    (
                        (0.8 - 0.4) * (0.5 + 0.4 + 0.5 + 0.5 + 0.4 + 0.5)
                        + (2.5 - 0.8) * (0.7 + 0.7 + 0.7 + 0.5 + 0.4 + 0.5)
                    )
                    / (2.5 - 0.4)
                ],
            ]
        )

        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       1           0
        4       0           0.4
        5       0           0.5
        6       0           0.7
        7       0           1.0
        8       0           0.4
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.0     2.5     8       1,3
        0.2     0.8     4       0,2
        0.0     0.2     5       8,2
        0.2     0.8     5       8,4
        0.8     2.5     5       8,2
        0.8     2.5     6       0,5
        0.0     0.2     7       0,5
        """
        )
        sites = io.StringIO(
            """\
        id  position    ancestral_state
        """
        )
        mutations = io.StringIO(
            """\
        site    node    derived_state   parent
        """
        )
        ts = tskit.load_text(
            nodes=nodes, edges=edges, sites=sites, mutations=mutations, strict=False
        )

        mode = "branch"
        A = [[0], [1], [2], [3]]
        self.assertAlmostEqual(branch_true_f4_0123, f4(ts, A, mode=mode)[0][0])
        self.assertAlmostEqual(branch_true_f4_0123, ts.f4(A, mode=mode))
        self.assertArrayAlmostEqual(
            branch_true_f4_0123_windowed, ts.f4(A, windows=windows, mode=mode).flatten()
        )
        A = [[0], [3], [2], [1]]
        self.assertAlmostEqual(
            branch_true_f4_0321, f4(ts, A, [(0, 1, 2, 3)], mode=mode)[0][0]
        )
        self.assertAlmostEqual(branch_true_f4_0321, ts.f4(A, mode=mode))
        A = [[0], [2], [1], [3]]
        self.assertAlmostEqual(0.0, f4(ts, A, [(0, 1, 2, 3)], mode=mode)[0])
        self.assertAlmostEqual(0.0, ts.f4(A, mode=mode))
        A = [[0, 2], [1, 3]]
        self.assertAlmostEqual(
            branch_true_f2_02_13, f2(ts, A, [(0, 1)], mode=mode)[0][0]
        )
        self.assertAlmostEqual(branch_true_f2_02_13, ts.f2(A, mode=mode))

        # diversity
        A = [[0, 1, 2, 3]]
        self.assertArrayAlmostEqual(
            branch_true_diversity_windowed, diversity(ts, A, windows=windows, mode=mode)
        )
        self.assertArrayAlmostEqual(
            branch_true_diversity_windowed, ts.diversity(A, windows=windows, mode=mode)
        )

    def test_case_recurrent_muts(self):
        # With mutations:
        #
        # 1.0          6
        # 0.7         / \                                    5
        #           (0)  \                                  /(6)
        # 0.5      (1)    4                4               /   4
        #          /     / \              / \             /  (7|8)
        # 0.4    (2)   (3)  \           (4)  3           /   /   \
        #        /     /     \          /   /(5)        /   /     \
        # 0.0   0     1       2        1   0   2       0   1       2
        #          (0.0, 0.2),        (0.2, 0.8),       (0.8, 1.0)
        # genotypes:
        #       0     2       0        1   0   1       0   2       3
        site_true_Y = 0 + 1 + 1

        nodes = io.StringIO(
            """\
        id      is_sample   time
        0       1           0
        1       1           0
        2       1           0
        3       0           0.4
        4       0           0.5
        5       0           0.7
        6       0           1.0
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.2     0.8     3       0,2
        0.0     0.2     4       1,2
        0.2     0.8     4       1,3
        0.8     1.0     4       1,2
        0.8     1.0     5       0,4
        0.0     0.2     6       0,4
        """
        )
        sites = io.StringIO(
            """\
        id  position    ancestral_state
        0   0.05        0
        1   0.3         0
        2   0.9         0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    derived_state   parent
        0       0       1               -1
        0       0       2               0
        0       0       0               1
        0       1       2               -1
        1       1       1               -1
        1       2       1               -1
        2       4       1               -1
        2       1       2               6
        2       2       3               6
        """
        )
        ts = tskit.load_text(
            nodes=nodes, edges=edges, sites=sites, mutations=mutations, strict=False
        )

        # Y3:
        site_tsc_Y = ts.Y3([[0], [1], [2]], mode="site")
        py_ssc_Y = Y3(ts, [[0], [1], [2]], [(0, 1, 2)], windows=[0.0, 1.0], mode="site")
        self.assertAlmostEqual(site_tsc_Y, site_true_Y)
        self.assertAlmostEqual(py_ssc_Y, site_true_Y)

    def test_case_2(self):
        # Here are the trees:
        # t                  |              |              |             |            |
        #
        # 0       --3--      |     --3--    |     --3--    |    --3--    |    --3--   |
        #        /  |  \     |    /  |  \   |    /     \   |   /     \   |   /     \  |
        # 1     4   |   5    |   4   |   5  |   4       5  |  4       5  |  4       5 |
        #       |\ / \ /|    |   |\   \     |   |\     /   |  |\     /   |  |\     /| |
        # 2     | 6   7 |    |   | 6   7    |   | 6   7    |  | 6   7    |  | 6   7 | |
        #       | |\ /| |    |   *  \  |    |   |  \  |    |  |  \       |  |  \    | |
        # 3     | | 8 | |    |   |   8 *    |   |   8 |    |  |   8      |  |   8   | |
        #       | |/ \| |    |   |  /  |    |   |  /  |    |  |  / \     |  |  / \  | |
        # 4     | 9  10 |    |   * 9  10    |   | 9  10    |  | 9  10    |  | 9  10 | |
        #       |/ \ / \|    |   |  \   \   |   |  \   \   |  |  \   \   |  |  \    | |
        # 5     0   1   2    |   0   1   2  |   0   1   2  |  0   1   2  |  0   1   2 |
        #
        #                    |   0.0 - 0.1  |   0.1 - 0.2  |  0.2 - 0.4  |  0.4 - 0.5 |
        # ... continued:
        # t                  |             |             |             |
        #
        # 0         --3--    |    --3--    |    --3--    |    --3--    |    --3--
        #          /     \   |   /     \   |   /     \   |   /     \   |   /  |  \
        # 1       4       5  |  4       5  |  4       5  |  4       5  |  4   |   5
        #         |\     /|  |   \     /|  |   \     /|  |   \     /|  |     /   /|
        # 2       | 6   7 |  |    6   7 |  |    6   7 |  |    6   7 |  |    6   7 |
        #         |  *    *  |     \    |  |       *  |  |    |  /  |  |    |  /  |
        # 3  ...  |   8   |  |      8   |  |      8   |  |    | 8   |  |    | 8   |
        #         |  / \  |  |     / \  |  |     * \  |  |    |  \  |  |    |  \  |
        # 4       | 9  10 |  |    9  10 |  |    9  10 |  |    9  10 |  |    9  10 |
        #         |    /  |  |   /   /  |  |   /   /  |  |   /   /  |  |   /   /  |
        # 5       0   1   2  |  0   1   2  |  0   1   2  |  0   1   2  |  0   1   2
        #
        #         0.5 - 0.6  |  0.6 - 0.7  |  0.7 - 0.8  |  0.8 - 0.9  |  0.9 - 1.0
        #
        # Above, subsequent mutations are backmutations.

        # divergence betw 0 and 1
        branch_true_diversity_01 = 2 * (0.6 * 4 + 0.2 * 2 + 0.2 * 5)
        # divergence betw 1 and 2
        branch_true_diversity_12 = 2 * (0.2 * 5 + 0.2 * 2 + 0.3 * 5 + 0.3 * 4)
        # divergence betw 0 and 2
        branch_true_diversity_02 = 2 * (0.2 * 5 + 0.2 * 4 + 0.3 * 5 + 0.1 * 4 + 0.2 * 5)
        # Y(0;1, 2)
        branch_true_Y = 0.2 * 4 + 0.2 * (4 + 2) + 0.2 * 4 + 0.2 * 2 + 0.2 * (5 + 1)

        # site stats
        # Y(0;1, 2)
        site_true_Y = 1

        nodes = io.StringIO(
            """\
        is_sample       time    population
        1       0.000000        0
        1       0.000000        0
        1       0.000000        0
        0       5.000000        0
        0       4.000000        0
        0       4.000000        0
        0       3.000000        0
        0       3.000000        0
        0       2.000000        0
        0       1.000000        0
        0       1.000000        0
        """
        )
        edges = io.StringIO(
            """\
        left    right   parent  child
        0.500000        1.000000        10      1
        0.000000        0.400000        10      2
        0.600000        1.000000        9       0
        0.000000        0.500000        9       1
        0.800000        1.000000        8       10
        0.200000        0.800000        8       9,10
        0.000000        0.200000        8       9
        0.700000        1.000000        7       8
        0.000000        0.200000        7       10
        0.800000        1.000000        6       9
        0.000000        0.700000        6       8
        0.400000        1.000000        5       2,7
        0.100000        0.400000        5       7
        0.600000        0.900000        4       6
        0.000000        0.600000        4       0,6
        0.900000        1.000000        3       4,5,6
        0.100000        0.900000        3       4,5
        0.000000        0.100000        3       4,5,7
        """
        )
        sites = io.StringIO(
            """\
        id  position    ancestral_state
        0   0.0         0
        1   0.55        0
        2   0.75        0
        3   0.85        0
        """
        )
        mutations = io.StringIO(
            """\
        site    node    derived_state   parent
        0       0       1               -1
        0       10      1               -1
        0       0       0               0
        1       8       1               -1
        1       2       1               -1
        2       8       1               -1
        2       9       0               5
        """
        )
        ts = tskit.load_text(
            nodes=nodes, edges=edges, sites=sites, mutations=mutations, strict=False
        )

        def f(x):
            return np.array([float(x[0] == 1) / 2.0])

        # divergence between 0 and 1
        mode = "branch"
        for A, truth in zip(
            [[[0, 1]], [[1, 2]], [[0, 2]]],
            [
                branch_true_diversity_01,
                branch_true_diversity_12,
                branch_true_diversity_02,
            ],
        ):
            self.assertAlmostEqual(diversity(ts, A, mode=mode)[0][0], truth)
            self.assertAlmostEqual(ts.sample_count_stat(A, f, 1, mode=mode)[0], truth)
            self.assertAlmostEqual(ts.diversity(A, mode="branch")[0], truth)

        # Y-statistic for (0/12)
        A = [[0], [1, 2]]

        def f(x):
            return np.array(
                [
                    float(
                        ((x[0] == 1) and (x[1] == 0)) or ((x[0] == 0) and (x[1] == 2))
                    )
                    / 2.0
                ]
            )

        # tree lengths:
        self.assertArrayAlmostEqual(
            Y3(ts, [[0], [1], [2]], [(0, 1, 2)], mode=mode), branch_true_Y
        )
        self.assertArrayAlmostEqual(
            ts.Y3([[0], [1], [2]], [(0, 1, 2)], mode=mode), branch_true_Y
        )
        self.assertArrayAlmostEqual(
            ts.sample_count_stat(A, f, 1, mode=mode)[0], branch_true_Y
        )

        # sites:
        mode = "site"
        site_tsc_Y = ts.Y3([[0], [1], [2]], mode=mode)
        py_ssc_Y = Y3(ts, [[0], [1], [2]], [(0, 1, 2)], windows=[0.0, 1.0])
        self.assertAlmostEqual(site_tsc_Y, site_true_Y)
        self.assertAlmostEqual(py_ssc_Y, site_true_Y)
        self.assertAlmostEqual(ts.sample_count_stat(A, f, 1, mode=mode)[0], site_true_Y)


class TestOutputDimensions(StatsTestCase):
    """
    Tests for the dimension stripping behaviour of the stats functions.
    """

    def get_example_ts(self, ts_10_mut_fixture):
        ts = ts_10_mut_fixture
        assert ts.num_sites > 1
        return ts

    def test_one_way_no_window_scalar_stat(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        x = ts.diversity()
        assert isinstance(x, np.floating)

    def test_one_way_one_list_scalar_stat(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        x = ts.diversity(sample_sets=list(ts.samples()))
        assert isinstance(x, np.floating)

    def test_one_way_nested_list_not_scalar_stat(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        x = ts.diversity(sample_sets=[list(ts.samples())])
        assert x.shape == (1,)

    def test_one_way_one_window_scalar_stat(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        x = ts.diversity(windows=[0, ts.sequence_length])
        assert x.shape == (1,)
        for samples in (None, list(ts.samples())):
            x = ts.diversity(sample_sets=samples, windows=[0, ts.sequence_length])
            assert x.shape == (1,)

    def test_multi_way_no_window_scalar_stat(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        n = ts.num_samples
        x = ts.f2(
            sample_sets=[
                [i for i in range(0, int(n / 2))],
                [i for i in range(int(n / 2), n)],
            ]
        )
        assert isinstance(x, np.floating)

    def test_multi_way_one_window_not_scalar_stat(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        n = ts.num_samples
        x = ts.f2(
            sample_sets=[
                [i for i in range(0, int(n / 2))],
                [i for i in range(int(n / 2), n)],
            ],
            windows=[0, ts.sequence_length],
        )
        assert x.shape == (1,)

    def test_multi_way_no_indexes_scalar_stat(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        n = ts.num_samples
        x = ts.f2(
            sample_sets=[
                [i for i in range(0, int(n / 2))],
                [i for i in range(int(n / 2), n)],
            ],
        )
        assert isinstance(x, np.floating)

    def test_multi_way_indexes_not_scalar_stat(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        n = ts.num_samples
        x = ts.f2(
            sample_sets=[
                [i for i in range(0, int(n / 2))],
                [i for i in range(int(n / 2), n)],
            ],
            indexes=[(0, 1)],
        )
        assert x.shape == (1,)

    def test_afs_default_windows(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        n = ts.num_samples
        A = ts.samples()[:4]
        B = ts.samples()[6:]
        for mode in ["site", "branch"]:
            x = ts.allele_frequency_spectrum(mode=mode)
            # x is a 1D numpy array with n + 1 values
            assert x.shape == (n + 1,)
            self.assertArrayEqual(
                x, ts.allele_frequency_spectrum([ts.samples()], mode=mode)
            )
            x = ts.allele_frequency_spectrum([A, B], mode=mode)
            assert x.shape == (len(A) + 1, len(B) + 1)

    def test_afs_windows(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        L = ts.sequence_length

        windows = [0, L / 4, L / 2, L]
        A = ts.samples()[:4]
        B = ts.samples()[6:]
        for mode in ["site", "branch"]:
            x = ts.allele_frequency_spectrum([A, B], windows=windows, mode=mode)
            assert x.shape == (3, len(A) + 1, len(B) + 1)

            x = ts.allele_frequency_spectrum([A], windows=windows, mode=mode)
            assert x.shape == (3, len(A) + 1)

            x = ts.allele_frequency_spectrum(windows=windows, mode=mode)
            # Default returns this for all samples
            assert x.shape == (3, ts.num_samples + 1)
            y = ts.allele_frequency_spectrum([ts.samples()], windows=windows, mode=mode)
            self.assertArrayEqual(x, y)

    def test_one_way_stat_default_windows(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        # Use diversity as the example one-way stat.
        for mode in ["site", "branch"]:
            x = ts.diversity(mode=mode)
            # x is a zero-d numpy value
            assert np.shape(x) == tuple()
            assert x == float(x)
            assert x == ts.diversity(ts.samples(), mode=mode)
            self.assertArrayEqual([x], ts.diversity([ts.samples()], mode=mode))

        mode = "node"
        x = ts.diversity(mode=mode)
        # x is a 1D numpy array with N values
        assert x.shape == (ts.num_nodes,)
        self.assertArrayEqual(x, ts.diversity(ts.samples(), mode=mode))
        y = ts.diversity([ts.samples()], mode=mode)
        # We're adding on the *last* dimension, so must reshape
        self.assertArrayEqual(x.reshape(ts.num_nodes, 1), y)

    def verify_one_way_stat_windows(self, ts, method):
        L = ts.sequence_length
        N = ts.num_nodes

        windows = [0, L / 4, L / 2, 0.75 * L, L]
        A = ts.samples()[:6]
        B = ts.samples()[6:]
        for mode in ["site", "branch"]:
            x = method([A, B], windows=windows, mode=mode)
            # Four windows, 2 sets.
            assert x.shape == (4, 2)

            x = method([A], windows=windows, mode=mode)
            # Four windows, 1 sets.
            assert x.shape == (4, 1)

            x = method(A, windows=windows, mode=mode)
            # Dropping the outer list removes the last dimension
            assert x.shape == (4,)

            x = method(windows=windows, mode=mode)
            # Default returns this for all samples
            assert x.shape == (4,)
            y = method(ts.samples(), windows=windows, mode=mode)
            self.assertArrayEqual(x, y)

        mode = "node"
        x = method([A, B], windows=windows, mode=mode)
        # Four windows, N nodes and 2 sets.
        assert x.shape == (4, N, 2)

        x = method([A], windows=windows, mode=mode)
        # Four windows, N nodes and 1 set.
        assert x.shape == (4, N, 1)

        x = method(A, windows=windows, mode=mode)
        # Drop the outer list, so we lose the last dimension
        assert x.shape == (4, N)

        x = method(windows=windows, mode=mode)
        # The default sample sets also drops the last dimension
        assert x.shape == (4, N)

        assert ts.num_trees == 1
        # In this example, we know that the trees are all the same so check this
        # for sanity.
        self.assertArrayEqual(x[0], x[1])
        self.assertArrayEqual(x[0], x[2])

    def test_diversity_windows(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        self.verify_one_way_stat_windows(ts, ts.diversity)

    def test_Tajimas_D_windows(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        self.verify_one_way_stat_windows(ts, ts.Tajimas_D)

    def test_segregating_sites_windows(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        self.verify_one_way_stat_windows(ts, ts.segregating_sites)

    def test_two_way_stat_default_windows(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        # Use divergence as the example one-way stat.
        A = ts.samples()[:6]
        B = ts.samples()[6:]
        for mode in ["site", "branch"]:
            x = ts.divergence([A, B], mode=mode)
            # x is a zero-d numpy value
            assert np.shape(x) == tuple()
            assert x == float(x)
            # If indexes is a 1D array, we also drop the outer dimension
            assert x == ts.divergence([A, B, A], indexes=[0, 1], mode=mode)
            # But, if it's a 2D array we keep the outer dimension
            assert [x] == ts.divergence([A, B], indexes=[[0, 1]], mode=mode)

        mode = "node"
        x = ts.divergence([A, B], mode=mode)
        # x is a 1D numpy array with N values
        assert x.shape == (ts.num_nodes,)
        self.assertArrayEqual(x, ts.divergence([A, B], indexes=[0, 1], mode=mode))
        y = ts.divergence([A, B], indexes=[[0, 1]], mode=mode)
        # We're adding on the *last* dimension, so must reshape
        self.assertArrayEqual(x.reshape(ts.num_nodes, 1), y)

    def verify_two_way_stat_windows(self, ts, method):
        L = ts.sequence_length
        N = ts.num_nodes

        windows = [0, L / 4, L / 2, L]
        A = ts.samples()[:7]
        B = ts.samples()[7:]
        for mode in ["site", "branch"]:
            x = method([A, B, A], indexes=[[0, 1], [0, 2]], windows=windows, mode=mode)
            # Three windows, 2 pairs
            assert x.shape == (3, 2)

            x = method([A, B], indexes=[[0, 1]], windows=windows, mode=mode)
            # Three windows, 1 pair
            assert x.shape == (3, 1)

            x = method([A, B], indexes=[0, 1], windows=windows, mode=mode)
            # Dropping the outer list removes the last dimension
            assert x.shape == (3,)

            y = method([A, B], windows=windows, mode=mode)
            assert y.shape == (3,)
            self.assertArrayEqual(x, y)

        mode = "node"
        x = method([A, B], indexes=[[0, 1], [0, 1]], windows=windows, mode=mode)
        # Three windows, N nodes and 2 pairs
        assert x.shape == (3, N, 2)

        x = method([A, B], indexes=[[0, 1]], windows=windows, mode=mode)
        # Three windows, N nodes and 1 pairs
        assert x.shape == (3, N, 1)

        x = method([A, B], indexes=[0, 1], windows=windows, mode=mode)
        # Drop the outer list, so we lose the last dimension
        assert x.shape == (3, N)

        x = method([A, B], windows=windows, mode=mode)
        # The default sample sets also drops the last dimension
        assert x.shape == (3, N)

        assert ts.num_trees == 1
        # In this example, we know that the trees are all the same so check this
        # for sanity.
        self.assertArrayEqual(x[0], x[1])
        self.assertArrayEqual(x[0], x[2])

    def test_divergence_windows(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        self.verify_two_way_stat_windows(ts, ts.divergence)

    def test_Fst_windows(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        self.verify_two_way_stat_windows(ts, ts.Fst)

    def test_f2_windows(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        self.verify_two_way_stat_windows(ts, ts.f2)

    def verify_three_way_stat_windows(self, ts, method):
        L = ts.sequence_length
        N = ts.num_nodes

        windows = [0, L / 4, L / 2, L]
        A = ts.samples()[:2]
        B = ts.samples()[2:4]
        C = ts.samples()[4:]
        for mode in ["site", "branch"]:
            x = method(
                [A, B, C], indexes=[[0, 1, 2], [0, 2, 1]], windows=windows, mode=mode
            )
            # Three windows, 2 triple
            assert x.shape == (3, 2)

            x = method([A, B, C], indexes=[[0, 1, 2]], windows=windows, mode=mode)
            # Three windows, 1 triple
            assert x.shape == (3, 1)

            x = method([A, B, C], indexes=[0, 1, 2], windows=windows, mode=mode)
            # Dropping the outer list removes the last dimension
            assert x.shape == (3,)

            y = method([A, B, C], windows=windows, mode=mode)
            assert y.shape == (3,)
            self.assertArrayEqual(x, y)

        mode = "node"
        x = method(
            [A, B, C], indexes=[[0, 1, 2], [0, 2, 1]], windows=windows, mode=mode
        )
        # Three windows, N nodes and 2 triples
        assert x.shape == (3, N, 2)

        x = method([A, B, C], indexes=[[0, 1, 2]], windows=windows, mode=mode)
        # Three windows, N nodes and 1 triples
        assert x.shape == (3, N, 1)

        x = method([A, B, C], indexes=[0, 1, 2], windows=windows, mode=mode)
        # Drop the outer list, so we lose the last dimension
        assert x.shape == (3, N)

        x = method([A, B, C], windows=windows, mode=mode)
        # The default sample sets also drops the last dimension
        assert x.shape == (3, N)

        assert ts.num_trees == 1
        # In this example, we know that the trees are all the same so check this
        # for sanity.
        self.assertArrayEqual(x[0], x[1])
        self.assertArrayEqual(x[0], x[2])

    def test_Y3_windows(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        self.verify_three_way_stat_windows(ts, ts.Y3)

    def test_f3_windows(self, ts_10_mut_fixture):
        ts = self.get_example_ts(ts_10_mut_fixture)
        self.verify_three_way_stat_windows(ts, ts.f3)


class TestTimeUncalibratedErrors:
    def test_uncalibrated_time_allele_frequency_spectrum(self, ts_fixture):
        ts_fixture.allele_frequency_spectrum(mode="branch")
        tables = ts_fixture.dump_tables()
        tables.time_units = tskit.TIME_UNITS_UNCALIBRATED
        ts_uncalibrated = tables.tree_sequence()
        ts_uncalibrated.allele_frequency_spectrum(mode="site")
        with pytest.raises(
            tskit.LibraryError,
            match="Statistics using branch lengths cannot be calculated when time_units"
            " is 'uncalibrated'",
        ):
            ts_uncalibrated.allele_frequency_spectrum(mode="branch")

    def test_uncalibrated_time_general_stat(self, ts_fixture):
        W = np.ones((ts_fixture.num_samples, 2))
        ts_fixture.general_stat(
            W, lambda x: x * (x < ts_fixture.num_samples), W.shape[1], mode="branch"
        )
        tables = ts_fixture.dump_tables()
        tables.time_units = tskit.TIME_UNITS_UNCALIBRATED
        ts_uncalibrated = tables.tree_sequence()
        ts_uncalibrated.general_stat(
            W, lambda x: x * (x < ts_uncalibrated.num_samples), W.shape[1], mode="site"
        )
        with pytest.raises(
            tskit.LibraryError,
            match="Statistics using branch lengths cannot be calculated when time_units"
            " is 'uncalibrated'",
        ):
            ts_uncalibrated.general_stat(
                W,
                lambda x: x * (x < ts_uncalibrated.num_samples),
                W.shape[1],
                mode="branch",
            )


class TestGeneralStatCallbackErrors:
    def test_zero_d(self, ts_fixture):
        def f_0d(_):
            return 0

        msg = "Array returned by general_stat callback is 0 dimensional; must be 1D"
        with pytest.raises(ValueError, match=msg):
            ts_fixture.sample_count_stat(
                sample_sets=[ts_fixture.samples()], f=f_0d, output_dim=1, strict=False
            )

    def test_two_d(self, ts_fixture):
        def f_2d(x):
            return np.array([x])

        msg = "Array returned by general_stat callback is 2 dimensional; must be 1D"
        with pytest.raises(ValueError, match=msg):
            ts_fixture.sample_count_stat(
                sample_sets=[ts_fixture.samples()], f=f_2d, output_dim=1, strict=False
            )

    def test_wrong_length(self, ts_fixture):
        def f_too_long(_):
            return np.array([0, 0])

        msg = "Array returned by general_stat callback is of length 2; must be 1"
        with pytest.raises(ValueError, match=msg):
            ts_fixture.sample_count_stat(
                sample_sets=[ts_fixture.samples()],
                f=f_too_long,
                output_dim=1,
                strict=False,
            )


class TestTimeWindows:

    def test_general_stat(self, four_taxa_test_case):
        # 1.00┊   7     ┊         ┊         ┊
        #     ┊ ┏━┻━┓   ┊         ┊         ┊
        # 0.70┊ ┃   ┃   ┊         ┊   6     ┊
        #     ┊ ┃   ┃   ┊         ┊ ┏━┻━┓   ┊
        # 0.50┊ ┃   5   ┊    5    ┊ ┃   5   ┊
        #     ┊ ┃  ┏┻━┓ ┊  ┏━┻━┓  ┊ ┃  ┏┻━┓ ┊
        # 0.40┊ ┃  8  ┃ ┊  4   8  ┊ ┃  8  ┃ ┊
        #     ┊ ┃ ┏┻┓ ┃ ┊ ┏┻┓ ┏┻┓ ┊ ┃ ┏┻┓ ┃ ┊
        # 0.00┊ 0 1 3 2 ┊ 0 2 1 3 ┊ 0 1 3 2 ┊
        #   0.00      0.20      0.80      2.50
        ts = four_taxa_test_case
        true_x = np.array(
            [
                [
                    [
                        0.2 * (1 + 0.5 + 0.4)
                        + (0.8 - 0.2) * (1 + 0.8)
                        + (2.5 - 0.8) * (1.0 + 0.5 + 0.4)
                    ],
                    [0.2 * 1.0 + 0 + (2.5 - 0.8) * 0.4],
                ]
            ]
        )

        n = ts.num_samples

        def f(x):
            return (x > 0) * (1 - x / n)

        W = np.ones((ts.num_samples, 1))
        x = naive_branch_general_stat(
            ts, W, f, time_windows=[0, 0.5, 2.0], span_normalise=False
        )
        np.testing.assert_allclose(x, true_x)

        x0 = branch_general_stat(ts, W, f, time_windows=None, span_normalise=False)
        x1 = naive_branch_general_stat(
            ts, W, f, time_windows=None, span_normalise=False
        )
        np.testing.assert_allclose(x0, x1)
        x_tw = branch_general_stat(
            ts, W, f, time_windows=[0, 0.5, 2.0], span_normalise=False
        )

        np.testing.assert_allclose(x, x_tw)