<commit_before># Software License Agreement (BSD License)
#
# Copyright (c) 2009-2011, Eucalyptus Systems, Inc.
# All rights reserved.
#
# Redistribution and use of this software 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.
#
# 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 THE COPYRIGHT OWNER OR CONTRIBUTORS 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.
#
# Author: vic.iglesias@eucalyptus.com

import re

class Eunode:
    def __init__(self, hostname, partition,tester = None):
        self.hostname = hostname
        self.partition = partition
        if tester is not None:
            self.tester = tester
            self.machine = self.tester.get_component_ip(hostname)
            self.hypervisor = self.machine.sys("cat " + self.tester.eucapath + "/etc/eucalyptus/eucalyptus.conf | grep hypervisor").split("=").strip('"')

class Euservice:
    def __init__(self, service_string, tester = None):
        values = service_string.split()
        self.type = values[1]
        self.partition = values[2]
        self.name = values[3]
        self.state = values[4]
        self.uri = values[6]
        self.fullname = values[7]
        self.hostname = self.uri.split(":")[1].split("/")[2]
        self.running = True
        if tester is not None:
            self.tester = tester
            self.machine = tester.get_machine_by_ip(self.hostname)

    def isEnabled(self):
        if re.search('ENABLED',self.state) :
            return True
        else:
            return False

    def isDisabled(self):
        if not re.search('ENABLED', self.state):
            return True
        else:
            return False

    def disable(self):
        self.tester.service_manager.disable(self)

    def enable(self):
        self.tester.service_manager.enable(self)

    def stop(self):
        self.tester.service_manager.stop(self)

    def start(self):
        self.tester.service_manager.start(self)

class Partition:
    name = ""
    ccs = []
    scs = []
    vbs = []
    ncs = []
    volumes = []
    instances = []

    def __init__(self, name, service_manager ):
        self.name = name
        self.service_manager = service_manager

    def get_enabled(self, list):
        self.service_manager.update()
        for service in list:
            if service.isEnabled():
                return service
        return None

    def get_disabled(self, list):
        self.service_manager.update()
        for service in list:
            if not service.isEnabled():
                return service
        return None

    def get_enabled_cc(self):
        return self.get_enabled(self.ccs)

    def get_enabled_sc(self):
        return self.get_enabled(self.scs)

    def get_enabled_vb(self):
        return self.get_enabled(self.vbs)

    def get_disabled_cc(self):
        return self.get_disabled(self.ccs)

    def get_disabled_sc(self):
        return self.get_disabled(self.scs)

    def get_disabled_vb(self):
        return self.get_disabled(self.vbs)

class EuserviceManager(object):

    def __init__(self, tester ):
        '''
        SERVICE    storage            PARTI00            SC_61              ENABLED       16      http://192.168.51.32:8773/services/Storage    arn:euca:eucalyptus:PARTI00:storage:SC_61/
        update this service based up on the information parsed out of the "describestring"
        '''
        ### Make sure i have the right connection to make first contact with euca-describe-services
        self.walruses= []
        self.clcs = []
        self.arbitrators = []
        self.partitions = {}
        self.internal_components = []
        self.dns = None
        self.tester = tester
        self.eucaprefix = ". " + self.tester.credpath + "/eucarc && " + self.tester.eucapath
        if self.tester.clc is None:
            raise AttributeError("Tester object does not have CLC machine to use for SSH")

        self.update()

    def get(self, type=None, partition=None, attempt_both=True):
        if type is not None:
            type = " -T " + str(type)
        else:
            type = ""
        #### This is a hack around the fact that the -P filter is not working need to fix this once that functionality is fixed
        if partition is not None:
            partition = " | grep " + str(partition)
        else:
            partition = ""
        try:
            describe_services = self.tester.clc.sys(self.eucaprefix + "/usr/sbin/euca-describe-services " + str(type)  +  "| grep SERVICE" + str(partition), timeout=15)
            if len(describe_services) < 1:
                raise IndexError("Did not receive proper response from describe services when looking for " + str(type))
        except Exception, e:
            if len(self.tester.get_component_machines("clc")) is 1:
                raise Exception("Unable to get service information from the only clc: " + self.tester.clc.hostname )
            if attempt_both:
                self.tester.swap_clc()
                describe_services = self.tester.clc.sys(self.eucaprefix + "/usr/sbin/euca-describe-services " + str(type)  + " | grep SERVICE "  + str(partition)  , timeout=15)
                if len(describe_services) < 1:
                    raise IndexError("Did not receive proper response from describe services when looking for " + str(type))
            raise e

        self.populate_nodes()

        services = []
        for service_line in describe_services:
            services.append(Euservice(service_line, self.tester))
        return services

    def populate_nodes(self):
        clc = self.tester.get_enabled_clc()
        nodes_list = clc.sys("euca_conf --list-nodes")
        for node_string in node_lists:
            split_string = node_string.split()
            node = Eunode(split_string[1], split_string[2])
            for part in self.partitions:
                if node.patition is part.name:
                    part.ncs.append(node)

    def reset(self):
        self.walruses= []
        self.clcs = []
        self.arbitrators = []
        self.internal_components = []
        for k, v in self.partitions.iteritems():
            self.partitions[k].ccs = []
            self.partitions[k].scs = []
            self.partitions[k].vbs = []

    def get_all_services(self):
        all_services = []
        self.update()
        if len(self.clcs) > 0:
            all_services = all_services + self.clcs
        if len(self.walruses) > 0:
            all_services = all_services + self.walruses
        for partition in self.partitions.keys():

            ccs = self.partitions[partition].ccs
            if len(ccs) > 0:
                all_services = all_services + ccs

            scs = self.partitions[partition].scs
            if len(scs) > 0:
                all_services = all_services + scs

            vbs = self.partitions[partition].vbs
            if len(vbs) > 0:
                all_services = all_services + vbs

        return all_services

    def start_all(self):
        all_services = self.get_all_services()
        for service in all_services:
            try:
                service.start()
            except Exception, e:
                self.tester.debug("Caught an exception when trying to turn up a service. Probably was already running")

    def get_conclusive_enabled_clc(self):
        self.reset()
        try:
            first_clc = self.get("eucalyptus")
        except:
            first_clc = None

        self.tester.swap_clc()

        try:
            second_clc = self.get("eucalyptus")
        except:
            second_clc = None

        ### Only one responded properly
        if second_clc is None:
            return first_clc
        if first_clc is None:
            return second_clc

        ### Inconsistency among 2 CLCs
        if second_clc.host is first_clc.host:
                return first_clc
        else:
            raise Exception("Could not find a CLC to connect to")

    def sync_credentials(self):
        self.reset()

    def update(self, name=None):
        ### Get all services
        self.reset()
        services = self.get(name)
        for current_euservice in services:
            ### If this is system wide component add it to the base level array
            if re.search("eucalyptus", current_euservice.type) :
                self.clcs.append(current_euservice)
                continue
            elif re.search("walrus", current_euservice.type):
                self.walruses.append(current_euservice)
                continue
            elif re.search("dns", current_euservice.type):
                self.dns = current_euservice
                continue
            elif re.search("arbitrator", current_euservice.type):
                self.arbitrators.append(current_euservice)
                continue
            ### If this is a partition specific component nest the service in its partition array
            if re.search("eucalyptus", current_euservice.partition) or re.search("bootstrap", current_euservice.partition):
                self.internal_components.append(current_euservice)
                continue
            else:
                append = False
                if current_euservice.partition in self.partitions:
                    my_partition = self.partitions[current_euservice.partition]
                else:
                    ### First time accessing this partition
                    my_partition = Partition(current_euservice.partition, self)
                    append = True
                if re.search("cluster", current_euservice.type):
                    my_partition.ccs.append(current_euservice)
                if re.search("storage", current_euservice.type):
                    my_partition.scs.append(current_euservice)
                if re.search("vmwarebroker", current_euservice.type):
                    my_partition.vbs.append(current_euservice)
                if append:
                    self.partitions[current_euservice.partition] = my_partition

    def isReachable(self, address):
        return self.tester.ping(address)

    def modify_service(self, euservice, state):
        if not self.isReachable(self.tester.clc.hostname):
            self.tester.clc = self.tester.get_component_machines("clc")[1]
        modify_response = self.tester.clc.sys(self.eucaprefix + "/usr/sbin/euca-modify-service -s " + str(state)  + " " + euservice.name)
        if re.search("true",modify_response[0]):
            return True
        else:
            raise AssertionError("Response to modify service was not true: " + str(modify_response))

    def modify_process(self, euservice, command):
        service_name = "eucalyptus-cloud"
        if re.search("cluster", euservice.type):
            service_name = "eucalyptus-cc"
        if not euservice.machine.found(self.tester.eucapath + "/etc/init.d/" + service_name + " " + command, "done"):
            self.tester.fail("Was unable to stop service: " + euservice.name + " on host " + euservice.machine.hostname)
            raise Exception("Did not properly modify service")

    def stop(self, euservice):
        self.modify_process(euservice, "stop")
        euservice.running = False

    def start(self, euservice):
        self.modify_process(euservice, "start")
        euservice.running = True

    def enable(self,euservice):
        self.modify_service(euservice, "ENABLED")

    def disable(self,euservice):
        self.modify_service(euservice, "DISABLED")

    def wait_for_service(self, euservice, state = "ENABLED", attempt_both = True, timeout=600):
        interval = 60
        poll_count = timeout / interval
        while (poll_count > 0):
            matching_services = []
            try:
                matching_services = self.get(euservice.type, euservice.partition, attempt_both)
                for service in matching_services:
                    if re.search(state, service.state):
                        return service
            except Exception, e:
                self.tester.debug("Caught " + str(e) + " when trying to get services. Retrying in " + str(interval) + "s")
            poll_count -= 1
            self.tester.sleep(interval)

        if poll_count is 0:
            self.tester.fail("Service: " + euservice.name + " did not enter "  + state + " state")
            raise Exception("Service: " + euservice.name + " did not enter "  + state + " state")

    def all_services_operational(self):
        all_services = self.get_all_services()
        for service in all_services:
            self.wait_for_service(service,"ENABLED")

    def get_enabled_clc(self):
        clc = self.get_enabled(self.clcs)
        if clc is None:
            raise Exception("Neither CLC is enabled")
        else:
            return clc

    def get_disabled_clc(self):
        clc = self.get_disabled(self.clcs)
        if clc is None:
            raise Exception("Neither CLC is disabled")
        else:
            return clc

    def get_enabled_walrus(self):
        walrus = self.get_enabled(self.walruses)
        if walrus is None:
            raise Exception("Neither Walrus is enabled")
        else:
            return walrus

    def get_disabled_walrus(self):
        walrus = self.get_enabled(self.walruses)
        if walrus is None:
            raise Exception("Neither Walrus is disabled")
        else:
            return walrus

    def get_enabled(self, list_of_services):
        self.update()
        for service in list_of_services:
            if service.isEnabled():
                return service
        return None

    def get_disabled(self, list_of_services):
        self.update()
        for service in list_of_services:
            if service.isDisabled():
                return service
        return None

    <commit_msg>Fix to euservice implementation populate_nodes implementation, do not populate nodes by default<commit_after># Software License Agreement (BSD License)
#
# Copyright (c) 2009-2011, Eucalyptus Systems, Inc.
# All rights reserved.
#
# Redistribution and use of this software 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.
#
# 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 THE COPYRIGHT OWNER OR CONTRIBUTORS 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.
#
# Author: vic.iglesias@eucalyptus.com

import re

class Eunode:
    def __init__(self, hostname, partition,tester = None):
        self.hostname = hostname
        self.partition = partition
        if tester is not None:
            self.tester = tester
            self.machine = self.tester.get_component_ip(hostname)
            self.hypervisor = self.machine.sys("cat " + self.tester.eucapath + "/etc/eucalyptus/eucalyptus.conf | grep hypervisor").split("=").strip('"')

class Euservice:
    def __init__(self, service_string, tester = None):
        values = service_string.split()
        self.type = values[1]
        self.partition = values[2]
        self.name = values[3]
        self.state = values[4]
        self.uri = values[6]
        self.fullname = values[7]
        self.hostname = self.uri.split(":")[1].split("/")[2]
        self.running = True
        if tester is not None:
            self.tester = tester
            self.machine = tester.get_machine_by_ip(self.hostname)

    def isEnabled(self):
        if re.search('ENABLED',self.state) :
            return True
        else:
            return False

    def isDisabled(self):
        if not re.search('ENABLED', self.state):
            return True
        else:
            return False

    def disable(self):
        self.tester.service_manager.disable(self)

    def enable(self):
        self.tester.service_manager.enable(self)

    def stop(self):
        self.tester.service_manager.stop(self)

    def start(self):
        self.tester.service_manager.start(self)

class Partition:
    name = ""
    ccs = []
    scs = []
    vbs = []
    ncs = []
    volumes = []
    instances = []

    def __init__(self, name, service_manager ):
        self.name = name
        self.service_manager = service_manager

    def get_enabled(self, list):
        self.service_manager.update()
        for service in list:
            if service.isEnabled():
                return service
        return None

    def get_disabled(self, list):
        self.service_manager.update()
        for service in list:
            if not service.isEnabled():
                return service
        return None

    def get_enabled_cc(self):
        return self.get_enabled(self.ccs)

    def get_enabled_sc(self):
        return self.get_enabled(self.scs)

    def get_enabled_vb(self):
        return self.get_enabled(self.vbs)

    def get_disabled_cc(self):
        return self.get_disabled(self.ccs)

    def get_disabled_sc(self):
        return self.get_disabled(self.scs)

    def get_disabled_vb(self):
        return self.get_disabled(self.vbs)

class EuserviceManager(object):

    def __init__(self, tester ):
        '''
        SERVICE    storage            PARTI00            SC_61              ENABLED       16      http://192.168.51.32:8773/services/Storage    arn:euca:eucalyptus:PARTI00:storage:SC_61/
        update this service based up on the information parsed out of the "describestring"
        '''
        ### Make sure i have the right connection to make first contact with euca-describe-services
        self.walruses= []
        self.clcs = []
        self.arbitrators = []
        self.partitions = {}
        self.internal_components = []
        self.dns = None
        self.tester = tester
        self.eucaprefix = ". " + self.tester.credpath + "/eucarc && " + self.tester.eucapath
        if self.tester.clc is None:
            raise AttributeError("Tester object does not have CLC machine to use for SSH")

        self.update()

    def get(self, type=None, partition=None, attempt_both=True):
        if type is not None:
            type = " -T " + str(type)
        else:
            type = ""
        #### This is a hack around the fact that the -P filter is not working need to fix this once that functionality is fixed
        if partition is not None:
            partition = " | grep " + str(partition)
        else:
            partition = ""
        try:
            describe_services = self.tester.clc.sys(self.eucaprefix + "/usr/sbin/euca-describe-services " + str(type)  +  "| grep SERVICE" + str(partition), timeout=15)
            if len(describe_services) < 1:
                raise IndexError("Did not receive proper response from describe services when looking for " + str(type))
        except Exception, e:
            if len(self.tester.get_component_machines("clc")) is 1:
                raise Exception("Unable to get service information from the only clc: " + self.tester.clc.hostname )
            if attempt_both:
                self.tester.swap_clc()
                describe_services = self.tester.clc.sys(self.eucaprefix + "/usr/sbin/euca-describe-services " + str(type)  + " | grep SERVICE "  + str(partition)  , timeout=15)
                if len(describe_services) < 1:
                    raise IndexError("Did not receive proper response from describe services when looking for " + str(type))
            raise e

        #self.populate_nodes()

        services = []
        for service_line in describe_services:
            services.append(Euservice(service_line, self.tester))
        return services

    def populate_nodes(self):
        clc = self.get_enabled_clc()
        nodes_list = clc.machine.sys("euca_conf --list-nodes")
        for node_string in node_lists:
            split_string = node_string.split()
            node = Eunode(split_string[1], split_string[2])
            for part in self.partitions:
                if node.patition is part.name:
                    part.ncs.append(node)

    def reset(self):
        self.walruses= []
        self.clcs = []
        self.arbitrators = []
        self.internal_components = []
        for k, v in self.partitions.iteritems():
            self.partitions[k].ccs = []
            self.partitions[k].scs = []
            self.partitions[k].vbs = []

    def get_all_services(self):
        all_services = []
        self.update()
        if len(self.clcs) > 0:
            all_services = all_services + self.clcs
        if len(self.walruses) > 0:
            all_services = all_services + self.walruses
        for partition in self.partitions.keys():

            ccs = self.partitions[partition].ccs
            if len(ccs) > 0:
                all_services = all_services + ccs

            scs = self.partitions[partition].scs
            if len(scs) > 0:
                all_services = all_services + scs

            vbs = self.partitions[partition].vbs
            if len(vbs) > 0:
                all_services = all_services + vbs

        return all_services

    def start_all(self):
        all_services = self.get_all_services()
        for service in all_services:
            try:
                service.start()
            except Exception, e:
                self.tester.debug("Caught an exception when trying to turn up a service. Probably was already running")

    def get_conclusive_enabled_clc(self):
        self.reset()
        try:
            first_clc = self.get("eucalyptus")
        except:
            first_clc = None

        self.tester.swap_clc()

        try:
            second_clc = self.get("eucalyptus")
        except:
            second_clc = None

        ### Only one responded properly
        if second_clc is None:
            return first_clc
        if first_clc is None:
            return second_clc

        ### Inconsistency among 2 CLCs
        if second_clc.host is first_clc.host:
                return first_clc
        else:
            raise Exception("Could not find a CLC to connect to")

    def sync_credentials(self):
        self.reset()

    def update(self, name=None):
        ### Get all services
        self.reset()
        services = self.get(name)
        for current_euservice in services:
            ### If this is system wide component add it to the base level array
            if re.search("eucalyptus", current_euservice.type) :
                self.clcs.append(current_euservice)
                continue
            elif re.search("walrus", current_euservice.type):
                self.walruses.append(current_euservice)
                continue
            elif re.search("dns", current_euservice.type):
                self.dns = current_euservice
                continue
            elif re.search("arbitrator", current_euservice.type):
                self.arbitrators.append(current_euservice)
                continue
            ### If this is a partition specific component nest the service in its partition array
            if re.search("eucalyptus", current_euservice.partition) or re.search("bootstrap", current_euservice.partition):
                self.internal_components.append(current_euservice)
                continue
            else:
                append = False
                if current_euservice.partition in self.partitions:
                    my_partition = self.partitions[current_euservice.partition]
                else:
                    ### First time accessing this partition
                    my_partition = Partition(current_euservice.partition, self)
                    append = True
                if re.search("cluster", current_euservice.type):
                    my_partition.ccs.append(current_euservice)
                if re.search("storage", current_euservice.type):
                    my_partition.scs.append(current_euservice)
                if re.search("vmwarebroker", current_euservice.type):
                    my_partition.vbs.append(current_euservice)
                if append:
                    self.partitions[current_euservice.partition] = my_partition

    def isReachable(self, address):
        return self.tester.ping(address)

    def modify_service(self, euservice, state):
        if not self.isReachable(self.tester.clc.hostname):
            self.tester.clc = self.tester.get_component_machines("clc")[1]
        modify_response = self.tester.clc.sys(self.eucaprefix + "/usr/sbin/euca-modify-service -s " + str(state)  + " " + euservice.name)
        if re.search("true",modify_response[0]):
            return True
        else:
            raise AssertionError("Response to modify service was not true: " + str(modify_response))

    def modify_process(self, euservice, command):
        service_name = "eucalyptus-cloud"
        if re.search("cluster", euservice.type):
            service_name = "eucalyptus-cc"
        if not euservice.machine.found(self.tester.eucapath + "/etc/init.d/" + service_name + " " + command, "done"):
            self.tester.fail("Was unable to stop service: " + euservice.name + " on host " + euservice.machine.hostname)
            raise Exception("Did not properly modify service")

    def stop(self, euservice):
        self.modify_process(euservice, "stop")
        euservice.running = False

    def start(self, euservice):
        self.modify_process(euservice, "start")
        euservice.running = True

    def enable(self,euservice):
        self.modify_service(euservice, "ENABLED")

    def disable(self,euservice):
        self.modify_service(euservice, "DISABLED")

    def wait_for_service(self, euservice, state = "ENABLED", attempt_both = True, timeout=600):
        interval = 60
        poll_count = timeout / interval
        while (poll_count > 0):
            matching_services = []
            try:
                matching_services = self.get(euservice.type, euservice.partition, attempt_both)
                for service in matching_services:
                    if re.search(state, service.state):
                        return service
            except Exception, e:
                self.tester.debug("Caught " + str(e) + " when trying to get services. Retrying in " + str(interval) + "s")
            poll_count -= 1
            self.tester.sleep(interval)

        if poll_count is 0:
            self.tester.fail("Service: " + euservice.name + " did not enter "  + state + " state")
            raise Exception("Service: " + euservice.name + " did not enter "  + state + " state")

    def all_services_operational(self):
        all_services = self.get_all_services()
        for service in all_services:
            self.wait_for_service(service,"ENABLED")

    def get_enabled_clc(self):
        clc = self.get_enabled(self.clcs)
        if clc is None:
            raise Exception("Neither CLC is enabled")
        else:
            return clc

    def get_disabled_clc(self):
        clc = self.get_disabled(self.clcs)
        if clc is None:
            raise Exception("Neither CLC is disabled")
        else:
            return clc

    def get_enabled_walrus(self):
        walrus = self.get_enabled(self.walruses)
        if walrus is None:
            raise Exception("Neither Walrus is enabled")
        else:
            return walrus

    def get_disabled_walrus(self):
        walrus = self.get_enabled(self.walruses)
        if walrus is None:
            raise Exception("Neither Walrus is disabled")
        else:
            return walrus

    def get_enabled(self, list_of_services):
        self.update()
        for service in list_of_services:
            if service.isEnabled():
                return service
        return None

    def get_disabled(self, list_of_services):
        self.update()
        for service in list_of_services:
            if service.isDisabled():
                return service
        return None

    <|endoftext|>


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

from __future__ import division, print_function

import os
import time
from functools import partial
from collections import Counter

import h5py
import numpy as np
from scipy.stats import ks_2samp

import tqdm
from schwimmbad import MPIPool

from exoabc import Simulator, data

__all__ = []

period_range = (50, 300)
prad_range = (0.75, 2.5)
depth_range = (0, 1000)
maxn = 3

prefix = "q1_q16"
stlr = data.get_burke_gk(prefix=prefix)
kois = data.get_candidates(stlr=stlr, prefix=prefix)

sim = Simulator(
    stlr,
    period_range[0], period_range[1], 0.0,
    prad_range[0], prad_range[1], -2.0,
    -3.0, np.zeros(maxn),
    min_log_multi=-10.0, max_log_multi=20.0,
    release=prefix,
    seed=int(os.getpid() + 1000*time.time()) % 20000,
)

def compute_stats(catalog):
    m = (period_range[0] <= catalog.koi_period)
    m &= (catalog.koi_period <= period_range[1])
    m &= (prad_range[0] <= catalog.koi_prad)
    m &= (catalog.koi_prad <= prad_range[1])
    m &= (depth_range[0] <= catalog.koi_depth)
    m &= (catalog.koi_depth <= depth_range[1])
    c = catalog[m]

    # Multiplicity
    h = Counter(Counter(c.kepid).values())
    hist = np.zeros(maxn+1, dtype=int)
    for i in range(1, maxn+1):
        hist[i] = h.get(i, 0)
    hist[0] = len(stlr) - np.sum(hist[1:])

    return (
        hist, np.array(c.koi_period), np.array(c.koi_depth),
        np.array(c.koi_duration)
    )
obs_stats = compute_stats(kois)

def compute_distance(ds1, ds2):
    multi_dist = np.sum((np.log(ds1[0]+1) - np.log(ds2[0]+1))**2.0)
    period_dist = ks_2samp(ds1[1], ds2[1]).statistic
    depth_dist = ks_2samp(ds1[2], ds2[2]).statistic
    return multi_dist + period_dist + depth_dist

def sample(initial):
    if initial is None:
        lp = sim.sample_parameters()
        if not np.isfinite(lp):
            return np.inf, sim.get_parameters(), sim.state
    else:
        lp = sim.set_parameters(initial)
        if not np.isfinite(lp):
            return np.inf, sim.get_parameters(), sim.state

    pars, state = sim.get_parameters(), sim.state
    df = sim.sample_population()
    if len(df) <= 1:
        return np.inf, pars, state
    return compute_distance(obs_stats, compute_stats(df)), pars, state

def pmc_sample_one(eps, tau, theta0, weights, initial=None):
    # Sample until a suitable sample is found.
    rho = np.inf
    while rho > eps or not np.isfinite(rho):
        theta_star = theta0[np.random.choice(np.arange(len(weights)),
                                             p=weights)]
        theta_i = theta_star + tau * np.random.randn(len(theta_star))
        p, _, state_i = sample(theta_i)
        rho = np.sum(p)

    # Re-weight the samples.
    log_prior = sim.log_pdf()
    norm = 0.5*((theta0 - theta_i)/tau[None, :])**2 + np.log(tau[None, :])
    norm = np.log(weights) - np.sum(norm, axis=1)
    log_weight = log_prior - np.logaddexp.reduce(norm)
    return rho, theta_i, state_i, log_weight

def parse_samples(samples):
    rho = np.array([s[0] for s in samples])
    m = np.isfinite(rho)
    rho = rho[m]
    params = np.array([s[1] for s in samples])[m]
    states = np.array([s[2] for s in samples])[m]
    if len(samples[0]) == 3:
        return rho, params, states
    log_w = np.array([s[3] for s in samples])[m]
    return rho, params, states, np.exp(log_w - np.logaddexp.reduce(log_w))

def update_target_density(rho, params, weights, percentile=25.0):
    norm = np.sum(weights)
    mu = np.sum(params * weights[:, None], axis=0) / norm
    tau = np.sqrt(2 * np.sum((params-mu)**2*weights[:, None], axis=0) / norm)
    eps = np.percentile(rho, percentile)
    return eps, tau

with MPIPool() as pool:
    pool.wait()

    # Run step 1 of PMC method.
    N = 1000
    rhos, thetas, states = parse_samples(list(pool.map(
        sample, tqdm.tqdm((None for N in range(N)), total=N))))
    weights = np.ones(len(rhos)) / len(rhos)

    os.makedirs("results", exist_ok=True)
    for it in range(10):
        eps, tau = update_target_density(rhos, thetas, weights)
        func = partial(pmc_sample_one, eps, tau, thetas, weights)
        rhos, thetas, states, weights = parse_samples(list(pool.map(
            func, tqdm.tqdm((None for N in range(N)), total=N))))

        with h5py.File("results/{0:03d}.h5".format(it), "w") as f:
            f.attrs["iteration"] = it
            f.attrs["eps"] = eps
            f.attrs["tau"] = tau
            f.create_dataset("rho", data=rhos)
            f.create_dataset("theta", data=thetas)
            f.create_dataset("weight", data=weights)
<commit_msg>better plots<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import division, print_function

import os
import time
from functools import partial
from collections import Counter

import h5py
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import ks_2samp

import tqdm
import corner
from schwimmbad import MPIPool

from exoabc import Simulator, data

__all__ = []

period_range = (50, 300)
prad_range = (0.75, 2.5)
depth_range = (0, 1000)
maxn = 3

prefix = "q1_q16"
stlr = data.get_burke_gk(prefix=prefix)
kois = data.get_candidates(stlr=stlr, prefix=prefix)

sim = Simulator(
    stlr,
    period_range[0], period_range[1], 0.0,
    prad_range[0], prad_range[1], -2.0,
    -3.0, np.zeros(maxn),
    min_log_sigma=-5.0, max_log_sigma=5.0,
    min_log_multi=-5.0, max_log_multi=2.0,
    release=prefix,
    seed=int(os.getpid() + 1000*time.time()) % 20000,
)

def compute_stats(catalog):
    m = (period_range[0] <= catalog.koi_period)
    m &= (catalog.koi_period <= period_range[1])
    m &= (prad_range[0] <= catalog.koi_prad)
    m &= (catalog.koi_prad <= prad_range[1])
    m &= (depth_range[0] <= catalog.koi_depth)
    m &= (catalog.koi_depth <= depth_range[1])
    c = catalog[m]

    # Multiplicity
    h = Counter(Counter(c.kepid).values())
    hist = np.zeros(maxn+1, dtype=int)
    for i in range(1, maxn+1):
        hist[i] = h.get(i, 0)
    hist[0] = len(stlr) - np.sum(hist[1:])

    return (
        hist, np.array(c.koi_period), np.array(c.koi_depth),
        np.array(c.koi_duration)
    )
obs_stats = compute_stats(kois)

def compute_distance(ds1, ds2):
    multi_dist = np.mean((np.log(ds1[0]+1) - np.log(ds2[0]+1))**2.0)
    period_dist = ks_2samp(ds1[1], ds2[1]).statistic
    depth_dist = ks_2samp(ds1[2], ds2[2]).statistic
    return multi_dist + period_dist + depth_dist

def sample(initial):
    if initial is None:
        lp = sim.sample_parameters()
        if not np.isfinite(lp):
            return np.inf, sim.get_parameters(), sim.state
    else:
        lp = sim.set_parameters(initial)
        if not np.isfinite(lp):
            return np.inf, sim.get_parameters(), sim.state

    pars, state = sim.get_parameters(), sim.state
    df = sim.sample_population()
    if len(df) <= 1:
        return np.inf, pars, state
    return compute_distance(obs_stats, compute_stats(df)), pars, state

def pmc_sample_one(eps, tau, theta0, weights, initial=None):
    # Sample until a suitable sample is found.
    rho = np.inf
    while rho > eps or not np.isfinite(rho):
        theta_star = theta0[np.random.choice(np.arange(len(weights)),
                                             p=weights)]
        theta_i = theta_star + tau * np.random.randn(len(theta_star))
        p, _, state_i = sample(theta_i)
        rho = np.sum(p)

    # Re-weight the samples.
    log_prior = sim.log_pdf()
    norm = 0.5*((theta0 - theta_i)/tau[None, :])**2 + np.log(tau[None, :])
    norm = np.log(weights) - np.sum(norm, axis=1)
    log_weight = log_prior - np.logaddexp.reduce(norm)
    return rho, theta_i, state_i, log_weight

def parse_samples(samples):
    rho = np.array([s[0] for s in samples])
    m = np.isfinite(rho)
    rho = rho[m]
    params = np.array([s[1] for s in samples])[m]
    states = np.array([s[2] for s in samples])[m]
    if len(samples[0]) == 3:
        return rho, params, states
    log_w = np.array([s[3] for s in samples])[m]
    return rho, params, states, np.exp(log_w - np.logaddexp.reduce(log_w))

def update_target_density(rho, params, weights, percentile=30.0):
    norm = np.sum(weights)
    mu = np.sum(params * weights[:, None], axis=0) / norm
    tau = np.sqrt(2 * np.sum((params-mu)**2*weights[:, None], axis=0) / norm)
    eps = np.percentile(rho, percentile)
    return eps, tau

with MPIPool() as pool:
    pool.wait()

    from cycler import cycler
    from matplotlib import rcParams

    rcParams["font.size"] = 16
    rcParams["font.family"] = "sans-serif"
    rcParams["font.sans-serif"] = ["Computer Modern Sans"]
    rcParams["text.usetex"] = True
    rcParams["text.latex.preamble"] = r"\usepackage{cmbright}"
    rcParams["axes.prop_cycle"] = cycler("color", (
        "#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd", "#8c564b",
        "#e377c2", "#7f7f7f", "#bcbd22", "#17becf",
    ))  # d3.js color cycle

    # Run step 1 of PMC method.
    N = 1500
    rhos, thetas, states = parse_samples(list(pool.map(
        sample, tqdm.tqdm((None for N in range(N)), total=N))))
    weights = np.ones(len(rhos)) / len(rhos)

    os.makedirs("results", exist_ok=True)
    stlr.to_hdf(os.path.join("results", "stlr.h5"), "stlr", format="t")
    kois.to_hdf(os.path.join("results", "kois.h5"), "kois", format="t")
    for it in range(10):
        eps, tau = update_target_density(rhos, thetas, weights)
        func = partial(pmc_sample_one, eps, tau, thetas, weights)
        rhos, thetas, states, weights = parse_samples(list(pool.map(
            func, tqdm.tqdm((None for N in range(N)), total=N))))

        with h5py.File(os.path.join("results", "{0:03d}.h5".format(it)),
                       "w") as f:
            f.attrs["iteration"] = it
            f.attrs["eps"] = eps
            f.attrs["tau"] = tau
            for i in range(len(obs_stats)):
                f.attrs["obs_stats_{0}".format(i)] = obs_stats[i]
            f.create_dataset("rho", data=rhos)
            f.create_dataset("theta", data=thetas)
            f.create_dataset("weight", data=weights)
            f.create_dataset("state", data=states)

        fig = corner.corner(thetas, weights=weights)
        fig.savefig(os.path.join("results", "corner-{0:03d}.png".format(it)))
        plt.close(fig)

        fig, axes = plt.subplots(2, 3, figsize=(14, 8))

        # Observed distributions
        for i in np.random.choice(len(weights), p=weights, size=100):
            p = thetas[i]
            sim.set_parameters(p)
            sim.state = states[i]
            pop = sim.sample_population()
            sim_stats = compute_stats(pop)

            axes[0, 0].hist(sim_stats[1], range=period_range, histtype="step",
                            color="k", alpha=0.2)
            axes[0, 1].hist(sim_stats[2], range=depth_range, histtype="step",
                            color="k", alpha=0.2)
            axes[0, 2].plot(sim_stats[0], color="k", alpha=0.2)

        axes[0, 0].hist(obs_stats[1], range=period_range, histtype="step",
                        color="g", lw=2)
        axes[0, 1].hist(obs_stats[2], range=depth_range, histtype="step",
                        color="g", lw=2)
        axes[0, 2].plot(obs_stats[0], color="g", lw=2)
        axes[0, 2].set_yscale("log")
        axes[0, 0].set_xlabel("period")
        axes[0, 1].set_xlabel("depth")
        axes[0, 2].set_xlabel("multiplicity")
        axes[0, 0].set_yticklabels([])
        axes[0, 1].set_yticklabels([])
        axes[0, 0].set_ylabel("observed distributions")

        # True distributions
        for n, rng, ax in zip(thetas[:, :2].T, (period_range, prad_range),
                              axes[1, :2]):
            x = np.linspace(rng[0], rng[1], 5000)
            norm = (n + 1) / (rng[1]**(n+1) - rng[0]**(n+1))
            d = x[:, None]**n[None, :] * norm[None, :]
            q = np.percentile(d, [16, 50, 84], axis=1)
            ax.fill_between(x, q[0], q[2], color="k", alpha=0.1)
            ax.plot(x, q[1], color="k", lw=2)
            ax.set_xlim(*rng)

        n = np.concatenate((np.ones((len(thetas), 1)),
                            np.exp(thetas[:, -maxn:])),
                           axis=1)
        n /= np.sum(n, axis=1)[:, None]
        q = np.percentile(n, [16, 50, 84], axis=0)
        ax = axes[1, 2]
        ax.fill_between(np.arange(maxn+1), q[0], q[2], color="k", alpha=0.1)
        ax.plot(np.arange(maxn+1), q[1], color="k", lw=2)
        ax.set_xlim(0, maxn)

        axes[1, 2].set_yscale("log")
        axes[1, 0].set_xlabel("period")
        axes[1, 1].set_xlabel("radius")
        axes[1, 2].set_xlabel("multiplicity")
        axes[1, 0].set_yticklabels([])
        axes[1, 1].set_yticklabels([])
        axes[1, 0].set_ylabel("underlying distributions")

        fig.tight_layout()
        fig.savefig(os.path.join("results", "params-{0:03d}.png".format(it)),
                    bbox_inches="tight")
        plt.close(fig)
<|endoftext|>


<commit_before>import peer, curses

class Demo:
    def run(self, commandline):
        """ Runs the demo
        """
        commandline.command("help")
        commandline.command("gettracker 'Missouri S&T.jpg'")

if __name__ == "__main__":
    demo = Demo()
    curses.wrapper(peer.main, demo, "clientThreadConfig.cfg")
<commit_msg>Removed demo.py<commit_after><|endoftext|>


<commit_before>#!/usr/bin/python3.4
# -*-coding:Utf-8 -*
'''module to manage working path'''
import xml.etree.ElementTree as xmlMod
import os, re
from shutil import rmtree as rmdir
from usefullFunctions import indexPrintList, XML

class Output:
	'''class to manage working path'''

	def __init__(self, xml= None):
		'''initialize output path with default value or values extracted from an xml object'''
		if xml is None:
			if not os.path.exists('/home/'+os.getlogin()+'/.BlenderRenderManager/render'):
				os.mkdir('/home/'+os.getlogin()+'/.BlenderRenderManager/work')
			self.path = '/home/'+os.getlogin()+'/.BlenderRenderManager/work/'
		else:
			self.fromXml(xml)

	def fromXml(self, xml):
		'''initialize output path with value extracted from an xml object'''
		self.path = xml.get('path')

	def toXml(self):
		'''export output path into xml syntaxed string'''
		return '<output path="'+XML.encode(self.path)+'" />\n'

	def menu(self, log):
		'''method to see and edit output path'''
		log.menuIn('Output')

		while True:
			log.print()
			print('\n')

			print('Work Path :\n'+self.path+'\n\n new path? (empty to keep current)')

			new = input().strip()

			if new in [ '', '0', 'q' ]:
				log.menuOut()
				return False
			elif self.set(new, log) :
				log.menuOut()
				return True

	def set(self, path, log):
		'''method to edit output path'''

		# remove ' and/or "
		if path[0] in ['\'', '"'] and path[-1] == path[0]:
			path  = path[1:len(path)-1]

		# check if it's absolute path
		if path[0] != '/':
			log.error('The path must be absolute (begin by «/»)!')
			return False

		# check path exist
		if not os.path.exists(path):
			log.error('This path correspond to nothing!')
			return False

		# check path is a directory
		if not os.path.isdir(path):
			log.error('This path don\'t correspond to a directory!')
			return False

		# check path is writable
		if not os.access(path, os.W_OK):
			log.error('You don\'t have the permission to write in this directory!')
			return False

		if path[-1] != '/':
			path += '/'

		# apply path settings and confirm
		self.path = path
		log.write('Work path set to : '+self.path)
		return True

<commit_msg>move old working directory content into the new one<commit_after>#!/usr/bin/python3.4
# -*-coding:Utf-8 -*
'''module to manage working path'''
import xml.etree.ElementTree as xmlMod
import os, re
from shutil import rmtree as rmdir
from shutil import move as movedir
from usefullFunctions import indexPrintList, XML

class Output:
	'''class to manage working path'''

	def __init__(self, xml= None):
		'''initialize output path with default value or values extracted from an xml object'''
		if xml is None:
			if not os.path.exists('/home/'+os.getlogin()+'/.BlenderRenderManager/render'):
				os.mkdir('/home/'+os.getlogin()+'/.BlenderRenderManager/work')
			self.path = '/home/'+os.getlogin()+'/.BlenderRenderManager/work/'
		else:
			self.fromXml(xml)

	def fromXml(self, xml):
		'''initialize output path with value extracted from an xml object'''
		self.path = xml.get('path')

	def toXml(self):
		'''export output path into xml syntaxed string'''
		return '<output path="'+XML.encode(self.path)+'" />\n'

	def menu(self, log):
		'''method to see and edit output path'''
		log.menuIn('Output')

		while True:
			log.print()
			print('\n')

			print('Work Path :\n'+self.path+'\n\n new path? (empty to keep current)')

			new = input().strip()

			if new in [ '', '0', 'q' ]:
				log.menuOut()
				return False
			elif self.set(new, log) :
				log.menuOut()
				return True

	def set(self, path, log):
		'''method to edit output path'''

		# remove ' and/or "
		if path[0] in ['\'', '"'] and path[-1] == path[0]:
			path  = path[1:len(path)-1]

		# check if it's absolute path
		if path[0] != '/':
			log.error('The path must be absolute (begin by «/»)!')
			return False

		# check path exist
		if not os.path.exists(path):
			log.error('This path correspond to nothing!')
			return False

		# check path is a directory
		if not os.path.isdir(path):
			log.error('This path don\'t correspond to a directory!')
			return False

		# check path is writable
		if not os.access(path, os.W_OK):
			log.error('You don\'t have the permission to write in this directory!')
			return False

		if path[-1] != '/':
			path += '/'

		# apply path settings and confirm
		old = self.path
		self.path = path
		log.write('Work path set to : '+self.path)

		confirm = input('move old working directorie contert into the new one? (empty or y to confirm)').strip()
		if confirm in [ '', 'y' ]:
			for f in os.listdir(old):
				movedir(old+f, path)
			log.write('«'+old+'» content have been moved to «'+self.path+'»')
		else:
			log.write('«'+old+'» content haven\'t been moved!')

		return True

<|endoftext|>


<commit_before>import numpy as np
np.seterr(divide='ignore') # these warnings are usually harmless for this code
from matplotlib import pyplot as plt

import hsmm
from basic_distributions.observations import gaussian
from basic_distributions.durations import poisson
from util.text import progprint_xrange

#### Data generation
# Set parameters
N = 4
T = 500
obs_dim = 2
# Set duration parameters to be (10,20,30,40)
durparams = [10.*(idx+1) for idx in xrange(N)]
# Set observation hyperparameters (which control the random mean and covariance
# matrices for each state)
obs_hypparams = {'mu_0':np.zeros(obs_dim),
                'lmbda_0':np.eye(obs_dim),
                'kappa_0':0.1,
                'nu_0':obs_dim+2}

# Construct the true observation and duration distributions
truth_obs_distns = [gaussian(**obs_hypparams) for state in xrange(N)]
truth_dur_distns = [poisson(lmbda=param) for param in durparams]

# Build the true HSMM model
truthmodel = hsmm.hsmm(T,truth_obs_distns,truth_dur_distns)

# Sample data from the true model
data, labels = truthmodel.generate()

#### Posterior inference
# Set the weak limit truncation level. This is essentially the maximum
# number of states that can be learned
Nmax = 10

# Construct the observation and duration distribution objects, which set
# priors over parameters and then infer parameter values.
obs_distns = [gaussian(**obs_hypparams) for state in xrange(Nmax)]
dur_distns = [poisson() for state in xrange(Nmax)]

# Build the HSMM model that will represent the posterior
posteriormodel = hsmm.hsmm(T,obs_distns,dur_distns)

# Resample the model 100 times, printing a dot at each iteration
plot_every = 50
for idx in progprint_xrange(101):
    if (idx % plot_every) == 0:
        posteriormodel.plot(data)
        plt.gcf().suptitle('inferred HSMM after %d Gibbs iterations' % idx)

    posteriormodel.resample(data)
<commit_msg>minor tweaks<commit_after>from __future__ import division
import numpy as np
np.seterr(divide='ignore') # these warnings are usually harmless for this code
from matplotlib import pyplot as plt

import hsmm
from basic_distributions.observations import gaussian
from basic_distributions.durations import poisson
from util.text import progprint_xrange

#### Data generation
# Set parameters
N = 4
T = 500
obs_dim = 2
# Set duration parameters to be (10,20,30,40)
durparams = [10.*(idx+1) for idx in xrange(N)]
# Set observation hyperparameters (which control the random mean and covariance
# matrices for each state)
obs_hypparams = {'mu_0':np.zeros(obs_dim),
                'lmbda_0':np.eye(obs_dim),
                'kappa_0':0.1,
                'nu_0':obs_dim+2}

# Construct the true observation and duration distributions
truth_obs_distns = [gaussian(**obs_hypparams) for state in xrange(N)]
truth_dur_distns = [poisson(lmbda=param) for param in durparams]

# Build the true HSMM model
truthmodel = hsmm.hsmm(T,truth_obs_distns,truth_dur_distns)

# Sample data from the true model
data, labels = truthmodel.generate()

# Plot the truth
truthmodel.plot(data)
plt.gcf().suptitle('True HSMM')

#### Posterior inference
# Set the weak limit truncation level. This is essentially the maximum
# number of states that can be learned
Nmax = 10

# Construct the observation and duration distribution objects, which set
# priors over parameters and then infer parameter values.
obs_distns = [gaussian(**obs_hypparams) for state in xrange(Nmax)]
dur_distns = [poisson() for state in xrange(Nmax)]

# Build the HSMM model that will represent the posterior
posteriormodel = hsmm.hsmm(T,obs_distns,dur_distns)

# Resample the model 100 times, printing a dot at each iteration
# and plotting every so often
plot_every = 50
for idx in progprint_xrange(101):
    if (idx % plot_every) == 0:
        posteriormodel.plot(data)
        plt.gcf().suptitle('inferred HSMM after %d iterations (arbitrary colors)' % idx)

    posteriormodel.resample(data)
<|endoftext|>


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

log = logging.getLogger(__name__)

def ext_str(value, rule_obj, path):
    log.debug("value: %s", value)
    log.debug("rule_obj: %s", rule_obj)
    log.debug("path: %s", path)

    # Either raise some exception that you have defined your self
    # raise AssertionError('Custom assertion error in jinja_function()')

    # Or you should return True/False that will tell if it validated
    return True

def ext_list(value, rule_obj, path):
    log.debug("value: %s", value)
    log.debug("rule_obj: %s", rule_obj)
    log.debug("path: %s", path)

    # Either raise some exception that you have defined your self
    # raise AssertionError('Custom assertion error in jinja_function()')

    # Or you should return True/False that will tell if it validated
    return True

def ext_map(value, rule_obj, path):
    log.debug("value: %s", value)
    log.debug("rule_obj: %s", rule_obj)
    log.debug("path: %s", path)

    # Either raise some exception that you have defined your self
    # raise AssertionError('Custom assertion error in jinja_function()')

    # Or you should return True/False that will tell if it validated
    return True
<commit_msg>remove file<commit_after><|endoftext|>


<commit_before>""" Define a grid

This grid is the basis of different used models.
"""
from __future__ import print_function, division

from numbers import Integral

import numpy as np

from .quaternion import Quaternion
from .supercell import SuperCell, SuperCellChild
from .atom import Atom
from .geometry import Geometry

__all__ = ['Grid']

# Default nsc variable
_nsc = np.array([1]*3,np.int32)
_size = np.array([0]*3,np.int32)
_origo = np.array([0]*3,np.float64)
_bc = 1 # Periodic (Grid.Periodic)
_dtype = np.float64

class Grid(SuperCellChild):
    """ Object to retain grid information

    This grid object handles cell vectors and divisions of said grid.

    A grid can be periodic and non-periodic.
    """

    # Constant (should never be changed)
    Periodic = 1
    Neumann = 2
    Dirichlet = 3

    def __init__(self,size=_size,bc=_bc,sc=None,dtype=_dtype,geom=None):
        """ Initialize a `Grid` object.

        Initialize a `Grid` object.
        """

        self.set_supercell(sc)

        # Create the grid
        self.set_grid(size,dtype=dtype)

        # Create the grid boundary conditions
        self.set_bc(bc)

        # Create the atomic structure in the grid, if possible
        self.set_geom(geom)

    def __getitem__(self,key):
        """ Returns the grid contained """
        return self.grid[key]

    def __setitem__(self,key,val):
        """ Updates the grid contained """
        self.grid[key] = val

    def interp(self,size,*args,**kwargs):
        """ Returns an interpolated version of the grid

        Parameters
        ----------
        size : int, array_like
            the new size of the grid
        *args, **kwargs :
            optional arguments passed to the interpolation algorithm
            The interpolation routine is `scipy.interpolate.interpn`
        """
        # Get grid spacing
        dold = []
        dnew = []
        for i in range(3):
            dold.append(np.linspace(0,1,self.size[i]))
            dnew.append(np.linspace(0,1,size[i]))

        # Interpolate
        from scipy.interpolate import interpn

        # Create new grid
        grid = self.__class__(size, bc=np.copy(self.bc), sc=self.sc.copy())
        grid.grid = interpn(dold, self.grid, dnew, *args, **kwargs)

        return grid

    @property
    def size(self):
        """ Returns size of the grid """
        return np.asarray(self.grid.shape,np.int32)

    def set_grid(self,size,dtype=_dtype):
        """ Create the internal grid of certain size.
        """
        self.grid = np.empty(size,dtype=dtype)

    def set_bc(self,boundary=None,a=None,b=None,c=None):
        """ Set the boundary conditions on the grid
        boundary: [3], integer, optional
           boundary condition for all boundaries (or the same for all)
        a: integer, optional
           boundary condition for the first unit-cell vector direction
        b: integer, optional
           boundary condition for the second unit-cell vector direction
        c: integer, optional
           boundary condition for the third unit-cell vector direction
        """
        if not boundary is None:
            if isinstance(boundary,Integral):
                self.bc = np.array([boundary]*3,np.int32)
            else:
                self.bc = np.asarray(boundary,np.int32)
        if not a is None: self.bc[0] = a
        if not b is None: self.bc[1] = b
        if not c is None: self.bc[2] = c

    # Aliases
    set_boundary = set_bc
    set_boundary_condition = set_bc

    def set_geom(self,geom):
        """ Sets the `Geometry` for the grid.

        Setting the `Geometry` for the grid is a possibility
        to attach atoms to the grid.

        It is not a necessary entity.
        """
        if geom is None:
            # Fake geometry
            self.set_geom(Geometry([0,0,0],Atom['H'],sc=self.sc))
        else:
            self.geom = geom
            self.set_sc(geom.sc)

    def copy(self):
        """
        Returns a copy of the object.
        """
        grid = self.__class__(np.copy(self.size), bc=np.copy(self.bc),
                              sc=self.sc.copy(),dtype=self.grid.dtype)
        grid.grid[:,:,:] = self.grid[:,:,:]
        return grid

    def swapaxes(self,a,b):
        """ Returns Grid with swapped axis

        If ``swapaxes(0,1)`` it returns the 0 in the 1 values.
        """
        # Create index vector
        idx = np.arange(3)
        idx[b] = a
        idx[a] = b
        s = np.copy(self.size)
        grid = self.__class__(s[idx], bc=self.bc[idx],
                              sc=self.sc.swapaxes(a,b),dtype=self.grid.dtype)
        # We need to force the C-order or we loose the contiguoity
        grid.grid = np.copy(np.swapaxes(self.grid,a,b),order='C')
        return grid

    @property
    def dcell(self):
        """ Returns the delta-cell """
        # Calculate the grid-distribution
        g_size = self.size
        dcell = np.empty([3,3],np.float64)
        for ix in range(3):
            dcell[ix,:] = self.cell[ix,:] / g_size[ix]
        return dcell

    @property
    def dvol(self):
        """ Returns the delta-volume """
        return self.sc.vol / np.prod(self.size)

    def cross_section(self,idx,axis):
        """ Takes a cross-section of the grid along axis `axis`

        Remark: This API entry might change to handle arbitrary
        cuts via rotation of the axis """

        # First calculate the new size
        size = self.size
        cell = np.copy(self.cell)
        # Down-scale cell
        cell[axis,:] /= size[axis]
        size[axis] = 1
        grid = self.__class__(size, bc=np.copy(self.bc),geom=self.geom.copy())
        # Update cell size (the cell is smaller now)
        grid.set_sc(cell)

        if axis == 0:
            grid.grid[:,:,:] = self.grid[idx,:,:]
        elif axis == 1:
            grid.grid[:,:,:] = self.grid[:,idx,:]
        elif axis == 2:
            grid.grid[:,:,:] = self.grid[:,:,idx]
        else:
            raise ValueError('Unknown axis specification in cross_section')

        return grid

    def sum(self,axis):
        """ Returns the grid summed along axis `axis`. """
        # First calculate the new size
        size = self.size
        cell = np.copy(self.cell)
        # Down-scale cell
        cell[axis,:] /= size[axis]
        size[axis] = 1

        grid = self.__class__(size, bc=np.copy(self.bc),geom=self.geom.copy())
        # Update cell size (the cell is smaller now)
        grid.set_sc(cell)

        # Calculate sum
        grid.grid[:,:,:] = np.sum(self.grid,axis=axis)
        return grid

    def average(self,axis):
        """ Returns the average grid along direction `axis` """
        n = self.size[axis]
        return self.sum(axis) / float(n)

    def sub(self,idx,axis):
        """ Retains certain indices from a specified axis.

        Parameters
        ----------
        idx : array_like
           the indices of the grid axis `axis` to be retained
        axis : int
           the axis segment from which we retain the indices `idx`
        """
        uidx = np.unique(np.clip(idx,0,self.size[axis]-1))

        # Calculate new size
        size = self.size
        cell = np.copy(self.cell)
        old_N = size[axis]

        # Calculate new size
        size[axis] = len(idx)
        if size[axis] < 1:
            raise ValueError('You cannot retain no indices.')

        # Down-scale cell
        cell[axis,:] = cell[axis,:] / old_N * size[axis]

        grid = self.__class__(size, bc=np.copy(self.bc), geom=self.geom.copy())
        # Update cell size (the cell is smaller now)
        grid.set_sc(cell)

        # Remove the indices
        # First create the opposite, index
        if axis == 0:
            grid.grid[:,:,:] = self.grid[idx,:,:]
        elif axis == 1:
            grid.grid[:,:,:] = self.grid[:,idx,:]
        elif axis == 2:
            grid.grid[:,:,:] = self.grid[:,:,idx]

        return grid

    def remove(self,idx,axis):
        """ Removes certain indices from a specified axis.

        Parameters
        ----------
        idx : array_like
           the indices of the grid axis `axis` to be removed
        axis : int
           the axis segment from which we remove all indices `idx`
        """
        uidx = np.unique(np.clip(idx,0,self.size[axis]-1))
        ret_idx = np.setdiff1d(np.arange(self.size[axis]), uidx, assume_unique=True)
        return self.sub(ret_idx, axis)

    def index(self,coord,axis=None):
        """ Returns the index along the axis `axis` where `coord` exists

        Parameters
        ----------
        coord : array_like / float
           the coordinate of the axis
        axis : int
           the axis direction of the index
        """

        # if the axis is none, we do this for all axes
        if axis is None:
            # Loop over each direction
            idx = np.empty([3],np.int32)
            for i in [0,1,2]:
                # Normalize cell vector
                ca = self.cell[i,:] / np.sum(self.cell[i,:]**2)**.5
                # get the coordinate along the direction of the cell vector
                c = np.dot(self.cell[i,:],coord) * coord
                # Calculate the index corresponding to this placement
                idx[i] = self.index(c,i)
            return idx

        # Ensure a 1D array
        ac = np.atleast_1d(np.asarray(coord,np.float64))

        # Calculate the index of the coord in the cell
        dax = self.dcell[axis,:]

        # Calculate how many indices are required to fulfil
        # the correct line cut
        return int(np.rint( (np.sum(ac ** 2) / np.sum(dax ** 2)) **.5 ))

    def append(self,other,axis):
        """ Appends other `Grid` to this grid along axis

        """
        size = np.copy(self.size)
        size[axis] += other.size[axis]
        return self.__class__(size, bc = np.copy(self.bc),
                              sc=self.sc.append(other.sc,axis))

    def write(self,sile):
        """ Writes grid to the ``Sile`` using ``sile.write_grid``

        Parameters
        ----------
        sile : Sile, str
            a ``Sile`` object which will be used to write the grid
            if it is a string it will create a new sile using ``get_sile``
        """

        # This only works because, they *must*
        # have been imported previously
        from sids.io import get_sile, BaseSile
        if isinstance(sile,BaseSile):
            sile.write_grid(self)
        else:
            get_sile(sile,'w').write_grid(self)

    def __repr__(self):
        """ Representation of object """
        return 'Grid[{} {} {}]'.format(*self.size)

    def _check_compatibility(self,other,msg):
        """ Internal check for asserting two grids are commensurable """
        if self == other:
            return True
        else:
            raise ValueError('Grids are not compatible, '+msg)

    def _compatible_copy(self,other,*args,**kwargs):
        """ Returns a copy of self with an additional check of commensurable """
        if isinstance(other,Grid):
            if self._check_compatibility(other,*args,**kwargs):
                return self.copy()
        else:
            return self.copy()

    def __eq__(self,other):
        """ Returns true if the two grids are commensurable

        There will be no check of the values _on_ the grid. """
        return bool(np.all(self.size == other.size))

    def __ne__(self,other):
        """ Returns whether two grids have the same size """
        return not (self == other)

    def __add__(self,other):
        """ Returns a new grid with the addition of two grids

        Returns same size with same cell as the first"""
        if self == other:
            # We can return a new grid
            grid = self.copy()
            grid.grid = self.grid + other.grid
            return grid
        raise ValueError('Grids are not compatible, they cannot be added.')

    def __sub__(self,other):
        """ Returns a new grid with the difference of two grids

        Returns same size with same cell as the first"""
        grid = self._compatible_copy(other,'they cannot be subtracted')
        if isinstance(other,Grid):
            grid.grid = self.grid - other.grid
        else:
            grid.grid = self.grid - other
        return grid

    def __truediv__(self,other):
        grid = self._compatible_copy(other,'they cannot be divided')
        if isinstance(other,Grid):
            grid.grid = self.grid / other.grid
        else:
            grid.grid = self.grid / other
        return grid

    def __mul__(self,other):
        grid = self._compatible_copy(other,'they cannot be multiplied')
        if isinstance(other,Grid):
            grid.grid = self.grid * other.grid
        else:
            grid.grid = self.grid * other
        return grid

if __name__ == "__main__":
    pass
<commit_msg>bug: Fixed copy of grid to retain geometry<commit_after>""" Define a grid

This grid is the basis of different used models.
"""
from __future__ import print_function, division

from numbers import Integral

import numpy as np

from .quaternion import Quaternion
from .supercell import SuperCell, SuperCellChild
from .atom import Atom
from .geometry import Geometry

__all__ = ['Grid']

# Default nsc variable
_nsc = np.array([1]*3,np.int32)
_size = np.array([0]*3,np.int32)
_origo = np.array([0]*3,np.float64)
_bc = 1 # Periodic (Grid.Periodic)
_dtype = np.float64

class Grid(SuperCellChild):
    """ Object to retain grid information

    This grid object handles cell vectors and divisions of said grid.

    A grid can be periodic and non-periodic.
    """

    # Constant (should never be changed)
    Periodic = 1
    Neumann = 2
    Dirichlet = 3

    def __init__(self,size=_size,bc=_bc,sc=None,dtype=_dtype,geom=None):
        """ Initialize a `Grid` object.

        Initialize a `Grid` object.
        """

        self.set_supercell(sc)

        # Create the grid
        self.set_grid(size,dtype=dtype)

        # Create the grid boundary conditions
        self.set_bc(bc)

        # Create the atomic structure in the grid, if possible
        self.set_geom(geom)

    def __getitem__(self,key):
        """ Returns the grid contained """
        return self.grid[key]

    def __setitem__(self,key,val):
        """ Updates the grid contained """
        self.grid[key] = val

    def interp(self,size,*args,**kwargs):
        """ Returns an interpolated version of the grid

        Parameters
        ----------
        size : int, array_like
            the new size of the grid
        *args, **kwargs :
            optional arguments passed to the interpolation algorithm
            The interpolation routine is `scipy.interpolate.interpn`
        """
        # Get grid spacing
        dold = []
        dnew = []
        for i in range(3):
            dold.append(np.linspace(0,1,self.size[i]))
            dnew.append(np.linspace(0,1,size[i]))

        # Interpolate
        from scipy.interpolate import interpn

        # Create new grid
        grid = self.__class__(size, bc=np.copy(self.bc), sc=self.sc.copy())
        grid.grid = interpn(dold, self.grid, dnew, *args, **kwargs)

        return grid

    @property
    def size(self):
        """ Returns size of the grid """
        return np.asarray(self.grid.shape,np.int32)

    def set_grid(self,size,dtype=_dtype):
        """ Create the internal grid of certain size.
        """
        self.grid = np.empty(size,dtype=dtype)

    def set_bc(self,boundary=None,a=None,b=None,c=None):
        """ Set the boundary conditions on the grid
        boundary: [3], integer, optional
           boundary condition for all boundaries (or the same for all)
        a: integer, optional
           boundary condition for the first unit-cell vector direction
        b: integer, optional
           boundary condition for the second unit-cell vector direction
        c: integer, optional
           boundary condition for the third unit-cell vector direction
        """
        if not boundary is None:
            if isinstance(boundary,Integral):
                self.bc = np.array([boundary]*3,np.int32)
            else:
                self.bc = np.asarray(boundary,np.int32)
        if not a is None: self.bc[0] = a
        if not b is None: self.bc[1] = b
        if not c is None: self.bc[2] = c

    # Aliases
    set_boundary = set_bc
    set_boundary_condition = set_bc

    def set_geom(self,geom):
        """ Sets the `Geometry` for the grid.

        Setting the `Geometry` for the grid is a possibility
        to attach atoms to the grid.

        It is not a necessary entity.
        """
        if geom is None:
            # Fake geometry
            self.set_geom(Geometry([0,0,0],Atom['H'],sc=self.sc))
        else:
            self.geom = geom
            self.set_sc(geom.sc)

    def copy(self):
        """
        Returns a copy of the object.
        """
        grid = self.__class__(np.copy(self.size), bc=np.copy(self.bc),
                              dtype=self.grid.dtype,
                              geom=self.geom.copy())
        grid.grid[:,:,:] = self.grid[:,:,:]
        return grid

    def swapaxes(self,a,b):
        """ Returns Grid with swapped axis

        If ``swapaxes(0,1)`` it returns the 0 in the 1 values.
        """
        # Create index vector
        idx = np.arange(3)
        idx[b] = a
        idx[a] = b
        s = np.copy(self.size)
        grid = self.__class__(s[idx], bc=self.bc[idx],
                              sc=self.sc.swapaxes(a,b),dtype=self.grid.dtype)
        # We need to force the C-order or we loose the contiguoity
        grid.grid = np.copy(np.swapaxes(self.grid,a,b),order='C')
        return grid

    @property
    def dcell(self):
        """ Returns the delta-cell """
        # Calculate the grid-distribution
        g_size = self.size
        dcell = np.empty([3,3],np.float64)
        for ix in range(3):
            dcell[ix,:] = self.cell[ix,:] / g_size[ix]
        return dcell

    @property
    def dvol(self):
        """ Returns the delta-volume """
        return self.sc.vol / np.prod(self.size)

    def cross_section(self,idx,axis):
        """ Takes a cross-section of the grid along axis `axis`

        Remark: This API entry might change to handle arbitrary
        cuts via rotation of the axis """

        # First calculate the new size
        size = self.size
        cell = np.copy(self.cell)
        # Down-scale cell
        cell[axis,:] /= size[axis]
        size[axis] = 1
        grid = self.__class__(size, bc=np.copy(self.bc),geom=self.geom.copy())
        # Update cell size (the cell is smaller now)
        grid.set_sc(cell)

        if axis == 0:
            grid.grid[:,:,:] = self.grid[idx,:,:]
        elif axis == 1:
            grid.grid[:,:,:] = self.grid[:,idx,:]
        elif axis == 2:
            grid.grid[:,:,:] = self.grid[:,:,idx]
        else:
            raise ValueError('Unknown axis specification in cross_section')

        return grid

    def sum(self,axis):
        """ Returns the grid summed along axis `axis`. """
        # First calculate the new size
        size = self.size
        cell = np.copy(self.cell)
        # Down-scale cell
        cell[axis,:] /= size[axis]
        size[axis] = 1

        grid = self.__class__(size, bc=np.copy(self.bc),geom=self.geom.copy())
        # Update cell size (the cell is smaller now)
        grid.set_sc(cell)

        # Calculate sum
        grid.grid[:,:,:] = np.sum(self.grid,axis=axis)
        return grid

    def average(self,axis):
        """ Returns the average grid along direction `axis` """
        n = self.size[axis]
        return self.sum(axis) / float(n)

    def sub(self,idx,axis):
        """ Retains certain indices from a specified axis.

        Parameters
        ----------
        idx : array_like
           the indices of the grid axis `axis` to be retained
        axis : int
           the axis segment from which we retain the indices `idx`
        """
        uidx = np.unique(np.clip(idx,0,self.size[axis]-1))

        # Calculate new size
        size = self.size
        cell = np.copy(self.cell)
        old_N = size[axis]

        # Calculate new size
        size[axis] = len(idx)
        if size[axis] < 1:
            raise ValueError('You cannot retain no indices.')

        # Down-scale cell
        cell[axis,:] = cell[axis,:] / old_N * size[axis]

        grid = self.__class__(size, bc=np.copy(self.bc), geom=self.geom.copy())
        # Update cell size (the cell is smaller now)
        grid.set_sc(cell)

        # Remove the indices
        # First create the opposite, index
        if axis == 0:
            grid.grid[:,:,:] = self.grid[idx,:,:]
        elif axis == 1:
            grid.grid[:,:,:] = self.grid[:,idx,:]
        elif axis == 2:
            grid.grid[:,:,:] = self.grid[:,:,idx]

        return grid

    def remove(self,idx,axis):
        """ Removes certain indices from a specified axis.

        Parameters
        ----------
        idx : array_like
           the indices of the grid axis `axis` to be removed
        axis : int
           the axis segment from which we remove all indices `idx`
        """
        uidx = np.unique(np.clip(idx,0,self.size[axis]-1))
        ret_idx = np.setdiff1d(np.arange(self.size[axis]), uidx, assume_unique=True)
        return self.sub(ret_idx, axis)

    def index(self,coord,axis=None):
        """ Returns the index along the axis `axis` where `coord` exists

        Parameters
        ----------
        coord : array_like / float
           the coordinate of the axis
        axis : int
           the axis direction of the index
        """

        # if the axis is none, we do this for all axes
        if axis is None:
            # Loop over each direction
            idx = np.empty([3],np.int32)
            for i in [0,1,2]:
                # Normalize cell vector
                ca = self.cell[i,:] / np.sum(self.cell[i,:]**2)**.5
                # get the coordinate along the direction of the cell vector
                c = np.dot(self.cell[i,:],coord) * coord
                # Calculate the index corresponding to this placement
                idx[i] = self.index(c,i)
            return idx

        # Ensure a 1D array
        ac = np.atleast_1d(np.asarray(coord,np.float64))

        # Calculate the index of the coord in the cell
        dax = self.dcell[axis,:]

        # Calculate how many indices are required to fulfil
        # the correct line cut
        return int(np.rint( (np.sum(ac ** 2) / np.sum(dax ** 2)) **.5 ))

    def append(self,other,axis):
        """ Appends other `Grid` to this grid along axis

        """
        size = np.copy(self.size)
        size[axis] += other.size[axis]
        return self.__class__(size, bc = np.copy(self.bc),
                              sc=self.sc.append(other.sc,axis))

    def write(self,sile):
        """ Writes grid to the ``Sile`` using ``sile.write_grid``

        Parameters
        ----------
        sile : Sile, str
            a ``Sile`` object which will be used to write the grid
            if it is a string it will create a new sile using ``get_sile``
        """

        # This only works because, they *must*
        # have been imported previously
        from sids.io import get_sile, BaseSile
        if isinstance(sile,BaseSile):
            sile.write_grid(self)
        else:
            get_sile(sile,'w').write_grid(self)

    def __repr__(self):
        """ Representation of object """
        return 'Grid[{} {} {}]'.format(*self.size)

    def _check_compatibility(self,other,msg):
        """ Internal check for asserting two grids are commensurable """
        if self == other:
            return True
        else:
            raise ValueError('Grids are not compatible, '+msg)

    def _compatible_copy(self,other,*args,**kwargs):
        """ Returns a copy of self with an additional check of commensurable """
        if isinstance(other,Grid):
            if self._check_compatibility(other,*args,**kwargs):
                return self.copy()
        else:
            return self.copy()

    def __eq__(self,other):
        """ Returns true if the two grids are commensurable

        There will be no check of the values _on_ the grid. """
        return bool(np.all(self.size == other.size))

    def __ne__(self,other):
        """ Returns whether two grids have the same size """
        return not (self == other)

    def __add__(self,other):
        """ Returns a new grid with the addition of two grids

        Returns same size with same cell as the first"""
        if self == other:
            # We can return a new grid
            grid = self.copy()
            grid.grid = self.grid + other.grid
            return grid
        raise ValueError('Grids are not compatible, they cannot be added.')

    def __sub__(self,other):
        """ Returns a new grid with the difference of two grids

        Returns same size with same cell as the first"""
        grid = self._compatible_copy(other,'they cannot be subtracted')
        if isinstance(other,Grid):
            grid.grid = self.grid - other.grid
        else:
            grid.grid = self.grid - other
        return grid

    def __truediv__(self,other):
        grid = self._compatible_copy(other,'they cannot be divided')
        if isinstance(other,Grid):
            grid.grid = self.grid / other.grid
        else:
            grid.grid = self.grid / other
        return grid

    def __mul__(self,other):
        grid = self._compatible_copy(other,'they cannot be multiplied')
        if isinstance(other,Grid):
            grid.grid = self.grid * other.grid
        else:
            grid.grid = self.grid * other
        return grid

if __name__ == "__main__":
    pass
<|endoftext|>