<commit_before>"""This module provides utility classes and functions to load spike sorting
data sets."""

# -----------------------------------------------------------------------------
# Imports
# -----------------------------------------------------------------------------
import os
import os.path
import shutil
import re
from collections import Counter

import numpy as np
import pandas as pd
import tables as tb

from loader import (Loader, default_group_info, reorder, renumber_clusters,
    default_cluster_info)
from klustersloader import (find_filenames, save_clusters, convert_to_clu,
    find_filename, find_filename_or_new)
from tools import (load_text, normalize,
    load_binary, load_pickle, save_text, get_array,
    first_row, load_binary_memmap)
from selection import (select, select_pairs, get_spikes_in_clusters,
    get_some_spikes_in_clusters, get_some_spikes, get_indices, pandaize)
from kwiklib.utils.logger import (debug, info, warn, exception, FileLogger,
    register, unregister)
from kwiklib.utils.colors import COLORS_COUNT, generate_colors
from kwiklib.dataio.kwik import add_cluster
from kwiklib.dataio.klusterskwik import klusters_to_kwik
from .experiment import Experiment


def add_missing_clusters(exp):

    shanks = sorted(exp.channel_groups.keys())

    for shank in shanks:
        cg = exp.channel_groups[shank]
        clusters = cg.clusters.main.keys()
        clusters_unique = np.unique(cg.spikes.clusters.main[:])
        # Find missing clusters in the kwik file.
        missing = sorted(set(clusters_unique)-set(clusters))

        # Add all missing clusters with a default color and "Unsorted" cluster group (group #3).
        for idx in missing:
            info("Adding missing cluster %d in shank %d." % (idx, shank))
            add_cluster(exp._files, channel_group_id='%d' % shank,
                        id=str(idx),
                        clustering='main',
                        cluster_group=3)


# -----------------------------------------------------------------------------
# HDF5 Loader
# -----------------------------------------------------------------------------
class KwikLoader(Loader):
    # TODO: change the clustering ('main' by default)

    def __init__(self, parent=None, filename=None, userpref=None):
        self.experiment = None
        super(KwikLoader, self).__init__(parent=parent, filename=filename, userpref=userpref)

    # Read functions.
    # ---------------
    def _report_progress_open(self, spike, nspikes, shank, nshanks):
        i = shank * 100 + float(spike)/nspikes*100
        n = nshanks * 100
        self.report_progress(i, n)

    def _consistency_check(self):
        exp = self.experiment
        chgrp = self.shank

        cg = exp.channel_groups[chgrp]
        clusters = cg.clusters.main.keys()
        clusters_unique = np.unique(cg.spikes.clusters.main[:])

        # Find missing clusters in the kwik file.
        missing = sorted(set(clusters_unique)-set(clusters))

        # Add all missing clusters with a default color and "Unsorted" cluster group (group #3).
        for idx in missing:
            warn("Consistency check: adding cluster %d in the kwik file" % idx)
            add_cluster(exp._files, channel_group_id='%d' % chgrp,
                        id=idx,
                        clustering='main',
                        cluster_group=3)

    def open(self, filename=None):
        """Open everything."""
        if filename is None:
            filename = self.filename
        else:
            self.filename = filename
        dir, basename = os.path.split(filename)

        # Converting to kwik if needed
        # kwik = find_filename(basename, 'kwik', dir=dir)
        # xml = find_filename(basename, 'xml', dir=dir)
        # self.filename_clu = find_filename(basename, 'clu', dir=dir)
        self._filenames = find_filenames(filename)
        kwik = find_filename(basename, 'kwik', dir=dir)
        xml = self._filenames['xml']
        clu = self._filenames['clu']

        self.log_filename = find_filename_or_new(filename, 'kvlog', dir=dir)


        # Backup the .clu file.
        clu_original = find_filename_or_new(filename, 'clu_original')
        if os.path.exists(clu) and not os.path.exists(clu_original):
            shutil.copyfile(clu, clu_original)

        if not kwik:
            assert xml, ValueError("I need a valid .kwik file")
            return

        self.experiment = Experiment(basename, dir=dir, mode='a')

        # CONSISTENCY CHECK
        # add missing clusters
        add_missing_clusters(self.experiment)

        # TODO
        # self.initialize_logfile()
        # Load the similarity measure chosen by the user in the preferences
        # file: 'gaussian' or 'kl'.
        # Refresh the preferences file when a new file is opened.
        # USERPREF.refresh()
        self.similarity_measure = self.userpref['similarity_measure'] or 'gaussian'
        debug("Similarity measure: {0:s}.".format(self.similarity_measure))
        info("Opening {0:s}.".format(self.experiment.name))
        self.shanks = sorted(self.experiment.channel_groups.keys())

        self.freq = self.experiment.application_data.spikedetekt.sample_rate

        self.fetdim = self.experiment.application_data.spikedetekt.n_features_per_channel
        self.nsamples = self.experiment.application_data.spikedetekt.extract_s_before + self.experiment.application_data.spikedetekt.extract_s_after

        self.set_shank(self.shanks[0])

    # Shank functions.
    # ----------------
    def get_shanks(self):
        """Return the list of shanks available in the file."""
        return self.shanks

    def set_shank(self, shank):
        """Change the current shank and read the corresponding tables."""
        if not shank in self.shanks:
            warn("Shank {0:d} is not in the list of shanks: {1:s}".format(
                shank, str(self.shanks)))
            return
        self.shank = shank
<commit_msg>Add shank option when opening a dataset<commit_after>"""This module provides utility classes and functions to load spike sorting
data sets."""

# -----------------------------------------------------------------------------
# Imports
# -----------------------------------------------------------------------------
import os
import os.path
import shutil
import re
from collections import Counter

import numpy as np
import pandas as pd
import tables as tb

from loader import (Loader, default_group_info, reorder, renumber_clusters,
    default_cluster_info)
from klustersloader import (find_filenames, save_clusters, convert_to_clu,
    find_filename, find_filename_or_new)
from tools import (load_text, normalize,
    load_binary, load_pickle, save_text, get_array,
    first_row, load_binary_memmap)
from selection import (select, select_pairs, get_spikes_in_clusters,
    get_some_spikes_in_clusters, get_some_spikes, get_indices, pandaize)
from kwiklib.utils.logger import (debug, info, warn, exception, FileLogger,
    register, unregister)
from kwiklib.utils.colors import COLORS_COUNT, generate_colors
from kwiklib.dataio.kwik import add_cluster
from kwiklib.dataio.klusterskwik import klusters_to_kwik
from .experiment import Experiment


def add_missing_clusters(exp):

    shanks = sorted(exp.channel_groups.keys())

    for shank in shanks:
        cg = exp.channel_groups[shank]
        clusters = cg.clusters.main.keys()
        clusters_unique = np.unique(cg.spikes.clusters.main[:])
        # Find missing clusters in the kwik file.
        missing = sorted(set(clusters_unique)-set(clusters))

        # Add all missing clusters with a default color and "Unsorted" cluster group (group #3).
        for idx in missing:
            info("Adding missing cluster %d in shank %d." % (idx, shank))
            add_cluster(exp._files, channel_group_id='%d' % shank,
                        id=str(idx),
                        clustering='main',
                        cluster_group=3)


# -----------------------------------------------------------------------------
# HDF5 Loader
# -----------------------------------------------------------------------------
class KwikLoader(Loader):
    # TODO: change the clustering ('main' by default)

    def __init__(self, parent=None, filename=None, userpref=None):
        self.experiment = None
        super(KwikLoader, self).__init__(parent=parent, filename=filename, userpref=userpref)

    # Read functions.
    # ---------------
    def _report_progress_open(self, spike, nspikes, shank, nshanks):
        i = shank * 100 + float(spike)/nspikes*100
        n = nshanks * 100
        self.report_progress(i, n)

    def _consistency_check(self):
        exp = self.experiment
        chgrp = self.shank

        cg = exp.channel_groups[chgrp]
        clusters = cg.clusters.main.keys()
        clusters_unique = np.unique(cg.spikes.clusters.main[:])

        # Find missing clusters in the kwik file.
        missing = sorted(set(clusters_unique)-set(clusters))

        # Add all missing clusters with a default color and "Unsorted" cluster group (group #3).
        for idx in missing:
            warn("Consistency check: adding cluster %d in the kwik file" % idx)
            add_cluster(exp._files, channel_group_id='%d' % chgrp,
                        id=idx,
                        clustering='main',
                        cluster_group=3)

    def open(self, filename=None, shank=None):
        """Open everything."""
        if filename is None:
            filename = self.filename
        else:
            self.filename = filename
        dir, basename = os.path.split(filename)

        # Converting to kwik if needed
        # kwik = find_filename(basename, 'kwik', dir=dir)
        # xml = find_filename(basename, 'xml', dir=dir)
        # self.filename_clu = find_filename(basename, 'clu', dir=dir)
        self._filenames = find_filenames(filename)
        kwik = find_filename(basename, 'kwik', dir=dir)
        xml = self._filenames['xml']
        clu = self._filenames['clu']

        self.log_filename = find_filename_or_new(filename, 'kvlog', dir=dir)


        # Backup the .clu file.
        clu_original = find_filename_or_new(filename, 'clu_original')
        if os.path.exists(clu) and not os.path.exists(clu_original):
            shutil.copyfile(clu, clu_original)

        if not kwik:
            assert xml, ValueError("I need a valid .kwik file")
            return

        self.experiment = Experiment(basename, dir=dir, mode='a')

        # CONSISTENCY CHECK
        # add missing clusters
        add_missing_clusters(self.experiment)

        # TODO
        # self.initialize_logfile()
        # Load the similarity measure chosen by the user in the preferences
        # file: 'gaussian' or 'kl'.
        # Refresh the preferences file when a new file is opened.
        # USERPREF.refresh()
        self.similarity_measure = self.userpref['similarity_measure'] or 'gaussian'
        debug("Similarity measure: {0:s}.".format(self.similarity_measure))
        info("Opening {0:s}.".format(self.experiment.name))
        self.shanks = sorted(self.experiment.channel_groups.keys())

        self.freq = self.experiment.application_data.spikedetekt.sample_rate

        self.fetdim = self.experiment.application_data.spikedetekt.n_features_per_channel
        self.nsamples = self.experiment.application_data.spikedetekt.extract_s_before + self.experiment.application_data.spikedetekt.extract_s_after

        self.set_shank(shank or self.shanks[0])

    # Shank functions.
    # ----------------
    def get_shanks(self):
        """Return the list of shanks available in the file."""
        return self.shanks

    def set_shank(self, shank):
        """Change the current shank and read the corresponding tables."""
        if not shank in self.shanks:
            warn("Shank {0:d} is not in the list of shanks: {1:s}".format(
                shank, str(self.shanks)))
            return
        self.shank = shank

{'fraction_non_alphanumeric': 0.14589412255106293, 'fraction_numerical': 0.003167986661108795, 'mean_word_length': 2.249187432286024, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 10}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '95022', 'n_tokens_mistral': 3512, 'n_tokens_neox': 3397, 'n_words': 869}