<commit_before>import math
import os
import shutil
import socket
import stat
import subprocess
import textwrap
import time

REDIS_SERVER_EXEC = 'redis-server'
REDIS_SLOT_COUNT = 16384
_MAX_RETRY_ERRORS = 4
_ATTEMPT_INTERVAL = 0.3

from aioredis.log import logger
__all__ = ['setup_test_cluster']

def setup_test_cluster(*args, **kwargs):
    cluster = TestCluster(*args, **kwargs)
    cluster.setup()
    return cluster

class TestCluster:
    """This class allows to create a local Redis cluster for test purposes. It also includes methods to stop and
    restart nodes to test failover behaviour.
    """
    def __init__(self, redis_count, start_port, directory, node_timeout=3000):
        self.redis_count = redis_count
        self.start_port = start_port
        self.directory = os.path.abspath(directory)
        self.node_timeout = node_timeout
        self.processes = {}

    def setup(self):
        self._setup_directory()
        self._create_redis_directories()
        for port in self.get_ports():
            self._start_redis(port)
        self._configure_cluster()

    def terminate(self):
        for process in self.processes.values():
            process.terminate()
        for process in self.processes.values():
            process.wait(1)

    def stop_redis(self, port):
        if port not in self.processes:
            raise ValueError('No Redis running at port {}.'.format(port))
        process = self.processes.pop(port)
        process.terminate()
        process.wait(1)

    def restart_redis(self, port):
        if port in self.processes:
            raise ValueError('Redis process at port {} is still running.'.format(port))

        self._start_redis(port)

    def get_ports(self):
        return [self.start_port + i for i in range(self.redis_count)]

    def _setup_directory(self):
        if not os.path.exists(self.directory):
            os.makedirs(self.directory)

    def _create_redis_directories(self):
        for port in self.get_ports():
            redis_directory = self._get_redis_directory(port)
            if not os.path.exists(redis_directory):
                os.mkdir(redis_directory)
            else:
                self._delete_directory_contents(redis_directory)

            self._write_redis_config_file(os.path.join(redis_directory, 'redis.conf'), port)

    def _start_redis(self, port):
        directory = self._get_redis_directory(port)
        self.processes[port] = subprocess.Popen([REDIS_SERVER_EXEC, 'redis.conf'], cwd=directory)

    def _configure_cluster(self):
        time.sleep(_ATTEMPT_INTERVAL)  # Give cluster some time to start up

        addresses = [('127.0.0.1', port) for port in self.get_ports()]
        sockets = self._connect_sockets(addresses)
        masters_count = math.ceil(self.redis_count / 2)
        masters = sockets[:masters_count]
        master_addresses = addresses[:masters_count]
        slaves = sockets[masters_count:]
        master_node_ids = [self._determine_node_id(master, address) for master, address in zip(masters, addresses)]

        self._assign_slots(masters, master_addresses)
        self._send_meet_messages_to_all(sockets, addresses)
        time.sleep(_ATTEMPT_INTERVAL)  # MEET messages need some time to propagate
        self._send_replicate_messages(slaves, master_node_ids)
        self._wait_until_cluster_state_ok(sockets)

        for socket in sockets:
            socket.close()

    def _connect_sockets(self, addresses):
        return self._retry(socket.create_connection, addresses, error_message='Could not connect to Redis.')

    def _determine_node_id(self, socket, address):
        socket.sendall(b'CLUSTER NODES\r\n')
        data = self._read_bulk_string_response(socket)
        node_id = data[:40].decode('utf-8')
        logger.debug("Master at {} has node id {}.".format(address, node_id))
        return node_id

    def _assign_slots(self, masters, addresses):
        slot_boundaries = [math.floor(i * REDIS_SLOT_COUNT / len(masters)) for i in range(len(masters) + 1)]
        slot_ranges = [range(b1, b2) for b1, b2 in zip(slot_boundaries, slot_boundaries[1:])]
        for master, slot_range, address in zip(masters, slot_ranges, addresses):
            logger.debug("Assigning master at {} slots {}-{}".format(address, slot_range.start, slot_range.stop - 1))
            slots = ' '.join(str(slot) for slot in slot_range)
            try:
                self._send_command_and_expect_ok(master, 'CLUSTER ADDSLOTS {}\r\n'.format(slots))
            except IOError as e:
                raise IOError('ADDSLOTS failed. Maybe a cluster is already running? ({}).'.format(str(e)))

    def _send_meet_messages_to_all(self, sockets, addresses):
        for i, socket in enumerate(sockets):
            for j, address in enumerate(addresses):
                if i != j:
                    self._send_command_and_expect_ok(socket, 'CLUSTER MEET {} {}\r\n'.format(*address))

    def _send_replicate_messages(self, slaves, master_node_ids):
        def _send_replicate_message(arg):
            slave, master_node_id = arg
            self._send_command_and_expect_ok(slave, 'CLUSTER REPLICATE {}\r\n'.format(master_node_id))

        self._retry(_send_replicate_message, list(zip(slaves, master_node_ids)), 'Replication failed.')

    def _wait_until_cluster_state_ok(self, sockets):
        def _check_state(socket):
            socket.sendall(b'CLUSTER INFO\r\n')
            data = self._read_bulk_string_response(socket).decode('utf-8')
            if 'cluster_state:ok' not in data:
                raise IOError('Cluster state not ok')

        self._retry(_check_state, sockets, error_message='Cluster state not ok.', max_errors=10)

    def _retry(self, method, arguments, error_message, max_errors=_MAX_RETRY_ERRORS, interval=_ATTEMPT_INTERVAL):
        results = [None] * len(arguments)
        successful_indexes = []
        errors = 0
        while len(successful_indexes) < len(arguments):
            for i, argument in enumerate(arguments):
                if i not in successful_indexes:
                    try:
                        results[i] = method(argument)
                        successful_indexes.append(i)
                    except (IOError, ConnectionRefusedError):
                        errors += 1
                        if errors >= max_errors:
                            raise IOError(error_message + ' Stop retrying after {} errors.'.format(errors))
                        else:
                            logger.info(error_message + ' Will retry after {}s.'.format(interval))
                            time.sleep(interval)

        return results

    def _recv_until(self, socket, delimiter):
        data = b''
        while delimiter not in data:
            data += socket.recv(1024)

        index = data.index(delimiter)
        return data[:index], data[index + len(delimiter):]

    def _recv_bytes(self, socket, byte_count):
        data = b''
        while len(data) < byte_count:
            received = socket.recv(min(1024, byte_count - len(data)))
            if len(received) == 0:
                raise IOError('Socket closed')
            else:
                data += received
        return data

    def _send_command_and_expect_ok(self, socket, command):
        socket.sendall(command.encode('utf-8'))
        response, _ = self._recv_until(socket, b'\r\n')
        if response != b'+OK':
            raise IOError(response.decode('utf-8'))

    def _read_bulk_string_response(self, socket):
        header, data = self._recv_until(socket, b'\r\n')
        if header[0] != ord('$'):
            raise ValueError('Expected bulk string response.')
        byte_count = int(header[1:].decode('utf-8'))
        remaining_data = self._recv_bytes(socket, byte_count - len(data) + 2)
        if not remaining_data[:-2] != 'b\r\n':
            raise ValueError('Invalid bulk string received.')
        data += remaining_data[:-2]
        return data

    def _get_redis_directory(self, port):
        return os.path.join(self.directory, 'redis-{}'.format(port))

    @staticmethod
    def _delete_directory_contents(directory):
        for name in os.listdir(directory):
            path = os.path.join(directory, name)
            if os.path.isdir(path):
                shutil.rmtree(path)
            else:
                os.remove(path)

    def _write_redis_config_file(self, path, port):
        with open(path, 'w') as file:
            file.write(textwrap.dedent("""
                port {}
                cluster-enabled yes
                cluster-config-file nodes.conf
                cluster-node-timeout {}
                """.format(port, self.node_timeout)
            ))

        # Protect against the CONFIG REWRITE test
        self._remove_write_access(path)

    @staticmethod
    def _remove_write_access(path):
        os.chmod(path, stat.S_IRUSR | stat.S_IRGRP | stat.S_IROTH)
<commit_msg>Added a docstring.<commit_after>import math
import os
import shutil
import socket
import stat
import subprocess
import textwrap
import time

REDIS_SERVER_EXEC = 'redis-server'
REDIS_SLOT_COUNT = 16384
_MAX_RETRY_ERRORS = 4
_ATTEMPT_INTERVAL = 0.3

from aioredis.log import logger
__all__ = ['setup_test_cluster']

def setup_test_cluster(*args, **kwargs):
    cluster = TestCluster(*args, **kwargs)
    cluster.setup()
    return cluster

class TestCluster:
    """This class allows to create a local Redis cluster for test purposes. It also includes methods to stop and
    restart nodes to test failover behaviour.

    Parameters:
        - *redis_count* is the number of processes the cluster should contain. The first half of them will be
        master nodes (rounding upwards).
        - *start_port*: The cluster will use the ports from start_port to start_port + redis_count.
        - *directory* is used to store the configuration files of all processes.
        - *node_timeout*: The cluster node timeout in millliseconds, see http://redis.io/topics/cluster-tutorial.

    """
    def __init__(self, redis_count, start_port, directory, node_timeout=3000):
        self.redis_count = redis_count
        self.start_port = start_port
        self.directory = os.path.abspath(directory)
        self.node_timeout = node_timeout
        self.processes = {}

    def setup(self):
        self._setup_directory()
        self._create_redis_directories()
        for port in self.get_ports():
            self._start_redis(port)
        self._configure_cluster()

    def terminate(self):
        for process in self.processes.values():
            process.terminate()
        for process in self.processes.values():
            process.wait(1)

    def stop_redis(self, port):
        if port not in self.processes:
            raise ValueError('No Redis running at port {}.'.format(port))
        process = self.processes.pop(port)
        process.terminate()
        process.wait(1)

    def restart_redis(self, port):
        if port in self.processes:
            raise ValueError('Redis process at port {} is still running.'.format(port))

        self._start_redis(port)

    def get_ports(self):
        return [self.start_port + i for i in range(self.redis_count)]

    def _setup_directory(self):
        if not os.path.exists(self.directory):
            os.makedirs(self.directory)

    def _create_redis_directories(self):
        for port in self.get_ports():
            redis_directory = self._get_redis_directory(port)
            if not os.path.exists(redis_directory):
                os.mkdir(redis_directory)
            else:
                self._delete_directory_contents(redis_directory)

            self._write_redis_config_file(os.path.join(redis_directory, 'redis.conf'), port)

    def _start_redis(self, port):
        directory = self._get_redis_directory(port)
        self.processes[port] = subprocess.Popen([REDIS_SERVER_EXEC, 'redis.conf'], cwd=directory)

    def _configure_cluster(self):
        time.sleep(_ATTEMPT_INTERVAL)  # Give cluster some time to start up

        addresses = [('127.0.0.1', port) for port in self.get_ports()]
        sockets = self._connect_sockets(addresses)
        masters_count = math.ceil(self.redis_count / 2)
        masters = sockets[:masters_count]
        master_addresses = addresses[:masters_count]
        slaves = sockets[masters_count:]
        master_node_ids = [self._determine_node_id(master, address) for master, address in zip(masters, addresses)]

        self._assign_slots(masters, master_addresses)
        self._send_meet_messages_to_all(sockets, addresses)
        time.sleep(_ATTEMPT_INTERVAL)  # MEET messages need some time to propagate
        self._send_replicate_messages(slaves, master_node_ids)
        self._wait_until_cluster_state_ok(sockets)

        for socket in sockets:
            socket.close()

    def _connect_sockets(self, addresses):
        return self._retry(socket.create_connection, addresses, error_message='Could not connect to Redis.')

    def _determine_node_id(self, socket, address):
        socket.sendall(b'CLUSTER NODES\r\n')
        data = self._read_bulk_string_response(socket)
        node_id = data[:40].decode('utf-8')
        logger.debug("Master at {} has node id {}.".format(address, node_id))
        return node_id

    def _assign_slots(self, masters, addresses):
        slot_boundaries = [math.floor(i * REDIS_SLOT_COUNT / len(masters)) for i in range(len(masters) + 1)]
        slot_ranges = [range(b1, b2) for b1, b2 in zip(slot_boundaries, slot_boundaries[1:])]
        for master, slot_range, address in zip(masters, slot_ranges, addresses):
            logger.debug("Assigning master at {} slots {}-{}".format(address, slot_range.start, slot_range.stop - 1))
            slots = ' '.join(str(slot) for slot in slot_range)
            try:
                self._send_command_and_expect_ok(master, 'CLUSTER ADDSLOTS {}\r\n'.format(slots))
            except IOError as e:
                raise IOError('ADDSLOTS failed. Maybe a cluster is already running? ({}).'.format(str(e)))

    def _send_meet_messages_to_all(self, sockets, addresses):
        for i, socket in enumerate(sockets):
            for j, address in enumerate(addresses):
                if i != j:
                    self._send_command_and_expect_ok(socket, 'CLUSTER MEET {} {}\r\n'.format(*address))

    def _send_replicate_messages(self, slaves, master_node_ids):
        def _send_replicate_message(arg):
            slave, master_node_id = arg
            self._send_command_and_expect_ok(slave, 'CLUSTER REPLICATE {}\r\n'.format(master_node_id))

        self._retry(_send_replicate_message, list(zip(slaves, master_node_ids)), 'Replication failed.')

    def _wait_until_cluster_state_ok(self, sockets):
        def _check_state(socket):
            socket.sendall(b'CLUSTER INFO\r\n')
            data = self._read_bulk_string_response(socket).decode('utf-8')
            if 'cluster_state:ok' not in data:
                raise IOError('Cluster state not ok')

        self._retry(_check_state, sockets, error_message='Cluster state not ok.', max_errors=10)

    def _retry(self, method, arguments, error_message, max_errors=_MAX_RETRY_ERRORS, interval=_ATTEMPT_INTERVAL):
        results = [None] * len(arguments)
        successful_indexes = []
        errors = 0
        while len(successful_indexes) < len(arguments):
            for i, argument in enumerate(arguments):
                if i not in successful_indexes:
                    try:
                        results[i] = method(argument)
                        successful_indexes.append(i)
                    except (IOError, ConnectionRefusedError):
                        errors += 1
                        if errors >= max_errors:
                            raise IOError(error_message + ' Stop retrying after {} errors.'.format(errors))
                        else:
                            logger.info(error_message + ' Will retry after {}s.'.format(interval))
                            time.sleep(interval)

        return results

    def _recv_until(self, socket, delimiter):
        data = b''
        while delimiter not in data:
            data += socket.recv(1024)

        index = data.index(delimiter)
        return data[:index], data[index + len(delimiter):]

    def _recv_bytes(self, socket, byte_count):
        data = b''
        while len(data) < byte_count:
            received = socket.recv(min(1024, byte_count - len(data)))
            if len(received) == 0:
                raise IOError('Socket closed')
            else:
                data += received
        return data

    def _send_command_and_expect_ok(self, socket, command):
        socket.sendall(command.encode('utf-8'))
        response, _ = self._recv_until(socket, b'\r\n')
        if response != b'+OK':
            raise IOError(response.decode('utf-8'))

    def _read_bulk_string_response(self, socket):
        header, data = self._recv_until(socket, b'\r\n')
        if header[0] != ord('$'):
            raise ValueError('Expected bulk string response.')
        byte_count = int(header[1:].decode('utf-8'))
        remaining_data = self._recv_bytes(socket, byte_count - len(data) + 2)
        if not remaining_data[:-2] != 'b\r\n':
            raise ValueError('Invalid bulk string received.')
        data += remaining_data[:-2]
        return data

    def _get_redis_directory(self, port):
        return os.path.join(self.directory, 'redis-{}'.format(port))

    @staticmethod
    def _delete_directory_contents(directory):
        for name in os.listdir(directory):
            path = os.path.join(directory, name)
            if os.path.isdir(path):
                shutil.rmtree(path)
            else:
                os.remove(path)

    def _write_redis_config_file(self, path, port):
        with open(path, 'w') as file:
            file.write(textwrap.dedent("""
                port {}
                cluster-enabled yes
                cluster-config-file nodes.conf
                cluster-node-timeout {}
                """.format(port, self.node_timeout)
            ))

        # Protect against the CONFIG REWRITE test
        self._remove_write_access(path)

    @staticmethod
    def _remove_write_access(path):
        os.chmod(path, stat.S_IRUSR | stat.S_IRGRP | stat.S_IROTH)
<|endoftext|>


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

# Copyright (c) 2010-2011, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License version 3
# only, as published by the Free Software Foundation.
#
# 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 Lesser General Public License version 3 for more details
# (a copy is included in the LICENSE file that accompanied this code).
#
# You should have received a copy of the GNU Lesser General Public License
# version 3 along with OpenQuake.  If not, see
# <http://www.gnu.org/licenses/lgpl-3.0.txt> for a copy of the LGPLv3 License.

# vim: tabstop=4 shiftwidth=4 softtabstop=4

"""
Probabilistic Event Mixin: defines the behaviour of a Job. Calls the
compute_risk task
"""

from numpy import zeros

from celery.exceptions import TimeoutError

from openquake import kvs
from openquake import logs
from openquake import shapes

from openquake.risk import common
from openquake.risk import probabilistic_event_based as prob
from openquake.parser import vulnerability

from openquake.risk.job import aggregate_loss_curve
from openquake.risk.job import general

from openquake.db import models

LOGGER = logs.LOG
DEFAULT_CONDITIONAL_LOSS_POE = 0.01

class ProbabilisticEventMixin(): # pylint: disable=W0232,W0201
    """Mixin for Probalistic Event Risk Job."""

    @general.preload
    @general.output
    def execute(self):
        """Execute the job."""

        aggregate_curve = prob.AggregateLossCurve()

        tasks = []
        for block_id in self.blocks_keys:
            LOGGER.debug("Starting task block, block_id = %s of %s"
                    % (block_id, len(self.blocks_keys)))

            tasks.append(general.compute_risk.delay(self.job_id, block_id))

        for task in tasks:
            try:
                # TODO (chris): Figure out where to put that timeout
                task.wait(timeout=None)

                aggregate_curve.append(task.result)
            except TimeoutError:
                # TODO(jmc): Cancel and respawn this task
                return

        curve = aggregate_curve.compute(self._tses(), self._time_span())
        aggregate_loss_curve.compute_aggregate_curve(self, curve)

        return results

# TODO (ac): Add doc!
# TODO (ac): Test it!
    def _tses(self):
# TODO (ac): Confirm this works regardless of the method of hazard calc
        histories = int(self["NUMBER_OF_SEISMICITY_HISTORIES"])
        realizations = int(self["NUMBER_OF_LOGIC_TREE_SAMPLES"])
        num_ses = histories * realizations

        return num_ses * self._time_span()

    def _time_span(self):
        return float(self["INVESTIGATION_TIME"])

    def _gmf_db_list(self, job_id):  # pylint: disable=R0201
        """Returns a list of the output IDs of all computed GMFs"""

        ids = models.Output.objects.filter(
            oq_job=job_id, output_type='gmf').values_list('id', flat=True)

        return list(ids)

    def _get_db_gmf(self, gmf_id):
        """Returns a field for the given GMF"""
        grid = self.region.grid
        field = zeros((grid.rows, grid.columns))

        gmf_sites = models.GmfData.objects.filter(output=gmf_id)

        for gmf_site in gmf_sites:
            loc = gmf_site.location
            site = shapes.Site(loc.x, loc.y)
            grid_point = grid.point_at(site)

            field[grid_point.row][grid_point.column] = gmf_site.ground_motion

        return shapes.Field(field)

    def _sites_to_gmf_keys(self, sites):
        """Returns the GMF keys "row!col" for the given site list"""
        keys = []

        for site in sites:
            risk_point = self.region.grid.point_at(site)
            key = "%s!%s" % (risk_point.row, risk_point.column)
            keys.append(key)

        return keys

    def _get_db_gmfs(self, sites, job_id):
        """Aggregates GMF data from the DB by site"""
        all_gmfs = self._gmf_db_list(job_id)
        gmf_keys = self._sites_to_gmf_keys(sites)
        gmfs = dict((k, []) for k in gmf_keys)

        for gmf_id in all_gmfs:
            field = self._get_db_gmf(gmf_id)

            for key in gmfs.keys():
                (row, col) = key.split("!")
                gmfs[key].append(field.get(int(row), int(col)))

        return gmfs

    def _get_kvs_gmfs(self, sites, histories, realizations):
        """Aggregates GMF data from the KVS by site"""
        gmf_keys = self._sites_to_gmf_keys(sites)
        gmfs = dict((k, []) for k in gmf_keys)

        for i in range(0, histories):
            for j in range(0, realizations):
                key = kvs.tokens.stochastic_set_key(self.job_id, i, j)
                fieldset = shapes.FieldSet.from_json(kvs.get(key),
                    self.region.grid)

                for field in fieldset:
                    for key in gmfs.keys():
                        (row, col) = key.split("!")
                        gmfs[key].append(field.get(int(row), int(col)))

        return gmfs

    def slice_gmfs(self, block_id):
        """Load and collate GMF values for all sites in this block. """
        block = general.Block.from_kvs(block_id)
        gmfs = self._get_db_gmfs(block.sites, self.job_id)

        for key, gmf_slice in gmfs.items():
            (row, col) = key.split("!")
            key_gmf = kvs.tokens.gmf_set_key(self.job_id, col, row)
            LOGGER.debug("GMF_SLICE for %s X %s : \n\t%s" % (
                    col, row, gmf_slice))
            timespan = float(self['INVESTIGATION_TIME'])
            gmf = {"IMLs": gmf_slice, "TSES": self._tses(),
                    "TimeSpan": self._timespan()}
            kvs.set_value_json_encoded(key_gmf, gmf)

# TODO (ac): Improve doc!
    def compute_risk(self, block_id, **kwargs):  # pylint: disable=W0613
        """This task computes risk for a block of sites. It requires to have
        pre-initialized in kvs:
         1) list of sites
         2) gmfs
         3) exposure portfolio (=assets)
         4) vulnerability
        """

        conditional_loss_poes = [float(x) for x in self.params.get(
                "CONDITIONAL_LOSS_POE", "0.01").split()]

        self.slice_gmfs(block_id)

        self.vuln_curves = vulnerability.load_vuln_model_from_kvs(
            self.job_id)

        block = general.Block.from_kvs(block_id)

        # aggregate the losses for this block
        aggregate_curve = prob.AggregateLossCurve()

        for point in block.grid(self.region):
            key = kvs.tokens.gmf_set_key(self.job_id, point.column, point.row)
            gmf_slice = kvs.get_value_json_decoded(key)

            asset_key = kvs.tokens.asset_key(
                self.job_id, point.row, point.column)

            for asset in kvs.get_list_json_decoded(asset_key):
                LOGGER.debug("Processing asset %s" % (asset))

                # loss ratios, used both to produce the curve
                # and to aggregate the losses
                loss_ratios = self.compute_loss_ratios(asset, gmf_slice)

                loss_ratio_curve = self.compute_loss_ratio_curve(
                    point.column, point.row, asset, gmf_slice, loss_ratios)

                aggregate_curve.append(loss_ratios * asset["assetValue"])

                if loss_ratio_curve is not None:
                    loss_curve = self.compute_loss_curve(
                        point.column, point.row, loss_ratio_curve, asset)

                    for loss_poe in conditional_loss_poes:
                        self.compute_conditional_loss(point.column, point.row,
                                loss_curve, asset, loss_poe)

        return aggregate_curve.losses

    def compute_loss_ratios(self, asset, gmf_slice):
        """For a given asset and ground motion field, computes the loss ratios
        used to obtain the related loss ratio curve and aggregate loss curve."""

        epsilon_provider = general.EpsilonProvider(self.params)

# TODO (ac): Extract to a method, and test it!
        vuln_function = self.vuln_curves.get(
                asset["vulnerabilityFunctionReference"], None)

        if not vuln_function:
            LOGGER.error(
                "Unknown vulnerability function %s for asset %s"
                % (asset["vulnerabilityFunctionReference"], asset["assetID"]))

            return None

        return prob.compute_loss_ratios(
            vuln_function, gmf_slice, epsilon_provider, asset)

    def compute_conditional_loss(self, col, row, loss_curve, asset, loss_poe):
        """Compute the conditional loss for a loss curve and Probability of
        Exceedance (PoE)."""

        loss_conditional = common.compute_conditional_loss(
            loss_curve, loss_poe)

        key = kvs.tokens.loss_key(
                self.job_id, row, col, asset["assetID"], loss_poe)

        LOGGER.debug("Conditional loss is %s, write to key %s" %
                (loss_conditional, key))

        kvs.set(key, loss_conditional)

    def compute_loss_ratio_curve(self, col, row, asset, gmf_slice, loss_ratios):
        """Compute the loss ratio curve for a single asset."""

        # fail if the asset has an unknown vulnerability code
        vuln_function = self.vuln_curves.get(
                asset["vulnerabilityFunctionReference"], None)

        if not vuln_function:
            LOGGER.error(
                "Unknown vulnerability function %s for asset %s"
                % (asset["vulnerabilityFunctionReference"], asset["assetID"]))

            return None

        epsilon_provider = general.EpsilonProvider(self.params)

        loss_ratio_curve = prob.compute_loss_ratio_curve(
                vuln_function, gmf_slice, epsilon_provider, asset,
                self._get_number_of_samples(), loss_ratios=loss_ratios)

        # NOTE (jmc): Early exit if the loss ratio is all zeros
        if not False in (loss_ratio_curve.ordinates == 0.0):
            return None

        key = kvs.tokens.loss_ratio_key(
            self.job_id, row, col, asset["assetID"])

        kvs.set(key, loss_ratio_curve.to_json())

        LOGGER.debug("Loss ratio curve is %s, write to key %s" %
                (loss_ratio_curve, key))

        return loss_ratio_curve

    def _get_number_of_samples(self):
        """Return the number of samples used to compute the loss ratio
        curve specified by the PROB_NUM_OF_SAMPLES parameter.

        Return None if the parameter is not specified, or empty or
        the value can't be casted to int.
        """

        number_of_samples = None
        raw_value = getattr(self, "PROB_NUM_OF_SAMPLES", None)

        if raw_value:
            try:
                number_of_samples = int(raw_value)
            except ValueError:
                LOGGER.error("PROB_NUM_OF_SAMPLES %s can't be converted "
                             "to int, using default value..." % raw_value)

        return number_of_samples

    def compute_loss_curve(self, column, row, loss_ratio_curve, asset):
        """Compute the loss curve for a single asset."""

        if asset is None:
            return None

        loss_curve = loss_ratio_curve.rescale_abscissae(asset["assetValue"])

        key = kvs.tokens.loss_curve_key(
            self.job_id, row, column, asset["assetID"])

        LOGGER.debug("Loss curve is %s, write to key %s" % (loss_curve, key))
        kvs.set(key, loss_curve.to_json())

        return loss_curve

general.RiskJobMixin.register("Event Based", ProbabilisticEventMixin)
<commit_msg>Updated doc, refactoring<commit_after># -*- coding: utf-8 -*-

# Copyright (c) 2010-2011, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License version 3
# only, as published by the Free Software Foundation.
#
# 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 Lesser General Public License version 3 for more details
# (a copy is included in the LICENSE file that accompanied this code).
#
# You should have received a copy of the GNU Lesser General Public License
# version 3 along with OpenQuake.  If not, see
# <http://www.gnu.org/licenses/lgpl-3.0.txt> for a copy of the LGPLv3 License.

# vim: tabstop=4 shiftwidth=4 softtabstop=4

"""
Probabilistic Event Mixin: defines the behaviour of a Job. Calls the
compute_risk task
"""

from numpy import zeros

from celery.exceptions import TimeoutError

from openquake import kvs
from openquake import logs
from openquake import shapes

from openquake.risk import common
from openquake.risk import probabilistic_event_based as prob
from openquake.parser import vulnerability

from openquake.risk.job import aggregate_loss_curve
from openquake.risk.job import general

from openquake.db import models

LOGGER = logs.LOG
DEFAULT_CONDITIONAL_LOSS_POE = 0.01

class ProbabilisticEventMixin(): # pylint: disable=W0232,W0201
    """Mixin for Probalistic Event Risk Job."""

    @general.preload
    @general.output
    def execute(self):
        """Execute the job."""

        aggregate_curve = prob.AggregateLossCurve()

        tasks = []
        for block_id in self.blocks_keys:
            LOGGER.debug("Starting task block, block_id = %s of %s"
                    % (block_id, len(self.blocks_keys)))

            tasks.append(general.compute_risk.delay(self.job_id, block_id))

        for task in tasks:
            try:
                # TODO (chris): Figure out where to put that timeout
                task.wait(timeout=None)

                aggregate_curve.append(task.result)
            except TimeoutError:
                # TODO(jmc): Cancel and respawn this task
                return

        curve = aggregate_curve.compute(self._tses(), self._time_span())
        aggregate_loss_curve.plot_aggregate_curve(self, curve)

    def _tses(self):
        """Return the time representative of the Stochastic Event Set
        specified for this job."""

# TODO (ac): Confirm this works regardless of the method of hazard calc
        histories = int(self["NUMBER_OF_SEISMICITY_HISTORIES"])
        realizations = int(self["NUMBER_OF_LOGIC_TREE_SAMPLES"])
        num_ses = histories * realizations

        return num_ses * self._time_span()

    def _time_span(self):
        """Return the time span specified for this job."""
        return float(self["INVESTIGATION_TIME"])

    def _gmf_db_list(self, job_id):  # pylint: disable=R0201
        """Returns a list of the output IDs of all computed GMFs"""

        ids = models.Output.objects.filter(
            oq_job=job_id, output_type='gmf').values_list('id', flat=True)

        return list(ids)

    def _get_db_gmf(self, gmf_id):
        """Returns a field for the given GMF"""
        grid = self.region.grid
        field = zeros((grid.rows, grid.columns))

        gmf_sites = models.GmfData.objects.filter(output=gmf_id)

        for gmf_site in gmf_sites:
            loc = gmf_site.location
            site = shapes.Site(loc.x, loc.y)
            grid_point = grid.point_at(site)

            field[grid_point.row][grid_point.column] = gmf_site.ground_motion

        return shapes.Field(field)

    def _sites_to_gmf_keys(self, sites):
        """Returns the GMF keys "row!col" for the given site list"""
        keys = []

        for site in sites:
            risk_point = self.region.grid.point_at(site)
            key = "%s!%s" % (risk_point.row, risk_point.column)
            keys.append(key)

        return keys

    def _get_db_gmfs(self, sites, job_id):
        """Aggregates GMF data from the DB by site"""
        all_gmfs = self._gmf_db_list(job_id)
        gmf_keys = self._sites_to_gmf_keys(sites)
        gmfs = dict((k, []) for k in gmf_keys)

        for gmf_id in all_gmfs:
            field = self._get_db_gmf(gmf_id)

            for key in gmfs.keys():
                (row, col) = key.split("!")
                gmfs[key].append(field.get(int(row), int(col)))

        return gmfs

    def _get_kvs_gmfs(self, sites, histories, realizations):
        """Aggregates GMF data from the KVS by site"""
        gmf_keys = self._sites_to_gmf_keys(sites)
        gmfs = dict((k, []) for k in gmf_keys)

        for i in range(0, histories):
            for j in range(0, realizations):
                key = kvs.tokens.stochastic_set_key(self.job_id, i, j)
                fieldset = shapes.FieldSet.from_json(kvs.get(key),
                    self.region.grid)

                for field in fieldset:
                    for key in gmfs.keys():
                        (row, col) = key.split("!")
                        gmfs[key].append(field.get(int(row), int(col)))

        return gmfs

    def slice_gmfs(self, block_id):
        """Load and collate GMF values for all sites in this block. """
        block = general.Block.from_kvs(block_id)
        gmfs = self._get_db_gmfs(block.sites, self.job_id)

        for key, gmf_slice in gmfs.items():
            (row, col) = key.split("!")
            key_gmf = kvs.tokens.gmf_set_key(self.job_id, col, row)
            LOGGER.debug("GMF_SLICE for %s X %s : \n\t%s" % (
                    col, row, gmf_slice))
            gmf = {"IMLs": gmf_slice, "TSES": self._tses(),
                    "TimeSpan": self._timespan()}
            kvs.set_value_json_encoded(key_gmf, gmf)

    def compute_risk(self, block_id, **kwargs):  # pylint: disable=W0613
        """Compute risk for a block of sites, that means:

        * loss ratio curves
        * loss curves
        * conditional losses
        * (partial) aggregate loss curve
        """

        self.slice_gmfs(block_id)

        self.vuln_curves = vulnerability.load_vuln_model_from_kvs(
            self.job_id)

        block = general.Block.from_kvs(block_id)

        # aggregate the losses for this block
        aggregate_curve = prob.AggregateLossCurve()

        for point in block.grid(self.region):
            key = kvs.tokens.gmf_set_key(self.job_id, point.column, point.row)
            gmf_slice = kvs.get_value_json_decoded(key)

            asset_key = kvs.tokens.asset_key(
                self.job_id, point.row, point.column)

            for asset in kvs.get_list_json_decoded(asset_key):
                LOGGER.debug("Processing asset %s" % (asset))

                # loss ratios, used both to produce the curve
                # and to aggregate the losses
                loss_ratios = self.compute_loss_ratios(asset, gmf_slice)

                loss_ratio_curve = self.compute_loss_ratio_curve(
                    point.column, point.row, asset, gmf_slice, loss_ratios)

                aggregate_curve.append(loss_ratios * asset["assetValue"])

                if loss_ratio_curve is not None:
                    loss_curve = self.compute_loss_curve(
                        point.column, point.row, loss_ratio_curve, asset)

                    for loss_poe in self._conditional_loss_poes():
                        self.compute_conditional_loss(point.column, point.row,
                                loss_curve, asset, loss_poe)

        return aggregate_curve.losses

    def _conditional_loss_poes(self):
        """Return the PoE(s) specified in the configuration file used to
        compute the conditional loss."""

        return [float(x) for x in self.params.get(
            "CONDITIONAL_LOSS_POE", "0.01").split()]

    def compute_loss_ratios(self, asset, gmf_slice):
        """For a given asset and ground motion field, computes
        the loss ratios used to obtain the related loss ratio curve
        and aggregate loss curve."""

        epsilon_provider = general.EpsilonProvider(self.params)

        vuln_function = self.vuln_curves.get(
            asset["vulnerabilityFunctionReference"], None)

        if not vuln_function:
            LOGGER.error(
                "Unknown vulnerability function %s for asset %s"
                % (asset["vulnerabilityFunctionReference"], asset["assetID"]))

            return None

        return prob.compute_loss_ratios(
            vuln_function, gmf_slice, epsilon_provider, asset)

    def compute_conditional_loss(self, col, row, loss_curve, asset, loss_poe):
        """Compute the conditional loss for a loss curve and Probability of
        Exceedance (PoE)."""

        loss_conditional = common.compute_conditional_loss(
            loss_curve, loss_poe)

        key = kvs.tokens.loss_key(
                self.job_id, row, col, asset["assetID"], loss_poe)

        LOGGER.debug("Conditional loss is %s, write to key %s" %
                (loss_conditional, key))

        kvs.set(key, loss_conditional)

    def compute_loss_ratio_curve(
            self, col, row, asset, gmf_slice, loss_ratios):
        """Compute the loss ratio curve for a single asset."""

        vuln_function = self.vuln_curves.get(
            asset["vulnerabilityFunctionReference"], None)

        if not vuln_function:
            LOGGER.error(
                "Unknown vulnerability function %s for asset %s"
                % (asset["vulnerabilityFunctionReference"], asset["assetID"]))

            return None

        epsilon_provider = general.EpsilonProvider(self.params)

        loss_ratio_curve = prob.compute_loss_ratio_curve(
                vuln_function, gmf_slice, epsilon_provider, asset,
                self._get_number_of_samples(), loss_ratios=loss_ratios)

        # NOTE (jmc): Early exit if the loss ratio is all zeros
        if not False in (loss_ratio_curve.ordinates == 0.0):
            return None

        key = kvs.tokens.loss_ratio_key(
            self.job_id, row, col, asset["assetID"])

        kvs.set(key, loss_ratio_curve.to_json())

        LOGGER.debug("Loss ratio curve is %s, write to key %s" %
                (loss_ratio_curve, key))

        return loss_ratio_curve

    def _get_number_of_samples(self):
        """Return the number of samples used to compute the loss ratio
        curve specified by the PROB_NUM_OF_SAMPLES parameter.

        Return None if the parameter is not specified, or empty or
        the value can't be casted to int.
        """

        number_of_samples = None
        raw_value = getattr(self, "PROB_NUM_OF_SAMPLES", None)

        if raw_value:
            try:
                number_of_samples = int(raw_value)
            except ValueError:
                LOGGER.error("PROB_NUM_OF_SAMPLES %s can't be converted "
                             "to int, using default value..." % raw_value)

        return number_of_samples

    def compute_loss_curve(self, column, row, loss_ratio_curve, asset):
        """Compute the loss curve for a single asset."""

        if asset is None:
            return None

        loss_curve = loss_ratio_curve.rescale_abscissae(asset["assetValue"])

        key = kvs.tokens.loss_curve_key(
            self.job_id, row, column, asset["assetID"])

        LOGGER.debug("Loss curve is %s, write to key %s" % (loss_curve, key))
        kvs.set(key, loss_curve.to_json())

        return loss_curve

general.RiskJobMixin.register("Event Based", ProbabilisticEventMixin)
<|endoftext|>


<commit_before>###############################################################################
# plot_distanceintegral_fewfields.py: make a useful plot of the distance
#                                     integration in the effective volume
#                                     for a few APOGEE-sized fields
###############################################################################
import sys
import os, os.path
import pickle
import numpy
from scipy import signal, integrate, interpolate
import healpy
import matplotlib
matplotlib.use('Agg')
from galpy.util import save_pickles, bovy_plot
import densprofiles
import dust
import gaia_rc
# magnitude limits for the survey
_GMIN= 3.
_GMAX= 20.
_NGSAMPLES= 10000
_NSIDE= 2048
_DEGTORAD= numpy.pi/180.
def plot_distanceintegral_fewfields(savename,plotname,apogee=False):
    #ls and bs of the patches
    ls= [280.,30.,12.,37.5,90.,127.5,30.,60.]
    bs= [45.,4.,2.,0.,0.,0.,0.,0.]
    radius= 1.49
    locs= [4214,4244,4378,4312,4242,4318,4240,4241] # for apogee
    # If we are doing APOGEE selection, load the APOGEE selection function
    if apogee:
        selectFile= '../savs/selfunc-nospdata.sav'
        if os.path.exists(selectFile):
            with open(selectFile,'rb') as savefile:
                apo= pickle.load(savefile)
    if os.path.exists(savename):
        with open(savename,'rb') as savefile:
            area= pickle.load(savefile)
            if apogee:
                area_cdens= pickle.load(savefile)
    else:
        if not apogee:
            # For samping over the absolute magnitude distribution
            iso= gaia_rc.load_iso()
            Gsamples= gaia_rc.sample_Gdist(iso,n=_NGSAMPLES)
        # Loop through each distance slice
        area= numpy.zeros((len(ls),len(dust._GREEN15DISTS)))
        if apogee:
            area_cdens= numpy.zeros((len(ls),len(dust._GREEN15DISTS)))
        for ii,dist in enumerate(dust._GREEN15DISTS):
            print "Working on distance %i: %f" % (ii,dist)
            G0= 0.68+dust.dist2distmod(dist)
            H0= -1.49+dust.dist2distmod(dist)
            # Load the dust map
            combinedmap= dust.load_combined(dist,nest=False,nside_out=_NSIDE)
            # Loop through a few patches
            theta, phi= healpy.pixelfunc.pix2ang(_NSIDE,
                                                 numpy.arange(healpy.pixelfunc.nside2npix(_NSIDE)),
                                                 nest=True)
            cosb= numpy.sin(theta)
            for ff,(ll,bb) in enumerate(zip(ls,bs)):
                vec= healpy.pixelfunc.ang2vec((90.-bb)*_DEGTORAD,ll*_DEGTORAD)
                ipixs= healpy.query_disc(_NSIDE,vec,radius*_DEGTORAD,
                                         inclusive=False,nest=True)
                ipixIndx= numpy.in1d(numpy.arange(healpy.pixelfunc.nside2npix(_NSIDE)),
                                     ipixs,assume_unique=True)
                # Calculate the density
                densmap= densprofiles.expdisk(phi[ipixIndx],
                                              numpy.pi/2.-theta[ipixIndx],
                                              dist*numpy.ones(len(ipixs)),
                                              glon=True,
                                              params=[1./3.,1./0.3])
                combinedmask= numpy.zeros(len(ipixs))
                tcombinedmap= combinedmap[ipixIndx]
                if apogee:
                    # Loop through the three cohorts
                    for c in ['short','medium','long']:
                        hmin= apo.Hmin(locs[ff],cohort=c)
                        hmax= apo.Hmax(locs[ff],cohort=c)
                        combinedmask+= ((tcombinedmap > (hmin-H0))\
                                            *(tcombinedmap < (hmax-H0))).astype('float')*apo(locs[ff],hmin+0.01)
                else:
                    for jj in range(_NGSAMPLES):
                        combinedmask+= ((tcombinedmap > (_GMIN-G0-Gsamples[jj]+0.68))\
                                            *(tcombinedmap < (_GMAX-G0-Gsamples[jj]+0.68))).astype('float')
                combinedmask/= _NGSAMPLES
                # Compute cross correlation
                area[ff,ii]= float(numpy.sum(cosb[ipixIndx]*densmap\
                                                 *combinedmask))
                if apogee:
                    # Also store the approximation that the density is constant
                    cdens= densprofiles.expdisk(ll*_DEGTORAD,bb*_DEGTORAD,dist,
                                                glon=True,
                                                params=[1./3.,1./0.3])
                    area_cdens[ff,ii]= float(numpy.sum(cosb[ipixIndx]*cdens\
                                                           *combinedmask))
        save_pickles(savename,area)
    # Plot the power spectrum
    if True:
        bovy_plot.bovy_print(fig_height=3.)
        matplotlib.rcParams['text.latex.preamble']=[r"\usepackage{yfonts}"]
        for ff in range(len(ls)):
            psdx, psd= signal.periodogram(area[ff]*dust._GREEN15DISTS**3./numpy.sum(area[ff]*dust._GREEN15DISTS**3.),
                                          fs=1./(dust._GREEN15DISTMODS[1]-dust._GREEN15DISTMODS[0]),
                                          detrend=lambda x: x,scaling='spectrum')
            bovy_plot.bovy_plot(psdx[1:],psd[1:],
                                '-',loglog=True,
                                xlabel=r'$2\pi\,k_\mu\,(\mathrm{mag}^{-1})$',
                                ylabel=r'$P_k$',
                                xrange=[0.04,4.],
                                overplot=ff > 0)
            if ff == 0:
                bovy_plot.bovy_text(r'$\mathrm{normalized}\ D^3\,\nu_*(\mu|\theta)\,\textswab{S}(\mu)$',
                                    bottom_left=True,size=16.)
        bovy_plot.bovy_end_print(plotname)
    else:
        bovy_plot.bovy_print(fig_height=3.)
        matplotlib.rcParams['text.latex.preamble']=[r"\usepackage{yfonts}"]
        for ff in range(len(ls)):
            bovy_plot.bovy_plot(dust._GREEN15DISTMODS,
                                area[ff]*dust._GREEN15DISTS**3.,
                                'k-',
                                xlabel=r'$\mu\,(\mathrm{mag}^{-1})$',
                                ylabel=r'$D^3\,\nu_*(\mu|\theta)\,\textswab{S}(\mu)$')
        bovy_plot.bovy_end_print(plotname)
    for ff in range(len(ls)):
        spl= interpolate.InterpolatedUnivariateSpline(dust._GREEN15DISTMODS,
                                                      area[ff]*dust._GREEN15DISTS**3.,
                                                  k=5)
        fthder= [spl.derivatives(dm)[4] for dm in dust._GREEN15DISTMODS]
        print "Simpson error (%f,%f)= %g" % (ls[ff],bs[ff],0.5**4./180.*numpy.mean(numpy.fabs(fthder))/integrate.simps(area[ff]*dust._GREEN15DISTS**3.,dx=0.5))
    return None

if __name__ == '__main__':
    plot_distanceintegral_fewfields(sys.argv[1], # savefilename
                                    sys.argv[2], # plotfilename
                                    len(sys.argv) > 3) # APOGEE selection?
<commit_msg>apogee-maps: also print volume<commit_after>###############################################################################
# plot_distanceintegral_fewfields.py: make a useful plot of the distance
#                                     integration in the effective volume
#                                     for a few APOGEE-sized fields
###############################################################################
import sys
import os, os.path
import pickle
import numpy
from scipy import signal, integrate, interpolate
import healpy
import matplotlib
matplotlib.use('Agg')
from galpy.util import save_pickles, bovy_plot
import densprofiles
import dust
import gaia_rc
# magnitude limits for the survey
_GMIN= 3.
_GMAX= 20.
_NGSAMPLES= 10000
_NSIDE= 2048
_DEGTORAD= numpy.pi/180.
def plot_distanceintegral_fewfields(savename,plotname,apogee=False):
    #ls and bs of the patches
    ls= [280.,30.,12.,37.5,90.,127.5,30.,60.]
    bs= [45.,4.,2.,0.,0.,0.,0.,0.]
    radius= 1.49
    locs= [4214,4244,4378,4312,4242,4318,4240,4241] # for apogee
    # If we are doing APOGEE selection, load the APOGEE selection function
    if apogee:
        selectFile= '../savs/selfunc-nospdata.sav'
        if os.path.exists(selectFile):
            with open(selectFile,'rb') as savefile:
                apo= pickle.load(savefile)
    if os.path.exists(savename):
        with open(savename,'rb') as savefile:
            area= pickle.load(savefile)
            if apogee:
                area_cdens= pickle.load(savefile)
    else:
        if not apogee:
            # For samping over the absolute magnitude distribution
            iso= gaia_rc.load_iso()
            Gsamples= gaia_rc.sample_Gdist(iso,n=_NGSAMPLES)
        # Loop through each distance slice
        area= numpy.zeros((len(ls),len(dust._GREEN15DISTS)))
        if apogee:
            area_cdens= numpy.zeros((len(ls),len(dust._GREEN15DISTS)))
        for ii,dist in enumerate(dust._GREEN15DISTS):
            print "Working on distance %i: %f" % (ii,dist)
            G0= 0.68+dust.dist2distmod(dist)
            H0= -1.49+dust.dist2distmod(dist)
            # Load the dust map
            combinedmap= dust.load_combined(dist,nest=False,nside_out=_NSIDE)
            # Loop through a few patches
            theta, phi= healpy.pixelfunc.pix2ang(_NSIDE,
                                                 numpy.arange(healpy.pixelfunc.nside2npix(_NSIDE)),
                                                 nest=True)
            cosb= numpy.sin(theta)
            for ff,(ll,bb) in enumerate(zip(ls,bs)):
                vec= healpy.pixelfunc.ang2vec((90.-bb)*_DEGTORAD,ll*_DEGTORAD)
                ipixs= healpy.query_disc(_NSIDE,vec,radius*_DEGTORAD,
                                         inclusive=False,nest=True)
                ipixIndx= numpy.in1d(numpy.arange(healpy.pixelfunc.nside2npix(_NSIDE)),
                                     ipixs,assume_unique=True)
                # Calculate the density
                densmap= densprofiles.expdisk(phi[ipixIndx],
                                              numpy.pi/2.-theta[ipixIndx],
                                              dist*numpy.ones(len(ipixs)),
                                              glon=True,
                                              params=[1./3.,1./0.3])
                combinedmask= numpy.zeros(len(ipixs))
                tcombinedmap= combinedmap[ipixIndx]
                if apogee:
                    # Loop through the three cohorts
                    for c in ['short','medium','long']:
                        hmin= apo.Hmin(locs[ff],cohort=c)
                        hmax= apo.Hmax(locs[ff],cohort=c)
                        combinedmask+= ((tcombinedmap > (hmin-H0))\
                                            *(tcombinedmap < (hmax-H0))).astype('float')*apo(locs[ff],hmin+0.01)
                else:
                    for jj in range(_NGSAMPLES):
                        combinedmask+= ((tcombinedmap > (_GMIN-G0-Gsamples[jj]+0.68))\
                                            *(tcombinedmap < (_GMAX-G0-Gsamples[jj]+0.68))).astype('float')
                combinedmask/= _NGSAMPLES
                # Compute cross correlation
                area[ff,ii]= float(numpy.sum(cosb[ipixIndx]*densmap\
                                                 *combinedmask))
                if apogee:
                    # Also store the approximation that the density is constant
                    cdens= densprofiles.expdisk(ll*_DEGTORAD,bb*_DEGTORAD,dist,
                                                glon=True,
                                                params=[1./3.,1./0.3])
                    area_cdens[ff,ii]= float(numpy.sum(cosb[ipixIndx]*cdens\
                                                           *combinedmask))
        save_pickles(savename,area)
    # Plot the power spectrum
    if True:
        bovy_plot.bovy_print(fig_height=3.)
        matplotlib.rcParams['text.latex.preamble']=[r"\usepackage{yfonts}"]
        for ff in range(len(ls)):
            psdx, psd= signal.periodogram(area[ff]*dust._GREEN15DISTS**3./numpy.sum(area[ff]*dust._GREEN15DISTS**3.),
                                          fs=1./(dust._GREEN15DISTMODS[1]-dust._GREEN15DISTMODS[0]),
                                          detrend=lambda x: x,scaling='spectrum')
            bovy_plot.bovy_plot(psdx[1:],psd[1:],
                                '-',loglog=True,
                                xlabel=r'$2\pi\,k_\mu\,(\mathrm{mag}^{-1})$',
                                ylabel=r'$P_k$',
                                xrange=[0.04,4.],
                                overplot=ff > 0)
            if ff == 0:
                bovy_plot.bovy_text(r'$\mathrm{normalized}\ D^3\,\nu_*(\mu|\theta)\,\textswab{S}(\mu)$',
                                    bottom_left=True,size=16.)
        bovy_plot.bovy_end_print(plotname)
    else:
        bovy_plot.bovy_print(fig_height=3.)
        matplotlib.rcParams['text.latex.preamble']=[r"\usepackage{yfonts}"]
        for ff in range(len(ls)):
            bovy_plot.bovy_plot(dust._GREEN15DISTMODS,
                                area[ff]*dust._GREEN15DISTS**3.,
                                'k-',
                                xlabel=r'$\mu\,(\mathrm{mag}^{-1})$',
                                ylabel=r'$D^3\,\nu_*(\mu|\theta)\,\textswab{S}(\mu)$')
        bovy_plot.bovy_end_print(plotname)
    for ff in range(len(ls)):
        spl= interpolate.InterpolatedUnivariateSpline(dust._GREEN15DISTMODS,
                                                      area[ff]*dust._GREEN15DISTS**3.,
                                                  k=5)
        fthder= [spl.derivatives(dm)[4] for dm in dust._GREEN15DISTMODS]
        print "Simpson error (%f,%f)= %g, volume = %g" % (ls[ff],bs[ff],0.5**4./180.*numpy.mean(numpy.fabs(fthder))/integrate.simps(area[ff]*dust._GREEN15DISTS**3.,dx=0.5),numpy.sum(area[ff]*dust._GREEN15DISTS**3.))
    return None

if __name__ == '__main__':
    plot_distanceintegral_fewfields(sys.argv[1], # savefilename
                                    sys.argv[2], # plotfilename
                                    len(sys.argv) > 3) # APOGEE selection?
<|endoftext|>


<commit_before>#!/usr/bin/python
#-*- coding: iso-8859-15 -*-
# SADR METEOLLSKY
# http://www.sadr.fr
# SEBASTIEN LECLERC 2017
# Inspired by :
# fehlfarbe
# http://indilib.org/forum/general/606-take-image-with-python-script.html
# allsky picture script

import sys, time, logging, os
import PyIndi
import pyfits
import cv2
import cStringIO
from PIL import Image, ImageDraw, ImageFont
from datetime import datetime
import numpy as np
from meteollskyconfig import *

if len(sys.argv) <> 3:
    print "Usage: python allsky_frame.py <<time exposure>> <<gain>>"
    sys.exit()

fi=open(CHARTPATH+"luminosity.js","r")
luminosity = fi.readline()
arg = luminosity.split('=')
lumValue = arg[1]
print("Luminosity: " + lumValue)
fi.close()
if float(lumValue) == 0:
    #EXP_TIME = 19. - (19. / 100 * float(lumValue))
    EXP_TIME = 45
    EXP_GAIN = 30
    COLOR = 0
else:
    EXP_TIME = 0.000001
    EXP_GAIN = 1
    COLOR = 1

EXP_TIME = str(EXP_TIME)
EXP_GAIN = str(EXP_GAIN)
USB_TRAFFIC = 0
print("EXP_TIME: " + str(EXP_TIME))
print("EXP_GAIN: " + str(EXP_GAIN))
#EXP_TIME = sys.argv[1]
#EXP_GAIN = str(50)
#EXP_GAIN = sys.argv[2]

class IndiClient(PyIndi.BaseClient):

    device = None
    roi = None

    def __init__(self):
        super(IndiClient, self).__init__()
        self.logger = logging.getLogger('PyQtIndi.IndiClient')
        self.logger.info('creating an instance of PyQtIndi.IndiClient')
    def newDevice(self, d):
        #self.logger.info("new device " + d.getDeviceName())
        if d.getDeviceName() == INDIDEVICE:
            self.logger.info("Set new device %s!" % INDIDEVICE)
            # save reference to the device in member variable
            self.device = d
    def newProperty(self, p):
        #self.logger.info("new property "+ p.getName() + " for device "+ p.getDeviceName())
        if self.device is not None and p.getName() == "CONNECTION" and p.getDeviceName() == self.device.getDeviceName():
            self.logger.info("Got property CONNECTION for %s!" % INDIDEVICE)
            # connect to device
            self.connectDevice(self.device.getDeviceName())
            # set BLOB mode to BLOB_ALSO
            self.setBLOBMode(1, self.device.getDeviceName(), None)
        if p.getName() == "CCD_EXPOSURE":
            # take first exposure
            self.takeExposure()
        if p.getName() == "CCD_GAIN":
            gain = self.device.getNumber("CCD_GAIN")
            gain[0].value = float(EXP_GAIN)
            self.sendNewNumber(gain)
        if p.getName() == "USB_TRAFFIC":
            traffic = self.device.getNumber("USB_TRAFFIC")
            print ("USB Traffic"+str(USB_TRAFFIC)+".")
            traffic[0].value = float(USB_TRAFFIC)
            self.sendNewNumber(traffic)
    def removeProperty(self, p):
        #self.logger.info("remove property "+ p.getName() + " for device "+ p.getDeviceName())
        pass
    def newBLOB(self, bp):
        print "new Blob arrived"
        self.logger.info("new BLOB "+ bp.name.decode())
        # get image data
        img = bp.getblobdata()
        # write image data to StringIO buffer
        blobfile = cStringIO.StringIO(img)
        hdulist = pyfits.open(blobfile)
        scidata = hdulist[0].data
        moyenne = scidata.mean()
        if self.roi is not None:
            scidata = scidata[self.roi[1]:self.roi[1]+self.roi[3], self.roi[0]:self.roi[0]+self.roi[2]]
        hdulist[0].data = scidata
        # process bayer pattern and grey image
        processedImage = cv2.cvtColor(scidata, cv2.COLOR_BAYER_GR2RGB)
        if COLOR == 0:
            processedImage = cv2.cvtColor(processedImage, cv2.COLOR_BGR2GRAY)
        cv2.imwrite(CHARTPATH+"allsky.jpg" , processedImage)

        ######### WATERMARK #############
        # Open the watermark image
        main = Image.open(CHARTPATH+"allsky.jpg")
        # Adding an alpha channel to the main image
        main = main.convert("RGBA")

        # Open the watermark image
        overlay = Image.open(WATERMARK)
        # Adding an alpha channel to the main image
        overlay = overlay.convert("RGBA")

        # Get an ImageDraw object so we can draw on the image
        waterdraw = ImageDraw.ImageDraw(overlay, "RGBA")
        # get a font
        fnt = ImageFont.truetype('font/Typography Times Bold.ttf', 15)
        # Place the text at (10, 10) in the upper left corner. Text will be white.
        waterdraw.text((5, 5), "%s" % datetime.now(), font=fnt, fill=(255,255,255,255))
        waterdraw.text((5, 25), "GAIN "+EXP_GAIN+" / EXP "+EXP_TIME+"s", font=fnt, fill=(255,255,255,255))
        if COLOR == 0:
            waterdraw.text((5, 45), "NO COLOR / LUM "+lumValue, font=fnt, fill=(255,255,255,255))
        else:
            waterdraw.text((5, 45), "COLOR / LUM "+lumValue, font=fnt, fill=(255,255,255,255))
        waterdraw.text((5, 65), "MEAN " + str(moyenne), font=fnt, fill=(255,255,255,255))

        # Make a final composite image
        result = Image.alpha_composite(main, overlay)
        ######### END OF WATERMARK #############

        # Paste the watermark (with alpha layer) onto the original image and save it
        result.save(CHARTPATH+"allsky.jpg", "JPEG")

        # start new exposure for timelapse images!
        time.sleep(10)
        self.takeExposure()

    def newSwitch(self, svp):
      self.logger.info ("new Switch "+ svp.name.decode() + " for device "+ svp.device.decode())
    def newNumber(self, nvp):
        self.logger.info("new Number "+ nvp.name.decode() + " for device "+ nvp.device.decode())
    def newText(self, tvp):
        self.logger.info("new Text "+ tvp.name.decode() + " for device "+ tvp.device.decode())
    def newLight(self, lvp):
        self.logger.info("new Light "+ lvp.name.decode() + " for device "+ lvp.device.decode())
    def newMessage(self, d, m):
        #self.logger.info("new Message "+ d.messageQueue(m).decode())
        pass
    def serverConnected(self):
        print("Server connected ("+self.getHost()+":"+str(self.getPort())+")")
        self.connected = True
    def serverDisconnected(self, code):
        self.logger.info("Server disconnected (exit code = "+str(code)+","+str(self.getHost())+":"+str(self.getPort())+")")
        # set connected to False
        self.connected = False
    def takeExposure(self):
        self.logger.info("<<<<<<<< Exposure >>>>>>>>>")
        #get current exposure time
        exp = self.device.getNumber("CCD_EXPOSURE")
        # set exposure time to 5 seconds
        print ("Exposition time"+EXP_TIME+" s.")
        exp[0].value = float(EXP_TIME)
        # send new exposure time to server/device
        self.sendNewNumber(exp)

logging.basicConfig(format='%(asctime)s %(message)s', level=logging.INFO)

############# MAIN #############

if __name__ == '__main__':
    # instantiate the client
    indiclient=IndiClient()
    # region of interest for my allsky cam
    indiclient.roi = (206, 30, 890, 820)
    # set indi server localhost and port 7624
    INDIPORT=int(INDIPORT)
    indiclient.setServer(INDISERVER, INDIPORT)
    # connect to indi server
    print("Connecting and waiting 2secs")
    if (not(indiclient.connectServer())):
         print("No indiserver running on "+indiclient.getHost()+":"+str(indiclient.getPort())+" - Try to run")
         sys.exit(1)
    #time.sleep(30)

    # start endless loop, client works asynchron in background, loop stops after disconnect
    #while indiclient.connected:
    while True:
        #indiclient.takeExposure()
        fi=open(CHARTPATH+"luminosity.js","r")
        luminosity = fi.readline()
        arg = luminosity.split('=')
        lumValue = arg[1]
        print("Luminosity: " + lumValue)
        fi.close()
        if mean > 128:
            if EXP_GAIN == 1:
                EXP_TIME /= 2
                if EXP_TIME < 0.000001:
                    EXP_TIME = 0.000001
                    COLOR =1
            else :
                EXP_GAIN /= 2
                if EXP_GAIN < 1:
                    EXP_GAIN = 1
                    COLOR = 1
        else:
            if EXP_GAIN == 30:
                EXP_TIME *= 2
                if EXP_TIME > 45:
                    EXP_TIME = 45
                    COLOR = 0
            else :
                EXP_GAIN *= 2
                if EXP_GAIN > 30:
                    EXP_GAIN = 30
                    COLOR = 1

        #if float(lumValue) == 0:
            #EXP_TIME = 19. - (19. / 100 * float(lumValue))
            #EXP_TIME = 45
            #EXP_GAIN = 30
            #COLOR = 0
        #else:
            #EXP_TIME = 0.000001
            #EXP_GAIN = 1
            #COLOR = 1

        EXP_TIME = str(EXP_TIME)
        EXP_GAIN = str(EXP_GAIN)
        print("EXP_TIME: " + str(EXP_TIME))
        print("EXP_GAIN: " + str(EXP_GAIN))
        time.sleep(15)
<commit_msg>fix variable mistake<commit_after>#!/usr/bin/python
#-*- coding: iso-8859-15 -*-
# SADR METEOLLSKY
# http://www.sadr.fr
# SEBASTIEN LECLERC 2017
# Inspired by :
# fehlfarbe
# http://indilib.org/forum/general/606-take-image-with-python-script.html
# allsky picture script

import sys, time, logging, os
import PyIndi
import pyfits
import cv2
import cStringIO
from PIL import Image, ImageDraw, ImageFont
from datetime import datetime
import numpy as np
from meteollskyconfig import *

if len(sys.argv) <> 3:
    print "Usage: python allsky_frame.py <<time exposure>> <<gain>>"
    sys.exit()

fi=open(CHARTPATH+"luminosity.js","r")
luminosity = fi.readline()
arg = luminosity.split('=')
lumValue = arg[1]
print("Luminosity: " + lumValue)
fi.close()
if float(lumValue) == 0:
    #EXP_TIME = 19. - (19. / 100 * float(lumValue))
    EXP_TIME = 45
    EXP_GAIN = 30
    COLOR = 0
else:
    EXP_TIME = 0.000001
    EXP_GAIN = 1
    COLOR = 1

EXP_TIME = str(EXP_TIME)
EXP_GAIN = str(EXP_GAIN)
USB_TRAFFIC = 0
print("EXP_TIME: " + str(EXP_TIME))
print("EXP_GAIN: " + str(EXP_GAIN))
#EXP_TIME = sys.argv[1]
#EXP_GAIN = str(50)
#EXP_GAIN = sys.argv[2]

class IndiClient(PyIndi.BaseClient):

    device = None
    roi = None

    def __init__(self):
        super(IndiClient, self).__init__()
        self.logger = logging.getLogger('PyQtIndi.IndiClient')
        self.logger.info('creating an instance of PyQtIndi.IndiClient')
    def newDevice(self, d):
        #self.logger.info("new device " + d.getDeviceName())
        if d.getDeviceName() == INDIDEVICE:
            self.logger.info("Set new device %s!" % INDIDEVICE)
            # save reference to the device in member variable
            self.device = d
    def newProperty(self, p):
        #self.logger.info("new property "+ p.getName() + " for device "+ p.getDeviceName())
        if self.device is not None and p.getName() == "CONNECTION" and p.getDeviceName() == self.device.getDeviceName():
            self.logger.info("Got property CONNECTION for %s!" % INDIDEVICE)
            # connect to device
            self.connectDevice(self.device.getDeviceName())
            # set BLOB mode to BLOB_ALSO
            self.setBLOBMode(1, self.device.getDeviceName(), None)
        if p.getName() == "CCD_EXPOSURE":
            # take first exposure
            self.takeExposure()
        if p.getName() == "CCD_GAIN":
            gain = self.device.getNumber("CCD_GAIN")
            gain[0].value = float(EXP_GAIN)
            self.sendNewNumber(gain)
        if p.getName() == "USB_TRAFFIC":
            traffic = self.device.getNumber("USB_TRAFFIC")
            print ("USB Traffic"+str(USB_TRAFFIC)+".")
            traffic[0].value = float(USB_TRAFFIC)
            self.sendNewNumber(traffic)
    def removeProperty(self, p):
        #self.logger.info("remove property "+ p.getName() + " for device "+ p.getDeviceName())
        pass
    def newBLOB(self, bp):
        print "new Blob arrived"
        self.logger.info("new BLOB "+ bp.name.decode())
        # get image data
        img = bp.getblobdata()
        # write image data to StringIO buffer
        blobfile = cStringIO.StringIO(img)
        hdulist = pyfits.open(blobfile)
        scidata = hdulist[0].data
        moyenne = scidata.mean()
        if self.roi is not None:
            scidata = scidata[self.roi[1]:self.roi[1]+self.roi[3], self.roi[0]:self.roi[0]+self.roi[2]]
        hdulist[0].data = scidata
        # process bayer pattern and grey image
        processedImage = cv2.cvtColor(scidata, cv2.COLOR_BAYER_GR2RGB)
        if COLOR == 0:
            processedImage = cv2.cvtColor(processedImage, cv2.COLOR_BGR2GRAY)
        cv2.imwrite(CHARTPATH+"allsky.jpg" , processedImage)

        ######### WATERMARK #############
        # Open the watermark image
        main = Image.open(CHARTPATH+"allsky.jpg")
        # Adding an alpha channel to the main image
        main = main.convert("RGBA")

        # Open the watermark image
        overlay = Image.open(WATERMARK)
        # Adding an alpha channel to the main image
        overlay = overlay.convert("RGBA")

        # Get an ImageDraw object so we can draw on the image
        waterdraw = ImageDraw.ImageDraw(overlay, "RGBA")
        # get a font
        fnt = ImageFont.truetype('font/Typography Times Bold.ttf', 15)
        # Place the text at (10, 10) in the upper left corner. Text will be white.
        waterdraw.text((5, 5), "%s" % datetime.now(), font=fnt, fill=(255,255,255,255))
        waterdraw.text((5, 25), "GAIN "+EXP_GAIN+" / EXP "+EXP_TIME+"s", font=fnt, fill=(255,255,255,255))
        if COLOR == 0:
            waterdraw.text((5, 45), "NO COLOR / LUM "+lumValue, font=fnt, fill=(255,255,255,255))
        else:
            waterdraw.text((5, 45), "COLOR / LUM "+lumValue, font=fnt, fill=(255,255,255,255))
        waterdraw.text((5, 65), "MEAN " + str(moyenne), font=fnt, fill=(255,255,255,255))

        # Make a final composite image
        result = Image.alpha_composite(main, overlay)
        ######### END OF WATERMARK #############

        # Paste the watermark (with alpha layer) onto the original image and save it
        result.save(CHARTPATH+"allsky.jpg", "JPEG")

        # start new exposure for timelapse images!
        time.sleep(10)
        self.takeExposure()

    def newSwitch(self, svp):
      self.logger.info ("new Switch "+ svp.name.decode() + " for device "+ svp.device.decode())
    def newNumber(self, nvp):
        self.logger.info("new Number "+ nvp.name.decode() + " for device "+ nvp.device.decode())
    def newText(self, tvp):
        self.logger.info("new Text "+ tvp.name.decode() + " for device "+ tvp.device.decode())
    def newLight(self, lvp):
        self.logger.info("new Light "+ lvp.name.decode() + " for device "+ lvp.device.decode())
    def newMessage(self, d, m):
        #self.logger.info("new Message "+ d.messageQueue(m).decode())
        pass
    def serverConnected(self):
        print("Server connected ("+self.getHost()+":"+str(self.getPort())+")")
        self.connected = True
    def serverDisconnected(self, code):
        self.logger.info("Server disconnected (exit code = "+str(code)+","+str(self.getHost())+":"+str(self.getPort())+")")
        # set connected to False
        self.connected = False
    def takeExposure(self):
        self.logger.info("<<<<<<<< Exposure >>>>>>>>>")
        #get current exposure time
        exp = self.device.getNumber("CCD_EXPOSURE")
        # set exposure time to 5 seconds
        print ("Exposition time"+EXP_TIME+" s.")
        exp[0].value = float(EXP_TIME)
        # send new exposure time to server/device
        self.sendNewNumber(exp)

logging.basicConfig(format='%(asctime)s %(message)s', level=logging.INFO)

############# MAIN #############

if __name__ == '__main__':
    # instantiate the client
    indiclient=IndiClient()
    # region of interest for my allsky cam
    indiclient.roi = (206, 30, 890, 820)
    # set indi server localhost and port 7624
    INDIPORT=int(INDIPORT)
    indiclient.setServer(INDISERVER, INDIPORT)
    # connect to indi server
    print("Connecting and waiting 2secs")
    if (not(indiclient.connectServer())):
         print("No indiserver running on "+indiclient.getHost()+":"+str(indiclient.getPort())+" - Try to run")
         sys.exit(1)
    #time.sleep(30)

    # start endless loop, client works asynchron in background, loop stops after disconnect
    #while indiclient.connected:
    while True:
        #indiclient.takeExposure()
        fi=open(CHARTPATH+"luminosity.js","r")
        luminosity = fi.readline()
        arg = luminosity.split('=')
        lumValue = arg[1]
        print("Luminosity: " + lumValue)
        fi.close()
        if moyenne > 128:
            if EXP_GAIN == 1:
                EXP_TIME /= 2
                if EXP_TIME < 0.000001:
                    EXP_TIME = 0.000001
                    COLOR =1
            else :
                EXP_GAIN /= 2
                if EXP_GAIN < 1:
                    EXP_GAIN = 1
                    COLOR = 1
        else:
            if EXP_GAIN == 30:
                EXP_TIME *= 2
                if EXP_TIME > 45:
                    EXP_TIME = 45
                    COLOR = 0
            else :
                EXP_GAIN *= 2
                if EXP_GAIN > 30:
                    EXP_GAIN = 30
                    COLOR = 1

        #if float(lumValue) == 0:
            #EXP_TIME = 19. - (19. / 100 * float(lumValue))
            #EXP_TIME = 45
            #EXP_GAIN = 30
            #COLOR = 0
        #else:
            #EXP_TIME = 0.000001
            #EXP_GAIN = 1
            #COLOR = 1

        EXP_TIME = str(EXP_TIME)
        EXP_GAIN = str(EXP_GAIN)
        print("EXP_TIME: " + str(EXP_TIME))
        print("EXP_GAIN: " + str(EXP_GAIN))
        time.sleep(15)
<|endoftext|>


<commit_before># Copyright 2008 Google Inc. All Rights Reserved.

"""Fake urllib HTTP connection objects."""

__author__ = 'Mark Roach (mrroach@google.com)'

from pprint import pformat
import six
from six import BytesIO
from six.moves import urllib

class Error(Exception):
    """Base exception type for this module."""

def initialize():
    """Install TestHandler as the only active handler for http requests."""
    opener = urllib.request.build_opener(TestHandler)
    urllib.request.install_opener(opener)

def create_response_key(method, url, request_headers):
    """Create the response key for a request.

    Args:
        method: The http method (e.g. 'get', 'put', etc.)
        url: The path being requested including site.
        request_headers: Dictionary of headers passed along with the request.
    Returns:
        The key as a string.
    """
    parsed = urllib.parse.urlsplit(url)
    qs = urllib.parse.parse_qs(parsed.query)
    query = urllib.parse.urlencode([(k, qs[k]) for k in sorted(qs.keys())])
    return str((
        method,
        urllib.parse.urlunsplit((
            parsed.scheme, parsed.netloc, parsed.path, query, parsed.fragment)),
        dictionary_to_canonical_str(request_headers)))

def dictionary_to_canonical_str(dictionary):
    """Create canonical string from a dictionary.

    Args:
        dictionary: The dictionary to convert.
    Returns:
        A string of the dictionary in canonical form.
    """
    return str([(k.capitalize(), dictionary[k]) for k in sorted(
        dictionary.keys())])

class TestHandler(urllib.request.HTTPHandler, urllib.request.HTTPSHandler):
    """A urllib handler object which returns a predefined response."""

    _response = None
    _response_map = {}
    request = None
    site = ''

    @classmethod
    def set_response(cls, response):
        """Set a static response to be returned for all requests.

        Args:
            response: A FakeResponse object to be returned.
        """
        cls._response_map = {}
        cls._response = response

    @classmethod
    def respond_to(cls, method, path,
                   request_headers, body, code=200, response_headers=None):
        """Build a response object to be used for a specific request.

        Args:
            method: The http method (e.g. 'get', 'put' etc.)
            path: The path being requested (e.g. '/collection/id.json')
            request_headers: Dictionary of headers passed along with the request
            body: The string that should be returned for a matching request
            code: The http response code to return
            response_headers: Dictionary of headers to return
        Returns:
            None
        """
        key = create_response_key(method, urllib.parse.urljoin(cls.site, path),
                                  request_headers)
        value = (code, body, response_headers)
        cls._response_map[str(key)] = value

    def do_open(self, http_class, request, **http_conn_args):
        """Return the response object for the given request.

        Overrides the HTTPHandler method of the same name to return a
        FakeResponse instead of creating any network connections.

        Args:
            http_class: The http protocol being used.
            request: A urllib.request.Request object.
        Returns:
            A FakeResponse object.
        """
        self.__class__.request = request  # Store the most recent request object
        if self._response_map:
            key = create_response_key(
                request.get_method(), request.get_full_url(), request.headers)
            if str(key) in self._response_map:
                (code, body, response_headers) = self._response_map[str(key)]
                return FakeResponse(code, body, response_headers)
            else:
                raise Error('Unknown request %s %s'
                            '\nrequest:%s\nresponse_map:%s' % (
                            request.get_method(), request.get_full_url(),
                            str(key), pformat(list(self._response_map.keys()))))
        elif isinstance(self._response, Exception):
            raise(self._response)
        else:
            return self._response

class FakeResponse(object):
    """A fake HTTPResponse object for testing."""

    def __init__(self, code, body, headers=None):
        self.code = code
        self.msg = str(code)
        if headers is None:
            headers = {}
        self.headers = headers
        self.info = lambda: self.headers
        if isinstance(body, six.text_type):
            body = body.encode('utf-8')
        self.body_file = BytesIO(body)

    def read(self):
        """Read the entire response body."""
        return self.body_file.read()

    def readline(self):
        """Read a single line from the response body."""
        return self.body_file.readline()

    def close(self):
        """Close the connection."""
        pass
<commit_msg>Add a _context instance method to the TestHandler in order to properly mirror HTTPSHandler behaviour.<commit_after># Copyright 2008 Google Inc. All Rights Reserved.

"""Fake urllib HTTP connection objects."""

__author__ = 'Mark Roach (mrroach@google.com)'

from pprint import pformat
import six
from six import BytesIO
from six.moves import urllib

class Error(Exception):
    """Base exception type for this module."""

def initialize():
    """Install TestHandler as the only active handler for http requests."""
    opener = urllib.request.build_opener(TestHandler)
    urllib.request.install_opener(opener)

def create_response_key(method, url, request_headers):
    """Create the response key for a request.

    Args:
        method: The http method (e.g. 'get', 'put', etc.)
        url: The path being requested including site.
        request_headers: Dictionary of headers passed along with the request.
    Returns:
        The key as a string.
    """
    parsed = urllib.parse.urlsplit(url)
    qs = urllib.parse.parse_qs(parsed.query)
    query = urllib.parse.urlencode([(k, qs[k]) for k in sorted(qs.keys())])
    return str((
        method,
        urllib.parse.urlunsplit((
            parsed.scheme, parsed.netloc, parsed.path, query, parsed.fragment)),
        dictionary_to_canonical_str(request_headers)))

def dictionary_to_canonical_str(dictionary):
    """Create canonical string from a dictionary.

    Args:
        dictionary: The dictionary to convert.
    Returns:
        A string of the dictionary in canonical form.
    """
    return str([(k.capitalize(), dictionary[k]) for k in sorted(
        dictionary.keys())])

class TestHandler(urllib.request.HTTPHandler, urllib.request.HTTPSHandler):
    """A urllib handler object which returns a predefined response."""

    _response = None
    _response_map = {}
    request = None
    site = ''

    def __init__(self, debuglevel=0, context=None):
        self._debuglevel = debuglevel
        self._context = context

    @classmethod
    def set_response(cls, response):
        """Set a static response to be returned for all requests.

        Args:
            response: A FakeResponse object to be returned.
        """
        cls._response_map = {}
        cls._response = response

    @classmethod
    def respond_to(cls, method, path,
                   request_headers, body, code=200, response_headers=None):
        """Build a response object to be used for a specific request.

        Args:
            method: The http method (e.g. 'get', 'put' etc.)
            path: The path being requested (e.g. '/collection/id.json')
            request_headers: Dictionary of headers passed along with the request
            body: The string that should be returned for a matching request
            code: The http response code to return
            response_headers: Dictionary of headers to return
        Returns:
            None
        """
        key = create_response_key(method, urllib.parse.urljoin(cls.site, path),
                                  request_headers)
        value = (code, body, response_headers)
        cls._response_map[str(key)] = value

    def do_open(self, http_class, request, **http_conn_args):
        """Return the response object for the given request.

        Overrides the HTTPHandler method of the same name to return a
        FakeResponse instead of creating any network connections.

        Args:
            http_class: The http protocol being used.
            request: A urllib.request.Request object.
        Returns:
            A FakeResponse object.
        """
        self.__class__.request = request  # Store the most recent request object
        if self._response_map:
            key = create_response_key(
                request.get_method(), request.get_full_url(), request.headers)
            if str(key) in self._response_map:
                (code, body, response_headers) = self._response_map[str(key)]
                return FakeResponse(code, body, response_headers)
            else:
                raise Error('Unknown request %s %s'
                            '\nrequest:%s\nresponse_map:%s' % (
                            request.get_method(), request.get_full_url(),
                            str(key), pformat(list(self._response_map.keys()))))
        elif isinstance(self._response, Exception):
            raise(self._response)
        else:
            return self._response

class FakeResponse(object):
    """A fake HTTPResponse object for testing."""

    def __init__(self, code, body, headers=None):
        self.code = code
        self.msg = str(code)
        if headers is None:
            headers = {}
        self.headers = headers
        self.info = lambda: self.headers
        if isinstance(body, six.text_type):
            body = body.encode('utf-8')
        self.body_file = BytesIO(body)

    def read(self):
        """Read the entire response body."""
        return self.body_file.read()

    def readline(self):
        """Read a single line from the response body."""
        return self.body_file.readline()

    def close(self):
        """Close the connection."""
        pass
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
import json
import time
from django.conf import settings
from django.http import StreamingHttpResponse
from django.utils.decorators import method_decorator
from django.views.decorators.csrf import csrf_exempt

from opps.views.generic.detail import DetailView
from opps.views.generic.list import ListView
from opps.views.generic.json_views import JSONPResponse
from opps.db import Db

from .models import Notification

class AsyncServer(DetailView):
    model = Notification

    def _queue(self):
        _db = Db(
            self.get_object().container.get_absolute_url(),
            self.get_object().id)
        pubsub = _db.object().pubsub()
        pubsub.subscribe(_db.key)

        while True:
            for m in pubsub.listen():
                if m['type'] == 'message':
                    yield u"data: {}\n\n".format(m['data'])
            yield u"data: {}\n\n".format(json.dumps({"action": "ping"}))
            time.sleep(0.5)

    @method_decorator(csrf_exempt)
    def dispatch(self, request, *args, **kwargs):
        response = StreamingHttpResponse(self._queue(),
                                         mimetype='text/event-stream')
        response['Cache-Control'] = 'no-cache'
        response['Software'] = 'opps-liveblogging'
        response['Access-Control-Allow-Origin'] = '*'
        response.flush()
        return response

class LongPullingServer(ListView, JSONPResponse):
    model = Notification

    def get_template_names(self):
        templates = []
        domain_folder = self.get_template_folder()

        if not self.long_slug:
            templates.append('{}/none.json'.format(domain_folder))
            return templates

        list_name = 'list'

        if self.template_name_suffix:
            list_name = "{}{}".format(list_name, self.template_name_suffix)

        if self.channel:
            # Check layout, change via admin
            if self.channel.layout != u'default':
                list_name = self.channel.layout

            if self.channel.group and self.channel.parent:
                templates.append('{}/{}/{}.json'.format(
                    domain_folder, self.channel.parent.long_slug, list_name))

                if self.request.GET.get('page') and\
                   self.__class__.__name__ not in\
                   settings.OPPS_PAGINATE_NOT_APP:
                    templates.append('{}/{}/{}_paginated.json'.format(
                        domain_folder, self.channel.parent.long_slug,
                        list_name))

            if self.request.GET.get('page') and\
               self.__class__.__name__ not in settings.OPPS_PAGINATE_NOT_APP:
                templates.append('{}/{}/{}_paginated.json'.format(
                    domain_folder, self.channel.long_slug, list_name))

            templates.append('{}/{}/{}.json'.format(
                domain_folder, self.channel.long_slug, list_name))

            for t in self.channel.get_ancestors()[::-1]:
                templates.append('{}/{}/{}.json'.format(
                    domain_folder, t.long_slug, list_name))
                if self.request.GET.get('page') and\
                   self.__class__.__name__ not in\
                   settings.OPPS_PAGINATE_NOT_APP:
                    templates.append('{}/{}/{}_paginated.json'.format(
                        domain_folder, t.long_slug, list_name))

        if self.request.GET.get('page') and\
           self.__class__.__name__ not in settings.OPPS_PAGINATE_NOT_APP:
            templates.append('{}/{}_paginated.json'.format(domain_folder,
                                                           list_name))

        templates.append('{}/{}.json'.format(domain_folder, list_name))
        return templates

    def get_queryset(self):
        query = super(LongPullingServer, self).get_queryset()
        old_id = self.request.GET.get('old_id', 0)
        return query.filter(id__gte=old_id)._clone()
<commit_msg>Fix bug, if not exist post return error on sse action<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
import json
import time
from django.conf import settings
from django.http import StreamingHttpResponse
from django.utils.decorators import method_decorator
from django.views.decorators.csrf import csrf_exempt

from opps.views.generic.list import ListView
from opps.views.generic.json_views import JSONPResponse
from opps.db import Db

from .models import Notification

class AsyncServer(ListView):
    model = Notification

    def _db(self, obj):
        _db = Db(
            obj.container.get_absolute_url(),
            obj.container.id)
        pubsub = _db.object().pubsub()
        pubsub.subscribe(_db.key)
        return pubsub

    def _queue(self):
        try:
            obj = self.get_queryset()[0]
        except:
            obj = False

        if not obj:
            while True:
                yield u"data: {}\n\n".format(
                    json.dumps({"action": "error"}))
                time.sleep(10)
        else:
            while True:
                pubsub = self._db(obj)
                for m in pubsub.listen():
                    if m['type'] == 'message':
                        yield u"data: {}\n\n".format(m['data'])
                yield u"data: {}\n\n".format(
                    json.dumps({"action": "ping"}))
                time.sleep(0.5)

    @method_decorator(csrf_exempt)
    def dispatch(self, request, *args, **kwargs):
        response = StreamingHttpResponse(self._queue(),
                                         mimetype='text/event-stream')
        response['Cache-Control'] = 'no-cache'
        response['Software'] = 'opps-liveblogging'
        response['Access-Control-Allow-Origin'] = '*'
        response.flush()
        return response

class LongPullingServer(ListView, JSONPResponse):
    model = Notification

    def get_template_names(self):
        templates = []
        domain_folder = self.get_template_folder()

        if not self.long_slug:
            templates.append('{}/none.json'.format(domain_folder))
            return templates

        list_name = 'list'

        if self.template_name_suffix:
            list_name = "{}{}".format(list_name, self.template_name_suffix)

        if self.channel:
            # Check layout, change via admin
            if self.channel.layout != u'default':
                list_name = self.channel.layout

            if self.channel.group and self.channel.parent:
                templates.append('{}/{}/{}.json'.format(
                    domain_folder, self.channel.parent.long_slug, list_name))

                if self.request.GET.get('page') and\
                   self.__class__.__name__ not in\
                   settings.OPPS_PAGINATE_NOT_APP:
                    templates.append('{}/{}/{}_paginated.json'.format(
                        domain_folder, self.channel.parent.long_slug,
                        list_name))

            if self.request.GET.get('page') and\
               self.__class__.__name__ not in settings.OPPS_PAGINATE_NOT_APP:
                templates.append('{}/{}/{}_paginated.json'.format(
                    domain_folder, self.channel.long_slug, list_name))

            templates.append('{}/{}/{}.json'.format(
                domain_folder, self.channel.long_slug, list_name))

            for t in self.channel.get_ancestors()[::-1]:
                templates.append('{}/{}/{}.json'.format(
                    domain_folder, t.long_slug, list_name))
                if self.request.GET.get('page') and\
                   self.__class__.__name__ not in\
                   settings.OPPS_PAGINATE_NOT_APP:
                    templates.append('{}/{}/{}_paginated.json'.format(
                        domain_folder, t.long_slug, list_name))

        if self.request.GET.get('page') and\
           self.__class__.__name__ not in settings.OPPS_PAGINATE_NOT_APP:
            templates.append('{}/{}_paginated.json'.format(domain_folder,
                                                           list_name))

        templates.append('{}/{}.json'.format(domain_folder, list_name))
        return templates

    def get_queryset(self):
        query = super(LongPullingServer, self).get_queryset()
        old_id = self.request.GET.get('old_id', 0)
        return query.filter(id__gte=old_id)._clone()
<|endoftext|>