<commit_before>import os
import numpy as np
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import tqdm
import parmap
import scipy
import networkx as nx
import logging

from yass import read_config

from statsmodels import robust
from scipy.signal import argrelmin
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA


from yass.deconvolve.match_pursuit import (MatchPursuit_objectiveUpsample, 
                                            Residual)
                                                                                        
from yass.cluster.util import (binary_reader,
                               global_merge_max_dist, PCA, 
                               load_waveforms_from_memory,
                               make_CONFIG2, upsample_parallel, 
                               clean_templates, find_clean_templates,
                               global_merge_max_dist)

from yass.cluster.cluster import Cluster

from yass.deconvolve.merge import TemplateMerge

from diptest import diptest as dp

import multiprocessing as mp

from matplotlib import colors as mcolors
colors = dict(mcolors.BASE_COLORS, **mcolors.CSS4_COLORS)
by_hsv = ((tuple(mcolors.rgb_to_hsv(mcolors.to_rgba(color)[:3])), name)
                 for name, color in colors.items())
sorted_colors = [name for hsv, name in by_hsv]


class Deconv(object):

    def __init__(self, spike_train, templates, output_directory, 
                 recordings_filename, CONFIG):

        self.spike_train = spike_train
        self.templates = templates.swapaxes(1,2).swapaxes(0,1)
        self.output_directory = output_directory
        self.recordings_filename = recordings_filename
        self.CONFIG = make_CONFIG2(CONFIG)
        
        # initialize 
        self.initialize()
        
        # run iterative deconv-clustering on initial chunk - OPTIONAL
        # Cat: TODO: maybe try to use a single deconv() function for both
        if self.n_iterations_deconv>0: 
            self.deconv_iterative()
        
        # run deconv on all data
        self.deconv_final()
        
        # save final output in main directory
        np.save(os.path.join(self.CONFIG.path_to_output_directory,
                            'spike_train_post_deconv_post_merge.npy'),
                    self.spike_train)
        np.save(os.path.join(self.CONFIG.path_to_output_directory,
                            'templates_post_deconv_post_merge.npy'),
                    self.templates)
                    
                
    def initialize(self): 
        
        # location of standardized .bin file
        self.standardized_filename = os.path.join(self.CONFIG.path_to_output_directory,
                                             'preprocess', 
                                             self.recordings_filename)
        self.n_channels = self.CONFIG.recordings.n_channels
        self.sampling_rate = self.CONFIG.recordings.sampling_rate
        self.n_processors = self.CONFIG.resources.n_processors
        self.n_sec_chunk = self.CONFIG.resources.n_sec_chunk
        self.buffer_size = 200
        self.spike_size_default = int(self.CONFIG.recordings.spike_size_ms*
                                      self.CONFIG.recordings.sampling_rate//1000*2+1)
                                      
        # compute segment list for parallel processing
        self.compute_idx_list()

        # Cat: TODO: read from CONFIG
        self.threshold = 10.
        self.conv_approx_rank = 10
        self.default_upsample_value=0
        self.upsample_max_val = 32.

        # Cat: TODO: max iteration not currently used
        self.max_iter = 5000
        
        # set the number of cluster iterations to be done
        # Cat: TODO: read from CONFIG
        self.n_iterations_merge = 1

        # iterative deconv parameters
        # Cat: TODO: read from CONFIG
        self.n_iterations_deconv = 0
        self.n_seconds_initial = 1200
        self.initial_chunk = int(self.n_seconds_initial//self.CONFIG.resources.n_sec_chunk) 

        # make deconv directory
        self.deconv_dir = os.path.join(self.CONFIG.path_to_output_directory,
                                  'deconv')
        if not os.path.isdir(self.deconv_dir):
            os.makedirs(self.deconv_dir)
            
        # if self.n_iterations_deconv>0: 
            # if not os.path.isdir(self.deconv_dir+'/initial/'):
                # os.makedirs(self.deconv_dir+'/initial/')
       
    def compute_idx_list(self):
    
        fp_len = np.memmap(self.standardized_filename, 
                   dtype='float32', mode='r').shape[0]

        # Generate indexes in chunks (units of seconds); usually 10-60sec
        indexes = np.arange(0, fp_len / self.n_channels, 
                            self.sampling_rate * self.n_sec_chunk)
        if indexes[-1] != fp_len/self.n_channels:
            indexes = np.hstack((indexes, fp_len/self.n_channels))

        # Make the 4 parameter list to be sent to deconvolution algorithm
        idx_list = []
        for k in range(len(indexes) - 1):
            idx_list.append([
                indexes[k], indexes[k + 1], self.buffer_size,
                indexes[k + 1] - indexes[k] + self.buffer_size
            ])

        self.idx_list = np.int64(np.vstack(idx_list)) #[:2]


    def deconv_iterative(self):

        self.idx_list_local = self.idx_list[:self.initial_chunk]
        deconv_chunk_dir = os.path.join(self.CONFIG.path_to_output_directory,
                                        'deconv',
                                        'initial/')
        if not os.path.isdir(deconv_chunk_dir):
            os.makedirs(deconv_chunk_dir)
        
        # loop over iterations
        for k in range(self.n_iterations_deconv): 

            # initialize recluster directory
            self.deconv_chunk_dir = os.path.join(deconv_chunk_dir,
                                                str(k))

            if not os.path.isdir(self.deconv_chunk_dir):
                os.makedirs(self.deconv_chunk_dir)

            # run match pursuit and return templates and spike_trains
            self.match_pursuit_function()

            # run post-deconv merge; recomputes self.templates and self.spike_train
            self.merge()

            # reclusters and loads a new self.templates structure
            self.reclustering_function()
            
            # post recluster merge
            _, templates = global_merge_max_dist(
                                        self.deconv_chunk_dir+'/recluster/',
                                        self.CONFIG,
                                        'deconv',
                                        np.arange(self.templates.shape[2]))
            
            # reshape templates
            self.templates = templates.transpose(1,2,0)

            
    def deconv_final(self):

        # Cat: TODO: no need to use chunks anymore, just run all data over single
        #      temporariy fix: run over rest of data in single chunk run:

        # select segments to be processed
        self.idx_list_local = self.idx_list
        self.deconv_chunk_dir = os.path.join(self.CONFIG.path_to_output_directory,
                                        'deconv',
                                        'final/')
        if not os.path.isdir(self.deconv_chunk_dir):
            os.makedirs(self.deconv_chunk_dir)

        # Cat: TODO: don't recomp temp_temp for final step if prev. computed
        self.match_pursuit_function()
        
        # run post-deconv merge
        # self.merge()


    def merge(self): 
                            
        print ("\nPost-deconv merge...")

        # first compute residual
        print ("  computing residual to generate clean spikes")
        self.compute_residual_function()

        # iterate over template-based merge step
        if not os.path.isdir(self.deconv_chunk_dir+'/merge'):
            os.makedirs(self.deconv_chunk_dir+'/merge')
        <commit_msg>turn on post deconv merge<commit_after>import os
import numpy as np
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import tqdm
import parmap
import scipy
import networkx as nx
import logging

from yass import read_config

from statsmodels import robust
from scipy.signal import argrelmin
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA


from yass.deconvolve.match_pursuit import (MatchPursuit_objectiveUpsample, 
                                            Residual)
                                                                                        
from yass.cluster.util import (binary_reader,
                               global_merge_max_dist, PCA, 
                               load_waveforms_from_memory,
                               make_CONFIG2, upsample_parallel, 
                               clean_templates, find_clean_templates,
                               global_merge_max_dist)

from yass.cluster.cluster import Cluster

from yass.deconvolve.merge import TemplateMerge

from diptest import diptest as dp

import multiprocessing as mp

from matplotlib import colors as mcolors
colors = dict(mcolors.BASE_COLORS, **mcolors.CSS4_COLORS)
by_hsv = ((tuple(mcolors.rgb_to_hsv(mcolors.to_rgba(color)[:3])), name)
                 for name, color in colors.items())
sorted_colors = [name for hsv, name in by_hsv]


class Deconv(object):

    def __init__(self, spike_train, templates, output_directory, 
                 recordings_filename, CONFIG):

        self.spike_train = spike_train
        self.templates = templates.swapaxes(1,2).swapaxes(0,1)
        self.output_directory = output_directory
        self.recordings_filename = recordings_filename
        self.CONFIG = make_CONFIG2(CONFIG)
        
        # initialize 
        self.initialize()
        
        # run iterative deconv-clustering on initial chunk - OPTIONAL
        # Cat: TODO: maybe try to use a single deconv() function for both
        if self.n_iterations_deconv>0: 
            self.deconv_iterative()
        
        # run deconv on all data
        self.deconv_final()
        
        # save final output in main directory
        np.save(os.path.join(self.CONFIG.path_to_output_directory,
                            'spike_train_post_deconv_post_merge.npy'),
                    self.spike_train)
        np.save(os.path.join(self.CONFIG.path_to_output_directory,
                            'templates_post_deconv_post_merge.npy'),
                    self.templates)
                    
                
    def initialize(self): 
        
        # location of standardized .bin file
        self.standardized_filename = os.path.join(self.CONFIG.path_to_output_directory,
                                             'preprocess', 
                                             self.recordings_filename)
        self.n_channels = self.CONFIG.recordings.n_channels
        self.sampling_rate = self.CONFIG.recordings.sampling_rate
        self.n_processors = self.CONFIG.resources.n_processors
        self.n_sec_chunk = self.CONFIG.resources.n_sec_chunk
        self.buffer_size = 200
        self.spike_size_default = int(self.CONFIG.recordings.spike_size_ms*
                                      self.CONFIG.recordings.sampling_rate//1000*2+1)
                                      
        # compute segment list for parallel processing
        self.compute_idx_list()

        # Cat: TODO: read from CONFIG
        self.threshold = 10.
        self.conv_approx_rank = 10
        self.default_upsample_value=0
        self.upsample_max_val = 32.

        # Cat: TODO: max iteration not currently used
        self.max_iter = 5000
        
        # set the number of cluster iterations to be done
        # Cat: TODO: read from CONFIG
        self.n_iterations_merge = 1

        # iterative deconv parameters
        # Cat: TODO: read from CONFIG
        self.n_iterations_deconv = 0
        self.n_seconds_initial = 1200
        self.initial_chunk = int(self.n_seconds_initial//self.CONFIG.resources.n_sec_chunk) 

        # make deconv directory
        self.deconv_dir = os.path.join(self.CONFIG.path_to_output_directory,
                                  'deconv')
        if not os.path.isdir(self.deconv_dir):
            os.makedirs(self.deconv_dir)
            
        # if self.n_iterations_deconv>0: 
            # if not os.path.isdir(self.deconv_dir+'/initial/'):
                # os.makedirs(self.deconv_dir+'/initial/')
       
    def compute_idx_list(self):
    
        fp_len = np.memmap(self.standardized_filename, 
                   dtype='float32', mode='r').shape[0]

        # Generate indexes in chunks (units of seconds); usually 10-60sec
        indexes = np.arange(0, fp_len / self.n_channels, 
                            self.sampling_rate * self.n_sec_chunk)
        if indexes[-1] != fp_len/self.n_channels:
            indexes = np.hstack((indexes, fp_len/self.n_channels))

        # Make the 4 parameter list to be sent to deconvolution algorithm
        idx_list = []
        for k in range(len(indexes) - 1):
            idx_list.append([
                indexes[k], indexes[k + 1], self.buffer_size,
                indexes[k + 1] - indexes[k] + self.buffer_size
            ])

        self.idx_list = np.int64(np.vstack(idx_list)) #[:2]


    def deconv_iterative(self):

        self.idx_list_local = self.idx_list[:self.initial_chunk]
        deconv_chunk_dir = os.path.join(self.CONFIG.path_to_output_directory,
                                        'deconv',
                                        'initial/')
        if not os.path.isdir(deconv_chunk_dir):
            os.makedirs(deconv_chunk_dir)
        
        # loop over iterations
        for k in range(self.n_iterations_deconv): 

            # initialize recluster directory
            self.deconv_chunk_dir = os.path.join(deconv_chunk_dir,
                                                str(k))

            if not os.path.isdir(self.deconv_chunk_dir):
                os.makedirs(self.deconv_chunk_dir)

            # run match pursuit and return templates and spike_trains
            self.match_pursuit_function()

            # run post-deconv merge; recomputes self.templates and self.spike_train
            self.merge()

            # reclusters and loads a new self.templates structure
            self.reclustering_function()
            
            # post recluster merge
            _, templates = global_merge_max_dist(
                                        self.deconv_chunk_dir+'/recluster/',
                                        self.CONFIG,
                                        'deconv',
                                        np.arange(self.templates.shape[2]))
            
            # reshape templates
            self.templates = templates.transpose(1,2,0)

            
    def deconv_final(self):

        # Cat: TODO: no need to use chunks anymore, just run all data over single
        #      temporariy fix: run over rest of data in single chunk run:

        # select segments to be processed
        self.idx_list_local = self.idx_list
        self.deconv_chunk_dir = os.path.join(self.CONFIG.path_to_output_directory,
                                        'deconv',
                                        'final/')
        if not os.path.isdir(self.deconv_chunk_dir):
            os.makedirs(self.deconv_chunk_dir)

        # Cat: TODO: don't recomp temp_temp for final step if prev. computed
        self.match_pursuit_function()
        
        # run post-deconv merge
        self.merge()


    def merge(self): 
                            
        print ("\nPost-deconv merge...")

        # first compute residual
        print ("  computing residual to generate clean spikes")
        self.compute_residual_function()

        # iterate over template-based merge step
        if not os.path.isdir(self.deconv_chunk_dir+'/merge'):
            os.makedirs(self.deconv_chunk_dir+'/merge')
        
{'fraction_non_alphanumeric': 0.07328767123287672, 'fraction_numerical': 0.006973848069738481, 'mean_word_length': 1.680407209612817, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 44, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3932145', 'n_tokens_mistral': 4525, 'n_tokens_neox': 4292, 'n_words': 940}