genetic value

Compute genetic values from effect size dataframe

README.md

@@ -2,7 +2,7 @@
 
 **tstrait** is a quantitative trait simulator of [tskit](https://tskit.readthedocs.io/) tree sequences.
 
-- Documentation:
+- Documentation: https://tskit.dev/tstrait/docs/
 - Source code: https://github.com/tskit-dev/tstrait
 - Tree sequence documentation: https://tskit.dev/learn/
 

tests/data.py

@@ -58,6 +58,10 @@ def diff_ind_tree():
     """
     ts = binary_tree()
     tables = ts.dump_tables()
+    tables.individuals.clear()
+    tables.individuals.add_row()
+    tables.individuals.add_row()
+    tables.individuals.add_row()
     individuals = tables.nodes.individual
     individuals[[0, 2]] = 1
     individuals[[1, 3]] = 2
@@ -106,6 +110,27 @@ def non_binary_tree():
     return ts
 
 
+def triploid_tree():
+    """Same mutation and tree structure as non_binary_tree, but individuals have
+    different nodes and are triploids. See the details of the tree sequence
+    in the doctring for `non_binary_tree`.
+
+    Individual0: Node0, Node2 and Node4
+    Individual1: Node1, Node3 and Node5
+    """
+    ts = non_binary_tree()
+    tables = ts.dump_tables()
+    tables.individuals.clear()
+    tables.individuals.add_row()
+    tables.individuals.add_row()
+    individuals = tables.nodes.individual
+    individuals[[0, 2, 4]] = 0
+    individuals[[1, 3, 5]] = 1
+    tables.nodes.individual = individuals
+    ts = tables.tree_sequence()
+    return ts
+
+
 def all_trees_ts(n):
     """
     Generate a tree sequence that corresponds to the lexicographic listing

tests/test_base.py

@@ -1,13 +1,16 @@
 import fractions
 
 import numpy as np
+import pandas as pd
 import pytest
 
 from tstrait.base import (
     _check_int,
     _check_val,
     _check_instance,
     _check_numeric_array,
+    _check_dataframe,
+    _check_non_decreasing,
 )  # noreorder
 
 
@@ -72,3 +75,36 @@ def test_true(self):
     def test_nonnumeric(self):
         with pytest.raises(TypeError, match="input must be numeric"):
             _check_numeric_array([1, "1"], "input")
+
+
+class TestDataFrame:
+    def test_type(self):
+        with pytest.raises(
+            TypeError,
+            match="df must be a <class 'pandas.core.frame.DataFrame'> instance",
+        ):
+            _check_dataframe(1, {"one"}, "df")
+
+    def test_dataframe(self):
+        df = pd.DataFrame({"first": [0, 1], "second": [3, 4]})
+
+        with pytest.raises(
+            ValueError, match="columns must be included in df dataframe"
+        ):
+            _check_dataframe(df, ["column"], "df")
+
+        pd.testing.assert_frame_equal(
+            _check_dataframe(df, ["first"], "df"), df[["first"]]
+        )
+
+        pd.testing.assert_frame_equal(
+            _check_dataframe(df, ["first", "second"], "df"), df
+        )
+
+
+class TestNonDecreasing:
+    def test_non_decreasing(self):
+        with pytest.raises(ValueError, match="array must be non-decreasing"):
+            _check_non_decreasing([0, 1, 0], "array")
+
+        _check_non_decreasing([0, 0, 1, 1, 4], "array")