adding time slice to ts.samples()

Author: mufernando

python/tests/test_highlevel.py

@@ -541,6 +541,22 @@ def test_samples(self):
             ]
         assert total == ts.num_samples
 
+    def test_samples_time(self):
+        rng = np.random.default_rng(seed=923)
+        ts = self.get_tree_sequence(2)
+        times = rng.exponential(size=10)
+        for time in times:
+            min_samples = []
+            max_samples = []
+            for node in ts.nodes():
+                if node.is_sample:
+                    if node.time <= time:
+                        min_samples.append(node.id)
+                    if node.time >= time:
+                        max_samples.append(node.id)
+            assert np.array.equal(np.array(min_samples), ts.sampels(min_time=time))
+            assert np.array.equal(np.array(max_samples), ts.sampels(max_time=time))
+
     def test_genotype_matrix_indexing(self):
         num_demes = 4
         ts = self.get_tree_sequence(num_demes)

python/tskit/trees.py

@@ -4800,7 +4800,9 @@ def get_samples(self, population_id=None):
         # Deprecated alias for samples()
         return self.samples(population_id)
 
-    def samples(self, population=None, population_id=None):
+    def samples(
+        self, population=None, population_id=None, max_time=None, min_time=None
+    ):
         """
         Returns an array of the sample node IDs in this tree sequence. If the
         ``population`` parameter is specified, only return sample IDs from that
@@ -4820,10 +4822,17 @@ def samples(self, population=None, population_id=None):
         if population_id is not None:
             population = population_id
         samples = self._ll_tree_sequence.get_samples()
+        keep = np.full(shape=samples.shape, fill_value=True)
         if population is not None:
             sample_population = self.tables.nodes.population[samples]
-            samples = samples[sample_population == population]
-        return samples
+            keep = np.logical_and(keep, sample_population == population)
+        if max_time is not None:
+            sample_times = self.tables.nodes.time[samples]
+            keep = np.logical_and(keep, sample_times <= max_time)
+        if min_time is not None:
+            sample_times = self.tables.nodes.time[samples]
+            keep = np.logical_and(keep, sample_times >= min_time)
+        return samples[keep]
 
     def write_fasta(self, output, sequence_ids=None, wrap_width=60):
         ""