<commit_before>#! /usr/bin/env python

##############################################################################
##  pyvolve: Python platform for simulating evolutionary sequences.
##
##  Written by Stephanie J. Spielman (stephanie.spielman@gmail.com)
##############################################################################

'''
Evolve sequences along a phylogeny.

TODO: Provide output file for gamma rate categories. In cases of branch het, the model for each set of categories should be indicated.
'''

import itertools
from copy import deepcopy
import numpy as np
from scipy import linalg
import random as rn
from misc import *
MOLECULES = Genetics()

class Evolver(object):
    '''
        Class to evolve sequences along a phylogeny.
        Currently supported:
            1. Site heterogeneity, among and within partitions
            2. Branch heterogeneity (via model flags in phylogeny, similar to approach used by Indelible)

        Coming soon:
            1. Indels
    '''

    def __init__(self, **kwargs):
        '''
            Required arguments:
                1. *tree* is the phylogeny along which we will evolve
                2. *partitions* is a list of Partition() objects to evolve

            Optional arguments:
                1. *seqfile* is an output file for saving final simulated sequences
                2. *seqfmt* is the format for seqfile (either fasta, nexus, phylip, phylip-relaxed, stockholm, etc. Anything that Biopython can accept!!) Default is FASTA.
                3. *write_anc* is a bool for whether ancestral sequences should be output. If so, they are output with the tip sequences in seqfile. Default is False.
                4. *ratefile* is an output file for saving rate information about each simulated column. For codon sequences, saves dN and dS. For nucleotide and amino acid sequences, saves heterogeneity info (if applicable)
        '''

        self.partitions = kwargs.get('partitions', None)
        self.full_tree  = kwargs.get('tree', Tree())
        self.seqfile    = kwargs.get('seqfile', None)
        self.seqfmt     = kwargs.get('seqfmt', 'fasta').lower()
        self.write_anc  = kwargs.get('write_anc', False)
        self.ratefile   = kwargs.get('ratefile', None)

        # These dictionaries enable convenient post-processing of the simulated alignment. Otherwise we'd have to always loop over full tree, which would be very slow.
        self.leaf_seqs = {} # Store final tip sequences only
        self.evolved_seqs = {} # Stores sequences from all nodes, including internal and tips

        # Setup and sanity checks
        self._setup_partitions()
        self._set_code()

    def _setup_partitions(self):
        '''
            Setup and various sanity checks.
        '''
        if isinstance(self.partitions, Partition):
            self.partitions = [self.partitions]
        else:
            assert(type(self.partitions) is list), "\n\nYou must provide either a single Partition object or list of Partition objects to evolver."
            for p in self.partitions:
                assert(isinstance(p, Partition)), "\n\nYou must provide either a single Partition object or list of Partition objects to evolver."

        self._root_seq_length = 0
        for part in self.partitions:

            ################ Set up branch heterogeneity, if specified ################
            # Yes branch heterogeneity -> sanity check the (hopefully) specified root_model, and yell if not assigned or assigned incorrectly.
            if part.branch_het():
                for m in part.models:
                    if m.name == part.root_model:
                        self.full_tree.model_flag = part.root_model
                        assert(self.full_tree.model_flag is not None), "\n\n Your root_model does not correspond to any of the Model()/CodonModel() objects provided to your Partition() objects."
            else:
                part.models = [part.models]
                part.root_model = None

            ################ Set up rate heterogeneity, if specified ################
            # Site-rate classes will not change, even with branch heterogeneity, so can simply pick first model.
            m = part.models[0]
            full = part.size
            # No rate heterogeneity. Turn part.size into list of length 1. No need to shuffle.
            if m.num_classes() == 1:
                part.size = [part.size]

            # Yes rate heterogeneity. Divvy up part.size into rate het chunks, and set shuffle to True
            else:
                part.shuffle = True
                remaining = part.size
                part.size = []
                for i in range(len(m.rates) - 1): # don't fill in last one yet since rounding issues will occur.
                    section = int( m.probs[i] * full )
                    part.size.append( section )
                    remaining -= section
                part.size.append(remaining)
            assert( sum(part.size) ==  full ), "\n\nImproperly divvied up rate heterogeneity."
            self._root_seq_length += full

        ################ Final check on size ################
        assert(self._root_seq_length > 0), "\n\nPartitions have no size!"

    def _set_code(self):
        '''
            Assign genetic code.
        '''
        dim = self.partitions[0].models[0].params['state_freqs'].shape[0]
        if dim == 4:
            self._code = MOLECULES.nucleotides
        elif dim == 20:
            self._code = MOLECULES.amino_acids
        elif dim == 61:
            self._code = MOLECULES.codons
        else:
            raise AssertionError("This should never be reached.")

    def __call__(self):
        '''
            Run evolver! Perform simulation, any necessary post-processing, and save sequences and/or other info to appropriate files.
        '''

        # Simulate recursively
        self.sim_subtree(self.full_tree)

        # Shuffle sequences?
        self._shuffle_sites()

        # Save rate info
        if self.ratefile is not None:
            self.write_ratefile()

        # Save sequences?
        if self.seqfile is not None:
            if self.write_anc:
                self.write_sequences(self.seqfile, self.seqfmt, self.evolved_seqs)
            else:
                self.write_sequences(self.seqfile, self.seqfmt, self.leaf_seqs)
    #########################################################################################

    ######################## FUNCTIONS TO PROCESS SIMULATED SEQUENCES #######################
    def _site_to_sequence(self, site):
        '''
            Convert a single Site() object or list of Site() objects into a sequence string.
            Argument *site* is either a Site() object or a list of them.
        '''
        if type(site) is list:
            sequence = ''
            for s in site:
                sequence += self._code[s.int_seq]
        else:
            sequence = self._code[site.int_seq]
        return sequence

    def _shuffle_sites(self):
        '''
            Shuffle evolved sequences within partitions, if specified.
            In particular, we shuffle sequences in the self.evolved_seqs dictionary, and then we copy over to the self.leaf_seqs dictionary.

            LATER: only shuffle sites which have Site.origin == "root". Inserted sites don't need to be shuffled.

        '''
        start = 0
        for part_index in range( len(self.partitions) ):
            part = self.partitions[part_index]
            if part.shuffle:
                size = sum( part.size )
                part_pos = np.arange( size ) + start
                np.random.shuffle(part_pos)
                for record in self.evolved_seqs:
                    temp = []
                    for pp in part_pos:
                        temp.append( self.evolved_seqs[record][part_index][pp] )
                    self.evolved_seqs[record][part_index][start:start + size] = temp

        # Apply shuffling to self.leaf_seqs
        for record in self.leaf_seqs:
            self.leaf_seqs[record] = self.evolved_seqs[record]

    def write_sequences(self, outfile, fmt, seqdict):
        '''
            Write resulting sequences (seqdict - this is either just the tips or all nodes) to a file in specified format.
            Arguments:
                1. "outfile" is the name of the file for saving the alignment
                2. "fmt" is the alignment output file format (either fasta, nexus, phylip, phylip-relaxed, stockholm, etc. Anything that Biopython can accept!!) If not provided, will output in fasta format.
        '''
        from Bio.Seq import Seq
        from Bio.SeqRecord import SeqRecord
        from Bio.Alphabet import generic_alphabet
        from Bio import SeqIO

        alignment = []
        for entry in seqdict:
            merged_entry = list( itertools.chain.from_iterable( seqdict[entry] ) )
            sequence = self._site_to_sequence( merged_entry )
            seq_object = SeqRecord( Seq( sequence , generic_alphabet ), id = entry, description = "")
            alignment.append(seq_object)
        try:
            SeqIO.write(alignment, outfile, fmt)
        except:
            raise AssertionError("\n Output file format is unknown. Consult with Biopython manual to see which I/O formats are accepted.")

    def write_ratefile(self):
        '''
            Write ratefile, a tab-delimited file containing site-specific rate information. Considers leaf sequences only.
            Writes -   Site_Index    Partition_Index     Rate_Category
            All indexing is from *1*.
        '''
        refseq = self.leaf_seqs.values()[0]
        with open(self.ratefile, 'w') as ratef:
            ratef.write("Site_Index\tPartition_Index\tRate_Category")
            site_index = 1
            for p in range(len(refseq)):
                part = refseq[p]
                for i in range(len(part)):
                    w = "\n" + str(site_index) + "\t" + str(p +  1) + "\t" + str(part[i].rate + 1)
                    ratef.write(w)
                    site_index += 1

    ######################### FUNCTIONS INVOLVED IN SEQUENCE EVOLUTION ############################
    def _obtain_model(self, part, flag):
        '''
            Obtain the appropriate Model()/CodonModel() for evolution along a particular branch.
        '''
        my_model = None
        if part.branch_het():
            my_model = None
            for m in part.models:
                if m.name == flag:
                    my_model = m
                    break
        else:
            my_model = part.models[0]
        assert( my_model is not None ), "\n\nCould not retrieve model a particular branch's evolution."
        return my_model

    def _generate_prob_from_unif(self, prob_array):
        '''
            Sample a sequence (nuc,aa,or codon), and return an integer for the sequence chosen from a uniform distribution.
            Arugment *prob_array* is any list and/or numpy array of probabilities which sum to 1.
        '''

        assert ( abs(np.sum(prob_array) - 1.) < ZERO), "Probabilities do not sum to 1. Cannot generate a new sequence."
        r = rn.uniform(0,1)
        i = 0
        sum = prob_array[i]
        while sum < r:
            i += 1
            sum += prob_array[i]
        return i

    def _generate_root_seq(self):
        '''
            Generate a root sequence based on the stationary frequencies or using a provided sequence
            Return a complete root sequence list of Site objects.
        '''

        root_sequence = [] # This will contain a list for each partition's sequence (which is itself a list of Site() objects)

        for part in self.partitions:

            # Grab model info for this partition to get frequency vector for root simulation
            root_model = self._obtain_model(part, self.full_tree.model_flag)

            # Generate root_sequence and assign the Site a rate class
            part_root = []
            for i in range( root_model.num_classes() ):
                for j in range( part.size[i] ):
                    new_site = Site()
                    new_site.rate = i
                    new_site.int_seq = self._generate_prob_from_unif( root_model.params['state_freqs'] )
                    part_root.append( new_site )
            assert( len(part_root) == sum(part.size) ), "\n\nRoot sequence improperly generated for a partition, evolution cannot happen."
            root_sequence.append(part_root)
        return root_sequence

    def sim_subtree(self, current_node, parent_node = None):
        '''
            Function to traverse a Tree object recursively and simulate sequences.
            Arguments:
                1. *current_node* is the node (either internal node or leaf) TO WHICH we evolving
                2. *parent_node* is the node we are evolving FROM. Default of None is only called when the root sequence is not yet made.
        '''

        # We are at the base and must generate root sequence
        if (parent_node is None):
            current_node.seq = self._generate_root_seq() # the .seq attribute is actually a list of Site() objects.
            self.evolved_seqs['root'] = current_node.seq
        else:
            current_node.seq = self.evolve_branch(current_node, parent_node)
            self.evolved_seqs[current_node.name] = current_node.seq

        # We are at an internal node. Keep evolving
        if len(current_node.children)>0:
            for child_node in current_node.children:
                self.sim_subtree(child_node, current_node)

        # We are at a leaf. Save the final sequence
        else:
            self.leaf_seqs[current_node.name] = current_node.seq

    def _check_parent_branch(self, parent_node, current_node):
        '''
            Function ensures that, for a given node we'd like to evolve to, an appropriate branch length exists.
            If the branch length is acceptable, an evolutionary model is then assigned to the node if one has yet been assigned.

            Arguments:
                1. *parent_node* is node FROM which we evolve
                2. *current_node* is the node (either internal node or leaf) TO WHICH we evolve
        '''
        assert (parent_node.seq != None), "\n\nThere is no parent sequence from which to evolve!"
        assert (current_node.branch_length > 0), "\n\n Your tree has a negative branch length. I'm going to quit now."
        if current_node.model_flag is None:
            current_node.model_flag = parent_node.model_flag

    def evolve_branch(self, current_node, parent_node):
        '''
            Function to evolve a given sequence during tree traversal.
            Arguments:
                1. *current_node* is the node (either internal node or leaf) we are evolving TO
                2. *parent_node* is the node we are evolving FROM.
        '''

        # Ensure parent sequence exists and branch length is acceptable. Return the model flag to use here.
        self._check_parent_branch(parent_node, current_node)

        # Evolve only if branch length is greater than 0 (1e-8).
        if current_node.branch_length <= ZERO:
            new_seq = deepcopy(parent_node.seq)

        else:
            new_seq = []
            for p in range( len(self.partitions) ):

                # Obtain current model
                part = self.partitions[p]
                current_model = self._obtain_model(part, current_node.model_flag)

                index = 0
                temp_new_seq = []
                for i in range( current_model.num_classes() ):
                    # Grab instantaneous rate matrix, which is done differently depending if codon (dN/dS) model or not. This is the rate het in the partition.
                    inst_matrix = None
                    if part.codon_model():
                        inst_matrix = current_model.matrices[i]
                    else:
                        inst_matrix = current_model.matrix * current_model.rates[i] # note that rates = [1.] if no site heterogeneity, so matrix unchanged
                    assert( inst_matrix is not None ), "\n\nCouldn't retrieve instantaneous rate matrix."

                    # Generate transition matrix and assert correct
                    prob_matrix = linalg.expm( np.multiply(inst_matrix, float(current_node.branch_length) ) )
                    assert( np.allclose( np.sum(prob_matrix, axis = 1), np.ones(len(self._code))) ), "Rows in transition matrix do not each sum to 1."

                    # Evolve branch
                    part_parent_seq = parent_node.seq[p][index : index + part.size[i]]
                    for j in range( part.size[i] ):
                        new_site = deepcopy( part_parent_seq[j] )
                        new_site.int_seq = self._generate_prob_from_unif( prob_matrix[ new_site.int_seq ] )
                        temp_new_seq.append( new_site )
                        index += 1
                new_seq.append( temp_new_seq )

        return new_seq

<commit_msg>an assertion was at the wrong indent level<commit_after>#! /usr/bin/env python

##############################################################################
##  pyvolve: Python platform for simulating evolutionary sequences.
##
##  Written by Stephanie J. Spielman (stephanie.spielman@gmail.com)
##############################################################################

'''
Evolve sequences along a phylogeny.

TODO: Provide output file for gamma rate categories. In cases of branch het, the model for each set of categories should be indicated.
'''

import itertools
from copy import deepcopy
import numpy as np
from scipy import linalg
import random as rn
from misc import *
MOLECULES = Genetics()

class Evolver(object):
    '''
        Class to evolve sequences along a phylogeny.
        Currently supported:
            1. Site heterogeneity, among and within partitions
            2. Branch heterogeneity (via model flags in phylogeny, similar to approach used by Indelible)

        Coming soon:
            1. Indels
    '''

    def __init__(self, **kwargs):
        '''
            Required arguments:
                1. *tree* is the phylogeny along which we will evolve
                2. *partitions* is a list of Partition() objects to evolve

            Optional arguments:
                1. *seqfile* is an output file for saving final simulated sequences
                2. *seqfmt* is the format for seqfile (either fasta, nexus, phylip, phylip-relaxed, stockholm, etc. Anything that Biopython can accept!!) Default is FASTA.
                3. *write_anc* is a bool for whether ancestral sequences should be output. If so, they are output with the tip sequences in seqfile. Default is False.
                4. *ratefile* is an output file for saving rate information about each simulated column. For codon sequences, saves dN and dS. For nucleotide and amino acid sequences, saves heterogeneity info (if applicable)
        '''

        self.partitions = kwargs.get('partitions', None)
        self.full_tree  = kwargs.get('tree', Tree())
        self.seqfile    = kwargs.get('seqfile', None)
        self.seqfmt     = kwargs.get('seqfmt', 'fasta').lower()
        self.write_anc  = kwargs.get('write_anc', False)
        self.ratefile   = kwargs.get('ratefile', None)

        # These dictionaries enable convenient post-processing of the simulated alignment. Otherwise we'd have to always loop over full tree, which would be very slow.
        self.leaf_seqs = {} # Store final tip sequences only
        self.evolved_seqs = {} # Stores sequences from all nodes, including internal and tips

        # Setup and sanity checks
        self._setup_partitions()
        self._set_code()

    def _setup_partitions(self):
        '''
            Setup and various sanity checks.
        '''
        if isinstance(self.partitions, Partition):
            self.partitions = [self.partitions]
        else:
            assert(type(self.partitions) is list), "\n\nYou must provide either a single Partition object or list of Partition objects to evolver."
            for p in self.partitions:
                assert(isinstance(p, Partition)), "\n\nYou must provide either a single Partition object or list of Partition objects to evolver."

        self._root_seq_length = 0
        for part in self.partitions:

            ################ Set up branch heterogeneity, if specified ################
            # Yes branch heterogeneity -> sanity check the (hopefully) specified root_model, and yell if not assigned or assigned incorrectly.
            if part.branch_het():
                for m in part.models:
                    if m.name == part.root_model:
                        self.full_tree.model_flag = part.root_model
                assert(self.full_tree.model_flag is not None), "\n\n Your root_model does not correspond to any of the Model()/CodonModel() objects provided to your Partition() objects."
            else:
                part.models = [part.models]
                part.root_model = None

            ################ Set up rate heterogeneity, if specified ################
            # Site-rate classes will not change, even with branch heterogeneity, so can simply pick first model.
            m = part.models[0]
            full = part.size
            # No rate heterogeneity. Turn part.size into list of length 1. No need to shuffle.
            if m.num_classes() == 1:
                part.size = [part.size]

            # Yes rate heterogeneity. Divvy up part.size into rate het chunks, and set shuffle to True
            else:
                part.shuffle = True
                remaining = part.size
                part.size = []
                for i in range(len(m.rates) - 1): # don't fill in last one yet since rounding issues will occur.
                    section = int( m.probs[i] * full )
                    part.size.append( section )
                    remaining -= section
                part.size.append(remaining)
            assert( sum(part.size) ==  full ), "\n\nImproperly divvied up rate heterogeneity."
            self._root_seq_length += full

        ################ Final check on size ################
        assert(self._root_seq_length > 0), "\n\nPartitions have no size!"

    def _set_code(self):
        '''
            Assign genetic code.
        '''
        dim = self.partitions[0].models[0].params['state_freqs'].shape[0]
        if dim == 4:
            self._code = MOLECULES.nucleotides
        elif dim == 20:
            self._code = MOLECULES.amino_acids
        elif dim == 61:
            self._code = MOLECULES.codons
        else:
            raise AssertionError("This should never be reached.")

    def __call__(self):
        '''
            Run evolver! Perform simulation, any necessary post-processing, and save sequences and/or other info to appropriate files.
        '''

        # Simulate recursively
        self.sim_subtree(self.full_tree)

        # Shuffle sequences?
        self._shuffle_sites()

        # Save rate info
        if self.ratefile is not None:
            self.write_ratefile()

        # Save sequences?
        if self.seqfile is not None:
            if self.write_anc:
                self.write_sequences(self.seqfile, self.seqfmt, self.evolved_seqs)
            else:
                self.write_sequences(self.seqfile, self.seqfmt, self.leaf_seqs)
    #########################################################################################

    ######################## FUNCTIONS TO PROCESS SIMULATED SEQUENCES #######################
    def _site_to_sequence(self, site):
        '''
            Convert a single Site() object or list of Site() objects into a sequence string.
            Argument *site* is either a Site() object or a list of them.
        '''
        if type(site) is list:
            sequence = ''
            for s in site:
                sequence += self._code[s.int_seq]
        else:
            sequence = self._code[site.int_seq]
        return sequence

    def _shuffle_sites(self):
        '''
            Shuffle evolved sequences within partitions, if specified.
            In particular, we shuffle sequences in the self.evolved_seqs dictionary, and then we copy over to the self.leaf_seqs dictionary.

            LATER: only shuffle sites which have Site.origin == "root". Inserted sites don't need to be shuffled.

        '''
        start = 0
        for part_index in range( len(self.partitions) ):
            part = self.partitions[part_index]
            if part.shuffle:
                size = sum( part.size )
                part_pos = np.arange( size ) + start
                np.random.shuffle(part_pos)
                for record in self.evolved_seqs:
                    temp = []
                    for pp in part_pos:
                        temp.append( self.evolved_seqs[record][part_index][pp] )
                    self.evolved_seqs[record][part_index][start:start + size] = temp

        # Apply shuffling to self.leaf_seqs
        for record in self.leaf_seqs:
            self.leaf_seqs[record] = self.evolved_seqs[record]

    def write_sequences(self, outfile, fmt, seqdict):
        '''
            Write resulting sequences (seqdict - this is either just the tips or all nodes) to a file in specified format.
            Arguments:
                1. "outfile" is the name of the file for saving the alignment
                2. "fmt" is the alignment output file format (either fasta, nexus, phylip, phylip-relaxed, stockholm, etc. Anything that Biopython can accept!!) If not provided, will output in fasta format.
        '''
        from Bio.Seq import Seq
        from Bio.SeqRecord import SeqRecord
        from Bio.Alphabet import generic_alphabet
        from Bio import SeqIO

        alignment = []
        for entry in seqdict:
            merged_entry = list( itertools.chain.from_iterable( seqdict[entry] ) )
            sequence = self._site_to_sequence( merged_entry )
            seq_object = SeqRecord( Seq( sequence , generic_alphabet ), id = entry, description = "")
            alignment.append(seq_object)
        try:
            SeqIO.write(alignment, outfile, fmt)
        except:
            raise AssertionError("\n Output file format is unknown. Consult with Biopython manual to see which I/O formats are accepted.")

    def write_ratefile(self):
        '''
            Write ratefile, a tab-delimited file containing site-specific rate information. Considers leaf sequences only.
            Writes -   Site_Index    Partition_Index     Rate_Category
            All indexing is from *1*.
        '''
        refseq = self.leaf_seqs.values()[0]
        with open(self.ratefile, 'w') as ratef:
            ratef.write("Site_Index\tPartition_Index\tRate_Category")
            site_index = 1
            for p in range(len(refseq)):
                part = refseq[p]
                for i in range(len(part)):
                    w = "\n" + str(site_index) + "\t" + str(p +  1) + "\t" + str(part[i].rate + 1)
                    ratef.write(w)
                    site_index += 1

    ######################### FUNCTIONS INVOLVED IN SEQUENCE EVOLUTION ############################
    def _obtain_model(self, part, flag):
        '''
            Obtain the appropriate Model()/CodonModel() for evolution along a particular branch.
        '''
        my_model = None
        if part.branch_het():
            my_model = None
            for m in part.models:
                if m.name == flag:
                    my_model = m
                    break
        else:
            my_model = part.models[0]
        assert( my_model is not None ), "\n\nCould not retrieve model a particular branch's evolution."
        return my_model

    def _generate_prob_from_unif(self, prob_array):
        '''
            Sample a sequence (nuc,aa,or codon), and return an integer for the sequence chosen from a uniform distribution.
            Arugment *prob_array* is any list and/or numpy array of probabilities which sum to 1.
        '''

        assert ( abs(np.sum(prob_array) - 1.) < ZERO), "Probabilities do not sum to 1. Cannot generate a new sequence."
        r = rn.uniform(0,1)
        i = 0
        sum = prob_array[i]
        while sum < r:
            i += 1
            sum += prob_array[i]
        return i

    def _generate_root_seq(self):
        '''
            Generate a root sequence based on the stationary frequencies or using a provided sequence
            Return a complete root sequence list of Site objects.
        '''

        root_sequence = [] # This will contain a list for each partition's sequence (which is itself a list of Site() objects)

        for part in self.partitions:

            # Grab model info for this partition to get frequency vector for root simulation
            root_model = self._obtain_model(part, self.full_tree.model_flag)

            # Generate root_sequence and assign the Site a rate class
            part_root = []
            for i in range( root_model.num_classes() ):
                for j in range( part.size[i] ):
                    new_site = Site()
                    new_site.rate = i
                    new_site.int_seq = self._generate_prob_from_unif( root_model.params['state_freqs'] )
                    part_root.append( new_site )
            assert( len(part_root) == sum(part.size) ), "\n\nRoot sequence improperly generated for a partition, evolution cannot happen."
            root_sequence.append(part_root)
        return root_sequence

    def sim_subtree(self, current_node, parent_node = None):
        '''
            Function to traverse a Tree object recursively and simulate sequences.
            Arguments:
                1. *current_node* is the node (either internal node or leaf) TO WHICH we evolving
                2. *parent_node* is the node we are evolving FROM. Default of None is only called when the root sequence is not yet made.
        '''

        # We are at the base and must generate root sequence
        if (parent_node is None):
            current_node.seq = self._generate_root_seq() # the .seq attribute is actually a list of Site() objects.
            self.evolved_seqs['root'] = current_node.seq
        else:
            current_node.seq = self.evolve_branch(current_node, parent_node)
            self.evolved_seqs[current_node.name] = current_node.seq

        # We are at an internal node. Keep evolving
        if len(current_node.children)>0:
            for child_node in current_node.children:
                self.sim_subtree(child_node, current_node)

        # We are at a leaf. Save the final sequence
        else:
            self.leaf_seqs[current_node.name] = current_node.seq

    def _check_parent_branch(self, parent_node, current_node):
        '''
            Function ensures that, for a given node we'd like to evolve to, an appropriate branch length exists.
            If the branch length is acceptable, an evolutionary model is then assigned to the node if one has yet been assigned.

            Arguments:
                1. *parent_node* is node FROM which we evolve
                2. *current_node* is the node (either internal node or leaf) TO WHICH we evolve
        '''
        assert (parent_node.seq != None), "\n\nThere is no parent sequence from which to evolve!"
        assert (current_node.branch_length > 0), "\n\n Your tree has a negative branch length. I'm going to quit now."
        if current_node.model_flag is None:
            current_node.model_flag = parent_node.model_flag

    def evolve_branch(self, current_node, parent_node):
        '''
            Function to evolve a given sequence during tree traversal.
            Arguments:
                1. *current_node* is the node (either internal node or leaf) we are evolving TO
                2. *parent_node* is the node we are evolving FROM.
        '''

        # Ensure parent sequence exists and branch length is acceptable. Return the model flag to use here.
        self._check_parent_branch(parent_node, current_node)

        # Evolve only if branch length is greater than 0 (1e-8).
        if current_node.branch_length <= ZERO:
            new_seq = deepcopy(parent_node.seq)

        else:
            new_seq = []
            for p in range( len(self.partitions) ):

                # Obtain current model
                part = self.partitions[p]
                current_model = self._obtain_model(part, current_node.model_flag)

                index = 0
                temp_new_seq = []
                for i in range( current_model.num_classes() ):
                    # Grab instantaneous rate matrix, which is done differently depending if codon (dN/dS) model or not. This is the rate het in the partition.
                    inst_matrix = None
                    if part.codon_model():
                        inst_matrix = current_model.matrices[i]
                    else:
                        inst_matrix = current_model.matrix * current_model.rates[i] # note that rates = [1.] if no site heterogeneity, so matrix unchanged
                    assert( inst_matrix is not None ), "\n\nCouldn't retrieve instantaneous rate matrix."

                    # Generate transition matrix and assert correct
                    prob_matrix = linalg.expm( np.multiply(inst_matrix, float(current_node.branch_length) ) )
                    assert( np.allclose( np.sum(prob_matrix, axis = 1), np.ones(len(self._code))) ), "Rows in transition matrix do not each sum to 1."

                    # Evolve branch
                    part_parent_seq = parent_node.seq[p][index : index + part.size[i]]
                    for j in range( part.size[i] ):
                        new_site = deepcopy( part_parent_seq[j] )
                        new_site.int_seq = self._generate_prob_from_unif( prob_matrix[ new_site.int_seq ] )
                        temp_new_seq.append( new_site )
                        index += 1
                new_seq.append( temp_new_seq )

        return new_seq

<|endoftext|>


<commit_before>from contextlib import contextmanager
from fabric.api import *
from fabric.contrib.files import upload_template, append
import os.path

PROJ_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
CONF_DIR = os.path.abspath(os.path.join(PROJ_DIR, 'conf'))

env.user = "oc"
env.gunicorn_port = 9000
env.code_dir = '~/OpenCommunity/'
env.venv_command = '. bin/activate'
env.log_dir = '/var/log/opencommunity/'
env.clone_url = "https://github.com/hasadna/OpenCommunity.git"
env.pip_version = "1.5.4"

@contextmanager
def virtualenv(path):
    with cd(path):
        with prefix(env.venv_command):
            yield

@task
def qa_old():
    env.user = "udi"
    env.hosts = ['oc-dev.modelarity.com']
    env.log_dir = '%slogs/' % env.code_dir
    env.pidfile = '/home/udi/OpenCommunity/run/masterpid'

def qa():
    env.vhost = '%s.qa.opencommunity.org.il' % env.user
    env.redirect_host = 'www.%s' % env.vhost
    env.hosts = ['qa.opencommunity.org.il']
    env.ocuser = "oc_" + env.user
    env.code_dir = '/home/%s/OpenCommunity/' % env.user
    env.log_dir = '%slogs/' % env.code_dir
    env.clone_url = "https://github.com/%s/OpenCommunity.git" % env.github_user
    env.venv_command = '. venv/bin/activate'
    env.venv_dir = '%svenv/' % env.code_dir
    env.pidfile = '/home/%s/opencommunity.pid' % env.ocuser

@task
def udi():
    env.gunicorn_port = 9010
    env.user = 'udi'
    env.github_user = 'nonZero'
    qa()

@task
def amir():
    env.gunicorn_port = 9011
    env.user = 'amir'
    env.github_user = 'amir99'
    qa()

@task
def yaniv():
    env.gunicorn_port = 9012
    env.user = 'yaniv'
    env.github_user = 'yaniv14'
    qa()

@task
def paul():
    env.gunicorn_port = 9013
    env.user = 'paul'
    env.github_user = 'pwalsh'
    qa()

@task
def prod():
    env.hosts = ['oc@ny1.opencommunity.org.il']
    env.redirect_host = 'opencommunity.org.il'
    env.venv_command = '. ~/.virtualenvs/oc/bin/activate'
    env.pidfile = '/home/oc/OpenCommunity/src/masterpid'

@task
def enprod():
    env.user = "en"

    env.code_dir = '/home/%s/OpenCommunity/' % env.user
    env.venv_dir = '%svenv/' % env.code_dir
    env.venv_command = '. venv/bin/activate'
    env.log_dir = '%slogs/' % env.code_dir

    env.vhost = 'en.demos.org.il'
    env.redirect_host = 'www.%s' % env.vhost
    env.hosts = [env.vhost]

    env.ocuser = "weben"
    env.pidfile = '/home/%s/web.pid' % env.ocuser

    env.gunicorn_port = 9001

    env.github_user = 'yaniv14'
    env.clone_url = "https://github.com/%s/OpenCommunity.git" % env.github_user

@task
def host_type():
    run('uname -s')

@task
def git_log():
    with virtualenv(env.code_dir):
        run("git log -n 1")

@task
def freeze():
    with virtualenv(env.code_dir):
        run("pip freeze")

@task
def reload_app():
    with virtualenv(env.code_dir):
        run("cd src && sudo kill -HUP `cat %s`" % env.pidfile)

@task
def upgrade_pip():
    with virtualenv(env.code_dir):
        run("pip install pip=={}".format(env.pip_version))

@task
def deploy(restart=True):
    upgrade_pip()
    with virtualenv(env.code_dir):
        run("git pull")
        run("pip install -r requirements.txt")
        run("pip install -r deploy-requirements.txt")
        run("cd src && python manage.py syncdb --noinput")
        run("cd src && python manage.py migrate --merge --noinput")
        run("cd src && python manage.py collectstatic --noinput")
        run("git log -n 1 --format=\"%ai %h\" > static/version.txt")
        run("git log -n 1 > static/version-full.txt")
        if restart:
            reload_app()

@task
def hard_reload():
    run("sudo supervisorctl restart opencommunity")

@task
def very_hard_reload():
    run("sudo service supervisor stop")
    run("sudo service supervisor start")

@task
def log():
    run("tail %s*" % env.log_dir)

APT_PACKAGES = [
    'postgresql',
    'nginx',
    'supervisor',
    'python',
    'virtualenvwrapper',
    'git',
    'python-dev',
    'libpq-dev',
    'libjpeg-dev',
    'libjpeg8',
    'zlib1g-dev',
    'libfreetype6',
    'libfreetype6-dev',
    'postfix',
]

@task
def server_setup():
    run("sudo apt-get update")
    run("sudo apt-get upgrade -y")
    run("sudo apt-get install -y %s" % " ".join(APT_PACKAGES))

@task
def create_ocuser_and_db():
    run("sudo adduser %s --gecos '' --disabled-password" % env.ocuser)
    run("sudo -iu postgres createuser %s -S -D -R" % env.ocuser)
    run("sudo -iu postgres createdb %s -O %s" % (env.ocuser, env.ocuser))
    run("sudo -iu postgres psql -c \"alter user %s with password '%s';\"" % (
    env.ocuser, env.ocuser))

@task
def clone_project():
    run("git clone %s %s" % (env.clone_url, env.code_dir))
    with cd(env.code_dir):
        run("virtualenv venv --prompt='(%s) '" % env.ocuser)

@task
def create_local_settings():
    with cd(env.code_dir):
        upload_template('ocd/local_settings.template',
                        env.code_dir + 'src/ocd/local_settings.py',
                        {'ocuser': env.ocuser, 'host': env.vhost},
                        use_jinja=True)

@task
def nginx_setup():
    with cd(env.code_dir):
        upload_template('nginx.conf.template',
                        env.code_dir + 'conf/nginx.conf',
                        {
                            'host': env.vhost,
                            'redirect_host': env.redirect_host,
                            'dir': env.code_dir,
                            'port': env.gunicorn_port,
                        }, use_jinja=True, template_dir=CONF_DIR)
        nginx_conf1 = '/etc/nginx/sites-available/%s.conf' % env.ocuser
        nginx_conf2 = '/etc/nginx/sites-enabled/%s.conf' % env.ocuser

        #       FIXME
        #       uncomment('/etc/nginx/nginx.conf',
        #                 'server_names_hash_bucket_size\s+64',
        #                 use_sudo=True)

        run('sudo ln -fs %sconf/nginx.conf %s' % (env.code_dir, nginx_conf1))
        run('sudo ln -fs %s %s' % (nginx_conf1, nginx_conf2))
        run('sudo service nginx configtest')
        run('sudo service nginx start')
        run('sudo service nginx reload')

@task
def gunicorn_setup():
    with cd(env.code_dir):
        upload_template('server.sh.template',
                        env.code_dir + 'server.sh',
                        {
                            'venv': env.venv_dir,
                            'port': env.gunicorn_port,
                            'pidfile': env.pidfile,
                        }, mode=0777, use_jinja=True, template_dir=PROJ_DIR)

@task
def supervisor_setup():
    with cd(env.code_dir):
        upload_template('supervisor.conf.template',
                        env.code_dir + 'conf/supervisor.conf',
                        {
                            'dir': env.code_dir,
                            'ocuser': env.ocuser,
                            'logdir': env.log_dir,
                        }, mode=0777, use_jinja=True, template_dir=CONF_DIR)

        run(
            'sudo ln -fs %sconf/supervisor.conf /etc/supervisor/conf.d/%s.conf'
            % (env.code_dir, env.ocuser))
        run("sudo supervisorctl reread")
        run("sudo supervisorctl update")
        run("sudo supervisorctl start %s" % env.ocuser)

@task
def project_setup():
    with cd(env.code_dir):
        run('mkdir -p uploads')
        run('sudo chown %s uploads' % env.ocuser)
        run('mkdir -p %s' % env.log_dir)
    create_local_settings()
    deploy(restart=False)
    gunicorn_setup()
    supervisor_setup()
    nginx_setup()

@task
def initial_project_setup():
    create_ocuser_and_db()
    clone_project()
    project_setup()

@task
def createsuperuser():
    """ Creates a Django superuser for the project """
    with virtualenv(env.code_dir):
        run("cd src && python manage.py createsuperuser")

@task
def supervisor_status():
    """ Show server's supoervisord status """
    run("sudo supervisorctl status")

@task
def switch(branch):
    """ fetches all branchs, and checkouts the specified git branch """
    with cd(env.code_dir):
        run('git fetch origin')
        run('git checkout {}'.format(branch))
        deploy()

@task
def showkeys():
    """ Displays authorized public ssh keys for user """
    with hide('stdout'):
        keys = run('cat .ssh/authorized_keys')
    print keys

@task
def push_key(key_file):
    """ Appends an ssh public key file from the specified file
    """
    with open(key_file) as f:
        key_text = f.read()
    append('~/.ssh/authorized_keys', key_text)

@task
def rebuild_index():
    """Rebuilds haystack search indexes for project"""
    with virtualenv(env.code_dir):
        run("cd src && python manage.py rebuild_index --noinput")
        run("sudo chown -v {} whoosh_index whoosh_index/*".format(env.ocuser))
<commit_msg>English production deploys from the HaSadna fork and the production branch.<commit_after>from contextlib import contextmanager
from fabric.api import *
from fabric.contrib.files import upload_template, append
import os.path

PROJ_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
CONF_DIR = os.path.abspath(os.path.join(PROJ_DIR, 'conf'))

env.user = "oc"
env.gunicorn_port = 9000
env.code_dir = '~/OpenCommunity/'
env.venv_command = '. bin/activate'
env.log_dir = '/var/log/opencommunity/'
env.clone_url = "https://github.com/hasadna/OpenCommunity.git"
env.pip_version = "1.5.4"

@contextmanager
def virtualenv(path):
    with cd(path):
        with prefix(env.venv_command):
            yield

@task
def qa_old():
    env.user = "udi"
    env.hosts = ['oc-dev.modelarity.com']
    env.log_dir = '%slogs/' % env.code_dir
    env.pidfile = '/home/udi/OpenCommunity/run/masterpid'

def qa():
    env.vhost = '%s.qa.opencommunity.org.il' % env.user
    env.redirect_host = 'www.%s' % env.vhost
    env.hosts = ['qa.opencommunity.org.il']
    env.ocuser = "oc_" + env.user
    env.code_dir = '/home/%s/OpenCommunity/' % env.user
    env.log_dir = '%slogs/' % env.code_dir
    env.clone_url = "https://github.com/%s/OpenCommunity.git" % env.github_user
    env.venv_command = '. venv/bin/activate'
    env.venv_dir = '%svenv/' % env.code_dir
    env.pidfile = '/home/%s/opencommunity.pid' % env.ocuser

@task
def udi():
    env.gunicorn_port = 9010
    env.user = 'udi'
    env.github_user = 'nonZero'
    qa()

@task
def amir():
    env.gunicorn_port = 9011
    env.user = 'amir'
    env.github_user = 'amir99'
    qa()

@task
def yaniv():
    env.gunicorn_port = 9012
    env.user = 'yaniv'
    env.github_user = 'yaniv14'
    qa()

@task
def paul():
    env.gunicorn_port = 9013
    env.user = 'paul'
    env.github_user = 'pwalsh'
    qa()

@task
def prod():
    env.hosts = ['oc@ny1.opencommunity.org.il']
    env.redirect_host = 'opencommunity.org.il'
    env.venv_command = '. ~/.virtualenvs/oc/bin/activate'
    env.pidfile = '/home/oc/OpenCommunity/src/masterpid'

@task
def enprod():
    env.user = "en"

    env.code_dir = '/home/%s/OpenCommunity/' % env.user
    env.venv_dir = '%svenv/' % env.code_dir
    env.venv_command = '. venv/bin/activate'
    env.log_dir = '%slogs/' % env.code_dir

    env.vhost = 'en.demos.org.il'
    env.redirect_host = 'www.%s' % env.vhost
    env.hosts = [env.vhost]

    env.ocuser = "weben"
    env.pidfile = '/home/%s/web.pid' % env.ocuser

    env.gunicorn_port = 9001

@task
def host_type():
    run('uname -s')

@task
def git_log():
    with virtualenv(env.code_dir):
        run("git log -n 1")

@task
def freeze():
    with virtualenv(env.code_dir):
        run("pip freeze")

@task
def reload_app():
    with virtualenv(env.code_dir):
        run("cd src && sudo kill -HUP `cat %s`" % env.pidfile)

@task
def upgrade_pip():
    with virtualenv(env.code_dir):
        run("pip install pip=={}".format(env.pip_version))

@task
def deploy(restart=True):
    upgrade_pip()
    with virtualenv(env.code_dir):
        run("git pull")
        run("pip install -r requirements.txt")
        run("pip install -r deploy-requirements.txt")
        run("cd src && python manage.py syncdb --noinput")
        run("cd src && python manage.py migrate --merge --noinput")
        run("cd src && python manage.py collectstatic --noinput")
        run("git log -n 1 --format=\"%ai %h\" > static/version.txt")
        run("git log -n 1 > static/version-full.txt")
        if restart:
            reload_app()

@task
def hard_reload():
    run("sudo supervisorctl restart opencommunity")

@task
def very_hard_reload():
    run("sudo service supervisor stop")
    run("sudo service supervisor start")

@task
def log():
    run("tail %s*" % env.log_dir)

APT_PACKAGES = [
    'postgresql',
    'nginx',
    'supervisor',
    'python',
    'virtualenvwrapper',
    'git',
    'python-dev',
    'libpq-dev',
    'libjpeg-dev',
    'libjpeg8',
    'zlib1g-dev',
    'libfreetype6',
    'libfreetype6-dev',
    'postfix',
]

@task
def server_setup():
    run("sudo apt-get update")
    run("sudo apt-get upgrade -y")
    run("sudo apt-get install -y %s" % " ".join(APT_PACKAGES))

@task
def create_ocuser_and_db():
    run("sudo adduser %s --gecos '' --disabled-password" % env.ocuser)
    run("sudo -iu postgres createuser %s -S -D -R" % env.ocuser)
    run("sudo -iu postgres createdb %s -O %s" % (env.ocuser, env.ocuser))
    run("sudo -iu postgres psql -c \"alter user %s with password '%s';\"" % (
    env.ocuser, env.ocuser))

@task
def clone_project():
    run("git clone %s %s" % (env.clone_url, env.code_dir))
    with cd(env.code_dir):
        run("virtualenv venv --prompt='(%s) '" % env.ocuser)

@task
def create_local_settings():
    with cd(env.code_dir):
        upload_template('ocd/local_settings.template',
                        env.code_dir + 'src/ocd/local_settings.py',
                        {'ocuser': env.ocuser, 'host': env.vhost},
                        use_jinja=True)

@task
def nginx_setup():
    with cd(env.code_dir):
        upload_template('nginx.conf.template',
                        env.code_dir + 'conf/nginx.conf',
                        {
                            'host': env.vhost,
                            'redirect_host': env.redirect_host,
                            'dir': env.code_dir,
                            'port': env.gunicorn_port,
                        }, use_jinja=True, template_dir=CONF_DIR)
        nginx_conf1 = '/etc/nginx/sites-available/%s.conf' % env.ocuser
        nginx_conf2 = '/etc/nginx/sites-enabled/%s.conf' % env.ocuser

        #       FIXME
        #       uncomment('/etc/nginx/nginx.conf',
        #                 'server_names_hash_bucket_size\s+64',
        #                 use_sudo=True)

        run('sudo ln -fs %sconf/nginx.conf %s' % (env.code_dir, nginx_conf1))
        run('sudo ln -fs %s %s' % (nginx_conf1, nginx_conf2))
        run('sudo service nginx configtest')
        run('sudo service nginx start')
        run('sudo service nginx reload')

@task
def gunicorn_setup():
    with cd(env.code_dir):
        upload_template('server.sh.template',
                        env.code_dir + 'server.sh',
                        {
                            'venv': env.venv_dir,
                            'port': env.gunicorn_port,
                            'pidfile': env.pidfile,
                        }, mode=0777, use_jinja=True, template_dir=PROJ_DIR)

@task
def supervisor_setup():
    with cd(env.code_dir):
        upload_template('supervisor.conf.template',
                        env.code_dir + 'conf/supervisor.conf',
                        {
                            'dir': env.code_dir,
                            'ocuser': env.ocuser,
                            'logdir': env.log_dir,
                        }, mode=0777, use_jinja=True, template_dir=CONF_DIR)

        run(
            'sudo ln -fs %sconf/supervisor.conf /etc/supervisor/conf.d/%s.conf'
            % (env.code_dir, env.ocuser))
        run("sudo supervisorctl reread")
        run("sudo supervisorctl update")
        run("sudo supervisorctl start %s" % env.ocuser)

@task
def project_setup():
    with cd(env.code_dir):
        run('mkdir -p uploads')
        run('sudo chown %s uploads' % env.ocuser)
        run('mkdir -p %s' % env.log_dir)
    create_local_settings()
    deploy(restart=False)
    gunicorn_setup()
    supervisor_setup()
    nginx_setup()

@task
def initial_project_setup():
    create_ocuser_and_db()
    clone_project()
    project_setup()

@task
def createsuperuser():
    """ Creates a Django superuser for the project """
    with virtualenv(env.code_dir):
        run("cd src && python manage.py createsuperuser")

@task
def supervisor_status():
    """ Show server's supoervisord status """
    run("sudo supervisorctl status")

@task
def switch(branch):
    """ fetches all branchs, and checkouts the specified git branch """
    with cd(env.code_dir):
        run('git fetch origin')
        run('git checkout {}'.format(branch))
        deploy()

@task
def showkeys():
    """ Displays authorized public ssh keys for user """
    with hide('stdout'):
        keys = run('cat .ssh/authorized_keys')
    print keys

@task
def push_key(key_file):
    """ Appends an ssh public key file from the specified file
    """
    with open(key_file) as f:
        key_text = f.read()
    append('~/.ssh/authorized_keys', key_text)

@task
def rebuild_index():
    """Rebuilds haystack search indexes for project"""
    with virtualenv(env.code_dir):
        run("cd src && python manage.py rebuild_index --noinput")
        run("sudo chown -v {} whoosh_index whoosh_index/*".format(env.ocuser))
<|endoftext|>


<commit_before>"""
"""

from wheezy.core.descriptors import attribute
from wheezy.core.introspection import import_name

from config import config

from membership.service.bridge import MembershipService
from posts.service.bridge import PostsService
from public.service.bridge import QuoteService

class Factory(object):

    def __init__(self, session_name, **context):
        self.context = context
        self.session = sessions[session_name]()
        self.factory = RepositoryFactory(self.session)

    def __enter__(self):
        self.session.__enter__()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.session.__exit__(exc_value, exc_value, traceback)

    @attribute
    def membership(self):
        return MembershipService(self.factory, self.context['errors'])

    @attribute
    def posts(self):
        return PostsService(self.factory,
                            self.context['errors'],
                            self.context['principal'])

    @attribute
    def quote(self):
        return QuoteService(self.factory, self.context['errors'])

class RepositoryFactory(object):

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

    @attribute
    def membership(self):
        return MembershipPersistence(self.session)

    @attribute
    def posts(self):
        return PostsPersistence(self.session)

    @attribute
    def quote(self):
        return QuotePersistence(self.session)

def mock_sessions():
    from wheezy.core.db import NullSession
    return {
        'ro': NullSession, 'rw': NullSession
    }

# region: configuration details
mode = config.get('runtime', 'mode')
MembershipPersistence = import_name('membership.repository.%s.'
                                    'MembershipRepository' % mode)
PostsPersistence = import_name('posts.repository.%s.'
                               'PostsRepository' % mode)
QuotePersistence = import_name('public.repository.%s.'
                               'QuoteRepository' % mode)
if mode == 'mock':
    sessions = mock_sessions()
else:
    raise NotImplementedError(mode)
del mode, config
<commit_msg>Minor refactoring with factory context arguments.<commit_after>"""
"""

from wheezy.core.descriptors import attribute
from wheezy.core.introspection import import_name

from config import config

from membership.service.bridge import MembershipService
from posts.service.bridge import PostsService
from public.service.bridge import QuoteService

class Factory(object):

    def __init__(self, session_name, errors, principal):
        session = sessions[session_name]()
        self.factory = RepositoryFactory(session)
        self.session = session
        self.errors = errors
        self.principal = principal

    def __enter__(self):
        self.session.__enter__()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.session.__exit__(exc_value, exc_value, traceback)

    @attribute
    def membership(self):
        return MembershipService(
            self.factory, self.errors, self.principal)

    @attribute
    def posts(self):
        return PostsService(
            self.factory, self.errors, self.principal)

    @attribute
    def quote(self):
        return QuoteService(self.factory)

class RepositoryFactory(object):

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

    @attribute
    def membership(self):
        return MembershipPersistence(self.session)

    @attribute
    def posts(self):
        return PostsPersistence(self.session)

    @attribute
    def quote(self):
        return QuotePersistence(self.session)

def mock_sessions():
    from wheezy.core.db import NullSession
    return {
        'ro': NullSession, 'rw': NullSession
    }

# region: configuration details
mode = config.get('runtime', 'mode')
MembershipPersistence = import_name('membership.repository.%s.'
                                    'MembershipRepository' % mode)
PostsPersistence = import_name('posts.repository.%s.'
                               'PostsRepository' % mode)
QuotePersistence = import_name('public.repository.%s.'
                               'QuoteRepository' % mode)
if mode == 'mock':
    sessions = mock_sessions()
else:
    raise NotImplementedError(mode)
del mode, config
<|endoftext|>


<commit_before>#----------------------------------------------------------------------
# Copyright (c) 2010 Raytheon BBN Technologies
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and/or hardware specification (the "Work") to
# deal in the Work without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Work, and to permit persons to whom the Work
# is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Work.
#
# THE WORK IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE WORK OR THE USE OR OTHER DEALINGS
# IN THE WORK.
#----------------------------------------------------------------------
"""
Reference GENI GCF Clearinghouse. Uses SFA Certificate and credential objects.
Run from gch.py
Will produce signed user credentials from a GID, return a
list of aggregates read from a config file, and create a new Slice Credential.

"""

import datetime
import logging
import os
import traceback
import uuid

from SecureXMLRPCServer import SecureXMLRPCServer
from credential import CredentialVerifier, create_credential, publicid_to_urn, is_valid_urn, string_to_urn_format, create_gid
import sfa.trust.gid as gid

# Substitute eg "openflow//stanford"
# Be sure this matches init-ca.py:CERT_AUTHORITY
# This is in publicid format
SLICE_AUTHORITY = "geni.net//gpo//gcf"

# Credential lifetimes in seconds
# Extend slice lifetimes to actually use the resources
USER_CRED_LIFE = 86400
SLICE_CRED_LIFE = 3600

# The list of Aggregates that this Clearinghouse knows about
# should be defined in the geni_aggregates file
# ListResources will refer the client to these aggregates
# Clearinghouse.runserver currently does the register_aggregate_pair
# calls for each row in that file
# but this should be doable dynamically
# Some sample pairs:
# GPOMYPLC = ('urn:publicid:IDN+plc:gpo1+authority+sa',
#             'http://myplc1.gpolab.bbn.com:12348')
# TESTGCFAM = ('urn:publicid:IDN+geni.net:gpo+authority+gcf',
#              'https://127.0.0.1:8001')
# OTHERGPOMYPLC = ('urn:publicid:IDN+plc:gpo+authority+site2',
#                    'http://128.89.81.74:12348')
# ELABINELABAM = ('urn:publicid:IDN+elabinelab.geni.emulab.net',
#                 'https://myboss.elabinelab.geni.emulab.net:443/protogeni/xmlrpc/am')

class SampleClearinghouseServer(object):
    """A sample clearinghouse with barebones functionality."""

    def __init__(self, delegate):
        self._delegate = delegate

    def GetVersion(self):
        return self._delegate.GetVersion()

    def CreateSlice(self, urn=None):
        return self._delegate.CreateSlice(urn_req=urn)

    def DeleteSlice(self, urn):
        return self._delegate.DeleteSlice(urn)

    def ListAggregates(self):
        return self._delegate.ListAggregates()

    def CreateUserCredential(self, cert):
        return self._delegate.CreateUserCredential(cert)

class Clearinghouse(object):

    def __init__(self):
        self.logger = logging.getLogger('gch')
        self.slices = {}
        self.aggs = []

    def load_aggregates(self, aggfile):
        """Loads aggregates from a file.

        The file has one aggregate per line. Each line contains a URN
        and a URL separated by a comma.

        Returns True if aggregates were loaded, False otherwise.
        """
        if not os.path.isfile(aggfile):
            self.logger.warn('Aggregate file %r does not exist.', aggfile)
            return

        line_num = 0
        for line in file(aggfile):
            line_num += 1
            spl = line.split(',')
            if len(spl) != 2:
                msg = ('File %s, line %d is malformed.'
                       + ' Expected "URN, URL", found %r')
                self.logger.warn(msg, aggfile, line_num, line)
                continue
            (urn, url) = spl
            urn = urn.strip()
            url = url.strip()
            if not urn:
                self.logger.warn('Empty URN on line %d of %s',
                                 line_num, aggfile)
                continue
            if not url:
                self.logger.warn('Empty URL on line %d of %s',
                                 line_num, aggfile)
                continue
            if urn in [x for (x, _) in self.aggs]:
                self.logger.warn('Duplicate URN %s at line %d of %s',
                                 urn, line_num, aggfile)
                continue
            self.logger.info("Registering AM %s at %s", urn, url)
            self.aggs.append((urn, url))

    def runserver(self, addr, keyfile=None, certfile=None,
                  ca_certs=None, aggfile=None):
        """Run the clearinghouse server."""
        # ca_certs is a file of 1 ch cert possibly (itself), or a dir of several for peering
        # If not supplied just use the certfile as the only trusted root
        self.keyfile = keyfile
        self.certfile = certfile

        # Error check the keyfile, certfile all exist
        if keyfile is None or not os.path.isfile(os.path.expanduser(keyfile)):
            raise Exception("Missing CH key file %s" % keyfile)
        if certfile is None or not os.path.isfile(os.path.expanduser(certfile)):
            raise Exception("Missing CH cert file %s" % certfile)

        if ca_certs is None:
            ca_certs = certfile
            self.logger.info("Using only my CH cert as a trusted root cert")

        self.trusted_root_files = CredentialVerifier(ca_certs).root_cert_files

        if not os.path.exists(os.path.expanduser(ca_certs)):
            raise Exception("Missing CA cert(s): %s" % ca_certs)

        # Load up the aggregates
        self.load_aggregates(aggfile)
        # FIXME: if there are no aggregates, should we continue?
        self.logger.info("%d Aggregate Managers registered from aggregates file %r", len(self.aggs), aggfile)

        # This is the arg to _make_server
        ca_certs_onefname = CredentialVerifier.getCAsFileFromDir(ca_certs)

        # This is used below by CreateSlice
        self.ca_cert_fnames = []
        if os.path.isfile(os.path.expanduser(ca_certs)):
            self.ca_cert_fnames = [os.path.expanduser(ca_certs)]
        elif os.path.isdir(os.path.expanduser(ca_certs)):
            self.ca_cert_fnames = [os.path.join(os.path.expanduser(ca_certs), name) for name in os.listdir(os.path.expanduser(ca_certs)) if name != CredentialVerifier.CATEDCERTSFNAME]

        # Create the xmlrpc server, load the rootkeys and do the ssl thing.
        self._server = self._make_server(addr, keyfile, certfile,
                                         ca_certs_onefname)
        self._server.register_instance(SampleClearinghouseServer(self))
        self.logger.info('GENI CH Listening on port %d...' % (addr[1]))
        self._server.serve_forever()

    def _make_server(self, addr, keyfile=None, certfile=None,
                     ca_certs=None):
        """Creates the XML RPC server."""
        # ca_certs is a file of concatenated certs
        return SecureXMLRPCServer(addr, keyfile=keyfile, certfile=certfile,
                                  ca_certs=ca_certs)

    def GetVersion(self):
        self.logger.info("Called GetVersion")
        version = dict()
        version['gch_api'] = 1
        return version

    # FIXME: Change that URN to be a name and non-optional
    # Currently client.py doesnt supply it, and
    # Omni takes a name and constructs a URN to supply
    def CreateSlice(self, urn_req = None):
        self.logger.info("Called CreateSlice URN REQ %r" % urn_req)
        slice_gid = None

        if urn_req and self.slices.has_key(urn_req):
            # If the Slice has expired, treat this as
            # a request to renew
            slice_cred = self.slices[urn_req]
            if slice_cred.expiration <= datetime.datetime.utcnow():
                # Need to renew this slice
                self.logger.info("CreateSlice on %r found existing cred that expired at %r - will renew", urn_req, slice_cred.expiration)
                slice_gid = slice_cred.get_gid_object()
            else:
                self.logger.debug("Slice cred is still valid at %r until %r - return it", datetime.datetime.utcnow(), slice_cred.expiration)
                return slice_cred.save_to_string()

        # First ensure we have a slice_urn
        if urn_req:
            # FIXME: Validate urn_req has the right form
            # to be issued by this CH
            if not is_valid_urn(urn_req):
                # FIXME: make sure it isnt empty, etc...
                urn = publicid_to_urn(urn_req)
            else:
                urn = urn_req
        else:
            # Generate a unique URN for the slice
            # based on this CH location and a UUID

            # Where was the slice created?
            (ipaddr, port) = self._server.socket._sock.getsockname()
            # FIXME: Get public_id start from a properties file
            # Create a unique name for the slice based on uuid
            slice_name = uuid.uuid4().__str__()[4:12]
            public_id = 'IDN %s slice %s//%s:%d' % (SLICE_AUTHORITY, slice_name,
                                                                   ipaddr,
                                                                   port)
            # this func adds the urn:publicid:
            # and converts spaces to +'s, and // to :
            urn = publicid_to_urn(public_id)

        # Now create a GID for the slice (signed credential)
        if slice_gid is None:
            try:
                slice_gid = create_gid(string_to_urn_format(SLICE_AUTHORITY + " slice"), urn, self.keyfile, self.certfile)[0]
            except Exception, exc:
                self.logger.error("Cant create slice gid for slice urn %s: %s", urn, traceback.format_exc())
                raise Exception("Failed to create slice %s. Cant create slice gid" % urn, exc)

        # Now get the user GID which will have permissions on this slice.
        # Get client x509 cert from the SSL connection
        # It doesnt have the chain but should be signed
        # by this CHs cert, which should also be a trusted
        # root at any federated AM. So everyone can verify it as is.
        # Note that if a user from a different CH (installed
        # as trusted by this CH for some reason) called this method,
        # that user would be used here - and can still get a valid slice
        try:
            user_gid = gid.GID(string=self._server.pem_cert)
        except Exception, exc:
            self.logger.error("CreateSlice failed to create user_gid from SSL client cert: %s", traceback.format_exc())
            raise Exception("Failed to create slice %s. Cant get user GID from SSL client certificate." % urn, exc)

        # OK have a user_gid so can get a slice credential
        # authorizing this user on the slice
        try:
            slice_cred = self.create_slice_credential(user_gid,
                                                      slice_gid)
        except Exception, exc:
            self.logger.error('CreateSlice failed to get slice credential for user %r, slice %r: %s', user_gid.get_hrn(), slice_gid.get_hrn(), traceback.format_exc())
            raise Exception('CreateSlice failed to get slice credential for user %r, slice %r' % (user_gid.get_hrn(), slice_gid.get_hrn()), exc)
        self.logger.info('Created slice %r' % (urn))

        self.slices[urn] = slice_cred

        return slice_cred.save_to_string()

    def DeleteSlice(self, urn_req):
        self.logger.info("Called DeleteSlice %r" % urn_req)
        if self.slices.has_key(urn_req):
            self.slices.pop(urn_req)
            self.logger.info("Deleted slice")
            return True
        self.logger.info('Slice was not found')
        # Slice not found!
        # FIXME: Raise an error so client knows why this failed?
        return False

    def ListAggregates(self):
        self.logger.info("Called ListAggregates")
        # TODO: Allow dynamic registration of aggregates
        return self.aggs

    def CreateUserCredential(self, user_gid):
        '''Return string representation of a user credential
        issued by this CH with caller/object this user_gid (string)
        with user privileges'''
        # FIXME: Validate arg - non empty, my user
        user_gid = gid.GID(string=user_gid)
        self.logger.info("Called CreateUserCredential for GID %s" % user_gid.get_hrn())
        try:
            ucred = create_credential(user_gid, user_gid, USER_CRED_LIFE, 'user', self.keyfile, self.certfile, self.trusted_root_files)
        except Exception, exc:
            self.logger.error("Failed to create user credential for %s: %s", user_gid.get_hrn(), traceback.format_exc())
            raise Exception("Failed to create user credential for %s" % user_gid.get_hrn(), exc)
        return ucred.save_to_string()

    def create_slice_credential(self, user_gid, slice_gid):
        '''Create a Slice credential object for this user_gid (object) on given slice gid (object)'''
        # FIXME: Validate the user_gid and slice_gid
        # are my user and slice
        return create_credential(user_gid, slice_gid, SLICE_CRED_LIFE, 'slice', self.keyfile, self.certfile, self.trusted_root_files )

<commit_msg>gam working<commit_after>#----------------------------------------------------------------------
# Copyright (c) 2010 Raytheon BBN Technologies
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and/or hardware specification (the "Work") to
# deal in the Work without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Work, and to permit persons to whom the Work
# is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Work.
#
# THE WORK IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE WORK OR THE USE OR OTHER DEALINGS
# IN THE WORK.
#----------------------------------------------------------------------
"""
Reference GENI GCF Clearinghouse. Uses SFA Certificate and credential objects.
Run from gch.py
Will produce signed user credentials from a GID, return a
list of aggregates read from a config file, and create a new Slice Credential.

"""

import datetime
import logging
import os
import traceback
import uuid

from SecureXMLRPCServer import SecureXMLRPCServer
from credential import CredentialVerifier, create_credential, publicid_to_urn, is_valid_urn, string_to_urn_format, create_gid
import sfa.trust.gid as gid

# Substitute eg "openflow//stanford"
# Be sure this matches init-ca.py:CERT_AUTHORITY
# This is in publicid format
SLICE_AUTHORITY = "geni//gpo//gcf"

# Credential lifetimes in seconds
# Extend slice lifetimes to actually use the resources
USER_CRED_LIFE = 86400
SLICE_CRED_LIFE = 3600

# The list of Aggregates that this Clearinghouse knows about
# should be defined in the geni_aggregates file
# ListResources will refer the client to these aggregates
# Clearinghouse.runserver currently does the register_aggregate_pair
# calls for each row in that file
# but this should be doable dynamically
# Some sample pairs:
# GPOMYPLC = ('urn:publicid:IDN+plc:gpo1+authority+sa',
#             'http://myplc1.gpolab.bbn.com:12348')
# TESTGCFAM = ('urn:publicid:IDN+geni.net:gpo+authority+gcf',
#              'https://127.0.0.1:8001')
# OTHERGPOMYPLC = ('urn:publicid:IDN+plc:gpo+authority+site2',
#                    'http://128.89.81.74:12348')
# ELABINELABAM = ('urn:publicid:IDN+elabinelab.geni.emulab.net',
#                 'https://myboss.elabinelab.geni.emulab.net:443/protogeni/xmlrpc/am')

class SampleClearinghouseServer(object):
    """A sample clearinghouse with barebones functionality."""

    def __init__(self, delegate):
        self._delegate = delegate

    def GetVersion(self):
        return self._delegate.GetVersion()

    def CreateSlice(self, urn=None):
        return self._delegate.CreateSlice(urn_req=urn)

    def DeleteSlice(self, urn):
        return self._delegate.DeleteSlice(urn)

    def ListAggregates(self):
        return self._delegate.ListAggregates()

    def CreateUserCredential(self, cert):
        return self._delegate.CreateUserCredential(cert)

class Clearinghouse(object):

    def __init__(self):
        self.logger = logging.getLogger('gch')
        self.slices = {}
        self.aggs = []

    def load_aggregates(self, aggfile):
        """Loads aggregates from a file.

        The file has one aggregate per line. Each line contains a URN
        and a URL separated by a comma.

        Returns True if aggregates were loaded, False otherwise.
        """
        if not os.path.isfile(aggfile):
            self.logger.warn('Aggregate file %r does not exist.', aggfile)
            return

        line_num = 0
        for line in file(aggfile):
            line_num += 1
            spl = line.split(',')
            if len(spl) != 2:
                msg = ('File %s, line %d is malformed.'
                       + ' Expected "URN, URL", found %r')
                self.logger.warn(msg, aggfile, line_num, line)
                continue
            (urn, url) = spl
            urn = urn.strip()
            url = url.strip()
            if not urn:
                self.logger.warn('Empty URN on line %d of %s',
                                 line_num, aggfile)
                continue
            if not url:
                self.logger.warn('Empty URL on line %d of %s',
                                 line_num, aggfile)
                continue
            if urn in [x for (x, _) in self.aggs]:
                self.logger.warn('Duplicate URN %s at line %d of %s',
                                 urn, line_num, aggfile)
                continue
            self.logger.info("Registering AM %s at %s", urn, url)
            self.aggs.append((urn, url))

    def runserver(self, addr, keyfile=None, certfile=None,
                  ca_certs=None, aggfile=None):
        """Run the clearinghouse server."""
        # ca_certs is a file of 1 ch cert possibly (itself), or a dir of several for peering
        # If not supplied just use the certfile as the only trusted root
        self.keyfile = keyfile
        self.certfile = certfile

        # Error check the keyfile, certfile all exist
        if keyfile is None or not os.path.isfile(os.path.expanduser(keyfile)):
            raise Exception("Missing CH key file %s" % keyfile)
        if certfile is None or not os.path.isfile(os.path.expanduser(certfile)):
            raise Exception("Missing CH cert file %s" % certfile)

        if ca_certs is None:
            ca_certs = certfile
            self.logger.info("Using only my CH cert as a trusted root cert")

        self.trusted_root_files = CredentialVerifier(ca_certs).root_cert_files

        if not os.path.exists(os.path.expanduser(ca_certs)):
            raise Exception("Missing CA cert(s): %s" % ca_certs)

        # Load up the aggregates
        self.load_aggregates(aggfile)
        # FIXME: if there are no aggregates, should we continue?
        self.logger.info("%d Aggregate Managers registered from aggregates file %r", len(self.aggs), aggfile)

        # This is the arg to _make_server
        ca_certs_onefname = CredentialVerifier.getCAsFileFromDir(ca_certs)

        # This is used below by CreateSlice
        self.ca_cert_fnames = []
        if os.path.isfile(os.path.expanduser(ca_certs)):
            self.ca_cert_fnames = [os.path.expanduser(ca_certs)]
        elif os.path.isdir(os.path.expanduser(ca_certs)):
            self.ca_cert_fnames = [os.path.join(os.path.expanduser(ca_certs), name) for name in os.listdir(os.path.expanduser(ca_certs)) if name != CredentialVerifier.CATEDCERTSFNAME]

        # Create the xmlrpc server, load the rootkeys and do the ssl thing.
        self._server = self._make_server(addr, keyfile, certfile,
                                         ca_certs_onefname)
        self._server.register_instance(SampleClearinghouseServer(self))
        self.logger.info('GENI CH Listening on port %d...' % (addr[1]))
        self._server.serve_forever()

    def _make_server(self, addr, keyfile=None, certfile=None,
                     ca_certs=None):
        """Creates the XML RPC server."""
        # ca_certs is a file of concatenated certs
        return SecureXMLRPCServer(addr, keyfile=keyfile, certfile=certfile,
                                  ca_certs=ca_certs)

    def GetVersion(self):
        self.logger.info("Called GetVersion")
        version = dict()
        version['gch_api'] = 1
        return version

    # FIXME: Change that URN to be a name and non-optional
    # Currently client.py doesnt supply it, and
    # Omni takes a name and constructs a URN to supply
    def CreateSlice(self, urn_req = None):
        self.logger.info("Called CreateSlice URN REQ %r" % urn_req)
        slice_gid = None

        if urn_req and self.slices.has_key(urn_req):
            # If the Slice has expired, treat this as
            # a request to renew
            slice_cred = self.slices[urn_req]
            if slice_cred.expiration <= datetime.datetime.utcnow():
                # Need to renew this slice
                self.logger.info("CreateSlice on %r found existing cred that expired at %r - will renew", urn_req, slice_cred.expiration)
                slice_gid = slice_cred.get_gid_object()
            else:
                self.logger.debug("Slice cred is still valid at %r until %r - return it", datetime.datetime.utcnow(), slice_cred.expiration)
                return slice_cred.save_to_string()

        # First ensure we have a slice_urn
        if urn_req:
            # FIXME: Validate urn_req has the right form
            # to be issued by this CH
            if not is_valid_urn(urn_req):
                # FIXME: make sure it isnt empty, etc...
                urn = publicid_to_urn(urn_req)
            else:
                urn = urn_req
        else:
            # Generate a unique URN for the slice
            # based on this CH location and a UUID

            # Where was the slice created?
            (ipaddr, port) = self._server.socket._sock.getsockname()
            # FIXME: Get public_id start from a properties file
            # Create a unique name for the slice based on uuid
            slice_name = uuid.uuid4().__str__()[4:12]
            public_id = 'IDN %s slice %s//%s:%d' % (SLICE_AUTHORITY, slice_name,
                                                                   ipaddr,
                                                                   port)
            # this func adds the urn:publicid:
            # and converts spaces to +'s, and // to :
            urn = publicid_to_urn(public_id)

        # Now create a GID for the slice (signed credential)
        if slice_gid is None:
            try:
                slice_gid = create_gid(string_to_urn_format(SLICE_AUTHORITY + " slice"), urn, self.keyfile, self.certfile)[0]
            except Exception, exc:
                self.logger.error("Cant create slice gid for slice urn %s: %s", urn, traceback.format_exc())
                raise Exception("Failed to create slice %s. Cant create slice gid" % urn, exc)

        # Now get the user GID which will have permissions on this slice.
        # Get client x509 cert from the SSL connection
        # It doesnt have the chain but should be signed
        # by this CHs cert, which should also be a trusted
        # root at any federated AM. So everyone can verify it as is.
        # Note that if a user from a different CH (installed
        # as trusted by this CH for some reason) called this method,
        # that user would be used here - and can still get a valid slice
        try:
            user_gid = gid.GID(string=self._server.pem_cert)
        except Exception, exc:
            self.logger.error("CreateSlice failed to create user_gid from SSL client cert: %s", traceback.format_exc())
            raise Exception("Failed to create slice %s. Cant get user GID from SSL client certificate." % urn, exc)

        # OK have a user_gid so can get a slice credential
        # authorizing this user on the slice
        try:
            slice_cred = self.create_slice_credential(user_gid,
                                                      slice_gid)
        except Exception, exc:
            self.logger.error('CreateSlice failed to get slice credential for user %r, slice %r: %s', user_gid.get_hrn(), slice_gid.get_hrn(), traceback.format_exc())
            raise Exception('CreateSlice failed to get slice credential for user %r, slice %r' % (user_gid.get_hrn(), slice_gid.get_hrn()), exc)
        self.logger.info('Created slice %r' % (urn))

        self.slices[urn] = slice_cred

        return slice_cred.save_to_string()

    def DeleteSlice(self, urn_req):
        self.logger.info("Called DeleteSlice %r" % urn_req)
        if self.slices.has_key(urn_req):
            self.slices.pop(urn_req)
            self.logger.info("Deleted slice")
            return True
        self.logger.info('Slice was not found')
        # Slice not found!
        # FIXME: Raise an error so client knows why this failed?
        return False

    def ListAggregates(self):
        self.logger.info("Called ListAggregates")
        # TODO: Allow dynamic registration of aggregates
        return self.aggs

    def CreateUserCredential(self, user_gid):
        '''Return string representation of a user credential
        issued by this CH with caller/object this user_gid (string)
        with user privileges'''
        # FIXME: Validate arg - non empty, my user
        user_gid = gid.GID(string=user_gid)
        self.logger.info("Called CreateUserCredential for GID %s" % user_gid.get_hrn())
        try:
            ucred = create_credential(user_gid, user_gid, USER_CRED_LIFE, 'user', self.keyfile, self.certfile, self.trusted_root_files)
        except Exception, exc:
            self.logger.error("Failed to create user credential for %s: %s", user_gid.get_hrn(), traceback.format_exc())
            raise Exception("Failed to create user credential for %s" % user_gid.get_hrn(), exc)
        return ucred.save_to_string()

    def create_slice_credential(self, user_gid, slice_gid):
        '''Create a Slice credential object for this user_gid (object) on given slice gid (object)'''
        # FIXME: Validate the user_gid and slice_gid
        # are my user and slice
        return create_credential(user_gid, slice_gid, SLICE_CRED_LIFE, 'slice', self.keyfile, self.certfile, self.trusted_root_files )

<|endoftext|>


<commit_before>from sys import stdin
from collections import defaultdict
from argparse import ArgumentParser

def parse_args():
    p = ArgumentParser('Reads paraphrase candidates from STDIN. Format: phrase1 <TAB> phrase2 [<TAB> gold]')
    p.add_argument('-l', '--lower', action='store_true', default=False,
                   help='lower all input')
    p.add_argument('--mode', choices=['word', 'char'], default='char',
                   help='word or character ngrams')
    p.add_argument('-N', type=int, default=3, help='N of the Ngram')
    p.add_argument('-t', '--threshold', type=float, default=0.5,
                   help='Jaccard threshold')
    p.add_argument('--silent', action='store_true', default=False,
                   help='do not print the Jaccard score for each pair, '
                   'print only the summary.')
    p.add_argument('-v', '--verbose', action='store_true', default=False,
                   help='print verbose summary (true/false positives and negatives)')
    p.add_argument('--strict', type=int, default=1,
                   help='good - bad = x, accept x>=strict answers. E.g.'
                   'if strict is 2, then 4-1=2 is accepted, 3-2=1 is not accepted.'
                   'This argument has no effect if there is no third column in the input')
    p.add_argument('--round-to', type=int, default=3,
                   help='round prec/rec/acc/F values to R')
    return p.parse_args()

cfg = parse_args()

def trim_text(text):
    if cfg.lower:
        return text.lower()
    return text

def get_ngrams(text):
    ngrams = set()
    trimmed = trim_text(text)
    N = cfg.N
    if cfg.mode == 'char':
        for i in xrange(0, len(trimmed) - N + 1):
            ngrams.add(trimmed[i:i + N])
    else:
        words = trimmed.split(' ')
        for i in xrange(0, len(words) - N + 1):
            ngrams.add(tuple(words[i:i + N]))
    return ngrams

def get_answer(field):
    good = int(field[1])
    bad = int(field[-2])
    return good - bad >= cfg.threshold

def jaccard(set1, set2):
    return float(len(set1 & set2)) / len(set1 | set2)

def store_stat(stat, jacc, gold):
    if jacc >= cfg.threshold:
        guess = True
    else:
        guess = False
    if gold and guess:
        stat['tp'] += 1
    elif gold and not guess:
        stat['fn'] += 1
    elif not gold and guess:
        stat['fp'] += 1
    else:
        stat['tn'] += 1
    stat['N'] += 1

def compute_and_print_stats(stat):
    if cfg.verbose:
        print('true pos: {0}, true neg: {1}, false pos: {2}, false neg: {3}'.format(
            stat['tp'], stat['tn'], stat['fp'], stat['fn']))
    acc = float(stat['tp'] + stat['tn']) / stat['N']
    try:
        prec = float(stat['tp']) / (stat['tp'] + stat['fp'])
    except ZeroDivisionError:
        prec = 0
    try:
        rec = float(stat['tp']) / (stat['tp'] + stat['fn'])
    except ZeroDivisionError:
        rec = 0
    try:
        F = 2 * prec * rec / (prec + rec)
    except ZeroDivisionError:
        F = 0
    r = cfg.round_to
    print('| Accuracy | Precision | Recall | F score |\n| --- | --- | --- | --- |')
    print('| {0} | {1} | {2} | {3} |'.format(round(acc, r), round(prec, r), round(rec, r), round(F, r)))

def main():
    stat = defaultdict(int)
    for l in stdin:
        fd = l.decode('utf8').strip().split('\t')
        feat1 = get_ngrams(fd[0])
        feat2 = get_ngrams(fd[1])
        jacc = jaccard(feat1, feat2)
        if not cfg.silent:
            print(u'{0}\t{1}\t{2}'.format(jacc, fd[0], fd[1]).encode('utf8'))
        if len(fd) > 2:
            answer = get_answer(fd[2])
            store_stat(stat, jacc, answer)
    if stat:
        compute_and_print_stats(stat)

if __name__ == '__main__':
    main()
<commit_msg>Generate threshold-F score (and other measures) table. 	modified:   src/jaccard.py<commit_after>from sys import stdin
from collections import defaultdict
from argparse import ArgumentParser

def parse_args():
    p = ArgumentParser('Reads paraphrase candidates from STDIN. Format: phrase1 <TAB> phrase2 [<TAB> gold]')
    p.add_argument('-l', '--lower', action='store_true', default=False,
                   help='lower all input')
    p.add_argument('--mode', choices=['word', 'char'], default='char',
                   help='word or character ngrams')
    p.add_argument('-N', type=int, default=3, help='N of the Ngram')
    p.add_argument('-t', '--threshold', type=float, default=0.5,
                   help='Jaccard threshold')
    p.add_argument('--silent', action='store_true', default=False,
                   help='do not print the Jaccard score for each pair, '
                   'print only the summary.')
    p.add_argument('--roc', type=str, default='',
                   help='produce a ROC table')
    p.add_argument('-v', '--verbose', action='store_true', default=False,
                   help='print verbose summary (true/false positives and negatives)')
    p.add_argument('--strict', type=int, default=1,
                   help='good - bad = x, accept x>=strict answers. E.g.'
                   'if strict is 2, then 4-1=2 is accepted, 3-2=1 is not accepted.'
                   'This argument has no effect if there is no third column in the input')
    p.add_argument('--round-to', type=int, default=3,
                   help='round prec/rec/acc/F values to R')
    return p.parse_args()

cfg = parse_args()
roc_thresholds = [float(i) / 100 for i in xrange(0, 101, 1)]

def trim_text(text):
    if cfg.lower:
        return text.lower()
    return text

def get_ngrams(text):
    ngrams = set()
    trimmed = trim_text(text)
    N = cfg.N
    if cfg.mode == 'char':
        for i in xrange(0, len(trimmed) - N + 1):
            ngrams.add(trimmed[i:i + N])
    else:
        words = trimmed.split(' ')
        for i in xrange(0, len(words) - N + 1):
            ngrams.add(tuple(words[i:i + N]))
    return ngrams

def get_answer(field):
    good = int(field[1])
    bad = int(field[-2])
    return good - bad >= cfg.strict

def jaccard(set1, set2):
    return float(len(set1 & set2)) / len(set1 | set2)

def store_stat(stat, jacc, gold):
    if cfg.roc:
        for th in roc_thresholds:
            store_with_threshold(jacc, gold, th, stat[th])
    else:
        store_with_threshold(jacc, gold, cfg.threshold, stat)

def store_with_threshold(jacc, gold, th, stat):
    if jacc >= th:
        guess = True
    else:
        guess = False
    if gold and guess:
        stat['tp'] += 1
    elif gold and not guess:
        stat['fn'] += 1
    elif not gold and guess:
        stat['fp'] += 1
    else:
        stat['tn'] += 1
    stat['N'] += 1

def get_measures(stat):
    acc = float(stat['tp'] + stat['tn']) / stat['N']
    try:
        prec = float(stat['tp']) / (stat['tp'] + stat['fp'])
    except ZeroDivisionError:
        prec = 0
    try:
        rec = float(stat['tp']) / (stat['tp'] + stat['fn'])
    except ZeroDivisionError:
        rec = 0
    try:
        F = 2 * prec * rec / (prec + rec)
    except ZeroDivisionError:
        F = 0
    return acc, prec, rec, F

def compute_and_print_stats(stat):
    if cfg.verbose:
        print('true pos: {0}, true neg: {1}, false pos: {2}, false neg: {3}'.format(
            stat['tp'], stat['tn'], stat['fp'], stat['fn']))
    acc, prec, rec, F = get_measures(stat)
    r = cfg.round_to
    print('| Accuracy | Precision | Recall | F score |\n| --- | --- | --- | --- |')
    print('| {0} | {1} | {2} | {3} |'.format(round(acc, r), round(prec, r), round(rec, r), round(F, r)))

def main():
    if cfg.roc:
        stat = defaultdict(lambda: defaultdict(int))
    else:
        stat = defaultdict(int)
    for l in stdin:
        fd = l.decode('utf8').strip().split('\t')
        feat1 = get_ngrams(fd[0])
        feat2 = get_ngrams(fd[1])
        jacc = jaccard(feat1, feat2)
        if not cfg.silent:
            print(u'{0}\t{1}\t{2}'.format(jacc, fd[0], fd[1]).encode('utf8'))
        if len(fd) > 2:
            answer = get_answer(fd[2])
            store_stat(stat, jacc, answer)
    if stat and not cfg.roc:
        compute_and_print_stats(stat)
    if cfg.roc:
        with open(cfg.roc, 'w') as f:
            for th, st in sorted(stat.iteritems(), key=lambda x: -x[0]):
                acc, prec, rec, F = get_measures(st)
                N = float(st['N'])
                tp_rate = st['tp'] / N
                fp_rate = st['fp'] / N
                f.write('{0}\t{1}\t{2}\t{3}\t{4}\t{5}\t{6}\n'.format(th, tp_rate, fp_rate,
                        acc, prec, rec, F))

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>from monoid import Monoid
from control import Control

class Machine:
    def __init__(self, base, control=None):

        # control will be an FST representation later
        if not isinstance(control, Control) and control is not None:
            raise TypeError("control should be a Control instance")
        self.control = control

        # base is a monoid
        if not isinstance(base, Monoid):
            raise TypeError("base should be a Monoid instance")
        self.base = base

    def __str__(self):
        """
        Returns machine's "printname"
        """
        return str(self.base)

    def __eq__(self, other):
        return self.control == other.control and self.base == other.base

    def __hash__(self):
        return hash(id(self))

    def append(self, which_partition, what):
        self.base.append(which_partition, what)

    def to_dot(self, toplevel=False):
        s = u"subgraph"
        if toplevel:
            s = u"graph"

        s += u" cluster_{0}_{1} {{\n".format(self.base.partitions[0], id(self))
        s += u"label={0};\n".format(self.base.partitions[0])

        if len(self.base.partitions) > 1:
            s += "color=black;\n"
            for p in reversed(self.base.partitions[1:]):
                s += u"subgraph cluster_{0}_{1} {{\n".format(self.base.partitions[0], id(p))
                s += "label=\"\"\n"
                s += "color=lightgrey;\n"
                for m in reversed(p):
                    if isinstance(m, Machine):
                        s += m.to_dot()
                s += "}\n"
        else:
            #s += "color=white;\n"
            s += u"{0}[color=white, fontcolor=white];\n".format(self.base.partitions[0])
        s += "}\n"

        return s
<commit_msg>unicode bug fix<commit_after>from monoid import Monoid
from control import Control

class Machine:
    def __init__(self, base, control=None):

        # control will be an FST representation later
        if not isinstance(control, Control) and control is not None:
            raise TypeError("control should be a Control instance")
        self.control = control

        # base is a monoid
        if not isinstance(base, Monoid):
            raise TypeError("base should be a Monoid instance")
        self.base = base

    def __str__(self):
        """
        Returns machine's "printname"
        """
        return str(self.base)

    def __unicode__(self):
        return unicode(self.base)

    def __eq__(self, other):
        return self.control == other.control and self.base == other.base

    def __hash__(self):
        return hash(id(self))

    def append(self, which_partition, what):
        self.base.append(which_partition, what)

    def to_dot(self, toplevel=False):
        s = u"subgraph"
        if toplevel:
            s = u"graph"

        s += u" cluster_{0}_{1} {{\n".format(self.base.partitions[0], id(self))
        s += u"label={0};\n".format(self.base.partitions[0])

        if len(self.base.partitions) > 1:
            s += "color=black;\n"
            for p in reversed(self.base.partitions[1:]):
                s += u"subgraph cluster_{0}_{1} {{\n".format(self.base.partitions[0], id(p))
                s += "label=\"\"\n"
                s += "color=lightgrey;\n"
                for m in reversed(p):
                    if isinstance(m, Machine):
                        s += m.to_dot()
                s += "}\n"
        else:
            #s += "color=white;\n"
            s += u"{0}[color=white, fontcolor=white];\n".format(self.base.partitions[0])
        s += "}\n"

        return s
<|endoftext|>