<commit_before># Mouse reviewer for Orca
#
# Copyright 2008 Eitan Isaacson
# Copyright 2016 Igalia, S.L.
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library 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
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the
# Free Software Foundation, Inc., Franklin Street, Fifth Floor,
# Boston MA  02110-1301 USA.

"""Mouse review mode."""

__id__        = "$Id$"
__version__   = "$Revision$"
__date__      = "$Date$"
__copyright__ = "Copyright (c) 2008 Eitan Isaacson" \
                "Copyright (c) 2016 Igalia, S.L."
__license__   = "LGPL"

import gi
import math
import pyatspi
import time

from gi.repository import Gdk
try:
    gi.require_version("Wnck", "3.0")
    from gi.repository import Wnck
    _mouseReviewCapable = True
except:
    _mouseReviewCapable = False

from . import debug
from . import event_manager
from . import messages
from . import orca_state
from . import script_manager
from . import settings_manager
from . import speech

_eventManager = event_manager.getManager()
_scriptManager = script_manager.getManager()
_settingsManager = settings_manager.getManager()

class _StringContext:
    """The textual information associated with an _ItemContext."""

    def __init__(self, obj, script=None, string="", start=0, end=0):
        """Initialize the _StringContext.

        Arguments:
        - string: The human-consumable string
        - obj: The accessible object associated with this string
        - start: The start offset with respect to entire text, if one exists
        - end: The end offset with respect to the entire text, if one exists
        - script: The script associated with the accessible object
        """

        self._obj = hash(obj)
        self._script = script
        self._string = string
        self._start = start
        self._end = end

    def __eq__(self, other):
        return other is not None \
            and self._obj == other._obj \
            and self._string == other._string \
            and self._start == other._start \
            and self._end == other._end

    def present(self):
        """Presents this context to the user."""

        if not (self._script and self._string):
            return False

        voice = self._script.speechGenerator.voice(string=self._string)
        string = self._script.utilities.adjustForRepeats(self._string)
        self._script.speakMessage(string, voice=voice, interrupt=False)
        self._script.displayBrailleMessage(self._string, -1)
        return True

class _ItemContext:
    """Holds all the information of the item at a specified point."""

    def __init__(self, x=0, y=0, obj=None, frame=None, script=None):
        """Initialize the _ItemContext.

        Arguments:
        - x: The X coordinate
        - y: The Y coordinate
        - obj: The accessible object of interest at that coordinate
        - frame: The containing accessible object (often a top-level window)
        - script: The script associated with the accessible object
        """

        self._x = x
        self._y = y
        self._obj = obj
        self._frame = frame
        self._script = script
        self._string = self._getStringContext()
        self._time = time.time()

    def __eq__(self, other):
        return other is not None \
            and self._frame == other._frame \
            and self._obj == other._obj \
            and self._string == other._string

    def _treatAsDuplicate(self, prior):
        if self._obj != prior._obj or self._frame != prior._frame:
            return False

        if self._time - prior._time > 0.1:
            return False

        msg = "MOUSE REVIEW: Treating as duplicate"
        debug.println(debug.LEVEL_INFO, msg, True)
        return True

    def _getStringContext(self):
        """Returns the _StringContext associated with the specified point."""

        if not (self._script and self._obj):
            return _StringContext(self._obj)

        interfaces = pyatspi.listInterfaces(self._obj)
        if "Text" not in interfaces:
            return _StringContext(self._obj, self._script)

        state = self._obj.getState()
        if not state.contains(pyatspi.STATE_SELECTABLE):
            boundary = pyatspi.TEXT_BOUNDARY_WORD_START
        else:
            boundary = pyatspi.TEXT_BOUNDARY_LINE_START

        string, start, end = self._script.utilities.textAtPoint(
            self._obj, self._x, self._y, boundary=boundary)
        if not string and self._script.utilities.isTextArea(self._obj):
            string = self._script.speechGenerator.getRoleName(self._obj)

        return _StringContext(self._obj, self._script, string, start, end)

    def present(self, prior):
        """Presents this context to the user."""

        if self == prior or self._treatAsDuplicate(prior):
            return False

        interrupt = self._obj and self._obj != prior._obj \
            or math.sqrt((self._x - prior._x)**2 + (self._y - prior._y)**2) > 25

        if interrupt:
            self._script.presentationInterrupt()

        if self._frame and self._frame != prior._frame:
            self._script.presentObject(self._frame, alreadyFocused=True)

        if self._string != prior._string and self._string.present():
            return True

        if self._obj and self._obj != prior._obj:
            self._script.presentObject(self._obj)

        return True

class MouseReviewer:
    """Main class for the mouse-review feature."""

    def __init__(self):
        self._active = _settingsManager.getSetting("enableMouseReview")
        self._currentMouseOver = _ItemContext()
        self._pointer = None
        self._windows = []
        self._handlerIds = {}

        if not _mouseReviewCapable:
            msg = "MOUSE REVIEW ERROR: Wnck is not available"
            debug.println(debug.LEVEL_INFO, msg, True)
            return

        display = Gdk.Display.get_default()
        try:
            seat = Gdk.Display.get_default_seat(display)
            self._pointer = seat.get_pointer()
        except AttributeError:
            msg = "MOUSE REVIEW ERROR: Gtk+ 3.20 is not available"
            debug.println(debug.LEVEL_INFO, msg, True)
            return

        if not self._active:
            return

        self.activate()

    def _get_listeners(self):
        """Returns the accessible-event listeners for mouse review."""

        return {"mouse:abs": self._listener}

    def activate(self):
        """Activates mouse review."""

        _eventManager.registerModuleListeners(self._get_listeners())
        screen = Wnck.Screen.get_default()
        if screen:
            i = screen.connect("window-stacking-changed", self._on_stacking_changed)
            self._handlerIds[i] = screen

    def deactivate(self):
        """Deactivates mouse review."""

        _eventManager.deregisterModuleListeners(self._get_listeners())
        for key, value in self._handlerIds.items():
            value.disconnect(key)
        self._handlerIds = {}

    def toggle(self, script=None, event=None):
        """Toggle mouse reviewing on or off."""

        if not _mouseReviewCapable:
            return

        self._active = not self._active
        _settingsManager.setSetting("enableMouseReview", self._active)

        if not self._active:
            self.deactivate()
            msg = messages.MOUSE_REVIEW_DISABLED
        else:
            self.activate()
            msg = messages.MOUSE_REVIEW_ENABLED

        if orca_state.activeScript:
            orca_state.activeScript.presentMessage(msg)

    def _on_stacking_changed(self, screen):
        """Callback for Wnck's window-stacking-changed signal."""

        stacked = screen.get_windows_stacked()
        stacked.reverse()
        self._windows = stacked

    def _contains_point(self, obj, x, y, coordType=None):
        if coordType is None:
            coordType = pyatspi.DESKTOP_COORDS

        try:
            return obj.queryComponent().contains(x, y, coordType)
        except:
            return False

    def _has_bounds(self, obj, bounds, coordType=None):
        """Returns True if the bounding box of obj is bounds."""

        if coordType is None:
            coordType = pyatspi.DESKTOP_COORDS

        try:
            extents = obj.queryComponent().getExtents(coordType)
        except:
            return False

        return list(extents) == list(bounds)

    def _accessible_window_at_point(self, pX, pY):
        """Returns the accessible window at the specified coordinates."""

        window = None
        for w in self._windows:
            x, y, width, height = w.get_geometry()
            if x <= pX <= x + width and y <= pY <= y + height:
                window = w
                break

        if not window:
            return None

        app = None
        windowApp = window.get_application()
        if not windowApp:
            return None

        pid = windowApp.get_pid()
        for a in pyatspi.Registry.getDesktop(0):
            if a.get_process_id() == pid:
                app = a
                break

        if not app:
            return None

        candidates = [o for o in app if self._contains_point(o, pX, pY)]
        if len(candidates) == 1:
            return candidates[0]

        name = window.get_name()
        matches = [o for o in candidates if o.name == name]
        if len(matches) == 1:
            return matches[0]

        bbox = window.get_client_window_geometry()
        matches = [o for o in candidates if self._has_bounds(o, bbox)]
        if len(matches) == 1:
            return matches[0]

        return None

    def _on_mouse_moved(self, event):
        """Callback for mouse:abs events."""

        screen, pX, pY = self._pointer.get_position()
        window = self._accessible_window_at_point(pX, pY)
        msg = "MOUSE REVIEW: Window at (%i, %i) is %s" % (pX, pY, window)
        debug.println(debug.LEVEL_INFO, msg, True)
        if not window:
            return

        script = orca_state.activeScript
        if not script:
            return

        obj = script.utilities.descendantAtPoint(window, pX, pY)
        msg = "MOUSE REVIEW: Object at (%i, %i) is %s" % (pX, pY, obj)
        debug.println(debug.LEVEL_INFO, msg, True)

        script = _scriptManager.getScript(window.getApplication(), obj)
        new = _ItemContext(pX, pY, obj, window, script)
        new.present(self._currentMouseOver)
        self._currentMouseOver = new

    def _listener(self, event):
        """Generic listener, mainly to output debugging info."""

        startTime = time.time()
        msg = "\nvvvvv PROCESS OBJECT EVENT %s vvvvv" % event.type
        debug.println(debug.LEVEL_INFO, msg, False)

        if event.type.startswith("mouse:abs"):
            self._on_mouse_moved(event)

        msg = "TOTAL PROCESSING TIME: %.4f\n" % (time.time() - startTime)
        msg += "^^^^^ PROCESS OBJECT EVENT %s ^^^^^\n" % event.type
        debug.println(debug.LEVEL_INFO, msg, False)

reviewer = MouseReviewer()
<commit_msg>Minimize context presented during mouse review<commit_after># Mouse reviewer for Orca
#
# Copyright 2008 Eitan Isaacson
# Copyright 2016 Igalia, S.L.
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library 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
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the
# Free Software Foundation, Inc., Franklin Street, Fifth Floor,
# Boston MA  02110-1301 USA.

"""Mouse review mode."""

__id__        = "$Id$"
__version__   = "$Revision$"
__date__      = "$Date$"
__copyright__ = "Copyright (c) 2008 Eitan Isaacson" \
                "Copyright (c) 2016 Igalia, S.L."
__license__   = "LGPL"

import gi
import math
import pyatspi
import time

from gi.repository import Gdk
try:
    gi.require_version("Wnck", "3.0")
    from gi.repository import Wnck
    _mouseReviewCapable = True
except:
    _mouseReviewCapable = False

from . import debug
from . import event_manager
from . import messages
from . import orca_state
from . import script_manager
from . import settings_manager
from . import speech

_eventManager = event_manager.getManager()
_scriptManager = script_manager.getManager()
_settingsManager = settings_manager.getManager()

class _StringContext:
    """The textual information associated with an _ItemContext."""

    def __init__(self, obj, script=None, string="", start=0, end=0):
        """Initialize the _StringContext.

        Arguments:
        - string: The human-consumable string
        - obj: The accessible object associated with this string
        - start: The start offset with respect to entire text, if one exists
        - end: The end offset with respect to the entire text, if one exists
        - script: The script associated with the accessible object
        """

        self._obj = hash(obj)
        self._script = script
        self._string = string
        self._start = start
        self._end = end

    def __eq__(self, other):
        return other is not None \
            and self._obj == other._obj \
            and self._string == other._string \
            and self._start == other._start \
            and self._end == other._end

    def present(self):
        """Presents this context to the user."""

        if not (self._script and self._string):
            return False

        voice = self._script.speechGenerator.voice(string=self._string)
        string = self._script.utilities.adjustForRepeats(self._string)
        self._script.speakMessage(string, voice=voice, interrupt=False)
        self._script.displayBrailleMessage(self._string, -1)
        return True

class _ItemContext:
    """Holds all the information of the item at a specified point."""

    def __init__(self, x=0, y=0, obj=None, frame=None, script=None):
        """Initialize the _ItemContext.

        Arguments:
        - x: The X coordinate
        - y: The Y coordinate
        - obj: The accessible object of interest at that coordinate
        - frame: The containing accessible object (often a top-level window)
        - script: The script associated with the accessible object
        """

        self._x = x
        self._y = y
        self._obj = obj
        self._frame = frame
        self._script = script
        self._string = self._getStringContext()
        self._time = time.time()

    def __eq__(self, other):
        return other is not None \
            and self._frame == other._frame \
            and self._obj == other._obj \
            and self._string == other._string

    def _treatAsDuplicate(self, prior):
        if self._obj != prior._obj or self._frame != prior._frame:
            return False

        if self._time - prior._time > 0.1:
            return False

        msg = "MOUSE REVIEW: Treating as duplicate"
        debug.println(debug.LEVEL_INFO, msg, True)
        return True

    def _getStringContext(self):
        """Returns the _StringContext associated with the specified point."""

        if not (self._script and self._obj):
            return _StringContext(self._obj)

        interfaces = pyatspi.listInterfaces(self._obj)
        if "Text" not in interfaces:
            return _StringContext(self._obj, self._script)

        state = self._obj.getState()
        if not state.contains(pyatspi.STATE_SELECTABLE):
            boundary = pyatspi.TEXT_BOUNDARY_WORD_START
        else:
            boundary = pyatspi.TEXT_BOUNDARY_LINE_START

        string, start, end = self._script.utilities.textAtPoint(
            self._obj, self._x, self._y, boundary=boundary)
        if not string and self._script.utilities.isTextArea(self._obj):
            string = self._script.speechGenerator.getRoleName(self._obj)

        return _StringContext(self._obj, self._script, string, start, end)

    def _getContainer(self):
        roles = [pyatspi.ROLE_DIALOG,
                 pyatspi.ROLE_FRAME,
                 pyatspi.ROLE_PAGE_TAB,
                 pyatspi.ROLE_WINDOW]
        isContainer = lambda x: x and x.getRole() in roles
        return pyatspi.findAncestor(self._obj, isContainer)

    def present(self, prior):
        """Presents this context to the user."""

        if self == prior or self._treatAsDuplicate(prior):
            return False

        interrupt = self._obj and self._obj != prior._obj \
            or math.sqrt((self._x - prior._x)**2 + (self._y - prior._y)**2) > 25

        if interrupt:
            self._script.presentationInterrupt()

        if self._frame and self._frame != prior._frame:
            self._script.presentObject(self._frame, alreadyFocused=True)

        if self._string != prior._string and self._string.present():
            return True

        if self._obj and self._obj != prior._obj:
            priorObj = prior._obj or self._getContainer()
            self._script.presentObject(self._obj, priorObj=priorObj)

        return True

class MouseReviewer:
    """Main class for the mouse-review feature."""

    def __init__(self):
        self._active = _settingsManager.getSetting("enableMouseReview")
        self._currentMouseOver = _ItemContext()
        self._pointer = None
        self._windows = []
        self._handlerIds = {}

        if not _mouseReviewCapable:
            msg = "MOUSE REVIEW ERROR: Wnck is not available"
            debug.println(debug.LEVEL_INFO, msg, True)
            return

        display = Gdk.Display.get_default()
        try:
            seat = Gdk.Display.get_default_seat(display)
            self._pointer = seat.get_pointer()
        except AttributeError:
            msg = "MOUSE REVIEW ERROR: Gtk+ 3.20 is not available"
            debug.println(debug.LEVEL_INFO, msg, True)
            return

        if not self._active:
            return

        self.activate()

    def _get_listeners(self):
        """Returns the accessible-event listeners for mouse review."""

        return {"mouse:abs": self._listener}

    def activate(self):
        """Activates mouse review."""

        _eventManager.registerModuleListeners(self._get_listeners())
        screen = Wnck.Screen.get_default()
        if screen:
            i = screen.connect("window-stacking-changed", self._on_stacking_changed)
            self._handlerIds[i] = screen

    def deactivate(self):
        """Deactivates mouse review."""

        _eventManager.deregisterModuleListeners(self._get_listeners())
        for key, value in self._handlerIds.items():
            value.disconnect(key)
        self._handlerIds = {}

    def toggle(self, script=None, event=None):
        """Toggle mouse reviewing on or off."""

        if not _mouseReviewCapable:
            return

        self._active = not self._active
        _settingsManager.setSetting("enableMouseReview", self._active)

        if not self._active:
            self.deactivate()
            msg = messages.MOUSE_REVIEW_DISABLED
        else:
            self.activate()
            msg = messages.MOUSE_REVIEW_ENABLED

        if orca_state.activeScript:
            orca_state.activeScript.presentMessage(msg)

    def _on_stacking_changed(self, screen):
        """Callback for Wnck's window-stacking-changed signal."""

        stacked = screen.get_windows_stacked()
        stacked.reverse()
        self._windows = stacked

    def _contains_point(self, obj, x, y, coordType=None):
        if coordType is None:
            coordType = pyatspi.DESKTOP_COORDS

        try:
            return obj.queryComponent().contains(x, y, coordType)
        except:
            return False

    def _has_bounds(self, obj, bounds, coordType=None):
        """Returns True if the bounding box of obj is bounds."""

        if coordType is None:
            coordType = pyatspi.DESKTOP_COORDS

        try:
            extents = obj.queryComponent().getExtents(coordType)
        except:
            return False

        return list(extents) == list(bounds)

    def _accessible_window_at_point(self, pX, pY):
        """Returns the accessible window at the specified coordinates."""

        window = None
        for w in self._windows:
            x, y, width, height = w.get_geometry()
            if x <= pX <= x + width and y <= pY <= y + height:
                window = w
                break

        if not window:
            return None

        app = None
        windowApp = window.get_application()
        if not windowApp:
            return None

        pid = windowApp.get_pid()
        for a in pyatspi.Registry.getDesktop(0):
            if a.get_process_id() == pid:
                app = a
                break

        if not app:
            return None

        candidates = [o for o in app if self._contains_point(o, pX, pY)]
        if len(candidates) == 1:
            return candidates[0]

        name = window.get_name()
        matches = [o for o in candidates if o.name == name]
        if len(matches) == 1:
            return matches[0]

        bbox = window.get_client_window_geometry()
        matches = [o for o in candidates if self._has_bounds(o, bbox)]
        if len(matches) == 1:
            return matches[0]

        return None

    def _on_mouse_moved(self, event):
        """Callback for mouse:abs events."""

        screen, pX, pY = self._pointer.get_position()
        window = self._accessible_window_at_point(pX, pY)
        msg = "MOUSE REVIEW: Window at (%i, %i) is %s" % (pX, pY, window)
        debug.println(debug.LEVEL_INFO, msg, True)
        if not window:
            return

        script = orca_state.activeScript
        if not script:
            return

        obj = script.utilities.descendantAtPoint(window, pX, pY)
        msg = "MOUSE REVIEW: Object at (%i, %i) is %s" % (pX, pY, obj)
        debug.println(debug.LEVEL_INFO, msg, True)

        script = _scriptManager.getScript(window.getApplication(), obj)
        new = _ItemContext(pX, pY, obj, window, script)
        new.present(self._currentMouseOver)
        self._currentMouseOver = new

    def _listener(self, event):
        """Generic listener, mainly to output debugging info."""

        startTime = time.time()
        msg = "\nvvvvv PROCESS OBJECT EVENT %s vvvvv" % event.type
        debug.println(debug.LEVEL_INFO, msg, False)

        if event.type.startswith("mouse:abs"):
            self._on_mouse_moved(event)

        msg = "TOTAL PROCESSING TIME: %.4f\n" % (time.time() - startTime)
        msg += "^^^^^ PROCESS OBJECT EVENT %s ^^^^^\n" % event.type
        debug.println(debug.LEVEL_INFO, msg, False)

reviewer = MouseReviewer()
<|endoftext|>


<commit_before>from threading import Event

import pytest
from mock import patch, MagicMock

from fuse import FuseOSError
from gitfs.views.current import CurrentView

class TestCurrentView(object):
    def test_rename(self):
        mocked_re = MagicMock()
        mocked_index = MagicMock()
        mocked_os = MagicMock()
        mocked_result = MagicMock()

        mocked_result.rename.return_value = True
        mocked_re.sub.return_value = "new"
        mocked_os.path.split.return_value = [1, 1]

        with patch.multiple('gitfs.views.current', re=mocked_re,
                            os=mocked_os):
            from gitfs.views import current
            current.PassthroughView.rename = lambda self, old, new: True

            current = CurrentView(regex="regex", repo_path="repo_path",
                                  read_only=Event(), want_to_merge=Event())
            current._index = mocked_index

            result = current.rename("old", "new")
            assert result is True
            mocked_index.assert_called_once_with(**{
                'remove': 1,
                'add': "new",
                "message": "Rename old to new"
            })
            mocked_os.path.split.assert_called_once_with("old")

    def test_rename_in_git_dir(self):
        current = CurrentView(repo_path="repo",
                              read_only=Event(), want_to_merge=Event())
        with pytest.raises(FuseOSError):
            current.rename(".git/", ".git/")

    def test_symlink(self):
        mocked_index = MagicMock()
        mocked_full_path = MagicMock()

        mocked_full_path.return_value = "full_path"

        with patch('gitfs.views.current.os') as mocked_os:
            mocked_os.symlink.return_value = "done"

            current = CurrentView(repo_path="repo",
                                  read_only=Event(), want_to_merge=Event())
            current._index = mocked_index
            current._full_path = mocked_full_path

            assert current.symlink("name", "target") == "done"
            mocked_os.symlink.assert_called_once_with("target", "full_path")
            mocked_full_path.assert_called_once_with("name")
            message = "Create symlink to target for name"
            mocked_index.assert_called_once_with(add="name", message=message)

    def test_readlink(self):
        mocked_full_path = MagicMock()
        mocked_full_path.return_value = "full path"

        with patch('gitfs.views.current.os') as mocked_os:
            mocked_os.readlink.return_value = "done"

            current = CurrentView(repo_path="repo",
                                  read_only=Event(), want_to_merge=Event())
            current._full_path = mocked_full_path

            assert current.readlink("path") == "done"

    def test_readlink_in_git_repo(self):
        with pytest.raises(FuseOSError):
            current = CurrentView(repo_path="repo",
                                  read_only=Event(), want_to_merge=Event())
            current.readlink(".git/")

    def test_getattr(self):
        mocked_full = MagicMock()
        mocked_os = MagicMock()
        mocked_stat = MagicMock()

        mocked_stat.simple = "stat"
        mocked_os.lstat.return_value = mocked_stat
        mocked_full.return_value = "full_path"

        with patch.multiple('gitfs.views.current', os=mocked_os,
                            STATS=['simple']):
            current = CurrentView(repo_path="repo", uid=1, gid=1)
            current._full_path = mocked_full

            result = current.getattr("path")
            asserted_result = {
                'st_uid': 1,
                'st_gid': 1,
                'simple': "stat"
            }
            assert result == asserted_result

            mocked_os.lstat.assert_called_once_with("full_path")
            mocked_full.assert_called_once_with("path")
<commit_msg>test: getattr from .git dir<commit_after>from threading import Event

import pytest
from mock import patch, MagicMock

from fuse import FuseOSError
from gitfs.views.current import CurrentView

class TestCurrentView(object):
    def test_rename(self):
        mocked_re = MagicMock()
        mocked_index = MagicMock()
        mocked_os = MagicMock()
        mocked_result = MagicMock()

        mocked_result.rename.return_value = True
        mocked_re.sub.return_value = "new"
        mocked_os.path.split.return_value = [1, 1]

        with patch.multiple('gitfs.views.current', re=mocked_re,
                            os=mocked_os):
            from gitfs.views import current
            current.PassthroughView.rename = lambda self, old, new: True

            current = CurrentView(regex="regex", repo_path="repo_path",
                                  read_only=Event(), want_to_merge=Event())
            current._index = mocked_index

            result = current.rename("old", "new")
            assert result is True
            mocked_index.assert_called_once_with(**{
                'remove': 1,
                'add': "new",
                "message": "Rename old to new"
            })
            mocked_os.path.split.assert_called_once_with("old")

    def test_rename_in_git_dir(self):
        current = CurrentView(repo_path="repo",
                              read_only=Event(), want_to_merge=Event())
        with pytest.raises(FuseOSError):
            current.rename(".git/", ".git/")

    def test_symlink(self):
        mocked_index = MagicMock()
        mocked_full_path = MagicMock()

        mocked_full_path.return_value = "full_path"

        with patch('gitfs.views.current.os') as mocked_os:
            mocked_os.symlink.return_value = "done"

            current = CurrentView(repo_path="repo",
                                  read_only=Event(), want_to_merge=Event())
            current._index = mocked_index
            current._full_path = mocked_full_path

            assert current.symlink("name", "target") == "done"
            mocked_os.symlink.assert_called_once_with("target", "full_path")
            mocked_full_path.assert_called_once_with("name")
            message = "Create symlink to target for name"
            mocked_index.assert_called_once_with(add="name", message=message)

    def test_readlink(self):
        mocked_full_path = MagicMock()
        mocked_full_path.return_value = "full path"

        with patch('gitfs.views.current.os') as mocked_os:
            mocked_os.readlink.return_value = "done"

            current = CurrentView(repo_path="repo",
                                  read_only=Event(), want_to_merge=Event())
            current._full_path = mocked_full_path

            assert current.readlink("path") == "done"

    def test_readlink_in_git_repo(self):
        with pytest.raises(FuseOSError):
            current = CurrentView(repo_path="repo",
                                  read_only=Event(), want_to_merge=Event())
            current.readlink(".git/")

    def test_getattr(self):
        mocked_full = MagicMock()
        mocked_os = MagicMock()
        mocked_stat = MagicMock()

        mocked_stat.simple = "stat"
        mocked_os.lstat.return_value = mocked_stat
        mocked_full.return_value = "full_path"

        with patch.multiple('gitfs.views.current', os=mocked_os,
                            STATS=['simple']):
            current = CurrentView(repo_path="repo", uid=1, gid=1)
            current._full_path = mocked_full

            result = current.getattr("path")
            asserted_result = {
                'st_uid': 1,
                'st_gid': 1,
                'simple': "stat"
            }
            assert result == asserted_result

            mocked_os.lstat.assert_called_once_with("full_path")
            mocked_full.assert_called_once_with("path")

    def test_getattr_in_git_dir(self):
        with pytest.raises(FuseOSError):
            current = CurrentView(repo_path="repo", uid=1, gid=1)
            current.getattr(".git/index")
<|endoftext|>


<commit_before>#!/usr/bin/python
"""
Ssh-key key handlers and associated routines
"""
import contextlib
import logging
import paramiko
import re

from cStringIO import StringIO
from teuthology import contextutil
import teuthology.misc as misc
from ..orchestra import run

log = logging.getLogger(__name__)
ssh_keys_user = 'ssh-keys-user'

def generate_keys():
    """
    Generatees a public and private key
    """
    key = paramiko.RSAKey.generate(2048)
    privateString = StringIO()
    key.write_private_key(privateString)
    return key.get_base64(), privateString.getvalue()

def particular_ssh_key_test(line_to_test, ssh_key):
    """
    Check the validity of the ssh_key
    """
    match = re.match('[\w-]+ {key} \S+@\S+'.format(key=re.escape(ssh_key)), line_to_test)

    if match:
        return False
    else:
        return True

def ssh_keys_user_line_test(line_to_test, username ):
    """
    Check the validity of the username
    """
    match = re.match('[\w-]+ \S+ {username}@\S+'.format(username=username), line_to_test)

    if match:
        return False
    else:
        return True

def cleanup_added_key(ctx):
    """
    Delete the keys and removes ~/.ssh/authorized_keys2 entries we added
    """
    log.info('cleaning up keys added for testing')

    for remote in ctx.cluster.remotes:
        username, hostname = str(remote).split('@')
        if "" == username or "" == hostname:
            continue
        else:
            log.info('  cleaning up keys for user {user} on {host}'.format(host=hostname, user=username))
            misc.delete_file(remote, '/home/{user}/.ssh/id_rsa'.format(user=username))
            misc.delete_file(remote, '/home/{user}/.ssh/id_rsa.pub'.format(user=username))
            misc.delete_file(remote, '/home/{user}/.ssh/authorized_keys2'.format(user=username))

@contextlib.contextmanager
def tweak_ssh_config(ctx, config):
    """
    Turn off StrictHostKeyChecking
    """
    run.wait(
        ctx.cluster.run(
            args=[
                'echo',
                'StrictHostKeyChecking no\n',
                run.Raw('>'),
                run.Raw('/home/ubuntu/.ssh/config'),
                run.Raw('&&'),
                'echo',
                'UserKnownHostsFile ',
                run.Raw('/dev/null'),
                run.Raw('>>'),
                run.Raw('/home/ubuntu/.ssh/config'),
                run.Raw('&&'),
                run.Raw('chmod 600 /home/ubuntu/.ssh/config'),
            ],
            wait=False,
        )
    )

    try:
        yield

    finally:
        run.wait(
            ctx.cluster.run(
                args=['rm',run.Raw('/home/ubuntu/.ssh/config')],
            wait=False
            ),
        )

@contextlib.contextmanager
def push_keys_to_host(ctx, config, public_key, private_key):
    """
    Push keys to all hosts
    """
    log.info('generated public key {pub_key}'.format(pub_key=public_key))

    # add an entry for all hosts in ctx to auth_keys_data
    auth_keys_data = ''

    for inner_host in ctx.cluster.remotes.iterkeys():
        inner_username, inner_hostname = str(inner_host).split('@')
        # create a 'user@hostname' string using our fake hostname
        fake_hostname = '{user}@{host}'.format(user=ssh_keys_user, host=str(inner_hostname))
        auth_keys_data += '\nssh-rsa {pub_key} {user_host}\n'.format(pub_key=public_key, user_host=fake_hostname)

    # for each host in ctx, add keys for all other hosts
    for remote in ctx.cluster.remotes:
        username, hostname = str(remote).split('@')
        if "" == username or "" == hostname:
            continue
        else:
            log.info('pushing keys to {host} for {user}'.format(host=hostname, user=username))

            # adding a private key
            priv_key_file = '/home/{user}/.ssh/id_rsa'.format(user=username)
            priv_key_data = '{priv_key}'.format(priv_key=private_key)
            misc.delete_file(remote, priv_key_file, force=True)
            # Hadoop requires that .ssh/id_rsa have permissions of '500'
            misc.create_file(remote, priv_key_file, priv_key_data, str(500))

            # then a private key
            pub_key_file = '/home/{user}/.ssh/id_rsa.pub'.format(user=username)
            pub_key_data = 'ssh-rsa {pub_key} {user_host}'.format(pub_key=public_key, user_host=str(remote))
            misc.delete_file(remote, pub_key_file, force=True)
            misc.create_file(remote, pub_key_file, pub_key_data)

            # adding appropriate entries to the authorized_keys2 file for this host
            auth_keys_file = '/home/{user}/.ssh/authorized_keys2'.format(user=username)

            # now add the list of keys for hosts in ctx to ~/.ssh/authorized_keys2
            misc.create_file(remote, auth_keys_file, auth_keys_data, str(600))

    try:
        yield

    finally:
        # cleanup the keys
        log.info("Cleaning up SSH keys")
        cleanup_added_key(ctx)

@contextlib.contextmanager
def task(ctx, config):
    """
    Creates a set of RSA keys, distributes the same key pair
    to all hosts listed in ctx.cluster, and adds all hosts
    to all others authorized_keys list.

    During cleanup it will delete .ssh/id_rsa, .ssh/id_rsa.pub
    and remove the entries in .ssh/authorized_keys while leaving
    pre-existing entries in place.
    """

    if config is None:
        config = {}
    assert isinstance(config, dict), \
        "task hadoop only supports a dictionary for configuration"

    # this does not need to do cleanup and does not depend on
    # ctx, so I'm keeping it outside of the nested calls
    public_key_string, private_key_string = generate_keys()

    with contextutil.nested(
        lambda: tweak_ssh_config(ctx, config),
        lambda: push_keys_to_host(ctx, config, public_key_string, private_key_string),
        #lambda: tweak_ssh_config(ctx, config),
        ):
        yield

<commit_msg>task.ssh_keys(): manage authorized_keys, not keys2<commit_after>#!/usr/bin/python
"""
Ssh-key key handlers and associated routines
"""
import contextlib
import logging
import paramiko
import re

from cStringIO import StringIO
from teuthology import contextutil
import teuthology.misc as misc
from ..orchestra import run

log = logging.getLogger(__name__)
ssh_keys_user = 'ssh-keys-user'

def generate_keys():
    """
    Generatees a public and private key
    """
    key = paramiko.RSAKey.generate(2048)
    privateString = StringIO()
    key.write_private_key(privateString)
    return key.get_base64(), privateString.getvalue()

def particular_ssh_key_test(line_to_test, ssh_key):
    """
    Check the validity of the ssh_key
    """
    match = re.match('[\w-]+ {key} \S+@\S+'.format(key=re.escape(ssh_key)), line_to_test)

    if match:
        return False
    else:
        return True

def ssh_keys_user_line_test(line_to_test, username ):
    """
    Check the validity of the username
    """
    match = re.match('[\w-]+ \S+ {username}@\S+'.format(username=username), line_to_test)

    if match:
        return False
    else:
        return True

def cleanup_added_key(ctx):
    """
    Delete the keys and removes ~/.ssh/authorized_keys entries we added
    """
    log.info('cleaning up keys added for testing')

    for remote in ctx.cluster.remotes:
        username, hostname = str(remote).split('@')
        if "" == username or "" == hostname:
            continue
        else:
            log.info('  cleaning up keys for user {user} on {host}'.format(host=hostname, user=username))
            misc.delete_file(remote, '/home/{user}/.ssh/id_rsa'.format(user=username))
            misc.delete_file(remote, '/home/{user}/.ssh/id_rsa.pub'.format(user=username))
            cmd = "sed -i /#TEUTHOLOGY_START/,/#TEUTHOLOGY_END/d /home/{user}/.ssh/authorized_keys".format(user=username).split()
            remote.run(args=cmd)

@contextlib.contextmanager
def tweak_ssh_config(ctx, config):
    """
    Turn off StrictHostKeyChecking
    """
    run.wait(
        ctx.cluster.run(
            args=[
                'echo',
                'StrictHostKeyChecking no\n',
                run.Raw('>'),
                run.Raw('/home/ubuntu/.ssh/config'),
                run.Raw('&&'),
                'echo',
                'UserKnownHostsFile ',
                run.Raw('/dev/null'),
                run.Raw('>>'),
                run.Raw('/home/ubuntu/.ssh/config'),
                run.Raw('&&'),
                run.Raw('chmod 600 /home/ubuntu/.ssh/config'),
            ],
            wait=False,
        )
    )

    try:
        yield

    finally:
        run.wait(
            ctx.cluster.run(
                args=['rm',run.Raw('/home/ubuntu/.ssh/config')],
            wait=False
            ),
        )

@contextlib.contextmanager
def push_keys_to_host(ctx, config, public_key, private_key):
    """
    Push keys to all hosts
    """
    log.info('generated public key {pub_key}'.format(pub_key=public_key))

    # add an entry for all hosts in ctx to auth_keys_data
    auth_keys_data = ''

    for inner_host in ctx.cluster.remotes.iterkeys():
        inner_username, inner_hostname = str(inner_host).split('@')
        # create a 'user@hostname' string using our fake hostname
        fake_hostname = '{user}@{host}'.format(user=ssh_keys_user, host=str(inner_hostname))
        auth_keys_data += '\nssh-rsa {pub_key} {user_host}\n'.format(pub_key=public_key, user_host=fake_hostname)

    # for each host in ctx, add keys for all other hosts
    for remote in ctx.cluster.remotes:
        username, hostname = str(remote).split('@')
        if "" == username or "" == hostname:
            continue
        else:
            log.info('pushing keys to {host} for {user}'.format(host=hostname, user=username))

            # adding a private key
            priv_key_file = '/home/{user}/.ssh/id_rsa'.format(user=username)
            priv_key_data = '{priv_key}'.format(priv_key=private_key)
            misc.delete_file(remote, priv_key_file, force=True)
            # Hadoop requires that .ssh/id_rsa have permissions of '500'
            misc.create_file(remote, priv_key_file, priv_key_data, str(500))

            # then a private key
            pub_key_file = '/home/{user}/.ssh/id_rsa.pub'.format(user=username)
            pub_key_data = 'ssh-rsa {pub_key} {user_host}'.format(pub_key=public_key, user_host=str(remote))
            misc.delete_file(remote, pub_key_file, force=True)
            misc.create_file(remote, pub_key_file, pub_key_data)

            # add appropriate entries to the authorized_keys file for this host
            auth_keys_file = '/home/{user}/.ssh/authorized_keys'.format(user=username)
            lines = '#TEUTHOLOGY_START\n' + auth_keys_data + '\n#TEUTHOLOGY_END\n'
            misc.append_lines_to_file(remote, auth_keys_file, lines)

    try:
        yield

    finally:
        # cleanup the keys
        log.info("Cleaning up SSH keys")
        cleanup_added_key(ctx)

@contextlib.contextmanager
def task(ctx, config):
    """
    Creates a set of RSA keys, distributes the same key pair
    to all hosts listed in ctx.cluster, and adds all hosts
    to all others authorized_keys list.

    During cleanup it will delete .ssh/id_rsa, .ssh/id_rsa.pub
    and remove the entries in .ssh/authorized_keys while leaving
    pre-existing entries in place.
    """

    if config is None:
        config = {}
    assert isinstance(config, dict), \
        "task hadoop only supports a dictionary for configuration"

    # this does not need to do cleanup and does not depend on
    # ctx, so I'm keeping it outside of the nested calls
    public_key_string, private_key_string = generate_keys()

    with contextutil.nested(
        lambda: tweak_ssh_config(ctx, config),
        lambda: push_keys_to_host(ctx, config, public_key_string, private_key_string),
        #lambda: tweak_ssh_config(ctx, config),
        ):
        yield

<|endoftext|>


<commit_before># Standard library
from collections import OrderedDict
import copy
import warnings

# Third-party
import astropy.units as u
from astropy.table import Table
from astropy.time import Time
import numpy as np
from twobody import KeplerOrbit, PolynomialRVTrend

# Package
from ..utils import quantity_to_hdf5, quantity_from_hdf5

__all__ = ['JokerSamples']

class JokerSamples(OrderedDict):
    _valid_keys = ['P', 'M0', 'e', 'omega', 'jitter', 'K']

    def __init__(self, t0=None, poly_trend=1, **kwargs):
        """A dictionary-like object for storing posterior samples from
        The Joker, with some extra functionality.

        Parameters
        ----------
        t0 : `astropy.time.Time`, numeric (optional)
            The reference time for the orbital parameters.
        poly_trend : int, optional
            If specified, sample over a polynomial velocity trend with the
            specified number of coefficients. For example, ``poly_trend=3`` will
            sample over parameters of a long-term quadratic velocity trend.
            Default is 1, just a constant velocity shift.
        **kwargs
            These are the orbital element names.
        """

        # initialize empty dictionary
        super(JokerSamples, self).__init__()
        self._setup(poly_trend, t0)

        for key, val in kwargs.items():
            self[key] = val  # calls __setitem__ below

    def _setup(self, poly_trend, t0):
        # reference time
        self.t0 = t0

        self._size = None
        self._shape = None

        self._cache = dict()

        self.poly_trend = int(poly_trend)
        self._trend_names = ['v{0}'.format(i)
                             for i in range(self.poly_trend)]
        for name in self._trend_names:
            if name not in self._valid_keys:
                self._valid_keys.append(name)

    def __getstate__(self):
        return (self.poly_trend, self.t0, dict(self))

    def __setstate__(self, state):
        poly_trend, t0, data = state
        self._setup(poly_trend, t0)
        self.update(data)

    def __reduce__(self):
        return (JokerSamples, (), self.__getstate__())

    def _validate_key(self, key):
        if key not in self._valid_keys:
            raise ValueError("Invalid key '{0}'.".format(key))

    def _validate_val(self, val):
        val = u.Quantity(val)
        if self._shape is not None and val.shape != self.shape:
            raise ValueError("Shape of new samples must match those already "
                             "stored! ({0}, expected {1})"
                             .format(len(val), self.shape))

        return val

    def __getitem__(self, slc):
        if isinstance(slc, str):
            return super(JokerSamples, self).__getitem__(slc)

        else:
            new = copy.copy(self)
            new._size = None  # reset number of samples
            new._shape = None  # reset number of samples

            for k in self.keys():
                new[k] = self[k][slc]

            return new

    def __setitem__(self, key, val):
        self._validate_key(key)
        val = self._validate_val(val)

        if self._shape is None:
            self._shape = val.shape
            self._size = val.size

        super(JokerSamples, self).__setitem__(key, val)

    @property
    def n_samples(self):
        warnings.warn(".n_samples is deprecated in favor of .size",
                      DeprecationWarning)
        return self.size

    @property
    def size(self):
        if self._size is None:
            raise ValueError("No samples stored!")
        return self._size

    @property
    def shape(self):
        if self._shape is None:
            raise ValueError("No samples stored!")
        return self._shape

    def __len__(self):
        return self.n_samples

    def __str__(self):
        return ("<JokerSamples in [{0}], {1} samples>"
                .format(','.join(self.keys()), len(self)))

    @classmethod
    def from_hdf5(cls, f, n=None, **kwargs):
        """
        Parameters
        ----------
        f : :class:`h5py.File`, :class:`h5py.Group`
        n : int (optional)
            The number of samples to load.
        **kwargs
            All other keyword arguments are passed to the class initializer.
        """

        if 't0_bmjd' in f.attrs:
            # Read the reference time:
            t0 = Time(f.attrs['t0_bmjd'], format='mjd', scale='tcb')
        else:
            t0 = None

        samples = cls(t0=t0, **kwargs)
        for key in cls._valid_keys:
            if key in f:
                samples[key] = quantity_from_hdf5(f, key, n=n)

        return samples

    def to_hdf5(self, f):
        """
        Parameters
        ----------
        f : :class:`h5py.File`, :class:`h5py.Group`
        """

        for key in self.keys():
            quantity_to_hdf5(f, key, self[key])

        if self.t0 is not None:
            f.attrs['t0_bmjd'] = self.t0.tcb.mjd

    @classmethod
    def from_table(cls, tbl_or_f):
        """Read a samples object from an Astropy table.

        Parameters
        ----------
        tbl_or_f : `~astropy.table.Table`, str
            Either a table instance or a string filename to be read with
            `astropy.table.Table.read()`.
        """
        if isinstance(tbl_or_f, str):
            tbl_or_f = Table.read(tbl_or_f)

        kwargs = dict()
        kwargs['poly_trend'] = tbl_or_f.meta.get('poly_trend', 1)
        if 't0_bmjd'.upper() in tbl_or_f.meta:
            kwargs['t0'] = Time(tbl_or_f.meta['t0_bmjd'.upper()], format='mjd',
                                scale='tcb')

        samples = cls(**kwargs)
        for key in cls._valid_keys:
            if key in tbl_or_f.colnames:
                samples[key] = u.Quantity(tbl_or_f[key])

        return samples

    def to_table(self):
        """Convert the samples to an Astropy table object.

        Returns
        -------
        tbl : `~astropy.table.Table`
        """
        tbl = Table()
        for k in self.keys():
            tbl[k] = self[k]

        tbl.meta['poly_trend'] = self.poly_trend

        if self.t0 is not None:
            tbl.meta['t0_bmjd'] = self.t0.tcb.mjd

        return tbl

    ##########################################################################
    # Interaction with TwoBody

    def get_orbit(self, index=None, **kwargs):
        """Get a `twobody.KeplerOrbit` object for the samples at the specified
        index.

        Parameters
        ----------
        index : int (optional)
            The index of the samples to turn into a `twobody.KeplerOrbit`
            instance. If the samples object is scalar, no index is necessary.
        **kwargs
            Other keyword arguments are passed to the `twobody.KeplerOrbit`
            initializer. For example, you can specify the inclination by passing
            ``i=...`, or  longitude of the ascending node by passing
            ``Omega=...``.

        Returns
        -------
        orbit : `twobody.KeplerOrbit`
            The samples converted to an orbit object. The barycenter position
            and distance are set to arbitrary values.
        """
        if 'orbit' not in self._cache:
            self._cache['orbit'] = KeplerOrbit(P=1*u.yr, e=0., omega=0*u.deg,
                                               Omega=0*u.deg, i=90*u.deg,
                                               a=1*u.au, t0=self.t0)

        # all of this to avoid the __init__ of KeplerOrbit / KeplerElements
        orbit = copy.copy(self._cache['orbit'])

        P = self['P']
        e = self['e']
        K = self['K']
        omega = self['omega']
        M0 = self['M0']
        a = kwargs.pop('a', P * K / (2*np.pi) * np.sqrt(1 - e**2))

        if len(self) == 1 and len(self.shape) == 0:
            if index > 0:
                raise ValueError('Samples are scalar-valued!')

            trend_coeffs = [self[x] for x in self._trend_names]

        else:
            P = P[index]
            e = e[index]
            a = a[index]
            omega = omega[index]
            M0 = M0[index]
            trend_coeffs = [self[x][index] for x in self._trend_names]

        orbit.elements._P = P
        orbit.elements._e = e * u.dimensionless_unscaled
        orbit.elements._a = a
        orbit.elements._omega = omega
        orbit.elements._M0 = M0
        orbit.elements._Omega = kwargs.pop('Omega', 0*u.deg)
        orbit.elements._i = kwargs.pop('i', 90*u.deg)
        orbit._vtrend = PolynomialRVTrend(trend_coeffs, t0=self.t0)
        orbit._barycenter = kwargs.pop('barycenter', None)

        if kwargs:
            raise ValueError("Unrecognized arguments {0}"
                             .format(', '.join(list(kwargs.keys()))))

        return orbit

    @property
    def orbits(self):
        """A generator that successively returns `twobody.KeplerOrbit` objects
        for each sample. See docstring for `thejoker.JokerSamples.get_orbit` for
        more information.

        """
        for i in range(len(self)):
            yield self.get_orbit(i)

    # Numpy reduce function
    def _apply(self, func):
        cls = self.__class__

        kw = dict()
        for k in self.keys():
            kw[k] = func(self[k])

        kw['t0'] = self.t0
        kw['poly_trend'] = self.poly_trend
        return cls(**kw)

    def mean(self):
        """Return a new scalar object by taking the mean across all samples"""
        return self._apply(np.mean)

    def median(self):
        """Return a new scalar object by taking the medin across all samples"""
        return self._apply(np.mean)

    def std(self):
        """Return a new scalar object by taking the medin across all samples"""
        return self._apply(np.std)
<commit_msg>compute keys on the fly<commit_after># Standard library
from collections import OrderedDict
import copy
import warnings

# Third-party
import astropy.units as u
from astropy.table import Table
from astropy.time import Time
import numpy as np
from twobody import KeplerOrbit, PolynomialRVTrend

# Package
from ..utils import quantity_to_hdf5, quantity_from_hdf5

__all__ = ['JokerSamples']

class JokerSamples(OrderedDict):
    _valid_keys = ['P', 'M0', 'e', 'omega', 'jitter', 'K']

    def __init__(self, t0=None, poly_trend=1, **kwargs):
        """A dictionary-like object for storing posterior samples from
        The Joker, with some extra functionality.

        Parameters
        ----------
        t0 : `astropy.time.Time`, numeric (optional)
            The reference time for the orbital parameters.
        poly_trend : int, optional
            If specified, sample over a polynomial velocity trend with the
            specified number of coefficients. For example, ``poly_trend=3`` will
            sample over parameters of a long-term quadratic velocity trend.
            Default is 1, just a constant velocity shift.
        **kwargs
            These are the orbital element names.
        """

        # initialize empty dictionary
        super(JokerSamples, self).__init__()
        self._setup(poly_trend, t0)

        for key, val in kwargs.items():
            self[key] = val  # calls __setitem__ below

    def _setup(self, poly_trend, t0):
        # reference time
        self.t0 = t0

        self._size = None
        self._shape = None

        self._cache = dict()

        self.poly_trend = int(poly_trend)

    @property
    def _valid_keys(self):
        self._trend_names = ['v{0}'.format(i)
                             for i in range(self.poly_trend)]
        for name in self._trend_names:
            if name not in self._valid_keys:
                self._valid_keys.append(name)

    def __getstate__(self):
        return (self.poly_trend, self.t0, dict(self))

    def __setstate__(self, state):
        poly_trend, t0, data = state
        self._setup(poly_trend, t0)
        self.update(data)

    def __reduce__(self):
        return (JokerSamples, (), self.__getstate__())

    def _validate_key(self, key):
        if key not in self._valid_keys:
            raise ValueError("Invalid key '{0}'.".format(key))

    def _validate_val(self, val):
        val = u.Quantity(val)
        if self._shape is not None and val.shape != self.shape:
            raise ValueError("Shape of new samples must match those already "
                             "stored! ({0}, expected {1})"
                             .format(len(val), self.shape))

        return val

    def __getitem__(self, slc):
        if isinstance(slc, str):
            return super(JokerSamples, self).__getitem__(slc)

        else:
            new = copy.copy(self)
            new._size = None  # reset number of samples
            new._shape = None  # reset number of samples

            for k in self.keys():
                new[k] = self[k][slc]

            return new

    def __setitem__(self, key, val):
        self._validate_key(key)
        val = self._validate_val(val)

        if self._shape is None:
            self._shape = val.shape
            self._size = val.size

        super(JokerSamples, self).__setitem__(key, val)

    @property
    def n_samples(self):
        warnings.warn(".n_samples is deprecated in favor of .size",
                      DeprecationWarning)
        return self.size

    @property
    def size(self):
        if self._size is None:
            raise ValueError("No samples stored!")
        return self._size

    @property
    def shape(self):
        if self._shape is None:
            raise ValueError("No samples stored!")
        return self._shape

    def __len__(self):
        return self.n_samples

    def __str__(self):
        return ("<JokerSamples in [{0}], {1} samples>"
                .format(','.join(self.keys()), len(self)))

    @classmethod
    def from_hdf5(cls, f, n=None, **kwargs):
        """
        Parameters
        ----------
        f : :class:`h5py.File`, :class:`h5py.Group`
        n : int (optional)
            The number of samples to load.
        **kwargs
            All other keyword arguments are passed to the class initializer.
        """

        if 't0_bmjd' in f.attrs:
            # Read the reference time:
            t0 = Time(f.attrs['t0_bmjd'], format='mjd', scale='tcb')
        else:
            t0 = None

        samples = cls(t0=t0, **kwargs)
        for key in cls._valid_keys:
            if key in f:
                samples[key] = quantity_from_hdf5(f, key, n=n)

        return samples

    def to_hdf5(self, f):
        """
        Parameters
        ----------
        f : :class:`h5py.File`, :class:`h5py.Group`
        """

        for key in self.keys():
            quantity_to_hdf5(f, key, self[key])

        if self.t0 is not None:
            f.attrs['t0_bmjd'] = self.t0.tcb.mjd

    @classmethod
    def from_table(cls, tbl_or_f):
        """Read a samples object from an Astropy table.

        Parameters
        ----------
        tbl_or_f : `~astropy.table.Table`, str
            Either a table instance or a string filename to be read with
            `astropy.table.Table.read()`.
        """
        if isinstance(tbl_or_f, str):
            tbl_or_f = Table.read(tbl_or_f)

        kwargs = dict()
        kwargs['poly_trend'] = tbl_or_f.meta.get('poly_trend', 1)
        if 't0_bmjd'.upper() in tbl_or_f.meta:
            kwargs['t0'] = Time(tbl_or_f.meta['t0_bmjd'.upper()], format='mjd',
                                scale='tcb')

        samples = cls(**kwargs)
        for key in cls._valid_keys:
            if key in tbl_or_f.colnames:
                samples[key] = u.Quantity(tbl_or_f[key])

        return samples

    def to_table(self):
        """Convert the samples to an Astropy table object.

        Returns
        -------
        tbl : `~astropy.table.Table`
        """
        tbl = Table()
        for k in self.keys():
            tbl[k] = self[k]

        tbl.meta['poly_trend'] = self.poly_trend

        if self.t0 is not None:
            tbl.meta['t0_bmjd'] = self.t0.tcb.mjd

        return tbl

    ##########################################################################
    # Interaction with TwoBody

    def get_orbit(self, index=None, **kwargs):
        """Get a `twobody.KeplerOrbit` object for the samples at the specified
        index.

        Parameters
        ----------
        index : int (optional)
            The index of the samples to turn into a `twobody.KeplerOrbit`
            instance. If the samples object is scalar, no index is necessary.
        **kwargs
            Other keyword arguments are passed to the `twobody.KeplerOrbit`
            initializer. For example, you can specify the inclination by passing
            ``i=...`, or  longitude of the ascending node by passing
            ``Omega=...``.

        Returns
        -------
        orbit : `twobody.KeplerOrbit`
            The samples converted to an orbit object. The barycenter position
            and distance are set to arbitrary values.
        """
        if 'orbit' not in self._cache:
            self._cache['orbit'] = KeplerOrbit(P=1*u.yr, e=0., omega=0*u.deg,
                                               Omega=0*u.deg, i=90*u.deg,
                                               a=1*u.au, t0=self.t0)

        # all of this to avoid the __init__ of KeplerOrbit / KeplerElements
        orbit = copy.copy(self._cache['orbit'])

        P = self['P']
        e = self['e']
        K = self['K']
        omega = self['omega']
        M0 = self['M0']
        a = kwargs.pop('a', P * K / (2*np.pi) * np.sqrt(1 - e**2))

        if len(self) == 1 and len(self.shape) == 0:
            if index > 0:
                raise ValueError('Samples are scalar-valued!')

            trend_coeffs = [self[x] for x in self._trend_names]

        else:
            P = P[index]
            e = e[index]
            a = a[index]
            omega = omega[index]
            M0 = M0[index]
            trend_coeffs = [self[x][index] for x in self._trend_names]

        orbit.elements._P = P
        orbit.elements._e = e * u.dimensionless_unscaled
        orbit.elements._a = a
        orbit.elements._omega = omega
        orbit.elements._M0 = M0
        orbit.elements._Omega = kwargs.pop('Omega', 0*u.deg)
        orbit.elements._i = kwargs.pop('i', 90*u.deg)
        orbit._vtrend = PolynomialRVTrend(trend_coeffs, t0=self.t0)
        orbit._barycenter = kwargs.pop('barycenter', None)

        if kwargs:
            raise ValueError("Unrecognized arguments {0}"
                             .format(', '.join(list(kwargs.keys()))))

        return orbit

    @property
    def orbits(self):
        """A generator that successively returns `twobody.KeplerOrbit` objects
        for each sample. See docstring for `thejoker.JokerSamples.get_orbit` for
        more information.

        """
        for i in range(len(self)):
            yield self.get_orbit(i)

    # Numpy reduce function
    def _apply(self, func):
        cls = self.__class__

        kw = dict()
        for k in self.keys():
            kw[k] = func(self[k])

        kw['t0'] = self.t0
        kw['poly_trend'] = self.poly_trend
        return cls(**kw)

    def mean(self):
        """Return a new scalar object by taking the mean across all samples"""
        return self._apply(np.mean)

    def median(self):
        """Return a new scalar object by taking the medin across all samples"""
        return self._apply(np.mean)

    def std(self):
        """Return a new scalar object by taking the medin across all samples"""
        return self._apply(np.std)
<|endoftext|>


<commit_before>"""
Sample mapper.
"""
from everest.repositories.rdb.utils import mapper
from sqlalchemy import String
from sqlalchemy import event
from sqlalchemy.orm import column_property
from sqlalchemy.orm import relationship
from sqlalchemy.orm.session import object_session
from sqlalchemy.orm.util import class_mapper
from sqlalchemy.sql import cast
from sqlalchemy.sql import literal
from sqlalchemy.sql import select
from sqlalchemy.sql.expression import and_
from sqlalchemy.sql.expression import func
from sqlalchemy.sql.expression import insert
from sqlalchemy.sql.expression import null
from sqlalchemy.sql.functions import coalesce
from thelma.db.utils import string_agg
from thelma.models.container import Container
from thelma.models.sample import SAMPLE_TYPES
from thelma.models.sample import Sample
from thelma.models.sample import SampleMolecule
from thelma.models.sample import StockSample

__docformat__ = 'reStructuredText en'
__all__ = ['create_mapper']

def create_mapper(sample_tbl, sample_molecule_tbl, molecule_tbl,
                  molecule_design_pool_tbl):
    "Mapper factory."
    s = sample_tbl
    sm = sample_molecule_tbl
    m = molecule_tbl
    mdp = molecule_design_pool_tbl
    s1 = sample_tbl.alias()
    # FIXME: The following construct introduces a dependency on string_agg
    #        in the SQL engine.
    mds_sel = select(
        [mdp.c.molecule_design_set_id],
        mdp.c.member_hash ==
          select([func.md5(
                    string_agg(cast(m.c.molecule_design_id, String),
                               literal(';'),
                               order_by=m.c.molecule_design_id))
                  ],
                 from_obj=[s1.join(sm,
                                   and_(sm.c.sample_id == s1.c.sample_id,
                                        s1.c.sample_id == s.c.sample_id))
                           .join(m,
                                 m.c.molecule_id == sm.c.molecule_id)
                           ]) \
                .group_by(sm.c.sample_id))
    m = mapper(Sample, sample_tbl,
        id_attribute='sample_id',
        properties=dict(
            molecule_design_pool_id=
                    column_property(coalesce(mds_sel.as_scalar(), null())),
            container=relationship(Container,
                                   uselist=False,
                                   back_populates='sample'),
            sample_molecules=
                    relationship(SampleMolecule,
                                 back_populates='sample',
                                 lazy='joined'
                                 ),
            ),
        polymorphic_on=sample_tbl.c.sample_type,
        polymorphic_identity=SAMPLE_TYPES.BASIC
        )
    # Listen to changes to the sample_type attribute.
    event.listen(Sample.sample_type, "set", check_set_sample_type) # pylint: disable=E1101
    return m

def check_set_sample_type(target, value, oldvalue, initiator): # pylint: disable=W0613
    if isinstance(target, Sample) \
       and value == SAMPLE_TYPES.STOCK \
       and oldvalue != SAMPLE_TYPES.STOCK:
        sess = object_session(target)
        if target.id is None:
            # We need a sample ID for the following execute statement to work.
            sess.flush()
        ss_tbl = class_mapper(StockSample).local_table
        sess.execute(insert(ss_tbl,
                            values=dict(sample_id=target.sample_id,
                                        molecule_design_set_id=
                                            target.molecule_design_pool.id,
                                        supplier_id=target.supplier.id,
                                        molecule_type_id=
                                            target.molecule_type.id,
                                        concentration=
                                            target.concentration)
                                  )
                           )
<commit_msg>Added a missing cascade option to the sample mapper.<commit_after>"""
Sample mapper.
"""
from everest.repositories.rdb.utils import mapper
from sqlalchemy import String
from sqlalchemy import event
from sqlalchemy.orm import column_property
from sqlalchemy.orm import relationship
from sqlalchemy.orm.session import object_session
from sqlalchemy.orm.util import class_mapper
from sqlalchemy.sql import cast
from sqlalchemy.sql import literal
from sqlalchemy.sql import select
from sqlalchemy.sql.expression import and_
from sqlalchemy.sql.expression import func
from sqlalchemy.sql.expression import insert
from sqlalchemy.sql.expression import null
from sqlalchemy.sql.functions import coalesce
from thelma.db.utils import string_agg
from thelma.models.container import Container
from thelma.models.sample import SAMPLE_TYPES
from thelma.models.sample import Sample
from thelma.models.sample import SampleMolecule
from thelma.models.sample import StockSample

__docformat__ = 'reStructuredText en'
__all__ = ['create_mapper']

def create_mapper(sample_tbl, sample_molecule_tbl, molecule_tbl,
                  molecule_design_pool_tbl):
    "Mapper factory."
    s = sample_tbl
    sm = sample_molecule_tbl
    m = molecule_tbl
    mdp = molecule_design_pool_tbl
    s1 = sample_tbl.alias()
    # FIXME: The following construct introduces a dependency on string_agg
    #        in the SQL engine.
    mds_sel = select(
        [mdp.c.molecule_design_set_id],
        mdp.c.member_hash ==
          select([func.md5(
                    string_agg(cast(m.c.molecule_design_id, String),
                               literal(';'),
                               order_by=m.c.molecule_design_id))
                  ],
                 from_obj=[s1.join(sm,
                                   and_(sm.c.sample_id == s1.c.sample_id,
                                        s1.c.sample_id == s.c.sample_id))
                           .join(m,
                                 m.c.molecule_id == sm.c.molecule_id)
                           ]) \
                .group_by(sm.c.sample_id))
    m = mapper(Sample, sample_tbl,
        id_attribute='sample_id',
        properties=dict(
            molecule_design_pool_id=
                    column_property(coalesce(mds_sel.as_scalar(), null())),
            container=relationship(Container,
                                   uselist=False,
                                   back_populates='sample'),
            sample_molecules=
                    relationship(SampleMolecule,
                                 back_populates='sample',
                                 cascade='all,delete,delete-orphan',
                                 lazy='joined'
                                 ),
            ),
        polymorphic_on=sample_tbl.c.sample_type,
        polymorphic_identity=SAMPLE_TYPES.BASIC
        )
    # Listen to changes to the sample_type attribute.
    event.listen(Sample.sample_type, "set", check_set_sample_type) # pylint: disable=E1101
    return m

def check_set_sample_type(target, value, oldvalue, initiator): # pylint: disable=W0613
    if isinstance(target, Sample) \
       and value == SAMPLE_TYPES.STOCK \
       and oldvalue != SAMPLE_TYPES.STOCK:
        sess = object_session(target)
        if target.id is None:
            # We need a sample ID for the following execute statement to work.
            sess.flush()
        ss_tbl = class_mapper(StockSample).local_table
        sess.execute(insert(ss_tbl,
                            values=dict(sample_id=target.sample_id,
                                        molecule_design_set_id=
                                            target.molecule_design_pool.id,
                                        supplier_id=target.supplier.id,
                                        molecule_type_id=
                                            target.molecule_type.id,
                                        concentration=
                                            target.concentration)
                                  )
                           )
<|endoftext|>


<commit_before>"""Models.

Changes affecting results or their presentation should also update
constants.py `change_date`,
"""

from __future__ import annotations

from datetime import date, datetime, timedelta
from logging import INFO, basicConfig, getLogger
from sys import stdout
from typing import Dict, Generator, Tuple, Sequence,Optional

import numpy as np
import pandas as pd

from .constants import EPSILON, CHANGE_DATE
from .parameters import Parameters

basicConfig(
    level=INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    stream=stdout,
)
logger = getLogger(__name__)

class SimSirModel:

    def __init__(self, p: Parameters):

        self.rates = {
            key: d.rate
            for key, d in p.dispositions.items()
        }

        self.days = {
            key: d.days
            for key, d in p.dispositions.items()
        }

        self.keys = ("susceptible", "infected", "recovered")

        # An estimate of the number of infected people on the day that
        # the first hospitalized case is seen
        #
        # Note: this should not be an integer.
        infected = (
            1.0 / p.market_share / p.hospitalized.rate
        )

        susceptible = p.population - infected

        gamma = 1.0 / p.infectious_days
        self.gamma = gamma

        self.susceptible = susceptible
        self.infected = infected
        self.recovered = p.recovered

        if p.date_first_hospitalized is None and p.doubling_time is not None:
            # Back-projecting to when the first hospitalized case would have
            logger.info('Using doubling_time: %s', p.doubling_time)

            intrinsic_growth_rate = get_growth_rate(p.doubling_time)

            self.beta = get_beta(intrinsic_growth_rate,  gamma, self.susceptible, 0.0)

            self.i_day = 0 # seed to the full length
            self.beta_t = self.beta
            self.run_projection(p)
            self.i_day = i_day = int(get_argmin_ds(self.census_df, p.current_hospitalized))

            self.beta_t = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, p.relative_contact_rate)
            self.run_projection(p)

            logger.info('Set i_day = %s', i_day)
            p.date_first_hospitalized = p.current_date - timedelta(days=i_day)
            logger.info(
                'Estimated date_first_hospitalized: %s; current_date: %s; i_day: %s',
                p.date_first_hospitalized,
                p.current_date,
                self.i_day)

        elif p.date_first_hospitalized is not None and p.doubling_time is None:
            # Fitting spread parameter to observed hospital census (dates of 1 patient and today)
            self.i_day = (p.current_date - p.date_first_hospitalized).days
            logger.info(
                'Using date_first_hospitalized: %s; current_date: %s; i_day: %s',
                p.date_first_hospitalized,
                p.current_date,
                self.i_day)
            min_loss = 2.0**99
            dts = np.linspace(1, 15, 29)
            losses = np.zeros(dts.shape[0])
            self.current_hospitalized = p.current_hospitalized
            for i, i_dt in enumerate(dts):
                intrinsic_growth_rate = get_growth_rate(i_dt)
                self.beta = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, 0.0)
                self.beta_t = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, p.relative_contact_rate)

                self.run_projection(p)
                loss = self.get_loss()
                losses[i] = loss

            p.doubling_time = dts[pd.Series(losses).argmin()]
            logger.info('Estimated doubling_time: %s', p.doubling_time)
            intrinsic_growth_rate = get_growth_rate(p.doubling_time)
            self.beta = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, 0.0)
            self.beta_t = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, p.relative_contact_rate)
            self.run_projection(p)

            self.population = p.population
        else:
            logger.info(
                'doubling_time: %s; date_first_hospitalized: %s',
                p.doubling_time,
                p.date_first_hospitalized,
            )
            raise AssertionError('doubling_time or date_first_hospitalized must be provided.')

        logger.info('len(np.arange(-i_day, n_days+1)): %s', len(np.arange(-self.i_day, p.n_days+1)))
        logger.info('len(raw_df): %s', len(self.raw_df))

        self.infected = self.raw_df['infected'].values[self.i_day]
        self.susceptible = self.raw_df['susceptible'].values[self.i_day]
        self.recovered = self.raw_df['recovered'].values[self.i_day]

        self.intrinsic_growth_rate = intrinsic_growth_rate

        # r_t is r_0 after distancing
        self.r_t = self.beta_t / gamma * susceptible
        self.r_naught = self.beta / gamma * susceptible

        doubling_time_t = 1.0 / np.log2(
            self.beta_t * susceptible - gamma + 1)
        self.doubling_time_t = doubling_time_t

        self.sim_sir_w_date_df = build_sim_sir_w_date_df(self.raw_df, p.current_date, self.keys)

        self.sim_sir_w_date_floor_df = build_floor_df(self.sim_sir_w_date_df, self.keys)
        self.admits_floor_df = build_floor_df(self.admits_df, p.dispositions.keys())
        self.census_floor_df = build_floor_df(self.census_df, p.dispositions.keys())

        self.daily_growth_rate = get_growth_rate(p.doubling_time)
        self.daily_growth_rate_t = get_growth_rate(self.doubling_time_t)

    def run_projection(self, p):
        self.raw_df = sim_sir_df(
            self.susceptible,
            self.infected,
            p.recovered,
            self.gamma,
            -self.i_day,
            self.beta,
            self.i_day,
            self.beta_t,
            p.n_days
        )
        self.dispositions_df = build_dispositions_df(self.raw_df, self.rates, p.market_share, p.current_date)
        self.admits_df = build_admits_df(self.dispositions_df)
        self.census_df = build_census_df(self.admits_df, self.days)
        self.current_infected = self.raw_df.infected.loc[self.i_day]

    def get_loss(self) -> float:
        """Squared error: predicted vs. actual current hospitalized."""
        predicted = self.census_df.hospitalized.loc[self.i_day]
        return (self.current_hospitalized - predicted) ** 2.0

def get_argmin_ds(census_df: pd.DataFrame, current_hospitalized: float) -> float:
    # By design, this forbids choosing a day after the peak
    # If that's a problem, see #381
    peak_day = census_df.hospitalized.argmax()
    losses_df = (census_df.hospitalized[:peak_day] - current_hospitalized) ** 2.0
    return losses_df.argmin()

def get_beta(
    intrinsic_growth_rate: float,
    gamma: float,
    susceptible: float,
    relative_contact_rate: float
) -> float:
    return (
        (intrinsic_growth_rate + gamma)
        / susceptible
        * (1.0 - relative_contact_rate)
    )

def get_growth_rate(doubling_time: Optional[float]) -> float:
    """Calculates average daily growth rate from doubling time."""
    if doubling_time is None or doubling_time == 0.0:
        return 0.0
    return (2.0 ** (1.0 / doubling_time) - 1.0)

def sir(
    s: float, i: float, r: float, beta: float, gamma: float, n: float
) -> Tuple[float, float, float]:
    """The SIR model, one time step."""
    s_n = (-beta * s * i) + s
    i_n = (beta * s * i - gamma * i) + i
    r_n = gamma * i + r

    # TODO:
    #   Post check dfs for negative values and
    #   warn the user that their input data is bad.
    #   JL: I suspect that these adjustments covered bugs.

    #if s_n < 0.0:
    #    s_n = 0.0
    #if i_n < 0.0:
    #    i_n = 0.0
    #if r_n < 0.0:
    #    r_n = 0.0
    scale = n / (s_n + i_n + r_n)
    return s_n * scale, i_n * scale, r_n * scale

def gen_sir(
    s: float, i: float, r: float, gamma: float, i_day: int, *args
) -> Generator[Tuple[int, float, float, float], None, None]:
    """Simulate SIR model forward in time yielding tuples.
    Parameter order has changed to allow multiple (beta, n_days)
    to reflect multiple changing social distancing policies.
    """
    s, i, r = (float(v) for v in (s, i, r))
    n = s + i + r
    d = i_day
    while args:
        beta, n_days, *args = args
        for _ in range(n_days):
            yield d, s, i, r
            s, i, r = sir(s, i, r, beta, gamma, n)
            d += 1
    yield d, s, i, r

def sim_sir_df(
    s: float, i: float, r: float,
    gamma: float, i_day: int, *args
) -> pd.DataFrame:
    """Simulate the SIR model forward in time."""
    return pd.DataFrame(
        data=gen_sir(s, i, r, gamma, i_day, *args),
        columns=("day", "susceptible", "infected", "recovered"),
    )

def build_sim_sir_w_date_df(
    raw_df: pd.DataFrame,
    current_date: datetime,
    keys: Sequence[str],
) -> pd.DataFrame:
    day = raw_df.day
    return pd.DataFrame({
        "day": day,
        "date": day.astype('timedelta64[D]') + np.datetime64(current_date),
        **{
            key: raw_df[key]
            for key in keys
        }
    })

def build_floor_df(df, keys):
    """Build floor sim sir w date."""
    return pd.DataFrame({
        "day": df.day,
        "date": df.date,
        **{
            key: np.floor(df[key])
            for key in keys
        }
    })

def build_dispositions_df(
    raw_df: pd.DataFrame,
    rates: Dict[str, float],
    market_share: float,
    current_date: datetime,
) -> pd.DataFrame:
    """Build dispositions dataframe of patients adjusted by rate and market_share."""
    patients = raw_df.infected + raw_df.recovered
    day = raw_df.day
    return pd.DataFrame({
        "day": day,
        "date": day.astype('timedelta64[D]') + np.datetime64(current_date),
        **{
            key: patients * rate * market_share
            for key, rate in rates.items()
        }
    })

def build_admits_df(dispositions_df: pd.DataFrame) -> pd.DataFrame:
    """Build admits dataframe from dispositions."""
    admits_df = dispositions_df.iloc[:, :] - dispositions_df.shift(1)
    admits_df.day = dispositions_df.day
    admits_df.date = dispositions_df.date
    return admits_df

def build_census_df(
    admits_df: pd.DataFrame,
    lengths_of_stay: Dict[str, int],
) -> pd.DataFrame:
    """Average Length of Stay for each disposition of COVID-19 case (total guesses)"""
    return pd.DataFrame({
        'day': admits_df.day,
        'date': admits_df.date,
        **{
            key: (
                admits_df[key].cumsum()
                - admits_df[key].cumsum().shift(los).fillna(0)
            )
            for key, los in lengths_of_stay.items()
        }
    })
<commit_msg>Fix truncated comment<commit_after>"""Models.

Changes affecting results or their presentation should also update
constants.py `change_date`,
"""

from __future__ import annotations

from datetime import date, datetime, timedelta
from logging import INFO, basicConfig, getLogger
from sys import stdout
from typing import Dict, Generator, Tuple, Sequence,Optional

import numpy as np
import pandas as pd

from .constants import EPSILON, CHANGE_DATE
from .parameters import Parameters

basicConfig(
    level=INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    stream=stdout,
)
logger = getLogger(__name__)

class SimSirModel:

    def __init__(self, p: Parameters):

        self.rates = {
            key: d.rate
            for key, d in p.dispositions.items()
        }

        self.days = {
            key: d.days
            for key, d in p.dispositions.items()
        }

        self.keys = ("susceptible", "infected", "recovered")

        # An estimate of the number of infected people on the day that
        # the first hospitalized case is seen
        #
        # Note: this should not be an integer.
        infected = (
            1.0 / p.market_share / p.hospitalized.rate
        )

        susceptible = p.population - infected

        gamma = 1.0 / p.infectious_days
        self.gamma = gamma

        self.susceptible = susceptible
        self.infected = infected
        self.recovered = p.recovered

        if p.date_first_hospitalized is None and p.doubling_time is not None:
            # Back-projecting to when the first hospitalized case would have been admitted
            logger.info('Using doubling_time: %s', p.doubling_time)

            intrinsic_growth_rate = get_growth_rate(p.doubling_time)

            self.beta = get_beta(intrinsic_growth_rate,  gamma, self.susceptible, 0.0)

            self.i_day = 0 # seed to the full length
            self.beta_t = self.beta
            self.run_projection(p)
            self.i_day = i_day = int(get_argmin_ds(self.census_df, p.current_hospitalized))

            self.beta_t = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, p.relative_contact_rate)
            self.run_projection(p)

            logger.info('Set i_day = %s', i_day)
            p.date_first_hospitalized = p.current_date - timedelta(days=i_day)
            logger.info(
                'Estimated date_first_hospitalized: %s; current_date: %s; i_day: %s',
                p.date_first_hospitalized,
                p.current_date,
                self.i_day)

        elif p.date_first_hospitalized is not None and p.doubling_time is None:
            # Fitting spread parameter to observed hospital census (dates of 1 patient and today)
            self.i_day = (p.current_date - p.date_first_hospitalized).days
            logger.info(
                'Using date_first_hospitalized: %s; current_date: %s; i_day: %s',
                p.date_first_hospitalized,
                p.current_date,
                self.i_day)
            min_loss = 2.0**99
            dts = np.linspace(1, 15, 29)
            losses = np.zeros(dts.shape[0])
            self.current_hospitalized = p.current_hospitalized
            for i, i_dt in enumerate(dts):
                intrinsic_growth_rate = get_growth_rate(i_dt)
                self.beta = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, 0.0)
                self.beta_t = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, p.relative_contact_rate)

                self.run_projection(p)
                loss = self.get_loss()
                losses[i] = loss

            p.doubling_time = dts[pd.Series(losses).argmin()]
            logger.info('Estimated doubling_time: %s', p.doubling_time)
            intrinsic_growth_rate = get_growth_rate(p.doubling_time)
            self.beta = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, 0.0)
            self.beta_t = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, p.relative_contact_rate)
            self.run_projection(p)

            self.population = p.population
        else:
            logger.info(
                'doubling_time: %s; date_first_hospitalized: %s',
                p.doubling_time,
                p.date_first_hospitalized,
            )
            raise AssertionError('doubling_time or date_first_hospitalized must be provided.')

        logger.info('len(np.arange(-i_day, n_days+1)): %s', len(np.arange(-self.i_day, p.n_days+1)))
        logger.info('len(raw_df): %s', len(self.raw_df))

        self.infected = self.raw_df['infected'].values[self.i_day]
        self.susceptible = self.raw_df['susceptible'].values[self.i_day]
        self.recovered = self.raw_df['recovered'].values[self.i_day]

        self.intrinsic_growth_rate = intrinsic_growth_rate

        # r_t is r_0 after distancing
        self.r_t = self.beta_t / gamma * susceptible
        self.r_naught = self.beta / gamma * susceptible

        doubling_time_t = 1.0 / np.log2(
            self.beta_t * susceptible - gamma + 1)
        self.doubling_time_t = doubling_time_t

        self.sim_sir_w_date_df = build_sim_sir_w_date_df(self.raw_df, p.current_date, self.keys)

        self.sim_sir_w_date_floor_df = build_floor_df(self.sim_sir_w_date_df, self.keys)
        self.admits_floor_df = build_floor_df(self.admits_df, p.dispositions.keys())
        self.census_floor_df = build_floor_df(self.census_df, p.dispositions.keys())

        self.daily_growth_rate = get_growth_rate(p.doubling_time)
        self.daily_growth_rate_t = get_growth_rate(self.doubling_time_t)

    def run_projection(self, p):
        self.raw_df = sim_sir_df(
            self.susceptible,
            self.infected,
            p.recovered,
            self.gamma,
            -self.i_day,
            self.beta,
            self.i_day,
            self.beta_t,
            p.n_days
        )
        self.dispositions_df = build_dispositions_df(self.raw_df, self.rates, p.market_share, p.current_date)
        self.admits_df = build_admits_df(self.dispositions_df)
        self.census_df = build_census_df(self.admits_df, self.days)
        self.current_infected = self.raw_df.infected.loc[self.i_day]

    def get_loss(self) -> float:
        """Squared error: predicted vs. actual current hospitalized."""
        predicted = self.census_df.hospitalized.loc[self.i_day]
        return (self.current_hospitalized - predicted) ** 2.0

def get_argmin_ds(census_df: pd.DataFrame, current_hospitalized: float) -> float:
    # By design, this forbids choosing a day after the peak
    # If that's a problem, see #381
    peak_day = census_df.hospitalized.argmax()
    losses_df = (census_df.hospitalized[:peak_day] - current_hospitalized) ** 2.0
    return losses_df.argmin()

def get_beta(
    intrinsic_growth_rate: float,
    gamma: float,
    susceptible: float,
    relative_contact_rate: float
) -> float:
    return (
        (intrinsic_growth_rate + gamma)
        / susceptible
        * (1.0 - relative_contact_rate)
    )

def get_growth_rate(doubling_time: Optional[float]) -> float:
    """Calculates average daily growth rate from doubling time."""
    if doubling_time is None or doubling_time == 0.0:
        return 0.0
    return (2.0 ** (1.0 / doubling_time) - 1.0)

def sir(
    s: float, i: float, r: float, beta: float, gamma: float, n: float
) -> Tuple[float, float, float]:
    """The SIR model, one time step."""
    s_n = (-beta * s * i) + s
    i_n = (beta * s * i - gamma * i) + i
    r_n = gamma * i + r

    # TODO:
    #   Post check dfs for negative values and
    #   warn the user that their input data is bad.
    #   JL: I suspect that these adjustments covered bugs.

    #if s_n < 0.0:
    #    s_n = 0.0
    #if i_n < 0.0:
    #    i_n = 0.0
    #if r_n < 0.0:
    #    r_n = 0.0
    scale = n / (s_n + i_n + r_n)
    return s_n * scale, i_n * scale, r_n * scale

def gen_sir(
    s: float, i: float, r: float, gamma: float, i_day: int, *args
) -> Generator[Tuple[int, float, float, float], None, None]:
    """Simulate SIR model forward in time yielding tuples.
    Parameter order has changed to allow multiple (beta, n_days)
    to reflect multiple changing social distancing policies.
    """
    s, i, r = (float(v) for v in (s, i, r))
    n = s + i + r
    d = i_day
    while args:
        beta, n_days, *args = args
        for _ in range(n_days):
            yield d, s, i, r
            s, i, r = sir(s, i, r, beta, gamma, n)
            d += 1
    yield d, s, i, r

def sim_sir_df(
    s: float, i: float, r: float,
    gamma: float, i_day: int, *args
) -> pd.DataFrame:
    """Simulate the SIR model forward in time."""
    return pd.DataFrame(
        data=gen_sir(s, i, r, gamma, i_day, *args),
        columns=("day", "susceptible", "infected", "recovered"),
    )

def build_sim_sir_w_date_df(
    raw_df: pd.DataFrame,
    current_date: datetime,
    keys: Sequence[str],
) -> pd.DataFrame:
    day = raw_df.day
    return pd.DataFrame({
        "day": day,
        "date": day.astype('timedelta64[D]') + np.datetime64(current_date),
        **{
            key: raw_df[key]
            for key in keys
        }
    })

def build_floor_df(df, keys):
    """Build floor sim sir w date."""
    return pd.DataFrame({
        "day": df.day,
        "date": df.date,
        **{
            key: np.floor(df[key])
            for key in keys
        }
    })

def build_dispositions_df(
    raw_df: pd.DataFrame,
    rates: Dict[str, float],
    market_share: float,
    current_date: datetime,
) -> pd.DataFrame:
    """Build dispositions dataframe of patients adjusted by rate and market_share."""
    patients = raw_df.infected + raw_df.recovered
    day = raw_df.day
    return pd.DataFrame({
        "day": day,
        "date": day.astype('timedelta64[D]') + np.datetime64(current_date),
        **{
            key: patients * rate * market_share
            for key, rate in rates.items()
        }
    })

def build_admits_df(dispositions_df: pd.DataFrame) -> pd.DataFrame:
    """Build admits dataframe from dispositions."""
    admits_df = dispositions_df.iloc[:, :] - dispositions_df.shift(1)
    admits_df.day = dispositions_df.day
    admits_df.date = dispositions_df.date
    return admits_df

def build_census_df(
    admits_df: pd.DataFrame,
    lengths_of_stay: Dict[str, int],
) -> pd.DataFrame:
    """Average Length of Stay for each disposition of COVID-19 case (total guesses)"""
    return pd.DataFrame({
        'day': admits_df.day,
        'date': admits_df.date,
        **{
            key: (
                admits_df[key].cumsum()
                - admits_df[key].cumsum().shift(los).fillna(0)
            )
            for key, los in lengths_of_stay.items()
        }
    })
<|endoftext|>