def klustakwik(
model=None,
spike_ids=None,
spike_clusters=None, # Initial spike clusters.
features=None,
masks=None,
iter_callback=None,
**kwargs):
"""Run the clustering algorithm on the model, or on any features
and masks.
Return the `spike_clusters` assignements.
"""
# Get the features and masks.
if model is not None:
if features is None:
features = PartialArray(model.all_features_masks, 0)
if masks is None:
masks = PartialArray(model.all_features_masks, 1)
# Select some spikes if needed.
if spike_ids is not None:
features = features[spike_ids]
masks = masks[spike_ids]
# Convert the features and masks to the sparse structure used
# by KK.
data = sparsify_features_masks(features, masks)
data = data.to_sparse_data()
# Instantiate KK.
from klustakwik2 import KK, __version__
kk = KK(data, **kwargs)
@kk.register_callback
def f(_):
# Skip split iterations.
if _.name != '':
return
if iter_callback:
iter_callback(kk.clusters)
# Launch KK.
if spike_clusters is None:
kk.cluster_mask_starts()
else:
kk.cluster_from(spike_clusters)
spike_clusters = kk.clusters
params = kk.params
params['version'] = __version__
return spike_clusters, params
def cluster(
self,
model=None,
spike_ids=None,
features=None,
masks=None,
):
"""Run the clustering algorithm on the model, or on any features
and masks.
Return the `spike_clusters` assignements.
Emit the `iter` event at every KlustaKwik iteration.
"""
# Get the features and masks.
if model is not None:
if features is None:
features = PartialArray(model.features_masks, 0)
if masks is None:
masks = PartialArray(model.features_masks, 1)
# Select some spikes if needed.
if spike_ids is not None:
features = features[spike_ids]
masks = masks[spike_ids]
# Convert the features and masks to the sparse structure used
# by KK.
data = sparsify_features_masks(features, masks)
data = data.to_sparse_data()
# Run KK.
from klustakwik2 import KK
kk = KK(data, **self._kwargs)
@kk.register_callback
def f(_):
# Skip split iterations.
if _.name != '':
return
self.emit('iter', kk.clusters)
self.params = kk.all_params
kk.cluster_mask_starts()
spike_clusters = kk.clusters
return spike_clusters
