<commit_before>from __future__ import absolute_import
from __future__ import unicode_literals

from app import r as redis

import collections

class Index:
    def __init__(self, key):
        if not hasattr(self, '__prefix__'):
            raise AttributeError("Inheriting classes must define the attribute: __prefix__")

        id = key
        if not key.startswith(self.__prefix__):
            id = self.__prefix__ + key
        else:
            key = key[len(self.__prefix__):]

        self.id = id
        self.key = key

    def get(self):
        value = redis.get(self.id)

        if value and hasattr(self, '__relation__'):
            # Create an object of the specified
            value = self.__relation__(value.decode())

        return value

    def set(self, value):
        if hasattr(value, 'id'):
            # If the valueue has an id, use the id
            value = value.id

        return redis.set(self.id, value)

    def delete(self):
        return redis.delete(self.id)

    @classmethod
    def deleteall(cls):
        if not hasattr(cls, '__prefix__'):
            raise AttributeError("Inheriting classes must define the attribute: __prefix__")

        keys = redis.keys(cls.__prefix__ + '*')
        if len(keys) > 0:
            return redis.delete(*keys)

    @classmethod
    def all(cls, *args, **kwargs):
        if not hasattr(cls, '__prefix__'):
            raise AttributeError("Inheriting classes must define the attribute: __prefix__")

        keys = redis.keys(cls.__prefix__ + '*')
        for key in keys:
            yield cls(key.decode(), *args, **kwargs)

    @classmethod
    def items(cls, *args, **kwargs):
        """Yield the key,object pairs in the index"""
        for index in cls.all(*args, **kwargs):
            yield index.key, index.get()

    @classmethod
    def contains(cls, key):
        if not hasattr(cls, '__prefix__'):
            raise AttributeError("Inheriting classes must define the attribute: __prefix__")

        if not key.startswith(cls.__prefix__):
            key = cls.__prefix__ + key

        return redis.exists(key)

class Model(collections.MutableMapping):
    """Defines a basic model for data storage using HASH inside
    redis.

    Implementing classes must define a __prefix__ attribute, to group keys of the
    defined model.
    """

    def __init__(self, id):
        if not hasattr(self, '__prefix__'):
            raise AttributeError("Inheriting classes must define the attribute: __prefix__")

        if not id.startswith(self.__prefix__):
            id = self.__prefix__ + id

        self.id = id

    def __getattr__(self, name):
        return self.__getitem__(name)

    def __setattr__(self, name, value):
        """Set the attribute with the value given by 'value'

        It works by calling self.__setitem__ with a couple of exceptions
        if the attribute name is 'id', it can only be set if it has not been set
        before.

        If a property setter has been defined for the inheriting class with the same
        name as the key, then the method calls the property setter first.

        Property setters must take care to set the actual value using self[name] to
        avoid infinite loops
        """
        if name == 'id':
            if name in self.__dict__:
                # Cannot change id (TODO: use redis.rename()?)
                raise AttributeError("The model id cannot be changed")
            else:
                self.__dict__[name] = value
                return True

        cls = self.__class__

        # Check if there are property setter.
        # This only works if the property setter has the same name as the property
        if hasattr(cls, name) and \
                isinstance(cls.__dict__[name], property):
            try:
                cls.__dict__[name].__set__(self, value)
                return True
            except AttributeError:
                # No setter with the same name defined
                pass

        return self.__setitem__(name, value)

    def __getitem__(self, key):
        value = redis.hget(self.id, self.__keytransform__(key))

        # If the value is none
        if not value:
            return None

        # Try to parse numeric types
        parsers = [int, float]
        for parse in parsers:
            try:
                return parse(value)
            except ValueError:
                pass

        return value.decode('utf-8')

    def __setitem__(self, key, value):
        return redis.hset(self.id, self.__keytransform__(key), value)

    def __delitem__(self, key):
        redis.hdel(self.id, self.__keytransform__(key))

    def __iter__(self):
        return iter(redis.hgetall(self.id))

    def __len__(self):
        return redis.hlen(self.id)

    def __keytransform__(self, key):
        return key

    def __repr__(self):
        return repr(redis.hgetall(self.id))

    def __contains__(self, key):
        return redis.hexists(self.id, self.__keytransform__(key))

    def incrby(self, key, amount=1):
        """Increment the provided key in the dictionary by the specified amount"""
        redis.hincrby(self.id, self.__keytransform__(key), amount)

    def delete(self):
        """Delete the entity from the database"""
        return redis.delete(self.id)

    @classmethod
    def exists(cls, id):
        if not hasattr(cls, '__prefix__'):
            raise AttributeError("Inheriting classes must define the attribute: __prefix__")

        if not id.startswith(cls.__prefix__):
            id = cls.__prefix__ + id

        return redis.exists(id)

    @classmethod
    def all(cls, *args, **kwargs):
        """Return an iterator over the objects matching the model in the database

        It performs the query by calling redis.keys() using the defined prefix.
        If extra arguments are given, they are passed to the contructor of the model.
        """
        if not hasattr(cls, '__prefix__'):
            raise AttributeError("Inheriting classes must define the attribute: __prefix__")

        keys = redis.keys(cls.__prefix__ + '*')
        for key in keys:
            yield cls(key.decode(), *args, **kwargs)

    @classmethod
    def deleteall(cls):
        if not hasattr(cls, '__prefix__'):
            raise AttributeError("Inheriting classes must define the attribute: __prefix__")

        keys = redis.keys(cls.__prefix__ + '*')
        if len(keys) > 0:
            return redis.delete(*keys)
<commit_msg>Return iterator over the model keys for consistency with MutableMapping<commit_after>from __future__ import absolute_import
from __future__ import unicode_literals

from app import r as redis

import collections

class Index:
    def __init__(self, key):
        if not hasattr(self, '__prefix__'):
            raise AttributeError("Inheriting classes must define the attribute: __prefix__")

        id = key
        if not key.startswith(self.__prefix__):
            id = self.__prefix__ + key
        else:
            key = key[len(self.__prefix__):]

        self.id = id
        self.key = key

    def get(self):
        value = redis.get(self.id)

        if value and hasattr(self, '__relation__'):
            # Create an object of the specified
            value = self.__relation__(value.decode())

        return value

    def set(self, value):
        if hasattr(value, 'id'):
            # If the valueue has an id, use the id
            value = value.id

        return redis.set(self.id, value)

    def delete(self):
        return redis.delete(self.id)

    @classmethod
    def deleteall(cls):
        if not hasattr(cls, '__prefix__'):
            raise AttributeError("Inheriting classes must define the attribute: __prefix__")

        keys = redis.keys(cls.__prefix__ + '*')
        if len(keys) > 0:
            return redis.delete(*keys)

    @classmethod
    def all(cls, *args, **kwargs):
        if not hasattr(cls, '__prefix__'):
            raise AttributeError("Inheriting classes must define the attribute: __prefix__")

        keys = redis.keys(cls.__prefix__ + '*')
        for key in keys:
            yield cls(key.decode(), *args, **kwargs)

    @classmethod
    def items(cls, *args, **kwargs):
        """Yield the key,object pairs in the index"""
        for index in cls.all(*args, **kwargs):
            yield index.key, index.get()

    @classmethod
    def contains(cls, key):
        if not hasattr(cls, '__prefix__'):
            raise AttributeError("Inheriting classes must define the attribute: __prefix__")

        if not key.startswith(cls.__prefix__):
            key = cls.__prefix__ + key

        return redis.exists(key)

class Model(collections.MutableMapping):
    """Defines a basic model for data storage using HASH inside
    redis.

    Implementing classes must define a __prefix__ attribute, to group keys of the
    defined model.
    """

    def __init__(self, id):
        if not hasattr(self, '__prefix__'):
            raise AttributeError("Inheriting classes must define the attribute: __prefix__")

        if not id.startswith(self.__prefix__):
            id = self.__prefix__ + id

        self.id = id

    def __getattr__(self, name):
        return self.__getitem__(name)

    def __setattr__(self, name, value):
        """Set the attribute with the value given by 'value'

        It works by calling self.__setitem__ with a couple of exceptions
        if the attribute name is 'id', it can only be set if it has not been set
        before.

        If a property setter has been defined for the inheriting class with the same
        name as the key, then the method calls the property setter first.

        Property setters must take care to set the actual value using self[name] to
        avoid infinite loops
        """
        if name == 'id':
            if name in self.__dict__:
                # Cannot change id (TODO: use redis.rename()?)
                raise AttributeError("The model id cannot be changed")
            else:
                self.__dict__[name] = value
                return True

        cls = self.__class__

        # Check if there are property setter.
        # This only works if the property setter has the same name as the property
        if hasattr(cls, name) and \
                isinstance(cls.__dict__[name], property):
            try:
                cls.__dict__[name].__set__(self, value)
                return True
            except AttributeError:
                # No setter with the same name defined
                pass

        return self.__setitem__(name, value)

    def __getitem__(self, key):
        value = redis.hget(self.id, self.__keytransform__(key))

        # If the value is none
        if not value:
            return None

        # Try to parse numeric types
        parsers = [int, float]
        for parse in parsers:
            try:
                return parse(value)
            except ValueError:
                pass

        return value.decode('utf-8')

    def __setitem__(self, key, value):
        return redis.hset(self.id, self.__keytransform__(key), value)

    def __delitem__(self, key):
        redis.hdel(self.id, self.__keytransform__(key))

    def __iter__(self):
        return iter(redis.hkeys(self.id))

    def __len__(self):
        return redis.hlen(self.id)

    def __keytransform__(self, key):
        return key

    def __repr__(self):
        return repr(redis.hgetall(self.id))

    def __contains__(self, key):
        return redis.hexists(self.id, self.__keytransform__(key))

    def incrby(self, key, amount=1):
        """Increment the provided key in the dictionary by the specified amount"""
        redis.hincrby(self.id, self.__keytransform__(key), amount)

    def delete(self):
        """Delete the entity from the database"""
        return redis.delete(self.id)

    @classmethod
    def exists(cls, id):
        if not hasattr(cls, '__prefix__'):
            raise AttributeError("Inheriting classes must define the attribute: __prefix__")

        if not id.startswith(cls.__prefix__):
            id = cls.__prefix__ + id

        return redis.exists(id)

    @classmethod
    def all(cls, *args, **kwargs):
        """Return an iterator over the objects matching the model in the database

        It performs the query by calling redis.keys() using the defined prefix.
        If extra arguments are given, they are passed to the contructor of the model.
        """
        if not hasattr(cls, '__prefix__'):
            raise AttributeError("Inheriting classes must define the attribute: __prefix__")

        keys = redis.keys(cls.__prefix__ + '*')
        for key in keys:
            yield cls(key.decode(), *args, **kwargs)

    @classmethod
    def deleteall(cls):
        if not hasattr(cls, '__prefix__'):
            raise AttributeError("Inheriting classes must define the attribute: __prefix__")

        keys = redis.keys(cls.__prefix__ + '*')
        if len(keys) > 0:
            return redis.delete(*keys)
<|endoftext|>


<commit_before>#!/usr/bin/env python2
# -*- coding: UTF-8 -*-

# use: python savvy.py -i samples/ANLJ01.1.fsa_nt

# S:erotype
# A:ntimicrobial Resistance
# V:irulence Factors
# -vy

import logging
# long function calls, cause datastruct_savvy is important
import datastruct_savvy

from rdflib import Graph
from turtle_utils import generate_uri as gu
from turtle_grapher import generate_output, generate_graph, generate_turtle_skeleton, generate_file_output

from os.path import basename

# bruteforce
from insert import upload_graph

def call_ectyper(graph, args_dict):
    # i don't intend to import anything from ECTyper (there are a lot of
    # imports in it - not sure if we'll use them all)
    import subprocess

    from rdflib import Literal
    from ast import literal_eval
    from os.path import splitext

    ectyper_dict = {}
    #logging.info('calling ectyper from fun call_ectyper')
    # concurrency is handled at the batch level, not here (note: this might change)
    # we only use ectyper for serotyping and vf, amr is handled by rgi directly
    if not args_dict['disable_serotype'] or not args_dict['disable_vf']:
        ectyper_dict = subprocess.check_output(['./app/ecoli_serotyping/src/Tools_Controller/tools_controller.py',
                                                '-in', args_dict['i'],
                                                '-s', str(
                                                    int(not args_dict['disable_serotype'])),
                                                '-vf', str(
                                                    int(not args_dict['disable_vf'])),
                                                '-pi', args_dict['pi']
                                                ])
        #logging.info('inner call completed')

        # because we are using check_output, this catches any print messages from tools_controller
        # TODO: switch to pipes
        if 'error' in ectyper_dict.lower():
            #logging.error('ectyper failed for' + args_dict['i'])
            print 'ECTyper failed for: ', args_dict['i']
            print 'returning graph w/o serotype'
            return graph

        #logging.info('evalulating ectyper output')
        # generating the dict
        ectyper_dict = literal_eval(ectyper_dict)
        # logging.info(ectyper_dict)
        #logging.info('evaluation okay')

        # TODO: edit ectyper sure were not using this ducktape approach
        # we are calling tools_controller on only one file, so grab that dict
        key, ectyper_dict = ectyper_dict.popitem()

        if not args_dict['disable_serotype']:
            # serotype parsing
            #logging.info('parsing Serotype')
            graph = datastruct_savvy.parse_serotype(
                graph, ectyper_dict['Serotype'], args_dict['uriIsolate'])
            #logging.info('serotype parsed okay')

        if not args_dict['disable_vf']:
            # vf
            #logging.info('parsing vf')
            graph = datastruct_savvy.parse_gene_dict(
                graph, ectyper_dict['Virulence Factors'], args_dict['uriGenome'])
            #logging.info('vf parsed okay')

    if not args_dict['disable_amr']:
        # amr
        #logging.info('generating amr')
        amr_result = generate_amr(
            graph, args_dict['uriGenome'], args_dict['i'])
        graph = amr_result['graph']
        ectyper_dict['Antimicrobial Resistance'] = amr_result['amr_dict']
        #logging.info('amr generation okay')

    return {'graph': graph, 'ectyper_dict': ectyper_dict}

def generate_amr(graph, uriGenome, fasta_file):
    import subprocess
    import pandas

    from os import rename
    from rdflib import BNode, Literal

    if '/' in fasta_file:
        outputname = fasta_file.split('/')[-1]
    else:
        outputname = fasta_file

    # differs from ectyper as we dont care about the temp results, just the final .tsv
    # direct (the main) call
    subprocess.call(['rgi',
                     '-i', fasta_file,
                     '-o', 'outputs/' + outputname])

    # the rgi_json call in rgitool.py isn't needed for us
    # this generates the .tsv we want
    subprocess.call(['rgi_jsontab',
                     '-i', 'outputs/' + outputname + '.json',
                     '-o', 'outputs/' + outputname])

    rename('outputs/' + outputname + '.txt', 'outputs/' + outputname + '.tsv')

    amr_results = pandas.read_table('outputs/' + outputname + '.tsv')
    amr_results = amr_results[
        ['ORF_ID', 'START', 'STOP', 'ORIENTATION', 'CUT_OFF', 'Best_Hit_ARO']]

    amr_results.rename(
        columns={'ORF_ID': 'contig_id', 'Best_Hit_ARO': 'GENE_NAME'}, inplace=True)

    # sometimes there are spaces at the end of the contig id, also we remove
    # the additional occurance tag that RGI adds to contig ids
    amr_results['contig_id'] = amr_results['contig_id'].apply(
        lambda n: n.strip().rsplit('_', 1)[0])

    # note: you might be tempted to prefix a set_index('contig_id') but
    # remember, the same contig might have multiple genes
    amr_results = amr_results.to_dict(orient='index')

    # we have to manually check for contigs with multiple genes
    # TODO: write something less horrendously slow and memory consuming
    amr_dict = {}
    for i in amr_results.keys():
        contig_id = amr_results[i]['contig_id']
        if contig_id not in amr_dict.keys():
            amr_dict[contig_id] = []
        amr_dict[contig_id].append(dict((keys, amr_results[i][keys]) for keys in (
            'START', 'STOP', 'GENE_NAME', 'ORIENTATION', 'CUT_OFF', 'GENE_NAME')))
    # wipe the amr_results early
    amr_results = None

    graph = datastruct_savvy.parse_gene_dict(graph, amr_dict, uriGenome)

    return {'graph': graph, 'amr_dict': amr_dict}

def json_return(args_dict, gene_dict):
    json_r = []

    # strip gene_dicts that user doesn't want to see
    # remember, we want to run all analysis on our end so we have that data in blazegraph
    for analysis in gene_dict:
        if analysis == 'Serotype' and args_dict['disable_serotype']:
            del gene_dict['Serotype']
        if analysis == 'Antimicrobial Resistance' and args_dict['disable_amr']:
            del gene_dict['Antimicrobial Resistance']
        if analysis == 'Virulence Factors' and args_dict['disable_vf']:
            del gene_dict['Virulence Factors']

    for analysis in gene_dict:
        if analysis == 'Serotype':
            instance_dict = {}
            instance_dict['filename'] = basename(args_dict['i'])[27:]
            instance_dict['hitname'] = str(gene_dict[analysis].values()).replace(',', ' ').strip("'").strip("[").strip("]")
            instance_dict['contigid'] = 'n/a'
            instance_dict['analysis'] = analysis
            instance_dict['hitorientation'] = 'n/a'
            instance_dict['hitstart'] = 'n/a'
            instance_dict['hitstop'] = 'n/a'
            instance_dict['hitcutoff'] = 'n/a'
            json_r.append(instance_dict)
        else:
            for contig_id in gene_dict[analysis]:
                # where gene_results is a list for amr/vf
                for item in gene_dict[analysis][contig_id]:
                    # for w/e reason vf, has a '0' int in the list of dicts
                    # TODO: bug fix^
                    if type(item) is dict:
                        instance_dict = {}
                        instance_dict['filename'] = basename(args_dict['i'])[27:]
                        instance_dict['contigid'] = contig_id
                        instance_dict['analysis'] = analysis
                        instance_dict['hitname'] = item['GENE_NAME']
                        instance_dict['hitorientation'] = item['ORIENTATION']
                        instance_dict['hitstart'] = item['START']
                        instance_dict['hitstop'] = item['STOP']
                        if analysis == 'Antimicrobial Resistance':
                            instance_dict['hitcutoff'] = item['CUT_OFF']
                        else:
                            instance_dict['hitcutoff'] = 'n/a'
                        json_r.append(instance_dict)
    return json_r

def savvy(args_dict):
    '''
    Args:
        args_dict(dict): i prefer working with args in a dictionary, rather than a namespace is all

    Returns:
        (rdflib.Graph): a graph object with the VF/AMR/Serotype added to it via ECTyper/RGI
    '''
    from os import remove  # for batch cleanup
    # starting #logging
    '''
    logging.basicConfig(
        filename='outputs/' + __name__ +
        args_dict['i'].split('/')[-1] + '.log',
        level=logging.INFO
    )
    '''

    print("Importing FASTA from: " + args_dict['i'])
    #logging.info('importing from' + args_dict['i'])

    # setting up graph
    graph = generate_graph()

    ectyper_result={}

    #logging.info('generating barebones ttl from file')
    graph = generate_turtle_skeleton(
        graph, args_dict['i'], args_dict['uriIsolate'], args_dict['uriGenome'])
    #logging.info('barebones ttl generated')

    if not (args_dict['disable_serotype'] and args_dict['disable_vf'] and args_dict['disable_amr']):
        #logging.info('calling ectyper')
        ectyper_result = call_ectyper(graph, args_dict)
        graph = ectyper_result['graph']
        #logging.info('ectyper call completed')

    # individual fasta logs are wiped on completion (or you'd have several
    # thousand of these)
    #remove('outputs/' + __name__ + args_dict['i'].split('/')[-1] + '.log')
    print upload_graph(graph)
    return json_return(args_dict, ectyper_result['ectyper_dict'])

if __name__ == "__main__":
    import argparse
    import os  # for batch cleanup

    from ConfigParser import SafeConfigParser
    from turtle_utils import generate_hash

    # parsing cli-input
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-i",
        help="FASTA file",
        required=True
    )
    parser.add_argument(
        "--disable-serotype",
        help="Disables use of the Serotyper. Serotyper is triggered by default.",
        action="store_true"
    )
    parser.add_argument(
        "--disable-vf",
        help="Disables use of ECTyper to get associated Virulence Factors. VFs are computed by default.",
        action="store_true"
    )
    parser.add_argument(
        "--disable-amr",
        help="Disables use of RGI to get Antimicrobial Resistance Factors.  AMR genes are computed by default.",
        action="store_true"
    )
    # note: by in large, we expect uri to be given as just the unique string
    # value  (be it the hash or the integer) without any prefixes, the actual
    # rdflib.URIRef object will be generated in this script
    # this is mainly for batch computation
    parser.add_argument(
        "--uriGenome",
        help="Allows the specification of the Genome URI separately. Expect just the hash (not an actual uri).",
    )
    # This is both for batch computation and for future extensions where there
    # are multiple sequencings per isolate (Campy)
    parser.add_argument(
        "--uriIsolate",
        help="Allows the specification of the Isolate URI separately. Expect just the integer (not the full :spfyID)",
        type=int
    )
    args = parser.parse_args()
    # we make a dictionary from the cli-inputs and add are uris to it
    # mainly used for when a func needs a lot of the args
    args_dict = vars(args)

    # starting#logging
    '''
    logging.basicConfig(
        filename='outputs/' + __name__ +
        args_dict['i'].split('/')[-1] + '.log',
        level=logging.INFO
    )
    '''

    # check if a genome uri isn't set yet
    if args_dict['uriIsolate'] is None:
        # this is temporary, TODO: include a spqarql query to the db
        uriIsolate = gu(':spfy' + str(hash(args_dict['i'].split('/')[-1])))
    else:
        uriIsolate = gu(':spfy' + args_dict['uriIsolate'])

    # if the fasta_file hash was not precomputed (batch scripts should
    # precompute it), we compute that now
    if args_dict['uriGenome'] is None:
        uriGenome = gu(':' + generate_hash(args_dict['i']))
    else:
        uriGenome = gu(':' + args_dict['uriGenome'])

    args_dict['uriIsolate'] = uriIsolate
    args_dict['uriGenome'] = uriGenome

    print savvy(args_dict)
<commit_msg>bug fix<commit_after>#!/usr/bin/env python2
# -*- coding: UTF-8 -*-

# use: python savvy.py -i samples/ANLJ01.1.fsa_nt

# S:erotype
# A:ntimicrobial Resistance
# V:irulence Factors
# -vy

import logging
# long function calls, cause datastruct_savvy is important
import datastruct_savvy

from rdflib import Graph
from turtle_utils import generate_uri as gu
from turtle_grapher import generate_output, generate_graph, generate_turtle_skeleton, generate_file_output

from os.path import basename

# bruteforce
from insert import upload_graph

def call_ectyper(graph, args_dict):
    # i don't intend to import anything from ECTyper (there are a lot of
    # imports in it - not sure if we'll use them all)
    import subprocess

    from rdflib import Literal
    from ast import literal_eval
    from os.path import splitext

    ectyper_dict = {}
    #logging.info('calling ectyper from fun call_ectyper')
    # concurrency is handled at the batch level, not here (note: this might change)
    # we only use ectyper for serotyping and vf, amr is handled by rgi directly
    if not args_dict['disable_serotype'] or not args_dict['disable_vf']:
        ectyper_dict = subprocess.check_output(['./app/ecoli_serotyping/src/Tools_Controller/tools_controller.py',
                                                '-in', args_dict['i'],
                                                '-s', str(
                                                    int(not args_dict['disable_serotype'])),
                                                '-vf', str(
                                                    int(not args_dict['disable_vf'])),
                                                '-pi', str(args_dict['pi'])
                                                ])
        #logging.info('inner call completed')

        # because we are using check_output, this catches any print messages from tools_controller
        # TODO: switch to pipes
        if 'error' in ectyper_dict.lower():
            #logging.error('ectyper failed for' + args_dict['i'])
            print 'ECTyper failed for: ', args_dict['i']
            print 'returning graph w/o serotype'
            return graph

        #logging.info('evalulating ectyper output')
        # generating the dict
        ectyper_dict = literal_eval(ectyper_dict)
        # logging.info(ectyper_dict)
        #logging.info('evaluation okay')

        # TODO: edit ectyper sure were not using this ducktape approach
        # we are calling tools_controller on only one file, so grab that dict
        key, ectyper_dict = ectyper_dict.popitem()

        if not args_dict['disable_serotype']:
            # serotype parsing
            #logging.info('parsing Serotype')
            graph = datastruct_savvy.parse_serotype(
                graph, ectyper_dict['Serotype'], args_dict['uriIsolate'])
            #logging.info('serotype parsed okay')

        if not args_dict['disable_vf']:
            # vf
            #logging.info('parsing vf')
            graph = datastruct_savvy.parse_gene_dict(
                graph, ectyper_dict['Virulence Factors'], args_dict['uriGenome'])
            #logging.info('vf parsed okay')

    if not args_dict['disable_amr']:
        # amr
        #logging.info('generating amr')
        amr_result = generate_amr(
            graph, args_dict['uriGenome'], args_dict['i'])
        graph = amr_result['graph']
        ectyper_dict['Antimicrobial Resistance'] = amr_result['amr_dict']
        #logging.info('amr generation okay')

    return {'graph': graph, 'ectyper_dict': ectyper_dict}

def generate_amr(graph, uriGenome, fasta_file):
    import subprocess
    import pandas

    from os import rename
    from rdflib import BNode, Literal

    if '/' in fasta_file:
        outputname = fasta_file.split('/')[-1]
    else:
        outputname = fasta_file

    # differs from ectyper as we dont care about the temp results, just the final .tsv
    # direct (the main) call
    subprocess.call(['rgi',
                     '-i', fasta_file,
                     '-o', 'outputs/' + outputname])

    # the rgi_json call in rgitool.py isn't needed for us
    # this generates the .tsv we want
    subprocess.call(['rgi_jsontab',
                     '-i', 'outputs/' + outputname + '.json',
                     '-o', 'outputs/' + outputname])

    rename('outputs/' + outputname + '.txt', 'outputs/' + outputname + '.tsv')

    amr_results = pandas.read_table('outputs/' + outputname + '.tsv')
    amr_results = amr_results[
        ['ORF_ID', 'START', 'STOP', 'ORIENTATION', 'CUT_OFF', 'Best_Hit_ARO']]

    amr_results.rename(
        columns={'ORF_ID': 'contig_id', 'Best_Hit_ARO': 'GENE_NAME'}, inplace=True)

    # sometimes there are spaces at the end of the contig id, also we remove
    # the additional occurance tag that RGI adds to contig ids
    amr_results['contig_id'] = amr_results['contig_id'].apply(
        lambda n: n.strip().rsplit('_', 1)[0])

    # note: you might be tempted to prefix a set_index('contig_id') but
    # remember, the same contig might have multiple genes
    amr_results = amr_results.to_dict(orient='index')

    # we have to manually check for contigs with multiple genes
    # TODO: write something less horrendously slow and memory consuming
    amr_dict = {}
    for i in amr_results.keys():
        contig_id = amr_results[i]['contig_id']
        if contig_id not in amr_dict.keys():
            amr_dict[contig_id] = []
        amr_dict[contig_id].append(dict((keys, amr_results[i][keys]) for keys in (
            'START', 'STOP', 'GENE_NAME', 'ORIENTATION', 'CUT_OFF', 'GENE_NAME')))
    # wipe the amr_results early
    amr_results = None

    graph = datastruct_savvy.parse_gene_dict(graph, amr_dict, uriGenome)

    return {'graph': graph, 'amr_dict': amr_dict}

def json_return(args_dict, gene_dict):
    json_r = []

    # strip gene_dicts that user doesn't want to see
    # remember, we want to run all analysis on our end so we have that data in blazegraph
    for analysis in gene_dict:
        if analysis == 'Serotype' and args_dict['disable_serotype']:
            del gene_dict['Serotype']
        if analysis == 'Antimicrobial Resistance' and args_dict['disable_amr']:
            del gene_dict['Antimicrobial Resistance']
        if analysis == 'Virulence Factors' and args_dict['disable_vf']:
            del gene_dict['Virulence Factors']

    for analysis in gene_dict:
        if analysis == 'Serotype':
            instance_dict = {}
            instance_dict['filename'] = basename(args_dict['i'])[27:]
            instance_dict['hitname'] = str(gene_dict[analysis].values()).replace(',', ' ').strip("'").strip("[").strip("]")
            instance_dict['contigid'] = 'n/a'
            instance_dict['analysis'] = analysis
            instance_dict['hitorientation'] = 'n/a'
            instance_dict['hitstart'] = 'n/a'
            instance_dict['hitstop'] = 'n/a'
            instance_dict['hitcutoff'] = 'n/a'
            json_r.append(instance_dict)
        else:
            for contig_id in gene_dict[analysis]:
                # where gene_results is a list for amr/vf
                for item in gene_dict[analysis][contig_id]:
                    # for w/e reason vf, has a '0' int in the list of dicts
                    # TODO: bug fix^
                    if type(item) is dict:
                        instance_dict = {}
                        instance_dict['filename'] = basename(args_dict['i'])[27:]
                        instance_dict['contigid'] = contig_id
                        instance_dict['analysis'] = analysis
                        instance_dict['hitname'] = item['GENE_NAME']
                        instance_dict['hitorientation'] = item['ORIENTATION']
                        instance_dict['hitstart'] = item['START']
                        instance_dict['hitstop'] = item['STOP']
                        if analysis == 'Antimicrobial Resistance':
                            instance_dict['hitcutoff'] = item['CUT_OFF']
                        else:
                            instance_dict['hitcutoff'] = 'n/a'
                        json_r.append(instance_dict)
    return json_r

def savvy(args_dict):
    '''
    Args:
        args_dict(dict): i prefer working with args in a dictionary, rather than a namespace is all

    Returns:
        (rdflib.Graph): a graph object with the VF/AMR/Serotype added to it via ECTyper/RGI
    '''
    from os import remove  # for batch cleanup
    # starting #logging
    '''
    logging.basicConfig(
        filename='outputs/' + __name__ +
        args_dict['i'].split('/')[-1] + '.log',
        level=logging.INFO
    )
    '''

    print("Importing FASTA from: " + args_dict['i'])
    #logging.info('importing from' + args_dict['i'])

    # setting up graph
    graph = generate_graph()

    ectyper_result={}

    #logging.info('generating barebones ttl from file')
    graph = generate_turtle_skeleton(
        graph, args_dict['i'], args_dict['uriIsolate'], args_dict['uriGenome'])
    #logging.info('barebones ttl generated')

    if not (args_dict['disable_serotype'] and args_dict['disable_vf'] and args_dict['disable_amr']):
        #logging.info('calling ectyper')
        ectyper_result = call_ectyper(graph, args_dict)
        graph = ectyper_result['graph']
        #logging.info('ectyper call completed')

    # individual fasta logs are wiped on completion (or you'd have several
    # thousand of these)
    #remove('outputs/' + __name__ + args_dict['i'].split('/')[-1] + '.log')
    print upload_graph(graph)
    return json_return(args_dict, ectyper_result['ectyper_dict'])

if __name__ == "__main__":
    import argparse
    import os  # for batch cleanup

    from ConfigParser import SafeConfigParser
    from turtle_utils import generate_hash

    # parsing cli-input
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-i",
        help="FASTA file",
        required=True
    )
    parser.add_argument(
        "--disable-serotype",
        help="Disables use of the Serotyper. Serotyper is triggered by default.",
        action="store_true"
    )
    parser.add_argument(
        "--disable-vf",
        help="Disables use of ECTyper to get associated Virulence Factors. VFs are computed by default.",
        action="store_true"
    )
    parser.add_argument(
        "--disable-amr",
        help="Disables use of RGI to get Antimicrobial Resistance Factors.  AMR genes are computed by default.",
        action="store_true"
    )
    # note: by in large, we expect uri to be given as just the unique string
    # value  (be it the hash or the integer) without any prefixes, the actual
    # rdflib.URIRef object will be generated in this script
    # this is mainly for batch computation
    parser.add_argument(
        "--uriGenome",
        help="Allows the specification of the Genome URI separately. Expect just the hash (not an actual uri).",
    )
    # This is both for batch computation and for future extensions where there
    # are multiple sequencings per isolate (Campy)
    parser.add_argument(
        "--uriIsolate",
        help="Allows the specification of the Isolate URI separately. Expect just the integer (not the full :spfyID)",
        type=int
    )
    args = parser.parse_args()
    # we make a dictionary from the cli-inputs and add are uris to it
    # mainly used for when a func needs a lot of the args
    args_dict = vars(args)

    # starting#logging
    '''
    logging.basicConfig(
        filename='outputs/' + __name__ +
        args_dict['i'].split('/')[-1] + '.log',
        level=logging.INFO
    )
    '''

    # check if a genome uri isn't set yet
    if args_dict['uriIsolate'] is None:
        # this is temporary, TODO: include a spqarql query to the db
        uriIsolate = gu(':spfy' + str(hash(args_dict['i'].split('/')[-1])))
    else:
        uriIsolate = gu(':spfy' + args_dict['uriIsolate'])

    # if the fasta_file hash was not precomputed (batch scripts should
    # precompute it), we compute that now
    if args_dict['uriGenome'] is None:
        uriGenome = gu(':' + generate_hash(args_dict['i']))
    else:
        uriGenome = gu(':' + args_dict['uriGenome'])

    args_dict['uriIsolate'] = uriIsolate
    args_dict['uriGenome'] = uriGenome

    print savvy(args_dict)
<|endoftext|>


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

__LICENSE__ = """
Copyright (c) 2011, Steven Tobin and Contributors.
All rights reserved.

Contributors: Steven Tobin,
              Rob Lanphier,
              Mithrandir <mithrandiragain@lavabit.com>

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
    * Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of the <organization> nor the
      names of its contributors may be used to endorse or promote products
      derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""

import random
import os
import sys
import optparse
import re
import math

def generate_wordlist(wordfile=None,
                      min_length=5,
                      max_length=9,
                      valid_chars='.'):
    """
    generate a word list from either a kwarg word_file, or a system default
    """

    if wordfile is None:
        if os.path.exists("/usr/share/dict/words"):
            wordfile = "/usr/share/dict/words"
        elif os.path.exists("/usr/dict/words"):
            wordfile = "/usr/dict/words"
        else:
            # if we get here wordfile is not set, the try...except block below
            # will catch it
            print("No default word file found, please supply custom list")

    wordfile = os.path.expanduser(wordfile)  # just to be sure

    words = []
    regexp = re.compile("^%s{%i,%i}$" % (valid_chars, min_length, max_length))
    try:
        with open(wordfile) as wlf:
            for line in wlf:
                thisword = line.strip()
                if regexp.match(thisword) is not None:
                    words.append(thisword)
    except:
            print("Word list not loaded")
            raise SystemExit
    return words

def report_entropy(length, numwords):
    """
    Report number of words and bits of entropy
    """
    bits = math.log(length, 2)
    if (int(bits) == bits):
        print("Your word list contains %i words, or 2^%i words. " %
             (length, bits))
    else:
        print("Your word list contains %i words, or 2^%0.2f words. " %
             (length, bits))

    print("A %i word password from this list will have roughly"
           "%i (%0.2f * %i) bits of entropy," %
           (numwords, int(bits * numwords), bits, numwords)),
    print("assuming truly random word selection.")

def generate_xkcdpassword(wordlist, n_words=4, interactive=False):
    """
    generate an XKCD-style password from the words in wordlist
    """

    # useful if driving the logic from other code
    if not interactive:
        try:
            # random.SystemRandom() should be cryptographically secure
            return " ".join(random.SystemRandom().sample(wordlist, n_words))
        except NotImplementedError:
            print("System does not support random number generator or Python"
                   "version < 2.4.")

    # else, interactive session
    custom_n_words = raw_input("Enter number of words (default 4): ")

    if custom_n_words:
        n_words = int(custom_n_words)

    accepted = "n"

    while accepted.lower() not in ["y", "yes"]:
        try:
            passwd = " ".join(random.SystemRandom().sample(wordlist, n_words))
        except NotImplementedError:
            print("System does not support random number generator or Python"
                   "version < 2.4.")
        print("Generated: ", passwd)
        accepted = raw_input("Accept? [yN] ")

    return passwd

if __name__ == '__main__':

    usage = "usage: %prog [options]"
    parser = optparse.OptionParser(usage)

    parser.add_option("-w", "--wordfile", dest="wordfile",
                      default=None,
                      help="List of valid words for password")
    parser.add_option("--min", dest="min_length",
                      default=5, type="int",
                      help="Minimum length of words to make password")
    parser.add_option("--max", dest="max_length",
                      default=9, type="int",
                      help="Maximum length of words to make password")
    parser.add_option("-n", "--numwords", dest="numwords",
                      default=4, type="int",
                      help="Number of words to make password")
    parser.add_option("-i", "--interactive", dest="interactive",
                      default=False, action="store_true",
                      help="Interactively select a password")
    parser.add_option("-v", "--valid_chars", dest="valid_chars",
                      default='.',
                      help="Valid chars, using regexp style (e.g. '[a-z]'")
    parser.add_option("-e", "--entropy", dest="entropy",
                      default=False, action="store_true",
                      help="Report entropy for given option")

    (options, args) = parser.parse_args()

    if len(args) > 1:
        parser.error("Too many arguments.")
    if len(args) == 1:
        # supporting either -w or args[0] for wordlist, but not both
        if options.wordfile is None:
            options.wordfile = args[0]
        else:
            parser.error("Conflicting values for wordlist: " + args[0] +
                         " and " + options.wordfile)

    my_wordlist = generate_wordlist(wordfile=options.wordfile,
                                    min_length=options.min_length,
                                    max_length=options.max_length,
                                    valid_chars=options.valid_chars)
    if options.entropy:
        report_entropy(length=len(my_wordlist), numwords=options.numwords)
    print(generate_xkcdpassword(my_wordlist, interactive=options.interactive,
        n_words=options.numwords))
<commit_msg>updated contibutor list<commit_after>#!/usr/bin/env python
# encoding: utf-8

__LICENSE__ = """
Copyright (c) 2011, Steven Tobin and Contributors.
All rights reserved.

Contributors: Steven Tobin,
              Rob Lanphier,
              Mithrandir <mithrandiragain@lavabit.com>,
              Daniel Beecham

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
    * Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of the <organization> nor the
      names of its contributors may be used to endorse or promote products
      derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""

import random
import os
import sys
import optparse
import re
import math

def generate_wordlist(wordfile=None,
                      min_length=5,
                      max_length=9,
                      valid_chars='.'):
    """
    generate a word list from either a kwarg word_file, or a system default
    """

    if wordfile is None:
        if os.path.exists("/usr/share/dict/words"):
            wordfile = "/usr/share/dict/words"
        elif os.path.exists("/usr/dict/words"):
            wordfile = "/usr/dict/words"
        else:
            # if we get here wordfile is not set, the try...except block below
            # will catch it
            print("No default word file found, please supply custom list")

    wordfile = os.path.expanduser(wordfile)  # just to be sure

    words = []
    regexp = re.compile("^%s{%i,%i}$" % (valid_chars, min_length, max_length))
    try:
        with open(wordfile) as wlf:
            for line in wlf:
                thisword = line.strip()
                if regexp.match(thisword) is not None:
                    words.append(thisword)
    except:
            print("Word list not loaded")
            raise SystemExit
    return words

def report_entropy(length, numwords):
    """
    Report number of words and bits of entropy
    """
    bits = math.log(length, 2)
    if (int(bits) == bits):
        print("Your word list contains %i words, or 2^%i words. " %
             (length, bits))
    else:
        print("Your word list contains %i words, or 2^%0.2f words. " %
             (length, bits))

    print("A %i word password from this list will have roughly"
           "%i (%0.2f * %i) bits of entropy," %
           (numwords, int(bits * numwords), bits, numwords)),
    print("assuming truly random word selection.")

def generate_xkcdpassword(wordlist, n_words=4, interactive=False):
    """
    generate an XKCD-style password from the words in wordlist
    """

    # useful if driving the logic from other code
    if not interactive:
        try:
            # random.SystemRandom() should be cryptographically secure
            return " ".join(random.SystemRandom().sample(wordlist, n_words))
        except NotImplementedError:
            print("System does not support random number generator or Python"
                   "version < 2.4.")

    # else, interactive session
    custom_n_words = raw_input("Enter number of words (default 4): ")

    if custom_n_words:
        n_words = int(custom_n_words)

    accepted = "n"

    while accepted.lower() not in ["y", "yes"]:
        try:
            passwd = " ".join(random.SystemRandom().sample(wordlist, n_words))
        except NotImplementedError:
            print("System does not support random number generator or Python"
                   "version < 2.4.")
        print("Generated: ", passwd)
        accepted = raw_input("Accept? [yN] ")

    return passwd

if __name__ == '__main__':

    usage = "usage: %prog [options]"
    parser = optparse.OptionParser(usage)

    parser.add_option("-w", "--wordfile", dest="wordfile",
                      default=None,
                      help="List of valid words for password")
    parser.add_option("--min", dest="min_length",
                      default=5, type="int",
                      help="Minimum length of words to make password")
    parser.add_option("--max", dest="max_length",
                      default=9, type="int",
                      help="Maximum length of words to make password")
    parser.add_option("-n", "--numwords", dest="numwords",
                      default=4, type="int",
                      help="Number of words to make password")
    parser.add_option("-i", "--interactive", dest="interactive",
                      default=False, action="store_true",
                      help="Interactively select a password")
    parser.add_option("-v", "--valid_chars", dest="valid_chars",
                      default='.',
                      help="Valid chars, using regexp style (e.g. '[a-z]'")
    parser.add_option("-e", "--entropy", dest="entropy",
                      default=False, action="store_true",
                      help="Report entropy for given option")

    (options, args) = parser.parse_args()

    if len(args) > 1:
        parser.error("Too many arguments.")
    if len(args) == 1:
        # supporting either -w or args[0] for wordlist, but not both
        if options.wordfile is None:
            options.wordfile = args[0]
        else:
            parser.error("Conflicting values for wordlist: " + args[0] +
                         " and " + options.wordfile)

    my_wordlist = generate_wordlist(wordfile=options.wordfile,
                                    min_length=options.min_length,
                                    max_length=options.max_length,
                                    valid_chars=options.valid_chars)
    if options.entropy:
        report_entropy(length=len(my_wordlist), numwords=options.numwords)
    print(generate_xkcdpassword(my_wordlist, interactive=options.interactive,
        n_words=options.numwords))
<|endoftext|>


<commit_before>import os

from flask import (
    Flask,
    got_request_exception,
    render_template,
)
from flask_assets import Environment, Bundle
from flask_sitemap import Sitemap
from syspath import git_root

from app.routes import handlers, sitemap_urls

import dotenv
dotenv.load_dotenv(os.path.join(git_root.path, '.env'))

app = Flask(__name__)
app.debug = os.environ['DEBUG'] == 'true'
if os.environ.get('SERVER_NAME', ''):  # pragma: no cover
    app.config['SERVER_NAME'] = os.environ['SERVER_NAME']

app.config['SITEMAP_INCLUDE_RULES_WITHOUT_PARAMS'] = True
app.config['SITEMAP_URL_SCHEME'] = 'https'
assets = Environment(app)
ext = Sitemap(app=app)
ext.register_generator(sitemap_urls)

js = Bundle(
    'js/jquery.js',
    'js/tether.js',
    'js/bootstrap.js',
    'js/pdfobject.js',
    'js/global.js',
    filters='rjsmin', output='gen/bundle.min.js'
)
assets.register('js_all', js)
css = Bundle(
    'css/normalize.css',
    'css/tether.css',
    'css/bootstrap.css',
    'css/syntax.css',
    'css/global.css',
    filters='pyscss', output='gen/bundle.min.css'
)
assets.register('css_all', css)

if os.environ['ENV'] == 'production':
    import rollbar
    import rollbar.contrib.flask

    @app.before_first_request
    def init_rollbar():
        """init rollbar module"""
        rollbar.init(
            os.environ['ROLLBAR_SERVER_TOKEN'],
            # environment name
            os.environ['ENV'],
            # server root directory, makes tracebacks prettier
            root=os.path.dirname(os.path.realpath(__file__)),
            # flask already sets up logging
            allow_logging_basic_config=False)

        # send exceptions from `app` to rollbar, using flask's signal system.
        got_request_exception.connect(
            rollbar.contrib.flask.report_exception, app)

@app.context_processor
def inject_envs():
    envs = {}
    envs['ROLLBAR_CLIENT_TOKEN'] = os.environ['ROLLBAR_CLIENT_TOKEN']
    envs['SEGMENT_TOKEN'] = os.environ['SEGMENT_TOKEN']
    envs['ENV'] = os.environ['ENV']
    envs['LOGFIT_CLIENT_TOKEN'] = os.environ['LOGFIT_CLIENT_TOKEN']
    return {'ENV': envs}

app.register_blueprint(handlers)

@app.errorhandler(404)
def page_not_found(e):
    return render_template("404.htm"), 404

if __name__ == "__main__":
    app.run(host="0.0.0.0")
<commit_msg>Reorder imports to get environment variables<commit_after>import os

import dotenv
from flask import (
    Flask,
    got_request_exception,
    render_template,
)
from flask_assets import Environment, Bundle
from flask_sitemap import Sitemap
from syspath import git_root
dotenv.load_dotenv(os.path.join(git_root.path, '.env'))

from app.routes import handlers, sitemap_urls

app = Flask(__name__)
app.debug = os.environ['DEBUG'] == 'true'
if os.environ.get('SERVER_NAME', ''):  # pragma: no cover
    app.config['SERVER_NAME'] = os.environ['SERVER_NAME']

app.config['SITEMAP_INCLUDE_RULES_WITHOUT_PARAMS'] = True
app.config['SITEMAP_URL_SCHEME'] = 'https'
assets = Environment(app)
ext = Sitemap(app=app)
ext.register_generator(sitemap_urls)

js = Bundle(
    'js/jquery.js',
    'js/tether.js',
    'js/bootstrap.js',
    'js/pdfobject.js',
    'js/global.js',
    filters='rjsmin', output='gen/bundle.min.js'
)
assets.register('js_all', js)
css = Bundle(
    'css/normalize.css',
    'css/tether.css',
    'css/bootstrap.css',
    'css/syntax.css',
    'css/global.css',
    filters='pyscss', output='gen/bundle.min.css'
)
assets.register('css_all', css)

if os.environ['ENV'] == 'production':
    import rollbar
    import rollbar.contrib.flask

    @app.before_first_request
    def init_rollbar():
        """init rollbar module"""
        rollbar.init(
            os.environ['ROLLBAR_SERVER_TOKEN'],
            # environment name
            os.environ['ENV'],
            # server root directory, makes tracebacks prettier
            root=os.path.dirname(os.path.realpath(__file__)),
            # flask already sets up logging
            allow_logging_basic_config=False)

        # send exceptions from `app` to rollbar, using flask's signal system.
        got_request_exception.connect(
            rollbar.contrib.flask.report_exception, app)

@app.context_processor
def inject_envs():
    envs = {}
    envs['ROLLBAR_CLIENT_TOKEN'] = os.environ['ROLLBAR_CLIENT_TOKEN']
    envs['SEGMENT_TOKEN'] = os.environ['SEGMENT_TOKEN']
    envs['ENV'] = os.environ['ENV']
    envs['LOGFIT_CLIENT_TOKEN'] = os.environ['LOGFIT_CLIENT_TOKEN']
    return {'ENV': envs}

app.register_blueprint(handlers)

@app.errorhandler(404)
def page_not_found(e):
    return render_template("404.htm"), 404

if __name__ == "__main__":
    app.run(host="0.0.0.0")
<|endoftext|>


<commit_before>import cProfile
import time

import pandas as pd
from pyomo.environ import Constraint
from pypsa import Network

class Disaggregation:
    def __init__(self, original_network, clustered_network, clustering):
        """
        :param original_network: Initial (unclustered) network structure
        :param clustered_network: Clustered network used for the optimization
        :param clustering: The clustering object as returned by
        `pypsa.networkclustering.get_clustering_from_busmap`
        """
        self.original_network = original_network
        self.clustered_network = clustered_network
        self.clustering = clustering

        self.buses = pd.merge(original_network.buses,
                              clustering.busmap.to_frame(name='cluster'),
                              left_index=True, right_index=True)

        self.idx_prefix = '_'

    def add_constraints(self, cluster, extra_functionality=None):
        """
        Dummy function that allows the extension of `extra_functionalites` by
        custom conditions.

        :param cluster: Index of the cluster to disaggregate
        :param extra_functionality: extra_functionalities to extend
        :return: unaltered `extra_functionalites`
        """
        return extra_functionality

    def reindex_with_prefix(self, dataframe, prefix=None):
        if prefix is None:
            prefix = self.idx_prefix
            dataframe.set_index(
                dataframe.index.map(lambda x: self.idx_prefix + x),
            inplace=True)

    def construct_partial_network(self, cluster, scenario):
        """
        Compute the partial network that has been merged into a single cluster.
        The resulting network retains the external cluster buses that
        share some line with the cluster identified by `cluster`.
        These external buses will be prefixed by self.id_prefix in order to
        prevent name clashes with buses in the disaggregation

        :param cluster: Index of the cluster to disaggregate
        :return: Tuple of (partial_network, external_buses) where
        `partial_network` is the result of the partial decomposition
        and `external_buses` represent clusters adjacent to `cluster` that may
        be influenced by calculations done on the partial network.
        """

        #Create an empty network
        partial_network = scenario.build_network()

        # find all lines that have at least one bus inside the cluster
        busflags = (self.buses['cluster'] == cluster)

        def is_bus_in_cluster(conn):
            return busflags[conn]

        # Copy configurations to new network
        partial_network.snapshots = self.original_network.snapshots
        partial_network.snapshot_weightings = (self.original_network
                                                   .snapshot_weightings)
        partial_network.carriers = self.original_network.carriers

        # Collect all connectors that have some node inside the cluster

        external_buses = pd.DataFrame()

        line_types = ['lines', 'links', 'transformers']
        for line_type in line_types:
            # Copy all lines that reside entirely inside the cluster ...
            setattr(partial_network, line_type,
                    filter_internal_connector(
                        getattr(self.original_network, line_type),
                        is_bus_in_cluster))

            # ... and their time series
            # TODO: These are all time series, not just the ones from lines
            #       residing entirely in side the cluster.
            #       Is this a problem?
            setattr(partial_network, line_type + '_t',
                    getattr(self.original_network, line_type + '_t'))

            # Copy all lines whose `bus0` lies within the cluster
            left_external_connectors = filter_left_external_connector(
                getattr(self.original_network, line_type),
                is_bus_in_cluster)

            if not left_external_connectors.empty:
                f = lambda x: self.idx_prefix + self.clustering.busmap[x]
                left_external_connectors.bus0 = (left_external_connectors
                                                 .bus0
                                                 .apply(f))
                external_buses = pd.concat((external_buses,
                                            left_external_connectors.bus0))

            # Copy all lines whose `bus1` lies within the cluster
            right_external_connectors = filter_right_external_connector(
                getattr(self.original_network, line_type),
                is_bus_in_cluster)
            if not right_external_connectors.empty:
                f = lambda x: self.idx_prefix + self.clustering.busmap[x]
                right_external_connectors.bus1 = (right_external_connectors
                                                  .bus1
                                                  .apply(f))
                external_buses = pd.concat((external_buses,
                                            right_external_connectors.bus1))

        # Collect all buses that are contained in or somehow connected to the
        # cluster

        buses_in_lines = self.buses[busflags].index

        bus_types = ['loads', 'generators', 'stores', 'storage_units',
                     'shunt_impedances']

        # Copy all values that are part of the cluster
        partial_network.buses = self.original_network.buses[
            self.original_network.buses.index.isin(buses_in_lines)]

        # Collect all buses that are external, but connected to the cluster ...
        externals_to_insert = self.clustered_network.buses[
            self.clustered_network.buses.index.isin(
                map(lambda x: x[0][len(self.idx_prefix):],
                    external_buses.values))]

        # ... prefix them to avoid name clashes with buses from the original
        # network ...
        self.reindex_with_prefix(externals_to_insert)

        # .. and insert them as well as their time series
        partial_network.buses = (partial_network.buses
                                                .append(externals_to_insert))
        partial_network.buses_t = self.original_network.buses_t

        # TODO: Rename `bustype` to on_bus_type
        for bustype in bus_types:
            # Copy loads, generators, ... from original network to network copy
            setattr(partial_network, bustype,
                    filter_buses(getattr(self.original_network, bustype),
                                 buses_in_lines))

            # Collect on-bus components from external, connected clusters
            buses_to_insert = filter_buses(
                getattr(self.clustered_network, bustype),
                map(lambda x: x[0][len(self.idx_prefix):],
                    external_buses.values))

            # Prefix their external bindings
            buses_to_insert.bus = self.idx_prefix + buses_to_insert.bus

            setattr(partial_network, bustype,
                    getattr(partial_network, bustype).append(buses_to_insert))

            # Also copy their time series
            setattr(partial_network,
                    bustype + '_t',
                    getattr(self.original_network, bustype + '_t'))

        # Just a simple sanity check
        # TODO: Remove when sure that disaggregation will not go insane anymore
        for line_type in line_types:
            assert (getattr(partial_network, line_type).bus0.isin(
                partial_network.buses.index).all())
            assert (getattr(partial_network, line_type).bus1.isin(
                partial_network.buses.index).all())

        return partial_network, external_buses

    def execute(self, scenario, solver=None):
        self.solve(scenario, solver)

    def solve(self, scenario, solver):
        """
        Decompose each cluster into separate units and try to optimize them
        separately
        :param scenario:
        :param solver: Solver that may be used to optimize partial networks
        """
        clusters = set(self.clustering.busmap.values)
        n = len(clusters)
        i = 0
        profile = cProfile.Profile()
        for cluster in clusters:
            i += 1

            print('---')
            print('Decompose cluster %s (%d/%d)' % (cluster, i, n))
            profile.enable()
            t = time.time()
            partial_network, externals = self.construct_partial_network(
                    cluster,
                    scenario)
            print('Decomposed in ', (time.time() - t))
            t = time.time()
            profile.disable()
            profile.enable()
            self.solve_partial_network(cluster, partial_network, scenario,
                                       solver)
            self.transfer_results(partial_network, externals)
            profile.disable()
            print('Decomposition optimized in ', (time.time() - t))

        profile.print_stats(sort='cumtime')

    def transfer_results(self, partial_network, externals):
        for bustype in ['loads', 'generators', 'stores', 'storage_units',
                        'shunt_impedances']:
            # TODO: This variable is never used again.
            #       Is this a problem?
            changed_buses = (getattr(partial_network, bustype)
                             .index
                             .intersection(getattr(self.original_network,
                                                   bustype).index))
            orig_buses = getattr(self.original_network, bustype + '_t')
            part_buses = getattr(partial_network, bustype + '_t')
            for key in orig_buses.keys():
                for snap in partial_network.snapshots:
                    orig_buses[key].loc[snap].update(part_buses[key].loc[snap])

    def solve_partial_network(self, cluster, partial_network, scenario,
                              solver=None):
        extras = self.add_constraints(cluster)
        partial_network.lopf(scenario.timeindex,
                             solver_name=solver,
                             extra_functionality=extras)

class MiniSolverDisaggregation(Disaggregation):
    def add_constraints(self, cluster, extra_functionality=None):
        if extra_functionality is None:
            extra_functionality = lambda network, snapshots: None
        extra_functionality = self._validate_disaggregation_generators(
                cluster,
                extra_functionality)
        return extra_functionality

    def _validate_disaggregation_generators(self, cluster, f):
        def extra_functionality(network, snapshots):
            f(network, snapshots)
            generators = self.original_network.generators.assign(
                bus=lambda df: df.bus.map(self.clustering.busmap))
            grouper = [generators.carrier]
            def construct_constraint(model, snapshot, carrier):
                # TODO: Optimize

                generator_p = [model.generator_p[(x, snapshot)]
                               for x in generators.loc[
                                   (generators.bus == cluster) &
                                   (generators.carrier == carrier)].index]
                if not generator_p:
                    return Constraint.Feasible
                sum_generator_p = sum(generator_p)

                cluster_generators = self.clustered_network.generators[
                    (self.clustered_network.generators.bus == cluster) &
                    (self.clustered_network.generators.carrier == carrier)]
                sum_clustered_p = sum(
                    self.clustered_network.generators_t['p'].loc[snapshot, c]
                    for c in cluster_generators.index)
                return sum_generator_p == sum_clustered_p

            # TODO: Generate a better name
            network.model.validate_generators = Constraint(
                    list(snapshots),
                    set(generators.carrier),
                    rule=construct_constraint)
        return extra_functionality

    # TODO: This function is never used.
    #       Is this a problem?
    def _validate_disaggregation_buses(self, cluster, f):
        def extra_functionality(network, snapshots):
            f(network, snapshots)

            for bustype, bustype_pypsa, suffixes in [
                ('storage', 'storage_units', ['_dispatch', '_spill', '_store']),
                ('store', 'stores',[''])]:
                generators = getattr(self.original_network, bustype_pypsa).assign(
                    bus=lambda df: df.bus.map(self.clustering.busmap))
                for suffix in suffixes:
                    def construct_constraint(model, snapshot):
                        # TODO: Optimize
                        buses_p = [getattr(model,bustype + '_p' + suffix)[(x, snapshot)] for x in
                                       generators.loc[(generators.bus == cluster)].index]
                        if not buses_p:
                            return Constraint.Feasible
                        sum_bus_p = sum(buses_p)
                        cluster_buses = getattr(self.clustered_network, bustype_pypsa)[
                            (getattr(self.clustered_network, bustype_pypsa).bus == cluster)]
                        sum_clustered_p = sum(
                            getattr(self.clustered_network, bustype_pypsa + '_t')['p'].loc[snapshot,c] for
                            c in cluster_buses.index)
                        return sum_bus_p == sum_clustered_p

                    # TODO: Generate a better name
                    network.model.add_component('validate_' + bustype + suffix,Constraint(list(snapshots), rule=construct_constraint))
        return extra_functionality

class UniformDisaggregation(Disaggregation):
    def solve_partial_network(self, cluster, partial_network, scenario,
                              solver=None):
        for carrier in set(partial_network.generators.carrier):
            cgs = (self.clustered_network.generators
                    [self.clustered_network.generators.bus == cluster]
                    [self.clustered_network.generators.carrier == carrier])
            assert len(cgs) == 1, (
                   "Cluster {} has {} generators for carrier {}. "
                   .format(cluster, len(cgs), carrier) +
                   "Should be exactly one.")
            pgs = (partial_network.generators
                    [partial_network.generators.carrier == carrier])
            column = lambda cluster, carrier: "{} {}".format(cluster, carrier)
            cluster_t = (self.clustered_network
                             .generators_t['p']
                             .loc[:, column(cluster, carrier)])
            pnmp = pgs.p_nom * pgs.p_max_pu
            psum = sum(pnmp)
            pgs_t = partial_network.generators_t
            for generator_id in pgs.index:
                pgs_t['p'].loc[:, column(generator_id, carrier)] = (
                        cluster_t * pnmp.loc[generator_id] / psum)

def swap_series(s):
    return pd.Series(s.index.values, index=s)

def filter_internal_connector(conn, is_bus_in_cluster):
    return conn[conn.bus0.apply(is_bus_in_cluster)
                & conn.bus1.apply(is_bus_in_cluster)]

def filter_left_external_connector(conn, is_bus_in_cluster):
    return conn[~ conn.bus0.apply(is_bus_in_cluster)
                & conn.bus1.apply(is_bus_in_cluster)]

def filter_right_external_connector(conn, is_bus_in_cluster):
    return conn[conn.bus0.apply(is_bus_in_cluster)
                & ~conn.bus1.apply(is_bus_in_cluster)]

def filter_buses(bus, buses):
    return bus[bus.bus.isin(buses)]

def filter_on_buses(connecitve, buses):
    return connecitve[connecitve.bus.isin(buses)]

def update_constraints(network, externals):
    pass
<commit_msg>Skip carriers not present in partial network<commit_after>import cProfile
import time

import pandas as pd
from pyomo.environ import Constraint
from pypsa import Network

class Disaggregation:
    def __init__(self, original_network, clustered_network, clustering):
        """
        :param original_network: Initial (unclustered) network structure
        :param clustered_network: Clustered network used for the optimization
        :param clustering: The clustering object as returned by
        `pypsa.networkclustering.get_clustering_from_busmap`
        """
        self.original_network = original_network
        self.clustered_network = clustered_network
        self.clustering = clustering

        self.buses = pd.merge(original_network.buses,
                              clustering.busmap.to_frame(name='cluster'),
                              left_index=True, right_index=True)

        self.idx_prefix = '_'

    def add_constraints(self, cluster, extra_functionality=None):
        """
        Dummy function that allows the extension of `extra_functionalites` by
        custom conditions.

        :param cluster: Index of the cluster to disaggregate
        :param extra_functionality: extra_functionalities to extend
        :return: unaltered `extra_functionalites`
        """
        return extra_functionality

    def reindex_with_prefix(self, dataframe, prefix=None):
        if prefix is None:
            prefix = self.idx_prefix
            dataframe.set_index(
                dataframe.index.map(lambda x: self.idx_prefix + x),
            inplace=True)

    def construct_partial_network(self, cluster, scenario):
        """
        Compute the partial network that has been merged into a single cluster.
        The resulting network retains the external cluster buses that
        share some line with the cluster identified by `cluster`.
        These external buses will be prefixed by self.id_prefix in order to
        prevent name clashes with buses in the disaggregation

        :param cluster: Index of the cluster to disaggregate
        :return: Tuple of (partial_network, external_buses) where
        `partial_network` is the result of the partial decomposition
        and `external_buses` represent clusters adjacent to `cluster` that may
        be influenced by calculations done on the partial network.
        """

        #Create an empty network
        partial_network = scenario.build_network()

        # find all lines that have at least one bus inside the cluster
        busflags = (self.buses['cluster'] == cluster)

        def is_bus_in_cluster(conn):
            return busflags[conn]

        # Copy configurations to new network
        partial_network.snapshots = self.original_network.snapshots
        partial_network.snapshot_weightings = (self.original_network
                                                   .snapshot_weightings)
        partial_network.carriers = self.original_network.carriers

        # Collect all connectors that have some node inside the cluster

        external_buses = pd.DataFrame()

        line_types = ['lines', 'links', 'transformers']
        for line_type in line_types:
            # Copy all lines that reside entirely inside the cluster ...
            setattr(partial_network, line_type,
                    filter_internal_connector(
                        getattr(self.original_network, line_type),
                        is_bus_in_cluster))

            # ... and their time series
            # TODO: These are all time series, not just the ones from lines
            #       residing entirely in side the cluster.
            #       Is this a problem?
            setattr(partial_network, line_type + '_t',
                    getattr(self.original_network, line_type + '_t'))

            # Copy all lines whose `bus0` lies within the cluster
            left_external_connectors = filter_left_external_connector(
                getattr(self.original_network, line_type),
                is_bus_in_cluster)

            if not left_external_connectors.empty:
                f = lambda x: self.idx_prefix + self.clustering.busmap[x]
                left_external_connectors.bus0 = (left_external_connectors
                                                 .bus0
                                                 .apply(f))
                external_buses = pd.concat((external_buses,
                                            left_external_connectors.bus0))

            # Copy all lines whose `bus1` lies within the cluster
            right_external_connectors = filter_right_external_connector(
                getattr(self.original_network, line_type),
                is_bus_in_cluster)
            if not right_external_connectors.empty:
                f = lambda x: self.idx_prefix + self.clustering.busmap[x]
                right_external_connectors.bus1 = (right_external_connectors
                                                  .bus1
                                                  .apply(f))
                external_buses = pd.concat((external_buses,
                                            right_external_connectors.bus1))

        # Collect all buses that are contained in or somehow connected to the
        # cluster

        buses_in_lines = self.buses[busflags].index

        bus_types = ['loads', 'generators', 'stores', 'storage_units',
                     'shunt_impedances']

        # Copy all values that are part of the cluster
        partial_network.buses = self.original_network.buses[
            self.original_network.buses.index.isin(buses_in_lines)]

        # Collect all buses that are external, but connected to the cluster ...
        externals_to_insert = self.clustered_network.buses[
            self.clustered_network.buses.index.isin(
                map(lambda x: x[0][len(self.idx_prefix):],
                    external_buses.values))]

        # ... prefix them to avoid name clashes with buses from the original
        # network ...
        self.reindex_with_prefix(externals_to_insert)

        # .. and insert them as well as their time series
        partial_network.buses = (partial_network.buses
                                                .append(externals_to_insert))
        partial_network.buses_t = self.original_network.buses_t

        # TODO: Rename `bustype` to on_bus_type
        for bustype in bus_types:
            # Copy loads, generators, ... from original network to network copy
            setattr(partial_network, bustype,
                    filter_buses(getattr(self.original_network, bustype),
                                 buses_in_lines))

            # Collect on-bus components from external, connected clusters
            buses_to_insert = filter_buses(
                getattr(self.clustered_network, bustype),
                map(lambda x: x[0][len(self.idx_prefix):],
                    external_buses.values))

            # Prefix their external bindings
            buses_to_insert.bus = self.idx_prefix + buses_to_insert.bus

            setattr(partial_network, bustype,
                    getattr(partial_network, bustype).append(buses_to_insert))

            # Also copy their time series
            setattr(partial_network,
                    bustype + '_t',
                    getattr(self.original_network, bustype + '_t'))

        # Just a simple sanity check
        # TODO: Remove when sure that disaggregation will not go insane anymore
        for line_type in line_types:
            assert (getattr(partial_network, line_type).bus0.isin(
                partial_network.buses.index).all())
            assert (getattr(partial_network, line_type).bus1.isin(
                partial_network.buses.index).all())

        return partial_network, external_buses

    def execute(self, scenario, solver=None):
        self.solve(scenario, solver)

    def solve(self, scenario, solver):
        """
        Decompose each cluster into separate units and try to optimize them
        separately
        :param scenario:
        :param solver: Solver that may be used to optimize partial networks
        """
        clusters = set(self.clustering.busmap.values)
        n = len(clusters)
        i = 0
        profile = cProfile.Profile()
        for cluster in clusters:
            i += 1

            print('---')
            print('Decompose cluster %s (%d/%d)' % (cluster, i, n))
            profile.enable()
            t = time.time()
            partial_network, externals = self.construct_partial_network(
                    cluster,
                    scenario)
            print('Decomposed in ', (time.time() - t))
            t = time.time()
            profile.disable()
            profile.enable()
            self.solve_partial_network(cluster, partial_network, scenario,
                                       solver)
            self.transfer_results(partial_network, externals)
            profile.disable()
            print('Decomposition optimized in ', (time.time() - t))

        profile.print_stats(sort='cumtime')

    def transfer_results(self, partial_network, externals):
        for bustype in ['loads', 'generators', 'stores', 'storage_units',
                        'shunt_impedances']:
            # TODO: This variable is never used again.
            #       Is this a problem?
            changed_buses = (getattr(partial_network, bustype)
                             .index
                             .intersection(getattr(self.original_network,
                                                   bustype).index))
            orig_buses = getattr(self.original_network, bustype + '_t')
            part_buses = getattr(partial_network, bustype + '_t')
            for key in orig_buses.keys():
                for snap in partial_network.snapshots:
                    orig_buses[key].loc[snap].update(part_buses[key].loc[snap])

    def solve_partial_network(self, cluster, partial_network, scenario,
                              solver=None):
        extras = self.add_constraints(cluster)
        partial_network.lopf(scenario.timeindex,
                             solver_name=solver,
                             extra_functionality=extras)

class MiniSolverDisaggregation(Disaggregation):
    def add_constraints(self, cluster, extra_functionality=None):
        if extra_functionality is None:
            extra_functionality = lambda network, snapshots: None
        extra_functionality = self._validate_disaggregation_generators(
                cluster,
                extra_functionality)
        return extra_functionality

    def _validate_disaggregation_generators(self, cluster, f):
        def extra_functionality(network, snapshots):
            f(network, snapshots)
            generators = self.original_network.generators.assign(
                bus=lambda df: df.bus.map(self.clustering.busmap))
            grouper = [generators.carrier]
            def construct_constraint(model, snapshot, carrier):
                # TODO: Optimize

                generator_p = [model.generator_p[(x, snapshot)]
                               for x in generators.loc[
                                   (generators.bus == cluster) &
                                   (generators.carrier == carrier)].index]
                if not generator_p:
                    return Constraint.Feasible
                sum_generator_p = sum(generator_p)

                cluster_generators = self.clustered_network.generators[
                    (self.clustered_network.generators.bus == cluster) &
                    (self.clustered_network.generators.carrier == carrier)]
                sum_clustered_p = sum(
                    self.clustered_network.generators_t['p'].loc[snapshot, c]
                    for c in cluster_generators.index)
                return sum_generator_p == sum_clustered_p

            # TODO: Generate a better name
            network.model.validate_generators = Constraint(
                    list(snapshots),
                    set(generators.carrier),
                    rule=construct_constraint)
        return extra_functionality

    # TODO: This function is never used.
    #       Is this a problem?
    def _validate_disaggregation_buses(self, cluster, f):
        def extra_functionality(network, snapshots):
            f(network, snapshots)

            for bustype, bustype_pypsa, suffixes in [
                ('storage', 'storage_units', ['_dispatch', '_spill', '_store']),
                ('store', 'stores',[''])]:
                generators = getattr(self.original_network, bustype_pypsa).assign(
                    bus=lambda df: df.bus.map(self.clustering.busmap))
                for suffix in suffixes:
                    def construct_constraint(model, snapshot):
                        # TODO: Optimize
                        buses_p = [getattr(model,bustype + '_p' + suffix)[(x, snapshot)] for x in
                                       generators.loc[(generators.bus == cluster)].index]
                        if not buses_p:
                            return Constraint.Feasible
                        sum_bus_p = sum(buses_p)
                        cluster_buses = getattr(self.clustered_network, bustype_pypsa)[
                            (getattr(self.clustered_network, bustype_pypsa).bus == cluster)]
                        sum_clustered_p = sum(
                            getattr(self.clustered_network, bustype_pypsa + '_t')['p'].loc[snapshot,c] for
                            c in cluster_buses.index)
                        return sum_bus_p == sum_clustered_p

                    # TODO: Generate a better name
                    network.model.add_component('validate_' + bustype + suffix,Constraint(list(snapshots), rule=construct_constraint))
        return extra_functionality

class UniformDisaggregation(Disaggregation):
    def solve_partial_network(self, cluster, partial_network, scenario,
                              solver=None):
        for carrier in set(partial_network.generators.carrier):
            cgs = (self.clustered_network.generators
                    [self.clustered_network.generators.bus == cluster]
                    [self.clustered_network.generators.carrier == carrier])
            if len(cgs) == 0:
                break
            assert len(cgs) == 1, (
                   "Cluster {} has {} generators for carrier {}. "
                   .format(cluster, len(cgs), carrier) +
                   "Should be exactly one.")
            pgs = (partial_network.generators
                    [partial_network.generators.carrier == carrier])
            column = lambda cluster, carrier: "{} {}".format(cluster, carrier)
            cluster_t = (self.clustered_network
                             .generators_t['p']
                             .loc[:, column(cluster, carrier)])
            pnmp = pgs.p_nom * pgs.p_max_pu
            psum = sum(pnmp)
            pgs_t = partial_network.generators_t
            for generator_id in pgs.index:
                pgs_t['p'].loc[:, column(generator_id, carrier)] = (
                        cluster_t * pnmp.loc[generator_id] / psum)

def swap_series(s):
    return pd.Series(s.index.values, index=s)

def filter_internal_connector(conn, is_bus_in_cluster):
    return conn[conn.bus0.apply(is_bus_in_cluster)
                & conn.bus1.apply(is_bus_in_cluster)]

def filter_left_external_connector(conn, is_bus_in_cluster):
    return conn[~ conn.bus0.apply(is_bus_in_cluster)
                & conn.bus1.apply(is_bus_in_cluster)]

def filter_right_external_connector(conn, is_bus_in_cluster):
    return conn[conn.bus0.apply(is_bus_in_cluster)
                & ~conn.bus1.apply(is_bus_in_cluster)]

def filter_buses(bus, buses):
    return bus[bus.bus.isin(buses)]

def filter_on_buses(connecitve, buses):
    return connecitve[connecitve.bus.isin(buses)]

def update_constraints(network, externals):
    pass
<|endoftext|>


<commit_before>from components import stateManagment, formControl, player, helpers
from functools import partial
from model.model import Model
import pygame as pg

SCENE_NAME = "SHOW_TOWER"
class Scene(stateManagment.BaseScene):
    def __init__(self, scene, helpers):
        super(Scene, self).__init__(helpers)
        self.current_state = SCENE_NAME
        self.vars = helpers['vars']
        self.assets = helpers['assets']
        self.game = game = self.vars['pygame']
        self.tower_upper = Grid(1, 5)
        self.tower_lower = Grid()
        self.players = None
        self.button = formControl.Button(
            [0, self.game['height']/5, 200, 50],
            (0, 0, 0),
            self.next_scene(),
            text="Stuff"
        )

    def next_scene(self):
#        self.persist['game_state']['current_player_index'] += 1

        self.done = True
        self.next_state = "CHOOSE_DIRECTION"

    def get_event(self, event):
        helpers.check_paused_event(self, event)
        self.button.check_event(event)

    def startup(self, persistent):
        self.next_state = 'CHOOSE_DIRECTION'
        self.persist = persistent
        self.players = self.persist['game_state']['players']
        game_state = self.persist['game_state']
        pindex = game_state['current_player_index']
        player = game_state['players'][pindex]
        player.canmove()
        player.update()

        if player.lower_done:
            self.tower_upper.addplayer(player)
            player.y = 4
            player.upper_done = True
            player.lower_done = False
        elif not player.lower_done and player.upper_done:
            self.tower_upper.addplayer(player)
            if player.y < 0:
                player.lower_done = True
        elif player.lower_done and player.upper_done:
            #TODO make victory screen
            pass
        else:
            self.tower_lower.addplayer(player)
            if player.y < 0:
                player.lower_done = True

    #formControl.Button([x,y,w,h],[color],function(where thebuttongoes to), text=something)
    def update(self, dt):
        pass

    def draw(self, surface):
        surface.fill((0,0,0))
        self.tower_upper.draw(surface, self.game['width'], self.game['height']/5, 0)
        self.tower_lower.draw(surface, self.game['width'], self.game['height']/2.5, self.game['height']/4 + 50)
        for p in self.players:
            p.draw(surface, self.game['width'], self.game['height'])
        self.button.update(surface)

class Grid:
    def __init__(self, grid_width=2, grid_height=10):
        self.points =[]
        self.players =[]

        self.grid_width = grid_width
        self.grid_height = grid_height

        self.colorlist = ((255,0,0), (0,0,255), (255, 255, 0), (0,255, 0))

    def addplayer(self, player):
        if not self.players.__contains__(player):
            self.players.append(player)

#draw the grid and update whilst checking if someone wins
#if someone wins, def returns True
    def draw(self, screen, width, height, repos):

        #draw backgroundcolors
        i = 1
        for counter in range(0,4):
            pg.draw.rect(screen, self.colorlist[counter], [i, repos, width / 4, height+repos], 0)
            i += width / 4

        for c in range(0,4):
            templist = []
            for x in range(0, self.grid_width):
                for y in range(0, self.grid_height):
                    if not self.players == []:
                        for player in self.players:
                            if (((player.x+2)%(x+2) == 0 or player.x/2%(x+2) == 0) and player.y == y and c == (player.x+2)/2) or \
                                    (((player.x+2)%(x+2) == 0 or player.x/2%(x+2) == 0) and player.y == y and c == (player.x-1)/2):
                                Point(x, y, c, 1).drawself(screen, width, height + repos, self.grid_height, repos)
                            else:
                                Point(x, y ,c, 0).drawself(screen, width, height+repos, self.grid_height, repos)
                                templist.append(Point(x, y ,c, 0))
                    else:
                        Point(x, y, c, 0).drawself(screen, width, height + repos, self.grid_height, repos)
                        templist.append(Point(x, y, c, 0))

                templist.append(Point(x, y ,c, 1))
            self.points.append(templist)
class Point:
    def __init__(self, x, y, category, highlight):
        self.x = x
        self.y = y
        self.category = category
        self.highlight = highlight

    def highlight(self):
        if self.highlight == 0:
            self.highlight = 1
        else:
            self.highlight = 0

    def drawself(self, screen, width, height, grid_height=10, repos=0):
        if self.x >= 0 and self.y >= 0:
            pg.draw.rect(screen, ((self.highlight*255),(self.highlight*255),(self.highlight*255)),
                         [width/20 + width/4*self.category + width/8*self.x,
                          height/grid_height *self.y + height/50 + repos,
                          8*(1+self.highlight),
                          8*(1+self.highlight)], 2)
        else:
            print("Player is not in game yet")
<commit_msg>fix bug<commit_after>from components import stateManagment, formControl, player, helpers
from functools import partial
from model.model import Model
import pygame as pg

SCENE_NAME = "SHOW_TOWER"
class Scene(stateManagment.BaseScene):
    def __init__(self, scene, helpers):
        super(Scene, self).__init__(helpers)
        self.current_state = SCENE_NAME
        self.vars = helpers['vars']
        self.assets = helpers['assets']
        self.game = game = self.vars['pygame']
        self.tower_upper = Grid(1, 5)
        self.tower_lower = Grid()
        self.players = None
        self.button = formControl.Button(
            (300, 300 , 200, 50),
            pg.Color('pink'),
            self.next_scene,
            text="Stuff"
        )

    def next_scene(self, id):
        game_state = self.persist['game_state']
        current_player_index = game_state['current_player_index']
        players = game_state['players']
        if len(players) == current_player_index + 1:
            self.persist['game_state']['current_player_index'] = 0
        else:
            self.persist['game_state']['current_player_index'] += 1
        self.done = True
        self.next_state = "CHOOSE_DIRECTION"

    def get_event(self, event):
        helpers.check_paused_event(self, event)
        self.button.check_event(event)

    def startup(self, persistent):
        self.next_state = 'CHOOSE_DIRECTION'
        self.persist = persistent
        self.players = self.persist['game_state']['players']
        print(self.persist)
        game_state = self.persist['game_state']
        pindex = game_state['current_player_index']
        player = game_state['players'][pindex]
        player.canmove()
        player.update()

        if player.lower_done:
            self.tower_upper.addplayer(player)
            player.y = 4
            player.upper_done = True
            player.lower_done = False
        elif not player.lower_done and player.upper_done:
            self.tower_upper.addplayer(player)
            if player.y < 0:
                player.lower_done = True
        elif player.lower_done and player.upper_done:
            #TODO make victory screen
            pass
        else:
            self.tower_lower.addplayer(player)
            if player.y < 0:
                player.lower_done = True

    #formControl.Button([x,y,w,h],[color],function(where thebuttongoes to), text=something)
    def update(self, dt):
        pass

    def draw(self, surface):
        surface.fill((0,0,0))
        self.tower_upper.draw(surface, self.game['width'], self.game['height']/5, 0)
        self.tower_lower.draw(surface, self.game['width'], self.game['height']/2.5, self.game['height']/4 + 50)
        for p in self.players:
            p.draw(surface, self.game['width'], self.game['height'])
        self.button.update(surface)

class Grid:
    def __init__(self, grid_width=2, grid_height=10):
        self.points =[]
        self.players =[]

        self.grid_width = grid_width
        self.grid_height = grid_height

        self.colorlist = ((255,0,0), (0,0,255), (255, 255, 0), (0,255, 0))

    def addplayer(self, player):
        if not self.players.__contains__(player):
            self.players.append(player)

#draw the grid and update whilst checking if someone wins
#if someone wins, def returns True
    def draw(self, screen, width, height, repos):

        #draw backgroundcolors
        i = 1
        for counter in range(0,4):
            pg.draw.rect(screen, self.colorlist[counter], [i, repos, width / 4, height+repos], 0)
            i += width / 4

        for c in range(0,4):
            templist = []
            for x in range(0, self.grid_width):
                for y in range(0, self.grid_height):
                    if not self.players == []:
                        for player in self.players:
                            if (((player.x+2)%(x+2) == 0 or player.x/2%(x+2) == 0) and player.y == y and c == (player.x+2)/2) or \
                                    (((player.x+2)%(x+2) == 0 or player.x/2%(x+2) == 0) and player.y == y and c == (player.x-1)/2):
                                Point(x, y, c, 1).drawself(screen, width, height + repos, self.grid_height, repos)
                            else:
                                Point(x, y ,c, 0).drawself(screen, width, height+repos, self.grid_height, repos)
                                templist.append(Point(x, y ,c, 0))
                    else:
                        Point(x, y, c, 0).drawself(screen, width, height + repos, self.grid_height, repos)
                        templist.append(Point(x, y, c, 0))

                templist.append(Point(x, y ,c, 1))
            self.points.append(templist)
class Point:
    def __init__(self, x, y, category, highlight):
        self.x = x
        self.y = y
        self.category = category
        self.highlight = highlight

    def highlight(self):
        if self.highlight == 0:
            self.highlight = 1
        else:
            self.highlight = 0

    def drawself(self, screen, width, height, grid_height=10, repos=0):
        if self.x >= 0 and self.y >= 0:
            pg.draw.rect(screen, ((self.highlight*255),(self.highlight*255),(self.highlight*255)),
                         [width/20 + width/4*self.category + width/8*self.x,
                          height/grid_height *self.y + height/50 + repos,
                          8*(1+self.highlight),
                          8*(1+self.highlight)], 2)
        else:
            print("Player is not in game yet")
<|endoftext|>