Added tests for BinaryFileGroup

Author: Jacob Pennington

tests/conftest.py

@@ -221,5 +221,3 @@ def bfile(saved_ops, torch_device, data_directory):
     bfile = io.bfile_from_ops(ops, filename=filename, device=torch_device)
 
     return bfile
-
-### End

tests/test_io.py

@@ -3,6 +3,7 @@
 from pathlib import Path
 
 import numpy as np
+import torch
 
 from kilosort import io
 
@@ -231,3 +232,42 @@ def test_tmax_only(torch_device, data_directory):
     finally:
         # Delete memmap file and re-raise exception
         path.unlink()
+
+
+def test_file_group(torch_device, data_directory, bfile):
+    file = data_directory / 'ZFM-02370_mini.imec0.ap.short.bin'
+    fs = bfile.fs
+    n_chans = bfile.n_chan_bin
+
+    # Test with file_objects option
+    # Load as three 15-second files instead of one 45-second file.
+    objs = [np.memmap(file, dtype='int16', shape=bfile.shape, mode='r')
+            for _ in range(3)]
+    objs[0] = objs[0][:int(15*fs),:]
+    objs[1] = objs[1][int(15*fs):int(30*fs),:]
+    objs[2] = objs[2][int(30*fs):,:]
+    bfg = io.BinaryFileGroup(file_objects=objs)
+    bfile2 = io.BinaryFiltered(
+        filename='test', n_chan_bin=n_chans, fs=fs, chan_map=bfile.chan_map,
+        device=torch_device, file_object=bfg, dtype='int16'
+        )
+
+    # First batch, overlapping batch, and last batch
+    # (assumes 45s test dataset with 2s batch size)
+    for i in [0, 7, 22]:
+        b1 = bfile.padded_batch_to_torch(i, skip_preproc=True)
+        b2 = bfile2.padded_batch_to_torch(i)
+        assert torch.allclose(b1, b2)
+
+    # Test with filenames option
+    files = [file]*3  # Load the same data three times
+    bfile3 = io.BinaryFiltered(
+        filename=files, n_chan_bin=n_chans, fs=fs, chan_map=bfile.chan_map,
+        device=torch_device, dtype='int16'
+    )
+
+    # First and first, last and last, last of original and last of concat
+    for i,j in [(0,0), (21,21), (22,67)]:
+        b1 = bfile.padded_batch_to_torch(i, skip_preproc=True)
+        b2 = bfile3.padded_batch_to_torch(j)
+        assert torch.allclose(b1, b2)