<commit_before>import os,sys,gzip
import pandas as pd
import numpy as np
import multiprocessing as mp

from variantAnnotations import *
from Vcf_metadata import *


class Vcf(object):
    '''
        Loads in a vcf file, aware of gzipped files.
        testing
        '''
    
    def __init__(self, filename, sample_id='', cols=['#CHROM', 'POS', 'REF', 'ALT', 'FORMAT'], \
                 chunksize=5000, n_cores=1):
        
        #Header
        header_parsed = Vcf_metadata(filename)
        self.header_df = self.get_header_df(header_parsed.header)  #header parsed into key/values dataframe
        
        
        #Sample IDs
        self.samples = list(self.header_df.ix['SampleIDs'])[0]
        if sample_id == 'all':
            self.sample_id = self.samples[:]
        else:
            self.sample_id = [sample_id]
        
        
        #Columns
        self.all_columns = list(self.header_df.ix['ColumnHeader'])[0]
        self.FORMAT = self.all_columns[8]
        
        
        assert len(set(cols) & set(['#CHROM', 'POS', 'REF', 'ALT', 'FORMAT'])) > 4, "cols requires the following columns: ['#CHROM', 'POS', 'REF', 'ALT', 'FORMAT']"
        self.cols = cols
        if len(cols) > 0:  #columns specified
            self.usecols = [c for c in self.all_columns if c in cols]
            if len(sample_id) > 0:
                self.usecols.extend(self.sample_id)
            else: 
                assert False, 'no sample IDs'
        else:  #columns not specified
            self.usecols = [s for s in self.cols if s not in self.samples]
            self.usecols.extend(self.sample_id)
        
        #Open pandas chunk object (TextReader)
        self.chunksize = chunksize
        self.vcf_chunks = pd.read_table(filename, sep="\t", compression=header_parsed.compression, \
                                        skiprows=(len(self.header_df)-2), usecols=self.usecols, chunksize=chunksize)
    
    
        self.n_cores = n_cores<commit_msg>Add docstring to VCF class<commit_after>import os,sys,gzip
import pandas as pd
import numpy as np
import multiprocessing as mp

from variantAnnotations import *
from Vcf_metadata import *


class Vcf(object):
    '''
    Loads in a vcf file, aware of gzipped files.
    
    
    Parameters
    --------------------------------------
    filename: str, required
        path to vcf file
        
    sample_id: str or list, default='all'
        specifies the sample column ids to read and parse
        
        'all' means all sample columns
         
         can use a str (e.g. 'NA12878')
         or
         can use a list (e.g. ['NA12878', 'NA12877'] 
        
        
    cols: list, default ['#CHROM', 'POS', 'REF', 'ALT', 'FORMAT']
        specifies the VCF column names, EXCEPT SAMPLE COLS, to read and parse
        
        Must include ['#CHROM', 'POS', 'REF', 'ALT', 'FORMAT']
        
        Additional columns such as QUAL, FILTER, INFO will be accepted
            e.g. ['#CHROM', 'POS', 'REF', 'ALT', 'FORMAT', 'INFO', 'QUAL']
    
    chunksize: int, default=5000
        specifies the number of VCF lines to read and parse in 1 chunk
        
        Note using a large chunksize with large n_cores requires LOTS OF RAM
        
        requires ~40 seconds to parse 1000 rows with 2500 samples
        
        
    n_cores: int, default=1
        specifies the number of cpus to use during variantAnnotation
        
        Note using a large chunksize with large n_cores requires LOTS OF RAM
        
    
    Methods
    -----------------------------------------
    get_vcf_df_chunk
        returns VCF pandasDF with chunksize
        
    
    add_variant_annotations
        Annotates each variant
        See docstring for details
    
    
    
    Returns VCF Obj with following attributes
    -----------------_________________________
    header_df: pandas df
        VCF header as a pandas df
        
    samples: list
        sample column IDs
        
    all_columns: list
        all sample column IDs in VCF
        
    vcf_chunks: pandas.io.parsers.TextFileReader chunk
        VCF chunk
        Access to chunk provided by get_vcf_df_chunk()
        
    df: pandas DF
        Index: CHROM, POS, REF, ALT
        Columns: CHROM, POS, REF, ALT, SAMPLE(S) +/- {QUAL, FILTER, INFO if specified}
        
    
    '''
    
    def __init__(self, filename, sample_id='all', cols=['#CHROM', 'POS', 'REF', 'ALT', 'FORMAT'], \
                 chunksize=5000, n_cores=1):
        
        #Header
        header_parsed = Vcf_metadata(filename)
        self.header_df = self.get_header_df(header_parsed.header)  #header parsed into key/values dataframe
        
        
        #Sample IDs
        self.samples = list(self.header_df.ix['SampleIDs'])[0]
        if sample_id == 'all':
            self.sample_id = self.samples[:]
        else:
            if type(sample_id) == str:
                self.sample_id = [sample_id]
            else:
                self.sample_id = sample_id
        
        
        #Columns
        self.all_columns = list(self.header_df.ix['ColumnHeader'])[0]
        self.FORMAT = self.all_columns[8]
        
        
        assert len(set(cols) & set(['#CHROM', 'POS', 'REF', 'ALT', 'FORMAT'])) > 4, "cols requires the following columns: ['#CHROM', 'POS', 'REF', 'ALT', 'FORMAT']"
        self.cols = cols
        if len(cols) > 0:  #columns specified
            self.usecols = [c for c in self.all_columns if c in cols]
            if len(sample_id) > 0:
                self.usecols.extend(self.sample_id)
                #print self.usecols
            else: 
                assert False, 'no sample IDs'
        else:  #columns not specified
            self.usecols = [s for s in self.cols if s not in self.samples]
            self.usecols.extend(self.sample_id)
        
        #Open pandas chunk object (TextReader)
        self.chunksize = chunksize
        self.vcf_chunks = pd.read_table(filename, sep="\t", compression=header_parsed.compression, \
                                        skiprows=(len(self.header_df)-2), usecols=self.usecols, chunksize=chunksize)
    
    
        self.n_cores = n_cores
{'fraction_non_alphanumeric': 0.11096711485501377, 'fraction_numerical': 0.008909768346023004, 'mean_word_length': 1.8153214774281805, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4963662', 'n_tokens_mistral': 1879, 'n_tokens_neox': 1805, 'n_words': 525}