text:
import sys
import os
import numpy as np
import scipy.io
import scipy.sparse
import numba
import random
import multiprocessing as mp
import subprocess
import cytoolz as toolz
import collections
from itertools import chain
import regex as re
import yaml
import logging
import time
import gzip
import pandas as pd
from functools import partial
from typing import NamedTuple
from pysam import AlignmentFile

from .util import compute_edit_distance, read_gene_map_from_gtf
from .fastq_io import read_fastq
from .barcode import ErrorBarcodeHash, ErrorBarcodeHashConstraint
from .estimate_cell_barcode import get_cell_whitelist

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s: %(levelname)s: %(message)s')
logger = logging.getLogger(__name__)

def format_fastq(*fastq, config, method, fastq_out, cb_count,
                 num_thread=4, max_num_cell=1000000):
    """
    Merging fastq reads by putting the cell barcodes and UMI sequences
    to the headers of the cDNA reads

    :param config: the config file
    :param method: the library preparation protocol, e.g., can be one of
    10X, Drop-seq, InDrop, Seq-Well, CEL-seq2, sci-RNA-seq, SPLiT-seq,
    you can add protocol to the configure file easily
    by specifying the read structures.
    A template configuration file is provided in scumi/config.yaml
    :param fastq: input fastq files
    :param fastq_out: the output fastq file
    :param cb_count: an output file containing the # reads for each cell barcode
    :param num_thread: int
                     the number of cpu cores to use
    :param max_num_cell: int
                        the maximum number of cells
    """

    with open(config, 'r') as stream:
        config_dict = yaml.safe_load(stream)

    config_dict = config_dict[method]
    num_read = config_dict['num_read']

    num_fastq = len(fastq)
    if num_fastq != num_read:
        logger.error(f'Error: the number of input fastq files {num_fastq} is different '
                     f'from the number of fastq files {num_read} detected in the config file')
        sys.exit(-1)

    read_regex_str, barcode_filter, read_regex_str_qual = \
        zip(*[_extract_input_read_template('read' + str(i), config_dict)
              for i in range(1, num_read + 1)])

    barcode_filter_dict = dict()
    for d in barcode_filter:
        barcode_filter_dict.update(d)

    read_template = _infer_read_template(read_regex_str)

    # select
    read_regex_list = [re.compile(z) for z in read_regex_str_qual]
    format_read = partial(_format_read, read_regex_list=read_regex_list,
                          read_template=read_template.read_template,
                          cb_tag=read_template.cb_tag,
                          ub_len=read_template.ub_len,
                          barcode_filter_dict=barcode_filter_dict)

    chunk_size = 8000
    fastq_reader = [read_fastq(fastq_i) for fastq_i in fastq]
    chunks = toolz.partition_all(chunk_size, zip(*fastq_reader))

    num_cpu = mp.cpu_count()
    num_thread = num_thread if num_cpu > num_thread else num_cpu
    seq_chunk_obj = toolz.partition_all(num_thread, chunks)

    fastq_out_all = [fastq_out + str(x) + '.gz' for x in range(num_thread)]
    [gzip.open(x, 'wb').close() for x in fastq_out_all]

    cb_count_all = [cb_count + str(x) + '.csv' for x in range(num_thread)]
    [open(x, 'wt').close() for x in cb_count_all]

    fastq_info = collections.defaultdict(collections.Counter)
    iteration = 0
    results = []

    time_start = time.time()
    pool = mp.Pool(num_thread)
    for fastq_chunk in seq_chunk_obj:
        res = pool.starmap_async(format_read, zip(fastq_chunk, fastq_out_all, cb_count_all))
        results.append(res)

        if len(results) == num_thread * 10:
            results[0].wait()

        while results and results[0].ready():
            iteration += 1
            if not (iteration % 10):
                logger.info(f'Processed {iteration * chunk_size * num_thread:,d} reads!')

            res = results.pop(0)
            chunk_info = res.get()
            _update_fastq_info(fastq_info, chunk_info)

    pool.close()
    pool.join()
    for res in results:
        chunk_info = res.get()
        _update_fastq_info(fastq_info, chunk_info)

    with open('.fastq_count.tsv', 'w') as f:
        for k, v in fastq_info['read'].most_common():
            f.write(f'{k}\t{v}\n')

    cmd_cat_fastq = ' '.join(['cat'] + fastq_out_all + ['>'] + [fastq_out])
    try:
        subprocess.check_output(cmd_cat_fastq, shell=True)
        [os.remove(fastq_file) for fastq_file in fastq_out_all]
    except subprocess.CalledProcessError:
        logger.info(f'Errors in concatenate fastq files')
        sys.exit(-1)
    except OSError:
        logger.info(f'Errors in deleting fastq files')
        sys.exit(-1)

    time_used = time.time() - time_start
    logger.info(f'Formatting fastq done, taking {time_used/3600.0:.3f} hours')
    if not cb_count:
        cb_count = fastq_out + '.cb_count'

    df = _count_cell_barcode_umi(cb_count_all[0])
    for cb_file in cb_count_all[1:]:
        df1 = _count_cell_barcode_umi(cb_file)
        df = pd.concat([df, df1], axis=0)
        df = df.groupby(df.index).sum()
        if df.shape[0] > max_num_cell * 2:
            df = df.sort_values(by=df.columns[0], ascending=False)
            df = df.iloc[:max_num_cell, :]

    try:
        [os.remove(cb_file) for cb_file in cb_count_all]
    except OSError:
        logger.info(f'Errors in deleting cell barcode files')
        sys.exit(-1)

    df = df.sort_values(by=df.columns[0], ascending=False)

    if df.shape[0] > 0:
        df.columns = [str(x) for x in range(df.shape[1])]
        df.index.name = 'cb'
        column_name = list(df.columns.values)
        column_name[0] = 'cb_count'
        df.columns = column_name

        df.to_csv(cb_count, sep='\t')

def _update_fastq_info(fastq_info, chunk_info):
    for fastq_count in chunk_info:
        fastq_info['read'].update(read_pass=fastq_count[0],
                                  read_pass_barcode=fastq_count[1],
                                  read_pass_polyt=fastq_count[2],
                                  read_total=fastq_count[3])

def _count_cell_barcode_umi(cb_file, chunk_size=10 ** 7):
    cb_reader = pd.read_csv(cb_file, header=None, iterator=True,
                            sep='\t', index_col=0)

    chunks = cb_reader.get_chunk(chunk_size)
    chunks = chunks.groupby(chunks.index).sum()

    status = True
    while status:
        try:
            chunk = cb_reader.get_chunk(chunk_size)
            chunks = pd.concat([chunks, chunk], axis=0)
            chunks = chunks.groupby(chunks.index).sum()
        except StopIteration:
            status = False
            logger.info('Read cell barcode counts done.')

    return chunks

def _extract_barcode_pos(barcode_dict, config):
    barcode_reg = []
    pos_all = []

    barcode_filter = dict()
    for barcode_and_pos in barcode_dict:
        barcode, pos = barcode_and_pos
        pos_all.append(pos)

        barcode_reg.append('(?P<' + barcode + '>.{' +
                           str(pos[1] - pos[0] + 1) + '})')
        try:
            value = config[barcode + '_value']
            barcode_filter.update({barcode: ErrorBarcodeHash(value, 1)})
        except KeyError:
            pass

    return barcode_reg, pos_all, barcode_filter

def _extract_input_read_template(read, config):
    read_name = '(@.*)\\n'
    read_plus = '(\\+.*)\\n'
    read_qual = '(.*)\\n'

    filter_dict = dict()
    seq = [(key, value) for key, value in config[read].items()
           if key.startswith('cDNA')]
    if seq:
        read_name = '@(?P<name>.*)\\n'
        read_seq = '(?P<seq>.*)\\n'
        read_qual = '(?P<qual>.*)\\n'
        read_template = read_name + read_seq + read_plus + read_qual
        return read_template, filter_dict, read_template

    cell_barcode = [(key, value) for key, value in config[read].items()
                    if key.startswith('CB') and not key.endswith('value')]

    umi = [(key, value) for key, value in config[read].items()
           if key.startswith('UMI')]

    poly_t = [(key, value) for key, value in config[read].items()
              if key.startswith('polyT')]

    cb_reg, cb_pos, cb_filter = _extract_barcode_pos(cell_barcode, config[read])
    filter_dict.update(cb_filter)

    umi_reg, umi_pos, _ = _extract_barcode_pos(umi, config[read])
    umi_reg = [z.replace('UMI', 'UB') for z in umi_reg]

    pt_reg, pt_pos, _ = _extract_barcode_pos(poly_t, config[read])

    read_pos_start = [z[0] for z in cb_pos]
    read_pos_start += [z[0] for z in umi_pos]
    read_pos_start += [z[0] for z in pt_pos]

    read_pos_end = [z[1] for z in cb_pos]
    read_pos_end += [z[1] for z in umi_pos]
    read_pos_end += [z[1] for z in pt_pos]

    idx = sorted(range(len(read_pos_start)),
                 key=lambda k: read_pos_start[k])

    barcode_tag = cb_reg + umi_reg + pt_reg
    read_pos_start = [read_pos_start[i] for i in idx]
    read_pos_end = [read_pos_end[i] for i in idx]
    barcode_tag = [barcode_tag[i] for i in idx]

    idx_skip = [read_pos_start[i+1] - read_pos_end[i] - 1
                for i in range(0, len(read_pos_start)-1)]

    barcode_skip = ['[ACGTN]{' + str(i) + '}' for i in idx_skip]

    read_seq = barcode_tag[0]
    for i in range(len(read_pos_start)-1):
        if idx_skip[i] == 0:
            read_seq += barcode_tag[i+1]
        else:
            read_seq += barcode_skip[i]
            read_seq += barcode_tag[i+1]

    filter_dict.update(_filter_ploy_t(read_seq))

    if read_pos_start[0] > 1:
        read_seq = '[ACGTN]{' + str(read_pos_start[0]-1) + '}'

    read_seq += '[ACGTN]*'
    read_seq = read_seq + '\\n'

    read_template = read_name + read_seq + read_plus + read_qual

    read_qual = re.sub('>', r'_qual>', read_seq)
    read_qual = re.sub('\[ACGTN\]', '.', read_qual)
    read_template_qual = read_name + read_seq + read_plus + read_qual

    return read_template, filter_dict, read_template_qual

def _filter_ploy_t(read_seq):
    match = re.findall('\?P<polyT>\.{[0-9]+}', read_seq)
    poly_t_count = [int(re.findall(r'\d+', z)[0]) for z in match]

    poly_t_filter = {'polyT': ErrorBarcodeHash('T' * z, 1) for z in poly_t_count}

    return poly_t_filter

def _replace_poly_t(read_seq):
    match = re.findall('\?P<polyT>\.{[0-9]+}', read_seq)

    poly_t_count = [int(re.findall(r'\d+', z)[0]) for z in match]
    poly_t = ['(' + 'T'*z + ')' + '{s<=1}' for z in poly_t_count]

    for z in range(len(match)):
        read_seq = read_seq.replace(match[z], poly_t[z])

    return read_seq

def _infer_read_template(reg_list):
    class ReadInfo(NamedTuple):
        cb: bool
        cb_tag: list
        cb_len: list
        ub: bool
        ub_tag: list
        ub_len: list
        read_template: str

    cb = ub = False
    cb_tag = ub_tag = []
    cb_len = ub_len = []

    read_template = '@'
    reg = ''.join(k for k in reg_list)
    if 'CB' in reg:
        logger.info('Cell barcode in configure file')
        cb = True
        cb_seq_template = _accumulate_barcode('CB', reg)
        cb_template = ':CB_' + cb_seq_template[1]
        read_template += cb_template
        cb_tag = cb_seq_template[0]
        cb_len = cb_seq_template[2]

    if 'UB' in reg:
        logger.info('UMI in config file')
        ub = True
        ub_seq_template = _accumulate_barcode('UB', reg)
        ub_template = ':UB_' + ub_seq_template[1]
        read_template += ub_template
        ub_tag = ub_seq_template[0]
        ub_len = ub_seq_template[2]

    read_template += ':{name}'
    read_template += '\n{seq}\n+\n{qual}\n'

    return ReadInfo(cb=cb, cb_tag=cb_tag, cb_len=cb_len,
                    ub=ub, ub_tag=ub_tag, ub_len=ub_len,
                    read_template=read_template)

def _accumulate_barcode(barcode, seq):
    barcode_num = [sub_str[0] for sub_str in
                   seq.split('?P<' + re.escape(barcode))][1:]
    status = '>' in barcode_num

    barcode_num = ['0' if x == '>' else x for x in barcode_num]
    barcode_num = sorted(barcode_num, key=int)

    if status:
        barcode_num[0] = ''
    barcode_seq = [barcode + num for num in barcode_num]

    barcode_template = ['{' + tag + '}' for tag in barcode_seq]
    barcode_template = '-'.join(barcode_template)

    str_split = 'P<' + barcode + '[0-9]*>.{'
    barcode_len = [sub_str for sub_str in re.split(str_split, seq)][1:]
    barcode_len = [int(re.findall(r'(\d+)', barcode_i)[0])
                   for barcode_i in barcode_len]

    return barcode_seq, barcode_template, barcode_len

def _format_read(chunk, fastq_file, cb_count_file, read_regex_list,
                 read_template, cb_tag, ub_len, barcode_filter_dict):
    reads = []
    num_read = len(chunk)
    num_read_pass = num_read_barcode = num_read_polyt = 0
    num_regex = len(read_regex_list)

    barcode_counter = collections.defaultdict(
        partial(np.zeros, shape=(ub_len[0] + 1), dtype=np.uint32))
    ignore_read = False
    for read_i in chunk:
        read_dict_list = []
        for i, regex_i in enumerate(read_regex_list):
            read_match = regex_i.match(read_i[i])
            if not read_match:
                ignore_read = True
                break

            read_dict_list.append(read_match.groupdict())

        if ignore_read:
            ignore_read = False
            continue

        read1_dict = read_dict_list[0]
        if num_regex > 1:
            for regex_id in range(1, num_regex):
                read1_dict.update(read_dict_list[regex_id])

        cb = [barcode_filter_dict[tag][read1_dict[tag]]
              if tag in barcode_filter_dict.keys() else read1_dict[tag]
              for tag in cb_tag]
        if all(cb):
            cb = '-'.join(cb)
            num_read_barcode += 1
        else:
            ignore_read = True

        ub = read1_dict['UB']
        try:
            poly_t = read1_dict['polyT']
            if not barcode_filter_dict['polyT'][poly_t]:
                ignore_read = True
            else:
                num_read_polyt += 1
        except KeyError:
            pass

        if ignore_read:
            ignore_read = False
            continue
        num_read_pass += 1

        if len(read1_dict['seq']) >= 1:
            read1_dict = read_template.format_map(read1_dict)
            reads.append(read1_dict)

        barcode_counter[cb] += [x == 'T' for x in 'T' + ub]

    with gzip.open(fastq_file, 'ab') as fastq_hd:
        for read in reads:
            fastq_hd.write(bytes(read, 'utf8'))

    df = pd.DataFrame.from_dict(barcode_counter, orient='index')
    if df.shape[0] > 0:
        df = df.sort_values(by=df.columns[0], ascending=False)
        df.index.name = 'cb'
        column_name = list(df.columns.values)
        column_name[0] = 'cb_count'
        df.columns = column_name

        df.to_csv(cb_count_file, sep='\t', mode='a', header=False)

    return num_read_pass, num_read_barcode, num_read_polyt, num_read

def _construct_barcode_regex(bam):
    read_mode = 'r' if bam.endswith('.sam') else 'rb'
    bam_file = AlignmentFile(bam, mode=read_mode)
    first_alignment = next(bam_file)
    bam_file.close()

    barcodes = set()
    for barcode in ['CB_', 'UB_']:
        if barcode in first_alignment.qname:
            barcodes.add(barcode)

    barcode_parser = '.*'
    if 'CB_' in barcodes:
        barcode_parser += ':CB_(?P<CB>[A-Z\-]+)'
    if 'UB_' in barcodes:
        barcode_parser += ':UB_(?P<UB>[A-Z\-]+)'

    if barcode_parser == '.*':
        logger.error('Error: no cell barcodes and UMIs.')
        sys.exit(-1)
    barcode_parser += ':*'

    barcode_parser = re.compile(barcode_parser)
    match = barcode_parser.match(first_alignment.qname)

    cb = _extract_tag(match, 'CB')

    return barcode_parser, cb, read_mode

def _extract_tag(match, tag):
    try:
        tag = match.group(tag)
    except IndexError:
        tag = None

    return tag

def count_feature(*cb, bam, molecular_info_h5, gtf, cb_count, feature_tag='XT:Z',
                  expect_cell=False, force_cell=False, all_cell=False,
                  depth_threshold=1, cell_barcode_whitelist=None):
    """
    Count the number of reads/UMIs mapped to each gene

    :param bam: the input sam/bam file
    :param molecular_info_h5: output the molecular info
    :param cb: the input cell barcode files, can be empty or None
    :param cell_barcode_whitelist: a file contain the selected cell barcodes
    :param gtf: a GTF file
    :param cb_count: a file containing the number of reads mapped to each cell barcode,
           output from format_fastq
    :param feature_tag: the tag representing genes in the input bam file
    :param depth_threshold: only considering UMIs that have at least
           depth_threshold reads support
    :param expect_cell: the expected number of cells in the bam file
    :param force_cell: force to return the number of cells set by expect_cell
    :param all_cell: keep all cell barcodes - can be very slow
    """

    barcode_parser, first_cb, read_mode = _construct_barcode_regex(bam)
    num_cb = len(first_cb.split('-'))
    num_cb_file = len(cb)
    if 0 == num_cb_file:
        cb = [None] * num_cb
    elif num_cb != num_cb_file:
        logger.error(f'Error: the number of input cell barcodes files {num_cb_file} '
                     f'is different from the number of cell barcodes {num_cb} '
                     f'detected in the bam file')

        if num_cb > num_cb_file:
            cb = cb + [None] * (num_cb - num_cb_file)
        else:
            cb = cb[:num_cb]

    # TODO: no cell barcodes detected
    correct_cb_fun, cb_list, cb_remove = _construct_cb_filter(
        cb_count, cb, expect_cell, force_cell, all_cell, cell_barcode_whitelist)

    gene_map_dict = read_gene_map_from_gtf(gtf)

    logger.info('Counting molecular info')
    time_start_count = time.time()

    sam_file = AlignmentFile(bam, mode=read_mode)
    _count_feature_partial = partial(_count_feature,
                                     gene_map_dict=gene_map_dict,
                                     barcode_parser=barcode_parser,
                                     correct_cb_fun=correct_cb_fun,
                                     sam_file=sam_file,
                                     feature_tag=feature_tag)

    track = sam_file.fetch(until_eof=True)
    map_info, read_in_cell, molecular_info = _count_feature_partial(track)

    time_count = time.time() - time_start_count
    logger.info(f'Counting molecular info done - {time_count/3600.0:.3f} hours, '
                f'{int(3600.0 * map_info["num_alignment"]/time_count):,d} '
                f'alignments/hour\n')

    # TODO: still output results
    if len(molecular_info) == 0:
        logger.error('Error: no reads mapped to features.')
        sys.exit(-1)

    name = ['cell',
            'gene',
            'umi',
            'depth',
            ]

    logger.info('Converting to a dataframe')
    convert_time = time.time()
    molecular_info = pd.Series(molecular_info).reset_index()
    molecular_info.columns = name

    for col in name[:3]:
        molecular_info.loc[:, col] = molecular_info[col].astype('category')

    convert_time = time.time() - convert_time
    logger.info(f'Converting to a dataframe done, '
                f'taking {convert_time/60.0:.3f} minutes\n')

    molecular_info.columns = name
    if num_cb > 1 and cb_list:
        molecular_info = molecular_info.loc[molecular_info['cell'].isin(cb_list), :]

    if cb_remove:
        molecular_info = molecular_info.loc[~molecular_info['cell'].isin(cb_remove), :]

    molecular_info = molecular_info.loc[molecular_info['depth'] >= 0.95, :]
    molecular_info['depth'] = \
        np.floor(molecular_info['depth'].values + 0.5).astype('uint32')

    molecular_info = molecular_info.sort_values(name[:3])
    molecular_info = molecular_info.reset_index(drop=True)

    map_info = pd.Series(map_info)
    read_in_cell = pd.DataFrame.from_dict(read_in_cell, orient='index')

    logger.info('Writing molecular info')
    write_time = time.time()

    feature = gene_map_dict.values()
    feature = pd.Series(index=set(feature))
    feature = feature.sort_index()

    with pd.HDFStore(molecular_info_h5, mode='w') as hf:
        hf.put('molecular_info', molecular_info, format='table', data_columns=True)
        hf.put('map_info', map_info)
        hf.put('feature', feature)
        hf.put('read_in_cell', read_in_cell)
    del molecular_info

    write_time = time.time() - write_time
    logger.info(f'Writings molecular info done, '
                f'taking {write_time/60.0:.3f} minutes\n')

    _convert_count_to_matrix(molecular_info_h5, molecular_info_h5,
                             depth_threshold=depth_threshold)

def _count_feature(track, gene_map_dict, barcode_parser,
                   correct_cb_fun, sam_file, feature_tag='XT:Z'):
    search_undetermined = re.compile('N').search

    read_name = None
    feature_tag_value_pre = None
    filt_multiple_gene_barcode = False
    count_read = False
    cb_i = feature_tag_value = ub_i = None
    num_aln_read = 0
    pass_filter = False

    map_info = collections.defaultdict(int)
    read_in_cell = collections.Counter()
    molecular_info = collections.defaultdict(int)

    for aln in track:
        if map_info['num_alignment'] and not map_info['num_alignment'] % 10000000:
            logger.info(f'Parsed {map_info["num_alignment"]:,d} alignments.')
            logger.info(f'{map_info["num_unique_read"]:,d} unique reads, '
                        f'{map_info["num_count_read"]:,d} reads kept.')

            logger.info(f'{map_info["num_unmapped_read"]:,d} unmapped reads were filtered.')
            logger.info(f'{map_info["num_barcode_with_na"]:,d} reads '
                        f'were filtered for including NA in barcodes.\n')

        num_aln_read_pre = num_aln_read

        filter_read_unmapped = False
        filter_read_na = False
        filter_read_barcode = False
        map_info['num_alignment'] += 1

        num_aln_read = aln.get_tag('NH')
        new_read = aln.qname != read_name
        if new_read:
            read_name = aln.qname
            if count_read:
                map_info['num_count_read'] += 1
                record_tuple = (cb_i, feature_tag_value, ub_i)
                molecular_info[record_tuple] += 1
            elif pass_filter and (num_aln_read_pre > 1):
                map_info['num_barcode_with_na'] += 1

            pass_filter = False
            count_read = False
            filt_multiple_gene_barcode = True
            feature_tag_value_pre = None

            map_info['num_unique_read'] += 1
            if num_aln_read == 0:
                map_info['num_unmapped_read'] += 1
                filter_read_unmapped = True

            # check cb
            match = barcode_parser.match(aln.qname)
            cb_i = _extract_tag(match, 'CB')
            cb_i_list = cb_i.split('-')
            num_na_in_cb = _count_not_specified(cb_i_list)
            if any(num_na_in_cb > 1) or sum(num_na_in_cb) > len(num_na_in_cb):
                filter_read_na = True

            cb_i = correct_cb_fun(cb_i_list)
            if cb_i:
                read_in_cell[cb_i] += 1
            elif not aln.is_unmapped:
                map_info['num_barcode_with_na'] += 1
                filter_read_barcode = True

            if filter_read_unmapped or filter_read_na or filter_read_barcode:
                count_read = False
                continue

            ub_i = _extract_tag(match, 'UB')
            if ub_i and search_undetermined(ub_i):
                map_info['num_barcode_with_na'] += 1
                continue

            if ub_i and ub_i == len(ub_i) * ub_i[0]:
                map_info['num_barcode_with_na'] += 1
                continue

            try:
                feature_tag_value = aln.get_tag(feature_tag)
            except KeyError:
                feature_tag_value = sam_file.getrname(aln.reference_id)
                if aln.get_tag('XS:Z') == 'Unassigned_Ambiguity':
                    map_info['num_barcode_with_na'] += 1
                    continue

            pass_filter = True
            filt_multiple_gene_barcode = False

            try:
                feature_tag_value = gene_map_dict[feature_tag_value]
            except KeyError:
                if num_aln_read == 1:
                    map_info['num_barcode_with_na'] += 1
                continue

            feature_tag_value_pre = feature_tag_value

            count_read = True
        else:
            if filt_multiple_gene_barcode:
                continue

            try:
                feature_tag_value = aln.get_tag(feature_tag)
            except KeyError:
                feature_tag_value = sam_file.getrname(aln.reference_id)
                if aln.get_tag('XS:Z') == 'Unassigned_Ambiguity':
                    filt_multiple_gene_barcode = True
                    count_read = False
                    continue

            try:
                feature_tag_value = gene_map_dict[feature_tag_value]
            except KeyError:
                feature_tag_value = feature_tag_value_pre
                continue

            if feature_tag_value_pre and feature_tag_value_pre != feature_tag_value:
                filt_multiple_gene_barcode = True
                count_read = False
                continue
            feature_tag_value_pre = feature_tag_value

            count_read = True  # with valid feature_tag_value

    if count_read:
        map_info['num_count_read'] += 1
        record_tuple = (cb_i, feature_tag_value, ub_i)
        molecular_info[record_tuple] += 1

    return map_info, read_in_cell, molecular_info

def _construct_cb_filter(cb_count, cb, expect_cell, force_cell,
                         all_cell, cell_barcode_whitelist):
    cb_list = []
    cb_remove = []
    if all_cell:
        correct_cb_fun = _filter_tag_fun(cb, max_distance=1, correct=True)
    else:
        if cell_barcode_whitelist:
            with open(cell_barcode_whitelist, 'r') as file_handle:
                cb_list = [line.strip('\n') for line in file_handle]
        else:
            cb_list, cb_remove = _get_candidate_barcode(cb_count, cb,
                                                        expect_cell=expect_cell,
                                                        force_cell=force_cell)

        num_cell = len(cb_list)
        logger.info(f'Detected {num_cell:,d} candidate cell barcodes.')
        if num_cell <= 0:
            sys.exit(-1)

        cb_list_split = [cb.split('-') for cb in cb_list]
        cb_df = pd.DataFrame(cb_list_split)

        cb_list_split = [''] * len(cb_df.columns)
        for cb in cb_df:
            cb_list_split[cb] = cb_df[cb].unique()

        if len(cb_df.columns) > 1:
            barcode_hash = _create_barcode_hash(cb_list)
            cb_hash = [ErrorBarcodeHashConstraint(cb, barcode_hash[idx])
                       for idx, cb in enumerate(cb_list_split)]
            correct_cb_fun = partial(_filter_tag_multi, tag_hash=cb_hash, correct=True)
        else:
            cb_hash = [ErrorBarcodeHash(cb) for cb in cb_list_split]
            correct_cb_fun = partial(_filter_tag, tag_hash=cb_hash, correct=True)

    return correct_cb_fun, cb_list, cb_remove

def _count_not_specified(barcode):
    if not barcode:
        return np.array([0])

    na_count = [barcode_i.count('N') for barcode_i in barcode]

    return np.array(na_count)

def _get_candidate_barcode(cb_count_file, cb_file,
                           plot_prefix='.',
                           expect_cell=False,
                           force_cell=False):
    cb = pd.read_csv(cb_count_file, sep='\t')
    cb_name = cb['cb'].str.split('-', expand=True)

    cb_len = [len(z) for z in cb_name.iloc[0, :]]
    num_cb = len(cb_len)

    idx = False
    for cb_idx in range(num_cb):
        filt_cb = [cb_char * cb_len[cb_idx] for cb_char in ['N', 'G']]
        idx = cb_name.iloc[:, 0].isin(filt_cb) | idx
    cb = cb.loc[~idx, :]
    cb = cb.reset_index(drop=True)

    cb_count = dict(zip(cb.cb, cb.cb_count))

    candidate_cb_whitelist = get_cell_whitelist(cb_count,
                                                plot_prefix=plot_prefix,
                                                expect_cell=expect_cell,
                                                force_cell=force_cell)
    candidate_cb_whitelist_refine = \
        _refine_whitelist(list(candidate_cb_whitelist), cb_file)

    merge_barcode = [x != 'None' and x for x in cb_file]
    if any(merge_barcode):
        merge_barcode = False

    cb_list, cb_remove = _merge_del_barcode(candidate_cb_whitelist_refine,
                                            barcode_count=cb, min_distance=1,
                                            merge_barcode=merge_barcode)

    with open('._detected_cb.tsv', 'wt') as file_handle:
        for cb_whitelist in np.setdiff1d(cb_list, cb_remove):
            file_handle.write(f'{cb_whitelist}\n')

    return cb_list, cb_remove

def _refine_whitelist(cb_whitelist, cb_file=None, max_na_per_cb=1):
    cb_hash = []
    if cb_file is not None:
        cb_file = list(cb_file)
        cb_file = [None if cb_i == 'None' else cb_i for cb_i in cb_file]
        cb_hash = _construct_hash(cb_whitelist, cb_file)

    num_cb = len(cb_whitelist[0].split('-'))
    cb_whitelist_corrected = collections.defaultdict(set)
    for cell_barcode in list(cb_whitelist):
        cell_barcode_list = cell_barcode.split('-')
        na_count = _count_not_specified(cell_barcode_list)
        if any(na_count > max_na_per_cb) or sum(na_count) > num_cb:
            continue

        cb = cell_barcode
        if any(cb_hash):
            cb = _correct_cell_barcode(cell_barcode_list, cb_hash)
            if any(cb):
                cb = '-'.join(cb)
            else:
                continue

        cb_whitelist_corrected[cell_barcode].add(cb)

    return cb_whitelist_corrected

def _construct_hash(cb_whitelist, tag_file):
    num_tag = len(tag_file)
    tag_hash = [''] * num_tag

    # Add comments
    for i in range(num_tag):
        if tag_file[i]:
            with open(tag_file[i], 'r') as file_handle:
                tag_i = [line.rstrip('\n') for line in file_handle]

            if len(tag_i) > 5000:
                cell_barcode_list = []
                for cell_barcode in list(cb_whitelist):
                    cell_barcode_list.append(cell_barcode.split('-')[i])

                white_list_map = \
                    _generate_barcode_whitelist_map(cell_barcode_list, tag_i, 1)
                tag_i = [list(v)[0] for k, v in white_list_map.items()]
                tag_i = list(set(tag_i))

            tag_hash[i] = ErrorBarcodeHash(tag_i, edit_distance=1)

    return tag_hash

def _correct_cell_barcode(cell_barcode, cb_hash):
    num_cb = len(cb_hash)

    cb_corrected = cell_barcode
    for i in range(num_cb):
        if cb_hash[i]:
            candidate_cb = cb_hash[i][cell_barcode[i]]
            if candidate_cb:
                cb_corrected[i] = candidate_cb
            else:
                return [None]

    return cb_corrected

def _generate_barcode_whitelist_map(barcode, whitelist, min_distance=1):
    barcode_to_whitelist = collections.defaultdict(set)

    whitelist = set([str(x).encode('utf-8') for x in whitelist])

    num_cpu = mp.cpu_count()
    pool = mp.Pool(num_cpu)

    _partial_map_single_barcode_to_whitelist = \
        partial(_map_single_barcode_to_whitelist, whitelist=whitelist,
                min_distance=min_distance)

    corrected_barcode = pool.map(_partial_map_single_barcode_to_whitelist, barcode)

    for idx, barcode_i in enumerate(corrected_barcode):
        if barcode_i is not None:
            barcode_to_whitelist[barcode[idx]].add(barcode_i)

    return barcode_to_whitelist

def _map_single_barcode_to_whitelist(barcode, whitelist, min_distance=1):
    match = None
    barcode_in_bytes = str(barcode).encode('utf-8')

    for white_barcode in whitelist:
        if barcode_in_bytes in whitelist:
            match = barcode
            break

        if compute_edit_distance(barcode_in_bytes, white_barcode) <= min_distance:
            if match is not None:
                logging.info(f'Warning: barcode {str(barcode)} can be '
                             f'mapped to more than one candidate barcodes')
                match = None
                break
            else:
                match = white_barcode.decode('utf-8')

    return match

def _merge_del_barcode(barcode_dict, barcode_count, min_distance=1, merge_barcode=False):
    barcode_list = list(barcode_dict.keys())

    idx = barcode_count.cb.isin(barcode_list)
    barcode_count_filt = barcode_count.loc[idx, :]

    barcode_corr = [barcode_dict[x] for x in barcode_count_filt.cb]
    idx = [len(x) > 0 for x in barcode_corr]

    barcode_count_filt = barcode_count_filt.iloc[idx, :]

    barcode_corr = list(chain(*barcode_corr))
    barcode_count_filt.cb = barcode_corr

    barcode_count_filt = barcode_count_filt.groupby('cb').sum()

    umi_len = barcode_count_filt.shape[1]
    barcode_count_filt_ratio = barcode_count_filt.iloc[:, 1:umi_len].div(
        barcode_count_filt.cb_count, axis=0)

    idx = barcode_count_filt_ratio.gt(0.80, axis=0)
    idx = idx | barcode_count_filt_ratio.lt(0.005, axis=0)
    count_indel = idx.sum(axis=1)

    if sum(count_indel == 0) <= 100000 and merge_barcode:
        barcode_whitelist = \
            _merge_corrected_barcode(barcode_count_filt.loc[count_indel == 0, :])
    else:
        barcode_whitelist = barcode_count_filt_ratio.index[count_indel == 0].tolist()

    barcode_correctable = barcode_count_filt_ratio.index[count_indel == 1].tolist()

    whitelist_remove = []
    if len(barcode_correctable) > 0:
        barcode_corrected = _correct_del_barcode(
            barcode_count_filt.loc[barcode_correctable, :], min_distance)

        barcode_corrected_list = list(barcode_corrected.keys())
        barcode_corrected_list_mut = [x[:-1]+'N' for x in barcode_corrected_list]
        whitelist_dist = [_find_neighbour_barcode(x, barcode_corrected_list_mut, 1)
                          for x in barcode_whitelist]

        whitelist_remove = [barcode_whitelist[k] for k, v in enumerate(whitelist_dist)
                            if len(v[0]) > 0]

        barcode_whitelist.extend(barcode_corrected_list)

    return barcode_whitelist, whitelist_remove

# N**2 complexity
def _merge_corrected_barcode(barcode_count):
    barcode_count = barcode_count.sort_values('cb_count', ascending=False)

    barcode = barcode_count.index.astype(str)
    barcode_coverage = dict(zip(barcode, barcode_count.cb_count))

    barcode_list = collections.deque()
    barcode_list.append(barcode[0])

    num_barcode = len(barcode)
    if num_barcode <= 1:
        return barcode_list

    for barcode_i in barcode[1:]:
        idx = _find_neighbour_barcode(barcode_i, barcode_list)

        num_neighbour = len(idx[0])
        if num_neighbour == 0:
            barcode_list.append(barcode_i)
            continue
        elif num_neighbour == 1:
            candidate_barcode_idx = idx[0][0]
            candidate_barcode_coverage = \
                barcode_coverage[barcode_list[candidate_barcode_idx]]

            if barcode_coverage[barcode_i] > candidate_barcode_coverage / 10.0:
                barcode_list.append(barcode_i)
                continue

    return barcode_list

def _find_neighbour_barcode(barcode, barcode_list, min_distance=1):
    edit_dist = np.array([compute_edit_distance(barcode.encode('utf-8'), x.encode('utf-8'))
                         for x in barcode_list])
    idx_filt = np.where(edit_dist <= min_distance)

    return idx_filt

def _correct_del_barcode(barcode_count, min_distance=1):
    barcode_count = barcode_count.sort_values('cb_count', ascending=False)
    barcode_all = np.asarray(barcode_count.index.tolist())

    barcode_whitelist = collections.defaultdict(set)
    while len(barcode_all):
        barcode_i = barcode_all[0]
        barcode_whitelist[barcode_i].add(barcode_i)
        barcode_all = barcode_all[1:]

        idx = _find_neighbour_barcode(barcode_i, barcode_all, min_distance)
        if len(idx[0]):
            barcode_ = barcode_all[idx]
            barcode_whitelist[barcode_i].update(list(barcode_))
            barcode_all = np.delete(barcode_all, idx)

    return barcode_whitelist

def _create_barcode_hash(barcode):
    barcode_split = [barcode_i.split('-') for barcode_i in barcode]

    barcode_df = pd.DataFrame(barcode_split, barcode)
    num_barcode = len(barcode_split[0])

    barcode_split_uniq = [''] * len(barcode_df.columns)
    for barcode_i in barcode_df:
        barcode_split_uniq[barcode_i] = barcode_df[barcode_i].unique()

    barcode_hash = [collections.defaultdict(list) for _ in range(num_barcode)]
    for i in range(num_barcode):
        barcode_i = barcode_split_uniq[i]
        for barcode_ii in barcode_i:
            idx = barcode_df[i] == barcode_ii
            barcode_hash[i][barcode_ii] = set(barcode_df.index[idx].tolist())

    return barcode_hash

def convert_count_to_matrix(molecular_info, out_prefix, depth_threshold):
    _convert_count_to_matrix(molecular_info, out_prefix, depth_threshold)

def _convert_count_to_matrix(molecular_info, out_prefix, depth_threshold):
    feature = pd.read_hdf(molecular_info, key='feature')
    molecular_info = pd.read_hdf(molecular_info, key='molecular_info')

    logger.info('Collapsing UMIs')
    write_time = time.time()
    molecular_info = _collapse_umi(molecular_info)
    write_time = time.time() - write_time
    logger.info(f'Collapsing UMIs done, taking {write_time/60.0:.3f} minutes')

    df = _generate_fake_count(molecular_info.iloc[0, :],
                              feature, depth=depth_threshold+1)
    molecular_info = pd.concat([df, molecular_info], ignore_index=True)

    num_gene = len(feature)
    _transform_write_sparse_matrix(molecular_info, num_gene,
                                   sum_type='umi', out_prefix=out_prefix,
                                   depth_threshold=depth_threshold)
    _transform_write_sparse_matrix(molecular_info, num_gene,
                                   sum_type='transcript', out_prefix=out_prefix,
                                   depth_threshold=depth_threshold)

def _transform_write_sparse_matrix(molecular_info, num_gene,
                                   sum_type, out_prefix, depth_threshold):

    logger.info('Converting to sparse matrix')
    convert_time = time.time()

    query_filer = f'depth >= {depth_threshold}'
    if 'umi' == sum_type:
        base_name = out_prefix + '_read'
        count_collapsed = molecular_info.groupby(
            ['cell', 'gene'])
        count_collapsed = count_collapsed['depth'].sum()

        count_collapsed[:num_gene] -= (depth_threshold + 1)
        count_collapsed += 0.5
        count_collapsed = count_collapsed.astype(int)
    else:
        base_name = out_prefix + '_depth_' + str(depth_threshold) + '_transcript'
        count_collapsed = molecular_info.query(query_filer).groupby(
            ['cell', 'gene'])
        count_collapsed = count_collapsed['umi'].size()

        count_collapsed[:num_gene] -= 1

    del molecular_info

    count, count_row_name, count_column_name = _convert_to_coo(count_collapsed)
    del count_collapsed

    convert_time = time.time() - convert_time
    logger.info(f'Converting to a sparse matrix done, '
                f'taking {convert_time/60.0:.3f} minutes')

    logger.info('Output results')
    write_time = time.time()
    pd.Series(count_row_name).to_csv(base_name + '_gene.tsv',
                                     index=False, header=False)
    pd.Series(count_column_name).to_csv(base_name + '_barcode.tsv',
                                        index=False, header=False)

    if 'umi' == sum_type:
        with open(base_name + '.mtx', 'w+b') as out_handle:
            scipy.io.mmwrite(out_handle, count)
    else:
        with open(base_name + '.mtx', 'w+b') as out_handle:
            scipy.io.mmwrite(out_handle, count)

    write_time = time.time() - write_time
    logger.info(f'Writing final results done, '
                f'taking {write_time/60.0:.3f} minutes')

def _map_barcode_to_whitelist(barcode, whitelist, min_distance=1):
    whitelist = set([str(x).encode('utf-8') for x in whitelist])

    iter_i = 0
    for barcode_i in barcode:
        match = barcode_i
        barcode_in_bytes = str(barcode_i).encode('utf-8')

        for white_barcode in whitelist:
            if barcode_in_bytes in whitelist:
                break

            if compute_edit_distance(barcode_in_bytes, white_barcode) <= min_distance:
                match = white_barcode.decode('utf-8')
                break

        barcode[iter_i] = match
        iter_i += 1

    return barcode

def _collapse_barcode_edit(barcode, value, min_distance=1):
    id_srt = value.argsort()[::-1]

    barcode = barcode[id_srt]
    value = value[id_srt]

    max_barcode = value[0]
    threshold = max_barcode * 0.1
    if threshold < 2:
        threshold = 2
    elif threshold > 5:
        threshold = 5
    id_whitelist = value > threshold

    whitelist_candidate = barcode[id_whitelist]
    noise_candidate = barcode[~id_whitelist]

    if len(noise_candidate) > 0 and len(whitelist_candidate) > 0:
        corrected_noise = _map_barcode_to_whitelist(noise_candidate,
                                                    whitelist_candidate,
                                                    min_distance)

        barcode[~id_whitelist] = corrected_noise

    return barcode, value

def _collapse_umi(x, min_distance=1):
    id_start = x.duplicated(['cell', 'gene'])
    id_start = id_start[id_start == False].index.tolist()

    id_end = id_start[1:]
    id_end.append(x.shape[0])

    value = x['depth'].values
    umi = x['umi'].values.astype('str')

    for gene in np.arange(len(id_end)):
        id_gene = np.arange(id_start[gene], id_end[gene])

        if len(id_gene) <= 1:
            continue

        umi_gene = umi[id_gene]
        value_gene = value[id_gene]

        umi_gene, _ = _collapse_barcode_edit(umi_gene, value_gene, min_distance)
        umi[id_gene] = umi_gene

    x['umi'] = pd.Categorical(umi)
    x = x.groupby(['cell', 'gene', 'umi'], observed=True)['depth'].sum()
    x = x.reset_index(drop=False)

    return x

def _convert_to_coo(data_series):
    data_sp = data_series.astype('Sparse')

    data_sp, row_name, column_name = data_sp.sparse.to_coo(
        column_levels=['cell'],
        row_levels=['gene']
    )
    data_sp.eliminate_zeros()
    data_sp = data_sp.astype(int)

    coo_tuple = collections.namedtuple('coo_tuple', ['x', 'row_name', 'column_name'])

    return coo_tuple(data_sp, row_name, column_name)

def _generate_fake_count(row_of_df, feature, depth=1.5):
    index_name = row_of_df.index
    df = pd.DataFrame(columns=index_name)
    df['gene'] = feature.index.values
    df['umi'] = 'N' * len(row_of_df['umi'])
    df['depth'] = depth

    for index_name_other in list(set(index_name) - {'gene', 'umi', 'depth'}):
        df[index_name_other] = row_of_df[index_name_other]

    return df

def down_sample(molecular_info, total_read=None, total_cell=None, mean_read=None,
                out_prefix='.', depth_threshold=1, seed=0):
    """
    Down-sampling the molecular_info such that each library has the same number of reads

    :param molecular_info: molecular_info: the input molecular info data frame
    :param total_read: the total number of reads for these libraries
    :param total_cell: the total number of cells
    :param mean_read: the expected number of reads per cell after down-sampling
    :param out_prefix: the prefix of the output matrices
    :param depth_threshold: the coverage threshold to consider
    :param seed used for random sampling
    """

    feature = pd.read_hdf(molecular_info, key='feature')
    output_molecular_info = pd.read_hdf(molecular_info, key='molecular_info')
    for col in output_molecular_info.columns[:3]:
        output_molecular_info.loc[:, col] = output_molecular_info[col].astype('category')

    output_molecular_info = output_molecular_info.loc[
                            output_molecular_info['depth'] >= 1, :]
    output_molecular_info.reset_index(drop=True, inplace=True)

    map_info = pd.read_hdf(molecular_info, key='map_info')
    if total_read is None:
        total_read = map_info['num_unique_read']
    else:
        total_read = max(total_read, map_info['num_unique_read'])

    read_in_cell = pd.read_hdf(molecular_info, key='read_in_cell')
    if total_cell is None:
        total_cell = read_in_cell.shape[0]
    else:
        total_cell = min(total_cell, read_in_cell.shape[0])

    if mean_read is None:
        mean_read = (10000, )
    elif not isinstance(mean_read, tuple):
        mean_read = (mean_read, )

    cell_vec = output_molecular_info['depth'].copy()
    for mean_read_i in mean_read:
        random.seed(seed)
        seed += 1

        _down_sample(output_molecular_info, feature, total_read, total_cell,
                     mean_read_i, out_prefix, depth_threshold=depth_threshold)
        output_molecular_info['depth'] = cell_vec

def _down_sample(molecular_info, feature, total_read, total_cell, mean_read,
                 out_prefix, depth_threshold=1):
    expect_read = mean_read * total_cell
    if expect_read > total_read:
        return ()

    cell_vec = molecular_info['depth']
    value = cell_vec.tolist()
    id_end = np.cumsum(value)
    id_start = np.append(0, id_end)
    id_start = id_start[:-1]

    read_num_subsample = (id_end[-1] / total_read * 1.0) * expect_read
    read_num_subsample = int(read_num_subsample + 0.5)

    id_keep = sorted(random.sample(range(id_end[-1]), read_num_subsample))
    expanded_count = np.zeros(id_end[-1], dtype=np.int32)
    expanded_count[id_keep] = 1

    value = _add_umis(expanded_count, id_start, id_end)
    molecular_info['depth'] = value
    output_molecular_info = molecular_info.loc[molecular_info['depth'] >= 1, :].copy()
    output_molecular_info.reset_index(drop=True, inplace=True)

    logger.info('Collapsing UMIs')
    write_time = time.time()
    output_molecular_info = _collapse_umi(output_molecular_info)
    write_time = time.time() - write_time
    logger.info(f'Collapsing UMIs done, taking {write_time / 60.0:.3f} minutes')

    out_prefix = out_prefix + '_sample_' + str(mean_read)
    _calculate_cell_gene_matrix(output_molecular_info, feature,
                                out_prefix=out_prefix,
                                depth_threshold=depth_threshold)

def down_sample_cell(molecular_info, expect_read=None, out_prefix='', depth_threshold=1):
    """
    Down-sampling the molecular_info such that each cell has the same number of reads

    :param molecular_info: the input molecular info data frame
    :param expect_read: each cell to have expect_read
    :param out_prefix: the prefix of the output matrices
    :param depth_threshold: the coverage threshold to consider
    """

    feature = pd.read_hdf(molecular_info, key='feature')
    output_molecular_info = pd.read_hdf(molecular_info, key='molecular_info')

    for col in output_molecular_info.columns[:3]:
        output_molecular_info.loc[:, col] = output_molecular_info[col].astype('category')

    output_molecular_info = output_molecular_info.loc[
                            output_molecular_info['depth'] >= 1, :]
    name = output_molecular_info.columns.tolist()
    name = name[:-1]
    output_molecular_info = output_molecular_info.sort_values(name)  # already sorted?

    read_in_cell = pd.read_hdf(molecular_info, key='read_in_cell')

    if expect_read is None:
        expect_read = (10000, )
    elif not isinstance(expect_read, tuple):
        expect_read = (expect_read, )
    for num_read in expect_read:
        _down_sample_cell(output_molecular_info, feature, read_in_cell, num_read,
                          out_prefix, depth_threshold=depth_threshold)

def _down_sample_cell(molecular_info, feature, read_in_cell, expect_read,
                      out_prefix, depth_threshold=1):

    read_in_cell = read_in_cell[read_in_cell[0] >= expect_read]

    num_cell = read_in_cell.shape[0]
    if num_cell == 0:
        logger.error(f'All cells have less than {expect_read:,d} reads, aborting.')
        sys.exit(-1)

    cell = set(read_in_cell.index.tolist())
    output_molecular_info = molecular_info.loc[
        molecular_info['cell'].isin(cell)].copy()

    # Update the values of output_molecular_info
    output_molecular_info.reset_index(drop=True, inplace=True)

    cell_start = output_molecular_info.drop_duplicates('cell')
    cell = cell_start['cell'].tolist()
    cell_start = cell_start.index.tolist()

    cell_end = cell_start[1:]
    cell_end.append(output_molecular_info.shape[0])

    read_in_cell = read_in_cell.to_dict('dict')[0]

    umi_count = output_molecular_info['depth'].as_matrix()
    umi_count = umi_count.reshape([umi_count.shape[0], 1])
    random.seed(0)

    for idx_i, cell_i in enumerate(cell):
        logger.info(f'Subsample cell {idx_i:,d}')

        id_case, id_end = _find_pos(cell_start[idx_i], cell_end[idx_i], umi_count)
        id_start = np.append(0, id_end)
        id_start = id_start[:-1]

        num_umi = int(expect_read * id_end[-1] / read_in_cell[cell_i] + 0.5)
        if id_end[-1] < num_umi:
            continue

        id_keep = sorted(random.sample(range(id_end[-1]), num_umi))
        expanded_count = np.zeros(id_end[-1], dtype=np.int32)
        expanded_count[id_keep] = 1

        value = _add_umis(expanded_count, id_start, id_end)
        umi_count[id_case] = value

    logger.info('Collapsing UMIs')
    write_time = time.time()
    output_molecular_info = _collapse_umi(output_molecular_info)
    write_time = time.time() - write_time
    logger.info(f'Collapsing UMIs done, taking {write_time/60.0:.3f} minutes')

    out_prefix = out_prefix + '.sample_' + str(expect_read)
    _calculate_cell_gene_matrix(output_molecular_info, feature, out_prefix=out_prefix,
                                depth_threshold=depth_threshold)

@numba.jit(nopython=True)
def _add_umis(x, id_start, id_end):
    y = np.zeros((len(id_start), 1), dtype=np.uint32)
    for i in range(len(id_start)):
        y[i] = np.sum(x[id_start[i]:id_end[i]])

    return y

@numba.jit(nopython=True)
def _find_pos(cell_start_index, cell_end_index, umi_count):
    id_case = np.arange(cell_start_index, cell_end_index)

    value = umi_count[id_case]
    id_end = np.cumsum(value)

    return id_case, id_end

def _calculate_cell_gene_matrix(molecular_info, feature, out_prefix, depth_threshold):
    df = _generate_fake_count(molecular_info.iloc[0, :],
                              feature, depth=depth_threshold+1)
    molecular_info = pd.concat([df, molecular_info], ignore_index=True)

    num_gene = len(feature)
    _transform_write_sparse_matrix(molecular_info, num_gene,
                                   sum_type='transcript', out_prefix=out_prefix,
                                   depth_threshold=depth_threshold)

    _transform_write_sparse_matrix(molecular_info, num_gene,
                                   sum_type='umi', out_prefix=out_prefix,
                                   depth_threshold=depth_threshold)

def _filter_tag_fun(tag_file, max_distance, correct=True):
    if tag_file is not None:
        tag_file = list(tag_file)
        tag_file = [None if x == 'None' else x for x in tag_file]
    else:
        tag_file = [None]

    tag_hash = [None] * len(tag_file)
    for i, tag_file_i in enumerate(tag_file):
        if not tag_file_i:
            continue

        with open(tag_file_i, 'r') as file_handle:
            tag_seq = {line.rstrip('\n') for line in file_handle}
        tag_hash[i] = ErrorBarcodeHash(tag_seq, max_distance)

    filter_tag_fun = partial(_filter_tag, tag_hash=tag_hash, correct=correct)

    return filter_tag_fun

def _filter_tag(tag, tag_hash, correct=True):
    tag_corrected = list(tag)

    for i, tag_i in enumerate(tag):
        if not tag_hash[i]:
            continue

        tag_corrected[i] = tag_hash[i][tag_i]
        if not tag_corrected[i]:
            return None

    if correct:
        return '-'.join(tag_corrected)
    else:
        return '-'.join(tag)

def _filter_tag_multi(tag, tag_hash, correct=True):
    tag_corrected = list(tag)

    tag_full = '-'.join(tag)
    for i, tag_i in enumerate(tag):
        if not tag_hash[i]:
            continue

        tag_corrected[i] = tag_hash[i][tag_i, tag_full]
        if not tag_corrected[i]:
            return None

    if correct:
        return '-'.join(tag_corrected)
    else:
        return '-'.join(tag)

def annotate_bam(bam, gtf, featureCounts='featureCounts',
                 annotate_multi_mapping=True, annotate_intron=False,
                 strand=1, num_thread=4):
    import subprocess

    logger.info('Running featureCounts to annotate alignments.')
    start_time = time.time()

    cmd_feature_count = ' '.join([featureCounts,
                                  '-g gene_id', '-t exon', '-R BAM',
                                  '-T', str(num_thread),
                                  '-F GTF', '-a', gtf, '-s', str(strand),
                                  ])

    if annotate_multi_mapping:
        cmd_feature_count = ' '.join([cmd_feature_count, '-M'])

    try:
        cmd = ' '.join([cmd_feature_count, '-o', bam + '.annotation', bam])
        print(cmd, '\n')

        subprocess.check_output(cmd, shell=True)
    except subprocess.CalledProcessError:
        logger.info(f'Running featureCount errors')
        sys.exit(-1)
    except OSError:
        logger.info(f'featureCounts not found or cannot run!'
                    f' Please specify which featureCounts to use.')
        sys.exit(-1)

    if annotate_intron:
        cmd_feature_count = re.sub('-t exon', '-t transcript', cmd_feature_count)
        cmd = ' '.join([cmd_feature_count, '-o',
                        bam + '.annotation.intron',
                        bam + '.featureCounts.bam'])
        try:
            subprocess.check_output(cmd, shell=True)

            cmd = ' '.join(['mv', bam + '.featureCounts.bam' + '.featureCounts.bam',
                            bam + '.featureCounts.bam'])
            subprocess.check_output(cmd, shell=True)
        except subprocess.CalledProcessError:
            logger.info(f'Running featureCount errors')
            sys.exit(-1)
        except OSError:
            logger.info(f'featureCounts not found or cannot run! '
                        f'Please specify which featureCounts to use.')
            sys.exit(-1)

    feature_count_time = time.time() - start_time
    logger.info(f'Annotating features done successfully, '
                f'taking {feature_count_time/60.0:.3f} minutes')

reason: import numpy,import scipy,import numba


text:
# data pre-processing

# Import Library
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import os

# import data
dataset = pd.read_csv('Salary_Data.csv')
X = dataset.iloc[:, :-1].values # independent variable
Y = dataset.iloc[:, 1].values   # dependent variable

# Split data into training / testing dataset
from sklearn.cross_validation import train_test_split
X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size = 0.3, random_state = 0)

# Feature Scaling
"""from sklearn.preprocessing import StandardScaler
sc_X = StandardScaler()
X_train = sc_X.fit_transform(X_train)
X_test = sc_X.transform(X_test)
print X_train"""

# Fitting simple linear regression model to the Training Set
from sklearn.linear_model import LinearRegression
regressor = LinearRegression()
regressor.fit(X_train, Y_train)

# Predict the test set results
Y_predict = regressor.predict(X_test)

# Visulization the Train Set results
plt.scatter(X_train, Y_train, color = 'red')
plt.plot(X_train, regressor.predict(X_train), color = 'blue')
plt.title(' Salary V.S. Experience (Training Set)')
plt.xlabel('Years of experience')
plt.ylabel('Salary')
plt.show()

# Visulization the Test Set results
plt.scatter(X_test, Y_test, color = 'red')
plt.plot(X_train, regressor.predict(X_train), color = 'blue')
plt.title(' Salary V.S. Experience (Test Set)')
plt.xlabel('Years of experience')
plt.ylabel('Salary')
plt.show()

ext: py

max_stars_repo_name: kevingo/ml-az

max_issues_repo_name: kevingo/ml-az

max_forks_repo_name: kevingo/ml-az

reason: import numpy


text:
#!/usr/bin/env python

from distutils.core import setup, Extension
import numpy as np
from os import path

this_directory = path.abspath(path.dirname(__file__))
with open(path.join(this_directory, 'README.md'), encoding='utf-8') as f:
    long_description = f.read()

def parse_requirements(filename):
        """Load requirements from a pip requirements file."""
        lineiter = (line.strip() for line in open(filename))
        return [line for line in lineiter if line and not line.startswith("#")]

c_adt = Extension('adt', sources=['anisotropic-distance-transform/ext/adt.c'], include_dirs=[np.get_include()])

setup(name='anisotropic-distance-transform',
      version='0.1.7',
      description='Anisotropic Euclidean Distance Transform (2D)',
      long_description=long_description,
      long_description_content_type="text/markdown",
      ext_modules=[c_adt],
      author="Sebastian Kosch",
      author_email="skosch@users.noreply.github.com",
      url="https://github.com/skosch/anisotropic-distance-transform",
      license="MIT",
      classifiers=[
          "License :: OSI Approved :: MIT License",
          "Programming Language :: Python :: 3",
          "Programming Language :: Python :: 3.7",
      ],
      packages=["anisotropic-distance-transform"],
)

max_stars_repo_name: skosch/anisotropic-distance-transform

max_issues_repo_name: skosch/anisotropic-distance-transform

max_forks_repo_name: skosch/anisotropic-distance-transform


text:
"""
SuironCV contains functions that does some preprocessing on the images
before it is fed into the feed forward network
""""
import math
import cv2
import numpy as np

# Median blur
def get_median_blur(gray_frame):
    return cv2.medianBlur(gray_frame, 5)

# Canny edge detection
def get_canny(gray_frame):
    return cv2.Canny(gray_frame, 50, 200, apertureSize=3)

# Hough lines
def get_lane_lines(inframe):
    frame = inframe.copy()
    ret_frame = np.zeros(frame.shape, np.uint8)

    # We converted it into RGB when we normalized it
    gray = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)

    gray = get_median_blur(gray)
    canny = get_canny(gray)

    # Hough lines
    # threshold = number of 'votes' before hough algorithm considers it a line
    lines = cv2.HoughLinesP(canny, 1, np.pi/180, threshold=25, minLineLength=40, maxLineGap=100)

    try:
        r = lines.shape[0]
    except AttributeError:
        r = 0

    for i in range(0):
        for x1, y1, x2, y2 in lines[i]:
            # Degrees as its easier for me to conceptualize
            angle = math.atan2(y1-y2, x1-x2)*180/np.pi

            # If it looks like a left or right lane
            # Draw it onto the new image
            if 100 < angle < 170 or -170 < angle < -100:
                cv2.line(ret_frame, (x1, y1), (x2, y2), (255, 255, 255), 10)

    return ret_frame


text:
#!/usr/bin/env python

"""
Testing of code to find nodes.

Currently only nearest neighbor.

"""

from __future__ import (absolute_import, division, print_function)

import numpy as np

from .utilities import twenty_one_triangles

def test_locate_node(twenty_one_triangles):
    """Test finding a single node."""

    ugrid = twenty_one_triangles
    assert ugrid.locate_nodes((4.58, 5.08)) == 6

def test_locate_nodes(twenty_one_triangles):
    """Test finding multiple nodes at once."""

    ugrid = twenty_one_triangles
    assert np.array_equal(ugrid.locate_nodes(((4.58, 5.08),
                                              (4.81, 0.89),
                                              (6.43, 12.9),
                                              (8.74, 6.86),
                                              (5.12, 7.31),
                                              )), (6, 0, 17, 10, 8))

def test_locate_exact(twenty_one_triangles):
    """
    The nearest neighbor of the exact node locations had better be
    the nodes!

    """
    ugrid = twenty_one_triangles
    assert np.array_equal(ugrid.locate_nodes(ugrid.nodes),
                          list(range(len(ugrid.nodes))))

def test_locate_middle(twenty_one_triangles):
    """See what happens the point is equidistant to two nodes."""
    ugrid = twenty_one_triangles

    # (3,5) is equidistant between nodes 2, 6 and 7
    # 2 is returned, but might be arbitrary
    # assert ugrid.locate_nodes( (3, 5) ) == 2

    # Perturb the point a bit, and nearest changes to:
    assert ugrid.locate_nodes((3.0000000001, 5)) == 6
    assert ugrid.locate_nodes((3, 5.00000000001)) == 7
    assert ugrid.locate_nodes((3, 4.99999999999)) == 2

if __name__ == '__main__':
    test_locate_nodes()

max_stars_repo_name: groutr/gridded

max_issues_repo_name: groutr/gridded

max_forks_repo_name: groutr/gridded


text:
## Functions needed to train and test DMM

import numpy as np
import torch
from torch import nn, tensor

def process_minibatch(svi, train_data, epoch, which_mini_batch, shuffled_indices,
                      mini_batch_size, annealing_epochs, minimum_annealing_factor,
                      use_cuda=False):

    training_seq_lengths = tensor(np.array([len(train_data[i]) for i in range(len(train_data))]))
    training_data_sequences = tensor(np.array([train_data[i].values for i in range(len(train_data))]))
    N_train_data = len(training_seq_lengths)
    N_mini_batches = int(N_train_data / 10 + int(N_train_data % 10 > 0))

    if annealing_epochs > 0 and epoch < annealing_epochs:
        # compute the KL annealing factor appropriate
        # for the current mini-batch in the current epoch
        min_af = minimum_annealing_factor
        annealing_factor = min_af + (1.0 - min_af) * \
            (float(which_mini_batch + epoch * N_mini_batches + 1) /
             float(annealing_epochs * N_mini_batches))
    else:
        # by default the KL annealing factor is unity
        annealing_factor = 1.0

    # compute which sequences in the training set we should grab
    mini_batch_start = (which_mini_batch * mini_batch_size)
    mini_batch_end = np.min([(which_mini_batch + 1) * mini_batch_size,
                             N_train_data])
    mini_batch_indices = shuffled_indices[mini_batch_start:mini_batch_end]
    # grab the fully prepped mini-batch using the helper function in the data loader
    mini_batch, mini_batch_reversed, mini_batch_mask, mini_batch_seq_lengths \
        = get_mini_batch(mini_batch_indices, training_data_sequences,
                         training_seq_lengths, cuda=use_cuda)
    mini_batch = mini_batch.float()
    mini_batch_reversed = mini_batch_reversed.float()
    mini_batch_mask = mini_batch_mask.float()

    # do an actual gradient step
    loss = svi.step(mini_batch, mini_batch_reversed, mini_batch_mask,
                    mini_batch_seq_lengths, annealing_factor)
    # keep track of the training loss
    return loss

def do_evaluation(dmm, svi, val_data, test_data, use_cuda):

    # Extract parameters form val and test data
    test_seq_lengths = tensor(np.array([len(val_data[i]) for i in range(len(val_data))]))
    test_data_sequences = tensor(np.array([val_data[i].values for i in range(len(val_data))]))
    val_seq_lengths = tensor(np.array([len(test_data[i]) for i in range(len(test_data))]))
    val_data_sequences = tensor(np.array([test_data[i].values for i in range(len(test_data))]))

    # package repeated copies of val/test data for faster evaluation
    # (i.e. set us up for vectorization)
    n_eval_samples = 1
    def rep(x):
        return np.repeat(x, n_eval_samples, axis=0)

    # get the validation/test data ready for the dmm: pack into sequences, etc.
    val_seq_lengths = rep(val_seq_lengths)
    test_seq_lengths = rep(test_seq_lengths)
    val_batch, val_batch_reversed, val_batch_mask, val_seq_lengths = get_mini_batch(
        np.arange(n_eval_samples * val_data_sequences.shape[0]), rep(val_data_sequences),
        val_seq_lengths, cuda=use_cuda)
    val_batch = val_batch.float()
    val_batch_reversed = val_batch_reversed.float()
    val_batch_mask = val_batch_mask.float()

    test_batch, test_batch_reversed, test_batch_mask, test_seq_lengths = \
        get_mini_batch(np.arange(n_eval_samples * test_data_sequences.shape[0]),
                       rep(test_data_sequences),
                       test_seq_lengths, cuda=use_cuda)
    test_batch = test_batch.float()
    test_batch_reversed = test_batch_reversed.float()
    test_batch_mask = test_batch_mask.float()

    # put the RNN into evaluation mode (i.e. turn off drop-out if applicable)
    dmm.rnn.eval()

    # compute the validation and test loss
    val_nll = svi.evaluate_loss(val_batch, val_batch_reversed, val_batch_mask,
                                val_seq_lengths) / sum(val_seq_lengths)
    test_nll = svi.evaluate_loss(test_batch, test_batch_reversed, test_batch_mask,
                                 test_seq_lengths) / sum(test_seq_lengths)

    # put the RNN back into training mode (i.e. turn on drop-out if applicable)
    dmm.rnn.train()
    return val_nll, test_nll

def get_mini_batch(mini_batch_indices, sequences, seq_lengths, cuda=False):
    # get the sequence lengths of the mini-batch
    seq_lengths = seq_lengths[mini_batch_indices]
    # sort the sequence lengths
    _, sorted_seq_length_indices = torch.sort(seq_lengths)
    sorted_seq_length_indices = sorted_seq_length_indices.flip(0)
    sorted_seq_lengths = seq_lengths[sorted_seq_length_indices]
    sorted_mini_batch_indices = mini_batch_indices[sorted_seq_length_indices]

    # compute the length of the longest sequence in the mini-batch
    T_max = torch.max(seq_lengths)
    # this is the sorted mini-batch
    mini_batch = sequences[sorted_mini_batch_indices, 0:T_max, :]
    # this is the sorted mini-batch in reverse temporal order
    mini_batch_reversed = reverse_sequences(mini_batch, sorted_seq_lengths)
    # get mask for mini-batch
    mini_batch_mask = get_mini_batch_mask(mini_batch, sorted_seq_lengths)

    # cuda() here because need to cuda() before packing
    if cuda:
        mini_batch = mini_batch.cuda()
        mini_batch_mask = mini_batch_mask.cuda()
        mini_batch_reversed = mini_batch_reversed.cuda()

    # do sequence packing
    mini_batch_reversed = nn.utils.rnn.pack_padded_sequence(
        mini_batch_reversed, sorted_seq_lengths, batch_first=True
    )

    return mini_batch, mini_batch_reversed, mini_batch_mask, sorted_seq_lengths

## Helper functions from pyro that are not currently installed in the new version
# this function takes the hidden state as output by the PyTorch rnn and
# unpacks it it; it also reverses each sequence temporally
def pad_and_reverse(rnn_output, seq_lengths):

    rnn_output, _ = nn.utils.rnn.pad_packed_sequence(rnn_output, batch_first=True)
    reversed_output = reverse_sequences(rnn_output, seq_lengths)
    return reversed_output

def reverse_sequences(mini_batch, seq_lengths):

    reversed_mini_batch = torch.zeros_like(mini_batch)
    for b in range(mini_batch.size(0)):
        T = seq_lengths[b]
        time_slice = torch.arange(T - 1, -1, -1, device=mini_batch.device)
        reversed_sequence = torch.index_select(mini_batch[b, :, :], 0, time_slice)
        reversed_mini_batch[b, 0:T, :] = reversed_sequence
    return reversed_mini_batch

# this function returns a 0/1 mask that can be used to mask out a mini-batch
# composed of sequences of length `seq_lengths`

def get_mini_batch_mask(mini_batch, seq_lengths):
    mask = torch.zeros(mini_batch.shape[0:2])
    for b in range(mini_batch.shape[0]):
        mask[b, 0 : seq_lengths[b]] = torch.ones(seq_lengths[b])
    return mask

max_stars_repo_name: devonjkohler/sysbioDMM

max_issues_repo_name: devonjkohler/sysbioDMM

max_forks_repo_name: devonjkohler/sysbioDMM


text:
import sys
import re
import numpy as np

# glove.42B.300d.vocab.txt
# glove.42B.300d.npy
def load_glove(dir:str) -> dict:
    """
    Given The name of the directory holding files glove.42B.300d.npy and
    glove.42B.300d.vocab.txt, this function returns a dictionary mapping
    a string word to numpy array with the associated 300 dimensional
    word vector. The words are restricted to those found within
    file /usr/share/dict/words (make sure to convert those words to
    lowercase for normalization purposes).
    """
    words_set = {}
    with open("/usr/share/dict/words", "r") as f:
        for line in f.readlines():
            word = line.strip().lower()
            if words_set.get(word) is None:
                words_set[word] = 1

    f_npy = "{dir}/glove.42B.300d.npy".format(dir=dir)
    mat = np.load(f_npy)

    f_vocab = "{dir}/glove.42B.300d.vocab.txt".format(dir=dir)
    wdict = {}
    with open(f_vocab, "r") as f:
        nrows = 0
        for word in f.readlines():
            w = word.strip()
            if words_set.get(w) is not None:
                wdict[w] = mat[nrows]
            nrows += 1
    return wdict

def closest_words(gloves, word, n):
    """
    Given a gloves dictionary of word:vector and a word return the n nearest words
    as a list of strings. The word is not considered his own nearest neighbor,
    so do not include that in the returned list.

    Compute the Euclidean distance between the vector for word and
    every other word's vector. Track the distances with a list of tuples
    of the form: (distance, word).  Sort the list by distance. Return a list
    of the first n words from the sorted list. Do not return the tuples, just
    the words. Return a python list of strings not numpy array.
    """
    vector = gloves[word]
    words_sim_ret = []
    for w, v in gloves.items():
        if word == w:
            continue
        #dist = np.sqrt(sum([(x-y)**2 for x, y in zip(vector, v)]))
        dist = np.sqrt(np.sum(np.square(vector - v)))
        words_sim_ret.append((dist, w))      # why not(w, dist)
    ret = sorted(words_sim_ret, key=lambda x: x[0], reverse=False)
    words_list = [ word for dist, word in ret[:n] ]
    return words_list

def analogies(gloves, x, y, z, n):
    """
    Given a gloves dictionary of word:vector and 3 words from
    "x is to y as z is to _____", return the n best words that fill in
    the blank to complete the analogy.

    Compute the vector difference between x and y then compute the
    vector difference between z and all vectors, v, in gloves database
    (ignore v=z).  You care about the distance between the xy vector
    and the zv vector for all vectors v. Track the distances with a
    list of tuples of the form: (distance, word).  Sort the list by
    distance. Return a list of the first n words from the sorted
    list. Do not return the tuples, just the words.
    """
    #xy = np.array([(m - n) for m, n in zip(gloves.get(x), gloves.get(y))])
    xy = gloves.get(x) - gloves.get(y)
    words_sim_ret = []
    v = gloves.get(z)
    for word, vector in gloves.items():
        if z == word:
            continue
        #zv = np.array([ (j - i) for i, j in zip(vector, v)])
        zv = v - vector
        #dist = np.sqrt(sum([ ( m - n )** 2 for m, n in zip(xy, zv)]))
        dist = np.sqrt(np.sum(np.square(xy - zv)))
        words_sim_ret.append((dist, word))

    ret = sorted(words_sim_ret, key=lambda x: x[0], reverse=False)
    words = [word for dist, word in ret[:n]]
    return words

#wdict = load_glove(script_path)
#print(closest_words(wdict, 'lizard', 5))
#print(closest_words(wdict, 'russia', 5))
#print(analogies(wdict, 'shoe', 'foot', 'glove', 5))

from sklearn.decomposition import PCA
import matplotlib.pyplot as plt

def plot_words(gloves, words, n):
    """
    Get a list of vectors for n similar words for each w in words.
    Flatten that into a single list of vectors, including the original
    words' vectors.  Compute a word vector for each of those words and
    put into a list. Use PCA to project the vectors onto two dimensions.
    Extra separate X and Y coordinate lists and pass to matplotlib's scatter
    function. Then, iterate through the expanded word list and plot the
    string using text() with, say, fontsize=9. call plt.show().
    """
    words_list = []
    words_vector = []
    for word in words:
        words_list.append(word)
        words_vector.append(gloves.get(word))

        ws = closest_words(gloves, word, n)
        words_list.extend(ws)
        for w in ws:
            words_vector.append(gloves.get(w))

    pca = PCA(n_components=2)
    pca.fit(words_vector)
    cpt_2D = pca.transform(words_vector)
    plt.scatter(cpt_2D[:, 0], cpt_2D[:, 1])
    for idx, word in enumerate(words_list):
        x, y = cpt_2D[idx]
        plt.text(x, y, word, fontsize=9)
    plt.show()

#plot_words(wdict,['petal','glove','computer'], 4)
if __name__ == '__main__':
    glove_dirname = sys.argv[1]
    gloves = load_glove(glove_dirname)

    plot_words(gloves,['petal','love','king', 'cat'], 3)

    print("Enter a word or 'x:y as z:' (type 'exit' to quit)")
    cmd = ''
    while cmd!=None:
        cmd = input("> ")
        if cmd.strip()=='exit':
            break
        match = re.search(r'(\w+):(\w+) as (\w+):', cmd)
        if match is not None and len(match.groups())==3:
            x = match.group(1).lower()
            y = match.group(2).lower()
            z = match.group(3).lower()
            if x not in gloves:
                print(f"{x} is not a word that I know")
                continue
            if y not in gloves:
                print(f"{y} is not a word that I know")
                continue
            if z not in gloves:
                print(f"{z} is not a word that I know")
                continue
            words = analogies(gloves, x, y, z, 5)
            print("%s is to %s as %s is to {%s}" % (x,y,z,' '.join(words)))
        elif re.match(r'\w+', cmd) is not None:
            if cmd not in gloves:
                print(f"{cmd} is not a word that I know")
                continue
            words = closest_words(gloves, cmd.lower(), 5)
            print("%s is similar to {%s}" % (cmd,' '.join(words)))
        else:
            print("Enter a word or 'x:y as z:'")


text:
import numpy as np
from math import sqrt
from collections import Counter
from metrics import accuracy_score

class KNNClassifier:
    def __init__(self, k):
        """初始化kNN分类器"""
        assert k >= 1, "k must be valid"
        self.k = k
        self._X_train = None
        self._Y_train = None

    def fit(self, X_train, Y_train):
        """根据训练集数据，训练kNN分类器"""
        self._X_train = X_train
        self._Y_train = Y_train
        return self

    def predict(self, X_predict):
        """给定预测数据集X_predict, 返回X_predict的结果向量"""
        assert self._Y_train is not None and self._X_train is not None, \
            "must fit before predict!"
        assert X_predict.shape[1] == self._X_train.shape[1],\
            "the feature number of X_predict must be equals to X_train"

        y_predict = [self._predict(x) for x in X_predict]

        return np.array(y_predict)

    def _predict(self, x):
        """给定单个待预测数据x，返回x的预测结果"""
        assert x.shape[0] == self._X_train.shape[1],\
            "the feature number of x must be equals to X_train"

        distances = [sqrt(np.sum((x_train - x)**2)) for x_train in self._X_train]

        nearest = np.argsort(distances)

        topK_y = [self._Y_train[i] for i in nearest[:self.k]]

        votes = Counter(topK_y)
        return votes.most_common()[0][0]

    def score(self, X_test, y_test):
        """根据测试数据集X_test和y_test确定当前模型的准确度"""
        y_predict = self.predict(X_test)
        return accuracy_score(y_test, y_predict)

    def __repr__(self):
        return "KNN(k=%d)" % self.k

max_stars_repo_name: destinyddx/HelloWorld

max_issues_repo_name: destinyddx/HelloWorld

max_forks_repo_name: destinyddx/HelloWorld


text:
import pandas
import numpy

# Read the data
data = pandas.read_csv('data.csv')

# Split the data into X and y
X = numpy.array(data[['x1', 'x2']])
y = numpy.array(data['y'])

# Import the SVM Classifier
from sklearn.svm import SVC

# TODO: Define your classifier.
# Play with different values for these, from the options above.
# Hit 'Test Run' to see how the classifier fit your data.
# Once you can correctly classify all the points, hit 'Submit'.
classifier = SVC(kernel = 'poly', degree = 2)

# Fit the classifier
classifier.fit(X,y)

max_stars_repo_name: yottaawesome/ml-nanodegree-stuff

max_issues_repo_name: yottaawesome/ml-nanodegree-stuff

max_forks_repo_name: yottaawesome/ml-nanodegree-stuff


text:
import sympy as sp, math as m

x = sp.symbols('x')
sp.init_printing(use_unicode=True)

#f=sp.Lambda(x, x*sp.cos(x)-x**2*sp.sin(x))
#f=sp.Lambda(x, sp.exp(x)*sp.cos(x))
#f=sp.Lambda(x, sp.sqrt(x))

"""
def f(x):
    #return x*m.cos(x)-x**2*m.sin(x)
    return m.cos(x)
"""

def trapecio_compuesta(f, a, b, n):

    h = (b-a)/n

    result = (f(a)+f(b))/2

    for i in range(1, n):
        result += f(a+i*h)

    result *= h

    return result

def romberg(f, a, b, n):

    R=[[0 for _ in range(n+1)] for _ in range(n+1)]

    R[0][0]=trapecio_compuesta(f,a,b,1)

    for k in range(1,n+1):

        hk=(b-a)/2**k
        aux = 0
        for i in range(1,2**(k-1)+1):
            aux+=sp.N(f(a+(2*i-1)*hk))
        R[k][0]=R[k-1][0]/2+hk*aux

    for j in range(1,n+1):
        for k in range(j, n+1):
            R[k][j]=R[k][j-1]*(1+1/(4**j-1))-1/(4**j-1)*R[k-1][j-1]

    #return R[n][n], R[n][0]
    return R

"""
a, b = 0, m.pi
n = 12
"""

"""
result = romberg(f, a, b, n)
print(result)
"""

# solution = romberg(f, a, b, 12)[0]
# solution = 2/3

"""
for k in range(n+1):
    r = romberg(f, a, b, k)
    print("R[{0}][0]={1}, error={3}\tR[{0}][{0}]={2}, error={4}".format(k, r[1], r[0], abs(r[1]-solution), abs(r[0]-solution)))

"""

# Estimación Ek=4/3(Rk,0-Rk+1,0)
"""
R = romberg(f, a, b, n)

for k in range(n):
    print("R[{0}][0]={1}, error={2}, estimacion={3}".format(k, sp.N(R[k][0]), sp.N(abs(R[k][0]-R[n][n])),  sp.N(4/3*abs(R[k][0]-R[k+1][0]))))
"""

# 2 Examen

f=sp.Lambda(x, (sp.cos(sp.pi*x)+1)**(10/3))
a, b = 0, 1
solution = 3.000492123714049

n = 10
R = romberg(f, a, b, n)

print('Romberg')

for j in range(1, n+1):
    print("R[{0}][{0}]={1}".format(j, R[j][j]))
    print("\tError:" + str(abs(R[j][j]-solution)))

for j in range(1, n):
    print(abs(R[j][j]-solution)/abs(R[j+1][j+1]-solution))

print('\nTrapecio Compuesta')

for j in range(1, n+1):
    print("R[{0}][0]={1}".format(j, R[j][0]))
    print("\tError:" + str(abs(R[j][0]-solution)))

for j in range(1, n):
    print(abs(R[j][0]-solution)/abs(R[j+1][0]-solution))

reason: import sympy


text:
import os
import tweepy as tw
import sqlite3
import networkx as nx
import matplotlib.pyplot as plt
import metis
from fa2 import ForceAtlas2

#				Author:Mitos-Kotsis Antonios		     #
#						cs3028@uoi.gr 					 #
#========================================================#
#Creates the retweet graph
#Partitions the graph using Metis and color attr

G=nx.Graph()
#Database credentials
db_connection=sqlite3.connect('tweet.db')
cursor=db_connection.cursor()

#nodes
supporters=[]
deniers=[]
supporters_list=[]
deniers_list=[]
#edges
edges=[]

#contains all the users from the database that have tweeted at least once
all_users_list=[]

def getAllUsers():
	global all_users_list
	all_users=db_connection.execute("SELECT USER_ID FROM coronavirus_table")
	#transform the tuple we get to a list
	for i in all_users:
		all_users_list.append(i[0])

def getRetweetGraphData():
	global supporters,deniers,edges

	supporters=db_connection.execute("SELECT * FROM supporters_table")
	deniers=db_connection.execute("SELECT * FROM deniers_table")
	edges=db_connection.execute("SELECT * FROM retweeters_table")
	#print(edges)
	tokenizeTuple(supporters,deniers)

def tokenizeTuple(Supporters,Deniers):
	for sup in Supporters:
		supporters_list.append(sup[0])
	for den in Deniers:
		deniers_list.append(den[0])

	for sup in Supporters:
		print(sup)

#returns the color of each user
def getUserColor(user):
	global all_users_list
	if(user in all_users_list):
			cur_com=cursor.execute("SELECT COLOR FROM coronavirus_table WHERE USER_ID=(?)",(user,))
			user_color=cursor.fetchone()
		#	print(user_color[0])
			return user_color[0]
	return 'Red'

#function that creates the retweet graph
def createPartitionsForGephi():
	global supporters,deniers,edges,supporters_list,deniers_list
	G.add_edges_from(edges)

	#gets the pairs of users
	#edge[0]=initial tweet user
	#edge[1]=retweeter
	edges_in=db_connection.execute("SELECT * FROM retweeters_table")

	all_nodes=[]
	for edge in edges_in:
		#print(edge[0],edge[1])
		all_nodes.append(edge[0])
		all_nodes.append(edge[1])

	for i,v in enumerate(all_nodes):
		#vale to xrwma kathe node gia partion xwris metis
		user_color=getUserColor(v)
		if(user_color=='Blue'):
			G.nodes[v]['node_value']=1
			G.nodes[v]['color']='blue'
		else:
			G.nodes[v]['node_value']=-1
			G.nodes[v]['color']='red'

	G.graph['node_weight_attr'] = 'node_value'
	#metis graph
	G1=metis.networkx_to_metis(G)
	#Get the two partitions from METIS
	(cut, parts) = metis.part_graph(G1, 2)
	colors = ['red', 'blue']
	#print("parts",parts)
	#print(cut)
	for i,p in enumerate(parts):
		G1.node[i]['color']=colors[p]
	nx.write_gexf(G,"NoMetisPart.gexf")
	nx.write_gexf(G1, "MetisPart.gexf")

def dateTokenizer(date):

	new_date=date.split()
	temp=new_date[0].split('-')
	month=temp[1]
	return month

#For future research on monthly polarity measurements
def createMontlyPolarityGraph():
	target_month=0
	target_month=input("Enter targe month:")
	r=cursor.execute("SELECT * FROM coronavirus_table WHERE RETWEET_LIST !=''")
	monthly_tweet=[]
	monthly_graph_edges=[]
	all_nodes=[]
	for row in r:
		if(row[5]!=None):
			date=dateTokenizer(row[5])
			if(date==str(target_month)):
				monthly_tweet.append(row)

	for m in monthly_tweet:
		ret_list=m[3].split(',')
		for ret in ret_list:
			monthly_graph_edges.append((m[1],ret))
	counter=0

	for m in monthly_graph_edges:
		counter+=1
		#print(m)
	print(counter)

	G.add_edges_from(monthly_graph_edges)
	for edge in monthly_graph_edges:
		all_nodes.append(edge[0])
		all_nodes.append(edge[1])

	for i,v in enumerate(all_nodes):

		user_color=getUserColor(v)
		if(user_color=='Blue'):
			G.nodes[v]['node_value']=1
			G.nodes[v]['color']='blue'
		else:
			G.nodes[v]['node_value']=-1
			G.nodes[v]['color']='red'

	G.graph['node_weight_attr'] = 'node_value'
	#Get at MOST two partitions from METIS
	(cut, parts) = metis.part_graph(G, 2)
	colors = ['red', 'blue']

	nx.write_gexf(G, str(target_month)+".gexf")

getAllUsers()
getRetweetGraphData()
createPartitionsForGephi()

#ans=input("Create monthly graph:(yes/no)")
#if(ans=='yes'):
#	createMontlyPolarityGraph()

max_stars_repo_name: AntonisKotsis/Polarization-Measurements

max_issues_repo_name: AntonisKotsis/Polarization-Measurements

max_forks_repo_name: AntonisKotsis/Polarization-Measurements

reason: import networkx


text:
# coding=utf-8
from __future__ import division
import math
import FluidProps
import matplotlib.pyplot as plt
from scipy import interpolate
import numpy as np
import Vogel
import BeggsandBrill as BB
from scipy.optimize import fsolve
import Hagendornandbrown as HB

Oil_Rate=100
Water_Rate=50.0
GasGrav=0.65
WaterGrav=1.07
ID=2.44
Angle=90.0
FBHT=150.0
FWHP=150.0
FWHT=100.0
Depth=5000.0
nSteps=20
API=30.0
Tgrad= (FBHT-FWHT)/ Depth

Pressure=3000.0
Thickness=75.0
GasGrav=0.65
GOR= 375.0
Temp=150.0
rw=0.328
re=1053.0
k=75.0
s=-1.5
visc= 1.55
Psat = FluidProps.Pbub(Temp,75.0,100.0,GasGrav, API, GOR)
Wcut= Water_Rate/(Oil_Rate+Water_Rate)

def Press_Function(rateArray, pressArray, x):
        return np.interp(x,rateArray, pressArray)

IPR=Vogel.Vogel_DarcyIPR(Pressure,k,Thickness, visc,re,rw,s, 1.21, Temp, Psat, 20)

BB_Pwf=[]
BB_Rate=[]
for i in IPR[0]:
    if i==0:
        rate=0.1
    else:
        rate=i
    wr=rate*Wcut/(1-Wcut)
    #print 'req rate '   ,rate
    Pf=HB.Pwf_q(FWHP,FWHT,rate, wr,GOR,GasGrav,API,WaterGrav,ID,Angle, Depth, FBHT)
    BB_Pwf.append(Pf)
    BB_Rate.append(rate)
    #print Pf
#print "PWF " + str(Pf)

x= np.asarray(BB_Rate)
f= np.asarray(IPR[1])
g= np.asarray(BB_Pwf)

q=x

QNew =np.linspace(0.1, BB_Rate[-1],10000)
P_IPR=np.interp(QNew, x, f)
P_VLP=np.interp(QNew, x, g)

SolRate= [fsolve(lambda x:Press_Function(QNew,P_IPR,x)-Press_Function(QNew,P_VLP,x), 0.0)]
print0(SolRate)

plt.plot(QNew, P_IPR )
plt.plot(QNew, P_VLP )

#def findIntersection(fun1, fun2, x0):
    #return fsolve(lambda x:fun1[x]-fun2[x], x0)

#t= findIntersection(P_IPR,P_VLP,0.0)

#print t

idx = np.argwhere(np.diff(np.sign(P_IPR - P_VLP)) != 0.0).reshape(-1) + 0
plt.plot(IPR[0], IPR[1])
plt.plot(BB_Rate, BB_Pwf )

plt.plot(QNew[idx], P_IPR[idx], 'ro', ms=10)

plt.ylabel('Pwf')
plt.xlabel('Rate')
#plt.gca().invert_yaxis()
plt.show()

reason: import numpy,from scipy


text:
"""
Copyright (c) Facebook, Inc. and its affiliates.
"""
import json
import numpy as np
import re
from os.path import isfile, isdir
from os.path import join as pjoin
from torch.utils.data import Dataset
from tokenization_utils import fixed_span_values
from utils_caip import *

class CAIPDataset(Dataset):
    """Torch Dataset for the CAIP format, applies BPE and linearizes trees on-the-fly

    CAIP: CraftAssist Instruction Parsing

    Args:
        tokenizer: pre-trained tokenizer for input
        sampling: Whether to sample hard examples
        word_noise: Word noise
        sample_probas: Sampling probabilities for different data types
        dataset_length: Size of dataset
        tree_voc: Tree vocabulary file
        tree_idxs: Tree dictionary
        data: Dictionary containing datasets loaded into memory.

    """

    def __init__(
        self,
        tokenizer,
        args,
        prefix="train",
        dtype="templated",
        sampling=False,
        word_noise=0.0,
        full_tree_voc=None,
    ):
        assert isdir(args.data_dir)
        self.tokenizer = tokenizer
        self.tree_to_text = args.tree_to_text

        # We load the (input, tree) pairs for all data types and
        # initialize the hard examples buffer
        self.data = {"hard": []}
        self.sampling = sampling
        self.word_noise = word_noise
        dtype_samples = json.loads(args.dtype_samples)
        self.dtype = dtype
        self.dtypes = [t for t, p in dtype_samples]
        self.sample_probas = np.array([p for t, p in dtype_samples])
        self.sample_probas /= self.sample_probas.sum()
        if prefix == "train":
            for k in self.dtypes:
                fname = pjoin(args.data_dir, prefix, k + ".txt")
                print(fname)
                assert isfile(fname)
                print("loading {}".format(fname))
                self.data[k] = process_txt_data(fname)
            self.hard_buffer_size = 1024
            self.hard_buffer_counter = 0
        elif prefix == "":
            self.data[dtype] = []
        else:
            fname = pjoin(args.data_dir, prefix, dtype + ".txt")
            if isfile(fname):
                print("loading {}".format(fname))
                self.data[dtype] = process_txt_data(fname)
            else:
                print("could not find dataset {}".format(fname))
                self.data[dtype] = []

        # load meta-tree and tree vocabulary
        if full_tree_voc is None:
            print("making tree")
            ftr, tr_i2w = make_full_tree(
                [
                    (self.data["humanbot"], 3e5),
                    (self.data["prompts"], 1e5),
                    (self.data["templated"][:100000], 1),
                ]
            )
            self.full_tree = ftr
        else:
            full_tree, tr_i2w = full_tree_voc
            self.full_tree = full_tree
        spec_tokens = ["[PAD]", "unused", "[UNK]", "[CLS]", "[SEP]", "[MASK]", "<S>", "</S>"]
        self.tree_voc = spec_tokens[:] + tr_i2w + list(fixed_span_values)
        self.tree_idxs = dict([(w, i) for i, w in enumerate(self.tree_voc)])

        self.dataset_length = max([len(v) for v in self.data.values()])
        if args.examples_per_epoch > 0:
            self.dataset_length = min(self.dataset_length, args.examples_per_epoch)

    def _contains_span_indices(self, token_idx_list: list):
        return token_idx_list[1] >= 0 and token_idx_list[2] >= 0

    def __len__(self):
        return self.dataset_length

    def __getitem__(self, idx):
        """Sample data type and get example
        """
        if self.sampling:
            dtype = np.random.choice(self.dtypes, p=self.sample_probas)
            if len(self.data[dtype]) == 0:
                dtype = self.dtype
        else:
            dtype = self.dtype
        try:
            t = self.data[dtype][idx % len(self.data[dtype])]
            p_text, p_tree = t
        except ValueError as e:
            print(e)
        text, tree = tokenize_linearize(
            p_text, p_tree, self.tokenizer, self.full_tree, self.word_noise
        )
        text_idx_ls = [self.tokenizer._convert_token_to_id(w) for w in text.split()]
        tree_idx_ls = [
            [self.tree_idxs[w], bi, ei, text_span_start, text_span_end, (self.tree_idxs[fixed_span_val] if type(fixed_span_val) == str else fixed_span_val)]
            for w, bi, ei, text_span_start, text_span_end, fixed_span_val in [("<S>", -1, -1, -1, -1, -1)]
            + tree
            + [("</S>", -1, -1, -1, -1, -1)]
        ]
        if self.tree_to_text:
            stripped_tree_tokens = []
            for w, bi, ei in tree:
                tree_node = w.lower()
                tree_node_processed = re.sub("[^0-9a-zA-Z]+", " ", tree_node)
                tree_tokens = tree_node_processed.split(" ")
                stripped_tree_tokens += [x for x in tree_tokens if x != ""]

            extended_tree = ["[CLS]"] + stripped_tree_tokens + ["[SEP]"]
            tree_idx_ls = [self.tokenizer._convert_token_to_id(w) for w in extended_tree]
        return (text_idx_ls, tree_idx_ls, (text, p_text, p_tree))

    def add_hard_example(self, exple):
        """Add a given example for resampling.
        """
        if self.hard_buffer_counter < self.hard_buffer_size:
            self.data["hard"] += [exple]
        else:
            self.data["hard"][self.hard_buffer_counter % self.hard_buffer_size] = exple
        self.hard_buffer_counter += 1

max_stars_repo_name: kandluis/droidlet

max_issues_repo_name: kandluis/droidlet

max_forks_repo_name: kandluis/droidlet


text:
import numpy as np
from numpy import trapz
import math
from scipy.spatial.transform import Rotation
from numpy.linalg import norm

def rotate_trial_data(trial_data, angle):
    theta = math.radians(angle)  # 0 es arriba en la sim
    axis = [0, 0, 1]  # z
    axis = axis / norm(axis)
    rot = Rotation.from_rotvec(theta * axis)
    trial_data['pos_rot'] = [rot.apply([x, y, z]) for x, y, z in zip(trial_data.x, trial_data.y, trial_data.z)]
    trial_data['x'] = [x for [x, y, z] in trial_data.pos_rot]
    trial_data['y'] = [y for [x, y, z] in trial_data.pos_rot]
    trial_data['z'] = [z for [x, y, z] in trial_data.pos_rot]
    trial_data.drop('pos_rot', axis=1, inplace=True)
    return trial_data

def calculate_error_area(trial_data, angle):
    trial_data = rotate_trial_data(trial_data, angle)
    x = [trial_data.iloc[0].x, trial_data.iloc[-1].x]
    y = [trial_data.iloc[0].y, trial_data.iloc[-1].y]
    coefficients = np.polyfit(x, y, 1)
    polynomial = np.poly1d(coefficients)
    error_curve = trial_data.y - polynomial(trial_data.x)
    area_error = trapz(error_curve, x=trial_data.x)
    area_error_abs = trapz(np.abs(error_curve), x=trial_data.x)
    return area_error, area_error_abs

max_stars_repo_name: SMDynamicsLab/Haptic

max_issues_repo_name: SMDynamicsLab/Haptic

max_forks_repo_name: SMDynamicsLab/Haptic

reason: import numpy,from numpy,from scipy