<commit_before>import warnings

import numpy

from simphony.testing.utils import dummy_cuba_value

class CUBADataAccumulator(object):
    """ Accumulate data information per CUBA key.

    A collector object that stores :class:``DataContainer`` data into
    a list of values per CUBA key. By appending :class:`DataContainer`
    instanced the user can effectively convert the per item mapping of
    data values in a CUDS container to a per CUBA key mapping of the
    data values (useful for coping data to vtk array containers).

    The Accumulator has two modes of operation ``fixed`` and
    ``expand``. ``fixed`` means that data will be stored for
    a predefined set of keys on every ``append`` call and missing
    values will be saved as ``None``. Where ``expand`` will extend
    the internal table of values whenever a new key is introduced.

    .. rubric:: expand operation

    >>> accumulator = CUBADataAccumulator():
    >>> accumulator.append(DataContainer(TEMPERATURE=34))
    >>> accumulator.keys()
    {CUBA.TEMPERATURE}
    >>> accumulator.append(DataContainer(VELOCITY=(0.1, 0.1, 0.1))
    >>> accumulator.append(DataContainer(TEMPERATURE=56))
    >>> accumulator.keys()
    {CUBA.TEMPERATURE, CUBA.VELOCITY}
    >>> accumulator[CUBA.TEMPERATURE]
    [34, None, 56]
    >>> accumulator[CUBA.VELOCITY]
    [None, (0.1, 0.1, 0.1), None]

    .. rubric:: fixed operation

    >>> accumulator = CUBADataAccumulator([CUBA.TEMPERATURE, CUBA.PRESSURE]):
    >>> accumulator.keys()
    {CUBA.TEMPERATURE, CUBA.PRESSURE}
    >>> accumulator.append(DataContainer(TEMPERATURE=34))
    >>> accumulator.append(DataContainer(VELOCITY=(0.1, 0.1, 0.1))
    >>> accumulator.append(DataContainer(TEMPERATURE=56))
    >>> accumulator.keys()
    {CUBA.TEMPERATURE, CUBA.PRESSURE}
    >>> accumulator[CUBA.TEMPERATURE]
    [34, None, 56]
    >>> accumulator[CUBA.PRESSURE]
    [None, None, None]
    >>> accumulator[CUBA.VELOCITY]
    KeyError(...)

    """
    def __init__(self, keys=()):
        """Constructor

        Parameters
        ----------
        keys : list

            The list of keys that the accumulator should care
            about. Providing this value at initialisation sets up the
            accumulator to operate in ``fixed`` mode. If no keys are
            provided then accumulator operates in ``expand`` mode.

        """
        self._keys = set(keys)
        self._expand_mode = len(keys) == 0
        self._data = {}
        self._record_size = 0
        self._expand(self._keys)

    @property
    def keys(self):
        """ The set of CUBA keys that this accumulator contains.

        """
        return set(self._keys)

    def append(self, data):
        """ Append info from a ``DataContainer``.

        Parameters
        ----------
        data : DataContainer
            The data information to append.

        If the accumulator operates in ``fixed`` mode:

        - Any keys in :code:`self.keys()` that have values in ``data``
          will be stored (appended to the related key lits).
        - Missing keys will be stored as ``None``

        If the accumulator operates in ``expand`` mode:

        - Any new keys in `Data` will be added to the :code:`self.keys()` list
          and the related list of values with length equal to the current
          record size will be initialised with values of ``None``.
        - Any keys in the modified :code:`self.keys()` that have values in
          ``data`` will be stored (appended to the list of the related key).
        - Missing keys will be store as ``None``.

        """
        if self._expand_mode:
            new_keys = set(data.keys()) - self._keys
            self._keys.update(new_keys)
            self._expand(new_keys)
        for key in self._keys:
            self._data[key].append(data.get(key, None))
        self._record_size += 1

    def load_onto_vtk(self, vtk_data):
        """ Load the stored information onto a vtk data container.

        Parameters
        ----------
        vtk_data : vtkPointData or vtkCellData
            The vtk container to load the value onto.

        Data are loaded onto the vtk container based on their data
        type. The name of the added array is the name of the CUBA key
        (i.e. :samp:`{CUBA}.name`). Currently only scalars and three
        dimensional vectors are supported.

        """
        def replacer(data):
            return nan if data is None else data

        for cuba in self.keys:
            default = dummy_cuba_value(cuba)
            if isinstance(default, (float, int)):
                data = numpy.array(self._data[cuba], dtype=float)
                index = vtk_data.add_array(data)
                vtk_data.get_array(index).name = cuba.name
            elif isinstance(default, numpy.ndarray) and len(default) == 3:
                nan = numpy.array([None, None, None], dtype=float)
                data = numpy.array(
                    tuple(replacer(data) for data in self._data[cuba]),
                    dtype=numpy.float)
                index = vtk_data.add_array(data)
                vtk_data.get_array(index).name = cuba.name
            else:
                message = 'property {!r} is currently ignored'
                warnings.warn(message.format(cuba))

    def __len__(self):
        """ The number of values that are stored per key

        .. note:: Behaviour is temporary and will probably change soon.

        """
        return self._record_size

    def __getitem__(self, key):
        """ Get the list of accumulated values for the CUBA key.

        Parameters
        ----------
        key : CUBA
            A CUBA Enum value

        Returns
        -------
        result : list
            A list of data values collected for ``key``. Missing values
            are designated with ``None``.

        """
        return self._data[key]

    def _expand(self, keys):
        for key in keys:
            self._data[key] = [None] * self._record_size

def gather_cells(
        iterable, vtk_mapping, point2index, counter, accumulator):
    """ Gather the vtk cell information from an element iterator.

    Arguments
    ---------
    iterable :
        The Element iterable object

    vtk_mapping : dict
        The mapping from points number to tvtk.Cell type.

    point2index: dict
        The mapping from points uid to the index of the vtk points array.

    counter : itertools.count
        The counter object to use when evaluating the ``elements2index``
        mapping.

    accumulator : CUBADataAccumulator
        The accumulator instance to use and collect the data information

    Returns
    -------
    cells : list
         The cell point information encoded in a one dimensional list.

    cells_size : list
         The list of points number per cell.

    cells_types : list
         The list of cell types in sequence.

    element2index : dict
         The mapping from element uid to iteration index.

    """
    cells = []
    cells_size = []
    cell_types = []
    element2index = {}

    for element in iter(iterable):
        element2index[element.uid] = counter.next()
        npoints = len(element.points)
        cells_size.append(npoints + 1)
        cells.append(npoints)
        cells.extend(point2index[uid] for uid in element.points)
        cell_types.append(vtk_mapping[npoints])
        accumulator.append(element.data)

    return cells, cells_size, cell_types, element2index
<commit_msg>delete trailing whitespace<commit_after>import warnings

import numpy

from simphony.testing.utils import dummy_cuba_value

class CUBADataAccumulator(object):
    """ Accumulate data information per CUBA key.

    A collector object that stores :class:``DataContainer`` data into
    a list of values per CUBA key. By appending :class:`DataContainer`
    instanced the user can effectively convert the per item mapping of
    data values in a CUDS container to a per CUBA key mapping of the
    data values (useful for coping data to vtk array containers).

    The Accumulator has two modes of operation ``fixed`` and
    ``expand``. ``fixed`` means that data will be stored for
    a predefined set of keys on every ``append`` call and missing
    values will be saved as ``None``. Where ``expand`` will extend
    the internal table of values whenever a new key is introduced.

    .. rubric:: expand operation

    >>> accumulator = CUBADataAccumulator():
    >>> accumulator.append(DataContainer(TEMPERATURE=34))
    >>> accumulator.keys()
    {CUBA.TEMPERATURE}
    >>> accumulator.append(DataContainer(VELOCITY=(0.1, 0.1, 0.1))
    >>> accumulator.append(DataContainer(TEMPERATURE=56))
    >>> accumulator.keys()
    {CUBA.TEMPERATURE, CUBA.VELOCITY}
    >>> accumulator[CUBA.TEMPERATURE]
    [34, None, 56]
    >>> accumulator[CUBA.VELOCITY]
    [None, (0.1, 0.1, 0.1), None]

    .. rubric:: fixed operation

    >>> accumulator = CUBADataAccumulator([CUBA.TEMPERATURE, CUBA.PRESSURE]):
    >>> accumulator.keys()
    {CUBA.TEMPERATURE, CUBA.PRESSURE}
    >>> accumulator.append(DataContainer(TEMPERATURE=34))
    >>> accumulator.append(DataContainer(VELOCITY=(0.1, 0.1, 0.1))
    >>> accumulator.append(DataContainer(TEMPERATURE=56))
    >>> accumulator.keys()
    {CUBA.TEMPERATURE, CUBA.PRESSURE}
    >>> accumulator[CUBA.TEMPERATURE]
    [34, None, 56]
    >>> accumulator[CUBA.PRESSURE]
    [None, None, None]
    >>> accumulator[CUBA.VELOCITY]
    KeyError(...)

    """
    def __init__(self, keys=()):
        """Constructor

        Parameters
        ----------
        keys : list

            The list of keys that the accumulator should care
            about. Providing this value at initialisation sets up the
            accumulator to operate in ``fixed`` mode. If no keys are
            provided then accumulator operates in ``expand`` mode.

        """
        self._keys = set(keys)
        self._expand_mode = len(keys) == 0
        self._data = {}
        self._record_size = 0
        self._expand(self._keys)

    @property
    def keys(self):
        """ The set of CUBA keys that this accumulator contains.

        """
        return set(self._keys)

    def append(self, data):
        """ Append info from a ``DataContainer``.

        Parameters
        ----------
        data : DataContainer
            The data information to append.

        If the accumulator operates in ``fixed`` mode:

        - Any keys in :code:`self.keys()` that have values in ``data``
          will be stored (appended to the related key lits).
        - Missing keys will be stored as ``None``

        If the accumulator operates in ``expand`` mode:

        - Any new keys in `Data` will be added to the :code:`self.keys()` list
          and the related list of values with length equal to the current
          record size will be initialised with values of ``None``.
        - Any keys in the modified :code:`self.keys()` that have values in
          ``data`` will be stored (appended to the list of the related key).
        - Missing keys will be store as ``None``.

        """
        if self._expand_mode:
            new_keys = set(data.keys()) - self._keys
            self._keys.update(new_keys)
            self._expand(new_keys)
        for key in self._keys:
            self._data[key].append(data.get(key, None))
        self._record_size += 1

    def load_onto_vtk(self, vtk_data):
        """ Load the stored information onto a vtk data container.

        Parameters
        ----------
        vtk_data : vtkPointData or vtkCellData
            The vtk container to load the value onto.

        Data are loaded onto the vtk container based on their data
        type. The name of the added array is the name of the CUBA key
        (i.e. :samp:`{CUBA}.name`). Currently only scalars and three
        dimensional vectors are supported.

        """
        def replacer(data):
            return nan if data is None else data

        for cuba in self.keys:
            default = dummy_cuba_value(cuba)
            if isinstance(default, (float, int)):
                data = numpy.array(self._data[cuba], dtype=float)
                index = vtk_data.add_array(data)
                vtk_data.get_array(index).name = cuba.name
            elif isinstance(default, numpy.ndarray) and len(default) == 3:
                nan = numpy.array([None, None, None], dtype=float)
                data = numpy.array(
                    tuple(replacer(data) for data in self._data[cuba]),
                    dtype=numpy.float)
                index = vtk_data.add_array(data)
                vtk_data.get_array(index).name = cuba.name
            else:
                message = 'property {!r} is currently ignored'
                warnings.warn(message.format(cuba))

    def __len__(self):
        """ The number of values that are stored per key

        .. note:: Behaviour is temporary and will probably change soon.

        """
        return self._record_size

    def __getitem__(self, key):
        """ Get the list of accumulated values for the CUBA key.

        Parameters
        ----------
        key : CUBA
            A CUBA Enum value

        Returns
        -------
        result : list
            A list of data values collected for ``key``. Missing values
            are designated with ``None``.

        """
        return self._data[key]

    def _expand(self, keys):
        for key in keys:
            self._data[key] = [None] * self._record_size

def gather_cells(
        iterable, vtk_mapping, point2index, counter, accumulator):
    """ Gather the vtk cell information from an element iterator.

    Arguments
    ---------
    iterable :
        The Element iterable object

    vtk_mapping : dict
        The mapping from points number to tvtk.Cell type.

    point2index: dict
        The mapping from points uid to the index of the vtk points array.

    counter : itertools.count
        The counter object to use when evaluating the ``elements2index``
        mapping.

    accumulator : CUBADataAccumulator
        The accumulator instance to use and collect the data information

    Returns
    -------
    cells : list
         The cell point information encoded in a one dimensional list.

    cells_size : list
         The list of points number per cell.

    cells_types : list
         The list of cell types in sequence.

    element2index : dict
         The mapping from element uid to iteration index.

    """
    cells = []
    cells_size = []
    cell_types = []
    element2index = {}

    for element in iter(iterable):
        element2index[element.uid] = counter.next()
        npoints = len(element.points)
        cells_size.append(npoints + 1)
        cells.append(npoints)
        cells.extend(point2index[uid] for uid in element.points)
        cell_types.append(vtk_mapping[npoints])
        accumulator.append(element.data)

    return cells, cells_size, cell_types, element2index
<|endoftext|>


<commit_before># Copyright 2013 IBM Corp.

import re
from eventlet import greenthread

import powervc.common.config as cfg
from powervc.common.gettextutils import _
from nova.compute import flavors
from nova import exception
from nova import db
from oslo_log import log as logging
from nova.openstack.common import loopingcall
from powervc.nova.driver.compute import constants

LOG = logging.getLogger(__name__)

CONF = cfg.CONF

def periodic_flavor_sync(ctx, driver, scg_id_list):
        """
        Periodically update the flavors from PowerVC.
        A default time of 300 seconds is specified for the refresh interval.
        if the refresh interval is set to 0, then flavors are not refreshed.
        """
        sync_interval = CONF.powervc.flavor_sync_interval

        if sync_interval is None or sync_interval == 0:
            return

        def flavors_sync(driver, scg_id_list):
            FlavorSync(driver, scg_id_list).synchronize_flavors(ctx)
            LOG.debug('Flavors synchronization completed')

        sync_flavors = loopingcall.FixedIntervalLoopingCall(flavors_sync,
                                                            driver,
                                                            scg_id_list)
        sync_flavors.start(interval=sync_interval, initial_delay=sync_interval)

class FlavorSync():
    """A class that synchorizes the flavors.
    The functionality provided here is called by the manager.
    The driver provided interfaces to the PowerVC.
    """

    def __init__(self, driver, scg_id_list):
        self.driver = driver
        self.prefix = CONF.powervc.flavor_prefix
        self.scg_id_list = scg_id_list

    def synchronize_flavors(self, ctx):
        """
        Get a list of all public flavors from PowerVC.
        If it is in configuration white list,
            and not in black list, insert it.
         if it is already in local tables, ignore it.
        """
        LOG.info(_("Flavors synchronization starts."))
        # Get all public flavors. By default, detail and public is set.
        pvcFlavors = self.driver.list_flavors()
        # Sync flavors in list
        for flavor in pvcFlavors:
            LOG.info(_("Flavor:%s") % str(flavor))
            greenthread.sleep(0)
            # This check is added to eliminate sync of private flavors
            # Can be removed once PowerVC fixes to return only public flavors
            # by default.
            if not(flavor.__dict__.get(constants.IS_PUBLIC)):
                continue

            if (self._check_for_sync(flavor.name)):
                response = self._check_for_extraspecs(flavor)
                if response is not None:
                    self._sync_flavor(ctx, flavor, response[1])
        LOG.info(_("Flavors synchronization ends."))

    def _sanitize(self, opts_list):
        """
        Remove any whitespace only list values
        """
        for opt in opts_list:
            if len(opt.strip()) == 0:
                opts_list.remove(opt)
        return opts_list

    def get_flavors_white_list(self):
        """
        Get the flavors to sync from the powervc conf file
        """
        return self._sanitize(CONF.powervc.flavor_white_list)

    def get_flavors_black_list(self):
        """
        Get the black listed flavors from the powervc conf file
        """
        return self._sanitize(CONF.powervc.flavor_black_list)

    def _check_for_sync(self, fl_name):
        """ Check the white/black lists to determine if sync candidate """
        fl_sync = True
        # Get the list of flavors names to sync.
        fl_wlist = self.get_flavors_white_list()
        fl_blist = self.get_flavors_black_list()

        if (len(fl_wlist) != 0):
            fl_sync = self._regex_comp(fl_name, fl_wlist)
        if (fl_sync and (len(fl_blist) != 0)):
            fl_sync = not(self._regex_comp(fl_name, fl_blist))
        return fl_sync

    def _regex_comp(self, name, flist):
        """
        Make a regex comparison for name in the list
        Return a boolean True if found in the list
        """
        if name in flist:
            return True
        for item in flist:
            p = re.compile(item)
            match = p.match(name)
            if (match is not None):
                return True
        return False

    def _sync_flavor(self, ctx, flavor, extra_specs):
        """
        Insert the flavor with extra specs if not in local database
        """
        flavor_in_local_db = None
        flavor_name = self.prefix + flavor.name
        try:
            flavor_in_local_db = db.flavor_get_by_name(ctx, flavor_name)
        except exception.FlavorNotFoundByName:
            self._insert_pvc_flavor_extraspecs(ctx, flavor, extra_specs)

        # Update the extra_speces of the flavor
        if flavor_in_local_db is not None:
            flavor_id = flavor_in_local_db.get('flavorid', '')
            if (flavor_id is not ''
                    and extra_specs):
                self._update_flavor_extraspecs(ctx,
                                               flavor_id,
                                               extra_specs)

    def _check_for_extraspecs(self, flavor):
        """
        Check for valid extraspecs defined and to be synced.
        The method returns the following values:
        (True, None) - flavor to be synced, and no extra specs defined.
        (True, extraspecs) - flavor to be synced with the extra specs defined.
        None - scg connectivity group defined in extraspecs is not supported,
               and flavor not to be synced.

        Checking for scg to be removed when powervc driver supports multiple
        scgs
        """
        flavor_extraspecs = self.driver.get_flavor_extraspecs(flavor)
        if flavor_extraspecs:
            scg_key = constants.SCG_KEY
            if scg_key in flavor_extraspecs:
                if not self.scg_id_list:
                    return None
                if not flavor_extraspecs[scg_key] in self.scg_id_list:
                    return None
        return (True, flavor_extraspecs)

    def _insert_pvc_flavor_extraspecs(self, context, flavor, extra_specs):
        """ Insert the flavor and extra specs if any """
        flavor_created = self._create_flavor(context, flavor)
        if extra_specs:
            self._update_flavor_extraspecs(context,
                                           flavor_created.get('flavorid'),
                                           extra_specs)

    def _update_flavor_extraspecs(self, context, flavorid, flavor_extraspecs):
        """ Insert the flavor extra specs """
        db.flavor_extra_specs_update_or_create(context,
                                               flavorid,
                                               flavor_extraspecs)

    def _create_flavor(self, context, flavor):
        """ Create and insert the flavor """
        flavor_dict = flavor.__dict__
        name = self.prefix + flavor.name
        flavorid = self.prefix + flavor.id
        memory = flavor.ram
        vcpus = flavor.vcpus
        root_gb = flavor.disk
        ephemeral_gb = flavor_dict.get('OS-FLV-EXT-DATA:ephemeral', 0)
        u_swap = flavor_dict.get('swap', 0)
        rxtx_factor = flavor_dict.get('rxtx_factor', 1.0)
        is_public = flavor_dict.get('os-flavor-access:is_public', True)
        if u_swap == "":
            swap = 0
        else:
            swap = int(u_swap)

        try:
            return flavors.create(name, memory, vcpus, root_gb,
                                  ephemeral_gb=ephemeral_gb,
                                  flavorid=flavorid, swap=swap,
                                  rxtx_factor=rxtx_factor,
                                  is_public=is_public)
        except exception.InstanceExists as err:
            raise err
<commit_msg>Fix flavorsync issue during nova-powervc startup<commit_after># Copyright 2013, 2019 IBM Corp.

import re
from eventlet import greenthread

import powervc.common.config as cfg
from powervc.common.gettextutils import _
from nova.compute import flavors
from nova import exception
from nova import db
from oslo_log import log as logging
from nova.openstack.common import loopingcall
from powervc.nova.driver.compute import constants

LOG = logging.getLogger(__name__)

CONF = cfg.CONF

def periodic_flavor_sync(ctx, driver, scg_id_list):
        """
        Periodically update the flavors from PowerVC.
        A default time of 300 seconds is specified for the refresh interval.
        if the refresh interval is set to 0, then flavors are not refreshed.
        """
        sync_interval = CONF.powervc.flavor_sync_interval

        if sync_interval is None or sync_interval == 0:
            return

        def flavors_sync(driver, scg_id_list):
            FlavorSync(driver, scg_id_list).synchronize_flavors(ctx)
            LOG.debug('Flavors synchronization completed')

        sync_flavors = loopingcall.FixedIntervalLoopingCall(flavors_sync,
                                                            driver,
                                                            scg_id_list)
        sync_flavors.start(interval=sync_interval, initial_delay=sync_interval)

class FlavorSync():
    """A class that synchorizes the flavors.
    The functionality provided here is called by the manager.
    The driver provided interfaces to the PowerVC.
    """

    def __init__(self, driver, scg_id_list):
        self.driver = driver
        self.prefix = CONF.powervc.flavor_prefix
        self.scg_id_list = scg_id_list

    def synchronize_flavors(self, ctx):
        """
        Get a list of all public flavors from PowerVC.
        If it is in configuration white list,
            and not in black list, insert it.
         if it is already in local tables, ignore it.
        """
        LOG.info(_("Flavors synchronization starts."))
        # Get all public flavors. By default, detail and public is set.
        pvcFlavors = self.driver.list_flavors()
        # Sync flavors in list
        for flavor in pvcFlavors:
            LOG.info(_("Flavor:%s") % str(flavor))
            greenthread.sleep(0)
            # This check is added to eliminate sync of private flavors
            # Can be removed once PowerVC fixes to return only public flavors
            # by default.
            if not(flavor.__dict__.get(constants.IS_PUBLIC)):
                continue

            if (self._check_for_sync(flavor.name)):
                response = self._check_for_extraspecs(flavor)
                if response is not None:
                    self._sync_flavor(ctx, flavor, response[1])
        LOG.info(_("Flavors synchronization ends."))

    def _sanitize(self, opts_list):
        """
        Remove any whitespace only list values
        """
        for opt in opts_list:
            if len(opt.strip()) == 0:
                opts_list.remove(opt)
        return opts_list

    def get_flavors_white_list(self):
        """
        Get the flavors to sync from the powervc conf file
        """
        return self._sanitize(CONF.powervc.flavor_white_list)

    def get_flavors_black_list(self):
        """
        Get the black listed flavors from the powervc conf file
        """
        return self._sanitize(CONF.powervc.flavor_black_list)

    def _check_for_sync(self, fl_name):
        """ Check the white/black lists to determine if sync candidate """
        fl_sync = True
        # Get the list of flavors names to sync.
        fl_wlist = self.get_flavors_white_list()
        fl_blist = self.get_flavors_black_list()

        if (len(fl_wlist) != 0):
            fl_sync = self._regex_comp(fl_name, fl_wlist)
        if (fl_sync and (len(fl_blist) != 0)):
            fl_sync = not(self._regex_comp(fl_name, fl_blist))
        return fl_sync

    def _regex_comp(self, name, flist):
        """
        Make a regex comparison for name in the list
        Return a boolean True if found in the list
        """
        if name in flist:
            return True
        for item in flist:
            p = re.compile(item)
            match = p.match(name)
            if (match is not None):
                return True
        return False

    def _sync_flavor(self, ctx, flavor, extra_specs):
        """
        Insert the flavor with extra specs if not in local database
        """
        flavor_in_local_db = None
        flavor_name = self.prefix + flavor.name
        try:
            flavor_in_local_db = db.flavor_get_by_name(ctx, flavor_name)
        except exception.FlavorNotFoundByName:
            self._insert_pvc_flavor_extraspecs(ctx, flavor, extra_specs)

        # Update the extra_speces of the flavor
        if flavor_in_local_db is not None:
            flavor_id = flavor_in_local_db.get('flavorid', '')
            if (flavor_id is not ''
                    and extra_specs):
                self._update_flavor_extraspecs(ctx,
                                               flavor_id,
                                               extra_specs)

    def _check_for_extraspecs(self, flavor):
        """
        Check for valid extraspecs defined and to be synced.
        The method returns the following values:
        (True, None) - flavor to be synced, and no extra specs defined.
        (True, extraspecs) - flavor to be synced with the extra specs defined.
        None - scg connectivity group defined in extraspecs is not supported,
               and flavor not to be synced.

        Checking for scg to be removed when powervc driver supports multiple
        scgs
        """
        flavor_extraspecs = self.driver.get_flavor_extraspecs(flavor)
        if flavor_extraspecs:
            scg_key = constants.SCG_KEY
            if scg_key in flavor_extraspecs:
                if not self.scg_id_list:
                    return None
                if not flavor_extraspecs[scg_key] in self.scg_id_list:
                    return None
        return (True, flavor_extraspecs)

    def _insert_pvc_flavor_extraspecs(self, context, flavor, extra_specs):
        """ Insert the flavor and extra specs if any """
        flavor_created = self._create_flavor(context, flavor)
        if flavor_created and extra_specs:
            self._update_flavor_extraspecs(context,
                                           flavor_created.get('flavorid'),
                                           extra_specs)

    def _update_flavor_extraspecs(self, context, flavorid, flavor_extraspecs):
        """ Insert the flavor extra specs """
        db.flavor_extra_specs_update_or_create(context,
                                               flavorid,
                                               flavor_extraspecs)

    def _create_flavor(self, context, flavor):
        """ Create and insert the flavor """
        flavor_dict = flavor.__dict__
        name = self.prefix + flavor.name
        flavorid = self.prefix + flavor.id
        memory = flavor.ram
        vcpus = flavor.vcpus
        root_gb = flavor.disk
        ephemeral_gb = flavor_dict.get('OS-FLV-EXT-DATA:ephemeral', 0)
        u_swap = flavor_dict.get('swap', 0)
        rxtx_factor = flavor_dict.get('rxtx_factor', 1.0)
        is_public = flavor_dict.get('os-flavor-access:is_public', True)
        if u_swap == "":
            swap = 0
        else:
            swap = int(u_swap)

        try:
            return flavors.create(name, memory, vcpus, root_gb,
                                  ephemeral_gb=ephemeral_gb,
                                  flavorid=flavorid, swap=swap,
                                  rxtx_factor=rxtx_factor,
                                  is_public=is_public)
        except Exception as exc:
            LOG.error(_("Unable to sync flavor "
                        + str(name) + ". " + str(exc.format_message())))
            return None
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# pylint: enable=W0511,W0142,I0011,E1101,E0611,F0401,E1103,R0801,W0232

# Copyright (c) 2010-2013, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""
Post processing functionality for the classical PSHA hazard calculator.
E.g. mean and quantile curves.
"""

import math
import numpy

from celery.task.sets import TaskSet
from django.db import transaction
from itertools import izip

import openquake.engine

from openquake.engine import logs
from openquake.engine.calculators.hazard.general import CURVE_CACHE_SIZE
from openquake.engine.db import models
from openquake.engine.utils import config
from openquake.engine.utils import tasks as utils_tasks
from openquake.engine.utils.general import block_splitter
from openquake.engine.writer import BulkInserter

# Number of locations considered by each task
DEFAULT_LOCATIONS_PER_TASK = 1000

def compute_hazard_maps(curves, imls, poes):
    """
    Given a set of hazard curve poes, interpolate a hazard map at the specified
    ``poe``.

    :param curves:
        2D array of floats. Each row represents a curve, where the values
        in the row are the PoEs (Probabilities of Exceedance) corresponding to
        ``imls``. Each curve corresponds to a geographical location.
    :param imls:
        Intensity Measure Levels associated with these hazard ``curves``. Type
        should be an array-like of floats.
    :param float poes:
        Value(s) on which to interpolate a hazard map from the input
        ``curves``. Can be an array-like or scalar value (for a single PoE).

    :returns:
        A 2D numpy array of hazard map data. Each element/row in the resulting
        array represents the interpolated map for each ``poes`` value
        specified. If ``poes`` is just a single scalar value, the result array
        will have a length of 1.

        The results are structure this way so it is easy to iterate over the
        hazard map results, and in a consistent way (no matter how many
        ``poes`` values are specified).
    """
    poes = numpy.array(poes)

    if len(poes.shape) == 0:
        # ``poes`` was passed in as a scalar;
        # convert it to 1D array of 1 element
        poes = poes.reshape(1)

    result = []
    imls = numpy.array(imls[::-1])

    for curve in curves:
        hmap_val = numpy.interp(poes, curve[::-1], imls)
        result.append(hmap_val)

    return numpy.array(result).transpose()

_HAZ_MAP_DISP_NAME_MEAN_FMT = 'hazard-map(%(poe)s)-%(imt)s-mean'
_HAZ_MAP_DISP_NAME_QUANTILE_FMT = (
    'hazard-map(%(poe)s)-%(imt)s-quantile(%(quantile)s)')
# Hazard maps for a specific end branch
_HAZ_MAP_DISP_NAME_FMT = 'hazard-map(%(poe)s)-%(imt)s-rlz-%(rlz)s'

_UHS_DISP_NAME_MEAN_FMT = 'uhs-(%(poe)s)-mean'
_UHS_DISP_NAME_QUANTILE_FMT = 'uhs-(%(poe)s)-quantile(%(quantile)s)'
_UHS_DISP_NAME_FMT = 'uhs-(%(poe)s)-rlz-%(rlz)s'

# Silencing 'Too many local variables'
# pylint: disable=R0914
def hazard_curves_to_hazard_map(job_id, hazard_curve_id, poes):
    """
    Function to process a set of hazard curves into 1 hazard map for each PoE
    in ``poes``.

    Hazard map results are written directly to the database.

    :param int job_id:
        ID of the current :class:`openquake.engine.db.models.OqJob`.
    :param int hazard_curve_id:
        ID of a set of
        :class:`hazard curves <openquake.engine.db.models.HazardCurve>`.
    :param list poes:
        List of PoEs for which we want to iterpolate hazard maps.
    """
    job = models.OqJob.objects.get(id=job_id)
    hc = models.HazardCurve.objects.get(id=hazard_curve_id)

    hcd = models.HazardCurveData.objects.all_curves_simple(
        filter_args=dict(hazard_curve=hc.id), order_by='location'
    )
    hcd = list(hcd)

    imt = hc.imt
    if imt == 'SA':
        # if it's SA, include the period using the standard notation
        imt = 'SA(%s)' % hc.sa_period

    # Gather all of the curves and compute the maps, for all PoEs
    curves = (poes for _, _, poes in hcd)
    hazard_maps = compute_hazard_maps(curves, hc.imls, poes)

    # Prepare the maps to be saved to the DB
    for i, poe in enumerate(poes):
        map_values = hazard_maps[i]
        lons = numpy.empty(map_values.shape)
        lats = numpy.empty(map_values.shape)

        for loc_idx, _ in enumerate(map_values):
            lons[loc_idx] = hcd[loc_idx][0]
            lats[loc_idx] = hcd[loc_idx][1]

        # Create 'Output' records for the map for this PoE
        if hc.statistics == 'mean':
            disp_name = _HAZ_MAP_DISP_NAME_MEAN_FMT % dict(poe=poe, imt=imt)
        elif hc.statistics == 'quantile':
            disp_name = _HAZ_MAP_DISP_NAME_QUANTILE_FMT % dict(
                poe=poe, imt=imt, quantile=hc.quantile)
        else:
            disp_name = _HAZ_MAP_DISP_NAME_FMT % dict(
                poe=poe, imt=imt, rlz=hc.lt_realization.id)

        output = models.Output.objects.create_output(
            job, disp_name, 'hazard_map'
        )
        # Save the complete hazard map
        models.HazardMap.objects.create(
            output=output,
            lt_realization=hc.lt_realization,
            investigation_time=hc.investigation_time,
            imt=hc.imt,
            statistics=hc.statistics,
            quantile=hc.quantile,
            sa_period=hc.sa_period,
            sa_damping=hc.sa_damping,
            poe=poe,
            lons=lons,
            lats=lats,
            imls=map_values,
        )

# Disabling 'invalid name'
# pylint: disable=C0103
hazard_curves_to_hazard_map_task = utils_tasks.oqtask(
    hazard_curves_to_hazard_map)
hazard_curves_to_hazard_map_task.ignore_result = False

def do_hazard_map_post_process(job):
    """
    Create and distribute tasks for processing hazard curves into hazard maps.

    :param job:
        A :class:`openquake.engine.db.models.OqJob` which has some hazard
        curves associated with it.
    """
    logs.LOG.debug('> Post-processing - Hazard Maps')
    block_size = int(config.get('hazard', 'concurrent_tasks'))

    poes = job.hazard_calculation.poes

    # Stats for debug logging:
    hazard_curve_ids = models.HazardCurve.objects.filter(
        output__oq_job=job).values_list('id', flat=True)
    logs.LOG.debug('num haz curves: %s' % len(hazard_curve_ids))

    # Limit the number of concurrent tasks to the configured concurrency level:
    block_gen = block_splitter(hazard_curve_ids, block_size)
    total_blocks = int(math.ceil(len(hazard_curve_ids) / float(block_size)))

    for i, block in enumerate(block_gen):
        logs.LOG.debug('> Hazard post-processing block, %s of %s'
                       % (i + 1, total_blocks))

        if openquake.engine.no_distribute():
            # just execute the post-processing using the plain function form of
            # the task
            for hazard_curve_id in block:
                hazard_curves_to_hazard_map_task(job.id, hazard_curve_id, poes)
        else:
            tasks = []
            for hazard_curve_id in block:
                tasks.append(hazard_curves_to_hazard_map_task.subtask(
                    (job.id, hazard_curve_id, poes)))
            results = TaskSet(tasks=tasks).apply_async()

            utils_tasks._check_exception(results)

        logs.LOG.debug('< Done Hazard Map post-processing block, %s of %s'
                       % (i + 1, total_blocks))
    logs.LOG.debug('< Done post-processing - Hazard Maps')

def do_uhs_post_proc(job):
    """
    Compute and save (to the DB) Uniform Hazard Spectra for all hazard maps for
    the given ``job``.

    :param job:
        Instance of :class:`openquake.engine.db.models.OqJob`.
    """
    hc = job.hazard_calculation

    rlzs = models.LtRealization.objects.filter(hazard_calculation=hc)

    for poe in hc.poes:
        maps_for_poe = models.HazardMap.objects.filter(
            output__oq_job=job, poe=poe
        )

        # mean (if defined)
        mean_maps = maps_for_poe.filter(statistics='mean')
        if mean_maps.count() > 0:
            mean_uhs = make_uhs(mean_maps)
            _save_uhs(job, mean_uhs, poe, statistics='mean')

        # quantiles (if defined)
        for quantile in hc.quantile_hazard_curves:
            quantile_maps = maps_for_poe.filter(
                statistics='quantile', quantile=quantile
            )
            quantile_uhs = make_uhs(quantile_maps)
            _save_uhs(job, quantile_uhs, poe, statistics='quantile',
                      quantile=quantile)

        # for each logic tree branch:
        for rlz in rlzs:
            rlz_maps = maps_for_poe.filter(
                statistics=None, lt_realization=rlz
            )
            rlz_uhs = make_uhs(rlz_maps)
            _save_uhs(job, rlz_uhs, poe, rlz=rlz)

def make_uhs(maps):
    """
    Make Uniform Hazard Spectra curves for each location.

    It is assumed that the `lons` and `lats` for each of the ``maps`` are
    uniform.

    :param maps:
        A sequence of :class:`openquake.engine.db.models.HazardMap` objects, or
        equivalent objects with the same fields attributes.
    :returns:
        A `dict` with two values:
        * 'periods': a list of the SA periods from all of the ``maps``, sorted
          ascendingly
        * 'uh_spectra': a list of triples (lon, lat, imls), where `imls` is a
          `tuple` of the IMLs from all maps for each of the `periods`
    """
    result = dict()
    result['periods'] = []

    # filter out non-PGA -SA maps
    maps = [x for x in maps if x.imt in ('PGA', 'SA')]

    # give PGA maps an sa_period of 0.0
    # this is slightly hackish, but makes the sorting simple
    for each_map in maps:
        if each_map.imt == 'PGA':
            each_map.sa_period = 0.0

    # sort the maps by period:
    sorted_maps = sorted(maps, key=lambda m: m.sa_period)

    # start constructing the results:
    result['periods'] = [x.sa_period for x in sorted_maps]

    # assume the `lons` and `lats` are uniform for all maps
    lons = sorted_maps[0].lons
    lats = sorted_maps[0].lats

    result['uh_spectra'] = []
    imls_list = izip(*(x.imls for x in sorted_maps))
    for i, (lon, lat, imls) in enumerate(izip(lons, lats, imls_list)):
        result['uh_spectra'].append((lon, lat, imls))

    return result

def _save_uhs(job, uhs_results, poe, rlz=None, statistics=None, quantile=None):
    """
    Save computed UHS data to the DB.

    UHS results can be either for an end branch or for mean or quantile
    statistics.

    :param job:
        :class:`openquake.engine.db.models.OqJob` instance to be associated
        with the results.
    :param uhs_results:
        UHS computation results structured like the output of :func:`make_uhs`.
    :param float poe:
        Probability of exceedance of the hazard maps from which these UH
        Spectra were produced.
    :param rlz:
        :class:`openquake.engine.db.models.LtRealization`. Specify only if
        these results are for an end branch.
    :param statistics:
        'mean' or 'quantile'. Specify only if these are statistical results.
    :param float quantile:
        Specify only if ``statistics`` == 'quantile'.
    """
    output = models.Output(
        oq_job=job,
        owner=job.owner,
        output_type='uh_spectra'
    )
    uhs = models.UHS(
        poe=poe,
        investigation_time=job.hazard_calculation.investigation_time,
        periods=uhs_results['periods'],
    )
    if rlz is not None:
        uhs.lt_realization = rlz
        output.display_name = _UHS_DISP_NAME_FMT % dict(poe=poe, rlz=rlz.id)
    elif statistics is not None:
        uhs.statistics = statistics
        if statistics == 'quantile':
            uhs.quantile = quantile
            output.display_name = (_UHS_DISP_NAME_QUANTILE_FMT
                                % dict(poe=poe, quantile=quantile))
        else:
            # mean
            output.display_name = _UHS_DISP_NAME_MEAN_FMT % dict(poe=poe)
    output.save()
    uhs.output = output
    # This should fail if neither `lt_realization` nor `statistics` is defined:
    uhs.save()

    with transaction.commit_on_success(using='reslt_writer'):
        inserter = BulkInserter(models.UHSData,
                                max_cache_size=CURVE_CACHE_SIZE)

        for lon, lat, imls in uhs_results['uh_spectra']:
            inserter.add_entry(
                uhs_id=uhs.id,
                imls='{%s}' % ','.join(str(x) for x in imls),
                location='POINT(%s %s)' % (lon, lat)
            )
        inserter.flush()
<commit_msg>calcs/hazard/classical/post_processing:<commit_after># -*- coding: utf-8 -*-
# pylint: enable=W0511,W0142,I0011,E1101,E0611,F0401,E1103,R0801,W0232

# Copyright (c) 2010-2013, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""
Post processing functionality for the classical PSHA hazard calculator.
E.g. mean and quantile curves.
"""

import math
import numpy

from celery.task.sets import TaskSet
from django.db import transaction
from itertools import izip

import openquake.engine

from openquake.engine import logs
from openquake.engine.calculators.hazard.general import CURVE_CACHE_SIZE
from openquake.engine.db import models
from openquake.engine.utils import config
from openquake.engine.utils import tasks as utils_tasks
from openquake.engine.utils.general import block_splitter
from openquake.engine.writer import BulkInserter

# Number of locations considered by each task
DEFAULT_LOCATIONS_PER_TASK = 1000

def compute_hazard_maps(curves, imls, poes):
    """
    Given a set of hazard curve poes, interpolate a hazard map at the specified
    ``poe``.

    :param curves:
        2D array of floats. Each row represents a curve, where the values
        in the row are the PoEs (Probabilities of Exceedance) corresponding to
        ``imls``. Each curve corresponds to a geographical location.
    :param imls:
        Intensity Measure Levels associated with these hazard ``curves``. Type
        should be an array-like of floats.
    :param float poes:
        Value(s) on which to interpolate a hazard map from the input
        ``curves``. Can be an array-like or scalar value (for a single PoE).

    :returns:
        A 2D numpy array of hazard map data. Each element/row in the resulting
        array represents the interpolated map for each ``poes`` value
        specified. If ``poes`` is just a single scalar value, the result array
        will have a length of 1.

        The results are structure this way so it is easy to iterate over the
        hazard map results, and in a consistent way (no matter how many
        ``poes`` values are specified).
    """
    poes = numpy.array(poes)

    if len(poes.shape) == 0:
        # ``poes`` was passed in as a scalar;
        # convert it to 1D array of 1 element
        poes = poes.reshape(1)

    result = []
    imls = numpy.array(imls[::-1])

    for curve in curves:
        hmap_val = numpy.interp(poes, curve[::-1], imls)
        result.append(hmap_val)

    return numpy.array(result).transpose()

_HAZ_MAP_DISP_NAME_MEAN_FMT = 'hazard-map(%(poe)s)-%(imt)s-mean'
_HAZ_MAP_DISP_NAME_QUANTILE_FMT = (
    'hazard-map(%(poe)s)-%(imt)s-quantile(%(quantile)s)')
# Hazard maps for a specific end branch
_HAZ_MAP_DISP_NAME_FMT = 'hazard-map(%(poe)s)-%(imt)s-rlz-%(rlz)s'

_UHS_DISP_NAME_MEAN_FMT = 'uhs-(%(poe)s)-mean'
_UHS_DISP_NAME_QUANTILE_FMT = 'uhs-(%(poe)s)-quantile(%(quantile)s)'
_UHS_DISP_NAME_FMT = 'uhs-(%(poe)s)-rlz-%(rlz)s'

# Silencing 'Too many local variables'
# pylint: disable=R0914
def hazard_curves_to_hazard_map(job_id, hazard_curve_id, poes):
    """
    Function to process a set of hazard curves into 1 hazard map for each PoE
    in ``poes``.

    Hazard map results are written directly to the database.

    :param int job_id:
        ID of the current :class:`openquake.engine.db.models.OqJob`.
    :param int hazard_curve_id:
        ID of a set of
        :class:`hazard curves <openquake.engine.db.models.HazardCurve>`.
    :param list poes:
        List of PoEs for which we want to iterpolate hazard maps.
    """
    job = models.OqJob.objects.get(id=job_id)
    hc = models.HazardCurve.objects.get(id=hazard_curve_id)

    hcd = models.HazardCurveData.objects.all_curves_simple(
        filter_args=dict(hazard_curve=hc.id), order_by='location'
    )
    hcd = list(hcd)

    imt = hc.imt
    if imt == 'SA':
        # if it's SA, include the period using the standard notation
        imt = 'SA(%s)' % hc.sa_period

    # Gather all of the curves and compute the maps, for all PoEs
    curves = (poes for _, _, poes in hcd)
    hazard_maps = compute_hazard_maps(curves, hc.imls, poes)

    # Prepare the maps to be saved to the DB
    for i, poe in enumerate(poes):
        map_values = hazard_maps[i]
        lons = numpy.empty(map_values.shape)
        lats = numpy.empty(map_values.shape)

        for loc_idx, _ in enumerate(map_values):
            lons[loc_idx] = hcd[loc_idx][0]
            lats[loc_idx] = hcd[loc_idx][1]

        # Create 'Output' records for the map for this PoE
        if hc.statistics == 'mean':
            disp_name = _HAZ_MAP_DISP_NAME_MEAN_FMT % dict(poe=poe, imt=imt)
        elif hc.statistics == 'quantile':
            disp_name = _HAZ_MAP_DISP_NAME_QUANTILE_FMT % dict(
                poe=poe, imt=imt, quantile=hc.quantile)
        else:
            disp_name = _HAZ_MAP_DISP_NAME_FMT % dict(
                poe=poe, imt=imt, rlz=hc.lt_realization.id)

        output = models.Output.objects.create_output(
            job, disp_name, 'hazard_map'
        )
        # Save the complete hazard map
        models.HazardMap.objects.create(
            output=output,
            lt_realization=hc.lt_realization,
            investigation_time=hc.investigation_time,
            imt=hc.imt,
            statistics=hc.statistics,
            quantile=hc.quantile,
            sa_period=hc.sa_period,
            sa_damping=hc.sa_damping,
            poe=poe,
            lons=lons,
            lats=lats,
            imls=map_values,
        )

# Disabling 'invalid name'
# pylint: disable=C0103
hazard_curves_to_hazard_map_task = utils_tasks.oqtask(
    hazard_curves_to_hazard_map)
hazard_curves_to_hazard_map_task.ignore_result = False

def do_hazard_map_post_process(job):
    """
    Create and distribute tasks for processing hazard curves into hazard maps.

    :param job:
        A :class:`openquake.engine.db.models.OqJob` which has some hazard
        curves associated with it.
    """
    logs.LOG.debug('> Post-processing - Hazard Maps')
    block_size = int(config.get('hazard', 'concurrent_tasks'))

    poes = job.hazard_calculation.poes

    # Stats for debug logging:
    hazard_curve_ids = models.HazardCurve.objects.filter(
        output__oq_job=job).values_list('id', flat=True)
    logs.LOG.debug('num haz curves: %s' % len(hazard_curve_ids))

    # Limit the number of concurrent tasks to the configured concurrency level:
    block_gen = block_splitter(hazard_curve_ids, block_size)
    total_blocks = int(math.ceil(len(hazard_curve_ids) / float(block_size)))

    for i, block in enumerate(block_gen):
        logs.LOG.debug('> Hazard post-processing block, %s of %s'
                       % (i + 1, total_blocks))

        if openquake.engine.no_distribute():
            # just execute the post-processing using the plain function form of
            # the task
            for hazard_curve_id in block:
                hazard_curves_to_hazard_map_task(job.id, hazard_curve_id, poes)
        else:
            tasks = []
            for hazard_curve_id in block:
                tasks.append(hazard_curves_to_hazard_map_task.subtask(
                    (job.id, hazard_curve_id, poes)))
            results = TaskSet(tasks=tasks).apply_async()

            utils_tasks._check_exception(results)

        logs.LOG.debug('< Done Hazard Map post-processing block, %s of %s'
                       % (i + 1, total_blocks))
    logs.LOG.debug('< Done post-processing - Hazard Maps')

def do_uhs_post_proc(job):
    """
    Compute and save (to the DB) Uniform Hazard Spectra for all hazard maps for
    the given ``job``.

    :param job:
        Instance of :class:`openquake.engine.db.models.OqJob`.
    """
    hc = job.hazard_calculation

    rlzs = models.LtRealization.objects.filter(hazard_calculation=hc)

    for poe in hc.poes:
        maps_for_poe = models.HazardMap.objects.filter(
            output__oq_job=job, poe=poe
        )

        # mean (if defined)
        mean_maps = maps_for_poe.filter(statistics='mean')
        if mean_maps.count() > 0:
            mean_uhs = make_uhs(mean_maps)
            _save_uhs(job, mean_uhs, poe, statistics='mean')

        # quantiles (if defined)
        for quantile in hc.quantile_hazard_curves:
            quantile_maps = maps_for_poe.filter(
                statistics='quantile', quantile=quantile
            )
            quantile_uhs = make_uhs(quantile_maps)
            _save_uhs(job, quantile_uhs, poe, statistics='quantile',
                      quantile=quantile)

        # for each logic tree branch:
        for rlz in rlzs:
            rlz_maps = maps_for_poe.filter(
                statistics=None, lt_realization=rlz
            )
            rlz_uhs = make_uhs(rlz_maps)
            _save_uhs(job, rlz_uhs, poe, rlz=rlz)

def make_uhs(maps):
    """
    Make Uniform Hazard Spectra curves for each location.

    It is assumed that the `lons` and `lats` for each of the ``maps`` are
    uniform.

    :param maps:
        A sequence of :class:`openquake.engine.db.models.HazardMap` objects, or
        equivalent objects with the same fields attributes.
    :returns:
        A `dict` with two values:
        * 'periods': a list of the SA periods from all of the ``maps``, sorted
          ascendingly
        * 'uh_spectra': a list of triples (lon, lat, imls), where `imls` is a
          `tuple` of the IMLs from all maps for each of the `periods`
    """
    result = dict()
    result['periods'] = []

    # filter out non-PGA -SA maps
    maps = [x for x in maps if x.imt in ('PGA', 'SA')]

    # give PGA maps an sa_period of 0.0
    # this is slightly hackish, but makes the sorting simple
    for each_map in maps:
        if each_map.imt == 'PGA':
            each_map.sa_period = 0.0

    # sort the maps by period:
    sorted_maps = sorted(maps, key=lambda m: m.sa_period)

    # start constructing the results:
    result['periods'] = [x.sa_period for x in sorted_maps]

    # assume the `lons` and `lats` are uniform for all maps
    lons = sorted_maps[0].lons
    lats = sorted_maps[0].lats

    result['uh_spectra'] = []
    imls_list = izip(*(x.imls for x in sorted_maps))
    for lon, lat, imls in izip(lons, lats, imls_list):
        result['uh_spectra'].append((lon, lat, imls))

    return result

def _save_uhs(job, uhs_results, poe, rlz=None, statistics=None, quantile=None):
    """
    Save computed UHS data to the DB.

    UHS results can be either for an end branch or for mean or quantile
    statistics.

    :param job:
        :class:`openquake.engine.db.models.OqJob` instance to be associated
        with the results.
    :param uhs_results:
        UHS computation results structured like the output of :func:`make_uhs`.
    :param float poe:
        Probability of exceedance of the hazard maps from which these UH
        Spectra were produced.
    :param rlz:
        :class:`openquake.engine.db.models.LtRealization`. Specify only if
        these results are for an end branch.
    :param statistics:
        'mean' or 'quantile'. Specify only if these are statistical results.
    :param float quantile:
        Specify only if ``statistics`` == 'quantile'.
    """
    output = models.Output(
        oq_job=job,
        owner=job.owner,
        output_type='uh_spectra'
    )
    uhs = models.UHS(
        poe=poe,
        investigation_time=job.hazard_calculation.investigation_time,
        periods=uhs_results['periods'],
    )
    if rlz is not None:
        uhs.lt_realization = rlz
        output.display_name = _UHS_DISP_NAME_FMT % dict(poe=poe, rlz=rlz.id)
    elif statistics is not None:
        uhs.statistics = statistics
        if statistics == 'quantile':
            uhs.quantile = quantile
            output.display_name = (_UHS_DISP_NAME_QUANTILE_FMT
                                % dict(poe=poe, quantile=quantile))
        else:
            # mean
            output.display_name = _UHS_DISP_NAME_MEAN_FMT % dict(poe=poe)
    output.save()
    uhs.output = output
    # This should fail if neither `lt_realization` nor `statistics` is defined:
    uhs.save()

    with transaction.commit_on_success(using='reslt_writer'):
        inserter = BulkInserter(models.UHSData,
                                max_cache_size=CURVE_CACHE_SIZE)

        for lon, lat, imls in uhs_results['uh_spectra']:
            inserter.add_entry(
                uhs_id=uhs.id,
                imls='{%s}' % ','.join(str(x) for x in imls),
                location='POINT(%s %s)' % (lon, lat)
            )
        inserter.flush()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# pylint: enable=W0511,W0142,I0011,E1101,E0611,F0401,E1103,R0801,W0232

# Copyright (c) 2010-2013, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""
Post processing functionality for the classical PSHA hazard calculator.
E.g. mean and quantile curves.
"""

import numpy

from django.db import transaction
from itertools import izip

from openquake.engine import logs
from openquake.engine.calculators.hazard.general import CURVE_CACHE_SIZE
from openquake.engine.db import models
from openquake.engine.utils import tasks
from openquake.engine.writer import CacheInserter

# minimum value for the poes: under this value use linear interpolation
EPSILON = 1E-30

# Number of locations considered by each task
DEFAULT_LOCATIONS_PER_TASK = 1000

def compute_hazard_maps(curves, imls, poes):
    """
    Given a set of hazard curve poes, interpolate a hazard map at the specified
    ``poe``.

    :param curves:
        2D array of floats. Each row represents a curve, where the values
        in the row are the PoEs (Probabilities of Exceedance) corresponding to
        ``imls``. Each curve corresponds to a geographical location.
    :param imls:
        Intensity Measure Levels associated with these hazard ``curves``. Type
        should be an array-like of floats.
    :param float poes:
        Value(s) on which to interpolate a hazard map from the input
        ``curves``. Can be an array-like or scalar value (for a single PoE).

    :returns:
        A 2D numpy array of hazard map data. Each element/row in the resulting
        array represents the interpolated map for each ``poes`` value
        specified. If ``poes`` is just a single scalar value, the result array
        will have a length of 1.

        The results are structure this way so it is easy to iterate over the
        hazard map results, and in a consistent way (no matter how many
        ``poes`` values are specified).
    """
    poes = numpy.array(poes)

    if len(poes.shape) == 0:
        # ``poes`` was passed in as a scalar;
        # convert it to 1D array of 1 element
        poes = poes.reshape(1)

    result = []
    imls = numpy.array(imls[::-1])

    for curve in curves:
        curve_cutoff = [max(p, EPSILON) for p in curve[::-1]]
        # exp-log interpolation, to reduce numerical errors
        # see https://bugs.launchpad.net/oq-engine/+bug/1252770
        hmap_val = numpy.exp(numpy.interp(
            numpy.log(poes), numpy.log(curve_cutoff), numpy.log(imls)))

        result.append(hmap_val)

    return numpy.array(result).transpose()

_HAZ_MAP_DISP_NAME_MEAN_FMT = 'Mean Hazard map(%(poe)s) %(imt)s'
_HAZ_MAP_DISP_NAME_QUANTILE_FMT = (
    '%(quantile)s Quantile Hazard Map(%(poe)s) %(imt)s')
# Hazard maps for a specific end branch
_HAZ_MAP_DISP_NAME_FMT = 'Hazard Map(%(poe)s) %(imt)s rlz-%(rlz)s'

_UHS_DISP_NAME_MEAN_FMT = 'Mean UHS (%(poe)s)'
_UHS_DISP_NAME_QUANTILE_FMT = '%(quantile)s Quantile UHS (%(poe)s)'
_UHS_DISP_NAME_FMT = 'UHS (%(poe)s) rlz-%(rlz)s'

# Silencing 'Too many local variables'
# pylint: disable=R0914
def hazard_curves_to_hazard_map(job_id, hazard_curve_id, poes):
    """
    Function to process a set of hazard curves into 1 hazard map for each PoE
    in ``poes``.

    Hazard map results are written directly to the database.

    :param int job_id:
        ID of the current :class:`openquake.engine.db.models.OqJob`.
    :param int hazard_curve_id:
        ID of a set of
        :class:`hazard curves <openquake.engine.db.models.HazardCurve>`.
    :param list poes:
        List of PoEs for which we want to iterpolate hazard maps.
    """
    job = models.OqJob.objects.get(id=job_id)
    hc = models.HazardCurve.objects.get(id=hazard_curve_id)

    hcd = models.HazardCurveData.objects.all_curves_simple(
        filter_args=dict(hazard_curve=hc.id), order_by='location'
    )
    hcd = list(hcd)

    imt = hc.imt
    if imt == 'SA':
        # if it's SA, include the period using the standard notation
        imt = 'SA(%s)' % hc.sa_period

    # Gather all of the curves and compute the maps, for all PoEs
    curves = (poes for _, _, poes in hcd)
    hazard_maps = compute_hazard_maps(curves, hc.imls, poes)

    # Prepare the maps to be saved to the DB
    for i, poe in enumerate(poes):
        map_values = hazard_maps[i]
        lons = numpy.empty(map_values.shape)
        lats = numpy.empty(map_values.shape)

        for loc_idx, _ in enumerate(map_values):
            lons[loc_idx] = hcd[loc_idx][0]
            lats[loc_idx] = hcd[loc_idx][1]

        # Create 'Output' records for the map for this PoE
        if hc.statistics == 'mean':
            disp_name = _HAZ_MAP_DISP_NAME_MEAN_FMT % dict(poe=poe, imt=imt)
        elif hc.statistics == 'quantile':
            disp_name = _HAZ_MAP_DISP_NAME_QUANTILE_FMT % dict(
                poe=poe, imt=imt, quantile=hc.quantile)
        else:
            disp_name = _HAZ_MAP_DISP_NAME_FMT % dict(
                poe=poe, imt=imt, rlz=hc.lt_realization.id)

        output = models.Output.objects.create_output(
            job, disp_name, 'hazard_map'
        )
        # Save the complete hazard map
        models.HazardMap.objects.create(
            output=output,
            lt_realization=hc.lt_realization,
            investigation_time=hc.investigation_time,
            imt=hc.imt,
            statistics=hc.statistics,
            quantile=hc.quantile,
            sa_period=hc.sa_period,
            sa_damping=hc.sa_damping,
            poe=poe,
            lons=lons.tolist(),
            lats=lats.tolist(),
            imls=map_values.tolist(),
        )

hazard_curves_to_hazard_map_task = tasks.oqtask(hazard_curves_to_hazard_map)

def hazard_curves_to_hazard_map_task_arg_gen(job):
    """
    Yield task arguments for processing hazard curves into hazard maps.

    :param job:
        A :class:`openquake.engine.db.models.OqJob` which has some hazard
        curves associated with it.
    """
    poes = job.hazard_calculation.poes

    hazard_curve_ids = models.HazardCurve.objects.filter(
        output__oq_job=job, imt__isnull=False).values_list('id', flat=True)
    logs.LOG.debug('num haz curves: %d', len(hazard_curve_ids))

    for hazard_curve_id in hazard_curve_ids:
        yield job.id, hazard_curve_id, poes

def do_uhs_post_proc(job):
    """
    Compute and save (to the DB) Uniform Hazard Spectra for all hazard maps for
    the given ``job``.

    :param job:
        Instance of :class:`openquake.engine.db.models.OqJob`.
    """
    hc = job.hazard_calculation

    rlzs = models.LtRealization.objects.filter(hazard_calculation=hc)

    for poe in hc.poes:
        maps_for_poe = models.HazardMap.objects.filter(
            output__oq_job=job, poe=poe
        )

        # mean (if defined)
        mean_maps = maps_for_poe.filter(statistics='mean')
        if mean_maps.count() > 0:
            mean_uhs = make_uhs(mean_maps)
            _save_uhs(job, mean_uhs, poe, statistics='mean')

        # quantiles (if defined)
        for quantile in hc.quantile_hazard_curves:
            quantile_maps = maps_for_poe.filter(
                statistics='quantile', quantile=quantile
            )
            quantile_uhs = make_uhs(quantile_maps)
            _save_uhs(job, quantile_uhs, poe, statistics='quantile',
                      quantile=quantile)

        # for each logic tree branch:
        for rlz in rlzs:
            rlz_maps = maps_for_poe.filter(
                statistics=None, lt_realization=rlz
            )
            rlz_uhs = make_uhs(rlz_maps)
            _save_uhs(job, rlz_uhs, poe, rlz=rlz)

def make_uhs(maps):
    """
    Make Uniform Hazard Spectra curves for each location.

    It is assumed that the `lons` and `lats` for each of the ``maps`` are
    uniform.

    :param maps:
        A sequence of :class:`openquake.engine.db.models.HazardMap` objects, or
        equivalent objects with the same fields attributes.
    :returns:
        A `dict` with two values::
            * periods: a list of the SA periods from all of the ``maps``,
              sorted ascendingly
            * uh_spectra: a list of triples (lon, lat, imls), where `imls`
              is a `tuple` of the IMLs from all maps for each of the `periods`
    """
    result = dict()
    result['periods'] = []

    # filter out non-PGA -SA maps
    maps = [x for x in maps if x.imt in ('PGA', 'SA')]

    # give PGA maps an sa_period of 0.0
    # this is slightly hackish, but makes the sorting simple
    for each_map in maps:
        if each_map.imt == 'PGA':
            each_map.sa_period = 0.0

    # sort the maps by period:
    sorted_maps = sorted(maps, key=lambda m: m.sa_period)

    # start constructing the results:
    result['periods'] = [x.sa_period for x in sorted_maps]

    # assume the `lons` and `lats` are uniform for all maps
    lons = sorted_maps[0].lons
    lats = sorted_maps[0].lats

    result['uh_spectra'] = []
    imls_list = izip(*(x.imls for x in sorted_maps))
    for lon, lat, imls in izip(lons, lats, imls_list):
        result['uh_spectra'].append((lon, lat, imls))

    return result

def _save_uhs(job, uhs_results, poe, rlz=None, statistics=None, quantile=None):
    """
    Save computed UHS data to the DB.

    UHS results can be either for an end branch or for mean or quantile
    statistics.

    :param job:
        :class:`openquake.engine.db.models.OqJob` instance to be associated
        with the results.
    :param uhs_results:
        UHS computation results structured like the output of :func:`make_uhs`.
    :param float poe:
        Probability of exceedance of the hazard maps from which these UH
        Spectra were produced.
    :param rlz:
        :class:`openquake.engine.db.models.LtRealization`. Specify only if
        these results are for an end branch.
    :param statistics:
        'mean' or 'quantile'. Specify only if these are statistical results.
    :param float quantile:
        Specify only if ``statistics`` == 'quantile'.
    """
    output = models.Output(
        oq_job=job,
        output_type='uh_spectra'
    )
    uhs = models.UHS(
        poe=poe,
        investigation_time=job.hazard_calculation.investigation_time,
        periods=uhs_results['periods'],
    )
    if rlz is not None:
        uhs.lt_realization = rlz
        output.display_name = _UHS_DISP_NAME_FMT % dict(poe=poe, rlz=rlz.id)
    elif statistics is not None:
        uhs.statistics = statistics
        if statistics == 'quantile':
            uhs.quantile = quantile
            output.display_name = (_UHS_DISP_NAME_QUANTILE_FMT
                                   % dict(poe=poe, quantile=quantile))
        else:
            # mean
            output.display_name = _UHS_DISP_NAME_MEAN_FMT % dict(poe=poe)
    output.save()
    uhs.output = output
    # This should fail if neither `lt_realization` nor `statistics` is defined:
    uhs.save()

    with transaction.commit_on_success(using='job_init'):
        inserter = CacheInserter(models.UHSData, CURVE_CACHE_SIZE)
        for lon, lat, imls in uhs_results['uh_spectra']:
            inserter.add(
                models.UHSData(
                    uhs_id=uhs.id,
                    imls='{%s}' % ','.join(str(x) for x in imls),
                    location='POINT(%s %s)' % (lon, lat))
            )
        inserter.flush()
<commit_msg>Fixed a comment<commit_after># -*- coding: utf-8 -*-
# pylint: enable=W0511,W0142,I0011,E1101,E0611,F0401,E1103,R0801,W0232

# Copyright (c) 2010-2013, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""
Post processing functionality for the classical PSHA hazard calculator.
E.g. mean and quantile curves.
"""

import numpy

from django.db import transaction
from itertools import izip

from openquake.engine import logs
from openquake.engine.calculators.hazard.general import CURVE_CACHE_SIZE
from openquake.engine.db import models
from openquake.engine.utils import tasks
from openquake.engine.writer import CacheInserter

# cutoff value for the poe
EPSILON = 1E-30

# Number of locations considered by each task
DEFAULT_LOCATIONS_PER_TASK = 1000

def compute_hazard_maps(curves, imls, poes):
    """
    Given a set of hazard curve poes, interpolate a hazard map at the specified
    ``poe``.

    :param curves:
        2D array of floats. Each row represents a curve, where the values
        in the row are the PoEs (Probabilities of Exceedance) corresponding to
        ``imls``. Each curve corresponds to a geographical location.
    :param imls:
        Intensity Measure Levels associated with these hazard ``curves``. Type
        should be an array-like of floats.
    :param float poes:
        Value(s) on which to interpolate a hazard map from the input
        ``curves``. Can be an array-like or scalar value (for a single PoE).

    :returns:
        A 2D numpy array of hazard map data. Each element/row in the resulting
        array represents the interpolated map for each ``poes`` value
        specified. If ``poes`` is just a single scalar value, the result array
        will have a length of 1.

        The results are structure this way so it is easy to iterate over the
        hazard map results, and in a consistent way (no matter how many
        ``poes`` values are specified).
    """
    poes = numpy.array(poes)

    if len(poes.shape) == 0:
        # ``poes`` was passed in as a scalar;
        # convert it to 1D array of 1 element
        poes = poes.reshape(1)

    result = []
    imls = numpy.array(imls[::-1])

    for curve in curves:
        curve_cutoff = [max(p, EPSILON) for p in curve[::-1]]
        # exp-log interpolation, to reduce numerical errors
        # see https://bugs.launchpad.net/oq-engine/+bug/1252770
        hmap_val = numpy.exp(numpy.interp(
            numpy.log(poes), numpy.log(curve_cutoff), numpy.log(imls)))

        result.append(hmap_val)

    return numpy.array(result).transpose()

_HAZ_MAP_DISP_NAME_MEAN_FMT = 'Mean Hazard map(%(poe)s) %(imt)s'
_HAZ_MAP_DISP_NAME_QUANTILE_FMT = (
    '%(quantile)s Quantile Hazard Map(%(poe)s) %(imt)s')
# Hazard maps for a specific end branch
_HAZ_MAP_DISP_NAME_FMT = 'Hazard Map(%(poe)s) %(imt)s rlz-%(rlz)s'

_UHS_DISP_NAME_MEAN_FMT = 'Mean UHS (%(poe)s)'
_UHS_DISP_NAME_QUANTILE_FMT = '%(quantile)s Quantile UHS (%(poe)s)'
_UHS_DISP_NAME_FMT = 'UHS (%(poe)s) rlz-%(rlz)s'

# Silencing 'Too many local variables'
# pylint: disable=R0914
def hazard_curves_to_hazard_map(job_id, hazard_curve_id, poes):
    """
    Function to process a set of hazard curves into 1 hazard map for each PoE
    in ``poes``.

    Hazard map results are written directly to the database.

    :param int job_id:
        ID of the current :class:`openquake.engine.db.models.OqJob`.
    :param int hazard_curve_id:
        ID of a set of
        :class:`hazard curves <openquake.engine.db.models.HazardCurve>`.
    :param list poes:
        List of PoEs for which we want to iterpolate hazard maps.
    """
    job = models.OqJob.objects.get(id=job_id)
    hc = models.HazardCurve.objects.get(id=hazard_curve_id)

    hcd = models.HazardCurveData.objects.all_curves_simple(
        filter_args=dict(hazard_curve=hc.id), order_by='location'
    )
    hcd = list(hcd)

    imt = hc.imt
    if imt == 'SA':
        # if it's SA, include the period using the standard notation
        imt = 'SA(%s)' % hc.sa_period

    # Gather all of the curves and compute the maps, for all PoEs
    curves = (poes for _, _, poes in hcd)
    hazard_maps = compute_hazard_maps(curves, hc.imls, poes)

    # Prepare the maps to be saved to the DB
    for i, poe in enumerate(poes):
        map_values = hazard_maps[i]
        lons = numpy.empty(map_values.shape)
        lats = numpy.empty(map_values.shape)

        for loc_idx, _ in enumerate(map_values):
            lons[loc_idx] = hcd[loc_idx][0]
            lats[loc_idx] = hcd[loc_idx][1]

        # Create 'Output' records for the map for this PoE
        if hc.statistics == 'mean':
            disp_name = _HAZ_MAP_DISP_NAME_MEAN_FMT % dict(poe=poe, imt=imt)
        elif hc.statistics == 'quantile':
            disp_name = _HAZ_MAP_DISP_NAME_QUANTILE_FMT % dict(
                poe=poe, imt=imt, quantile=hc.quantile)
        else:
            disp_name = _HAZ_MAP_DISP_NAME_FMT % dict(
                poe=poe, imt=imt, rlz=hc.lt_realization.id)

        output = models.Output.objects.create_output(
            job, disp_name, 'hazard_map'
        )
        # Save the complete hazard map
        models.HazardMap.objects.create(
            output=output,
            lt_realization=hc.lt_realization,
            investigation_time=hc.investigation_time,
            imt=hc.imt,
            statistics=hc.statistics,
            quantile=hc.quantile,
            sa_period=hc.sa_period,
            sa_damping=hc.sa_damping,
            poe=poe,
            lons=lons.tolist(),
            lats=lats.tolist(),
            imls=map_values.tolist(),
        )

hazard_curves_to_hazard_map_task = tasks.oqtask(hazard_curves_to_hazard_map)

def hazard_curves_to_hazard_map_task_arg_gen(job):
    """
    Yield task arguments for processing hazard curves into hazard maps.

    :param job:
        A :class:`openquake.engine.db.models.OqJob` which has some hazard
        curves associated with it.
    """
    poes = job.hazard_calculation.poes

    hazard_curve_ids = models.HazardCurve.objects.filter(
        output__oq_job=job, imt__isnull=False).values_list('id', flat=True)
    logs.LOG.debug('num haz curves: %d', len(hazard_curve_ids))

    for hazard_curve_id in hazard_curve_ids:
        yield job.id, hazard_curve_id, poes

def do_uhs_post_proc(job):
    """
    Compute and save (to the DB) Uniform Hazard Spectra for all hazard maps for
    the given ``job``.

    :param job:
        Instance of :class:`openquake.engine.db.models.OqJob`.
    """
    hc = job.hazard_calculation

    rlzs = models.LtRealization.objects.filter(hazard_calculation=hc)

    for poe in hc.poes:
        maps_for_poe = models.HazardMap.objects.filter(
            output__oq_job=job, poe=poe
        )

        # mean (if defined)
        mean_maps = maps_for_poe.filter(statistics='mean')
        if mean_maps.count() > 0:
            mean_uhs = make_uhs(mean_maps)
            _save_uhs(job, mean_uhs, poe, statistics='mean')

        # quantiles (if defined)
        for quantile in hc.quantile_hazard_curves:
            quantile_maps = maps_for_poe.filter(
                statistics='quantile', quantile=quantile
            )
            quantile_uhs = make_uhs(quantile_maps)
            _save_uhs(job, quantile_uhs, poe, statistics='quantile',
                      quantile=quantile)

        # for each logic tree branch:
        for rlz in rlzs:
            rlz_maps = maps_for_poe.filter(
                statistics=None, lt_realization=rlz
            )
            rlz_uhs = make_uhs(rlz_maps)
            _save_uhs(job, rlz_uhs, poe, rlz=rlz)

def make_uhs(maps):
    """
    Make Uniform Hazard Spectra curves for each location.

    It is assumed that the `lons` and `lats` for each of the ``maps`` are
    uniform.

    :param maps:
        A sequence of :class:`openquake.engine.db.models.HazardMap` objects, or
        equivalent objects with the same fields attributes.
    :returns:
        A `dict` with two values::
            * periods: a list of the SA periods from all of the ``maps``,
              sorted ascendingly
            * uh_spectra: a list of triples (lon, lat, imls), where `imls`
              is a `tuple` of the IMLs from all maps for each of the `periods`
    """
    result = dict()
    result['periods'] = []

    # filter out non-PGA -SA maps
    maps = [x for x in maps if x.imt in ('PGA', 'SA')]

    # give PGA maps an sa_period of 0.0
    # this is slightly hackish, but makes the sorting simple
    for each_map in maps:
        if each_map.imt == 'PGA':
            each_map.sa_period = 0.0

    # sort the maps by period:
    sorted_maps = sorted(maps, key=lambda m: m.sa_period)

    # start constructing the results:
    result['periods'] = [x.sa_period for x in sorted_maps]

    # assume the `lons` and `lats` are uniform for all maps
    lons = sorted_maps[0].lons
    lats = sorted_maps[0].lats

    result['uh_spectra'] = []
    imls_list = izip(*(x.imls for x in sorted_maps))
    for lon, lat, imls in izip(lons, lats, imls_list):
        result['uh_spectra'].append((lon, lat, imls))

    return result

def _save_uhs(job, uhs_results, poe, rlz=None, statistics=None, quantile=None):
    """
    Save computed UHS data to the DB.

    UHS results can be either for an end branch or for mean or quantile
    statistics.

    :param job:
        :class:`openquake.engine.db.models.OqJob` instance to be associated
        with the results.
    :param uhs_results:
        UHS computation results structured like the output of :func:`make_uhs`.
    :param float poe:
        Probability of exceedance of the hazard maps from which these UH
        Spectra were produced.
    :param rlz:
        :class:`openquake.engine.db.models.LtRealization`. Specify only if
        these results are for an end branch.
    :param statistics:
        'mean' or 'quantile'. Specify only if these are statistical results.
    :param float quantile:
        Specify only if ``statistics`` == 'quantile'.
    """
    output = models.Output(
        oq_job=job,
        output_type='uh_spectra'
    )
    uhs = models.UHS(
        poe=poe,
        investigation_time=job.hazard_calculation.investigation_time,
        periods=uhs_results['periods'],
    )
    if rlz is not None:
        uhs.lt_realization = rlz
        output.display_name = _UHS_DISP_NAME_FMT % dict(poe=poe, rlz=rlz.id)
    elif statistics is not None:
        uhs.statistics = statistics
        if statistics == 'quantile':
            uhs.quantile = quantile
            output.display_name = (_UHS_DISP_NAME_QUANTILE_FMT
                                   % dict(poe=poe, quantile=quantile))
        else:
            # mean
            output.display_name = _UHS_DISP_NAME_MEAN_FMT % dict(poe=poe)
    output.save()
    uhs.output = output
    # This should fail if neither `lt_realization` nor `statistics` is defined:
    uhs.save()

    with transaction.commit_on_success(using='job_init'):
        inserter = CacheInserter(models.UHSData, CURVE_CACHE_SIZE)
        for lon, lat, imls in uhs_results['uh_spectra']:
            inserter.add(
                models.UHSData(
                    uhs_id=uhs.id,
                    imls='{%s}' % ','.join(str(x) for x in imls),
                    location='POINT(%s %s)' % (lon, lat))
            )
        inserter.flush()
<|endoftext|>