<commit_before>import os
import pandas as pd
from cyvcf2 import VCF
import numpy as np

v = VCF("WI.20170531.snpeff.vcf.gz")



data = []
attrs = ['CHROM',
         'POS',
         'REF',
         'ALT',
         'FILTER',
         'aaf',
         'nucl_diversity',
         'is_snp',
         'is_indel',
         'call_rate',
         'num_called',
         'num_het',
         'num_hom_ref',
         'num_hom_alt',
         'ploidy',
         'start',
         'end',
         'is_transition']

for line in v("I:1-1000000"):
    var_line = {attr:getattr(line, attr) for attr in attrs if hasattr(line, attr)}
    var_line['samples'] = v.samples
    var_line['FT'] = list(line.format("FT"))
    ANN = line.INFO.get("ANN")
    if ANN:
        var_line['ANN'] = [x.split("|") for x in ANN.split(",")]
    var_line['DP'] = list([int(x) for x in np.ndarray.flatten(line.format("DP"))])
    #var_line.update({f"INFO_{k}": v for k, v in dict(line.INFO).items()})
    data.append(var_line)

df = pd.DataFrame(data)

df.to_parquet('df.parquet', engine='pyarrow')

ANN_header = ["allele",
                  "effect",
                  "impact",
                  "gene_name",
                  "gene_id",
                  "feature_type",
                  "feature_id",
                  "transcript_biotype",
                  "exon_intron_rank",
                  "nt_change",
                  "aa_change",
                  "cDNA_position/cDNA_len",
                  "protein_position",
                  "distance_to_feature",
                  "error"]

class ANN:
    ANN_header = ["allele",
                  "effect",
                  "impact",
                  "gene_name",
                  "gene_id",
                  "feature_type",
                  "feature_id",
                  "transcript_biotype",
                  "exon_intron_rank",
                  "nt_change",
                  "aa_change",
                  "cDNA_position/cDNA_len",
                  "protein_position",
                  "distance_to_feature",
                  "error"]

def fetch_annotation_column(annotations, column):
    if annotations is not np.nan:
        ann_column_index = ANN_header.index(column)
        return [x[ann_column_index] for x in annotations]
    return np.nan


class pd_vcf:

    def __init__(self, df):
        self.df = df

    def n_variants(self, chrom, start, end):
        query_string = f"CHROM == '{chrom}' & POS > {start} & POS < {end}"
        return self.df.query(query_string).CHROM.count()


    class ANNO:
        
        def fetch_annotation_column(annotations, )

        def __init__(self):
            for ann in ANN_header:
                setattr(self, ann, property(lambda x: fetch_annotation_column(x, ann) if type(x) == list else x))



    ANN = ANNO()


pdv = pd_vcf(df)


pdv.n_variants(chrom="I", start=0, end=10000)


<commit_msg>Develop VCF Pandas frame<commit_after>import os
import json
import re
import pandas as pd
import numpy as np
from collections import defaultdict, Counter
from cyvcf2 import VCF
from pandas import DataFrame, Series
from logzero import logger
from click import secho


def infinite_dict():
    return defaultdict(infinite_dict)



class AnnotationSeries(Series):

    our_column_names = ('_ANN', )

    ANN_fields = ["allele",
                  "effect",
                  "impact",
                  "gene_name",
                  "gene_id",
                  "feature_type",
                  "feature_id",
                  "transcript_biotype",
                  "exon_intron_rank",
                  "nt_change",
                  "aa_change",
                  "cDNA_position/cDNA_len",
                  "protein_position",
                  "distance_to_feature",
                  "error"]

    def __new__(cls, *args, **kwargs):
        if kwargs.get('name', None) in cls.our_column_names:
            print("INITIALIZING COLUMN")
            obj = object.__new__(cls)
            obj.__init__(*args, **kwargs)
            return obj
        return pd.Series(*args, **kwargs)

    def __eq__(self, other):
        logger.info(other)
        return self.apply(lambda row: other in row)

    def equal(self, other):
        logger.info(other)
        return self.apply(lambda row: other in row)


    @property
    def _constructor(self):
        return AnnotationSeries

    @property
    def _constructor_expanddim(self):
        return VCF_DataFrame

    @property
    def impact(self):
        ann_column_index = self.ANN_fields.index('impact')
        result = self.apply(lambda row: [x[ann_column_index] for x in row] if type(row) == list else np.nan)
        return AnnotationSeries(data=result)



class VCF_DataFrame(DataFrame):

    _metadata = ['samples']

    ANN_fields = ["allele",
                  "effect",
                  "impact",
                  "gene_name",
                  "gene_id",
                  "feature_type",
                  "feature_id",
                  "transcript_biotype",
                  "exon_intron_rank",
                  "nt_change",
                  "aa_change",
                  "cDNA_position/cDNA_len",
                  "protein_position",
                  "distance_to_feature",
                  "error"]


    def __init__(self, *args, **kwargs):
        super(VCF_DataFrame, self).__init__(*args, **kwargs)

    @property
    def _constructor(self):
        return VCF_DataFrame


    @property
    def _constructor_sliced(self):
        return AnnotationSeries


    @classmethod
    def from_vcf(cls, filename, interval=None):
        """
            Create a numpy-array VCF object.

            filename:
                Name of the VCF
            interval:
                An interval of the VCF to use (chrom:start-end)
        """
        attrs = ['CHROM',
                 'POS',
                 'REF',
                 'ALT',
                 'FILTER',
                 'aaf',
                 'nucl_diversity',
                 'is_snp',
                 'is_indel',
                 'call_rate',
                 'num_called',
                 'num_het',
                 'num_hom_ref',
                 'num_hom_alt',
                 'ploidy',
                 'start',
                 'end',
                 'is_transition']

        rows = []
        vcf = VCF(filename, gts012=True)
        for line in vcf(interval):
            var_line = {attr: getattr(line, attr) for attr in attrs if hasattr(line, attr)}
            var_line['FT'] = np.array(line.format("FT"))
            var_line['DP'] = np.array(line.format("DP").flatten())
            var_line['gt_types'] = line.gt_types
            var_line['gt_bases'] = line.gt_bases
            ANN = line.INFO.get("ANN")
            if ANN:
                var_line['_ANN'] = [x.split("|") for x in ANN.split(",")]
            rows.append(var_line)
        dataset = VCF_DataFrame.from_dict(rows)
        # Add num missing column
        dataset['num_missing'] = dataset.gt_bases.apply(lambda row: np.sum(np.isin(row, ['./.', '.|.'])))

        # Add samples
        dataset.samples = np.array(vcf.samples)

        return dataset


    def n_variants(self, chrom, start, end):
        query_string = f"CHROM == '{chrom}' & POS > {start} & POS < {end}"
        return self.query(query_string).CHROM.count()



    def _prune_non_snps(self):
        """
            Remove snps not present in the VCF (monomorphic sites)
            Also will remove sites that are all missing.
        """
        non_snps = self.gt_types.apply(lambda x: len(np.unique(x)) > 1)
        logger.info(non_snps.size)
        return self[non_snps]


    def subset_samples(self, samples, prune_non_snps=True, inplace=False):
        sample_bool_keep = np.isin(self.samples, samples)
        df = self.copy()
        # Subset gt_types
        df.gt_types = df.gt_types.apply(lambda row: row[sample_bool_keep])
        df.gt_bases = df.gt_bases.apply(lambda row: row[sample_bool_keep])
        df.DP = df.DP.apply(lambda row: row[sample_bool_keep])
        df.FT = df.FT.apply(lambda row: row[sample_bool_keep])

        # Update variables
        df.num_hom_ref = df.gt_types.apply(lambda row: np.sum(row == 0))
        df.num_het = df.gt_types.apply(lambda row: np.sum(row == 1))
        df.num_hom_alt = df.gt_types.apply(lambda row: np.sum(row == 2))
        df.num_missing = df.gt_bases.apply(lambda row: np.sum(np.isin(row, ['./.', '.|.'])))
        # Do not change '==' to 'is'; numpy doesn't use 'in'.
        df.num_called = df.gt_bases.apply(lambda row: np.sum(np.isin(row, ['./.', '.|.']) == False))
        df.call_rate = df.gt_types.apply(lambda row: np.sum(row != 3)/row.size)

        if prune_non_snps:
            original_size = df.size
            df = df._prune_non_snps()
            pruned_snps = original_size - df.size
            secho(f"Pruned SNPs: {pruned_snps}", fg='green')

        # Update samples
        df.samples = self.samples[np.isin(self.samples, samples)]

        if inplace:
            self.samples = df.samples
            self = df
        else:
            return VCF_DataFrame(df)


    def _parse_interval(interval):
        """
            Parses an interval
        """
        chrom, *pos = re.split(":-", interval)
        if len(pos) not in [0, 2]:
            raise Exception("Invalid interval")
        elif len(pos) == 2:
            pos = list(map(int, pos))
            return chrom, pos[0], pos[1]
        return chrom, None, None


    def ANN(self, field, eq=None):
        """
            Constructs a nested ndarray([field, field], ...) of 
            the specified field.

            Args:
                field - The fieldname to pull out of all variant annotations from the ANN field.
                eq - If specified, returns a vector True/False if eq is found in any of the ANN field annotations.
        """
        ann_column_index = self.ANN_fields.index(field)
        field_result = self._ANN.apply(lambda row: [x[ann_column_index] for x in row] if type(row) == list else np.nan)
        if eq:
            return field_result.apply(lambda row: eq in row if type(row) == list else False)
        else:
            return field_result

    def ANN_count(self, field, unique=True):
        """
            Outputs the distribution of values for each label

            Args:
                unique - If set to False, will output the total number of field labels
                         for annotations.
        """
        self.ANN(field)


    def interval(self, interval):
        """
            Filters a VCF on an interval
        """
        chrom, start, end = self._parse_interval(interval)
        if chrom and start and end:
            query_string = f"CHROM == '{chrom}' & POS > {start} & POS < {end}"
        elif chrom:
            query_string = f"CHROM == '{chrom}'"
        return self.query(query_string)


    def genome_summary(self, interval=None):
        """
            Generates a comprehensive interval summary
        """
        if interval:
            df = self.interval(interval)
        else:
            df = self

        results = infinite_dict()

        # Impact
        impact = results['variants']['impact']
        impact['total'] = Counter(sum(df.ANN("impact").dropna(), []))
        impact['unique'] = Counter(sum(df.ANN("impact").dropna().apply(lambda x: list(set(x))), []))
        
        # FILTER summary
        impact = results['variants']['impact']
        impact['total'] = Counter(sum(df.ANN("impact").dropna(), []))
        impact['unique'] = Counter(sum(df.ANN("impact").dropna().apply(lambda x: list(set(x))), []))

        # Summary
        summary = results['variants']


        summary['filters']['FILTER'] = Counter(df.FILTER.dropna())
        FT_vals = np.concatenate(df.FT.values)
        summary['filters']['FT']['combined'] = Counter(FT_vals)
        summary['filters']['FT']['separate'] = Counter(np.concatenate(Series(FT_vals).apply(lambda x: x.split(";")).values))


        # snp
        summary['snp']['records'] = sum(df.is_snp)
        summary['snp']['num_missing'] = sum(df[df.is_snp].num_missing)
        summary['snp']['avg_call_rate'] = np.average(df[df.is_snp].call_rate)
        summary['snp']['transition'] = sum(df[df.is_snp].is_transition)
        summary['snp']['transversion'] = sum(df[df.is_snp].is_transition)
        summary['snp']['num_hom_ref'] = sum(df[df.is_snp].num_hom_ref)
        summary['snp']['num_het'] = sum(df[df.is_snp].num_het)
        summary['snp']['num_hom_alt'] = sum(df[df.is_snp].num_hom_alt)

        # indel
        summary['indel']['records'] = sum(df.is_indel)
        summary['indel']['num_missing'] = sum(df[df.is_indel].num_missing)
        summary['indel']['avg_call_rate'] = np.average(df[df.is_indel].call_rate)
        summary['indel']['transition'] = sum(df[df.is_indel].is_transition)
        summary['indel']['transversion'] = sum(df[df.is_indel].is_transition)
        summary['indel']['num_hom_ref'] = sum(df[df.is_indel].num_hom_ref)
        summary['indel']['num_het'] = sum(df[df.is_indel].num_het)
        summary['indel']['num_hom_alt'] = sum(df[df.is_indel].num_hom_alt)

        print(json.dumps(results, indent=4, sort_keys=True))
        # Genes
        return json.dumps(results)




v = VCF_DataFrame.from_vcf("WI.20170531.snpeff.vcf.gz", "I:1-10000")

df = v

v.ANN("impact")
v.ANN("impact", "HIGH")
v.ANN("gene_name")


v.genome_summary()



{'fraction_non_alphanumeric': 0.13221296847114813, 'fraction_numerical': 0.003866745984533016, 'mean_word_length': 1.911038961038961, 'pattern_counts': {'":': 2, '<': 6, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1260168', 'n_tokens_mistral': 4353, 'n_tokens_neox': 4068, 'n_words': 856}