From 2483eb18cd8dd8a6d568b56b67582db4e0d7f189 Mon Sep 17 00:00:00 2001 From: Jacob Pennington Date: Tue, 26 Mar 2024 14:43:03 -0400 Subject: [PATCH 1/5] fixed whitening matrix inv for phy --- kilosort/io.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/kilosort/io.py b/kilosort/io.py index 83640da8..9d0dac23 100644 --- a/kilosort/io.py +++ b/kilosort/io.py @@ -158,7 +158,9 @@ def save_to_phy(st, clu, tF, Wall, probe, ops, imin, results_dir=None, # whitening matrix ** saving real whitening matrix doesn't work with phy currently whitening_mat = ops['Wrot'].cpu().numpy() np.save((results_dir / 'whitening_mat_dat.npy'), whitening_mat) - whitening_mat = 0.005 * np.eye(len(chan_map), dtype='float32') + # NOTE: commented out for reference, this was different in KS 2.5 because + # the binary file was already whitened. + # whitening_mat = 0.005 * np.eye(len(chan_map), dtype='float32') whitening_mat_inv = np.linalg.inv(whitening_mat + 1e-5 * np.eye(whitening_mat.shape[0])) np.save((results_dir / 'whitening_mat.npy'), whitening_mat) np.save((results_dir / 'whitening_mat_inv.npy'), whitening_mat_inv) From c00f68e1c5e331484d41c0cb56e5138b7500b89c Mon Sep 17 00:00:00 2001 From: Jacob Pennington Date: Wed, 27 Mar 2024 11:59:30 -0400 Subject: [PATCH 2/5] unwhitened template amplitudes, saved normalized templates --- kilosort/io.py | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/kilosort/io.py b/kilosort/io.py index 9d0dac23..71034285 100644 --- a/kilosort/io.py +++ b/kilosort/io.py @@ -156,14 +156,18 @@ def save_to_phy(st, clu, tF, Wall, probe, ops, imin, results_dir=None, np.save((results_dir / 'channel_positions.npy'), channel_positions) # whitening matrix ** saving real whitening matrix doesn't work with phy currently - whitening_mat = ops['Wrot'].cpu().numpy() - np.save((results_dir / 'whitening_mat_dat.npy'), whitening_mat) + whitening_mat = ops['Wrot'] + np.save((results_dir / 'whitening_mat_dat.npy'), whitening_mat.cpu()) # NOTE: commented out for reference, this was different in KS 2.5 because # the binary file was already whitened. # whitening_mat = 0.005 * np.eye(len(chan_map), dtype='float32') - whitening_mat_inv = np.linalg.inv(whitening_mat + 1e-5 * np.eye(whitening_mat.shape[0])) - np.save((results_dir / 'whitening_mat.npy'), whitening_mat) - np.save((results_dir / 'whitening_mat_inv.npy'), whitening_mat_inv) + whitening_mat_inv = torch.inverse( + whitening_mat + + 1e-5 * torch.eye(whitening_mat.shape[0]).to(whitening_mat.device) + ) + #whitening_mat_inv = np.linalg.inv(whitening_mat + 1e-5 * np.eye(whitening_mat.shape[0])) + np.save((results_dir / 'whitening_mat.npy'), whitening_mat.cpu()) + np.save((results_dir / 'whitening_mat_inv.npy'), whitening_mat_inv.cpu()) # spike properties spike_times = st[:,0].astype('int64') + imin # shift by minimum sample index @@ -191,13 +195,19 @@ def save_to_phy(st, clu, tF, Wall, probe, ops, imin, results_dir=None, # template properties similar_templates = CCG.similarity(Wall, ops['wPCA'].contiguous(), nt=ops['nt']) - template_amplitudes = ((Wall**2).sum(axis=(-2,-1))**0.5).cpu().numpy() + temp_amplitudes = ((Wall**2).sum(axis=(-2,-1))**0.5).cpu().numpy() templates = (Wall.unsqueeze(-1).cpu() * ops['wPCA'].cpu()).sum(axis=-2).numpy() templates = templates.transpose(0,2,1) + # normalize templates by amplitude + templates = templates / temp_amplitudes[:, np.newaxis, np.newaxis] templates_ind = np.tile(np.arange(Wall.shape[1])[np.newaxis, :], (templates.shape[0],1)) np.save((results_dir / 'similar_templates.npy'), similar_templates) np.save((results_dir / 'templates.npy'), templates) np.save((results_dir / 'templates_ind.npy'), templates_ind) + # get unwhitened template amplitudes to use as cluster_Amplitudes + iwrot = whitening_mat_inv.to(Wall.device) + unwhitened = torch.einsum('jk, ikl -> ijl', iwrot, Wall) + template_amplitudes = ((unwhitened**2).sum(axis=(-2,-1))**0.5).cpu().numpy() # contamination ratio acg_threshold = ops['settings']['acg_threshold'] From ebdaf5de4c3c4c19ba8f138184cb705ed8720e57 Mon Sep 17 00:00:00 2001 From: Jacob Pennington Date: Tue, 2 Apr 2024 15:01:50 -0400 Subject: [PATCH 3/5] reverted cluster amp calculation to use whitened templates --- kilosort/io.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/kilosort/io.py b/kilosort/io.py index 71034285..855fc686 100644 --- a/kilosort/io.py +++ b/kilosort/io.py @@ -165,7 +165,6 @@ def save_to_phy(st, clu, tF, Wall, probe, ops, imin, results_dir=None, whitening_mat + 1e-5 * torch.eye(whitening_mat.shape[0]).to(whitening_mat.device) ) - #whitening_mat_inv = np.linalg.inv(whitening_mat + 1e-5 * np.eye(whitening_mat.shape[0])) np.save((results_dir / 'whitening_mat.npy'), whitening_mat.cpu()) np.save((results_dir / 'whitening_mat_inv.npy'), whitening_mat_inv.cpu()) @@ -195,19 +194,19 @@ def save_to_phy(st, clu, tF, Wall, probe, ops, imin, results_dir=None, # template properties similar_templates = CCG.similarity(Wall, ops['wPCA'].contiguous(), nt=ops['nt']) - temp_amplitudes = ((Wall**2).sum(axis=(-2,-1))**0.5).cpu().numpy() + template_amplitudes = ((Wall**2).sum(axis=(-2,-1))**0.5).cpu().numpy() templates = (Wall.unsqueeze(-1).cpu() * ops['wPCA'].cpu()).sum(axis=-2).numpy() templates = templates.transpose(0,2,1) # normalize templates by amplitude - templates = templates / temp_amplitudes[:, np.newaxis, np.newaxis] + # TODO: check if this helps / hurts going between snippets and templates + # scale should not change when switching between + # TODO: post issue on phy github asking where 'amp' is actually coming from, + # other issue answers are old and don't seem to point anywhere relevant + templates = templates / template_amplitudes[:, np.newaxis, np.newaxis] templates_ind = np.tile(np.arange(Wall.shape[1])[np.newaxis, :], (templates.shape[0],1)) np.save((results_dir / 'similar_templates.npy'), similar_templates) np.save((results_dir / 'templates.npy'), templates) np.save((results_dir / 'templates_ind.npy'), templates_ind) - # get unwhitened template amplitudes to use as cluster_Amplitudes - iwrot = whitening_mat_inv.to(Wall.device) - unwhitened = torch.einsum('jk, ikl -> ijl', iwrot, Wall) - template_amplitudes = ((unwhitened**2).sum(axis=(-2,-1))**0.5).cpu().numpy() # contamination ratio acg_threshold = ops['settings']['acg_threshold'] From 91dfa45bd60a7bc909c1c7bbdb069ea7d99128f6 Mon Sep 17 00:00:00 2001 From: jacobpennington Date: Fri, 3 May 2024 11:07:07 -0700 Subject: [PATCH 4/5] removed outdated comment --- kilosort/io.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kilosort/io.py b/kilosort/io.py index 855fc686..b5931044 100644 --- a/kilosort/io.py +++ b/kilosort/io.py @@ -155,7 +155,7 @@ def save_to_phy(st, clu, tF, Wall, probe, ops, imin, results_dir=None, np.save((results_dir / 'channel_map.npy'), chan_map) np.save((results_dir / 'channel_positions.npy'), channel_positions) - # whitening matrix ** saving real whitening matrix doesn't work with phy currently + # whitening matrix whitening_mat = ops['Wrot'] np.save((results_dir / 'whitening_mat_dat.npy'), whitening_mat.cpu()) # NOTE: commented out for reference, this was different in KS 2.5 because From be0ef920807a0a402b0bd7eddba6a5a105674507 Mon Sep 17 00:00:00 2001 From: jacobpennington Date: Fri, 3 May 2024 12:22:55 -0700 Subject: [PATCH 5/5] Removed template normalization --- kilosort/io.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/kilosort/io.py b/kilosort/io.py index b5931044..19940f10 100644 --- a/kilosort/io.py +++ b/kilosort/io.py @@ -197,12 +197,6 @@ def save_to_phy(st, clu, tF, Wall, probe, ops, imin, results_dir=None, template_amplitudes = ((Wall**2).sum(axis=(-2,-1))**0.5).cpu().numpy() templates = (Wall.unsqueeze(-1).cpu() * ops['wPCA'].cpu()).sum(axis=-2).numpy() templates = templates.transpose(0,2,1) - # normalize templates by amplitude - # TODO: check if this helps / hurts going between snippets and templates - # scale should not change when switching between - # TODO: post issue on phy github asking where 'amp' is actually coming from, - # other issue answers are old and don't seem to point anywhere relevant - templates = templates / template_amplitudes[:, np.newaxis, np.newaxis] templates_ind = np.tile(np.arange(Wall.shape[1])[np.newaxis, :], (templates.shape[0],1)) np.save((results_dir / 'similar_templates.npy'), similar_templates) np.save((results_dir / 'templates.npy'), templates)