Merge pull request #51 from tskit-dev/first_gen

added first gen function

pyslim/slim_tree_sequence.py

@@ -501,6 +501,13 @@ def individual_ages_at(self, time):
         ages[alive] = self.individual_times[alive] - time
         return ages
 
+    def first_generation_individuals(self):
+        """
+        Returns the IDs of the individuals remembered as part of the first SLiM generation,
+        as determined by their flags.
+        """
+        return np.where(self.tables.individuals.flags & INDIVIDUAL_FIRST_GEN > 0)[0]
+
 
 def _set_nodes_individuals(
         tables, node_ind=None, location=(0, 0, 0), age=0, ind_id=None,

tests/test_tree_sequence.py

@@ -103,6 +103,14 @@ def test_individual_embellishments(self):
                 self.assertEqual(ts.individual_populations[j], ind.population)
                 self.assertArrayEqual(ts.individual_locations[j], ind.location)
 
+    def test_first_gen(self):
+        for ts in self.get_slim_examples():
+            firstgen = ts.first_generation_individuals()
+            firstgen_times = ts.individual_times[firstgen]
+            self.assertEqual(len(set(firstgen_times)), 1)
+            self.assertEqual(firstgen_times[0], np.max(ts.individual_times))
+            for i in firstgen:
+                self.assertGreaterEqual(ts.individual(i).flags & pyslim.INDIVIDUAL_FIRST_GEN, 0)
 
 class TestNodeMetadata(tests.PyslimTestCase):
     '''