<commit_before>############################################################
#
#    Copyright (c) 2013, Southpaw Technology
#                        All Rights Reserved
#
#    PROPRIETARY INFORMATION.  This software is proprietary to
#    Southpaw Technology, and is not to be reproduced, transmitted,
#    or disclosed in any way without written permission.
#
#

import tacticenv

from pyasm.common import Common, TacticException, Config

import os, subprocess, time, re, os

class RSyncConnectionException(Exception):
    pass

class RSync(object):

    def __init__(my, **kwargs):
        my.kwargs = kwargs
        my.paths = []
        my.data = {}

    def get_paths(my):
        return my.paths

    def get_data(my):
        return my.data

    def execute(my):
        base_dir = my.kwargs.get("base_dir")
        relative_dir = my.kwargs.get("relative_dir")
        file_name = my.kwargs.get("file_name")
        to_path = my.kwargs.get("to_path")

        from_path = "%s/./%s/%s" % (base_dir, relative_dir, file_name)

        tries = 1
        success = False
        value = ""
        start = time.time()

        message_key = my.kwargs.get("message_key")

        while 1:
            try:
                value = my.sync_paths(from_path, to_path)
                success = True
                break
            except RSyncConnectionException, e:
                time.sleep(60)
                server.log_message(message_key, {"message": str(e)}, "error_retry")
                continue

            except Exception, e:
                print "Failed on try [%s]..." % tries
                raise

                # ping the server to see

                if message_key:
                    from tactic_client_lib import TacticServerStub
                    server = TacticServerStub.get()
                    server.log_message(message_key, {"message": str(e)}, "error_retry")

                time.sleep(tries)
                if tries == 3:
                    break

                tries += 1

        #print "success: ", success
        #print time.time() - start, " seconds"
        #print value

    def get_current_data(my):
        return my.current_data

    def get_data(my):
        return my.data

    def sync_paths(my, from_path, to_path):

        server = my.kwargs.get("server")
        login = my.kwargs.get("login")
        paths = my.kwargs.get("paths")
        base_dir = my.kwargs.get("base_dir")

        paths_sizes = []
        if paths:
            for path in paths:
                full_path = "%s/%s" % (base_dir, path)
                print "full_path: ", full_path
                size = os.path.getsize(full_path)
                paths_sizes.append(size)

        if server:
            to_path = "%s@%s:%s" % (login, server, to_path)

        if os.name == "nt":
            # This assumes we are using the cygwin implementation of RSync
            if to_path[1] == ":":
                # then this is a drive letter
                drive_letter = to_path[0]
                to_path = "/cygdrive/%s/%s" % (drive_letter, to_path[3:])
            if from_path[1] == ":":
                # then this is a drive letter
                drive_letter = from_path[0]
                from_path = "/cygdrive/%s/%s" % (drive_letter, from_path[3:])

            to_path = to_path.replace("\\", '/')
            from_path = from_path.replace("\\", '/')

        rsync = Config.get_value("services", "rsync")
        if not rsync:
            rsync = Common.which("rsync")

        if not rsync:
            raise Exception("RSync executable could not be found")

        cmd_list = []
        cmd_list.append(rsync)
        flags = "-az"
        cmd_list.append(flags)
        if server:
            cmd_list.append("-e ssh")
        cmd_list.append("-v")
        cmd_list.append("--progress")

        delete = False
        if delete:
            cmd_list.append("--delete")

        dry_run = my.kwargs.get("dry_run")
        if dry_run in [True, 'true']:
            cmd_list.append("--dry-run")

        relative = True
        if relative:
            cmd_list.append("--relative")

        partial = True
        if partial:
            cmd_list.append("--partial")

        cmd_list.append('%s' % from_path)
        cmd_list.append('%s' % to_path)

        print "exec: ", " ".join(cmd_list)

        program = subprocess.Popen(cmd_list, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True)
        #program.wait()

        message_key = my.kwargs.get("message_key")

        progress = RSyncProgress(
                message_key=message_key,
                paths=paths,
                paths_sizes=paths_sizes
        )

        on_update = my.kwargs.get("on_update")
        if not on_update:
            on_update = progress.on_update
        assert(on_update)

        on_complete = my.kwargs.get("on_complete")
        if not on_complete:
            on_complete = progress.on_complete
        assert(on_complete)

        on_error = my.kwargs.get("on_error")
        if not on_error:
            on_error = progress.on_error
        assert(on_error)

        data = []
        lines = []
        path = None
        my.paths = []
        error = []
        #while program.poll() is None:
        buffer = []
        line = ""
        while 1:
            char = program.stdout.read(1)
            if not char:
                break

            if char == "\n":
                line = "".join(buffer)

            else:
                buffer.append(char)

            if line:
                lines.append(line)
                if line.startswith(" "):

                    line = line.strip()
                    parts = re.split( re.compile("\ +"), line )

                    my.current_data = {
                        "bytes": int(parts[0]),
                        "percent": parts[1],
                        "rate": parts[2],
                        "time_left": parts[3]
                    }
                    #print my.current_data
                    #print "status: ", line

                    if on_update:
                        on_update(path, my.current_data)

                elif line.startswith("rsync "):
                    error.append(line)
                elif line.startswith("rsync: "):
                    error.append(line)
                elif line.startswith("deleting "):
                    my.handle_data_line(line)
                elif line.startswith("sent "):
                    my.handle_data_line(line)
                elif line.startswith("total "):
                    my.handle_data_line(line)
                elif line.startswith("sending incremental "):
                    my.handle_data_line(line)
                elif line.startswith("created directory "):
                    my.handle_data_line(line)
                else:
                    line = line.strip()
                    path = line
                    if not line.endswith("/"):
                        if line.find(" -> ") != -1:
                            parts = line.split(" -> ")
                            line = parts[0]

                        my.paths.append(line)

                # reset the line
                line = None
                buffer = []

        if error:
            error = "\n".join(error)
            if error.startswith("rsync: connection unexpectedly closed"):
                raise RSyncConnectionException(error)

            if on_error:
                on_error(path, {"message": error})
            raise Exception("Sync Error\n%s" % error)

        if on_complete:
            on_complete(path, {"message": "All files synchronized"})

        return "success"

    def handle_data_line(my, line):

        if line.startswith("sent "):
            # sent 520 bytes  received 22 bytes  361.33 bytes/sec
            pass

        elif line.startswith("total "):
            #total size is 431666  speedup is 796.43 (DRY RUN)
            parts = line.split()
            total_size = parts[3]
            total_size = int(total_size)
            my.data['total_size'] = total_size

class RSyncProgress(object):
    def __init__(my, **kwargs):
        from tactic_client_lib import TacticServerStub
        my.server = TacticServerStub.get()
        my.message_key = kwargs.get("message_key")
        my.paths = kwargs.get("paths")
        my.paths_sizes = kwargs.get("paths_sizes")
        if not my.paths:
            my.paths = []

    def on_update(my, path, data):
        data['path'] = path

        bytes = data.get("bytes")

        try:
            index = my.paths.index(path)
        except:
            index = 0
        data["path_index"] = index+1
        data["paths_count"] = len(my.paths)
        data["paths_sizes"] = my.paths_sizes

        total_size = 0
        current_size = 0
        for i in range(0, len(my.paths)):
            if i < index:
                current_size += my.paths_sizes[i]
            elif i == index:
                current_size += bytes

            total_size += my.paths_sizes[i]

        #print "current_size: ", current_size
        #print "total_size: ", total_size
        #print "percent: ", (current_size*1000)/(total_size*10)
        if total_size:
            total_percent = (current_size*1000)/(total_size*10)
        else:
            total_percent = 0
        data['total_percent'] = total_percent

        if my.message_key:
            my.server.log_message(my.message_key, data, status="in_progress")

    def on_error(my, path, data):
        if my.message_key:
            my.server.log_message(my.message_key, data, status="error")

    def on_complete(my, path, data):
        if my.message_key:
            my.server.log_message(my.message_key, data, status="complete")

if __name__ == '__main__':

    import time

    from_path = "/home/tactic/svg"

    start = time.time()

    dir_info = Common.get_dir_info(from_path)
    print "dir_info: ", dir_info
    total_size = dir_info.get("size")

    from pyasm.security import Batch
    Batch()

    class Progress(object):
        def __init__(my):
            my.total_sent = 0

        def on_update(my, path, data):
            from tactic_client_lib import TacticServerStub
            server = TacticServerStub.get()
            data['path'] = path

            bytes = data.get("bytes")
            my.total_sent += bytes

            print "path: ", path
            print "total: ", my.total_sent

            server.log_message("wow", data)
    progress = Progress()

    kwargs = {
            "login": "root",
            "server": "sync1.southpawtech.com",
            "from_path": from_path,
            "to_path": "/spt/test/svg",
            "on_update": progress.on_update
    }
    cmd = RSync(**kwargs)

    cmd.execute()

<commit_msg>BRANCH 4.2-105: fixed undefined variable<commit_after>############################################################
#
#    Copyright (c) 2013, Southpaw Technology
#                        All Rights Reserved
#
#    PROPRIETARY INFORMATION.  This software is proprietary to
#    Southpaw Technology, and is not to be reproduced, transmitted,
#    or disclosed in any way without written permission.
#
#

import tacticenv

from pyasm.common import Common, TacticException, Config

import os, subprocess, time, re, os

class RSyncConnectionException(Exception):
    pass

class RSync(object):

    def __init__(my, **kwargs):
        my.kwargs = kwargs
        my.paths = []
        my.data = {}

    def get_paths(my):
        return my.paths

    def get_data(my):
        return my.data

    def execute(my):
        base_dir = my.kwargs.get("base_dir")
        relative_dir = my.kwargs.get("relative_dir")
        file_name = my.kwargs.get("file_name")
        to_path = my.kwargs.get("to_path")

        from_path = "%s/./%s/%s" % (base_dir, relative_dir, file_name)

        tries = 1
        success = False
        value = ""
        start = time.time()

        message_key = my.kwargs.get("message_key")

        while 1:
            try:
                value = my.sync_paths(from_path, to_path)
                success = True
                break
            except RSyncConnectionException, e:
                time.sleep(60)

                from tactic_client_lib import TacticServerStub
                server = TacticServerStub.get()
                server.log_message(message_key, {"message": str(e)}, "error_retry")
                continue

            except Exception, e:
                print "Failed on try [%s]..." % tries
                raise

                # ping the server to see

                if message_key:
                    from tactic_client_lib import TacticServerStub
                    server = TacticServerStub.get()
                    server.log_message(message_key, {"message": str(e)}, "error_retry")

                time.sleep(tries)
                if tries == 3:
                    break

                tries += 1

        #print "success: ", success
        #print time.time() - start, " seconds"
        #print value

    def get_current_data(my):
        return my.current_data

    def get_data(my):
        return my.data

    def sync_paths(my, from_path, to_path):

        server = my.kwargs.get("server")
        login = my.kwargs.get("login")
        paths = my.kwargs.get("paths")
        base_dir = my.kwargs.get("base_dir")

        paths_sizes = []
        if paths:
            for path in paths:
                full_path = "%s/%s" % (base_dir, path)
                print "full_path: ", full_path
                size = os.path.getsize(full_path)
                paths_sizes.append(size)

        if server:
            to_path = "%s@%s:%s" % (login, server, to_path)

        if os.name == "nt":
            # This assumes we are using the cygwin implementation of RSync
            if to_path[1] == ":":
                # then this is a drive letter
                drive_letter = to_path[0]
                to_path = "/cygdrive/%s/%s" % (drive_letter, to_path[3:])
            if from_path[1] == ":":
                # then this is a drive letter
                drive_letter = from_path[0]
                from_path = "/cygdrive/%s/%s" % (drive_letter, from_path[3:])

            to_path = to_path.replace("\\", '/')
            from_path = from_path.replace("\\", '/')

        rsync = Config.get_value("services", "rsync")
        if not rsync:
            rsync = Common.which("rsync")

        if not rsync:
            raise Exception("RSync executable could not be found")

        cmd_list = []
        cmd_list.append(rsync)
        flags = "-az"
        cmd_list.append(flags)
        if server:
            cmd_list.append("-e ssh")
        cmd_list.append("-v")
        cmd_list.append("--progress")

        delete = False
        if delete:
            cmd_list.append("--delete")

        dry_run = my.kwargs.get("dry_run")
        if dry_run in [True, 'true']:
            cmd_list.append("--dry-run")

        relative = True
        if relative:
            cmd_list.append("--relative")

        partial = True
        if partial:
            cmd_list.append("--partial")

        cmd_list.append('%s' % from_path)
        cmd_list.append('%s' % to_path)

        print "exec: ", " ".join(cmd_list)

        program = subprocess.Popen(cmd_list, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True)
        #program.wait()

        message_key = my.kwargs.get("message_key")

        progress = RSyncProgress(
                message_key=message_key,
                paths=paths,
                paths_sizes=paths_sizes
        )

        on_update = my.kwargs.get("on_update")
        if not on_update:
            on_update = progress.on_update
        assert(on_update)

        on_complete = my.kwargs.get("on_complete")
        if not on_complete:
            on_complete = progress.on_complete
        assert(on_complete)

        on_error = my.kwargs.get("on_error")
        if not on_error:
            on_error = progress.on_error
        assert(on_error)

        data = []
        lines = []
        path = None
        my.paths = []
        error = []
        #while program.poll() is None:
        buffer = []
        line = ""
        while 1:
            char = program.stdout.read(1)
            if not char:
                break

            if char == "\n":
                line = "".join(buffer)

            else:
                buffer.append(char)

            if line:
                lines.append(line)
                if line.startswith(" "):

                    line = line.strip()
                    parts = re.split( re.compile("\ +"), line )

                    my.current_data = {
                        "bytes": int(parts[0]),
                        "percent": parts[1],
                        "rate": parts[2],
                        "time_left": parts[3]
                    }
                    #print my.current_data
                    #print "status: ", line

                    if on_update:
                        on_update(path, my.current_data)

                elif line.startswith("rsync "):
                    error.append(line)
                elif line.startswith("rsync: "):
                    error.append(line)
                elif line.startswith("deleting "):
                    my.handle_data_line(line)
                elif line.startswith("sent "):
                    my.handle_data_line(line)
                elif line.startswith("total "):
                    my.handle_data_line(line)
                elif line.startswith("sending incremental "):
                    my.handle_data_line(line)
                elif line.startswith("created directory "):
                    my.handle_data_line(line)
                else:
                    line = line.strip()
                    path = line
                    if not line.endswith("/"):
                        if line.find(" -> ") != -1:
                            parts = line.split(" -> ")
                            line = parts[0]

                        my.paths.append(line)

                # reset the line
                line = None
                buffer = []

        if error:
            error = "\n".join(error)
            if error.startswith("rsync: connection unexpectedly closed"):
                raise RSyncConnectionException(error)

            if on_error:
                on_error(path, {"message": error})
            raise Exception("Sync Error\n%s" % error)

        if on_complete:
            on_complete(path, {"message": "All files synchronized"})

        return "success"

    def handle_data_line(my, line):

        if line.startswith("sent "):
            # sent 520 bytes  received 22 bytes  361.33 bytes/sec
            pass

        elif line.startswith("total "):
            #total size is 431666  speedup is 796.43 (DRY RUN)
            parts = line.split()
            total_size = parts[3]
            total_size = int(total_size)
            my.data['total_size'] = total_size

class RSyncProgress(object):
    def __init__(my, **kwargs):
        from tactic_client_lib import TacticServerStub
        my.server = TacticServerStub.get()
        my.message_key = kwargs.get("message_key")
        my.paths = kwargs.get("paths")
        my.paths_sizes = kwargs.get("paths_sizes")
        if not my.paths:
            my.paths = []

    def on_update(my, path, data):
        data['path'] = path

        bytes = data.get("bytes")

        try:
            index = my.paths.index(path)
        except:
            index = 0
        data["path_index"] = index+1
        data["paths_count"] = len(my.paths)
        data["paths_sizes"] = my.paths_sizes

        total_size = 0
        current_size = 0
        for i in range(0, len(my.paths)):
            if i < index:
                current_size += my.paths_sizes[i]
            elif i == index:
                current_size += bytes

            total_size += my.paths_sizes[i]

        #print "current_size: ", current_size
        #print "total_size: ", total_size
        #print "percent: ", (current_size*1000)/(total_size*10)
        if total_size:
            total_percent = (current_size*1000)/(total_size*10)
        else:
            total_percent = 0
        data['total_percent'] = total_percent

        if my.message_key:
            my.server.log_message(my.message_key, data, status="in_progress")

    def on_error(my, path, data):
        if my.message_key:
            my.server.log_message(my.message_key, data, status="error")

    def on_complete(my, path, data):
        if my.message_key:
            my.server.log_message(my.message_key, data, status="complete")

if __name__ == '__main__':

    import time

    from_path = "/home/tactic/svg"

    start = time.time()

    dir_info = Common.get_dir_info(from_path)
    print "dir_info: ", dir_info
    total_size = dir_info.get("size")

    from pyasm.security import Batch
    Batch()

    class Progress(object):
        def __init__(my):
            my.total_sent = 0

        def on_update(my, path, data):
            from tactic_client_lib import TacticServerStub
            server = TacticServerStub.get()
            data['path'] = path

            bytes = data.get("bytes")
            my.total_sent += bytes

            print "path: ", path
            print "total: ", my.total_sent

            server.log_message("wow", data)
    progress = Progress()

    kwargs = {
            "login": "root",
            "server": "sync1.southpawtech.com",
            "from_path": from_path,
            "to_path": "/spt/test/svg",
            "on_update": progress.on_update
    }
    cmd = RSync(**kwargs)

    cmd.execute()

<|endoftext|>


<commit_before>import os
from itertools import zip_longest
from collections import defaultdict, Counter
import epride as ep
import pandas as pd

bridges_dict = {"16S": "GWATTACCGCGGCKGCTGCATCTTCTCCAAATGGGTCATGATC",
                "18S": "AAGAACGGCCATGCACCACCACATCTTCTCCAAATGGGTCATGATC",
                "narG2": "ACCGACATGCCGWTSCTGGTCATCTTCTCCAAATGGGTCATGATC",
                "nosZ2": "AACAAGTTCTCSAAGGACCGCATCTTCTCCAAATGGGTCATGATC",
                "nosZ3": "CTCMAAGGACCGGTTCMTSMCATCTTCTCCAAATGGGTCATGATC",
                "norB2": "GCACCGGYCACCAYTAYTWCCATCTTCTCCAAATGGGTCATGATC"}

size_filter_dict = {'narG2': 100, 'norB2': 90, 'nosZ3': 100, '18S': 70, '16S': 90, 'nosZ2': 100}

def move_barcodes_and_type_to_fasta_id(bc_seq, bridge_dict):
    bridge_dict = {key: ep.expand_primers(ep.reverse_complement(val)) for key, val in bridge_dict.items()}
    for seq_id, seq in bc_seq:
        for bridge_id, bridges in bridge_dict.items():
            for bridge in bridges:
                if bridge in seq:
                    bc, rest = seq.split(bridge)
                    if len(bc) == 20:
                        seq_id = "{} barcode={} barcode_type={}".format(seq_id.strip(), bc, bridge_id)
                        yield([seq_id, rest])

def process_barcode_info(input_seq_file, output_seq_file, bridge_dict):
    seqs = ep.read_fasta(input_seq_file)
    fasta_iter = move_barcodes_and_type_to_fasta_id(seqs, bridge_dict)
    ep.write_fasta(fasta_iter, output_seq_file)

def get_len_distr(seqs):
    len_distr_dict = defaultdict(list)
    for seq_id, seq in seqs:
        seq_type_id = seq_id.split("_")[-1]
        len_distr_dict[seq_type_id].append(int(len(seq) / 10) * 10)
    return len_distr_dict

class BCSeq(object):
    def __init__(self, seq, bridges, length_dict=None):
        self.bridges = bridges
        rc_bridges = {key: ep.expand_primers(ep.reverse_complement(val))
                      for key, val in self.bridges.items()}
        self.seq_list = list(ep.read_fasta(seq))
        self.processed_seqs = list(move_barcodes_and_type_to_fasta_id(self.seq_list, rc_bridges))
        self.seq_dict = defaultdict(list)
        for seq_id, seq in self.processed_seqs:
            seq_short_id = seq_id.split("_")[0].split(">")[1]
            self.seq_dict[seq_short_id].append([seq_id, seq])
        self.bc_data = {}
        self.li = LenDist({key: Counter(map(lambda x: int(x / 10) * 10, val)) for key, val in
                           get_len_distr(self.processed_seqs).items()}, length_dict)
        self.length_dict = length_dict

    def __getitem__(self, ix):
        if ix in self.bc_data:
            return self.bc_data[ix]
        else:
            self.bc_data[ix] = BCData(self.seq_dict[ix], self.length_dict)
            return self.bc_data[ix]

    def __repr__(self):
        repr_lst = ["Sequence group: {:15} number of sequences: {}".format(seq_name, len(seqs))
                    for seq_name, seqs in self.seq_dict.items()]
        repr_str = "\n".join(repr_lst)
        return repr_str

class BCData(object):
    def __init__(self, seq, length_dict=None):
        self.seq_list = seq
        self.li = LenDist({key: Counter(map(lambda x: int(x / 10) * 10, val))
                           for key, val in get_len_distr(self.seq_list).items()}, length_dict)
        self.seq_filt = []
        self.length_dict = length_dict
        if self.length_dict:
            self.filter_seqs()
            bc_type_table = pd.DataFrame([seq_id.strip().split("_")[2:]
                                          for seq_id, seq in self.seq_filt],
                                         columns=["Barcode", "Seq_type"])
        else:
            bc_type_table = pd.DataFrame([seq_id.strip().split("_")[2:]
                                          for seq_id, seq in self.seq_list],
                                         columns=["Barcode", "Seq_type"])
        self.type_table = bc_type_table.groupby(["Barcode", "Seq_type"])\
                                       .size().reset_index().rename(columns={0: "Count"})
        self.bc_table = pd.pivot_table(self.type_table, values="Count", index="Barcode",
                                       columns="Seq_type", fill_value=0)

    def __repr__(self):
        abund_table = self.bc_table.groupby(list(self.bc_table.columns))\
            .size().reset_index().sort_values(0, ascending=False)\
                                 .rename(columns={0: 'Count'}).head(20)
        abund_table.index = range(len(abund_table.index))
        repr_string = abund_table.to_string()
        return repr_string

    def filter_seqs(self):
        for seq_id, seq in self.seq_list:
            seq_type = seq_id.split("_")[-1]
            min_len = self.length_dict[seq_type]
            if len(seq) >= min_len:
                self.seq_filt.append([seq_id, seq])

class LenDist(object):
    def __init__(self, distr_info, length_dict=None):
        self.distr_info = distr_info
        self.length_dict = length_dict

    def process_repr(self):
        for item, val in self.distr_info.items():
            yield([item, ''])
            yield(['LI', 'Count'])
            count_per_item = sorted([[len_interval, count]
                                     for len_interval, count in val.items()],
                                    key=lambda x: x[0])
            for count_item, count in count_per_item:
                try:
                    if self.length_dict[item] == count_item:
                        yield(["---", "---"])
                    yield([count_item, count])
                except TypeError:
                    yield([count_item, count])

    def koe_repr(self):
        tmp = self.process_repr()
        return(list(tmp))
        # acc = []
        # for i in zip(*tmp):
        #     acc.append(i)
        # return acc
        # for i in zip(*self.process_repr()):
        #     acc.append("\t\t".join(i))
        # repr_str = "\n".join(acc)
        # return repr_str

    # def __repr__(self):
    #     # return str(self.process_repr())
    #     acc = []
    #     for i in zip(*self.process_repr()):
    #         acc.append("\t\t".join(i))
    #     repr_str = "\n".join(acc)
    #     return repr_str

# os.chdir("/Users/tamminma/gits/epic_barcoder/epic_barcoder/")

# test_fasta = list(ep.read_fasta("../test_seqs.fasta"))

# a = BCSeq("../test_seqs.fasta", bridges_dict, size_filter_dict)

# print(a)

# print(a['WW'])

# print(a.li.koe_repr())

# b = a.li.koe_repr()
<commit_msg>Tweaks<commit_after>import os
from itertools import zip_longest
from collections import defaultdict, Counter
import epride as ep
import pandas as pd

bridges_dict = {"16S": "GWATTACCGCGGCKGCTGCATCTTCTCCAAATGGGTCATGATC",
                "18S": "AAGAACGGCCATGCACCACCACATCTTCTCCAAATGGGTCATGATC",
                "narG2": "ACCGACATGCCGWTSCTGGTCATCTTCTCCAAATGGGTCATGATC",
                "nosZ2": "AACAAGTTCTCSAAGGACCGCATCTTCTCCAAATGGGTCATGATC",
                "nosZ3": "CTCMAAGGACCGGTTCMTSMCATCTTCTCCAAATGGGTCATGATC",
                "norB2": "GCACCGGYCACCAYTAYTWCCATCTTCTCCAAATGGGTCATGATC"}

size_filter_dict = {'narG2': 100, 'norB2': 90, 'nosZ3': 100, '18S': 70, '16S': 90, 'nosZ2': 100}

def move_barcodes_and_type_to_fasta_id(bc_seq, bridge_dict):
    bridge_dict = {key: ep.expand_primers(ep.reverse_complement(val))
                   for key, val in bridge_dict.items()}
    for seq_id, seq in bc_seq:
        for bridge_id, bridges in bridge_dict.items():
            for bridge in bridges:
                if bridge in seq:
                    bc, rest = seq.split(bridge)
                    if len(bc) == 20:
                        seq_id = "{} barcode={} sequence_type={}".format(seq_id.strip(),
                                                                        bc, bridge_id)
                        yield([seq_id, rest])

def process_barcode_info(input_seq_file, output_seq_file, bridge_dict):
    seqs = ep.read_fasta(input_seq_file)
    fasta_iter = move_barcodes_and_type_to_fasta_id(seqs, bridge_dict)
    ep.write_fasta(fasta_iter, output_seq_file)

def get_len_distr(seqs):
    len_distr_dict = defaultdict(list)
    for seq_id, seq in seqs:
        seq_type_id = seq_id.split("_")[-1]
        len_distr_dict[seq_type_id].append(int(len(seq) / 10) * 10)
    return len_distr_dict

class BCSeq(object):
    def __init__(self, seq, bridges, length_dict=None):
        self.bridges = bridges
        rc_bridges = {key: ep.expand_primers(ep.reverse_complement(val))
                      for key, val in self.bridges.items()}
        self.seq_list = list(ep.read_fasta(seq))
        self.processed_seqs = list(move_barcodes_and_type_to_fasta_id(self.seq_list, rc_bridges))
        self.seq_dict = defaultdict(list)
        for seq_id, seq in self.processed_seqs:
            seq_short_id = seq_id.split("_")[0].split(">")[1]
            self.seq_dict[seq_short_id].append([seq_id, seq])
        self.bc_data = {}
        self.li = LenDist({key: Counter(map(lambda x: int(x / 10) * 10, val)) for key, val in
                           get_len_distr(self.processed_seqs).items()}, length_dict)
        self.length_dict = length_dict

    def __getitem__(self, ix):
        if ix in self.bc_data:
            return self.bc_data[ix]
        else:
            self.bc_data[ix] = BCData(self.seq_dict[ix], self.length_dict)
            return self.bc_data[ix]

    def __repr__(self):
        repr_lst = ["Sequence group: {:15} number of sequences: {}".format(seq_name, len(seqs))
                    for seq_name, seqs in self.seq_dict.items()]
        repr_str = "\n".join(repr_lst)
        return repr_str

class BCData(object):
    def __init__(self, seq, length_dict=None):
        self.seq_list = seq
        self.li = LenDist({key: Counter(map(lambda x: int(x / 10) * 10, val))
                           for key, val in get_len_distr(self.seq_list).items()}, length_dict)
        self.seq_filt = []
        self.length_dict = length_dict
        if self.length_dict:
            self.filter_seqs()
            bc_type_table = pd.DataFrame([seq_id.strip().split("_")[2:]
                                          for seq_id, seq in self.seq_filt],
                                         columns=["Barcode", "Seq_type"])
        else:
            bc_type_table = pd.DataFrame([seq_id.strip().split("_")[2:]
                                          for seq_id, seq in self.seq_list],
                                         columns=["Barcode", "Seq_type"])
        self.type_table = bc_type_table.groupby(["Barcode", "Seq_type"])\
                                       .size().reset_index().rename(columns={0: "Count"})
        self.bc_table = pd.pivot_table(self.type_table, values="Count", index="Barcode",
                                       columns="Seq_type", fill_value=0)

    def __repr__(self):
        abund_table = self.bc_table.groupby(list(self.bc_table.columns))\
            .size().reset_index().sort_values(0, ascending=False)\
                                 .rename(columns={0: 'Count'}).head(20)
        abund_table.index = range(len(abund_table.index))
        repr_string = abund_table.to_string()
        return repr_string

    def filter_seqs(self):
        for seq_id, seq in self.seq_list:
            seq_type = seq_id.split("_")[-1]
            min_len = self.length_dict[seq_type]
            if len(seq) >= min_len:
                self.seq_filt.append([seq_id, seq])

class LenDist(object):
    def __init__(self, distr_info, length_dict=None):
        self.distr_info = distr_info
        self.length_dict = length_dict

    def process_repr(self):
        for item, val in self.distr_info.items():
            yield([item, ''])
            yield(['LI', 'Count'])
            count_per_item = sorted([[len_interval, count]
                                     for len_interval, count in val.items()],
                                    key=lambda x: x[0])
            for count_item, count in count_per_item:
                try:
                    if self.length_dict[item] == count_item:
                        yield(["---", "---"])
                    yield([count_item, count])
                except TypeError:
                    yield([count_item, count])

    def koe_repr(self):
        tmp = self.process_repr()
        return(list(tmp))
        # acc = []
        # for i in zip(*tmp):
        #     acc.append(i)
        # return acc
        # for i in zip(*self.process_repr()):
        #     acc.append("\t\t".join(i))
        # repr_str = "\n".join(acc)
        # return repr_str

    # def __repr__(self):
    #     # return str(self.process_repr())
    #     acc = []
    #     for i in zip(*self.process_repr()):
    #         acc.append("\t\t".join(i))
    #     repr_str = "\n".join(acc)
    #     return repr_str

# os.chdir("/Users/tamminma/gits/epic_barcoder/epic_barcoder/")

# test_fasta = list(ep.read_fasta("../test_seqs.fasta"))

# a = BCSeq("../test_seqs.fasta", bridges_dict, size_filter_dict)

# print(a)

# print(a['WW'])

# print(a.li.koe_repr())

# b = a.li.koe_repr()
<|endoftext|>


<commit_before>import os
import subprocess
import string
import random
import time
from itertools import zip_longest
from collections import defaultdict, Counter
import epride as ep
import pandas as pd

bridges_dict = {"16S": "GWATTACCGCGGCKGCTGCATCTTCTCCAAATGGGTCATGATC",
                "18S": "AAGAACGGCCATGCACCACCACATCTTCTCCAAATGGGTCATGATC",
                "narG2": "ACCGACATGCCGWTSCTGGTCATCTTCTCCAAATGGGTCATGATC",
                "nosZ2": "AACAAGTTCTCSAAGGACCGCATCTTCTCCAAATGGGTCATGATC",
                "nosZ3": "CTCMAAGGACCGGTTCMTSMCATCTTCTCCAAATGGGTCATGATC",
                "norB2": "GCACCGGYCACCAYTAYTWCCATCTTCTCCAAATGGGTCATGATC"}

size_filter_dict = {'narG2': 100, 'norB2': 90, 'nosZ3': 100, '18S': 70, '16S': 90, 'nosZ2': 100}

array_dict = {'lsf': '''#!/bin/bash
#BSUB -n 1
#BSUB -R "rusage[mem={}]"
#BSUB -J {}[1-{}]
#BSUB -o wrf.%I
#BSUB -e wrf.%I
#BSUB -W {}

cd {}

name=$(sed -n "$LSB_JOBINDEX"p {})

{}
''', 'slurm': '''#!/bin/bash
#SBATCH --mem-per-cpu={}
#SBATCH -J {}
#SBATCH --array=1-{}
#SBATCH -t {}
#SBATCH -o array_job_out_%j.txt
#SBATCH -e array_job_err_%j.txt
#SBATCH -n 1
#SBATCH -p serial

cd {}

name=$(sed -n "$SLURM_ARRAY_TASK_ID"p {})

{}
'''}

def move_barcodes_and_type_to_fasta_id(bc_seq, bridge_dict):
    bridge_dict = {key: ep.expand_primers(ep.reverse_complement(val))
                   for key, val in bridge_dict.items()}
    for seq_id, seq in bc_seq:
        for bridge_id, bridges in bridge_dict.items():
            for bridge in bridges:
                if bridge in seq:
                    bc, rest = seq.split(bridge)
                    if len(bc) == 20:
                        seq_id = "{} barcode={} sequence_type={}".format(seq_id.strip(),
                                                                        bc, bridge_id)
                        yield([seq_id, rest])

def process_barcode_info(input_seq_file, output_seq_file, bridge_dict):
    seqs = ep.read_fasta(input_seq_file)
    fasta_iter = move_barcodes_and_type_to_fasta_id(seqs, bridge_dict)
    ep.write_fasta(fasta_iter, output_seq_file)

def get_len_distr(seqs):
    len_distr_dict = defaultdict(list)
    for seq_id, seq in seqs:
        seq_type_id = seq_id.split("_")[-1]
        len_distr_dict[seq_type_id].append(int(len(seq) / 10) * 10)
    return len_distr_dict

def get_seed_dict(uc_file):
    seed_dict = {}
    with open(uc_file) as f:
        for line in f:
            if line.split()[0] == "H":
                seq_id = line.split()[8]
                seed_id = line.split()[9]
                seed_dict[seq_id] = seed_id
            if line.split()[0] == "S":
                seq_id = line.split()[8]
                seed_dict[seq_id] = seq_id
    return seed_dict

def add_otus_to_fasta(seq_file, uc_file, output_file):
    seeds = get_seed_dict(uc_file)
    seq_acc = []
    for seq_id, seq in ep.read_fasta(seq_file):
        short_id = seq_id[1:].split()[0]
        seed_id = seeds[short_id]
        new_seq_id = "{} OTU={}".format(seq_id.strip(), seed_id)
        seq_acc.append([new_seq_id, seq])
    ep.write_fasta(seq_acc, output_file)

def generate_id(size=8):
    """ Generate random sequences of characters for temporary file names.
    """
    chars = string.ascii_uppercase + string.digits
    return ''.join(random.choice(chars) for _ in range(size))

def make_split_seq_dict(seq_iter, no_splits):
    seq_list = list(seq_iter)
    seq_len = len(seq_list)
    chunk_size = int(seq_len/no_splits)
    split_dict = defaultdict(list)
    for ix, (seq_id, seq) in enumerate(seq_list):
        if ix % chunk_size == 0:
            chunk_id = generate_id()
        split_dict[chunk_id].append([seq_id, seq])
    return split_dict

def split_seqs(seq_file, no_splits):
    seqs = ep.read_fasta(seq_file)
    split_dict = make_split_seq_dict(seqs, no_splits)
    for key, val in split_dict.items():
        seq_name = key + "_tmp.fasta"
        ep.write_fasta(val, seq_name)
    return list(split_dict.keys())

def make_array_job(seqs, batch_command, post_command=None, no_splits=1000, scheduler='slurm', memory=2048, run_time='02:00'):
    job_name = generate_id()
    user = subprocess.check_output('whoami', universal_newlines=True).strip()
    namelist = job_name + "_tmp.namelist"
    seq_ids = split_seqs(seqs, no_splits)
    job_no = len(seq_ids)
    home_dir = os.getcwd()
    array = array_dict[scheduler].format(memory, job_name, job_no,
                                         run_time, home_dir, namelist,
                                         batch_command)
    array_file_name = generate_id() + "_tmp.sh"
    with open(array_file_name, "w") as f:
        for line in array:
            f.write(line)
    with open(namelist, "w") as f:
        for item in seq_ids:
            f.write(item + "_tmp.fasta\n")
    if scheduler == 'slurm':
        subprocess.call(['sbatch', array_file_name])
        while True:
            jobs = subprocess.check_output(['squeue', '-u', user],
                                           universal_newlines=True).split("\n")
            if len(jobs) == 1:
                break
            print("{} jobs left".format(len(jobs) - 1))
            time.sleep(5)

class BCSeq(object):
    def __init__(self, seq, bridges, length_dict=None):
        self.bridges = bridges
        rc_bridges = {key: ep.expand_primers(ep.reverse_complement(val))
                      for key, val in self.bridges.items()}
        self.seq_list = list(ep.read_fasta(seq))
        self.processed_seqs = list(move_barcodes_and_type_to_fasta_id(self.seq_list, rc_bridges))
        self.seq_dict = defaultdict(list)
        for seq_id, seq in self.processed_seqs:
            seq_short_id = seq_id.split("_")[0].split(">")[1]
            self.seq_dict[seq_short_id].append([seq_id, seq])
        self.bc_data = {}
        self.li = LenDist({key: Counter(map(lambda x: int(x / 10) * 10, val)) for key, val in
                           get_len_distr(self.processed_seqs).items()}, length_dict)
        self.length_dict = length_dict

    def __getitem__(self, ix):
        if ix in self.bc_data:
            return self.bc_data[ix]
        else:
            self.bc_data[ix] = BCData(self.seq_dict[ix], self.length_dict)
            return self.bc_data[ix]

    def __repr__(self):
        repr_lst = ["Sequence group: {:15} number of sequences: {}".format(seq_name, len(seqs))
                    for seq_name, seqs in self.seq_dict.items()]
        repr_str = "\n".join(repr_lst)
        return repr_str

class BCData(object):
    def __init__(self, seq, length_dict=None):
        self.seq_list = seq
        self.li = LenDist({key: Counter(map(lambda x: int(x / 10) * 10, val))
                           for key, val in get_len_distr(self.seq_list).items()}, length_dict)
        self.seq_filt = []
        self.length_dict = length_dict
        if self.length_dict:
            self.filter_seqs()
            bc_type_table = pd.DataFrame([seq_id.strip().split("_")[2:]
                                          for seq_id, seq in self.seq_filt],
                                         columns=["Barcode", "Seq_type"])
        else:
            bc_type_table = pd.DataFrame([seq_id.strip().split("_")[2:]
                                          for seq_id, seq in self.seq_list],
                                         columns=["Barcode", "Seq_type"])
        self.type_table = bc_type_table.groupby(["Barcode", "Seq_type"])\
                                       .size().reset_index().rename(columns={0: "Count"})
        self.bc_table = pd.pivot_table(self.type_table, values="Count", index="Barcode",
                                       columns="Seq_type", fill_value=0)

    def __repr__(self):
        abund_table = self.bc_table.groupby(list(self.bc_table.columns))\
            .size().reset_index().sort_values(0, ascending=False)\
                                 .rename(columns={0: 'Count'}).head(20)
        abund_table.index = range(len(abund_table.index))
        repr_string = abund_table.to_string()
        return repr_string

    def filter_seqs(self):
        for seq_id, seq in self.seq_list:
            seq_type = seq_id.split("_")[-1]
            min_len = self.length_dict[seq_type]
            if len(seq) >= min_len:
                self.seq_filt.append([seq_id, seq])

class LenDist(object):
    def __init__(self, distr_info, length_dict=None):
        self.distr_info = distr_info
        self.length_dict = length_dict

    def process_repr(self):
        for item, val in self.distr_info.items():
            yield([item, ''])
            yield(['LI', 'Count'])
            count_per_item = sorted([[len_interval, count]
                                     for len_interval, count in val.items()],
                                    key=lambda x: x[0])
            for count_item, count in count_per_item:
                try:
                    if self.length_dict[item] == count_item:
                        yield(["---", "---"])
                    yield([count_item, count])
                except TypeError:
                    yield([count_item, count])

    def koe_repr(self):
        tmp = self.process_repr()
        return(list(tmp))
        # acc = []
        # for i in zip(*tmp):
        #     acc.append(i)
        # return acc
        # for i in zip(*self.process_repr()):
        #     acc.append("\t\t".join(i))
        # repr_str = "\n".join(acc)
        # return repr_str

    # def __repr__(self):
    #     # return str(self.process_repr())
    #     acc = []
    #     for i in zip(*self.process_repr()):
    #         acc.append("\t\t".join(i))
    #     repr_str = "\n".join(acc)
    #     return repr_str

# os.chdir("/Users/tamminma/gits/epic_barcoder/epic_barcoder/")

# test_fasta = list(ep.read_fasta("../test_seqs.fasta"))

# a = BCSeq("../test_seqs.fasta", bridges_dict, size_filter_dict)

# print(a)

# print(a['WW'])

# print(a.li.koe_repr())

# b = a.li.koe_repr()
<commit_msg>Tweaks<commit_after>import os
import subprocess
import string
import random
import time
from itertools import zip_longest
from collections import defaultdict, Counter
import epride as ep
import pandas as pd

bridges_dict = {"16S": "GWATTACCGCGGCKGCTGCATCTTCTCCAAATGGGTCATGATC",
                "18S": "AAGAACGGCCATGCACCACCACATCTTCTCCAAATGGGTCATGATC",
                "narG2": "ACCGACATGCCGWTSCTGGTCATCTTCTCCAAATGGGTCATGATC",
                "nosZ2": "AACAAGTTCTCSAAGGACCGCATCTTCTCCAAATGGGTCATGATC",
                "nosZ3": "CTCMAAGGACCGGTTCMTSMCATCTTCTCCAAATGGGTCATGATC",
                "norB2": "GCACCGGYCACCAYTAYTWCCATCTTCTCCAAATGGGTCATGATC"}

size_filter_dict = {'narG2': 100, 'norB2': 90, 'nosZ3': 100, '18S': 70, '16S': 90, 'nosZ2': 100}

array_dict = {'lsf': '''#!/bin/bash
#BSUB -n 1
#BSUB -R "rusage[mem={}]"
#BSUB -J {}[1-{}]
#BSUB -o wrf.%I
#BSUB -e wrf.%I
#BSUB -W {}

cd {}

name=$(sed -n "$LSB_JOBINDEX"p {})

{}
''', 'slurm': '''#!/bin/bash
#SBATCH --mem-per-cpu={}
#SBATCH -J {}
#SBATCH --array=1-{}
#SBATCH -t {}
#SBATCH -o array_job_out_%j.txt
#SBATCH -e array_job_err_%j.txt
#SBATCH -n 1
#SBATCH -p serial

cd {}

name=$(sed -n "$SLURM_ARRAY_TASK_ID"p {})

{}
'''}

def move_barcodes_and_type_to_fasta_id(bc_seq, bridge_dict):
    bridge_dict = {key: ep.expand_primers(ep.reverse_complement(val))
                   for key, val in bridge_dict.items()}
    for seq_id, seq in bc_seq:
        for bridge_id, bridges in bridge_dict.items():
            for bridge in bridges:
                if bridge in seq:
                    bc, rest = seq.split(bridge)
                    if len(bc) == 20:
                        seq_id = "{} barcode={} sequence_type={}".format(seq_id.strip(),
                                                                        bc, bridge_id)
                        yield([seq_id, rest])

def process_barcode_info(input_seq_file, output_seq_file, bridge_dict):
    seqs = ep.read_fasta(input_seq_file)
    fasta_iter = move_barcodes_and_type_to_fasta_id(seqs, bridge_dict)
    ep.write_fasta(fasta_iter, output_seq_file)

def get_len_distr(seqs):
    len_distr_dict = defaultdict(list)
    for seq_id, seq in seqs:
        seq_type_id = seq_id.split("_")[-1]
        len_distr_dict[seq_type_id].append(int(len(seq) / 10) * 10)
    return len_distr_dict

def get_seed_dict(uc_file):
    seed_dict = {}
    with open(uc_file) as f:
        for line in f:
            if line.split()[0] == "H":
                seq_id = line.split()[8]
                seed_id = line.split()[9]
                seed_dict[seq_id] = seed_id
            if line.split()[0] == "S":
                seq_id = line.split()[8]
                seed_dict[seq_id] = seq_id
    return seed_dict

def add_otus_to_fasta(seq_file, uc_file, output_file):
    seeds = get_seed_dict(uc_file)
    seq_acc = []
    for seq_id, seq in ep.read_fasta(seq_file):
        short_id = seq_id[1:].split()[0]
        seed_id = seeds[short_id]
        new_seq_id = "{} OTU={}".format(seq_id.strip(), seed_id)
        seq_acc.append([new_seq_id, seq])
    ep.write_fasta(seq_acc, output_file)

def generate_id(size=8):
    """ Generate random sequences of characters for temporary file names.
    """
    chars = string.ascii_uppercase + string.digits
    return ''.join(random.choice(chars) for _ in range(size))

def make_split_seq_dict(seq_iter, no_splits):
    seq_list = list(seq_iter)
    seq_len = len(seq_list)
    chunk_size = int(seq_len/no_splits)
    split_dict = defaultdict(list)
    for ix, (seq_id, seq) in enumerate(seq_list):
        if ix % chunk_size == 0:
            chunk_id = generate_id()
        split_dict[chunk_id].append([seq_id, seq])
    return split_dict

def split_seqs(seq_file, no_splits):
    seqs = ep.read_fasta(seq_file)
    split_dict = make_split_seq_dict(seqs, no_splits)
    for key, val in split_dict.items():
        seq_name = key + "_tmp.fasta"
        ep.write_fasta(val, seq_name)
    return list(split_dict.keys())

def make_array_job(seqs, batch_command, post_command=None, no_splits=1000, scheduler='slurm', memory=2048, run_time='02:00'):
    job_name = generate_id()
    user = subprocess.check_output('whoami', universal_newlines=True).strip()
    namelist = job_name + "_tmp.namelist"
    seq_ids = split_seqs(seqs, no_splits)
    job_no = len(seq_ids)
    home_dir = os.getcwd()
    array = array_dict[scheduler].format(memory, job_name, job_no,
                                         run_time, home_dir, namelist,
                                         batch_command)
    array_file_name = generate_id() + "_tmp.sh"
    with open(array_file_name, "w") as f:
        for line in array:
            f.write(line)
    with open(namelist, "w") as f:
        for item in seq_ids:
            f.write(item + "_tmp.fasta\n")
    if scheduler == 'slurm':
        subprocess.call(['sbatch', array_file_name])
        print("Submitting {} jobs.".format(job_no))
        time.sleep(10)
        while True:
            jobs = subprocess.check_output(['squeue', '-u', user],
                                           universal_newlines=True).split("\n")
            if len(jobs) == 2:
                break
            print("{} jobs left".format(len(jobs) - 2))
            time.sleep(5)

class BCSeq(object):
    def __init__(self, seq, bridges, length_dict=None):
        self.bridges = bridges
        rc_bridges = {key: ep.expand_primers(ep.reverse_complement(val))
                      for key, val in self.bridges.items()}
        self.seq_list = list(ep.read_fasta(seq))
        self.processed_seqs = list(move_barcodes_and_type_to_fasta_id(self.seq_list, rc_bridges))
        self.seq_dict = defaultdict(list)
        for seq_id, seq in self.processed_seqs:
            seq_short_id = seq_id.split("_")[0].split(">")[1]
            self.seq_dict[seq_short_id].append([seq_id, seq])
        self.bc_data = {}
        self.li = LenDist({key: Counter(map(lambda x: int(x / 10) * 10, val)) for key, val in
                           get_len_distr(self.processed_seqs).items()}, length_dict)
        self.length_dict = length_dict

    def __getitem__(self, ix):
        if ix in self.bc_data:
            return self.bc_data[ix]
        else:
            self.bc_data[ix] = BCData(self.seq_dict[ix], self.length_dict)
            return self.bc_data[ix]

    def __repr__(self):
        repr_lst = ["Sequence group: {:15} number of sequences: {}".format(seq_name, len(seqs))
                    for seq_name, seqs in self.seq_dict.items()]
        repr_str = "\n".join(repr_lst)
        return repr_str

class BCData(object):
    def __init__(self, seq, length_dict=None):
        self.seq_list = seq
        self.li = LenDist({key: Counter(map(lambda x: int(x / 10) * 10, val))
                           for key, val in get_len_distr(self.seq_list).items()}, length_dict)
        self.seq_filt = []
        self.length_dict = length_dict
        if self.length_dict:
            self.filter_seqs()
            bc_type_table = pd.DataFrame([seq_id.strip().split("_")[2:]
                                          for seq_id, seq in self.seq_filt],
                                         columns=["Barcode", "Seq_type"])
        else:
            bc_type_table = pd.DataFrame([seq_id.strip().split("_")[2:]
                                          for seq_id, seq in self.seq_list],
                                         columns=["Barcode", "Seq_type"])
        self.type_table = bc_type_table.groupby(["Barcode", "Seq_type"])\
                                       .size().reset_index().rename(columns={0: "Count"})
        self.bc_table = pd.pivot_table(self.type_table, values="Count", index="Barcode",
                                       columns="Seq_type", fill_value=0)

    def __repr__(self):
        abund_table = self.bc_table.groupby(list(self.bc_table.columns))\
            .size().reset_index().sort_values(0, ascending=False)\
                                 .rename(columns={0: 'Count'}).head(20)
        abund_table.index = range(len(abund_table.index))
        repr_string = abund_table.to_string()
        return repr_string

    def filter_seqs(self):
        for seq_id, seq in self.seq_list:
            seq_type = seq_id.split("_")[-1]
            min_len = self.length_dict[seq_type]
            if len(seq) >= min_len:
                self.seq_filt.append([seq_id, seq])

class LenDist(object):
    def __init__(self, distr_info, length_dict=None):
        self.distr_info = distr_info
        self.length_dict = length_dict

    def process_repr(self):
        for item, val in self.distr_info.items():
            yield([item, ''])
            yield(['LI', 'Count'])
            count_per_item = sorted([[len_interval, count]
                                     for len_interval, count in val.items()],
                                    key=lambda x: x[0])
            for count_item, count in count_per_item:
                try:
                    if self.length_dict[item] == count_item:
                        yield(["---", "---"])
                    yield([count_item, count])
                except TypeError:
                    yield([count_item, count])

    def koe_repr(self):
        tmp = self.process_repr()
        return(list(tmp))
        # acc = []
        # for i in zip(*tmp):
        #     acc.append(i)
        # return acc
        # for i in zip(*self.process_repr()):
        #     acc.append("\t\t".join(i))
        # repr_str = "\n".join(acc)
        # return repr_str

    # def __repr__(self):
    #     # return str(self.process_repr())
    #     acc = []
    #     for i in zip(*self.process_repr()):
    #         acc.append("\t\t".join(i))
    #     repr_str = "\n".join(acc)
    #     return repr_str

# os.chdir("/Users/tamminma/gits/epic_barcoder/epic_barcoder/")

# test_fasta = list(ep.read_fasta("../test_seqs.fasta"))

# a = BCSeq("../test_seqs.fasta", bridges_dict, size_filter_dict)

# print(a)

# print(a['WW'])

# print(a.li.koe_repr())

# b = a.li.koe_repr()
<|endoftext|>


<commit_before>from __future__ import absolute_import

from django.conf import settings
from typing import Optional

import logging
import traceback
import platform

from django.core import mail
from django.http import HttpRequest
from django.utils.log import AdminEmailHandler
from django.views.debug import ExceptionReporter, get_exception_reporter_filter

from zerver.lib.queue import queue_json_publish

class AdminZulipHandler(logging.Handler):
    """An exception log handler that sends the exception to the queue to be
       sent to the Zulip feedback server.
    """

    # adapted in part from django/utils/log.py

    def __init__(self):
        # type: () -> None
        logging.Handler.__init__(self)

    def emit(self, record):
        # type: (ExceptionReporter) -> None
        try:
            request = record.request  # type: HttpRequest

            exception_filter = get_exception_reporter_filter(request)

            if record.exc_info:
                stack_trace = ''.join(traceback.format_exception(*record.exc_info))  # type: Optional[str]
            else:
                stack_trace = None

            report = dict(
                node = platform.node(),
                message = record.getMessage(),
                stack_trace = stack_trace,
            )

            report['path'] = request.path
            report['method'] = request.method
            report['remote_addr'] = request.META.get('REMOTE_ADDR', None),
            report['query_string'] = request.META.get('QUERY_STRING', None),
            report['server_name'] = request.META.get('SERVER_NAME', None),
            try:
                from django.contrib.auth.models import AnonymousUser
                user_profile = request.user
                if isinstance(user_profile, AnonymousUser):
                    user_full_name = None
                    user_email = None
                else:
                    user_full_name = user_profile.full_name
                    user_email = user_profile.email
            except Exception:
                # Unexpected exceptions here should be handled gracefully
                traceback.print_exc()
                user_full_name = None
                user_email = None
            report['user_email'] = user_email
            report['user_full_name'] = user_full_name

            try:
                report['data'] = request.GET if request.method == 'GET' else \
                    exception_filter.get_post_parameters(request)
            except Exception:
                # exception_filter.get_post_parameters will throw
                # RequestDataTooBig if there's a really big file uploaded
                report['data'] = {}

            try:
                report['host'] = request.get_host().split(':')[0]
            except Exception:
                # request.get_host() will throw a DisallowedHost
                # exception if the host is invalid
                report['host'] = platform.node()

        except Exception:
            traceback.print_exc()
            report = dict(
                node = platform.node(),
                host = platform.node(),
                message = record.getMessage(),
            )

        try:
            if settings.STAGING_ERROR_NOTIFICATIONS:
                # On staging, process the report directly so it can happen inside this
                # try/except to prevent looping
                from zilencer.error_notify import notify_server_error
                notify_server_error(report)
            else:
                queue_json_publish('error_reports', dict(
                    type = "server",
                    report = report,
                ), lambda x: None)
        except Exception:
            # If this breaks, complain loudly but don't pass the traceback up the stream
            # However, we *don't* want to use logging.exception since that could trigger a loop.
            logging.warning("Reporting an exception triggered an exception!", exc_info=True)
<commit_msg>AdminZulipHandler: Extract add_request_metadata.<commit_after>from __future__ import absolute_import

from django.conf import settings
from typing import Any, Dict, Optional

import logging
import traceback
import platform

from django.core import mail
from django.http import HttpRequest
from django.utils.log import AdminEmailHandler
from django.views.debug import ExceptionReporter, get_exception_reporter_filter

from zerver.lib.queue import queue_json_publish

def add_request_metadata(report, request):
    # type: (Dict[str, Any], HttpRequest) -> None
    report['path'] = request.path
    report['method'] = request.method
    report['remote_addr'] = request.META.get('REMOTE_ADDR', None),
    report['query_string'] = request.META.get('QUERY_STRING', None),
    report['server_name'] = request.META.get('SERVER_NAME', None),
    try:
        from django.contrib.auth.models import AnonymousUser
        user_profile = request.user
        if isinstance(user_profile, AnonymousUser):
            user_full_name = None
            user_email = None
        else:
            user_full_name = user_profile.full_name
            user_email = user_profile.email
    except Exception:
        # Unexpected exceptions here should be handled gracefully
        traceback.print_exc()
        user_full_name = None
        user_email = None
    report['user_email'] = user_email
    report['user_full_name'] = user_full_name

    exception_filter = get_exception_reporter_filter(request)
    try:
        report['data'] = request.GET if request.method == 'GET' else \
            exception_filter.get_post_parameters(request)
    except Exception:
        # exception_filter.get_post_parameters will throw
        # RequestDataTooBig if there's a really big file uploaded
        report['data'] = {}

    try:
        report['host'] = request.get_host().split(':')[0]
    except Exception:
        # request.get_host() will throw a DisallowedHost
        # exception if the host is invalid
        report['host'] = platform.node()

class AdminZulipHandler(logging.Handler):
    """An exception log handler that sends the exception to the queue to be
       sent to the Zulip feedback server.
    """

    # adapted in part from django/utils/log.py

    def __init__(self):
        # type: () -> None
        logging.Handler.__init__(self)

    def emit(self, record):
        # type: (logging.LogRecord) -> None
        try:
            if record.exc_info:
                stack_trace = ''.join(traceback.format_exception(*record.exc_info))  # type: Optional[str]
            else:
                stack_trace = None

            report = dict(
                node = platform.node(),
                message = record.getMessage(),
                stack_trace = stack_trace,
            )
            add_request_metadata(report, record.request)  # type: ignore  # record.request is added dynamically
        except Exception:
            traceback.print_exc()
            report = dict(
                node = platform.node(),
                host = platform.node(),
                message = record.getMessage(),
            )

        try:
            if settings.STAGING_ERROR_NOTIFICATIONS:
                # On staging, process the report directly so it can happen inside this
                # try/except to prevent looping
                from zilencer.error_notify import notify_server_error
                notify_server_error(report)
            else:
                queue_json_publish('error_reports', dict(
                    type = "server",
                    report = report,
                ), lambda x: None)
        except Exception:
            # If this breaks, complain loudly but don't pass the traceback up the stream
            # However, we *don't* want to use logging.exception since that could trigger a loop.
            logging.warning("Reporting an exception triggered an exception!", exc_info=True)
<|endoftext|>


<commit_before>import datetime
import pytz
import math
import pandas

# from gevent.select import select
from zmq.core.poll import select

import zipline.messaging as qmsg
import zipline.util as qutil
import zipline.protocol as zp
import zipline.finance.performance as perf

class TradeSimulationClient(qmsg.Component):

    def __init__(self, trading_environment):
        qmsg.Component.__init__(self)
        self.received_count         = 0
        self.prev_dt                = None
        self.event_queue            = None
        self.txn_count              = 0
        self.trading_environment    = trading_environment
        self.current_dt             = trading_environment.period_start
        self.last_iteration_dur     = datetime.timedelta(seconds=0)
        self.algorithm              = None

        assert self.trading_environment.frame_index != None
        self.event_frame = pandas.DataFrame(
            index=self.trading_environment.frame_index
        )

        self.perf = perf.PerformanceTracker(self.trading_environment)

    @property
    def get_id(self):
        return str(zp.FINANCE_COMPONENT.TRADING_CLIENT)

    def set_algorithm(self, algorithm):
        """
        :param algorithm: must implement the algorithm protocol. See
        :py:mod:`zipline.test.algorithm`
        """
        self.algorithm = algorithm
        #register the trading_client's order method with the algorithm
        self.algorithm.set_order(self.order)

    def open(self):
        self.result_feed = self.connect_result()
        self.order_socket = self.connect_order()

    def do_work(self):
        # poll all the sockets
        socks = dict(self.poll.poll(self.heartbeat_timeout))

        # see if the poller has results for the result_feed
        if self.result_feed in socks and \
            socks[self.result_feed] == self.zmq.POLLIN:

            # get the next message from the result feed
            msg = self.result_feed.recv()

            # if the feed is done, shut 'er down
            if msg == str(zp.CONTROL_PROTOCOL.DONE):
                qutil.LOGGER.info("Client is DONE!")
                # signal the performance tracker that the simulation has
                # ended. Perf will internally calculate the full risk report.
                self.perf.handle_simulation_end()
                # shutdown the feedback loop to the OrderDataSource
                self.signal_order_done()
                # signal Simulator, our ComponentHost, that this component is
                # done and Simulator needn't block exit on this component.
                self.signal_done()
                return

            # result_feed is a merge component, so unframe accordingly
            event = zp.MERGE_UNFRAME(msg)

            # update performance and relay the event to the algorithm
            self.process_event(event)

             # signal done to order source.
            self.order_socket.send(str(zp.ORDER_PROTOCOL.BREAK))

    def process_event(self, event):
        # track the number of transactions, for testing purposes.
        if(event.TRANSACTION != None):
            self.txn_count += 1

        #filter order flow out of the events sent to callbacks
        if event.source_id != zp.FINANCE_COMPONENT.ORDER_SOURCE:

            # the performance class needs to process each event, without
            # skipping. Algorithm should wait until the performance has been
            # updated, so that down stream components can safely assume that
            # performance is up to date. Note that this is done before we
            # mark the time for the algorithm's processing, thereby not
            # running the algo's clock for performance book keeping.
            self.perf.process_event(event)

            # mark the start time for client's processing of this event.
            event_start = datetime.datetime.utcnow()

            # queue the event.
            self.queue_event(event)

            # if the event is later than our current time, run the algo
            # otherwise, the algorithm has fallen behind the feed
            # and processing per event is longer than time between events.
            if event.dt >= self.current_dt:
                self.run_algorithm()

            # tally the time spent on this iteration
            self.last_iteration_dur = datetime.datetime.utcnow() - event_start
            # move the algorithm's clock forward to include iteration time
            self.current_dt = self.current_dt + self.last_iteration_dur

    def run_algorithm(self):
        """
        As per the algorithm protocol:

        - Set the current portfolio for the algorithm as per protocol.
        - Construct frame based on backlog of events, send to algorithm.
        """
        current_portfolio = self.perf.get_portfolio()
        self.algorithm.set_portfolio(current_portfolio)
        frame = self.get_frame()
        if len(frame) > 0:
            self.algorithm.handle_frame(frame)

    def connect_order(self):
        return self.connect_push_socket(self.addresses['order_address'])

    def order(self, sid, amount):
        order = zp.namedict({
            'dt':self.current_dt,
            'sid':sid,
            'amount':amount
        })

        self.order_socket.send(zp.ORDER_FRAME(order))

    def signal_order_done(self):
        self.order_socket.send(str(zp.ORDER_PROTOCOL.DONE))

    def queue_event(self, event):
        if self.event_queue == None:
            self.event_queue = []
        series = event.as_series()
        self.event_queue.append(series)

    def get_frame(self):
        for event in self.event_queue:
            self.event_frame[event['sid']] = event
        self.event_queue = []
        return self.event_frame

class OrderDataSource(qmsg.DataSource):
    """DataSource that relays orders from the client"""

    def __init__(self):
        """
        :param simulation_time: datetime in UTC timezone, sets the start
        time of simulation. orders
            will be timestamped relative to this datetime.
                event = {
                    'sid'    : an integer for security id,
                    'dt'     : datetime object,
                    'price'  : float for price,
                    'volume' : integer for volume
                }
        """
        qmsg.DataSource.__init__(self, zp.FINANCE_COMPONENT.ORDER_SOURCE)
        self.sent_count         = 0

    @property
    def get_type(self):
        return zp.DATASOURCE_TYPE.ORDER

    #
    @property
    def is_blocking(self):
        """
        This datasource is in a loop with the TradingSimulationClient
        """
        return False

    def open(self):
        qmsg.DataSource.open(self)
        self.order_socket = self.bind_order()

    def bind_order(self):
        return self.bind_pull_socket(self.addresses['order_address'])

    def do_work(self):

        #TODO: if this is the first iteration, break deadlock by sending a dummy order
        if(self.sent_count == 0):
            self.send(zp.namedict({}))

        #pull all orders from client.
        orders = []
        count = 0

        # TODO : this can be written in a concurrency agnostic
        # way... have a chat with Fawce about this ~Steve
        while True:

            (rlist, wlist, xlist) = select(
                [self.order_socket],
                [],
                [self.order_socket],
                #allow half the time of a heartbeat for the order
                #timeout, so we have time to signal we are done.
                timeout=self.heartbeat_timeout/2000
            )

            #no more orders, should this be an error condition?
            if len(rlist) == 0 or len(xlist) > 0:
                # no order message means there was a timeout above.
                # this is an indeterminant case, we don't know the cause.
                # the safest move is to break out of this loop and try again
                break

            order_msg = rlist[0].recv()

            if order_msg == str(zp.ORDER_PROTOCOL.DONE):
                self.signal_done()
                return

            if order_msg == str(zp.ORDER_PROTOCOL.BREAK):
                break

            order = zp.ORDER_UNFRAME(order_msg)
            #send the order along
            self.send(order)
            count += 1
            self.sent_count += 1

        #TODO: we have to send at least one dummy order per do_work iteration
        # or the feed will block waiting for our messages.
        if(count == 0):
            self.send(zp.namedict({}))

class TransactionSimulator(qmsg.BaseTransform):

    def __init__(self):
        qmsg.BaseTransform.__init__(self, zp.TRANSFORM_TYPE.TRANSACTION)
        self.open_orders                = {}
        self.order_count                = 0
        self.txn_count                  = 0
        self.trade_window                = datetime.timedelta(seconds=30)
        self.orderTTL                   = datetime.timedelta(days=1)
        self.volume_share               = 0.05
        self.commission                 = 0.03

    def transform(self, event):
        """
        Pulls one message from the event feed, then
        loops on orders until client sends DONE message.
        """
        if(event.type == zp.DATASOURCE_TYPE.ORDER):
            self.add_open_order(event)
            self.state['value'] = None
        elif(event.type == zp.DATASOURCE_TYPE.TRADE):
            txn = self.apply_trade_to_open_orders(event)
            self.state['value'] = txn
        else:
            self.state['value'] = None
            qutil.LOGGER.info("unexpected event type in transform: {etype}".format(etype=event.type))
        #TODO: what to do if we get another kind of datasource event.type?

        return self.state

    def add_open_order(self, event):
        """Orders are captured in a buffer by sid. No calculations are done here.
            Amount is explicitly converted to an int.
            Orders of amount zero are ignored.
        """
        event.amount = int(event.amount)
        if event.amount == 0:
            qutil.LOGGER.debug("requested to trade zero shares of {sid}".format(sid=event.sid))
            return

        self.order_count += 1

        if(not self.open_orders.has_key(event.sid)):
            self.open_orders[event.sid] = []
        self.open_orders[event.sid].append(event)

    def apply_trade_to_open_orders(self, event):

        if(event.volume == 0):
            #there are zero volume events bc some stocks trade
            #less frequently than once per minute.
            return self.create_dummy_txn(event.dt)

        if self.open_orders.has_key(event.sid):
            orders = self.open_orders[event.sid]
        else:
            return None

        remaining_orders = []
        total_order = 0
        dt = event.dt

        for order in orders:
            #we're using minute bars, so allow orders within
            #30 seconds of the trade
            if((order.dt - event.dt) < self.trade_window):
                total_order += order.amount
                if(order.dt > dt):
                    dt = order.dt
            #if the order still has time to live (TTL) keep track
            elif((self.algo_time - order.dt) < self.orderTTL):
                remaining_orders.append(order)

        self.open_orders[event.sid] = remaining_orders

        if(total_order != 0):
            direction = total_order / math.fabs(total_order)
        else:
            direction = 1

        volume_share = (direction * total_order) / event.volume
        if volume_share > .25:
            volume_share = .25
        amount = volume_share * event.volume * direction
        impact = (volume_share)**2 * .1 * direction * event.price
        return self.create_transaction(
            event.sid,
            amount,
            event.price + impact,
            dt.replace(tzinfo = pytz.utc),
            direction
        )

    def create_transaction(self, sid, amount, price, dt, direction):
        self.txn_count += 1
        txn = {'sid'                : sid,
                'amount'             : int(amount),
                'dt'                 : dt,
                'price'              : price,
                'commission'          : self.commission * amount * direction,
                'source_id'          : zp.FINANCE_COMPONENT.TRANSACTION_SIM
                }
        return zp.namedict(txn)

class TradingEnvironment(object):

    def __init__(
        self,
        benchmark_returns,
        treasury_curves,
        period_start    = None,
        period_end      = None,
        capital_base    = None,
        max_drawdown    = None
    ):

        self.trading_days = []
        self.trading_day_map = {}
        self.treasury_curves = treasury_curves
        self.benchmark_returns = benchmark_returns
        self.frame_index = ['sid', 'volume', 'dt', 'price', 'changed']
        self.period_start = period_start
        self.period_end = period_end
        self.capital_base = capital_base
        self.period_trading_days = None
        self.max_drawdown = max_drawdown

        for bm in benchmark_returns:
            self.trading_days.append(bm.date)
            self.trading_day_map[bm.date] = bm

        self.first_open = self.calculate_first_open()
        self.last_close = self.calculate_last_close()

    def calculate_first_open(self):
        """
        Finds the first trading day on or after self.period_start.
        """
        first_open  = self.period_start
        one_day      = datetime.timedelta(days=1)

        while not self.is_trading_day(first_open):
            first_open = first_open + one_day

        first_open = self.set_NYSE_time(first_open, 9, 30)
        return first_open

    def calculate_last_close(self):
        """
        Finds the last trading day on or before self.period_end
        """
        last_close  = self.period_end
        one_day     = datetime.timedelta(days=1)

        while not self.is_trading_day(last_close):
            last_close = last_close - one_day

        last_close = self.set_NYSE_time(last_close, 16, 00)

        return last_close

    #TODO: add other exchanges and timezones...
    def set_NYSE_time(self, dt, hour, minute):
        naive = datetime.datetime(
            year=dt.year,
            month=dt.month,
            day=dt.day
        )
        local = pytz.timezone ('US/Eastern')
        local_dt = naive.replace (tzinfo = local)
        # set the clock to the opening bell in NYC time.
        local_dt = local_dt.replace(hour=hour, minute=minute)
        # convert to UTC
        utc_dt = local_dt.astimezone (pytz.utc)
        return utc_dt

    def normalize_date(self, test_date):
        return datetime.datetime(
            year=test_date.year,
            month=test_date.month,
            day=test_date.day,
            tzinfo=pytz.utc
        )

    @property
    def days_in_period(self):
        """return the number of trading days within the period [start, end)"""
        assert(self.period_start != None)
        assert(self.period_end != None)

        if self.period_trading_days == None:
            self.period_trading_days = []
            for date in self.trading_days:
                if date > self.period_end:
                    break
                if date >= self.period_start:
                    self.period_trading_days.append(date)

        return len(self.period_trading_days)

    def is_trading_day(self, test_date):
        dt = self.normalize_date(test_date)
        return self.trading_day_map.has_key(dt)

    def get_benchmark_daily_return(self, test_date):
        date = self.normalize_date(test_date)
        if self.trading_day_map.has_key(date):
            return self.trading_day_map[date].returns
        else:
            return 0.0

    def add_to_frame(self, name):
        """
        Add an entry to the frame index.
        :param name: new index entry name. Used by TradingSimulationClient
        to
        """
        self.frame_index.append(name)

<commit_msg>Updated ordersource to be concurrency agnostic!<commit_after>import datetime
import pytz
import math
import pandas

# from gevent.select import select
from zmq.core.poll import select

import zipline.messaging as qmsg
import zipline.util as qutil
import zipline.protocol as zp
import zipline.finance.performance as perf

class TradeSimulationClient(qmsg.Component):

    def __init__(self, trading_environment):
        qmsg.Component.__init__(self)
        self.received_count         = 0
        self.prev_dt                = None
        self.event_queue            = None
        self.txn_count              = 0
        self.trading_environment    = trading_environment
        self.current_dt             = trading_environment.period_start
        self.last_iteration_dur     = datetime.timedelta(seconds=0)
        self.algorithm              = None

        assert self.trading_environment.frame_index != None
        self.event_frame = pandas.DataFrame(
            index=self.trading_environment.frame_index
        )

        self.perf = perf.PerformanceTracker(self.trading_environment)

    @property
    def get_id(self):
        return str(zp.FINANCE_COMPONENT.TRADING_CLIENT)

    def set_algorithm(self, algorithm):
        """
        :param algorithm: must implement the algorithm protocol. See
        :py:mod:`zipline.test.algorithm`
        """
        self.algorithm = algorithm
        #register the trading_client's order method with the algorithm
        self.algorithm.set_order(self.order)

    def open(self):
        self.result_feed = self.connect_result()
        self.order_socket = self.connect_order()

    def do_work(self):
        # poll all the sockets
        socks = dict(self.poll.poll(self.heartbeat_timeout))

        # see if the poller has results for the result_feed
        if self.result_feed in socks and \
            socks[self.result_feed] == self.zmq.POLLIN:

            # get the next message from the result feed
            msg = self.result_feed.recv()

            # if the feed is done, shut 'er down
            if msg == str(zp.CONTROL_PROTOCOL.DONE):
                qutil.LOGGER.info("Client is DONE!")
                # signal the performance tracker that the simulation has
                # ended. Perf will internally calculate the full risk report.
                self.perf.handle_simulation_end()
                # shutdown the feedback loop to the OrderDataSource
                self.signal_order_done()
                # signal Simulator, our ComponentHost, that this component is
                # done and Simulator needn't block exit on this component.
                self.signal_done()
                return

            # result_feed is a merge component, so unframe accordingly
            event = zp.MERGE_UNFRAME(msg)

            # update performance and relay the event to the algorithm
            self.process_event(event)

             # signal done to order source.
            self.order_socket.send(str(zp.ORDER_PROTOCOL.BREAK))

    def process_event(self, event):
        # track the number of transactions, for testing purposes.
        if(event.TRANSACTION != None):
            self.txn_count += 1

        #filter order flow out of the events sent to callbacks
        if event.source_id != zp.FINANCE_COMPONENT.ORDER_SOURCE:

            # the performance class needs to process each event, without
            # skipping. Algorithm should wait until the performance has been
            # updated, so that down stream components can safely assume that
            # performance is up to date. Note that this is done before we
            # mark the time for the algorithm's processing, thereby not
            # running the algo's clock for performance book keeping.
            self.perf.process_event(event)

            # mark the start time for client's processing of this event.
            event_start = datetime.datetime.utcnow()

            # queue the event.
            self.queue_event(event)

            # if the event is later than our current time, run the algo
            # otherwise, the algorithm has fallen behind the feed
            # and processing per event is longer than time between events.
            if event.dt >= self.current_dt:
                self.run_algorithm()

            # tally the time spent on this iteration
            self.last_iteration_dur = datetime.datetime.utcnow() - event_start
            # move the algorithm's clock forward to include iteration time
            self.current_dt = self.current_dt + self.last_iteration_dur

    def run_algorithm(self):
        """
        As per the algorithm protocol:

        - Set the current portfolio for the algorithm as per protocol.
        - Construct frame based on backlog of events, send to algorithm.
        """
        current_portfolio = self.perf.get_portfolio()
        self.algorithm.set_portfolio(current_portfolio)
        frame = self.get_frame()
        if len(frame) > 0:
            self.algorithm.handle_frame(frame)

    def connect_order(self):
        return self.connect_push_socket(self.addresses['order_address'])

    def order(self, sid, amount):
        order = zp.namedict({
            'dt':self.current_dt,
            'sid':sid,
            'amount':amount
        })

        self.order_socket.send(zp.ORDER_FRAME(order))

    def signal_order_done(self):
        self.order_socket.send(str(zp.ORDER_PROTOCOL.DONE))

    def queue_event(self, event):
        if self.event_queue == None:
            self.event_queue = []
        series = event.as_series()
        self.event_queue.append(series)

    def get_frame(self):
        for event in self.event_queue:
            self.event_frame[event['sid']] = event
        self.event_queue = []
        return self.event_frame

class OrderDataSource(qmsg.DataSource):
    """DataSource that relays orders from the client"""

    def __init__(self):
        """
        :param simulation_time: datetime in UTC timezone, sets the start
        time of simulation. orders
            will be timestamped relative to this datetime.
                event = {
                    'sid'    : an integer for security id,
                    'dt'     : datetime object,
                    'price'  : float for price,
                    'volume' : integer for volume
                }
        """
        qmsg.DataSource.__init__(self, zp.FINANCE_COMPONENT.ORDER_SOURCE)
        self.sent_count         = 0

    @property
    def get_type(self):
        return zp.DATASOURCE_TYPE.ORDER

    #
    @property
    def is_blocking(self):
        """
        This datasource is in a loop with the TradingSimulationClient
        """
        return False

    def open(self):
        qmsg.DataSource.open(self)
        self.order_socket = self.bind_order()

    def bind_order(self):
        return self.bind_pull_socket(self.addresses['order_address'])

    def do_work(self):

        #TODO: if this is the first iteration, break deadlock by sending a dummy order
        if(self.sent_count == 0):
            self.send(zp.namedict({}))

        #pull all orders from client.
        orders = []
        count = 0

        # TODO : this can be written in a concurrency agnostic
        # way... have a chat with Fawce about this ~Steve

        while True:
            # poll all the sockets
            # we reduce the timeout here by a factor of 2, because we need
            # to potentially receive the client's done message before the
            # controller or heartbeat times out.
            socks = dict(self.poll.poll(self.heartbeat_timeout/2))

            # see if the poller has results for the result_feed
            if self.order_socket in socks and \
                socks[self.order_socket] == self.zmq.POLLIN:

                order_msg = self.order_socket.recv()

                if order_msg == str(zp.ORDER_PROTOCOL.DONE):
                    self.signal_done()
                    return

                if order_msg == str(zp.ORDER_PROTOCOL.BREAK):
                    break

                order = zp.ORDER_UNFRAME(order_msg)
                #send the order along
                self.send(order)
                count += 1
                self.sent_count += 1

            else:
                # no orders, break out
                break

        #TODO: we have to send at least one dummy order per do_work iteration
        # or the feed will block waiting for our messages.
        if(count == 0):
            self.send(zp.namedict({}))

class TransactionSimulator(qmsg.BaseTransform):

    def __init__(self):
        qmsg.BaseTransform.__init__(self, zp.TRANSFORM_TYPE.TRANSACTION)
        self.open_orders                = {}
        self.order_count                = 0
        self.txn_count                  = 0
        self.trade_window                = datetime.timedelta(seconds=30)
        self.orderTTL                   = datetime.timedelta(days=1)
        self.volume_share               = 0.05
        self.commission                 = 0.03

    def transform(self, event):
        """
        Pulls one message from the event feed, then
        loops on orders until client sends DONE message.
        """
        if(event.type == zp.DATASOURCE_TYPE.ORDER):
            self.add_open_order(event)
            self.state['value'] = None
        elif(event.type == zp.DATASOURCE_TYPE.TRADE):
            txn = self.apply_trade_to_open_orders(event)
            self.state['value'] = txn
        else:
            self.state['value'] = None
            qutil.LOGGER.info("unexpected event type in transform: {etype}".format(etype=event.type))
        #TODO: what to do if we get another kind of datasource event.type?

        return self.state

    def add_open_order(self, event):
        """Orders are captured in a buffer by sid. No calculations are done here.
            Amount is explicitly converted to an int.
            Orders of amount zero are ignored.
        """
        event.amount = int(event.amount)
        if event.amount == 0:
            qutil.LOGGER.debug("requested to trade zero shares of {sid}".format(sid=event.sid))
            return

        self.order_count += 1

        if(not self.open_orders.has_key(event.sid)):
            self.open_orders[event.sid] = []
        self.open_orders[event.sid].append(event)

    def apply_trade_to_open_orders(self, event):

        if(event.volume == 0):
            #there are zero volume events bc some stocks trade
            #less frequently than once per minute.
            return self.create_dummy_txn(event.dt)

        if self.open_orders.has_key(event.sid):
            orders = self.open_orders[event.sid]
        else:
            return None

        remaining_orders = []
        total_order = 0
        dt = event.dt

        for order in orders:
            #we're using minute bars, so allow orders within
            #30 seconds of the trade
            if((order.dt - event.dt) < self.trade_window):
                total_order += order.amount
                if(order.dt > dt):
                    dt = order.dt
            #if the order still has time to live (TTL) keep track
            elif((self.algo_time - order.dt) < self.orderTTL):
                remaining_orders.append(order)

        self.open_orders[event.sid] = remaining_orders

        if(total_order != 0):
            direction = total_order / math.fabs(total_order)
        else:
            direction = 1

        volume_share = (direction * total_order) / event.volume
        if volume_share > .25:
            volume_share = .25
        amount = volume_share * event.volume * direction
        impact = (volume_share)**2 * .1 * direction * event.price
        return self.create_transaction(
            event.sid,
            amount,
            event.price + impact,
            dt.replace(tzinfo = pytz.utc),
            direction
        )

    def create_transaction(self, sid, amount, price, dt, direction):
        self.txn_count += 1
        txn = {'sid'                : sid,
                'amount'             : int(amount),
                'dt'                 : dt,
                'price'              : price,
                'commission'          : self.commission * amount * direction,
                'source_id'          : zp.FINANCE_COMPONENT.TRANSACTION_SIM
                }
        return zp.namedict(txn)

class TradingEnvironment(object):

    def __init__(
        self,
        benchmark_returns,
        treasury_curves,
        period_start    = None,
        period_end      = None,
        capital_base    = None,
        max_drawdown    = None
    ):

        self.trading_days = []
        self.trading_day_map = {}
        self.treasury_curves = treasury_curves
        self.benchmark_returns = benchmark_returns
        self.frame_index = ['sid', 'volume', 'dt', 'price', 'changed']
        self.period_start = period_start
        self.period_end = period_end
        self.capital_base = capital_base
        self.period_trading_days = None
        self.max_drawdown = max_drawdown

        for bm in benchmark_returns:
            self.trading_days.append(bm.date)
            self.trading_day_map[bm.date] = bm

        self.first_open = self.calculate_first_open()
        self.last_close = self.calculate_last_close()

    def calculate_first_open(self):
        """
        Finds the first trading day on or after self.period_start.
        """
        first_open  = self.period_start
        one_day      = datetime.timedelta(days=1)

        while not self.is_trading_day(first_open):
            first_open = first_open + one_day

        first_open = self.set_NYSE_time(first_open, 9, 30)
        return first_open

    def calculate_last_close(self):
        """
        Finds the last trading day on or before self.period_end
        """
        last_close  = self.period_end
        one_day     = datetime.timedelta(days=1)

        while not self.is_trading_day(last_close):
            last_close = last_close - one_day

        last_close = self.set_NYSE_time(last_close, 16, 00)

        return last_close

    #TODO: add other exchanges and timezones...
    def set_NYSE_time(self, dt, hour, minute):
        naive = datetime.datetime(
            year=dt.year,
            month=dt.month,
            day=dt.day
        )
        local = pytz.timezone ('US/Eastern')
        local_dt = naive.replace (tzinfo = local)
        # set the clock to the opening bell in NYC time.
        local_dt = local_dt.replace(hour=hour, minute=minute)
        # convert to UTC
        utc_dt = local_dt.astimezone (pytz.utc)
        return utc_dt

    def normalize_date(self, test_date):
        return datetime.datetime(
            year=test_date.year,
            month=test_date.month,
            day=test_date.day,
            tzinfo=pytz.utc
        )

    @property
    def days_in_period(self):
        """return the number of trading days within the period [start, end)"""
        assert(self.period_start != None)
        assert(self.period_end != None)

        if self.period_trading_days == None:
            self.period_trading_days = []
            for date in self.trading_days:
                if date > self.period_end:
                    break
                if date >= self.period_start:
                    self.period_trading_days.append(date)

        return len(self.period_trading_days)

    def is_trading_day(self, test_date):
        dt = self.normalize_date(test_date)
        return self.trading_day_map.has_key(dt)

    def get_benchmark_daily_return(self, test_date):
        date = self.normalize_date(test_date)
        if self.trading_day_map.has_key(date):
            return self.trading_day_map[date].returns
        else:
            return 0.0

    def add_to_frame(self, name):
        """
        Add an entry to the frame index.
        :param name: new index entry name. Used by TradingSimulationClient
        to
        """
        self.frame_index.append(name)

<|endoftext|>


<commit_before># coding=utf-8
"""
本模块为支持 `connection_keep_alive` 选项而存在
提供了一个线程安全的 keep-alive 链接池

requests的连接在每个session中是自动 keep-alive 的,
    在 `connection_keep_alive` 关闭时, 每次请求都会创建一个新的session,
        并发起一次新的请求, 则会带来相当大的连接开销(时间上)
    通过保持并复用 requests 的session, 可以极大地减少requests在请求远程服务器时的连接延迟

以前的版本是线程不安全, 当并发数大时会出现 ConnectionResetError
"""
from time import time
import requests
import threading

SESSION_TTL = 180  # 在清除过期session时, 会丢弃所有180秒未活动的session

# session池
pool = {
    "example.com": [
        # 每个域名下都有一堆session,
        # session的获取遵循 LIFO(后进先出) 原则,
        #    即优先获取最近使用过的 session
        # 这样可以增加 keep-alive 的存活几率
        {
            "domain": "example.com",
            "session": requests.Session(),
            "active": time(),
        },
    ],
}

locked_session = threading.local()  # 这是一个 thread-local 变量

def get_session(domain):
    """
    获取一个此域名的 keep-alive 的session
    :param domain: 域名
    :type domain: str
    :rtype: requests.Session
    """
    if domain not in pool:
        pool[domain] = []

    if not hasattr(locked_session, "session"):
        # 这个变量用于存储本线程中被锁定的session
        # 当一个session被拿出来使用时, 会从 pool 中被移除, 加入到下面这个变量中
        # 当线程结束后, 需要调用 release_lock() 来释放被锁定的session
        #    此时被锁定的session会重新进入session池
        locked_session.session = []

    if not pool[domain]:
        # 线程池空, 新建一个 session
        session = {
            "domain": domain,
            "session": requests.Session(),
        }
    else:
        # 从线程池中取出最近的一个
        session = pool[domain].pop()

    session["active"] = time()

    locked_session.session.append(session)

    return session["session"]

def release_lock():
    for session in locked_session.session:  # type: dict
        pool[session["domain"]].append(session)

def clear(force_flush=False):
    if force_flush:
        pool.clear()
    else:
        for domain in pool.keys():
            pool[domain] = [s for s in pool[domain] if s["active"] > time() - SESSION_TTL]
<commit_msg>[/] 修复另一处 locked_session.session 未初始化的bug<commit_after># coding=utf-8
"""
本模块为支持 `connection_keep_alive` 选项而存在
提供了一个线程安全的 keep-alive 链接池

requests的连接在每个session中是自动 keep-alive 的,
    在 `connection_keep_alive` 关闭时, 每次请求都会创建一个新的session,
        并发起一次新的请求, 则会带来相当大的连接开销(时间上)
    通过保持并复用 requests 的session, 可以极大地减少requests在请求远程服务器时的连接延迟

以前的版本是线程不安全, 当并发数大时会出现 ConnectionResetError
"""
from time import time
import requests
import threading

SESSION_TTL = 180  # 在清除过期session时, 会丢弃所有180秒未活动的session

# session池
pool = {
    "example.com": [
        # 每个域名下都有一堆session,
        # session的获取遵循 LIFO(后进先出) 原则,
        #    即优先获取最近使用过的 session
        # 这样可以增加 keep-alive 的存活几率
        {
            "domain": "example.com",
            "session": requests.Session(),
            "active": time(),
        },
    ],
}

locked_session = threading.local()  # 这是一个 thread-local 变量

def get_session(domain):
    """
    获取一个此域名的 keep-alive 的session
    :param domain: 域名
    :type domain: str
    :rtype: requests.Session
    """
    if domain not in pool:
        pool[domain] = []

    if not hasattr(locked_session, "session"):
        # 这个变量用于存储本线程中被锁定的session
        # 当一个session被拿出来使用时, 会从 pool 中被移除, 加入到下面这个变量中
        # 当线程结束后, 需要调用 release_lock() 来释放被锁定的session
        #    此时被锁定的session会重新进入session池
        locked_session.session = []

    if not pool[domain]:
        # 线程池空, 新建一个 session
        session = {
            "domain": domain,
            "session": requests.Session(),
        }
    else:
        # 从线程池中取出最近的一个
        session = pool[domain].pop()

    session["active"] = time()

    locked_session.session.append(session)

    return session["session"]

def release_lock():
    if not hasattr(locked_session, "session"):
        return
    for session in locked_session.session:  # type: dict
        pool[session["domain"]].append(session)

def clear(force_flush=False):
    if force_flush:
        pool.clear()
    else:
        for domain in pool.keys():
            pool[domain] = [s for s in pool[domain] if s["active"] > time() - SESSION_TTL]
<|endoftext|>