<commit_before>from __future__ import division, print_function
import os, re, time, csv, sys
from io_util import make_dir, remove_dir, write_json
# from io_util import myopen, make_dir, remove_dir, tree_to_json, write_json
# from collections import defaultdict
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio.Alphabet import generic_dna
from Bio.SeqFeature import FeatureLocation
import numpy as np
from seq_util import pad_nucleotide_sequences, nuc_alpha, aa_alpha
from datetime import datetime
import json
from pdb import set_trace
from utils import fix_names, num_date, parse_date
from pprint import pprint
from Bio import AlignIO
from Bio.Align import MultipleSeqAlignment
import random

TINY = 1e-10

def calc_af(aln, alpha):
    aln_array = np.array(aln)
    af = np.zeros((len(alpha), aln_array.shape[1]))
    for ai, state in enumerate(alpha):
        af[ai] += (aln_array==state).mean(axis=0)
    af[-1] = 1.0 - af[:-1].sum(axis=0)
    return af

class sequence_set(object):

    def __init__(self, logger, sequences, reference, dateFormat):
        super(sequence_set, self).__init__()
        self.log = logger

        # load sequences from the (parsed) JSON - don't forget to sort out dates
        self.seqs = {}
        for name, data in sequences.iteritems():
            self.seqs[name] = SeqRecord(Seq(data["seq"], generic_dna),
                   id=name, name=name, description=name)
            self.seqs[name].attributes = data["attributes"]
            # tidy up dates
            date_struc = parse_date(self.seqs[name].attributes["raw_date"], dateFormat)
            self.seqs[name].attributes["num_date"] = date_struc[1]
            self.seqs[name].attributes["date"] = date_struc[2]

        # if the reference is to be analysed it'll already be in the (filtered & subsampled)
        # sequences, so no need to add it here, and no need to care about attributes etc
        # we do, however, need it for alignment
        self.reference_in_dataset = reference["included"]
        name = reference["strain"]
        self.reference_seq = SeqRecord(Seq(reference["seq"], generic_dna),
               id=name, name=name, description=name)
        if "genes" in reference and len(reference["genes"]):
            self.proteins = {k:FeatureLocation(start=v["start"], end=v["end"], strand=v["strand"]) for k, v in reference["genes"].iteritems()}
        else:
            self.proteins = None

        # other things:
        self.run_dir = '_'.join(['temp', time.strftime('%Y%m%d-%H%M%S',time.gmtime()), str(random.randint(0,1000000))])
        self.nthreads = 2 # should load from config file

    def convert_trait_to_numerical_date(self, trait, dateFormat):
        for name, seq in self.seqs.iteritems():
            try:
                date_struc = parse_date(seq.attributes[trait], dateFormat)
                seq.attributes[trait] = date_struc[1]
            except KeyError:
                self.log.warn("Attribute {} not found for sequence {}. Ignoring".format(trait, seq.name))

    def codon_align(self):
        self.log.fatal("Codon align not yet implemented")

    def align(self, fname, verbose, debug=False):
        '''
        align sequences using mafft

        side-effects:
            self.aln {MultipleSeqAlignment} reference not present if not in self.seqs
            self.reference_aln {SeqRecord} always set, even if the reference is subsequently discarded
            self.sequence_lookup {dict} map linking seq.id to the alignment, potentialy without the reference
            saves the alignment (always including reference) to fname
        '''
        make_dir(self.run_dir)
        os.chdir(self.run_dir)
        if self.reference_in_dataset:
            out_seqs = self.seqs.values()
        else:
            self.log.notify("Adding reference for alignment step")
            out_seqs = self.seqs.values() + [self.reference_seq]

        SeqIO.write(out_seqs, "temp_in.fasta", "fasta")
        self.log.notify("Running alignment")
        if verbose == 0:
            os.system("mafft --anysymbol --thread " + str(self.nthreads) + " temp_in.fasta 1> temp_out.fasta 2>mafft_stderr")
        else:
            os.system("mafft --anysymbol --thread " + str(self.nthreads) + " temp_in.fasta 1> temp_out.fasta")
        self.aln = AlignIO.read('temp_out.fasta', 'fasta')
        os.chdir("..")
        os.rename(os.path.join(self.run_dir, "temp_out.fasta"), fname)
        if not debug: remove_dir(self.run_dir)

        self.set_reference_alignment()
        if not self.reference_in_dataset:
            self.remove_reference_from_alignment()
        self.set_sequence_lookup()
        self.add_attributes_to_aln()

    def remove_reference_from_alignment(self):
        count = len(self.aln)
        self.aln = MultipleSeqAlignment([s for s in self.aln if s.name!=self.reference_seq.name])
        assert(count == (len(self.aln)+1))

    def set_reference_alignment(self):
        self.reference_aln = [x for x in list(self.aln) if x.name==self.reference_seq.name][0]

    def set_sequence_lookup(self):
        self.sequence_lookup = {seq.id:seq for seq in self.aln}

    def add_attributes_to_aln(self):
        for seqid, seq in self.seqs.iteritems():
            self.sequence_lookup[seqid].attributes = seq.attributes

    def try_restore_align_from_disk(self, fname):
        try:
            self.aln = AlignIO.read(fname, "fasta")
        except IOError:
            return
        except Exception as e:
            self.log.notify("Error restoring from alignment... re-doing")
            print(e)
            return

        try:
            self.set_reference_alignment()
        except IndexError:
            self.log.notify("Reference not found in alignment... re-doing")
            del self.aln
            return

        if not self.reference_in_dataset:
            self.remove_reference_from_alignment()

        if len({x.id for x in self.aln} ^ set(self.seqs.keys())) != 0:
            self.log.notify("Alignment on disk had different sequences... re-doing")
            del self.aln
            del self.reference_aln
            return

        # at this stage we are happy with the alignment
        self.set_sequence_lookup()
        self.add_attributes_to_aln()
        self.log.notify("Alignment restored from disk")

    def strip_non_reference(self):
        '''
        remove insertions relative to the reference from the alignment
        '''
        ungapped = np.array(self.reference_aln)!='-'
        for seq in self.aln:
            seq.seq = Seq("".join(np.array(seq)[ungapped]))

    def make_gaps_ambiguous(self):
        '''
        replace all gaps by 'N' in all sequences in the alignment. TreeTime will treat them
        as fully ambiguous and replace then with the most likely state
        '''
        for seq in self.aln:
            seq_array = np.array(seq)
            gaps = seq_array=='-'
            seq_array[gaps]='N'
            seq.seq = Seq("".join(seq_array))

    def translate(self):
        '''
        make alignments of translations.
        '''
        from Bio.Seq import CodonTable
        codon_table  = CodonTable.ambiguous_dna_by_name['Standard'].forward_table
        self.translations={}
        if not hasattr(self, "proteins"): # ensure dictionary to hold annotation
            self.proteins={}

        # add a default translation of the entire sequence unless otherwise specified
        if len(self.proteins)==0:
            self.proteins.update({'cds':FeatureLocation(start=0,
                end=self.aln.get_alignment_length(), strand=1)})

        # loop over all proteins and create one MSA for each
        for prot in self.proteins:
            aa_seqs = []
            for seq in self.aln:
                tmpseq = self.proteins[prot].extract(seq)
                try:
                    # attempt translation by extracting the sequence according to the BioPhython SeqFeature
                    # in frame gaps of three will translate as '-'
                    tmpseq.seq = tmpseq.seq.translate(gap='-')
                except:
                    # any other codon like '-AA' or 'NNT' etc will fail. translate codons one by one
                    str_seq = str(tmpseq.seq)
                    codons = np.fromstring(str_seq[:len(str_seq)-len(str_seq)%3], dtype='S3')
                    aas = []
                    for c in codons:
                        try: #parse result of single codon translation, add amino acids as appropriate
                            aa = codon_table.get(c)
                            if aa is None:
                                if c=='---':
                                    aas.append('-')
                                else:
                                    aa.append('X')
                            else:
                                aas.append(aa)
                        except:
                            aas.append('X')

                    tmpseq.seq = Seq("".join(aas))
                    self.log.notify("Trouble translating because of invalid codons %s"%seq.id)

                # copy attributes
                tmpseq.attributes = seq.attributes
                aa_seqs.append(tmpseq)
            self.translations[prot] = MultipleSeqAlignment(aa_seqs)

    def clock_filter(self, root_seq=None, n_iqd=3, max_gaps = 1.0, plot=False):
        '''
        remove sequences form the set that are that evolve much faster or slower
        compared the majority. Regions with predominantly gaps can be removed since
        this can skew the evolutionary rates.
        '''
        if root_seq is None: # use consensus
            af = calc_af(self.aln, nuc_alpha)
            root_seq = np.fromstring(nuc_alpha, 'S1')[af.argmax(axis=0)]
        if type(root_seq)==str and root_seq in self.sequence_lookup:
            root_seq = np.array(self.sequence_lookup[root_seq])
        if max_gaps<1.0:
            af=calc_af(self.aln, nuc_alpha)
            good_pos = af[nuc_alpha.index('-')]<max_gaps
        else:
            good_pos = np.ones(self.aln.get_alignment_length(), dtype=bool)
        date_vs_distance = {}
        # self.reference_aln = None already set at alignment step
        for seq in self.aln:
            date_vs_distance[seq.id] = (seq.attributes['num_date'],
                np.mean((np.array(seq)!=root_seq)[(np.array(seq)!='-')&(root_seq!='-')&good_pos]))
            # if seq.id==self.reference.id:
            #     self.reference_aln = seq
        date_vs_distance_array=np.array(date_vs_distance.values())
        from scipy.stats import linregress, scoreatpercentile
        slope, intercept, rval, pval, stderr = linregress(date_vs_distance_array[:,0], date_vs_distance_array[:,1])
        print("distance vs time regression:",slope)
        residuals = (intercept + slope*date_vs_distance_array[:,0]) - date_vs_distance_array[:,1]
        IQD = scoreatpercentile(residuals, 75) - scoreatpercentile(residuals,25)
        if plot:
            import matplotlib.pyplot as plt
            plt.ion()
            plt.scatter(date_vs_distance_array[:,0], date_vs_distance_array[:,1], c='g')
            bad_points = abs(intercept+slope*date_vs_distance_array[:,0] - date_vs_distance_array[:,1])>n_iqd*IQD
            plt.scatter(date_vs_distance_array[bad_points,0], date_vs_distance_array[bad_points,1], c='r')

        print("before clock filter:",len(self.aln))
        tmp = {seq.id:seq for seq in self.aln
                if abs(intercept+slope*date_vs_distance[seq.id][0] - date_vs_distance[seq.id][1])<n_iqd*IQD}
        if self.reference.id not in tmp and self.reference.reference_in_dataset:
            self.log.notify('adding reference again after clock filter')
            tmp[self.reference.id] = self.reference_aln
        self.aln = MultipleSeqAlignment(tmp.values())
        print("after clock filter:",len(self.aln))

    def diversity_statistics(self):
        ''' calculate alignment entropy of nucleotide and optionally protein alignments '''
        if not hasattr(self, "aln"):
            self.log.fatal("Diversity statistics calculated before alignment generated.")
            return
        aln_array = np.array(self.aln)
        self.af = {'nuc': calc_af(self.aln, nuc_alpha)}
        tmp_af = self.af['nuc'][:-2]/self.af['nuc'][:-2].sum(axis=0)
        self.entropy ={'nuc': -(tmp_af*np.log(tmp_af+TINY)).sum(axis=0)}

        if hasattr(self, "translations"):
            for prot, aln in self.translations.iteritems():
                self.af[prot] = calc_af(aln, aa_alpha)
                tmp_af = self.af[prot][:-2]/self.af[prot][:-2].sum(axis=0)
                self.entropy[prot] = -(tmp_af*np.log(tmp_af+TINY)).sum(axis=0)

    def export_diversity(self, fname = 'entropy.json', indent=None):
        '''
        write the alignment entropy of each alignment (nucleotide and translations) to file
        '''
        if not hasattr(self, "entropy"):
            self.diversity_statistics()
        entropy_json = {}
        for feat in self.entropy:
            S = [max(0,round(x,4)) for x in self.entropy[feat]]
            n = len(S)
            if feat=='nuc':
                entropy_json[feat] = {'pos':range(0,n), 'codon':[x//3 for x in range(0,n)], 'val':S}
            else:
                entropy_json[feat] = {'pos':[x for x in self.proteins[feat]][::3],
                                      'codon':[(x-self.proteins[feat].start)//3 for x in self.proteins[feat]][::3], 'val':S}
        write_json(entropy_json, fname, indent=indent)
<commit_msg>Moved safe translation code into its own function, added tests, and checked for specific exceptions instead of catching all errors.<commit_after>from __future__ import division, print_function
import os, re, time, csv, sys
from io_util import make_dir, remove_dir, write_json
# from io_util import myopen, make_dir, remove_dir, tree_to_json, write_json
# from collections import defaultdict
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio.Alphabet import generic_dna
from Bio.SeqFeature import FeatureLocation
from Bio.Data.CodonTable import TranslationError
from Bio.Seq import CodonTable
import numpy as np
from seq_util import pad_nucleotide_sequences, nuc_alpha, aa_alpha
from datetime import datetime
import json
from pdb import set_trace
from utils import fix_names, num_date, parse_date
from pprint import pprint
from Bio import AlignIO
from Bio.Align import MultipleSeqAlignment
import random

TINY = 1e-10

def calc_af(aln, alpha):
    aln_array = np.array(aln)
    af = np.zeros((len(alpha), aln_array.shape[1]))
    for ai, state in enumerate(alpha):
        af[ai] += (aln_array==state).mean(axis=0)
    af[-1] = 1.0 - af[:-1].sum(axis=0)
    return af

def safe_translate(sequence):
    """Returns an amino acid translation of the given nucleotide sequence accounting
    for gaps in the given sequence.

    >>> safe_translate("ATG")
    'M'
    >>> safe_translate("ATGGT-")
    'MX'
    >>> safe_translate("ATG---")
    'M-'
    >>> safe_translate("ATGTAG")
    'M*'
    >>> safe_translate("")
    ''
    >>> safe_translate("ATGT")
    'M'
    """
    try:
        # Attempt translation by extracting the sequence according to the
        # BioPhython SeqFeature in frame gaps of three will translate as '-'
        translated_sequence = str(Seq(sequence).translate(gap='-'))
    except TranslationError:
        # Any other codon like '-AA' or 'NNT' etc will fail. Translate codons
        # one by one.
        codon_table  = CodonTable.ambiguous_dna_by_name['Standard'].forward_table
        str_seq = str(sequence)
        codons = np.fromstring(str_seq[:len(str_seq) - len(str_seq) % 3], dtype='S3')
        assert len(codons) > 0
        aas = []

        for c in codons:
            # Parse result of single codon translation, add amino acids as
            # appropriate.
            try:
                aa = codon_table.get(c)
                if aa is None:
                    if c == '---':
                        aas.append('-')
                    else:
                        aas.append('X')
                else:
                    aas.append(aa)
            except (TranslationError, ValueError):
                aas.append('X')

        translated_sequence = "".join(aas)

    return translated_sequence

class sequence_set(object):

    def __init__(self, logger, sequences, reference, dateFormat):
        super(sequence_set, self).__init__()
        self.log = logger

        # load sequences from the (parsed) JSON - don't forget to sort out dates
        self.seqs = {}
        for name, data in sequences.iteritems():
            self.seqs[name] = SeqRecord(Seq(data["seq"], generic_dna),
                   id=name, name=name, description=name)
            self.seqs[name].attributes = data["attributes"]
            # tidy up dates
            date_struc = parse_date(self.seqs[name].attributes["raw_date"], dateFormat)
            self.seqs[name].attributes["num_date"] = date_struc[1]
            self.seqs[name].attributes["date"] = date_struc[2]

        # if the reference is to be analysed it'll already be in the (filtered & subsampled)
        # sequences, so no need to add it here, and no need to care about attributes etc
        # we do, however, need it for alignment
        self.reference_in_dataset = reference["included"]
        name = reference["strain"]
        self.reference_seq = SeqRecord(Seq(reference["seq"], generic_dna),
               id=name, name=name, description=name)
        if "genes" in reference and len(reference["genes"]):
            self.proteins = {k:FeatureLocation(start=v["start"], end=v["end"], strand=v["strand"]) for k, v in reference["genes"].iteritems()}
        else:
            self.proteins = None

        # other things:
        self.run_dir = '_'.join(['temp', time.strftime('%Y%m%d-%H%M%S',time.gmtime()), str(random.randint(0,1000000))])
        self.nthreads = 2 # should load from config file

    def convert_trait_to_numerical_date(self, trait, dateFormat):
        for name, seq in self.seqs.iteritems():
            try:
                date_struc = parse_date(seq.attributes[trait], dateFormat)
                seq.attributes[trait] = date_struc[1]
            except KeyError:
                self.log.warn("Attribute {} not found for sequence {}. Ignoring".format(trait, seq.name))

    def codon_align(self):
        self.log.fatal("Codon align not yet implemented")

    def align(self, fname, verbose, debug=False):
        '''
        align sequences using mafft

        side-effects:
            self.aln {MultipleSeqAlignment} reference not present if not in self.seqs
            self.reference_aln {SeqRecord} always set, even if the reference is subsequently discarded
            self.sequence_lookup {dict} map linking seq.id to the alignment, potentialy without the reference
            saves the alignment (always including reference) to fname
        '''
        make_dir(self.run_dir)
        os.chdir(self.run_dir)
        if self.reference_in_dataset:
            out_seqs = self.seqs.values()
        else:
            self.log.notify("Adding reference for alignment step")
            out_seqs = self.seqs.values() + [self.reference_seq]

        SeqIO.write(out_seqs, "temp_in.fasta", "fasta")
        self.log.notify("Running alignment")
        if verbose == 0:
            os.system("mafft --anysymbol --thread " + str(self.nthreads) + " temp_in.fasta 1> temp_out.fasta 2>mafft_stderr")
        else:
            os.system("mafft --anysymbol --thread " + str(self.nthreads) + " temp_in.fasta 1> temp_out.fasta")
        self.aln = AlignIO.read('temp_out.fasta', 'fasta')
        os.chdir("..")
        os.rename(os.path.join(self.run_dir, "temp_out.fasta"), fname)
        if not debug: remove_dir(self.run_dir)

        self.set_reference_alignment()
        if not self.reference_in_dataset:
            self.remove_reference_from_alignment()
        self.set_sequence_lookup()
        self.add_attributes_to_aln()

    def remove_reference_from_alignment(self):
        count = len(self.aln)
        self.aln = MultipleSeqAlignment([s for s in self.aln if s.name!=self.reference_seq.name])
        assert(count == (len(self.aln)+1))

    def set_reference_alignment(self):
        self.reference_aln = [x for x in list(self.aln) if x.name==self.reference_seq.name][0]

    def set_sequence_lookup(self):
        self.sequence_lookup = {seq.id:seq for seq in self.aln}

    def add_attributes_to_aln(self):
        for seqid, seq in self.seqs.iteritems():
            self.sequence_lookup[seqid].attributes = seq.attributes

    def try_restore_align_from_disk(self, fname):
        try:
            self.aln = AlignIO.read(fname, "fasta")
        except IOError:
            return
        except Exception as e:
            self.log.notify("Error restoring from alignment... re-doing")
            print(e)
            return

        try:
            self.set_reference_alignment()
        except IndexError:
            self.log.notify("Reference not found in alignment... re-doing")
            del self.aln
            return

        if not self.reference_in_dataset:
            self.remove_reference_from_alignment()

        if len({x.id for x in self.aln} ^ set(self.seqs.keys())) != 0:
            self.log.notify("Alignment on disk had different sequences... re-doing")
            del self.aln
            del self.reference_aln
            return

        # at this stage we are happy with the alignment
        self.set_sequence_lookup()
        self.add_attributes_to_aln()
        self.log.notify("Alignment restored from disk")

    def strip_non_reference(self):
        '''
        remove insertions relative to the reference from the alignment
        '''
        ungapped = np.array(self.reference_aln)!='-'
        for seq in self.aln:
            seq.seq = Seq("".join(np.array(seq)[ungapped]))

    def make_gaps_ambiguous(self):
        '''
        replace all gaps by 'N' in all sequences in the alignment. TreeTime will treat them
        as fully ambiguous and replace then with the most likely state
        '''
        for seq in self.aln:
            seq_array = np.array(seq)
            gaps = seq_array=='-'
            seq_array[gaps]='N'
            seq.seq = Seq("".join(seq_array))

    def translate(self):
        '''
        make alignments of translations.
        '''
        from Bio.Seq import CodonTable
        codon_table  = CodonTable.ambiguous_dna_by_name['Standard'].forward_table
        self.translations={}
        if not hasattr(self, "proteins"): # ensure dictionary to hold annotation
            self.proteins={}

        # add a default translation of the entire sequence unless otherwise specified
        if len(self.proteins)==0:
            self.proteins.update({'cds':FeatureLocation(start=0,
                end=self.aln.get_alignment_length(), strand=1)})

        # loop over all proteins and create one MSA for each
        for prot in self.proteins:
            aa_seqs = []
            for seq in self.aln:
                tmpseq = self.proteins[prot].extract(seq)
                translated_seq = safe_translate(str(tmpseq.seq))
                tmpseq.seq = Seq(translated_seq)

                if "-" in translated_seq or "X" in translated_seq:
                    self.log.notify("Trouble translating because of invalid codons %s" % seq.id)

                # copy attributes
                tmpseq.attributes = seq.attributes
                aa_seqs.append(tmpseq)
            self.translations[prot] = MultipleSeqAlignment(aa_seqs)

    def clock_filter(self, root_seq=None, n_iqd=3, max_gaps = 1.0, plot=False):
        '''
        remove sequences form the set that are that evolve much faster or slower
        compared the majority. Regions with predominantly gaps can be removed since
        this can skew the evolutionary rates.
        '''
        if root_seq is None: # use consensus
            af = calc_af(self.aln, nuc_alpha)
            root_seq = np.fromstring(nuc_alpha, 'S1')[af.argmax(axis=0)]
        if type(root_seq)==str and root_seq in self.sequence_lookup:
            root_seq = np.array(self.sequence_lookup[root_seq])
        if max_gaps<1.0:
            af=calc_af(self.aln, nuc_alpha)
            good_pos = af[nuc_alpha.index('-')]<max_gaps
        else:
            good_pos = np.ones(self.aln.get_alignment_length(), dtype=bool)
        date_vs_distance = {}
        # self.reference_aln = None already set at alignment step
        for seq in self.aln:
            date_vs_distance[seq.id] = (seq.attributes['num_date'],
                np.mean((np.array(seq)!=root_seq)[(np.array(seq)!='-')&(root_seq!='-')&good_pos]))
            # if seq.id==self.reference.id:
            #     self.reference_aln = seq
        date_vs_distance_array=np.array(date_vs_distance.values())
        from scipy.stats import linregress, scoreatpercentile
        slope, intercept, rval, pval, stderr = linregress(date_vs_distance_array[:,0], date_vs_distance_array[:,1])
        print("distance vs time regression:",slope)
        residuals = (intercept + slope*date_vs_distance_array[:,0]) - date_vs_distance_array[:,1]
        IQD = scoreatpercentile(residuals, 75) - scoreatpercentile(residuals,25)
        if plot:
            import matplotlib.pyplot as plt
            plt.ion()
            plt.scatter(date_vs_distance_array[:,0], date_vs_distance_array[:,1], c='g')
            bad_points = abs(intercept+slope*date_vs_distance_array[:,0] - date_vs_distance_array[:,1])>n_iqd*IQD
            plt.scatter(date_vs_distance_array[bad_points,0], date_vs_distance_array[bad_points,1], c='r')

        print("before clock filter:",len(self.aln))
        tmp = {seq.id:seq for seq in self.aln
                if abs(intercept+slope*date_vs_distance[seq.id][0] - date_vs_distance[seq.id][1])<n_iqd*IQD}
        if self.reference.id not in tmp and self.reference.reference_in_dataset:
            self.log.notify('adding reference again after clock filter')
            tmp[self.reference.id] = self.reference_aln
        self.aln = MultipleSeqAlignment(tmp.values())
        print("after clock filter:",len(self.aln))

    def diversity_statistics(self):
        ''' calculate alignment entropy of nucleotide and optionally protein alignments '''
        if not hasattr(self, "aln"):
            self.log.fatal("Diversity statistics calculated before alignment generated.")
            return
        aln_array = np.array(self.aln)
        self.af = {'nuc': calc_af(self.aln, nuc_alpha)}
        tmp_af = self.af['nuc'][:-2]/self.af['nuc'][:-2].sum(axis=0)
        self.entropy ={'nuc': -(tmp_af*np.log(tmp_af+TINY)).sum(axis=0)}

        if hasattr(self, "translations"):
            for prot, aln in self.translations.iteritems():
                self.af[prot] = calc_af(aln, aa_alpha)
                tmp_af = self.af[prot][:-2]/self.af[prot][:-2].sum(axis=0)
                self.entropy[prot] = -(tmp_af*np.log(tmp_af+TINY)).sum(axis=0)

    def export_diversity(self, fname = 'entropy.json', indent=None):
        '''
        write the alignment entropy of each alignment (nucleotide and translations) to file
        '''
        if not hasattr(self, "entropy"):
            self.diversity_statistics()
        entropy_json = {}
        for feat in self.entropy:
            S = [max(0,round(x,4)) for x in self.entropy[feat]]
            n = len(S)
            if feat=='nuc':
                entropy_json[feat] = {'pos':range(0,n), 'codon':[x//3 for x in range(0,n)], 'val':S}
            else:
                entropy_json[feat] = {'pos':[x for x in self.proteins[feat]][::3],
                                      'codon':[(x-self.proteins[feat].start)//3 for x in self.proteins[feat]][::3], 'val':S}
        write_json(entropy_json, fname, indent=indent)
<|endoftext|>


<commit_before>__version__ = '2.28.1'
<commit_msg>Bump version<commit_after>__version__ = '2.28.2'
<|endoftext|>


<commit_before># All game related code

import json
import random

import card_parser

class Game:
    def __init__(self):
        self.players = {}
        self.picked_cards = []

    def add_player(self, conn, data):
        player = Player(conn, data, self.picked_cards)
        self.players[player.get_name()] = player
        conn.sendMessage(json.dumps({'action': 'accepted', 'data': ''}))

    def check_for_start(self):
        if len(self.players) == 3:
            self.start_game()

    def attribute_selected(self, data):
        # Get all cards with players
        player_cards = self.players.items()
        cards = [(player, player.get_card()) for _, player in player_cards]
        all_card_values = [card.get_values() for player, card in cards]

        # Get current card and compare it
        attr = data['data']['attributeToCompare']
        winner, winner_card = Card.compare(attr, cards)
        data = {'all_cards': all_card_values, 'winner_card': winner_card.get_values()}
        winner.add_cards([card for _, card in cards])

        # Check if game has ended
        has_ended, player = self.check_game_end()
        if has_ended:
            data['loser'] = player
            has_next_round = True if len(self.players) > 1 else False
            self.broadcast(data, 'playerLost', next_card=has_next_round)
            if not has_next_round:
                self.end_connections()
            else:
                for p in player:
                    self.players[p].get_connection().sendClose()
                    del self.players[player]

        # Update after comparison
        data['turn'] = winner.get_name()
        data['attribute_compared'] = attr
        self.broadcast(data, 'next')

    def start_game(self):
        # Randomly select player to start
        random.seed(None)
        name, player = random.choice(self.players.items())
        plys = [plyr for _, plyr in self.players.items()]
        data = {'turn': name, 'players': plys}
        self.broadcast(data, 'start')

    def check_game_end(self):
        # Check if all players still have cards left
        has_ended = False
        players = []
        for name, player in self.players.items():
            if not player.has_cards():
                has_ended = True
                players.append(name)
        return has_ended, players

    def broadcast(self, data, action, next_card=True):
        # Send data to all players
        for name, player in self.players.items():
            if next_card:
                data['card'] = player.next_card()
            conn = player.get_connection()
            conn.sendMessage(json.dumps({'action': action, 'data': data}))

    def end_connections(self):
        # Close all connections
        for _, player in self.players.items():
            player.get_connection().sendClose()

class Player:
    def __init__(self, conn, data, picked_cards):
        self.name = data['name']
        self.connection = conn
        self.cards = self.receive_cards(data, picked_cards)
        self.current_card = None

    def get_name(self):
        return self.name

    def get_card(self):
        return self.current_card

    def get_connection(self):
        return self.connection

    def receive_cards(self, player_data, picked_cards):
        # Get card information from mobile.de API based on location
        long = player_data['data']['long']
        lat = player_data['data']['lat']
        cards = card_parser.main(lat, long, picked_cards)
        return [Card(json.loads(values)) for values in cards]

    def next_card(self):
        try:
            self.current_card = self.cards.pop(0)
        except:
            return {}
        return self.current_card.get_values()

    def add_cards(self, cards):
        try:
            self.cards.extend(cards)
        except:
            print "SOMETHING WENT WRONG!"

    def has_cards(self):
        return len(self.cards) > 0

class Card:
    comparisons = {'price': min,
                   'power': max,
                   'registration': max,
                   'mileage': min,
                   'consumption': min}

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

    @staticmethod
    def compare(attr, player_cards):
        # Select min or max according to selected attribute
        comp = Card.comparisons[attr]
        # Compare card value from (player, card) as float
        return comp(player_cards, key=lambda pair: float(pair[1].get(attr)))

    def get_values(self):
        return self.values

    def get(self, attr):
        return self.values[attr]
<commit_msg>Fix player list<commit_after># All game related code

import json
import random

import card_parser

class Game:
    def __init__(self):
        self.players = {}
        self.picked_cards = []

    def add_player(self, conn, data):
        player = Player(conn, data, self.picked_cards)
        self.players[player.get_name()] = player
        conn.sendMessage(json.dumps({'action': 'accepted', 'data': ''}))

    def check_for_start(self):
        if len(self.players) == 3:
            self.start_game()

    def attribute_selected(self, data):
        # Get all cards with players
        player_cards = self.players.items()
        cards = [(player, player.get_card()) for _, player in player_cards]
        all_card_values = [card.get_values() for player, card in cards]

        # Get current card and compare it
        attr = data['data']['attributeToCompare']
        winner, winner_card = Card.compare(attr, cards)
        data = {'all_cards': all_card_values, 'winner_card': winner_card.get_values()}
        winner.add_cards([card for _, card in cards])

        # Check if game has ended
        has_ended, player = self.check_game_end()
        if has_ended:
            data['loser'] = player
            has_next_round = True if len(self.players) > 1 else False
            self.broadcast(data, 'playerLost', next_card=has_next_round)
            if not has_next_round:
                self.end_connections()
            else:
                for p in player:
                    self.players[p].get_connection().sendClose()
                    del self.players[player]

        # Update after comparison
        data['turn'] = winner.get_name()
        data['attribute_compared'] = attr
        self.broadcast(data, 'next')

    def start_game(self):
        # Randomly select player to start
        random.seed(None)
        name, player = random.choice(self.players.items())
        names = [name for name, _ in self.players.items()]
        data = {'turn': name, 'players': names}
        self.broadcast(data, 'start')

    def check_game_end(self):
        # Check if all players still have cards left
        has_ended = False
        players = []
        for name, player in self.players.items():
            if not player.has_cards():
                has_ended = True
                players.append(name)
        return has_ended, players

    def broadcast(self, data, action, next_card=True):
        # Send data to all players
        for name, player in self.players.items():
            if next_card:
                data['card'] = player.next_card()
            conn = player.get_connection()
            conn.sendMessage(json.dumps({'action': action, 'data': data}))

    def end_connections(self):
        # Close all connections
        for _, player in self.players.items():
            player.get_connection().sendClose()

class Player:
    def __init__(self, conn, data, picked_cards):
        self.name = data['name']
        self.connection = conn
        self.cards = self.receive_cards(data, picked_cards)
        self.current_card = None

    def get_name(self):
        return self.name

    def get_card(self):
        return self.current_card

    def get_connection(self):
        return self.connection

    def receive_cards(self, player_data, picked_cards):
        # Get card information from mobile.de API based on location
        long = player_data['data']['long']
        lat = player_data['data']['lat']
        cards = card_parser.main(lat, long, picked_cards)
        return [Card(json.loads(values)) for values in cards]

    def next_card(self):
        try:
            self.current_card = self.cards.pop(0)
        except:
            return {}
        return self.current_card.get_values()

    def add_cards(self, cards):
        try:
            self.cards.extend(cards)
        except:
            print "SOMETHING WENT WRONG!"

    def has_cards(self):
        return len(self.cards) > 0

class Card:
    comparisons = {'price': min,
                   'power': max,
                   'registration': max,
                   'mileage': min,
                   'consumption': min}

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

    @staticmethod
    def compare(attr, player_cards):
        # Select min or max according to selected attribute
        comp = Card.comparisons[attr]
        # Compare card value from (player, card) as float
        return comp(player_cards, key=lambda pair: float(pair[1].get(attr)))

    def get_values(self):
        return self.values

    def get(self, attr):
        return self.values[attr]
<|endoftext|>


<commit_before># ----- Info ------------------------------------------------------------------

__author__ = 'Michael Montero <mcmontero@gmail.com>'

# ----- Imports ---------------------------------------------------------------

import glob
import json
import os
import random
import re
import string
import time

# ----- Public Classes --------------------------------------------------------

class FileSystemQueue(object):
    '''Allows for queueing data to a file system queue.'''

    def __init__(self, queue_dir, domain=None):
        self.__queue_dir = queue_dir
        self.__domain = domain

    def enqueue(self, data=tuple()):
        file_name = self.__get_file_name()

        with open(file_name + '.writing', 'w') as f:
            f.write(json.dumps(data))

        os.rename(file_name + '.writing', file_name)

        return file_name

    def dequeue(self, queue_file_name):
        try:
            os.unlink(queue_file_name)
        except FileNotFoundError:
            pass

        return self

    def get(self, remove_queue_file=True):
        files = glob.glob(os.path.join(self.__queue_dir, '*'))
        if files is None or len(files) == 0:
            return None

        queue = []
        for file_path in sorted(files):
            if re.search('\.writing$', file_path):
                continue

            if self.__domain is not None:
                if not re.search(
                        '^' + self.__domain + '-',
                        os.path.basename(file_path)
                       ):
                    continue

            with open(file_path, 'r') as f:
                queue.append({
                    'file': file_path,
                    'data': json.loads(f.read())
                })

            if remove_queue_file is True:
                os.unlink(file_path)

        return queue if len(queue) > 0 else None

    def __get_file_name(self):
        file = '{}-{}'.format(int(time.time()), self.__random_string(32))
        if self.__domain is not None:
            file = self.__domain + '-' + file

        return '{}/{}'.format(self.__queue_dir, file)

    def __random_string(self, length):
        return \
            ''.join(
                random.choice(
                    string.ascii_lowercase +
                    string.ascii_uppercase +
                    string.digits
                )
                    for _ in range(length)
            )
<commit_msg>Proper handling of null payloads.<commit_after># ----- Info ------------------------------------------------------------------

__author__ = 'Michael Montero <mcmontero@gmail.com>'

# ----- Imports ---------------------------------------------------------------

import glob
import json
import os
import random
import re
import string
import time

# ----- Public Classes --------------------------------------------------------

class FileSystemQueue(object):
    '''Allows for queueing data to a file system queue.'''

    def __init__(self, queue_dir, domain=None):
        self.__queue_dir = queue_dir
        self.__domain = domain

    def enqueue(self, data=tuple()):
        file_name = self.__get_file_name()

        with open(file_name + '.writing', 'w') as f:
            f.write(json.dumps(data))

        os.rename(file_name + '.writing', file_name)

        return file_name

    def dequeue(self, queue_file_name):
        try:
            os.unlink(queue_file_name)
        except FileNotFoundError:
            pass

        return self

    def get(self, remove_queue_file=True):
        files = glob.glob(os.path.join(self.__queue_dir, '*'))
        if files is None or len(files) == 0:
            return None

        queue = []
        for file_path in sorted(files):
            if re.search('\.writing$', file_path):
                continue

            if self.__domain is not None:
                if not re.search(
                        '^' + self.__domain + '-',
                        os.path.basename(file_path)
                       ):
                    continue

            with open(file_path, 'r') as f:
                payload = f.read()
                if payload is not None and len(payload) > 0:
                    queue.append({
                        'file': file_path,
                        'data': json.loads(payload)
                    })

            if remove_queue_file is True:
                os.unlink(file_path)

        return queue if len(queue) > 0 else None

    def __get_file_name(self):
        file = '{}-{}'.format(int(time.time()), self.__random_string(32))
        if self.__domain is not None:
            file = self.__domain + '-' + file

        return '{}/{}'.format(self.__queue_dir, file)

    def __random_string(self, length):
        return \
            ''.join(
                random.choice(
                    string.ascii_lowercase +
                    string.ascii_uppercase +
                    string.digits
                )
                    for _ in range(length)
            )
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Created on Mon Mar 20 00:17:56 2017

@author: Michael
"""

# TURN OFF AUTHORIZATION OF YOUR FIREBASE PROJECT BEFORE RUNNING THE CODE.
# THE TWO PACKAGE NAMES ARE INCLUDED BELOW
# ALSO THIS WAS DONE IN PYTHON 2.7

from faker import Factory # pip install Faker
from random import randint
from firebase import firebase # pip install python-firebase

faker = Factory.create()

carMake = ["Nissan", "Chevy", "Honda", "Cadillac", "Ferrari",
           "Ford", "Lexus", "Hyundai", "Dodge", "Mitsubishi"]

carModel = ["Maxima", "Cobalt", "Civic", "ATS", "Enzo",
            "F-150", "IS-350", "Accent", "Dart", "Lancer"]

spotTypelist = ["student", "staff", "private", "handicapped"]

index = 0
vehiclesNum = setsNum = 0
userURL = projectName = firebaseDirectory = header = vehicleString = []
lot = spotID = vacancy = spotType = occupant = []
userEmail = userID = firstName = lastName = phone = permitType = expDate = purchaseDate = []
make = model = color = licensePlate = []
authorized = False

projectName = raw_input("Firebase project = example for https://example.firebaseio.com\n"
                        "Please enter the name of your Firebase project: ")
userURL = "https://"
userURL = userURL + projectName
userURL = userURL + ".firebaseio.com"

setsNum = input("Please enter how many sets of data you want to create: ")

firebase = firebase.FirebaseApplication(userURL, None)

for i in range(setsNum):
        # Spots
    lot = faker.state_abbr()            # Lot
    firebaseDirectory = "/Spots/"
    firebaseDirectory = firebaseDirectory + lot
    spotID = randint(1, 250)            # Spot ID
    vacancy = faker.boolean()           # Vacant
    occupant = faker.email()
    if vacancy != True:
        authorized = faker.boolean()    # Authorized
        occupant = faker.email()        # User Email
    else:
        authorized = False
        occupant = 'none'

    index = randint(0,3)
    spotType = spotTypelist[index]

    firebase.put(firebaseDirectory, spotID, params={'print': 'silent'},
                         data={'vacancy': vacancy,
                               'type': spotType,
                               'occupant': occupant,
                               'authorized': authorized})

        # UserAccounts
    userID = faker.ean8()               # User ID
    firebaseDirectory = "/UserAccounts/"
    userEmail = faker.email()           # User Email
    firstName = faker.first_name()      # First Name
    lastName = faker.last_name()        # Last Name
    phone = faker.phone_number()        # Phone Number
            # permitInfo
    index = randint(0,3)
    permitType = spotTypelist[index]    # Permit Type
    expDate = faker.date()              # Purchase Date
    purchaseDate = faker.date()         # Expiration Date

    firebase.put(firebaseDirectory, userID, params={'print': 'silent'},
                         data={'userEmail': userEmail,
                               'firstName': firstName,
                               'lastName': lastName,
                               'phone': phone,
                               'permit': {
                                   'type': permitType,
                                   'expDate': expDate,
                                   'purchaseDate': purchaseDate}})

    firebaseDirectory = firebaseDirectory + userID
    vehiclesNum = randint(1,4)
    for j in range(vehiclesNum, 0, -1):
            # vehicleInfo
        vehicleString = str(j)
        header = "vehicle"
        header = header + vehicleString
        index = randint(0,9)
        make = carMake[index]           # Vehicle Make
        index = randint(0,9)
        model = carModel[index]         # Vehicle Model
        color = faker.safe_color_name() # Vehicle Color
        licensePlate = faker.ean8()     # License Plate

        firebase.put(firebaseDirectory, header, params={'print': 'silent'},
                         data={'make': make,
                               'model': model,
                               'color': color,
                               'licensePlate': licensePlate})

<commit_msg>fixed exp/pur dates to make sense and vehicles section<commit_after># -*- coding: utf-8 -*-
"""
Created on Mon Mar 20 00:17:56 2017

@author: Michael
"""

# TURN OFF AUTHORIZATION OF YOUR FIREBASE PROJECT BEFORE RUNNING THE CODE.
# THE TWO PACKAGE NAMES ARE INCLUDED BELOW
# ALSO THIS WAS DONE IN PYTHON 2.7

from faker import Factory # pip install Faker
from random import randint
from firebase import firebase # pip install python-firebase
from datetime import timedelta
import os
import datetime

faker = Factory.create()

carMake = ["Nissan", "Chevy", "Honda", "Cadillac", "Ferrari",
           "Ford", "Lexus", "Hyundai", "Dodge", "Mitsubishi"]

carModel = ["Maxima", "Cobalt", "Civic", "ATS", "Enzo",
            "F-150", "IS-350", "Accent", "Dart", "Lancer"]

spotTypelist = ["student", "staff", "private", "handicapped"]

index = 0
vehiclesNum = setsNum = 0
userURL = projectName = firebaseDirectory = header = vehicleString = []
lot = spotID = vacancy = spotType = occupant = []
userEmail = userID = firstName = lastName = phone = permitType = expDate = purchaseDate = []
make = model = color = licensePlate = []
authorized = False

projectName = raw_input("Firebase project = example for https://example.firebaseio.com\n"
                        "Please enter the name of your Firebase project: ")
userURL = "https://"
userURL = userURL + projectName
userURL = userURL + ".firebaseio.com"

setsNum = input("Please enter how many sets of data you want to create: ")

firebase = firebase.FirebaseApplication(userURL, None)

for i in range(setsNum):
        # Spots
    lot = faker.state_abbr()            # Lot
    firebaseDirectory = "/Spots/"
    firebaseDirectory = firebaseDirectory + lot
    spotID = randint(1, 250)            # Spot ID
    vacancy = faker.boolean()           # Vacant
    occupant = faker.email()
    if vacancy != True:
        authorized = faker.boolean()    # Authorized
        occupant = faker.email()        # User Email
    else:
        authorized = False
        occupant = 'none'

    index = randint(0,3)
    spotType = spotTypelist[index]

    firebase.put(firebaseDirectory, spotID, params={'print': 'silent'},
                         data={'vacancy': vacancy,
                               'type': spotType,
                               'occupant': occupant,
                               'authorized': authorized})

        # UserAccounts
    userID = faker.ean8()               # User ID
    firebaseDirectory = "/UserAccounts/"
    userEmail = faker.email()           # User Email
    firstName = faker.first_name()      # First Name
    lastName = faker.last_name()        # Last Name
    phone = faker.phone_number()        # Phone Number
            # permitInfo
    index = randint(0,3)
    permitType = spotTypelist[index]    # Permit Type
                                        # Purchase Date
    purchaseDate = datetime.datetime.date(faker.date_time_between(start_date="-1y", end_date="now", tzinfo=None))
                                        # Expiration Date (exactly one year added for permits)
    expDate = purchaseDate + timedelta(days=365)

    firebase.put(firebaseDirectory, userID, params={'print': 'silent'},
                         data={'userEmail': userEmail,
                               'firstName': firstName,
                               'lastName': lastName,
                               'phone': phone,
                               'permit': {
                                   'type': permitType,
                                   'expDate': expDate,
                                   'purchaseDate': purchaseDate}})

    firebaseDirectory = firebaseDirectory + userID + "/vehicles"
    vehiclesNum = randint(1,4)
    for j in range(vehiclesNum, 0, -1):
            # vehicleInfo
        vehicleString = str(j)
        header = "vehicle"
        header = header + vehicleString
        index = randint(0,9)
        make = carMake[index]           # Vehicle Make
        index = randint(0,9)
        model = carModel[index]         # Vehicle Model
        color = faker.safe_color_name() # Vehicle Color
        licensePlate = faker.ean8()     # License Plate

        firebase.put(firebaseDirectory, header, params={'print': 'silent'},
                         data={'make': make,
                               'model': model,
                               'color': color,
                               'licensePlate': licensePlate})

os._exit(0)
<|endoftext|>


<commit_before>#ipc_lista1.2
#Professor: Jucimar Junior
#Any Mendes Carvalho - 165310044
#
#
#
#

#Faça um Programa que peça um número e então mostre a mensagem O número informado foi [número].

number = input("Digite um número: ")
print "O número digitado foi ",number
<commit_msg>Update ipc_lista1.2.py<commit_after>#ipc_lista1.2
#Professor: Jucimar Junior
#Any Mendes Carvalho - 1615310044
#
#
#
#

#Faça um Programa que peça um número e então mostre a mensagem O número informado foi [número].

number = input("Digite um número: ")
print "O número digitado foi ",number
<|endoftext|>


<commit_before>#ipc_lista1.8
#Professor: Jucimar Junior
#Any Mendes Carvalho -
<commit_msg>Update ipc_lista1.8.py<commit_after>#ipc_lista1.8
#Professor: Jucimar Junior
#Any Mendes Carvalho - 16
<|endoftext|>


<commit_before># Copyright (C) 2015 SINTEF ICT,
# Applied Mathematics, Norway.
#
# Contact information:
# E-mail: eivind.fonn@sintef.no
# SINTEF ICT, Department of Applied Mathematics,
# P.O. Box 4760 Sluppen,
# 7045 Trondheim, Norway.
#
# This file is part of BADGER.
#
# BADGER is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# BADGER is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public
# License along with GoTools. If not, see
# <http://www.gnu.org/licenses/>.
#
# In accordance with Section 7(b) of the GNU Affero General Public
# License, a covered work must retain the producer line in every data
# file that is created or manipulated using GoTools.
#
# Other Usage
# You can be released from the requirements of the license by purchasing
# a commercial license. Buying such a license is mandatory as soon as you
# develop commercial activities involving BADGER without disclosing the
# source code of your own applications.
#
# This file may be used in accordance with the terms contained in a
# written agreement between you and SINTEF ICT.

import argparse
import shlex
import sys
import yaml

from collections import OrderedDict
from jinja2 import Template

import badger.output as output
import badger.log as log

def coerce_list(dictionary, key, split=None, required=False):
    if not required:
        if key not in dictionary:
            dictionary[key] = []
    if isinstance(dictionary[key], str):
        if isinstance(split, str):
            dictionary[key] = dictionary[key].split(split)
        elif split:
            dictionary[key] = split(dictionary[key])
        else:
            dictionary[key] = [dictionary[key]]

def parse_args(input=None):
    parser = argparse.ArgumentParser(description='Batch job runner.')
    parser.add_argument('-o', '--output', required=False, default='output.yaml',
                        help='The output file')
    parser.add_argument('-f', '--format', required=False, default=None,
                        choices=output.FORMATS, help='The output format')
    parser.add_argument('-d', '--dry', required=False, default=False,
                        action='store_true', help='Dry run')
    parser.add_argument('-v', '--verbosity', required=False, default=1, type=int,
                        choices=range(0, 5), help='Verbosity level for stdout')
    parser.add_argument('-l', '--logverbosity', required=False, default=2, type=int,
                        choices=range(0, 5), help='Verbosity level for log file')
    parser.add_argument('file', help='Configuration file for the batch job')
    args = parser.parse_args(input)

    if args.format is None:
        try:
            args.format = args.output.split('.')[-1]
            assert args.format in output.FORMATS
        except (AssertionError, IndexError):
            print('Unable to determine output format from filename "{}"'.format(args.output),
                  file=sys.stderr)
            sys.exit(1)

    log.stdout_verbosity = args.verbosity
    log.log_verbosity = args.logverbosity
    log.log_file = args.file + '.log'
    if args.logverbosity > 1:
        with open(log.log_file, 'w') as f: pass

    return args

# YAML is unordered by default, this is an ordered loader
# Thanks http://stackoverflow.com/a/21912744/2729168
def ordered_load(stream, Loader=yaml.Loader, object_pairs_hook=OrderedDict):
    class OrderedLoader(Loader):
        pass

    def construct_mapping(loader, node):
        loader.flatten_mapping(node)
        return object_pairs_hook(loader.construct_pairs(node))

    OrderedLoader.add_constructor(
        yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG,
        construct_mapping)
    return yaml.load(stream, OrderedLoader)

def treat_setup(setup):
    coerce_list(setup, 'templates')
    coerce_list(setup, 'files')
    coerce_list(setup, 'cmdargs', split=shlex.split)
    coerce_list(setup, 'executable', split=shlex.split, required=True)
    coerce_list(setup, 'parse')
    for key in ['dependencies', 'types', 'parameters']:
        if key not in setup:
            setup[key] = {}

    kwargs = {'variable_start_string': '$',
              'variable_end_string': '$'}
    for k in ['templates', 'files', 'executable', 'cmdargs']:
        setup[k] = [Template(v, **kwargs) for v in setup[k]]

    for key in setup['dependencies']:
        setup['dependencies'][key] = str(setup['dependencies'][key])

def load_setup(fn):
    with open(fn, 'r') as f:
        setup = ordered_load(f, yaml.SafeLoader)
    setup = setup or {}

    treat_setup(setup)

    return setup

def empty_setup(executable='', **kwargs):
    setup = {
        'templates': [],
        'files': [],
        'executable': executable,
        'cmdargs': [],
        'parameters': OrderedDict(),
        'dependencies': OrderedDict(),
        'parse': [],
        'types': OrderedDict(),
        }

    setup.update(kwargs)
    treat_setup(setup)
    return setup
<commit_msg>Write log in output directory<commit_after># Copyright (C) 2015 SINTEF ICT,
# Applied Mathematics, Norway.
#
# Contact information:
# E-mail: eivind.fonn@sintef.no
# SINTEF ICT, Department of Applied Mathematics,
# P.O. Box 4760 Sluppen,
# 7045 Trondheim, Norway.
#
# This file is part of BADGER.
#
# BADGER is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# BADGER is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public
# License along with GoTools. If not, see
# <http://www.gnu.org/licenses/>.
#
# In accordance with Section 7(b) of the GNU Affero General Public
# License, a covered work must retain the producer line in every data
# file that is created or manipulated using GoTools.
#
# Other Usage
# You can be released from the requirements of the license by purchasing
# a commercial license. Buying such a license is mandatory as soon as you
# develop commercial activities involving BADGER without disclosing the
# source code of your own applications.
#
# This file may be used in accordance with the terms contained in a
# written agreement between you and SINTEF ICT.

import argparse
import shlex
import sys
import yaml

from os.path import dirname, join
from collections import OrderedDict
from jinja2 import Template

import badger.output as output
import badger.log as log

def coerce_list(dictionary, key, split=None, required=False):
    if not required:
        if key not in dictionary:
            dictionary[key] = []
    if isinstance(dictionary[key], str):
        if isinstance(split, str):
            dictionary[key] = dictionary[key].split(split)
        elif split:
            dictionary[key] = split(dictionary[key])
        else:
            dictionary[key] = [dictionary[key]]

def parse_args(input=None):
    parser = argparse.ArgumentParser(description='Batch job runner.')
    parser.add_argument('-o', '--output', required=False, default='output.yaml',
                        help='The output file')
    parser.add_argument('-f', '--format', required=False, default=None,
                        choices=output.FORMATS, help='The output format')
    parser.add_argument('-d', '--dry', required=False, default=False,
                        action='store_true', help='Dry run')
    parser.add_argument('-v', '--verbosity', required=False, default=1, type=int,
                        choices=range(0, 5), help='Verbosity level for stdout')
    parser.add_argument('-l', '--logverbosity', required=False, default=2, type=int,
                        choices=range(0, 5), help='Verbosity level for log file')
    parser.add_argument('file', help='Configuration file for the batch job')
    args = parser.parse_args(input)

    if args.format is None:
        try:
            args.format = args.output.split('.')[-1]
            assert args.format in output.FORMATS
        except (AssertionError, IndexError):
            print('Unable to determine output format from filename "{}"'.format(args.output),
                  file=sys.stderr)
            sys.exit(1)

    log.stdout_verbosity = args.verbosity
    log.log_verbosity = args.logverbosity
    log.log_file = join(dirname(args.output), args.file + '.log')
    log.log_file = args.file + '.log'
    if args.logverbosity > 1:
        with open(log.log_file, 'w') as f: pass

    return args

# YAML is unordered by default, this is an ordered loader
# Thanks http://stackoverflow.com/a/21912744/2729168
def ordered_load(stream, Loader=yaml.Loader, object_pairs_hook=OrderedDict):
    class OrderedLoader(Loader):
        pass

    def construct_mapping(loader, node):
        loader.flatten_mapping(node)
        return object_pairs_hook(loader.construct_pairs(node))

    OrderedLoader.add_constructor(
        yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG,
        construct_mapping)
    return yaml.load(stream, OrderedLoader)

def treat_setup(setup):
    coerce_list(setup, 'templates')
    coerce_list(setup, 'files')
    coerce_list(setup, 'cmdargs', split=shlex.split)
    coerce_list(setup, 'executable', split=shlex.split, required=True)
    coerce_list(setup, 'parse')
    for key in ['dependencies', 'types', 'parameters']:
        if key not in setup:
            setup[key] = {}

    kwargs = {'variable_start_string': '$',
              'variable_end_string': '$'}
    for k in ['templates', 'files', 'executable', 'cmdargs']:
        setup[k] = [Template(v, **kwargs) for v in setup[k]]

    for key in setup['dependencies']:
        setup['dependencies'][key] = str(setup['dependencies'][key])

def load_setup(fn):
    with open(fn, 'r') as f:
        setup = ordered_load(f, yaml.SafeLoader)
    setup = setup or {}

    treat_setup(setup)

    return setup

def empty_setup(executable='', **kwargs):
    setup = {
        'templates': [],
        'files': [],
        'executable': executable,
        'cmdargs': [],
        'parameters': OrderedDict(),
        'dependencies': OrderedDict(),
        'parse': [],
        'types': OrderedDict(),
        }

    setup.update(kwargs)
    treat_setup(setup)
    return setup
<|endoftext|>


<commit_before><commit_msg>surfaces now as well<commit_after><|endoftext|>


<commit_before>import pandas as pd
import numpy as np
from bambi.external.six import string_types
from collections import OrderedDict, defaultdict
from bambi.utils import listify
from patsy import dmatrices, dmatrix
import warnings
from bambi.priors import PriorFactory
from copy import deepcopy
import re

def listify(obj):
    return obj if isinstance(obj, (list, tuple)) else [obj]

class Model(object):

    def __init__(self, data=None, intercept=False, backend='pymc3',
                 default_priors=None):
        '''
        Args:
            dataset (DataFrame): the pandas DF containing the data to use.
        '''

        self.default_priors = PriorFactory(default_priors)

        obj_cols = data.select_dtypes(['object']).columns
        data[obj_cols] = data[obj_cols].apply(lambda x: x.astype('category'))
        self.data = data
        # Some random effects stuff later requires us to make guesses about
        # column groupings into terms based on patsy's naming scheme.
        if re.search("[\[\]]+", ''.join(data.columns)):
            warnings.warn("At least one of the column names in the specified "
                          "dataset contain square brackets ('[' or ']')."
                          "This may cause unexpected behavior if you specify "
                          "models with random effects. You are encouraged to "
                          "rename your columns to avoid square brackets.")
        self.reset()

        if backend.lower() == 'pymc3':
            from bambi.backends import PyMC3BackEnd
            self.backend = PyMC3BackEnd()
        else:
            raise ValueError("At the moment, only the PyMC3 backend is supported.")

        if intercept:
            self.add_intercept()

    def reset(self):
        # self.cache = OrderedDict()
        self.contrasts = OrderedDict()
        self.terms = OrderedDict()
        self.y = None

    def build(self):
        ''' Build the PyMC3 model. '''
        if self.y is None:
            raise ValueError("No outcome (y) variable is set! Please call "
                             "set_y() before build() or fit().")
        for t in self.terms.values():
            t._setup()
        self.backend.build(self)
        self.built = True

    def fit(self, fixed=None, random=None, family='gaussian', link=None,
            run=True, **kwargs):
        if fixed is not None or random is not None:
            self.add_formula(fixed=fixed, random=random, append=False,
                             family=family, link=link)
        ''' Run the BackEnd to fit the model. '''
        if not self.built:
            warnings.warn("Current Bayesian model has not been built yet; "
              "building it first before sampling begins.")
            self.build()
        if run:
            return self.backend.run(self, **kwargs)

    def add_intercept(self):
        n = len(self.data)
        df = pd.DataFrame(np.ones((n, 1)), columns=['Intercept'])
        self.add_term('Intercept', df)

    def add_formula(self, fixed=None, random=None, append=False, priors=None,
                    categorical=None, family='gaussian', link=None):

        data = self.data

        if not append:
            self.reset()

        if fixed is not None:
            # Explicitly convert columns to category if desired--though this can
            # also be done within the formula using C().
            if categorical is not None:
                data = data.copy()
                categorical = listify(categorical)
                data[categorical] = data[categorical].apply(lambda x: x.astype('category'))

            if '~' in fixed:
                y, X = dmatrices(fixed, data=data)
                y_label = y.design_info.term_names[0]
                self.add_y(y_label, family=family, link=link)
            else:
                X = dmatrix(fixed, data=data)

            # Loop over predictor terms
            for _name, _slice in X.design_info.term_name_slices.items():
                cols = X.design_info.column_names[_slice]
                term_data = pd.DataFrame(X[:, _slice], columns=cols)
                self.add_term(_name, data=term_data)

        # Random effects
        if random is not None:
            random = listify(random)
            for f in random:
                kwargs = {'random': True}
                if re.search('[\*\(\)\+\-]+', f):
                    raise ValueError("Random term '%s' contains an invalid "
                        "character. Note that formula-style operators other "
                        "than | are not currently supported in random effects "
                        "specifications.")
                # '1|factor' is considered identical to 'factor'
                f = re.sub(r'^1\s+\|(.*)', r'\1', f).strip()
                if '|' not in f:
                    kwargs['categorical'] = True
                    kwargs['drop_first'] = False
                    variable = f
                else:
                    variable, split_by = re.split('\s*\|\s*', f)
                    kwargs['split_by'] = split_by
                self.add_term(variable=variable, label=f, **kwargs)

    def add_y(self, variable, family='gaussian', link=None, *args,
              **kwargs):

        if isinstance(family, string_types):
            family = self.default_priors.get(family=family)
        self.family = family

        # Override family's link if another is explicitly passed
        if link is not None:
            self.family.link = link

        prior = self.family.prior
        self.add_term(variable, prior=prior, *args, **kwargs)
        # use last-added term name b/c it could have been changed in add_term
        name = list(self.terms.values())[-1].name
        self.set_y(name)

    def add_term(self, variable, data=None, label=None, categorical=False,
                 random=False, split_by=None, prior=None, drop_first=True):
        ''' Create a new Term and add it to the current Model. All positional
        and keyword arguments are passed directly to the Term initializer. '''

        if data is None:
            data = self.data.copy()

        # Make sure user didn't forget to set categorical=True
        elif variable in data.columns and \
             data.loc[:, variable].dtype.name in ['object', 'category']:
             categorical = True

        if categorical:
            data[variable] = data[variable].astype('category')

        # Extract splitting variable
        if split_by is not None:
            data[split_by] = data[split_by].astype('category')
            split_by = listify(split_by)
            group_term = ':'.join(split_by)
            f = '0 + %s : (%s)' % (variable, group_term)
            data = dmatrix(f, data=data)
            cols = data.design_info.column_names
            data = pd.DataFrame(data, columns=cols)

            # For categorical random effects, one variance term per split_by level
            if random and categorical:
                split_data = {}
                groups = list(set([re.sub(r'^.*?\:', '', c) for c in cols]))
                for g in groups:
                    patt = re.escape(r':%s' % g) + '$'
                    level_data = data.filter(regex=patt)
                    level_data.columns = [c.split(':')[0] for c in level_data.columns]
                    level_data = level_data.loc[:, (level_data!=0).any(axis=0)]
                    split_data[g] = level_data.values
                data = split_data

        elif categorical or (variable in data.columns and \
                             data[variable].dtype.name in ['object', 'category']):
            data = pd.get_dummies(data[variable], drop_first=drop_first)
        else:
            # If all columns have identical names except for levels in [],
            # assume they've already been contrast-coded, and pass data as-is
            cols = [re.sub('\[.*?\]', '', c) for c in data.columns]
            if len(set(cols)) > 1:
                data = data[[variable]]

        if label is None:
            label = variable

        if random:
            term = RandomTerm(self, label, data, categorical=categorical,
                              prior=prior)
        else:
            term = Term(self, label, data, categorical=categorical,
                        prior=prior)
        self.terms[term.name] = term
        self.built = False

    def set_y(self, label):
        ''' Set the outcome variable. '''
        if self.y is not None:
            self.terms[self.y.label] = self.y
        self.y = self.terms.pop(label)
        self.built = False

class Term(object):

    type_ = 'fixed'

    def __init__(self, model, name, data, categorical=False, prior=None,
                 **kwargs):
        '''
        Args:
            model (Model): The associated Model instance.
            name (str): Name of the term.
            data (DataFrame, ndarray): The pandas DataFrame or numpy array from
                containing the data. If a DF is passed, the variable names
                are used to extract the target columns. If a numpy array is
                passed, all columns of the array are used, without selection.
            categorical (bool): If True, the source variable is interpreted as
                nominal/categorical. If False, the source variable is treated
                as continuous.
            prior (dict): A specification of the prior(s) to use.
                Must have keys for 'name' and 'args'; optionally, can also
                pass 'sigma', which is another dict with name/arg keys.
            kwargs: Optional keyword arguments passed to the model-building
                back-end.
        '''
        self.model = model
        self.name = name
        self.categorical = categorical
        self.prior = prior

        if isinstance(data, pd.Series):
            data = data.to_frame()
        if isinstance(data, pd.DataFrame):
            self.levels = list(data.columns)
            data = data.values
        elif isinstance(data, dict):
            pass   # Random effects pass through here
        else:
            data = np.atleast_2d(data)
            self.levels = list(range(data.shape[1]))

        self.data = data
        self.kwargs = kwargs

    def _setup(self):
        # TODO: come up with a more sensible way of getting/setting default priors
        if self.prior is None:
            # y_data = self.model.y.data
            # if self.name == 'Intercept':
            #     self.prior = deepcopy(default_priors['term']['intercept'])
            #     self.prior['args']['alpha'] = y_data.mean()
            #     self.prior['args']['beta'] *= y_data.std()
            # else:
            #     self.prior = deepcopy(default_priors['fixed'])
            #     max_std = self.data.std(0).max()
            #     self.prior['args']['sd'] *= max_std * 2 * y_data.std()
            term_type = 'intercept' if self.name == 'Intercept' else 'fixed'
            self.prior = self.model.default_priors.get(term=term_type)

class RandomTerm(Term):

    type_ = 'random'

    def __init__(self, model, name, data, yoke=False, prior=None, **kwargs):

        self.yoke = yoke
        super(RandomTerm, self).__init__(model, name, data, prior=prior, **kwargs)

    def _setup(self):
        if self.prior is None:
            self.prior = self.model.default_priors.get(term='random')

            # # Rescale prior sd--need to implement better heuristic
            # data = self.data

            # # nested terms are in dicts, so put non-nested terms in dummy dict
            # max_range = 0.
            # if not isinstance(data, dict):
            #     data = {'dummy': data}
            # for level in data.values():
            #     lev_range = level.mean(0).max() - level.mean(0).min()
            #     if lev_range > max_range:
            #         max_range = lev_range
            # scl = max(max_range, 1)
            # self.prior['sigma']['args']['beta'] *= (scl * 2 * self.model.y.data.std())
<commit_msg>Fix regular expression for stripping '1|' from random effects<commit_after>import pandas as pd
import numpy as np
from bambi.external.six import string_types
from collections import OrderedDict, defaultdict
from bambi.utils import listify
from patsy import dmatrices, dmatrix
import warnings
from bambi.priors import PriorFactory
from copy import deepcopy
import re

def listify(obj):
    return obj if isinstance(obj, (list, tuple)) else [obj]

class Model(object):

    def __init__(self, data=None, intercept=False, backend='pymc3',
                 default_priors=None):
        '''
        Args:
            dataset (DataFrame): the pandas DF containing the data to use.
        '''

        self.default_priors = PriorFactory(default_priors)

        obj_cols = data.select_dtypes(['object']).columns
        data[obj_cols] = data[obj_cols].apply(lambda x: x.astype('category'))
        self.data = data
        # Some random effects stuff later requires us to make guesses about
        # column groupings into terms based on patsy's naming scheme.
        if re.search("[\[\]]+", ''.join(data.columns)):
            warnings.warn("At least one of the column names in the specified "
                          "dataset contain square brackets ('[' or ']')."
                          "This may cause unexpected behavior if you specify "
                          "models with random effects. You are encouraged to "
                          "rename your columns to avoid square brackets.")
        self.reset()

        if backend.lower() == 'pymc3':
            from bambi.backends import PyMC3BackEnd
            self.backend = PyMC3BackEnd()
        else:
            raise ValueError("At the moment, only the PyMC3 backend is supported.")

        if intercept:
            self.add_intercept()

    def reset(self):
        # self.cache = OrderedDict()
        self.contrasts = OrderedDict()
        self.terms = OrderedDict()
        self.y = None

    def build(self):
        ''' Build the PyMC3 model. '''
        if self.y is None:
            raise ValueError("No outcome (y) variable is set! Please call "
                             "set_y() before build() or fit().")
        for t in self.terms.values():
            t._setup()
        self.backend.build(self)
        self.built = True

    def fit(self, fixed=None, random=None, family='gaussian', link=None,
            run=True, **kwargs):
        if fixed is not None or random is not None:
            self.add_formula(fixed=fixed, random=random, append=False,
                             family=family, link=link)
        ''' Run the BackEnd to fit the model. '''
        if not self.built:
            warnings.warn("Current Bayesian model has not been built yet; "
              "building it first before sampling begins.")
            self.build()
        if run:
            return self.backend.run(self, **kwargs)

    def add_intercept(self):
        n = len(self.data)
        df = pd.DataFrame(np.ones((n, 1)), columns=['Intercept'])
        self.add_term('Intercept', df)

    def add_formula(self, fixed=None, random=None, append=False, priors=None,
                    categorical=None, family='gaussian', link=None):

        data = self.data

        if not append:
            self.reset()

        if fixed is not None:
            # Explicitly convert columns to category if desired--though this can
            # also be done within the formula using C().
            if categorical is not None:
                data = data.copy()
                categorical = listify(categorical)
                data[categorical] = data[categorical].apply(lambda x: x.astype('category'))

            if '~' in fixed:
                y, X = dmatrices(fixed, data=data)
                y_label = y.design_info.term_names[0]
                self.add_y(y_label, family=family, link=link)
            else:
                X = dmatrix(fixed, data=data)

            # Loop over predictor terms
            for _name, _slice in X.design_info.term_name_slices.items():
                cols = X.design_info.column_names[_slice]
                term_data = pd.DataFrame(X[:, _slice], columns=cols)
                self.add_term(_name, data=term_data)

        # Random effects
        if random is not None:
            random = listify(random)
            for f in random:
                kwargs = {'random': True}
                if re.search('[\*\(\)\+\-]+', f):
                    raise ValueError("Random term '%s' contains an invalid "
                        "character. Note that formula-style operators other "
                        "than | are not currently supported in random effects "
                        "specifications.")
                # '1|factor' is considered identical to 'factor'
                f = re.sub(r'^1\s*\|(.*)', r'\1', f).strip()
                if '|' not in f:
                    kwargs['categorical'] = True
                    kwargs['drop_first'] = False
                    variable = f
                else:
                    variable, split_by = re.split('\s*\|\s*', f)
                    kwargs['split_by'] = split_by
                self.add_term(variable=variable, label=f, **kwargs)

    def add_y(self, variable, family='gaussian', link=None, *args,
              **kwargs):

        if isinstance(family, string_types):
            family = self.default_priors.get(family=family)
        self.family = family

        # Override family's link if another is explicitly passed
        if link is not None:
            self.family.link = link

        prior = self.family.prior
        self.add_term(variable, prior=prior, *args, **kwargs)
        # use last-added term name b/c it could have been changed in add_term
        name = list(self.terms.values())[-1].name
        self.set_y(name)

    def add_term(self, variable, data=None, label=None, categorical=False,
                 random=False, split_by=None, prior=None, drop_first=True):
        ''' Create a new Term and add it to the current Model. All positional
        and keyword arguments are passed directly to the Term initializer. '''

        if data is None:
            data = self.data.copy()

        # Make sure user didn't forget to set categorical=True
        elif variable in data.columns and \
             data.loc[:, variable].dtype.name in ['object', 'category']:
             categorical = True

        if categorical:
            data[variable] = data[variable].astype('category')

        # Extract splitting variable
        if split_by is not None:
            data[split_by] = data[split_by].astype('category')
            split_by = listify(split_by)
            group_term = ':'.join(split_by)
            f = '0 + %s : (%s)' % (variable, group_term)
            data = dmatrix(f, data=data)
            cols = data.design_info.column_names
            data = pd.DataFrame(data, columns=cols)

            # For categorical random effects, one variance term per split_by level
            if random and categorical:
                split_data = {}
                groups = list(set([re.sub(r'^.*?\:', '', c) for c in cols]))
                for g in groups:
                    patt = re.escape(r':%s' % g) + '$'
                    level_data = data.filter(regex=patt)
                    level_data.columns = [c.split(':')[0] for c in level_data.columns]
                    level_data = level_data.loc[:, (level_data!=0).any(axis=0)]
                    split_data[g] = level_data.values
                data = split_data

        elif categorical or (variable in data.columns and \
                             data[variable].dtype.name in ['object', 'category']):
            data = pd.get_dummies(data[variable], drop_first=drop_first)
        else:
            # If all columns have identical names except for levels in [],
            # assume they've already been contrast-coded, and pass data as-is
            cols = [re.sub('\[.*?\]', '', c) for c in data.columns]
            if len(set(cols)) > 1:
                data = data[[variable]]

        if label is None:
            label = variable

        if random:
            term = RandomTerm(self, label, data, categorical=categorical,
                              prior=prior)
        else:
            term = Term(self, label, data, categorical=categorical,
                        prior=prior)
        self.terms[term.name] = term
        self.built = False

    def set_y(self, label):
        ''' Set the outcome variable. '''
        if self.y is not None:
            self.terms[self.y.label] = self.y
        self.y = self.terms.pop(label)
        self.built = False

class Term(object):

    type_ = 'fixed'

    def __init__(self, model, name, data, categorical=False, prior=None,
                 **kwargs):
        '''
        Args:
            model (Model): The associated Model instance.
            name (str): Name of the term.
            data (DataFrame, ndarray): The pandas DataFrame or numpy array from
                containing the data. If a DF is passed, the variable names
                are used to extract the target columns. If a numpy array is
                passed, all columns of the array are used, without selection.
            categorical (bool): If True, the source variable is interpreted as
                nominal/categorical. If False, the source variable is treated
                as continuous.
            prior (dict): A specification of the prior(s) to use.
                Must have keys for 'name' and 'args'; optionally, can also
                pass 'sigma', which is another dict with name/arg keys.
            kwargs: Optional keyword arguments passed to the model-building
                back-end.
        '''
        self.model = model
        self.name = name
        self.categorical = categorical
        self.prior = prior

        if isinstance(data, pd.Series):
            data = data.to_frame()
        if isinstance(data, pd.DataFrame):
            self.levels = list(data.columns)
            data = data.values
        elif isinstance(data, dict):
            pass   # Random effects pass through here
        else:
            data = np.atleast_2d(data)
            self.levels = list(range(data.shape[1]))

        self.data = data
        self.kwargs = kwargs

    def _setup(self):
        # TODO: come up with a more sensible way of getting/setting default priors
        if self.prior is None:
            # y_data = self.model.y.data
            # if self.name == 'Intercept':
            #     self.prior = deepcopy(default_priors['term']['intercept'])
            #     self.prior['args']['alpha'] = y_data.mean()
            #     self.prior['args']['beta'] *= y_data.std()
            # else:
            #     self.prior = deepcopy(default_priors['fixed'])
            #     max_std = self.data.std(0).max()
            #     self.prior['args']['sd'] *= max_std * 2 * y_data.std()
            term_type = 'intercept' if self.name == 'Intercept' else 'fixed'
            self.prior = self.model.default_priors.get(term=term_type)

class RandomTerm(Term):

    type_ = 'random'

    def __init__(self, model, name, data, yoke=False, prior=None, **kwargs):

        self.yoke = yoke
        super(RandomTerm, self).__init__(model, name, data, prior=prior, **kwargs)

    def _setup(self):
        if self.prior is None:
            self.prior = self.model.default_priors.get(term='random')

            # # Rescale prior sd--need to implement better heuristic
            # data = self.data

            # # nested terms are in dicts, so put non-nested terms in dummy dict
            # max_range = 0.
            # if not isinstance(data, dict):
            #     data = {'dummy': data}
            # for level in data.values():
            #     lev_range = level.mean(0).max() - level.mean(0).min()
            #     if lev_range > max_range:
            #         max_range = lev_range
            # scl = max(max_range, 1)
            # self.prior['sigma']['args']['beta'] *= (scl * 2 * self.model.y.data.std())
<|endoftext|>


<commit_before>import os
import sys
import toolz as tz
from bcbio import utils
from bcbio.pipeline import config_utils
import bcbio.pipeline.datadict as dd
from bcbio.ngsalign import bowtie2, bwa
from bcbio.distributed.transaction import file_transaction
from bcbio.provenance import do
from bcbio.log import logger

def clean_chipseq_alignment(data):
    aligner = dd.get_aligner(data)
    data["align_bam"] = dd.get_work_bam(data)
    if aligner:
        if aligner == "bowtie2":
            filterer = bowtie2.filter_multimappers
        elif aligner == "bwa":
            filterer = bwa.filter_multimappers
        else:
            logger.error("ChIP-seq only supported for bowtie2 and bwa.")
            sys.exit(-1)
        unique_bam = filterer(dd.get_work_bam(data), data)
        data["work_bam"] = unique_bam
    else:
        logger.info("Warning: When BAM file is given as input, bcbio skips multimappers removal."
                    "If BAM is not cleaned for peak calling, can result in downstream errors.")
    # lcr_bed = utils.get_in(data, ("genome_resources", "variation", "lcr"))
    encode_bed = tz.get_in(["genome_resources", "variation", "encode_blacklist"], data)
    if encode_bed:
        data["work_bam"] = _prepare_bam(data["work_bam"], encode_bed, data['config'])
    data["bigwig"] = _bam_coverage(dd.get_sample_name(data), dd.get_work_bam(data), data)
    return [[data]]

def _prepare_bam(bam_file, bed_file, config):
    """Remove regions from bed files"""
    if not bam_file or not bed_file:
        return bam_file
    out_file = utils.append_stem(bam_file, '_filter')
    bedtools = config_utils.get_program("bedtools", config)
    if not utils.file_exists(out_file):
        with file_transaction(out_file) as tx_out:
            cmd = "{bedtools} subtract -nonamecheck -A -a {bam_file} -b {bed_file} > {tx_out}"
            do.run(cmd.format(**locals()), "Clean %s" % bam_file)
    return out_file

def  get_genome(genome):
    from bcbio.chipseq import macs2
    loaded = macs2.HS
    if genome in loaded:
        return loaded[genome]

def _bam_coverage(name, bam_input, data):
    """Run bamCoverage from deeptools"""
    cmd = ("{bam_coverage} -b {bam_input} -o {bw_output} "
          "--binSize 20 --effectiveGenomeSize {size} "
          "--smoothLength 60 --extendReads 150 --centerReads -p {cores}")
    size = int(get_genome(dd.get_genome_build(data)))
    cores = dd.get_num_cores(data)
    try:
        bam_coverage = config_utils.get_program("bamCoverage", data)
    except config_utils.CmdNotFound:
        logger.info("No bamCoverage found, skipping bamCoverage.")
        return None
    bw_output = os.path.join(os.path.dirname(bam_input), "%s.bw" % name)
    if utils.file_exists(bw_output):
        return bw_output
    with file_transaction(bw_output) as out_tx:
        do.run(cmd.format(**locals()), "Run bamCoverage in %s" % name)
    return bw_output
<commit_msg>CHIP:fix index generation after filtering<commit_after>import os
import sys
import toolz as tz
from bcbio import utils
from bcbio import bam
from bcbio.pipeline import config_utils
import bcbio.pipeline.datadict as dd
from bcbio.ngsalign import bowtie2, bwa
from bcbio.distributed.transaction import file_transaction
from bcbio.provenance import do
from bcbio.log import logger

def clean_chipseq_alignment(data):
    aligner = dd.get_aligner(data)
    data["align_bam"] = dd.get_work_bam(data)
    if aligner:
        if aligner == "bowtie2":
            filterer = bowtie2.filter_multimappers
        elif aligner == "bwa":
            filterer = bwa.filter_multimappers
        else:
            logger.error("ChIP-seq only supported for bowtie2 and bwa.")
            sys.exit(-1)
        unique_bam = filterer(dd.get_work_bam(data), data)
        data["work_bam"] = unique_bam
    else:
        logger.info("Warning: When BAM file is given as input, bcbio skips multimappers removal."
                    "If BAM is not cleaned for peak calling, can result in downstream errors.")
    # lcr_bed = utils.get_in(data, ("genome_resources", "variation", "lcr"))
    encode_bed = tz.get_in(["genome_resources", "variation", "encode_blacklist"], data)
    if encode_bed:
        data["work_bam"] = _prepare_bam(data["work_bam"], encode_bed, data['config'])
        bam.index(data["work_bam"], data['config'])
    data["bigwig"] = _bam_coverage(dd.get_sample_name(data), dd.get_work_bam(data), data)
    return [[data]]

def _prepare_bam(bam_file, bed_file, config):
    """Remove regions from bed files"""
    if not bam_file or not bed_file:
        return bam_file
    out_file = utils.append_stem(bam_file, '_filter')
    bedtools = config_utils.get_program("bedtools", config)
    if not utils.file_exists(out_file):
        with file_transaction(out_file) as tx_out:
            cmd = "{bedtools} subtract -nonamecheck -A -a {bam_file} -b {bed_file} > {tx_out}"
            do.run(cmd.format(**locals()), "Clean %s" % bam_file)
    return out_file

def  get_genome(genome):
    from bcbio.chipseq import macs2
    loaded = macs2.HS
    if genome in loaded:
        return loaded[genome]

def _bam_coverage(name, bam_input, data):
    """Run bamCoverage from deeptools"""
    cmd = ("{bam_coverage} -b {bam_input} -o {bw_output} "
          "--binSize 20 --effectiveGenomeSize {size} "
          "--smoothLength 60 --extendReads 150 --centerReads -p {cores}")
    size = int(get_genome(dd.get_genome_build(data)))
    cores = dd.get_num_cores(data)
    try:
        bam_coverage = config_utils.get_program("bamCoverage", data)
    except config_utils.CmdNotFound:
        logger.info("No bamCoverage found, skipping bamCoverage.")
        return None
    bw_output = os.path.join(os.path.dirname(bam_input), "%s.bw" % name)
    if utils.file_exists(bw_output):
        return bw_output
    with file_transaction(bw_output) as out_tx:
        do.run(cmd.format(**locals()), "Run bamCoverage in %s" % name)
    return bw_output
<|endoftext|>