<commit_before># -*- coding: utf-8 -*-
# Copyright (C) 2012 Sun Ning<classicning@gmail.com>
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software 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 Software.

# THE SOFTWARE 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 SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# clojure literals => python types
#
# clojure vector [1 2 3 4] => python list [1 2 3 4] *coll
# clojure list (1 2 3 4) => python list [1 2 3 4] *coll
# clojure set #{1 2 3 4} => python set set(1 2 3 4) *coll
# clojure map {:a 1 :b 2} => python dict dict(a=1,b=2) *coll
# clojure string "a" => python unicode "a"
# clojure character \a => python unicode "a"
# clojure keyword :a => python unicode "a"
# clojure integer 123 => python integer 123
# clojure float 12.3 => python float 12.3
# clojure boolean true => python boolean true
# clojure nil => python None
#

__all__ = ["dump", "dumps", "load", "loads"]

import os
from cStringIO import StringIO

def number(v):
    if '.' in v:
        return float(v)
    else:
        return int(v)

_STOP_CHARS = [" ", ",", "\n", "\r", "\t"]
_COLL_OPEN_CHARS = ["#", "[", "{", "("]
_COLL_CLOSE_CHARS = ["]", "}", ")"]
_EXTRA_NUM_CHARS = ["-", "+", ".", "e", "E"]

class CljDecoder(object):
    def __init__(self, fd):
        self.fd = fd
        self.value_stack = []
        self.terminator = None ## for collection type

    def decode(self):
        while True:
            v = self.__read_token()
            if len(self.value_stack) == 0:
                return v

    def __get_type_from_char(self, c):
        """return a tuple of type information
        * type name
        * a flag to indicate if it's a collection
        """
        if c.isdigit() or c =='-':
            return ("number", False, None)
        elif c == 't' or c == 'f': ## true/false
            return ("boolean", False, None)
        elif c == 'n': ## nil
            return ("nil", False, None)
        elif c == '\\' :
            return ("char", False, None)
        elif c == ':':
            return ("keyword", False, None)
        elif c == '"':
            return ("string", False, None)
        elif c == '#':
            return ("set", True, "}")
        elif c == '{':
            return ("dict", True, "}")
        elif c == '(':
            return ("list", True, ")")
        elif c == '[':
            return ('list', True, "]")
        else:
            return (None, False, None)

    def __read_token(self):
        fd = self.fd

        c = fd.read(1)

        ## skip all stop chars if necessary
        while c in _STOP_CHARS:
            c = fd.read(1)

        ## raise exception when unexpected EOF found
        if c == '':
            raise ValueError("Unexpected EOF")

        t, coll, term = self.__get_type_from_char(c)
        if coll:
            ## move cursor
            if t == "set":
                ## skip {
                fd.read(1)

            self.terminator = term

            self.value_stack.append(([], self.terminator, t))
            return None
        else:
            v = None ## token value
            e = None ## end char
            r = True ## the token contains data or not

            if t == "boolean":
                if c == 't':
                    chars = fd.read(4)
                    if chars[:3] != 'rue':
                        raise ValueError('Expect true, got t'+chars[:3])
                    e = chars[-1]
                    v = True
                else:
                    chars = fd.read(5)
                    if chars[:4] != 'alse':
                        raise ValueError('Expect false, got f'+chars[:4])
                    e = chars[-1]
                    v = False

            elif t == "char":
                buf = []
                while c is not self.terminator and c is not "" and c not in _STOP_CHARS:
                    c = fd.read(1)
                    buf.append(c)

                e = c
                v = ''.join(buf[:-1])

            elif t == "nil":
                chars = fd.read(3)
                if chars[:2] != 'il':
                    raise ValueError('Expect nil, got n'+chars[:2])
                e = chars[-1]
                v = None

            elif t == "number":
                buf = []
                while c.isdigit() or (c in _EXTRA_NUM_CHARS):
                    buf.append(c)
                    c = fd.read(1)
                e = c
                numstr = ''.join(buf)
                v = number(numstr)

                ## special case for
                ## [23[12]]
                ## this is a valid clojure form
                if e in _COLL_OPEN_CHARS:
                    fd.seek(-1, os.SEEK_CUR)

            elif t == "keyword":
                buf = []    ##skip the leading ":"
                while c is not self.terminator and c is not "" and c not in _STOP_CHARS:
                    c = fd.read(1)
                    buf.append(c)

                e = c
                v = ''.join(buf[:-1])

            elif t == "string":
                buf = []
                cp = c = fd.read(1) ## to check escaping character \

                while not(c == '"' and cp != '\\'):
                    buf.append(c)
                    cp = c
                    c = fd.read(1)
                e = c
                #v = u''.join(buf).decode('unicode-escape')
                v = ''.join(buf).decode('string-escape')
            else:
                if c not in _COLL_CLOSE_CHARS:
                    raise ValueError('Unexpected char: '+c)
                r = False
                e = c

            if e is self.terminator:
                current_scope, _, container = self.value_stack.pop()

                if r:
                    current_scope.append(v)

                if container == "set":
                    v = set(current_scope)
                elif container == "list":
                    v = current_scope
                elif container == "dict":
                    v = {}
                    for i in range(0, len(current_scope), 2):
                        v[current_scope[i]] = current_scope[i+1]

            if len(self.value_stack) > 0:
                self.value_stack[-1][0].append(v)
                self.terminator = self.value_stack[-1][1]

            return v

class CljEncoder(object):
    def __init__(self, data, fd):
        self.data = data
        self.fd = fd
        self.circular = {}

    def encode(self):
        self.__do_encode(self.data)

    def get_type(self,t):
        if t is None:
            return ("None", False)
        elif isinstance(t, str):
            return ("string", False)
        elif isinstance(t, bool):
            return ("boolean", False)
        elif isinstance(t, (int,float,long)):
            return ("number", False)
        elif isinstance(t, dict):
            return ("dict", True)
        elif isinstance(t, list):
            return ("list", True)
        elif isinstance(t, set):
            return ("set", True)
        else:
            return ("unknown", False)

    def __do_encode(self, d):
        fd = self.fd
        t,coll = self.get_type(d)

        if coll:
            refid = id(d)
            if refid in self.circular:
                raise ValueError('Circular reference detected')
            else:
                self.circular[refid] = True

            if t == "dict":
                fd.write("{")
                for k,v in d.items():
                    self.__do_encode(k)
                    fd.write(" ")
                    self.__do_encode(v)
                    fd.write(" ")
                fd.seek(-1, os.SEEK_CUR)
                fd.write("}")
            elif t == "list":
                fd.write("[")
                for v in d:
                    self.__do_encode(v)
                    fd.write(" ")
                fd.seek(-1, os.SEEK_CUR)
                fd.write("]")
            elif t == "set":
                fd.write("#{")
                for v in d:
                    self.__do_encode(v)
                    fd.write(" ")
                fd.seek(-1, os.SEEK_CUR)
                fd.write("}")
        else:
            if t == "number":
                fd.write(str(d))
            elif t == "string":
                s = d.encode("string-escape").replace('"', '\\"')
                fd.write('"'+s+'"')
            elif t == "boolean":
                if d:
                    fd.write('true')
                else:
                    fd.write('false')
            elif t == 'None':
                fd.write('nil')
            else:
                fd.write('"'+str(d)+'"')

def dump(obj, fp):
    return CljEncoder(obj, fp).encode()

def dumps(obj):
    buf = StringIO()
    dump(obj, buf)
    result = buf.getvalue()
    buf.close()
    return result

def load(fp):
    decoder = CljDecoder(fp)
    return decoder.decode()

def loads(s):
    buf = StringIO(s)
    result = load(buf)
    buf.close()
    return result

<commit_msg>add line/col number for user friendly exception message<commit_after># -*- coding: utf-8 -*-
# Copyright (C) 2012 Sun Ning<classicning@gmail.com>
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software 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 Software.

# THE SOFTWARE 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 SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# clojure literals => python types
#
# clojure vector [1 2 3 4] => python list [1 2 3 4] *coll
# clojure list (1 2 3 4) => python list [1 2 3 4] *coll
# clojure set #{1 2 3 4} => python set set(1 2 3 4) *coll
# clojure map {:a 1 :b 2} => python dict dict(a=1,b=2) *coll
# clojure string "a" => python unicode "a"
# clojure character \a => python unicode "a"
# clojure keyword :a => python unicode "a"
# clojure integer 123 => python integer 123
# clojure float 12.3 => python float 12.3
# clojure boolean true => python boolean true
# clojure nil => python None
#

__all__ = ["dump", "dumps", "load", "loads"]

import os
from cStringIO import StringIO

def number(v):
    if '.' in v:
        return float(v)
    else:
        return int(v)

_STOP_CHARS = [" ", ",", "\n", "\r", "\t"]
_COLL_OPEN_CHARS = ["#", "[", "{", "("]
_COLL_CLOSE_CHARS = ["]", "}", ")"]
_EXTRA_NUM_CHARS = ["-", "+", ".", "e", "E"]

class CljDecoder(object):
    def __init__(self, fd):
        self.fd = fd
        self.cur_line = 1
        self.cur_pos = 1
        self.value_stack = []
        self.terminator = None ## for collection type

    def decode(self):
        while True:
            v = self.__read_token()
            if len(self.value_stack) == 0:
                return v

    def __get_type_from_char(self, c):
        """return a tuple of type information
        * type name
        * a flag to indicate if it's a collection
        """
        if c.isdigit() or c =='-':
            return ("number", False, None)
        elif c == 't' or c == 'f': ## true/false
            return ("boolean", False, None)
        elif c == 'n': ## nil
            return ("nil", False, None)
        elif c == '\\' :
            return ("char", False, None)
        elif c == ':':
            return ("keyword", False, None)
        elif c == '"':
            return ("string", False, None)
        elif c == '#':
            return ("set", True, "}")
        elif c == '{':
            return ("dict", True, "}")
        elif c == '(':
            return ("list", True, ")")
        elif c == '[':
            return ('list', True, "]")
        else:
            return (None, False, None)

    def __read_fd(self, size):
        self.cur_pos = self.cur_pos + size
        if size == 1:
            c = self.fd.read(size)
            if  c == '\n':
                self.cur_line = self.cur_line + 1
            return c
        else:
            return self.fd.read(size)

    def __read_token(self):
        c = self.__read_fd(1)

        ## skip all stop chars if necessary
        while c in _STOP_CHARS:
            c = self.__read_fd(1)

        ## raise exception when unexpected EOF found
        if c == '':
            raise ValueError("Unexpected EOF")

        t, coll, term = self.__get_type_from_char(c)
        if coll:
            ## move cursor
            if t == "set":
                ## skip {
                self.__read_fd(1)

            self.terminator = term

            self.value_stack.append(([], self.terminator, t))
            return None
        else:
            v = None ## token value
            e = None ## end char
            r = True ## the token contains data or not

            if t == "boolean":
                if c == 't':
                    chars = self.__read_fd(4)
                    if chars[:3] != 'rue':
                        raise ValueError('Expect true, got t%s at line %d, col %d' % (chars[:3], self.cur_line, self.cur_pos))
                    e = chars[-1]
                    v = True
                else:
                    chars = self.__read_fd(5)
                    if chars[:4] != 'alse':
                        raise ValueError('Expect true, got t%s at line %d, col %d' % (chars[:3], self.cur_line, self.cur_pos))
                    e = chars[-1]
                    v = False

            elif t == "char":
                buf = []
                while c is not self.terminator and c is not "" and c not in _STOP_CHARS:
                    c = self.__read_fd(1)
                    buf.append(c)

                e = c
                v = ''.join(buf[:-1])

            elif t == "nil":
                chars = self.__read_fd(3)
                if chars[:2] != 'il':
                    raise ValueError('Expect nil, got n%s at line %d, col %d' % (chars[:2], self.cur_line, self.cur_pos))
                e = chars[-1]
                v = None

            elif t == "number":
                buf = []
                while c.isdigit() or (c in _EXTRA_NUM_CHARS):
                    buf.append(c)
                    c = self.__read_fd(1)
                e = c
                numstr = ''.join(buf)
                v = number(numstr)

                ## special case for
                ## [23[12]]
                ## this is a valid clojure form
                if e in _COLL_OPEN_CHARS:
                    self.fd.seek(-1, os.SEEK_CUR)

            elif t == "keyword":
                buf = []    ##skip the leading ":"
                while c is not self.terminator and c is not "" and c not in _STOP_CHARS:
                    c = self.__read_fd(1)
                    buf.append(c)

                e = c
                v = ''.join(buf[:-1])

            elif t == "string":
                buf = []
                cp = c = self.__read_fd(1) ## to check escaping character \

                while not(c == '"' and cp != '\\'):
                    buf.append(c)
                    cp = c
                    c = self.__read_fd(1)
                e = c
                #v = u''.join(buf).decode('unicode-escape')
                v = ''.join(buf).decode('string-escape')
            else:
                if c not in _COLL_CLOSE_CHARS:
                    raise ValueError('Unexpected char: "%s" at line %d, col %d' % (c, self.cur_line, self.cur_pos))
                r = False
                e = c

            if e is self.terminator:
                current_scope, _, container = self.value_stack.pop()

                if r:
                    current_scope.append(v)

                if container == "set":
                    v = set(current_scope)
                elif container == "list":
                    v = current_scope
                elif container == "dict":
                    v = {}
                    for i in range(0, len(current_scope), 2):
                        v[current_scope[i]] = current_scope[i+1]

            if len(self.value_stack) > 0:
                self.value_stack[-1][0].append(v)
                self.terminator = self.value_stack[-1][1]

            return v

class CljEncoder(object):
    def __init__(self, data, fd):
        self.data = data
        self.fd = fd
        self.circular = {}

    def encode(self):
        self.__do_encode(self.data)

    def get_type(self,t):
        if t is None:
            return ("None", False)
        elif isinstance(t, str):
            return ("string", False)
        elif isinstance(t, bool):
            return ("boolean", False)
        elif isinstance(t, (int,float,long)):
            return ("number", False)
        elif isinstance(t, dict):
            return ("dict", True)
        elif isinstance(t, list):
            return ("list", True)
        elif isinstance(t, set):
            return ("set", True)
        else:
            return ("unknown", False)

    def __do_encode(self, d):
        fd = self.fd
        t,coll = self.get_type(d)

        if coll:
            refid = id(d)
            if refid in self.circular:
                raise ValueError('Circular reference detected')
            else:
                self.circular[refid] = True

            if t == "dict":
                fd.write("{")
                for k,v in d.items():
                    self.__do_encode(k)
                    fd.write(" ")
                    self.__do_encode(v)
                    fd.write(" ")
                fd.seek(-1, os.SEEK_CUR)
                fd.write("}")
            elif t == "list":
                fd.write("[")
                for v in d:
                    self.__do_encode(v)
                    fd.write(" ")
                fd.seek(-1, os.SEEK_CUR)
                fd.write("]")
            elif t == "set":
                fd.write("#{")
                for v in d:
                    self.__do_encode(v)
                    fd.write(" ")
                fd.seek(-1, os.SEEK_CUR)
                fd.write("}")
        else:
            if t == "number":
                fd.write(str(d))
            elif t == "string":
                s = d.encode("string-escape").replace('"', '\\"')
                fd.write('"'+s+'"')
            elif t == "boolean":
                if d:
                    fd.write('true')
                else:
                    fd.write('false')
            elif t == 'None':
                fd.write('nil')
            else:
                fd.write('"'+str(d)+'"')

def dump(obj, fp):
    return CljEncoder(obj, fp).encode()

def dumps(obj):
    buf = StringIO()
    dump(obj, buf)
    result = buf.getvalue()
    buf.close()
    return result

def load(fp):
    decoder = CljDecoder(fp)
    return decoder.decode()

def loads(s):
    buf = StringIO(s)
    result = load(buf)
    buf.close()
    return result

<|endoftext|>


<commit_before>"""
Translate gene regions from nucleotides to amino acids.
"""

import os, sys
import numpy as np
from Bio import SeqIO, SeqFeature, Seq, SeqRecord, Phylo
from .utils import read_node_data, load_features, write_json, write_VCF_translation
from treetime.vcf_utils import read_vcf

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

    Optionally, returns a tuple of the translated sequence and whether an
    exception was raised during initial translation.

    >>> safe_translate("ATG")
    'M'
    >>> safe_translate("ATGGT-")
    'MX'
    >>> safe_translate("ATG---")
    'M-'
    >>> safe_translate("ATGTAG")
    'M*'
    >>> safe_translate("")
    ''
    >>> safe_translate("ATGT")
    'MX'
    >>> safe_translate("ATG", report_exceptions=True)
    ('M', False)
    >>> safe_translate("ATGA-G", report_exceptions=True)
    ('MX', True)
    """
    from Bio.Data.CodonTable import TranslationError
    from Bio.Seq import CodonTable
    translation_exception = False

    #sequences not mod 3 give messy BiopythonWarning, so avoid by padding.
    if len(sequence)%3:
        sequence_padded = sequence + "N"*(3-len(sequence)%3)
    else:
        sequence_padded = sequence
    try:
        # Attempt translation by extracting the sequence according to the
        # BioPhython SeqFeature in frame gaps of three will translate as '-'
        translated_sequence = str(Seq.Seq(sequence_padded).translate(gap='-'))
    except TranslationError:
        translation_exception = True
        # Any other codon like '-AA' or 'NNT' etc will fail. Translate codons
        # one by one.
        codon_table  = CodonTable.ambiguous_dna_by_name['Standard'].forward_table
        str_seq = str(sequence_padded)
        codons = np.frombuffer(str_seq[:len(str_seq) - len(str_seq) % 3].encode(), dtype='S3').astype("U")
        assert len(codons) > 0
        aas = []

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

        translated_sequence = "".join(aas)

    if report_exceptions:
        return translated_sequence, translation_exception
    else:
        return translated_sequence

def translate_feature(aln, feature):
    '''
    Translates a subsequence of input nucleotide sequences.

    Parameters
    ----------
    aln : dict
        sequences indexed by node name

    feature : Bio.Seq.Seq
        BioPython sequence feature

    Returns
    -------
    dict :
        translated sequences indexed by node name

    '''
    translations = {}
    for sname, seq in aln.items():
        aa_seq = safe_translate(str(feature.extract(seq)))
        translations[sname] = aa_seq

    return translations

def translate_vcf_feature(sequences, ref, feature):
    '''Translates a subsequence of input nucleotide sequences.

    Parameters
    ----------
    sequences : dict
        TreeTime format dictionary from VCF-input of sequences indexed by node name

    ref :
        reference alignment the VCF was mapped to

    feature : Bio.Seq.Seq
        BioPython sequence feature

    Returns
    -------
    dict :
        translated reference gene, positions of AA differences, and AA
        differences indexed by node name

    '''
    def str_reverse_comp(str_seq):
        #gets reverse-compliment of a string and returns it as a string
        seq_str = Seq(str_seq)
        return str(seq_str.reverse_complement())

    from Bio.Seq import Seq
    from Bio.SeqRecord import SeqRecord
    prot = {}
    prot['sequences'] = {}
    prot['positions'] = []

    refNuc = str(feature.extract( SeqRecord(seq=Seq(ref)) ).seq)
    # Need to get ref translation to store. check if multiple of 3 for sanity.
    # will be padded in safe_translate if not
    if len(refNuc)%3:
        print("Gene length of {} is not a multiple of 3. will pad with N".format(feature.qualifiers['Name'][0]), file=sys.stderr)

    ref_aa_seq = safe_translate(refNuc)
    prot['reference'] = ref_aa_seq

    start = int(feature.location.start) #1 less than GFF file; Bio 'pythons' the numbering
    end = int(feature.location.end)

    for seqk in sequences.keys():
        varSite = np.array(list(sequences[seqk].keys()))  #get positions where nuc diffs
        #reduce to only those within current gene
        geneVarSites = np.logical_and(varSite >= start, varSite < end)
        nucVarSites = varSite[geneVarSites] #translate this back to nuc position
        genNucSites = nucVarSites-start #get it in position within the gene to ensure in frame

        #Translate just the codon this nuc diff is in, and find out which AA loc
        #But need numbering to be w/in protin, not whole genome
        if feature.strand == -1:
            aaRepLocs = {(end-start-i-1)//3:safe_translate( str_reverse_comp( "".join([sequences[seqk][key+start]
                                    if key+start in sequences[seqk].keys() else ref[key+start]
                                for key in range(i-i%3,i+3-i%3)]) ))
                            for i in genNucSites}
        else:
            aaRepLocs = {i//3:safe_translate( "".join([sequences[seqk][key+start]
                                if key+start in sequences[seqk].keys() else ref[key+start]
                            for key in range(i-i%3,i+3-i%3)]) )
                        for i in genNucSites}

        aaRepLocsFinal = {}
        #remove if is a synonymous mutation!
        for key,val in aaRepLocs.items():
            if ref_aa_seq[key] != val:
                aaRepLocsFinal[key] = val
        aaRepLocs = aaRepLocsFinal

        #store the dict of differences
        prot['sequences'][seqk] = aaRepLocs
        #add to list of positions if needed
        for key in aaRepLocs.keys():
            if key not in prot['positions']:
                prot['positions'].append(key)

    prot['positions'].sort()

    #if no variable sites, exclude this gene
    if len(prot['positions']) == 0:
        return None
    else:
        return prot

def construct_mut(start, pos, end):
    return str(start) + str(pos) + str(end)

def assign_aa_vcf(tree, translations):
    aa_muts = {}

    #get mutations on the root
    root = tree.root
    aa_muts[root.name]={"aa_muts":{}}
    for fname, prot in translations.items():
        root_muts = prot['sequences'][root.name]
        tmp = []
        for pos in prot['positions']:
            if pos in root_muts:
                tmp.append(construct_mut(prot['reference'][pos], int(pos+1), root_muts[pos]))
        aa_muts[root.name]["aa_muts"][fname] = tmp

    for n in tree.get_nonterminals():
        for c in n:
            aa_muts[c.name]={"aa_muts":{}}
        for fname, prot in translations.items():
            n_muts = prot['sequences'][n.name]
            for c in n:
                tmp = []
                c_muts = prot['sequences'][c.name]
                for pos in prot['positions']:
                    #if pos in both, check if same
                    if pos in n_muts and pos in c_muts:
                        if n_muts[pos] != c_muts[pos]:
                            tmp.append(construct_mut(n_muts[pos], int(pos+1), c_muts[pos]))
                    elif pos in n_muts:
                        tmp.append(construct_mut(n_muts[pos], int(pos+1), prot['reference'][pos]))
                    elif pos in c_muts:
                        tmp.append(construct_mut(prot['reference'][pos], int(pos+1), c_muts[pos]))

                aa_muts[c.name]["aa_muts"][fname] = tmp

    return aa_muts

def get_genes_from_file(fname):
    genes = []
    if os.path.isfile(fname):
        with open(fname) as ifile:
            for line in ifile:
                fields = line.strip().split('#')
                if fields[0].strip():
                    genes.append(fields[0].strip())
    else:
        print("File with genes not found. Looking for", fname)

    unique_genes = np.unique(np.array(genes))
    if len(unique_genes) != len(genes):
        print("You have duplicates in your genes file. They are being ignored.")
    print("Read in {} specified genes to translate.".format(len(unique_genes)))

    return unique_genes

def register_arguments(parser):
    parser.add_argument('--tree', help="prebuilt Newick -- no tree will be built if provided")
    parser.add_argument('--ancestral-sequences', type=str, help='JSON (fasta input) or VCF (VCF input) containing ancestral and tip sequences')
    parser.add_argument('--reference-sequence', required=True,
                        help='GenBank or GFF file containing the annotation')
    parser.add_argument('--genes', nargs='+', help="genes to translate (list or file containing list)")
    parser.add_argument('--output', type=str, help="name of JSON files for aa mutations")
    parser.add_argument('--alignment-output', type=str, help="write out translated gene alignments. "
                                   "If a VCF-input, a .vcf or .vcf.gz will be output here (depending on file ending). If fasta-input, specify the file name "
                                   "like so: 'my_alignment_%%GENE.fasta', where '%%GENE' will be replaced by the name of the gene")
    parser.add_argument('--vcf-reference-output', type=str, help="fasta file where reference sequence translations for VCF input will be written")
    parser.add_argument('--vcf-reference', type=str, help='fasta file of the sequence the VCF was mapped to')

def run(args):
    ## read tree and data, if reading data fails, return with error code
    tree = Phylo.read(args.tree, 'newick')

    # If genes is a file, read in the genes to translate
    if args.genes and len(args.genes) == 1 and os.path.isfile(args.genes[0]):
        genes = get_genes_from_file(args.genes[0])
    else:
        genes = args.genes

    ## check file format and read in sequences
    is_vcf = False
    if any([args.ancestral_sequences.lower().endswith(x) for x in ['.vcf', '.vcf.gz']]):
        if not args.vcf_reference:
            print("ERROR: a reference Fasta is required with VCF-format input")
            return 1
        compress_seq = read_vcf(args.ancestral_sequences, args.vcf_reference)
        sequences = compress_seq['sequences']
        ref = compress_seq['reference']
        is_vcf = True
    else:
        node_data = read_node_data(args.ancestral_sequences, args.tree)
        if node_data is None:
            print("ERROR: could not read node data (incl sequences)")
            return 1
        # extract sequences from node meta data
        sequences = {}
        for k,v in node_data['nodes'].items():
            if 'sequence' in v:
                sequences[k] = v['sequence']

    ## load features; only requested features if genes given
    features = load_features(args.reference_sequence, genes)
    print("Read in {} features from reference sequence file".format(len(features)))
    if features is None:
        print("ERROR: could not read features of reference sequence file")
        return 1

    ### translate every feature - but not 'nuc'!
    translations = {}
    deleted = []
    for fname, feat in features.items():
        if is_vcf:
            trans = translate_vcf_feature(sequences, ref, feat)
            if trans:
                translations[fname] = trans
            else:
                deleted.append(fname)
        else:
            if feat.type != 'source':
                translations[fname] = translate_feature(sequences, feat)

    if len(deleted) != 0:
        print("{} genes had no mutations and so have been be excluded.".format(len(deleted)))

    ## glob the annotations for later auspice export
    #
    # Note that both our JSON schema and BioPython FeatureLocations use
    # "Pythonic" coordinates: [zero-origin, half-open).
    annotations = {}
    for fname, feat in features.items():
        annotations[fname] = {'seqid':ref.id,
                              'type':feat.type,
                              'start':int(feat.location.start)+1,
                              'end':int(feat.location.end),
                              'strand': '+' if feat.location.strand else '-'}
    if is_vcf: #need to add our own nuc
        annotations['nuc'] = {'seqid':"reference",
                              'type':ref.id,
                              'start': 1,
                              'end': len(ref),
                              'strand': '+'}

    ## determine amino acid mutations for each node
    if is_vcf:
        aa_muts = assign_aa_vcf(tree, translations)
    else:
        aa_muts = {}

        #fasta input shouldn't have mutations on root, so give empty entry
        root = tree.get_nonterminals()[0]
        aa_muts[root.name]={"aa_muts":{}}

        for n in tree.get_nonterminals():
            for c in n:
                aa_muts[c.name]={"aa_muts":{}}
            for fname, aln in translations.items():
                for c in n:
                    if c.name in aln and n.name in aln:
                        tmp = [construct_mut(a, int(pos+1), d) for pos, (a,d) in
                                enumerate(zip(aln[n.name], aln[c.name])) if a!=d]
                        aa_muts[c.name]["aa_muts"][fname] = tmp
                    else:
                        print("no sequence pair for nodes %s-%s"%(c.name, n.name))
            if n==tree.root:
                aa_muts[n.name]={"aa_muts":{}, "aa_sequences":{}}
                for fname, aln in translations.items():
                    if n.name in aln:
                        aa_muts[n.name]["aa_sequences"][fname] = "".join(aln[n.name])

    write_json({'annotations':annotations, 'nodes':aa_muts}, args.output)
    print("amino acid mutations written to",args.output, file=sys.stdout)

    ## write alignments to file is requested
    if args.alignment_output:
        if is_vcf:
            ## write VCF-style output if requested
            fileEndings = -1
            if args.alignment_output.lower().endswith('.gz'):
                fileEndings = -2
            vcf_out_ref = args.vcf_reference_output or '.'.join(args.alignment_output.split('.')[:fileEndings]) + '_reference.fasta'
            write_VCF_translation(translations, args.alignment_output, vcf_out_ref)
        else:
            ## write fasta-style output if requested
            if '%GENE' in args.alignment_output:
                for fname, seqs in translations.items():
                    SeqIO.write([SeqRecord.SeqRecord(seq=Seq.Seq(s), id=sname, name=sname, description='')
                                 for sname, s in seqs.items()],
                                 args.alignment_output.replace('%GENE', fname), 'fasta')
            else:
                print("ERROR: alignment output file does not contain '%GENE', so will not be written.")
<commit_msg>augur.translate: use reference file name as ref id, should be replaced by the accession number<commit_after>"""
Translate gene regions from nucleotides to amino acids.
"""

import os, sys
import numpy as np
from Bio import SeqIO, SeqFeature, Seq, SeqRecord, Phylo
from .utils import read_node_data, load_features, write_json, write_VCF_translation
from treetime.vcf_utils import read_vcf

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

    Optionally, returns a tuple of the translated sequence and whether an
    exception was raised during initial translation.

    >>> safe_translate("ATG")
    'M'
    >>> safe_translate("ATGGT-")
    'MX'
    >>> safe_translate("ATG---")
    'M-'
    >>> safe_translate("ATGTAG")
    'M*'
    >>> safe_translate("")
    ''
    >>> safe_translate("ATGT")
    'MX'
    >>> safe_translate("ATG", report_exceptions=True)
    ('M', False)
    >>> safe_translate("ATGA-G", report_exceptions=True)
    ('MX', True)
    """
    from Bio.Data.CodonTable import TranslationError
    from Bio.Seq import CodonTable
    translation_exception = False

    #sequences not mod 3 give messy BiopythonWarning, so avoid by padding.
    if len(sequence)%3:
        sequence_padded = sequence + "N"*(3-len(sequence)%3)
    else:
        sequence_padded = sequence
    try:
        # Attempt translation by extracting the sequence according to the
        # BioPhython SeqFeature in frame gaps of three will translate as '-'
        translated_sequence = str(Seq.Seq(sequence_padded).translate(gap='-'))
    except TranslationError:
        translation_exception = True
        # Any other codon like '-AA' or 'NNT' etc will fail. Translate codons
        # one by one.
        codon_table  = CodonTable.ambiguous_dna_by_name['Standard'].forward_table
        str_seq = str(sequence_padded)
        codons = np.frombuffer(str_seq[:len(str_seq) - len(str_seq) % 3].encode(), dtype='S3').astype("U")
        assert len(codons) > 0
        aas = []

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

        translated_sequence = "".join(aas)

    if report_exceptions:
        return translated_sequence, translation_exception
    else:
        return translated_sequence

def translate_feature(aln, feature):
    '''
    Translates a subsequence of input nucleotide sequences.

    Parameters
    ----------
    aln : dict
        sequences indexed by node name

    feature : Bio.Seq.Seq
        BioPython sequence feature

    Returns
    -------
    dict :
        translated sequences indexed by node name

    '''
    translations = {}
    for sname, seq in aln.items():
        aa_seq = safe_translate(str(feature.extract(seq)))
        translations[sname] = aa_seq

    return translations

def translate_vcf_feature(sequences, ref, feature):
    '''Translates a subsequence of input nucleotide sequences.

    Parameters
    ----------
    sequences : dict
        TreeTime format dictionary from VCF-input of sequences indexed by node name

    ref :
        reference alignment the VCF was mapped to

    feature : Bio.Seq.Seq
        BioPython sequence feature

    Returns
    -------
    dict :
        translated reference gene, positions of AA differences, and AA
        differences indexed by node name

    '''
    def str_reverse_comp(str_seq):
        #gets reverse-compliment of a string and returns it as a string
        seq_str = Seq(str_seq)
        return str(seq_str.reverse_complement())

    from Bio.Seq import Seq
    from Bio.SeqRecord import SeqRecord
    prot = {}
    prot['sequences'] = {}
    prot['positions'] = []

    refNuc = str(feature.extract( SeqRecord(seq=Seq(ref)) ).seq)
    # Need to get ref translation to store. check if multiple of 3 for sanity.
    # will be padded in safe_translate if not
    if len(refNuc)%3:
        print("Gene length of {} is not a multiple of 3. will pad with N".format(feature.qualifiers['Name'][0]), file=sys.stderr)

    ref_aa_seq = safe_translate(refNuc)
    prot['reference'] = ref_aa_seq

    start = int(feature.location.start) #1 less than GFF file; Bio 'pythons' the numbering
    end = int(feature.location.end)

    for seqk in sequences.keys():
        varSite = np.array(list(sequences[seqk].keys()))  #get positions where nuc diffs
        #reduce to only those within current gene
        geneVarSites = np.logical_and(varSite >= start, varSite < end)
        nucVarSites = varSite[geneVarSites] #translate this back to nuc position
        genNucSites = nucVarSites-start #get it in position within the gene to ensure in frame

        #Translate just the codon this nuc diff is in, and find out which AA loc
        #But need numbering to be w/in protin, not whole genome
        if feature.strand == -1:
            aaRepLocs = {(end-start-i-1)//3:safe_translate( str_reverse_comp( "".join([sequences[seqk][key+start]
                                    if key+start in sequences[seqk].keys() else ref[key+start]
                                for key in range(i-i%3,i+3-i%3)]) ))
                            for i in genNucSites}
        else:
            aaRepLocs = {i//3:safe_translate( "".join([sequences[seqk][key+start]
                                if key+start in sequences[seqk].keys() else ref[key+start]
                            for key in range(i-i%3,i+3-i%3)]) )
                        for i in genNucSites}

        aaRepLocsFinal = {}
        #remove if is a synonymous mutation!
        for key,val in aaRepLocs.items():
            if ref_aa_seq[key] != val:
                aaRepLocsFinal[key] = val
        aaRepLocs = aaRepLocsFinal

        #store the dict of differences
        prot['sequences'][seqk] = aaRepLocs
        #add to list of positions if needed
        for key in aaRepLocs.keys():
            if key not in prot['positions']:
                prot['positions'].append(key)

    prot['positions'].sort()

    #if no variable sites, exclude this gene
    if len(prot['positions']) == 0:
        return None
    else:
        return prot

def construct_mut(start, pos, end):
    return str(start) + str(pos) + str(end)

def assign_aa_vcf(tree, translations):
    aa_muts = {}

    #get mutations on the root
    root = tree.root
    aa_muts[root.name]={"aa_muts":{}}
    for fname, prot in translations.items():
        root_muts = prot['sequences'][root.name]
        tmp = []
        for pos in prot['positions']:
            if pos in root_muts:
                tmp.append(construct_mut(prot['reference'][pos], int(pos+1), root_muts[pos]))
        aa_muts[root.name]["aa_muts"][fname] = tmp

    for n in tree.get_nonterminals():
        for c in n:
            aa_muts[c.name]={"aa_muts":{}}
        for fname, prot in translations.items():
            n_muts = prot['sequences'][n.name]
            for c in n:
                tmp = []
                c_muts = prot['sequences'][c.name]
                for pos in prot['positions']:
                    #if pos in both, check if same
                    if pos in n_muts and pos in c_muts:
                        if n_muts[pos] != c_muts[pos]:
                            tmp.append(construct_mut(n_muts[pos], int(pos+1), c_muts[pos]))
                    elif pos in n_muts:
                        tmp.append(construct_mut(n_muts[pos], int(pos+1), prot['reference'][pos]))
                    elif pos in c_muts:
                        tmp.append(construct_mut(prot['reference'][pos], int(pos+1), c_muts[pos]))

                aa_muts[c.name]["aa_muts"][fname] = tmp

    return aa_muts

def get_genes_from_file(fname):
    genes = []
    if os.path.isfile(fname):
        with open(fname) as ifile:
            for line in ifile:
                fields = line.strip().split('#')
                if fields[0].strip():
                    genes.append(fields[0].strip())
    else:
        print("File with genes not found. Looking for", fname)

    unique_genes = np.unique(np.array(genes))
    if len(unique_genes) != len(genes):
        print("You have duplicates in your genes file. They are being ignored.")
    print("Read in {} specified genes to translate.".format(len(unique_genes)))

    return unique_genes

def register_arguments(parser):
    parser.add_argument('--tree', help="prebuilt Newick -- no tree will be built if provided")
    parser.add_argument('--ancestral-sequences', type=str, help='JSON (fasta input) or VCF (VCF input) containing ancestral and tip sequences')
    parser.add_argument('--reference-sequence', required=True,
                        help='GenBank or GFF file containing the annotation')
    parser.add_argument('--genes', nargs='+', help="genes to translate (list or file containing list)")
    parser.add_argument('--output', type=str, help="name of JSON files for aa mutations")
    parser.add_argument('--alignment-output', type=str, help="write out translated gene alignments. "
                                   "If a VCF-input, a .vcf or .vcf.gz will be output here (depending on file ending). If fasta-input, specify the file name "
                                   "like so: 'my_alignment_%%GENE.fasta', where '%%GENE' will be replaced by the name of the gene")
    parser.add_argument('--vcf-reference-output', type=str, help="fasta file where reference sequence translations for VCF input will be written")
    parser.add_argument('--vcf-reference', type=str, help='fasta file of the sequence the VCF was mapped to')

def run(args):
    ## read tree and data, if reading data fails, return with error code
    tree = Phylo.read(args.tree, 'newick')

    # If genes is a file, read in the genes to translate
    if args.genes and len(args.genes) == 1 and os.path.isfile(args.genes[0]):
        genes = get_genes_from_file(args.genes[0])
    else:
        genes = args.genes

    ## check file format and read in sequences
    is_vcf = False
    if any([args.ancestral_sequences.lower().endswith(x) for x in ['.vcf', '.vcf.gz']]):
        if not args.vcf_reference:
            print("ERROR: a reference Fasta is required with VCF-format input")
            return 1
        compress_seq = read_vcf(args.ancestral_sequences, args.vcf_reference)
        sequences = compress_seq['sequences']
        ref = compress_seq['reference']
        is_vcf = True
    else:
        node_data = read_node_data(args.ancestral_sequences, args.tree)
        if node_data is None:
            print("ERROR: could not read node data (incl sequences)")
            return 1
        # extract sequences from node meta data
        sequences = {}
        for k,v in node_data['nodes'].items():
            if 'sequence' in v:
                sequences[k] = v['sequence']

    ## load features; only requested features if genes given
    features = load_features(args.reference_sequence, genes)
    print("Read in {} features from reference sequence file".format(len(features)))
    if features is None:
        print("ERROR: could not read features of reference sequence file")
        return 1

    ### translate every feature - but not 'nuc'!
    translations = {}
    deleted = []
    for fname, feat in features.items():
        if is_vcf:
            trans = translate_vcf_feature(sequences, ref, feat)
            if trans:
                translations[fname] = trans
            else:
                deleted.append(fname)
        else:
            if feat.type != 'source':
                translations[fname] = translate_feature(sequences, feat)

    if len(deleted) != 0:
        print("{} genes had no mutations and so have been be excluded.".format(len(deleted)))

    ## glob the annotations for later auspice export
    #
    # Note that both our JSON schema and BioPython FeatureLocations use
    # "Pythonic" coordinates: [zero-origin, half-open).
    annotations = {}
    for fname, feat in features.items():
        annotations[fname] = {'seqid':args.reference_sequence,
                              'type':feat.type,
                              'start':int(feat.location.start)+1,
                              'end':int(feat.location.end),
                              'strand': '+' if feat.location.strand else '-'}
    if is_vcf: #need to add our own nuc
        annotations['nuc'] = {'seqid':args.reference_sequence,
                              'type':ref.id,
                              'start': 1,
                              'end': len(ref),
                              'strand': '+'}

    ## determine amino acid mutations for each node
    if is_vcf:
        aa_muts = assign_aa_vcf(tree, translations)
    else:
        aa_muts = {}

        #fasta input shouldn't have mutations on root, so give empty entry
        root = tree.get_nonterminals()[0]
        aa_muts[root.name]={"aa_muts":{}}

        for n in tree.get_nonterminals():
            for c in n:
                aa_muts[c.name]={"aa_muts":{}}
            for fname, aln in translations.items():
                for c in n:
                    if c.name in aln and n.name in aln:
                        tmp = [construct_mut(a, int(pos+1), d) for pos, (a,d) in
                                enumerate(zip(aln[n.name], aln[c.name])) if a!=d]
                        aa_muts[c.name]["aa_muts"][fname] = tmp
                    else:
                        print("no sequence pair for nodes %s-%s"%(c.name, n.name))
            if n==tree.root:
                aa_muts[n.name]={"aa_muts":{}, "aa_sequences":{}}
                for fname, aln in translations.items():
                    if n.name in aln:
                        aa_muts[n.name]["aa_sequences"][fname] = "".join(aln[n.name])

    write_json({'annotations':annotations, 'nodes':aa_muts}, args.output)
    print("amino acid mutations written to",args.output, file=sys.stdout)

    ## write alignments to file is requested
    if args.alignment_output:
        if is_vcf:
            ## write VCF-style output if requested
            fileEndings = -1
            if args.alignment_output.lower().endswith('.gz'):
                fileEndings = -2
            vcf_out_ref = args.vcf_reference_output or '.'.join(args.alignment_output.split('.')[:fileEndings]) + '_reference.fasta'
            write_VCF_translation(translations, args.alignment_output, vcf_out_ref)
        else:
            ## write fasta-style output if requested
            if '%GENE' in args.alignment_output:
                for fname, seqs in translations.items():
                    SeqIO.write([SeqRecord.SeqRecord(seq=Seq.Seq(s), id=sname, name=sname, description='')
                                 for sname, s in seqs.items()],
                                 args.alignment_output.replace('%GENE', fname), 'fasta')
            else:
                print("ERROR: alignment output file does not contain '%GENE', so will not be written.")
<|endoftext|>


<commit_before># coding: utf-8

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Third-party
import astropy.units as u
from astropy.coordinates.angles import rotation_matrix
import gary.integrate as gi
import numpy as np
from scipy.signal import argrelmin
from sklearn.neighbors import KernelDensity

from .fast_ensemble import ensemble_integrate

__all__ = ['create_ball', 'nearest_pericenter', 'align_ensemble']

def create_ball(w0, potential, N=1000, m_scale=1E4):
    menc = potential.mass_enclosed(w0)
    rscale = (m_scale / menc)**(1/3.) * np.sqrt(np.sum(w0[:3]**2))
    vscale = (m_scale / menc)**(1/3.) * np.sqrt(np.sum(w0[3:]**2))

    ball_w0 = np.zeros((N,6))
    ball_w0[:,:3] = np.random.normal(w0[:3], rscale, size=(N,3))
    ball_w0[:,3:] = np.random.normal(w0[3:], vscale, size=(N,3))

    return np.vstack((w0,ball_w0))

def nearest_pericenter(w0, potential, dt, T):
    """
    Find the nearest pericenter to w0
    """

    t,w = potential.integrate_orbit(w0, dt=dt, nsteps=T*10,
                                    Integrator=gi.DOPRI853Integrator)

    r = np.sqrt(np.sum(w[:,0,:3]**2, axis=-1))
    peris = argrelmin(r)[0]

    return w[peris[0], 0]

def compute_align_matrix(w):
    """
    Given a single phase-space position, compute the rotation matrix that
    orients the angular momentum with the z axis and places the point
    along the x axis.

    Parameters
    ----------
    w : array_like
    """

    if w.ndim > 1:
        raise ValueError("Input phase-space position should be 1D.")

    x = w[:3].copy()
    v = w[3:].copy()

    # first rotate about z to put on x-z plane
    theta = np.arctan2(x[1], x[0]) * u.radian
    R1 = rotation_matrix(theta, 'z')
    x = np.asarray(R1.dot(x))[0]
    v = np.asarray(R1.dot(v))[0]

    # now rotate about y to put on x axis
    theta = np.arctan2(x[2], x[0]) * u.radian
    R2 = rotation_matrix(-theta, 'y')
    x = np.asarray(R2.dot(x))[0]
    v = np.asarray(R2.dot(v))[0]

    # now align L with z axis
    # theta = np.arccos(L[2] / np.sqrt(np.sum(L**2))) * u.radian
    L = np.cross(x, v)
    theta = np.arctan2(L[2], L[1]) * u.radian
    R3 = rotation_matrix(theta - 90*u.deg, 'x')
    x = np.asarray(R3.dot(x))[0]
    v = np.asarray(R3.dot(v))[0]

    return R3*R2*R1

def align_ensemble(ws):
    R = compute_align_matrix(ws[-1,0])
    new_x = np.array(R.dot(ws[-1,:,:3].T).T)
    new_v = np.array(R.dot(ws[-1,:,3:].T).T)
    return new_x, new_v

def do_the_kld(nkld, ball_w0, potential, dt, nsteps, kde_bandwidth,
               density_thresholds):
    """
    Parameters
    ----------
    nkld : int
        Number of times to evaluate the KLD.
    ...TODO
    """

    ww = np.ascontiguousarray(ball_w0.copy())
    nensemble = ww.shape[0]

    # energy of parent orbit
    E0 = float(np.squeeze(potential.total_energy(ww[0,:3],ww[0,3:])))
    predicted_density = lambda x, E0: np.sqrt(E0 - potential(x))

    kld_idx = np.linspace(0, nsteps, nkld+1).astype(int)

    # container to store fraction of stars with density above each threshold
    frac_above_dens = np.zeros((nkld,len(density_thresholds)))

    t = np.empty(nkld+1)
    t[0] = 0.
    kld = np.empty(nkld)
    for i in range(nkld):
        # number of steps to advance the ensemble
        dstep = kld_idx[i+1] - kld_idx[i]
        www = ensemble_integrate(potential.c_instance, ww, dt, dstep, 0.)

        # store the time
        t[i+1] = t[i] + dt*dstep

        # build an estimate of the configuration space density of the ensemble
        kde = KernelDensity(kernel='epanechnikov', bandwidth=kde_bandwidth)
        kde.fit(www[:,:3])
        kde_densy = np.exp(kde.score_samples(www[:,:3]))

        # number of stars with density > threshold
        for j,h in enumerate(density_thresholds):
            frac_above_dens[i,j] = (kde_densy > h).sum() / float(nensemble)

        p_densy = predicted_density(www[:,:3], E0)
        D = np.log(kde_densy / p_densy)
        KLD = D[np.isfinite(D)].sum() / float(nensemble)
        kld[i] = KLD

        ww = www.copy()

    print(frac_above_dens)
    return t[1:], kld, frac_above_dens
<commit_msg>remove print<commit_after># coding: utf-8

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Third-party
import astropy.units as u
from astropy.coordinates.angles import rotation_matrix
import gary.integrate as gi
import numpy as np
from scipy.signal import argrelmin
from sklearn.neighbors import KernelDensity

from .fast_ensemble import ensemble_integrate

__all__ = ['create_ball', 'nearest_pericenter', 'align_ensemble']

def create_ball(w0, potential, N=1000, m_scale=1E4):
    menc = potential.mass_enclosed(w0)
    rscale = (m_scale / menc)**(1/3.) * np.sqrt(np.sum(w0[:3]**2))
    vscale = (m_scale / menc)**(1/3.) * np.sqrt(np.sum(w0[3:]**2))

    ball_w0 = np.zeros((N,6))
    ball_w0[:,:3] = np.random.normal(w0[:3], rscale, size=(N,3))
    ball_w0[:,3:] = np.random.normal(w0[3:], vscale, size=(N,3))

    return np.vstack((w0,ball_w0))

def nearest_pericenter(w0, potential, dt, T):
    """
    Find the nearest pericenter to w0
    """

    t,w = potential.integrate_orbit(w0, dt=dt, nsteps=T*10,
                                    Integrator=gi.DOPRI853Integrator)

    r = np.sqrt(np.sum(w[:,0,:3]**2, axis=-1))
    peris = argrelmin(r)[0]

    return w[peris[0], 0]

def compute_align_matrix(w):
    """
    Given a single phase-space position, compute the rotation matrix that
    orients the angular momentum with the z axis and places the point
    along the x axis.

    Parameters
    ----------
    w : array_like
    """

    if w.ndim > 1:
        raise ValueError("Input phase-space position should be 1D.")

    x = w[:3].copy()
    v = w[3:].copy()

    # first rotate about z to put on x-z plane
    theta = np.arctan2(x[1], x[0]) * u.radian
    R1 = rotation_matrix(theta, 'z')
    x = np.asarray(R1.dot(x))[0]
    v = np.asarray(R1.dot(v))[0]

    # now rotate about y to put on x axis
    theta = np.arctan2(x[2], x[0]) * u.radian
    R2 = rotation_matrix(-theta, 'y')
    x = np.asarray(R2.dot(x))[0]
    v = np.asarray(R2.dot(v))[0]

    # now align L with z axis
    # theta = np.arccos(L[2] / np.sqrt(np.sum(L**2))) * u.radian
    L = np.cross(x, v)
    theta = np.arctan2(L[2], L[1]) * u.radian
    R3 = rotation_matrix(theta - 90*u.deg, 'x')
    x = np.asarray(R3.dot(x))[0]
    v = np.asarray(R3.dot(v))[0]

    return R3*R2*R1

def align_ensemble(ws):
    R = compute_align_matrix(ws[-1,0])
    new_x = np.array(R.dot(ws[-1,:,:3].T).T)
    new_v = np.array(R.dot(ws[-1,:,3:].T).T)
    return new_x, new_v

def do_the_kld(nkld, ball_w0, potential, dt, nsteps, kde_bandwidth,
               density_thresholds):
    """
    Parameters
    ----------
    nkld : int
        Number of times to evaluate the KLD.
    ...TODO
    """

    ww = np.ascontiguousarray(ball_w0.copy())
    nensemble = ww.shape[0]

    # energy of parent orbit
    E0 = float(np.squeeze(potential.total_energy(ww[0,:3],ww[0,3:])))
    predicted_density = lambda x, E0: np.sqrt(E0 - potential(x))

    kld_idx = np.linspace(0, nsteps, nkld+1).astype(int)

    # container to store fraction of stars with density above each threshold
    frac_above_dens = np.zeros((nkld,len(density_thresholds)))

    t = np.empty(nkld+1)
    t[0] = 0.
    kld = np.empty(nkld)
    for i in range(nkld):
        # number of steps to advance the ensemble
        dstep = kld_idx[i+1] - kld_idx[i]
        www = ensemble_integrate(potential.c_instance, ww, dt, dstep, 0.)

        # store the time
        t[i+1] = t[i] + dt*dstep

        # build an estimate of the configuration space density of the ensemble
        kde = KernelDensity(kernel='epanechnikov', bandwidth=kde_bandwidth)
        kde.fit(www[:,:3])
        kde_densy = np.exp(kde.score_samples(www[:,:3]))

        # number of stars with density > threshold
        for j,h in enumerate(density_thresholds):
            frac_above_dens[i,j] = (kde_densy > h).sum() / float(nensemble)

        p_densy = predicted_density(www[:,:3], E0)
        D = np.log(kde_densy / p_densy)
        KLD = D[np.isfinite(D)].sum() / float(nensemble)
        kld[i] = KLD

        ww = www.copy()

    return t[1:], kld, frac_above_dens
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

from cmd2 import Cmd
import argparse
import sys
import os
import math
import types
import traceback
import pid
import time
import util
import errors
import textwrap
import openvpn
from blessed import Terminal
from consts import *
from core import Core
from config import Config, EBSettings
from registration import Registration, InfoLoader
from softhsm import SoftHsmV1Config
from ejbca import Ejbca
from ebsysconfig import SysConfig
from letsencrypt import LetsEncrypt
from ebclient.registration import ENVIRONMENT_PRODUCTION, ENVIRONMENT_DEVELOPMENT, ENVIRONMENT_TEST
from pkg_resources import get_distribution, DistributionNotFound
from cli import Installer
import logging
import coloredlogs

logger = logging.getLogger(__name__)
coloredlogs.install(level=logging.ERROR)

class VpnInstaller(Installer):
    """
    Extended installer - with VPN.
    """

    def __init__(self, *args, **kwargs):
        """
        Init core
        :param args:
        :param kwargs:
        :return:
        """
        Installer.__init__(self, *args, **kwargs)
        self.ovpn = None

    def init_argparse(self):
        """
        Adding new VPN related arguments
        :return:
        """
        parser = Installer.init_argparse(self)
        return parser

    def ask_for_email_reason(self, is_required=None):
        """
        Reason why we need email - required in VPN case.
        :param is_required:
        :return:
        """
        if is_required:
            self.tprint('We need your email address for:\n'
                        '   a) identity verification for EnigmaBridge account \n'
                        '   b) LetsEncrypt certificate registration\n'
                        '   c) PKI setup - VPN configuration')
            self.tprint('We will send you a verification email.')
            self.tprint('Without a valid e-mail address you won\'t be able to continue with the installation\n')
        else:
            raise ValueError('Email is required in VPN case')

    def do_init(self, line):
        self.tprint('Going to install VPN server backed by Enigma Bridge FIPS140-2 encryption service.\n')

        # EJBCA installation
        init_res = Installer.do_init(self, line)
        return init_res

    def get_args_intro(self, parser):
        """
        Argument parser intro text
        :return:
        """
        parser.description = 'EnigmaBridge Private Space installer'

    def init_main_try(self):
        """
        Main installer block, called from the global try:
        :return:
        """
        self.init_services()
        self.ejbca.do_vpn = True
        self.ovpn = openvpn.OpenVpn(sysconfig=self.syscfg)

        # Get registration options and choose one - network call.
        self.reg_svc.load_auth_types()

        # Show email prompt and intro text only for new initializations.
        res = self.init_prompt_user()
        if res != 0:
            self.return_code(res)

        # System check proceeds (mem, network).
        # We do this even if we continue with previous registration, to have fresh view on the system.
        # Check if we have EJBCA resources on the drive
        res = self.init_test_environment()
        if res != 0:
            self.return_code(res)

        # Determine if we have enough RAM for the work.
        # If not, a new swap file is created so the system has at least 2GB total memory space
        # for compilation & deployment.
        res = self.install_check_memory(syscfg=self.syscfg)
        if res != 0:
            return self.return_code(res)

        # Preferred LE method? If set...
        self.last_is_vpc = False

        # Lets encrypt reachability test, if preferred method is DNS - do only one attempt.
        # We test this to detect VPC also. If 443 is reachable, we are not in VPC
        res, args_le_preferred_method = self.init_le_vpc_check(self.get_args_le_verification(),
                                                               self.get_args_vpc(), reg_svc=self.reg_svc)
        if res != 0:
            return self.return_code(res)

        # User registration may be multi-step process.
        res, new_config = self.init_enigma_registration()
        if res != 0:
            return self.return_code(res)

        # Custom hostname for EJBCA - not yet supported
        new_config.ejbca_hostname_custom = False
        new_config.is_private_network = self.last_is_vpc
        new_config.le_preferred_verification = args_le_preferred_method

        # Assign a new dynamic domain for the host
        res, self.domain_is_ok = self.init_domains_check(reg_svc=self.reg_svc)
        new_config = self.reg_svc.config
        if res != 0:
            return self.return_code(res)

        # Install to the OS - cron job & on boot service
        res = self.init_install_os_hooks()
        if res != 0:
            return self.return_code(res)

        # Dump config & SoftHSM
        conf_file = Core.write_configuration(new_config)
        self.tprint('New configuration was written to: %s\n' % conf_file)

        # SoftHSMv1 reconfigure
        res = self.init_softhsm(new_config=new_config)
        if res != 0:
            return self.return_code(res)

        # EJBCA configuration
        res = self.init_install_ejbca(new_config=new_config)
        if res != 0:
            return self.return_code(res)

        # Generate new keys
        res = self.init_create_vpn_eb_keys()
        if res != 0:
            return self.return_code(res)

        # JBoss restart is needed - so it sees the new keys
        self.ejbca.jboss_restart()

        # VPN setup
        self.ejbca.vpn_create_ca()
        self.ejbca.vpn_create_profiles()
        self.ejbca.vpn_create_server_certs()
        self.ejbca.vpn_create_crl()
        vpn_ca, vpn_cert, vpn_key = self.ejbca.vpn_get_server_cert_paths()

        # VPN server
        self.tprint('Installing & configuring VPN server')
        self.ovpn.install()
        self.ovpn.generate_dh_group()
        self.ovpn.configure_server()
        self.ovpn.store_server_cert(ca=vpn_ca, cert=vpn_cert, key=vpn_key)

        # VPN CRL
        crl_path = self.ejbca.vpn_get_crl_path()
        self.ovpn.configure_crl(crl_path=crl_path)

        # Starting VPN server
        self.ovpn.enable()
        self.ovpn.switch(start=True)
        self.ejbca.vpn_install_cron()

        # LetsEncrypt enrollment
        res = self.init_le_install()
        if res != 0:
            return self.return_code(res)

        self.tprint('')
        self.init_celebrate()
        self.cli_sleep(3)
        self.cli_separator()

        # Finalize, P12 file & final instructions
        new_p12 = self.ejbca.copy_p12_file()
        self.init_show_p12_info(new_p12=new_p12, new_config=new_config)

        # Generate VPN client for the admin
        self.ejbca.vpn_create_user(self.config.email, 'default')

        # Test if main admin port of EJBCA is reachable.
        self.init_test_admin_port_reachability()

        self.cli_sleep(5)
        return self.return_code(0)

    def init_create_vpn_eb_keys(self):
        """
        Creates a new keys in the SoftHSM token -> EB.
        :return:
        """
        self.tprint('\nEnigma Bridge service will generate new keys:')
        ret, out, err = self.ejbca.pkcs11_generate_default_key_set(softhsm=self.soft_config)

        if ret != 0:
            self.tprint(self.t.red('\nError generating new keys'))
            self.tprint('The installation has to be repeated later')

            self.tprint('\nError from the command:')
            self.tprint(''.join(out))
            self.tprint('\n')
            self.tprint(''.join(err))
            return 1
        return 0

def main():
    app = VpnInstaller()
    app.app_main()

if __name__ == '__main__':
    main()

<commit_msg>minor fmting fix<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

from cmd2 import Cmd
import argparse
import sys
import os
import math
import types
import traceback
import pid
import time
import util
import errors
import textwrap
import openvpn
from blessed import Terminal
from consts import *
from core import Core
from config import Config, EBSettings
from registration import Registration, InfoLoader
from softhsm import SoftHsmV1Config
from ejbca import Ejbca
from ebsysconfig import SysConfig
from letsencrypt import LetsEncrypt
from ebclient.registration import ENVIRONMENT_PRODUCTION, ENVIRONMENT_DEVELOPMENT, ENVIRONMENT_TEST
from pkg_resources import get_distribution, DistributionNotFound
from cli import Installer
import logging
import coloredlogs

logger = logging.getLogger(__name__)
coloredlogs.install(level=logging.ERROR)

class VpnInstaller(Installer):
    """
    Extended installer - with VPN.
    """

    def __init__(self, *args, **kwargs):
        """
        Init core
        :param args:
        :param kwargs:
        :return:
        """
        Installer.__init__(self, *args, **kwargs)
        self.ovpn = None

    def init_argparse(self):
        """
        Adding new VPN related arguments
        :return:
        """
        parser = Installer.init_argparse(self)
        return parser

    def ask_for_email_reason(self, is_required=None):
        """
        Reason why we need email - required in VPN case.
        :param is_required:
        :return:
        """
        if is_required:
            self.tprint('We need your email address for:\n'
                        '   a) identity verification for EnigmaBridge account \n'
                        '   b) LetsEncrypt certificate registration\n'
                        '   c) PKI setup - VPN configuration')
            self.tprint('We will send you a verification email.')
            self.tprint('Without a valid e-mail address you won\'t be able to continue with the installation\n')
        else:
            raise ValueError('Email is required in VPN case')

    def do_init(self, line):
        self.tprint('Going to install VPN server backed by Enigma Bridge FIPS140-2 encryption service.\n')

        # EJBCA installation
        init_res = Installer.do_init(self, line)
        return init_res

    def get_args_intro(self, parser):
        """
        Argument parser intro text
        :return:
        """
        parser.description = 'EnigmaBridge Private Space installer'

    def init_main_try(self):
        """
        Main installer block, called from the global try:
        :return:
        """
        self.init_services()
        self.ejbca.do_vpn = True
        self.ovpn = openvpn.OpenVpn(sysconfig=self.syscfg)

        # Get registration options and choose one - network call.
        self.reg_svc.load_auth_types()

        # Show email prompt and intro text only for new initializations.
        res = self.init_prompt_user()
        if res != 0:
            self.return_code(res)

        # System check proceeds (mem, network).
        # We do this even if we continue with previous registration, to have fresh view on the system.
        # Check if we have EJBCA resources on the drive
        res = self.init_test_environment()
        if res != 0:
            self.return_code(res)

        # Determine if we have enough RAM for the work.
        # If not, a new swap file is created so the system has at least 2GB total memory space
        # for compilation & deployment.
        res = self.install_check_memory(syscfg=self.syscfg)
        if res != 0:
            return self.return_code(res)

        # Preferred LE method? If set...
        self.last_is_vpc = False

        # Lets encrypt reachability test, if preferred method is DNS - do only one attempt.
        # We test this to detect VPC also. If 443 is reachable, we are not in VPC
        res, args_le_preferred_method = self.init_le_vpc_check(self.get_args_le_verification(),
                                                               self.get_args_vpc(), reg_svc=self.reg_svc)
        if res != 0:
            return self.return_code(res)

        # User registration may be multi-step process.
        res, new_config = self.init_enigma_registration()
        if res != 0:
            return self.return_code(res)

        # Custom hostname for EJBCA - not yet supported
        new_config.ejbca_hostname_custom = False
        new_config.is_private_network = self.last_is_vpc
        new_config.le_preferred_verification = args_le_preferred_method

        # Assign a new dynamic domain for the host
        res, self.domain_is_ok = self.init_domains_check(reg_svc=self.reg_svc)
        new_config = self.reg_svc.config
        if res != 0:
            return self.return_code(res)

        # Install to the OS - cron job & on boot service
        res = self.init_install_os_hooks()
        if res != 0:
            return self.return_code(res)

        # Dump config & SoftHSM
        conf_file = Core.write_configuration(new_config)
        self.tprint('New configuration was written to: %s\n' % conf_file)

        # SoftHSMv1 reconfigure
        res = self.init_softhsm(new_config=new_config)
        if res != 0:
            return self.return_code(res)

        # EJBCA configuration
        res = self.init_install_ejbca(new_config=new_config)
        if res != 0:
            return self.return_code(res)

        # Generate new keys
        res = self.init_create_vpn_eb_keys()
        if res != 0:
            return self.return_code(res)

        # JBoss restart is needed - so it sees the new keys
        self.ejbca.jboss_restart()

        # VPN setup
        self.ejbca.vpn_create_ca()
        self.ejbca.vpn_create_profiles()
        self.ejbca.vpn_create_server_certs()
        self.ejbca.vpn_create_crl()
        vpn_ca, vpn_cert, vpn_key = self.ejbca.vpn_get_server_cert_paths()
        crl_path = self.ejbca.vpn_get_crl_path()

        # VPN server
        self.tprint('Installing & configuring VPN server')
        self.ovpn.install()
        self.ovpn.generate_dh_group()
        self.ovpn.configure_server()
        self.ovpn.store_server_cert(ca=vpn_ca, cert=vpn_cert, key=vpn_key)
        self.ovpn.configure_crl(crl_path=crl_path)

        # Starting VPN server
        self.ovpn.enable()
        self.ovpn.switch(start=True)
        self.ejbca.vpn_install_cron()

        # LetsEncrypt enrollment
        res = self.init_le_install()
        if res != 0:
            return self.return_code(res)

        self.tprint('')
        self.init_celebrate()
        self.cli_sleep(3)
        self.cli_separator()

        # Finalize, P12 file & final instructions
        new_p12 = self.ejbca.copy_p12_file()
        self.init_show_p12_info(new_p12=new_p12, new_config=new_config)

        # Generate VPN client for the admin
        self.ejbca.vpn_create_user(self.config.email, 'default')

        # Test if main admin port of EJBCA is reachable.
        self.init_test_admin_port_reachability()

        self.cli_sleep(5)
        return self.return_code(0)

    def init_create_vpn_eb_keys(self):
        """
        Creates a new keys in the SoftHSM token -> EB.
        :return:
        """
        self.tprint('\nEnigma Bridge service will generate new keys:')
        ret, out, err = self.ejbca.pkcs11_generate_default_key_set(softhsm=self.soft_config)

        if ret != 0:
            self.tprint(self.t.red('\nError generating new keys'))
            self.tprint('The installation has to be repeated later')

            self.tprint('\nError from the command:')
            self.tprint(''.join(out))
            self.tprint('\n')
            self.tprint(''.join(err))
            return 1
        return 0

def main():
    app = VpnInstaller()
    app.app_main()

if __name__ == '__main__':
    main()

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

# Exercício de avaliação de expressão aritmética.
# Só podem ser usadas as estruturas Pilha e Fila implementadas em aulas anteriores.
# Deve ter análise de tempo e espaço para função avaliação

from aula5.fila import Fila
from aula4.pilha import Pilha

class ErroLexico(Exception):
    pass

class ErroSintatico(Exception):
    pass

def analise_lexica(expressao):
    """
    Executa análise lexica transformando a expressao em fila de objetos:
    Transforma inteiros em ints
    Flutuantes em floats
    e verificar se demais caracteres são validos: +-*/(){}[]
    :param expressao: string com expressao a ser analisada
    :return: fila com tokens
    """
    fila = Fila()

    #numeros e tokens
    nt = R"0123456789.-+*/{}[]()"

    if expressao:

        valor = ''

        for i in expressao:
            if i in nt:
                if i in '.-+*/{}[]()':
                    if valor:
                        fila.enfileirar(valor)
                        valor = ''
                    fila.enfileirar(i)
                else:
                    valor = valor + i
            else:
                raise ErroLexico()

        if valor:
            fila.enfileirar(valor)

    return fila

def analise_sintatica(fila):
    """
    Função que realiza analise sintática de tokens produzidos por analise léxica.
    Executa validações sintáticas e se não houver erro retorn fila_sintatica para avaliacao

    :param fila: fila proveniente de análise lexica
    :return: fila_sintatica com elementos tokens de numeros
    """

    if fila.__len__():

        # Cria um novo objeto fila para adicionar os valores
        fila_sintetica = Fila()

        #Variavel de apoio para juntar a string
        valor = ''

        for i in range(fila.__len__()):

            if fila._deque[i] in '-+/*(){}[]':
                if valor:
                    if '.' in valor:
                        fila_sintetica.enfileirar(float(valor))
                    else:
                        fila_sintetica.enfileirar(int(valor))

                valor = ''
                fila_sintetica.enfileirar(fila._deque[i])
            else:
                valor = valor + fila._deque[i]

        if valor:
            if '.' in valor:
                fila_sintetica.enfileirar(float(valor))
            else:
                fila_sintetica.enfileirar(int(valor))

        return fila_sintetica
    else:
        raise ErroSintatico

def avaliar(expressao):
    """
    Função que avalia expressão aritmetica retornando se valor se não houver nenhum erro
    :param expressao: string com expressão aritmética
    :return: valor númerico com resultado

    Complexidade*:

    tempo: O(n)
    Memoria: O(n)
    """

    if expressao:

        fila = analise_sintatica(analise_lexica(expressao))

        # print(fila._deque)

        if fila.__len__() == 1:
            return fila.primeiro()
        else:
            pilha = Pilha()

            for i in range(fila.__len__()):

                pilha.empilhar(fila._deque[i])

                if pilha.__len__() >= 3 and str(pilha.topo()) not in '-+*/(){}[]':

                    valor = pilha.topo()
                    pilha.desempilhar()

                    if pilha.topo() == '+':
                        pilha.desempilhar()
                        valor = pilha.desempilhar() + valor
                        pilha.empilhar(valor)
                        valor = ''
                    elif pilha.topo() == '-':
                        pilha.desempilhar()
                        valor = pilha.desempilhar() - valor
                        pilha.empilhar(valor)
                        valor = ''
                    elif pilha.topo() == '*':
                        pilha.desempilhar()
                        valor = pilha.desempilhar() * valor
                        pilha.empilhar(valor)
                        valor = ''
                    elif pilha.topo() == '/':
                        pilha.desempilhar()
                        valor = pilha.desempilhar() / valor
                        pilha.empilhar(valor)
                        valor = ''
                    else:
                        pilha.empilhar(valor)

                elif str(pilha.topo()) in ')}]' and i == fila.__len__() - 1:
                    pilha.desempilhar()

                    # print(pilha._lista)

                    while len(pilha) > 1:

                        if str(pilha.topo()) not in '-+*/(){}[]':
                            valor = pilha.topo()
                            pilha.desempilhar()

                            if pilha.topo() == '+':
                                pilha.desempilhar()
                                valor = pilha.desempilhar() + valor
                                pilha.empilhar(valor)
                                valor = ''
                            elif pilha.topo() == '-':
                                pilha.desempilhar()
                                valor = pilha.desempilhar() - valor
                                pilha.empilhar(valor)
                                valor = ''
                            elif pilha.topo() == '*':
                                pilha.desempilhar()
                                valor = pilha.desempilhar() * valor
                                pilha.empilhar(valor)
                                valor = ''
                            elif pilha.topo() == '/':
                                pilha.desempilhar()
                                valor = pilha.desempilhar() / valor
                                pilha.empilhar(valor)
                                valor = ''
                            elif str(pilha.topo()) in '(){}[]':
                                pilha.desempilhar()
                                pilha.empilhar(valor)
                            else:
                                pilha.empilhar(valor)
                        else:
                            pilha.desempilhar()

            # print('Nova pilha: ',pilha._lista, 'Tamanho: ', len(pilha))

            return pilha.topo()

    raise ErroSintatico()

import unittest

class AnaliseLexicaTestes(unittest.TestCase):
    def test_expressao_vazia(self):
        fila = analise_lexica('')
        self.assertTrue(fila.vazia())

    def test_caracter_estranho(self):
        self.assertRaises(ErroLexico, analise_lexica, 'a')
        self.assertRaises(ErroLexico, analise_lexica, 'ab')

    def test_inteiro_com_um_algarismo(self):
        fila = analise_lexica('1')
        self.assertEqual('1', fila.desenfileirar())
        self.assertTrue(fila.vazia())

    def test_inteiro_com_vários_algarismos(self):
        fila = analise_lexica('1234567890')
        self.assertEqual('1234567890', fila.desenfileirar())
        self.assertTrue(fila.vazia())

    def test_float(self):
        fila = analise_lexica('1234567890.34')
        self.assertEqual('1234567890', fila.desenfileirar())
        self.assertEqual('.', fila.desenfileirar())
        self.assertEqual('34', fila.desenfileirar())
        self.assertTrue(fila.vazia())

    def test_parenteses(self):
        fila = analise_lexica('(1)')
        self.assertEqual('(', fila.desenfileirar())
        self.assertEqual('1', fila.desenfileirar())
        self.assertEqual(')', fila.desenfileirar())
        self.assertTrue(fila.vazia())

    def test_chaves(self):
        fila = analise_lexica('{(1)}')
        self.assertEqual('{', fila.desenfileirar())
        self.assertEqual('(', fila.desenfileirar())
        self.assertEqual('1', fila.desenfileirar())
        self.assertEqual(')', fila.desenfileirar())
        self.assertEqual('}', fila.desenfileirar())
        self.assertTrue(fila.vazia())

    def test_colchetes(self):
        fila = analise_lexica('[{(1.0)}]')
        self.assertEqual('[', fila.desenfileirar())
        self.assertEqual('{', fila.desenfileirar())
        self.assertEqual('(', fila.desenfileirar())
        self.assertEqual('1', fila.desenfileirar())
        self.assertEqual('.', fila.desenfileirar())
        self.assertEqual('0', fila.desenfileirar())
        self.assertEqual(')', fila.desenfileirar())
        self.assertEqual('}', fila.desenfileirar())
        self.assertEqual(']', fila.desenfileirar())
        self.assertTrue(fila.vazia())

    def test_adicao(self):
        fila = analise_lexica('1+2.0')
        self.assertEqual('1', fila.desenfileirar())
        self.assertEqual('+', fila.desenfileirar())
        self.assertEqual('2', fila.desenfileirar())
        self.assertEqual('.', fila.desenfileirar())
        self.assertEqual('0', fila.desenfileirar())
        self.assertTrue(fila.vazia())

    def test_subtracao(self):
        fila = analise_lexica('1-2.0')
        self.assertEqual('1', fila.desenfileirar())
        self.assertEqual('-', fila.desenfileirar())
        self.assertEqual('2', fila.desenfileirar())
        self.assertEqual('.', fila.desenfileirar())
        self.assertEqual('0', fila.desenfileirar())
        self.assertTrue(fila.vazia())

    def test_multiplicacao(self):
        fila = analise_lexica('1*2.0')
        self.assertEqual('1', fila.desenfileirar())
        self.assertEqual('*', fila.desenfileirar())
        self.assertEqual('2', fila.desenfileirar())
        self.assertEqual('.', fila.desenfileirar())
        self.assertEqual('0', fila.desenfileirar())
        self.assertTrue(fila.vazia())

    def test_divisao(self):
        fila = analise_lexica('1/2.0')
        self.assertEqual('1', fila.desenfileirar())
        self.assertEqual('/', fila.desenfileirar())
        self.assertEqual('2', fila.desenfileirar())
        self.assertEqual('.', fila.desenfileirar())
        self.assertEqual('0', fila.desenfileirar())
        self.assertTrue(fila.vazia())

    def test_expresao_com_todos_simbolos(self):
        expressao = '1/{2.0+3*[7-(5-3)]}'
        fila = analise_lexica(expressao)
        self.assertListEqual(list(expressao), [e for e in fila])
        self.assertTrue(fila.vazia())

class AnaliseSintaticaTestes(unittest.TestCase):
    def test_fila_vazia(self):
        fila = Fila()
        self.assertRaises(ErroSintatico, analise_sintatica, fila)

    def test_int(self):
        fila = Fila()
        fila.enfileirar('1234567890')
        fila_sintatica = analise_sintatica(fila)
        self.assertEqual(1234567890, fila_sintatica.desenfileirar())
        self.assertTrue(fila_sintatica.vazia())

    def test_float(self):
        fila = Fila()
        fila.enfileirar('1234567890')
        fila.enfileirar('.')
        fila.enfileirar('4')
        fila_sintatica = analise_sintatica(fila)
        self.assertEqual(1234567890.4, fila_sintatica.desenfileirar())
        self.assertTrue(fila_sintatica.vazia())

    def test_expressao_com_todos_elementos(self):
        fila = analise_lexica('1000/{222.125+3*[7-(5-3)]}')
        fila_sintatica = analise_sintatica(fila)
        self.assertListEqual([1000, '/', '{', 222.125, '+', 3, '*', '[', 7, '-', '(', 5, '-', 3, ')', ']', '}'],[e for e in fila_sintatica])

class AvaliacaoTestes(unittest.TestCase):
    def test_expressao_vazia(self):
        self.assertRaises(ErroSintatico, avaliar, '')

    def test_inteiro(self):
        self.assert_avaliacao('1')

    def test_float(self):
        self.assert_avaliacao('2.1')

    def test_soma(self):
        self.assert_avaliacao('2+1')

    def test_subtracao_e_parenteses(self):
        self.assert_avaliacao('(2-1)')

    def test_expressao_com_todos_elementos(self):
        self.assertEqual(1.0, avaliar('2.0/[4*3+1-{15-(1+3)}]'))

    def assert_avaliacao(self, expressao):
        self.assertEqual(eval(expressao), avaliar(expressao))

if __name__ == '__main__':
    unittest.main()
<commit_msg> Exe. avaliação de expressão aritmética finalizado<commit_after># -*- coding: utf-8 -*-

# Exercício de avaliação de expressão aritmética.
# Só podem ser usadas as estruturas Pilha e Fila implementadas em aulas anteriores.
# Deve ter análise de tempo e espaço para função avaliação

from aula4.pilha import Pilha
from aula5.fila import Fila

class ErroLexico(Exception):
    pass

class ErroSintatico(Exception):
    pass

def analise_lexica(expressao):
    """
    Executa análise lexica transformando a expressao em fila de objetos:
    Transforma inteiros em ints
    Flutuantes em floats
    e verificar se demais caracteres são validos: +-*/(){}[]
    :param expressao: string com expressao a ser analisada
    :return: fila com tokens
    """
    fila = Fila()

    #numeros e tokens
    nt = R"0123456789.-+*/{}[]()"

    if expressao:

        valor = ''

        for i in expressao:
            if i in nt:
                if i in '.-+*/{}[]()':
                    if valor:
                        fila.enfileirar(valor)
                        valor = ''
                    fila.enfileirar(i)
                else:
                    valor = valor + i
            else:
                raise ErroLexico()

        if valor:
            fila.enfileirar(valor)

    return fila

def analise_sintatica(fila):
    """
    Função que realiza analise sintática de tokens produzidos por analise léxica.
    Executa validações sintáticas e se não houver erro retorn fila_sintatica para avaliacao

    :param fila: fila proveniente de análise lexica
    :return: fila_sintatica com elementos tokens de numeros
    """

    if fila.__len__():

        # Cria um novo objeto fila para adicionar os valores
        fila_sintetica = Fila()

        #Variavel de apoio para juntar a string
        valor = ''

        for i in range(fila.__len__()):

            if fila._deque[i] in '-+/*(){}[]':
                if valor:
                    if '.' in valor:
                        fila_sintetica.enfileirar(float(valor))
                    else:
                        fila_sintetica.enfileirar(int(valor))

                valor = ''
                fila_sintetica.enfileirar(fila._deque[i])
            else:
                valor = valor + fila._deque[i]

        if valor:
            if '.' in valor:
                fila_sintetica.enfileirar(float(valor))
            else:
                fila_sintetica.enfileirar(int(valor))

        return fila_sintetica
    else:
        raise ErroSintatico

def avaliar(expressao):
    """
    Função que avalia expressão aritmetica retornando se valor se não houver nenhum erro
    :param expressao: string com expressão aritmética
    :return: valor númerico com resultado

    Complexidade*:

    tempo: O(n)
    Memoria: O(n)
    """

    if expressao:

        fila = analise_sintatica(analise_lexica(expressao))

        # print(fila._deque)

        if fila.__len__() == 1:
            return fila.primeiro()
        else:
            pilha = Pilha()

            for i in range(fila.__len__()):

                pilha.empilhar(fila._deque[i])

                if pilha.__len__() >= 3 and str(pilha.topo()) not in '-+*/(){}[]':

                    valor = pilha.topo()
                    pilha.desempilhar()

                    if pilha.topo() == '+':
                        pilha.desempilhar()
                        valor = pilha.desempilhar() + valor
                        pilha.empilhar(valor)
                        valor = ''
                    elif pilha.topo() == '-':
                        pilha.desempilhar()
                        valor = pilha.desempilhar() - valor
                        pilha.empilhar(valor)
                        valor = ''
                    elif pilha.topo() == '*':
                        pilha.desempilhar()
                        valor = pilha.desempilhar() * valor
                        pilha.empilhar(valor)
                        valor = ''
                    elif pilha.topo() == '/':
                        pilha.desempilhar()
                        valor = pilha.desempilhar() / valor
                        pilha.empilhar(valor)
                        valor = ''
                    else:
                        pilha.empilhar(valor)

                elif str(pilha.topo()) in ')}]' and i == fila.__len__() - 1:
                    pilha.desempilhar()

                    # print(pilha._lista)

                    while len(pilha) > 1:

                        if str(pilha.topo()) not in '-+*/(){}[]':
                            valor = pilha.topo()
                            pilha.desempilhar()

                            if pilha.topo() == '+':
                                pilha.desempilhar()
                                valor = pilha.desempilhar() + valor
                                pilha.empilhar(valor)
                                valor = ''
                            elif pilha.topo() == '-':
                                pilha.desempilhar()
                                valor = pilha.desempilhar() - valor
                                pilha.empilhar(valor)
                                valor = ''
                            elif pilha.topo() == '*':
                                pilha.desempilhar()
                                valor = pilha.desempilhar() * valor
                                pilha.empilhar(valor)
                                valor = ''
                            elif pilha.topo() == '/':
                                pilha.desempilhar()
                                valor = pilha.desempilhar() / valor
                                pilha.empilhar(valor)
                                valor = ''
                            elif str(pilha.topo()) in '(){}[]':
                                pilha.desempilhar()
                                pilha.empilhar(valor)
                            else:
                                pilha.empilhar(valor)
                        else:
                            pilha.desempilhar()

            # print('Nova pilha: ',pilha._lista, 'Tamanho: ', len(pilha))

            return pilha.topo()

    raise ErroSintatico()

import unittest

class AnaliseLexicaTestes(unittest.TestCase):
    def test_expressao_vazia(self):
        fila = analise_lexica('')
        self.assertTrue(fila.vazia())

    def test_caracter_estranho(self):
        self.assertRaises(ErroLexico, analise_lexica, 'a')
        self.assertRaises(ErroLexico, analise_lexica, 'ab')

    def test_inteiro_com_um_algarismo(self):
        fila = analise_lexica('1')
        self.assertEqual('1', fila.desenfileirar())
        self.assertTrue(fila.vazia())

    def test_inteiro_com_vários_algarismos(self):
        fila = analise_lexica('1234567890')
        self.assertEqual('1234567890', fila.desenfileirar())
        self.assertTrue(fila.vazia())

    def test_float(self):
        fila = analise_lexica('1234567890.34')
        self.assertEqual('1234567890', fila.desenfileirar())
        self.assertEqual('.', fila.desenfileirar())
        self.assertEqual('34', fila.desenfileirar())
        self.assertTrue(fila.vazia())

    def test_parenteses(self):
        fila = analise_lexica('(1)')
        self.assertEqual('(', fila.desenfileirar())
        self.assertEqual('1', fila.desenfileirar())
        self.assertEqual(')', fila.desenfileirar())
        self.assertTrue(fila.vazia())

    def test_chaves(self):
        fila = analise_lexica('{(1)}')
        self.assertEqual('{', fila.desenfileirar())
        self.assertEqual('(', fila.desenfileirar())
        self.assertEqual('1', fila.desenfileirar())
        self.assertEqual(')', fila.desenfileirar())
        self.assertEqual('}', fila.desenfileirar())
        self.assertTrue(fila.vazia())

    def test_colchetes(self):
        fila = analise_lexica('[{(1.0)}]')
        self.assertEqual('[', fila.desenfileirar())
        self.assertEqual('{', fila.desenfileirar())
        self.assertEqual('(', fila.desenfileirar())
        self.assertEqual('1', fila.desenfileirar())
        self.assertEqual('.', fila.desenfileirar())
        self.assertEqual('0', fila.desenfileirar())
        self.assertEqual(')', fila.desenfileirar())
        self.assertEqual('}', fila.desenfileirar())
        self.assertEqual(']', fila.desenfileirar())
        self.assertTrue(fila.vazia())

    def test_adicao(self):
        fila = analise_lexica('1+2.0')
        self.assertEqual('1', fila.desenfileirar())
        self.assertEqual('+', fila.desenfileirar())
        self.assertEqual('2', fila.desenfileirar())
        self.assertEqual('.', fila.desenfileirar())
        self.assertEqual('0', fila.desenfileirar())
        self.assertTrue(fila.vazia())

    def test_subtracao(self):
        fila = analise_lexica('1-2.0')
        self.assertEqual('1', fila.desenfileirar())
        self.assertEqual('-', fila.desenfileirar())
        self.assertEqual('2', fila.desenfileirar())
        self.assertEqual('.', fila.desenfileirar())
        self.assertEqual('0', fila.desenfileirar())
        self.assertTrue(fila.vazia())

    def test_multiplicacao(self):
        fila = analise_lexica('1*2.0')
        self.assertEqual('1', fila.desenfileirar())
        self.assertEqual('*', fila.desenfileirar())
        self.assertEqual('2', fila.desenfileirar())
        self.assertEqual('.', fila.desenfileirar())
        self.assertEqual('0', fila.desenfileirar())
        self.assertTrue(fila.vazia())

    def test_divisao(self):
        fila = analise_lexica('1/2.0')
        self.assertEqual('1', fila.desenfileirar())
        self.assertEqual('/', fila.desenfileirar())
        self.assertEqual('2', fila.desenfileirar())
        self.assertEqual('.', fila.desenfileirar())
        self.assertEqual('0', fila.desenfileirar())
        self.assertTrue(fila.vazia())

    def test_expresao_com_todos_simbolos(self):
        expressao = '1/{2.0+3*[7-(5-3)]}'
        fila = analise_lexica(expressao)
        self.assertListEqual(list(expressao), [e for e in fila])
        self.assertTrue(fila.vazia())

class AnaliseSintaticaTestes(unittest.TestCase):
    def test_fila_vazia(self):
        fila = Fila()
        self.assertRaises(ErroSintatico, analise_sintatica, fila)

    def test_int(self):
        fila = Fila()
        fila.enfileirar('1234567890')
        fila_sintatica = analise_sintatica(fila)
        self.assertEqual(1234567890, fila_sintatica.desenfileirar())
        self.assertTrue(fila_sintatica.vazia())

    def test_float(self):
        fila = Fila()
        fila.enfileirar('1234567890')
        fila.enfileirar('.')
        fila.enfileirar('4')
        fila_sintatica = analise_sintatica(fila)
        self.assertEqual(1234567890.4, fila_sintatica.desenfileirar())
        self.assertTrue(fila_sintatica.vazia())

    def test_expressao_com_todos_elementos(self):
        fila = analise_lexica('1000/{222.125+3*[7-(5-3)]}')
        fila_sintatica = analise_sintatica(fila)
        self.assertListEqual([1000, '/', '{', 222.125, '+', 3, '*', '[', 7, '-', '(', 5, '-', 3, ')', ']', '}'],[e for e in fila_sintatica])

class AvaliacaoTestes(unittest.TestCase):
    def test_expressao_vazia(self):
        self.assertRaises(ErroSintatico, avaliar, '')

    def test_inteiro(self):
        self.assert_avaliacao('1')

    def test_float(self):
        self.assert_avaliacao('2.1')

    def test_soma(self):
        self.assert_avaliacao('2+1')

    def test_subtracao_e_parenteses(self):
        self.assert_avaliacao('(2-1)')

    def test_expressao_com_todos_elementos(self):
        self.assertEqual(1.0, avaliar('2.0/[4*3+1-{15-(1+3)}]'))

    def assert_avaliacao(self, expressao):
        self.assertEqual(eval(expressao), avaliar(expressao))

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