<commit_before>#  Copyright (c) 2021, Manfred Moitzi
#  License: MIT License
from typing import Optional, Dict, List, Tuple, Iterable
from pathlib import Path
import enum
from ezdxf.lldxf.loader import SectionDict
from ezdxf.addons.browser.loader import load_section_dict
from ezdxf.lldxf.types import DXFVertex, tag_type
from ezdxf.lldxf.tags import Tags

__all__ = [
    "DXFDocument",
    "get_row_from_line_number",
    "dxfstr",
    "EntityHistory",
    "SearchIndex",
]

class DXFDocument:
    def __init__(self, sections: SectionDict = None):
        # Important: the section dict has to store the raw string tags
        # else an association of line numbers to entities is not possible.
        # Comment tags (999) are ignored, because the load_section_dict()
        # function can not handle and store comments.
        # Therefore comments causes error in the line number associations and
        # should be stripped off before processing for precise debugging of
        # DXF files.
        # TODO: ezdxf strip-comments FILE
        self.sections: SectionDict = dict()
        self.handle_index: Optional[HandleIndex] = None
        self.line_index: Optional[LineIndex] = None
        self.valid_handles = None
        self.filename = ""
        if sections:
            self.update(sections)

    @property
    def filepath(self):
        return Path(self.filename)

    @property
    def max_line_number(self) -> int:
        if self.line_index:
            return self.line_index.max_line_number
        else:
            return 1

    def load(self, filename: str):
        self.filename = filename
        self.update(load_section_dict(filename))

    def update(self, sections: SectionDict):
        self.sections = sections
        self.handle_index = HandleIndex(self.sections)
        self.line_index = LineIndex(self.sections)

    def absolute_filepath(self):
        return self.filepath.absolute()

    def get_header_section(self):
        return self.sections.get("HEADER")

    def get_entity(self, handle: str) -> Optional[Tags]:
        if self.handle_index:
            return self.handle_index.get(handle)
        return None

    def get_line_number(self, entity: Tags) -> Optional[int]:
        if self.line_index:
            return self.line_index.get_start_line_for_entity(entity)
        return None

    def get_entity_at_line(self, number: int) -> Optional[Tags]:
        if self.line_index:
            return self.line_index.get_entity_at_line(number)
        return None

    def next_entity(self, entity: Tags) -> Optional[Tags]:
        return self.handle_index.next_entity(entity)

    def previous_entity(self, entity: Tags) -> Optional[Tags]:
        return self.handle_index.previous_entity(entity)

class HandleIndex:
    def __init__(self, sections: SectionDict):
        self._index = HandleIndex.build(sections)

    def __contains__(self, handle: str) -> bool:
        return handle.upper() in self._index

    def get(self, handle: str):
        return self._index.get(handle.upper())

    @staticmethod
    def build(sections: SectionDict) -> Dict:
        entity_index = dict()
        for section in sections.values():
            for entity in section:
                try:
                    handle = entity.get_handle()
                    entity_index[handle.upper()] = entity
                except ValueError:
                    pass
        return entity_index

    def next_entity(self, entity: Tags) -> Tags:
        return_next = False
        for e in self._index.values():
            if return_next:
                return e
            if e is entity:
                return_next = True
        return entity

    def previous_entity(self, entity: Tags) -> Tags:
        prev = entity
        for e in self._index.values():
            if e is entity:
                return prev
            prev = e
        return entity

class LineIndex:
    def __init__(self, sections: SectionDict):
        # id, (start_line_number, entity tags)
        self._entity_index: Dict[
            int, Tuple[int, Tags]
        ] = LineIndex.build_entity_index(sections)

        # entity index of sorted (start_line_number, entity) tuples
        index = LineIndex.build_line_index(sections)
        self._line_index: List[Tuple[int, Tags]] = index
        self._max_line_number = 1
        if index:
            last_start_number, e = self._line_index[-1]
            self._max_line_number = last_start_number + len(e) * 2 - 1

    @property
    def max_line_number(self) -> int:
        return self._max_line_number

    @staticmethod
    def build_entity_index(sections: SectionDict) -> Dict:
        index: Dict[int, Tuple[int, Tags]] = dict()
        line_number = 1
        for section in sections.values():
            # the section dict contain raw string tags
            for entity in section:
                index[id(entity)] = line_number, entity
                line_number += len(entity) * 2  # group code, value
            line_number += 2  # for missing ENDSEC tag
        return index

    @staticmethod
    def build_line_index(sections: SectionDict) -> List:
        index: List[Tuple[int, Tags]] = list()
        start_line_number = 1
        for name, section in sections.items():
            # the section dict contain raw string tags
            for entity in section:
                index.append((start_line_number, entity))
                # add 2 lines for each tag: group code, value
                start_line_number += len(entity) * 2
            start_line_number += 2  # for missing ENDSEC tag
        index.sort()  # sort index by line number
        return index

    def get_start_line_for_entity(self, entity: Tags) -> Optional[int]:
        entry = self._entity_index.get(id(entity))
        if entry:
            return entry[0]
        return None

    def get_entity_at_line(self, number: int) -> Optional[Tags]:
        index = self._line_index
        if len(index) == 0:
            return None

        _, entity = index[0]  # first entity
        for start, e in index:
            if start > number:
                return entity
            entity = e
        return entity

def get_row_from_line_number(
    entity: Tags, start_line_number: int, select_line_number: int
) -> int:
    count = select_line_number - start_line_number
    lines = 0
    row = 0
    for tag in entity:
        if lines >= count:
            return row
        if isinstance(tag, DXFVertex):
            lines += len(tag.value) * 2
        else:
            lines += 2
        row += 1
    return row

def dxfstr(tags: Tags) -> str:
    return "".join(tag.dxfstr() for tag in tags)

class EntityHistory:
    def __init__(self):
        self._history: List[Tags] = list()
        self._index: int = 0
        self._time_travel: List[Tags] = list()

    def __len__(self):
        return len(self._history)

    @property
    def index(self):
        return self._index

    def clear(self):
        self._history.clear()
        self._time_travel.clear()
        self._index = 0

    def append(self, entity: Tags):
        if self._time_travel:
            self._history.extend(self._time_travel)
            self._time_travel.clear()

        self._index = len(self._history)
        self._history.append(entity)

    def back(self) -> Optional[Tags]:
        entity = None
        if self._history:
            index = self._index - 1
            if index >= 0:
                entity = self._time_wrap(index)
            else:
                entity = self._history[0]
        return entity

    def forward(self) -> Tags:
        entity = None
        history = self._history
        if history:
            index = self._index + 1
            if index < len(history):
                entity = self._time_wrap(index)
            else:
                entity = history[-1]
        return entity

    def _time_wrap(self, index) -> Tags:
        self._index = index
        entity = self._history[index]
        self._time_travel.append(entity)
        return entity

    def content(self) -> List[Tags]:
        return list(self._history)

class SearchIndex:
    NOT_FOUND = None, -1

    def __init__(self, entities: Iterable[Tags]):
        self.entities: List[Tags] = list(entities)
        self._index: int = 0
        self._search_term: Optional[str] = None
        self._backward = False
        self.case_insensitive = True
        self.whole_words = False
        self.numbers = False
        self.regex = False  # False = normal mode

    def set_current_entity(self, entity: Tags):
        try:
            self._index = self.entities.index(entity)
        except ValueError:
            self._index = 0

    def reset(self, search_term: str, backward=False):
        self._index = 0
        self._search_term = str(search_term)
        count = len(self.entities)
        if backward and count:
            self._index = count - 1

    def find(self, term: str, backward=False) -> Tuple[Optional[Tags], int]:
        self.reset(term, backward)
        if len(self.entities):
            if backward:
                return self.find_backward()
            else:
                return self.find_next()
        else:
            return self.NOT_FOUND

    def find_next(self) -> Tuple[Optional[Tags], int]:
        count = len(self.entities)
        while self._index < count:
            entity = self.entities[self._index]
            self._index += 1
            tag_index = self.match(entity)
            if tag_index >= 0:
                return entity, tag_index
        return self.NOT_FOUND

    def find_backward(self) -> Tuple[Optional[Tags], int]:
        while self._index >= 0:
            entity = self.entities[self._index]
            self._index -= 1
            tag_index = self.match(entity)
            if tag_index >= 0:
                return entity, tag_index
        return self.NOT_FOUND

    def match(self, entity: Tags) -> int:
        if self._search_term:
            regex = self.regex
            whole_words = self.whole_words
            case_insensitive = self.case_insensitive
            if regex:
                case_insensitive = False
                whole_words = False

            if self.case_insensitive:
                search_term = self._search_term.lower()
            else:
                search_term = self._search_term
            not_numbers = not self.numbers
            for index, (code, value) in enumerate(entity):
                if tag_type(code) is not str:
                    if not_numbers:
                        continue
                    value = str(value)

                if case_insensitive:
                    value = value.lower()

                if regex:
                    pass
                elif whole_words:
                    for word in value.split():
                        if search_term == word:
                            return index
                else:
                    if search_term in value:
                        return index
        return -1
<commit_msg>refactor SearchIndex.reset()<commit_after>#  Copyright (c) 2021, Manfred Moitzi
#  License: MIT License
from typing import Optional, Dict, List, Tuple, Iterable
from pathlib import Path
import enum
from ezdxf.lldxf.loader import SectionDict
from ezdxf.addons.browser.loader import load_section_dict
from ezdxf.lldxf.types import DXFVertex, tag_type
from ezdxf.lldxf.tags import Tags

__all__ = [
    "DXFDocument",
    "get_row_from_line_number",
    "dxfstr",
    "EntityHistory",
    "SearchIndex",
]

class DXFDocument:
    def __init__(self, sections: SectionDict = None):
        # Important: the section dict has to store the raw string tags
        # else an association of line numbers to entities is not possible.
        # Comment tags (999) are ignored, because the load_section_dict()
        # function can not handle and store comments.
        # Therefore comments causes error in the line number associations and
        # should be stripped off before processing for precise debugging of
        # DXF files.
        # TODO: ezdxf strip-comments FILE
        self.sections: SectionDict = dict()
        self.handle_index: Optional[HandleIndex] = None
        self.line_index: Optional[LineIndex] = None
        self.valid_handles = None
        self.filename = ""
        if sections:
            self.update(sections)

    @property
    def filepath(self):
        return Path(self.filename)

    @property
    def max_line_number(self) -> int:
        if self.line_index:
            return self.line_index.max_line_number
        else:
            return 1

    def load(self, filename: str):
        self.filename = filename
        self.update(load_section_dict(filename))

    def update(self, sections: SectionDict):
        self.sections = sections
        self.handle_index = HandleIndex(self.sections)
        self.line_index = LineIndex(self.sections)

    def absolute_filepath(self):
        return self.filepath.absolute()

    def get_header_section(self):
        return self.sections.get("HEADER")

    def get_entity(self, handle: str) -> Optional[Tags]:
        if self.handle_index:
            return self.handle_index.get(handle)
        return None

    def get_line_number(self, entity: Tags) -> Optional[int]:
        if self.line_index:
            return self.line_index.get_start_line_for_entity(entity)
        return None

    def get_entity_at_line(self, number: int) -> Optional[Tags]:
        if self.line_index:
            return self.line_index.get_entity_at_line(number)
        return None

    def next_entity(self, entity: Tags) -> Optional[Tags]:
        return self.handle_index.next_entity(entity)

    def previous_entity(self, entity: Tags) -> Optional[Tags]:
        return self.handle_index.previous_entity(entity)

class HandleIndex:
    def __init__(self, sections: SectionDict):
        self._index = HandleIndex.build(sections)

    def __contains__(self, handle: str) -> bool:
        return handle.upper() in self._index

    def get(self, handle: str):
        return self._index.get(handle.upper())

    @staticmethod
    def build(sections: SectionDict) -> Dict:
        entity_index = dict()
        for section in sections.values():
            for entity in section:
                try:
                    handle = entity.get_handle()
                    entity_index[handle.upper()] = entity
                except ValueError:
                    pass
        return entity_index

    def next_entity(self, entity: Tags) -> Tags:
        return_next = False
        for e in self._index.values():
            if return_next:
                return e
            if e is entity:
                return_next = True
        return entity

    def previous_entity(self, entity: Tags) -> Tags:
        prev = entity
        for e in self._index.values():
            if e is entity:
                return prev
            prev = e
        return entity

class LineIndex:
    def __init__(self, sections: SectionDict):
        # id, (start_line_number, entity tags)
        self._entity_index: Dict[
            int, Tuple[int, Tags]
        ] = LineIndex.build_entity_index(sections)

        # entity index of sorted (start_line_number, entity) tuples
        index = LineIndex.build_line_index(sections)
        self._line_index: List[Tuple[int, Tags]] = index
        self._max_line_number = 1
        if index:
            last_start_number, e = self._line_index[-1]
            self._max_line_number = last_start_number + len(e) * 2 - 1

    @property
    def max_line_number(self) -> int:
        return self._max_line_number

    @staticmethod
    def build_entity_index(sections: SectionDict) -> Dict:
        index: Dict[int, Tuple[int, Tags]] = dict()
        line_number = 1
        for section in sections.values():
            # the section dict contain raw string tags
            for entity in section:
                index[id(entity)] = line_number, entity
                line_number += len(entity) * 2  # group code, value
            line_number += 2  # for missing ENDSEC tag
        return index

    @staticmethod
    def build_line_index(sections: SectionDict) -> List:
        index: List[Tuple[int, Tags]] = list()
        start_line_number = 1
        for name, section in sections.items():
            # the section dict contain raw string tags
            for entity in section:
                index.append((start_line_number, entity))
                # add 2 lines for each tag: group code, value
                start_line_number += len(entity) * 2
            start_line_number += 2  # for missing ENDSEC tag
        index.sort()  # sort index by line number
        return index

    def get_start_line_for_entity(self, entity: Tags) -> Optional[int]:
        entry = self._entity_index.get(id(entity))
        if entry:
            return entry[0]
        return None

    def get_entity_at_line(self, number: int) -> Optional[Tags]:
        index = self._line_index
        if len(index) == 0:
            return None

        _, entity = index[0]  # first entity
        for start, e in index:
            if start > number:
                return entity
            entity = e
        return entity

def get_row_from_line_number(
    entity: Tags, start_line_number: int, select_line_number: int
) -> int:
    count = select_line_number - start_line_number
    lines = 0
    row = 0
    for tag in entity:
        if lines >= count:
            return row
        if isinstance(tag, DXFVertex):
            lines += len(tag.value) * 2
        else:
            lines += 2
        row += 1
    return row

def dxfstr(tags: Tags) -> str:
    return "".join(tag.dxfstr() for tag in tags)

class EntityHistory:
    def __init__(self):
        self._history: List[Tags] = list()
        self._index: int = 0
        self._time_travel: List[Tags] = list()

    def __len__(self):
        return len(self._history)

    @property
    def index(self):
        return self._index

    def clear(self):
        self._history.clear()
        self._time_travel.clear()
        self._index = 0

    def append(self, entity: Tags):
        if self._time_travel:
            self._history.extend(self._time_travel)
            self._time_travel.clear()

        self._index = len(self._history)
        self._history.append(entity)

    def back(self) -> Optional[Tags]:
        entity = None
        if self._history:
            index = self._index - 1
            if index >= 0:
                entity = self._time_wrap(index)
            else:
                entity = self._history[0]
        return entity

    def forward(self) -> Tags:
        entity = None
        history = self._history
        if history:
            index = self._index + 1
            if index < len(history):
                entity = self._time_wrap(index)
            else:
                entity = history[-1]
        return entity

    def _time_wrap(self, index) -> Tags:
        self._index = index
        entity = self._history[index]
        self._time_travel.append(entity)
        return entity

    def content(self) -> List[Tags]:
        return list(self._history)

class SearchIndex:
    NOT_FOUND = None, -1

    def __init__(self, entities: Iterable[Tags]):
        self.entities: List[Tags] = list(entities)
        self._index: int = 0
        self._search_term: Optional[str] = None
        self._backward = False
        self.case_insensitive = True
        self.whole_words = False
        self.numbers = False
        self.regex = False  # False = normal mode

    def set_current_entity(self, entity: Tags):
        try:
            self._index = self.entities.index(entity)
        except ValueError:
            self._index = 0

    def reset_index(self, backward=False):
        self._index = 0
        count = len(self.entities)
        if backward and count:
            self._index = count - 1

    def reset_search_term(self, term: str):
        self._search_term = str(term)

    def find(
        self, term: str, backward=False, reset_index=True
    ) -> Tuple[Optional[Tags], int]:
        self.reset_search_term(term)
        if reset_index:
            self.reset_index(backward)
        if len(self.entities):
            if backward:
                return self.find_backward()
            else:
                return self.find_next()
        else:
            return self.NOT_FOUND

    def find_next(self) -> Tuple[Optional[Tags], int]:
        count = len(self.entities)
        while self._index < count:
            entity = self.entities[self._index]
            self._index += 1
            tag_index = self.match(entity)
            if tag_index >= 0:
                return entity, tag_index
        return self.NOT_FOUND

    def find_backward(self) -> Tuple[Optional[Tags], int]:
        while self._index >= 0:
            entity = self.entities[self._index]
            self._index -= 1
            tag_index = self.match(entity)
            if tag_index >= 0:
                return entity, tag_index
        return self.NOT_FOUND

    def match(self, entity: Tags) -> int:
        if self._search_term:
            regex = self.regex
            whole_words = self.whole_words
            case_insensitive = self.case_insensitive
            if regex:
                case_insensitive = False
                whole_words = False

            if self.case_insensitive:
                search_term = self._search_term.lower()
            else:
                search_term = self._search_term
            not_numbers = not self.numbers
            for index, (code, value) in enumerate(entity):
                if tag_type(code) is not str:
                    if not_numbers:
                        continue
                    value = str(value)

                if case_insensitive:
                    value = value.lower()

                if regex:
                    pass
                elif whole_words:
                    for word in value.split():
                        if search_term == word:
                            return index
                else:
                    if search_term in value:
                        return index
        return -1
<|endoftext|>


<commit_before>from __future__ import absolute_import
from nose.tools import (assert_equal, assert_not_equal, assert_almost_equal,
                        raises)
from nose.plugins.skip import Skip, SkipTest

import pytest
try:
    from unittest.mock import MagicMock
except ImportError:
    from mock import MagicMock

from ..test_helpers import (
    true_func, assert_equal_array_array, make_1d_traj, data_filename,
    assert_items_equal
)

import openpathsampling as paths

from openpathsampling.beta.hooks import *

import logging
logging.getLogger('openpathsampling.initialization').setLevel(logging.CRITICAL)
logging.getLogger('openpathsampling.ensemble').setLevel(logging.CRITICAL)
logging.getLogger('openpathsampling.storage').setLevel(logging.CRITICAL)
logging.getLogger('openpathsampling.netcdfplus').setLevel(logging.CRITICAL)

class TrivialPathSimulator(paths.PathSimulator):
    """
    use this to check that the hooks get attached and run correctly
    """
    def run(self, n_steps):
        self.run_hooks('before_simulation', sim=self)
        state = ""
        hook_state = {}
        for step in range(n_steps):
            self.run_hooks('before_step', sim=self, step_number=step,
                           step_info=(step, n_steps), state=state)
            result = str(step)
            state += result
            hook_state = self.run_hooks('after_step', sim=self,
                                        step_number=step,
                                        step_info=(step, n_steps),
                                        state=state, results=result,
                                        hook_state=hook_state)
        self.run_hooks('after_simulation', sim=self)

class NoStepHookPathSimulator(paths.PathSimulator):
    """use this to check that only correct hooks get assigned"""
    hook_names = ['before_simulation', 'extra_hook']
    def run(self, n_steps):
        self.run_hooks('before_simulation', sim=self)
        self.run_hooks('extra_hook', sim=self, extra="foo")

class StupidHook(PathSimulatorHook):
    def __init__(self):
        self.began_simulation = False
        self.began_steps = 0
        self.finished_steps = 0
        self.finished_simulation = False

    def before_simulation(self, sim):
        self.began_simulation = True

    def before_step(self, sim, step_number, step_info, state):
        self.began_steps +=1

    def after_step(self, sim, step_number, step_info, state, results,
                   hook_state):
        self.finished_steps += 1

    def after_simulation(self, sim):
        self.finished_simulation = True

class StupiderHook(PathSimulatorHook):
    def __init__(self):
        self.foo = None

    def hook_method(self, sim, extra=None):
        self.foo = extra

class TestPathSimulatorHook(object):
    # includes tests for attach_hook and run_hook
    def test_attach_hooks_class(self):
        hook = StupidHook()
        sim = TrivialPathSimulator(storage=None)
        sim.attach_hook(hook)
        for k in sim.hook_names:
            assert_equal(len(sim.hooks[k]), 1)

    def test_attach_hooks_method(self):
        hook_instance = StupiderHook()
        sim = NoStepHookPathSimulator(storage=None)
        sim.attach_hook(hook_instance.hook_method, hook_for='extra_hook')
        assert_not_equal(hook_instance.foo, "foo")
        sim.run(1)
        assert_equal(hook_instance.foo, "foo")

    def test_trivial_simulation_hooks(self):
        sim = TrivialPathSimulator(storage=None)
        stupid = StupidHook()
        stupider = StupiderHook()
        sim.attach_hook(stupid)
        extra = lambda sim: stupider.hook_method(sim, "foo")
        sim.attach_hook(extra, hook_for='after_simulation')
        # before running, check that state is as expected
        assert_equal(stupid.began_simulation, False)
        assert_equal(stupid.finished_simulation, False)
        assert_equal(stupid.began_steps, 0)
        assert_equal(stupid.finished_steps, 0)
        assert_not_equal(stupider.foo, "foo")
        # run
        sim.run(3)
        # check results
        assert_equal(stupid.began_simulation, True)
        assert_equal(stupid.finished_simulation, True)
        assert_equal(stupid.began_steps, 3)
        assert_equal(stupid.finished_steps, 3)
        assert_equal(stupider.foo, "foo")

    @raises(TypeError)
    def test_bad_extra_hook_method(self):
        sim = NoStepHookPathSimulator(storage=None)
        stupid = StupidHook()
        sim.attach_hook(stupid)

    @raises(TypeError)
    def test_bad_hook_name(self):
        sim = NoStepHookPathSimulator(storage=None)
        stupider = StupiderHook()
        sim.attach_hook(stupider.hook_method, "blah")

    def test_no_step_hook_simulation_hooks(self):
        sim = NoStepHookPathSimulator(storage=None)
        stupid = StupidHook()
        stupider = StupiderHook()
        sim.attach_hook(stupid.before_simulation,
                        hook_for='before_simulation')
        sim.attach_hook(stupider.hook_method, hook_for='extra_hook')
        assert_equal(len(sim.hooks), len(sim.hook_names))
        for hook in sim.hooks:
            assert_equal(len(sim.hooks[hook]), 1)

        assert_equal(stupid.began_simulation, False)
        assert_not_equal(stupider.foo, "foo")
        sim.run(1)
        assert_equal(stupid.began_simulation, True)
        assert_equal(stupider.foo, "foo")

class TestStorageHook(object):
    def setup(self):
        self.storage = MagicMock()
        self.empty_hook = StorageHook()
        self.hook = StorageHook(storage=self.storage,
                                frequency=10)
        self.simulation = MagicMock(storage=self.storage,
                                    save_frequency=10)

    @pytest.mark.parametrize('hook', ['empty', 'std'])
    def test_before_simulation(self, hook):
        hook = {'empty': self.empty_hook,
                'std': self.hook}[hook]
        hook.before_simulation(self.simulation)
        assert hook.storage == self.storage
        assert hook.frequency == 10

    @pytest.mark.parametrize('step_num', [0, 5, 10])
    def test_after_step(self, step_num):
        self.hook.after_step(self.simulation, step_num, ('step', 'info'),
                             ('state'), "results", "hook_state")
        self.storage.save.assert_called_once()
        if step_num in [0, 10]:
            self.storage.sync_all.assert_called_once()

    def test_after_simulation(self):
        self.hook.after_simulation(self.simulation)
        self.storage.sync_all.assert_called_once()

class TestShootFromSnapshotsOutputHook(object):
    def setup(self):
        pass

    def test_before_simulation(self):
        raise SkipTest

    def test_before_step(self):
        raise SkipTest
<commit_msg>tests for ShootFromSnapshotsOutputHook<commit_after>from __future__ import absolute_import
from nose.tools import (assert_equal, assert_not_equal, assert_almost_equal,
                        raises)
from nose.plugins.skip import Skip, SkipTest

import io

import pytest
try:
    from unittest.mock import MagicMock
except ImportError:
    from mock import MagicMock

from ..test_helpers import (
    true_func, assert_equal_array_array, make_1d_traj, data_filename,
    assert_items_equal
)

import openpathsampling as paths

from openpathsampling.beta.hooks import *

import logging
logging.getLogger('openpathsampling.initialization').setLevel(logging.CRITICAL)
logging.getLogger('openpathsampling.ensemble').setLevel(logging.CRITICAL)
logging.getLogger('openpathsampling.storage').setLevel(logging.CRITICAL)
logging.getLogger('openpathsampling.netcdfplus').setLevel(logging.CRITICAL)

class TrivialPathSimulator(paths.PathSimulator):
    """
    use this to check that the hooks get attached and run correctly
    """
    def run(self, n_steps):
        self.run_hooks('before_simulation', sim=self)
        state = ""
        hook_state = {}
        for step in range(n_steps):
            self.run_hooks('before_step', sim=self, step_number=step,
                           step_info=(step, n_steps), state=state)
            result = str(step)
            state += result
            hook_state = self.run_hooks('after_step', sim=self,
                                        step_number=step,
                                        step_info=(step, n_steps),
                                        state=state, results=result,
                                        hook_state=hook_state)
        self.run_hooks('after_simulation', sim=self)

class NoStepHookPathSimulator(paths.PathSimulator):
    """use this to check that only correct hooks get assigned"""
    hook_names = ['before_simulation', 'extra_hook']
    def run(self, n_steps):
        self.run_hooks('before_simulation', sim=self)
        self.run_hooks('extra_hook', sim=self, extra="foo")

class StupidHook(PathSimulatorHook):
    def __init__(self):
        self.began_simulation = False
        self.began_steps = 0
        self.finished_steps = 0
        self.finished_simulation = False

    def before_simulation(self, sim):
        self.began_simulation = True

    def before_step(self, sim, step_number, step_info, state):
        self.began_steps +=1

    def after_step(self, sim, step_number, step_info, state, results,
                   hook_state):
        self.finished_steps += 1

    def after_simulation(self, sim):
        self.finished_simulation = True

class StupiderHook(PathSimulatorHook):
    def __init__(self):
        self.foo = None

    def hook_method(self, sim, extra=None):
        self.foo = extra

class TestPathSimulatorHook(object):
    # includes tests for attach_hook and run_hook
    def test_attach_hooks_class(self):
        hook = StupidHook()
        sim = TrivialPathSimulator(storage=None)
        sim.attach_hook(hook)
        for k in sim.hook_names:
            assert_equal(len(sim.hooks[k]), 1)

    def test_attach_hooks_method(self):
        hook_instance = StupiderHook()
        sim = NoStepHookPathSimulator(storage=None)
        sim.attach_hook(hook_instance.hook_method, hook_for='extra_hook')
        assert_not_equal(hook_instance.foo, "foo")
        sim.run(1)
        assert_equal(hook_instance.foo, "foo")

    def test_trivial_simulation_hooks(self):
        sim = TrivialPathSimulator(storage=None)
        stupid = StupidHook()
        stupider = StupiderHook()
        sim.attach_hook(stupid)
        extra = lambda sim: stupider.hook_method(sim, "foo")
        sim.attach_hook(extra, hook_for='after_simulation')
        # before running, check that state is as expected
        assert_equal(stupid.began_simulation, False)
        assert_equal(stupid.finished_simulation, False)
        assert_equal(stupid.began_steps, 0)
        assert_equal(stupid.finished_steps, 0)
        assert_not_equal(stupider.foo, "foo")
        # run
        sim.run(3)
        # check results
        assert_equal(stupid.began_simulation, True)
        assert_equal(stupid.finished_simulation, True)
        assert_equal(stupid.began_steps, 3)
        assert_equal(stupid.finished_steps, 3)
        assert_equal(stupider.foo, "foo")

    @raises(TypeError)
    def test_bad_extra_hook_method(self):
        sim = NoStepHookPathSimulator(storage=None)
        stupid = StupidHook()
        sim.attach_hook(stupid)

    @raises(TypeError)
    def test_bad_hook_name(self):
        sim = NoStepHookPathSimulator(storage=None)
        stupider = StupiderHook()
        sim.attach_hook(stupider.hook_method, "blah")

    def test_no_step_hook_simulation_hooks(self):
        sim = NoStepHookPathSimulator(storage=None)
        stupid = StupidHook()
        stupider = StupiderHook()
        sim.attach_hook(stupid.before_simulation,
                        hook_for='before_simulation')
        sim.attach_hook(stupider.hook_method, hook_for='extra_hook')
        assert_equal(len(sim.hooks), len(sim.hook_names))
        for hook in sim.hooks:
            assert_equal(len(sim.hooks[hook]), 1)

        assert_equal(stupid.began_simulation, False)
        assert_not_equal(stupider.foo, "foo")
        sim.run(1)
        assert_equal(stupid.began_simulation, True)
        assert_equal(stupider.foo, "foo")

class TestStorageHook(object):
    def setup(self):
        self.storage = MagicMock()
        self.empty_hook = StorageHook()
        self.hook = StorageHook(storage=self.storage,
                                frequency=10)
        self.simulation = MagicMock(storage=self.storage,
                                    save_frequency=10)

    @pytest.mark.parametrize('hook', ['empty', 'std'])
    def test_before_simulation(self, hook):
        hook = {'empty': self.empty_hook,
                'std': self.hook}[hook]
        hook.before_simulation(self.simulation)
        assert hook.storage == self.storage
        assert hook.frequency == 10

    @pytest.mark.parametrize('step_num', [0, 5, 10])
    def test_after_step(self, step_num):
        self.hook.after_step(self.simulation, step_num, ('step', 'info'),
                             ('state'), "results", "hook_state")
        self.storage.save.assert_called_once()
        if step_num in [0, 10]:
            self.storage.sync_all.assert_called_once()

    def test_after_simulation(self):
        self.hook.after_simulation(self.simulation)
        self.storage.sync_all.assert_called_once()

class TestShootFromSnapshotsOutputHook(object):
    def setup(self):
        self.stream = io.StringIO()
        self.simulation = MagicMock(output_stream=self.stream,
                                    allow_refresh=False)
        self.empty_hook = ShootFromSnapshotsOutputHook()
        self.hook = ShootFromSnapshotsOutputHook(output_stream=self.stream,
                                                 allow_refresh=False)

    @pytest.mark.parametrize('hook_name', ['empty', 'std'])
    def test_before_simulation(self, hook_name):
        hook = {'empty': self.empty_hook,
                'std': self.hook}[hook_name]
        hook.before_simulation(self.simulation)
        assert hook.output_stream == self.stream
        assert hook.allow_refresh is False

    def test_before_step(self):
        step_info = (0, 100, 1, 10)  # snap_num, n_snap, shot_num, n_shots
        self.hook.before_step(self.simulation, step_number=2,
                              step_info=step_info, state=None)
        contents = self.stream.getvalue()
        assert contents == "Working on snapshot 1 / 100; shot 2 / 10"
<|endoftext|>


<commit_before>import argparse
import sys
import logging
import os
import csv
import gzip

DEBUG=False
NotDEBUG=not DEBUG

parser = argparse.ArgumentParser(description="Convert TOPMed database to annovar format.",
                                 formatter_class=argparse.ArgumentDefaultsHelpFormatter)

parser.add_argument('-i', '--input', action='store', nargs='?', help='Input TOPMed file', required=NotDEBUG)
parser.add_argument('-o', '--output', action='store', nargs='?', help="Output Output file", required=NotDEBUG)

args = parser.parse_args()
if DEBUG:
  args.input = "/scratch/cqs/references/TOPMed/bravo-dbsnp-all.vcf.gz"
  args.output = "/scratch/cqs/references/annovar/humandb/hg38_topmed05.txt"

logger = logging.getLogger('topmed2annovar')
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)-8s - %(message)s')

if args.input.endswith(".gz"):
  fin = gzip.open(args.input, 'rt')
else:
  fin = open(args.input, "rt")

with fin:
  tmpFile = args.output + ".tmp"
  with open(tmpFile, "w") as fout:
    fout.write("#Chr\tStart\tEnd\tRef\tAlt\tTOPMed\n")
    for line in fin:
      if line.startswith('#'):
        continue
      parts = line.rstrip().split("\t")
      chr = parts[0]
      if chr.startswith("chr"):
        chr = chr[3:]
      start = parts[1]
      end = parts[1]
      ref = parts[3]
      alt = parts[4]
      info = parts[7]
      infoparts = info.split(';')
      af = ""
      for ip in infoparts:
        ipparts = ip.split("=")
        if ipparts[0] == "AF":
          af = ipparts[1]
          break
      if af == "":
        continue
      fout.write("%s\t%s\t%s\t%s\t%s\t%s\n" % (chr, start, end, ref, alt, af))

  if os.path.isfile(args.output):
    os.remove(args.output)
  os.rename(tmpFile, args.output)

  realpath = os.path.dirname(os.path.realpath(__file__))
  rPath = realpath + "/annovar_idx.pl"
  cmd = "perl " + rPath + " " + args.output + " 1000 > " + args.output + ".idx"
  logger.info(cmd)
  os.system(cmd)

  logger.info("done.")

<commit_msg>add python3 header<commit_after>#!/usr/bin/env python3

import argparse
import sys
import logging
import os
import csv
import gzip

DEBUG=False
NotDEBUG=not DEBUG

parser = argparse.ArgumentParser(description="Convert TOPMed database to annovar format.",
                                 formatter_class=argparse.ArgumentDefaultsHelpFormatter)

parser.add_argument('-i', '--input', action='store', nargs='?', help='Input TOPMed file', required=NotDEBUG)
parser.add_argument('-o', '--output', action='store', nargs='?', help="Output Output file", required=NotDEBUG)

args = parser.parse_args()
if DEBUG:
  args.input = "/scratch/cqs/references/TOPMed/bravo-dbsnp-all.vcf.gz"
  args.output = "/scratch/cqs/references/annovar/humandb/hg38_topmed05.txt"

logger = logging.getLogger('topmed2annovar')
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)-8s - %(message)s')

if args.input.endswith(".gz"):
  fin = gzip.open(args.input, 'rt')
else:
  fin = open(args.input, "rt")

with fin:
  tmpFile = args.output + ".tmp"
  with open(tmpFile, "w") as fout:
    fout.write("#Chr\tStart\tEnd\tRef\tAlt\tTOPMed\n")
    for line in fin:
      if line.startswith('#'):
        continue
      parts = line.rstrip().split("\t")
      chr = parts[0]
      if chr.startswith("chr"):
        chr = chr[3:]
      start = parts[1]
      end = parts[1]
      ref = parts[3]
      alt = parts[4]
      info = parts[7]
      infoparts = info.split(';')
      af = ""
      for ip in infoparts:
        ipparts = ip.split("=")
        if ipparts[0] == "AF":
          af = ipparts[1]
          break
      if af == "":
        continue
      fout.write("%s\t%s\t%s\t%s\t%s\t%s\n" % (chr, start, end, ref, alt, af))

  if os.path.isfile(args.output):
    os.remove(args.output)
  os.rename(tmpFile, args.output)

  realpath = os.path.dirname(os.path.realpath(__file__))
  rPath = realpath + "/annovar_idx.pl"
  cmd = "perl " + rPath + " " + args.output + " 1000 > " + args.output + ".idx"
  logger.info(cmd)
  os.system(cmd)

  logger.info("done.")

<|endoftext|>


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

"""
Core functionality for the classical PSHA hazard calculator.
"""

import numpy
import os
import random
import StringIO

from django.db import transaction
from nrml import parsers as nrml_parsers

from openquake import engine2
from openquake import writer
from openquake.calculators import base
from openquake.calculators.hazard import general
from openquake.db import models
from openquake.input import logictree
from openquake.input import source
from openquake.job.validation import MIN_SINT_32
from openquake.job.validation import MAX_SINT_32

class ClassicalHazardCalculator(base.CalculatorNext):
    """
    Classical PSHA hazard calculator. Computes hazard curves for a given set of
    points.

    For each realization of the calculation, we randomly sample source models
    and GMPEs (Ground Motion Prediction Equations) from logic trees.
    """

    def initialize_site_model(self):
        """
        If a site model is specified in the calculation configuration. parse
        it and load it into the `hzrdi.site_model` table. This includes a
        validation step to ensure that the area covered by the site model
        completely envelops the calculation geometry. (If this requirement is
        not satisfied, an exception will be raised. See
        :func:`openquake.calculators.hazard.general.validate_site_model`.)

        Then, take all of the points/locations of interest defined by the
        calculation geometry. For each point, do distance queries on the site
        model and get the site parameters which are closest to the point of
        interest. This aggregation of points to the closest site parameters
        is what we store in `htemp.site_data`. (Computing this once prior to
        starting the calculation is optimal, since each task will need to
        consider all sites.)
        """
        hc_id = self.job.hazard_calculation.id

        site_model_inp = general.get_site_model(hc_id)
        if site_model_inp is not None:
            # Explicit cast to `str` here because the XML parser doesn't like
            # unicode. (More specifically, lxml doesn't like unicode.)
            site_model_content = str(site_model_inp.model_content.raw_content)

            # Store `site_model` records:
            general.store_site_model(
                site_model_inp, StringIO.StringIO(site_model_content))

            mesh = self.job.hazard_calculation.points_to_compute()

            # Get the site model records we stored:
            site_model_data = models.SiteModel.objects.filter(
                input=site_model_inp)

            general.validate_site_model(site_model_data, mesh)

            general.store_site_data(hc_id, site_model_inp, mesh)

    def initialize_sources(self):
        """
        Parse and validation logic trees (source and gsim). Then get all
        sources referenced in the the source model logic tree, create
        :class:`~openquake.db.models.Input` records for all of them, parse
        then, and save the parsed sources to the `parsed_source` table
        (see :class:`openquake.db.models.ParsedSource`).
        """
        hc = self.job.hazard_calculation

        [smlt] = models.inputs4hcalc(hc.id, input_type='lt_source')
        [gsimlt] = models.inputs4hcalc(hc.id, input_type='lt_gsim')
        source_paths = logictree.read_logic_trees(
            hc.base_path, smlt.path, gsimlt.path)

        src_inputs = []
        for src_path in source_paths:
            full_path = os.path.join(hc.base_path, src_path)

            # Get or reuse the 'source' Input:
            inp = engine2.get_input(
                full_path, 'source', hc.owner, hc.force_inputs)
            src_inputs.append(inp)

            # Associate the source input to the calculation:
            models.Input2hcalc.objects.get_or_create(
                input=inp, hazard_calculation=hc)

            # Associate the source input to the source model logic tree input:
            models.Src2ltsrc.objects.get_or_create(
                hzrd_src=inp, lt_src=smlt, filename=src_path)

        # Now parse the source models and store `pared_source` records:
        for src_inp in src_inputs:
            src_content = StringIO.StringIO(src_inp.model_content.raw_content)
            sm_parser = nrml_parsers.SourceModelParser(src_content)
            src_db_writer = source.SourceDBWriter(
                src_inp, sm_parser.parse(), hc.rupture_mesh_spacing,
                hc.width_of_mfd_bin, hc.area_source_discretization)
            src_db_writer.serialize()

    # Silencing 'Too many local variables'
    # pylint: disable=R0914
    @transaction.commit_on_success(using='reslt_writer')
    def initialize_realizations(self):
        """
        Create records for the `hzrdr.lt_realization` and
        `htemp.source_progress` records. To do this, we sample the source model
        logic tree to choose a source model for the realization, then we sample
        the GSIM logic tree. We record the logic tree paths for both trees in
        the `lt_realization` record.

        Then we create `htemp.source_progress` records for each source in the
        source model chosen for each realization.
        """
        hc = self.job.hazard_calculation

        # Each realization will have two seeds:
        # One for source model logic tree, one for GSIM logic tree.
        rnd = random.Random()
        seed = hc.random_seed
        rnd.seed(seed)

        [smlt] = models.inputs4hcalc(hc.id, input_type='lt_source')

        ltp = logictree.LogicTreeProcessor(hc.id)

        # The first realization gets the seed we specified in the config file.
        for i in xrange(hc.number_of_logic_tree_samples):
            lt_rlz = models.LtRealization(hazard_calculation=hc)
            lt_rlz.ordinal = i
            lt_rlz.seed = seed

            # Sample source model logic tree branch paths:
            sm_name, _, sm_lt_branch_ids = ltp.sample_source_model_logictree(
                rnd.randint(MIN_SINT_32, MAX_SINT_32))
            lt_rlz.sm_lt_path = sm_lt_branch_ids

            # Sample GSIM logic tree branch paths:
            _, gsim_branch_ids = ltp.sample_gmpe_logictree(
                rnd.randint(MIN_SINT_32, MAX_SINT_32))
            lt_rlz.gsim_lt_path = gsim_branch_ids

            # Get the source model for this sample:
            hzrd_src = models.Src2ltsrc.objects.get(
                lt_src=smlt.id, filename=sm_name).hzrd_src
            parsed_sources = models.ParsedSource.objects.filter(input=hzrd_src)

            lt_rlz.total_sources = len(parsed_sources)
            lt_rlz.save()

            # Create source_progress for this realization
            # A bulk insert is more efficient because there could be lots of
            # of individual sources.
            sp_inserter = writer.BulkInserter(models.SourceProgress)
            for ps in parsed_sources:
                sp_inserter.add_entry(
                    lt_realization_id=lt_rlz.id, parsed_source_id=ps.id)
            sp_inserter.flush()

            # Now stub out the curve result records for this realization:
            self.initialize_hazard_curve_progress(lt_rlz)

            # update the seed for the next realization
            seed = rnd.randint(MIN_SINT_32, MAX_SINT_32)

    def initialize_hazard_curve_progress(self, lt_rlz):
        """
        As a calculation progresses, workers will periodically update the
        intermediate results. These results will be stored in
        `htemp.hazard_curve_progress` until the calculation is completed.

        Before the core calculation begins, we need to initalize these records,
        one data set per IMT. Each dataset will be stored in the database as a
        pickled 2D numpy array (with number of rows == calculation points of
        interest and number of columns == number of IML values for a given
        IMT).

        We will create 1 `hazard_curve_progress` record per IMT per
        realization.

        :param lt_rlz:
            :class:`openquake.db.models.LtRealization` object to associate
            with these inital hazard curve values.
        """
        hc = self.job.hazard_calculation

        num_points = len(hc.points_to_compute())

        im_data = hc.intensity_measure_types_and_levels
        for imt, imls in im_data.items():
            hc_prog = models.HazardCurveProgress()
            hc_prog.lt_realization = lt_rlz
            hc_prog.imt = imt
            hc_prog.result_matrix = numpy.zeros((num_points, len(imls)))
            hc_prog.save()

    def pre_execute(self):
        """
        Do pre-execution work. At the moment, this work entails: parsing and
        initializing sources, parsing and initializing the site model (if there
        is one), and generating logic tree realizations. (The latter piece
        basically defines the work to be done in the `execute` phase.)
        """

        # Parse logic trees and create source Inputs.
        self.initialize_sources()

        # Deal with the site model and compute site data for the calculation
        # (if a site model was specified, that is).
        self.initialize_site_model()

        # Now bootstrap the logic tree realizations and related data.
        # This defines for us the "work" that needs to be done when we reach
        # the `execute` phase.
        # This will also stub out hazard curve result records. Workers will
        # update these periodically with partial results (partial meaning,
        # result curves for just a subset of the overall sources) when some
        # work is complete.
        self.initialize_realizations()
<commit_msg>Added missing update to random generator<commit_after># Copyright (c) 2010-2012, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""
Core functionality for the classical PSHA hazard calculator.
"""

import numpy
import os
import random
import StringIO

from django.db import transaction
from nrml import parsers as nrml_parsers

from openquake import engine2
from openquake import writer
from openquake.calculators import base
from openquake.calculators.hazard import general
from openquake.db import models
from openquake.input import logictree
from openquake.input import source
from openquake.job.validation import MIN_SINT_32
from openquake.job.validation import MAX_SINT_32

class ClassicalHazardCalculator(base.CalculatorNext):
    """
    Classical PSHA hazard calculator. Computes hazard curves for a given set of
    points.

    For each realization of the calculation, we randomly sample source models
    and GMPEs (Ground Motion Prediction Equations) from logic trees.
    """

    def initialize_site_model(self):
        """
        If a site model is specified in the calculation configuration. parse
        it and load it into the `hzrdi.site_model` table. This includes a
        validation step to ensure that the area covered by the site model
        completely envelops the calculation geometry. (If this requirement is
        not satisfied, an exception will be raised. See
        :func:`openquake.calculators.hazard.general.validate_site_model`.)

        Then, take all of the points/locations of interest defined by the
        calculation geometry. For each point, do distance queries on the site
        model and get the site parameters which are closest to the point of
        interest. This aggregation of points to the closest site parameters
        is what we store in `htemp.site_data`. (Computing this once prior to
        starting the calculation is optimal, since each task will need to
        consider all sites.)
        """
        hc_id = self.job.hazard_calculation.id

        site_model_inp = general.get_site_model(hc_id)
        if site_model_inp is not None:
            # Explicit cast to `str` here because the XML parser doesn't like
            # unicode. (More specifically, lxml doesn't like unicode.)
            site_model_content = str(site_model_inp.model_content.raw_content)

            # Store `site_model` records:
            general.store_site_model(
                site_model_inp, StringIO.StringIO(site_model_content))

            mesh = self.job.hazard_calculation.points_to_compute()

            # Get the site model records we stored:
            site_model_data = models.SiteModel.objects.filter(
                input=site_model_inp)

            general.validate_site_model(site_model_data, mesh)

            general.store_site_data(hc_id, site_model_inp, mesh)

    def initialize_sources(self):
        """
        Parse and validation logic trees (source and gsim). Then get all
        sources referenced in the the source model logic tree, create
        :class:`~openquake.db.models.Input` records for all of them, parse
        then, and save the parsed sources to the `parsed_source` table
        (see :class:`openquake.db.models.ParsedSource`).
        """
        hc = self.job.hazard_calculation

        [smlt] = models.inputs4hcalc(hc.id, input_type='lt_source')
        [gsimlt] = models.inputs4hcalc(hc.id, input_type='lt_gsim')
        source_paths = logictree.read_logic_trees(
            hc.base_path, smlt.path, gsimlt.path)

        src_inputs = []
        for src_path in source_paths:
            full_path = os.path.join(hc.base_path, src_path)

            # Get or reuse the 'source' Input:
            inp = engine2.get_input(
                full_path, 'source', hc.owner, hc.force_inputs)
            src_inputs.append(inp)

            # Associate the source input to the calculation:
            models.Input2hcalc.objects.get_or_create(
                input=inp, hazard_calculation=hc)

            # Associate the source input to the source model logic tree input:
            models.Src2ltsrc.objects.get_or_create(
                hzrd_src=inp, lt_src=smlt, filename=src_path)

        # Now parse the source models and store `pared_source` records:
        for src_inp in src_inputs:
            src_content = StringIO.StringIO(src_inp.model_content.raw_content)
            sm_parser = nrml_parsers.SourceModelParser(src_content)
            src_db_writer = source.SourceDBWriter(
                src_inp, sm_parser.parse(), hc.rupture_mesh_spacing,
                hc.width_of_mfd_bin, hc.area_source_discretization)
            src_db_writer.serialize()

    # Silencing 'Too many local variables'
    # pylint: disable=R0914
    @transaction.commit_on_success(using='reslt_writer')
    def initialize_realizations(self):
        """
        Create records for the `hzrdr.lt_realization` and
        `htemp.source_progress` records. To do this, we sample the source model
        logic tree to choose a source model for the realization, then we sample
        the GSIM logic tree. We record the logic tree paths for both trees in
        the `lt_realization` record.

        Then we create `htemp.source_progress` records for each source in the
        source model chosen for each realization.
        """
        hc = self.job.hazard_calculation

        # Each realization will have two seeds:
        # One for source model logic tree, one for GSIM logic tree.
        rnd = random.Random()
        seed = hc.random_seed
        rnd.seed(seed)

        [smlt] = models.inputs4hcalc(hc.id, input_type='lt_source')

        ltp = logictree.LogicTreeProcessor(hc.id)

        # The first realization gets the seed we specified in the config file.
        for i in xrange(hc.number_of_logic_tree_samples):
            lt_rlz = models.LtRealization(hazard_calculation=hc)
            lt_rlz.ordinal = i
            lt_rlz.seed = seed

            # Sample source model logic tree branch paths:
            sm_name, _, sm_lt_branch_ids = ltp.sample_source_model_logictree(
                rnd.randint(MIN_SINT_32, MAX_SINT_32))
            lt_rlz.sm_lt_path = sm_lt_branch_ids

            # Sample GSIM logic tree branch paths:
            _, gsim_branch_ids = ltp.sample_gmpe_logictree(
                rnd.randint(MIN_SINT_32, MAX_SINT_32))
            lt_rlz.gsim_lt_path = gsim_branch_ids

            # Get the source model for this sample:
            hzrd_src = models.Src2ltsrc.objects.get(
                lt_src=smlt.id, filename=sm_name).hzrd_src
            parsed_sources = models.ParsedSource.objects.filter(input=hzrd_src)

            lt_rlz.total_sources = len(parsed_sources)
            lt_rlz.save()

            # Create source_progress for this realization
            # A bulk insert is more efficient because there could be lots of
            # of individual sources.
            sp_inserter = writer.BulkInserter(models.SourceProgress)
            for ps in parsed_sources:
                sp_inserter.add_entry(
                    lt_realization_id=lt_rlz.id, parsed_source_id=ps.id)
            sp_inserter.flush()

            # Now stub out the curve result records for this realization:
            self.initialize_hazard_curve_progress(lt_rlz)

            # update the seed for the next realization
            seed = rnd.randint(MIN_SINT_32, MAX_SINT_32)
            rnd.seed(seed)

    def initialize_hazard_curve_progress(self, lt_rlz):
        """
        As a calculation progresses, workers will periodically update the
        intermediate results. These results will be stored in
        `htemp.hazard_curve_progress` until the calculation is completed.

        Before the core calculation begins, we need to initalize these records,
        one data set per IMT. Each dataset will be stored in the database as a
        pickled 2D numpy array (with number of rows == calculation points of
        interest and number of columns == number of IML values for a given
        IMT).

        We will create 1 `hazard_curve_progress` record per IMT per
        realization.

        :param lt_rlz:
            :class:`openquake.db.models.LtRealization` object to associate
            with these inital hazard curve values.
        """
        hc = self.job.hazard_calculation

        num_points = len(hc.points_to_compute())

        im_data = hc.intensity_measure_types_and_levels
        for imt, imls in im_data.items():
            hc_prog = models.HazardCurveProgress()
            hc_prog.lt_realization = lt_rlz
            hc_prog.imt = imt
            hc_prog.result_matrix = numpy.zeros((num_points, len(imls)))
            hc_prog.save()

    def pre_execute(self):
        """
        Do pre-execution work. At the moment, this work entails: parsing and
        initializing sources, parsing and initializing the site model (if there
        is one), and generating logic tree realizations. (The latter piece
        basically defines the work to be done in the `execute` phase.)
        """

        # Parse logic trees and create source Inputs.
        self.initialize_sources()

        # Deal with the site model and compute site data for the calculation
        # (if a site model was specified, that is).
        self.initialize_site_model()

        # Now bootstrap the logic tree realizations and related data.
        # This defines for us the "work" that needs to be done when we reach
        # the `execute` phase.
        # This will also stub out hazard curve result records. Workers will
        # update these periodically with partial results (partial meaning,
        # result curves for just a subset of the overall sources) when some
        # work is complete.
        self.initialize_realizations()
<|endoftext|>


<commit_before>from django.contrib.auth import get_user_model

from rest_framework import serializers

from .models import Contribution

contributor_cls = get_user_model()

class ContributionCSVSerializer(serializers.ModelSerializer):

    class Meta:
        model = Contribution

    def to_representation(self, obj):
        return {
            'first_name': obj.contributor.first_name,
            'last_name': obj.contributor.last_name,
            'title': obj.content.title,
            'feature_type': obj.content.feature_type,
            'publish_date': obj.content.published,
            'rate': obj.get_pay,
            'payroll_name': '{0} {1}'.format(obj.contributor.first_name, obj.contributor.last_name)
        }
<commit_msg>get payroll_name<commit_after>from django.contrib.auth import get_user_model

from rest_framework import serializers

from .models import Contribution

contributor_cls = get_user_model()

class ContributionCSVSerializer(serializers.ModelSerializer):

    class Meta:
        model = Contribution

    def to_representation(self, obj):
        full_name = obj.contributor.get_full_name()
        data = {
            'first_name': obj.contributor.first_name,
            'last_name': obj.contributor.last_name,
            'title': obj.content.title,
            'feature_type': obj.content.feature_type,
            'publish_date': obj.content.published,
            'rate': obj.get_pay,
            'payroll_name': full_name
        }
        profile = getattr(obj.contributor, 'freelanceprofile', None)
        if profile:
            payroll_name = getattr(profile, 'payroll_name', None)
            if payroll_name:
                data['payroll_name'] = payroll_name
        return data
<|endoftext|>


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

import os
import numpy

from openquake.hazardlib import InvalidFile
from openquake.commonlib.writers import write_csv
from openquake.qa_tests_data.scenario_damage import (
    case_1, case_1c, case_1h, case_2, case_3, case_4, case_4b, case_5, case_5a,
    case_6, case_7, case_8, case_9)
from openquake.calculators.tests import CalculatorTestCase, strip_calc_id
from openquake.calculators.extract import extract
from openquake.calculators.export import export
from openquake.calculators.views import view

aae = numpy.testing.assert_almost_equal

class ScenarioDamageTestCase(CalculatorTestCase):
    def assert_ok(self, pkg, job_ini, exports='csv', kind='dmg'):
        test_dir = os.path.dirname(pkg.__file__)
        out = self.run_calc(test_dir, job_ini, exports=exports)
        got = out[kind + '_by_asset', exports]
        expected_dir = os.path.join(test_dir, 'expected')
        expected = sorted(f for f in os.listdir(expected_dir)
                          if f.endswith(exports) and 'by_taxon' not in f)
        self.assertEqual(len(got), len(expected))
        for fname, actual in zip(expected, got):
            self.assertEqualFiles('expected/%s' % fname, actual)

    def test_case_1(self):
        # test with a single event
        self.assert_ok(case_1, 'job_risk.ini')
        got = view('num_units', self.calc.datastore)
        self.assertEqual('''\
======== =========
taxonomy num_units
======== =========
RC       2,000
RM       4,000
*ALL*    6,000
======== =========''', got)

        # test agg_damages, 1 realization x 3 damage states
        [dmg] = extract(self.calc.datastore, 'agg_damages/structural?'
                        'taxonomy=RC&CRESTA=01.1')
        numpy.testing.assert_almost_equal(
            [1498.0121, 472.96616, 29.021801], dmg, decimal=4)
        # test no intersection
        dmg = extract(self.calc.datastore, 'agg_damages/structural?'
                      'taxonomy=RM&CRESTA=01.1')
        self.assertEqual(dmg.shape, ())

    def test_case_1c(self):
        # this is a case with more hazard sites than exposure sites
        test_dir = os.path.dirname(case_1c.__file__)
        self.run_calc(test_dir, 'job.ini', exports='csv')
        total = extract(self.calc.datastore, 'agg_damages/structural')
        aae([[0.4906653, 0.3249882, 0.0708492, 0.0211334, 0.092364]],
            total)  # shape (R, D) = (1, 5)

        # check extract gmf_data works with a filtered site collection
        gmf_data = dict(extract(self.calc.datastore, 'gmf_data'))
        self.assertEqual(gmf_data['rlz-000'].shape, (1,))

    def test_case_1h(self):
        # test for consequences with a single asset
        self.assert_ok(case_1h, 'job_risk.ini', exports='csv', kind='losses')

    def test_case_2(self):
        self.assert_ok(case_2, 'job_risk.ini')

    def test_case_3(self):
        self.assert_ok(case_3, 'job_risk.ini')

    def test_case_4(self):
        self.assert_ok(case_4, 'job_haz.ini,job_risk.ini')

    def test_case_4b(self):
        self.run_calc(case_4b.__file__, 'job_haz.ini,job_risk.ini')

        [fname] = export(('dmg_by_event', 'csv'), self.calc.datastore)
        self.assertEqualFiles('expected/' + strip_calc_id(fname), fname)

        [fname] = export(('losses_by_event', 'csv'), self.calc.datastore)
        self.assertEqualFiles('expected/' + strip_calc_id(fname), fname)

        fnames = export(('losses_by_asset', 'csv'), self.calc.datastore)
        self.assertEqual(len(fnames), 2)  # one per realization
        for fname in fnames:
            self.assertEqualFiles('expected/' + strip_calc_id(fname), fname)

    def test_wrong_gsim_lt(self):
        with self.assertRaises(InvalidFile) as ctx:
            self.run_calc(os.path.dirname(case_4b.__file__), 'job_err.ini')
        self.assertIn('must contain a single branchset, found 2!',
                      str(ctx.exception))

    def test_case_5(self):
        # this is a test for the rupture filtering
        # NB: the exposure file is imported twice on purpose, to make
        # sure that nothing changes; the case is very tricky since the
        # hazard site collection is filtered by the maximum_distance,
        # there is no region_constraint in hazard and there is in risk
        self.assert_ok(case_5, 'job_haz.ini,job_risk.ini')

    def test_case_5a(self):
        # this is a case with two gsims and one asset
        self.assert_ok(case_5a, 'job_haz.ini,job_risk.ini')
        dmg = extract(self.calc.datastore, 'agg_damages/structural?taxonomy=*')
        tmpname = write_csv(None, dmg)  # shape (T, R, D) == (1, 2, 5)
        self.assertEqualFiles('expected/dmg_by_taxon.csv', tmpname)

    def test_case_6(self):
        # this is a case with 5 assets on the same point
        self.assert_ok(case_6, 'job_h.ini,job_r.ini')
        dmg = extract(self.calc.datastore, 'agg_damages/structural?taxonomy=*')
        tmpname = write_csv(None, dmg)  # shape (T, R, D) == (5, 1, 5)
        self.assertEqualFiles('expected/dmg_by_taxon.csv', tmpname)

    def test_case_7(self):
        # this is a case with three loss types
        self.assert_ok(case_7, 'job_h.ini,job_r.ini', exports='csv')

        # just run the npz export
        [npz] = export(('dmg_by_asset', 'npz'), self.calc.datastore)
        self.assertEqual(strip_calc_id(npz), 'dmg_by_asset.npz')

    def test_case_8(self):
        # case with a shakemap
        self.run_calc(case_8.__file__, 'prejob.ini')
        self.run_calc(case_8.__file__, 'job.ini',
                      hazard_calculation_id=str(self.calc.datastore.calc_id))
        [fname] = export(('dmg_by_event', 'csv'), self.calc.datastore)
        self.assertEqualFiles('expected/dmg_by_event.csv', fname)

    def test_case_9(self):
        # case with volcanic multiperil ASH, LAVA, LAHARS, PYRO
        self.run_calc(case_9.__file__, 'job.ini')
        [fname] = export(('dmg_by_asset', 'csv'), self.calc.datastore)
        self.assertEqualFiles('expected/dmg_by_asset.csv', fname)
        fnames = export(('gmf_data', 'csv'), self.calc.datastore)
        self.assertEqual(len(fnames), 2)  # gmfs and sites, no sigma_epsilon

        [fname] = export(('losses_by_asset', 'csv'), self.calc.datastore)
        self.assertEqualFiles('expected/losses_by_asset.csv', fname)

        w = 'collapsed?kind=rlz-2&tag=name_1&tag=name_2'
        [fname] = export(('aggregate_by/' + w, 'csv'), self.calc.datastore)
        self.assertEqualFiles('expected/lahars_by_name_12.csv', fname)

        w = 'collapsed?kind=rlz-3&tag=name_1&tag=name_2'
        [fname] = export(('aggregate_by/' + w, 'csv'), self.calc.datastore)
        self.assertEqualFiles('expected/pyro_by_name_12.csv', fname)

    def test_case_9_bis(self):
        # case with volcanic lava
        self.run_calc(case_9.__file__, 'job.ini',
                      multi_peril_csv="{'LAVA': 'lava_flow.csv'}")
        w = 'collapsed?kind=rlz-0&tag=name_1&tag=name_2'
        [fname] = export(('aggregate_by/' + w, 'csv'), self.calc.datastore)
        self.assertEqualFiles('expected/lava_by_name_12.csv', fname)

        # check invalid key structura_fragility_file
        with self.assertRaises(ValueError):
            self.run_calc(case_9.__file__, 'job.ini',
                          structura_fragility_file='fragility_model.xml')

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

import os
import numpy

from openquake.hazardlib import InvalidFile
from openquake.commonlib.writers import write_csv
from openquake.qa_tests_data.scenario_damage import (
    case_1, case_1c, case_1h, case_2, case_3, case_4, case_4b, case_5, case_5a,
    case_6, case_7, case_8, case_9)
from openquake.calculators.tests import CalculatorTestCase, strip_calc_id
from openquake.calculators.extract import extract
from openquake.calculators.export import export
from openquake.calculators.views import view

aae = numpy.testing.assert_almost_equal

class ScenarioDamageTestCase(CalculatorTestCase):
    def assert_ok(self, pkg, job_ini, exports='csv', kind='dmg'):
        test_dir = os.path.dirname(pkg.__file__)
        out = self.run_calc(test_dir, job_ini, exports=exports)
        got = out[kind + '_by_asset', exports]
        expected_dir = os.path.join(test_dir, 'expected')
        expected = sorted(f for f in os.listdir(expected_dir)
                          if f.endswith(exports) and 'by_taxon' not in f)
        self.assertEqual(len(got), len(expected))
        for fname, actual in zip(expected, got):
            self.assertEqualFiles('expected/%s' % fname, actual)

    def test_case_1(self):
        # test with a single event
        self.assert_ok(case_1, 'job_risk.ini')
        got = view('num_units', self.calc.datastore)
        self.assertEqual('''\
======== =========
taxonomy num_units
======== =========
RC       2,000
RM       4,000
*ALL*    6,000
======== =========''', got)

        # test agg_damages, 1 realization x 3 damage states
        [dmg] = extract(self.calc.datastore, 'agg_damages/structural?'
                        'taxonomy=RC&CRESTA=01.1')
        numpy.testing.assert_almost_equal(
            [1498.0121, 472.96616, 29.021801], dmg, decimal=4)
        # test no intersection
        dmg = extract(self.calc.datastore, 'agg_damages/structural?'
                      'taxonomy=RM&CRESTA=01.1')
        self.assertEqual(dmg.shape, ())

    def test_case_1c(self):
        # this is a case with more hazard sites than exposure sites
        test_dir = os.path.dirname(case_1c.__file__)
        self.run_calc(test_dir, 'job.ini', exports='csv')
        total = extract(self.calc.datastore, 'agg_damages/structural')
        aae([[0.4906653, 0.3249882, 0.0708492, 0.0211334, 0.092364]],
            total)  # shape (R, D) = (1, 5)

        # check extract gmf_data works with a filtered site collection
        gmf_data = dict(extract(self.calc.datastore, 'gmf_data'))
        self.assertEqual(gmf_data['rlz-000'].shape, (1,))

    def test_case_1h(self):
        # test for consequences with a single asset
        self.assert_ok(case_1h, 'job_risk.ini', exports='csv', kind='losses')

    def test_case_2(self):
        self.assert_ok(case_2, 'job_risk.ini')

    def test_case_3(self):
        self.assert_ok(case_3, 'job_risk.ini')

    def test_case_4(self):
        self.assert_ok(case_4, 'job_haz.ini,job_risk.ini')

    def test_case_4b(self):
        self.run_calc(case_4b.__file__, 'job_haz.ini,job_risk.ini')

        [fname] = export(('dmg_by_event', 'csv'), self.calc.datastore)
        self.assertEqualFiles('expected/' + strip_calc_id(fname), fname)

        [fname] = export(('losses_by_event', 'csv'), self.calc.datastore)
        self.assertEqualFiles('expected/' + strip_calc_id(fname), fname)

        fnames = export(('losses_by_asset', 'csv'), self.calc.datastore)
        self.assertEqual(len(fnames), 2)  # one per realization
        for fname in fnames:
            self.assertEqualFiles('expected/' + strip_calc_id(fname), fname)

    def test_wrong_gsim_lt(self):
        with self.assertRaises(InvalidFile) as ctx:
            self.run_calc(os.path.dirname(case_4b.__file__), 'job_err.ini')
        self.assertIn('must contain a single branchset, found 2!',
                      str(ctx.exception))

    def test_case_5(self):
        # this is a test for the rupture filtering
        # NB: the exposure file is imported twice on purpose, to make
        # sure that nothing changes; the case is very tricky since the
        # hazard site collection is filtered by the maximum_distance,
        # there is no region_constraint in hazard and there is in risk
        self.assert_ok(case_5, 'job_haz.ini,job_risk.ini')

    def test_case_5a(self):
        # this is a case with two gsims and one asset
        self.assert_ok(case_5a, 'job_haz.ini,job_risk.ini')
        dmg = extract(self.calc.datastore, 'agg_damages/structural?taxonomy=*')
        tmpname = write_csv(None, dmg)  # shape (T, R, D) == (1, 2, 5)
        self.assertEqualFiles('expected/dmg_by_taxon.csv', tmpname)

    def test_case_6(self):
        # this is a case with 5 assets on the same point
        self.assert_ok(case_6, 'job_h.ini,job_r.ini')
        dmg = extract(self.calc.datastore, 'agg_damages/structural?taxonomy=*')
        tmpname = write_csv(None, dmg)  # shape (T, R, D) == (5, 1, 5)
        self.assertEqualFiles('expected/dmg_by_taxon.csv', tmpname)

    def test_case_7(self):
        # this is a case with three loss types
        self.assert_ok(case_7, 'job_h.ini,job_r.ini', exports='csv')

        # just run the npz export
        [npz] = export(('dmg_by_asset', 'npz'), self.calc.datastore)
        self.assertEqual(strip_calc_id(npz), 'dmg_by_asset.npz')

    def test_case_8(self):
        # case with a shakemap
        self.run_calc(case_8.__file__, 'prejob.ini')
        self.run_calc(case_8.__file__, 'job.ini',
                      hazard_calculation_id=str(self.calc.datastore.calc_id))
        [fname] = export(('dmg_by_event', 'csv'), self.calc.datastore)
        self.assertEqualFiles('expected/dmg_by_event.csv', fname)

    def test_case_9(self):
        # case with volcanic multiperil ASH, LAVA, LAHARS, PYRO
        self.run_calc(case_9.__file__, 'job.ini')
        [fname] = export(('dmg_by_asset', 'csv'), self.calc.datastore)
        self.assertEqualFiles('expected/dmg_by_asset.csv', fname)
        fnames = export(('gmf_data', 'csv'), self.calc.datastore)
        self.assertEqual(len(fnames), 2)  # gmfs and sites, no sigma_epsilon

        [fname] = export(('losses_by_asset', 'csv'), self.calc.datastore)
        self.assertEqualFiles('expected/losses_by_asset.csv', fname)

        w = 'collapsed?kind=rlz-2&tag=name_1&tag=name_2'
        [fname] = export(('aggregate_by/' + w, 'csv'), self.calc.datastore)
        self.assertEqualFiles('expected/lahars_by_name_12.csv', fname)

        w = 'collapsed?kind=rlz-3&tag=name_1&tag=name_2'
        [fname] = export(('aggregate_by/' + w, 'csv'), self.calc.datastore)
        self.assertEqualFiles('expected/pyro_by_name_12.csv', fname)

    def test_case_9_bis(self):
        # case with volcanic lava
        self.run_calc(case_9.__file__, 'job.ini',
                      multi_peril_csv="{'LAVA': 'lava_flow.csv'}")
        w = 'collapsed?kind=rlz-0&tag=name_1&tag=name_2'
        [fname] = export(('aggregate_by/' + w, 'csv'), self.calc.datastore)
        self.assertEqualFiles('expected/lava_by_name_12.csv', fname)

        # check invalid key structura_fragility_file
        with self.assertRaises(ValueError):
            self.run_calc(case_9.__file__, 'job.ini',
                          structura_fragility_file='fragility_model.xml')

        # check invalid key structura_consequence_file
        with self.assertRaises(ValueError):
            self.run_calc(case_9.__file__, 'job.ini',
                          structura_consequence_file='consequence_model.xml')
<|endoftext|>