<commit_before>import logging
import os
import numpy as np
import torch
from tqdm import tqdm

from sklearn.decomposition import PCA
from sklearn.mixture import GaussianMixture
from diptest import diptest as dp

import yass
from yass import read_config
from yass.reader import READER
from yass.residual.residual_gpu import RESIDUAL_GPU2
from yass.deconvolve.utils import shift_svd_denoise
from yass import postprocess

def run_post_deconv_split(output_directory,
                          fname_templates,
                          fname_spike_train,
                          fname_shifts,
                          fname_scales,
                          fname_raw,
                          raw_dtype,
                          fname_residual,
                          residual_dtype,
                          residual_offset=0,
                          initial_batch=False):

    CONFIG = read_config()
    
    # output folder
    if not os.path.exists(output_directory):
        os.makedirs(output_directory)

    # reader
    if CONFIG.deconvolution.deconv_gpu:
        n_sec_chunk = CONFIG.resources.n_sec_chunk_gpu_deconv
    else:
        n_sec_chunk = CONFIG.resources.n_sec_chunk
    reader_residual = READER(fname_residual,
                             residual_dtype,
                             CONFIG,
                             n_sec_chunk,
                             offset=residual_offset)

    reader_raw = READER(fname_raw,
                        raw_dtype,
                        CONFIG)

    # load input data
    templates = np.load(fname_templates)
    spike_train = np.load(fname_spike_train)
    shifts = np.load(fname_shifts)
    scales = np.load(fname_scales)

    # get cleaned ptp
    fname_cleaned_ptp = os.path.join(output_directory, 'cleaned_ptp.npy')
    fname_spike_times = os.path.join(output_directory, 'spike_times.npy')
    fname_shifts = os.path.join(output_directory, 'shifts_list.npy')
    fname_scales = os.path.join(output_directory, 'scales_list.npy')
    fname_vis_chans = os.path.join(output_directory, 'vis_chans.npy')
    if os.path.exists(fname_cleaned_ptp) and os.path.exists(fname_spike_times):
        
        cleaned_ptp = np.load(fname_cleaned_ptp, allow_pickle=True)
        spike_times_list = np.load(fname_spike_times, allow_pickle=True)
        shifts_list = np.load(fname_shifts, allow_pickle=True)
        scales_list = np.load(fname_scales, allow_pickle=True)
        vis_chans = np.load(fname_vis_chans, allow_pickle=True)
    else:
        print('get cleaned ptp')
        (cleaned_ptp, spike_times_list,
         shifts_list, scales_list, vis_chans) = get_cleaned_ptp(
            templates, spike_train, shifts, scales,
            reader_residual, fname_templates, CONFIG)

        np.save(fname_shifts, shifts_list, allow_pickle=True)
        np.save(fname_scales, scales_list, allow_pickle=True)
        np.save(fname_vis_chans, vis_chans, allow_pickle=True)
        np.save(fname_cleaned_ptp, cleaned_ptp, allow_pickle=True)
        np.save(fname_spike_times, spike_times_list, allow_pickle=True)

    # split units
    fname_templates_updated = os.path.join(
        output_directory, 'templated_updated.npy')
    fname_spike_train_updated = os.path.join(
        output_directory, 'spike_train_updated.npy')
    fname_shifts_updated = os.path.join(
        output_directory, 'shifts_updated.npy')
    fname_scales_updated = os.path.join(
        output_directory, 'scales_updated.npy')
    if not(os.path.exists(fname_templates_updated) and 
           os.path.exists(fname_spike_train_updated)):

        print('run split')
        if initial_batch:
            update_original_templates = True
            min_fr_accept = 0
            min_ptp_accept = 0
            min_fraction_accept = 0
        else:
            update_original_templates = False
            min_fr_accept = 1
            min_ptp_accept = 10000
            min_fraction_accept = 0.15
        (templates_updated, spike_train_updated,
         shifts_updated, scales_updated) = run_split(
            cleaned_ptp, spike_times_list,
            shifts_list, scales_list, vis_chans,
            templates, spike_train, reader_raw,<commit_msg>lower template drfit split accept firing rate<commit_after>import logging
import os
import numpy as np
import torch
from tqdm import tqdm

from sklearn.decomposition import PCA
from sklearn.mixture import GaussianMixture
from diptest import diptest as dp

import yass
from yass import read_config
from yass.reader import READER
from yass.residual.residual_gpu import RESIDUAL_GPU2
from yass.deconvolve.utils import shift_svd_denoise
from yass import postprocess

def run_post_deconv_split(output_directory,
                          fname_templates,
                          fname_spike_train,
                          fname_shifts,
                          fname_scales,
                          fname_raw,
                          raw_dtype,
                          fname_residual,
                          residual_dtype,
                          residual_offset=0,
                          initial_batch=False):

    CONFIG = read_config()
    
    # output folder
    if not os.path.exists(output_directory):
        os.makedirs(output_directory)

    # reader
    if CONFIG.deconvolution.deconv_gpu:
        n_sec_chunk = CONFIG.resources.n_sec_chunk_gpu_deconv
    else:
        n_sec_chunk = CONFIG.resources.n_sec_chunk
    reader_residual = READER(fname_residual,
                             residual_dtype,
                             CONFIG,
                             n_sec_chunk,
                             offset=residual_offset)

    reader_raw = READER(fname_raw,
                        raw_dtype,
                        CONFIG)

    # load input data
    templates = np.load(fname_templates)
    spike_train = np.load(fname_spike_train)
    shifts = np.load(fname_shifts)
    scales = np.load(fname_scales)

    # get cleaned ptp
    fname_cleaned_ptp = os.path.join(output_directory, 'cleaned_ptp.npy')
    fname_spike_times = os.path.join(output_directory, 'spike_times.npy')
    fname_shifts = os.path.join(output_directory, 'shifts_list.npy')
    fname_scales = os.path.join(output_directory, 'scales_list.npy')
    fname_vis_chans = os.path.join(output_directory, 'vis_chans.npy')
    if os.path.exists(fname_cleaned_ptp) and os.path.exists(fname_spike_times):
        
        cleaned_ptp = np.load(fname_cleaned_ptp, allow_pickle=True)
        spike_times_list = np.load(fname_spike_times, allow_pickle=True)
        shifts_list = np.load(fname_shifts, allow_pickle=True)
        scales_list = np.load(fname_scales, allow_pickle=True)
        vis_chans = np.load(fname_vis_chans, allow_pickle=True)
    else:
        print('get cleaned ptp')
        (cleaned_ptp, spike_times_list,
         shifts_list, scales_list, vis_chans) = get_cleaned_ptp(
            templates, spike_train, shifts, scales,
            reader_residual, fname_templates, CONFIG)

        np.save(fname_shifts, shifts_list, allow_pickle=True)
        np.save(fname_scales, scales_list, allow_pickle=True)
        np.save(fname_vis_chans, vis_chans, allow_pickle=True)
        np.save(fname_cleaned_ptp, cleaned_ptp, allow_pickle=True)
        np.save(fname_spike_times, spike_times_list, allow_pickle=True)

    # split units
    fname_templates_updated = os.path.join(
        output_directory, 'templated_updated.npy')
    fname_spike_train_updated = os.path.join(
        output_directory, 'spike_train_updated.npy')
    fname_shifts_updated = os.path.join(
        output_directory, 'shifts_updated.npy')
    fname_scales_updated = os.path.join(
        output_directory, 'scales_updated.npy')
    if not(os.path.exists(fname_templates_updated) and 
           os.path.exists(fname_spike_train_updated)):

        print('run split')
        if initial_batch:
            update_original_templates = True
            min_fr_accept = 0
            min_ptp_accept = 0
            min_fraction_accept = 0
        else:
            update_original_templates = False
            min_fr_accept = 0.5
            min_ptp_accept = 10000
            min_fraction_accept = 0.15
        (templates_updated, spike_train_updated,
         shifts_updated, scales_updated) = run_split(
            cleaned_ptp, spike_times_list,
            shifts_list, scales_list, vis_chans,
            templates, spike_train, reader_raw,
{'fraction_non_alphanumeric': 0.06920702754036087, 'fraction_numerical': 0.003442545109211776, 'mean_word_length': 2.348569157392687, '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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5133370', 'n_tokens_mistral': 2732, 'n_tokens_neox': 2629, 'n_words': 453}