<commit_before>import asyncio
import warnings
from inspect import isawaitable
from typing import (
    Any,
    Awaitable,
    Callable,
    Dict,
    List,
    MutableMapping,
    Optional,
    Tuple,
    Union,
)
from urllib.parse import quote

import sanic.app  # noqa
from sanic.compat import Header
from sanic.exceptions import InvalidUsage, ServerError
from sanic.log import logger
from sanic.request import Request
from sanic.response import HTTPResponse, StreamingHTTPResponse
from sanic.server import ConnInfo, StreamBuffer
from sanic.websocket import WebSocketConnection

ASGIScope = MutableMapping[str, Any]
ASGIMessage = MutableMapping[str, Any]
ASGISend = Callable[[ASGIMessage], Awaitable[None]]
ASGIReceive = Callable[[], Awaitable[ASGIMessage]]

class MockProtocol:
    def __init__(self, transport: "MockTransport", loop):
        self.transport = transport
        self._not_paused = asyncio.Event(loop=loop)
        self._not_paused.set()
        self._complete = asyncio.Event(loop=loop)

    def pause_writing(self) -> None:
        self._not_paused.clear()

    def resume_writing(self) -> None:
        self._not_paused.set()

    async def complete(self) -> None:
        self._not_paused.set()
        await self.transport.send(
            {"type": "http.response.body", "body": b"", "more_body": False}
        )

    @property
    def is_complete(self) -> bool:
        return self._complete.is_set()

    async def push_data(self, data: bytes) -> None:
        if not self.is_complete:
            await self.transport.send(
                {"type": "http.response.body", "body": data, "more_body": True}
            )

    async def drain(self) -> None:
        await self._not_paused.wait()

class MockTransport:
    _protocol: Optional[MockProtocol]

    def __init__(self, scope: ASGIScope, receive: ASGIReceive, send: ASGISend) -> None:
        self.scope = scope
        self._receive = receive
        self._send = send
        self._protocol = None
        self.loop = None

    def get_protocol(self) -> MockProtocol:
        if not self._protocol:
            self._protocol = MockProtocol(self, self.loop)
        return self._protocol

    def get_extra_info(self, info: str) -> Union[str, bool, None]:
        if info == "peername":
            return self.scope.get("server")
        elif info == "sslcontext":
            return self.scope.get("scheme") in ["https", "wss"]
        return None

    def get_websocket_connection(self) -> WebSocketConnection:
        try:
            return self._websocket_connection
        except AttributeError:
            raise InvalidUsage("Improper websocket connection.")

    def create_websocket_connection(
        self, send: ASGISend, receive: ASGIReceive
    ) -> WebSocketConnection:
        self._websocket_connection = WebSocketConnection(send, receive)
        return self._websocket_connection

    def add_task(self) -> None:
        raise NotImplementedError

    async def send(self, data) -> None:
        # TODO:
        # - Validation on data and that it is formatted properly and is valid
        await self._send(data)

    async def receive(self) -> ASGIMessage:
        return await self._receive()

class Lifespan:
    def __init__(self, asgi_app: "ASGIApp") -> None:
        self.asgi_app = asgi_app

        print(self.asgi_app.sanic_app.listeners)

        if "before_server_start" in self.asgi_app.sanic_app.listeners:
            warnings.warn(
                'You have set a listener for "before_server_start" '
                "in ASGI mode. "
                "It will be executed as early as possible, but not before "
                "the ASGI server is started."
            )
        if "after_server_stop" in self.asgi_app.sanic_app.listeners:
            warnings.warn(
                'You have set a listener for "after_server_stop" '
                "in ASGI mode. "
                "It will be executed as late as possible, but not after "
                "the ASGI server is stopped."
            )

    async def startup(self) -> None:
        """
        Gather the listeners to fire on server start.
        Because we are using a third-party server and not Sanic server, we do
        not have access to fire anything BEFORE the server starts.
        Therefore, we fire before_server_start and after_server_start
        in sequence since the ASGI lifespan protocol only supports a single
        startup event.
        """
        listeners = self.asgi_app.sanic_app.listeners.get(
            "before_server_start", []
        ) + self.asgi_app.sanic_app.listeners.get("after_server_start", [])

        for handler in listeners:
            response = handler(self.asgi_app.sanic_app, self.asgi_app.sanic_app.loop)
            if isawaitable(response):
                await response

    async def shutdown(self) -> None:
        """
        Gather the listeners to fire on server stop.
        Because we are using a third-party server and not Sanic server, we do
        not have access to fire anything AFTER the server stops.
        Therefore, we fire before_server_stop and after_server_stop
        in sequence since the ASGI lifespan protocol only supports a single
        shutdown event.
        """
        listeners = self.asgi_app.sanic_app.listeners.get(
            "before_server_stop", []
        ) + self.asgi_app.sanic_app.listeners.get("after_server_stop", [])

        for handler in listeners:
            response = handler(self.asgi_app.sanic_app, self.asgi_app.sanic_app.loop)
            if isawaitable(response):
                await response

    async def __call__(
        self, scope: ASGIScope, receive: ASGIReceive, send: ASGISend
    ) -> None:
        message = await receive()
        if message["type"] == "lifespan.startup":
            await self.startup()
            await send({"type": "lifespan.startup.complete"})

        message = await receive()
        if message["type"] == "lifespan.shutdown":
            await self.shutdown()
            await send({"type": "lifespan.shutdown.complete"})

class ASGIApp:
    sanic_app: "sanic.app.Sanic"
    request: Request
    transport: MockTransport
    do_stream: bool
    lifespan: Lifespan
    ws: Optional[WebSocketConnection]

    def __init__(self) -> None:
        self.ws = None

    @classmethod
    async def create(
        cls, sanic_app, scope: ASGIScope, receive: ASGIReceive, send: ASGISend
    ) -> "ASGIApp":
        instance = cls()
        instance.sanic_app = sanic_app
        instance.transport = MockTransport(scope, receive, send)
        instance.transport.loop = sanic_app.loop
        setattr(instance.transport, "add_task", sanic_app.loop.create_task)

        headers = Header(
            [
                (key.decode("latin-1"), value.decode("latin-1"))
                for key, value in scope.get("headers", [])
            ]
        )
        instance.do_stream = True if headers.get("expect") == "100-continue" else False
        instance.lifespan = Lifespan(instance)

        if scope["type"] == "lifespan":
            await instance.lifespan(scope, receive, send)
        else:
            path = scope["path"][1:] if scope["path"].startswith("/") else scope["path"]
            url = "/".join([scope.get("root_path", ""), quote(path)])
            url_bytes = url.encode("latin-1")
            url_bytes += b"?" + scope["query_string"]

            if scope["type"] == "http":
                version = scope["http_version"]
                method = scope["method"]
            elif scope["type"] == "websocket":
                version = "1.1"
                method = "GET"

                instance.ws = instance.transport.create_websocket_connection(
                    send, receive
                )
                await instance.ws.accept()
            else:
                pass
                # TODO:
                # - close connection

            request_class = sanic_app.request_class or Request
            instance.request = request_class(
                url_bytes, headers, version, method, instance.transport, sanic_app,
            )
            instance.request.conn_info = ConnInfo(instance.transport)

            if sanic_app.is_request_stream:
                is_stream_handler = sanic_app.router.is_stream_handler(instance.request)
                if is_stream_handler:
                    instance.request.stream = StreamBuffer(
                        sanic_app.config.REQUEST_BUFFER_QUEUE_SIZE
                    )
                    instance.do_stream = True

        return instance

    async def read_body(self) -> bytes:
        """
        Read and return the entire body from an incoming ASGI message.
        """
        body = b""
        more_body = True
        while more_body:
            message = await self.transport.receive()
            body += message.get("body", b"")
            more_body = message.get("more_body", False)

        return body

    async def stream_body(self) -> None:
        """
        Read and stream the body in chunks from an incoming ASGI message.
        """
        more_body = True

        while more_body:
            message = await self.transport.receive()
            chunk = message.get("body", b"")
            await self.request.stream.put(chunk)

            more_body = message.get("more_body", False)

        await self.request.stream.put(None)

    async def __call__(self) -> None:
        """
        Handle the incoming request.
        """
        if not self.do_stream:
            self.request.body = await self.read_body()
        else:
            self.sanic_app.loop.create_task(self.stream_body())

        handler = self.sanic_app.handle_request
        callback = None if self.ws else self.stream_callback
        await handler(self.request, None, callback)

    async def stream_callback(self, response: HTTPResponse) -> None:
        """
        Write the response.
        """
        headers: List[Tuple[bytes, bytes]] = []
        cookies: Dict[str, str] = {}
        try:
            cookies = {
                v.key: v
                for _, v in list(
                    filter(
                        lambda item: item[0].lower() == "set-cookie",
                        response.headers.items(),
                    )
                )
            }
            headers += [
                (str(name).encode("latin-1"), str(value).encode("latin-1"))
                for name, value in response.headers.items()
                if name.lower() not in ["set-cookie"]
            ]
        except AttributeError:
            logger.error(
                "Invalid response object for url %s, "
                "Expected Type: HTTPResponse, Actual Type: %s",
                self.request.url,
                type(response),
            )
            exception = ServerError("Invalid response type")
            response = self.sanic_app.error_handler.response(self.request, exception)
            headers = [
                (str(name).encode("latin-1"), str(value).encode("latin-1"))
                for name, value in response.headers.items()
                if name not in (b"Set-Cookie",)
            ]

        if "content-length" not in response.headers and not isinstance(
            response, StreamingHTTPResponse
        ):
            headers += [(b"content-length", str(len(response.body)).encode("latin-1"))]

        if "content-type" not in response.headers:
            headers += [(b"content-type", str(response.content_type).encode("latin-1"))]

        if response.cookies:
            cookies.update(
                {
                    v.key: v
                    for _, v in response.cookies.items()
                    if v.key not in cookies.keys()
                }
            )

        headers += [
            (b"set-cookie", cookie.encode("utf-8")) for k, cookie in cookies.items()
        ]

        await self.transport.send(
            {
                "type": "http.response.start",
                "status": response.status,
                "headers": headers,
            }
        )

        if isinstance(response, StreamingHTTPResponse):
            response.protocol = self.transport.get_protocol()
            await response.stream()
            await response.protocol.complete()

        else:
            await self.transport.send(
                {
                    "type": "http.response.body",
                    "body": response.body,
                    "more_body": False,
                }
            )
<commit_msg>squash<commit_after>import asyncio
import warnings
from inspect import isawaitable
from typing import (
    Any,
    Awaitable,
    Callable,
    Dict,
    List,
    MutableMapping,
    Optional,
    Tuple,
    Union,
)
from urllib.parse import quote

import sanic.app  # noqa
from sanic.compat import Header
from sanic.exceptions import InvalidUsage, ServerError
from sanic.log import logger
from sanic.request import Request
from sanic.response import HTTPResponse, StreamingHTTPResponse
from sanic.server import ConnInfo, StreamBuffer
from sanic.websocket import WebSocketConnection

ASGIScope = MutableMapping[str, Any]
ASGIMessage = MutableMapping[str, Any]
ASGISend = Callable[[ASGIMessage], Awaitable[None]]
ASGIReceive = Callable[[], Awaitable[ASGIMessage]]

class MockProtocol:
    def __init__(self, transport: "MockTransport", loop):
        self.transport = transport
        self._not_paused = asyncio.Event(loop=loop)
        self._not_paused.set()
        self._complete = asyncio.Event(loop=loop)

    def pause_writing(self) -> None:
        self._not_paused.clear()

    def resume_writing(self) -> None:
        self._not_paused.set()

    async def complete(self) -> None:
        self._not_paused.set()
        await self.transport.send(
            {"type": "http.response.body", "body": b"", "more_body": False}
        )

    @property
    def is_complete(self) -> bool:
        return self._complete.is_set()

    async def push_data(self, data: bytes) -> None:
        if not self.is_complete:
            await self.transport.send(
                {"type": "http.response.body", "body": data, "more_body": True}
            )

    async def drain(self) -> None:
        await self._not_paused.wait()

class MockTransport:
    _protocol: Optional[MockProtocol]

    def __init__(self, scope: ASGIScope, receive: ASGIReceive, send: ASGISend) -> None:
        self.scope = scope
        self._receive = receive
        self._send = send
        self._protocol = None
        self.loop = None

    def get_protocol(self) -> MockProtocol:
        if not self._protocol:
            self._protocol = MockProtocol(self, self.loop)
        return self._protocol

    def get_extra_info(self, info: str) -> Union[str, bool, None]:
        if info == "peername":
            return self.scope.get("server")
        elif info == "sslcontext":
            return self.scope.get("scheme") in ["https", "wss"]
        return None

    def get_websocket_connection(self) -> WebSocketConnection:
        try:
            return self._websocket_connection
        except AttributeError:
            raise InvalidUsage("Improper websocket connection.")

    def create_websocket_connection(
        self, send: ASGISend, receive: ASGIReceive
    ) -> WebSocketConnection:
        self._websocket_connection = WebSocketConnection(send, receive)
        return self._websocket_connection

    def add_task(self) -> None:
        raise NotImplementedError

    async def send(self, data) -> None:
        # TODO:
        # - Validation on data and that it is formatted properly and is valid
        await self._send(data)

    async def receive(self) -> ASGIMessage:
        return await self._receive()

class Lifespan:
    def __init__(self, asgi_app: "ASGIApp") -> None:
        self.asgi_app = asgi_app

        if "before_server_start" in self.asgi_app.sanic_app.listeners:
            warnings.warn(
                'You have set a listener for "before_server_start" '
                "in ASGI mode. "
                "It will be executed as early as possible, but not before "
                "the ASGI server is started."
            )
        if "after_server_stop" in self.asgi_app.sanic_app.listeners:
            warnings.warn(
                'You have set a listener for "after_server_stop" '
                "in ASGI mode. "
                "It will be executed as late as possible, but not after "
                "the ASGI server is stopped."
            )

    async def startup(self) -> None:
        """
        Gather the listeners to fire on server start.
        Because we are using a third-party server and not Sanic server, we do
        not have access to fire anything BEFORE the server starts.
        Therefore, we fire before_server_start and after_server_start
        in sequence since the ASGI lifespan protocol only supports a single
        startup event.
        """
        listeners = self.asgi_app.sanic_app.listeners.get(
            "before_server_start", []
        ) + self.asgi_app.sanic_app.listeners.get("after_server_start", [])

        for handler in listeners:
            response = handler(self.asgi_app.sanic_app, self.asgi_app.sanic_app.loop)
            if isawaitable(response):
                await response

    async def shutdown(self) -> None:
        """
        Gather the listeners to fire on server stop.
        Because we are using a third-party server and not Sanic server, we do
        not have access to fire anything AFTER the server stops.
        Therefore, we fire before_server_stop and after_server_stop
        in sequence since the ASGI lifespan protocol only supports a single
        shutdown event.
        """
        listeners = self.asgi_app.sanic_app.listeners.get(
            "before_server_stop", []
        ) + self.asgi_app.sanic_app.listeners.get("after_server_stop", [])

        for handler in listeners:
            response = handler(self.asgi_app.sanic_app, self.asgi_app.sanic_app.loop)
            if isawaitable(response):
                await response

    async def __call__(
        self, scope: ASGIScope, receive: ASGIReceive, send: ASGISend
    ) -> None:
        message = await receive()
        if message["type"] == "lifespan.startup":
            await self.startup()
            await send({"type": "lifespan.startup.complete"})

        message = await receive()
        if message["type"] == "lifespan.shutdown":
            await self.shutdown()
            await send({"type": "lifespan.shutdown.complete"})

class ASGIApp:
    sanic_app: "sanic.app.Sanic"
    request: Request
    transport: MockTransport
    do_stream: bool
    lifespan: Lifespan
    ws: Optional[WebSocketConnection]

    def __init__(self) -> None:
        self.ws = None

    @classmethod
    async def create(
        cls, sanic_app, scope: ASGIScope, receive: ASGIReceive, send: ASGISend
    ) -> "ASGIApp":
        instance = cls()
        instance.sanic_app = sanic_app
        instance.transport = MockTransport(scope, receive, send)
        instance.transport.loop = sanic_app.loop
        setattr(instance.transport, "add_task", sanic_app.loop.create_task)

        headers = Header(
            [
                (key.decode("latin-1"), value.decode("latin-1"))
                for key, value in scope.get("headers", [])
            ]
        )
        instance.do_stream = True if headers.get("expect") == "100-continue" else False
        instance.lifespan = Lifespan(instance)

        if scope["type"] == "lifespan":
            await instance.lifespan(scope, receive, send)
        else:
            path = scope["path"][1:] if scope["path"].startswith("/") else scope["path"]
            url = "/".join([scope.get("root_path", ""), quote(path)])
            url_bytes = url.encode("latin-1")
            url_bytes += b"?" + scope["query_string"]

            if scope["type"] == "http":
                version = scope["http_version"]
                method = scope["method"]
            elif scope["type"] == "websocket":
                version = "1.1"
                method = "GET"

                instance.ws = instance.transport.create_websocket_connection(
                    send, receive
                )
                await instance.ws.accept()
            else:
                pass
                # TODO:
                # - close connection

            request_class = sanic_app.request_class or Request
            instance.request = request_class(
                url_bytes, headers, version, method, instance.transport, sanic_app,
            )
            instance.request.conn_info = ConnInfo(instance.transport)

            if sanic_app.is_request_stream:
                is_stream_handler = sanic_app.router.is_stream_handler(instance.request)
                if is_stream_handler:
                    instance.request.stream = StreamBuffer(
                        sanic_app.config.REQUEST_BUFFER_QUEUE_SIZE
                    )
                    instance.do_stream = True

        return instance

    async def read_body(self) -> bytes:
        """
        Read and return the entire body from an incoming ASGI message.
        """
        body = b""
        more_body = True
        while more_body:
            message = await self.transport.receive()
            body += message.get("body", b"")
            more_body = message.get("more_body", False)

        return body

    async def stream_body(self) -> None:
        """
        Read and stream the body in chunks from an incoming ASGI message.
        """
        more_body = True

        while more_body:
            message = await self.transport.receive()
            chunk = message.get("body", b"")
            await self.request.stream.put(chunk)

            more_body = message.get("more_body", False)

        await self.request.stream.put(None)

    async def __call__(self) -> None:
        """
        Handle the incoming request.
        """
        if not self.do_stream:
            self.request.body = await self.read_body()
        else:
            self.sanic_app.loop.create_task(self.stream_body())

        handler = self.sanic_app.handle_request
        callback = None if self.ws else self.stream_callback
        await handler(self.request, None, callback)

    async def stream_callback(self, response: HTTPResponse) -> None:
        """
        Write the response.
        """
        headers: List[Tuple[bytes, bytes]] = []
        cookies: Dict[str, str] = {}
        try:
            cookies = {
                v.key: v
                for _, v in list(
                    filter(
                        lambda item: item[0].lower() == "set-cookie",
                        response.headers.items(),
                    )
                )
            }
            headers += [
                (str(name).encode("latin-1"), str(value).encode("latin-1"))
                for name, value in response.headers.items()
                if name.lower() not in ["set-cookie"]
            ]
        except AttributeError:
            logger.error(
                "Invalid response object for url %s, "
                "Expected Type: HTTPResponse, Actual Type: %s",
                self.request.url,
                type(response),
            )
            exception = ServerError("Invalid response type")
            response = self.sanic_app.error_handler.response(self.request, exception)
            headers = [
                (str(name).encode("latin-1"), str(value).encode("latin-1"))
                for name, value in response.headers.items()
                if name not in (b"Set-Cookie",)
            ]

        if "content-length" not in response.headers and not isinstance(
            response, StreamingHTTPResponse
        ):
            headers += [(b"content-length", str(len(response.body)).encode("latin-1"))]

        if "content-type" not in response.headers:
            headers += [(b"content-type", str(response.content_type).encode("latin-1"))]

        if response.cookies:
            cookies.update(
                {
                    v.key: v
                    for _, v in response.cookies.items()
                    if v.key not in cookies.keys()
                }
            )

        headers += [
            (b"set-cookie", cookie.encode("utf-8")) for k, cookie in cookies.items()
        ]

        await self.transport.send(
            {
                "type": "http.response.start",
                "status": response.status,
                "headers": headers,
            }
        )

        if isinstance(response, StreamingHTTPResponse):
            response.protocol = self.transport.get_protocol()
            await response.stream()
            await response.protocol.complete()

        else:
            await self.transport.send(
                {
                    "type": "http.response.body",
                    "body": response.body,
                    "more_body": False,
                }
            )
<|endoftext|>


<commit_before>from matplotlib import pyplot as plt
from matplotlib.patches import Wedge
import numpy as np

def path_figure(theta_matrix, robot_arm, show=True):
    """
    Arguments:
    theta_matrix - A set of theta column vectors
    robot_arm - An object of the RobotArm class

    Returns:
    None, but plots the configuration of each theta vector as subplots
    """
    # Check input arguments
    num_of_destinations = robot_arm.destinations.shape[1]
    if not theta_matrix.shape == (robot_arm.n, num_of_destinations):
        raise ValueError('''
                        The number of joint positions does not match the
                         number of destination points
                         ''')

    # Set up plot style options
    plt.style.use('ggplot')
    fig, axes = plt.subplots(nrows=1, ncols=num_of_destinations)
    for ax in np.ravel(axes):
        set_axis_options(ax, robot_arm)

    # Plotting content of each subplot
    for index, theta in enumerate(theta_matrix.T):
        plot_position(axes[index], theta, robot_arm)

    if show is True:
        plt.show()

    return fig

def set_axis_options(ax, robot_arm):
    ax.set_autoscale_on(False)

    ax.axhline(y=0, color='grey')
    ax.axvline(x=0, color='grey')

    # Padding
    a = 1.1
    max_x = abs(max(robot_arm.destinations, key=lambda p: abs(p[0]))[0])
    max_y = abs(max(robot_arm.destinations, key=lambda p: abs(p[1]))[1])
    m = max(max_x, max_y, robot_arm.reach)

    ax.set_xlim(-a * m, a * m)
    ax.set_ylim(-a * m, a * m)

def plot_position(axis, theta, robot_arm):
    joint_positions = robot_arm.joint_positions(theta)
    x = np.hstack((0, joint_positions[0, :]))
    y = np.hstack((0, joint_positions[1, :]))
    axis.plot(x, y, '-o')

    # Plot all the points that shall be reached
    for p in robot_arm.destinations.T:
        axis.plot(p[0], p[1], 'x')

    # Plot configuration space of robot
    configuration_space = Wedge(
        (0, 0),
        r=robot_arm.reach,
        theta1=0,
        theta2=360,
        width=robot_arm.reach - robot_arm.inner_reach,
        facecolor='grey',
        alpha=0.3,
        edgecolor='black',
        linewidth=0.6
    )

    axis.add_patch(configuration_space)
<commit_msg>Enumerate destinations<commit_after>from matplotlib import pyplot as plt
from matplotlib.patches import Wedge
import numpy as np

def path_figure(theta_matrix, robot_arm, show=True):
    """
    Arguments:
    theta_matrix - A set of theta column vectors
    robot_arm - An object of the RobotArm class

    Returns:
    None, but plots the configuration of each theta vector as subplots
    """
    # Check input arguments
    num_of_destinations = robot_arm.destinations.shape[1]
    if not theta_matrix.shape == (robot_arm.n, num_of_destinations):
        raise ValueError('''
                        The number of joint positions does not match the
                         number of destination points
                         ''')

    # Set up plot style options
    plt.style.use('ggplot')
    fig, axes = plt.subplots(nrows=1, ncols=num_of_destinations)
    for ax in np.ravel(axes):
        set_axis_options(ax, robot_arm)

    # Plotting content of each subplot
    for index, theta in enumerate(theta_matrix.T):
        plot_position(axes[index], theta, robot_arm)

    if show is True:
        plt.show()

    return fig

def set_axis_options(ax, robot_arm):
    ax.set_autoscale_on(False)

    ax.axhline(y=0, color='grey')
    ax.axvline(x=0, color='grey')

    # Padding
    a = 1.1
    max_x = abs(max(robot_arm.destinations, key=lambda p: abs(p[0]))[0])
    max_y = abs(max(robot_arm.destinations, key=lambda p: abs(p[1]))[1])
    m = max(max_x, max_y, robot_arm.reach)

    ax.set_xlim(-a * m, a * m)
    ax.set_ylim(-a * m, a * m)

def plot_position(axis, theta, robot_arm):
    joint_positions = robot_arm.joint_positions(theta)
    x = np.hstack((0, joint_positions[0, :]))
    y = np.hstack((0, joint_positions[1, :]))
    axis.plot(x, y, '-o')

    # Plot all the points that shall be reached
    for index, p in enumerate(robot_arm.destinations.T):
        point, = axis.plot(p[0], p[1], 'x')
        axis.text(p[0], p[1], str(index + 1), fontsize=14, color=point.get_color())

    # Plot configuration space of robot
    configuration_space = Wedge(
        (0, 0),
        r=robot_arm.reach,
        theta1=0,
        theta2=360,
        width=robot_arm.reach - robot_arm.inner_reach,
        facecolor='grey',
        alpha=0.3,
        edgecolor='black',
        linewidth=0.6
    )

    axis.add_patch(configuration_space)
<|endoftext|>


<commit_before>import numpy as np
import mne
from warnings import warn
from sarna.utils import group, _invert_selection, _transfer_selection_to_raw
# from numba import jit

def dB(x):
    return 10 * np.log10(x)

# - [ ] add detrending (1 + x + 1/x) or FOOOF cooperation
def transform_spectrum(spectrum, dB=False, normalize=False, detrend=False):
    """Common spectrum transformations.

    Parameters
    ----------
    spectrum : numpy array
        channels x frequencies
    """

    if dB:
        spectrum = 10 * np.log10(spectrum)
    if normalize:
        if dB:
            # move whole spectrum up, so that normalization
            # does not return weird results
            min_val = spectrum.min() - 0.01
            spectrum -= min_val
        spectrum /= spectrum.sum(axis=1)[:, np.newaxis]
    return spectrum

# - [ ] consider moving to utils
# - [x] warn if sfreq not given and some values are float
# - [x] treat floats as time and int as samples
# - [ ] maybe a smarter API or a class...
def window_steps(window_length, window_step, signal_len, sfreq=None):
    is_float = [isinstance(x, float)
                for x in [window_length, window_step, signal_len]]
    any_float = any(is_float)
    if any_float and sfreq is None:
        raise TypeError('Some variables are float but sfreq was not given.')

    if any_float and sfreq is not None:
        if is_float[0]:
            window_length = int(np.round(window_length * sfreq))
        if is_float[1]:
            window_step = int(np.round(window_step * sfreq))
        if is_float[2]:
            signal_len = int(np.round(signal_len * sfreq))

    num_steps = int(np.floor((signal_len - window_length) / window_step)) + 1
    for w in range(num_steps):
        yield slice(w * window_step, window_length + w * window_step)

def plot_topo_and_psd(inst, mean_psd, freqs, channels):
    from matplotlib import gridspec
    import matplotlib.pyplot as plt
    from mne.viz.topomap import plot_psds_topomap

    fig = plt.figure(figsize=(8, 3))
    gs = gridspec.GridSpec(1, 2, width_ratios=[1, 2])
    ax = [plt.subplot(g) for g in gs]

    plot_psds_topomap(psds=mean_psd, freqs=freqs, pos=inst.info,
                      dB=False, axes=[ax[0]], bands=[(4., 8., 'Theta')],
                      normalize=False, cmap='inferno', show=False)

    # highlight channels
    circles = ax[0].findobj(plt.Circle)
    for ch in channels:
        circles[ch].set_color('r')
        circles[ch].set_radius(0.025)

    plot_freq = (freqs > 1.) & (freqs < 15.)
    ax[1].plot(freqs[plot_freq], mean_psd[channels, :][:, plot_freq].T)
    chan_avg = mean_psd[channels, :].mean(axis=0)
    ax[1].plot(freqs[plot_freq], chan_avg[plot_freq], color='k', lw=2)
    return fig

def _correct_overlap(periods):
    '''

    Parameters
    ----------
    periods : np.ndarray
        Numpy array of (n_periods, 3) shape. The columns are: epoch index,
        within-epoch sample index of period start, within-epoch sample index of
        period end.

    Returns
    -------
    periods : np.ndarray
        Corrected numpy array of (n_periods, 3) shape. The columns are: epoch
        index, within-epoch sample index of period start, within-epoch sample
        index of period end.

    '''
    n_rows = periods.shape[0]
    current_period = periods[0, :].copy()
    correct = list()

    for idx in range(1, n_rows):
        overlap = ((periods[idx, 0] == current_period[0])
                   and (periods[idx, 1] <= current_period[2]))

        if overlap:
            current_period[-1] = periods[idx, -1]
        else:
            correct.append(current_period)
            current_period = periods[idx, :].copy()

    correct.append(current_period)
    periods = np.stack(correct, axis=0)

    return periods

def _find_sel_amplitude_periods(epochs, threshold=2.5, min_period=0.1,
                                 periods = 'high', extend=None):
    '''
    Find segments of high or low amplitude in filtered, hilbert-transformed signal.

    Parameters
    ----------
    epochs : mne.Epochs
        Epoched data. Must be filtered and hilbert-transformed.
    threshold : float, str
        Threshold defining high amplitude periods to select: if float, it
        is interpreted as a z value threshold; if str, as percentage of
        fragments with in the highest amplitude in form of ``'xx%'``
        (for example with ``'25%'`` 25% of singal with highest amplitude will
        be selected ). Defaults to ``2.5``.
    min_period : float
        Minimum length of high amplitude period in seconds.
        Defaults to ``0.1``.
    periods : str
        Sepcification of perionds to find. Might be 'high' or 'low' amplitude.
    extend : float | None
        Extend each period by this many seconds on both sides (before and
        after). Defaults to ``None`` which does not extend the periods.

    Returns
    -------
    periods : np.ndarray
        Numpy array of (n_periods, 3) shape. The columns are: epoch index,
        within-epoch sample index of period start, within-epoch sample index of
        period end.
    '''
    from scipy.stats import zscore

    # amplitude periods
    n_epochs, n_channels, n_samples = epochs._data.shape
    comp_data = epochs._data.transpose([1, 0, 2]).reshape((n_channels, -1))

    # find segments with elevated amplitude
    envelope = np.nanmean(comp_data, axis=0)

    if isinstance(threshold, str) and '%' in threshold:
        perc = 100 - float(threshold.replace('%', ''))
        threshold = np.nanpercentile(envelope, perc)
    else:
        envelope = zscore(envelope, nan_policy='omit')

    if periods == 'high':
        grp = group(envelope > threshold)
    else:
        grp = group(envelope < threshold)

    if len(grp) == 0:
        raise ValueError('No high amplitude periods were found.')
    # check if there are some segments that start at one epoch
    # and end in another
    # -> if so, they could be split, but we will ignore them for now
    epoch_idx = np.floor(grp / n_samples)
    epoch_diff = np.diff(epoch_idx, axis=0)
    epochs_joint = epoch_diff > 0
    if epochs_joint.any():
        msg = ('{:d} high-amplitude segments will be ignored because'
               ' the developer was lazy.')
        warn(msg.format(epochs_joint.sum()))
        epoch_diff = epoch_diff[~epochs_joint]

    segment_len = np.diff(grp, axis=1)
    good_length = segment_len[:, 0] * (1 / epochs.info['sfreq']) > min_period
    grp = grp[good_length, :]
    epoch_idx = np.floor(grp[:, [0]] / n_samples).astype('int')
    grp -= epoch_idx * n_samples

    if extend is not None:
        extend_samples = int(np.round(extend * epochs.info['sfreq']))
        extend_samples = np.array([-extend_samples, extend_samples])
        grp += extend_samples[np.newaxis, :]

        # check for limits
        msk1 = grp[:, 0] < 0
        msk2 = grp[:, 1] >= n_samples
        grp[msk1, 0] = 0
        grp[msk2, 1] = n_samples - 1

    periods = np.append(epoch_idx, grp, axis=1)
    periods = periods if extend is None else _correct_overlap(periods)

    return periods

def create_amplitude_annotations(raw, freq=None, events=None, event_id=None,
                                 picks=None, tmin=-0.2, tmax=0.5,
                                 threshold=2., min_period=0.1, periods = 'high',
                                 extend=None):
    '''
    Parameters
    ----------
    raw : mne.Raw
        Raw file to use.
    events: numpy array | None
        Mne events array of shape (n_events, 3). If None (default) `tmin` and
        `tmax` are not calculated with respect to events but the whole time
        range of the `raw` file.
    event_id: list | numpy array
        Event types (IDs) to use in defining segments for which psd is
        computed. If None (default) and events were passed all event types are
        used.
    freq : list | numpy array
        Frequency limits defining a range for which low amplitude periods will
        be calculated.
    picks : list
        List of channels for which low amplitude periods will be calculated.
    tmin : float
        Start time before event.
    tmax : float
        End time after event.
    threshold : float, str
        Threshold defining high amplitude periods to select: if float, it
        is interpreted as a z value threshold; if str, as percentage of
        fragments with in the highest amplitude in form of ``'xx%'``
        (for example with ``'25%'`` 25% of singal with highest amplitude will
        be selected ). Defaults to ``2.5``.
    min_period : float
        Minimum length of high amplitude period in seconds.
        Defaults to ``0.1``.
    periods : str
        Sepcification of perionds to find. Might be 'high' or 'low' amplitude.
    extend : float | None
        Extend each period by this many seconds on both sides (before and
        after). Defaults to ``None`` which does not extend the periods.

    Returns
    -------
    raw_annot : mne.Raw
        Raw files with annotations.
    '''

    if freq is None:
        raise TypeError('Frequencies have to be defined')
    if events is None:
        raise TypeError('Events have to be defined')
    if event_id is None:
        event_id = np.unique(events[:, 2]).tolist()

    filt_raw = raw.copy().filter(freq[0], freq[1])

    filt_raw_nan = filt_raw.copy()
    filt_raw_nan._data = raw.get_data(reject_by_annotation='NaN')

    epochs_nan = mne.Epochs(filt_raw_nan, events=events, event_id=event_id,
                            tmin=tmin, tmax=tmax, baseline=None, preload=True,
                            reject_by_annotation=False)
    epochs = mne.Epochs(filt_raw, events=events, event_id=event_id, tmin=tmin,
                        tmax=tmax, baseline=None, preload=True,
                        reject_by_annotation=False)

    del filt_raw_nan

    filt_hilb_data = np.abs(epochs.copy().pick(picks).apply_hilbert()._data)
    filt_hilb_data[np.isnan(epochs_nan.copy().pick(picks)._data)] = np.nan
    epochs_nan._data = filt_hilb_data

    hi_amp_epochs = _find_sel_amplitude_periods(epochs_nan,
                                                 threshold=threshold,
                                                 min_period=min_period,
                                                 extend=extend)

    hi_amp_raw = _transfer_selection_to_raw(epochs, raw,
                                            selection=hi_amp_epochs)
    amp_inv_samples = _invert_selection(raw, selection=hi_amp_raw)

    sfreq = raw.info['sfreq']
    amp_inv_annot_sec = amp_inv_samples / sfreq

    n_segments = amp_inv_samples.shape[0]
    if periods == 'high':
        amp_annot = mne.Annotations(amp_inv_annot_sec[:, 0],
                                    amp_inv_annot_sec[:, 1],
                                    ['BAD_lowamp'] * n_segments)
    else:
        amp_annot = mne.Annotations(amp_inv_annot_sec[:, 0],
                                    amp_inv_annot_sec[:, 1],
                                    ['BAD_highamp'] * n_segments)
    return amp_annot

def grand_average_psd(psd_list):
    '''Perform grand average on a list of PSD objects.

    Parameters
    ----------
    psd_list : list
        List of ``borsar.freq.PSD`` objects.

    Returns
    -------
    grand_psd : borsar.freq.PSD
        Grand averaged spectrum.
    '''
    assert isinstance(psd_list, list)

    # make sure that all psds have the same number and order of channels
    # and the same frequencies
    freq1 = psd_list[0].freqs
    ch_names1 = psd_list[0].ch_names
    n_channels = len(ch_names1)
    for psd in psd_list:
        assert len(psd.freqs) == len(freq1)
        assert (psd.freqs == freq1).all()
        assert len(psd.ch_names) == n_channels
        assert all([ch_names1[idx] == psd.ch_names[idx]
                    for idx in range(n_channels)])

    all_psds = list()
    for this_psd in psd_list:
        # upewniamy się, że epoki są uśrednione
        this_psd = this_psd.copy().average()
        all_psds.append(this_psd.data)

    # łączymy widma w macierz (osoby x kanały x częstotliwości)
    all_psds = np.stack(all_psds, axis=0)

    # uśredniamy wymiar osób, zostają nam kanały x częstotliwości
    all_psds = all_psds.mean(axis=0)

    # kopiujemy wybrane psd z listy
    grand_psd = this_psd.copy()
    # i wypełniamy wartościamy średniej po osobach
    grand_psd._data = all_psds

    return grand_psd
<commit_msg>ENH: Apply suggestions from code review<commit_after>import numpy as np
import mne
from warnings import warn
from sarna.utils import group, _invert_selection, _transfer_selection_to_raw
# from numba import jit

def dB(x):
    return 10 * np.log10(x)

# - [ ] add detrending (1 + x + 1/x) or FOOOF cooperation
def transform_spectrum(spectrum, dB=False, normalize=False, detrend=False):
    """Common spectrum transformations.

    Parameters
    ----------
    spectrum : numpy array
        channels x frequencies
    """

    if dB:
        spectrum = 10 * np.log10(spectrum)
    if normalize:
        if dB:
            # move whole spectrum up, so that normalization
            # does not return weird results
            min_val = spectrum.min() - 0.01
            spectrum -= min_val
        spectrum /= spectrum.sum(axis=1)[:, np.newaxis]
    return spectrum

# - [ ] consider moving to utils
# - [x] warn if sfreq not given and some values are float
# - [x] treat floats as time and int as samples
# - [ ] maybe a smarter API or a class...
def window_steps(window_length, window_step, signal_len, sfreq=None):
    is_float = [isinstance(x, float)
                for x in [window_length, window_step, signal_len]]
    any_float = any(is_float)
    if any_float and sfreq is None:
        raise TypeError('Some variables are float but sfreq was not given.')

    if any_float and sfreq is not None:
        if is_float[0]:
            window_length = int(np.round(window_length * sfreq))
        if is_float[1]:
            window_step = int(np.round(window_step * sfreq))
        if is_float[2]:
            signal_len = int(np.round(signal_len * sfreq))

    num_steps = int(np.floor((signal_len - window_length) / window_step)) + 1
    for w in range(num_steps):
        yield slice(w * window_step, window_length + w * window_step)

def plot_topo_and_psd(inst, mean_psd, freqs, channels):
    from matplotlib import gridspec
    import matplotlib.pyplot as plt
    from mne.viz.topomap import plot_psds_topomap

    fig = plt.figure(figsize=(8, 3))
    gs = gridspec.GridSpec(1, 2, width_ratios=[1, 2])
    ax = [plt.subplot(g) for g in gs]

    plot_psds_topomap(psds=mean_psd, freqs=freqs, pos=inst.info,
                      dB=False, axes=[ax[0]], bands=[(4., 8., 'Theta')],
                      normalize=False, cmap='inferno', show=False)

    # highlight channels
    circles = ax[0].findobj(plt.Circle)
    for ch in channels:
        circles[ch].set_color('r')
        circles[ch].set_radius(0.025)

    plot_freq = (freqs > 1.) & (freqs < 15.)
    ax[1].plot(freqs[plot_freq], mean_psd[channels, :][:, plot_freq].T)
    chan_avg = mean_psd[channels, :].mean(axis=0)
    ax[1].plot(freqs[plot_freq], chan_avg[plot_freq], color='k', lw=2)
    return fig

def _correct_overlap(periods):
    '''

    Parameters
    ----------
    periods : np.ndarray
        Numpy array of (n_periods, 3) shape. The columns are: epoch index,
        within-epoch sample index of period start, within-epoch sample index of
        period end.

    Returns
    -------
    periods : np.ndarray
        Corrected numpy array of (n_periods, 3) shape. The columns are: epoch
        index, within-epoch sample index of period start, within-epoch sample
        index of period end.

    '''
    n_rows = periods.shape[0]
    current_period = periods[0, :].copy()
    correct = list()

    for idx in range(1, n_rows):
        overlap = ((periods[idx, 0] == current_period[0])
                   and (periods[idx, 1] <= current_period[2]))

        if overlap:
            current_period[-1] = periods[idx, -1]
        else:
            correct.append(current_period)
            current_period = periods[idx, :].copy()

    correct.append(current_period)
    periods = np.stack(correct, axis=0)

    return periods

def _find_sel_amplitude_periods(epochs, threshold=2.5, min_period=0.1,
                                 periods='high', extend=None):
    '''
    Find segments of high or low amplitude in filtered, hilbert-transformed signal.

    Parameters
    ----------
    epochs : mne.Epochs
        Epoched data. Must be filtered and hilbert-transformed.
    threshold : float, str
        Threshold defining high amplitude periods to select: if float, it
        is interpreted as a z value threshold; if str, as percentage of
        fragments with in the highest amplitude in form of ``'xx%'``
        (for example with ``'25%'`` 25% of singal with highest amplitude will
        be selected ). Defaults to ``2.5``.
    min_period : float
        Minimum length of high amplitude period in seconds.
        Defaults to ``0.1``.
    periods : str
        Sepcification of perionds to find. Might be 'high' or 'low' amplitude.
    extend : float | None
        Extend each period by this many seconds on both sides (before and
        after). Defaults to ``None`` which does not extend the periods.

    Returns
    -------
    periods : np.ndarray
        Numpy array of (n_periods, 3) shape. The columns are: epoch index,
        within-epoch sample index of period start, within-epoch sample index of
        period end.
    '''
    from scipy.stats import zscore

    # amplitude periods
    n_epochs, n_channels, n_samples = epochs._data.shape
    comp_data = epochs._data.transpose([1, 0, 2]).reshape((n_channels, -1))

    # find segments with elevated amplitude
    envelope = np.nanmean(comp_data, axis=0)

    if isinstance(threshold, str) and '%' in threshold:
        perc = 100 - float(threshold.replace('%', ''))
        threshold = np.nanpercentile(envelope, perc)
    else:
        envelope = zscore(envelope, nan_policy='omit')

    if periods == 'high':
        grp = group(envelope > threshold)
    elif periods == 'low':
        grp = group(envelope < threshold)
    else:
        raise ValueError('Unrecognised `periods` option "{}".'.format(periods))

    if len(grp) == 0:
        raise ValueError('No {} amplitude periods were found.'.format(periods))
    # check if there are some segments that start at one epoch
    # and end in another
    # -> if so, they could be split, but we will ignore them for now
    epoch_idx = np.floor(grp / n_samples)
    epoch_diff = np.diff(epoch_idx, axis=0)
    epochs_joint = epoch_diff > 0
    if epochs_joint.any():
        msg = ('{:d} high-amplitude segments will be ignored because'
               ' the developer was lazy.')
        warn(msg.format(epochs_joint.sum()))
        epoch_diff = epoch_diff[~epochs_joint]

    segment_len = np.diff(grp, axis=1)
    good_length = segment_len[:, 0] * (1 / epochs.info['sfreq']) > min_period
    grp = grp[good_length, :]
    epoch_idx = np.floor(grp[:, [0]] / n_samples).astype('int')
    grp -= epoch_idx * n_samples

    if extend is not None:
        extend_samples = int(np.round(extend * epochs.info['sfreq']))
        extend_samples = np.array([-extend_samples, extend_samples])
        grp += extend_samples[np.newaxis, :]

        # check for limits
        msk1 = grp[:, 0] < 0
        msk2 = grp[:, 1] >= n_samples
        grp[msk1, 0] = 0
        grp[msk2, 1] = n_samples - 1

    periods = np.append(epoch_idx, grp, axis=1)
    periods = periods if extend is None else _correct_overlap(periods)

    return periods

def create_amplitude_annotations(raw, freq=None, events=None, event_id=None,
                                 picks=None, tmin=-0.2, tmax=0.5,
                                 threshold=2., min_period=0.1, periods='high',
                                 extend=None):
    '''
    Parameters
    ----------
    raw : mne.Raw
        Raw file to use.
    events: numpy array | None
        Mne events array of shape (n_events, 3). If None (default) `tmin` and
        `tmax` are not calculated with respect to events but the whole time
        range of the `raw` file.
    event_id: list | numpy array
        Event types (IDs) to use in defining segments for which psd is
        computed. If None (default) and events were passed all event types are
        used.
    freq : list | numpy array
        Frequency limits defining a range for which low amplitude periods will
        be calculated.
    picks : list
        List of channels for which low amplitude periods will be calculated.
    tmin : float
        Start time before event.
    tmax : float
        End time after event.
    threshold : float, str
        Threshold defining high amplitude periods to select: if float, it
        is interpreted as a z value threshold; if str, as percentage of
        fragments with in the highest amplitude in form of ``'xx%'``
        (for example with ``'25%'`` 25% of singal with highest amplitude will
        be selected ). Defaults to ``2.5``.
    min_period : float
        Minimum length of high amplitude period in seconds.
        Defaults to ``0.1``.
    periods : str
        Sepcification of perionds to find. Might be 'high' or 'low' amplitude.
    extend : float | None
        Extend each period by this many seconds on both sides (before and
        after). Defaults to ``None`` which does not extend the periods.

    Returns
    -------
    raw_annot : mne.Raw
        Raw files with annotations.
    '''

    if freq is None:
        raise TypeError('Frequencies have to be defined')
    if events is None:
        raise TypeError('Events have to be defined')
    if event_id is None:
        event_id = np.unique(events[:, 2]).tolist()

    filt_raw = raw.copy().filter(freq[0], freq[1])

    filt_raw_nan = filt_raw.copy()
    filt_raw_nan._data = raw.get_data(reject_by_annotation='NaN')

    epochs_nan = mne.Epochs(filt_raw_nan, events=events, event_id=event_id,
                            tmin=tmin, tmax=tmax, baseline=None, preload=True,
                            reject_by_annotation=False)
    epochs = mne.Epochs(filt_raw, events=events, event_id=event_id, tmin=tmin,
                        tmax=tmax, baseline=None, preload=True,
                        reject_by_annotation=False)

    del filt_raw_nan

    filt_hilb_data = np.abs(epochs.copy().pick(picks).apply_hilbert()._data)
    filt_hilb_data[np.isnan(epochs_nan.copy().pick(picks)._data)] = np.nan
    epochs_nan._data = filt_hilb_data

    hi_amp_epochs = _find_sel_amplitude_periods(epochs_nan,
                                                 threshold=threshold,
                                                 min_period=min_period,
                                                 extend=extend)

    hi_amp_raw = _transfer_selection_to_raw(epochs, raw,
                                            selection=hi_amp_epochs)
    amp_inv_samples = _invert_selection(raw, selection=hi_amp_raw)

    sfreq = raw.info['sfreq']
    amp_inv_annot_sec = amp_inv_samples / sfreq

    n_segments = amp_inv_samples.shape[0]
    annot_name = 'BAD_{}amp'.format(periods)
    amp_annot = mne.Annotations(amp_inv_annot_sec[:, 0],
                                amp_inv_annot_sec[:, 1],
                                annot_name * n_segments)
    return amp_annot

def grand_average_psd(psd_list):
    '''Perform grand average on a list of PSD objects.

    Parameters
    ----------
    psd_list : list
        List of ``borsar.freq.PSD`` objects.

    Returns
    -------
    grand_psd : borsar.freq.PSD
        Grand averaged spectrum.
    '''
    assert isinstance(psd_list, list)

    # make sure that all psds have the same number and order of channels
    # and the same frequencies
    freq1 = psd_list[0].freqs
    ch_names1 = psd_list[0].ch_names
    n_channels = len(ch_names1)
    for psd in psd_list:
        assert len(psd.freqs) == len(freq1)
        assert (psd.freqs == freq1).all()
        assert len(psd.ch_names) == n_channels
        assert all([ch_names1[idx] == psd.ch_names[idx]
                    for idx in range(n_channels)])

    all_psds = list()
    for this_psd in psd_list:
        # upewniamy się, że epoki są uśrednione
        this_psd = this_psd.copy().average()
        all_psds.append(this_psd.data)

    # łączymy widma w macierz (osoby x kanały x częstotliwości)
    all_psds = np.stack(all_psds, axis=0)

    # uśredniamy wymiar osób, zostają nam kanały x częstotliwości
    all_psds = all_psds.mean(axis=0)

    # kopiujemy wybrane psd z listy
    grand_psd = this_psd.copy()
    # i wypełniamy wartościamy średniej po osobach
    grand_psd._data = all_psds

    return grand_psd
<|endoftext|>


<commit_before>from random import shuffle

from model.card import Card
from model.player import Player

class Game:
    def __init__(self, players):
        self.players = players
        self.card_deck = list()
        self.skat = list()
        self.create_deck()

    def create_deck(self):
        for suit in Card.Suit.__members__.items():
            for face in Card.Face.__members__.items():
                self.card_deck.append(Card(suit, face))

    def start_new(self):
        self.clear_cards()
        shuffle(self.card_deck)
        self.give_out_cards()

    def clear_cards(self):
        del self.skat[:]
        for player in self.players:
            del player.cards[:]

    def give_out_cards(self):
        # give cards P1:3, P2:3, P3:3, S:2, P1:4, P2:4, P3:4, P1:3, P2:3, P3:3 --> skat alike
        # P1 gets: 1, 2, 3, 12, 13, 14, 15, 24, 25, 26
        # P2 gets: 4, 5, 6, 16, 17, 18, 19, 27, 28, 29
        # P3 gets: 7, 8, 9, 20, 21, 22, 23, 30, 31, 32
        # Skat gets 10, 11
        round1 = (0, 2)
        round2 = (11, 14)
        round3 = (23, 25)
        for idx, card in enumerate(self.card_deck):
            # Player 1
            if (round1[0] <= idx <= round1[1]) \
                    or (round2[0] <= idx <= round2[1]) \
                    or (round3[0] <= idx <= round3[1]):
                self.players[0].cards.append(card)
            # Player 2
            elif (round1[0] + 3 <= idx <= round1[1] + 3) \
                    or (round2[0] + 4 <= idx <= round2[1] + 4) \
                    or (round3[0] + 3 <= idx <= round3[1] + 3):
                self.players[1].cards.append(card)
            # Player 3
            elif (round1[0] + 6 <= idx <= round1[1] + 6) \
                    or (round2[0] + 8 <= idx <= round2[1] + 8) \
                    or (round3[0] + 6 <= idx <= round3[1] + 6):
                self.players[2].cards.append(card)
            # Skat
            else:
                self.skat.append(card)

game = Game([Player("Player1"), Player("Player2"), Player("Player3")])
game.start_new()

print("Skat:")
for card in game.skat:
    print(card)

print("\n" + game.players[0].name + ":")
for card in game.players[0].cards:
    print(card)

print("\n" + game.players[1].name + ":")
for card in game.players[1].cards:
    print(card)

print("\n" + game.players[2].name + ":")
for card in game.players[2].cards:
    print(card)<commit_msg>add dealer and seats information<commit_after>from random import shuffle

from model.card import Card
from model.player import Player

PLAYER_COUNT = 3

class Game:
    def __init__(self, players):
        self.players = players
        self.card_deck = list()
        self.skat = list()
        self.dealer = -1

        self.create_deck()

    def get_dealer(self):
        return self.players[self.dealer]

    def get_first_seat(self):
        return self.players[(self.dealer + 1) % PLAYER_COUNT]

    def get_second_seat(self):
        return self.players[(self.dealer + 2) % PLAYER_COUNT]

    def get_third_seat(self):
        return self.players[(self.dealer + 3) % PLAYER_COUNT]

    def create_deck(self):
        for suit in Card.Suit.__members__.items():
            for face in Card.Face.__members__.items():
                self.card_deck.append(Card(suit, face))

    def start_new(self):
        self.clear_cards()
        shuffle(self.card_deck)
        self.give_out_cards()
        self.dealer = (self.dealer + 1) % PLAYER_COUNT

    def clear_cards(self):
        del self.skat[:]
        for player in self.players:
            del player.cards[:]

    def give_out_cards(self):
        # give cards P1:3, P2:3, P3:3, S:2, P1:4, P2:4, P3:4, P1:3, P2:3, P3:3 --> skat alike
        # P1 gets: 1, 2, 3, 12, 13, 14, 15, 24, 25, 26
        # P2 gets: 4, 5, 6, 16, 17, 18, 19, 27, 28, 29
        # P3 gets: 7, 8, 9, 20, 21, 22, 23, 30, 31, 32
        # Skat gets 10, 11
        round1 = (0, 2)
        round2 = (11, 14)
        round3 = (23, 25)
        for idx, card in enumerate(self.card_deck):
            # Player 1
            if (round1[0] <= idx <= round1[1]) \
                    or (round2[0] <= idx <= round2[1]) \
                    or (round3[0] <= idx <= round3[1]):
                self.players[0].cards.append(card)
            # Player 2
            elif (round1[0] + 3 <= idx <= round1[1] + 3) \
                    or (round2[0] + 4 <= idx <= round2[1] + 4) \
                    or (round3[0] + 3 <= idx <= round3[1] + 3):
                self.players[1].cards.append(card)
            # Player 3
            elif (round1[0] + 6 <= idx <= round1[1] + 6) \
                    or (round2[0] + 8 <= idx <= round2[1] + 8) \
                    or (round3[0] + 6 <= idx <= round3[1] + 6):
                self.players[2].cards.append(card)
            # Skat
            else:
                self.skat.append(card)

game = Game([Player("Player1"), Player("Player2"), Player("Player3")])
game.start_new()

print("\nSkat:")
for card in game.skat:
    print(card)

print("\n" + game.players[0].name + ":")
for card in game.players[0].cards:
    print(card)

print("\n" + game.players[1].name + ":")
for card in game.players[1].cards:
    print(card)

print("\n" + game.players[2].name + ":")
for card in game.players[2].cards:
    print(card)

print("\nDealer=" + game.get_dealer().name)
print("First_Seat=" + game.get_first_seat().name)
print("Second_Seat=" + game.get_second_seat().name)
print("Third_Seat=" + game.get_third_seat().name)<|endoftext|>


<commit_before>import matplotlib.pyplot as pp
import matplotlib as mpl
import numpy as np
from scipy import stats as st
import datetime

import stats

COLORS = ['#08C43A', '#6B2747', 'b']
FONT_SIZE = 18

def MinuteFormatter(x, p):
    return "%02.d:%02.d" % (int(x / 60), x % 60)

def PlotAlt(workout):
    pp.figure()

    alt = workout.GetAltArray()
    dist = workout.GetDistArray()
    #print alt
    pp.plot(dist, alt)
    pp.show()

def PlotPace(workouts):
    pp.figure()
    ax = pp.subplot(111)

    for i,workout in enumerate(workouts):
        pace = workout.GetPace()
        lap_starts = workout.GetLapEnds()
        pp.plot(lap_starts, pace, "o--", markersize=12,
                    mfc=COLORS[i], label=workout.name,
                    color='k')

    pp.ylabel("Run pace (min / km)", fontsize=FONT_SIZE)
    pp.xlabel("Distance (km)", fontsize=FONT_SIZE)

    for tick in ax.xaxis.get_major_ticks():
        tick.label.set_fontsize(FONT_SIZE)
    for tick in ax.yaxis.get_major_ticks():
        tick.label.set_fontsize(FONT_SIZE)

    ax.yaxis.set_major_formatter(mpl.ticker.FuncFormatter(MinuteFormatter))

    pp.grid(color='b')
    pp.ylim([200, 450])
    pp.legend(loc='upper left', fancybox=True, fontsize=FONT_SIZE)

    pp.show()

def PlotPaceVsDistance(workouts):
    nitems = len(workouts)
    dist = np.zeros(nitems)
    pace = np.zeros(nitems)

    recent_dist = []
    recent_pace = []
    # Grab dist and pace data for all and lates runs
    for i, workout in enumerate(workouts):
        dist[i] = workout.GetTotalDist() / 1000.0 # in km
        pace[i] = workout.GetTotalTime() / dist[i]

        if datetime.datetime.now() - workout.GetStartTime() < datetime.timedelta(days=60):
            recent_dist.append(dist[i])
            recent_pace.append(pace[i])

    # Finally, do the scatters
    f = pp.figure()
    ax = pp.subplot(111)
    pp.scatter(dist, pace)
    pp.scatter(recent_dist, recent_pace, color='#08C43A')

    pp.ylabel("Run pace (min / km)")
    pp.xlabel("Run length (km)")

    pp.xlim([0, 47.0])
    pp.ylim([220, 370])

    ax.yaxis.set_major_formatter(mpl.ticker.FuncFormatter(MinuteFormatter))

    # Fit for all runs
    slope, intercept, r_value, p_value, std_err = st.linregress(dist, pace)
    regr_x = np.array([5, 42])
    regr_y = intercept + regr_x * slope
    pp.plot(regr_x, regr_y, '--', color='k')
    pp.annotate("%.2f s/km^2" % slope, (regr_x[1] - 10.0, regr_y[1]))

    # Fit for recent runs
    slope, intercept, r_value, p_value, std_err = st.linregress(recent_dist, recent_pace)
    regr_x = np.array([5, 42])
    regr_y = intercept + regr_x * slope
    pp.plot(regr_x, regr_y, '--', color='#08C43A')
    pp.annotate("%.2f s/km^2" % slope, (regr_x[1] - 10.0, regr_y[1]))

    pp.show()

def PlotDistanceAtPace(workouts):
    paces = []
    dists = []

    recent_dist = []
    recent_pace = []

    for workout in workouts:
        for lap in workout.laps:
            if lap.dist == 0.0:
                continue

            pace = lap.time / (lap.dist / 1000.0)
            paces.append(pace)
            dists.append(lap.dist / 1000.0)

            if datetime.datetime.now() - workout.GetStartTime() < datetime.timedelta(days=60):
                recent_dist.append(lap.dist / 1000.0)
                recent_pace.append(pace)

    pp.figure()
    ax = pp.subplot(111)

    pp.hist(paces, weights=dists, bins=30, range=[180, 480], color='#52730B')

    pp.hist(recent_pace, weights=recent_dist, bins=30, range=[180, 480], color='r')

    ax.xaxis.set_major_formatter(mpl.ticker.FuncFormatter(MinuteFormatter))

    pp.xlabel("Pace (min / km)")
    pp.ylabel("Total distance at pace (km)")

    pp.legend(["Total", "Last 60 days"])

    pp.show()

def PlotMonthlyDist(workouts):
    monthly = {}

    for workout in workouts:
        month = workout.GetStartTime().month
        year = workout.GetStartTime().year

        month = workout.GetStartTime().replace(day=1, hour=0, minute=0, second=0, microsecond=0)
        if not month in monthly:
            monthly[month] = 0.0
        monthly[month] += workout.GetTotalDist() / 1000.0

    pp.figure()
    ax = pp.subplot(111, aspect=0.4)

    n_items = len(monthly.keys())
    sort_ind = sorted(range(n_items), key=lambda i: monthly.keys()[i])
    months =  [monthly.keys()[ind] for ind in sort_ind]
    totals =  [monthly.values()[ind] for ind in sort_ind]

    AV_MONTHS = 3
    totals_av = []
    # get moving average
    for i in range(AV_MONTHS-1, n_items):
        total = 0.0
        for j in range(AV_MONTHS):
            total += totals[i-j]
        totals_av.append(total / AV_MONTHS)

    pp.plot(months, totals, 'o--', markersize=14, color='k', mfc='b')

    pp.plot(months[AV_MONTHS-1:], totals_av, 'o--', color='k', markersize=14, mfc='g')

    pp.legend(["Monthly", "3-month average"], loc='upper left')
    pp.ylim([0, 360])
    pp.xlim([months[0] - datetime.timedelta(days=31), months[-1] + datetime.timedelta(days=31)])

    ax.yaxis.grid(color='b')

    pp.ylabel("Total distance (km)")

    pp.show()
<commit_msg>beautify plots for exporting<commit_after>import matplotlib.pyplot as pp
import matplotlib.dates
import matplotlib as mpl
import numpy as np
from scipy import stats as st
import datetime

import stats

COLORS = ['#08C43A', '#6B2747', 'b']
FONT_SIZE = 18

def MinuteFormatter(x, p):
    return "%02.d:%02.d" % (int(x / 60), x % 60)

def PlotAlt(workout):
    pp.figure()

    alt = workout.GetAltArray()
    dist = workout.GetDistArray()
    #print alt
    pp.plot(dist, alt)
    pp.show()

def PlotPace(workouts):
    pp.figure()
    ax = pp.subplot(111)

    for i,workout in enumerate(workouts):
        pace = workout.GetPace()
        lap_starts = workout.GetLapEnds()
        pp.plot(lap_starts, pace, "o--", markersize=12,
                    mfc=COLORS[i], label=workout.name,
                    color='k')

    pp.ylabel("Run pace (min / km)", fontsize=FONT_SIZE)
    pp.xlabel("Distance (km)", fontsize=FONT_SIZE)

    for tick in ax.xaxis.get_major_ticks():
        tick.label.set_fontsize(FONT_SIZE)
    for tick in ax.yaxis.get_major_ticks():
        tick.label.set_fontsize(FONT_SIZE)

    ax.yaxis.set_major_formatter(mpl.ticker.FuncFormatter(MinuteFormatter))

    pp.grid(color='b')
    #pp.ylim([200, 450])
    pp.legend(loc='upper left', fancybox=True, fontsize=FONT_SIZE)

    pp.show()

def PlotPaceVsDistance(workouts, fname=None):
    nitems = len(workouts)
    dist = np.zeros(nitems)
    pace = np.zeros(nitems)

    recent_dist = []
    recent_pace = []
    # Grab dist and pace data for all and lates runs
    for i, workout in enumerate(workouts):
        dist[i] = workout.GetTotalDist() / 1000.0 # in km
        pace[i] = workout.GetTotalTime() / dist[i]

        if datetime.datetime.now() - workout.GetStartTime() < datetime.timedelta(days=60):
            recent_dist.append(dist[i])
            recent_pace.append(pace[i])

    # Finally, do the scatters
    f = pp.figure()
    ax = pp.subplot(111)
    pp.plot(dist, pace, "o", markersize=6)
    pp.plot(recent_dist, recent_pace, "o", mfc=COLORS[0], markersize=6)

    pp.ylabel("Run pace (min / km)", fontsize=FONT_SIZE)
    pp.xlabel("Run length (km)", fontsize=FONT_SIZE)

    for tick in ax.xaxis.get_major_ticks():
        tick.label.set_fontsize(FONT_SIZE)
    for tick in ax.yaxis.get_major_ticks():
        tick.label.set_fontsize(FONT_SIZE)

    pp.xlim([0, 47.0])
    pp.ylim([220, 370])

    ax.yaxis.set_major_formatter(mpl.ticker.FuncFormatter(MinuteFormatter))

    # Fit for all runs
    slope, intercept, r_value, p_value, std_err = st.linregress(dist, pace)
    regr_x = np.array([5, 42])
    regr_y = intercept + regr_x * slope
    pp.plot(regr_x, regr_y, '--', color='k')
    pp.annotate("%.2f s/km^2" % slope, (regr_x[1] - 10.0, regr_y[1] - 10), fontsize=FONT_SIZE-2)

    # Fit for recent runs
    slope, intercept, r_value, p_value, std_err = st.linregress(recent_dist, recent_pace)
    regr_x = np.array([5, 42])
    regr_y = intercept + regr_x * slope
    pp.plot(regr_x, regr_y, '--', color=COLORS[0])
    pp.annotate("%.2f s/km^2" % slope, (regr_x[1] - 10.0, regr_y[1] - 10), color=COLORS[0], fontsize=FONT_SIZE-2)

    ax.yaxis.grid(color='b')
    pp.legend(["Total", "Last 60 days"], loc='lower right', fancybox=True, fontsize=FONT_SIZE)

    pp.tight_layout()
    if fname:
        pp.savefig(fname)
    pp.show()

def PlotDistanceAtPace(workouts, fname=None):
    paces = []
    dists = []

    recent_dist = []
    recent_pace = []

    for workout in workouts:
        for lap in workout.laps:
            if lap.dist == 0.0:
                continue

            pace = lap.time / (lap.dist / 1000.0)
            paces.append(pace)
            dists.append(lap.dist / 1000.0)

            if datetime.datetime.now() - workout.GetStartTime() < datetime.timedelta(days=60):
                recent_dist.append(lap.dist / 1000.0)
                recent_pace.append(pace)

    pp.figure()
    ax = pp.subplot(111)

    pp.hist(paces, weights=dists, bins=30, range=[180, 480], color='#52730B')

    pp.hist(recent_pace, weights=recent_dist, bins=30, range=[180, 480], color='r')

    ax.xaxis.set_major_formatter(mpl.ticker.FuncFormatter(MinuteFormatter))

    pp.xlabel("Pace (min / km)", fontsize=FONT_SIZE)
    pp.ylabel("Total distance at pace (km)", fontsize=FONT_SIZE)

    for tick in ax.xaxis.get_major_ticks():
        tick.label.set_fontsize(FONT_SIZE)
    for tick in ax.yaxis.get_major_ticks():
        tick.label.set_fontsize(FONT_SIZE)

    pp.legend(["Total", "Last 60 days"], loc='upper right', fancybox=True, fontsize=FONT_SIZE)
    pp.grid(color='b')

    pp.tight_layout()
    if fname:
        pp.savefig(fname)
    pp.show()

def PlotMonthlyDist(workouts, fname=None):
    monthly = {}

    for workout in workouts:
        month = workout.GetStartTime().month
        year = workout.GetStartTime().year

        month = workout.GetStartTime().replace(day=1, hour=0, minute=0, second=0, microsecond=0)
        if not month in monthly:
            monthly[month] = 0.0
        monthly[month] += workout.GetTotalDist() / 1000.0

    pp.figure()
    ax = pp.subplot(111, aspect=0.4)

    n_items = len(monthly.keys())
    sort_ind = sorted(range(n_items), key=lambda i: monthly.keys()[i])
    months =  [monthly.keys()[ind] for ind in sort_ind]
    totals =  [monthly.values()[ind] for ind in sort_ind]

    AV_MONTHS = 3
    totals_av = []
    # get moving average
    for i in range(AV_MONTHS-1, n_items):
        total = 0.0
        for j in range(AV_MONTHS):
            total += totals[i-j]
        totals_av.append(total / AV_MONTHS)

    pp.plot(months, totals, 'o--', markersize=14, color='k', mfc='b')

    pp.plot(months[AV_MONTHS-1:], totals_av, 'o--', color='k', markersize=14, mfc='g')

    pp.legend(["Monthly", "3-month average"], loc='upper left', fancybox=True, fontsize=FONT_SIZE)
    pp.xlim([months[0] - datetime.timedelta(days=14), months[-1] + datetime.timedelta(days=14)])

    ax.yaxis.grid(color='b')

    pp.ylabel("Total distance (km)", fontsize=FONT_SIZE)

    locator = mpl.dates.MonthLocator()
    formatter = mpl.dates.AutoDateFormatter(locator)
    formatter.scaled[30.0] = '%b' # only show month
    ax.xaxis.set_major_locator(locator)
    ax.xaxis.set_major_formatter(formatter)

    for tick in ax.xaxis.get_major_ticks():
        tick.label.set_fontsize(FONT_SIZE)
    for tick in ax.yaxis.get_major_ticks():
        tick.label.set_fontsize(FONT_SIZE)

    pp.tight_layout()
    if fname:
        pp.savefig(fname)
    pp.show()
<|endoftext|>


<commit_before>from __future__ import division
import os
import thread
import subprocess
import numpy as np

### python RSA.py --wppl BDA --perception human multimodal_fc6 multimodal_conv42 multimodal_pool1 --pragmatics combined S0 --production cost nocost --split_type balancedavg1 balancedavg2 balancedavg3 balancedavg4 balancedavg5
### python RSA.py --wppl BDA --perception human multimodal_fc6 multimodal_conv42 multimodal_pool1 --pragmatics combined S0 --production cost nocost --split_type balancedavg1
### python RSA.py --wppl BDA --perception multimodal_conv42 --pragmatics S0 --production nocost --split_type balancedavg1 balancedavg2 balancedavg3 balancedavg4 balancedavg5
### python RSA.py --wppl evaluate --perception human --pragmatics combined --production cost --split_type balancedavg5
### python RSA.py --wppl evaluate --perception multimodal_fc6 --pragmatics combined --production cost --split_type balancedavg3
### python RSA.py --wppl evaluate --perception multimodal_conv42 --pragmatics combined --production cost --split_type balancedavg1 balancedavg2 balancedavg3 balancedavg4 balancedavg5
### python RSA.py --wppl evaluate --perception multimodal_fc6 --pragmatics S0 --production cost --split_type balancedavg1 balancedavg2 balancedavg3 balancedavg4 balancedavg5
### python RSA.py --wppl evaluate --perception multimodal_fc6 --pragmatics combined --production nocost --split_type balancedavg1 balancedavg2 balancedavg3 balancedavg4 balancedavg5
### python RSA.py --wppl BDA-enumerate --sim_scaling_lb 1 --sim_scaling_ub 200 --step_size 2 --split_type balancedavg5
### python RSA.py --wppl AIS --perception multimodal_fc6 --production cost --pragmatics combined --num_ais_samples 2 --split_type balancedavg1

def run_bda(perception, pragmatics, production, split_type):
    if not os.path.exists('./bdaOutput'):
        os.makedirs('./bdaOutput')
    if not os.path.exists('./bdaOutput/{}_{}'.format(perception,split_type)):
        os.makedirs('./bdaOutput/{}_{}'.format(perception,split_type))
        os.makedirs('./bdaOutput/{}_{}/raw'.format(perception,split_type))
    # check to make sure we do not already have output
    if not os.path.exists('./bdaOutput/{}_{}/raw/{}_{}_{}_{}Params.csv'.format(perception,\
                                                                                split_type,\
                                                                                perception,\
                                                                                pragmatics,\
                                                                                production,\
                                                                                split_type)):
        #sample: models/bdaOutput/human_balancedavg1/raw/human_combined_cost_balancedavg1Params.csv
        cmd_string = 'webppl BDA.wppl --require ./refModule/ -- --perception {} --pragmatics {} --production {} --splitType {}'.format(perception, pragmatics, production, split_type)
        print 'Running: {}'.format(cmd_string)
        thread.start_new_thread(os.system,(cmd_string,))
    else:
        print 'Already have BDA output for model {} {} {} {}. Not proceeding unless files moved/renamed.'.format(perception,pragmatics,production,split_type)

def run_bda_enumerate(simScaling, split_type):
    if not os.path.exists('./enumerateOutput'):
        os.makedirs('./enumerateOutput')
    if not os.path.exists(os.path.join('./enumerateOutput',split_type)):
        os.makedirs(os.path.join('./enumerateOutput',split_type))
    cmd_string = 'webppl BDA-enumerate.wppl --require ./refModule/ -- --simScaling {} --splitType {}'.format(simScaling, split_type)
    print 'Running: {}'.format(cmd_string)
    thread.start_new_thread(os.system,(cmd_string,))

def run_evaluate(perception, pragmatics, production, split_type):
    if not os.path.exists('./evaluateOutput'):
        os.makedirs('./evaluateOutput')
    if not os.path.exists('./evaluateOutput/{}_{}_{}_{}'.format(perception,\
                                                                pragmatics,\
                                                                production,\
                                                                split_type)):

        cmd_string = 'webppl evaluate.wppl --require ./refModule/ -- --paramSetting {}_{}_{} --adaptorType {} --splitType {}'.format(perception, pragmatics, production, perception, split_type)
        print 'Running: {}'.format(cmd_string)
        thread.start_new_thread(os.system,(cmd_string,))
    else:
        print 'Already have evaluation output for model {} {} {} {}. Not proceeding unless files moved/renamed.'.format(perception,pragmatics,production,split_type)

def run_ais(perception, pragmatics, production, split_type, num_samp):
    if not os.path.exists('./aisOutput'):
        os.makedirs('./aisOutput')
    cmd_string = 'webppl BF.wppl --require ./refModule/ -- --perception {} --pragmatics {} --production {} --splitType {}'.format(perception, pragmatics, production, split_type)
    print '{} | Running: {}'.format(num_samp,cmd_string)
    thread.start_new_thread(os.system,(cmd_string,))

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--wppl', type=str, help='options: BDA | BDA-enumerate | evaluate | AIS', default='BDA')
    parser.add_argument('--perception', nargs='+', type=str, \
                        help='option: options: human| multimodal_conv42 | multimodal_pool1 | multimodal_fc6',\
                        default = 'multimodal_conv42')
    parser.add_argument('--pragmatics', nargs='+', type=str, \
                        help='option: combined | S1 | S0',\
                        default = 'combined')
    parser.add_argument('--production', nargs='+', type=str, \
                        help='option: cost | nocost',\
                        default = 'cost')
    parser.add_argument('--split_type', nargs='+', type=str, \
                        help='option: splitbyobject | alldata | balancedavg',\
                        default = 'balancedavg')
    parser.add_argument('--sim_scaling_lb', type=float, \
                        help='for BDA-enumerate only: this is the LOWER bound for the simScaling param. \
                              We will sweep through values from (sim_scaling_lb, sim_scaling_ub) in step_size sized steps',\
                        default = 1.0)
    parser.add_argument('--sim_scaling_ub', type=float, \
                        help='for BDA-enumerate only: this is the UPPER bound for the simScaling param. \
                              We will sweep through values from (sim_scaling_lb, sim_scaling_ub) in step_size sized steps',\
                        default = 200.0)
    parser.add_argument('--step_size', type=float, \
                        help='for BDA-enumerate only: this is the step size we will use to march through \
                              the simScaling range',\
                        default = 2.0)
    parser.add_argument('--num_ais_samples', type=int, \
                        help='how many AIS samples do you want to take in parallel?',
                        default = 1)

    args = parser.parse_args()
    print args.split_type
    perception = args.perception
    production = args.production
    pragmatics = args.pragmatics
    split_type = args.split_type
    lb = args.sim_scaling_lb
    ub = args.sim_scaling_ub
    if lb==ub: ## to get last value in range, make sure that np.arange has an interval of non-zero length to work with
        ub = ub + 2
    step_size = args.step_size

    assert args.wppl in ['BDA','evaluate', 'BDA-enumerate', 'AIS']

    ## first run BDA-enumerate.wppl
    if 'BDA-enumerate' in args.wppl:
        ss_range = np.arange(lb,ub,step_size)
        for i,ss in enumerate(ss_range):
            for split in split_type:
                run_bda_enumerate(ss,split)

    ## first run BDA.wppl
    elif 'BDA' in args.wppl:
        for perc in perception:
            for prag in pragmatics:
                for prod in production:
                    for split in split_type:
                        run_bda(perc,prag,prod,split)

    elif 'evaluate' in args.wppl:
        ## then on output, run evaluate.wppl
        for perc in perception:
            print perc
            for prag in pragmatics:
                for prod in production:
                    for split in split_type:
                        run_evaluate(perc,prag,prod,split)

    elif 'AIS' in args.wppl:
        for perc in perception:
            for prag in pragmatics:
                for prod in production:
                    for split in split_type:
                        for num_samp in np.arange(args.num_ais_samples):
                            run_ais(perc,prag,prod,split,num_samp)

    else:
        print '{} wppl command not recognized'.format(args.wppl)
<commit_msg>add flattening func to master rsa script<commit_after>from __future__ import division
import os
import thread
import subprocess
import numpy as np
import analysis_helpers as h

### python RSA.py --wppl BDA --perception human multimodal_fc6 multimodal_conv42 multimodal_pool1 --pragmatics combined S0 --production cost nocost --split_type balancedavg1 balancedavg2 balancedavg3 balancedavg4 balancedavg5
### python RSA.py --wppl BDA --perception human multimodal_fc6 multimodal_conv42 multimodal_pool1 --pragmatics combined S0 --production cost nocost --split_type balancedavg1
### python RSA.py --wppl BDA --perception multimodal_conv42 --pragmatics S0 --production nocost --split_type balancedavg1 balancedavg2 balancedavg3 balancedavg4 balancedavg5
### python RSA.py --wppl evaluate --perception human --pragmatics combined --production cost --split_type balancedavg1
### python RSA.py --wppl evaluate --perception multimodal_fc6 --pragmatics combined --production cost --split_type balancedavg3
### python RSA.py --wppl evaluate --perception multimodal_conv42 --pragmatics combined --production cost --split_type balancedavg1 balancedavg2 balancedavg3 balancedavg4 balancedavg5
### python RSA.py --wppl evaluate --perception multimodal_fc6 --pragmatics S0 --production cost --split_type balancedavg1 balancedavg2 balancedavg3 balancedavg4 balancedavg5
### python RSA.py --wppl evaluate --perception multimodal_fc6 --pragmatics combined --production nocost --split_type balancedavg1 balancedavg2 balancedavg3 balancedavg4 balancedavg5
### python RSA.py --wppl BDA-enumerate --sim_scaling_lb 1 --sim_scaling_ub 200 --step_size 2 --split_type balancedavg5
### python RSA.py --wppl AIS --perception multimodal_fc6 --production cost --pragmatics combined --num_ais_samples 2 --split_type balancedavg1

def run_bda(perception, pragmatics, production, split_type):
    if not os.path.exists('./bdaOutput'):
        os.makedirs('./bdaOutput')
    if not os.path.exists('./bdaOutput/{}_{}'.format(perception,split_type)):
        os.makedirs('./bdaOutput/{}_{}'.format(perception,split_type))
        os.makedirs('./bdaOutput/{}_{}/raw'.format(perception,split_type))
    # check to make sure we do not already have output
    if not os.path.exists('./bdaOutput/{}_{}/raw/{}_{}_{}_{}Params.csv'.format(perception,\
                                                                                split_type,\
                                                                                perception,\
                                                                                pragmatics,\
                                                                                production,\
                                                                                split_type)):
        #sample: models/bdaOutput/human_balancedavg1/raw/human_combined_cost_balancedavg1Params.csv
        cmd_string = 'webppl BDA.wppl --require ./refModule/ -- --perception {} --pragmatics {} --production {} --splitType {}'.format(perception, pragmatics, production, split_type)
        print 'Running: {}'.format(cmd_string)
        thread.start_new_thread(os.system,(cmd_string,))
    else:
        print 'Already have BDA output for model {} {} {} {}. Not proceeding unless files moved/renamed.'.format(perception,pragmatics,production,split_type)

def flatten_bda_output(adaptor_types = ['multimodal_pool1','multimodal_conv42','multimodal_fc6', 'human'],
                        verbosity=1)
    h.flatten_param_posterior(adaptor_types = adaptor_types,
                              verbosity=verbosity)

def run_bda_enumerate(simScaling, split_type):
    if not os.path.exists('./enumerateOutput'):
        os.makedirs('./enumerateOutput')
    if not os.path.exists(os.path.join('./enumerateOutput',split_type)):
        os.makedirs(os.path.join('./enumerateOutput',split_type))
    cmd_string = 'webppl BDA-enumerate.wppl --require ./refModule/ -- --simScaling {} --splitType {}'.format(simScaling, split_type)
    print 'Running: {}'.format(cmd_string)
    thread.start_new_thread(os.system,(cmd_string,))

def run_evaluate(perception, pragmatics, production, split_type):
    if not os.path.exists('./evaluateOutput'):
        os.makedirs('./evaluateOutput')
    out_dir = './evaluateOutput/{}_{}_{}_{}'.format(perception,\
                                                    pragmatics,\
                                                    production,\
                                                    split_type)
    if not os.path.exists(out_dir):
        # os.makedirs(out_dir)
        cmd_string = 'webppl evaluate.wppl --require ./refModule/ -- --paramSetting {}_{}_{} --adaptorType {} --splitType {}'.format(perception, pragmatics, production, perception, split_type)
        print 'Running: {}'.format(cmd_string)
        thread.start_new_thread(os.system,(cmd_string,))
    else:
        print 'Already have evaluation output for model {} {} {} {}. Not proceeding unless files moved/renamed.'.format(perception,pragmatics,production,split_type)

def run_ais(perception, pragmatics, production, split_type, num_samp):
    if not os.path.exists('./aisOutput'):
        os.makedirs('./aisOutput')
    cmd_string = 'webppl BF.wppl --require ./refModule/ -- --perception {} --pragmatics {} --production {} --splitType {}'.format(perception, pragmatics, production, split_type)
    print '{} | Running: {}'.format(num_samp,cmd_string)
    thread.start_new_thread(os.system,(cmd_string,))

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--wppl', type=str, help='options: BDA | BDA-enumerate | evaluate | AIS', default='BDA')
    parser.add_argument('--perception', nargs='+', type=str, \
                        help='option: options: human| multimodal_conv42 | multimodal_pool1 | multimodal_fc6',\
                        default = 'multimodal_conv42')
    parser.add_argument('--pragmatics', nargs='+', type=str, \
                        help='option: combined | S1 | S0',\
                        default = 'combined')
    parser.add_argument('--production', nargs='+', type=str, \
                        help='option: cost | nocost',\
                        default = 'cost')
    parser.add_argument('--split_type', nargs='+', type=str, \
                        help='option: splitbyobject | alldata | balancedavg',\
                        default = 'balancedavg')
    parser.add_argument('--sim_scaling_lb', type=float, \
                        help='for BDA-enumerate only: this is the LOWER bound for the simScaling param. \
                              We will sweep through values from (sim_scaling_lb, sim_scaling_ub) in step_size sized steps',\
                        default = 1.0)
    parser.add_argument('--sim_scaling_ub', type=float, \
                        help='for BDA-enumerate only: this is the UPPER bound for the simScaling param. \
                              We will sweep through values from (sim_scaling_lb, sim_scaling_ub) in step_size sized steps',\
                        default = 200.0)
    parser.add_argument('--step_size', type=float, \
                        help='for BDA-enumerate only: this is the step size we will use to march through \
                              the simScaling range',\
                        default = 2.0)
    parser.add_argument('--num_ais_samples', type=int, \
                        help='how many AIS samples do you want to take in parallel?',
                        default = 1)

    args = parser.parse_args()
    print args.split_type
    perception = args.perception
    production = args.production
    pragmatics = args.pragmatics
    split_type = args.split_type
    lb = args.sim_scaling_lb
    ub = args.sim_scaling_ub
    if lb==ub: ## to get last value in range, make sure that np.arange has an interval of non-zero length to work with
        ub = ub + 2
    step_size = args.step_size

    assert args.wppl in ['BDA','evaluate', 'BDA-enumerate', 'AIS']

    ## first run BDA-enumerate.wppl
    if 'BDA-enumerate' in args.wppl:
        ss_range = np.arange(lb,ub,step_size)
        for i,ss in enumerate(ss_range):
            for split in split_type:
                run_bda_enumerate(ss,split)

    ## first run BDA.wppl
    elif 'BDA' in args.wppl:
        for perc in perception:
            for prag in pragmatics:
                for prod in production:
                    for split in split_type:
                        run_bda(perc,prag,prod,split)

    elif 'evaluate' in args.wppl:
        ## then on output, run evaluate.wppl
        for perc in perception:
            print perc
            for prag in pragmatics:
                for prod in production:
                    for split in split_type:
                        run_evaluate(perc,prag,prod,split)

    elif 'AIS' in args.wppl:
        for perc in perception:
            for prag in pragmatics:
                for prod in production:
                    for split in split_type:
                        for num_samp in np.arange(args.num_ais_samples):
                            run_ais(perc,prag,prod,split,num_samp)

    else:
        print '{} wppl command not recognized'.format(args.wppl)
<|endoftext|>


<commit_before>import sys
import yaml
from . import _logger

class YAMLLoader(object):
    def __init__(self, **kwargs):
        kwargs.setdefault('log_level', 'warning')
        for k, v in kwargs.iteritems():
            setattr(self, k, v)

    @staticmethod
    def read_yaml(file_name):
        try:
            _logger.debug("Try to open yaml file %s" % file_name)
            config_file = open(file_name)
            _logger.debug("Try to load content of %s" % file_name)
            config = yaml.load(config_file)
            config_file.close()
            return config
        except IOError as e:
            _logger.error("Could not find or open yaml file %s" % file_name)
            _logger.error(e.strerror)
            raise e
        except yaml.YAMLError as e:
            _logger.error("Failed to load yaml file %s" % file_name)
            raise e
        except Exception as e:
            _logger.error("Unexpected error for %s" % file_name)
            raise e

class YAMLDumper(object):
    def __init__(self, **kwargs):
        kwargs.setdefault('log_level', 'warning')
        for k, v in kwargs.iteritems():
            setattr(self, k, v)

    @staticmethod
    def dump_yaml(data, file_name):
        try:
            _logger.debug("Try to open file %s" % file_name)
            out_file = open(file_name, "w+")
            _logger.debug("Try to dump data to file %s" % file_name)
            yaml.dump(data, out_file)
            out_file.close()
        except Exception as e:
            _logger.error("Failed to dump data to file %s" % file_name)
            raise e<commit_msg>add accept_none option for reader<commit_after>import sys
import yaml
from . import _logger

class YAMLLoader(object):
    def __init__(self, **kwargs):
        kwargs.setdefault('log_level', 'warning')
        for k, v in kwargs.iteritems():
            setattr(self, k, v)

    @staticmethod
    def read_yaml(file_name, accept_none=False):
        if accept_none and file_name is None:
            return None
        try:
            _logger.debug("Try to open yaml file %s" % file_name)
            config_file = open(file_name)
            _logger.debug("Try to load content of %s" % file_name)
            config = yaml.load(config_file)
            config_file.close()
            return config
        except IOError as e:
            _logger.error("Could not find or open yaml file %s" % file_name)
            _logger.error(e.strerror)
            raise e
        except yaml.YAMLError as e:
            _logger.error("Failed to load yaml file %s" % file_name)
            raise e
        except Exception as e:
            _logger.error("Unexpected error for %s" % file_name)
            raise e

class YAMLDumper(object):
    def __init__(self, **kwargs):
        kwargs.setdefault('log_level', 'warning')
        for k, v in kwargs.iteritems():
            setattr(self, k, v)

    @staticmethod
    def dump_yaml(data, file_name):
        try:
            _logger.debug("Try to open file %s" % file_name)
            out_file = open(file_name, "w+")
            _logger.debug("Try to dump data to file %s" % file_name)
            yaml.dump(data, out_file)
            out_file.close()
        except Exception as e:
            _logger.error("Failed to dump data to file %s" % file_name)
            raise e<|endoftext|>


<commit_before>import click

from nsone.config import ConfigException
from ns1cli import __version__
from ns1cli.cli import cli as root_cli
from ns1cli.util import Formatter

class HelpFormatter(Formatter):

    def print_config(self, config):
        try:
            click.secho('Current Key: %s' % config.getCurrentKeyID(), bold=True)
            self.pretty_print(config.getKeyConfig())
        except ConfigException as e:
            pass

        self.out(config)

def solve(s):
    s = s.split('\n', 1)[-1]
    if s.find('\n') == -1:
        return s
    return s.rsplit('\n', 1)[0]

@click.command('help',
               short_help='displays help for a sequence of commands')
@click.argument('subcommands', required=False, nargs=-1)
@click.pass_context
def cli(ctx, subcommands):
    """Displays usage/help for a given sequence of commands.

    \b
    Examples:
        ns1 help record
        ns1 help record answer
        ns1 help record region
        ns1 help record answer add
    """
    cmd = root_cli
    for cmd_name in subcommands:
        cmd = cmd.get_command(ctx, cmd_name)

    if not cmd:
        raise click.BadParameter('Unknown command sequence')

    help_text = cmd.get_help(ctx)
    temp = solve(help_text)

    formatter = ctx.make_formatter()
    formatter.write_heading('NS1Cli v%s' % __version__)
    formatter.indent()
    formatter.write(help_text)
    help_text = formatter.getvalue()
    click.echo_via_pager(help_text)
<commit_msg>Removes help command(-h/--help option still exists)<commit_after><|endoftext|>