text:
#!/usr/bin/env python2.6

"""
08/30/2011 nknguyen at soe ucsc edu

Input: location of vjUsage matrices where rows = Vs and columns= Js.
Each cell X_ij contains the (clonotype/read) count that have the combination V_i and J_j.
Each file corresponds to one sample.

Output: j usage plot, v usage plot (x = genes, y = usage,  different lines = samples)
        vjUsage heatmap for each sample
        vjUsage heatmaps for all samples in one file
"""

import os, sys, re
from optparse import OptionParser

#import libPlotting as libplot
import matplotlib as mpl
import matplotlib.pyplot as pyplot
from matplotlib.ticker import *
from matplotlib.font_manager import FontProperties
#from numpy import array
import immunoseq.lib.immunoseqLib as immunoseqLib
from scipy.stats.stats import pearsonr
from scipy.stats import ttest_ind
import numpy as np
import mdp
#from mpl_toolkits.axes_grid1.inset_locator import inset_axes

#from mpl_toolkits.mplot3d import Axes3D

class Sample:
    def __init__(self, name):
        self.name = name
        self.vjusage = [] #list of list (2D matrix: rows = Vs, Columns = Js)
        self.vcounts = []
        self.jcounts = []
        self.totalCount = 0.0

    def setUsage(self, matrix, rownames, colnames):
        self.vjusage = matrix
        self.vgenes = rownames
        self.jgenes = colnames
        if len(self.vgenes) != len(self.vjusage):
            sys.stderr.write("Row labels vector has different length from the number of rows.\n")
            sys.exit(1)
        for row in self.vjusage:
            if len(row) != len(self.jgenes):
                sys.stderr.write("Column labels vector has different length (%d) from the number of columns (%d).\n" %(len(self.jgenes), len(row)))
                sys.stderr.write("Column labels: %s\n" %('\t'.join(self.jgenes)))
                sys.stderr.write("Sample: %s\n" %self.name)
                print row
                sys.exit(1)

    def setVcounts( self ):
        self.vcounts = [ sum(v) for v in self.vjusage ]

    def setJcounts( self ):
        self.jcounts = [ sum([v[j] for v in self.vjusage]) for j in range( len(self.jgenes) ) ]

    def setTotalCount( self ):
        self.totalCount = sum( [sum(v) for v in self.vjusage] )

    def getVfreqs( self ):
        if self.totalCount == 0:
            return []
        return [ c/self.totalCount for c in self.vcounts ]

    def getVfreqs2( self, genes, abs ):
        if self.totalCount == 0:
            return []
        freqs = []
        for g in genes:
            if g not in self.vgenes:
                raise KeyError("Gene %s should be in vgenes, but were not found of %s. Vgenes: %s\n" %(g, self.name, ",".join(self.vgenes)))
            for i, currgene in enumerate(self.vgenes):
                if g == currgene:
                    if abs:
                        freqs.append(self.vcounts[i])
                    else:
                        freqs.append( self.vcounts[i]/self.totalCount )
        return freqs

    def getJfreqs( self ):
        if self.totalCount == 0:
            return []
        return [ c/self.totalCount for c in self.jcounts ]

    def getJfreqs2( self, genes, abs ):
        if self.totalCount == 0:
            return []
        freqs = []
        for g in genes:
            if g not in self.jgenes:
                raise KeyError("Gene %s should be in jgenes, but were not found of %s. Vgenes: %s\n" %(g, self.name, ",".join(self.jgenes)))
            for i, currgene in enumerate(self.jgenes):
                if g == currgene:
                    if abs:
                        freqs.append(self.jcounts[i])
                    else:
                        freqs.append( self.jcounts[i]/self.totalCount )
        return freqs

    def normalizeVJusage( self ):
        m = []
        for r in self.vjusage:
            row = [ float(c)/self.totalCount for c in r ]
            m.append( row )
        return m

def getGeneIndices(genes, subgenes):
    indices = []
    for g in subgenes:
        for i, gene in enumerate(genes):
            if g == gene:
                indices.append(i)
                break
    return indices

def intersectVJ(samples, type2intersectGenes):
    ivs = type2intersectGenes['v']
    ijs = type2intersectGenes['j']
    for sample in samples:
        sample.intersectVJusage = intersectVJsample( sample.vjusage, sample.vgenes, sample.jgenes, ivs, ijs)
        sample.normIntersectVJusage = normalizeVJusage( sample.intersectVJusage )
    return

def intersectVJsample(vjusage, vs, js, ivs, ijs):
    #vs: current V genes (or vjusage rownames), js: current J genes (or vjusage colnames), ivs: itersect v genes, ijs: intersect j genes
    vindices = getGeneIndices(vs, ivs)
    jindices = getGeneIndices(js, ijs)
    m = []
    for vindex in vindices: #rows:
        row = vjusage[vindex]
        newrow = []
        for jindex in jindices:
            newrow.append( row[jindex] )
        m.append(newrow)
    return m

def getTotalCount( vjusage ):
    return sum( [sum(v) for v in vjusage] )

def normalizeVJusage(vjusage):
    m = []
    totalCount = getTotalCount(vjusage)
    if totalCount == 0:
        return vjusage

    for r in vjusage:
        row = [float(c)/totalCount for c in r]
        m.append(row)
    return m

def setUsageDistAxes( fig, numSamples, samplesPerPlot ):
    numaxes = numSamples / samplesPerPlot
    if numSamples % samplesPerPlot != 0:
        numaxes += 1

    axesList = []
    axleft = 0.1
    axright = 0.95
    axwidth = axright - axleft
    axbottom = 0.15
    axtop = 0.92
    axheight = axtop - axbottom
    margin = 0.1

    h = ( axheight - ( margin * (numaxes - 1) ) )/float( numaxes )

    bottom = axtop - h
    for i in range( numaxes ):#add axes from top down
        axesList.append( fig.add_axes( [axleft, bottom, axwidth, h] ) )
        bottom = bottom - (margin + h)

    return axesList

def getGeneNumber(genename):
    items = re.split('\D+', genename)
    nums = []
    for i in items:
        if i != '':
            nums.append( int(i) )
    return nums

def filterGenes(genes, keepGenes):
    #Get the intersection of genes & keepGenes:
    intersectGenes = []
    for g in keepGenes:
        if g in genes:
            intersectGenes.append(g)
    return intersectGenes

def getSelectedGenes(type2genes, vs, js):
    newtype2genes = {}
    types = {'v': vs, 'j':js}
    for type, selectedGenes in types.iteritems():
        if selectedGenes:
            genes = filterGenes(type2genes[type], selectedGenes)
            newtype2genes[type] = genes
        else:
            newtype2genes[type] = type2genes[type]
    return newtype2genes

def getIntersect(samples):
    type2genes = {}
    types=['v', 'j']
    for type in types:
        genes = getIntersectGenes(samples, type)
        genes = sorted( genes, key=lambda g: getGeneNumber(g) )
        type2genes[type] = genes
        #type2genes[type] = sorted( getIntersectGenes(samples, type) )
    return type2genes

def getIntersectGenes(samples, genetype):
    if genetype not in ['v', 'j']:
        raise ValueError("genetype must be either 'v' or 'j'\n")
    if len(samples) == 0:
        raise ValueError("len(samples) = 0! (zero sample!)")
    genes = samples[0].vgenes
    filterGenes = []
    if genetype == 'j':
        genes = samples[0].jgenes

    for i in xrange(1, len(samples)):#gene must be in all other samples, otherwise remove from list
        s = samples[i]
        for j,g in enumerate(genes):
            if (genetype == 'j' and g not in s.jgenes) or (genetype == 'v' and g not in s.vgenes):
                if g not in filterGenes:
                    filterGenes.append(g)
    intersectGenes = []
    for g in genes:
        if g not in filterGenes:
            intersectGenes.append(g)
    return intersectGenes

def sortGenes(genes, samples, genetype):
    if genetype not in ['v', 'j']:
        raise KeyError("genetype must be either 'v' or 'j'\n")
    sortedgenes = []
    gc = []
    if genetype == 'v':
        for g in genes:
            c = 0
            for sample in samples:
                for i, sg in enumerate(sample.vgenes):
                    if sg == g:
                        c += sample.vcounts[i]
                        break
            gc.append( (g, c) )
    else:
        for g in genes:
            c = 0
            for sample in samples:
                for i, sg in enumerate(sample.jgenes):
                    if sg == g:
                        c += sample.jcounts[i]
                        break
            gc.append( (g, c) )
    #if genetype == 'v':
    #    for i,g in enumerate(sample.vgenes):
    #        gc.append( (g, sample.vcounts[i]) )
    #else:
    #    for i,g in enumerate(sample.jgenes):
    #        gc.append( (g, sample.jcounts[i]) )
    sortedgc = sorted(gc, key=lambda item:item[1], reverse=True)
    sortedgenes = [item[0] for item in sortedgc]
    #for item in sortedgc:
    #    if item[0] in genes:
    #        sortedgenes.append( item[0] )
    return sortedgenes

def getStdData(std, sample, genetype, genes, abs):
    stdData = getGeneUsage(std, genetype, genes, abs)
    ydata = getGeneUsage(sample, genetype, genes, abs)
    total = sum(ydata)
    if len(stdData) == 0 or total == 0:
        return [], []
    higherYdata = [ float(y + stdData[i])/total for i, y in enumerate(ydata) ]
    lowerYdata = [ float(y - stdData[i])/total for i, y in enumerate(ydata) ]

    return higherYdata, lowerYdata

def getGeneUsage( sample, genetype, genes, abs ):
    if genetype == "v":
        return sample.getVfreqs2(genes, abs)
    elif genetype == "j":
        return sample.getJfreqs2(genes, abs)
    else:
        sys.stderr.write( "Genetype '%s' is unrecognized. Genetype can only be 'v' or 'j'.\n" %genetype )
        sys.exit(1)

#def getGeneUsage( sample, genetype ):
#    if genetype == "v":
#        return sample.getVfreqs()
#    elif genetype == "j":
#        return sample.getJfreqs()
#    else:
#        sys.stderr.write( "Genetype '%s' is unrecognized. Genetype can only be 'v' or 'j'.\n" %genetype )
#        sys.exit(1)

#AS color:
def getColors3():
    colors = [ "#00611C", "#0E8C3A", #green
               "#00CD00", "#00FF66", # bright green
               "#551A8B", "#72587F", # purple
               "#B6316C", "#CD6889", # cranberry
               "#FFFF33", "#FFFFB3", #yellow
               "#EE8833", "#EEC591", # orange
               "#00008B", "#4372AA", #blue
               "#0147FA", "#0198E1", # bright blue
               "#984EA3", "#F6BDFE", #violet
               "#8B4513", "#8B5742", # brown
               "#BC8F8F", "#C76114", #
               "#AA6600", "#B5A642",
               "#E31A1C", "#FE8E8F",
               "#984EA3", "#F6BDFE",
               ] #
    return colors

def drawUsageData( axesList, samples, stds, options, genetype, intersectGenes ):
    if len( samples ) <= 0:
        return
    samplesPerPlot = options.samplesPerPlot

    #HACK COLORS
    #colors = []
    #if options.sampleColor:
    #    colors = options.sampleColor
    #if options.sampleMarker:
    #    markers = options.sampleMarker

    #if options.sampleOrder and not options.sampleColor:
    #    colors = [ "#E31A1C", "#FE8E8F"]*len(options.sampleOrder) + ["#377EB8", "#A6D7FE"]*(len(samples) - len(options.sampleOrder))
    #    markers = ["o"]*len(options.sampleOrder) + ["^"]*(len(samples) - len(options.sampleOrder))
    #if not options.groupDrawInfo and len(colors) < len(samples):
    #    raise ValueError("Too many samples (%d), not enough colors (%d). Please specify the options sampleColor and sampleMarker to fix the problem." %(len(samples), len(colors)))
    #END HACK COLORS

    c = -2
    textsize = 'large'
    fontP = FontProperties()
    fontP.set_size(textsize)
    linesDict = {}
    labelsDict = {}

    xtickLabels = intersectGenes

    #get x location
    xdata = range( 0, len(xtickLabels), 1 )
    maxy = 0

    #HACK COLORS:
    s2group = {'as1D': 'B27+,AS', 'as8D': 'B27+,AS', 'as11D': 'B27+,AS', 'as16D': 'B27+,AS',
               'as15D': 'B27-,AS',
               'asBD': 'B27+,Healthy', 'as20D': 'B27+,Healthy',
               'adaptMA': 'B27-,Healthy', 'adaptM35': 'B27-,Healthy', 'adaptF57': 'B27-,Healthy', 'adaptAS': 'B27-,Healthy', 'adaptCD': 'B27-,Healthy', 'adaptF28':'B27-,Healthy',
               'irep1D': 'B27+,AS', 'as10R': 'B27+,AS', 'as11R': 'B27+,AS', 'as12R': 'B27+,AS', 'as13R': 'B27+,AS', 'as1R': 'B27+,AS', 'irepBD': 'B27+,Healthy', 'asBR': 'B27+,Healthy',
               'sameGroupShared': 'AS', 'diffGroupShared': 'Healthy', 'uniq': 'Unique',
               'b27pos': 'B27+', 'b27neg':'B27-'}

    #B27 +-
    #g2c = {'B27+,AS': '#E31A1C', 'B27+,Healthy': '#E31A1C', 'B27-,AS':'#377EB8', 'B27-,Healthy':'#377EB8', 'B27+':'#E31A1C', 'B27-': '#377EB8'}
    #g2lc = {'B27+,AS': '#FE8E8F', 'B27+,Healthy': '#FE8E8F', 'B27-,AS':'#A6D7FE', 'B27-,Healthy':'#A6D7FE', 'B27+': '#FE8E8F', 'B27-': '#A6D7FE'}
    #g2m = {'B27+,AS': 'o', 'B27+,Healthy': '^', 'B27-,AS':'o', 'B27-,Healthy':'^', 'B27+':'o', 'B27-':'^'}

    #AS vs Healthy
    #g2c = {'B27+,AS': '#E31A1C', 'B27+,Healthy': '#9E1114', 'AS':'#377EB8', 'Healthy':'#275880', 'Unique': '#4DAF4A'}
    #g2lc = {'B27+,AS': '#FE8E8F', 'B27+,Healthy': '#FE8E8F', 'AS':'#A6D7FE', 'Healthy':'#A6D7FE', 'Unique': '#B8FEB5'}
    #g2m = {'B27+,AS': 'o', 'B27+,Healthy': '^', 'AS':'o', 'Healthy':'^', 'Unique': 's'}

    #AS vs Healthy, no color (gray scale)
    g2c = {'B27+,AS': '#353535', 'B27+,Healthy': '#848484', 'AS':'#353535', 'Healthy':'#848484', 'Unique': '#4DAF4A'}
    g2lc = {'B27+,AS': '#BDBDBD', 'B27+,Healthy': '#BDBDBD', 'AS':'#BDBDBD', 'Healthy':'#BDBDBD', 'Unique': '#B8FEB5'}
    g2m = {'B27+,AS': 'o', 'B27+,Healthy': '^', 'AS':'o', 'Healthy':'^', 'Unique': 's'}
    #END HACK COLORS

    if options.groupDrawInfo:
        s2group = options.s2g
        g2c = options.g2c
        g2lc = options.g2lc
        g2m = options.g2m

    for i in range( len(axesList) ):
        lines = []
        sampleNames = []
        axes = axesList[i]
        startIndex = i*samplesPerPlot
        endIndex = min( [startIndex + samplesPerPlot, len(samples)] )
        groups = []
        for j in range( startIndex, endIndex ):
            sample = samples[j]
            sampleNames.append( "%s" % (sample.name))
            ydata = getGeneUsage( sample, genetype, xtickLabels, options.abs )
            maxy = max( [maxy, max(ydata)] )

            #c += 2
            #l = axes.plot( xdata, ydata, color=colors[c], marker=markers[c/2], markersize=8.0, markeredgecolor=colors[c], linestyle='none' )
            #axes.plot( xdata, ydata, color=colors[c +1], linestyle='-', linewidth=0.01 )
            #lines.append( l )

            group = s2group[sample.name]
            l = axes.plot( xdata, ydata, color=g2c[group], marker=g2m[group], markersize=8.0, markeredgecolor=g2c[group], linestyle='none' )
            axes.plot( xdata, ydata, color=g2lc[group], linestyle='-', linewidth=0.01 )
            if group not in groups:
                lines.append( l )
                groups.append(group)

            #Draw standard deviation if the data is available:
            if sample.name in options.avr2std and options.avr2std[sample.name] in stds:
                std = stds[ options.avr2std[sample.name] ]
                higherYdata, lowerYdata = getStdData( std, sample, genetype, xtickLabels, options.abs )
                if len(higherYdata) == 0 or len(lowerYdata) == 0:
                    continue
                for i in xrange( len(ydata) ):
                    #axes.plot( [xdata[i], xdata[i]], [higherYdata[i], lowerYdata[i]], linestyle='-', marker='_', markersize=8.0, linewidth=0.02, color=colors[c+1] )
                    axes.plot( [xdata[i], xdata[i]], [higherYdata[i], lowerYdata[i]], linestyle='-', marker='_', markersize=8.0, linewidth=0.02, color=g2lc[group] )
                maxy = max([maxy, max(higherYdata)])

        linesDict[i] = lines
        #labelsDict[i] = sampleNames
        labelsDict[i] = groups
        #axes.set_title( 'TRB%s Usage Distribution' % genetype.upper(), size='xx-large', weight='bold' )

    #HACK
    for i, xlabel in enumerate(xtickLabels):
        if xlabel == 'TRBV6-5/TRBV6-6':
            xtickLabels[i] = 'TRBV6-5/6-6'
    xtickLabels = [xlabel.lstrip('TRB') for xlabel in xtickLabels]
    for i, xlabel in enumerate(xtickLabels):
        items = xlabel.split('|')
        items = [item.lstrip('TRB') for item in items]
        xtickLabels[i] = '|'.join(items)
    #END HACK

    for i in range( len(axesList) ):
        axes = axesList[ i ]
        immunoseqLib.editSpine( axes )
        axes.set_xlabel('Gene', size = 'xx-large', weight='bold')
        axes.set_ylabel('Frequency (% of total sequences)', size='xx-large', weight='bold')

        #Legend
        if not options.std:
            #HACK
            cat2index = {}
            for j, label in enumerate(labelsDict[i]):
                cat2index[label] = j
            labels = cat2index.keys()
            lines = [linesDict[i][cat2index[label]] for label in labels]
            legend = pyplot.legend( lines, labels, numpoints=1, prop=fontP, loc="best" )
            #END HACK

            #legend = pyplot.legend( linesDict[i], labelsDict[i], numpoints=1, prop=fontP, loc="best" )
            #legend = axes.legend( linesDict[i], labelsDict[i], numpoints = 1, "upper right", ncol=3)
            #for t in legend.get_texts():
            #    t.set_fontsize('xx-small')
            legend._drawFrame = False

        axes.xaxis.set_ticklabels( xtickLabels )
        axes.xaxis.set_ticks( xdata )
        axes.set_xlim(-0.5, len(xtickLabels) + 0.5)

        numTicks = 20
        yticks = [ float(t)/numTicks for t in range(numTicks +1) ]
        ytickLabels = []
        for y in yticks:
            ytickLabels.append( "%d" %(y*100) )
        axes.yaxis.set_ticklabels( ytickLabels )
        axes.yaxis.set_ticks( yticks )
        axes.set_ylim(-0.01, maxy + 0.01)
        #axes.set_ylim(-0.01, 0.22)
        for label in axes.get_xticklabels():
            label.set_fontsize( 'small' )
            label.set_fontweight( 'bold' )
            label.set_rotation( 75 )
        for label in axes.get_yticklabels():
            label.set_fontsize( textsize )
            label.set_fontweight( 'bold' )
        axes.yaxis.grid(b=True, color="#CCCCCC", linestyle='-', linewidth=0.005)
        axes.xaxis.grid(b=True, color="#CCCCCC", linestyle='-', linewidth=0.005)
    return

def drawUsageDist( samples, stds, options, genetype, intersectGenes ):
    #Draw V usage distribution of each sample onto the same pdf
    options.out = os.path.join( options.outdir, "%sUsage" %genetype)
    fig, pdf = immunoseqLib.initImage( 10.0, 12.0, options )
    axesList = setUsageDistAxes( fig, len(samples), options.samplesPerPlot )
    drawUsageData( axesList, samples, stds, options, genetype, intersectGenes )
    immunoseqLib.writeImage( fig, pdf, options )

################# TWO-SAMPLE STUDENT'S T TEST TO SEE IF THERE IS BIAS AMONG DIFFERENT GROUPS ######################
#In this section, for each genetype (V, J, or VJ), we will calculate the Pearson Correlation of the relative V usage (frequencies instead of raw counts) of all pairs of samples that are WITHIN the same technology and all pairs of samples from two DIFFERENT technologies. And do a t test to see whether to accept for reject the null hypothesis (there is no bias among the technologies, i.e the mean PearsonCorrelations of pairs with same tech is the same with the mean PearsonCorrelation of pairs with different techs.
def readGroup2samples(file):
    group2samples = {}
    if not file:
        return group2samples
    f = open(file, 'r')
    for line in f:
        items = line.strip().split()
        group2samples[items[0]] = items[1].split(',')
    f.close()
    return group2samples

#========== ttest for each gene (targetGroup has higher/smaller usage than otherGroup) ============
def getGene2countVorJ(sample, genetype, genes, abs):
    gene2count = {}
    if genetype == 'v':
        for i, gene in enumerate(sample.vgenes):
            if gene in genes:
                gene2count[gene] = sample.vcounts[i]
                if not abs:
                    gene2count[gene] /= sample.totalCount
    else:
        for i, gene in enumerate(sample.jgenes):
            if gene in genes:
                gene2count[gene] = sample.jcounts[i]
                if not abs:
                    gene2count[gene] /= sample.totalCount
    if len(gene2count.keys()) != len(genes):
        raise ValueError("Some gene in the gene list %s are not present in sample genes %s\n" %(','.join(genes), ','.join(sample.vgenes)))
    return gene2count

def getGene2countVJ(sample, selectedvjs, vs, js, abs):
    gene2count = {}
    if abs:
        vjusage = sample.intersectVJusage
    else:
        vjusage = sample.normIntersectVJusage
    for i, r in  enumerate( vjusage ) :
        for j, c in enumerate(r):
            vname = vs[i]
            jname = js[j]
            vjname = "%s|%s" %(vname, jname)
            if not selectedvjs or vjname in selectedvjs:
                gene2count[vjname] = c
    return gene2count

def pairTtestSingleGene(samples, group2samples, genetype, type2intersectGenes, selectedvjs, targetGroup, abs):
    '''ttest for each single gene
    '''
    #genes = getIntersectGenes(samples, genetype)
    gene2results = {} #key = gene, val = [tval, pval, meanGroup1, stdGroup1, meanGroup2, stdGroup2]
    gene2counts1 = {} #key = gene, val = list of counts of targetGroup's samples
    gene2counts2 = {} #key = gene, val = list of counts of other samples (non-targetGroup)

    targetSamples = group2samples[targetGroup]

    for sample in samples:
        if genetype == 'vj':
            gene2count = getGene2countVJ(sample, selectedvjs, type2intersectGenes['v'], type2intersectGenes['j'], abs)
        else:
            gene2count = getGene2countVorJ(sample, genetype, type2intersectGenes[genetype], abs)
        if sample.name in targetSamples:
            for gene, count in gene2count.iteritems():
                if gene not in gene2counts1:
                    gene2counts1[gene] = [count]
                else:
                    gene2counts1[gene].append(count)
        else:
            for gene, count in gene2count.iteritems():
                if gene not in gene2counts2:
                    gene2counts2[gene] = [count]
                else:
                    gene2counts2[gene].append(count)

    #t-test:
    for gene, vec1 in gene2counts1.iteritems():
        vec2 = gene2counts2[gene]
        tval, pval = ttest_ind(vec1, vec2)
        gene2results[gene] = [tval, pval, np.mean(vec1), np.std(vec1), np.mean(vec2), np.std(vec2)]
    return gene2results

#=========== ttests of correlation of gene usage among samples (hypothesis: same group has higher correlation) =================
def getCorrVorJ(sample1, sample2, genetype, genes, abs):
    #abs = False
    vec1 = getGeneUsage(sample1, genetype, genes, abs)
    vec2 = getGeneUsage(sample2, genetype, genes, abs)
    corr, pval = pearsonr(vec1, vec2)
    return corr

def getVJusage(sample, selectedvjs, vs, js, abs):
    #vs = rownames, js = column names of sample.normIntersectVJusage
    vec = []
    if abs:
        vjusage = sample.intersectVJusage
    else:
        vjusage = sample.normIntersectVJusage
    #for i, r in  enumerate( sample.normIntersectVJusage ) :
    for i, r in  enumerate( vjusage ) :
        for j, c in enumerate(r):
            vname = vs[i]
            jname = js[j]
            vjname = "%s|%s" %(vname, jname)
            if not selectedvjs or vjname in selectedvjs:
                vec.append(c)
    return vec

def getCorrVJ(sample1, sample2, type2intersectGenes, selectedvjs, abs):
    vs = type2intersectGenes['v']
    js = type2intersectGenes['j']
    vec1 = getVJusage(sample1, selectedvjs, vs, js, abs)
    vec2 = getVJusage(sample2, selectedvjs, vs, js, abs)
    corr, pval = pearsonr(vec1, vec2)
    return corr

def getCorr(sample1, sample2, genetype, type2intersectGenes, selectedvjs, abs):
    if genetype == 'vj':
        return getCorrVJ(sample1, sample2, type2intersectGenes, selectedvjs, abs)
    else:
        return getCorrVorJ(sample1, sample2, genetype, type2intersectGenes[genetype], abs)

def getSampleByName(samples, name):
    for s in samples:
        if s.name == name:
            return s
    return None

def pairTtest(samples, group2samples, genetype, type2intersectGenes, selectedvjs, targetGroup, abs):
    #genes = getIntersectGenes(samples, genetype)
    samePairs = [] #all pairs of samples in the targetGroup
    diffPairs = [] #all pairs of samples in the other groups (all intra-group pairs, no inter-group)

    #Same pairs:
    for g, gsamples in group2samples.iteritems():
        for i in xrange( len(gsamples) - 1 ):
            s1 = getSampleByName(samples, gsamples[i])
            if not s1:
                print [s.name for s in samples]
                raise KeyError("Could not found sample %s in the sample list\n" %gsamples[i])
            for j in xrange( i+1, len(gsamples) ):
                s2 = getSampleByName(samples, gsamples[j])
                if not s2:
                    print [s.name for s in samples]
                    raise KeyError("Could not found sample %s in the sample list\n" %gsamples[j])
                if targetGroup and targetGroup == g:
                    samePairs.append( getCorr(s1, s2, genetype, type2intersectGenes, selectedvjs, abs) )
                else:
                    diffPairs.append( getCorr(s1, s2, genetype, type2intersectGenes, selectedvjs, abs) )

    #t-test:
    tval, pval = ttest_ind(samePairs, diffPairs)
    return tval, pval, np.mean(samePairs), np.std(samePairs), np.mean(diffPairs), np.std(diffPairs)

def pairTtest_intraVsInterGroups(samples, group2samples, genetype, type2intersectGenes, selectedvjs, targetGroup, abs):
    #genes = getIntersectGenes(samples, genetype)
    samePairs = []
    diffPairs = []

    #Same pairs:
    for g, gsamples in group2samples.iteritems():
        issame = True
        if targetGroup and g != targetGroup: #if targetGroup is specified, samePairs includes pairs of the target group only
            issame = False
        for i in xrange( len(gsamples) - 1 ):
            s1 = getSampleByName(samples, gsamples[i])
            if not s1:
                print [s.name for s in samples]
                raise KeyError("Could not found sample %s in the sample list\n" %gsamples[i])
            for j in xrange( i+1, len(gsamples) ):
                s2 = getSampleByName(samples, gsamples[j])
                if not s2:
                    print [s.name for s in samples]
                    raise KeyError("Could not found sample %s in the sample list\n" %gsamples[j])
                if issame:
                    samePairs.append( getCorr(s1, s2, genetype, type2intersectGenes, selectedvjs, abs) )
                #else:
                #    diffPairs.append( getCorr(s1, s2, genetype, type2intersectGenes, selectedvjs, abs) )
    #Diff pairs:
    groups = group2samples.keys()
    for i in xrange( len(groups) - 1 ):
        for gs1 in group2samples[ groups[i] ]:
            s1 = getSampleByName(samples, gs1)
            for j in xrange( i+1, len(groups) ):
                for gs2 in group2samples[ groups[j] ]:
                    s2 = getSampleByName(samples, gs2)
                    diffPairs.append( getCorr(s1, s2, genetype, type2intersectGenes, selectedvjs, abs) )

    #t-test:
    tval, pval = ttest_ind(samePairs, diffPairs)
    return tval, pval, np.mean(samePairs), np.std(samePairs), np.mean(diffPairs), np.std(diffPairs)

def ttests(samples, group2samples, outfile, type2intersectGenes, selectedvjs, targetGroup, abs, pvalcutoff):
    #group2samples = readGroup2samples(group2samplesFile)
    genetypes = ['v', 'j', 'vj']
    f = open(outfile, 'w')
    f.write("Genetype\tPval\tSameMean\tSameStd\tDiffMean\tDiffStd\n")
    for genetype in genetypes:
        tval, pval, sameMean, sameStd, diffMean, diffStd = pairTtest(samples, group2samples, genetype, type2intersectGenes, selectedvjs, targetGroup, abs)
        f.write("%s\t%f\t%f\t%f\t%f\t%f\n" %(genetype, pval, sameMean, sameStd, diffMean, diffStd) )

    #Separate genes:
    if targetGroup:
        for genetype in genetypes:
            f.write("#\n#%s\n" %genetype)
            gene2results = pairTtestSingleGene(samples, group2samples, genetype, type2intersectGenes, selectedvjs, targetGroup, abs)

            pfile = os.path.join( os.path.split(outfile)[0], "%s-pvals.txt" %genetype) #also print all pvalues of (separate) ttests of all genes:
            pfh = open( pfile, 'w' )
            for gene, results in gene2results.iteritems():
                pval = results[1]
                if pval <= pvalcutoff:
                    f.write("%s\t%s\n" %(gene, '\t'.join([str(c) for c in results[1:]])) )
                #Hack:
                #pfh.write("%s\t%f\n" %(gene, pval))
                maxMean = max([ results[2], results[4] ])
                #if (maxMean <= 1 and maxMean >= 0.01) or maxMean >= 1000 :
                if (maxMean <= 1 and maxMean >= 0.0001) or maxMean >= 1000 :
                    pfh.write("%s\t%s\n" %(gene, '\t'.join([str(c) for c in results[1:]])) )
            pfh.close()
    f.close()

############# DRAW VJ HEATMAPS #####################
def setCompareAxes( fig ):
    """
    Set axes for the VJheatmap Plot. There are 3 subplots total:
    Plot 0: Heatmap plot, where rows = Vs, cols = Js
    Plot 1: Right at the top of plot 0, shows J usage
    Plot 2: To the right of plot 1, shows V usage
    """
    axesList = []
    axleft = 0.1
    axright = 0.98
    axwidth = axright - axleft
    axbottom = 0.08
    axtop = 0.95
    axheight = axtop - axbottom
    margin = 0.07 #space between plots

    plot0wh = 0.7 #plot1 width = plot1 height
    plot1h = axheight - (plot0wh + margin)
    plot2w = plot1h

    axesList.append( fig.add_axes( [axleft, axbottom, plot0wh, plot0wh] ) ) #Plot 0
    axesList.append( fig.add_axes( [axleft, axbottom + plot0wh + margin, plot0wh, plot1h] ) ) #Plot 1
    axesList.append( fig.add_axes( [axleft + plot0wh + margin, axbottom, plot2w, plot0wh] ) ) #Plot 2
    return axesList

def getMinMaxVJusage(samples, abs):
    minvj = float('inf')
    maxvj = 0
    for sample in samples:
        if abs:
            data = sample.intersectVJusage
        else:
            data = sample.normIntersectVJusage
        for r in data:
            minvj = min([min(r), minvj])
            maxvj = max([max(r), maxvj])
    return minvj, maxvj

#def drawVJdata( axesList, sample, options ):
def drawVJdata( fig, axes, sample, vgenes, jgenes, options, minvj, maxvj ):
    #Draw heatmap:
    #data = sample.normalizeVJusage()
    if options.abs:
        data = sample.intersectVJusage
    else:
        data = sample.normIntersectVJusage

    #Normalize data to the range minvj-maxvj:
    if options.heatmapNoScale:
        hmaxes = axes.imshow( data, interpolation='nearest' )
    else:
        norm = mpl.colors.Normalize(vmin=minvj, vmax=maxvj)
        cmap = mpl.cm.get_cmap('rainbow', 20)
        hmaxes = axes.imshow( data, interpolation='nearest', cmap=cmap, norm=norm )

    #Colorbar:
    #from mpl_toolkits.axes_grid1.inset_locator import inset_axes
    #axins = inset_axes(axes, width="5%", height="10%", loc=3, bbox_to_anchor=(1.05, 0, 1, 1), bbox_transform=axes.transAxes, borderpad=0)
    #matplotlib.pyplot.colorbar(hmaxes, cax=axins, ticks=[0, 0.5, 1])
    cbar = fig.colorbar(hmaxes, shrink=0.3)
    #cbar.ax.set_yticklabels(['0', '0.5', '1'])

    immunoseqLib.editSpine( axes )
    axes.set_title( "VJ usage of sample %s" % immunoseqLib.properName(sample.name) )
    xticks = [ x for x in range( len(jgenes) ) ]
    axes.xaxis.set_ticks( xticks )
    for i, xlabel in enumerate(jgenes):
        items = xlabel.split('|')
        items = [item.lstrip('TRB') for item in items]
        jgenes[i] = '|'.join(items)
    axes.xaxis.set_ticklabels( jgenes )

    yticks = [ y for y in range( len(vgenes) ) ]
    axes.yaxis.set_ticks( yticks )
    for i, ylabel in enumerate(vgenes):
        items = ylabel.split('|')
        items = [item.lstrip('TRB') for item in items]
        vgenes[i] = '|'.join(items)
    axes.yaxis.set_ticklabels( vgenes )

    textsize = 'x-small'
    for label in axes.get_xticklabels():
        label.set_fontsize( textsize )
        label.set_rotation( 80 )
    for label in axes.get_yticklabels():
        label.set_fontsize( textsize )

    #Draw Jusage histogram:
    #p1axes = axesList[1]
    #p1axes.bar( range( len(sample.jcounts) ), sample.jcounts )

    #Draw Vusage histogram:
    #p2axes = axesList[2]
    #bottoms = range( len(sample.vgenes) )
    #p2axes.barh( bottoms, sample.vcounts )

def drawVJ( sample, vgenes, jgenes, options, minvj, maxvj ):
    options.out = os.path.join( options.outdir, "%s-vjUsage" %sample.name )
    fig, pdf = immunoseqLib.initImage( 8.0, 10.0, options )
    axes = immunoseqLib.setAxes(fig)
    #axesList = setCompareAxes( fig )

    drawVJdata( fig, axes, sample, vgenes, jgenes, options, minvj, maxvj )
    immunoseqLib.writeImage( fig, pdf, options )
#################### DONE DRAWING VJ HEATMAPS ##################

################### Principal Component Analysis PCA ###########
def getSample2group(group2samples):
    s2g = {}
    for g, samples in group2samples.iteritems():
        for s in samples:
            s2g[s] = g
    return s2g

def drawPcaData(axes, rownames, transformedM, options):
    group2xdata = {}
    group2ydata = {}

    for i,r in enumerate(transformedM):
        sample = rownames[i][0]
        group = rownames[i][1]
        x = r[0]
        y = r[1]
        if group not in group2xdata:
            group2xdata[group] = [x]
            group2ydata[group] = [y]
        else:
            group2xdata[group].append(x)
            group2ydata[group].append(y)

    groups = sorted(group2xdata.keys())
    if options.groupDrawInfo:
        g2c = options.g2c
        g2m = options.g2m
        colors = [g2c[g] for g in groups]
        markers = [g2m[g] for g in groups]
    else:
        colors = immunoseqLib.getColors6()
        markers = ['^', 'o', 'd', 'p', 'v', '*', 's']

    lines = []
    minx = float('inf')
    maxx = float('-inf')
    miny = float('inf')
    maxy = float('-inf')
    for i, group in enumerate(groups):
        if i >= len(colors):
            raise ValueError("drawPcaData: Need more color!")
        color = colors[i]
        marker = markers[i]
        xdata = group2xdata[group]
        ydata = group2ydata[group]

        minx = min(minx, min(xdata))
        miny = min(miny, min(ydata))
        maxx = max(maxx, max(xdata))
        maxy = max(maxy, max(ydata))

        l = axes.plot(xdata, ydata, color=color, marker=marker, markersize=15.0, markeredgecolor=color, linestyle='none' )
        lines.append(l)

    for label in axes.get_xticklabels():
        label.set_fontsize( 'large' )
        label.set_fontweight( 'bold' )
    for label in axes.get_yticklabels():
        label.set_fontsize( 'large' )
        label.set_fontweight( 'bold' )

    rangex = maxx - minx
    axes.set_xlim(minx - rangex*0.1, maxx + rangex*0.1)
    rangey = maxy - miny
    axes.set_ylim(miny - rangey*0.1, maxy + rangey*0.1 )

    linenames = []
    for g in groups:
        linenames.append( immunoseqLib.properName(g) )
    #legend = pyplot.legend( lines, sorted(group2xdata.keys()), numpoints=1, loc="lower left" )
    legend = pyplot.legend( lines, linenames, numpoints=1, loc="best" )
    #legend._drawFrame = False
    #axes.set_title( 'PCA', size='xx-large', weight='bold' )
    axes.set_xlabel('PC1', size='xx-large', weight='bold')
    axes.set_ylabel('PC2', size='xx-large', weight='bold')
    axes.yaxis.grid(b=True, color="#3F3F3F", linestyle='-', linewidth=0.5)
    axes.xaxis.grid(b=True, color="#3F3F3F", linestyle='-', linewidth=0.5)

def drawPca(rownames, transformedM, outfile, options):
    #Draw V usage distribution of each sample onto the same pdf
    options.out = outfile
    fig, pdf = immunoseqLib.initImage( 12.5, 8.0, options )
    axes = immunoseqLib.setAxes(fig)
    drawPcaData( axes, rownames, transformedM, options )
    immunoseqLib.writeImage( fig, pdf, options )

def preparePcaMatrix(samples, group2samples, type2intersectGenes, selectedvjs, genetype, abs):
    m = []
    rownames = [] #sample name of each row
    #abs = False
    for g, names in group2samples.iteritems():
        for n in names:
            rownames.append( (n, g) )
            sample = getSampleByName(samples, n)
            if genetype == 'vj':
                rowdata = getVJusage(sample, selectedvjs, type2intersectGenes['v'], type2intersectGenes['j'], abs)
            else:
                rowdata = getGeneUsage( sample, genetype, type2intersectGenes[genetype], abs )
            m.append(rowdata)
    return np.array(m), rownames

def getPcaTransformedMatrix(samples, group2samples, type2intersectGenes, selectedvjs, genetype, abs, outfile, options):
    m, rownames = preparePcaMatrix(samples, group2samples, type2intersectGenes, selectedvjs, genetype, abs)
    transformedM = mdp.pca(m, output_dim=4)

    #Write to text file
    f = open("%s.txt" %outfile, 'w')
    for i,r in enumerate(transformedM):
        f.write("%s\t%s\t%s\n" %(rownames[i][0], rownames[i][1], "\t".join( [str(c) for c in r] )))
    f.close()

    #Draw plot:
    drawPca(rownames, transformedM, outfile, options)

################### END OF Principal Component Analysis PCA ###########

def readfiles( indir, noheader ):
    samples = []
    stds = {} #key = name, val = Sample

    for file in os.listdir( indir ):
        if os.path.isdir(file):
            continue
        #name = os.path.basename( file )
        #name = file.split('-')[0]
        filenameItems = file.split('-')
        name = '-'.join( filenameItems[:len(filenameItems) -1] )
        if name == '':
            name = file.split('.')[0]
        #name = name.rstrip('_uniq').rstrip('_abs')
        name = name.split('_')[0]
        sample = Sample( name )
        fh = open( os.path.join(indir, file), 'r' )
        m = []
        for line in fh.readlines():
            if re.search('/OR', line) or re.search('undefined', line):
                continue
            if len(line.strip()) == 0:
                continue
            m.append( line.strip().split('\t') )

        if not noheader and len(m) > 0:
            colnames = m.pop(0)
            rownames = []
            #for i,row in enumerate(m):
            i = 0
            while i < len(m):
                row = m[i]
                if len(row) != len(colnames) + 1:
                    m.remove(m[i])
                    continue

                rownames.append( m[i].pop(0) )
                i += 1
        else:
            rownames = [ i for i in range( len(m) ) ]
            colnames = []
            if len(rownames) > 0:
                rownames = [ j for j in range( len(m[0]) ) ]

        #Convert value to float:
        for i in range( len(m) ):
            m[i] = [ float(c) for c in m[i] ]

        sample.setUsage(m, rownames, colnames)
        sample.setVcounts()
        sample.setJcounts()
        sample.setTotalCount()
        if re.match('std', sample.name):
            stds[sample.name] = sample
        else:
            samples.append( sample )

    return samples, stds

def checkSamples( samples ):
    if len(samples) == 0:
        return
    vgenes = samples[0].vgenes
    jgenes = samples[0].jgenes
    for s in samples:
        if vgenes != s.vgenes:
            sys.stderr.write("Sample %s has different number of Vs compared with sample %s\n" %(s.name, samples[0].name) )
            sys.exit(1)
        elif jgenes != s.jgenes:
            sys.stderr.write("Sample %s has different number of Js compared with sample %s\n" %(s.name, samples[0].name) )
            sys.exit(1)

def avr2std(file):
    avr2std = {}
    f = open(file, 'r')
    for line in f:
        items = line.strip().split()
        avr2std[items[0]] = items[1]
    f.close()
    return avr2std

def getGroupDrawInfo(file):
    # Format: <groupname> <comma,sep,sample,list> <groupcolor> <groupLightColor> <groupmarker>
    s2g = {}
    g2c = {}
    g2lc = {}
    g2m = {}
    f = open(file, 'r')
    for line in f:
        if len(line) == 0 or line[0] == '#':
            continue
        items = line.strip().split()
        if len(items) < 5:
            raise ValueError("Wrong groupDrawInfo format. Expected 5 fields, only has %d fields. Line %s\n" %(len(items), line))

        group = items[0]
        samples = items[1].split(',')
        color = items[2]
        lightcolor = items[3]
        marker = items[4]
        for s in samples:
            s2g[s] = group
        g2c[group] = color
        g2lc[group] = lightcolor
        g2m[group] = marker
    f.close()
    return s2g, g2c, g2lc, g2m

def checkOptions( args, options, parser ):
    if not options.indir or not os.path.isdir( options.indir ):
        parser.error("Input directory is required. Please specify a valid input directory.\n")
    if not os.path.isdir( options.outdir ):
        parser.error("Output directory '%s' does not exist.\n" %options.outdir)
    if options.sampleOrder:
        options.sampleOrder = options.sampleOrder.split(',')
    if options.sampleColor:
        options.sampleColor = options.sampleColor.split(',')
    if options.sampleMarker:
        options.sampleMarker = options.sampleMarker.split(',')
    if options.groupDrawInfo:
        options.s2g , options.g2c, options.g2lc, options.g2m = getGroupDrawInfo(options.groupDrawInfo)

    options.avr2std = {}
    if options.std:
        options.avr2std = avr2std(options.std)
    sep = ','
    if options.vs:
        options.vs = immunoseqLib.readList(options.vs, sep)
    if options.js:
        options.js = immunoseqLib.readList(options.js, sep)
    if options.vjs:
        options.vjs = immunoseqLib.readList(options.vjs, sep)

def initOptions( parser ):
    parser.add_option('-i', '--indir', dest='indir', help="Required argument: Input directory")
    parser.add_option('-o', '--outdir', dest='outdir', default='.', help="Output directory. Default = current directory")
    parser.add_option('--noheader', dest='noheader', action="store_true", default=False, help="If NOT specified (default), it is assumed that the first row has the matrix column names (J genes), and first column has the row names (vgenes).")
    parser.add_option('-a', '--abs', dest='abs', default=False, action='store_true', help='If specified, will plot the absolute counts instead of pecentage')
    parser.add_option('--samplesPerPlot', dest='samplesPerPlot', type='int', default=15)
    parser.add_option('--std', dest='std', help='Required if wish to add standard deviation for the average usage. File must have the format: <averageUsageFile> <stdFile>. Default=None')
    parser.add_option('--group2samples', dest='group2samples', help='Format: <groupname> <Comma separated list of samples of this group>')
    parser.add_option('--ttest', dest='ttest', action = 'store_true', default=False, help='If specified, will perform ttest')
    parser.add_option('--ttestTargetGroup', dest='ttestTargetGroup', help='If specified, will compare samples belong to this group versus all other samples')
    parser.add_option('--vj', dest='vj', action = 'store_true', default=False, help='If specified, will draw VJ usage heatmap for each sample')
    parser.add_option('--vs', dest='vs', help='File containing a list of V genes of interest. If specified, only look at usage of these genes')
    parser.add_option('--js', dest='js', help='File containing a list of J genes of interest. If specified, only look at usage of these genes')
    parser.add_option('--vjs', dest='vjs', help='File containing a list of VJ gene combinations of interest. If specified, only look at usage of these genes')
    parser.add_option('-p', '--pval', dest='pval', default=1.0, type='float', help='pvalue cutoff. Only report ttests with pval <= this value')
    parser.add_option('--heatmapNoScale', dest='heatmapNoScale', action='store_true', default=False, help='Normally, the VJ heatmap plots are scale to the min and the max values of all samples. If this option is specified, heatmap of each sample is scale to its own range')
    parser.add_option('--noplot', dest='noplot', action='store_true', default=False, help='If specified, do not draw geneusage plot')
    parser.add_option('--pca', dest='pca', action='store_true', default=False, help='If specified, print out the transform pca matrices')
    parser.add_option('--sampleOrder', dest='sampleOrder', type='string', help='Optional. Comma seperated list of samples in the order of interest')
    parser.add_option('--sampleColor', dest='sampleColor', type='string', help='Optional. Comma seperated list of sample color in the order of sampleOrder. Default=None')
    parser.add_option('--sampleMarker', dest='sampleMarker', type='string', help='Optional. Comma seperated list of sample Marker in the order of sampleOrder. Default=None')
    parser.add_option('--groupDrawInfo', dest='groupDrawInfo', help='Optional. File with plot drawing instructions. Format: <groupname> <comma,sep,sample,list> <groupcolor> <groupLightColor> <groupmarker>')
    #parser.add_option('--vjPlotPerPage', dest='vjPlotPerPage', type='int', default=8)

def main():
    parser = immunoseqLib.initOptions()
    initOptions( parser )
    immunoseqLib.initPlotOptions( parser )
    options, args = parser.parse_args()
    checkOptions( args, options, parser )
    immunoseqLib.checkPlotOptions( options, parser )

    samples, stds = readfiles( options.indir, options.noheader )
    type2intersectGenes = getIntersect(samples) #Get genes that are present in all samples
    type2selectedGenes = getSelectedGenes(type2intersectGenes, options.vs, options.js) #Select only genes of interest that are present in all samples

    group2samples = readGroup2samples(options.group2samples)

    #Intersect VJ:
    if options.ttest or options.vj or options.pca:
        #intersectVJ(samples, type2intersectGenes, options.vjs) #added intersectVJusage, normIntersectVJusage
        intersectVJ(samples, type2selectedGenes) #added intersectVJusage, normIntersectVJusage

    #Pca:
    if options.pca:
        vPcaOut = os.path.join(options.outdir, "vPca")
        getPcaTransformedMatrix(samples, group2samples, type2intersectGenes, options.vjs, 'v', options.abs, vPcaOut, options)
        vjPcaOut = os.path.join(options.outdir, "vjPca")
        getPcaTransformedMatrix(samples, group2samples, type2intersectGenes, options.vjs, 'vj', options.abs, vjPcaOut, options)
    #return

    #Student's t test:
    if options.ttest:
        outfile = os.path.join(options.outdir, "ttest.txt")
        ttests(samples, group2samples, outfile, type2selectedGenes, options.vjs, options.ttestTargetGroup, options.abs, options.pval)

    #HACK: Sort samples:
    #sampleOrder = 'SBC1,SBC7,SBC3,SBC4,SBC5,SBC6,SBC2,SBC8,Patient-01-R,Patient-01-D,Patient-10,Patient-11,Patient-12,Patient-13,Patient-B-R,Patient-B-D,Patient-8-D,Patient-15-D'
    if options.sampleOrder:
        orders = options.sampleOrder
        sortedSamples = []
        for s in orders:
            for sample in samples:
                if sample.name == s:
                    sortedSamples.append(sample)
        for s in samples:
            if s not in sortedSamples:
                sortedSamples.append(s)
    else:
        sortedSamples = samples

    #checkSamples( samples )
    if not options.noplot:
        drawUsageDist( sortedSamples, stds, options, 'v', type2selectedGenes['v'] )
        drawUsageDist( sortedSamples, stds, options, 'j', type2selectedGenes['j'] )

    if options.vj:
        minvj, maxvj = getMinMaxVJusage(samples, options.abs)
        for sample in samples:
            drawVJ( sample, type2selectedGenes['v'], type2selectedGenes['j'], options, minvj, maxvj )

if __name__ == "__main__":
    main()

max_stars_repo_name: ngannguyen/immunoseq

max_issues_repo_name: ngannguyen/immunoseq

max_forks_repo_name: ngannguyen/immunoseq

reason: import numpy,from numpy,from scipy


text:
#
# Copyright (c) 2017 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import operator
import random
from enum import Enum
from typing import List, Tuple, Any

import numpy as np

from rl_coach.core_types import Transition
from rl_coach.memories.memory import MemoryGranularity
from rl_coach.memories.non_episodic.experience_replay import ExperienceReplayParameters, ExperienceReplay
from rl_coach.schedules import Schedule, ConstantSchedule

class PrioritizedExperienceReplayParameters(ExperienceReplayParameters):
    def __init__(self):
        super().__init__()
        self.max_size = (MemoryGranularity.Transitions, 1000000)
        self.alpha = 0.6
        self.beta = ConstantSchedule(0.4)
        self.epsilon = 1e-6

    @property
    def path(self):
        return 'rl_coach.memories.non_episodic.prioritized_experience_replay:PrioritizedExperienceReplay'

class SegmentTree(object):
    """
    A tree which can be used as a min/max heap or a sum tree
    Add or update item value - O(log N)
    Sampling an item - O(log N)
    """
    class Operation(Enum):
        MAX = {"operator": max, "initial_value": -float("inf")}
        MIN = {"operator": min, "initial_value": float("inf")}
        SUM = {"operator": operator.add, "initial_value": 0}

    def __init__(self, size: int, operation: Operation):
        self.next_leaf_idx_to_write = 0
        self.size = size
        if not (size > 0 and size & (size - 1) == 0):
            raise ValueError("A segment tree size must be a positive power of 2. The given size is {}".format(self.size))
        self.operation = operation
        self.tree = np.ones(2 * size - 1) * self.operation.value['initial_value']
        self.data = [None] * size

    def _propagate(self, node_idx: int) -> None:
        """
        Propagate an update of a node's value to its parent node
        :param node_idx: the index of the node that was updated
        :return: None
        """
        parent = (node_idx - 1) // 2

        self.tree[parent] = self.operation.value['operator'](self.tree[parent * 2 + 1], self.tree[parent * 2 + 2])

        if parent != 0:
            self._propagate(parent)

    def _retrieve(self, root_node_idx: int, val: float)-> int:
        """
        Retrieve the first node that has a value larger than val and is a child of the node at index idx
        :param root_node_idx: the index of the root node to search from
        :param val: the value to query for
        :return: the index of the resulting node
        """
        left = 2 * root_node_idx + 1
        right = left + 1

        if left >= len(self.tree):
            return root_node_idx

        if val <= self.tree[left]:
            return self._retrieve(left, val)
        else:
            return self._retrieve(right, val-self.tree[left])

    def total_value(self) -> float:
        """
        Return the total value of the tree according to the tree operation. For SUM for example, this will return
        the total sum of the tree. for MIN, this will return the minimal value
        :return: the total value of the tree
        """
        return self.tree[0]

    def add(self, val: float, data: Any) -> None:
        """
        Add a new value to the tree with data assigned to it
        :param val: the new value to add to the tree
        :param data: the data that should be assigned to this value
        :return: None
        """
        self.data[self.next_leaf_idx_to_write] = data
        self.update(self.next_leaf_idx_to_write, val)

        self.next_leaf_idx_to_write += 1
        if self.next_leaf_idx_to_write >= self.size:
            self.next_leaf_idx_to_write = 0

    def update(self, leaf_idx: int, new_val: float) -> None:
        """
        Update the value of the node at index idx
        :param leaf_idx: the index of the node to update
        :param new_val: the new value of the node
        :return: None
        """
        node_idx = leaf_idx + self.size - 1
        if not 0 <= node_idx < len(self.tree):
            raise ValueError("The given left index ({}) can not be found in the tree. The available leaves are: 0-{}"
                             .format(leaf_idx, self.size - 1))

        self.tree[node_idx] = new_val
        self._propagate(node_idx)

    def get(self, val: float) -> Tuple[int, float, Any]:
        """
        Given a value between 0 and the tree sum, return the object which this value is in it's range.
        For example, if we have 3 leaves: 10, 20, 30, and val=35, this will return the 3rd leaf, by accumulating
        leaves by their order until getting to 35. This allows sampling leaves according to their proportional
        probability.
        :param val: a value within the range 0 and the tree sum
        :return: the index of the resulting leaf in the tree, it's probability and
                 the object itself
        """
        node_idx = self._retrieve(0, val)
        leaf_idx = node_idx - self.size + 1
        data_value = self.tree[node_idx]
        data = self.data[leaf_idx]

        return leaf_idx, data_value, data

    def __str__(self):
        result = ""
        start = 0
        size = 1
        while size <= self.size:
            result += "{}\n".format(self.tree[start:(start + size)])
            start += size
            size *= 2
        return result

class PrioritizedExperienceReplay(ExperienceReplay):
    """
    This is the proportional sampling variant of the prioritized experience replay as described
    in https://arxiv.org/pdf/1511.05952.pdf.
    """
    def __init__(self, max_size: Tuple[MemoryGranularity, int], alpha: float=0.6, beta: Schedule=ConstantSchedule(0.4),
                 epsilon: float=1e-6, allow_duplicates_in_batch_sampling: bool=True):
        """
        :param max_size: the maximum number of transitions or episodes to hold in the memory
        :param alpha: the alpha prioritization coefficient
        :param beta: the beta parameter used for importance sampling
        :param epsilon: a small value added to the priority of each transition
        :param allow_duplicates_in_batch_sampling: allow having the same transition multiple times in a batch
        """
        if max_size[0] != MemoryGranularity.Transitions:
            raise ValueError("Prioritized Experience Replay currently only support setting the memory size in "
                             "transitions granularity.")
        self.power_of_2_size = 1
        while self.power_of_2_size < max_size[1]:
            self.power_of_2_size *= 2
        super().__init__((MemoryGranularity.Transitions, self.power_of_2_size), allow_duplicates_in_batch_sampling)
        self.sum_tree = SegmentTree(self.power_of_2_size, SegmentTree.Operation.SUM)
        self.min_tree = SegmentTree(self.power_of_2_size, SegmentTree.Operation.MIN)
        self.max_tree = SegmentTree(self.power_of_2_size, SegmentTree.Operation.MAX)
        self.alpha = alpha
        self.beta = beta
        self.epsilon = epsilon
        self.maximal_priority = 1.0

    def _update_priority(self, leaf_idx: int, error: float) -> None:
        """
        Update the priority of a given transition, using its index in the tree and its error
        :param leaf_idx: the index of the transition leaf in the tree
        :param error: the new error value
        :return: None
        """
        if error < 0:
            raise ValueError("The priorities must be non-negative values")
        priority = (error + self.epsilon)
        self.sum_tree.update(leaf_idx, priority ** self.alpha)
        self.min_tree.update(leaf_idx, priority ** self.alpha)
        self.max_tree.update(leaf_idx, priority)
        self.maximal_priority = self.max_tree.total_value()

    def update_priorities(self, indices: List[int], error_values: List[float]) -> None:
        """
        Update the priorities of a batch of transitions using their indices and their new TD error terms
        :param indices: the indices of the transitions to update
        :param error_values: the new error values
        :return: None
        """
        self.reader_writer_lock.lock_writing_and_reading()

        if len(indices) != len(error_values):
            raise ValueError("The number of indexes requested for update don't match the number of error values given")
        for transition_idx, error in zip(indices, error_values):
            self._update_priority(transition_idx, error)

        self.reader_writer_lock.release_writing_and_reading()

    def sample(self, size: int) -> List[Transition]:
        """
        Sample a batch of transitions form the replay buffer. If the requested size is larger than the number
        of samples available in the replay buffer then the batch will return empty.
        :param size: the size of the batch to sample
        :return: a batch (list) of selected transitions from the replay buffer
        """

        self.reader_writer_lock.lock_writing()

        if self.num_transitions() >= size:
            # split the tree leaves to equal segments and sample one transition from each segment
            batch = []
            segment_size = self.sum_tree.total_value() / size

            # get the maximum weight in the memory
            min_probability = self.min_tree.total_value() / self.sum_tree.total_value()  # min P(j) = min p^a / sum(p^a)
            max_weight = (min_probability * self.num_transitions()) ** -self.beta.current_value  # max wi

            # sample a batch
            for i in range(size):
                start_probability = segment_size * i
                end_probability = segment_size * (i + 1)

                # sample leaf and calculate its weight
                val = random.uniform(start_probability, end_probability)
                leaf_idx, priority, transition = self.sum_tree.get(val)
                priority /= self.sum_tree.total_value()   # P(j) = p^a / sum(p^a)
                weight = (self.num_transitions() * priority) ** -self.beta.current_value  # (N * P(j)) ^ -beta
                normalized_weight = weight / max_weight  # wj = ((N * P(j)) ^ -beta) / max wi

                transition.info['idx'] = leaf_idx
                transition.info['weight'] = normalized_weight

                batch.append(transition)

            self.beta.step()

        else:
            raise ValueError("The replay buffer cannot be sampled since there are not enough transitions yet. "
                             "There are currently {} transitions".format(self.num_transitions()))

        self.reader_writer_lock.release_writing()
        return batch

    def store(self, transition: Transition) -> None:
        """
        Store a new transition in the memory.
        :param transition: a transition to store
        :return: None
        """
        self.reader_writer_lock.lock_writing_and_reading()

        transition_priority = self.maximal_priority
        self.sum_tree.add(transition_priority ** self.alpha, transition)
        self.min_tree.add(transition_priority ** self.alpha, transition)
        self.max_tree.add(transition_priority, transition)
        super().store(transition, False)

        self.reader_writer_lock.release_writing_and_reading()

    def clean(self) -> None:
        """
        Clean the memory by removing all the episodes
        :return: None
        """
        self.reader_writer_lock.lock_writing_and_reading()

        super().clean(lock=False)
        self.sum_tree = SegmentTree(self.power_of_2_size, SegmentTree.Operation.SUM)
        self.min_tree = SegmentTree(self.power_of_2_size, SegmentTree.Operation.MIN)
        self.max_tree = SegmentTree(self.power_of_2_size, SegmentTree.Operation.MAX)

        self.reader_writer_lock.release_writing_and_reading()

ext: py

reason: import numpy


text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Make plots using using the *seaborn-paper* style
"""

# See: http://matplotlib.org/users/style_sheets.html

import numpy as np
import matplotlib.pyplot as plt

plt.style.use('seaborn-paper')  # <- SETUP STYLE

# Build datas ###############

x = np.arange(-10, 10, 0.01)
y = np.sin(2. * 2. * np.pi * x) * 1. / np.sqrt(2. * np.pi) * np.exp(-(x**2.)/2.)

# Plot data #################

fig = plt.figure()
ax = fig.add_subplot(111)

ax.plot(x, y, "-", label="Test")

# Set title and labels ######

ax.set_title("Test")
ax.set_xlabel("x")
ax.set_ylabel("f(x)")

# Set legend ################

ax.legend(loc='lower right')

# Save file #################

plt.savefig("plot_style_seaborn_paper.png")

# Plot ######################

plt.show()

max_stars_repo_name: jeremiedecock/snippets

max_issues_repo_name: jeremiedecock/snippets

max_forks_repo_name: jeremiedecock/snippets


text:
import os
import json
import string
from collections import defaultdict
import numpy as np
from scipy.spatial import distance
# import cPickle as pkl
# import pdb
import math

from params import save_features_path, data_format, vgg, weights_dir

# img_featuer_path = "/home/zhaoxin/chenminghai/mRNN-CR/image_features_mRNN/VGG_feat_mRNN_refine_dct_mscoco_2014.npy"
# annotation_path = '/home/zhaoxin/chenminghai/data/coco/annotations/captions_split_train2014.json'
# img_feature_path = 'E:\\User\\freelancer\\image_cap\\wr\\feats\\VGG_feat_mRNN_refine_dct_mscoco_2014.npy'
# img_feature_path = 'E:\\User\\freelancer\\image_cap\\features_coco\\all_npy_dict.npy'
# annotation_path = 'E:\\User\\freelancer\\image_captioning\\annotations\\captions_train2014.json'
# annotation_path = 'E:\\User\\freelancer\\image_cap\\features_coco\\train_captions_train2014.json'

print('Starting weight process on {} data using {} architecture.'.format(data_format, 'vgg' if vgg else 'inception'))

img_feature_path = os.path.join(save_features_path, 'all_npy_dict.npy')
annotation_path = os.path.join(save_features_path, 'train_captions_to_weigth.json')

sigma = 37.382034917914233

# with open(img_featuer_path) as f:
#     features = np.load(f).tolist()

with open(img_feature_path, 'rb') as f:
    features = np.load(f, encoding='bytes').tolist()

with open(annotation_path, 'rt') as f:
    data = json.load(f)

# type(features)
#
# for k, v in features.items():
#     break
#
# a = np.array([[1,1,1], [2,2,2], [3,3,3], [1,1,1]])
#
# distances = distance.cdist(a, a[:2], 'euclidean')

tag2img_list = defaultdict(list)
for x in data['annotations']:
    # tokens = str(x['caption']).lower().translate(str.maketrans(dict.fromkeys(string.punctuation))).strip().split()
    # tags = list(set(str(x['caption']).lower().translate(str.maketrans(string.punctuation, ' '*len(string.punctuation))).strip().split()))
    tags = list(set(str(x['caption']).lower().translate(str.maketrans('!"#$%&()*+.,-/:;=?@[\\]^_`{|}~ ', ' '*len('!"#$%&()*+.,-/:;=?@[\\]^_`{|}~ '))).strip().split()))
    # for tag in tags:
    #     tag2img_list[tag].append(x['image_id'])
    [tag2img_list[tag].append(x['image_id']) for tag in tags]

for x in tag2img_list:
    tag2img_list[x] = list(set(tag2img_list[x]))

tag2score_list_1 = {}
tag2score_list_2 = {}
num = 0

for tag in tag2img_list:
    num += 1
    img_list = tag2img_list[tag]
    print(tag, ": ", len(img_list), '-', num, '/', len(tag2img_list))

    #if tag == 'a':
    #    continue
    # a = features.reshape(1,-1)
    feature_list = np.array([features[x] for x in img_list])
    # feature_list = features.reshape(1,-1)
    tag2score_list_1[tag] = dict()
    tag2score_list_2[tag] = dict()
    k = 30000 if len(feature_list) < 30000 else 10000
    if len(feature_list) == 1:
        print('ehh:', feature_list)
        feature_list.reshape(1, -1)
    print(feature_list.shape)
    distances = distance.cdist(feature_list, feature_list[0:k], 'euclidean')
    for i in range(len(img_list)):
        print('Todo: ', len(img_list) - i)
        dists = [math.exp(-x/sigma) for x in distances[i]]
        dists = [math.exp(-x / sigma) for x in distances[0]]
        tag2score_list_1[tag][img_list[i]] = float(sum(dists))/len(dists)/math.exp(-1)
        d2 = [math.exp(-x/sigma) for x in distances[i] if x != 0]
        if d2 == []:
            d2 = [math.exp(-1)]
        tag2score_list_2[tag][img_list[i]] = np.mean(d2)/math.exp(-1)
    #for i in range(len(img_list)):
    #   distances = distance.cdist([feature_list[i]], feature_list, "euclidean").tolist()
    #   tag2score_list_1[tag][img_list[i]] = np.mean([math.exp(-x/sigma) for x in distances[0]])
    #    d2 = [math.exp(-x/sigma) for x in distances[0] if x!=0 ]
    #    if d2 == []:
    #        d2 = [math.exp(-1)]
    #    tag2score_list_2[tag][img_list[i]] = np.mean(d2)

    if num % 1000 == 0:
        print('Saving json.')
        with open(os.path.join(weights_dir, 'tag2score_list_1.json'), 'wt') as f:
            json.dump(tag2score_list_1, f)
        with open(os.path.join(weights_dir, 'tag2score_list_2.json'), 'wt') as f:
            json.dump(tag2score_list_2, f)

# pdb.set_trace()

with open(os.path.join(weights_dir, 'tag2score_list_1.json'), 'wt') as f:
    json.dump(tag2score_list_1, f)

with open(os.path.join(weights_dir, 'tag2score_list_2.json'), 'wt') as f:
    json.dump(tag2score_list_2, f)

reason: import numpy,from scipy


text:
# -*- coding: utf-8 -*-

import os, logging, argparse
import pandas as pd
import numpy as np
from time import time
from sklearn.utils import shuffle

def main(interim_path = "data/interim",
         processed_path = "data/processed"):

    """Cleans aggregated data and saves it to <processed_path>.

    Drops unnecessary columns, and construct new, interesting ones.
    Also drops rows, which has to high error rates during feature extraction.
    Then drops rows with NaNs or with outliers as values, like Score > 1."""

    # logging
    logger = logging.getLogger(__name__)

    # normalize paths
    interim_path = os.path.normpath(interim_path)
    logger.debug("Path to interim data normalized: {}"
                 .format(interim_path))
    processed_path = os.path.normpath(processed_path)
    logger.debug("Path to processed data normalized: {}"
                 .format(processed_path))

    # load aggregated_df
    aggregated_df = pd.read_pickle(os.path.join(interim_path,
                                                'aggregated_df.pkl'))
    logger.info("Loaded aggregated_df.pkl. Shape of df: {}"
                .format(aggregated_df.shape))

    #%% start cleaning
    start = time()

    # split aggregated_df
    teams_df, features_df = aggregated_df.iloc[:,:9], aggregated_df.iloc[:,9:]
    logger.info("agrgregated_df splitted to teams_df {} & features_df {}."
                .format(teams_df.shape, features_df.shape))

    # create multi-index
    # make features names unique for which both sum and mean are kept
    features_df.rename(mapper={('radon_mi', 'sum'):('radon_mi_sum', 'sum'),
                               ('radon_mi', 'mean'):('radon_mi_mean', 'mean')},
            axis=1, inplace=True)
    logger.info("Renamed col ('radon_mi', 'sum'): ('radon_mi_sum', 'sum').")
    logger.info("Renamed col ('radon_mi', 'mean'): ('radon_mi_mean', 'mean').")
    multi_index = pd.MultiIndex.from_tuples(features_df.columns,
                                            names=['feature', 'stat'])
    features_df.columns = multi_index
    logger.info("Created MultiIndex of features_df to split sums and means.")

    #%% cleaning columns

    # drop mean columns since mostly irrelevant
    X = features_df.drop('mean', axis=1, level=1)
    logger.info("Created X by dropping all 'mean' features. Shape: {}"
                .format(X.shape))

    # re-include relevant columns from means
    X = pd.concat([features_df[('radon_cc_mean', 'mean')],
                            features_df[('radon_mi_mean', 'mean')],
                            features_df[('radon_raw_is_error', 'mean')],
                            features_df[('radon_cc_is_error', 'mean')],
                            features_df[('radon_h_is_error', 'mean')],
                            features_df[('radon_mi_is_error', 'mean')],
                            features_df[('pylint_is_error', 'mean')],
                            X], axis=1)
    logger.info("Re-included 7 relevant 'mean' features. Shape: {}"
                .format(X.shape))

    # drop column (radon_cc_mean, sum) since it doesn't make sense
    X.drop([('radon_cc_mean', 'sum')], axis=1, inplace=True)
    logger.info("Dropped column (radon_cc_mean, sum). Shape: {}"
                .format(X.shape))

    # drop halstead metrics since they are not additive thus wrong
    X.drop([('radon_h_calculated_length', 'sum'),
            ('radon_h_volume', 'sum'),
            ('radon_h_difficulty', 'sum'),
            ('radon_h_effort', 'sum'),
            ('radon_h_time', 'sum'),
            ('radon_h_bugs', 'sum')],
            axis=1, inplace=True)
    logger.info("Dropped 6 wrong calculated Halstead metrics. Shape: {}"
                .format(X.shape))

    # construct correct halstead metrics
    # see: https://radon.readthedocs.io/en/latest/intro.html#halstead-metrics
    X[('radon_h_calculated_length', 'custom')] = X[('radon_h_h1', 'sum')] \
            * np.log2(X[('radon_h_h1', 'sum')]) \
            + X[('radon_h_h2', 'sum')] \
            * np.log2(X[('radon_h_h2', 'sum')])
    X[('radon_h_volume', 'custom')] = X[('radon_h_length', 'sum')] \
            * np.log2(X[('radon_h_vocabulary', 'sum')])
    X[('radon_h_difficulty', 'custom')] = X[('radon_h_h1', 'sum')] / 2 \
            * X[('radon_h_N2', 'sum')] / X[('radon_h_h2', 'sum')]
    X[('radon_h_effort', 'custom')] = X[('radon_h_difficulty', 'custom')] \
            * X[('radon_h_volume', 'custom')]
    # Following two Halstead metrics are skipped since propotional to other
    # X[('radon_h_time', 'custom')] = X[('radon_h_effort', 'custom')] / 18
    # X[('radon_h_bugs', 'custom')] = X[('radon_h_volume', 'custom')] / 3000
    logger.info("Re-constructed 4 relevant Halstead metrics. Shape: {}"
                .format(X.shape))

    # transforms features to become ratios of loc to avoid multi-collinearity
    # calculate maximum number for lines of codes, since inconsistent
    loc_max = pd.concat([
            X[('radon_raw_loc', 'sum')],
            X[('radon_raw_lloc', 'sum')],
            X[('radon_raw_sloc', 'sum')]
                    + X[('radon_raw_multi', 'sum')]
                    + X[('radon_raw_single_comments', 'sum')]
                    + X[('radon_raw_blank', 'sum')],
            X[('pylint_code', 'sum')]
                    + X[('pylint_docstring', 'sum')]
                    + X[('pylint_comment', 'sum')]
                    + X[('pylint_empty', 'sum')]], axis=1)
    loc_max = loc_max.max(axis=1)
    X.drop(columns=[('radon_raw_loc', 'sum'), ('radon_raw_lloc', 'sum')],
            inplace=True)
    logger.info("Dropped 'radon_raw_loc' and 'radon_raw_lloc'. Shape: {}"
                .format(X.shape))
    # create list of columns, which are correlated/dependent to max_loc
    dependent = []
    for col in X.columns:
        if col[1] == 'mean' or '_is_error' in col[0] or col[0] == 'loc_max':
            logger.debug("Skipped independent feature: {}".format(col))
            continue
        elif 'pylint' in col[0] or 'radon' in col[0]:
            dependent.append(col)
            logger.debug("Included dependent feature:  {}".format(col))
        elif 'uses_module_' in col[0]:
            # set every value > 0 to 1 because you only want indication
            X[col].where(X[col]==0, other=1, inplace=True)
            logger.debug("Skipped independent feature: {}".format(col))
            continue
        else:
            logger.error("{} not catched during creation of list 'dependent'."
                         .format(col))
    for col in dependent:
        X[(col[0]+'_ratio', 'ratio')] = X[col]/loc_max
        logger.debug("Created new ratio feature for: {}".format(col))
    X.drop(columns=dependent, inplace=True)
    logger.info("Transformed {} features to ratios of lines-of-code. Shape: {}"
                .format(len(dependent), X.shape))
    # create new column as log-transformed loc_max
    X[('loc_max_log', 'max')] = np.log(loc_max)
    logger.info("Added column 'loc_max_log'. Shape: {}"
                .format(X.shape))

    #%% cleaning rows

    # delete repos with too many errors during feature extraction
    X = X[X[('radon_raw_is_error', 'mean')] != 1]
    X = X[X[('radon_cc_is_error', 'mean')] != 1]
    X = X[X[('radon_h_is_error', 'mean')] != 1]
    X = X[X[('radon_mi_is_error', 'mean')] != 1]
    X = X[X[('pylint_is_error', 'mean')] != 1]
    logger.info("Dropped all rows with error rates = 1. Shape: {}"
                .format(X.shape))

    # drop columns of error rates
    X.drop(columns=[('radon_raw_is_error', 'mean'),
                     ('radon_cc_is_error', 'mean'),
                     ('radon_h_is_error', 'mean'),
                     ('radon_mi_is_error', 'mean'),
                     ('pylint_is_error', 'mean'),
                     ('radon_raw_is_error', 'sum'),
                     ('radon_cc_is_error', 'sum'),
                     ('radon_h_is_error', 'sum'),
                     ('radon_mi_is_error', 'sum'),
                     ('pylint_is_error', 'sum'),], inplace=True)
    logger.info("Dropped 10 is_error columns. Shape: {}"
                .format(X.shape))

    # drop second level of MultiIndex
    X.columns = X.columns.droplevel(level=1)

    # concatenate X with log-transfomated Ranking
    X = pd.concat(
            [np.log(pd.to_numeric(teams_df.Ranking)).rename('ranking_log'), X],
            join='inner', axis=1)
    logger.info("Concatenated log-transformed 'ranking_log'. Shape: {}"
                .format(X.shape))

    # drop repos with outliers
    X = X[pd.to_numeric(aggregated_df.Score) <= 1] # 16
    logger.info("Dropped rows with score > 1. Shape: {}"
                .format(X.shape))
    X = X[(4 <= X.loc_max_log) & (X.loc_max_log <= 10)] # 2
    logger.info("Dropped rows with lines of code > 10.000. Shape: {}"
                .format(X.shape))
    X = X[X.pylint_warning_ratio <= 1] # 3
    logger.info("Dropped rows with pylint_warning_ratio > 1. Shape: {}"
                .format(X.shape))
    X = X[X.radon_h_effort_ratio <= 1000] # 6
    logger.info("Dropped rows with radon_h_effort_ratio > 1000. Shape: {}"
                .format(X.shape))
    X = X[X.radon_h_difficulty_ratio <= .07] # 2
    logger.info("Dropped rows with radon_h_difficulty_ratio > 0.1. Shape: {}"
                .format(X.shape))
    X = X[X.radon_cc_mean <= 9] # 4
    logger.info("Dropped rows with radon_cc_mean > 7.5. Shape: {}"
                .format(X.shape))

    # delete rows with NaNs
    X.dropna(inplace=True)
    logger.info("Dropped all rows with NaNs. Shape: {}"
                .format(X.shape))

    #%% concat with teams_df shuffle the data
    cleaned_df = pd.concat([teams_df, X], join='inner', axis=1)
    cleaned_df = shuffle(cleaned_df, random_state=0)

    #%% export cleaned_df as pickle file to processed folder
    cleaned_df.to_pickle(os.path.join(processed_path, 'cleaned_df.pkl'))
    logger.info("Saved cleaned_df to {}."
            .format(os.path.join(processed_path, 'cleaned_df.pkl')))

    # logging time passed
    end = time()
    time_passed = pd.Timedelta(seconds=end-start).round(freq='s')
    logger.info("Time needed to clean the features: {}".format(time_passed))

#%%
if __name__ == '__main__':

    # configure logging
    logging.basicConfig(
        level = logging.DEBUG,
        format = "%(asctime)s %(name)-20s %(levelname)-8s %(message)s",
        filename = "logs/clean.log",
        datefmt = "%a, %d %b %Y %H:%M:%S")

    # parse arguments
    parser = argparse.ArgumentParser(
            description = "Cleans aggregated data.")
    parser.add_argument(
            '--interim_path',
            default = "data/interim",
            help = "path to aggregated features aggregated_df.pkl \
                    (default: data/interim)")
    parser.add_argument(
            '--processed_path',
            default = "data/processed",
            help = "path to store the cleaned output cleaned_df.pkl \
                    (default: data/processed)")
    args = parser.parse_args()

    # run main
    main(args.interim_path, args.processed_path)

max_stars_repo_name: Keesiu/meta-kaggle

max_issues_repo_name: Keesiu/meta-kaggle

max_forks_repo_name: Keesiu/meta-kaggle


text:
# Demo script to analyse data from Stripe 82
#
# This is minimal version of the procedure carried out in the
#  corresponding notebook (sdss_demo.ipynb).

import pyBA
import sdss
import numpy as np

np.set_printoptions(linewidth=125,suppress=True,precision=3)

# Load data
data = sdss.csv2rec("match_astrom_342.1914_-0.90194_run4198.dat")
nties = len(data)
print nties

# Parse data array into objects
# the A objects will be the fiducial positions from the deep coadd
# the B objects will be source positions in this epoch
# both are given in arcseconds relative to a fiducial center
objectsA = np.array( [ pyBA.Bivarg(mu=[x["master_dra"],x["master_ddec"]],sigma=np.array([x["master_raerr"],x["master_decerr"]])) for x in data ] )
objectsB = np.array( [ pyBA.Bivarg(mu=[x["dra"],x["ddec"]], sigma=np.array([x["raerr"],x["decerr"]])) for x in data ] )

# Suggest starting point for background mapping
S = pyBA.background.suggest_mapping(objectsA,objectsB)

# Get maximum a posteriori background mapping parameters
P = pyBA.background.MAP( objectsA, objectsB, mu0=S.mu, prior=pyBA.Bgmap(), norm_approx=True )
print P.mu

pass

# Create astrometric mapping object
D = pyBA.Amap(P,objectsA, objectsB)

# Condition GP
D.condition()
print D.scale
print D.amp.ravel()

# Regression of a single point
print 'Single point regression'
xy = np.array([0,0])
R = D.regression(xy)
print R[0]

# Regression of a single distribution
print 'Single distribution regression'
BV = pyBA.Bivarg()
R = D.regression(BV)
print R[0]

# Regression of an array of points
print 'Multiple point regression'
xy = 100*np.random.rand(10,2) - 50
R = D.regression(xy)
print R[0]

# Regression of an array of distributions
print 'Multiple distribution regression'
BV = [pyBA.Bivarg(mu=xyi, sigma=np.random.rand(2)) for xyi in xy]
R = D.regression(BV)
print R[0]

max_stars_repo_name: berianjames/pyBAST

max_issues_repo_name: berianjames/pyBAST

max_forks_repo_name: berianjames/pyBAST


text:
import numpy as np
from PIL import Image
a = np.load('data180k/imageset.npy')[5923]*255
b = np.loadtxt('data180k/selfmadetxt/0/00005924.txt')
A = np.loadtxt('A_8192.txt')
print(np.dot(A,a)-b)

max_stars_repo_name: Polarbeartnt/SP-ILC

max_issues_repo_name: Polarbeartnt/SP-ILC

max_forks_repo_name: Polarbeartnt/SP-ILC


text:
# coding=UTF-8
# This Python file uses the following encoding: utf-8

import numpy as np
import matplotlib.pyplot as plt
from numpy import linalg
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.cm as cmx
import os
from matplotlib.pyplot import MultipleLocator

def visualize_3d_gmm(points, w, mu, stdev, type, index, export=True):
    '''
    plots points and their corresponding gmm model in 3D
    Input:
        points: N X 3, sampled points
        w: n_gaussians, gmm weights
        mu: 3 X n_gaussians, gmm means
        stdev: gmm.covariances_ (assuming diagonal covariance matrix)
    Output:
        None
    '''

    n_gaussians = mu.shape[1]
    N = int(np.round(points.shape[0] / n_gaussians))
    # Visualize data
    fig = plt.figure(figsize=(8, 8))
    axes = fig.add_subplot(111, projection='3d')
    #axes.set_xlim([-1, 1])
    #axes.set_ylim([-1, 1])
    #axes.set_zlim([-1, 1])
    plt.grid()
    #plt.gca().set_aspect("equal")
    for i in range(n_gaussians):
        if type == 'full':
            covariances = stdev[i][:3, :3]
            filename = 'SDSS full'
        elif type == 'tied':
            covariances = stdev[:3, :3]
            filename = 'VIPERS_GMM_W4a1_tied'
        elif type == 'diag':
            covariances = np.diag(stdev[i][:3])
            filename = 'VIPERS_GMM_W4a1_diag'
        elif type == 'spherical':
            covariances = np.eye(n_gaussians) * stdev[i]
            filename = 'VIPERS_GMM_W4a1_spherical'
        v, u = np.linalg.eigh(covariances)
        #r = 2. * np.sqrt(2.) * np.sqrt(v)
        r = np.sqrt(v)
        plot_sphere(w=w[i], center=mu[:3, i], r=r, rotation=u, ax=axes)
        #[:, i]取所有行（即三个维度）的第i个数据

    # for n in range(n_gaussians):
    #     data = points[np.where(index == n)]
    #     plt.set_cmap('Set1')
    #     colors = cmx.Set1(np.linspace(0, 1, n_gaussians))
    #     axes.scatter(data[:, 0], data[:, 1], data[:, 2], s = 1.0, alpha = 0.8, color = colors[n])

    plt.title(filename)
    axes.set_xlabel('X /Mpc')
    axes.set_ylabel('Y /Mpc')
    axes.set_zlabel('Z /MPc')
    x_major_locator = MultipleLocator(20)
    # 把x轴的刻度间隔设置为1，并存在变量里
    y_major_locator = MultipleLocator(20)
    z_major_locator = MultipleLocator(20)
    # 把y轴的刻度间隔设置为10，并存在变量里
    ax = plt.gca()
    # ax为两条坐标轴的实例
    ax.xaxis.set_major_locator(x_major_locator)
    # 把x轴的主刻度设置为1的倍数
    ax.yaxis.set_major_locator(y_major_locator)
    ax.zaxis.set_major_locator(z_major_locator)
    plt.savefig(filename, dpi=100, format='png')
    plt.show()

def plot_sphere(w=0, center=[0,0,0], r=[1, 1, 1], rotation=[1,1,1], ax=None):
    '''
        plot a sphere surface
        Input:
            c: 3 elements list, sphere center
            r: 3 element list, sphere original scale in each axis ( allowing to draw elipsoids)
            subdiv: scalar, number of subdivisions (subdivision^2 points sampled on the surface)
                    是椭球的分辨率
            ax: optional pyplot axis object to plot the sphere in.
            sigma_multiplier: sphere additional scale (choosing an std value when plotting gaussians)
        Output:
            ax: pyplot axis object
    '''

    if ax is None:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
    u = np.linspace(0, 2 * np.pi, 30)   #np.linspace 取等差数列
    v = np.linspace(0, np.pi, 30)
    x = r[0] * np.outer(np.cos(u), np.sin(v))
    y = r[1] * np.outer(np.sin(u), np.sin(v))
    z = r[2] * np.outer(np.ones(np.size(u)), np.cos(v))

    for i in range(len(x)):
        for j in range(len(x)):
            # [x[i, j], y[i, j], z[i, j]] = [x[i, j], y[i, j], z[i, j]] + center #spherical专用
            [x[i, j], y[i, j], z[i, j]] = np.dot([x[i, j], y[i, j], z[i, j]], rotation) + center

    ax.plot_surface(x, y, z, alpha=0.1)

    return ax

max_stars_repo_name: umikanero/XDGMM

max_issues_repo_name: umikanero/XDGMM

max_forks_repo_name: umikanero/XDGMM

reason: import numpy,from numpy