<commit_before>#!/usr/bin/env python
"""Python helper script for virt-install(1)"""

import subprocess
import os
import platform
import argparse
import yaml
import re
import string
import random
import xml.etree.ElementTree as ET

__author__ = 'Jason Callaway'
__email__ = 'jason@jasoncallaway.com'
__license__ = 'Apache License Version 2.0'
__version__ = '0.1'
__status__ = 'alpha'

# TODO: add support for other platforms
SUPPORTED_PLATFORMS = ['rhel', 'centos', 'fedora']
CONFIG_PATH = os.path.expanduser('~') + '/.config/kvm-install/config.yaml'

class KVMInstall(object):
    def setup_tmp(self, random8):
        self.tmpdir = '/tmp/kvm-install-' + random8
        try:
            os.makedirs(self.tmpdir)
        except Exception, e:
            raise e
        self.stdout = self.tmpdir + '/stdout.txt'
        self.stderr = self.tmpdir + '/stderr.txt'
        self.virsh_netdumpxml = self.tmpdir + '/netdump.xml'

    def get_random(self, domain, length):
        return ''.join(random.SystemRandom().choice(domain) for _ in range(length))

    def parse_config(self):
        """Parse home dir .config file"""

        if self.args.configfile is None:
            config_path = CONFIG_PATH
        else:
            config_path = self.args.configfile
        if self.args.verbose is True:
            print '  using config file: ' + config_path
        try:
            # If the config file doesn't exist, let's create and populate it
            if not os.path.isfile(config_path):
                os.makedirs(os.path.split(config_path)[0])
                with open(config_path, 'w') as config_file:
                    config_file.write('---\n' +
                                      'vcpus: 1\n' +
                                      'ram: 1024\n' +
                                      'disk: 10\n' +
                                      'domain: example.com\n' +
                                      'network: default\n' +
                                      'mac: 5c:e0:c5:c4:26\n' +
                                      'type: linux\n' +
                                      'variant: rhel7\n')
        except Exception, e:
            raise Exception('unable to create config file at ' + config_path + ': ' + str(e))

        try:
            # Now read and parse it
            config_string = open(config_path).read()
            self.config = yaml.load(config_string)
        except Exception, e:
            raise Exception('unable to read config file at ' + config_path + ': ' + str(e))

        # Now iterate over the arguments to build the config.
        # Remember, self.args is a Namespace.
        for k in self.args.__dict__:
            if self.args.__dict__[k] is not None:
                self.config[k] = self.args.__dict__[k]

    def run_command(self, command, stdout, stderr):
        if self.args.verbose is True:
            print '  running command: ' + ' '.join(command)
            print '  stdout: ' + stdout
            print '  stderr: ' + stderr
        out = open(stdout, 'a')
        err = open(stderr, 'a')
        exit_signal = subprocess.call(command, stdout=out, stderr=err)
        if exit_signal != 0:
            raise Exception('command failed with exit signal ' + str(exit_signal) + ': ' + ' '.join(command))
        out.close()
        err.close()

    def setup_lvm(self):
        """Setup the VMs root volume with LVM"""

        # Grab the config values we need
        from_lvm = self.config['clone']
        size = self.config['disk']
        name = self.config['name']

        command = ['lvcreate', '-s', from_lvm, '-L', size + 'G', '-n', name]
        try:
            self.run_command(command, self.stdout, self.stderr)
        except Exception, e:
            raise e

    def setup_image(self):
        """Setup the VMs root volume with an image file"""

        # Grab the config values we need
        from_image = self.config['image']
        path = os.path.split(from_image)[0]
        extension = os.path.splitext(from_image)[1]
        size = self.config['disk']
        name = self.config['name']

        command = ['cp', from_image, path + '/' + name + extension]
        try:
            self.run_command(command, self.stdout, self.stderr)
        except Exception, e:
            raise e

    def net_dumpxml(self, network):
        network = self.config['network']
        command = ['virsh', 'net-dumpxml', network]
        try:
            self.run_command(command, self.virsh_netdumpxml, self.stderr)
        except Exception, e:
            raise e

    def get_etree_elements(self, xmlfile, element):
        tree = ET.parse(xmlfile)
        l = []
        for elem in tree.getiterator():
            if elem.get(element) is not None:
                l.append(elem.get(element))
        return l

    def get_mac_addresses(self):
        return self.get_etree_elements(self.virsh_netdumpxml, 'mac')

    def generate_mac(self, prefix):
        generated_mac = ''
        # Determine how long our prefix is
        num_colons = prefix.count(':')
        # Add that number of hex substrings
        for _ in range(5 - num_colons):
            # This is a little big funky. I wanted to be sure we have only a-f,0-9, but the string.hexdigits
            # string includes a-f,A-F, so we have to convert to lower case and strip out duplicates
            hex_domain = ''.join(set(string.hexdigits.lower()))
            print '+++ hex_domain: ' + hex_domain
            new_hex = self.get_random(hex_domain, 2)
            print '+++ new_hex: ' + new_hex
            generated_mac = generated_mac.join(':' + new_hex)
        print '+++ generated_mac: ' + generated_mac
        return self.config['mac'] + generated_mac

    def get_ip_addresses(self):
        return self.get_etree_elements(self.virsh_netdumpxml, 'ip')

    def get_ip_range(self, xmlfile):
        tree = ET.parse(xmlfile)
        root = tree.getroot()
        start = root.find('ip').find('dhcp').find('range').get('start')
        end = root.find('ip').find('dhcp').find('range').get('end')
        return [start, end]

    def generate_ip(self):
        ip_start, ip_end = self.get_ip_range(self.virsh_netdumpxml)
        start = re.sub('^\d{1,3}\.\d{1,3}\.\d{1,3}\.', '', ip_start)
        end = re.sub('^\d{1,3}\.\d{1,3}\.\d{1,3}\.', '', ip_end)
        first_three_octets = re.sub('\.\d{1,3}$', '', ip_start)
        return first_three_octets + '.' + str(random.randint(int(start), int(end)))

    def setup_network(self):
        """Setup the virsh network settings for the VM"""

        self.net_dumpxml(self.config['network'])

        # TODO: Add IPv6 support

        try:
            # First, find a new mac address
            mac_addresses = self.get_mac_addresses()
            print '+++ mac_addresses: ' + ' '.join(mac_addresses)
            new_mac = ''
            good_mac = False
            while good_mac is False:
                new_mac = self.generate_mac(self.config['mac'])
                if new_mac not in mac_addresses:
                    good_mac = True
                    if self.args.verbose is True:
                        print '  new mac found: ' + new_mac
        except Exception, e:
            raise Exception('setup_network failed to generate a new mac address: ' + str(e))

        try:
            # Then find an IP address in range that doesn't already exist
            ip_addresses = self.get_ip_addresses()
            new_ip = ''
            good_ip = False
            while good_ip is False:
                new_ip = self.generate_ip()
                if new_ip not in ip_addresses:
                    good_ip = True
                    if self.args.verbose is True:
                        print '  new ip found: ' + new_ip
        except Exception, e:
            raise Exception('setup_network failed to generate a new ip address: ' + str(e))

        # Record the new IP for other functions' use
        self.config['new_ip'] = new_ip

        # Now generate the virst net-update command
        command = ['virsh',
                   'net-update',
                   self.config['network'],
                   'add-last',
                   'ip-dhcp-host']
        host_xml = '\\"<host mac=\'' + new_mac + '\' name=\'' + self.config['name'] + '.' + self.config['domain'] + \
                   '\' ip=\'' + new_ip + '\'/>\\"'
        command.append(host_xml)

        print '+++ command: ' + ' '.join(command)

        # First, update the persistent config
        try:
            config_command = command
            config_command.append('--config')
            self.run_command(config_command, self.stdout, self.stderr)
        except Exception, e:
            raise Exception('virsh net-update --config failed: ' + str(e))

        # Now, update the current config
        try:
            current_command = command
            current_command.append('--current')
            self.run_command(current_command, self.stdout, self.stderr)
        except Exception, e:
            raise Exception('virsh net-update --current failed: ' + str(e))

    def update_etchosts(self):
        try:
            etchosts = open('/etc/hosts', 'r+')
            hosts = etchosts.read()
            hosts = hosts + self.config['new_ip'] + '\t' + self.config['name'] + '.' + self.config['domain'] + ' ' + \
                    self.config['name'] + '\n'
            etchosts.seek(0)
            etchosts.truncate()
            etchosts.write(hosts)
            etchosts.close()
        except Exception, e:
            raise e

    def restart_dnsmasq(self):
        command = ['systemctl', 'restart', 'dnsmasq.service']
        try:
            self.run_command(command, self.stdout, self.stderr)
        except Exception, e:
            raise e

    def do_virtinstall(self):
        command = ['virt-install',
                   '--noautoconsole',
                   '--hvm',
                   '--vnc',
                   '--name', self.config['name'],
                   '--vcpus', self.config['vcpus'],
                   '--ram', self.config['ram'],
                   '--network', self.config['network'],
                   '--os-type', self.config['type'],
                   '--os-variant', self.config['variant'],
                   '--boot', 'hd']
        if self.config.has_key('clone'):
            devpath = os.path.split(self.config['clone'])[0]
            command.append['--disk', 'path=' + devpath + '/' + self.config['name']]
        else:
            imgpath = os.path.split(self.config['image'])[0]
            command.append['--disk', 'path=' + imgpath + '/' + self.config['name'] + '.img' +
                           ',size=' + self.config['size'] + ',format=qcow2']
        try:
            self.run_command(command, self.stdout, self.stderr)
        except Exception, e:
            raise e

    def __init__(self, parsed_args):
        # TODO: put in environemnt checks, i.e., does virt-install exist, etc.

        # Grab the parsed command line arguments
        self.args = parsed_args

        # This make my IDE happy
        self.config = {}

        # Set up our random string and temp directory
        random8 = self.get_random(string.ascii_letters + string.digits, 8)
        self.setup_tmp(random8)

        # Check to see if we're on a supported platform.
        if platform.dist()[0] not in SUPPORTED_PLATFORMS:
            raise Exception('unsupported platform: ' + platform.dist()[0])

        # Parse the config file and build our config object
        if self.args.verbose is True:
            print ' parsing config file'
        if self.args.configfile is None:
            self.args.configfile = CONFIG_PATH
        self.parse_config()

        # If we have both a clone and image config directive, prefer LVM
        if self.config.has_key('clone'):
            if self.args.verbose is True:
                print ' setting up lvm'
            self.setup_lvm()
        else:
            if self.args.verbose is True:
                print ' setting up image'
            if self.config.has_key('image'):
                self.setup_image()
            else:
                raise Exception('you must specify either an LVM or file base image with -c or -i')

        try:
            # Now set up the new network
            if self.args.verbose is True:
                print ' setting up network'
            self.setup_network()
        except Exception, e:
            raise Exception('setup network failed: ' + str(e))

        try:
            # Update /etc/hosts
            if self.args.verbose is True:
                print ' updating /etc/hosts'
            self.update_etchosts()
        except Exception, e:
            raise Exception('update /etc/hosts failed: ' + str(e))

        try:
            # Restart the dnsmasq service
            if self.args.verbose is True:
                print ' restarting dnsmasq'
            self.restart_dnsmasq()
        except Exception, e:
            raise Exception('restart dnsmasq failed: ' + str(e))

        try:
            # Finally, we can install the VM
            if self.args.verbose is True:
                print ' doing virt-install'
            self.do_virtinstall()
        except Exception, e:
            raise Exception('virt-install failed: ' + str(e))

if __name__ == "__main__":
    # Note that we want all of the arguments to be parsed as Strings.
    # This makes building the virsh and virt-install commands easier.
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--clone',
                        help='name of the source logical volume to be cloned')
    parser.add_argument('-i', '--image',
                        help='image file to duplicate')
    parser.add_argument('-v', '--vcpus',
                        help='number of virtual CPUs')
    parser.add_argument('-r', '--ram',
                        help='amount of RAM in MB')
    parser.add_argument('-d', '--disk',
                        help='disk size in GB')
    parser.add_argument('-D', '--domain',
                        help='domainname for dhcp / dnsmasq')
    parser.add_argument('-N', '--network',
                        help='libvirt network')
    parser.add_argument('--type',
                        help='os type, i.e., linux')
    parser.add_argument('--variant',
                        help='os variant, i.e., rhel7')
    parser.add_argument('-f', '--configfile',
                        help='specify an alternate config file, default=~/.config/kvm-install/config.yaml')
    parser.add_argument('--verbose', dest='verbose', action='store_true',
                        help='verbose output')
    parser.add_argument('name',
                        help='name of the new virtual machine')
    parser.set_defaults(verbose=False)
    args = parser.parse_args()
    KVMInstall(args)<commit_msg>almost works, just a few problems with the virt-install logic<commit_after>#!/usr/bin/env python
"""Python helper script for virt-install(1)"""

import subprocess
import os
import platform
import argparse
import yaml
import re
import string
import random
import xml.etree.ElementTree as ET

__author__ = 'Jason Callaway'
__email__ = 'jason@jasoncallaway.com'
__license__ = 'Apache License Version 2.0'
__version__ = '0.1'
__status__ = 'alpha'

# TODO: add support for other platforms
SUPPORTED_PLATFORMS = ['rhel', 'centos', 'fedora']
CONFIG_PATH = os.path.expanduser('~') + '/.config/kvm-install/config.yaml'

class KVMInstall(object):
    def setup_tmp(self, random8):
        self.tmpdir = '/tmp/kvm-install-' + random8
        try:
            os.makedirs(self.tmpdir)
        except Exception, e:
            raise e
        self.stdout = self.tmpdir + '/stdout.txt'
        self.stderr = self.tmpdir + '/stderr.txt'
        self.virsh_netdumpxml = self.tmpdir + '/netdump.xml'

    def get_random(self, domain, length):
        return ''.join(random.SystemRandom().choice(domain) for _ in range(length))

    def parse_config(self):
        """Parse home dir .config file"""

        if self.args.configfile is None:
            config_path = CONFIG_PATH
        else:
            config_path = self.args.configfile
        if self.args.verbose is True:
            print '  using config file: ' + config_path
        try:
            # If the config file doesn't exist, let's create and populate it
            if not os.path.isfile(config_path):
                os.makedirs(os.path.split(config_path)[0])
                with open(config_path, 'w') as config_file:
                    config_file.write('---\n' +
                                      'vcpus: 1\n' +
                                      'ram: 1024\n' +
                                      'disk: 10\n' +
                                      'domain: example.com\n' +
                                      'network: default\n' +
                                      'mac: 5c:e0:c5:c4:26\n' +
                                      'type: linux\n' +
                                      'variant: rhel7\n')
        except Exception, e:
            raise Exception('unable to create config file at ' + config_path + ': ' + str(e))

        try:
            # Now read and parse it
            config_string = open(config_path).read()
            self.config = yaml.load(config_string)
        except Exception, e:
            raise Exception('unable to read config file at ' + config_path + ': ' + str(e))

        # Now iterate over the arguments to build the config.
        # Remember, self.args is a Namespace.
        for k in self.args.__dict__:
            if self.args.__dict__[k] is not None:
                self.config[k] = self.args.__dict__[k]

    def run_command(self, command, stdout, stderr):
        if self.args.verbose is True:
            print '  running command: ' + ' '.join(command)
            print '  stdout: ' + stdout
            print '  stderr: ' + stderr
        out = open(stdout, 'a')
        err = open(stderr, 'a')
        exit_signal = subprocess.call(' '.join(command), stdout=out, stderr=err, shell=True)
        if exit_signal != 0:
            raise Exception('command failed with exit signal ' + str(exit_signal) + ': ' + ' '.join(command))
        out.close()
        err.close()

    def setup_lvm(self):
        """Setup the VMs root volume with LVM"""

        # Grab the config values we need
        from_lvm = self.config['clone']
        size = self.config['disk']
        name = self.config['name']

        command = ['lvcreate', '-s', from_lvm, '-L', size + 'G', '-n', name]
        try:
            self.run_command(command, self.stdout, self.stderr)
        except Exception, e:
            raise e

    def setup_image(self):
        """Setup the VMs root volume with an image file"""

        # Grab the config values we need
        from_image = self.config['image']
        path = os.path.split(from_image)[0]
        extension = os.path.splitext(from_image)[1]
        size = self.config['disk']
        name = self.config['name']

        command = ['cp', from_image, path + '/' + name + extension]
        try:
            self.run_command(command, self.stdout, self.stderr)
        except Exception, e:
            raise e

    def net_dumpxml(self, network):
        network = self.config['network']
        command = ['virsh', 'net-dumpxml', network]
        try:
            self.run_command(command, self.virsh_netdumpxml, self.stderr)
        except Exception, e:
            raise e

    def get_etree_elements(self, xmlfile, element):
        tree = ET.parse(xmlfile)
        l = []
        for elem in tree.getiterator():
            if elem.get(element) is not None:
                l.append(elem.get(element))
        return l

    def get_mac_addresses(self):
        return self.get_etree_elements(self.virsh_netdumpxml, 'mac')

    def generate_mac(self, prefix):
        generated_mac = ''
        # Determine how long our prefix is
        num_colons = prefix.count(':')
        # Add that number of hex substrings
        for _ in range(5 - num_colons):
            # This is a little big funky. I wanted to be sure we have only a-f,0-9, but the string.hexdigits
            # string includes a-f,A-F, so we have to convert to lower case and strip out duplicates
            hex_domain = ''.join(set(string.hexdigits.lower()))
            print '+++ hex_domain: ' + hex_domain
            new_hex = self.get_random(hex_domain, 2)
            print '+++ new_hex: ' + new_hex
            generated_mac = generated_mac.join(':' + new_hex)
        print '+++ generated_mac: ' + generated_mac
        return self.config['mac'] + generated_mac

    def get_ip_addresses(self):
        return self.get_etree_elements(self.virsh_netdumpxml, 'ip')

    def get_ip_range(self, xmlfile):
        tree = ET.parse(xmlfile)
        root = tree.getroot()
        start = root.find('ip').find('dhcp').find('range').get('start')
        end = root.find('ip').find('dhcp').find('range').get('end')
        return [start, end]

    def generate_ip(self):
        ip_start, ip_end = self.get_ip_range(self.virsh_netdumpxml)
        start = re.sub('^\d{1,3}\.\d{1,3}\.\d{1,3}\.', '', ip_start)
        end = re.sub('^\d{1,3}\.\d{1,3}\.\d{1,3}\.', '', ip_end)
        first_three_octets = re.sub('\.\d{1,3}$', '', ip_start)
        return first_three_octets + '.' + str(random.randint(int(start), int(end)))

    def setup_network(self):
        """Setup the virsh network settings for the VM"""

        self.net_dumpxml(self.config['network'])

        # TODO: Add IPv6 support

        try:
            # First, find a new mac address
            mac_addresses = self.get_mac_addresses()
            print '+++ mac_addresses: ' + ' '.join(mac_addresses)
            new_mac = ''
            good_mac = False
            while good_mac is False:
                new_mac = self.generate_mac(self.config['mac'])
                if new_mac not in mac_addresses:
                    good_mac = True
                    if self.args.verbose is True:
                        print '  new mac found: ' + new_mac
        except Exception, e:
            raise Exception('setup_network failed to generate a new mac address: ' + str(e))

        try:
            # Then find an IP address in range that doesn't already exist
            ip_addresses = self.get_ip_addresses()
            new_ip = ''
            good_ip = False
            while good_ip is False:
                new_ip = self.generate_ip()
                if new_ip not in ip_addresses:
                    good_ip = True
                    if self.args.verbose is True:
                        print '  new ip found: ' + new_ip
        except Exception, e:
            raise Exception('setup_network failed to generate a new ip address: ' + str(e))

        # Record the new IP for other functions' use
        self.config['new_ip'] = new_ip

        # Now generate the virst net-update command
        command = ['virsh',
                   'net-update',
                   self.config['network'],
                   'add-last',
                   'ip-dhcp-host']
        host_xml = '"<host mac=\'' + new_mac + '\' name=\'' + self.config['name'] + '.' + self.config['domain'] + \
                   '\' ip=\'' + new_ip + '\'/>"'
        command.append(host_xml)

        print '+++ command: ' + ' '.join(command)

        config_command = list(command)
        current_command = list(command)

        # Now, update the current config
        try:
            current_command.append('--current')
            self.run_command(current_command, self.stdout, self.stderr)
        except Exception, e:
            raise Exception('virsh net-update --current failed: ' + str(e))

        # First, update the persistent config
        try:
            config_command.append('--config')
            self.run_command(config_command, self.stdout, self.stderr)
        except Exception, e:
            raise Exception('virsh net-update --config failed: ' + str(e))

        # Now do the same for DNS
        command = list()
        command = ['virsh',
                   'net-update',
                   self.config['network'],
                   'add-last',
                   'dns-host']
        host_xml = '"<host ip=\'' + new_ip + '\'><hostname>' + self.config['name'] + '.' + self.config['domain'] + \
                   '</hostname></host>"'
        command.append(host_xml)

        print '+++ command: ' + ' '.join(command)

        config_command = list(command)
        current_command = list(command)

        # Now, update the current config
        try:
            current_command.append('--current')
            self.run_command(current_command, self.stdout, self.stderr)
        except Exception, e:
            raise Exception('virsh net-update --current failed: ' + str(e))

        # First, update the persistent config
        try:
            config_command.append('--config')
            self.run_command(config_command, self.stdout, self.stderr)
        except Exception, e:
            raise Exception('virsh net-update --config failed: ' + str(e))

    def update_etchosts(self):
        try:
            etchosts = open('/etc/hosts', 'r+')
            hosts = etchosts.read()
            hosts = hosts + self.config['new_ip'] + '\t' + self.config['name'] + '.' + self.config['domain'] + ' ' + \
                    self.config['name'] + '\n'
            etchosts.seek(0)
            etchosts.truncate()
            etchosts.write(hosts)
            etchosts.close()
        except Exception, e:
            raise e

    def restart_dnsmasq(self):
        command = ['systemctl', 'restart', 'dnsmasq.service']
        try:
            self.run_command(command, self.stdout, self.stderr)
        except Exception, e:
            raise e

    def do_virtinstall(self):
        command = ['virt-install',
                   '--noautoconsole',
                   '--hvm',
                   '--vnc',
                   '--name', self.config['name'],
                   '--vcpus', self.config['vcpus'],
                   '--ram', self.config['ram'],
                   '--network', self.config['network'],
                   '--os-type', self.config['type'],
                   '--os-variant', self.config['variant'],
                   '--boot', 'hd']
        if self.config.has_key('clone'):
            devpath = os.path.split(self.config['clone'])[0]
            command.append['--disk', 'path=' + devpath + '/' + self.config['name']]
        else:
            imgpath = os.path.split(self.config['image'])[0]
            command.append['--disk', 'path=' + imgpath + '/' + self.config['name'] + '.img' +
                           ',size=' + str(self.config['disk']) + ',format=qcow2']
        try:
            self.run_command(command, self.stdout, self.stderr)
        except Exception, e:
            raise e

    def __init__(self, parsed_args):
        # TODO: put in environemnt checks, i.e., does virt-install exist, etc.

        # Grab the parsed command line arguments
        self.args = parsed_args

        # This make my IDE happy
        self.config = {}

        # Set up our random string and temp directory
        random8 = self.get_random(string.ascii_letters + string.digits, 8)
        self.setup_tmp(random8)

        # Check to see if we're on a supported platform.
        if platform.dist()[0] not in SUPPORTED_PLATFORMS:
            raise Exception('unsupported platform: ' + platform.dist()[0])

        # Parse the config file and build our config object
        if self.args.verbose is True:
            print ' parsing config file'
        if self.args.configfile is None:
            self.args.configfile = CONFIG_PATH
        self.parse_config()

        # If we have both a clone and image config directive, prefer LVM
        if self.config.has_key('clone'):
            if self.args.verbose is True:
                print ' setting up lvm'
            self.setup_lvm()
        else:
            if self.args.verbose is True:
                print ' setting up image'
            if self.config.has_key('image'):
                self.setup_image()
            else:
                raise Exception('you must specify either an LVM or file base image with -c or -i')

        try:
            # Now set up the new network
            if self.args.verbose is True:
                print ' setting up network'
            self.setup_network()
        except Exception, e:
            raise Exception('setup network failed: ' + str(e))

        try:
            # Update /etc/hosts
            if self.args.verbose is True:
                print ' updating /etc/hosts'
            self.update_etchosts()
        except Exception, e:
            raise Exception('update /etc/hosts failed: ' + str(e))

        try:
            # Restart the dnsmasq service
            if self.args.verbose is True:
                print ' restarting dnsmasq'
            self.restart_dnsmasq()
        except Exception, e:
            raise Exception('restart dnsmasq failed: ' + str(e))

        try:
            # Finally, we can install the VM
            if self.args.verbose is True:
                print ' doing virt-install'
            self.do_virtinstall()
        except Exception, e:
            raise Exception('virt-install failed: ' + str(e))

if __name__ == "__main__":
    # Note that we want all of the arguments to be parsed as Strings.
    # This makes building the virsh and virt-install commands easier.
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--clone',
                        help='name of the source logical volume to be cloned')
    parser.add_argument('-i', '--image',
                        help='image file to duplicate')
    parser.add_argument('-v', '--vcpus',
                        help='number of virtual CPUs')
    parser.add_argument('-r', '--ram',
                        help='amount of RAM in MB')
    parser.add_argument('-d', '--disk',
                        help='disk size in GB')
    parser.add_argument('-D', '--domain',
                        help='domainname for dhcp / dnsmasq')
    parser.add_argument('-N', '--network',
                        help='libvirt network')
    parser.add_argument('--type',
                        help='os type, i.e., linux')
    parser.add_argument('--variant',
                        help='os variant, i.e., rhel7')
    parser.add_argument('-f', '--configfile',
                        help='specify an alternate config file, default=~/.config/kvm-install/config.yaml')
    parser.add_argument('--verbose', dest='verbose', action='store_true',
                        help='verbose output')
    parser.add_argument('name',
                        help='name of the new virtual machine')
    parser.set_defaults(verbose=False)
    args = parser.parse_args()
    KVMInstall(args)<|endoftext|>


<commit_before># -*- encoding: utf-8 -*-
###############################################################################
#                                                                             #
# Copyright (C) 2013  Raphaël Valyi - Akretion                                #
#                                                                             #
#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/>.        #
###############################################################################

from openerp.osv import orm
from openerp.tools.translate import _

class SaleOrder(orm.Model):
    _inherit = 'sale.order'

    def _make_invoice(self, cr, uid, order, lines, context=None):
        obj_invoice_line = self.pool.get('account.invoice.line')
        lines_service = []
        lines_product = []
        inv_product_ids = []
        inv_service_ids = []
        inv_id_product = False
        inv_id_service = False

        if not context:
            context = {}

        obj_company = self.pool.get('res.company').browse(
            cr, uid, order.shop_id.company_id.id)

        if not obj_company.product_invoice_id:
            raise orm.except_orm(
                _('No fiscal document serie found !'),
                _("No fiscal document serie found for selected company %s") % (
                    order.company_id.name))

        for inv_line in obj_invoice_line.browse(cr, uid, lines, context=context):
            if inv_line.product_id.fiscal_type == 'service':
                lines_service.append(inv_line.id)

            if inv_line.product_id.fiscal_type == 'product':
                lines_product.append(inv_line.id)

        if lines_product:
            context['fiscal_type'] = 'product'
            inv_id_product = super(SaleOrder, self)._make_invoice(
                cr, uid, order, lines_product, context=context)
            inv_product_ids.append(inv_id_product)

        if lines_service:
            context['fiscal_type'] = 'service'
            inv_id_service = super(SaleOrder, self)._make_invoice(
                cr, uid, order, lines_service, context=context)
            inv_service_ids.append(inv_id_service)

        return inv_id_product or inv_id_service

    def _fiscal_comment(self, cr, uid, order, context=None):

        fp_comment = []
        fc_comment = []
        fp_ids = []
        fc_ids = []

        for line in order.order_line:
            if line.fiscal_position and \
            line.fiscal_position.inv_copy_note and \
            line.fiscal_position.note:
                if not line.fiscal_position.id in fp_ids:
                    fp_comment.append(line.fiscal_position.note)
                    fp_ids.append(line.fiscal_position.id)

            if line.product_id.ncm_id:
                fc = line.product_id.ncm_id
                if fc.inv_copy_note and fc.note:
                    if not fc.id in fc_ids:
                        fc_comment.append(fc.note)
                        fc_ids.append(fc.id)

        return fp_comment + fc_comment

    def _prepare_order_picking(self, cr, uid, order, context=None):
        result = super(SaleOrder, self)._prepare_order_picking(cr, uid,
            order, context)
        result['fiscal_category_id'] = order.fiscal_category_id and \
        order.fiscal_category_id.id
        result['fiscal_position'] = order.fiscal_position and \
        order.fiscal_position.id
        return result

    def _prepare_invoice(self, cr, uid, order, lines, context=None):
        """Prepare the dict of values to create the new invoice for a
           sale order. This method may be overridden to implement custom
           invoice generation (making sure to call super() to establish
           a clean extension chain).

           :param browse_record order: sale.order record to invoice
           :param list(int) line: list of invoice line IDs that must be
                                  attached to the invoice
           :return: dict of value to create() the invoice
        """
        result = super(SaleOrder, self)._prepare_invoice(
            cr, uid, order, lines, context)
        #TODO - Testar se só sobrescrevendo o metodo _fiscal_comment nao precisa fazer isso

        comment = []
        fiscal_comment = self._fiscal_comment(cr, uid, order, context=context)
        result['comment'] = " - ".join(comment + fiscal_comment)
        return result

class SaleOrderLine(orm.Model):
    _inherit = 'sale.order.line'

    def _prepare_order_line_invoice_line(self, cr, uid, line,
                                         account_id=False, context=None):
        result = super(SaleOrderLine, self)._prepare_order_line_invoice_line(
            cr, uid, line, account_id, context)

        if line.product_id.fiscal_type == 'product':
            cfop = self.pool.get("account.fiscal.position").read(
                cr, uid, [result['fiscal_position']], ['cfop_id'],
                context=context)
            if cfop[0]['cfop_id']:
                result['cfop_id'] = cfop[0]['cfop_id'][0]
        return result<commit_msg>minor optimization<commit_after># -*- encoding: utf-8 -*-
###############################################################################
#                                                                             #
# Copyright (C) 2013  Raphaël Valyi - Akretion                                #
#                                                                             #
#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/>.        #
###############################################################################

from openerp.osv import orm
from openerp.tools.translate import _

class SaleOrder(orm.Model):
    _inherit = 'sale.order'

    def _make_invoice(self, cr, uid, order, lines, context=None):
        obj_invoice_line = self.pool.get('account.invoice.line')
        lines_service = []
        lines_product = []
        inv_product_ids = []
        inv_service_ids = []
        inv_id_product = False
        inv_id_service = False

        if not context:
            context = {}

        obj_company = self.pool.get('res.company').browse(
            cr, uid, order.shop_id.company_id.id)

        if not obj_company.product_invoice_id:
            raise orm.except_orm(
                _('No fiscal document serie found !'),
                _("No fiscal document serie found for selected company %s") % (
                    order.company_id.name))

        for inv_line in obj_invoice_line.browse(cr, uid, lines, context=context):
            if inv_line.product_id.fiscal_type == 'service':
                lines_service.append(inv_line.id)
            elif inv_line.product_id.fiscal_type == 'product':
                lines_product.append(inv_line.id)

        if lines_product:
            context['fiscal_type'] = 'product'
            inv_id_product = super(SaleOrder, self)._make_invoice(
                cr, uid, order, lines_product, context=context)
            inv_product_ids.append(inv_id_product)

        if lines_service:
            context['fiscal_type'] = 'service'
            inv_id_service = super(SaleOrder, self)._make_invoice(
                cr, uid, order, lines_service, context=context)
            inv_service_ids.append(inv_id_service)

        return inv_id_product or inv_id_service

    def _fiscal_comment(self, cr, uid, order, context=None):
        fp_comment = []
        fc_comment = []
        fp_ids = []
        fc_ids = []

        for line in order.order_line:
            if line.fiscal_position and \
            line.fiscal_position.inv_copy_note and \
            line.fiscal_position.note:
                if not line.fiscal_position.id in fp_ids:
                    fp_comment.append(line.fiscal_position.note)
                    fp_ids.append(line.fiscal_position.id)

            if line.product_id.ncm_id:
                fc = line.product_id.ncm_id
                if fc.inv_copy_note and fc.note:
                    if not fc.id in fc_ids:
                        fc_comment.append(fc.note)
                        fc_ids.append(fc.id)

        return fp_comment + fc_comment

    def _prepare_order_picking(self, cr, uid, order, context=None):
        result = super(SaleOrder, self)._prepare_order_picking(cr, uid,
            order, context)
        result['fiscal_category_id'] = order.fiscal_category_id and \
        order.fiscal_category_id.id
        result['fiscal_position'] = order.fiscal_position and \
        order.fiscal_position.id
        return result

    def _prepare_invoice(self, cr, uid, order, lines, context=None):
        """Prepare the dict of values to create the new invoice for a
           sale order. This method may be overridden to implement custom
           invoice generation (making sure to call super() to establish
           a clean extension chain).

           :param browse_record order: sale.order record to invoice
           :param list(int) line: list of invoice line IDs that must be
                                  attached to the invoice
           :return: dict of value to create() the invoice
        """
        result = super(SaleOrder, self)._prepare_invoice(
            cr, uid, order, lines, context)
        #TODO - Testar se só sobrescrevendo o metodo _fiscal_comment nao precisa fazer isso

        comment = []
        fiscal_comment = self._fiscal_comment(cr, uid, order, context=context)
        result['comment'] = " - ".join(comment + fiscal_comment)
        return result

class SaleOrderLine(orm.Model):
    _inherit = 'sale.order.line'

    def _prepare_order_line_invoice_line(self, cr, uid, line,
                                         account_id=False, context=None):
        result = super(SaleOrderLine, self)._prepare_order_line_invoice_line(
            cr, uid, line, account_id, context)

        if line.product_id.fiscal_type == 'product':
            cfop = self.pool.get("account.fiscal.position").read(
                cr, uid, [result['fiscal_position']], ['cfop_id'],
                context=context)
            if cfop[0]['cfop_id']:
                result['cfop_id'] = cfop[0]['cfop_id'][0]
        return result
<|endoftext|>


<commit_before>#! /usr/bin/env python
"""Calculate vector divergence and related quantities at nodes or cells."""
import numpy as np
from landlab.utils.decorators import use_field_name_or_array

@use_field_name_or_array('face')
def calc_net_face_flux_at_cells(grid, unit_flux_at_faces, out=None):
    """Calculate net face fluxes at cells.

    Given a flux per unit width across each face in the grid, calculate the net
    outflux (or influx, if negative) at each cell.

    Construction::

        calc_net_face_flux_at_cells(grid, unit_flux_at_faces, out=None)

    Parameters
    ----------
    grid : ModelGrid
        A ModelGrid.
    unit_flux_at_faces : ndarray or field name
        Flux per unit width associated with faces.
    out : ndarray, optional
        Buffer to hold the result.

    Returns
    -------
    ndarray (x number of cells)
        Net flux at cells.

    Examples
    --------
    >>> from landlab import RasterModelGrid, CLOSED_BOUNDARY
    >>> rg = RasterModelGrid(3, 4, 10.0)
    >>> z = rg.add_zeros('node', 'topographic__elevation')
    >>> z[5] = 50.0
    >>> z[6] = 36.0
    >>> fg = rg.calculate_gradients_at_faces(z)  # there are 7 faces
    >>> fg
    array([ 5. ,  3.6,  5. , -1.4, -3.6, -5. , -3.6])
    >>> calc_net_face_flux_at_cells(rg, -fg)
    array([-164.,  -94.])
    >>> rg.set_status_at_node_on_edges(right=CLOSED_BOUNDARY)
    >>> rg.set_status_at_node_on_edges(top=CLOSED_BOUNDARY)
    >>> unit_flux_at_faces = np.zeros(rg.number_of_faces)
    >>> unit_flux_at_faces[rg.active_faces] = -fg[rg.active_faces]
    >>> calc_net_face_flux_at_cells(rg, unit_flux_at_faces)
    array([-114.,  -22.])

    >>> from landlab import HexModelGrid
    >>> hg = HexModelGrid(3, 3, 10.0)
    >>> z = rg.add_zeros('node', 'topographic__elevation', noclobber=False)
    >>> z[4] = 50.0
    >>> z[5] = 36.0
    >>> fg = hg.calculate_gradients_at_faces(z)  # there are 11 faces
    >>> fg
    array([ 5. ,  5. ,  3.6,  3.6,  5. , -1.4, -3.6, -5. , -5. , -3.6, -3.6])
    >>> nffc = calc_net_face_flux_at_cells(hg, -fg)
    >>> np.round(nffc)
    array([-152.,  -96.])

    Notes
    -----
    This is essentially a line integral for the fluxes along the boundaries of
    each cell. Hence, the resulting output has dimensions of total flux (so,
    if the unit flux happens to be mass per time per face width, the output
    will be in mass per unit time).
    """
    if out is None:
        out = grid.empty(centering='cell')
    total_flux = unit_flux_at_faces * grid.face_width
    out = np.zeros(grid.number_of_cells)
    fac = grid.faces_at_cell
    for c in range(grid.link_dirs_at_node.shape[1]):
        out += total_flux[fac[:,c]] \
        * grid.link_dirs_at_node[grid.node_at_cell,c]
    return out

@use_field_name_or_array('face')
def calc_face_flux_divergence_at_cells(grid, unit_flux_at_faces):
    """Calculate divergence of face-based fluxes at cells.

    Given a flux per unit width across each face in the grid, calculate the net
    outflux (or influx, if negative) divided by cell area, at each cell.

    Construction::

        calc_face_flux_divergence_at_cells(grid, unit_flux_at_faces, out=None)

    Parameters
    ----------
    grid : ModelGrid
        A ModelGrid.
    unit_flux_at_faces : ndarray or field name
        Flux per unit width associated with faces.

    Returns
    -------
    ndarray (x number of cells)
        Net flux at cells.

    Examples
    --------
    >>> from landlab import RasterModelGrid, CLOSED_BOUNDARY
    >>> rg = RasterModelGrid(3, 4, 10.0)
    >>> z = rg.add_zeros('node', 'topographic__elevation')
    >>> z[5] = 50.0
    >>> z[6] = 36.0
    >>> fg = rg.calculate_gradients_at_faces(z)  # there are 7 faces
    >>> fg
    array([ 5. ,  3.6,  5. , -1.4, -3.6, -5. , -3.6])
    >>> calc_face_flux_divergence_at_cells(rg, -fg)
    array([-1.64, -0.94])
    >>> rg.set_status_at_node_on_edges(right=CLOSED_BOUNDARY)
    >>> rg.set_status_at_node_on_edges(top=CLOSED_BOUNDARY)
    >>> unit_flux_at_faces = np.zeros(rg.number_of_faces)
    >>> unit_flux_at_faces[rg.active_faces] = -fg[rg.active_faces]
    >>> calc_face_flux_divergence_at_cells(rg, unit_flux_at_faces)
    array([-1.14, -0.22])

    Notes
    -----
    .
    """
    return calc_net_face_flux_at_cells(grid, unit_flux_at_faces) \
           / grid.area_of_cell
<commit_msg>add text to docstring in divergence.py<commit_after>#! /usr/bin/env python
"""Calculate vector divergence and related quantities at nodes or cells."""
import numpy as np
from landlab.utils.decorators import use_field_name_or_array

@use_field_name_or_array('face')
def calc_net_face_flux_at_cells(grid, unit_flux_at_faces, out=None):
    """Calculate net face fluxes at cells.

    Given a flux per unit width across each face in the grid, calculate the net
    outflux (or influx, if negative) at each cell.

    Construction::

        calc_net_face_flux_at_cells(grid, unit_flux_at_faces, out=None)

    Parameters
    ----------
    grid : ModelGrid
        A ModelGrid.
    unit_flux_at_faces : ndarray or field name
        Flux per unit width associated with faces.
    out : ndarray, optional
        Buffer to hold the result.

    Returns
    -------
    ndarray (x number of cells)
        Net flux at cells.

    Examples
    --------
    >>> from landlab import RasterModelGrid, CLOSED_BOUNDARY
    >>> rg = RasterModelGrid(3, 4, 10.0)
    >>> z = rg.add_zeros('node', 'topographic__elevation')
    >>> z[5] = 50.0
    >>> z[6] = 36.0
    >>> fg = rg.calculate_gradients_at_faces(z)  # there are 7 faces
    >>> fg
    array([ 5. ,  3.6,  5. , -1.4, -3.6, -5. , -3.6])
    >>> calc_net_face_flux_at_cells(rg, -fg)
    array([-164.,  -94.])
    >>> rg.set_status_at_node_on_edges(right=CLOSED_BOUNDARY)
    >>> rg.set_status_at_node_on_edges(top=CLOSED_BOUNDARY)
    >>> unit_flux_at_faces = np.zeros(rg.number_of_faces)
    >>> unit_flux_at_faces[rg.active_faces] = -fg[rg.active_faces]
    >>> calc_net_face_flux_at_cells(rg, unit_flux_at_faces)
    array([-114.,  -22.])

    >>> from landlab import HexModelGrid
    >>> hg = HexModelGrid(3, 3, 10.0)
    >>> z = rg.add_zeros('node', 'topographic__elevation', noclobber=False)
    >>> z[4] = 50.0
    >>> z[5] = 36.0
    >>> fg = hg.calculate_gradients_at_faces(z)  # there are 11 faces
    >>> fg
    array([ 5. ,  5. ,  3.6,  3.6,  5. , -1.4, -3.6, -5. , -5. , -3.6, -3.6])
    >>> nffc = calc_net_face_flux_at_cells(hg, -fg)
    >>> np.round(nffc)
    array([-152.,  -96.])

    Notes
    -----
    This is essentially a line integral for the fluxes along the boundaries of
    each cell. Hence, the resulting output has dimensions of total flux (so,
    if the unit flux happens to be mass per time per face width, the output
    will be in mass per unit time).
    """
    if out is None:
        out = grid.empty(centering='cell')
    total_flux = unit_flux_at_faces * grid.face_width
    out = np.zeros(grid.number_of_cells)
    fac = grid.faces_at_cell
    for c in range(grid.link_dirs_at_node.shape[1]):
        out += total_flux[fac[:,c]] \
        * grid.link_dirs_at_node[grid.node_at_cell,c]
    return out

@use_field_name_or_array('face')
def calc_face_flux_divergence_at_cells(grid, unit_flux_at_faces):
    """Calculate divergence of face-based fluxes at cells.

    Given a flux per unit width across each face in the grid, calculate the net
    outflux (or influx, if negative) divided by cell area, at each cell.

    Construction::

        calc_face_flux_divergence_at_cells(grid, unit_flux_at_faces, out=None)

    Parameters
    ----------
    grid : ModelGrid
        A ModelGrid.
    unit_flux_at_faces : ndarray or field name
        Flux per unit width associated with faces.

    Returns
    -------
    ndarray (x number of cells)
        Flux divergence at cells.

    Examples
    --------
    >>> from landlab import RasterModelGrid, CLOSED_BOUNDARY
    >>> rg = RasterModelGrid(3, 4, 10.0)
    >>> z = rg.add_zeros('node', 'topographic__elevation')
    >>> z[5] = 50.0
    >>> z[6] = 36.0
    >>> fg = rg.calculate_gradients_at_faces(z)  # there are 7 faces
    >>> fg
    array([ 5. ,  3.6,  5. , -1.4, -3.6, -5. , -3.6])
    >>> calc_face_flux_divergence_at_cells(rg, -fg)
    array([-1.64, -0.94])
    >>> rg.set_status_at_node_on_edges(right=CLOSED_BOUNDARY)
    >>> rg.set_status_at_node_on_edges(top=CLOSED_BOUNDARY)
    >>> unit_flux_at_faces = np.zeros(rg.number_of_faces)
    >>> unit_flux_at_faces[rg.active_faces] = -fg[rg.active_faces]
    >>> calc_face_flux_divergence_at_cells(rg, unit_flux_at_faces)
    array([-1.14, -0.22])

    Notes
    -----
    Performs a numerical flux divergence operation on cells.
    """
    return calc_net_face_flux_at_cells(grid, unit_flux_at_faces) \
           / grid.area_of_cell
<|endoftext|>


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

from pyqtgraph.Qt import QtGui, QtCore
from PyQt4.QtWebKit import QWebView
from PyQt4.QtCore import Qt
from PyQt4.QtGui import QApplication, QCursor
import time
import os
import yaml
import sys

if sys.version_info[0] == 3:
    import urllib.request
else:
    import urllib
# http://stackoverflow.com/questions/17960942/attributeerror-module-object-has-no-attribute-urlretrieve

from ..drivers import Oscillo, Spectrum
from ..drivers import OscilloSimu, SpectrumSimu
from .oscillo_widget import OscilloWidget
from .spectrum_widget import SpectrumWidget
from .connect_widget import ConnectWidget

class WelcomeWidget(QtGui.QWidget):
    """ This widget allows to connect to one of the available drivers.

    Note:
        When the connection is established, the drivers are `Oscillo` and
        `Spectrum`. Otherwise the drivers are in simulation mode (`OscilloSimu`
        and `SpectrumSimu`). A web page is also displayed on the left side.

    """
    def __init__(self, parent, ip_path):
        super(WelcomeWidget, self).__init__()

        with open('config.yaml') as config_file:
            self.config = yaml.load(config_file)

        self.parent = parent
        self.ip_path = ip_path
        self.opened = True
        self.select_opened = True

        # Define layouts
        self.lay = QtGui.QHBoxLayout()
        self.left_layout = QtGui.QVBoxLayout()
        self.right_layout = QtGui.QVBoxLayout()

        # Connection (ip address and password)
        self.connect_layout = QtGui.QVBoxLayout()
        self.connect_widget = ConnectWidget(self, self.ip_path)
        self.connect_layout.addWidget(self.connect_widget)

        # Select between drivers
        self.drivers_layout = QtGui.QVBoxLayout()

        self.instrument_buttons = []
        for i, instrument in enumerate(self.parent.instrument_list):
            self.instrument_buttons.append(self.set_button(instrument.capitalize() +' (Simu)'))
            self.drivers_layout.addWidget(self.instrument_buttons[i], 1, QtCore.Qt.AlignCenter)
            def make_callback(instrument):
                return lambda : self.instrument_onclick(instrument)
            self.instrument_buttons[i].clicked.connect(make_callback(instrument))
        #self.oscillo_button = self.set_button('Oscillo (Simu)')
        #self.spectrum_button = self.set_button('Spectrum (Simu)')
        #self.drivers_layout.addWidget(self.oscillo_button,1,QtCore.Qt.AlignCenter)
        #self.drivers_layout.addWidget(self.spectrum_button,1,QtCore.Qt.AlignCenter)

        # Left Layout
        self.view = QWebView()
        self.view.load(QtCore.QUrl.fromLocalFile(
                       os.path.join(self.parent.static_path, 'welcome.html')))
        self.left_layout.addWidget(self.view)

        # Right layout
        self.right_layout.addLayout(self.connect_layout)
        self.right_layout.addLayout(self.drivers_layout)
        self.right_layout.addStretch(1)
        self.right_frame = QtGui.QFrame(self)
        self.right_frame.setFrameShape(QtGui.QFrame.StyledPanel)
        self.right_frame.setLayout(self.right_layout)

        # Add layouts to main layout
        self.lay.addLayout(self.left_layout, 1)
        self.lay.addWidget(self.right_frame)
        self.setLayout(self.lay)

        # Connections
        #self.oscillo_button.clicked.connect(self.oscillo_onclick)
        #self.spectrum_button.clicked.connect(self.spectrum_onclick)

    def update(self):
        pass

    def set_button(self, name):
        button = QtGui.QPushButton(name)
        button.setStyleSheet('QPushButton {color: green;}')
        button.setFixedWidth(200)
        button.setFixedHeight(150)
        return button

    def set_connected(self):
        for i, button in enumerate(self.instrument_buttons):
            button.setText(self.parent.instrument_list[i].capitalize())

    def set_disconnected(self):
        for i, button in enumerate(self.instrument_buttons):
            button.setText(self.parent.instrument_list[i].capitalize()+' (Simu)')

    def instrument_onclick(self, instrument):
        print 'Instrument = ', instrument
        if self.connect_widget.is_connected:
            QApplication.setOverrideCursor(QCursor(Qt.WaitCursor))
            self.connect_widget.install_instrument(instrument)
            driver = globals()[instrument.capitalize()](self.connect_widget.client)
            driver.set_led(driver.client.host.split('.')[-1])
            QApplication.restoreOverrideCursor()
        else:
            driver = globals()[instrument.capitalize()+'Simu']()
        index = self.parent.stacked_widget.addWidget(globals()[instrument.capitalize()+'Widget'](driver, self.parent))
        self.parent.stacked_widget.setCurrentIndex(index)

    # def oscillo_onclick(self):
    #     if self.connect_widget.is_connected:
    #         QApplication.setOverrideCursor(QCursor(Qt.WaitCursor))
    #         self.connect_widget.install_instrument("oscillo")
    #         driver = Oscillo(self.connect_widget.client)
    #         driver.set_led(driver.client.host.split('.')[-1])
    #         QApplication.restoreOverrideCursor()
    #     else:
    #         driver = OscilloSimu()
    #     index = self.parent.stacked_widget.addWidget(OscilloWidget(driver, self.parent))
    #     self.parent.stacked_widget.setCurrentIndex(index)

    # def spectrum_onclick(self):
    #     if self.connect_widget.is_connected:
    #         QApplication.setOverrideCursor(QCursor(Qt.WaitCursor))
    #         self.connect_widget.install_instrument("spectrum")
    #         driver = Spectrum(self.connect_widget.client)
    #         driver.set_led(driver.client.host.split('.')[-1])
    #         QApplication.restoreOverrideCursor()
    #     else:
    #         driver = SpectrumSimu()
    #     index = self.parent.stacked_widget.addWidget(SpectrumWidget(driver, self.parent))
    #     self.parent.stacked_widget.setCurrentIndex(index)

<commit_msg>some cleaning<commit_after># -*- coding: utf-8 -*-

from pyqtgraph.Qt import QtGui, QtCore
from PyQt4.QtWebKit import QWebView
from PyQt4.QtCore import Qt
from PyQt4.QtGui import QApplication, QCursor
import time
import os
import yaml
import sys

if sys.version_info[0] == 3:
    import urllib.request
else:
    import urllib
# http://stackoverflow.com/questions/17960942/attributeerror-module-object-has-no-attribute-urlretrieve

from ..drivers import Oscillo, Spectrum
from ..drivers import OscilloSimu, SpectrumSimu
from .oscillo_widget import OscilloWidget
from .spectrum_widget import SpectrumWidget
from .connect_widget import ConnectWidget

class WelcomeWidget(QtGui.QWidget):
    """ This widget allows to connect to one of the available drivers.
    """
    def __init__(self, parent, ip_path):
        super(WelcomeWidget, self).__init__()

        with open('config.yaml') as config_file:
            self.config = yaml.load(config_file)

        self.parent = parent
        self.ip_path = ip_path
        self.opened = True
        self.select_opened = True

        # Define layouts
        self.lay = QtGui.QHBoxLayout()
        self.left_layout = QtGui.QVBoxLayout()
        self.right_layout = QtGui.QVBoxLayout()

        # Connection (ip address and password)
        self.connect_layout = QtGui.QVBoxLayout()
        self.connect_widget = ConnectWidget(self, self.ip_path)
        self.connect_layout.addWidget(self.connect_widget)

        # Select between drivers
        self.drivers_layout = QtGui.QVBoxLayout()

        self.instrument_buttons = []
        for i, instrument in enumerate(self.parent.instrument_list):
            self.instrument_buttons.append(self.set_button(instrument.capitalize() +' (Simu)'))
            self.drivers_layout.addWidget(self.instrument_buttons[i], 1, QtCore.Qt.AlignCenter)
            def make_callback(instrument):
                return lambda : self.instrument_onclick(instrument)
            self.instrument_buttons[i].clicked.connect(make_callback(instrument))

        # Left Layout
        self.view = QWebView()
        self.view.load(QtCore.QUrl.fromLocalFile(
                       os.path.join(self.parent.static_path, 'welcome.html')))
        self.left_layout.addWidget(self.view)

        # Right layout
        self.right_layout.addLayout(self.connect_layout)
        self.right_layout.addLayout(self.drivers_layout)
        self.right_layout.addStretch(1)
        self.right_frame = QtGui.QFrame(self)
        self.right_frame.setFrameShape(QtGui.QFrame.StyledPanel)
        self.right_frame.setLayout(self.right_layout)

        # Add layouts to main layout
        self.lay.addLayout(self.left_layout, 1)
        self.lay.addWidget(self.right_frame)
        self.setLayout(self.lay)

    def update(self):
        pass

    def set_button(self, name):
        button = QtGui.QPushButton(name)
        button.setStyleSheet('QPushButton {color: green;}')
        button.setFixedWidth(200)
        button.setFixedHeight(150)
        return button

    def set_connected(self):
        for i, button in enumerate(self.instrument_buttons):
            button.setText(self.parent.instrument_list[i].capitalize())

    def set_disconnected(self):
        for i, button in enumerate(self.instrument_buttons):
            button.setText(self.parent.instrument_list[i].capitalize()+' (Simu)')

    def instrument_onclick(self, instrument):
        print 'Instrument = ', instrument
        if self.connect_widget.is_connected:
            QApplication.setOverrideCursor(QCursor(Qt.WaitCursor))
            self.connect_widget.install_instrument(instrument)
            driver = globals()[instrument.capitalize()](self.connect_widget.client)
            driver.set_led(driver.client.host.split('.')[-1])
            QApplication.restoreOverrideCursor()
        else:
            driver = globals()[instrument.capitalize()+'Simu']()
        index = self.parent.stacked_widget.addWidget(globals()[instrument.capitalize()+'Widget'](driver, self.parent))
        self.parent.stacked_widget.setCurrentIndex(index)<|endoftext|>


<commit_before>from __future__ import annotations

import typing
from typing import List, Optional, Union

import numpy

from .. import tracers
from openfisca_core.indexed_enums import EnumArray

if typing.TYPE_CHECKING:
    from numpy.typing import ArrayLike

    Array = Union[EnumArray, ArrayLike]

class ComputationLog:

    _full_tracer: tracers.FullTracer

    def __init__(self, full_tracer: tracers.FullTracer) -> None:
        self._full_tracer = full_tracer

    def display(
            self,
            value: Optional[Array],
            ) -> str:
        if isinstance(value, EnumArray):
            value = value.decode_to_str()

        return numpy.array2string(value, max_line_width = float("inf"))

    def lines(
            self,
            aggregate: bool = False,
            max_depth: Optional[int] = None,
            ) -> List[str]:
        depth = 1

        lines_by_tree = [
            self._get_node_log(node, depth, aggregate, max_depth)
            for node
            in self._full_tracer.trees
            ]

        return self._compact(self._flatten(lines_by_tree))

    def print_log(self, aggregate = False, max_depth = None) -> None:
        """
        Print the computation log of a simulation.

        If ``aggregate`` is ``False`` (default), print the value of each
        computed vector.

        If ``aggregate`` is ``True``, only print the minimum, maximum, and
        average value of each computed vector.

        This mode is more suited for simulations on a large population.

        If ``max_depth`` is ``None`` (default), print the entire computation.

        If ``max_depth`` is set, for example to ``3``, only print computed
        vectors up to a depth of ``max_depth``.
        """
        for line in self.lines(aggregate, max_depth):
            print(line)  # noqa T001

    def _get_node_log(
            self,
            node: tracers.TraceNode,
            depth: int,
            aggregate: bool,
            max_depth: Optional[int],
            ) -> List[Optional[str]]:

        node_log = [self._print_line(depth, node, aggregate, max_depth)]

        children_logs = [
            self._get_node_log(child, depth + 1, aggregate, max_depth)
            for child
            in node.children
            ]

        return node_log + self._flatten(children_logs)

    def _print_line(
            self,
            depth: int,
            node: tracers.TraceNode,
            aggregate: bool,
            max_depth: Optional[int],
            ) -> Optional[str]:
        indent = '  ' * depth
        value = node.value

        if max_depth is not None and (depth > max_depth):
            return None

        if value is None:
            formatted_value = "{'avg': '?', 'max': '?', 'min': '?'}"

        elif aggregate:
            try:
                formatted_value = str({
                    'avg': numpy.mean(value),
                    'max': numpy.max(value),
                    'min': numpy.min(value),
                    })

            except TypeError:
                formatted_value = "{'avg': '?', 'max': '?', 'min': '?'}"

        else:
            formatted_value = self.display(value)

        return f"{indent}{node.name}<{node.period}> >> {formatted_value}"

    def _compact(
            self,
            list_: List[Optional[str]],
            ) -> List[str]:
        return [item for item in list_ if item is not None]

    def _flatten(
            self,
            lists: List[List[Optional[str]]],
            ) -> List[Optional[str]]:
        return [item for list_ in lists for item in list_]
<commit_msg>Do not use None<commit_after>from __future__ import annotations

import typing
from typing import List, Optional, Union

import numpy

from .. import tracers
from openfisca_core.indexed_enums import EnumArray

if typing.TYPE_CHECKING:
    from numpy.typing import ArrayLike

    Array = Union[EnumArray, ArrayLike]

class ComputationLog:

    _full_tracer: tracers.FullTracer

    def __init__(self, full_tracer: tracers.FullTracer) -> None:
        self._full_tracer = full_tracer

    def display(
            self,
            value: Optional[Array],
            ) -> str:
        if isinstance(value, EnumArray):
            value = value.decode_to_str()

        return numpy.array2string(value, max_line_width = float("inf"))

    def lines(
            self,
            aggregate: bool = False,
            max_depth: Optional[int] = None,
            ) -> List[str]:
        depth = 1

        lines_by_tree = [
            self._get_node_log(node, depth, aggregate, max_depth)
            for node
            in self._full_tracer.trees
            ]

        return self._flatten(lines_by_tree)

    def print_log(self, aggregate = False, max_depth = None) -> None:
        """
        Print the computation log of a simulation.

        If ``aggregate`` is ``False`` (default), print the value of each
        computed vector.

        If ``aggregate`` is ``True``, only print the minimum, maximum, and
        average value of each computed vector.

        This mode is more suited for simulations on a large population.

        If ``max_depth`` is ``None`` (default), print the entire computation.

        If ``max_depth`` is set, for example to ``3``, only print computed
        vectors up to a depth of ``max_depth``.
        """
        for line in self.lines(aggregate, max_depth):
            print(line)  # noqa T001

    def _get_node_log(
            self,
            node: tracers.TraceNode,
            depth: int,
            aggregate: bool,
            max_depth: Optional[int],
            ) -> List[str]:

        if max_depth is not None and (depth > max_depth):
            return []

        node_log = [self._print_line(depth, node, aggregate, max_depth)]

        children_logs = [
            self._get_node_log(child, depth + 1, aggregate, max_depth)
            for child
            in node.children
            ]

        return node_log + self._flatten(children_logs)

    def _print_line(
            self,
            depth: int,
            node: tracers.TraceNode,
            aggregate: bool,
            max_depth: Optional[int],
            ) -> str:
        indent = '  ' * depth
        value = node.value

        if value is None:
            formatted_value = "{'avg': '?', 'max': '?', 'min': '?'}"

        elif aggregate:
            try:
                formatted_value = str({
                    'avg': numpy.mean(value),
                    'max': numpy.max(value),
                    'min': numpy.min(value),
                    })

            except TypeError:
                formatted_value = "{'avg': '?', 'max': '?', 'min': '?'}"

        else:
            formatted_value = self.display(value)

        return f"{indent}{node.name}<{node.period}> >> {formatted_value}"

    def _flatten(
            self,
            lists: List[List[str]],
            ) -> List[str]:
        return [item for list_ in lists for item in list_]
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2015-2018 GEM Foundation
#
# OpenQuake 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.
#
# OpenQuake 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 OpenQuake. If not, see <http://www.gnu.org/licenses/>.
import logging
import operator
import numpy

from openquake.baselib.python3compat import zip, encode
from openquake.hazardlib.stats import set_rlzs_stats
from openquake.hazardlib.calc.stochastic import TWO32
from openquake.risklib import riskinput
from openquake.calculators import base
from openquake.calculators.export.loss_curves import get_loss_builder

U8 = numpy.uint8
U16 = numpy.uint16
U32 = numpy.uint32
F32 = numpy.float32
F64 = numpy.float64
U64 = numpy.uint64
getweight = operator.attrgetter('weight')

def build_loss_tables(dstore):
    """
    Compute the total losses by rupture and losses by rlzi.
    """
    oq = dstore['oqparam']
    L = len(oq.loss_dt().names)
    R = dstore['csm_info'].get_num_rlzs()
    serials = dstore['ruptures']['serial']
    idx_by_ser = dict(zip(serials, range(len(serials))))
    tbl = numpy.zeros((len(serials), L), F32)
    lbr = numpy.zeros((R, L), F32)  # losses by rlz
    for rec in dstore['losses_by_event'].value:  # call .value for speed
        idx = idx_by_ser[rec['eid'] // TWO32]
        tbl[idx] += rec['loss']
        lbr[rec['rlzi']] += rec['loss']
    return tbl, lbr

def event_based_risk(riskinputs, riskmodel, param, monitor):
    """
    :param riskinputs:
        :class:`openquake.risklib.riskinput.RiskInput` objects
    :param riskmodel:
        a :class:`openquake.risklib.riskinput.CompositeRiskModel` instance
    :param param:
        a dictionary of parameters
    :param monitor:
        :class:`openquake.baselib.performance.Monitor` instance
    :returns:
        a dictionary of numpy arrays of shape (L, R)
    """
    I = param['insured_losses'] + 1
    L = len(riskmodel.lti)
    param['lrs_dt'] = numpy.dtype([('rlzi', U16), ('ratios', (F32, (L * I,)))])
    for ri in riskinputs:
        with monitor('getting hazard'):
            ri.hazard_getter.init()
            hazard = ri.hazard_getter.get_hazard()
        mon = monitor('build risk curves', measuremem=False)
        A = len(ri.aids)
        R = ri.hazard_getter.num_rlzs
        agg = []  # triples (eids, li, losses)
        try:
            avg = numpy.zeros((A, R, L * I), F32)
        except MemoryError:
            raise MemoryError(
                'Building array avg of shape (%d, %d, %d)' % (A, R, L*I))
        result = dict(aids=ri.aids, avglosses=avg)
        aid2idx = {aid: idx for idx, aid in enumerate(ri.aids)}
        if 'builder' in param:
            builder = param['builder']
            P = len(builder.return_periods)
            all_curves = numpy.zeros((A, R, P), builder.loss_dt)
        # update the result dictionary and the agg array with each output
        for out in riskmodel.gen_outputs(ri, monitor, hazard):
            if len(out.eids) == 0:  # this happens for sites with no events
                continue
            r = out.rlzi
            for l, loss_ratios in enumerate(out):
                if loss_ratios is None:  # for GMFs below the minimum_intensity
                    continue
                loss_type = riskmodel.loss_types[l]
                agglosses = numpy.zeros((len(out.eids), I), F32)
                for a, asset in enumerate(out.assets):
                    ratios = loss_ratios[a]  # shape (E, I)
                    aid = asset.ordinal
                    idx = aid2idx[aid]
                    aval = asset.value(loss_type)
                    losses = aval * ratios
                    if 'builder' in param:
                        with mon:  # this is the heaviest part
                            for i in range(I):
                                lt = loss_type + '_ins' * i
                                all_curves[idx, r][lt] = builder.build_curve(
                                    aval, ratios[:, i], r)

                    # average losses
                    if param['avg_losses']:
                        rat = ratios.sum(axis=0) * param['ses_ratio'] * aval
                        for i in range(I):
                            avg[idx, r, l + L * i] = rat[i]

                    # agglosses
                    agglosses += losses

                # agglosses
                for i in range(I):
                    agg.append((out.eids, l + L * i, agglosses[:, i]))

        result['agglosses'] = agg
        if 'builder' in param:
            clp = param['conditional_loss_poes']
            result['curves-rlzs'], result['curves-stats'] = builder.pair(
                all_curves, param['stats'])
            if R > 1 and param['individual_curves'] is False:
                del result['curves-rlzs']
            if clp:
                result['loss_maps-rlzs'], result['loss_maps-stats'] = (
                    builder.build_maps(all_curves, clp, param['stats']))
                if R > 1 and param['individual_curves'] is False:
                    del result['loss_maps-rlzs']

        # store info about the GMFs, must be done at the end
        result['gmdata'] = ri.gmdata
        yield result

@base.calculators.add('event_based_risk')
class EbrCalculator(base.RiskCalculator):
    """
    Event based PSHA calculator generating the total losses by taxonomy
    """
    core_task = event_based_risk
    is_stochastic = True

    def pre_execute(self):
        oq = self.oqparam
        super().pre_execute('event_based')
        parent = self.datastore.parent
        if not self.oqparam.ground_motion_fields:
            return  # this happens in the reportwriter

        self.L = len(self.riskmodel.lti)
        self.T = len(self.assetcol.tagcol)
        self.A = len(self.assetcol)
        self.I = oq.insured_losses + 1
        if parent:
            self.datastore['csm_info'] = parent['csm_info']
            self.eidrlz = parent['events'].value[['eid', 'rlz']].copy()
        else:
            self.eidrlz = self.datastore['events'].value[['eid', 'rlz']].copy()
        self.eidrlz.sort()
        if oq.return_periods != [0]:
            # setting return_periods = 0 disable loss curves and maps
            eff_time = oq.investigation_time * oq.ses_per_logic_tree_path
            if eff_time < 2:
                logging.warn('eff_time=%s is too small to compute loss curves',
                             eff_time)
            else:
                self.param['builder'] = get_loss_builder(
                    parent if parent else self.datastore,
                    oq.return_periods, oq.loss_dt())
        # sorting the eids is essential to get the epsilons in the right
        # order (i.e. consistent with the one used in ebr from ruptures)
        self.E = len(self.eidrlz)
        eps = self.epsilon_getter()()
        self.riskinputs = self.build_riskinputs('gmf', eps, self.E)
        self.param['num_events'] = self.E
        self.param['insured_losses'] = oq.insured_losses
        self.param['avg_losses'] = oq.avg_losses
        self.param['ses_ratio'] = oq.ses_ratio
        self.param['stats'] = oq.risk_stats()
        self.param['conditional_loss_poes'] = oq.conditional_loss_poes
        self.taskno = 0
        self.start = 0
        avg_losses = self.oqparam.avg_losses
        if avg_losses:
            self.dset = self.datastore.create_dset(
                'avg_losses-rlzs', F32, (self.A, self.R, self.L * self.I))
        self.agglosses = numpy.zeros((self.E, self.L * self.I), F32)
        if 'builder' in self.param:
            self.build_datasets(self.param['builder'])
        if parent:
            parent.close()  # avoid fork issues

    def build_datasets(self, builder):
        oq = self.oqparam
        R = len(builder.weights)
        stats = oq.risk_stats()
        A = self.A
        S = len(stats)
        P = len(builder.return_periods)
        C = len(self.oqparam.conditional_loss_poes)
        self.loss_maps_dt = oq.loss_dt((F32, (C,)))
        if oq.individual_curves:
            self.datastore.create_dset(
                'curves-rlzs', builder.loss_dt, (A, R, P), fillvalue=None)
            self.datastore.set_attrs(
                'curves-rlzs', return_periods=builder.return_periods)
        if oq.conditional_loss_poes:
            self.datastore.create_dset(
                'loss_maps-rlzs', self.loss_maps_dt, (A, R), fillvalue=None)
        if R > 1:
            self.datastore.create_dset(
                'curves-stats', builder.loss_dt, (A, S, P), fillvalue=None)
            self.datastore.set_attrs(
                'curves-stats', return_periods=builder.return_periods,
                stats=[encode(name) for (name, func) in stats])
            if oq.conditional_loss_poes:
                self.datastore.create_dset(
                    'loss_maps-stats', self.loss_maps_dt, (A, S),
                    fillvalue=None)
                self.datastore.set_attrs(
                    'loss_maps-stats',
                    stats=[encode(name) for (name, func) in stats])

    def epsilon_getter(self):
        """
        :returns: a callable (start, stop) producing a slice of epsilons
        """
        return riskinput.make_epsilon_getter(
            len(self.assetcol), self.E,
            self.oqparam.asset_correlation,
            self.oqparam.master_seed,
            self.oqparam.ignore_covs or not self.riskmodel.covs)

    def save_losses(self, dic):
        """
        Save the event loss tables incrementally.

        :param dic:
            dictionary with agglosses, avglosses
        """
        aids = dic.pop('aids')
        for eids, li, agglosses in dic.pop('agglosses'):
            self.agglosses[eids, li] += agglosses
        if self.oqparam.avg_losses:
            self.dset[aids, :, :] = dic.pop('avglosses')
        self._save_curves(dic, aids)
        self._save_maps(dic, aids)
        self.taskno += 1

    def _save_curves(self, dic, aids):
        for key in ('curves-rlzs', 'curves-stats'):
            array = dic.get(key)  # shape (A, S, P)
            if array is not None:
                self.datastore[key][aids, :, :] = array

    def _save_maps(self, dic, aids):
        for key in ('loss_maps-rlzs', 'loss_maps-stats'):
            array = dic.get(key)  # shape (A, S)
            if array is not None:
                loss_maps = numpy.zeros(array.shape[:2], self.loss_maps_dt)
                for lti, lt in enumerate(self.loss_maps_dt.names):
                    loss_maps[lt] = array[:, :, :, lti]
                self.datastore[key][aids, :] = loss_maps

    def combine(self, dummy, res):
        """
        :param dummy: unused parameter
        :param res: a result dictionary
        """
        with self.monitor('saving losses', measuremem=True):
            self.save_losses(res)
        return 1

    def post_execute(self, result):
        """
        Save risk data and build the aggregate loss curves
        """
        logging.info('Saving event loss table')
        elt_dt = numpy.dtype(
            [('eid', U64), ('rlzi', U16), ('loss', (F32, (self.L * self.I,)))])
        with self.monitor('saving event loss table', measuremem=True):
            agglosses = numpy.fromiter(
                ((eid, rlz, losses)
                 for (eid, rlz), losses in zip(self.eidrlz, self.agglosses)
                 if losses.any()), elt_dt)
            self.datastore['losses_by_event'] = agglosses
            loss_types = ' '.join(self.oqparam.loss_dt().names)
            self.datastore.set_attrs('losses_by_event', loss_types=loss_types)
        self.postproc()

    def postproc(self):
        """
        Build aggregate loss curves in process
        """
        dstore = self.datastore
        self.before_export()  # set 'realizations'
        oq = self.oqparam
        stats = oq. risk_stats()
        # store avg_losses-stats
        if oq.avg_losses:
            set_rlzs_stats(self.datastore, 'avg_losses')
        try:
            b = self.param['builder']
        except KeyError:  # don't build auxiliary tables
            return
        if dstore.parent:
            dstore.parent.open('r')  # to read the ruptures
        if 'ruptures' in self.datastore and len(self.datastore['ruptures']):
            logging.info('Building loss tables')
            with self.monitor('building loss tables', measuremem=True):
                rlt, lbr = build_loss_tables(dstore)
                dstore['rup_loss_table'] = rlt
                dstore['losses_by_rlzi'] = lbr
                ridx = [rlt[:, lti].argmax() for lti in range(self.L)]
                dstore.set_attrs('rup_loss_table', ridx=ridx)
        logging.info('Building aggregate loss curves')
        with self.monitor('building agg_curves', measuremem=True):
            array, arr_stats = b.build(dstore['losses_by_event'].value, stats)
        units = self.assetcol.cost_calculator.get_units(
            loss_types=array.dtype.names)
        if oq.individual_curves:
            self.datastore['agg_curves-rlzs'] = array
            self.datastore.set_attrs(
                'agg_curves-rlzs',
                return_periods=b.return_periods, units=units)
        if arr_stats is not None:
            self.datastore['agg_curves-stats'] = arr_stats
            self.datastore.set_attrs(
                'agg_curves-stats', return_periods=b.return_periods,
                stats=[encode(name) for (name, func) in stats], units=units)
<commit_msg>Cleanup [skip hazardlib]<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2015-2018 GEM Foundation
#
# OpenQuake 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.
#
# OpenQuake 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 OpenQuake. If not, see <http://www.gnu.org/licenses/>.
import logging
import operator
import numpy

from openquake.baselib.python3compat import zip, encode
from openquake.hazardlib.stats import set_rlzs_stats
from openquake.hazardlib.calc.stochastic import TWO32
from openquake.risklib import riskinput
from openquake.calculators import base
from openquake.calculators.export.loss_curves import get_loss_builder

U8 = numpy.uint8
U16 = numpy.uint16
U32 = numpy.uint32
F32 = numpy.float32
F64 = numpy.float64
U64 = numpy.uint64
getweight = operator.attrgetter('weight')

def build_loss_tables(dstore):
    """
    Compute the total losses by rupture and losses by rlzi.
    """
    oq = dstore['oqparam']
    L = len(oq.loss_dt().names)
    R = dstore['csm_info'].get_num_rlzs()
    serials = dstore['ruptures']['serial']
    idx_by_ser = dict(zip(serials, range(len(serials))))
    tbl = numpy.zeros((len(serials), L), F32)
    lbr = numpy.zeros((R, L), F32)  # losses by rlz
    for rec in dstore['losses_by_event'].value:  # call .value for speed
        idx = idx_by_ser[rec['eid'] // TWO32]
        tbl[idx] += rec['loss']
        lbr[rec['rlzi']] += rec['loss']
    return tbl, lbr

def event_based_risk(riskinputs, riskmodel, param, monitor):
    """
    :param riskinputs:
        :class:`openquake.risklib.riskinput.RiskInput` objects
    :param riskmodel:
        a :class:`openquake.risklib.riskinput.CompositeRiskModel` instance
    :param param:
        a dictionary of parameters
    :param monitor:
        :class:`openquake.baselib.performance.Monitor` instance
    :returns:
        a dictionary of numpy arrays of shape (L, R)
    """
    I = param['insured_losses'] + 1
    L = len(riskmodel.lti)
    param['lrs_dt'] = numpy.dtype([('rlzi', U16), ('ratios', (F32, (L * I,)))])
    for ri in riskinputs:
        with monitor('getting hazard'):
            ri.hazard_getter.init()
            hazard = ri.hazard_getter.get_hazard()
        mon = monitor('build risk curves', measuremem=False)
        A = len(ri.aids)
        R = ri.hazard_getter.num_rlzs
        agg = []  # triples (eids, li, losses)
        try:
            avg = numpy.zeros((A, R, L * I), F32)
        except MemoryError:
            raise MemoryError(
                'Building array avg of shape (%d, %d, %d)' % (A, R, L*I))
        result = dict(aids=ri.aids, avglosses=avg)
        aid2idx = {aid: idx for idx, aid in enumerate(ri.aids)}
        if 'builder' in param:
            builder = param['builder']
            P = len(builder.return_periods)
            all_curves = numpy.zeros((A, R, P), builder.loss_dt)
        # update the result dictionary and the agg array with each output
        for out in riskmodel.gen_outputs(ri, monitor, hazard):
            if len(out.eids) == 0:  # this happens for sites with no events
                continue
            r = out.rlzi
            agglosses = numpy.zeros((len(out.eids), L*I), F32)
            for l, loss_ratios in enumerate(out):
                if loss_ratios is None:  # for GMFs below the minimum_intensity
                    continue
                loss_type = riskmodel.loss_types[l]
                for a, asset in enumerate(out.assets):
                    ratios = loss_ratios[a]  # shape (E, I)
                    aid = asset.ordinal
                    idx = aid2idx[aid]
                    aval = asset.value(loss_type)
                    losses = aval * ratios
                    if 'builder' in param:
                        with mon:  # this is the heaviest part
                            for i in range(I):
                                lt = loss_type + '_ins' * i
                                all_curves[idx, r][lt] = builder.build_curve(
                                    aval, ratios[:, i], r)

                    # average losses
                    if param['avg_losses']:
                        rat = ratios.sum(axis=0) * param['ses_ratio'] * aval
                        for i in range(I):
                            avg[idx, r, l + L * i] = rat[i]

                    # agglosses
                    for i in range(I):
                        agglosses[:, l + L * i] += losses[:, i]

            # agglosses
            agg.append((out.eids, agglosses))

        result['agglosses'] = agg
        if 'builder' in param:
            clp = param['conditional_loss_poes']
            result['curves-rlzs'], result['curves-stats'] = builder.pair(
                all_curves, param['stats'])
            if R > 1 and param['individual_curves'] is False:
                del result['curves-rlzs']
            if clp:
                result['loss_maps-rlzs'], result['loss_maps-stats'] = (
                    builder.build_maps(all_curves, clp, param['stats']))
                if R > 1 and param['individual_curves'] is False:
                    del result['loss_maps-rlzs']

        # store info about the GMFs, must be done at the end
        result['gmdata'] = ri.gmdata
        yield result

@base.calculators.add('event_based_risk')
class EbrCalculator(base.RiskCalculator):
    """
    Event based PSHA calculator generating the total losses by taxonomy
    """
    core_task = event_based_risk
    is_stochastic = True

    def pre_execute(self):
        oq = self.oqparam
        super().pre_execute('event_based')
        parent = self.datastore.parent
        if not self.oqparam.ground_motion_fields:
            return  # this happens in the reportwriter

        self.L = len(self.riskmodel.lti)
        self.T = len(self.assetcol.tagcol)
        self.A = len(self.assetcol)
        self.I = oq.insured_losses + 1
        if parent:
            self.datastore['csm_info'] = parent['csm_info']
            self.eidrlz = parent['events'].value[['eid', 'rlz']].copy()
        else:
            self.eidrlz = self.datastore['events'].value[['eid', 'rlz']].copy()
        self.eidrlz.sort()
        if oq.return_periods != [0]:
            # setting return_periods = 0 disable loss curves and maps
            eff_time = oq.investigation_time * oq.ses_per_logic_tree_path
            if eff_time < 2:
                logging.warn('eff_time=%s is too small to compute loss curves',
                             eff_time)
            else:
                self.param['builder'] = get_loss_builder(
                    parent if parent else self.datastore,
                    oq.return_periods, oq.loss_dt())
        # sorting the eids is essential to get the epsilons in the right
        # order (i.e. consistent with the one used in ebr from ruptures)
        self.E = len(self.eidrlz)
        eps = self.epsilon_getter()()
        self.riskinputs = self.build_riskinputs('gmf', eps, self.E)
        self.param['num_events'] = self.E
        self.param['insured_losses'] = oq.insured_losses
        self.param['avg_losses'] = oq.avg_losses
        self.param['ses_ratio'] = oq.ses_ratio
        self.param['stats'] = oq.risk_stats()
        self.param['conditional_loss_poes'] = oq.conditional_loss_poes
        self.taskno = 0
        self.start = 0
        avg_losses = self.oqparam.avg_losses
        if avg_losses:
            self.dset = self.datastore.create_dset(
                'avg_losses-rlzs', F32, (self.A, self.R, self.L * self.I))
        self.agglosses = numpy.zeros((self.E, self.L * self.I), F32)
        if 'builder' in self.param:
            self.build_datasets(self.param['builder'])
        if parent:
            parent.close()  # avoid fork issues

    def build_datasets(self, builder):
        oq = self.oqparam
        R = len(builder.weights)
        stats = oq.risk_stats()
        A = self.A
        S = len(stats)
        P = len(builder.return_periods)
        C = len(self.oqparam.conditional_loss_poes)
        self.loss_maps_dt = oq.loss_dt((F32, (C,)))
        if oq.individual_curves:
            self.datastore.create_dset(
                'curves-rlzs', builder.loss_dt, (A, R, P), fillvalue=None)
            self.datastore.set_attrs(
                'curves-rlzs', return_periods=builder.return_periods)
        if oq.conditional_loss_poes:
            self.datastore.create_dset(
                'loss_maps-rlzs', self.loss_maps_dt, (A, R), fillvalue=None)
        if R > 1:
            self.datastore.create_dset(
                'curves-stats', builder.loss_dt, (A, S, P), fillvalue=None)
            self.datastore.set_attrs(
                'curves-stats', return_periods=builder.return_periods,
                stats=[encode(name) for (name, func) in stats])
            if oq.conditional_loss_poes:
                self.datastore.create_dset(
                    'loss_maps-stats', self.loss_maps_dt, (A, S),
                    fillvalue=None)
                self.datastore.set_attrs(
                    'loss_maps-stats',
                    stats=[encode(name) for (name, func) in stats])

    def epsilon_getter(self):
        """
        :returns: a callable (start, stop) producing a slice of epsilons
        """
        return riskinput.make_epsilon_getter(
            len(self.assetcol), self.E,
            self.oqparam.asset_correlation,
            self.oqparam.master_seed,
            self.oqparam.ignore_covs or not self.riskmodel.covs)

    def save_losses(self, dic):
        """
        Save the event loss tables incrementally.

        :param dic:
            dictionary with agglosses, avglosses
        """
        aids = dic.pop('aids')
        for eids, agglosses in dic.pop('agglosses'):
            self.agglosses[eids] += agglosses  # shape (E, LI)
        if self.oqparam.avg_losses:
            self.dset[aids, :, :] = dic.pop('avglosses')
        self._save_curves(dic, aids)
        self._save_maps(dic, aids)
        self.taskno += 1

    def _save_curves(self, dic, aids):
        for key in ('curves-rlzs', 'curves-stats'):
            array = dic.get(key)  # shape (A, S, P)
            if array is not None:
                self.datastore[key][aids, :, :] = array

    def _save_maps(self, dic, aids):
        for key in ('loss_maps-rlzs', 'loss_maps-stats'):
            array = dic.get(key)  # shape (A, S)
            if array is not None:
                loss_maps = numpy.zeros(array.shape[:2], self.loss_maps_dt)
                for lti, lt in enumerate(self.loss_maps_dt.names):
                    loss_maps[lt] = array[:, :, :, lti]
                self.datastore[key][aids, :] = loss_maps

    def combine(self, dummy, res):
        """
        :param dummy: unused parameter
        :param res: a result dictionary
        """
        with self.monitor('saving losses', measuremem=True):
            self.save_losses(res)
        return 1

    def post_execute(self, result):
        """
        Save risk data and build the aggregate loss curves
        """
        logging.info('Saving event loss table')
        elt_dt = numpy.dtype(
            [('eid', U64), ('rlzi', U16), ('loss', (F32, (self.L * self.I,)))])
        with self.monitor('saving event loss table', measuremem=True):
            agglosses = numpy.fromiter(
                ((eid, rlz, losses)
                 for (eid, rlz), losses in zip(self.eidrlz, self.agglosses)
                 if losses.any()), elt_dt)
            self.datastore['losses_by_event'] = agglosses
            loss_types = ' '.join(self.oqparam.loss_dt().names)
            self.datastore.set_attrs('losses_by_event', loss_types=loss_types)
        self.postproc()

    def postproc(self):
        """
        Build aggregate loss curves in process
        """
        dstore = self.datastore
        self.before_export()  # set 'realizations'
        oq = self.oqparam
        stats = oq. risk_stats()
        # store avg_losses-stats
        if oq.avg_losses:
            set_rlzs_stats(self.datastore, 'avg_losses')
        try:
            b = self.param['builder']
        except KeyError:  # don't build auxiliary tables
            return
        if dstore.parent:
            dstore.parent.open('r')  # to read the ruptures
        if 'ruptures' in self.datastore and len(self.datastore['ruptures']):
            logging.info('Building loss tables')
            with self.monitor('building loss tables', measuremem=True):
                rlt, lbr = build_loss_tables(dstore)
                dstore['rup_loss_table'] = rlt
                dstore['losses_by_rlzi'] = lbr
                ridx = [rlt[:, lti].argmax() for lti in range(self.L)]
                dstore.set_attrs('rup_loss_table', ridx=ridx)
        logging.info('Building aggregate loss curves')
        with self.monitor('building agg_curves', measuremem=True):
            array, arr_stats = b.build(dstore['losses_by_event'].value, stats)
        units = self.assetcol.cost_calculator.get_units(
            loss_types=array.dtype.names)
        if oq.individual_curves:
            self.datastore['agg_curves-rlzs'] = array
            self.datastore.set_attrs(
                'agg_curves-rlzs',
                return_periods=b.return_periods, units=units)
        if arr_stats is not None:
            self.datastore['agg_curves-stats'] = arr_stats
            self.datastore.set_attrs(
                'agg_curves-stats', return_periods=b.return_periods,
                stats=[encode(name) for (name, func) in stats], units=units)
<|endoftext|>