tests for merged output

Author: mayofaulkner

phylib/io/alf.py

@@ -270,7 +270,6 @@ def make_template_and_spikes_objects(self):
             templates_inds = np.zeros((n_clusters, ncw), dtype=np.int32)
             # for each template, find the nearest channels to keep (one the same probe...)
             for t in np.arange(n_clusters):
-                # here we need to fill with nans if it doesn't exists, but then can no longet be int (sorry) # or have it all 0
                 channels = self.model.clusters_channels
 
                 current_probe = self.model.channel_probes[channels[t]]

phylib/io/model.py

@@ -413,16 +413,14 @@ def _load_data(self):
             self.n_channels_loc = 0
 
         # Clusters waveforms
-        if np.all(self.spike_clusters == self.spike_templates):
+        if not np.all(self.spike_clusters == self.spike_templates) and self.sparse_templates.cols is None:
+            self.merge_map, _ = self.get_merge_map()
+            self.sparse_clusters = self.cluster_waveforms()
+            self.n_clusters = self.spike_clusters.max() + 1
+        else:
             self.merge_map = {}
-            self.nan_clusters = []
             self.sparse_clusters = self.sparse_templates
             self.n_clusters = self.spike_templates.max() + 1
-        else:
-            if self.sparse_templates.cols is None:
-                self.merge_map, self.nan_clusters = self.get_merge_map()
-                self.sparse_clusters = self.cluster_waveforms()
-                self.n_clusters = self.spike_clusters.max() + 1
 
         # Spike waveforms (optional, otherwise fetched from raw data as needed).
         self.spike_waveforms = self._load_spike_waveforms()

phylib/io/tests/conftest.py

@@ -98,6 +98,19 @@ def _make_dataset(tempdir, param='dense', has_spike_attributes=True):
         _remove(tempdir / 'whitening_mat_inv.npy')
         _remove(tempdir / 'sim_binary.dat')
 
+    if param == 'merged':
+        # remove this file to make templates dense
+        _remove(tempdir / 'template_ind.npy')
+        clus = np.load(tempdir / 'spike_clusters.npy')
+        max_clus = np.max(clus)
+        # merge cluster 0 and 1
+        clus[np.bitwise_or(clus == 0, clus == 1)] = max_clus + 1
+        # split cluster 9 into two clusters
+        idx = np.where(clus == 9)[0]
+        clus[idx[0:3]] = max_clus + 2
+        clus[idx[3:]] = max_clus + 3
+        np.save(tempdir / 'spike_clusters.npy', clus)
+
     # Spike attributes.
     if has_spike_attributes:
         write_array(tempdir / 'spike_fail.npy', np.full(10, np.nan))  # wrong number of spikes
@@ -120,7 +133,7 @@ def _make_dataset(tempdir, param='dense', has_spike_attributes=True):
     return template_path
 
 
-@fixture(scope='function', params=('dense', 'sparse', 'misc'))
+@fixture(scope='function', params=('dense', 'sparse', 'misc', 'merged'))
 def template_path_full(tempdir, request):
     return _make_dataset(tempdir, request.param)
 

phylib/io/tests/test_alf.py

@@ -210,3 +210,72 @@ def read_after_write():
     c.convert(out_path, label='probe00')
     check_conversion_output()
     read_after_write()
+
+
+def test_merger(dataset):
+
+    path = Path(dataset.tmp_dir)
+    out_path = path / 'alf'
+
+    model = TemplateModel(
+        dir_path=path, dat_path=dataset.dat_path, sample_rate=2000, n_channels_dat=dataset.nc)
+
+    c = EphysAlfCreator(model)
+    c.convert(out_path)
+
+    model.close()
+
+    # path.joinpath('_phy_spikes_subset.channels.npy').unlink()
+    # path.joinpath('_phy_spikes_subset.waveforms.npy').unlink()
+    # path.joinpath('_phy_spikes_subset.spikes.npy').unlink()
+
+    out_path_merge = path / 'alf_merge'
+    spike_clusters = dataset._load('spike_clusters.npy')
+    clu, n_clu = np.unique(spike_clusters, return_counts=True)
+
+    # merge the first two clusters
+    merge_clu = clu[0:2]
+    spike_clusters[np.bitwise_or(spike_clusters == clu[0], spike_clusters == clu[1])] = np.max(clu) + 1
+    # split the cluster with the most spikes
+    split_clu = clu[-1]
+    idx = np.where(spike_clusters == split_clu)[0]
+    spike_clusters[idx[0:int(n_clu[-1] / 2)]] = np.max(clu) + 2
+    spike_clusters[idx[int(n_clu[-1] / 2):]] = np.max(clu) + 3
+
+    np.save(path / 'spike_clusters.npy', spike_clusters)
+
+    model = TemplateModel(
+        dir_path=path, dat_path=dataset.dat_path, sample_rate=2000, n_channels_dat=dataset.nc)
+    print(model.merge_map)
+    c = EphysAlfCreator(model)
+    c.convert(out_path_merge)
+
+    # Test that the split are the same for the expected datasets
+    clu_old = np.load(next(out_path.glob('clusters.peakToTrough.npy')))
+    clu_new = np.load(next(out_path_merge.glob('clusters.peakToTrough.npy')))
+    assert clu_old[split_clu] == clu_new[np.max(clu) + 2]
+    assert clu_old[split_clu] == clu_new[np.max(clu) + 3]
+    assert clu_new[split_clu] == 0
+    assert clu_new[merge_clu[0]] == 0
+    assert clu_new[merge_clu[1]] == 0
+
+    clu_old = np.load(next(out_path.glob('clusters.channels.npy')))
+    clu_new = np.load(next(out_path_merge.glob('clusters.channels.npy')))
+    assert clu_old[split_clu] == clu_new[np.max(clu) + 2]
+    assert clu_old[split_clu] == clu_new[np.max(clu) + 3]
+    assert clu_new[split_clu] == 0
+    assert clu_new[merge_clu[0]] == 0
+    assert clu_new[merge_clu[1]] == 0
+
+    clu_old = np.load(next(out_path.glob('clusters.depths.npy')))
+    clu_new = np.load(next(out_path_merge.glob('clusters.depths.npy')))
+    assert clu_old[split_clu] == clu_new[np.max(clu) + 2]
+    assert clu_old[split_clu] == clu_new[np.max(clu) + 3]
+    assert clu_new[split_clu] == 0
+    assert clu_new[merge_clu[0]] == 0
+    assert clu_new[merge_clu[1]] == 0
+
+    clu_old = np.load(next(out_path.glob('clusters.waveformsChannels.npy')))
+    clu_new = np.load(next(out_path_merge.glob('clusters.waveformsChannels.npy')))
+    assert np.array_equal(clu_old[split_clu], clu_new[np.max(clu) + 2])
+    assert np.array_equal(clu_old[split_clu], clu_new[np.max(clu) + 3])

phylib/io/tests/test_model.py

@@ -14,7 +14,8 @@
 
 # from phylib.utils import Bunch
 from phylib.utils.testing import captured_output
-from ..model import from_sparse, load_model
+# from ..model import from_sparse, load_model
+from phylib.io.model import from_sparse, load_model
 
 logger = logging.getLogger(__name__)
 
@@ -113,6 +114,22 @@ def test_model_depth(template_model):
     assert depths.shape == (template_model.n_spikes,)
 
 
+def test_model_merge(template_model_full):
+    m = template_model_full
+
+    # This is the case where we can do the merging
+    if not np.all(m.spike_templates == m.spike_clusters) and m.sparse_clusters.cols is None:
+        assert len(m.merge_map) > 0
+        assert not np.array_equal(m.sparse_clusters.data, m.sparse_templates.data)
+        assert m.sparse_clusters.data.shape[0] == m.n_clusters
+        assert m.sparse_templates.data.shape[0] == m.n_templates
+
+    else:
+        assert len(m.merge_map) == 0
+        assert np.array_equal(m.sparse_clusters.data, m.sparse_templates.data)
+        assert np.array_equal(m.n_templates, m.n_clusters)
+
+
 def test_model_save(template_model_full):
     m = template_model_full
     m.save_metadata('test', {1: 1})