<commit_before># -*- coding: utf-8 -*-

"""Tests of Kwik store items."""

#------------------------------------------------------------------------------
# Imports
#------------------------------------------------------------------------------

import numpy as np
from numpy.testing import assert_array_equal as ae
from pytest import raises

from ....utils.array import _spikes_per_cluster
from ....utils.tempdir import TemporaryDirectory
from ..model import (KwikModel,
                     )
from ..mock import create_mock_kwik
from ...store import ClusterStore
from ..store_items import FeatureMasks, Waveforms, ClusterStatistics


#------------------------------------------------------------------------------
# Tests
#------------------------------------------------------------------------------

_N_CLUSTERS = 5
_N_SPIKES = 100
_N_CHANNELS = 28
_N_FETS = 2
_N_SAMPLES_TRACES = 10000


def test_kwik_store():

    with TemporaryDirectory() as tempdir:
        # Create the test HDF5 file in the temporary directory.
        filename = create_mock_kwik(tempdir,
                                    n_clusters=_N_CLUSTERS,
                                    n_spikes=_N_SPIKES,
                                    n_channels=_N_CHANNELS,
                                    n_features_per_channel=_N_FETS,
                                    n_samples_traces=_N_SAMPLES_TRACES)

        nc = _N_CHANNELS - 2
        nf = _N_FETS

        model = KwikModel(filename)
        spc = _spikes_per_cluster(np.arange(_N_SPIKES), model.spike_clusters)
        clusters = sorted(spc.keys())

        # We initialize the ClusterStore.
        cs = ClusterStore(model=model,
                          path=tempdir,
                          spikes_per_cluster=spc,
                          )

        cs.register_item(FeatureMasks, chunk_size=15)
        cs.register_item(Waveforms,
                         n_spikes_max=5,
                         excerpt_size=2,
                         )
        stats = cs.register_item(ClusterStatistics)

        # We add a custom statistic function.
        def mean_features(cluster):
            fet = cs.features(cluster)
            cs.memory_store.store(cluster, mean_features=fet.mean(axis=0))

        stats.add('mean_features', mean_features)
        cs.register_field('mean_features', 'memory')

        # Now we generate the store.
        cs.generate()

        for cluster in clusters:
            # Check features.
            fet_store = cs.features(cluster)
            fet_expected = model.features[spc[cluster]].reshape((-1, nc, nf))
            ae(fet_store, fet_expected)

            # Check masks.
            masks_store = cs.masks(cluster)
            masks_expected = model.masks[spc[cluster]]
            ae(masks_store, masks_expected)

            # Check waveforms.
            waveforms_store = cs.waveforms(cluster)
            waveforms_spikes = cs.waveforms_spikes(cluster)
            waveforms_expected = model.waveforms[waveforms_spikes]
            ae(waveforms_store, waveforms_expected)

            # Check some statistics.
            ae(cs.mean_masks(cluster),
               masks_expected.mean(axis=0))
            ae(cs.mean_waveforms(cluster),
               waveforms_expected.mean(axis=0))
            ae(cs.mean_features(cluster),
               fet_expected.mean(axis=0))
<commit_msg>Tweaks.<commit_after># -*- coding: utf-8 -*-

"""Tests of Kwik store items."""

#------------------------------------------------------------------------------
# Imports
#------------------------------------------------------------------------------

import numpy as np
from numpy.testing import assert_array_equal as ae

from ....utils.array import _spikes_per_cluster
from ....utils.tempdir import TemporaryDirectory
from ..model import (KwikModel,
                     )
from ...store import ClusterStore
from ..mock import create_mock_kwik
from ..store_items import FeatureMasks, Waveforms, ClusterStatistics


#------------------------------------------------------------------------------
# Tests
#------------------------------------------------------------------------------

_N_CLUSTERS = 5
_N_SPIKES = 100
_N_CHANNELS = 28
_N_FETS = 2
_N_SAMPLES_TRACES = 10000


def test_kwik_store():

    with TemporaryDirectory() as tempdir:
        # Create the test HDF5 file in the temporary directory.
        filename = create_mock_kwik(tempdir,
                                    n_clusters=_N_CLUSTERS,
                                    n_spikes=_N_SPIKES,
                                    n_channels=_N_CHANNELS,
                                    n_features_per_channel=_N_FETS,
                                    n_samples_traces=_N_SAMPLES_TRACES)

        nc = _N_CHANNELS - 2
        nf = _N_FETS

        model = KwikModel(filename)
        spc = _spikes_per_cluster(np.arange(_N_SPIKES), model.spike_clusters)
        clusters = sorted(spc.keys())

        # We initialize the ClusterStore.
        cs = ClusterStore(model=model,
                          path=tempdir,
                          spikes_per_cluster=spc,
                          )

        cs.register_item(FeatureMasks, chunk_size=15)
        cs.register_item(Waveforms,
                         n_spikes_max=5,
                         excerpt_size=2,
                         )
        stats = cs.register_item(ClusterStatistics)

        # We add a custom statistic function.
        def mean_features(cluster):
            fet = cs.features(cluster)
            cs.memory_store.store(cluster, mean_features=fet.mean(axis=0))

        stats.add('mean_features', mean_features)
        cs.register_field('mean_features', 'memory')

        # Now we generate the store.
        cs.generate()

        for cluster in clusters:
            # Check features.
            fet_store = cs.features(cluster)
            fet_expected = model.features[spc[cluster]].reshape((-1, nc, nf))
            ae(fet_store, fet_expected)

            # Check masks.
            masks_store = cs.masks(cluster)
            masks_expected = model.masks[spc[cluster]]
            ae(masks_store, masks_expected)

            # Check waveforms.
            waveforms_store = cs.waveforms(cluster)
            waveforms_spikes = cs.waveforms_spikes(cluster)
            waveforms_expected = model.waveforms[waveforms_spikes]
            ae(waveforms_store, waveforms_expected)

            # Check some statistics.
            ae(cs.mean_masks(cluster),
               masks_expected.mean(axis=0))
            ae(cs.mean_waveforms(cluster),
               waveforms_expected.mean(axis=0))
            ae(cs.mean_features(cluster),
               fet_expected.mean(axis=0))

{'fraction_non_alphanumeric': 0.1677101967799642, 'fraction_numerical': 0.007155635062611807, 'mean_word_length': 1.4115744069015097, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4316628', 'n_tokens_mistral': 1847, 'n_tokens_neox': 1747, 'n_words': 340}