<commit_before>import numpy as np
import os
import parmap
import scipy

from yass.postprocess.util import run_deconv, partition_spike_time

def remove_high_mad(fname_templates,
                    fname_spike_train,
                    fname_weights,
                    reader,
                    neigh_channels,
                    save_dir,
                    units_in=None,
                    multi_processing=False,
                    n_processors=1):

    # output folder
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    # units_in is all units if none
    if units_in is None:
        n_units = np.load(fname_templates).shape[0]
        units_in = np.arange(units_in)

    # partition spike tImes
    spt_save_dir = os.path.join(save_dir, 'spike_times')
    fnames_spike_times = partition_spike_time(
        spt_save_dir, fname_spike_train, units_in)

    # collect save file name
    fnames_out = []
    for unit in units_in:
        fname = os.path.join(save_dir, 'unit_{}.npz'.format(unit))
        fnames_out.append(fname)

    if multi_processing:
        parmap.starmap(find_high_mad_unit, 
                       list(zip(fnames_spike_times, fnames_out)),
                       reader,
                       neigh_channels,
                       processes=n_processors,
                       pm_pbar=True)                
    else:
        for ctr in range(len(units_in)):
            find_high_mad_unit(
                fnames_spike_times[ctr],
                fnames_out[ctr],
                reader,
                neigh_channels)

    collision_units = []
    for ctr, fname in enumerate(fnames_out):
        tmp = np.load(fname)
        if tmp['kill']:
            collision_units.append(units_in[ctr])

    idx_keep = units_in[~np.in1d(units_in, collision_units)]

    return idx_keep


def find_high_mad_unit(
    fname_spike_time,
    fname_out,
    reader,
    neigh_channels,
    min_var=2,
    breach=10,
    up_factor=2,
    min_ptp=2):

    if os.path.exists(fname_out):
        return

    # load spike times
    spt = np.load(fname_spike_time)

    # 1000 is enough for mad calculation
    max_spikes = 1000
    if len(spt) > max_spikes:
        spt = np.random.choice(a=spt,
                               size=max_spikes,
                               replace=False)

    # get waveforms
    wf, _ = reader.read_waveforms(spt)
    _, n_times, n_channels = wf.shape

    # get mean waveform
    mean_wf = np.mean(wf, axis=0)

    # limit to visible channels
    ptps = mean_wf.ptp(0)
    visch = np.where(ptps > min_ptp)[0]
    if len(visch) == 0:
        visch = [ptps.argmax()]
    wf = wf[:, :, visch]
    ptps = ptps[visch]
    neigh_channels = neigh_channels[visch][:, visch]

    # upsample
    wf_up = scipy.signal.resample(wf, n_times*up_factor,
                                  axis=1)

    # max channel within visible channels
    e_var, t_var = get_mad(wf_up, up_factor, ptps.argmax())
    violation_per_channel = np.sum(e_var > (t_var+min_var), axis=0)

    channels_checked = np.zeros(len(visch), 'bool')
    mc = ptps.argmax()
    channels_checked[neigh_channels[mc]] = True
    
    while ((violation_per_channel.sum() > breach) and 
           (not np.all(channels_checked)) and 
           (violation_per_channel[channels_checked].sum() <= breach)
          ):
        
        cands = np.where(~channels_checked)[0]
        channel = cands[ptps[cands].argmax()]
        e_var, t_var = get_mad(wf_up, up_factor, channel)

        violation_temp = np.sum(e_var > (t_var+min_var), axis=0)
        violation_per_channel = np.vstack((
            violation_per_channel,
            violation_temp)).min(0)
        
        channels_checked[neigh_channels[channel]] = True

    if violation_per_channel.sum() > breach:
        kill = True
    else:
        kill = False

    # save result
    np.savez(fname_out,
             kill=kill,
             violation_per_channel=violation_per_channel,
             visch=visch)

def get_mad(wf_up, up_factor, channel):
    
    best_shifts = find_shifts(wf_up[:, :, channel], up_factor)

    # shift
    wf_aligned_up = np.zeros(wf_up.shape, 'float32')
    for j in range(wf_up.shape[0]):
        wf_aligned_up[j] = np.roll(wf_up[j], -best_shifts[j], axis=0)

    # get t_var
    t_var = get_t_var(wf_aligned_up)

    wf_aligned = wf_aligned_up[:,up_factor//2:-up_factor//2]
    wf_aligned = wf_aligned[:,np.arange(0, wf_aligned.shape[1], up_factor)]
    t_var = t_var[up_factor//2:-up_factor//2]
    t_var = t_var[np.arange(0, len(t_var), up_factor)]

    # mad value for aligned waveforms
    e_var = np.square(np.median(np.abs(np.median(wf_aligned, axis=0)[None] - wf_aligned), axis=0)/0.67449)
    
    return e_var, t_var

<commit_msg>limit mad to active area<commit_after>import numpy as np
import os
import parmap
import scipy

from yass.postprocess.util import run_deconv, partition_spike_time

def remove_high_mad(fname_templates,
                    fname_spike_train,
                    fname_weights,
                    reader,
                    neigh_channels,
                    save_dir,
                    units_in=None,
                    multi_processing=False,
                    n_processors=1):

    # output folder
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    # units_in is all units if none
    if units_in is None:
        n_units = np.load(fname_templates).shape[0]
        units_in = np.arange(units_in)

    # partition spike tImes
    spt_save_dir = os.path.join(save_dir, 'spike_times')
    fnames_spike_times = partition_spike_time(
        spt_save_dir, fname_spike_train, units_in)

    # collect save file name
    fnames_out = []
    for unit in units_in:
        fname = os.path.join(save_dir, 'unit_{}.npz'.format(unit))
        fnames_out.append(fname)

    if multi_processing:
        parmap.starmap(find_high_mad_unit, 
                       list(zip(fnames_spike_times, fnames_out)),
                       reader,
                       neigh_channels,
                       processes=n_processors,
                       pm_pbar=True)                
    else:
        for ctr in range(len(units_in)):
            find_high_mad_unit(
                fnames_spike_times[ctr],
                fnames_out[ctr],
                reader,
                neigh_channels)

    collision_units = []
    for ctr, fname in enumerate(fnames_out):
        tmp = np.load(fname)
        if tmp['kill']:
            collision_units.append(units_in[ctr])

    idx_keep = units_in[~np.in1d(units_in, collision_units)]

    return idx_keep


def find_high_mad_unit(
    fname_spike_time,
    fname_out,
    reader,
    neigh_channels,
    min_var=2,
    breach=10,
    up_factor=2,
    min_ptp=2):

    if os.path.exists(fname_out):
        return

    # load spike times
    spt = np.load(fname_spike_time)

    # 1000 is enough for mad calculation
    max_spikes = 1000
    if len(spt) > max_spikes:
        spt = np.random.choice(a=spt,
                               size=max_spikes,
                               replace=False)

    # get waveforms
    wf, _ = reader.read_waveforms(spt)
    _, n_times, n_channels = wf.shape

    # get mean waveform
    mean_wf = np.median(wf, axis=0)

    # limit to visible channels
    ptps = mean_wf.ptp(0)
    visch = np.where(ptps > min_ptp)[0]
    if len(visch) == 0:
        visch = [ptps.argmax()]
    wf = wf[:, :, visch]
    ptps = ptps[visch]
    neigh_channels = neigh_channels[visch][:, visch]

    # upsample
    wf_up = scipy.signal.resample(wf, n_times*up_factor,
                                  axis=1)

    # max channel within visible channels
    e_var, t_var, active_area = get_mad(wf_up, up_factor, ptps.argmax())
    e_var[~active_area] = 0
    violation_per_channel = np.sum(e_var > (t_var + min_var), axis=0)

    channels_checked = np.zeros(len(visch), 'bool')
    mc = ptps.argmax()
    channels_checked[neigh_channels[mc]] = True
    
    while ((violation_per_channel.sum() > breach) and 
           (not np.all(channels_checked)) and 
           (violation_per_channel[channels_checked].sum() <= breach)
          ):
        
        cands = np.where(~channels_checked)[0]
        channel = cands[ptps[cands].argmax()]
        e_var, t_var, active_area = get_mad(wf_up, up_factor, channel)
        e_var[~active_area] = 0

        violation_temp = np.sum(e_var > (t_var + min_var), axis=0)
        violation_per_channel = np.vstack((
            violation_per_channel,
            violation_temp)).min(0)
        
        channels_checked[neigh_channels[channel]] = True

    if violation_per_channel.sum() > breach:
        kill = True
    else:
        kill = False

    # save result
    np.savez(fname_out,
             kill=kill,
             violation_per_channel=violation_per_channel,
             visch=visch)

def get_mad(wf_up, up_factor, channel):
    
    best_shifts = find_shifts(wf_up[:, :, channel], up_factor)

    # shift
    wf_aligned_up = np.zeros(wf_up.shape, 'float32')
    for j in range(wf_up.shape[0]):
        wf_aligned_up[j] = np.roll(wf_up[j], -best_shifts[j], axis=0)

    # get t_var
    t_var = get_t_var(wf_aligned_up)

    wf_aligned = wf_aligned_up[:,up_factor//2:-up_factor//2]
    wf_aligned = wf_aligned[:,np.arange(0, wf_aligned.shape[1], up_factor)]
    t_var = t_var[up_factor//2:-up_factor//2]
    t_var = t_var[np.arange(0, len(t_var), up_factor)]

    active_area = np.abs(np.median(wf_aligned, 0)) > 0.5

    # mad value for aligned waveforms
    e_var = np.square(np.median(np.abs(np.median(wf_aligned, axis=0)[None] - wf_aligned), axis=0)/0.67449)
    
    return e_var, t_var, active_area


{'fraction_non_alphanumeric': 0.10220709570957096, 'fraction_numerical': 0.009591584158415841, 'mean_word_length': 2.146333549643089, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 16, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3922231', 'n_tokens_mistral': 3489, 'n_tokens_neox': 3337, 'n_words': 595}