DEBUG: track sampler was just yielding one time. Now infinite.

bigwig_loader/sampler/track_sampler.py

@@ -8,4 +8,5 @@ def __init__(self, total_number_of_tracks: int, sample_size: int):
         self.sample_size = sample_size
 
     def __iter__(self) -> Iterator[list[int]]:
-        yield sorted(sample(range(self.total_number_of_tracks), self.sample_size))
+        while True:
+            yield sorted(sample(range(self.total_number_of_tracks), self.sample_size))

tests/test_new_dataset.py

@@ -114,5 +114,6 @@ def test_batch_return_type(bigwig_path, reference_genome_path, merged_intervals)
         sub_sample_tracks=1,
     )
 
+    print("start")
     test_output_shape_sub_sampled_tracks(ds)
     print("done")

tests/test_track_sampler.py

@@ -0,0 +1,13 @@
+from bigwig_loader.sampler.track_sampler import TrackSampler
+
+
+def test_track_samler():
+    sampler = TrackSampler(total_number_of_tracks=40, sample_size=10)
+    samples = []
+    for i, sample in enumerate(sampler):
+        samples.append(sample)
+        assert len(sample) == 10
+        assert all([0 <= s < 40 for s in sample])
+        if i == 100:
+            break
+    assert len(samples) == 101