<commit_before>import numpy as np

from mock import *

from atrcopy import SegmentData, AtariDosDiskImage, Dos33DiskImage,InvalidBinaryFile
from atrcopy.errors import *

class BaseFilesystemModifyTest(object):
    diskimage_type = None
    sample_file = None
    num_files_in_sample = 0

    def setup(self):
        data = np.fromfile(self.sample_file, dtype=np.uint8)
        rawdata = SegmentData(data)
        self.image = self.diskimage_type(rawdata)

    def check_entries(self, entries, prefix="TEST", save=None):
        orig_num_files = len(self.image.files)
        filenames = []
        count = 1
        for data in entries:
            filename = "%s%d.BIN" % (prefix, count)
            self.image.write_file(filename, None, data)
            assert len(self.image.files) == orig_num_files + count
            data2 = self.image.find_file(filename)
            assert data.tostring() == data2
            count += 1

        # loop over them again to make sure data wasn't overwritten
        count = 1
        for data in entries:
            filename = "%s%d.BIN" % (prefix, count)
            data2 = self.image.find_file(filename)
            assert data.tostring() == data2
            count += 1
            filenames.append(filename)

        if save is not None:
            self.image.save(save)

        return filenames

    def test_small(self):
        assert len(self.image.files) == self.num_files_in_sample

        data = np.asarray([0xff, 0xff, 0x00, 0x60, 0x01, 0x60, 1, 2], dtype=np.uint8)
        self.image.write_file("TEST.XEX", None, data)
        assert len(self.image.files) == self.num_files_in_sample + 1

        data2 = self.image.find_file("TEST.XEX")
        assert data.tostring() == data2

    def test_50k(self):
        assert len(self.image.files) == self.num_files_in_sample

        data = np.arange(50*1024, dtype=np.uint8)
        self.image.write_file("RAMP50K.BIN", None, data)
        assert len(self.image.files) == self.num_files_in_sample + 1

        data2 = self.image.find_file("RAMP50K.BIN")
        assert data.tostring() == data2

    def test_many_small(self):
        entries = [
            np.asarray([0xff, 0xff, 0x00, 0x60, 0x01, 0x60, 1, 2], dtype=np.uint8),
            np.arange(1*1024, dtype=np.uint8),
            np.arange(2*1024, dtype=np.uint8),
            np.arange(3*1024, dtype=np.uint8),
            np.arange(4*1024, dtype=np.uint8),
            np.arange(5*1024, dtype=np.uint8),
            np.arange(6*1024, dtype=np.uint8),
            np.arange(7*1024, dtype=np.uint8),
            np.arange(8*1024, dtype=np.uint8),
            np.arange(9*1024, dtype=np.uint8),
            np.arange(10*1024, dtype=np.uint8),
            ]
        self.check_entries(entries, save="many_small.atr")

    def test_big_failure(self):
        assert len(self.image.files) == self.num_files_in_sample

        data = np.arange(50*1024, dtype=np.uint8)
        self.image.write_file("RAMP50K.BIN", None, data)
        assert len(self.image.files) == self.num_files_in_sample + 1
        with pytest.raises(NotEnoughSpaceOnDisk):
            self.image.write_file("RAMP50K2.BIN", None, data)
        assert len(self.image.files) == self.num_files_in_sample + 1

    def test_delete(self):
        entries1 = [
            np.arange(3*1024, dtype=np.uint8),
            np.arange(3*1024, dtype=np.uint8),
            np.arange(3*1024, dtype=np.uint8),
            np.arange(3*1024, dtype=np.uint8),
            np.arange(3*1024, dtype=np.uint8),
            np.arange(3*1024, dtype=np.uint8),
            np.arange(3*1024, dtype=np.uint8),
            np.arange(3*1024, dtype=np.uint8),
            np.arange(3*1024, dtype=np.uint8),
            np.arange(10*1024, dtype=np.uint8),
            np.arange(10*1024, dtype=np.uint8),
            ]
        entries2 = [
            np.arange(10*1024, dtype=np.uint8),
            np.arange(11*1024, dtype=np.uint8),
        ]

        filenames = self.check_entries(entries1, "FIRST")
        assert len(self.image.files) == self.num_files_in_sample + 11
        self.image.delete_file(filenames[2])
        self.image.delete_file(filenames[5])
        self.image.delete_file(filenames[0])
        self.image.delete_file(filenames[8])
        assert len(self.image.files) == self.num_files_in_sample + 7

        filename = self.check_entries(entries2, "SECOND", save="test_delete.atr")
        assert len(self.image.files) == self.num_files_in_sample + 9

# class TestAtariDosSDImage(BaseFilesystemModifyTest):
#     diskimage_type = AtariDosDiskImage
#     sample_file = "../test_data/dos_sd_test1.atr"
#     num_files_in_sample = 5

class TestDos33Image(BaseFilesystemModifyTest):
    diskimage_type = Dos33DiskImage
    sample_file = "../test_data/dos33_master.dsk"
    num_files_in_sample = 19

if __name__ == "__main__":
    t = TestAtariDosFile()
    t.setup()
    t.test_segment()
<commit_msg>Changed file comparisons to numpy array_equals<commit_after>import numpy as np

from mock import *

from atrcopy import SegmentData, AtariDosDiskImage, Dos33DiskImage,InvalidBinaryFile
from atrcopy.errors import *

class BaseFilesystemModifyTest(object):
    diskimage_type = None
    sample_file = None
    num_files_in_sample = 0

    def setup(self):
        data = np.fromfile(self.sample_file, dtype=np.uint8)
        rawdata = SegmentData(data)
        self.image = self.diskimage_type(rawdata)

    def check_entries(self, entries, prefix="TEST", save=None):
        orig_num_files = len(self.image.files)
        filenames = []
        count = 1
        for data in entries:
            filename = "%s%d.BIN" % (prefix, count)
            self.image.write_file(filename, None, data)
            assert len(self.image.files) == orig_num_files + count
            data2 = np.fromstring(self.image.find_file(filename), dtype=np.uint8)
            assert np.array_equal(data, data2[0:len(data)])
            count += 1

        # loop over them again to make sure data wasn't overwritten
        count = 1
        for data in entries:
            filename = "%s%d.BIN" % (prefix, count)
            data2 = np.fromstring(self.image.find_file(filename), dtype=np.uint8)
            assert np.array_equal(data, data2[0:len(data)])
            count += 1
            filenames.append(filename)

        if save is not None:
            self.image.save(save)

        return filenames

    def test_small(self):
        assert len(self.image.files) == self.num_files_in_sample

        data = np.asarray([0xff, 0xff, 0x00, 0x60, 0x01, 0x60, 1, 2], dtype=np.uint8)
        self.image.write_file("TEST.XEX", None, data)
        assert len(self.image.files) == self.num_files_in_sample + 1

        data2 = np.fromstring(self.image.find_file("TEST.XEX"), dtype=np.uint8)
        assert np.array_equal(data, data2[0:len(data)])

    def test_50k(self):
        assert len(self.image.files) == self.num_files_in_sample

        data = np.arange(50*1024, dtype=np.uint8)
        self.image.write_file("RAMP50K.BIN", None, data)
        assert len(self.image.files) == self.num_files_in_sample + 1

        data2 = self.image.find_file("RAMP50K.BIN")
        assert data.tostring() == data2

    def test_many_small(self):
        entries = [
            np.asarray([0xff, 0xff, 0x00, 0x60, 0x01, 0x60, 1, 2], dtype=np.uint8),
            np.arange(1*1024, dtype=np.uint8),
            np.arange(2*1024, dtype=np.uint8),
            np.arange(3*1024, dtype=np.uint8),
            np.arange(4*1024, dtype=np.uint8),
            np.arange(5*1024, dtype=np.uint8),
            np.arange(6*1024, dtype=np.uint8),
            np.arange(7*1024, dtype=np.uint8),
            np.arange(8*1024, dtype=np.uint8),
            np.arange(9*1024, dtype=np.uint8),
            np.arange(10*1024, dtype=np.uint8),
            ]
        self.check_entries(entries, save="many_small.atr")

    def test_big_failure(self):
        assert len(self.image.files) == self.num_files_in_sample

        data = np.arange(50*1024, dtype=np.uint8)
        self.image.write_file("RAMP50K.BIN", None, data)
        assert len(self.image.files) == self.num_files_in_sample + 1
        with pytest.raises(NotEnoughSpaceOnDisk):
            self.image.write_file("RAMP50K2.BIN", None, data)
        assert len(self.image.files) == self.num_files_in_sample + 1

    def test_delete(self):
        entries1 = [
            np.arange(3*1024, dtype=np.uint8),
            np.arange(3*1024, dtype=np.uint8),
            np.arange(3*1024, dtype=np.uint8),
            np.arange(3*1024, dtype=np.uint8),
            np.arange(3*1024, dtype=np.uint8),
            np.arange(3*1024, dtype=np.uint8),
            np.arange(3*1024, dtype=np.uint8),
            np.arange(3*1024, dtype=np.uint8),
            np.arange(3*1024, dtype=np.uint8),
            np.arange(10*1024, dtype=np.uint8),
            np.arange(10*1024, dtype=np.uint8),
            ]
        entries2 = [
            np.arange(10*1024, dtype=np.uint8),
            np.arange(11*1024, dtype=np.uint8),
        ]

        filenames = self.check_entries(entries1, "FIRST")
        assert len(self.image.files) == self.num_files_in_sample + 11
        self.image.delete_file(filenames[2])
        self.image.delete_file(filenames[5])
        self.image.delete_file(filenames[0])
        self.image.delete_file(filenames[8])
        assert len(self.image.files) == self.num_files_in_sample + 7

        filename = self.check_entries(entries2, "SECOND", save="test_delete.atr")
        assert len(self.image.files) == self.num_files_in_sample + 9

class TestAtariDosSDImage(BaseFilesystemModifyTest):
    diskimage_type = AtariDosDiskImage
    sample_file = "../test_data/dos_sd_test1.atr"
    num_files_in_sample = 5

class TestDos33Image(BaseFilesystemModifyTest):
    diskimage_type = Dos33DiskImage
    sample_file = "../test_data/dos33_master.dsk"
    num_files_in_sample = 19

if __name__ == "__main__":
    t = TestAtariDosFile()
    t.setup()
    t.test_segment()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2012-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/>.

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

from openquake.hazardlib.const import StdDev
from openquake.hazardlib.gsim.base import ContextMaker
from openquake.hazardlib.imt import from_string

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__)

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` sites:
        Sites of interest to calculate GMFs.

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

    :param gsims:
        a set of GSIM instances

    :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.
    """
    # 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)` 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, sites, imts, gsims,
                 truncation_level=None, correlation_model=None, samples=0):
        if len(sites) == 0:
            raise ValueError('No sites')
        elif len(imts) == 0:
            raise ValueError('No IMTs')
        elif len(gsims) == 0:
            raise ValueError('No GSIMs')
        self.rupture = rupture
        self.sites = sites
        self.imts = [from_string(imt) for imt in imts]
        self.gsims = sorted(gsims)
        self.truncation_level = truncation_level
        self.correlation_model = correlation_model
        self.samples = samples
        # `rupture` can be a high level rupture object containing a low
        # level hazardlib rupture object as a .rupture attribute
        if hasattr(rupture, 'rupture'):
            rupture = rupture.rupture
        self.ctx = ContextMaker(gsims).make_contexts(sites, rupture)

    def compute(self, gsim, num_events, seed=None):
        """
        :param gsim: a GSIM instance
        :param num_events: the number of seismic events
        :param seed: a random seed or None
        :returns: a 32 bit array of shape (num_imts, num_sites, num_events)
        """
        try:  # read the seed from self.rupture.rupture if possible
            seed = seed or self.rupture.rupture.seed
        except AttributeError:
            pass
        if seed is not None:
            numpy.random.seed(seed)
        result = numpy.zeros(
            (len(self.imts), len(self.sites), num_events), numpy.float32)
        for imti, imt in enumerate(self.imts):
            result[imti] = self._compute(None, gsim, num_events, imt)
        return result

    def _compute(self, seed, gsim, num_events, imt):
        """
        :param seed: a random seed or None if the seed is already set
        :param gsim: a GSIM instance
        :param num_events: the number of seismic events
        :param imt: an IMT instance
        :returns: a 32 bit array of shape (num_sites, num_events)
        """
        if seed is not None:
            numpy.random.seed(seed)
        sctx, rctx, dctx = self.ctx
        if self.truncation_level == 0:
            assert self.correlation_model is None
            mean, _stddevs = gsim.get_mean_and_stddevs(
                sctx, rctx, dctx, imt, stddev_types=[])
            mean = gsim.to_imt_unit_values(mean)
            mean.shape += (1, )
            mean = mean.repeat(num_events, axis=1)
            return mean
        elif self.truncation_level is None:
            distribution = scipy.stats.norm()
        else:
            assert self.truncation_level > 0
            distribution = scipy.stats.truncnorm(
                - self.truncation_level, self.truncation_level)

        if gsim.DEFINED_FOR_STANDARD_DEVIATION_TYPES == \
           set([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] = gsim.get_mean_and_stddevs(
                sctx, rctx, dctx, imt, [StdDev.TOTAL])
            stddev_total = stddev_total.reshape(stddev_total.shape + (1, ))
            mean = mean.reshape(mean.shape + (1, ))

            total_residual = stddev_total * distribution.rvs(
                size=(len(self.sites), num_events))
            gmf = gsim.to_imt_unit_values(mean + total_residual)
        else:
            mean, [stddev_inter, stddev_intra] = gsim.get_mean_and_stddevs(
                sctx, rctx, dctx, imt,
                [StdDev.INTER_EVENT, StdDev.INTRA_EVENT])
            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 * distribution.rvs(
                size=(len(self.sites), num_events))

            if self.correlation_model is not None:
                ir = self.correlation_model.apply_correlation(
                    self.sites, imt, intra_residual)
                # this fixes a mysterious bug: ir[row] is actually
                # a matrix of shape (E, 1) and not a vector of size E
                intra_residual = numpy.zeros(ir.shape)
                for i, val in numpy.ndenumerate(ir):
                    intra_residual[i] = val

            inter_residual = stddev_inter * distribution.rvs(
                size=num_events)

            gmf = gsim.to_imt_unit_values(
                mean + intra_residual + inter_residual)

        return gmf

# 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.
    """
    gc = GmfComputer(rupture, sites, [str(imt) for imt in imts], [gsim],
                     truncation_level, correlation_model)
    res = gc.compute(gsim, realizations, seed)
    return {imt: res[imti] for imti, imt in enumerate(gc.imts)}
<commit_msg>Now the GmfComputer accepts an IntegrationDistance object in input<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2012-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/>.

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

from openquake.hazardlib.const import StdDev
from openquake.hazardlib.gsim.base import ContextMaker
from openquake.hazardlib.imt import from_string
from openquake.hazardlib.calc.filters import IntegrationDistance

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__)

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` sites:
        Sites of interest to calculate GMFs.

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

    :param gsims:
        a set of GSIM instances

    :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.
    """
    # 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)` 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, sites, imts, gsims,
                 truncation_level=None, correlation_model=None, samples=0,
                 maximum_distance=IntegrationDistance(None)):
        if len(sites) == 0:
            raise ValueError('No sites')
        elif len(imts) == 0:
            raise ValueError('No IMTs')
        elif len(gsims) == 0:
            raise ValueError('No GSIMs')
        self.rupture = rupture
        self.sites = sites
        self.imts = [from_string(imt) for imt in imts]
        self.gsims = sorted(gsims)
        self.truncation_level = truncation_level
        self.correlation_model = correlation_model
        self.samples = samples
        self.maximum_distance = maximum_distance
        # `rupture` can be a high level rupture object containing a low
        # level hazardlib rupture object as a .rupture attribute
        if hasattr(rupture, 'rupture'):
            rupture = rupture.rupture
        self.ctx = ContextMaker(gsims, maximum_distance).make_contexts(
            sites, rupture)

    def compute(self, gsim, num_events, seed=None):
        """
        :param gsim: a GSIM instance
        :param num_events: the number of seismic events
        :param seed: a random seed or None
        :returns: a 32 bit array of shape (num_imts, num_sites, num_events)
        """
        try:  # read the seed from self.rupture.rupture if possible
            seed = seed or self.rupture.rupture.seed
        except AttributeError:
            pass
        if seed is not None:
            numpy.random.seed(seed)
        result = numpy.zeros(
            (len(self.imts), len(self.sites), num_events), numpy.float32)
        for imti, imt in enumerate(self.imts):
            result[imti] = self._compute(None, gsim, num_events, imt)
        return result

    def _compute(self, seed, gsim, num_events, imt):
        """
        :param seed: a random seed or None if the seed is already set
        :param gsim: a GSIM instance
        :param num_events: the number of seismic events
        :param imt: an IMT instance
        :returns: a 32 bit array of shape (num_sites, num_events)
        """
        if seed is not None:
            numpy.random.seed(seed)
        sctx, rctx, dctx = self.ctx
        if self.truncation_level == 0:
            assert self.correlation_model is None
            mean, _stddevs = gsim.get_mean_and_stddevs(
                sctx, rctx, dctx, imt, stddev_types=[])
            mean = gsim.to_imt_unit_values(mean)
            mean.shape += (1, )
            mean = mean.repeat(num_events, axis=1)
            return mean
        elif self.truncation_level is None:
            distribution = scipy.stats.norm()
        else:
            assert self.truncation_level > 0
            distribution = scipy.stats.truncnorm(
                - self.truncation_level, self.truncation_level)

        if gsim.DEFINED_FOR_STANDARD_DEVIATION_TYPES == \
           set([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] = gsim.get_mean_and_stddevs(
                sctx, rctx, dctx, imt, [StdDev.TOTAL])
            stddev_total = stddev_total.reshape(stddev_total.shape + (1, ))
            mean = mean.reshape(mean.shape + (1, ))

            total_residual = stddev_total * distribution.rvs(
                size=(len(self.sites), num_events))
            gmf = gsim.to_imt_unit_values(mean + total_residual)
        else:
            mean, [stddev_inter, stddev_intra] = gsim.get_mean_and_stddevs(
                sctx, rctx, dctx, imt,
                [StdDev.INTER_EVENT, StdDev.INTRA_EVENT])
            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 * distribution.rvs(
                size=(len(self.sites), num_events))

            if self.correlation_model is not None:
                ir = self.correlation_model.apply_correlation(
                    self.sites, imt, intra_residual)
                # this fixes a mysterious bug: ir[row] is actually
                # a matrix of shape (E, 1) and not a vector of size E
                intra_residual = numpy.zeros(ir.shape)
                for i, val in numpy.ndenumerate(ir):
                    intra_residual[i] = val

            inter_residual = stddev_inter * distribution.rvs(
                size=num_events)

            gmf = gsim.to_imt_unit_values(
                mean + intra_residual + inter_residual)

        return gmf

# 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.
    """
    gc = GmfComputer(rupture, sites, [str(imt) for imt in imts], [gsim],
                     truncation_level, correlation_model)
    res = gc.compute(gsim, realizations, seed)
    return {imt: res[imti] for imti, imt in enumerate(gc.imts)}
<|endoftext|>


<commit_before>__author__ = 'Gareth Coles'

"""Tests for the command manager"""

import logging
import nose

import nose.tools as nosetools
from mock import MagicMock as Mock

from system.command_manager import CommandManager
from system.enums import CommandState

class test_commands:

    def __init__(self):
        self.manager = CommandManager()
        self.manager.logger.setLevel(logging.CRITICAL)  # Shut up, logger

        self.factory_manager = Mock(name="factory_manager")
        self.plugin = Mock(name="plugin")

    @nosetools.nottest
    def teardown(self):
        # Clean up
        self.manager.commands = {}
        self.manager.aliases = {}
        self.manager.auth_handlers = []
        self.manager.perm_handler = None

        self.plugin.reset_mock()
        self.plugin.handler.reset_mock()

        self.factory_manager.reset_mock()

    @nose.with_setup(teardown=teardown)
    def test_singleton(self):
        """CMNDS | Test Singleton metaclass"""
        nosetools.assert_true(self.manager is CommandManager())

    @nose.with_setup(teardown=teardown)
    def test_set_factory_manager(self):
        """CMNDS | Test setting factory manager"""
        self.manager.set_factory_manager(self.factory_manager)
        nosetools.assert_true(self.factory_manager)

    @nose.with_setup(teardown=teardown)
    def test_add_command(self):
        """CMNDS | Test adding commands"""
        r = self.manager.register_command("test", self.plugin.handler,
                                          self.plugin, "test.test",
                                          ["test2"], True)

        nosetools.assert_true(r)
        nosetools.assert_true("test" in self.manager.commands)

        command = self.manager.commands.get("test", None)

        if command:
            nosetools.assert_true("f" in command)
            nosetools.assert_true(command.get("f") is self.plugin.handler)

            nosetools.assert_true("permission" in command)
            nosetools.assert_true(command.get("permission") == "test.test")

            nosetools.assert_true("owner" in command)
            nosetools.assert_true(command.get("owner") is self.plugin)

            nosetools.assert_true("default" in command)
            nosetools.assert_true(command.get("default"))

        nosetools.assert_true("test2" in self.manager.aliases)

        alias = self.manager.aliases.get("test2", None)

        if alias:
            nosetools.assert_true(alias == "test")

        r = self.manager.register_command("test", self.plugin.handler,
                                          self.plugin, "test.test",
                                          ["test2"], True)

        nosetools.assert_false(r)

    @nose.with_setup(teardown=teardown)
    def test_unregister_commands(self):
        """CMNDS | Test unregistering commands"""
        self.manager.register_command("test1", self.plugin.handler,
                                      self.plugin, aliases=["test11"])
        self.manager.register_command("test2", self.plugin.handler,
                                      self.plugin, aliases=["test22"])
        self.manager.register_command("test3", self.plugin.handler,
                                      self.plugin, aliases=["test33"])

        nosetools.assert_equals(len(self.manager.commands), 3)
        nosetools.assert_equals(len(self.manager.aliases), 3)

        self.manager.unregister_commands_for_owner(self.plugin)

        nosetools.assert_equals(len(self.manager.commands), 0)
        nosetools.assert_equals(len(self.manager.aliases), 0)

    @nose.with_setup(teardown=teardown)
    def test_run_commands_defaults(self):
        """CMNDS | Test running commands directly | Defaults"""

        self.manager.register_command("test4", self.plugin.handler,
                                      self.plugin, aliases=["test5"],
                                      default=True)

        caller = Mock(name="caller")
        source = Mock(name="source")
        protocol = Mock(name="protocol")

        # Testing defaults

        r = self.manager.run_command("test4", caller, source, protocol, "")
        nosetools.assert_equals(r, (CommandState.Success, None))

        r = self.manager.run_command("test5", caller, source, protocol, "")
        nosetools.assert_equals(r, (CommandState.Success, None))

        nosetools.assert_equals(self.plugin.handler.call_count, 2)

    @nose.with_setup(teardown=teardown)
    def test_run_commands_aliases(self):
        """CMNDS | Test running commands directly | Aliases"""

        self.manager.register_command("test4", self.plugin.handler,
                                      self.plugin, aliases=["test5"],
                                      default=True)

        caller = Mock(name="caller")
        source = Mock(name="source")
        protocol = Mock(name="protocol")

        # Testing defaults

        r = self.manager.run_command("test4", caller, source, protocol, "")
        nosetools.assert_equals(r, (CommandState.Success, None))

        self.plugin.handler.assert_called_with(protocol, caller, source,
                                               "test4", "", [])

        # Reset mock
        self.plugin.handler.reset_mock()

        r = self.manager.run_command("test5", caller, source, protocol, "")
        nosetools.assert_equals(r, (CommandState.Success, None))

        self.plugin.handler.assert_called_with(protocol, caller, source,
                                               "test4", "", [])

        # Reset mock
        self.plugin.handler.reset_mock()

        self.manager.register_command("test5", self.plugin.handler,
                                      self.plugin, default=True)

        r = self.manager.run_command("test5", caller, source, protocol, "")
        nosetools.assert_equals(r, (CommandState.Success, None))

        self.plugin.handler.assert_called_with(protocol, caller, source,
                                               "test5", "", [])

    @nose.with_setup(teardown=teardown)
    def test_run_commands_auth(self):
        """CMNDS | Test running commands directly | Auth"""

        caller = Mock(name="caller")
        source = Mock(name="source")
        protocol = Mock(name="protocol")

        auth = Mock(name="auth")
        perms = Mock(name="perms")

        self.manager.add_auth_handler(auth)
        self.manager.set_permissions_handler(perms)

        ### COMMAND WITH PERMISSION ###
        auth.authorized.return_value = True
        perms.check.return_value = True

        self.manager.register_command("test5", self.plugin.handler,
                                      self.plugin, "test.test",
                                      aliases=["test6"])

        r = self.manager.run_command("test5", caller, source, protocol, "")

        nosetools.assert_equals(r, (CommandState.Success, None))
        nosetools.assert_equals(self.plugin.handler.call_count, 1)

        auth.authorized.reset_mock()
        perms.check.reset_mock()
        self.plugin.handler.reset_mock()

        ### ALIAS WITH PERMISSION ###

        r = self.manager.run_command("test6", caller, source, protocol, "")

        nosetools.assert_equals(r, (CommandState.Success, None))
        nosetools.assert_equals(self.plugin.handler.call_count, 1)

        auth.authorized.reset_mock()
        perms.check.reset_mock()

        auth.authorized.return_value = True
        perms.check.return_value = False
        self.plugin.handler.reset_mock()

        ### COMMAND WITHOUT PERMISSION ###

        r = self.manager.run_command("test5", caller, source, protocol, "")

        nosetools.assert_equals(r, (CommandState.NoPermission, None))
        nosetools.assert_equals(self.plugin.handler.call_count, 0)

        perms.check.return_value = True

        auth.authorized.reset_mock()
        perms.check.reset_mock()
        self.plugin.handler.reset_mock()

        ### COMMAND WITH EXCEPTION ###

        self.plugin.handler = Mock(side_effect=Exception('Boom!'))
        self.manager.unregister_commands_for_owner(self.plugin)

        self.manager.register_command("test5", self.plugin.handler,
                                      self.plugin, "test.test")

        r = self.manager.run_command("test5", caller, source, protocol, "")

        nosetools.assert_equals(r[0], CommandState.Error)
        nosetools.assert_true(isinstance(r[1], Exception))
        nosetools.assert_equals(self.plugin.handler.call_count, 1)

        auth.authorized.reset_mock()
        perms.check.reset_mock()
        self.plugin.handler.reset_mock()

        ### UNKNOWN COMMAND ###

        r = self.manager.run_command("test7", caller, source, protocol, "")
        nosetools.assert_equals(r, (CommandState.Unknown, None))
        nosetools.assert_equals(self.plugin.handler.call_count, 0)
<commit_msg>Fix nosetests<commit_after>__author__ = 'Gareth Coles'

"""Tests for the command manager"""

import logging
import nose

import nose.tools as nosetools
from mock import MagicMock as Mock

from system.command_manager import CommandManager
from system.enums import CommandState

class test_commands:

    def __init__(self):
        self.manager = CommandManager()
        self.manager.logger.setLevel(logging.CRITICAL)  # Shut up, logger

        self.factory_manager = Mock(name="factory_manager")
        self.plugin = Mock(name="plugin")

    @nosetools.nottest
    def teardown(self):
        # Clean up
        self.manager.commands = {}
        self.manager.aliases = {}
        self.manager.auth_handler = None
        self.manager.perm_handler = None

        self.plugin.reset_mock()
        self.plugin.handler.reset_mock()

        self.factory_manager.reset_mock()

    @nose.with_setup(teardown=teardown)
    def test_singleton(self):
        """CMNDS | Test Singleton metaclass"""
        nosetools.assert_true(self.manager is CommandManager())

    @nose.with_setup(teardown=teardown)
    def test_set_factory_manager(self):
        """CMNDS | Test setting factory manager"""
        self.manager.set_factory_manager(self.factory_manager)
        nosetools.assert_true(self.factory_manager)

    @nose.with_setup(teardown=teardown)
    def test_add_command(self):
        """CMNDS | Test adding commands"""
        r = self.manager.register_command("test", self.plugin.handler,
                                          self.plugin, "test.test",
                                          ["test2"], True)

        nosetools.assert_true(r)
        nosetools.assert_true("test" in self.manager.commands)

        command = self.manager.commands.get("test", None)

        if command:
            nosetools.assert_true("f" in command)
            nosetools.assert_true(command.get("f") is self.plugin.handler)

            nosetools.assert_true("permission" in command)
            nosetools.assert_true(command.get("permission") == "test.test")

            nosetools.assert_true("owner" in command)
            nosetools.assert_true(command.get("owner") is self.plugin)

            nosetools.assert_true("default" in command)
            nosetools.assert_true(command.get("default"))

        nosetools.assert_true("test2" in self.manager.aliases)

        alias = self.manager.aliases.get("test2", None)

        if alias:
            nosetools.assert_true(alias == "test")

        r = self.manager.register_command("test", self.plugin.handler,
                                          self.plugin, "test.test",
                                          ["test2"], True)

        nosetools.assert_false(r)

    @nose.with_setup(teardown=teardown)
    def test_unregister_commands(self):
        """CMNDS | Test unregistering commands"""
        self.manager.register_command("test1", self.plugin.handler,
                                      self.plugin, aliases=["test11"])
        self.manager.register_command("test2", self.plugin.handler,
                                      self.plugin, aliases=["test22"])
        self.manager.register_command("test3", self.plugin.handler,
                                      self.plugin, aliases=["test33"])

        nosetools.assert_equals(len(self.manager.commands), 3)
        nosetools.assert_equals(len(self.manager.aliases), 3)

        self.manager.unregister_commands_for_owner(self.plugin)

        nosetools.assert_equals(len(self.manager.commands), 0)
        nosetools.assert_equals(len(self.manager.aliases), 0)

    @nose.with_setup(teardown=teardown)
    def test_run_commands_defaults(self):
        """CMNDS | Test running commands directly | Defaults"""

        self.manager.register_command("test4", self.plugin.handler,
                                      self.plugin, aliases=["test5"],
                                      default=True)

        caller = Mock(name="caller")
        source = Mock(name="source")
        protocol = Mock(name="protocol")

        # Testing defaults

        r = self.manager.run_command("test4", caller, source, protocol, "")
        nosetools.assert_equals(r, (CommandState.Success, None))

        r = self.manager.run_command("test5", caller, source, protocol, "")
        nosetools.assert_equals(r, (CommandState.Success, None))

        nosetools.assert_equals(self.plugin.handler.call_count, 2)

    @nose.with_setup(teardown=teardown)
    def test_run_commands_aliases(self):
        """CMNDS | Test running commands directly | Aliases"""

        self.manager.register_command("test4", self.plugin.handler,
                                      self.plugin, aliases=["test5"],
                                      default=True)

        caller = Mock(name="caller")
        source = Mock(name="source")
        protocol = Mock(name="protocol")

        # Testing defaults

        r = self.manager.run_command("test4", caller, source, protocol, "")
        nosetools.assert_equals(r, (CommandState.Success, None))

        self.plugin.handler.assert_called_with(protocol, caller, source,
                                               "test4", "", [])

        # Reset mock
        self.plugin.handler.reset_mock()

        r = self.manager.run_command("test5", caller, source, protocol, "")
        nosetools.assert_equals(r, (CommandState.Success, None))

        self.plugin.handler.assert_called_with(protocol, caller, source,
                                               "test4", "", [])

        # Reset mock
        self.plugin.handler.reset_mock()

        self.manager.register_command("test5", self.plugin.handler,
                                      self.plugin, default=True)

        r = self.manager.run_command("test5", caller, source, protocol, "")
        nosetools.assert_equals(r, (CommandState.Success, None))

        self.plugin.handler.assert_called_with(protocol, caller, source,
                                               "test5", "", [])

    @nose.with_setup(teardown=teardown)
    def test_run_commands_auth(self):
        """CMNDS | Test running commands directly | Auth"""

        caller = Mock(name="caller")
        source = Mock(name="source")
        protocol = Mock(name="protocol")

        auth = Mock(name="auth")
        perms = Mock(name="perms")

        self.manager.set_auth_handler(auth)
        self.manager.set_permissions_handler(perms)

        ### COMMAND WITH PERMISSION ###
        auth.authorized.return_value = True
        perms.check.return_value = True

        self.manager.register_command("test5", self.plugin.handler,
                                      self.plugin, "test.test",
                                      aliases=["test6"])

        r = self.manager.run_command("test5", caller, source, protocol, "")

        nosetools.assert_equals(r, (CommandState.Success, None))
        nosetools.assert_equals(self.plugin.handler.call_count, 1)

        auth.authorized.reset_mock()
        perms.check.reset_mock()
        self.plugin.handler.reset_mock()

        ### ALIAS WITH PERMISSION ###

        r = self.manager.run_command("test6", caller, source, protocol, "")

        nosetools.assert_equals(r, (CommandState.Success, None))
        nosetools.assert_equals(self.plugin.handler.call_count, 1)

        auth.authorized.reset_mock()
        perms.check.reset_mock()

        auth.authorized.return_value = True
        perms.check.return_value = False
        self.plugin.handler.reset_mock()

        ### COMMAND WITHOUT PERMISSION ###

        r = self.manager.run_command("test5", caller, source, protocol, "")

        nosetools.assert_equals(r, (CommandState.NoPermission, None))
        nosetools.assert_equals(self.plugin.handler.call_count, 0)

        perms.check.return_value = True

        auth.authorized.reset_mock()
        perms.check.reset_mock()
        self.plugin.handler.reset_mock()

        ### COMMAND WITH EXCEPTION ###

        self.plugin.handler = Mock(side_effect=Exception('Boom!'))
        self.manager.unregister_commands_for_owner(self.plugin)

        self.manager.register_command("test5", self.plugin.handler,
                                      self.plugin, "test.test")

        r = self.manager.run_command("test5", caller, source, protocol, "")

        nosetools.assert_equals(r[0], CommandState.Error)
        nosetools.assert_true(isinstance(r[1], Exception))
        nosetools.assert_equals(self.plugin.handler.call_count, 1)

        auth.authorized.reset_mock()
        perms.check.reset_mock()
        self.plugin.handler.reset_mock()

        ### UNKNOWN COMMAND ###

        r = self.manager.run_command("test7", caller, source, protocol, "")
        nosetools.assert_equals(r, (CommandState.Unknown, None))
        nosetools.assert_equals(self.plugin.handler.call_count, 0)
<|endoftext|>


<commit_before>from .exceptions import GrammarError
from .lexer import Token
from .tree import Tree
from .visitors import InlineTransformer # XXX Deprecated
from .visitors import Transformer_InPlace
from .visitors import _vargs_meta, _vargs_meta_inline

###{standalone
from functools import partial, wraps
from itertools import repeat, product

class ExpandSingleChild:
    def __init__(self, node_builder):
        self.node_builder = node_builder

    def __call__(self, children):
        if len(children) == 1:
            return children[0]
        else:
            return self.node_builder(children)

class PropagatePositions:
    def __init__(self, node_builder):
        self.node_builder = node_builder

    def __call__(self, children):
        res = self.node_builder(children)

        if isinstance(res, Tree):
            for c in children:
                if isinstance(c, Tree) and not c.meta.empty:
                    res.meta.line = c.meta.line
                    res.meta.column = c.meta.column
                    res.meta.start_pos = c.meta.start_pos
                    res.meta.empty = False
                    break
                elif isinstance(c, Token):
                    res.meta.line = c.line
                    res.meta.column = c.column
                    res.meta.start_pos = c.pos_in_stream
                    res.meta.empty = False
                    break

            for c in reversed(children):
                if isinstance(c, Tree) and not c.meta.empty:
                    res.meta.end_line = c.meta.end_line
                    res.meta.end_column = c.meta.end_column
                    res.meta.end_pos = c.meta.end_pos
                    res.meta.empty = False
                    break
                elif isinstance(c, Token):
                    res.meta.end_line = c.end_line
                    res.meta.end_column = c.end_column
                    res.meta.end_pos = c.end_pos
                    res.meta.empty = False
                    break

        return res

class ChildFilter:
    def __init__(self, to_include, append_none, node_builder):
        self.node_builder = node_builder
        self.to_include = to_include
        self.append_none = append_none

    def __call__(self, children):
        filtered = []

        for i, to_expand, add_none in self.to_include:
            if add_none:
                filtered += [None] * add_none
            if to_expand:
                filtered += children[i].children
            else:
                filtered.append(children[i])

        if self.append_none:
            filtered += [None] * self.append_none

        return self.node_builder(filtered)

class ChildFilterLALR(ChildFilter):
    "Optimized childfilter for LALR (assumes no duplication in parse tree, so it's safe to change it)"

    def __call__(self, children):
        filtered = []
        for i, to_expand, add_none in self.to_include:
            if add_none:
                filtered += [None] * add_none
            if to_expand:
                if filtered:
                    filtered += children[i].children
                else:   # Optimize for left-recursion
                    filtered = children[i].children
            else:
                filtered.append(children[i])

        if self.append_none:
            filtered += [None] * self.append_none

        return self.node_builder(filtered)

class ChildFilterLALR_NoPlaceholders(ChildFilter):
    "Optimized childfilter for LALR (assumes no duplication in parse tree, so it's safe to change it)"
    def __init__(self, to_include, node_builder):
        self.node_builder = node_builder
        self.to_include = to_include

    def __call__(self, children):
        filtered = []
        for i, to_expand in self.to_include:
            if to_expand:
                if filtered:
                    filtered += children[i].children
                else:   # Optimize for left-recursion
                    filtered = children[i].children
            else:
                filtered.append(children[i])
        return self.node_builder(filtered)

def _should_expand(sym):
    return not sym.is_term and sym.name.startswith('_')

def maybe_create_child_filter(expansion, keep_all_tokens, ambiguous, _empty_indices):
    # Prepare empty_indices as: How many Nones to insert at each index?
    if _empty_indices:
        assert _empty_indices.count(False) == len(expansion)
        s = ''.join(str(int(b)) for b in _empty_indices)
        empty_indices = [len(ones) for ones in s.split('0')]
        assert len(empty_indices) == len(expansion)+1, (empty_indices, len(expansion))
    else:
        empty_indices = [0] * (len(expansion)+1)

    to_include = []
    nones_to_add = 0
    for i, sym in enumerate(expansion):
        nones_to_add += empty_indices[i]
        if keep_all_tokens or not (sym.is_term and sym.filter_out):
            to_include.append((i, _should_expand(sym), nones_to_add))
            nones_to_add = 0

    nones_to_add += empty_indices[len(expansion)]

    if _empty_indices or len(to_include) < len(expansion) or any(to_expand for i, to_expand,_ in to_include):
        if _empty_indices or ambiguous:
            return partial(ChildFilter if ambiguous else ChildFilterLALR, to_include, nones_to_add)
        else:
            # LALR without placeholders
            return partial(ChildFilterLALR_NoPlaceholders, [(i, x) for i,x,_ in to_include])

class AmbiguousExpander:
    """Deal with the case where we're expanding children ('_rule') into a parent but the children
       are ambiguous. i.e. (parent->_ambig->_expand_this_rule). In this case, make the parent itself
       ambiguous with as many copies as their are ambiguous children, and then copy the ambiguous children
       into the right parents in the right places, essentially shifting the ambiguiuty up the tree."""
    def __init__(self, to_expand, tree_class, node_builder):
        self.node_builder = node_builder
        self.tree_class = tree_class
        self.to_expand = to_expand

    def __call__(self, children):
        def _is_ambig_tree(child):
            return hasattr(child, 'data') and child.data == '_ambig'

        #### When we're repeatedly expanding ambiguities we can end up with nested ambiguities.
        #    All children of an _ambig node should be a derivation of that ambig node, hence
        #    it is safe to assume that if we see an _ambig node nested within an ambig node
        #    it is safe to simply expand it into the parent _ambig node as an alternative derivation.
        ambiguous = []
        for i, child in enumerate(children):
            if _is_ambig_tree(child):
                if i in self.to_expand:
                    ambiguous.append(i)

                to_expand = [j for j, grandchild in enumerate(child.children) if _is_ambig_tree(grandchild)]
                child.expand_kids_by_index(*to_expand)

        if not ambiguous:
            return self.node_builder(children)

        expand = [ iter(child.children) if i in ambiguous else repeat(child) for i, child in enumerate(children) ]
        return self.tree_class('_ambig', [self.node_builder(list(f[0])) for f in product(zip(*expand))])

def maybe_create_ambiguous_expander(tree_class, expansion, keep_all_tokens):
    to_expand = [i for i, sym in enumerate(expansion)
                 if keep_all_tokens or ((not (sym.is_term and sym.filter_out)) and _should_expand(sym))]
    if to_expand:
        return partial(AmbiguousExpander, to_expand, tree_class)

def ptb_inline_args(func):
    @wraps(func)
    def f(children):
        return func(*children)
    return f

def inplace_transformer(func):
    @wraps(func)
    def f(children):
        # function name in a Transformer is a rule name.
        tree = Tree(func.__name__, children)
        return func(tree)
    return f

def apply_visit_wrapper(func, name, wrapper):
    if wrapper is _vargs_meta or wrapper is _vargs_meta_inline:
        raise NotImplementedError("Meta args not supported for internal transformer")
    @wraps(func)
    def f(children):
        return wrapper(func, name, children, None)
    return f

class ParseTreeBuilder:
    def __init__(self, rules, tree_class, propagate_positions=False, keep_all_tokens=False, ambiguous=False, maybe_placeholders=False):
        self.tree_class = tree_class
        self.propagate_positions = propagate_positions
        self.always_keep_all_tokens = keep_all_tokens
        self.ambiguous = ambiguous
        self.maybe_placeholders = maybe_placeholders

        self.rule_builders = list(self._init_builders(rules))

    def _init_builders(self, rules):
        for rule in rules:
            options = rule.options
            keep_all_tokens = self.always_keep_all_tokens or options.keep_all_tokens
            expand_single_child = options.expand1

            wrapper_chain = list(filter(None, [
                (expand_single_child and not rule.alias) and ExpandSingleChild,
                maybe_create_child_filter(rule.expansion, keep_all_tokens, self.ambiguous, options.empty_indices if self.maybe_placeholders else None),
                self.propagate_positions and PropagatePositions,
                self.ambiguous and maybe_create_ambiguous_expander(self.tree_class, rule.expansion, keep_all_tokens),
            ]))

            yield rule, wrapper_chain

    def create_callback(self, transformer=None):
        callbacks = {}

        for rule, wrapper_chain in self.rule_builders:

            user_callback_name = rule.alias or rule.options.template_source or rule.origin.name
            try:
                f = getattr(transformer, user_callback_name)
                # XXX InlineTransformer is deprecated!
                wrapper = getattr(f, 'visit_wrapper', None)
                if wrapper is not None:
                    f = apply_visit_wrapper(f, user_callback_name, wrapper)
                else:
                    if isinstance(transformer, InlineTransformer):
                        f = ptb_inline_args(f)
                    elif isinstance(transformer, Transformer_InPlace):
                        f = inplace_transformer(f)
            except AttributeError:
                f = partial(self.tree_class, user_callback_name)

            for w in wrapper_chain:
                f = w(f)

            if rule in callbacks:
                raise GrammarError("Rule '%s' already exists" % (rule,))

            callbacks[rule] = f

        return callbacks

###}
<commit_msg>Reduce number of calls to Tree.meta in PropagatePositions<commit_after>from .exceptions import GrammarError
from .lexer import Token
from .tree import Tree
from .visitors import InlineTransformer # XXX Deprecated
from .visitors import Transformer_InPlace
from .visitors import _vargs_meta, _vargs_meta_inline

###{standalone
from functools import partial, wraps
from itertools import repeat, product

class ExpandSingleChild:
    def __init__(self, node_builder):
        self.node_builder = node_builder

    def __call__(self, children):
        if len(children) == 1:
            return children[0]
        else:
            return self.node_builder(children)

class PropagatePositions:
    def __init__(self, node_builder):
        self.node_builder = node_builder

    def __call__(self, children):
        res = self.node_builder(children)

        # local reference to Tree.meta reduces number of presence checks
        if isinstance(res, Tree):
            res_meta = res.meta
            for c in children:
                if isinstance(c, Tree):
                    child_meta = c.meta
                    if not child_meta.empty:
                        res_meta.line = child_meta.line
                        res_meta.column = child_meta.column
                        res_meta.start_pos = child_meta.start_pos
                        res_meta.empty = False
                        break
                elif isinstance(c, Token):
                    res_meta.line = c.line
                    res_meta.column = c.column
                    res_meta.start_pos = c.pos_in_stream
                    res_meta.empty = False
                    break

            for c in reversed(children):
                if isinstance(c, Tree):
                    child_meta = c.meta
                    if not child_meta.empty:
                        res_meta.end_line = child_meta.end_line
                        res_meta.end_column = child_meta.end_column
                        res_meta.end_pos = child_meta.end_pos
                        res_meta.empty = False
                        break
                elif isinstance(c, Token):
                    res_meta.end_line = c.end_line
                    res_meta.end_column = c.end_column
                    res_meta.end_pos = c.end_pos
                    res_meta.empty = False
                    break

        return res

class ChildFilter:
    def __init__(self, to_include, append_none, node_builder):
        self.node_builder = node_builder
        self.to_include = to_include
        self.append_none = append_none

    def __call__(self, children):
        filtered = []

        for i, to_expand, add_none in self.to_include:
            if add_none:
                filtered += [None] * add_none
            if to_expand:
                filtered += children[i].children
            else:
                filtered.append(children[i])

        if self.append_none:
            filtered += [None] * self.append_none

        return self.node_builder(filtered)

class ChildFilterLALR(ChildFilter):
    "Optimized childfilter for LALR (assumes no duplication in parse tree, so it's safe to change it)"

    def __call__(self, children):
        filtered = []
        for i, to_expand, add_none in self.to_include:
            if add_none:
                filtered += [None] * add_none
            if to_expand:
                if filtered:
                    filtered += children[i].children
                else:   # Optimize for left-recursion
                    filtered = children[i].children
            else:
                filtered.append(children[i])

        if self.append_none:
            filtered += [None] * self.append_none

        return self.node_builder(filtered)

class ChildFilterLALR_NoPlaceholders(ChildFilter):
    "Optimized childfilter for LALR (assumes no duplication in parse tree, so it's safe to change it)"
    def __init__(self, to_include, node_builder):
        self.node_builder = node_builder
        self.to_include = to_include

    def __call__(self, children):
        filtered = []
        for i, to_expand in self.to_include:
            if to_expand:
                if filtered:
                    filtered += children[i].children
                else:   # Optimize for left-recursion
                    filtered = children[i].children
            else:
                filtered.append(children[i])
        return self.node_builder(filtered)

def _should_expand(sym):
    return not sym.is_term and sym.name.startswith('_')

def maybe_create_child_filter(expansion, keep_all_tokens, ambiguous, _empty_indices):
    # Prepare empty_indices as: How many Nones to insert at each index?
    if _empty_indices:
        assert _empty_indices.count(False) == len(expansion)
        s = ''.join(str(int(b)) for b in _empty_indices)
        empty_indices = [len(ones) for ones in s.split('0')]
        assert len(empty_indices) == len(expansion)+1, (empty_indices, len(expansion))
    else:
        empty_indices = [0] * (len(expansion)+1)

    to_include = []
    nones_to_add = 0
    for i, sym in enumerate(expansion):
        nones_to_add += empty_indices[i]
        if keep_all_tokens or not (sym.is_term and sym.filter_out):
            to_include.append((i, _should_expand(sym), nones_to_add))
            nones_to_add = 0

    nones_to_add += empty_indices[len(expansion)]

    if _empty_indices or len(to_include) < len(expansion) or any(to_expand for i, to_expand,_ in to_include):
        if _empty_indices or ambiguous:
            return partial(ChildFilter if ambiguous else ChildFilterLALR, to_include, nones_to_add)
        else:
            # LALR without placeholders
            return partial(ChildFilterLALR_NoPlaceholders, [(i, x) for i,x,_ in to_include])

class AmbiguousExpander:
    """Deal with the case where we're expanding children ('_rule') into a parent but the children
       are ambiguous. i.e. (parent->_ambig->_expand_this_rule). In this case, make the parent itself
       ambiguous with as many copies as their are ambiguous children, and then copy the ambiguous children
       into the right parents in the right places, essentially shifting the ambiguiuty up the tree."""
    def __init__(self, to_expand, tree_class, node_builder):
        self.node_builder = node_builder
        self.tree_class = tree_class
        self.to_expand = to_expand

    def __call__(self, children):
        def _is_ambig_tree(child):
            return hasattr(child, 'data') and child.data == '_ambig'

        #### When we're repeatedly expanding ambiguities we can end up with nested ambiguities.
        #    All children of an _ambig node should be a derivation of that ambig node, hence
        #    it is safe to assume that if we see an _ambig node nested within an ambig node
        #    it is safe to simply expand it into the parent _ambig node as an alternative derivation.
        ambiguous = []
        for i, child in enumerate(children):
            if _is_ambig_tree(child):
                if i in self.to_expand:
                    ambiguous.append(i)

                to_expand = [j for j, grandchild in enumerate(child.children) if _is_ambig_tree(grandchild)]
                child.expand_kids_by_index(*to_expand)

        if not ambiguous:
            return self.node_builder(children)

        expand = [ iter(child.children) if i in ambiguous else repeat(child) for i, child in enumerate(children) ]
        return self.tree_class('_ambig', [self.node_builder(list(f[0])) for f in product(zip(*expand))])

def maybe_create_ambiguous_expander(tree_class, expansion, keep_all_tokens):
    to_expand = [i for i, sym in enumerate(expansion)
                 if keep_all_tokens or ((not (sym.is_term and sym.filter_out)) and _should_expand(sym))]
    if to_expand:
        return partial(AmbiguousExpander, to_expand, tree_class)

def ptb_inline_args(func):
    @wraps(func)
    def f(children):
        return func(*children)
    return f

def inplace_transformer(func):
    @wraps(func)
    def f(children):
        # function name in a Transformer is a rule name.
        tree = Tree(func.__name__, children)
        return func(tree)
    return f

def apply_visit_wrapper(func, name, wrapper):
    if wrapper is _vargs_meta or wrapper is _vargs_meta_inline:
        raise NotImplementedError("Meta args not supported for internal transformer")
    @wraps(func)
    def f(children):
        return wrapper(func, name, children, None)
    return f

class ParseTreeBuilder:
    def __init__(self, rules, tree_class, propagate_positions=False, keep_all_tokens=False, ambiguous=False, maybe_placeholders=False):
        self.tree_class = tree_class
        self.propagate_positions = propagate_positions
        self.always_keep_all_tokens = keep_all_tokens
        self.ambiguous = ambiguous
        self.maybe_placeholders = maybe_placeholders

        self.rule_builders = list(self._init_builders(rules))

    def _init_builders(self, rules):
        for rule in rules:
            options = rule.options
            keep_all_tokens = self.always_keep_all_tokens or options.keep_all_tokens
            expand_single_child = options.expand1

            wrapper_chain = list(filter(None, [
                (expand_single_child and not rule.alias) and ExpandSingleChild,
                maybe_create_child_filter(rule.expansion, keep_all_tokens, self.ambiguous, options.empty_indices if self.maybe_placeholders else None),
                self.propagate_positions and PropagatePositions,
                self.ambiguous and maybe_create_ambiguous_expander(self.tree_class, rule.expansion, keep_all_tokens),
            ]))

            yield rule, wrapper_chain

    def create_callback(self, transformer=None):
        callbacks = {}

        for rule, wrapper_chain in self.rule_builders:

            user_callback_name = rule.alias or rule.options.template_source or rule.origin.name
            try:
                f = getattr(transformer, user_callback_name)
                # XXX InlineTransformer is deprecated!
                wrapper = getattr(f, 'visit_wrapper', None)
                if wrapper is not None:
                    f = apply_visit_wrapper(f, user_callback_name, wrapper)
                else:
                    if isinstance(transformer, InlineTransformer):
                        f = ptb_inline_args(f)
                    elif isinstance(transformer, Transformer_InPlace):
                        f = inplace_transformer(f)
            except AttributeError:
                f = partial(self.tree_class, user_callback_name)

            for w in wrapper_chain:
                f = w(f)

            if rule in callbacks:
                raise GrammarError("Rule '%s' already exists" % (rule,))

            callbacks[rule] = f

        return callbacks

###}
<|endoftext|>


<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 exp(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  # sig and eps have shapes (M, E) instead
            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
        """
        M = len(self.imts)
        result = numpy.zeros(
            (len(self.imts), len(self.ctx.sids), num_events), F32)
        sig = numpy.zeros((M, num_events), F32)
        eps = numpy.zeros((M, num_events), F32)
        numpy.random.seed(self.seed)
        num_sids = len(self.ctx.sids)
        if self.distribution:
            # build M arrays of random numbers of shape (N+1, E)
            rnd = [rvs(self.distribution, num_sids + 1, num_events)
                   for _ in range(M)]
        else:
            rnd = [None] * M
        for imti, imt in enumerate(self.imts):
            try:
                result[imti], sig[imti], eps[imti] = self._compute(
                     mean_stds[:, imti], num_events, imt, gsim, rnd[imti])
            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, rnd):
        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
            gmf = exp(mean, imt)[:, None]
            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
            sig = sig[:, None]
            mean = mean[:, None]

            total_res = sig * rnd[:-1]
            gmf = exp(mean + total_res, imt)
            stdi = numpy.empty(num_events, F32)
            stdi.fill(numpy.nan)
            epsilons = numpy.empty(num_events, F32)
            epsilons.fill(numpy.nan)
        else:
            mean, sig, tau, phi = mean_stds  # shape N -> (N, 1)
            phi = phi[:, None]
            tau = tau[:, None]
            mean = mean[:, None]
            # the [:, None] is used to implement multiplication by row;
            # for instance if  a = [1 2], b = [[1 2] [3 4]] then
            # a[:, None] * b = [[1 2] [6 8]] which is the expected result;
            # otherwise one would get multiplication by column [[1 4] [3 8]]
            intra_res = phi * rnd[:-1]  # shape (N, E)
            epsilons = rnd[-1]  # shape E

            if self.correlation_model is not None:
                intra_res = self.correlation_model.apply_correlation(
                    self.sites, imt, intra_res, phi)
                if len(intra_res.shape) == 1:  # a vector
                    intra_res = intra_res[:, None]

            inter_res = tau * epsilons  # shape (N, 1) * E => (N, E)
            gmf = exp(mean + intra_res + inter_res, imt)
            stdi = tau.max(axis=0)  # from shape (N, E) => E
        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>Small renaming [ci skip]<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 exp(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  # sig and eps have shapes (M, E) instead
            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
        """
        M = len(self.imts)
        result = numpy.zeros(
            (len(self.imts), len(self.ctx.sids), num_events), F32)
        sig = numpy.zeros((M, num_events), F32)  # same for all events
        eps = numpy.zeros((M, num_events), F32)  # not the same
        numpy.random.seed(self.seed)
        num_sids = len(self.ctx.sids)
        if self.distribution:
            # build M arrays of random numbers of shape (N+1, E)
            rnd = [rvs(self.distribution, num_sids + 1, num_events)
                   for _ in range(M)]
        else:
            rnd = [None] * M
        for imti, imt in enumerate(self.imts):
            try:
                result[imti], sig[imti], eps[imti] = self._compute(
                     mean_stds[:, imti], num_events, imt, gsim, rnd[imti])
            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, rnd):
        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
            gmf = exp(mean, imt)[:, None]
            gmf = gmf.repeat(num_events, axis=1)
            sigma = 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
            sig = sig[:, None]
            mean = mean[:, None]

            total_res = sig * rnd[:-1]
            gmf = exp(mean + total_res, imt)
            sigma = numpy.empty(num_events, F32)
            sigma.fill(numpy.nan)
            epsilons = numpy.empty(num_events, F32)
            epsilons.fill(numpy.nan)
        else:
            mean, sig, tau, phi = mean_stds  # shape N -> (N, 1)
            phi = phi[:, None]
            tau = tau[:, None]
            mean = mean[:, None]
            # the [:, None] is used to implement multiplication by row;
            # for instance if  a = [1 2], b = [[1 2] [3 4]] then
            # a[:, None] * b = [[1 2] [6 8]] which is the expected result;
            # otherwise one would get multiplication by column [[1 4] [3 8]]
            intra_res = phi * rnd[:-1]  # shape (N, E)
            epsilons = rnd[-1]  # shape E

            if self.correlation_model is not None:
                intra_res = self.correlation_model.apply_correlation(
                    self.sites, imt, intra_res, phi)
                if len(intra_res.shape) == 1:  # a vector
                    intra_res = intra_res[:, None]

            inter_res = tau * epsilons  # shape (N, 1) * E => (N, E)
            gmf = exp(mean + intra_res + inter_res, imt)
            sigma = tau.max()  # from shape (N, 1) => scalar
        return gmf, sigma, 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|>