<commit_before>import numpy as np
import numexpr as ne
import torch as tc
from collections import OrderedDict
from phy.plot import View, base, visuals
from vispy.util.event import Event
from .color_scheme import palette
from ..utils.utils import Timer, EventEmitter
from ..utils import conf 
from ..utils.conf import error, warning
from ..base.CLU import CLU
from ._core import _get_array, _accumulate
from ._core import _spkNo2maskNo_numba, _cache_out, _cache_in_vector, _cache_in_scalar, _representsInt, _get_box_index


class spike_view(View):
   
    def __init__(self, interactive=True):
        super(spike_view, self).__init__('grid')
        self.palette = palette
        self._transparency = 0.2
        self._highlight_color = np.array([1,0,0,1]).astype('float32')
        self._data_bound = (-2, -1, 2, 1)
        self.interactive = interactive
        self._selected = {}
        self._view_lock = True
        self._magnet_const = 1 # k in knn when emit magnet event to absorb spikes from cluster 0
        self._magnet_mode = False # set magnet_mode to True will lock the magnet_sink_id until magnet_mode is set False
        self._sink_id = 0

        self.event = EventEmitter() 
        # self.events.add(model_modified=Event)
        
    def attach(self, gui):
        gui.add_view(self)

    def _affine_transform(self, x):
        # important!, this map the spk to view space through a affine transformation: y = ax + b
        # self.y = self._affine_transform(self.spk.transpose(2,0,1).ravel())
        # that gives the affine transformation y = ax + b
        a,b = self._get_affine_params(x.min(), x.max())
        y = a*x + b
        return y

    def _get_affine_params(self, min, max):
        '''
            get the affine param which transfer [min, max] to [-1, 1]
        '''
        a = np.ones((2,2), dtype=np.int32)
        a[0,0], a[1,0] = min, max
        b = np.array([-1,1])
        r = np.linalg.solve(a,b)
        return r[0], r[1]

    def _get_data(self, data_bound):
        # new_data_list=[]
        # self._y=[]
        # self._box_=[]
        # self._color_=[]
        i = 0
        # self.box_index   = np.zeros((np.prod(self.spk.shape), 2)).astype(np.float32)
        # self.color_index   = np.zeros((np.prod(self.spk.shape), 4)).astype(np.float32)
        self.box_index   = np.zeros((self.spk.shape[0]*self.spk.shape[2], 2)).astype(np.float32)   # the box index for every pts
        self.color_index = np.zeros((self.spk.shape[0]*self.spk.shape[2], 4)).astype(np.float32)   # the color for every pts
        self._y          = np.zeros((self.spk.shape[0]*self.spk.shape[2], self.spk.shape[1])).astype(np.float32)  # the value for every pts
        for chNo in range(self.n_ch):
            for cluNo in self.clu.index_id:
                s = self.spk[self.clu.index[cluNo],:,chNo]
                if s.ndim == 1:
                    s = s.reshape(1, -1)
                nspk = s.shape[0]
                # self._y.append(np.asarray(s))
                self._y[i:i+nspk, :] = s
                # self._depth[i:i+nspk*s.shape[1], 1] = tc.from_numpy(s.ravel())
                self.box_index[i:i+nspk, :] = np.array([chNo, cluNo])
                self.color_index[i:i+nspk, :3] = np.array(palette[cluNo])
                self.color_index[i:i+nspk,  3] = self._transparency
                i+=nspk
                
                # color_len = nspk*self.n_samples
                # self.color_index[j:j+color_len, :3] = np.array(palette[cluNo])
                # self.color_index[j:j+color_len,  3] = self._transparency
                # j+=color_len

                # transparency= np.tile(self._transparency, (n,1))
                # color       = np.hstack((np.asarray([self.palette[cluNo] for i in s]), transparency)).astype(np.float32)
                # box_index   = np.tile((chNo, cluNo), ((n*self.n_samples,1))).astype(np.float32)
                # self._box_.append(box_index)
                # self._color_.append(color)
                # new_data_list.append({# 'y': y,
                #                       # 'x': x,
                #                       'color': color,
                #                       # 'depth': depth,
                #                       # 'data_bounds': data_bounds,
                #                       'box_index': box_index})
        # self.data = _accumulate(new_data_list)
        # self.data = {}
        # self.data['box_index'] = np.vstack(self._box_)
        # self.data['color']     = np.vstack(self._color_)
        
        self.y = self._y.ravel()

    def _build(self):
        self.grid.shape = (self.n_ch, self.clu.nclu)
        # data = self.data
        # self.box_index = np.vstack(self._box_)
        # self.box_index = np.zeros((np.prod(self.spk.shape), 2)).astype(np.float32)
        # _get_box_index(self.box_index, self.n_ch, self.clu.nclu, self.n_samples, list(self.clu.index.values()))

        # self.depth = np.c_[self.x, self.y, np.zeros(*self.x.shape)].astype(np.float32)

        _depth = tc.zeros((self.x.shape[0], 3)).float()
        _depth[:, 0] = tc.from_numpy(self.x)
        _depth[:, 1] = tc.from_numpy(self.y)
        self.depth = _depth.numpy()
        # self._depth[:, 1] = tc.from_numpy(self.y)
        # self.depth = self._depth.numpy()
        
        
        # self.color = np.repeat(self.color_index, self.n_samples, axis=0)
        n = self.color_index.shape[0]
        color_index = tc.from_numpy(self.color_index)
        color_index.unsqueeze_(-1); # create the third dim, so you can expand on that dim
        self.color = color_index.expand(n, 4, self.n_samples).transpose_(1,2).reshape(-1, 4).numpy()

        # self.box_index = np.repeat(self.box_index, self.n_samples, axis=0)
        n = self.box_index.shape[0]
        box_index = tc.from_numpy(self.box_index)
        box_index.unsqueeze_(-1);
        self.box_index = box_index.expand(n, 2, self.n_samples).transpose_(1,2).reshape(-1, 2).numpy()

        # self.color = np.repeat(np.vstack(self._color_), self.n_samples, axis=0)
        # self.color = self.color_index
        self._cache_depth = self.depth.copy()
        self._cache_color = self.color.copy()
        self._cache_mask_ = np.array([])

    def set_data(self, spk, clu=None):
        #################################
        # this init block take about 1ms
        self.spk = self._affine_transform(spk)
        if clu is None:
            self.clu = CLU(np.zeros(spk.shape[0]).astype('int'))
        else:
            self.clu = clu
        self.n_signals = spk.shape[0]
        self.n_samples = spk.shape[1]
        self.n_ch      = spk.shape[2]
        self._xsig = np.linspace(-0.5, 0.5, self.n_samples)
        self.x = np.tile(self._xsig, self.n_ch*self.n_signals)
        # self._depth = tc.zeros((self.x.shape[0], 3)).float()
        # self._depth[:, 0] = tc.from_numpy(self.x)
        # self.depth = np.zeros((self.x.shape[0], 3)).astype(np.float32)

        #################################
        self.clear()

        ###################################
        # render take 400ms for 5 million points
        self.render()
        
        @self.clu.connect
        def on_cluster(*args, **kwargs):
            with Timer('[VIEW] Spikeview -- rerender', verbose=conf.ENABLE_PROFILER):
                self._selected = {}
                self._magnet_const = 1
                self._magnet_mode = False
                self.rerender()
        
        @self.clu.connect
        def on_select(*args, **kwargs):
            self._selected = self.clu.global2local(self.clu.selectlist)
            self.highlight(self._selected, external=True)
            if self.clu.selectlist.shape[0]>0:
                self.view_lock = True

    def render(self, update=False):

        visual = visuals.PlotVisual()
        self.add_visual(visual)
        
        with Timer('[VIEW] Spikeview -- render - step 0: get data', verbose=conf.ENABLE_PROFILER):
            self._get_data(self._data_bound)

        with Timer('[VIEW] Spikeview -- render - step 1: set data', verbose=conf.ENABLE_PROFILER):
            self._build()
            self.signal_index = np.repeat(np.arange(len(self.y)/len(self._xsig)), 
                                              len(self._xsig)).astype(np.float32)


        with Timer('[VIEW] Spikeview -- render - step 2: gsgl update', verbose=conf.ENABLE_PROFILER):
            visual.program['a_position'] = self.depth
            visual.program['a_color'] = self.color<commit_msg>Fixes #24<commit_after>import numpy as np
import numexpr as ne
import torch as tc
from collections import OrderedDict
from phy.plot import View, base, visuals
from vispy.util.event import Event
from .color_scheme import palette
from ..utils.utils import Timer, EventEmitter
from ..utils import conf 
from ..utils.conf import error, warning
from ..base.CLU import CLU
from ._core import _get_array, _accumulate
from ._core import _spkNo2maskNo_numba, _cache_out, _cache_in_vector, _cache_in_scalar, _representsInt, _get_box_index


class spike_view(View):
   
    def __init__(self, interactive=True):
        super(spike_view, self).__init__('grid')
        self.palette = palette
        self._transparency = 0.2
        self._highlight_color = np.array([1,0,0,1]).astype('float32')
        self._data_bound = (-2, -1, 2, 1)
        self.interactive = interactive
        self._selected = {}
        self._view_lock = True
        self._magnet_const = 1 # k in knn when emit magnet event to absorb spikes from cluster 0
        self._magnet_mode = False # set magnet_mode to True will lock the magnet_sink_id until magnet_mode is set False
        self._sink_id = 0

        self.event = EventEmitter() 
        # self.events.add(model_modified=Event)

        
    def attach(self, gui):
        gui.add_view(self)


    def set_data(self, spk, clu=None):
        #################################
        # this init block take about 1ms
        self.spk = self._affine_transform(spk)
        if clu is None:
            self.clu = CLU(np.zeros(spk.shape[0]).astype('int'))
        else:
            self.clu = clu
        self.n_signals = spk.shape[0]
        self.n_samples = spk.shape[1]
        self.n_ch      = spk.shape[2]
        self._xsig = np.linspace(-0.5, 0.5, self.n_samples)
        self.x = np.tile(self._xsig, self.n_ch*self.n_signals)
        self._depth = tc.zeros((self.x.shape[0], 3)).float()
        self._depth[:, 0] = tc.from_numpy(self.x)        
        self.signal_index = np.repeat(np.arange(len(self.x)/len(self._xsig)), 
                                  len(self._xsig)).astype(np.float32)
        # self._depth = tc.zeros((self.x.shape[0], 3)).float()
        # self._depth[:, 0] = tc.from_numpy(self.x)
        # self.depth = np.zeros((self.x.shape[0], 3)).astype(np.float32)

        #################################
        self.clear()

        ###################################
        # render take 400ms for 5 million points
        self.render()
        
        @self.clu.connect
        def on_cluster(*args, **kwargs):
            with Timer('[VIEW] Spikeview -- rerender', verbose=conf.ENABLE_PROFILER):
                self._selected = {}
                self._magnet_const = 1
                self._magnet_mode = False
                self.rerender()
        
        @self.clu.connect
        def on_select(*args, **kwargs):
            self._selected = self.clu.global2local(self.clu.selectlist)
            self.highlight(self._selected, external=True)
            if self.clu.selectlist.shape[0]>0:
                self.view_lock = True


    def _affine_transform(self, x):
        # important!, this map the spk to view space through a affine transformation: y = ax + b
        # self.y = self._affine_transform(self.spk.transpose(2,0,1).ravel())
        # that gives the affine transformation y = ax + b
        a,b = self._get_affine_params(x.min(), x.max())
        y = a*x + b
        return y

    def _get_affine_params(self, min, max):
        '''
            get the affine param which transfer [min, max] to [-1, 1]
        '''
        a = np.ones((2,2), dtype=np.int32)
        a[0,0], a[1,0] = min, max
        b = np.array([-1,1])
        r = np.linalg.solve(a,b)
        return r[0], r[1]

    def _get_data(self, data_bound):
        # new_data_list=[]
        # self._y=[]
        # self._box_=[]
        # self._color_=[]
        i = 0
        # self.box_index   = np.zeros((np.prod(self.spk.shape), 2)).astype(np.float32)
        # self.color_index   = np.zeros((np.prod(self.spk.shape), 4)).astype(np.float32)
        self.box_index   = np.zeros((self.spk.shape[0]*self.spk.shape[2], 2)).astype(np.float32)   # the box index for every pts
        self.color_index = np.zeros((self.spk.shape[0]*self.spk.shape[2], 4)).astype(np.float32)   # the color for every pts
        self._y          = np.zeros((self.spk.shape[0]*self.spk.shape[2], self.spk.shape[1])).astype(np.float32)  # the value for every pts
        for chNo in range(self.n_ch):
            for cluNo in self.clu.index_id:
                s = self.spk[self.clu.index[cluNo],:,chNo]
                if s.ndim == 1:
                    s = s.reshape(1, -1)
                nspk = s.shape[0]
                # self._y.append(np.asarray(s))
                self._y[i:i+nspk, :] = s
                # self._depth[i:i+nspk*s.shape[1], 1] = tc.from_numpy(s.ravel())
                self.box_index[i:i+nspk, :] = np.array([chNo, cluNo])
                self.color_index[i:i+nspk, :3] = np.array(palette[cluNo])
                self.color_index[i:i+nspk,  3] = self._transparency
                i+=nspk
                
                # color_len = nspk*self.n_samples
                # self.color_index[j:j+color_len, :3] = np.array(palette[cluNo])
                # self.color_index[j:j+color_len,  3] = self._transparency
                # j+=color_len

                # transparency= np.tile(self._transparency, (n,1))
                # color       = np.hstack((np.asarray([self.palette[cluNo] for i in s]), transparency)).astype(np.float32)
                # box_index   = np.tile((chNo, cluNo), ((n*self.n_samples,1))).astype(np.float32)
                # self._box_.append(box_index)
                # self._color_.append(color)
                # new_data_list.append({# 'y': y,
                #                       # 'x': x,
                #                       'color': color,
                #                       # 'depth': depth,
                #                       # 'data_bounds': data_bounds,
                #                       'box_index': box_index})
        # self.data = _accumulate(new_data_list)
        # self.data = {}
        # self.data['box_index'] = np.vstack(self._box_)
        # self.data['color']     = np.vstack(self._color_)
        
        self.y = self._y.ravel()

    def _build(self):
        self.grid.shape = (self.n_ch, self.clu.nclu)
        # data = self.data
        # self.box_index = np.vstack(self._box_)
        # self.box_index = np.zeros((np.prod(self.spk.shape), 2)).astype(np.float32)
        # _get_box_index(self.box_index, self.n_ch, self.clu.nclu, self.n_samples, list(self.clu.index.values()))

        # self.depth = np.c_[self.x, self.y, np.zeros(*self.x.shape)].astype(np.float32)
        self._depth[:, 1] = tc.from_numpy(self.y)
        self._depth[:, 2] = 0
        self.depth = self._depth.numpy()
        # self._depth[:, 1] = tc.from_numpy(self.y)
        # self.depth = self._depth.numpy()
        
        # self.color = np.repeat(self.color_index, self.n_samples, axis=0)
        n = self.color_index.shape[0]
        color_index = tc.from_numpy(self.color_index)
        color_index.unsqueeze_(-1); # create the third dim, so you can expand on that dim
        self.color = color_index.expand(n, 4, self.n_samples).transpose_(1,2).reshape(-1, 4).numpy()

        # self.box_index = np.repeat(self.box_index, self.n_samples, axis=0)
        n = self.box_index.shape[0]
        box_index = tc.from_numpy(self.box_index)
        box_index.unsqueeze_(-1);
        self.box_index = box_index.expand(n, 2, self.n_samples).transpose_(1,2).reshape(-1, 2).numpy()

        # self.color = np.repeat(np.vstack(self._color_), self.n_samples, axis=0)
        # self.color = self.color_index
        self._cache_depth = self.depth.copy()
        self._cache_color = self.color.copy()
        self._cache_mask_ = np.array([])


    def render(self, update=False):
        visual = visuals.PlotVisual()
        self.add_visual(visual)
        
        with Timer('[VIEW] Spikeview -- render - step 0: get data', verbose=conf.ENABLE_PROFILER):
            self._get_data(self._data_bound)

        with Timer('[VIEW] Spikeview -- render - step 1: set data', verbose=conf.ENABLE_PROFILER):
            self._build()

        with Timer('[VIEW] Spikeview -- render - step 2: gsgl update', verbose=conf.ENABLE_PROFILER):
            visual.program['a_position'] = self.depth
            visual.program['a_color'] = self.color
{'fraction_non_alphanumeric': 0.14801974380068164, 'fraction_numerical': 0.015160418380538254, 'mean_word_length': 2.048361096184847, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 13, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4198831', 'n_tokens_mistral': 5982, 'n_tokens_neox': 5725, 'n_words': 1041}