<jupyter_start><jupyter_text><NAME> Demo for Gomez Lab, Department of Biomedical Engineering @ UNC - Chapel Hill For discovering genes with overlapping expression in barcodes Sept 2018When analyzing single-cell mRNA expression data, it's often handy to examine the occurence of genes mRNA sqeuence analysis. Current state of the art techniques for single-cell sequencing reveal only few hundred genes for each cell, so the overall frequency and co-occurence of genes is an interesting point of study. This notebook develops some functions to identify dublets, triplets, etc. of genes that exhibit presence in at least $min_shared_cells$ number of cells.This demo code was developed to help find insights for speedups in existing code in use by the Gomez Lab in the Department of Biomedical Engineering @ UNC - Chapel Hill.<jupyter_code>import multiprocessing as mp
import itertools as it
import functools as ft
import pickle
import sys
import numpy as np<jupyter_output><empty_output><jupyter_text>We rely heavily on the frozensets, which we can use to represent combinations of genes handily. This allows us to wrap these sets into other data structures in Python for merging/accessing multiple sets at the same time.<jupyter_code>(frozenset(('acv',)) | frozenset(('abv',))) | (frozenset(('acv',)) | frozenset(('afg',)))<jupyter_output><empty_output><jupyter_text>A couple of helper functions for gather_gene_sets<jupyter_code>def is_valid_candidate(tuple_of_previous_candidates, size_k):
    return len(tuple_of_previous_candidates[0] | tuple_of_previous_candidates[1]) == size_k

def big_valid_helper(chunk_of_candidates, size_k, q):
    is_v_cand = ft.partial(is_valid_candidate, size_k = size_k)
    out = set(filter(is_v_cand, chunk_of_candidates))
    if(len(out) == 0):
        q.put(None)
    else:
        q.put(out)
    return
   

def get_next_valid_candidates(set_of_previous_candidates, n, size_k):
    ''' Input: set of candidates from last iteration.
        Returns: list of tuples of new valid candidates for this iteration.
        Something something multiprocessing
    '''

    next_candidates = list(it.combinations(set_of_previous_candidates, 2))
    
    is_v_cand = ft.partial(is_valid_candidate, size_k = size_k)

    if(n > pow(mp.cpu_count(),3)):
        p = mp.Pool(processes=max(mp.cpu_count()-1,1))
        valid_candidates = p.map(is_v_cand, next_candidates)
        valid_candidates = it.compress(next_candidates, valid_candidates)
        p.close()
    else:
        valid_candidates = filter(is_v_cand, next_candidates)

    return valid_candidates

def check_new_candidate_count(tuple_of_new_candidates, last_candidate_cell_dict):
    '''takes in a tuple of new candidates, then evaluates if the intersection would have enough cells'''
    return len(intersect_candidate_cells(tuple_of_new_candidates, last_candidate_cell_dict))

def union_candidates(tuple_of_new_candidates):
    '''takes in a tuple of new candidates, then evaluates if the intersection would have enough cells'''
    return tuple_of_new_candidates[0] | tuple_of_new_candidates[1]

def check_greater(count, min_shared_cells):
    return count > min_shared_cells

def make_candidate_cell_dict(filtered_candidates, new_candidates, last_candidate_cell_dict):
    '''combines dict entries for the filtered candidates set and returns a new candidate dict matching the new_candidates'''
    combiner = ft.partial(intersect_candidate_cells, last_candidate_cell_dict)
    if(len(new_candidates) > pow(mp.cpu_count(),2)):
        p = mp.Pool(processes=max(mp.cpu_count()-1,1))
        out = dict(zip(new_candidates, list(map(combiner, filtered_candidates))))
        p.close()
        return out
    else:
        return dict(zip(new_candidates, list(map(combiner, filtered_candidates))))
    
def intersect_candidate_cells(filtered_candidates, last_candidate_cell_dict):
    '''just a helper for make_candidate_cell_dict for mapping'''
    return last_candidate_cell_dict[filtered_candidates[0]] & last_candidate_cell_dict[filtered_candidates[1]]
    
def expected_counter(tuple_of_new_candidates, last_candidate_cell_dict, total_cells):
    '''takes in a tuple of new candidates, then evaluates if the intersection would have enough cells'''
    return last_candidate_cell_dict[tuple_of_new_candidates[0]] * last_candidate_cell_dict[tuple_of_new_candidates[1]] / total_cells

def generate_new_candidates(valid_candidates, last_candidate_cell_dict, min_shared_cells, total_cells):
    '''heavy lifter. returns new candidates, new counts, new expected counts, and new candidate cell dicts'''
    counter = ft.partial(check_new_candidate_count, last_candidate_cell_dict=last_candidate_cell_dict)
    count_checker = ft.partial(check_greater,  min_shared_cells=min_shared_cells)

    if(len(valid_candidates) >  pow(mp.cpu_count(),2)):
        p = mp.Pool(processes=max(mp.cpu_count()-1,1))
        candidate_counts = list(p.map(counter, valid_candidates))
        p.close()
    else:
        candidate_counts = list(map(counter, valid_candidates))
 
    # it.compress and check_greater are fast enough
    candidate_filter = list(map(count_checker, candidate_counts))
    filtered_candidates = list(it.compress(valid_candidates, candidate_filter))
    
    if(len(filtered_candidates)==0):
        print('Ran out of candidates! ')
        return set(), set(), set(), {}
    
    # it.compress is fast enough
    new_counts = list(it.compress(candidate_counts, candidate_filter))
    
    if(len(filtered_candidates) >  pow(mp.cpu_count(),2)):
        p = mp.Pool(processes=max(mp.cpu_count()-1,1))
        new_expected_counts = list(p.map(expected_counter, filtered_candidates))
        new_candidates = list(p.map(union_candidates, filtered_candidates))
        p.close()
    else:
        new_expected_counts = list(map(expected_counter, filtered_candidates))
        new_candidates = list(map(union_candidates, filtered_candidates))
    
    # parallelized locally
    new_candidate_cell_dict = make_candidate_cell_dict(filtered_candidates, new_candidates, last_candidate_cell_dict)
        
    return new_candidates, new_counts, new_expected_counts, new_candidate_cell_dict

### storage helpers ###
def pickle_candidates(new_candidates, new_counts, new_expected_counts, new_candidate_cell_dict, size_k):
    '''These files are gonna be decently big. Do not want to keep them in memory.'''
    with open('candidates_' + str(size_k) + '.pickle', 'wb') as f:
        pickle.dump(zip(list(new_candidates), list(new_counts), list(new_expected_counts)), f, pickle.HIGHEST_PROTOCOL)
    with open('cell_dict_' + str(size_k) + '.pickle', 'wb') as f:
        pickle.dump(new_candidate_cell_dict, f, pickle.HIGHEST_PROTOCOL)
        
def unpickle_candidates(size_k):
    '''unpickler to unpickle the last one'''
    with open('candidates_' + str(size_k) + '.pickle', 'rb') as f:
        candidates, counts, expected_counts = zip(*pickle.load(f))
    with open('cell_dict_' + str(size_k) + '.pickle', 'rb') as f:
        candidate_cell_dict = pickle.load(f)
        
    return candidates, counts, expected_counts, candidate_cell_dict 

### these functions are for n_combinations large, e.g. when there are more candidates for evaluation than we'd like to run at the same time #
def giant_filter(tuple_of_candidates, size_k, last_cell_dict, min_shared_cells):
    if(not is_valid_candidate(tuple_of_candidates, size_k)):
        return None
    elif(len(intersect_candidate_cells(tuple_of_candidates, last_cell_dict)) <= min_shared_cells): 
        return None
    else:
        return tuple_of_candidates
    
def giant_candidater(tuple_of_valid_candidates, last_cell_dict, total_cells):
    cand = union_candidates(tuple_of_valid_candidates)
    new_cell_dict = {cand: intersect_candidate_cells(tuple_of_valid_candidates, last_cell_dict)}
    count = len(new_cell_dict)
    expected_count = last_cell_dict[tuple_of_valid_candidates[0]] * last_cell_dict[tuple_of_valid_candidates[1]] / total_cells
def gather_gene_sets(tidy_df, min_shared_cells = 100, min_percent_cells = None, max_cluster_size = sys.maxsize):
    '''
        Input: tidy_df, a df with columns of ['symbol','barcode','gene name', 'class']
        Output: Writes multiple files of format "candidates_k.pickle" and "cell_dict_k.pickle" for k in 1:max_cluster_size
            
        Assumptions:    barcode corresponds to cell_id
                        gene_names <-> symbol (1:1 relationship)
    '''
    total_cells = tidy_df['barcode'].nunique()
    
    if(min_percent_cells is not None):
        min_shared_cells = int(min_percent_cells * total_cells)
    
    # get the counts for each gene
    candidate_cell_dict = tidy_df.groupby('symbol')
    counts = candidate_cell_dict.count()['barcode'].to_dict()
    candidate_cell_dict = candidate_cell_dict['barcode'].apply(frozenset)
    
    # store the symbols
    next_candidates = set(counts.keys())

    # filter stuff -- first iteration!
    candidate_filter = list(map(lambda x: x > min_shared_cells, counts.values()))
    candidates = list(map(frozenset, map(lambda x: (x,), it.compress(next_candidates, candidate_filter))))
    counts = list(it.compress(counts.values(), candidate_filter))
    
    # clean the groups dict to get candidate cell sets for each input -- may be slow!
    # print(frozenset(candidate_cell_dict.groups[next(iter(next(iter(candidates))))][0]),)
    candidate_cell_dict = dict(zip(candidates,[candidate_cell_dict[next(iter(x))] for x in candidates]))
    
    # 0 expected counts for first iteration
    expected_counts = [0] * len(candidate_cell_dict)
    
    # store our first entry!
    size_k=1
    pickle_candidates(candidates, counts, expected_counts, candidate_cell_dict, size_k)
    
    n = len(candidate_cell_dict)
    print(str(n) + ' first candidates!')
    
    # loop!
    while(n > 0 and size_k < max_cluster_size):
        size_k += 1
        
        total_combos = pow(n,2) - n 
        
        if(total_combos > pow(4096, 2)): # see if there would be > ~17 million combos
            # break into chunks of ~ 4 million for evaluation
            m = total_combos // pow(2048, 2)
            carry_over = total_combos % pow(2048, 2)
            
            total_done = 0
            ell = (m if carry_over == 0 else m + 1)
            results = []
            
            valid_candidates_iter = it.combinations(candidates, 2)
            print('created iter of size ' + str(total_combos))
            stack = 0
   
            next_candidates = filter(None, [None])
            
            p = mp.Pool(processes = max(mp.cpu_count()-1,1))
            
            big_helper = ft.partial(giant_filter, size_k = size_k, last_cell_dict = candidate_cell_dict, min_shared_cells = min_shared_cells)
            
            while total_done < ell - 1:
                next_candidates = it.chain(next_candidates, filter(lambda x: x is not None, p.map(big_helper, it.islice(valid_candidates_iter, pow(2048,2)))))
                #print('iter chunk: ' + str(total_done) + ' of ' + str(ell-1))
                total_done += 1
                
            next_candidates = it.chain(next_candidates, filter(lambda x: x is not None, p.map(big_helper, valid_candidates_iter)))
            total_done += 1
            
            #print('generated next candidates!')
            
            big_cander = ft.partial(giant_candidater, last_cell_dict=candidate_cell_dict, total_cells=total_cells)
            
            chainer = iter([])
            
            while(total_done > 0):
                chainer = it.chain(chainer, iter(p.map(big_cander, next_candidates)))
                #print('remaining chunk: ' + str(total_done) + ' of ' + str(ell))
                total_done -=  1
            
            try:
                candidates, counts, expected_counts, candidate_cell_dict = zip(*chainer)
            except:
                print('Ran out of candidates here!')
                break
            candidate_cell_dict = {x:candidate_cell_dict[x] for x[0] in candidate_cell_dict}
            
            p.close()
        else:
            ### this is the easy logic section in case we have fewer combinations. Much easier to read
            next_candidates = set(get_next_valid_candidates(candidates, n, size_k))
            print('Generated next valid candidates! Size: ' + str(size_k))
            print(len(next_candidates))
            print(next(iter(next_candidates)))
            candidates, counts, expected_counts, candidate_cell_dict = generate_new_candidates(next_candidates, candidate_cell_dict, min_shared_cells, total_cells)
        
        # calculate the new length of cells for the next iteration
        n = len(candidate_cell_dict)
        print('Evaluated '+ str(n) +' valid candidates!')
            
        # store our candidates at each stage. This allows us to reduce our RAM usage
        if(n > 0):
            pickle_candidates(candidates, counts, expected_counts, candidate_cell_dict, size_k)
import random
import pandas as pd
def generate_toy_data(n_cells, k_genes = 250):
    ''' 26**3 genes ~ 17,500 unique genes
        Lots of different cells - probably no collisions
        
        Inputs: k_genes -- the number of genes generated for each barcode (cell)
    '''
    join = lambda x: ''.join(x)
    
    genes = [0] * k_genes * n_cells
    bars = [0] * k_genes * n_cells
    classes = [0] * k_genes * n_cells
    symbols = [0] * k_genes * n_cells
    local_genes = [0] * k_genes
    
    for i in range(n_cells):
        local_genes = [join(random.choices('abcdefghijklmnopqrstuvwxyz', k=3)) for x in range(k_genes)]
        genes[i*k_genes:(i+1)*k_genes] = local_genes
        bars[i*k_genes:(i+1)*k_genes] = [join(random.choices('ABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890', k=8))] * k_genes # sample a barcode
        classes[i*k_genes:(i+1)*k_genes] = [join(random.choices('ABCD', k=1))] * k_genes # sample a class -- not accurate
        symbols[i*k_genes:(i+1)*k_genes] = genes[i*k_genes:(i+1)*k_genes]
    
    dat = pd.DataFrame()
    dat['genes'] = genes
    dat['barcode'] = bars
    dat['symbol'] = symbols
    dat['class'] = classes
         
    return dat
# generate a toy dataset
tidy_df = generate_toy_data(5000).drop_duplicates()
tidy_df['barcode'].nunique()
tidy_df['genes'].nunique()
%%timeit
gather_gene_sets(tidy_df, min_shared_cells=75, max_cluster_size=5)<jupyter_output>4721 first candidates!
created iter of size 22283120
Ran out of candidates here!
4721 first candidates!
created iter of size 22283120
Ran out of candidates here!
4721 first candidates!
created iter of size 22283120
Ran out of candidates here!
4721 first candidates!
created iter of size 22283120
Ran out of candidates here!
4721 first candidates!
created iter of size 22283120
Ran out of candidates here!
4721 first candidates!
created iter of size 22283120
Ran out of candidates here!
4721 first candidates!
created iter of size 22283120
Ran out of candidates here!
19.8 s ± 532 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
{'fraction_non_alphanumeric': 0.09372089929493149, 'fraction_numerical': 0.014167886124783823, 'mean_word_length': 3.6247308520455244, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '140986', 'license': 'MIT', 'n_tokens_mistral': 5033, 'n_tokens_neox': 4691, 'n_words': 1186, 'path': 'demos/single_cell_barcode_overlap_demo.ipynb', 'repo_name': 'isrobson/pyDataTools', 'size': '15040'}