CODE: sim_trait

Update the sim_trait function.

tests/data.py

@@ -6,10 +6,10 @@ def binary_tree():
     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
+    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
@@ -162,7 +162,7 @@ def binary_tree_seq():
     """Sample tree sequence
     Ancestral State: A
     Genotype: [Node0, Node1, Node2, Node3]
-    Site0: [T, A, T, T], has 1 mutation
+    Site0: [A, T, 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
 
@@ -194,3 +194,66 @@ def binary_tree_seq():
     ts = tables.tree_sequence()
 
     return ts
+
+
+def simple_tree_seq():
+    """
+    Tree sequence data with a single mutation in each tree.
+    Ancestral State: A
+    Genotype: [Node0, Node1, Node2, Node3]
+    Site0: [A, T, T, T]
+    Site1: [A, A, A, A]
+    Site2: [C, A, A, A]
+
+    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=4, derived_state="A")
+    tables.mutations.add_row(site=2, node=0, derived_state="C")
+
+    ts = tables.tree_sequence()
+
+    return ts
+
+
+def allele_freq_one():
+    """
+    Sample tree sequence with allele frequence one at a single site
+    with ancestral state A.
+    """
+    ts = tskit.Tree.generate_balanced(4, span=10).tree_sequence
+    tables = ts.dump_tables()
+    tables.sites.add_row(0, "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=0, node=5, derived_state="A", parent=0)
+
+    ts = tables.tree_sequence()
+
+    return ts