test + error

requirements/CI-complete/requirements.txt

@@ -4,4 +4,5 @@ numpy
 pandas
 pytest==6.2.5
 pytest-cov==3.0.0
+scipy
 tskit==0.5.3

requirements/CI-tests-pip/requirements.txt

@@ -2,4 +2,5 @@ msprime
 numba
 numpy
 pandas
+scipy
 tskit==0.5.3

setup.cfg

@@ -58,6 +58,7 @@ setup_requires =
 tests = 
     msprime
     pytest
+    scipy
 
 [tool:pytest]
 testpaths = 

tests/data.py

@@ -0,0 +1,171 @@
+import tskit
+
+
+def binary_tree():
+    """Sample tree sequence. The genotype is the following:
+    Genotype: [Node0, Node1, Node2, Node3, Node4, Node5]
+    Node4 and Node5 are internal nodes
+    Ancestral State: A
+    Site0: [A, A, T, T, A, T], has 1 mutation
+    Site1: [C, T, A, A, T, A], has 2 mutations of the same type
+    Site2: [C, A, A, A, A, A], has multiple mutations in the same node
+    Site3: [C, C, T, T, C, T], has reverse mutation and mutation in head node
+
+    Individual1: Node0 and Node1
+    Individual2: Node2 and Node3
+    Individual0: Node4 and Node5 (Internal node)
+    """
+    ts = tskit.Tree.generate_balanced(4, span=10).tree_sequence
+    tables = ts.dump_tables()
+    for j in range(4):
+        tables.sites.add_row(j, "A")
+
+    tables.individuals.add_row()
+    tables.individuals.add_row()
+    tables.individuals.add_row()
+    individuals = tables.nodes.individual
+    individuals[[0, 1]] = 1
+    individuals[[2, 3]] = 2
+    individuals[[4, 5]] = 0
+    tables.nodes.individual = individuals
+
+    tables.mutations.add_row(site=0, node=5, derived_state="T")
+
+    tables.mutations.add_row(site=1, node=4, derived_state="T")
+    tables.mutations.add_row(site=1, node=0, derived_state="C", parent=1)
+
+    tables.mutations.add_row(site=2, node=0, derived_state="T")
+    tables.mutations.add_row(site=2, node=0, derived_state="G", parent=3)
+    tables.mutations.add_row(site=2, node=0, derived_state="T", parent=4)
+    tables.mutations.add_row(site=2, node=0, derived_state="C", parent=5)
+
+    tables.mutations.add_row(site=3, node=6, derived_state="T")
+    tables.mutations.add_row(site=3, node=4, derived_state="C", parent=7)
+
+    ts = tables.tree_sequence()
+
+    return ts
+
+
+def diff_ind_tree():
+    """Same mutation and tree structure as binary_tree, but individuals will
+    have different nodes. See the details of the tree sequence in the docstring
+    for `binary_tree`.
+
+    Individual1: Node0 and Node2
+    Individual2: Node1 and Node3
+    Individual0: Node4 and Node5 (Internal Node)
+    """
+    ts = binary_tree()
+    tables = ts.dump_tables()
+    individuals = tables.nodes.individual
+    individuals[[0, 2]] = 1
+    individuals[[1, 3]] = 2
+    individuals[[4, 5]] = 0
+    tables.nodes.individual = individuals
+    ts = tables.tree_sequence()
+    return ts
+
+
+def non_binary_tree():
+    """Non-binary tree sequence.
+    Genotype: [Node0, Node1, Node2, Node3, Node4, Node5]
+    Ancestral State: A
+    Site0: [A, A, A, T, T, T], has 1 mutation
+    Site1: [A, A, A, C, C, T], has reverse mutation and multiple mutations
+        in the same node
+
+    Individual0: Node0 and Node1
+    Individual1: Node2 and Node3
+    Individual2: Node4 and Node5
+    """
+    ts = tskit.Tree.generate_balanced(6, arity=4, span=10).tree_sequence
+
+    tables = ts.dump_tables()
+    for j in range(2):
+        tables.sites.add_row(j, "A")
+
+    tables.individuals.add_row()
+    tables.individuals.add_row()
+    tables.individuals.add_row()
+    individuals = tables.nodes.individual
+    individuals[0] = 0
+    individuals[1] = 0
+    individuals[2] = 1
+    individuals[3] = 1
+    individuals[4] = 2
+    individuals[5] = 2
+    tables.nodes.individual = individuals
+
+    tables.mutations.add_row(site=0, node=6, derived_state="T")
+    tables.mutations.add_row(site=1, node=6, derived_state="T")
+    tables.mutations.add_row(site=1, node=6, derived_state="C", parent=1)
+    tables.mutations.add_row(site=1, node=5, derived_state="T", parent=2)
+
+    ts = tables.tree_sequence()
+    return ts
+
+
+def all_trees_ts(n):
+    """
+    Generate a tree sequence that corresponds to the lexicographic listing
+    of all trees with n leaves (i.e. from tskit.all_trees(n)).
+
+    Note: it would be nice to include a version of this in the combinatorics
+    module at some point but the implementation is quite inefficient. Also
+    it's not entirely clear that the way we're allocating node times is
+    guaranteed to work.
+    """
+    tables = tskit.TableCollection(0)
+    for _ in range(n):
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
+    for j in range(1, n):
+        tables.nodes.add_row(flags=0, time=j)
+
+    L = 0
+    for tree in tskit.all_trees(n):
+        for u in tree.preorder()[1:]:
+            tables.edges.add_row(L, L + 1, tree.parent(u), u)
+        L += 1
+    tables.sequence_length = L
+    tables.sort()
+    tables.simplify()
+    return tables.tree_sequence()
+
+
+def binary_tree_seq():
+    """Sample tree sequence
+    Ancestral State: A
+    Genotype: [Node0, Node1, Node2, Node3]
+    Site0: [T, A, T, T], has 1 mutation
+    Site1: [A, T, G, G], has reverse mutation
+    Site2: [A, C, C, A], has multiple mutation on the same edge
+
+    Individual0: Node0 and Node1
+    Individual1: Node2 and Node3
+    """
+    ts = all_trees_ts(4)
+    tables = ts.dump_tables()
+
+    tables.individuals.add_row()
+    tables.individuals.add_row()
+    individuals = tables.nodes.individual
+    individuals[0] = 0
+    individuals[1] = 0
+    individuals[2] = 1
+    individuals[3] = 1
+    tables.nodes.individual = individuals
+
+    tables.sites.add_row(7, "A")
+    tables.sites.add_row(11, "A")
+    tables.sites.add_row(12, "A")
+
+    tables.mutations.add_row(site=0, node=4, derived_state="T")
+    tables.mutations.add_row(site=1, node=5, derived_state="T")
+    tables.mutations.add_row(site=1, node=4, derived_state="G", parent=1)
+    tables.mutations.add_row(site=2, node=4, derived_state="T")
+    tables.mutations.add_row(site=2, node=4, derived_state="C", parent=3)
+
+    ts = tables.tree_sequence()
+
+    return ts

tests/test_base.py

@@ -0,0 +1,74 @@
+import fractions
+
+import numpy as np
+import pytest
+
+from tstrait.base import (
+    _check_int,
+    _check_val,
+    _check_instance,
+    _check_numeric_array,
+)  # noreorder
+
+
+class TestInstance:
+    def test_instance(self):
+        n = _check_instance(1, "n", int)
+        assert n == 1
+
+    @pytest.mark.parametrize("n", [4.0, np.array([4]), fractions.Fraction(4, 1)])
+    def test_instance_bad(self, n):
+        with pytest.raises(TypeError, match=f"n must be a {int} instance"):
+            _check_instance(n, "n", int)
+
+
+class TestVal:
+    @pytest.mark.parametrize(
+        "n", [1.1, np.float16(1.1), np.int16(1), np.array([1.1])[0]]
+    )
+    def test_val(self, n):
+        n = _check_val(n, "n")
+        assert type(n) == float
+
+    @pytest.mark.parametrize("n", [4j, np.array([4]), "1"])
+    def test_val_bad(self, n):
+        with pytest.raises(TypeError, match="n must be numeric"):
+            _check_val(n, "n")
+
+    @pytest.mark.parametrize("n", [0, -1])
+    def test_val_below_min(self, n):
+        with pytest.raises(ValueError, match="n must be a number greater than 0"):
+            _check_val(n, "n", 0)
+
+    def test_val_inclusive(self):
+        _check_val(0, "n", 0, inclusive=True)
+
+    def test_val_below_min_inclusive(self):
+        with pytest.raises(ValueError, match="n must be a number not " "less than 0"):
+            _check_val(-1, "n", 0, inclusive=True)
+
+
+class TestInt:
+    @pytest.mark.parametrize("n", [1, np.uint8(1), np.int16(1), np.array(1)])
+    def test_int(self, n):
+        n = _check_int(n, "n")
+        assert n == 1
+
+    @pytest.mark.parametrize("n", [4.0, np.array([4]), fractions.Fraction(4, 1)])
+    def test_int_bad(self, n):
+        with pytest.raises(TypeError, match="n must be an integer"):
+            _check_int(n, "n")
+
+    def test_int_below_min(self):
+        with pytest.raises(ValueError, match="n must be an integer not " "less than 0"):
+            _check_int(-1, "n", 0)
+
+
+class TestNumericArray:
+    def test_true(self):
+        _check_numeric_array([1, 2.1], "input")
+        _check_numeric_array(np.array([1, 2.1]), "input")
+
+    def test_nonnumeric(self):
+        with pytest.raises(TypeError, match="input must be numeric"):
+            _check_numeric_array([1, "1"], "input")

tests/test_init.py

@@ -0,0 +1,13 @@
+try:
+    from tstrait import *  # noqa
+
+    _top_import_error = None
+except Exception as e:
+    _top_import_error = e
+
+
+def test_import_tstrait():
+    # Test either above import has failed for some reason
+    # "import *" is discouraged outside of the module level, hence we
+    # rely on setting up the variable above
+    assert _top_import_error is None