defined individual_ages for WF models

tskit-dev · Jul 1, 2019 · 6121876 · 6121876
1 parent 5542cbe
commit 6121876
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 4 deletions.
diff --git a/pyslim/slim_tree_sequence.py b/pyslim/slim_tree_sequence.py
@@ -145,8 +145,9 @@ def __init__(self, ts, reference_sequence=None):
         self.individual_ages = np.zeros(ts.num_individuals, dtype='int')
         self.individual_populations = np.zeros(ts.num_individuals, dtype='int')
         for j, ind in enumerate(ts.individuals()):
-            md = decode_individual(ind.metadata)
-            self.individual_ages[j] = md.age
+            if self.slim_provenance.model_type != "WF":
+                md = decode_individual(ind.metadata)
+                self.individual_ages[j] = md.age
             populations = [self.node(n).population for n in ind.nodes]
             if len(set(populations)) > 1:
                 raise ValueError("Individual has nodes from more than one population.")

diff --git a/tests/__init__.py b/tests/__init__.py
@@ -13,7 +13,8 @@
 import os
 
 # recipes that record everyone ever
-_everyone_example_files = ["tests/examples/recipe_record_everyone"]
+_everyone_example_files = ["tests/examples/recipe_record_everyone",
+                           "tests/examples/recipe_record_everyone_WF"]
 
 _wf_example_files = ["tests/examples/recipe_{}".format(x)
                   for x in ['WF', 'nucleotides', 'long_nucleotides']]

diff --git a/tests/examples/recipe_record_everyone_WF.slim b/tests/examples/recipe_record_everyone_WF.slim
@@ -0,0 +1,29 @@
+initialize()
+{
+    setSeed(23);
+    initializeSLiMOptions(dimensionality="xy");
+    initializeTreeSeq();
+    initializeMutationRate(1e-2);
+    initializeMutationType("m1", 0.5, "f", -0.1);
+    initializeGenomicElementType("g1", m1, 1.0);
+    initializeGenomicElement(g1, 0, 99);
+    initializeRecombinationRate(1e-2);
+    defineConstant("K", 10);
+}
+
+1 early() {
+    sim.addSubpop("p1", 10);
+    for (ind in p1.individuals)
+        ind.setSpatialPosition(p1.pointUniform());
+}
+
+early() {
+    sim.treeSeqRememberIndividuals(p1.individuals);
+}
+
+10 {
+    sim.treeSeqOutput("recipe_record_everyone_WF.trees");
+    catn("Done.");
+    sim.simulationFinished();
+}
+
diff --git a/tests/test_tree_sequence.py b/tests/test_tree_sequence.py
@@ -171,6 +171,7 @@ class TestEveryone(tests.PyslimTestCase):
 
     def test_alive_ages(self):
         for ts in self.get_slim_everyone_examples():
+            is_WF = (ts.slim_provenance.model_type == "WF")
             ages_mat = np.zeros((ts.num_individuals, ts.slim_generation))
             alive_mat = np.zeros((ts.num_individuals, ts.slim_generation))
             alive_now = ts.individuals_alive_at(0)
@@ -180,9 +181,13 @@ def test_alive_ages(self):
                     alive_mat[j, time] = 1
 
             for j, ind in enumerate(ts.individuals()):
+                if is_WF:
+                    age = 0
+                else:
+                    age = ind.metadata.age
                 self.assertEqual(j in alive_now,
                                  ind.flags & pyslim.INDIVIDUAL_ALIVE > 0)
-                self.assertEqual(sum(alive_mat[j, :]), 1 + ind.metadata.age)
+                self.assertEqual(sum(alive_mat[j, :]), 1 + age)
                 self.assertEqual(alive_mat[j, int(ind.time)], 1)
                 self.assertEqual(ages_mat[j, int(ind.time)], 0)
                 for t in range(ts.slim_generation):