<commit_before># (C) Datadog, Inc. 2018-present
# All rights reserved
# Licensed under a 3-clause BSD style license (see LICENSE)
from __future__ import absolute_import

import json
import os
import re
from base64 import urlsafe_b64encode
from typing import Dict, List, Optional, Tuple

import pytest

from .._env import (
    E2E_FIXTURE_NAME,
    E2E_PARENT_PYTHON,
    SKIP_ENVIRONMENT,
    TESTING_PLUGIN,
    e2e_active,
    e2e_testing,
    format_config,
    get_env_vars,
    get_state,
    replay_check_run,
    save_state,
    serialize_data,
)

__aggregator = None
__datadog_agent = None
MockResponse = None

@pytest.fixture
def aggregator():
    """This fixture returns a mocked Agent aggregator with state cleared."""
    global __aggregator

    # Since this plugin is loaded before pytest-cov, we need to import lazily so coverage
    # can see imports, class/function definitions, etc. of the base package.
    if __aggregator is None:
        try:
            from datadog_checks.base.stubs import aggregator as __aggregator
        except ImportError:
            raise ImportError('datadog-checks-base is not installed!')

    __aggregator.reset()
    return __aggregator

@pytest.fixture
def datadog_agent():
    global __datadog_agent

    if __datadog_agent is None:
        try:
            from datadog_checks.base.stubs import datadog_agent as __datadog_agent
        except ImportError:
            raise ImportError('datadog-checks-base is not installed!')

    __datadog_agent.reset()
    return __datadog_agent

@pytest.fixture(scope='session', autouse=True)
def dd_environment_runner(request):
    # Skip the runner if the skip environment variable is specified
    do_skip = os.getenv(SKIP_ENVIRONMENT) == 'true'

    testing_plugin = os.getenv(TESTING_PLUGIN) == 'true'

    # Do nothing if no e2e action is triggered and continue with tests
    if not testing_plugin and not e2e_active() and not do_skip:  # no cov
        return
    # If e2e tests are being run it means the environment has
    # already been spun up so we prevent another invocation
    elif e2e_testing() or do_skip:  # no cov
        # Since the scope is `session` there should only ever be one definition
        fixture_def = request._fixturemanager._arg2fixturedefs[E2E_FIXTURE_NAME][0]

        # Make the underlying function a no-op
        fixture_def.func = lambda *args, **kwargs: None
        return

    try:
        config = request.getfixturevalue(E2E_FIXTURE_NAME)
    except Exception as e:
        # pytest doesn't export this exception class so we have to do some introspection
        if e.__class__.__name__ == 'FixtureLookupError':
            # Make it explicit for our command
            pytest.exit('NO E2E FIXTURE AVAILABLE')

        raise

    metadata = {}

    # Environment fixture also returned some metadata
    if isinstance(config, tuple):
        config, possible_metadata = config

        # Support only defining the env_type for ease-of-use
        if isinstance(possible_metadata, str):
            metadata['env_type'] = possible_metadata
        else:
            metadata.update(possible_metadata)

    # Default to Docker as that is the most common
    metadata.setdefault('env_type', 'docker')

    # Save any environment variables
    metadata.setdefault('env_vars', {})
    metadata['env_vars'].update(get_env_vars(raw=True))

    # Inject any log configuration
    logs_config = get_state('logs_config', [])
    if logs_config:
        config = format_config(config)
        config['logs'] = logs_config

    # Mount any volumes for Docker
    if metadata['env_type'] == 'docker':
        docker_volumes = get_state('docker_volumes', [])
        if docker_volumes:
            metadata.setdefault('docker_volumes', []).extend(docker_volumes)

    data = {'config': config, 'metadata': metadata}

    message = serialize_data(data)

    message = 'DDEV_E2E_START_MESSAGE {} DDEV_E2E_END_MESSAGE'.format(message)

    if testing_plugin:
        return message
    else:  # no cov
        # Exit testing and pass data back up to command
        pytest.exit(message)

@pytest.fixture
def dd_agent_check(request, aggregator, datadog_agent):
    if not e2e_testing():
        pytest.skip('Not running E2E tests')

    # Lazily import to reduce plugin load times for everyone
    from datadog_checks.dev import TempDir, run_command

    def run_check(config=None, **kwargs):
        root = os.path.dirname(request.module.__file__)
        while True:
            if os.path.isfile(os.path.join(root, 'pyproject.toml')) or os.path.isfile(os.path.join(root, 'setup.py')):
                check = os.path.basename(root)
                break

            new_root = os.path.dirname(root)
            if new_root == root:
                raise OSError('No Datadog Agent check found')

            root = new_root

        python_path = os.environ[E2E_PARENT_PYTHON]
        env = os.environ.get('TOX_ENV_NAME') or os.environ['HATCH_ENV_ACTIVE']

        # TODO: switch to `ddev` when the old CLI is gone
        check_command = [python_path, '-m', 'datadog_checks.dev', 'env', 'check', check, env, '--json']

        if config:
            config = format_config(config)
            config_file = os.path.join(temp_dir, '{}-{}-{}.json'.format(check, env, urlsafe_b64encode(os.urandom(6))))

            with open(config_file, 'wb') as f:
                output = json.dumps(config).encode('utf-8')
                f.write(output)
            check_command.extend(['--config', config_file])

        for key, value in kwargs.items():
            if value is not False:
                check_command.append('--{}'.format(key.replace('_', '-')))

                if value is not True:
                    check_command.append(str(value))

        result = run_command(check_command, capture=True)

        matches = re.findall(r'((?:\{ \[|\[).*?\n(?:\} \]|\]))', result.stdout, re.DOTALL)

        if not matches:
            raise ValueError(
                '{}{}\nCould not find valid check output'.format(
                    result.stdout,
                    result.stderr,
                )
            )

        for raw_json in matches:
            try:
                collector = json.loads(raw_json)
            except Exception as e:
                raise Exception("Error loading json: {}\nCollector Json Output:\n{}".format(e, raw_json))
            replay_check_run(collector, aggregator, datadog_agent)

        return aggregator

    # Give an explicit name so we don't shadow other uses
    with TempDir('dd_agent_check') as temp_dir:
        yield run_check

@pytest.fixture
def dd_run_check():
    checks = {}

    def run_check(check, extract_message=False, cancel=True):
        if cancel:
            checks[id(check)] = check

        error = check.run()

        if error:
            error = json.loads(error)[0]
            if extract_message:
                raise Exception(error['message'])
            else:
                exc_lines = ['']
                exc_lines.extend(error['traceback'].splitlines())
                raise Exception('\n'.join(exc_lines))

        return ''

    yield run_check

    for c in checks.values():
        try:
            c.cancel()
        except Exception:
            pass

@pytest.fixture(scope='session')
def dd_get_state():
    return get_state

@pytest.fixture(scope='session')
def dd_save_state():
    return save_state

@pytest.fixture(scope='session')
def dd_default_hostname():
    import socket

    return socket.gethostname().lower()

@pytest.fixture
def mock_http_response(mocker):
    # Lazily import `requests` as it may be costly under certain conditions
    global MockResponse
    if MockResponse is None:
        from ..http import MockResponse

    yield lambda *args, **kwargs: mocker.patch(
        kwargs.pop('method', 'requests.get'), return_value=MockResponse(*args, **kwargs)
    )

@pytest.fixture
def mock_performance_objects(mocker, dd_default_hostname):
    def mock_perf_objects(
        perf_objects,  # type: Dict[str, Tuple[List[Optional[str]], Dict[str, List[float]]]]
        server=dd_default_hostname,  # type: str
    ):
        import win32pdh

        mocker.patch('win32pdh.OpenQuery')
        mocker.patch('win32pdh.CollectQueryData')
        mocker.patch('win32pdh.AddCounter', side_effect=lambda _, path: path)
        mocker.patch('win32pdh.AddEnglishCounter', side_effect=lambda _, path: path)
        mocker.patch('win32pdh.RemoveCounter')
        mocker.patch('win32pdh.EnumObjects', return_value=list(perf_objects))

        def enum_object_items(data_source, machine_name, object_name, detail_level):
            instances, counter_values = perf_objects[object_name]
            return list(counter_values), instances if instances != [None] else []

        mocker.patch('win32pdh.EnumObjectItems', side_effect=enum_object_items)

        counters = {}
        for object_name, data in perf_objects.items():
            instances, counter_values = data
            instance_counts = {instance_name: 0 for instance_name in instances}
            instance_indices = []
            for instance_name in instances:
                instance_indices.append(instance_counts[instance_name])
                instance_counts[instance_name] += 1

            for counter_name, values in counter_values.items():
                for instance_name, index, value in zip(instances, instance_indices, values):
                    counters[
                        win32pdh.MakeCounterPath((server, object_name, instance_name, None, index, counter_name))
                    ] = value

        mocker.patch('win32pdh.ValidatePath', side_effect=lambda path: 0 if path in counters else 1)
        mocker.patch('win32pdh.GetFormattedCounterValue', side_effect=lambda path, _: (None, counters[path]))

    return mock_perf_objects

def pytest_configure(config):
    # pytest will emit warnings if these aren't registered ahead of time
    config.addinivalue_line('markers', 'unit: marker for unit tests')
    config.addinivalue_line('markers', 'integration: marker for integration tests')
    config.addinivalue_line('markers', 'e2e: marker for end-to-end Datadog Agent tests')
    config.addinivalue_line("markers", "latest_metrics: marker for verifying support of new metrics")

def pytest_addoption(parser):
    parser.addoption("--run-latest-metrics", action="store_true", default=False, help="run check_metrics tests")

def pytest_collection_modifyitems(config, items):
    # at test collection time, this function gets called by pytest, see:
    # https://docs.pytest.org/en/latest/example/simple.html#control-skipping-of-tests-according-to-command-line-option
    # if the particular option is not present, it will skip all tests marked `latest_metrics`
    if config.getoption("--run-latest-metrics"):
        # --run-check-metrics given in cli: do not skip slow tests
        return
    skip_latest_metrics = pytest.mark.skip(reason="need --run-latest-metrics option to run")
    for item in items:
        if "latest_metrics" in item.keywords:
            item.add_marker(skip_latest_metrics)
<commit_msg>Add a dummy --memray option to the pytest plugin (#13352)<commit_after># (C) Datadog, Inc. 2018-present
# All rights reserved
# Licensed under a 3-clause BSD style license (see LICENSE)
from __future__ import absolute_import

import json
import os
import re
import warnings
from base64 import urlsafe_b64encode
from typing import Dict, List, Optional, Tuple

import pytest
from six import PY2

from .._env import (
    E2E_FIXTURE_NAME,
    E2E_PARENT_PYTHON,
    SKIP_ENVIRONMENT,
    TESTING_PLUGIN,
    e2e_active,
    e2e_testing,
    format_config,
    get_env_vars,
    get_state,
    replay_check_run,
    save_state,
    serialize_data,
)

__aggregator = None
__datadog_agent = None
MockResponse = None

@pytest.fixture
def aggregator():
    """This fixture returns a mocked Agent aggregator with state cleared."""
    global __aggregator

    # Since this plugin is loaded before pytest-cov, we need to import lazily so coverage
    # can see imports, class/function definitions, etc. of the base package.
    if __aggregator is None:
        try:
            from datadog_checks.base.stubs import aggregator as __aggregator
        except ImportError:
            raise ImportError('datadog-checks-base is not installed!')

    __aggregator.reset()
    return __aggregator

@pytest.fixture
def datadog_agent():
    global __datadog_agent

    if __datadog_agent is None:
        try:
            from datadog_checks.base.stubs import datadog_agent as __datadog_agent
        except ImportError:
            raise ImportError('datadog-checks-base is not installed!')

    __datadog_agent.reset()
    return __datadog_agent

@pytest.fixture(scope='session', autouse=True)
def dd_environment_runner(request):
    # Skip the runner if the skip environment variable is specified
    do_skip = os.getenv(SKIP_ENVIRONMENT) == 'true'

    testing_plugin = os.getenv(TESTING_PLUGIN) == 'true'

    # Do nothing if no e2e action is triggered and continue with tests
    if not testing_plugin and not e2e_active() and not do_skip:  # no cov
        return
    # If e2e tests are being run it means the environment has
    # already been spun up so we prevent another invocation
    elif e2e_testing() or do_skip:  # no cov
        # Since the scope is `session` there should only ever be one definition
        fixture_def = request._fixturemanager._arg2fixturedefs[E2E_FIXTURE_NAME][0]

        # Make the underlying function a no-op
        fixture_def.func = lambda *args, **kwargs: None
        return

    try:
        config = request.getfixturevalue(E2E_FIXTURE_NAME)
    except Exception as e:
        # pytest doesn't export this exception class so we have to do some introspection
        if e.__class__.__name__ == 'FixtureLookupError':
            # Make it explicit for our command
            pytest.exit('NO E2E FIXTURE AVAILABLE')

        raise

    metadata = {}

    # Environment fixture also returned some metadata
    if isinstance(config, tuple):
        config, possible_metadata = config

        # Support only defining the env_type for ease-of-use
        if isinstance(possible_metadata, str):
            metadata['env_type'] = possible_metadata
        else:
            metadata.update(possible_metadata)

    # Default to Docker as that is the most common
    metadata.setdefault('env_type', 'docker')

    # Save any environment variables
    metadata.setdefault('env_vars', {})
    metadata['env_vars'].update(get_env_vars(raw=True))

    # Inject any log configuration
    logs_config = get_state('logs_config', [])
    if logs_config:
        config = format_config(config)
        config['logs'] = logs_config

    # Mount any volumes for Docker
    if metadata['env_type'] == 'docker':
        docker_volumes = get_state('docker_volumes', [])
        if docker_volumes:
            metadata.setdefault('docker_volumes', []).extend(docker_volumes)

    data = {'config': config, 'metadata': metadata}

    message = serialize_data(data)

    message = 'DDEV_E2E_START_MESSAGE {} DDEV_E2E_END_MESSAGE'.format(message)

    if testing_plugin:
        return message
    else:  # no cov
        # Exit testing and pass data back up to command
        pytest.exit(message)

@pytest.fixture
def dd_agent_check(request, aggregator, datadog_agent):
    if not e2e_testing():
        pytest.skip('Not running E2E tests')

    # Lazily import to reduce plugin load times for everyone
    from datadog_checks.dev import TempDir, run_command

    def run_check(config=None, **kwargs):
        root = os.path.dirname(request.module.__file__)
        while True:
            if os.path.isfile(os.path.join(root, 'pyproject.toml')) or os.path.isfile(os.path.join(root, 'setup.py')):
                check = os.path.basename(root)
                break

            new_root = os.path.dirname(root)
            if new_root == root:
                raise OSError('No Datadog Agent check found')

            root = new_root

        python_path = os.environ[E2E_PARENT_PYTHON]
        env = os.environ.get('TOX_ENV_NAME') or os.environ['HATCH_ENV_ACTIVE']

        # TODO: switch to `ddev` when the old CLI is gone
        check_command = [python_path, '-m', 'datadog_checks.dev', 'env', 'check', check, env, '--json']

        if config:
            config = format_config(config)
            config_file = os.path.join(temp_dir, '{}-{}-{}.json'.format(check, env, urlsafe_b64encode(os.urandom(6))))

            with open(config_file, 'wb') as f:
                output = json.dumps(config).encode('utf-8')
                f.write(output)
            check_command.extend(['--config', config_file])

        for key, value in kwargs.items():
            if value is not False:
                check_command.append('--{}'.format(key.replace('_', '-')))

                if value is not True:
                    check_command.append(str(value))

        result = run_command(check_command, capture=True)

        matches = re.findall(r'((?:\{ \[|\[).*?\n(?:\} \]|\]))', result.stdout, re.DOTALL)

        if not matches:
            raise ValueError(
                '{}{}\nCould not find valid check output'.format(
                    result.stdout,
                    result.stderr,
                )
            )

        for raw_json in matches:
            try:
                collector = json.loads(raw_json)
            except Exception as e:
                raise Exception("Error loading json: {}\nCollector Json Output:\n{}".format(e, raw_json))
            replay_check_run(collector, aggregator, datadog_agent)

        return aggregator

    # Give an explicit name so we don't shadow other uses
    with TempDir('dd_agent_check') as temp_dir:
        yield run_check

@pytest.fixture
def dd_run_check():
    checks = {}

    def run_check(check, extract_message=False, cancel=True):
        if cancel:
            checks[id(check)] = check

        error = check.run()

        if error:
            error = json.loads(error)[0]
            if extract_message:
                raise Exception(error['message'])
            else:
                exc_lines = ['']
                exc_lines.extend(error['traceback'].splitlines())
                raise Exception('\n'.join(exc_lines))

        return ''

    yield run_check

    for c in checks.values():
        try:
            c.cancel()
        except Exception:
            pass

@pytest.fixture(scope='session')
def dd_get_state():
    return get_state

@pytest.fixture(scope='session')
def dd_save_state():
    return save_state

@pytest.fixture(scope='session')
def dd_default_hostname():
    import socket

    return socket.gethostname().lower()

@pytest.fixture
def mock_http_response(mocker):
    # Lazily import `requests` as it may be costly under certain conditions
    global MockResponse
    if MockResponse is None:
        from ..http import MockResponse

    yield lambda *args, **kwargs: mocker.patch(
        kwargs.pop('method', 'requests.get'), return_value=MockResponse(*args, **kwargs)
    )

@pytest.fixture
def mock_performance_objects(mocker, dd_default_hostname):
    def mock_perf_objects(
        perf_objects,  # type: Dict[str, Tuple[List[Optional[str]], Dict[str, List[float]]]]
        server=dd_default_hostname,  # type: str
    ):
        import win32pdh

        mocker.patch('win32pdh.OpenQuery')
        mocker.patch('win32pdh.CollectQueryData')
        mocker.patch('win32pdh.AddCounter', side_effect=lambda _, path: path)
        mocker.patch('win32pdh.AddEnglishCounter', side_effect=lambda _, path: path)
        mocker.patch('win32pdh.RemoveCounter')
        mocker.patch('win32pdh.EnumObjects', return_value=list(perf_objects))

        def enum_object_items(data_source, machine_name, object_name, detail_level):
            instances, counter_values = perf_objects[object_name]
            return list(counter_values), instances if instances != [None] else []

        mocker.patch('win32pdh.EnumObjectItems', side_effect=enum_object_items)

        counters = {}
        for object_name, data in perf_objects.items():
            instances, counter_values = data
            instance_counts = {instance_name: 0 for instance_name in instances}
            instance_indices = []
            for instance_name in instances:
                instance_indices.append(instance_counts[instance_name])
                instance_counts[instance_name] += 1

            for counter_name, values in counter_values.items():
                for instance_name, index, value in zip(instances, instance_indices, values):
                    counters[
                        win32pdh.MakeCounterPath((server, object_name, instance_name, None, index, counter_name))
                    ] = value

        mocker.patch('win32pdh.ValidatePath', side_effect=lambda path: 0 if path in counters else 1)
        mocker.patch('win32pdh.GetFormattedCounterValue', side_effect=lambda path, _: (None, counters[path]))

    return mock_perf_objects

def pytest_configure(config):
    # pytest will emit warnings if these aren't registered ahead of time
    config.addinivalue_line('markers', 'unit: marker for unit tests')
    config.addinivalue_line('markers', 'integration: marker for integration tests')
    config.addinivalue_line('markers', 'e2e: marker for end-to-end Datadog Agent tests')
    config.addinivalue_line("markers", "latest_metrics: marker for verifying support of new metrics")

def pytest_addoption(parser):
    parser.addoption("--run-latest-metrics", action="store_true", default=False, help="run check_metrics tests")

    if PY2:
        # Add a dummy memray option to make it possible to run memray with `ddev test --memray <integration>` only on py3 environments
        parser.addoption("--memray", action="store_true", default=False, help="Dummy parameter for memray")

def pytest_collection_modifyitems(config, items):
    # at test collection time, this function gets called by pytest, see:
    # https://docs.pytest.org/en/latest/example/simple.html#control-skipping-of-tests-according-to-command-line-option
    # if the particular option is not present, it will skip all tests marked `latest_metrics`
    if config.getoption("--run-latest-metrics"):
        # --run-check-metrics given in cli: do not skip slow tests
        return

    if PY2 and config.getoption("--memray"):
        warnings.warn("--memray option ignored as it's not supported for py2 environments.")

    skip_latest_metrics = pytest.mark.skip(reason="need --run-latest-metrics option to run")
    for item in items:
        if "latest_metrics" in item.keywords:
            item.add_marker(skip_latest_metrics)
<|endoftext|>


<commit_before># coding: utf8

# Copyright 2013-2015 Vincent Jacques <vincent@vincent-jacques.net>

import collections
import textwrap
import unittest

from ActionTree.drawings import make_graph

MockAction = collections.namedtuple("MockAction", "label, dependencies")

class GraphTestCase(unittest.TestCase):
    def test_single_action(self):
        a = MockAction("a", [])

        self.assertEqual(
            make_graph(a).source,
            textwrap.dedent("""\
                digraph action {
                \tnode [shape=box]
                \t\t0 [label=a]
                }"""
            )
        )

    def test_dependency(self):
        b = MockAction("b", [])
        a = MockAction("a", [b])

        self.assertEqual(
            make_graph(a).source,
            textwrap.dedent("""\
                digraph action {
                \tnode [shape=box]
                \t\t0 [label=a]
                \t\t1 [label=b]
                \t\t\t0 -> 1
                }"""
            )
        )

    def test_diamond(self):
        a = MockAction("a", [])
        b = MockAction("b", [a])
        c = MockAction("c", [a])
        d = MockAction("d", [b, c])

        self.assertEqual(
            make_graph(d).source,
            textwrap.dedent("""\
                digraph action {
                \tnode [shape=box]
                \t\t0 [label=d]
                \t\t1 [label=b]
                \t\t2 [label=a]
                \t\t\t1 -> 2
                \t\t\t0 -> 1
                \t\t3 [label=c]
                \t\t\t3 -> 2
                \t\t\t0 -> 3
                }"""
            )
        )
<commit_msg>Add test<commit_after># coding: utf8

# Copyright 2013-2015 Vincent Jacques <vincent@vincent-jacques.net>

import collections
import textwrap
import unittest

from ActionTree.drawings import make_graph

MockAction = collections.namedtuple("MockAction", "label, dependencies")

class GraphTestCase(unittest.TestCase):
    def test_single_action(self):
        a = MockAction("a", [])

        self.assertEqual(
            make_graph(a).source,
            textwrap.dedent("""\
                digraph action {
                \tnode [shape=box]
                \t\t0 [label=a]
                }"""
            )
        )

    def test_dependency(self):
        b = MockAction("b", [])
        a = MockAction("a", [b])

        self.assertEqual(
            make_graph(a).source,
            textwrap.dedent("""\
                digraph action {
                \tnode [shape=box]
                \t\t0 [label=a]
                \t\t1 [label=b]
                \t\t\t0 -> 1
                }"""
            )
        )

    def test_diamond(self):
        a = MockAction("a", [])
        b = MockAction("b", [a])
        c = MockAction("c", [a])
        d = MockAction("d", [b, c])

        self.assertEqual(
            make_graph(d).source,
            textwrap.dedent("""\
                digraph action {
                \tnode [shape=box]
                \t\t0 [label=d]
                \t\t1 [label=b]
                \t\t2 [label=a]
                \t\t\t1 -> 2
                \t\t\t0 -> 1
                \t\t3 [label=c]
                \t\t\t3 -> 2
                \t\t\t0 -> 3
                }"""
            )
        )

    def test_weird_label(self):
        a = MockAction("spaces and; semi=columns", [])

        self.assertEqual(
            make_graph(a).source,
            textwrap.dedent("""\
                digraph action {
                \tnode [shape=box]
                \t\t0 [label="spaces and; semi=columns"]
                }"""
            )
        )
<|endoftext|>


<commit_before>from ethereum.utils import sha3, privtoaddr, int_to_addr, to_string, big_endian_to_int, checksum_encode, int_to_big_endian, encode_hex
from ethereum.genesis_helpers import mk_basic_state
from ethereum.transactions import Transaction
from ethereum.consensus_strategy import get_consensus_strategy
from ethereum.config import config_homestead, config_tangerine, config_spurious, config_metropolis, default_config, Env
from ethereum.pow.ethpow import Miner
from ethereum.messages import apply_transaction, apply_message
from ethereum.common import verify_execution_results, mk_block_from_prevstate, set_execution_results
from ethereum.meta import make_head_candidate
from ethereum.abi import ContractTranslator
import rlp
# Initialize accounts
accounts = []
keys = []

for account_number in range(10):
    keys.append(sha3(to_string(account_number)))
    accounts.append(privtoaddr(keys[-1]))

k0, k1, k2, k3, k4, k5, k6, k7, k8, k9 = keys[:10]
a0, a1, a2, a3, a4, a5, a6, a7, a8, a9 = accounts[:10]

base_alloc = {}
minimal_alloc = {}
for a in accounts:
    base_alloc[a] = {'balance': 10**24}
for i in range(1, 9):
    base_alloc[int_to_addr(i)] = {'balance': 1}
    minimal_alloc[int_to_addr(i)] = {'balance': 1}
minimal_alloc[accounts[0]] = {'balance': 10**18}

# Initialize languages
languages = {}

try:
    import serpent
    languages['serpent'] = serpent
except ImportError:
    pass

from ethereum.tools._solidity import get_solidity
_solidity = get_solidity()
if _solidity:
    languages['solidity'] = _solidity

try:
    from viper import compiler
    languages['viper'] = compiler
except (ImportError, TypeError):
    pass

class TransactionFailed(Exception):
    pass

from ethereum.abi import ContractTranslator
import types

STARTGAS = 3141592
GASPRICE = 0

from ethereum.slogging import configure_logging
config_string = ':info'
# configure_logging(config_string=config_string)

class ABIContract(object):  # pylint: disable=too-few-public-methods

    def __init__(self, _tester, _abi, address):
        self.address = address

        if isinstance(_abi, ContractTranslator):
            abi_translator = _abi
        else:
            abi_translator = ContractTranslator(_abi)

        self.translator = abi_translator

        for function_name in self.translator.function_data:
            if self.translator.function_data[function_name]['is_constant']:
                function = self.method_factory(_tester.call, function_name)
            else:
                function = self.method_factory(_tester.tx, function_name)
            method = types.MethodType(function, self)
            setattr(self, function_name, method)

    @staticmethod
    def method_factory(tx_or_call, function_name):
        """ Return a proxy for calling a contract method with automatic encoding of
        argument and decoding of results.
        """

        def kall(self, *args, **kwargs):
            key = kwargs.get('sender', k0)

            result = tx_or_call(  # pylint: disable=protected-access
                sender=key,
                to=self.address,
                value=kwargs.get('value', 0),
                data=self.translator.encode(function_name, args),
                startgas=kwargs.get('startgas', STARTGAS),
                gasprice=kwargs.get('gasprice', GASPRICE)
            )

            if result is False:
                return result
            if result == b'':
                return None
            o = self.translator.decode(function_name, result)
            return o[0] if len(o) == 1 else o
        return kall

def get_env(env):
    d = {
        None: config_metropolis,
        'mainnet': default_config,
        'homestead': config_homestead,
        'tangerine': config_tangerine,
        'spurious': config_spurious,
        'metropolis': config_metropolis,
    }
    return env if isinstance(env, Env) else Env(config=d[env])

class State(object):
    def __init__(self, genesis):
        self.state = genesis

    def tx(self, sender=k0, to=b'\x00' * 20, value=0,
           data=b'', startgas=STARTGAS, gasprice=GASPRICE):
        sender_addr = privtoaddr(sender)
        transaction = Transaction(self.state.get_nonce(sender_addr), gasprice, startgas,
                                  to, value, data).sign(sender)
        success, output = apply_transaction(self.state, transaction)
        if not success:
            raise TransactionFailed()
        return output

    def call(self, sender=k0, to=b'\x00' * 20, value=0,
             data=b'', startgas=STARTGAS, gasprice=GASPRICE):
        self.state.commit()
        sender_addr = privtoaddr(sender)
        result = apply_message(
            self.state.ephemeral_clone(),
            sender=sender_addr,
            to=to,
            code_address=to,
            value=value,
            data=data,
            gas=startgas)
        if result is None:
            raise TransactionFailed()
        return result

class Chain(object):
    def __init__(self, alloc=base_alloc, env=None, genesis=None):
        from ethereum.pow import chain as pow_chain
        if genesis:
            if genesis.env.config['CONSENSUS_STRATEGY'] == 'hybrid_casper':
                from ethereum.hybrid_casper import chain as hybrid_casper_chain
                self.chain = hybrid_casper_chain.Chain(genesis, reset_genesis=True)
            else:
                self.chain = pow_chain.Chain(genesis, reset_genesis=True)
        else:
            self.chain = pow_chain.Chain(mk_basic_state(alloc, None, get_env(env)), reset_genesis=True)
        self.cs = get_consensus_strategy(self.chain.env.config)
        self.block = mk_block_from_prevstate(
            self.chain, timestamp=self.chain.state.timestamp + 1)
        self.head_state = self.chain.state.ephemeral_clone()
        self.cs.initialize(self.head_state, self.block)
        self.last_sender = None
        self.last_tx = None

    def direct_tx(self, transaction):
        self.last_tx = transaction
        if self.last_sender is not None and privtoaddr(
                self.last_sender) != transaction.sender:
            self.last_sender = None
        success, output = apply_transaction(self.head_state, transaction)
        self.block.transactions.append(transaction)
        if not success:
            raise TransactionFailed()
        return output

    def tx(self, sender=k0, to=b'\x00' * 20, value=0,
           data=b'', startgas=STARTGAS, gasprice=GASPRICE):
        sender_addr = privtoaddr(sender)
        self.last_sender = sender
        transaction = Transaction(self.head_state.get_nonce(sender_addr), gasprice, startgas,
                                  to, value, data).sign(sender)
        output = self.direct_tx(transaction)
        return output

    def call(self, sender=k0, to=b'\x00' * 20, value=0,
             data=b'', startgas=STARTGAS, gasprice=GASPRICE):
        self.head_state.commit()
        sender_addr = privtoaddr(sender)
        result = apply_message(
            self.head_state.ephemeral_clone(),
            sender=sender_addr,
            to=to,
            code_address=to,
            value=value,
            data=data,
            gas=startgas)
        if result is None:
            raise TransactionFailed()
        return result

    def last_gas_used(self, with_tx=False):
        if len(self.head_state.receipts) == 1:
            diff = self.head_state.receipts[-1].gas_used
        else:
            diff = self.head_state.receipts[-1].gas_used - \
                self.head_state.receipts[-2].gas_used
        return diff - (not with_tx) * self.last_tx.intrinsic_gas_used

    def contract(self, sourcecode, args=[], sender=k0, value=0,
                 language=None, l=None, startgas=STARTGAS, gasprice=GASPRICE):
        assert not (l and language)
        language = l or language
        if language == 'evm':
            assert len(args) == 0
            return self.tx(sender=sender, to=b'', value=value,
                           data=sourcecode, startgas=startgas, gasprice=gasprice)
        else:
            compiler = languages[language]
            interface = compiler.mk_full_signature(sourcecode)
            ct = ContractTranslator(interface)
            code = compiler.compile(
                sourcecode) + (ct.encode_constructor_arguments(args) if args else b'')
            addr = self.tx(
                sender=sender,
                to=b'',
                value=value,
                data=code,
                startgas=startgas,
                gasprice=gasprice)
            return ABIContract(self, ct, addr)

    def mine(self, number_of_blocks=1, coinbase=a0):
        self.cs.finalize(self.head_state, self.block)
        set_execution_results(self.head_state, self.block)
        self.block = Miner(self.block).mine(rounds=100, start_nonce=0)
        assert self.chain.add_block(self.block)
        b = self.block
        for i in range(1, number_of_blocks):
            b, _ = make_head_candidate(
                self.chain, parent=b, timestamp=self.chain.state.timestamp + 14, coinbase=coinbase)
            b = Miner(b).mine(rounds=100, start_nonce=0)
            assert self.chain.add_block(b)
        self.change_head(b.header.hash, coinbase)
        return b

    def change_head(self, parent, coinbase=a0):
        self.head_state = self.chain.mk_poststate_of_blockhash(
            parent).ephemeral_clone()
        self.block = mk_block_from_prevstate(
            self.chain,
            self.head_state,
            timestamp=self.chain.state.timestamp,
            coinbase=coinbase)
        self.cs.initialize(self.head_state, self.block)

    def snapshot(self):
        self.head_state.commit()
        return self.head_state.snapshot(), len(
            self.block.transactions), self.block.number

    def revert(self, snapshot):
        state_snapshot, txcount, blknum = snapshot
        assert blknum == self.block.number, "Cannot revert beyond block boundaries!"
        self.block.transactions = self.block.transactions[:txcount]
        self.head_state.revert(state_snapshot)

def int_to_0x_hex(v):
    o = encode_hex(int_to_big_endian(v))
    if o and o[0] == '0':
        return '0x' + o[1:]
    else:
        return '0x' + o

def mk_state_test_prefill(c):
    env = {
        "currentCoinbase": checksum_encode(c.head_state.block_coinbase),
        "currentDifficulty": int_to_0x_hex(c.head_state.block_difficulty),
        "currentGasLimit": int_to_0x_hex(c.head_state.gas_limit),
        "currentNumber": int_to_0x_hex(c.head_state.block_number),
        "currentTimestamp": int_to_0x_hex(c.head_state.timestamp),
        "previousHash": "0x" + encode_hex(c.head_state.prev_headers[0].hash),
    }
    pre = c.head_state.to_dict()
    return {"env": env, "pre": pre}

def mk_state_test_postfill(c, prefill, filler_mode=False):
    txdata = c.last_tx.to_dict()
    modified_tx_data = {
        "data": [txdata["data"]],
        "gasLimit": [int_to_0x_hex(txdata["startgas"])],
        "gasPrice": int_to_0x_hex(txdata["gasprice"]),
        "nonce": int_to_0x_hex(txdata["nonce"]),
        "secretKey": '0x' + encode_hex(c.last_sender),
        "to": txdata["to"],
        "value": [int_to_0x_hex(txdata["value"])],
    }
    c.head_state.commit()
    postStateHash = '0x' + encode_hex(c.head_state.trie.root_hash)
    if c.chain.config == config_homestead:
        config = 'Homestead'
    elif c.chain.config == config_tangerine:
        config = 'EIP150'
    elif c.chain.config == config_spurious:
        config = 'EIP158'
    elif c.chain.config == config_metropolis:
        config = 'Metropolis'
    else:
        raise Exception("Cannot get config")
    o = {
        "env": prefill["env"],
        "pre": prefill["pre"],
        "transaction": modified_tx_data,
    }
    if not filler_mode:
        o["post"] = {config: [{"hash": postStateHash,
                               "indexes": {"data": 0, "gas": 0, "value": 0}}]}
    else:
        o["expect"] = [{"indexes": {"data": 0, "gas": 0, "value": 0}, "network": [
            "Metropolis"], "result": c.head_state.to_dict()}]
    return o
<commit_msg>enabled snapshot usage again<commit_after>from ethereum.utils import sha3, privtoaddr, int_to_addr, to_string, big_endian_to_int, checksum_encode, int_to_big_endian, encode_hex
from ethereum.genesis_helpers import mk_basic_state
from ethereum.transactions import Transaction
from ethereum.consensus_strategy import get_consensus_strategy
from ethereum.config import config_homestead, config_tangerine, config_spurious, config_metropolis, default_config, Env
from ethereum.pow.ethpow import Miner
from ethereum.messages import apply_transaction, apply_message
from ethereum.common import verify_execution_results, mk_block_from_prevstate, set_execution_results
from ethereum.meta import make_head_candidate
from ethereum.abi import ContractTranslator
import rlp
# Initialize accounts
accounts = []
keys = []

for account_number in range(10):
    keys.append(sha3(to_string(account_number)))
    accounts.append(privtoaddr(keys[-1]))

k0, k1, k2, k3, k4, k5, k6, k7, k8, k9 = keys[:10]
a0, a1, a2, a3, a4, a5, a6, a7, a8, a9 = accounts[:10]

base_alloc = {}
minimal_alloc = {}
for a in accounts:
    base_alloc[a] = {'balance': 10**24}
for i in range(1, 9):
    base_alloc[int_to_addr(i)] = {'balance': 1}
    minimal_alloc[int_to_addr(i)] = {'balance': 1}
minimal_alloc[accounts[0]] = {'balance': 10**18}

# Initialize languages
languages = {}

try:
    import serpent
    languages['serpent'] = serpent
except ImportError:
    pass

from ethereum.tools._solidity import get_solidity
_solidity = get_solidity()
if _solidity:
    languages['solidity'] = _solidity

try:
    from viper import compiler
    languages['viper'] = compiler
except (ImportError, TypeError):
    pass

class TransactionFailed(Exception):
    pass

from ethereum.abi import ContractTranslator
import types

STARTGAS = 3141592
GASPRICE = 0

from ethereum.slogging import configure_logging
config_string = ':info'
# configure_logging(config_string=config_string)

class ABIContract(object):  # pylint: disable=too-few-public-methods

    def __init__(self, _tester, _abi, address):
        self.address = address

        if isinstance(_abi, ContractTranslator):
            abi_translator = _abi
        else:
            abi_translator = ContractTranslator(_abi)

        self.translator = abi_translator

        for function_name in self.translator.function_data:
            if self.translator.function_data[function_name]['is_constant']:
                function = self.method_factory(_tester.call, function_name)
            else:
                function = self.method_factory(_tester.tx, function_name)
            method = types.MethodType(function, self)
            setattr(self, function_name, method)

    @staticmethod
    def method_factory(tx_or_call, function_name):
        """ Return a proxy for calling a contract method with automatic encoding of
        argument and decoding of results.
        """

        def kall(self, *args, **kwargs):
            key = kwargs.get('sender', k0)

            result = tx_or_call(  # pylint: disable=protected-access
                sender=key,
                to=self.address,
                value=kwargs.get('value', 0),
                data=self.translator.encode(function_name, args),
                startgas=kwargs.get('startgas', STARTGAS),
                gasprice=kwargs.get('gasprice', GASPRICE)
            )

            if result is False:
                return result
            if result == b'':
                return None
            o = self.translator.decode(function_name, result)
            return o[0] if len(o) == 1 else o
        return kall

def get_env(env):
    d = {
        None: config_metropolis,
        'mainnet': default_config,
        'homestead': config_homestead,
        'tangerine': config_tangerine,
        'spurious': config_spurious,
        'metropolis': config_metropolis,
    }
    return env if isinstance(env, Env) else Env(config=d[env])

class State(object):
    def __init__(self, genesis):
        self.state = genesis

    def tx(self, sender=k0, to=b'\x00' * 20, value=0,
           data=b'', startgas=STARTGAS, gasprice=GASPRICE):
        sender_addr = privtoaddr(sender)
        transaction = Transaction(self.state.get_nonce(sender_addr), gasprice, startgas,
                                  to, value, data).sign(sender)
        success, output = apply_transaction(self.state, transaction)
        if not success:
            raise TransactionFailed()
        return output

    def call(self, sender=k0, to=b'\x00' * 20, value=0,
             data=b'', startgas=STARTGAS, gasprice=GASPRICE):
        self.state.commit()
        sender_addr = privtoaddr(sender)
        result = apply_message(
            self.state.ephemeral_clone(),
            sender=sender_addr,
            to=to,
            code_address=to,
            value=value,
            data=data,
            gas=startgas)
        if result is None:
            raise TransactionFailed()
        return result

class Chain(object):
    def __init__(self, alloc=base_alloc, env=None, genesis=None):
        from ethereum.pow import chain as pow_chain
        if genesis:
            if type(genesis)!=dict and genesis.env.config['CONSENSUS_STRATEGY'] == 'hybrid_casper':
                from ethereum.hybrid_casper import chain as hybrid_casper_chain
                self.chain = hybrid_casper_chain.Chain(genesis, reset_genesis=True)
            else:
                self.chain = pow_chain.Chain(genesis, reset_genesis=True)
        else:
            self.chain = pow_chain.Chain(mk_basic_state(alloc, None, get_env(env)), reset_genesis=True)
        self.cs = get_consensus_strategy(self.chain.env.config)
        self.block = mk_block_from_prevstate(
            self.chain, timestamp=self.chain.state.timestamp + 1)
        self.head_state = self.chain.state.ephemeral_clone()
        self.cs.initialize(self.head_state, self.block)
        self.last_sender = None
        self.last_tx = None

    def direct_tx(self, transaction):
        self.last_tx = transaction
        if self.last_sender is not None and privtoaddr(
                self.last_sender) != transaction.sender:
            self.last_sender = None
        success, output = apply_transaction(self.head_state, transaction)
        self.block.transactions.append(transaction)
        if not success:
            raise TransactionFailed()
        return output

    def tx(self, sender=k0, to=b'\x00' * 20, value=0,
           data=b'', startgas=STARTGAS, gasprice=GASPRICE):
        sender_addr = privtoaddr(sender)
        self.last_sender = sender
        transaction = Transaction(self.head_state.get_nonce(sender_addr), gasprice, startgas,
                                  to, value, data).sign(sender)
        output = self.direct_tx(transaction)
        return output

    def call(self, sender=k0, to=b'\x00' * 20, value=0,
             data=b'', startgas=STARTGAS, gasprice=GASPRICE):
        self.head_state.commit()
        sender_addr = privtoaddr(sender)
        result = apply_message(
            self.head_state.ephemeral_clone(),
            sender=sender_addr,
            to=to,
            code_address=to,
            value=value,
            data=data,
            gas=startgas)
        if result is None:
            raise TransactionFailed()
        return result

    def last_gas_used(self, with_tx=False):
        if len(self.head_state.receipts) == 1:
            diff = self.head_state.receipts[-1].gas_used
        else:
            diff = self.head_state.receipts[-1].gas_used - \
                self.head_state.receipts[-2].gas_used
        return diff - (not with_tx) * self.last_tx.intrinsic_gas_used

    def contract(self, sourcecode, args=[], sender=k0, value=0,
                 language=None, l=None, startgas=STARTGAS, gasprice=GASPRICE):
        assert not (l and language)
        language = l or language
        if language == 'evm':
            assert len(args) == 0
            return self.tx(sender=sender, to=b'', value=value,
                           data=sourcecode, startgas=startgas, gasprice=gasprice)
        else:
            compiler = languages[language]
            interface = compiler.mk_full_signature(sourcecode)
            ct = ContractTranslator(interface)
            code = compiler.compile(
                sourcecode) + (ct.encode_constructor_arguments(args) if args else b'')
            addr = self.tx(
                sender=sender,
                to=b'',
                value=value,
                data=code,
                startgas=startgas,
                gasprice=gasprice)
            return ABIContract(self, ct, addr)

    def mine(self, number_of_blocks=1, coinbase=a0):
        self.cs.finalize(self.head_state, self.block)
        set_execution_results(self.head_state, self.block)
        self.block = Miner(self.block).mine(rounds=100, start_nonce=0)
        assert self.chain.add_block(self.block)
        b = self.block
        for i in range(1, number_of_blocks):
            b, _ = make_head_candidate(
                self.chain, parent=b, timestamp=self.chain.state.timestamp + 14, coinbase=coinbase)
            b = Miner(b).mine(rounds=100, start_nonce=0)
            assert self.chain.add_block(b)
        self.change_head(b.header.hash, coinbase)
        return b

    def change_head(self, parent, coinbase=a0):
        self.head_state = self.chain.mk_poststate_of_blockhash(
            parent).ephemeral_clone()
        self.block = mk_block_from_prevstate(
            self.chain,
            self.head_state,
            timestamp=self.chain.state.timestamp,
            coinbase=coinbase)
        self.cs.initialize(self.head_state, self.block)

    def snapshot(self):
        self.head_state.commit()
        return self.head_state.snapshot(), len(
            self.block.transactions), self.block.number

    def revert(self, snapshot):
        state_snapshot, txcount, blknum = snapshot
        assert blknum == self.block.number, "Cannot revert beyond block boundaries!"
        self.block.transactions = self.block.transactions[:txcount]
        self.head_state.revert(state_snapshot)

def int_to_0x_hex(v):
    o = encode_hex(int_to_big_endian(v))
    if o and o[0] == '0':
        return '0x' + o[1:]
    else:
        return '0x' + o

def mk_state_test_prefill(c):
    env = {
        "currentCoinbase": checksum_encode(c.head_state.block_coinbase),
        "currentDifficulty": int_to_0x_hex(c.head_state.block_difficulty),
        "currentGasLimit": int_to_0x_hex(c.head_state.gas_limit),
        "currentNumber": int_to_0x_hex(c.head_state.block_number),
        "currentTimestamp": int_to_0x_hex(c.head_state.timestamp),
        "previousHash": "0x" + encode_hex(c.head_state.prev_headers[0].hash),
    }
    pre = c.head_state.to_dict()
    return {"env": env, "pre": pre}

def mk_state_test_postfill(c, prefill, filler_mode=False):
    txdata = c.last_tx.to_dict()
    modified_tx_data = {
        "data": [txdata["data"]],
        "gasLimit": [int_to_0x_hex(txdata["startgas"])],
        "gasPrice": int_to_0x_hex(txdata["gasprice"]),
        "nonce": int_to_0x_hex(txdata["nonce"]),
        "secretKey": '0x' + encode_hex(c.last_sender),
        "to": txdata["to"],
        "value": [int_to_0x_hex(txdata["value"])],
    }
    c.head_state.commit()
    postStateHash = '0x' + encode_hex(c.head_state.trie.root_hash)
    if c.chain.config == config_homestead:
        config = 'Homestead'
    elif c.chain.config == config_tangerine:
        config = 'EIP150'
    elif c.chain.config == config_spurious:
        config = 'EIP158'
    elif c.chain.config == config_metropolis:
        config = 'Metropolis'
    else:
        raise Exception("Cannot get config")
    o = {
        "env": prefill["env"],
        "pre": prefill["pre"],
        "transaction": modified_tx_data,
    }
    if not filler_mode:
        o["post"] = {config: [{"hash": postStateHash,
                               "indexes": {"data": 0, "gas": 0, "value": 0}}]}
    else:
        o["expect"] = [{"indexes": {"data": 0, "gas": 0, "value": 0}, "network": [
            "Metropolis"], "result": c.head_state.to_dict()}]
    return o
<|endoftext|>


<commit_before>#!/usr/bin/env python

"""
RS 2017/02/17:  Bayesian quadrature compressions of filter transmission curves

These routines take as input an empirical transmission curve and use Bayesian
quadrature (O'Hagan 1991; Huszar & Duvenaud 2012) as a means for approximating
the integral of that transmission curve against a Gaussian process with known
covariance function.  This will be an important step of making tractable the
problem of inferring a spectral time series against photometry.
"""

import sys
import glob
import dill as pickle
import numpy as np
from scipy import linalg, interpolate, integrate, optimize

class BQFilter(object):
    """
    Implements compression of a filter transfer function using Bayesian
    quadrature.  Chooses evaluation points and weights using a greedy
    algorithm that successively minimizes, at each stage, the mean square
    difference between the original filter function and an approximate
    version as evaluated on functions under a Gaussian process prior.
    Uses the algorithm as outlined in section 3 of
        "Optimally-Weighted Herding is Bayesian Quadrature",
            F. Huszar & D. Duvenaud, Proc. UAI 2012, p. 377.
    """
    # In general, underscores represent internal variables associated
    # with the training points.

    def _vmsg(self, msg):
        """
        Writes unbuffered status messages to stdout.
        """
        if self.verbose:
            print msg
        sys.stdout.flush()

    def __init__(self, _x, _fx, kcov, khyp=[ ], verbose=False):
        """
        Parameters:
            _x:    points at which original filter transfer function is
                   empirically defined (e.g. wavelengths); this will form
                   the probability distribution p for Bayesian quadrature
            _fx:   transfer function values (e.g. filter transmission)
            kcov:  covariance kernel for GP integrand (callable),
                   defined to take two arguments x1, x2
            khyp:  (fixed) hyperparameters for kcov (np.array of floats)
            verbose:  print status messages? (bool)
        """
        # Normalize the training points to zero mean and unit variance,
        # to help stabilize the scipy quadrature routines
        self.verbose = verbose
        self._x, self._fx = _x, _fx
        self._xmu, self._xsig = np.mean(_x, axis=0), np.std(_x, axis=0)
        self._u = self._x2u(self._x)
        self._ulo, self._uhi = self._u.min(), self._u.max()
        # Internal callables for filter transfer function and covariance
        integZu = interpolate.interp1d(self._u, _fx)
        self.Zu, Zu_err = integrate.quad(integZu, self._ulo, self._uhi)
        self._vmsg("__init__:  Zu = {:.3g} +/- {:.3g}".format(self.Zu, Zu_err))
        self._vmsg("__init__:  orig. filter norm = {:.3g}".format(
                   self.Zu * self._xsig))
        self.pu = interpolate.interp1d(self._u, _fx/self.Zu)
        # self.kcov = kcov
        self.kcov, self.khyp = kcov, khyp
        # Internal state
        self.u = np.array([ ])      # quadrature points
        self.zu = np.array([ ])     # quadrature weights
        self.K = np.array([[ ]])    # covariance at quadrature points
        # Starting variance
        self._calc_base_variance_integral()

    def _x2u(self, x):
        return (x - self._xmu)/self._xsig

    def _u2x(self, u):
        return u*self._xsig + self._xmu

    def kxx(self, x1, x2):
        return self.kcov(x1, x2, *(self.khyp))

    def kuu(self, u1, u2):
        return self.kcov(self._u2x(u1), self._u2x(u2), *(self.khyp))

    def _calc_base_variance_integral(self):
        """
        Calculates the integral
            int_dx int_dx' k(x,x') * p(x) * p(x')
        that forms the baseline variance estimate for a BQ filter.
        """
        if False:
            V0, V0_err = 0.101403302569, 1.83139792831e-06
            self.Vn = self.V0 = V0
            return
        # Define a lot of throwaway functions to wrap parts of the problem
        # to match requirements for scipy.integrate.quad.  Rescale x-axis
        # to improve convergence of the integral.
        g = lambda const: (lambda x: const)
        integ_u = lambda u1, u2: self.kuu(u1, u2) * self.pu(u1) * self.pu(u2)
        # Run dblquad -- this should take about 1 min to complete
        self._vmsg("_calc_base_variance_integral:  Calculating...")
        V0, V0_err = integrate.dblquad(
                integ_u, self._ulo, self._uhi, g(self._ulo), g(self._uhi))
        self.Vn = self.V0 = V0
        self._vmsg("_calc_base_variance_integral: V0 = {} +/- {}"
                   .format(V0, V0_err))
        self._vmsg("_calc_base_variance_integral: V0[xval] = {}"
                   .format(V0 * (self._xsig * self.Zu)**2))

    def add_one_point(self):
        """
        Runs optimization for adding a single point to the BQ filter.
        """
        Ktmp = np.array(self.K)
        def u_var(u_n):
            # Wrapper functions for integrands
            integ_u = lambda u: self.kuu(u, u_n) * self.pu(u)
            _zn = lambda u: integrate.quad(integ_u, self._ulo, self._uhi)[0]
            # Update internal state
            self.u[-1] = u_n
            self.zu[-1] = _zn(u_n)
            kuu_n = self.kuu(self.u, u_n)
            Ktmp[-1,:] = Ktmp[:,-1] = kuu_n
            # Calculate and return variance = V0 - z.T * inv(K) * z
            self.Kchol = linalg.cholesky(Ktmp, lower=True)
            zeta = linalg.solve_triangular(self.Kchol, self.zu, lower=True)
            self.Vn = self.V0 - np.dot(zeta, zeta)
            return self.Vn
        # Enlarge internal state and optimize over location of new point
        # Since doing this in u, initial guess for new point is 0.0
        n = len(self.u)
        KX = np.atleast_2d([self.u])
        self.u = np.concatenate([self.u, [0.0]])
        self.zu = np.concatenate([self.zu, [0.0]])
        if Ktmp.shape[1] == 0:
            Ktmp = np.array([[1.0]])
        else:
            Ktmp = np.vstack([np.hstack([self.K, KX.T   ]),
                              np.hstack([KX,     [[1.0]]])])
        Ktmp += 1e-8 * np.eye(n+1)
        # As we add more points the Cholesky factor may become more unstable,
        # so add a small nugget -- as small as we can get away with.
        nugget = 1e-07
        self._vmsg("add_one_point:  Optimizing over point #{}...".format(n+1))
        while nugget <= 0.1:
            # try:
                u0 = np.array([0.0])
                bounds = np.array([(self._ulo, self._uhi)])
                cons = [{ 'type': 'ineq', 'fun': lambda u: u - self._ulo },
                        { 'type': 'ineq', 'fun': lambda u: self._uhi - u }]
                result = optimize.minimize(
                        u_var, u0, method='COBYLA', constraints=cons)
                break
            # except Exception as e:
                self._vmsg("add_one_point:  Cholesky factorization failed")
                self._vmsg(str(e))
                self._vmsg("add_one_point:  adding nugget = {} "
                           "to stabilize Cholesky factor".format(nugget))
                Ktmp += nugget * np.eye(n+1)
                nugget *= 10
        if nugget > 0.1:
            self._vmsg("add_one_point:  total Cholesky factorization fail")
            self.u, self.zu = self.u[:-1], self.zu[:-1]
            return
        elif result.success:
            self._vmsg("add_one_point:  Added new point (wt) {} ({}); Vn = {}"
                       .format(self.u[-1], self.zu[-1], self.Vn))
            self.K = Ktmp
        else:
            self._vmsg("add_one_point:  Optimization failed, don't trust me!")
            self._vmsg("Failure message:  {}".format(result.message))
        self.wbq = linalg.cho_solve((self.Kchol, True), self.zu)
        # Return quadrature points and weights for integration transformed
        # back to original x-axis, along with the renormalized variance.
        self.x = self._u2x(self.u)
        self.zx = self.zu * self._xsig * self.Zu
        self.Vxn = self.Vn * (self._xsig * self.Zu)**2

    def add_n_points(self, n=0):
        """
        What it says on the tin:  runs self.add_one_point() n times.
        """
        for i in range(n):
            self.add_one_point()

    def int_quadz(self, f):
        """
        Uses straight-up quadrature to evaluate integral of f.  In most
        interesting cases f will be an interpolate.interp1d over some
        set of points (for example, an observed supernova spectrum).
        Parameters:
            f: 1-D callable
        """
        integ_u = lambda u: f(self._u2x(u)) * self.pu(u)
        pnorm = self._xsig * self.Zu
        Fu, Fu_err = integrate.quad(integ_u, self._ulo, self._uhi)
        Fx, Fx_err = Fu * pnorm, Fu_err * pnorm
        self._vmsg('int_quadz: F = {} +/- {}'.format(Fx, Fx_err))
        return Fx

    def int_bayes(self, f):
        """
        Uses Bayesian quadrature rule to evaluate integral of f.  The rule
        is derived assuming f is a Gaussian process with a given covariance
        kernel (i.e. fixed hyperparameters).
        Parameters:
            f: 1-D callable
        """
        pnorm = self._xsig * self.Zu
        Fx = np.dot(self.wbq, f(self.x)) * pnorm
        self._vmsg('int_bayes: F = {}'.format(Fx))
        return Fx

def sqexp(x1, x2, l):
    """
    Kernel for the GP, in this case an isotropic square exponential
    """
    return np.exp(-0.5*((x1-x2)/l)**2)

def sqlogexp(x1, x2, logl):
    """
    GP kernel, square exponential in log of variable
    """
    return np.exp(-0.5*((np.log(x1)-np.log(x2))/logl)**2)

def compress_filter(fname, kcov, khyp, n_points):
    """
    Reads in a transfer curve for a filter, and computes an optimal
    Bayesian quadrature rule for a square exponential covariance kernel.
    Parameters:
        fname:  name of two-column text file with (x, y) pairs
        kcov:   covariance kernel for GP integrand (callable),
                defined to take two arguments x1, x2
        khyp:  (fixed) hyperparameters for kcov (np.array of floats)
        n_points:  number of quadrature points desired
    """
    _x, _fx = np.loadtxt(fname, unpack=True)
    bquad = BQFilter(_x, _fx, kcov, khyp, verbose=True)
    bquad.add_n_points(n_points)
    return bquad

def integrate_test_suite(bquad):
    """
    Uses Bayesian quadrature to integrate a bunch of spectra, and compares
    with results from integrating straight against filter.
    """
    fquadz, fbayes = [ ], [ ]
    for fn in glob.glob("testdata/spec*.txt"):
        _x, _fx = np.loadtxt(fn, unpack=True)
        f = interpolate.interp1d(_x, _fx)
        print "Integrating", fn
        try:
            fquadz.append(bquad.int_quadz(f))
            fbayes.append(bquad.int_bayes(f))
        except Exception as e:
            print "...failed:", e
    delta_f = np.array(fbayes)/np.array(fquadz)
    print "bayes/quadz ratio over test data = {:.3f} +/- {:.3f}".format(
            np.mean(delta_f, axis=0), np.std(delta_f, axis=0))

def test_compress_filter():
    """
    Tests against a given dataset
    """
    for filt in ['u', 'B', 'V_9844', 'g', 'r', 'i']:
        filtfname = "CSP_filter_curves/CSP_{}.txt".format(filt)
        pklfname = filtfname.replace('.txt', '_bquad.pkl')
        print "*** Compressing: {} ***".format(filtfname)
        bquad = compress_filter(filtfname, sqlogexp, [0.01], 30)
        print "Writing to", pklfname
        with open(pklfname, 'w') as pklfile:
            pickle.dump(bquad, pklfile, -1)
        integrate_test_suite(bquad)

if __name__ == "__main__":
    test_compress_filter()
<commit_msg>Added brute force optimization (really just for comparison)<commit_after>#!/usr/bin/env python

"""
RS 2017/02/17:  Bayesian quadrature compressions of filter transmission curves

These routines take as input an empirical transmission curve and use Bayesian
quadrature (O'Hagan 1991; Huszar & Duvenaud 2012) as a means for approximating
the integral of that transmission curve against a Gaussian process with known
covariance function.  This will be an important step of making tractable the
problem of inferring a spectral time series against photometry.
"""

import sys
import glob
import dill as pickle
import numpy as np
from scipy import linalg, interpolate, integrate, optimize

class BQFilter(object):
    """
    Implements compression of a filter transfer function using Bayesian
    quadrature.  Chooses evaluation points and weights using a greedy
    algorithm that successively minimizes, at each stage, the mean square
    difference between the original filter function and an approximate
    version as evaluated on functions under a Gaussian process prior.
    Uses the algorithm as outlined in section 3 of
        "Optimally-Weighted Herding is Bayesian Quadrature",
            F. Huszar & D. Duvenaud, Proc. UAI 2012, p. 377.
    """
    # In general, underscores represent internal variables associated
    # with the training points.

    def _vmsg(self, msg):
        """
        Writes unbuffered status messages to stdout.
        """
        if self.verbose:
            print msg
        sys.stdout.flush()

    def __init__(self, _x, _fx, kcov, khyp=[ ], verbose=False):
        """
        Parameters:
            _x:    points at which original filter transfer function is
                   empirically defined (e.g. wavelengths); this will form
                   the probability distribution p for Bayesian quadrature
            _fx:   transfer function values (e.g. filter transmission)
            kcov:  covariance kernel for GP integrand (callable),
                   defined to take two arguments x1, x2
            khyp:  (fixed) hyperparameters for kcov (np.array of floats)
            verbose:  print status messages? (bool)
        """
        # Normalize the training points to zero mean and unit variance,
        # to help stabilize the scipy quadrature routines
        self.verbose = verbose
        self._x, self._fx = _x, _fx
        self._xmu, self._xsig = np.mean(_x, axis=0), np.std(_x, axis=0)
        self._u = self._x2u(self._x)
        self._ulo, self._uhi = self._u.min(), self._u.max()
        # Internal callables for filter transfer function and covariance
        integZu = interpolate.interp1d(self._u, _fx)
        self.Zu, Zu_err = integrate.quad(integZu, self._ulo, self._uhi)
        self._vmsg("__init__:  Zu = {:.3g} +/- {:.3g}".format(self.Zu, Zu_err))
        self._vmsg("__init__:  orig. filter norm = {:.3g}".format(
                   self.Zu * self._xsig))
        self.pu = interpolate.interp1d(self._u, _fx/self.Zu)
        # self.kcov = kcov
        self.kcov, self.khyp = kcov, khyp
        # Internal state
        self.u = np.array([ ])      # quadrature points
        self.zu = np.array([ ])     # quadrature weights
        self.K = np.array([[ ]])    # covariance at quadrature points
        # Starting variance
        self._calc_base_variance_integral()

    def _x2u(self, x):
        return (x - self._xmu)/self._xsig

    def _u2x(self, u):
        return u*self._xsig + self._xmu

    def kxx(self, x1, x2):
        return self.kcov(x1, x2, *(self.khyp))

    def kuu(self, u1, u2):
        return self.kcov(self._u2x(u1), self._u2x(u2), *(self.khyp))

    def _calc_base_variance_integral(self):
        """
        Calculates the integral
            int_dx int_dx' k(x,x') * p(x) * p(x')
        that forms the baseline variance estimate for a BQ filter.
        """
        if False:
            V0, V0_err = 0.101403302569, 1.83139792831e-06
            self.Vn = self.V0 = V0
            return
        # Define a lot of throwaway functions to wrap parts of the problem
        # to match requirements for scipy.integrate.quad.  Rescale x-axis
        # to improve convergence of the integral.
        g = lambda const: (lambda x: const)
        integ_u = lambda u1, u2: self.kuu(u1, u2) * self.pu(u1) * self.pu(u2)
        # Run dblquad -- this should take about 1 min to complete
        self._vmsg("_calc_base_variance_integral:  Calculating...")
        V0, V0_err = integrate.dblquad(
                integ_u, self._ulo, self._uhi, g(self._ulo), g(self._uhi))
        self.Vn = self.V0 = V0
        self._vmsg("_calc_base_variance_integral: V0 = {} +/- {}"
                   .format(V0, V0_err))
        self._vmsg("_calc_base_variance_integral: V0[xval] = {}"
                   .format(V0 * (self._xsig * self.Zu)**2))

    def add_one_point(self):
        """
        Runs optimization for adding a single point to the BQ filter.
        """
        Ktmp = np.array(self.K)
        def u_var(u_n):
            # Wrapper functions for integrands
            integ_u = lambda u: self.kuu(u, u_n) * self.pu(u)
            _zn = integrate.quad(integ_u, self._ulo, self._uhi)[0]
            # Update internal state
            self.u[-1] = u_n
            self.zu[-1] = _zn
            kuu_n = self.kuu(self.u, u_n)
            Ktmp[-1,:] = Ktmp[:,-1] = kuu_n
            # Calculate and return variance = V0 - z.T * inv(K) * z
            self.Kchol = linalg.cholesky(Ktmp, lower=True)
            zeta = linalg.solve_triangular(self.Kchol, self.zu, lower=True)
            self.Vn = self.V0 - np.dot(zeta, zeta)
            return self.Vn
        # Enlarge internal state and optimize over location of new point
        # Since doing this in u, initial guess for new point is 0.0
        n = len(self.u)
        KX = np.atleast_2d([self.u])
        self.u = np.concatenate([self.u, [0.0]])
        self.zu = np.concatenate([self.zu, [0.0]])
        if Ktmp.shape[1] == 0:
            Ktmp = np.array([[1.0]])
        else:
            Ktmp = np.vstack([np.hstack([self.K, KX.T   ]),
                              np.hstack([KX,     [[1.0]]])])
        Ktmp += 1e-8 * np.eye(n+1)
        # As we add more points the Cholesky factor may become more unstable,
        # so add a small nugget -- as small as we can get away with.
        nugget = 1e-07
        self._vmsg("add_one_point:  Optimizing over point #{}...".format(n+1))
        while nugget <= 0.1:
            # try:
                u0 = np.array([0.0])
                bounds = np.array([(self._ulo, self._uhi)])
                cons = [{ 'type': 'ineq', 'fun': lambda u: u - self._ulo },
                        { 'type': 'ineq', 'fun': lambda u: self._uhi - u }]
                result = optimize.minimize(
                        u_var, u0, method='COBYLA', constraints=cons)
                break
            # except Exception as e:
                self._vmsg("add_one_point:  Cholesky factorization failed")
                self._vmsg(str(e))
                self._vmsg("add_one_point:  adding nugget = {} "
                           "to stabilize Cholesky factor".format(nugget))
                Ktmp += nugget * np.eye(n+1)
                nugget *= 10
        if nugget > 0.1:
            self._vmsg("add_one_point:  total Cholesky factorization fail")
            self.u, self.zu = self.u[:-1], self.zu[:-1]
            return
        elif result.success:
            self._vmsg("add_one_point:  Added new point (wt) {} ({}); Vn = {}"
                       .format(self.u[-1], self.zu[-1], self.Vn))
            self.K = Ktmp
        else:
            self._vmsg("add_one_point:  Optimization failed, don't trust me!")
            self._vmsg("Failure message:  {}".format(result.message))
        self.wbq = linalg.cho_solve((self.Kchol, True), self.zu)
        # Return quadrature points and weights for integration transformed
        # back to original x-axis, along with the renormalized variance.
        self.x = self._u2x(self.u)
        self.zx = self.zu * self._xsig * self.Zu
        self.Vxn = self.Vn * (self._xsig * self.Zu)**2

    def add_n_points(self, n=0):
        """
        What it says on the tin:  runs self.add_one_point() n times.
        """
        for i in range(n):
            self.add_one_point()

    def solve_n_points(self, n=0):
        """
        Runs ab initio optimization for an n-point Bayesian quadrature.
        """
        print "self._ulo, self._uhi =", self._ulo, self._uhi
        def u_var(uvec):
            self.u = np.array(uvec)
            self.zu = np.zeros(len(uvec))
            for i, ui in enumerate(uvec):
                integ_u = lambda u: self.kuu(u, ui) * self.pu(u)
                self.zu[i] = integrate.quad(integ_u, self._ulo, self._uhi)[0]
            Ktmp = self.kuu(self.u[:,np.newaxis], self.u[np.newaxis,:])
            Ktmp += 1e-12*np.eye(len(self.u))
            # Calculate and return variance = V0 - z.T * inv(K) * z
            self.Kchol = linalg.cholesky(Ktmp, lower=True)
            zeta = linalg.solve_triangular(self.Kchol, self.zu, lower=True)
            self.Vn = self.V0 - np.dot(zeta, zeta)
            self._vmsg("*** u_var:  uvec =" +
                       ("{:.3f} " * len(uvec)).format(*uvec) +
                       "Vn = {}".format(self.Vn))
            return self.Vn
        # Optimize all the points
        u0 = np.linspace(self._ulo, self._uhi, n+2)[1:-1]
        bounds = np.array([(self._ulo, self._uhi)])
        cons =  [{ 'type': 'ineq', 'fun': lambda u: u[i] - self._ulo }
                 for i in range(n)]
        cons += [{ 'type': 'ineq', 'fun': lambda u: self._uhi - u[i] }
                 for i in range(n)]
        result = optimize.minimize(
                u_var, u0, method='COBYLA', constraints=cons)
        self.wbq = linalg.cho_solve((self.Kchol, True), self.zu)
        # Return quadrature points and weights for integration transformed
        # back to original x-axis, along with the renormalized variance.
        self.x = self._u2x(self.u)
        self.zx = self.zu * self._xsig * self.Zu
        self.Vxn = self.Vn * (self._xsig * self.Zu)**2

    def int_quadz(self, f):
        """
        Uses straight-up quadrature to evaluate integral of f.  In most
        interesting cases f will be an interpolate.interp1d over some
        set of points (for example, an observed supernova spectrum).
        Parameters:
            f: 1-D callable
        """
        integ_u = lambda u: f(self._u2x(u)) * self.pu(u)
        pnorm = self._xsig * self.Zu
        Fu, Fu_err = integrate.quad(integ_u, self._ulo, self._uhi)
        Fx, Fx_err = Fu * pnorm, Fu_err * pnorm
        self._vmsg('int_quadz: F = {} +/- {}'.format(Fx, Fx_err))
        return Fx

    def int_bayes(self, f):
        """
        Uses Bayesian quadrature rule to evaluate integral of f.  The rule
        is derived assuming f is a Gaussian process with a given covariance
        kernel (i.e. fixed hyperparameters).
        Parameters:
            f: 1-D callable
        """
        pnorm = self._xsig * self.Zu
        Fx = np.dot(self.wbq, f(self.x)) * pnorm
        self._vmsg('int_bayes: F = {}'.format(Fx))
        return Fx

def sqexp(x1, x2, l):
    """
    Kernel for the GP, in this case an isotropic square exponential
    """
    return np.exp(-0.5*((x1-x2)/l)**2)

def sqlogexp(x1, x2, logl):
    """
    GP kernel, square exponential in log of variable
    """
    return np.exp(-0.5*((np.log(x1)-np.log(x2))/logl)**2)

def compress_filter(fname, kcov, khyp, n_points):
    """
    Reads in a transfer curve for a filter, and computes an optimal
    Bayesian quadrature rule for a square exponential covariance kernel.
    Parameters:
        fname:  name of two-column text file with (x, y) pairs
        kcov:   covariance kernel for GP integrand (callable),
                defined to take two arguments x1, x2
        khyp:  (fixed) hyperparameters for kcov (np.array of floats)
        n_points:  number of quadrature points desired
    """
    _x, _fx = np.loadtxt(fname, unpack=True)
    bquad = BQFilter(_x, _fx, kcov, khyp, verbose=True)
    bquad.add_n_points(n_points)
    return bquad

def integrate_test_suite(bquad):
    """
    Uses Bayesian quadrature to integrate a bunch of spectra, and compares
    with results from integrating straight against filter.
    """
    fquadz, fbayes = [ ], [ ]
    for fn in glob.glob("testdata/spec*.txt"):
        _x, _fx = np.loadtxt(fn, unpack=True)
        f = interpolate.interp1d(_x, _fx)
        print "Integrating", fn
        try:
            fquadz.append(bquad.int_quadz(f))
            fbayes.append(bquad.int_bayes(f))
        except Exception as e:
            print "...failed:", e
    delta_f = np.array(fbayes)/np.array(fquadz)
    print "bayes/quadz ratio over test data = {:.3f} +/- {:.3f}".format(
            np.mean(delta_f, axis=0), np.std(delta_f, axis=0))

def test_compress_filter():
    """
    Tests against a given dataset
    """
    for filt in ['u', 'B', 'V_9844', 'g', 'r', 'i']:
        filtfname = "CSP_filter_curves/CSP_{}.txt".format(filt)
        pklfname = filtfname.replace('.txt', '_bquad.pkl')
        print "*** Compressing: {} ***".format(filtfname)
        bquad = compress_filter(filtfname, sqlogexp, [0.01], 30)
        print "Writing to", pklfname
        with open(pklfname, 'w') as pklfile:
            pickle.dump(bquad, pklfile, -1)
        integrate_test_suite(bquad)

if __name__ == "__main__":
    test_compress_filter()
<|endoftext|>


<commit_before># The Hazard Library
# Copyright (C) 2016-2018 GEM Foundation
#
# This program 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.
#
# This program 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 this program.  If not, see <http://www.gnu.org/licenses/>.

import os
import numpy as np
import unittest

from openquake.hazardlib import nrml
from openquake.hazardlib.const import TRT
from openquake.hazardlib.geo import Point
from openquake.baselib.general import DictArray
from openquake.hazardlib.calc.filters import SourceFilter
from openquake.hazardlib.site import Site, SiteCollection
from openquake.hazardlib.gsim.sadigh_1997 import SadighEtAl1997
from openquake.hazardlib.sourceconverter import SourceConverter
from openquake.hazardlib.calc.hazard_curve import calc_hazard_curves

DATA = os.path.join(os.path.dirname(__file__), '../source_model')

class HazardCurvesClusterTestCase01(unittest.TestCase):

    def setUp(self):
        testfile = os.path.join(DATA, 'source_group_cluster.xml')
        sc = SourceConverter(area_source_discretization=10.,
                             investigation_time=1.)
        # This provides a SourceModel
        self.sg = getattr(nrml.to_python(testfile, sc), 'src_groups')
        self.imtls = DictArray({'PGA': [0.01, 0.1, 0.2, 0.3, 1.0]})
        gsim = SadighEtAl1997()
        self.gsim_by_trt = {TRT.ACTIVE_SHALLOW_CRUST: gsim}
        site = Site(Point(1.0, -0.1), 800, z1pt0=30., z2pt5=1.)
        s_filter = SourceFilter(SiteCollection([site]), {})
        self.sites = s_filter

    def test_hazard_curve(self):
        """ Classical PSHA with cluster source """
        curves = calc_hazard_curves(self.sg,
                                    self.sites,
                                    self.imtls,
                                    self.gsim_by_trt,
                                    truncation_level=None)
        crv = curves[0][0]
        rates = np.array([1.00000000e-03, 9.98565030e-04, 8.42605169e-04,
                          4.61559062e-04, 1.10100503e-06])
        expected = 1 - np.exp(-rates)
        np.testing.assert_almost_equal(crv, expected)
<commit_msg>Now test passing<commit_after># The Hazard Library
# Copyright (C) 2016-2018 GEM Foundation
#
# This program 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.
#
# This program 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 this program.  If not, see <http://www.gnu.org/licenses/>.

import os
import numpy as np
import unittest

from openquake.hazardlib import nrml
from openquake.hazardlib.const import TRT
from openquake.hazardlib.geo import Point
from openquake.baselib.general import DictArray
from openquake.hazardlib.calc.filters import SourceFilter
from openquake.hazardlib.site import Site, SiteCollection
from openquake.hazardlib.gsim.sadigh_1997 import SadighEtAl1997
from openquake.hazardlib.sourceconverter import SourceConverter
from openquake.hazardlib.calc.hazard_curve import calc_hazard_curves

DATA = os.path.join(os.path.dirname(__file__), '../source_model')

class HazardCurvesClusterTestCase01(unittest.TestCase):

    def setUp(self):
        testfile = os.path.join(DATA, 'source_group_cluster.xml')
        sc = SourceConverter(area_source_discretization=10.,
                             investigation_time=1.)
        # This provides a SourceModel
        self.sg = getattr(nrml.to_python(testfile, sc), 'src_groups')
        self.imtls = DictArray({'PGA': [0.01, 0.1, 0.2, 0.3, 1.0]})
        gsim = SadighEtAl1997()
        self.gsim_by_trt = {TRT.ACTIVE_SHALLOW_CRUST: gsim}
        site = Site(Point(1.0, -0.1), 800, z1pt0=30., z2pt5=1.)
        s_filter = SourceFilter(SiteCollection([site]), {})
        self.sites = s_filter

    def test_hazard_curve(self):
        """ Classical PSHA with cluster source """
        curves = calc_hazard_curves(self.sg,
                                    self.sites,
                                    self.imtls,
                                    self.gsim_by_trt,
                                    truncation_level=None)
        crv = curves[0][0]
        # Expected results computed with a notebook using the original USGS
        # formulation as described in Appendix F of Petersen et al. (2008).
        # The rates of exceedance were converted a posteriori into
        # probabilities.
        rates = np.array([1.00000000e-03, 9.98565030e-04, 8.42605169e-04,
                          4.61559062e-04, 1.10100503e-06])
        expected = 1 - np.exp(-rates)
        np.testing.assert_almost_equal(crv, expected)
<|endoftext|>


<commit_before># Copyright (c) 2015 Ultimaker B.V.
# Uranium is released under the terms of the AGPLv3 or higher.

from PyQt5.QtCore import Qt, pyqtSlot, pyqtProperty, pyqtSignal

from UM.Qt.ListModel import ListModel
from UM.Application import Application
from UM.Settings.MachineInstance import MachineInstance

class MachineDefinitionsModel(ListModel):
    IdRole = Qt.UserRole + 1
    NameRole = Qt.UserRole + 2
    VariantNameRole = Qt.UserRole + 3
    ManufacturerRole = Qt.UserRole + 4
    AuthorRole = Qt.UserRole + 5
    PagesRole = Qt.UserRole + 6

    def __init__(self, parent = None):
        super().__init__(parent)

        self.addRoleName(self.IdRole, "id")
        self.addRoleName(self.NameRole, "name")
        self.addRoleName(self.VariantNameRole, "variantName")
        self.addRoleName(self.ManufacturerRole, "manufacturer")
        self.addRoleName(self.AuthorRole, "author")
        self.addRoleName(self.PagesRole, "pages")

        self._manager = Application.getInstance().getMachineManager()

        self._show_variants = True

        self._manager.machineDefinitionsChanged.connect(self._onMachinesChanged)
        self._onMachinesChanged()

    showVariantsChanged = pyqtSignal()

    def setShowVariants(self, show):
        if show != self._show_variants:
            self._show_variants = show
            self._onMachinesChanged()
            self.showVariantsChanged.emit()

    @pyqtProperty(bool, fset = setShowVariants, notify = showVariantsChanged)
    def showVariants(self):
        return self._show_variants

    @pyqtSlot(str, str)
    def createInstance(self, name, definition_id):
        definition = self._manager.findMachineDefinition(definition_id)

        instance = MachineInstance(self._manager, name = name, definition = definition)
        self._manager.addMachineInstance(instance)
        self._manager.setActiveMachineInstance(instance)

    def _onMachinesChanged(self):
        self.clear()

        definitions = self._manager.getMachineDefinitions(include_variants = self._show_variants)
        definitions = filter(lambda s: s.isVisible(), definitions)
        definitions = sorted(definitions)

        for machine in definitions:
            self.appendItem({
                "id": machine.getId(),
                "name": machine.getName(),
                "variantName": machine.getVariantName(),
                "manufacturer": machine.getManufacturer(),
                "author": machine.getAuthor(),
                "pages": machine.getPages()
            })
<commit_msg>Fix a crash when adding a new machine on OSX<commit_after># Copyright (c) 2015 Ultimaker B.V.
# Uranium is released under the terms of the AGPLv3 or higher.

from PyQt5.QtCore import Qt, pyqtSlot, pyqtProperty, pyqtSignal

from UM.Qt.ListModel import ListModel
from UM.Application import Application
from UM.Settings.MachineInstance import MachineInstance
from UM.Event import CallFunctionEvent

class MachineDefinitionsModel(ListModel):
    IdRole = Qt.UserRole + 1
    NameRole = Qt.UserRole + 2
    VariantNameRole = Qt.UserRole + 3
    ManufacturerRole = Qt.UserRole + 4
    AuthorRole = Qt.UserRole + 5
    PagesRole = Qt.UserRole + 6

    def __init__(self, parent = None):
        super().__init__(parent)

        self.addRoleName(self.IdRole, "id")
        self.addRoleName(self.NameRole, "name")
        self.addRoleName(self.VariantNameRole, "variantName")
        self.addRoleName(self.ManufacturerRole, "manufacturer")
        self.addRoleName(self.AuthorRole, "author")
        self.addRoleName(self.PagesRole, "pages")

        self._manager = Application.getInstance().getMachineManager()

        self._show_variants = True

        self._manager.machineDefinitionsChanged.connect(self._onMachinesChanged)
        self._onMachinesChanged()

    showVariantsChanged = pyqtSignal()

    def setShowVariants(self, show):
        if show != self._show_variants:
            self._show_variants = show
            self._onMachinesChanged()
            self.showVariantsChanged.emit()

    @pyqtProperty(bool, fset = setShowVariants, notify = showVariantsChanged)
    def showVariants(self):
        return self._show_variants

    @pyqtSlot(str, str)
    def createInstance(self, name, definition_id):
        definition = self._manager.findMachineDefinition(definition_id)

        instance = MachineInstance(self._manager, name = name, definition = definition)
        self._manager.addMachineInstance(instance)

        # Workaround for an issue on OSX where directly calling setActiveMachineInstance would
        # crash in the QML garbage collector.
        event = CallFunctionEvent(self._manager.setActiveMachineInstance, [instance], {})
        Application.getInstance().functionEvent(event)

    def _onMachinesChanged(self):
        self.clear()

        definitions = self._manager.getMachineDefinitions(include_variants = self._show_variants)
        definitions = filter(lambda s: s.isVisible(), definitions)
        definitions = sorted(definitions)

        for machine in definitions:
            self.appendItem({
                "id": machine.getId(),
                "name": machine.getName(),
                "variantName": machine.getVariantName(),
                "manufacturer": machine.getManufacturer(),
                "author": machine.getAuthor(),
                "pages": machine.getPages()
            })
<|endoftext|>


<commit_before>#
# This sets up how models are displayed
# in the web admin interface.
#
from builtins import object

from django import forms
from django.conf import settings
from django.contrib import admin
from evennia.typeclasses.admin import AttributeInline, TagInline
from evennia.objects.models import ObjectDB
from django.contrib.admin.utils import flatten_fieldsets

class ObjectAttributeInline(AttributeInline):
    """
    Defines inline descriptions of Attributes (experimental)

    """
    model = ObjectDB.db_attributes.through
    related_field = "objectdb"

class ObjectTagInline(TagInline):
    """
    Defines inline descriptions of Tags (experimental)

    """
    model = ObjectDB.db_tags.through
    related_field = "objectdb"

class ObjectCreateForm(forms.ModelForm):
    """
    This form details the look of the fields.

    """
    class Meta(object):
        model = ObjectDB
        fields = '__all__'
    db_key = forms.CharField(label="Name/Key",
                             widget=forms.TextInput(attrs={'size': '78'}),
                             help_text="Main identifier, like 'apple', 'strong guy', 'Elizabeth' etc. If creating a Character, check so the name is unique among characters!",)
    db_typeclass_path = forms.CharField(label="Typeclass",
                                        initial=settings.BASE_OBJECT_TYPECLASS,
                                        widget=forms.TextInput(attrs={'size': '78'}),
                                        help_text="This defines what 'type' of entity this is. This variable holds a Python path to a module with a valid Evennia Typeclass. If you are creating a Character you should use the typeclass defined by settings.BASE_CHARACTER_TYPECLASS or one derived from that.")
    db_cmdset_storage = forms.CharField(label="CmdSet",
                                        initial="",
                                        required=False,
                                        widget=forms.TextInput(attrs={'size': '78'}),
                                        help_text="Most non-character objects don't need a cmdset and can leave this field blank.")
    raw_id_fields = ('db_destination', 'db_location', 'db_home')

class ObjectEditForm(ObjectCreateForm):
    """
    Form used for editing. Extends the create one with more fields

    """

    class Meta(object):
        fields = '__all__'
    db_lock_storage = forms.CharField(label="Locks",
                                      required=False,
                                      widget=forms.Textarea(attrs={'cols':'100', 'rows':'2'}),
                                      help_text="In-game lock definition string. If not given, defaults will be used. This string should be on the form <i>type:lockfunction(args);type2:lockfunction2(args);...")

class ObjectDBAdmin(admin.ModelAdmin):
    """
    Describes the admin page for Objects.

    """

    inlines = [ObjectTagInline, ObjectAttributeInline]
    list_display = ('id', 'db_key', 'db_player', 'db_typeclass_path')
    list_display_links = ('id', 'db_key')
    ordering = ['db_player', 'db_typeclass_path', 'id']
    search_fields = ['^db_key', 'db_typeclass_path']
    raw_id_fields = ('db_destination', 'db_location', 'db_home')

    save_as = True
    save_on_top = True
    list_select_related = True
    list_filter = ('db_typeclass_path',)

    # editing fields setup

    form = ObjectEditForm
    fieldsets = (
        (None, {
                'fields': (('db_key','db_typeclass_path'), ('db_lock_storage', ),
                           ('db_location', 'db_home'), 'db_destination','db_cmdset_storage'
                           )}),
        )

    add_form = ObjectCreateForm
    add_fieldsets = (
        (None, {
                'fields': (('db_key','db_typeclass_path'),
                           ('db_location', 'db_home'), 'db_destination', 'db_cmdset_storage'
                           )}),
        )

    def get_fieldsets(self, request, obj=None):
        """
        Return fieldsets.

        Args:
            request (Request): Incoming request.
            obj (ObjectDB, optional): Database object.
        """
        if not obj:
            return self.add_fieldsets
        return super(ObjectDBAdmin, self).get_fieldsets(request, obj)

    def get_form(self, request, obj=None, **kwargs):
        """
        Use special form during creation.

        Args:
            request (Request): Incoming request.
            obj (Object, optional): Database object.

        """
        defaults = {}
        if obj is None:
            defaults.update({
                'form': self.add_form,
                'fields': flatten_fieldsets(self.add_fieldsets),
            })
            defaults.update(kwargs)
        return super(ObjectDBAdmin, self).get_form(request, obj, **defaults)

    def save_model(self, request, obj, form, change):
        """
        Model-save hook.

        Args:
            request (Request): Incoming request.
            obj (Object): Database object.
            form (Form): Form instance.
            change (bool): If this is a change or a new object.

        """
        obj.save()
        if not change:
            # adding a new object
            # have to call init with typeclass passed to it
            obj.set_class_from_typeclass(typeclass_path=obj.db_typeclass_path)
            obj.basetype_setup()
            obj.basetype_posthook_setup()
            obj.at_object_creation()
        obj.at_init()

admin.site.register(ObjectDB, ObjectDBAdmin)
<commit_msg>Fix the save and continue editing option<commit_after>#
# This sets up how models are displayed
# in the web admin interface.
#
from builtins import object

from django import forms
from django.conf import settings
from django.contrib import admin
from evennia.typeclasses.admin import AttributeInline, TagInline
from evennia.objects.models import ObjectDB
from django.contrib.admin.utils import flatten_fieldsets

class ObjectAttributeInline(AttributeInline):
    """
    Defines inline descriptions of Attributes (experimental)

    """
    model = ObjectDB.db_attributes.through
    related_field = "objectdb"

class ObjectTagInline(TagInline):
    """
    Defines inline descriptions of Tags (experimental)

    """
    model = ObjectDB.db_tags.through
    related_field = "objectdb"

class ObjectCreateForm(forms.ModelForm):
    """
    This form details the look of the fields.

    """
    class Meta(object):
        model = ObjectDB
        fields = '__all__'
    db_key = forms.CharField(label="Name/Key",
                             widget=forms.TextInput(attrs={'size': '78'}),
                             help_text="Main identifier, like 'apple', 'strong guy', 'Elizabeth' etc. If creating a Character, check so the name is unique among characters!",)
    db_typeclass_path = forms.CharField(label="Typeclass",
                                        initial=settings.BASE_OBJECT_TYPECLASS,
                                        widget=forms.TextInput(attrs={'size': '78'}),
                                        help_text="This defines what 'type' of entity this is. This variable holds a Python path to a module with a valid Evennia Typeclass. If you are creating a Character you should use the typeclass defined by settings.BASE_CHARACTER_TYPECLASS or one derived from that.")
    db_cmdset_storage = forms.CharField(label="CmdSet",
                                        initial="",
                                        required=False,
                                        widget=forms.TextInput(attrs={'size': '78'}),
                                        help_text="Most non-character objects don't need a cmdset and can leave this field blank.")
    raw_id_fields = ('db_destination', 'db_location', 'db_home')

class ObjectEditForm(ObjectCreateForm):
    """
    Form used for editing. Extends the create one with more fields

    """

    class Meta(object):
        fields = '__all__'
    db_lock_storage = forms.CharField(label="Locks",
                                      required=False,
                                      widget=forms.Textarea(attrs={'cols':'100', 'rows':'2'}),
                                      help_text="In-game lock definition string. If not given, defaults will be used. This string should be on the form <i>type:lockfunction(args);type2:lockfunction2(args);...")

class ObjectDBAdmin(admin.ModelAdmin):
    """
    Describes the admin page for Objects.

    """

    inlines = [ObjectTagInline, ObjectAttributeInline]
    list_display = ('id', 'db_key', 'db_player', 'db_typeclass_path')
    list_display_links = ('id', 'db_key')
    ordering = ['db_player', 'db_typeclass_path', 'id']
    search_fields = ['^db_key', 'db_typeclass_path']
    raw_id_fields = ('db_destination', 'db_location', 'db_home')

    save_as = True
    save_on_top = True
    list_select_related = True
    list_filter = ('db_typeclass_path',)

    # editing fields setup

    form = ObjectEditForm
    fieldsets = (
        (None, {
                'fields': (('db_key','db_typeclass_path'), ('db_lock_storage', ),
                           ('db_location', 'db_home'), 'db_destination','db_cmdset_storage'
                           )}),
        )

    add_form = ObjectCreateForm
    add_fieldsets = (
        (None, {
                'fields': (('db_key','db_typeclass_path'),
                           ('db_location', 'db_home'), 'db_destination', 'db_cmdset_storage'
                           )}),
        )

    def get_urls(self):
        urls = super(ObjectDBAdmin, self).get_urls()
        print urls
        return urls

    def get_fieldsets(self, request, obj=None):
        """
        Return fieldsets.

        Args:
            request (Request): Incoming request.
            obj (ObjectDB, optional): Database object.
        """
        if not obj:
            return self.add_fieldsets
        return super(ObjectDBAdmin, self).get_fieldsets(request, obj)

    def get_form(self, request, obj=None, **kwargs):
        """
        Use special form during creation.

        Args:
            request (Request): Incoming request.
            obj (Object, optional): Database object.

        """
        defaults = {}
        if obj is None:
            defaults.update({
                'form': self.add_form,
                'fields': flatten_fieldsets(self.add_fieldsets),
            })
            defaults.update(kwargs)
        return super(ObjectDBAdmin, self).get_form(request, obj, **defaults)

    def save_model(self, request, obj, form, change):
        """
        Model-save hook.

        Args:
            request (Request): Incoming request.
            obj (Object): Database object.
            form (Form): Form instance.
            change (bool): If this is a change or a new object.

        """
        obj.save()
        if not change:
            # adding a new object
            # have to call init with typeclass passed to it
            obj.set_class_from_typeclass(typeclass_path=obj.db_typeclass_path)
            obj.basetype_setup()
            obj.basetype_posthook_setup()
            obj.at_object_creation()
        obj.at_init()

    def response_add(self, request, obj, post_url_continue=None):
        if '_continue' in request.POST:
            from django.http import HttpResponseRedirect
            from django.core.urlresolvers import reverse
            return HttpResponseRedirect(reverse("admin:objects_objectdb_change", args=[obj.id]))
        return super(ObjectDBAdmin, self).response_add(request, obj, post_url_continue)

admin.site.register(ObjectDB, ObjectDBAdmin)
<|endoftext|>