<commit_before>#!/usr/bin/python3
import apsw as lite
import os as os
import time as time
import sys
import numpy as np
from scipy.stats import pearsonr

from lowess import lowess
from Camoco import Camoco
from Locus import SNP
from Chrom import Chrom

from matplotlib import pylab

class HapMap(Camoco):
    def __init__(self,name,basedir="~/.camoco"):
        # initialize super class
        super(self.__class__,self).__init__(name,type='HapMap',basedir=basedir)

    def accessions(self):
        ''' returns genotypes for HapMap object '''
        return [ x[0] for x in self.db.cursor().execute('SELECT * FROM accessions ORDER BY rowid').fetchall()]

    def snps(self,chromo,start,end):
        ''' Returns SNPs within range '''
        return self.db.cursor().execute(
            '''SELECT rowid,* FROM snps WHERE chromosome = ? AND position >= ? AND position <= ?''',
            (chromo,start,end)
        )
    def genotypes(self,snp,accessions=None):
        ''' returns the genotypes as a list based on position or snpid,
            returns None if snp not in database '''
        genos = [list(map(int,list(x[0]))) for x in self.db.cursor().execute(
        ''' SELECT alleles FROM snps JOIN genotypes ON snps.rowid = genotypes.snpRowID
            WHERE snps.chromosome = ? AND snps.position = ?''', (snp.chrom,snp.pos)
        )]
        if len(genos) == 0:
            return None
        if accessions is not None:
            try:
                return np.array(genos[0])[list(map(self.accessions().index,accessions))]
            except ValueError as e:
                self.log(e)
        return np.array(genos[0])

    def genotype_missing_mask(self,snp):
        ''' returns a boolean mask on missing genotypes '''
        return self.genotypes(snp) == 9

    def ld(self,snp1,snp2,nearest=False):
        ''' returns the pearson r2 between two snps, if snps are not in the database,
            will use the closest snp if nearest=True (default False)'''
        if nearest:
            snp1 = self.nearest_snp(snp1)
            snp2 = self.nearest_snp(snp2)
        geno1 = self.genotypes(snp1)
        geno2 = self.genotypes(snp2)
        if geno1 is None or geno1 is None:
            return (np.nan,np.nan)
        non_missing = (geno1!=9)&(geno2!=9)
        return (pearsonr(geno1[non_missing],geno2[non_missing])[0]**2,sum(non_missing),snp2-snp1)

    def downstream_snps(self,snp,snp_limit=1000,pos_limit=1e8):
        ''' returns snps downstream of snp '''
        return [SNP(*x) for x in self.db.cursor().execute(''' 
            SELECT chromosome,position,id FROM snps 
            WHERE chromosome = ?
            AND position > ?
            AND position < ?
            ORDER BY position ASC
            LIMIT ?
            ''',(snp.chrom,snp.start,snp.start+pos_limit,int(snp_limit)))]
    def upstream_snps(self,snp,snp_limit=1000,pos_limit=1e8):<commit_msg>fixed SQL bug and changed return behavior of genotypes method<commit_after>#!/usr/bin/python3
import apsw as lite
import os as os
import time as time
import sys
import numpy as np
from scipy.stats import pearsonr

from lowess import lowess
from Camoco import Camoco
from Locus import SNP
from Chrom import Chrom

from matplotlib import pylab

class HapMap(Camoco):
    def __init__(self,name,basedir="~/.camoco"):
        # initialize super class
        super(self.__class__,self).__init__(name,type='HapMap',basedir=basedir)

    def accessions(self):
        ''' returns genotypes for HapMap object '''
        return [ x[0] for x in self.db.cursor().execute('SELECT * FROM accessions ORDER BY rowid').fetchall()]

    def snps(self,chromo,start,end):
        ''' Returns SNPs within range '''
        return self.db.cursor().execute(
            '''SELECT rowid,* FROM snps WHERE chromosome = ? AND position >= ? AND position <= ?''',
            (chromo,start,end)
        )
    def genotypes(self,snp,accessions=None):
        ''' returns the genotypes as a list based on position or snpid,
            returns None if snp not in database '''
        genos = [list(map(int,list(x[0]))) for x in self.db.cursor().execute(
        ''' SELECT alleles FROM snps JOIN genotypes ON genotypes.snpRowID = snps.rowid
            WHERE snps.chromosome = ? AND snps.position = ?''', (snp.chrom,snp.pos)
        )]
        if len(genos) == 0:
            return None
        if accessions is not None:
            try:
                return np.array(genos[0])[list(map(self.accessions().index,accessions))]
            except ValueError as e:
                self.log(e)
        else:
            return np.array(genos[0])

    def genotype_missing_mask(self,snp):
        ''' returns a boolean mask on missing genotypes '''
        return self.genotypes(snp) == 9

    def ld(self,snp1,snp2,nearest=False):
        ''' returns the pearson r2 between two snps, if snps are not in the database,
            will use the closest snp if nearest=True (default False)'''
        if nearest:
            snp1 = self.nearest_snp(snp1)
            snp2 = self.nearest_snp(snp2)
        geno1 = self.genotypes(snp1)
        geno2 = self.genotypes(snp2)
        if geno1 is None or geno1 is None:
            return (np.nan,np.nan)
        non_missing = (geno1!=9)&(geno2!=9)
        return (pearsonr(geno1[non_missing],geno2[non_missing])[0]**2,sum(non_missing),snp2-snp1)

    def downstream_snps(self,snp,snp_limit=1000,pos_limit=1e8):
        ''' returns snps downstream of snp '''
        return [SNP(*x) for x in self.db.cursor().execute(''' 
            SELECT chromosome,position,id FROM snps 
            WHERE chromosome = ?
            AND position > ?
            AND position < ?
            ORDER BY position ASC
            LIMIT ?
            ''',(snp.chrom,snp.start,snp.start+pos_limit,int(snp_limit)))]
    def upstream_snps(self,snp,snp_limit=1000,pos_limit=1e8):
{'fraction_non_alphanumeric': 0.11362862010221465, 'fraction_numerical': 0.014310051107325384, 'mean_word_length': 2.5517241379310347, 'pattern_counts': {'":': 0, '<': 7, '<?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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '860982', 'n_tokens_mistral': 1946, 'n_tokens_neox': 1821, 'n_words': 489}