<s> import sys <EOL> from collections import Counter , defaultdict <EOL> import GeminiQuery <EOL> import sql_utils <EOL> import compiler <EOL> from gemini_constants import * <EOL> from . gemini_bcolz import filter , NoGTIndexException <EOL> from . mendelianerror import mendelian_error <EOL> import itertools as it <EOL> import operator as op <EOL> from inheritance import Family <EOL> class GeminiInheritanceModel ( object ) : <EOL> required_columns = ( "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , <EOL> "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" ) <EOL> def __init__ ( self , args ) : <EOL> self . args = args <EOL> self . gq = GeminiQuery . GeminiQuery ( args . db , include_gt_cols = True ) <EOL> self . added = [ ] <EOL> self . gt_cols = self . gq . gt_cols <EOL> if not args . columns : <EOL> args . columns = "<STR_LIT>" + "<STR_LIT:U+002CU+0020>" . join ( self . gt_cols ) <EOL> self . set_family_info ( ) <EOL> @ property <EOL> def query ( self ) : <EOL> if self . args . columns is not None : <EOL> query = "<STR_LIT>" + str ( self . args . columns ) + "<STR_LIT>" <EOL> else : <EOL> query = "<STR_LIT>" + "<STR_LIT>" % "<STR_LIT:U+002CU+0020>" . join ( self . gt_cols ) <EOL> query = sql_utils . ensure_columns ( query , [ '<STR_LIT>' , '<STR_LIT>' ] ) <EOL> if self . args . filter : <EOL> query += "<STR_LIT>" + self . args . filter <EOL> if hasattr ( self . args , '<STR_LIT:X>' ) : <EOL> if self . args . X == [ ] : <EOL> self . args . X = [ '<STR_LIT>' , '<STR_LIT:X>' ] <EOL> part = "<STR_LIT>" % "<STR_LIT:U+002CU+0020>" . join ( "<STR_LIT>" % x for x in self . args . X ) <EOL> if "<STR_LIT>" in query : <EOL> query += "<STR_LIT>" + part <EOL> else : <EOL> query += "<STR_LIT>" + part <EOL> if self . model in ( "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" ) or ( self . model == "<STR_LIT>" and self . args . min_kindreds is not None ) : <EOL> if "<STR_LIT>" in query : <EOL> query += "<STR_LIT>" <EOL> else : <EOL> query += "<STR_LIT>" <EOL> if self . model == "<STR_LIT>" or not ( <EOL> self . args . min_kindreds in ( None , <NUM_LIT:1> ) <EOL> and ( self . args . families is None <EOL> or not "<STR_LIT:U+002C>" in self . args . families ) ) : <EOL> query += "<STR_LIT>" <EOL> return query <EOL> def bcolz_candidates ( self ) : <EOL> """<STR_LIT>""" <EOL> variant_ids = set ( ) <EOL> try : <EOL> for i , family_id in enumerate ( self . family_ids ) : <EOL> gt_filter = self . family_masks [ i ] <EOL> if isinstance ( gt_filter , dict ) : <EOL> for k , flt in gt_filter . items ( ) : <EOL> variant_ids . update ( filter ( self . args . db , flt , { } ) ) <EOL> else : <EOL> if gt_filter == '<STR_LIT:False>' : continue <EOL> variant_ids . update ( filter ( self . args . db , gt_filter , { } ) ) <EOL> return sorted ( set ( variant_ids ) ) <EOL> except NoGTIndexException : <EOL> return None <EOL> def gen_candidates ( self , group_key ) : <EOL> if isinstance ( group_key , basestring ) : <EOL> group_key = op . itemgetter ( group_key ) <EOL> q = self . query <EOL> vids = self . bcolz_candidates ( ) <EOL> if vids is None : <EOL> self . gq . run ( q , needs_genotypes = True ) <EOL> elif len ( vids ) > <NUM_LIT:0> : <EOL> q = GeminiQuery . add_variant_ids_to_query ( q , vids ) <EOL> self . gq . run ( q , needs_genotypes = True ) <EOL> else : <EOL> raise StopIteration <EOL> for grp_key , grp in it . groupby ( self . gq , group_key ) : <EOL> yield grp_key , grp <EOL> def all_candidates ( self ) : <EOL> _ , candidates = self . gen_candidates ( group_key = None ) <EOL> for candidate in candidates : <EOL> yield candidate <EOL> def gene_candidates ( self ) : <EOL> for gene , candidates in self . gen_candidates ( group_key = "<STR_LIT>" ) : <EOL> yield gene , candidates <EOL> def set_family_info ( self ) : <EOL> """<STR_LIT>""" <EOL> self . families = families = Family . from_cursor ( self . gq . conn ) . values ( ) <EOL> args = self . args <EOL> self . family_ids = [ ] <EOL> self . family_masks = [ ] <EOL> kwargs = { '<STR_LIT>' : not getattr ( self . args , "<STR_LIT>" , False ) , <EOL> '<STR_LIT>' : args . min_gq } <EOL> if self . model == "<STR_LIT>" : <EOL> kwargs = { '<STR_LIT>' : self . args . only_affected } <EOL> if self . model != "<STR_LIT>" and self . model != "<STR_LIT>" : <EOL> if hasattr ( self . args , '<STR_LIT>' ) : <EOL> kwargs [ '<STR_LIT:strict>' ] = not self . args . lenient <EOL> elif self . model == "<STR_LIT>" : <EOL> kwargs [ '<STR_LIT>' ] = self . args . pattern_only <EOL> if hasattr ( self . args , '<STR_LIT>' ) : <EOL> kwargs [ '<STR_LIT>' ] = self . args . gt_phred_ll <EOL> if self . model in ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ) : <EOL> kwargs . pop ( '<STR_LIT>' ) <EOL> requested_fams = None if not args . families else set ( args . families . split ( "<STR_LIT:U+002C>" ) ) <EOL> for family in families : <EOL> if requested_fams is None or family . family_id in requested_fams : <EOL> family_filter = getattr ( family , self . model ) ( <EOL> min_depth = self . args . min_sample_depth , <EOL> ** kwargs ) <EOL> else : <EOL> family_filter = '<STR_LIT:False>' <EOL> self . family_masks . append ( family_filter ) <EOL> self . family_ids . append ( family . family_id ) <EOL> def report_candidates ( self ) : <EOL> args = self . args <EOL> req_cols = [ '<STR_LIT>' , '<STR_LIT>' ] <EOL> if args . min_sample_depth and self . args . min_sample_depth > <NUM_LIT:0> : <EOL> req_cols . append ( '<STR_LIT>' ) <EOL> if getattr ( args , '<STR_LIT>' , False ) or self . model == "<STR_LIT>" : <EOL> for col in [ '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' ] : <EOL> if col in self . gt_cols : <EOL> req_cols . append ( col ) <EOL> if args . min_gq != <NUM_LIT:0> and '<STR_LIT>' in self . gt_cols : <EOL> req_cols . append ( '<STR_LIT>' ) <EOL> is_mendel = False <EOL> if any ( isinstance ( m , dict ) for m in self . family_masks ) : <EOL> assert self . model == "<STR_LIT>" <EOL> is_mendel = True <EOL> masks = [ ] <EOL> for mdict in self . family_masks : <EOL> if isinstance ( mdict , basestring ) : <EOL> masks . append ( mdict ) <EOL> continue <EOL> m = { } <EOL> for k , mask in mdict . items ( ) : <EOL> m [ k ] = '<STR_LIT:False>' if mask is None or mask . strip ( "<STR_LIT:(>" ) . strip ( "<STR_LIT:)>" ) == '<STR_LIT:False>' else mask <EOL> if m [ k ] != '<STR_LIT:False>' : <EOL> m [ k ] = compiler . compile ( m [ k ] , m [ k ] , '<STR_LIT>' ) <EOL> masks . append ( m ) <EOL> else : <EOL> masks = [ '<STR_LIT:False>' if m is None or m is False or m . strip ( '<STR_LIT:(>' ) . strip ( '<STR_LIT:)>' ) in <EOL> ( '<STR_LIT>' , '<STR_LIT:False>' ) else m for m in self . family_masks ] <EOL> masks = [ compiler . compile ( m , m , '<STR_LIT>' ) if m != '<STR_LIT:False>' else '<STR_LIT:False>' for m in masks ] <EOL> requested_fams = None if not args . families else set ( args . families . split ( "<STR_LIT:U+002C>" ) ) <EOL> for gene , li in self . candidates ( ) : <EOL> kindreds = set ( ) <EOL> to_yield = [ ] <EOL> seen = set ( ) <EOL> for row in li : <EOL> cur_fam = row . print_fields . get ( '<STR_LIT>' ) <EOL> for col in self . added : <EOL> try : <EOL> del row . print_fields [ col ] <EOL> except ValueError : <EOL> pass <EOL> cols = dict ( ( col , row [ col ] ) for col in req_cols ) <EOL> fams_to_test = enumerate ( self . families ) if cur_fam is None else [ ( i , f ) for i , f in enumerate ( self . families ) if f . family_id == cur_fam ] <EOL> if requested_fams is not None : <EOL> fams_to_test = ( ( i , f ) for i , f in fams_to_test if f . family_id in requested_fams ) <EOL> fams , models = [ ] , [ ] <EOL> if is_mendel : <EOL> for i , f in fams_to_test : <EOL> mask_dict = masks [ i ] <EOL> for inh_model , mask in mask_dict . items ( ) : <EOL> if masks [ i ] != '<STR_LIT:False>' and eval ( mask , cols ) : <EOL> if f in fams : <EOL> models [ - <NUM_LIT:1> ] += "<STR_LIT:;>" + inh_model <EOL> else : <EOL> fams . append ( f ) <EOL> models . append ( inh_model ) <EOL> else : <EOL> fams = [ f for i , f in fams_to_test <EOL> if masks [ i ] != '<STR_LIT:False>' and eval ( masks [ i ] , cols ) ] <EOL> kindreds . update ( f . family_id for f in fams ) <EOL> for fam in fams : <EOL> pdict = row . print_fields . copy ( ) <EOL> kindreds . add ( fam . family_id ) <EOL> pdict [ "<STR_LIT>" ] = fam . family_id <EOL> pdict [ "<STR_LIT>" ] = "<STR_LIT:U+002C>" . join ( "<STR_LIT:%s>" % m for m in fam . subjects ) <EOL> pdict [ "<STR_LIT>" ] = "<STR_LIT:U+002C>" . join ( eval ( str ( s ) , cols ) for s in fam . gts ) <EOL> pdict [ "<STR_LIT>" ] = "<STR_LIT:U+002C>" . join ( x . name or x . sample_id for x in fam . subjects if x . affected ) <EOL> pdict [ "<STR_LIT>" ] = len ( fams ) <EOL> if is_mendel : <EOL> pdict [ "<STR_LIT>" ] = "<STR_LIT:;>" . join ( models ) <EOL> pdict [ '<STR_LIT>' ] = fam . affecteds_with_parent if args . only_affected else fam . samples_with_parent <EOL> vs = [ ] <EOL> if all ( c in self . gt_cols for c in ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ) ) : <EOL> pdict [ "<STR_LIT>" ] = "<STR_LIT>" <EOL> for s in pdict [ '<STR_LIT>' ] : <EOL> mdk = str ( s . mom . genotype_lls + s . dad . genotype_lls + s . genotype_lls ) <EOL> vals = eval ( mdk , cols ) <EOL> vs . append ( mendelian_error ( vals [ : <NUM_LIT:3> ] , vals [ <NUM_LIT:3> : <NUM_LIT:6> ] , vals [ <NUM_LIT:6> : ] , pls = True ) ) <EOL> pdict [ "<STR_LIT>" ] = "<STR_LIT:U+002C>" . join ( "<STR_LIT>" % v for v in vs if v is not None ) <EOL> pdict [ '<STR_LIT>' ] = "<STR_LIT:U+002C>" . join ( s . name or str ( s . sample_id ) for s in pdict [ '<STR_LIT>' ] ) <EOL> s = str ( pdict ) <EOL> if s in seen : <EOL> continue <EOL> seen . add ( s ) <EOL> to_yield . append ( pdict ) <EOL> if <NUM_LIT:0> < len ( kindreds ) >= self . args . min_kindreds : <EOL> if '<STR_LIT>' in to_yield [ <NUM_LIT:0> ] : <EOL> counts = Counter ( ( item [ '<STR_LIT>' ] , item [ '<STR_LIT>' ] ) for item in to_yield ) <EOL> to_yield = [ item for item in to_yield if counts [ ( item [ '<STR_LIT>' ] , item [ '<STR_LIT>' ] ) ] > <NUM_LIT:1> ] <EOL> if len ( set ( item [ '<STR_LIT>' ] for item in to_yield ) ) < self . args . min_kindreds : <EOL> continue <EOL> for item in to_yield : <EOL> item [ '<STR_LIT>' ] = len ( kindreds ) <EOL> yield item <EOL> def run ( self ) : <EOL> has_gts = False <EOL> from . gemini_bcolz import gt_cols_types <EOL> for i , s in enumerate ( self . report_candidates ( ) ) : <EOL> if i == <NUM_LIT:0> : <EOL> has_gts = [ x [ <NUM_LIT:0> ] for x in gt_cols_types if x [ <NUM_LIT:0> ] in s ] or False <EOL> print "<STR_LIT:\t>" . join ( s . keys ( ) ) <EOL> if has_gts : <EOL> for col in has_gts : <EOL> s [ col ] = str ( s [ col ] ) . replace ( '<STR_LIT:\n>' , '<STR_LIT>' ) <EOL> print "<STR_LIT:\t>" . join ( map ( str , s . values ( ) ) ) <EOL> class XRec ( GeminiInheritanceModel ) : <EOL> model = "<STR_LIT>" <EOL> def candidates ( self ) : <EOL> for g , li in self . gen_candidates ( '<STR_LIT>' ) : <EOL> yield g , li <EOL> class XDenovo ( XRec ) : <EOL> model = "<STR_LIT>" <EOL> class XDom ( XRec ) : <EOL> model = "<STR_LIT>" <EOL> class AutoDom ( GeminiInheritanceModel ) : <EOL> model = "<STR_LIT>" <EOL> def candidates ( self ) : <EOL> for g , li in self . gen_candidates ( '<STR_LIT>' ) : <EOL> yield g , li <EOL> class AutoRec ( AutoDom ) : <EOL> model = "<STR_LIT>" <EOL> class DeNovo ( GeminiInheritanceModel ) : <EOL> model = "<STR_LIT>" <EOL> def candidates ( self ) : <EOL> kins = self . args . min_kindreds <EOL> for g , li in self . gen_candidates ( '<STR_LIT>' if kins is not None else None ) : <EOL> yield g , li <EOL> class MendelViolations ( GeminiInheritanceModel ) : <EOL> model = "<STR_LIT>" <EOL> def candidates ( self ) : <EOL> for g , li in self . gen_candidates ( None ) : <EOL> yield g , li <EOL> class CompoundHet ( GeminiInheritanceModel ) : <EOL> model = "<STR_LIT>" <EOL> @ property <EOL> def query ( self ) : <EOL> args = self . args <EOL> if args . columns is not None : <EOL> custom_columns = self . _add_necessary_columns ( str ( args . columns ) ) <EOL> query = "<STR_LIT>" + custom_columns + "<STR_LIT>" + "<STR_LIT>" <EOL> else : <EOL> query = "<STR_LIT>" + "<STR_LIT>" + "<STR_LIT>" + "<STR_LIT>" <EOL> if args . filter : query += "<STR_LIT>" + args . filter <EOL> return query + "<STR_LIT>" <EOL> def _add_necessary_columns ( self , custom_columns ) : <EOL> """<STR_LIT>""" <EOL> custom_columns = [ x . strip ( ) for x in custom_columns . split ( "<STR_LIT:U+002C>" ) ] <EOL> if "<STR_LIT:*>" in custom_columns : <EOL> return "<STR_LIT:U+002C>" . join ( custom_columns ) <EOL> self . added = [ ] <EOL> for col in ( "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT:start>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" ) : <EOL> if not col in custom_columns : <EOL> custom_columns . append ( col ) <EOL> if col != "<STR_LIT>" : <EOL> self . added . append ( col ) <EOL> return "<STR_LIT:U+002C>" . join ( custom_columns ) <EOL> def filter_candidates ( self , candidates , <EOL> comp_het_counter = [ <NUM_LIT:0> ] ) : <EOL> """<STR_LIT>""" <EOL> args = self . args <EOL> requested_fams = None if args . families is None else set ( args . families . split ( "<STR_LIT:U+002C>" ) ) <EOL> for idx , comp_het in enumerate ( candidates ) : <EOL> comp_het_counter [ <NUM_LIT:0> ] += <NUM_LIT:1> <EOL> for fam_ch in candidates [ comp_het ] : <EOL> if fam_ch [ '<STR_LIT>' ] > args . max_priority : <EOL> continue <EOL> for subject in ( fam_ch [ '<STR_LIT>' ] if args . pattern_only else fam_ch [ '<STR_LIT>' ] + fam_ch [ '<STR_LIT>' ] ) : <EOL> family_id = subject . family_id <EOL> if requested_fams is not None and not family_id in requested_fams : <EOL> continue <EOL> ch_id = str ( comp_het_counter [ <NUM_LIT:0> ] ) <EOL> cid = "<STR_LIT>" % ( ch_id , comp_het [ <NUM_LIT:0> ] . row [ '<STR_LIT>' ] , <EOL> comp_het [ <NUM_LIT:1> ] . row [ '<STR_LIT>' ] ) <EOL> for i in ( <NUM_LIT:0> , <NUM_LIT:1> ) : <EOL> pdict = comp_het [ i ] . row . print_fields . copy ( ) <EOL> pdict [ "<STR_LIT>" ] = family_id <EOL> pdict [ "<STR_LIT>" ] = None <EOL> pdict [ "<STR_LIT>" ] = None <EOL> pdict [ "<STR_LIT>" ] = None <EOL> pdict [ "<STR_LIT>" ] = None <EOL> pdict [ "<STR_LIT>" ] = cid <EOL> pdict [ '<STR_LIT>' ] = fam_ch [ '<STR_LIT>' ] <EOL> comp_het [ i ] . row . print_fields = pdict <EOL> yield comp_het [ i ] . row <EOL> def candidates ( self ) : <EOL> args = self . args <EOL> self . gq . _connect_to_database ( ) <EOL> fams = self . fams = Family . from_cursor ( self . gq . conn ) <EOL> if args . families : <EOL> fams = { f : fam for f , fam in fams . items ( ) if f in set ( args . families . split ( "<STR_LIT:U+002C>" ) ) } <EOL> for grp , li in self . gen_candidates ( '<STR_LIT>' ) : <EOL> samples_w_hetpair = defaultdict ( list ) <EOL> sites = [ ] <EOL> for row in li : <EOL> gt_types , gt_bases , gt_phases = row [ '<STR_LIT>' ] , row [ '<STR_LIT>' ] , row [ '<STR_LIT>' ] <EOL> site = Site ( row ) <EOL> site . gt_phases , site . gt_bases , site . gt_types = gt_phases , gt_bases , gt_types <EOL> sites . append ( site ) <EOL> for i , site1 in enumerate ( sites [ : - <NUM_LIT:1> ] , start = <NUM_LIT:1> ) : <EOL> for site2 in sites [ i : ] : <EOL> for family_id , fam in fams . items ( ) : <EOL> ch = fam . comp_het_pair ( site1 . gt_types , site1 . gt_bases , <EOL> site2 . gt_types , site2 . gt_bases , <EOL> site1 . gt_phases , site2 . gt_phases , <EOL> ref1 = site1 . row [ '<STR_LIT>' ] , <EOL> alt1 = site1 . row [ '<STR_LIT>' ] , <EOL> ref2 = site2 . row [ '<STR_LIT>' ] , <EOL> alt2 = site2 . row [ '<STR_LIT>' ] , <EOL> allow_unaffected = args . allow_unaffected , <EOL> fast_mode = True , <EOL> pattern_only = args . pattern_only ) <EOL> if not ch [ '<STR_LIT>' ] : continue <EOL> samples_w_hetpair [ ( site1 , site2 ) ] . append ( ch ) <EOL> yield grp , self . filter_candidates ( samples_w_hetpair ) <EOL> class Site ( object ) : <EOL> __slots__ = ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ) <EOL> def __init__ ( self , row ) : <EOL> self . row = row <EOL> self . gt_phases = None <EOL> self . gt_bases = None <EOL> self . gt_types = None <EOL> def __eq__ ( self , other ) : <EOL> return self . row [ '<STR_LIT>' ] == other . row [ '<STR_LIT>' ] and self . row [ '<STR_LIT:start>' ] == other . row [ '<STR_LIT:start>' ] <EOL> def __repr__ ( self ) : <EOL> return "<STR_LIT:U+002C>" . join ( [ self . row [ '<STR_LIT>' ] , <EOL> str ( self . row [ '<STR_LIT:start>' ] ) , <EOL> str ( self . row [ '<STR_LIT:end>' ] ) ] ) <EOL> def __hash__ ( self ) : <EOL> "<STR_LIT>" <EOL> return sum ( ord ( c ) for c in self . row [ '<STR_LIT>' ] ) + int ( self . row [ '<STR_LIT:start>' ] ) </s>
