fixed error in tests and docs

tskit-dev · Mar 28, 2023 · 2a492a6 · 2a492a6
1 parent d847276
commit 2a492a6
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 12 deletions.
diff --git a/docs/api.md b/docs/api.md
@@ -87,6 +87,26 @@ for discussion and examples of individual features.
     :members:
 ```
 
+```{eval-rst}
+.. autoclass:: msprime.NodeType
+    :members:
+
+    .. autoattribute:: RECOMBINANT
+         :annotation: = NodeType(1<<17)
+
+    .. autoattribute:: COMMON_ANCESTOR
+         :annotation: = NodeType(1<<18)
+
+    .. autoattribute:: MIGRANT
+         :annotation: = NodeType(1<<19)
+
+    .. autoattribute:: GENE_CONVERSION
+         :annotation: = NodeType(1<<21)
+
+    .. autoattribute:: PASS_THROUGH
+         :annotation: = NodeType(1<<22)
+```
+
 #### Models
 
 ```{eval-rst}

diff --git a/msprime/ancestry.py b/msprime/ancestry.py
@@ -549,11 +549,12 @@ def simulate(
         arising from common ancestor and recombination events in the output
         tree sequence. This will result in unary nodes (i.e., nodes in marginal
         trees that have only one child). Defaults to False.
-    :param msprime.NodeType additional_nodes: Retain all ancestry for any nodes of
-    the specified type. This will result in unary nodes. Defaults to msprime.NodeType(0).
+    :param NodeType additional_nodes: Retain all ancestry for any nodes of
+        the specified type. This will result in unary nodes.
+        Defaults to :class: `.NodeType(0).`
     :param bool coalescing_segments_only: If False, retain all ancestry for any
-    nodes that are coalescent nodes anywhere in the sequence. This will result
-    in unary nodes. Defaults to True.
+        nodes that are coalescent nodes anywhere in the sequence. This will result
+        in unary nodes. Defaults to True.
     :param model: The simulation model to use.
         This can either be a string (e.g., ``"smc_prime"``) or an instance of
         a ancestry model class (e.g, ``msprime.DiscreteTimeWrightFisher()``.
@@ -1178,12 +1179,12 @@ def sim_ancestry(
         tree sequence. This will result in unary nodes (i.e., nodes in marginal
         trees that have only one child). Defaults to False.
         See the :ref:`sec_ancestry_full_arg` section for examples.
-    :param msprime.NodeType additional_nodes: Retain all ancestry for any node
-    of the specified type. This will result in unary nodes.
-    Defaults to msprime.NodeType(0).
+    :param NodeType additional_nodes: Retain all ancestry for any node
+        of the specified type. This will result in unary nodes.
+        Defaults to class `msprime.NodeType(0)`.
     :param bool coalescing_segments_only: If False, retain all ancestry for any
-    nodes that are coalescent nodes anywhere in the sequence. This will result in
-    unary nodes. Defaults to True.
+        nodes that are coalescent nodes anywhere in the sequence. This will result in
+        unary nodes. Defaults to True.
     :param bool record_migrations: If True, record all migration events
         that occur in the :ref:`tskit:sec_migration_table_definition` of
         the output tree sequence. Defaults to False.
@@ -1325,7 +1326,7 @@ class ExitReason(enum.IntEnum):
 
 class NodeType(enum.Flag):
     """
-    The different node types associated with all possible events.
+    Specify the type of node for which you want to track ancestry.
     """
 
     RECOMBINANT = 1 << 17

diff --git a/tests/test_mutations.py b/tests/test_mutations.py
@@ -2395,7 +2395,7 @@ def test_record_unary(self):
             20,
             sequence_length=10,
             random_seed=2,
-            record_unary=True,
+            coalescing_segments_only=False,
             recombination_rate=1,
         )
         assert ts.num_trees > 1
@@ -2408,6 +2408,7 @@ def test_full_arg(self):
             sequence_length=10,
             random_seed=2,
             record_full_arg=True,
+            coalescing_segments_only=False,
             recombination_rate=1,
         )
         assert ts.num_trees > 1

diff --git a/tests/test_simulate_from.py b/tests/test_simulate_from.py
@@ -1295,7 +1295,7 @@ def test_record_unary(self):
             random_seed=4,
             recombination_rate=0.1,
             population_size=1,
-            record_unary=True,
+            coalescing_segments_only=False,
         )
         assert rts.num_trees - ts.num_trees > 0
         rts_simplified = rts.simplify()
@@ -1312,6 +1312,7 @@ def test_full_arg(self):
             recombination_rate=1,
             population_size=1,
             record_full_arg=True,
+            coalescing_segments_only=False,
         )
         assert rts.num_trees - ts.num_trees > 0
         flags = rts.tables.nodes.flags