<commit_before>import logging

from six import string_types
import constants

class BucketFull(Exception):
    """Raised when the bucket is full."""
    pass

class KBucket(object):
    """FILLME"""

    def __init__(self, rangeMin, rangeMax, market_id=1):
        """
        Initialize a new KBucket with a range and a market_id.

        @param rangeMin: The lower boundary for the range in the 160-bit ID
                         space covered by this k-bucket
        @param rangeMax: The upper boundary for the range in the ID space
                         covered by this k-bucket
        @param market_id: FILLME
        """

        self.lastAccessed = 0
        self.rangeMin = rangeMin
        self.rangeMax = rangeMax
        self.contacts = []
        self.market_id = market_id

        self.log = logging.getLogger(
            '[%s] %s' % (market_id, self.__class__.__name__)
        )

    def __len__(self):
        return len(self.contacts)

    def addContact(self, contact):
        """
        Add a contact to the contact list.

        The new contact is always appended to the contact list after removing
        any prior occurences of the same contact.

        @param contact: The contact to add or a string containing the
                        contact's node ID
        @type contact: connection.CryptoPeerConnection or str

        @raise node.kbucket.BucketFull: The bucket is full and the contact
                                        to add is not already in it.
        """
        try:
            # Assume contact exists. Attempt to remove the old one...
            self.contacts.remove(contact)
            # ... and add the new one at the end of the list.
            self.contacts.append(contact)

            # The code above works as follows:
            # Assume C1 is the existing contact and C2 is the new contact.
            # Iff C1 is equal to C2, it will be removed from the list.
            # Since Contact.__eq__ compares only GUIDs, contact C1 will
            # be replaced even if it's not exactly the same as C2.
            # This is the intended behaviour; the fresh contact may have
            # updated add-on data (e.g. optimization-specific stuff).
        except ValueError:
            # The contact wasn't there after all, so add it.
            if len(self.contacts) < constants.k:
                self.contacts.append(contact)
            else:
                raise BucketFull('No space in bucket to insert contact')

    def getContact(self, contactID):
        """
        Return the contact with the specified ID or None if not present.

        @param contactID: The ID to search.
        @type contact: connection.CryptoPeerConnection or str

        @rtype: connection.CryptoPeerConnection or None
        """
        self.log.debug('[getContact] %s' % contactID)
        for contact in self.contacts:
            if contact == contactID:
                self.log.debug('[getContact] Found %s' % contact)
                return contact
        self.log.debug('[getContact] No Results')
        return None

    def getContacts(self, count=-1, excludeContact=None):
        """
        Return a list containing up to the first `count` number of contacts.

        @param count: The amount of contacts to return;
                      if 0 or less, return all contacts.
        @type count: int
        @param excludeContact: A contact to exclude; if this contact is in
                               the list of returned values, it will be
                               discarded before returning. If a C{str} is
                               passed as this argument, it must be the
                               contact's ID.
        @type excludeContact: connection.CryptoPeerConnection or str

        @return: The first `count` contacts in the contact list.
                 This amount is capped by the available contacts
                 and the bucket size, of course. If no contacts
                 are present, an empty list is returned.
        @rtype:  list of connection.CryptoPeerConnection
        """

        currentLen = len(self)
        if not currentLen:
            return []

        if count <= 0:
            count = currentLen
        else:
            count = min(count, currentLen)

        # Return no more contacts than bucket size.
        count = min(count, constants.k)

        contactList = self.contacts[:count]
        if excludeContact is not None:
            try:
                # NOTE: If the excludeContact is removed, the resulting
                # list has one less contact than expected. Not sure if
                # this is a bug.
                contactList.remove(excludeContact)
            except ValueError:
                self.log.debug(
                    '[kbucket.getContacts() warning] '
                    'tried to exclude non-existing contact '
                    '(%s)' % excludeContact
                )
        return contactList

    def removeContact(self, contact):
        """
        Remove given contact from contact list.

        @param contact: The contact to remove, or a string containing the
                        contact's node ID
        @type contact: connection.CryptoPeerConnection or str

        @raise ValueError: The specified contact is not in this bucket
        """
        self.contacts.remove(contact)

    def keyInRange(self, key):
        """
        Tests whether the specified key (i.e. node ID) is in the range
        of the 160-bit ID space covered by this k-bucket (in other words, it
        returns whether or not the specified key should be placed in this
        k-bucket)

        @param key: The key to test
        @type key: str or int

        @return: C{True} if key is in this k-bucket's range,
                 C{False} otherwise.
        @rtype: bool
        """
        if isinstance(key, string_types):
            key = long(key, 16)
        return self.rangeMin <= key < self.rangeMax
<commit_msg>Do not use string_types from six.<commit_after>import logging

import constants

class BucketFull(Exception):
    """Raised when the bucket is full."""
    pass

class KBucket(object):
    """FILLME"""

    def __init__(self, rangeMin, rangeMax, market_id=1):
        """
        Initialize a new KBucket with a range and a market_id.

        @param rangeMin: The lower boundary for the range in the 160-bit ID
                         space covered by this k-bucket
        @param rangeMax: The upper boundary for the range in the ID space
                         covered by this k-bucket
        @param market_id: FILLME
        """

        self.lastAccessed = 0
        self.rangeMin = rangeMin
        self.rangeMax = rangeMax
        self.contacts = []
        self.market_id = market_id

        self.log = logging.getLogger(
            '[%s] %s' % (market_id, self.__class__.__name__)
        )

    def __len__(self):
        return len(self.contacts)

    def addContact(self, contact):
        """
        Add a contact to the contact list.

        The new contact is always appended to the contact list after removing
        any prior occurences of the same contact.

        @param contact: The contact to add or a string containing the
                        contact's node ID
        @type contact: connection.CryptoPeerConnection or str

        @raise node.kbucket.BucketFull: The bucket is full and the contact
                                        to add is not already in it.
        """
        try:
            # Assume contact exists. Attempt to remove the old one...
            self.contacts.remove(contact)
            # ... and add the new one at the end of the list.
            self.contacts.append(contact)

            # The code above works as follows:
            # Assume C1 is the existing contact and C2 is the new contact.
            # Iff C1 is equal to C2, it will be removed from the list.
            # Since Contact.__eq__ compares only GUIDs, contact C1 will
            # be replaced even if it's not exactly the same as C2.
            # This is the intended behaviour; the fresh contact may have
            # updated add-on data (e.g. optimization-specific stuff).
        except ValueError:
            # The contact wasn't there after all, so add it.
            if len(self.contacts) < constants.k:
                self.contacts.append(contact)
            else:
                raise BucketFull('No space in bucket to insert contact')

    def getContact(self, contactID):
        """
        Return the contact with the specified ID or None if not present.

        @param contactID: The ID to search.
        @type contact: connection.CryptoPeerConnection or str

        @rtype: connection.CryptoPeerConnection or None
        """
        self.log.debug('[getContact] %s' % contactID)
        for contact in self.contacts:
            if contact == contactID:
                self.log.debug('[getContact] Found %s' % contact)
                return contact
        self.log.debug('[getContact] No Results')
        return None

    def getContacts(self, count=-1, excludeContact=None):
        """
        Return a list containing up to the first `count` number of contacts.

        @param count: The amount of contacts to return;
                      if 0 or less, return all contacts.
        @type count: int
        @param excludeContact: A contact to exclude; if this contact is in
                               the list of returned values, it will be
                               discarded before returning. If a C{str} is
                               passed as this argument, it must be the
                               contact's ID.
        @type excludeContact: connection.CryptoPeerConnection or str

        @return: The first `count` contacts in the contact list.
                 This amount is capped by the available contacts
                 and the bucket size, of course. If no contacts
                 are present, an empty list is returned.
        @rtype:  list of connection.CryptoPeerConnection
        """

        currentLen = len(self)
        if not currentLen:
            return []

        if count <= 0:
            count = currentLen
        else:
            count = min(count, currentLen)

        # Return no more contacts than bucket size.
        count = min(count, constants.k)

        contactList = self.contacts[:count]
        if excludeContact is not None:
            try:
                # NOTE: If the excludeContact is removed, the resulting
                # list has one less contact than expected. Not sure if
                # this is a bug.
                contactList.remove(excludeContact)
            except ValueError:
                self.log.debug(
                    '[kbucket.getContacts() warning] '
                    'tried to exclude non-existing contact '
                    '(%s)' % excludeContact
                )
        return contactList

    def removeContact(self, contact):
        """
        Remove given contact from contact list.

        @param contact: The contact to remove, or a string containing the
                        contact's node ID
        @type contact: connection.CryptoPeerConnection or str

        @raise ValueError: The specified contact is not in this bucket
        """
        self.contacts.remove(contact)

    def keyInRange(self, key):
        """
        Tests whether the specified key (i.e. node ID) is in the range
        of the 160-bit ID space covered by this k-bucket (in other words, it
        returns whether or not the specified key should be placed in this
        k-bucket)

        @param key: The key to test
        @type key: str or int

        @return: C{True} if key is in this k-bucket's range,
                 C{False} otherwise.
        @rtype: bool
        """
        if isinstance(key, basestring):
            key = long(key, 16)
        return self.rangeMin <= key < self.rangeMax
<|endoftext|>


<commit_before>import arcpy

sdeStream = r'C:\inetpub\wwwroot\serverprojects\MapData\SGID10.gdb\StreamsNHDHighRes'
waterQualityFGDB = r'C:\inetpub\wwwroot\serverprojects\MapData\SurfaceWaterQuality\Data_1.gdb'
temp = waterQualityFGDB + '\\temp'
temp2 = waterQualityFGDB + '\\temp2'
temp3 = waterQualityFGDB + '\\temp3'
streamsLocal = waterQualityFGDB + '\\StreamsDissolved'
GNIS_Name = 'GNIS_Name'
COUNTY = 'COUNTY'
NAME = 'NAME'
counties = r'Database Connections\SGID10_DirectConnect.sde\SGID10.BOUNDARIES.Counties'
dissolveTempField = 'tempDissolve'

if arcpy.Exists(streamsLocal):
    print 'deleting old streams local data'
    arcpy.Delete_management(streamsLocal)
    print arcpy.GetMessages()

print '\n\ncreating layer'
lyr = arcpy.MakeFeatureLayer_management(sdeStream, 'streamLayer', "\"%s\" is not null AND \"%s\" <> ''" % (GNIS_Name, GNIS_Name))
print arcpy.GetMessages()

print '\n\n' + str(arcpy.GetCount_management(lyr))

print '\n\ndissolving'
arcpy.Dissolve_management(lyr, temp, GNIS_Name)
print arcpy.GetMessages()

print '\n\nExploding lines'
arcpy.MultipartToSinglepart_management(temp, temp2)
print arcpy.GetMessages()

print '\n\nidentity'
arcpy.Identity_analysis(temp2, counties, temp3)
print arcpy.GetMessages()

print '\n\nadding dissolve field'
arcpy.AddField_management(temp3, dissolveTempField, "TEXT", '', '', 50)
print arcpy.GetMessages()

print '\n\ncalculating dissolve field'
arcpy.CalculateField_management(temp3, dissolveTempField, '!%s! + !%s!' % (GNIS_Name, NAME), "PYTHON")
print arcpy.GetMessages()

print '\n\ncreating layer to filer outside of state'
lyr2 = arcpy.MakeFeatureLayer_management(temp3, 'tempLayer', "\"%s\" <> ''" % (NAME))
print arcpy.GetMessages()

print '\n\ndissolving again'
arcpy.Dissolve_management(lyr2, streamsLocal, dissolveTempField)
print arcpy.GetMessages()

print '\n\njoining fields'
arcpy.JoinField_management(streamsLocal, dissolveTempField, temp3, dissolveTempField, '%s;%s' % (GNIS_Name, NAME))
print arcpy.GetMessages()

print '\n\nadding ' + COUNTY
arcpy.AddField_management(streamsLocal, COUNTY, 'TEXT', '', '', 25)
print arcpy.GetMessages()

print '\n\ncalculating ' + COUNTY
arcpy.CalculateField_management(streamsLocal, COUNTY, '!%s!' % (NAME), 'PYTHON')
print arcpy.GetMessages()

print '\n\ndeleting ' + NAME
arcpy.DeleteField_management(streamsLocal, NAME)
print arcpy.GetMessages()

print '\n\ndelete'
#arcpy.Delete_management(temp)
#print arcpy.GetMessages()
#arcpy.Delete_management(temp2)
#print arcpy.GetMessages()
#arcpy.Delete_management(temp3)
#print arcpy.GetMessages()

print 'done'<commit_msg>chore: modernize nhd stream data import script<commit_after>import arcpy
import sys
from os.path import join

sde = sys.argv[1]
sdeStream = join(sde, 'SGID10.WATER.StreamsNHDHighRes')
waterQualityFGDB = sys.argv[2]
temp = join(waterQualityFGDB, 'temp')
temp2 = join(waterQualityFGDB, 'temp2')
temp3 = join(waterQualityFGDB, 'temp3')
streamsLocal = join(waterQualityFGDB, 'StreamsDissolved')
GNIS_Name = 'GNIS_Name'
COUNTY = 'COUNTY'
NAME = 'NAME'
counties = join(sde, 'SGID10.BOUNDARIES.Counties')
dissolveTempField = 'tempDissolve'

if arcpy.Exists(streamsLocal):
    print('deleting old streams local data')
    arcpy.Delete_management(streamsLocal)
    print(arcpy.GetMessages())

print('creating layer')
lyr = arcpy.MakeFeatureLayer_management(sdeStream, 'streamLayer', "\"%s\" is not null AND \"%s\" <> ''" % (GNIS_Name, GNIS_Name))
print(arcpy.GetMessages())

print('' + str(arcpy.GetCount_management(lyr)))

print('dissolving')
arcpy.Dissolve_management(lyr, temp, GNIS_Name)
print(arcpy.GetMessages())

print('Exploding lines')
arcpy.MultipartToSinglepart_management(temp, temp2)
print(arcpy.GetMessages())

print('identity')
arcpy.Identity_analysis(temp2, counties, temp3)
print(arcpy.GetMessages())

print('adding dissolve field')
arcpy.AddField_management(temp3, dissolveTempField, "TEXT", '', '', 50)
print(arcpy.GetMessages())

print('calculating dissolve field')
arcpy.CalculateField_management(temp3, dissolveTempField, '!%s! + !%s!' % (GNIS_Name, NAME), "PYTHON")
print(arcpy.GetMessages())

print('creating layer to filer outside of state')
lyr2 = arcpy.MakeFeatureLayer_management(temp3, 'tempLayer', "\"%s\" <> ''" % (NAME))
print(arcpy.GetMessages())

print('dissolving again')
arcpy.Dissolve_management(lyr2, streamsLocal, dissolveTempField)
print(arcpy.GetMessages())

print('joining fields')
arcpy.JoinField_management(streamsLocal, dissolveTempField, temp3, dissolveTempField, '%s;%s' % (GNIS_Name, NAME))
print(arcpy.GetMessages())

print('adding ' + COUNTY)
arcpy.AddField_management(streamsLocal, COUNTY, 'TEXT', '', '', 25)
print(arcpy.GetMessages())

print('calculating ' + COUNTY)
arcpy.CalculateField_management(streamsLocal, COUNTY, '!%s!' % (NAME), 'PYTHON')
print(arcpy.GetMessages())

print('deleting ' + NAME)
arcpy.DeleteField_management(streamsLocal, NAME)
print(arcpy.GetMessages())

print('delete')
arcpy.Delete_management(temp)
print(arcpy.GetMessages())
arcpy.Delete_management(temp2)
print(arcpy.GetMessages())
arcpy.Delete_management(temp3)
print(arcpy.GetMessages())

print('done')
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
This is the application file for the tool eTraGo.

Define your connection parameters and power flow settings before executing the function etrago.

This program is free software; you can redistribute it and/or
modify it under the terms of the GNU Affero General Public License as
published by the Free Software Foundation; either version 3 of the
License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""

__copyright__ = "Flensburg University of Applied Sciences, Europa-Universität Flensburg, Centre for Sustainable Energy Systems, DLR-Institute for Networked Energy Systems"
__license__ = "GNU Affero General Public License Version 3 (AGPL-3.0)"
__author__ = "ulfmueller, lukasol, wolfbunke, mariusves, s3pp"

import numpy as np
from numpy import genfromtxt
np.random.seed()
import time

import os

if not 'READTHEDOCS' in os.environ:
    # Sphinx does not run this code.
    # Do not import internal packages directly
    from etrago.tools.io import NetworkScenario, results_to_oedb
    from etrago.tools.plot import (plot_line_loading, plot_stacked_gen,
                                     add_coordinates, curtailment, gen_dist,
                                     storage_distribution)
    from etrago.tools.utilities import (load_shedding, data_manipulation_sh,
                                    results_to_csv, parallelisation, pf_post_lopf,
                                    loading_minimization, calc_line_losses, group_parallel_lines)
    from etrago.cluster.networkclustering import busmap_from_psql, cluster_on_extra_high_voltage, kmean_clustering
    from egoio.tools import db
    from sqlalchemy.orm import sessionmaker

args = {# Setup and Configuration:
        'db': 'oedb', # db session
        'gridversion': None, # None for model_draft or Version number (e.g. v0.2.11) for grid schema
        'method': 'lopf', # lopf or pf
        'pf_post_lopf': False, # state whether you want to perform a pf after a lopf simulation
        'start_snapshot': 1,
        'end_snapshot' : 8760,
        'scn_name': 'NEP 2035', # state which scenario you want to run: Status Quo, NEP 2035, eGo100
        'solver': 'gurobi', # glpk, cplex or gurobi
        # Export options:
        'lpfile': False, # state if and where you want to save pyomo's lp file: False or /path/tofolder
        'results': '/home/lukas_wienholt/results/NEP2035_k400_t3_gurobi_noise', # state if and where you want to save results as csv: False or /path/tofolder
        'export': False, # state if you want to export the results back to the database
        # Settings:
        'storage_extendable':True, # state if you want storages to be installed at each node if necessary.
        'generator_noise':True, # state if you want to apply a small generator noise
        'reproduce_noise': '/home/lukas_wienholt/noise_values.csv', # state if you want to use a predefined set of random noise for the given scenario. if so, provide path, e.g. 'noise_values.csv'
        'minimize_loading':False,
        # Clustering:
        'k_mean_clustering': 400, # state if you want to perform a k-means clustering on the given network. State False or the value k (e.g. 20).
        'network_clustering': False, # state if you want to perform a clustering of HV buses to EHV buses.
        # Simplifications:
        'parallelisation':False, # state if you want to run snapshots parallely.
        'skip_snapshots':3,
        'line_grouping': False, # state if you want to group lines running between the same buses.
        'branch_capacity_factor': 0.7, # globally extend or lower branch capacities
        'load_shedding':False, # meet the demand at very high cost; for debugging purposes.
        'comments':None }

def etrago(args):
    """The etrago function works with following arguments:

    Parameters
    ----------

    db (str):
    	'oedb',
        Name of Database session setting stored in config.ini of oemof.db

    gridversion (str):
        'v0.2.11',
        Name of the data version number of oedb: state 'None' for
        model_draft (sand-box) or an explicit version number
        (e.g. 'v0.2.10') for the grid schema.

    method (str):
        'lopf',
        Choose between a non-linear power flow ('pf') or
        a linear optimal power flow ('lopf').

    pf_post_lopf (bool):
        False,
        Option to run a non-linear power flow (pf) directly after the
        linear optimal power flow (and thus the dispatch) has finished.

    start_snapshot (int):
    	1,
        Start hour of the scenario year to be calculated.

    end_snapshot (int) :
    	2,
        End hour of the scenario year to be calculated.

    scn_name (str):
    	'Status Quo',
	Choose your scenario. Currently, there are three different
	scenarios: 'Status Quo', 'NEP 2035', 'eGo100'. If you do not
	want to use the full German dataset, you can use the excerpt of
	Schleswig-Holstein by adding the acronym SH to the scenario
	name (e.g. 'SH Status Quo').

    solver (str):
        'glpk',
        Choose your preferred solver. Current options: 'glpk' (open-source),
        'cplex' or 'gurobi'.

    lpfile (obj):
        False,
        State if and where you want to save pyomo's lp file. Options:
        False or '/path/tofolder'.

    results (obj):
        False,
        State if and where you want to save results as csv files.Options:
        False or '/path/tofolder'.

    export (bool):
        False,
        State if you want to export the results of your calculation
        back to the database.

    storage_extendable (bool):
        True,
        Choose if you want to allow to install extendable storages
        (unlimited in size) at each grid node in order to meet the flexibility demand.

    generator_noise (bool):
        True,
        Choose if you want to apply a small random noise to the marginal
        costs of each generator in order to prevent an optima plateau.

    reproduce_noise (obj):
        False,
        State if you want to use a predefined set of random noise for
        the given scenario. If so, provide path to the csv file,
        e.g. 'noise_values.csv'.

    minimize_loading (bool):
        False,

    k_mean_clustering (bool):
        False,
        State if you want to apply a clustering of all network buses down to
        only 'k' buses. The weighting takes place considering generation and load
        at each node. If so, state the number of k you want to apply. Otherwise
        put False. This function doesn't work together with 'line_grouping = True'
	    or 'network_clustering = True'.

    network_clustering (bool):
        False,
        Choose if you want to cluster the full HV/EHV dataset down to only the EHV
        buses. In that case, all HV buses are assigned to their closest EHV sub-station,
        taking into account the shortest distance on power lines.

    parallelisation (bool):
        False,
        Choose if you want to calculate a certain number of snapshots in parallel. If
        yes, define the respective amount in the if-clause execution below. Otherwise
        state False here.

    line_grouping (bool):
        True,
        State if you want to group lines that connect the same two buses into one system.

    branch_capacity_factor (numeric):
        1,
        Add a factor here if you want to globally change line capacities (e.g. to "consider"
        an (n-1) criterion or for debugging purposes.

    load_shedding (bool):
        False,
        State here if you want to make use of the load shedding function which is helpful when
        debugging: a very expensive generator is set to each bus and meets the demand when regular
        generators cannot do so.

    comments (str):
        None

    Result:
    -------

    """
    conn = db.connection(section=args['db'])
    Session = sessionmaker(bind=conn)
    session = Session()

    # additional arguments cfgpath, version, prefix
    if args['gridversion'] == None:
        args['ormcls_prefix'] = 'EgoGridPfHv'
    else:
        args['ormcls_prefix'] = 'EgoPfHv'

    scenario = NetworkScenario(session,
                               version=args['gridversion'],
                               prefix=args['ormcls_prefix'],
                               method=args['method'],
                               start_snapshot=args['start_snapshot'],
                               end_snapshot=args['end_snapshot'],
                               scn_name=args['scn_name'])

    network = scenario.build_network()

    # add coordinates
    network = add_coordinates(network)

    # set SOC at the beginning and end of the period to equal values
    network.storage_units.cyclic_state_of_charge = True

    # TEMPORARY vague adjustment due to transformer bug in data processing
    if args['gridversion'] == 'v0.2.11':
        network.transformers.x=network.transformers.x*0.0001

    if args['branch_capacity_factor']:
        network.lines.s_nom = network.lines.s_nom*args['branch_capacity_factor']
        network.transformers.s_nom = network.transformers.s_nom*args['branch_capacity_factor']

    if args['generator_noise']:
        # create or reproduce generator noise
        if not args['reproduce_noise'] == False:
            noise_values = genfromtxt('noise_values.csv', delimiter=',')
            # add random noise to all generator
            network.generators.marginal_cost = noise_values
        else:
            noise_values = network.generators.marginal_cost + abs(np.random.normal(0,0.001,len(network.generators.marginal_cost)))
            np.savetxt("noise_values.csv", noise_values, delimiter=",")
            noise_values = genfromtxt('noise_values.csv', delimiter=',')
            # add random noise to all generator
            network.generators.marginal_cost = noise_values

    if args['storage_extendable']:
        # set virtual storages to be extendable
        if network.storage_units.carrier[network.storage_units.carrier== 'extendable_storage'].any() == 'extendable_storage':
            network.storage_units.loc[network.storage_units.carrier=='extendable_storage','p_nom_extendable'] = True
        # set virtual storage costs with regards to snapshot length
            network.storage_units.capital_cost = (network.storage_units.capital_cost /
            (8760//(args['end_snapshot']-args['start_snapshot']+1)))

    # for SH scenario run do data preperation:
    if args['scn_name'] == 'SH Status Quo' or args['scn_name'] == 'SH NEP 2035':
        data_manipulation_sh(network)

    # grouping of parallel lines
    if args['line_grouping']:
        group_parallel_lines(network)

    #load shedding in order to hunt infeasibilities
    if args['load_shedding']:
    	load_shedding(network)

    # network clustering
    if args['network_clustering']:
        network.generators.control="PV"
        busmap = busmap_from_psql(network, session, scn_name=args['scn_name'])
        network = cluster_on_extra_high_voltage(network, busmap, with_time=True)

    # k-mean clustering
    if not args['k_mean_clustering'] == False:
        network = kmean_clustering(network, n_clusters=args['k_mean_clustering'])

    # Branch loading minimization
    if args['minimize_loading']:
        extra_functionality = loading_minimization
    else:
        extra_functionality=None

    if args['skip_snapshots']:
        network.snapshots=network.snapshots[::args['skip_snapshots']]
        network.snapshot_weightings=network.snapshot_weightings[::args['skip_snapshots']]*args['skip_snapshots']

    # parallisation
    if args['parallelisation']:
        parallelisation(network, start_snapshot=args['start_snapshot'], end_snapshot=args['end_snapshot'],group_size=1, solver_name=args['solver'], extra_functionality=extra_functionality)
    # start linear optimal powerflow calculations
    elif args['method'] == 'lopf':
        x = time.time()
        network.lopf(network.snapshots, solver_name=args['solver'], extra_functionality=extra_functionality, solver_options={'threads':4, 'method':2, 'crossover':0, 'BarConvTol':1.e-5,'FeasibilityTol':1.e-6})
        y = time.time()
        z = (y - x) / 60 # z is time for lopf in minutes
    # start non-linear powerflow simulation
    elif args['method'] == 'pf':
        network.pf(scenario.timeindex)
       # calc_line_losses(network)

    if args['pf_post_lopf']:
        pf_post_lopf(network, scenario)
        calc_line_losses(network)

       # provide storage installation costs
    if sum(network.storage_units.p_nom_opt) != 0:
        installed_storages = network.storage_units[ network.storage_units.p_nom_opt!=0]
        storage_costs = sum(installed_storages.capital_cost * installed_storages.p_nom_opt)
        print("Investment costs for all storages in selected snapshots [EUR]:",round(storage_costs,2))

    # write lpfile to path
    if not args['lpfile'] == False:
        network.model.write(args['lpfile'], io_options={'symbolic_solver_labels':
                                                     True})
    # write PyPSA results back to database
    if args['export']:
        username = str(conn.url).split('//')[1].split(':')[0]
        args['user_name'] = username
        safe_results=False #default is False. If it is set to 'True' the result set will be safed
                           #to the versioned grid schema eventually apart from
                           #being saved to the model_draft.
                           #ONLY set to True if you know what you are doing.
        results_to_oedb(session, network, args, grid='hv', safe_results = safe_results)

    # write PyPSA results to csv to path
    if not args['results'] == False:
        results_to_csv(network, args['results'])

    # close session
    #session.close()

    return network

if __name__ == '__main__':
    # execute etrago function
    network = etrago(args)
    # plots
    # make a line loading plot
    #plot_line_loading(network)
    # plot stacked sum of nominal power for each generator type and timestep
    #plot_stacked_gen(network, resolution="MW")
    # plot to show extendable storages
    #storage_distribution(network)
<commit_msg>change feasibility tol<commit_after># -*- coding: utf-8 -*-
"""
This is the application file for the tool eTraGo.

Define your connection parameters and power flow settings before executing the function etrago.

This program is free software; you can redistribute it and/or
modify it under the terms of the GNU Affero General Public License as
published by the Free Software Foundation; either version 3 of the
License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""

__copyright__ = "Flensburg University of Applied Sciences, Europa-Universität Flensburg, Centre for Sustainable Energy Systems, DLR-Institute for Networked Energy Systems"
__license__ = "GNU Affero General Public License Version 3 (AGPL-3.0)"
__author__ = "ulfmueller, lukasol, wolfbunke, mariusves, s3pp"

import numpy as np
from numpy import genfromtxt
np.random.seed()
import time

import os

if not 'READTHEDOCS' in os.environ:
    # Sphinx does not run this code.
    # Do not import internal packages directly
    from etrago.tools.io import NetworkScenario, results_to_oedb
    from etrago.tools.plot import (plot_line_loading, plot_stacked_gen,
                                     add_coordinates, curtailment, gen_dist,
                                     storage_distribution)
    from etrago.tools.utilities import (load_shedding, data_manipulation_sh,
                                    results_to_csv, parallelisation, pf_post_lopf,
                                    loading_minimization, calc_line_losses, group_parallel_lines)
    from etrago.cluster.networkclustering import busmap_from_psql, cluster_on_extra_high_voltage, kmean_clustering
    from egoio.tools import db
    from sqlalchemy.orm import sessionmaker

args = {# Setup and Configuration:
        'db': 'oedb', # db session
        'gridversion': None, # None for model_draft or Version number (e.g. v0.2.11) for grid schema
        'method': 'lopf', # lopf or pf
        'pf_post_lopf': False, # state whether you want to perform a pf after a lopf simulation
        'start_snapshot': 1,
        'end_snapshot' : 8760,
        'scn_name': 'NEP 2035', # state which scenario you want to run: Status Quo, NEP 2035, eGo100
        'solver': 'gurobi', # glpk, cplex or gurobi
        # Export options:
        'lpfile': False, # state if and where you want to save pyomo's lp file: False or /path/tofolder
        'results': '/home/openego/pf_results/storage_paper/NEP2035_k300_t3_gurobi_noise', # state if and where you want to save results as csv: False or /path/tofolder
        'export': False, # state if you want to export the results back to the database
        # Settings:
        'storage_extendable':True, # state if you want storages to be installed at each node if necessary.
        'generator_noise':True, # state if you want to apply a small generator noise
        'reproduce_noise': '/home/openego/noise_values_lukastorage.csv', # state if you want to use a predefined set of random noise for the given scenario. if so, provide path, e.g. 'noise_values.csv'
        'minimize_loading':False,
        # Clustering:
        'k_mean_clustering': 300, # state if you want to perform a k-means clustering on the given network. State False or the value k (e.g. 20).
        'network_clustering': False, # state if you want to perform a clustering of HV buses to EHV buses.
        # Simplifications:
        'parallelisation':False, # state if you want to run snapshots parallely.
        'skip_snapshots':3,
        'line_grouping': False, # state if you want to group lines running between the same buses.
        'branch_capacity_factor': 0.7, # globally extend or lower branch capacities
        'load_shedding':False, # meet the demand at very high cost; for debugging purposes.
        'comments':None }

def etrago(args):
    """The etrago function works with following arguments:

    Parameters
    ----------

    db (str):
    	'oedb',
        Name of Database session setting stored in config.ini of oemof.db

    gridversion (str):
        'v0.2.11',
        Name of the data version number of oedb: state 'None' for
        model_draft (sand-box) or an explicit version number
        (e.g. 'v0.2.10') for the grid schema.

    method (str):
        'lopf',
        Choose between a non-linear power flow ('pf') or
        a linear optimal power flow ('lopf').

    pf_post_lopf (bool):
        False,
        Option to run a non-linear power flow (pf) directly after the
        linear optimal power flow (and thus the dispatch) has finished.

    start_snapshot (int):
    	1,
        Start hour of the scenario year to be calculated.

    end_snapshot (int) :
    	2,
        End hour of the scenario year to be calculated.

    scn_name (str):
    	'Status Quo',
	Choose your scenario. Currently, there are three different
	scenarios: 'Status Quo', 'NEP 2035', 'eGo100'. If you do not
	want to use the full German dataset, you can use the excerpt of
	Schleswig-Holstein by adding the acronym SH to the scenario
	name (e.g. 'SH Status Quo').

    solver (str):
        'glpk',
        Choose your preferred solver. Current options: 'glpk' (open-source),
        'cplex' or 'gurobi'.

    lpfile (obj):
        False,
        State if and where you want to save pyomo's lp file. Options:
        False or '/path/tofolder'.

    results (obj):
        False,
        State if and where you want to save results as csv files.Options:
        False or '/path/tofolder'.

    export (bool):
        False,
        State if you want to export the results of your calculation
        back to the database.

    storage_extendable (bool):
        True,
        Choose if you want to allow to install extendable storages
        (unlimited in size) at each grid node in order to meet the flexibility demand.

    generator_noise (bool):
        True,
        Choose if you want to apply a small random noise to the marginal
        costs of each generator in order to prevent an optima plateau.

    reproduce_noise (obj):
        False,
        State if you want to use a predefined set of random noise for
        the given scenario. If so, provide path to the csv file,
        e.g. 'noise_values.csv'.

    minimize_loading (bool):
        False,

    k_mean_clustering (bool):
        False,
        State if you want to apply a clustering of all network buses down to
        only 'k' buses. The weighting takes place considering generation and load
        at each node. If so, state the number of k you want to apply. Otherwise
        put False. This function doesn't work together with 'line_grouping = True'
	    or 'network_clustering = True'.

    network_clustering (bool):
        False,
        Choose if you want to cluster the full HV/EHV dataset down to only the EHV
        buses. In that case, all HV buses are assigned to their closest EHV sub-station,
        taking into account the shortest distance on power lines.

    parallelisation (bool):
        False,
        Choose if you want to calculate a certain number of snapshots in parallel. If
        yes, define the respective amount in the if-clause execution below. Otherwise
        state False here.

    line_grouping (bool):
        True,
        State if you want to group lines that connect the same two buses into one system.

    branch_capacity_factor (numeric):
        1,
        Add a factor here if you want to globally change line capacities (e.g. to "consider"
        an (n-1) criterion or for debugging purposes.

    load_shedding (bool):
        False,
        State here if you want to make use of the load shedding function which is helpful when
        debugging: a very expensive generator is set to each bus and meets the demand when regular
        generators cannot do so.

    comments (str):
        None

    Result:
    -------

    """
    conn = db.connection(section=args['db'])
    Session = sessionmaker(bind=conn)
    session = Session()

    # additional arguments cfgpath, version, prefix
    if args['gridversion'] == None:
        args['ormcls_prefix'] = 'EgoGridPfHv'
    else:
        args['ormcls_prefix'] = 'EgoPfHv'

    scenario = NetworkScenario(session,
                               version=args['gridversion'],
                               prefix=args['ormcls_prefix'],
                               method=args['method'],
                               start_snapshot=args['start_snapshot'],
                               end_snapshot=args['end_snapshot'],
                               scn_name=args['scn_name'])

    network = scenario.build_network()

    # add coordinates
    network = add_coordinates(network)

    # set SOC at the beginning and end of the period to equal values
    network.storage_units.cyclic_state_of_charge = True

    # TEMPORARY vague adjustment due to transformer bug in data processing
    if args['gridversion'] == 'v0.2.11':
        network.transformers.x=network.transformers.x*0.0001

    if args['branch_capacity_factor']:
        network.lines.s_nom = network.lines.s_nom*args['branch_capacity_factor']
        network.transformers.s_nom = network.transformers.s_nom*args['branch_capacity_factor']

    if args['generator_noise']:
        # create or reproduce generator noise
        if not args['reproduce_noise'] == False:
            noise_values = genfromtxt('noise_values.csv', delimiter=',')
            # add random noise to all generator
            network.generators.marginal_cost = noise_values
        else:
            noise_values = network.generators.marginal_cost + abs(np.random.normal(0,0.001,len(network.generators.marginal_cost)))
            np.savetxt("noise_values.csv", noise_values, delimiter=",")
            noise_values = genfromtxt('noise_values.csv', delimiter=',')
            # add random noise to all generator
            network.generators.marginal_cost = noise_values

    if args['storage_extendable']:
        # set virtual storages to be extendable
        if network.storage_units.carrier[network.storage_units.carrier== 'extendable_storage'].any() == 'extendable_storage':
            network.storage_units.loc[network.storage_units.carrier=='extendable_storage','p_nom_extendable'] = True
        # set virtual storage costs with regards to snapshot length
            network.storage_units.capital_cost = (network.storage_units.capital_cost /
            (8760//(args['end_snapshot']-args['start_snapshot']+1)))

    # for SH scenario run do data preperation:
    if args['scn_name'] == 'SH Status Quo' or args['scn_name'] == 'SH NEP 2035':
        data_manipulation_sh(network)

    # grouping of parallel lines
    if args['line_grouping']:
        group_parallel_lines(network)

    #load shedding in order to hunt infeasibilities
    if args['load_shedding']:
    	load_shedding(network)

    # network clustering
    if args['network_clustering']:
        network.generators.control="PV"
        busmap = busmap_from_psql(network, session, scn_name=args['scn_name'])
        network = cluster_on_extra_high_voltage(network, busmap, with_time=True)

    # k-mean clustering
    if not args['k_mean_clustering'] == False:
        network = kmean_clustering(network, n_clusters=args['k_mean_clustering'])

    # Branch loading minimization
    if args['minimize_loading']:
        extra_functionality = loading_minimization
    else:
        extra_functionality=None

    if args['skip_snapshots']:
        network.snapshots=network.snapshots[::args['skip_snapshots']]
        network.snapshot_weightings=network.snapshot_weightings[::args['skip_snapshots']]*args['skip_snapshots']

    # parallisation
    if args['parallelisation']:
        parallelisation(network, start_snapshot=args['start_snapshot'], end_snapshot=args['end_snapshot'],group_size=1, solver_name=args['solver'], extra_functionality=extra_functionality)
    # start linear optimal powerflow calculations
    elif args['method'] == 'lopf':
        x = time.time()
        network.lopf(network.snapshots, solver_name=args['solver'], extra_functionality=extra_functionality, solver_options={'threads':4, 'method':2, 'crossover':0, 'BarConvTol':1.e-5,'FeasibilityTol':1.e-5})
        y = time.time()
        z = (y - x) / 60 # z is time for lopf in minutes
    # start non-linear powerflow simulation
    elif args['method'] == 'pf':
        network.pf(scenario.timeindex)
       # calc_line_losses(network)

    if args['pf_post_lopf']:
        pf_post_lopf(network, scenario)
        calc_line_losses(network)

       # provide storage installation costs
    if sum(network.storage_units.p_nom_opt) != 0:
        installed_storages = network.storage_units[ network.storage_units.p_nom_opt!=0]
        storage_costs = sum(installed_storages.capital_cost * installed_storages.p_nom_opt)
        print("Investment costs for all storages in selected snapshots [EUR]:",round(storage_costs,2))

    # write lpfile to path
    if not args['lpfile'] == False:
        network.model.write(args['lpfile'], io_options={'symbolic_solver_labels':
                                                     True})
    # write PyPSA results back to database
    if args['export']:
        username = str(conn.url).split('//')[1].split(':')[0]
        args['user_name'] = username
        safe_results=False #default is False. If it is set to 'True' the result set will be safed
                           #to the versioned grid schema eventually apart from
                           #being saved to the model_draft.
                           #ONLY set to True if you know what you are doing.
        results_to_oedb(session, network, args, grid='hv', safe_results = safe_results)

    # write PyPSA results to csv to path
    if not args['results'] == False:
        results_to_csv(network, args['results'])

    # close session
    #session.close()

    return network

if __name__ == '__main__':
    # execute etrago function
    network = etrago(args)
    # plots
    # make a line loading plot
    #plot_line_loading(network)
    # plot stacked sum of nominal power for each generator type and timestep
    #plot_stacked_gen(network, resolution="MW")
    # plot to show extendable storages
    #storage_distribution(network)
<|endoftext|>


<commit_before>#!/usr/bin/env python
############################################################################################
# This example creates LIDAR like map using an ultrasonic sensor and a servo with the GoPiGo
#
# http://www.dexterindustries.com/GoPiGo/
# History
# ------------------------------------------------
# Author     	Date      		Comments
# Karan		  	13 June 14  	Initial Authoring
#
'''
## License
 GoPiGo for the Raspberry Pi: an open source robotics platform for the Raspberry Pi.
 Copyright (C) 2015  Dexter Industries

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/gpl-3.0.txt>.
'''
#
############################################################################################
#
# ! Attach Ultrasonic sensor to A1 Port.
#
############################################################################################
from gopigo import *
import sys
from collections import Counter
import math

def us_map():
	delay=.02
	debug =0					#True to print all raw values
	num_of_readings=45			#Number of readings to take
	incr=180/num_of_readings	#increment of angle in servo
	ang_l=[0]*(num_of_readings+1)	#list to hold the angle's of readings
	dist_l=[0]*(num_of_readings+1)	#list to hold the distance at each angle
	x=[0]*(num_of_readings+1)	#list to hold the x coordinate of each point
	y=[0]*(num_of_readings+1)	#list to hold the y coordinate of each point

	buf=[0]*40
	ang=0
	lim=250		#maximum limit of distance measurement (any value over this which be initialized to the limit value)
	index=0
	sample=2	#Number of samples for each angle (more the samples, better the data but more the time taken)
	print "Getting the data"

	while True:
		#Take the readings from the Ultrasonic sensor and process them to get the correct values
		for i in range(sample):
			dist=us_dist(15)
			if dist<lim and dist>=0:
				buf[i]=dist
			else:
				buf[i]=lim

		#Find the sample that is most common among all the samples for a particular angle
		max=Counter(buf).most_common()
		rm=-1
		for i in range (len(max)):
			if max[i][0] <> lim and max[i][0] <> 0:
				rm=max[i][0]
				break
		if rm==-1:
			rm=lim

		if debug==1:
			print index,ang,rm
		ang_l[index]=ang
		dist_l[index]=rm
		index+=1

		#Move the servo to the next angle
		servo(ang)
		time.sleep(delay)
		ang+=incr
		#print ang
		if ang>180:
			break

	#Print the values in a grid of 51x51 on the terminal
	grid_size=51

	#Convert the distance and angle to (x,y) coordinates and scale it down
	for i in range(num_of_readings+1):
		x[i]=(int(dist_l[i]*math.cos(math.pi*(ang_l[i])/180))/10)
		y[i]=int(dist_l[i]*math.sin(math.pi*ang_l[i]/180))/10

	#Rotate the readings so that it is printed in the correct manner
	for i in range(num_of_readings+1):
		x[i]=(grid_size/2)-x[i]
		y[i]=(grid_size/2)-y[i]
		print(str(x[i]) + " " + str(y[i]))
	#Create a grid
	grid = [[0 for a in xrange(grid_size)] for a in xrange(grid_size)]
	for i in range (num_of_readings+1):
		if dist_l[i]<>lim:
			grid[y[i]][x[i]]=1
	fence='-'*(grid_size+1)

	#Print the map
	print "Map:"
	print fence*2
	for i in range(grid_size/2):
		print "|",
		for j in range (grid_size):
			if (j==grid_size/2)and i==(grid_size/2)-1:
				print 'x',
			elif grid[i][j]==0:
				print ' ',
			else:
				print 'o',
		print "|"
	print fence*2
	return min(dist_l) #Return the closest distance in all directions

stop()
while True:
	#GoPiGo moves forward, stops makes a map and again moves forward
	enable_com_timeout(2000)
	enc_tgt(1,1,18)	#Set encoder targetting. Stop after 4 rotations of both the wheels
	fwd()
	time.sleep(.2)
	while True:
		enc=read_enc_status()
		ts=read_timeout_status()
		time.sleep(.05)
		print enc,ts
		if enc == 0:	#Stop when target is reached
			break
		if  ts==0:
			break

	time.sleep(2)
	enable_servo()
	if us_map() <20:	#If any obstacle is closer than 20 cm, stop
		break
	disable_servo()<commit_msg>print<commit_after>#!/usr/bin/env python
############################################################################################
# This example creates LIDAR like map using an ultrasonic sensor and a servo with the GoPiGo
#
# http://www.dexterindustries.com/GoPiGo/
# History
# ------------------------------------------------
# Author     	Date      		Comments
# Karan		  	13 June 14  	Initial Authoring
#
'''
## License
 GoPiGo for the Raspberry Pi: an open source robotics platform for the Raspberry Pi.
 Copyright (C) 2015  Dexter Industries

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/gpl-3.0.txt>.
'''
#
############################################################################################
#
# ! Attach Ultrasonic sensor to A1 Port.
#
############################################################################################
from gopigo import *
import sys
from collections import Counter
import math

def us_map():
	delay=.02
	debug =0					#True to print all raw values
	num_of_readings=45			#Number of readings to take
	incr=180/num_of_readings	#increment of angle in servo
	ang_l=[0]*(num_of_readings+1)	#list to hold the angle's of readings
	dist_l=[0]*(num_of_readings+1)	#list to hold the distance at each angle
	x=[0]*(num_of_readings+1)	#list to hold the x coordinate of each point
	y=[0]*(num_of_readings+1)	#list to hold the y coordinate of each point

	buf=[0]*40
	ang=0
	lim=250		#maximum limit of distance measurement (any value over this which be initialized to the limit value)
	index=0
	sample=2	#Number of samples for each angle (more the samples, better the data but more the time taken)
	print "Getting the data"

	while True:
		#Take the readings from the Ultrasonic sensor and process them to get the correct values
		for i in range(sample):
			dist=us_dist(15)
			if dist<lim and dist>=0:
				buf[i]=dist
			else:
				buf[i]=lim

		#Find the sample that is most common among all the samples for a particular angle
		max=Counter(buf).most_common()
		rm=-1
		for i in range (len(max)):
			if max[i][0] <> lim and max[i][0] <> 0:
				rm=max[i][0]
				break
		if rm==-1:
			rm=lim

		if debug==1:
			print index,ang,rm
		ang_l[index]=ang
		dist_l[index]=rm
		index+=1

		#Move the servo to the next angle
		servo(ang)
		time.sleep(delay)
		ang+=incr
		#print ang
		if ang>180:
			break

	#Print the values in a grid of 51x51 on the terminal
	grid_size=51

	#Convert the distance and angle to (x,y) coordinates and scale it down
	for i in range(num_of_readings+1):
		x[i]=(int(dist_l[i]*math.cos(math.pi*(ang_l[i])/180))/10)
		y[i]=int(dist_l[i]*math.sin(math.pi*ang_l[i]/180))/10

	#Rotate the readings so that it is printed in the correct manner
	for i in range(num_of_readings+1):
		x[i]=(grid_size/2)-x[i]
		y[i]=(grid_size/2)-y[i]
		print(str(x[i]) + " " + str(y[i]))
	#Create a grid
	grid = [[0 for a in xrange(grid_size)] for a in xrange(grid_size)]
	for i in range (num_of_readings+1):
		if dist_l[i]<>lim:
			grid[y[i]][x[i]]=1
	fence='-'*(grid_size+1)

	print(grid)

	#Print the map
	print "Map:"
	print fence*2
	for i in range(grid_size/2):
		print "|",
		for j in range (grid_size):
			if (j==grid_size/2)and i==(grid_size/2)-1:
				print 'x',
			elif grid[i][j]==0:
				print ' ',
			else:
				print 'o',
		print "|"
	print fence*2
	return min(dist_l) #Return the closest distance in all directions

stop()
while True:
	#GoPiGo moves forward, stops makes a map and again moves forward
	enable_com_timeout(2000)
	enc_tgt(1,1,18)	#Set encoder targetting. Stop after 4 rotations of both the wheels
	fwd()
	time.sleep(.2)
	while True:
		enc=read_enc_status()
		ts=read_timeout_status()
		time.sleep(.05)
		print enc,ts
		if enc == 0:	#Stop when target is reached
			break
		if  ts==0:
			break

	time.sleep(2)
	enable_servo()
	if us_map() <20:	#If any obstacle is closer than 20 cm, stop
		break
	disable_servo()<|endoftext|>


<commit_before>import nibabel

import nibabel as nib

import os
from nipype.interfaces.ants import Registration

import numpy as np
from nibabel.testing import data_path
import pandas as pd
from libtiff import TIFF
from medpy.io import save
from PIL import Image
import nibabel as nb
import operator
import connectLib as cLib
import scipy.sparse as sparse
from cluster import Cluster

def ANTs(fixedImg, movingImg, lowerFence, upperFence):
    img2 = nib.Nifti1Image(fixedImg, np.eye(4))
    nb.save(img2, 'fixed.nii')
    img3 = nib.Nifti1Image(fixedImg, np.eye(4))
    nb.save(img3, 'moving.nii')
    reg = Registration()
    reg.inputs.fixed_image = 'fixed.nii'
    reg.inputs.moving_image = 'moving.nii'
    reg.inputs.output_transform_prefix = 'thisTransform'
    reg.inputs.output_warped_image = 'registered.nii.gz'
    reg.inputs.output_transform_prefix = "output_"
    reg.inputs.transforms = ['Translation', 'Rigid', 'Affine']
    reg.inputs.transform_parameters = [(0.1,), (0.1,), (0.1,)]
    reg.inputs.number_of_iterations = ([[10000, 111110, 11110]] * 3)
    reg.inputs.dimension = 3
    reg.inputs.write_composite_transform = True
    reg.inputs.collapse_output_transforms = False
    reg.inputs.metric = ['MeanSquares'] * 3
    reg.inputs.metric_weight = [1] * 3
    reg.inputs.radius_or_number_of_bins = [32] * 3
    reg.inputs.sampling_strategy = ['Regular'] * 3
    reg.inputs.sampling_percentage = [0.3] * 3
    reg.inputs.convergence_threshold = [1.e-8] * 3
    reg.inputs.convergence_window_size = [20] * 3
    reg.inputs.smoothing_sigmas = [[4, 2, 1]] * 3
    reg.inputs.sigma_units = ['vox'] * 3
    reg.inputs.shrink_factors = [[6, 4, 2]] + [[3, 2, 1]] * 2
    reg.inputs.use_estimate_learning_rate_once = [True] * 3
    reg.inputs.use_histogram_matching = [False] * 3
    reg.inputs.initial_moving_transform_com = True
    reg.run()
    real_registered = os.path.join('registered.nii.gz')
    img = nib.load(real_registered)
    real_registered_img = img.get_data()

    registeredClusters = cLib.clusterThresh(real_registered_img, lowerFence, upperFence)
    fixedClusters = cLib.clusterThresh(fixedImg, lowerFence, upperFence)
    movingClusters = cLib.clusterThresh(movingImg, lowerFence, upperFence)

    for i in range(len(fixedClusters)):
        fixedCluster = fixedClusters[i]
        distances = []
        for j in range(len(registeredClusters)):
            registeredCluster = registeredClusters[j]
            distances.append(np.linalg.norm([x1 - x2 for (x1, x2) in zip(registeredCluster.getCentroid(), fixedCluster.getCentroid())]))
        min_index, _ = min(enumerate(distances), key=operator.itemgetter(1))
        centroid = registeredClusters[min_index].getCentroid()
        members = [np.ceil(i) for i in centroid]
        value = np.floor(real_registered_img[members[0], members[1], members[2]])
        i = 0
        while value == 0:
            if i < len(registeredClusters[min_index].members):
                newMembers = registeredClusters[min_index].members[i]
                value = np.floor(real_registered_img[newMembers[0], newMembers[1], newMembers[2]])
                i = i + 1
            if i >= len(registeredClusters[min_index].members):
                value = 2000

        #convert labeled to Sparse
        sparseLabeledMoving = np.empty(len(movingImg), dtype=object)
        for i in range(len(movingImg)):
            sparseLabeledMoving[i] = sparse.csr_matrix(movingImg[i])

        memberList = []
        memberListWithZ = []
        while(len(memberListWithZ) == 0):
            for z in range(len(sparseLabeledMoving)):
                memberListWithoutZ = np.argwhere(sparseLabeledMoving[z] == value)
                memberListWithZ = [[z] + list(tup) for tup in memberListWithoutZ]
                memberList.extend(memberListWithZ)
            value = value - 1

        fixedCluster.timeRegistration = Cluster(memberList)

    return fixedClusters
<commit_msg>made matching fast<commit_after>import nibabel

import nibabel as nib

import os
from nipype.interfaces.ants import Registration

import numpy as np
from scipy.spatial import KDTree
from nibabel.testing import data_path
import pandas as pd
from medpy.io import save
from PIL import Image
import nibabel as nb
import operator
import connectLib as cLib
import scipy.sparse as sparse
from cluster import Cluster

def ANTs(fixedImg, movingImg, lowerFence, upperFence, r = 5000):
    img2 = nib.Nifti1Image(fixedImg, np.eye(4))
    nb.save(img2, 'fixed.nii')
    img3 = nib.Nifti1Image(fixedImg, np.eye(4))
    nb.save(img3, 'moving.nii')
    reg = Registration()
    reg.inputs.fixed_image = 'fixed.nii'
    reg.inputs.moving_image = 'moving.nii'
    reg.inputs.output_transform_prefix = 'thisTransform'
    reg.inputs.output_warped_image = 'registered.nii.gz'
    reg.inputs.output_transform_prefix = "output_"
    reg.inputs.transforms = ['Translation', 'Rigid', 'Affine']
    reg.inputs.transform_parameters = [(0.1,), (0.1,), (0.1,)]
    reg.inputs.number_of_iterations = ([[10000, 111110, 11110]] * 3)
    reg.inputs.dimension = 3
    reg.inputs.write_composite_transform = True
    reg.inputs.collapse_output_transforms = False
    reg.inputs.metric = ['MeanSquares'] * 3
    reg.inputs.metric_weight = [1] * 3
    reg.inputs.radius_or_number_of_bins = [32] * 3
    reg.inputs.sampling_strategy = ['Regular'] * 3
    reg.inputs.sampling_percentage = [0.3] * 3
    reg.inputs.convergence_threshold = [1.e-8] * 3
    reg.inputs.convergence_window_size = [20] * 3
    reg.inputs.smoothing_sigmas = [[4, 2, 1]] * 3
    reg.inputs.sigma_units = ['vox'] * 3
    reg.inputs.shrink_factors = [[6, 4, 2]] + [[3, 2, 1]] * 2
    reg.inputs.use_estimate_learning_rate_once = [True] * 3
    reg.inputs.use_histogram_matching = [False] * 3
    reg.inputs.initial_moving_transform_com = True
    reg.run()
    real_registered = os.path.join('registered.nii.gz')
    img = nib.load(real_registered)
    real_registered_img = img.get_data()

    registeredClusters = cLib.clusterThresh(real_registered_img, lowerFence, upperFence)
    fixedClusters = cLib.clusterThresh(fixedImg, lowerFence, upperFence)
    movingClusters = cLib.clusterThresh(movingImg, lowerFence, upperFence)

    #l2 centroid match, capped at r
    A = [elem.getCentroid() for elem in fixedClusters]
    B = [elem.getCentroid() for elem in movingClusters]

    tree = KDTree(B)
    for baseIdx, a in enumerate(A):
        dist, idx = tree.query(a, k=1, distance_upper_bound = r)
        if dist == float('Inf'):
            fixedClusters[baseIdx].timeRegistration=Cluster([[-1, -1, -1]])
        else:
            fixedClusters[baseIdx].timeRegistration=movingClusters[idx]

    '''
    for i in range(len(fixedClusters)):
        fixedCluster = fixedClusters[i]
        distances = []
        for j in range(len(registeredClusters)):
            registeredCluster = registeredClusters[j]
            distances.append(np.linalg.norm([x1 - x2 for (x1, x2) in zip(registeredCluster.getCentroid(), fixedCluster.getCentroid())]))
        min_index, _ = min(enumerate(distances), key=operator.itemgetter(1))
        centroid = registeredClusters[min_index].getCentroid()
        members = [np.ceil(i) for i in centroid]
        value = np.floor(real_registered_img[members[0], members[1], members[2]])
        i = 0
        while value == 0:
            if i < len(registeredClusters[min_index].members):
                newMembers = registeredClusters[min_index].members[i]
                value = np.floor(real_registered_img[newMembers[0], newMembers[1], newMembers[2]])
                i = i + 1
            if i >= len(registeredClusters[min_index].members):
                value = 2000

        #convert labeled to Sparse
        sparseLabeledMoving = np.empty(len(movingImg), dtype=object)
        for i in range(len(movingImg)):
            sparseLabeledMoving[i] = sparse.csr_matrix(movingImg[i])

        memberList = []
        memberListWithZ = []
        while(len(memberListWithZ) == 0):
            for z in range(len(sparseLabeledMoving)):
                memberListWithoutZ = np.argwhere(sparseLabeledMoving[z] == value)
                memberListWithZ = [[z] + list(tup) for tup in memberListWithoutZ]
                memberList.extend(memberListWithZ)
            value = value - 1

        fixedCluster.timeRegistration = Cluster(memberList)
    '''

    return fixedClusters
<|endoftext|>


<commit_before>""" EPYNET Classes """
from . import epanet2 as ep
from .objectcollection import ObjectCollection
from .baseobject import BaseObject
import pudb

class Node(BaseObject):
    """ Base EPANET Node class """

    static_properties = {'elevation': ep.EN_ELEVATION}
    properties = {'head': ep.EN_HEAD, 'pressure': ep.EN_PRESSURE}

    def __init__(self, uid):
        super(Node, self).__init__(uid)
        self.links = ObjectCollection()

    def get_index(self, uid):
        return ep.ENgetnodeindex(uid)

    def set_object_value(self, uid, code, value):
        index = self.get_index(uid)
        return ep.ENsetnodevalue(index, code, value)

    def get_object_value(self, uid, code):
        index = self.get_index(uid)
        return ep.ENgetnodevalue(index, code)

    @property
    def index(self):
        return self.get_index(self.uid)

    @property
    def coordinates(self):
        if "coords" not in self._static_values.keys():
            self._static_values["coords"] = ep.ENgetcoord(self.index)
        return self._static_values["coords"]

    # extra functionality
    @property
    def upstream_links(self):
        """ return a list of upstream links """
        if self.results != {}:
            raise ValueError("This method is only supported for steady state simulations")
        links = ObjectCollection()
        for link in self.links:
            if (link.to_node == self and link.flow >= 0) or (link.from_node == self and link.flow < 0):
                links[link.uid] = link

        return links

    @property
    def downstream_links(self):
        """ return a list of downstream nodes """
        if self.results != {}:
            raise ValueError("This method is only supported for steady state simulations")

        links = ObjectCollection()
        for link in self.links:
            if (link.from_node == self and link.flow >= 1e-3) or (link.to_node == self and link.flow < -1e-3):
                links[link.uid] = link
        return links

    @property
    def inflow(self):
        outflow = 0
        for link in self.upstream_links:
            outflow += abs(link.flow)
        return outflow

    @property
    def outflow(self):
        outflow = 0
        for link in self.downstream_links:
            outflow += abs(link.flow)
        return outflow
        """ calculates all the water flowing out of the node """

class Reservoir(Node):
    """ EPANET Reservoir Class """
    node_type = "Reservoir"

class Junction(Node):
    """ EPANET Junction Class """
    static_properties = {'elevation': ep.EN_ELEVATION, 'basedemand': ep.EN_BASEDEMAND}
    properties = {'head': ep.EN_HEAD, 'pressure': ep.EN_PRESSURE, 'demand': ep.EN_DEMAND}
    node_type = "Junction"

    @property
    def pattern(self):
        pattern_index = self.get_property(ep.EN_PATTERN)
        uid = ep.ENgetpatternid(pattern_index)
        return Pattern(uid)

    @pattern.setter
    def pattern(self, value):
        if isinstance(value, int):
            pattern_index = value
        elif isinstance(value, str):
            pattern_index = ep.ENgetpatternindex(value)
        else:
            pattern_index = value.index

        self.set_object_value(self.uid, ep.EN_PATTERN, pattern_index)

class Tank(Node):
    """ EPANET Tank Class """
    node_type = "Tank"

    static_properties = {'elevation': ep.EN_ELEVATION, 'basedemand': ep.EN_BASEDEMAND,
                         'initvolume': ep.EN_INITVOLUME, 'diameter': ep.EN_TANKDIAM,
                         'minvolume': ep.EN_MINVOLUME, 'minlevel': ep.EN_MINLEVEL,
                         'maxlevel': ep.EN_MAXLEVEL, 'maxvolume': 25, 'tanklevel': ep.EN_TANKLEVEL}
    properties = {'head': ep.EN_HEAD, 'pressure': ep.EN_PRESSURE,
                  'demand': ep.EN_DEMAND, 'volume': 24, 'level': ep.EN_TANKLEVEL}
<commit_msg>removed pudb import<commit_after>""" EPYNET Classes """
from . import epanet2 as ep
from .objectcollection import ObjectCollection
from .baseobject import BaseObject

class Node(BaseObject):
    """ Base EPANET Node class """

    static_properties = {'elevation': ep.EN_ELEVATION}
    properties = {'head': ep.EN_HEAD, 'pressure': ep.EN_PRESSURE}

    def __init__(self, uid):
        super(Node, self).__init__(uid)
        self.links = ObjectCollection()

    def get_index(self, uid):
        return ep.ENgetnodeindex(uid)

    def set_object_value(self, uid, code, value):
        index = self.get_index(uid)
        return ep.ENsetnodevalue(index, code, value)

    def get_object_value(self, uid, code):
        index = self.get_index(uid)
        return ep.ENgetnodevalue(index, code)

    @property
    def index(self):
        return self.get_index(self.uid)

    @property
    def coordinates(self):
        if "coords" not in self._static_values.keys():
            self._static_values["coords"] = ep.ENgetcoord(self.index)
        return self._static_values["coords"]

    # extra functionality
    @property
    def upstream_links(self):
        """ return a list of upstream links """
        if self.results != {}:
            raise ValueError("This method is only supported for steady state simulations")
        links = ObjectCollection()
        for link in self.links:
            if (link.to_node == self and link.flow >= 0) or (link.from_node == self and link.flow < 0):
                links[link.uid] = link

        return links

    @property
    def downstream_links(self):
        """ return a list of downstream nodes """
        if self.results != {}:
            raise ValueError("This method is only supported for steady state simulations")

        links = ObjectCollection()
        for link in self.links:
            if (link.from_node == self and link.flow >= 1e-3) or (link.to_node == self and link.flow < -1e-3):
                links[link.uid] = link
        return links

    @property
    def inflow(self):
        outflow = 0
        for link in self.upstream_links:
            outflow += abs(link.flow)
        return outflow

    @property
    def outflow(self):
        outflow = 0
        for link in self.downstream_links:
            outflow += abs(link.flow)
        return outflow
        """ calculates all the water flowing out of the node """

class Reservoir(Node):
    """ EPANET Reservoir Class """
    node_type = "Reservoir"

class Junction(Node):
    """ EPANET Junction Class """
    static_properties = {'elevation': ep.EN_ELEVATION, 'basedemand': ep.EN_BASEDEMAND}
    properties = {'head': ep.EN_HEAD, 'pressure': ep.EN_PRESSURE, 'demand': ep.EN_DEMAND}
    node_type = "Junction"

    @property
    def pattern(self):
        pattern_index = self.get_property(ep.EN_PATTERN)
        uid = ep.ENgetpatternid(pattern_index)
        return Pattern(uid)

    @pattern.setter
    def pattern(self, value):
        if isinstance(value, int):
            pattern_index = value
        elif isinstance(value, str):
            pattern_index = ep.ENgetpatternindex(value)
        else:
            pattern_index = value.index

        self.set_object_value(self.uid, ep.EN_PATTERN, pattern_index)

class Tank(Node):
    """ EPANET Tank Class """
    node_type = "Tank"

    static_properties = {'elevation': ep.EN_ELEVATION, 'basedemand': ep.EN_BASEDEMAND,
                         'initvolume': ep.EN_INITVOLUME, 'diameter': ep.EN_TANKDIAM,
                         'minvolume': ep.EN_MINVOLUME, 'minlevel': ep.EN_MINLEVEL,
                         'maxlevel': ep.EN_MAXLEVEL, 'maxvolume': 25, 'tanklevel': ep.EN_TANKLEVEL}
    properties = {'head': ep.EN_HEAD, 'pressure': ep.EN_PRESSURE,
                  'demand': ep.EN_DEMAND, 'volume': 24, 'level': ep.EN_TANKLEVEL}
<|endoftext|>


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

import datetime
from uuid import UUID
import argparse

import elasticsearch

from bhs_api import create_app
from bhs_api import phonetic
from bhs_api.utils import uuids_to_str, SEARCHABLE_COLLECTIONS
from bhs_api.item import SHOW_FILTER

completion_field = {
                    "type": "keyword",
                    "fields": {
                        "suggest": {
                            "type": "completion",
                            "max_input_length": 20,
                            "contexts": [{
                                "name": "collection",
                                "type": "category",
                                "path": "_type"
                            }]
                        }
                    },
                }

header_mapping = {
                    "properties": {
                        "En": completion_field,
                        "He": completion_field,
                    }
                 }

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--collection',
                        help='run only on collection')
    parser.add_argument('-r', '--remove', action = "store_true",
                        help='remove the current index')
    parser.add_argument('--db',
                        help='the db to run on defaults to the value in /etc/bhs/config.yml')
    return parser.parse_args()

def add_phonetics(doc):
    if doc['Header']['En']:
        s = phonetic.get_english_dms(doc['Header']['En'])
    elif doc['Header']['He']:
        s = phonetic.get_hebrew_dms(doc['Header']['He'])
    else:
        s = 'BADWOLF'
    options = s.split(' ')
    doc['dm_soundex'] = options

if __name__ == '__main__':

    args = parse_args()
    app, conf = create_app()
    if args.db:
        db = app.client_data_db[args.db]
    else:
        db = app.data_db

    index_name = db.name
    # start with a clean index
    if args.remove:
        if app.es.indices.exists(index_name):
            app.es.indices.delete(index_name)
        # set the mapping to support completion fields
        app.es.indices.create(index_name, body={
            "mappings": {
                "places": { "properties": {
                    "Header": header_mapping,
                }},
                "familyNames": { "properties": {
                    "Header": header_mapping,
                    "dm_soundex": {
                        "type": "completion",
                        "max_input_length": 20,
                        "contexts": [{
                            "name": "collection",
                            "type": "category",
                            "path": "_type"
                        }]
                    }
                }},
            }
        })

    if args.collection:
        collections = [args.collection]
    else:
        collections = SEARCHABLE_COLLECTIONS

    for collection in collections:
        started = datetime.datetime.now()
        for doc in db[collection].find(SHOW_FILTER):
            _id = doc['_id']
            del doc['_id']
            del doc['UnitHeaderDMSoundex']
            # un null the fields that are used for completion
            if collection in ('places', 'familyNames'):
                add_phonetics(doc)
            # fill empty headers as es completion fails on null values
            header = doc['Header']
            for lang in ('En', 'He'):
                if not header[lang]:
                    header[lang] = '1234567890'
            try:
                res = app.es.index(index=index_name, doc_type=collection, id=_id, body=doc)
            except elasticsearch.exceptions.SerializationError:
                # UUID fields are causing es to crash, turn them to strings
                uuids_to_str(doc)
                try:
                    res = app.es.index(index=index_name, doc_type=collection, id=_id, body=doc)
                except elasticsearch.exceptions.SerializationError as e:
                    import pdb; pdb.set_trace()
            except elasticsearch.exceptions.RequestError as e:
                import pdb; pdb.set_trace()
        finished = datetime.datetime.now()
        print 'Collection {} took {}'.format(collection, finished-started)
<commit_msg>header should be text for better search results<commit_after>#!/usr/bin/env python

import datetime
from uuid import UUID
import argparse

import elasticsearch

from bhs_api import create_app
from bhs_api import phonetic
from bhs_api.utils import uuids_to_str, SEARCHABLE_COLLECTIONS
from bhs_api.item import SHOW_FILTER

completion_field = {
                    "type": "text",
                    "fields": {
                        "suggest": {
                            "type": "completion",
                            "max_input_length": 20,
                            "contexts": [{
                                "name": "collection",
                                "type": "category",
                                "path": "_type"
                            }]
                        }
                    },
                }

header_mapping = {
                    "properties": {
                        "En": completion_field,
                        "He": completion_field,
                    }
                 }

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--collection',
                        help='run only on collection')
    parser.add_argument('-r', '--remove', action = "store_true",
                        help='remove the current index')
    parser.add_argument('--db',
                        help='the db to run on defaults to the value in /etc/bhs/config.yml')
    return parser.parse_args()

def add_phonetics(doc):
    if doc['Header']['En']:
        s = phonetic.get_english_dms(doc['Header']['En'])
    elif doc['Header']['He']:
        s = phonetic.get_hebrew_dms(doc['Header']['He'])
    else:
        s = 'BADWOLF'
    options = s.split(' ')
    doc['dm_soundex'] = options

if __name__ == '__main__':

    args = parse_args()
    app, conf = create_app()
    if args.db:
        db = app.client_data_db[args.db]
    else:
        db = app.data_db

    index_name = db.name
    # start with a clean index
    if args.remove:
        if app.es.indices.exists(index_name):
            app.es.indices.delete(index_name)
        # set the mapping to support completion fields
        app.es.indices.create(index_name, body={
            "mappings": {
                "places": { "properties": {
                    "Header": header_mapping,
                }},
                "familyNames": { "properties": {
                    "Header": header_mapping,
                    "dm_soundex": {
                        "type": "completion",
                        "max_input_length": 20,
                        "contexts": [{
                            "name": "collection",
                            "type": "category",
                            "path": "_type"
                        }]
                    }
                }},
            }
        })

    if args.collection:
        collections = [args.collection]
    else:
        collections = SEARCHABLE_COLLECTIONS

    for collection in collections:
        started = datetime.datetime.now()
        for doc in db[collection].find(SHOW_FILTER):
            _id = doc['_id']
            del doc['_id']
            del doc['UnitHeaderDMSoundex']
            # un null the fields that are used for completion
            if collection in ('places', 'familyNames'):
                add_phonetics(doc)
            # fill empty headers as es completion fails on null values
            header = doc['Header']
            for lang in ('En', 'He'):
                if not header[lang]:
                    header[lang] = '1234567890'
            try:
                res = app.es.index(index=index_name, doc_type=collection, id=_id, body=doc)
            except elasticsearch.exceptions.SerializationError:
                # UUID fields are causing es to crash, turn them to strings
                uuids_to_str(doc)
                try:
                    res = app.es.index(index=index_name, doc_type=collection, id=_id, body=doc)
                except elasticsearch.exceptions.SerializationError as e:
                    import pdb; pdb.set_trace()
            except elasticsearch.exceptions.RequestError as e:
                import pdb; pdb.set_trace()
        finished = datetime.datetime.now()
        print 'Collection {} took {}'.format(collection, finished-started)
<|endoftext|>


<commit_before>from operator import __not__

def identity(x):
    return x

def constantly(x):
    return lambda *a, **kw: x

# an operator.methodcaller() brother
def caller(*a, **kw):
    return lambda f: f(*a, **kw)

# not using functools.partial to get real function
def partial(func, *args, **kwargs):
    def new_func(*fargs, **fkwargs):
        return func(*(args + fargs), **dict(kwargs, **fkwargs))
    new_func.func = func
    new_func.args = args
    new_func.kwargs = kwargs
    return new_func

def compose(*fs):
    pair = lambda f, g: lambda *a, **kw: f(g(*a, **kw))
    return reduce(pair, fs, identity)

def complement(f):
    return compose(__not__, f)

def juxt(*fs):
    return lambda *a, **kw: [f(*a, **kw) for f in fs]

from operator import __add__, __sub__
from whatever import _

def test_caller():
    assert caller([1, 2])(sum) == 3

def test_compose():
    double = _ * 2
    inc    = _ + 1
    assert compose(inc, double)(10) == 21
    assert compose(str, inc, double)(10) == '21'

def test_complement():
    assert complement(identity)(0) == True
    assert complement(identity)([1, 2]) == False

def test_partial():
    assert partial(__add__, 10)(1) == 11
    assert partial(__add__, 'abra')('cadabra') == 'abracadabra'

    merge = lambda a=None, b=None: a + b
    assert partial(merge, a='abra')(b='cadabra') == 'abracadabra'
    assert partial(merge, b='abra')(a='cadabra') == 'cadabraabra'

def test_juxt():
    assert juxt(__add__, __sub__)(10, 2) == [12, 8]
<commit_msg>Added a pair of funcs colls utils<commit_after>from operator import __not__
from itertools import ifilter

def identity(x):
    return x

def constantly(x):
    return lambda *a, **kw: x

# an operator.methodcaller() brother
def caller(*a, **kw):
    return lambda f: f(*a, **kw)

# not using functools.partial to get real function
def partial(func, *args, **kwargs):
    def new_func(*fargs, **fkwargs):
        return func(*(args + fargs), **dict(kwargs, **fkwargs))
    new_func.func = func
    new_func.args = args
    new_func.kwargs = kwargs
    return new_func

def compose(*fs):
    pair = lambda f, g: lambda *a, **kw: f(g(*a, **kw))
    return reduce(pair, fs, identity)

def complement(f):
    return compose(__not__, f)

def juxt(*fs):
    return lambda *a, **kw: [f(*a, **kw) for f in fs]

def ijuxt(*fs):
    return lambda *a, **kw: (f(*a, **kw) for f in fs)

# fnil ?

# NOTE: Should these ones be moved somewhere? func_colls? fcolls?
#       Should I call them every_pred, all_pred, any_pred ... ?
# every-pred in clojure
def all_fn(*fs):
    return compose(all, ijuxt(*fs))

def first(pred, coll=None):
    if coll is None:
        return first(None, pred)
    return next(ifilter(pred, coll), None)

# some-fn in clojure
def first_fn(*fs):
    return compose(first, ijuxt(*fs))

def any_fn(*fs):
    return compose(any, ijuxt(*fs))

from operator import __add__, __sub__
from whatever import _

def test_caller():
    assert caller([1, 2])(sum) == 3

def test_compose():
    double = _ * 2
    inc    = _ + 1
    assert compose(inc, double)(10) == 21
    assert compose(str, inc, double)(10) == '21'

def test_complement():
    assert complement(identity)(0) == True
    assert complement(identity)([1, 2]) == False

def test_partial():
    assert partial(__add__, 10)(1) == 11
    assert partial(__add__, 'abra')('cadabra') == 'abracadabra'

    merge = lambda a=None, b=None: a + b
    assert partial(merge, a='abra')(b='cadabra') == 'abracadabra'
    assert partial(merge, b='abra')(a='cadabra') == 'cadabraabra'

def test_juxt():
    assert juxt(__add__, __sub__)(10, 2) == [12, 8]

def test_first_fn():
    assert first_fn(_-1, _*0, _+1, _*2)(1) == 2
<|endoftext|>