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

"""
Module :mod:`~openquake.hazardlib.calc.gmf` exports
:func:`ground_motion_fields`.
"""
import time
import numpy
import scipy.stats

from openquake.baselib.general import AccumDict
from openquake.baselib.python3compat import decode
from openquake.hazardlib.const import StdDev
from openquake.hazardlib.gsim.base import ContextMaker, FarAwayRupture
from openquake.hazardlib.imt import from_string

U32 = numpy.uint32
F32 = numpy.float32

class CorrelationButNoInterIntraStdDevs(Exception):
    def __init__(self, corr, gsim):
        self.corr = corr
        self.gsim = gsim

    def __str__(self):
        return '''\
You cannot use the correlation model %s with the GSIM %s, \
that defines only the total standard deviation. If you want to use a \
correlation model you have to select a GMPE that provides the inter and \
intra event standard deviations.''' % (
            self.corr.__class__.__name__, self.gsim.__class__.__name__)

def rvs(distribution, *size):
    array = distribution.rvs(size)
    return array

def to_imt_unit_values(vals, imt):
    """
    Exponentiate the values unless the IMT is MMI
    """
    if str(imt) == 'MMI':
        return vals
    return numpy.exp(vals)

class GmfComputer(object):
    """
    Given an earthquake rupture, the ground motion field computer computes
    ground shaking over a set of sites, by randomly sampling a ground
    shaking intensity model.

    :param rupture:
        Rupture to calculate ground motion fields radiated from.

    :param :class:`openquake.hazardlib.site.SiteCollection` sitecol:
        a complete SiteCollection

    :param imts:
        a sorted list of Intensity Measure Type strings

    :param cmaker:
        a :class:`openquake.hazardlib.gsim.base.ContextMaker` instance

    :param truncation_level:
        Float, number of standard deviations for truncation of the intensity
        distribution, or ``None``.

    :param correlation_model:
        Instance of correlation model object. See
        :mod:`openquake.hazardlib.correlation`. Can be ``None``, in which
        case non-correlated ground motion fields are calculated.
        Correlation model is not used if ``truncation_level`` is zero.

    :param amplifier:
        None or an instance of Amplifier
    """
    # The GmfComputer is called from the OpenQuake Engine. In that case
    # the rupture is an higher level containing a
    # :class:`openquake.hazardlib.source.rupture.Rupture` instance as an
    # attribute. Then the `.compute(gsim, num_events, ms)` method is called and
    # a matrix of size (I, N, E) is returned, where I is the number of
    # IMTs, N the number of affected sites and E the number of events. The
    # seed is extracted from the underlying rupture.
    def __init__(self, rupture, sitecol, cmaker, correlation_model=None,
                 amplifier=None, sec_perils=()):
        if len(sitecol) == 0:
            raise ValueError('No sites')
        elif len(cmaker.imtls) == 0:
            raise ValueError('No IMTs')
        elif len(cmaker.gsims) == 0:
            raise ValueError('No GSIMs')
        self.cmaker = cmaker
        self.imts = [from_string(imt) for imt in cmaker.imtls]
        self.cmaker = cmaker
        self.gsims = sorted(cmaker.gsims)
        self.correlation_model = correlation_model
        self.amplifier = amplifier
        self.sec_perils = sec_perils
        # `rupture` is an EBRupture instance in the engine
        if hasattr(rupture, 'source_id'):
            self.ebrupture = rupture
            self.source_id = rupture.source_id  # the underlying source
            rupture = rupture.rupture  # the underlying rupture
        else:  # in the hazardlib tests
            self.source_id = '?'
        self.seed = rupture.rup_id
        ctxs = cmaker.get_ctxs([rupture], sitecol, self.source_id)
        if not ctxs:
            raise FarAwayRupture
        self.ctx = ctxs[0]
        if correlation_model:  # store the filtered sitecol
            self.sites = sitecol.complete.filtered(self.ctx.sids)
        if cmaker.trunclevel is None:
            self.distribution = scipy.stats.norm()
        elif cmaker.trunclevel == 0:
            self.distribution = None
        else:
            assert cmaker.trunclevel > 0, cmaker.trunclevel
            self.distribution = scipy.stats.truncnorm(
                - cmaker.trunclevel, cmaker.trunclevel)

    def compute_all(self, sig_eps=None):
        """
        :returns: (dict with fields eid, sid, gmv_X, ...), dt
        """
        min_iml = self.cmaker.min_iml
        rlzs_by_gsim = self.cmaker.gsims
        t0 = time.time()
        sids = self.ctx.sids
        eids_by_rlz = self.ebrupture.get_eids_by_rlz(rlzs_by_gsim)
        mag = self.ebrupture.rupture.mag
        data = AccumDict(accum=[])
        mean_stds = self.cmaker.get_mean_stds([self.ctx])  # (4, G, M, N)
        for g, (gs, rlzs) in enumerate(rlzs_by_gsim.items()):
            num_events = sum(len(eids_by_rlz[rlz]) for rlz in rlzs)
            if num_events == 0:  # it may happen
                continue
            # NB: the trick for performance is to keep the call to
            # .compute outside of the loop over the realizations;
            # it is better to have few calls producing big arrays
            array, sig, eps = self.compute(gs, num_events, mean_stds[:, g])
            M, N, E = array.shape
            for n in range(N):
                for e in range(E):
                    if (array[:, n, e] < min_iml).all():
                        array[:, n, e] = 0
            array = array.transpose(1, 0, 2)  # from M, N, E to N, M, E
            n = 0
            for rlz in rlzs:
                eids = eids_by_rlz[rlz]
                for ei, eid in enumerate(eids):
                    gmfa = array[:, :, n + ei]  # shape (N, M)
                    if sig_eps is not None:
                        tup = tuple([eid, rlz] + list(sig[:, n + ei]) +
                                    list(eps[:, n + ei]))
                        sig_eps.append(tup)
                    items = []
                    for sp in self.sec_perils:
                        o = sp.compute(mag, zip(self.imts, gmfa.T), self.ctx)
                        for outkey, outarr in zip(sp.outputs, o):
                            items.append((outkey, outarr))
                    for i, gmv in enumerate(gmfa):
                        if gmv.sum() == 0:
                            continue
                        data['sid'].append(sids[i])
                        data['eid'].append(eid)
                        data['rlz'].append(rlz)  # used in compute_gmfs_curves
                        for m in range(M):
                            data[f'gmv_{m}'].append(gmv[m])
                        for outkey, outarr in items:
                            data[outkey].append(outarr[i])
                        # gmv can be zero due to the minimum_intensity, coming
                        # from the job.ini or from the vulnerability functions
                n += len(eids)
        return data, time.time() - t0

    def compute(self, gsim, num_events, mean_stds):
        """
        :param gsim: GSIM used to compute mean_stds
        :param num_events: the number of seismic events
        :param mean_stds: array of shape O, M, N
        :returns:
            a 32 bit array of shape (num_imts, num_sites, num_events) and
            two arrays with shape (num_imts, num_events): sig for stddev_inter
            and eps for the random part
        """
        result = numpy.zeros(
            (len(self.imts), len(self.ctx.sids), num_events), F32)
        sig = numpy.zeros((len(self.imts), num_events), F32)
        eps = numpy.zeros((len(self.imts), num_events), F32)
        numpy.random.seed(self.seed)
        for imti, imt in enumerate(self.imts):
            try:
                result[imti], sig[imti], eps[imti] = self._compute(
                     mean_stds[:, imti], num_events, imt, gsim)
            except Exception as exc:
                raise RuntimeError(
                    '(%s, %s, source_id=%r) %s: %s' %
                    (gsim, imt, decode(self.source_id),
                     exc.__class__.__name__, exc)
                ).with_traceback(exc.__traceback__)
        if self.amplifier:
            self.amplifier.amplify_gmfs(
                self.ctx.ampcode, result, self.imts, self.seed)
        return result, sig, eps

    def _compute(self, mean_stds, num_events, imt, gsim):
        """
        :param mean_stds: array of shape (4, N)
        :param num_events: the number of seismic events
        :param imt: an IMT instance
        :param gsim: a GSIM instance
        :returns: (gmf(num_sites, num_events), stddev_inter(num_events),
                   epsilons(num_events))
        """
        num_sids = len(self.ctx.sids)
        if self.distribution is None:
            # for truncation_level = 0 there is only mean, no stds
            if self.correlation_model:
                raise ValueError('truncation_level=0 requires '
                                 'no correlation model')
            mean = mean_stds[0]
            gmf = to_imt_unit_values(mean, imt)
            gmf.shape += (1, )
            gmf = gmf.repeat(num_events, axis=1)
            return (gmf,
                    numpy.zeros(num_events, F32),
                    numpy.zeros(num_events, F32))
        elif gsim.DEFINED_FOR_STANDARD_DEVIATION_TYPES == {StdDev.TOTAL}:
            # If the GSIM provides only total standard deviation, we need
            # to compute mean and total standard deviation at the sites
            # of interest.
            # In this case, we also assume no correlation model is used.
            if self.correlation_model:
                raise CorrelationButNoInterIntraStdDevs(
                    self.correlation_model, gsim)

            mean, stddev_total = mean_stds[:2]
            stddev_total = stddev_total.reshape(stddev_total.shape + (1, ))
            mean = mean.reshape(mean.shape + (1, ))

            total_residual = stddev_total * rvs(
                self.distribution, num_sids, num_events)
            gmf = to_imt_unit_values(mean + total_residual, imt)
            stdi = numpy.nan
            epsilons = numpy.empty(num_events, F32)
            epsilons.fill(numpy.nan)
        else:
            mean, stddev_total, stddev_inter, stddev_intra = mean_stds
            stddev_intra = stddev_intra.reshape(stddev_intra.shape + (1, ))
            stddev_inter = stddev_inter.reshape(stddev_inter.shape + (1, ))
            mean = mean.reshape(mean.shape + (1, ))
            intra_residual = stddev_intra * rvs(
                self.distribution, num_sids, num_events)

            if self.correlation_model is not None:
                intra_residual = self.correlation_model.apply_correlation(
                    self.sites, imt, intra_residual, stddev_intra)
                sh = intra_residual.shape
                if len(sh) == 1:  # a vector
                    intra_residual = intra_residual.reshape(sh + (1,))

            epsilons = rvs(self.distribution, num_events)
            inter_residual = stddev_inter * epsilons

            gmf = to_imt_unit_values(
                mean + intra_residual + inter_residual, imt)
            stdi = stddev_inter.max(axis=0)
        return gmf, stdi, epsilons

# this is not used in the engine; it is still useful for usage in IPython
# when demonstrating hazardlib capabilities
def ground_motion_fields(rupture, sites, imts, gsim, truncation_level,
                         realizations, correlation_model=None, seed=None):
    """
    Given an earthquake rupture, the ground motion field calculator computes
    ground shaking over a set of sites, by randomly sampling a ground shaking
    intensity model. A ground motion field represents a possible 'realization'
    of the ground shaking due to an earthquake rupture.

    .. note::

     This calculator is using random numbers. In order to reproduce the
     same results numpy random numbers generator needs to be seeded, see
     http://docs.scipy.org/doc/numpy/reference/generated/numpy.random.seed.html

    :param openquake.hazardlib.source.rupture.Rupture rupture:
        Rupture to calculate ground motion fields radiated from.
    :param openquake.hazardlib.site.SiteCollection sites:
        Sites of interest to calculate GMFs.
    :param imts:
        List of intensity measure type objects (see
        :mod:`openquake.hazardlib.imt`).
    :param gsim:
        Ground-shaking intensity model, instance of subclass of either
        :class:`~openquake.hazardlib.gsim.base.GMPE` or
        :class:`~openquake.hazardlib.gsim.base.IPE`.
    :param truncation_level:
        Float, number of standard deviations for truncation of the intensity
        distribution, or ``None``.
    :param realizations:
        Integer number of GMF realizations to compute.
    :param correlation_model:
        Instance of correlation model object. See
        :mod:`openquake.hazardlib.correlation`. Can be ``None``, in which case
        non-correlated ground motion fields are calculated. Correlation model
        is not used if ``truncation_level`` is zero.
    :param int seed:
        The seed used in the numpy random number generator
    :returns:
        Dictionary mapping intensity measure type objects (same
        as in parameter ``imts``) to 2d numpy arrays of floats,
        representing different realizations of ground shaking intensity
        for all sites in the collection. First dimension represents
        sites and second one is for realizations.
    """
    cmaker = ContextMaker(rupture.tectonic_region_type, [gsim],
                          dict(truncation_level=truncation_level,
                               imtls={str(imt): [1] for imt in imts}))
    rupture.rup_id = seed
    gc = GmfComputer(rupture, sites, cmaker, correlation_model)
    mean_stds = cmaker.get_mean_stds([gc.ctx])[:, 0]
    res, _sig, _eps = gc.compute(gsim, realizations, mean_stds)
    return {imt: res[imti] for imti, imt in enumerate(gc.imts)}
<commit_msg>IMT-correlation for the epsilons in calc.gmf<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2012-2021 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.

"""
Module :mod:`~openquake.hazardlib.calc.gmf` exports
:func:`ground_motion_fields`.
"""
import time
import numpy
import scipy.stats

from openquake.baselib.general import AccumDict
from openquake.baselib.python3compat import decode
from openquake.hazardlib.const import StdDev
from openquake.hazardlib.gsim.base import ContextMaker, FarAwayRupture
from openquake.hazardlib.imt import from_string

U32 = numpy.uint32
F32 = numpy.float32

class CorrelationButNoInterIntraStdDevs(Exception):
    def __init__(self, corr, gsim):
        self.corr = corr
        self.gsim = gsim

    def __str__(self):
        return '''\
You cannot use the correlation model %s with the GSIM %s, \
that defines only the total standard deviation. If you want to use a \
correlation model you have to select a GMPE that provides the inter and \
intra event standard deviations.''' % (
            self.corr.__class__.__name__, self.gsim.__class__.__name__)

def rvs(distribution, *size):
    array = distribution.rvs(size)
    return array

def to_imt_unit_values(vals, imt):
    """
    Exponentiate the values unless the IMT is MMI
    """
    if str(imt) == 'MMI':
        return vals
    return numpy.exp(vals)

class GmfComputer(object):
    """
    Given an earthquake rupture, the ground motion field computer computes
    ground shaking over a set of sites, by randomly sampling a ground
    shaking intensity model.

    :param rupture:
        Rupture to calculate ground motion fields radiated from.

    :param :class:`openquake.hazardlib.site.SiteCollection` sitecol:
        a complete SiteCollection

    :param imts:
        a sorted list of Intensity Measure Type strings

    :param cmaker:
        a :class:`openquake.hazardlib.gsim.base.ContextMaker` instance

    :param truncation_level:
        Float, number of standard deviations for truncation of the intensity
        distribution, or ``None``.

    :param correlation_model:
        Instance of correlation model object. See
        :mod:`openquake.hazardlib.correlation`. Can be ``None``, in which
        case non-correlated ground motion fields are calculated.
        Correlation model is not used if ``truncation_level`` is zero.

    :param amplifier:
        None or an instance of Amplifier
    """
    # The GmfComputer is called from the OpenQuake Engine. In that case
    # the rupture is an higher level containing a
    # :class:`openquake.hazardlib.source.rupture.Rupture` instance as an
    # attribute. Then the `.compute(gsim, num_events, ms)` method is called and
    # a matrix of size (I, N, E) is returned, where I is the number of
    # IMTs, N the number of affected sites and E the number of events. The
    # seed is extracted from the underlying rupture.
    def __init__(self, rupture, sitecol, cmaker, correlation_model=None,
                 amplifier=None, sec_perils=()):
        if len(sitecol) == 0:
            raise ValueError('No sites')
        elif len(cmaker.imtls) == 0:
            raise ValueError('No IMTs')
        elif len(cmaker.gsims) == 0:
            raise ValueError('No GSIMs')
        self.cmaker = cmaker
        self.imts = [from_string(imt) for imt in cmaker.imtls]
        self.cmaker = cmaker
        self.gsims = sorted(cmaker.gsims)
        self.correlation_model = correlation_model
        self.amplifier = amplifier
        self.sec_perils = sec_perils
        # `rupture` is an EBRupture instance in the engine
        if hasattr(rupture, 'source_id'):
            self.ebrupture = rupture
            self.source_id = rupture.source_id  # the underlying source
            rupture = rupture.rupture  # the underlying rupture
        else:  # in the hazardlib tests
            self.source_id = '?'
        self.seed = rupture.rup_id
        ctxs = cmaker.get_ctxs([rupture], sitecol, self.source_id)
        if not ctxs:
            raise FarAwayRupture
        self.ctx = ctxs[0]
        if correlation_model:  # store the filtered sitecol
            self.sites = sitecol.complete.filtered(self.ctx.sids)
        if cmaker.trunclevel is None:
            self.distribution = scipy.stats.norm()
        elif cmaker.trunclevel == 0:
            self.distribution = None
        else:
            assert cmaker.trunclevel > 0, cmaker.trunclevel
            self.distribution = scipy.stats.truncnorm(
                - cmaker.trunclevel, cmaker.trunclevel)

    def compute_all(self, sig_eps=None):
        """
        :returns: (dict with fields eid, sid, gmv_X, ...), dt
        """
        min_iml = self.cmaker.min_iml
        rlzs_by_gsim = self.cmaker.gsims
        t0 = time.time()
        sids = self.ctx.sids
        eids_by_rlz = self.ebrupture.get_eids_by_rlz(rlzs_by_gsim)
        mag = self.ebrupture.rupture.mag
        data = AccumDict(accum=[])
        mean_stds = self.cmaker.get_mean_stds([self.ctx])  # (4, G, M, N)
        for g, (gs, rlzs) in enumerate(rlzs_by_gsim.items()):
            num_events = sum(len(eids_by_rlz[rlz]) for rlz in rlzs)
            if num_events == 0:  # it may happen
                continue
            # NB: the trick for performance is to keep the call to
            # .compute outside of the loop over the realizations;
            # it is better to have few calls producing big arrays
            array, sig, eps = self.compute(gs, num_events, mean_stds[:, g])
            M, N, E = array.shape
            for n in range(N):
                for e in range(E):
                    if (array[:, n, e] < min_iml).all():
                        array[:, n, e] = 0
            array = array.transpose(1, 0, 2)  # from M, N, E to N, M, E
            n = 0
            for rlz in rlzs:
                eids = eids_by_rlz[rlz]
                for ei, eid in enumerate(eids):
                    gmfa = array[:, :, n + ei]  # shape (N, M)
                    if sig_eps is not None:
                        tup = tuple([eid, rlz] + list(sig[:, n + ei]) +
                                    list(eps[:, n + ei]))
                        sig_eps.append(tup)
                    items = []
                    for sp in self.sec_perils:
                        o = sp.compute(mag, zip(self.imts, gmfa.T), self.ctx)
                        for outkey, outarr in zip(sp.outputs, o):
                            items.append((outkey, outarr))
                    for i, gmv in enumerate(gmfa):
                        if gmv.sum() == 0:
                            continue
                        data['sid'].append(sids[i])
                        data['eid'].append(eid)
                        data['rlz'].append(rlz)  # used in compute_gmfs_curves
                        for m in range(M):
                            data[f'gmv_{m}'].append(gmv[m])
                        for outkey, outarr in items:
                            data[outkey].append(outarr[i])
                        # gmv can be zero due to the minimum_intensity, coming
                        # from the job.ini or from the vulnerability functions
                n += len(eids)
        return data, time.time() - t0

    def compute(self, gsim, num_events, mean_stds):
        """
        :param gsim: GSIM used to compute mean_stds
        :param num_events: the number of seismic events
        :param mean_stds: array of shape O, M, N
        :returns:
            a 32 bit array of shape (num_imts, num_sites, num_events) and
            two arrays with shape (num_imts, num_events): sig for tau
            and eps for the random part
        """
        result = numpy.zeros(
            (len(self.imts), len(self.ctx.sids), num_events), F32)
        sig = numpy.zeros((len(self.imts), num_events), F32)
        eps = numpy.zeros((len(self.imts), num_events), F32)
        numpy.random.seed(self.seed)
        for imti, imt in enumerate(self.imts):
            try:
                result[imti], sig[imti], eps[imti] = self._compute(
                     mean_stds[:, imti], num_events, imt, gsim)
            except Exception as exc:
                raise RuntimeError(
                    '(%s, %s, source_id=%r) %s: %s' %
                    (gsim, imt, decode(self.source_id),
                     exc.__class__.__name__, exc)
                ).with_traceback(exc.__traceback__)
        if self.amplifier:
            self.amplifier.amplify_gmfs(
                self.ctx.ampcode, result, self.imts, self.seed)
        return result, sig, eps

    def _compute(self, mean_stds, num_events, imt, gsim):
        """
        :param mean_stds: array of shape (4, N)
        :param num_events: the number of seismic events
        :param imt: an IMT instance
        :param gsim: a GSIM instance
        :returns: (gmf(num_sites, num_events), tau(num_events),
                   epsilons(num_events))
        """
        num_sids = len(self.ctx.sids)
        if self.distribution is None:
            # for truncation_level = 0 there is only mean, no stds
            if self.correlation_model:
                raise ValueError('truncation_level=0 requires '
                                 'no correlation model')
            mean = mean_stds[0]
            gmf = to_imt_unit_values(mean, imt)
            gmf.shape += (1, )
            gmf = gmf.repeat(num_events, axis=1)
            stdi = numpy.zeros(num_events, F32)
            epsilons = numpy.zeros(num_events, F32)
        elif gsim.DEFINED_FOR_STANDARD_DEVIATION_TYPES == {StdDev.TOTAL}:
            # If the GSIM provides only total standard deviation, we need
            # to compute mean and total standard deviation at the sites
            # of interest.
            # In this case, we also assume no correlation model is used.
            if self.correlation_model:
                raise CorrelationButNoInterIntraStdDevs(
                    self.correlation_model, gsim)

            mean, sig = mean_stds[:2]
            sig = sig.reshape(sig.shape + (1, ))
            mean = mean.reshape(mean.shape + (1, ))

            total_residual = sig * rvs(
                self.distribution, num_sids, num_events)
            gmf = to_imt_unit_values(mean + total_residual, imt)
            stdi = numpy.nan
            epsilons = numpy.empty(num_events, F32)
            epsilons.fill(numpy.nan)
        else:
            mean, sig, tau, phi = mean_stds
            phi = phi.reshape(phi.shape + (1, ))
            tau = tau.reshape(tau.shape + (1, ))
            mean = mean.reshape(mean.shape + (1, ))
            intra_residual = phi * rvs(
                self.distribution, num_sids, num_events)

            if self.correlation_model is not None:
                intra_residual = self.correlation_model.apply_correlation(
                    self.sites, imt, intra_residual, phi)
                sh = intra_residual.shape
                if len(sh) == 1:  # a vector
                    intra_residual = intra_residual.reshape(sh + (1,))

            epsilons = rvs(self.distribution, num_events)
            inter_residual = tau * epsilons

            gmf = to_imt_unit_values(
                mean + intra_residual + inter_residual, imt)
            stdi = tau.max(axis=0)
        return gmf, stdi, epsilons

# this is not used in the engine; it is still useful for usage in IPython
# when demonstrating hazardlib capabilities
def ground_motion_fields(rupture, sites, imts, gsim, truncation_level,
                         realizations, correlation_model=None, seed=None):
    """
    Given an earthquake rupture, the ground motion field calculator computes
    ground shaking over a set of sites, by randomly sampling a ground shaking
    intensity model. A ground motion field represents a possible 'realization'
    of the ground shaking due to an earthquake rupture.

    .. note::

     This calculator is using random numbers. In order to reproduce the
     same results numpy random numbers generator needs to be seeded, see
     http://docs.scipy.org/doc/numpy/reference/generated/numpy.random.seed.html

    :param openquake.hazardlib.source.rupture.Rupture rupture:
        Rupture to calculate ground motion fields radiated from.
    :param openquake.hazardlib.site.SiteCollection sites:
        Sites of interest to calculate GMFs.
    :param imts:
        List of intensity measure type objects (see
        :mod:`openquake.hazardlib.imt`).
    :param gsim:
        Ground-shaking intensity model, instance of subclass of either
        :class:`~openquake.hazardlib.gsim.base.GMPE` or
        :class:`~openquake.hazardlib.gsim.base.IPE`.
    :param truncation_level:
        Float, number of standard deviations for truncation of the intensity
        distribution, or ``None``.
    :param realizations:
        Integer number of GMF realizations to compute.
    :param correlation_model:
        Instance of correlation model object. See
        :mod:`openquake.hazardlib.correlation`. Can be ``None``, in which case
        non-correlated ground motion fields are calculated. Correlation model
        is not used if ``truncation_level`` is zero.
    :param int seed:
        The seed used in the numpy random number generator
    :returns:
        Dictionary mapping intensity measure type objects (same
        as in parameter ``imts``) to 2d numpy arrays of floats,
        representing different realizations of ground shaking intensity
        for all sites in the collection. First dimension represents
        sites and second one is for realizations.
    """
    cmaker = ContextMaker(rupture.tectonic_region_type, [gsim],
                          dict(truncation_level=truncation_level,
                               imtls={str(imt): [1] for imt in imts}))
    rupture.rup_id = seed
    gc = GmfComputer(rupture, sites, cmaker, correlation_model)
    mean_stds = cmaker.get_mean_stds([gc.ctx])[:, 0]
    res, _sig, _eps = gc.compute(gsim, realizations, mean_stds)
    return {imt: res[imti] for imti, imt in enumerate(gc.imts)}
<|endoftext|>


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

"""
Here there are some real functional tests starting an engine server and
running computations.
"""
from __future__ import print_function
import io
import os
import re
import sys
import json
import time
import unittest
import numpy
from django.test import Client
from openquake.baselib.general import writetmp
from openquake.engine.export import core
from openquake.server.db import actions
from openquake.server.dbserver import db, get_status
from openquake.commands.abort import abort

class EngineServerTestCase(unittest.TestCase):
    datadir = os.path.join(os.path.dirname(__file__), 'data')

    # general utilities

    @classmethod
    def post(cls, path, data=None):
        return cls.c.post('/v1/calc/%s' % path, data)

    @classmethod
    def post_nrml(cls, data):
        return cls.c.post('/v1/valid/', dict(xml_text=data))

    @classmethod
    def get(cls, path, **data):
        resp = cls.c.get('/v1/calc/%s' % path, data,
                         HTTP_HOST='127.0.0.1')
        assert resp.content
        try:
            return json.loads(resp.content.decode('utf8'))
        except:
            print('Invalid JSON, see %s' % writetmp(resp.content),
                  file=sys.stderr)
            return {}

    @classmethod
    def get_text(cls, path, **data):
        resp = cls.c.get('/v1/calc/%s' % path, data)
        if resp.status_code == 500:
            raise Exception(resp.content.decode('utf8'))
        return b''.join(resp.streaming_content)

    @classmethod
    def wait(cls):
        # wait until all calculations stop
        while True:
            running_calcs = cls.get('list', is_running='true')
            if not running_calcs:
                break
            time.sleep(0.5)

    def postzip(self, archive):
        with open(os.path.join(self.datadir, archive), 'rb') as a:
            resp = self.post('run', dict(archive=a))
        try:
            js = json.loads(resp.content.decode('utf8'))
        except:
            raise ValueError(b'Invalid JSON response: %r' % resp.content)
        if resp.status_code == 200:  # ok case
            job_id = js['job_id']
            self.job_ids.append(job_id)
            time.sleep(1)  # wait a bit for the calc to start
            return job_id
        else:  # error case
            return ''.join(js)  # traceback string

    # start/stop server utilities

    @classmethod
    def setUpClass(cls):
        assert get_status() == 'running'
        cls.job_ids = []
        env = os.environ.copy()
        env['OQ_DISTRIBUTE'] = 'no'
        # let's impersonate the user openquake, the one running the WebUI:
        # we need to set LOGNAME on Linux and USERNAME on Windows
        env['LOGNAME'] = env['USERNAME'] = 'openquake'
        cls.c = Client()

    @classmethod
    def tearDownClass(cls):
        cls.wait()

    def setUp(self):
        if sys.version_info[0] == 2:
            # python 2 will die
            raise unittest.SkipTest('Python 2')

    # tests

    def test_404(self):
        # looking for a missing calc_id
        resp = self.c.get('/v1/calc/0')
        assert resp.status_code == 404, resp

    def test_ok(self):
        job_id = self.postzip('archive_ok.zip')
        self.wait()
        log = self.get('%s/log/:' % job_id)
        self.assertGreater(len(log), 0)
        results = self.get('%s/results' % job_id)
        self.assertGreater(len(results), 0)
        for res in results:
            for etype in res['outtypes']:  # test all export types
                text = self.get_text(
                    'result/%s' % res['id'], export_type=etype)
                print('downloading result/%s' % res['id'], res['type'], etype)
                self.assertGreater(len(text), 0)

        # test no filtering in actions.get_calcs
        all_jobs = self.get('list')
        self.assertGreater(len(all_jobs), 0)

        extract_url = '/v1/calc/%s/extract/' % job_id

        # check extract/composite_risk_model.attrs
        url = extract_url + 'composite_risk_model.attrs'
        self.assertEqual(self.c.get(url).status_code, 200)

        # check asset_values
        resp = self.c.get(extract_url + 'asset_values/0')
        data = b''.join(ln for ln in resp.streaming_content)
        got = numpy.load(io.BytesIO(data))  # load npz file
        self.assertEqual(len(got['array']), 0)  # there are 0 assets on site 0
        self.assertEqual(resp.status_code, 200)

        # check avg_losses-rlzs
        resp = self.c.get(
            extract_url + 'agglosses/structural?taxonomy=W-SLFB-1')
        data = b''.join(ln for ln in resp.streaming_content)
        got = numpy.load(io.BytesIO(data))  # load npz file
        self.assertEqual(len(got['array']), 1)  # expected 1 aggregate value
        self.assertEqual(resp.status_code, 200)

        # TODO: check aggcurves

        # there is some logic in `core.export_from_db` that it is only
        # exercised when the export fails
        datadir, dskeys = actions.get_results(db, job_id)
        # try to export a non-existing output
        with self.assertRaises(core.DataStoreExportError) as ctx:
            core.export_from_db(('XXX', 'csv'), job_id, datadir, '/tmp')
        self.assertIn('Could not export XXX in csv', str(ctx.exception))

    def test_classical(self):
        job_id = self.postzip('classical.zip')
        self.wait()

        # check that we get the expected outputs
        results = self.get('%s/results' % job_id)
        self.assertEqual(['fullreport', 'hcurves', 'hmaps', 'realizations',
                          'sourcegroups'], [r['name'] for r in results])

        # check the filename of the hmaps
        hmaps_id = results[2]['id']
        resp = self.c.head('/v1/calc/result/%s?export_type=csv' % hmaps_id)
        # remove output ID digits from the filename
        cd = re.sub(r'\d', '', resp._headers['content-disposition'][1])
        self.assertEqual(
            cd, 'attachment; filename=output--hazard_map-mean_.csv')

        # check oqparam
        resp = self.get('%s/oqparam' % job_id)  # dictionary of parameters
        self.assertEqual(resp['calculation_mode'], 'classical')

        # check the /extract endpoint
        url = '/v1/calc/%s/extract/hazard/rlzs' % job_id
        resp = self.c.get(url)
        self.assertEqual(resp.status_code, 200)

    def test_abort(self):
        job_id = self.postzip('archive_ok.zip')
        time.sleep(1)  # give time
        abort(job_id)
        self.wait()

    def test_err_1(self):
        # the rupture XML file has a syntax error
        job_id = self.postzip('archive_err_1.zip')
        self.wait()

        # download the datastore, even if incomplete
        resp = self.c.get('/v1/calc/%s/datastore' % job_id)
        self.assertEqual(resp.status_code, 200)

        tb = self.get('%s/traceback' % job_id)
        if not tb:
            sys.stderr.write('Empty traceback, please check!\n')

        self.post('%s/remove' % job_id)
        # make sure job_id is no more in the list of relevant jobs
        job_ids = [job['id'] for job in self.get('list', relevant=True)]
        self.assertFalse(job_id in job_ids)
        # NB: the job is invisible but still there

    def test_err_2(self):
        # the file logic-tree-source-model.xml is missing
        tb_str = self.postzip('archive_err_2.zip')
        self.assertIn('No such file', tb_str)

    def test_err_3(self):
        # there is no file job.ini, job_hazard.ini or job_risk.ini
        tb_str = self.postzip('archive_err_3.zip')
        self.assertIn('Could not find any file of the form', tb_str)

    def test_available_gsims(self):
        resp = self.c.get('/v1/available_gsims')
        self.assertIn(b'ChiouYoungs2014PEER', resp.content)

    # tests for nrml validation

    def test_validate_nrml_valid(self):
        valid_file = os.path.join(self.datadir, 'vulnerability_model.xml')
        with open(valid_file) as vf:
            valid_content = vf.read()
        resp = self.post_nrml(valid_content)
        self.assertEqual(resp.status_code, 200)
        resp_text_dict = json.loads(resp.content.decode('utf8'))
        self.assertTrue(resp_text_dict['valid'])
        self.assertIsNone(resp_text_dict['error_msg'])
        self.assertIsNone(resp_text_dict['error_line'])

    def test_validate_nrml_invalid(self):
        invalid_file = os.path.join(self.datadir,
                                    'vulnerability_model_invalid.xml')
        with open(invalid_file) as vf:
            invalid_content = vf.read()
        resp = self.post_nrml(invalid_content)
        self.assertEqual(resp.status_code, 200)
        resp_text_dict = json.loads(resp.content.decode('utf8'))
        self.assertFalse(resp_text_dict['valid'])
        self.assertIn(u'Could not convert lossRatio->positivefloats:'
                      ' float -0.018800826 < 0',
                      resp_text_dict['error_msg'])
        self.assertEqual(resp_text_dict['error_line'], 7)

    def test_validate_nrml_unclosed_tag(self):
        invalid_file = os.path.join(self.datadir,
                                    'vulnerability_model_unclosed_tag.xml')
        with open(invalid_file) as vf:
            invalid_content = vf.read()
        resp = self.post_nrml(invalid_content)
        self.assertEqual(resp.status_code, 200)
        resp_text_dict = json.loads(resp.content.decode('utf8'))
        self.assertFalse(resp_text_dict['valid'])
        self.assertIn(u'mismatched tag', resp_text_dict['error_msg'])
        self.assertEqual(resp_text_dict['error_line'], 9)

    def test_validate_nrml_missing_parameter(self):
        # passing a wrong parameter, instead of the required 'xml_text'
        resp = self.c.post('/v1/valid/', foo='bar')
        self.assertEqual(resp.status_code, 400)
        self.assertEqual(resp.content,
                         b'Please provide the "xml_text" parameter')
<commit_msg>Skipped test_abort [skip hazardlib] [demos]<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2015-2017 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.

"""
Here there are some real functional tests starting an engine server and
running computations.
"""
from __future__ import print_function
import io
import os
import re
import sys
import json
import time
import unittest
import numpy
from django.test import Client
from openquake.baselib.general import writetmp
from openquake.engine.export import core
from openquake.server.db import actions
from openquake.server.dbserver import db, get_status
from openquake.commands.abort import abort

class EngineServerTestCase(unittest.TestCase):
    datadir = os.path.join(os.path.dirname(__file__), 'data')

    # general utilities

    @classmethod
    def post(cls, path, data=None):
        return cls.c.post('/v1/calc/%s' % path, data)

    @classmethod
    def post_nrml(cls, data):
        return cls.c.post('/v1/valid/', dict(xml_text=data))

    @classmethod
    def get(cls, path, **data):
        resp = cls.c.get('/v1/calc/%s' % path, data,
                         HTTP_HOST='127.0.0.1')
        assert resp.content
        try:
            return json.loads(resp.content.decode('utf8'))
        except:
            print('Invalid JSON, see %s' % writetmp(resp.content),
                  file=sys.stderr)
            return {}

    @classmethod
    def get_text(cls, path, **data):
        resp = cls.c.get('/v1/calc/%s' % path, data)
        if resp.status_code == 500:
            raise Exception(resp.content.decode('utf8'))
        return b''.join(resp.streaming_content)

    @classmethod
    def wait(cls):
        # wait until all calculations stop
        while True:
            running_calcs = cls.get('list', is_running='true')
            if not running_calcs:
                break
            time.sleep(0.5)

    def postzip(self, archive):
        with open(os.path.join(self.datadir, archive), 'rb') as a:
            resp = self.post('run', dict(archive=a))
        try:
            js = json.loads(resp.content.decode('utf8'))
        except:
            raise ValueError(b'Invalid JSON response: %r' % resp.content)
        if resp.status_code == 200:  # ok case
            job_id = js['job_id']
            self.job_ids.append(job_id)
            time.sleep(1)  # wait a bit for the calc to start
            return job_id
        else:  # error case
            return ''.join(js)  # traceback string

    # start/stop server utilities

    @classmethod
    def setUpClass(cls):
        assert get_status() == 'running'
        cls.job_ids = []
        env = os.environ.copy()
        env['OQ_DISTRIBUTE'] = 'no'
        # let's impersonate the user openquake, the one running the WebUI:
        # we need to set LOGNAME on Linux and USERNAME on Windows
        env['LOGNAME'] = env['USERNAME'] = 'openquake'
        cls.c = Client()

    @classmethod
    def tearDownClass(cls):
        cls.wait()

    def setUp(self):
        if sys.version_info[0] == 2:
            # python 2 will die
            raise unittest.SkipTest('Python 2')

    # tests

    def test_404(self):
        # looking for a missing calc_id
        resp = self.c.get('/v1/calc/0')
        assert resp.status_code == 404, resp

    def test_ok(self):
        job_id = self.postzip('archive_ok.zip')
        self.wait()
        log = self.get('%s/log/:' % job_id)
        self.assertGreater(len(log), 0)
        results = self.get('%s/results' % job_id)
        self.assertGreater(len(results), 0)
        for res in results:
            for etype in res['outtypes']:  # test all export types
                text = self.get_text(
                    'result/%s' % res['id'], export_type=etype)
                print('downloading result/%s' % res['id'], res['type'], etype)
                self.assertGreater(len(text), 0)

        # test no filtering in actions.get_calcs
        all_jobs = self.get('list')
        self.assertGreater(len(all_jobs), 0)

        extract_url = '/v1/calc/%s/extract/' % job_id

        # check extract/composite_risk_model.attrs
        url = extract_url + 'composite_risk_model.attrs'
        self.assertEqual(self.c.get(url).status_code, 200)

        # check asset_values
        resp = self.c.get(extract_url + 'asset_values/0')
        data = b''.join(ln for ln in resp.streaming_content)
        got = numpy.load(io.BytesIO(data))  # load npz file
        self.assertEqual(len(got['array']), 0)  # there are 0 assets on site 0
        self.assertEqual(resp.status_code, 200)

        # check avg_losses-rlzs
        resp = self.c.get(
            extract_url + 'agglosses/structural?taxonomy=W-SLFB-1')
        data = b''.join(ln for ln in resp.streaming_content)
        got = numpy.load(io.BytesIO(data))  # load npz file
        self.assertEqual(len(got['array']), 1)  # expected 1 aggregate value
        self.assertEqual(resp.status_code, 200)

        # TODO: check aggcurves

        # there is some logic in `core.export_from_db` that it is only
        # exercised when the export fails
        datadir, dskeys = actions.get_results(db, job_id)
        # try to export a non-existing output
        with self.assertRaises(core.DataStoreExportError) as ctx:
            core.export_from_db(('XXX', 'csv'), job_id, datadir, '/tmp')
        self.assertIn('Could not export XXX in csv', str(ctx.exception))

    def test_classical(self):
        job_id = self.postzip('classical.zip')
        self.wait()

        # check that we get the expected outputs
        results = self.get('%s/results' % job_id)
        self.assertEqual(['fullreport', 'hcurves', 'hmaps', 'realizations',
                          'sourcegroups'], [r['name'] for r in results])

        # check the filename of the hmaps
        hmaps_id = results[2]['id']
        resp = self.c.head('/v1/calc/result/%s?export_type=csv' % hmaps_id)
        # remove output ID digits from the filename
        cd = re.sub(r'\d', '', resp._headers['content-disposition'][1])
        self.assertEqual(
            cd, 'attachment; filename=output--hazard_map-mean_.csv')

        # check oqparam
        resp = self.get('%s/oqparam' % job_id)  # dictionary of parameters
        self.assertEqual(resp['calculation_mode'], 'classical')

        # check the /extract endpoint
        url = '/v1/calc/%s/extract/hazard/rlzs' % job_id
        resp = self.c.get(url)
        self.assertEqual(resp.status_code, 200)

    def test_abort(self):
        raise unittest.SkipTest('abort does not work with zmq')
        job_id = self.postzip('archive_ok.zip')
        time.sleep(1)  # give time
        abort(job_id)
        self.wait()

    def test_err_1(self):
        # the rupture XML file has a syntax error
        job_id = self.postzip('archive_err_1.zip')
        self.wait()

        # download the datastore, even if incomplete
        resp = self.c.get('/v1/calc/%s/datastore' % job_id)
        self.assertEqual(resp.status_code, 200)

        tb = self.get('%s/traceback' % job_id)
        if not tb:
            sys.stderr.write('Empty traceback, please check!\n')

        self.post('%s/remove' % job_id)
        # make sure job_id is no more in the list of relevant jobs
        job_ids = [job['id'] for job in self.get('list', relevant=True)]
        self.assertFalse(job_id in job_ids)
        # NB: the job is invisible but still there

    def test_err_2(self):
        # the file logic-tree-source-model.xml is missing
        tb_str = self.postzip('archive_err_2.zip')
        self.assertIn('No such file', tb_str)

    def test_err_3(self):
        # there is no file job.ini, job_hazard.ini or job_risk.ini
        tb_str = self.postzip('archive_err_3.zip')
        self.assertIn('Could not find any file of the form', tb_str)

    def test_available_gsims(self):
        resp = self.c.get('/v1/available_gsims')
        self.assertIn(b'ChiouYoungs2014PEER', resp.content)

    # tests for nrml validation

    def test_validate_nrml_valid(self):
        valid_file = os.path.join(self.datadir, 'vulnerability_model.xml')
        with open(valid_file) as vf:
            valid_content = vf.read()
        resp = self.post_nrml(valid_content)
        self.assertEqual(resp.status_code, 200)
        resp_text_dict = json.loads(resp.content.decode('utf8'))
        self.assertTrue(resp_text_dict['valid'])
        self.assertIsNone(resp_text_dict['error_msg'])
        self.assertIsNone(resp_text_dict['error_line'])

    def test_validate_nrml_invalid(self):
        invalid_file = os.path.join(self.datadir,
                                    'vulnerability_model_invalid.xml')
        with open(invalid_file) as vf:
            invalid_content = vf.read()
        resp = self.post_nrml(invalid_content)
        self.assertEqual(resp.status_code, 200)
        resp_text_dict = json.loads(resp.content.decode('utf8'))
        self.assertFalse(resp_text_dict['valid'])
        self.assertIn(u'Could not convert lossRatio->positivefloats:'
                      ' float -0.018800826 < 0',
                      resp_text_dict['error_msg'])
        self.assertEqual(resp_text_dict['error_line'], 7)

    def test_validate_nrml_unclosed_tag(self):
        invalid_file = os.path.join(self.datadir,
                                    'vulnerability_model_unclosed_tag.xml')
        with open(invalid_file) as vf:
            invalid_content = vf.read()
        resp = self.post_nrml(invalid_content)
        self.assertEqual(resp.status_code, 200)
        resp_text_dict = json.loads(resp.content.decode('utf8'))
        self.assertFalse(resp_text_dict['valid'])
        self.assertIn(u'mismatched tag', resp_text_dict['error_msg'])
        self.assertEqual(resp_text_dict['error_line'], 9)

    def test_validate_nrml_missing_parameter(self):
        # passing a wrong parameter, instead of the required 'xml_text'
        resp = self.c.post('/v1/valid/', foo='bar')
        self.assertEqual(resp.status_code, 400)
        self.assertEqual(resp.content,
                         b'Please provide the "xml_text" parameter')
<|endoftext|>


<commit_before>import StringIO
import json
import logging
import urlparse

from django.http import HttpResponse

from openquake.engine.db import models as oqe_models
from openquake.engine import engine
from openquake.engine.export import hazard as hazard_export
from openquake.engine.export import risk as risk_export

METHOD_NOT_ALLOWED = 405
NOT_IMPLEMENTED = 501
JSON = 'application/json'

IGNORE_FIELDS = ('base_path', 'export_dir', 'owner')
GEOM_FIELDS = ('region', 'sites', 'region_constraint', 'sites_disagg')
RISK_INPUTS = ('hazard_calculation', 'hazard_output')

#: For exporting calculation outputs, the client can request a specific format
#: (xml, geojson, csv, etc.). If the client does not specify give them (NRML)
#: XML by default.
DEFAULT_EXPORT_TYPE = 'xml'

LOGGER = logging.getLogger('openquakeserver')

class allowed_methods(object):
    """
    Use as a view decorator to strictly enforce HTTP request types.

    It works like this:

    .. code-block:: python

        @allowed_methods(('GET', 'POST'))
        def some_view(request):
            # code goes here

        @ allowed_methods(('DELETE', ))
        def some_other_view(request):
            # code goes here
    """

    def __init__(self, methods):
        self.methods = methods

    def __call__(self, func):
        def wrapped(request, *args, **kwargs):
            if not request.method in self.methods:
                return HttpResponse(status=METHOD_NOT_ALLOWED)
            else:
                return func(request, *args, **kwargs)
        return wrapped

def _get_base_url(request):
    """
    Construct a base URL, given a request object.

    This comprises the protocol prefix (http:// or https://) and the host,
    which can include the port number. For example:
    http://www.openquake.org or https://www.openquake.org:8000.
    """
    if request.is_secure():
        base_url = 'https://%s'
    else:
        base_url = 'http://%s'
    base_url %= request.META['HTTP_HOST']
    return base_url

def _calc_to_response_data(calc):
    """
    Extract the calculation parameters into a dictionary.
    """
    fields = [x.name for x in calc._meta.fields if x.name not in IGNORE_FIELDS]
    response_data = {}
    for field_name in fields:
        try:
            value = getattr(calc, field_name)
            if value is not None:
                if field_name in GEOM_FIELDS:
                    response_data[field_name] = json.loads(value.geojson)
                elif field_name in RISK_INPUTS:
                    response_data[field_name] = value.id
                else:
                    response_data[field_name] = value
        except AttributeError:
            # Better that we miss an attribute than crash.
            pass
    return response_data

@allowed_methods(('GET', ))
def calc_hazard(request):
    """
    Get a list of risk calculations and report their id, status, description,
    and a url where more detailed information can be accessed.

    Responses are in JSON.
    """
    base_url = _get_base_url(request)

    haz_calc_data = _get_haz_calcs()

    response_data = []
    for hc_id, status, desc in haz_calc_data:
        url = urlparse.urljoin(base_url, 'v1/calc/hazard/%d' % hc_id)
        response_data.append(
            dict(id=hc_id, status=status, description=desc, url=url)
        )

    return HttpResponse(content=json.dumps(response_data),
                        content_type=JSON)

def _get_haz_calcs():
    """
    Helper function for get job+calculation data from the oq-engine database.

    Gets all hazard calculation records available.
    """
    return oqe_models.OqJob.objects\
        .select_related()\
        .filter(hazard_calculation__isnull=False)\
        .values_list('hazard_calculation',
                     'status',
                     'hazard_calculation__description')

@allowed_methods(('GET', ))
def calc_hazard_info(request, calc_id):
    """
    Get a JSON blob containing all of parameters for the given calculation
    (specified by ``calc_id``). Also includes the current job status (
    executing, complete, etc.).
    """
    response_data = _get_haz_calc_info(calc_id)

    return HttpResponse(content=json.dumps(response_data), content_type=JSON)

def _get_haz_calc_info(calc_id):
    """
    Helper function to get job info and hazard calculation params from the
    oq-engine DB, as a dictionary.
    """
    job = oqe_models.OqJob.objects\
        .select_related()\
        .get(hazard_calculation=calc_id)

    hc = job.hazard_calculation
    response_data = _calc_to_response_data(hc)

    response_data['status'] = job.status
    return response_data

@allowed_methods(('GET', ))
def calc_hazard_results(request, calc_id):
    """
    Get a summarized list of hazard calculation results for a given
    ``calc_id``. Result is a JSON array of objects containing the following
    attributes:

        * id
        * name
        * type (hazard_curve, hazard_map, etc.)
        * url (the exact url where the full result can be accessed)
    """
    base_url = _get_base_url(request)

    results = engine.get_hazard_outputs(calc_id)

    response_data = []
    for result in results:
        url = urlparse.urljoin(base_url,
                               'v1/calc/hazard/result/%d' % result.id)
        datum = dict(
            id=result.id,
            name=result.display_name,
            type=result.output_type,
            url=url,
        )
        response_data.append(datum)

    return HttpResponse(content=json.dumps(response_data))

@allowed_methods(('GET', ))
def get_hazard_result(request, result_id):
    """
    Download a specific hazard result, by ``result_id``.

    Parameters for the GET request can include an `export_type`, such as 'xml',
    'geojson', 'csv', etc.
    """
    return _get_result(request, result_id, hazard_export.export)

@allowed_methods(('GET', ))
def calc_risk(request):
    """
    Get a list of risk calculations and report their id, status, description,
    and a url where more detailed information can be accessed.

    Responses are in JSON.
    """
    base_url = _get_base_url(request)

    risk_calc_data = _get_risk_calcs()

    response_data = []
    for hc_id, status, desc in risk_calc_data:
        url = urlparse.urljoin(base_url, 'v1/calc/risk/%d' % hc_id)
        response_data.append(
            dict(id=hc_id, status=status, description=desc, url=url)
        )

    return HttpResponse(content=json.dumps(response_data),
                        content_type=JSON)

def _get_risk_calcs():
    """
    Helper function for get job+calculation data from the oq-engine database.

    Gets all risk calculation records available.
    """
    return oqe_models.OqJob.objects\
        .select_related()\
        .filter(risk_calculation__isnull=False)\
        .values_list('risk_calculation',
                     'status',
                     'risk_calculation__description')

@allowed_methods(('GET', ))
def calc_risk_info(request, calc_id):
    """
    Get a JSON blob containing all of parameters for the given calculation
    (specified by ``calc_id``). Also includes the current job status (
    executing, complete, etc.).
    """
    response_data = _get_risk_calc_info(calc_id)

    return HttpResponse(content=json.dumps(response_data), content_type=JSON)

def _get_risk_calc_info(calc_id):
    """
    Helper function to get job info and hazard calculation params from the
    oq-engine DB, as a dictionary.
    """
    job = oqe_models.OqJob.objects\
        .select_related()\
        .get(risk_calculation=calc_id)

    rc = job.risk_calculation
    response_data = _calc_to_response_data(rc)

    response_data['status'] = job.status
    return response_data

@allowed_methods(('GET', ))
def calc_risk_results(request, calc_id):
    """
    Get a summarized list of risk calculation results for a given
    ``calc_id``. Result is a JSON array of objects containing the following
    attributes:

        * id
        * name
        * type (hazard_curve, hazard_map, etc.)
        * url (the exact url where the full result can be accessed)
    """
    base_url = _get_base_url(request)

    results = engine.get_risk_outputs(calc_id)

    response_data = []
    for result in results:
        url = urlparse.urljoin(base_url,
                               'v1/calc/risk/result/%d' % result.id)
        datum = dict(
            id=result.id,
            name=result.display_name,
            type=result.output_type,
            url=url,
        )
        response_data.append(datum)

    return HttpResponse(content=json.dumps(response_data))

@allowed_methods(('GET', ))
def get_risk_result(request, result_id):
    """
    Download a specific hazard result, by ``result_id``.

    Parameters for the GET request can include an `export_type`, such as 'xml',
    'geojson', 'csv', etc.
    """
    return _get_result(request, result_id, risk_export.export)

def _get_result(request, result_id, export_fn):
    """
    The common abstracted functionality for getting hazard or risk results.
    The functionality is the same, except for the hazard/risk specific
    ``export_fn``.

    :param request:
        `django.http.HttpRequest` object. Can contain a `export_type` GET
        param (the default is 'xml' if no param is specified).
    :param result_id:
        The id of the requested artifact.
    :param export_fn:
        Export function, which accepts the following params:

            * result_id (int)
            * target (a file path or file-like)
            * export_type (option kwarg)

    :returns:
        If the requested ``result_id`` is not available in the format
        designated by the `export_type`.

        Otherwise, return a `django.http.HttpResponse` containing the content
        of the requested artifact.
    """
    export_type = request.GET.get('export_type', DEFAULT_EXPORT_TYPE)

    content = StringIO.StringIO()
    try:
        content = export_fn(result_id, content, export_type=export_type)
    except NotImplementedError, err:
        # Throw back a 501 if the exact export parameters are not supported
        return HttpResponse(err.message, status=NOT_IMPLEMENTED)

    # Just in case the original StringIO object was closed:
    resp_content = StringIO.StringIO()
    resp_content.writelines(content.buflist)
    del content

    # TODO(LB): A possible necessary optimization--in the future--would be to
    # iteratively stream large files.
    # TODO(LB): Large files could pose a memory consumption problem.
    resp_value = resp_content.getvalue()
    resp_content.close()
    return HttpResponse(resp_value)
<commit_msg>views:<commit_after>import StringIO
import json
import logging
import urlparse

from django.http import HttpResponse
from openquake.engine import engine as oq_engine
from openquake.engine.db import models as oqe_models
from openquake.engine.export import hazard as hazard_export
from openquake.engine.export import risk as risk_export

METHOD_NOT_ALLOWED = 405
NOT_IMPLEMENTED = 501
JSON = 'application/json'

IGNORE_FIELDS = ('base_path', 'export_dir', 'owner')
GEOM_FIELDS = ('region', 'sites', 'region_constraint', 'sites_disagg')
RISK_INPUTS = ('hazard_calculation', 'hazard_output')

#: For exporting calculation outputs, the client can request a specific format
#: (xml, geojson, csv, etc.). If the client does not specify give them (NRML)
#: XML by default.
DEFAULT_EXPORT_TYPE = 'xml'

LOGGER = logging.getLogger('openquakeserver')

class allowed_methods(object):
    """
    Use as a view decorator to strictly enforce HTTP request types.

    It works like this:

    .. code-block:: python

        @allowed_methods(('GET', 'POST'))
        def some_view(request):
            # code goes here

        @ allowed_methods(('DELETE', ))
        def some_other_view(request):
            # code goes here
    """

    def __init__(self, methods):
        self.methods = methods

    def __call__(self, func):
        def wrapped(request, *args, **kwargs):
            if not request.method in self.methods:
                return HttpResponse(status=METHOD_NOT_ALLOWED)
            else:
                return func(request, *args, **kwargs)
        return wrapped

def _get_base_url(request):
    """
    Construct a base URL, given a request object.

    This comprises the protocol prefix (http:// or https://) and the host,
    which can include the port number. For example:
    http://www.openquake.org or https://www.openquake.org:8000.
    """
    if request.is_secure():
        base_url = 'https://%s'
    else:
        base_url = 'http://%s'
    base_url %= request.META['HTTP_HOST']
    return base_url

def _calc_to_response_data(calc):
    """
    Extract the calculation parameters into a dictionary.
    """
    fields = [x.name for x in calc._meta.fields if x.name not in IGNORE_FIELDS]
    response_data = {}
    for field_name in fields:
        try:
            value = getattr(calc, field_name)
            if value is not None:
                if field_name in GEOM_FIELDS:
                    response_data[field_name] = json.loads(value.geojson)
                elif field_name in RISK_INPUTS:
                    response_data[field_name] = value.id
                else:
                    response_data[field_name] = value
        except AttributeError:
            # Better that we miss an attribute than crash.
            pass
    return response_data

@allowed_methods(('GET', ))
def calc_hazard(request):
    """
    Get a list of risk calculations and report their id, status, description,
    and a url where more detailed information can be accessed.

    Responses are in JSON.
    """
    base_url = _get_base_url(request)

    haz_calc_data = _get_haz_calcs()

    response_data = []
    for hc_id, status, desc in haz_calc_data:
        url = urlparse.urljoin(base_url, 'v1/calc/hazard/%d' % hc_id)
        response_data.append(
            dict(id=hc_id, status=status, description=desc, url=url)
        )

    return HttpResponse(content=json.dumps(response_data),
                        content_type=JSON)

def _get_haz_calcs():
    """
    Helper function for get job+calculation data from the oq-engine database.

    Gets all hazard calculation records available.
    """
    return oqe_models.OqJob.objects\
        .select_related()\
        .filter(hazard_calculation__isnull=False)\
        .values_list('hazard_calculation',
                     'status',
                     'hazard_calculation__description')

@allowed_methods(('GET', ))
def calc_hazard_info(request, calc_id):
    """
    Get a JSON blob containing all of parameters for the given calculation
    (specified by ``calc_id``). Also includes the current job status (
    executing, complete, etc.).
    """
    response_data = _get_haz_calc_info(calc_id)

    return HttpResponse(content=json.dumps(response_data), content_type=JSON)

def _get_haz_calc_info(calc_id):
    """
    Helper function to get job info and hazard calculation params from the
    oq-engine DB, as a dictionary.
    """
    job = oqe_models.OqJob.objects\
        .select_related()\
        .get(hazard_calculation=calc_id)

    hc = job.hazard_calculation
    response_data = _calc_to_response_data(hc)

    response_data['status'] = job.status
    return response_data

@allowed_methods(('GET', ))
def calc_hazard_results(request, calc_id):
    """
    Get a summarized list of hazard calculation results for a given
    ``calc_id``. Result is a JSON array of objects containing the following
    attributes:

        * id
        * name
        * type (hazard_curve, hazard_map, etc.)
        * url (the exact url where the full result can be accessed)
    """
    base_url = _get_base_url(request)

    results = oq_engine.get_hazard_outputs(calc_id)

    response_data = []
    for result in results:
        url = urlparse.urljoin(base_url,
                               'v1/calc/hazard/result/%d' % result.id)
        datum = dict(
            id=result.id,
            name=result.display_name,
            type=result.output_type,
            url=url,
        )
        response_data.append(datum)

    return HttpResponse(content=json.dumps(response_data))

@allowed_methods(('GET', ))
def get_hazard_result(request, result_id):
    """
    Download a specific hazard result, by ``result_id``.

    Parameters for the GET request can include an `export_type`, such as 'xml',
    'geojson', 'csv', etc.
    """
    return _get_result(request, result_id, hazard_export.export)

@allowed_methods(('GET', ))
def calc_risk(request):
    """
    Get a list of risk calculations and report their id, status, description,
    and a url where more detailed information can be accessed.

    Responses are in JSON.
    """
    base_url = _get_base_url(request)

    risk_calc_data = _get_risk_calcs()

    response_data = []
    for hc_id, status, desc in risk_calc_data:
        url = urlparse.urljoin(base_url, 'v1/calc/risk/%d' % hc_id)
        response_data.append(
            dict(id=hc_id, status=status, description=desc, url=url)
        )

    return HttpResponse(content=json.dumps(response_data),
                        content_type=JSON)

def _get_risk_calcs():
    """
    Helper function for get job+calculation data from the oq-engine database.

    Gets all risk calculation records available.
    """
    return oqe_models.OqJob.objects\
        .select_related()\
        .filter(risk_calculation__isnull=False)\
        .values_list('risk_calculation',
                     'status',
                     'risk_calculation__description')

@allowed_methods(('GET', ))
def calc_risk_info(request, calc_id):
    """
    Get a JSON blob containing all of parameters for the given calculation
    (specified by ``calc_id``). Also includes the current job status (
    executing, complete, etc.).
    """
    response_data = _get_risk_calc_info(calc_id)

    return HttpResponse(content=json.dumps(response_data), content_type=JSON)

def _get_risk_calc_info(calc_id):
    """
    Helper function to get job info and hazard calculation params from the
    oq-engine DB, as a dictionary.
    """
    job = oqe_models.OqJob.objects\
        .select_related()\
        .get(risk_calculation=calc_id)

    rc = job.risk_calculation
    response_data = _calc_to_response_data(rc)

    response_data['status'] = job.status
    return response_data

@allowed_methods(('GET', ))
def calc_risk_results(request, calc_id):
    """
    Get a summarized list of risk calculation results for a given
    ``calc_id``. Result is a JSON array of objects containing the following
    attributes:

        * id
        * name
        * type (hazard_curve, hazard_map, etc.)
        * url (the exact url where the full result can be accessed)
    """
    base_url = _get_base_url(request)

    results = oq_engine.get_risk_outputs(calc_id)

    response_data = []
    for result in results:
        url = urlparse.urljoin(base_url,
                               'v1/calc/risk/result/%d' % result.id)
        datum = dict(
            id=result.id,
            name=result.display_name,
            type=result.output_type,
            url=url,
        )
        response_data.append(datum)

    return HttpResponse(content=json.dumps(response_data))

@allowed_methods(('GET', ))
def get_risk_result(request, result_id):
    """
    Download a specific hazard result, by ``result_id``.

    Parameters for the GET request can include an `export_type`, such as 'xml',
    'geojson', 'csv', etc.
    """
    return _get_result(request, result_id, risk_export.export)

def _get_result(request, result_id, export_fn):
    """
    The common abstracted functionality for getting hazard or risk results.
    The functionality is the same, except for the hazard/risk specific
    ``export_fn``.

    :param request:
        `django.http.HttpRequest` object. Can contain a `export_type` GET
        param (the default is 'xml' if no param is specified).
    :param result_id:
        The id of the requested artifact.
    :param export_fn:
        Export function, which accepts the following params:

            * result_id (int)
            * target (a file path or file-like)
            * export_type (option kwarg)

    :returns:
        If the requested ``result_id`` is not available in the format
        designated by the `export_type`.

        Otherwise, return a `django.http.HttpResponse` containing the content
        of the requested artifact.
    """
    export_type = request.GET.get('export_type', DEFAULT_EXPORT_TYPE)

    content = StringIO.StringIO()
    try:
        content = export_fn(result_id, content, export_type=export_type)
    except NotImplementedError, err:
        # Throw back a 501 if the exact export parameters are not supported
        return HttpResponse(err.message, status=NOT_IMPLEMENTED)

    # Just in case the original StringIO object was closed:
    resp_content = StringIO.StringIO()
    resp_content.writelines(content.buflist)
    del content

    # TODO(LB): A possible necessary optimization--in the future--would be to
    # iteratively stream large files.
    # TODO(LB): Large files could pose a memory consumption problem.
    resp_value = resp_content.getvalue()
    resp_content.close()
    return HttpResponse(resp_value)
<|endoftext|>


<commit_before># This module is for formatting and writing unit-tests in python. The general format is as follows
# 1. Use start to start a test and give it, as an argument, the name of the test
# 2. Use whatever check functions are relevant to test the run
# 3. Use finish to signal the end of the test
# 4. Repeat stages 1-3 as the number of tests you want to run in the file
# 5. Use print_results_and_exit to print the number of tests and assertions that passed/failed in the correct format
#    before exiting with 0 if all tests passed or with 1 if there was a failed test

import os, sys, subprocess, traceback, platform
import pyrealsense2 as rs

n_assertions = 0
n_failed_assertions = 0
n_tests = 0
n_failed_tests = 0
test_failed = False
test_in_progress = False

# If this is the first time running this script we set the wanted environment, however it is impossible to change the current running
# environment so we rerun the script in a child process that inherits the environment we set
def set_env_vars(env_vars):
    if len(sys.argv) < 2:
        for env_var, val in env_vars.items():
            os.environ[env_var] = val
        sys.argv.append("rerun")
        if platform.system() == 'Linux' and "microsoft" not in platform.uname()[3].lower():
            cmd = ["python3"] + sys.argv
        else:
            cmd = ["py", "-3"] + sys.argv
        p = subprocess.run( cmd, stderr=subprocess.PIPE, universal_newlines=True )
        exit(p.returncode)

# Returns the first device that was found, if no device is found the test is skipped. That way we can still run
# the unit-tests when no device is connected and not fail the tests that check a connected device
def find_first_device_or_exit():
    c = rs.context()
    if not c.devices.size():  # if no device is connected we skip the test
        print("No device found, skipping test")
        exit(0)
    return c.devices[0]

# Function for printing the current call stack. Used when an assertion fails
def print_stack():
    for line in traceback.format_stack():
        print(line)

# Functions for asserting test cases:
# The check family of functions tests an expression and continues the test whether the assertion succeeded or failed.
# The require family are equivalent but execution is aborted if the assertion fails.

# Function for when a check fails
def check_failed():
    global n_failed_assertions, test_failed
    n_failed_assertions += 1
    test_failed = True

def abort():
    print("A was specified in a failed check. Aborting test")
    exit(1)

# Receive an expression which is an assertion. If false the assertion failed.
def check(exp, abort_if_failed = False):
    global n_assertions
    n_assertions += 1
    if not exp:
        print("Check failed, received", exp)
        check_failed()
        print_stack()
        if abort_if_failed:
            abort()
        return False
    return True

# Receives the resulted value and the expected value and asserts they are equal
def check_equal(result, expected, abort_if_failed = False):
    check(type(expected) != list)
    global n_assertions
    n_assertions += 1
    if result != expected:
        print("Result was:" + result + "\nBut we expected: " + expected)
        check_failed()
        print_stack()
        if abort_if_failed:
            abort()
        return False
    return True

# This function should never be reached
def check_not_reached( abort_if_failed = False ):
    check(False, abort_if_failed)

# This function should be put in except blocks that should not be reached.
# It's different from check_not_reached because it expects to be in an except block and prints the stack of the error
# and not the call-stack for this function
def check_no_exception( abort_if_failed = False ):
    global n_assertions
    n_assertions += 1
    check_failed()
    traceback.print_exc( file = sys.stdout )
    if abort_if_failed:
        abort()
        return False
    return True

# Receives 2 lists and asserts they are identical. python "equality" (using ==) requires same length & elements
# but not necessarily same ordering. Here we require exactly the same, including ordering.
def check_equal_lists(result, expected, abort_if_failed = False):
    global n_assertions
    n_assertions += 1
    failed = False
    if len(result) != len(expected):
        failed = True
        print("Check equal lists failed due to lists of different sizes:")
        print("The resulted list has", len(result), "elements, but the expected list has", len(expected), "elements")
    i = 0
    for res, exp in zip(result, expected):
        if res != exp:
            failed = True
            print("Check equal lists failed due to unequal elements:")
            print("The element of index", i, "in both lists was not equal")
        i += 1
    if failed:
        print("Result list:", result)
        print("Expected list:", expected)
        check_failed()
        print_stack()
        if abort_if_failed:
            abort()
        return False
    return True

# Receives an exception and asserts its type and message is as expected. This function is called with a caught exception,
def check_exception(exception, expected_type, expected_msg = None, abort_if_failed = False):
    check_equal(type(exception), expected_type, abort_if_failed)
    if expected_msg:
        check_equal(str(exception), expected_msg, abort_if_failed)

# Function for manually failing a test
def fail():
    global test_in_progress, n_failed_tests, test_failed
    if test_in_progress and not test_failed:
        n_failed_tests += 1
        test_failed = True

# Functions for formatting test cases
def start(test_name):
    global n_tests, test_failed, test_in_progress
    if test_in_progress:
        print("Tried to start test before previous test finished. Aborting test")
        exit(1)
    n_tests += 1
    test_failed = False
    test_in_progress = True
    print(test_name)

def finish():
    global test_failed, n_failed_tests, test_in_progress
    if not test_in_progress:
        print("Tried to finish a test without starting one")
        return
    if test_failed:
        n_failed_tests += 1
        print("Test failed")
    else:
        print("Test passed")
    print()
    test_in_progress = False

# The format has to agree with the expected format in check_log() in run-unit-tests and with the C++ format using Catch
def print_results_and_exit():
    global n_assertions, n_tests, n_failed_assertions, n_failed_tests
    if n_failed_assertions:
        passed = n_assertions - n_failed_assertions
        print("test cases:", n_tests, "|" , n_failed_tests,  "failed")
        print("assertions:", n_assertions, "|", passed, "passed |", n_failed_assertions, "failed")
        exit(1)
    print("All tests passed (" + str(n_assertions) + " assertions in " + str(n_tests) + " test cases)")
    exit(0)
<commit_msg>check no exception<commit_after># This module is for formatting and writing unit-tests in python. The general format is as follows
# 1. Use start to start a test and give it, as an argument, the name of the test
# 2. Use whatever check functions are relevant to test the run
# 3. Use finish to signal the end of the test
# 4. Repeat stages 1-3 as the number of tests you want to run in the file
# 5. Use print_results_and_exit to print the number of tests and assertions that passed/failed in the correct format
#    before exiting with 0 if all tests passed or with 1 if there was a failed test

import os, sys, subprocess, traceback, platform
import pyrealsense2 as rs

n_assertions = 0
n_failed_assertions = 0
n_tests = 0
n_failed_tests = 0
test_failed = False
test_in_progress = False

# If this is the first time running this script we set the wanted environment, however it is impossible to change the current running
# environment so we rerun the script in a child process that inherits the environment we set
def set_env_vars(env_vars):
    if len(sys.argv) < 2:
        for env_var, val in env_vars.items():
            os.environ[env_var] = val
        sys.argv.append("rerun")
        if platform.system() == 'Linux' and "microsoft" not in platform.uname()[3].lower():
            cmd = ["python3"] + sys.argv
        else:
            cmd = ["py", "-3"] + sys.argv
        p = subprocess.run( cmd, stderr=subprocess.PIPE, universal_newlines=True )
        exit(p.returncode)

# Returns the first device that was found, if no device is found the test is skipped. That way we can still run
# the unit-tests when no device is connected and not fail the tests that check a connected device
def find_first_device_or_exit():
    c = rs.context()
    if not c.devices.size():  # if no device is connected we skip the test
        print("No device found, skipping test")
        exit(0)
    return c.devices[0]

# Function for printing the current call stack. Used when an assertion fails
def print_stack():
    for line in traceback.format_stack():
        print(line)

# Functions for asserting test cases:
# The check family of functions tests an expression and continues the test whether the assertion succeeded or failed.
# The require family are equivalent but execution is aborted if the assertion fails.

# Function for when a check fails
def check_failed():
    global n_failed_assertions, test_failed
    n_failed_assertions += 1
    test_failed = True

def abort():
    print("A was specified in a failed check. Aborting test")
    exit(1)

# Receive an expression which is an assertion. If false the assertion failed.
def check(exp, abort_if_failed = False):
    global n_assertions
    n_assertions += 1
    if not exp:
        print("Check failed, received", exp)
        check_failed()
        print_stack()
        if abort_if_failed:
            abort()
        return False
    return True

# Receives the resulted value and the expected value and asserts they are equal
def check_equal(result, expected, abort_if_failed = False):
    check(type(expected) != list)
    global n_assertions
    n_assertions += 1
    if result != expected:
        print("Result was:" + result + "\nBut we expected: " + expected)
        check_failed()
        print_stack()
        if abort_if_failed:
            abort()
        return False
    return True

# This function should never be reached
def check_not_reached( abort_if_failed = False ):
    check(False, abort_if_failed)

# This function should be put in except blocks that should not be reached.
# It's different from check_not_reached because it expects to be in an except block and prints the stack of the error
# and not the call-stack for this function
def check_no_exception( abort_if_failed = False ):
    global n_assertions
    n_assertions += 1
    check_failed()
    traceback.print_exc( file = sys.stdout )
    if abort_if_failed:
        abort()

# Receives 2 lists and asserts they are identical. python "equality" (using ==) requires same length & elements
# but not necessarily same ordering. Here we require exactly the same, including ordering.
def check_equal_lists(result, expected, abort_if_failed = False):
    global n_assertions
    n_assertions += 1
    failed = False
    if len(result) != len(expected):
        failed = True
        print("Check equal lists failed due to lists of different sizes:")
        print("The resulted list has", len(result), "elements, but the expected list has", len(expected), "elements")
    i = 0
    for res, exp in zip(result, expected):
        if res != exp:
            failed = True
            print("Check equal lists failed due to unequal elements:")
            print("The element of index", i, "in both lists was not equal")
        i += 1
    if failed:
        print("Result list:", result)
        print("Expected list:", expected)
        check_failed()
        print_stack()
        if abort_if_failed:
            abort()
        return False
    return True

# Receives an exception and asserts its type and message is as expected. This function is called with a caught exception,
def check_exception(exception, expected_type, expected_msg = None, abort_if_failed = False):
    check_equal(type(exception), expected_type, abort_if_failed)
    if expected_msg:
        check_equal(str(exception), expected_msg, abort_if_failed)

# Function for manually failing a test
def fail():
    global test_in_progress, n_failed_tests, test_failed
    if test_in_progress and not test_failed:
        n_failed_tests += 1
        test_failed = True

# Functions for formatting test cases
def start(test_name):
    global n_tests, test_failed, test_in_progress
    if test_in_progress:
        print("Tried to start test before previous test finished. Aborting test")
        exit(1)
    n_tests += 1
    test_failed = False
    test_in_progress = True
    print(test_name)

def finish():
    global test_failed, n_failed_tests, test_in_progress
    if not test_in_progress:
        print("Tried to finish a test without starting one")
        return
    if test_failed:
        n_failed_tests += 1
        print("Test failed")
    else:
        print("Test passed")
    print()
    test_in_progress = False

# The format has to agree with the expected format in check_log() in run-unit-tests and with the C++ format using Catch
def print_results_and_exit():
    global n_assertions, n_tests, n_failed_assertions, n_failed_tests
    if n_failed_assertions:
        passed = n_assertions - n_failed_assertions
        print("test cases:", n_tests, "|" , n_failed_tests,  "failed")
        print("assertions:", n_assertions, "|", passed, "passed |", n_failed_assertions, "failed")
        exit(1)
    print("All tests passed (" + str(n_assertions) + " assertions in " + str(n_tests) + " test cases)")
    exit(0)
<|endoftext|>