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

# 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/>.

"""
Hazard getters for Risk calculators.

A HazardGetter is responsible fo getting hazard outputs needed by a risk
calculation.
"""

import numpy

from openquake.engine import logs
from openquake.hazardlib import geo
from openquake.engine.db import models
from openquake.engine.performance import DummyMonitor

#: Scaling constant do adapt to the postgis functions (that work with
#: meters)
KILOMETERS_TO_METERS = 1000

class HazardGetter(object):
    """
    Base abstract class of an Hazard Getter.

    An Hazard Getter is used to query for the closest hazard data for
    each given asset. An Hazard Getter must be pickable such that it
    should be possible to use different strategies (e.g. distributed
    or not, using postgis or not).

    :attr int hazard_output:
        A Hazard Output object :class:`openquake.engine.db.models.Output`

    :attr int hazard_id:
        The ID of an Hazard Output container object
        (e.g. :class:`openquake.engine.db.models.HazardCurve`)

    :attr assets:
        The assets for which we wants to compute.

    :attr max_distance:
        The maximum distance, in kilometers, to use.

    :attr imt:
        The imt (in long form) for which data have to be retrieved

    :attr float weight:
        The weight (if applicable) to be given to the retrieved data
    """
    def __init__(self, hazard_output, assets, max_distance, imt):
        self.hazard_output = hazard_output
        hazard = self.container(hazard_output)
        self.hazard_id = hazard.id
        self.assets = assets
        self.max_distance = max_distance
        self.imt = imt

        if hasattr(hazard, 'lt_realization') and hazard.lt_realization:
            self.weight = hazard.lt_realization.weight
        else:
            self.weight = None

        # FIXME(lp). It is better to directly store the convex hull
        # instead of the mesh. We are not doing it because
        # hazardlib.Polygon is not (yet) pickeable
        self._assets_mesh = geo.mesh.Mesh.from_points_list([
            geo.point.Point(asset.site.x, asset.site.y)
            for asset in self.assets])
        self.asset_dict = dict((asset.id, asset) for asset in self.assets)
        self.all_asset_ids = set(self.asset_dict)

    def container(self, hazard_output):
        """
        Returns the corresponding output container object from an
        Hazard :class:`openquake.engine.db.models.Output` instance
        """
        raise NotImplementedError

    def __repr__(self):
        return "<%s max_distance=%s assets=%s>" % (
            self.__class__.__name__, self.max_distance,
            [a.id for a in self.assets])

    def get_data(self, imt):
        """
        Subclasses must implement this.

        :param str imt: a string representation of the intensity
        measure type (e.g. SA(0.1)) in which the hazard data should be
        returned

        :returns:
            An OrderedDict mapping ID of
            :class:`openquake.engine.db.models.ExposureData` objects to
            hazard_data (e.g. an array with the poes, or an array with the
            ground motion values). Bear in mind that the returned data could
            lack some assets being filtered out by the ``maximum_distance``
            criteria.
        """
        raise NotImplementedError

    def __call__(self):
        """
        :returns:
            A tuple with two elements. The first is an array of instances of
            :class:`openquake.engine.db.models.ExposureData`, the second is an
            array with the corresponding hazard data.
        """
        # data is a gmf or a set of hazard curves
        asset_ids, data = self.get_data(self.imt)

        missing_asset_ids = self.all_asset_ids - set(asset_ids)

        for missing_asset_id in missing_asset_ids:
            # please don't remove this log: it was required by Vitor since
            # this is a case that should NOT happen and must raise a warning
            logs.LOG.warn(
                "No hazard has been found for the asset %s within %s km" % (
                    self.asset_dict[missing_asset_id], self.max_distance))

        ret = ([self.asset_dict[asset_id] for asset_id in asset_ids], data)

        return ret

class HazardCurveGetterPerAsset(HazardGetter):
    """
    Simple HazardCurve Getter that performs a spatial query for each
    asset.

    :attr imls:
        The intensity measure levels of the curves we are going to get.

    :attr dict _cache:
        A cache of the computed hazard curve object on a per-location basis.
    """

    def __init__(self, hazard, assets, max_distance, imt):
        super(HazardCurveGetterPerAsset, self).__init__(
            hazard, assets, max_distance, imt)
        self._cache = {}

    def container(self, hazard_output):
        return hazard_output.hazardcurve

    def get_data(self, imt):
        """
        Calls ``get_by_site`` for each asset and pack the results as
        requested by the :meth:`HazardGetter.get_data` interface.
        """
        imt_type, sa_period, sa_damping = models.parse_imt(imt)

        hc = models.HazardCurve.objects.get(pk=self.hazard_id)

        if hc.output.output_type == 'hazard_curve':
            imls = hc.imls
            hazard_id = self.hazard_id
        elif hc.output.output_type == 'hazard_curve_multi':
            hc = models.HazardCurve.objects.get(
                output__oq_job=hc.output.oq_job,
                output__output_type='hazard_curve',
                statistics=hc.statistics,
                lt_realization=hc.lt_realization,
                imt=imt_type,
                sa_period=sa_period,
                sa_damping=sa_damping)
            imls = hc.imls
            hazard_id = hc.id

        hazard_assets = [(asset.id, self.get_by_site(
            asset.site, hazard_id, imls))
            for asset in self.assets]

        assets = []
        curves = []
        for asset_id, (hazard_curve, distance) in hazard_assets:
            if distance < self.max_distance * KILOMETERS_TO_METERS:
                assets.append(asset_id)
                curves.append(hazard_curve)

        return assets, curves

    def get_by_site(self, site, hazard_id, imls):
        """
        :param site:
            An instance of :class:`django.contrib.gis.geos.point.Point`
            corresponding to the location of an asset.
        """
        if site.wkt in self._cache:
            return self._cache[site.wkt]

        cursor = models.getcursor('job_init')

        query = """
        SELECT
            hzrdr.hazard_curve_data.poes,
            min(ST_Distance(location::geography,
                            ST_GeographyFromText(%s), false))
                AS min_distance
        FROM hzrdr.hazard_curve_data
        WHERE hazard_curve_id = %s
        GROUP BY id
        ORDER BY min_distance
        LIMIT 1;"""

        args = (site.wkt, hazard_id)

        cursor.execute(query, args)
        poes, distance = cursor.fetchone()

        hazard = zip(imls, poes)

        self._cache[site.wkt] = (hazard, distance)

        return hazard, distance

class GroundMotionValuesGetter(HazardGetter):
    """
    Hazard getter for loading ground motion values.
    """

    def container(self, hazard_output):
        return hazard_output.gmfcollection

    def __call__(self, monitor=DummyMonitor()):
        """
        :param monitor: an instance of :class:`openquake.engine.performance.EnginePerformanceMonitor`
        :returns:
            A tuple with two elements. The first is an array of instances of
            :class:`openquake.engine.db.models.ExposureData`, the second is a
            pair (gmfs, ruptures) with the closest ground motion value for each
            asset.
        """
        with monitor('associating asset_ids <-> gmf_ids'):
            asset_ids, gmf_ids = self.get_data(self.imt)
        missing_asset_ids = self.all_asset_ids - set(asset_ids)

        for missing_asset_id in missing_asset_ids:
            # please dont' remove this log: it was required by Vitor since
            # this is a case that should NOT happen and must raise a warning
            logs.LOG.warn(
                "No hazard has been found for the asset %s within %s km" % (
                    self.asset_dict[missing_asset_id], self.max_distance))

        distinct_gmf_ids = tuple(set(gmf_ids))
        cursor = models.getcursor('job_init')

        if self.hazard_output.output_type == 'gmf_scenario':
            gmfs = dict((i, models.GmfAgg.objects.get(pk=i).gmvs)
                        for i in distinct_gmf_ids)
            return ([self.asset_dict[asset_id] for asset_id in asset_ids],
                    [gmfs[i] for i in gmf_ids])

        elif not gmf_ids:  # all missing
            return [], ([], [])

        cursor = models.getcursor('job_init')

        # get the sorted ruptures from all the distinct GMFs
        with monitor('getting ruptures'):
            cursor.execute('''\
        SELECT distinct unnest(array_concat(rupture_ids)) FROM hzrdr.gmf_agg
        WHERE id in %s ORDER BY unnest''', (distinct_gmf_ids,))
            # TODO: in principle it should be possible to remove the ORDER BY;
            # qa_tests.risk.event_based.case_3.test.EventBasedRiskCase3TestCase
            # breaks if I do so (MS)
            sorted_ruptures = numpy.array([r[0] for r in cursor.fetchall()])

        # get the data from the distinct GMFs
        with monitor('getting gmvs'):
            cursor.execute('''\
            SELECT id, gmvs, rupture_ids FROM hzrdr.gmf_agg
            WHERE id in %s''', (distinct_gmf_ids,))
            gmfs = {}
            for gmf_id, gmvs, ruptures in cursor.fetchall():
                gmvd = dict(zip(ruptures, gmvs))
                gmvs = numpy.array([gmvd.get(r, 0.) for r in sorted_ruptures])
                gmfs[gmf_id] = gmvs

        ret = ([self.asset_dict[asset_id] for asset_id in asset_ids],
               ([gmfs[i] for i in gmf_ids], sorted_ruptures))
        return ret

    def get_data(self, imt):
        cursor = models.getcursor('job_init')

        imt_type, sa_period, sa_damping = models.parse_imt(imt)
        spectral_filters = ""
        args = (imt_type, self.hazard_id)

        if imt_type == "SA":
            spectral_filters = "AND sa_period = %s AND sa_damping = %s"
            args += (sa_period, sa_damping)

        # Query explanation. We need to get for each asset the closest
        # ground motion for a given logic tree realization and a given imt.

        # To this aim, we perform a spatial join with the exposure table that
        # is previously filtered by the assets extent, exposure model
        # and taxonomy. We are not filtering with an IN statement on
        # the ids of the assets for perfomance reasons.

        # The ``distinct ON (exposure_data.id)`` combined by the
        # ``ORDER BY ST_Distance`` does the job to select the closest
        # gmvs
        query = """
  SELECT DISTINCT ON (e.id) e.id, g.id
  FROM riski.exposure_data AS e
  JOIN htemp.site_data AS s
  ON ST_DWithin(e.site::geography, s.location::geography, %s)
  JOIN hzrdr.gmf_agg AS g
  ON g.site_id = s.id
  WHERE taxonomy = %s AND exposure_model_id = %s AND
        e.site && %s AND imt = %s AND gmf_collection_id = %s {}
  ORDER BY e.id, ST_Distance(e.site::geography, s.location::geography, false)
           """.format(spectral_filters)  # this will fill in the {}

        assets_extent = self._assets_mesh.get_convex_hull()
        args = (self.max_distance * KILOMETERS_TO_METERS,
                self.assets[0].taxonomy,
                self.assets[0].exposure_model_id,
                assets_extent.wkt) + args

        cursor.execute(query, args)
        # print cursor.mogrify(query, args)

        data = cursor.fetchall()

        assets, gmf_ids = [], []
        for asset_id, gmf_id in data:
            # the query may return spurious assets outside the considered block
            if asset_id in self.asset_dict:  # in block
                assets.append(asset_id)
                gmf_ids.append(gmf_id)
        return assets, gmf_ids
<commit_msg>Fixed the hazard_getters<commit_after># -*- coding: utf-8 -*-

# 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/>.

"""
Hazard getters for Risk calculators.

A HazardGetter is responsible fo getting hazard outputs needed by a risk
calculation.
"""

import numpy

from openquake.engine import logs
from openquake.hazardlib import geo
from openquake.engine.db import models
from openquake.engine.performance import DummyMonitor

#: Scaling constant do adapt to the postgis functions (that work with
#: meters)
KILOMETERS_TO_METERS = 1000

class HazardGetter(object):
    """
    Base abstract class of an Hazard Getter.

    An Hazard Getter is used to query for the closest hazard data for
    each given asset. An Hazard Getter must be pickable such that it
    should be possible to use different strategies (e.g. distributed
    or not, using postgis or not).

    :attr int hazard_output:
        A Hazard Output object :class:`openquake.engine.db.models.Output`

    :attr int hazard_id:
        The ID of an Hazard Output container object
        (e.g. :class:`openquake.engine.db.models.HazardCurve`)

    :attr assets:
        The assets for which we wants to compute.

    :attr max_distance:
        The maximum distance, in kilometers, to use.

    :attr imt:
        The imt (in long form) for which data have to be retrieved

    :attr float weight:
        The weight (if applicable) to be given to the retrieved data
    """
    def __init__(self, hazard_output, assets, max_distance, imt):
        self.hazard_output = hazard_output
        hazard = self.container(hazard_output)
        self.hazard_id = hazard.id
        self.assets = assets
        self.max_distance = max_distance
        self.imt = imt

        if hasattr(hazard, 'lt_realization') and hazard.lt_realization:
            self.weight = hazard.lt_realization.weight
        else:
            self.weight = None

        # FIXME(lp). It is better to directly store the convex hull
        # instead of the mesh. We are not doing it because
        # hazardlib.Polygon is not (yet) pickeable
        self._assets_mesh = geo.mesh.Mesh.from_points_list([
            geo.point.Point(asset.site.x, asset.site.y)
            for asset in self.assets])
        self.asset_dict = dict((asset.id, asset) for asset in self.assets)
        self.all_asset_ids = set(self.asset_dict)

    def container(self, hazard_output):
        """
        Returns the corresponding output container object from an
        Hazard :class:`openquake.engine.db.models.Output` instance
        """
        raise NotImplementedError

    def __repr__(self):
        return "<%s max_distance=%s assets=%s>" % (
            self.__class__.__name__, self.max_distance,
            [a.id for a in self.assets])

    def get_data(self, imt):
        """
        Subclasses must implement this.

        :param str imt: a string representation of the intensity
        measure type (e.g. SA(0.1)) in which the hazard data should be
        returned

        :returns:
            An OrderedDict mapping ID of
            :class:`openquake.engine.db.models.ExposureData` objects to
            hazard_data (e.g. an array with the poes, or an array with the
            ground motion values). Bear in mind that the returned data could
            lack some assets being filtered out by the ``maximum_distance``
            criteria.
        """
        raise NotImplementedError

    def __call__(self):
        """
        :returns:
            A tuple with two elements. The first is an array of instances of
            :class:`openquake.engine.db.models.ExposureData`, the second is an
            array with the corresponding hazard data.
        """
        # data is a gmf or a set of hazard curves
        asset_ids, data = self.get_data(self.imt)

        missing_asset_ids = self.all_asset_ids - set(asset_ids)

        for missing_asset_id in missing_asset_ids:
            # please don't remove this log: it was required by Vitor since
            # this is a case that should NOT happen and must raise a warning
            logs.LOG.warn(
                "No hazard has been found for the asset %s within %s km" % (
                    self.asset_dict[missing_asset_id], self.max_distance))

        ret = ([self.asset_dict[asset_id] for asset_id in asset_ids], data)

        return ret

class HazardCurveGetterPerAsset(HazardGetter):
    """
    Simple HazardCurve Getter that performs a spatial query for each
    asset.

    :attr imls:
        The intensity measure levels of the curves we are going to get.

    :attr dict _cache:
        A cache of the computed hazard curve object on a per-location basis.
    """

    def __init__(self, hazard, assets, max_distance, imt):
        super(HazardCurveGetterPerAsset, self).__init__(
            hazard, assets, max_distance, imt)
        self._cache = {}

    def container(self, hazard_output):
        return hazard_output.hazardcurve

    def get_data(self, imt):
        """
        Calls ``get_by_site`` for each asset and pack the results as
        requested by the :meth:`HazardGetter.get_data` interface.
        """
        imt_type, sa_period, sa_damping = models.parse_imt(imt)

        hc = models.HazardCurve.objects.get(pk=self.hazard_id)

        if hc.output.output_type == 'hazard_curve':
            imls = hc.imls
            hazard_id = self.hazard_id
        elif hc.output.output_type == 'hazard_curve_multi':
            hc = models.HazardCurve.objects.get(
                output__oq_job=hc.output.oq_job,
                output__output_type='hazard_curve',
                statistics=hc.statistics,
                lt_realization=hc.lt_realization,
                imt=imt_type,
                sa_period=sa_period,
                sa_damping=sa_damping)
            imls = hc.imls
            hazard_id = hc.id

        hazard_assets = [(asset.id, self.get_by_site(
            asset.site, hazard_id, imls))
            for asset in self.assets]

        assets = []
        curves = []
        for asset_id, (hazard_curve, distance) in hazard_assets:
            if distance < self.max_distance * KILOMETERS_TO_METERS:
                assets.append(asset_id)
                curves.append(hazard_curve)

        return assets, curves

    def get_by_site(self, site, hazard_id, imls):
        """
        :param site:
            An instance of :class:`django.contrib.gis.geos.point.Point`
            corresponding to the location of an asset.
        """
        if site.wkt in self._cache:
            return self._cache[site.wkt]

        cursor = models.getcursor('job_init')

        query = """
        SELECT
            hzrdr.hazard_curve_data.poes,
            min(ST_Distance(location::geography,
                            ST_GeographyFromText(%s), false))
                AS min_distance
        FROM hzrdr.hazard_curve_data
        WHERE hazard_curve_id = %s
        GROUP BY id
        ORDER BY min_distance
        LIMIT 1;"""

        args = (site.wkt, hazard_id)

        cursor.execute(query, args)
        poes, distance = cursor.fetchone()

        hazard = zip(imls, poes)

        self._cache[site.wkt] = (hazard, distance)

        return hazard, distance

class GroundMotionValuesGetter(HazardGetter):
    """
    Hazard getter for loading ground motion values.
    """

    def container(self, hazard_output):
        return hazard_output.gmfcollection

    def __call__(self, monitor=DummyMonitor()):
        """
        :param monitor: an instance of :class:`openquake.engine.performance.EnginePerformanceMonitor`
        :returns:
            A tuple with two elements. The first is an array of instances of
            :class:`openquake.engine.db.models.ExposureData`, the second is a
            pair (gmfs, ruptures) with the closest ground motion value for each
            asset.
        """
        with monitor('associating asset_ids <-> gmf_ids'):
            asset_ids, gmf_ids = self.get_data(self.imt)
        missing_asset_ids = self.all_asset_ids - set(asset_ids)

        for missing_asset_id in missing_asset_ids:
            # please dont' remove this log: it was required by Vitor since
            # this is a case that should NOT happen and must raise a warning
            logs.LOG.warn(
                "No hazard has been found for the asset %s within %s km" % (
                    self.asset_dict[missing_asset_id], self.max_distance))

        distinct_gmf_ids = tuple(set(gmf_ids))
        cursor = models.getcursor('job_init')

        if self.hazard_output.output_type == 'gmf_scenario':
            gmfs = dict((i, models.GmfAgg.objects.get(pk=i).gmvs)
                        for i in distinct_gmf_ids)
            return ([self.asset_dict[asset_id] for asset_id in asset_ids],
                    [gmfs[i] for i in gmf_ids])

        elif not gmf_ids:  # all missing
            return [], ([], [])

        cursor = models.getcursor('job_init')

        # get the sorted ruptures from all the distinct GMFs
        with monitor('getting ruptures'):
            cursor.execute('''\
        SELECT distinct unnest(array_concat(rupture_ids)) FROM hzrdr.gmf_agg
        WHERE id in %s ORDER BY unnest''', (distinct_gmf_ids,))
            # TODO: in principle it should be possible to remove the ORDER BY;
            # qa_tests.risk.event_based.case_3.test.EventBasedRiskCase3TestCase
            # breaks if I do so (MS)
            sorted_ruptures = numpy.array([r[0] for r in cursor.fetchall()])

        # get the data from the distinct GMFs
        with monitor('getting gmvs'):
            cursor.execute('''\
            SELECT id, gmvs, rupture_ids FROM hzrdr.gmf_agg
            WHERE id in %s''', (distinct_gmf_ids,))
            gmfs = {}
            for gmf_id, gmvs, ruptures in cursor.fetchall():
                gmvd = dict(zip(ruptures, gmvs))
                gmvs = numpy.array([gmvd.get(r, 0.) for r in sorted_ruptures])
                gmfs[gmf_id] = gmvs

        ret = ([self.asset_dict[asset_id] for asset_id in asset_ids],
               ([gmfs[i] for i in gmf_ids], sorted_ruptures))
        return ret

    def get_data(self, imt):
        cursor = models.getcursor('job_init')

        imt_type, sa_period, sa_damping = models.parse_imt(imt)
        spectral_filters = ""
        args = (imt_type, self.hazard_id)

        if imt_type == "SA":
            spectral_filters = "AND sa_period = %s AND sa_damping = %s"
            args += (sa_period, sa_damping)

        # Query explanation. We need to get for each asset the closest
        # ground motion for a given logic tree realization and a given imt.

        # To this aim, we perform a spatial join with the exposure table that
        # is previously filtered by the assets extent, exposure model
        # and taxonomy. We are not filtering with an IN statement on
        # the ids of the assets for perfomance reasons.

        # The ``distinct ON (exposure_data.id)`` combined by the
        # ``ORDER BY ST_Distance`` does the job to select the closest
        # gmvs
        query = """
  SELECT DISTINCT ON (e.id) e.id, g.id
  FROM riski.exposure_data AS e
  JOIN htemp.site_data AS s
  ON ST_DWithin(e.site::geography, s.location::geography, %s)
  JOIN hzrdr.gmf_agg AS g
  ON g.site_id = s.id
  WHERE s.hazard_job_id = %s AND taxonomy = %s AND exposure_model_id = %s
        AND e.site && %s AND imt = %s AND gmf_collection_id = %s {}
  ORDER BY e.id, ST_Distance(e.site::geography, s.location::geography, false)
           """.format(spectral_filters)  # this will fill in the {}

        assets_extent = self._assets_mesh.get_convex_hull()
        args = (self.max_distance * KILOMETERS_TO_METERS,
                self.hazard_output.oq_job.id,
                self.assets[0].taxonomy,
                self.assets[0].exposure_model_id,
                assets_extent.wkt) + args

        cursor.execute(query, args)
        # print cursor.mogrify(query, args)

        data = cursor.fetchall()

        assets, gmf_ids = [], []
        for asset_id, gmf_id in data:
            # the query may return spurious assets outside the considered block
            if asset_id in self.asset_dict:  # in block
                assets.append(asset_id)
                gmf_ids.append(gmf_id)
        return assets, gmf_ids
<|endoftext|>


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

# 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/>.

"""
Hazard getters for Risk calculators.

An HazardGetter is responsible to get hazard outputs needed by a risk
calculation.
"""

from collections import OrderedDict
from openquake.engine import logs
from openquake.hazardlib import geo
from openquake.engine.db import models
from django.db import connection

#: Scaling constant do adapt to the postgis functions (that work with
#: meters)
KILOMETERS_TO_METERS = 1000

class HazardGetter(object):
    """
    Base abstract class of an Hazard Getter.

    An Hazard Getter is used to query for the closest hazard data for
    each given asset. An Hazard Getter must be pickable such that it
    should be possible to use different strategies (e.g. distributed
    or not, using postgis or not).

    :attr hazard_id:
      the ID of an Hazard output container (e.g.
      :class:`openquake.engine.db.models.HazardCurve`)

    :attr imt: the imt of the hazard considered by the getter

    :attr assets: the assets for which we wants to compute

    :attr max_distance: the maximum distance, in kilometers, to use
    """
    def __init__(self, hazard_id, imt, assets, max_distance):
        self.hazard_id = hazard_id
        self.imt = imt
        self.assets = assets
        self.max_distance = max_distance

        self._assets_mesh = geo.mesh.Mesh.from_points_list([
            geo.point.Point(asset.site.x, asset.site.y)
            for asset in self.assets])
        self._imt, self._sa_period, self._sa_damping = (
            models.parse_imt(self.imt))
        self.asset_dict = dict((asset.id, asset) for asset in self.assets)
        self._cache = {}

    def get_data(self):
        """
        :returns: an OrderedDict mapping ID of
        :class:`openquake.engine.db.models.ExposureData` objects to
        hazard_data (e.g. an array with the poes, or an array with the
        ground motion values). Mind that the returned data could lack
        some assets being filtered out by the ``maximum_distance``
        criteria.

        Subclasses must implement this.
        """
        raise NotImplementedError

    def __call__(self):
        """
        :returns: a tuple with three elements. The first is an array
        of instances of
        :class:`openquake.engine.db.models.ExposureData`, the second
        is an array with the corresponding hazard data, the third is
        the array of IDs of assets that has been filtered out by the
        getter by the ``maximum_distance`` criteria.
        """
        data = self.get_data()

        filtered_asset_ids = set(data.keys())
        all_asset_ids = set(self.asset_dict.keys())
        missing_asset_ids = all_asset_ids - filtered_asset_ids
        extra_asset_ids = filtered_asset_ids - all_asset_ids

        # FIXME(lp). It might happens that the convex hull contains
        # more assets than the requested ones. At this moment, we just
        # ignore them.
        if extra_asset_ids:
            logs.LOG.debug("Extra asset have been computed ids: %s" % (
                models.ExposureData.objects.filter(
                    pk__in=extra_asset_ids)))

        for missing_asset_id in missing_asset_ids:
            logs.LOG.warn(
                "No hazard has been found for the asset %s within %s km" % (
                    self.asset_dict[missing_asset_id], self.max_distance))

        return ([self.asset_dict[asset_id] for asset_id in data
                 if asset_id in self.asset_dict],
                [data[asset_id] for asset_id in data
                 if asset_id in self.asset_dict],
                missing_asset_ids)

class HazardCurveGetterPerAsset(HazardGetter):
    """
    Simple HazardCurve Getter that performs a spatial query for each
    asset.

    :attr imls: the intensity measure levels of the curves we are
    going to get.

    :attr dict _cache: a cache of the computed hazard curve object on
    a per-location basis.
    """

    def __init__(self, hazard_id, imt, assets, max_distance):
        super(HazardCurveGetterPerAsset, self).__init__(
            hazard_id, imt, assets, max_distance)
        self.imls = models.HazardCurve.objects.get(
            pk=self.hazard_id).imls

    def get_data(self):
        """
        Calls ``get_by_site`` for each asset and pack the results as
        requested by the :method:`HazardGetter.get_data` interface.
        """
        hazard_assets = [(asset.id, self.get_by_site(asset.site))
                         for asset in self.assets]

        return OrderedDict(
            [(asset_id, hazard_curve)
             for asset_id, (hazard_curve, distance) in hazard_assets
             if distance < self.max_distance * KILOMETERS_TO_METERS])

    def get_by_site(self, site):
        """
        :param site: an instance of
        :class:`django.contrib.gis.geos.point.Point` corresponding to
        the location of an asset.
        """
        if site.wkt in self._cache:
            return self._cache[site.wkt]

        cursor = connection.cursor()

        query = """
        SELECT
            hzrdr.hazard_curve_data.poes,
            min(ST_Distance(location::geography,
                            ST_GeographyFromText(%s), false))
                AS min_distance
        FROM hzrdr.hazard_curve_data
        WHERE hazard_curve_id = %s
        GROUP BY id
        ORDER BY min_distance
        LIMIT 1;"""

        args = (site.wkt, self.hazard_id)

        cursor.execute(query, args)
        poes, distance = cursor.fetchone()

        hazard = zip(self.imls, poes)

        self._cache[site.wkt] = (hazard, distance)

        return hazard, distance

class GroundMotionValuesGetter(HazardGetter):
    """
    Hazard getter for loading ground motion values.
    """

    def get_data(self):
        gmf_collection = models.GmfCollection.objects.get(
            id=self.hazard_id)

        if gmf_collection.output.output_type == "gmf":
            gmf_set_ids = tuple(
                gmf_collection.gmfset_set.values_list('id', flat=True))
        else:
            raise ValueError(
                "Output must be of type `gmf`. "
                "At the moment, we only support computation of loss curves "
                "for a specific logic tree branch.")

        cursor = connection.cursor()

        spectral_filters = ""
        args = (self._imt, gmf_set_ids)

        if self._imt == "SA":
            spectral_filters = "AND sa_period = %s AND sa_damping = %s"
            args += (self._sa_period, self._sa_damping)

        # Query explanation. We need to get for each asset the closest
        # ground motion values (and the corresponding rupture ids from
        # which they have been generated) for a given logic tree
        # realization and a given imt.

        # We first concatenate ground motion values grouped by
        # location (so for each location we have all the ground motion
        # values found in different gmf_sets, collected in a column
        # called ``allgmvs_arr``).

        # For performance reasons, we help the query by filtering also
        # on a polygon built by dilating the assets extent of the
        # maximum distance

        # Then, we perform a spatial join with the exposure table that
        # is previously filtered by the assets extent, exposure model
        # and taxonomy. We are not filtering with an IN statement on
        # the ids of the assets for perfomance reasons.

        # The ``distinct ON (exposure_data.id)`` combined by the
        # ``ORDER BY ST_Distance`` does the job to select the closest
        # gmvs
        query = """
  SELECT DISTINCT ON (oqmif.exposure_data.id)
  oqmif.exposure_data.id, gmf_table.allgmvs_arr, gmf_table.allrupture_ids
  FROM oqmif.exposure_data JOIN
    (SELECT location,
            array_concat(gmvs ORDER BY gmf_set_id, result_grp_ordinal)
            AS allgmvs_arr,
            array_concat(rupture_ids ORDER BY gmf_set_id, result_grp_ordinal)
            AS allrupture_ids
     FROM hzrdr.gmf
     WHERE imt = %s AND gmf_set_id IN %s {}
     AND location && %s
     GROUP BY location) AS gmf_table
  ON ST_DWithin(oqmif.exposure_data.site, gmf_table.location, %s)
  WHERE oqmif.exposure_data.site && %s
  AND taxonomy = %s AND exposure_model_id = %s
  AND array_length(gmf_table.allgmvs_arr, 1) > 0
  ORDER BY oqmif.exposure_data.id,
           ST_Distance(oqmif.exposure_data.site, gmf_table.location, false)
           """.format(spectral_filters)  # this will fill in the {}

        assets_extent = self._assets_mesh.get_convex_hull()
        args += (assets_extent.dilate(self.max_distance).wkt,
                 self.max_distance * KILOMETERS_TO_METERS,
                 assets_extent.wkt,
                 self.assets[0].taxonomy,
                 self.assets[0].exposure_model_id)

        cursor.execute(query, args)

        data = cursor.fetchall()

        # nested dicts with structure: asset_id -> (rupture_id -> gmv)
        assets_ruptures_gmvs = OrderedDict()

        # store all the ruptures returned by the query
        ruptures = []

        for asset_id, gmvs, rupture_ids in data:
            assets_ruptures_gmvs[asset_id] = dict(zip(rupture_ids, gmvs))
            ruptures.extend(rupture_ids)
        ruptures = set(ruptures)

        # We expect that the query may return a different number of
        # gmvs and ruptures for each asset (because only the ruptures
        # that gives a positive ground shaking are stored). Here on,
        # we finalize `assets_ruptures_gmvs` by filling in with zero
        # values for each rupture that has not given a contribute.

        # for each asset, we look for missing ruptures
        for asset_id, ruptures_gmvs_dict in assets_ruptures_gmvs.items():

            # all the ruptures producing a positive ground shaking for
            # `asset`
            asset_ruptures = set(ruptures_gmvs_dict)

            missing_ruptures = ruptures - asset_ruptures

            # we finalize the asset data with 0
            for rupture_id in missing_ruptures:
                ruptures_gmvs_dict[rupture_id] = 0.

        # The OrderedDict is needed by __call__ in order to scan
        # multiple times the data and still get corresponding values.
        # See the return statement of the __call__ method.

        # maps asset_id -> to a 2-tuple (rupture_ids, gmvs)
        return OrderedDict([
            (asset_id, zip(
                *sorted(zip(asset_data.values(), asset_data.keys()),
                        key=lambda x: x[1])))
            for asset_id, asset_data in assets_ruptures_gmvs.items()])

class GroundMotionScenarioGetter(HazardGetter):
    """
    Hazard getter for loading ground motion values. It uses the same
    approach used in :class:`GroundMotionValuesGetter`.
    """

    def setup(self):
        super(GroundMotionScenarioGetter, self).setup()
        if self._sa_period:
            self._imt = '%s(%s)' % (self._imt, self._sa_period)

    def get_data(self):
        cursor = connection.cursor()

        # See the comment in `GroundMotionValuesGetter.get_data` for
        # an explanation of the query
        query = """
  SELECT DISTINCT ON (oqmif.exposure_data.id) oqmif.exposure_data.id,
         gmf_table.gmvs
  FROM oqmif.exposure_data JOIN (
    SELECT location, gmvs
           FROM hzrdr.gmf_scenario
           WHERE hzrdr.gmf_scenario.imt = %s
           AND hzrdr.gmf_scenario.output_id = %s
           AND hzrdr.gmf_scenario.location && %s) gmf_table
  ON ST_DWithin(oqmif.exposure_data.site, gmf_table.location, %s)
  WHERE oqmif.exposure_data.site && %s
    AND taxonomy = %s AND exposure_model_id = %s
  ORDER BY oqmif.exposure_data.id,
    ST_Distance(oqmif.exposure_data.site, gmf_table.location, false)
           """

        assets_extent = self._assets_mesh.get_convex_hull()
        args = (self._imt, self.hazard_id,
                assets_extent.dilate(self.max_distance).wkt,
                self.max_distance * KILOMETERS_TO_METERS,
                assets_extent.wkt,
                self.assets[0].taxonomy,
                self.assets[0].exposure_model_id)

        cursor.execute(query, args)

        return OrderedDict(cursor.fetchall())
<commit_msg>massaging codebase<commit_after># -*- coding: utf-8 -*-

# 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/>.

"""
Hazard getters for Risk calculators.

An HazardGetter is responsible to get hazard outputs needed by a risk
calculation.
"""

from collections import OrderedDict
from openquake.engine import logs
from openquake.hazardlib import geo
from openquake.engine.db import models
from django.db import connection

#: Scaling constant do adapt to the postgis functions (that work with
#: meters)
KILOMETERS_TO_METERS = 1000

class HazardGetter(object):
    """
    Base abstract class of an Hazard Getter.

    An Hazard Getter is used to query for the closest hazard data for
    each given asset. An Hazard Getter must be pickable such that it
    should be possible to use different strategies (e.g. distributed
    or not, using postgis or not).

    :attr hazard_id:
      the ID of an Hazard output container (e.g.
      :class:`openquake.engine.db.models.HazardCurve`)

    :attr imt: the imt of the hazard considered by the getter

    :attr assets: the assets for which we wants to compute

    :attr max_distance: the maximum distance, in kilometers, to use
    """
    def __init__(self, hazard_id, imt, assets, max_distance):
        self.hazard_id = hazard_id
        self.imt = imt
        self.assets = assets
        self.max_distance = max_distance

        self._assets_mesh = geo.mesh.Mesh.from_points_list([
            geo.point.Point(asset.site.x, asset.site.y)
            for asset in self.assets])
        self._imt, self._sa_period, self._sa_damping = (
            models.parse_imt(self.imt))
        self.asset_dict = dict((asset.id, asset) for asset in self.assets)
        self._cache = {}

    def get_data(self):
        """
        :returns: an OrderedDict mapping ID of
        :class:`openquake.engine.db.models.ExposureData` objects to
        hazard_data (e.g. an array with the poes, or an array with the
        ground motion values). Mind that the returned data could lack
        some assets being filtered out by the ``maximum_distance``
        criteria.

        Subclasses must implement this.
        """
        raise NotImplementedError

    def __call__(self):
        """
        :returns: a tuple with three elements. The first is an array
        of instances of
        :class:`openquake.engine.db.models.ExposureData`, the second
        is an array with the corresponding hazard data, the third is
        the array of IDs of assets that has been filtered out by the
        getter by the ``maximum_distance`` criteria.
        """
        data = self.get_data()

        filtered_asset_ids = set(data.keys())
        all_asset_ids = set(self.asset_dict.keys())
        missing_asset_ids = all_asset_ids - filtered_asset_ids
        extra_asset_ids = filtered_asset_ids - all_asset_ids

        # FIXME(lp). It might happens that the convex hull contains
        # more assets than the requested ones. At this moment, we just
        # ignore them.
        if extra_asset_ids:
            logs.LOG.debug("Extra asset have been computed ids: %s" % (
                models.ExposureData.objects.filter(
                    pk__in=extra_asset_ids)))

        for missing_asset_id in missing_asset_ids:
            logs.LOG.warn(
                "No hazard has been found for the asset %s within %s km" % (
                    self.asset_dict[missing_asset_id], self.max_distance))

        return ([self.asset_dict[asset_id] for asset_id in data
                 if asset_id in self.asset_dict],
                [data[asset_id] for asset_id in data
                 if asset_id in self.asset_dict],
                missing_asset_ids)

class HazardCurveGetterPerAsset(HazardGetter):
    """
    Simple HazardCurve Getter that performs a spatial query for each
    asset.

    :attr imls: the intensity measure levels of the curves we are
    going to get.

    :attr dict _cache: a cache of the computed hazard curve object on
    a per-location basis.
    """

    def __init__(self, hazard_id, imt, assets, max_distance):
        super(HazardCurveGetterPerAsset, self).__init__(
            hazard_id, imt, assets, max_distance)
        self.imls = models.HazardCurve.objects.get(
            pk=self.hazard_id).imls

    def get_data(self):
        """
        Calls ``get_by_site`` for each asset and pack the results as
        requested by the :method:`HazardGetter.get_data` interface.
        """
        hazard_assets = [(asset.id, self.get_by_site(asset.site))
                         for asset in self.assets]

        return OrderedDict(
            [(asset_id, hazard_curve)
             for asset_id, (hazard_curve, distance) in hazard_assets
             if distance < self.max_distance * KILOMETERS_TO_METERS])

    def get_by_site(self, site):
        """
        :param site: an instance of
        :class:`django.contrib.gis.geos.point.Point` corresponding to
        the location of an asset.
        """
        if site.wkt in self._cache:
            return self._cache[site.wkt]

        cursor = connection.cursor()

        query = """
        SELECT
            hzrdr.hazard_curve_data.poes,
            min(ST_Distance(location::geography,
                            ST_GeographyFromText(%s), false))
                AS min_distance
        FROM hzrdr.hazard_curve_data
        WHERE hazard_curve_id = %s
        GROUP BY id
        ORDER BY min_distance
        LIMIT 1;"""

        args = (site.wkt, self.hazard_id)

        cursor.execute(query, args)
        poes, distance = cursor.fetchone()

        hazard = zip(self.imls, poes)

        self._cache[site.wkt] = (hazard, distance)

        return hazard, distance

class GroundMotionValuesGetter(HazardGetter):
    """
    Hazard getter for loading ground motion values.
    """

    def get_data(self):
        gmf_collection = models.GmfCollection.objects.get(
            id=self.hazard_id)

        if gmf_collection.output.output_type == "gmf":
            gmf_set_ids = tuple(
                gmf_collection.gmfset_set.values_list('id', flat=True))
        else:
            raise ValueError(
                "Output must be of type `gmf`. "
                "At the moment, we only support computation of loss curves "
                "for a specific logic tree branch.")

        cursor = connection.cursor()

        spectral_filters = ""
        args = (self._imt, gmf_set_ids)

        if self._imt == "SA":
            spectral_filters = "AND sa_period = %s AND sa_damping = %s"
            args += (self._sa_period, self._sa_damping)

        # Query explanation. We need to get for each asset the closest
        # ground motion values (and the corresponding rupture ids from
        # which they have been generated) for a given logic tree
        # realization and a given imt.

        # We first concatenate ground motion values grouped by
        # location (so for each location we have all the ground motion
        # values found in different gmf_sets, collected in a column
        # called ``allgmvs_arr``).

        # For performance reasons, we help the query by filtering also
        # on a polygon built by dilating the assets extent of the
        # maximum distance

        # Then, we perform a spatial join with the exposure table that
        # is previously filtered by the assets extent, exposure model
        # and taxonomy. We are not filtering with an IN statement on
        # the ids of the assets for perfomance reasons.

        # The ``distinct ON (exposure_data.id)`` combined by the
        # ``ORDER BY ST_Distance`` does the job to select the closest
        # gmvs
        query = """
  SELECT DISTINCT ON (oqmif.exposure_data.id)
  oqmif.exposure_data.id, gmf_table.allgmvs_arr, gmf_table.allrupture_ids
  FROM oqmif.exposure_data JOIN
    (SELECT location,
            array_concat(gmvs ORDER BY gmf_set_id, result_grp_ordinal)
            AS allgmvs_arr,
            array_concat(rupture_ids ORDER BY gmf_set_id, result_grp_ordinal)
            AS allrupture_ids
     FROM hzrdr.gmf
     WHERE imt = %s AND gmf_set_id IN %s {}
     AND location && %s
     GROUP BY location) AS gmf_table
  ON ST_DWithin(oqmif.exposure_data.site, gmf_table.location, %s)
  WHERE oqmif.exposure_data.site && %s
  AND taxonomy = %s AND exposure_model_id = %s
  AND array_length(gmf_table.allgmvs_arr, 1) > 0
  ORDER BY oqmif.exposure_data.id,
           ST_Distance(oqmif.exposure_data.site, gmf_table.location, false)
           """.format(spectral_filters)  # this will fill in the {}

        assets_extent = self._assets_mesh.get_convex_hull()
        args += (assets_extent.dilate(self.max_distance).wkt,
                 self.max_distance * KILOMETERS_TO_METERS,
                 assets_extent.wkt,
                 self.assets[0].taxonomy,
                 self.assets[0].exposure_model_id)

        cursor.execute(query, args)

        data = cursor.fetchall()

        # nested dicts with structure: asset_id -> (rupture_id -> gmv)
        assets_ruptures_gmvs = OrderedDict()

        # store all the ruptures returned by the query
        ruptures = []

        for asset_id, gmvs, rupture_ids in data:
            assets_ruptures_gmvs[asset_id] = dict(zip(rupture_ids, gmvs))
            ruptures.extend(rupture_ids)
        ruptures = set(ruptures)

        # We expect that the query may return a different number of
        # gmvs and ruptures for each asset (because only the ruptures
        # that gives a positive ground shaking are stored). Here on,
        # we finalize `assets_ruptures_gmvs` by filling in with zero
        # values for each rupture that has not given a contribute.

        # for each asset, we look for missing ruptures
        for asset_id, ruptures_gmvs_dict in assets_ruptures_gmvs.items():

            # all the ruptures producing a positive ground shaking for
            # `asset`
            asset_ruptures = set(ruptures_gmvs_dict)

            missing_ruptures = ruptures - asset_ruptures

            # we finalize the asset data with 0
            for rupture_id in missing_ruptures:
                ruptures_gmvs_dict[rupture_id] = 0.

        # maps asset_id -> to a 2-tuple (rupture_ids, gmvs)
        return OrderedDict([
            (asset_id, zip(
                *sorted(zip(asset_data.values(), asset_data.keys()),
                        key=lambda x: x[1])))
            for asset_id, asset_data in assets_ruptures_gmvs.items()])

class GroundMotionScenarioGetter(HazardGetter):
    """
    Hazard getter for loading ground motion values. It uses the same
    approach used in :class:`GroundMotionValuesGetter`.
    """

    def setup(self):
        super(GroundMotionScenarioGetter, self).setup()
        if self._sa_period:
            self._imt = '%s(%s)' % (self._imt, self._sa_period)

    def get_data(self):
        cursor = connection.cursor()

        # See the comment in `GroundMotionValuesGetter.get_data` for
        # an explanation of the query
        query = """
  SELECT DISTINCT ON (oqmif.exposure_data.id) oqmif.exposure_data.id,
         gmf_table.gmvs
  FROM oqmif.exposure_data JOIN (
    SELECT location, gmvs
           FROM hzrdr.gmf_scenario
           WHERE hzrdr.gmf_scenario.imt = %s
           AND hzrdr.gmf_scenario.output_id = %s
           AND hzrdr.gmf_scenario.location && %s) gmf_table
  ON ST_DWithin(oqmif.exposure_data.site, gmf_table.location, %s)
  WHERE oqmif.exposure_data.site && %s
    AND taxonomy = %s AND exposure_model_id = %s
  ORDER BY oqmif.exposure_data.id,
    ST_Distance(oqmif.exposure_data.site, gmf_table.location, false)
           """

        assets_extent = self._assets_mesh.get_convex_hull()
        args = (self._imt, self.hazard_id,
                assets_extent.dilate(self.max_distance).wkt,
                self.max_distance * KILOMETERS_TO_METERS,
                assets_extent.wkt,
                self.assets[0].taxonomy,
                self.assets[0].exposure_model_id)

        cursor.execute(query, args)

        return OrderedDict(cursor.fetchall())
<|endoftext|>


<commit_before>import os.path
import unittest
import numpy
from openquake.hazardlib import geo, imt
from openquake.hazardlib.shakemap import (
    get_shakemap_array, get_sitecol_shakemap, to_gmfs, amplify_ground_shaking,
    spatial_correlation_array, spatial_covariance_array,
    cross_correlation_matrix, cholesky)

aae = numpy.testing.assert_almost_equal
F64 = numpy.float64
imts = [imt.from_string(x) for x in ['PGA', 'SA(0.3)', 'SA(1.0)', 'SA(3.0)']]
imt_dt = numpy.dtype([(str(imt), float) for imt in imts])
shakemap_dt = numpy.dtype([('lon', float), ('lat', float), ('val', imt_dt),
                           ('std', imt_dt), ('vs30', float)])
CDIR = os.path.dirname(__file__)

def mean_std(shakemap, site_effects):
    _, gmfs = to_gmfs(
        shakemap, 'yes', 'yes', site_effects, trunclevel=3,
        num_gmfs=1000, seed=42)
    return gmfs.mean(axis=1), numpy.log(gmfs).std(axis=1)

class ShakemapTestCase(unittest.TestCase):

    def test_gmfs(self):
        f1 = os.path.join(CDIR, 'ghorka_grid.xml')
        f2 = os.path.join(CDIR, 'ghorka_uncertainty.xml')
        array = get_shakemap_array(f1, f2)
        sitecol, shakemap = get_sitecol_shakemap(array, imt_dt.names)
        n = 4  # number of sites
        self.assertEqual(len(sitecol), n)
        gmf_by_imt, _ = mean_std(shakemap, site_effects=True)
        aae(gmf_by_imt, [[0.005391, 0.0223217, 0.0399937, 0.0183143],
                         [0.0061, 0.025619, 0.0487997, 0.0225788],
                         [0.0060717, 0.0253156, 0.0478506, 0.0219296],
                         [0.007087, 0.0298716, 0.0622145, 0.0290721]])

    def test_amplify(self):
        gmvs = numpy.array([0.1, 0.2, 0.3])
        res = amplify_ground_shaking(T=3.0, vs30=780, gmvs=gmvs)
        aae(res, [0.09832577, 0.19690711, 0.2958982])

        res = amplify_ground_shaking(T=0.3, vs30=780, gmvs=gmvs)
        aae(res, [0.09909498, 0.19870543, 0.29922175])

    def test_matrices(self):

        # distance matrix
        lons = numpy.array([84., 84., 84., 85.5, 85.5, 85.5, 87., 87., 87.])
        lats = numpy.array([26., 27.5, 29., 26., 27.5, 29., 26., 27.5, 29.])
        dmatrix = geo.geodetic.distance_matrix(lons, lats)
        aae(dmatrix.sum(), 18539.6131407)

        # spatial correlation
        sca = spatial_correlation_array(dmatrix, imts, 'yes')
        aae(sca.sum(), 36.000370229)

        # spatial covariance
        std = numpy.array([(0.5, 0.52, 0.64, 0.73)] * 9, imt_dt)  # 9 sites
        scov = spatial_covariance_array([std[n] for n in imt_dt.names], sca)
        aae(scov.sum(), 13.166200147)

        # cross correlation
        ccor = cross_correlation_matrix(imts, 'yes')
        aae(ccor.sum(), 10.49124788)

        # cholesky decomposition
        L = cholesky(scov, ccor)
        self.assertEqual(L.shape, (36, 36))
        aae(L.sum(), 30.5121263)

        # intensity
        val = numpy.array(
            [(5.38409665, 3.9383686, 3.55435415, 4.37692394)] * 9, imt_dt)

        shakemap = numpy.zeros(9, shakemap_dt)  # 9 sites
        shakemap['lon'] = lons
        shakemap['lat'] = lats
        shakemap['vs30'] = numpy.array([301.17] * 9)
        shakemap['val'] = val
        shakemap['std'] = std
        _, gmfs = to_gmfs(
            shakemap, 'yes', 'no', site_effects=False, trunclevel=3,
            num_gmfs=2, seed=42)
        # shape (N, E, M)
        aae(gmfs[..., 0].sum(axis=0), [0.3708301, 0.5671011])  # PGA

        _, gmfs = to_gmfs(
            shakemap, 'yes', 'yes', site_effects=True, trunclevel=3,
            num_gmfs=2, seed=42)
        aae(gmfs[..., 0].sum(axis=0), [0.5127146, 0.7800232])  # PGA
        aae(gmfs[..., 2].sum(axis=0), [0.4932519, 0.6731384])  # SA(1.0)

        # disable spatial correlation
        _, gmfs = to_gmfs(
            shakemap, 'no', 'no', site_effects=False,
            trunclevel=3, num_gmfs=2, seed=42)
        # shape (N, E, M)
        aae(gmfs[..., 0].sum(axis=0), [0.370832, 0.5670994])  # PGA

        _, gmfs = to_gmfs(
            shakemap, 'no', 'yes', site_effects=True,
            trunclevel=3, num_gmfs=2, seed=42)
        aae(gmfs[..., 0].sum(axis=0), [0.5127171, 0.7800206])  # PGA
        aae(gmfs[..., 2].sum(axis=0), [0.4932519, 0.6731384])  # SA(1.0)

    def test_from_files(self):
        # files provided by Vitor Silva, without site amplification
        f1 = os.path.join(CDIR, 'test_shaking.xml')
        f2 = os.path.join(CDIR, 'test_uncertainty.xml')
        array = get_shakemap_array(f1, f2)
        sitecol, shakemap = get_sitecol_shakemap(array, imt_dt.names)
        n = 4  # number of sites
        self.assertEqual(len(sitecol), n)
        gmf_by_imt, std_by_imt = mean_std(shakemap, site_effects=False)
        # PGA, SA(0.3), SA(1.0), SA(3.0)
        aae(gmf_by_imt, [[0.0975815, 0.2442196, 0.0286512, 0.6358019],
                         [0.2023841, 0.5013746, 0.0297236, 0.6544367],
                         [0.3010831, 0.5986038, 0.0305651, 0.6575208],
                         [0.386838, 0.9331248, 0.0296789, 0.6393688]])
        aae(std_by_imt, [[0.592238, 0.672398, 0.6325073, 0.6445988],
                         [0.6077153, 0.6661571, 0.6296381, 0.668559],
                         [0.6146356, 0.674883, 0.6714424, 0.6613612],
                         [0.5815353, 0.6460007, 0.6491335, 0.6603457]])
<commit_msg>[skip CI]<commit_after>import os.path
import unittest
import numpy
from openquake.hazardlib import geo, imt
from openquake.hazardlib.shakemap import (
    get_shakemap_array, get_sitecol_shakemap, to_gmfs, amplify_ground_shaking,
    spatial_correlation_array, spatial_covariance_array,
    cross_correlation_matrix, cholesky)

aae = numpy.testing.assert_almost_equal
F64 = numpy.float64
imts = [imt.from_string(x) for x in ['PGA', 'SA(0.3)', 'SA(1.0)', 'SA(3.0)']]
imt_dt = numpy.dtype([(str(imt), float) for imt in imts])
shakemap_dt = numpy.dtype([('lon', float), ('lat', float), ('val', imt_dt),
                           ('std', imt_dt), ('vs30', float)])
CDIR = os.path.dirname(__file__)

def mean_std(shakemap, site_effects):
    _, gmfs = to_gmfs(
        shakemap, 'yes', 'yes', site_effects, trunclevel=3,
        num_gmfs=1000, seed=42)
    return gmfs.mean(axis=1), numpy.log(gmfs).std(axis=1)

class ShakemapTestCase(unittest.TestCase):

    def test_gmfs(self):
        f1 = os.path.join(CDIR, 'ghorka_grid.xml')
        f2 = os.path.join(CDIR, 'ghorka_uncertainty.xml')
        array = get_shakemap_array(f1, f2)
        sitecol, shakemap = get_sitecol_shakemap(array, imt_dt.names)
        n = 4  # number of sites
        self.assertEqual(len(sitecol), n)
        gmf_by_imt, _ = mean_std(shakemap, site_effects=True)
        aae(gmf_by_imt, [[0.005391, 0.0223217, 0.0399937, 0.0183143],
                         [0.0061, 0.025619, 0.0487997, 0.0225788],
                         [0.0060717, 0.0253156, 0.0478506, 0.0219296],
                         [0.007087, 0.0298716, 0.0622145, 0.0290721]])

    def test_amplify(self):
        gmvs = numpy.array([0.1, 0.2, 0.3])
        res = amplify_ground_shaking(T=3.0, vs30=780, gmvs=gmvs)
        aae(res, [0.09832577, 0.19690711, 0.2958982])

        res = amplify_ground_shaking(T=0.3, vs30=780, gmvs=gmvs)
        aae(res, [0.09909498, 0.19870543, 0.29922175])

    def test_matrices(self):

        # distance matrix
        lons = numpy.array([84., 84., 84., 85.5, 85.5, 85.5, 87., 87., 87.])
        lats = numpy.array([26., 27.5, 29., 26., 27.5, 29., 26., 27.5, 29.])
        dmatrix = geo.geodetic.distance_matrix(lons, lats)
        aae(dmatrix.sum(), 18539.6131407)

        # spatial correlation
        sca = spatial_correlation_array(dmatrix, imts, 'yes')
        aae(sca.sum(), 36.000370229)

        # spatial covariance
        std = numpy.array([(0.5, 0.52, 0.64, 0.73)] * 9, imt_dt)  # 9 sites
        scov = spatial_covariance_array([std[n] for n in imt_dt.names], sca)
        aae(scov.sum(), 13.166200147)

        # cross correlation
        ccor = cross_correlation_matrix(imts, 'yes')
        aae(ccor.sum(), 10.49124788)

        # cholesky decomposition
        L = cholesky(scov, ccor)
        self.assertEqual(L.shape, (36, 36))
        aae(L.sum(), 30.5121263)

        # intensity
        val = numpy.array(
            [(5.38409665, 3.9383686, 3.55435415, 4.37692394)] * 9, imt_dt)

        shakemap = numpy.zeros(9, shakemap_dt)  # 9 sites
        shakemap['lon'] = lons
        shakemap['lat'] = lats
        shakemap['vs30'] = numpy.array([301.17] * 9)
        shakemap['val'] = val
        shakemap['std'] = std
        _, gmfs = to_gmfs(
            shakemap, 'yes', 'no', site_effects=False, trunclevel=3,
            num_gmfs=2, seed=42)
        # shape (N, E, M)
        aae(gmfs[..., 0].sum(axis=0), [0.3708301, 0.5671011])  # PGA

        _, gmfs = to_gmfs(
            shakemap, 'yes', 'yes', site_effects=True, trunclevel=3,
            num_gmfs=2, seed=42)
        aae(gmfs[..., 0].sum(axis=0), [0.5127146, 0.7800232])  # PGA
        aae(gmfs[..., 2].sum(axis=0), [0.4932519, 0.6731384])  # SA(1.0)

        # disable spatial correlation
        _, gmfs = to_gmfs(
            shakemap, 'no', 'no', site_effects=False,
            trunclevel=3, num_gmfs=2, seed=42)
        # shape (N, E, M)
        aae(gmfs[..., 0].sum(axis=0), [0.370832, 0.5670994])  # PGA

        _, gmfs = to_gmfs(
            shakemap, 'no', 'yes', site_effects=True,
            trunclevel=3, num_gmfs=2, seed=42)
        aae(gmfs[..., 0].sum(axis=0), [0.5127171, 0.7800206])  # PGA
        aae(gmfs[..., 2].sum(axis=0), [0.4932519, 0.6731384])  # SA(1.0)

    def test_from_files(self):
        # files provided by Vitor Silva, without site amplification
        f1 = os.path.join(CDIR, 'test_shaking.xml')
        f2 = os.path.join(CDIR, 'test_uncertainty.xml')
        array = get_shakemap_array(f1, f2)
        sitecol, shakemap = get_sitecol_shakemap(array, imt_dt.names)
        n = 4  # number of sites
        self.assertEqual(len(sitecol), n)
        gmf_by_imt, std_by_imt = mean_std(shakemap, site_effects=False)
        #                 PGA,       SA(0.3),   SA(1.0),   SA(3.0)
        aae(gmf_by_imt, [[0.0975815, 0.2442196, 0.0286512, 0.6358019],
                         [0.2023841, 0.5013746, 0.0297236, 0.6544367],
                         [0.3010831, 0.5986038, 0.0305651, 0.6575208],
                         [0.3868380, 0.9331248, 0.0296789, 0.6393688]])
        aae(std_by_imt, [[0.5922380, 0.6723980, 0.6325073, 0.6445988],
                         [0.6077153, 0.6661571, 0.6296381, 0.668559],
                         [0.6146356, 0.6748830, 0.6714424, 0.6613612],
                         [0.5815353, 0.6460007, 0.6491335, 0.6603457]])
<|endoftext|>


<commit_before># vim: expandtab:tabstop=4:shiftwidth=4

"""This module comprises Aws specific utility functions."""

import os
import re
import subprocess
from distutils.version import LooseVersion
import ruamel.yaml as yaml

# Buildbot does not have multi_inventory installed
#pylint: disable=no-name-in-module
from openshift_tools.inventory_clients import multi_inventory

class ArgumentError(Exception):
    """This class is raised when improper arguments are passed."""

    def __init__(self, message):
        """Initialize an ArgumentError.

        Keyword arguments:
        message -- the exact error message being raised
        """
        super(ArgumentError, self).__init__()
        self.message = message

# pylint: disable=too-many-public-methods
class InventoryUtil(object):
    """This class contains the Inventory utility functions."""

    def __init__(self, host_type_aliases=None, use_cache=True):
        """Initialize the Inventory utility class.

        Keyword arguments:
        host_type_aliases -- a list of aliases to common host-types (e.g. ex-node)
        use_cache -- rely on cached inventory instead of querying for new inventory
        """

        self.cached = use_cache
        self.alias_lookup = {}
        host_type_aliases = host_type_aliases or {}

        self.host_type_aliases = host_type_aliases
        self.file_path = os.path.join(os.path.dirname(os.path.realpath(__file__)))

        self.setup_host_type_alias_lookup()
        self._inventory = None

    @property
    def inventory(self):
        """ Sets up a class property named inventory, this is the getter
        It will pull in the file if the property is empty, otherwise just returns the variable
        """
        if self._inventory is None:
            self._inventory = multi_inventory.MultiInventory(None).run()
        return self._inventory

    @staticmethod
    def get_cluster(name):
        """ return a cluster object """

        return Cluster(name)

    def setup_host_type_alias_lookup(self):
        """Sets up the alias to host-type lookup table."""
        for key, values in self.host_type_aliases.iteritems():
            for value in values:
                self.alias_lookup[value] = key

    def _get_tags_(self, regex):
        """ Searches for tags in the inventory and returns all of the tags
            found.

            Param: a compiled regular expression
            Returns: a List of tags
        """

        tags = []
        for key in self.inventory:
            matched = regex.match(key)
            if matched:
                tags.append(matched.group(1))

        tags.sort()
        return tags

    def get_clusters(self):
        """Searches for cluster tags in the inventory and returns all of the clusters found."""
        pattern = re.compile(r'^oo_clusterid_(.*)')
        return self._get_tags_(pattern)

    def get_environments(self):
        """Searches for env tags in the inventory and returns all of the envs found."""
        pattern = re.compile(r'^oo_environment_(.*)')
        return self._get_tags_(pattern)

    def get_host_types(self):
        """Searches for host-type tags in the inventory and returns all host-types found."""
        pattern = re.compile(r'^oo_hosttype_(.*)')
        return self._get_tags_(pattern)

    def get_sub_host_types(self):
        """Searches for sub-host-type tags in the inventory and returns all sub-host-types found."""
        pattern = re.compile(r'^oo_subhosttype_(.*)')
        return self._get_tags_(pattern)

    def get_security_groups(self):
        """Searches for security_groups in the inventory and returns all SGs found."""
        pattern = re.compile(r'^security_group_(.*)')
        return self._get_tags_(pattern)

    def build_host_dict_by_env(self, args=None):
        """Searches the inventory for hosts in an env and returns their hostvars."""
        args = args or []

        inst_by_env = {}
        for _, host in self.inventory['_meta']['hostvars'].items():
            # If you don't have an environment tag, we're going to ignore you
            if 'oo_environment' not in host:
                continue

            if host['oo_environment'] not in inst_by_env:
                inst_by_env[host['oo_environment']] = {}
            host_id = "%s:%s" % (host['oo_name'], host['oo_id'])
            inst_by_env[host['oo_environment']][host_id] = host

        return inst_by_env

    def print_host_types(self):
        """Gets the list of host types and aliases and outputs them in columns."""

        host_types = self.get_host_types()
        ht_format_str = "%35s"
        alias_format_str = "%-20s"
        combined_format_str = ht_format_str + "    " + alias_format_str

        print
        print combined_format_str % ('Host Types', 'Aliases')
        print combined_format_str % ('----------', '-------')

        for host_type in host_types:
            aliases = []
            if host_type in self.host_type_aliases:
                aliases = self.host_type_aliases[host_type]
                print combined_format_str % (host_type, ", ".join(aliases))
            else:
                print  ht_format_str % host_type
        print

    def print_cluster_list(self):
        """Gets the list of clusters and outputs them"""
        clusters = self.get_clusters()

        for cluster in clusters:
            print cluster

    def resolve_host_type(self, host_type):
        """Converts a host-type alias into a host-type.

        Keyword arguments:
        host_type -- The alias or host_type to look up.

        Example (depends on aliases defined in config file):
            host_type = ex-node
            returns: openshift-node
        """
        if self.alias_lookup.has_key(host_type):
            return self.alias_lookup[host_type]
        return host_type

    @staticmethod
    def gen_version_tag(ver):
        """Generate the version tag
        """
        return "oo_version_%s" % ver

    @staticmethod
    def gen_clusterid_tag(clu):
        """Generate the clusterid tag
        """
        return "oo_clusterid_%s" % clu

    @staticmethod
    def gen_env_tag(env):
        """Generate the environment tag
        """
        return "oo_environment_%s" % env

    def gen_host_type_tag(self, host_type, version):
        """Generate the host type tag
        """
        if version == '2':
            host_type = self.resolve_host_type(host_type)
        return "oo_hosttype_%s" % host_type

    @staticmethod
    def gen_sub_host_type_tag(sub_host_type):
        """Generate the host type tag
        """
        return "oo_subhosttype_%s" % sub_host_type

    # This function uses all of these params to perform a filters on our host inventory.
    # pylint: disable=too-many-arguments
    def get_host_list(self, clusters=None, host_type=None, sub_host_type=None, envs=None, version=None):
        """Get the list of hosts from the inventory using host-type and environment
        """
        retval = set([])
        envs = envs or []

        retval.update(self.inventory.get('all_hosts', []))

        if clusters:
            cluster_hosts = set([])
            if len(clusters) > 1:
                for cluster in clusters:
                    clu_tag = InventoryUtil.gen_clusterid_tag(cluster)
                    cluster_hosts.update(self.inventory.get(clu_tag, []))
            else:
                cluster_hosts.update(self.inventory.get(InventoryUtil.gen_clusterid_tag(clusters[0]), []))

            retval.intersection_update(cluster_hosts)

        if envs:
            env_hosts = set([])
            if len(envs) > 1:
                for env in envs:
                    env_tag = InventoryUtil.gen_env_tag(env)
                    env_hosts.update(self.inventory.get(env_tag, []))
            else:
                env_hosts.update(self.inventory.get(InventoryUtil.gen_env_tag(envs[0]), []))

            retval.intersection_update(env_hosts)

        if host_type:
            retval.intersection_update(self.inventory.get(self.gen_host_type_tag(host_type, version), []))

        if sub_host_type:
            retval.intersection_update(self.inventory.get(self.gen_sub_host_type_tag(sub_host_type), []))

        if version != 'all':
            retval.intersection_update(self.inventory.get(InventoryUtil.gen_version_tag(version), []))

        return list(retval)

    def convert_to_ip(self, hosts):
        """convert a list of host names to ip addresses"""

        if not isinstance(hosts, list):
            hosts = [hosts]

        ips = []
        for host in hosts:
            ips.append(self.inventory['_meta']['hostvars'][host]['oo_public_ip'])

        return ips

    @staticmethod
    def get_cluster_variable(cluster, variable):
        """ return an inventory variable that is common to a cluster"""

        cluster = InventoryUtil.get_cluster(cluster)

        return cluster.get_variable(variable)

    def get_node_variable(self, host, variable):
        """ return an inventory variable from a host"""

        if host in self.inventory['_meta']['hostvars'] and variable in self.inventory['_meta']['hostvars'][host]:
            return self.inventory['_meta']['hostvars'][host][variable]

        return None

class Cluster(object):
    """ This is a class to acces data about an Ops cluster """

    def __init__(self, name):
        """ Init the cluster class """

        self._name = name
        self._openshift_version = None
        self._docker_version = None
        self.inventory = multi_inventory.MultiInventory(None).run()
        self._master_config = None

    def __str__(self):
        """ str representation of Cluster """

        return self._name

    def __repr__(self):
        """ repr representation of Cluster """

        return self._name

    @property
    def name(self):
        """ cluster name property """

        return self._name

    @property
    def location(self):
        """ cluster location property """

        return self.get_variable('oo_location')

    @property
    def sublocation(self):
        """ cluster sublocation property """

        return self.get_variable('oo_sublocation')

    @property
    def environment(self):
        """ cluster environment property """

        return self.get_variable('oo_environment')

    @property
    def deployment(self):
        """ cluster deployment property """

        return self.get_variable('oo_deployment')

    @property
    def account(self):
        """ cluster account property """

        return self.get_variable('oo_account')

    @property
    def accountid(self):
        """ cluster account property """

        return self.get_variable('oo_accountid')

    @property
    def test_cluster(self):
        """ cluster cluster property """

        return bool(self.get_variable('oo_test_cluster'))

    @property
    def primary_master(self):
        """ return the first master """

        primary_master = list(set(self.inventory["oo_master_primary"]) &
                              set(self.inventory["oo_clusterid_" + self._name]))[0]

        return primary_master

    @property
    def cluster_nodes(self):
        """ return the number of nodes - infra and compute """

        cluster_nodes = self.inventory["oo_clusterid_" + self._name]

        return cluster_nodes

    @property
    def node_count(self):
        """ return the number of nodes - infra and compute """

        cluster_compute_nodes = list(set(self.inventory["oo_hosttype_node"]) &
                             set(self.inventory["oo_clusterid_" + self._name]))

        return len(cluster_compute_nodes)

    @property
    def scalegroup_node_count(self):
        """ return the number of scalegroup nodes - infra and compute """

        sg_cluster_nodes = list(set(self.inventory["oo_hosttype_node"]) &
                                set(self.inventory["oo_scalegroup_True"]) &
                                set(self.inventory["oo_clusterid_" + self._name]))

        return len(sg_cluster_nodes)

    @property
    def master_config(self):
        """ This is the master_config.yaml file stored as a dict """

        if not self._master_config:
            master_config_yaml = self.run_cmd_on_master("/usr/bin/cat /etc/origin/master/master-config.yaml",
                                                        strip=False)
            self._master_config = yaml.load(master_config_yaml)

        return self._master_config

    @staticmethod
    def set_version(version):
        """ manipulate the version variable """

        os_version = {}
        version_split = version.split('.')

        os_version['version_release'] = version
        os_version['version'] = version.split('-')[0]
        os_version['short'] = '.'.join(version_split[0:2])
        os_version['short_underscore'] = os_version['short'].replace(".", "_")
        os_version['release'] = version.split('-')[1]
        os_version['version_release_no_git'] = version.split('.git')[0]

        # Let's do the wierdness to set full version begin
        if LooseVersion(os_version['short']) < LooseVersion('3.6'):
            os_version['full'] = os_version['short']
        elif os_version['short'] == '3.6':
            os_version['full'] = os_version['version']
        else:
            if "-0.git" in os_version['version_release']:
                os_version['full'] = os_version['version_release_no_git']
            else:
                os_version['full'] = os_version['version']

        return os_version

    @property
    def openshift_version(self):
        """ return a dict of openshift_version """

        if not self._openshift_version:
            self._openshift_version = {}
            version = self.run_cmd_on_master("rpm -q --queryformat '%{VERSION}-%{RELEASE}' \
                                                           atomic-openshift")
            self._openshift_version = Cluster.set_version(version)

        return self._openshift_version

    def get_variable(self, variable):
        """ return an inventory variable that is common to a cluster"""

        variables = []
        for host in self.inventory['oo_clusterid_' + self.name]:
            if variable in self.inventory['_meta']['hostvars'][host]:
                variables.append(self.inventory['_meta']['hostvars'][host][variable])

        if len(list(set(variables))) == 1:
            return variables[0]

        return None

    def run_cmd_on_master(self, command, strip=True):
        """
            Run a command on the primary master and return the output
            command: string of the command to run on the primary master
            e.g: command = "echo 'hello world'"
        """

        command_output = subprocess.check_output(["/usr/bin/ossh", "root@" + self.primary_master, "-c", command])

        if strip:
            return command_output.strip()

        return command_output

    def convert_inv_to_os_name(self, hostname):
        """ convert ops hostname to the openshift hostname
            example: free-stg-master-03fb6 -> ip-172-31-78-254.us-east-2.compute.internal
        """

        if hostname not in self.inventory["_meta"]["hostvars"]:
            return None

        host_vars = self.inventory["_meta"]["hostvars"][hostname]

        if host_vars['oo_location'] == 'gcp':
            return hostname
        elif host_vars['oo_location'] == 'aws':
            return host_vars['ec2_private_dns_name']

    def convert_list_to_os_names(self, hostnames):
        """ convert a list of ops hostname to the openshift hostname
            example: ['free-stg-node-infra-a18ed', 'free-stg-node-compute-006ef'] ->
                     [u'ip-172-31-73-38.us-east-2.compute.internal',
                      u'ip-172-31-74-247.us-east-2.compute.internal']
        """

        if hostnames:
            converted_hosts = []

            for host in hostnames:
                os_host = self.convert_inv_to_os_name(host)
                converted_hosts.append(os_host)

            return converted_hosts

    def convert_os_to_inv_name(self, hostname):
        """ convert openshift name  to inventory name:
            example: 'ip-172-31-69-53.us-east-2.compute.internal' => free-stg-node-infra-70a4e
        """

        if hostname in self.cluster_nodes:
            return hostname

        for node in self.cluster_nodes:
            if hostname == self.inventory["_meta"]["hostvars"][node]["ec2_private_dns_name"]:
                return node

        return None
<commit_msg>Adding vfull back in -- accidentally removed?<commit_after># vim: expandtab:tabstop=4:shiftwidth=4

"""This module comprises Aws specific utility functions."""

import os
import re
import subprocess
from distutils.version import LooseVersion
import ruamel.yaml as yaml

# Buildbot does not have multi_inventory installed
#pylint: disable=no-name-in-module
from openshift_tools.inventory_clients import multi_inventory

class ArgumentError(Exception):
    """This class is raised when improper arguments are passed."""

    def __init__(self, message):
        """Initialize an ArgumentError.

        Keyword arguments:
        message -- the exact error message being raised
        """
        super(ArgumentError, self).__init__()
        self.message = message

# pylint: disable=too-many-public-methods
class InventoryUtil(object):
    """This class contains the Inventory utility functions."""

    def __init__(self, host_type_aliases=None, use_cache=True):
        """Initialize the Inventory utility class.

        Keyword arguments:
        host_type_aliases -- a list of aliases to common host-types (e.g. ex-node)
        use_cache -- rely on cached inventory instead of querying for new inventory
        """

        self.cached = use_cache
        self.alias_lookup = {}
        host_type_aliases = host_type_aliases or {}

        self.host_type_aliases = host_type_aliases
        self.file_path = os.path.join(os.path.dirname(os.path.realpath(__file__)))

        self.setup_host_type_alias_lookup()
        self._inventory = None

    @property
    def inventory(self):
        """ Sets up a class property named inventory, this is the getter
        It will pull in the file if the property is empty, otherwise just returns the variable
        """
        if self._inventory is None:
            self._inventory = multi_inventory.MultiInventory(None).run()
        return self._inventory

    @staticmethod
    def get_cluster(name):
        """ return a cluster object """

        return Cluster(name)

    def setup_host_type_alias_lookup(self):
        """Sets up the alias to host-type lookup table."""
        for key, values in self.host_type_aliases.iteritems():
            for value in values:
                self.alias_lookup[value] = key

    def _get_tags_(self, regex):
        """ Searches for tags in the inventory and returns all of the tags
            found.

            Param: a compiled regular expression
            Returns: a List of tags
        """

        tags = []
        for key in self.inventory:
            matched = regex.match(key)
            if matched:
                tags.append(matched.group(1))

        tags.sort()
        return tags

    def get_clusters(self):
        """Searches for cluster tags in the inventory and returns all of the clusters found."""
        pattern = re.compile(r'^oo_clusterid_(.*)')
        return self._get_tags_(pattern)

    def get_environments(self):
        """Searches for env tags in the inventory and returns all of the envs found."""
        pattern = re.compile(r'^oo_environment_(.*)')
        return self._get_tags_(pattern)

    def get_host_types(self):
        """Searches for host-type tags in the inventory and returns all host-types found."""
        pattern = re.compile(r'^oo_hosttype_(.*)')
        return self._get_tags_(pattern)

    def get_sub_host_types(self):
        """Searches for sub-host-type tags in the inventory and returns all sub-host-types found."""
        pattern = re.compile(r'^oo_subhosttype_(.*)')
        return self._get_tags_(pattern)

    def get_security_groups(self):
        """Searches for security_groups in the inventory and returns all SGs found."""
        pattern = re.compile(r'^security_group_(.*)')
        return self._get_tags_(pattern)

    def build_host_dict_by_env(self, args=None):
        """Searches the inventory for hosts in an env and returns their hostvars."""
        args = args or []

        inst_by_env = {}
        for _, host in self.inventory['_meta']['hostvars'].items():
            # If you don't have an environment tag, we're going to ignore you
            if 'oo_environment' not in host:
                continue

            if host['oo_environment'] not in inst_by_env:
                inst_by_env[host['oo_environment']] = {}
            host_id = "%s:%s" % (host['oo_name'], host['oo_id'])
            inst_by_env[host['oo_environment']][host_id] = host

        return inst_by_env

    def print_host_types(self):
        """Gets the list of host types and aliases and outputs them in columns."""

        host_types = self.get_host_types()
        ht_format_str = "%35s"
        alias_format_str = "%-20s"
        combined_format_str = ht_format_str + "    " + alias_format_str

        print
        print combined_format_str % ('Host Types', 'Aliases')
        print combined_format_str % ('----------', '-------')

        for host_type in host_types:
            aliases = []
            if host_type in self.host_type_aliases:
                aliases = self.host_type_aliases[host_type]
                print combined_format_str % (host_type, ", ".join(aliases))
            else:
                print  ht_format_str % host_type
        print

    def print_cluster_list(self):
        """Gets the list of clusters and outputs them"""
        clusters = self.get_clusters()

        for cluster in clusters:
            print cluster

    def resolve_host_type(self, host_type):
        """Converts a host-type alias into a host-type.

        Keyword arguments:
        host_type -- The alias or host_type to look up.

        Example (depends on aliases defined in config file):
            host_type = ex-node
            returns: openshift-node
        """
        if self.alias_lookup.has_key(host_type):
            return self.alias_lookup[host_type]
        return host_type

    @staticmethod
    def gen_version_tag(ver):
        """Generate the version tag
        """
        return "oo_version_%s" % ver

    @staticmethod
    def gen_clusterid_tag(clu):
        """Generate the clusterid tag
        """
        return "oo_clusterid_%s" % clu

    @staticmethod
    def gen_env_tag(env):
        """Generate the environment tag
        """
        return "oo_environment_%s" % env

    def gen_host_type_tag(self, host_type, version):
        """Generate the host type tag
        """
        if version == '2':
            host_type = self.resolve_host_type(host_type)
        return "oo_hosttype_%s" % host_type

    @staticmethod
    def gen_sub_host_type_tag(sub_host_type):
        """Generate the host type tag
        """
        return "oo_subhosttype_%s" % sub_host_type

    # This function uses all of these params to perform a filters on our host inventory.
    # pylint: disable=too-many-arguments
    def get_host_list(self, clusters=None, host_type=None, sub_host_type=None, envs=None, version=None):
        """Get the list of hosts from the inventory using host-type and environment
        """
        retval = set([])
        envs = envs or []

        retval.update(self.inventory.get('all_hosts', []))

        if clusters:
            cluster_hosts = set([])
            if len(clusters) > 1:
                for cluster in clusters:
                    clu_tag = InventoryUtil.gen_clusterid_tag(cluster)
                    cluster_hosts.update(self.inventory.get(clu_tag, []))
            else:
                cluster_hosts.update(self.inventory.get(InventoryUtil.gen_clusterid_tag(clusters[0]), []))

            retval.intersection_update(cluster_hosts)

        if envs:
            env_hosts = set([])
            if len(envs) > 1:
                for env in envs:
                    env_tag = InventoryUtil.gen_env_tag(env)
                    env_hosts.update(self.inventory.get(env_tag, []))
            else:
                env_hosts.update(self.inventory.get(InventoryUtil.gen_env_tag(envs[0]), []))

            retval.intersection_update(env_hosts)

        if host_type:
            retval.intersection_update(self.inventory.get(self.gen_host_type_tag(host_type, version), []))

        if sub_host_type:
            retval.intersection_update(self.inventory.get(self.gen_sub_host_type_tag(sub_host_type), []))

        if version != 'all':
            retval.intersection_update(self.inventory.get(InventoryUtil.gen_version_tag(version), []))

        return list(retval)

    def convert_to_ip(self, hosts):
        """convert a list of host names to ip addresses"""

        if not isinstance(hosts, list):
            hosts = [hosts]

        ips = []
        for host in hosts:
            ips.append(self.inventory['_meta']['hostvars'][host]['oo_public_ip'])

        return ips

    @staticmethod
    def get_cluster_variable(cluster, variable):
        """ return an inventory variable that is common to a cluster"""

        cluster = InventoryUtil.get_cluster(cluster)

        return cluster.get_variable(variable)

    def get_node_variable(self, host, variable):
        """ return an inventory variable from a host"""

        if host in self.inventory['_meta']['hostvars'] and variable in self.inventory['_meta']['hostvars'][host]:
            return self.inventory['_meta']['hostvars'][host][variable]

        return None

class Cluster(object):
    """ This is a class to acces data about an Ops cluster """

    def __init__(self, name):
        """ Init the cluster class """

        self._name = name
        self._openshift_version = None
        self._docker_version = None
        self.inventory = multi_inventory.MultiInventory(None).run()
        self._master_config = None

    def __str__(self):
        """ str representation of Cluster """

        return self._name

    def __repr__(self):
        """ repr representation of Cluster """

        return self._name

    @property
    def name(self):
        """ cluster name property """

        return self._name

    @property
    def location(self):
        """ cluster location property """

        return self.get_variable('oo_location')

    @property
    def sublocation(self):
        """ cluster sublocation property """

        return self.get_variable('oo_sublocation')

    @property
    def environment(self):
        """ cluster environment property """

        return self.get_variable('oo_environment')

    @property
    def deployment(self):
        """ cluster deployment property """

        return self.get_variable('oo_deployment')

    @property
    def account(self):
        """ cluster account property """

        return self.get_variable('oo_account')

    @property
    def accountid(self):
        """ cluster account property """

        return self.get_variable('oo_accountid')

    @property
    def test_cluster(self):
        """ cluster cluster property """

        return bool(self.get_variable('oo_test_cluster'))

    @property
    def primary_master(self):
        """ return the first master """

        primary_master = list(set(self.inventory["oo_master_primary"]) &
                              set(self.inventory["oo_clusterid_" + self._name]))[0]

        return primary_master

    @property
    def cluster_nodes(self):
        """ return the number of nodes - infra and compute """

        cluster_nodes = self.inventory["oo_clusterid_" + self._name]

        return cluster_nodes

    @property
    def node_count(self):
        """ return the number of nodes - infra and compute """

        cluster_compute_nodes = list(set(self.inventory["oo_hosttype_node"]) &
                             set(self.inventory["oo_clusterid_" + self._name]))

        return len(cluster_compute_nodes)

    @property
    def scalegroup_node_count(self):
        """ return the number of scalegroup nodes - infra and compute """

        sg_cluster_nodes = list(set(self.inventory["oo_hosttype_node"]) &
                                set(self.inventory["oo_scalegroup_True"]) &
                                set(self.inventory["oo_clusterid_" + self._name]))

        return len(sg_cluster_nodes)

    @property
    def master_config(self):
        """ This is the master_config.yaml file stored as a dict """

        if not self._master_config:
            master_config_yaml = self.run_cmd_on_master("/usr/bin/cat /etc/origin/master/master-config.yaml",
                                                        strip=False)
            self._master_config = yaml.load(master_config_yaml)

        return self._master_config

    @staticmethod
    def set_version(version):
        """ manipulate the version variable """

        os_version = {}
        version_split = version.split('.')

        os_version['version_release'] = version
        os_version['version'] = version.split('-')[0]
        os_version['short'] = '.'.join(version_split[0:2])
        os_version['short_underscore'] = os_version['short'].replace(".", "_")
        os_version['release'] = version.split('-')[1]
        os_version['version_release_no_git'] = version.split('.git')[0]

        # Let's do the wierdness to set full version begin
        if LooseVersion(os_version['short']) < LooseVersion('3.6'):
            os_version['full'] = os_version['short']
        elif os_version['short'] == '3.6':
            os_version['full'] = os_version['version']
        else:
            if "-0.git" in os_version['version_release']:
                os_version['full'] = os_version['version_release_no_git']
            else:
                os_version['full'] = os_version['version']

        os_version['vfull'] = 'v' + os_version['full']

        return os_version

    @property
    def openshift_version(self):
        """ return a dict of openshift_version """

        if not self._openshift_version:
            self._openshift_version = {}
            version = self.run_cmd_on_master("rpm -q --queryformat '%{VERSION}-%{RELEASE}' \
                                                           atomic-openshift")
            self._openshift_version = Cluster.set_version(version)

        return self._openshift_version

    def get_variable(self, variable):
        """ return an inventory variable that is common to a cluster"""

        variables = []
        for host in self.inventory['oo_clusterid_' + self.name]:
            if variable in self.inventory['_meta']['hostvars'][host]:
                variables.append(self.inventory['_meta']['hostvars'][host][variable])

        if len(list(set(variables))) == 1:
            return variables[0]

        return None

    def run_cmd_on_master(self, command, strip=True):
        """
            Run a command on the primary master and return the output
            command: string of the command to run on the primary master
            e.g: command = "echo 'hello world'"
        """

        command_output = subprocess.check_output(["/usr/bin/ossh", "root@" + self.primary_master, "-c", command])

        if strip:
            return command_output.strip()

        return command_output

    def convert_inv_to_os_name(self, hostname):
        """ convert ops hostname to the openshift hostname
            example: free-stg-master-03fb6 -> ip-172-31-78-254.us-east-2.compute.internal
        """

        if hostname not in self.inventory["_meta"]["hostvars"]:
            return None

        host_vars = self.inventory["_meta"]["hostvars"][hostname]

        if host_vars['oo_location'] == 'gcp':
            return hostname
        elif host_vars['oo_location'] == 'aws':
            return host_vars['ec2_private_dns_name']

    def convert_list_to_os_names(self, hostnames):
        """ convert a list of ops hostname to the openshift hostname
            example: ['free-stg-node-infra-a18ed', 'free-stg-node-compute-006ef'] ->
                     [u'ip-172-31-73-38.us-east-2.compute.internal',
                      u'ip-172-31-74-247.us-east-2.compute.internal']
        """

        if hostnames:
            converted_hosts = []

            for host in hostnames:
                os_host = self.convert_inv_to_os_name(host)
                converted_hosts.append(os_host)

            return converted_hosts

    def convert_os_to_inv_name(self, hostname):
        """ convert openshift name  to inventory name:
            example: 'ip-172-31-69-53.us-east-2.compute.internal' => free-stg-node-infra-70a4e
        """

        if hostname in self.cluster_nodes:
            return hostname

        for node in self.cluster_nodes:
            if hostname == self.inventory["_meta"]["hostvars"][node]["ec2_private_dns_name"]:
                return node

        return None
<|endoftext|>